https://github.com/victor-eds created https://github.com/llvm/llvm-project/pull/228769
The previous patch added `OpFoldResults fold(FoldAdaptor)` for ops. The fold traits and `DialectFoldInterface` still use only the legacy form. A trait can now define `static OpFoldResults foldTrait(Operation *, ArrayRef<Attribute>)`. The legacy trait form still works through an adapter. `IsCommutative` and `CastOpInterface` now use the new form. `DialectFoldInterface` gets a new `fold` method that returns `OpFoldResults`. Its default implementation calls the legacy `fold` method through the adapter, so existing dialects keep their behaviour. `Operation::fold` now calls this new method. >From 85ecfc775c634d4e8b29cc744b6fa5d7f72bd78b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADctor=20P=C3=A9rez=20Carrasco?= <[email protected]> Date: Thu, 1 Oct 2026 08:00:02 -0700 Subject: [PATCH] [mlir] Add the OpFoldResults form to fold traits and DialectFoldInterface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous patch added `OpFoldResults fold(FoldAdaptor)` for ops. The fold traits and `DialectFoldInterface` still use only the legacy vector form, which cannot express a partial fold. A trait can now define `static OpFoldResults foldTrait(Operation *, ArrayRef<Attribute>)`. The trait fold dispatch detects this form by its return type and calls it directly. This form has the same parameters as the single-result trait form, so the single-result detector excludes it. The legacy trait form still works through the strict adapter. `IsCommutative` and `CastOpInterface` now use the new form. `DialectFoldInterface` gets a new `fold` method that returns `OpFoldResults`. Its default implementation calls the legacy `fold` method through the strict adapter, so existing dialects keep their behavior. `Operation::fold` now calls the new method. The rule that a replacement does not name a result that the same fold replaces still applies only to the op's own `OpFoldResults` fold. `CastOpInterface` forwards its operands, and in a graph region an operand can be another result of the same op. Upstream folds that case, so trait and `DialectFoldInterface` results keep that behavior. The new unit tests cover the trait form, the stop at the first trait that folds, the detector, and the new `DialectFoldInterface` method. Signed-off-by: Víctor Pérez Carrasco <[email protected]> --- mlir/include/mlir/IR/OpDefinition.h | 38 ++- mlir/include/mlir/Interfaces/CastInterfaces.h | 5 +- .../include/mlir/Interfaces/CastInterfaces.td | 6 +- .../mlir/Interfaces/DialectFoldInterface.td | 17 ++ mlir/include/mlir/Interfaces/FoldInterfaces.h | 12 +- mlir/lib/IR/Operation.cpp | 18 +- mlir/lib/Interfaces/CastInterfaces.cpp | 11 +- mlir/unittests/IR/OpFoldResultsTest.cpp | 284 +++++++++++++++++- 8 files changed, 358 insertions(+), 33 deletions(-) diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h index b09f51e217f39..74c5d1ee85860 100644 --- a/mlir/include/mlir/IR/OpDefinition.h +++ b/mlir/include/mlir/IR/OpDefinition.h @@ -284,8 +284,7 @@ namespace OpTrait { // corresponding trait classes. This avoids them being template // instantiated/duplicated. namespace impl { -LogicalResult foldCommutative(Operation *op, ArrayRef<Attribute> operands, - SmallVectorImpl<OpFoldResult> &results); +OpFoldResults foldCommutative(Operation *op, ArrayRef<Attribute> operands); OpFoldResult foldIdempotent(Operation *op); OpFoldResult foldInvolution(Operation *op); LogicalResult verifyZeroOperands(Operation *op); @@ -1120,9 +1119,8 @@ class ResultsAreSignlessIntegerLike template <typename ConcreteType> class IsCommutative : public TraitBase<ConcreteType, IsCommutative> { public: - static LogicalResult foldTrait(Operation *op, ArrayRef<Attribute> operands, - SmallVectorImpl<OpFoldResult> &results) { - return impl::foldCommutative(op, operands, results); + static OpFoldResults foldTrait(Operation *op, ArrayRef<Attribute> operands) { + return impl::foldCommutative(op, operands); } }; @@ -1534,14 +1532,26 @@ inline bool hasTrait<>(TypeID traitID) { // Trait Folding //===----------------------------------------------------------------------===// +/// Trait to check if T provides a general 'foldTrait' method that returns +/// OpFoldResults. +template <typename T, typename... Args> +using has_fold_results_trait = std::enable_if_t< + std::is_same_v<decltype(T::foldTrait(std::declval<Operation *>(), + std::declval<ArrayRef<Attribute>>())), + OpFoldResults>>; +template <typename T> +using detect_has_fold_results_trait = + llvm::is_detected<has_fold_results_trait, T>; /// Trait to check if T provides a 'foldTrait' method for single result -/// operations. +/// operations. The OpFoldResults form has the same parameters, so it is +/// excluded. template <typename T, typename... Args> using has_single_result_fold_trait = decltype(T::foldTrait( std::declval<Operation *>(), std::declval<ArrayRef<Attribute>>())); template <typename T> using detect_has_single_result_fold_trait = - llvm::is_detected<has_single_result_fold_trait, T>; + std::conjunction<llvm::is_detected<has_single_result_fold_trait, T>, + std::negation<detect_has_fold_results_trait<T>>>; /// Trait to check if T provides a general 'foldTrait' method. template <typename T, typename... Args> using has_fold_trait = @@ -1553,7 +1563,7 @@ using detect_has_fold_trait = llvm::is_detected<has_fold_trait, T>; /// Trait to check if T provides any `foldTrait` method. template <typename T> using detect_has_any_fold_trait = - std::disjunction<detect_has_fold_trait<T>, + std::disjunction<detect_has_fold_results_trait<T>, detect_has_fold_trait<T>, detect_has_single_result_fold_trait<T>>; /// Returns the result of folding a trait that implements a `foldTrait` function @@ -1573,8 +1583,16 @@ foldTrait(Operation *op, ArrayRef<Attribute> operands) { return success(); return result; } -/// Returns the result of folding a trait that implements the generalized -/// `foldTrait(op, operands, results)` function, with the strict legacy +/// Returns the result of folding a trait that implements a generalized +/// `foldTrait` function that supports any operation type. +template <typename Trait> +std::enable_if_t<detect_has_fold_results_trait<Trait>::value, OpFoldResults> +foldTrait(Operation *op, ArrayRef<Attribute> operands) { + return Trait::foldTrait(op, operands); +} +/// Returns the result of folding a trait that implements the legacy +/// generalized `foldTrait` function. This is the folder for traits that do +/// not implement the `OpFoldResults` form. It keeps the strict legacy /// contract. template <typename Trait> std::enable_if_t<detect_has_fold_trait<Trait>::value, OpFoldResults> diff --git a/mlir/include/mlir/Interfaces/CastInterfaces.h b/mlir/include/mlir/Interfaces/CastInterfaces.h index 991b084b8b44d..b163c649e2f33 100644 --- a/mlir/include/mlir/Interfaces/CastInterfaces.h +++ b/mlir/include/mlir/Interfaces/CastInterfaces.h @@ -21,9 +21,8 @@ class DialectRegistry; namespace impl { /// Attempt to fold the given cast operation. -LogicalResult foldCastInterfaceOp(Operation *op, - ArrayRef<Attribute> attrOperands, - SmallVectorImpl<OpFoldResult> &foldResults); +OpFoldResults foldCastInterfaceOp(Operation *op, + ArrayRef<Attribute> attrOperands); /// Attempt to verify the given cast operation. LogicalResult verifyCastInterfaceOp(Operation *op); diff --git a/mlir/include/mlir/Interfaces/CastInterfaces.td b/mlir/include/mlir/Interfaces/CastInterfaces.td index 3b81eef8b864b..c32bc2946bf24 100644 --- a/mlir/include/mlir/Interfaces/CastInterfaces.td +++ b/mlir/include/mlir/Interfaces/CastInterfaces.td @@ -38,9 +38,9 @@ def CastOpInterface : OpInterface<"CastOpInterface"> { let extraTraitClassDeclaration = [{ /// Attempt to fold the given cast operation. - static LogicalResult foldTrait(Operation *op, ArrayRef<Attribute> operands, - SmallVectorImpl<OpFoldResult> &results) { - return impl::foldCastInterfaceOp(op, operands, results); + static OpFoldResults foldTrait(Operation *op, + ArrayRef<Attribute> operands) { + return impl::foldCastInterfaceOp(op, operands); } }]; let verify = [{ diff --git a/mlir/include/mlir/Interfaces/DialectFoldInterface.td b/mlir/include/mlir/Interfaces/DialectFoldInterface.td index 138be794f605c..37a3d4953caa9 100644 --- a/mlir/include/mlir/Interfaces/DialectFoldInterface.td +++ b/mlir/include/mlir/Interfaces/DialectFoldInterface.td @@ -11,6 +11,23 @@ def DialectFoldInterface : DialectInterface<"DialectFoldInterface"> { let cppNamespace = "::mlir"; let methods = [ + InterfaceMethod<[{ + Registered fallback fold for the dialect. Like the fold hook of each + operation, it attempts to fold the operation with the specified constant + operand values - the elements in "operands" will correspond directly to + the operands of the operation, but may be null if non-constant. The + result has the same contract as the fold hook of an operation. The + default implementation calls the `fold` method with a `results` vector, + with its strict contract. + }], + "::mlir::OpFoldResults", "fold", + (ins "::mlir::Operation *":$op, + "::mlir::ArrayRef<::mlir::Attribute>":$operands), + [{ + return ::mlir::detail::foldWithLegacyDialectFoldInterface( + this, op, operands); + }] + >, InterfaceMethod<[{ Registered fallback fold for the dialect. Like the fold hook of each operation, it attempts to fold the operation with the specified constant diff --git a/mlir/include/mlir/Interfaces/FoldInterfaces.h b/mlir/include/mlir/Interfaces/FoldInterfaces.h index 5a4247d848d18..52be6d164119f 100644 --- a/mlir/include/mlir/Interfaces/FoldInterfaces.h +++ b/mlir/include/mlir/Interfaces/FoldInterfaces.h @@ -9,13 +9,23 @@ #define MLIR_INTERFACES_FOLDINTERFACES_H_ #include "mlir/IR/DialectInterface.h" +#include "mlir/IR/OpFoldResult.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/SmallVector.h" namespace mlir { class Attribute; -class OpFoldResult; +class DialectFoldInterface; +class Operation; class Region; + +namespace detail { +/// Fold `op` with the legacy vector `fold` method of `iface`, with the strict +/// legacy contract. +OpFoldResults +foldWithLegacyDialectFoldInterface(const DialectFoldInterface *iface, + Operation *op, ArrayRef<Attribute> operands); +} // namespace detail } // namespace mlir #include "mlir/Interfaces/DialectFoldInterface.h.inc" diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp index e771d3f194bcc..dd216e1932d04 100644 --- a/mlir/lib/IR/Operation.cpp +++ b/mlir/lib/IR/Operation.cpp @@ -600,6 +600,15 @@ void Operation::setSuccessor(Block *block, unsigned index) { getBlockOperands()[index].set(block); } +OpFoldResults +detail::foldWithLegacyDialectFoldInterface(const DialectFoldInterface *iface, + Operation *op, + ArrayRef<Attribute> operands) { + SmallVector<OpFoldResult> results; + LogicalResult status = iface->fold(op, operands, results); + return convertLegacyFoldResults(status, results); +} + /// Attempt to fold this operation using the Op's registered foldHook. OpFoldResults Operation::fold(ArrayRef<Attribute> operands) { // If we have a registered operation definition matching this one, use it to @@ -618,9 +627,7 @@ OpFoldResults Operation::fold(ArrayRef<Attribute> operands) { if (!interface) return results; - SmallVector<OpFoldResult> legacyResults; - LogicalResult status = interface->fold(this, operands, legacyResults); - results = detail::convertLegacyFoldResults(status, legacyResults); + results = interface->fold(this, operands); results.normalize(this); return results; } @@ -849,9 +856,8 @@ InFlightDiagnostic OpState::emitRemark(const Twine &message) { // Op Trait implementations //===----------------------------------------------------------------------===// -LogicalResult -OpTrait::impl::foldCommutative(Operation *op, ArrayRef<Attribute> operands, - SmallVectorImpl<OpFoldResult> &results) { +OpFoldResults OpTrait::impl::foldCommutative(Operation *op, + ArrayRef<Attribute> operands) { // Nothing to fold if there are not at least 2 operands. if (op->getNumOperands() < 2) return failure(); diff --git a/mlir/lib/Interfaces/CastInterfaces.cpp b/mlir/lib/Interfaces/CastInterfaces.cpp index 05c872daf5dab..7e32cebdf739b 100644 --- a/mlir/lib/Interfaces/CastInterfaces.cpp +++ b/mlir/lib/Interfaces/CastInterfaces.cpp @@ -18,19 +18,16 @@ using namespace mlir; //===----------------------------------------------------------------------===// /// Attempt to fold the given cast operation. -LogicalResult -impl::foldCastInterfaceOp(Operation *op, ArrayRef<Attribute> attrOperands, - SmallVectorImpl<OpFoldResult> &foldResults) { +OpFoldResults impl::foldCastInterfaceOp(Operation *op, + ArrayRef<Attribute> attrOperands) { OperandRange operands = op->getOperands(); if (operands.empty()) return failure(); ResultRange results = op->getResults(); // Check for the case where the input and output types match 1-1. - if (operands.getTypes() == results.getTypes()) { - foldResults.append(operands.begin(), operands.end()); - return success(); - } + if (operands.getTypes() == results.getTypes()) + return operands; return failure(); } diff --git a/mlir/unittests/IR/OpFoldResultsTest.cpp b/mlir/unittests/IR/OpFoldResultsTest.cpp index 1d1991ba86641..aad6a6f90095b 100644 --- a/mlir/unittests/IR/OpFoldResultsTest.cpp +++ b/mlir/unittests/IR/OpFoldResultsTest.cpp @@ -28,14 +28,19 @@ namespace op_fold_results_test { using LegacyFoldFn = std::function<LogicalResult(Operation *, SmallVectorImpl<OpFoldResult> &)>; -/// Per-test behavior of the ops, the trait, and the dialect interface below. +/// Per-test behavior of the ops, traits, and dialect interfaces below. struct FoldState { std::function<OpFoldResults(Operation *)> opFoldFn; LegacyFoldFn traitFoldFn; + std::function<OpFoldResults(Operation *)> resultsTraitFoldFn; LegacyFoldFn dialectFoldFn; + std::function<OpFoldResults(Operation *)> resultsDialectFoldFn; SmallVector<Attribute> seenOperands; unsigned traitCalls = 0; + unsigned resultsTraitCalls = 0; unsigned dialectCalls = 0; + unsigned resultsDialectCalls = 0; + unsigned resultsDialectLegacyCalls = 0; }; /// The state of the running test. The test fixture owns it. @@ -52,6 +57,16 @@ struct LegacyFoldTrait } }; +template <typename ConcreteType> +struct ResultsFoldTrait + : public OpTrait::TraitBase<ConcreteType, ResultsFoldTrait> { + static OpFoldResults foldTrait(Operation *op, ArrayRef<Attribute>) { + ++foldState->resultsTraitCalls; + return foldState->resultsTraitFoldFn ? foldState->resultsTraitFoldFn(op) + : failure(); + } +}; + template <typename OpT> struct FoldAdaptorImpl { FoldAdaptorImpl(ArrayRef<Attribute> operands, OpT) : operands(operands) {} @@ -75,6 +90,22 @@ struct PartialFoldOp : public Op<PartialFoldOp, OpTrait::NResults<2>::Impl, } }; +/// An op with two results and no fold of its own. ResultsFoldTrait comes +/// first, so LegacyFoldTrait runs only if ResultsFoldTrait fails. +struct TraitFoldOp + : public Op<TraitFoldOp, OpTrait::NResults<2>::Impl, + OpTrait::VariadicOperands, ResultsFoldTrait, LegacyFoldTrait> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TraitFoldOp) + using Op::Op; + static ArrayRef<StringRef> getAttributeNames() { return {}; } + static StringRef getOperationName() { return "fold_test.trait_fold"; } +}; + +static_assert(op_definition_impl::detect_has_fold_results_trait< + ResultsFoldTrait<TraitFoldOp>>::value); +static_assert(!op_definition_impl::detect_has_single_result_fold_trait< + ResultsFoldTrait<TraitFoldOp>>::value); + struct ConstantOp : public Op<ConstantOp, OpTrait::OneResult, OpTrait::ZeroOperands, OpTrait::ConstantLike> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConstantOp) @@ -104,10 +135,49 @@ struct FoldTestDialect : public Dialect { explicit FoldTestDialect(MLIRContext *context) : Dialect(getDialectNamespace(), context, TypeID::get<FoldTestDialect>()) { - addOperations<PartialFoldOp, ConstantOp>(); + addOperations<PartialFoldOp, TraitFoldOp, ConstantOp>(); addInterfaces<TestFoldInterface>(); } }; + +/// An op with two results, no fold, and no fold traits. +struct ResultsFallbackOp + : public Op<ResultsFallbackOp, OpTrait::NResults<2>::Impl, + OpTrait::ZeroOperands> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ResultsFallbackOp) + using Op::Op; + static ArrayRef<StringRef> getAttributeNames() { return {}; } + static StringRef getOperationName() { return "results_fold_test.op"; } +}; + +/// Overrides the OpFoldResults form of the fallback. The legacy method reports +/// an in-place fold, so a call to it changes the fold result. +struct TestResultsFoldInterface : public DialectFoldInterface { + using DialectFoldInterface::DialectFoldInterface; + OpFoldResults fold(Operation *op, ArrayRef<Attribute>) const final { + ++foldState->resultsDialectCalls; + return foldState->resultsDialectFoldFn ? foldState->resultsDialectFoldFn(op) + : failure(); + } + LogicalResult fold(Operation *, ArrayRef<Attribute>, + SmallVectorImpl<OpFoldResult> &) const final { + ++foldState->resultsDialectLegacyCalls; + return success(); + } +}; + +struct ResultsFoldTestDialect : public Dialect { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ResultsFoldTestDialect) + static constexpr StringLiteral getDialectNamespace() { + return "results_fold_test"; + } + explicit ResultsFoldTestDialect(MLIRContext *context) + : Dialect(getDialectNamespace(), context, + TypeID::get<ResultsFoldTestDialect>()) { + addOperations<ResultsFallbackOp>(); + addInterfaces<TestResultsFoldInterface>(); + } +}; } // namespace op_fold_results_test using namespace op_fold_results_test; @@ -117,7 +187,7 @@ class OpFoldResultsTest : public ::testing::Test { protected: OpFoldResultsTest() : builder(&context) { context.allowUnregisteredDialects(); - context.loadDialect<FoldTestDialect>(); + context.loadDialect<FoldTestDialect, ResultsFoldTestDialect>(); i32 = builder.getI32Type(); f32 = builder.getF32Type(); foldState = &state; @@ -921,6 +991,154 @@ TEST_F(OpFoldResultsTest, OwnResultsNormalizeToFailureSoTraitsRun) { EXPECT_EQ(foldState->traitCalls, 2u); } +TEST_F(OpFoldResultsTest, ResultsTraitPartialFoldSkipsLaterTraits) { + Operation *op = createOp({i32, i32}, "fold_test.trait_fold"); + Attribute attr = builder.getI32IntegerAttr(1); + foldState->resultsTraitFoldFn = [&](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)); + EXPECT_EQ(foldState->resultsTraitCalls, 1u); + + SmallVector<OpFoldResult> results; + EXPECT_TRUE(failed(op->fold(results))); + EXPECT_TRUE(results.empty()); + EXPECT_EQ(foldState->resultsTraitCalls, 2u); + EXPECT_EQ(foldState->traitCalls, 0u); + EXPECT_EQ(foldState->dialectCalls, 0u); +} + +TEST_F(OpFoldResultsTest, ResultsTraitInPlaceSkipsLaterTraits) { + Operation *op = createOp({i32, i32}, "fold_test.trait_fold"); + Attribute attr = builder.getI32IntegerAttr(1); + foldState->resultsTraitFoldFn = [](Operation *) -> OpFoldResults { + return success(); + }; + foldState->traitFoldFn = [&](Operation *, + SmallVectorImpl<OpFoldResult> &results) { + results.append(2, attr); + return success(); + }; + + OpFoldResults result = op->fold(); + EXPECT_TRUE(result.modifiedInPlace()); + EXPECT_FALSE(result.replacesAny()); + EXPECT_EQ(foldState->resultsTraitCalls, 1u); + + SmallVector<OpFoldResult> results; + EXPECT_TRUE(succeeded(op->fold(results))); + EXPECT_TRUE(results.empty()); + + // A partial fold keeps the in-place bit of the trait. + foldState->resultsTraitFoldFn = [&](Operation *foldedOp) { + OpFoldResults partial(foldedOp); + partial.replace(0u, attr); + partial.setModifiedInPlace(); + return partial; + }; + result = op->fold(); + EXPECT_TRUE(result.modifiedInPlace()); + EXPECT_FALSE(result.replacesAll()); + ASSERT_EQ(result.size(), 2u); + EXPECT_EQ(result[0], OpFoldResult(attr)); + EXPECT_FALSE(result[1]); + EXPECT_TRUE(succeeded(op->fold(results))); + EXPECT_TRUE(results.empty()); + EXPECT_EQ(foldState->traitCalls, 0u); +} + +TEST_F(OpFoldResultsTest, ResultsTraitOwnResultsNormalizeToFailure) { + Operation *op = createOp({i32, i32}, "fold_test.trait_fold"); + Attribute attr = builder.getI32IntegerAttr(1); + foldState->resultsTraitFoldFn = [](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->resultsTraitCalls, 1u); + 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, ResultsTraitFailureThenLegacyTrait) { + Operation *op = createOp({i32, i32}, "fold_test.trait_fold"); + foldState->resultsTraitFoldFn = [](Operation *) -> OpFoldResults { + return failure(); + }; + + // Both traits fail, so the fold falls back on the dialect. + EXPECT_TRUE(op->fold().failed()); + EXPECT_EQ(foldState->resultsTraitCalls, 1u); + EXPECT_EQ(foldState->traitCalls, 1u); + EXPECT_EQ(foldState->dialectCalls, 1u); + + foldState->traitFoldFn = [](Operation *, SmallVectorImpl<OpFoldResult> &) { + return success(); + }; + OpFoldResults result = op->fold(); + EXPECT_TRUE(result.modifiedInPlace()); + EXPECT_FALSE(result.replacesAny()); + EXPECT_EQ(foldState->resultsTraitCalls, 2u); + EXPECT_EQ(foldState->traitCalls, 2u); + + SmallVector<OpFoldResult> results; + EXPECT_TRUE(succeeded(op->fold(results))); + EXPECT_TRUE(results.empty()); + EXPECT_EQ(foldState->dialectCalls, 1u); +} + +TEST_F(OpFoldResultsTest, ResultsTraitMayForwardAnotherResult) { + Operation *op = createOp({i32, i32}, "fold_test.trait_fold"); + Attribute attr = builder.getI32IntegerAttr(1); + // Replacement 0 is result 1, which the fold also replaces. A cast in a graph + // region can return such a chain, so the replaced-result check does not + // apply to trait folds. + foldState->resultsTraitFoldFn = [&](Operation *foldedOp) -> OpFoldResults { + return {foldedOp->getResult(1), attr}; + }; + + OpFoldResults result = op->fold(); + EXPECT_TRUE(result.replacesAll()); + ASSERT_EQ(result.size(), 2u); + EXPECT_EQ(result[0], OpFoldResult(op->getResult(1))); + EXPECT_EQ(result[1], OpFoldResult(attr)); + + SmallVector<OpFoldResult> results; + EXPECT_TRUE(succeeded(op->fold(results))); + ASSERT_EQ(results.size(), 2u); + EXPECT_EQ(results[0], OpFoldResult(op->getResult(1))); + EXPECT_EQ(results[1], OpFoldResult(attr)); + EXPECT_EQ(foldState->traitCalls, 0u); +} + TEST_F(OpFoldResultsTest, DialectFoldInterfaceFallback) { Operation *op = createOp({i32, i32}, "fold_test.partial"); Attribute attr = builder.getI32IntegerAttr(1); @@ -971,6 +1189,66 @@ TEST_F(OpFoldResultsTest, DialectFoldInterfaceFallback) { EXPECT_EQ(foldState->dialectCalls, 1u); } +TEST_F(OpFoldResultsTest, ResultsDialectFoldInterfacePartialFold) { + Operation *op = createOp({i32, i32}, "results_fold_test.op"); + Attribute attr = builder.getI32IntegerAttr(1); + foldState->resultsDialectFoldFn = [&](Operation *foldedOp) { + OpFoldResults partial(foldedOp); + partial.replace(0u, attr); + return partial; + }; + + OpFoldResults result = op->fold(); + EXPECT_TRUE(result.succeeded()); + EXPECT_FALSE(result.modifiedInPlace()); + EXPECT_FALSE(result.replacesAll()); + ASSERT_EQ(result.size(), 2u); + EXPECT_EQ(result[0], OpFoldResult(attr)); + EXPECT_FALSE(result[1]); + EXPECT_EQ(foldState->resultsDialectCalls, 1u); + + SmallVector<OpFoldResult> results; + EXPECT_TRUE(failed(op->fold(results))); + EXPECT_TRUE(results.empty()); + EXPECT_EQ(foldState->resultsDialectCalls, 2u); + EXPECT_EQ(foldState->resultsDialectLegacyCalls, 0u); +} + +TEST_F(OpFoldResultsTest, ResultsDialectFoldInterfaceInPlace) { + Operation *op = createOp({i32, i32}, "results_fold_test.op"); + foldState->resultsDialectFoldFn = [](Operation *) -> OpFoldResults { + return success(); + }; + + OpFoldResults result = op->fold(); + EXPECT_TRUE(result.modifiedInPlace()); + EXPECT_FALSE(result.replacesAny()); + EXPECT_EQ(foldState->resultsDialectCalls, 1u); + + SmallVector<OpFoldResult> results; + EXPECT_TRUE(succeeded(op->fold(results))); + EXPECT_TRUE(results.empty()); + EXPECT_EQ(foldState->resultsDialectLegacyCalls, 0u); +} + +TEST_F(OpFoldResultsTest, ResultsDialectFoldInterfaceSkipsLegacyMethod) { + Operation *op = createOp({i32, i32}, "results_fold_test.op"); + + EXPECT_TRUE(op->fold().failed()); + SmallVector<OpFoldResult> results; + EXPECT_TRUE(failed(op->fold(results))); + EXPECT_TRUE(results.empty()); + EXPECT_EQ(foldState->resultsDialectCalls, 2u); + + // The fallback result is normalized. + foldState->resultsDialectFoldFn = [](Operation *foldedOp) -> OpFoldResults { + return foldedOp->getResults(); + }; + EXPECT_TRUE(op->fold().failed()); + EXPECT_EQ(foldState->resultsDialectCalls, 3u); + EXPECT_EQ(foldState->resultsDialectLegacyCalls, 0u); +} + TEST_F(OpFoldResultsTest, FoldComputesConstantOperands) { Operation *constant = createOp({i32}, "fold_test.constant"); Attribute attr = builder.getI32IntegerAttr(7); _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
