[MLIR] Fix op folding to not run pre-replace when not constant folding

OperationFolder::tryToFold was running the pre-replacement
action even when there was no constant folding, i.e., when the operation
was just being updated in place but was not going to be replaced. This
led to nested ops being unnecessarily removed from the worklist and only
being processed in the next outer iteration of the greedy pattern
rewriter, which is also why this didn't affect the final output IR but
only the convergence rate. It also led to an op's results' users to be
unnecessarily added to the worklist.

Signed-off-by: Uday Bondhugula <uday@polymagelabs.com>

Differential Revision: https://reviews.llvm.org/D76268

mlir/lib/Transforms/Utils/FoldUtils.cpp

@@ -85,17 +85,17 @@ LogicalResult OperationFolder::tryToFold(
   if (failed(tryToFold(op, results, processGeneratedConstants)))
     return failure();
 
-  // Constant folding succeeded. We will start replacing this op's uses and
-  // eventually erase this op. Invoke the callback provided by the caller to
-  // perform any pre-replacement action.
-  if (preReplaceAction)
-    preReplaceAction(op);
-
   // Check to see if the operation was just updated in place.
   if (results.empty())
     return success();
 
-  // Otherwise, replace all of the result values and erase the operation.
+  // 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.
   for (unsigned i = 0, e = results.size(); i != e; ++i)
     op->getResult(i).replaceAllUsesWith(results[i]);
   op->erase();