From f178c13fa89960c7247a6367269919acf87fd1b3 Mon Sep 17 00:00:00 2001
From: Andrew Young <youngar17@gmail.com>
Date: Thu, 18 Mar 2021 20:06:02 -0700
Subject: [PATCH] [mlir] Support use-def cycles in graph regions during
 regionDCE

When deleting operations in DCE, the algorithm uses a post-order walk of
the IR to ensure that value uses were erased before value defs. Graph
regions do not have the same structural invariants as SSA CFG, and this
post order walk could delete value defs before uses.  This problem is
guaranteed to occur when there is a cycle in the use-def graph.

This change stops DCE from visiting the operations and blocks in any
meaningful order.  Instead, we rely on explicitly dropping all uses of a
value before deleting it.

Reviewed By: mehdi_amini, rriddle

Differential Revision: https://reviews.llvm.org/D98919
---
 mlir/lib/Transforms/Utils/RegionUtils.cpp  | 15 ++++++---------
 mlir/test/Transforms/canonicalize-dce.mlir | 17 +++++++++++++++++
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp
index 7dd064ef0341..21d0ff53fdc8 100644
--- a/mlir/lib/Transforms/Utils/RegionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/RegionUtils.cpp
@@ -312,21 +312,18 @@ static LogicalResult deleteDeadness(MutableArrayRef<Region> regions,
     if (region.empty())
       continue;
 
-    // We do the deletion in an order that deletes all uses before deleting
-    // defs.
-    // MLIR's SSA structural invariants guarantee that except for block
-    // arguments, the use-def graph is acyclic, so this is possible with a
-    // single walk of ops and then a final pass to clean up block arguments.
-    //
-    // To do this, we visit ops in an order that visits domtree children
-    // before domtree parents. A CFG post-order (with reverse iteration with a
-    // block) satisfies that without needing an explicit domtree calculation.
+    // Delete every operation that is not live. Graph regions may have cycles
+    // in the use-def graph, so we must explicitly dropAllUses() from each
+    // operation as we erase it. Visiting the operations in post-order
+    // guarantees that in SSA CFG regions value uses are removed before defs,
+    // which makes dropAllUses() a no-op.
     for (Block *block : llvm::post_order(&region.front())) {
       eraseTerminatorSuccessorOperands(block->getTerminator(), liveMap);
       for (Operation &childOp :
            llvm::make_early_inc_range(llvm::reverse(block->getOperations()))) {
         if (!liveMap.wasProvenLive(&childOp)) {
           erasedAnything = true;
+          childOp.dropAllUses();
           childOp.erase();
         } else {
           erasedAnything |=
diff --git a/mlir/test/Transforms/canonicalize-dce.mlir b/mlir/test/Transforms/canonicalize-dce.mlir
index 4a351dca426e..e96bd65d389a 100644
--- a/mlir/test/Transforms/canonicalize-dce.mlir
+++ b/mlir/test/Transforms/canonicalize-dce.mlir
@@ -156,3 +156,20 @@ func @f(
   "foo.print"(%t4) : (tensor<4xf32>) -> ()
   return
 }
+
+// -----
+
+// Test case: Test values with use-def cycles are deleted properly.
+
+// CHECK:      func @f()
+// CHECK-NEXT:   test.graph_region
+// CHECK-NEXT:     "test.terminator"() : () -> ()
+
+func @f() {
+  test.graph_region {
+    %0 = "math.exp"(%1) : (f32) -> f32
+    %1 = "math.exp"(%0) : (f32) -> f32
+    "test.terminator"() : ()->()
+  }
+  return
+}
-- 
GitLab