Allow redefining builtin custom derivatives

Lookup now prioritizes custom_derivatives::overrides and falls back to builtin only when no overrides declaration exists. Overrides mismatch suppresses fallback to throw  user errors.

Closes: #1376.

Author: MdSaifAliMolla

lib/Differentiator/DiffPlanner.cpp

@@ -978,15 +978,18 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
     });
   }
 
-  static Expr* getOverloadExpr(Sema& S, DeclContext* DC, DiffRequest& R) {
+  static Expr* getOverloadExpr(Sema& S, DeclContext* DC, DiffRequest& R,
+                               bool* foundAnyDecl = nullptr) {
+    if (foundAnyDecl)
+      *foundAnyDecl = false;
     // Error estimation only uses forward mode derivatives if they are
     // user-prodived to handle builtin derivatives. If found, we have to change
     // the mode of the request.
     if (R.EnableErrorEstimation && R.Mode == DiffMode::pullback &&
         utils::canUsePushforwardInRevMode(R.Function)) {
       R.Mode = DiffMode::pushforward;
       R.EnableErrorEstimation = false;
-      if (Expr* overload = getOverloadExpr(S, DC, R)) {
+      if (Expr* overload = getOverloadExpr(S, DC, R, foundAnyDecl)) {
         R.DVI.clear();
         return overload;
       }
@@ -1026,6 +1029,9 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
     if (Found.empty())
       return nullptr; // Nothing found.
 
+    if (foundAnyDecl)
+      *foundAnyDecl = true;
+
     TemplateSpecCandidateSet FailedCandidates(R.CallContext->getBeginLoc(),
                                               /*ForTakingAddress=*/false);
     if (Expr* overload =
@@ -1078,21 +1084,43 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
       }
     } else
       fnDecl = request.Function;
-    DeclContext* DC = customDerNS;
-
-    if (isa<CXXMethodDecl>(fnDecl))
-      DC = utils::LookupNSD(m_Sema, "class_functions", /*shouldExist=*/false,
-                            DC);
-    else
-      DC = utils::FindDeclContext(m_Sema, DC, fnDecl->getDeclContext());
-
-    if (!DC)
-      return false;
-
     assert(request.Mode != DiffMode::unknown &&
            "Called lookup without specified DiffMode");
 
-    if (Expr* overload = getOverloadExpr(m_Sema, DC, request)) {
+    auto LookupOverload = [this, fnDecl, &request](NamespaceDecl* NSD,
+                                                   bool* foundAnyDecl) {
+      DeclContext* DC = NSD;
+      if (isa<CXXMethodDecl>(fnDecl))
+        DC = utils::LookupNSD(m_Sema, "class_functions",
+                              /*shouldExist=*/false, DC);
+      else
+        DC = utils::FindDeclContext(m_Sema, DC, fnDecl->getDeclContext());
+      if (!DC)
+        return static_cast<Expr*>(nullptr);
+      return getOverloadExpr(m_Sema, DC, request, foundAnyDecl);
+    };
+
+    // Look for the user overrides namespace nested inside custom_derivatives
+    NamespaceDecl* overridesNS = utils::LookupNSD(
+        m_Sema, "overrides", /*shouldExist=*/false, customDerNS);
+
+    // clad::custom_derivatives::overrides (user overrides, higher priority).
+    bool foundOverridesDecl = false;
+    if (overridesNS) {
+      if (Expr* overload = LookupOverload(overridesNS, &foundOverridesDecl)) {
+        // Overload found. Mark the request as custom derivative overrides and
+        // save the set of overloads to process later.
+        request.CustomDerivative = overload;
+        return true;
+      }
+      // Overrides declaration found but signature mismatch; do not fall back to
+      // builtin.
+      if (foundOverridesDecl)
+        return false;
+    }
+
+    // clad::custom_derivatives (builtins, fallback).
+    if (Expr* overload = LookupOverload(customDerNS, nullptr)) {
       // Overload found. Mark the request as custom derivative and save
       // the set of overloads to process later.
       request.CustomDerivative = overload;

test/Regressions/issue-1376.cpp

@@ -0,0 +1,33 @@
+// RUN: %cladclang -I%S/../../include %s -o %t 2>&1 | %filecheck %s
+// RUN: %t | %filecheck_exec %s
+
+#include "clad/Differentiator/Differentiator.h"
+
+extern "C" int printf(const char*, ...);
+
+namespace clad::custom_derivatives::overrides::std {
+template <typename T, typename dT>
+CUDA_HOST_DEVICE ValueAndPushforward<T, dT> exp_pushforward(T x, dT d_x) {
+  return {::std::exp(x), 2 * ::std::exp(x) * d_x};
+}
+}
+double differentiable_code(double x) {
+    return std::exp(x);
+}
+
+// CHECK: void differentiable_code_grad(double x, double *_d_x) {
+// CHECK-NEXT:     {
+// CHECK-NEXT:         double _r0 = 0.;
+// CHECK-NEXT:         _r0 += 1 * clad::custom_derivatives::overrides::std::exp_pushforward(x, 1.).pushforward;
+// CHECK-NEXT:         *_d_x += _r0;
+// CHECK-NEXT:     }
+// CHECK-NEXT: }
+
+int main() {
+  auto df = clad::gradient(differentiable_code);
+  df.dump();
+  double x = 1.0, d_x = 0;
+  df.execute(x, &d_x);
+  printf("d_x = %.5f\n", d_x);
+  // CHECK-EXEC: d_x = 5.43656
+}