Compute and process Differentiation Request graph #873

vaithak · 2024-04-26T16:16:13Z

Plan for dynamic graph

The relations between different differentiation requests can be modelled as a graph.
For example, if f_a calls f_b, there will be two differentiation requests df_a and df_b, the edge between them can be understood as created_because_of.
This also means that the functions called by the users to be explicitly differentiated (or DiffRequests created because of these) are the source nodes, i.e. no incoming edges. In most cases, this graph aligns with the call graph, but in some cases, the graph depends on the internal implementation, like the Hessian computation, which requires creating multiple fwd_mode requests followed by a rev_mode request.
We can use this graph to order the computation of differentiation requests. This was already being done implicitly in the initial recursive implementation. Whenever we encountered a call expression, we started differentiation of the
called function; this was sort of like a depth-first search strategy.
- This had problems, as Clang reported errors when it encountered a new function scope (of the derivative of the called function) in the middle of the old function scope (of the derivative of the callee function). It treated the nested one like a lambda expression. The issue regarding this: Possible memory leak when differentiating call expressions in forward mode #745.
To fix this, an initial strategy was to eliminate the recursive approach. Hence, a queue-based breadth-first approach was implemented in this PR: Restructure differentiation schedule into a breadth first traversal #848.
- Although it fixed the problem, the graph traversal was still implicit. We needed some way to compute/store the complete graph and possibly optimize it, such as converting edges to model the requires_derivative_of relation. Using this, we could proceed with differentiation in a topologically sorted ordering.
- It also required one caveat: although we don't differentiate the called function completely in a recursive way, we still need to declare it so that we can have the call expression completed (i.e. auto t0 = f_pushforward(...)).
To move towards the final stage of having the complete graph computed before starting the differentiation, we need the complete information on how the DiffRequest will be formed inside the visitors (including arguments or DVI info).
This whole approach will require activity analysis in the first pass.
- As an incremental improvement, the first requirement was to implement infrastructure to support explicit modelling of the graph and use that to have a breadth-first traversal (and eventually topological ordering).

This is the initial PR for capturing the differentiation plan in a graphical format.

However, the traversal order is still breadth-first, as we don't have the complete graph in the first pass - mainly because of a lack of information about the args required for pushforward and pullbacks.

This can be improved with the help of activity analysis to capture the complete graph in the first pass, processing the plan in a topologically sorted manner and pruning the graph for user-defined functions. I started this with this approach, and the initial experimental commit is available here for future reference: vaithak@82c0b42.

vaithak · 2024-04-26T16:18:17Z

@vgvassilev One improvement that we can make in this PR itself is to fix the printing order of generated functions; this can be done because we have dynamically created the graph during traversals.
This will help us minimize (or, in some cases, completely remove) the printing of declarations first, and definitions later.
Do you think it would be a good idea?

codecov · 2024-04-26T16:23:19Z

Codecov Report

Attention: Patch coverage is 99.15966% with 1 lines in your changes are missing coverage. Please review.

Project coverage is 94.81%. Comparing base (922bd1c) to head (ed6a89a).

Additional details and impacted files

@@           Coverage Diff           @@
##           master     #873   +/-   ##
=======================================
  Coverage   94.81%   94.81%           
=======================================
  Files          51       52    +1     
  Lines        7620     7701   +81     
=======================================
+ Hits         7225     7302   +77     
- Misses        395      399    +4

Files	Coverage Δ
include/clad/Differentiator/DerivativeBuilder.h	`100.00% <ø> (ø)`
include/clad/Differentiator/DerivedFnCollector.h	`100.00% <ø> (ø)`
include/clad/Differentiator/DiffPlanner.h	`100.00% <100.00%> (ø)`
include/clad/Differentiator/Differentiator.h	`100.00% <ø> (ø)`
lib/Differentiator/BaseForwardModeVisitor.cpp	`98.90% <100.00%> (+0.09%)`	⬆️
lib/Differentiator/DerivativeBuilder.cpp	`99.00% <100.00%> (+0.02%)`	⬆️
lib/Differentiator/DerivedFnCollector.cpp	`100.00% <100.00%> (ø)`
lib/Differentiator/DiffPlanner.cpp	`98.64% <100.00%> (+<0.01%)`	⬆️
lib/Differentiator/HessianModeVisitor.cpp	`99.49% <100.00%> (+0.02%)`	⬆️
lib/Differentiator/ReverseModeVisitor.cpp	`97.18% <100.00%> (ø)`
... and 3 more

