Incorrect derivative result when nested void functions and recursive nature functions are used. about enzyme HOT 5 OPEN

#include <cuda_runtime_api.h>
#include <cuda_runtime.h>
#include <iostream>
#include <cstdio>
#include <stdlib.h>
#include <cmath>
#include <device_launch_parameters.h>
#include <random>

int __device__ enzyme_dup;
int __device__ enzyme_out;
int __device__ enzyme_const;
int __device__ enzyme_dupnoneed;

void __enzyme_autodiff(...);

template<typename RT, typename ... Args>
RT __enzyme_autodiff(void*, Args ...);

#define N 12

inline
cudaError_t cudaCheck(cudaError_t result) {

    if(result != cudaSuccess){
        fprintf(stderr, "CUDA runtime error: %s \n", cudaGetErrorString(result));
        assert(result == cudaSuccess);
    }

    return result;
}

inline
__host__ __device__ void gbmSimulation(double* short_rate, double* gbm , double* rand, int maturity) {

    int size = maturity;

    double dt = 1.0/12.0;

    double s = 1.15;

    double vol = 0.02;

    gbm[0] = s;

    for (int i = 1; i < size; i++){
        gbm[i] = gbm[i-1] * exp((short_rate[i] - 0.5*vol*vol)*dt + vol*sqrt(dt)*rand[i]);
    }

}

inline
__host__ __device__ double sumtil( double* vec, double* gbms, double* randomNumbers, int size) {

    double ret = 0.0;
    double volatility = 0.023;
    double dt  = 1.0/12.0;
    double s = 1.15;

    gbms[0] = s;

    //the following line of code, if called and used in the calculations that follow produce the incorrect derivative result
    //when compared with the finite difference approximation

    gbmSimulation(vec,gbms,randomNumbers, size);

    for (int i = 0; i < size; i++) {

        ret += gbms[i];
    }

    //the following commented section highlights the second issue (please comment out the void function call as well as the previous for loop),
     // which is that the CPU and GPU implementations produce different results with the CPU based implementation producing the correct result when compared with the FDA
    //for (int i = 1; i < size; i++) {

            //correct CPU derivative result, the GPU implementation however does not !!

            //gbms[i] = gbms[i - 1] * exp((vec[i] - 0.5 * volatility * volatility) * dt + volatility * (0.5 * dt) * randomNumbers[i]);

            //ret += gbms[i];
    //}

    return ret;
}

__host__ __device__ double sumtilFda(double* vec, double* gbms, double* randomNumber, int size, int index) {

    double bump = 0.000000085;

    vec[index] += bump;

    double upVal = sumtil(vec, gbms, randomNumber, size);

    vec[index] -= 2*bump;

    double downVal = sumtil(vec, gbms,  randomNumber, size);

    double sensi = (upVal - downVal)/(2*bump);

    vec[index] += bump;

    return sensi;
}

typedef double (*f_ptr)(double*, double*, double*, int);

extern void __device__ __enzyme_autodiffCuda(f_ptr,
                                         int, double*, double*,
                                         int, double*,
                                         int, double*,
                                         int, int
);


__global__ void computeEnzymeGrad( double* d_vec, double* d_vec_res, double* d_gbms, double* d_randomNumbers, int maturity) {

    __enzyme_autodiffCuda(sumtil,enzyme_dup, d_vec, d_vec_res, enzyme_const, d_gbms, enzyme_const,d_randomNumbers, enzyme_const, maturity);

}


int main() {

    std::random_device rd{};

    std::mt19937 gen(42);

    std::normal_distribution<double> norm(5.15, 2.85);
    std::normal_distribution<double> normTwo(1.15, 0.65);
    std::normal_distribution<double> normThree(5.0, 3.0);

    size_t bytes = N*sizeof(double);

    double *vec = (double*)malloc(bytes);
    double *gbms = (double*)malloc(bytes);
    double *rands = (double*)malloc(bytes);
    double *results_x = (double*)malloc(bytes);

    double *device_vec, *device_gbms, *device_rands, *device_der_vec;

    cudaCheck(cudaMalloc(&device_vec, bytes));
    cudaCheck(cudaMalloc(&device_gbms, bytes));
    cudaCheck(cudaMalloc(&device_rands, bytes));
    cudaCheck(cudaMalloc(&device_der_vec, bytes));

    for (int i = 0; i < N; i++){
        vec[i] = norm(gen);
        gbms[i] = 0.0;
        rands[i] = normThree(gen);
        results_x[i] = 0.0;
    }

    int n = N;

    cudaCheck(cudaMemcpy(device_vec, vec, bytes, cudaMemcpyHostToDevice));
    cudaCheck(cudaMemcpy(device_gbms, gbms, bytes, cudaMemcpyHostToDevice));
    cudaCheck(cudaMemcpy(device_rands, rands, bytes, cudaMemcpyHostToDevice));
    cudaCheck(cudaMemcpy(device_der_vec, results_x, bytes, cudaMemcpyHostToDevice));

    computeEnzymeGrad<<<1,1>>>(device_vec, device_der_vec, device_gbms, device_rands,n);
    cudaCheck(cudaDeviceSynchronize());

    cudaCheck(cudaMemcpy(results_x, device_der_vec, bytes, cudaMemcpyDeviceToHost));

    //=================================================== FDA & HOST CHECK ====================================================

    double* fdaResults = (double*)malloc(bytes);
    double* EnzymeHost = (double*)malloc(bytes);

    for (int i = 0; i < N; i++) {
        fdaResults[i] = 0.0;
    }

    for (int i = 0; i < N; i++) {
        fdaResults[i] = sumtilFda(vec, gbms, rands, n, i);
    }

    __enzyme_autodiff((void*)sumtil,enzyme_dup, vec, EnzymeHost, enzyme_const, gbms, enzyme_const, rands, enzyme_const, n);

    for (int i = 0; i < N; i++) {
        printf("\nx[%d]='%f';\n", i, vec[i]);
        printf("FDA for grad_x[%d]='%.18f';\n", i, fdaResults[i]);
        printf("AAD CUDA Enzyme for grad_x[%d]='%.18f';\n", i, results_x[i]);
        printf("AAD Enzyme Host for grad_x[%d]='%.18f';\n", i, EnzymeHost[i]);
    }

    free(vec);
    free(gbms);
    free(rands);
    free(results_x);
    free(fdaResults);

    cudaFree(device_vec);
    cudaFree(device_gbms);
    cudaFree(device_rands);
    cudaFree(device_der_vec);
}

//=======================================================================================================================