Fix ref-return pullback lookup fallback (#1832)

Author: fogsong233

lib/Differentiator/DiffPlanner.cpp

@@ -1139,6 +1139,7 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
       return true;
 
     bool nonDiff = false;
+    bool hasNoMemoryInputForPointerOrRefReturn = false;
     // FIXME: We might want to support nested calls to differentiate/gradient
     // inside differentiated functions.
     if (!m_TopMostReq) {
@@ -1221,10 +1222,13 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
       if (!(MD && MD->isInstance()) && !hasPointerOrRefReturn &&
           allArgumentsAreLiterals(E->arguments(), m_ParentReq))
         nonDiff = true;
-      // In the reverse mode, such functions don't have dfdx()
-      if (!utils::hasMemoryTypeParams(FD) && hasPointerOrRefReturn &&
-          m_TopMostReq->Mode == DiffMode::reverse)
-        nonDiff = true;
+      // In reverse mode, calls without memory parameters normally have no
+      // adjoint destination for a pointer or reference return. Defer the final
+      // decision for instance calls until custom derivatives are known because
+      // their implicit object can carry the adjoint.
+      hasNoMemoryInputForPointerOrRefReturn =
+          !utils::hasMemoryTypeParams(FD) && hasPointerOrRefReturn &&
+          m_TopMostReq->Mode == DiffMode::reverse;
       // Skip reverse-mode scheduling for integral-return helper calls that
       // cannot accumulate through memory arguments. Keep this narrow to avoid
       // suppressing diagnostics on variadic/non-helper calls.
@@ -1371,7 +1375,36 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
         E->getDirectCallee();
     bool shouldUseRestoreTracker =
         utils::shouldUseRestoreTracker(request.Function);
-    if (!(LookupCustomDerivativeDecl(request) || nonDiff) || requestTBR) {
+    bool hasCustomPullback = LookupCustomDerivativeDecl(request);
+    // Share one request between early classification and final scheduling.
+    DiffRequest forwPassRequest;
+    if (request.Mode == DiffMode::pullback) {
+      forwPassRequest.Function = request.Function;
+      forwPassRequest.BaseFunctionName = request.BaseFunctionName;
+      forwPassRequest.Mode = DiffMode::reverse_mode_forward_pass;
+      forwPassRequest.CallContext = request.CallContext;
+      forwPassRequest.UseRestoreTracker = shouldUseRestoreTracker;
+    }
+
+    if (hasNoMemoryInputForPointerOrRefReturn) {
+      const auto* calledMethod = dyn_cast<CXXMethodDecl>(FD);
+      bool isRefReturningInstanceCall =
+          request.Mode == DiffMode::pullback && calledMethod &&
+          calledMethod->isInstance() &&
+          utils::isNonConstReferenceType(request->getReturnType());
+      bool hasCustomReverseForw = false;
+      if (isRefReturningInstanceCall)
+        hasCustomReverseForw = LookupCustomDerivativeDecl(forwPassRequest);
+
+      // A reverse_forw still needs a matching pullback. Preserve the existing
+      // forward-only contract only when a custom reverse_forw has no custom
+      // pullback, as with smart-pointer and reference-wrapper helpers.
+      if (!isRefReturningInstanceCall ||
+          (hasCustomReverseForw && !hasCustomPullback))
+        nonDiff = true;
+    }
+
+    if (!(hasCustomPullback || nonDiff) || requestTBR) {
       clang::CFG::BuildOptions Options;
       std::unique_ptr<AnalysisDeclContext> AnalysisDC =
           std::make_unique<AnalysisDeclContext>(
@@ -1445,14 +1478,10 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
     }
 
     if (request.Mode == DiffMode::pullback) {
-      DiffRequest forwPassRequest;
-      forwPassRequest.Function = request.Function;
-      forwPassRequest.BaseFunctionName = request.BaseFunctionName;
-      forwPassRequest.Mode = DiffMode::reverse_mode_forward_pass;
-      forwPassRequest.CallContext = request.CallContext;
+      // TBR can prove that no state needs restoring, so refresh this before
+      // scheduling the request.
       forwPassRequest.UseRestoreTracker = shouldUseRestoreTracker;
       QualType returnType = request->getReturnType();
-      bool hasCustomPullback = request.CustomDerivative != nullptr;
       bool hasCustomReverseForw = LookupCustomDerivativeDecl(forwPassRequest);
 
       if (hasCustomReverseForw ||

lib/Differentiator/ReverseModeVisitor.cpp

@@ -2129,7 +2129,17 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     // derivatives even if there is no `dfdx()` and thus we should call the
     // derived function. In the case of member functions, `implicit`
     // this object is always passed by reference.
-    if (!nonDiff && !dfdx() && !utils::hasMemoryTypeParams(FD))
+    bool isRefReturningInstanceCall =
+        MD && MD->isInstance() &&
+        utils::isNonConstReferenceType(FD->getReturnType());
+    bool needsPullbackForRefReturningInstance =
+        isRefReturningInstanceCall && !elideReverseForw;
+    // A non-elidable reverse_forw for a reference-returning member function
+    // leaves the adjoint on the returned reference. Keep the call
+    // differentiable so its statically scheduled pullback can propagate it
+    // through the implicit object.
+    if (!nonDiff && !dfdx() && !utils::hasMemoryTypeParams(FD) &&
+        !needsPullbackForRefReturningInstance)
       nonDiff = true;
 
     // If all arguments are constant literals, then this does not contribute to
@@ -2188,7 +2198,12 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
         }
       }
       pullbackFD = FindDerivedFunction(pullbackRequest);
-      if (pullbackFD && utils::hasEmptyBody(pullbackFD))
+      if (pullbackFD && utils::hasEmptyBody(pullbackFD) &&
+          !needsPullbackForRefReturningInstance)
+        nonDiff = true;
+      // Custom elidable reverse_forw helpers propagate their adjoint directly
+      // and intentionally have no pullback, for example smart pointers.
+      if (!pullbackFD && isRefReturningInstanceCall && elideReverseForw)
         nonDiff = true;
     }
 

test/Gradient/UserDefinedTypes.C

@@ -961,7 +961,10 @@ double fn26(double x, double y) {
 // CHECK-NEXT:      ptrClass p(&x);
 // CHECK-NEXT:      ptrClass _d_p(_d_x);
 // CHECK-NEXT:      clad::ValueAndAdjoint<double &, double &> _t0 = p.operator_star_reverse_forw(&_d_p);
+// CHECK-NEXT:      {
 // CHECK-NEXT:      _t0.adjoint += 1;
+// CHECK-NEXT:      p.operator_star_pullback(&_d_p);
+// CHECK-NEXT:      }
 // CHECK-NEXT:  }
 
 struct MyStructWrapper {
@@ -1264,6 +1267,7 @@ double fn34(double x, double y) {
 // CHECK-NEXT:    clad::ValueAndAdjoint<double &, double &> _t1 = obj_x.conversion_operator_reverse_forw(clad::Tag<double &>(), &_d_obj_x);
 // CHECK-NEXT:    {
 // CHECK-NEXT:        _t1.adjoint += 1;
+// CHECK-NEXT:        obj_x.conversion_operator_pullback(&_d_obj_x);
 // CHECK-NEXT:        _d_conv.val += 1;
 // CHECK-NEXT:    }
 // CHECK-NEXT:    obj_y.conversion_operator_pullback(_d_conv, &_d_obj_y);

test/Regressions/issue-1827.cpp

@@ -0,0 +1,84 @@
+// RUN: %cladclang %s -I%S/../../include -o %t 2>&1 | %filecheck %s
+// RUN: %t | %filecheck_exec %s
+
+#include "clad/Differentiator/Differentiator.h"
+
+#include <cstdio>
+
+struct Ref {
+  double* p;
+  Ref(double& x) : p(&x) {}
+  operator double&() const { return *p; }
+};
+
+struct IndexedRef {
+  double* p;
+  double& at(int i) const { return p[i]; }
+};
+
+int customPullbackCalls = 0;
+
+namespace clad {
+namespace custom_derivatives {
+namespace class_functions {
+clad::ValueAndAdjoint<double&, double&>
+at_reverse_forw(const ::IndexedRef* x, int i, const ::IndexedRef* d_x,
+                int /*d_i*/) {
+  return {x->at(i), d_x->at(i)};
+}
+
+void at_pullback(const ::IndexedRef* /*x*/, int /*i*/, ::IndexedRef* /*d_x*/,
+                 int* /*d_i*/) {
+  ++customPullbackCalls;
+}
+} // namespace class_functions
+} // namespace custom_derivatives
+} // namespace clad
+
+template <class T> double product(const T* x) {
+  return x[0] * x[1] * x[2] * x[3];
+}
+
+double custom_ref_case(double* params) {
+  Ref x[]{params[3], params[2], params[1], params[0]};
+  return product(x);
+}
+
+double custom_derivative_case(double* params) {
+  IndexedRef x{params};
+  return x.at(1);
+}
+
+int main() {
+  double params[]{1, 2, 3, 4};
+  double dParams[4]{};
+  auto grad = clad::gradient(custom_ref_case);
+  grad.execute(params, dParams);
+  std::printf("{%.0f, %.0f, %.0f, %.0f}\n", dParams[0], dParams[1],
+              dParams[2], dParams[3]);
+
+  double dCustomParams[4]{};
+  auto customGrad = clad::gradient(custom_derivative_case);
+  customGrad.execute(params, dCustomParams);
+  std::printf("{%.0f, %.0f, %.0f, %.0f}\n", dCustomParams[0],
+              dCustomParams[1], dCustomParams[2], dCustomParams[3]);
+  std::printf("%d\n", customPullbackCalls);
+}
+
+// CHECK: void product_pullback(const Ref *x, double _d_y, Ref *_d_x) {
+// CHECK:     clad::ValueAndAdjoint<double &, double &> _t3 = x[0].conversion_operator_reverse_forw(clad::Tag<double &>(), &_d_x[0]);
+// CHECK:     clad::ValueAndAdjoint<double &, double &> _t4 = x[1].conversion_operator_reverse_forw(clad::Tag<double &>(), &_d_x[1]);
+// CHECK:     clad::ValueAndAdjoint<double &, double &> _t5 = x[2].conversion_operator_reverse_forw(clad::Tag<double &>(), &_d_x[2]);
+// CHECK:     clad::ValueAndAdjoint<double &, double &> _t6 = x[3].conversion_operator_reverse_forw(clad::Tag<double &>(), &_d_x[3]);
+// CHECK:         x[0].conversion_operator_pullback(&_d_x[0]);
+// CHECK:         x[1].conversion_operator_pullback(&_d_x[1]);
+// CHECK:         x[2].conversion_operator_pullback(&_d_x[2]);
+// CHECK:         x[3].conversion_operator_pullback(&_d_x[3]);
+
+// CHECK: void custom_derivative_case_grad(double *params, double *_d_params) {
+// CHECK:     clad::ValueAndAdjoint<double &, double &> _t0 = clad::custom_derivatives::class_functions::at_reverse_forw(&x, 1, &_d_x, 0);
+// CHECK:         clad::custom_derivatives::class_functions::at_pullback(&x, 1, &_d_x, &_r0);
+
+// CHECK-EXEC: {24, 12, 8, 6}
+// CHECK-EXEC-NEXT: {0, 1, 0, 0}
+// CHECK-EXEC-NEXT: 1