[Mlir-commits] [mlir] [mlir] Keep existing values when fold materialization fails (PR #228739)
Victor Perez
llvmlistbot at llvm.org
Sat Oct 3 09:42:05 PDT 2026
https://github.com/victor-eds created https://github.com/llvm/llvm-project/pull/228739
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.
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.
Running the new tests without the fix:
```console
$ mlir-opt block_arg.mlir -canonicalize
Stack dump:
...
mlir::RewriterBase::eraseOp(mlir::Operation*) + 152
...
mlir::applyPatternsGreedily(mlir::Region&, ...) + 1733
Segmentation fault (core dumped)
```
```console
$ mlir-opt existing_op.mlir -canonicalize
existing_op.mlir:5:10: error: null operand found
%1:3 = "test.op_fold_unmaterializable"(%0) : (i32) -> (i32, i32, i32)
^
note: see current operation: %0:3 = "test.op_fold_unmaterializable"(<<NULL VALUE>>) : (<<NULL TYPE>>) -> (i32, i32, i32)
```
In the first case, the forwarded value is a block argument, so the rollback erases a null op. In the second one, the rollback erases `test.addi`, which the fold did not create and which is still in use.
Assisted by: Claude Code.
>From 2f29ee679d41883d08559d08b292abbd089b9f57 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.
Without the fix, 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, i32, i32)
^
note: see current operation: %0:3 = "test.op_fold_unmaterializable"(<<NULL VALUE>>) : (<<NULL TYPE>>) -> (i32, i32, i32)
The inputs are the two new test functions. In the first, the forwarded
value is a block argument, so the rollback erases a null op. In the
second, 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 b71194e783ba06..dc6484b05ee820 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 8cad6b98441df1..81538e88beae43 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 57bba2e5749a86..cfc45509f2fc89 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 9fe2168e117374..d20761468463df 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 slot 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 a6b05b55804b7e..5617617cf0faa1 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