[mlir][linalg] Add generalization to CodegenStrategy.

[gysit]

Summary

Add a generalization pass and integrate it with CodegenStrategy.

This patch depends on https://reviews.llvm.org/D110728.

[gysit]

Rebase.

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

 
 /// Create a LinalgStrategyPromotePass.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgStrategyPromotePass(
     StringRef opName = "",
     linalg::LinalgPromotionOptions opt = linalg::LinalgPromotionOptions(),
     linalg::LinalgTransformationFilter filter =
         linalg::LinalgTransformationFilter());
 
+/// Create a LinalgStrategyGeneralizePass.
+std::unique_ptr<OperationPass<FuncOp>>
+createLinalgStrategyGeneralizePass(StringRef opName = "",
+                                   linalg::LinalgTransformationFilter filter =
+                                       linalg::LinalgTransformationFilter());
+
 /// Create a LinalgStrategyVectorizePass.
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgStrategyVectorizePass(StringRef opName = "",
                                   linalg::LinalgVectorizationOptions opt =
                                       linalg::LinalgVectorizationOptions(),
                                   linalg::LinalgTransformationFilter filter =
                                       linalg::LinalgTransformationFilter());
 

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

def LinalgStrategyPromotePass
   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 LinalgStrategyGeneralizePass
+    : FunctionPass<"linalg-strategy-generalize-pass"> {
+  let summary = "Configurable pass to apply pattern-based generalization.";
+  let constructor = "mlir::createLinalgStrategyGeneralizePass()";
+  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 LinalgStrategyVectorizePass
     : FunctionPass<"linalg-strategy-vectorize-pass"> {
   let summary = "Configurable pass to apply pattern-based linalg vectorization.";
   let constructor = "mlir::createLinalgStrategyVectorizePass()";
   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(createLinalgStrategyPromotePass(opName, options, m));
   }
 
 private:
   std::string opName;
   linalg::LinalgPromotionOptions options;
 };
 
