llvm / llvm/llvm-project

[mlir][linalg] Passes that rewrite a structured op's iteration domain do not remap linalg.index (4 sites identified)

Open
#220,505 3 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

## 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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.