Enzyme Linker plugin doesn't recognize the custom gradient.

[zihay]

Hi, I am using the Enzyme linker plugin LLDEnzyme-12.so to test the custom gradient functionality. It seems that it can't recognize the custom gradient I defined. I am not sure if I am using it correctly. Below is the code. By calling this, it gives the correct result.

make main
./main

This gives the wrong result.

make main2
./main2

Could you take a look? Thanks!

main.cpp

#include <stdio.h>

double __enzyme_autodiff(void*, double);

__attribute__((optnone))
void square_(const double* src, double* dest) {
    *dest = *src * *src;
}

int augment = 0;
void* augment_square_(const double* src, const double *d_src, double* dest, double* d_dest) {
    augment++;
    // intentionally incorrect for debugging
    *dest = 7.0;
    *d_dest = 11.0;
    return nullptr;
}

int gradient = 0;
void gradient_square_(const double* src, double *d_src, const double* dest, const double* d_dest, void* tape) {
    gradient++;
    // intentionally incorrect for debugging
    *d_src = 13.0;
}

void* __enzyme_register_gradient_square[] = {
    (void*)square_,
    (void*)augment_square_,
    (void*)gradient_square_,
};


double square(double x) {
    double y;
    square_(&x, &y);
    return y;
}

double dsquare(double x) {
    return __enzyme_autodiff((void*)square, x);
}


int main() {
    double res = dsquare(3.0);
    printf("res=%f augment=%d gradient=%d\n", res, augment, gradient);
`}`

Makefile

input.ll: main.cpp
	clang main.cpp -S -emit-llvm -o input.ll -O2 -fno-vectorize -fno-slp-vectorize -fno-unroll-loops

output.ll: input.ll
	opt input.ll -load=/usr/local/lib/LLVMEnzyme-12.so -enzyme -o output.ll -S

main: output.ll
	clang output.ll -o main

main2: main.cpp
	clang++ -fuse-ld=lld -flto -o main2 main.cpp -Wl,--lto-legacy-pass-manager -Wl,-mllvm=-load=/usr/local/lib/LLDEnzyme-12.so

[wsmoses]

@LeilaGhaffari this appears to be very similar to how at O0 the enzyme_like functionality we added didn't work until we ensured the code wasn't inlined until after Enzyme.

Would you be able to investigate this?

[laylagi]

Interesting! Sure, I can give a look in the next couple of days.