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

[Uday Bondhugula via Phabricator via llvm-commits]

This revision was automatically updated to reflect the committed changes.
Closed by commit rGcbcb12fd44df: [MLIR] Handle in-place folding properly in greedy pattern rewrite driver (authored by bondhugula).

Changed prior to commit:
  https://reviews.llvm.org/D77485?vs=255939&id=256771#toc

Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D77485/new/

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
@@ -104,7 +104,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.
@@ -133,7 +134,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.
@@ -183,9 +185,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, `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.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D77485.256771.patch
Type: text/x-patch
Size: 3731 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200411/64a1e60f/attachment.bin>