[mlir][NFC] Migrate rest of the dialects to the new fold API

[zero9178]

Summary

See https://discourse.llvm.org/t/psa-new-improved-fold-method-signature-has-landed-please-update-your-downstream-projects/67618 for context

This is the last required patch in my series of patches migrating in-tree uses of the old fold API. Previous patches have been per dialect (sorry for the review "spam", hope it wasn't too bad), while this one contains remaining dialects that have only one or two fold implementations or are otherwise trivial.

[Mogball]

nice

[zero9178]

Thought I editted the commit message, but seemingly forgot 🙃
Landed in https://reviews.llvm.org/rG7df761217cd7d0026ffff23c4bdac846bb60f185

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def Affine_Dialect : Dialect {
   let name = "affine";
   let cppNamespace = "mlir";
   let hasConstantMaterializer = 1;
   let dependentDialects = ["arith::ArithDialect"];
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 // Base class for Affine dialect ops.
 class Affine_Op<string mnemonic, list<Trait> traits = []> :
     Op<Affine_Dialect, mnemonic, traits>;
 
 // Require regions to have affine.yield.
 def ImplicitAffineTerminator

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td

let extraClassDeclaration = [{
     /// BufferDeallocation pass.
     ///
     /// Note: Only ops with at least one OpResult that bufferizes to a buffer
     /// allocation (as per BufferizableOpInterface) may have this attribute.
     constexpr const static ::llvm::StringLiteral
         kEscapeAttrName = "bufferization.escape";
   }];
   let hasOperationAttrVerify = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 #endif // BUFFERIZATION_BASE

mlir/include/mlir/Dialect/Complex/IR/ComplexBase.td

def Complex_Dialect : Dialect {
   let description = [{
     The complex dialect is intended to hold complex numbers creation and
     arithmetic ops.
   }];
 
   let dependentDialects = ["arith::ArithDialect"];
   let hasConstantMaterializer = 1;
   let useDefaultAttributePrinterParser = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 #endif // COMPLEX_BASE

mlir/include/mlir/Dialect/EmitC/IR/EmitCBase.td

def EmitC_Dialect : Dialect {
   let summary = "Dialect to generate C/C++ from MLIR.";
   let description = [{
     [include "Dialects/emitc.md"]
   }];
 
   let hasConstantMaterializer = 1;
   let useDefaultTypePrinterParser = 1;
   let useDefaultAttributePrinterParser = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 #endif // MLIR_DIALECT_EMITC_IR_EMITCBASE

mlir/include/mlir/Dialect/Func/IR/FuncOps.td

 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def Func_Dialect : Dialect {
   let name = "func";
   let cppNamespace = "::mlir::func";
   let dependentDialects = ["cf::ControlFlowDialect"];
   let hasConstantMaterializer = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 // Base class for Func dialect ops.
 class Func_Op<string mnemonic, list<Trait> traits = []> :
     Op<Func_Dialect, mnemonic, traits>;
 
 //===----------------------------------------------------------------------===//
 // CallOp

mlir/include/mlir/Dialect/GPU/IR/GPUBase.td

let extraClassDeclaration = [{
     /// Returns the numeric value used to identify the private memory address
     /// space.
     static unsigned getPrivateAddressSpace() { return 5; }
   }];
 
   let dependentDialects = ["arith::ArithDialect"];
   let useDefaultAttributePrinterParser = 1;
   let useDefaultTypePrinterParser = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 def GPU_AsyncToken : DialectType<
   GPU_Dialect, CPred<"$_self.isa<::mlir::gpu::AsyncTokenType>()">, "async token type">,
              BuildableType<"mlir::gpu::AsyncTokenType::get($_builder.getContext())">;
 
 // Predicat to check if type is gpu::MMAMatrixType.
 def IsMMAMatrixTypePred : CPred<"$_self.isa<::mlir::gpu::MMAMatrixType>()">;

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

 def LLVM_Dialect : Dialect {
   let name = "llvm";
   let cppNamespace = "::mlir::LLVM";
 
   let useDefaultAttributePrinterParser = 1;
   let hasRegionArgAttrVerify = 1;
   let hasRegionResultAttrVerify = 1;
   let hasOperationAttrVerify = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 
   let extraClassDeclaration = [{
     /// Name of the data layout attributes.
     static StringRef getDataLayoutAttrName() { return "llvm.data_layout"; }
     static StringRef getAlignAttrName() { return "llvm.align"; }
     static StringRef getNoAliasAttrName() { return "llvm.noalias"; }
     static StringRef getReadonlyAttrName() { return "llvm.readonly"; }
     static StringRef getNoAliasScopesAttrName() { return "noalias_scopes"; }

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

let dependentDialects = [
     "math::MathDialect",
     "memref::MemRefDialect",
     "tensor::TensorDialect",
   ];
   let useDefaultAttributePrinterParser = 1;
   let hasCanonicalizer = 1;
   let hasOperationAttrVerify = 1;
   let hasConstantMaterializer = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
   let extraClassDeclaration = [{
     /// Attribute name used to to memoize indexing maps for named ops.
     constexpr const static ::llvm::StringLiteral
         kMemoizedIndexingMapsAttrName = "linalg.memoized_indexing_maps";
 
     using RegionBuilderFunType = llvm::function_ref<
       void(ImplicitLocOpBuilder &b, Block &, ArrayRef<NamedAttribute>)>;
     RegionBuilderFunType getRegionBuilder(StringRef name) {

mlir/include/mlir/Dialect/Quant/QuantOpsBase.td

 
 include "mlir/IR/OpBase.td"
 
 def Quantization_Dialect : Dialect {
   let name = "quant";
   let cppNamespace = "::mlir::quant";
 
   let useDefaultTypePrinterParser = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 //===----------------------------------------------------------------------===//
 // Quantization type definitions
 //===----------------------------------------------------------------------===//
 
 class quant_TypedPrimitiveOrContainer<Type etype> :
     Type<Or<[etype.predicate,

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

 include "mlir/Interfaces/ParallelCombiningOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 
 def SCF_Dialect : Dialect {
   let name = "scf";
   let cppNamespace = "::mlir::scf";
   let dependentDialects = ["arith::ArithDialect"];
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 // Base class for SCF dialect ops.
 class SCF_Op<string mnemonic, list<Trait> traits = []> :
     Op<SCF_Dialect, mnemonic, traits>;
 
 //===----------------------------------------------------------------------===//
 // ConditionOp

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorBase.td

let description = [{
     Proceedings of the IEEE/ACM International Symposium on Code Generation
     and Optimization, 2019.
     * [Kjolstad20] Fredrik Berg Kjolstad. Sparse Tensor Algebra Compilation.
     PhD thesis, MIT, February, 2020.
   }];
 
   let useDefaultAttributePrinterParser = 1;
   let useDefaultTypePrinterParser = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 #endif // SPARSETENSOR_BASE

mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td

def Transform_Dialect : Dialect {
   let name = "transform";
   let cppNamespace = "::mlir::transform";
 
   let dependentDialects = [
     "::mlir::pdl::PDLDialect",
     "::mlir::pdl_interp::PDLInterpDialect",
   ];
 
+  let useFoldAPI = kEmitFoldAdaptorFolder;
+
   let extraClassDeclaration = [{
       /// Returns the named PDL constraint functions available in the dialect
       /// as a map from their name to the function.
       const ::llvm::StringMap<::mlir::PDLConstraintFunction> &
       getPDLConstraintHooks() const;
 
       /// Parses a type registered by this dialect or one of its extensions.
       ::mlir::Type parseType(::mlir::DialectAsmParser &parser) const override;

mlir/include/mlir/IR/BuiltinDialect.td

private:
     void registerAttributes();
     // Register the builtin Location Attributes.
     void registerLocationAttributes();
     // Register the builtin Types.
     void registerTypes();
 
   public:
   }];
+
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 #endif // BUILTIN_BASE

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 // The result of the affine apply operation can be used as a symbol in `region`
 // if all its operands are symbols in `region`.
 bool AffineApplyOp::isValidSymbol(Region *region) {
   return llvm::all_of(getOperands(), [&](Value operand) {
     return mlir::isValidSymbol(operand, region);
   });
 }
 
-OpFoldResult AffineApplyOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult AffineApplyOp::fold(FoldAdaptor adaptor) {
   auto map = getAffineMap();
 
   // Fold dims and symbols to existing values.
   auto expr = map.getResult(0);
   if (auto dim = expr.dyn_cast<AffineDimExpr>())
     return getOperand(dim.getPosition());
   if (auto sym = expr.dyn_cast<AffineSymbolExpr>())
     return getOperand(map.getNumDims() + sym.getPosition());
 
   // Otherwise, default to folding the map.
   SmallVector<Attribute, 1> result;
-  if (failed(map.constantFold(operands, result)))
+  if (failed(map.constantFold(adaptor.getMapOperands(), result)))
     return {};
   return result[0];
 }
 
 /// Returns the largest known divisor of `e`. Exploits information from the
 /// values in `operands`.
 static int64_t getLargestKnownDivisor(AffineExpr e, ArrayRef<Value> operands) {
   // This method isn't aware of `operands`.

 }
 
 /// Returns true if the affine.for has zero iterations in trivial cases.
 static bool hasTrivialZeroTripCount(AffineForOp op) {
   std::optional<uint64_t> tripCount = getTrivialConstantTripCount(op);
   return tripCount && *tripCount == 0;
 }
 
-LogicalResult AffineForOp::fold(ArrayRef<Attribute> operands,
+LogicalResult AffineForOp::fold(FoldAdaptor adaptor,
                                 SmallVectorImpl<OpFoldResult> &results) {
   bool folded = succeeded(foldLoopBounds(*this));
   folded |= succeeded(canonicalizeLoopBounds(*this));
   if (hasTrivialZeroTripCount(*this)) {
     // The initial values of the loop-carried variables (iter_args) are the
     // results of the op.
     results.assign(getIterOperands().begin(), getIterOperands().end());
     folded = true;
if (llvm::none_of(operands,
     return;
 
   composeAffineMapAndOperands(&map, &operands);
   set = IntegerSet::get(map.getNumDims(), map.getNumSymbols(), map.getResults(),
                         set.getEqFlags());
 }
 
 /// Canonicalize an affine if op's conditional (integer set + operands).
-LogicalResult AffineIfOp::fold(ArrayRef<Attribute>,
+LogicalResult AffineIfOp::fold(FoldAdaptor,
                                SmallVectorImpl<OpFoldResult> &) {
   auto set = getIntegerSet();
   SmallVector<Value, 4> operands(getOperands());
   composeSetAndOperands(set, operands);
   canonicalizeSetAndOperands(&set, &operands);
 
   // Check if the canonicalization or composition led to any change.
   if (getIntegerSet() == set && llvm::equal(operands, getOperands()))
LogicalResult AffineLoadOp::verify() {
   return success();
 }
 
 void AffineLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
   results.add<SimplifyAffineOp<AffineLoadOp>>(context);
 }
 
-OpFoldResult AffineLoadOp::fold(ArrayRef<Attribute> cstOperands) {
+OpFoldResult AffineLoadOp::fold(FoldAdaptor adaptor) {
   /// load(memrefcast) -> load
   if (succeeded(memref::foldMemRefCast(*this)))
     return getResult();
 
   // Fold load from a global constant memref.
   auto getGlobalOp = getMemref().getDefiningOp<memref::GetGlobalOp>();
   if (!getGlobalOp)
     return {};
LogicalResult AffineStoreOp::verify() {
   return success();
 }
 
 void AffineStoreOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
   results.add<SimplifyAffineOp<AffineStoreOp>>(context);
 }
 
-LogicalResult AffineStoreOp::fold(ArrayRef<Attribute> cstOperands,
+LogicalResult AffineStoreOp::fold(FoldAdaptor adaptor,
                                   SmallVectorImpl<OpFoldResult> &results) {
   /// store(memrefcast) -> store
   return memref::foldMemRefCast(*this, getValueToStore());
 }
 
 //===----------------------------------------------------------------------===//
 // AffineMinMaxOpBase
 //===----------------------------------------------------------------------===//

 
 //===----------------------------------------------------------------------===//
 // AffineMinOp
 //===----------------------------------------------------------------------===//
 //
 //   %0 = affine.min (d0) -> (1000, d0 + 512) (%i0)
 //
 
-OpFoldResult AffineMinOp::fold(ArrayRef<Attribute> operands) {
-  return foldMinMaxOp(*this, operands);
+OpFoldResult AffineMinOp::fold(FoldAdaptor adaptor) {
+  return foldMinMaxOp(*this, adaptor.getOperands());
 }
 
 void AffineMinOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                               MLIRContext *context) {
   patterns.add<CanonicalizeSingleResultAffineMinMaxOp<AffineMinOp>,
                DeduplicateAffineMinMaxExpressions<AffineMinOp>,
                MergeAffineMinMaxOp<AffineMinOp>, SimplifyAffineOp<AffineMinOp>,
                CanonicalizeAffineMinMaxOpExprAndTermOrder<AffineMinOp>>(

 
 //===----------------------------------------------------------------------===//
 // AffineMaxOp
 //===----------------------------------------------------------------------===//
 //
 //   %0 = affine.max (d0) -> (1000, d0 + 512) (%i0)
 //
 
-OpFoldResult AffineMaxOp::fold(ArrayRef<Attribute> operands) {
-  return foldMinMaxOp(*this, operands);
+OpFoldResult AffineMaxOp::fold(FoldAdaptor adaptor) {
+  return foldMinMaxOp(*this, adaptor.getOperands());
 }
 
 void AffineMaxOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                               MLIRContext *context) {
   patterns.add<CanonicalizeSingleResultAffineMinMaxOp<AffineMaxOp>,
                DeduplicateAffineMinMaxExpressions<AffineMaxOp>,
                MergeAffineMinMaxOp<AffineMaxOp>, SimplifyAffineOp<AffineMaxOp>,
                CanonicalizeAffineMinMaxOpExprAndTermOrder<AffineMaxOp>>(

 }
 
 void AffinePrefetchOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                    MLIRContext *context) {
   // prefetch(memrefcast) -> prefetch
   results.add<SimplifyAffineOp<AffinePrefetchOp>>(context);
 }
 
-LogicalResult AffinePrefetchOp::fold(ArrayRef<Attribute> cstOperands,
+LogicalResult AffinePrefetchOp::fold(FoldAdaptor adaptor,
                                      SmallVectorImpl<OpFoldResult> &results) {
   /// prefetch(memrefcast) -> prefetch
   return memref::foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // AffineParallelOp
 //===----------------------------------------------------------------------===//
static LogicalResult canonicalizeLoopBounds(AffineParallelOp op) {
   if (lbCanonicalized)
     op.setLowerBounds(lb.getOperands(), lb.getAffineMap());
   if (ubCanonicalized)
     op.setUpperBounds(ub.getOperands(), ub.getAffineMap());
 
   return success();
 }
 
-LogicalResult AffineParallelOp::fold(ArrayRef<Attribute> operands,
+LogicalResult AffineParallelOp::fold(FoldAdaptor adaptor,
                                      SmallVectorImpl<OpFoldResult> &results) {
   return canonicalizeLoopBounds(*this);
 }
 
 /// Prints a lower(upper) bound of an affine parallel loop with max(min)
 /// conditions in it. `mapAttr` is a flat list of affine expressions and `group`
 /// identifies which of the those expressions form max/min groups. `operands`
 /// are the SSA values of dimensions and symbols and `keyword` is either "min"

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

void CloneOp::getEffects(
   effects.emplace_back(MemoryEffects::Read::get(), getInput(),
                        SideEffects::DefaultResource::get());
   effects.emplace_back(MemoryEffects::Write::get(), getOutput(),
                        SideEffects::DefaultResource::get());
   effects.emplace_back(MemoryEffects::Allocate::get(), getOutput(),
                        SideEffects::DefaultResource::get());
 }
 
-OpFoldResult CloneOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult CloneOp::fold(FoldAdaptor adaptor) {
   return succeeded(memref::foldMemRefCast(*this)) ? getResult() : Value();
 }
 
 namespace {
 
 /// Merge the clone and its source (by converting the clone to a cast) when
 /// possible.
 struct SimplifyClones : public OpRewritePattern<CloneOp> {
LogicalResult DeallocTensorOp::bufferize(RewriterBase &rewriter,
   rewriter.eraseOp(getOperation());
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ToTensorOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult ToTensorOp::fold(ArrayRef<Attribute>) {
+OpFoldResult ToTensorOp::fold(FoldAdaptor) {
   if (auto toMemref = getMemref().getDefiningOp<ToMemrefOp>())
     // Approximate alias analysis by conservatively folding only when no there
     // is no interleaved operation.
     if (toMemref->getBlock() == this->getOperation()->getBlock() &&
         toMemref->getNextNode() == this->getOperation())
       return toMemref.getTensor();
   return {};
 }
void ToTensorOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
   results.add<DimOfToTensorFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ToMemrefOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult ToMemrefOp::fold(ArrayRef<Attribute>) {
+OpFoldResult ToMemrefOp::fold(FoldAdaptor) {
   if (auto memrefToTensor = getTensor().getDefiningOp<ToTensorOp>())
     if (memrefToTensor.getMemref().getType() == getType())
       return memrefToTensor.getMemref();
   return {};
 }
 
 namespace {
 

mlir/lib/Dialect/Complex/IR/ComplexOps.cpp

 
 using namespace mlir;
 using namespace mlir::complex;
 
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult ConstantOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.empty() && "constant has no operands");
+OpFoldResult ConstantOp::fold(FoldAdaptor adaptor) {
   return getValue();
 }
 
 void ConstantOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "cst");
 }
 
LogicalResult ConstantOp::verify() {
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // CreateOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult CreateOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2 && "binary op takes two operands");
+OpFoldResult CreateOp::fold(FoldAdaptor adaptor) {
   // Fold complex.create(complex.re(op), complex.im(op)).
   if (auto reOp = getOperand(0).getDefiningOp<ReOp>()) {
     if (auto imOp = getOperand(1).getDefiningOp<ImOp>()) {
       if (reOp.getOperand() == imOp.getOperand()) {
         return reOp.getOperand();
       }
     }
   }
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // ImOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult ImOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 1 && "unary op takes 1 operand");
-  ArrayAttr arrayAttr = operands[0].dyn_cast_or_null<ArrayAttr>();
+OpFoldResult ImOp::fold(FoldAdaptor adaptor) {
+  ArrayAttr arrayAttr = adaptor.getComplex().dyn_cast_or_null<ArrayAttr>();
   if (arrayAttr && arrayAttr.size() == 2)
     return arrayAttr[1];
   if (auto createOp = getOperand().getDefiningOp<CreateOp>())
     return createOp.getOperand(1);
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // ReOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult ReOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 1 && "unary op takes 1 operand");
-  ArrayAttr arrayAttr = operands[0].dyn_cast_or_null<ArrayAttr>();
+OpFoldResult ReOp::fold(FoldAdaptor adaptor) {
+  ArrayAttr arrayAttr = adaptor.getComplex().dyn_cast_or_null<ArrayAttr>();
   if (arrayAttr && arrayAttr.size() == 2)
     return arrayAttr[0];
   if (auto createOp = getOperand().getDefiningOp<CreateOp>())
     return createOp.getOperand(0);
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // AddOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult AddOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2 && "binary op takes 2 operands");
-
+OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
   // complex.add(complex.sub(a, b), b) -> a
   if (auto sub = getLhs().getDefiningOp<SubOp>())
     if (getRhs() == sub.getRhs())
       return sub.getLhs();
 
   // complex.add(b, complex.sub(a, b)) -> a
   if (auto sub = getRhs().getDefiningOp<SubOp>())
     if (getLhs() == sub.getRhs())
OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // SubOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult SubOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2 && "binary op takes 2 operands");
-
+OpFoldResult SubOp::fold(FoldAdaptor adaptor) {
   // complex.sub(complex.add(a, b), b) -> a
   if (auto add = getLhs().getDefiningOp<AddOp>())
     if (getRhs() == add.getRhs())
       return add.getLhs();
 
   // complex.sub(a, complex.constant<0.0, 0.0>) -> a
   if (auto constantOp = getRhs().getDefiningOp<ConstantOp>()) {
     auto arrayAttr = constantOp.getValue();
     if (arrayAttr[0].cast<FloatAttr>().getValue().isZero() &&
         arrayAttr[1].cast<FloatAttr>().getValue().isZero()) {
       return getLhs();
     }
   }
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // NegOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult NegOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 1 && "unary op takes 1 operand");
-
+OpFoldResult NegOp::fold(FoldAdaptor adaptor) {
   // complex.neg(complex.neg(a)) -> a
   if (auto negOp = getOperand().getDefiningOp<NegOp>())
     return negOp.getOperand();
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // LogOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult LogOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 1 && "unary op takes 1 operand");
-
+OpFoldResult LogOp::fold(FoldAdaptor adaptor) {
   // complex.log(complex.exp(a)) -> a
   if (auto expOp = getOperand().getDefiningOp<ExpOp>())
     return expOp.getOperand();
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // ExpOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult ExpOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 1 && "unary op takes 1 operand");
-
+OpFoldResult ExpOp::fold(FoldAdaptor adaptor) {
   // complex.exp(complex.log(a)) -> a
   if (auto logOp = getOperand().getDefiningOp<LogOp>())
     return logOp.getOperand();
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // ConjOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult ConjOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 1 && "unary op takes 1 operand");
-
+OpFoldResult ConjOp::fold(FoldAdaptor adaptor) {
   // complex.conj(complex.conj(a)) -> a
   if (auto conjOp = getOperand().getDefiningOp<ConjOp>())
     return conjOp.getOperand();
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Complex/IR/ComplexOps.cpp.inc"

mlir/lib/Dialect/EmitC/IR/EmitC.cpp

LogicalResult emitc::ConstantOp::verify() {
   TypedAttr value = getValueAttr();
   Type type = getType();
   if (!value.getType().isa<NoneType>() && type != value.getType())
     return emitOpError() << "requires attribute's type (" << value.getType()
                          << ") to match op's return type (" << type << ")";
   return success();
 }
 
-OpFoldResult emitc::ConstantOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.empty() && "constant has no operands");
+OpFoldResult emitc::ConstantOp::fold(FoldAdaptor adaptor) {
   return getValue();
 }
 
 //===----------------------------------------------------------------------===//
 // IncludeOp
 //===----------------------------------------------------------------------===//
 
 void IncludeOp::print(OpAsmPrinter &p) {

mlir/lib/Dialect/Func/IR/FuncOps.cpp

LogicalResult ConstantOp::verify() {
 
   // Check that the referenced function has the correct type.
   if (fn.getFunctionType() != type)
     return emitOpError("reference to function with mismatched type");
 
   return success();
 }
 
-OpFoldResult ConstantOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.empty() && "constant has no operands");
+OpFoldResult ConstantOp::fold(FoldAdaptor adaptor) {
   return getValueAttr();
 }
 
 void ConstantOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "f");
 }
 

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

LogicalResult SubgroupMmaComputeOp::verify() {
 
   if (aShape[1] != bShape[0] || aShape[0] != cShape[0] ||
       bShape[1] != cShape[1])
     return emitError("operand shapes do not satisfy matmul constraints");
 
   return success();
 }
 
-LogicalResult MemcpyOp::fold(ArrayRef<Attribute> operands,
+LogicalResult MemcpyOp::fold(FoldAdaptor adaptor,
                              SmallVectorImpl<::mlir::OpFoldResult> &results) {
   return memref::foldMemRefCast(*this);
 }
 
-LogicalResult MemsetOp::fold(ArrayRef<Attribute> operands,
+LogicalResult MemsetOp::fold(FoldAdaptor adaptor,
                              SmallVectorImpl<::mlir::OpFoldResult> &results) {
   return memref::foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // GPU_WaitOp
 //===----------------------------------------------------------------------===//
 

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

for (int64_t idx : position) {
     if (auto structType = llvmType.dyn_cast<LLVMStructType>())
       llvmType = structType.getBody()[idx];
     else
       llvmType = llvmType.cast<LLVMArrayType>().getElementType();
   }
   return llvmType;
 }
 
-OpFoldResult LLVM::ExtractValueOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult LLVM::ExtractValueOp::fold(FoldAdaptor adaptor) {
   auto insertValueOp = getContainer().getDefiningOp<InsertValueOp>();
   OpFoldResult result = {};
   while (insertValueOp) {
     if (getPosition() == insertValueOp.getPosition())
       return insertValueOp.getValue();
     unsigned min =
         std::min(getPosition().size(), insertValueOp.getPosition().size());
     // If one is fully prefix of the other, stop propagating back as it will
LogicalResult LLVM::ConstantOp::verify() {
   }
   if (!getValue().isa<IntegerAttr, ArrayAttr, FloatAttr, ElementsAttr>())
     return emitOpError()
            << "only supports integer, float, string or elements attributes";
   return success();
 }
 
 // Constant op constant-folds to its value.
-OpFoldResult LLVM::ConstantOp::fold(ArrayRef<Attribute>) { return getValue(); }
+OpFoldResult LLVM::ConstantOp::fold(FoldAdaptor) { return getValue(); }
 
 //===----------------------------------------------------------------------===//
 // Utility functions for parsing atomic ops
 //===----------------------------------------------------------------------===//
 
 // Helper function to parse a keyword into the specified attribute named by
 // `attrName`. The keyword must match one of the string values defined by the
 // AtomicBinOp enum. The resulting I64 attribute is added to the `result`
return emitOpError("can be given only acquire, release, acq_rel, "
                        "and seq_cst orderings");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Folder for LLVM::BitcastOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult LLVM::BitcastOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult LLVM::BitcastOp::fold(FoldAdaptor adaptor) {
   // bitcast(x : T0, T0) -> x
   if (getArg().getType() == getType())
     return getArg();
   // bitcast(bitcast(x : T0, T1), T0) -> x
   if (auto prev = getArg().getDefiningOp<BitcastOp>())
     if (prev.getArg().getType() == getType())
       return prev.getArg();
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // Folder for LLVM::AddrSpaceCastOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult LLVM::AddrSpaceCastOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult LLVM::AddrSpaceCastOp::fold(FoldAdaptor adaptor) {
   // addrcast(x : T0, T0) -> x
   if (getArg().getType() == getType())
     return getArg();
   // addrcast(addrcast(x : T0, T1), T0) -> x
   if (auto prev = getArg().getDefiningOp<AddrSpaceCastOp>())
     if (prev.getArg().getType() == getType())
       return prev.getArg();
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // Folder for LLVM::GEPOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult LLVM::GEPOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult LLVM::GEPOp::fold(FoldAdaptor adaptor) {
   GEPIndicesAdaptor<ArrayRef<Attribute>> indices(getRawConstantIndicesAttr(),
-                                                 operands.drop_front());
+                                                 adaptor.getDynamicIndices());
 
   // gep %x:T, 0 -> %x
   if (getBase().getType() == getType() && indices.size() == 1)
     if (auto integer = indices[0].dyn_cast_or_null<IntegerAttr>())
       if (integer.getValue().isZero())
         return getBase();
 
   // Canonicalize any dynamic indices of constant value to constant indices.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 
 } // namespace
 
 void GenericOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
   results.add<EraseIdentityGenericOp>(context);
 }
 
-LogicalResult GenericOp::fold(ArrayRef<Attribute>,
-                              SmallVectorImpl<OpFoldResult> &) {
+LogicalResult GenericOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) {
   return memref::foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // MapOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseDstStyleOp(

mlir/lib/Dialect/Quant/IR/QuantOps.cpp

addTypes<AnyQuantizedType, CalibratedQuantizedType, UniformQuantizedType,
            UniformQuantizedPerAxisType>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Quant/QuantOps.cpp.inc"
       >();
   addBytecodeInterface(this);
 }
 
-OpFoldResult StorageCastOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult StorageCastOp::fold(FoldAdaptor adaptor) {
   // Matches x -> [scast -> scast] -> y, replacing the second scast with the
   // value of x if the casts invert each other.
   auto srcScastOp = getArg().getDefiningOp<StorageCastOp>();
   if (!srcScastOp || srcScastOp.getArg().getType() != getType())
     return OpFoldResult();
   return srcScastOp.getArg();
 }
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Quant/QuantOps.cpp.inc"

mlir/lib/Dialect/SCF/IR/SCF.cpp

if (elseRegion)
       regions.push_back(RegionSuccessor(elseRegion));
     return;
   }
 
   // Add the successor regions using the condition.
   regions.push_back(RegionSuccessor(condition ? &getThenRegion() : elseRegion));
 }
 
-LogicalResult IfOp::fold(ArrayRef<Attribute> operands,
+LogicalResult IfOp::fold(FoldAdaptor adaptor,
                          SmallVectorImpl<OpFoldResult> &results) {
   // if (!c) then A() else B() -> if c then B() else A()
   if (getElseRegion().empty())
     return failure();
 
   arith::XOrIOp xorStmt = getCondition().getDefiningOp<arith::XOrIOp>();
   if (!xorStmt)
     return failure();

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

if (auto tp2 = getDest().getType().dyn_cast<RankedTensorType>()) {
         if (shape1[d] != shape2[d] && shape2[d] != ShapedType::kDynamic)
           return emitError("unexpected conversion mismatch in dimension ") << d;
       return success();
     }
   }
   return emitError("unexpected type in convert");
 }
 
-OpFoldResult ConvertOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult ConvertOp::fold(FoldAdaptor adaptor) {
   Type dstType = getType();
   // Fold trivial dense-to-dense convert and leave trivial sparse-to-sparse
   // convert for codegen to remove. This is because we use trivial
   // sparse-to-sparse convert to tell bufferization that the sparse codegen
   // will expand the tensor buffer into sparse tensor storage.
   if (!getSparseTensorEncoding(dstType) && dstType == getSource().getType())
     return getSource();
   return {};
LogicalResult GetStorageSpecifierOp::verify() {
   return success();
 }
 
 template <typename SpecifierOp>
 static SetStorageSpecifierOp getSpecifierSetDef(SpecifierOp op) {
   return op.getSpecifier().template getDefiningOp<SetStorageSpecifierOp>();
 }
 
-OpFoldResult GetStorageSpecifierOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult GetStorageSpecifierOp::fold(FoldAdaptor adaptor) {
   StorageSpecifierKind kind = getSpecifierKind();
   std::optional<APInt> dim = getDim();
   for (auto op = getSpecifierSetDef(*this); op; op = getSpecifierSetDef(op))
     if (kind == op.getSpecifierKind() && dim == op.getDim())
       return op.getValue();
   return {};
 }
 

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

void transform::MergeHandlesOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getHandles(), effects);
   producesHandle(getResult(), effects);
 
   // There are no effects on the Payload IR as this is only a handle
   // manipulation.
 }
 
-OpFoldResult transform::MergeHandlesOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult transform::MergeHandlesOp::fold(FoldAdaptor adaptor) {
   if (getDeduplicate() || getHandles().size() != 1)
     return {};
 
   // If deduplication is not required and there is only one operand, it can be
   // used directly instead of merging.
   return getHandles().front();
 }
 

mlir/lib/IR/BuiltinDialect.cpp

LogicalResult ModuleOp::verify() {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // UnrealizedConversionCastOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
-UnrealizedConversionCastOp::fold(ArrayRef<Attribute> attrOperands,
+UnrealizedConversionCastOp::fold(FoldAdaptor adaptor,
                                  SmallVectorImpl<OpFoldResult> &foldResults) {
   OperandRange operands = getInputs();
   ResultRange results = getOutputs();
 
   if (operands.getType() == results.getType()) {
     foldResults.append(operands.begin(), operands.end());
     return success();
   }

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

 };
 } // namespace
 
 void TestOpWithRegionPattern::getCanonicalizationPatterns(
     RewritePatternSet &results, MLIRContext *context) {
   results.add<TestRemoveOpWithInnerOps>(context);
 }
 
-OpFoldResult TestOpWithRegionFold::fold(ArrayRef<Attribute> operands) {
+OpFoldResult TestOpWithRegionFold::fold(FoldAdaptor adaptor) {
   return getOperand();
 }
 
-OpFoldResult TestOpConstant::fold(ArrayRef<Attribute> operands) {
+OpFoldResult TestOpConstant::fold(FoldAdaptor adaptor) {
   return getValue();
 }
 
 LogicalResult TestOpWithVariadicResultsAndFolder::fold(
-    ArrayRef<Attribute> operands, SmallVectorImpl<OpFoldResult> &results) {
+    FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &results) {
   for (Value input : this->getOperands()) {
     results.push_back(input);
   }
   return success();
 }
 
-OpFoldResult TestOpInPlaceFold::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 1);
-  if (operands.front()) {
-    (*this)->setAttr("attr", operands.front());
+OpFoldResult TestOpInPlaceFold::fold(FoldAdaptor adaptor) {
+  if (adaptor.getOp()) {
+    (*this)->setAttr("attr", adaptor.getOp());
     return getResult();
   }
   return {};
 }
 
-OpFoldResult TestPassthroughFold::fold(ArrayRef<Attribute> operands) {
+OpFoldResult TestPassthroughFold::fold(FoldAdaptor adaptor) {
   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();
 

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

def Test_Dialect : Dialect {
   let hasConstantMaterializer = 1;
   let hasOperationAttrVerify = 1;
   let hasRegionArgAttrVerify = 1;
   let hasRegionResultAttrVerify = 1;
   let hasOperationInterfaceFallback = 1;
   let hasNonDefaultDestructor = 1;
   let useDefaultTypePrinterParser = 0;
   let useDefaultAttributePrinterParser = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
   let isExtensible = 1;
   let dependentDialects = ["::mlir::DLTIDialect"];
 
   let extraClassDeclaration = [{
     void registerAttributes();
     void registerTypes();
 
     // Provides a custom printing/parsing for some operations.

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

def TestOpInPlaceFold : TEST_Op<"op_in_place_fold"> {
   let hasFolder = 1;
 }
 
 // Test op that simply returns success.
 def TestOpInPlaceFoldSuccess : TEST_Op<"op_in_place_fold_success"> {
   let results = (outs Variadic<I1>);
   let hasFolder = 1;
   let extraClassDefinition = [{
-    ::mlir::LogicalResult $cppClass::fold(ArrayRef<Attribute> operands,
+    ::mlir::LogicalResult $cppClass::fold(FoldAdaptor adaptor,
         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);
-  }];
+  let hasFolder = 1;
 }
 
 // An op that always fold itself.
 def TestPassthroughFold : TEST_Op<"passthrough_fold"> {
   let arguments = (ins AnyType:$op);
   let results = (outs AnyType);
   let hasFolder = 1;
 }

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

 using namespace mlir;
 using namespace test;
 
 //===----------------------------------------------------------------------===//
 // Trait Folder.
 //===----------------------------------------------------------------------===//
 
 OpFoldResult TestInvolutionTraitFailingOperationFolderOp::fold(
-    ArrayRef<Attribute> operands) {
+    FoldAdaptor adaptor) {
   // This failure should cause the trait fold to run instead.
   return {};
 }
 
 OpFoldResult TestInvolutionTraitSuccesfulOperationFolderOp::fold(
-    ArrayRef<Attribute> operands) {
+    FoldAdaptor adaptor) {
   auto argumentOp = getOperand();
   // The success case should cause the trait fold to be supressed.
   return argumentOp.getDefiningOp() ? argumentOp : OpFoldResult{};
 }
 
 namespace {
 struct TestTraitFolder
     : public PassWrapper<TestTraitFolder, OperationPass<func::FuncOp>> {

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

ArrayAttr {0}::getIndexingMaps() {{
   return Builder(getContext()).getAffineMapArrayAttr(indexingMaps);
 }
 )FMT";
 
 // Implementations of fold and getEffects.
 // Parameters:
 // {0}: Class name
 const char structuredOpFoldersFormat[] = R"FMT(
-LogicalResult {0}::fold(ArrayRef<Attribute>,
+LogicalResult {0}::fold(FoldAdaptor,
                         SmallVectorImpl<OpFoldResult> &) {{
   return memref::foldMemRefCast(*this);
 }
 void {0}::getEffects(SmallVectorImpl<
     SideEffects::EffectInstance<MemoryEffects::Effect> >&effects) {{
       if (hasTensorSemantics()) return;
       getGenericEffectsImpl(effects,
         getOperation()->getResults(), getDpsInputOperands(), getDpsInitOperands());