support stream on cupy kernel functions
@@ -8,6 +8,7 @@ from libcpp cimport vector
from cupy.cuda cimport driver
from cupy.core cimport core
+from cupy.cuda import stream as stream_module
cdef extern from "cupy_stdint.h" nogil:
@@ -104,8 +105,11 @@ cdef inline CPointer _pointer(x):
raise TypeError('Unsupported type %s. (size=%d)', type(x), itemsize)
-cdef inline size_t _get_stream(strm) except *:
- return 0 if strm is None else strm.ptr
+cdef inline size_t _get_stream(stream) except *:
+ if stream is None:
+ return stream_module.get_current_stream().ptr
+ else:
+ return stream.ptr
cdef void _launch(size_t func, Py_ssize_t grid0, int grid1, int grid2,
