[Mlir-commits] [mlir] [MLIR] Update convert-gpu-to-spirv pass to prepare using GPU compilat… (PR #69941)

Sang Ik Lee llvmlistbot at llvm.org
Mon Oct 23 09:35:51 PDT 2023


https://github.com/silee2 created https://github.com/llvm/llvm-project/pull/69941

…ion pipeline for OpenCL kernels.

This includes a couple of changes to pass behavior for OpenCL kernels. Vulkan shaders are not impacted by the changes.

1. SPIRV module is placed inside GPU module. This change is required for gpu-module-to-binary to work correctly as it expects kernel function to be inside the GPU module.
2. A dummy func.func with same kernel name as gpu.func is created. GPU compilation pipeline defers lowering of gpu launch kernel op. Since spirv.func is not directly tied to gpu launch kernel, a dummy func.func is required to avoid legalization issues.
3. Use correct mapping when mapping MemRef memory space to SPIR-V storage class for OpenCL kernels.

>From 7ca3f97b5ee6e5cefd94afd3b090d0dba2120cea Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 23 Oct 2023 16:25:15 +0000
Subject: [PATCH] [MLIR] Update convert-gpu-to-spirv pass to prepare using GPU
 compilation pipeline for OpenCL kernels.

This includes a couple of changes to pass behavior for OpenCL kernels.
Vulkan shaders are not impacted by the changes.

1. SPIRV module is placed inside GPU module. This change is required for
gpu-module-to-binary to work correctly as it expects kernel function to be
inside the GPU module.
2. A dummy func.func with same kernel name as gpu.func is created.
GPU compilation pipeline defers lowering of gpu launch kernel op.
Since spirv.func is not directly tied to gpu launch kernel,
a dummy func.func is required to avoid legalization issues.
3. Use correct mapping when mapping MemRef memory space to SPIR-V storage class for OpenCL kernels.
---
 mlir/include/mlir/Conversion/Passes.td        |  5 +-
 .../Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp  | 76 ++++++++++++++++---
 .../Conversion/GPUToSPIRV/module-opencl.mlir  |  4 +
 3 files changed, 75 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 274784fe4a7b29c..652ef5ad95158ca 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -578,7 +578,10 @@ def ConvertGPUToSPIRV : Pass<"convert-gpu-to-spirv", "ModuleOp"> {
     to control the set and binding if wanted.
   }];
   let constructor = "mlir::createConvertGPUToSPIRVPass()";
-  let dependentDialects = ["spirv::SPIRVDialect"];
+  let dependentDialects = [
+    "spirv::SPIRVDialect",
+    "func::FuncDialect",
+  ];
   let options = [
     Option<"use64bitIndex", "use-64bit-index",
            "bool", /*default=*/"false",
diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
index 272e3de8723aeb6..35ee0d7038a2c9a 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
@@ -17,6 +17,7 @@
 #include "mlir/Conversion/FuncToSPIRV/FuncToSPIRV.h"
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h"
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
@@ -54,22 +55,63 @@ void GPUToSPIRVPass::runOnOperation() {
 
   SmallVector<Operation *, 1> gpuModules;
   OpBuilder builder(context);
+
+  auto getTargetEnvFromGPUModuleOp = [=](gpu::GPUModuleOp moduleOp) {
+    Operation *gpuModule = moduleOp.getOperation();
+    auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule);
+    std::unique_ptr<ConversionTarget> target =
+        SPIRVConversionTarget::get(targetAttr);
+
+    SPIRVConversionOptions options;
+    options.use64bitIndex = this->use64bitIndex;
+    SPIRVTypeConverter typeConverter(targetAttr, options);
+    const spirv::TargetEnv &targetEnv = typeConverter.getTargetEnv();
+    return targetEnv;
+  };
+
   module.walk([&](gpu::GPUModuleOp moduleOp) {
     // Clone each GPU kernel module for conversion, given that the GPU
     // launch op still needs the original GPU kernel module.
-    builder.setInsertionPoint(moduleOp.getOperation());
+    // SPIRV module insertion point by is after original GPU module.
+    // This works fine for Vulkan shader that has a dedicated runner.
+    // But OpenCL kernel needs SPIRV module placed inside original GPU module as
+    // OpenCL uses GPU compilation pipeline.
+    auto targetEnv = getTargetEnvFromGPUModuleOp(moduleOp);
+    FailureOr<spirv::MemoryModel> memoryModel =
+        spirv::getMemoryModel(targetEnv);
+    if (failed(memoryModel))
+      return signalPassFailure();
+    (memoryModel == spirv::MemoryModel::OpenCL)
+        ? builder.setInsertionPoint(moduleOp.getBody(),
+                                    moduleOp.getBody()->begin())
+        : builder.setInsertionPoint(moduleOp.getOperation());
     gpuModules.push_back(builder.clone(*moduleOp.getOperation()));
   });
 
   // Run conversion for each module independently as they can have different
   // TargetEnv attributes.
   for (Operation *gpuModule : gpuModules) {
+    auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule);
+    std::unique_ptr<ConversionTarget> target =
+        SPIRVConversionTarget::get(targetAttr);
+
+    SPIRVConversionOptions options;
+    options.use64bitIndex = this->use64bitIndex;
+    SPIRVTypeConverter typeConverter(targetAttr, options);
+    const spirv::TargetEnv &targetEnv = typeConverter.getTargetEnv();
+    FailureOr<spirv::MemoryModel> memoryModel =
+        spirv::getMemoryModel(targetEnv);
+    if (failed(memoryModel))
+      return signalPassFailure();
+
     // Map MemRef memory space to SPIR-V storage class first if requested.
     if (mapMemorySpace) {
       std::unique_ptr<ConversionTarget> target =
           spirv::getMemorySpaceToStorageClassTarget(*context);
       spirv::MemorySpaceToStorageClassMap memorySpaceMap =
-          spirv::mapMemorySpaceToVulkanStorageClass;
+          (memoryModel == spirv::MemoryModel::OpenCL)
+              ? spirv::mapMemorySpaceToOpenCLStorageClass
+              : spirv::mapMemorySpaceToVulkanStorageClass;
       spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap);
 
       RewritePatternSet patterns(context);
@@ -79,13 +121,6 @@ void GPUToSPIRVPass::runOnOperation() {
         return signalPassFailure();
     }
 
-    auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule);
-    std::unique_ptr<ConversionTarget> target =
-        SPIRVConversionTarget::get(targetAttr);
-
-    SPIRVConversionOptions options;
-    options.use64bitIndex = this->use64bitIndex;
-    SPIRVTypeConverter typeConverter(targetAttr, options);
     populateMMAToSPIRVCoopMatrixTypeConversion(typeConverter,
                                                this->useCoopMatrixNV);
 
