#define NUM_THREADS (1<< NUM_LEVELS)

// 32 channels means shift five and add
#define BANK_ADDRESS(i) (i + (i>> 5))  

#define LOCAL_SIZE (BANK_ADDRESS(NUM_THREADS))

void ThreadSum(uint tid, __local uint sharedSum[LOCAL_SIZE]) {
   uint tid2 = BANK_ADDRESS(tid);
   
   for(uint d = 0; d < NUM_LEVELS - 1; ++d) {
      barrier(CLK_LOCAL_MEM_FENCE);
      uint mask = (2<< d) - 1;
      uint offset = 1<< d;
      if(mask == (mask & tid))
         sharedSum[tid2] += sharedSum[BANK_ADDRESS(tid - offset)];
   }
   barrier(CLK_LOCAL_MEM_FENCE);

   if(0 == tid) {
      uint ai = BANK_ADDRESS(NUM_THREADS / 2 - 1);
      uint bi = BANK_ADDRESS(NUM_THREADS - 1);
      
      uint at = sharedSum[ai];

      sharedSum[ai] += sharedSum[bi];
      sharedSum[bi] += at + at;
   }
    
   for(uint d = NUM_LEVELS - 1; d; --d) {
      barrier(CLK_LOCAL_MEM_FENCE);
      uint mask = (1<< d) - 1;
      uint offset = 1<< (d - 1);
      if(mask == (mask & tid)) {
         uint t = sharedSum[tid2];
         uint r = BANK_ADDRESS(tid - offset);
         sharedSum[tid2] += sharedSum[r];
         sharedSum[r] = t;
      }
   }
   barrier(CLK_LOCAL_MEM_FENCE);
}


uint4 Inclusive4Sum(uint4 vec) {
   vec.y += vec.x;
   vec.z += vec.y;
   vec.w += vec.z;
   return vec;
}
uint8 Inclusive8Sum(uint8 vec) {
   uint8 result;
   result.lo = Inclusive4Sum(vec.lo);
   result.hi = Inclusive4Sum(vec.hi) + result.lo.w;
   return result;
}
       
__kernel __attribute__((reqd_work_group_size(NUM_THREADS, 1, 1)))
void PrefixSumBlock_Pass1(
   __global uint* pass1_values,
   __global uint* pass1_partialSums) {

   __local uint sharedSum[LOCAL_SIZE];
   global uint8* pass1_values_vec8 = (__global uint8*)pass1_values;

   uint tid = get_local_id(0);
   uint gid = get_group_id(0);
 
   uint index = NUM_THREADS * gid + tid;

   uint8 a = pass1_values_vec8[index];
   uint8 aInc = Inclusive8Sum(a);

   uint tid2 = BANK_ADDRESS(tid);
   sharedSum[tid2] = aInc.s7;

   ThreadSum(tid, sharedSum);
   
   uint total = sharedSum[BANK_ADDRESS(0)];
   uint aExc = sharedSum[tid2] - total;
   
   uint8 sum = aInc - a + aExc;

   pass1_values_vec8[index] = sum;

   if(0 == tid)
      pass1_partialSums[gid] = total;   
}
///////////////////////////////////////////////////////////////////////////////////////////////////
// Finalize
   
__kernel __attribute__((reqd_work_group_size(NUM_THREADS, 1, 1)))
void PrefixSumBlock_Finalize(
   __global uint* finalize_values) {
   
   __local uint sharedSum[LOCAL_SIZE]; 
   __global uint8* finalize_values_vec8 = (__global uint8*)finalize_values;
   
   uint tid = get_local_id(0);
   uint gid = get_group_id(0);
 
   uint index = NUM_THREADS * gid + tid;
      
   uint8 a = finalize_values_vec8[index];
   uint8 aInc = Inclusive8Sum(a);
   
   uint tid2 = BANK_ADDRESS(tid);
   sharedSum[tid2] = aInc.s7;
   
   ThreadSum(tid, sharedSum);

   uint total = sharedSum[BANK_ADDRESS(0)];
   uint exc = sharedSum[tid2] - total;
   
   finalize_values_vec8[index] = aInc - a + (uint8)exc;
}

   
///////////////////////////////////////////////////////////////////////////////////////////////////
// Pass 2
__kernel __attribute__((reqd_work_group_size(NUM_THREADS, 1, 1)))
void PrefixSumBlock_Pass2(
   __global uint* pass2_values,
   __global const uint* pass2_offsets) {

   __global uint8* pass2_values_vec8 = (__global uint8*)pass2_values;

   uint tid = get_local_id(0);
   uint gid = get_group_id(0);

   uint index = NUM_THREADS * gid + tid;
   uint8 a = pass2_values_vec8[index];

   uint partialSum = pass2_offsets[gid];

   pass2_values_vec8[index] = a + (uint8)partialSum;
}

