set-spirv-api-attrs: Walk and find gpu module. (#1100)

Author: silee2

include/imex/Transforms/Passes.td

@@ -65,7 +65,7 @@ def SetSPIRVCapabilities : Pass<"set-spirv-capabilities"> {
   ];
 }
 
-def SetSPIRVAbiAttribute : Pass<"set-spirv-abi-attrs", "::mlir::gpu::GPUModuleOp"> {
+def SetSPIRVAbiAttribute : Pass<"set-spirv-abi-attrs"> {
   let summary = "Sets Spirv Abi attribute";
   let constructor = "imex::createSetSPIRVAbiAttributePass()";
   let dependentDialects = ["::mlir::gpu::GPUDialect",

lib/Transforms/SetSPIRVAbiAttribute.cpp

@@ -53,31 +53,34 @@ class SetSPIRVAbiAttributePass
   }
 
   void runOnOperation() override {
-    auto gpuModule = getOperation();
     auto *context = &getContext();
     auto attrName =
         mlir::StringAttr::get(context, mlir::spirv::getEntryPointABIAttrName());
-    if (m_clientAPI == "opencl") {
-      auto abi = mlir::spirv::getEntryPointABIAttr(context);
-      for (const auto &gpuFunc : gpuModule.getOps<mlir::gpu::GPUFuncOp>()) {
-        if (!mlir::gpu::GPUDialect::isKernel(gpuFunc) ||
-            gpuFunc->getAttr(attrName))
-          continue;
+    auto op = getOperation();
+    mlir::StringRef clientAPI = m_clientAPI;
+    op->walk([&](mlir::gpu::GPUModuleOp gpuModule) {
+      if (clientAPI == "opencl") {
+        auto abi = mlir::spirv::getEntryPointABIAttr(context);
+        for (const auto &gpuFunc : gpuModule.getOps<mlir::gpu::GPUFuncOp>()) {
+          if (!mlir::gpu::GPUDialect::isKernel(gpuFunc) ||
+              gpuFunc->getAttr(attrName))
+            continue;
 
-        gpuFunc->setAttr("VectorComputeFunctionINTEL",
-                         mlir::UnitAttr::get(context));
-        gpuFunc->setAttr(attrName, abi);
-      }
-    } else if (m_clientAPI == "vulkan") {
-      auto abi = mlir::spirv::getEntryPointABIAttr(context, {1, 1, 1});
-      for (const auto &gpuFunc : gpuModule.getOps<mlir::gpu::GPUFuncOp>()) {
-        if (!mlir::gpu::GPUDialect::isKernel(gpuFunc) ||
-            gpuFunc->getAttr(attrName))
-          continue;
+          gpuFunc->setAttr("VectorComputeFunctionINTEL",
+                           mlir::UnitAttr::get(context));
+          gpuFunc->setAttr(attrName, abi);
+        }
+      } else if (clientAPI == "vulkan") {
+        auto abi = mlir::spirv::getEntryPointABIAttr(context, {1, 1, 1});
+        for (const auto &gpuFunc : gpuModule.getOps<mlir::gpu::GPUFuncOp>()) {
+          if (!mlir::gpu::GPUDialect::isKernel(gpuFunc) ||
+              gpuFunc->getAttr(attrName))
+            continue;
 
-        gpuFunc->setAttr(attrName, abi);
+          gpuFunc->setAttr(attrName, abi);
+        }
       }
-    }
+    });
   }
 
 private:

test/Transforms/set-spirv-abi-attr.mlir

@@ -1,12 +1,12 @@
-// RUN: imex-opt --set-spirv-abi-attrs='client-api=opencl' %s | FileCheck %s --check-prefix=OPENCL
-// RUN: imex-opt --set-spirv-abi-attrs='client-api=vulkan' %s | FileCheck %s --check-prefix=VULKAN
+// RUN: imex-opt --split-input-file --set-spirv-abi-attrs='client-api=opencl' %s | FileCheck %s --check-prefix=OPENCL
+// RUN: imex-opt --split-input-file --set-spirv-abi-attrs='client-api=vulkan' %s | FileCheck %s --check-prefix=VULKAN
 
 gpu.module @main_kernel {
   gpu.func @main_kernel(%arg0: memref<8xf32>, %arg1: memref<8xf32>, %arg2: memref<8xf32>) kernel {
-
-  // OPENCL: gpu.func @main_kernel(%arg0: memref<8xf32>, %arg1: memref<8xf32>, %arg2: memref<8xf32>) kernel attributes {VectorComputeFunctionINTEL, spirv.entry_point_abi = #spirv.entry_point_abi<>} {
-  // VULKAN: gpu.func @main_kernel(%arg0: memref<8xf32>, %arg1: memref<8xf32>, %arg2: memref<8xf32>) kernel attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
-
+  // OPENCL-LABEL: gpu.func @main_kernel(%arg0: memref<8xf32>, %arg1: memref<8xf32>, %arg2: memref<8xf32>) kernel attributes
+  // OPENCL-SAME:    {VectorComputeFunctionINTEL, spirv.entry_point_abi = #spirv.entry_point_abi<>} {
+  // VULKAN-LABEL: gpu.func @main_kernel(%arg0: memref<8xf32>, %arg1: memref<8xf32>, %arg2: memref<8xf32>) kernel attributes
+  // VULKAN-SAME:    {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
     cf.br ^bb1
   ^bb1:  // pred: ^bb0
     %0 = gpu.block_id  x
@@ -17,3 +17,26 @@ gpu.module @main_kernel {
     gpu.return
   }
 }
+
+// -----
+
+module {
+  module attributes {gpu.container_module} {
+    func.func @run(%arg0: memref<4096x4096xf16>) -> memref<4096x4096xf16> attributes {llvm.emit_c_interface} {
+      %c1 = arith.constant 1 : index
+      gpu.launch_func  @run_kernel::@run_kernel blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1)  args(%arg0 : memref<4096x4096xf16>)
+      return %arg0 : memref<4096x4096xf16>
+    }
+    gpu.module @run_kernel {
+      // OPENCL-LABEL: gpu.func @run_kernel(%arg0: memref<4096x4096xf16>) kernel attributes
+      // OPENCL-SAME:    {VectorComputeFunctionINTEL, known_block_size = array<i32: 1, 1, 1>, known_grid_size = array<i32: 1, 1, 1>,
+      // OPENCL-SAME:     spirv.entry_point_abi = #spirv.entry_point_abi<>} {
+      // VULKAN-LABEL: gpu.func @run_kernel(%arg0: memref<4096x4096xf16>) kernel attributes
+      // VULKAN-SAME:    {known_block_size = array<i32: 1, 1, 1>, known_grid_size = array<i32: 1, 1, 1>,
+      // VULKAN-SAME:     spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
+      gpu.func @run_kernel(%arg0: memref<4096x4096xf16>) kernel attributes {known_block_size = array<i32: 1, 1, 1>, known_grid_size = array<i32: 1, 1, 1>} {
+      gpu.return
+      }
+    }
+  }
+}