Reimplement jacobians using the vectorized forward mode #1121

PetroZarytskyi · 2024-10-22T22:30:52Z

Previously, jacobians were based on the non-vectorized reverse mode, which was mostly incapable of capturing multiple outputs. The implementation worked in a few particular cases. For example, it was not possible to differentiate function calls or declare variables inside the original function body.
This PR implements jacobians using the vectorized forward mode. At the very least, this will solve the issues described above and give a way forward to solve other ones. This also means introducing features to the vectorized fwd mode will introduce the same features to jacobians.
Let's take a look at the new signature of jacobians. Since the function to be differentiated is expected to have multiple outputs, we should expect the output in the form of array/pointer/reference parameters (just like before). And for every output parameter, we should generate a corresponding adjoint parameter for the user to acquire the results. Since there is no way to specify which parameters are used as output and which are not, adjoints are generated for all array/pointer/reference parameters. For example:

void f(double a, double b, double* c)  
 --> 
void f_jac(double a, double b, double* c, array_ref<matrix<double>> _d_c)

or

void f(double a, double b, double* c, double[7] t) 
 -->
void f_jac(double a, double b, double* c, double[7] t,
 array_ref<matrix<double>> _d_c, array_ref<matrix<double>> _d_t)

array_ref is necessary for compatibility with the existing infrastructure for vectorized fwd mode overloads generation.
This signature is also similar to the old one. e.g.

df.execute(a, b, c, result); // old behavior
df.execute(a, b, c, &result); // new behavior

However, the behavior differs for multiple output parameters, which the old jacobians did not support.

Note: the same functionality can be achieved by using the vectorized reverse mode, which should probably be implemented at some point. However, the old code for jacobians is unlikely to be useful for that, and there is not much point in keeping it.

codecov · 2024-10-22T22:50:06Z

Codecov Report

Attention: Patch coverage is 88.26291% with 25 lines in your changes missing coverage. Please review.

Project coverage is 94.03%. Comparing base (d82f7fd) to head (70fcedc).

Files with missing lines	Patch %	Lines
lib/Differentiator/JacobianModeVisitor.cpp	86.62%	23 Missing ⚠️
lib/Differentiator/BaseForwardModeVisitor.cpp	90.00%	1 Missing ⚠️
lib/Differentiator/VectorForwardModeVisitor.cpp	83.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1121      +/-   ##
==========================================
- Coverage   94.37%   94.03%   -0.34%     
==========================================
  Files          50       51       +1     
  Lines        8366     8798     +432     
==========================================
+ Hits         7895     8273     +378     
- Misses        471      525      +54

Files with missing lines	Coverage Δ
lib/Differentiator/DerivativeBuilder.cpp	`100.00% <100.00%> (ø)`
lib/Differentiator/DiffPlanner.cpp	`98.59% <100.00%> (+0.03%)`	⬆️
lib/Differentiator/ReverseModeVisitor.cpp	`95.06% <100.00%> (-0.45%)`	⬇️
lib/Differentiator/VisitorBase.cpp	`97.28% <100.00%> (+0.14%)`	⬆️
lib/Differentiator/BaseForwardModeVisitor.cpp	`98.68% <90.00%> (-0.06%)`	⬇️
lib/Differentiator/VectorForwardModeVisitor.cpp	`99.69% <83.33%> (-0.31%)`	⬇️
lib/Differentiator/JacobianModeVisitor.cpp	`86.62% <86.62%> (ø)`

... and 27 files with indirect coverage changes

Files with missing lines	Coverage Δ
lib/Differentiator/DerivativeBuilder.cpp	`100.00% <100.00%> (ø)`
lib/Differentiator/DiffPlanner.cpp	`98.59% <100.00%> (+0.03%)`	⬆️
lib/Differentiator/ReverseModeVisitor.cpp	`95.06% <100.00%> (-0.45%)`	⬇️
lib/Differentiator/VisitorBase.cpp	`97.28% <100.00%> (+0.14%)`	⬆️
lib/Differentiator/BaseForwardModeVisitor.cpp	`98.68% <90.00%> (-0.06%)`	⬇️
lib/Differentiator/VectorForwardModeVisitor.cpp	`99.69% <83.33%> (-0.31%)`	⬇️
lib/Differentiator/JacobianModeVisitor.cpp	`86.62% <86.62%> (ø)`

