[MLIR] Introduce applyOpPatternsAndFold for op local rewrites

[bondhugula]

Summary

Introduce mlir::applyOpPatternsAndFold which applies patterns as well as
any folding only on a specified op (in contrast to
applyPatternsAndFoldGreedily which applies patterns only on the regions
of an op isolated from above). The caller is made aware of the op being
folded away or erased.

Depends on D77485.

[rriddle]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

250

This is wrong, you don't handle when the op is erased by a pattern. This really feels like should be its own driver, it doesn't really share anything.

[bondhugula]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

250

Hmm... if the op was erased in a pattern by just erase(), there is no way to know if it was erased? If the op is erased via eraseOp or replaceOp, checking for null at the op's position via worklistMap will address that (we'll have to just keep the op in the worklist before matchAndRewrite). The other issue is that the erased op's operands will also get into the worklist and we don't want them to, but this can be addressed by clearing out the worklist. This could use a separate driver as you say, but it doesn't immediately address the first issue - how do we detect erasure if done via op.erase()?

[rriddle]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

250

how do we detect erasure if done via op.erase()?

We don't have to. It isn't valid to do IR mutations outside of the PatternRewriter.

I thought this was documented somewhere, but https://mlir.llvm.org/docs/QuickstartRewrites/ seems to be extremely out of date. I have an in-progress update to https://mlir.llvm.org/docs/Canonicalization/ I'll add that in as part of it.

[bondhugula]

Implementing this functionality has hit a dead-end; there's no easy way out.

[bondhugula]

Update to use worklist free driver. This isn't still well tested - so could
wait till it could be used at a few places.

[bondhugula]

Fix some gratuitous changes. This isn't still well-tested and could wait for
some of the pending revisions to get in.

[bondhugula]

rebase and permute commit stack

[bondhugula]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

250

I don't recall that requirement ever being documented, although at one point some instances of op.erase() were replaced by rewriter.eraseOp(op). The document itself hasn't changed core content wise since Mar 2019. I couldn't find it in your recent update to the canonicalization doc. Is another one pending?

If that's the constraint on how one can erase, then I think I have it implemented correctly now. It needs better testing though.

[bondhugula]

Rebase.

[rriddle]

Can you add test coverage?

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

250

Yes, there is another one pending.

253

Can you use a comment block to split these?

256

"bottom up" doesn't really make sense here.

265

Please use /// for the top-level comments.

329

while (!erased && changed && ...)

?

[bondhugula]

Address review comments.

[bondhugula]

Thanks for the review.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

253

Not sure I understood it right - but PTAL at what I did.

329

That won't be accurate since although you'll correctly exit the iteration and return, you will be returning !changed, which will be false instead of being true (i.e., the driver has converged).

[rriddle]

Can you add test coverage?

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

329

Ah yeah, duh. Thanks.

[bondhugula]

In D77487#1975842, @rriddle wrote:

Can you add test coverage?

Sure; I'm trying to find a good set of places to call this from. Thanks for the review.

[bondhugula]

Add testing.

[bondhugula]

Drop unrelated changes.

[bondhugula]

Ping @rriddle - do you have more comments here? I added test cases and the output parameter on the method has been made optional. (A test case that exercises the op being erased is in the child revision. )

[nicolasvasilache]

Nice!

I expect this to be very useful for more easily expressing fused passes.

[bondhugula]

Comment updates - reset erased to false in driver's simplifyLocally.

[rriddle]

I have a growing concern that rebuilding the pattern list is going to become extremely costly for compile time. I don't have a concrete suggestion at the moment, but we should avoid rebuilding if possible.

That being said, I think this is a great start to doing local canonicalization. Thank you for pushing this!

mlir/include/mlir/IR/PatternMatch.h

455

I'm of the opinion that we should remove the return value from these methods. It is not used anywhere, and it isn't clear what the user would use it for in any of the existing usages.

mlir/lib/Dialect/Affine/Transforms/AffineDataCopyGeneration.cpp

236 ↗

