How to register gradient metadata?

[davidedelvento]

According to https://enzyme.mit.edu/getting_started/CallingConvention/#custom-gradients one could give functions a custom gradient by attaching two pieces of metadata, however the example is in IR only and I would like to do it in C.

This could be extremely useful in various circumstances, obviously. In my case, I am using external libraries such as GSL which already have corresponding derivative invocations which I can use to define the appropriate gradients. I tried defining a __enzyme_register_gradient_myfunc containing myfunc (simply a wrapper around a GSL function) and the corresponding derivatives (again, simple wrapper, just to make the arguments Enzyme-friendly). Unsurprisingly (since I was just guessing at the API) that still failed with No augmented forward pass found and No reverse pass found error. I'm using LLDEnzyme with LTO.

Is this at all possible in C or should I resign to either manually edit the IR or compile GSL with Enzyme?

[wsmoses]

It is indeed (and the magic named global is the easiest way to do so right now in C++, though we are hoping to add function attribute in the future).

Maybe share your code and we can see what's going awry?

[davidedelvento]

Thanks!!

Let me provide an even simpler example, without external dependencies (again just to show that I don't understand how to use the magic named global)

#include <stdio.h>
#include <stdlib.h>

// Function to differentiate
double square(double x) {
  return x * x;
}
double custom_diff_square(double x) {
  return x * x;
}
double my_custom_diff(double x) {
  // let's pretend that this is what I want
  return 3 * x;
}

void* __enzyme_register_gradient_custom_diff_square[3] = { custom_diff_square,
                                                      my_custom_diff,
                                                      my_custom_diff };

double __enzyme_autodiff(void*, ...);
int enzyme_const, enzyme_dup, enzyme_out;

int main(int argc, char *argv[]) {
  double x = 20;
  }

  double grad_x = __enzyme_autodiff((void*)square, x);
  // expected and actual output:
  // Gradient square(20.000000) = 40.000000
  printf("Gradient square(%f) = %f\n", x, grad_x);

  grad_x = __enzyme_autodiff((void*)custom_diff_square, x);
  // expected output (given the custom thing)
  // Gradient square(20.000000) = 60.000000
  // actual output (looks like __enzyme_register_gradient has been ignored)
  // Gradient square(20.000000) = 40.000000
  printf("Custom gradient square(%f) = %f\n", x, grad_x);

  return 0;
}

compiled with

clang -flto a.c -c
clang -fuse-ld=lld  -flto a.o -o a.exe -lgsl -v -Wl,--load-pass-plugin=/path/to/build-v0.0.102-with-llvm-v17.0.6/Enzyme/LLDEnzyme-17.so

[wsmoses]

Hey, yeah the issue here is that you registered a gradient of the wrong type. You should've gotten an error message for that, but apparently that only happened on assertion. builds. Fixed that here: #1764

The error message you get should now be something like

tmp.c:22:5: error: Enzyme: Bad function type of custom reverse pass for function custom_diff_square of type double (double)
  expected gradient function to have argument types [double, double]
  Instead found my_custom_diff of type double (double)
 at context:   %15 = call double (ptr, ...) @__enzyme_autodiff(ptr noundef @custom_diff_square, double noundef %14)