... and 2 files with indirect coverage changes

Files	Coverage Δ
include/clad/Differentiator/DerivativeBuilder.h	`100.00% <ø> (ø)`
include/clad/Differentiator/DerivedFnCollector.h	`100.00% <ø> (ø)`
include/clad/Differentiator/DiffPlanner.h	`100.00% <100.00%> (ø)`
include/clad/Differentiator/Differentiator.h	`100.00% <ø> (ø)`
lib/Differentiator/BaseForwardModeVisitor.cpp	`98.90% <100.00%> (+0.09%)`	⬆️
lib/Differentiator/DerivativeBuilder.cpp	`99.00% <100.00%> (+0.02%)`	⬆️
lib/Differentiator/DerivedFnCollector.cpp	`100.00% <100.00%> (ø)`
lib/Differentiator/DiffPlanner.cpp	`98.64% <100.00%> (+<0.01%)`	⬆️
lib/Differentiator/HessianModeVisitor.cpp	`99.49% <100.00%> (+0.02%)`	⬆️
lib/Differentiator/ReverseModeVisitor.cpp	`97.18% <100.00%> (ø)`
... and 3 more

... and 2 files with indirect coverage changes

github-actions

clang-tidy made some suggestions

There were too many comments to post at once. Showing the first 10 out of 18. Check the log or trigger a new build to see more.

include/clad/Differentiator/DiffPlanner.h

include/clad/Differentiator/DynamicGraph.h

lib/Differentiator/BaseForwardModeVisitor.cpp

include/clad/Differentiator/DiffPlanner.h

vgvassilev · 2024-04-26T17:19:01Z

@vgvassilev One improvement that we can make in this PR itself is to fix the printing order of generated functions; this can be done because we have dynamically created the graph during traversals. This will help us minimize (or, in some cases, completely remove) the printing of declarations first, and definitions later. Do you think it would be a good idea?

I think so, is there a catch?

vaithak · 2024-04-27T17:20:31Z

I think so, is there a catch?

The only issue I see is that diagnostic messages won't align with the sequence of generated functions in the code dump. So, maybe it's a little bit tough for debugging.

vgvassilev · 2024-04-27T18:11:03Z

I think so, is there a catch?

The only issue I see is that diagnostic messages won't align with the sequence of generated functions in the code dump. So, maybe it's a little bit tough for debugging.

Hm... I am confused - how is it different from now?

vaithak · 2024-04-27T18:45:43Z

Hm... I am confused - how is it different from now?

Currently, if the order of differentiation is f1 followed by f2. Then, any diagnostic warning or error in f2 will appear after f1 has been dumped/printed and before f2 is dumped.
This won't be true anymore.

vgvassilev · 2024-04-27T18:56:26Z

Hm... I am confused - how is it different from now?

Currently, if the order of differentiation is f1 followed by f2. Then, any diagnostic warning or error in f2 will appear after f1 has been dumped/printed and before f2 is dumped. This won't be true anymore.

Let's have this improvement in a separate PR and then we can discuss further. What do you think?

vaithak · 2024-04-28T09:32:26Z

Let's have this improvement in a separate PR and then we can discuss further. What do you think?

Sounds good 👍🏼

github-actions

clang-tidy made some suggestions

github-actions · 2024-04-30T12:47:19Z

tools/ClangPlugin.cpp

+        m_DerivativeBuilder.reset(
+            new DerivativeBuilder(S, *this, m_DFC, m_DiffRequestGraph));


warning: use std::make_unique instead [modernize-make-unique]

Suggested change

m_DerivativeBuilder.reset(

new DerivativeBuilder(S, *this, m_DFC, m_DiffRequestGraph));

m_DerivativeBuilder = std::make_unique<DerivativeBuilder>(

S, *this, m_DFC, m_DiffRequestGraph);

github-actions · 2024-04-30T12:47:19Z

tools/ClangPlugin.cpp

@@ -122,11 +122,13 @@
      Sema& S = m_CI.getSema();

      if (!m_DerivativeBuilder)
-        m_DerivativeBuilder.reset(new DerivativeBuilder(S, *this, m_DFC));
+        m_DerivativeBuilder.reset(
+            new DerivativeBuilder(S, *this, m_DFC, m_DiffRequestGraph));


