[Mlir-commits] [mlir] [mlir][Vector] Reject scalable outer dim in multi_reduction unrolling (PR #224174)

[Cheng-Chieh Liu]

https://github.com/jay87406 updated https://github.com/llvm/llvm-project/pull/224174

>From cff5c6387deafa346d73dc89b193ef7ccbd4b60a Mon Sep 17 00:00:00 2001
From: Cheng-Chieh Liu <jay87406 at gmail.com>
Date: Wed, 16 Sep 2026 04:49:31 -0500
Subject: [PATCH 1/2] [mlir][Vector] Reject scalable outer dim in
 multi_reduction unrolling

TwoDimMultiReductionToReduction and TwoDimMultiReductionToElementWise
lower a rank-2 vector.multi_reduction into a sequence of per-lane
vector.extract/vector.reduction (or arith) ops, using
getSourceVectorType().getShape()[0] as the loop trip count. For a
scalable outer dimension like vector<[2]x4xf32>, that's only the
static minimum, not the actual vscale-scaled extent, so lanes past
the fixed trip count were silently dropped (or excluded from the
reduction) whenever vscale > 1 -- no error, just a wrong result.

Reject a scalable outer dimension in both patterns, and add negative
lit coverage for each lowering strategy.
---
 .../Transforms/LowerVectorMultiReduction.cpp   |  8 ++++++++
 .../vector-multi-reduction-pass-lowering.mlir  | 18 ++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp
index 76599822fbfe4..02aa8df70a7d2 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp
@@ -354,6 +354,10 @@ struct TwoDimMultiReductionToElementWise
     if (multiReductionOp.isReducedDim(1) || !multiReductionOp.isReducedDim(0))
       return failure();
 
+    if (multiReductionOp.getSourceVectorType().getScalableDims()[0])
+      return rewriter.notifyMatchFailure(
+          multiReductionOp, "cannot scalarize a scalable outer dimension");
+
     Value mask = maskingOp ? maskingOp.getMask() : Value();
 
     auto loc = multiReductionOp.getLoc();
@@ -411,6 +415,10 @@ struct TwoDimMultiReductionToReduction
     if (multiReductionOp.isReducedDim(0) || !multiReductionOp.isReducedDim(1))
       return failure();
 
+    if (multiReductionOp.getSourceVectorType().getScalableDims()[0])
+      return rewriter.notifyMatchFailure(
+          multiReductionOp, "cannot scalarize a scalable outer dimension");
+
     Value mask = maskingOp ? maskingOp.getMask() : nullptr;
 
     auto loc = multiReductionOp.getLoc();
diff --git a/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
index ddbc5c7bdb2c0..453684535a1a6 100644
--- a/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
@@ -61,3 +61,21 @@ func.func @vector_multi_reduction_masked(%arg0: vector<2x4xf32>, %acc: vector<2x
 //  INNER-PARALLEL: %[[INNERMASK:.+]] = vector.extract %[[TPMASK]][0] : vector<2xi1> from vector<4x2xi1>
 //  INNER-PARALLEL: %[[REDUCED:.+]] = arith.mulf %[[INNERVEC]], %[[ACC]] : vector<2xf32>
 //  INNER-PARALLEL: arith.select %[[INNERMASK]], %[[REDUCED]], %[[ACC]] : vector<2xi1>, vector<2xf32>
+
+// -----
+
+// INNER-REDUCTION-LABEL: @negative_multi_reduction_scalable_outer
+//       INNER-REDUCTION: vector.multi_reduction <mul>, %{{.*}}, %{{.*}} [1] : vector<[2]x4xf32> to vector<[2]xf32>
+func.func @negative_multi_reduction_scalable_outer(%arg0: vector<[2]x4xf32>, %acc: vector<[2]xf32>) -> vector<[2]xf32> {
+    %0 = vector.multi_reduction <mul>, %arg0, %acc [1] : vector<[2]x4xf32> to vector<[2]xf32>
+    return %0 : vector<[2]xf32>
+}
+
+// -----
+
+// INNER-PARALLEL-LABEL: @negative_multi_reduction_scalable_outer_reduce
+//       INNER-PARALLEL: vector.multi_reduction <mul>, %{{.*}}, %{{.*}} [0] : vector<[2]x4xf32> to vector<4xf32>
+func.func @negative_multi_reduction_scalable_outer_reduce(%arg0: vector<[2]x4xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
+    %0 = vector.multi_reduction <mul>, %arg0, %acc [0] : vector<[2]x4xf32> to vector<4xf32>
+    return %0 : vector<4xf32>
+}

>From 85eb4531a16b86826f409044f3e699041f16a4f8 Mon Sep 17 00:00:00 2001
From: Cheng-Chieh Liu <jay87406 at gmail.com>
Date: Sat, 19 Sep 2026 04:22:23 -0500
Subject: [PATCH 2/2] Check the other lowering strategy succeeds, add masked
 negative tests

Per review: show that rejecting the scalable outer dim in one
lowering strategy doesn't drop functionality, since the other
strategy still lowers the same input correctly. Also add masked
variants of both negative tests.
---
 .../vector-multi-reduction-pass-lowering.mlir | 34 ++++++++++++++++---
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
index 453684535a1a6..0496e9fa750ce 100644
--- a/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-multi-reduction-pass-lowering.mlir
@@ -64,8 +64,10 @@ func.func @vector_multi_reduction_masked(%arg0: vector<2x4xf32>, %acc: vector<2x
 
 // -----
 
-// INNER-REDUCTION-LABEL: @negative_multi_reduction_scalable_outer
-//       INNER-REDUCTION: vector.multi_reduction <mul>, %{{.*}}, %{{.*}} [1] : vector<[2]x4xf32> to vector<[2]xf32>
+//       ALL-LABEL: @negative_multi_reduction_scalable_outer
+//        ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<[2]xf32>
+// INNER-REDUCTION: vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [1] : vector<[2]x4xf32> to vector<[2]xf32>
+//  INNER-PARALLEL: vector.transpose %[[INPUT]], [1, 0] : vector<[2]x4xf32> to vector<4x[2]xf32>
 func.func @negative_multi_reduction_scalable_outer(%arg0: vector<[2]x4xf32>, %acc: vector<[2]xf32>) -> vector<[2]xf32> {
     %0 = vector.multi_reduction <mul>, %arg0, %acc [1] : vector<[2]x4xf32> to vector<[2]xf32>
     return %0 : vector<[2]xf32>
@@ -73,9 +75,33 @@ func.func @negative_multi_reduction_scalable_outer(%arg0: vector<[2]x4xf32>, %ac
 
 // -----
 
-// INNER-PARALLEL-LABEL: @negative_multi_reduction_scalable_outer_reduce
-//       INNER-PARALLEL: vector.multi_reduction <mul>, %{{.*}}, %{{.*}} [0] : vector<[2]x4xf32> to vector<4xf32>
+//       ALL-LABEL: @negative_multi_reduction_scalable_outer_reduce
+//        ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<4xf32>
+// INNER-REDUCTION: vector.transpose %[[INPUT]], [1, 0] : vector<[2]x4xf32> to vector<4x[2]xf32>
+//  INNER-PARALLEL: vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [0] : vector<[2]x4xf32> to vector<4xf32>
 func.func @negative_multi_reduction_scalable_outer_reduce(%arg0: vector<[2]x4xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
     %0 = vector.multi_reduction <mul>, %arg0, %acc [0] : vector<[2]x4xf32> to vector<4xf32>
     return %0 : vector<4xf32>
 }
+
+// -----
+
+//       ALL-LABEL: @negative_multi_reduction_scalable_outer_masked
+//        ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<[2]xf32>, %[[MASK:.+]]: vector<[2]x4xi1>
+// INNER-REDUCTION: vector.mask %[[MASK]] { vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [1] : vector<[2]x4xf32> to vector<[2]xf32> } : vector<[2]x4xi1> -> vector<[2]xf32>
+//  INNER-PARALLEL: vector.transpose %[[MASK]], [1, 0] : vector<[2]x4xi1> to vector<4x[2]xi1>
+func.func @negative_multi_reduction_scalable_outer_masked(%arg0: vector<[2]x4xf32>, %acc: vector<[2]xf32>, %mask: vector<[2]x4xi1>) -> vector<[2]xf32> {
+    %0 = vector.mask %mask { vector.multi_reduction <mul>, %arg0, %acc [1] : vector<[2]x4xf32> to vector<[2]xf32> } : vector<[2]x4xi1> -> vector<[2]xf32>
+    return %0 : vector<[2]xf32>
+}
+
+// -----
+
+//       ALL-LABEL: @negative_multi_reduction_scalable_outer_reduce_masked
+//        ALL-SAME: %[[INPUT:.+]]: vector<[2]x4xf32>, %[[ACC:.+]]: vector<4xf32>, %[[MASK:.+]]: vector<[2]x4xi1>
+// INNER-REDUCTION: vector.transpose %[[MASK]], [1, 0] : vector<[2]x4xi1> to vector<4x[2]xi1>
+//  INNER-PARALLEL: vector.mask %[[MASK]] { vector.multi_reduction <mul>, %[[INPUT]], %[[ACC]] [0] : vector<[2]x4xf32> to vector<4xf32> } : vector<[2]x4xi1> -> vector<4xf32>
+func.func @negative_multi_reduction_scalable_outer_reduce_masked(%arg0: vector<[2]x4xf32>, %acc: vector<4xf32>, %mask: vector<[2]x4xi1>) -> vector<4xf32> {
+    %0 = vector.mask %mask { vector.multi_reduction <mul>, %arg0, %acc [0] : vector<[2]x4xf32> to vector<4xf32> } : vector<[2]x4xi1> -> vector<4xf32>
+    return %0 : vector<4xf32>
+}