[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 d2a77bf27938bbfebe3fef20f457a20f79122a03 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 | 50 +++++++++++++++++++
 2 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index d58aedcc9d02342..dfb92c78eabc29b 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 d1f0961bbb735c6..5b10f2308c49391 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -888,3 +888,53 @@ 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>
+
+// -----
+
+func.func @vector_broadcast_rank_mismatch(%arg0: vector<2x4xi16>) -> vector<1x2x4xi16> {
+  %0 = vector.broadcast %arg0 : vector<2x4xi16> to vector<1x2x4xi16>
+  return %0 : vector<1x2x4xi16>
+}
+
+// CHECK-LABEL: func @vector_broadcast_rank_mismatch
+// CHECK-SAME: (%[[SRC:.*]]: vector<2x4xi16>) -> vector<1x2x4xi16>
+// CHECK-NOT: vector.broadcast
+// CHECK-NEXT: %[[INIT:.*]] = arith.constant dense<0> : vector<1x2x4xi16>
+// CHECK-NEXT: %[[SLICE0:.*]] = vector.extract_strided_slice %[[SRC]] offsets = [0, 0], sizes = [2, 2], strides = [1, 1] : vector<2x4xi16> to vector<2x2xi16>
+// CHECK-NEXT: %[[BCAST0:.*]] = vector.broadcast %[[SLICE0]] : vector<2x2xi16> to vector<1x2x2xi16>
+// CHECK-NEXT: %[[RESULT0:.*]] = vector.insert_strided_slice %[[BCAST0]], %[[INIT]] offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x2x2xi16> into vector<1x2x4xi16>
+// CHECK-NEXT: %[[SLICE1:.*]] = vector.extract_strided_slice %[[SRC]] offsets = [0, 2], sizes = [2, 2], strides = [1, 1] : vector<2x4xi16> to vector<2x2xi16>
+// CHECK-NEXT: %[[BCAST1:.*]] = vector.broadcast %[[SLICE1]] : vector<2x2xi16> to vector<1x2x2xi16>
+// CHECK-NEXT: %[[RESULT1:.*]] = vector.insert_strided_slice %[[BCAST1]], %[[RESULT0]] offsets = [0, 0, 2], strides = [1, 1, 1] : vector<1x2x2xi16> into vector<1x2x4xi16>
+// CHECK-NOT: vector.broadcast
+// CHECK-NEXT: return %[[RESULT1]] : vector<1x2x4xi16>
+// CHECK-NOT: vector.broadcast
+// CHECK-NEXT: }
+
+// -----
+
+func.func @vector_broadcast_rank_mismatch_unit_dim(%arg0: vector<1x4xi16>) -> vector<1x4x4xi16> {
+  %0 = vector.broadcast %arg0 : vector<1x4xi16> to vector<1x4x4xi16>
+  return %0 : vector<1x4x4xi16>
+}
+
+// CHECK-LABEL: func @vector_broadcast_rank_mismatch_unit_dim
+// CHECK-SAME: (%[[SRC:.*]]: vector<1x4xi16>) -> vector<1x4x4xi16>
+// CHECK-NOT: vector.broadcast
+// CHECK-NEXT: %[[INIT:.*]] = arith.constant dense<0> : vector<1x4x4xi16>
+// CHECK-NEXT: %[[SLICE0:.*]] = vector.extract_strided_slice %[[SRC]] offsets = [0, 0], sizes = [1, 2], strides = [1, 1] : vector<1x4xi16> to vector<1x2xi16>
+// CHECK-NEXT: %[[BCAST0:.*]] = vector.broadcast %[[SLICE0]] : vector<1x2xi16> to vector<1x2x2xi16>
+// CHECK-NEXT: %[[RESULT0:.*]] = vector.insert_strided_slice %[[BCAST0]], %[[INIT]] offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x2x2xi16> into vector<1x4x4xi16>
+// CHECK-NEXT: %[[SLICE1:.*]] = vector.extract_strided_slice %[[SRC]] offsets = [0, 2], sizes = [1, 2], strides = [1, 1] : vector<1x4xi16> to vector<1x2xi16>
+// CHECK-NEXT: %[[BCAST1:.*]] = vector.broadcast %[[SLICE1]] : vector<1x2xi16> to vector<1x2x2xi16>
+// CHECK-NEXT: %[[RESULT1:.*]] = vector.insert_strided_slice %[[BCAST1]], %[[RESULT0]] offsets = [0, 0, 2], strides = [1, 1, 1] : vector<1x2x2xi16> into vector<1x4x4xi16>
+// CHECK-NEXT: %[[SLICE2:.*]] = vector.extract_strided_slice %[[SRC]] offsets = [0, 0], sizes = [1, 2], strides = [1, 1] : vector<1x4xi16> to vector<1x2xi16>
+// CHECK-NEXT: %[[BCAST2:.*]] = vector.broadcast %[[SLICE2]] : vector<1x2xi16> to vector<1x2x2xi16>
+// CHECK-NEXT: %[[RESULT2:.*]] = vector.insert_strided_slice %[[BCAST2]], %[[RESULT1]] offsets = [0, 2, 0], strides = [1, 1, 1] : vector<1x2x2xi16> into vector<1x4x4xi16>
+// CHECK-NEXT: %[[SLICE3:.*]] = vector.extract_strided_slice %[[SRC]] offsets = [0, 2], sizes = [1, 2], strides = [1, 1] : vector<1x4xi16> to vector<1x2xi16>
+// CHECK-NEXT: %[[BCAST3:.*]] = vector.broadcast %[[SLICE3]] : vector<1x2xi16> to vector<1x2x2xi16>
+// CHECK-NEXT: %[[RESULT3:.*]] = vector.insert_strided_slice %[[BCAST3]], %[[RESULT2]] offsets = [0, 2, 2], strides = [1, 1, 1] : vector<1x2x2xi16> into vector<1x4x4xi16>
+// CHECK-NOT: vector.broadcast
+// CHECK-NEXT: return %[[RESULT3]] : vector<1x4x4xi16>
+// CHECK-NOT: vector.broadcast
+// CHECK-NEXT: }