[MLIR] GPU dialect `-gpu-kernel-outlining` optimization drops argument attributes
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
In the following example, `gpu-kernel-outlining` drops argument attributes when values captured by `gpu.launch` are converted into kernel arguments, which may lead to missed optimization opportunities in subsequent passes.
### Reproducer
**input.mlir**
```llvm
module {
llvm.func @kernel_argument(%arg0: !llvm.ptr {llvm.align = 16 : i64, llvm.noalias}) {
%c1 = arith.constant 1 : index
gpu.launch blocks(%arg1, %arg2, %arg3) in (%arg7 = %c1, %arg8 = %c1, %arg9 = %c1) threads(%arg4, %arg5, %arg6) in (%arg10 = %c1, %arg11 = %c1, %arg12 = %c1) {
"test.consume"(%arg0) : (!llvm.ptr) -> ()
gpu.terminator
}
llvm.return
}
}
```
Run: `mlir-opt -gpu-launch-sink-index-computations -gpu-kernel-outlining input.mlir`
The original argument has llvm.align and llvm.noalias, but the outlined kernel argument has no corresponding attributes:
```llvm
module attributes {gpu.container_module} {
llvm.func @kernel_argument(%arg0: !llvm.ptr {llvm.align = 16 : i64, llvm.noalias}) {
%c1 = arith.constant 1 : index
gpu.launch_func @kernel_argument_kernel::@kernel_argument_kernel blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1) args(%arg0 : !llvm.ptr)
llvm.return
}
gpu.module @kernel_argument_kernel {
gpu.func @kernel_argument_kernel(%arg0: !llvm.ptr) kernel attributes {known_block_size = array, known_grid_size = array} {
%block_id_x = gpu.block_id x
%block_id_y = gpu.block_id y
%block_id_z = gpu.block_id z
%thread_id_x = gpu.thread_id x
%thread_id_y = gpu.thread_id y
%thread_id_z = gpu.thread_id z
%grid_dim_x = gpu.grid_dim x
%grid_dim_y = gpu.grid_dim y
%grid_dim_z = gpu.grid_dim z
%block_dim_x = gpu.block_dim x
%block_dim_y = gpu.block_dim y
%block_dim_z = gpu.block_dim z
"test.consume"(%arg0) : (!llvm.ptr) -> ()
gpu.return
}
}
}
```
### Expected behavior
Applicable argument attributes should be preserved in the outlined kernel argument so that downstream optimizations can still use this information.
Version: b1e0aa39763434060ab6e4dc87af3630ac9a04ac
Contributor guide
Research direction
Start at the gpu-kernel-outlining pass and reproduce the supplied input.mlir with mlir-opt -gpu-launch-sink-index-computations -gpu-kernel-outlining. Done means the outlined kernel argument retains applicable llvm.align and llvm.noalias attributes, preserving the information for downstream optimizations.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 58/100