[flang-commits] [flang] [mlir] [mlir][CSE] Add pruneDeadOps to CSE pass (PR #191394)

[Mehdi Amini via flang-commits]

================
@@ -388,7 +390,45 @@ void CSEDriver::simplifyRegion(ScopedMapTy &knownValues, Region &region) {
   }
 }
 
+void CSEDriver::pruneDeadOps(Operation *op) {
+  SmallVector<Operation *> deadOps;
+  op->walk([&](Operation *op) {
+    if (isOpTriviallyDead(op))
+      deadOps.push_back(op);
+  });
+
+  // Use an erased set to avoid double erasure. An operand may be used
+  // multiple times by same dead ops, causing the same defining op
+  // to be added to the worklist more than once. The erased set guards
+  // against processing the same op twice.
+  DenseSet<Operation *> erased;
+  for (Operation *op : deadOps) {
+    SmallVector<Operation *> worklist;
+    worklist.push_back(op);
+    while (!worklist.empty()) {
+      Operation *op = worklist.pop_back_val();
+      if (erased.contains(op))
+        continue;
+      if (!isOpTriviallyDead(op))
+        continue;
+      for (Value arg : op->getOperands())
+        if (Operation *argOp = arg.getDefiningOp())
+          worklist.push_back(argOp);
+      erased.insert(op);
+      rewriter.eraseOp(op);
+      ++numDCE;
+    }
+  }
+}
+
 void CSEDriver::simplify(Operation *op, bool *changed) {
+  /// Eagerly erase trivially dead operations to prevent them from interfering
+  /// with the CSE pass. Retaining dead operations in a region can affect the
+  /// equivalence check between two region ops, causing CSE to miss
+  /// optimization opportunities, and may also trigger unnecessary calls to
+  /// simplifyRegion on dead region ops.
+  pruneDeadOps(op);
----------------
joker-eph wrote:

Seems like you're doing this once ahead of CSE but there aren't any direct interaction with CSE during the algorithm. So it more looks to me like running canonicalization before CSE should solve these cases.

https://github.com/llvm/llvm-project/pull/191394