[MLIR] `affine-loop-tile` tiles a no-op outer loop and misses inner matmul loop nest
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
`affine-loop-tile` with `tile-size=16` strip-mines an outer loop whose static trip count is only 8, producing a single-tile loop nest, but it does not tile the inner `128 x 128 x 128` matmul-like loop nest, resulting in a missed optimization.
The outer loop is not profitable to tile because its trip count is smaller than the tile size. The inner perfect loop nest has large constant trip counts and is the loop nest that should be tiled.
### Input Program
```llvm
module {
func.func @matmul_suffix(%arg0: memref<8x128x128xf32>, %arg1: memref<128x128xf32>, %arg2: memref<8x128x128xf32>, %arg3: memref<8xf32>) {
%cst = arith.constant 0.000000e+00 : f32
affine.for %arg4 = 0 to 8 {
affine.store %cst, %arg3[%arg4] : memref<8xf32>
affine.for %arg5 = 0 to 128 {
affine.for %arg6 = 0 to 128 {
affine.for %arg7 = 0 to 128 {
%0 = affine.load %arg0[%arg4, %arg5, %arg7] : memref<8x128x128xf32>
%1 = affine.load %arg1[%arg7, %arg6] : memref<128x128xf32>
%2 = affine.load %arg2[%arg4, %arg5, %arg6] : memref<8x128x128xf32>
%3 = arith.mulf %0, %1 : f32
%4 = arith.addf %2, %3 : f32
affine.store %4, %arg2[%arg4, %arg5, %arg6] : memref<8x128x128xf32>
}
}
}
}
return
}
}
```
### Command
```bash
mlir-opt input.mlir --affine-loop-tile="tile-size=16"
```
### Actual Output
The pass only tiles the outer loop with trip count 8. The outer tiled loop has only one iteration because the original trip count is 8 and the tile size is 16. This is effectively a no-op strip-mining transformation and can even make the IR worse before canonicalization. The inner 128 x 128 x 128 loop nest is left untiled.
```llvm
#map = affine_map<(d0) -> (d0)>
#map1 = affine_map<(d0) -> (d0 + 8)>
module {
func.func @matmul_suffix(%arg0: memref<8x128x128xf32>, %arg1: memref<128x128xf32>, %arg2: memref<8x128x128xf32>, %arg3: memref<8xf32>) {
%cst = arith.constant 0.000000e+00 : f32
affine.for %arg4 = 0 to 8 step 16 {
affine.for %arg5 = #map(%arg4) to #map1(%arg4) {
affine.store %cst, %arg3[%arg5] : memref<8xf32>
affine.for %arg6 = 0 to 128 {
affine.for %arg7 = 0 to 128 {
affine.for %arg8 = 0 to 128 {
%0 = affine.load %arg0[%arg5, %arg6, %arg8] : memref<8x128x128xf32>
%1 = affine.load %arg1[%arg8, %arg7] : memref<128x128xf32>
%2 = affine.load %arg2[%arg5, %arg6, %arg7] : memref<8x128x128xf32>
%3 = arith.mulf %0, %1 : f32
%4 = arith.addf %2, %3 : f32
affine.store %4, %arg2[%arg5, %arg6, %arg7] : memref<8x128x128xf32>
}
}
}
}
}
return
}
}
```
Version: b37a8a70cc98915d157cac2f049e9f46d6da8fe5
Contributor guide
Research direction
Start by running mlir-opt input.mlir --affine-loop-tile="tile-size=16" with the provided input.mlir and inspect the affine-loop-tile entry point and its handling of nested loops. Done means the trip-count-8 outer loop is not needlessly strip-mined and the inner 128 x 128 x 128 matmul-like loop nest is tiled.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100