diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td --- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td +++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td @@ -1438,6 +1438,7 @@ let arguments = (ins PDL_Operation:$target, Variadic:$vector_sizes, + UnitAttr:$vectorize_nd_extract, DefaultValuedOptionalAttr: $static_vector_sizes); let results = (outs); diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp --- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp +++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp @@ -2442,11 +2442,10 @@ return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag)); } - if (failed(linalg::vectorize(rewriter, linalgOp, vectorSizes))) { - Diagnostic diag(target->getLoc(), DiagnosticSeverity::Error); - diag << "failed to vectorize op"; - return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag)); - } + if (failed(linalg::vectorize(rewriter, linalgOp, vectorSizes, + getVectorizeNdExtract()))) { + return DiagnosedSilenceableFailure::success(); + } } return DiagnosedSilenceableFailure::success(); 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 @@ -1936,6 +1936,35 @@ // ----- +// Masked vectorisation requires of `tensor.extract` requires the `{ vectorize_nd_extract }` attribute. +// TOOD: Implement masked vectorization for `tensor.extract` + +#map1 = affine_map<(d0, d1) -> (d0, d1)> +func.func @extract_masked_vectorize(%arg0: tensor, %arg1: tensor) -> tensor { + %c0 = arith.constant 1 : index + %c1 = arith.constant 2 : index + %2 = linalg.generic { + indexing_maps = [#map1], + iterator_types = ["parallel", "parallel"] + } outs(%arg1 : tensor) { + ^bb0(%arg3: f32): + %7 = tensor.extract %arg0[%c0, %c1] : tensor + linalg.yield %7 : f32 + } -> tensor + return %2 : tensor +} + +// CHECK-LABEL: @extract_masked_vectorize +// CHECK-NOT: vector + +transform.sequence failures(propagate) { + ^bb1(%arg1: !pdl.operation): + %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 + transform.structured.masked_vectorize %0 vector_sizes [3, 3] + } + +// ----- + // This is a regression test. This IR cannot be vectorized, but // structured.vectorize should nevertheless succeed.