Move TBR call to DiffCollector

Author: Max Andriychuk

include/clad/Differentiator/DiffPlanner.h

@@ -10,6 +10,7 @@
 #include "clang/AST/Decl.h"
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/Analysis/AnalysisDeclContext.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/Compiler.h"
@@ -56,6 +57,8 @@ struct DiffRequest {
 public:
   /// Function to be differentiated.
   const clang::FunctionDecl* Function = nullptr;
+  /// Stores info relevant to the analysis(CFG, ASTContext).
+  clang::AnalysisDeclContext* AnalysisDC = nullptr;
   /// Name of the base function to be differentiated. Can be different from
   /// function->getNameAsString() when higher-order derivatives are computed.
   std::string BaseFunctionName = {};
@@ -178,6 +181,10 @@ struct DiffRequest {
   std::string ComputeDerivativeName() const;
   bool HasIndependentParameter(const clang::ParmVarDecl* PVD) const;
 
+  std::set<clang::SourceLocation>& getToBeRecorded() const {
+    m_TbrRunInfo.HasAnalysisRun = true;
+    return m_TbrRunInfo.ToBeRecorded;
+  }
   void addVariedDecl(const clang::VarDecl* init) {
     m_ActivityRunInfo.VariedDecls.insert(init);
   }

lib/Differentiator/ActivityAnalyzer.cpp

@@ -6,9 +6,9 @@ using namespace clang;
 namespace clad {
 
 void VariedAnalyzer::Analyze(const FunctionDecl* FD) {
-  m_BlockData.resize(m_ADContext.getCFG()->size());
+  m_BlockData.resize(m_AnalysisDC->getCFG()->size());
   // Set current block ID to the ID of entry the block.
-  CFGBlock* entry = &m_ADContext.getCFG()->getEntry();
+  CFGBlock* entry = &m_AnalysisDC->getCFG()->getEntry();
   m_CurBlockID = entry->getBlockID();
   m_BlockData[m_CurBlockID] = createNewVarsData({});
   for (const VarDecl* i : m_VariedDecls)
@@ -34,7 +34,7 @@ void mergeVarsData(std::set<const clang::VarDecl*>* targetData,
 }
 
 CFGBlock* VariedAnalyzer::getCFGBlockByID(unsigned ID) {
-  return *(m_ADContext.getCFG()->begin() + ID);
+  return *(m_AnalysisDC->getCFG()->begin() + ID);
 }
 
 void VariedAnalyzer::AnalyzeCFGBlock(const CFGBlock& block) {

lib/Differentiator/ActivityAnalyzer.h

@@ -36,7 +36,7 @@ class VariedAnalyzer : public clang::RecursiveASTVisitor<VariedAnalyzer> {
 
   clang::CFGBlock* getCFGBlockByID(unsigned ID);
 
-  clang::AnalysisDeclContext& m_ADContext;
+  clang::AnalysisDeclContext* m_AnalysisDC;
   std::vector<std::unique_ptr<VarsData>> m_BlockData;
   unsigned m_CurBlockID{};
   std::set<unsigned> m_CFGQueue;
@@ -55,9 +55,9 @@ class VariedAnalyzer : public clang::RecursiveASTVisitor<VariedAnalyzer> {
 
 public:
   /// Constructor
-  VariedAnalyzer(clang::AnalysisDeclContext& ADContext,
+  VariedAnalyzer(clang::AnalysisDeclContext* AnalysisDC,
                  std::set<const clang::VarDecl*>& Decls)
-      : m_VariedDecls(Decls), m_ADContext(ADContext) {}
+      : m_VariedDecls(Decls), m_AnalysisDC(AnalysisDC) {}
 
   /// Destructor
   ~VariedAnalyzer() = default;

lib/Differentiator/DiffPlanner.cpp

@@ -649,12 +649,10 @@ DeclRefExpr* getArgFunction(CallExpr* call, Sema& SemaRef) {
 
     if (!m_TbrRunInfo.HasAnalysisRun) {
       TimedAnalysisRegion R("TBR " + BaseFunctionName);
-
-      TBRAnalyzer analyzer(Function->getASTContext(),
-                           m_TbrRunInfo.ToBeRecorded);
+      TBRAnalyzer analyzer(AnalysisDC, getToBeRecorded());
       analyzer.Analyze(Function);
-      m_TbrRunInfo.HasAnalysisRun = true;
     }
+
     auto found = m_TbrRunInfo.ToBeRecorded.find(E->getBeginLoc());
     return found != m_TbrRunInfo.ToBeRecorded.end();
   }
@@ -1214,23 +1212,31 @@ DeclRefExpr* getArgFunction(CallExpr* call, Sema& SemaRef) {
 
     if (!LookupCustomDerivativeDecl(request)) {
       // Create analysis manager once for all analyses.
-      AnalysisDeclContextManager ADCM(request.Function->getASTContext());
+      AnalysisDeclContextManager* ADCM =
+          new AnalysisDeclContextManager(request.Function->getASTContext());
       clang::CFG::BuildOptions Options;
-      AnalysisDeclContext ADC(&ADCM, request.Function, Options);
-      ADC.getCFG();
+      request.AnalysisDC =
+          new AnalysisDeclContext(ADCM, request.Function, Options);
 
       if (m_TopMostReq->EnableVariedAnalysis) {
         TimedAnalysisRegion R("VA " + request.BaseFunctionName);
-        VariedAnalyzer analyzer(ADC, request.getVariedDecls());
+        VariedAnalyzer analyzer(request.AnalysisDC, request.getVariedDecls());
         analyzer.Analyze(request.Function);
       }
 
       if (m_TopMostReq->EnableUsefulAnalysis) {
         TimedAnalysisRegion R("UA " + request.BaseFunctionName);
-        UsefulAnalyzer analyzer(ADC, request.getUsefulDecls());
+        UsefulAnalyzer analyzer(request.AnalysisDC, request.getUsefulDecls());
         analyzer.Analyze(request.Function);
       }
 
+      if (request.Function->isDefined() && m_TopMostReq->EnableTBRAnalysis &&
+          (request.Mode == DiffMode::reverse ||
+           request.Mode == DiffMode::pullback)) {
+        TimedAnalysisRegion R("TBR " + request.BaseFunctionName);
+        TBRAnalyzer analyzer(request.AnalysisDC, request.getToBeRecorded());
+        analyzer.Analyze(request.Function);
+      }
       // Recurse into call graph.
       TraverseFunctionDeclOnce(request.Function);
     }
@@ -1314,9 +1320,23 @@ DeclRefExpr* getArgFunction(CallExpr* call, Sema& SemaRef) {
     if (m_Sema.isStdInitializerList(recordTy, /*elemType=*/nullptr))
       return true;
 
-    if (!LookupCustomDerivativeDecl(request))
+    if (!LookupCustomDerivativeDecl(request)) {
+      AnalysisDeclContextManager* ADCM =
+          new AnalysisDeclContextManager(request.Function->getASTContext());
+      clang::CFG::BuildOptions Options;
+      request.AnalysisDC =
+          new AnalysisDeclContext(ADCM, request.Function, Options);
+
+      if (m_TopMostReq->EnableTBRAnalysis &&
+          (request.Mode == DiffMode::reverse ||
+           request.Mode == DiffMode::pullback)) {
+        TimedAnalysisRegion R("TBR " + request.BaseFunctionName);
+        TBRAnalyzer analyzer(request.AnalysisDC, request.getToBeRecorded());
+        analyzer.Analyze(request.Function);
+      }
       // Recurse into call graph.
       TraverseFunctionDeclOnce(request.Function);
+    }
     m_DiffRequestGraph.addNode(request, /*isSource=*/true);
 
     return true;

lib/Differentiator/TBRAnalyzer.cpp

@@ -198,10 +198,20 @@ TBRAnalyzer::VarData::VarData(QualType QT, const ASTContext& C,
     const auto* recordDecl = recordType->getDecl();
     auto& newArrMap = m_Val.m_ArrData;
     newArrMap = std::unique_ptr<ArrMap>(new ArrMap());
-    for (const auto* field : recordDecl->fields()) {
-      const auto varType = field->getType();
-      (*newArrMap)[getProfileID(field)] = VarData(varType, C);
-    }
+
+    // FIXME: For some reason if a variable is of kokkos-type we start creating
+    // infinite amount of VarData objects.
+    bool isInKokkosNS = false;
+    if (const auto* ns =
+            llvm::dyn_cast<clang::NamespaceDecl>(recordDecl->getDeclContext()))
+      if (ns->getName() == "Kokkos")
+        isInKokkosNS = true;
+
+    if (!isInKokkosNS)
+      for (const auto* field : recordDecl->fields()) {
+        const auto varType = field->getType();
+        (*newArrMap)[getProfileID(field)] = VarData(varType, C);
+      }
   }
 }
 
@@ -250,11 +260,13 @@ void TBRAnalyzer::addVar(const clang::VarDecl* VD, bool forceNonRefType) {
     varType = arrayParam->getOriginalType();
   else
     varType = VD->getType();
+
   // If varType represents auto or auto*, get the type of init.
   if (utils::IsAutoOrAutoPtrType(varType))
     varType = VD->getInit()->getType();
 
-  curBranch[VD] = VarData(varType, m_Context, forceNonRefType);
+  curBranch[VD] =
+      VarData(varType, m_AnalysisDC->getASTContext(), forceNonRefType);
 }
 
 void TBRAnalyzer::markLocation(const clang::Expr* E) {
@@ -301,15 +313,11 @@ TBRAnalyzer::getVarDataFromDecl(const clang::VarDecl* VD) {
 }
 
 void TBRAnalyzer::Analyze(const FunctionDecl* FD) {
-  // Build the CFG (control-flow graph) of FD.
-  clang::CFG::BuildOptions Options;
-  m_CFG = clang::CFG::buildCFG(FD, FD->getBody(), &m_Context, Options);
-
-  m_BlockData.resize(m_CFG->size());
-  m_BlockPassCounter.resize(m_CFG->size(), 0);
+  m_BlockData.resize(m_AnalysisDC->getCFG()->size());
+  m_BlockPassCounter.resize(m_AnalysisDC->getCFG()->size(), 0);
 
   // Set current block ID to the ID of entry the block.
-  auto* entry = &m_CFG->getEntry();
+  auto* entry = &m_AnalysisDC->getCFG()->getEntry();
   m_CurBlockID = entry->getBlockID();
   m_BlockData[m_CurBlockID] = std::unique_ptr<VarsData>(new VarsData());
 
@@ -319,7 +327,8 @@ void TBRAnalyzer::Analyze(const FunctionDecl* FD) {
   if (MD && !MD->isStatic()) {
     const Type* recordType = MD->getParent()->getTypeForDecl();
     VarData& thisData = getCurBlockVarsData()[nullptr];
-    thisData = VarData(QualType::getFromOpaquePtr(recordType), m_Context);
+    thisData = VarData(QualType::getFromOpaquePtr(recordType),
+                       m_AnalysisDC->getASTContext());
     // We have to set all pointer/reference parameters to tbr
     // since method pullbacks aren't supposed to change objects.
     setIsRequired(thisData);
@@ -347,7 +356,7 @@ void TBRAnalyzer::Analyze(const FunctionDecl* FD) {
         LLVM_DEBUG(llvm::dbgs() << "successor: " << succ->getBlockID() << "\n");
   }
 
-  clang::SourceManager& SM = m_Context.getSourceManager();
+  clang::SourceManager& SM = m_AnalysisDC->getASTContext().getSourceManager();
   for (SourceLocation Loc : m_TBRLocs) {
     unsigned line = SM.getPresumedLoc(Loc).getLine();
     unsigned column = SM.getPresumedLoc(Loc).getColumn();
@@ -415,7 +424,7 @@ void TBRAnalyzer::VisitCFGBlock(const CFGBlock& block) {
 }
 
 CFGBlock* TBRAnalyzer::getCFGBlockByID(unsigned ID) {
-  return *(m_CFG->begin() + ID);
+  return *(m_AnalysisDC->getCFG()->begin() + ID);
 }
 
 TBRAnalyzer::VarsData*
@@ -571,9 +580,11 @@ bool TBRAnalyzer::TraverseDeclStmt(DeclStmt* DS) {
     if (auto* VD = dyn_cast<VarDecl>(D)) {
       addVar(VD);
       if (clang::Expr* init = VD->getInit()) {
+
         setMode(Mode::kMarkingMode);
         TraverseStmt(init);
         resetMode();
+
         auto& VDExpr = getCurBlockVarsData()[VD];
         // if the declared variable is ref type attach its VarData to the
         // VarData of the RHS variable.
@@ -622,9 +633,10 @@ bool TBRAnalyzer::TraverseBinaryOperator(BinaryOperator* BinOp) {
     // Multiplication results in a linear expression if and only if one of the
     // factors is constant.
     Expr::EvalResult dummy;
-    bool nonLinear =
-        !clad_compat::Expr_EvaluateAsConstantExpr(R, dummy, m_Context) &&
-        !clad_compat::Expr_EvaluateAsConstantExpr(L, dummy, m_Context);
+    bool nonLinear = !clad_compat::Expr_EvaluateAsConstantExpr(
+                         R, dummy, m_AnalysisDC->getASTContext()) &&
+                     !clad_compat::Expr_EvaluateAsConstantExpr(
+                         L, dummy, m_AnalysisDC->getASTContext());
     if (nonLinear)
       startNonLinearMode();
 
@@ -637,8 +649,8 @@ bool TBRAnalyzer::TraverseBinaryOperator(BinaryOperator* BinOp) {
     // Division normally only results in a linear expression when the
     // denominator is constant.
     Expr::EvalResult dummy;
-    bool nonLinear =
-        !clad_compat::Expr_EvaluateAsConstantExpr(R, dummy, m_Context);
+    bool nonLinear = !clad_compat::Expr_EvaluateAsConstantExpr(
+        R, dummy, m_AnalysisDC->getASTContext());
     if (nonLinear)
       startNonLinearMode();
 
@@ -663,8 +675,8 @@ bool TBRAnalyzer::TraverseBinaryOperator(BinaryOperator* BinOp) {
       // represents the same operation as 'x = x * y' ('x = x / y') and,
       // therefore, LHS has to be visited in kMarkingMode|kNonLinearMode.
       Expr::EvalResult dummy;
-      bool RisNotConst =
-          !clad_compat::Expr_EvaluateAsConstantExpr(R, dummy, m_Context);
+      bool RisNotConst = !clad_compat::Expr_EvaluateAsConstantExpr(
+          R, dummy, m_AnalysisDC->getASTContext());
       if (RisNotConst)
         setMode(Mode::kMarkingMode | Mode::kNonLinearMode);
       TraverseStmt(L);

lib/Differentiator/TBRAnalyzer.h

@@ -4,6 +4,7 @@
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/Stmt.h"
+#include "clang/Analysis/AnalysisDeclContext.h"
 #include "clang/Analysis/CFG.h"
 
 #include "clad/Differentiator/CladUtils.h"
@@ -31,7 +32,7 @@ class TBRAnalyzer : public clang::RecursiveASTVisitor<TBRAnalyzer> {
 
   ProfileID getProfileID(const Expr* E) const {
     ProfileID profID;
-    E->Profile(profID, m_Context, /* Canonical */ true);
+    E->Profile(profID, m_AnalysisDC->getASTContext(), /* Canonical */ true);
     return profID;
   }
 
@@ -229,10 +230,7 @@ class TBRAnalyzer : public clang::RecursiveASTVisitor<TBRAnalyzer> {
   /// a new one).
   std::vector<int> m_ModeStack;
 
-  ASTContext& m_Context;
-
-  /// clang::CFG of the function being analysed.
-  std::unique_ptr<clang::CFG> m_CFG;
+  clang::AnalysisDeclContext* m_AnalysisDC;
 
   /// Stores VarsData structures for CFG blocks (the indices in
   /// the vector correspond to CFG blocks' IDs)
@@ -282,8 +280,9 @@ class TBRAnalyzer : public clang::RecursiveASTVisitor<TBRAnalyzer> {
 
 public:
   /// Constructor
-  TBRAnalyzer(ASTContext& Context, std::set<clang::SourceLocation>& Locs)
-      : m_TBRLocs(Locs), m_Context(Context) {
+  TBRAnalyzer(clang::AnalysisDeclContext* AnalysisDC,
+              std::set<clang::SourceLocation>& Locs)
+      : m_TBRLocs(Locs), m_AnalysisDC(AnalysisDC) {
     m_ModeStack.push_back(0);
   }
 

lib/Differentiator/UsefulAnalyzer.cpp

@@ -6,9 +6,9 @@ namespace clad {
 
 void UsefulAnalyzer::Analyze(const FunctionDecl* FD) {
   // Build the CFG (control-flow graph) of FD.
-  m_BlockData.resize(m_ADContext.getCFG()->size());
+  m_BlockData.resize(m_AnalysisDC->getCFG()->size());
   // Set current block ID to the ID of entry the block.
-  CFGBlock* exit = &m_ADContext.getCFG()->getExit();
+  CFGBlock* exit = &m_AnalysisDC->getCFG()->getExit();
   m_CurBlockID = exit->getBlockID();
   m_BlockData[m_CurBlockID] = createNewVarsData({});
   // Add the entry block to the queue.
@@ -25,7 +25,7 @@ void UsefulAnalyzer::Analyze(const FunctionDecl* FD) {
 }
 
 CFGBlock* UsefulAnalyzer::getCFGBlockByID(unsigned ID) {
-  return *(m_ADContext.getCFG()->begin() + ID);
+  return *(m_AnalysisDC->getCFG()->begin() + ID);
 }
 
 bool UsefulAnalyzer::isUseful(const VarDecl* VD) const {

lib/Differentiator/UsefulAnalyzer.h

@@ -31,7 +31,7 @@ class UsefulAnalyzer : public clang::RecursiveASTVisitor<UsefulAnalyzer> {
 
   clang::CFGBlock* getCFGBlockByID(unsigned ID);
 
-  clang::AnalysisDeclContext& m_ADContext;
+  clang::AnalysisDeclContext* m_AnalysisDC;
   std::unique_ptr<clang::CFG> m_CFG;
   std::vector<std::unique_ptr<VarsData>> m_BlockData;
   unsigned m_CurBlockID{};
@@ -46,9 +46,9 @@ class UsefulAnalyzer : public clang::RecursiveASTVisitor<UsefulAnalyzer> {
 
 public:
   /// Constructor
-  UsefulAnalyzer(clang::AnalysisDeclContext& ADContext,
+  UsefulAnalyzer(clang::AnalysisDeclContext* AnalysisDC,
                  std::set<const clang::VarDecl*>& Decls)
-      : m_UsefulDecls(Decls), m_ADContext(ADContext) {}
+      : m_UsefulDecls(Decls), m_AnalysisDC(AnalysisDC) {}
 
   /// Destructor
   ~UsefulAnalyzer() = default;