[mlir][linalg] Add padding pass to strategy passes.

[gysit]

Summary

Add a strategy pass that pads and hoists after tiling and fusion.

Depends On D112412

[gysit]

Rebase.

[gysit]

Fix padIf method to pass filter to the pad method.

mlir/include/mlir/Dialect/Linalg/Passes.h

 //===----------------------------------------------------------------------===//
 /// Create a LinalgStrategyTilePass.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgStrategyTilePass(
     StringRef opName = "",
     linalg::LinalgTilingOptions opt = linalg::LinalgTilingOptions(),
     linalg::LinalgTransformationFilter filter =
         linalg::LinalgTransformationFilter());
 
+/// Create a LinalgStrategyPadPass.
+std::unique_ptr<OperationPass<FuncOp>> createLinalgStrategyPadPass(
+    StringRef opName = "",
+    linalg::LinalgPaddingOptions opt = linalg::LinalgPaddingOptions(),
+    linalg::LinalgTransformationFilter filter =
+        linalg::LinalgTransformationFilter());
+
 /// Create a LinalgStrategyPromotePass.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgStrategyPromotePass(
     StringRef opName = "",
     linalg::LinalgPromotionOptions opt = linalg::LinalgPromotionOptions(),
     linalg::LinalgTransformationFilter filter =
         linalg::LinalgTransformationFilter());
 
 /// Create a LinalgStrategyGeneralizePass.

mlir/include/mlir/Dialect/Linalg/Passes.td

def LinalgStrategyTilePass
   let options = [
     Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
       "Which func op is the anchor to latch on.">,
     Option<"anchorOpName", "anchor-op", "std::string", /*default=*/"",
       "Which linalg op within the func is the anchor to latch on.">,
   ];
 }
 
