[Mlir-commits] [mlir] [mlir][gpu] Avoid duplicate attrs when lowering gpu.func (PR #205035)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu

@llvm/pr-subscribers-mlir

Author: lijinpei-amd

<details>
<summary>Changes</summary>

Fixes: https://github.com/llvm/llvm-project/issues/204586

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


2 Files Affected:

- (modified) mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp (+18-14) 
- (modified) mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir (+20) 


``````````diff
diff --git a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
index eeb90a7ff8150..c40b7c02739f8 100644
--- a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
+++ b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
@@ -97,31 +97,35 @@ FailureOr<LoweredLLVMFuncAttrs> GPUFuncOpLowering::buildLoweredGPULLVMFuncAttrs(
   NamedAttrList &discardable = loweredAttrs->discardableAttrs;
   auto *gpuDialect = cast<gpu::GPUDialect>(gpuFuncOp->getDialect());
 
-  auto appendIfNameAndValue = [&](StringAttr name, Attribute value) {
-    if (name && value)
-      discardable.append(name, value);
+  auto setIfValue = [&](StringAttr name, Attribute value) {
+    assert(name && "expected non-null attribute name");
+    if (value)
+      discardable.set(name, value);
+  };
+  auto setIfNameAndValue = [&](StringAttr name, Attribute value) {
+    if (name)
+      setIfValue(name, value);
   };
 
   DenseI32ArrayAttr knownBlockSize = gpuFuncOp.getKnownBlockSizeAttr();
   DenseI32ArrayAttr knownGridSize = gpuFuncOp.getKnownGridSizeAttr();
   DenseI32ArrayAttr knownClusterSize = gpuFuncOp.getKnownClusterSizeAttr();
 
-  appendIfNameAndValue(gpuDialect->getKnownBlockSizeAttrHelper().getName(),
-                       knownBlockSize);
-  appendIfNameAndValue(gpuDialect->getKnownGridSizeAttrHelper().getName(),
-                       knownGridSize);
-  appendIfNameAndValue(gpuDialect->getKnownClusterSizeAttrHelper().getName(),
-                       knownClusterSize);
+  setIfValue(gpuDialect->getKnownBlockSizeAttrHelper().getName(),
+             knownBlockSize);
+  setIfValue(gpuDialect->getKnownGridSizeAttrHelper().getName(), knownGridSize);
+  setIfValue(gpuDialect->getKnownClusterSizeAttrHelper().getName(),
+             knownClusterSize);
 
   if (isKernelFunc) {
-    discardable.append(gpuDialect->getKernelFuncAttrName(),
-                       rewriter.getUnitAttr());
+    discardable.set(gpuDialect->getKernelFuncAttrName(),
+                    rewriter.getUnitAttr());
     // Add a dialect specific kernel attribute in addition to GPU kernel
     // attribute. The former is necessary for further translation while the
     // latter is expected by gpu.launch_func.
-    appendIfNameAndValue(kernelAttributeName, rewriter.getUnitAttr());
-    appendIfNameAndValue(kernelBlockSizeAttributeName, knownBlockSize);
-    appendIfNameAndValue(kernelClusterSizeAttributeName, knownClusterSize);
+    setIfNameAndValue(kernelAttributeName, rewriter.getUnitAttr());
+    setIfNameAndValue(kernelBlockSizeAttributeName, knownBlockSize);
+    setIfNameAndValue(kernelClusterSizeAttributeName, knownClusterSize);
   }
 
   return loweredAttrs;
diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
index b96069ac41a44..9b451095f5500 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
@@ -785,6 +785,26 @@ gpu.module @test_module_33 {
   }
 }
 
+// CHECK-LABEL: gpu.module @test_module_preexisting_kernel_attrs
+"gpu.module"() <{sym_name = "test_module_preexisting_kernel_attrs"}> ({
+  // CHECK: llvm.func @legacy_kernel_attr() attributes {gpu.kernel, nvvm.kernel}
+  "gpu.func"() <{function_type = () -> ()}> ({
+    "gpu.return"() : () -> ()
+  }) {gpu.kernel, sym_name = "legacy_kernel_attr"} : () -> ()
+
+  // CHECK-LABEL: llvm.func @preexisting_gpu_known_block_size()
+  // CHECK: attributes {gpu.kernel, gpu.known_block_size = array<i32: 32, 1, 1>, nvvm.kernel, nvvm.maxntid = array<i32: 32, 1, 1>}
+  "gpu.func"() <{function_type = () -> (), known_block_size = array<i32: 32, 1, 1>, kernel}> ({
+    "gpu.return"() : () -> ()
+  }) {gpu.known_block_size = array<i32: 64, 1, 1>, sym_name = "preexisting_gpu_known_block_size"} : () -> ()
+
+  // CHECK-LABEL: llvm.func @preexisting_nvvm_maxntid()
+  // CHECK: attributes {gpu.kernel, gpu.known_block_size = array<i32: 32, 1, 1>, nvvm.kernel, nvvm.maxntid = array<i32: 32, 1, 1>}
+  "gpu.func"() <{function_type = () -> (), known_block_size = array<i32: 32, 1, 1>, kernel}> ({
+    "gpu.return"() : () -> ()
+  }) {nvvm.maxntid = array<i32: 64, 1, 1>, sym_name = "preexisting_nvvm_maxntid"} : () -> ()
+}) : () -> ()
+
 
 gpu.module @test_module_34 {
   // CHECK-LABEL: llvm.func @memref_signature(

``````````

</details>


https://github.com/llvm/llvm-project/pull/205035