[mlir] Add vectorize_nd_extract attribute to masked_vectorize

[awarzynski]

Summary

This patch simply adds vectorize_nd_extract (that's currently 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.

[awarzynski]

@dcaballe Is this what you had in mind? The change is rather trivial, but exposes a small design/implementation mismatch between transform.structured.vectorize and transform.structured.masked_vectorize. Note how:

• the former goes via applyToOne, which does not issue any diagnostics,
• the latter implements apply, which does issue any diagnostics.

This is significant in this case, because a diagnostic leads to a non-zero return code from the driver (e.g. mlir-opt), which means that e.g. "vectorization.mlir" fails (non-zero return code means is interpreted as a test failure by LIT). Perhaps my update is against what's intended here, but then we should look for some other way to make sure that we get consistent behavior across Transfer dialect Ops (at least in the context of the Linalg vectorizer).

[nicolasvasilache]

transform.vectorize is a blanket op that applies a bunch of patterns greedily, it is unfortunate but it is what gets us an unsurprising behavior for now, it will be improved over time.
transform.masked_vectorize is targeted to 1 op because it needs the vector_sizes information which comes from the user injecting static information.

I am not clear what this PR aims to achieve.

If you wanted to just turn off the silenceable error that vectorization emits, you can simply wrap the transform into a

transform.sequence failures(suppress) {

instead of a

transform.sequence failures(propagate) {

[awarzynski]

Thanks for taking a look!

In D142634#4086006, @nicolasvasilache wrote:

If you wanted to just turn off the silenceable error that vectorization emits, you can simply wrap the transform into a

transform.sequence failures(suppress) {

instead of a

transform.sequence failures(propagate) {

This is exactly what I was after. Let me removed all other changes that are not needed and update the summary.

[awarzynski]

Replace failures(propagate) with failures(suppress), remove unrelated changes that are not needed.

[ftynse]

If you wanted to just turn off the silenceable error that vectorization emits, you can simply wrap the transform into a

We can also check for emitted diagnostics using -verify-diagnostics if we expect the transform op to fail, silencing failures in tests isn't necessarily the best idea. We could also split out this test into a separate file instead of appending to a 2k line test that becomes unmanageable.

[awarzynski]

Incorporate suggestions from @ftynse

• Moved the test to a newly created file: vectorization-unsupported.mlir
• Added --verify-diagnostics

[awarzynski]

Ping

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

let description = [{
     This operation produces a definite failure if the dynamic vector sizes (SSA
     values) do not satify the constraints mentioned above. It produces a
     silenceable failure if at least one target op is not a Linalg op or fails to
     vectorize.
   }];
 
   let arguments = (ins PDL_Operation:$target,
                        Variadic<PDL_Operation>:$vector_sizes,
+                       UnitAttr:$vectorize_nd_extract,
                        DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:
                           $static_vector_sizes);
   let results = (outs);
   let assemblyFormat = [{
       $target
       `vector_sizes` custom<DynamicIndexList>($vector_sizes,
                                               $static_vector_sizes)
       attr-dict

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

DiagnosedSilenceableFailure transform::MaskedVectorizeOp::apply(
   for (Operation *target : targets) {
     auto linalgOp = dyn_cast<LinalgOp>(target);
     if (!linalgOp) {
       Diagnostic diag(target->getLoc(), DiagnosticSeverity::Error);
       diag << "cannot vectorize non-Linalg op";
       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();
 }
 
 void transform::MaskedVectorizeOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);

mlir/test/Dialect/Linalg/vectorization.mlir

 // CHECK:           %[[VAL_11:.*]] = vector.mask %[[VAL_10]] { vector.transfer_read %[[VAL_0]]{{.*}} {in_bounds = [true, true, true]} : tensor<?x?x?xf32>, vector<4x8x16xf32> } : vector<4x8x16xi1> -> vector<4x8x16xf32>
 // CHECK:           %[[VAL_13:.*]] = vector.create_mask %[[VAL_7]], %[[VAL_5]] : vector<16x8xi1>
 // CHECK:           %[[VAL_14:.*]] = vector.mask %[[VAL_13]] { vector.transfer_read %[[VAL_1]]{{.*}} {in_bounds = [true, true], permutation_map = #{{.*}}} : tensor<?x?xf32>, vector<8x16xf32> } : vector<16x8xi1> -> vector<8x16xf32>
 // CHECK:           %[[VAL_15:.*]] = vector.mask %[[VAL_10]] { vector.multi_reduction <add>, %[[VAL_11]], %[[VAL_14]] [0] : vector<4x8x16xf32> to vector<8x16xf32> } : vector<4x8x16xi1> -> vector<8x16xf32>
 // CHECK:           %[[VAL_17:.*]] = vector.mask %[[VAL_13]] { vector.transfer_write %[[VAL_15]], %{{.*}} {in_bounds = [true, true], permutation_map = #{{.*}}} : vector<8x16xf32>, tensor<?x?xf32> } : vector<16x8xi1> -> tensor<?x?xf32>
 
 // -----
 
+// 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<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %c0 = arith.constant 1 : index
+  %c1 = arith.constant 2 : index
+  %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>
+}
+
+// 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.
 
 #map = affine_map<(d0) -> (d0)>
 // CHECK-LABEL:   @not_vectorizable
 func.func @not_vectorizable(%arg0: tensor<1x?xf32>, %arg1: index, %arg2: index, %arg3: index) -> tensor<1x128xf32> {
   %0 = tensor.empty() : tensor<1x128xf32>
   %1 = scf.for %arg5 = %arg2 to %arg1 step %arg3 iter_args(%arg6 = %0) -> (tensor<1x128xf32>) {