[mlir][Tosa][NFC] Migrate Tosa dialect 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

[jpienaar]

Failure in https://buildkite.com/llvm-project/premerge-checks/builds/130112#0185a244-e641-456e-a1f4-31f7c6c1a8ef seems due formatting.

[zero9178]

clang-format

mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td

let description = [{
     there will be tools to lower from the ML frameworks into TOSA.
   }];
 
   let dependentDialects = ["tensor::TensorDialect"];
 
   let cppNamespace = "mlir::tosa";
   let hasConstantMaterializer = 1;
   let useDefaultAttributePrinterParser = 1;
+  let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Attributes.
 //===----------------------------------------------------------------------===//
 
 class Tosa_Attr<string attrName, string attrMnemonic, list<Trait> traits = []>
     : AttrDef<Tosa_Dialect, attrName, traits> {

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

if (lETy.isa<FloatType>()) {
       auto result = FloatFolder()(l, r);
       return DenseElementsAttr::get(returnTy, result);
     }
   }
 
   return {};
 }
 
-OpFoldResult AddOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
   auto lhsTy = getInput1().getType().dyn_cast<RankedTensorType>();
   auto rhsTy = getInput2().getType().dyn_cast<RankedTensorType>();
   auto resultTy = getType().dyn_cast<RankedTensorType>();
   if (!lhsTy || !rhsTy || !resultTy)
     return {};
   if (lhsTy != rhsTy)
     return {};
 
   auto resultETy = resultTy.getElementType();
-  auto lhsAttr = operands[0].dyn_cast_or_null<DenseElementsAttr>();
-  auto rhsAttr = operands[1].dyn_cast_or_null<DenseElementsAttr>();
+  auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
+  auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
 
   if (lhsAttr && lhsAttr.isSplat() && resultETy.isa<FloatType>()) {
     if (lhsAttr.getSplatValue<APFloat>().isZero())
       return getInput2();
   }
 
   if (rhsAttr && rhsAttr.isSplat() && resultETy.isa<FloatType>()) {
     if (rhsAttr.getSplatValue<APFloat>().isZero())
OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
 
   if (!lhsAttr || !rhsAttr)
     return {};
 
   return binaryFolder<std::plus<APInt>, std::plus<APFloat>>(lhsAttr, rhsAttr,
                                                             lhsTy);
 }
 
-OpFoldResult DivOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult DivOp::fold(FoldAdaptor adaptor) {
   auto lhsTy = getInput1().getType().dyn_cast<RankedTensorType>();
   auto rhsTy = getInput2().getType().dyn_cast<RankedTensorType>();
   auto resultTy = getType().dyn_cast<RankedTensorType>();
   if (!lhsTy || !rhsTy || !resultTy)
     return {};
   if (lhsTy != rhsTy)
     return {};
 
   auto resultETy = resultTy.getElementType();
-  auto lhsAttr = operands[0].dyn_cast_or_null<DenseElementsAttr>();
-  auto rhsAttr = operands[1].dyn_cast_or_null<DenseElementsAttr>();
+  auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
+  auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
   if (lhsAttr && lhsAttr.isSplat()) {
     if (resultETy.isa<IntegerType>() && lhsAttr.getSplatValue<APInt>().isZero())
       return lhsAttr;
   }
 
   if (rhsAttr && rhsAttr.isSplat()) {
     if (resultETy.isa<IntegerType>() && rhsAttr.getSplatValue<APInt>().isOne())
       return getInput1();
if (ty.getElementType().isa<FloatType>()) {
       return DenseElementsAttr::get(ty, result);
     }
   }
 
   return {};
 }
 } // namespace
 
-OpFoldResult MulOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
   auto lhs = getInput1();
   auto rhs = getInput2();
   auto lhsTy = lhs.getType().dyn_cast<RankedTensorType>();
   auto rhsTy = rhs.getType().dyn_cast<RankedTensorType>();
   auto resultTy = getType().dyn_cast<RankedTensorType>();
   if (!lhsTy || !rhsTy || !resultTy)
     return {};
   if (lhsTy != rhsTy)
     return {};
 
   auto resultETy = resultTy.getElementType();
