#include <cuda_runtime_api.h>
#include <memory.h>
#include <cstdlib>
#include <ctime>
#include <stdio.h>
#include <cuda/cmath>

__global__ void vecAdd(float* A, float* B, float* C, int vectorLength) {
  int workIndex = threadIdx.x + blockIdx.x*blockDim.x;
  if(workIndex < vectorLength) {
    C[workIndex] = A[workIndex] + B[workIndex];
  }
}

void initArray(float* A, int length) {
  std::srand(std::time({}));
  for(int i = 0; i < length; i++) {
    A[i] = rand() / (float)RAND_MAX;
  }
}

void serialVecAdd(float* A, float* B, float* C, int length) {
  for(int i = 0; i < length; i++) {
    C[i] = A[i] + B[i];
  }
}

bool vectorApproximatelyEqual(float* A, float* B, int length, float epsilon = 0.00001) {
  for(int i = 0; i < length; i++) {
    if(fabs(A[i] - B[i]) > epsilon) {
      printf("Index %d mismatch: %f != %f", i, A[i], B[i]);
      return false;
    }
  }
  return true;
}

//explicit-memory-begin
void explicitMemExample(int vectorLength) {
  // Pointers for host memory
  float* A = nullptr;
  float* B = nullptr;
  float* C = nullptr;
  float* comparisonResult = (float*)malloc(vectorLength * sizeof(float));
    
  // Pointers for device memory
  float* devA = nullptr;
  float* devB = nullptr;
  float* devC = nullptr;

  // Allocate Host Memory using cudaMallocHost API. This is best practice
  // when buffers will be used for copies between CPU and GPU memory
  cudaMallocHost(&A, vectorLength * sizeof(float));
  cudaMallocHost(&B, vectorLength * sizeof(float));
  cudaMallocHost(&C, vectorLength * sizeof(float));

  // Initialize vectors on the host
  initArray(A, vectorLength);
  initArray(B, vectorLength);

  // start-allocate-and-copy
  // Allocate memory on the GPU
  cudaMalloc(&devA, vectorLength * sizeof(float));
  cudaMalloc(&devB, vectorLength * sizeof(float));
  cudaMalloc(&devC, vectorLength * sizeof(float));

  // Copy data to the GPU
  cudaMemcpy(devA, A, vectorLength * sizeof(float), cudaMemcpyDefault);
  cudaMemcpy(devB, B, vectorLength * sizeof(float), cudaMemcpyDefault);
  cudaMemset(devC, 0, vectorLength * sizeof(float));
  // end-allocate-and-copy

  // Launch the kernel
  int threads = 256;
  int blocks = cuda::ceil_div(vectorLength, threads);
  vecAdd<<<blocks, threads>>>(devA, devB, devC, vectorLength);

  // wait for kernel execution to complete
  cudaDeviceSynchronize();

  // Copy results back to host
  cudaMemcpy(C, devC, vectorLength * sizeof(float), cudaMemcpyDefault);

  // Perform computation serially on CPU for comparison
  serialVecAdd(A, B, comparisonResult, vectorLength);

  // Confirm that CPU and GPU got the same answer
  if(vectorApproximatelyEqual(C, comparisonResult, vectorLength)) {
    printf("Explicit Memory: CPU and GPU answers match\n");
  }
  else {
    printf("Explicit Memory: Error - CPU and GPU answers to not match\n");
  }

  // clean up
  cudaFree(devA);
  cudaFree(devB);
  cudaFree(devC);
  cudaFreeHost(A);
  cudaFreeHost(B);
  cudaFreeHost(C);
  free(comparisonResult);
}
//explicit-memory-end


int main(int argc, char** argv) {
  int vectorLength = 1024;
  if(argc >= 2) {
    vectorLength = std::atoi(argv[1]);
  }
  explicitMemExample(vectorLength);		
  return 0;
}

// ----- example.cu -----

#include <stdio.h>

__global__ void kernel() {
  printf("Hello from kernel\n");
}

void kernel_launcher() {
  kernel<<<1, 1>>>();
  cudaDeviceSynchronize();
}

int main() {
  kernel_launcher();
  return 0;
}