[mlir][transform] Extend `MaskedVectorizeOp` to work for regular vectorization too

[awarzynski]

Summary

This patch extends MaskedVectorizeOp so that it can be used for "regular" (as opposed to "masked") vectorization as well. While we can already use VectorizeOp for "regular" vectorization, that Op will also apply various patterns on top of vectorization. That means that at the moment, when testing the vectorizer with VectorizeOp, we are effectively testing "vectorization + patterns", i.e. 2 things at a time.

With these updates, you can trigger "regular" vectorization with MaskedVectorizeOp by simply skipping the vector sizes:

transform.structured.masked_vectorize %target : !transform.any_op

Following this change we should probably also rename this Op.

[awarzynski]

@nicolasvasilache , mostly following-up on https://reviews.llvm.org/D153687#inline-1486109. Happy to refine this further if it's still the direction that you'd like for us to take.

[christopherbate]

I also have related https://reviews.llvm.org/D150006 to handle the unroll. Together they allow vectorize + unroll in one go, which is a massive improvement over the current state of things where vector unroll must be applied with a call back post-vectorization. In particular, if we have a Transform op to extract LinalgOp static iteration space sizes then that can be used to inform unroll shape in a natural manner.

[dcaballe]

Nice! Would it make sense for this op to optionally take vector_sizes? That would unify unmasked and masked vectorization into one op. There's no reason to keep them separate.

[awarzynski]

In D157774#4611631, @dcaballe wrote:

Nice! Would it make sense for this op to optionally take vector_sizes? That would unify unmasked and masked vectorization into one op. There's no reason to keep them separate.

Yup, we can do that. That basically means merging VectorizeOneOp with MaskedVectorizeOp. Let me do that first. If we are happy with the result, we can rename MaskedVectorizeOp as something more generic - perhaps VectorizeOneOp :)

[awarzynski]

Merge VectorizeOneOp into MaskedVectorizeOp

[dcaballe]

LGTM and the change is minor but probably good to give some time to someone more familiar with TD.

[awarzynski]

In D157774#4620404, @dcaballe wrote:

LGTM and the change is minor but probably good to give some time to someone more familiar with TD.

Thanks! I've just updated the summary to match the current version. Will wait for more reviews before landing this.

[nicolasvasilache]

In principle this looks good to me, thanks for this!

Can we update masked_vectorize to vectorize_one_op and be consistent with the doc or will this happen in a followup ?

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

2086

typo "vectorization" (according to the other spellings)

2092

can we call this attribute masked_vector_sizes to clarify the intent?

2094

The example does not match the op.

Also, can we drop the vector_sizes spec when inferred ?

2129–2130

Can we make vectorize_nd_extract part of the oilist ?

Free dictionary attributes are always tricky to debug when misspelled.

[nicolasvasilache]

In D157774#4610706, @christopherbate wrote:

I also have related https://reviews.llvm.org/D150006 to handle the unroll. Together they allow vectorize + unroll in one go, which is a massive improvement over the current state of things where vector unroll must be applied with a call back post-vectorization. In particular, if we have a Transform op to extract LinalgOp static iteration space sizes then that can be used to inform unroll shape in a natural manner.

@christopherbate I apologize, I completely missed this .. looking

[awarzynski]

In D157774#4624442, @nicolasvasilache wrote:

Can we update masked_vectorize to vectorize_one_op and be consistent with the doc or will this happen in a followup ?

Let me do that in a follow-up patch. I will also address all the other comments before landing this.

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

2129–2130

See: https://reviews.llvm.org/D159360

[ingomueller-net]

As discussed with @nicolasvasilache, I have started to rename the ops here. That PR is still work in progress but feedback on the new names is already welcome. (Current suggestion: Vectorize => VectorizeChildren and MaskedVectorize => Vectorize.)

[awarzynski]

In D157774#4646894, @ingomueller-net wrote:

As discussed with @nicolasvasilache, I have started to rename the ops here. That PR is still work in progress but feedback on the new names is already welcome. (Current suggestion: Vectorize => VectorizeChildren and MaskedVectorize => Vectorize.

Apologies to Nicolas for the delay with the renaming - it was the next item on my bucket list :(

@ingomueller-net , thanks for taking this on! MaskedVectorize -> Vectorize makes sense to me. Not sure about VectorizeChildren - that op does quite a bit more than plain vectorization.

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

::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::transform::TransformState &state);
   }];
 }
 
 def MaskedVectorizeOp : Op<Transform_Dialect, "structured.masked_vectorize",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      TransformOpInterface, ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
-    Vectorize the target ops, which must be Linalg ops, with masked vectors
-    of the specified size.
+    Vectorize the target ops, which must be Linalg ops. 
+
+    Use the optional vector sizes to specify exactly what configuration the
+    vectorizer should use. It will then use masked vectors of the specified
+    size to enforce this configuration ("masked vectorization"). If no vector
+    sizes are specified, the vectorizer will infer the shapes to use from the
+    target Linalg ops ("regular vectorization"). More specifically:
+
+    ```mlir
+    # Masked vectorization - vector sizes are specified explicitly
+    transform.structured.masked_vectorize %target vector_sizes [1, 4] : !transform.any_op
+    # Regular vectorization - vector sizes are inferred from the target Op
+    transform.structured.masked_vectorize %target : !transform.any_op
+    ```
 
     The vector sizes can be either static or dynamic (SSA values). In case of
     SSA values, the handle must be mapped to exactly one payload op with
     exactly one index-typed result.
 
     Note: The input vector sizes must be bigger than or equal to their
     counterpart iteration space sizes.
 
     Typically this operator should be applied to linalg operations that have
