[mlir] Add peeling xform to Codegen Strategy

[dcaballe]

Summary

This patch adds the knobs to use peeling in the codegen strategy
infrastructure.

[springerm]

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

975–976

I would add a comment here saying that the order in which nested loops appear in the SmallVector affects the peeling. I.e., whether partial iterations will be peeled or not. (That's why it is a SmallVector and not a DenseSet.)

983

Can this struct go into an anonymous namespace in the .cpp file?

1003

nit: Both fields can be const.

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

290

if (failed(apply....))
  signalPassFailure();

[dcaballe]

Addressed feedback

[dcaballe]

Thanks. I'll commit it in a while.

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

983

I'm not sure but probably better to keep it here for consistency with the other transformations that have a similar ones in this file?

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

 
 /// Create a LinalgStrategyInterchangePass.
 std::unique_ptr<OperationPass<func::FuncOp>>
 createLinalgStrategyInterchangePass(
     ArrayRef<int64_t> iteratorInterchange = {},
     const linalg::LinalgTransformationFilter &filter =
         linalg::LinalgTransformationFilter());
 
+/// Create a LinalgStrategyPeelPass.
+std::unique_ptr<OperationPass<func::FuncOp>> createLinalgStrategyPeelPass(
+    StringRef opName = "",
+    linalg::LinalgPeelOptions opt = linalg::LinalgPeelOptions(),
+    const linalg::LinalgTransformationFilter &filter =
+        linalg::LinalgTransformationFilter());
+
 /// Create a LinalgStrategyVectorizePass.
 std::unique_ptr<OperationPass<func::FuncOp>> createLinalgStrategyVectorizePass(
     StringRef opName = "",
     linalg::LinalgVectorizationOptions opt =
         linalg::LinalgVectorizationOptions(),
     const linalg::LinalgTransformationFilter &filter =
         linalg::LinalgTransformationFilter(),
     bool padVectorize = false);

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

def LinalgStrategyInterchangePass
   let constructor = "mlir::createLinalgStrategyInterchangePass()";
   let dependentDialects = ["linalg::LinalgDialect"];
   let options = [
     Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
       "Which func op is the anchor to latch on.">,
   ];
 }
 
