[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 99f2691838b96c91826394a9763b4f0444a47c28 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..e39cff66c5494 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -423,6 +423,48 @@ func.func @vector_broadcast_with_tailing_unit_dim(%v: vector<4x1xf32>) -> vector
 
 // -----
 
+/// 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_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_unit_dim(%arg: vector<1x4xi16>) -> vector<1x4x4xi16> {
+  %0 = vector.broadcast %arg : vector<1x4xi16> to vector<1x4x4xi16>
+  return %0 : vector<1x4x4xi16>
+}
+
+// -----
+
 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>