Eliminates the reverse_forw pass when the custom pullback exists and it removes unwanted tape overhead (#1810)

Fixes issue #1808

Author: Vedant2005goyal

lib/Differentiator/DiffPlanner.cpp

@@ -1420,9 +1420,14 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
       forwPassRequest.CallContext = request.CallContext;
       forwPassRequest.UseRestoreTracker = shouldUseRestoreTracker;
       QualType returnType = request->getReturnType();
-      if (LookupCustomDerivativeDecl(forwPassRequest) ||
-          utils::isMemoryType(returnType) || shouldUseRestoreTracker)
+      bool hasCustomPullback = request.CustomDerivative != nullptr;
+      bool hasCustomReverseForw = LookupCustomDerivativeDecl(forwPassRequest);
+
+      if (hasCustomReverseForw ||
+          (!hasCustomPullback &&
+           (utils::isMemoryType(returnType) || shouldUseRestoreTracker))) {
         m_DiffRequestGraph.addNode(forwPassRequest, /*isSource=*/true);
+      }
     }
 
     if (!nonDiff && request.Mode != DiffMode::unknown)

test/Gradient/FunctionCalls.C

@@ -1098,16 +1098,6 @@ void inner_function(double *out, bool flag, const double *C) {
    if (flag)
       out[0] = C[0];
 }
-// CHECK: void inner_function_reverse_forw(double *out, bool flag, const double *C, double *_d_out, bool _d_flag, {{(const )?}}double *_d_C, clad::restore_tracker &_tracker0) {
-// CHECK-NEXT:     {
-// CHECK-NEXT:         bool _cond0 = flag;
-// CHECK-NEXT:         if (_cond0) {
-// CHECK-NEXT:             _tracker0.store(out[0]);
-// CHECK-NEXT:             out[0] = C[0];
-// CHECK-NEXT:         }
-// CHECK-NEXT:     }
-// CHECK-NEXT: }
-
 double fn31(double *variables) {
    double out = 0.;
    inner_function(&out, true, variables);

test/Gradient/STLCustomDerivatives.C

@@ -714,23 +714,23 @@ int main() {
 // CHECK-NEXT:              _t2++;
 // CHECK-NEXT:              res += v.at(i0);
 // CHECK-NEXT:          }
-// CHECK-NEXT:          clad::restore_tracker _tracker0 = {};
-// CHECK-NEXT:          v.assign_reverse_forw(3, 0, &_d_v, 0, 0, _tracker0);
-// CHECK-NEXT:          clad::restore_tracker _tracker1 = {};
-// CHECK-NEXT:          v.assign_reverse_forw(2, y, &_d_v, 0, *_d_y, _tracker1);
+// CHECK-NEXT:          std::vector<double> _t3 = v;
+// CHECK-NEXT:          v.assign(3, 0);
+// CHECK-NEXT:          std::vector<double> _t4 = v;
+// CHECK-NEXT:          v.assign(2, y);
 // CHECK-NEXT:          {
 // CHECK-NEXT:              _d_res += 1;
 // CHECK-NEXT:              _d_v[0] += 1;
 // CHECK-NEXT:              _d_v[1] += 1;
 // CHECK-NEXT:              _d_v[2] += 1;
 // CHECK-NEXT:          }
 // CHECK-NEXT:          {
-// CHECK-NEXT:              _tracker1.restore();
+// CHECK-NEXT:              v = _t4;
 // CHECK-NEXT:              {{.*size_type|size_t}} _r2 = {{0U|0UL|0}};
 // CHECK-NEXT:              {{.*}}assign_pullback(&v, 2, y, &_d_v, &_r2, _d_y);
 // CHECK-NEXT:          }
 // CHECK-NEXT:          {
-// CHECK-NEXT:              _tracker0.restore();
+// CHECK-NEXT:              v = _t3;
 // CHECK-NEXT:              {{.*size_type|size_t}} _r0 = {{0U|0UL|0}};
 // CHECK-NEXT:              {{.*}}value_type _r1 = 0.;
 // CHECK-NEXT:              {{.*}}assign_pullback(&v, 3, 0, &_d_v, &_r0, &_r1);