35 lines
793 B
Plaintext
35 lines
793 B
Plaintext
#pragma once
|
|
|
|
namespace cacheflow {
|
|
|
|
template<typename T>
|
|
__inline__ __device__ T warpReduceSum(T val) {
|
|
#pragma unroll
|
|
for (int mask = 16; mask > 0; mask >>= 1)
|
|
val += __shfl_xor_sync(0xffffffff, val, mask, 32);
|
|
return val;
|
|
}
|
|
|
|
/* Calculate the sum of all elements in a block */
|
|
template<typename T>
|
|
__inline__ __device__ T blockReduceSum(T val) {
|
|
static __shared__ T shared[32];
|
|
int lane = threadIdx.x & 0x1f;
|
|
int wid = threadIdx.x >> 5;
|
|
|
|
val = warpReduceSum<T>(val);
|
|
|
|
if (lane == 0)
|
|
shared[wid] = val;
|
|
|
|
__syncthreads();
|
|
|
|
// Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
|
|
// blockDim.x is not divided by 32
|
|
val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f);
|
|
val = warpReduceSum<T>(val);
|
|
return val;
|
|
}
|
|
|
|
} // namespace cacheflow
|