spdl.io.cuda_config¶
- cuda_config(device_index: int, **kwargs) CUDAConfig [source]¶
Sprcify the CUDA device and memory management.
- Parameters:
device_index (int) – The device to move the data to.
stream (int) –
Optional: Pointer to a custom CUDA stream. By default, per-thread default stream is used.
The value corresponds to
uintptr_t
of CUDA API.Using PyTorch default CUDA stream
It is possible to provide the same stream as the one used in Python’s main thread. For example, you can fetch the default CUDA stream that PyTorch is using as follow.
>>> stream = torch.cuda.Stream() >>> cuda_stream = stream.cuda_stream
Warning
Using the same stream as a model is running might introduce undesired synchronization.
allocator (tuple[Callable[[int, int, int], int], Callable[[int], None]]) –
Optional: A pair of custom CUDA memory allcoator and deleter functions.
Allocator
The allocator function takes the following arguments, and return the address of the allocated memory.
Size:
int
CUDA device index:
int
CUDA stream address:
int
(uintptr_t
)
Deleter
The deleter takes the address of memory allocated by the allocator and free the memory.
An example of such functions are PyTorch’s
caching_allocator_alloc()
andcaching_allocator_delete()
.