[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

857 ↗

(On Diff #346556)

s/linalg.generics/linalg.generic/

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

657 ↗

(On Diff #346556)

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.

662 ↗

(On Diff #346556)

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

689 ↗

(On Diff #346556)

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

684 ↗

(On Diff #347693)

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

649 ↗

(On Diff #346556)

s/with with/with/

657 ↗

(On Diff #346556)

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.

662 ↗

(On Diff #346556)

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.

[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/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestConstantFold.cpp
   TestInlining.cpp
   TestLoopFusion.cpp
-  TestLinalgLowerPadTensor.cpp
   TestLoopMapping.cpp
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   LINK_LIBS PUBLIC
   MLIRTestDialect
   MLIRTransforms
   )
 
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test)

mlir/test/lib/Transforms/TestLinalgLowerPadTensor.cpp

-//===- TestLinalgLowerPadTensor.cpp - Test Linalg hoisting functions ------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-//
-// This file implements logic for testing Linalg hoisting functions.
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
-#include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
-#include "mlir/Pass/Pass.h"
-
-using namespace mlir;
-using namespace mlir::linalg;
-
-namespace {
-
-struct TestLinalgLowerPadTensor
-    : public PassWrapper<TestLinalgLowerPadTensor, FunctionPass> {
-  TestLinalgLowerPadTensor() = default;
-  TestLinalgLowerPadTensor(const TestLinalgLowerPadTensor &pass) {}
-  void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<AffineDialect>();
-  }
-
-  void runOnFunction() override;
-
-};
-} // end anonymous namespace
-
-void TestLinalgLowerPadTensor::runOnFunction() {
-
-  auto func = getFunction();
-  func.emitWarning() << "Test";
-  bool changed = true;
-
-  while (changed) {
-    changed = false;
-
-    func.walk([&](linalg::PadTensorOp padOp) {
-      {
-        padOp.emitWarning() << "Found a PadTensorOp";
-        // if (does not have a prior load tensor, )
-        //   // leave it alone
-        //   return WalkResult::advance();
-        changed = true;
-        // padOp.erase();
-
-        // Need to interrupt and restart: erasing the loop messes up the walk.
-        return WalkResult::interrupt();
-      }
-      return WalkResult::advance();
-    });
-
-    // Find if a tensor.load originated this op, continue if true
-
-    // Create the new MemRef with padded dimensions
-    // and manually copy the data there
-    // Replaces all
-    // op.replaceAllUsesWith(affineForOp.getResults().take_back()[0]);
-    return;
-  }
-}
-
-
-namespace mlir {
-namespace test {
-void registerTestLinalgLowerPadTensor() {
-  PassRegistration<TestLinalgLowerPadTensor> testTestLinalgLowerPadTensorPass(
-      "test-linalg-lower-pad-tensor",
-      "Test lowering tensor load and pad_tensor to affine.");
-}
-} // namespace test
-} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgElementwiseFusion();
 void registerTestPushExpandingReshape();
 void registerTestLinalgFusionTransforms();
 void registerTestLinalgTensorFusionTransforms();
 void registerTestLinalgGreedyFusion();
 void registerTestLinalgHoisting();
-void registerTestLinalgLowerPadTensor();
 void registerTestLinalgTileAndFuseSequencePass();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();
 void registerTestLoopFusion();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
 void registerTestMathPolynomialApproximationPass();
 void registerTestMemRefDependenceCheck();
#endif
   test::registerTestInterfaces();
   test::registerTestLinalgCodegenStrategy();
   test::registerTestLinalgElementwiseFusion();
   test::registerTestPushExpandingReshape();
   test::registerTestLinalgFusionTransforms();
   test::registerTestLinalgTensorFusionTransforms();
   test::registerTestLinalgGreedyFusion();
   test::registerTestLinalgHoisting();
-  test::registerTestLinalgLowerPadTensor();
   test::registerTestLinalgTileAndFuseSequencePass();
   test::registerTestLinalgTransforms();
   test::registerTestLivenessPass();
   test::registerTestLoopFusion();
   test::registerTestLoopMappingPass();
   test::registerTestLoopUnrollingPass();
   test::registerTestMathPolynomialApproximationPass();
   test::registerTestMemRefDependenceCheck();