[mlir] Add vectorize_nd_extract attribute to the masked_vectorize Op

This patch simply adds `vectorize_nd_extract` (that's currently only
used for the `vectorize` Op) to `masked_vectorize`. A test is added to
verify that it works as expected - it prevents the masked vectorisation
of `tensor.extract`, which is currently not supported.

Differential Revision: https://reviews.llvm.org/D142634

Author: banach-space

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

@@ -1594,6 +1594,7 @@ def MaskedVectorizeOp : Op<Transform_Dialect, "structured.masked_vectorize",
 
   let arguments = (ins PDL_Operation:$target,
                        Variadic<PDL_Operation>:$vector_sizes,
+                       UnitAttr:$vectorize_nd_extract,
                        DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:
                           $static_vector_sizes);
   let results = (outs);

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

@@ -2986,7 +2986,8 @@ DiagnosedSilenceableFailure transform::MaskedVectorizeOp::apply(
              << "cannot vectorize non-Linalg op";
     }
 
-    if (failed(linalg::vectorize(rewriter, linalgOp, vectorSizes))) {
+    if (failed(linalg::vectorize(rewriter, linalgOp, vectorSizes,
+                                 getVectorizeNdExtract()))) {
       return mlir::emitSilenceableFailure(target->getLoc())
              << "failed to vectorize op";
     }

mlir/test/Dialect/Linalg/vectorization-unsupported.mlir

@@ -0,0 +1,29 @@
+// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics
+
+// 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<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %c0 = arith.constant 1 : index
+  %c1 = arith.constant 2 : index
+  // expected-error@+1 {{failed to vectorize op}}
+  %2 = linalg.generic {
+    indexing_maps = [#map1],
+    iterator_types = ["parallel", "parallel"]
+  } outs(%arg1 : tensor<?x?xf32>) {
+  ^bb0(%arg3: f32):
+    %7 = tensor.extract %arg0[%c0, %c1] : tensor<?x?xf32>
+    linalg.yield %7 : f32
+  } -> tensor<?x?xf32>
+  return %2 : tensor<?x?xf32>
+}
+
+
+transform.sequence failures(propagate) {
+ ^bb1(%arg1: !pdl.operation):
+   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+   transform.structured.masked_vectorize %0 vector_sizes [3, 3]
+ }