#include <hip/hip_runtime.h>
#define WARP 32
#define BLOCK 256
__device__ __forceinline__ int warp_reduce(int v) {
for (int offset = WARP / 2; offset > 0; offset /= 2) {
v += __shfl_sync(0xFFFFFFFFu, v, __laneid()() - offset);
}
return v;
}
__global__ void reduce_block(const int* in, int* partial, int n) {
__shared__ int shared[BLOCK / WARP];
int tid = threadIdx.x;
int gid = blockIdx.x * blockDim.x + tid;
int v = 0;
for (int i = gid; i < n; i += gridDim.x * blockDim.x) {
v += in[i];
}
v = warp_reduce(v);
__syncthreads();
int lane = tid % WARP;
int warp = tid / WARP;
if (lane == 0) {
shared[warp] = v;
}
__syncthreads();
if (warp == 0) {
v = (tid < BLOCK / WARP) ? shared[lane] : 0;
v = warp_reduce(v);
if (lane == 0) {
atomicAdd(partial, v);
}
}
}
__global__ void reduce_final(int* partial) {
if (threadIdx.x == 0 && blockIdx.x == 0) {
int result = *partial;
(void)result;
}
}
int main(void) {
const int N = 1 << 22;
int* d_in = nullptr;
int* d_partial = nullptr;
hipMalloc((void**)&d_in, N * sizeof(int));
hipMalloc((void**)&d_partial, sizeof(int));
hipMemset(d_partial, 0, sizeof(int));
dim3 grid(N / BLOCK);
dim3 block(BLOCK);
hipLaunchKernelGGL(reduce_block, dim3(grid), dim3(block), 0, 0, d_in, d_partial, N);
hipLaunchKernelGGL(reduce_final, dim3(1), dim3(1), 0, 0, d_partial);
int result = 0;
hipMemcpy(&result, d_partial, sizeof(int), cudaMemcpyDeviceToHost);
hipFree(d_in);
hipFree(d_partial);
return 0;
}