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

using namespace std;

void  initVector(float *u, int n, float c) {
  int i;
  for (i=0; i<n; i++)
      u[i] = c;
}

__global__ void gpuVectAdd(float *u, float *v,
  float *z, int N) 
{
  int i = blockIdx.x * blockDim.x + threadIdx.x;

  if (i < N ) {
    int index = i; 
    z[index] = u[index] + v[index];
  }
}

int main(int argc, char *argv[]) {
 
  const int N = 1000;
  cudaError_t error;
  int gpuNumber, gpuID;

  float *u, *v, *z;

  u = (float *) malloc(N*sizeof(float));
  v = (float *) malloc(N*sizeof(float));
  z = (float *) malloc(N*sizeof(float));

  error = cudaGetDeviceCount( &gpuNumber );
  if (error) cerr << "Error with cudaDeviceCount" << endl;
  cout << "Available GPUs: " << gpuNumber << endl;

  for (gpuID = 0; gpuID < gpuNumber; gpuID++) {

    float *u_dev, *v_dev, *z_dev;
    error = cudaMalloc((void **) &u_dev, N*sizeof(float));
    cout << "From GPU: " << gpuID << " " << cudaGetErrorString(error) << endl;
    error = cudaMalloc((void **) &v_dev, N*sizeof(float));
    cout << "From GPU: " << gpuID << " " << cudaGetErrorString(error) << endl;
    error = cudaMalloc((void **) &z_dev, N*sizeof(float));
    cout << "From GPU: " << gpuID << " " << cudaGetErrorString(error) << endl;

    initVector((float *) u, N, 1.0);
    initVector((float *) v, N, 2.0);
    initVector((float *) z, N, 0.0);
 
    error = cudaMemcpy(u_dev, u, N*sizeof(float), cudaMemcpyHostToDevice);
    cout << "From GPU: " << gpuID << " " << cudaGetErrorString(error) << endl;
    error = cudaMemcpy(v_dev, v, N*sizeof(float), cudaMemcpyHostToDevice);
    cout << "From GPU: " << gpuID << " " << cudaGetErrorString(error) << endl;

    dim3 mblockDim(32);
    dim3 mgridDim(N/mblockDim.x + 1);

    gpuVectAdd<<<mgridDim, mblockDim>>>(u_dev, v_dev, z_dev, N);
    cudaDeviceSynchronize(); 
    cout << "From GPU: " << gpuID << " " << cudaGetLastError() << endl; 
  
    error = cudaMemcpy(z, z_dev, N*sizeof(float), cudaMemcpyDeviceToHost);
    cout << "From GPU: " << gpuID << " " << cudaGetErrorString(error) << endl;

    cout << z[0] << " " << z[1] << " " << z[2] << endl;

    cudaFree(u_dev);
    cudaFree(v_dev);
    cudaFree(z_dev);
  }

  return 0;
}