[Mlir-commits] [mlir] [mlir][VectorToSCF] Decline transfer ops masked by vector.mask (PR #221307)

[Hung-Kuan Tseng]

================
@@ -923,3 +923,75 @@ func.func private @transfer_write_no_alloc_scope()
 %cst = arith.constant dense<0.0> : vector<2x3xf32>
 %m = memref.alloc() : memref<2x3xf32>
 vector.transfer_write %cst, %m[%c0, %c0] : vector<2x3xf32>, memref<2x3xf32>
+
+// -----
+
+// Negative tests: a transfer op inside `vector.mask` is left alone, since the
+// mask region can only hold that one op. The mask has to be lowered first.
+
+// CHECK-LABEL: @masked_transfer_read
+// CHECK:       vector.mask %{{.*}} { vector.transfer_read
+// CHECK-NOT:   scf.for
+// FULL-UNROLL-LABEL: @masked_transfer_read
+// FULL-UNROLL:       vector.mask %{{.*}} { vector.transfer_read
+// FULL-UNROLL-NOT:   vector.extract
+// TARGET-RANK-ZERO-LABEL: @masked_transfer_read
+// TARGET-RANK-ZERO:       vector.mask %{{.*}} { vector.transfer_read
+// TARGET-RANK-ZERO-NOT:   vector.extract
+func.func @masked_transfer_read(%mem: memref<?x?xf32>, %mask: vector<4x8xi1>) -> vector<4x8xf32> {
+  %c0 = arith.constant 0 : index
+  %pad = arith.constant 0.0 : f32
+  %0 = vector.mask %mask { vector.transfer_read %mem[%c0, %c0], %pad : memref<?x?xf32>, vector<4x8xf32> } : vector<4x8xi1> -> vector<4x8xf32>
+  return %0 : vector<4x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @masked_transfer_write
+// CHECK:       vector.mask %{{.*}} { vector.transfer_write
+// CHECK-NOT:   scf.for
+// FULL-UNROLL-LABEL: @masked_transfer_write
+// FULL-UNROLL:       vector.mask %{{.*}} { vector.transfer_write
+// FULL-UNROLL-NOT:   vector.extract
+// TARGET-RANK-ZERO-LABEL: @masked_transfer_write
+// TARGET-RANK-ZERO:       vector.mask %{{.*}} { vector.transfer_write
+// TARGET-RANK-ZERO-NOT:   vector.extract
+func.func @masked_transfer_write(%vec: vector<4x8xf32>, %mem: memref<?x?xf32>, %mask: vector<4x8xi1>) {
+  %c0 = arith.constant 0 : index
+  vector.mask %mask { vector.transfer_write %vec, %mem[%c0, %c0] : vector<4x8xf32>, memref<?x?xf32> } : vector<4x8xi1>
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @masked_transfer_read_1d_strided
+// CHECK:       vector.mask %{{.*}} { vector.transfer_read
+// CHECK-NOT:   scf.for
+// FULL-UNROLL-LABEL: @masked_transfer_read_1d_strided
+// FULL-UNROLL:       vector.mask %{{.*}} { vector.transfer_read
+// FULL-UNROLL-NOT:   scf.for
+// TARGET-RANK-ZERO-LABEL: @masked_transfer_read_1d_strided
+// TARGET-RANK-ZERO:       vector.mask %{{.*}} { vector.transfer_read
+// TARGET-RANK-ZERO-NOT:   scf.for
+func.func @masked_transfer_read_1d_strided(%mem: memref<?x?xf32>, %mask: vector<8xi1>) -> vector<8xf32> {
----------------
Tim096 wrote:

It targets a different pattern. The `(d0, d1) -> (d0)` map makes it a strided 1-D read, which goes through `TransferOp1dConversion` under the default and `FULL-UNROLL` RUN lines; `@masked_transfer_read` goes through the progressive and unroll patterns. Without this patch they also fail differently: the 1-D pattern replaces the op with an `scf.for`, so the verifier reports `expects a MaskableOpInterface within the mask region` instead of `expects only one operation to mask`. (`TARGET-RANK-ZERO` unrolls it to 0-D first, so there it takes the unroll path.) I've added a comment saying this above the test.


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