[Mlir-commits] [mlir] [mlir][SPIRV] Fix lookup logic `spirv.target_env` for `gpu.module` (PR #147262)

[Jaeho Kim]

https://github.com/oojahooo updated https://github.com/llvm/llvm-project/pull/147262

>From 5d05b68214f9166dae933b6313852fa8f079e155 Mon Sep 17 00:00:00 2001
From: oojahooo <oojahooo at gmail.com>
Date: Mon, 7 Jul 2025 18:15:11 +0900
Subject: [PATCH 1/3] [mlir][SPIRV] Fix lookup logic `spirv.target_env` for
 `gpu.module`

The `gpu.module` operation can contain `spirv.target_env` attributes
within an array attribute named `"targets"`. So it accounts for that
case by iterating over the `"targets"` attribute, if present, and
looking up `spirv.target_env`.
---
 mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
index 5ecbd5d7c59d5..dbaa10e89bd42 100644
--- a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
@@ -184,6 +184,15 @@ spirv::TargetEnvAttr spirv::lookupTargetEnv(Operation *op) {
     if (!op)
       break;
 
+    if (auto arrAttr = op->getAttrOfType<ArrayAttr>("targets")) {
+      for (auto attr : arrAttr) {
+        if (auto spirvTargetEnvAttr =
+                llvm::dyn_cast<spirv::TargetEnvAttr>(attr)) {
+          return spirvTargetEnvAttr;
+        }
+      }
+    }
+
     if (auto attr = op->getAttrOfType<spirv::TargetEnvAttr>(
             spirv::getTargetEnvAttrName()))
       return attr;

>From 71d3d92b442a332b20187e2c01698b062033a993 Mon Sep 17 00:00:00 2001
From: oojahooo <oojahooo at gmail.com>
Date: Tue, 8 Jul 2025 18:24:35 +0900
Subject: [PATCH 2/3] Add test

---
 .../GPUToSPIRV/lookup-target-env.mlir           | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)
 create mode 100644 mlir/test/Conversion/GPUToSPIRV/lookup-target-env.mlir

diff --git a/mlir/test/Conversion/GPUToSPIRV/lookup-target-env.mlir b/mlir/test/Conversion/GPUToSPIRV/lookup-target-env.mlir
new file mode 100644
index 0000000000000..7c74bbad49101
--- /dev/null
+++ b/mlir/test/Conversion/GPUToSPIRV/lookup-target-env.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt --spirv-attach-target='caps=Shader exts=SPV_KHR_storage_buffer_storage_class' --convert-gpu-to-spirv %s -o - | FileCheck %s
+
+module attributes {gpu.container_module} {
+  // CHECK-LABEL: spirv.module @{{.*}} GLSL450 requires #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]> {
+  gpu.module @kernels {
+    // CHECK: spirv.func @load_kernel
+    // CHECK-SAME: %[[ARG:.*]]: !spirv.ptr<!spirv.struct<(!spirv.array<48 x f32, stride=4> [0])>, StorageBuffer> {spirv.interface_var_abi = #spirv.interface_var_abi<(0, 0)>})
+    gpu.func @load_kernel(%arg0: memref<12x4xf32>) kernel attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+      %c0 = arith.constant 0 : index
+      // CHECK:             %[[PTR:.*]] = spirv.AccessChain %[[ARG]]{{\[}}{{%.*}}, {{%.*}}{{\]}}
+      // CHECK-NEXT:        {{%.*}} = spirv.Load "StorageBuffer" %[[PTR]] : f32
+      %0 = memref.load %arg0[%c0, %c0] : memref<12x4xf32>
+      // CHECK:             spirv.Return
+      gpu.return
+    }
+  }
+}

>From ae6ba9cc1395906b28b849260fdd71024c31bab8 Mon Sep 17 00:00:00 2001
From: oojahooo <oojahooo at gmail.com>
Date: Tue, 8 Jul 2025 18:33:15 +0900
Subject: [PATCH 3/3] Delete braces for complying with conding standard

It resolves
https://github.com/llvm/llvm-project/pull/147262#discussion_r2189731690
---
 mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
index dbaa10e89bd42..c90db80b7b970 100644
--- a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
@@ -185,12 +185,10 @@ spirv::TargetEnvAttr spirv::lookupTargetEnv(Operation *op) {
       break;
 
     if (auto arrAttr = op->getAttrOfType<ArrayAttr>("targets")) {
-      for (auto attr : arrAttr) {
+      for (auto attr : arrAttr)
         if (auto spirvTargetEnvAttr =
-                llvm::dyn_cast<spirv::TargetEnvAttr>(attr)) {
+                llvm::dyn_cast<spirv::TargetEnvAttr>(attr))
           return spirvTargetEnvAttr;
-        }
-      }
     }
 
     if (auto attr = op->getAttrOfType<spirv::TargetEnvAttr>(