C++ interface TODO list

[samuelpmishLLNL]

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]

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]

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

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

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

[wsmoses]

sret issues should now be fixed by #2051

[thelfer]

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

yup all of this is on main

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

[samuelpmish]

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]

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

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

[jandrej]

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

[thelfer]

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

[jandrej]

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