... and 27 files with indirect coverage changes

github-actions

clang-tidy made some suggestions

github-actions · 2024-10-22T22:58:55Z

include/clad/Differentiator/FunctionTraits.h

-  struct JacobianDerivedFnTraits<R (C::*)(Args...) cv vol ref noex> {          \
-    using type = void (C::*)(Args..., SelectLast_t<Args...>) cv vol ref noex;  \
-  };
+#define JacobianDerivedFnTraits_AddSPECS(var, cv, vol, ref, noex)            \


warning: function-like macro 'JacobianDerivedFnTraits_AddSPECS' used; consider a 'constexpr' template function [cppcoreguidelines-macro-usage]

#define JacobianDerivedFnTraits_AddSPECS(var, cv, vol, ref, noex) \ ^

github-actions · 2024-10-22T22:58:55Z

lib/Differentiator/BaseForwardModeVisitor.cpp

+        clad::utils::ComputeEffectiveFnName(FD) + "_pushforward";
+    callDiff = m_Builder.BuildCallToCustomDerivativeOrNumericalDiff(
+        customPushforward, customDerivativeArgs, getCurrentScope(),
+        const_cast<DeclContext*>(FD->getDeclContext()));


warning: do not use const_cast [cppcoreguidelines-pro-type-const-cast]

const_cast<DeclContext*>(FD->getDeclContext())); ^

github-actions · 2024-10-22T22:58:56Z

lib/Differentiator/JacobianModeVisitor.cpp

+  assert(m_DiffReq.Mode == DiffMode::jacobian);
+
+  DiffParams args{};
+  for (auto dParam : m_DiffReq.DVI)


warning: loop variable is copied but only used as const reference; consider making it a const reference [performance-for-range-copy]

Suggested change

for (auto dParam : m_DiffReq.DVI)

for (const auto& dParam : m_DiffReq.DVI)

github-actions · 2024-10-22T22:58:56Z

lib/Differentiator/JacobianModeVisitor.cpp

