[Mlir-commits] [mlir] [MLIR][Vector] Fix ShapeCastBroadcastFolder when shape_cast changes non-unit dim strides (PR #227574)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 14 11:48:14 PDT 2026
https://github.com/yueshengys updated https://github.com/llvm/llvm-project/pull/227574
>From ee25df2d1c8d18fb72bfcd3f7222e93cee3f5b2e Mon Sep 17 00:00:00 2001
From: yueshengys <135203619+yueshengys at users.noreply.github.com>
Date: Wed, 30 Sep 2026 06:01:08 +0000
Subject: [PATCH 1/4] [MLIR][Vector] Fix ShapeCastBroadcastFolder when
shape_cast changes non-unit dim strides
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 35 +++++----------
.../canonicalize/vector-to-shape-cast.mlir | 44 +++++++++++++++++++
2 files changed, 54 insertions(+), 25 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index d5fd5095b18b00..678d06e2a8db19 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -7145,22 +7145,6 @@ class ShapeCastCreateMaskFolderBoundaryUnitDim final
}
};
-// vector.broadcast has two distinct semantic modes: duplication across leading
-// dimensions, and stretching across inner dimensions. This helper returns the
-// product of the inner-dimension stretching factors.
-int64_t getBroadcastStretchingFactor(ArrayRef<int64_t> srcShape,
- ArrayRef<int64_t> dstShape) {
- int stretchingFactor = 1;
- int numLeadingDims = dstShape.size() - srcShape.size();
- for (int i = 0, e = srcShape.size(); i < e; i++) {
- int64_t dstDim = dstShape[numLeadingDims + i];
- if (srcShape[i] == 1 && dstDim != 1) {
- stretchingFactor *= dstDim;
- }
- }
- return stretchingFactor;
-}
-
/// Pattern to rewrite Y = ShapeCast(Broadcast(X)) as Y = Broadcast(X)
class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
public:
@@ -7194,16 +7178,17 @@ class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
BroadcastableToResult::Success) {
return failure();
}
- // Avoid folding if this would result in switching between the two
- // distinct semantic modes of vector.broadcast (duplication vs
- // stretching). See https://github.com/llvm/llvm-project/issues/190614.
- // This is detected by a change in the stretching factor. However if the
- // source has a single element, there is no ambiguity.
- if (srcVectorType.getNumElements() != 1) {
- if (getBroadcastStretchingFactor(srcShape, dstShape) !=
- getBroadcastStretchingFactor(srcShape, broadcastShape)) {
+ // `shape_cast` preserves the flattened 1-D (row-major) element order.
+ // Thus, `shape_cast(broadcast(src, broadcastShape), dstShape)` produces
+ // the exact same elements as `broadcast(src, dstShape)` if and only if
+ // every non-unit dimension of `src` has the same linear stride in both
+ // `broadcastShape` and `dstShape`.
+ int64_t bcastStride = 1, dstStride = 1;
+ for (size_t i = 1, e = srcShape.size(); i <= e; ++i) {
+ if (srcShape[e - i] != 1 && bcastStride != dstStride)
return failure();
- }
+ bcastStride *= broadcastShape.take_back(i).front();
+ dstStride *= dstShape.take_back(i).front();
}
}
diff --git a/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir b/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
index 7ae190d0b3c56c..bb9ca038fce783 100644
--- a/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
@@ -144,3 +144,47 @@ func.func @no_fold_bcast_mode_switch(%arg0: vector<2x1xf32>) -> vector<2x2xf32>
%1 = vector.shape_cast %0 : vector<2x2x1xf32> to vector<2x2xf32>
return %1 : vector<2x2xf32>
}
+
+// -----
+
+// CHECK-LABEL: @no_fold_bcast_axis_shift
+// CHECK: vector.broadcast %{{.*}} : vector<1x4x1xf32> to vector<1x4x4xf32>
+// CHECK-NEXT: vector.shape_cast %{{.*}} : vector<1x4x4xf32> to vector<4x4x1xf32>
+func.func @no_fold_bcast_axis_shift(%arg0: vector<1x4x1xf32>) -> vector<4x4x1xf32> {
+ %0 = vector.broadcast %arg0 : vector<1x4x1xf32> to vector<1x4x4xf32>
+ %1 = vector.shape_cast %0 : vector<1x4x4xf32> to vector<4x4x1xf32>
+ return %1 : vector<4x4x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_leading_dims
+// CHECK: %[[RES:.*]] = vector.broadcast %{{.*}} : vector<3xf32> to vector<8x3xf32>
+// CHECK-NEXT: return %[[RES]] : vector<8x3xf32>
+func.func @fold_bcast_leading_dims(%arg0: vector<3xf32>) -> vector<8x3xf32> {
+ %0 = vector.broadcast %arg0 : vector<3xf32> to vector<2x4x3xf32>
+ %1 = vector.shape_cast %0 : vector<2x4x3xf32> to vector<8x3xf32>
+ return %1 : vector<8x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_src_leading_unit_dims
+// CHECK: %[[RES:.*]] = vector.broadcast %{{.*}} : vector<1x3xf32> to vector<4x3xf32>
+// CHECK-NEXT: return %[[RES]] : vector<4x3xf32>
+func.func @fold_bcast_src_leading_unit_dims(%arg0: vector<1x3xf32>) -> vector<4x3xf32> {
+ %0 = vector.broadcast %arg0 : vector<1x3xf32> to vector<4x1x3xf32>
+ %1 = vector.shape_cast %0 : vector<4x1x3xf32> to vector<4x3xf32>
+ return %1 : vector<4x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_consecutive_unit_dims
+// CHECK: %[[RES:.*]] = vector.broadcast %{{.*}} : vector<2x1x1x3xf32> to vector<2x2x2x3xf32>
+// CHECK-NEXT: return %[[RES]] : vector<2x2x2x3xf32>
+func.func @fold_bcast_consecutive_unit_dims(%arg0: vector<2x1x1x3xf32>) -> vector<2x2x2x3xf32> {
+ %0 = vector.broadcast %arg0 : vector<2x1x1x3xf32> to vector<2x1x4x3xf32>
+ %1 = vector.shape_cast %0 : vector<2x1x4x3xf32> to vector<2x2x2x3xf32>
+ return %1 : vector<2x2x2x3xf32>
+}
>From f90f67207b9cb7b8778f802c189b357d69d84859 Mon Sep 17 00:00:00 2001
From: yueshengys <135203619+yueshengys at users.noreply.github.com>
Date: Wed, 30 Sep 2026 22:07:53 +0000
Subject: [PATCH 2/4] [mlir][vector] Address review comments: handle scalable
dims and scalar sources
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 36 ++++++-
.../canonicalize/vector-to-shape-cast.mlir | 94 +++++++++++++++++--
2 files changed, 120 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 678d06e2a8db19..963faa50e5d230 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -7183,12 +7183,40 @@ class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
// the exact same elements as `broadcast(src, dstShape)` if and only if
// every non-unit dimension of `src` has the same linear stride in both
// `broadcastShape` and `dstShape`.
- int64_t bcastStride = 1, dstStride = 1;
+ //
+ // In physical memory, a scalable dimension of static size D has runtime
+ // size D * vscale. A linear stride through trailing dimensions is thus:
+ // stride = staticStride * (vscale ^ numScalableDims)
+ // Since vscale is an unknown runtime quantity (vscale >= 1), both
+ // staticStride and numScalableDims must match to guarantee stride
+ // equivalence across all targets.
+ struct Stride {
+ int64_t staticStride = 1;
+ int numScalableDims = 0;
+
+ bool operator==(const Stride &other) const {
+ return staticStride == other.staticStride &&
+ numScalableDims == other.numScalableDims;
+ }
+ bool operator!=(const Stride &other) const { return !(*this == other); }
+ };
+
+ Stride bcastStride, dstStride;
+ VectorType broadcastVectorType = broadcastOp.getResultVectorType();
+ ArrayRef<bool> bcastScalable = broadcastVectorType.getScalableDims();
+ ArrayRef<bool> dstScalable = dstVectorType.getScalableDims();
+ ArrayRef<bool> srcScalable = srcVectorType.getScalableDims();
+
for (size_t i = 1, e = srcShape.size(); i <= e; ++i) {
- if (srcShape[e - i] != 1 && bcastStride != dstStride)
+ bool isUnitDim = srcShape[e - i] == 1 && !srcScalable[e - i];
+ if (!isUnitDim && bcastStride != dstStride)
return failure();
- bcastStride *= broadcastShape.take_back(i).front();
- dstStride *= dstShape.take_back(i).front();
+ bcastStride.staticStride *= broadcastShape[broadcastShape.size() - i];
+ if (bcastScalable[broadcastShape.size() - i])
+ bcastStride.numScalableDims++;
+ dstStride.staticStride *= dstShape[dstShape.size() - i];
+ if (dstScalable[dstShape.size() - i])
+ dstStride.numScalableDims++;
}
}
diff --git a/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir b/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
index bb9ca038fce783..bffcc012cb5b0f 100644
--- a/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
@@ -130,12 +130,6 @@ func.func @extract_from_broadcast(%src: vector<1x1x1xf32>) -> vector<1xf32> {
// -----
-// https://github.com/llvm/llvm-project/issues/190614 — ShapeCastBroadcastFolder
-// must not rewrite shape_cast(broadcast) into a lower-rank broadcast when that
-// changes duplication (new leading dims) vs stretching (source unit dim).
-// For %arg0 = [[1], [2]]: broadcast to 2x2x1 then shape_cast to 2x2 yields
-// [[1, 2], [1, 2]]; folding to broadcast 2x1 -> 2x2 would incorrectly yield
-// [[1, 1], [2, 2]].
// CHECK-LABEL: @no_fold_bcast_mode_switch
// CHECK: vector.broadcast %{{.*}} : vector<2x1xf32> to vector<2x2x1xf32>
// CHECK-NEXT: vector.shape_cast %{{.*}} : vector<2x2x1xf32> to vector<2x2xf32>
@@ -188,3 +182,91 @@ func.func @fold_bcast_consecutive_unit_dims(%arg0: vector<2x1x1x3xf32>) -> vecto
%1 = vector.shape_cast %0 : vector<2x1x4x3xf32> to vector<2x2x2x3xf32>
return %1 : vector<2x2x2x3xf32>
}
+
+// -----
+
+// CHECK-LABEL: @no_fold_bcast_scalable_vs_fixed
+// CHECK: vector.broadcast %{{.*}} : vector<2x1x1xf32> to vector<2x[4]x2xf32>
+// CHECK-NEXT: vector.shape_cast %{{.*}} : vector<2x[4]x2xf32> to vector<[1]x2x4x2xf32>
+func.func @no_fold_bcast_scalable_vs_fixed(%arg0: vector<2x1x1xf32>) -> vector<[1]x2x4x2xf32> {
+ %0 = vector.broadcast %arg0 : vector<2x1x1xf32> to vector<2x[4]x2xf32>
+ %1 = vector.shape_cast %0 : vector<2x[4]x2xf32> to vector<[1]x2x4x2xf32>
+ return %1 : vector<[1]x2x4x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_bcast_scalable_axis_shift
+// CHECK: vector.broadcast %{{.*}} : vector<1x[4]x1xf32> to vector<1x[4]x4xf32>
+// CHECK-NEXT: vector.shape_cast %{{.*}} : vector<1x[4]x4xf32> to vector<4x[4]x1xf32>
+func.func @no_fold_bcast_scalable_axis_shift(%arg0: vector<1x[4]x1xf32>) -> vector<4x[4]x1xf32> {
+ %0 = vector.broadcast %arg0 : vector<1x[4]x1xf32> to vector<1x[4]x4xf32>
+ %1 = vector.shape_cast %0 : vector<1x[4]x4xf32> to vector<4x[4]x1xf32>
+ return %1 : vector<4x[4]x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_leading_scalable_dims
+// CHECK: %[[RES:.*]] = vector.broadcast %{{.*}} : vector<[4]xf32> to vector<8x[4]xf32>
+// CHECK-NEXT: return %[[RES]] : vector<8x[4]xf32>
+func.func @fold_bcast_leading_scalable_dims(%arg0: vector<[4]xf32>) -> vector<8x[4]xf32> {
+ %0 = vector.broadcast %arg0 : vector<[4]xf32> to vector<2x4x[4]xf32>
+ %1 = vector.shape_cast %0 : vector<2x4x[4]xf32> to vector<8x[4]xf32>
+ return %1 : vector<8x[4]xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_leading_scalable_dim_fixed_src
+// CHECK: %[[RES:.*]] = vector.broadcast %{{.*}} : vector<4xf32> to vector<[8]x4xf32>
+// CHECK-NEXT: return %[[RES]] : vector<[8]x4xf32>
+func.func @fold_bcast_leading_scalable_dim_fixed_src(%arg0: vector<4xf32>) -> vector<[8]x4xf32> {
+ %0 = vector.broadcast %arg0 : vector<4xf32> to vector<2x[4]x4xf32>
+ %1 = vector.shape_cast %0 : vector<2x[4]x4xf32> to vector<[8]x4xf32>
+ return %1 : vector<[8]x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_scalable_consecutive_unit_dims
+// CHECK: %[[RES:.*]] = vector.broadcast %{{.*}} : vector<[2]x1x1x3xf32> to vector<[2]x2x2x3xf32>
+// CHECK-NEXT: return %[[RES]] : vector<[2]x2x2x3xf32>
+func.func @fold_bcast_scalable_consecutive_unit_dims(%arg0: vector<[2]x1x1x3xf32>) -> vector<[2]x2x2x3xf32> {
+ %0 = vector.broadcast %arg0 : vector<[2]x1x1x3xf32> to vector<[2]x1x4x3xf32>
+ %1 = vector.shape_cast %0 : vector<[2]x1x4x3xf32> to vector<[2]x2x2x3xf32>
+ return %1 : vector<[2]x2x2x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_bcast_scalable_unit_dim_axis_shift
+// CHECK: vector.broadcast %{{.*}} : vector<[1]x1xf32> to vector<[1]x4xf32>
+// CHECK-NEXT: vector.shape_cast %{{.*}} : vector<[1]x4xf32> to vector<4x[1]xf32>
+func.func @no_fold_bcast_scalable_unit_dim_axis_shift(%arg0: vector<[1]x1xf32>) -> vector<4x[1]xf32> {
+ %0 = vector.broadcast %arg0 : vector<[1]x1xf32> to vector<[1]x4xf32>
+ %1 = vector.shape_cast %0 : vector<[1]x4xf32> to vector<4x[1]xf32>
+ return %1 : vector<4x[1]xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_scalar
+// CHECK: %[[RES:.*]] = vector.broadcast %{{.*}} : f32 to vector<2x4xf32>
+// CHECK-NEXT: return %[[RES]] : vector<2x4xf32>
+func.func @fold_bcast_scalar(%arg0: f32) -> vector<2x4xf32> {
+ %0 = vector.broadcast %arg0 : f32 to vector<8xf32>
+ %1 = vector.shape_cast %0 : vector<8xf32> to vector<2x4xf32>
+ return %1 : vector<2x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_scalar_scalable
+// CHECK: %[[RES:.*]] = vector.broadcast %{{.*}} : f32 to vector<2x[4]xf32>
+// CHECK-NEXT: return %[[RES]] : vector<2x[4]xf32>
+func.func @fold_bcast_scalar_scalable(%arg0: f32) -> vector<2x[4]xf32> {
+ %0 = vector.broadcast %arg0 : f32 to vector<[8]xf32>
+ %1 = vector.shape_cast %0 : vector<[8]xf32> to vector<2x[4]xf32>
+ return %1 : vector<2x[4]xf32>
+}
>From 81ad696f1e4a302ffc1b5ca4d4f24dcf12687add Mon Sep 17 00:00:00 2001
From: yueshengys <135203619+yueshengys at users.noreply.github.com>
Date: Thu, 1 Oct 2026 21:26:24 +0000
Subject: [PATCH 3/4] [mlir][vector] Use computeStrides and llvm::reverse in
ShapeCastBroadcastFolder
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 61 +++++++++----------
.../canonicalize/vector-to-shape-cast.mlir | 7 +++
2 files changed, 37 insertions(+), 31 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 963faa50e5d230..fe3d07d79dadbe 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -7157,21 +7157,15 @@ class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
if (!broadcastOp)
return failure();
- auto srcVectorType = dyn_cast<VectorType>(broadcastOp.getSourceType());
- bool srcIsScalar = !srcVectorType;
-
// Replace Y = ShapeCast(Broadcast(X)) with Y = Broadcast(X)
// Example
// %0 = vector.broadcast %in : vector<3xf32> to vector<2x4x3xf32>
// %1 = vector.shape_cast %0 : vector<2x4x3xf32> to vector<8x3xf32>
// to
// %1 = vector.broadcast %in : vector<3xf32> to vector<8x3xf32>
+ auto srcVectorType = dyn_cast<VectorType>(broadcastOp.getSourceType());
+ bool srcIsScalar = !srcVectorType;
VectorType dstVectorType = shapeCastOp.getResultVectorType();
- ArrayRef<int64_t> dstShape = dstVectorType.getShape();
- ArrayRef<int64_t> srcShape =
- srcIsScalar ? ArrayRef<int64_t>{} : srcVectorType.getShape();
- ArrayRef<int64_t> broadcastShape =
- broadcastOp.getResultVectorType().getShape();
if (!srcIsScalar) {
if (isBroadcastableTo(srcVectorType, dstVectorType) !=
@@ -7184,39 +7178,44 @@ class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
// every non-unit dimension of `src` has the same linear stride in both
// `broadcastShape` and `dstShape`.
//
- // In physical memory, a scalable dimension of static size D has runtime
- // size D * vscale. A linear stride through trailing dimensions is thus:
+ // A scalable dimension of static size D has runtime size D * vscale.
+ // A linear stride through trailing dimensions is thus:
// stride = staticStride * (vscale ^ numScalableDims)
// Since vscale is an unknown runtime quantity (vscale >= 1), both
// staticStride and numScalableDims must match to guarantee stride
// equivalence across all targets.
- struct Stride {
- int64_t staticStride = 1;
- int numScalableDims = 0;
-
- bool operator==(const Stride &other) const {
- return staticStride == other.staticStride &&
- numScalableDims == other.numScalableDims;
- }
- bool operator!=(const Stride &other) const { return !(*this == other); }
- };
-
- Stride bcastStride, dstStride;
VectorType broadcastVectorType = broadcastOp.getResultVectorType();
+ SmallVector<int64_t> bcastStrides =
+ computeStrides(broadcastVectorType.getShape());
+ SmallVector<int64_t> dstStrides =
+ computeStrides(dstVectorType.getShape());
ArrayRef<bool> bcastScalable = broadcastVectorType.getScalableDims();
ArrayRef<bool> dstScalable = dstVectorType.getScalableDims();
ArrayRef<bool> srcScalable = srcVectorType.getScalableDims();
- for (size_t i = 1, e = srcShape.size(); i <= e; ++i) {
- bool isUnitDim = srcShape[e - i] == 1 && !srcScalable[e - i];
- if (!isUnitDim && bcastStride != dstStride)
+ // Iterate from the innermost (trailing) dimension outward. `llvm::zip`
+ // stops once all `src` dimensions are visited, aligning each `src`
+ // dimension with the corresponding trailing dimension of `broadcast` and
+ // `dst`.
+ for (auto [idx, dims] : llvm::enumerate(llvm::zip(
+ llvm::reverse(srcVectorType.getShape()),
+ llvm::reverse(srcScalable), llvm::reverse(bcastStrides),
+ llvm::reverse(dstStrides)))) {
+ auto [srcDim, isSrcScalable, bcastStride, dstStride] = dims;
+ // Unit dimensions (fixed size 1; `[1]` has runtime size `vscale` and is
+ // not unit) only ever have index 0, so their stride does not contribute
+ // to the linear element offset.
+ bool isUnitDim = srcDim == 1 && !isSrcScalable;
+ if (isUnitDim)
+ continue;
+ // For the dimension at reverse index `idx`, its stride is the product
+ // of the `idx` more minor (trailing) dimensions. Verify that both the
+ // static stride factor and the `vscale` exponent (the number of
+ // scalable dimensions in the trailing `idx` dimensions) match.
+ if (bcastStride != dstStride ||
+ llvm::count(bcastScalable.take_back(idx), true) !=
+ llvm::count(dstScalable.take_back(idx), true))
return failure();
- bcastStride.staticStride *= broadcastShape[broadcastShape.size() - i];
- if (bcastScalable[broadcastShape.size() - i])
- bcastStride.numScalableDims++;
- dstStride.staticStride *= dstShape[dstShape.size() - i];
- if (dstScalable[dstShape.size() - i])
- dstStride.numScalableDims++;
}
}
diff --git a/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir b/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
index bffcc012cb5b0f..b625a8d0be8b78 100644
--- a/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
@@ -130,6 +130,13 @@ func.func @extract_from_broadcast(%src: vector<1x1x1xf32>) -> vector<1xf32> {
// -----
+// https://github.com/llvm/llvm-project/issues/190614 — ShapeCastBroadcastFolder
+// must not rewrite shape_cast(broadcast) into a lower-rank broadcast when that
+// changes the linear stride of a non-unit source dimension (duplication across
+// a new leading dim vs stretching an inner unit dim).
+// For %arg0 = [[1], [2]]: broadcast to 2x2x1 then shape_cast to 2x2 yields
+// [[1, 2], [1, 2]]; folding to broadcast 2x1 -> 2x2 would incorrectly yield
+// [[1, 1], [2, 2]].
// CHECK-LABEL: @no_fold_bcast_mode_switch
// CHECK: vector.broadcast %{{.*}} : vector<2x1xf32> to vector<2x2x1xf32>
// CHECK-NEXT: vector.shape_cast %{{.*}} : vector<2x2x1xf32> to vector<2x2xf32>
>From aba2af0fbda613a0fe7bebac17b41445be71a87e Mon Sep 17 00:00:00 2001
From: yueshengys <135203619+yueshengys at users.noreply.github.com>
Date: Wed, 14 Oct 2026 18:47:44 +0000
Subject: [PATCH 4/4] [mlir][vector] Address review comments: use running
stride products
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 62 +++++++++----------
.../canonicalize/vector-to-shape-cast.mlir | 11 ++++
2 files changed, 39 insertions(+), 34 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index fe3d07d79dadbe..895459de388cc3 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -7172,50 +7172,44 @@ class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
BroadcastableToResult::Success) {
return failure();
}
- // `shape_cast` preserves the flattened 1-D (row-major) element order.
- // Thus, `shape_cast(broadcast(src, broadcastShape), dstShape)` produces
- // the exact same elements as `broadcast(src, dstShape)` if and only if
- // every non-unit dimension of `src` has the same linear stride in both
- // `broadcastShape` and `dstShape`.
- //
- // A scalable dimension of static size D has runtime size D * vscale.
- // A linear stride through trailing dimensions is thus:
- // stride = staticStride * (vscale ^ numScalableDims)
- // Since vscale is an unknown runtime quantity (vscale >= 1), both
- // staticStride and numScalableDims must match to guarantee stride
- // equivalence across all targets.
+ // `shape_cast` preserves the flattened 1-D (row-major) element order, so
+ // folding `shape_cast(broadcast(src, broadcastShape), dstShape)` into
+ // `broadcast(src, dstShape)` is valid when every non-unit dimension of
+ // `src` has the same linear stride in both `broadcastShape` and
+ // `dstShape`. For scalable vectors, a dimension's linear stride is
+ // `staticStride * (vscale ^ scalableDims)`, so we check both the static
+ // stride product and the count of trailing scalable dimensions.
VectorType broadcastVectorType = broadcastOp.getResultVectorType();
- SmallVector<int64_t> bcastStrides =
- computeStrides(broadcastVectorType.getShape());
- SmallVector<int64_t> dstStrides =
- computeStrides(dstVectorType.getShape());
- ArrayRef<bool> bcastScalable = broadcastVectorType.getScalableDims();
- ArrayRef<bool> dstScalable = dstVectorType.getScalableDims();
- ArrayRef<bool> srcScalable = srcVectorType.getScalableDims();
// Iterate from the innermost (trailing) dimension outward. `llvm::zip`
// stops once all `src` dimensions are visited, aligning each `src`
// dimension with the corresponding trailing dimension of `broadcast` and
// `dst`.
- for (auto [idx, dims] : llvm::enumerate(llvm::zip(
- llvm::reverse(srcVectorType.getShape()),
- llvm::reverse(srcScalable), llvm::reverse(bcastStrides),
- llvm::reverse(dstStrides)))) {
- auto [srcDim, isSrcScalable, bcastStride, dstStride] = dims;
+ int64_t bcastStride = 1, dstStride = 1;
+ int64_t bcastScalableDims = 0, dstScalableDims = 0;
+ for (auto [srcDim, isSrcScalable, bcastDim, isBcastScalable, dstDim,
+ isDstScalable] :
+ llvm::zip(llvm::reverse(srcVectorType.getShape()),
+ llvm::reverse(srcVectorType.getScalableDims()),
+ llvm::reverse(broadcastVectorType.getShape()),
+ llvm::reverse(broadcastVectorType.getScalableDims()),
+ llvm::reverse(dstVectorType.getShape()),
+ llvm::reverse(dstVectorType.getScalableDims()))) {
// Unit dimensions (fixed size 1; `[1]` has runtime size `vscale` and is
// not unit) only ever have index 0, so their stride does not contribute
- // to the linear element offset.
+ // to the linear element offset. For non-unit dimensions, verify that
+ // both the static stride factor and the `vscale` exponent accumulated
+ // from the more minor (trailing) dimensions match.
bool isUnitDim = srcDim == 1 && !isSrcScalable;
- if (isUnitDim)
- continue;
- // For the dimension at reverse index `idx`, its stride is the product
- // of the `idx` more minor (trailing) dimensions. Verify that both the
- // static stride factor and the `vscale` exponent (the number of
- // scalable dimensions in the trailing `idx` dimensions) match.
- if (bcastStride != dstStride ||
- llvm::count(bcastScalable.take_back(idx), true) !=
- llvm::count(dstScalable.take_back(idx), true))
+ if (!isUnitDim && (bcastStride != dstStride ||
+ bcastScalableDims != dstScalableDims)) {
return failure();
+ }
+
+ bcastStride *= bcastDim;
+ dstStride *= dstDim;
+ bcastScalableDims += isBcastScalable;
+ dstScalableDims += isDstScalable;
}
}
diff --git a/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir b/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
index b625a8d0be8b78..eacfa7f701d9bd 100644
--- a/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize/vector-to-shape-cast.mlir
@@ -277,3 +277,14 @@ func.func @fold_bcast_scalar_scalable(%arg0: f32) -> vector<2x[4]xf32> {
%1 = vector.shape_cast %0 : vector<[8]xf32> to vector<2x[4]xf32>
return %1 : vector<2x[4]xf32>
}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_0d_vector
+// CHECK: %[[RES:.*]] = vector.broadcast %{{.*}} : vector<f32> to vector<2x4xf32>
+// CHECK-NEXT: return %[[RES]] : vector<2x4xf32>
+func.func @fold_bcast_0d_vector(%arg0: vector<f32>) -> vector<2x4xf32> {
+ %0 = vector.broadcast %arg0 : vector<f32> to vector<8xf32>
+ %1 = vector.shape_cast %0 : vector<8xf32> to vector<2x4xf32>
+ return %1 : vector<2x4xf32>
+}
More information about the Mlir-commits
mailing list