[mlir] Add a new fold API using Generic Adaptors

[zero9178]

Summary

This is part of the RFC for a better fold API: https://discourse.llvm.org/t/rfc-a-better-fold-api-using-more-generic-adaptors/67374

This patch implements the required foldHook changes and the TableGen machinery for generating fold method signatures using FoldAdaptor for ops, based on the value of useFoldAPI. It may be one of 2 values, with convenient named constants to create a quasi enum. The new fold method will then be generated if kEmitFoldAdaptorFolder is used.

Since the new FoldAdaptor approach is strictly better than the old signature, part of this patch updates the documentation and all example to encourage use of the new fold signature.
Included are also tests exercising the new API, ensuring proper construction of the FoldAdaptor and proper generation by TableGen.

[Mogball]

mlir/docs/Canonicalization.md

141

I'd honestly prefer this to be a dialect setting instead of a per-operation setting.

mlir/docs/Tutorials/Toy/Ch-7.md

464

please add an argument name. Omitting argument names is unusual in the MLIR coding style

mlir/include/mlir/IR/OpDefinition.h

1714

can you add a has_fold_adaptor_v? This one is used a few times

mlir/test/IR/test-fold-adaptor.mlir

5

it doesn't seem correct to me to have a test under IR/ contain arith ops

17

please add a newline

[zero9178]

Address review comments:

• address style issues
• use test dialect ops
• change the fold API to be selectable per dialect instead of per Op

[zero9178]

mlir/docs/Canonicalization.md

141

Thank you for telling me. I tend to agree and tentatively updated the patch, making it a dialect setting. I have also added a reply to the RFC thread to document the new behaviour and how it changes migrations.

mlir/include/mlir/IR/OpDefinition.h

1714

I have gone ahead and have done this for all the detect_* typedefs since there is no real benefit of having the struct version around in existing code. I have also changed the naming to be more consistent with how it is used in LLVM, that is _t suffix for the detector using declaration, and the _v boolean you wanted for the actual boolean.

[zero9178]

clang-format

[Mogball]

You know, if you make kEmitRawAttributesFolder emit a warning, users will migrate a lot faster 😂

Though that means you might at least be on the hook for updating all upstream dialects.

[zero9178]

In D140886#4032203, @Mogball wrote:

You know, if you make kEmitRawAttributesFolder emit a warning, users will migrate a lot faster 😂

Though that means you might at least be on the hook for updating all upstream dialects.

If I was poetic enough to formulate the warning message in a way that users would be moved with enthusiams to switch to the new API I would!
I fear I'd instead receive compliants about being spammed by 500 warnings messages however 😅

[Mogball]

In D140886#4032311, @zero9178 wrote:

In D140886#4032203, @Mogball wrote:

You know, if you make kEmitRawAttributesFolder emit a warning, users will migrate a lot faster 😂

Though that means you might at least be on the hook for updating all upstream dialects.

If I was poetic enough to formulate the warning message in a way that users would be moved with enthusiams to switch to the new API I would!
I fear I'd instead receive compliants about being spammed by 500 warnings messages however 😅

Just hold your ground.

LGTM but please give others a change to take a look at this.

mlir/include/mlir/IR/OpDefinition.h

1815

please add braces here

mlir/test/mlir-tblgen/has-fold-invalid-values.td

16

can you emit a more detailed error message? Like "invalid value for dialect useFoldAPI"? This is the kind of error message that typically results in me grepping the code base for where it came from

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

2342–2347

please add braces here. If one branch has braces, they all should.

[zero9178]

Address review comments

[mehdi_amini]

In D140886#4032311, @zero9178 wrote:

In D140886#4032203, @Mogball wrote:

You know, if you make kEmitRawAttributesFolder emit a warning, users will migrate a lot faster 😂

Though that means you might at least be on the hook for updating all upstream dialects.

If I was poetic enough to formulate the warning message in a way that users would be moved with enthusiams to switch to the new API I would!
I fear I'd instead receive compliants about being spammed by 500 warnings messages however 😅

I would implement it such that the warning is emitted at most once per tablegen invocation, that should make it not too noisy I think?
(in-tree dialects need to be clean first of course)

[zero9178]

In D140886#4036905, @mehdi_amini wrote:

In D140886#4032311, @zero9178 wrote:

