[llvm-branch-commits] [mlir] [mlir] Add the `OpFoldResults` op fold form (PR #228766)
Victor Perez via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Sat Oct 3 13:30:17 PDT 2026
https://github.com/victor-eds created https://github.com/llvm/llvm-project/pull/228766
An op can now define `OpFoldResults fold(FoldAdaptor)`, enabling partial folding. The drivers still call the legacy `Operation::fold` overloads, so they cannot apply a partial fold yet. As this form has the same parameters as the single-result form, this patch ads some template logic to avoid collisions.
A replacement in the new API can be another result of the op only if the fold keeps that result. If a replacement uses a result that the same fold also replaces, the outcome would depend on the order of the replacements, so the hook drops the replacements.
A `static_assert` rejects each of these cases: an op that defines both the new form and the legacy vector form, and an op with the OneResult trait that defines the new form. The `static_assert` for a single-result fold on an op without the OneResult trait now also names the new form.
RFC: https://discourse.llvm.org/t/rfc-partial-folding-for-multi-result-ops/91954
>From a59bc0aa01e870390e2b1d46be10de2c17e33270 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20P=C3=A9rez=20Carrasco?=
<victor.pc.upm at gmail.com>
Date: Fri, 2 Oct 2026 14:57:03 -0700
Subject: [PATCH] [mlir] Add the OpFoldResults op fold form
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
An op can now define `OpFoldResults fold(FoldAdaptor)`. With this form,
a fold can replace some results of a multi-result op and keep the
others. The drivers still call the legacy `Operation::fold` overloads,
so they do not apply a partial fold yet. This form has the same
parameters as the single-result form, so the single-result detectors
also match it, and `getFoldHookFn` excludes it from the single-result
forms. `getFoldHookFn` selects the new form if the op defines it, else
the single-result form, else the legacy vector form. The new form then
folds the traits if its own fold replaces no result.
A replacement in a new-form result can be another result of the op
only if the fold keeps that result. The hook normalizes the result
first, so the op's own result in its own slot still means "keep". If a
replacement names a result that the same fold also replaces, the outcome
would depend on the order of the replacements, so the hook drops the
replacements. A forwarding fold can return such a result where an op can
use its own results: in a graph region or in an unreachable block.
A `static_assert` rejects each of these cases: an op that defines both
the new form and the legacy vector form, and an op with the OneResult
trait that defines the new form. The `static_assert` for a
single-result fold on an op without the OneResult trait now also names
the new form.
The new unit tests cover the dispatch of the new form, the trait folding
after it, the partial fold with the `DialectFoldInterface` fallback, the
drop of a fold that names a replaced result, in and outside a graph
region, and the constant operands of the new form. The two `normalize`
death tests now reach `normalize` through the new form.
RFC: https://discourse.llvm.org/t/rfc-partial-folding-for-multi-result-ops/91954
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
mlir/include/mlir/IR/OpDefinition.h | 50 ++++-
mlir/include/mlir/IR/OpFoldResult.h | 18 +-
mlir/include/mlir/IR/OperationSupport.h | 3 +-
mlir/lib/IR/OpFoldResult.cpp | 14 ++
mlir/unittests/IR/OpFoldResultsTest.cpp | 272 ++++++++++++++++++++++--
5 files changed, 334 insertions(+), 23 deletions(-)
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index 79ccbab2bc99b..b09f51e217f39 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -1821,6 +1821,15 @@ class Op : public OpState, public Traits<ConcreteType>... {
template <class T>
constexpr static bool has_fold_adaptor_v =
llvm::is_detected<has_fold_adaptor_fold_t, T>::value;
+ /// Trait to check if T provides a 'fold' method with a FoldAdaptor that
+ /// returns OpFoldResults.
+ template <typename T, typename... Args>
+ using has_fold_adaptor_results_fold_t = std::enable_if_t<std::is_same_v<
+ decltype(std::declval<T>().fold(std::declval<typename T::FoldAdaptor>())),
+ OpFoldResults>>;
+ template <class T>
+ constexpr static bool has_fold_adaptor_results_v =
+ llvm::is_detected<has_fold_adaptor_results_fold_t, T>::value;
/// Trait to check if T provides a 'print' method.
template <typename T, typename... Args>
@@ -1921,22 +1930,41 @@ class Op : public OpState, public Traits<ConcreteType>... {
/// Return the internal implementations of each of the OperationName hooks.
static constexpr FoldHookFn getFoldHookFn() {
constexpr bool hasOneResult = hasTrait<OpTrait::OneResult>();
+ constexpr bool hasResultsFold = has_fold_adaptor_results_v<ConcreteType>;
+ // `OpFoldResults fold(FoldAdaptor)` also matches the single-result
+ // detectors.
constexpr bool hasSingleResultFold =
- has_single_result_fold_v<ConcreteType> ||
- has_fold_adaptor_single_result_v<ConcreteType>;
+ !hasResultsFold && (has_single_result_fold_v<ConcreteType> ||
+ has_fold_adaptor_single_result_v<ConcreteType>);
+ constexpr bool hasLegacyFold =
+ has_fold_v<ConcreteType> || has_fold_adaptor_v<ConcreteType>;
// Without this check, the dispatch below silently skips such a `fold`.
static_assert(hasOneResult || !hasSingleResultFold,
"a single-result `OpFoldResult fold(...)` requires the "
"OneResult trait; ops with zero, several, or variadic "
- "results must use `LogicalResult fold(FoldAdaptor, "
+ "results must use `OpFoldResults fold(FoldAdaptor)` or "
+ "`LogicalResult fold(FoldAdaptor, "
+ "SmallVectorImpl<OpFoldResult> &)`");
+ static_assert(!(hasResultsFold && hasLegacyFold),
+ "an op must not define both `OpFoldResults "
+ "fold(FoldAdaptor)` and a legacy `LogicalResult fold(..., "
"SmallVectorImpl<OpFoldResult> &)`");
+ static_assert(!(hasResultsFold && hasOneResult),
+ "an op with the OneResult trait must use the single-result "
+ "`OpFoldResult fold(FoldAdaptor)` form instead of "
+ "`OpFoldResults fold(FoldAdaptor)`");
+ // The operation defines an `OpFoldResults fold(FoldAdaptor)` method.
+ if constexpr (hasResultsFold)
+ return [](Operation *op, ArrayRef<Attribute> operands) {
+ return foldResultsHook<ConcreteType>(op, operands);
+ };
// If the operation is single result and defines a `fold` method.
if constexpr (hasOneResult && hasSingleResultFold)
return [](Operation *op, ArrayRef<Attribute> operands) {
return foldSingleResultHook<ConcreteType>(op, operands);
};
- // The operation is not single result and defines a `fold` method.
- if constexpr (has_fold_v<ConcreteType> || has_fold_adaptor_v<ConcreteType>)
+ // The operation defines a legacy multi-result `fold` method.
+ if constexpr (hasLegacyFold)
return [](Operation *op, ArrayRef<Attribute> operands) {
return foldHook<ConcreteType>(op, operands);
};
@@ -1947,6 +1975,18 @@ class Op : public OpState, public Traits<ConcreteType>... {
operands);
};
}
+ /// Return the result of folding an operation that defines an
+ /// `OpFoldResults fold(FoldAdaptor)` method.
+ template <typename ConcreteOpT>
+ static OpFoldResults foldResultsHook(Operation *op,
+ ArrayRef<Attribute> operands) {
+ OpFoldResults result = cast<ConcreteOpT>(op).fold(
+ typename ConcreteOpT::FoldAdaptor(operands, cast<ConcreteOpT>(op)));
+ result.normalize(op);
+ ::mlir::detail::dropReplacementsOfReplacedResults(op, result);
+ return op_definition_impl::foldTraits<Traits<ConcreteType>...>(
+ op, operands, std::move(result));
+ }
/// Return the result of folding a single result operation that defines a
/// `fold` method.
template <typename ConcreteOpT>
diff --git a/mlir/include/mlir/IR/OpFoldResult.h b/mlir/include/mlir/IR/OpFoldResult.h
index f65b703206cdf..df6c3eabffd18 100644
--- a/mlir/include/mlir/IR/OpFoldResult.h
+++ b/mlir/include/mlir/IR/OpFoldResult.h
@@ -67,10 +67,12 @@ namespace mlir {
/// Allow printing to a stream.
raw_ostream &operator<<(raw_ostream &os, OpFoldResult ofr);
-/// The result of a fold of any op. Replacement i is an Attribute (replace
-/// result i with a constant), a Value (replace result i with that value), or
-/// null / op->getResult(i) (keep result i). A separate bit records an in-place
-/// change of the op.
+/// The result of a fold of any op. An op with exactly one fixed result defines
+/// its folder with OpFoldResult; any other op can return OpFoldResults
+/// directly. Replacement i is an Attribute (replace result i with a
+/// constant), a Value (replace result i with that value), or null /
+/// op->getResult(i) (keep result i). A separate bit records an in-place change
+/// of the op.
class [[nodiscard]] OpFoldResults {
public:
/// Failure: the fold did not apply and the IR is unchanged.
@@ -161,6 +163,14 @@ namespace detail {
/// then normalized.
OpFoldResults convertLegacyFoldResults(LogicalResult status,
ArrayRef<OpFoldResult> results);
+
+/// Drop the replacements of the normalized `result` if a replacement names
+/// another result of `op` that `result` also replaces, because the outcome
+/// would depend on the order in which a driver replaces the results. A
+/// forwarding fold can return such a result where an op can use its own
+/// results: in a graph region or in an unreachable block. The in-place bit
+/// stays.
+void dropReplacementsOfReplacedResults(Operation *op, OpFoldResults &result);
} // namespace detail
} // namespace mlir
diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index 7248903860788..15a7fa0ec563a 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -302,7 +302,8 @@ class OperationName {
/// 1. A replacement is an Attribute (replace the result with a constant),
/// a Value (replace the result with that value), or null or the result
/// itself (keep the result). Replacement i may be result j only if
- /// result j is kept.
+ /// result j is kept. A forwarding fold can break this rule in a graph
+ /// region or an unreachable block; then its replacements are dropped.
/// 2. A failure replaces no result and has no in-place mark. The IR must be
/// unchanged.
/// 3. The hook can mutate the operation in place, without changing anything
diff --git a/mlir/lib/IR/OpFoldResult.cpp b/mlir/lib/IR/OpFoldResult.cpp
index 185d771e773b9..2c6c1ca8ecafc 100644
--- a/mlir/lib/IR/OpFoldResult.cpp
+++ b/mlir/lib/IR/OpFoldResult.cpp
@@ -149,3 +149,17 @@ OpFoldResults detail::convertLegacyFoldResults(LogicalResult status,
"legacy fold returned a null result");
return results;
}
+
+void detail::dropReplacementsOfReplacedResults(Operation *op,
+ OpFoldResults &result) {
+ bool namesReplacedResult =
+ llvm::any_of(result.getReplacements(), [&](OpFoldResult replacement) {
+ auto opResult = dyn_cast_if_present<OpResult>(
+ dyn_cast_if_present<Value>(replacement));
+ return opResult && opResult.getOwner() == op &&
+ result[opResult.getResultNumber()];
+ });
+ if (!namesReplacedResult)
+ return;
+ result = success(result.modifiedInPlace());
+}
diff --git a/mlir/unittests/IR/OpFoldResultsTest.cpp b/mlir/unittests/IR/OpFoldResultsTest.cpp
index bd118d54a2738..1d1991ba86641 100644
--- a/mlir/unittests/IR/OpFoldResultsTest.cpp
+++ b/mlir/unittests/IR/OpFoldResultsTest.cpp
@@ -30,8 +30,10 @@ using LegacyFoldFn =
/// Per-test behavior of the ops, the trait, and the dialect interface below.
struct FoldState {
+ std::function<OpFoldResults(Operation *)> opFoldFn;
LegacyFoldFn traitFoldFn;
LegacyFoldFn dialectFoldFn;
+ SmallVector<Attribute> seenOperands;
unsigned traitCalls = 0;
unsigned dialectCalls = 0;
};
@@ -50,13 +52,40 @@ struct LegacyFoldTrait
}
};
-/// An op with two results and no fold method. Only its trait folds.
+template <typename OpT>
+struct FoldAdaptorImpl {
+ FoldAdaptorImpl(ArrayRef<Attribute> operands, OpT) : operands(operands) {}
+ ArrayRef<Attribute> getOperands() const { return operands; }
+ ArrayRef<Attribute> operands;
+};
+
+/// An op with two results that defines `OpFoldResults fold(FoldAdaptor)`.
struct PartialFoldOp : public Op<PartialFoldOp, OpTrait::NResults<2>::Impl,
OpTrait::VariadicOperands, LegacyFoldTrait> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PartialFoldOp)
using Op::Op;
+ using FoldAdaptor = FoldAdaptorImpl<PartialFoldOp>;
static ArrayRef<StringRef> getAttributeNames() { return {}; }
static StringRef getOperationName() { return "fold_test.partial"; }
+ OpFoldResults fold(FoldAdaptor adaptor) {
+ foldState->seenOperands.assign(adaptor.getOperands().begin(),
+ adaptor.getOperands().end());
+ return foldState->opFoldFn ? foldState->opFoldFn(getOperation())
+ : failure();
+ }
+};
+
+struct ConstantOp : public Op<ConstantOp, OpTrait::OneResult,
+ OpTrait::ZeroOperands, OpTrait::ConstantLike> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConstantOp)
+ using Op::Op;
+ using FoldAdaptor = FoldAdaptorImpl<ConstantOp>;
+ static ArrayRef<StringRef> getAttributeNames() {
+ static StringRef names[] = {"value"};
+ return names;
+ }
+ static StringRef getOperationName() { return "fold_test.constant"; }
+ OpFoldResult fold(FoldAdaptor) { return getOperation()->getAttr("value"); }
};
struct TestFoldInterface : public DialectFoldInterface {
@@ -75,7 +104,7 @@ struct FoldTestDialect : public Dialect {
explicit FoldTestDialect(MLIRContext *context)
: Dialect(getDialectNamespace(), context,
TypeID::get<FoldTestDialect>()) {
- addOperations<PartialFoldOp>();
+ addOperations<PartialFoldOp, ConstantOp>();
addInterfaces<TestFoldInterface>();
}
};
@@ -648,8 +677,8 @@ TEST_F(OpFoldResultsTest, LegacyDynamicFoldHookMayForwardAnotherResult) {
Operation *op = createOp({i32, i32}, "test_fold.legacy_op");
Attribute attr = builder.getI32IntegerAttr(1);
- // Replacement 0 is result 1, which the fold also replaces. The legacy
- // adapter does not reject this.
+ // Replacement 0 is result 1, which the fold also replaces. The
+ // replaced-result check applies only to `OpFoldResults fold(FoldAdaptor)`.
legacyFoldFn = [&](Operation *foldedOp,
SmallVectorImpl<OpFoldResult> &results) {
results.push_back(foldedOp->getResult(1));
@@ -669,6 +698,61 @@ TEST_F(OpFoldResultsTest, LegacyDynamicFoldHookMayForwardAnotherResult) {
EXPECT_EQ(results[1], OpFoldResult(attr));
}
+// In a graph region, a forwarding fold can name another result of its op. If
+// the same fold also replaces that result, the fold does not apply.
+TEST_F(OpFoldResultsTest, GraphRegionFoldNamingReplacedResultDoesNotApply) {
+ // The region of an unregistered op may be a graph region.
+ OperationState parentState(UnknownLoc::get(&context), "foo.graph");
+ parentState.addRegion();
+ Operation *parent = Operation::create(parentState);
+ ops.push_back(parent);
+ Block *block = new Block();
+ parent->getRegion(0).push_back(block);
+ OperationState state(UnknownLoc::get(&context), "fold_test.partial");
+ state.addTypes({i32, i32});
+ Operation *op = Operation::create(state);
+ block->push_back(op);
+ Attribute attr = builder.getI32IntegerAttr(1);
+
+ foldState->opFoldFn = [&](Operation *foldedOp) -> OpFoldResults {
+ return {foldedOp->getResult(1), attr};
+ };
+ EXPECT_TRUE(op->fold().failed());
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+
+ // The in-place bit stays.
+ foldState->opFoldFn = [&](Operation *foldedOp) {
+ OpFoldResults result = {foldedOp->getResult(1), attr};
+ result.setModifiedInPlace();
+ return result;
+ };
+ OpFoldResults result = op->fold();
+ EXPECT_TRUE(result.modifiedInPlace());
+ EXPECT_FALSE(result.replacesAny());
+}
+
+// Outside a graph region, an op can also name its own results, for example in
+// an unreachable block. The fold does not apply there either.
+TEST_F(OpFoldResultsTest, NonGraphRegionFoldNamingReplacedResultDoesNotApply) {
+ Operation *producer = createOp({i32});
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ foldState->opFoldFn = [&](Operation *foldedOp) -> OpFoldResults {
+ return {foldedOp->getResult(1), attr};
+ };
+ EXPECT_TRUE(op->fold().failed());
+
+ // Replacement 0 is result 1, but the fold also replaces result 1.
+ OpFoldResults result(op);
+ result.replace(0u, op->getResult(1));
+ result.replace(1u, producer->getResult(0));
+ result.normalize(op);
+ detail::dropReplacementsOfReplacedResults(op, result);
+ EXPECT_TRUE(result.failed());
+}
+
TEST_F(OpFoldResultsTest, LegacyDynamicOpDefinitionGet) {
loadDynamicDialect();
Operation *op = createOp({i32}, "test_fold.get_op");
@@ -713,6 +797,130 @@ TEST_F(OpFoldResultsTest, FreeHelpersMatchMembers) {
EXPECT_TRUE(failed(failedResult));
}
+TEST_F(OpFoldResultsTest, OwnPartialFoldSkipsTraits) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ foldState->opFoldFn = [&](Operation *foldedOp) {
+ OpFoldResults result(foldedOp);
+ result.replace(1u, attr);
+ return result;
+ };
+ foldState->traitFoldFn = [](Operation *, SmallVectorImpl<OpFoldResult> &) {
+ return success();
+ };
+
+ OpFoldResults result = op->fold();
+ EXPECT_TRUE(result.succeeded());
+ EXPECT_FALSE(result.modifiedInPlace());
+ EXPECT_FALSE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_FALSE(result[0]);
+ EXPECT_EQ(result[1], OpFoldResult(attr));
+
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+ EXPECT_EQ(foldState->traitCalls, 0u);
+}
+
+TEST_F(OpFoldResultsTest, OwnInPlaceThenTraitReplacesAll) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute lhs = builder.getI32IntegerAttr(1);
+ Attribute rhs = builder.getI32IntegerAttr(2);
+ foldState->opFoldFn = [](Operation *) -> OpFoldResults { return success(); };
+ foldState->traitFoldFn = [&](Operation *,
+ SmallVectorImpl<OpFoldResult> &results) {
+ results.push_back(lhs);
+ results.push_back(rhs);
+ return success();
+ };
+
+ OpFoldResults result = op->fold();
+ EXPECT_TRUE(result.modifiedInPlace());
+ EXPECT_TRUE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_EQ(result[0], OpFoldResult(lhs));
+ EXPECT_EQ(result[1], OpFoldResult(rhs));
+ EXPECT_EQ(foldState->traitCalls, 1u);
+
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ ASSERT_EQ(results.size(), 2u);
+ EXPECT_EQ(results[0], OpFoldResult(lhs));
+ EXPECT_EQ(results[1], OpFoldResult(rhs));
+ EXPECT_EQ(foldState->traitCalls, 2u);
+}
+
+TEST_F(OpFoldResultsTest, OwnInPlaceThenTraitReplacesSome) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ foldState->opFoldFn = [](Operation *) -> OpFoldResults { return success(); };
+ foldState->traitFoldFn = [&](Operation *foldedOp,
+ SmallVectorImpl<OpFoldResult> &results) {
+ results.push_back(foldedOp->getResult(0));
+ results.push_back(attr);
+ return success();
+ };
+
+ OpFoldResults result = op->fold();
+ EXPECT_TRUE(result.modifiedInPlace());
+ EXPECT_TRUE(result.replacesAny());
+ EXPECT_FALSE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_FALSE(result[0]);
+ EXPECT_EQ(result[1], OpFoldResult(attr));
+ EXPECT_EQ(foldState->traitCalls, 1u);
+
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+}
+
+TEST_F(OpFoldResultsTest, OwnFailureThenTraitInPlace) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ foldState->opFoldFn = [](Operation *) -> OpFoldResults { return failure(); };
+ foldState->traitFoldFn = [](Operation *, SmallVectorImpl<OpFoldResult> &) {
+ return success();
+ };
+
+ OpFoldResults result = op->fold();
+ EXPECT_TRUE(result.modifiedInPlace());
+ EXPECT_FALSE(result.replacesAny());
+ EXPECT_EQ(foldState->traitCalls, 1u);
+
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+}
+
+TEST_F(OpFoldResultsTest, OwnResultsNormalizeToFailureSoTraitsRun) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ foldState->opFoldFn = [](Operation *foldedOp) -> OpFoldResults {
+ return foldedOp->getResults();
+ };
+ foldState->traitFoldFn = [&](Operation *,
+ SmallVectorImpl<OpFoldResult> &results) {
+ results.append(2, attr);
+ return success();
+ };
+
+ OpFoldResults result = op->fold();
+ EXPECT_FALSE(result.modifiedInPlace());
+ EXPECT_TRUE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_EQ(result[0], OpFoldResult(attr));
+ EXPECT_EQ(result[1], OpFoldResult(attr));
+ EXPECT_EQ(foldState->traitCalls, 1u);
+
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ ASSERT_EQ(results.size(), 2u);
+ EXPECT_EQ(results[0], OpFoldResult(attr));
+ EXPECT_EQ(results[1], OpFoldResult(attr));
+ EXPECT_EQ(foldState->traitCalls, 2u);
+}
+
TEST_F(OpFoldResultsTest, DialectFoldInterfaceFallback) {
Operation *op = createOp({i32, i32}, "fold_test.partial");
Attribute attr = builder.getI32IntegerAttr(1);
@@ -722,7 +930,7 @@ TEST_F(OpFoldResultsTest, DialectFoldInterfaceFallback) {
return success();
};
- // The traits fail.
+ // The op fold and the traits fail.
OpFoldResults result = op->fold();
EXPECT_TRUE(result.replacesAll());
ASSERT_EQ(result.size(), 2u);
@@ -734,8 +942,26 @@ TEST_F(OpFoldResultsTest, DialectFoldInterfaceFallback) {
EXPECT_EQ(results.size(), 2u);
EXPECT_EQ(foldState->dialectCalls, 2u);
- // The fallback result is normalized.
+ // A partial fold does not fall back.
+ foldState->opFoldFn = [&](Operation *foldedOp) {
+ OpFoldResults partial(foldedOp);
+ partial.replace(0u, attr);
+ return partial;
+ };
foldState->dialectCalls = 0;
+ result = op->fold();
+ EXPECT_TRUE(result.succeeded());
+ EXPECT_FALSE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_EQ(result[0], OpFoldResult(attr));
+ EXPECT_FALSE(result[1]);
+ EXPECT_EQ(foldState->dialectCalls, 0u);
+ results.clear();
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_EQ(foldState->dialectCalls, 0u);
+
+ // The fallback result is normalized.
+ foldState->opFoldFn = nullptr;
foldState->dialectFoldFn = [](Operation *foldedOp,
SmallVectorImpl<OpFoldResult> &results) {
llvm::append_range(results, foldedOp->getResults());
@@ -745,6 +971,21 @@ TEST_F(OpFoldResultsTest, DialectFoldInterfaceFallback) {
EXPECT_EQ(foldState->dialectCalls, 1u);
}
+TEST_F(OpFoldResultsTest, FoldComputesConstantOperands) {
+ Operation *constant = createOp({i32}, "fold_test.constant");
+ Attribute attr = builder.getI32IntegerAttr(7);
+ constant->setAttr("value", attr);
+ Operation *producer = createOp({i32});
+ Operation *op =
+ createOp({i32, i32}, "fold_test.partial",
+ ValueRange{constant->getResult(0), producer->getResult(0)});
+
+ EXPECT_TRUE(op->fold().failed());
+ ASSERT_EQ(foldState->seenOperands.size(), 2u);
+ EXPECT_EQ(foldState->seenOperands[0], attr);
+ EXPECT_FALSE(foldState->seenOperands[1]);
+}
+
#ifdef GTEST_HAS_DEATH_TEST
#ifndef NDEBUG
namespace {
@@ -753,17 +994,22 @@ class OpFoldResultsDeathTest : public OpFoldResultsTest {};
TEST_F(OpFoldResultsDeathTest, ValueReplacementOfIncorrectType) {
Operation *producer = createOp({f32});
- Operation *op = createOp({i32, i32});
- OpFoldResults result(op);
- result.replace(0u, producer->getResult(0));
- EXPECT_DEATH(result.normalize(op), "incorrect fold result type");
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ foldState->opFoldFn = [&](Operation *foldedOp) {
+ OpFoldResults result(foldedOp);
+ result.replace(0u, producer->getResult(0));
+ return result;
+ };
+ EXPECT_DEATH((void)op->fold(), "incorrect fold result type");
}
TEST_F(OpFoldResultsDeathTest, ReplacementCountMismatch) {
- Operation *op = createOp({i32, i32});
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
Attribute attr = builder.getI32IntegerAttr(1);
- OpFoldResults result = {attr, attr, attr};
- EXPECT_DEATH(result.normalize(op),
+ foldState->opFoldFn = [&](Operation *) -> OpFoldResults {
+ return {attr, attr, attr};
+ };
+ EXPECT_DEATH((void)op->fold(),
"expected one replacement per operation result");
}
More information about the llvm-branch-commits
mailing list