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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Aayush Shrivastava (iamaayushrivastava)

<details>
<summary>Changes</summary>

Fixes #<!-- -->216646

---
Full diff: https://github.com/llvm/llvm-project/pull/216664.diff


2 Files Affected:

- (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+11) 
- (modified) mlir/test/Transforms/remove-dead-values.mlir (+30) 


``````````diff
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.

``````````

</details>


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