diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -1282,7 +1282,7 @@ static LogicalResult vectorizeDynamicLinalgOpPrecondition(linalg::LinalgOp op) { // TODO: Masking only supports dynamic generic ops for now. - if (!isa(op)) + if (!isa(op)) return failure(); LDBG("Dynamically-shaped op meets vectorization pre-conditions\n"); diff --git a/mlir/test/Dialect/Linalg/vectorization.mlir b/mlir/test/Dialect/Linalg/vectorization.mlir --- a/mlir/test/Dialect/Linalg/vectorization.mlir +++ b/mlir/test/Dialect/Linalg/vectorization.mlir @@ -2737,3 +2737,23 @@ 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, %B : memref) { + // CHECK: %[[c0:.*]] = arith.constant 0 : index + // CHECK: %[[d0:.*]] = memref.dim %{{.*}}, %[[c0]] : memref + // CHECK: %[[c1:.*]] = arith.constant 1 : index + // CHECK: %[[d1:.*]] = memref.dim %{{.*}}, %[[c1]] : memref + // CHECK: %[[mask:.*]] = vector.create_mask %[[d0]], %[[d1]] : vector<2x4xi1> + // CHECK: vector.mask %[[mask]] {{.*}} vector.transfer_read %{{.*}} {in_bounds = [true, true]} : memref, vector<2x4xf32> } : vector<2x4xi1> -> vector<2x4xf32> + // CHECK: vector.mask %[[mask]] {{.*}} vector.transfer_write %{{.*}} {in_bounds = [true, true]} : vector<2x4xf32>, memref } : vector<2x4xi1> + linalg.copy ins(%A : memref) outs(%B : memref) + 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] +}