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, <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, 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.

Fixes #472

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

Codecov Report

Attention: Patch coverage is 96.09929% with 11 lines in your changes missing coverage. Please review.

Project coverage is 94.39%. Comparing base (cddc21d) to head (6b97654).

Files with missing lines	Patch %	Lines
lib/Differentiator/JacobianModeVisitor.cpp	93.97%	10 Missing ⚠️
lib/Differentiator/VectorForwardModeVisitor.cpp	83.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1121      +/-   ##
==========================================
- Coverage   94.44%   94.39%   -0.06%     
==========================================
  Files          50       51       +1     
  Lines        8770     8849      +79     
==========================================
+ Hits         8283     8353      +70     
- Misses        487      496       +9

Files with missing lines	Coverage Δ
include/clad/Differentiator/ReverseModeVisitor.h	`96.55% <ø> (-0.87%)`	⬇️
include/clad/Differentiator/VisitorBase.h	`100.00% <ø> (ø)`
lib/Differentiator/BaseForwardModeVisitor.cpp	`98.76% <100.00%> (+<0.01%)`	⬆️
lib/Differentiator/DerivativeBuilder.cpp	`100.00% <100.00%> (ø)`
lib/Differentiator/DiffPlanner.cpp	`98.63% <100.00%> (-0.01%)`	⬇️
lib/Differentiator/ReverseModeVisitor.cpp	`95.52% <100.00%> (-0.14%)`	⬇️
lib/Differentiator/VisitorBase.cpp	`97.63% <100.00%> (+0.36%)`	⬆️
lib/Differentiator/VectorForwardModeVisitor.cpp	`99.69% <83.33%> (-0.31%)`	⬇️
lib/Differentiator/JacobianModeVisitor.cpp	`93.97% <93.97%> (ø)`

... and 1 file with indirect coverage changes

Files with missing lines	Coverage Δ
include/clad/Differentiator/ReverseModeVisitor.h	`96.55% <ø> (-0.87%)`	⬇️
include/clad/Differentiator/VisitorBase.h	`100.00% <ø> (ø)`
lib/Differentiator/BaseForwardModeVisitor.cpp	`98.76% <100.00%> (+<0.01%)`	⬆️
lib/Differentiator/DerivativeBuilder.cpp	`100.00% <100.00%> (ø)`
lib/Differentiator/DiffPlanner.cpp	`98.63% <100.00%> (-0.01%)`	⬇️
lib/Differentiator/ReverseModeVisitor.cpp	`95.52% <100.00%> (-0.14%)`	⬇️
lib/Differentiator/VisitorBase.cpp	`97.63% <100.00%> (+0.36%)`	⬆️
lib/Differentiator/VectorForwardModeVisitor.cpp	`99.69% <83.33%> (-0.31%)`	⬇️
lib/Differentiator/JacobianModeVisitor.cpp	`93.97% <93.97%> (ø)`

... and 1 file 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())); ^

lib/Differentiator/JacobianModeVisitor.cpp

lib/Differentiator/VectorForwardModeVisitor.cpp

github-actions

clang-tidy made some suggestions

lib/Differentiator/JacobianModeVisitor.cpp

github-actions

clang-tidy made some suggestions

lib/Differentiator/JacobianModeVisitor.cpp

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

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

github-actions

clang-tidy made some suggestions

lib/Differentiator/VisitorBase.cpp

vgvassilev · 2024-10-25T07:11:08Z

include/clad/Differentiator/Array.h

+/// Function to define element wise adding of an arrays and an array_ref.
+template <typename T, typename U>
+CUDA_HOST_DEVICE
+    array_expression<const array<T>&, BinaryAdd, const array_ref<U>&>
+    operator+(const array<T>& arr1, const array_ref<U>& arr2) {
+  assert(arr1.size() == arr2.size());
+  return array_expression<const array<T>&, BinaryAdd, const array_ref<U>&>(
+      arr1, arr2);
+}
+
+/// Function to define element wise adding of an arrays_ref and an array.
+template <typename T, typename U>
+CUDA_HOST_DEVICE
+    array_expression<const array_ref<T>&, BinaryAdd, const array<U>&>
+    operator+(const array_ref<T>& arr1, const array<U>& arr2) {
+  assert(arr1.size() == arr2.size());
+  return array_expression<const array_ref<T>&, BinaryAdd, const array<U>&>(
+      arr1, arr2);
+}
+
+/// Function to define element wise adding of an arrays and an array_ref.
+template <typename T, typename U>
+CUDA_HOST_DEVICE
+    array_expression<const array<T>&, BinarySub, const array_ref<U>&>
+    operator-(const array<T>& arr1, const array_ref<U>& arr2) {
+  assert(arr1.size() == arr2.size());
+  return array_expression<const array<T>&, BinarySub, const array_ref<U>&>(
+      arr1, arr2);
+}
+
+/// Function to define element wise adding of an arrays_ref and an array.
+template <typename T, typename U>
+CUDA_HOST_DEVICE
+    array_expression<const array_ref<T>&, BinarySub, const array<U>&>
+    operator-(const array_ref<T>& arr1, const array<U>& arr2) {
+  assert(arr1.size() == arr2.size());
+  return array_expression<const array_ref<T>&, BinarySub, const array<U>&>(
+      arr1, arr2);
+}


