Looking for help registering custom derivatives of function template

[samuelpmish]

Hi,

I've got an application that solves a lot of different nonlinear equations in one unknown, and I'd like to use enzyme to help implement the derivative expressions needed by the nonlinear solver. Using newton's method as an example, I'm trying something like

template <auto f, typename T>
double newtons_method(double x0, T p)
{
    auto fprime = [&](double x) {
        double dx = 1.0;
        return __enzyme_fwddiff<double>((void*)+f, enzyme_dup, x, dx, enzyme_const, p);
    };

    const int KMAX_ITERS = 25;
    double tolerance = 1e-10;

    double x = x0;
    [[maybe_unused]] bool converged = false;
    for (int i = 0; i < KMAX_ITERS; i++) {
        double r = f(x, p);
        if (std::abs(r) < tolerance) {
            converged = true;
            break;
        }
        double J = fprime(x);
        x -= r/J;
    }
    return x;
}

That is, newtons_method is a function template, so that it can support root-finding for arbitrary functions and "parameters" p.

note: the function that we try to find a root of is passed by template argument per the recommendation in (#1186 (comment)).

This seems to be working well, but I'd like to implement custom rules for the forward/reverse mode differentiation (the implicit function theorem tells us we don't need to track derivative info through the for loop, we only need to linearize about the solution). For example,

template <auto f, typename T>
double newtons_method_fwddiff(double x0, T p, T dp)
{
    double x = newtons_method<f>(x0, p);
    double dfdx = __enzyme_fwddiff<double>((void*)+f, enzyme_dup, x, 1.0, enzyme_const, p);
    double dfdp = __enzyme_fwddiff<double>((void*)+f, enzyme_const, x, enzyme_dup, p, dp);
    std::cout << "Custom diff is being called" << std::endl;
    return -dfdp/dfdx;
}

This code also seems to be working as-expected, but I don't know how to communicate to enzyme that it should be preferred when writing

  // I'd like enzyme to use my custom fwddiff rule here, but how?
  auto dx = __enzyme_fwddiff<double>(
    (void*)newtons_method<residual, ParameterType>, 
    enzyme_const, x0, 
    enzyme_dup, p, dp
  );

I haven't found much in the way of explanation for how to register custom derivative rules. My first attempt is:

template <auto f, typename T>
void* __enzyme_register_derivative_newtons_method[] = { 
    (void*) newtons_method<f, T>, 
    (void*) newtons_method_fwddiff<f, T> 
};

but it doesn't seem to be working. It's possible I'm misunderstanding the requirements, or that this variable-template approach isn't getting instantiated so it isn't noticed by enzyme's passes.

Can anyone help me figure out how to make progress toward custom derivatives? I am also interested in writing a custom reverse pass, but I haven't found any examples of how to do that either.

Thank you,

Sam

---

Full code available here

[wsmoses]

Ah counter-intuitively, right now in C/C++ custom rules [or even differentiation of known functions like sin] doesn't work if it's the outermost function that you pass in.

e.g.

you can't presently do __enzyme_fwddiff(sin). but you can do f(x) = sin(x); __enzyme_fwddiff(f)

[samuelpmish]

Hmm, I think we might be talking past each other.

To try and clarify, I'm saying that I have written a bunch of nonlinear functions I'd like to solve, say foo(x,p), bar(x,p), baz(x,p). And I can solve for x (with fixed p) by invoking newtons_method<foo>, newtons_method<bar>, and newtons_method<baz>. Now, how can I get register a custom derivative of this implicit function template newtons_method<f> with enzyme?

Or is the implication that there's no way to define a custom derivative rule for a user-written function template?

[wsmoses]

Ah sorry, that is one of two issues you'll face (and seemingly not the one you were looking for).

You may be able to force the template registration like this perhaps?

template<auto f, typename T>
struct RegisterDerivative {
__attribute__((used))
static __enzyme_register_derivative_newtons_method[] = { 
    (void*) newtons_method<f, T>, 
    (void*) newtons_method_fwddiff<f, T> 
};
}

template<> struct RegisterDerivative<f1, float>;
....
// explicit registration needed for all types you may need

