Extract DiffScheduler to own the request graph and collector. NFC

vgvassilev · vgvassilev · commit 076976686ef8 · 2026-07-19T19:16:31.000-04:00
The request graph, the analysis-context pool, the derived-function map
and the collector that fills them all lived as separate CladPlugin
members, threaded individually into DerivativeBuilder and reconstructed
at every walk site.

Group them into a DiffScheduler that CladPlugin owns and hands to
DerivativeBuilder by reference. The four are one concept -- the
collector needs the other three to plan a request -- so naming it leaves
the plugin with its actual job of driving Sema and CodeGen. The
scheduler is built on first use because it needs Sema, which is not
available when the plugin is constructed.

Owning RequestOptions by value rather than referencing a walk-site local
is what lets the collector outlive a single walk; note in passing that
it models invocation-wide defaults, not per-request state.
diff --git a/include/clad/Differentiator/DerivativeBuilder.h b/include/clad/Differentiator/DerivativeBuilder.h
@@ -12,6 +12,7 @@
 #include "clad/Differentiator/CladUtils.h"
 #include "clad/Differentiator/DerivedFnCollector.h"
 #include "clad/Differentiator/DiffPlanner.h"
+#include "clad/Differentiator/DiffScheduler.h"
 
 #include "clang/AST/Decl.h"
 #include "clang/AST/RecursiveASTVisitor.h"
