[mlir][linalg] Vectorize linalg.pad_op source copying (static source shape)

[springerm]

Summary

If the source operand of a linalg.pad_op operation has static shape, vectorize the copying of the source.

Depends On D103679

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

fix bug

[springerm]

rebase

[springerm]

no change

[springerm]

rebase

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

static bool hasSameTensorSize(Value beforePadding, SubTensorOp afterTrimming) {
   }
 
   // All tests passed.
   return true;
 }
 
 /// Rewrite a PadTensorOp into a sequence of InitTensorOp, FillOp and
 /// SubTensorInsertOp. For now, only constant padding values are supported.
-/// Note: This rewrite is not yet a vectorization, but some of the generated ops
-/// may be vectorized down the line (e.g., FillOp).
-/// TODO: If there is enough static shape information, generate TransferReadOps
-/// and TransferWriteOps instead of SubTensorInsertOp.
+/// If there is enough static type information, TransferReadOps and
+/// TransferWriteOps may be generated instead of SubTensorInsertOps.
 struct GenericPadTensorOpVectorizationPattern
     : public OpRewritePattern<PadTensorOp> {
   using OpRewritePattern<PadTensorOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(PadTensorOp padOp,
                                 PatternRewriter &rewriter) const final {
     // Given an OpFoldResult, return an index-typed value.
     auto getIdxValue = [&](OpFoldResult ofr) {
LogicalResult matchAndRewrite(PadTensorOp padOp,
     }
 
     Value init = rewriter.create<InitTensorOp>(
         padOp.getLoc(), dynSizes, staticSizes, resultType.getElementType());
     Value fill =
         rewriter.create<FillOp>(padOp.getLoc(), init, padValue).result();
 
     auto sourceType = padOp.getSourceType();
+
+    // Copy of source with static shape can be vectorized.
+    if (sourceType.hasStaticShape()) {
+      auto vecType =
+          VectorType::get(sourceType.getShape(), sourceType.getElementType());
+      vectorizeStaticShapeSource(rewriter, padOp, init, vecType);
+      return success();
+    }
+
+    // TODO: Vectorize dynamic source but static destination.
+
+    // Neither source type nor PadTensorOp result type have static shape. Such
+    // PadTensorOps cannot be vectorized. Generate a SubTensorInsertOp instead.
+
     // Compute size of source of PadTensorOp.
     SmallVector<OpFoldResult> srcSizes;
     for (unsigned dim = 0; dim < sourceType.getRank(); ++dim) {
       if (sourceType.isDynamicDim(dim)) {
         srcSizes.push_back(rewriter.createOrFold<memref::DimOp>(
             padOp.getLoc(), padOp.source(), dim));
       } else {
         srcSizes.push_back(rewriter.getIndexAttr(sourceType.getDimSize(dim)));
       }
     }
     // Strides of SubTensorInsertOp are all 1.
     SmallVector<OpFoldResult> strides(sourceType.getRank(),
                                       rewriter.getIndexAttr(1));
     rewriter.replaceOpWithNewOp<SubTensorInsertOp>(
         padOp, padOp.source(), fill, padOp.getMixedLowPad(), srcSizes, strides);
 
     return success();
   }
+
+  /// Vectorize the copying of a PadTensorOp's source that has static shape.
+  void vectorizeStaticShapeSource(PatternRewriter &rewriter, PadTensorOp padOp,
+                                  Value dest, VectorType vecType) const {
+    // Generate TransferReadOp.
+    SmallVector<Value> readIndices(
+        vecType.getRank(), rewriter.create<ConstantIndexOp>(padOp.getLoc(), 0));
+    auto read = rewriter.create<vector::TransferReadOp>(
+        padOp.getLoc(), vecType, padOp.source(), readIndices);
+
+    // Generate TransferWriteOp. The destination dimensions may be dynamic, but
+    // the write cannot be out-of-bounds. (A large enough destination tensor is
+    // allocated in this pattern.)
+    auto writeIndices =
+        ofrToIndexValues(rewriter, padOp.getLoc(), padOp.getMixedLowPad());
+    SmallVector<bool> inBounds(vecType.getRank(), true);
+    rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
+        padOp, read, dest, writeIndices, inBounds);
+  }
 };
 
 /// Base pattern for rewriting PadTensorOps whose result is consumed by a given
 /// operation type OpTy.
 template <typename OpTy>
 struct VectorizePadTensorOpUserPattern : public OpRewritePattern<PadTensorOp> {
   using OpRewritePattern<PadTensorOp>::OpRewritePattern;
 

mlir/test/Dialect/Linalg/vectorization.mlir

%0 = linalg.pad_tensor %arg0 low[0, 0, 2] high[0, 1, 0] {
     ^bb0(%arg1: index, %arg2: index, %arg3: index):
       linalg.yield %pad_value : f32
     } : tensor<2x?x2xf32> to tensor<2x3x4xf32>
   return %0 : tensor<2x3x4xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @pad_static_source(
+//  CHECK-SAME:                  %[[ARG0:.*]]: tensor<2x5x2xf32>, %[[PAD:.*]]: f32
+//   CHECK-NOT:   linalg.pad_tensor
+//   CHECK-DAG:   %[[C0:.*]] = constant 0 : index
+//   CHECK-DAG:   %[[C2:.*]] = constant 2 : index
+//       CHECK:   %[[INIT:.*]] = linalg.init_tensor [2, 6, 4] : tensor<2x6x4xf32>
+//       CHECK:   %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : tensor<2x5x2xf32>, vector<2x5x2xf32>
+//       CHECK:   %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[INIT]][%[[C0]], %[[C0]], %[[C2]]] {in_bounds = [true, true, true]} : vector<2x5x2xf32>, tensor<2x6x4xf32>
+//       CHECK:   return %[[WRITE]]
+func @pad_static_source(%arg0: tensor<2x5x2xf32>, %pad_value: f32) -> tensor<2x6x4xf32> {
+  %0 = linalg.pad_tensor %arg0 low[0, 0, 2] high[0, 1, 0] {
+    ^bb0(%arg1: index, %arg2: index, %arg3: index):
+      linalg.yield %pad_value : f32
+    } : tensor<2x5x2xf32> to tensor<2x6x4xf32>
+  return %0 : tensor<2x6x4xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @pad_static_dynamic(
 //  CHECK-SAME:                          %[[SRC:.*]]: tensor<1x2x2x?xf32>, %[[LOW:.*]]: index, %[[HIGH:.*]]: index
 //   CHECK-NOT:   linalg.pad_tensor
 //   CHECK-DAG:   %[[C2:.*]] = constant 2 : index
 //   CHECK-DAG:   %[[C3:.*]] = constant 3 : index
 //   CHECK-DAG:   %[[C5:.*]] = constant 5 : index
 //       CHECK:   %[[V0:.*]] = addi %[[LOW]], %[[C2]] : index
 //       CHECK:   %[[V1:.*]] = addi %[[V0]], %[[C3]] : index