[mlir][linalg] Lower PadTensorOp to InitTensorOp + FillOp + SubTensorInitOp

[springerm]

Summary

Currently limited to constant pad values. Any combination of dynamic/static tensor sizes and padding sizes is supported.

Depends On D103590

[nicolasvasilache]

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

677

The key reason why this pattern is needed here is that it triggers vectorization on subtensor_insert (which is generally undesirable because bufferization generally makes them go away for free).
Please document this.

739

I would immediately vectorize SubTensorInsertOp.
I wouldn't implement a vectorization pattern for SubTensorInsertOp right now because it is generally dangerous (see bufferization comment above).
This way your GenericPadTensorOpVectorizationPattern becomes a true vectorization pattern and you can drop the note.

If / when multiple vectorization patterns occur we can factor it out and be very loud about how it should / should not be used.

[springerm]

rebase

[springerm]

address comments

[springerm]

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

677

This functionality is coming the next commit, see comment below.

739

I'm putting that in a separate commit if that's OK. (To keep the change small.) Added a comment at the beginning of the function to explain what's still missing.

[nicolasvasilache]

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

739

ah yes, I missed the intent that the "not yet a vectorization" conveys SGTM then!

[springerm]

rebase

[springerm]

small change

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

[springerm]

no change

[springerm]

small improvements

[springerm]

no change

[springerm]

rebase

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

llvm::for_each(ofrs, [&](auto o) {
     } else {
       result.push_back(builder.create<ConstantIndexOp>(
           loc, getIntFromAttr(o.template get<Attribute>())));
     }
   });
   return result;
 }
 