+  // Generate name for the derivative function.
+  std::string derivedFnName = m_DiffReq.BaseFunctionName + "_jac";
+  if (args.size() != FD->getNumParams()) {
+    for (auto arg : args) {


warning: 'auto arg' can be declared as 'const auto *arg' [llvm-qualified-auto]

Suggested change

for (auto arg : args) {

for (const auto *arg : args) {

github-actions · 2024-10-22T22:58:56Z

lib/Differentiator/JacobianModeVisitor.cpp

+  std::string derivedFnName = m_DiffReq.BaseFunctionName + "_jac";
+  if (args.size() != FD->getNumParams()) {
+    for (auto arg : args) {
+      auto it = std::find(FD->param_begin(), FD->param_end(), arg);


warning: 'auto it' can be declared as 'const auto *it' [llvm-qualified-auto]

Suggested change

auto it = std::find(FD->param_begin(), FD->param_end(), arg);

const auto *it = std::find(FD->param_begin(), FD->param_end(), arg);

github-actions · 2024-10-22T22:58:56Z

lib/Differentiator/JacobianModeVisitor.cpp

+  for (size_t i = 0; i < numParamsOriginalFn; ++i) {
+    bool is_array =
+        utils::isArrayOrPointerType(m_DiffReq->getParamDecl(i)->getType());
+    auto param = params[i];


warning: 'auto param' can be declared as 'auto *param' [llvm-qualified-auto]

Suggested change

auto param = params[i];

auto *param = params[i];

github-actions · 2024-10-22T22:58:56Z

lib/Differentiator/JacobianModeVisitor.cpp

+
+  // Traverse the function body and generate the derivative.
+  Stmt* BodyDiff = Visit(FD->getBody()).getStmt();
+  if (auto CS = dyn_cast<CompoundStmt>(BodyDiff))


warning: 'auto CS' can be declared as 'auto *CS' [llvm-qualified-auto]

Suggested change

if (auto CS = dyn_cast<CompoundStmt>(BodyDiff))

if (auto *CS = dyn_cast<CompoundStmt>(BodyDiff))

github-actions · 2024-10-22T22:58:56Z

lib/Differentiator/VectorForwardModeVisitor.cpp

-  DiffInputVarsInfo DVI;
-  DVI = request.DVI;
-  for (auto dParam : DVI)
+  for (auto dParam : m_DiffReq.DVI)


warning: loop variable is copied but only used as const reference; consider making it a const reference [performance-for-range-copy]

Suggested change

for (auto dParam : m_DiffReq.DVI)

for (const auto& dParam : m_DiffReq.DVI)

…orized frw mode and jacobians

github-actions

clang-tidy made some suggestions

github-actions · 2024-10-23T07:36:06Z

lib/Differentiator/JacobianModeVisitor.cpp

+  assert(m_DiffReq.Mode == DiffMode::jacobian);
+
+  DiffParams args{};
+  for (const DiffInputVarInfo& dParam : m_DiffReq.DVI)


warning: loop variable is copied but only used as const reference; consider making it a const reference [performance-for-range-copy]

Suggested change

for (const DiffInputVarInfo& dParam : m_DiffReq.DVI)

for (const auto& dParam : m_DiffReq.DVI)

github-actions · 2024-10-23T07:36:06Z

lib/Differentiator/JacobianModeVisitor.cpp

+  // Generate name for the derivative function.
+  std::string derivedFnName = m_DiffReq.BaseFunctionName + "_jac";
+  if (args.size() != FD->getNumParams()) {
+    for (const ValueDecl* arg : args) {


warning: 'auto arg' can be declared as 'const auto *arg' [llvm-qualified-auto]

Suggested change

for (const ValueDecl* arg : args) {

for (const auto *arg : args) {

github-actions · 2024-10-23T07:36:06Z

lib/Differentiator/JacobianModeVisitor.cpp

+  std::string derivedFnName = m_DiffReq.BaseFunctionName + "_jac";
+  if (args.size() != FD->getNumParams()) {
+    for (const ValueDecl* arg : args) {
+      auto* it = std::find(FD->param_begin(), FD->param_end(), arg);


warning: 'auto it' can be declared as 'const auto *it' [llvm-qualified-auto]

Suggested change

auto* it = std::find(FD->param_begin(), FD->param_end(), arg);

const auto *it = std::find(FD->param_begin(), FD->param_end(), arg);

github-actions · 2024-10-23T07:36:06Z

lib/Differentiator/JacobianModeVisitor.cpp

+  for (size_t i = 0; i < numParamsOriginalFn; ++i) {
+    bool is_array =
+        utils::isArrayOrPointerType(m_DiffReq->getParamDecl(i)->getType());
+    ParmVarDecl* param = params[i];


warning: 'auto param' can be declared as 'auto *param' [llvm-qualified-auto]

Suggested change

ParmVarDecl* param = params[i];

auto *param = params[i];

github-actions · 2024-10-23T07:36:06Z

lib/Differentiator/JacobianModeVisitor.cpp

+
+  // Traverse the function body and generate the derivative.
+  Stmt* BodyDiff = Visit(FD->getBody()).getStmt();
+  if (auto* CS = dyn_cast<CompoundStmt>(BodyDiff))


warning: 'auto CS' can be declared as 'auto *CS' [llvm-qualified-auto]

Suggested change

if (auto* CS = dyn_cast<CompoundStmt>(BodyDiff))

if (auto *CS = dyn_cast<CompoundStmt>(BodyDiff))

github-actions

clang-tidy made some suggestions

github-actions · 2024-10-23T07:54:48Z

lib/Differentiator/JacobianModeVisitor.cpp

+  std::string derivedFnName = m_DiffReq.BaseFunctionName + "_jac";
+  if (args.size() != FD->getNumParams()) {
+    for (const ValueDecl* arg : args) {
+      auto* it = std::find(FD->param_begin(), FD->param_end(), arg);


warning: 'auto *it' can be declared as 'const auto *it' [readability-qualified-auto]

Suggested change

auto* it = std::find(FD->param_begin(), FD->param_end(), arg);

const auto* it = std::find(FD->param_begin(), FD->param_end(), arg);

vgvassilev · 2024-10-23T08:59:34Z

Could we write a few benchmarks to make sure we do not regress in performance?

PetroZarytskyi added 2 commits October 22, 2024 23:41

Remove excessive FD and request parameters from DeriveVectorMode

40e4d88

Improve support for operations between clad::array and clad::array_ref

3ec9af2

PetroZarytskyi force-pushed the jac branch from a5af14c to f97dff1 Compare October 22, 2024 22:39

github-actions bot reviewed Oct 22, 2024

View reviewed changes

PetroZarytskyi added 5 commits October 23, 2024 09:33

Implement forward-mode jacobians

d6c48a5

Generalize custom derivative templates to be compatible with the vect…

a8aa5df

…orized frw mode and jacobians

Update the tests

14d9e58

Remove the old jacobian code from RMV

4a1749d

Always use valid location when generating operators

70fcedc

PetroZarytskyi force-pushed the jac branch from 42cbdd3 to 70fcedc Compare October 23, 2024 07:35

github-actions bot reviewed Oct 23, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Reimplement jacobians using the vectorized forward mode #1121

Reimplement jacobians using the vectorized forward mode #1121

PetroZarytskyi commented Oct 22, 2024

codecov bot commented Oct 22, 2024 •

edited

Loading

github-actions bot left a comment

github-actions bot Oct 22, 2024

github-actions bot Oct 22, 2024

github-actions bot Oct 22, 2024

github-actions bot Oct 22, 2024

github-actions bot Oct 22, 2024

github-actions bot Oct 22, 2024

github-actions bot Oct 22, 2024

github-actions bot Oct 22, 2024

github-actions bot left a comment

github-actions bot Oct 23, 2024

github-actions bot Oct 23, 2024

github-actions bot Oct 23, 2024

github-actions bot Oct 23, 2024

github-actions bot Oct 23, 2024

github-actions bot left a comment

github-actions bot Oct 23, 2024

vgvassilev commented Oct 23, 2024

	for (auto dParam : m_DiffReq.DVI)
	for (const auto& dParam : m_DiffReq.DVI)

	auto it = std::find(FD->param_begin(), FD->param_end(), arg);
	const auto *it = std::find(FD->param_begin(), FD->param_end(), arg);

	if (auto CS = dyn_cast<CompoundStmt>(BodyDiff))
	if (auto *CS = dyn_cast<CompoundStmt>(BodyDiff))

	for (const DiffInputVarInfo& dParam : m_DiffReq.DVI)
	for (const auto& dParam : m_DiffReq.DVI)

	for (const ValueDecl* arg : args) {
	for (const auto *arg : args) {

	auto* it = std::find(FD->param_begin(), FD->param_end(), arg);
	const auto *it = std::find(FD->param_begin(), FD->param_end(), arg);

	if (auto* CS = dyn_cast<CompoundStmt>(BodyDiff))
	if (auto *CS = dyn_cast<CompoundStmt>(BodyDiff))

	auto* it = std::find(FD->param_begin(), FD->param_end(), arg);
	const auto* it = std::find(FD->param_begin(), FD->param_end(), arg);

Reimplement jacobians using the vectorized forward mode #1121

Are you sure you want to change the base?

Reimplement jacobians using the vectorized forward mode #1121

Conversation

PetroZarytskyi commented Oct 22, 2024

codecov bot commented Oct 22, 2024 • edited Loading

Codecov Report

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot Oct 22, 2024

Choose a reason for hiding this comment

github-actions bot Oct 22, 2024

Choose a reason for hiding this comment

github-actions bot Oct 22, 2024

Choose a reason for hiding this comment

github-actions bot Oct 22, 2024

Choose a reason for hiding this comment

github-actions bot Oct 22, 2024

Choose a reason for hiding this comment

github-actions bot Oct 22, 2024

Choose a reason for hiding this comment

github-actions bot Oct 22, 2024

Choose a reason for hiding this comment

github-actions bot Oct 22, 2024

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot Oct 23, 2024

Choose a reason for hiding this comment

github-actions bot Oct 23, 2024

Choose a reason for hiding this comment

github-actions bot Oct 23, 2024

Choose a reason for hiding this comment

github-actions bot Oct 23, 2024

Choose a reason for hiding this comment

github-actions bot Oct 23, 2024

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot Oct 23, 2024

Choose a reason for hiding this comment

vgvassilev commented Oct 23, 2024

codecov bot commented Oct 22, 2024 •

edited

Loading