#include "../test_utils.h"
extern void __enzyme_autodiff(void*, ...);
void square(double** p_delv, double** p_e, int ** idx, int numReg, int numElemReg) {
double* delv = *p_delv;
double* e = *p_e;
for (int r = 0; r < numReg; r++) {
double* tmp = (double*)malloc(numElemReg * sizeof(double));
for (int i=0; i<numElemReg; i++) {
int off = idx[r][i];
tmp[i] = delv[off];
}
for (int i=0; i<numElemReg; i++) {
int off = idx[r][i];
e[off] = tmp[i] * tmp[i];
}
free(tmp);
}
}
int main() {
int numReg = 100;
double *delv = (double*)malloc(sizeof(double)*numReg);
double *e = (double*)malloc(sizeof(double)*numReg);
double *d_delv = (double*)malloc(sizeof(double)*numReg);
double *d_e = (double*)malloc(sizeof(double)*numReg);
int* idxs[numReg];
int numRegElem = 200;
for (int i=0; i<numReg; i++) {
int* data = (int*)malloc(sizeof(int)*numRegElem);
for (int j=0; j<numRegElem; j++) {
data[j] = j % numReg;
}
idxs[i] = data;
delv[i] = i;
d_delv[i] = 0;
e[i] = 0;
d_e[i] = 1;
}
square(&delv, &e, idxs, numReg, numRegElem);
for (int i=0; i<numReg; i++) {
printf("e=%f delv=%f\n", e[i], delv[i]);
}
__enzyme_autodiff((void*)square, enzyme_dup, &delv, &d_delv, enzyme_dup, &e, &d_e, enzyme_const, idxs, enzyme_const, numReg, enzyme_const, numRegElem);
for (int i=0; i<numReg; i++) {
printf("d_e=%f d_delv=%f\n", d_e[i], d_delv[i]);
APPROX_EQ(d_e[i], 0.0, 1e-10);
APPROX_EQ(d_delv[i], 2.0 * i, 1e-10);
}
}