+def LinalgStrategyPeelPass
+    : Pass<"linalg-strategy-peel-pass", "func::FuncOp"> {
+  let summary = "Configurable pass to apply pattern-based linalg peeling.";
+  let constructor = "mlir::createLinalgStrategyPeelPass()";
+  let dependentDialects = [
+    "linalg::LinalgDialect",
+    "scf::SCFDialect"
+  ];
+  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
     : Pass<"linalg-strategy-vectorize-pass", "func::FuncOp"> {
   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

explicit Decompose(LinalgTransformationFilter::FilterFunction f = nullptr)
       : Transformation(std::move(f)) {}
 
   void addToPassPipeline(OpPassManager &pm,
                          LinalgTransformationFilter m) const override {
     pm.addPass(createLinalgStrategyDecomposePass(m));
   }
 };
 
+/// Represent one application of createLinalgStrategyPeelPass.
+struct Peel : public Transformation {
+  explicit Peel(linalg::LinalgPeelOptions options,
+                LinalgTransformationFilter::FilterFunction f = nullptr)
+      : Transformation(std::move(f)), opName(), options(options) {}
+
+  Peel(StringRef name, linalg::LinalgPeelOptions options,
+       LinalgTransformationFilter::FilterFunction f = nullptr)
+      : Transformation(std::move(f)), opName(name), options(options) {}
+
+  void addToPassPipeline(OpPassManager &pm,
+                         LinalgTransformationFilter m) const override {
+    pm.addPass(createLinalgStrategyPeelPass(opName, options, m));
+  }
+
+private:
+  std::string opName;
+  linalg::LinalgPeelOptions options;
+};
+
 /// Represent one application of createLinalgStrategyVectorizePass.
 struct Vectorize : public Transformation {
   explicit Vectorize(linalg::LinalgVectorizationOptions options,
                      LinalgTransformationFilter::FilterFunction f = nullptr,
                      bool padVectorize = false)
       : Transformation(std::move(f)), opName(), options(options),
         vectorizePadding(padVectorize) {}
 
decompose(const LinalgTransformationFilter::FilterFunction &f = nullptr) {
     transformationSequence.emplace_back(std::make_unique<Decompose>(f));
     return *this;
   }
   /// Conditionally append patterns to decompose convolutions.
   CodegenStrategy &
   decomposeIf(bool b, LinalgTransformationFilter::FilterFunction f = nullptr) {
     return b ? decompose(std::move(f)) : *this;
   }
+  /// Append a pattern to peel 'LinalgOpType'.
+  CodegenStrategy &
+  peel(StringRef opName, const LinalgPeelOptions &options,
+       const LinalgTransformationFilter::FilterFunction &f = nullptr) {
+    transformationSequence.emplace_back(
+        std::make_unique<Peel>(opName, options, f));
+    return *this;
+  }
+  /// Conditionally append a pattern to peel 'LinalgOpType'.
+  CodegenStrategy &
+  peelIf(bool b, StringRef opName, const LinalgPeelOptions &options,
+         LinalgTransformationFilter::FilterFunction f = nullptr) {
+    return b ? peel(opName, options, std::move(f)) : *this;
+  }
   /// Append a pattern to rewrite `LinalgOpType` as a vector operation.
   CodegenStrategy &
   vectorize(StringRef opName,
             const LinalgTransformationFilter::FilterFunction &f = nullptr,
             bool vectorizePadding = false) {
     transformationSequence.emplace_back(std::make_unique<Vectorize>(
         opName, linalg::LinalgVectorizationOptions(), f, vectorizePadding));
     return *this;

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

 struct TiledLinalgOp {
   LinalgOp op;
   SmallVector<Operation *, 8> loops;
   SmallVector<Value, 4> tensorResults;
 };
 FailureOr<TiledLinalgOp> tileLinalgOp(RewriterBase &b, LinalgOp op,
                                       const LinalgTilingOptions &options);
 
+/// Peel and canonicalize 'loops'.
+void peelLoops(RewriterBase &rewriter, ArrayRef<scf::ForOp> loops);
+
 /// Peel the loops of a TiledLinalgOp.
 void peelTiledLinalgOp(RewriterBase &rewriter, TiledLinalgOp &res,
                        ArrayRef<int64_t> peeledLoops,
                        LinalgTilingLoopType loopType);
 
 /// Fuse a sequence of linalg operations (`ops`) using tile-and-fuse. This
 /// proceeds as follows:
 /// - Find outer parallel loops in these ops that can be fused.
struct LinalgPromotionPattern : public LinalgBasePromotionPattern {
   LinalgPromotionPattern(
       StringRef opName, MLIRContext *context, LinalgPromotionOptions options,
       LinalgTransformationFilter f = LinalgTransformationFilter(),
       PatternBenefit benefit = 1)
       : LinalgBasePromotionPattern(opName, context, options, f, benefit) {}
 };
 
 ///
+/// Linalg peeling patterns.
+///
+
+/// Compute the loops to peel and return them in a SmallVector. Loops will be
+/// peeled in order of appearance in the SmallVector. This order will impact the
+/// output IR. If an inner-to-outer order is provided, the peeled iterations of
+/// the outer loops will also contain the peeled inner loops. If an
+/// outer-to-inner order is provided, the peeled iterations of the outer loops
+/// will not contain any peeled inner loops.
+using LoopsToPeelComputationFunction = std::function<void(
+    OpBuilder &, Operation *, SmallVectorImpl<scf::ForOp> &)>;
+
+struct LinalgPeelOptions {
+  LoopsToPeelComputationFunction loopsToPeelComputationFunction = nullptr;
+};
+
+/// `filter` controls LinalgTransformMarker matching and update when specified.
+struct LinalgPeelingPattern : public OpInterfaceRewritePattern<LinalgOp> {
+  /// Construct a generic pattern applied to all LinalgOp that verify `filter`.
+  LinalgPeelingPattern(
+      MLIRContext *context,
+      LinalgTransformationFilter f = LinalgTransformationFilter(),
+      LinalgPeelOptions options = LinalgPeelOptions(),
+      PatternBenefit benefit = 1);
+
+  /// Construct a pattern specifically applied to `opName`.
+  LinalgPeelingPattern(
+      StringRef opName, MLIRContext *context,
+      LinalgPeelOptions options = LinalgPeelOptions(),
+      LinalgTransformationFilter f = LinalgTransformationFilter(),
+      PatternBenefit benefit = 1);
+
+  LogicalResult matchAndRewrite(LinalgOp linalgOp,
+                                PatternRewriter &rewriter) const override;
+
+private:
+  /// LinalgTransformMarker handles special attribute manipulations.
+  const LinalgTransformationFilter filter;
+  /// Peeling options.
+  const LinalgPeelOptions options;
+};
+
+///
 /// Linalg vectorization patterns.
 ///
 /// Empty for now, used for SFINAE purposes only.
 struct LinalgVectorizationOptions {};
 
 /// `filter` controls LinalgTransformMarker matching and update when specified.
 /// See `vectorizeLinalgOp` for more details.
 struct LinalgVectorizationPattern : public OpInterfaceRewritePattern<LinalgOp> {

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

void runOnOperation() override {
     }
     (void)applyPatternsAndFoldGreedily(funcOp, std::move(promotionPattern));
   }
 
   LinalgPromotionOptions options;
   LinalgTransformationFilter filter;
 };
 
+/// Configurable pass to apply pattern-based linalg peeling.
+struct LinalgStrategyPeelPass
+    : public LinalgStrategyPeelPassBase<LinalgStrategyPeelPass> {
+
+  LinalgStrategyPeelPass() = default;
+
+  LinalgStrategyPeelPass(StringRef opName, LinalgPeelOptions opt,
+                         LinalgTransformationFilter filt)
+      : options(opt), filter(std::move(filt)) {
+    this->anchorOpName.setValue(opName.str());
+  }
+
+  void runOnOperation() override {
+    auto funcOp = getOperation();
+    if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
+      return;
+
+    RewritePatternSet peelingPatterns(funcOp.getContext());
+    if (!anchorOpName.empty()) {
+      peelingPatterns.add<LinalgPeelingPattern>(
+          anchorOpName, funcOp.getContext(), options, filter);
+    } else {
+      peelingPatterns.add<LinalgPeelingPattern>(funcOp.getContext(), filter,
+                                                options);
+    }
+    if (failed(
+            applyPatternsAndFoldGreedily(funcOp, std::move(peelingPatterns))))
+      return signalPassFailure();
+  }
+
+  LinalgPeelOptions options;
+  LinalgTransformationFilter filter;
+};
+
 /// Configurable pass to apply pattern-based linalg vectorization.
 struct LinalgStrategyVectorizePass
     : public LinalgStrategyVectorizePassBase<LinalgStrategyVectorizePass> {
 
   LinalgStrategyVectorizePass() = default;
 
   LinalgStrategyVectorizePass(StringRef opName, LinalgVectorizationOptions opt,
                               LinalgTransformationFilter filt,

 std::unique_ptr<OperationPass<func::FuncOp>>
 mlir::createLinalgStrategyInterchangePass(
     ArrayRef<int64_t> iteratorInterchange,
     const LinalgTransformationFilter &filter) {
   return std::make_unique<LinalgStrategyInterchangePass>(iteratorInterchange,
                                                          filter);
 }
 
+/// Create a LinalgStrategyPeelPass.
+std::unique_ptr<OperationPass<func::FuncOp>>
+mlir::createLinalgStrategyPeelPass(StringRef opName, LinalgPeelOptions opt,
+                                   const LinalgTransformationFilter &filter) {
+  return std::make_unique<LinalgStrategyPeelPass>(opName, opt, filter);
+}
+
 /// Create a LinalgStrategyVectorizePass.
 std::unique_ptr<OperationPass<func::FuncOp>>
 mlir::createLinalgStrategyVectorizePass(
     StringRef opName, LinalgVectorizationOptions opt,
     const LinalgTransformationFilter &filter, bool padVectorize) {
   return std::make_unique<LinalgStrategyVectorizePass>(opName, opt, filter,
                                                        padVectorize);
 }

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

return llvm::TypeSwitch<Operation *, SmallVector<Value, 4>>(op)
                                                       partialIteration)))
           return partialIteration->getResults();
         assert(!partialIteration && "expected that loop was not peeled");
         return forOp->getResults();
       })
       .Default([&](Operation *op) { return op->getResults(); });
 }
 
