[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

251

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

251

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

251

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

251

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

251

Yes, there is another one pending.

295

Can you use a comment block to split these?

298

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

307

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

371

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

?

[bondhugula]

Address review comments.

[bondhugula]

Thanks for the review.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

295

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

371

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

371

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

316

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);
+  bool simplifyLocally(Operation *op, int maxIterations, bool &erased);
 
   void addToWorklist(Operation *op) {
     // Check to see if the worklist already contains this op.
     if (worklistMap.count(op))
       return;
 
     worklistMap[op] = worklist.size();
     worklist.push_back(op);
if (succeeded(simplifyRegions(regions))) {
       folder.clear();
       changed = true;
     }
   } while (changed && ++i < maxIterations);
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return !changed;
 }
 
+/// Performs the rewrites and folding only on `op`. `erased` is set to true if
+/// the op is folded or erased as a result of being dead. Returns true if the
+/// rewrite converges in `maxIterations`.
+bool GreedyPatternRewriteDriver::simplifyLocally(Operation *op,
+                                                 int maxIterations,
+                                                 bool &erased) {
+  bool changed = false;
+  int i = 0;
+  erased = false;
+  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);
+  } while (changed && ++i < maxIterations);
+
+  // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
+  return !changed;
+}
+
 /// Rewrite the regions of the specified operation, which must be isolated from
 /// above, by repeatedly applying the highest benefit patterns in a greedy
 /// work-list driven manner. Return true if no more patterns can be matched in
 /// the result operation regions.
 /// Note: This does not apply patterns to the top-level operation itself.
 ///
 bool mlir::applyPatternsAndFoldGreedily(
     Operation *op, const OwningRewritePatternList &patterns) {
bool mlir::applyPatternsAndFoldGreedily(
   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;
 }
+
+/// Rewrites `op` alone using the supplied canonicalization patterns and
+/// folding.
+bool mlir::applyOpPatternsAndFold(Operation *op,
+                                  const OwningRewritePatternList &patterns,
+                                  bool &erased) {
+  // Start the pattern driver.
+  GreedyPatternRewriteDriver driver(op->getContext(), patterns);
+  bool converged =
+      driver.simplifyLocally(op, maxPatternMatchIterations, erased);
+  LLVM_DEBUG(if (!converged) {
+    llvm::dbgs() << "The pattern rewrite doesn't converge after scanning "
+                 << maxPatternMatchIterations << " times";
+  });
+  return converged;
+}