[Mlir-commits] [mlir] f73f875 - [mlir][CSE] Delete dead code immediately in CSE pass (#190926)

[llvmlistbot at llvm.org]

Author: lonely eagle
Date: 2026-04-08T21:18:18+08:00
New Revision: f73f8754cabd4caf3c9989d8f542d2b9a9f85d07

URL: https://github.com/llvm/llvm-project/commit/f73f8754cabd4caf3c9989d8f542d2b9a9f85d07
DIFF: https://github.com/llvm/llvm-project/commit/f73f8754cabd4caf3c9989d8f542d2b9a9f85d07.diff

LOG: [mlir][CSE] Delete dead code immediately in CSE pass (#190926)

This PR addresses an issue where dead code within a region prevents
essentially identical region ops from being CSE'd. Currently, CSE only
removes dead ops at the end of the pass. This PR fixes the issue by
deleting dead code immediately upon encounter.

Added: 
    

Modified: 
    mlir/lib/Transforms/CSE.cpp
    mlir/test/Transforms/cse.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/CSE.cpp b/mlir/lib/Transforms/CSE.cpp
index 4d25e5e7c92b6..f6f6ccb75e5d8 100644
--- a/mlir/lib/Transforms/CSE.cpp
+++ b/mlir/lib/Transforms/CSE.cpp
@@ -258,9 +258,11 @@ LogicalResult CSEDriver::simplifyOperation(ScopedMapTy &knownValues,
   if (op->hasTrait<OpTrait::IsTerminator>())
     return failure();
 
-  // If the operation is already trivially dead just add it to the erase list.
+  // If the operation is already trivially dead, erase it immediately.
+  // Retaining dead operations in a region can affect the equivalence check
+  // between two region ops, causing CSE to miss optimization opportunities.
   if (isOpTriviallyDead(op)) {
-    opsToErase.push_back(op);
+    rewriter.eraseOp(op);
     ++numDCE;
     return success();
   }
@@ -308,7 +310,7 @@ LogicalResult CSEDriver::simplifyOperation(ScopedMapTy &knownValues,
 
 void CSEDriver::simplifyBlock(ScopedMapTy &knownValues, Block *bb,
                               bool hasSSADominance) {
-  for (auto &op : *bb) {
+  for (auto &op : llvm::make_early_inc_range(*bb)) {
     // Most operations don't have regions, so fast path that case.
     if (op.getNumRegions() != 0) {
       // If this operation is isolated above, we can't process nested regions
@@ -397,8 +399,10 @@ void CSEDriver::simplify(Operation *op, bool *changed) {
   /// Erase any operations that were marked as dead during simplification.
   for (auto *op : opsToErase)
     rewriter.eraseOp(op);
-  if (changed)
-    *changed = !opsToErase.empty();
+  if (changed) {
+    assert(opsToErase.empty() || numDCE || numCSE);
+    *changed = numDCE || numCSE;
+  }
 
   // Note: CSE does currently not remove ops with regions, so DominanceInfo
   // does not have to be invalidated.

diff  --git a/mlir/test/Transforms/cse.mlir b/mlir/test/Transforms/cse.mlir
index 4b2907287d89e..9c2af700259c7 100644
--- a/mlir/test/Transforms/cse.mlir
+++ b/mlir/test/Transforms/cse.mlir
@@ -511,6 +511,38 @@ func.func @cse_multiple_regions(%c: i1, %t: tensor<5xf32>) -> (tensor<5xf32>, te
 
 // -----
 
+// This test is used to verify whether CSE can correctly
+// handle region operations that contain dead ops.
+
+func.func @cse_multiple_regions_with_dead_op(%c: i1, %t: tensor<5xf32>) -> (tensor<5xf32>, tensor<5xf32>) {
+  %r1 = scf.if %c -> (tensor<5xf32>) {
+    %0 = tensor.empty() : tensor<5xf32>
+    %1 = arith.constant 1: index
+    scf.yield %0 : tensor<5xf32>
+  } else {
+    scf.yield %t : tensor<5xf32>
+  }
+  %r2 = scf.if %c -> (tensor<5xf32>) {
+    %0 = tensor.empty() : tensor<5xf32>
+    scf.yield %0 : tensor<5xf32>
+  } else {
+    scf.yield %t : tensor<5xf32>
+  }
+  return %r1, %r2 : tensor<5xf32>, tensor<5xf32>
+}
+
+// CHECK-LABEL: func @cse_multiple_regions_with_dead_op
+//       CHECK:   %[[if:.*]] = scf.if {{.*}} {
+//       CHECK:     tensor.empty
+//       CHECK:     scf.yield
+//       CHECK:   } else {
+//       CHECK:     scf.yield
+//       CHECK:   }
+//   CHECK-NOT:   scf.if
+//       CHECK:   return %[[if]], %[[if]]
+
+// -----
+
 // CHECK-LABEL: @cse_recursive_effects_success
 func.func @cse_recursive_effects_success() -> (i32, i32, i32) {
   // CHECK-NEXT: %[[READ_VALUE:.*]] = "test.op_with_memread"() : () -> i32