[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
llvmlistbot at llvm.org
Thu Mar 26 01:09:57 PDT 2026
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/4] 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 ®ion,
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/4] 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/4] 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 ®ion,
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(®ion.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
>From 99f5e70753a07e4c9ec2be78a93fe9005dc71835 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Thu, 26 Mar 2026 08:09:36 +0000
Subject: [PATCH 4/4] fix reduce delete op logic.
---
mlir/lib/Reducer/ReductionTreePass.cpp | 37 +++++++++++------------
mlir/test/mlir-reduce/reduction-tree.mlir | 27 +++++++++++++++++
2 files changed, 44 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Reducer/ReductionTreePass.cpp b/mlir/lib/Reducer/ReductionTreePass.cpp
index d137a76f2f38a..de0ff2339f678 100644
--- a/mlir/lib/Reducer/ReductionTreePass.cpp
+++ b/mlir/lib/Reducer/ReductionTreePass.cpp
@@ -14,7 +14,6 @@
//
//===----------------------------------------------------------------------===//
-#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/IR/DialectInterface.h"
#include "mlir/Reducer/Passes.h"
#include "mlir/Reducer/ReductionNode.h"
@@ -57,19 +56,6 @@ static void applyPatterns(Region ®ion,
opsInRange.push_back(&op.value());
}
- // `applyOpPatternsGreedily` with folding may erase the ops so we can't do the
- // pattern matching in above iteration. Besides, erase op not-in-range may end
- // up in invalid module, so `applyOpPatternsGreedily` with folding should come
- // before that transform.
- for (Operation *op : opsInRange) {
- // `applyOpPatternsGreedily` with folding returns whether the op is
- // converted. Omit it because we don't have expectation this reduction will
- // be success or not.
- (void)applyOpPatternsGreedily(op, patterns,
- GreedyRewriteConfig().setStrictness(
- GreedyRewriteStrictness::ExistingOps));
- }
-
if (eraseOpNotInRange) {
// clang-format off
@@ -84,7 +70,7 @@ static void applyPatterns(Region ®ion,
// 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
+ // `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;
@@ -102,8 +88,8 @@ static void applyPatterns(Region ®ion,
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 an `operand` is a key in the map, it indicates that the operand
+ // was defined within `opsNotInRange`.
if (valueToMaterializationMap.contains(operand))
--valueToMaterializationMap[operand];
}
@@ -131,9 +117,6 @@ static void applyPatterns(Region ®ion,
}
}
- // 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;
@@ -142,6 +125,20 @@ static void applyPatterns(Region ®ion,
}
LDBG() << "after erase ops not in ranges:\n" << region;
}
+
+ // After removing `opsNotInRange`, we apply applyOpPatternsGreedily both to
+ // run specific patterns and to eliminate operations that have no users. The
+ // reason we do not directly delete all userless operations is that some may
+ // be `interesting` ops. Therefore, we utilize `applyOpPatternsGreedily` here
+ // instead.
+ for (Operation *op : opsInRange) {
+ // `applyOpPatternsGreedily` with folding returns whether the op is
+ // converted. Omit it because we don't have expectation this reduction will
+ // be success or not.
+ (void)applyOpPatternsGreedily(op, patterns,
+ GreedyRewriteConfig().setStrictness(
+ GreedyRewriteStrictness::ExistingOps));
+ }
}
/// We will apply the reducer patterns to the operations in the ranges specified
diff --git a/mlir/test/mlir-reduce/reduction-tree.mlir b/mlir/test/mlir-reduce/reduction-tree.mlir
index 88becc45f5626..539cee2ba3584 100644
--- a/mlir/test/mlir-reduce/reduction-tree.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree.mlir
@@ -83,3 +83,30 @@ func.func @simple2(%arg0: i32, %arg1: i32, %arg2: i32) {
func.func @simple5() {
return
}
+
+// -----
+
+// CHECK-LABEL: func @materialization
+// CHECK-SAME: %[[ARG0:.*]]: i32
+func.func @materialization(%arg0: i32) -> (i32) {
+ %0 = "test.op_crash_long" (%arg0, %arg0, %arg0) : (i32, i32, i32) -> i32
+ %1 = arith.addi %0, %0 : i32
+ %2 = arith.addi %1, %1 : i32
+ return %2 : i32
+}
+// CHECK-NEXT: %[[CAST:.*]] = builtin.unrealized_conversion_cast to i32
+// CHECK-NEXT: %{{.*}} = "test.op_crash_short"() : () -> i32
+// CHECK-NEXT: return %[[CAST]] : i32
+
+// -----
+
+// CHECK-LABEL: func @no_materialization
+// CHECK-SAME: %[[ARG0:.*]]: i32
+func.func @no_materialization(%arg0: i32) -> (i32) {
+ %0 = "test.op_crash_long" (%arg0, %arg0, %arg0) : (i32, i32, i32) -> i32
+ %1 = arith.addi %0, %0 : i32
+ return %1 : i32
+}
+// CHECK-NEXT: %[[CRASH:.*]] = "test.op_crash_short"() : () -> i32
+// CHECK-NEXT: %[[ADDI:.*]] = arith.addi %[[CRASH]], %[[CRASH]] : i32
+// CHECK-NEXT: return %[[ADDI]] : i32
More information about the Mlir-commits
mailing list