Simplify reduction core code
@@ -27,20 +27,19 @@ cpdef _get_simple_reduction_kernel(
typedef ${reduce_type} _type_reduce;
extern "C" __global__ void ${name}(${params}) {
- extern __shared__ _type_reduce _sdata_raw[];
- _type_reduce *_sdata = _sdata_raw;
+ __shared__ _type_reduce _sdata[${block_size}];
unsigned int _tid = threadIdx.x;
- int _J_offset = _tid / _block_stride;
+ int _J_offset = _tid >> __popc(_block_stride - 1); // _tid / _block_stride
int _j_offset = _J_offset * _out_ind.size();
- int _J_stride = ${block_size} / _block_stride;
+ int _J_stride = ${block_size} >> __popc(_block_stride - 1);
long long _j_stride = (long long)_J_stride * _out_ind.size();
for (int _i_base = blockIdx.x * _block_stride;
_i_base < _out_ind.size();
_i_base += gridDim.x * _block_stride) {
_type_reduce _s = _type_reduce(${identity});
- int _i = _i_base + _tid % _block_stride;
+ int _i = _i_base + (_tid & (_block_stride - 1)); // _tid % _block_stride
int _J = _J_offset;
for (long long _j = _i + _j_offset; _j < _in_ind.size();
_j += _j_stride, _J += _J_stride) {
@@ -52,34 +51,17 @@ cpdef _get_simple_reduction_kernel(
if (_block_stride < ${block_size}) {
_sdata[_tid] = _s;
__syncthreads();
- if (_block_stride <= 256) {
- _REDUCE(256);
- __syncthreads();
- if (_block_stride <= 128) {
- _REDUCE(128);
- __syncthreads();
- if (_block_stride <= 64) {
- _REDUCE(64);
- __syncthreads();
- if (_block_stride <= 32) {
- _REDUCE(32);
- if (_block_stride <= 16) {
- _REDUCE(16);
- if (_block_stride <= 8) {
- _REDUCE(8);
- if (_block_stride <= 4) {
- _REDUCE(4);
- if (_block_stride <= 2) {
- _REDUCE(2);
- if (_block_stride <= 1) {
- _REDUCE(1);
- }
- }
- }
- }
- }
+ for (unsigned int _block = ${block_size} / 2;
+ _block >= max(64, _block_stride); _block >>= 1) {
+ if (_tid < _block) {
+ _REDUCE(_block);
}
+ __syncthreads();
}
+ if (_tid < 32) {
+ for (unsigned int _block = 32;
+ _block >= _block_stride; _block >>= 1) {
+ _REDUCE(_block);
}
}
_s = _sdata[_tid];
