Add support for call expressions in vector forward mode AD

include/clad/Differentiator/BaseForwardModeVisitor.h

@@ -93,6 +93,11 @@ class BaseForwardModeVisitor
   VisitUnaryExprOrTypeTraitExpr(const clang::UnaryExprOrTypeTraitExpr* UE);
   StmtDiff VisitPseudoObjectExpr(const clang::PseudoObjectExpr* POE);
 
+  virtual clang::QualType
+  GetPushForwardDerivativeType(clang::QualType ParamType);
+  virtual std::string GetPushForwardFunctionSuffix();
+  virtual DiffMode GetPushForwardMode();
+
 protected:
   /// Helper function for differentiating the switch statement body.
   ///

include/clad/Differentiator/DerivativeBuilder.h

@@ -72,8 +72,9 @@ namespace clad {
   private:
     friend class VisitorBase;
     friend class BaseForwardModeVisitor;
-    friend class ForwardModeVisitor;
+    friend class PushForwardModeVisitor;
     friend class VectorForwardModeVisitor;
+    friend class VectorPushForwardModeVisitor;
     friend class ReverseModeVisitor;
     friend class HessianModeVisitor;
     friend class JacobianModeVisitor;

include/clad/Differentiator/DiffMode.h

@@ -8,6 +8,7 @@ enum class DiffMode {
   vector_forward_mode,
   experimental_pushforward,
   experimental_pullback,
+  experimental_vector_pushforward,
   reverse,
   hessian,
   jacobian,

include/clad/Differentiator/ForwardModeVisitor.h → include/clad/Differentiator/PushForwardModeVisitor.h

@@ -10,17 +10,19 @@
 #include "BaseForwardModeVisitor.h"
 
 namespace clad {
-  /// A visitor for processing the function code in forward mode.
-  /// Used to compute derivatives by clad::differentiate.
-class ForwardModeVisitor : public BaseForwardModeVisitor {
+/// A visitor for processing the function code in forward mode.
+/// Used to compute derivatives by clad::differentiate.
+class PushForwardModeVisitor : virtual public BaseForwardModeVisitor {
 
 public:
-  ForwardModeVisitor(DerivativeBuilder& builder);
-  ~ForwardModeVisitor();
+  PushForwardModeVisitor(DerivativeBuilder& builder);
+  ~PushForwardModeVisitor();
 
   DerivativeAndOverload DerivePushforward(const clang::FunctionDecl* FD,
                                           const DiffRequest& request);
 
+  virtual void ExecuteInsideFunctionBlock();
+
   /// Returns the return type for the pushforward function of the function
   /// `m_Function`.
   /// \note `m_Function` field should be set before using this function.

include/clad/Differentiator/VectorForwardModeVisitor.h

@@ -8,7 +8,7 @@
 namespace clad {
 /// A visitor for processing the function code in vector forward mode.
 /// Used to compute derivatives by clad::vector_forward_differentiate.
-class VectorForwardModeVisitor : public BaseForwardModeVisitor {
+class VectorForwardModeVisitor : virtual public BaseForwardModeVisitor {
 private:
   llvm::SmallVector<const clang::ValueDecl*, 16> m_IndependentVars;
   /// Map used to keep track of parameter variables w.r.t which the
@@ -78,6 +78,14 @@ class VectorForwardModeVisitor : public BaseForwardModeVisitor {
   StmtDiff VisitReturnStmt(const clang::ReturnStmt* RS) override;
   // Decl is not Stmt, so it cannot be visited directly.
   VarDeclDiff DifferentiateVarDecl(const clang::VarDecl* VD) override;
+
+  clang::QualType
+  GetPushForwardDerivativeType(clang::QualType ParamType) override;
+  std::string GetPushForwardFunctionSuffix() override;
+  DiffMode GetPushForwardMode() override;
+
+  // Function for setting the independent variables for vector mode.
+  void SetIndependentVarsExpr(clang::Expr* IndVarsExpr);
 };
 } // end namespace clad
 

include/clad/Differentiator/VectorPushForwardModeVisitor.h

@@ -0,0 +1,21 @@
+#ifndef CLAD_DIFFERENTIATOR_VECTORPUSHFORWARDMODEVISITOR_H
+#define CLAD_DIFFERENTIATOR_VECTORPUSHFORWARDMODEVISITOR_H
+
+#include "PushForwardModeVisitor.h"
+#include "VectorForwardModeVisitor.h"
+
+namespace clad {
+class VectorPushForwardModeVisitor : public PushForwardModeVisitor,
+                                     public VectorForwardModeVisitor {
+
+public:
+  VectorPushForwardModeVisitor(DerivativeBuilder& builder);
+  ~VectorPushForwardModeVisitor();
+
+  void ExecuteInsideFunctionBlock() override;
+
+  StmtDiff VisitReturnStmt(const clang::ReturnStmt* RS) override;
+};
+} // end namespace clad
+
+#endif // CLAD_DIFFERENTIATOR_VECTORPUSHFORWARDMODEVISITOR_H

lib/Differentiator/BaseForwardModeVisitor.cpp

@@ -914,6 +914,19 @@ Expr* DerivativeBuilder::BuildCallToCustomDerivativeOrNumericalDiff(
   return OverloadedFn;
 }
 
+QualType
+BaseForwardModeVisitor::GetPushForwardDerivativeType(QualType ParamType) {
+  return ParamType;
+}
+
+std::string BaseForwardModeVisitor::GetPushForwardFunctionSuffix() {
+  return "_pushforward";
+}
+
+DiffMode BaseForwardModeVisitor::GetPushForwardMode() {
+  return DiffMode::experimental_pushforward;
+}
+
 StmtDiff BaseForwardModeVisitor::VisitCallExpr(const CallExpr* CE) {
   const FunctionDecl* FD = CE->getDirectCallee();
   if (!FD) {
@@ -1008,10 +1021,6 @@ StmtDiff BaseForwardModeVisitor::VisitCallExpr(const CallExpr* CE) {
     CallArgs.push_back(argDiff.getExpr());
     if (BaseForwardModeVisitor::IsDifferentiableType(arg->getType())) {
       Expr* dArg = argDiff.getExpr_dx();
-      QualType CallArgTy = CallArgs.back()->getType();
-      assert((!dArg || m_Context.hasSameType(CallArgTy, dArg->getType())) &&
-             "Type mismatch, we might fail to instantiate a pullback");
-      (void)CallArgTy;
       // FIXME: What happens when dArg is nullptr?
       diffArgs.push_back(dArg);
     }
@@ -1034,14 +1043,13 @@ StmtDiff BaseForwardModeVisitor::VisitCallExpr(const CallExpr* CE) {
 
   // Try to find a user-defined overloaded derivative.
   std::string customPushforward =
-      clad::utils::ComputeEffectiveFnName(FD) + "_pushforward";
+      clad::utils::ComputeEffectiveFnName(FD) + GetPushForwardFunctionSuffix();
   Expr* callDiff = m_Builder.BuildCallToCustomDerivativeOrNumericalDiff(
       customPushforward, customDerivativeArgs, getCurrentScope(),
       const_cast<DeclContext*>(FD->getDeclContext()));
 
   // Check if it is a recursive call.
-  if (!callDiff && (FD == m_Function) &&
-      m_Mode == DiffMode::experimental_pushforward) {
+  if (!callDiff && (FD == m_Function) && m_Mode == GetPushForwardMode()) {
     // The differentiated function is called recursively.
     Expr* derivativeRef =
         m_Sema
@@ -1060,7 +1068,7 @@ StmtDiff BaseForwardModeVisitor::VisitCallExpr(const CallExpr* CE) {
     // derive the called function.
     DiffRequest pushforwardFnRequest;
     pushforwardFnRequest.Function = FD;
-    pushforwardFnRequest.Mode = DiffMode::experimental_pushforward;
+    pushforwardFnRequest.Mode = GetPushForwardMode();
     pushforwardFnRequest.BaseFunctionName = FD->getNameAsString();
     // pushforwardFnRequest.RequestedDerivativeOrder = m_DerivativeOrder;
     // Silence diag outputs in nested derivation process.

lib/Differentiator/CMakeLists.txt

@@ -27,14 +27,15 @@ add_llvm_library(cladDifferentiator
   DiffPlanner.cpp
   ErrorEstimator.cpp
   EstimationModel.cpp
-  ForwardModeVisitor.cpp
   HessianModeVisitor.cpp
   JacobianModeVisitor.cpp
   MultiplexExternalRMVSource.cpp
+  PushForwardModeVisitor.cpp
   ReverseModeForwPassVisitor.cpp
   ReverseModeVisitor.cpp
   StmtClone.cpp
   VectorForwardModeVisitor.cpp
+  VectorPushForwardModeVisitor.cpp
   Version.cpp
   VisitorBase.cpp
   ${version_inc}

lib/Differentiator/DerivativeBuilder.cpp

@@ -18,13 +18,14 @@
 #include "clad/Differentiator/CladUtils.h"
 #include "clad/Differentiator/DiffPlanner.h"
 #include "clad/Differentiator/ErrorEstimator.h"
-#include "clad/Differentiator/ForwardModeVisitor.h"
 #include "clad/Differentiator/HessianModeVisitor.h"
 #include "clad/Differentiator/JacobianModeVisitor.h"
+#include "clad/Differentiator/PushForwardModeVisitor.h"
 #include "clad/Differentiator/ReverseModeForwPassVisitor.h"
 #include "clad/Differentiator/ReverseModeVisitor.h"
 #include "clad/Differentiator/StmtClone.h"
 #include "clad/Differentiator/VectorForwardModeVisitor.h"
+#include "clad/Differentiator/VectorPushForwardModeVisitor.h"
 
 #include <algorithm>
 
@@ -209,11 +210,14 @@ namespace clad {
       BaseForwardModeVisitor V(*this);
       result = V.Derive(FD, request);
     } else if (request.Mode == DiffMode::experimental_pushforward) {
-      ForwardModeVisitor V(*this);
+      PushForwardModeVisitor V(*this);
       result = V.DerivePushforward(FD, request);
     } else if (request.Mode == DiffMode::vector_forward_mode) {
       VectorForwardModeVisitor V(*this);
       result = V.DeriveVectorMode(FD, request);
+    } else if (request.Mode == DiffMode::experimental_vector_pushforward) {
+      VectorPushForwardModeVisitor V(*this);
+      result = V.DerivePushforward(FD, request);
     } else if (request.Mode == DiffMode::reverse) {
       ReverseModeVisitor V(*this);
       result = V.Derive(FD, request);