Add initial support for diff of ref return types in rev mode

include/clad/Differentiator/DerivativeBuilder.h

@@ -70,7 +70,8 @@ namespace clad {
     friend class ReverseModeVisitor;
     friend class HessianModeVisitor;
     friend class JacobianModeVisitor;
-
+    friend class ReverseModeForwPassVisitor;
+
     clang::Sema& m_Sema;
     plugin::CladPlugin& m_CladPlugin;
     clang::ASTContext& m_Context;

include/clad/Differentiator/DiffMode.h

@@ -8,6 +8,7 @@ namespace clad {
     experimental_pushforward,
     experimental_pullback,
     reverse,
+    reverse_mode_forward_pass,
     hessian,
     jacobian,
     error_estimation

include/clad/Differentiator/Differentiator.h

@@ -32,6 +32,12 @@ extern "C" {
 }
 
 namespace clad {
+  template<typename T, typename U>
+  struct ValueAndAdjoint {
+    T value;
+    U adjoint;
+  };
+
   /// \returns the size of a c-style string
   CUDA_HOST_DEVICE unsigned int GetLength(const char* code) {
     unsigned int count;

include/clad/Differentiator/ReverseModeForwPassVisitor.h

@@ -0,0 +1,39 @@
+#ifndef CLAD_TRANSFORM_SOURCE_FN_VISITOR_H
+#define CLAD_TRANSFORM_SOURCE_FN_VISITOR_H
+
+#include "clad/Differentiator/ParseDiffArgsTypes.h"
+#include "clad/Differentiator/ReverseModeVisitor.h"
+
+#include "clang/AST/StmtVisitor.h"
+#include "clang/Sema/Sema.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+namespace clad {
+class ReverseModeForwPassVisitor : public ReverseModeVisitor {
+private:
+  Stmts m_Globals;
+
+  llvm::SmallVector<clang::QualType, 8>
+  ComputeParamTypes(const DiffParams& diffParams);
+  clang::QualType ComputeReturnType();
+  llvm::SmallVector<clang::ParmVarDecl*, 8> BuildParams(DiffParams& diffParams);
+  clang::QualType GetParameterDerivativeType(clang::QualType yType,
+                                             clang::QualType xType);
+
+public:
+  ReverseModeForwPassVisitor(DerivativeBuilder& builder);
+  OverloadedDeclWithContext Derive(const clang::FunctionDecl* FD,
+                                   const DiffRequest& request);
+
+  StmtDiff ProcessSingleStmt(const clang::Stmt* S);
+
+  StmtDiff VisitStmt(const clang::Stmt* S) override;
+  StmtDiff VisitCompoundStmt(const clang::CompoundStmt* CS) override;
+  StmtDiff VisitDeclRefExpr(const clang::DeclRefExpr* DRE) override;
+  // StmtDiff VisitDeclStmt(const clang::DeclStmt* DS);
+  StmtDiff VisitReturnStmt(const clang::ReturnStmt* RS) override;
+};
+} // namespace clad
+
+#endif

include/clad/Differentiator/ReverseModeVisitor.h

@@ -31,7 +31,7 @@ namespace clad {
       : public clang::ConstStmtVisitor<ReverseModeVisitor, StmtDiff>,
         public VisitorBase {
 
-  private:
+  protected:
     // FIXME: We should remove friend-dependency of the plugin classes here.
     // For this we will need to separate out AST related functions in
     // a separate namespace, as well as add getters/setters function of
@@ -312,11 +312,11 @@ namespace clad {
     StmtDiff VisitArraySubscriptExpr(const clang::ArraySubscriptExpr* ASE);
     StmtDiff VisitBinaryOperator(const clang::BinaryOperator* BinOp);
     StmtDiff VisitCallExpr(const clang::CallExpr* CE);
-    StmtDiff VisitCompoundStmt(const clang::CompoundStmt* CS);
+    virtual StmtDiff VisitCompoundStmt(const clang::CompoundStmt* CS);
     StmtDiff VisitConditionalOperator(const clang::ConditionalOperator* CO);
     StmtDiff VisitCXXBoolLiteralExpr(const clang::CXXBoolLiteralExpr* BL);
     StmtDiff VisitCXXDefaultArgExpr(const clang::CXXDefaultArgExpr* DE);
-    StmtDiff VisitDeclRefExpr(const clang::DeclRefExpr* DRE);
+    virtual StmtDiff VisitDeclRefExpr(const clang::DeclRefExpr* DRE);
     StmtDiff VisitDeclStmt(const clang::DeclStmt* DS);
     StmtDiff VisitFloatingLiteral(const clang::FloatingLiteral* FL);
     StmtDiff VisitForStmt(const clang::ForStmt* FS);
@@ -326,8 +326,8 @@ namespace clad {
     StmtDiff VisitIntegerLiteral(const clang::IntegerLiteral* IL);
     StmtDiff VisitMemberExpr(const clang::MemberExpr* ME);
     StmtDiff VisitParenExpr(const clang::ParenExpr* PE);
-    StmtDiff VisitReturnStmt(const clang::ReturnStmt* RS);
-    StmtDiff VisitStmt(const clang::Stmt* S);
+    virtual StmtDiff VisitReturnStmt(const clang::ReturnStmt* RS);
+    virtual StmtDiff VisitStmt(const clang::Stmt* S);
     StmtDiff VisitUnaryOperator(const clang::UnaryOperator* UnOp);
     StmtDiff VisitExprWithCleanups(const clang::ExprWithCleanups* EWC);
     /// Decl is not Stmt, so it cannot be visited directly.

lib/Differentiator/CMakeLists.txt

@@ -37,6 +37,7 @@ add_llvm_library(cladDifferentiator
   HessianModeVisitor.cpp
   JacobianModeVisitor.cpp
   MultiplexExternalRMVSource.cpp
+  ReverseModeForwPassVisitor.cpp
   ReverseModeVisitor.cpp
   ErrorEstimator.cpp
   EstimationModel.cpp

lib/Differentiator/DerivativeBuilder.cpp

@@ -11,6 +11,7 @@
 #include "clad/Differentiator/HessianModeVisitor.h"
 #include "clad/Differentiator/JacobianModeVisitor.h"
 #include "clad/Differentiator/ReverseModeVisitor.h"
+#include "clad/Differentiator/ReverseModeForwPassVisitor.h"
 
 #include "clad/Differentiator/DiffPlanner.h"
 #include "clad/Differentiator/StmtClone.h"
@@ -172,7 +173,11 @@ namespace clad {
     } else if (request.Mode == DiffMode::experimental_pullback) {
       ReverseModeVisitor V(*this);
       result = V.DerivePullback(FD, request);
-    } else if (request.Mode == DiffMode::hessian) {
+    } else if (request.Mode == DiffMode::reverse_mode_forward_pass) {
+      ReverseModeForwPassVisitor V(*this);
+      result = V.Derive(FD, request);
+    }
+    else if (request.Mode == DiffMode::hessian) {
       HessianModeVisitor H(*this);
       result = H.Derive(FD, request);
     } else if (request.Mode == DiffMode::jacobian) {