[Mlir-commits] [mlir] [mlir][vector] Add multi_reduction_rank_reducing_unrolling (PR #182301)

[Andrzej Warzyński]

banach-space wrote:

I had to take a step back from this.

I've tried creating a mental map of the available patterns and how they complement each other. To this end, I looked at the other patterns that are currently labelled as unrolling, namely those collected under [`populateVectorMultiReductionUnrollingPatterns`](https://github.com/llvm/llvm-project/blob/8f378ea7e6fa31179266c69368be56a866b631e1/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp?plain=1#L540-L549). That gave me two patterns:

* `TwoDimMultiReductionToElementWise`,
* `TwoDimMultiReductionToReduction`

To make this a bit easier for me (and in the spirit of keeping things consistent), I prepared https://github.com/llvm/llvm-project/pull/183048/. It’s inspired by your implementation - thanks for prompting me to prepare that!

Below is a consolidated summary of all the unrolling patterns (including the ones implemented in this PR):

| Pattern                               | Matches                                  | # Reduction Dims | Reduction Dim Position | Rewrites To     |
|----------------------------------------|-------------------------------------------|------------------|-------------------------|--------------------------------------|
| `MultiReductionToArithOps`   **(new)**     | `vector.multi_reduction`  | 1              | Outer-most  | `vector.extract` + `arith` ops |
| `UnrollMultiReductionInnerParallel` **(new)** | `vector.multi_reduction` (innerparallel)  | > 1              | Outer-most | Sequence of `vector.multi_reduction`|
| `TwoDimMultiReductionToElementWise`    | 2D `vector.multi_reduction`               | 1                | Outer-most              | `vector.extract` + `arith` ops |
| `TwoDimMultiReductionToReduction`      | 2D `vector.multi_reduction`               | 1                | Inner-most              | `vector.reduction` ops              |

Let me know if something here is off and I will update it in place.

---

Two high-level observations from that exercise:

* There is a large overlap between `MultiReductionToArithOps` and `TwoDimMultiReductionToElementWise`. Why do we need both? 
* `UnrollMultiReductionInnerParallel` and `TwoDimMultiReductionToReduction` are clearly different in intent - that’s good.

---

Below are some notes on naming (inspired by your comment):

> This can be visualized as a three-stage process: everything before unrolling, unrolling, everything after unrolling.

I agree, and I think we should make that distinction explicit in the names and in the comments. To me, unrolling means:

> Take an N-D operation `A` and replace it with a sequence of (N−1)-D operations `A` (i.e., decompose by dimension).

Replacing `A` with a different operation `B` is lowering/rewrite. Replacing `A` with a sequence of `B` is the grey area between unrolling and lowering. For clarity I'd prefer names that show intent, for example:
* `UnrollXToX` (if we're truly decomposing same-op across dimensions)
* `RewriteXAsY` or `LowerXToY` (if we're changing op kind)
* `UnrollXAsYSequence` (if we do both: decomposition + op-kind change)

Concretely, `TwoDimMultiReductionToReduction` looks like a lowering (it rewrites to `vector.reduction`), not an unrolling, so it might be clearer not to group it under an unrolling heading.

---

Btw,  is the goal of `UnrollMultiReductionInnerParallel` to "reduce" the number of "reduction dims"? So its basically "unrolling outer reduction dim"? I am trying to find a clear label for this pattern.

https://github.com/llvm/llvm-project/pull/182301