[Mlir-commits] [mlir] [mlir] Add OpFoldResults for partial folds of multi-result ops (PR #227430)
Victor Perez
llvmlistbot at llvm.org
Fri Oct 2 12:48:52 PDT 2026
=?utf-8?q?Víctor_Pérez?= Carrasco
Message-ID:
In-Reply-To: <llvm.org/llvm/llvm-project/pull/227430 at github.com>
https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/227430
>From 89493160d065bde9a70df5acd736eabb8f4ed4c4 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 21:01:18 -0700
Subject: [PATCH 1/2] [mlir] Keep existing values when fold materialization
fails
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
A fold can return an existing value for one result and an attribute for
another result. When the greedy driver cannot materialize a constant for
the attribute, it rolls back the fold. The old rollback erased the
defining op of every replacement value. This erased ops that the fold did
not create, and it tried to erase a null op when the value was a block
argument. In builds with assertions, an assert fired instead.
The driver now records the constants that it materializes for each fold
call. When a materialization fails, it erases only those constants
through the rewriter and applies no replacement.
The test dialect now fails to materialize the string attribute
"unmaterializable". A new test op folds to a constant, its operand, and
this attribute. The new canonicalize tests check that the fold does not
apply and that the operand and its defining op stay.
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
.../Utils/GreedyPatternRewriteDriver.cpp | 20 +++++++---------
mlir/test/Transforms/test-canonicalize.mlir | 24 +++++++++++++++++++
mlir/test/lib/Dialect/Test/TestDialect.cpp | 4 ++++
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 16 +++++++++++++
mlir/test/lib/Dialect/Test/TestOps.td | 9 +++++++
5 files changed, 62 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index b71194e783ba067..dc6484b05ee820a 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -519,6 +519,7 @@ bool GreedyPatternRewriteDriver::processWorklist() {
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(op);
SmallVector<Value> replacements;
+ SmallVector<Operation *> materializedConstants;
bool materializationSucceeded = true;
for (auto [ofr, resultType] :
llvm::zip_equal(foldResults, op->getResultTypes())) {
@@ -533,17 +534,13 @@ bool GreedyPatternRewriteDriver::processWorklist() {
rewriter, cast<Attribute>(ofr), resultType, op->getLoc());
if (!constOp) {
- // If materialization fails, cleanup any operations generated for
- // the previous results.
- llvm::SmallDenseSet<Operation *> replacementOps;
- for (Value replacement : replacements) {
- assert(replacement.use_empty() &&
- "folder reused existing op for one result but constant "
- "materialization failed for another result");
- replacementOps.insert(replacement.getDefiningOp());
- }
- for (Operation *op : replacementOps) {
- rewriter.eraseOp(op);
+ // If materialization fails, erase only the constants that were
+ // materialized for the previous results. Values that the folder
+ // returned existed before the fold, so they must stay.
+ for (Operation *cst : materializedConstants) {
+ assert(cst->use_empty() &&
+ "materialized constant has uses before replacement");
+ rewriter.eraseOp(cst);
}
materializationSucceeded = false;
@@ -554,6 +551,7 @@ bool GreedyPatternRewriteDriver::processWorklist() {
"materializeConstant produced op that is not a ConstantLike");
assert(constOp->getResultTypes()[0] == resultType &&
"materializeConstant produced incorrect result type");
+ materializedConstants.push_back(constOp);
replacements.push_back(constOp->getResult(0));
}
diff --git a/mlir/test/Transforms/test-canonicalize.mlir b/mlir/test/Transforms/test-canonicalize.mlir
index 8cad6b98441df14..81538e88beae434 100644
--- a/mlir/test/Transforms/test-canonicalize.mlir
+++ b/mlir/test/Transforms/test-canonicalize.mlir
@@ -87,3 +87,27 @@ func.func @test_region_simplify(%input1 : i32, %cond : i1) -> i32 {
^bb1(%used_arg : i32, %unused_arg : i32):
return %used_arg : i32
}
+
+// When constant materialization fails for one result, the fold must not apply.
+// The driver erases only the constant that it materialized for the first
+// result, and keeps the op that defines the forwarded operand.
+// CHECK-LABEL: func @fold_unmaterializable_existing_op
+// CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: i32, %[[ARG_1:[a-z0-9]*]]: i32)
+func.func @fold_unmaterializable_existing_op(%arg0 : i32, %arg1 : i32) -> (i32, i32, i32) {
+ // CHECK-NEXT: %[[ADD:[a-z0-9]+]] = "test.addi"(%[[ARG_0]], %[[ARG_1]])
+ %0 = "test.addi"(%arg0, %arg1) : (i32, i32) -> i32
+ // CHECK-NEXT: %[[RES:[a-z0-9]+]]:3 = "test.op_fold_unmaterializable"(%[[ADD]])
+ %1:3 = "test.op_fold_unmaterializable"(%0) : (i32) -> (i32, i32, i32)
+ // CHECK-NEXT: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2
+ return %1#0, %1#1, %1#2 : i32, i32, i32
+}
+
+// Same as above, but the forwarded operand is a block argument.
+// CHECK-LABEL: func @fold_unmaterializable_block_arg
+// CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: i32)
+func.func @fold_unmaterializable_block_arg(%arg0 : i32) -> (i32, i32, i32) {
+ // CHECK-NEXT: %[[RES:[a-z0-9]+]]:3 = "test.op_fold_unmaterializable"(%[[ARG_0]])
+ %0:3 = "test.op_fold_unmaterializable"(%arg0) : (i32) -> (i32, i32, i32)
+ // CHECK-NEXT: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2
+ return %0#0, %0#1, %0#2 : i32, i32, i32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 57bba2e5749a86a..cfc45509f2fc895 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -369,6 +369,10 @@ TestDialect::~TestDialect() {
Operation *TestDialect::materializeConstant(OpBuilder &builder, Attribute value,
Type type, Location loc) {
+ // Tests use this marker to make constant materialization fail.
+ if (auto str = dyn_cast<StringAttr>(value);
+ str && str.getValue() == "unmaterializable")
+ return nullptr;
return TestOpConstant::create(builder, loc, type, value);
}
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 9fe2168e1173744..d20761468463dfc 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -236,6 +236,22 @@ LogicalResult TestOpWithVariadicResultsAndFolder::fold(
return success();
}
+//===----------------------------------------------------------------------===//
+// TestOpFoldUnmaterializable
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+TestOpFoldUnmaterializable::fold(FoldAdaptor adaptor,
+ SmallVectorImpl<OpFoldResult> &results) {
+ // The unmaterializable slot comes last, so that the rollback sees both a
+ // materialized constant and an existing value.
+ Builder b(getContext());
+ results.push_back(b.getI32IntegerAttr(42));
+ results.push_back(getOperand());
+ results.push_back(b.getStringAttr("unmaterializable"));
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// TestOpInPlaceFold
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index a6b05b55804b7e4..5617617cf0faa17 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1676,6 +1676,15 @@ def TestOpWithVariadicResultsAndFolder: TEST_Op<"op_with_variadic_results_and_fo
let hasFolder = 1;
}
+// Op for testing the rollback of a fold when constant materialization fails
+// for one result. The folder returns a constant, the operand, and an attribute
+// that the test dialect cannot materialize.
+def TestOpFoldUnmaterializable : TEST_Op<"op_fold_unmaterializable"> {
+ let arguments = (ins I32:$operand);
+ let results = (outs I32, I32, I32);
+ let hasFolder = 1;
+}
+
def TestAddIOp : TEST_Op<"addi"> {
let arguments = (ins AnyTypeOf<[I32, TestI32]>:$op1,
AnyTypeOf<[I32, TestI32]>:$op2);
>From 60d11e1443e35d7f1b0818d8105e02ca0384ede8 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 2/2] [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. For
example, a loop that yields one of its iteration arguments unchanged
could replace the matching result with the initial value and keep the
other results. 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 either no slot or
one slot per result, and an in-place bit. The plural name says that the
object holds one `OpFoldResult` per result of the op. A null slot keeps
its result. `{}`, `nullptr`, 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 clears the slots if no slot
replaces its result. In assert builds, `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 single-result
detectors also match it, and `getFoldHookFn` excludes it from the
single-result forms. `getFoldHookFn` selects the new form if the op
defines it, else the single-result form, else the legacy vector form.
Each form then folds the traits if its own fold replaces no result. An
op without a fold method folds only the traits. The legacy trait form
goes through a strict adapter. A slot of a new-form result must not
name a result of the op that the same fold replaces, because the
outcome would depend on the order of the replacements. In a graph
region, where a forwarding fold can return such a result, the
replacements are dropped. Outside a graph region, it asserts.
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 ignored
such a single-result 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. For folds
and fold traits that follow the documented contract, the legacy
overloads return the same result as before.
Downstream impact:
- Code that overrides `OperationName::Impl::foldHook`, calls
`OperationName::foldHook`, or names `OperationName::FoldHookFn` must
use the new signature `OpFoldResults(Operation *, ArrayRef<Attribute>)`.
A dynamic op can keep a legacy fold hook through
`DynamicOpDefinition::LegacyFoldHookFn`.
- An op without the OneResult trait that defines a single-result `fold`
no longer compiles. It must use one of the multi-result forms.
- `DynamicOpDefinition::get(...)` with `nullptr` or `{}` as the fold
hook is ambiguous now. `setFoldHookFn(nullptr)` and
`setFoldHookFn({})` remove the fold hook, so that folding the op
fails.
- A legacy fold or fold trait that breaks the contract can behave
differently. A slot that holds result i of the op in slot i now keeps
that result. If the fold keeps every result, it fails instead of
replacing the op with itself. If it keeps only some results, the
legacy `Operation::fold` overloads drop the partial fold. A fold that
returns `failure()` after it writes results no longer blocks the
later traits.
The new unit tests in `OpFoldResultsTest.cpp` cover the constructors,
`normalize`, null safety, the new op form and the trait folding after
it, the `DialectFoldInterface` fallback, the legacy adapters, and the
debug asserts.
RFC: https://discourse.llvm.org/t/rfc-partial-folding-for-multi-result-ops/91954
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
mlir/include/mlir/IR/ExtensibleDialect.h | 23 +-
mlir/include/mlir/IR/OpDefinition.h | 251 ++---
mlir/include/mlir/IR/OpFoldResult.h | 187 ++++
mlir/include/mlir/IR/Operation.h | 21 +
mlir/include/mlir/IR/OperationSupport.h | 54 +-
mlir/lib/IR/CMakeLists.txt | 1 +
mlir/lib/IR/ExtensibleDialect.cpp | 41 +-
mlir/lib/IR/MLIRContext.cpp | 5 +-
mlir/lib/IR/OpFoldResult.cpp | 153 ++++
mlir/lib/IR/Operation.cpp | 76 +-
mlir/unittests/IR/CMakeLists.txt | 1 +
mlir/unittests/IR/OpFoldResultsTest.cpp | 1059 ++++++++++++++++++++++
12 files changed, 1680 insertions(+), 192 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 79c07cbefcc5ce9..7a1ee4b71e7c7fe 100644
--- a/mlir/include/mlir/IR/ExtensibleDialect.h
+++ b/mlir/include/mlir/IR/ExtensibleDialect.h
@@ -440,6 +440,12 @@ class DynamicOpDefinition : public OperationName::Impl {
public:
using GetCanonicalizationPatternsFn =
llvm::unique_function<void(RewritePatternSet &, MLIRContext *) const>;
+ /// The legacy fold hook signature. Kept for backwards compatibility until the
+ /// switch to the signature returning OpFoldResults is completed. It has the
+ /// strict legacy contract: failure, success with an empty vector (in place),
+ /// or success with one 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 +468,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 +509,10 @@ class DynamicOpDefinition : public OperationName::Impl {
void setFoldHookFn(OperationName::FoldHookFn &&foldHook) {
foldHookFn = std::move(foldHook);
}
+ /// Same as above, but with a legacy fold hook.
+ void setFoldHookFn(LegacyFoldHookFn &&foldHook);
+ /// Remove the fold hook, so that folding the op always fails.
+ void setFoldHookFn(std::nullptr_t);
/// Set the hook returning any canonicalization pattern rewrites that the op
/// supports, for use by the canonicalization pass.
@@ -514,9 +532,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 0ab6679187c8eb5..751bf5fe3c3db77 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());
@@ -1604,46 +1560,68 @@ 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 generalized
+/// `foldTrait(op, operands, results)` function, with 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();
}
/// Given a tuple type containing a set of traits, return the result of folding
-/// the given operation.
+/// the given operation. `own` is the normalized result of the op's own fold.
+/// The traits run only when `own` replaced no result, and folding stops at the
+/// first trait that does not fail.
+template <typename... Ts>
+std::enable_if_t<(detect_has_any_fold_trait<Ts>::value || ...), OpFoldResults>
+foldTraits(Operation *op, ArrayRef<Attribute> operands,
+ OpFoldResults own = {}) {
+ if (own.replacesAny())
+ return own;
+ OpFoldResults trait;
+ // Stop at the first trait whose normalized result does not fail.
+ auto tryTrait = [&](OpFoldResults result) {
+ result.normalize(op);
+ if (result.failed())
+ return false;
+ trait = std::move(result);
+ return true;
+ };
+ (void)((detect_has_any_fold_trait<Ts>::value &&
+ tryTrait(foldTrait<Ts>(op, operands))) ||
+ ...);
+ bool inPlace = own.modifiedInPlace() || trait.modifiedInPlace();
+ OpFoldResults result =
+ trait.replacesAny() ? std::move(trait) : std::move(own);
+ result.setModifiedInPlace(inPlace);
+ return result;
+}
+/// Same as above, for a set of traits in which no trait can fold.
template <typename... Ts>
-LogicalResult foldTraits(Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
- return success((succeeded(foldTrait<Ts>(op, operands, results)) || ...));
+std::enable_if_t<!(detect_has_any_fold_trait<Ts>::value || ...), OpFoldResults>
+foldTraits(Operation *, ArrayRef<Attribute>, OpFoldResults own = {}) {
+ return own;
}
//===----------------------------------------------------------------------===//
@@ -1841,6 +1819,15 @@ class Op : public OpState, public Traits<ConcreteType>... {
template <class T>
constexpr static bool has_fold_adaptor_v =
llvm::is_detected<has_fold_adaptor_fold_t, T>::value;
+ /// Trait to check if T provides a 'fold' method with a FoldAdaptor that
+ /// returns OpFoldResults.
+ template <typename T, typename... Args>
+ using has_fold_adaptor_results_fold_t = std::enable_if_t<std::is_same_v<
+ decltype(std::declval<T>().fold(std::declval<typename T::FoldAdaptor>())),
+ OpFoldResults>>;
+ template <class T>
+ constexpr static bool has_fold_adaptor_results_v =
+ llvm::is_detected<has_fold_adaptor_results_fold_t, T>::value;
/// Trait to check if T provides a 'print' method.
template <typename T, typename... Args>
@@ -1933,43 +1920,76 @@ 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 = hasTrait<OpTrait::OneResult>();
+ constexpr bool hasResultsFold = has_fold_adaptor_results_v<ConcreteType>;
+ // `OpFoldResults fold(FoldAdaptor)` also matches the single-result
+ // detectors.
+ constexpr bool hasSingleResultFold =
+ !hasResultsFold && (has_single_result_fold_v<ConcreteType> ||
+ has_fold_adaptor_single_result_v<ConcreteType>);
+ constexpr bool hasLegacyFold =
+ has_fold_v<ConcreteType> || has_fold_adaptor_v<ConcreteType>;
+ // Without this check, the dispatch below silently skips such a `fold`.
+ static_assert(hasOneResult || !hasSingleResultFold,
+ "a single-result `OpFoldResult fold(...)` requires the "
+ "OneResult trait; ops with zero, several, or variadic "
+ "results must use `OpFoldResults fold(FoldAdaptor)` or "
+ "`LogicalResult fold(FoldAdaptor, "
+ "SmallVectorImpl<OpFoldResult> &)`");
+ static_assert(!(hasResultsFold && hasLegacyFold),
+ "an op must not define both `OpFoldResults "
+ "fold(FoldAdaptor)` and a legacy `LogicalResult fold(..., "
+ "SmallVectorImpl<OpFoldResult> &)`");
+ static_assert(!(hasResultsFold && hasOneResult),
+ "an op with the OneResult trait must use the single-result "
+ "`OpFoldResult fold(FoldAdaptor)` form instead of "
+ "`OpFoldResults fold(FoldAdaptor)`");
+ // The operation defines an `OpFoldResults fold(FoldAdaptor)` method.
+ if constexpr (hasResultsFold)
+ return [](Operation *op, ArrayRef<Attribute> operands) {
+ return foldResultsHook<ConcreteType>(op, operands);
+ };
// If the operation is single result and defines a `fold` method.
- if constexpr (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 foldHook<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);
+ ::mlir::detail::dropReplacementsOfReplacedResults(op, result);
+ return op_definition_impl::foldTraits<Traits<ConcreteType>...>(
+ op, operands, std::move(result));
+ }
/// Return the result of folding a single result operation that defines a
/// `fold` method.
template <typename ConcreteOpT>
- 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 +1998,34 @@ 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);
+ return op_definition_impl::foldTraits<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 foldHook(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::foldTraits<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 000000000000000..749ef8a6bfc0bb2
--- /dev/null
+++ b/mlir/include/mlir/IR/OpFoldResult.h
@@ -0,0 +1,187 @@
+//===- 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
+// 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/Compiler.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 any op. An op with exactly one fixed result defines
+/// its folder with OpFoldResult; any other op can return OpFoldResults
+/// directly. 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;
+ /// Failure, the same as the default constructor.
+ 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, or the op's own result, keeps the result, so the fold
+ /// fails. Unlike `OpFoldResult fold`, the op's own result does not mean in
+ /// place; use success() or setModifiedInPlace() for an in-place change.
+ 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);
+ }
+ /// 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. This function
+ /// checks only that `result` is an OpResult, not which op owns it. A null
+ /// replacement, or `replacement == result`, resets the slot to keep the
+ /// result. The last write wins.
+ void replace(Value result, OpFoldResult replacement);
+ /// Same as above, but the object must already have one slot per result.
+ /// normalize() maps a replacement with the op's own result to "keep".
+ void replace(unsigned resultIndex, OpFoldResult replacement);
+ /// Set whether the fold changed the op in place (operands, attributes,
+ /// properties, or regions).
+ void setModifiedInPlace(bool modified = true);
+
+ // Queries for drivers. They are valid after normalize().
+
+ /// Return true if the fold changed the op in place or replaces a result.
+ bool succeeded() const;
+ /// Return true if the fold did not apply.
+ bool failed() const;
+ /// Return true if the fold changed the op in place.
+ bool modifiedInPlace() const;
+ /// Return true if at least one slot replaces its result.
+ 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;
+ /// Return the number of slots: 0, or one per result of the op.
+ unsigned size() const;
+ /// Return the slot of result `i`; null means keep. If size() == 0, every
+ /// index reads as keep; otherwise `i` must be less than size().
+ OpFoldResult operator[](unsigned i) const;
+ /// Return the slots. The returned range points into this object.
+ ArrayRef<OpFoldResult> getSlots() const LLVM_LIFETIME_BOUND;
+
+ /// Map a slot that holds the op's own result to "keep", and clear the slots
+ /// if no slot replaces its result. In debug builds, check that a non-empty
+ /// result has one slot per op result, and that each Value slot has the type
+ /// of its result. Idempotent.
+ void normalize(Operation *op);
+
+private:
+ SmallVector<OpFoldResult, 2> slots;
+ bool inPlace = false;
+};
+
+/// Return true if `result` changed the op in place or replaces a result.
+inline bool succeeded(const OpFoldResults &result) {
+ return result.succeeded();
+}
+
+/// Return true if `result` is a failure.
+inline bool failed(const OpFoldResults &result) { return result.failed(); }
+
+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 has one slot per result. normalize() then maps an
+/// entry that is the op's own result to keep. A null entry is a bug and asserts
+/// in debug builds.
+OpFoldResults convertLegacyFoldResults(LogicalResult status,
+ ArrayRef<OpFoldResult> results);
+
+/// Drop the replacements of the normalized `result` if a slot names another
+/// result of `op` that `result` also replaces, because the outcome would depend
+/// on the order in which a driver replaces the results. A forwarding fold can
+/// return such a result in a graph region, where an op can use its own
+/// results. Outside a graph region, it is a bug in the fold and asserts. The
+/// in-place bit stays.
+void dropReplacementsOfReplacedResults(Operation *op, OpFoldResults &result);
+} // 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 3ea9ad8ca0cfc43..086bd0a8ad81a96 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -767,6 +767,21 @@ class alignas(8) Operation final
// Accessors for various properties of operations
//===--------------------------------------------------------------------===//
+ /// 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.
+ ///
+ /// 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.
@@ -776,6 +791,9 @@ class alignas(8) Operation final
/// `results` is empty.
/// * Otherwise, `results` is filled with the folded results.
/// If folding was unsuccessful, this function returns "failure".
+ /// If the fold replaces only some results, those replacements are dropped:
+ /// this function returns "success" with an empty `results` if this operation
+ /// was modified in place, and "failure" otherwise.
LogicalResult fold(ArrayRef<Attribute> operands,
SmallVectorImpl<OpFoldResult> &results);
@@ -786,6 +804,9 @@ class alignas(8) Operation final
/// `results` is empty.
/// * Otherwise, `results` is filled with the folded results.
/// If folding was unsuccessful, this function returns "failure".
+ /// If the fold replaces only some results, those replacements are dropped:
+ /// this function returns "success" with an empty `results` if this operation
+ /// was modified in place, and "failure" otherwise.
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 991d9847672e2e9..7e9c7c60858542c 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,35 @@ 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 either no slot or 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. In a graph region, a forwarding fold can break this
+ /// rule; then its replacements are dropped.
+ /// 2. A failure replaces no result and has no in-place mark. The IR must be
+ /// unchanged.
+ /// 3. The hook can mutate the operation in place, without changing anything
+ /// 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 but not all results and also mutate the
+ /// operation in place. The operation stays.
+ /// 5. If the operation has results and every slot is replaced, the caller
+ /// removes the operation and uses the replacements instead, even if the
+ /// hook also mutated it in place.
+ ///
+ /// 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 +609,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 01c80dafcd21a78..2578cc086c3c664 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 b51d3301606468e..c6b36a765973b55 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,32 @@ 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));
+}
+
+void DynamicOpDefinition::setFoldHookFn(std::nullptr_t) {
+ foldHookFn = [](Operation *, ArrayRef<Attribute>) -> OpFoldResults {
+ return failure();
+ };
+}
+
//===----------------------------------------------------------------------===//
// Extensible dialect
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 94cf420070131c9..f64240bb0a583ac 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 000000000000000..d047ac913163920
--- /dev/null
+++ b/mlir/lib/IR/OpFoldResult.cpp
@@ -0,0 +1,153 @@
+//===- 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"
+#include "mlir/IR/RegionKindInterface.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()) {}
+
+OpFoldResults::OpFoldResults(OpFoldResult replacement) {
+ if (replacement)
+ slots.push_back(replacement);
+}
+
+OpFoldResults::OpFoldResults(Value replacement)
+ : OpFoldResults(OpFoldResult(replacement)) {}
+
+OpFoldResults::OpFoldResults(Attribute replacement)
+ : OpFoldResults(OpFoldResult(replacement)) {}
+
+OpFoldResults::OpFoldResults(std::initializer_list<OpFoldResult> replacements)
+ : slots(replacements) {}
+
+OpFoldResults::OpFoldResults(Operation *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(opResult && "expected an OpResult");
+ unsigned numResults = opResult.getOwner()->getNumResults();
+ if (slots.empty())
+ slots.resize(numResults);
+ assert(slots.size() == numResults && "expected one slot per result");
+ if (dyn_cast_if_present<Value>(replacement) == result)
+ replacement = OpFoldResult();
+ replace(opResult.getResultNumber(), replacement);
+}
+
+void OpFoldResults::replace(unsigned resultIndex, OpFoldResult replacement) {
+ assert(resultIndex < slots.size() && "result index out of range");
+ slots[resultIndex] = replacement;
+}
+
+void OpFoldResults::setModifiedInPlace(bool modified) { inPlace = modified; }
+
+bool OpFoldResults::succeeded() const { return inPlace || replacesAny(); }
+
+bool OpFoldResults::failed() const { return !succeeded(); }
+
+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())
+ 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");
+ if (!slots.empty()) {
+ assert(slots.size() == op->getNumResults() &&
+ "expected one fold result slot per operation result");
+ for (auto [slot, result] : llvm::zip_equal(slots, op->getResults()))
+ if (dyn_cast_if_present<Value>(slot) == result)
+ slot = OpFoldResult();
+ if (replacesNone(slots))
+ slots.clear();
+ }
+
+#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");
+ }
+ }
+#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;
+}
+
+void detail::dropReplacementsOfReplacedResults(Operation *op,
+ OpFoldResults &result) {
+ bool namesReplacedResult =
+ llvm::any_of(result.getSlots(), [&](OpFoldResult slot) {
+ auto opResult =
+ dyn_cast_if_present<OpResult>(dyn_cast_if_present<Value>(slot));
+ return opResult && opResult.getOwner() == op &&
+ result[opResult.getResultNumber()];
+ });
+ if (!namesReplacedResult)
+ return;
+ assert(op->getParentRegion() && mayBeGraphRegion(*op->getParentRegion()) &&
+ "fold result names a result of the folded operation that the same "
+ "fold replaces");
+ result = success(result.modifiedInPlace());
+}
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index d262e844603a1a1..52b92c3d57dfba6 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 ce9b230222f6212..e6858585aaf6c49 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 000000000000000..3183760b53b37a8
--- /dev/null
+++ b/mlir/unittests/IR/OpFoldResultsTest.cpp
@@ -0,0 +1,1059 @@
+//===- 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/Dialect.h"
+#include "mlir/IR/ExtensibleDialect.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpFoldResult.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/FoldInterfaces.h"
+#include "gtest/gtest.h"
+
+#include <functional>
+
+using namespace mlir;
+
+// The fallback TypeID resolver rejects a trait template in an anonymous
+// namespace, so the test dialect lives in a named namespace.
+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.
+struct FoldState {
+ std::function<OpFoldResults(Operation *)> opFoldFn;
+ LegacyFoldFn traitFoldFn;
+ LegacyFoldFn dialectFoldFn;
+ SmallVector<Attribute> seenOperands;
+ unsigned traitCalls = 0;
+ unsigned dialectCalls = 0;
+};
+
+/// The state of the running test. The test fixture owns it.
+static FoldState *foldState = nullptr;
+
+template <typename ConcreteType>
+struct LegacyFoldTrait
+ : public OpTrait::TraitBase<ConcreteType, LegacyFoldTrait> {
+ static LogicalResult foldTrait(Operation *op, ArrayRef<Attribute>,
+ SmallVectorImpl<OpFoldResult> &results) {
+ ++foldState->traitCalls;
+ return foldState->traitFoldFn ? foldState->traitFoldFn(op, results)
+ : failure();
+ }
+};
+
+template <typename OpT>
+struct FoldAdaptorImpl {
+ FoldAdaptorImpl(ArrayRef<Attribute> operands, OpT) : operands(operands) {}
+ ArrayRef<Attribute> getOperands() const { return operands; }
+ ArrayRef<Attribute> operands;
+};
+
+/// An op with two results that defines `OpFoldResults fold(FoldAdaptor)`.
+struct PartialFoldOp : public Op<PartialFoldOp, OpTrait::NResults<2>::Impl,
+ OpTrait::VariadicOperands, LegacyFoldTrait> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PartialFoldOp)
+ using Op::Op;
+ using FoldAdaptor = FoldAdaptorImpl<PartialFoldOp>;
+ static ArrayRef<StringRef> getAttributeNames() { return {}; }
+ static StringRef getOperationName() { return "fold_test.partial"; }
+ OpFoldResults fold(FoldAdaptor adaptor) {
+ foldState->seenOperands.assign(adaptor.getOperands().begin(),
+ adaptor.getOperands().end());
+ return foldState->opFoldFn ? foldState->opFoldFn(getOperation())
+ : failure();
+ }
+};
+
+struct ConstantOp : public Op<ConstantOp, OpTrait::OneResult,
+ OpTrait::ZeroOperands, OpTrait::ConstantLike> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConstantOp)
+ using Op::Op;
+ using FoldAdaptor = FoldAdaptorImpl<ConstantOp>;
+ static ArrayRef<StringRef> getAttributeNames() {
+ static StringRef names[] = {"value"};
+ return names;
+ }
+ static StringRef getOperationName() { return "fold_test.constant"; }
+ OpFoldResult fold(FoldAdaptor) { return getOperation()->getAttr("value"); }
+};
+
+struct TestFoldInterface : public DialectFoldInterface {
+ using DialectFoldInterface::DialectFoldInterface;
+ LogicalResult fold(Operation *op, ArrayRef<Attribute>,
+ SmallVectorImpl<OpFoldResult> &results) const final {
+ ++foldState->dialectCalls;
+ return foldState->dialectFoldFn ? foldState->dialectFoldFn(op, results)
+ : failure();
+ }
+};
+
+struct FoldTestDialect : public Dialect {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FoldTestDialect)
+ static constexpr StringLiteral getDialectNamespace() { return "fold_test"; }
+ explicit FoldTestDialect(MLIRContext *context)
+ : Dialect(getDialectNamespace(), context,
+ TypeID::get<FoldTestDialect>()) {
+ addOperations<PartialFoldOp, ConstantOp>();
+ addInterfaces<TestFoldInterface>();
+ }
+};
+} // namespace op_fold_results_test
+
+using namespace op_fold_results_test;
+
+namespace {
+class OpFoldResultsTest : public ::testing::Test {
+protected:
+ OpFoldResultsTest() : builder(&context) {
+ context.allowUnregisteredDialects();
+ context.loadDialect<FoldTestDialect>();
+ i32 = builder.getI32Type();
+ f32 = builder.getF32Type();
+ foldState = &state;
+ }
+
+ ~OpFoldResultsTest() override {
+ // Destroy users before the ops that define their operands.
+ for (Operation *op : llvm::reverse(ops))
+ op->destroy();
+ foldState = nullptr;
+ }
+
+ /// Create an op with the given result types and operands. The fixture
+ /// destroys it.
+ Operation *createOp(TypeRange resultTypes, StringRef name = "foo.bar",
+ ValueRange operands = {}) {
+ OperationState state(UnknownLoc::get(&context), name);
+ state.addTypes(resultTypes);
+ state.addOperands(operands);
+ Operation *op = Operation::create(state);
+ ops.push_back(op);
+ return op;
+ }
+
+ /// Load a dynamic dialect. The fold hook of `test_fold.op` calls `foldFn`.
+ /// The legacy fold hooks of `test_fold.legacy_op` and `test_fold.get_op`
+ /// call `legacyFoldFn`; `test_fold.get_op` uses the legacy overload of
+ /// DynamicOpDefinition::get. `test_fold.no_fold_op` has no fold hook.
+ 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));
+
+ DynamicOpDefinition::LegacyFoldHookFn legacyFold =
+ [this](Operation *op, ArrayRef<Attribute>,
+ SmallVectorImpl<OpFoldResult> &results) {
+ return legacyFoldFn(op, results);
+ };
+ dialect->registerDynamicOp(DynamicOpDefinition::get(
+ "get_op", dialect, verify, verify,
+ [](OpAsmParser &, OperationState &) -> ParseResult {
+ return failure();
+ },
+ [](Operation *, OpAsmPrinter &, StringRef) {}, std::move(legacyFold),
+ [](RewritePatternSet &, MLIRContext *) {},
+ [](const OperationName &, NamedAttrList &) {}));
+
+ std::unique_ptr<DynamicOpDefinition> noFoldOpDef =
+ DynamicOpDefinition::get("no_fold_op", dialect, verify, verify);
+ noFoldOpDef->setFoldHookFn(
+ [this](Operation *op, ArrayRef<Attribute>) { return foldFn(op); });
+ noFoldOpDef->setFoldHookFn(nullptr);
+ dialect->registerDynamicOp(std::move(noFoldOpDef));
+ });
+ }
+
+ MLIRContext context;
+ Builder builder;
+ Type i32;
+ Type f32;
+ SmallVector<Operation *> ops;
+ std::function<OpFoldResults(Operation *)> foldFn;
+ std::function<LogicalResult(Operation *, SmallVectorImpl<OpFoldResult> &)>
+ legacyFoldFn;
+ FoldState state;
+};
+} // 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));
+
+ // A null replacement, or the result itself, resets the slot to keep the
+ // result.
+ result.replace(op->getResult(0), OpFoldResult());
+ EXPECT_FALSE(result[0]);
+ result.replace(op->getResult(1), op->getResult(1));
+ EXPECT_FALSE(result[1]);
+ EXPECT_TRUE(result.failed());
+ EXPECT_FALSE(result.replacesAny());
+
+ result.setModifiedInPlace();
+ EXPECT_TRUE(result.succeeded());
+ EXPECT_TRUE(result.modifiedInPlace());
+ result.setModifiedInPlace(false);
+ EXPECT_TRUE(result.failed());
+ EXPECT_FALSE(result.modifiedInPlace());
+
+ // normalize() maps the op's own result in an indexed slot to "keep".
+ OpFoldResults ownResult(op);
+ ownResult.replace(1u, op->getResult(1));
+ ownResult.normalize(op);
+ EXPECT_TRUE(ownResult.failed());
+ EXPECT_EQ(ownResult.size(), 0u);
+
+ OpFoldResults partial(op);
+ partial.replace(1u, value);
+ partial.normalize(op);
+ EXPECT_TRUE(partial.succeeded());
+ ASSERT_EQ(partial.size(), 2u);
+ EXPECT_FALSE(partial[0]);
+ EXPECT_EQ(partial[1], OpFoldResult(value));
+ ArrayRef<OpFoldResult> slots = partial.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, ReplaceAfterNormalizeCreatesOneSlotPerResult) {
+ 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() with a result on a normalized object creates one slot per
+ // result.
+ inPlace.replace(op->getResult(1), 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(op->getResult(0), 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.setModifiedInPlace();
+ 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.setModifiedInPlace();
+ 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.setModifiedInPlace();
+ 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());
+}
+
+TEST_F(OpFoldResultsTest, LegacyDynamicFoldHookMayForwardAnotherResult) {
+ loadDynamicDialect();
+ Operation *op = createOp({i32, i32}, "test_fold.legacy_op");
+ Attribute attr = builder.getI32IntegerAttr(1);
+
+ // Slot 0 forwards result 1, which slot 1 replaces. The replaced-result check
+ // applies only to `OpFoldResults fold(FoldAdaptor)`.
+ legacyFoldFn = [&](Operation *foldedOp,
+ SmallVectorImpl<OpFoldResult> &results) {
+ results.push_back(foldedOp->getResult(1));
+ results.push_back(attr);
+ return success();
+ };
+ 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));
+}
+
+// In a graph region, a forwarding fold can name another result of its op. If
+// the same fold also replaces that result, the fold does not apply.
+TEST_F(OpFoldResultsTest, GraphRegionFoldNamingReplacedResultDoesNotApply) {
+ // The region of an unregistered op may be a graph region.
+ OperationState parentState(UnknownLoc::get(&context), "foo.graph");
+ parentState.addRegion();
+ Operation *parent = Operation::create(parentState);
+ ops.push_back(parent);
+ Block *block = new Block();
+ parent->getRegion(0).push_back(block);
+ OperationState state(UnknownLoc::get(&context), "fold_test.partial");
+ state.addTypes({i32, i32});
+ Operation *op = Operation::create(state);
+ block->push_back(op);
+ Attribute attr = builder.getI32IntegerAttr(1);
+
+ foldState->opFoldFn = [&](Operation *foldedOp) -> OpFoldResults {
+ return {foldedOp->getResult(1), attr};
+ };
+ EXPECT_TRUE(op->fold().failed());
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+
+ // The in-place bit stays.
+ foldState->opFoldFn = [&](Operation *foldedOp) {
+ OpFoldResults result = {foldedOp->getResult(1), attr};
+ result.setModifiedInPlace();
+ return result;
+ };
+ OpFoldResults result = op->fold();
+ EXPECT_TRUE(result.modifiedInPlace());
+ EXPECT_FALSE(result.replacesAny());
+}
+
+TEST_F(OpFoldResultsTest, LegacyDynamicOpDefinitionGet) {
+ loadDynamicDialect();
+ Operation *op = createOp({i32}, "test_fold.get_op");
+ Attribute attr = builder.getI32IntegerAttr(1);
+
+ legacyFoldFn = [&](Operation *, SmallVectorImpl<OpFoldResult> &results) {
+ results.push_back(attr);
+ return success();
+ };
+ OpFoldResults result = op->fold();
+ expectOneSlot(result, attr);
+
+ legacyFoldFn = [](Operation *, SmallVectorImpl<OpFoldResult> &) {
+ return failure();
+ };
+ EXPECT_TRUE(op->fold().failed());
+}
+
+TEST_F(OpFoldResultsTest, NullFoldHookFails) {
+ loadDynamicDialect();
+ Operation *op = createOp({i32}, "test_fold.no_fold_op");
+ // The removed hook would report an in-place fold.
+ foldFn = [](Operation *) -> OpFoldResults { return success(); };
+ EXPECT_TRUE(op->fold().failed());
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+}
+
+TEST_F(OpFoldResultsTest, FreeHelpersMatchMembers) {
+ Operation *op = createOp({i32, i32});
+ OpFoldResults partial(op);
+ partial.replace(1u, builder.getI32IntegerAttr(1));
+ OpFoldResults inPlace = success();
+ OpFoldResults failedResult = failure();
+ for (const OpFoldResults *result : {&partial, &inPlace, &failedResult}) {
+ EXPECT_EQ(succeeded(*result), result->succeeded());
+ EXPECT_EQ(failed(*result), result->failed());
+ }
+ EXPECT_TRUE(succeeded(partial));
+ EXPECT_TRUE(succeeded(inPlace));
+ EXPECT_TRUE(failed(failedResult));
+}
+
+TEST_F(OpFoldResultsTest, OwnPartialFoldSkipsTraits) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ foldState->opFoldFn = [&](Operation *foldedOp) {
+ OpFoldResults result(foldedOp);
+ result.replace(1u, attr);
+ return result;
+ };
+ foldState->traitFoldFn = [](Operation *, SmallVectorImpl<OpFoldResult> &) {
+ return success();
+ };
+
+ OpFoldResults result = op->fold();
+ EXPECT_TRUE(result.succeeded());
+ EXPECT_FALSE(result.modifiedInPlace());
+ EXPECT_FALSE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_FALSE(result[0]);
+ EXPECT_EQ(result[1], OpFoldResult(attr));
+
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+ EXPECT_EQ(foldState->traitCalls, 0u);
+}
+
+TEST_F(OpFoldResultsTest, OwnInPlaceThenTraitReplacesAll) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute lhs = builder.getI32IntegerAttr(1);
+ Attribute rhs = builder.getI32IntegerAttr(2);
+ foldState->opFoldFn = [](Operation *) -> OpFoldResults { return success(); };
+ foldState->traitFoldFn = [&](Operation *,
+ SmallVectorImpl<OpFoldResult> &results) {
+ results.push_back(lhs);
+ results.push_back(rhs);
+ return success();
+ };
+
+ OpFoldResults result = op->fold();
+ EXPECT_TRUE(result.modifiedInPlace());
+ EXPECT_TRUE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_EQ(result[0], OpFoldResult(lhs));
+ EXPECT_EQ(result[1], OpFoldResult(rhs));
+ EXPECT_EQ(foldState->traitCalls, 1u);
+
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ ASSERT_EQ(results.size(), 2u);
+ EXPECT_EQ(results[0], OpFoldResult(lhs));
+ EXPECT_EQ(results[1], OpFoldResult(rhs));
+ EXPECT_EQ(foldState->traitCalls, 2u);
+}
+
+TEST_F(OpFoldResultsTest, OwnInPlaceThenTraitReplacesSome) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ foldState->opFoldFn = [](Operation *) -> OpFoldResults { return success(); };
+ foldState->traitFoldFn = [&](Operation *foldedOp,
+ SmallVectorImpl<OpFoldResult> &results) {
+ results.push_back(foldedOp->getResult(0));
+ results.push_back(attr);
+ return success();
+ };
+
+ OpFoldResults result = op->fold();
+ EXPECT_TRUE(result.modifiedInPlace());
+ EXPECT_TRUE(result.replacesAny());
+ EXPECT_FALSE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_FALSE(result[0]);
+ EXPECT_EQ(result[1], OpFoldResult(attr));
+ EXPECT_EQ(foldState->traitCalls, 1u);
+
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+}
+
+TEST_F(OpFoldResultsTest, OwnFailureThenTraitInPlace) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ foldState->opFoldFn = [](Operation *) -> OpFoldResults { return failure(); };
+ foldState->traitFoldFn = [](Operation *, SmallVectorImpl<OpFoldResult> &) {
+ return success();
+ };
+
+ OpFoldResults result = op->fold();
+ EXPECT_TRUE(result.modifiedInPlace());
+ EXPECT_FALSE(result.replacesAny());
+ EXPECT_EQ(foldState->traitCalls, 1u);
+
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ EXPECT_TRUE(results.empty());
+}
+
+TEST_F(OpFoldResultsTest, OwnResultsNormalizeToFailureSoTraitsRun) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ foldState->opFoldFn = [](Operation *foldedOp) -> OpFoldResults {
+ return foldedOp->getResults();
+ };
+ foldState->traitFoldFn = [&](Operation *,
+ SmallVectorImpl<OpFoldResult> &results) {
+ results.append(2, attr);
+ return success();
+ };
+
+ OpFoldResults result = op->fold();
+ EXPECT_FALSE(result.modifiedInPlace());
+ EXPECT_TRUE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_EQ(result[0], OpFoldResult(attr));
+ EXPECT_EQ(result[1], OpFoldResult(attr));
+ EXPECT_EQ(foldState->traitCalls, 1u);
+
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ ASSERT_EQ(results.size(), 2u);
+ EXPECT_EQ(results[0], OpFoldResult(attr));
+ EXPECT_EQ(results[1], OpFoldResult(attr));
+ EXPECT_EQ(foldState->traitCalls, 2u);
+}
+
+TEST_F(OpFoldResultsTest, DialectFoldInterfaceFallback) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ foldState->dialectFoldFn = [&](Operation *,
+ SmallVectorImpl<OpFoldResult> &results) {
+ results.append(2, attr);
+ return success();
+ };
+
+ // The op fold and the traits fail.
+ OpFoldResults result = op->fold();
+ EXPECT_TRUE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_EQ(result[0], OpFoldResult(attr));
+ EXPECT_EQ(foldState->traitCalls, 1u);
+ EXPECT_EQ(foldState->dialectCalls, 1u);
+ SmallVector<OpFoldResult> results;
+ EXPECT_TRUE(succeeded(op->fold(results)));
+ EXPECT_EQ(results.size(), 2u);
+ EXPECT_EQ(foldState->dialectCalls, 2u);
+
+ // A partial fold does not fall back.
+ foldState->opFoldFn = [&](Operation *foldedOp) {
+ OpFoldResults partial(foldedOp);
+ partial.replace(0u, attr);
+ return partial;
+ };
+ foldState->dialectCalls = 0;
+ result = op->fold();
+ EXPECT_TRUE(result.succeeded());
+ EXPECT_FALSE(result.replacesAll());
+ ASSERT_EQ(result.size(), 2u);
+ EXPECT_EQ(result[0], OpFoldResult(attr));
+ EXPECT_FALSE(result[1]);
+ EXPECT_EQ(foldState->dialectCalls, 0u);
+ results.clear();
+ EXPECT_TRUE(failed(op->fold(results)));
+ EXPECT_EQ(foldState->dialectCalls, 0u);
+
+ // The fallback result is normalized.
+ foldState->opFoldFn = nullptr;
+ foldState->dialectFoldFn = [](Operation *foldedOp,
+ SmallVectorImpl<OpFoldResult> &results) {
+ llvm::append_range(results, foldedOp->getResults());
+ return success();
+ };
+ EXPECT_TRUE(op->fold().failed());
+ EXPECT_EQ(foldState->dialectCalls, 1u);
+}
+
+TEST_F(OpFoldResultsTest, FoldComputesConstantOperands) {
+ Operation *constant = createOp({i32}, "fold_test.constant");
+ Attribute attr = builder.getI32IntegerAttr(7);
+ constant->setAttr("value", attr);
+ Operation *producer = createOp({i32});
+ Operation *op =
+ createOp({i32, i32}, "fold_test.partial",
+ ValueRange{constant->getResult(0), producer->getResult(0)});
+
+ EXPECT_TRUE(op->fold().failed());
+ ASSERT_EQ(foldState->seenOperands.size(), 2u);
+ EXPECT_EQ(foldState->seenOperands[0], attr);
+ EXPECT_FALSE(foldState->seenOperands[1]);
+}
+
+#ifdef GTEST_HAS_DEATH_TEST
+#ifndef NDEBUG
+namespace {
+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.
+ OpFoldResults result(op);
+ result.replace(0u, op->getResult(1));
+ result.replace(1u, producer->getResult(0));
+ result.normalize(op);
+ EXPECT_DEATH(detail::dropReplacementsOfReplacedResults(op, result),
+ "fold result names a result of the folded operation");
+}
+
+TEST_F(OpFoldResultsDeathTest, FoldNamesReplacedResult) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ foldState->opFoldFn = [&](Operation *foldedOp) -> OpFoldResults {
+ return {foldedOp->getResult(1), attr};
+ };
+ EXPECT_DEATH((void)op->fold(),
+ "fold result names a result of the folded operation");
+}
+
+TEST_F(OpFoldResultsDeathTest, ValueSlotOfIncorrectType) {
+ Operation *producer = createOp({f32});
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ foldState->opFoldFn = [&](Operation *foldedOp) {
+ OpFoldResults result(foldedOp);
+ result.replace(0u, producer->getResult(0));
+ return result;
+ };
+ EXPECT_DEATH((void)op->fold(), "incorrect fold result type");
+}
+
+TEST_F(OpFoldResultsDeathTest, SlotCountMismatch) {
+ Operation *op = createOp({i32, i32}, "fold_test.partial");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ foldState->opFoldFn = [&](Operation *) -> OpFoldResults {
+ return {attr, attr, attr};
+ };
+ EXPECT_DEATH((void)op->fold(),
+ "expected one fold result slot per operation result");
+}
+
+TEST_F(OpFoldResultsDeathTest, LegacyNullResult) {
+ loadDynamicDialect();
+ Operation *op = createOp({i32, i32}, "test_fold.legacy_op");
+ Attribute attr = builder.getI32IntegerAttr(1);
+ legacyFoldFn = [&](Operation *, SmallVectorImpl<OpFoldResult> &results) {
+ results.push_back(attr);
+ results.push_back(OpFoldResult());
+ return success();
+ };
+ EXPECT_DEATH((void)op->fold(), "legacy fold returned a null result");
+}
+#endif // NDEBUG
+#endif // GTEST_HAS_DEATH_TEST
More information about the Mlir-commits
mailing list