From 7b635121d54bfa684e67896ea08fad23dab939a4 Mon Sep 17 00:00:00 2001
From: Mihail Mihov <mihovmihailp+github@gmail.com>
Date: Fri, 5 Jul 2024 16:19:53 +0300
Subject: [PATCH] Add support for non-differentiable attribute in reverse mode

---
 lib/Differentiator/ReverseModeVisitor.cpp | 53 +++++++++++++++++------
 test/Gradient/NonDifferentiable.C         |  8 ++--
 2 files changed, 44 insertions(+), 17 deletions(-)

diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp
index 6394ee9dd..7cee0814e 100644
--- a/lib/Differentiator/ReverseModeVisitor.cpp
+++ b/lib/Differentiator/ReverseModeVisitor.cpp
@@ -1369,6 +1369,27 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
       return StmtDiff(Clone(CE));
     }
 
+    SourceLocation validLoc{CE->getBeginLoc()};
+
+    // If the function is non_differentiable, return zero derivative.
+    if (clad::utils::hasNonDifferentiableAttribute(CE)) {
+      // Calling the function without computing derivatives
+      llvm::SmallVector<Expr*, 4> ClonedArgs;
+      for (unsigned i = 0, e = CE->getNumArgs(); i < e; ++i)
+        ClonedArgs.push_back(Clone(CE->getArg(i)));
+
+      Expr* Call = m_Sema
+                       .ActOnCallExpr(getCurrentScope(), Clone(CE->getCallee()),
+                                      validLoc, ClonedArgs, validLoc)
+                       .get();
+      // Creating a zero derivative
+      auto* zero = ConstantFolder::synthesizeLiteral(m_Context.IntTy, m_Context,
+                                                     /*val=*/0);
+
+      // Returning the function call and zero derivative
+      return StmtDiff(Call, zero);
+    }
+
     auto NArgs = FD->getNumParams();
     // If the function has no args and is not a member function call then we
     // assume that it is not related to independent variables and does not
@@ -2755,19 +2776,21 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
     if (declsBegin != DS->decls().end() && isa<VarDecl>(*declsBegin)) {
       auto* VD = dyn_cast<VarDecl>(*declsBegin);
       QualType QT = VD->getType();
-      if (!QT->isPointerType()) {
-        auto* typeDecl = QT->getAsCXXRecordDecl();
-        // We should also simply copy the original lambda. The differentiation
-        // of lambdas is happening in the `VisitCallExpr`. For now, only the
-        // declarations with lambda expressions without captures are supported.
-        isLambda = typeDecl && typeDecl->isLambda();
-        if (isLambda) {
-          for (auto* D : DS->decls())
-            if (auto* VD = dyn_cast<VarDecl>(D))
-              decls.push_back(VD);
-          Stmt* DSClone = BuildDeclStmt(decls);
-          return StmtDiff(DSClone, nullptr);
-        }
+      if (QT->isPointerType())
+        QT = QT->getPointeeType();
+
+      auto* typeDecl = QT->getAsCXXRecordDecl();
+      // We should also simply copy the original lambda. The differentiation
+      // of lambdas is happening in the `VisitCallExpr`. For now, only the
+      // declarations with lambda expressions without captures are supported.
+      isLambda = typeDecl && typeDecl->isLambda();
+      if (isLambda ||
+          (typeDecl && clad::utils::hasNonDifferentiableAttribute(typeDecl))) {
+        for (auto* D : DS->decls())
+          if (auto* VD = dyn_cast<VarDecl>(D))
+            decls.push_back(VD);
+        Stmt* DSClone = BuildDeclStmt(decls);
+        return StmtDiff(DSClone, nullptr);
       }
     }
 
@@ -2954,6 +2977,10 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
            "CXXMethodDecl nodes not supported yet!");
     MemberExpr* clonedME = utils::BuildMemberExpr(
         m_Sema, getCurrentScope(), baseDiff.getExpr(), field->getName());
+    auto* zero = ConstantFolder::synthesizeLiteral(m_Context.DoubleTy,
+                                                   m_Context, /*val=*/0);
+    if (clad::utils::hasNonDifferentiableAttribute(ME))
+      return {clonedME, zero};
     if (!baseDiff.getExpr_dx())
       return {clonedME, nullptr};
     MemberExpr* derivedME = utils::BuildMemberExpr(
diff --git a/test/Gradient/NonDifferentiable.C b/test/Gradient/NonDifferentiable.C
index b25d659d6..b0811c622 100644
--- a/test/Gradient/NonDifferentiable.C
+++ b/test/Gradient/NonDifferentiable.C
@@ -74,15 +74,15 @@ double fn_s2_operator(double i, double j) {
 
 #define TEST_CLASS(classname, name, i, j)                                      \
   auto d_##name = clad::gradient(&classname::name);                            \
-  double result_##name[2] = {};						       \
+  double result_##name[2];						       \
   d_##name.execute(expr_1, i, j, &result_##name[0], &result_##name[1]);		       \
-  printf("%.2f %.2f\n", result_##name[0], result_##name[1]);
+  printf("%.2f %.2f\n\n", result_##name[0], result_##name[1]);
 
 #define TEST_FUNC(name, i, j)                                                  \
   auto d_##name = clad::gradient(&name);	                               \
-  double result_##name[2] = {};						       \
+  double result_##name[2];						       \
   d_##name.execute(i, j, &result_##name[0], &result_##name[1]);			       \
-  printf("%.2f\n", result_##name[0], result_##name[1]);
+  printf("%.2f %.2f\n\n", result_##name[0], result_##name[1]);
 
 int main() {
   // FIXME: The parts of this test that are commented out are currently not working, due to bugs