Implement forward-mode jacobians

include/clad/Differentiator/JacobianModeVisitor.h

@@ -0,0 +1,22 @@
+#ifndef CLAD_DIFFERENTIATOR_JACOBIANMODEVISITOR_H
+#define CLAD_DIFFERENTIATOR_JACOBIANMODEVISITOR_H
+
+#include "VectorPushForwardModeVisitor.h"
+
+namespace clad {
+class JacobianModeVisitor : public VectorPushForwardModeVisitor {
+
+public:
+  JacobianModeVisitor(DerivativeBuilder& builder, const DiffRequest& request);
+
+  DiffMode GetPushForwardMode() override;
+
+  std::string GetPushForwardFunctionSuffix() override;
+
+  DerivativeAndOverload DeriveJacobian();
+
+  StmtDiff VisitReturnStmt(const clang::ReturnStmt* RS) override;
+};
+} // end namespace clad
+
+#endif // CLAD_DIFFERENTIATOR_JACOBIANMODEVISITOR_H

lib/Differentiator/BaseForwardModeVisitor.cpp

@@ -378,6 +378,8 @@ DerivativeAndOverload BaseForwardModeVisitor::Derive() {
 clang::QualType BaseForwardModeVisitor::ComputePushforwardFnReturnType() {
   assert(m_DiffReq.Mode == GetPushForwardMode());
   QualType originalFnRT = m_DiffReq->getReturnType();
+  if (m_DiffReq.Mode == DiffMode::jacobian)
+    return GetPushForwardDerivativeType(originalFnRT);
   if (originalFnRT->isVoidType())
     return m_Context.VoidTy;
   TemplateDecl* valueAndPushforward =
@@ -1445,7 +1447,9 @@ BaseForwardModeVisitor::VisitBinaryOperator(const BinaryOperator* BinOp) {
       derivedR = BuildParens(derivedR);
     opDiff = BuildOp(opCode, derivedL, derivedR);
   } else if (BinOp->isAssignmentOp()) {
-    if (Ldiff.getExpr_dx()->isModifiableLvalue(m_Context) != Expr::MLV_Valid) {
+    if ((Ldiff.getExpr_dx()->isModifiableLvalue(m_Context) !=
+         Expr::MLV_Valid) &&
+        !isCladArrayType(Ldiff.getExpr_dx()->getType())) {
       diag(DiagnosticsEngine::Warning, BinOp->getEndLoc(),
            "derivative of an assignment attempts to assign to unassignable "
            "expr, assignment ignored");

lib/Differentiator/CMakeLists.txt

@@ -30,6 +30,7 @@ llvm_add_library(cladDifferentiator
   DiffPlanner.cpp
   ErrorEstimator.cpp
   EstimationModel.cpp
+  JacobianModeVisitor.cpp
   HessianModeVisitor.cpp
   MultiplexExternalRMVSource.cpp
   PushForwardModeVisitor.cpp

lib/Differentiator/DerivativeBuilder.cpp

@@ -19,6 +19,7 @@
 #include "clad/Differentiator/DiffPlanner.h"
 #include "clad/Differentiator/ErrorEstimator.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"
@@ -429,8 +430,8 @@ static void registerDerivative(FunctionDecl* derivedFD, Sema& semaRef) {
       HessianModeVisitor H(*this, request);
       result = H.Derive();
     } else if (request.Mode == DiffMode::jacobian) {
-      ReverseModeVisitor R(*this, request);
-      result = R.Derive();
+      JacobianModeVisitor J(*this, request);
+      result = J.DeriveJacobian();
     } else if (request.Mode == DiffMode::error_estimation) {
       ReverseModeVisitor R(*this, request);
       InitErrorEstimation(m_ErrorEstHandler, m_EstModel, *this, request);

lib/Differentiator/JacobianModeVisitor.cpp

@@ -0,0 +1,43 @@
+#include "clad/Differentiator/JacobianModeVisitor.h"
+
+#include "ConstantFolder.h"
+#include "clad/Differentiator/CladUtils.h"
+
+#include "llvm/Support/SaveAndRestore.h"
+
+using namespace clang;
+
+namespace clad {
+JacobianModeVisitor::JacobianModeVisitor(DerivativeBuilder& builder,
+                                         const DiffRequest& request)
+    : VectorPushForwardModeVisitor(builder, request) {}
+
+DiffMode JacobianModeVisitor::GetPushForwardMode() {
+  return DiffMode::jacobian;
+}
+
+std::string JacobianModeVisitor::GetPushForwardFunctionSuffix() {
+  return "_jac";
+}
+
+DerivativeAndOverload JacobianModeVisitor::DeriveJacobian() {
+  return DerivePushforward(m_DiffReq.Function, m_DiffReq);
+}
+
+StmtDiff JacobianModeVisitor::VisitReturnStmt(const clang::ReturnStmt* RS) {
+  // If there is no return value, we must not attempt to differentiate
+  if (!RS->getRetValue())
+    return nullptr;
+
+  StmtDiff retValDiff = Visit(RS->getRetValue());
+  // return StmtDiff(retValDiff.getExpr_dx());
+  // This can instantiate as part of the move or copy initialization and
+  // needs a fake source location.
+  SourceLocation fakeLoc = utils::GetValidSLoc(m_Sema);
+  Stmt* returnStmt =
+      m_Sema
+          .ActOnReturnStmt(fakeLoc, retValDiff.getExpr_dx(), getCurrentScope())
+          .get();
+  return StmtDiff(returnStmt);
+}
+} // end namespace clad