llvm / llvm/llvm-project

[MLIR] `affine-loop-tile` tiles a no-op outer loop and misses inner matmul loop nest

Open
#220,514 2 comments 0 reactions 0 assignees View on GitHub
mlir
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.