[llvm] [mlir] [mlir][spirv] Integrate `convert-to-spirv` into `mlir-vulkan-runner` (PR #106082)

Angel Zhang via llvm-commits llvm-commits at lists.llvm.org
Mon Aug 26 08:32:19 PDT 2024


https://github.com/angelz913 updated https://github.com/llvm/llvm-project/pull/106082

>From 7e2d3046a60a849b60bd6a087336e8c76715ba1e Mon Sep 17 00:00:00 2001
From: Angel Zhang <angel.zhang at amd.com>
Date: Sat, 24 Aug 2024 03:10:36 +0000
Subject: [PATCH 1/3] [mlir][spirv] Integrate convert-to-spirv into
 mlir-vulkan-runner

---
 mlir/include/mlir/Conversion/Passes.td        |  5 +-
 .../Conversion/ConvertToSPIRV/CMakeLists.txt  |  1 +
 .../ConvertToSPIRV/ConvertToSPIRVPass.cpp     | 81 +++++++++++++------
 .../mlir-vulkan-runner/mlir-vulkan-runner.cpp |  5 +-
 .../llvm-project-overlay/mlir/BUILD.bazel     |  1 +
 5 files changed, 66 insertions(+), 27 deletions(-)

diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 7bde9e490e4f4e..244827ade66be6 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -50,7 +50,10 @@ def ConvertToSPIRVPass : Pass<"convert-to-spirv"> {
     "Run function signature conversion to convert vector types">,
     Option<"runVectorUnrolling", "run-vector-unrolling", "bool",
     /*default=*/"true",
-    "Run vector unrolling to convert vector types in function bodies">
+    "Run vector unrolling to convert vector types in function bodies">,
+    Option<"runOnGPUModules", "run-on-gpu-modules", "bool",
+    /*default=*/"false",
+    "Clone and convert only the GPU modules for integration testing">
   ];
 }
 
diff --git a/mlir/lib/Conversion/ConvertToSPIRV/CMakeLists.txt b/mlir/lib/Conversion/ConvertToSPIRV/CMakeLists.txt
index 863ef9603da385..124a4c453e75c5 100644
--- a/mlir/lib/Conversion/ConvertToSPIRV/CMakeLists.txt
+++ b/mlir/lib/Conversion/ConvertToSPIRV/CMakeLists.txt
@@ -15,6 +15,7 @@ add_mlir_conversion_library(MLIRConvertToSPIRVPass
   MLIRArithToSPIRV
   MLIRArithTransforms
   MLIRFuncToSPIRV
+  MLIRGPUDialect
   MLIRGPUToSPIRV
   MLIRIndexToSPIRV
   MLIRIR
diff --git a/mlir/lib/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.cpp b/mlir/lib/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.cpp
index 9e57b923ea6894..f624999decac48 100644
--- a/mlir/lib/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.cpp
+++ b/mlir/lib/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.cpp
@@ -16,6 +16,7 @@
 #include "mlir/Conversion/UBToSPIRV/UBToSPIRV.h"
 #include "mlir/Conversion/VectorToSPIRV/VectorToSPIRV.h"
 #include "mlir/Dialect/Arith/Transforms/Passes.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
@@ -40,6 +41,35 @@ using namespace mlir;
 
 namespace {
 
+/// Map memRef memory space to SPIR-V storage class.
+void mapToMemRef(Operation *op, spirv::TargetEnvAttr &targetAttr) {
+  spirv::TargetEnv targetEnv(targetAttr);
+  bool targetEnvSupportsKernelCapability =
+      targetEnv.allows(spirv::Capability::Kernel);
+  spirv::MemorySpaceToStorageClassMap memorySpaceMap =
+      targetEnvSupportsKernelCapability
+          ? spirv::mapMemorySpaceToOpenCLStorageClass
+          : spirv::mapMemorySpaceToVulkanStorageClass;
+  spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap);
+  spirv::convertMemRefTypesAndAttrs(op, converter);
+}
+
+/// Populate patterns for each dialect.
+void populateConvertToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
+                                    ScfToSPIRVContext &scfToSPIRVContext,
+                                    RewritePatternSet &patterns) {
+  arith::populateCeilFloorDivExpandOpsPatterns(patterns);
+  arith::populateArithToSPIRVPatterns(typeConverter, patterns);
+  populateBuiltinFuncToSPIRVPatterns(typeConverter, patterns);
+  populateFuncToSPIRVPatterns(typeConverter, patterns);
+  populateGPUToSPIRVPatterns(typeConverter, patterns);
+  index::populateIndexToSPIRVPatterns(typeConverter, patterns);
+  populateMemRefToSPIRVPatterns(typeConverter, patterns);
+  populateVectorToSPIRVPatterns(typeConverter, patterns);
+  populateSCFToSPIRVPatterns(typeConverter, scfToSPIRVContext, patterns);
+  ub::populateUBToSPIRVConversionPatterns(typeConverter, patterns);
+}
+
 /// A pass to perform the SPIR-V conversion.
 struct ConvertToSPIRVPass final
     : impl::ConvertToSPIRVPassBase<ConvertToSPIRVPass> {
@@ -64,31 +94,32 @@ struct ConvertToSPIRVPass final
     RewritePatternSet patterns(context);
     ScfToSPIRVContext scfToSPIRVContext;
 
-    // Map MemRef memory space to SPIR-V storage class.
-    spirv::TargetEnv targetEnv(targetAttr);
-    bool targetEnvSupportsKernelCapability =
-        targetEnv.allows(spirv::Capability::Kernel);
-    spirv::MemorySpaceToStorageClassMap memorySpaceMap =
-        targetEnvSupportsKernelCapability
-            ? spirv::mapMemorySpaceToOpenCLStorageClass
-            : spirv::mapMemorySpaceToVulkanStorageClass;
-    spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap);
-    spirv::convertMemRefTypesAndAttrs(op, converter);
-
-    // Populate patterns for each dialect.
-    arith::populateCeilFloorDivExpandOpsPatterns(patterns);
-    arith::populateArithToSPIRVPatterns(typeConverter, patterns);
-    populateBuiltinFuncToSPIRVPatterns(typeConverter, patterns);
-    populateFuncToSPIRVPatterns(typeConverter, patterns);
-    populateGPUToSPIRVPatterns(typeConverter, patterns);
-    index::populateIndexToSPIRVPatterns(typeConverter, patterns);
-    populateMemRefToSPIRVPatterns(typeConverter, patterns);
-    populateVectorToSPIRVPatterns(typeConverter, patterns);
-    populateSCFToSPIRVPatterns(typeConverter, scfToSPIRVContext, patterns);
-    ub::populateUBToSPIRVConversionPatterns(typeConverter, patterns);
-
-    if (failed(applyPartialConversion(op, *target, std::move(patterns))))
-      return signalPassFailure();
+    if (runOnGPUModules) {
+      SmallVector<Operation *, 1> gpuModules;
+      OpBuilder builder(context);
+      op->walk([&](gpu::GPUModuleOp gpuModule) {
+        builder.setInsertionPoint(gpuModule);
+        gpuModules.push_back(builder.clone(*gpuModule));
+      });
+      // Run conversion for each module independently as they can have
+      // different TargetEnv attributes.
+      for (Operation *gpuModule : gpuModules) {
+        spirv::TargetEnvAttr targetAttr =
+            spirv::lookupTargetEnvOrDefault(gpuModule);
+        mapToMemRef(gpuModule, targetAttr);
+        populateConvertToSPIRVPatterns(typeConverter, scfToSPIRVContext,
+                                       patterns);
+        if (failed(
+                applyFullConversion(gpuModule, *target, std::move(patterns))))
+          return signalPassFailure();
+      }
+    } else {
+      mapToMemRef(op, targetAttr);
+      populateConvertToSPIRVPatterns(typeConverter, scfToSPIRVContext,
+                                     patterns);
+      if (failed(applyPartialConversion(op, *target, std::move(patterns))))
+        return signalPassFailure();
+    }
   }
 };
 
