
#ifndef _SCAN_BEST_KERNEL_CU_
#define _SCAN_BEST_KERNEL_CU_

#include "my_cutil.h"


////////////////////////////////////////////////////////////////////////////////
//   Notes for running in PUG:
//   Bitvector size: 14 bits
//   Need to turn on automatic loop refinement
////////////////////////////////////////////////////////////////////////////////


// Define this to more rigorously avoid bank conflicts, 
// even at the lower (root) levels of the tree
// Note that due to the higher addressing overhead, performance 
// is lower with ZERO_BANK_CONFLICTS enabled.  It is provided
// as an example.
//#define ZERO_BANK_CONFLICTS 

// 16 banks on G80
#define NUM_BANKS 16
#define LOG_NUM_BANKS 4

#ifdef ZERO_BANK_CONFLICTS
#define CONFLICT_FREE_OFFSET(index) ((index) >> LOG_NUM_BANKS + (index) >> (2*LOG_NUM_BANKS))
#else
#define CONFLICT_FREE_OFFSET(index) ((index) >> LOG_NUM_BANKS)
#endif

///////////////////////////////////////////////////////////////////////////////
// Work-efficient compute implementation of scan, one thread per 2 elements
// Work-efficient: O(log(n)) steps, and O(n) adds.
// Also shared storage efficient: Uses n + n/NUM_BANKS shared memory -- no ping-ponging
// Also avoids most bank conflicts using single-element offsets every NUM_BANKS elements.
//
// In addition, If ZERO_BANK_CONFLICTS is defined, uses 
//     n + n/NUM_BANKS + n/(NUM_BANKS*NUM_BANKS) 
// shared memory. If ZERO_BANK_CONFLICTS is defined, avoids ALL bank conflicts using 
// single-element offsets every NUM_BANKS elements, plus additional single-element offsets 
// after every NUM_BANKS^2 elements.
//
// Uses a balanced tree type algorithm.  See Blelloch, 1990 "Prefix Sums 
// and Their Applications", or Prins and Chatterjee PRAM course notes:
// http://www.cs.unc.edu/~prins/Classes/203/Handouts/pram.pdf
// 
// This work-efficient version is based on the algorithm presented in Guy Blelloch's
// excellent paper "Prefix sums and their applications".
// http://www-2.cs.cmu.edu/afs/cs.cmu.edu/project/scandal/public/papers/CMU-CS-90-190.html
//
// Pro: Work Efficient, very few bank conflicts (or zero if ZERO_BANK_CONFLICTS is defined)
// Con: More instructions to compute bank-conflict-free shared memory addressing,
// and slightly more shared memory storage used.
//


__device__ void loadSharedChunkFromMem(float *s_data,
                                       const float *g_idata, 
                                       int n, int baseIndex,
                                       int& ai, int& bi, 
                                       int& mem_ai, int& mem_bi, 
                                       int& bankOffsetA, int& bankOffsetB, int isNP2)
{
    int thid = threadIdx.x;
    mem_ai = baseIndex + threadIdx.x;
    mem_bi = mem_ai + blockDim.x;

    ai = thid;
    bi = thid + blockDim.x;

    // compute spacing to avoid bank conflicts
    bankOffsetA = CONFLICT_FREE_OFFSET(ai);
    bankOffsetB = CONFLICT_FREE_OFFSET(bi);

    // Cache the computational window in shared memory
    // pad values beyond n with zeros

    s_data[ai + bankOffsetA] = g_idata[mem_ai];
    
    if (isNP2 != 0) // compile-time decision
    {
        s_data[bi + bankOffsetB] = (bi < n) ? g_idata[mem_bi] : 0;
    }
    else
    {
        s_data[bi + bankOffsetB] = g_idata[mem_bi];
    }
}

__device__ void storeSharedChunkToMem(float* g_odata, 
                                      const float* s_data,
                                      int n, 
                                      int ai, int bi, 
                                      int mem_ai, int mem_bi,
                                      int bankOffsetA, int bankOffsetB, int isNP2)
{
    __syncthreads();

    // write results to global memory
#ifdef MUTATION
    g_odata[0]      = s_data[ai + bankOffsetA];
#else
    g_odata[mem_ai] = s_data[ai + bankOffsetA];
#endif
    if (isNP2 != 0) // compile-time decision
    {
        if (bi < n)
            g_odata[mem_bi] = s_data[bi + bankOffsetB]; 
    }
    else
    {
        g_odata[mem_bi] = s_data[bi + bankOffsetB]; 
    }
}

__device__ void clearLastElement(float* s_data, 
                                 float *g_blockSums, 
                                 int blockIndex, int storeSum)
{
    if (threadIdx.x == 0)
    {
        int index = (blockDim.x << 1) - 1;
        index += CONFLICT_FREE_OFFSET(index);
        
        if (storeSum != 0) // compile-time decision
        {
            // write this block's total sum to the corresponding index in the blockSums array
            g_blockSums[blockIndex] = s_data[index];
        }

        // zero the last element in the scan so it will propagate back to the front
        s_data[index] = 0;
    }
}



