[Mlir-commits] [mlir] Implement Multi-Path traversal mode for mlir-reduce (PR #215628)
Daria Sukhonina
llvmlistbot at llvm.org
Tue Aug 11 10:46:08 PDT 2026
https://github.com/zetanumbers created https://github.com/llvm/llvm-project/pull/215628
Just like suggestion proposed I've implemented Multi-Path traversal mode with a minimal change. The new `multipath-better(-baseline).mlir` tests have been added to showcase a potential improvement to minified mlir code.
This PR is rebased on #214738
>From 43f8303a443a74a656b5b24ce21735dea495a26a Mon Sep 17 00:00:00 2001
From: Daria Sukhonina <dariasukhonina at gmail.com>
Date: Fri, 7 Aug 2026 16:49:06 +0300
Subject: [PATCH 1/2] Fix mlir-reduce splitting smallest range instead of the
largest one
---
mlir/lib/Reducer/ReductionNode.cpp | 2 +-
.../mlir-reduce/reduction-tree/doc-example.mlir | 9 +++++----
.../mlir-reduce/reduction-tree/except-last.mlir | 17 +++++++++++++++++
mlir/test/mlir-reduce/script/except-last.sh | 7 +++++++
4 files changed, 30 insertions(+), 5 deletions(-)
create mode 100644 mlir/test/mlir-reduce/reduction-tree/except-last.mlir
create mode 100755 mlir/test/mlir-reduce/script/except-last.sh
diff --git a/mlir/lib/Reducer/ReductionNode.cpp b/mlir/lib/Reducer/ReductionNode.cpp
index 11aeaf77b4642..f298c12665b83 100644
--- a/mlir/lib/Reducer/ReductionNode.cpp
+++ b/mlir/lib/Reducer/ReductionNode.cpp
@@ -75,7 +75,7 @@ ArrayRef<ReductionNode *> ReductionNode::generateNewVariants() {
// final ranges vector will be {{1, 3}, {4, 6}, {6, 9}}.
auto maxElement =
llvm::max_element(ranges, [](const Range &lhs, const Range &rhs) {
- return (lhs.second - lhs.first) > (rhs.second - rhs.first);
+ return (lhs.second - lhs.first) < (rhs.second - rhs.first);
});
// The length of range is less than 1, we can't split it to create new
diff --git a/mlir/test/mlir-reduce/reduction-tree/doc-example.mlir b/mlir/test/mlir-reduce/reduction-tree/doc-example.mlir
index 7f194025a7ef0..0593db2168afe 100644
--- a/mlir/test/mlir-reduce/reduction-tree/doc-example.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree/doc-example.mlir
@@ -26,7 +26,8 @@ func.func @func2(%arg0: i1) -> f32 {
// CHECK-LABEL: func @func2
// CHECK-SAME: (%arg0: i1) -> f32
// CHECK-DAG: %[[C22:.*]] = arith.constant 2.200000e+00 : f32
-// CHECK-DAG: %[[C75:.*]] = arith.constant 7.500000e+00 : f32
-// CHECK: %[[SEL:.*]] = arith.select %arg0, %[[C75]], %[[C22]] : f32
-// CHECK: %[[ADD:.*]] = arith.addf %[[SEL]], %[[C22]] : f32
-// CHECK: return %[[ADD]] : f32
+// CHECK-DAG: %[[C53:.*]] = arith.constant 5.300000e+00 : f32
+// CHECK: %[[ADD1:.*]] = arith.addf %[[C22]], %[[C53]] : f32
+// CHECK: %[[SEL:.*]] = arith.select %arg0, %[[ADD1]], %[[C22]] : f32
+// CHECK: %[[ADD2:.*]] = arith.addf %[[C22]], %[[SEL]] : f32
+// CHECK: return %[[ADD2]] : f32
diff --git a/mlir/test/mlir-reduce/reduction-tree/except-last.mlir b/mlir/test/mlir-reduce/reduction-tree/except-last.mlir
new file mode 100644
index 0000000000000..fa5c59e33177a
--- /dev/null
+++ b/mlir/test/mlir-reduce/reduction-tree/except-last.mlir
@@ -0,0 +1,17 @@
+// UNSUPPORTED: system-windows
+// RUN: mlir-reduce %s -reduction-tree='traversal-mode=0 test=%S/../script/except-last.sh' | FileCheck %s
+// We are testing the ability of keeping trivially-dead yet interesting code
+
+func.func @except_last() {
+ %0 = arith.constant 1 : i32
+ %1 = arith.constant 2 : i32
+ %2 = arith.constant 3 : i32
+ %3 = arith.constant 4 : i32
+ return
+}
+
+// CHECK-LABEL: func @except_last
+// CHECK-NEXT: arith.constant 1 : i32
+// CHECK-NEXT: arith.constant 2 : i32
+// CHECK-NEXT: arith.constant 3 : i32
+// CHECK-NEXT: return
diff --git a/mlir/test/mlir-reduce/script/except-last.sh b/mlir/test/mlir-reduce/script/except-last.sh
new file mode 100755
index 0000000000000..b6dce52a9403f
--- /dev/null
+++ b/mlir/test/mlir-reduce/script/except-last.sh
@@ -0,0 +1,7 @@
+#!/bin/sh
+
+file=$1
+
+if grep -q "arith.constant 1 : i32" $file && grep -q "arith.constant 2 : i32" $file && grep -q "arith.constant 3 : i32" $file; then
+ exit 1
+fi
>From 6233b721285fc5ae2a0800ea8bbb01e2657109c5 Mon Sep 17 00:00:00 2001
From: Daria Sukhonina <dariasukhonina at gmail.com>
Date: Tue, 11 Aug 2026 20:33:09 +0300
Subject: [PATCH 2/2] Implement MultiPath traversal mode for mlir-reduce
---
mlir/docs/Tools/mlir-reduce.md | 4 +-
mlir/include/mlir/Reducer/ReductionNode.h | 19 +++++++---
mlir/lib/Reducer/ReductionNode.cpp | 37 +++++++++++++++++--
mlir/lib/Reducer/ReductionTreePass.cpp | 3 ++
.../reduction-tree/doc-example.mlir | 1 +
.../mlir-reduce/reduction-tree/invalid.mlir | 1 +
.../multipath-better-baseline.mlir | 19 ++++++++++
.../reduction-tree/multipath-better.mlir | 18 +++++++++
.../reduction-tree/reduction-tree.mlir | 1 +
.../reduction-tree/simple-test.mlir | 1 +
.../reduction-tree/trivially-dead.mlir | 1 +
.../mlir-reduce/script/multipath-better.sh | 11 ++++++
12 files changed, 104 insertions(+), 12 deletions(-)
create mode 100644 mlir/test/mlir-reduce/reduction-tree/multipath-better-baseline.mlir
create mode 100644 mlir/test/mlir-reduce/reduction-tree/multipath-better.mlir
create mode 100755 mlir/test/mlir-reduce/script/multipath-better.sh
diff --git a/mlir/docs/Tools/mlir-reduce.md b/mlir/docs/Tools/mlir-reduce.md
index 9b5c26257f5c7..31cc82fbd6647 100644
--- a/mlir/docs/Tools/mlir-reduce.md
+++ b/mlir/docs/Tools/mlir-reduce.md
@@ -177,8 +177,8 @@ int main(int argc, char **argv) {
`mlir-reduce` is missing several features,
-* `-reduction-tree` now only supports `Single-Path` traversal mode, extends it
-with different traversal strategies may reduce the input better.
+* `-reduction-tree` now only supports `Single-Path` and `Multi-Path` traversal modes,
+extending it with different traversal strategies may reduce the input better.
* Produce the optimal result when interrupted. The reduction process may take
a quite long time, it'll be better to get an optimal result so far while an
interrupt is triggered.
diff --git a/mlir/include/mlir/Reducer/ReductionNode.h b/mlir/include/mlir/Reducer/ReductionNode.h
index 6ca4e13d159ac..4724a3ba97f12 100644
--- a/mlir/include/mlir/Reducer/ReductionNode.h
+++ b/mlir/include/mlir/Reducer/ReductionNode.h
@@ -116,8 +116,7 @@ class ReductionNode {
BaseIterator &operator++() {
ReductionNode *top = visitQueue.front();
visitQueue.pop();
- for (ReductionNode *node : getNeighbors(top))
- visitQueue.push(node);
+ pushNeighbors(top);
return *this;
}
@@ -131,11 +130,10 @@ class ReductionNode {
ReductionNode *operator->() const { return visitQueue.front(); }
protected:
- ArrayRef<ReductionNode *> getNeighbors(ReductionNode *node) {
- return static_cast<T *>(this)->getNeighbors(node);
+ void pushNeighbors(ReductionNode *node) {
+ return static_cast<T *>(this)->pushNeighbors(node);
}
- private:
std::queue<ReductionNode *> visitQueue;
};
@@ -188,7 +186,16 @@ class ReductionNode::iterator<SinglePath>
: public BaseIterator<iterator<SinglePath>> {
friend BaseIterator<iterator<SinglePath>>;
using BaseIterator::BaseIterator;
- ArrayRef<ReductionNode *> getNeighbors(ReductionNode *node);
+ void pushNeighbors(ReductionNode *node);
+};
+
+// Specialized iterator for MultiPath traversal
+template <>
+class ReductionNode::iterator<MultiPath>
+ : public BaseIterator<iterator<MultiPath>> {
+ friend BaseIterator<iterator<MultiPath>>;
+ using BaseIterator::BaseIterator;
+ void pushNeighbors(ReductionNode *node);
};
} // namespace mlir
diff --git a/mlir/lib/Reducer/ReductionNode.cpp b/mlir/lib/Reducer/ReductionNode.cpp
index f298c12665b83..c178bd2edde76 100644
--- a/mlir/lib/Reducer/ReductionNode.cpp
+++ b/mlir/lib/Reducer/ReductionNode.cpp
@@ -115,8 +115,8 @@ void ReductionNode::update(std::pair<Tester::Interestingness, size_t> result) {
}
}
-ArrayRef<ReductionNode *>
-ReductionNode::iterator<SinglePath>::getNeighbors(ReductionNode *node) {
+void
+ReductionNode::iterator<SinglePath>::pushNeighbors(ReductionNode *node) {
// Single Path: Traverses the smallest successful variant at each level until
// no new successful variants can be created at that level.
ArrayRef<ReductionNode *> variantsFromParent =
@@ -129,7 +129,7 @@ ReductionNode::iterator<SinglePath>::getNeighbors(ReductionNode *node) {
if (!llvm::all_of(variantsFromParent, [](ReductionNode *node) {
return node->isInteresting() != Tester::Interestingness::Untested;
})) {
- return {};
+ return;
}
ReductionNode *smallest = nullptr;
@@ -150,5 +150,34 @@ ReductionNode::iterator<SinglePath>::getNeighbors(ReductionNode *node) {
node = node->getParent();
}
- return node->generateNewVariants();
+ for (ReductionNode *newVariant : node->generateNewVariants()) {
+ visitQueue.push(newVariant);
+ }
+}
+
+void
+ReductionNode::iterator<MultiPath>::pushNeighbors(ReductionNode *node) {
+ // MultiPath: Traverses every successful variant at each level until
+ // no new successful variants can be created.
+ ArrayRef<ReductionNode *> variantsFromParent =
+ node->getParent()->getVariants();
+
+ if (node->isInteresting() == Tester::Interestingness::True && node->getSize() < node->getParent()->getSize()) {
+ for (ReductionNode *newVariant : node->generateNewVariants()) {
+ visitQueue.push(newVariant);
+ }
+ }
+
+ // The parent node created several variants and they may be waiting for
+ // examing interestingness. In Single Path approach, we will select the
+ // smallest variant to continue our exploration. Thus we should wait until the
+ // last variant to be examed then do the following traversal decision.
+ if (llvm::all_of(variantsFromParent, [](ReductionNode *node) {
+ return node->isInteresting() != Tester::Interestingness::Untested;
+ })) {
+
+ for (ReductionNode *newVariant : node->getParent()->generateNewVariants()) {
+ visitQueue.push(newVariant);
+ }
+ }
}
diff --git a/mlir/lib/Reducer/ReductionTreePass.cpp b/mlir/lib/Reducer/ReductionTreePass.cpp
index 2244475e268fe..9fb9d39f038d0 100644
--- a/mlir/lib/Reducer/ReductionTreePass.cpp
+++ b/mlir/lib/Reducer/ReductionTreePass.cpp
@@ -299,6 +299,9 @@ LogicalResult ReductionTreePass::reduceOp(ModuleOp module, Region ®ion) {
case TraversalMode::SinglePath:
return findOptimal<ReductionNode::iterator<TraversalMode::SinglePath>>(
module, region, reducerPatterns, tester);
+ case TraversalMode::MultiPath:
+ return findOptimal<ReductionNode::iterator<TraversalMode::MultiPath>>(
+ module, region, reducerPatterns, tester);
default:
return module.emitError() << "unsupported traversal mode detected";
}
diff --git a/mlir/test/mlir-reduce/reduction-tree/doc-example.mlir b/mlir/test/mlir-reduce/reduction-tree/doc-example.mlir
index 0593db2168afe..a2c674956721f 100644
--- a/mlir/test/mlir-reduce/reduction-tree/doc-example.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree/doc-example.mlir
@@ -1,5 +1,6 @@
// UNSUPPORTED: system-windows
// RUN: mlir-reduce %s -reduction-tree='traversal-mode=0 test=%S/../script/grep-select.sh' | FileCheck %s
+// RUN: mlir-reduce %s -reduction-tree='traversal-mode=2 test=%S/../script/grep-select.sh' | FileCheck %s
// This test case is referenced on the website (mlir/docs/Tools/mlir-reduce.md).
diff --git a/mlir/test/mlir-reduce/reduction-tree/invalid.mlir b/mlir/test/mlir-reduce/reduction-tree/invalid.mlir
index 7e973520a657a..dd8dcfd32422e 100644
--- a/mlir/test/mlir-reduce/reduction-tree/invalid.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree/invalid.mlir
@@ -1,6 +1,7 @@
// UNSUPPORTED: system-windows
// RUN: not mlir-reduce -reduction-tree --no-implicit-module %s 2>&1 | FileCheck %s --check-prefix=CHECK-TREE
// RUN: not mlir-reduce -reduction-tree='traversal-mode=0 test=%S/../script/false.sh' %s 2>&1 | FileCheck %s --check-prefix=CHECK-INTERESTING
+// RUN: not mlir-reduce -reduction-tree='traversal-mode=2 test=%S/../script/false.sh' %s 2>&1 | FileCheck %s --check-prefix=CHECK-INTERESTING
// The reduction passes are currently restricted to 'builtin.module'.
// CHECK-TREE: error: top-level op must be 'builtin.module'
diff --git a/mlir/test/mlir-reduce/reduction-tree/multipath-better-baseline.mlir b/mlir/test/mlir-reduce/reduction-tree/multipath-better-baseline.mlir
new file mode 100644
index 0000000000000..97606747347c0
--- /dev/null
+++ b/mlir/test/mlir-reduce/reduction-tree/multipath-better-baseline.mlir
@@ -0,0 +1,19 @@
+// UNSUPPORTED: system-windows
+// RUN: mlir-reduce %s -reduction-tree='traversal-mode=0 test=%S/../script/multipath-better.sh' | FileCheck %s
+// We are testing the ability of keeping trivially-dead yet interesting code
+
+func.func @multipath_better() {
+ %0 = arith.constant 1 : i32
+ %1 = arith.constant 2 : i32
+ %2 = arith.constant 3 : i32
+ %3 = arith.constant 4 : i32
+ %4 = arith.constant 5 : i32
+ %5 = arith.constant 6 : i32
+ return
+}
+
+// CHECK-LABEL: func @multipath_better
+// CHECK-NEXT: arith.constant 1 : i32
+// CHECK-NEXT: arith.constant 2 : i32
+// CHECK-NEXT: arith.constant 3 : i32
+// CHECK-NEXT: return
diff --git a/mlir/test/mlir-reduce/reduction-tree/multipath-better.mlir b/mlir/test/mlir-reduce/reduction-tree/multipath-better.mlir
new file mode 100644
index 0000000000000..395149a45f9f4
--- /dev/null
+++ b/mlir/test/mlir-reduce/reduction-tree/multipath-better.mlir
@@ -0,0 +1,18 @@
+// UNSUPPORTED: system-windows
+// RUN: mlir-reduce %s -reduction-tree='traversal-mode=2 test=%S/../script/multipath-better.sh' | FileCheck %s
+// We are testing the ability of keeping trivially-dead yet interesting code
+
+func.func @multipath_better() {
+ %0 = arith.constant 1 : i32
+ %1 = arith.constant 2 : i32
+ %2 = arith.constant 3 : i32
+ %3 = arith.constant 4 : i32
+ %4 = arith.constant 5 : i32
+ %5 = arith.constant 6 : i32
+ return
+}
+
+// CHECK-LABEL: func @multipath_better
+// CHECK-NEXT: arith.constant 1 : i32
+// CHECK-NEXT: arith.constant 4 : i32
+// CHECK-NEXT: return
diff --git a/mlir/test/mlir-reduce/reduction-tree/reduction-tree.mlir b/mlir/test/mlir-reduce/reduction-tree/reduction-tree.mlir
index b235ca14d693a..c067925310799 100644
--- a/mlir/test/mlir-reduce/reduction-tree/reduction-tree.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree/reduction-tree.mlir
@@ -1,5 +1,6 @@
// UNSUPPORTED: system-windows
// RUN: mlir-reduce %s -split-input-file -reduction-tree='traversal-mode=0 test=%S/../script/failure-test.sh' | FileCheck %s
+// RUN: mlir-reduce %s -split-input-file -reduction-tree='traversal-mode=2 test=%S/../script/failure-test.sh' | FileCheck %s
// "test.op_crash_long" should be replaced with a shorter form "test.op_crash_short".
// CHECK-NOT: func @simple1() {
diff --git a/mlir/test/mlir-reduce/reduction-tree/simple-test.mlir b/mlir/test/mlir-reduce/reduction-tree/simple-test.mlir
index c9c62fec28234..7a0ff3bd7ca24 100644
--- a/mlir/test/mlir-reduce/reduction-tree/simple-test.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree/simple-test.mlir
@@ -1,5 +1,6 @@
// UNSUPPORTED: system-windows
// RUN: mlir-reduce %s -reduction-tree='traversal-mode=0 test=%S/../script/true.sh' | FileCheck %s
+// RUN: mlir-reduce %s -reduction-tree='traversal-mode=2 test=%S/../script/true.sh' | FileCheck %s
// Since the test.sh always returns 1 (interesting),
// all operations within the ModuleOp should be erased.
diff --git a/mlir/test/mlir-reduce/reduction-tree/trivially-dead.mlir b/mlir/test/mlir-reduce/reduction-tree/trivially-dead.mlir
index b75021226afd0..9e028b6a77490 100644
--- a/mlir/test/mlir-reduce/reduction-tree/trivially-dead.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree/trivially-dead.mlir
@@ -1,5 +1,6 @@
// UNSUPPORTED: system-windows
// RUN: mlir-reduce %s -reduction-tree='traversal-mode=0 test=%S/../script/trivially-dead.sh' | FileCheck %s
+// RUN: mlir-reduce %s -reduction-tree='traversal-mode=2 test=%S/../script/trivially-dead.sh' | FileCheck %s
// We are testing the ability of keeping trivially-dead yet interesting code
func.func @trivially_dead() {
diff --git a/mlir/test/mlir-reduce/script/multipath-better.sh b/mlir/test/mlir-reduce/script/multipath-better.sh
new file mode 100755
index 0000000000000..0d44bb47d0f48
--- /dev/null
+++ b/mlir/test/mlir-reduce/script/multipath-better.sh
@@ -0,0 +1,11 @@
+#!/bin/sh
+
+file=$1
+
+if grep -q "arith.constant 1 : i32" $file && grep -q "arith.constant 2 : i32" $file && grep -q "arith.constant 3 : i32" $file; then
+ exit 1
+fi
+
+if grep -q "arith.constant 1 : i32" $file && grep -q "arith.constant 4 : i32" $file; then
+ exit 1
+fi
More information about the Mlir-commits
mailing list