pool_cuda_graphs

CudaGraphPool

Memory pool for CUDA graphs.

__init__(size, device='cuda')

:param size: size of the pool in bytes.

can_store(t)

Check if the tensor t can be stored in the pool. :param t: tensor to check :return: True if the tensor can be stored in the pool

copy_to_pool(t)

Copy the tensor t in the pool and return a tensor that is a view of the pool. :param t: tensor to copy in the pool :return: tensor copy (that is a view of the pool)

reset()

Reset the pool offset to 0.

get_aligned_size(t, alignment=8)

Get the aligned size of the tensor t. :param t: tensor to get the aligned size of :param alignment: alignment size :return: aligned size