[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:

```
func.func @cse_dead_ops(%arg0: i1) {
  %c0_i32 = arith.constant 0 : i32
  %0 = arith.select %arg0, %c0_i32, %c0_i32 : i32
  %1 = scf.if %arg0 -> (i32) {
    %c0_i32_0 = arith.constant 0 : i32
    scf.yield %c0_i32_0 : i32
  } else {
    %c0_i32_0 = arith.constant 0 : i32
    scf.yield %c0_i32_0 : i32
  }
  return
}
```
it run cse, would be
```
func.func @cse_dead_ops(%arg0: i1) {
  %c0_i32 = arith.constant 0 : i32
  return
}
```
then I run canonicalization, it would be 
```
func.func @cse_dead_ops(%arg0: i1) {
  return
}
```
Overall, running canonicalization beforehand ensures that CSE remains highly effective. To be honest, I feel that CSE’s approach to removing dead ops is somewhat lacking, as it doesn't fully cover that functionality—though it doesn't seem to be a major issue.

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