llvm / llvm/llvm-project

[MLIR] GPU dialect `-gpu-kernel-outlining` optimization drops argument attributes

Open
#221,441 1 comment 0 reactions 0 assignees View on GitHub
mlir
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.