[mlir][linalg] Transform PadTensorOp into InitOp, FillOp, GenericOp

[agostini01]

Summary

Introduces a test pass that rewrites PadTensorOps with static shapes as a sequence of:

linalg.init_tensor // to create output
linalg.fill        // to initialize with padding value
linalg.generic     // to copy the original contents to the padded tensor

The pass can be triggered with:

• --test-linalg-transform-patterns="test-transform-pad-tensor"

[agostini01]

Added missing files for the path

[nicolasvasilache]

Very cool, thanks for your contribution!
@hanchung @springerm FYI

I'll review tomorrow.

[agostini01]

Thank you!
Really appreciate it!
Let me know if any changes are needed.

[nicolasvasilache]

Approved modulo moving the piece of code that looks like a pad helper.
Thanks for your contribution!

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

860

s/linalg.generics/linalg.generic/

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

654

Not sure whether you'd need this in a (near) future but I think this can also be generalized by using subtensor_insert:

linalg.init_tensor // to create output
linalg.fill // to initialize with padding value
linalg.subtensor_insert // to copy the original contents to the padded tensor

You could add a simple affine.apply op that does dim(i) + low(i) + high(i) and handle all cases ?
If you did this with makeComposedAffineApply, you'd even get the static shape out of the box and you could branch to the static impl.

This could even turn into a fill + subtensor_insert canonicalization down the line.

659

This seems like a useful function to have as a helper on the pad op ?

686

Ah ok I see why you limit yourself to static cases in this revision, makes sense.

[agostini01]

• Fix typos addressing comments
• Add subtensor_insert comment on where it can replace linalg.generic

[agostini01]

Thank you for the suggestions. I have made small changes to code comments.

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

646

s/with with/with/

654

The use of subtensor_insert is an interesting suggestion. I will experiment with it and submit a new patch if my personal end-to-end examples work with it.

659

I decided against implementing it as a helper function because I see this pattern (cast to shape, hasStaticShape) being used all the time in other parts of the codebase.

684

Added this comment as a reminder that code below can be replaced by subtensor_insert.

[agostini01]

@nicolasvasilache, @mehdi_amini it is the first time I will land a patch through phabricator.

After checking out the branch for this feature I executed:

git pull --rebase https://github.com/llvm/llvm-project.git main
cmake --build /working_dir/llvm-project/build/ --target mlir-opt --target check-mlir # with no failed tests
arc land

Which produce the following output

 STRATEGY  Merging with "squash" strategy, the default strategy.
 SOURCE  Landing the current branch, "lower-pad-tensor".
 ONTO REMOTE  Landing onto remote "origin", the default remote under Git.
 ONTO TARGET  Refs were selected by reading "arc.land.onto" configuration: main.
 INTO REMOTE  Will merge into remote "origin" by default, because this is the remote the change is landing onto.
 INTO TARGET  Will merge into target "main" by default, because this is the "onto" target.
 TARGET  No local copy of ref "main" in remote "origin" exists, attempting fetch...

  $   git fetch --no-tags --quiet -- origin main


 INTO COMMIT  Preparing merge into the empty state to create target "main" in remote "origin".

 <!> LARGE WORKING SET 