-  auto lhsAttr = operands[0].dyn_cast_or_null<DenseElementsAttr>();
-  auto rhsAttr = operands[1].dyn_cast_or_null<DenseElementsAttr>();
+  auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
+  auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
 
   if (lhsAttr && lhsAttr.isSplat() && resultETy.isa<FloatType>()) {
     auto val = lhsAttr.getSplatValue<APFloat>();
     if (val.isZero())
       return lhsAttr;
     if (val.isExactlyValue(1.0))
       return rhs;
   }
if (val.isZero())
       return rhsAttr;
     if (val.getSExtValue() == shifted)
       return lhs;
   }
 
   return mulBinaryFolder(lhsAttr, rhsAttr, lhsTy, getShift());
 }
 
-OpFoldResult SubOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult SubOp::fold(FoldAdaptor adaptor) {
   auto lhsTy = getInput1().getType().dyn_cast<RankedTensorType>();
   auto rhsTy = getInput2().getType().dyn_cast<RankedTensorType>();
   auto resultTy = getType().dyn_cast<RankedTensorType>();
   if (!lhsTy || !rhsTy || !resultTy)
     return {};
   if (lhsTy != rhsTy)
     return {};
 
   auto resultETy = resultTy.getElementType();
-  auto lhsAttr = operands[0].dyn_cast_or_null<DenseElementsAttr>();
-  auto rhsAttr = operands[1].dyn_cast_or_null<DenseElementsAttr>();
+  auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
+  auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
 
   if (rhsAttr && rhsAttr.isSplat() && resultETy.isa<FloatType>()) {
     if (rhsAttr.getSplatValue<APFloat>().isZero())
       return getInput1();
   }
 
   if (rhsAttr && rhsAttr.isSplat() && resultETy.isa<IntegerType>()) {
     if (rhsAttr.getSplatValue<APInt>().isZero())

 struct APIntFoldGreaterEqual {
   APIntFoldGreaterEqual() = default;
   APInt operator()(const APInt &l, const APInt &r) {
     return APInt(1, l.sge(r));
   }
 };
 } // namespace
 
-OpFoldResult GreaterOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult GreaterOp::fold(FoldAdaptor adaptor) {
   auto resultTy = getType().dyn_cast<RankedTensorType>();
-  auto lhsAttr = operands[0].dyn_cast_or_null<DenseElementsAttr>();
-  auto rhsAttr = operands[1].dyn_cast_or_null<DenseElementsAttr>();
+  auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
+  auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
 
   if (!lhsAttr || !rhsAttr)
     return {};
 
   return binaryFolder<APIntFoldGreater, ComparisonFold<std::greater<APFloat>>>(
       lhsAttr, rhsAttr, resultTy);
 }
 
-OpFoldResult GreaterEqualOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult GreaterEqualOp::fold(FoldAdaptor adaptor) {
   auto resultTy = getType().dyn_cast<RankedTensorType>();
-  auto lhsAttr = operands[0].dyn_cast_or_null<DenseElementsAttr>();
-  auto rhsAttr = operands[1].dyn_cast_or_null<DenseElementsAttr>();
+  auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
+  auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
 
   if (!lhsAttr || !rhsAttr)
     return {};
 
   return binaryFolder<APIntFoldGreaterEqual,
                       ComparisonFold<std::greater_equal<APFloat>>>(
       lhsAttr, rhsAttr, resultTy);
 }
 
-OpFoldResult EqualOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult EqualOp::fold(FoldAdaptor adaptor) {
   auto resultTy = getType().dyn_cast<RankedTensorType>();
-  auto lhsAttr = operands[0].dyn_cast_or_null<DenseElementsAttr>();
-  auto rhsAttr = operands[1].dyn_cast_or_null<DenseElementsAttr>();
+  auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
+  auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
   Value lhs = getInput1();
   Value rhs = getInput2();
   auto lhsTy = lhs.getType().cast<ShapedType>();
 
   // If we are comparing an integer value to itself it is always true. We can
   // not do this with float due to float values.
   if (lhsTy.getElementType().isa<IntegerType>() && resultTy.hasStaticShape() &&
       lhs == rhs) {
     return DenseElementsAttr::get(resultTy, true);
   }
 
   if (!lhsAttr || !rhsAttr)
     return {};
 
   return binaryFolder<ComparisonFold<std::equal_to<APInt>>,
                       ComparisonFold<std::equal_to<APFloat>>>(lhsAttr, rhsAttr,
                                                               resultTy);
 }
 
-OpFoldResult CastOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult CastOp::fold(FoldAdaptor adaptor) {
   if (getInput().getType() == getType())
     return getInput();
 
-  auto operand = operands[0].dyn_cast_or_null<ElementsAttr>();
+  auto operand = adaptor.getInput().dyn_cast_or_null<ElementsAttr>();
   if (!operand)
     return {};
 
   auto inTy = getInput().getType().cast<ShapedType>();
   auto outTy = getType().cast<ShapedType>();
   auto inETy = inTy.getElementType();
   auto outETy = outTy.getElementType();
 
if (inETy.isa<IntegerType>() && outETy.isa<IntegerType>()) {
 
       return SplatElementsAttr::get(outTy, intVal);
     }
   }
 
   return {};
 }
 
