[llvm-branch-commits] [mlir] [mlir] Apply partial folds in DialectConversion (PR #229171)
Victor Perez via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Mon Oct 5 10:58:16 PDT 2026
https://github.com/victor-eds created https://github.com/llvm/llvm-project/pull/229171
`OperationLegalizer::legalizeWithFold` now uses the new `OpBuilder::tryFold` overload and `materializeFoldResults` enabling partial fold support.
This enables us to apply partial folds when pattern rollback is disabled. We cannot do it when pattern rollback is enabled as this may lead to cycles in the value mapping. Note we can always revisit if we think it's worth it. A partial fold only succeeds if in-place, in which case only the in-place changes would be applied (backwards compatibility).
No rollback mode also enables more powerful modes as skipping constants with no uses, so we don't materialized unused operations we'd need to legalize right away, or that may even fail legalization. In this case, a fold that gives no replacement and no in-place change fails.
>From 3401a303b1627be2df020072e701bf772c86b33b 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:13:38 -0700
Subject: [PATCH] [mlir] Apply partial folds in DialectConversion
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
`OperationLegalizer::legalizeWithFold` now uses the new
`OpBuilder::tryFold` overload and `materializeFoldResults`, and it
applies partial folds without pattern rollback. The legalizer replaces
the uses of each replaced result that has uses, legalizes the new
constants, and tries to legalize the op again.
The legalizer materializes only the replaced results that have uses,
also for a full fold, because it must legalize each new op. So it
creates no constant that nothing uses, and a result without uses can
hold an attribute that does not materialize. In a graph region, the
replacement skips the uses that come before the replacement value, so
the legalizer counts progress by the use counts of the replaced
results. A fold that replaces all results still erases the op, also
when all of its results have no uses.
With pattern rollback, `replaceAllUsesWith` only records a value
mapping, and a partial fold can map a result twice or in a cycle. So
in this mode, a partial fold applies only its in-place change, and
fails without one. A full fold with rollback still materializes every
result.
A fold that gives no replacement and no in-place change fails.
The new tests `test-legalizer-partial-fold.mlir` and
`test-legalizer-fold-dead-result.mlir` run both modes.
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
.../Transforms/Utils/DialectConversion.cpp | 156 +++++++++++++++---
.../test-legalizer-fold-dead-result.mlir | 36 ++++
.../test-legalizer-partial-fold.mlir | 117 +++++++++++++
3 files changed, 283 insertions(+), 26 deletions(-)
create mode 100644 mlir/test/Transforms/test-legalizer-fold-dead-result.mlir
create mode 100644 mlir/test/Transforms/test-legalizer-partial-fold.mlir
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 507f0650169ca..ccd0bdbad990f 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -2447,6 +2447,24 @@ class OperationLegalizer {
/// Attempt to legalize the given operation by folding it.
LogicalResult legalizeWithFold(Operation *op);
+ /// Apply only the in-place change of a fold of the given operation, and
+ /// legalize the changed operation. Fails if the fold did not change the
+ /// operation in place.
+ LogicalResult legalizeInPlaceFold(Operation *op, bool modifiedInPlace);
+
+ /// Replace the uses of the results that a partial fold of the given
+ /// operation replaces, and legalize the changed operation.
+ LogicalResult legalizePartialFold(Operation *op, bool modifiedInPlace,
+ ArrayRef<Value> replacementValues,
+ ArrayRef<Operation *> constants,
+ const RewriterState &curState);
+
+ /// Legalize the constants that a fold of an `opName` operation created. On
+ /// failure, roll back to `curState`.
+ LogicalResult legalizeFoldConstants(StringRef opName,
+ ArrayRef<Operation *> constants,
+ const RewriterState &curState);
+
/// Attempt to legalize the given operation by applying a pattern. Returns
/// success if the operation was legalized, failure otherwise.
LogicalResult legalizeWithPattern(Operation *op);
@@ -2648,6 +2666,18 @@ static T moveAndReset(T &obj) {
return result;
}
+/// Return the constants that define the replacements of attributes.
+static SmallVector<Operation *, 2>
+getFoldConstants(const OpFoldResults &foldResults,
+ ArrayRef<Value> replacementValues) {
+ SmallVector<Operation *, 2> constants;
+ for (auto [foldResult, replacement] :
+ llvm::zip_equal(foldResults.getReplacements(), replacementValues))
+ if (replacement && isa<Attribute>(foldResult))
+ constants.push_back(replacement.getDefiningOp());
+ return constants;
+}
+
LogicalResult OperationLegalizer::legalizeWithFold(Operation *op) {
auto &rewriterImpl = rewriter.getImpl();
LLVM_DEBUG({
@@ -2667,43 +2697,117 @@ LogicalResult OperationLegalizer::legalizeWithFold(Operation *op) {
// Try to fold the operation.
StringRef opName = op->getName().getStringRef();
- SmallVector<Value, 2> replacementValues;
- SmallVector<Operation *, 2> newOps;
+ bool allowRollback = rewriter.getConfig().allowPatternRollback;
rewriter.setInsertionPoint(op);
rewriter.startOpModification(op);
- if (failed(rewriter.tryFold(op, replacementValues, &newOps))) {
+ OpFoldResults foldResults = rewriter.tryFold(op);
+ // With rollback, a replacement only updates a value mapping, and a partial
+ // fold can map a result twice or in a cycle. So a partial fold applies only
+ // its in-place change.
+ if (!foldResults.replacesAny() ||
+ (allowRollback && !foldResults.replacesAll()))
+ return legalizeInPlaceFold(op, foldResults.modifiedInPlace());
+
+ // A full fold with rollback materializes every result. Otherwise, a replaced
+ // result without uses gets no constant.
+ FailureOr<SmallVector<Value>> replacementValues =
+ rewriter.materializeFoldResults(op, foldResults,
+ /*liveOnly=*/!allowRollback ||
+ !foldResults.replacesAll());
+ // With rollback, a constant that fails to materialize fails the fold, and
+ // the rollback also undoes the in-place change.
+ if (failed(replacementValues) && allowRollback) {
LLVM_DEBUG(logFailure(rewriterImpl.logger, "unable to fold"));
rewriter.cancelOpModification(op);
return failure();
}
+ // If no result with uses gets a new value, only the in-place change applies.
+ if (failed(replacementValues) ||
+ (!foldResults.replacesAll() &&
+ llvm::all_of(*replacementValues, [](Value value) { return !value; })))
+ return legalizeInPlaceFold(op, foldResults.modifiedInPlace());
+
+ SmallVector<Operation *, 2> constants =
+ getFoldConstants(foldResults, *replacementValues);
+ if (!foldResults.replacesAll())
+ return legalizePartialFold(op, foldResults.modifiedInPlace(),
+ *replacementValues, constants, curState);
+
+ // Insert a replacement for 'op' with the folded replacement values.
rewriter.finalizeOpModification(op);
+ rewriter.replaceOp(op, *replacementValues);
+ if (failed(legalizeFoldConstants(opName, constants, curState)))
+ return failure();
- // An empty list of replacement values indicates that the fold was in-place.
- // As the operation changed, a new legalization needs to be attempted.
- if (replacementValues.empty())
- return legalize(op);
+ LLVM_DEBUG(logSuccess(rewriterImpl.logger, ""));
+ return success();
+}
- // Insert a replacement for 'op' with the folded replacement values.
- rewriter.replaceOp(op, replacementValues);
-
- // Recursively legalize any new constant operations.
- for (Operation *newOp : newOps) {
- if (failed(legalize(newOp))) {
- LLVM_DEBUG(logFailure(rewriterImpl.logger,
- "failed to legalize generated constant '{0}'",
- newOp->getName()));
- if (!rewriter.getConfig().allowPatternRollback) {
- // Rolling back a folder is like rolling back a pattern.
- llvm::reportFatalInternalError(
- "op '" + opName +
- "' folder rollback of IR modifications requested");
- }
- rewriterImpl.resetState(curState, std::string(opName) + " folder");
- return failure();
- }
+LogicalResult OperationLegalizer::legalizeInPlaceFold(Operation *op,
+ bool modifiedInPlace) {
+ if (!modifiedInPlace) {
+ LLVM_DEBUG(logFailure(rewriter.getImpl().logger, "unable to fold"));
+ rewriter.cancelOpModification(op);
+ return failure();
}
+ rewriter.finalizeOpModification(op);
+ return legalize(op);
+}
- LLVM_DEBUG(logSuccess(rewriterImpl.logger, ""));
+LogicalResult OperationLegalizer::legalizePartialFold(
+ Operation *op, bool modifiedInPlace, ArrayRef<Value> replacementValues,
+ ArrayRef<Operation *> constants, const RewriterState &curState) {
+ auto &rewriterImpl = rewriter.getImpl();
+ // The replacement skips the uses that come before the replacement value in a
+ // graph region, so a result counts as progress only if it loses uses.
+ bool changed = modifiedInPlace;
+ for (auto [result, replacement] :
+ llvm::zip_equal(op->getResults(), replacementValues)) {
+ if (!replacement)
+ continue;
+ unsigned numUsesBefore = result.getNumUses();
+ rewriter.replaceAllUsesWith(result, replacement);
+ if (result.getNumUses() < numUsesBefore)
+ changed = true;
+ }
+ if (!changed) {
+ LLVM_DEBUG(logFailure(rewriterImpl.logger, "fold made no progress"));
+ for (Operation *constant : constants)
+ rewriter.eraseOp(constant);
+ rewriter.cancelOpModification(op);
+ return failure();
+ }
+ rewriter.finalizeOpModification(op);
+ if (failed(legalizeFoldConstants(op->getName().getStringRef(), constants,
+ curState)))
+ return failure();
+
+ // The op changed, so a new legalization needs to be attempted. Clear the
+ // pattern state of this fold first.
+ rewriterImpl.patternNewOps.clear();
+ rewriterImpl.patternModifiedOps.clear();
+ return legalize(op);
+}
+
+LogicalResult
+OperationLegalizer::legalizeFoldConstants(StringRef opName,
+ ArrayRef<Operation *> constants,
+ const RewriterState &curState) {
+ auto &rewriterImpl = rewriter.getImpl();
+ for (Operation *constant : constants) {
+ if (succeeded(legalize(constant)))
+ continue;
+ LLVM_DEBUG(logFailure(rewriterImpl.logger,
+ "failed to legalize generated constant '{0}'",
+ constant->getName()));
+ if (!rewriter.getConfig().allowPatternRollback) {
+ // Rolling back a folder is like rolling back a pattern.
+ llvm::reportFatalInternalError(
+ "op '" + opName + "' folder rollback of IR modifications requested");
+ }
+ rewriterImpl.resetState(curState, std::string(opName) + " folder");
+ return failure();
+ }
return success();
}
diff --git a/mlir/test/Transforms/test-legalizer-fold-dead-result.mlir b/mlir/test/Transforms/test-legalizer-fold-dead-result.mlir
new file mode 100644
index 0000000000000..75fe855c307dd
--- /dev/null
+++ b/mlir/test/Transforms/test-legalizer-fold-dead-result.mlir
@@ -0,0 +1,36 @@
+// RUN: mlir-opt -split-input-file -test-legalize-patterns="allow-pattern-rollback=1" %s | FileCheck %s --check-prefix=ROLLBACK
+// RUN: mlir-opt -split-input-file -test-legalize-patterns="allow-pattern-rollback=0" %s | FileCheck %s --check-prefix=NO-ROLLBACK
+
+// Constant materialization fails for result 2. The "no rollback" mode skips
+// result 2, because it has no uses, and erases the op. The "rollback" mode
+// materializes every result, so the fold fails and the op stays. The remarks
+// differ between the modes, so this file does not verify them.
+
+// ROLLBACK-LABEL: func @fold_dead_unmaterializable_result
+// ROLLBACK-SAME: (%[[ARG0:[a-z0-9]+]]: i32)
+// ROLLBACK-NEXT: %[[RES:[a-z0-9]+]]:3 = "test.op_fold_unmaterializable"(%[[ARG0]])
+// ROLLBACK-NEXT: "test.return"(%[[RES]]#0, %[[RES]]#1)
+// NO-ROLLBACK-LABEL: func @fold_dead_unmaterializable_result
+// NO-ROLLBACK-SAME: (%[[ARG0:[a-z0-9]+]]: i32)
+// NO-ROLLBACK-NEXT: %[[C42:[a-z0-9]+]] = "test.constant"() <{value = 42 : i32}> : () -> i32
+// NO-ROLLBACK-NEXT: "test.return"(%[[C42]], %[[ARG0]])
+func.func @fold_dead_unmaterializable_result(%arg0: i32) -> (i32, i32) {
+ %0:3 = "test.op_fold_unmaterializable"(%arg0) : (i32) -> (i32, i32, i32)
+ "test.return"(%0#0, %0#1) : (i32, i32) -> ()
+}
+
+// -----
+
+// No result has uses. The "no rollback" mode materializes nothing and erases
+// the op.
+
+// ROLLBACK-LABEL: func @fold_all_results_dead
+// ROLLBACK-SAME: (%[[ARG0:[a-z0-9]+]]: i32)
+// ROLLBACK-NEXT: "test.op_fold_unmaterializable"(%[[ARG0]])
+// ROLLBACK-NEXT: "test.return"()
+// NO-ROLLBACK-LABEL: func @fold_all_results_dead
+// NO-ROLLBACK-NEXT: "test.return"()
+func.func @fold_all_results_dead(%arg0: i32) {
+ %0:3 = "test.op_fold_unmaterializable"(%arg0) : (i32) -> (i32, i32, i32)
+ "test.return"() : () -> ()
+}
diff --git a/mlir/test/Transforms/test-legalizer-partial-fold.mlir b/mlir/test/Transforms/test-legalizer-partial-fold.mlir
new file mode 100644
index 0000000000000..4c784ef74ebdb
--- /dev/null
+++ b/mlir/test/Transforms/test-legalizer-partial-fold.mlir
@@ -0,0 +1,117 @@
+// RUN: mlir-opt -split-input-file -test-legalize-patterns="allow-pattern-rollback=1" -verify-diagnostics %s | FileCheck %s --check-prefixes=CHECK,ROLLBACK
+// RUN: mlir-opt -split-input-file -test-legalize-patterns="allow-pattern-rollback=0" -verify-diagnostics %s | FileCheck %s --check-prefixes=CHECK,NO-ROLLBACK
+
+// Without rollback, the legalizer applies a partial fold. With rollback, it
+// applies only the in-place change of a partial fold. The test ops are not in
+// the conversion target, so each op stays and gets a remark in both modes.
+
+// CHECK-LABEL: func @partial_fold
+// CHECK-SAME: (%[[ARG0:[a-z0-9]+]]: i32)
+func.func @partial_fold(%arg0: i32) -> (i32, i32, i32) {
+ // NO-ROLLBACK-NEXT: %[[C42:[a-z0-9]+]] = "test.constant"() <{value = 42 : i32}> : () -> i32
+ // NO-ROLLBACK-NEXT: %[[RES:[a-z0-9]+]]:3 = "test.op_partial_fold"(%[[ARG0]])
+ // NO-ROLLBACK-NEXT: "test.return"(%[[C42]], %[[ARG0]], %[[RES]]#2)
+ // ROLLBACK-NEXT: %[[RES:[a-z0-9]+]]:3 = "test.op_partial_fold"(%[[ARG0]])
+ // ROLLBACK-NEXT: "test.return"(%[[RES]]#0, %[[RES]]#1, %[[RES]]#2)
+ // expected-remark at +1 {{op 'test.op_partial_fold' is not legalizable}}
+ %0:3 = "test.op_partial_fold"(%arg0) : (i32) -> (i32, i32, i32)
+ "test.return"(%0#0, %0#1, %0#2) : (i32, i32, i32) -> ()
+}
+
+// -----
+
+// CHECK-LABEL: func @partial_fold_in_place
+// CHECK-SAME: (%[[ARG0:[a-z0-9]+]]: i32)
+func.func @partial_fold_in_place(%arg0: i32) -> (i32, i32) {
+ // CHECK-NEXT: %[[FWD:[a-z0-9_]+]], %[[KEPT:[a-z0-9_]+]] = "test.op_partial_fold_in_place"(%[[ARG0]]) <{folded}>
+ // NO-ROLLBACK-NEXT: "test.return"(%[[ARG0]], %[[KEPT]])
+ // ROLLBACK-NEXT: "test.return"(%[[FWD]], %[[KEPT]])
+ // expected-remark at +1 {{op 'test.op_partial_fold_in_place' is not legalizable}}
+ %0:2 = "test.op_partial_fold_in_place"(%arg0) : (i32) -> (i32, i32)
+ "test.return"(%0#0, %0#1) : (i32, i32) -> ()
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_keep_all
+func.func @fold_keep_all() -> (i32, i32) {
+ // CHECK-NEXT: %[[RES:[a-z0-9]+]]:2 = "test.op_fold_keep_all"()
+ // CHECK-NEXT: "test.return"(%[[RES]]#0, %[[RES]]#1)
+ // expected-remark at +1 {{op 'test.op_fold_keep_all' is not legalizable}}
+ %0:2 = "test.op_fold_keep_all"() : () -> (i32, i32)
+ "test.return"(%0#0, %0#1) : (i32, i32) -> ()
+}
+
+// -----
+
+// The replaced results have no uses, so the fold makes no progress.
+
+// CHECK-LABEL: func @partial_fold_dead_replaced_results
+// CHECK-SAME: (%[[ARG0:[a-z0-9]+]]: i32)
+func.func @partial_fold_dead_replaced_results(%arg0: i32) -> i32 {
+ // CHECK-NEXT: %[[RES:[a-z0-9]+]]:3 = "test.op_partial_fold"(%[[ARG0]])
+ // CHECK-NEXT: "test.return"(%[[RES]]#2)
+ // expected-remark at +1 {{op 'test.op_partial_fold' is not legalizable}}
+ %0:3 = "test.op_partial_fold"(%arg0) : (i32) -> (i32, i32, i32)
+ "test.return"(%0#2) : (i32) -> ()
+}
+
+// -----
+
+// CHECK-LABEL: func @partial_fold_unmaterializable
+func.func @partial_fold_unmaterializable() -> (i32, i32, i32) {
+ // CHECK-NEXT: %[[RES:[a-z0-9]+]]:3 = "test.op_partial_fold_unmaterializable"()
+ // CHECK-NEXT: "test.return"(%[[RES]]#0, %[[RES]]#1, %[[RES]]#2)
+ // expected-remark at +1 {{op 'test.op_partial_fold_unmaterializable' is not legalizable}}
+ %0:3 = "test.op_partial_fold_unmaterializable"() : () -> (i32, i32, i32)
+ "test.return"(%0#0, %0#1, %0#2) : (i32, i32, i32) -> ()
+}
+
+// -----
+
+// In a graph region, an operand of the op can be a result of the same op. A
+// fold that forwards such an operand keeps the result.
+
+// CHECK-LABEL: func @partial_fold_graph_region
+// CHECK-SAME: (%[[ARG0:[a-z0-9]+]]: i32)
+func.func @partial_fold_graph_region(%arg0: i32) {
+ // CHECK-NEXT: test.graph_region {
+ // CHECK-NEXT: %[[RES_0:[a-z0-9]+]]:2 = "test.op_fold_forward_operands"(%[[RES_0]]#0, %[[RES_0]]#1)
+ // CHECK-NEXT: %[[RES_1:[a-z0-9]+]]:2 = "test.op_fold_forward_operands"(%[[ARG0]], %[[RES_1]]#1)
+ // NO-ROLLBACK-NEXT: "test.valid"(%[[RES_0]]#0, %[[RES_0]]#1, %[[ARG0]], %[[RES_1]]#1)
+ // ROLLBACK-NEXT: "test.valid"(%[[RES_0]]#0, %[[RES_0]]#1, %[[RES_1]]#0, %[[RES_1]]#1)
+ // expected-remark at +1 {{op 'test.graph_region' is not legalizable}}
+ test.graph_region {
+ // expected-remark at +1 {{op 'test.op_fold_forward_operands' is not legalizable}}
+ %0:2 = "test.op_fold_forward_operands"(%0#0, %0#1) : (i32, i32) -> (i32, i32)
+ // expected-remark at +1 {{op 'test.op_fold_forward_operands' is not legalizable}}
+ %1:2 = "test.op_fold_forward_operands"(%arg0, %1#1) : (i32, i32) -> (i32, i32)
+ "test.valid"(%0#0, %0#1, %1#0, %1#1) : (i32, i32, i32, i32) -> ()
+ }
+ "test.return"() : () -> ()
+}
+
+// -----
+
+// The replacement does not replace a use in the op that defines the
+// replacement value. Without rollback, the fold of the first op replaces no
+// use, so it makes no progress, and the fold of the second op replaces the use
+// in "test.valid". With rollback, neither partial fold applies.
+
+// CHECK-LABEL: func @partial_fold_graph_region_skipped_use
+func.func @partial_fold_graph_region_skipped_use() {
+ // CHECK-NEXT: test.graph_region {
+ // CHECK-NEXT: %[[RES_0:[a-z0-9]+]]:2 = "test.op_fold_forward_operands"(%[[RES_1:[a-z0-9]+]]#0, %[[RES_0]]#1)
+ // CHECK-NEXT: %[[RES_1]]:2 = "test.op_fold_forward_operands"(%[[RES_0]]#0, %[[RES_1]]#1)
+ // ROLLBACK-NEXT: "test.valid"(%[[RES_1]]#0, %[[RES_1]]#1, %[[RES_0]]#1)
+ // NO-ROLLBACK-NEXT: "test.valid"(%[[RES_0]]#0, %[[RES_1]]#1, %[[RES_0]]#1)
+ // expected-remark at +1 {{op 'test.graph_region' is not legalizable}}
+ test.graph_region {
+ // expected-remark at +1 {{op 'test.op_fold_forward_operands' is not legalizable}}
+ %0:2 = "test.op_fold_forward_operands"(%1#0, %0#1) : (i32, i32) -> (i32, i32)
+ // expected-remark at +1 {{op 'test.op_fold_forward_operands' is not legalizable}}
+ %1:2 = "test.op_fold_forward_operands"(%0#0, %1#1) : (i32, i32) -> (i32, i32)
+ "test.valid"(%1#0, %1#1, %0#1) : (i32, i32, i32) -> ()
+ }
+ "test.return"() : () -> ()
+}
More information about the llvm-branch-commits
mailing list