[Mlir-commits] [mlir] 0ddd043 - [MLIR] Fix op folding to not run pre-replace when not constant folding

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2020-03-20T07:49:49+05:30
New Revision: 0ddd04391d281b1d8679ff95ddfce79a61253250

URL: https://github.com/llvm/llvm-project/commit/0ddd04391d281b1d8679ff95ddfce79a61253250
DIFF: https://github.com/llvm/llvm-project/commit/0ddd04391d281b1d8679ff95ddfce79a61253250.diff

LOG: [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 at polymagelabs.com>

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

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/FoldUtils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp
index 66535ec67897..7d209b2231a2 100644
--- a/mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ b/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();