INNER CODE UNIT · Python
PinnedHostMem
onnx/onnx-tensorrt · onnx_tensorrt/tensorrt_engine.py:60
class PinnedHostMem:
"""Pinned host memory allocation for faster GPU transfers"""
def __init__(self, size: int, dtype: Optional[np.dtype] = None):
if dtype is None:
dtype = np.dtype(np.uint8)
else:
dtype = np.dtype(dtype)
nbytes = size * dtype.itemsize
host_mem = cuda_call(cudart.cudaMallocHost(nbytes))
self._host_ptr = host_mem
self._host_size = size
self._nbytes = nbytes
self._dtype = dtype
@property
def array(self) -> np.ndarray:
# Create view with proper memory ownership