[Mlir-commits] [mlir] [mlir] Add OpFoldResults for partial folds of multi-result ops (PR #227430)
Victor Perez
llvmlistbot at llvm.org
Thu Oct 1 11:00:20 PDT 2026
https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/227430
>From 0e3113a7f0abeceb3ec6a21f430c647f3600f2c0 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: Mon, 28 Sep 2026 23:02:01 -0700
Subject: [PATCH] [mlir] Add OpFoldResults for partial folds of multi-result
ops
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
A fold of a multi-result op can only replace all results or change the
op in place. It cannot replace some results and keep the others. This
patch adds the core API for such a partial fold. The drivers do not use
it yet.
The new header `mlir/IR/OpFoldResult.h` holds `OpFoldResult` and the new
class `OpFoldResults`. An `OpFoldResults` object holds one slot per
result, an in-place bit, and a failure flag. A null slot keeps its
result. `{}` and `failure()` are failures, and `success()` is an
in-place change. `normalize` turns a slot that holds its own result into
a null slot, and it turns a result without a change into a failure.
`normalize` also checks the type of each `Value` slot. This check
replaces `checkFoldResultTypes`.
An op can now define `OpFoldResults fold(FoldAdaptor)`. This form has
the same parameters as the single-result form, so the fold hook
detectors now also check the return type of `fold` and `foldTrait`. The
single-result forms must return a type that converts to OpFoldResult,
and the legacy vector forms must return LogicalResult. `getFoldHookFn`
tries the new form first, then the single-result form, then the legacy
vector form, and then the fold traits. The legacy trait form goes
through a strict adapter.
A `static_assert` rejects each of these cases: an op that defines both
the new form and the legacy vector form, an op with the OneResult trait
that defines the new form, and an op without the OneResult trait that
defines a single-result fold. Before this patch, the dispatch skipped
the last fold without a diagnostic.
The type-erased fold hooks return `OpFoldResults`. The legacy hooks of
`DynamicOpDefinition` and `DialectFoldInterface` stay and go through the
same strict adapter. The new `Operation::fold` overloads return a
normalized `OpFoldResults`. The legacy `Operation::fold` overloads
return the full vector only when the fold replaces all results.
Otherwise they return the in-place state and an empty vector.
The new unit tests in `OpFoldResultsTest.cpp` cover the constructors,
`normalize`, null safety, the legacy adapter, and the debug assert for a
slot that holds another result of the op when that result is replaced.
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
mlir/include/mlir/IR/ExtensibleDialect.h | 19 +-
mlir/include/mlir/IR/OpDefinition.h | 302 ++++++------
mlir/include/mlir/IR/OpFoldResult.h | 156 ++++++
mlir/include/mlir/IR/Operation.h | 32 +-
mlir/include/mlir/IR/OperationSupport.h | 51 +-
mlir/lib/IR/CMakeLists.txt | 1 +
mlir/lib/IR/ExtensibleDialect.cpp | 35 +-
mlir/lib/IR/MLIRContext.cpp | 5 +-
mlir/lib/IR/OpFoldResult.cpp | 158 ++++++
mlir/lib/IR/Operation.cpp | 76 ++-
mlir/unittests/IR/CMakeLists.txt | 1 +
mlir/unittests/IR/OpFoldResultsTest.cpp | 588 +++++++++++++++++++++++
12 files changed, 1203 insertions(+), 221 deletions(-)
create mode 100644 mlir/include/mlir/IR/OpFoldResult.h
create mode 100644 mlir/lib/IR/OpFoldResult.cpp
create mode 100644 mlir/unittests/IR/OpFoldResultsTest.cpp
diff --git a/mlir/include/mlir/IR/ExtensibleDialect.h b/mlir/include/mlir/IR/ExtensibleDialect.h
index 79c07cbefcc5c..f0e221c8c0ce6 100644
--- a/mlir/include/mlir/IR/ExtensibleDialect.h
+++ b/mlir/include/mlir/IR/ExtensibleDialect.h
@@ -440,6 +440,11 @@ class DynamicOpDefinition : public OperationName::Impl {
public:
using GetCanonicalizationPatternsFn =
llvm::unique_function<void(RewritePatternSet &, MLIRContext *) const>;
+ /// The legacy fold hook signature. It has the strict legacy contract:
+ /// failure, success with an empty vector (in place), or success with one
+ /// non-null entry per result.
+ using LegacyFoldHookFn = llvm::unique_function<LogicalResult(
+ Operation *, ArrayRef<Attribute>, SmallVectorImpl<OpFoldResult> &) const>;
/// Create a new op at runtime. The op is registered only after passing it to
/// the dialect using registerDynamicOp.
@@ -462,6 +467,14 @@ class DynamicOpDefinition : public OperationName::Impl {
OperationName::FoldHookFn &&foldHookFn,
GetCanonicalizationPatternsFn &&getCanonicalizationPatternsFn,
OperationName::PopulateDefaultAttrsFn &&populateDefaultAttrsFn);
+ static std::unique_ptr<DynamicOpDefinition>
+ get(StringRef name, ExtensibleDialect *dialect,
+ OperationName::VerifyInvariantsFn &&verifyFn,
+ OperationName::VerifyRegionInvariantsFn &&verifyRegionFn,
+ OperationName::ParseAssemblyFn &&parseFn,
+ OperationName::PrintAssemblyFn &&printFn, LegacyFoldHookFn &&foldHookFn,
+ GetCanonicalizationPatternsFn &&getCanonicalizationPatternsFn,
+ OperationName::PopulateDefaultAttrsFn &&populateDefaultAttrsFn);
/// Returns the op typeID.
TypeID getTypeID() { return typeID; }
@@ -495,6 +508,7 @@ class DynamicOpDefinition : public OperationName::Impl {
void setFoldHookFn(OperationName::FoldHookFn &&foldHook) {
foldHookFn = std::move(foldHook);
}
+ void setFoldHookFn(LegacyFoldHookFn &&foldHook);
/// Set the hook returning any canonicalization pattern rewrites that the op
/// supports, for use by the canonicalization pass.
@@ -514,9 +528,8 @@ class DynamicOpDefinition : public OperationName::Impl {
return traits.insert(std::move(trait));
}
- LogicalResult foldHook(Operation *op, ArrayRef<Attribute> attrs,
- SmallVectorImpl<OpFoldResult> &results) final {
- return foldHookFn(op, attrs, results);
+ OpFoldResults foldHook(Operation *op, ArrayRef<Attribute> attrs) final {
+ return foldHookFn(op, attrs);
}
void getCanonicalizationPatterns(RewritePatternSet &set,
MLIRContext *context) final {
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index 0ab6679187c8e..a62c36467a650 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -21,6 +21,7 @@
#include "mlir/IR/Dialect.h"
#include "mlir/IR/ODSSupport.h"
+#include "mlir/IR/OpFoldResult.h"
#include "mlir/IR/Operation.h"
#include "llvm/Support/PointerLikeTypeTraits.h"
@@ -267,51 +268,6 @@ inline bool operator!=(OpState lhs, OpState rhs) {
return lhs.getOperation() != rhs.getOperation();
}
-raw_ostream &operator<<(raw_ostream &os, OpFoldResult ofr);
-
-/// This class represents a single result from folding an operation.
-class OpFoldResult : public PointerUnion<Attribute, Value> {
- using PointerUnion<Attribute, Value>::PointerUnion;
-
-public:
- LLVM_DUMP_METHOD void dump() const { llvm::errs() << *this << "\n"; }
-
- MLIRContext *getContext() const {
- PointerUnion pu = *this;
- return isa<Attribute>(pu) ? cast<Attribute>(pu).getContext()
- : cast<Value>(pu).getContext();
- }
-};
-
-// Temporarily exit the MLIR namespace to add casting support as later code in
-// this uses it. The CastInfo must come after the OpFoldResult definition and
-// before any cast function calls depending on CastInfo.
-
-} // namespace mlir
-
-namespace llvm {
-
-// Allow llvm::cast style functions.
-template <typename To>
-struct CastInfo<To, mlir::OpFoldResult>
- : public CastInfo<To, mlir::OpFoldResult::PointerUnion> {};
-
-template <typename To>
-struct CastInfo<To, const mlir::OpFoldResult>
- : public CastInfo<To, const mlir::OpFoldResult::PointerUnion> {};
-
-} // namespace llvm
-
-namespace mlir {
-
-/// Allow printing to a stream.
-inline raw_ostream &operator<<(raw_ostream &os, OpFoldResult ofr) {
- if (Value value = llvm::dyn_cast_if_present<Value>(ofr))
- value.print(os);
- else
- llvm::dyn_cast_if_present<Attribute>(ofr).print(os);
- return os;
-}
/// Allow printing to a stream.
inline raw_ostream &operator<<(raw_ostream &os, OpState op) {
op.print(os, OpPrintingFlags().useLocalScope());
@@ -1579,19 +1535,23 @@ inline bool hasTrait<>(TypeID traitID) {
//===----------------------------------------------------------------------===//
/// Trait to check if T provides a 'foldTrait' method for single result
-/// operations.
+/// operations that returns a type convertible to OpFoldResult.
template <typename T, typename... Args>
-using has_single_result_fold_trait = decltype(T::foldTrait(
- std::declval<Operation *>(), std::declval<ArrayRef<Attribute>>()));
+using has_single_result_fold_trait = std::enable_if_t<std::is_convertible_v<
+ decltype(T::foldTrait(std::declval<Operation *>(),
+ std::declval<ArrayRef<Attribute>>())),
+ OpFoldResult>>;
template <typename T>
using detect_has_single_result_fold_trait =
llvm::is_detected<has_single_result_fold_trait, T>;
-/// Trait to check if T provides a general 'foldTrait' method.
+/// Trait to check if T provides a general 'foldTrait' method that returns
+/// LogicalResult.
template <typename T, typename... Args>
-using has_fold_trait =
+using has_fold_trait = std::enable_if_t<std::is_same_v<
decltype(T::foldTrait(std::declval<Operation *>(),
std::declval<ArrayRef<Attribute>>(),
- std::declval<SmallVectorImpl<OpFoldResult> &>()));
+ std::declval<SmallVectorImpl<OpFoldResult> &>())),
+ LogicalResult>>;
template <typename T>
using detect_has_fold_trait = llvm::is_detected<has_fold_trait, T>;
/// Trait to check if T provides any `foldTrait` method.
@@ -1604,46 +1564,66 @@ using detect_has_any_fold_trait =
/// that is specialized for operations that have a single result.
template <typename Trait>
std::enable_if_t<detect_has_single_result_fold_trait<Trait>::value,
- LogicalResult>
-foldTrait(Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
+ OpFoldResults>
+foldTrait(Operation *op, ArrayRef<Attribute> operands) {
assert(op->hasTrait<OpTrait::OneResult>() &&
"expected trait on non single-result operation to implement the "
"general `foldTrait` method");
- // If a previous trait has already been folded and replaced this operation, we
- // fail to fold this trait.
- if (!results.empty())
+ OpFoldResult result = Trait::foldTrait(op, operands);
+ if (!result)
return failure();
-
- if (OpFoldResult result = Trait::foldTrait(op, operands)) {
- if (llvm::dyn_cast_if_present<Value>(result) != op->getResult(0))
- results.push_back(result);
+ // The op's own result means that the trait changed the op in place.
+ if (llvm::dyn_cast_if_present<Value>(result) == op->getResult(0))
return success();
- }
- return failure();
+ return result;
}
-/// Returns the result of folding a trait that implements a generalized
-/// `foldTrait` function that is supports any operation type.
+/// 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, LogicalResult>
-foldTrait(Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
- // If a previous trait has already been folded and replaced this operation, we
- // fail to fold this trait.
- return results.empty() ? Trait::foldTrait(op, operands, results) : failure();
+std::enable_if_t<detect_has_fold_trait<Trait>::value, OpFoldResults>
+foldTrait(Operation *op, ArrayRef<Attribute> operands) {
+ SmallVector<OpFoldResult, 2> results;
+ LogicalResult status = Trait::foldTrait(op, operands, results);
+ return ::mlir::detail::convertLegacyFoldResults(status, results);
}
template <typename Trait>
-inline std::enable_if_t<!detect_has_any_fold_trait<Trait>::value, LogicalResult>
-foldTrait(Operation *, ArrayRef<Attribute>, SmallVectorImpl<OpFoldResult> &) {
+inline std::enable_if_t<!detect_has_any_fold_trait<Trait>::value, OpFoldResults>
+foldTrait(Operation *, ArrayRef<Attribute>) {
return failure();
}
+/// Fold `op` with the trait `Trait` into `result`, and return true if the fold
+/// did not fail.
+template <typename Trait>
+bool foldTraitInto(Operation *op, ArrayRef<Attribute> operands,
+ OpFoldResults &result) {
+ result = foldTrait<Trait>(op, operands);
+ result.normalize(op);
+ return result.succeeded();
+}
+
/// Given a tuple type containing a set of traits, return the result of folding
-/// the given operation.
+/// the given operation. Folding stops at the first trait that does not fail.
template <typename... Ts>
-LogicalResult foldTraits(Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
- return success((succeeded(foldTrait<Ts>(op, operands, results)) || ...));
+OpFoldResults foldTraits(Operation *op, ArrayRef<Attribute> operands) {
+ OpFoldResults result;
+ (void)(foldTraitInto<Ts>(op, operands, result) || ...);
+ return result;
+}
+
+/// Combine the result `own` of the op's own fold with the result of its traits.
+/// The traits run only when the op's own fold replaced no result.
+template <typename... Ts>
+OpFoldResults foldTraitsAfterOp(Operation *op, ArrayRef<Attribute> operands,
+ OpFoldResults own) {
+ if (own.replacesAny())
+ return own;
+ OpFoldResults trait = foldTraits<Ts...>(op, operands);
+ OpFoldResults &result = trait.replacesAny() ? trait : own;
+ result.markModifiedInPlace(own.modifiedInPlace() || trait.modifiedInPlace());
+ return std::move(result);
}
//===----------------------------------------------------------------------===//
@@ -1811,36 +1791,55 @@ class Op : public OpState, public Traits<ConcreteType>... {
}
private:
- /// Trait to check if T provides a 'fold' method for a single result op.
+ /// Trait to check if T provides a 'fold' method for a single result op that
+ /// returns a type convertible to OpFoldResult.
template <typename T, typename... Args>
- using has_single_result_fold_t =
- decltype(std::declval<T>().fold(std::declval<ArrayRef<Attribute>>()));
+ using has_single_result_fold_t = std::enable_if_t<std::is_convertible_v<
+ decltype(std::declval<T>().fold(std::declval<ArrayRef<Attribute>>())),
+ OpFoldResult>>;
template <typename T>
constexpr static bool has_single_result_fold_v =
llvm::is_detected<has_single_result_fold_t, T>::value;
- /// Trait to check if T provides a general 'fold' method.
+ /// Trait to check if T provides a general 'fold' method that returns
+ /// LogicalResult.
template <typename T, typename... Args>
- using has_fold_t = decltype(std::declval<T>().fold(
- std::declval<ArrayRef<Attribute>>(),
- std::declval<SmallVectorImpl<OpFoldResult> &>()));
+ using has_fold_t = std::enable_if_t<
+ std::is_same_v<decltype(std::declval<T>().fold(
+ std::declval<ArrayRef<Attribute>>(),
+ std::declval<SmallVectorImpl<OpFoldResult> &>())),
+ LogicalResult>>;
template <typename T>
constexpr static bool has_fold_v = llvm::is_detected<has_fold_t, T>::value;
/// Trait to check if T provides a 'fold' method with a FoldAdaptor for a
- /// single result op.
+ /// single result op that returns a type convertible to OpFoldResult.
template <typename T, typename... Args>
- using has_fold_adaptor_single_result_fold_t =
- decltype(std::declval<T>().fold(std::declval<typename T::FoldAdaptor>()));
+ using has_fold_adaptor_single_result_fold_t = std::enable_if_t<
+ std::is_convertible_v<decltype(std::declval<T>().fold(
+ std::declval<typename T::FoldAdaptor>())),
+ OpFoldResult>>;
template <class T>
constexpr static bool has_fold_adaptor_single_result_v =
llvm::is_detected<has_fold_adaptor_single_result_fold_t, T>::value;
- /// Trait to check if T provides a general 'fold' method with a FoldAdaptor.
+ /// Trait to check if T provides a general 'fold' method with a FoldAdaptor
+ /// that returns LogicalResult.
template <typename T, typename... Args>
- using has_fold_adaptor_fold_t = decltype(std::declval<T>().fold(
- std::declval<typename T::FoldAdaptor>(),
- std::declval<SmallVectorImpl<OpFoldResult> &>()));
+ using has_fold_adaptor_fold_t = std::enable_if_t<
+ std::is_same_v<decltype(std::declval<T>().fold(
+ std::declval<typename T::FoldAdaptor>(),
+ std::declval<SmallVectorImpl<OpFoldResult> &>())),
+ LogicalResult>>;
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>
@@ -1933,43 +1932,73 @@ class Op : public OpState, public Traits<ConcreteType>... {
return detail::InterfaceMap::template get<Traits<ConcreteType>...>();
}
- using FoldHookFn = LogicalResult (*)(Operation *, ArrayRef<Attribute>,
- SmallVectorImpl<OpFoldResult> &);
+ using FoldHookFn = OpFoldResults (*)(Operation *, ArrayRef<Attribute>);
using HasTraitFn = bool (*)(TypeID);
using PrintAssemblyFn = void (*)(Operation *, OpAsmPrinter &, StringRef);
using VerifyInvariantsFn = LogicalResult (*)(Operation *);
/// Return the internal implementations of each of the OperationName hooks.
static constexpr FoldHookFn getFoldHookFn() {
+ constexpr bool hasOneResult =
+ llvm::is_one_of<OpTrait::OneResult<ConcreteType>,
+ Traits<ConcreteType>...>::value;
+ constexpr bool hasResultsFold = has_fold_adaptor_results_v<ConcreteType>;
+ constexpr bool hasSingleResultFold =
+ 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 the multi-result fold form");
+ 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 (llvm::is_one_of<OpTrait::OneResult<ConcreteType>,
- Traits<ConcreteType>...>::value &&
- (has_single_result_fold_v<ConcreteType> ||
- has_fold_adaptor_single_result_v<ConcreteType>))
- return [](Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
- return foldSingleResultHook<ConcreteType>(op, operands, results);
+ 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>)
- return [](Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
- return foldHook<ConcreteType>(op, operands, results);
+ // The operation defines a legacy multi-result `fold` method.
+ if constexpr (hasLegacyFold)
+ return [](Operation *op, ArrayRef<Attribute> operands) {
+ return foldLegacyHook<ConcreteType>(op, operands);
};
// The operation does not define a `fold` method.
- return [](Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
+ return [](Operation *op, ArrayRef<Attribute> operands) {
// In this case, we only need to fold the traits of the operation.
- return op_definition_impl::foldTraits<Traits<ConcreteType>...>(
- op, operands, results);
+ return op_definition_impl::foldTraits<Traits<ConcreteType>...>(op,
+ 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);
+ return op_definition_impl::foldTraitsAfterOp<Traits<ConcreteType>...>(
+ op, operands, std::move(result));
+ }
/// Return the result of folding a single result operation that defines a
/// `fold` method.
template <typename ConcreteOpT>
- static LogicalResult
- foldSingleResultHook(Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
+ static OpFoldResults foldSingleResultHook(Operation *op,
+ ArrayRef<Attribute> operands) {
OpFoldResult result;
if constexpr (has_fold_adaptor_single_result_v<ConcreteOpT>) {
result = cast<ConcreteOpT>(op).fold(
@@ -1978,39 +2007,36 @@ class Op : public OpState, public Traits<ConcreteType>... {
result = cast<ConcreteOpT>(op).fold(operands);
}
- // If the fold failed or was in-place, try to fold the traits of the
- // operation.
- if (!result ||
- llvm::dyn_cast_if_present<Value>(result) == op->getResult(0)) {
- if (succeeded(op_definition_impl::foldTraits<Traits<ConcreteType>...>(
- op, operands, results)))
- return success();
- return success(static_cast<bool>(result));
- }
- results.push_back(result);
- return success();
- }
- /// Return the result of folding an operation that defines a `fold` method.
+ // A null result is a failure. The op's own result means that the fold
+ // changed the op in place.
+ OpFoldResults own =
+ llvm::dyn_cast_if_present<Value>(result) == op->getResult(0)
+ ? OpFoldResults(success())
+ : OpFoldResults(result);
+ own.normalize(op);
+ return op_definition_impl::foldTraitsAfterOp<Traits<ConcreteType>...>(
+ op, operands, std::move(own));
+ }
+ /// Return the result of folding an operation that defines a legacy
+ /// multi-result `fold` method, with the strict legacy contract.
template <typename ConcreteOpT>
- static LogicalResult foldHook(Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
- auto result = LogicalResult::failure();
+ static OpFoldResults foldLegacyHook(Operation *op,
+ ArrayRef<Attribute> operands) {
+ SmallVector<OpFoldResult, 2> results;
+ auto status = LogicalResult::failure();
if constexpr (has_fold_adaptor_v<ConcreteOpT>) {
- result = cast<ConcreteOpT>(op).fold(
+ status = cast<ConcreteOpT>(op).fold(
typename ConcreteOpT::FoldAdaptor(operands, cast<ConcreteOpT>(op)),
results);
} else {
- result = cast<ConcreteOpT>(op).fold(operands, results);
+ status = cast<ConcreteOpT>(op).fold(operands, results);
}
- // If the fold failed or was in-place, try to fold the traits of the
- // operation.
- if (failed(result) || results.empty()) {
- if (succeeded(op_definition_impl::foldTraits<Traits<ConcreteType>...>(
- op, operands, results)))
- return success();
- }
- return result;
+ OpFoldResults own =
+ ::mlir::detail::convertLegacyFoldResults(status, results);
+ own.normalize(op);
+ return op_definition_impl::foldTraitsAfterOp<Traits<ConcreteType>...>(
+ op, operands, std::move(own));
}
static constexpr HasTraitFn getHasTraitFn() {
diff --git a/mlir/include/mlir/IR/OpFoldResult.h b/mlir/include/mlir/IR/OpFoldResult.h
new file mode 100644
index 0000000000000..4af7c54690b21
--- /dev/null
+++ b/mlir/include/mlir/IR/OpFoldResult.h
@@ -0,0 +1,156 @@
+//===- OpFoldResult.h - Results of an operation fold ------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines OpFoldResult, the result of a fold of an operation with
+// exactly one result, and OpFoldResults, the result of a fold of any other
+// operation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_OPFOLDRESULT_H
+#define MLIR_IR_OPFOLDRESULT_H
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/PointerUnion.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/raw_ostream.h"
+#include <cstddef>
+#include <initializer_list>
+#include <type_traits>
+
+namespace mlir {
+class OpFoldResult;
+class Operation;
+
+raw_ostream &operator<<(raw_ostream &os, OpFoldResult ofr);
+
+/// This class represents a single result from folding an operation.
+class OpFoldResult : public PointerUnion<Attribute, Value> {
+ using PointerUnion<Attribute, Value>::PointerUnion;
+
+public:
+ LLVM_DUMP_METHOD void dump() const { llvm::errs() << *this << "\n"; }
+
+ MLIRContext *getContext() const {
+ PointerUnion pu = *this;
+ return isa<Attribute>(pu) ? cast<Attribute>(pu).getContext()
+ : cast<Value>(pu).getContext();
+ }
+};
+
+// Temporarily exit the MLIR namespace to add casting support as later code in
+// this uses it. The CastInfo must come after the OpFoldResult definition and
+// before any cast function calls depending on CastInfo.
+
+} // namespace mlir
+
+namespace llvm {
+
+// Allow llvm::cast style functions.
+template <typename To>
+struct CastInfo<To, mlir::OpFoldResult>
+ : public CastInfo<To, mlir::OpFoldResult::PointerUnion> {};
+
+template <typename To>
+struct CastInfo<To, const mlir::OpFoldResult>
+ : public CastInfo<To, const mlir::OpFoldResult::PointerUnion> {};
+
+} // namespace llvm
+
+namespace mlir {
+
+/// Allow printing to a stream.
+inline raw_ostream &operator<<(raw_ostream &os, OpFoldResult ofr) {
+ if (Value value = llvm::dyn_cast_if_present<Value>(ofr))
+ value.print(os);
+ else
+ llvm::dyn_cast_if_present<Attribute>(ofr).print(os);
+ return os;
+}
+
+/// The result of a fold of an op that does not have exactly one fixed result.
+/// Slot i holds 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.
+ OpFoldResults() = default;
+ OpFoldResults(std::nullptr_t);
+ /// success(): the op changed in place. failure(): failure.
+ OpFoldResults(LogicalResult status);
+ /// One slot. Valid only when the op has exactly one result at run time.
+ /// A null replacement is a failure.
+ OpFoldResults(OpFoldResult replacement);
+ OpFoldResults(Value replacement);
+ OpFoldResults(Attribute replacement);
+ /// One slot per result.
+ OpFoldResults(std::initializer_list<OpFoldResult> replacements);
+ /// One slot per range element. An empty range is a failure.
+ template <typename RangeT,
+ typename = std::enable_if_t<
+ !std::is_convertible_v<RangeT, Attribute> &&
+ !std::is_convertible_v<RangeT, Value> &&
+ std::is_convertible_v<llvm::detail::ValueOfRange<RangeT>,
+ OpFoldResult>>>
+ OpFoldResults(RangeT &&replacements) {
+ llvm::append_range(slots, replacements);
+ isFailure = llvm::none_of(
+ slots, [](OpFoldResult slot) { return static_cast<bool>(slot); });
+ }
+ /// Incremental form: every slot is kept and the op is not changed in place.
+ explicit OpFoldResults(Operation *op);
+
+ /// Set the slot of `result`, which must be a result of the op. A null
+ /// replacement, or `replacement == result`, keeps the slot. The last write
+ /// wins.
+ void replace(Value result, OpFoldResult replacement);
+ void replace(unsigned resultIndex, OpFoldResult replacement);
+ /// Record that the fold changed the op in place (operands, attributes,
+ /// properties, or regions).
+ void markModifiedInPlace(bool modified = true);
+
+ /// Queries for drivers. They are valid after normalize().
+ bool succeeded() const; // modifiedInPlace() || replacesAny()
+ bool failed() const;
+ bool modifiedInPlace() const;
+ bool replacesAny() const;
+ /// replacesAny() and every slot is non-null. False when size() == 0, so an
+ /// in-place fold of a zero-result op never erases the op.
+ bool replacesAll() const;
+ unsigned size() const; // 0 or op->getNumResults()
+ OpFoldResult operator[](unsigned i) const; // Null means "keep".
+ ArrayRef<OpFoldResult> getSlots() const;
+
+ /// Expand a one-slot or list form to op->getNumResults() slots, map
+ /// own-result entries to null, and collapse "nothing replaced and not in
+ /// place" to failure. In debug builds, check the slot count, the Value
+ /// types, and rule D17. Idempotent.
+ void normalize(Operation *op);
+
+private:
+ SmallVector<OpFoldResult, 2> slots;
+ Operation *op = nullptr;
+ bool inPlace = false;
+ bool isFailure = true;
+};
+
+namespace detail {
+/// Convert the result of a legacy vector fold with the strict legacy contract:
+/// failure stays failure, success with an empty vector means "in place", and
+/// success with a full vector replaces every result. A null entry asserts.
+OpFoldResults convertLegacyFoldResults(LogicalResult status,
+ ArrayRef<OpFoldResult> results);
+} // namespace detail
+
+} // namespace mlir
+
+#endif // MLIR_IR_OPFOLDRESULT_H
diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 3ea9ad8ca0cfc..4719054bec455 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -771,21 +771,35 @@ class alignas(8) Operation final
/// - the elements in "operands" will correspond directly to the operands of
/// the operation, but may be null if non-constant.
///
- /// If folding was successful, this function returns "success".
+ /// The result may replace some, all, or none of the results of this
+ /// operation, and may also record that this operation was modified in place.
+ /// The result is normalized: it has either no slot or one slot per result,
+ /// and a null slot means "keep the result". See `OperationName::foldHook`
+ /// for the contract.
+ OpFoldResults fold(ArrayRef<Attribute> operands);
+
+ /// Attempt to fold this operation. Same as above, but computes the constant
+ /// operand values.
+ OpFoldResults fold();
+
+ /// Attempt to fold this 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.
+ ///
+ /// This legacy form has a strict contract. If folding was successful, this
+ /// function returns "success".
/// * If this operation was modified in-place (but not folded away),
/// `results` is empty.
/// * Otherwise, `results` is filled with the folded results.
- /// If folding was unsuccessful, this function returns "failure".
+ /// If folding was unsuccessful, this function returns "failure". A fold that
+ /// replaces only some results is reported as in place if this operation was
+ /// modified in place, and otherwise as a failure. The replacements of such a
+ /// fold are not applied, so the IR is unchanged.
LogicalResult fold(ArrayRef<Attribute> operands,
SmallVectorImpl<OpFoldResult> &results);
- /// Attempt to fold this operation.
- ///
- /// If folding was successful, this function returns "success".
- /// * If this operation was modified in-place (but not folded away),
- /// `results` is empty.
- /// * Otherwise, `results` is filled with the folded results.
- /// If folding was unsuccessful, this function returns "failure".
+ /// Attempt to fold this operation. Same as above, but computes the constant
+ /// operand values.
LogicalResult fold(SmallVectorImpl<OpFoldResult> &results);
/// Returns true if `InterfaceT` has been promised by the dialect or
diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index 991d9847672e2..ccfd693c8c171 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -20,6 +20,7 @@
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/DialectRegistry.h"
#include "mlir/IR/Location.h"
+#include "mlir/IR/OpFoldResult.h"
#include "mlir/IR/TypeRange.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
@@ -52,7 +53,6 @@ class OpAsmParser;
class OpAsmPrinter;
class OperandRange;
class OperandRangeRange;
-class OpFoldResult;
class Pattern;
class Region;
class ResultRange;
@@ -131,8 +131,8 @@ class PropertyRef {
class OperationName {
public:
- using FoldHookFn = llvm::unique_function<LogicalResult(
- Operation *, ArrayRef<Attribute>, SmallVectorImpl<OpFoldResult> &) const>;
+ using FoldHookFn = llvm::unique_function<OpFoldResults(
+ Operation *, ArrayRef<Attribute>) const>;
using HasTraitFn = llvm::unique_function<bool(TypeID) const>;
using ParseAssemblyFn =
llvm::unique_function<ParseResult(OpAsmParser &, OperationState &)>;
@@ -154,8 +154,7 @@ class OperationName {
/// may not be populated.
struct InterfaceConcept {
virtual ~InterfaceConcept() = default;
- virtual LogicalResult foldHook(Operation *, ArrayRef<Attribute>,
- SmallVectorImpl<OpFoldResult> &) = 0;
+ virtual OpFoldResults foldHook(Operation *, ArrayRef<Attribute>) = 0;
virtual void getCanonicalizationPatterns(RewritePatternSet &,
MLIRContext *) = 0;
virtual bool hasTrait(TypeID) = 0;
@@ -247,8 +246,7 @@ class OperationName {
: Impl(name, dialect, typeID, std::move(interfaceMap)) {
propertiesTypeID = TypeID::get<Attribute>();
}
- LogicalResult foldHook(Operation *, ArrayRef<Attribute>,
- SmallVectorImpl<OpFoldResult> &) final;
+ OpFoldResults foldHook(Operation *, ArrayRef<Attribute>) final;
void getCanonicalizationPatterns(RewritePatternSet &, MLIRContext *) final;
bool hasTrait(TypeID) final;
OperationName::ParseAssemblyFn getParseAssemblyFn() final;
@@ -297,24 +295,32 @@ class OperationName {
/// can implement this to provide simplifications rules that are applied by
/// the Builder::createOrFold API and the canonicalization pass.
///
- /// This is an intentionally limited interface - implementations of this
- /// hook can only perform the following changes to the operation:
+ /// This is an intentionally limited interface. The returned OpFoldResults
+ /// holds one slot per result of the operation, and an in-place bit:
///
- /// 1. They can leave the operation alone and without changing the IR, and
- /// return failure.
- /// 2. They can mutate the operation in place, without changing anything
- /// else in the IR. In this case, return success.
- /// 3. They can return a list of existing values that can be used instead
- /// of the operation. In this case, fill in the results list and return
- /// success. The caller will remove the operation and use those results
- /// instead.
+ /// 1. A slot holds 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). Slot i may hold result j only if slot j
+ /// keeps its result.
+ /// 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
+ /// else in the IR. In this case, it marks the result as modified in
+ /// place. The operation must still verify.
+ /// 4. The hook can replace some results and also mutate the operation in
+ /// place. The operation stays.
+ /// 5. If every slot is replaced, the caller removes the operation and uses
+ /// the replacements instead.
+ ///
+ /// The hook creates no operations and changes no IR outside the operation. A
+ /// replacement Value must exist before the fold, must have the type of the
+ /// result that it replaces, and must dominate the operation.
///
/// This allows expression of some simple in-place canonicalizations (e.g.
/// "x+0 -> x", "min(x,y,x,z) -> min(x,y,z)", "x+y-x -> y", etc), as well as
/// generalized constant folding.
- LogicalResult foldHook(Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) const {
- return getImpl()->foldHook(op, operands, results);
+ OpFoldResults foldHook(Operation *op, ArrayRef<Attribute> operands) const {
+ return getImpl()->foldHook(op, operands);
}
/// This hook returns any canonicalization pattern rewrites that the
@@ -600,9 +606,8 @@ class RegisteredOperationName : public OperationName {
TypeID::get<ConcreteOp>(), ConcreteOp::getInterfaceMap()) {
propertiesTypeID = TypeID::get<Properties>();
}
- LogicalResult foldHook(Operation *op, ArrayRef<Attribute> attrs,
- SmallVectorImpl<OpFoldResult> &results) final {
- return ConcreteOp::getFoldHookFn()(op, attrs, results);
+ OpFoldResults foldHook(Operation *op, ArrayRef<Attribute> attrs) final {
+ return ConcreteOp::getFoldHookFn()(op, attrs);
}
void getCanonicalizationPatterns(RewritePatternSet &set,
MLIRContext *context) final {
diff --git a/mlir/lib/IR/CMakeLists.txt b/mlir/lib/IR/CMakeLists.txt
index 01c80dafcd21a..2578cc086c3c6 100644
--- a/mlir/lib/IR/CMakeLists.txt
+++ b/mlir/lib/IR/CMakeLists.txt
@@ -27,6 +27,7 @@ add_mlir_library(MLIRIR
Location.cpp
MLIRContext.cpp
ODSSupport.cpp
+ OpFoldResult.cpp
Operation.cpp
OperationSupport.cpp
PatternLoggingListener.cpp
diff --git a/mlir/lib/IR/ExtensibleDialect.cpp b/mlir/lib/IR/ExtensibleDialect.cpp
index b51d330160646..1e85653892c46 100644
--- a/mlir/lib/IR/ExtensibleDialect.cpp
+++ b/mlir/lib/IR/ExtensibleDialect.cpp
@@ -285,6 +285,17 @@ void DynamicAttr::print(AsmPrinter &printer) {
// Dynamic operation
//===----------------------------------------------------------------------===//
+/// Wrap a legacy fold hook into a fold hook that returns OpFoldResults.
+static OperationName::FoldHookFn
+adaptLegacyFoldHookFn(DynamicOpDefinition::LegacyFoldHookFn &&foldHookFn) {
+ return [foldHookFn = std::move(foldHookFn)](
+ Operation *op, ArrayRef<Attribute> operands) -> OpFoldResults {
+ SmallVector<OpFoldResult> results;
+ LogicalResult status = foldHookFn(op, operands, results);
+ return detail::convertLegacyFoldResults(status, results);
+ };
+}
+
DynamicOpDefinition::DynamicOpDefinition(
StringRef name, ExtensibleDialect *dialect,
OperationName::VerifyInvariantsFn &&verifyFn,
@@ -333,8 +344,8 @@ std::unique_ptr<DynamicOpDefinition> DynamicOpDefinition::get(
OperationName::VerifyRegionInvariantsFn &&verifyRegionFn,
OperationName::ParseAssemblyFn &&parseFn,
OperationName::PrintAssemblyFn &&printFn) {
- auto foldHookFn = [](Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
+ auto foldHookFn = [](Operation *op,
+ ArrayRef<Attribute> operands) -> OpFoldResults {
return failure();
};
@@ -366,6 +377,26 @@ std::unique_ptr<DynamicOpDefinition> DynamicOpDefinition::get(
std::move(populateDefaultAttrsFn)));
}
+std::unique_ptr<DynamicOpDefinition> DynamicOpDefinition::get(
+ StringRef name, ExtensibleDialect *dialect,
+ OperationName::VerifyInvariantsFn &&verifyFn,
+ OperationName::VerifyRegionInvariantsFn &&verifyRegionFn,
+ OperationName::ParseAssemblyFn &&parseFn,
+ OperationName::PrintAssemblyFn &&printFn, LegacyFoldHookFn &&foldHookFn,
+ GetCanonicalizationPatternsFn &&getCanonicalizationPatternsFn,
+ OperationName::PopulateDefaultAttrsFn &&populateDefaultAttrsFn) {
+ return DynamicOpDefinition::get(name, dialect, std::move(verifyFn),
+ std::move(verifyRegionFn), std::move(parseFn),
+ std::move(printFn),
+ adaptLegacyFoldHookFn(std::move(foldHookFn)),
+ std::move(getCanonicalizationPatternsFn),
+ std::move(populateDefaultAttrsFn));
+}
+
+void DynamicOpDefinition::setFoldHookFn(LegacyFoldHookFn &&foldHook) {
+ foldHookFn = adaptLegacyFoldHookFn(std::move(foldHook));
+}
+
//===----------------------------------------------------------------------===//
// Extensible dialect
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 94cf420070131..f64240bb0a583 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -949,9 +949,8 @@ StringRef OperationName::getDialectNamespace() const {
return getStringRef().split('.').first;
}
-LogicalResult
-OperationName::UnregisteredOpModel::foldHook(Operation *, ArrayRef<Attribute>,
- SmallVectorImpl<OpFoldResult> &) {
+OpFoldResults
+OperationName::UnregisteredOpModel::foldHook(Operation *, ArrayRef<Attribute>) {
return failure();
}
void OperationName::UnregisteredOpModel::getCanonicalizationPatterns(
diff --git a/mlir/lib/IR/OpFoldResult.cpp b/mlir/lib/IR/OpFoldResult.cpp
new file mode 100644
index 0000000000000..a0c108963c73b
--- /dev/null
+++ b/mlir/lib/IR/OpFoldResult.cpp
@@ -0,0 +1,158 @@
+//===- OpFoldResult.cpp - Results of an operation fold --------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/OpFoldResult.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Operation.h"
+
+using namespace mlir;
+
+/// Return true if no slot holds a replacement.
+static bool replacesNone(ArrayRef<OpFoldResult> slots) {
+ return llvm::none_of(
+ slots, [](OpFoldResult slot) { return static_cast<bool>(slot); });
+}
+
+//===----------------------------------------------------------------------===//
+// OpFoldResults
+//===----------------------------------------------------------------------===//
+
+OpFoldResults::OpFoldResults(std::nullptr_t) {}
+
+OpFoldResults::OpFoldResults(LogicalResult status)
+ : inPlace(status.succeeded()), isFailure(status.failed()) {}
+
+OpFoldResults::OpFoldResults(OpFoldResult replacement) {
+ if (!replacement)
+ return;
+ slots.push_back(replacement);
+ isFailure = false;
+}
+
+OpFoldResults::OpFoldResults(Value replacement)
+ : OpFoldResults(OpFoldResult(replacement)) {}
+
+OpFoldResults::OpFoldResults(Attribute replacement)
+ : OpFoldResults(OpFoldResult(replacement)) {}
+
+OpFoldResults::OpFoldResults(std::initializer_list<OpFoldResult> replacements)
+ : slots(replacements), isFailure(replacesNone(slots)) {}
+
+OpFoldResults::OpFoldResults(Operation *op) : op(op) {
+ assert(op && "expected a non-null operation");
+ slots.resize(op->getNumResults());
+}
+
+void OpFoldResults::replace(Value result, OpFoldResult replacement) {
+ auto opResult = dyn_cast_if_present<OpResult>(result);
+ assert(op && opResult && opResult.getOwner() == op &&
+ "expected a result of the folded operation");
+ replace(opResult.getResultNumber(), replacement);
+}
+
+void OpFoldResults::replace(unsigned resultIndex, OpFoldResult replacement) {
+ assert(op && "replace() needs the Operation * form or a normalized object");
+ unsigned numResults = op->getNumResults();
+ if (slots.empty())
+ slots.resize(numResults);
+ assert(slots.size() == numResults && "expected one slot per result");
+ assert(resultIndex < numResults && "result index out of range");
+ // The op's own result keeps the slot.
+ if (dyn_cast_if_present<Value>(replacement) == op->getResult(resultIndex))
+ replacement = OpFoldResult();
+ slots[resultIndex] = replacement;
+ isFailure = !inPlace && replacesNone(slots);
+}
+
+void OpFoldResults::markModifiedInPlace(bool modified) {
+ inPlace = modified;
+ isFailure = !inPlace && replacesNone(slots);
+}
+
+bool OpFoldResults::succeeded() const { return !isFailure; }
+
+bool OpFoldResults::failed() const { return isFailure; }
+
+bool OpFoldResults::modifiedInPlace() const { return inPlace; }
+
+bool OpFoldResults::replacesAny() const { return !replacesNone(slots); }
+
+bool OpFoldResults::replacesAll() const {
+ return !slots.empty() && llvm::all_of(slots, [](OpFoldResult slot) {
+ return static_cast<bool>(slot);
+ });
+}
+
+unsigned OpFoldResults::size() const { return slots.size(); }
+
+OpFoldResult OpFoldResults::operator[](unsigned i) const {
+ if (slots.empty()) {
+ assert((!op || i < op->getNumResults()) && "result index out of range");
+ return OpFoldResult();
+ }
+ assert(i < slots.size() && "result index out of range");
+ return slots[i];
+}
+
+ArrayRef<OpFoldResult> OpFoldResults::getSlots() const { return slots; }
+
+void OpFoldResults::normalize(Operation *op) {
+ assert(op && "expected a non-null operation");
+ assert((!this->op || this->op == op) &&
+ "normalize() called with a different operation");
+ this->op = op;
+
+ if (!slots.empty()) {
+ assert(slots.size() == op->getNumResults() &&
+ "expected one fold result slot per operation result");
+ for (auto [slot, result] : llvm::zip(slots, op->getResults()))
+ if (dyn_cast_if_present<Value>(slot) == result)
+ slot = OpFoldResult();
+ if (replacesNone(slots))
+ slots.clear();
+ }
+ isFailure = !inPlace && slots.empty();
+
+#ifndef NDEBUG
+ for (auto [index, slot] : llvm::enumerate(slots)) {
+ auto value = dyn_cast_if_present<Value>(slot);
+ if (!value)
+ continue;
+ Type expectedType = op->getResult(index).getType();
+ if (value.getType() != expectedType) {
+ op->emitOpError() << "folder produced a value of incorrect type: "
+ << value.getType() << ", expected: " << expectedType;
+ assert(false && "incorrect fold result type");
+ }
+ // Rule D17: a slot may name another result of the op only if that result
+ // is kept.
+ auto opResult = dyn_cast<OpResult>(value);
+ assert((!opResult || opResult.getOwner() != op ||
+ !slots[opResult.getResultNumber()]) &&
+ "fold result names a result of the folded operation that the same "
+ "fold replaces");
+ }
+#endif // NDEBUG
+}
+
+//===----------------------------------------------------------------------===//
+// Legacy fold results
+//===----------------------------------------------------------------------===//
+
+OpFoldResults detail::convertLegacyFoldResults(LogicalResult status,
+ ArrayRef<OpFoldResult> results) {
+ if (failed(status))
+ return failure();
+ if (results.empty())
+ return success();
+ assert(llvm::all_of(
+ results,
+ [](OpFoldResult result) { return static_cast<bool>(result); }) &&
+ "legacy fold returned a null result");
+ return results;
+}
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index d262e844603a1..52b92c3d57dfb 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -600,63 +600,53 @@ void Operation::setSuccessor(Block *block, unsigned index) {
getBlockOperands()[index].set(block);
}
-#ifndef NDEBUG
-/// Assert that the folded results (in case of values) have the same type as
-/// the results of the given op.
-static void checkFoldResultTypes(Operation *op,
- SmallVectorImpl<OpFoldResult> &results) {
- if (results.empty())
- return;
-
- for (auto [ofr, opResult] : llvm::zip_equal(results, op->getResults())) {
- if (auto value = dyn_cast<Value>(ofr)) {
- if (value.getType() != opResult.getType()) {
- op->emitOpError() << "folder produced a value of incorrect type: "
- << value.getType()
- << ", expected: " << opResult.getType();
- assert(false && "incorrect fold result type");
- }
- }
- }
-}
-#endif // NDEBUG
-
/// Attempt to fold this operation using the Op's registered foldHook.
-LogicalResult Operation::fold(ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
+OpFoldResults Operation::fold(ArrayRef<Attribute> operands) {
// If we have a registered operation definition matching this one, use it to
// try to constant fold the operation.
- if (succeeded(name.foldHook(this, operands, results))) {
-#ifndef NDEBUG
- checkFoldResultTypes(this, results);
-#endif // NDEBUG
- return success();
- }
+ OpFoldResults results = name.foldHook(this, operands);
+ results.normalize(this);
+ if (results.succeeded())
+ return results;
// Otherwise, fall back on the dialect hook to handle it.
Dialect *dialect = getDialect();
if (!dialect)
- return failure();
+ return results;
auto *interface = dyn_cast<DialectFoldInterface>(dialect);
if (!interface)
- return failure();
+ return results;
- LogicalResult status = interface->fold(this, operands, results);
-#ifndef NDEBUG
- if (succeeded(status))
- checkFoldResultTypes(this, results);
-#endif // NDEBUG
- return status;
+ SmallVector<OpFoldResult> legacyResults;
+ LogicalResult status = interface->fold(this, operands, legacyResults);
+ results = detail::convertLegacyFoldResults(status, legacyResults);
+ results.normalize(this);
+ return results;
+}
+
+/// Compute the constant operand values of `op`.
+static SmallVector<Attribute> getConstantOperands(Operation *op) {
+ SmallVector<Attribute> constants(op->getNumOperands(), Attribute());
+ for (unsigned i = 0, e = op->getNumOperands(); i != e; ++i)
+ matchPattern(op->getOperand(i), m_Constant(&constants[i]));
+ return constants;
+}
+
+OpFoldResults Operation::fold() { return fold(getConstantOperands(this)); }
+
+LogicalResult Operation::fold(ArrayRef<Attribute> operands,
+ SmallVectorImpl<OpFoldResult> &results) {
+ OpFoldResults foldResults = fold(operands);
+ if (foldResults.replacesAll()) {
+ llvm::append_range(results, foldResults.getSlots());
+ return success();
+ }
+ return success(foldResults.modifiedInPlace());
}
LogicalResult Operation::fold(SmallVectorImpl<OpFoldResult> &results) {
- // Check if any operands are constants.
- SmallVector<Attribute> constants;
- constants.assign(getNumOperands(), Attribute());
- for (unsigned i = 0, e = getNumOperands(); i != e; ++i)
- matchPattern(getOperand(i), m_Constant(&constants[i]));
- return fold(constants, results);
+ return fold(getConstantOperands(this), results);
}
/// Emit an error with the op name prefixed, like "'dim' op " which is
diff --git a/mlir/unittests/IR/CMakeLists.txt b/mlir/unittests/IR/CMakeLists.txt
index ce9b230222f62..e6858585aaf6c 100644
--- a/mlir/unittests/IR/CMakeLists.txt
+++ b/mlir/unittests/IR/CMakeLists.txt
@@ -14,6 +14,7 @@ add_mlir_unittest(MLIRIRTests
LocationTest.cpp
MLIRContextResetTest.cpp
MemrefLayoutTest.cpp
+ OpFoldResultsTest.cpp
OpImplementationTest.cpp
OperationSupportTest.cpp
PatternMatchTest.cpp
diff --git a/mlir/unittests/IR/OpFoldResultsTest.cpp b/mlir/unittests/IR/OpFoldResultsTest.cpp
new file mode 100644
index 0000000000000..3e8bd4a3afc27
--- /dev/null
+++ b/mlir/unittests/IR/OpFoldResultsTest.cpp
@@ -0,0 +1,588 @@
+//===- OpFoldResultsTest.cpp - OpFoldResults unit tests -------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/ExtensibleDialect.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OpFoldResult.h"
+#include "mlir/IR/Operation.h"
+#include "gtest/gtest.h"
+
+#include <functional>
+
+using namespace mlir;
+
+namespace {
+class OpFoldResultsTest : public ::testing::Test {
+protected:
+ OpFoldResultsTest() : builder(&context) {
+ context.allowUnregisteredDialects();
+ i32 = builder.getI32Type();
+ f32 = builder.getF32Type();
+ }
+
+ ~OpFoldResultsTest() override {
+ for (Operation *op : ops)
+ op->destroy();
+ }
+
+ /// Create an op with the given result types. The fixture destroys it.
+ Operation *createOp(TypeRange resultTypes, StringRef name = "foo.bar") {
+ OperationState state(UnknownLoc::get(&context), name);
+ state.addTypes(resultTypes);
+ Operation *op = Operation::create(state);
+ ops.push_back(op);
+ return op;
+ }
+
+ /// Load a dynamic dialect with two ops. The fold hook of `test_fold.op`
+ /// calls `foldFn`. The legacy fold hook of `test_fold.legacy_op` calls
+ /// `legacyFoldFn`.
+ void loadDynamicDialect() {
+ context.getOrLoadDynamicDialect("test_fold", [&](DynamicDialect *dialect) {
+ auto verify = [](Operation *) { return success(); };
+ std::unique_ptr<DynamicOpDefinition> opDef =
+ DynamicOpDefinition::get("op", dialect, verify, verify);
+ opDef->setFoldHookFn(
+ [this](Operation *op, ArrayRef<Attribute>) { return foldFn(op); });
+ dialect->registerDynamicOp(std::move(opDef));
+
+ std::unique_ptr<DynamicOpDefinition> legacyOpDef =
+ DynamicOpDefinition::get("legacy_op", dialect, verify, verify);
+ legacyOpDef->setFoldHookFn(
+ [this](Operation *op, ArrayRef<Attribute>,
+ SmallVectorImpl<OpFoldResult> &results) {
+ return legacyFoldFn(op, results);
+ });
+ dialect->registerDynamicOp(std::move(legacyOpDef));
+ });
+ }
+
+ MLIRContext context;
+ Builder builder;
+ Type i32;
+ Type f32;
+ SmallVector<Operation *> ops;
+ std::function<OpFoldResults(Operation *)> foldFn;
+ std::function<LogicalResult(Operation *, SmallVectorImpl<OpFoldResult> &)>
+ legacyFoldFn;
+};
+} // namespace
+
+static void expectOneSlot(const OpFoldResults &result, OpFoldResult expected) {
+ EXPECT_TRUE(result.succeeded());
+ EXPECT_FALSE(result.modifiedInPlace());
+ EXPECT_TRUE(result.replacesAny());
+ EXPECT_TRUE(result.replacesAll());
+ ASSERT_EQ(result.size(), 1u);
+ EXPECT_EQ(result[0], expected);
+}
+
+// The helpers below return the way a folder returns, so they use copy
+// initialization.
+static OpFoldResults foldToEmptyBraces() { return {}; }
+static OpFoldResults foldToTypedValue(TypedValue<IntegerType> value) {
+ return value;
+}
+static OpFoldResults foldToIntegerAttr(IntegerAttr attr) { return attr; }
+static OpFoldResults foldToArrayAttr(ArrayAttr attr) { return attr; }
+static OpFoldResults foldToList(OpFoldResult lhs, OpFoldResult rhs) {
+ return {lhs, rhs};
+}
+
+TEST_F(OpFoldResultsTest, DefaultIsFailure) {
+ OpFoldResults result;
+ EXPECT_TRUE(result.failed());
+ EXPECT_FALSE(result.succeeded());
+ EXPECT_FALSE(result.modifiedInPlace());
+ EXPECT_FALSE(result.replacesAny());
+ EXPECT_FALSE(result.replacesAll());
+ EXPECT_EQ(result.size(), 0u);
+ EXPECT_TRUE(result.getSlots().empty());
+}
+
+TEST_F(OpFoldResultsTest, EmptyBracesAreFailure) {
+ OpFoldResults result = foldToEmptyBraces();
+ EXPECT_TRUE(result.failed());
+ EXPECT_FALSE(result.modifiedInPlace());
+ EXPECT_EQ(result.size(), 0u);
+}
+
+TEST_F(OpFoldResultsTest, NullptrIsFailure) {
+ OpFoldResults result = nullptr;
+ EXPECT_TRUE(result.failed());
+ EXPECT_FALSE(result.modifiedInPlace());
+ EXPECT_EQ(result.size(), 0u);
+}
+
+TEST_F(OpFoldResultsTest, LogicalResult) {
+ OpFoldResults inPlace = success();
+ EXPECT_TRUE(inPlace.succeeded());
+ EXPECT_TRUE(inPlace.modifiedInPlace());
+ EXPECT_FALSE(inPlace.replacesAny());
+ EXPECT_FALSE(inPlace.replacesAll());
+ EXPECT_EQ(inPlace.size(), 0u);
+
+ OpFoldResults failedResult = failure();
+ EXPECT_TRUE(failedResult.failed());
+ EXPECT_FALSE(failedResult.modifiedInPlace());
+ EXPECT_FALSE(failedResult.replacesAny());
+ EXPECT_EQ(failedResult.size(), 0u);
+}
+
+TEST_F(OpFoldResultsTest, SingleSlot) {
+ Operation *producer = createOp({i32});
+ Operation *op = createOp({i32});
+ Value value = producer->getResult(0);
+ IntegerAttr intAttr = builder.getI32IntegerAttr(7);
+ Attribute attr = intAttr;
+
+ OpFoldResults fromOpFoldResult = OpFoldResult(attr);
+ expectOneSlot(fromOpFoldResult, attr);
+
+ OpFoldResults fromValue = value;
+ expectOneSlot(fromValue, value);
+
+ OpFoldResults fromAttr = attr;
+ expectOneSlot(fromAttr, attr);
+
+ // TypedValue and IntegerAttr use the Value and Attribute constructors.
+ OpFoldResults fromTypedValue =
+ foldToTypedValue(cast<TypedValue<IntegerType>>(value));
+ expectOneSlot(fromTypedValue, value);
+
+ OpFoldResults fromIntegerAttr = foldToIntegerAttr(intAttr);
+ expectOneSlot(fromIntegerAttr, attr);
+
+ // The one slot matches the single result.
+ fromValue.normalize(op);
+ expectOneSlot(fromValue, value);
+ EXPECT_EQ(fromValue.size(), op->getNumResults());
+}
+
+TEST_F(OpFoldResultsTest, ArrayAttrFillsOneSlot) {
+ Operation *op = createOp({i32});
+ ArrayAttr arrayAttr = builder.getI32ArrayAttr({1, 2, 3});
+ OpFoldResults result = foldToArrayAttr(arrayAttr);
+ expectOneSlot(result, arrayAttr);
+ result.normalize(op);
+ expectOneSlot(result, arrayAttr);
+}
+
+TEST_F(OpFoldResultsTest, InitializerList) {
+ Operation *producer = createOp({i32});
+ Operation *op = createOp({i32, f32});
+ Value value = producer->getResult(0);
+ Attribute attr = builder.getF32FloatAttr(1.0);
+
+ OpFoldResults result = {value, attr};
+ EXPECT_TRUE(result.succeeded());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_EQ(result[0], OpFoldResult(value));
+ EXPECT_EQ(result[1], OpFoldResult(attr));
+ result.normalize(op);
+ EXPECT_TRUE(result.succeeded());
+ EXPECT_FALSE(result.modifiedInPlace());
+ EXPECT_TRUE(result.replacesAll());
+
+ OpFoldResults partial = foldToList(OpFoldResult(), attr);
+ partial.normalize(op);
+ EXPECT_TRUE(partial.succeeded());
+ EXPECT_TRUE(partial.replacesAny());
+ EXPECT_FALSE(partial.replacesAll());
+ ASSERT_EQ(partial.size(), 2u);
+ EXPECT_FALSE(partial[0]);
+ EXPECT_EQ(partial[1], OpFoldResult(attr));
+}
+
+TEST_F(OpFoldResultsTest, Range) {
+ Operation *producer = createOp({i32, f32});
+ Operation *op = createOp({i32, f32});
+ Attribute attr = builder.getI32IntegerAttr(3);
+
+ SmallVector<OpFoldResult> vector = {attr, OpFoldResult()};
+ OpFoldResults fromVector = vector;
+ EXPECT_TRUE(fromVector.succeeded());
+ ASSERT_EQ(fromVector.size(), 2u);
+ EXPECT_EQ(fromVector[0], OpFoldResult(attr));
+ EXPECT_FALSE(fromVector[1]);
+ fromVector.normalize(op);
+ EXPECT_TRUE(fromVector.replacesAny());
+ EXPECT_FALSE(fromVector.replacesAll());
+
+ ValueRange range = producer->getResults();
+ OpFoldResults fromRange = range;
+ EXPECT_TRUE(fromRange.succeeded());
+ ASSERT_EQ(fromRange.size(), 2u);
+ EXPECT_EQ(fromRange[0], OpFoldResult(producer->getResult(0)));
+ EXPECT_EQ(fromRange[1], OpFoldResult(producer->getResult(1)));
+ fromRange.normalize(op);
+ EXPECT_TRUE(fromRange.replacesAll());
+
+ OpFoldResults fromEmptyRange = ValueRange();
+ EXPECT_TRUE(fromEmptyRange.failed());
+ EXPECT_EQ(fromEmptyRange.size(), 0u);
+
+ OpFoldResults fromEmptyVector = SmallVector<OpFoldResult>();
+ EXPECT_TRUE(fromEmptyVector.failed());
+
+ SmallVector<OpFoldResult> nulls(2);
+ OpFoldResults fromNulls = nulls;
+ EXPECT_TRUE(fromNulls.failed());
+}
+
+TEST_F(OpFoldResultsTest, IncrementalForm) {
+ Operation *producer = createOp({i32});
+ Operation *op = createOp({i32, i32});
+ Value value = producer->getResult(0);
+ Attribute attr = builder.getI32IntegerAttr(5);
+
+ OpFoldResults result(op);
+ EXPECT_TRUE(result.failed());
+ EXPECT_FALSE(result.modifiedInPlace());
+ EXPECT_FALSE(result.replacesAny());
+ EXPECT_EQ(result.size(), 2u);
+
+ result.replace(op->getResult(1), attr);
+ EXPECT_TRUE(result.succeeded());
+ EXPECT_TRUE(result.replacesAny());
+ EXPECT_FALSE(result.replacesAll());
+ EXPECT_FALSE(result[0]);
+ EXPECT_EQ(result[1], OpFoldResult(attr));
+
+ // A literal 0 picks the unsigned overload.
+ result.replace(0, value);
+ EXPECT_TRUE(result.replacesAll());
+ EXPECT_EQ(result[0], OpFoldResult(value));
+
+ // The last write wins.
+ result.replace(0u, attr);
+ EXPECT_EQ(result[0], OpFoldResult(attr));
+
+ // The op's own result and a null replacement keep the slot.
+ result.replace(1u, op->getResult(1));
+ EXPECT_FALSE(result[1]);
+ result.replace(op->getResult(0), OpFoldResult());
+ EXPECT_FALSE(result[0]);
+ EXPECT_TRUE(result.failed());
+ EXPECT_FALSE(result.replacesAny());
+
+ result.markModifiedInPlace();
+ EXPECT_TRUE(result.succeeded());
+ EXPECT_TRUE(result.modifiedInPlace());
+ result.markModifiedInPlace(false);
+ EXPECT_TRUE(result.failed());
+ EXPECT_FALSE(result.modifiedInPlace());
+
+ result.replace(1u, value);
+ result.normalize(op);
+ EXPECT_TRUE(result.succeeded());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_FALSE(result[0]);
+ EXPECT_EQ(result[1], OpFoldResult(value));
+ ArrayRef<OpFoldResult> slots = result.getSlots();
+ ASSERT_EQ(slots.size(), 2u);
+ EXPECT_FALSE(slots[0]);
+ EXPECT_EQ(slots[1], OpFoldResult(value));
+}
+
+TEST_F(OpFoldResultsTest, NullSafety) {
+ Operation *op = createOp({i32, i32});
+
+ EXPECT_TRUE(OpFoldResults(Value()).failed());
+ EXPECT_TRUE(OpFoldResults(Attribute()).failed());
+ EXPECT_TRUE(OpFoldResults(OpFoldResult()).failed());
+
+ OpFoldResults nullList = {OpFoldResult(), OpFoldResult()};
+ EXPECT_TRUE(nullList.failed());
+ nullList.normalize(op);
+ EXPECT_TRUE(nullList.failed());
+ EXPECT_EQ(nullList.size(), 0u);
+ EXPECT_FALSE(nullList[0]);
+ EXPECT_FALSE(nullList[1]);
+
+ Attribute attr = builder.getI32IntegerAttr(1);
+ OpFoldResults mixed = {OpFoldResult(), attr};
+ mixed.normalize(op);
+ EXPECT_TRUE(mixed.succeeded());
+ EXPECT_FALSE(mixed[0]);
+ EXPECT_EQ(mixed[1], OpFoldResult(attr));
+
+ OpFoldResults incremental(op);
+ incremental.replace(0u, OpFoldResult());
+ incremental.normalize(op);
+ EXPECT_TRUE(incremental.failed());
+ EXPECT_EQ(incremental.size(), 0u);
+}
+
+TEST_F(OpFoldResultsTest, NormalizeExpandsToResultCount) {
+ Operation *op = createOp({i32, i32});
+ Attribute attr = builder.getI32IntegerAttr(2);
+
+ // An in-place result has no slots. Every slot reads as "keep".
+ OpFoldResults inPlace = success();
+ inPlace.normalize(op);
+ EXPECT_EQ(inPlace.size(), 0u);
+ EXPECT_FALSE(inPlace[1]);
+
+ // replace() on a normalized object creates one slot per result.
+ inPlace.replace(1u, attr);
+ EXPECT_TRUE(inPlace.succeeded());
+ EXPECT_TRUE(inPlace.modifiedInPlace());
+ ASSERT_EQ(inPlace.size(), op->getNumResults());
+ EXPECT_FALSE(inPlace[0]);
+ EXPECT_EQ(inPlace[1], OpFoldResult(attr));
+
+ OpFoldResults failedResult = failure();
+ failedResult.normalize(op);
+ failedResult.replace(0u, attr);
+ EXPECT_TRUE(failedResult.succeeded());
+ EXPECT_FALSE(failedResult.modifiedInPlace());
+ ASSERT_EQ(failedResult.size(), op->getNumResults());
+ EXPECT_EQ(failedResult[0], OpFoldResult(attr));
+ EXPECT_FALSE(failedResult[1]);
+}
+
+TEST_F(OpFoldResultsTest, NormalizeMapsOwnResultsToKeep) {
+ Operation *op = createOp({i32, i32});
+ Attribute attr = builder.getI32IntegerAttr(4);
+
+ OpFoldResults result = {op->getResult(0), attr};
+ result.normalize(op);
+ EXPECT_TRUE(result.succeeded());
+ EXPECT_TRUE(result.replacesAny());
+ EXPECT_FALSE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_FALSE(result[0]);
+ EXPECT_EQ(result[1], OpFoldResult(attr));
+
+ // A slot may name another result of the op if that result is kept.
+ OpFoldResults forward = {op->getResult(1), OpFoldResult()};
+ forward.normalize(op);
+ EXPECT_TRUE(forward.succeeded());
+ ASSERT_EQ(forward.size(), 2u);
+ EXPECT_EQ(forward[0], OpFoldResult(op->getResult(1)));
+ EXPECT_FALSE(forward[1]);
+}
+
+TEST_F(OpFoldResultsTest, NormalizeCollapsesToFailure) {
+ Operation *op = createOp({i32, i32});
+ Operation *oneResultOp = createOp({i32});
+
+ OpFoldResults ownResults = {op->getResult(0), op->getResult(1)};
+ ownResults.normalize(op);
+ EXPECT_TRUE(ownResults.failed());
+ EXPECT_FALSE(ownResults.replacesAny());
+ EXPECT_EQ(ownResults.size(), 0u);
+ EXPECT_TRUE(ownResults.getSlots().empty());
+
+ OpFoldResults ownRange = op->getResults();
+ ownRange.normalize(op);
+ EXPECT_TRUE(ownRange.failed());
+
+ OpFoldResults ownResult = oneResultOp->getResult(0);
+ ownResult.normalize(oneResultOp);
+ EXPECT_TRUE(ownResult.failed());
+
+ OpFoldResults nothingReplaced(op);
+ nothingReplaced.normalize(op);
+ EXPECT_TRUE(nothingReplaced.failed());
+ EXPECT_EQ(nothingReplaced.size(), 0u);
+
+ // An in-place change keeps the result successful.
+ OpFoldResults inPlace = {op->getResult(0), op->getResult(1)};
+ inPlace.markModifiedInPlace();
+ inPlace.normalize(op);
+ EXPECT_TRUE(inPlace.succeeded());
+ EXPECT_TRUE(inPlace.modifiedInPlace());
+ EXPECT_FALSE(inPlace.replacesAny());
+ EXPECT_FALSE(inPlace.replacesAll());
+ EXPECT_EQ(inPlace.size(), 0u);
+}
+
+TEST_F(OpFoldResultsTest, NormalizeIsIdempotent) {
+ Operation *op = createOp({i32, i32});
+ Attribute attr = builder.getI32IntegerAttr(6);
+
+ OpFoldResults result = {op->getResult(0), attr};
+ result.markModifiedInPlace();
+ result.normalize(op);
+ SmallVector<OpFoldResult> slots(result.getSlots());
+ result.normalize(op);
+ EXPECT_TRUE(result.succeeded());
+ EXPECT_TRUE(result.modifiedInPlace());
+ EXPECT_TRUE(result.replacesAny());
+ EXPECT_FALSE(result.replacesAll());
+ ASSERT_EQ(result.size(), slots.size());
+ for (unsigned i = 0, e = slots.size(); i < e; ++i)
+ EXPECT_EQ(result[i], slots[i]);
+
+ OpFoldResults failedResult = failure();
+ failedResult.normalize(op);
+ failedResult.normalize(op);
+ EXPECT_TRUE(failedResult.failed());
+ EXPECT_EQ(failedResult.size(), 0u);
+
+ OpFoldResults inPlace = success();
+ inPlace.normalize(op);
+ inPlace.normalize(op);
+ EXPECT_TRUE(inPlace.succeeded());
+ EXPECT_TRUE(inPlace.modifiedInPlace());
+ EXPECT_EQ(inPlace.size(), 0u);
+}
+
+TEST_F(OpFoldResultsTest, ZeroResultInPlaceDoesNotReplaceAll) {
+ Operation *op = createOp({});
+
+ OpFoldResults inPlace = success();
+ inPlace.normalize(op);
+ EXPECT_TRUE(inPlace.succeeded());
+ EXPECT_TRUE(inPlace.modifiedInPlace());
+ EXPECT_FALSE(inPlace.replacesAny());
+ EXPECT_FALSE(inPlace.replacesAll());
+ EXPECT_EQ(inPlace.size(), 0u);
+
+ OpFoldResults failedResult = failure();
+ failedResult.normalize(op);
+ EXPECT_TRUE(failedResult.failed());
+ EXPECT_FALSE(failedResult.replacesAll());
+}
+
+TEST_F(OpFoldResultsTest, UnregisteredOpFoldFails) {
+ Operation *op = createOp({i32});
+ EXPECT_TRUE(op->fold().failed());
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+}
+
+TEST_F(OpFoldResultsTest, LegacyOperationFoldKeepsStrictContract) {
+ loadDynamicDialect();
+ Operation *producer = createOp({i32, i32});
+ Operation *op = createOp({i32, i32}, "test_fold.op");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ SmallVector<OpFoldResult> results;
+
+ foldFn = [](Operation *) -> OpFoldResults { return failure(); };
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+ EXPECT_TRUE(op->fold().failed());
+
+ foldFn = [](Operation *) -> OpFoldResults { return success(); };
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+
+ foldFn = [&](Operation *) -> OpFoldResults {
+ return {attr, producer->getResult(1)};
+ };
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ ASSERT_EQ(results.size(), 2u);
+ EXPECT_EQ(results[0], OpFoldResult(attr));
+ EXPECT_EQ(results[1], OpFoldResult(producer->getResult(1)));
+ results.clear();
+
+ // The legacy overloads do not apply a partial fold. Without an in-place
+ // change, the partial fold is a failure.
+ foldFn = [&](Operation *foldedOp) {
+ OpFoldResults partial(foldedOp);
+ partial.replace(1u, attr);
+ return partial;
+ };
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+ OpFoldResults partialResult = op->fold();
+ EXPECT_TRUE(partialResult.succeeded());
+ EXPECT_FALSE(partialResult.modifiedInPlace());
+ ASSERT_EQ(partialResult.size(), 2u);
+ EXPECT_FALSE(partialResult[0]);
+ EXPECT_EQ(partialResult[1], OpFoldResult(attr));
+
+ // With an in-place change, the partial fold is reported as in place.
+ foldFn = [&](Operation *foldedOp) {
+ OpFoldResults partial(foldedOp);
+ partial.replace(1u, attr);
+ partial.markModifiedInPlace();
+ return partial;
+ };
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+
+ // A fold that keeps every slot and is not in place is a failure.
+ foldFn = [](Operation *foldedOp) -> OpFoldResults {
+ return foldedOp->getResults();
+ };
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+ EXPECT_TRUE(op->fold().failed());
+}
+
+TEST_F(OpFoldResultsTest, LegacyDynamicFoldHookKeepsStrictContract) {
+ loadDynamicDialect();
+ Operation *producer = createOp({i32, i32});
+ Operation *op = createOp({i32, i32}, "test_fold.legacy_op");
+ Attribute attr = builder.getI32IntegerAttr(1);
+
+ legacyFoldFn = [](Operation *, SmallVectorImpl<OpFoldResult> &) {
+ return failure();
+ };
+ EXPECT_TRUE(op->fold().failed());
+
+ legacyFoldFn = [](Operation *, SmallVectorImpl<OpFoldResult> &) {
+ return success();
+ };
+ OpFoldResults inPlace = op->fold();
+ EXPECT_TRUE(inPlace.succeeded());
+ EXPECT_TRUE(inPlace.modifiedInPlace());
+ EXPECT_FALSE(inPlace.replacesAny());
+
+ legacyFoldFn = [&](Operation *, SmallVectorImpl<OpFoldResult> &results) {
+ results.push_back(attr);
+ results.push_back(producer->getResult(0));
+ return success();
+ };
+ OpFoldResults all = op->fold();
+ EXPECT_TRUE(all.succeeded());
+ EXPECT_FALSE(all.modifiedInPlace());
+ EXPECT_TRUE(all.replacesAll());
+ ASSERT_EQ(all.size(), 2u);
+ EXPECT_EQ(all[0], OpFoldResult(attr));
+ EXPECT_EQ(all[1], OpFoldResult(producer->getResult(0)));
+
+ // A legacy fold that forwards the op's own results keeps every slot, so it
+ // is a failure.
+ legacyFoldFn = [](Operation *foldedOp,
+ SmallVectorImpl<OpFoldResult> &results) {
+ llvm::append_range(results, foldedOp->getResults());
+ return success();
+ };
+ EXPECT_TRUE(op->fold().failed());
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+}
+
+#ifdef GTEST_HAS_DEATH_TEST
+#ifndef NDEBUG
+namespace {
+class OpFoldResultsDeathTest : public OpFoldResultsTest {};
+} // namespace
+
+TEST_F(OpFoldResultsDeathTest, SlotNamesReplacedResult) {
+ Operation *producer = createOp({i32});
+ Operation *op = createOp({i32, i32});
+ // Slot 0 forwards result 1, but slot 1 replaces result 1 (rule D17).
+ OpFoldResults result(op);
+ result.replace(0u, op->getResult(1));
+ result.replace(1u, producer->getResult(0));
+ EXPECT_DEATH(result.normalize(op),
+ "fold result names a result of the folded operation");
+}
+#endif // NDEBUG
+#endif // GTEST_HAS_DEATH_TEST
More information about the Mlir-commits
mailing list