fix last segment size corruption in HIP
On CUDA the size is always correct as all threads writing to the same
location will be reduced to one thread writng. Such a behavior is not
guaranteed in HIP, though, so we better copy the value to the register.
@@ -131,21 +131,22 @@ __global__ void ${name}(${params}) {
int tile_size = (BLOCK_SIZE * ITEMS_PER_THREAD < _segment_size ?
BLOCK_SIZE * ITEMS_PER_THREAD :
_segment_size);
+ sizeT _seg_size = _segment_size;
#if defined FIRST_PASS
// for two-pass reduction only: "last segment" is special
if (_array_size > 0) {
if (_array_size - segment_idx <= _segment_size) {
- _segment_size = _array_size - segment_idx;
+ _seg_size = _array_size - segment_idx;
}
}
#endif
// loop over tiles within 1 segment
_type_reduce aggregate = _type_reduce(${identity});
- for (i = 0; i < _segment_size; i += BLOCK_SIZE * ITEMS_PER_THREAD) {
+ for (i = 0; i < _seg_size; i += BLOCK_SIZE * ITEMS_PER_THREAD) {
// for the last tile
- if (_segment_size - i <= tile_size) { tile_size = _segment_size - i; }
+ if (_seg_size - i <= tile_size) { tile_size = _seg_size - i; }
'''
if pre_map_expr == 'in0':
@@ -166,7 +167,7 @@ __global__ void ${name}(${params}) {
#if defined FIRST_PASS
int _J = (segment_idx + i + e_idx);
#else // only one pass
- int _J = (segment_idx + i + e_idx) % _segment_size;
+ int _J = (segment_idx + i + e_idx) % _seg_size;
#endif
if (e_idx < tile_size) {
@@ -636,9 +637,9 @@ cdef bint _try_to_call_cub_reduction(
size_type = 'uint64'
else:
size_type = 'int32'
+ type_map = _kernel._TypeMap(type_map._pairs + (('sizeT', size_type),))
params = (self._params[0:2]
- + _get_param_info(
- size_type + '_segment_size', not full_reduction)
+ + _get_param_info(size_type + ' _segment_size', True)
+ _get_param_info(size_type + ' _array_size', True))
# HACK for ReductionKernel:
