[mlir][spirv] Migrate to new fold API

[kuhar]

Summary

Fixes: https://github.com/llvm/llvm-project/issues/59938

[zero9178]

Thank you for the help!

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td

let description = [{
     benefits transformations. The dialect also aims to maintain straightforward
     serialization into and deserialization from the SPIR-V binary format.
     See https://mlir.llvm.org/docs/Dialects/SPIR-V/ for more details regarding
     high-level designs and implementation structures of the SPIR-V dialect.
   }];
 
   let cppNamespace = "::mlir::spirv";
   let useDefaultTypePrinterParser = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
   let hasConstantMaterializer = 1;
   let hasOperationAttrVerify = 1;
   let hasRegionArgAttrVerify = 1;
   let hasRegionResultAttrVerify = 1;
 
   let extraClassDeclaration = [{
     void registerAttributes();
     void registerTypes();

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

void spirv::AccessChainOp::getCanonicalizationPatterns(
     RewritePatternSet &results, MLIRContext *context) {
   results.add<CombineChainedAccessChain>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.BitcastOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult spirv::BitcastOp::fold(ArrayRef<Attribute> /*operands*/) {
+OpFoldResult spirv::BitcastOp::fold(FoldAdaptor /*adaptor*/) {
   Value curInput = getOperand();
   if (getType() == curInput.getType())
     return curInput;
 
   // Look through nested bitcasts.
   if (auto prevCast = curInput.getDefiningOp<spirv::BitcastOp>()) {
     Value prevInput = prevCast.getOperand();
     if (prevInput.getType() == getType())
       return prevInput;
 
     getOperandMutable().assign(prevInput);
     return getResult();
   }
 
   // TODO(kuhar): Consider constant-folding the operand attribute.
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.CompositeExtractOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult spirv::CompositeExtractOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult spirv::CompositeExtractOp::fold(FoldAdaptor adaptor) {
   if (auto insertOp =
           getComposite().getDefiningOp<spirv::CompositeInsertOp>()) {
     if (getIndices() == insertOp.getIndices())
       return insertOp.getObject();
   }
 
   if (auto constructOp =
           getComposite().getDefiningOp<spirv::CompositeConstructOp>()) {
     auto type = constructOp.getType().cast<spirv::CompositeType>();
     if (getIndices().size() == 1 &&
         constructOp.getConstituents().size() == type.getNumElements()) {
       auto i = getIndices().begin()->cast<IntegerAttr>();
       return constructOp.getConstituents()[i.getValue().getSExtValue()];
     }
   }
 
   auto indexVector =
       llvm::to_vector<8>(llvm::map_range(getIndices(), [](Attribute attr) {
         return static_cast<unsigned>(attr.cast<IntegerAttr>().getInt());
       }));
-  return extractCompositeElement(operands[0], indexVector);
+  return extractCompositeElement(adaptor.getComposite(), indexVector);
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.Constant
 //===----------------------------------------------------------------------===//
 
-OpFoldResult spirv::ConstantOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.empty() && "spirv.Constant has no operands");
+OpFoldResult spirv::ConstantOp::fold(FoldAdaptor /*adaptor*/) {
   return getValue();
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.IAdd
 //===----------------------------------------------------------------------===//
 
-OpFoldResult spirv::IAddOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2 && "spirv.IAdd expects two operands");
+OpFoldResult spirv::IAddOp::fold(FoldAdaptor adaptor) {
   // x + 0 = x
   if (matchPattern(getOperand2(), m_Zero()))
     return getOperand1();
 
   // According to the SPIR-V spec:
   //
   // The resulting value will equal the low-order N bits of the correct result
   // R, where N is the component width and R is computed with enough precision
   // to avoid overflow and underflow.
   return constFoldBinaryOp<IntegerAttr>(
-      operands, [](APInt a, const APInt &b) { return std::move(a) + b; });
+      adaptor.getOperands(),
+      [](APInt a, const APInt &b) { return std::move(a) + b; });
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.IMul
 //===----------------------------------------------------------------------===//
 
-OpFoldResult spirv::IMulOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2 && "spirv.IMul expects two operands");
+OpFoldResult spirv::IMulOp::fold(FoldAdaptor adaptor) {
   // x * 0 == 0
   if (matchPattern(getOperand2(), m_Zero()))
     return getOperand2();
   // x * 1 = x
   if (matchPattern(getOperand2(), m_One()))
     return getOperand1();
 
   // According to the SPIR-V spec:
   //
   // The resulting value will equal the low-order N bits of the correct result
   // R, where N is the component width and R is computed with enough precision
   // to avoid overflow and underflow.
   return constFoldBinaryOp<IntegerAttr>(
-      operands, [](const APInt &a, const APInt &b) { return a * b; });
+      adaptor.getOperands(),
+      [](const APInt &a, const APInt &b) { return a * b; });
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.ISub
 //===----------------------------------------------------------------------===//
 
-OpFoldResult spirv::ISubOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult spirv::ISubOp::fold(FoldAdaptor adaptor) {
   // x - x = 0
   if (getOperand1() == getOperand2())
     return Builder(getContext()).getIntegerAttr(getType(), 0);
 
   // According to the SPIR-V spec:
   //
   // The resulting value will equal the low-order N bits of the correct result
   // R, where N is the component width and R is computed with enough precision
   // to avoid overflow and underflow.
   return constFoldBinaryOp<IntegerAttr>(
-      operands, [](APInt a, const APInt &b) { return std::move(a) - b; });
+      adaptor.getOperands(),
+      [](APInt a, const APInt &b) { return std::move(a) - b; });
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.LogicalAnd
 //===----------------------------------------------------------------------===//
 
-OpFoldResult spirv::LogicalAndOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2 && "spirv.LogicalAnd should take two operands");
-
-  if (Optional<bool> rhs = getScalarOrSplatBoolAttr(operands.back())) {
+OpFoldResult spirv::LogicalAndOp::fold(FoldAdaptor adaptor) {
+  if (Optional<bool> rhs = getScalarOrSplatBoolAttr(adaptor.getOperand2())) {
     // x && true = x
     if (*rhs)
       return getOperand1();
 
     // x && false = false
     if (!*rhs)
-      return operands.back();
+      return adaptor.getOperand2();
   }
 
   return Attribute();
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.LogicalNotEqualOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult spirv::LogicalNotEqualOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2 &&
-         "spirv.LogicalNotEqual should take two operands");
-
-  if (Optional<bool> rhs = getScalarOrSplatBoolAttr(operands.back())) {
+OpFoldResult spirv::LogicalNotEqualOp::fold(FoldAdaptor adaptor) {
+  if (Optional<bool> rhs = getScalarOrSplatBoolAttr(adaptor.getOperand2())) {
     // x && false = x
     if (!rhs.value())
       return getOperand1();
   }
 
   return Attribute();
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.LogicalNot
 //===----------------------------------------------------------------------===//
 
 void spirv::LogicalNotOp::getCanonicalizationPatterns(
     RewritePatternSet &results, MLIRContext *context) {
   results
       .add<ConvertLogicalNotOfIEqual, ConvertLogicalNotOfINotEqual,
            ConvertLogicalNotOfLogicalEqual, ConvertLogicalNotOfLogicalNotEqual>(
           context);
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.LogicalOr
 //===----------------------------------------------------------------------===//
 
-OpFoldResult spirv::LogicalOrOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2 && "spirv.LogicalOr should take two operands");
-
-  if (auto rhs = getScalarOrSplatBoolAttr(operands.back())) {
+OpFoldResult spirv::LogicalOrOp::fold(FoldAdaptor adaptor) {
+  if (auto rhs = getScalarOrSplatBoolAttr(adaptor.getOperand2())) {
     if (*rhs)
       // x || true = true
-      return operands.back();
+      return adaptor.getOperand2();
 
     // x || false = x
     if (!*rhs)
       return getOperand1();
   }
 
   return Attribute();
 }