Generate constructor_reverse_forw automatically

Author: PetroZarytskyi

include/clad/Differentiator/CladUtils.h

@@ -429,6 +429,9 @@ namespace clad {
     bool isLinearConstructor(const clang::CXXConstructorDecl* CD,
                              const clang::ASTContext& C);
 
+    bool isElidableConstructor(const clang::CXXConstructorDecl* CD,
+                               const clang::ASTContext& C);
+
     /// Returns true if T allows to edit any memory.
     bool isMemoryType(clang::QualType T);
 

include/clad/Differentiator/ReverseModeVisitor.h

@@ -447,8 +447,10 @@ namespace clad {
     ///
     /// \param[in] thisTy `this` type.
     ///
+    /// \param[in] isDerivedThis if true, will build `_d_this`.
+    ///
     /// \returns {_this, free(_this)}
-    StmtDiff BuildThisExpr(clang::QualType thisTy);
+    StmtDiff BuildThisExpr(clang::QualType thisTy, bool isDerivedThis = false);
 
     /// Helper function that checks whether the function to be derived
     /// is meant to be executed only by the GPU

include/clad/Differentiator/STLBuiltins.h

@@ -466,6 +466,21 @@ size_pushforward(const ::std::array<T, N>* a,
 // vector reverse mode
 // more can be found in tests: test/Gradient/STLCustomDerivatives.C
 
+template <typename T, typename... Args>
+clad::ValueAndAdjoint<::std::vector<T>, ::std::vector<T>>
+constructor_reverse_forw(clad::Tag<::std::vector<T>>,
+                         Args...) elidable_reverse_forw;
+
+template <typename T, typename... Args>
+clad::ValueAndAdjoint<::std::allocator<T>, ::std::allocator<T>>
+constructor_reverse_forw(clad::Tag<::std::allocator<T>>,
+                         Args...) elidable_reverse_forw;
+
+template <typename T, typename U, typename... Args>
+clad::ValueAndAdjoint<::std::pair<T, U>, ::std::pair<T, U>>
+constructor_reverse_forw(clad::Tag<::std::pair<T, U>>,
+                         Args...) elidable_reverse_forw;
+
 template <typename T, typename U, typename pU>
 void push_back_reverse_forw(::std::vector<T>* v, U val, ::std::vector<T>* d_v,
                             pU /*d_val*/) {

include/clad/Differentiator/ThrustBuiltins.h

@@ -8,6 +8,8 @@
 
 namespace clad {
 
+#define elidable_reverse_forw __attribute__((annotate("elidable_reverse_forw")))
+
 // zero_init for thrust::device_vector
 template <class T> inline void zero_init(::thrust::device_vector<T>& v) {
   ::thrust::fill(v.begin(), v.end(), T(0)); // NOLINT(misc-include-cleaner)
@@ -17,6 +19,14 @@ namespace custom_derivatives::class_functions {
 
 // Constructors (reverse mode pullbacks)
 
+template <typename T>
+clad::ValueAndAdjoint<::thrust::device_vector<T>, ::thrust::device_vector<T>>
+constructor_reverse_forw(clad::Tag<::thrust::device_vector<T>>,
+                         typename ::thrust::device_vector<T>::size_type count,
+                         const T& val,
+                         typename ::thrust::device_vector<T>::size_type d_count,
+                         const T& d_val) elidable_reverse_forw;
+
 template <typename T>
 inline void
 constructor_pullback(typename ::thrust::device_vector<T>::size_type count,
@@ -29,6 +39,13 @@ constructor_pullback(typename ::thrust::device_vector<T>::size_type count,
   }
 }
 
+template <typename T>
+clad::ValueAndAdjoint<::thrust::device_vector<T>, ::thrust::device_vector<T>>
+constructor_reverse_forw(clad::Tag<::thrust::device_vector<T>>,
+                         typename ::thrust::device_vector<T>::size_type count,
+                         typename ::thrust::device_vector<T>::size_type d_count)
+    elidable_reverse_forw;
+
 template <typename T>
 inline void
 constructor_pullback(typename ::thrust::device_vector<T>::size_type count,

lib/Differentiator/CladUtils.cpp

@@ -358,6 +358,44 @@ namespace clad {
       return true;
     }
 
+    bool isElidableConstructor(const clang::CXXConstructorDecl* CD,
+                               const clang::ASTContext& C) {
+      const CXXRecordDecl* RD = CD->getParent();
+      if (CD->isCopyOrMoveConstructor() && CD->isTrivial())
+        return true;
+      if (RD->isAggregate())
+        return true;
+      if (utils::unwrapIfSingleStmt(CD->getBody()))
+        return false;
+      bool isZeroOrCopyCtor = true;
+      for (CXXCtorInitializer* CI : CD->inits()) {
+        if (CI->isMemberInitializer()) {
+          QualType memberTy = CI->getMember()->getType();
+          if (memberTy->isPointerType() || memberTy->isLValueReferenceType())
+            continue;
+        }
+        Expr* init = CI->getInit()->IgnoreImplicit();
+        Expr::EvalResult value;
+        bool isZero = false;
+        if (clad_compat::Expr_EvaluateAsConstantExpr(init, value, C)) {
+          const auto* val = &value.Val;
+          if (val->isArray() && val->hasArrayFiller())
+            val = &val->getArrayFiller();
+          if (val->isInt())
+            isZero = val->getInt() == 0;
+          else if (val->isFloat())
+            isZero = val->getFloat().isZero();
+        }
+
+        if (!(isa<DeclRefExpr>(init) || isa<CXXConstructExpr>(init) ||
+              isZero)) {
+          isZeroOrCopyCtor = false;
+          break;
+        }
+      }
+      return isZeroOrCopyCtor;
+    }
+
     void getRecordDeclFields(
         const clang::RecordDecl* RD,
         llvm::SmallVectorImpl<const clang::FieldDecl*>& fields) {

lib/Differentiator/DiffPlanner.cpp

@@ -1456,7 +1456,10 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
     forwPassRequest.BaseFunctionName = "constructor";
     forwPassRequest.Mode = DiffMode::reverse_mode_forward_pass;
     forwPassRequest.CallContext = E;
-    if (LookupCustomDerivativeDecl(forwPassRequest))
+    QualType recordTy = CD->getThisType()->getPointeeType();
+    bool elideRevForw =
+        utils::isElidableConstructor(CD, m_Sema.getASTContext());
+    if (LookupCustomDerivativeDecl(forwPassRequest) || !elideRevForw)
       m_DiffRequestGraph.addNode(forwPassRequest, /*isSource=*/true);
 
     // Don't build propagators for calls that do not contribute in
@@ -1488,7 +1491,6 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
     if (request.Function->getDefinition())
       request.Function = request.Function->getDefinition();
 
-    QualType recordTy = CD->getThisType()->getPointeeType();
     if (m_Sema.isStdInitializerList(recordTy, /*elemType=*/nullptr))
       return true;
 

lib/Differentiator/ReverseModeForwPassVisitor.cpp

@@ -11,6 +11,7 @@
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/Expr.h"
+#include "clang/AST/OperationKinds.h"
 #include "clang/Basic/LLVM.h"
 
 #include <algorithm>
@@ -67,15 +68,43 @@ DerivativeAndOverload ReverseModeForwPassVisitor::Derive() {
     beginBlock();
     beginBlock(direction::reverse);
 
+    // If we the differentiated function is a constructor, generate `this`
+    // object and differentiate its inits.
+    Stmt* ctorReturnStmt = nullptr;
+    if (const auto* CD = dyn_cast<CXXConstructorDecl>(m_DiffReq.Function)) {
+      QualType thisTy = CD->getThisType();
+      StmtDiff thisObj = BuildThisExpr(thisTy);
+      StmtDiff dthisObj = BuildThisExpr(thisTy, /*isDerivedThis=*/true);
+      m_ThisExprDerivative = dthisObj.getExpr();
+      addToCurrentBlock(dthisObj.getStmt_dx());
+      for (CXXCtorInitializer* CI : CD->inits()) {
+        StmtDiff CI_diff = DifferentiateCtorInit(CI, thisObj.getExpr());
+        addToCurrentBlock(CI_diff.getStmt(), direction::forward);
+        addToCurrentBlock(CI_diff.getStmt_dx(), direction::forward);
+      }
+      // Build `return {*_this, *_d_this};`
+      SourceLocation validLoc{CD->getBeginLoc()};
+      llvm::SmallVector<Expr*, 2> returnArgs = {
+          BuildOp(UO_Deref, thisObj.getExpr()),
+          BuildOp(UO_Deref, dthisObj.getExpr())};
+      Expr* returnInitList =
+          m_Sema.ActOnInitList(validLoc, returnArgs, validLoc).get();
+      ctorReturnStmt = m_Sema.BuildReturnStmt(validLoc, returnInitList).get();
+    }
+
     StmtDiff bodyDiff = Visit(m_DiffReq->getBody());
     Stmt* forward = bodyDiff.getStmt();
 
     for (Stmt* S : ReverseModeVisitor::m_Globals)
       addToCurrentBlock(S);
 
-    if (auto* CS = dyn_cast<CompoundStmt>(forward))
+    if (auto* CS = dyn_cast_or_null<CompoundStmt>(forward))
       for (Stmt* S : CS->body())
         addToCurrentBlock(S);
+    else
+      addToCurrentBlock(forward);
+
+    addToCurrentBlock(ctorReturnStmt);
 
     Stmt* fnBody = endBlock();
     m_Derivative->setBody(fnBody);
@@ -97,8 +126,13 @@ ReverseModeForwPassVisitor::BuildParams(DiffParams& diffParams) {
 
   std::size_t dParamTypesIdx = m_DiffReq->getNumParams();
 
-  if (const auto* CD = dyn_cast<CXXConversionDecl>(m_DiffReq.Function)) {
-    QualType typeTag = utils::GetCladTagOfType(m_Sema, CD->getConversionType());
+  QualType tagType;
+  if (const auto* CD = dyn_cast<CXXConversionDecl>(m_DiffReq.Function))
+    tagType = CD->getConversionType();
+  else if (const auto* CD = dyn_cast<CXXConstructorDecl>(m_DiffReq.Function))
+    tagType = CD->getThisType()->getPointeeType();
+  if (!tagType.isNull()) {
+    QualType typeTag = utils::GetCladTagOfType(m_Sema, tagType);
     IdentifierInfo* emptyII = &m_Context.Idents.get("");
     ParmVarDecl* typeTagPVD =
         utils::BuildParmVarDecl(m_Sema, m_Derivative, emptyII, typeTag);
@@ -110,7 +144,7 @@ ReverseModeForwPassVisitor::BuildParams(DiffParams& diffParams) {
 
   if (const auto* MD = dyn_cast<CXXMethodDecl>(m_DiffReq.Function)) {
     const CXXRecordDecl* RD = MD->getParent();
-    if (MD->isInstance() && !RD->isLambda()) {
+    if (MD->isInstance() && !RD->isLambda() && !isa<CXXConstructorDecl>(MD)) {
       auto* thisDerivativePVD = utils::BuildParmVarDecl(
           m_Sema, m_Derivative, CreateUniqueIdentifier("_d_this"),
           derivativeFnType->getParamType(dParamTypesIdx));