Add nogil to CUB
@@ -43,7 +43,7 @@ CUB_support_dtype = [numpy.int8, numpy.uint8,
# Extern
###############################################################################
-cdef extern from 'cupy_cub.h':
+cdef extern from 'cupy_cub.h' with nogil:
void cub_device_reduce(void*, size_t&, void*, void*, int, Stream_t,
int, int)
void cub_device_segmented_reduce(void*, size_t&, void*, void*, int, void*,
@@ -122,14 +122,28 @@ def device_reduce(ndarray x, op, out=None, bint keepdims=False):
y = ndarray((kv_bytes,), numpy.int8)
x_ptr = <void *>x.data.ptr
y_ptr = <void *>y.data.ptr
- dtype_id = _get_dtype_id(x.dtype)
+ dtype_id = <int>_get_dtype_id(x.dtype)
s = <Stream_t>stream.get_current_stream_ptr()
-
- ws_size = cub_device_reduce_get_workspace_size(x_ptr, y_ptr, x.size, s,
- op, dtype_id)
+ x_size = <int>x.size
+
+ with nogil:
+ ws_size = cub_device_reduce_get_workspace_size(x_ptr,
+ y_ptr,
+ x_size,
+ s,
+ op,
+ dtype_id)
ws = memory.alloc(ws_size)
ws_ptr = <void *>ws.ptr
- cub_device_reduce(ws_ptr, ws_size, x_ptr, y_ptr, x.size, s, op, dtype_id)
+ with nogil:
+ cub_device_reduce(ws_ptr,
+ ws_size,
+ x_ptr,
+ y_ptr,
+ x_size,
+ s,
+ op,
+ dtype_id)
if op == CUPY_CUB_ARGMIN or op == CUPY_CUB_ARGMAX:
# get key from KeyValuePair: need to reinterpret the first 4 bytes
# and then cast it
@@ -197,11 +211,13 @@ def device_segmented_reduce(ndarray x, op, axis, out=None,
dtype_id = _get_dtype_id(x.dtype)
# get workspace size and then fire up
+ with nogil:
ws_size = cub_device_segmented_reduce_get_workspace_size(
x_ptr, y_ptr, n_segments, offset_start_ptr, offset_end_ptr, s,
op, dtype_id)
ws = memory.alloc(ws_size)
ws_ptr = <void*>ws.ptr
+ with nogil:
cub_device_segmented_reduce(ws_ptr, ws_size, x_ptr, y_ptr, n_segments,
offset_start_ptr, offset_end_ptr, s,
op, dtype_id)
@@ -253,11 +269,13 @@ def device_csrmv(int n_rows, int n_cols, int nnz, ndarray values,
dtype_id = _get_dtype_id(dtype)
# get workspace size and then fire up
+ with nogil:
ws_size = cub_device_spmv_get_workspace_size(
values_ptr, row_offsets_ptr, col_indices_ptr, x_ptr, y_ptr, n_rows,
n_cols, nnz, s, dtype_id)
ws = memory.alloc(ws_size)
ws_ptr = <void *>ws.ptr
+ with nogil:
cub_device_spmv(ws_ptr, ws_size, values_ptr, row_offsets_ptr,
col_indices_ptr, x_ptr, y_ptr, n_rows, n_cols, nnz, s,
dtype_id)
