Add cudaLaunchHostFunc

cupy/cuda/stream.pyx

@@ -214,16 +214,41 @@ class BaseStream(object):
                 object), and returns nothing.
             arg (object): Argument to the callback.
 
-        """
-        if runtime._is_hip_environment and self.ptr == 0:
-            raise RuntimeError('HIP does not allow adding callbacks to the '
-                               'default (null) stream')
+        .. note::
+            Whenever possible, use the :meth:`launch_host_func` method
+            instead of this one, as it may be deprecated and removed from
+            CUDA at some point.
 
+        """
         def f(stream, status, dummy):
             callback(self, status, arg)
 
         runtime.streamAddCallback(self.ptr, f, 0)
 
+    def launch_host_func(self, callback, arg):
+        """Launch a callback on host when all queued work is done.
+
+        Args:
+            callback (function): Callback function. It must take only one
+                argument (user data object), and returns nothing.
+            arg (object): Argument to the callback.
+
+        .. note::
+            Whenever possible, this method is recommended over
+            :meth:`add_callback`, which may be deprecated and removed from
+            CUDA at some point.
+
+        .. seealso:: `cudaLaunchHostFunc()`_
+
+        .. _cudaLaunchHostFunc():
+            https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__EXECUTION.html#group__CUDART__EXECUTION_1g05841eaa5f90f27124241baafb3e856f
+
+        """
+        def f(dummy):
+            callback(arg)
+
+        runtime.launchHostFunc(self.ptr, f, 0)
+
     def record(self, event=None):
         """Records an event on the stream.
 

cupy_backends/cuda/api/runtime.pxd

@@ -800,6 +800,7 @@ cpdef streamDestroy(intptr_t stream)
 cpdef streamSynchronize(intptr_t stream)
 cpdef streamAddCallback(intptr_t stream, callback, intptr_t arg,
                         unsigned int flags=*)
+cpdef launchHostFunc(intptr_t stream, callback, intptr_t arg)
 cpdef streamQuery(intptr_t stream)
 cpdef streamWaitEvent(intptr_t stream, intptr_t event, unsigned int flags=*)
 cpdef intptr_t eventCreate() except? 0

cupy_backends/cuda/api/runtime.pyx

@@ -61,6 +61,9 @@ cdef extern from *:
         driver.Stream stream, Error status, void* userData)
     ctypedef StreamCallbackDef* StreamCallback 'cudaStreamCallback_t'
 
+    ctypedef void HostFnDef(void* userData)
+    ctypedef HostFnDef* HostFn 'cudaHostFn_t'
+
 
 cdef extern from '../cupy_cuda_runtime.h' nogil:
 
@@ -170,6 +173,7 @@ cdef extern from '../cupy_cuda_runtime.h' nogil:
     int cudaStreamSynchronize(driver.Stream stream)
     int cudaStreamAddCallback(driver.Stream stream, StreamCallback callback,
                               void* userData, unsigned int flags)
+    int cudaLaunchHostFunc(driver.Stream stream, HostFn fn, void* userData)
     int cudaStreamQuery(driver.Stream stream)
     int cudaStreamWaitEvent(driver.Stream stream, driver.Event event,
                             unsigned int flags)
@@ -798,8 +802,18 @@ cdef _streamCallbackFunc(driver.Stream hStream, int status,
     cpython.Py_DECREF(obj)
 
 
+cdef _HostFnFunc(void* func_arg) with gil:
+    obj = <object>func_arg
+    func, arg = obj
+    func(arg)
+    cpython.Py_DECREF(obj)
+
+
 cpdef streamAddCallback(intptr_t stream, callback, intptr_t arg,
                         unsigned int flags=0):
+    if _is_hip_environment and stream == 0:
+        raise RuntimeError('HIP does not allow adding callbacks to the '
+                           'default (null) stream')
     func_arg = (callback, arg)
     cpython.Py_INCREF(func_arg)
     with nogil:
@@ -809,6 +823,21 @@ cpdef streamAddCallback(intptr_t stream, callback, intptr_t arg,
     check_status(status)
 
 
+cpdef launchHostFunc(intptr_t stream, callback, intptr_t arg):
+    if _is_hip_environment:
+        raise RuntimeError('This feature is not supported on HIP')
+    if CUDA_VERSION < 10000:
+        raise RuntimeError('This feature is only supported on CUDA 10.0+')
+
+    func_arg = (callback, arg)
+    cpython.Py_INCREF(func_arg)
+    with nogil:
+        status = cudaLaunchHostFunc(
+            <driver.Stream>stream, <HostFn>_HostFnFunc,
+            <void*>func_arg)
+    check_status(status)
+
+
 cpdef streamQuery(intptr_t stream):
     return cudaStreamQuery(<driver.Stream>stream)
 

cupy_backends/cuda/hip/cupy_common.h

@@ -69,6 +69,7 @@ typedef hipDeviceProp_t cudaDeviceProp;
 
 
 typedef hipStreamCallback_t cudaStreamCallback_t;
+typedef void (*cudaHostFn_t)(void* userData);
 typedef hipPointerAttribute_t cudaPointerAttributes;
 
 typedef hipChannelFormatKind cudaChannelFormatKind;

cupy_backends/cuda/hip/cupy_runtime.h

@@ -256,6 +256,10 @@ cudaError_t cudaStreamAddCallback(cudaStream_t stream,
     return hipStreamAddCallback(stream, callback, userData, flags);
 }
 
+cudaError_t cudaLaunchHostFunc(cudaStream_t stream, cudaHostFn_t fn, void* userData) {
+    return hipErrorUnknown;
+}
+
 cudaError_t cudaStreamQuery(cudaStream_t stream) {
     return hipStreamQuery(stream);
 }

cupy_backends/cuda/stub/cupy_cuda_runtime.h

@@ -237,6 +237,10 @@ cudaError_t cudaStreamAddCallback(...) {
     return cudaSuccess;
 }
 
+cudaError_t cudaLaunchHostFunc(...) {
+    return cudaSuccess;
+}
+
 cudaError_t cudaStreamQuery(...) {
     return cudaSuccess;
 }

docs/source/reference/cuda.rst

@@ -172,6 +172,7 @@ to use these functions.
    cupy.cuda.runtime.streamAddCallback
    cupy.cuda.runtime.streamQuery
    cupy.cuda.runtime.streamWaitEvent
+   cupy.cuda.runtime.launchHostFunc
    cupy.cuda.runtime.eventCreate
    cupy.cuda.runtime.eventCreateWithFlags
    cupy.cuda.runtime.eventDestroy

tests/cupy_tests/cuda_tests/test_stream.py

@@ -59,6 +59,25 @@ def test_get_and_add_callback(self):
         stream.synchronize()
         assert out == list(range(N))
 
+    @attr.gpu
+    @unittest.skipIf(cuda.runtime.is_hip, 'HIP does not support this')
+    @unittest.skipIf(cuda.driver.get_build_version() < 10000,
+                     'Only CUDA 10.0+ supports this')
+    def test_launch_host_func(self):
+        N = 100
+        cupy_arrays = [testing.shaped_random((2, 3)) for _ in range(N)]
+
+        stream = cuda.Stream.null
+
+        out = []
+        for i in range(N):
+            numpy_array = cupy_arrays[i].get(stream=stream)
+            stream.launch_host_func(
+                lambda t: out.append(t[0]), (i, numpy_array))
+
+        stream.synchronize()
+        assert out == list(range(N))
+
     @attr.gpu
     def test_with_statement(self):
         stream1 = cuda.Stream()
@@ -93,11 +112,7 @@ def test_get_and_add_callback(self):
         N = 100
         cupy_arrays = [testing.shaped_random((2, 3)) for _ in range(N)]
 
-        if not cuda.runtime.is_hip:
-            stream = cuda.Stream.null
-        else:
-            # adding callbacks to the null stream in HIP would segfault...
-            stream = cuda.Stream()
+        stream = self.stream
 
         out = []
         for i in range(N):
@@ -108,3 +123,22 @@ def test_get_and_add_callback(self):
 
         stream.synchronize()
         assert out == list(range(N))
+
+    @attr.gpu
+    @unittest.skipIf(cuda.runtime.is_hip, 'HIP does not support this')
+    @unittest.skipIf(cuda.driver.get_build_version() < 10000,
+                     'Only CUDA 10.0+ supports this')
+    def test_launch_host_func(self):
+        N = 100
+        cupy_arrays = [testing.shaped_random((2, 3)) for _ in range(N)]
+
+        stream = self.stream
+
+        out = []
+        for i in range(N):
+            numpy_array = cupy_arrays[i].get(stream=stream)
+            stream.launch_host_func(
+                lambda t: out.append(t[0]), (i, numpy_array))
+
+        stream.synchronize()
+        assert out == list(range(N))