Example of using cudaStreamAddCallback

[wensimin]

I try to register a callback on a stream, and the following code reports an error

def stream_over(*argv):
    print(f"over !")

 cudart.cudaStreamAddCallback(stream, stream_over, None, 0)

nvidia-dev-1  |   cudart.cudaStreamAddCallback(stream._stream, stream_over, None, 0)
nvidia-dev-1  |   File "cuda/bindings/runtime.pyx", line 15344, in cuda.bindings.runtime.cudaStreamAddCallback
nvidia-dev-1  |   File "cuda/bindings/runtime.pyx", line 4054, in cuda.bindings.runtime.cudaStreamCallback_t.__cinit__
nvidia-dev-1  | TypeError: an integer is required

In the documentation I found this class cudaStreamCallback_t，but I don't know how to use it correctly
Any help？

[leofang]

Hi @wensimin this is how cudaStreamAddCallback (deprecated) and cudaLaunchHostFunc are tested:

cuda-python/cuda_bindings/tests/test_cudart.py

Lines 1293 to 1342 in eeda1ab

	def cudart_func_stream_callback(use_host_api):
	class testStruct(ctypes.Structure):
	_fields_ = [
	("a", ctypes.c_int),
	("b", ctypes.c_int),
	("c", ctypes.c_int),
	]
	def task_callback_host(userData):
	data = testStruct.from_address(userData)
	assert data.a == 1
	assert data.b == 2
	assert data.c == 3
	return 0
	def task_callback_stream(stream, status, userData):
	data = testStruct.from_address(userData)
	assert data.a == 1
	assert data.b == 2
	assert data.c == 3
	return 0
	if use_host_api:
	callback_type = ctypes.PYFUNCTYPE(ctypes.c_int, ctypes.c_void_p)
	target_task = task_callback_host
	else:
	callback_type = ctypes.PYFUNCTYPE(ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p)
	target_task = task_callback_stream
	# Construct ctype data
	c_callback = callback_type(target_task)
	c_data = testStruct(1, 2, 3)
	# ctypes is managing the pointer value for us
	if use_host_api:
	callback = cudart.cudaHostFn_t(_ptr=ctypes.addressof(c_callback))
	else:
	callback = cudart.cudaStreamCallback_t(_ptr=ctypes.addressof(c_callback))
	# Run
	err, stream = cudart.cudaStreamCreate()
	assertSuccess(err)
	if use_host_api:
	(err,) = cudart.cudaLaunchHostFunc(stream, callback, ctypes.addressof(c_data))
	assertSuccess(err)
	else:
	(err,) = cudart.cudaStreamAddCallback(stream, callback, ctypes.addressof(c_data), 0)
	assertSuccess(err)
	(err,) = cudart.cudaDeviceSynchronize()
	assertSuccess(err)