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

[lonely eagle]

================
@@ -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);
----------------
linuxlonelyeagle wrote:

By the way, it would be great if we could remove inserted values from the scoped hash table to maintain its balance. What do you think of this idea?🤨

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