Add support for SourceLocExpr during differentiation to enable code containing assert and other source-location builtins.

Author: Vedant2005goyal

include/clad/Differentiator/BaseForwardModeVisitor.h

@@ -99,6 +99,7 @@ class BaseForwardModeVisitor
   StmtDiff VisitWhileStmt(const clang::WhileStmt* WS);
   StmtDiff VisitDoStmt(const clang::DoStmt* DS);
   StmtDiff VisitContinueStmt(const clang::ContinueStmt* ContStmt);
+  StmtDiff VisitSourceLocExpr(const clang::SourceLocExpr* E);
 
   StmtDiff VisitSwitchStmt(const clang::SwitchStmt* SS);
   StmtDiff VisitBreakStmt(const clang::BreakStmt* BS);

include/clad/Differentiator/ReverseModeVisitor.h

@@ -485,6 +485,7 @@ namespace clad {
     StmtDiff VisitImplicitCastExpr(const clang::ImplicitCastExpr* ICE);
     StmtDiff VisitGNUNullExpr(const clang::GNUNullExpr* E);
     StmtDiff VisitPredefinedExpr(const clang::PredefinedExpr* E);
+    StmtDiff VisitSourceLocExpr(const clang::SourceLocExpr* SLE);
 
 #if CLANG_VERSION_MAJOR > 16
     StmtDiff VisitLambdaExpr(const clang::LambdaExpr* LE);

include/clad/Differentiator/StmtClone.h

@@ -79,6 +79,7 @@ namespace utils {
     DECLARE_CLONE_FN(CXXCatchStmt)
     DECLARE_CLONE_FN(CXXTryStmt)
     DECLARE_CLONE_FN(PredefinedExpr)
+    DECLARE_CLONE_FN(SourceLocExpr)
     DECLARE_CLONE_FN(CharacterLiteral)
     DECLARE_CLONE_FN(FloatingLiteral)
     DECLARE_CLONE_FN(ImaginaryLiteral)

lib/Differentiator/BaseForwardModeVisitor.cpp

@@ -581,7 +581,8 @@ StmtDiff BaseForwardModeVisitor::VisitConditionalOperator(
           .ActOnConditionalOp(noLoc, noLoc, cond, ifTrueDiff.getExpr(),
                               ifFalseDiff.getExpr())
           .get();
-
+  if (condExpr->getType()->isVoidType())
+    return StmtDiff(condExpr, nullptr);
   // cond is already used by the value conditional above; clone it for the
   // derivative conditional so the two do not share the stored condition.
   Expr* condExprDiff =
@@ -1262,7 +1263,7 @@ StmtDiff BaseForwardModeVisitor::VisitCallExpr(const CallExpr* CE) {
     if (!isa<CXXOperatorCallExpr>(CE) && !isa<CXXMemberCallExpr>(CE) &&
         !needsForwPass) {
       bool allArgsHaveZeroDerivatives = true;
-      for (unsigned i = 0, e = CE->getNumArgs(); i < e; ++i) {
+      for (unsigned i = 0, e = diffArgs.size(); i < e; ++i) {
         Expr* dArg = diffArgs[i];
         // If argDiff.expr_dx is nullptr or is a constant 0, then the derivative
         // of the function call is 0.
@@ -1321,7 +1322,7 @@ StmtDiff BaseForwardModeVisitor::VisitCallExpr(const CallExpr* CE) {
   // If clad failed to derive it, try finding its derivative using
   // numerical diff.
   if (!callDiff) {
-    Multiplier = diffArgs[0];
+    Multiplier = diffArgs.empty() ? nullptr : diffArgs[0];
     Expr* call =
         m_Sema
             .ActOnCallExpr(getCurrentScope(), Clone(CE->getCallee()), validLoc,
@@ -1533,7 +1534,7 @@ BaseForwardModeVisitor::VisitBinaryOperator(const BinaryOperator* BinOp) {
   } else if (opCode == BO_Comma) {
     // if expression is (E1, E2) then derivative is (E1', E1, E2')
     // because E1 may change some variables that E2 depends on.
-    if (!isUnusedResult(Ldiff.getExpr_dx())) {
+    if (Ldiff.getExpr_dx() && !isUnusedResult(Ldiff.getExpr_dx())) {
       opDiff = BuildOp(BO_Comma, BuildParens(Ldiff.getExpr_dx()),
                        BuildParens(Ldiff.getExpr()));
       opDiff = BuildOp(BO_Comma, BuildParens(opDiff),
@@ -1874,6 +1875,13 @@ StmtDiff BaseForwardModeVisitor::VisitStringLiteral(const StringLiteral* SL) {
                                  SL->getType(), utils::GetValidSLoc(m_Sema)));
 }
 
+StmtDiff
+BaseForwardModeVisitor::VisitSourceLocExpr(const clang::SourceLocExpr* E) {
+  auto* Constant0 =
+      ConstantFolder::synthesizeLiteral(m_Context.IntTy, m_Context, /*val=*/0);
+  return StmtDiff(CloneNode(E), Constant0);
+}
+
 StmtDiff BaseForwardModeVisitor::VisitWhileStmt(const WhileStmt* WS) {
   // Scope for the whole while loop.
   ScopeRAII whileScope(*this, Scope::ContinueScope | Scope::BreakScope |

lib/Differentiator/ReverseModeVisitor.cpp

@@ -2957,6 +2957,10 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
       for (Stmt* S : revBlock)
         addToCurrentBlock(S, direction::reverse);
       return {cloneE, derivedE};
+    } else if (opCode == UnaryOperatorKind::UO_Extension) {
+      diff = Visit(E);
+      ResultRef = diff.getExpr_dx();
+      valueForRevPass = diff.getRevSweepAsExpr();
     } else {
       if (opCode != UO_LNot)
         // We should only output warnings on visiting boolean conditions
@@ -3919,6 +3923,13 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
     return StmtDiff(CloneNode(E), Constant0);
   }
 
+  StmtDiff
+  ReverseModeVisitor::VisitSourceLocExpr(const clang::SourceLocExpr* E) {
+    auto* Constant0 = ConstantFolder::synthesizeLiteral(m_Context.IntTy,
+                                                        m_Context, /*val=*/0);
+    return StmtDiff(CloneNode(E), Constant0);
+  }
+
   StmtDiff ReverseModeVisitor::VisitCXXFunctionalCastExpr(
       const clang::CXXFunctionalCastExpr* FCE) {
     StmtDiff castExprDiff = Visit(FCE->getSubExpr(), dfdx());

lib/Differentiator/StmtClone.cpp

@@ -92,6 +92,16 @@ Stmt* StmtClone::VisitPredefinedExpr(PredefinedExpr* Node) {
   clad_compat::ExprSetDeps(result, Node);
   return result;
 }
+#if CLANG_VERSION_MAJOR >= 15
+DEFINE_CLONE_EXPR(SourceLocExpr,
+                  (Ctx, Node->getIdentKind(), CloneType(Node->getType()),
+                   Node->getBeginLoc(), Node->getEndLoc(),
+                   Node->getParentContext()))
+#else
+DEFINE_CLONE_EXPR(SourceLocExpr,
+                  (Ctx, Node->getIdentKind(), Node->getBeginLoc(),
+                   Node->getEndLoc(), Node->getParentContext()))
+#endif
 DEFINE_CLONE_EXPR(CharacterLiteral,
                   (Node->getValue(), Node->getKind(),
                    CloneType(Node->getType()), Node->getLocation()))

test/Features/DiffAssert.cpp

@@ -0,0 +1,52 @@
+// RUN: %cladclang -std=c++17 -I%S/../../include %s -o %t
+// RUN: %t | %filecheck_exec %s
+
+#include <cassert>
+#include "clad/Differentiator/Differentiator.h"
+
+double func(double x){
+    assert(x>0.0);
+    return x*x*x;
+}
+double func1(double x) {
+    const char* file = __builtin_FILE();
+    int line = __builtin_LINE();
+    return x*x*x;
+}
+
+void temp() {}
+double temp_func(double x) {
+    x > 0.0 ? temp() : temp();
+    return x * x * x;
+}
+
+int main(){
+    auto d_func=clad::gradient(func);
+    double dx=0;
+    d_func.execute(3.0,&dx);
+    printf("Diff result: %.2f\n", dx);
+    //CHECK-EXEC: Diff result: 27.00
+    
+    auto d_func_forw=clad::differentiate(func);
+    auto res_forw=d_func_forw.execute(3.0);
+    printf("Diff result: %.2f\n", res_forw);
+    // CHECK-EXEC: Diff result: 27.00
+
+    auto d_func1=clad::gradient(func1);
+    double dx1=0;
+    d_func1.execute(3.0,&dx1);
+    printf("Diff result: %.2f\n", dx1);
+    //CHECK-EXEC: Diff result: 27.00
+
+    auto d_func1_forw = clad::differentiate(func1);
+    auto res1_forw = d_func1_forw.execute(3.0);
+    printf("Diff result: %.2f\n", res1_forw);
+    //CHECK-EXEC: Diff result: 27.00
+
+    auto d_temp_func = clad::differentiate(temp_func);
+    auto res_void = d_temp_func.execute(3.0);
+    printf("Diff result: %.2f\n", res_void);
+    //CHECK-EXEC: Diff result: 27.00
+
+    return 0;
+}