[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 updated https://github.com/llvm/llvm-project/pull/185445

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

---
 mlir/lib/Reducer/ReductionTreePass.cpp | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/mlir/lib/Reducer/ReductionTreePass.cpp b/mlir/lib/Reducer/ReductionTreePass.cpp
index 83497143d9669..912318540e991 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();
     }

>From 03dc05a7aaa7f4590f522c1e263e8b025c2909d6 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 9 Mar 2026 15:49:22 +0000
Subject: [PATCH 2/3] add test.

---
 mlir/test/mlir-reduce/reduction-tree.mlir | 25 +++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/mlir/test/mlir-reduce/reduction-tree.mlir b/mlir/test/mlir-reduce/reduction-tree.mlir
index 2aee89741b42b..88becc45f5626 100644
--- a/mlir/test/mlir-reduce/reduction-tree.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree.mlir
@@ -58,3 +58,28 @@ func.func @simple4(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
 func.func @simple5() {
   return
 }
+
+// -----
+
+// This test checks the ability to remove useless ops interspersed between
+// an 'interesting' op and the return op.
+
+// CHECK: func @simple1() {
+func.func @simple1() {
+  return
+}
+
+// CHECK-LABEL: func @simple2(%arg0: i32, %arg1: i32, %arg2: i32) {
+func.func @simple2(%arg0: i32, %arg1: i32, %arg2: i32) {
+  call @simple1() : () -> ()
+  %0 = "test.op_crash_long" (%arg0, %arg1, %arg2) : (i32, i32, i32) -> i32
+  call @simple5() : ()-> ()
+  // CHECK: %0 = "test.op_crash_short"() : () -> i32
+  // CHECK-NEXT: return
+  return
+}
+
+// CHECK: func @simple5() {
+func.func @simple5() {
+  return
+}

>From 9e4739e8e402ef5954cb9ce629f7c2da62cfa4cc Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Wed, 25 Mar 2026 13:26:06 +0000
Subject: [PATCH 3/3] fix mlir-reduce delete ops logic.

---
 mlir/lib/Reducer/ReductionTreePass.cpp | 74 +++++++++++++++++++++++---
 1 file changed, 68 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Reducer/ReductionTreePass.cpp b/mlir/lib/Reducer/ReductionTreePass.cpp
index 912318540e991..d137a76f2f38a 100644
--- a/mlir/lib/Reducer/ReductionTreePass.cpp
+++ b/mlir/lib/Reducer/ReductionTreePass.cpp
@@ -25,6 +25,9 @@
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/Allocator.h"
+#include "llvm/Support/DebugLog.h"
+
+#define DEBUG_TYPE "reduction-tree"
 
 namespace mlir {
 #define GEN_PASS_DEF_REDUCTIONTREEPASS
@@ -67,19 +70,78 @@ static void applyPatterns(Region &region,
                                       GreedyRewriteStrictness::ExistingOps));
   }
 
-  if (eraseOpNotInRange)
+  if (eraseOpNotInRange) {
+
+    // clang-format off
+    LLVM_DEBUG(
+      LDBG() << "before erase ops not in ranges, keep the ranges:";
+      for (ReductionNode::Range range : rangeToKeep) {
+        LDBG() << "[" << range.first << " " << range.second << ")";
+      } 
+      LDBG() << "region:\n" << region;
+    );
+    // clang-format on
+
+    // The map uses the results of the operations as keys, while the values
+    // represent the remaining user count for each result. We iterate through
+    // 'opsNotInRange' to update this map; if a key's value remains greater than
+    // zero, it indicates that materialization is required for that specific
+    // value.
+    DenseMap<Value, int64_t> valueToMaterializationMap;
+
     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());
+
+      for (Value result : op->getResults())
+        valueToMaterializationMap[result] = result.getNumUses();
+
+      // Use a set to store all operands to prevent the map value from being
+      // decremented multiple times if an operation uses the same operand more
+      // than once.
+      SmallPtrSet<Value, 4> operandSet(op->getOperands().begin(),
+                                       op->getOperands().end());
+      for (Value operand : operandSet)
+        // If an 'operand' is a key in the map, it indicates that the operand
+        // was defined within 'opsNotInRange'.
+        if (valueToMaterializationMap.contains(operand))
+          --valueToMaterializationMap[operand];
+    }
+
+    SmallVector<Type, 4> materializationTypes;
+    SmallVector<Value, 4> valueNeedMaterialization;
+    for (auto mapValue : valueToMaterializationMap) {
+      // If a key in the map has a value greater than zero, it indicates that
+      // there are still operations in the remaining IR using this key.
+      // Therefore, we should materialize it.
+      if (mapValue.second > 0) {
+        materializationTypes.push_back(mapValue.first.getType());
+        valueNeedMaterialization.push_back(mapValue.first);
+      }
+    }
+
+    if (!materializationTypes.empty()) {
+      OpBuilder b(region.getContext());
+      b.setInsertionPointToStart(&region.front());
+      auto castOp = UnrealizedConversionCastOp::create(
+          b, b.getUnknownLoc(), materializationTypes, {});
+      for (auto [src, res] :
+           llvm::zip_equal(valueNeedMaterialization, castOp.getResults())) {
+        src.replaceAllUsesWith(res);
       }
+    }
+
+    // TODO: We remove operations that have no users. However, another issue is
+    // that an op without users might still be an 'interesting' op; therefore,
+    // we must perform an additional 'interestingness' test before deleting it.
+    for (Operation *op : opsNotInRange) {
+      if (op->hasTrait<mlir::OpTrait::IsTerminator>())
+        continue;
       op->dropAllUses();
       op->erase();
     }
+    LDBG() << "after erase ops not in ranges:\n" << region;
+  }
 }
 
 /// We will apply the reducer patterns to the operations in the ranges specified