#include <stdio.h>

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

const int N = 33 * 1024;
const int threadsPerBlock = 256;
const int blocksPerGrid =
            imin( 32, (N+threadsPerBlock-1) / threadsPerBlock );

__global__ void gpusum(const float *a, float *partial_sum) {
	__shared__ float cache[threadsPerBlock];

	int tid = threadIdx.x + blockIdx.x * blockDim.x;
	int cacheIndex = threadIdx.x;

	float temp = 0;
	while (tid < N) {
		temp += a[tid];
		tid += blockDim.x * gridDim.x;
	}

    // set the cache values
	cache[cacheIndex] = temp;
    // synchronize threads in this block
	__syncthreads();

    // for reductions, threadsPerBlock must be a power of 2
    // because of the following code
	int i = blockDim.x / 2;
	while (i != 0) {
		if (cacheIndex < i)
			cache[cacheIndex] += cache[cacheIndex + i];
		__syncthreads();
		i /= 2;
	}

	if (cacheIndex == 0)
		partial_sum[blockIdx.x] = cache[0];
}

int main(void) {
	float *a, *partial, total, expected_total;
	float *dev_a, *dev_partial;

    // allocate memory on the cpu side
	a = (float *)malloc(N * sizeof(float));
	partial = (float *)malloc(blocksPerGrid * sizeof(float));

    // allocate the memory on the GPU
	cudaMalloc((void **)&dev_a, N * sizeof(float));
	cudaMalloc((void **)&dev_partial, blocksPerGrid * sizeof(float));

    // fill in the host memory with data
	for (int i = 0; i < N; i++) {
		a[i] = i;
	}

    // copy the arrays 'a' to the GPU
	cudaMemcpy(dev_a, a, N * sizeof(float), cudaMemcpyHostToDevice);

	gpusum<<<blocksPerGrid, threadsPerBlock>>>(dev_a, dev_partial);

    // copy the array back from the GPU to the CPU
	cudaMemcpy(partial, dev_partial, blocksPerGrid * sizeof(float), 
               cudaMemcpyDeviceToHost);

    // finish up on the CPU side
    total = 0;
	for (int i = 0; i < blocksPerGrid; i++) {
		total += partial[i];
	}

    // sum of array [0, 1, ..., N-1]
	expected_total = (0 + N-1) * N / 2;
	printf("GPU sum = %f, expected = %f\n", total, expected_total);

	cudaFree(dev_a);
	cudaFree(dev_partial);
	free(a);
	free(partial);

	return 0;
}
