Example of using cudaStreamAddCallback #309

Activity

1. leofang commented on Dec 17, 2024

   

   leofang

   on Dec 17, 2024

   Member

   More actions

   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)
2. NVIDIA

   locked and limited conversation to collaborators on Dec 17, 2024
3. leofang

   converted this issue into a discussion #310 on Dec 17, 2024