Can we write tests for these? I think we generally test in Misc.

vgvassilev · 2024-10-25T07:11:21Z

include/clad/Differentiator/ArrayExpression.h

+  operator/(const array_expression<L2, BinOp2, R2>& r) const {
+    return array_expression<const array_expression<L1, BinOp1, R1>&, BinaryDiv,
+                            const array_expression<L2, BinOp2, R2>&>(*this, r);
+  }


vgvassilev · 2024-10-25T07:11:50Z

include/clad/Differentiator/ArrayRef.h

+    for (std::size_t i = 0; i < m_size; i++)
+      m_arr[i] = arr_exp[i];
+    return *this;
+  }


vgvassilev · 2024-10-25T07:14:24Z

include/clad/Differentiator/BuiltinDerivatives.h

+  // or equal to 0, then log(base) is undefined, and therefore if user only
+  // requested directional derivative of base^exp w.r.t base -- which is valid
+  // --, the result would be undefined because as per C++ valid number + NaN * 0
+  // = NaN.


Why do we need this overload?

vgvassilev · 2024-10-25T07:14:40Z

include/clad/Differentiator/BuiltinDerivatives.h

-template <typename T>
-CUDA_HOST_DEVICE ValueAndPushforward<T, T> acos_pushforward(T x, T d_x) {
+template <typename T, typename dT>
+CUDA_HOST_DEVICE ValueAndPushforward<T, dT> acos_pushforward(T x, dT d_x) {


Are these changes not good for a separate PR?

They can exist on their own but they will have no use. The types don't match in the vectorized fwd mode because those will be T and clad::array<T> respectively.

include/clad/Differentiator/Differentiator.h

vgvassilev · 2024-10-26T10:19:23Z

include/clad/Differentiator/JacobianModeVisitor.h

@@ -0,0 +1,20 @@
+#ifndef CLAD_DIFFERENTIATOR_JACOBIANMODEVISITOR_H


Can we make this a private header?

…in jacobians

…orized frw mode and jacobians

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 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

lib/Differentiator/JacobianModeVisitor.cpp Outdated Show resolved Hide resolved

PetroZarytskyi force-pushed the jac branch 4 times, most recently from ac0c084 to 8b40e97 Compare October 25, 2024 18:06

github-actions bot reviewed Oct 25, 2024

View reviewed changes

lib/Differentiator/VisitorBase.cpp Outdated Show resolved Hide resolved

vgvassilev reviewed Oct 26, 2024

View reviewed changes

PetroZarytskyi added 3 commits October 30, 2024 22:26

Remove excessive FD and request parameters from DeriveVectorMode

3101eef

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

2300e36

Move CreateGradientOverload to VisitorBase and rename it to reuse it …

e28b551

…in jacobians

PetroZarytskyi force-pushed the jac branch from 47c6992 to 6cad9f3 Compare October 30, 2024 20:30

PetroZarytskyi added 6 commits October 30, 2024 22:39

Implement forward-mode jacobians

ad97b85

Generalize custom derivative templates to be compatible with the vect…

35b8a20

…orized frw mode and jacobians

Update the tests

a499dde

Remove the old jacobian code from RMV

b14cfeb

Always use valid location when generating operators

edf5cf1

Add tests for new features

6b97654

PetroZarytskyi force-pushed the jac branch from 6cad9f3 to 6b97654 Compare October 30, 2024 20:39

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 •

edited

Loading

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 left a comment

github-actions bot left a comment

vgvassilev commented Oct 23, 2024

github-actions bot left a comment

vgvassilev Oct 25, 2024

vgvassilev Oct 25, 2024

vgvassilev Oct 25, 2024

vgvassilev Oct 25, 2024

vgvassilev Oct 25, 2024

PetroZarytskyi Oct 30, 2024

vgvassilev Oct 26, 2024

		@@ -0,0 +1,20 @@
		#ifndef CLAD_DIFFERENTIATOR_JACOBIANMODEVISITOR_H

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 • edited Loading

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 left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

vgvassilev commented Oct 23, 2024

github-actions bot left a comment

Choose a reason for hiding this comment

vgvassilev Oct 25, 2024

Choose a reason for hiding this comment

vgvassilev Oct 25, 2024

Choose a reason for hiding this comment

vgvassilev Oct 25, 2024

Choose a reason for hiding this comment

vgvassilev Oct 25, 2024

Choose a reason for hiding this comment

vgvassilev Oct 25, 2024

Choose a reason for hiding this comment

PetroZarytskyi Oct 30, 2024

Choose a reason for hiding this comment

vgvassilev Oct 26, 2024

Choose a reason for hiding this comment

PetroZarytskyi commented Oct 22, 2024 •

edited

Loading

codecov bot commented Oct 22, 2024 •

edited

Loading