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

[Federico Bruzzone]

================
@@ -888,3 +888,16 @@ 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<4x8xi16>) -> vector<1x4x8xi16> {
+  %0 = vector.broadcast %arg0 : vector<4x8xi16> to vector<1x4x8xi16>
+  return %0 : vector<1x4x8xi16>
+}
+
+// CHECK-LABEL: func @vector_broadcast_rank_mismatch
+// CHECK-SAME: (%[[SRC:.*]]: vector<4x8xi16>) -> vector<1x4x8xi16>
+// CHECK: %[[SLICE:.*]] = vector.extract_strided_slice %[[SRC]] offsets = [0, 0], sizes = [2, 2], strides = [1, 1] : vector<4x8xi16> to vector<2x2xi16>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[SLICE]] : vector<2x2xi16> to vector<1x2x2xi16>
+// CHECK: vector.insert_strided_slice %[[BCAST]], {{.*}} offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x2x2xi16> into vector<1x4x8xi16>
----------------
FedericoBruzzone wrote:

The CHECK lines only cover the first tile. Please add checks for the remaining tiles and a `CHECK-NOT: vector.broadcast` or something :D

https://github.com/llvm/llvm-project/pull/227351