-OpFoldResult ConstOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.empty() && "constant has no operands");
-  return getValueAttr();
-}
+OpFoldResult ConstOp::fold(FoldAdaptor adaptor) { return getValueAttr(); }
 
 #define REDUCE_FOLDER(OP)                                                      \
-  OpFoldResult OP::fold(ArrayRef<Attribute> operands) {                        \
+  OpFoldResult OP::fold(FoldAdaptor adaptor) {                                 \
     ShapedType inputTy = getInput().getType().cast<ShapedType>();              \
     if (!inputTy.hasRank())                                                    \
       return {};                                                               \
     if (inputTy.getDimSize(getAxis()) == 1)                                    \
       return getInput();                                                       \
     return {};                                                                 \
   }
 
 REDUCE_FOLDER(ReduceAllOp)
 REDUCE_FOLDER(ReduceAnyOp)
 REDUCE_FOLDER(ReduceMaxOp)
 REDUCE_FOLDER(ReduceMinOp)
 REDUCE_FOLDER(ReduceProdOp)
 REDUCE_FOLDER(ReduceSumOp)
 #undef REDUCE_FOLDER
 
-OpFoldResult ReshapeOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
   auto inputTy = getInput1().getType().dyn_cast<RankedTensorType>();
   auto outputTy = getType().dyn_cast<RankedTensorType>();
 
   if (!inputTy || !outputTy)
     return {};
 
   if (inputTy == outputTy)
     return getInput1();
 
-  auto operand = operands[0].dyn_cast_or_null<DenseElementsAttr>();
+  auto operand = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
   if (operand && outputTy.hasStaticShape() && operand.isSplat()) {
     return SplatElementsAttr::get(outputTy, operand.getSplatValue<Attribute>());
   }
 
   return {};
 }
 
