#include <hip/hip_runtime.h>
#define NBINS 256
__constant__ int c_binMax[1];
__device__ __forceinline__ int clamp_bin(int v) {
if (v < 0) return 0;
if (v > c_binMax[0]) return c_binMax[0];
return v;
}
__global__ void histogram(const int* data, int n, int* bins) {
__shared__ int local_bins[NBINS];
int tid = threadIdx.x;
int bid = blockIdx.x * blockDim.x + tid;
for (int i = tid; i < NBINS; i += blockDim.x) {
local_bins[i] = 0;
}
__syncthreads();
for (int i = bid; i < n; i += gridDim.x * blockDim.x) {
int b = clamp_bin(data[i]);
atomicAdd(&local_bins[b], 1);
}
__syncthreads();
for (int i = tid; i < NBINS; i += blockDim.x) {
if (local_bins[i] > 0) {
atomicAdd(&bins[i], local_bins[i]);
atomicMin(&bins[i], c_binMax[0]);
atomicMax(&bins[i], 0);
}
}
}
__global__ void warp_reduce(const int* in, int* out) {
int lane = __laneid()();
int v = in[threadIdx.x];
for (int offset = warpSize / 2; offset > 0; offset /= 2) {
int t = __shfl_sync(0xFFFFFFFFu, v, lane - offset);
v += t;
}
__syncwarp(0xFFFFFFFFu);
if (lane == 0) {
out[blockIdx.x] = v;
}
}
int main(void) {
const int N = 1 << 20;
int* d_data = nullptr;
int* d_bins = nullptr;
int* d_out = nullptr;
hipMalloc((void**)&d_data, N * sizeof(int));
hipMalloc((void**)&d_bins, NBINS * sizeof(int));
hipMalloc((void**)&d_out, 1024 * sizeof(int));
hipMemset(d_bins, 0, NBINS * sizeof(int));
hipMemcpy(d_data, d_data, N * sizeof(int), cudaMemcpyDeviceToDevice);
int maxbin = NBINS - 1;
hipMemcpyToSymbol(c_binMax, &maxbin, sizeof(int));
dim3 grid(N / 256);
dim3 block(256);
hipLaunchKernelGGL(histogram, dim3(grid), dim3(block), 0, 0, d_data, N, d_bins);
hipLaunchKernelGGL(warp_reduce, dim3(N / 32), dim3(32), 0, 0, d_data, d_out);
hipDeviceSynchronize();
hipFree(d_data);
hipFree(d_bins);
hipFree(d_out);
return 0;
}