[MLIR] Handle in-place folding properly in greedy pattern rewrite driver

[bondhugula]

Summary

OperatioFolder::tryToFold performs both true folding and in a few
instances in-place updates through op rewrites. In the latter case, we
should still be applying the supplied pattern rewrites in the same
iteration; however this wasn't the case since tryToFold returned
success() for both true folding and in-place updates, and the patterns
for the in-place updated ops were being applied only in the next
iteration of the driver's outer loop. This fix would make it converge
faster.

Depends on D77478.

[bondhugula]

Rebase and reorder stack.

[bondhugula]

Rebase

[rriddle]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

188–190

We could just set this inside of preReplaceAction instead of changing the folder API for now?

[bondhugula]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

188–190

Two quick questions

1. preReplaceAction is also not called when the folding fails. Did you mean to set 'foldedAndErased = true' in preReplaceAction? Leaving inPlaceUpdate to true when preReplaceAction is not called would be incorrect.

2. Also, wouldn't a tryToFold client want to know if the op was erased? I had a use case for that but that'll anyway not be needed once the child revision which adds 'op local fold and apply patterns' is available. But otherwise there would be no way to know if the fold erased the op.

[rriddle]

mlir/lib/Transforms/Utils/FoldUtils.cpp

78

Can we just have this as a precondition? That avoids uninitialized variables being passed in.

[bondhugula]

mlir/lib/Transforms/Utils/FoldUtils.cpp

78

Sorry, I didn't get you. Did you mean making inPlaceUpdate a compulsory parameter and asserting inPlaceUpdate (in which case we can instead just pass by reference) or adding assert(!*inPlaceUpdate && 'should be initialized to false');? Thanks.

[rriddle]

mlir/lib/Transforms/Utils/FoldUtils.cpp

78

I meant the second one.

[bondhugula]

mlir/lib/Transforms/Utils/FoldUtils.cpp

78

Hmm... that would be weird to have a part of its logic outside given inPlaceUpdate is an output parameter (and the assertion would still crash if they just sent in an uninitialized pointer).

mlir/include/mlir/Transforms/FoldUtils.h

 public:
   OperationFolder(MLIRContext *ctx) : interfaces(ctx) {}
 
   /// Tries to perform folding on the given `op`, including unifying
   /// deduplicated constants. If successful, replaces `op`'s uses with
   /// folded results, and returns success. `preReplaceAction` is invoked on `op`
   /// before it is replaced. 'processGeneratedConstants' is invoked for any new
   /// operations generated when folding. If the op was completely folded it is
-  /// erased.
+  /// erased. If it is just updated in place, `inPlaceUpdate` is set to true.
   LogicalResult
   tryToFold(Operation *op,
             function_ref<void(Operation *)> processGeneratedConstants = nullptr,
-            function_ref<void(Operation *)> preReplaceAction = nullptr);
+            function_ref<void(Operation *)> preReplaceAction = nullptr,
+            bool *inPlaceUpdate = nullptr);
 
   /// Notifies that the given constant `op` should be remove from this
   /// OperationFolder's internal bookkeeping.
   ///
   /// Note: this method must be called if a constant op is to be deleted
   /// externally to this OperationFolder. `op` must be a constant op.
   void notifyRemoval(Operation *op);
 

mlir/lib/Transforms/Utils/FoldUtils.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // OperationFolder
 //===----------------------------------------------------------------------===//
 
 LogicalResult OperationFolder::tryToFold(
     Operation *op, function_ref<void(Operation *)> processGeneratedConstants,
-    function_ref<void(Operation *)> preReplaceAction) {
+    function_ref<void(Operation *)> preReplaceAction, bool *inPlaceUpdate) {
+  if (inPlaceUpdate)
+    *inPlaceUpdate = false;
+
   // If this is a unique'd constant, return failure as we know that it has
   // already been folded.
   if (referencedDialects.count(op))
     return failure();
 
   // Try to fold the operation.
   SmallVector<Value, 8> results;
   OpBuilder builder(op);
   if (failed(tryToFold(builder, op, results, processGeneratedConstants)))
     return failure();
 
   // Check to see if the operation was just updated in place.
-  if (results.empty())
+  if (results.empty()) {
+    if (inPlaceUpdate)
+      *inPlaceUpdate = true;
     return success();
+  }
 
   // Constant folding succeeded. We will start replacing this op's uses and
   // erase this op. Invoke the callback provided by the caller to perform any
   // pre-replacement action.
   if (preReplaceAction)
     preReplaceAction(op);
 
   // Replace all of the result values and erase the operation.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

for (auto result : op->getResults())
         addToWorklist(user);
   }
 
 private:
   // Look over the provided operands for any defining operations that should
   // be re-added to the worklist. This function should be called when an
   // operation is modified or removed, as it may trigger further
   // simplifications.
-  template <typename Operands> void addToWorklist(Operands &&operands) {
+  template <typename Operands>
+  void addToWorklist(Operands &&operands) {
     for (Value operand : operands) {
       // If the use count of this operand is now < 2, we re-add the defining
       // operation to the worklist.
       // TODO(riverriddle) This is based on the fact that zero use operations
       // may be deleted, and that single use values often have more
       // canonicalization opportunities.
       if (!operand.use_empty() && !operand.hasOneUse())
         continue;
private:
   std::vector<Operation *> worklist;
   DenseMap<Operation *, unsigned> worklistMap;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 };
 } // end anonymous namespace
 
-/// Perform the rewrites while folding and erasing any dead ops.
+/// Performs the rewrites while folding and erasing any dead ops. Returns true
+/// if the rewrite converges in `maxIterations`.
 bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions,
                                           int maxIterations) {
   // Add the given operation to the worklist.
   auto collectOps = [this](Operation *op) { addToWorklist(op); };
 
   bool changed = false;
   int i = 0;
   do {
while (!worklist.empty()) {
         for (auto result : op->getResults())
           for (auto *userOp : result.getUsers())
             addToWorklist(userOp);
 
         notifyOperationRemoved(op);
       };
 
       // Try to fold this op.
-      if (succeeded(folder.tryToFold(op, collectOps, preReplaceAction))) {
+      bool inPlaceUpdate;
+      if ((succeeded(folder.tryToFold(op, collectOps, preReplaceAction,
+                                      &inPlaceUpdate)))) {
         changed = true;
+        if (!inPlaceUpdate)
           continue;
       }
 
       // 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);