[LLVMCPU] Populate ArmSVE to LLVM conversion patterns (iree-org#17665)

The ArmSME dialect will soon make use of some ArmSVE operations within its lowerings, so these need to be populated. Signed-off-by: Benjamin Maxwell <[email protected]>
nod-ai · Jun 14, 2024 · 3428231 · 3428231
1 parent 045bf32
commit 3428231
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 1 deletion.
diff --git a/compiler/plugins/target/LLVMCPU/BUILD.bazel b/compiler/plugins/target/LLVMCPU/BUILD.bazel
@@ -64,6 +64,8 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:ArmNeonDialect",
         "@llvm-project//mlir:ArmSMEDialect",
         "@llvm-project//mlir:ArmSMEToLLVMIRTranslation",
+        "@llvm-project//mlir:ArmSVEDialect",
+        "@llvm-project//mlir:ArmSVEToLLVMIRTranslation",
         "@llvm-project//mlir:BuiltinToLLVMIRTranslation",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LLVMDialect",

diff --git a/compiler/plugins/target/LLVMCPU/CMakeLists.txt b/compiler/plugins/target/LLVMCPU/CMakeLists.txt
@@ -42,6 +42,8 @@ iree_cc_library(
     MLIRArmNeonDialect
     MLIRArmSMEDialect
     MLIRArmSMEToLLVMIRTranslation
+    MLIRArmSVEDialect
+    MLIRArmSVEToLLVMIRTranslation
     MLIRBuiltinToLLVMIRTranslation
     MLIRIR
     MLIRLLVMDialect

diff --git a/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp b/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
@@ -34,13 +34,15 @@
 #include "llvm/Support/TargetSelect.h"
 #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
 #include "mlir/Dialect/ArmSME/IR/ArmSME.h"
+#include "mlir/Dialect/ArmSVE/IR/ArmSVEDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/DialectResourceBlobManager.h"
 #include "mlir/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
@@ -212,6 +214,7 @@ class LLVMCPUTargetBackend final : public TargetBackend {
     mlir::registerBuiltinDialectTranslation(registry);
     mlir::registerLLVMDialectTranslation(registry);
     mlir::registerArmSMEDialectTranslation(registry);
+    mlir::registerArmSVEDialectTranslation(registry);
     // TODO: make inclusion of ArmNeon conditional?
     // clang-format off
     registry.insert<IREE::Codegen::IREECodegenDialect,
@@ -220,7 +223,8 @@ class LLVMCPUTargetBackend final : public TargetBackend {
                     pdl::PDLDialect,
                     pdl_interp::PDLInterpDialect,
                     arm_neon::ArmNeonDialect,
-                    arm_sme::ArmSMEDialect>();
+                    arm_sme::ArmSMEDialect,
+                    arm_sve::ArmSVEDialect>();
     // clang-format on
   }
 

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel
@@ -128,6 +128,8 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:ArmSMEToLLVMIRTranslation",
         "@llvm-project//mlir:ArmSMEToSCF",
         "@llvm-project//mlir:ArmSMETransforms",
+        "@llvm-project//mlir:ArmSVEDialect",
+        "@llvm-project//mlir:ArmSVETransforms",
         "@llvm-project//mlir:BufferizationDialect",
         "@llvm-project//mlir:ComplexToLLVM",
         "@llvm-project//mlir:ComplexToStandard",

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt
@@ -102,6 +102,8 @@ iree_cc_library(
     MLIRArmSMEToLLVMIRTranslation
     MLIRArmSMEToSCF
     MLIRArmSMETransforms
+    MLIRArmSVEDialect
+    MLIRArmSVETransforms
     MLIRBufferizationDialect
     MLIRComplexToLLVM
     MLIRComplexToStandard

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/ConvertToLLVM.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/ConvertToLLVM.cpp
@@ -40,6 +40,8 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
+#include "mlir/Dialect/ArmSVE/IR/ArmSVEDialect.h"
+#include "mlir/Dialect/ArmSVE/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/Passes.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
@@ -1050,6 +1052,12 @@ void ConvertToLLVMPass::runOnOperation() {
       typeConverter, patterns, targetReassociateFpReductions.getValue());
   populateReconcileUnrealizedCastsPatterns(patterns);
 
+  if (isAArch64(targetAttr) &&
+      (hasAnySVEFeature(targetAttr) || hasSMEFeature(targetAttr))) {
+    populateArmSVELegalizeForLLVMExportPatterns(typeConverter, patterns);
+    configureArmSVELegalizeForExportTarget(target);
+  }
+
   HALDispatchABI abi(&typeConverter);
   // clang-format off
   patterns.insert<