[Mlir-commits] [mlir] [mlir][gpu] Correctly retrieve kernel function name for nested references (PR #152106)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Longsheng Mou (CoTinker)

<details>
<summary>Changes</summary>

This PR replaces `getRootReference()` with `getLeafReference()` to correctly retrieve the actual kernel function name when `kernelFunc` has nested symbol references.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp (+1-1) 
- (modified) mlir/test/Dialect/GPU/outlining.mlir (+18) 


``````````diff
diff --git a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
index 99f5c5b0cf139..346cce0f08b74 100644
--- a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
@@ -357,7 +357,7 @@ class GpuKernelOutliningPass
         SetVector<Value> operands;
         std::string kernelFnName;
         if (op.getKernelFunc()) {
-          kernelFnName = op.getKernelFunc()->getRootReference().str();
+          kernelFnName = op.getKernelFunc()->getLeafReference().str();
         } else {
           kernelFnName =
               Twine(op->getParentOfType<SymbolOpInterface>().getName(),
diff --git a/mlir/test/Dialect/GPU/outlining.mlir b/mlir/test/Dialect/GPU/outlining.mlir
index d48fa054432d1..fda9877f58fe2 100644
--- a/mlir/test/Dialect/GPU/outlining.mlir
+++ b/mlir/test/Dialect/GPU/outlining.mlir
@@ -530,6 +530,24 @@ func.func @testKernelAttributes() {
   return
 }
 
+// -----
+
+// This test tests the kernelFunc has nested references.
+
+// CHECK-LABEL: func.func @testKernelAttributesWithNestedFunc
+//       CHECK:   gpu.launch_func  @test_module::@test_kernel_func
+//       CHECK: gpu.module @test_module
+//       CHECK:   gpu.func @test_kernel_func()
+//   CHECK-NOT:   gpu.func @test_module()
+func.func @testKernelAttributesWithNestedFunc(%arg0 : index) {
+  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %arg0, %grid_y = %arg0, %grid_z = %arg0)
+             threads(%tx, %ty, %tz) in (%block_x = %arg0, %block_y = %arg0, %block_z = %arg0) {
+    "some_op"(%bx, %tx) : (index, index) -> ()
+    gpu.terminator
+  } {kernelModule = @test_module, kernelFunc = @test_module::@test_kernel_func}
+  return
+}
+
 // -----
 // This test tests the two optional attributes kernelModule and kernelFunc for gpu.launch, when kernelModule already exists.
 

``````````

</details>


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