[mlir][linalg] Passes that rewrite a structured op's iteration domain do not remap linalg.index (4 sites identified)
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
## Summary
`linalg.index d` returns the iteration index at position `d` of **the enclosing
structured operation's indexing-map domain**. Its meaning is therefore defined
by the surrounding `indexing_maps` / `iterator_types`. Any transformation that
changes that domain — splitting, dropping, adding or reordering loops — must
rewrite the `linalg.index` operations in the region it moves. Tiling does this, via `offsetIndices` (added by
[#188261](https://github.com/llvm/llvm-project/pull/188261).
Four other in-tree transformations move a region into a new domain and leave
`linalg.index` untouched. Each produces verifier-clean IR that silently computes
wrong region-visible index values. The operand data is remapped correctly in
every case, which is what makes the defect silent.
## Reproducers
### transform.structured.split_reduction
```mlir
func.func @sr(%in: tensor<8xi32>, %out: tensor) -> tensor {
%r = linalg.generic {indexing_maps = [affine_map<(k) -> (k)>, affine_map<(k) -> ()>],
iterator_types = ["reduction"]}
ins(%in : tensor<8xi32>) outs(%out : tensor) {
^bb0(%a: i32, %acc: i32):
%k = linalg.index 0 : index
%c = arith.index_cast %k : index to i32
%s = arith.addi %acc, %c : i32
linalg.yield %s : i32
} -> tensor
return %r : tensor
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
%g = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op
%a, %b, %c, %d = transform.structured.split_reduction %g
{ split_factor = 4, insert_split_dimension = 0 }
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
transform.yield
}
}
```
And we run `mlir-opt -transform-interpreter repro.mlir`
The operand is expanded `tensor.expand_shape %arg0 [[0, 1]] output_shape [4, 2]`,
so the source coordinate is `k = p*2 + r` with `p ∈ [0,4)`, `r ∈ [0,2)`. The new
generic has `iterator_types = ["parallel", "reduction"]` and still contains
`linalg.index 0`, which is now `p`.
This reduction sums the index itself, so the source computes
`0+1+...+7 = 28`. After splitting, the first generic computes
`partial[p] = Σ_r p = 2p`, giving `[0, 2, 4, 6]` and a final result of **12**.
With the correct remap `k = 2p + r` it would be `partial[p] = 4p + 1`, i.e.
`[1, 5, 9, 13]`, summing to **28**.
## The same defect appears in three more places
| Transformation | Domain change | `linalg.index` now returns | Correct remap |
|---|---|---|---|
| `transform.structured.pack` | `["parallel"]` → `["parallel","parallel"]` over `2x4`; `i = outer*4 + inner` | the outer tile coordinate | `index 0 * 4 + index 1` |
| `transform.structured.decompose` (`downscaleSizeOneWindowedConvolution`) | 6-D `(oh,ow,f,kh,kw,c)` → 4-D `(ow,f,kw,c)`; unit loops dropped | a *different surviving* loop — `index 1` was `ow`, is now `f` | `index 0` |
| `-sparse-reinterpret-map` | `["parallel"]` → `["parallel","parallel"]` over `4x2`; `i = 2*block + intra` | the storage-level coordinate `block` | `index 0 * 2 + index 1` |
Contributor guide
Research direction
Start with the four transformation entry points named in the issue: transform.structured.split_reduction, transform.structured.pack, transform.structured.decompose, and -sparse-reinterpret-map. Run the supplied split_reduction reproducer with mlir-opt -transform-interpreter, then verify that each transformation remaps linalg.index correctly and preserves the expected results for all four cases.
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
- 45/100