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

[Cheng-Chieh Liu]

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

`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 represents only the static minimum rather than the actual vscale-scaled extent. Consequently, whenever `vscale > 1`, lanes beyond the fixed trip count were either silently excluded from the reduction or left unwritten in the result — both cases producing silent data corruption instead of an error.

Reject scalable outer dimensions in both patterns and add negative lit test coverage for each lowering strategy.

>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] [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>
+}