[mlir][Linalg] Allow linalg.copy to be vectorized with masking

[nicolasvasilache]

Add test

[nicolasvasilache]

Rebase

[qcolombet]

mlir/test/Dialect/Linalg/vectorization.mlir

2743

Could you check for the vector.create_mask?

[nicolasvasilache]

Check mask

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

if (!reduceOp || !getCombinerOpKind(reduceOp)) {
       return failure();
     }
   }
   return success();
 }
 
 static LogicalResult vectorizeDynamicLinalgOpPrecondition(linalg::LinalgOp op) {
   // TODO: Masking only supports dynamic generic ops for now.
-  if (!isa<linalg::GenericOp, linalg::FillOp>(op))
+  if (!isa<linalg::GenericOp, linalg::FillOp, linalg::CopyOp>(op))
     return failure();
 
   LDBG("Dynamically-shaped op meets vectorization pre-conditions\n");
   return success();
 }
 
 LogicalResult
 mlir::linalg::vectorizeLinalgOpPrecondition(LinalgOp linalgOp,

mlir/test/Dialect/Linalg/vectorization.mlir

 //   CHECK: vector.mask %[[MASK]] { vector.transfer_write %[[BCAST]], {{.*}} {in_bounds = [true, true]} : vector<8x16xf32>, tensor<?x?xf32> } : vector<8x16xi1>
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   transform.structured.masked_vectorize %0 vector_sizes [8, 16]
 }
 
+// -----
+
+// CHECK-LABEL: func @test_masked_vectorize_linalg_copy
+func.func @test_masked_vectorize_linalg_copy(%A : memref<?x?xf32>, %B : memref<?x?xf32>) {
+  //       CHECK: vector.mask {{.*}} vector.transfer_read %{{.*}} {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x4xf32> } : vector<2x4xi1> -> vector<2x4xf32>
+  //       CHECK: vector.mask {{.*}} vector.transfer_write %{{.*}} {in_bounds = [true, true]} : vector<2x4xf32>, memref<?x?xf32> } : vector<2x4xi1>
+  linalg.copy ins(%A : memref<?x?xf32>) outs(%B : memref<?x?xf32>)
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  transform.structured.masked_vectorize %0 vector_sizes [2, 4]
+}