In D140886#4032203, @Mogball wrote:

You know, if you make kEmitRawAttributesFolder emit a warning, users will migrate a lot faster 😂

Though that means you might at least be on the hook for updating all upstream dialects.

If I was poetic enough to formulate the warning message in a way that users would be moved with enthusiams to switch to the new API I would!
I fear I'd instead receive compliants about being spammed by 500 warnings messages however 😅

I would implement it such that the warning is emitted at most once per tablegen invocation, that should make it not too noisy I think?
(in-tree dialects need to be clean first of course)

Taking upstream MLIR as benchmark, this leads to "just" 42 warnings, which is definitely less noisey (and the repo probably contains more dialects than most downstream projects).
I was previously under the impression that warnings emitted by TableGen were rather unpopular, especially since there is no way to silence them, except fixing the issue.
I suppose that is what we want in this case though and migrating should be easy, even if it has to be done manually.

That said, what time do you think this warning should be introduced? Immediately right after in-tree is clean, but before changing the default to the new one?

[Mogball]

There's a build option to disable TableGen warnings. That should work, of course.

[rriddle]

mlir/examples/toy/Ch7/mlir/ToyCombine.cpp

34

Do we still need the dyn_cast_or_null here?

[zero9178]

mlir/examples/toy/Ch7/mlir/ToyCombine.cpp

34

Yes. The FoldAdaptor (which is really just a GenericAdaptor<ArrayRef<Attribute>>) currently doesn't treat Attributes any special. So all the getters currently simply return the element type of the input range, which is Attribute in this case. (just like the Adaptors always return Value as element type of ValueRange).

Changing this so that the return types of FoldAdaptor returns the concrete Attribute type is something I'd like to look into in the future as the next improvement of fold, but this will require a few more changes, the most important one in my mind being that we'll probably need some field in TableGen Type, which tells us what the corresponding Attribute type is. We could then generate a template specialization using these as return types instead.

mlir/docs/Canonicalization.md

 rewriter.
 
 In [ODS](DefiningDialects/Operations.md), an operation can set the `hasFolder` bit to generate
 a declaration for the `fold` method. This method takes on a different form,
 depending on the structure of the operation.
 
 ```tablegen
 def MyOp : ... {
   let hasFolder = 1;
 }
 ```
 
 If the operation has a single result the following will be generated:
 
 ```c++
 /// Implementations of this hook can only perform the following changes to the
 /// operation:
 ///
 ///  1. They can leave the operation alone and without changing the IR, and
 ///     return nullptr.
 ///  2. They can mutate the operation in place, without changing anything else
 ///     in the IR. In this case, return the operation itself.
 ///  3. They can return an existing value or attribute that can be used instead
 ///     of the operation. The caller will remove the operation and use that
 ///     result instead.
 ///
-OpFoldResult MyOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult MyOp::fold(FoldAdaptor adaptor) {
   ...
 }
 ```
 
 Otherwise, the following is generated:
 
 ```c++
 /// Implementations of this hook can only perform the following changes to the
 /// operation:
 ///
 ///  1. They can leave the operation alone and without changing the IR, and
 ///     return failure.
 ///  2. They can mutate the operation in place, without changing anything else
 ///     in the IR. In this case, return success.
 ///  3. They can return a list of existing values or attribute that can be used
 ///     instead of the operation. In this case, fill in the results list and
 ///     return success. The results list must correspond 1-1 with the results of
 ///     the operation, partial folding is not supported. The caller will remove
 ///     the operation and use those results instead.
 ///
 /// Note that this mechanism cannot be used to remove 0-result operations.
-LogicalResult MyOp::fold(ArrayRef<Attribute> operands,
+LogicalResult MyOp::fold(FoldAdaptor adaptor,
                          SmallVectorImpl<OpFoldResult> &results) {
   ...
 }
 ```
 
-In the above, for each method an `ArrayRef<Attribute>` is provided that
-corresponds to the constant attribute value of each of the operands. These
+In the above, for each method a `FoldAdaptor` is provided with getters for
+each of the operands, returning the corresponding constant attribute. These
 operands are those that implement the `ConstantLike` trait. If any of the
 operands are non-constant, a null `Attribute` value is provided instead. For
 example, if MyOp provides three operands [`a`, `b`, `c`], but only `b` is
-constant then `operands` will be of the form [Attribute(), b-value,
-Attribute()].
+constant then `adaptor` will return Attribute() for `getA()` and `getC()`, 
+and b-value for `getB()`.
 
 Also above, is the use of `OpFoldResult`. This class represents the possible
 result of folding an operation result: either an SSA `Value`, or an
 `Attribute`(for a constant result). If an SSA `Value` is provided, it *must*
 correspond to an existing value. The `fold` methods are not permitted to
 generate new `Value`s. There are no specific restrictions on the form of the
 `Attribute` value returned, but it is important to ensure that the `Attribute`
 representation of a specific `Type` is consistent.

mlir/docs/DefiningDialects/_index.md

 /// Generate verification for the given attribute, whose name is prefixed by the namespace
 /// of this dialect, that was used on the attribute dictionary of a region result.
 /// Note: As described above, when a region entry block has a dictionary is up to the individual
 /// operation to define. 
 LogicalResult MyDialect::verifyRegionResultAttribute(Operation *op, unsigned regionIndex,
                                                      unsigned argIndex, NamedAttribute attribute);
 ```
 
+#### `useFoldAPI`
+
+There are currently two possible values that are allowed to be assigned to this
+field:
+* `kEmitFoldAdaptorFolder` generates a `fold` method making use of the op's 
+  `FoldAdaptor` to allow access of operands via convenient getter.
+
+  Generated code example:
+  ```cpp
+  OpFoldResult fold(FoldAdaptor adaptor);
+  // or
+  LogicalResult fold(FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult>& results);
+  ```
+* `kEmitRawAttributesFolder` generates the deprecated legacy `fold`
+  method, containing `ArrayRef<Attribute>` in the parameter list instead of
+  the op's `FoldAdaptor`. This API is scheduled for removal and should not be 
+  used by new dialects.
+
+  Generated code example:
+  ```cpp
+  OpFoldResult fold(ArrayRef<Attribute> operands);
+  // or
+  LogicalResult fold(ArrayRef<Attribute> operands, SmallVectorImpl<OpFoldResult>& results);
+  ```
+
 ### Operation Interface Fallback
 
 Some dialects have an open ecosystem and don't register all of the possible operations. In such
 cases it is still possible to provide support for implementing an `OpInterface` for these 
 operations. When an operation isn't registered or does not provide an implementation for an 
 interface, the query will fallback to the dialect itself. The `hasOperationInterfaceFallback`
 field may be used to declare this fallback for operations:
 

mlir/docs/Tutorials/Toy/Ch-7.md

 We have several `toy.struct_access` operations that access into a
 `toy.struct_constant`. As detailed in [chapter 3](Ch-3.md) (FoldConstantReshape),
 we can add folders for these `toy` operations by setting the `hasFolder` bit
 on the operation definition and providing a definition of the `*Op::fold`
 method.
 
 ```c++
 /// Fold constants.
-OpFoldResult ConstantOp::fold(ArrayRef<Attribute> operands) { return value(); }
+OpFoldResult ConstantOp::fold(FoldAdaptor adaptor) { return value(); }
 
 /// Fold struct constants.
-OpFoldResult StructConstantOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult StructConstantOp::fold(FoldAdaptor adaptor) {
   return value();
 }
 
 /// Fold simple struct access operations that access into a constant.
