Skip to content

[mlir][gpu] Update attribute definitions in gpu::LaunchOp #152106

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Aug 8, 2025

Conversation

CoTinker
Copy link
Contributor

@CoTinker CoTinker commented Aug 5, 2025

gpu::LaunchOp is updated the following way:

  • Change the attribute type of kernel function and module from SymbolRefAttr to FlatSymbolRefAttr to avoid nested symbol references.
  • Rename variables from camel case (kernelFunc, kernelModule) to lower case (function, module) and update the syntax.
  • LaunchOp::build support passing module and function attributes.

@llvmbot
Copy link
Member

llvmbot commented Aug 5, 2025

@llvm/pr-subscribers-mlir-gpu

Author: Longsheng Mou (CoTinker)

Changes

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 --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.
 

@llvmbot
Copy link
Member

llvmbot commented Aug 5, 2025

@llvm/pr-subscribers-mlir

Author: Longsheng Mou (CoTinker)

Changes

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 --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.
 

@wangzpgi wangzpgi requested a review from clementval August 6, 2025 01:49
@CoTinker CoTinker changed the title [mlir][gpu] Correctly retrieve kernel function name for nested references [mlir][gpu] Update attribute definitions in gpu::LaunchOp Aug 6, 2025
This PR makes two updates to `gpu.launch`:
- Change the attribute type for kernel function and module from
  `SymbolRefAttr` to `FlatSymbolRefAttr` to avoid nested symbol references.
- Rename variables from camel case (kernelFunc, kernelModule) to lower case
  (function, module).
Copy link
Collaborator

@joker-eph joker-eph left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LG, thanks!

@CoTinker CoTinker requested a review from grypp August 7, 2025 15:37
Copy link
Member

@grypp grypp left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nice thanks

Copy link
Contributor

@clementval clementval left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@CoTinker CoTinker merged commit 7d886fa into llvm:main Aug 8, 2025
9 checks passed
@CoTinker CoTinker deleted the leaf branch August 8, 2025 03:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants