[Mlir-commits] [mlir] [mlir][SPIRV] Fix lookup logic `spirv.target_env` for `gpu.module` (PR #147262)
Jaeho Kim
llvmlistbot at llvm.org
Wed Jul 9 19:11:32 PDT 2025
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/6] [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/6] 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/6] 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>(
>From fa3b44e7027b91bffb97deb2172be0d2f574abe4 Mon Sep 17 00:00:00 2001
From: oojahooo <oojahooo at gmail.com>
Date: Thu, 10 Jul 2025 11:03:37 +0900
Subject: [PATCH 4/6] fixup! [mlir][SPIRV] Fix lookup logic `spirv.target_env`
for `gpu.module`
Add lookup target env in "targets" attr logic to GPUToSPIRV pass
---
.../Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp | 27 ++++++++++++++++---
mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp | 7 -----
2 files changed, 23 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
index 0b2c06a08db2d..d182198b82f47 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
@@ -48,9 +48,29 @@ struct GPUToSPIRVPass final : impl::ConvertGPUToSPIRVBase<GPUToSPIRVPass> {
void runOnOperation() override;
private:
+ spirv::TargetEnvAttr lookupTargetEnvInTargets(gpu::GPUModuleOp moduleOp);
+ spirv::TargetEnvAttr lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp);
bool mapMemorySpace;
};
+spirv::TargetEnvAttr
+GPUToSPIRVPass::lookupTargetEnvInTargets(gpu::GPUModuleOp moduleOp) {
+ for (auto &targetAttr : moduleOp.getTargetsAttr())
+ if (auto spirvTargetEnvAttr =
+ llvm::dyn_cast<spirv::TargetEnvAttr>(targetAttr))
+ return spirvTargetEnvAttr;
+
+ return {};
+}
+
+spirv::TargetEnvAttr
+GPUToSPIRVPass::lookupTargetEnvOrDefault(gpu::GPUModuleOp moduleOp) {
+ if (auto targetEnvAttr = lookupTargetEnvInTargets(moduleOp))
+ return targetEnvAttr;
+
+ return spirv::lookupTargetEnvOrDefault(moduleOp);
+}
+
void GPUToSPIRVPass::runOnOperation() {
MLIRContext *context = &getContext();
ModuleOp module = getOperation();
@@ -58,9 +78,8 @@ void GPUToSPIRVPass::runOnOperation() {
SmallVector<Operation *, 1> gpuModules;
OpBuilder builder(context);
- auto targetEnvSupportsKernelCapability = [](gpu::GPUModuleOp moduleOp) {
- Operation *gpuModule = moduleOp.getOperation();
- auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule);
+ auto targetEnvSupportsKernelCapability = [this](gpu::GPUModuleOp moduleOp) {
+ auto targetAttr = lookupTargetEnvOrDefault(moduleOp);
spirv::TargetEnv targetEnv(targetAttr);
return targetEnv.allows(spirv::Capability::Kernel);
};
@@ -86,7 +105,7 @@ void GPUToSPIRVPass::runOnOperation() {
// TargetEnv attributes.
for (Operation *gpuModule : gpuModules) {
spirv::TargetEnvAttr targetAttr =
- spirv::lookupTargetEnvOrDefault(gpuModule);
+ lookupTargetEnvOrDefault(llvm::cast<gpu::GPUModuleOp>(gpuModule));
// Map MemRef memory space to SPIR-V storage class first if requested.
if (mapMemorySpace) {
diff --git a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
index c90db80b7b970..5ecbd5d7c59d5 100644
--- a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
@@ -184,13 +184,6 @@ 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 7d43c590b32fcd45d17239ba0d4c763609fd9e4d Mon Sep 17 00:00:00 2001
From: oojahooo <oojahooo at gmail.com>
Date: Thu, 10 Jul 2025 11:05:11 +0900
Subject: [PATCH 5/6] Remove unnecessary check in test
---
mlir/test/Conversion/GPUToSPIRV/lookup-target-env.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Conversion/GPUToSPIRV/lookup-target-env.mlir b/mlir/test/Conversion/GPUToSPIRV/lookup-target-env.mlir
index 7c74bbad49101..a6a8e36a8642e 100644
--- a/mlir/test/Conversion/GPUToSPIRV/lookup-target-env.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/lookup-target-env.mlir
@@ -1,7 +1,7 @@
// 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]> {
+ // CHECK-LABEL: spirv.module @{{.*}} GLSL450
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)>})
>From 40b6f07634f3a34fc0bc21bad0fc647fd52b3ab1 Mon Sep 17 00:00:00 2001
From: oojahooo <oojahooo at gmail.com>
Date: Thu, 10 Jul 2025 11:09:03 +0900
Subject: [PATCH 6/6] Add same lookup target env logic in GPUModuleConversion
---
mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
index b99ed261ecfa3..a2b86c35ea99d 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
@@ -385,6 +385,10 @@ LogicalResult GPUModuleConversion::matchAndRewrite(
if (auto attr = moduleOp->getAttrOfType<spirv::TargetEnvAttr>(
spirv::getTargetEnvAttrName()))
spvModule->setAttr(spirv::getTargetEnvAttrName(), attr);
+ for (auto targetAttr : moduleOp.getTargetsAttr())
+ if (auto spirvTargetEnvAttr =
+ llvm::dyn_cast<spirv::TargetEnvAttr>(targetAttr))
+ spvModule->setAttr(spirv::getTargetEnvAttrName(), spirvTargetEnvAttr);
rewriter.eraseOp(moduleOp);
return success();
More information about the Mlir-commits
mailing list