[Mlir-commits] [mlir] [MLIR][Vector] Relax shape_cast unrolling to per-reassociation-group contiguity (PR #205684)
Jianhui Li
llvmlistbot at llvm.org
Thu Jul 9 22:02:11 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/205684
>From e7ea262a41af79d3f18f497a20fff69ba40bf6a1 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 24 Jun 2026 21:22:36 +0000
Subject: [PATCH 1/3] relaxed UnrollShapeCastPattern in VectorUnroll.cpp so it
unrolls shape_casts whose target tile is contiguous per reassociation group
rather than requiring it to be contiguous across the entire result vector
---
.../Vector/Transforms/VectorUnroll.cpp | 119 +++++++++++++++---
.../Dialect/Vector/vector-unroll-options.mlir | 44 +++++++
.../Dialect/Vector/TestVectorTransforms.cpp | 11 ++
3 files changed, 154 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index b33828d5d5867..c3dc132195ef5 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -1285,20 +1285,63 @@ calculateSourceOffsets(ArrayRef<int64_t> resultOffsets,
return delinearize(linearIndex, computeStrides(sourceShape));
}
+/// A maximal aligned range of source dims [srcBegin, srcEnd) and result dims
+/// [resBegin, resEnd) of a `vector.shape_cast` that hold equal element counts.
+struct ShapeCastReassociationGroup {
+ int64_t srcBegin, srcEnd;
+ int64_t resBegin, resEnd;
+};
+
+/// Splits a shape_cast from `sourceShape` to `resultShape` into reassociation
+/// groups (trailing unit dims absorbed). Returns nullopt if shapes misalign.
+/// E.g. [8, 32, 32] -> [256, 32] ==> {[0,2)->[0,1)}, {[2,3)->[1,2)}
+static std::optional<SmallVector<ShapeCastReassociationGroup>>
+computeShapeCastGroups(ArrayRef<int64_t> sourceShape,
+ ArrayRef<int64_t> resultShape) {
+ SmallVector<ShapeCastReassociationGroup> groups;
+ int64_t si = 0, ri = 0;
+ int64_t srcRank = sourceShape.size(), resRank = resultShape.size();
+ while (si < srcRank && ri < resRank) {
+ int64_t srcBegin = si, resBegin = ri;
+ int64_t srcProd = sourceShape[si++];
+ int64_t resProd = resultShape[ri++];
+ // Grow the smaller side until both groups span the same element count.
+ while (srcProd != resProd) {
+ if (srcProd < resProd) {
+ if (si >= srcRank)
+ return std::nullopt;
+ srcProd *= sourceShape[si++];
+ } else {
+ if (ri >= resRank)
+ return std::nullopt;
+ resProd *= resultShape[ri++];
+ }
+ }
+ // Absorb trailing unit dimensions into the current group.
+ while (si < srcRank && sourceShape[si] == 1)
+ ++si;
+ while (ri < resRank && resultShape[ri] == 1)
+ ++ri;
+ groups.push_back({srcBegin, si, resBegin, ri});
+ }
+ if (si != srcRank || ri != resRank)
+ return std::nullopt;
+ return groups;
+}
+
/// This pattern unrolls `vector.shape_cast` operations according to the
/// provided target unroll shape. It unrolls a large shape cast into smaller
/// shape casts by extracting contiguous slices from the source vector, casting
/// each slice to the target shape, and assembling the result by inserting each
/// computed segment into the appropriate offset of the result vector.
///
-/// This pattern only applies when contiguous slices can be extracted from the
-/// source vector and inserted into the result vector such that each slice
-/// remains a valid vector (and not decompose to scalars). In these cases, the
-/// unrolling proceeds as:
+/// The target tile need only be contiguous within each reassociation group of
+/// the cast (not in the whole result vector), so that each extracted slice
+/// remains a valid vector. The unrolling proceeds as:
/// vector.extract_strided_slice -> vector.shape_cast (on the slice) ->
/// vector.insert_strided_slice.
///
-/// Example:
+/// Example (single group):
/// Given a shape cast operation:
/// %0 = vector.shape_cast %src : vector<8x2xf32> to vector<4x4xf32>
///
@@ -1316,6 +1359,18 @@ calculateSourceOffsets(ArrayRef<int64_t> resultOffsets,
/// %i1 = vector.insert_strided_slice %sc1, %i0 [2, 0], [1, 1]
/// : vector<2x4xf32> into vector<4x4xf32>
///
+/// Example (multiple groups): with target tile <8x1x4>, the tile is strided in
+/// the result <8x1x32> but contiguous per group (8|32 -> 8x1|32), so the
+/// matching strided box <8x4> is extracted from the source:
+/// %0 = vector.shape_cast %src : vector<8x32xf32> to vector<8x1x32xf32>
+///
+/// %s0 = vector.extract_strided_slice %src [0, 0], [8, 4], [1, 1]
+/// : vector<8x32xf32> to vector<8x4xf32>
+/// %sc0 = vector.shape_cast %s0 : vector<8x4xf32> to vector<8x1x4xf32>
+/// %i0 = vector.insert_strided_slice %sc0, %zero [0, 0, 0], [1, 1, 1]
+/// : vector<8x1x4xf32> into vector<8x1x32xf32>
+/// // ... repeat for the remaining slices.
+///
struct UnrollShapeCastPattern : public OpRewritePattern<vector::ShapeCastOp> {
UnrollShapeCastPattern(MLIRContext *context,
const vector::UnrollVectorOptions &options,
@@ -1335,20 +1390,44 @@ struct UnrollShapeCastPattern : public OpRewritePattern<vector::ShapeCastOp> {
ArrayRef<int64_t> sourceShape = sourceType.getShape();
ArrayRef<int64_t> resultShape = resultType.getShape();
- if (!isContiguous(*targetShape, resultShape))
+ // The cast factors into reassociation groups; the target tile only needs to
+ // be contiguous within each group, not in the whole result vector.
+ std::optional<SmallVector<ShapeCastReassociationGroup>> groups =
+ computeShapeCastGroups(sourceShape, resultShape);
+ if (!groups)
return rewriter.notifyMatchFailure(
- shapeCastOp, "Only supports cases where target shape is "
- "contiguous in result vector shape");
-
- int64_t targetElements = ShapedType::getNumElements(*targetShape);
-
- // Calculate the shape to extract from source.
- std::optional<SmallVector<int64_t>> extractShape =
- calculateSourceExtractShape(sourceShape, targetElements);
- if (!extractShape)
- return rewriter.notifyMatchFailure(
- shapeCastOp,
- "cannot extract target number of elements contiguously from source");
+ shapeCastOp, "cannot align source and result reassociation groups");
+
+ // The tile is right-aligned against the result; left-pad with 1s so it can
+ // be indexed per group.
+ SmallVector<int64_t> paddedTarget(resultShape.size(), 1);
+ llvm::copy(*targetShape,
+ paddedTarget.end() - static_cast<int64_t>(targetShape->size()));
+
+ // Validate per-group contiguity and build the source extract shape.
+ SmallVector<int64_t> extractShapeStorage;
+ for (const ShapeCastReassociationGroup &g : *groups) {
+ ArrayRef<int64_t> resSub =
+ resultShape.slice(g.resBegin, g.resEnd - g.resBegin);
+ ArrayRef<int64_t> tgtSub = ArrayRef<int64_t>(paddedTarget)
+ .slice(g.resBegin, g.resEnd - g.resBegin);
+ if (!isContiguous(tgtSub, resSub))
+ return rewriter.notifyMatchFailure(
+ shapeCastOp, "target shape is not contiguous within a "
+ "reassociation group of the result vector shape");
+
+ ArrayRef<int64_t> srcSub =
+ sourceShape.slice(g.srcBegin, g.srcEnd - g.srcBegin);
+ int64_t groupTargetElements = ShapedType::getNumElements(tgtSub);
+ std::optional<SmallVector<int64_t>> groupExtract =
+ calculateSourceExtractShape(srcSub, groupTargetElements);
+ if (!groupExtract)
+ return rewriter.notifyMatchFailure(
+ shapeCastOp, "cannot extract the target number of elements "
+ "contiguously from a source reassociation group");
+ extractShapeStorage.append(groupExtract->begin(), groupExtract->end());
+ }
+ ArrayRef<int64_t> extractShape = extractShapeStorage;
Location loc = shapeCastOp.getLoc();
@@ -1359,7 +1438,7 @@ struct UnrollShapeCastPattern : public OpRewritePattern<vector::ShapeCastOp> {
VectorType targetType =
VectorType::get(*targetShape, sourceType.getElementType());
- SmallVector<int64_t> extractStrides(extractShape->size(), 1);
+ SmallVector<int64_t> extractStrides(extractShape.size(), 1);
SmallVector<int64_t> insertStrides(targetShape->size(), 1);
for (SmallVector<int64_t> resultOffsets :
@@ -1367,7 +1446,7 @@ struct UnrollShapeCastPattern : public OpRewritePattern<vector::ShapeCastOp> {
SmallVector<int64_t> sourceOffsets =
calculateSourceOffsets(resultOffsets, sourceShape, resultShape);
Value sourceChunk = rewriter.createOrFold<vector::ExtractStridedSliceOp>(
- loc, shapeCastOp.getSource(), sourceOffsets, *extractShape,
+ loc, shapeCastOp.getSource(), sourceOffsets, extractShape,
extractStrides);
Value targetChunk = rewriter.createOrFold<vector::ShapeCastOp>(
loc, targetType, sourceChunk);
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index bb6fc4e38813d..58ca27b50e2ee 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -668,6 +668,50 @@ func.func @shape_cast_with_all_unit_target_shape(%v: vector<2xf32>) -> vector<2x
// CHECK: %[[I1:.*]] = vector.insert_strided_slice %[[SC1]], %[[I0]] {offsets = [1, 0], strides = [1, 1]} : vector<1x1xf32> into vector<2x1xf32>
// CHECK: return %[[I1]] : vector<2x1xf32>
+
+// Target tile [8, 1, 4] is strided in result <8x1x32> but contiguous per
+// reassociation group (8|32 -> 8x1|32). TargetShape is [8, 1, 4].
+func.func @shape_cast_multi_group_rank_increasing(%v: vector<8x32xf32>) -> vector<8x1x32xf32> {
+ %0 = vector.shape_cast %v : vector<8x32xf32> to vector<8x1x32xf32>
+ return %0 : vector<8x1x32xf32>
+}
+
+// CHECK-LABEL: func @shape_cast_multi_group_rank_increasing
+// CHECK-SAME: (%[[V:.*]]: vector<8x32xf32>) -> vector<8x1x32xf32> {
+// CHECK: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<8x1x32xf32>
+// CHECK: %[[S0:.*]] = vector.extract_strided_slice %[[V]] {offsets = [0, 0], sizes = [8, 4], strides = [1, 1]} : vector<8x32xf32> to vector<8x4xf32>
+// CHECK: %[[SC0:.*]] = vector.shape_cast %[[S0]] : vector<8x4xf32> to vector<8x1x4xf32>
+// CHECK: %[[I0:.*]] = vector.insert_strided_slice %[[SC0]], %[[CST]] {offsets = [0, 0, 0], strides = [1, 1, 1]} : vector<8x1x4xf32> into vector<8x1x32xf32>
+// CHECK: %[[S1:.*]] = vector.extract_strided_slice %[[V]] {offsets = [0, 4], sizes = [8, 4], strides = [1, 1]} : vector<8x32xf32> to vector<8x4xf32>
+// CHECK: %[[SC1:.*]] = vector.shape_cast %[[S1]] : vector<8x4xf32> to vector<8x1x4xf32>
+// CHECK: %[[I1:.*]] = vector.insert_strided_slice %[[SC1]], %[[I0]] {offsets = [0, 0, 4], strides = [1, 1, 1]} : vector<8x1x4xf32> into vector<8x1x32xf32>
+// CHECK: return
+
+
+// Target tile [2, 2] is strided in result <4x4> but contiguous per
+// reassociation group (2x2|4 -> 4|4). TargetShape is [2, 2].
+func.func @shape_cast_multi_group_rank_decreasing(%v: vector<2x2x4xf32>) -> vector<4x4xf32> {
+ %0 = vector.shape_cast %v : vector<2x2x4xf32> to vector<4x4xf32>
+ return %0 : vector<4x4xf32>
+}
+
+// CHECK-LABEL: func @shape_cast_multi_group_rank_decreasing
+// CHECK-SAME: (%[[V:.*]]: vector<2x2x4xf32>) -> vector<4x4xf32> {
+// CHECK: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<4x4xf32>
+// CHECK: %[[S0:.*]] = vector.extract_strided_slice %[[V]] {offsets = [0, 0, 0], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[SC0:.*]] = vector.shape_cast %[[S0]] : vector<1x2x2xf32> to vector<2x2xf32>
+// CHECK: %[[I0:.*]] = vector.insert_strided_slice %[[SC0]], %[[CST]] {offsets = [0, 0], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
+// CHECK: %[[S1:.*]] = vector.extract_strided_slice %[[V]] {offsets = [0, 0, 2], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[SC1:.*]] = vector.shape_cast %[[S1]] : vector<1x2x2xf32> to vector<2x2xf32>
+// CHECK: %[[I1:.*]] = vector.insert_strided_slice %[[SC1]], %[[I0]] {offsets = [0, 2], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
+// CHECK: %[[S2:.*]] = vector.extract_strided_slice %[[V]] {offsets = [1, 0, 0], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[SC2:.*]] = vector.shape_cast %[[S2]] : vector<1x2x2xf32> to vector<2x2xf32>
+// CHECK: %[[I2:.*]] = vector.insert_strided_slice %[[SC2]], %[[I1]] {offsets = [2, 0], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
+// CHECK: %[[S3:.*]] = vector.extract_strided_slice %[[V]] {offsets = [1, 0, 2], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[SC3:.*]] = vector.shape_cast %[[S3]] : vector<1x2x2xf32> to vector<2x2xf32>
+// CHECK: %[[I3:.*]] = vector.insert_strided_slice %[[SC3]], %[[I2]] {offsets = [2, 2], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
+// CHECK: return %[[I3]] : vector<4x4xf32>
+
// -----
// Test BitCastOp unrolling - target shape [4, 4]
diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
index 043181c16c759..af0e7e59eab63 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -206,6 +206,17 @@ struct TestVectorUnrollingPatterns
resultShape[1] == 1) {
return SmallVector<int64_t>{1, 1};
}
+ // Multi-group cases: tile contiguous per reassociation group
+ // but strided in the whole result.
+ auto sourceShape = shapeCast.getSourceVectorType().getShape();
+ if (resultShape.size() == 3 && resultShape[0] == 8 &&
+ resultShape[1] == 1 && resultShape[2] == 32) {
+ return SmallVector<int64_t>{8, 1, 4};
+ }
+ if (sourceShape.size() == 3 && resultShape.size() == 2 &&
+ resultShape[0] == 4 && resultShape[1] == 4) {
+ return SmallVector<int64_t>{2, 2};
+ }
// Default case: [2,4] for all tests.
return SmallVector<int64_t>{2, 4};
})
>From f14d24ad2e43bc60719a7329c7e12b80848a7ecb Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 10 Jul 2026 04:55:22 +0000
Subject: [PATCH 2/3] [MLIR][Vector] Add negative shape_cast unroll tests and
NOP note
Add two negative tests for the new per-reassociation-group checks in
UnrollShapeCastPattern: one where the target tile is not contiguous
within a result group, and one where the tile cannot be extracted
contiguously from a source group. Also document that this unrolling
assumes the introduced strided slices lower to a NOP, per review.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../Vector/Transforms/VectorUnroll.cpp | 5 ++++
.../Dialect/Vector/vector-unroll-options.mlir | 26 +++++++++++++++++++
.../Dialect/Vector/TestVectorTransforms.cpp | 4 +++
3 files changed, 35 insertions(+)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index c3dc132195ef5..62869111496d1 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -1341,6 +1341,11 @@ computeShapeCastGroups(ArrayRef<int64_t> sourceShape,
/// vector.extract_strided_slice -> vector.shape_cast (on the slice) ->
/// vector.insert_strided_slice.
///
+/// NOTE: This replaces a NOP `vector.shape_cast` with strided slices. Per-group
+/// contiguity keeps those slices contiguous, so they are expected to lower to a
+/// NOP too. Targets where strided slices do not lower to a NOP should not use
+/// this pattern, or should pick a tile that avoids introducing such slices.
+///
/// Example (single group):
/// Given a shape cast operation:
/// %0 = vector.shape_cast %src : vector<8x2xf32> to vector<4x4xf32>
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 139794aeaf59b..14edc15d9e666 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -747,6 +747,32 @@ func.func @shape_cast_multi_group_rank_decreasing(%v: vector<2x2x4xf32>) -> vect
// CHECK: %[[I3:.*]] = vector.insert_strided_slice %[[SC3]], %[[I2]] {offsets = [2, 2], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
// CHECK: return %[[I3]] : vector<4x4xf32>
+// Negative multi-group case: target tile [2, 2, 2] is not contiguous within
+// the result reassociation group [8, 4], so the cast is left un-unrolled.
+func.func @negative_shape_cast_multi_group_target_not_contiguous(%v: vector<2x32xf32>) -> vector<2x8x4xf32> {
+ %0 = vector.shape_cast %v : vector<2x32xf32> to vector<2x8x4xf32>
+ return %0 : vector<2x8x4xf32>
+}
+
+// CHECK-LABEL: func @negative_shape_cast_multi_group_target_not_contiguous
+// CHECK-SAME: (%[[V:.*]]: vector<2x32xf32>) -> vector<2x8x4xf32> {
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[V]] : vector<2x32xf32> to vector<2x8x4xf32>
+// CHECK: return %[[SC]] : vector<2x8x4xf32>
+
+
+// Negative multi-group case: the target tile is contiguous within the result
+// group [24], but its elements cannot be extracted contiguously from the
+// source group [8, 3], so the cast is left un-unrolled.
+func.func @negative_shape_cast_multi_group_source_not_determinable(%v: vector<2x8x3xf32>) -> vector<2x24xf32> {
+ %0 = vector.shape_cast %v : vector<2x8x3xf32> to vector<2x24xf32>
+ return %0 : vector<2x24xf32>
+}
+
+// CHECK-LABEL: func @negative_shape_cast_multi_group_source_not_determinable
+// CHECK-SAME: (%[[V:.*]]: vector<2x8x3xf32>) -> vector<2x24xf32> {
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[V]] : vector<2x8x3xf32> to vector<2x24xf32>
+// CHECK: return %[[SC]] : vector<2x24xf32>
+
// -----
// Test BitCastOp unrolling - target shape [4, 4]
diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
index 52e3444fe6c3b..4523a4cd3c486 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -222,6 +222,10 @@ struct TestVectorUnrollingPatterns
resultShape[0] == 4 && resultShape[1] == 4) {
return SmallVector<int64_t>{2, 2};
}
+ if (resultShape.size() == 3 && resultShape[0] == 2 &&
+ resultShape[1] == 8 && resultShape[2] == 4) {
+ return SmallVector<int64_t>{2, 2, 2};
+ }
// Default case: [2,4] for all tests.
return SmallVector<int64_t>{2, 4};
})
>From ddce50d6943b3bed54d6e3456abb578cf6ffea38 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 10 Jul 2026 05:01:55 +0000
Subject: [PATCH 3/3] polish comment in test
---
mlir/test/Dialect/Vector/vector-unroll-options.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 14edc15d9e666..da058d5d2410d 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -760,7 +760,7 @@ func.func @negative_shape_cast_multi_group_target_not_contiguous(%v: vector<2x32
// CHECK: return %[[SC]] : vector<2x8x4xf32>
-// Negative multi-group case: the target tile is contiguous within the result
+// Negative multi-group case: the target tile [2, 4] is contiguous within the result
// group [24], but its elements cannot be extracted contiguously from the
// source group [8, 3], so the cast is left un-unrolled.
func.func @negative_shape_cast_multi_group_source_not_determinable(%v: vector<2x8x3xf32>) -> vector<2x24xf32> {
More information about the Mlir-commits
mailing list