-OpFoldResult StructAccessOp::fold(ArrayRef<Attribute> operands) {
-  auto structAttr = operands.front().dyn_cast_or_null<mlir::ArrayAttr>();
+OpFoldResult StructAccessOp::fold(FoldAdaptor adaptor) {
+  auto structAttr = adaptor.getInput().dyn_cast_or_null<mlir::ArrayAttr>();
   if (!structAttr)
     return nullptr;
 
   size_t elementIndex = index().getZExtValue();
   return structAttr[elementIndex];
 }
 ```
 

mlir/examples/toy/Ch7/include/toy/Ops.td

def Toy_Dialect : Dialect {
 
   // We set this bit to generate a declaration of the `materializeConstant`
   // method so that we can materialize constants for our toy operations.
   let hasConstantMaterializer = 1;
 
   // We set this bit to generate the declarations for the dialect's type parsing
   // and printing hooks.
   let useDefaultTypePrinterParser = 1;
+
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 // Base class for toy dialect operations. This operation inherits from the base
 // `Op` class in OpBase.td, and provides:
 //   * The parent dialect of the operation.
 //   * The mnemonic for the operation, or the name without the dialect prefix.
 //   * A list of traits for the operation.
 class Toy_Op<string mnemonic, list<Trait> traits = []> :

mlir/examples/toy/Ch7/mlir/ToyCombine.cpp

 using namespace toy;
 
 namespace {
 /// Include the patterns defined in the Declarative Rewrite framework.
 #include "ToyCombine.inc"
 } // namespace
 
 /// Fold constants.
-OpFoldResult ConstantOp::fold(ArrayRef<Attribute> operands) {
-  return getValue();
-}
+OpFoldResult ConstantOp::fold(FoldAdaptor adaptor) { return getValue(); }
 
 /// Fold struct constants.
-OpFoldResult StructConstantOp::fold(ArrayRef<Attribute> operands) {
-  return getValue();
-}
+OpFoldResult StructConstantOp::fold(FoldAdaptor adaptor) { return getValue(); }
 
 /// Fold simple struct access operations that access into a constant.
-OpFoldResult StructAccessOp::fold(ArrayRef<Attribute> operands) {
-  auto structAttr = operands.front().dyn_cast_or_null<mlir::ArrayAttr>();
+OpFoldResult StructAccessOp::fold(FoldAdaptor adaptor) {
+  auto structAttr = adaptor.getInput().dyn_cast_or_null<mlir::ArrayAttr>();
   if (!structAttr)
     return nullptr;
 
   size_t elementIndex = getIndex();
   return structAttr[elementIndex];
 }
 
 /// This is an example of a c++ rewrite pattern for the TransposeOp. It

mlir/include/mlir/IR/DialectBase.td

 
 #ifndef DIALECTBASE_TD
 #define DIALECTBASE_TD
 
 //===----------------------------------------------------------------------===//
 // Dialect definitions
 //===----------------------------------------------------------------------===//
 
+// Generate 'fold' method with 'ArrayRef<Attribute>' parameter.
+// New code should prefer using 'kEmitFoldAdaptorFolder' and
+// consider 'kEmitRawAttributesFolder' deprecated and to be
+// removed in the future.
+defvar kEmitRawAttributesFolder = 0;
+// Generate 'fold' method with 'FoldAdaptor' parameter.
+defvar kEmitFoldAdaptorFolder = 1;
+
 class Dialect {
   // The name of the dialect.
   string name = ?;
 
   // Short summary of the dialect.
   string summary = ?;
 
   // The description of the dialect.
class Dialect {
   // individual type directly.
   bit useDefaultTypePrinterParser = 0;
 
   // If this dialect overrides the hook for canonicalization patterns.
   bit hasCanonicalizer = 0;
 
   // If this dialect can be extended at runtime with new operations or types.
   bit isExtensible = 0;
+
+  // Fold API to use for operations in this dialect.
+  int useFoldAPI = kEmitRawAttributesFolder;
 }
 
 #endif // DIALECTBASE_TD

mlir/include/mlir/IR/OpDefinition.h

if (!info)
           ConcreteType::getOperationName() + ".");
     (checkInterfaceTarget<Models>(), ...);
     info->attachInterface<Models...>();
   }
 
 private:
   /// Trait to check if T provides a 'fold' method for a single result op.
   template <typename T, typename... Args>
-  using has_single_result_fold =
+  using has_single_result_fold_t =
       decltype(std::declval<T>().fold(std::declval<ArrayRef<Attribute>>()));
   template <typename T>
-  using detect_has_single_result_fold =
-      llvm::is_detected<has_single_result_fold, T>;
+  constexpr static bool has_single_result_fold_v =
+      llvm::is_detected<has_single_result_fold_t, T>::value;
   /// Trait to check if T provides a general 'fold' method.
   template <typename T, typename... Args>
-  using has_fold = decltype(std::declval<T>().fold(
+  using has_fold_t = decltype(std::declval<T>().fold(
       std::declval<ArrayRef<Attribute>>(),
       std::declval<SmallVectorImpl<OpFoldResult> &>()));
   template <typename T>
-  using detect_has_fold = llvm::is_detected<has_fold, T>;
+  constexpr static bool has_fold_v = llvm::is_detected<has_fold_t, T>::value;
+  /// Trait to check if T provides a 'fold' method with a FoldAdaptor for a
+  /// single result op.
+  template <typename T, typename... Args>
+  using has_fold_adaptor_single_result_fold_t =
+      decltype(std::declval<T>().fold(std::declval<typename T::FoldAdaptor>()));
+  template <class T>
+  constexpr static bool has_fold_adaptor_single_result_v =
+      llvm::is_detected<has_fold_adaptor_single_result_fold_t, T>::value;
+  /// Trait to check if T provides a general 'fold' method with a FoldAdaptor.
+  template <typename T, typename... Args>
+  using has_fold_adaptor_fold_t = decltype(std::declval<T>().fold(
+      std::declval<typename T::FoldAdaptor>(),
+      std::declval<SmallVectorImpl<OpFoldResult> &>()));
+  template <class T>
+  constexpr static bool has_fold_adaptor_v =
+      llvm::is_detected<has_fold_adaptor_fold_t, T>::value;
+
   /// Trait to check if T provides a 'print' method.
   template <typename T, typename... Args>
   using has_print =
       decltype(std::declval<T>().print(std::declval<OpAsmPrinter &>()));
   template <typename T>
   using detect_has_print = llvm::is_detected<has_print, T>;
 
   /// Trait to check if T provides a 'ConcreteEntity' type alias.
private:
 
   /// Return the internal implementations of each of the OperationName
   /// hooks.
   /// Implementation of `FoldHookFn` OperationName hook.
   static OperationName::FoldHookFn getFoldHookFn() {
     // If the operation is single result and defines a `fold` method.
     if constexpr (llvm::is_one_of<OpTrait::OneResult<ConcreteType>,
                                   Traits<ConcreteType>...>::value &&
-                  detect_has_single_result_fold<ConcreteType>::value)
+                  (has_single_result_fold_v<ConcreteType> ||
+                   has_fold_adaptor_single_result_v<ConcreteType>))
       return [](Operation *op, ArrayRef<Attribute> operands,
                 SmallVectorImpl<OpFoldResult> &results) {
         return foldSingleResultHook<ConcreteType>(op, operands, results);
       };
     // The operation is not single result and defines a `fold` method.
-    if constexpr (detect_has_fold<ConcreteType>::value)
+    if constexpr (has_fold_v<ConcreteType> || has_fold_adaptor_v<ConcreteType>)
       return [](Operation *op, ArrayRef<Attribute> operands,
                 SmallVectorImpl<OpFoldResult> &results) {
         return foldHook<ConcreteType>(op, operands, results);
       };
     // The operation does not define a `fold` method.
     return [](Operation *op, ArrayRef<Attribute> operands,
               SmallVectorImpl<OpFoldResult> &results) {
       // In this case, we only need to fold the traits of the operation.
       return op_definition_impl::foldTraits<Traits<ConcreteType>...>(
           op, operands, results);
     };
   }
   /// Return the result of folding a single result operation that defines a
   /// `fold` method.
   template <typename ConcreteOpT>
   static LogicalResult
   foldSingleResultHook(Operation *op, ArrayRef<Attribute> operands,
                        SmallVectorImpl<OpFoldResult> &results) {
-    OpFoldResult result = cast<ConcreteOpT>(op).fold(operands);
+    OpFoldResult result;
+    if constexpr (has_fold_adaptor_single_result_v<ConcreteOpT>)
+      result = cast<ConcreteOpT>(op).fold(typename ConcreteOpT::FoldAdaptor(
+          operands, op->getAttrDictionary(), op->getRegions()));
+    else
+      result = cast<ConcreteOpT>(op).fold(operands);
 
     // If the fold failed or was in-place, try to fold the traits of the
     // operation.
     if (!result || result.template dyn_cast<Value>() == op->getResult(0)) {
       if (succeeded(op_definition_impl::foldTraits<Traits<ConcreteType>...>(
               op, operands, results)))
         return success();
       return success(static_cast<bool>(result));
     }
     results.push_back(result);
     return success();
   }
   /// Return the result of folding an operation that defines a `fold` method.
   template <typename ConcreteOpT>
   static LogicalResult foldHook(Operation *op, ArrayRef<Attribute> operands,
                                 SmallVectorImpl<OpFoldResult> &results) {
-    LogicalResult result = cast<ConcreteOpT>(op).fold(operands, results);
+    auto result = LogicalResult::failure();
+    if constexpr (has_fold_adaptor_v<ConcreteOpT>) {
+      result = cast<ConcreteOpT>(op).fold(
+          typename ConcreteOpT::FoldAdaptor(operands, op->getAttrDictionary(),
+                                            op->getRegions()),
+          results);
+    } else {
+      result = cast<ConcreteOpT>(op).fold(operands, results);
+    }
 
     // If the fold failed or was in-place, try to fold the traits of the
     // operation.
     if (failed(result) || results.empty()) {
       if (succeeded(op_definition_impl::foldTraits<Traits<ConcreteType>...>(
               op, operands, results)))
         return success();
     }

mlir/include/mlir/TableGen/Dialect.h

public:
   /// Returns true if this dialect should generate the default dispatch for
   /// type printing/parsing.
   bool useDefaultTypePrinterParser() const;
 
   /// Returns true if this dialect can be extended at runtime with new
   /// operations or types.
   bool isExtensible() const;
 
+  enum class FolderAPI {
+    RawAttributes = 0, /// fold method with ArrayRef<Attribute>.
+    FolderAdaptor = 1, /// fold method with the operation's FoldAdaptor.
+  };
+
+  /// Returns the folder API that should be emitted for operations in this
+  /// dialect.
+  FolderAPI getFolderAPI() const;
+
   // Returns whether two dialects are equal by checking the equality of the
   // underlying record.
   bool operator==(const Dialect &other) const;
 
   bool operator!=(const Dialect &other) const { return !(*this == other); }
 
   // Compares two dialects by comparing the names of the dialects.
   bool operator<(const Dialect &other) const;

mlir/include/mlir/TableGen/Operator.h

public:
   std::string getGetterName(StringRef name) const;
 
   /// Returns the setter name for the accessor of `name`.
   std::string getSetterName(StringRef name) const;
 
   /// Returns the remove name for the accessor of `name`.
   std::string getRemoverName(StringRef name) const;
 
+  bool hasFolder() const;
+
 private:
   /// Populates the vectors containing operands, attributes, results and traits.
   void populateOpStructure();
 
   /// Populates type inference info (mostly equality) with input a mapping from
   /// names to indices for arguments and results.
   void populateTypeInferenceInfo(
       const llvm::StringMap<int> &argumentsAndResultsIndex);

mlir/lib/TableGen/Dialect.cpp

 bool Dialect::useDefaultTypePrinterParser() const {
   return def->getValueAsBit("useDefaultTypePrinterParser");
 }
 
 bool Dialect::isExtensible() const {
   return def->getValueAsBit("isExtensible");
 }
 
+Dialect::FolderAPI Dialect::getFolderAPI() const {
+  int64_t value = def->getValueAsInt("useFoldAPI");
+  if (value < static_cast<int64_t>(FolderAPI::RawAttributes) ||
+      value > static_cast<int64_t>(FolderAPI::FolderAdaptor))
+    llvm::PrintFatalError(def->getLoc(),
+                          "Invalid value for dialect field `useFoldAPI`");
+
+  return static_cast<FolderAPI>(value);
+}
+
 bool Dialect::operator==(const Dialect &other) const {
   return def == other.def;
 }
 
 bool Dialect::operator<(const Dialect &other) const {
   return getName() < other.getName();
 }

mlir/lib/TableGen/Operator.cpp

 
 std::string Operator::getSetterName(StringRef name) const {
   return "set" + convertToCamelFromSnakeCase(name, /*capitalizeFirst=*/true);
 }
 
 std::string Operator::getRemoverName(StringRef name) const {
   return "remove" + convertToCamelFromSnakeCase(name, /*capitalizeFirst=*/true);
 }
+
+bool Operator::hasFolder() const { return def.getValueAsBit("hasFolder"); }

mlir/test/IR/test-fold-adaptor.mlir

+// RUN: mlir-opt %s -canonicalize -split-input-file | FileCheck %s
+
+func.func @test() -> i32 {
+  %c5 = "test.constant"() {value = 5 : i32} : () -> i32
+  %c1 = "test.constant"() {value = 1 : i32} : () -> i32
+  %c2 = "test.constant"() {value = 2 : i32} : () -> i32
+  %c3 = "test.constant"() {value = 3 : i32} : () -> i32
+  %res = test.fold_with_fold_adaptor %c5, [ %c1, %c2], { (%c3), (%c3) } {
+    %c0 = "test.constant"() {value = 0 : i32} : () -> i32
+  }
+  return %res : i32
+}
+
+// CHECK-LABEL: func.func @test
+// CHECK-NEXT: %[[C:.*]] = "test.constant"() {value = 33 : i32}
+// CHECK-NEXT: return %[[C]]

mlir/test/lib/Dialect/Test/TestDialect.cpp

 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
 
+#include <numeric>
+
 // Include this before the using namespace lines below to
 // test that we don't have namespace dependencies.
 #include "TestOpsDialect.cpp.inc"
 
 using namespace mlir;
 using namespace test;
 
 void test::registerTestDialect(DialectRegistry &registry) {
OpFoldResult TestOpInPlaceFold::fold(ArrayRef<Attribute> operands) {
   }
   return {};
 }
 
 OpFoldResult TestPassthroughFold::fold(ArrayRef<Attribute> operands) {
   return getOperand();
 }
 
+OpFoldResult TestOpFoldWithFoldAdaptor::fold(FoldAdaptor adaptor) {
+  int64_t sum = 0;
+  if (auto value = dyn_cast_or_null<IntegerAttr>(adaptor.getOp()))
+    sum += value.getValue().getSExtValue();
+
+  for (Attribute attr : adaptor.getVariadic())
+    if (auto value = dyn_cast_or_null<IntegerAttr>(attr))
+      sum += 2 * value.getValue().getSExtValue();
+
+  for (ArrayRef<Attribute> attrs : adaptor.getVarOfVar())
+    for (Attribute attr : attrs)
+      if (auto value = dyn_cast_or_null<IntegerAttr>(attr))
+        sum += 3 * value.getValue().getSExtValue();
+
+  sum += 4 * std::distance(adaptor.getBody().begin(), adaptor.getBody().end());
+
+  return IntegerAttr::get(getType(), sum);
+}
+
 LogicalResult OpWithInferTypeInterfaceOp::inferReturnTypes(
     MLIRContext *, std::optional<Location> location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {
   if (operands[0].getType() != operands[1].getType()) {
     return emitOptionalError(location, "operand type mismatch ",
                              operands[0].getType(), " vs ",
                              operands[1].getType());

mlir/test/lib/Dialect/Test/TestOps.td

def TestOpInPlaceFoldSuccess : TEST_Op<"op_in_place_fold_success"> {
   let extraClassDefinition = [{
     ::mlir::LogicalResult $cppClass::fold(ArrayRef<Attribute> operands,
         SmallVectorImpl<OpFoldResult> &results) {
       return success();
     }
   }];
 }
 
+def TestOpFoldWithFoldAdaptor
+  : TEST_Op<"fold_with_fold_adaptor",
+      [AttrSizedOperandSegments, NoTerminator]> {
+  let arguments = (ins
+    I32:$op,
+    DenseI32ArrayAttr:$attr,
+    Variadic<I32>:$variadic,
+    VariadicOfVariadic<I32, "attr">:$var_of_var
+  );
+
+  let results = (outs I32:$res);
+
+  let regions = (region AnyRegion:$body);
+
+  let assemblyFormat = [{
+    $op `,` `[` $variadic `]` `,` `{` $var_of_var `}` $body attr-dict-with-keyword
+  }];
+
+  let hasFolder = 0;
+
+  let extraClassDeclaration = [{
+    ::mlir::OpFoldResult fold(FoldAdaptor adaptor);
+  }];
+}
+
 // An op that always fold itself.
 def TestPassthroughFold : TEST_Op<"passthrough_fold"> {
   let arguments = (ins AnyType:$op);
   let results = (outs AnyType);
   let hasFolder = 1;
 }
 
 def TestDialectCanonicalizerOp : TEST_Op<"dialect_canonicalizable"> {

mlir/test/mlir-tblgen/has-fold-invalid-values.td

+// RUN: not mlir-tblgen -gen-op-decls -I %S/../../include %s 2>&1 | FileCheck %s
+
+include "mlir/IR/OpBase.td"
+
+def Test_Dialect : Dialect {
+  let name = "test";
+  let cppNamespace = "NS";
+  let useFoldAPI = 3;
+}
+
+def InvalidValue_Op : Op<Test_Dialect, "invalid_op"> {
+  let hasFolder = 1;
+}
+
+// CHECK: Invalid value for dialect field `useFoldAPI`

mlir/test/mlir-tblgen/op-decl-and-defs.td

 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b, ::mlir::IntegerAttr attr1);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b, ::mlir::IntegerAttr attr1);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type r, ::mlir::Value a, ::mlir::Value b, uint32_t attr1);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b, uint32_t attr1);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b, uint32_t attr1);
 // CHECK: static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
 
+def TestWithNewFold_Dialect : Dialect {
+  let name = "test";
+  let cppNamespace = "::mlir::testWithFold";
+  let useFoldAPI = kEmitFoldAdaptorFolder;
+}
+
+def NS_MOp : Op<TestWithNewFold_Dialect, "op_with_single_result_and_fold_adaptor_fold", []> {
+  let results = (outs AnyType:$res);
+
+  let hasFolder = 1;
+}
+
+// CHECK-LABEL: class MOp :
+// CHECK: ::mlir::OpFoldResult fold(FoldAdaptor adaptor);
+
+def NS_NOp : Op<TestWithNewFold_Dialect, "op_with_multiple_results_and_fold_adaptor_fold", []> {
+  let results = (outs AnyType:$res1, AnyType:$res2);
+
+  let hasFolder = 1;
+}
+
+// CHECK-LABEL: class NOp :
+// CHECK: ::mlir::LogicalResult fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult> &results);
 
 // Test that type defs have the proper namespaces when used as a constraint.
 // ---
 
 def Test_Dialect2 : Dialect {
   let name = "test";
   let cppNamespace = "::mlir::dialect2";
 }

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

void OpEmitter::genCanonicalizerDecls() {
   // If synthesizing the method, fill it it.
   if (hasBody) {
     ERROR_IF_PRUNED(method, "getCanonicalizationPatterns", op);
     method->body() << "  results.add(canonicalize);\n";
   }
 }
 
 void OpEmitter::genFolderDecls() {
+  if (!op.hasFolder())
+    return;
+
+  Dialect::FolderAPI folderApi = op.getDialect().getFolderAPI();
+  SmallVector<MethodParameter> paramList;
+  if (folderApi == Dialect::FolderAPI::RawAttributes)
+    paramList.emplace_back("::llvm::ArrayRef<::mlir::Attribute>", "operands");
+  else
+    paramList.emplace_back("FoldAdaptor", "adaptor");
+
+  StringRef retType;
   bool hasSingleResult =
       op.getNumResults() == 1 && op.getNumVariableLengthResults() == 0;
-
-  if (def.getValueAsBit("hasFolder")) {
-    if (hasSingleResult) {
-      auto *m = opClass.declareMethod(
-          "::mlir::OpFoldResult", "fold",
-          MethodParameter("::llvm::ArrayRef<::mlir::Attribute>", "operands"));
-      ERROR_IF_PRUNED(m, "operands", op);
-    } else {
-      SmallVector<MethodParameter> paramList;
-      paramList.emplace_back("::llvm::ArrayRef<::mlir::Attribute>", "operands");
-      paramList.emplace_back("::llvm::SmallVectorImpl<::mlir::OpFoldResult> &",
-                             "results");
-      auto *m = opClass.declareMethod("::mlir::LogicalResult", "fold",
-                                      std::move(paramList));
-      ERROR_IF_PRUNED(m, "fold", op);
-    }
+    retType = "::mlir::OpFoldResult";
+    retType = "::mlir::LogicalResult";
   }
+
+  auto *m = opClass.declareMethod(retType, "fold", std::move(paramList));
+  ERROR_IF_PRUNED(m, "fold", op);
 }
 
 void OpEmitter::genOpInterfaceMethods(const tblgen::InterfaceTrait *opTrait) {
   Interface interface = opTrait->getInterface();
 
   // Get the set of methods that should always be declared.
   auto alwaysDeclaredMethodsVec = opTrait->getAlwaysDeclaredMethods();
   llvm::StringSet<> alwaysDeclaredMethods;