Skip to content

Commit

Permalink
generalize pass gpu-kernel-outlining for symbol op (#72074)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
fengxie authored Nov 13, 2023
1 parent abeffc9 commit 9a3d3c7
Showing 2 changed files with 38 additions and 3 deletions.
4 changes: 2 additions & 2 deletions mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
Original file line number Diff line number Diff line change
@@ -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 =
37 changes: 36 additions & 1 deletion mlir/test/Dialect/GPU/outlining.mlir
Original file line number Diff line number Diff line change
@@ -37,7 +37,6 @@ func.func @launch() {
}

// CHECK-DL-LABEL: gpu.module @launch_kernel attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<index, 32 : i32>>}

// CHECK-LABEL: gpu.module @launch_kernel
// CHECK-NEXT: gpu.func @launch_kernel
// CHECK-SAME: (%[[KERNEL_ARG0:.*]]: f32, %[[KERNEL_ARG1:.*]]: memref<?xf32, 1>)
@@ -63,6 +62,42 @@ func.func @launch() {

// -----

// This test checks gpu-out-lining can handle gpu.launch kernel from an llvm.func
// CHECK-LABEL: @launch_from_llvm_func
llvm.func @launch_from_llvm_func() {
// CHECK: %[[ARG0:.*]] = "op"() : () -> f32
%0 = "op"() : () -> (f32)
// CHECK: %[[ARG1:.*]] = "op"() : () -> memref<?xf32, 1>
%1 = "op"() : () -> (memref<?xf32, 1>)

// CHECK: %[[DIM:.*]] = arith.constant 1
%dim = arith.constant 1 : index

// CHECK: gpu.launch_func @launch_from_llvm_func_kernel::@launch_from_llvm_func_kernel
// CHECK-SAME: (%[[DIM]], %[[DIM]], %[[DIM]])
// CHECK-SAME: (%[[DIM]], %[[DIM]], %[[DIM]]) args(%[[ARG0]] : f32, %[[ARG1]] : memref<?xf32, 1>)
// CHECK-NEXT: llvm.return

// CHECK: gpu.func {{.*}} kernel attributes
// CHECK-SAME: gpu.known_block_size = array<i32: 1, 1, 1>
// CHECK-SAME: gpu.known_grid_size = array<i32: 1, 1, 1>
// CHECK: gpu.return
gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %dim, %grid_y = %dim,
%grid_z = %dim)
threads(%tx, %ty, %tz) in (%block_x = %dim, %block_y = %dim,
%block_z = %dim) {
"use"(%0): (f32) -> ()
"some_op"(%bx, %block_x) : (index, index) -> ()
%2 = memref.load %1[%tx] : memref<?xf32, 1>
gpu.terminator
}
llvm.return
}

// CHECK-DL-LABLE: gpu.module @launch_from_llvm_func_kernel attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<index, 32 : i32>>}

// -----

// CHECK: module attributes {gpu.container_module}
// CHECK-LABEL: @multiple_launches
func.func @multiple_launches() {

0 comments on commit 9a3d3c7

Please sign in to comment.