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?= <[email protected]> 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 <[email protected]> --- .../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@+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@+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@+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@+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@+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@+1 {{op 'test.graph_region' is not legalizable}} + test.graph_region { + // expected-remark@+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@+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@+1 {{op 'test.graph_region' is not legalizable}} + test.graph_region { + // expected-remark@+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@+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"() : () -> () +} _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
