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

/* number of blocks (max 65,536) */
#define BLOCKS 50000

/* number of threads per block (max 1,024) */
#define THREADS 200

/* N is now the number of blocks times the number of threads 
   in this case it is 10 million */
#define N (BLOCKS * THREADS)

/* the kernel fills an array up with square roots */
__global__ void square_roots(double* array) {
    /* find my array index
       this now uses three values:
       blockIdx.x which is the block ID we are on
       blockDim.x which is the number of blocks in the X dimension
       threadIdx.x which is the thread ID we are */
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;

    /* compute the square root of this number and store it in the array
       CUDA provides a sqrt function as part of its math library, so this
       call uses that - not the one from math.h which is a CPU function */
    array[idx] = sqrt((double) idx);
}

int main() {
    /* store the square roots of 0 to (N-1) on the CPU
     * stored on the heap since it's too big for the stack for large values of N */
    double* roots = (double*) malloc(N * sizeof(double));

    /* allocate a GPU array to hold the square roots */
    double* gpu_roots;
    cudaMalloc((void**) &gpu_roots, N * sizeof(double));

    /* invoke the GPU to calculate the square roots
       now, we don't run all N blocks, because that may be more than 65,535 */
    square_roots<<<BLOCKS, THREADS>>>(gpu_roots);

    /* copy the data back */
    cudaMemcpy(roots, gpu_roots, N * sizeof(double), cudaMemcpyDeviceToHost);

    /* free the memory */
    cudaFree(gpu_roots);

    /* print out 100 evenly spaced square roots just to see that it worked */
    unsigned int i;
    for (i = 0; i < N; i += (N / 100)) {
        printf("sqrt(%d) = %lf\n", i, roots[i]);
    }

    /* free the CPU memory */
    free(roots);

    return 0;
}