warning: initializing non-owner argument of type 'std::unique_ptrclad::DerivativeBuilder::pointer' (aka 'clad::DerivativeBuilder *') with a newly created 'gsl::owner<>' [cppcoreguidelines-owning-memory]

new DerivativeBuilder(S, *this, m_DFC, m_DiffRequestGraph)); ^

include/clad/Differentiator/DynamicGraph.h

vgvassilev · 2024-05-01T13:05:13Z

lib/Differentiator/DiffPlanner.cpp

@@ -684,7 +685,7 @@ namespace clad {
      llvm::SaveAndRestore<const FunctionDecl*> saveTopMost = m_TopMostFD;
      m_TopMostFD = FD;
      TraverseDecl(derivedFD);
-      m_DiffPlans.push_back(std::move(request));
+      m_DiffRequestGraph.addNode(request, true /*isSource*/);


Suggested change

m_DiffRequestGraph.addNode(request, true /*isSource*/);

m_DiffRequestGraph.addNode(request, /*isSource=*/true);

There must be a clang-tidy check, why it did not warn here - can you check?

Even after a lot of trials, I can't figure it out. I tried this option with multiple configurations of its options: https://clang.llvm.org/extra/clang-tidy/checks/bugprone/argument-comment.html.
Still doesn't work, changed it manually for now.

test/Misc/DynamicGraph.C

vgvassilev · 2024-05-01T13:09:19Z

lib/Differentiator/HessianModeVisitor.cpp

+    if (!firstDerivative) {
+      // Derive declaration of the the forward mode derivative.
+      IndependentArgRequest.DeclarationOnly = true;
+      firstDerivative = plugin::ProcessDiffRequest(CP, IndependentArgRequest);


This PR does not get rid of plugin::ProcessDiffRequest, but do we have a plan to remove it. I think that's a layering violation.

The only reason we need this is for the forward declaration of derivative signatures to complete the call expressions. Once we get the graph in the first pass itself (with the help of activity analysis), we can proceed with topologically ordered differentiation (instead of breadth-first as we are doing now), and then this forward declaration will not be required.

Can we keep track of that discussion as an issue?

We already have a related issue here: #857, added a comment in it referencing this PR.

include/clad/Differentiator/DynamicGraph.h

vgvassilev

LGTM!

Plan for dynamic graph - The relations between different differentiation requests can be modelled as a graph. For example, if `f_a` calls `f_b`, there will be two differentiation requests `df_a` and `df_b`, the edge between them can be understood as `created_because_of`. This also means that the functions called by the users to be explicitly differentiated (or `DiffRequests` created because of these) are the source nodes, i.e. no incoming edges. In most cases, this graph aligns with the call graph, but in some cases, the graph depends on the internal implementation, like the Hessian computation, which requires creating multiple `fwd_mode` requests followed by a `rev_mode` request. - We can use this graph to order the computation of differentiation requests. This was already being done implicitly in the initial recursive implementation. Whenever we encountered a call expression, we started differentiation of the called function; this was sort of like a depth-first search strategy. - This had problems, as `Clang` reported errors when it encountered a new function scope (of the derivative of the called function) in the middle of the old function scope (of the derivative of the callee function). It treated the nested one like a lambda expression. The issue regarding this: vgvassilev#745. - To fix this, an initial strategy was to eliminate the recursive approach. Hence, a queue-based breadth-first approach was implemented in this PR: vgvassilev#848. - Although it fixed the problem, the graph traversal was still implicit. We needed some way to compute/store the complete graph and possibly optimize it, such as converting edges to model the `requires_derivative_of` relation. Using this, we could proceed with differentiation in a topologically sorted ordering. - It also required one caveat: although we don't differentiate the called function completely in a recursive way, we still need to declare it so that we can have the call expression completed (i.e. `auto t0 = f_pushforward(...)`). - To move towards the final stage of having the complete graph computed before starting the differentiation, we need the complete information on how the `DiffRequest` will be formed inside the visitors (including arguments or `DVI` info). This whole approach will require activity analysis in the first pass. - As an incremental improvement, the first requirement was to implement infrastructure to support explicit modelling of the graph and use that to have a breadth-first traversal (and eventually topological ordering). This is the initial PR for capturing the differentiation plan in a graphical format. However, the traversal order is still breadth-first, as we don't have the complete graph in the first pass - mainly because of a lack of information about the args required for `pushforward` and `pullbacks`. This can be improved with the help of activity analysis to capture the complete graph in the first pass, processing the plan in a topologically sorted manner and pruning the graph for user-defined functions. I started this with this approach, and the initial experimental commit is available here for future reference: 82c0b42.

github-actions bot reviewed Apr 26, 2024

View reviewed changes

vgvassilev reviewed Apr 26, 2024

View reviewed changes

include/clad/Differentiator/DiffPlanner.h Show resolved Hide resolved

vaithak force-pushed the improved-diff-plan branch 3 times, most recently from 153c3a1 to 66beec3 Compare April 30, 2024 12:27

github-actions bot reviewed Apr 30, 2024

View reviewed changes

vaithak force-pushed the improved-diff-plan branch from 66beec3 to b2ed85b Compare April 30, 2024 12:57

vgvassilev reviewed May 1, 2024

View reviewed changes

vgvassilev reviewed May 3, 2024

View reviewed changes

include/clad/Differentiator/DynamicGraph.h Outdated Show resolved Hide resolved

include/clad/Differentiator/DynamicGraph.h Outdated Show resolved Hide resolved

vaithak force-pushed the improved-diff-plan branch from 1a553a5 to 1ee7b35 Compare May 3, 2024 14:11

vgvassilev approved these changes May 3, 2024

View reviewed changes

vaithak mentioned this pull request May 3, 2024

Avoid access of plugin methods from Visitors #857

Open

vaithak added 2 commits May 3, 2024 17:02

Bring back DerivativeSet and use it to fix Hessian ordering

ed6a89a

vaithak force-pushed the improved-diff-plan branch from 1ee7b35 to ed6a89a Compare May 3, 2024 15:03

vgvassilev merged commit d32154c into vgvassilev:master May 3, 2024
89 checks passed

mcbarton mentioned this pull request Jul 31, 2024

Add llvm 18 for osx to ci #889

Draft

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Compute and process Differentiation Request graph #873

Compute and process Differentiation Request graph #873

vaithak commented Apr 26, 2024 •

edited

Loading

vaithak commented Apr 26, 2024 •

edited

Loading

codecov bot commented Apr 26, 2024 •

edited

Loading

github-actions bot left a comment

vgvassilev commented Apr 26, 2024

vaithak commented Apr 27, 2024

vgvassilev commented Apr 27, 2024

vaithak commented Apr 27, 2024

vgvassilev commented Apr 27, 2024

vaithak commented Apr 28, 2024

github-actions bot left a comment

github-actions bot Apr 30, 2024

github-actions bot Apr 30, 2024

vgvassilev May 1, 2024

vaithak May 3, 2024

vgvassilev May 1, 2024

vaithak May 3, 2024

vgvassilev May 3, 2024

vaithak May 3, 2024

vgvassilev left a comment

		m_DerivativeBuilder.reset(
		new DerivativeBuilder(S, *this, m_DFC, m_DiffRequestGraph));

	m_DiffRequestGraph.addNode(request, true /isSource/);
	m_DiffRequestGraph.addNode(request, /isSource=/true);

Compute and process Differentiation Request graph #873

Compute and process Differentiation Request graph #873

Conversation

vaithak commented Apr 26, 2024 • edited Loading

Plan for dynamic graph

vaithak commented Apr 26, 2024 • edited Loading

codecov bot commented Apr 26, 2024 • edited Loading

Codecov Report

github-actions bot left a comment

Choose a reason for hiding this comment

vgvassilev commented Apr 26, 2024

vaithak commented Apr 27, 2024

vgvassilev commented Apr 27, 2024

vaithak commented Apr 27, 2024

vgvassilev commented Apr 27, 2024

vaithak commented Apr 28, 2024

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot Apr 30, 2024

Choose a reason for hiding this comment

github-actions bot Apr 30, 2024

Choose a reason for hiding this comment

vgvassilev May 1, 2024

Choose a reason for hiding this comment

vaithak May 3, 2024

Choose a reason for hiding this comment

vgvassilev May 1, 2024

Choose a reason for hiding this comment

vaithak May 3, 2024

Choose a reason for hiding this comment

vgvassilev May 3, 2024

Choose a reason for hiding this comment

vaithak May 3, 2024

Choose a reason for hiding this comment

vgvassilev left a comment

Choose a reason for hiding this comment

vaithak commented Apr 26, 2024 •

edited

Loading

vaithak commented Apr 26, 2024 •

edited

Loading

codecov bot commented Apr 26, 2024 •

edited

Loading