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

[Longsheng Mou]

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

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

>From 88781d0d2940381937b6cfc393244e9a8e83f416 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Tue, 5 Aug 2025 17:10:27 +0800
Subject: [PATCH 1/2] [mlir][gpu] Correctly retrieve kernel function name for
 nested references

This PR replaces `getRootReference()` with `getLeafReference()` to correctly retrieve the actual kernel function name when `kernelFunc` has nested symbol references.
---
 mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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(),

>From fa19500cca3289cf8ffa4bbfd520c2db49615419 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Tue, 5 Aug 2025 17:13:53 +0800
Subject: [PATCH 2/2] add test

---
 mlir/test/Dialect/GPU/outlining.mlir | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

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.