[Mlir-commits] [mlir] [MLIR][Vector] Fix transferOps optimization inside maskOp (PR #90835)

[Andrzej Warzyński]

================
@@ -111,3 +111,18 @@ func.func @insert_slice_of_transfer_write_rank_extending(%t1 : tensor<?x?x12xf32
   %1 = tensor.insert_slice %0 into %t1[4, 3, %s] [1, 5, 6] [1, 1, 1] : tensor<5x6xf32> into tensor<?x?x12xf32>
   return %1 : tensor<?x?x12xf32>
 }
+
+// CHECK-LABEL: func @masked_transfer_read_of_extract_slice
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index
+//       CHECK-DAG: %[[m:.*]] = vector.create_mask{{.*}} : vector<5x6xi1>
+//       CHECK-DAG: %[[a:.*]] = affine.apply {{.*}}[[s1]]
+//       CHECK: vector.mask %[[m]] { vector.transfer_read %[[t]]{{.*}}: tensor<?x?xf32>, vector<5x6xf32> } : vector<5x6xi1> -> vector<5x6xf32>
+func.func @masked_transfer_read_of_extract_slice(%t : tensor<?x?xf32>, %s1 : index, %s2 : index) -> vector<5x6xf32> {
+  %c3 = arith.constant 3 : index
+  %c4 = arith.constant 4 : index
+  %cst = arith.constant 0.0 : f32
+  %0 = tensor.extract_slice %t[5, %s1] [10, %s2] [1, 1] : tensor<?x?xf32> to tensor<10x?xf32>
+  %mask = vector.create_mask %c3, %c4 : vector<5x6xi1>
+  %1 = vector.mask %mask {vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true, true]} : tensor<10x?xf32>, vector<5x6xf32>} : vector<5x6xi1> -> vector<5x6xf32>
+  return %1 : vector<5x6xf32>
+}
----------------
banach-space wrote:

nit: Could you move it next to `@transfer_read_of_extract_slice`? I think that those two should be next to another (one complements the other). If tests are not grouped together it's hard to identify the cases that are covered.

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