-OpFoldResult PadOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult PadOp::fold(FoldAdaptor adaptor) {
   // If the pad is all zeros we can fold this operation away.
-  if (operands[1]) {
-    auto densePad = operands[1].cast<DenseElementsAttr>();
+  if (adaptor.getPadding()) {
+    auto densePad = adaptor.getPadding().cast<DenseElementsAttr>();
     if (densePad.isSplat() && densePad.getSplatValue<APInt>().isZero()) {
       return getInput1();
     }
   }
 
   return {};
 }
 
 // Fold away cases where a tosa.resize operation returns a copy
 // of the input image.
-OpFoldResult ResizeOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult ResizeOp::fold(FoldAdaptor adaptor) {
   ArrayRef<int64_t> offset = getOffset();
   ArrayRef<int64_t> border = getBorder();
   ArrayRef<int64_t> scale = getScale();
 
   // Check unit scaling.
   if (scale[0] != scale[1] || scale[2] != scale[3]) {
     return {};
   }
OpFoldResult ResizeOp::fold(FoldAdaptor adaptor) {
   auto inputTy = input.getType().cast<RankedTensorType>();
   auto resultTy = getType().cast<RankedTensorType>();
   if (inputTy != resultTy)
     return {};
 
   return input;
 }
 
-OpFoldResult ReverseOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult ReverseOp::fold(FoldAdaptor adaptor) {
   auto operand = getInput();
   auto operandTy = operand.getType().cast<ShapedType>();
   auto axis = getAxis();
-  auto operandAttr = operands[0].dyn_cast_or_null<SplatElementsAttr>();
+  auto operandAttr = adaptor.getInput().dyn_cast_or_null<SplatElementsAttr>();
   if (operandAttr)
     return operandAttr;
 
   // If the dim-length is 1, tosa.reverse is a no-op.
   if (operandTy.hasRank() && operandTy.getDimSize(axis) == 1)
     return operand;
 
   return {};
 }
 
-OpFoldResult SliceOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult SliceOp::fold(FoldAdaptor adaptor) {
   auto inputTy = getInput().getType().dyn_cast<RankedTensorType>();
   auto outputTy = getType().dyn_cast<RankedTensorType>();
 
   if (!inputTy || !outputTy)
     return {};
 
   if (inputTy == outputTy && inputTy.hasStaticShape())
     return getInput();
 
-  if (!operands[0])
+  if (!adaptor.getInput())
     return {};
 
-  auto operand = operands[0].cast<ElementsAttr>();
+  auto operand = adaptor.getInput().cast<ElementsAttr>();
   if (operand.isSplat() && outputTy.hasStaticShape()) {
     return SplatElementsAttr::get(outputTy, operand.getSplatValue<Attribute>());
   }
 
   if (inputTy.hasStaticShape() && outputTy.hasStaticShape() &&
       outputTy.getNumElements() == 1) {
     llvm::SmallVector<uint64_t> indices(getStart());
     auto value = operand.getValues<Attribute>()[indices];
     return SplatElementsAttr::get(outputTy, value);
   }
 
   return {};
 }
 
-OpFoldResult tosa::SelectOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult tosa::SelectOp::fold(FoldAdaptor adaptor) {
   if (getOnTrue() == getOnFalse())
     return getOnTrue();
 
-  auto predicate = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
+  auto predicate = adaptor.getPred().dyn_cast_or_null<DenseIntElementsAttr>();
   if (!predicate)
     return {};
 
   if (!predicate.isSplat())
     return {};
   return predicate.getSplatValue<APInt>().getBoolValue() ? getOnTrue()
                                                          : getOnFalse();
 }
 
-OpFoldResult TileOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult TileOp::fold(FoldAdaptor adaptor) {
   bool allOnes = llvm::all_of(getMultiples(), [](int64_t v) { return v == 1; });
   if (allOnes && getInput1().getType() == getType())
     return getInput1();
   return {};
 }
 
-OpFoldResult TransposeOp::fold(ArrayRef<Attribute> operands) {
+OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) {
   auto inputTy = getInput1().getType().cast<ShapedType>();
   auto resultTy = getType().cast<ShapedType>();
 
   // Transposing splat values just means reshaping.
-  if (auto input = operands[0].dyn_cast_or_null<DenseElementsAttr>()) {
+  if (auto input = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>()) {
     if (input.isSplat() && resultTy.hasStaticShape() &&
         inputTy.getElementType() == resultTy.getElementType())
       return input.reshape(resultTy);
   }
 
   // Transpose does not change the input type.
   if (getInput1().getType() != getType())
     return {};