[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
Sun Oct 4 15:04:59 PDT 2026
https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/228766
>From daed177fe7fc554a3330811c6ffeaf738610f789 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 the `OpFoldResults` of the new form can be another
result of the op only if the fold keeps that result. The hook normalizes
the `OpFoldResults` first, so replacement i that is result i 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
replacement where an op can use its own results: in a graph region or
in an unreachable block. Only the new op form drops such replacements,
and the documentation of `OperationName::foldHook` says so.
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 test op `test.fold_dispatch` has a fold in the new form and a
legacy fold trait, and both read their result from attributes. New
cases in `IR/fold-dispatch.mlir` cover the dispatch of the new form, the
trait folds after it, the partial fold with the `DialectFoldInterface`
fallback, the drop of a fold that names a replaced result, in a plain
block, in a graph region, and in an unreachable block, 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 | 4 +-
mlir/lib/IR/OpFoldResult.cpp | 14 ++++
mlir/test/IR/fold-dispatch.mlir | 78 ++++++++++++++++++++++-
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 24 +++++++
mlir/test/lib/Dialect/Test/TestOps.h | 24 +++++++
mlir/test/lib/Dialect/Test/TestOps.td | 16 +++++
mlir/unittests/IR/OpFoldResultsTest.cpp | 72 ++++++++++++++++++---
9 files changed, 279 insertions(+), 21 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 c6186068d3243..ca368681d1d10 100644
--- a/mlir/include/mlir/IR/OpFoldResult.h
+++ b/mlir/include/mlir/IR/OpFoldResult.h
@@ -66,10 +66,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.
@@ -160,6 +162,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 5e3588d001940..9cdf02c07dee5 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -302,7 +302,9 @@ 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 in an unreachable block. If an `OpFoldResults fold` method
+ /// of the operation breaks it, its replacements are dropped.
/// 2. A failure replaces no result and does not set the in-place bit. 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/test/IR/fold-dispatch.mlir b/mlir/test/IR/fold-dispatch.mlir
index c77f8e89f12ee..4141882bda878 100644
--- a/mlir/test/IR/fold-dispatch.mlir
+++ b/mlir/test/IR/fold-dispatch.mlir
@@ -2,7 +2,7 @@
// `transform.test_fold` folds each test op with `Operation::fold` and reports
// the result as a remark. The test ops fold as their attributes say; see
-// `getConfiguredLegacyFoldResults`.
+// `getConfiguredLegacyFoldResults` and `getConfiguredFoldResults`.
// `test.fold_dispatch_fallback` has no fold and no fold trait, so only the
// fallback fold of the test dialect can fold it.
@@ -28,11 +28,85 @@ func.func @dialect_fallback(%arg0: i32) {
return
}
+// `test.fold_dispatch` has a fold in the `OpFoldResults` form and a legacy
+// fold trait. The trait folds only when the fold of the op replaces no result.
+func.func @own_fold_and_trait(%arg0: i32) {
+ // A partial fold of the op skips the trait.
+ // expected-remark @below {{fold: [keep, 42 : i32]}}
+ %0:2 = "test.fold_dispatch"()
+ {fold = [unit, 42 : i32], legacy_trait_fold_in_place_count = 1}
+ : () -> (i32, i32)
+ // After an in-place fold of the op, the trait replaces every result...
+ // expected-remark @below {{fold: [1 : i32, 2 : i32] in place}}
+ %1:2 = "test.fold_dispatch"()
+ {fold_in_place_count = 1, legacy_trait_fold = [1 : i32, 2 : i32]}
+ : () -> (i32, i32)
+ // ... or some results.
+ // expected-remark @below {{fold: [keep, 1 : i32] in place}}
+ %2:2 = "test.fold_dispatch"()
+ {fold_in_place_count = 1, legacy_trait_fold = ["result:0", 1 : i32]}
+ : () -> (i32, i32)
+ // After a failure of the op, the trait changes the op in place.
+ // expected-remark @below {{fold: in place}}
+ %3:2 = "test.fold_dispatch"() {legacy_trait_fold_in_place_count = 1}
+ : () -> (i32, i32)
+ // The fold of the op keeps every result, which normalizes to a failure, so
+ // the trait folds.
+ // expected-remark @below {{fold: [1 : i32, 1 : i32]}}
+ %4:2 = "test.fold_dispatch"()
+ {fold = ["result:0", "result:1"], legacy_trait_fold = [1 : i32, 1 : i32]}
+ : () -> (i32, i32)
+ // A partial fold of the op does not fall back on the dialect.
+ // expected-remark @below {{fold: [42 : i32, keep]}}
+ %5:2 = "test.fold_dispatch"()
+ {fold = [42 : i32, unit], dialect_fold = [7 : i32, 7 : i32]}
+ : () -> (i32, i32)
+ // The fold of the op gets the constant operands.
+ %c7 = "test.constant"() <{value = 7 : i32}> : () -> i32
+ // expected-remark @below {{fold: [7 : i32, keep]}}
+ %6:2 = "test.fold_dispatch"(%c7, %arg0)
+ {fold = ["operand_attr:0", "operand_attr:1"]} : (i32, i32) -> (i32, i32)
+ return
+}
+
+// A replacement can name another result of the op only if the fold keeps that
+// result. Otherwise the fold hook drops the replacements, and only the
+// in-place change stays.
+func.func @fold_naming_replaced_result() {
+ // expected-remark @below {{fold: failure}}
+ %0:2 = "test.fold_dispatch"() {fold = ["result:1", 1 : i32]}
+ : () -> (i32, i32)
+ // expected-remark @below {{fold: in place}}
+ %1:2 = "test.fold_dispatch"()
+ {fold = ["result:1", 1 : i32], fold_with_in_place} : () -> (i32, i32)
+ // expected-remark @below {{fold: [result 1, keep]}}
+ %2:2 = "test.fold_dispatch"() {fold = ["result:1", unit]}
+ : () -> (i32, i32)
+ // An op can use its own results in a graph region.
+ test.graph_region {
+ // expected-remark @below {{fold: failure}}
+ %3:2 = "test.fold_dispatch"(%3#1, %3#0) {fold = ["operand:0", "operand:1"]}
+ : (i32, i32) -> (i32, i32)
+ "test.valid"(%3#0, %3#1) : (i32, i32) -> ()
+ }
+ return
+}
+
+// An op can also use its own results in an unreachable block.
+func.func @fold_naming_replaced_result_unreachable() {
+ return
+^bb1:
+ // expected-remark @below {{fold: failure}}
+ %0:2 = "test.fold_dispatch"(%0#1, %0#0) {fold = ["operand:0", "operand:1"]}
+ : (i32, i32) -> (i32, i32)
+ return
+}
+
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(
%root: !transform.any_op {transform.readonly}) {
%ops = transform.structured.match
- ops{["test.fold_dispatch_fallback"]} in %root
+ ops{["test.fold_dispatch_fallback", "test.fold_dispatch"]} in %root
: (!transform.any_op) -> !transform.any_op
transform.test_fold %ops : !transform.any_op
transform.yield
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 4d6b6b226dec4..fc3bbf8035065 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -2377,3 +2377,27 @@ LogicalResult test::getConfiguredLegacyFoldResults(
results.push_back(getConfiguredReplacement(op, operands, element));
return success();
}
+
+OpFoldResults test::getConfiguredFoldResults(Operation *op,
+ ArrayRef<Attribute> operands,
+ StringRef name) {
+ if (consumeInPlaceStep(op, name))
+ return success();
+ auto config = op->getAttrOfType<ArrayAttr>(name);
+ if (!config)
+ return failure();
+ OpFoldResults results = llvm::map_to_vector(config, [&](Attribute element) {
+ return getConfiguredReplacement(op, operands, element);
+ });
+ std::string withInPlaceName = (name + "_with_in_place").str();
+ if (op->hasAttr(withInPlaceName)) {
+ op->removeAttr(withInPlaceName);
+ results.setModifiedInPlace();
+ }
+ return results;
+}
+
+OpFoldResults TestFoldDispatchOp::fold(FoldAdaptor adaptor) {
+ return getConfiguredFoldResults(getOperation(), adaptor.getOperands(),
+ "fold");
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.h b/mlir/test/lib/Dialect/Test/TestOps.h
index 82c0bd8d437f5..1912db4185b21 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.h
+++ b/mlir/test/lib/Dialect/Test/TestOps.h
@@ -263,6 +263,30 @@ llvm::LogicalResult getConfiguredLegacyFoldResults(
mlir::Operation *op, llvm::ArrayRef<mlir::Attribute> operands,
llvm::StringRef name, llvm::SmallVectorImpl<mlir::OpFoldResult> &results);
+/// Return the fold result in the `OpFoldResults` form that the attributes of
+/// `op` with the prefix `name` describe. The attributes are the ones of
+/// `getConfiguredLegacyFoldResults`. In addition, the unit attribute
+/// `<name>_with_in_place`, which the fold drops, marks the replacements as an
+/// in-place change.
+mlir::OpFoldResults
+getConfiguredFoldResults(mlir::Operation *op,
+ llvm::ArrayRef<mlir::Attribute> operands,
+ llvm::StringRef name);
+
+/// A legacy fold trait for the fold tests. It reads its result from the
+/// `legacy_trait_fold` attributes.
+template <typename ConcreteType>
+class TestLegacyFoldTrait
+ : public mlir::OpTrait::TraitBase<ConcreteType, TestLegacyFoldTrait> {
+public:
+ static llvm::LogicalResult
+ foldTrait(mlir::Operation *op, llvm::ArrayRef<mlir::Attribute> operands,
+ llvm::SmallVectorImpl<mlir::OpFoldResult> &results) {
+ return getConfiguredLegacyFoldResults(op, operands, "legacy_trait_fold",
+ results);
+ }
+};
+
} // namespace test
#define GET_OP_CLASSES
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 1c7f4e32dfa94..6bbfa93f818be 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1809,6 +1809,22 @@ def TestFoldDispatchFallbackOp : TEST_Op<"fold_dispatch_fallback"> {
let results = (outs Variadic<AnyType>);
}
+def TestLegacyFoldTrait : NativeOpTrait<"TestLegacyFoldTrait"> {
+ let cppNamespace = "::test";
+}
+
+// Op for the fold tests in IR/fold-dispatch.mlir. Its fold, in the
+// `OpFoldResults` form, reads its result from the `fold` attributes; see
+// `getConfiguredFoldResults`. Its legacy fold trait reads its result from the
+// `legacy_trait_fold` attributes.
+def TestFoldDispatchOp : TEST_Op<"fold_dispatch", [TestLegacyFoldTrait]> {
+ let arguments = (ins Variadic<AnyType>);
+ let results = (outs Variadic<AnyType>);
+ let extraClassDeclaration = [{
+ ::mlir::OpFoldResults fold(FoldAdaptor adaptor);
+ }];
+}
+
//===----------------------------------------------------------------------===//
// Test Patterns (Symbol Binding)
//===----------------------------------------------------------------------===//
diff --git a/mlir/unittests/IR/OpFoldResultsTest.cpp b/mlir/unittests/IR/OpFoldResultsTest.cpp
index 19d5b88fc7f09..bc3bb99851143 100644
--- a/mlir/unittests/IR/OpFoldResultsTest.cpp
+++ b/mlir/unittests/IR/OpFoldResultsTest.cpp
@@ -21,19 +21,67 @@
using namespace mlir;
+// The fallback TypeID resolver rejects a trait instantiated with a type from
+// an anonymous namespace, so the test dialect lives in a named namespace.
+namespace op_fold_results_test {
+/// Per-test behavior of the op below.
+struct FoldState {
+ std::function<OpFoldResults(Operation *)> opFoldFn;
+};
+
+/// The state of the running test. The test fixture owns it.
+static FoldState *foldState = nullptr;
+
+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> {
+ 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) {
+ return foldState->opFoldFn ? foldState->opFoldFn(getOperation())
+ : failure();
+ }
+};
+
+struct FoldTestDialect : public Dialect {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FoldTestDialect)
+ static constexpr StringLiteral getDialectNamespace() { return "fold_test"; }
+ explicit FoldTestDialect(MLIRContext *context)
+ : Dialect(getDialectNamespace(), context,
+ TypeID::get<FoldTestDialect>()) {
+ addOperations<PartialFoldOp>();
+ }
+};
+} // namespace op_fold_results_test
+
+using namespace op_fold_results_test;
+
namespace {
class OpFoldResultsTest : public ::testing::Test {
protected:
OpFoldResultsTest() : builder(&context) {
context.allowUnregisteredDialects();
+ context.loadDialect<FoldTestDialect>();
i32 = builder.getI32Type();
f32 = builder.getF32Type();
+ foldState = &state;
}
~OpFoldResultsTest() override {
// Destroy users before the ops that define their operands.
for (Operation *op : llvm::reverse(ops))
op->destroy();
+ foldState = nullptr;
}
/// Create an op with the given result types and operands. The fixture
@@ -101,6 +149,7 @@ class OpFoldResultsTest : public ::testing::Test {
std::function<OpFoldResults(Operation *)> foldFn;
std::function<LogicalResult(Operation *, SmallVectorImpl<OpFoldResult> &)>
legacyFoldFn;
+ FoldState state;
};
} // namespace
@@ -582,8 +631,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));
@@ -655,17 +704,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