#ifndef prefix_sum
#define prefix_sum
template <typename T>
static __device__ inline T block_prefix_sum(T val, void* buffer) {
const int lane = threadIdx.x % WS;
const int warp = threadIdx.x / WS;
const int warps = TPB / WS;
T* const carry = (T*)buffer;
assert(WS >= warps);
T tmp = __shfl_up(val, 1);
if (lane >= 1) val += tmp;
tmp = __shfl_up(val, 2);
if (lane >= 2) val += tmp;
tmp = __shfl_up(val, 4);
if (lane >= 4) val += tmp;
tmp = __shfl_up(val, 8);
if (lane >= 8) val += tmp;
tmp = __shfl_up(val, 16);
if (lane >= 16) val += tmp;
#if defined(WS) && (WS == 64)
tmp = __shfl_up(val, 32);
if (lane >= 32) val += tmp;
#endif
if (lane == WS - 1) carry[warp] = val;
__syncthreads();
if constexpr (warps > 1) {
if (warp == 0) {
T sum = carry[lane];
T tmp = __shfl_up(sum, 1);
if (lane >= 1) sum += tmp;
if constexpr (warps > 2) {
tmp = __shfl_up(sum, 2);
if (lane >= 2) sum += tmp;
if constexpr (warps > 4) {
tmp = __shfl_up(sum, 4);
if (lane >= 4) sum += tmp;
if constexpr (warps > 8) {
tmp = __shfl_up(sum, 8);
if (lane >= 8) sum += tmp;
if constexpr (warps > 16) {
tmp = __shfl_up(sum, 16);
if (lane >= 16) sum += tmp;
#if defined(WS) && (WS == 64)
if constexpr (warps > 32) {
tmp = __shfl_up(sum, 32);
if (lane >= 32) sum += tmp;
}
#endif
}
}
}
}
carry[lane] = sum;
}
__syncthreads();
if (warp > 0) val += carry[warp - 1];
}
return val;
}
#endif