+def LinalgStrategyPadPass
+    : FunctionPass<"linalg-strategy-pad-pass"> {
+  let summary = "Configurable pass to apply padding and hoisting.";
+  let constructor = "mlir::createLinalgStrategyPadPass()";
+  let dependentDialects = ["linalg::LinalgDialect"];
+  let options = [
+    Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
+      "Which func op is the anchor to latch on.">,
+    Option<"anchorOpName", "anchor-op", "std::string", /*default=*/"",
+      "Which linalg op within the func is the anchor to latch on.">,
+  ];
+}
+
 def LinalgStrategyPromotePass
     : FunctionPass<"linalg-strategy-promote-pass"> {
   let summary = "Configurable pass to apply pattern-based linalg promotion.";
   let constructor = "mlir::createLinalgStrategyPromotePass()";
   let dependentDialects = ["linalg::LinalgDialect"];
   let options = [
     Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
       "Which func op is the anchor to latch on.">,

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

void addToPassPipeline(OpPassManager &pm,
     pm.addPass(createLinalgStrategyTilePass(opName, options, m));
   }
 
 private:
   std::string opName;
   linalg::LinalgTilingOptions options;
 };
 
+/// Represent one application of LinalgStrategyPadPass.
+struct Pad : public Transformation {
+  Pad(StringRef name, linalg::LinalgPaddingOptions options,
+      LinalgTransformationFilter::FilterFunction f = nullptr)
+      : Transformation(f), opName(name), options(options) {}
+
+  void addToPassPipeline(OpPassManager &pm,
+                         LinalgTransformationFilter m) const override {
+    pm.addPass(createLinalgStrategyPadPass(opName, options, m));
+  }
+
+private:
+  std::string opName;
+  linalg::LinalgPaddingOptions options;
+};
+
 /// Represent one application of createLinalgStrategyPromotePass.
 struct Promote : public Transformation {
   Promote(StringRef name, linalg::LinalgPromotionOptions options,
           LinalgTransformationFilter::FilterFunction f = nullptr)
       : Transformation(f), opName(name), options(options) {}
 
   void addToPassPipeline(OpPassManager &pm,
                          LinalgTransformationFilter m) const override {
struct CodegenStrategy {
   }
   /// Conditionally append a pattern to add a level of tiling for
   /// `LinalgOpType` with tiling `options`.
   CodegenStrategy &
   tileIf(bool b, StringRef opName, linalg::LinalgTilingOptions options,
          LinalgTransformationFilter::FilterFunction f = nullptr) {
     return b ? tile(opName, options) : *this;
   }
+  /// Append a pattern to pad and hoist the operands of Op `opName` with padding
+  /// `options`.
+  CodegenStrategy &pad(StringRef opName, linalg::LinalgPaddingOptions options,
+                       LinalgTransformationFilter::FilterFunction f = nullptr) {
+    transformationSequence.emplace_back(
+        std::make_unique<Pad>(opName, options, f));
+    return *this;
+  }
+  /// Conditionally append a pattern to pad and hoist the operands of Op
+  /// `opName` with padding `options`.
+  CodegenStrategy &
+  padIf(bool b, StringRef opName, linalg::LinalgPaddingOptions options,
+        LinalgTransformationFilter::FilterFunction f = nullptr) {
+    return b ? pad(opName, options, f) : *this;
+  }
   /// Append a pattern to add a level of promotion for `LinalgOpType` with
   /// promotion `options`.
   CodegenStrategy &
   promote(StringRef opName, linalg::LinalgPromotionOptions options,
           LinalgTransformationFilter::FilterFunction f = nullptr) {
     transformationSequence.emplace_back(
         std::make_unique<Promote>(opName, options, f));
     return *this;

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

void runOnFunction() override {
     }
     (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
   }
 
   LinalgTilingOptions options;
   LinalgTransformationFilter filter;
 };
 
+/// Configurable pass to apply hoisting and padding.
+struct LinalgStrategyPadPass
+    : public LinalgStrategyPadPassBase<LinalgStrategyPadPass> {
+
+  LinalgStrategyPadPass() = default;
+
+  LinalgStrategyPadPass(StringRef opName, LinalgPaddingOptions opt,
+                        LinalgTransformationFilter filt)
+      : options(opt), filter(filt) {
+    this->anchorOpName.setValue(opName.str());
+  }
+
+  void runOnFunction() override {
+    auto funcOp = getFunction();
+    if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
+      return;
+
+    RewritePatternSet paddingPattern(funcOp.getContext());
+    if (!anchorOpName.empty()) {
+      paddingPattern.add<LinalgPaddingPattern>(
+          anchorOpName, funcOp.getContext(), options, filter);
+    } else {
+      paddingPattern.add<LinalgPaddingPattern>(funcOp.getContext(), options,
+                                               filter);
+    }
+    if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(paddingPattern))))
+      signalPassFailure();
+  }
+
+  LinalgPaddingOptions options;
+  LinalgTransformationFilter filter;
+};
+
 /// Configurable pass to apply pattern-based linalg generalization.
 struct LinalgStrategyGeneralizePass
     : public LinalgStrategyGeneralizePassBase<LinalgStrategyGeneralizePass> {
 
   LinalgStrategyGeneralizePass() = default;
 
   LinalgStrategyGeneralizePass(StringRef opName,
                                LinalgTransformationFilter filter)

 
 /// Create a LinalgStrategyTilePass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyTilePass(StringRef opName, LinalgTilingOptions opt,
                                    LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyTilePass>(opName, opt, filter);
 }
 
+/// Create a LinalgStrategyPadPass.
+std::unique_ptr<OperationPass<FuncOp>>
+mlir::createLinalgStrategyPadPass(StringRef opName, LinalgPaddingOptions opt,
+                                  LinalgTransformationFilter filter) {
+  return std::make_unique<LinalgStrategyPadPass>(opName, opt, filter);
+}
+
 /// Create a LinalgStrategyPromotePass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyPromotePass(StringRef opName,
                                       LinalgPromotionOptions opt,
                                       LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyPromotePass>(opName, opt, filter);
 }
 

mlir/test/Dialect/Linalg/codegen-strategy.mlir

