[Mlir-commits] [mlir] b4cf9c7 - [MLIR][Vector] Relax shape_cast unrolling to per-reassociation-group contiguity (#205684)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 10 08:51:51 PDT 2026
Author: Jianhui Li
Date: 2026-07-10T08:51:46-07:00
New Revision: b4cf9c7c89ec541495c38c4e34fa8967713bf843
URL: https://github.com/llvm/llvm-project/commit/b4cf9c7c89ec541495c38c4e34fa8967713bf843
DIFF: https://github.com/llvm/llvm-project/commit/b4cf9c7c89ec541495c38c4e34fa8967713bf843.diff
LOG: [MLIR][Vector] Relax shape_cast unrolling to per-reassociation-group contiguity (#205684)
UnrollShapeCastPattern previously only unrolled a vector.shape_cast when
the target unroll tile was contiguous in the whole result vector
(isContiguous(targetShape, resultShape)). This rejected valid cases such
as:
```mlir
%0 = vector.shape_cast %src : vector<8x32xf8> to vector<8x1x32xf8> // target tile [8, 1, 4]
%1 = vector.shape_cast %src : vector<8x32x32xf8> to vector<256x32xf8> // target tile [16, 4]
```
This PR factors the source and result shapes of a shape_cast into
independent reassociation groups — maximal aligned source/result dim
ranges that hold equal element counts:
```
- 8x32 → 8x1x32 ⇒ groups {8 ↔ 8x1}, {32 ↔ 32}
- 8x32x32 → 256x32 ⇒ groups {8x32 ↔ 256}, {32 ↔ 32}
```
Instead of requiring the target tile to be contiguous in the whole
result vector, it only checks that the tile is contiguous within each
group.
This criterion is correct because a shape_cast preserves row-major
linear order. The groups partition the dimensions into blocks whose
element ranges don't overlap, so the global linear index breaks down
into one independent sub-index per group. A tile that is contiguous
within each group therefore corresponds to a slice that is contiguous
within each group of both the source and result.
Assist-by-Claude
---------
Co-authored-by: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
Added:
Modified:
mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
mlir/test/Dialect/Vector/vector-unroll-options.mlir
mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index b33828d5d5867f..62869111496d19 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -1285,20 +1285,68 @@ 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:
+/// 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>
///
@@ -1316,6 +1364,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 +1395,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 +1443,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 +1451,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 0ec71620c5324d..da058d5d2410d5 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -703,6 +703,76 @@ 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>
+
+// 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 [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> {
+ %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 1a4bef664fbe1f..4523a4cd3c4867 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -211,6 +211,21 @@ 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};
+ }
+ 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};
})
More information about the Mlir-commits
mailing list