[llvm-branch-commits] [mlir] [mlir] Add an `OpBuilder::tryFold` overload that returns `OpFoldResults` (PR #228801)
Victor Perez via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Sun Oct 4 15:04:59 PDT 2026
https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/228801
>From 17b93037ad0cd9b0070f34627bc28e1783d6be36 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: Sat, 3 Oct 2026 18:11:31 -0700
Subject: [PATCH] [mlir] Add an OpBuilder::tryFold overload that returns
OpFoldResults
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
`OpBuilder::tryFold(op, results, constants)` applies a fold only if the
fold replaces every result, so it cannot give a partial fold to the
caller. This patch adds two `OpBuilder` functions that let a caller
apply a partial fold:
- `OpFoldResults tryFold(Operation *op)` folds the op and returns the
fold result. Like the old overload, it skips constants and repeats a
fold that only changes the op in place. It creates no constant.
- `materializeFoldResults(op, foldResults, liveOnly)` returns one value
per result: the replacement of a replaced result, or null for a kept
result. It turns each attribute into a new constant at the insertion
point. If `liveOnly` is set, a replaced result without uses gets null
and no constant. If one constant fails to materialize, the function
inserts no constant and fails.
The old overload now uses the two new functions. A fold that replaces
only some results, which only the new fold form can return, counts as
an in-place fold if it changed the op in place, and as a failure
otherwise. Before this patch, the in-place loop folded the op again
after a partial fold that changed the op in place. Now the old overload
stops after such a fold.
The test op `transform.test_fold` gets an `api` attribute: "try_fold"
folds with the new overload, and "materialize" also materializes the
replacements, optionally only the ones with uses. The new lit test
`IR/op-builder-fold.mlir` covers the new functions. A new unit test in
`OpFoldResultsTest.cpp` covers the old overload.
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
mlir/include/mlir/IR/Builders.h | 28 ++++-
mlir/lib/IR/Builders.cpp | 92 +++++++++------
mlir/test/IR/op-builder-fold.mlir | 110 ++++++++++++++++++
.../TestTransformDialectExtension.cpp | 28 ++++-
.../TestTransformDialectExtension.td | 21 +++-
mlir/unittests/IR/OpFoldResultsTest.cpp | 64 +++++++++-
6 files changed, 300 insertions(+), 43 deletions(-)
create mode 100644 mlir/test/IR/op-builder-fold.mlir
diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index 0bf41530b14b0..413b9c54ea94a 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -572,10 +572,36 @@ class OpBuilder : public Builder {
return op;
}
+ /// Try to fold `op`. A fold that only changes `op` in place repeats until a
+ /// fold replaces a result or fails. Return one of:
+ /// - failure if `op` is a constant or does not fold. After 64 folds that
+ /// only change `op` in place, this function also returns failure, and the
+ /// changes to `op` stay;
+ /// - success() if the folds only changed `op` in place;
+ /// - the replacements of the last fold, with the in-place bit set if a fold
+ /// changed `op` in place.
+ /// This function creates no constant and replaces no result; see
+ /// materializeFoldResults.
+ OpFoldResults tryFold(Operation *op);
+
+ /// Materialize the replacements in `foldResults`, a fold result of `op`. On
+ /// success, return one value per result of `op`:
+ /// - a replaced result gets its replacement. An attribute becomes a new
+ /// constant at the insertion point, and the value is its result;
+ /// - a kept result gets null;
+ /// - if `liveOnly` is set, a replaced result without uses also gets null and
+ /// no constant.
+ /// If a constant fails to materialize, insert no constant and return failure.
+ FailureOr<SmallVector<Value>>
+ materializeFoldResults(Operation *op, const OpFoldResults &foldResults,
+ bool liveOnly);
+
/// Attempts to fold the given operation and places new results within
/// `results`. Returns success if the operation was folded, failure otherwise.
/// If the fold was in-place, `results` will not be filled. Optionally, newly
- /// materialized constant operations can be returned to the caller.
+ /// materialized constant operations can be returned to the caller. A fold
+ /// that replaces only some results counts as an in-place fold if it changed
+ /// the operation in place, and as a failure otherwise.
///
/// Note: This function does not erase the operation on a successful fold.
LogicalResult
diff --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp
index 47e09faf0e04f..d6c210b8e9f1e 100644
--- a/mlir/lib/IR/Builders.cpp
+++ b/mlir/lib/IR/Builders.cpp
@@ -478,92 +478,118 @@ Operation *OpBuilder::create(Location loc, StringAttr opName,
return create(state);
}
-LogicalResult
-OpBuilder::tryFold(Operation *op, SmallVectorImpl<Value> &results,
- SmallVectorImpl<Operation *> *materializedConstants) {
- assert(results.empty() && "expected empty results");
- ResultRange opResults = op->getResults();
-
- results.reserve(opResults.size());
- auto cleanupFailure = [&] {
- results.clear();
- return failure();
- };
-
+OpFoldResults OpBuilder::tryFold(Operation *op) {
// If this operation is already a constant, there is nothing to do.
if (matchPattern(op, m_Constant()))
- return cleanupFailure();
+ return failure();
// Try to fold the operation.
- SmallVector<OpFoldResult, 4> foldResults;
LDBG() << "Trying to fold: "
<< OpWithFlags(op, OpPrintingFlags().skipRegions());
- if (failed(op->fold(foldResults)))
- return cleanupFailure();
+ OpFoldResults foldResults = op->fold();
+ if (foldResults.failed())
+ return failure();
// Bound the number of in-place fold iterations. Legitimate chains are very
// short (e.g. foldCommutative swaps once, then the op folds to a value).
// Without a bound, circular SSA uses in graph regions can cause an infinite
// loop (e.g. addi(x, 0) where x is the op's own result).
constexpr int kMaxInPlaceFolds = 64;
+ bool modifiedInPlace = false;
int count = 0;
do {
+ modifiedInPlace |= foldResults.modifiedInPlace();
LDBG() << "Folded in place #" << count
<< " times: " << OpWithFlags(op, OpPrintingFlags().skipRegions());
if (++count >= kMaxInPlaceFolds) {
LDBG() << "Aborting after " << kMaxInPlaceFolds
<< " in-place fold iterations: "
<< OpWithFlags(op, OpPrintingFlags().skipRegions());
- return cleanupFailure();
+ return failure();
}
- } while (foldResults.empty() && succeeded(op->fold(foldResults)));
+ } while (!foldResults.replacesAny() &&
+ (foldResults = op->fold()).succeeded());
// An in-place fold does not require generation of any constants.
- if (foldResults.empty())
+ if (!foldResults.replacesAny())
return success();
+ foldResults.setModifiedInPlace(modifiedInPlace);
+ return foldResults;
+}
+
+FailureOr<SmallVector<Value>> OpBuilder::materializeFoldResults(
+ Operation *op, const OpFoldResults &foldResults, bool liveOnly) {
+ SmallVector<Value> replacements(op->getNumResults());
+
// A temporary builder used for creating constants during folding.
OpBuilder cstBuilder(context);
SmallVector<Operation *, 1> generatedConstants;
- // Populate the results with the folded results.
Dialect *dialect = op->getDialect();
- for (auto [foldResult, expectedType] :
- llvm::zip_equal(foldResults, opResults.getTypes())) {
+ for (auto [index, result] : llvm::enumerate(op->getResults())) {
+ OpFoldResult foldResult = foldResults[index];
+ if (!foldResult || (liveOnly && result.use_empty()))
+ continue;
- // Normal values get pushed back directly.
- if (auto value = llvm::dyn_cast_if_present<Value>(foldResult)) {
- results.push_back(value);
+ // Normal values are used directly.
+ if (auto value = llvm::dyn_cast<Value>(foldResult)) {
+ replacements[index] = value;
continue;
}
// Otherwise, try to materialize a constant operation.
if (!dialect)
- return cleanupFailure();
+ return failure();
// Ask the dialect to materialize a constant operation for this value.
Attribute attr = cast<Attribute>(foldResult);
- auto *constOp = dialect->materializeConstant(cstBuilder, attr, expectedType,
- op->getLoc());
+ auto *constOp = dialect->materializeConstant(
+ cstBuilder, attr, result.getType(), op->getLoc());
if (!constOp) {
// Erase any generated constants.
for (Operation *cst : generatedConstants)
cst->erase();
- return cleanupFailure();
+ return failure();
}
assert(matchPattern(constOp, m_Constant()));
+ assert(constOp->getResultTypes()[0] == result.getType() &&
+ "materializeConstant produced incorrect result type");
generatedConstants.push_back(constOp);
- results.push_back(constOp->getResult(0));
+ replacements[index] = constOp->getResult(0);
}
// If we were successful, insert any generated constants.
for (Operation *cst : generatedConstants)
insert(cst);
- // Return materialized constant operations.
- if (materializedConstants)
- *materializedConstants = std::move(generatedConstants);
+ return replacements;
+}
+
+LogicalResult
+OpBuilder::tryFold(Operation *op, SmallVectorImpl<Value> &results,
+ SmallVectorImpl<Operation *> *materializedConstants) {
+ assert(results.empty() && "expected empty results");
+ OpFoldResults foldResults = tryFold(op);
+ if (!foldResults.replacesAll())
+ return success(foldResults.modifiedInPlace());
+
+ FailureOr<SmallVector<Value>> replacements =
+ materializeFoldResults(op, foldResults, /*liveOnly=*/false);
+ if (failed(replacements))
+ return failure();
+ llvm::append_range(results, *replacements);
+
+ // Return materialized constant operations. Each one defines the replacement
+ // of an attribute.
+ if (materializedConstants) {
+ materializedConstants->clear();
+ for (auto [foldResult, replacement] :
+ llvm::zip_equal(foldResults.getReplacements(), *replacements))
+ if (isa<Attribute>(foldResult))
+ materializedConstants->push_back(replacement.getDefiningOp());
+ }
return success();
}
diff --git a/mlir/test/IR/op-builder-fold.mlir b/mlir/test/IR/op-builder-fold.mlir
new file mode 100644
index 0000000000000..a30ffcd78acb0
--- /dev/null
+++ b/mlir/test/IR/op-builder-fold.mlir
@@ -0,0 +1,110 @@
+// RUN: mlir-opt %s --transform-interpreter --verify-diagnostics \
+// RUN: -split-input-file | FileCheck %s
+
+// `transform.test_fold` with `api = "try_fold"` folds each test op with
+// `OpBuilder::tryFold`. With `api = "materialize"`, it also materializes the
+// replacements with `OpBuilder::materializeFoldResults`. The test ops fold as
+// their attributes say; see `getConfiguredFoldResults`.
+
+// CHECK-LABEL: func @try_fold
+func.func @try_fold() {
+ // `tryFold` creates no constant.
+ // CHECK-NEXT: "test.fold_dispatch"() {fold = [unit, 42 : i32]}
+ // expected-remark @below {{fold: [keep, 42 : i32]}}
+ %0:2 = "test.fold_dispatch"() {fold = [unit, 42 : i32]} : () -> (i32, i32)
+ // A fold that changes the op in place repeats, and the result keeps the
+ // in-place bit.
+ // CHECK-NEXT: "test.fold_dispatch"() {fold = [1 : i32, 1 : i32]}
+ // expected-remark @below {{fold: [1 : i32, 1 : i32] in place}}
+ %1:2 = "test.fold_dispatch"() {fold_in_place_count = 1, fold = [1 : i32, 1 : i32]}
+ : () -> (i32, i32)
+ // An in-place fold followed by a failure is an in-place fold.
+ // CHECK-NEXT: "test.fold_dispatch"() : () -> (i32, i32)
+ // expected-remark @below {{fold: in place}}
+ %2:2 = "test.fold_dispatch"() {fold_in_place_count = 1} : () -> (i32, i32)
+ // After 64 in-place folds, `tryFold` stops with a failure.
+ // CHECK-NEXT: "test.fold_dispatch"() {fold_in_place_count = 36 : i64}
+ // expected-remark @below {{fold: failure}}
+ %3:2 = "test.fold_dispatch"() {fold_in_place_count = 100} : () -> (i32, i32)
+ // `tryFold` does not fold a constant.
+ // CHECK-NEXT: "test.constant"
+ // expected-remark @below {{fold: failure}}
+ %4 = "test.constant"() <{value = 7 : i32}> : () -> i32
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(
+ %root: !transform.any_op {transform.readonly}) {
+ %ops = transform.structured.match
+ ops{["test.fold_dispatch", "test.constant"]} in %root
+ : (!transform.any_op) -> !transform.any_op
+ transform.test_fold %ops api = "try_fold" : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-LABEL: func @materialize
+// CHECK-SAME: (%[[ARG0:.*]]: i32)
+func.func @materialize(%arg0: i32) {
+ // An attribute becomes a new constant before the op. A value is used
+ // directly.
+ // CHECK-NEXT: "test.constant"() <{value = 42 : i32}>
+ // CHECK-NEXT: "test.fold_dispatch"(%[[ARG0]])
+ // expected-remark @below {{fold: [42 : i32, operand 0]}}
+ // expected-remark @below {{materialized: [test.constant, operand 0]}}
+ %0:2 = "test.fold_dispatch"(%arg0) {fold = [42 : i32, "operand:0"]}
+ : (i32) -> (i32, i32)
+ // A kept result gets no value.
+ // CHECK-NEXT: "test.constant"() <{value = 43 : i32}>
+ // CHECK-NEXT: "test.fold_dispatch"()
+ // expected-remark @below {{fold: [keep, 43 : i32]}}
+ // expected-remark @below {{materialized: [none, test.constant]}}
+ %1:2 = "test.fold_dispatch"() {fold = [unit, 43 : i32]} : () -> (i32, i32)
+ // If a constant fails to materialize, no constant is inserted.
+ // CHECK-NEXT: "test.fold_dispatch"()
+ // CHECK-NEXT: return
+ // expected-remark @below {{fold: [44 : i32, "unmaterializable"]}}
+ // expected-remark @below {{materialized: failure}}
+ %2:2 = "test.fold_dispatch"() {fold = [44 : i32, "unmaterializable"]}
+ : () -> (i32, i32)
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(
+ %root: !transform.any_op {transform.readonly}) {
+ %ops = transform.structured.match ops{["test.fold_dispatch"]} in %root
+ : (!transform.any_op) -> !transform.any_op
+ transform.test_fold %ops api = "materialize" : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// With `liveOnly`, a replaced result without uses gets no value and no
+// constant.
+// CHECK-LABEL: func @materialize_live_only
+func.func @materialize_live_only() -> i32 {
+ // CHECK-NEXT: %[[C:.*]] = "test.constant"() <{value = 2 : i32}>
+ // CHECK-NEXT: "test.fold_dispatch"()
+ // CHECK-NEXT: return
+ // expected-remark @below {{fold: [1 : i32, 2 : i32]}}
+ // expected-remark @below {{materialized: [none, test.constant]}}
+ %0:2 = "test.fold_dispatch"() {fold = [1 : i32, 2 : i32]} : () -> (i32, i32)
+ return %0#1 : i32
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(
+ %root: !transform.any_op {transform.readonly}) {
+ %ops = transform.structured.match ops{["test.fold_dispatch"]} in %root
+ : (!transform.any_op) -> !transform.any_op
+ transform.test_fold %ops api = "materialize" live_only
+ : !transform.any_op
+ transform.yield
+ }
+}
diff --git a/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp b/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp
index 303da175b2744..6923414fb702d 100644
--- a/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp
+++ b/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp
@@ -995,8 +995,32 @@ DiagnosedSilenceableFailure
mlir::test::TestFoldOp::apply(transform::TransformRewriter &rewriter,
transform::TransformResults &results,
transform::TransformState &state) {
- for (Operation *op : state.getPayloadOps(getTarget()))
- emitFoldRemark(op, op->fold());
+ StringRef api = getApi();
+ if (api != "fold" && api != "try_fold" && api != "materialize")
+ return emitDefiniteFailure() << "unknown api: " << api;
+ for (Operation *op : state.getPayloadOps(getTarget())) {
+ if (api == "fold") {
+ emitFoldRemark(op, op->fold());
+ continue;
+ }
+ OpBuilder builder(op);
+ OpFoldResults foldResults = builder.tryFold(op);
+ emitFoldRemark(op, foldResults);
+ if (api == "try_fold")
+ continue;
+ FailureOr<SmallVector<Value>> values =
+ builder.materializeFoldResults(op, foldResults, getLiveOnly());
+ InFlightDiagnostic remark = op->emitRemark() << "materialized: ";
+ if (failed(values)) {
+ remark << "failure";
+ continue;
+ }
+ remark << "[";
+ llvm::interleaveComma(*values, remark, [&](Value value) {
+ remark << (value ? describeFoldReplacement(op, value) : "none");
+ });
+ remark << "]";
+ }
return DiagnosedSilenceableFailure::success();
}
diff --git a/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td b/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td
index 345f3f36d4951..5c763cc080d86 100644
--- a/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td
+++ b/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td
@@ -589,13 +589,22 @@ def TestFoldOp
[DeclareOpInterfaceMethods<TransformOpInterface>,
DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
let description = [{
- Folds each target op with `Operation::fold` and emits the fold result as a
- remark on the target op. The remark says "failure", "in place", or lists
- one entry per result, followed by "in place" if the fold also changed the
- op in place.
+ Folds each target op and emits the fold result as a remark on the target
+ op. The remark says "failure", "in place", or lists one entry per result,
+ followed by "in place" if the fold also changed the op in place. `api`
+ selects the fold function: "fold" (`Operation::fold`), "try_fold"
+ (`OpBuilder::tryFold`), or "materialize" (`OpBuilder::tryFold`, then
+ `OpBuilder::materializeFoldResults` before the target op, with `live_only`
+ as the `liveOnly` argument). "materialize" emits a second remark with the
+ materialized values.
+ }];
+ let arguments = (ins TransformHandleTypeInterface:$target,
+ DefaultValuedStrAttr<StrAttr, "fold">:$api,
+ UnitAttr:$live_only);
+ let assemblyFormat = [{
+ $target (`api` `=` $api^)? (`live_only` $live_only^)? attr-dict `:`
+ type($target)
}];
- let arguments = (ins TransformHandleTypeInterface:$target);
- let assemblyFormat = "$target attr-dict `:` type($target)";
let cppNamespace = "::mlir::test";
}
diff --git a/mlir/unittests/IR/OpFoldResultsTest.cpp b/mlir/unittests/IR/OpFoldResultsTest.cpp
index bc3bb99851143..69dc92a922cd2 100644
--- a/mlir/unittests/IR/OpFoldResultsTest.cpp
+++ b/mlir/unittests/IR/OpFoldResultsTest.cpp
@@ -53,13 +53,34 @@ struct PartialFoldOp : public Op<PartialFoldOp, OpTrait::NResults<2>::Impl,
}
};
+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 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>();
+ addOperations<PartialFoldOp, ConstantOp>();
+ }
+
+ Operation *materializeConstant(OpBuilder &builder, Attribute value, Type type,
+ Location loc) final {
+ OperationState state(loc, ConstantOp::getOperationName());
+ state.addAttribute("value", value);
+ state.addTypes(type);
+ return builder.create(state);
}
};
} // namespace op_fold_results_test
@@ -696,6 +717,47 @@ TEST_F(OpFoldResultsTest, FreeHelpersMatchMembers) {
EXPECT_TRUE(failed(failedResult));
}
+TEST_F(OpFoldResultsTest, OpBuilderLegacyTryFold) {
+ Block block;
+ OpBuilder b(&context);
+ b.setInsertionPointToEnd(&block);
+ Location loc = b.getUnknownLoc();
+ Value producer =
+ b.create(loc, b.getStringAttr("foo.producer"), {}, {i32})->getResult(0);
+ Operation *op =
+ b.create(loc, b.getStringAttr("fold_test.partial"), {}, {i32, i32});
+ Attribute attr = b.getI32IntegerAttr(1);
+ SmallVector<Value> results;
+ SmallVector<Operation *> constants;
+
+ // A partial fold counts only for its in-place change, and no fold follows
+ // it.
+ unsigned calls = 0;
+ bool inPlace = false;
+ foldState->opFoldFn = [&](Operation *) {
+ ++calls;
+ OpFoldResults result{attr, nullptr};
+ result.setModifiedInPlace(inPlace);
+ return result;
+ };
+ EXPECT_TRUE(failed(b.tryFold(op, results, &constants)));
+ inPlace = true;
+ EXPECT_TRUE(succeeded(b.tryFold(op, results, &constants)));
+ EXPECT_EQ(calls, 2u);
+ EXPECT_TRUE(results.empty());
+ EXPECT_TRUE(constants.empty());
+
+ // A full fold materializes its constants.
+ foldState->opFoldFn = [&](Operation *) -> OpFoldResults {
+ return {attr, producer};
+ };
+ ASSERT_TRUE(succeeded(b.tryFold(op, results, &constants)));
+ ASSERT_EQ(results.size(), 2u);
+ ASSERT_EQ(constants.size(), 1u);
+ EXPECT_EQ(results[0], constants[0]->getResult(0));
+ EXPECT_EQ(results[1], producer);
+}
+
#ifdef GTEST_HAS_DEATH_TEST
#ifndef NDEBUG
namespace {
More information about the llvm-branch-commits
mailing list