[Mlir-commits] [mlir] [mlir] Keep existing values when fold materialization fails (PR #228739)
Victor Perez
llvmlistbot at llvm.org
Sun Oct 4 13:28:04 PDT 2026
https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/228739
>From 1fec8813ad6e8ff90d316f7bcf96c3582b64640d 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] [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.
With only the change in `GreedyPatternRewriteDriver.cpp` reverted, a
build without assertions fails both new tests:
$ mlir-opt block_arg.mlir -canonicalize
Stack dump:
...
mlir::RewriterBase::eraseOp(mlir::Operation*) + 152
...
mlir::applyPatternsGreedily(mlir::Region&, ...) + 1733
Segmentation fault (core dumped)
$ mlir-opt existing_op.mlir -canonicalize
existing_op.mlir:5:10: error: null operand found
%1:3 = "test.op_fold_unmaterializable"(%0) : (i32) -> (i32, ...
^
note: see current operation: %0:3 =
"test.op_fold_unmaterializable"(<<NULL VALUE>>) : (<<NULL TYPE>>)
-> (i32, i32, i32)
Each input holds one of the two new test functions. In block_arg.mlir,
the existing value is a block argument, so the rollback erases a null
op. In existing_op.mlir, the rollback erases `test.addi`, which the
fold did not create and which is still in use.
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 b71194e783ba0..dc6484b05ee82 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 8cad6b98441df..81538e88beae4 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 57bba2e5749a8..cfc45509f2fc8 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 9fe2168e11737..18611ab59d1be 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 replacement 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 a6b05b55804b7..5617617cf0faa 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);
More information about the Mlir-commits
mailing list