[llvm-branch-commits] [mlir] [mlir] Add the `OpFoldResults` form to fold traits and `DialectFoldInterface` (PR #228769)
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/228769
>From bd833d9022ab34c3c5ee5b92e3be31000aa7a69f 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: Thu, 1 Oct 2026 08:00:02 -0700
Subject: [PATCH] [mlir] Add the OpFoldResults form to foldTrait and
DialectFoldInterface
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
The patch "[mlir] Add the OpFoldResults op fold form" added
`OpFoldResults fold(FoldAdaptor)` for ops. The fold traits and
`DialectFoldInterface` still use only the legacy vector form, which
cannot express a partial fold.
A trait can now define `static OpFoldResults foldTrait(Operation *,
ArrayRef<Attribute>)`. The trait fold dispatch detects this form by its
return type and calls it directly. This form has the same parameters as
the single-result trait form, so the single-result detector excludes it.
The legacy trait form still works through the strict adapter.
`IsCommutative` and `CastOpInterface` now use the new form.
`DialectFoldInterface` gets a new `fold` method that returns
`OpFoldResults`. Its default implementation calls the legacy `fold`
method through the strict adapter, so existing dialects keep their
behavior. `Operation::fold` now calls the new method.
The rule that a replacement does not name a result that the same fold
replaces still applies only to the op's own `OpFoldResults` fold.
`CastOpInterface` forwards its operands, and in a graph region an
operand can be another result of the same op. Upstream folds that case,
so trait and `DialectFoldInterface` results keep that behavior. The
documentation of the new `DialectFoldInterface` method says so.
Downstream impact: `impl::foldCommutative` and
`impl::foldCastInterfaceOp` now return `OpFoldResults` and take no
results vector.
The new test op `test.fold_dispatch_traits` has a fold trait in the new
form and a legacy fold trait, and the fallback fold of the test dialect
gets the new form too. All of them read their result from attributes.
New cases in `IR/fold-dispatch.mlir` cover the trait form, the stop at
the first trait that folds, the new `DialectFoldInterface` method, and
the replacements of a trait fold that names another replaced result.
Static asserts in the test dialect check the detector.
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
mlir/include/mlir/IR/OpDefinition.h | 38 ++++++++---
mlir/include/mlir/Interfaces/CastInterfaces.h | 5 +-
.../include/mlir/Interfaces/CastInterfaces.td | 6 +-
.../mlir/Interfaces/DialectFoldInterface.td | 18 +++++
mlir/include/mlir/Interfaces/FoldInterfaces.h | 12 +++-
mlir/lib/IR/Operation.cpp | 18 +++--
mlir/lib/Interfaces/CastInterfaces.cpp | 11 ++-
mlir/test/IR/fold-dispatch.mlir | 67 ++++++++++++++++++-
.../Dialect/Test/TestDialectInterfaces.cpp | 11 ++-
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 5 ++
mlir/test/lib/Dialect/Test/TestOps.h | 12 ++++
mlir/test/lib/Dialect/Test/TestOps.td | 14 ++++
12 files changed, 185 insertions(+), 32 deletions(-)
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index b09f51e217f39..74c5d1ee85860 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -284,8 +284,7 @@ namespace OpTrait {
// corresponding trait classes. This avoids them being template
// instantiated/duplicated.
namespace impl {
-LogicalResult foldCommutative(Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results);
+OpFoldResults foldCommutative(Operation *op, ArrayRef<Attribute> operands);
OpFoldResult foldIdempotent(Operation *op);
OpFoldResult foldInvolution(Operation *op);
LogicalResult verifyZeroOperands(Operation *op);
@@ -1120,9 +1119,8 @@ class ResultsAreSignlessIntegerLike
template <typename ConcreteType>
class IsCommutative : public TraitBase<ConcreteType, IsCommutative> {
public:
- static LogicalResult foldTrait(Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
- return impl::foldCommutative(op, operands, results);
+ static OpFoldResults foldTrait(Operation *op, ArrayRef<Attribute> operands) {
+ return impl::foldCommutative(op, operands);
}
};
@@ -1534,14 +1532,26 @@ inline bool hasTrait<>(TypeID traitID) {
// Trait Folding
//===----------------------------------------------------------------------===//
+/// Trait to check if T provides a general 'foldTrait' method that returns
+/// OpFoldResults.
+template <typename T, typename... Args>
+using has_fold_results_trait = std::enable_if_t<
+ std::is_same_v<decltype(T::foldTrait(std::declval<Operation *>(),
+ std::declval<ArrayRef<Attribute>>())),
+ OpFoldResults>>;
+template <typename T>
+using detect_has_fold_results_trait =
+ llvm::is_detected<has_fold_results_trait, T>;
/// Trait to check if T provides a 'foldTrait' method for single result
-/// operations.
+/// operations. The OpFoldResults form has the same parameters, so it is
+/// excluded.
template <typename T, typename... Args>
using has_single_result_fold_trait = decltype(T::foldTrait(
std::declval<Operation *>(), std::declval<ArrayRef<Attribute>>()));
template <typename T>
using detect_has_single_result_fold_trait =
- llvm::is_detected<has_single_result_fold_trait, T>;
+ std::conjunction<llvm::is_detected<has_single_result_fold_trait, T>,
+ std::negation<detect_has_fold_results_trait<T>>>;
/// Trait to check if T provides a general 'foldTrait' method.
template <typename T, typename... Args>
using has_fold_trait =
@@ -1553,7 +1563,7 @@ using detect_has_fold_trait = llvm::is_detected<has_fold_trait, T>;
/// Trait to check if T provides any `foldTrait` method.
template <typename T>
using detect_has_any_fold_trait =
- std::disjunction<detect_has_fold_trait<T>,
+ std::disjunction<detect_has_fold_results_trait<T>, detect_has_fold_trait<T>,
detect_has_single_result_fold_trait<T>>;
/// Returns the result of folding a trait that implements a `foldTrait` function
@@ -1573,8 +1583,16 @@ foldTrait(Operation *op, ArrayRef<Attribute> operands) {
return success();
return result;
}
-/// Returns the result of folding a trait that implements the generalized
-/// `foldTrait(op, operands, results)` function, with the strict legacy
+/// Returns the result of folding a trait that implements a generalized
+/// `foldTrait` function that supports any operation type.
+template <typename Trait>
+std::enable_if_t<detect_has_fold_results_trait<Trait>::value, OpFoldResults>
+foldTrait(Operation *op, ArrayRef<Attribute> operands) {
+ return Trait::foldTrait(op, operands);
+}
+/// Returns the result of folding a trait that implements the legacy
+/// generalized `foldTrait` function. This is the folder for traits that do
+/// not implement the `OpFoldResults` form. It keeps the strict legacy
/// contract.
template <typename Trait>
std::enable_if_t<detect_has_fold_trait<Trait>::value, OpFoldResults>
diff --git a/mlir/include/mlir/Interfaces/CastInterfaces.h b/mlir/include/mlir/Interfaces/CastInterfaces.h
index 991b084b8b44d..b163c649e2f33 100644
--- a/mlir/include/mlir/Interfaces/CastInterfaces.h
+++ b/mlir/include/mlir/Interfaces/CastInterfaces.h
@@ -21,9 +21,8 @@ class DialectRegistry;
namespace impl {
/// Attempt to fold the given cast operation.
-LogicalResult foldCastInterfaceOp(Operation *op,
- ArrayRef<Attribute> attrOperands,
- SmallVectorImpl<OpFoldResult> &foldResults);
+OpFoldResults foldCastInterfaceOp(Operation *op,
+ ArrayRef<Attribute> attrOperands);
/// Attempt to verify the given cast operation.
LogicalResult verifyCastInterfaceOp(Operation *op);
diff --git a/mlir/include/mlir/Interfaces/CastInterfaces.td b/mlir/include/mlir/Interfaces/CastInterfaces.td
index 3b81eef8b864b..c32bc2946bf24 100644
--- a/mlir/include/mlir/Interfaces/CastInterfaces.td
+++ b/mlir/include/mlir/Interfaces/CastInterfaces.td
@@ -38,9 +38,9 @@ def CastOpInterface : OpInterface<"CastOpInterface"> {
let extraTraitClassDeclaration = [{
/// Attempt to fold the given cast operation.
- static LogicalResult foldTrait(Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
- return impl::foldCastInterfaceOp(op, operands, results);
+ static OpFoldResults foldTrait(Operation *op,
+ ArrayRef<Attribute> operands) {
+ return impl::foldCastInterfaceOp(op, operands);
}
}];
let verify = [{
diff --git a/mlir/include/mlir/Interfaces/DialectFoldInterface.td b/mlir/include/mlir/Interfaces/DialectFoldInterface.td
index 138be794f605c..4c75ac6e27043 100644
--- a/mlir/include/mlir/Interfaces/DialectFoldInterface.td
+++ b/mlir/include/mlir/Interfaces/DialectFoldInterface.td
@@ -11,6 +11,24 @@ def DialectFoldInterface : DialectInterface<"DialectFoldInterface"> {
let cppNamespace = "::mlir";
let methods = [
+ InterfaceMethod<[{
+ Registered fallback fold for the dialect. Like the fold hook of each
+ operation, it attempts to fold the operation with the specified constant
+ operand values - the elements in "operands" will correspond directly to
+ the operands of the operation, but may be null if non-constant. The
+ result has the same contract as the fold hook of an operation, but the
+ caller does not drop a replacement that names another replaced result.
+ The default implementation calls the `fold` method with a `results`
+ vector, with its strict contract.
+ }],
+ "::mlir::OpFoldResults", "fold",
+ (ins "::mlir::Operation *":$op,
+ "::mlir::ArrayRef<::mlir::Attribute>":$operands),
+ [{
+ return ::mlir::detail::foldWithLegacyDialectFoldInterface(
+ this, op, operands);
+ }]
+ >,
InterfaceMethod<[{
Registered fallback fold for the dialect. Like the fold hook of each
operation, it attempts to fold the operation with the specified constant
diff --git a/mlir/include/mlir/Interfaces/FoldInterfaces.h b/mlir/include/mlir/Interfaces/FoldInterfaces.h
index 5a4247d848d18..52be6d164119f 100644
--- a/mlir/include/mlir/Interfaces/FoldInterfaces.h
+++ b/mlir/include/mlir/Interfaces/FoldInterfaces.h
@@ -9,13 +9,23 @@
#define MLIR_INTERFACES_FOLDINTERFACES_H_
#include "mlir/IR/DialectInterface.h"
+#include "mlir/IR/OpFoldResult.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
namespace mlir {
class Attribute;
-class OpFoldResult;
+class DialectFoldInterface;
+class Operation;
class Region;
+
+namespace detail {
+/// Fold `op` with the legacy vector `fold` method of `iface`, with the strict
+/// legacy contract.
+OpFoldResults
+foldWithLegacyDialectFoldInterface(const DialectFoldInterface *iface,
+ Operation *op, ArrayRef<Attribute> operands);
+} // namespace detail
} // namespace mlir
#include "mlir/Interfaces/DialectFoldInterface.h.inc"
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index e771d3f194bcc..dd216e1932d04 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -600,6 +600,15 @@ void Operation::setSuccessor(Block *block, unsigned index) {
getBlockOperands()[index].set(block);
}
+OpFoldResults
+detail::foldWithLegacyDialectFoldInterface(const DialectFoldInterface *iface,
+ Operation *op,
+ ArrayRef<Attribute> operands) {
+ SmallVector<OpFoldResult> results;
+ LogicalResult status = iface->fold(op, operands, results);
+ return convertLegacyFoldResults(status, results);
+}
+
/// Attempt to fold this operation using the Op's registered foldHook.
OpFoldResults Operation::fold(ArrayRef<Attribute> operands) {
// If we have a registered operation definition matching this one, use it to
@@ -618,9 +627,7 @@ OpFoldResults Operation::fold(ArrayRef<Attribute> operands) {
if (!interface)
return results;
- SmallVector<OpFoldResult> legacyResults;
- LogicalResult status = interface->fold(this, operands, legacyResults);
- results = detail::convertLegacyFoldResults(status, legacyResults);
+ results = interface->fold(this, operands);
results.normalize(this);
return results;
}
@@ -849,9 +856,8 @@ InFlightDiagnostic OpState::emitRemark(const Twine &message) {
// Op Trait implementations
//===----------------------------------------------------------------------===//
-LogicalResult
-OpTrait::impl::foldCommutative(Operation *op, ArrayRef<Attribute> operands,
- SmallVectorImpl<OpFoldResult> &results) {
+OpFoldResults OpTrait::impl::foldCommutative(Operation *op,
+ ArrayRef<Attribute> operands) {
// Nothing to fold if there are not at least 2 operands.
if (op->getNumOperands() < 2)
return failure();
diff --git a/mlir/lib/Interfaces/CastInterfaces.cpp b/mlir/lib/Interfaces/CastInterfaces.cpp
index 05c872daf5dab..7e32cebdf739b 100644
--- a/mlir/lib/Interfaces/CastInterfaces.cpp
+++ b/mlir/lib/Interfaces/CastInterfaces.cpp
@@ -18,19 +18,16 @@ using namespace mlir;
//===----------------------------------------------------------------------===//
/// Attempt to fold the given cast operation.
-LogicalResult
-impl::foldCastInterfaceOp(Operation *op, ArrayRef<Attribute> attrOperands,
- SmallVectorImpl<OpFoldResult> &foldResults) {
+OpFoldResults impl::foldCastInterfaceOp(Operation *op,
+ ArrayRef<Attribute> attrOperands) {
OperandRange operands = op->getOperands();
if (operands.empty())
return failure();
ResultRange results = op->getResults();
// Check for the case where the input and output types match 1-1.
- if (operands.getTypes() == results.getTypes()) {
- foldResults.append(operands.begin(), operands.end());
- return success();
- }
+ if (operands.getTypes() == results.getTypes())
+ return operands;
return failure();
}
diff --git a/mlir/test/IR/fold-dispatch.mlir b/mlir/test/IR/fold-dispatch.mlir
index 4141882bda878..b29962b0eee22 100644
--- a/mlir/test/IR/fold-dispatch.mlir
+++ b/mlir/test/IR/fold-dispatch.mlir
@@ -102,11 +102,76 @@ func.func @fold_naming_replaced_result_unreachable() {
return
}
+// `test.fold_dispatch_traits` has no fold of its own and two fold traits: one
+// in the `OpFoldResults` form, then a legacy one. Folding stops at the first
+// trait that does not fail.
+func.func @trait_forms() {
+ // A partial fold of the first trait skips the legacy trait and the dialect.
+ // expected-remark @below {{fold: [keep, 1 : i32]}}
+ %0:2 = "test.fold_dispatch_traits"()
+ {trait_fold = [unit, 1 : i32], legacy_trait_fold_in_place_count = 1,
+ dialect_fold = [7 : i32, 7 : i32]} : () -> (i32, i32)
+ // An in-place fold of the first trait also skips the legacy trait.
+ // expected-remark @below {{fold: in place}}
+ %1:2 = "test.fold_dispatch_traits"()
+ {trait_fold_in_place_count = 1, legacy_trait_fold = [1 : i32, 1 : i32]}
+ : () -> (i32, i32)
+ // A partial fold keeps the in-place bit of the trait.
+ // expected-remark @below {{fold: [1 : i32, keep] in place}}
+ %2:2 = "test.fold_dispatch_traits"()
+ {trait_fold = [1 : i32, unit], trait_fold_with_in_place,
+ legacy_trait_fold = [7 : i32, 7 : i32]} : () -> (i32, i32)
+ // The first trait keeps every result, which normalizes to a failure, so the
+ // legacy trait folds.
+ // expected-remark @below {{fold: [1 : i32, 1 : i32]}}
+ %3:2 = "test.fold_dispatch_traits"()
+ {trait_fold = ["result:0", "result:1"],
+ legacy_trait_fold = [1 : i32, 1 : i32]} : () -> (i32, i32)
+ // After a failure of the first trait, the legacy trait changes the op in
+ // place.
+ // expected-remark @below {{fold: in place}}
+ %4:2 = "test.fold_dispatch_traits"() {legacy_trait_fold_in_place_count = 1}
+ : () -> (i32, i32)
+ // When both traits fail, the fold falls back on the dialect.
+ // expected-remark @below {{fold: [3 : i32, 3 : i32]}}
+ %5:2 = "test.fold_dispatch_traits"() {dialect_fold = [3 : i32, 3 : i32]}
+ : () -> (i32, i32)
+ // A trait fold can name another result that it also replaces: a cast in a
+ // graph region can return such a chain. The fold hook does not drop the
+ // replacements of a trait fold.
+ // expected-remark @below {{fold: [result 1, 1 : i32]}}
+ %6:2 = "test.fold_dispatch_traits"() {trait_fold = ["result:1", 1 : i32]}
+ : () -> (i32, i32)
+ return
+}
+
+// The fallback fold of the test dialect also has the `OpFoldResults` form,
+// which replaces its legacy method.
+func.func @dialect_fallback_results_form() {
+ // expected-remark @below {{fold: [1 : i32, keep]}}
+ %0:2 = "test.fold_dispatch_fallback"()
+ {dialect_results_fold = [1 : i32, unit]} : () -> (i32, i32)
+ // expected-remark @below {{fold: in place}}
+ %1:2 = "test.fold_dispatch_fallback"()
+ {dialect_results_fold_in_place_count = 1} : () -> (i32, i32)
+ // The legacy method does not run.
+ // expected-remark @below {{fold: [2 : i32, 2 : i32]}}
+ %2:2 = "test.fold_dispatch_fallback"()
+ {dialect_results_fold = [2 : i32, 2 : i32],
+ dialect_fold = [7 : i32, 7 : i32]} : () -> (i32, i32)
+ // The result of the fallback is normalized.
+ // expected-remark @below {{fold: failure}}
+ %3:2 = "test.fold_dispatch_fallback"()
+ {dialect_results_fold = ["result:0", "result:1"]} : () -> (i32, i32)
+ return
+}
+
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(
%root: !transform.any_op {transform.readonly}) {
%ops = transform.structured.match
- ops{["test.fold_dispatch_fallback", "test.fold_dispatch"]} in %root
+ ops{["test.fold_dispatch_fallback", "test.fold_dispatch",
+ "test.fold_dispatch_traits"]} in %root
: (!transform.any_op) -> !transform.any_op
transform.test_fold %ops : !transform.any_op
transform.yield
diff --git a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
index 6ddf8cadf2e90..73aa7ef767c01 100644
--- a/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp
@@ -306,7 +306,16 @@ struct TestOpAsmInterface : public OpAsmDialectInterface {
struct TestDialectFoldInterface : public DialectFoldInterface {
using DialectFoldInterface::DialectFoldInterface;
- /// The fallback fold of the fold tests.
+ /// The fallback fold of the fold tests. Without `dialect_results_fold`
+ /// attributes, it uses the legacy method.
+ OpFoldResults fold(Operation *op, ArrayRef<Attribute> operands) const final {
+ if (op->hasAttr("dialect_results_fold") ||
+ op->hasAttr("dialect_results_fold_in_place_count"))
+ return getConfiguredFoldResults(op, operands, "dialect_results_fold");
+ return DialectFoldInterface::fold(op, operands);
+ }
+
+ /// The legacy fallback fold of the fold tests.
LogicalResult fold(Operation *op, ArrayRef<Attribute> operands,
SmallVectorImpl<OpFoldResult> &results) const final {
return getConfiguredLegacyFoldResults(op, operands, "dialect_fold",
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index fc3bbf8035065..855b5b1b7798f 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -2401,3 +2401,8 @@ OpFoldResults TestFoldDispatchOp::fold(FoldAdaptor adaptor) {
return getConfiguredFoldResults(getOperation(), adaptor.getOperands(),
"fold");
}
+
+static_assert(op_definition_impl::detect_has_fold_results_trait<
+ TestResultsFoldTrait<TestFoldDispatchTraitsOp>>::value);
+static_assert(!op_definition_impl::detect_has_single_result_fold_trait<
+ TestResultsFoldTrait<TestFoldDispatchTraitsOp>>::value);
diff --git a/mlir/test/lib/Dialect/Test/TestOps.h b/mlir/test/lib/Dialect/Test/TestOps.h
index 1912db4185b21..804550518b3c3 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.h
+++ b/mlir/test/lib/Dialect/Test/TestOps.h
@@ -287,6 +287,18 @@ class TestLegacyFoldTrait
}
};
+/// A fold trait in the `OpFoldResults` form for the fold tests. It reads its
+/// result from the `trait_fold` attributes.
+template <typename ConcreteType>
+class TestResultsFoldTrait
+ : public mlir::OpTrait::TraitBase<ConcreteType, TestResultsFoldTrait> {
+public:
+ static mlir::OpFoldResults
+ foldTrait(mlir::Operation *op, llvm::ArrayRef<mlir::Attribute> operands) {
+ return getConfiguredFoldResults(op, operands, "trait_fold");
+ }
+};
+
} // namespace test
#define GET_OP_CLASSES
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 6bbfa93f818be..5e48a5b60185c 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1825,6 +1825,20 @@ def TestFoldDispatchOp : TEST_Op<"fold_dispatch", [TestLegacyFoldTrait]> {
}];
}
+def TestResultsFoldTrait : NativeOpTrait<"TestResultsFoldTrait"> {
+ let cppNamespace = "::test";
+}
+
+// Op for the fold tests in IR/fold-dispatch.mlir. It has no fold of its own.
+// Its fold trait in the `OpFoldResults` form reads its result from the
+// `trait_fold` attributes, and runs before its legacy fold trait.
+def TestFoldDispatchTraitsOp
+ : TEST_Op<"fold_dispatch_traits",
+ [TestResultsFoldTrait, TestLegacyFoldTrait]> {
+ let arguments = (ins Variadic<AnyType>);
+ let results = (outs Variadic<AnyType>);
+}
+
//===----------------------------------------------------------------------===//
// Test Patterns (Symbol Binding)
//===----------------------------------------------------------------------===//
More information about the llvm-branch-commits
mailing list