@@ -108,6 +143,29 @@ void GPUToSPIRVPass::runOnOperation() {
     if (failed(applyFullConversion(gpuModule, *target, std::move(patterns))))
       return signalPassFailure();
   }
+  // In case of OpenCL, gpu.func in original gpu.module needs to replaced with
+  // an empty func.func with same arguments as gpu.func. And it also needs
+  // gpu.kernel attribute set.
+  module.walk([&](gpu::GPUModuleOp moduleOp) {
+    auto targetEnv = getTargetEnvFromGPUModuleOp(moduleOp);
+    FailureOr<spirv::MemoryModel> memoryModel =
+        spirv::getMemoryModel(targetEnv);
+    if (failed(memoryModel))
+      return signalPassFailure();
+    if (memoryModel == spirv::MemoryModel::OpenCL) {
+      moduleOp.walk([&](gpu::GPUFuncOp funcOp) {
+        builder.setInsertionPoint(funcOp);
+        auto newFuncOp = builder.create<func::FuncOp>(
+            funcOp.getLoc(), funcOp.getName(), funcOp.getFunctionType());
+        auto entryBlock = newFuncOp.addEntryBlock();
+        builder.setInsertionPointToEnd(entryBlock);
+        builder.create<func::ReturnOp>(funcOp.getLoc());
+        newFuncOp->setAttr(gpu::GPUDialect::getKernelFuncAttrName(),
+                           builder.getUnitAttr());
+        funcOp.erase();
+      });
+    }
+  });
 }
 
 } // namespace
diff --git a/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir b/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir
index 0aa50cc1e25294d..4b8d17cd6449389 100644
--- a/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/module-opencl.mlir
@@ -12,6 +12,8 @@ module attributes {
     //  CHECK-SAME:     {{%.*}}: !spirv.ptr<!spirv.array<12 x f32>, CrossWorkgroup>
     //   CHECK-NOT:     spirv.interface_var_abi
     //  CHECK-SAME:     spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
+    // CHECK-LABEL:   func.func @basic_module_structure
+    //  CHECK-SAME:     attributes {gpu.kernel}
     gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, #spirv.storage_class<CrossWorkgroup>>) kernel
         attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       gpu.return
@@ -45,6 +47,8 @@ module attributes {
     //  CHECK-SAME:     {{%.*}}: !spirv.ptr<!spirv.array<12 x f32>, CrossWorkgroup>
     //   CHECK-NOT:     spirv.interface_var_abi
     //  CHECK-SAME:     spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>
+    // CHECK-LABEL:   func.func @basic_module_structure
+    //  CHECK-SAME:     attributes {gpu.kernel}
     gpu.func @basic_module_structure(%arg0 : f32, %arg1 : memref<12xf32, #spirv.storage_class<CrossWorkgroup>>) kernel
         attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
       gpu.return



More information about the Mlir-commits mailing list