-// Test that both anchor-op name and MatmulOp-based codegen strategy produce the same result.
-// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=2,4,8 vectorize vectorize-contraction-to=matrixintrinsics unroll-vector-transfers=true" | FileCheck %s
-// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=2,4,8 tile-interchange=1,2,0 vectorize vectorize-contraction-to=matrixintrinsics unroll-vector-transfers=true" | FileCheck %s
-// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 promote promote-full-tile-pad register-tile-sizes=2,4,8 vectorize vectorize-contraction-to=outerproduct split-transfers=true unroll-vector-transfers=false" | FileCheck %s --check-prefix=OUTER
-// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=2,4,8 vectorize vectorize-contraction-to=matrixintrinsics unroll-vector-transfers=true" | FileCheck %s
-// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 promote promote-full-tile-pad register-tile-sizes=2,4,8 vectorize vectorize-contraction-to=outerproduct split-transfers=true unroll-vector-transfers=false" | FileCheck %s --check-prefix=OUTER
-// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 generalize iterator-interchange=0,2,1" | FileCheck %s --check-prefix=GENER
-
-
-// CHECK-LABEL: func @matmul(
-// OUTER-LABEL: func @matmul(
-// GENER-LABEL: func @matmul(
-func @matmul(%A: memref<1584x1584xf32>, %B: memref<1584x1584xf32>, %C: memref<1584x1584xf32>) {
-  linalg.matmul
-   ins(%A, %B: memref<1584x1584xf32>, memref<1584x1584xf32>)
-   outs(%C: memref<1584x1584xf32>)
-
-  //      CHECK: vector.matrix_multiply
-  // CHECK-SAME: {lhs_columns = 8 : i32, lhs_rows = 2 : i32, rhs_columns = 4 : i32}
-  // CHECK-SAME: (vector<16xf32>, vector<32xf32>) -> vector<8xf32>
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=2,4,8 vectorize vectorize-contraction-to=matrixintrinsics unroll-vector-transfers=true" -split-input-file | FileCheck %s --check-prefix=CHECK-INTRINSIC
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 promote promote-full-tile-pad register-tile-sizes=2,4,8 vectorize vectorize-contraction-to=outerproduct split-transfers=true unroll-vector-transfers=false" -split-input-file | FileCheck %s --check-prefix=CHECK-OUTER
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 tile-interchange=1,2,0 generalize iterator-interchange=0,2,1" -split-input-file | FileCheck %s --check-prefix=CHECK-INTERCHANGE
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 pad pack-paddings=1,1,0 hoist-paddings=3,3,0" -split-input-file | FileCheck %s --check-prefix=CHECK-PAD
+
+// CHECK-INTRINSIC: func @matmul(
+//     CHECK-OUTER: func @matmul(
+func @matmul(%arg0: memref<72x72xf32>, %arg1: memref<72x72xf32>, %arg2: memref<72x72xf32>) {
+
+  // Check the matrix intrinsic lowering is triggered.
+  //      CHECK-INTRINSIC: vector.matrix_multiply
+  // CHECK-INTRINSIC-SAME: {lhs_columns = 8 : i32, lhs_rows = 2 : i32, rhs_columns = 4 : i32}
+  // CHECK-INTRINSIC-SAME: (vector<16xf32>, vector<32xf32>) -> vector<8xf32>
+
+  // Check the outer product lowering is triggered.
+  //          CHECK-OUTER: vector.outerproduct {{.*}} : vector<2xf32>, vector<4xf32>
+  linalg.matmul ins(%arg0, %arg1: memref<72x72xf32>, memref<72x72xf32>) outs(%arg2: memref<72x72xf32>)
+  return
+}
 
-  //      OUTER: vector.outerproduct {{.*}} : vector<2xf32>, vector<4xf32>
+// -----
 
-  //      GENER: linalg.generic
-  // GENER-SAME: iterator_types = ["parallel", "reduction", "parallel"]
-  return
+// CHECK-INTERCHANGE: func @matmul(
+func @matmul(%arg0: tensor<72x72xf32>, %arg1: tensor<72x72xf32>, %arg2: tensor<72x72xf32>) -> tensor<72x72xf32> {
+  //  CHECK-INTERCHANGE-DAG: %[[C16:.*]] = arith.constant 16
+  //  CHECK-INTERCHANGE-DAG: %[[C32:.*]] = arith.constant 32
+  //  CHECK-INTERCHANGE-DAG: %[[C64:.*]] = arith.constant 64
+
+  // Check the tile loops are interchanged.
+  //      CHECK-INTERCHANGE: scf.for {{.*}} step %[[C32]]
+  //      CHECK-INTERCHANGE:   scf.for {{.*}} step %[[C64]]
+  //      CHECK-INTERCHANGE:    scf.for {{.*}} step %[[C16]]
+
+  // Check the operation has been generalized and interchanged.
+  //      CHECK-INTERCHANGE:      linalg.generic
+  // CHECK-INTERCHANGE-SAME:      iterator_types = ["parallel", "reduction", "parallel"]
+  %0 = linalg.matmul ins(%arg0, %arg1: tensor<72x72xf32>, tensor<72x72xf32>) outs(%arg2: tensor<72x72xf32>) -> tensor<72x72xf32>
+  return %0 : tensor<72x72xf32>
+}
+
+// -----
+
+//         CHECK-PAD: func @matmul(
+func @matmul(%arg0: tensor<72x72xf32>, %arg1: tensor<72x72xf32>, %arg2: tensor<72x72xf32>) -> tensor<72x72xf32> {
+
+  // Check the padding of the input operands has been hoisted out of the tile loop nest.
+  //      CHECK-PAD-COUNT=2: linalg.pad_tensor %{{.*}} nofold
+  //      CHECK-PAD-COUNT=3: scf.for
+  //              CHECK-PAD: linalg.matmul
+  %0 = linalg.matmul ins(%arg0, %arg1: tensor<72x72xf32>, tensor<72x72xf32>) outs(%arg2: tensor<72x72xf32>) -> tensor<72x72xf32>
+  return %0 : tensor<72x72xf32>
 }
 

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

struct TestLinalgCodegenStrategy
 
   template <typename LinalgNamedOp>
   void applyStrategyToNamedLinalgOp();
 
   void runOnFunction() override;
 
   void runStrategy(LinalgTilingOptions tilingOptions,
                    LinalgTilingOptions registerTilingOptions,
+                   LinalgPaddingOptions paddingOptions,
                    vector::VectorContractLowering vectorContractLowering,
                    vector::VectorTransferSplit vectorTransferSplit);
 
   ListOption<int64_t> tileSizes{*this, "tile-sizes",
                                 llvm::cl::MiscFlags::CommaSeparated,
                                 llvm::cl::desc("Specifies the tile sizes.")};
   ListOption<unsigned> tileInterchange{
       *this, "tile-interchange", llvm::cl::MiscFlags::CommaSeparated,
Option<bool> registerPromote{
       *this, "register-promote",
       llvm::cl::desc(
           "Promote the register tile into a small aligned memory buffer."),
       llvm::cl::init(false)};
   Option<bool> registerPromoteFullTile{
       *this, "register-promote-full-tile-pad",
       llvm::cl::desc("Pad the small aligned memory buffer to the tile sizes."),
       llvm::cl::init(false)};
+  Option<bool> pad{*this, "pad", llvm::cl::desc("Pad the operands."),
+                   llvm::cl::init(false)};
+  ListOption<int64_t> packPaddings{
+      *this, "pack-paddings",
+      llvm::cl::desc("Operand packing flags when test-pad-pattern"),
+      llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
+  ListOption<int64_t> hoistPaddings{
+      *this, "hoist-paddings",
+      llvm::cl::desc("Operand hoisting depths when test-pad-pattern"),
+      llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
   Option<bool> generalize{*this, "generalize",
                           llvm::cl::desc("Generalize named operations."),
                           llvm::cl::init(false)};
   ListOption<int64_t> iteratorInterchange{
       *this, "iterator-interchange", llvm::cl::MiscFlags::CommaSeparated,
       llvm::cl::desc("Specifies the iterator interchange.")};
   Option<bool> vectorize{
       *this, "vectorize",
struct TestLinalgCodegenStrategy
   Option<std::string> anchorFuncOpName{
       *this, "anchor-func",
       llvm::cl::desc(
           "Which single func op is the anchor for the codegen strategy to "
           "latch on."),
       llvm::cl::init("")};
 };
 
+// For now, just assume it is the zero of type.
+// In the future, it should be the zero of type + op.
+static Value getNeutralOfLinalgOp(OpBuilder &b, OpOperand &op) {
+  auto t = getElementTypeOrSelf(op.get());
+  return b.create<arith::ConstantOp>(op.getOwner()->getLoc(), t,
+                                     b.getZeroAttr(t));
+}
+
 void TestLinalgCodegenStrategy::runStrategy(
     LinalgTilingOptions tilingOptions,
     LinalgTilingOptions registerTilingOptions,
+    LinalgPaddingOptions paddingOptions,
     vector::VectorContractLowering vectorContractLowering,
     vector::VectorTransferSplit vectorTransferSplit) {
   assert(!anchorOpName.empty());
   CodegenStrategy strategy;
   StringRef genericOpName = GenericOp::getOperationName();
   strategy.tileIf(!tileSizes.empty(), anchorOpName, tilingOptions)
       .promoteIf(promote, anchorOpName,
                  LinalgPromotionOptions()
                      .setAlignment(16)
                      .setUseFullTileBuffersByDefault(promoteFullTile))
       .tileIf(!registerTileSizes.empty(), anchorOpName, registerTilingOptions)
       .promoteIf(registerPromote, anchorOpName,
                  LinalgPromotionOptions()
                      .setAlignment(16)
                      .setUseFullTileBuffersByDefault(registerPromoteFullTile))
+      .padIf(pad, anchorOpName, paddingOptions)
       .generalizeIf(generalize, anchorOpName)
       .interchangeIf(!iteratorInterchange.empty(), iteratorInterchange)
       .vectorizeIf(vectorize, generalize ? genericOpName : anchorOpName)
       .vectorLowering(
           LinalgVectorLoweringOptions()
               .setVectorTransformsOptions(
                   vector::VectorTransformsOptions()
                       .setVectorTransformsOptions(vectorContractLowering)
void TestLinalgCodegenStrategy::runOnFunction() {
   if (!tileInterchange.empty())
     tilingOptions = tilingOptions.setInterchange(tileInterchange);
 
   LinalgTilingOptions registerTilingOptions;
   if (!registerTileSizes.empty())
     registerTilingOptions =
         registerTilingOptions.setTileSizes(registerTileSizes);
 
+  LinalgPaddingOptions paddingOptions;
+  auto packFunc = [&](OpOperand &opOperand) {
+    return opOperand.getOperandNumber() < packPaddings.size()
+               ? packPaddings[opOperand.getOperandNumber()]
+               : false;
+  };
+  auto hoistingFunc = [&](OpOperand &opOperand) {
+    return opOperand.getOperandNumber() < hoistPaddings.size()
+               ? hoistPaddings[opOperand.getOperandNumber()]
+               : 0;
+  };
+  paddingOptions.setPaddingValueComputationFunction(getNeutralOfLinalgOp);
+  paddingOptions.setPaddingNoFoldComputationFunction(packFunc);
+  paddingOptions.setPaddingHoistComputationFunction(hoistingFunc);
+
   vector::VectorContractLowering vectorContractLowering =
       llvm::StringSwitch<vector::VectorContractLowering>(
           vectorizeContractionTo.getValue())
           .Case("matrixintrinsics", vector::VectorContractLowering::Matmul)
           .Case("dot", vector::VectorContractLowering::Dot)
           .Case("outerproduct", vector::VectorContractLowering::OuterProduct)
           .Default(vector::VectorContractLowering::OuterProduct);
   vector::VectorTransferSplit vectorTransferSplit =
       llvm::StringSwitch<vector::VectorTransferSplit>(
           splitVectorTransfersTo.getValue())
           .Case("none", vector::VectorTransferSplit::None)
           .Case("linalg-copy", vector::VectorTransferSplit::LinalgCopy)
           .Case("vector-transfers", vector::VectorTransferSplit::VectorTransfer)
           .Default(vector::VectorTransferSplit::None);
 
-  runStrategy(tilingOptions, registerTilingOptions, vectorContractLowering,
-              vectorTransferSplit);
+  runStrategy(tilingOptions, registerTilingOptions, paddingOptions,
+              vectorContractLowering, vectorTransferSplit);
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgCodegenStrategy() {
   PassRegistration<TestLinalgCodegenStrategy>();
 }
 } // namespace test
 } // namespace mlir