[Mlir-commits] [mlir] [mlir][reducer] add replace operands pattern to mlir-reduce (PR #192874)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Apr 19 14:17:43 PDT 2026
https://github.com/aidint updated https://github.com/llvm/llvm-project/pull/192874
>From eaf386d11467d80ebb2eb0ae079461ca472a9d34 Mon Sep 17 00:00:00 2001
From: aidint <at.aidin at gmail.com>
Date: Sun, 19 Apr 2026 23:01:54 +0200
Subject: [PATCH 1/2] add replace operands pattern to mlir-reduce
---
.../mlir/Reducer/Patterns/ReplaceOperands.h | 61 +++++++
mlir/lib/Reducer/CMakeLists.txt | 1 +
mlir/lib/Reducer/Patterns/ReplaceOperands.cpp | 151 ++++++++++++++++++
.../Dialect/Test/TestDialectInterfaces.cpp | 6 +
mlir/test/mlir-reduce/reduction-tree.mlir | 19 +++
5 files changed, 238 insertions(+)
create mode 100644 mlir/include/mlir/Reducer/Patterns/ReplaceOperands.h
create mode 100644 mlir/lib/Reducer/Patterns/ReplaceOperands.cpp
diff --git a/mlir/include/mlir/Reducer/Patterns/ReplaceOperands.h b/mlir/include/mlir/Reducer/Patterns/ReplaceOperands.h
new file mode 100644
index 0000000000000..2e71b7bf00799
--- /dev/null
+++ b/mlir/include/mlir/Reducer/Patterns/ReplaceOperands.h
@@ -0,0 +1,61 @@
+//===- ReplaceOperands.h - Replacing Operands Reduction Pattern -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_REDUCER_REPLACEOPERANDS_H
+#define MLIR_REDUCER_REPLACEOPERANDS_H
+
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Reducer/Tester.h"
+
+namespace mlir {
+
+struct OperandReductionNode {
+ using Range = std::pair<size_t, size_t>;
+
+ OperandReductionNode(Operation *reductionOp, ArrayRef<Range> ranges);
+
+ ArrayRef<OpOperand *> getOperandForType(Type type) {
+ return operandMap[type];
+ }
+
+ const DenseSet<Value> &getUsedValues() const { return usedValues; }
+
+ auto getNeededTypes() const { return operandMap.keys(); }
+
+ ModuleOp getModule() const { return module.get(); }
+ Operation *getOperation() const { return op; }
+
+ Value getMappedValue(Value value) { return mapping.lookup(value); }
+
+private:
+ Operation *op;
+ OwningOpRef<ModuleOp> module;
+ IRMapping mapping;
+ DenseSet<Value> usedValues;
+ SmallVector<Range, 0> startRanges;
+ SmallVector<Range, 0> discardRanges;
+ llvm::MapVector<Type, SmallVector<OpOperand *, 2>> operandMap;
+};
+
+struct ReplaceOperandsPattern : public mlir::RewritePattern {
+private:
+ Tester &tester;
+
+public:
+ ReplaceOperandsPattern(mlir::MLIRContext *context, Tester &tester)
+ : mlir::RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context),
+ tester(tester) {}
+
+ mlir::LogicalResult
+ matchAndRewrite(mlir::Operation *op,
+ mlir::PatternRewriter &rewriter) const override;
+};
+} // namespace mlir
+
+#endif
diff --git a/mlir/lib/Reducer/CMakeLists.txt b/mlir/lib/Reducer/CMakeLists.txt
index 68864e373c993..4212f821534f8 100644
--- a/mlir/lib/Reducer/CMakeLists.txt
+++ b/mlir/lib/Reducer/CMakeLists.txt
@@ -3,6 +3,7 @@ add_mlir_library(MLIRReduce
ReductionNode.cpp
ReductionTreePass.cpp
Tester.cpp
+ Patterns/ReplaceOperands.cpp
LINK_LIBS PUBLIC
MLIRIR
diff --git a/mlir/lib/Reducer/Patterns/ReplaceOperands.cpp b/mlir/lib/Reducer/Patterns/ReplaceOperands.cpp
new file mode 100644
index 0000000000000..9285dbee6d329
--- /dev/null
+++ b/mlir/lib/Reducer/Patterns/ReplaceOperands.cpp
@@ -0,0 +1,151 @@
+//===- ReplaceOperands.h - Replacing Operands Reduction Pattern -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Reducer/Patterns/ReplaceOperands.h"
+#include <cstddef>
+
+using namespace mlir;
+
+OperandReductionNode::OperandReductionNode(Operation *reductionOp,
+ ArrayRef<Range> ranges)
+ : startRanges(ranges), discardRanges(ranges) {
+
+ ModuleOp moduleOp = reductionOp->getParentOfType<ModuleOp>();
+ module = cast<ModuleOp>(moduleOp->clone(mapping));
+ op = mapping.lookup(reductionOp);
+ usedValues = DenseSet<Value>(op->operand_begin(), op->operand_end());
+
+ size_t rangeIndex = 0;
+ for (const auto &[index, operand] : enumerate(op->getOpOperands())) {
+ if (rangeIndex < discardRanges.size() &&
+ index == discardRanges[rangeIndex].second)
+ ++rangeIndex;
+ if (rangeIndex == discardRanges.size() ||
+ index < discardRanges[rangeIndex].first)
+ operandMap[operand.get().getType()].push_back(&operand);
+ }
+}
+
+LogicalResult
+ReplaceOperandsPattern::matchAndRewrite(Operation *op,
+ PatternRewriter &rewriter) const {
+
+ Block *block = op->getBlock();
+ // If Operation has no parent block, then return failure
+ if (!block)
+ return failure();
+
+ // If operation has no operands, we don't have anything to do
+ if (op->getNumOperands() == 0)
+ return failure();
+
+ // candidateValues stores suitable replacements per each type
+ DenseMap<Type, SmallVector<Value, 0>> candidateValues;
+
+ // Add block arguments first (they come first in program order)
+ for (BlockArgument arg : block->getArguments())
+ candidateValues[arg.getType()].push_back(arg);
+
+ // Walk operations in the block to find remaining types
+ for (auto &blockOp : block->getOperations()) {
+ // Stop before reaching the current operation
+ if (&blockOp == op)
+ break;
+
+ for (Value result : blockOp.getResults())
+ candidateValues[result.getType()].push_back(result);
+ }
+
+ OperandReductionNode node(op, {{0, 0}});
+
+ auto types = llvm::filter_to_vector(node.getNeededTypes(), [&](Type type) {
+ return candidateValues.contains(type);
+ });
+ auto values = llvm::map_to_vector(types, [&](Type type) -> ArrayRef<Value> {
+ return candidateValues.find(type)->second;
+ });
+
+ size_t typesLen = types.size();
+ if (typesLen == 0)
+ return failure();
+
+ // We'll use a gray-code like algorithm on an arbitrary radix to iterate over
+ // different combinations of values
+ SmallVector<size_t, 0> idx(typesLen, 0);
+ SmallVector<int8_t, 0> dir(typesLen, 1);
+
+ bool replacementFound = false, changed = false;
+ auto applyIndex = [&](size_t index) {
+ for (auto *operand : node.getOperandForType(types[index])) {
+ auto value = node.getMappedValue(values[index][idx[index]]);
+ if (operand->get() == value)
+ continue;
+ operand->set(value);
+ changed = true;
+ }
+ };
+
+ // The caller guarantees the input module is interesting, so we skip
+ // testing the initial configuration and go straight to the first
+ // Gray-code neighbor.
+ for (auto typeIndex : llvm::seq(typesLen))
+ applyIndex(typeIndex);
+
+ while (true) {
+
+ if (changed) {
+ auto [isInteresting, size] = tester.isInteresting(node.getModule());
+ if (isInteresting == Tester::Interestingness::True) {
+ replacementFound = true;
+ break;
+ }
+ }
+
+ ptrdiff_t next = 0;
+ size_t index = 0;
+ changed = false;
+ bool exhausted = true;
+
+ while (index < typesLen) {
+ next = (ptrdiff_t)idx[index] + dir[index];
+ if ((next >= 0) && ((size_t)next < values[index].size())) {
+ idx[index] = next;
+ applyIndex(index);
+ exhausted = false;
+ break;
+ }
+ dir[index] = -dir[index];
+ ++index;
+ }
+ if (exhausted)
+ break;
+ }
+
+ if (replacementFound) {
+ DenseMap<Type, unsigned> typeIdx;
+ typeIdx.reserve(typesLen);
+
+ for (auto [k, v] : llvm::zip(types, idx))
+ typeIdx[k] = v;
+
+ // Instead of replacing operands inplace, we'll replace the operation
+ // completely. If we don't replace the operation with a new one, the greedy
+ // driver will add the operation to the worklist again.
+ auto *newOp = rewriter.clone(*op);
+ for (auto &operand : newOp->getOpOperands()) {
+ auto type = operand.get().getType();
+ auto it = typeIdx.find(type);
+ if (it != typeIdx.end())
+ operand.set(candidateValues.find(type)->second[it->second]);
+ }
+ rewriter.replaceOp(op, newOp);
+ return success();
+ }
+
+ return failure();
+}
diff --git a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
index 1c9dbe1640687..107dd55f60a9d 100644
--- a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
@@ -10,7 +10,9 @@
#include "TestOps.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Interfaces/FoldInterfaces.h"
+#include "mlir/Reducer/Patterns/ReplaceOperands.h"
#include "mlir/Reducer/ReductionPatternInterface.h"
+#include "mlir/Reducer/Tester.h"
#include "mlir/Transforms/InliningUtils.h"
using namespace mlir;
@@ -430,6 +432,10 @@ struct TestReductionPatternInterface : public DialectReductionPatternInterface {
void populateReductionPatterns(RewritePatternSet &patterns) const final {
populateTestReductionPatterns(patterns);
}
+
+ void populateReductionPatternsWithTester(RewritePatternSet &patterns, Tester &tester) const final {
+ patterns.add<ReplaceOperandsPattern>(getContext(), tester);
+ }
};
} // namespace
diff --git a/mlir/test/mlir-reduce/reduction-tree.mlir b/mlir/test/mlir-reduce/reduction-tree.mlir
index 2aee89741b42b..8d046fb8177d5 100644
--- a/mlir/test/mlir-reduce/reduction-tree.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree.mlir
@@ -58,3 +58,22 @@ func.func @simple4(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
func.func @simple5() {
return
}
+
+// -----
+
+// This test case will be reduced by ReplaceOperandsPattern.
+// crash's i32 operands should be replaced by %arg0.
+// Extra dead operations should be removed.
+
+// CHECK-LABEL: func.func @replace(%arg0: i32) -> i32 {
+func.func @replace(%arg0: i32) -> i32 {
+ // CHECK-NOT: {{.*}} = arith.constant 1 : i32
+ %0 = arith.constant 1 : i32
+ // CHECK: {{.*}} = arith.constant 2.000000e+00 : f32
+ %1 = arith.constant 2.0 : f32
+ // CHECK-NOT: {{.*}} = arith.constant 3.0 : i32
+ %2 = arith.constant 3 : i32
+ // CHECK: {{.*}} = "test.op_crash"(%arg0, {{.*}}, %arg0) : (i32, f32, i32) -> i32
+ %crash = "test.op_crash"(%0, %1, %2) : (i32, f32, i32) -> i32
+ return %crash : i32
+}
>From bf0fda467455c819a099181a3835d3ae3e5a4681 Mon Sep 17 00:00:00 2001
From: aidint <at.aidin at gmail.com>
Date: Sun, 19 Apr 2026 23:17:31 +0200
Subject: [PATCH 2/2] resolve clang format issue
---
mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
index 107dd55f60a9d..2d6fa87b71fd3 100644
--- a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
@@ -433,7 +433,8 @@ struct TestReductionPatternInterface : public DialectReductionPatternInterface {
populateTestReductionPatterns(patterns);
}
- void populateReductionPatternsWithTester(RewritePatternSet &patterns, Tester &tester) const final {
+ void populateReductionPatternsWithTester(RewritePatternSet &patterns,
+ Tester &tester) const final {
patterns.add<ReplaceOperandsPattern>(getContext(), tester);
}
};
More information about the Mlir-commits
mailing list