+/// Represent one application of createLinalgStrategyGeneralizePass.
+struct Generalize : public Transformation {
+  explicit Generalize(StringRef name,
+                      LinalgTransformationFilter::FilterFunction f = nullptr)
+      : Transformation(f), opName(name) {}
+
+  void addToPassPipeline(OpPassManager &pm,
+                         LinalgTransformationFilter m) const override {
+    pm.addPass(createLinalgStrategyGeneralizePass(opName, m));
+  }
+
+private:
+  std::string opName;
+};
+
 /// Represent one application of createLinalgStrategyVectorizePass.
 struct Vectorize : public Transformation {
   explicit Vectorize(linalg::LinalgVectorizationOptions options,
                      LinalgTransformationFilter::FilterFunction f = nullptr)
       : Transformation(f), opName(), options(options) {}
 
   Vectorize(StringRef name, linalg::LinalgVectorizationOptions options,
             LinalgTransformationFilter::FilterFunction f = nullptr)
struct CodegenStrategy {
   /// Conditionally append a pattern to add a level of promotion for
   /// `LinalgOpType` with promotion `options`.
   CodegenStrategy &
   promoteIf(bool b, StringRef opName, linalg::LinalgPromotionOptions options,
             LinalgTransformationFilter::FilterFunction f = nullptr) {
     return b ? promote(opName, options, f) : *this;
     return *this;
   }
+  /// Append a pattern to generalize named operations.
+  CodegenStrategy &
+  generalize(StringRef opName,
+             LinalgTransformationFilter::FilterFunction f = nullptr) {
+    transformationSequence.emplace_back(
+        std::make_unique<Generalize>(opName, f));
+    return *this;
+  }
+  /// Conditionally append a pattern to generalize named operations.
+  CodegenStrategy &
+  generalizeIf(bool b, StringRef opName,
+               LinalgTransformationFilter::FilterFunction f = nullptr) {
+    return b ? generalize(opName, f) : *this;
+    return *this;
+  }
   /// Append a pattern to rewrite `LinalgOpType` as a vector operation.
   CodegenStrategy &
   vectorize(StringRef opName,
             LinalgTransformationFilter::FilterFunction f = nullptr) {
     assert(!opName.empty() && "expected an op name");
     transformationSequence.emplace_back(std::make_unique<Vectorize>(
         opName, linalg::LinalgVectorizationOptions(), 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 pattern-based linalg generalization.
+struct LinalgStrategyGeneralizePass
+    : public LinalgStrategyGeneralizePassBase<LinalgStrategyGeneralizePass> {
+
+  LinalgStrategyGeneralizePass() = default;
+
+  LinalgStrategyGeneralizePass(StringRef opName,
+                               LinalgTransformationFilter filter)
+      : filter(filter) {
+    this->anchorOpName.setValue(opName.str());
+  }
+
+  void runOnFunction() override {
+    auto funcOp = getFunction();
+    if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
+      return;
+
+    RewritePatternSet generalizationPattern(funcOp.getContext());
+    if (!anchorOpName.empty()) {
+      generalizationPattern.add<LinalgGeneralizationPattern>(
+          anchorOpName, funcOp.getContext(), filter);
+    } else {
+      generalizationPattern.add<LinalgGeneralizationPattern>(
+          funcOp.getContext(), filter);
+    }
+    if (failed(applyPatternsAndFoldGreedily(funcOp,
+                                            std::move(generalizationPattern))))
+      signalPassFailure();
+  }
+
+  LinalgTransformationFilter filter;
+};
+
 /// Configurable pass to apply pattern-based linalg promotion.
 struct LinalgStrategyPromotePass
     : public LinalgStrategyPromotePassBase<LinalgStrategyPromotePass> {
 
   LinalgStrategyPromotePass() = default;
 
   LinalgStrategyPromotePass(StringRef opName, LinalgPromotionOptions opt,
                             LinalgTransformationFilter filt)

 /// Create a LinalgStrategyPromotePass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyPromotePass(StringRef opName,
                                       LinalgPromotionOptions opt,
                                       LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyPromotePass>(opName, opt, filter);
 }
 
+/// Create a LinalgStrategyGeneralizePass.
+std::unique_ptr<OperationPass<FuncOp>>
+mlir::createLinalgStrategyGeneralizePass(StringRef opName,
+                                         LinalgTransformationFilter filter) {
+  return std::make_unique<LinalgStrategyGeneralizePass>(opName, filter);
+}
+
 /// Create a LinalgStrategyVectorizePass.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgStrategyVectorizePass(StringRef opName,
                                         LinalgVectorizationOptions opt,
                                         LinalgTransformationFilter filter) {
   return std::make_unique<LinalgStrategyVectorizePass>(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" | 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>
 
   // OUTER: vector.outerproduct {{.*}} : vector<2xf32>, vector<4xf32>
+  // GENER: linalg.generic
   return
 }
 

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

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> generalize{*this, "generalize",
+                          llvm::cl::desc("Generalize named operations."),
+                          llvm::cl::init(false)};
   Option<bool> vectorize{
       *this, "vectorize",
       llvm::cl::desc("Rewrite the linalg op as a vector operation."),
       llvm::cl::init(false)};
   Option<std::string> splitVectorTransfersTo{
       *this, "split-transfers",
       llvm::cl::desc(
           "Split vector transfers between slow (masked) and fast "

 
 void TestLinalgCodegenStrategy::runStrategy(
     LinalgTilingOptions tilingOptions,
     LinalgTilingOptions registerTilingOptions,
     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))
-      .vectorizeIf(vectorize, anchorOpName)
+      .generalizeIf(generalize, anchorOpName)
+      .vectorizeIf(vectorize, generalize ? genericOpName : anchorOpName)
       .setEnableVectorTransferPartialRewrite(true)
       .setEnableVectorContractLowering(true)
       .setEnableVectorToSCFConversion(true)
       .setVectorTransformsOptions(
           vector::VectorTransformsOptions()
               .setVectorTransformsOptions(vectorContractLowering)
               .setVectorTransferSplit(vectorTransferSplit))
       .setVectorTransferToSCFOptions(