[mlir][linalg] Design a contract for coupled multi-result partial reductions
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
### Problem
`PartialReductionOpInterface` and the SCF tilers can already carry multiple
partial results. The Linalg external model, however, assigns those results
independent algebra:
1. run `matchReduction` for each DPS init;
2. require one recognized scalar combiner;
3. infer that operation's neutral element;
4. emit one merge reduction per result.
That is sufficient for independent results such as `(sum, max)`, but it
correctly fails closed for a coupled state. For stable LogSumExp, for example,
the new sum depends on the newly merged maximum:
```text
merge((m1, l1), (m2, l2)):
m = max(m1, m2)
l = exp(m1 - m) * l1 + exp(m2 - m) * l2
return (m, l)
```
A variadic `linalg.reduce` can already represent this as one combiner region.
Applying `transform.structured.tile_reduction_using_for` currently stops in
the Linalg external model with:
```text
'linalg.reduce' op Failed to anaysis the reduction operation.
failed to tile using partial reduction
```
This is not a correctness bug. Merging `m` and `l` independently would be
invalid; the missing piece is an explicit contract under which the complete
region may be used as one partial-reduction algebra.
### Existing machinery
Most of the required scheduling path is already upstream:
- [#92624](https://github.com/llvm/llvm-project/pull/92624) made
`PartialReductionOpInterface` multi-result;
- [#120118](https://github.com/llvm/llvm-project/pull/120118) and
[#157932](https://github.com/llvm/llvm-project/pull/157932) let the existing
`for`/`forall` Transform operations target any implementation of that
interface;
- [#147593](https://github.com/llvm/llvm-project/pull/147593) connected the
parallel partial-reduction strategy to the existing tile-and-fuse path.
Source operations can therefore own their algebra and implement the existing
interface without an MLIR core change. IREE does this for coupled
`online_attention` and `arg_compare` state.
### Design question
Should coupled semantics remain source-op-owned, or should `linalg.reduce`
have a typed, opt-in contract for using its complete variadic combiner during
partial-reduction tiling?
If a shared Linalg contract is desirable, it needs to state separately:
- whether changing parenthesization is permitted;
- whether changing merge order is permitted;
- how partial identities are supplied, or that all destination inits are
two-sided identities;
- the floating-point and exceptional-value authority for those changes.
The implementation can then remain mechanical:
1. reject an uncontracted or memory-effecting coupled region;
2. construct one identity-initialized partial tensor per state component;
3. reuse the existing partial-tile implementation;
4. emit one variadic `linalg.reduce` for the merge;
5. clone the original combiner region once.
No new reduction operation, scheduler, Transform operation, or target-specific
pass is required.
### Prototype
An executable prototype is available at
[a8077f34c840](https://github.com/takatodo/llvm-project/commit/a8077f34c840).
It adds 101 C++ lines to the existing Linalg external model and 167 test lines.
The string attribute used there is only a placeholder, not a proposed API.
The prototype demonstrates:
- coupled LogSumExp `(m, l)` with the existing `scf.for` schedule;
- a non-floating coupled ring state with the existing `scf.forall` schedule;
- fail-closed behavior when the contract is absent;
- rejection of a memory-effecting combiner;
- preservation of the existing independent-result path.
The focused test and the complete Linalg dialect suite pass (`172/172`). A
three-way application onto LLVM `main` at `8d292a7c4b95` is conflict-free.
### Questions
1. Should the contract live on `linalg.reduce`, or should source operations
continue to own `PartialReductionOpInterface`?
2. Should identities be asserted destination inits or explicit values distinct
from the final destination inits?
3. Must the first contract require commutativity/reordering, or can an
order-preserving strategy require associativity only?
4. What IR-level authority should be required for floating-point
reassociation?
Automatic softmax/Attention recognition, profitability, target mapping, and
sliding-window state transitions are intentionally out of scope.
Assisted-by: OpenAI Codex
Contributor guide
Research direction
Read the existing Linalg external model, PartialReductionOpInterface path, and prototype commit a8077f34c840 first. Compare its focused tests and complete Linalg dialect suite results, then resolve where the coupled-reduction contract belongs, how identities and reordering are authorized, and what fail-closed behavior is required.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100