[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
Sat Oct 3 18:26:07 PDT 2026
https://github.com/victor-eds created https://github.com/llvm/llvm-project/pull/228801
The current `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 is now implemented using the 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 such a fold. Now the old overload stops after it.
New unit tests in `OpFoldResultsTest.cpp` cover the new functions and the old overload.
>From 7b6250e520a4ee1294ac843b36c19739d3533894 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 such a fold. Now the old overload stops after it.
New unit tests in `OpFoldResultsTest.cpp` cover the new functions and
the old overload.
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
mlir/include/mlir/IR/Builders.h | 26 +++-
mlir/lib/IR/Builders.cpp | 90 +++++++++-----
mlir/unittests/IR/OpFoldResultsTest.cpp | 155 ++++++++++++++++++++++++
3 files changed, 237 insertions(+), 34 deletions(-)
diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index 0bf41530b14b0..7c26958801283 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -572,10 +572,34 @@ 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;
+ /// - 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..aef89f43f1a10 100644
--- a/mlir/lib/IR/Builders.cpp
+++ b/mlir/lib/IR/Builders.cpp
@@ -478,92 +478,116 @@ 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()));
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/unittests/IR/OpFoldResultsTest.cpp b/mlir/unittests/IR/OpFoldResultsTest.cpp
index aad6a6f90095b..e143520c23056 100644
--- a/mlir/unittests/IR/OpFoldResultsTest.cpp
+++ b/mlir/unittests/IR/OpFoldResultsTest.cpp
@@ -138,6 +138,17 @@ struct FoldTestDialect : public Dialect {
addOperations<PartialFoldOp, TraitFoldOp, ConstantOp>();
addInterfaces<TestFoldInterface>();
}
+
+ /// A unit attribute does not materialize.
+ Operation *materializeConstant(OpBuilder &builder, Attribute value, Type type,
+ Location loc) final {
+ if (isa<UnitAttr>(value))
+ return nullptr;
+ OperationState state(loc, ConstantOp::getOperationName());
+ state.addAttribute("value", value);
+ state.addTypes(type);
+ return builder.create(state);
+ }
};
/// An op with two results, no fold, and no fold traits.
@@ -1264,6 +1275,150 @@ TEST_F(OpFoldResultsTest, FoldComputesConstantOperands) {
EXPECT_FALSE(foldState->seenOperands[1]);
}
+TEST_F(OpFoldResultsTest, OpBuilderTryFoldCreatesNoConstant) {
+ Block block;
+ OpBuilder b(&context);
+ b.setInsertionPointToEnd(&block);
+ Operation *op = b.create(
+ b.getUnknownLoc(), b.getStringAttr("fold_test.partial"), {}, {i32, i32});
+ Attribute attr = b.getI32IntegerAttr(1);
+ foldState->opFoldFn = [&](Operation *) -> OpFoldResults {
+ return {attr, nullptr};
+ };
+ OpFoldResults result = b.tryFold(op);
+ EXPECT_TRUE(result.replacesAny());
+ EXPECT_FALSE(result.replacesAll());
+ EXPECT_FALSE(result.modifiedInPlace());
+ EXPECT_EQ(result[0], OpFoldResult(attr));
+ EXPECT_EQ(block.getOperations().size(), 1u);
+}
+
+TEST_F(OpFoldResultsTest, OpBuilderTryFoldRepeatsInPlaceFolds) {
+ Block block;
+ OpBuilder b(&context);
+ b.setInsertionPointToEnd(&block);
+ Operation *op = b.create(
+ b.getUnknownLoc(), b.getStringAttr("fold_test.partial"), {}, {i32, i32});
+ Attribute attr = b.getI32IntegerAttr(1);
+ unsigned calls = 0;
+
+ // The result of the second fold keeps the in-place bit of the first fold.
+ foldState->opFoldFn = [&](Operation *) -> OpFoldResults {
+ if (calls++ == 0)
+ return success();
+ return {attr, attr};
+ };
+ OpFoldResults result = b.tryFold(op);
+ EXPECT_EQ(calls, 2u);
+ EXPECT_TRUE(result.replacesAll());
+ EXPECT_TRUE(result.modifiedInPlace());
+
+ // A failure after an in-place fold gives an in-place fold.
+ calls = 0;
+ foldState->opFoldFn = [&](Operation *) -> OpFoldResults {
+ return success(calls++ == 0);
+ };
+ result = b.tryFold(op);
+ EXPECT_EQ(calls, 2u);
+ EXPECT_TRUE(result.modifiedInPlace());
+ EXPECT_FALSE(result.replacesAny());
+
+ // A fold that always changes the op in place stops with a failure.
+ foldState->opFoldFn = [](Operation *) -> OpFoldResults { return success(); };
+ EXPECT_TRUE(b.tryFold(op).failed());
+}
+
+TEST_F(OpFoldResultsTest, OpBuilderTryFoldSkipsConstants) {
+ Operation *constant = createOp({i32}, "fold_test.constant");
+ constant->setAttr("value", builder.getI32IntegerAttr(7));
+ OpBuilder b(&context);
+ EXPECT_TRUE(b.tryFold(constant).failed());
+}
+
+TEST_F(OpFoldResultsTest, OpBuilderMaterializeFoldResults) {
+ 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);
+
+ // An attribute becomes a new constant. A value is used directly.
+ FailureOr<SmallVector<Value>> replacements =
+ b.materializeFoldResults(op, {attr, producer}, /*liveOnly=*/false);
+ ASSERT_TRUE(succeeded(replacements));
+ auto constant = (*replacements)[0].getDefiningOp<ConstantOp>();
+ ASSERT_TRUE(constant);
+ EXPECT_EQ(constant->getAttr("value"), attr);
+ EXPECT_EQ(constant->getBlock(), &block);
+ EXPECT_EQ((*replacements)[1], producer);
+
+ // A kept result gets null.
+ replacements =
+ b.materializeFoldResults(op, {nullptr, producer}, /*liveOnly=*/false);
+ ASSERT_TRUE(succeeded(replacements));
+ EXPECT_FALSE((*replacements)[0]);
+ EXPECT_EQ((*replacements)[1], producer);
+
+ // With `liveOnly`, a replaced result without uses gets null and no constant.
+ b.create(loc, b.getStringAttr("foo.user"), op->getResult(1));
+ size_t numOps = block.getOperations().size();
+ replacements = b.materializeFoldResults(op, {attr, attr}, /*liveOnly=*/true);
+ ASSERT_TRUE(succeeded(replacements));
+ EXPECT_FALSE((*replacements)[0]);
+ EXPECT_TRUE((*replacements)[1].getDefiningOp<ConstantOp>());
+ EXPECT_EQ(block.getOperations().size(), numOps + 1);
+
+ // If a constant fails to materialize, no constant is inserted.
+ EXPECT_TRUE(failed(b.materializeFoldResults(op, {attr, b.getUnitAttr()},
+ /*liveOnly=*/false)));
+ EXPECT_EQ(block.getOperations().size(), numOps + 1);
+}
+
+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