Enzyme icon indicating copy to clipboard operation
Enzyme copied to clipboard

C++ interface TODO list

Open samuelpmishLLNL opened this issue 1 year ago • 14 comments

I'm writing this issue to capture some of the ideas mentioned in a private discussion with @wsmoses about the future C++ interface. I'm hoping that this will increase visibility and allow other contributors to weigh in with design suggestions and considerations, as well as track progress on these topics.


  1. ~(fix WIP) Address an issue where C++ functions that return composite types by-value are mishandled after sret transformations (and produce error messages that report the functions return void). This has been observed in some larger projects, but still needs a minimal reproducer.~
    edit: fixed by https://github.com/EnzymeAD/Enzyme/pull/2051

  2. Add a make_zero function for basic types (and allow users to specialize it for custom types), in order to better support differentiation over custom containers.

  3. Add higher-level tools for Jacobian, Hessian, JVP, VJP, HVP (actual interface still TBD, maybe look to Julia implementation for inspiration here?)

  4. Support for batch-duplicated arguments (what container / layout should be used to pass the batch?)

  5. Custom Derivative rule interface (likely also requires some effort in front end)

samuelpmishLLNL avatar Jul 25 '24 20:07 samuelpmishLLNL

Dear @samuelpmishLLNL, @wsmoses,

Is there any preview of the expected API of this C++ interface ?

I am currently working on my own C++ wrapper which aims at translating this code code:

__enzyme_fwddiff<double>(reinterpret_cast<void*>(f),
                                          enzyme_const, 2., enzyme_dup, 2, 1)

into:

const auto v = VariableValueAndIncrement<double>{2, 1};
diff(f2, 2., v);

where the special type VariableValueAndIncrement indicates with respect to which variable differentiation is made.

I also plan to implement a get_derivative function that would return a callable computing the derivative of a given function with the following syntax:

auto d2f_dxdy = get_derivative<0, 1>(f);

Have you something similar in mind ?

thelfer avatar Aug 07 '24 10:08 thelfer

@thelfer yeah, sorry for the delay, see these for example:

https://github.com/EnzymeAD/Enzyme/blob/main/enzyme/test/Integration/ReverseMode/sugar.cpp https://github.com/EnzymeAD/Enzyme/blob/main/enzyme/test/Integration/CppSugar/gh_issue_1785.cpp https://github.com/EnzymeAD/Enzyme/tree/main/enzyme/test/Integration/CppSugar

This essentially adds type checking and other nice things to the current autodiff mechanic.

Subsequently we aim to add gradient/jacobian/etc wrappers atop this as well.

wsmoses avatar Aug 25 '24 15:08 wsmoses

@samuelpmishLLNL can you make a MWE for 1)

I just finished adding support for Enzyme to automatically interleave args, which I presume @jandrej will like regardless -- but also enables us to pass in tensors/structs by value without issue. see https://github.com/EnzymeAD/Enzyme/pull/2048

wsmoses avatar Aug 25 '24 15:08 wsmoses

This may be such an example: https://fwd.gymni.ch/yRcfkC

wsmoses avatar Aug 25 '24 15:08 wsmoses

sret issues should now be fixed by https://github.com/EnzymeAD/Enzyme/pull/2051

wsmoses avatar Aug 25 '24 19:08 wsmoses

@wsmoses Thanks for keeping in touch. We currently make progresses on our side: https://github.com/thelfer/TFELMathEnzyme/tree/master/tests

We made it worked for our own tensorial library, but it could easily be adapted to work with others, like Eigen.

For example:

    constexpr auto eps = double{1e-14};
    constexpr auto E = double{70e9};
    constexpr auto nu = double{0.3};
    constexpr auto lambda = computeLambda(E, nu);
    constexpr auto mu = computeMu(E, nu);
    const auto hooke_potential = [](const Stensor& e) {
      return (lambda / 2) * power<2>(trace(e)) + mu * (e | e);
    };
    // second derivative with respect to the first variable
    // m is either REVERSE or FORWARD
    const auto stress = getDerivativeFunction<m, 0>(hooke_potential);
     // second derivative with respect to the first variable
    const auto stiffness = getDerivativeFunction<m, 0, 0>(hooke_potential);
    const auto e = Stensor{0.01, 0, 0, 0};
    const auto s = stress(e);
    const auto K = stiffness(e);

We implemented a wrapper around Enzyme to hide many gory details.

However, the code of the TFELMathEnzyme library is very complicated and painful to write du to the fact that the symbols enzyme_const, enzyme_dup etc.. must be explicitely in the arguments list. The library thus only handles functions up to 3 variables right now (in reverse mode).

If the objects like enzyme::Const could remove this limitation, that would be fantastic, because the machinery of our library already handles on the fly conversion of objects to enzyme (for instance refences to pointers, or integer to double if required).

Are your examples already in master ?

thelfer avatar Aug 26 '24 11:08 thelfer

yup all of this is on main

And yes already the c++ autodiff syntax should abstract out all of the magic constants

wsmoses avatar Aug 26 '24 14:08 wsmoses

However, the code of the TFELMathEnzyme library is very complicated and painful to write du to the fact that the symbols enzyme_const, enzyme_dup etc.. must be explicitely in the arguments list. The library thus only handles functions up to 3 variables right now (in reverse mode).

This is one of the main things the C++ interface hopes to address. By communicating the activity (const, duplicated, etc) through the type system, it's possible to write variadic functions that use enzyme under the hood.

Are your examples already in master ?

There are only a few examples of the C++ interface right now, but I don't think we have a variadic one.

samuelpmish avatar Aug 26 '24 14:08 samuelpmish

@wsmoses Perfect ! I'll give it a try and make feed-backs !

thelfer avatar Aug 26 '24 14:08 thelfer

This is one of the main things the C++ interface hopes to address. By communicating the activity (const, duplicated, etc) through the type system, then it's possible to write variadic functions that use enzyme under the hood.

That's exactly what caused me pain and headaches !

thelfer avatar Aug 26 '24 14:08 thelfer

@thelfer related, you might want to have a look at https://github.com/mfem/mfem/blob/dfem-coefficient/examples/dfem/dfem.hpp#L1719

jandrej avatar Aug 26 '24 15:08 jandrej

@jandrej Thanks for the pointer. Have you published anything about this, just to get an overview of the project ?

thelfer avatar Aug 26 '24 20:08 thelfer

@jandrej Thanks for the pointer. Have you published anything about this, just to get an overview of the project ?

Not yet, but feel free to ping me if you want to talk (e.g. MFEM workshop slack)

jandrej avatar Aug 26 '24 20:08 jandrej

Hi. Would someone give me an update/example for hessian in computation in c++?

Thanks Andy

a-jp avatar Feb 23 '25 09:02 a-jp

Hi, the C++ interface looks great. I was wondering what the current status of it is? For example is it likely to change much, or considered not ready for use? I couldn't find any documentation for it. Thanks!

jeffpollock9 avatar Sep 23 '25 18:09 jeffpollock9

we definitely need more/any docs on it but it should be usable (looking in the headers for info on).

The big TBD is we need to add the nice c++ custom derivative registration

wsmoses avatar Sep 23 '25 18:09 wsmoses

we definitely need more/any docs on it but it should be usable (looking in the headers for info on).

The big TBD is we need to add the nice c++ custom derivative registration

@wsmoses thanks for your response.

jeffpollock9 avatar Sep 23 '25 18:09 jeffpollock9