github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

block_reduce.h fix format (#581)

pull/673/head
DouJS 2022-03-31 17:13:09 +08:00 committed by binmakeswell
parent d2dc6049b5
commit f08fc17f2b
1 changed files with 20 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
colossalai/kernel/cuda_native/csrc/kernels/include/block_reduce.h
View File

@ -13,23 +13,22 @@ const float REDUCE_FLOAT_INF_NEG = -100000000.f;
const float REDUCE_FLOAT_INF_POS = 100000000.f; const float REDUCE_FLOAT_INF_POS = 100000000.f;
const unsigned int WARP_REDUCE_SIZE = 32; const unsigned int WARP_REDUCE_SIZE = 32;
template <typename T> template <typename T> __forceinline__ __device__ T warpReduceSum(T val) {
__forceinline__ __device__ T warpReduceSum(T val) {
for (int mask = (WARP_REDUCE_SIZE >> 1); mask > 0; mask >>= 1) for (int mask = (WARP_REDUCE_SIZE >> 1); mask > 0; mask >>= 1)
val += __shfl_xor_sync(WARP_REDUCE_MASK, val, mask, WARP_REDUCE_SIZE); val += __shfl_xor_sync(WARP_REDUCE_MASK, val, mask, WARP_REDUCE_SIZE);
return val; return val;
} }
/* Calculate the sum of all elements in a block */ /* Calculate the sum of all elements in a block */
template <typename T> template <typename T> __forceinline__ __device__ T blockReduceSum(T val) {
__forceinline__ __device__ T blockReduceSum(T val) {
static __shared__ T shared[32]; static __shared__ T shared[32];
int lane = threadIdx.x & 0x1f; int lane = threadIdx.x & 0x1f;
int wid = threadIdx.x >> 5; int wid = threadIdx.x >> 5;
val = warpReduceSum<T>(val); val = warpReduceSum<T>(val);
if (lane == 0) shared[wid] = val; if (lane == 0)
shared[wid] = val;
__syncthreads(); __syncthreads();
val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : (T)0.0f; val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : (T)0.0f;