diff --git a/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp b/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp
index 032f5760361f4b..93ca922c57084f 100644
--- a/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp
+++ b/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp
@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.h"
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
@@ -64,7 +65,9 @@ static LogicalResult runMLIRPasses(Operation *op,
   passManager.addPass(createGpuKernelOutliningPass());
   passManager.addPass(memref::createFoldMemRefAliasOpsPass());
 
-  passManager.addPass(createConvertGPUToSPIRVPass(/*mapMemorySpace=*/true));
+  ConvertToSPIRVPassOptions convertToSPIRVOptions{};
+  convertToSPIRVOptions.runOnGPUModules = true;
+  passManager.addPass(createConvertToSPIRVPass(convertToSPIRVOptions));
   OpPassManager &modulePM = passManager.nest<spirv::ModuleOp>();
   modulePM.addPass(spirv::createSPIRVLowerABIAttributesPass());
   modulePM.addPass(spirv::createSPIRVUpdateVCEPass());
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index de069daf603f1e..35b3ef47603d99 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -8401,6 +8401,7 @@ cc_library(
         ":ArithTransforms",
         ":ConversionPassIncGen",
         ":FuncToSPIRV",
+        ":GPUDialect",
         ":GPUToSPIRV",
         ":IR",
         ":IndexToSPIRV",

>From 4d11fe0b4215d11e2b40ad79432495476e49c28a Mon Sep 17 00:00:00 2001
From: Angel Zhang <anzhouzhang913 at gmail.com>
Date: Mon, 26 Aug 2024 11:12:23 -0400
Subject: [PATCH 2/3] Update mlir/include/mlir/Conversion/Passes.td

Co-authored-by: Jakub Kuderski <kubakuderski at gmail.com>
---
 mlir/include/mlir/Conversion/Passes.td | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 244827ade66be6..e6190c525bd6a0 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -53,7 +53,7 @@ def ConvertToSPIRVPass : Pass<"convert-to-spirv"> {
     "Run vector unrolling to convert vector types in function bodies">,
     Option<"runOnGPUModules", "run-on-gpu-modules", "bool",
     /*default=*/"false",
-    "Clone and convert only the GPU modules for integration testing">
+    "Clone and convert GPU modules">
   ];
 }
 

>From 668478a5806792e83c995c50abfafba15f368c47 Mon Sep 17 00:00:00 2001
From: Angel Zhang <angel.zhang at amd.com>
Date: Mon, 26 Aug 2024 15:32:06 +0000
Subject: [PATCH 3/3] Code refactoring and comments

---
 mlir/include/mlir/Conversion/Passes.td        |  2 +-
 .../ConvertToSPIRV/ConvertToSPIRVPass.cpp     | 42 ++++++++++---------
 .../mlir-vulkan-runner/mlir-vulkan-runner.cpp |  2 +-
 3 files changed, 24 insertions(+), 22 deletions(-)

diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index e6190c525bd6a0..f43d8aa08aadde 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -51,7 +51,7 @@ def ConvertToSPIRVPass : Pass<"convert-to-spirv"> {
     Option<"runVectorUnrolling", "run-vector-unrolling", "bool",
     /*default=*/"true",
     "Run vector unrolling to convert vector types in function bodies">,
-    Option<"runOnGPUModules", "run-on-gpu-modules", "bool",
+    Option<"convertGPUModules", "convert-gpu-modules", "bool",
     /*default=*/"false",
     "Clone and convert GPU modules">
   ];
diff --git a/mlir/lib/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.cpp b/mlir/lib/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.cpp
index f624999decac48..9d6885182c5bd9 100644
--- a/mlir/lib/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.cpp
+++ b/mlir/lib/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.cpp
@@ -94,31 +94,33 @@ struct ConvertToSPIRVPass final
     RewritePatternSet patterns(context);
     ScfToSPIRVContext scfToSPIRVContext;
 
-    if (runOnGPUModules) {
-      SmallVector<Operation *, 1> gpuModules;
-      OpBuilder builder(context);
-      op->walk([&](gpu::GPUModuleOp gpuModule) {
-        builder.setInsertionPoint(gpuModule);
-        gpuModules.push_back(builder.clone(*gpuModule));
-      });
-      // Run conversion for each module independently as they can have
-      // different TargetEnv attributes.
-      for (Operation *gpuModule : gpuModules) {
-        spirv::TargetEnvAttr targetAttr =
-            spirv::lookupTargetEnvOrDefault(gpuModule);
-        mapToMemRef(gpuModule, targetAttr);
-        populateConvertToSPIRVPatterns(typeConverter, scfToSPIRVContext,
-                                       patterns);
-        if (failed(
-                applyFullConversion(gpuModule, *target, std::move(patterns))))
-          return signalPassFailure();
-      }
-    } else {
+    if (!convertGPUModules) {
       mapToMemRef(op, targetAttr);
       populateConvertToSPIRVPatterns(typeConverter, scfToSPIRVContext,
                                      patterns);
       if (failed(applyPartialConversion(op, *target, std::move(patterns))))
         return signalPassFailure();
+      return;
+    }
+
+    // Clone each GPU kernel module for conversion, given that the GPU
+    // launch op still needs the original GPU kernel module.
+    SmallVector<Operation *, 1> gpuModules;
+    OpBuilder builder(context);
+    op->walk([&](gpu::GPUModuleOp gpuModule) {
+      builder.setInsertionPoint(gpuModule);
+      gpuModules.push_back(builder.clone(*gpuModule));
+    });
+    // Run conversion for each module independently as they can have
+    // different TargetEnv attributes.
+    for (Operation *gpuModule : gpuModules) {
+      spirv::TargetEnvAttr targetAttr =
+          spirv::lookupTargetEnvOrDefault(gpuModule);
+      mapToMemRef(gpuModule, targetAttr);
+      populateConvertToSPIRVPatterns(typeConverter, scfToSPIRVContext,
+                                     patterns);
+      if (failed(applyFullConversion(gpuModule, *target, std::move(patterns))))
+        return signalPassFailure();
     }
   }
 };
diff --git a/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp b/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp
index 93ca922c57084f..2dd539ef83481f 100644
--- a/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp
+++ b/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp
@@ -66,7 +66,7 @@ static LogicalResult runMLIRPasses(Operation *op,
   passManager.addPass(memref::createFoldMemRefAliasOpsPass());
 
   ConvertToSPIRVPassOptions convertToSPIRVOptions{};
-  convertToSPIRVOptions.runOnGPUModules = true;
+  convertToSPIRVOptions.convertGPUModules = true;
   passManager.addPass(createConvertToSPIRVPass(convertToSPIRVOptions));
   OpPassManager &modulePM = passManager.nest<spirv::ModuleOp>();
   modulePM.addPass(spirv::createSPIRVLowerABIAttributesPass());



More information about the llvm-commits mailing list