[Mlir-commits] [mlir] [mlir][reducer] Make mlir-reduce don't delete terminator and use ub.poison to repalce op's results (PR #185445)

[lonely eagle]

https://github.com/linuxlonelyeagle created https://github.com/llvm/llvm-project/pull/185445

To ensure the IR remains valid after deleting ops, some ops can be deleted, but doing so might make the IR invalid, which prevents further reduction, make mlir-reduce don't delete terminator and use ub.poison to repalce op's results.

>From b2e9e32a9d57bd7d47c96eeeec8ad25b1ffcafaa Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 9 Mar 2026 15:35:59 +0000
Subject: [PATCH] make mlir-reduce don't delete terminator and use ub.poison to
 repalce op's results.

---
 mlir/lib/Reducer/ReductionTreePass.cpp |  9 +++++++++
 mlir/test/mlir-reduce/simple-test.mlir | 14 +++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Reducer/ReductionTreePass.cpp b/mlir/lib/Reducer/ReductionTreePass.cpp
index 1e00ed645f71e..bede177e9e58e 100644
--- a/mlir/lib/Reducer/ReductionTreePass.cpp
+++ b/mlir/lib/Reducer/ReductionTreePass.cpp
@@ -14,6 +14,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/Reducer/Passes.h"
 #include "mlir/Reducer/ReductionNode.h"
@@ -68,6 +69,14 @@ static void applyPatterns(Region &region,
 
   if (eraseOpNotInRange)
     for (Operation *op : opsNotInRange) {
+      if (op->hasTrait<mlir::OpTrait::IsTerminator>())
+        continue;
+      OpBuilder b(op);
+      for (auto result : op->getResults()) {
+        OpBuilder b(op);
+        auto p = ub::PoisonOp::create(b, op->getLoc(), result.getType());
+        result.replaceAllUsesWith(p.getResult());
+      }
       op->dropAllUses();
       op->erase();
     }
diff --git a/mlir/test/mlir-reduce/simple-test.mlir b/mlir/test/mlir-reduce/simple-test.mlir
index 1cc414946a592..7efa29c227166 100644
--- a/mlir/test/mlir-reduce/simple-test.mlir
+++ b/mlir/test/mlir-reduce/simple-test.mlir
@@ -1,6 +1,10 @@
 // UNSUPPORTED: system-windows
-// RUN: mlir-reduce %s -reduction-tree='traversal-mode=0 test=%S/test.sh'
+// RUN: mlir-reduce %s -reduction-tree='traversal-mode=0 test=%S/test.sh' | FileCheck %s
 
+// CHECK-LABEL: func @simple1(
+//  CHECK-SAME:   %[[ARG0:.*]]: i1,
+//  CHECK-SAME:   %[[ARG1:.*]]: memref<2xf32>,
+//  CHECK-SAME:   %[[ARG2:.*]]: memref<2xf32>)
 func.func @simple1(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   cf.cond_br %arg0, ^bb1, ^bb2
 ^bb1:
@@ -11,3 +15,11 @@ func.func @simple1(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
 ^bb3(%1: memref<2xf32>):
   return
 }
+// CHECK:   cf.cond_br %[[ARG0]], ^bb1, ^bb2
+// CHECK: ^bb1:
+// CHECK:   cf.br ^bb3(%[[ARG1]] : memref<2xf32>)
+// CHECK: ^bb2:
+// CHECK:   %[[POISON_0:.*]] = ub.poison : memref<2xf32>
+// CHECK:   cf.br ^bb3(%[[POISON_0]] : memref<2xf32>)
+// CHECK: ^bb3(%[[VAL_0:.*]]: memref<2xf32>):
+// CHECK:   return