(On Diff #257602)

This seems extremely costly. Can we create these once and just apply them if necessary?

mlir/lib/Dialect/Affine/Transforms/SimplifyAffineStructures.cpp

95 ↗

(On Diff #257602)

Same here. Can we avoid rebuilding a pattern set for each operation?

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

274

Please use /// here an below.

mlir/lib/Transforms/Utils/LoopUtils.cpp

318 ↗

(On Diff #257602)

nit: Use // here.

mlir/test/lib/Dialect/Affine/TestAffineDataCopy.cpp

100 ↗

(On Diff #257602)

nit: isa<>() || isa<>()

[bondhugula]

NFC - clean up erasure logic

[bondhugula]

clean up and fix erasure further

[bondhugula]

Address review comments

mlir/include/mlir/IR/PatternMatch.h

455

The child revision has the use case where it's essential - let's discuss there if there are alternatives.

mlir/lib/Dialect/Affine/Transforms/SimplifyAffineStructures.cpp

95 ↗

(On Diff #257602)

Sure - done.

mlir/include/mlir/IR/PatternMatch.h

 //===- PatternMatch.h - PatternMatcher classes -------==---------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 

 /// Note: These methods also perform folding and simple dead-code elimination
 ///       before attempting to match any of the provided patterns.
 ///
 bool applyPatternsAndFoldGreedily(Operation *op,
                                   const OwningRewritePatternList &patterns);
 /// Rewrite the given regions, which must be isolated from above.
 bool applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,
                                   const OwningRewritePatternList &patterns);
+
+/// Applies the specified patterns on `op` alone while also trying to fold it,
+/// by select the highest benefits patterns in a greedy manner. Returns true if
+/// no more patterns can be matched. `erased` is set to true if `op` is folded
+/// away or erased as a result of becoming dead.
+/// Note: This does not apply pattern to any regions of `op`.
+bool applyOpPatternsAndFold(Operation *op,
+                            const OwningRewritePatternList &patterns,
+                            bool &erased);
 } // end namespace mlir
 
 #endif // MLIR_PATTERN_MATCH_H

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 //===- GreedyPatternRewriteDriver.cpp - A greedy rewriter -----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //

 class GreedyPatternRewriteDriver : public PatternRewriter {
 public:
   explicit GreedyPatternRewriteDriver(MLIRContext *ctx,
                                       const OwningRewritePatternList &patterns)
       : PatternRewriter(ctx), matcher(patterns), folder(ctx) {
     worklist.reserve(64);
   }
 
-  /// Perform the rewrites while folding and erasing any dead ops. Return true
-  /// if the rewrite converges in `maxIterations`.
   bool simplify(MutableArrayRef<Region> regions, int maxIterations);
 
   void addToWorklist(Operation *op) {
     // Check to see if the worklist already contains this op.
     if (worklistMap.count(op))
       return;
 
     worklistMap[op] = worklist.size();
bool mlir::applyPatternsAndFoldGreedily(
 
   // Start the pattern driver.
   GreedyPatternRewriteDriver driver(regions[0].getContext(), patterns);
   bool converged = driver.simplify(regions, maxPatternMatchIterations);
   LLVM_DEBUG(if (!converged) {
     llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
                  << maxPatternMatchIterations << " times";
   });
   return converged;
 }
+
+namespace {
+/// This is a simple driver for the PatternMatcher to apply patterns and perform
+/// folding on a single op. It repeatedly applies the locally optimal patterns
+/// in a roughly "bottom up" way.
+class OpPatternRewriteDriver : public PatternRewriter {
+public:
+  explicit OpPatternRewriteDriver(MLIRContext *ctx,
+                                  const OwningRewritePatternList &patterns)
+      : PatternRewriter(ctx), matcher(patterns), folder(ctx) {}
+
+  bool simplifyLocally(Operation *op, int maxIterations);
+
+  // This is no additional action to perform other than to just insert the op.
+  Operation *insert(Operation *op) override { return OpBuilder::insert(op); }
+
+  bool isErased() { return erased; }
+
+  // These are hooks implemented for PatternRewriter.
+protected:
+  // If an operation is about to be removed, mark it so that we can let clients
+  // know.
+  void notifyOperationRemoved(Operation *op) override { erased = true; }
+
+  // When a root is going to be replaced, its removal will be notified as well.
+  // So there is nothing to do here.
+  void notifyRootReplaced(Operation *op) override {}
+
+private:
+  /// The low-level pattern matcher.
+  RewritePatternMatcher matcher;
+
+  /// Non-pattern based folder for operations.
+  OperationFolder folder;
+
+  /// True if the operation has been removed.
+  bool erased = false;
+};
+
+} // anonymous namespace
+
+/// Performs the rewrites and folding only on `op`. The simplification stops
+/// if the op is erased as a result of being folded, replaced, or dead. Returns
+/// true if the rewrite converges in `maxIterations`.
+bool OpPatternRewriteDriver::simplifyLocally(Operation *op, int maxIterations) {
+  bool changed = false;
+  int i = 0;
+  // Iterate until convergence or until maxIterations. Deletion of the op as
+  // a result of being dead or folded is convergence.
+  do {
+    // If the operation is trivially dead - remove it.
+    if (isOpTriviallyDead(op)) {
+      op->erase();
+      erased = true;
+      return true;
+    }
+
+    // Try to fold this op.
+    bool inPlaceUpdate;
+    if (succeeded(folder.tryToFold(op, /*processGeneratedConstants=*/nullptr,
+                                   /*preReplaceAction=*/nullptr,
+                                   &inPlaceUpdate))) {
+      changed = true;
+      if (!inPlaceUpdate) {
+        erased = true;
+        return true;
+      }
+    }
+
+    // Make sure that any new operations are inserted at this point.
+    setInsertionPoint(op);
+
+    // Try to match one of the patterns. The rewriter is automatically
+    // notified of any necessary changes, so there is nothing else to do here.
+    changed |= matcher.matchAndRewrite(op, *this);
+    if (erased)
+      return true;
+  } while (changed && ++i < maxIterations);
+
+  // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
+  return !changed;
+}
+
+/// Rewrites only `op` using the supplied canonicalization patterns and
+/// folding. `erased` is set to true if the op is erased as a result of being
+/// folded, replaced, or dead.
+bool mlir::applyOpPatternsAndFold(Operation *op,
+                                  const OwningRewritePatternList &patterns,
+                                  bool &erased) {
+  // Start the pattern driver.
+  OpPatternRewriteDriver driver(op->getContext(), patterns);
+  bool converged = driver.simplifyLocally(op, maxPatternMatchIterations);
+  LLVM_DEBUG(if (!converged) {
+    llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
+                 << maxPatternMatchIterations << " times";
+  });
+  erased = driver.isErased();
+  return converged;
+}