[llvm-branch-commits] [mlir] [mlir] Add the `OpFoldResults` form to fold traits and `DialectFoldInterface` (PR #228769)
Victor Perez via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Sat Oct 3 15:39:01 PDT 2026
https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/228769
>From f94b1c9e1b861bfcd800fc296810d4f0b1d0fe01 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: 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 <victor.pc.upm at gmail.com>
---
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 b09f51e217f391..74c5d1ee85860f 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 991b084b8b44d5..b163c649e2f335 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 3b81eef8b864bc..c32bc2946bf244 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 138be794f605c0..37a3d4953caa99 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 5a4247d848d183..52be6d164119f3 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 e771d3f194bcc4..dd216e1932d049 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 05c872daf5dab2..7e32cebdf739b8 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 1d1991ba86641e..aad6a6f90095bc 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);
More information about the llvm-branch-commits
mailing list