-    already be tiled to the appropriate sizes.
+    already been tiled to the appropriate sizes.
 
     #### Return modes:
 
-    This operation produces a definite failure if the dynamic vector sizes (SSA
-    values) do not satisfy the constraints mentioned above. It produces a
-    silenceable failure if at least one target op is not a Linalg op or fails to
-    vectorize.
+    This operation produces a silenceable failure if at least one target op is
+    not a Linalg op or fails to vectorize. It produces a definite failure if
+    the dynamic vector sizes (SSA values) do not satisfy the constraints
+    mentioned above. 
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
                        Variadic<TransformHandleTypeInterface>:$vector_sizes,
                        UnitAttr:$vectorize_nd_extract,
                        DefaultValuedOptionalAttr<DenseBoolArrayAttr, "{}">:
                           $scalable_sizes,
                        DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:
                           $static_vector_sizes);
 
   let results = (outs);
   let assemblyFormat = [{
-      $target
+      $target oilist(
       `vector_sizes` custom<DynamicIndexList>($vector_sizes,
                                               $static_vector_sizes,
                                               type($vector_sizes),
                                               $scalable_sizes)
+      )
       attr-dict
       `:` type($target)
   }];
 
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     // TODO: applyToOne.

mlir/test/Dialect/Linalg/vectorization.mlir

 // RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file | FileCheck %s
 
 // CHECK-LABEL: contraction_dot
 func.func @contraction_dot(%A: memref<1584xf32>, %B: memref<1584xf32>, %C: memref<f32>) {
 
 // CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584xf32>
 // CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [0] : vector<1584xf32> to f32
   linalg.dot ins(%A, %B: memref<1584xf32>, memref<1584xf32>)
             outs(%C: memref<f32>)
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.dot"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  %1 = get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
-  %2 = transform.structured.vectorize %1  { disable_multi_reduction_to_contract_patterns } : (!transform.any_op) -> !transform.any_op
+  transform.structured.masked_vectorize %0  : !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: contraction_matvec
 func.func @contraction_matvec(%A: memref<1584x1584xf32>, %B: memref<1584xf32>, %C: memref<1584xf32>) {
 
 // CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584x1584xf32>

mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
   %2 = transform.structured.vectorize %1 : (!transform.any_op) -> !transform.any_op
 }
 
 // -----
-
-#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 func.func @vectorize_nd_tensor_extract_constant_idx(%arg0: tensor<3x3xf32>, %arg2: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {
   %c0 = arith.constant 1 : index
   %c1 = arith.constant 2 : index
   %2 = linalg.generic {
-    indexing_maps = [#map1],
+    indexing_maps = [#map],
     iterator_types = ["parallel", "parallel", "parallel"]
   } outs(%arg2 : tensor<1x1x3xf32>) {
   ^bb0(%arg4: f32):
     %7 = tensor.extract %arg0[%c0, %c1] : tensor<3x3xf32>
     linalg.yield %7 : f32
   } -> tensor<1x1x3xf32>
   return %2 : tensor<1x1x3xf32>
 }
 
+// CHECK: #[[$MAP:.*]] = affine_map<(d0, d1) -> (0, 0, 0)>
 // CHECK-LABEL:   func.func @vectorize_nd_tensor_extract_constant_idx(
 // CHECK-SAME:      %[[ARG_0:.*]]: tensor<3x3xf32>,
 // CHECK-SAME:      %[[ARG_1:.*]]: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {
-// CHECK:           %[[C0:.*]] = arith.constant 0 : index
-// CHECK:           %[[C1:.*]] = arith.constant 1 : index
-// CHECK:           %[[C2:.*]] = arith.constant 2 : index
-// CHECK:           %[[EXTRACT:.*]] = tensor.extract %[[ARG_0]]{{\[}}%[[C1]], %[[C2]]] : tensor<3x3xf32>
-// CHECK:           %[[BCAST:.*]] = vector.broadcast %[[EXTRACT]] : f32 to vector<1x1x3xf32>
-// CHECK:           %[[VAL_7:.*]] = vector.transfer_write %[[BCAST]], %[[ARG_1]][%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x1x3xf32>, tensor<1x1x3xf32>
-// CHECK:           return %[[VAL_7]] : tensor<1x1x3xf32>
-// CHECK:         }
+// CHECK-DAG:       %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG:       %[[C2:.*]] = arith.constant 2 : index
+// CHECK-DAG:       arith.constant 0.000000e+00 : f32
+// CHECK-DAG:       %[[C0_f32:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[READ:.*]] = vector.transfer_read  %[[ARG_0]][%[[C1]], %[[C2]]], %[[C0_f32]] {in_bounds = [true, true, true], permutation_map = #[[$MAP]]} : tensor<3x3xf32>, vector<1x1x3xf32>
+// CHECK:           %[[C0_4:.*]] = arith.constant 0 : index
+// CHECK:           vector.transfer_write %[[READ]], %[[ARG_1]][%[[C0_4]], %[[C0_4]], %[[C0_4]]]  : vector<1x1x3xf32>, tensor<1x1x3xf32>
 
 transform.sequence failures(propagate) {
  ^bb1(%arg1: !transform.any_op):
    %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-   %1 = get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
-   %2 = transform.structured.vectorize %1 { vectorize_nd_extract } : (!transform.any_op) -> !transform.any_op
+  transform.structured.masked_vectorize %0 { vectorize_nd_extract }  : !transform.any_op
  }
 
 // -----
 
 #map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 func.func @vectorize_nd_tensor_extract_transfer_read_basic(%arg0: tensor<3x3x3xf32>, %arg2: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {
   %1 = linalg.generic {
     indexing_maps = [#map1],