+/// 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.
+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) {
+      if (auto val = ofr.dyn_cast<Value>())
+        return val;
+      return rewriter.create<ConstantIndexOp>(
+          padOp.getLoc(), getIntFromAttr(ofr.get<Attribute>())).getResult();
+    };
+
+    // Pad value must be a constant.
+    auto padValue = padOp.getConstantPaddingValue();
+    if (!padValue) return failure();
+
+    auto resultType = padOp.getResultType();
+    // Compute size of InitTensorOp. Any combination of static/dynamic is
+    // supported.
+    SmallVector<Value> dynSizes;
+    SmallVector<int64_t> staticSizes;
+    for (unsigned dim = 0; dim < resultType.getRank(); ++dim) {
+      if (resultType.isDynamicDim(dim)) {
+        auto srcSize = rewriter.createOrFold<memref::DimOp>(
+            padOp.getLoc(), padOp.source(), dim);
+        // Add low and high padding value.
+        auto plusLow = rewriter.createOrFold<AddIOp>(
+            padOp.getLoc(), srcSize, getIdxValue(padOp.getMixedLowPad()[dim]));
+        auto plusHigh = rewriter.createOrFold<AddIOp>(
+            padOp.getLoc(), plusLow, getIdxValue(padOp.getMixedHighPad()[dim]));
+        dynSizes.push_back(plusHigh);
+      }
+      staticSizes.push_back(resultType.getDimSize(dim));
+    }
+
+    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();
+    // 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();
+  }
+};
+
 /// 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;
 
   LogicalResult matchAndRewrite(PadTensorOp padOp,
                                 PatternRewriter &rewriter) const final {
rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
         insertOp, read, insertOp.dest(), writeIndices, inBounds);
 
     return success();
   }
 };
 
 void mlir::linalg::populatePadTensorOpVectorizationPatterns(
     RewritePatternSet &patterns, PatternBenefit baseBenefit) {
-  // TODO: Canonicalizer handles simple cases where low = 0 and high = 0, but a
-  // generic vectorization pattern is still missing.
-
+  patterns.add<GenericPadTensorOpVectorizationPattern>(
+      patterns.getContext(), baseBenefit);
   // Try these specialized patterns first before resorting to the generic one.
   patterns.add<PadTensorOpVectorizationWithTransferReadPattern,
                PadTensorOpVectorizationWithTransferWritePattern,
                PadTensorOpVectorizationWithSubTensorInsertPattern>(
-      patterns.getContext(), baseBenefit);
+      patterns.getContext(), baseBenefit.getBenefit() + 1);
 }
 
 // TODO: cleanup all the convolution vectorization patterns.
 template <class ConvOp, int N>
 LogicalResult ConvOpVectorization<ConvOp, N>::matchAndRewrite(
     ConvOp op, PatternRewriter &rewriter) const {
   Location loc = op.getLoc();
   MLIRContext *context = op.getContext();

mlir/test/Dialect/Linalg/vectorization.mlir

func @matmul_i8_i8_i32(%a: memref<4x6xi8>, %b: memref<6x12xi8>, %c: memref<4x12xi32>) {
   //  CHECK-SAME:     vector<4x12xi32>, memref<4x12xi32>
   linalg.matmul_i8_i8_i32 ins(%a, %b : memref<4x6xi8>, memref<6x12xi8>)
     outs(%c: memref<4x12xi32>)
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @pad_static_high_padding
-//       CHECK:   linalg.pad_tensor
-func @pad_static_high_padding(%arg0: tensor<?x?x?xf32>, %pad_value: f32) -> tensor<2x3x4xf32> {
-  %0 = linalg.pad_tensor %arg0 low[0, 0, 0] high[0, 1, 0] {
+// CHECK-LABEL: func @pad_static(
+//  CHECK-SAME:                  %[[ARG0:.*]]: tensor<2x?x2xf32>, %[[PAD:.*]]: f32
+//   CHECK-NOT:   linalg.pad_tensor
+//   CHECK-DAG:   %[[C1:.*]] = constant 1 : index
+//   CHECK-DAG:   %[[INIT:.*]] = linalg.init_tensor [2, 3, 4] : tensor<2x3x4xf32>
+//   CHECK-DAG:   %[[VEC:.*]] = vector.broadcast %[[PAD]] : f32 to vector<2x3x4xf32>
+//       CHECK:   %[[FILL:.*]] = vector.transfer_write %[[VEC]], %[[INIT]]{{.*}} : vector<2x3x4xf32>, tensor<2x3x4xf32>
+//   CHECK-DAG:   %[[DIM1:.*]] = memref.dim %[[ARG0]], %[[C1]]
+//       CHECK:   %[[RESULT:.*]] = subtensor_insert %[[ARG0]] into %2[0, 0, 2] [2, %[[DIM1]], 2] [1, 1, 1] : tensor<2x?x2xf32> into tensor<2x3x4xf32>
+//       CHECK:   return %[[RESULT]]
+func @pad_static(%arg0: tensor<2x?x2xf32>, %pad_value: f32) -> tensor<2x3x4xf32> {
+  %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<?x?x?xf32> to tensor<2x3x4xf32>
+    } : tensor<2x?x2xf32> to tensor<2x3x4xf32>
   return %0 : tensor<2x3x4xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @pad_dynamic
-//       CHECK:   linalg.pad_tensor
-func @pad_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index,
+// 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
+//       CHECK:   %[[V2:.*]] = addi %[[HIGH]], %[[C5]] : index
+//       CHECK:   %[[DIM3:.*]] = memref.dim %[[SRC]], %[[C3]] : tensor<1x2x2x?xf32>
+//       CHECK:   %[[V4:.*]] = addi %[[DIM3]], %[[C3]] : index
+//       CHECK:   %[[V5:.*]] = addi %[[V4]], %[[C2]] : index
+//       CHECK:   %[[INIT:.*]] = linalg.init_tensor [6, %[[V1]], %[[V2]], %[[V5]]] : tensor<6x?x?x?xf32>
+//       CHECK:   %[[FILL:.*]] = linalg.fill(%[[INIT]], %{{.*}}) : tensor<6x?x?x?xf32>, f32 -> tensor<6x?x?x?xf32>
+//       CHECK:   %[[SRCDIM:.*]] = memref.dim %[[SRC]], %[[C3]] : tensor<1x2x2x?xf32>
+//       CHECK:   %[[RESULT:.*]] = subtensor_insert %[[SRC]] into %[[FILL]][2, %[[LOW]], 3, 3] [1, 2, 2, %[[SRCDIM]]] [1, 1, 1, 1] : tensor<1x2x2x?xf32> into tensor<6x?x?x?xf32>
+//       CHECK:   return %[[RESULT]]
+func @pad_static_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index,
                   %pad_value: f32) -> tensor<6x?x?x?xf32> {
   %0 = linalg.pad_tensor %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] {
     ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
       linalg.yield %pad_value : f32
     } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32>
   return %0 : tensor<6x?x?x?xf32>
 }