#include <stdio.h>

__global__ void kernel(int * arraydevice, int length)
{ 
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  int numthreads = blockDim.x * gridDim.x;

  for(int i=index; i<length; i+=numthreads)
    arraydevice[i] *= arraydevice[i];
}

void check(int * array, int length)
{
  for(int i=0; i<length; i++)
  {
    if(array[i] != i*i)
    {
      printf("error at i=%d\n",i);
      exit(1);
    }
  }
}

int main()
{
  int * arrayhostA;
  int * arrayhostB;
  int * arraydeviceA;
  int * arraydeviceB;
  int length = 1<<10;
  int size = length*sizeof(int);

  // Pinned memory allocation on host.
  cudaMallocHost(&arrayhostA,size);
  cudaMallocHost(&arrayhostB,size);
  cudaMalloc(&arraydeviceA,size);
  cudaMalloc(&arraydeviceB,size);

  for(int i=0; i<length; i++)
  {
    arrayhostA[i] = i;
    arrayhostB[i] = i;
  }

  cudaStream_t streamA, streamB;
  cudaStreamCreate(&streamA);
  cudaStreamCreate(&streamB);

  // Launch streamA.
  cudaMemcpyAsync(arraydeviceA, arrayhostA, size, cudaMemcpyHostToDevice, streamA);
  kernel<<<3, 128, 0, streamA>>>(arraydeviceA, length);
  cudaMemcpyAsync(arrayhostA, arraydeviceA, size, cudaMemcpyDeviceToHost, streamA);

  // Launch streamB.
  cudaMemcpyAsync(arraydeviceB, arrayhostB, size, cudaMemcpyHostToDevice, streamB);
  kernel<<<3, 128, 0, streamB>>>(arraydeviceB, length);
  cudaMemcpyAsync(arrayhostB, arraydeviceB, size, cudaMemcpyDeviceToHost, streamB);

  // Wait for streamA to finish.
  cudaStreamSynchronize(streamA);
  check(arrayhostA,length);

  // Wait for streamB to finish.
  cudaStreamSynchronize(streamB);
  check(arrayhostB,length);

  cudaStreamDestroy(streamA);
  cudaStreamDestroy(streamB);

  return 0;
}