Enable some cases of functor calls in custom pushforwards

Previously, if a user wanted to provide a custom pushforward for a
function that uses functors in it, it was impossible to use generated
pushforwards for that functors' call operators. This commit aims to fix
this for basic functors that don't have multiple call operator overloads.

Fixes: vgvassilev#1023

include/clad/Differentiator/BaseForwardModeVisitor.h

@@ -47,6 +47,8 @@ class BaseForwardModeVisitor
 
   virtual void ExecuteInsidePushforwardFunctionBlock();
 
+  virtual void DifferentiateCallOperatorIfFunctor(clang::QualType QT);
+
   static bool IsDifferentiableType(clang::QualType T);
 
   virtual StmtDiff

lib/Differentiator/BaseForwardModeVisitor.cpp

@@ -1029,6 +1029,8 @@ BaseForwardModeVisitor::VisitArraySubscriptExpr(const ArraySubscriptExpr* ASE) {
 }
 
 StmtDiff BaseForwardModeVisitor::VisitDeclRefExpr(const DeclRefExpr* DRE) {
+  DifferentiateCallOperatorIfFunctor(DRE->getType());
+
   DeclRefExpr* clonedDRE = nullptr;
   // Check if referenced Decl was "replaced" with another identifier inside
   // the derivative
@@ -1594,6 +1596,7 @@ StmtDiff BaseForwardModeVisitor::VisitDeclStmt(const DeclStmt* DS) {
   // If the DeclStmt is not empty, check the first declaration.
   if (declsBegin != declsEnd && isa<VarDecl>(*declsBegin)) {
     auto* VD = dyn_cast<VarDecl>(*declsBegin);
+    DifferentiateCallOperatorIfFunctor(VD->getType());
     // Check for non-differentiable types.
     QualType QT = VD->getType();
     if (QT->isPointerType())
@@ -2057,8 +2060,65 @@ StmtDiff BaseForwardModeVisitor::VisitBreakStmt(const BreakStmt* stmt) {
   return StmtDiff(Clone(stmt));
 }
 
+void BaseForwardModeVisitor::DifferentiateCallOperatorIfFunctor(
+    clang::QualType QT) {
+  // Identify if the constructed type is a functor. For functors, we need to
+  // differentiate their call operator once an object has been constructed, to
+  // allow user calls to pushforwards inside user-provided custom derivatives.
+  // FIXME: A much more scalable solution would be to create pushforwards once
+  // they're called from user-provided custom derivatives. This could then be
+  // applied to other operators besides operator() to avoid compilation errors
+  // in such cases.
+  if (auto* RD = QT->getAsCXXRecordDecl()) {
+    CXXRecordDecl* constructedType = RD->getDefinition();
+    bool isFunctor = constructedType && !constructedType->isLambda();
+    std::vector<const CXXMethodDecl*> callMethods;
+    if (isFunctor) {
+      for (const auto* method : constructedType->methods()) {
+        if (const auto* cxxMethod = dyn_cast<CXXMethodDecl>(method)) {
+          if (cxxMethod->isOverloadedOperator() &&
+              cxxMethod->getOverloadedOperator() == OO_Call) {
+            callMethods.push_back(cxxMethod);
+          }
+        }
+      }
+      isFunctor = isFunctor && !callMethods.empty();
+    }
+
+    if (isFunctor) {
+      for (const auto* FD : callMethods) {
+        CXXScopeSpec SS;
+        bool hasCustomDerivative =
+            !m_Builder
+                 .LookupCustomDerivativeOrNumericalDiff(
+                     clad::utils::ComputeEffectiveFnName(FD) +
+                         GetPushForwardFunctionSuffix(),
+                     const_cast<DeclContext*>(FD->getDeclContext()), SS)
+                 .empty();
+
+        if (!hasCustomDerivative) {
+          // Request Clad to diff it.
+          DiffRequest pushforwardFnRequest;
+          pushforwardFnRequest.Function = FD;
+          pushforwardFnRequest.Mode = GetPushForwardMode();
+          pushforwardFnRequest.BaseFunctionName =
+              utils::ComputeEffectiveFnName(FD);
+          // Silence diag outputs in nested derivation process.
+          pushforwardFnRequest.VerboseDiags = false;
+
+          // Check if request already derived in DerivedFunctions.
+          m_Builder.HandleNestedDiffRequest(pushforwardFnRequest);
+        }
+      }
+    }
+  }
+}
+
 StmtDiff
 BaseForwardModeVisitor::VisitCXXConstructExpr(const CXXConstructExpr* CE) {
+  DifferentiateCallOperatorIfFunctor(CE->getType());
+
+  // Now continue differentiating the constructor itself:
   llvm::SmallVector<Expr*, 4> clonedArgs, derivedArgs;
   for (auto arg : CE->arguments()) {
     auto argDiff = Visit(arg);

test/ForwardMode/Functors.C

@@ -386,6 +386,36 @@ struct WidgetPointer {
   }
 };
 
+namespace clad {
+namespace custom_derivatives {
+  template <typename F>
+  void use_functor_pushforward(double x, F& f, double d_x, F& d_f) {
+      f.operator_call_pushforward(x, &d_f, d_x);
+  }
+}
+}
+template <typename F>
+void use_functor(double x, F& f) {
+    f(x);
+}
+
+struct Foo {
+    double &y;
+    Foo(double &y): y(y) {} 
+
+    double operator()(double x) {
+        y = 2*x;
+
+        return x;
+    }
+};
+
+double fn0(double x) {
+    Foo func = Foo{x};
+    use_functor(x, func);
+    return x;
+}
+
 #define INIT(E, ARG)\
 auto d_##E = clad::differentiate(&E, ARG);\
 auto d_##E##Ref = clad::differentiate(E, ARG);
@@ -504,4 +534,21 @@ int main() {
   TEST_2(W_Arr_5, 6, 5);          // CHECK-EXEC: 6.00 6.00
   TEST_2(W_Pointer_3, 6, 5);      // CHECK-EXEC: 37.00 37.00
   TEST_2(W_Pointer_5, 6, 5);      // CHECK-EXEC: 51.00 51.00
+
+  auto dfn0 = clad::differentiate(fn0, "x");
+  printf("RES: %f\n", dfn0.execute(3.0)); // CHECK-EXEC: RES: 2
 }
+
+// CHECK: clad::ValueAndPushforward<double, double> operator_call_pushforward(double x, Foo *_d_this, double _d_x);
+// CHECK: double fn0_darg0(double x) {
+// CHECK-NEXT:     double _d_x = 1;
+// CHECK-NEXT:    Foo _d_func = Foo{{[{]*_d_x[}]*}};
+// CHECK-NEXT:    Foo func = Foo{{[{]*x[}]*}};
+// CHECK-NEXT:    clad::custom_derivatives::use_functor_pushforward(x, func, _d_x, _d_func);
+// CHECK-NEXT:    return _d_x;
+// CHECK-NEXT:}
+// CHECK: clad::ValueAndPushforward<double, double> operator_call_pushforward(double x, Foo *_d_this, double _d_x) {
+// CHECK-NEXT:    _d_this->y = 0 * x + 2 * _d_x;
+// CHECK-NEXT:    this->y = 2 * x;
+// CHECK-NEXT:    return {x, _d_x};
+// CHECK-NEXT:}

test/ForwardMode/ReferenceArguments.C

@@ -4,6 +4,31 @@
 
 #include "clad/Differentiator/Differentiator.h"
 
+namespace clad {
+namespace custom_derivatives {
+  template <typename F>
+  void use_functor_pushforward(double &x, F& f, double &d_x, F& d_f) {
+      f.operator_call_pushforward(x, &d_f, d_x);
+  }
+}
+}
+template <typename F>
+void use_functor(double &x, F& f) {
+    f(x);
+}
+
+struct Foo {
+    double operator()(double& x) {
+        x = 2*x*x;
+        return x;
+    }
+};
+
+double fn0(double x, Foo& func) {
+    use_functor(x, func);
+    return x;
+}
+
 double fn1(double& i, double& j) {
   double res = i * i * j;
   return res;
@@ -21,12 +46,14 @@ double fn1(double& i, double& j) {
 #define INIT(fn, ...) auto d_##fn = clad::differentiate(fn, __VA_ARGS__);
 
 #define TEST(fn, ...)                                                          \
-  auto res = d_##fn.execute(__VA_ARGS__);                                      \
-  printf("{%.2f}\n", res)
+  printf("{%.2f}\n", d_##fn.execute(__VA_ARGS__))
 
 int main() {
+    INIT(fn0, "x");
     INIT(fn1, "i");
 
     double i = 3, j = 5;
     TEST(fn1, i, j);    // CHECK-EXEC: {30.00}
+    Foo fff;
+    TEST(fn0, i, fff);    // CHECK-EXEC: {12.00}
 }

test/ForwardMode/UserDefinedTypes.C

@@ -420,6 +420,8 @@ Tensor<double, 5> fn5(double i, double j) {
   return T;
 } 
 
+// CHECK: void operator_call_pushforward(double val, Tensor<double, 5> *_d_this, double _d_val);
+
 // CHECK: Tensor<double, 5> fn5_darg0(double i, double j) {
 // CHECK-NEXT:     double _d_i = 1;
 // CHECK-NEXT:     double _d_j = 0;
@@ -593,8 +595,6 @@ TensorD5 fn11(double i, double j) {
   return res1;
 }
 
-// CHECK: void operator_call_pushforward(double val, Tensor<double, 5> *_d_this, double _d_val);
-
 // CHECK: clad::ValueAndPushforward<double &, double &> operator_subscript_pushforward(std::size_t idx, Tensor<double, 5> *_d_this, std::size_t _d_idx);
 
 // CHECK: clad::ValueAndPushforward<Tensor<double, 5U>, Tensor<double, 5U> > operator_plus_pushforward(const Tensor<double, 5U> &a, const Tensor<double, 5U> &b, const Tensor<double, 5U> &_d_a, const Tensor<double, 5U> &_d_b);
@@ -965,6 +965,16 @@ double fn18(double i, double j) {
 // CHECK-NEXT:     return _d_v[0].mem;
 // CHECK-NEXT: }
 
+// CHECK: void operator_call_pushforward(double val, Tensor<double, 5> *_d_this, double _d_val) {
+// CHECK-NEXT:     {
+// CHECK-NEXT:         unsigned int _d_i = 0;
+// CHECK-NEXT:         for (unsigned int i = 0; i < 5U; ++i) {
+// CHECK-NEXT:             _d_this->data[i] = _d_val;
+// CHECK-NEXT:             this->data[i] = val;
+// CHECK-NEXT:         }
+// CHECK-NEXT:     }
+// CHECK-NEXT: }
+
 template<unsigned N>
 void print(const Tensor<double, N>& t) {
   for (int i=0; i<N; ++i) {
@@ -1071,16 +1081,6 @@ int main() {
 // CHECK-NEXT:     return {{[{](__imag )?}}this->[[_M_value:[a-zA-Z_]+]],{{( __imag)?}} _d_this->[[_M_value:[a-zA-Z_]+]]};
 // CHECK-NEXT: }
 
-// CHECK: void operator_call_pushforward(double val, Tensor<double, 5> *_d_this, double _d_val) {
-// CHECK-NEXT:     {
-// CHECK-NEXT:         unsigned int _d_i = 0;
-// CHECK-NEXT:         for (unsigned int i = 0; i < 5U; ++i) {
-// CHECK-NEXT:             _d_this->data[i] = _d_val;
-// CHECK-NEXT:             this->data[i] = val;
-// CHECK-NEXT:         }
-// CHECK-NEXT:     }
-// CHECK-NEXT: }
-
 // CHECK: clad::ValueAndPushforward<double &, double &> operator_subscript_pushforward(std::size_t idx, Tensor<double, 5> *_d_this, std::size_t _d_idx) {
 // CHECK-NEXT:     return {(double &)this->data[idx], (double &)_d_this->data[idx]};
 // CHECK-NEXT: }

test/Functors/Simple.C

@@ -58,6 +58,50 @@ float f(float x) {
   return x;
 }
 
+namespace clad {
+namespace custom_derivatives {
+  template <typename F>
+  void use_functor_pushforward(double x, F& f, double d_x, F& d_f) {
+      f.operator_call_pushforward(x, &d_f, d_x);
+  }
+}
+}
+template <typename F>
+void use_functor(double x, F& f) {
+    f(x);
+}
+
+struct Foo {
+    double &y;
+    Foo(double &y): y(y) {} 
+
+    double operator()(double x) {
+        y = 2*x;
+
+        return x;
+    }
+};
+
+double fn0(double x) {
+    Foo func = Foo({x});
+    use_functor(x, func);
+    return x;
+}
+
+// CHECK: clad::ValueAndPushforward<double, double> operator_call_pushforward(double x, Foo *_d_this, double _d_x);
+// CHECK: double fn0_darg0(double x) {
+// CHECK-NEXT:     double _d_x = 1;
+// CHECK-NEXT:    Foo _d_func = Foo({_d_x});
+// CHECK-NEXT:    Foo func = Foo({x});
+// CHECK-NEXT:    clad::custom_derivatives::use_functor_pushforward(x, func, _d_x, _d_func);
+// CHECK-NEXT:    return _d_x;
+// CHECK-NEXT:}
+// CHECK: clad::ValueAndPushforward<double, double> operator_call_pushforward(double x, Foo *_d_this, double _d_x) {
+// CHECK-NEXT:    _d_this->y = 0 * x + 2 * _d_x;
+// CHECK-NEXT:    this->y = 2 * x;
+// CHECK-NEXT:    return {x, _d_x};
+// CHECK-NEXT:}
+
 int main() {
   AFunctor doubler;
   int x = doubler(5);
@@ -73,5 +117,8 @@ int main() {
   auto f1_darg1 = clad::differentiate(&SimpleExpression::operator(), 1);
   printf("Result is = %f\n", f1_darg1.execute(expr, 3.5, 4.5)); // CHECK-EXEC: Result is = 9
 
+  auto dfn0 = clad::differentiate(fn0, "x");
+  printf("RES: %f\n", dfn0.execute(3.0)); // CHECK-EXEC: RES: 2
+
   return 0;
 }