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,7 +2442,8 @@ return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag)); } - if (failed(linalg::vectorize(rewriter, linalgOp, vectorSizes))) { + if (failed(linalg::vectorize(rewriter, linalgOp, vectorSizes, + getVectorizeNdExtract()))) { Diagnostic diag(target->getLoc(), DiagnosticSeverity::Error); diag << "failed to vectorize op"; return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag)); 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,38 @@ // ----- +// Masked vectorisation of `tensor.extract`: +// * requires the `{ vectorize_nd_extract }` attribute, +// * has not been implemented yet (hence the attribute is absent). +// 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. +// CHECK: transform.sequence + +transform.sequence failures(suppress) { + ^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.