Use CUB by default
import os
+from cupy_backends.cuda.api cimport runtime
+
+
cdef list _elementwise_accelerators = []
cdef list _reduction_accelerators = []
cdef list _routine_accelerators = []
@@ -43,7 +46,8 @@ def get_routine_accelerators():
cdef _set_default_accelerators():
- cdef str b, accelerator_names = os.getenv('CUPY_ACCELERATORS', '')
+ cdef str b, accelerator_names = os.getenv(
+ 'CUPY_ACCELERATORS', '' if runtime._is_hip_environment else 'cub')
cdef list accelerators = [b for b in accelerator_names.split(',') if b]
set_elementwise_accelerators(accelerators)
set_reduction_accelerators(accelerators)
