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?= <[email protected]> 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 <[email protected]> --- 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 b09f51e217f3917..74c5d1ee85860f9 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 991b084b8b44d5e..b163c649e2f3350 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 3b81eef8b864bc1..c32bc2946bf2441 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 138be794f605c02..4c75ac6e270437f 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 5a4247d848d1834..52be6d164119f35 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 e771d3f194bcc49..dd216e1932d0491 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 05c872daf5dab28..7e32cebdf739b8c 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 4141882bda878ff..b29962b0eee22f2 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 6ddf8cadf2e9088..73aa7ef767c019c 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 fc3bbf803506525..855b5b1b7798fb2 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 1912db4185b21fa..804550518b3c38f 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 6bbfa93f818be49..5e48a5b60185c53 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) //===----------------------------------------------------------------------===// _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
