[Mlir-commits] [mlir] generalize pass gpu-kernel-outlining for symbol op (PR #72074)

[llvmlistbot at llvm.org]

https://github.com/fengxie created https://github.com/llvm/llvm-project/pull/72074

This PR generalize gpu-out-lining pass to take care of ops `SymbolOpInterface` instead of just `func::FuncOp`.

Before this change, gpu-out-lining pass will skip `llvm.func`.
```mlir
module {
  llvm.func @main() {
    %c1 = arith.constant 1 : index
    gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1) threads(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) {
      gpu.terminator
    }
    llvm.return
  }
}
```

After this change, gpu-out-lining pass can handle llvm.func as well.

```bash
cmake --build . -- -j24 check-mlir
[281/282] Running the MLIR regression tests

Testing Time: 5.83s

Total Discovered Tests: 2418
  Skipped          :    1 (0.04%)
  Unsupported      :  354 (14.64%)
  Passed           : 2062 (85.28%)
  Expectedly Failed:    1 (0.04%)
```

>From 33c60ccee1075f0602f6916baff3d5cbbe97abde Mon Sep 17 00:00:00 2001
From: Fung Xie <ftse at nvidia.com>
Date: Mon, 13 Nov 2023 09:46:49 +0800
Subject: [PATCH] generalize pass gpu-kernel-outlining for symbol op

---
 mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
index b1e2f914db4cb9b..7432a58f18b4422 100644
--- a/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
@@ -349,13 +349,13 @@ class GpuKernelOutliningPass
   void runOnOperation() override {
     SymbolTable symbolTable(getOperation());
     bool modified = false;
-    for (auto func : getOperation().getOps<func::FuncOp>()) {
+    for (auto func : getOperation().getOps<SymbolOpInterface>()) {
       // Insert just after the function.
       Block::iterator insertPt(func->getNextNode());
       auto funcWalkResult = func.walk([&](gpu::LaunchOp op) {
         SetVector<Value> operands;
         std::string kernelFnName =
-            Twine(op->getParentOfType<func::FuncOp>().getName(), "_kernel")
+            Twine(op->getParentOfType<SymbolOpInterface>().getName(), "_kernel")
                 .str();
 
         gpu::GPUFuncOp outlinedFunc =