[Mlir-commits] [mlir] [MLIR] Fix remove-dead-values leaving stray poison after removing scf.if results (PR #216664)

[Aayush Shrivastava]

https://github.com/iamaayushrivastava created https://github.com/llvm/llvm-project/pull/216664

Fixes #216646

>From e25cb12800c348ebd5ce5503dca32dead3733f45 Mon Sep 17 00:00:00 2001
From: iamaayushrivastava <iamaayushrivastava at gmail.com>
Date: Mon, 17 Aug 2026 14:40:14 +0530
Subject: [PATCH] [MLIR] Fix remove-dead-values leaving stray poison after
 removing scf.if results

---
 mlir/lib/Transforms/RemoveDeadValues.cpp     | 11 +++++++
 mlir/test/Transforms/remove-dead-values.mlir | 30 ++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 6e55bc390be23..acc17f842a6b8 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -46,6 +46,7 @@
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/DebugLog.h"
@@ -835,4 +836,14 @@ void RemoveDeadValues::runOnOperation() {
     module->emitError("greedy pattern rewrite failed to converge");
     signalPassFailure();
   }
+
+  // Canonicalizing a region branch op (e.g. dropping a now-unused result and
+  // the corresponding terminator operand) erases operands in place without
+  // revisiting the defining ops of the values that were erased. This can
+  // strand now-dead values that this pass itself introduced, such as the
+  // `ub.poison` ops created above to replace operands forwarded only to dead
+  // successor inputs. Sweep those away here.
+  IRRewriter rewriter(context);
+  for (Region &region : module->getRegions())
+    eliminateTriviallyDeadOps(rewriter, region);
 }
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 390a448060b7f..99756d7983f37 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -392,6 +392,36 @@ func.func @clean_region_branch_op_remove_result(%arg0 : index, %arg1 : memref<i3
 
 // -----
 
+// The scf.if result is dead, but both branches have side-effecting ops, so
+// scf.if itself must be preserved. The dead result and the yielded values
+// that only fed that dead result (and any ub.poison value that replaces
+// them) must be removed by the `canonicalize` pass.
+//
+// CHECK-CANONICALIZE:       func.func @clean_region_branch_op_if_remove_result(%[[cond:.*]]: i1, %[[m:.*]]: memref<i32>) {
+// CHECK-CANONICALIZE-NEXT:    scf.if %[[cond]] {
+// CHECK-CANONICALIZE-NEXT:      %[[c1:.*]] = arith.constant 1
+// CHECK-CANONICALIZE-NEXT:      memref.store %[[c1]], %[[m]][]
+// CHECK-CANONICALIZE-NEXT:    } else {
+// CHECK-CANONICALIZE-NEXT:      %[[c2:.*]] = arith.constant 2
+// CHECK-CANONICALIZE-NEXT:      memref.store %[[c2]], %[[m]][]
+// CHECK-CANONICALIZE-NEXT:    }
+// CHECK-CANONICALIZE-NEXT:    return
+// CHECK-CANONICALIZE-NEXT:  }
+func.func @clean_region_branch_op_if_remove_result(%cond: i1, %m: memref<i32>) {
+  %r = scf.if %cond -> (i32) {
+    %c1 = arith.constant 1 : i32
+    memref.store %c1, %m[] : memref<i32>
+    scf.yield %c1 : i32
+  } else {
+    %c2 = arith.constant 2 : i32
+    memref.store %c2, %m[] : memref<i32>
+    scf.yield %c2 : i32
+  }
+  return
+}
+
+// -----
+
 // The simple ops which don't have memory effects or live results get removed.
 // %arg5 doesn't get removed from the @main even though it isn't live because
 // the signature of a public function is always left untouched.