fixup! [mlir] Add apply_patterns.linalg.pad_vectorization TD Op

Address comment from Javed

Author: banach-space

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

@@ -88,14 +88,17 @@ def ApplyPadVectorizationPatternsOp : Op<Transform_Dialect,
     "apply_patterns.linalg.pad_vectorization",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
-    Apply patterns that take tensor.pad and rewrites it as
-    vector.transfer_read/vector.transfer_write Ops.
-
-    These patterns will either fold tensor.pad with an existing
-    vector.transfer_read or vector.transfer_write producer/consumers (requires
-    other surrounding Ops to be already vectorised) or rewrite it, together
-    with tensor.insert_slice consumer, as a vector.transfer_read +
-    vector.transfer_write pair.
+    Apply patterns that vectorize tensor.pad.
+
+    These patterns rewrite tensor.pad Ops using vector.transfer_read and
+    vector.transfer_write operations. This is done either by:
+      1. Folding tensor.pad with an existing vector.transfer_read /
+      vector.transfer_write Op (generated prior to running these patterns). 
+      2. Rewriting it (when matched together with q tensor.insert_slice
+      consumer Op) as a vector.transfer_read + vector.transfer_write pair.
+
+    In both cases, these patterns look at producers and consumers for the
+    matched tensor.pad Op to find opportunities for vectorization.
   }];
 
   let assemblyFormat = "attr-dict";

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

@@ -2304,7 +2304,7 @@ static SmallVector<Value> ofrToIndexValues(RewriterBase &rewriter, Location loc,
   return result;
 }
 
-/// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp/GenerateOp and
+/// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and
 /// InsertSliceOp. For now, only constant padding values are supported.
 /// If there is enough static type information, TransferReadOps and
 /// TransferWriteOps may be generated instead of InsertSliceOps.
@@ -2712,6 +2712,9 @@ struct PadOpVectorizationWithInsertSlicePattern
 
 void mlir::linalg::populatePadOpVectorizationPatterns(
     RewritePatternSet &patterns, PatternBenefit baseBenefit) {
+  // TODO: The following pattern implements "decomposition" and
+  // optional "vectorization". Seperate "decomposition" into a sepereate
+  // pre-processing pattern group.
   patterns.add<GenericPadOpVectorizationPattern>(patterns.getContext(),
                                                  baseBenefit);
   // Try these specialized patterns first before resorting to the generic one.

mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir

@@ -37,18 +37,18 @@ module attributes {transform.with_named_sequence} {
 // -----
 
 ///----------------------------------------------------------------------------------------
-/// [Pattern: PadOpVectorizationWithTransferReadPattern
+/// [Pattern: PadOpVectorizationWithTransferWritePattern]
 ///----------------------------------------------------------------------------------------
 func.func private @make_vector() -> vector<7x9xf32>
 
-// CHECK-LABEL: func @pad_and_transfer_write_static
+// CHECK-LABEL: func @pad_and_transfer_write_static_low_and_high
 //  CHECK-SAME:     %[[ARG0:.*]]: tensor<5x6xf32>
 //   CHECK-NOT:   tensor.pad
 //       CHECK:   %[[C0:.*]] = arith.constant 0 : index
 //       CHECK:   %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32>
 //       CHECK:   %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[ARG0]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor<5x6xf32>
 //       CHECK:   return %[[RESULT]]
-func.func @pad_and_transfer_write_static(
+func.func @pad_and_transfer_write_static_low_and_high(
     %arg0: tensor<5x6xf32>) -> tensor<5x6xf32> {
   %c0 = arith.constant 0 : index
   %c5 = arith.constant 5.0 : f32
@@ -78,15 +78,15 @@ module attributes {transform.with_named_sequence} {
 
 func.func private @make_vector() -> vector<7x9xf32>
 
-// CHECK-LABEL: func @pad_and_transfer_write_dynamic_static
+// CHECK-LABEL: func @pad_and_transfer_write_static_low_dynamic_high
 //  CHECK-SAME:     %[[ARG0:.*]]: tensor<?x?xf32>, %[[SIZE:.*]]: index, %[[PADDING:.*]]: index
 //   CHECK-NOT:   tensor.pad
 //       CHECK:   %[[C0:.*]] = arith.constant 0 : index
 //       CHECK:   %[[SUB:.*]] = tensor.extract_slice %[[ARG0]][0, 0] [%[[SIZE]], 6] [1, 1] : tensor<?x?xf32> to tensor<?x6xf32>
 //       CHECK:   %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32>
 //       CHECK:   %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[SUB]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor<?x6xf32>
 //       CHECK:   return %[[RESULT]]
-func.func @pad_and_transfer_write_dynamic_static(
+func.func @pad_and_transfer_write_static_low_dynamic_high(
     %arg0: tensor<?x?xf32>, %size: index, %padding: index) -> tensor<?x6xf32> {
   %c0 = arith.constant 0 : index
   %c5 = arith.constant 5.0 : f32
@@ -166,7 +166,9 @@ module attributes {transform.with_named_sequence} {
 
 func.func private @make_vector() -> tensor<12x13xf32>
 
-// Same as @pad_and_insert_slice_dest in vectorization-wit-patterns.mlir, but 
+// Same as @pad_and_insert_slice_dest in vectorization-with-patterns.mlir, but
+// over here linalg::fill is not vectorized (patterns for linalg.fill are not
+// included here)
 // CHECK-LABEL:   func.func @pad_and_insert_slice_dest(
 // CHECK-SAME:      %[[ARG_0:.*]]: tensor<1x5x6xf32>) -> tensor<1x12x13xf32> {
 //  CHECK-NOT:     tensor.pad