enzymead / www Goto Github PK

#include <math.h>
#include <stdio.h>

#define nullptr ((void*)0)

// f(x, y) -> { sin(x), sin(x)^2 + y }
void mysimulation(double *data) {
    data[0] = sin(data[0]);
    data[1] += data[0] * data[0];
}

void __enzyme_autodiff(void (*)(double *), ...);
void* __enzyme_augmentfwd(void (*)(double *), ...);
void __enzyme_reverse(void (*)(double *), ...);

int enzyme_nofree;
int enzyme_tape;

int main()
{
    double data[2] = {2.0, 3.0};
    mysimulation(data);
    printf("vanilla run %f %f\n", data[0], data[1]);

    data[0] = 2.0; data[1] = 3.0;

    // Take the derivative wrt x, on fresh data.
    double d_data[2] = {1.0, 0.0};
    __enzyme_autodiff(mysimulation, data, d_data);
    printf("result run %f %f | derivative %f %f\n", data[0], data[1], d_data[0], d_data[1]);

    d_data[0] = 1.0; d_data[1] = 0.0;
    __enzyme_autodiff(mysimulation, data, d_data);
    printf("rerun derivative without resetting data result %f %f | derivative %f %f\n", data[0], data[1], d_data[0], d_data[1]);


    data[0] = 2.0; data[1] = 3.0;

    // Compute the original result, and also save any information
    // needed for the derivative. We don't need the shadow (derivative)
    // array here so we can pass nullptr.
    void* tape = __enzyme_augmentfwd(mysimulation, data, nullptr);
    printf("augmented run %f %f\n", data[0], data[1]);

    // Just for extra fun, overwrite the data
    data[0] = 0.0/0.0;
    data[1] = 0.0/0.0;

    // Compute the original result, and also save any information
    // needed for the derivative. We don't need the primal (original)
    // array here so we can pass nullptr.
    d_data[0] = 1.0; d_data[1] = 0.0;
    __enzyme_reverse(mysimulation, enzyme_nofree, nullptr, d_data, tape);
    printf("rerun derivative without resetting data result %f %f | derivative %f %f\n", data[0], data[1], d_data[0], d_data[1]);

    // We can in fact compute the derivative now a second time, using the same cached
    // data and again get the same (correct) result. Now without enzyme_nofree, we'll
    // free the preserved data.
    d_data[0] = 1.0; d_data[1] = 0.0;
    __enzyme_reverse(mysimulation, nullptr, d_data, tape);
    printf("second run derivative without resetting data result %f %f | derivative %f %f\n", data[0], data[1], d_data[0], d_data[1]);

    return 0;
}