Reimplement jacobians using the vectorized forward mode #1121
Conversation
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #1121 +/- ##
==========================================
- Coverage 94.48% 94.43% -0.06%
==========================================
Files 50 51 +1
Lines 8850 8928 +78
==========================================
+ Hits 8362 8431 +69
- Misses 488 497 +9
... and 1 file with indirect coverage changes
|
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| 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) \ |
There was a problem hiding this comment.
warning: function-like macro 'JacobianDerivedFnTraits_AddSPECS' used; consider a 'constexpr' template function [cppcoreguidelines-macro-usage]
#define JacobianDerivedFnTraits_AddSPECS(var, cv, vol, ref, noex) \
^There was a problem hiding this comment.
@vgvassilev This is tricky to address and it's not related to the PR. Do we want to fix this now?
There was a problem hiding this comment.
No but we should probably open an issue to track it?
| clad::utils::ComputeEffectiveFnName(FD) + "_pushforward"; | ||
| callDiff = m_Builder.BuildCallToCustomDerivativeOrNumericalDiff( | ||
| customPushforward, customDerivativeArgs, getCurrentScope(), | ||
| const_cast<DeclContext*>(FD->getDeclContext())); |
There was a problem hiding this comment.
warning: do not use const_cast [cppcoreguidelines-pro-type-const-cast]
const_cast<DeclContext*>(FD->getDeclContext()));
^There was a problem hiding this comment.
@vgvassilev I don't think there's a way to avoid this easily right now. The same thing is done in all visitors. Should I silence the warning?
There was a problem hiding this comment.
No, that needs deeper refactoring
|
Could we write a few benchmarks to make sure we do not regress in performance? |
PetroZarytskyi
force-pushed
the
jac
branch
4 times, most recently
from
ac0c084 to
8b40e97
Compare
| // 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. |
There was a problem hiding this comment.
Why do we need this overload?
There was a problem hiding this comment.
It was tricky to make this pushforward work as both vectorized and non-vectorized. The main reason is that the output type is hard to deduce. Also the line auto derivative = ... was created as array_expression, which was a problem for a couple of reasons. I reimplemented this with better templates in the last commit. Does this look better?
| 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) { |
There was a problem hiding this comment.
Are these changes not good for a separate PR?
There was a problem hiding this comment.
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.
PetroZarytskyi
force-pushed
the
jac
branch
2 times, most recently
from
6cad9f3 to
6b97654
Compare
PetroZarytskyi
force-pushed
the
jac
branch
2 times, most recently
from
fc48ba5 to
74d70de
Compare
| 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) \ |
There was a problem hiding this comment.
No but we should probably open an issue to track it?
PetroZarytskyi
force-pushed
the
jac
branch
2 times, most recently
from
f25aac9 to
a42eceb
Compare
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) ``` 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 vgvassilev#472, Fixes vgvassilev#1057, Fixes vgvassilev#480, Fixes vgvassilev#527
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:
or
This signature is also similar to the old one. e.g.
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, Fixes #1057, Fixes #480, Fixes #527