__device__ unsigned int buildSum(float *s_data)
{
    unsigned int thid = threadIdx.x;
    unsigned int stride = 1;
    
    // build the sum in place up the tree
    for (int d = blockDim.x; d > 0; d >>= 1)
    {
        __syncthreads();

        stride *= 2;

        if (thid < d)
        {
	    int i  = stride * thid;
            int ai = i + stride/2 - 1;
            int bi = ai + stride/2;

            ai += CONFLICT_FREE_OFFSET(ai);
            bi += CONFLICT_FREE_OFFSET(bi);

#ifdef MUTATION
            s_data[0] += s_data[ai];
#else
            s_data[bi] += s_data[ai];
#endif
	    // RACE(s_data);

        }
    }

    return stride;
}


// for testing
__device__ unsigned int buildSumKnel(float *s_data)
{
    unsigned int thid = threadIdx.x;
    unsigned int stride = 1;
    
    // build the sum in place up the tree
    for (int d = blockDim.x; d > 0; d >>= 1)
    {
        __syncthreads();

        stride *= 2;

        if (thid < d)
        {
	    int i  = stride * thid;
            int ai = i + stride/2 - 1;
            int bi = ai + stride/2;

            ai += CONFLICT_FREE_OFFSET(ai);
            bi += CONFLICT_FREE_OFFSET(bi);

            s_data[bi] += s_data[ai];
	    s_data[bi] = 1;
        }
    }

    return stride;
}


__device__ void scanRootToLeaves(float *s_data, unsigned int stride)
{
     unsigned int thid = threadIdx.x;

     // assume (stride == blockDim.x);
    // traverse down the tree building the scan in place

    for (int d = 1; d <= blockDim.x; d *= 2)
    {
        stride >>= 1;

        __syncthreads();

        if (thid < d)
        {
           int i  =  2 * stride * thid;
	   int ai = i + stride - 1;
            int bi = ai + stride;
	    ai += CONFLICT_FREE_OFFSET(ai);
	    bi += CONFLICT_FREE_OFFSET(bi);

	    float t  = s_data[ai];
	    s_data[ai] = s_data[bi];
	    s_data[bi] += t; 
        }
	// RACE(s_data);
    }
}


__device__ void prescanBlock(float *data, int blockIndex, float *blockSums, int storeSum)
{
    int stride = buildSum(data);               // build the sum in place up the tree
    clearLastElement(data, blockSums, 
		     (blockIndex == 0) ? blockIdx.x : blockIndex, storeSum);
    scanRootToLeaves(data, stride);            // traverse down tree to build the scan
}


// __device__ void prescanBlockKernel(float *data, int blockIndex, float *blockSums, int storeSum)
// {
//     int stride = buildSum(data);               // build the sum in place up the tree
//     clearLastElement(data, blockSums, 
// 		     (blockIndex == 0) ? blockIdx.x : blockIndex, storeSum);
//     scanRootToLeaves(data, stride);            // traverse down tree to build the scan
// }


__global__ void uniformAdd_kernel(float *g_data, 
                           float *uniforms, 
                           int n, 
                           int blockOffset, 
                           int baseIndex)
{
    __shared__ float uni;
    if (threadIdx.x == 0)
        uni = uniforms[blockIdx.x + blockOffset];
    
    unsigned int address = blockIdx.x * (blockDim.x << 1) + baseIndex + threadIdx.x; 

    __syncthreads();
    
    // note two adds per thread
#ifdef MUTATION
    g_data[0]                    += uni;
#else
    g_data[address]              += uni;
#endif
    g_data[address + blockDim.x] += (threadIdx.x + blockDim.x < n) * uni;
}


__global__ void prescankernel_kernel(float *g_odata, 
                        const float *g_idata, 
                        float *g_blockSums, 
                        int n, 
                        int blockIndex, 
                        int baseIndex,
                        int storeSum, int isNP2)
{
    int ai, bi, mem_ai, mem_bi, bankOffsetA, bankOffsetB;
    extern __shared__ float s_data[];

    // load data into shared memory
    loadSharedChunkFromMem((float*) s_data, (float*) g_idata, n,
			   (baseIndex == 0) ?
			   blockIdx.x * (blockDim.x << 1) : baseIndex,
			   ai, bi, mem_ai, mem_bi,
			   bankOffsetA, bankOffsetB,
			   isNP2);
     
    // scan the data in each block
    prescanBlock((float*) s_data, blockIndex, g_blockSums, storeSum);

    // write results to device memory 
    storeSharedChunkToMem((float*) g_odata, (float*) s_data, n, 
 			  ai, bi, mem_ai, mem_bi,  
 			  bankOffsetA, bankOffsetB, isNP2);
}



#endif // #ifndef _SCAN_BEST_KERNEL_CU_