@@ -111,8 +112,7 @@ struct DerivativeAndOverload {
     clang::Sema& m_Sema;
     plugin::CladPlugin& m_CladPlugin;
     clang::ASTContext& m_Context;
-    DerivedFnCollector& m_DFC;
-    clad::DynamicGraph<DiffRequest>& m_DiffRequestGraph;
+    DiffScheduler& m_Scheduler;
     std::unique_ptr<utils::StmtClone> m_NodeCloner;
     clang::NamespaceDecl* m_BuiltinDerivativesNSD;
     clang::NamespaceDecl* m_NumericalDiffNSD;
@@ -177,8 +177,7 @@ struct DerivativeAndOverload {
 
   public:
     DerivativeBuilder(clang::Sema& S, plugin::CladPlugin& P,
-                      DerivedFnCollector& DFC,
-                      clad::DynamicGraph<DiffRequest>& DRG);
+                      DiffScheduler& Scheduler);
     ~DerivativeBuilder();
     /// Fuction to set the error diagnostic printing value for numerical
     /// differentiation.
diff --git a/include/clad/Differentiator/DiffPlanner.h b/include/clad/Differentiator/DiffPlanner.h
@@ -245,6 +245,8 @@ struct DiffRequest {
 
   using DiffInterval = std::vector<clang::SourceRange>;
 
+  // FIXME: These are translation-unit-wide defaults taken from the compiler
+  // invocation, not the options of a request; rename to InvocationOptions.
   struct RequestOptions {
     /// This is a flag to indicate the default behaviour to enable/disable
     /// TBR analysis during reverse-mode differentiation.
diff --git a/include/clad/Differentiator/DiffScheduler.h b/include/clad/Differentiator/DiffScheduler.h
@@ -0,0 +1,45 @@
+//--------------------------------------------------------------------*- C++ -*-
+// clad - the C++ Clang-based Automatic Differentiator
+//------------------------------------------------------------------------------
+
+#ifndef CLAD_DIFFERENTIATOR_DIFFSCHEDULER_H
+#define CLAD_DIFFERENTIATOR_DIFFSCHEDULER_H
+
+#include "clad/Differentiator/DerivedFnCollector.h"
+#include "clad/Differentiator/DiffPlanner.h"
+#include "clad/Differentiator/DynamicGraph.h"
+
+namespace clang {
+class DeclGroupRef;
+class Sema;
+} // namespace clang
+
+namespace clad {
+
+/// Owns the differentiation request graph and everything that builds it: the
+/// collector, the analysis-context pool and the derived-function map.
+class DiffScheduler {
+  clang::Sema& m_Sema;
+  RequestOptions m_Options;
+  DiffInterval& m_Interval;
+  DynamicGraph<DiffRequest> m_Graph;
+  OwnedAnalysisContexts m_AllAnalysisDC;
+  DerivedFnCollector m_DFC;
+  DiffCollector m_Collector;
+
+public:
+  DiffScheduler(clang::Sema& S, const RequestOptions& Opts,
+                DiffInterval& Interval)
+      : m_Sema(S), m_Options(Opts), m_Interval(Interval),
+        m_Collector(m_Interval, m_Graph, m_Sema, m_Options, m_AllAnalysisDC) {}
+
+  DynamicGraph<DiffRequest>& getGraph() { return m_Graph; }
+  DerivedFnCollector& getDerivedFns() { return m_DFC; }
+
+  /// Static planning pass over a group of top-level declarations.
+  void Plan(clang::DeclGroupRef DGR) { m_Collector.Walk(DGR); }
+};
+
+} // namespace clad
+
+#endif // CLAD_DIFFERENTIATOR_DIFFSCHEDULER_H
diff --git a/lib/Differentiator/BaseForwardModeVisitor.cpp b/lib/Differentiator/BaseForwardModeVisitor.cpp
@@ -218,7 +218,7 @@ DerivativeAndOverload BaseForwardModeVisitor::Derive() {
     FunctionDecl* FoundFD =
         R.empty() ? nullptr : dyn_cast<FunctionDecl>(R.front());
     if (!RD->isLambda() && !R.empty() &&
-        !m_Builder.m_DFC.IsCladDerivative(FoundFD)) {
+        !m_Builder.m_Scheduler.getDerivedFns().IsCladDerivative(FoundFD)) {
       Sema::NestedNameSpecInfo IdInfo(RD->getIdentifier(), noLoc, noLoc,
                                       /*ObjectType=*/nullptr);
       // FIXME: Address nested classes where SS should be set.
diff --git a/lib/Differentiator/DerivativeBuilder.cpp b/lib/Differentiator/DerivativeBuilder.cpp
@@ -14,6 +14,7 @@
 #include "clad/Differentiator/Compatibility.h"
 #include "clad/Differentiator/DiffMode.h"
 #include "clad/Differentiator/DiffPlanner.h"
+#include "clad/Differentiator/DiffScheduler.h"
 #include "clad/Differentiator/DynamicGraph.h"
 #include "clad/Differentiator/ErrorEstimator.h"
 #include "clad/Differentiator/HessianModeVisitor.h"
@@ -58,10 +59,9 @@ using namespace clang;
 namespace clad {
 
 DerivativeBuilder::DerivativeBuilder(clang::Sema& S, plugin::CladPlugin& P,
-                                     DerivedFnCollector& DFC,
-                                     clad::DynamicGraph<DiffRequest>& G)
-    : m_Sema(S), m_CladPlugin(P), m_Context(S.getASTContext()), m_DFC(DFC),
-      m_DiffRequestGraph(G),
+                                     DiffScheduler& Scheduler)
+    : m_Sema(S), m_CladPlugin(P), m_Context(S.getASTContext()),
+      m_Scheduler(Scheduler),
       m_NodeCloner(new utils::StmtClone(m_Sema, m_Context)),
       m_BuiltinDerivativesNSD(nullptr), m_NumericalDiffNSD(nullptr) {}
 
@@ -390,7 +390,7 @@ static void registerDerivative(Decl* D, Sema& S, const DiffRequest& R) {
       // differentiation due to unavailable definition.
       if (auto* CE = dyn_cast_or_null<CallExpr>(OverloadedFn))
         if (FunctionDecl* FD = CE->getDirectCallee())
-          m_DFC.AddToCustomDerivativeSet(FD);
+          m_Scheduler.getDerivedFns().AddToCustomDerivativeSet(FD);
     }
     return OverloadedFn;
   }
@@ -425,7 +425,8 @@ static void registerDerivative(Decl* D, Sema& S, const DiffRequest& R) {
       // derivative function. In that case, we should not derive the definition
       // again.
       if (derivative &&
-          (derivative->isDefined() || m_DFC.IsCustomDerivative(derivative)))
+          (derivative->isDefined() ||
+           m_Scheduler.getDerivedFns().IsCustomDerivative(derivative)))
         alreadyDerived = true;
 
       // Add the request to derive the definition of the forward mode derivative
@@ -535,7 +536,8 @@ static void registerDerivative(Decl* D, Sema& S, const DiffRequest& R) {
         // If only declaration is requested, allow this for clad-generated
         // functions or custom derivatives.
         if (!request.DeclarationOnly ||
-            !(m_DFC.IsCladDerivative(FD) || m_DFC.IsCustomDerivative(FD))) {
+            !(m_Scheduler.getDerivedFns().IsCladDerivative(FD) ||
+              m_Scheduler.getDerivedFns().IsCustomDerivative(FD))) {
           const auto& name = FD->getName();
           // FIXME: Currently, these functions cannot be covered with custom
           // derivatives because templates are not well-supported in custom
@@ -647,7 +649,7 @@ static void registerDerivative(Decl* D, Sema& S, const DiffRequest& R) {
     //   derivative is a member function it goes into an infinite loop
     bool isCustomDerivative = false;
     if (auto* FD = dyn_cast_or_null<FunctionDecl>(result.derivative))
-      isCustomDerivative = m_DFC.IsCustomDerivative(FD);
+      isCustomDerivative = m_Scheduler.getDerivedFns().IsCustomDerivative(FD);
     if (!isCustomDerivative) {
       if (auto* FD = result.derivative)
         registerDerivative(FD, m_Sema, request);
@@ -718,14 +720,14 @@ static void registerDerivative(Decl* D, Sema& S, const DiffRequest& R) {
 
   FunctionDecl*
   DerivativeBuilder::FindDerivedFunction(const DiffRequest& request) {
-    auto DFI = m_DFC.Find(request);
+    auto DFI = m_Scheduler.getDerivedFns().Find(request);
     if (DFI.IsValid())
       return DFI.DerivedFn();
     return nullptr;
   }
 
   void DerivativeBuilder::AddEdgeToGraph(const DiffRequest& request,
                                          bool alreadyDerived /*=false*/) {
-    m_DiffRequestGraph.addEdgeToCurrentNode(request, alreadyDerived);
+    m_Scheduler.getGraph().addEdgeToCurrentNode(request, alreadyDerived);
   }
   } // end namespace clad
diff --git a/lib/Differentiator/HessianModeVisitor.cpp b/lib/Differentiator/HessianModeVisitor.cpp
@@ -183,13 +183,13 @@ DerivativeAndOverload HessianModeVisitor::Derive() {
               utils::CreateStringLiteral(m_Context, independentArgString);
           FunctionDecl* DFD = nullptr;
           if (m_DiffReq.Mode == DiffMode::hessian_diagonal)
-            DFD = DeriveUsingForwardModeTwice(m_Sema, m_CladPlugin, m_Builder,
-                                              m_DiffReq, ForwardModeIASL,
-                                              m_Builder.m_DFC);
+            DFD = DeriveUsingForwardModeTwice(
+                m_Sema, m_CladPlugin, m_Builder, m_DiffReq, ForwardModeIASL,
+                m_Builder.m_Scheduler.getDerivedFns());
           else
             DFD = DeriveUsingForwardAndReverseMode(
                 m_Sema, m_CladPlugin, m_Builder, m_DiffReq, ForwardModeIASL,
-                m_DiffReq.Args, m_Builder.m_DFC);
+                m_DiffReq.Args, m_Builder.m_Scheduler.getDerivedFns());
           secondDerivativeFuncs.push_back(DFD);
         }
       } else {
@@ -201,13 +201,13 @@ DerivativeAndOverload HessianModeVisitor::Derive() {
             utils::CreateStringLiteral(m_Context, PVD->getNameAsString());
         FunctionDecl* DFD = nullptr;
         if (m_DiffReq.Mode == DiffMode::hessian_diagonal)
-          DFD = DeriveUsingForwardModeTwice(m_Sema, m_CladPlugin, m_Builder,
-                                            m_DiffReq, ForwardModeIASL,
-                                            m_Builder.m_DFC);
+          DFD = DeriveUsingForwardModeTwice(
+              m_Sema, m_CladPlugin, m_Builder, m_DiffReq, ForwardModeIASL,
+              m_Builder.m_Scheduler.getDerivedFns());
         else
           DFD = DeriveUsingForwardAndReverseMode(
               m_Sema, m_CladPlugin, m_Builder, m_DiffReq, ForwardModeIASL,
-              m_DiffReq.Args, m_Builder.m_DFC);
+              m_DiffReq.Args, m_Builder.m_Scheduler.getDerivedFns());
         secondDerivativeFuncs.push_back(DFD);
       }
     }
diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp
@@ -378,7 +378,7 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
       FunctionDecl* FoundFD =
           R.empty() ? nullptr : dyn_cast<FunctionDecl>(R.front());
       if (!RD->isLambda() && !R.empty() &&
-          !m_Builder.m_DFC.IsCladDerivative(FoundFD)) {
+          !m_Builder.m_Scheduler.getDerivedFns().IsCladDerivative(FoundFD)) {
         Sema::NestedNameSpecInfo IdInfo(RD->getIdentifier(), noLoc, noLoc,
                                         /*ObjectType=*/nullptr);
         // FIXME: Address nested classes where SS should be set.
diff --git a/tools/ClangPlugin.cpp b/tools/ClangPlugin.cpp
@@ -206,9 +206,6 @@ void InitTimers();
       if (!CheckBuiltins())
         return;
 #if CLANG_VERSION_MAJOR > 16
-      Sema& S = m_CI.getSema();
-      RequestOptions opts{};
-      SetRequestOptions(opts);
       // Traverse all constexpr FunctionDecls for the static graph only once to
       // differentiate them immeditely.
       {
@@ -218,19 +215,17 @@ void InitTimers();
             continue;
           auto* FD = cast<FunctionDecl>(D);
           if (FD->isConstexpr() || !m_Multiplexer) {
-            DiffCollector collector(CladEnabledRange, m_DiffRequestGraph, S,
-                                    opts, m_AllAnalysisDC);
-            collector.Walk(DGR);
+            getScheduler().Plan(DGR);
             break;
           }
         }
       }
 
-      for (DiffRequest& request : m_DiffRequestGraph.getNodes()) {
+      for (DiffRequest& request : getScheduler().getGraph().getNodes()) {
         if (request.ImmediateMode && request.Function->isConstexpr()) {
-          m_DiffRequestGraph.setCurrentProcessingNode(request);
+          getScheduler().getGraph().setCurrentProcessingNode(request);
           ProcessDiffRequest(request);
-          m_DiffRequestGraph.markCurrentNodeProcessed();
+          getScheduler().getGraph().markCurrentNodeProcessed();
         }
       }
 #endif
@@ -349,8 +344,8 @@ void InitTimers();
     FunctionDecl* CladPlugin::ProcessDiffRequest(DiffRequest& request) {
       Sema& S = m_CI.getSema();
       if (!m_DerivativeBuilder)
-        m_DerivativeBuilder = std::make_unique<DerivativeBuilder>(
-            S, *this, m_DFC, m_DiffRequestGraph);
+        m_DerivativeBuilder =
+            std::make_unique<DerivativeBuilder>(S, *this, getScheduler());
 
       if (request.Global) {
         auto deriveResult = m_DerivativeBuilder->Derive(request);
@@ -366,7 +361,7 @@ void InitTimers();
       // FIXME: These requests are not fully generated in the diffplanner and we
       // have to update diff params on this stage.
       if (request.CurrentDerivativeOrder > 1 ||
-          m_DFC.IsCladDerivative(request.Function))
+          getScheduler().getDerivedFns().IsCladDerivative(request.Function))
         request.UpdateDiffParamsInfo(m_CI.getSema());
       const FunctionDecl* FD = request.Function;
       ASTContext& C = S.getASTContext();
@@ -397,7 +392,7 @@ void InitTimers();
       {
         llvm::SaveAndRestore<unsigned> Saved(request.RequestedDerivativeOrder,
                                              1);
-        auto DFI = m_DFC.Find(request);
+        auto DFI = getScheduler().getDerivedFns().Find(request);
         if (DFI.IsValid()) {
           DerivativeDecl = DFI.DerivedFn();
           OverloadedDerivativeDecl = DFI.OverloadedDerivedFn();
@@ -413,8 +408,8 @@ void InitTimers();
               utils::hasEmptyBody(DerivativeDecl))
             return nullptr;
           if (DerivativeDecl)
-            m_DFC.Add(DerivedFnInfo(request, DerivativeDecl,
-                                    OverloadedDerivativeDecl));
+            getScheduler().getDerivedFns().Add(DerivedFnInfo(
+                request, DerivativeDecl, OverloadedDerivativeDecl));
         }
       }
 
@@ -595,6 +590,16 @@ void InitTimers();
       SetUsefulAnalysisOptions(m_DO, opts);
     }
 
+    DiffScheduler& CladPlugin::getScheduler() {
+      if (!m_Scheduler) {
+        RequestOptions Opts{};
+        SetRequestOptions(Opts);
+        m_Scheduler = std::make_unique<DiffScheduler>(m_CI.getSema(), Opts,
+                                                      CladEnabledRange);
+      }
+      return *m_Scheduler;
+    }
+
     void CladPlugin::FinalizeTranslationUnit() {
       Sema& S = m_CI.getSema();
       // Restore the TUScope that became a 0 in Sema::ActOnEndOfTranslationUnit.
@@ -606,16 +611,16 @@ void InitTimers();
       Sema::LocalEagerInstantiationScope LocalInstantiations(
           S CLAD_COMPAT_CLANG21_AtEndOfTUParam);
 
-      if (!m_DiffRequestGraph.isProcessingNode()) {
+      if (!getScheduler().getGraph().isProcessingNode()) {
         // This check is to avoid recursive processing of the graph, as
         // HandleTopLevelDecl can be called recursively in non-standard
         // setup for code generation.
-        DiffRequest request = m_DiffRequestGraph.getNextToProcessNode();
+        DiffRequest request = getScheduler().getGraph().getNextToProcessNode();
         while (request.Function || request.Global) {
-          m_DiffRequestGraph.setCurrentProcessingNode(request);
+          getScheduler().getGraph().setCurrentProcessingNode(request);
           ProcessDiffRequest(request);
-          m_DiffRequestGraph.markCurrentNodeProcessed();
-          request = m_DiffRequestGraph.getNextToProcessNode();
+          getScheduler().getGraph().markCurrentNodeProcessed();
+          request = getScheduler().getGraph().getNextToProcessNode();
         }
       }
 
@@ -631,9 +636,6 @@ void InitTimers();
     void CladPlugin::HandleTranslationUnit(ASTContext& C) {
       // In case of diagnostics, don't bother, just let the compiler finish.
       if (!m_CI.getDiagnostics().hasErrorOccurred()) {
-        Sema& S = m_CI.getSema();
-        RequestOptions opts{};
-        SetRequestOptions(opts);
         // Traverse all collected DeclGroupRef only once to create the static
         // graph.
         TimedAnalysisRegion R("Rest of TU");
@@ -642,16 +644,14 @@ void InitTimers();
             if (const auto* FD = dyn_cast<FunctionDecl>(D))
               if (FD->isConstexpr())
                 continue;
-            DiffCollector collector(CladEnabledRange, m_DiffRequestGraph, S,
-                                    opts, m_AllAnalysisDC);
-            collector.Walk(DCI.m_DGR);
+            getScheduler().Plan(DCI.m_DGR);
             break;
           }
 
         if (m_CI.getFrontendOpts().ShowStats) {
           // Print the graph of the diff requests.
           llvm::errs() << "\n*** INFORMATION ABOUT THE DIFF REQUESTS\n";
-          m_DiffRequestGraph.dump();
+          getScheduler().getGraph().dump();
         }
 
         FinalizeTranslationUnit();
diff --git a/tools/ClangPlugin.h b/tools/ClangPlugin.h
@@ -11,6 +11,7 @@
 #include "clad/Differentiator/DerivedFnCollector.h"
 #include "clad/Differentiator/DiffMode.h"
 #include "clad/Differentiator/DiffPlanner.h"
+#include "clad/Differentiator/DiffScheduler.h"
 #include "clad/Differentiator/Version.h"
 
 #include "clang/AST/Decl.h"
@@ -97,9 +98,9 @@ struct DifferentiationOptions {
     DifferentiationOptions m_DO;
     std::unique_ptr<DerivativeBuilder> m_DerivativeBuilder;
     bool m_HasRuntime = false;
-    DerivedFnCollector m_DFC;
-    DynamicGraph<DiffRequest> m_DiffRequestGraph;
-    OwnedAnalysisContexts m_AllAnalysisDC;
+    /// Lazily constructed because it needs Sema, which is not available
+    /// until InitializeSema; reach it through getScheduler().
+    std::unique_ptr<DiffScheduler> m_Scheduler;
     enum class CallKind {
       HandleCXXStaticMemberVarInstantiation,
       HandleTopLevelDecl,
@@ -168,8 +169,9 @@ struct DifferentiationOptions {
           // setup, we exit early to give control to the non-standard setup for
           // code generation.
           // FIXME: This should go away if Cling starts using the clang driver.
-          if (!m_Multiplexer &&
-              (m_DFC.IsCladDerivative(FD) || m_DFC.IsCustomDerivative(FD)))
+          if (!m_Multiplexer && m_Scheduler &&
+              (m_Scheduler->getDerivedFns().IsCladDerivative(FD) ||
+               m_Scheduler->getDerivedFns().IsCustomDerivative(FD)))
             return true;
 
       HandleTopLevelDeclForClad(D);
@@ -252,6 +254,7 @@ struct DifferentiationOptions {
     clang::FunctionDecl* ProcessDiffRequest(DiffRequest& request);
 
   private:
+    DiffScheduler& getScheduler();
     void AppendDelayed(DelayedCallInfo DCI) {
       // Incremental processing handles the translation unit in chunks and it is
       // expected to have multiple calls to this functionality.
