Keep track of whether a request is immediate in DiffRequest

vgvassilev · Oct 23, 2024 · fa19625 · fa19625
1 parent 767eb13
commit fa19625
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/include/clad/Differentiator/DiffPlanner.h b/include/clad/Differentiator/DiffPlanner.h
@@ -65,6 +65,9 @@ struct DiffRequest {
   /// A flag to enable TBR analysis during reverse-mode differentiation.
   bool EnableTBRAnalysis = false;
   bool EnableVariedAnalysis = false;
+  /// A flag specifying whether this differentiation is to be used
+  /// in immediate contexts.
+  bool ImmediateMode = false;
   /// Puts the derived function and its code in the diff call
   void updateCall(clang::FunctionDecl* FD, clang::FunctionDecl* OverloadedFD,
                   clang::Sema& SemaRef);

diff --git a/lib/Differentiator/DiffPlanner.cpp b/lib/Differentiator/DiffPlanner.cpp
@@ -747,6 +747,8 @@ namespace clad {
         request.RequestedDerivativeOrder = derivative_order;
         if (clad::HasOption(bitmasked_opts_value, clad::opts::use_enzyme))
           request.use_enzyme = true;
+        if (clad::HasOption(bitmasked_opts_value, clad::opts::immediate_mode))
+          request.ImmediateMode = true;
         if (enable_tbr_in_req) {
           utils::EmitDiag(m_Sema, DiagnosticsEngine::Error, endLoc,
                           "TBR analysis is not meant for forward mode AD.");

diff --git a/tools/ClangPlugin.cpp b/tools/ClangPlugin.cpp
@@ -138,15 +138,18 @@ namespace clad {
       DiffCollector collector(DGR, CladEnabledRange, m_DiffRequestGraph, S,
                               opts);
 
+#if CLANG_VERSION_MAJOR > 16
       for (DiffRequest& request : m_DiffRequestGraph.getNodes()) {
-        if (!request.Function->isImmediateFunction() &&
-            !request.Function->isConstexpr())
+        if (!request.ImmediateMode ||
+            (!request.Function->isConstexpr() &&
+             !request.Function->isImmediateFunction()))
           continue;
 
         m_DiffRequestGraph.setCurrentProcessingNode(request);
         ProcessDiffRequest(request);
         m_DiffRequestGraph.markCurrentNodeProcessed();
       }
+#endif
 
       // We could not delay the processing of derivatives, inform act as if each
       // call is final. That would still have vgvassilev/clad#248 unresolved.