Refactor GenericPadTensorOpVectorizationPattern

[cathyzhyi]

Summary

Refactor the original code to rewrite a PadTensorOp into a
sequence of InitTensorOp, FillOp and InsertSliceOp without
vectorization by default. GenericPadTensorOpVectorizationPattern
provides a customized OptimizeCopyFn to vectorize the
copying step.

[cathyzhyi]

delete redundant comments

[cathyzhyi]

remove redundant type check

[silvas]

Can you add an integration test analogous to mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert.mlir ?

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

223

nit: Prefer PadTensorOp::Adaptor

237

nit: use llvm::hasSingleElement

256

nit: Use . at the end of the comment. Same with the other comments.

https://llvm.org/docs/CodingStandards.html#commenting

mlir/test/Dialect/Linalg/bufferize.mlir

270

can you give meaningful names to the CHECK values?

[nicolasvasilache]

Note that I asked for a similar rewrite in https://reviews.llvm.org/D102804 but it was eventually not done.
I am wondering whether the pattern you add here should be more generally be the rewrite on tensors I mentioned and then just let existing bufferization kick in?

[silvas]

In D105293#2854047, @nicolasvasilache wrote:

Note that I asked for a similar rewrite in https://reviews.llvm.org/D102804 but it was eventually not done.
I am wondering whether the pattern you add here should be more generally be the rewrite on tensors I mentioned and then just let existing bufferization kick in?

That makes sense to me!

[cathyzhyi]

address comments

[nicolasvasilache]

@cathyzhyi I copyhacked your code into a tensor rewrite here: https://reviews.llvm.org/D105317

It seems fine modulo canonicalizations missing on memref::DimOp + tensor (which also interferes with tensor::DimOp that @springerm is looking at).

Feel free to take whatever makes sense to you and refactor / land.
Depending on where you go with this I'll adapt.

Will pick it up again tomorrow.

Thanks for pushing this!

[cathyzhyi]

In D105293#2854084, @silvas wrote:

In D105293#2854047, @nicolasvasilache wrote:

Note that I asked for a similar rewrite in https://reviews.llvm.org/D102804 but it was eventually not done.
I am wondering whether the pattern you add here should be more generally be the rewrite on tensors I mentioned and then just let existing bufferization kick in?

That makes sense to me!

+1 thanks for the suggestion!

[cathyzhyi]

use rewrite pattern instead

[cathyzhyi]

@nicolasvasilache seems with the rewrite pattern version the integration test passed but there are some failures in the unit test. Will need to take a closer look tmr. Here is the unit test.

func @pad_tensor(%arg0: tensor<4x?x2x?xf32>, %arg1: index) -> tensor<4x?x?x?xf32> {
  %cst = constant 0.0 : f32
  %out = linalg.pad_tensor %arg0 low[0, 0, %arg1, 0] high[0, 0, 0, %arg1]  {
  ^bb0(%gen_arg1: index, %gen_arg2: index, %gen_arg3: index, %gen_arg4: index):  // no predecessors
    linalg.yield %cst : f32
  } : tensor<4x?x2x?xf32> to tensor<4x?x?x?xf32>
  return %out : tensor<4x?x?x?xf32>
}

[springerm]

The functionality of this pattern is already provided by GenericPadTensorOpVectorizationPattern. It also translates a PadTensorOp into InitTensorOp + FillOp + InsertSliceOp. However, GenericPadTensorOpVectorizationPattern also does a bit more: It tries to generate vectorized alternatives to FillOp and InsertSliceOp. Only if there is not enough static type information, it generates FillOp and InsertSliceOp.

See https://reviews.llvm.org/D103679 for details. (The pattern was extended in subsequent commits, so check on Github for the most recent version.)

Can you run the vectorization pass or is vectorization not desired in your use case? If you want just InitTensorOp + FillOp + InsertSliceOp and no vectorization, I think there should be a way to avoid duplicating that functionality.

[nicolasvasilache]

In D105293#2854620, @cathyzhyi wrote:

@nicolasvasilache seems with the rewrite pattern version the integration test passed but there are some failures in the unit test. Will need to take a closer look tmr. Here is the unit test.

func @pad_tensor(%arg0: tensor<4x?x2x?xf32>, %arg1: index) -> tensor<4x?x?x?xf32> {
  %cst = constant 0.0 : f32
  %out = linalg.pad_tensor %arg0 low[0, 0, %arg1, 0] high[0, 0, 0, %arg1]  {
  ^bb0(%gen_arg1: index, %gen_arg2: index, %gen_arg3: index, %gen_arg4: index):  // no predecessors
    linalg.yield %cst : f32
  } : tensor<4x?x2x?xf32> to tensor<4x?x?x?xf32>
  return %out : tensor<4x?x?x?xf32>
}

It seems the version @springerm has for vectorization is already fitting the bill and would need to be split out from the vectorization part.
@springerm , indeed if we get to bufferization and still have this form, we should bufferize without relying on vectorization.

[springerm]

After looking at this in more detail, the easiest way would be to add a boolean (template) parameter to GenericPadTensorOpVectorizationPattern, which enables or disables vectorization (and maybe rename the pattern). Then register the pattern in applyEnablingTransformations. However, then we would have a dependency of ComprehensiveBufferize.cpp on Vectorization.cpp. Not sure if that is a good idea. I don't see a better way of factoring out common functionality.

Alternatively, just duplicate the pattern (without vectorization) as this revision does at the moment. With the above code suggestion, we would be duplicating around 50 lines of code. Maybe not too bad.

@nicolasvasilache What do you think?

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

2490–2507 ↗