int main(int argc, char *argv[]) {

[wsmoses]

Here is a corrected version of your code:

#include <stdio.h>
#include <stdlib.h>

// Function to differentiate
double square(double x) {
  return x * x;
}
double custom_diff_square(double x) {
  return x * x;
}

double square_aug(double x) {
    return square(x);
}

double my_custom_diff(double x, double diffret) {
  // let's pretend that this is what I want
  return 3 * x * diffret;
}

void* __enzyme_register_gradient_custom_diff_square[3] = { custom_diff_square,
                                                      square_aug,
                                                      my_custom_diff };

double __enzyme_autodiff(void*, ...);
int enzyme_const, enzyme_dup, enzyme_out;

int main(int argc, char *argv[]) {
  double x = 20;

  double grad_x = __enzyme_autodiff((void*)square, x);
  // expected and actual output:
  // Gradient square(20.000000) = 40.000000
  printf("Gradient square(%f) = %f\n", x, grad_x);

  grad_x = __enzyme_autodiff((void*)custom_diff_square, x);
  // expected output (given the custom thing)
  // Gradient square(20.000000) = 60.000000
  // actual output (looks like __enzyme_register_gradient has been ignored)
  // Gradient square(20.000000) = 40.000000
  printf("Custom gradient square(%f) = %f\n", x, grad_x);

  return 0;
}

[davidedelvento]

@wsmoses not sure if you have seen my comment below, since you were busy with the (great!! watched online!) Enzyme Conference and github has this question marked answered, so it might not send notifications.
In short, I appreciate your help, but this is not working for me (see below for details)

[davidedelvento]

Thanks a lot for providing this example and sorry for how long it took me to respond.

I am very puzzled by this example.

First, I am not sure why you are defining square_aug. What is that supposed to be? The function I want to specify a custom derivative is custom_diff_square already, why defining another one? I guess part of the confusion may be because I confused you with my definitions, so let me try to clarify that.

• My function square is there just for reference (to show the difference between a custom derivative and an automatic one) but for a pure custom derivative example it should not be there.
• My function custom_diff_square is the function I want to define a custom derivative for (in my real life code that would be a library function, not something I define manually as here)
• My function my_custom_diff is the function I want to invoke as a custom derivative for custom_diff_square

And that leads me to my second puzzlement: what is this additional argument double diffret which I am supposed to add to my custom derivative? I mean, I'd expect the derivative to have the same arguments as the custom_diff_square....

And finally, the __enzyme_register_gradient_custom_diff_square has three elements. The first is custom_diff_square, the function I want to register the derivative for. The last is my_custom_diff, i.e. said derivative. But what is supposed to be the second one? I guess this ties back to my question above.

Finally, even doing all of this, the program still returns 40 in both cases, rather than 60 as I would expect from my custom derivative in which I am intentionally using a wrong derivative... so that puzzles me even more. This is with LLVM v18.1.1 and Enzyme v0.0.103 and the compilation lines are

clang -flto a.c -c
clang -fuse-ld=lld -flto a.o -v -Wl,-load-pass-plugin=/path/to/LLDEnzyme-18.so

PS: not sure I should, but with the same versions, my original version still compiles and works incorrectly, rather than spitting the error message you mentioned

[wsmoses]

Sorry, yeah completely missed this. And custom rules definitely need better (or rather any C++) docs. We're currently adding a bunch of nice c++ sugar, so hopefully we'll add after that.

The Julia custom rules docs is similar (but obviously julia not c++), but still can explain a bit what's happening: https://enzyme.mit.edu/index.fcgi/julia/stable/generated/custom_rule/#Defining-a-reverse-mode-rule . The CUDA guide also does this for CUDA/C++ but is more about the heterogeneity than the custom rules (https://enzyme.mit.edu/getting_started/CUDAGuide/).

In reverse mode (where you're defining your rule), the three functions are: the function to overload, a function to replace the original function with in the forward pass, the backwards pass, in that order. The reason why the second function is necessary, is to save any data that may be overwritten but required in the reverse pass.

For example, from the CUDA guide the aug_kern function cudaMalloc's memory to save data needed for the backwards pass:

void* aug_kern(float* src, float* dsrc, float* dst, float* ddst) {
    void** tape;
    cudaMalloc(&tape, sizeof(void*) * /*total number of threads*/100);
    aug_collide<<<1, 100>>>(src, dsrc, dst, ddst, tape);
    return (void*)tape;
}

This is what square_aug does here. However, since there's no overwritten data (and all registers like double's are passed in as is), it can look the same as the original function. If the function loaded memory from a pointer, and we were conservative and assumed the data could be overwritten later in the code being differentiated -- we may need to save it to compute the correct derivative.

As for the additional argument to the reverse pass. For functions that return an immutable, differentiated value (e.g. double) that additional variable is the partial derivative seed which we are back propagating to the reverse pass. For example, consider performing d/dx f(g(x)) and implementing a custom derivative for f. The total derivative is f'(y=g(x)) * g'(x). If we define a derivative for a generic f(y), in this case the derivative of the return will contain g'(x), such that when we compute the derivative propagation rule f'(y) * dy, we will in fact return the correct total derivative f'(g(x)) * g'(x).

[wsmoses]

and yeah @davidedelvento I didn't get a notifiaction so feel free to @ me in the future (@ 'ing you now for the same reason).

[davidedelvento]

@wsmoses Thank you very much for your patience, understanding and helpful explanation. I am still an Enzyme newbie, but I am learning its abilities and how to use it and I really like it. Everything you said makes perfect sense, however even the corrected example did not work as expected when I compiled as I described it above. I was able to figure out the missing piece myself digging into the integration tests and their builing scripts (particularly invsqrt.c has been enlightening).

The program started working when I compiled as follows

clang -flto a.c -c -fpass-plugin=/path/to/ClangEnzyme-18.so -Xclang -load -Xclang /path/to/ClangEnzyme-18.so
clang -fuse-ld=lld -flto a.o -v -Wl,-load-pass-plugin=/path/to/LLDEnzyme-18.so

Of course in this example I could have avoided the separate linking step and using only ClangEnzyme, but I need that since I will move to multisource immediately after.

I will be happy to make a PR including this contrived example, or the invsqrt in https://enzyme.mit.edu/getting_started/CallingConvention/#custom-gradients (or somewhere else) if you think it's appropriate.