[mlir] Partially revert removal of old `fold` method

[zero9178]

Summary

Mehdi noted in https://reviews.llvm.org/D144391 that given the low cost of keeping the old fold method signature working and the difficulty of writing a FoldAdaptor oneself, it'd be nice to keep the support for the sake of Ops written manually in C++.
This patch therefore partially reverts the removal of the old fold method by still allowing the old signature to be used. The active use of it is still discouraged and ODS will always generate the new method using FoldAdaptors.

I'd also like to note that the previous ought to have broken some manually defined fold methods in-tree that are defined here: https://github.com/llvm/llvm-project/blob/23bcd6b86271f1c219a69183a5d90654faca64b8/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h#L245 It seems like these are not part of the regressions tests however...

(also feel free to discuss whether this should be done or not in this review, I don't mind either way)

[mehdi_amini]

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

76

Before we do that, what would it take to write a FoldAdaptor here? Maybe my concern is misplaced if this isn't a lot of boilerplate.

[mehdi_amini]

I'd also like to note that the previous ought to have broken some manually defined fold methods in-tree that are defined here: https://github.com/llvm/llvm-project/blob/23bcd6b86271f1c219a69183a5d90654faca64b8/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h#L245 It seems like these are not part of the regressions tests however...

Oh so we don't even have a build failure right now... ouch!

[zero9178]

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

76

The minimal hacky useable version would be

cpp
struct FoldAdaptorHack
{
    ArrayRef<Attributes> operands;

    FoldAdaptorHack(ArrayRef<Attributes> operands,DictionaryAttr, ArrayRef<Region>) : operands(operands) {}
}

The only requirement the implementation currently enforces is that this constructor call is valid. All other methods of FoldAdaptors within fold are actually only consumed by the users implementation of fold. The con of the above approach is that it does not meet user expectations and documentation of what adaptors are and how they work. One could of course add getters, but then you'd have to write the above boilerplate for each C++ ops.

But arguably, if one writes the Op in C++ one is seemingly fine with some boilerplate.

[mehdi_amini]

Right now I think we should land this: the way the API break is not a build break but a subtle runtime issue (fold not being called) can be costly in terms of debugging. That adds to the discussion of how to write a FoldAdaptor and so I don't feel we need to weigh further on the adaptor itself.

mlir/include/mlir/IR/OpDefinition.h

static void attachInterface(MLIRContext &context) {
     (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_t =
-      decltype(std::declval<T>().fold(std::declval<typename T::FoldAdaptor>()));
+      decltype(std::declval<T>().fold(std::declval<ArrayRef<Attribute>>()));
   template <typename 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_t = decltype(std::declval<T>().fold(
-      std::declval<typename T::FoldAdaptor>(),
+      std::declval<ArrayRef<Attribute>>(),
       std::declval<SmallVectorImpl<OpFoldResult> &>()));
   template <typename 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>;
 
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 &&
-                  has_single_result_fold_v<ConcreteType>)
+                  (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 (has_fold_v<ConcreteType>)
+    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(typename ConcreteOpT::FoldAdaptor(
-            operands, op->getAttrDictionary(), op->getRegions()));
+    OpFoldResult result;
+    if constexpr (has_fold_adaptor_single_result_v<ConcreteOpT>)
+      result = cast<ConcreteOpT>(op).fold(typename ConcreteOpT::FoldAdaptor(
+    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(
+    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/test/IR/test-manual-cpp-fold.mlir

+// RUN: mlir-opt %s -canonicalize -split-input-file | FileCheck %s
+
+func.func @test() -> i32 {
+  %c5 = "test.constant"() {value = 5 : i32} : () -> i32
+  %res = "test.manual_cpp_op_with_fold"(%c5) : (i32) -> i32
+  return %res : i32
+}
+
+// CHECK-LABEL: func.func @test
+// CHECK-NEXT: %[[C:.*]] = "test.constant"() {value = 5 : i32}
+// CHECK-NEXT: return %[[C]]

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

 
 #include "TestOpInterfaces.h.inc"
 #include "TestOpsDialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.h.inc"
 
 namespace test {
+
+// Op deliberately defined in C++ code rather than ODS to test that C++
+// Ops can still use the old `fold` method.
+class ManualCppOpWithFold
+    : public mlir::Op<ManualCppOpWithFold, mlir::OpTrait::OneResult> {
+public:
+  using Op::Op;
+
+  static llvm::StringRef getOperationName() {
+    return "test.manual_cpp_op_with_fold";
+  }
+
+  static llvm::ArrayRef<llvm::StringRef> getAttributeNames() { return {}; }
+
+  mlir::OpFoldResult fold(llvm::ArrayRef<mlir::Attribute> attributes);
+};
+
 void registerTestDialect(::mlir::DialectRegistry &registry);
 void populateTestReductionPatterns(::mlir::RewritePatternSet &patterns);
 } // namespace test
 
 #endif // MLIR_TESTDIALECT_H

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

 
 void TestDialect::initialize() {
   registerAttributes();
   registerTypes();
   addOperations<
 #define GET_OP_LIST
 #include "TestOps.cpp.inc"
       >();
+  addOperations<ManualCppOpWithFold>();
   registerDynamicOp(getDynamicGenericOp(this));
   registerDynamicOp(getDynamicOneOperandTwoResultsOp(this));
   registerDynamicOp(getDynamicCustomParserPrinterOp(this));
 
   auto &blobInterface = addInterface<TestResourceBlobManagerInterface>();
   addInterface<TestOpAsmInterface>(blobInterface);
 
   addInterfaces<TestDialectFoldInterface, TestInlinerInterface,
void TestReflectBoundsOp::inferResultRanges(
   Builder b(ctx);
   setUminAttr(b.getIndexAttr(range.umin().getZExtValue()));
   setUmaxAttr(b.getIndexAttr(range.umax().getZExtValue()));
   setSminAttr(b.getIndexAttr(range.smin().getSExtValue()));
   setSmaxAttr(b.getIndexAttr(range.smax().getSExtValue()));
   setResultRanges(getResult(), range);
 }
 
+OpFoldResult ManualCppOpWithFold::fold(ArrayRef<Attribute> attributes) {
+  // Just a simple fold for testing purposes that reads an operands constant
+  // value and returns it.
+  if (!attributes.empty())
+    return attributes.front();
+  return nullptr;
+}
+
 #include "TestOpEnums.cpp.inc"
 #include "TestOpInterfaces.cpp.inc"
 #include "TestTypeInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"