Optimize pullback calls

include/clad/Differentiator/CladUtils.h

@@ -331,6 +331,12 @@ namespace clad {
 
     bool IsMemoryFunction(const clang::FunctionDecl* FD);
     bool IsMemoryDeallocationFunction(const clang::FunctionDecl* FD);
+
+    /// Checks if the argument arg involves type conversion
+    /// (e.g. float to double). This is used to find out
+    /// a similar cast can be performed for the derived
+    /// arg (for corresponding pointers).
+    bool isTriviallyCastableArg(const clang::Expr* arg);
     } // namespace utils
     } // namespace clad
 

include/clad/Differentiator/ErrorEstimator.h

@@ -119,10 +119,11 @@ class ErrorEstimationHandler : public ExternalRMVSource {
   /// \param[in] CallArgs The orignal call arguments of the function call.
   /// \param[in] ArgResultDecls The differentiated call arguments.
   /// \param[in] numArgs The number of call args.
-  void EmitNestedFunctionParamError(
-      clang::FunctionDecl* fnDecl,
-      llvm::SmallVectorImpl<clang::Expr*>& CallArgs,
-      llvm::SmallVectorImpl<clang::VarDecl*>& ArgResultDecls, size_t numArgs);
+  void
+  EmitNestedFunctionParamError(clang::FunctionDecl* fnDecl,
+                               llvm::SmallVectorImpl<clang::Expr*>& CallArgs,
+                               llvm::SmallVectorImpl<clang::Expr*>& ArgResult,
+                               size_t numArgs);
 
   /// Save values of registered variables so that they can be replaced
   /// properly in case of re-assignments.
@@ -265,8 +266,7 @@ class ErrorEstimationHandler : public ExternalRMVSource {
   void ActBeforeFinalizingVisitCallExpr(
       const clang::CallExpr*& CE, clang::Expr*& fnDecl,
       llvm::SmallVectorImpl<clang::Expr*>& derivedCallArgs,
-      llvm::SmallVectorImpl<clang::VarDecl*>& ArgResultDecls,
-      bool asGrad) override;
+      llvm::SmallVectorImpl<clang::Expr*>& ArgResult, bool asGrad) override;
   void
   ActAfterCloningLHSOfAssignOp(clang::Expr*& LCloned, clang::Expr*& R,
                                clang::BinaryOperator::Opcode& opCode) override;
@@ -275,8 +275,7 @@ class ErrorEstimationHandler : public ExternalRMVSource {
   void ActBeforeFinalizingDifferentiateSingleExpr(const direction& d) override;
   void ActBeforeDifferentiatingCallExpr(
       llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs,
-      llvm::SmallVectorImpl<clang::DeclStmt*>& ArgDecls,
-      bool hasAssignee) override;
+      llvm::SmallVectorImpl<clang::Stmt*>& ArgDecls, bool hasAssignee) override;
   void ActBeforeFinalizingVisitDeclStmt(
       llvm::SmallVectorImpl<clang::Decl*>& decls,
       llvm::SmallVectorImpl<clang::Decl*>& declsDiff) override;

include/clad/Differentiator/ExternalRMVSource.h

@@ -127,7 +127,7 @@ class ExternalRMVSource {
   virtual void ActBeforeFinalizingVisitCallExpr(
       const clang::CallExpr*& CE, clang::Expr*& OverloadedDerivedFn,
       llvm::SmallVectorImpl<clang::Expr*>& derivedCallArgs,
-      llvm::SmallVectorImpl<clang::VarDecl*>& ArgResultDecls, bool asGrad) {}
+      llvm::SmallVectorImpl<clang::Expr*>& ArgResult, bool asGrad) {}
 
   /// This is called just before finalising processing of post and pre
   /// increment and decrement operations.
@@ -155,7 +155,7 @@ class ExternalRMVSource {
 
   virtual void ActBeforeDifferentiatingCallExpr(
       llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs,
-      llvm::SmallVectorImpl<clang::DeclStmt*>& ArgDecls, bool hasAssignee) {}
+      llvm::SmallVectorImpl<clang::Stmt*>& ArgDecls, bool hasAssignee) {}
 
   virtual void ActBeforeFinalizingVisitDeclStmt(
       llvm::SmallVectorImpl<clang::Decl*>& decls,

include/clad/Differentiator/MultiplexExternalRMVSource.h

@@ -48,8 +48,7 @@ class MultiplexExternalRMVSource : public ExternalRMVSource {
   void ActBeforeFinalizingVisitCallExpr(
       const clang::CallExpr*& CE, clang::Expr*& OverloadedDerivedFn,
       llvm::SmallVectorImpl<clang::Expr*>& derivedCallArgs,
-      llvm::SmallVectorImpl<clang::VarDecl*>& ArgResultDecls,
-      bool asGrad) override;
+      llvm::SmallVectorImpl<clang::Expr*>& ArgResult, bool asGrad) override;
   void ActBeforeFinalisingPostIncDecOp(StmtDiff& diff) override;
   void ActAfterCloningLHSOfAssignOp(clang::Expr*&, clang::Expr*&,
                                     clang::BinaryOperatorKind& opCode) override;
@@ -59,8 +58,7 @@ class MultiplexExternalRMVSource : public ExternalRMVSource {
   void ActBeforeFinalizingDifferentiateSingleExpr(const direction& d) override;
   void ActBeforeDifferentiatingCallExpr(
       llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs,
-      llvm::SmallVectorImpl<clang::DeclStmt*>& ArgDecls,
-      bool hasAssignee) override;
+      llvm::SmallVectorImpl<clang::Stmt*>& ArgDecls, bool hasAssignee) override;
   void ActBeforeFinalizingVisitDeclStmt(
       llvm::SmallVectorImpl<clang::Decl*>& decls,
       llvm::SmallVectorImpl<clang::Decl*>& declsDiff) override;

include/clad/Differentiator/VisitorBase.h

@@ -583,16 +583,20 @@ namespace clad {
     ///
     /// \param[in] targetFuncCall The function to get the derivative for.
     /// \param[in] retType The return type of the target call expression.
+    /// \param[in] dfdx The dfdx corresponding to this call expression.
     /// \param[in] numArgs The total number of 'args'.
-    /// \param[in] NumericalDiffMultiArg The built statements to add to block
-    /// later.
+    /// \param[in] PreCallStmts The built statements to add to block
+    /// before the call to the derived function.
+    /// \param[in] PostCallStmts The built statements to add to block
+    /// after the call to the derived function.
     /// \param[in] args All the arguments to the target function.
     /// \param[in] outputArgs The output gradient arguments.
     ///
     /// \returns The derivative function call.
     clang::Expr* GetMultiArgCentralDiffCall(
         clang::Expr* targetFuncCall, clang::QualType retType, unsigned numArgs,
-        llvm::SmallVectorImpl<clang::Stmt*>& NumericalDiffMultiArg,
+        clang::Expr* dfdx, llvm::SmallVectorImpl<clang::Stmt*>& PreCallStmts,
+        llvm::SmallVectorImpl<clang::Stmt*>& PostCallStmts,
         llvm::SmallVectorImpl<clang::Expr*>& args,
         llvm::SmallVectorImpl<clang::Expr*>& outputArgs);
     /// Emits diagnostic messages on differentiation (or lack thereof) for

lib/Differentiator/CladUtils.cpp

@@ -680,5 +680,16 @@ namespace clad {
       return FD->getNameAsString() == "free";
 #endif
     }
+
+    bool isTriviallyCastableArg(const clang::Expr* arg) {
+      // `arg` is the of the type that the function expects.
+      const QualType expectedType = arg->getType().getCanonicalType();
+      // By removing all implicit casts, we can get the type of the original
+      // argument.
+      const QualType argType =
+          arg->IgnoreImplicit()->getType().getCanonicalType();
+      // Compare the unqualified types.
+      return expectedType.getTypePtr() == argType.getTypePtr();
+    }
   } // namespace utils
 } // namespace clad

lib/Differentiator/ErrorEstimator.cpp

@@ -134,7 +134,7 @@ void ErrorEstimationHandler::SaveReturnExpr(Expr* retExpr) {
 
 void ErrorEstimationHandler::EmitNestedFunctionParamError(
     FunctionDecl* fnDecl, llvm::SmallVectorImpl<Expr*>& derivedCallArgs,
-    llvm::SmallVectorImpl<VarDecl*>& ArgResultDecls, size_t numArgs) {
+    llvm::SmallVectorImpl<Expr*>& ArgResult, size_t numArgs) {
   assert(fnDecl && "Must have a value");
   for (size_t i = 0; i < numArgs; i++) {
     if (!fnDecl->getParamDecl(0)->getType()->isLValueReferenceType())
@@ -145,7 +145,7 @@ void ErrorEstimationHandler::EmitNestedFunctionParamError(
     // if (utils::IsReferenceOrPointerType(fnDecl->getParamDecl(i)->getType()))
     //   continue;
     Expr* errorExpr = m_EstModel->AssignError(
-        {derivedCallArgs[i], m_RMV->BuildDeclRef(ArgResultDecls[i])},
+        {derivedCallArgs[i], m_RMV->Clone(ArgResult[i])},
         fnDecl->getNameInfo().getAsString() + "_param_" + std::to_string(i));
     Expr* errorStmt = m_RMV->BuildOp(BO_AddAssign, m_FinalError, errorExpr);
     m_ReverseErrorStmts.push_back(errorStmt);
@@ -605,7 +605,7 @@ void ErrorEstimationHandler::ActBeforeFinalisingPostIncDecOp(StmtDiff& diff) {
 void ErrorEstimationHandler::ActBeforeFinalizingVisitCallExpr(
     const clang::CallExpr*& CE, clang::Expr*& OverloadedDerivedFn,
     llvm::SmallVectorImpl<Expr*>& derivedCallArgs,
-    llvm::SmallVectorImpl<VarDecl*>& ArgResultDecls, bool asGrad) {
+    llvm::SmallVectorImpl<Expr*>& ArgResult, bool asGrad) {
   if (OverloadedDerivedFn && asGrad) {
     // Derivative was found.
     FunctionDecl* fnDecl =
@@ -615,7 +615,7 @@ void ErrorEstimationHandler::ActBeforeFinalizingVisitCallExpr(
     // in the input prameters (if of reference type) to call and save to
     // emit them later.
 
-    EmitNestedFunctionParamError(fnDecl, derivedCallArgs, ArgResultDecls,
+    EmitNestedFunctionParamError(fnDecl, derivedCallArgs, ArgResult,
                                  CE->getNumArgs());
   }
 }
@@ -650,7 +650,7 @@ void ErrorEstimationHandler::ActBeforeFinalizingDifferentiateSingleExpr(
 
 void ErrorEstimationHandler::ActBeforeDifferentiatingCallExpr(
     llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs,
-    llvm::SmallVectorImpl<DeclStmt*>& ArgDecls, bool hasAssignee) {
+    llvm::SmallVectorImpl<Stmt*>& ArgDecls, bool hasAssignee) {
   auto errorRef =
       m_RMV->BuildVarDecl(m_RMV->m_Context.DoubleTy, "_t",
                           m_RMV->getZeroInit(m_RMV->m_Context.DoubleTy));

lib/Differentiator/MultiplexExternalRMVSource.cpp

@@ -149,10 +149,10 @@ void MultiplexExternalRMVSource::ActBeforeFinalisingVisitReturnStmt(
 void MultiplexExternalRMVSource::ActBeforeFinalizingVisitCallExpr(
     const clang::CallExpr*& CE, clang::Expr*& OverloadedDerivedFn,
     llvm::SmallVectorImpl<clang::Expr*>& derivedCallArgs,
-    llvm::SmallVectorImpl<clang::VarDecl*>& ArgResultDecls, bool asGrad) {
+    llvm::SmallVectorImpl<clang::Expr*>& ArgResult, bool asGrad) {
   for (auto source : m_Sources) {
-    source->ActBeforeFinalizingVisitCallExpr(CE, OverloadedDerivedFn, derivedCallArgs,
-                                             ArgResultDecls, asGrad);
+    source->ActBeforeFinalizingVisitCallExpr(
+        CE, OverloadedDerivedFn, derivedCallArgs, ArgResult, asGrad);
   }
 }
 
@@ -198,7 +198,7 @@ void MultiplexExternalRMVSource::ActBeforeFinalizingDifferentiateSingleExpr(
 
 void MultiplexExternalRMVSource::ActBeforeDifferentiatingCallExpr(
     llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs,
-    llvm::SmallVectorImpl<clang::DeclStmt*>& ArgDecls, bool hasAssignee) {
+    llvm::SmallVectorImpl<clang::Stmt*>& ArgDecls, bool hasAssignee) {
   for (auto source : m_Sources)
     source->ActBeforeDifferentiatingCallExpr(pullbackArgs, ArgDecls,
                                              hasAssignee);