There are 23,664 commit(s) reachable from the specified sources
("lower-pad-tensor"). You are landing into the empty state, so all of these
commits will land:

  7d26d6a1b062 Sema: add support for `__attribute__…
  c3fd2a50ba13 [libc] Remove special case for 8 and 16 bytes
  609f5e050cea [mlir] Rename 'setInsertionPointAfter' to avoid ambiguity
  9e3842d60351 [OPENMP]Fix codegen for is_device_ptr component, captured by…
  c3e6054b07be [OpenMP] Additional Information for Libomptarget Mappings
  < ... 23,654 more commits ... >
  0d6e32dcc207 [MLIR][LINALG] Transform PadTensorOp in TensorInitOp and…
  7c89360749fc [MLIR][LINALG] Replace init with GenericOp by FillOp on…
  cd5c8e72c5ab [mlir][linalg] Transform PadTensorOp into InitOp, FillOp…
  fc9fd24e2ba5 Fix typos addressing comments
  db8041a7b8f5 Add subtensor_insert comment on where it can replace linalg.
  >>>  Land 23,664 commit(s)? [y/N/?]

I interrupted this command.
Is arc land the correct command to land these changes?
LLVM docs mention git push.
Should I use git push? Are there arguments I should pass so that the commit message that describes this differential revision are sent in squash mode?
Sorry for the trivial question, just wanted to make sure.

[springerm]

Is arc land the correct command to land these changes?
LLVM docs mention git push.
Should I use git push? Are there arguments I should pass so that the commit message that describes this differential revision are sent in squash mode?
Sorry for the trivial question, just wanted to make sure.

I never tried arc land. I just checkout the main branch, do a git pull origin main, checkout a new branch (git checkout -b br1), cherry-pick over the commit that I want to land, do a git log to see if everything is OK, then push to main (git push origin HEAD:main).

Also, do you have push rights to the repo yet? If this is your first commit, then probably not. In that case, someone else must push the commit on your behalf and you can ask for push rights after showing a good commits. See https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access.

[agostini01]

Thank you for the explanation, now I know what I must do.

In D102804#2795585, @springerm wrote:

Also, do you have push rights to the repo yet? If this is your first commit, then probably not. In that case, someone else must push the commit on your behalf and you can ask for push rights after showing a good commits. See https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access.

But you are right, this is my first commit to the llvm-project. It would be great if you or another subscriber for this patch could land it for me.

Below is my name and email as suggested on the instructions:
Nicolas Agostini
[e-mail]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 void populateLinalgConvGeneralizationPatterns(
     RewritePatternSet &patterns,
     LinalgTransformationFilter filter = LinalgTransformationFilter());
 
 //===----------------------------------------------------------------------===//
 // Op-specific patterns.
 //===----------------------------------------------------------------------===//
 
+/// PadTensorOp is not canonicalized away yet, so we provide a transformation to
+/// `linalg.generic`.
+struct PadTensorOpTransformationPattern : public OpRewritePattern<PadTensorOp> {
+  using OpRewritePattern<PadTensorOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(PadTensorOp padOp,
+                                PatternRewriter &rewriter) const override;
+};
+
 /// PadTensorOp does not implement the LinalgStructuredOpInterface `LinalgOp`,
 /// it needs a specific pattern to vectorize.
 struct PadTensorOpVectorizationPattern : public OpRewritePattern<PadTensorOp> {
   using OpRewritePattern<PadTensorOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(PadTensorOp padOp,
                                 PatternRewriter &rewriter) const override;
 };

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

if (auto cst = e.dyn_cast<AffineConstantExpr>()) {
       rewriter.replaceOpWithNewOp<AffineApplyOp>(minOp, resultMap,
                                                  resultOperands);
     }
     return success();
   }
 
   return failure();
 }
+
+static SmallVector<StringRef> getNParallelLoopsAttrs(unsigned nParallelLoops) {
+  return SmallVector<StringRef>(nParallelLoops, getParallelIteratorTypeName());
+}
+
+/// Rewrite a PadTensorOp into a sequence of InitTensorOp, FillOp (to initialize
+/// with pad_val) and GenericOp (to copy contents).
+LogicalResult PadTensorOpTransformationPattern::matchAndRewrite(
+    linalg::PadTensorOp padOp, PatternRewriter &rewriter) const {
+
+  auto inputShapedType = padOp.source().getType().cast<ShapedType>();
+  auto resultShapedType = padOp.result().getType().cast<ShapedType>();
+
+  // Bail on non-static shapes.
+  if (!inputShapedType.hasStaticShape())
+    return failure();
+  if (!resultShapedType.hasStaticShape())
+    return failure();
+
+  // Only support padding with a constant for now, i.e. either:
+  //   1. A BBarg from a different block.
+  //   2. A value defined outside of the current block.
+  Block &block = padOp.region().front();
+  auto yieldOp = cast<YieldOp>(block.getTerminator());
+  assert(yieldOp.getNumOperands() == 1 && "expected single operand yield");
+  Value padValue = yieldOp.values().front();
+  Operation *definingOp = padValue.getDefiningOp();
+  if (definingOp && definingOp->getBlock() == &block)
+    return failure();
+  if (!definingOp && padValue.cast<BlockArgument>().getOwner() == &block)
+    return failure();
+
+  // Create tensor with the padded shape
+  Location loc = padOp.getLoc();
+  SmallVector<Value> indices(resultShapedType.getRank(),
+                             rewriter.create<ConstantIndexOp>(loc, 0));
+  Value initTensor = rewriter.create<InitTensorOp>(
+      loc, resultShapedType.getShape(), resultShapedType.getElementType());
+
+  // Initialize tensor with the pad value
+  Value tmpTensor =
+      rewriter.create<linalg::FillOp>(loc, initTensor, padValue).result();
+
+  // Copy original contents into new tensor
+  // Uses linalg.generic, but could be done with std.subtensor_insert
+  SmallVector<AffineExpr, 4> outputExprs;
+  for (unsigned i = 0; i < resultShapedType.getRank(); ++i) {
+    outputExprs.push_back(getAffineDimExpr(i, rewriter.getContext()) +
+                          padOp.static_low()[i].cast<IntegerAttr>().getInt());
+  }
+
+  SmallVector<AffineMap, 2> transferMaps = {
+      rewriter.getMultiDimIdentityMap(inputShapedType.getRank()),
+      AffineMap::get(resultShapedType.getRank(),
+                     /*symbolCount=*/0, outputExprs, rewriter.getContext())};
+
+  rewriter.replaceOpWithNewOp<linalg::GenericOp>(
+      padOp, resultShapedType, padOp.source(), tmpTensor, transferMaps,
+      getNParallelLoopsAttrs(resultShapedType.getRank()),
+      [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
+        nestedBuilder.create<linalg::YieldOp>(nestedLoc, args[0]);
+      });
+
+  return success();
+}

mlir/test/Dialect/Linalg/lower-pad-tensor.mlir

+// RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-transform-pad-tensor"  %s | FileCheck --check-prefix=CHECK %s
+
+// CHECK-DAG:   #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+// CHECK-DAG:   #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0 + 1, d1 + 1, d2 + 1, d3 + 2)>
+// CHECK-LABEL: func @pad_tensor_with_memrefs
+func @pad_tensor_with_memrefs(%arg0: memref<1x28x28x1xf32>) -> memref<2x31x31x3xf32> {
+  %cst = constant 0.000000e+00 : f32
+  %0 = memref.tensor_load %arg0 : memref<1x28x28x1xf32>
+  %1 = linalg.pad_tensor %0 low[1, 1, 1, 2] high[0, 2, 2, 0]  {
+  ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):  // no predecessors
+    linalg.yield %cst : f32
+  } : tensor<1x28x28x1xf32> to tensor<2x31x31x3xf32>
+  %2 = memref.buffer_cast %1 : memref<2x31x31x3xf32>
+  return %2 : memref<2x31x31x3xf32>
+}
+
+// CHECK:       linalg.fill
+// CHECK:       linalg.generic
+// CHECK-SAME:     indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
+
+// -----
+
+// CHECK-DAG:   #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG:   #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0 + 1, d1 + 2, d2 + 2)>
+// CHECK-LABEL: func @pad_tensor_no_memrefs
+func @pad_tensor_no_memrefs(%arg0: tensor<1x28x28xf32>) -> tensor<2x32x32xf32> {
+  %cst = constant 0.000000e+00 : f32
+  %0 = linalg.pad_tensor %arg0 low[1, 2, 2] high[0, 2, 2]  {
+  ^bb0(%arg1: index, %arg2: index, %arg3: index):  // no predecessors
+    linalg.yield %cst : f32
+  } : tensor<1x28x28xf32> to tensor<2x32x32xf32>
+  return %0 : tensor<2x32x32xf32>
+}
+
+// CHECK:       linalg.fill
+// CHECK:       linalg.generic
+// CHECK-SAME:      indexing_maps = [#[[$MAP2]], #[[$MAP3]]]
+
+// -----
+
+// CHECK-DAG:   #[[$MAP4:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+// CHECK-DAG:   #[[$MAP5:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 + 2, d2 + 2, d3)>
+// CHECK-LABEL: func @pad_tensor_detailed
+func @pad_tensor_detailed(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {
+  %cst = constant 0.000000e+00 : f32
+  %0 = linalg.pad_tensor %arg0 low[0, 2, 2, 0] high[0, 2, 2, 0]  {
+  ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):  // no predecessors
+    linalg.yield %cst : f32
+  } : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32>
+  return %0 : tensor<1x32x32x1xf32>
+}
+
+// CHECK:      %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32>
+// CHECK:      %[[CTE:.+]] = constant 0.000000e+00 : f32
+// CHECK:      %[[TMP:.+]] = linalg.init_tensor [1, 32, 32, 1] : tensor<1x32x32x1xf32>
+// CHECK:      %[[R1c:.+]] = linalg.fill
+// CHECK:      %[[R2c:.+]] = linalg.generic
+// CHECK-SAME:   indexing_maps = [#[[$MAP4]], #[[$MAP5]]]
+// CHECK-SAME:   iterator_types = ["parallel", "parallel", "parallel", "parallel"] 
+// CHECK:        ins(%arg0 : tensor<1x28x28x1xf32>) outs(%1 : tensor<1x32x32x1xf32>)
+// CHECK:      ^bb0(%[[VAL:.+]]: f32, %arg2: f32)
+// CHECK:        linalg.yield %[[VAL]] : f32
+// CHECK:      return %[[R2c:.+]]

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

Option<bool> testAffineMinSCFCanonicalizationPatterns{
       llvm::cl::desc("Test affine-min + scf canonicalization patterns."),
       llvm::cl::init(false)};
   Option<bool> testTileAndPadPattern{
       *this, "test-tile-and-pad-pattern",
       llvm::cl::desc("Test tile and pad pattern"), llvm::cl::init(false)};
   Option<int> testHoistPadding{*this, "test-hoist-padding",
                                llvm::cl::desc("Test hoist padding"),
                                llvm::cl::init(0)};
+  Option<bool> testTransformPadTensor{
+      *this, "test-transform-pad-tensor",
+      llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
+      llvm::cl::init(false)};
   ListOption<int64_t> tileSizesForPadding{
       *this, "tile-sizes-for-padding",
       llvm::cl::desc("Linalg tile sizes when tile+pad"), llvm::cl::ZeroOrMore,
       llvm::cl::MiscFlags::CommaSeparated};
   ListOption<unsigned> testInterchangePattern{
       *this, "test-interchange-pattern", llvm::cl::MiscFlags::CommaSeparated,
       llvm::cl::desc("Test the interchange pattern.")};
 };
static void applyLinalgToVectorPatterns(FuncOp funcOp) {
   patterns.add<LinalgVectorizationPattern>(
       funcOp.getContext(),
       LinalgTransformationFilter()
           .addOpFilter<ContractionOpInterface, FillOp, CopyOp, GenericOp>());
   patterns.add<PadTensorOpVectorizationPattern>(funcOp.getContext());
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
+static void applyPadTensorToGenericPatterns(FuncOp funcOp) {
+  RewritePatternSet patterns(funcOp.getContext());
+  patterns.add<PadTensorOpTransformationPattern>(funcOp.getContext());
+  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+}
+
 static void applyAffineMinSCFCanonicalizationPatterns(FuncOp funcOp) {
   RewritePatternSet foldPattern(funcOp.getContext());
   foldPattern.add<AffineMinSCFCanonicalizationPattern>(funcOp.getContext());
   FrozenRewritePatternSet frozenPatterns(std::move(foldPattern));
 
   // Explicitly walk and apply the pattern locally to avoid more general folding
   // on the rest of the IR.
   funcOp.walk([&frozenPatterns](AffineMinOp minOp) {
void TestLinalgTransforms::runOnFunction() {
   if (testMatmulToVectorPatterns1dTiling || testMatmulToVectorPatterns2dTiling)
     return applyMatmulToVectorPatterns(getFunction(),
                                        testMatmulToVectorPatterns1dTiling,
                                        testMatmulToVectorPatterns2dTiling);
   if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getFunction());
   if (testGenericToVectorPattern)
     return applyLinalgToVectorPatterns(getFunction());
+  if (testTransformPadTensor)
+    return applyPadTensorToGenericPatterns(getFunction());
   if (testAffineMinSCFCanonicalizationPatterns)
     return applyAffineMinSCFCanonicalizationPatterns(getFunction());
   if (testTileAndPadPattern)
     return applyTileAndPadPattern(getFunction(), tileSizesForPadding);
   if (testHoistPadding) {
     getFunction().walk([&](linalg::PadTensorOp padTensorOp) {
       (void)linalg::hoistPaddingOnTensors(padTensorOp, testHoistPadding);
     });