[MLIR][Vector] Fix multi_reduction fold to handle empty reduction dims for any rank

[joker-eph]

The fold for vector.multi_reduction only handled the rank-1 case with no reduction dimensions. For higher-rank vectors (e.g., vector<2x3xf32>) with empty reduction dims [], the fold returned null, allowing ElideUnitDimsInMultiDimReduction to fire incorrectly. That canonicalization pattern checks that all reduced dims have size 1, but with zero reduction dims the check trivially passes, and the pattern then computes acc op source (e.g., acc + source) instead of the correct no-op result (source).

This caused --canonicalize to produce a different value than --lower-vector-multi-reduction for the same program:

vector.mask %m { vector.multi_reduction , %src, %src [] :
vector<3x3xi32> to vector<3x3xi32> } : vector<3x3xi1> -> vector<3x3xi32>

• Without --lower-vector-multi-reduction: src + src (e.g., 2)
• With --lower-vector-multi-reduction: src (e.g., 1)

Fix the fold to return source for any rank when reduction_dims is empty. This makes the empty-dims case consistent: the operation is a noop regardless of rank, and ElideUnitDimsInMultiDimReduction no longer gets a chance to mishandle it.

Fixes #129415

Assisted-by: Claude Code

[llvmbot]

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

Changes

The fold for vector.multi_reduction only handled the rank-1 case with no reduction dimensions. For higher-rank vectors (e.g., vector<2x3xf32>) with empty reduction dims [], the fold returned null, allowing ElideUnitDimsInMultiDimReduction to fire incorrectly. That canonicalization pattern checks that all reduced dims have size 1, but with zero reduction dims the check trivially passes, and the pattern then computes acc op source (e.g., acc + source) instead of the correct no-op result (source).

This caused --canonicalize to produce a different value than --lower-vector-multi-reduction for the same program:

vector.mask %m { vector.multi_reduction <add>, %src, %src [] :
vector<3x3xi32> to vector<3x3xi32> } : vector<3x3xi1> -> vector<3x3xi32>

• Without --lower-vector-multi-reduction: src + src (e.g., 2)
• With --lower-vector-multi-reduction: src (e.g., 1)

Fix the fold to return source for any rank when reduction_dims is empty. This makes the empty-dims case consistent: the operation is a noop regardless of rank, and ElideUnitDimsInMultiDimReduction no longer gets a chance to mishandle it.

Fixes #129415

Assisted-by: Claude Code

---

Full diff: https://github.com/llvm/llvm-project/pull/188983.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+2-2)
• (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+20)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 73632875ca9e2..85f579e4b984d 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -528,8 +528,8 @@ void vector::MultiDimReductionOp::build(OpBuilder &builder,
 }
 
 OpFoldResult MultiDimReductionOp::fold(FoldAdaptor adaptor) {
-  // Single parallel dim, this is a noop.
-  if (getSourceVectorType().getRank() == 1 && !isReducedDim(0))
+  // No reduction dims: this is a noop regardless of rank.
+  if (getReductionDims().empty())
     return getSource();
   return {};
 }
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index e2045fc526ec1..beda984aec5b6 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2202,6 +2202,26 @@ func.func @masked_vector_multi_reduction_single_parallel(%arg0: vector<2xf32>, %
 
 // -----
 
+// CHECK-LABEL: func @vector_multi_reduction_no_reduction_dims_nd(
+//  CHECK-SAME:     %[[v:.*]]: vector<2x3xf32>,
+func.func @vector_multi_reduction_no_reduction_dims_nd(%arg0: vector<2x3xf32>, %acc: vector<2x3xf32>) -> vector<2x3xf32> {
+    %0 = vector.multi_reduction <add>, %arg0, %acc [] : vector<2x3xf32> to vector<2x3xf32>
+//       CHECK:   return %[[v]] : vector<2x3xf32>
+    return %0 : vector<2x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @masked_vector_multi_reduction_no_reduction_dims_nd(
+//  CHECK-SAME:     %[[VAL_0:.*]]: vector<2x3xf32>, %{{.*}}: vector<2x3xf32>,
+func.func @masked_vector_multi_reduction_no_reduction_dims_nd(%arg0: vector<2x3xf32>, %acc: vector<2x3xf32>, %mask: vector<2x3xi1>) -> vector<2x3xf32> {
+    %0 = vector.mask %mask { vector.multi_reduction <add>, %arg0, %acc [] : vector<2x3xf32> to vector<2x3xf32> } : vector<2x3xi1> -> vector<2x3xf32>
+//       CHECK:   return %[[VAL_0]] : vector<2x3xf32>
+    return %0 : vector<2x3xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @vector_multi_reduction_unit_dimensions(
 //  CHECK-SAME: %[[SOURCE:.+]]: vector<5x1x4x1x20xf32>, %[[ACC:.+]]: vector<5x4x20xf32>
 func.func @vector_multi_reduction_unit_dimensions(%source: vector<5x1x4x1x20xf32>, %acc: vector<5x4x20xf32>) -> vector<5x4x20xf32> {