[Mlir-commits] [mlir] [mlir][reducer] Introduce the materialization mechanism in the reduction-tree and fix the logic for deleting operations (PR #185445)
lonely eagle
llvmlistbot at llvm.org
Mon May 4 06:07:14 PDT 2026
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/185445
>From 8cd2423bdad47f0a94cb8a3ff6d0f0ce07982f2c Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 6 Apr 2026 05:56:10 +0000
Subject: [PATCH 1/2] rebase main.
---
mlir/lib/Reducer/ReductionTreePass.cpp | 87 ++++++++++++++++++++---
mlir/test/mlir-reduce/reduction-tree.mlir | 31 ++++++++
2 files changed, 108 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Reducer/ReductionTreePass.cpp b/mlir/lib/Reducer/ReductionTreePass.cpp
index 12358f7d71688..25d41c8b815af 100644
--- a/mlir/lib/Reducer/ReductionTreePass.cpp
+++ b/mlir/lib/Reducer/ReductionTreePass.cpp
@@ -59,10 +59,83 @@ 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.
+ 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;
+
+ 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);
+ }
+ }
+
+ for (Operation *op : opsNotInRange) {
+ if (op->hasTrait<mlir::OpTrait::IsTerminator>())
+ continue;
+ op->dropAllUses();
+ op->erase();
+ }
+ 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. It is essential to further eliminate redundant operations here;
+ // otherwise, the reduction will fail if the size of the deleted ops is
+ // smaller than the newly introduced `unrealized_conversion_cast`.
for (Operation *op : opsInRange) {
// `applyOpPatternsGreedily` with folding returns whether the op is
// converted. Omit it because we don't have expectation this reduction will
@@ -71,12 +144,6 @@ static void applyPatterns(Region ®ion,
GreedyRewriteConfig().setStrictness(
GreedyRewriteStrictness::ExistingOps));
}
-
- if (eraseOpNotInRange)
- for (Operation *op : opsNotInRange) {
- op->dropAllUses();
- op->erase();
- }
}
/// 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 b053a111e9a16..3775e9107cca7 100644
--- a/mlir/test/mlir-reduce/reduction-tree.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree.mlir
@@ -123,3 +123,34 @@ func.func @switch_reduction(%arg0: i32, %arg1: memref<2xf32>, %arg2: memref<2xf3
return
}
// CHECK-NEXT: "test.op_crash"(%[[ARG1]], %[[ARG2]])
+
+// -----
+
+// 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
+
+// -----
+
+// In this case, when the add operation was replaced by an unrealized_conversion_cast,
+// the file size actually increased, leading to a failure in materialization.
+
+// 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
+
>From 1653c6164df416dcbab6d4f71419ade0ff61925c Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 4 May 2026 13:06:11 +0000
Subject: [PATCH 2/2] add reduce.barrier op.
---
mlir/include/mlir/Reducer/CMakeLists.txt | 1 +
mlir/include/mlir/Reducer/IR/CMakeLists.txt | 2 ++
mlir/include/mlir/Reducer/IR/ReducerOps.h | 12 ++++++++++
mlir/include/mlir/Reducer/IR/ReducerOps.td | 25 +++++++++++++++++++++
mlir/include/mlir/Reducer/Passes.td | 3 +++
mlir/include/mlir/Reducer/ReductionNode.h | 4 ++++
mlir/lib/Reducer/CMakeLists.txt | 1 +
mlir/lib/Reducer/IR/CMakeLists.txt | 7 ++++++
mlir/lib/Reducer/IR/ReducerOps.cpp | 15 +++++++++++++
mlir/lib/Reducer/ReductionNode.cpp | 7 ++++++
mlir/lib/Reducer/ReductionTreePass.cpp | 14 ++++++++----
mlir/lib/RegisterAllDialects.cpp | 2 ++
mlir/test/mlir-reduce/failure-test.sh | 2 +-
mlir/test/mlir-reduce/reduction-tree.mlir | 14 ++++++------
14 files changed, 97 insertions(+), 12 deletions(-)
create mode 100644 mlir/include/mlir/Reducer/IR/CMakeLists.txt
create mode 100644 mlir/include/mlir/Reducer/IR/ReducerOps.h
create mode 100644 mlir/include/mlir/Reducer/IR/ReducerOps.td
create mode 100644 mlir/lib/Reducer/IR/CMakeLists.txt
create mode 100644 mlir/lib/Reducer/IR/ReducerOps.cpp
diff --git a/mlir/include/mlir/Reducer/CMakeLists.txt b/mlir/include/mlir/Reducer/CMakeLists.txt
index 37a19a481adc4..2221309cc6026 100644
--- a/mlir/include/mlir/Reducer/CMakeLists.txt
+++ b/mlir/include/mlir/Reducer/CMakeLists.txt
@@ -1,3 +1,4 @@
+add_subdirectory(IR)
set(LLVM_TARGET_DEFINITIONS Passes.td)
mlir_tablegen(Passes.h.inc -gen-pass-decls -name Reducer)
add_mlir_generic_tablegen_target(MLIRReducerIncGen)
diff --git a/mlir/include/mlir/Reducer/IR/CMakeLists.txt b/mlir/include/mlir/Reducer/IR/CMakeLists.txt
new file mode 100644
index 0000000000000..a55d0e9d638fb
--- /dev/null
+++ b/mlir/include/mlir/Reducer/IR/CMakeLists.txt
@@ -0,0 +1,2 @@
+add_mlir_dialect(ReducerOps reducer)
+add_mlir_doc(ReducerOps ReducerOps Dialects/ -gen-op-doc -dialect=reducer)
diff --git a/mlir/include/mlir/Reducer/IR/ReducerOps.h b/mlir/include/mlir/Reducer/IR/ReducerOps.h
new file mode 100644
index 0000000000000..6423d3fb0cef8
--- /dev/null
+++ b/mlir/include/mlir/Reducer/IR/ReducerOps.h
@@ -0,0 +1,12 @@
+#ifndef ReducerOps_H
+#define ReducerOps_H
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpImplementation.h"
+
+#include "mlir/Reducer/IR/ReducerOpsDialect.h.inc"
+#define GET_OP_CLASSES
+#include "mlir/Reducer/IR/ReducerOps.h.inc"
+
+#endif // ReducerOps_H
diff --git a/mlir/include/mlir/Reducer/IR/ReducerOps.td b/mlir/include/mlir/Reducer/IR/ReducerOps.td
new file mode 100644
index 0000000000000..f89b27e8f0e79
--- /dev/null
+++ b/mlir/include/mlir/Reducer/IR/ReducerOps.td
@@ -0,0 +1,25 @@
+#ifndef REDUCER_OPS
+#define REDUCER_OPS
+
+include "mlir/IR/DialectBase.td"
+include "mlir/IR/Traits.td"
+include "mlir/IR/OpBase.td"
+
+def Reducer_Dialect : Dialect {
+ let name = "reducer";
+ let cppNamespace = "::mlir::reducer";
+}
+
+// Base class for Reducer dialect ops.
+class Reducer_Op<string mnemonic, list<Trait> traits = []>
+ : Op<Reducer_Dialect, mnemonic, traits>;
+
+def BarrierOp : Reducer_Op<"barrier", []> {
+ let arguments = (ins Variadic<AnyType>:$inputs);
+ let results = (outs Variadic<AnyType>:$outputs);
+ let assemblyFormat = [{
+ ($inputs^ `:` type($inputs))? `to` type($outputs) attr-dict
+ }];
+}
+
+#endif // REDUCER_OPS
\ No newline at end of file
diff --git a/mlir/include/mlir/Reducer/Passes.td b/mlir/include/mlir/Reducer/Passes.td
index cce5c7570d4d9..53047bfeef568 100644
--- a/mlir/include/mlir/Reducer/Passes.td
+++ b/mlir/include/mlir/Reducer/Passes.td
@@ -14,6 +14,7 @@
#define MLIR_REDUCER_PASSES
include "mlir/Pass/PassBase.td"
+include "mlir/Reducer/IR/ReducerOps.td"
def CommonReductionPassOptions {
list<Option> options = [
@@ -32,6 +33,8 @@ def ReductionTreePass : Pass<"reduction-tree"> {
/* default */"0",
"The graph traversal mode, the default is single-path mode">,
] # CommonReductionPassOptions.options;
+
+ let dependentDialects = ["reducer::ReducerDialect"];
}
def OptReductionPass : Pass<"opt-reduction-pass"> {
diff --git a/mlir/include/mlir/Reducer/ReductionNode.h b/mlir/include/mlir/Reducer/ReductionNode.h
index 125a7c6f6f5e7..6dc2e655b4d65 100644
--- a/mlir/include/mlir/Reducer/ReductionNode.h
+++ b/mlir/include/mlir/Reducer/ReductionNode.h
@@ -65,6 +65,8 @@ class ReductionNode {
/// Return the size of the module.
size_t getSize() const { return size; }
+ size_t getNumOperands();
+
/// Returns true if the module exhibits the interesting behavior.
Tester::Interestingness isInteresting() const { return interesting; }
@@ -157,6 +159,8 @@ class ReductionNode {
/// constraints. This is only valid while the interestingness has been tested.
size_t size = 0;
+ size_t numOperand = -1;
+
/// This is true if the module has been evaluated and it exhibits the
/// interesting behavior.
Tester::Interestingness interesting = Tester::Interestingness::Untested;
diff --git a/mlir/lib/Reducer/CMakeLists.txt b/mlir/lib/Reducer/CMakeLists.txt
index b18a4bca04fcb..a81b85cca72fd 100644
--- a/mlir/lib/Reducer/CMakeLists.txt
+++ b/mlir/lib/Reducer/CMakeLists.txt
@@ -1,3 +1,4 @@
+add_subdirectory(IR)
add_mlir_library(MLIRReduce
OptReductionPass.cpp
ReductionNode.cpp
diff --git a/mlir/lib/Reducer/IR/CMakeLists.txt b/mlir/lib/Reducer/IR/CMakeLists.txt
new file mode 100644
index 0000000000000..925511e8ad376
--- /dev/null
+++ b/mlir/lib/Reducer/IR/CMakeLists.txt
@@ -0,0 +1,7 @@
+add_mlir_dialect_library(MLIRReducerDialect
+ ReducerOps.cpp
+
+ LINK_LIBS PUBLIC
+ MLIRDialectUtils
+ MLIRIR
+ )
diff --git a/mlir/lib/Reducer/IR/ReducerOps.cpp b/mlir/lib/Reducer/IR/ReducerOps.cpp
new file mode 100644
index 0000000000000..27cf022acc3c9
--- /dev/null
+++ b/mlir/lib/Reducer/IR/ReducerOps.cpp
@@ -0,0 +1,15 @@
+#include "mlir/Reducer/IR/ReducerOps.h"
+
+using namespace mlir;
+using namespace reducer;
+
+#include "mlir/Reducer/IR/ReducerOpsDialect.cpp.inc"
+#define GET_OP_CLASSES
+#include "mlir/Reducer/IR/ReducerOps.cpp.inc"
+
+void ReducerDialect::initialize() {
+ addOperations<
+#define GET_OP_LIST
+#include "mlir/Reducer/IR/ReducerOps.cpp.inc"
+ >();
+}
diff --git a/mlir/lib/Reducer/ReductionNode.cpp b/mlir/lib/Reducer/ReductionNode.cpp
index 897aae0becf33..ea5f66f9dd278 100644
--- a/mlir/lib/Reducer/ReductionNode.cpp
+++ b/mlir/lib/Reducer/ReductionNode.cpp
@@ -162,3 +162,10 @@ ReductionNode::iterator<SinglePath>::getNeighbors(ReductionNode *node) {
return node->generateNewVariants();
}
+
+size_t ReductionNode::getNumOperands() {
+ size_t num = 0;
+ for (auto &op : region->getOps())
+ num += op.getNumOperands();
+ return num;
+}
diff --git a/mlir/lib/Reducer/ReductionTreePass.cpp b/mlir/lib/Reducer/ReductionTreePass.cpp
index 25d41c8b815af..9743091293fb6 100644
--- a/mlir/lib/Reducer/ReductionTreePass.cpp
+++ b/mlir/lib/Reducer/ReductionTreePass.cpp
@@ -18,6 +18,7 @@
#include "mlir/IR/DialectInterface.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Reducer/IR/ReducerOps.h"
#include "mlir/Reducer/Passes.h"
#include "mlir/Reducer/ReductionNode.h"
#include "mlir/Reducer/ReductionPatternInterface.h"
@@ -112,10 +113,10 @@ static void applyPatterns(Region ®ion,
if (!materializationTypes.empty()) {
OpBuilder b(region.getContext());
b.setInsertionPointToStart(®ion.front());
- auto castOp = UnrealizedConversionCastOp::create(
- b, b.getUnknownLoc(), materializationTypes, {});
+ auto barrierOp = reducer::BarrierOp::create(b, b.getUnknownLoc(),
+ materializationTypes, {});
for (auto [src, res] :
- llvm::zip_equal(valueNeedMaterialization, castOp.getResults())) {
+ llvm::zip_equal(valueNeedMaterialization, barrierOp.getResults())) {
src.replaceAllUsesWith(res);
}
}
@@ -185,10 +186,15 @@ static LogicalResult findOptimal(ModuleOp module, Region ®ion,
applyPatterns(curRegion, patterns, currentNode.getRanges(),
eraseOpNotInRange);
+ LDBG() << currentNode.getRegion() << "\n";
currentNode.update(test.isInteresting(currentNode.getModule()));
+ LDBG() << currentNode.getNumOperands() << " "
+ << smallestNode->getNumOperands();
+ LDBG() << currentNode.getRegion() << "\n";
if (currentNode.isInteresting() == Tester::Interestingness::True &&
- currentNode.getSize() < smallestNode->getSize())
+ (currentNode.getSize() < smallestNode->getSize() ||
+ currentNode.getNumOperands() < smallestNode->getNumOperands()))
smallestNode = ¤tNode;
++iter;
diff --git a/mlir/lib/RegisterAllDialects.cpp b/mlir/lib/RegisterAllDialects.cpp
index ea5698f39c0b0..9acef6429cfb0 100644
--- a/mlir/lib/RegisterAllDialects.cpp
+++ b/mlir/lib/RegisterAllDialects.cpp
@@ -99,6 +99,7 @@
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/IR/Dialect.h"
#include "mlir/Interfaces/CastInterfaces.h"
+#include "mlir/Reducer/IR/ReducerOps.h"
#include "mlir/Target/LLVM/NVVM/Target.h"
#include "mlir/Target/LLVM/ROCDL/Target.h"
#include "mlir/Target/LLVM/XeVM/Target.h"
@@ -138,6 +139,7 @@ void mlir::registerAllDialects(DialectRegistry ®istry) {
pdl_interp::PDLInterpDialect,
ptr::PtrDialect,
quant::QuantDialect,
+ reducer::ReducerDialect,
ROCDL::ROCDLDialect,
scf::SCFDialect,
shape::ShapeDialect,
diff --git a/mlir/test/mlir-reduce/failure-test.sh b/mlir/test/mlir-reduce/failure-test.sh
index db6a4720d2d56..00b49e9270761 100755
--- a/mlir/test/mlir-reduce/failure-test.sh
+++ b/mlir/test/mlir-reduce/failure-test.sh
@@ -6,7 +6,7 @@ stdout_file=$(mktemp /tmp/stdout.XXXXXX)
stderr_file=$(mktemp /tmp/stderr.XXXXXX)
# Tests for the keyword "failure" in the stderr of the optimization pass
-mlir-opt $1 -test-mlir-reducer > $stdout_file 2> $stderr_file
+mlir-opt $1 -test-mlir-reducer -allow-unregistered-dialect > $stdout_file 2> $stderr_file
if [ $? -ne 0 ] && grep 'failure' $stderr_file; then
exit 1
diff --git a/mlir/test/mlir-reduce/reduction-tree.mlir b/mlir/test/mlir-reduce/reduction-tree.mlir
index 3775e9107cca7..a9ec7f99fafce 100644
--- a/mlir/test/mlir-reduce/reduction-tree.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree.mlir
@@ -131,12 +131,11 @@ func.func @switch_reduction(%arg0: i32, %arg1: memref<2xf32>, %arg2: memref<2xf3
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
+ return %1 : i32
}
-// CHECK-NEXT: %[[CAST:.*]] = builtin.unrealized_conversion_cast to i32
-// CHECK-NEXT: %{{.*}} = "test.op_crash_short"() : () -> i32
-// CHECK-NEXT: return %[[CAST]] : i32
+// CHECK-NEXT: %[[BARRIER_0:.*]] = reducer.barrier to i32
+// CHECK-NEXT: %[[VAL_0:.*]] = "test.op_crash_short"() : () -> i32
+// CHECK-NEXT: return %[[BARRIER_0]] : i32
// -----
@@ -145,10 +144,11 @@ func.func @materialization(%arg0: i32) -> (i32) {
// CHECK-LABEL: func @no_materialization
// CHECK-SAME: %[[ARG0:.*]]: i32
-func.func @no_materialization(%arg0: i32) -> (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
- return %1 : i32
+ %2 = arith.addi %1, %1 : i32
+ return %2 : i32
}
// CHECK-NEXT: %[[CRASH:.*]] = "test.op_crash_short"() : () -> i32
// CHECK-NEXT: %[[ADDI:.*]] = arith.addi %[[CRASH]], %[[CRASH]] : i32
More information about the Mlir-commits
mailing list