Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

C++ interface TODO list #2012

Open
samuelpmishLLNL opened this issue Jul 25, 2024 · 13 comments
Open

C++ interface TODO list #2012

samuelpmishLLNL opened this issue Jul 25, 2024 · 13 comments

Comments

@samuelpmishLLNL
Copy link
Collaborator

samuelpmishLLNL commented Jul 25, 2024

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 Sugar: fix sret fwd mode #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)

@thelfer
Copy link

thelfer commented Aug 7, 2024

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 ?

@wsmoses
Copy link
Member

wsmoses commented Aug 25, 2024

@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
Copy link
Member

wsmoses commented Aug 25, 2024

@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 #2048

@wsmoses
Copy link
Member

wsmoses commented Aug 25, 2024

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

@wsmoses
Copy link
Member

wsmoses commented Aug 25, 2024

sret issues should now be fixed by #2051

@thelfer
Copy link

thelfer commented Aug 26, 2024

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

@wsmoses
Copy link
Member

wsmoses commented Aug 26, 2024

yup all of this is on main

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

@samuelpmish
Copy link
Collaborator

samuelpmish commented Aug 26, 2024

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.

@thelfer
Copy link

thelfer commented Aug 26, 2024

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

@thelfer
Copy link

thelfer commented Aug 26, 2024

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 !

@jandrej
Copy link
Contributor

jandrej commented Aug 26, 2024

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

@thelfer
Copy link

thelfer commented Aug 26, 2024

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

@jandrej
Copy link
Contributor

jandrej commented Aug 26, 2024

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

5 participants