You will, however, seemingly still need to remedy the other issue I mention above since you want to call

  // I'd like enzyme to use my custom fwddiff rule here, but how?
  auto dx = __enzyme_fwddiff<double>(
    (void*)newtons_method<residual, ParameterType>, 
    ```
    
    but the custom derivative is defined for that function. You'll also need to make a wrapper function like I mention above.

[samuelpmishLLNL]

Thank you for the suggestions!

I tried using the struct to help manually instantiate an appropriate template specialization of the derivative rules. Code is available here: https://fwd.gymni.ch/kzNz0b

On the Enzyme compiler explorer instance, it compiles without issue but doesn't call the custom derivative expression.
On my workstation, I get the following error message when compiling:

[build] /path/to/project/tests/enzyme_test.cpp:19:16: error: Enzyme: No forward mode derivative found for _Z16__enzyme_fwddiffIdJiddiNSt3__15tupleIJddEEEEET_PvDpT0_
[build]  at context:   %call = call noundef double @_Z16__enzyme_fwddiffIdJiddiNSt3__15tupleIJddEEEEET_PvDpT0_(ptr noundef @_Z8residualdNSt3__15tupleIJddEEE, i32 noundef %0, double noundef %x, double noundef 1.000000e+00, i32 noundef %1, [2 x double] %.fca.1.insert) #13, !dbg !3005
[build]         return __enzyme_fwddiff<double>((void*)+f, enzyme_dup, x, dx, enzyme_const, p);
[build]                ^
[build] /path/to/project/tests/enzyme_test.cpp:19:16: error: Enzyme: Cannot cast __enzyme_autodiff primal argument 1, found @"_enzyme_forward__Z8residualdNSt3__15tupleIJddEEE'" = internal constant ptr @fwddiffe_Z8residualdNSt3__15tupleIJddEEE.2, type ptr - to arg 0 double

Compiler/LLVM version: 16.0.0
Enzyme version: 304e21b (Jan 23, 2024)

---

As I write this, I'm wondering: what should the function signature be for the custom derivative? My custom derivative is implicitly enzyme_const wr.t. x0, but enzyme_dup for p. Maybe I'm pointing enzyme at a function with the wrong inputs/outputs.

[wsmoses]

https://fwd.gymni.ch/vn6OBj

[wsmoses]

and with the noinline: https://fwd.gymni.ch/wm1Rrc

[samuelpmish]

Thanks for helping us out yesterday afternoon! I'll try to summarize some of the things we learned in that meeting:

• my custom fwddiff call had the wrong signature, it needed to return a tuple of {value, derivative}, rather than just the derivative
• it's helpful to put the ((noinline)) attribute on the newtons_method_impl definition (but I don't understand why)
• we were able to manually register a custom derivative for a function template, but couldn't find a way to automatically register that derivative. For example, right now to use that newtons_method function, the user has to do

// define the function we'd like to find a root for 
double myfunc(double x, double p) {
  ...
}

// find x satisfying myfunc(x, p) == 0
auto x = newtons_method<myfunc>(x0, p);

// I want to use enzyme to use my custom derivative rule for newtons_method() here
auto dx = __enzyme_fwddiff<double>(
  (void*)newtons_method<myfunc, double>, 
  enzyme_const, x0, 
  enzyme_dup, p, dp
);

// but I currently have to manually instantiate the custom derivative rule
__attribute__((used))
void * __enzyme_register_derivative_newtons_method_impl[2] = { 
  (void*) newtons_method_impl<myfunc, double>, 
  (void*) newtons_method_impl_fwddiff<myfunc, double> 
};

The downsides to this approach are:

• the user has to manually instantiate the custom derivative rule (although this could maybe be simplified with a macro)
• it's not clear to me that it's possible to register custom derivatives for more than 1 function. e.g. if I had two functions that I wanted to use with newtons_method, I couldn't write

__attribute__((used))
void * __enzyme_register_derivative_newtons_method_impl[2] = { 
  (void*) newtons_method_impl<myfunc1, double>, 
  (void*) newtons_method_impl_fwddiff<myfunc1, double> 
};

__attribute__((used))
void * __enzyme_register_derivative_newtons_method_impl[2] = { 
  (void*) newtons_method_impl<myfunc2, double>, 
  (void*) newtons_method_impl_fwddiff<myfunc2, double> 
};

because the symbol __enzyme_register_derivative_newtons_method_impl would be defined twice.

Near the end of the meeting, you also mentioned another potential path forward: decorating a function (or maybe function template) with an attribute that expresses that enzyme should use another function for the derivative.

@btalamini

[wsmoses]

Enzyme will recognize any global containing the string __enzyme_register_derivative so add an extra prefix/suffix/name mangling will ensure multiple symbols are possible. In this case to avoid defining derivatives for multiple templates it may be easier to make an inner function which you pass the templates function arg into and then defining a single derivative for the function that takes the function pointer. Noinline is required to ensure that the function with the custom derivative isn’t optimized into its caller (and thus we cannot rexognize a call to replace). Our attributes deal with this automatically but the global here doesn’t seem to do a priori. I’ll work on making this easier

…

On Fri, May 17, 2024 at 7:30 AM Sam ***@***.***> wrote: Thanks for helping us out yesterday afternoon! I'll try to summarize some of the things we learned in that meeting: - my custom fwddiff call had the wrong signature, it needed to return a tuple of {value, derivative}, rather than just the derivative - it's helpful to put the ((noinline)) attribute on the newtons_method_impl definition (but I don't understand why) - we were able to manually register a custom derivative for a function template, but couldn't find a way to automatically register that derivative. For example, right now to use that newtons_method function, the user has to do // define the function we'd like to find a root for double myfunc(double x, double p) { ... } // find x satisfying myfunc(x, p) == 0auto x = newtons_method<myfunc>(x0, p); // I want to use enzyme to use my custom derivative rule for newtons_method() hereauto dx = __enzyme_fwddiff<double>( (void*)newtons_method<myfunc, double>, enzyme_const, x0, enzyme_dup, p, dp ); // but I currently have to manually instantiate the custom derivative rule__attribute__((used)) void * __enzyme_register_derivative_newtons_method_impl[2] = { (void*) newtons_method_impl<myfunc, double>, (void*) newtons_method_impl_fwddiff<myfunc, double> }; The downsides to this approach are: - the user has to manually instantiate the custom derivative rule (although this could maybe be simplified with a macro) - it's not clear to me that it's possible to register custom derivatives for more than 1 function. e.g. if I had two functions that I wanted to use with newtons_method, I couldn't write __attribute__((used)) void * __enzyme_register_derivative_newtons_method_impl[2] = { (void*) newtons_method_impl<myfunc1, double>, (void*) newtons_method_impl_fwddiff<myfunc1, double> }; __attribute__((used)) void * __enzyme_register_derivative_newtons_method_impl[2] = { (void*) newtons_method_impl<myfunc2, double>, (void*) newtons_method_impl_fwddiff<myfunc2, double> }; because the symbol __enzyme_register_derivative_newtons_method_impl would be defined twice. Near the end of the meeting, you also mentioned another potential path forward: decorating a function (or maybe function template) with an attribute that expresses that enzyme should use another function for the derivative. @btalamini <https://github.com/btalamini> — Reply to this email directly, view it on GitHub <#1640 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXBYXTOWTFXHASEMPEDZCYIBBAVCNFSM6AAAAABCMWWPD6VHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM4TINZRGY3TC> . You are receiving this because you commented.Message ID: ***@***.***>

[samuelpmish]

Enzyme will recognize any global containing the string
__enzyme_register_derivative so add an extra prefix/suffix/name mangling
will ensure multiple symbols are possible.

Ah, I thought the suffix was important for the __enzyme_register_derivative to associate the derivative with the right symbol name, but I think you're saying that it really just looks at the two entries as a sort of key-value pair.

[wsmoses]

Yup exactly. I had put that suffix in the example as a way to ensure non collision if we added something else

…

On Fri, May 17, 2024 at 7:45 AM Sam ***@***.***> wrote: Enzyme will recognize any global containing the string __enzyme_register_derivative so add an extra prefix/suffix/name mangling will ensure multiple symbols are possible. Ah, I thought the suffix was important for the __enzyme_register_derivative to associate the derivative with the right symbol name, but I think you're saying that it really just looks at the two entries as a sort of key-value pair. — Reply to this email directly, view it on GitHub <#1640 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXDWEW5FJC5KAHVP3CTZCYJZZAVCNFSM6AAAAABCMWWPD6VHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM4TINZRHAZTC> . You are receiving this because you commented.Message ID: ***@***.***>