[Mlir-commits] [mlir] [mlir][vector] Handle rank mismatch when unrolling broadcasts (PR #227351)

[purnima shrivastava]

https://github.com/purnima-nlp updated https://github.com/llvm/llvm-project/pull/227351

>From 363eef5002357bdc739c9a30cc764291d25e1b14 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 | 19 +++++++++++++++++++
 2 files changed, 30 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..0ee71c741f9eb 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -423,6 +423,25 @@ func.func @vector_broadcast_with_tailing_unit_dim(%v: vector<4x1xf32>) -> vector
 
 // -----
 
+/// Verify unrolling when the broadcast result rank is greater than the
+/// native shape rank.
+// CHECK-LABEL: func.func @vector_broadcast_rank_mismatch(
+//  CHECK-SAME: [[arg0:%.+]]: vector<2x4xi16>
+//       CHECK: [[c:%.+]] = arith.constant dense<0> : vector<1x2x4xi16>
+//       CHECK: [[s0:%.+]] = vector.extract_strided_slice [[arg0]] offsets = [0, 0], sizes = [2, 2], strides = [1, 1] : vector<2x4xi16> to vector<2x2xi16>
+//       CHECK: [[b0:%.+]] = vector.broadcast [[s0]] : vector<2x2xi16> to vector<1x2x2xi16>
+//       CHECK: [[r0:%.+]] = vector.insert_strided_slice [[b0]], [[c]] offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x2x2xi16> into vector<1x2x4xi16>
+//       CHECK: [[s1:%.+]] = vector.extract_strided_slice [[arg0]] offsets = [0, 2], sizes = [2, 2], strides = [1, 1] : vector<2x4xi16> to vector<2x2xi16>
+//       CHECK: [[b1:%.+]] = vector.broadcast [[s1]] : vector<2x2xi16> to vector<1x2x2xi16>
+//       CHECK: [[r1:%.+]] = vector.insert_strided_slice [[b1]], [[r0]] offsets = [0, 0, 2], strides = [1, 1, 1] : vector<1x2x2xi16> into vector<1x2x4xi16>
+//       CHECK: return [[r1]] : 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>
+}
+
+// -----
+
 func.func @vector_load_2D(%mem: memref<4x4xf16>) -> vector<4x4xf16> {
   %c0 = arith.constant 0 : index
   %0 = vector.load %mem[%c0, %c0] : memref<4x4xf16>, vector<4x4xf16>