/* dot product of two arrays.
 *  
 */
#include <civlc.h>
#include <stdio.h>

#define imin(a,b) (a<b?a:b)

$input int BNPG;
$input int K;
$input int TNPB; // thread number per block: must be a power of 2
int N = (BNPG + 1) * K;
int threadsPerBlock = TNPB;
int blocksPerGrid =
            imin(BNPG, (N+threadsPerBlock-1) / threadsPerBlock );     
double *a, *b, c, *partial_c;

void gpu(){
  void gpuBlock(int blockID){
    #include "barrier.cvh"
    
    double cache[threadsPerBlock];
    int in_barrier[threadsPerBlock];
    
    void gpuThread(int threadID){
      int tid = threadID + blockID * threadsPerBlock;
      int cacheIndex = threadID;
      double temp = 0;
      
      $atomic {
        while (tid < N) {
          temp += a[tid] * b[tid];
          tid += threadsPerBlock * blocksPerGrid;
        }
        // set cache values
        cache[cacheIndex] = temp;
      }
      // synchronize threads
      barrier(in_barrier, threadID);
      int i = threadsPerBlock/2;
      while (i != 0) {
        if (cacheIndex < i)
          cache[cacheIndex] += cache[cacheIndex + i];
        // synchronize threads
        barrier(in_barrier, threadID);
        i /= 2;
      }
      if (cacheIndex == 0)
        partial_c[blockID] = cache[0];
    }
    
    $proc threads[threadsPerBlock];

    barrier_init(in_barrier, threadsPerBlock);
    for(int i = 0; i < threadsPerBlock; i++) {
      threads[i] = $spawn gpuThread(i);
    }

    for(int i = 0; i < threadsPerBlock; i++) {
      $wait threads[i];
    }
  }
  
  $proc blocks[blocksPerGrid];

  for(int i = 0; i < blocksPerGrid; i++) {
    blocks[i] = $spawn gpuBlock(i);  
  }

  for(int i = 0; i < blocksPerGrid; i++) {
    $wait blocks[i];
  }
}

int main( void ) {
  $heap host;
    
  // allocate memory on the cpu side
  a = (double *) $malloc(&host, N*sizeof(double));
  b = (double *) $malloc(&host, N*sizeof(double));
  partial_c = (double *) $malloc(&host, blocksPerGrid*sizeof(double));
  // fill in the host memory with data
  for (int i=0; i<N; i++) {
    a[i] = i;
    b[i] = i*2;
  }
  gpu();
  // finish up on the CPU side
  c = 0;
  for (int i=0; i<blocksPerGrid; i++) {
    c += partial_c[i];
  }
  #define sum_squares(x)  (x*(x+1)*(2*x+1)/6)
  // check result
  $assert(c == 2 * sum_squares( (double)(N - 1) ));
}