(On Diff #356084)

Can be replaced with:

Value paddingValue = op.getConstantPaddingValue();
if (!paddingValue)
  return failure();

[nicolasvasilache]

In D105293#2854818, @springerm wrote:

After looking at this in more detail, the easiest way would be to add a boolean (template) parameter to GenericPadTensorOpVectorizationPattern, which enables or disables vectorization (and maybe rename the pattern). Then register the pattern in applyEnablingTransformations. However, then we would have a dependency of ComprehensiveBufferize.cpp on Vectorization.cpp. Not sure if that is a good idea. I don't see a better way of factoring out common functionality.

Alternatively, just duplicate the pattern (without vectorization) as this revision does at the moment. With the above code suggestion, we would be duplicating around 50 lines of code. Maybe not too bad.

@nicolasvasilache What do you think?

Indeed, so what I have in my local branch is quite close:

1. the FillOp is not a "tryVectorize" but just a "createFillOrGenerateOp", it is unconditionally always there
2. I just deleted the copy vectorization.

It seems this could be refactored and moved to Linalg/Transforms/Transforms.h and take a lambda for the copy.
By default it would take nothing (to just lower to other tensor ops).
The vectorization pattern would just reconfigure the pattern to add its own copy vectorization mechanism.

[cathyzhyi]

refactor and reuse existing code as suggested.

[cathyzhyi]

update comments that's no longer relevant.

[cathyzhyi]

@silvas @nicolasvasilache @springerm Any suggestion on which pass to put this pattern in? There is a LinalgGeneralizationPass which currently only applies to patterns lowering to Linalg.generic.

[silvas]

In D105293#2855869, @cathyzhyi wrote:

@silvas @nicolasvasilache @springerm Any suggestion on which pass to put this pattern in? There is a LinalgGeneralizationPass which currently only applies to patterns lowering to Linalg.generic.

The pattern will be used by whichever passes need it, such as ComprehensiveBufferize, Bufferize. (you should be able to just add this pattern into Bufferize.cpp in your original patch and let dialect conversion just work.

[cathyzhyi]

add the pattern into linalg bufferization pass

[silvas]

Oh, sorry, I thought you meant which woudl use it. It would be nice if it was in a generic helper file. You can probably put it next to PadTensorOpTransformationPattern for consistency.

[cathyzhyi]

move the new pattern's source code next to PadTensorOpTransformationPattern for
consistency

[silvas]

LGTM. Wait for @springerm / @nicolasvasilache approval too.

[springerm]

Nice refactoring!

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

689–690 ↗

(On Diff #356269)

can be deleted

[cathyzhyi]

delete redundant comments as suggested

[nicolasvasilache]

This has reached consensus, I suspect @cathyzhyi does not yet have write privilege to github, committing on her behalf.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

let extraClassDeclaration = [{
 
     RankedTensorType getSourceType() {
       return source().getType().cast<RankedTensorType>();
     }
     RankedTensorType getResultType() {
       return getResult().getType().cast<RankedTensorType>();
     }
 
+    // Infer the dynamic shape of the result tensor along each dim
+    SmallVector<Value> getResultTypeShapes(OpBuilder &b);
+
     // Infer the shape of the result tensor given the static shapes
     // and element type of the result tensor.
     static RankedTensorType inferResultType(RankedTensorType sourceType,
                                 ArrayRef<int64_t> staticLow,
                                 ArrayRef<int64_t> staticHigh);
 
     // Return a PadTensorOp that pads `source` to `type` size where the static
     // sizes are assumed to be greater than the dynamic sizes. The op performs

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

for (int i = 0; i < rank; ++i) {
     auto highValue = builder.createOrFold<SubIOp>(loc, resultDimSize, dimOp);
     high.push_back(highValue);
     low.push_back(builder.createOrFold<ConstantIndexOp>(loc, 0));
   }
   return PadTensorOp::createPadScalarOp(type, source, pad, low, high, loc,
                                         builder);
 }
 
-LogicalResult PadTensorOp::reifyReturnTypeShapesPerResultDim(
-    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+SmallVector<Value> PadTensorOp::getResultTypeShapes(OpBuilder &b) {
   Location loc = getLoc();
   auto lowPad = getMixedLowPad();
   auto highPad = getMixedHighPad();
   SmallVector<Value> shapes;
   for (auto dim : llvm::seq<int64_t>(0, getSourceType().getRank())) {
     // Shape along each dimension is source dim + low pad + high pad.
     SmallVector<Value> mapOperands;
     mapOperands.push_back(b.createOrFold<memref::DimOp>(loc, source(), dim));
auto addOpFoldResult = [&](OpFoldResult valueOrAttr) {
           valueOrAttr.get<Attribute>().cast<IntegerAttr>().getInt();
       expr = expr + staticValue;
     };
     addOpFoldResult(lowPad[dim]);
     addOpFoldResult(highPad[dim]);
     shapes.push_back(applyMapToValues(
         b, loc, AffineMap::get(1, numSymbols, expr), mapOperands)[0]);
   }
-  reifiedReturnShapes.emplace_back(std::move(shapes));
+  return shapes;
+}
+
+LogicalResult PadTensorOp::reifyReturnTypeShapesPerResultDim(
+    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+  reifiedReturnShapes.emplace_back(getResultTypeShapes(b));
   return success();
 }
 
 namespace {
 // Folds linalg.pad_tensor when padding is static zeros.
 struct FoldStaticZeroPadding : public OpRewritePattern<PadTensorOp> {
   using OpRewritePattern<PadTensorOp>::OpRewritePattern;
 

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

matchAndRewrite(FillOp op, ArrayRef<Value> operands,
 
     rewriter.create<FillOp>(op.getLoc(), adaptor.value(), adaptor.output());
     rewriter.replaceOp(op, adaptor.output());
 
     return success();
   }
 };
 
+/// Returns the static/dynamic mixed sizes of the memref.
+static SmallVector<OpFoldResult> getMixedSizes(OpBuilder &b, Location loc,
+                                               Value memref) {
+  auto inputType = memref.getType().cast<ShapedType>();
+  auto inputShape = inputType.getShape();
+  SmallVector<OpFoldResult> sizeMixedValues;
+  for (int64_t i = 0; i < inputType.getRank(); ++i) {
+    if (inputShape[i] == ShapedType::kDynamicSize) {
+      Value dim = b.create<memref::DimOp>(loc, memref, i);
+      sizeMixedValues.push_back(dim);
+    } else {
+      sizeMixedValues.push_back(b.getI64IntegerAttr(inputShape[i]));
+    }
+  }
+  return sizeMixedValues;
+}
+
+/// Conversion pattern that bufferizes `linalg.pad_tensor` operation.
+class BufferizePadTensorOp : public OpConversionPattern<PadTensorOp> {
+public:
+  using OpConversionPattern<PadTensorOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(PadTensorOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    Location loc = op->getLoc();
+    linalg::PadTensorOpAdaptor adaptor(operands, op->getAttrDictionary());
+    Value sourceMemRef = adaptor.source();
+    assert(sourceMemRef.getType().isa<MemRefType>());
+
+    auto sourceType = sourceMemRef.getType().cast<ShapedType>();
+    // Allocate the destination buffer
+    SmallVector<Value> shapes = op.getResultTypeShapes(rewriter);
+    SmallVector<int64_t> dynShapes(sourceType.getRank(), -1);
+    auto memrefType = MemRefType::get(dynShapes, sourceType.getElementType());
+    Value destMemRef =
+        rewriter.create<memref::AllocOp>(loc, memrefType, shapes);
+
+    // Get padding value and fill the destination buffer.
+    auto yieldOps = op.region().getOps<linalg::YieldOp>();
+    if (std::distance(yieldOps.begin(), yieldOps.end()) != 1) {
+      return rewriter.notifyMatchFailure(op,
+                                         "linalg.pad_tensor with more than one "
+                                         "padding value is not supported");
+    }
+    Value paddingValue = (*yieldOps.begin()).values()[0];
+    auto constOp = paddingValue.getDefiningOp<ConstantOp>();
+    if (!constOp) {
+      return rewriter.notifyMatchFailure(
+          op,
+          "linalg.pad_tensor with non-constant padding value is not supported");
+    }
+    if (constOp.getValue().isa<DenseElementsAttr>()) {
+      return rewriter.notifyMatchFailure(
+          op, "linalg.pad_tensor with non-scalar constant padding value is not "
+              "supported");
+    }
+    rewriter.create<linalg::FillOp>(loc, paddingValue, destMemRef);
+
+    // Get the interior region
+    SmallVector<OpFoldResult> sizes =
+        getMixedSizes(rewriter, loc, sourceMemRef);
+    SmallVector<OpFoldResult> strides(sourceType.getRank(),
+                                      rewriter.getI64IntegerAttr(1));
+    auto resultSubView = rewriter.create<memref::SubViewOp>(
+        loc, destMemRef, op.getMixedLowPad(), sizes, strides);
+    // Copy input into the interior region
+    rewriter.create<linalg::CopyOp>(loc, sourceMemRef, resultSubView);
+    auto newResultType = getTypeConverter()->convertType(op.getType());
+    rewriter.replaceOpWithNewOp<memref::CastOp>(op, newResultType, destMemRef);
+    return success();
+  }
+};
+
 /// Generic conversion pattern that matches any LinalgOp. This avoids template
 /// instantiating one pattern for each LinalgOp.
 class BufferizeAnyLinalgOp : public OpInterfaceConversionPattern<LinalgOp> {
 public:
   using OpInterfaceConversionPattern<LinalgOp>::OpInterfaceConversionPattern;
 
   LogicalResult
   matchAndRewrite(LinalgOp op, ArrayRef<Value> operands,
void mlir::linalg::populateLinalgBufferizePatterns(
   // clang-format off
   patterns.add<
       BufferizeAnyLinalgOp,
       BufferizeFillOp,
       BufferizeInitTensorOp,
       BufferizeTensorReshapeOp<TensorExpandShapeOp>,
       BufferizeTensorReshapeOp<TensorCollapseShapeOp>,
       ExtractSliceOpConverter,
-      InsertSliceOpConverter
+      InsertSliceOpConverter,
+      BufferizePadTensorOp
     >(typeConverter, patterns.getContext());
   // clang-format on
 }

mlir/test/Dialect/Linalg/bufferize.mlir

%out = linalg.tensor_collapse_shape %arg0 [[0, 1]] :
      tensor<4x5xf32> into tensor<20xf32>
   return %out : tensor<20xf32>
 }
 // CHECK: %[[MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<4x5xf32>
 // CHECK: %[[RESHAPE:.*]] = linalg.collapse_shape %[[MEMREF]] {{\[}}[0, 1]]
 // CHECK-SAME: : memref<4x5xf32> into memref<20xf32>
 // CHECK: %[[TENSOR:.*]] = memref.tensor_load %[[RESHAPE]] : memref<20xf32>
 // CHECK: return %[[TENSOR]]
+
+// CHECK-LABEL:   func @bufferize_pad_tensor(
+// CHECK-SAME:                               %[[VAL_0:.*]]: tensor<4x?x2x?xf32>,
+// CHECK-SAME:                               %[[VAL_1:.*]]: index) -> tensor<4x?x?x?xf32> {
+// CHECK:           %[[VAL_2:.*]] = constant 3 : index
+// CHECK:           %[[VAL_3:.*]] = constant 1 : index
+// CHECK:           %[[VAL_4:.*]] = constant 0.000000e+00 : f32
+// CHECK:           %[[VAL_5:.*]] = memref.buffer_cast %[[VAL_0]] : memref<4x?x2x?xf32>
+// CHECK:           %[[VAL_6:.*]] = memref.dim %[[VAL_0]], %[[VAL_3]] : tensor<4x?x2x?xf32>
+// CHECK:           %[[VAL_7:.*]] = affine.apply #map0(){{\[}}%[[VAL_1]]]
+// CHECK:           %[[VAL_8:.*]] = memref.dim %[[VAL_0]], %[[VAL_2]] : tensor<4x?x2x?xf32>
+// CHECK:           %[[VAL_9:.*]] = affine.apply #map1(){{\[}}%[[VAL_1]], %[[VAL_8]]]
+// CHECK:           %[[VAL_10:.*]] = memref.alloc(%[[VAL_6]], %[[VAL_7]], %[[VAL_9]]) : memref<4x?x?x?xf32>
+// CHECK:           linalg.fill(%[[VAL_4]], %[[VAL_10]]) : f32, memref<4x?x?x?xf32>
+// CHECK:           %[[VAL_11:.*]] = memref.subview %[[VAL_10]][0, 0, %[[VAL_1]], 0] [4, %[[VAL_6]], 2, %[[VAL_8]]] [1, 1, 1, 1] : memref<4x?x?x?xf32> to memref<4x?x2x?xf32, #map2>
+// CHECK:           linalg.copy(%[[VAL_5]], %[[VAL_11]]) : memref<4x?x2x?xf32>, memref<4x?x2x?xf32, #map2>
+// CHECK:           %[[VAL_12:.*]] = memref.tensor_load %[[VAL_10]] : memref<4x?x?x?xf32>
+// CHECK:           return %[[VAL_12]] : tensor<4x?x?x?xf32>
+func @bufferize_pad_tensor(%arg0: tensor<4x?x2x?xf32>, %arg1: index) -> tensor<4x?x?x?xf32> {
+  %c0 = constant 0 : index
+  %cst = constant 0.0 : f32
+  %out = linalg.pad_tensor %arg0 low[%c0, %c0, %arg1, %c0] high[%c0, %c0, %c0, %arg1]  {
+  ^bb0(%gen_arg1: index, %gen_arg2: index, %gen_arg3: index, %gen_arg4: index):  // no predecessors
+    linalg.yield %cst : f32
+  } : tensor<4x?x2x?xf32> to tensor<4x?x?x?xf32>
+  return %out : tensor<4x?x?x?xf32>
+}