[Mlir-commits] [mlir] [mlir][vector] Fix mask rank extension for TransferWriteNonPermutationLowering (PR #144893)

[Andrzej Warzyński]

================
@@ -160,6 +160,31 @@ func.func @xfer_write_non_minor_identity_with_mask_out_of_bounds(
   return
 }
 
+// CHECK-LABEL:   func.func @xfer_write_non_minor_identity_with_mask_broadcast(
+// CHECK-SAME:      %[[MEM:.*]]: memref<?x?x?xf32>,
+// CHECK-SAME:      %[[VEC:.*]]: vector<7x8xf32>,
+// CHECK-SAME:      %[[MASK:.*]]: vector<7x8xi1>,
+// CHECK-SAME:      %[[IDX_1:.*]]: index, %[[IDX_2:.*]]: index, %[[IDX_3:.*]]: index) {
+// CHECK:           %[[BC:.*]] = vector.broadcast %[[VEC]] : vector<7x8xf32> to vector<1x7x8xf32>
+// CHECK:           %[[MBC:.*]] = vector.broadcast %[[MASK]] : vector<7x8xi1> to vector<1x7x8xi1>
+// CHECK:           %[[MTR:.*]] = vector.transpose %[[MBC]], [1, 0, 2] : vector<1x7x8xi1> to vector<7x1x8xi1>
+// CHECK:           %[[TR:.*]] = vector.transpose %[[BC]], [1, 0, 2] : vector<1x7x8xf32> to vector<7x1x8xf32>
+// CHECK:           vector.transfer_write %[[TR]], %[[MEM]]{{\[}}%[[IDX_1]], %[[IDX_2]], %[[IDX_3]]], %[[MTR]] {in_bounds = [false, true, false]} : vector<7x1x8xf32>, memref<?x?x?xf32>
+func.func @xfer_write_non_minor_identity_with_mask_broadcast_and_transpose(
----------------
banach-space wrote:

Based on the name (`@xfer_write_non_minor_identity_with_mask_broadcast_and_transpose`), I would expect to see explicit broadcast when comparing this to `@xfer_write_non_minor_identity_with_mask_out_of_bounds`. However, the only difference I can see are input ranks.

Don't want to come across as nit-picking, just trying to figure out what cases this PR fixes. Thanks!

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