[Mlir-commits] [mlir] [mlir][vector] Fix multi-reduction unrolling with a lower rank target shape (PR #216799)
Hamza Qureshi
llvmlistbot at llvm.org
Fri Aug 21 04:05:55 PDT 2026
https://github.com/hamzaqureshi5 updated https://github.com/llvm/llvm-project/pull/216799
>From 23564178f7d0c9714419fde472054ec7b906cb0b Mon Sep 17 00:00:00 2001
From: hamzaqureshi5 <hamza7771.861 at gmail.com>
Date: Mon, 17 Aug 2026 23:06:58 +0500
Subject: [PATCH 1/3] [mlir][vector] Fix multi-reduction unrolling with a lower
rank target shape
The unroll target shape may have fewer dimensions than the vector being
unrolled, in which case it applies to the trailing dimensions;
computeShapeRatio accepts that and UnrollElementwisePattern handles it by
padding the shape with leading unit dimensions.
UnrollMultiReductionPattern used the target shape directly, both as the sizes
of the slice taken from the source and to index the dimensions of the
reduction. With a target shape of a lower rank than the source, the slice
sizes had fewer entries than its offsets, which tripped an assertion in
inferStridedSliceOpResultType.
Pad the target shape the same way before using it.
Fixes #216640
---
.../Vector/Transforms/VectorUnroll.cpp | 19 +++++++++-----
.../Dialect/Vector/vector-unroll-options.mlir | 26 +++++++++++++++++++
2 files changed, 39 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 62869111496d1..565373b1866e5 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -383,17 +383,24 @@ struct UnrollMultiReductionPattern
Location loc = reductionOp.getLoc();
auto resultType = reductionOp->getResult(0).getType();
+ // A target shape with fewer dimensions than the source vector applies to
+ // its trailing dimensions. Add leading unit dimensions so that it can be
+ // used to slice the source and to index its dimensions, the same way
+ // UnrollElementwisePattern does.
+ SmallVector<int64_t> adjustedTargetShape(originalSize.size(), 1);
+ llvm::copy(*targetShape, adjustedTargetShape.end() - targetShape->size());
+
// Handle scalar result case: all dimensions are reduced.
// Each source tile is reduced to a scalar, and partial results are
// chained through the accumulator operand.
if (resultType.isIntOrFloat()) {
Value accumulator = reductionOp.getAcc();
for (SmallVector<int64_t> offsets :
- StaticTileOffsetRange(originalSize, *targetShape)) {
+ StaticTileOffsetRange(originalSize, adjustedTargetShape)) {
SmallVector<int64_t> operandStrides(offsets.size(), 1);
Value slicedOperand =
rewriter.createOrFold<vector::ExtractStridedSliceOp>(
- loc, reductionOp.getSource(), offsets, *targetShape,
+ loc, reductionOp.getSource(), offsets, adjustedTargetShape,
operandStrides);
Operation *newOp = cloneOpWithOperandsAndTypes(
rewriter, loc, reductionOp, {slicedOperand, accumulator},
@@ -413,20 +420,20 @@ struct UnrollMultiReductionPattern
// Stride of the ratios, this gives us the offsets of sliceCount in a basis
// of multiples of the targetShape.
for (SmallVector<int64_t> offsets :
- StaticTileOffsetRange(originalSize, *targetShape)) {
+ StaticTileOffsetRange(originalSize, adjustedTargetShape)) {
SmallVector<Value> operands;
SmallVector<int64_t> operandStrides(offsets.size(), 1);
Value slicedOperand =
rewriter.createOrFold<vector::ExtractStridedSliceOp>(
- loc, reductionOp.getSource(), offsets, *targetShape,
+ loc, reductionOp.getSource(), offsets, adjustedTargetShape,
operandStrides);
operands.push_back(slicedOperand);
SmallVector<int64_t> dstShape;
SmallVector<int64_t> destOffset;
- for (size_t i : llvm::seq(size_t(0), targetShape->size())) {
+ for (size_t i : llvm::seq(size_t(0), adjustedTargetShape.size())) {
if (!reductionOp.isReducedDim(i)) {
destOffset.push_back(offsets[i]);
- dstShape.push_back((*targetShape)[i]);
+ dstShape.push_back(adjustedTargetShape[i]);
}
}
Value acc;
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index da058d5d2410d..a0b39c798abd2 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -864,3 +864,29 @@ func.func @deinterleave_2d(%v: vector<4x8xi32>) -> (vector<4x4xi32>, vector<4x4x
// CHECK: {{.*}} = vector.insert_strided_slice {{.*}}, {{.*}} {offsets = [2, 0], strides = [1, 1]} : vector<2x4xi32> into vector<4x4xi32>
// CHECK: {{.*}} = vector.insert_strided_slice {{.*}}, {{.*}} {offsets = [2, 0], strides = [1, 1]} : vector<2x4xi32> into vector<4x4xi32>
// CHECK: return {{.*}}, {{.*}} : vector<4x4xi32>, vector<4x4xi32>
+
+// -----
+
+// The native unroll shape {2, 2} has fewer dimensions than the source vector.
+// It applies to the trailing dimensions, so the leading ones are sliced one
+// element at a time.
+
+func.func @vector_multi_reduction_rank_mismatch(%v : vector<2x2x4xf32>, %acc: vector<2x2xf32>) -> vector<2x2xf32> {
+ %0 = vector.multi_reduction #vector.kind<add>, %v, %acc [2] : vector<2x2x4xf32> to vector<2x2xf32>
+ return %0 : vector<2x2xf32>
+}
+// CHECK-LABEL: func @vector_multi_reduction_rank_mismatch
+// CHECK: %[[V0:.*]] = arith.constant dense<0.000000e+00> : vector<2x2xf32>
+// CHECK: %[[E0:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0, 0], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[ACC0:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0], sizes = [1, 2], strides = [1, 1]} : vector<2x2xf32> to vector<1x2xf32>
+// CHECK: %[[R0:.*]] = vector.multi_reduction <add>, %[[E0]], %[[ACC0]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// CHECK: %[[E1:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0, 2], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[R1:.*]] = vector.multi_reduction <add>, %[[E1]], %[[R0]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// CHECK: %[[E2:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 0, 0], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[ACC1:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 0], sizes = [1, 2], strides = [1, 1]} : vector<2x2xf32> to vector<1x2xf32>
+// CHECK: %[[R2:.*]] = vector.multi_reduction <add>, %[[E2]], %[[ACC1]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// CHECK: %[[E3:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 0, 2], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[R3:.*]] = vector.multi_reduction <add>, %[[E3]], %[[R2]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// CHECK: %[[V1:.*]] = vector.insert_strided_slice %[[R1]], %[[V0]] {offsets = [0, 0], strides = [1, 1]} : vector<1x2xf32> into vector<2x2xf32>
+// CHECK: %[[V2:.*]] = vector.insert_strided_slice %[[R3]], %[[V1]] {offsets = [1, 0], strides = [1, 1]} : vector<1x2xf32> into vector<2x2xf32>
+// CHECK: return %[[V2]] : vector<2x2xf32>
>From 6dbf6d167a6b9a55c92d64cfa376130a64a44d42 Mon Sep 17 00:00:00 2001
From: Hamza Qureshi <63870077+hamzaqureshi5 at users.noreply.github.com>
Date: Tue, 18 Aug 2026 10:53:39 +0500
Subject: [PATCH 2/3] Update
mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
Co-authored-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 565373b1866e5..a95bfba0814bd 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -385,8 +385,7 @@ struct UnrollMultiReductionPattern
// A target shape with fewer dimensions than the source vector applies to
// its trailing dimensions. Add leading unit dimensions so that it can be
- // used to slice the source and to index its dimensions, the same way
- // UnrollElementwisePattern does.
+ // used to slice the source and to index its dimensions.
SmallVector<int64_t> adjustedTargetShape(originalSize.size(), 1);
llvm::copy(*targetShape, adjustedTargetShape.end() - targetShape->size());
>From 6846cb17905493119253859ae0583994637d3b67 Mon Sep 17 00:00:00 2001
From: hamzaqureshi5 <hamza7771.861 at gmail.com>
Date: Fri, 21 Aug 2026 16:05:14 +0500
Subject: [PATCH 3/3] Clarify the rank mismatch comment in the test
---
mlir/test/Dialect/Vector/vector-unroll-options.mlir | 6 ++----
1 file changed, 2 insertions(+), 4 deletions(-)
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index a0b39c798abd2..cc84e1f52ccbb 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -867,10 +867,8 @@ func.func @deinterleave_2d(%v: vector<4x8xi32>) -> (vector<4x4xi32>, vector<4x4x
// -----
-// The native unroll shape {2, 2} has fewer dimensions than the source vector.
-// It applies to the trailing dimensions, so the leading ones are sliced one
-// element at a time.
-
+// TargetShape [2, 2] has a lower rank than the reduced source <2x2x4>, so it
+// applies to the trailing dimensions and is padded to [1, 2, 2].
func.func @vector_multi_reduction_rank_mismatch(%v : vector<2x2x4xf32>, %acc: vector<2x2xf32>) -> vector<2x2xf32> {
%0 = vector.multi_reduction #vector.kind<add>, %v, %acc [2] : vector<2x2x4xf32> to vector<2x2xf32>
return %0 : vector<2x2xf32>
More information about the Mlir-commits
mailing list