[WIP] Add Useful Analysis to the forward mode

include/clad/Differentiator/CladConfig.h

@@ -31,6 +31,8 @@ enum opts : unsigned {
   disable_tbr = 1 << (ORDER_BITS + 3),
   enable_va = 1 << (ORDER_BITS + 5),
   disable_va = 1 << (ORDER_BITS + 6),
+  enable_ua = 1 << (ORDER_BITS + 7),
+  disable_ua = 1 << (ORDER_BITS + 8),
 
   // Specifying whether we only want the diagonal of the hessian.
   diagonal_only = 1 << (ORDER_BITS + 4),

include/clad/Differentiator/DiffPlanner.h

@@ -38,9 +38,16 @@ struct DiffRequest {
     bool HasAnalysisRun = false;
   } m_ActivityRunInfo;
 
+  mutable struct UsefulRunInfo {
+    std::set<const clang::VarDecl*> UsefulDecls;
+    std::set<const clang::FunctionDecl*> UsefulFuncs;
+    bool HasAnalysisRun = false;
+  } m_UsefulRunInfo;
+
 public:
   /// Function to be differentiated.
   const clang::FunctionDecl* Function = nullptr;
+  bool ReqAdj = true;
   /// Name of the base function to be differentiated. Can be different from
   /// function->getNameAsString() when higher-order derivatives are computed.
   std::string BaseFunctionName = {};
@@ -65,6 +72,7 @@ struct DiffRequest {
   /// A flag to enable TBR analysis during reverse-mode differentiation.
   bool EnableTBRAnalysis = false;
   bool EnableVariedAnalysis = false;
+  bool EnableUsefulAnalysis = false;
   /// Puts the derived function and its code in the diff call
   void updateCall(clang::FunctionDecl* FD, clang::FunctionDecl* OverloadedFD,
                   clang::Sema& SemaRef);
@@ -123,6 +131,7 @@ struct DiffRequest {
            CallContext == other.CallContext && Args == other.Args &&
            Mode == other.Mode && EnableTBRAnalysis == other.EnableTBRAnalysis &&
            EnableVariedAnalysis == other.EnableVariedAnalysis &&
+           EnableUsefulAnalysis == other.EnableUsefulAnalysis &&
            DVI == other.DVI && use_enzyme == other.use_enzyme &&
            DeclarationOnly == other.DeclarationOnly;
   }
@@ -141,6 +150,7 @@ struct DiffRequest {
 
   bool shouldBeRecorded(clang::Expr* E) const;
   bool shouldHaveAdjoint(const clang::VarDecl* VD) const;
+  bool shouldHaveAdjointForw(const clang::VarDecl* VD) const;
 };
 
   using DiffInterval = std::vector<clang::SourceRange>;
@@ -150,6 +160,7 @@ struct DiffRequest {
     /// TBR analysis during reverse-mode differentiation.
     bool EnableTBRAnalysis = false;
     bool EnableVariedAnalysis = false;
+    bool EnableUsefulAnalysis = false;
   };
 
   class DiffCollector: public clang::RecursiveASTVisitor<DiffCollector> {

lib/Differentiator/BaseForwardModeVisitor.cpp

@@ -1063,7 +1063,14 @@
   // If DRE is of type pointer, then the derivative is a null pointer.
   if (clonedDRE->getType()->isPointerType())
     return StmtDiff(clonedDRE, nullptr);
+
+  if (auto* i = dyn_cast<VarDecl>(DRE->getDecl())) {
+    if (!m_DiffReq.shouldHaveAdjointForw(i))
+      return StmtDiff(clonedDRE, nullptr);
+  }
+
   QualType literalTy = utils::GetValueType(clonedDRE->getType());
+
   return StmtDiff(clonedDRE, ConstantFolder::synthesizeLiteral(
                                  literalTy, m_Context, /*val=*/0));
 }
@@ -1208,6 +1215,7 @@
   }
 
   llvm::SmallVector<Expr*, 16> pushforwardFnArgs;
+  //
   pushforwardFnArgs.insert(pushforwardFnArgs.end(), CallArgs.begin(),
                            CallArgs.end());
   pushforwardFnArgs.insert(pushforwardFnArgs.end(), diffArgs.begin(),
@@ -1284,6 +1292,7 @@
     pushforwardFnRequest.BaseFunctionName = utils::ComputeEffectiveFnName(FD);
     // Silence diag outputs in nested derivation process.
     pushforwardFnRequest.VerboseDiags = false;
+    pushforwardFnRequest.EnableUsefulAnalysis = m_DiffReq.EnableUsefulAnalysis;
 
     // Check if request already derived in DerivedFunctions.
     FunctionDecl* pushforwardFD =
@@ -1446,7 +1455,8 @@
       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() &&
+        Ldiff.getExpr_dx()->isModifiableLvalue(m_Context) != Expr::MLV_Valid) {
       diag(DiagnosticsEngine::Warning, BinOp->getEndLoc(),
            "derivative of an assignment attempts to assign to unassignable "
            "expr, assignment ignored");
@@ -1575,11 +1585,16 @@
   VarDecl* VDDerived =
       BuildVarDecl(VD->getType(), "_d_" + VD->getNameAsString(), initDx,
                    VD->isDirectInit(), /*TSI=*/nullptr, VD->getInitStyle());
-  m_Variables.emplace(VDClone, BuildDeclRef(VDDerived));
+
+  if (!m_DiffReq.shouldHaveAdjointForw(VD))
+    VDDerived = nullptr;
+  else
+    m_Variables.emplace(VDClone, BuildDeclRef(VDDerived));
   return DeclDiff<VarDecl>(VDClone, VDDerived);
 }
 
 StmtDiff BaseForwardModeVisitor::VisitDeclStmt(const DeclStmt* DS) {
+  // llvm::errs() << "\nVisitDeclStmt";
   llvm::SmallVector<Decl*, 4> decls;
   llvm::SmallVector<Decl*, 4> declsDiff;
   // If the type is marked as non_differentiable, skip generating its derivative
@@ -1642,7 +1657,8 @@
       if (VDDiff.getDecl()->getDeclName() != VD->getDeclName())
         m_DeclReplacements[VD] = VDDiff.getDecl();
       decls.push_back(VDDiff.getDecl());
-      declsDiff.push_back(VDDiff.getDecl_dx());
+      if (m_DiffReq.shouldHaveAdjointForw(VD))
+        declsDiff.push_back(VDDiff.getDecl_dx());
     } else if (auto* SAD = dyn_cast<StaticAssertDecl>(D)) {
       DeclDiff<StaticAssertDecl> SADDiff = DifferentiateStaticAssertDecl(SAD);
       if (SADDiff.getDecl())
@@ -1661,6 +1677,8 @@
     DSClone = BuildDeclStmt(decls);
   if (!declsDiff.empty())
     DSDiff = BuildDeclStmt(declsDiff);
+  // llvm::errs() << "\n=====";
+  // DSDiff->dump();
   return StmtDiff(DSClone, DSDiff);
 }
 

lib/Differentiator/CMakeLists.txt

@@ -38,6 +38,7 @@ llvm_add_library(cladDifferentiator
   ReverseModeVisitor.cpp
   TBRAnalyzer.cpp
   StmtClone.cpp
+  UsefulAnalyzer.cpp
   VectorForwardModeVisitor.cpp
   VectorPushForwardModeVisitor.cpp
   Version.cpp

lib/Differentiator/DiffPlanner.cpp

@@ -2,6 +2,7 @@
 
 #include "ActivityAnalyzer.h"
 #include "TBRAnalyzer.h"
+#include "UsefulAnalyzer.h"
 
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/RecursiveASTVisitor.h"
@@ -636,6 +637,26 @@
     return found != m_ActivityRunInfo.ToBeRecorded.end();
   }
 
+  bool DiffRequest::shouldHaveAdjointForw(const VarDecl* VD) const {
+    if (!EnableUsefulAnalysis)
+      return true;
+
+    if (!m_UsefulRunInfo.HasAnalysisRun) {
+
+      UsefulAnalyzer analyzer(Function->getASTContext(),
+                              m_UsefulRunInfo.UsefulDecls,
+                              m_UsefulRunInfo.UsefulFuncs);
+      analyzer.Analyze(Function);
+      m_UsefulRunInfo.HasAnalysisRun = true;
+      // llvm::errs() << "ToBeRecorded:  ";
+      // for (auto* i : m_UsefulRunInfo.UsefulDecls)
+      //   llvm::errs() << i->getNameAsString() << "  ";
+      // llvm::errs() << "\n";
+    }
+    auto found = m_UsefulRunInfo.UsefulDecls.find(VD);
+    return found != m_UsefulRunInfo.UsefulDecls.end();
+  }
+
   bool DiffCollector::VisitCallExpr(CallExpr* E) {
     // Check if we should look into this.
     // FIXME: Generated code does not usually have valid source locations.
@@ -669,6 +690,8 @@
       bool disable_tbr_in_req = false;
       bool enable_va_in_req = false;
       bool disable_va_in_req = false;
+      bool enable_ua_in_req = false;
+      bool disable_ua_in_req = false;
       if (!A->getAnnotation().equals("E") &&
           FD->getTemplateSpecializationArgs()) {
         const auto template_arg = FD->getTemplateSpecializationArgs()->get(0);
@@ -689,6 +712,10 @@
             clad::HasOption(bitmasked_opts_value, clad::opts::enable_va);
         disable_va_in_req =
             clad::HasOption(bitmasked_opts_value, clad::opts::disable_va);
+        enable_ua_in_req =
+            clad::HasOption(bitmasked_opts_value, clad::opts::enable_ua);
+        disable_ua_in_req =
+            clad::HasOption(bitmasked_opts_value, clad::opts::disable_ua);
         if (enable_tbr_in_req && disable_tbr_in_req) {
           utils::EmitDiag(m_Sema, DiagnosticsEngine::Error, endLoc,
                           "Both enable and disable TBR options are specified.");
@@ -711,6 +738,10 @@
         } else {
           request.EnableVariedAnalysis = m_Options.EnableVariedAnalysis;
         }
+        if (enable_ua_in_req || disable_ua_in_req)
+          request.EnableUsefulAnalysis = enable_ua_in_req && !disable_ua_in_req;
+        else
+          request.EnableUsefulAnalysis = m_Options.EnableUsefulAnalysis;
         if (clad::HasOption(bitmasked_opts_value, clad::opts::diagonal_only)) {
           if (!A->getAnnotation().equals("H")) {
             utils::EmitDiag(m_Sema, DiagnosticsEngine::Error, endLoc,

lib/Differentiator/PushForwardModeVisitor.cpp

@@ -24,7 +24,6 @@ StmtDiff PushForwardModeVisitor::VisitReturnStmt(const ReturnStmt* RS) {
   // If there is no return value, we must not attempt to differentiate
   if (!RS->getRetValue())
     return nullptr;
-
   StmtDiff retValDiff = Visit(RS->getRetValue());
   Expr* retVal = retValDiff.getExpr();
   Expr* retVal_dx = retValDiff.getExpr_dx();

lib/Differentiator/UsefulAnalyzer.cpp

@@ -0,0 +1,156 @@
+#include "UsefulAnalyzer.h"
+
+using namespace clang;
+
+namespace clad {
+
+void UsefulAnalyzer::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());
+  // Set current block ID to the ID of entry the block.
+  CFGBlock* exit = &m_CFG->getExit();
+  m_CurBlockID = exit->getBlockID();
+  m_BlockData[m_CurBlockID] = createNewVarsData({});
+  for (const VarDecl* i : m_UsefulDecls)
+    m_BlockData[m_CurBlockID]->insert(i);
+  // Add the entry block to the queue.
+  m_CFGQueue.insert(m_CurBlockID);
+
+  // Visit CFG blocks in the queue until it's empty.
+  while (!m_CFGQueue.empty()) {
+    auto IDIter = m_CFGQueue.begin();
+    m_CurBlockID = *IDIter;
+    m_CFGQueue.erase(IDIter);
+    CFGBlock& nextBlock = *getCFGBlockByID(m_CurBlockID);
+    AnalyzeCFGBlock(nextBlock);
+  }
+}
+
+CFGBlock* UsefulAnalyzer::getCFGBlockByID(unsigned ID) {
+  return *(m_CFG->begin() + ID);
+}
+
+bool UsefulAnalyzer::isUseful(const VarDecl* VD) const {
+  const VarsData& curBranch = getCurBlockVarsData();
+  return curBranch.find(VD) != curBranch.end();
+}
+
+void UsefulAnalyzer::copyVarToCurBlock(const clang::VarDecl* VD) {
+  VarsData& curBranch = getCurBlockVarsData();
+  curBranch.insert(VD);
+}
+
+static void mergeVarsData(std::set<const clang::VarDecl*>* targetData,
+                          std::set<const clang::VarDecl*>* mergeData) {
+  for (const clang::VarDecl* i : *mergeData)
+    targetData->insert(i);
+  *mergeData = *targetData;
+}
+
+void UsefulAnalyzer::AnalyzeCFGBlock(const CFGBlock& block) {
+
+  for (auto ib = block.end(); ib != block.begin() - 1; ib--) {
+    if (ib->getKind() == clang::CFGElement::Statement) {
+
+      const clang::Stmt* S = ib->castAs<clang::CFGStmt>().getStmt();
+      // The const_cast is inevitable, since there is no
+      // ConstRecusiveASTVisitor.
+      // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
+      TraverseStmt(const_cast<clang::Stmt*>(S));
+    }
+  }
+
+  for (const clang::CFGBlock::AdjacentBlock pred : block.preds()) {
+    if (!pred)
+      continue;
+    auto& predData = m_BlockData[pred->getBlockID()];
+    if (!predData)
+      predData = createNewVarsData(*m_BlockData[block.getBlockID()]);
+
+    bool shouldPushPred = true;
+    if (pred->getBlockID() < block.getBlockID()) {
+      if (m_LoopMem == *m_BlockData[block.getBlockID()])
+        shouldPushPred = false;
+
+      for (const VarDecl* i : *m_BlockData[block.getBlockID()])
+        m_LoopMem.insert(i);
+    }
+
+    if (shouldPushPred)
+      m_CFGQueue.insert(pred->getBlockID());
+
+    mergeVarsData(predData.get(), m_BlockData[block.getBlockID()].get());
+  }
+
+  for (const VarDecl* i : *m_BlockData[block.getBlockID()])
+    m_UsefulDecls.insert(i);
+}
+
+bool UsefulAnalyzer::VisitBinaryOperator(BinaryOperator* BinOp) {
+  Expr* L = BinOp->getLHS();
+  Expr* R = BinOp->getRHS();
+  const auto opCode = BinOp->getOpcode();
+  if (BinOp->isAssignmentOp()) {
+    m_Useful = false;
+    TraverseStmt(L);
+    m_Marking = m_Useful;
+    TraverseStmt(R);
+    m_Marking = false;
+  } else if (opCode == BO_Add || opCode == BO_Sub || opCode == BO_Mul ||
+             opCode == BO_Div) {
+    for (auto* subexpr : BinOp->children())
+      if (!isa<BinaryOperator>(subexpr))
+        TraverseStmt(subexpr);
+  }
+  return true;
+}
+
+bool UsefulAnalyzer::VisitDeclStmt(DeclStmt* DS) {
+  for (Decl* D : DS->decls()) {
+    if (auto* VD = dyn_cast<VarDecl>(D)) {
+      if (isUseful(VD)) {
+        m_Useful = true;
+        m_Marking = true;
+      }
+      if (Expr* init = cast<VarDecl>(D)->getInit())
+        TraverseStmt(init);
+      m_Marking = false;
+    }
+  }
+  return true;
+}
+
+bool UsefulAnalyzer::VisitReturnStmt(ReturnStmt* RS) {
+  m_Useful = true;
+  m_Marking = true;
+  auto* rv = RS->getRetValue();
+  TraverseStmt(rv);
+  return true;
+}
+
+bool UsefulAnalyzer::VisitCallExpr(CallExpr* CE) {
+  if (m_Useful)
+    return true;
+  FunctionDecl* FD = CE->getDirectCallee();
+  m_UsefulFuncs.insert(FD);
+  return true;
+}
+
+bool UsefulAnalyzer::VisitDeclRefExpr(DeclRefExpr* DRE) {
+  auto* VD = dyn_cast<VarDecl>(DRE->getDecl());
+  if (!VD)
+    return true;
+
+  if (isUseful(VD))
+    m_Useful = true;
+
+  if (m_Useful && m_Marking)
+    copyVarToCurBlock(VD);
+
+  return true;
+}
+
+} // namespace clad