[Mlir-commits] [mlir] [mlir][vector] Handle rank mismatch when unrolling broadcasts (PR #227351)
purnima shrivastava
llvmlistbot at llvm.org
Mon Oct 5 22:57:32 PDT 2026
https://github.com/purnima-nlp updated https://github.com/llvm/llvm-project/pull/227351
>From 53bd98ee363e854d797845176fe3ab664c87450e Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Tue, 29 Sep 2026 21:09:10 +0530
Subject: [PATCH] [mlir][vector] Handle rank mismatch when unrolling broadcasts
---
.../Vector/Transforms/VectorUnroll.cpp | 16 ++++---
.../Dialect/Vector/vector-unroll-options.mlir | 42 +++++++++++++++++++
2 files changed, 53 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index d58aedcc9d023..dfb92c78eabc2 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -808,16 +808,22 @@ struct UnrollBroadcastPattern : public OpRewritePattern<vector::BroadcastOp> {
Location loc = broadcastOp.getLoc();
VectorType srcType = dyn_cast<VectorType>(broadcastOp.getSourceType());
VectorType resType = broadcastOp.getResultVectorType();
+
+ SmallVector<int64_t> originalShape = *broadcastOp.getShapeForUnroll();
+ // Pad the target shape with leading unit dimensions when the broadcast
+ // result has a higher rank than the native shape.
+ SmallVector<int64_t> adjustedTargetShape(originalShape.size(), 1);
+ llvm::copy(*targetShape, adjustedTargetShape.end() - targetShape->size());
+
VectorType targetType =
- resType.cloneWith(*targetShape, resType.getElementType());
+ resType.cloneWith(adjustedTargetShape, resType.getElementType());
Value result = arith::ConstantOp::create(rewriter, loc, resType,
rewriter.getZeroAttr(resType));
- SmallVector<int64_t> originalShape = *broadcastOp.getShapeForUnroll();
SmallVector<int64_t> strides(originalShape.size(), 1);
for (SmallVector<int64_t> offsets :
- StaticTileOffsetRange(originalShape, *targetShape)) {
+ StaticTileOffsetRange(originalShape, adjustedTargetShape)) {
Value newSrc;
if (!srcType) {
// Scalar to vector broadcast.
@@ -826,8 +832,8 @@ struct UnrollBroadcastPattern : public OpRewritePattern<vector::BroadcastOp> {
// Vector to vector broadcast.
int64_t rank = srcType.getRank();
SmallVector<int64_t> srcOffsets(offsets.end() - rank, offsets.end());
- SmallVector<int64_t> srcShape(targetShape->end() - rank,
- targetShape->end());
+ SmallVector<int64_t> srcShape(adjustedTargetShape.end() - rank,
+ adjustedTargetShape.end());
SmallVector<int64_t> srcStrides(strides.end() - rank, strides.end());
// adjust the offset and shape for src if the corresponding dim is 1.
for (int64_t i = 0; i < rank; ++i) {
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index d1f0961bbb735..43caf248371b6 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -888,3 +888,45 @@ func.func @vector_multi_reduction_rank_mismatch(%v : vector<2x2x4xf32>, %acc: ve
// 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>
+
+// -----
+
+/// A rank-3 broadcast result must pad the rank-2 native shape [2, 2]
+/// to [1, 2, 2] and generate both trailing-dimension tiles.
+// CHECK-LABEL: func.func @vector_broadcast_rank_mismatch(
+// CHECK-SAME: %[[ARG:.*]]: vector<2x4xi16>
+// CHECK: %[[EXTRACT_0:.*]] = vector.extract_strided_slice %[[ARG]] offsets = [0, 0], sizes = [2, 2]
+// CHECK: %[[BCAST_0:.*]] = vector.broadcast %[[EXTRACT_0]] : vector<2x2xi16> to vector<1x2x2xi16>
+// CHECK: %[[INSERT_0:.*]] = vector.insert_strided_slice %[[BCAST_0]], %{{.*}} offsets = [0, 0, 0]
+// CHECK: %[[EXTRACT_1:.*]] = vector.extract_strided_slice %[[ARG]] offsets = [0, 2], sizes = [2, 2]
+// CHECK: %[[BCAST_1:.*]] = vector.broadcast %[[EXTRACT_1]] : vector<2x2xi16> to vector<1x2x2xi16>
+// CHECK: %[[RESULT:.*]] = vector.insert_strided_slice %[[BCAST_1]], %[[INSERT_0]] offsets = [0, 0, 2]
+// CHECK: return %[[RESULT]] : vector<1x2x4xi16>
+func.func @vector_broadcast_rank_mismatch(%arg: vector<2x4xi16>) -> vector<1x2x4xi16> {
+ %0 = vector.broadcast %arg : vector<2x4xi16> to vector<1x2x4xi16>
+ return %0 : vector<1x2x4xi16>
+}
+
+// -----
+
+/// A leading unit source dimension must reuse source slices across result tiles
+/// when the native shape [2, 2] is padded to [1, 2, 2].
+// CHECK-LABEL: func.func @vector_broadcast_rank_mismatch_with_leading_unit_dim(
+// CHECK-SAME: %[[ARG:.*]]: vector<1x4xi16>
+// CHECK: %[[EXTRACT_0:.*]] = vector.extract_strided_slice %[[ARG]] offsets = [0, 0], sizes = [1, 2]
+// CHECK: %[[BCAST_0:.*]] = vector.broadcast %[[EXTRACT_0]] : vector<1x2xi16> to vector<1x2x2xi16>
+// CHECK: %[[INSERT_0:.*]] = vector.insert_strided_slice %[[BCAST_0]], %{{.*}} offsets = [0, 0, 0]
+// CHECK: %[[EXTRACT_1:.*]] = vector.extract_strided_slice %[[ARG]] offsets = [0, 2], sizes = [1, 2]
+// CHECK: %[[BCAST_1:.*]] = vector.broadcast %[[EXTRACT_1]] : vector<1x2xi16> to vector<1x2x2xi16>
+// CHECK: %[[INSERT_1:.*]] = vector.insert_strided_slice %[[BCAST_1]], %[[INSERT_0]] offsets = [0, 0, 2]
+// CHECK: %[[EXTRACT_2:.*]] = vector.extract_strided_slice %[[ARG]] offsets = [0, 0], sizes = [1, 2]
+// CHECK: %[[BCAST_2:.*]] = vector.broadcast %[[EXTRACT_2]] : vector<1x2xi16> to vector<1x2x2xi16>
+// CHECK: %[[INSERT_2:.*]] = vector.insert_strided_slice %[[BCAST_2]], %[[INSERT_1]] offsets = [0, 2, 0]
+// CHECK: %[[EXTRACT_3:.*]] = vector.extract_strided_slice %[[ARG]] offsets = [0, 2], sizes = [1, 2]
+// CHECK: %[[BCAST_3:.*]] = vector.broadcast %[[EXTRACT_3]] : vector<1x2xi16> to vector<1x2x2xi16>
+// CHECK: %[[RESULT:.*]] = vector.insert_strided_slice %[[BCAST_3]], %[[INSERT_2]] offsets = [0, 2, 2]
+// CHECK: return %[[RESULT]] : vector<1x4x4xi16>
+func.func @vector_broadcast_rank_mismatch_with_leading_unit_dim(%arg: vector<1x4xi16>) -> vector<1x4x4xi16> {
+ %0 = vector.broadcast %arg : vector<1x4xi16> to vector<1x4x4xi16>
+ return %0 : vector<1x4x4xi16>
+}
More information about the Mlir-commits
mailing list