[mlir][spirv] Integrate `convert-to-spirv` into `mlir-vulkan-runner` #106082

obtuseangleAZ · 2024-08-26T14:29:33Z

Description
This PR adds a new option for convert-to-spirv pass to clone and convert only GPU kernel modules for integration testing. The reason for using pass options instead of two separate passes is that they both consist of memref types conversion and individual dialect patterns, except they run on different scopes. The PR also replaces the gpu-to-spirv pass with the convert-to-spirv pass (with the new option) in mlir-vulkan-runner.

Future Plan
Use nesting pass pipelines in mlir-vulkan-runner instead of adding this option.

llvmbot · 2024-08-26T14:30:07Z

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir-spirv

Author: Angel Zhang (angelz913)

Changes

This PR adds a new option for convert-to-spirv pass to clone and convert only GPU kernel modules for integration testing. The PR also replaces the gpu-to-spirv pass with the convert-to-spirv pass (with the new option) in mlir-vulkan-runner.

Full diff: https://github.com/llvm/llvm-project/pull/106082.diff

5 Files Affected:

(modified) mlir/include/mlir/Conversion/Passes.td (+4-1)
(modified) mlir/lib/Conversion/ConvertToSPIRV/CMakeLists.txt (+1)
(modified) mlir/lib/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.cpp (+56-25)
(modified) mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp (+4-1)
(modified) utils/bazel/llvm-project-overlay/mlir/BUILD.bazel (+1)

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",

mlir/include/mlir/Conversion/Passes.td

mlir/lib/Conversion/ConvertToSPIRV/ConvertToSPIRVPass.cpp

mlir/include/mlir/Conversion/Passes.td

Co-authored-by: Jakub Kuderski <kubakuderski@gmail.com>

kuhar

Could we add a couple of LIt tests for this? Right now the testing depends on the mlir runner, but most folks probably have it disabled and only run cpu tests in their daily development flow.

obtuseangleAZ · 2024-08-26T20:06:00Z

Could we add a couple of LIt tests for this? Right now the testing depends on the mlir runner, but most folks probably have it disabled and only run cpu tests in their daily development flow.

Added

github-actions · 2024-08-26T20:09:07Z

✅ With the latest revision this PR passed the C/C++ code formatter.

kuhar

Thanks. Having to use a new flag is suboptimal, but overall this seems like a big improvement over using the gpu-to-spirv pass that accumulated unrelated patterns.

joker-eph · 2024-08-29T18:36:53Z

Use nesting pass pipelines in mlir-vulkan-runner instead of adding this option.

I'm slightly concerned about seeing any "future plans" for mlir-vulkan-runner, the status is that this runner should be removed entirely (just like we removed to Cuda runner) to align the testing on mlir-cpu-runner.

It seems to me that it would be more fruitful to concentrate the efforts to actually remove it instead, it's a long-time pending TODO.

obtuseangleAZ requested review from rupprecht, keith, antiagainst and kuhar as code owners August 26, 2024 14:29

llvmbot added mlir:core MLIR Core Infrastructure mlir:spirv mlir bazel "Peripheral" support tier build system: utils/bazel labels Aug 26, 2024

kuhar reviewed Aug 26, 2024

View reviewed changes

obtuseangleAZ and others added 3 commits August 26, 2024 15:52

[mlir][spirv] Integrate convert-to-spirv into mlir-vulkan-runner

410aab8

Update mlir/include/mlir/Conversion/Passes.td

2359648

Co-authored-by: Jakub Kuderski <kubakuderski@gmail.com>

Code refactoring and comments

67f9799

kuhar reviewed Aug 26, 2024

View reviewed changes

obtuseangleAZ added 2 commits August 26, 2024 20:04

Fix bug

ec0f3f3

Add LIT tests

e1c87f6

obtuseangleAZ force-pushed the convert-to-spirv-e2e branch from 668478a to e1c87f6 Compare August 26, 2024 20:05

Code formatting

b42b9dd

kuhar approved these changes Aug 27, 2024

View reviewed changes

Hardcode84 approved these changes Aug 27, 2024

View reviewed changes

kuhar merged commit 2bf2468 into llvm:main Aug 27, 2024
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][spirv] Integrate `convert-to-spirv` into `mlir-vulkan-runner` #106082

[mlir][spirv] Integrate `convert-to-spirv` into `mlir-vulkan-runner` #106082

Uh oh!

obtuseangleAZ commented Aug 26, 2024 •

edited

Loading

Uh oh!

llvmbot commented Aug 26, 2024 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kuhar left a comment

Uh oh!

obtuseangleAZ commented Aug 26, 2024

Uh oh!

github-actions bot commented Aug 26, 2024 •

edited

Loading

Uh oh!

kuhar left a comment

Uh oh!

Uh oh!

joker-eph commented Aug 29, 2024

Uh oh!

Uh oh!

[mlir][spirv] Integrate convert-to-spirv into mlir-vulkan-runner #106082

[mlir][spirv] Integrate convert-to-spirv into mlir-vulkan-runner #106082

Uh oh!

Conversation

obtuseangleAZ commented Aug 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Aug 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

Uh oh!

obtuseangleAZ commented Aug 26, 2024

Uh oh!

github-actions bot commented Aug 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

joker-eph commented Aug 29, 2024

Uh oh!

Uh oh!

[mlir][spirv] Integrate `convert-to-spirv` into `mlir-vulkan-runner` #106082

[mlir][spirv] Integrate `convert-to-spirv` into `mlir-vulkan-runner` #106082

obtuseangleAZ commented Aug 26, 2024 •

edited

Loading

llvmbot commented Aug 26, 2024 •

edited

Loading

github-actions bot commented Aug 26, 2024 •

edited

Loading