+/// Peel and canonicalize 'loops'.
+void mlir::linalg::peelLoops(RewriterBase &rewriter,
+                             ArrayRef<scf::ForOp> loops) {
+  for (auto loopOp : loops) {
+    SmallVector<Value, 4> loopResults;
+    loopResults = peelLoop(rewriter, loopOp);
+  }
+}
+
 /// Peel loops after tiling.
 void mlir::linalg::peelTiledLinalgOp(RewriterBase &rewriter, TiledLinalgOp &res,
                                      ArrayRef<int64_t> peeledLoops,
                                      LinalgTilingLoopType loopType) {
   for (int64_t loop : peeledLoops) {
     assert(loop < static_cast<int64_t>(res.loops.size()) &&
            "requested peeling of non-existing loop");
     SmallVector<Value, 4> loopResults;
if (!promotedOp) {
     rewriter.cancelRootUpdate(op);
     return op->emitError("subview promotion failed");
   }
   rewriter.finalizeRootUpdate(op);
   filter.replaceLinalgTransformationFilter(rewriter, op);
   return success();
 }
 
+mlir::linalg::LinalgPeelingPattern::LinalgPeelingPattern(
+    MLIRContext *context, LinalgTransformationFilter f,
+    LinalgPeelOptions options, PatternBenefit benefit)
+    : OpInterfaceRewritePattern<LinalgOp>(context, benefit),
+      filter(std::move(f)), options(std::move(options)) {}
+
+mlir::linalg::LinalgPeelingPattern::LinalgPeelingPattern(
+    StringRef opName, MLIRContext *context, LinalgPeelOptions options,
+    LinalgTransformationFilter f, PatternBenefit benefit)
+    : OpInterfaceRewritePattern<LinalgOp>(context, benefit),
+      filter(f.addOpNameFilter(opName)), options(std::move(options)) {}
+
+LogicalResult mlir::linalg::LinalgPeelingPattern::matchAndRewrite(
+    LinalgOp linalgOp, PatternRewriter &rewriter) const {
+  if (failed(filter.checkAndNotify(rewriter, linalgOp)))
+    return failure();
+
+  // Increase marker counter even if peeling doesn't happen for this op.
+  filter.replaceLinalgTransformationFilter(rewriter, linalgOp);
+
+  if (!options.loopsToPeelComputationFunction)
+    return failure();
+
+  SmallVector<scf::ForOp, 4> loopsToPeel;
+  options.loopsToPeelComputationFunction(rewriter, linalgOp, loopsToPeel);
+  peelLoops(rewriter, loopsToPeel);
+  return success();
+}
+
 mlir::linalg::LinalgVectorizationPattern::LinalgVectorizationPattern(
     MLIRContext *context, LinalgTransformationFilter f,
     LinalgVectorizationOptions options, PatternBenefit benefit)
     : OpInterfaceRewritePattern<LinalgOp>(context, benefit),
       filter(std::move(f)) {}
 
 mlir::linalg::LinalgVectorizationPattern::LinalgVectorizationPattern(
     StringRef opName, MLIRContext *context, LinalgVectorizationOptions options,