[Mlir-commits] [mlir] [mlir][vector] Separate multi_reduce lowering into transformations, flattening, and unrolling. (PR #178974)

[Andrzej Warzyński]

banach-space wrote:

Thank you for sending this!

> The pass `lower-vector-multi-reduction` used flattening and unrolling along with several other patterns to lower multi_reduce operations into simpler operations. This makes it unrolling difficult to control. Therefore, we split these pattern sets into three stages.
> 
> 1. Transformations useful for `lower-vector-multi-reduction` in `populateVectorMultiReductionInnerOuterDimPatterns` which populates `InnerOuterDimReductionConversion` and `OneDimMultiReductionToTwoDim`.
> 2. Set the pattern `ReduceMultiDimReductionRank` in `populateVectorMultiReductionFlatteningPatterns`.
> 3. Add unrolling patterns in `populateVectorUnrollMultiReduction` which are rank reducing.
> 
> The original population of patterns `populateVectorMultiReductionLoweringPatterns` is preserved to allow external projects to migrate to these new pattern population functions at their own pace.
> 
> The vector unrolling has been generalized. Previously only vectors of rank two (which were expected due to `ReduceMultiDimReductionRank` being run earlier) were unrolled. With this new structure, the unrolling patterns have been generalized to allow unrolling of n-dimensional vectors either with the innermost or outermost dimensions are reduced.
> 
> The generalization works in the following pattern:
> 
> 1. Rank-1 multi_reduction -> convert directly to vector.reduction
> 2. OuterReduction with single reduction dim -> unroll to elementwise ops (base case).
> 3. OuterReduction with multiple reduction dims -> extract along outermost (reduced) dim, chain multi_reductions sequentially.
> 4. InnerReduction -> extract along outermost (parallel) dim, create independent multi_reductions, reassemble results.
> 
> The `lower-vector-multi-reduction` pass is now split into three different stages. Generalizing the unrolling patterns forces us to establish an order in which patterns must be applied. Otherwise, different patterns may match the same operations introducing non-determinism.
> 
> The first couple of commits were handwritten. The next couple of commits were based on other handwritten commits and just copy pasted by claude. The rest were mostly generated by claude.
> 
> Assisted-by: claude

There’s too much going on in a single PR here. The title says “Separate multi_reduce lowering into transformations, flattening, and unrolling” but I’m having difficulty finding the actual separation/refactor that performs that decomposition - most of the changes look like new code additions:
<img width="148" height="50" alt="Screenshot 2026-02-10 at 13 29 19" src="https://github.com/user-attachments/assets/86740102-9895-48e2-aaf9-356602a473a2" />

These patterns look like completely new code:
* `OneDimMultiReductionToReduction`
* `UnrollMultiReductionOuterBaseCase`
* `UnrollMultiReductionOuterGeneralCase`
* `UnrollInnerReductionAlongOuterParallel`

? Why do we need them here if this PR is merely about splitting existing logic? Also, are the newly added tests for these new or pre-existing patterns? How do I check which test is for which pattern? And why not re-use pre-existing tests for unrolling?

All in all I feel confused about the goal of these changes:
* Are you decomposing `populateVectorMultiReductionLoweringPatterns` into finer-granularity `populate-` hooks? If yes then the former should be refactored accordingly.
* Are you adding new patterns? Why? Doesn't `lower-vector-multi-reduction` already implement everything that is required?
* Are you also generalising unrolling? That's orthogonal to "separating" and should be implemented in a dedicated PR.

All in all, this PR should be split into smaller PRs. I can help identifying the right path, but would need to understand the overall goal a bit better.

Thanks!

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