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

[Uday Bondhugula via Phabricator via llvm-commits]

bondhugula created this revision.
bondhugula added reviewers: lattner, mehdi_amini, rriddle.
Herald added subscribers: llvm-commits, grosul1, Joonsoo, liufengdb, lucyrfox, mgester, arpith-jacob, nicolasvasilache, antiagainst, shauheen, burmako, jpienaar.
Herald added a project: LLVM.
bondhugula added a child revision: D77486: [MLIR][NFC] fix name operand -> op.

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 D77483 <https://reviews.llvm.org/D77483>.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D77485

Files:
  mlir/include/mlir/Transforms/FoldUtils.h
  mlir/lib/Transforms/Utils/FoldUtils.cpp
  mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp


Index: mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
===================================================================
--- mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -107,7 +107,8 @@
   // 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.
@@ -136,7 +137,8 @@
 };
 } // 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.
@@ -186,9 +188,12 @@
       };
 
       // Try to fold this op.
-      if (succeeded(folder.tryToFold(op, collectOps, preReplaceAction))) {
+      bool inPlaceUpdate;
+      if ((succeeded(folder.tryToFold(op, collectOps, preReplaceAction,
+                                      &inPlaceUpdate)))) {
         changed = true;
-        continue;
+        if (!inPlaceUpdate)
+          continue;
       }
 
       // Make sure that any new operations are inserted at this point.
Index: mlir/lib/Transforms/Utils/FoldUtils.cpp
===================================================================
--- mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ mlir/lib/Transforms/Utils/FoldUtils.cpp
@@ -74,7 +74,10 @@
 
 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))
@@ -87,8 +90,11 @@
     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
Index: mlir/include/mlir/Transforms/FoldUtils.h
===================================================================
--- mlir/include/mlir/Transforms/FoldUtils.h
+++ mlir/include/mlir/Transforms/FoldUtils.h
@@ -56,11 +56,12 @@
   /// 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, `inPlaceUpdated` 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.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D77485.255114.patch
Type: text/x-patch
Size: 3732 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200405/3efc86c7/attachment.bin>