[llvm-branch-commits] [mlir] [mlir] Add the `useOpFoldResults` ODS dialect bit (PR #228770)
Victor Perez via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Sun Oct 4 15:04:59 PDT 2026
https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/228770
>From 78ee09e330a65aa2e381f9b65c2c99205f0f4e5e 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 23:48:26 -0700
Subject: [PATCH] [mlir] Add the useOpFoldResults ODS dialect bit
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
The new `OpFoldResults fold(FoldAdaptor)` form needs a matching ODS
declaration. This patch adds the dialect bit `useOpFoldResults`. The bit
is off by default. When a dialect sets it, `genFolderDecls` declares
`OpFoldResults fold(FoldAdaptor adaptor)` for each op that does not have
exactly one fixed result. An op with exactly one fixed result keeps the
`OpFoldResult` form. A dialect without the bit keeps the legacy
declaration.
The `test` dialect sets the bit. Its three folds of ops that do not have
exactly one fixed result now use the new form. The behavior of these
folds does not change, except in two cases. In a graph region or in an
unreachable block, the new form drops a forwarding fold that names a
result that the same fold replaces. And the fold of
`TestOpWithVariadicResultsAndFolder` with no operands now fails. Before,
it reported an in-place change, but it changed nothing. The greedy
driver then put the op on the worklist again and folded it again, so
canonicalize did not stop. A new test in `test-canonicalize.mlir`
covers this case. With the new test and without the fold change,
canonicalize does not stop, and `timeout` ends it:
$ timeout 120 mlir-opt mlir/test/Transforms/test-canonicalize.mlir \
-pass-pipeline='builtin.module(func.func(canonicalize))'
$ echo $?
124
The patch also adds test ops for the driver tests of later patches: a
partial fold, a partial fold with an in-place change, a fold that keeps
all results, own-result forwarding in a graph region, a partial fold
with a replacement that does not materialize, an in-place change without
the in-place mark, and an in-place fold of an op without results. No
driver uses the partial result yet.
The new ODS test `op-fold-results.td` checks the three declarations.
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
mlir/include/mlir/IR/DialectBase.td | 5 +
mlir/include/mlir/TableGen/Dialect.h | 4 +
mlir/lib/TableGen/Dialect.cpp | 4 +
mlir/test/Transforms/test-canonicalize.mlir | 9 ++
mlir/test/lib/Dialect/Test/TestDialect.td | 1 +
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 87 +++++++++++++---
mlir/test/lib/Dialect/Test/TestOps.td | 62 +++++++++++-
mlir/test/mlir-tblgen/op-fold-results.td | 107 ++++++++++++++++++++
mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp | 2 +
9 files changed, 267 insertions(+), 14 deletions(-)
create mode 100644 mlir/test/mlir-tblgen/op-fold-results.td
diff --git a/mlir/include/mlir/IR/DialectBase.td b/mlir/include/mlir/IR/DialectBase.td
index da669c6ada639fc..c3a591a173e8f29 100644
--- a/mlir/include/mlir/IR/DialectBase.td
+++ b/mlir/include/mlir/IR/DialectBase.td
@@ -99,6 +99,11 @@ class Dialect {
// If this dialect can be extended at runtime with new operations or types.
bit isExtensible = 0;
+
+ // If set, operations of this dialect that do not have exactly one fixed
+ // result declare `OpFoldResults fold(FoldAdaptor)` instead of the legacy
+ // `LogicalResult fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &)`.
+ bit useOpFoldResults = 0;
}
#endif // DIALECTBASE_TD
diff --git a/mlir/include/mlir/TableGen/Dialect.h b/mlir/include/mlir/TableGen/Dialect.h
index 68230196b17cc0f..3e77e1bfa4b170e 100644
--- a/mlir/include/mlir/TableGen/Dialect.h
+++ b/mlir/include/mlir/TableGen/Dialect.h
@@ -92,6 +92,10 @@ class Dialect {
/// operations or types.
bool isExtensible() const;
+ /// Returns true if operations of this dialect that do not have exactly one
+ /// fixed result declare `OpFoldResults fold(FoldAdaptor)`.
+ bool useOpFoldResults() const;
+
const llvm::DagInit *getDiscardableAttributes() const;
const llvm::Record *getDef() const { return def; }
diff --git a/mlir/lib/TableGen/Dialect.cpp b/mlir/lib/TableGen/Dialect.cpp
index 84530bfa4f6fd21..af24dacfb0f58c1 100644
--- a/mlir/lib/TableGen/Dialect.cpp
+++ b/mlir/lib/TableGen/Dialect.cpp
@@ -106,6 +106,10 @@ bool Dialect::isExtensible() const {
return def->getValueAsBit("isExtensible");
}
+bool Dialect::useOpFoldResults() const {
+ return def->getValueAsBit("useOpFoldResults");
+}
+
const llvm::DagInit *Dialect::getDiscardableAttributes() const {
return def->getValueAsDag("discardableAttrs");
}
diff --git a/mlir/test/Transforms/test-canonicalize.mlir b/mlir/test/Transforms/test-canonicalize.mlir
index 81538e88beae434..e039fa1917b738a 100644
--- a/mlir/test/Transforms/test-canonicalize.mlir
+++ b/mlir/test/Transforms/test-canonicalize.mlir
@@ -37,6 +37,15 @@ func.func @remove_op_with_variadic_results_and_folder(%arg0 : i32, %arg1 : i32)
return %0, %1 : i32, i32
}
+// Without operands, the fold replaces no result, so it fails.
+// CHECK-LABEL: func @keep_op_with_variadic_results_and_folder_no_operands
+func.func @keep_op_with_variadic_results_and_folder_no_operands() {
+ // CHECK-NEXT: "test.op_with_variadic_results_and_folder"() : () -> ()
+ // CHECK-NEXT: return
+ "test.op_with_variadic_results_and_folder"() : () -> ()
+ return
+}
+
// CHECK-LABEL: func @test_commutative_multi
// CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: i32, %[[ARG_1:[a-z0-9]*]]: i32)
func.func @test_commutative_multi(%arg0: i32, %arg1: i32) -> (i32, i32) {
diff --git a/mlir/test/lib/Dialect/Test/TestDialect.td b/mlir/test/lib/Dialect/Test/TestDialect.td
index 37a263f1d10b8b1..04e6edf945dd9bb 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.td
+++ b/mlir/test/lib/Dialect/Test/TestDialect.td
@@ -24,6 +24,7 @@ def Test_Dialect : Dialect {
let useDefaultTypePrinterParser = 0;
let useDefaultAttributePrinterParser = 1;
let isExtensible = 1;
+ let useOpFoldResults = 1;
let dependentDialects = [
"::mlir::DLTIDialect",
"::mlir::bufferization::BufferizationDialect"
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 855b5b1b7798fb2..46df5e3268f56d9 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -228,27 +228,90 @@ OpFoldResult TestOpConstant::fold(FoldAdaptor adaptor) { return getValue(); }
// TestOpWithVariadicResultsAndFolder
//===----------------------------------------------------------------------===//
-LogicalResult TestOpWithVariadicResultsAndFolder::fold(
- FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &results) {
- for (Value input : this->getOperands()) {
- results.push_back(input);
- }
- return success();
+OpFoldResults TestOpWithVariadicResultsAndFolder::fold(FoldAdaptor adaptor) {
+ return getOperands();
}
//===----------------------------------------------------------------------===//
// TestOpFoldUnmaterializable
//===----------------------------------------------------------------------===//
-LogicalResult
-TestOpFoldUnmaterializable::fold(FoldAdaptor adaptor,
- SmallVectorImpl<OpFoldResult> &results) {
+OpFoldResults TestOpFoldUnmaterializable::fold(FoldAdaptor adaptor) {
// 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 {b.getI32IntegerAttr(42), getOperand(),
+ b.getStringAttr("unmaterializable")};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpPartialFold
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpPartialFold::fold(FoldAdaptor adaptor) {
+ Builder b(getContext());
+ return {b.getI32IntegerAttr(42), getOperand(), OpFoldResult()};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpPartialFoldInPlace
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpPartialFoldInPlace::fold(FoldAdaptor adaptor) {
+ if (getFolded())
+ return {};
+ setFolded(true);
+ OpFoldResults results(getOperation());
+ results.replace(getForwarded(), getOperand());
+ results.setModifiedInPlace();
+ return results;
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpFoldKeepAll
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpFoldKeepAll::fold(FoldAdaptor adaptor) {
+ return {getResult(0), getResult(1)};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpFoldForwardOperands
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpFoldForwardOperands::fold(FoldAdaptor adaptor) {
+ return {getFirst(), getSecond()};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpPartialFoldUnmaterializable
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpPartialFoldUnmaterializable::fold(FoldAdaptor adaptor) {
+ Builder b(getContext());
+ return {b.getI32IntegerAttr(42), b.getStringAttr("unmaterializable"),
+ OpFoldResult()};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpPartialFoldUnmarkedInPlace
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpPartialFoldUnmarkedInPlace::fold(FoldAdaptor adaptor) {
+ // The missing setModifiedInPlace() is the defect under test.
+ if (!getFolded())
+ setFolded(true);
+ return {getOperand(), OpFoldResult()};
+}
+
+//===----------------------------------------------------------------------===//
+// TestOpZeroResultsFoldInPlace
+//===----------------------------------------------------------------------===//
+
+OpFoldResults TestOpZeroResultsFoldInPlace::fold(FoldAdaptor adaptor) {
+ if (getFolded())
+ return failure();
+ setFolded(true);
return success();
}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 5e48a5b60185c53..a592257ca227b6d 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1685,6 +1685,65 @@ def TestOpFoldUnmaterializable : TEST_Op<"op_fold_unmaterializable"> {
let hasFolder = 1;
}
+// Op for testing a partial fold. The folder replaces result 0 with a constant
+// and result 1 with the operand, and keeps result 2.
+def TestOpPartialFold : TEST_Op<"op_partial_fold", [NoMemoryEffect]> {
+ let arguments = (ins I32:$operand);
+ let results = (outs I32, I32, I32);
+ let hasFolder = 1;
+}
+
+// Op for testing a partial fold together with an in-place change. The first
+// fold sets `folded`, marks the in-place change, replaces `forwarded` with the
+// operand, and keeps `kept`. Later folds fail.
+def TestOpPartialFoldInPlace : TEST_Op<"op_partial_fold_in_place"> {
+ let arguments = (ins I32:$operand, UnitAttr:$folded);
+ let results = (outs I32:$forwarded, I32:$kept);
+ let hasFolder = 1;
+}
+
+// Op for testing that a fold fails when it keeps every result and does not
+// change the op in place. The folder returns the results of the op.
+def TestOpFoldKeepAll : TEST_Op<"op_fold_keep_all"> {
+ let results = (outs I32, I32);
+ let hasFolder = 1;
+}
+
+// Op for testing a fold that forwards its operands. In a graph region where
+// operand i is result i of the same op, the folder returns only the results of
+// the op, and the fold must fail.
+def TestOpFoldForwardOperands : TEST_Op<"op_fold_forward_operands"> {
+ let arguments = (ins I32:$first, I32:$second);
+ let results = (outs I32, I32);
+ let hasFolder = 1;
+}
+
+// Op for testing a partial fold that fails constant materialization. The
+// folder replaces result 0 with a constant, replaces result 1 with an attribute
+// that the test dialect cannot materialize, and keeps result 2.
+def TestOpPartialFoldUnmaterializable
+ : TEST_Op<"op_partial_fold_unmaterializable"> {
+ let results = (outs I32, I32, I32);
+ let hasFolder = 1;
+}
+
+// Op for testing the check for an in-place change without an in-place mark.
+// The first fold sets `folded` but does not mark the in-place change. Every
+// fold replaces `forwarded` with the operand and keeps `kept`.
+def TestOpPartialFoldUnmarkedInPlace
+ : TEST_Op<"op_partial_fold_unmarked_in_place"> {
+ let arguments = (ins I32:$operand, UnitAttr:$folded);
+ let results = (outs I32:$forwarded, I32:$kept);
+ let hasFolder = 1;
+}
+
+// Op for testing an in-place fold of an op without results. The first fold
+// sets `folded`. Later folds fail.
+def TestOpZeroResultsFoldInPlace : TEST_Op<"op_zero_results_fold_in_place"> {
+ let arguments = (ins UnitAttr:$folded);
+ let hasFolder = 1;
+}
+
def TestAddIOp : TEST_Op<"addi"> {
let arguments = (ins AnyTypeOf<[I32, TestI32]>:$op1,
AnyTypeOf<[I32, TestI32]>:$op2);
@@ -1767,8 +1826,7 @@ def TestOpInPlaceFoldSuccess : TEST_Op<"op_in_place_fold_success"> {
let results = (outs Variadic<I1>);
let hasFolder = 1;
let extraClassDefinition = [{
- ::llvm::LogicalResult $cppClass::fold(FoldAdaptor adaptor,
- SmallVectorImpl<OpFoldResult> &results) {
+ ::mlir::OpFoldResults $cppClass::fold(FoldAdaptor adaptor) {
return success();
}
}];
diff --git a/mlir/test/mlir-tblgen/op-fold-results.td b/mlir/test/mlir-tblgen/op-fold-results.td
new file mode 100644
index 000000000000000..3b5d06c4e6012b8
--- /dev/null
+++ b/mlir/test/mlir-tblgen/op-fold-results.td
@@ -0,0 +1,107 @@
+// RUN: mlir-tblgen -gen-op-decls -I %S/../../include %s | FileCheck %s
+
+// Check the `fold` declaration that ODS emits with and without the dialect bit
+// `useOpFoldResults`.
+
+include "mlir/IR/OpBase.td"
+
+def Legacy_Dialect : Dialect {
+ let name = "legacy";
+ let cppNamespace = "legacy";
+}
+class Legacy_Op<string mnemonic, list<Trait> traits = []> :
+ Op<Legacy_Dialect, mnemonic, traits>;
+
+def Results_Dialect : Dialect {
+ let name = "results";
+ let cppNamespace = "results";
+ let useOpFoldResults = 1;
+}
+class Results_Op<string mnemonic, list<Trait> traits = []> :
+ Op<Results_Dialect, mnemonic, traits>;
+
+//===----------------------------------------------------------------------===//
+// Dialect without the bit
+//===----------------------------------------------------------------------===//
+
+def Legacy_SingleResultOp : Legacy_Op<"single_result"> {
+ let results = (outs AnyType:$res);
+ let hasFolder = 1;
+}
+
+// CHECK-LABEL: legacy::SingleResultOp declarations
+// CHECK: ::mlir::OpFoldResult fold(FoldAdaptor adaptor);
+
+def Legacy_TwoResultsOp : Legacy_Op<"two_results"> {
+ let results = (outs AnyType:$lhs, AnyType:$rhs);
+ let hasFolder = 1;
+}
+
+// CHECK-LABEL: legacy::TwoResultsOp declarations
+// CHECK: ::llvm::LogicalResult fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult> &results);
+
+def Legacy_VariadicResultOp : Legacy_Op<"variadic_result"> {
+ let results = (outs Variadic<AnyType>:$res);
+ let hasFolder = 1;
+}
+
+// CHECK-LABEL: legacy::VariadicResultOp declarations
+// CHECK: ::llvm::LogicalResult fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult> &results);
+
+//===----------------------------------------------------------------------===//
+// Dialect with the bit
+//===----------------------------------------------------------------------===//
+
+def Results_MixedResultsOp : Results_Op<"mixed_results"> {
+ let results = (outs AnyType:$fixed, Variadic<AnyType>:$variadic);
+ let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::MixedResultsOp declarations
+// CHECK-NOT: SmallVectorImpl<::mlir::OpFoldResult>
+// CHECK: ::mlir::OpFoldResults fold(FoldAdaptor adaptor);
+
+def Results_OptionalResultOp : Results_Op<"optional_result"> {
+ let results = (outs Optional<AnyType>:$res);
+ let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::OptionalResultOp declarations
+// CHECK-NOT: SmallVectorImpl<::mlir::OpFoldResult>
+// CHECK: ::mlir::OpFoldResults fold(FoldAdaptor adaptor);
+
+// The bit does not change the declaration of an op with exactly one fixed
+// result.
+def Results_SingleResultOp : Results_Op<"single_result"> {
+ let results = (outs AnyType:$res);
+ let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::SingleResultOp declarations
+// CHECK: ::mlir::OpFoldResult fold(FoldAdaptor adaptor);
+
+def Results_TwoResultsOp : Results_Op<"two_results"> {
+ let results = (outs AnyType:$lhs, AnyType:$rhs);
+ let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::TwoResultsOp declarations
+// CHECK-NOT: SmallVectorImpl<::mlir::OpFoldResult>
+// CHECK: ::mlir::OpFoldResults fold(FoldAdaptor adaptor);
+
+def Results_VariadicResultOp : Results_Op<"variadic_result"> {
+ let results = (outs Variadic<AnyType>:$res);
+ let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::VariadicResultOp declarations
+// CHECK-NOT: SmallVectorImpl<::mlir::OpFoldResult>
+// CHECK: ::mlir::OpFoldResults fold(FoldAdaptor adaptor);
+
+def Results_ZeroResultsOp : Results_Op<"zero_results"> {
+ let hasFolder = 1;
+}
+
+// CHECK-LABEL: results::ZeroResultsOp declarations
+// CHECK-NOT: SmallVectorImpl<::mlir::OpFoldResult>
+// CHECK: ::mlir::OpFoldResults fold(FoldAdaptor adaptor);
diff --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
index 48158b6b148aab8..0f7665df45ecf1b 100644
--- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
@@ -3430,6 +3430,8 @@ void OpEmitter::genFolderDecls() {
op.getNumResults() == 1 && op.getNumVariableLengthResults() == 0;
if (hasSingleResult) {
retType = "::mlir::OpFoldResult";
+ } else if (op.getDialect().useOpFoldResults()) {
+ retType = "::mlir::OpFoldResults";
} else {
paramList.emplace_back("::llvm::SmallVectorImpl<::mlir::OpFoldResult> &",
"results");
More information about the llvm-branch-commits
mailing list