-
Notifications
You must be signed in to change notification settings - Fork 85
Description
Some KvikIO Python API functions have CUDA stream parameters.
def raw_read_async(
self,
buf,
raw_stream: int,
size: Optional[int] = None,
file_offset: int = 0,
dev_offset: int = 0,
) -> IOFutureStream:
def stream_register(raw_stream: int, flags: int) -> None:Currently, the stream argument passed to these functions has to be raw handles of integer type (usually the underlying CUDA stream, which is a pointer). To make the API more flexible and robust, we may consider supporting common types of CUDA stream objects, on top of the existing raw CUDA stream support:
- cupy CUDA stream
- numba CUDA stream
- RMM CUDA stream
- cuda-core CUDA stream
- cuda-bindings CUDA stream (driver)
- cuda-bindings CUDA stream (runtime)
It is likely that we do not need to include these packages as KvikIO dependencies. Instead we only need to require certain methods to be available with which to retrieve the underlying CUDA stream and pass it to Cython.
Essentially KvikIO is expected to have a stream wrapper, say kvikio.Stream(obj=external_stream_obj_or_raw_stream) that retrieves the underlying CUDA stream from the abovementioned stream objects.