wehub-resource-sync
475 lines
17 KiB
Python
475 lines
17 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import contextlib
|
|
import ctypes
|
|
from typing import Optional, List, Union
|
|
|
|
import numpy as np
|
|
import tensorrt as trt
|
|
|
|
from cuda.bindings import driver as cuda, runtime as cudart, nvrtc
|
|
|
|
|
|
class ArrayWithOwner(np.ndarray):
|
|
"""Numpy array that holds a reference to its owner object"""
|
|
def __new__(cls, input_array, owner):
|
|
obj = np.asarray(input_array).view(cls)
|
|
obj._owner = owner
|
|
return obj
|
|
|
|
def __array_finalize__(self, obj):
|
|
if obj is None:
|
|
return
|
|
self._owner = getattr(obj, '_owner', None)
|
|
|
|
|
|
|
|
def cuda_call(call):
|
|
"""Helper function to make CUDA calls and check for errors"""
|
|
def _cudaGetErrorEnum(error):
|
|
if isinstance(error, cuda.CUresult):
|
|
err, name = cuda.cuGetErrorName(error)
|
|
return name if err == cuda.CUresult.CUDA_SUCCESS else "<unknown>"
|
|
elif isinstance(error, cudart.cudaError_t):
|
|
return cudart.cudaGetErrorName(error)[1]
|
|
elif isinstance(error, nvrtc.nvrtcResult):
|
|
return nvrtc.nvrtcGetErrorString(error)[1]
|
|
else:
|
|
raise RuntimeError("Unknown error type: {}".format(error))
|
|
|
|
err, res = call[0], call[1:]
|
|
if err.value:
|
|
raise RuntimeError(
|
|
"CUDA error code={}({})".format(
|
|
err.value, _cudaGetErrorEnum(err)
|
|
)
|
|
)
|
|
if len(res) == 1:
|
|
return res[0]
|
|
elif len(res) == 0:
|
|
return None
|
|
else:
|
|
return res
|
|
|
|
|
|
def create_cuda_context(device):
|
|
"""
|
|
Create CUDA context with version-aware API handling.
|
|
|
|
Handles different CUDA API versions based on actual documented signatures:
|
|
- CUDA 11.8-12.9: cuCtxCreate(flags, device) - 2 arguments
|
|
- CUDA 13.0+: cuCtxCreate(ctxCreateParams, flags, device) - 3 arguments
|
|
|
|
Args:
|
|
device: CUDA device handle from cuDeviceGet
|
|
|
|
Returns:
|
|
CUDA context handle
|
|
"""
|
|
# Try different API versions
|
|
try:
|
|
# Try CUDA 13.0+ API first (3 arguments with ctxCreateParams)
|
|
# cuCtxCreate(ctxCreateParams, flags, device)
|
|
return cuda_call(cuda.cuCtxCreate(None, 0, device))
|
|
except TypeError:
|
|
# CUDA 11.8-12.9 API: cuCtxCreate(flags, device)
|
|
return cuda_call(cuda.cuCtxCreate(0, device))
|
|
|
|
|
|
|
|
class HostDeviceMem:
|
|
"""Pair of host and device memory using RAII composition"""
|
|
def __init__(self, size: int, dtype: Optional[np.dtype] = None):
|
|
if dtype is None:
|
|
dtype = np.dtype(np.uint8)
|
|
else:
|
|
dtype = np.dtype(dtype)
|
|
self._size = size
|
|
self._dtype = dtype
|
|
|
|
# Use RAII classes for memory management
|
|
self._host_mem = PinnedHostMem(size, dtype)
|
|
self._device_mem = DeviceMem(size * dtype.itemsize)
|
|
|
|
@property
|
|
def host(self) -> np.ndarray:
|
|
# Return the array directly - ArrayWithOwner ensures proper lifetime management
|
|
return self._host_mem.array
|
|
|
|
@host.setter
|
|
def host(self, data: Union[np.ndarray, bytes]):
|
|
# Delegate to PinnedHostMem for proper data handling
|
|
self._host_mem.array = data
|
|
|
|
@property
|
|
def device_ptr(self) -> int:
|
|
"""Device memory pointer"""
|
|
return self._device_mem.device_ptr
|
|
|
|
@property
|
|
def nbytes(self) -> int:
|
|
return self._host_mem.nbytes
|
|
|
|
|
|
def __str__(self):
|
|
return f"Host:\n{self.host}\nDevice:\n{self.device_ptr}\nSize:\n{self.nbytes}\n"
|
|
|
|
def __repr__(self):
|
|
return self.__str__()
|
|
|
|
|
|
class DeviceMem:
|
|
"""Device-only memory allocation for cases where host memory is not needed"""
|
|
def __init__(self, size: int):
|
|
self._device_ptr = cuda_call(cudart.cudaMalloc(size))
|
|
self._nbytes = size
|
|
|
|
@property
|
|
def device_ptr(self) -> int:
|
|
"""Device memory pointer"""
|
|
return self._device_ptr
|
|
|
|
@property
|
|
def nbytes(self) -> int:
|
|
return self._nbytes
|
|
|
|
def free(self):
|
|
"""Explicitly free device memory"""
|
|
if self._device_ptr is not None:
|
|
try:
|
|
cuda_call(cudart.cudaFree(self._device_ptr))
|
|
self._device_ptr = None
|
|
except Exception:
|
|
# Log but don't raise - cleanup should be best effort
|
|
pass
|
|
|
|
def __str__(self):
|
|
return f"Device:\n{self.device_ptr}\nSize:\n{self.nbytes}\n"
|
|
|
|
def __repr__(self):
|
|
return self.__str__()
|
|
|
|
def __del__(self):
|
|
# Fallback cleanup - not guaranteed to be called
|
|
self.free()
|
|
|
|
|
|
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:
|
|
# Wrap the host buffer as uint8 first, then view as the real dtype.
|
|
# np.ctypeslib.as_ctypes_type has no mapping for dtypes without a ctypes
|
|
# equivalent (float16/bfloat16, FP8, INT4), and would raise NotImplementedError
|
|
# there; the byte-view round-trip avoids that call entirely.
|
|
ptr = ctypes.cast(ctypes.c_void_p(self._host_ptr), ctypes.POINTER(ctypes.c_ubyte))
|
|
host_u8 = np.ctypeslib.as_array(ptr, (self._nbytes,))
|
|
host_array = host_u8.view(self._dtype).reshape(self._host_size)
|
|
return ArrayWithOwner(host_array, self)
|
|
|
|
@array.setter
|
|
def array(self, data: Union[np.ndarray, bytes]):
|
|
"""Set the array data with proper bounds checking"""
|
|
host_array = self.array # Get the numpy array view
|
|
if isinstance(data, np.ndarray):
|
|
if data.size > self._host_size:
|
|
raise ValueError(
|
|
f"Tried to fit an array of size {data.size} into host memory of size {self._host_size}"
|
|
)
|
|
np.copyto(host_array[:data.size], data.flat, casting='safe')
|
|
else:
|
|
assert self._dtype == np.uint8
|
|
host_array[:self.nbytes] = np.frombuffer(data, dtype=np.uint8)
|
|
|
|
|
|
@property
|
|
def nbytes(self) -> int:
|
|
return self._nbytes
|
|
|
|
def free(self):
|
|
"""Explicitly free pinned host memory"""
|
|
if self._host_ptr is not None:
|
|
try:
|
|
cuda_call(cudart.cudaFreeHost(self._host_ptr))
|
|
self._host_ptr = None
|
|
except Exception:
|
|
# Log but don't raise - cleanup should be best effort
|
|
pass
|
|
|
|
def __str__(self):
|
|
return f"PinnedHost:\n{self.array}\nSize:\n{self.nbytes}\n"
|
|
|
|
def __repr__(self):
|
|
return self.__str__()
|
|
|
|
def __del__(self):
|
|
# Fallback cleanup - not guaranteed to be called
|
|
self.free()
|
|
|
|
class CudaStreamContext:
|
|
"""CUDA stream lifecycle management with context manager support"""
|
|
def __init__(self):
|
|
"""Initialize CUDA stream"""
|
|
self._stream = cuda_call(cudart.cudaStreamCreate())
|
|
|
|
def __enter__(self):
|
|
"""Create CUDA stream when entering context (if not already created)"""
|
|
if self._stream is None:
|
|
self._stream = cuda_call(cudart.cudaStreamCreate())
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Destroy CUDA stream when exiting context"""
|
|
if self._stream is not None:
|
|
try:
|
|
cuda_call(cudart.cudaStreamDestroy(self._stream))
|
|
except Exception:
|
|
# Silently handle cleanup failures
|
|
pass
|
|
self._stream = None
|
|
|
|
@property
|
|
def stream(self) -> cudart.cudaStream_t:
|
|
if self._stream is None:
|
|
raise RuntimeError("Stream not created. Use 'with' statement.")
|
|
return self._stream
|
|
|
|
def synchronize(self):
|
|
"""Synchronize the stream"""
|
|
if self._stream is None:
|
|
raise RuntimeError("Stream not created. Use 'with' statement.")
|
|
cuda_call(cudart.cudaStreamSynchronize(self._stream))
|
|
|
|
def free(self):
|
|
"""Explicitly free the CUDA stream"""
|
|
if self._stream is not None:
|
|
try:
|
|
cuda_call(cudart.cudaStreamDestroy(self._stream))
|
|
self._stream = None
|
|
except Exception:
|
|
# Log but don't raise - cleanup should be best effort
|
|
pass
|
|
|
|
def __del__(self):
|
|
"""Cleanup stream on destruction"""
|
|
if hasattr(self, '_stream') and self._stream is not None:
|
|
self.free()
|
|
|
|
def __str__(self):
|
|
return f"CudaStreamContext: {self._stream}"
|
|
|
|
def __repr__(self):
|
|
return self.__str__()
|
|
|
|
|
|
# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
|
|
# Shape precedence: context.infer_shapes() (when context is given and yields fully
|
|
# concrete shapes) > profile_idx max shape > engine.get_tensor_shape(). Using the
|
|
# context branch avoids allocating profile-max-sized output buffers when the actual
|
|
# input shapes are smaller.
|
|
#
|
|
# engine.get_tensor_profile_shape() is an INPUT-only API; for outputs it returns a
|
|
# Dims sentinel with nbDims=-1, which would crash tuple()/__len__() downstream.
|
|
# When the caller passes profile_idx without a configured context, this function
|
|
# builds a temporary IExecutionContext, pins each input to its profile MAX, runs
|
|
# infer_shapes() to derive output MAX shapes, allocates buffers, then discards the
|
|
# temporary context.
|
|
def allocate_buffers(
|
|
engine: trt.ICudaEngine,
|
|
profile_idx: Optional[int] = None,
|
|
context: Optional[trt.IExecutionContext] = None,
|
|
):
|
|
inputs = []
|
|
outputs = []
|
|
bindings = []
|
|
tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)]
|
|
|
|
def is_input(n):
|
|
return engine.get_tensor_mode(n) == trt.TensorIOMode.INPUT
|
|
|
|
# Decide a shape source. Prefer caller's context (if it yields concrete shapes);
|
|
# otherwise build a private one pinned to profile MAX inputs.
|
|
shape_ctx = None
|
|
owns_ctx = False
|
|
owns_stream = None
|
|
|
|
if context is not None:
|
|
try:
|
|
if not context.infer_shapes():
|
|
shape_ctx = context
|
|
except Exception:
|
|
# Any failure to infer shapes (RuntimeError, AttributeError, TRT-specific
|
|
# exceptions) just means we fall back to profile/static shapes below.
|
|
shape_ctx = None
|
|
|
|
try:
|
|
if shape_ctx is None and profile_idx is not None:
|
|
shape_ctx = engine.create_execution_context()
|
|
owns_ctx = True
|
|
if engine.num_optimization_profiles > 1:
|
|
owns_stream = cuda_call(cudart.cudaStreamCreate())
|
|
shape_ctx.set_optimization_profile_async(profile_idx, owns_stream)
|
|
cuda_call(cudart.cudaStreamSynchronize(owns_stream))
|
|
for name in tensor_names:
|
|
if is_input(name):
|
|
max_shape = engine.get_tensor_profile_shape(name, profile_idx)[-1]
|
|
shape_ctx.set_input_shape(name, max_shape)
|
|
if shape_ctx.infer_shapes():
|
|
raise RuntimeError(
|
|
"infer_shapes failed; cannot size output buffers from profile_idx."
|
|
)
|
|
|
|
for binding in tensor_names:
|
|
if shape_ctx is not None:
|
|
shape = tuple(shape_ctx.get_tensor_shape(binding))
|
|
else:
|
|
shape = tuple(engine.get_tensor_shape(binding))
|
|
if any(s < 0 for s in shape):
|
|
raise ValueError(
|
|
f"Binding {binding} has dynamic shape, but no profile was specified."
|
|
)
|
|
size = trt.volume(shape)
|
|
trt_type = engine.get_tensor_dtype(binding)
|
|
|
|
# Allocate host and device buffers
|
|
try:
|
|
dtype = np.dtype(trt.nptype(trt_type))
|
|
binding_memory = HostDeviceMem(size, dtype)
|
|
except TypeError: # no numpy support: create a byte array instead (BF16, FP8, INT4)
|
|
nbytes = int(size * trt_type.itemsize)
|
|
binding_memory = HostDeviceMem(nbytes)
|
|
|
|
# Append the device buffer to device bindings.
|
|
bindings.append(int(binding_memory.device_ptr))
|
|
|
|
# Append to the appropriate list.
|
|
if is_input(binding):
|
|
inputs.append(binding_memory)
|
|
else:
|
|
outputs.append(binding_memory)
|
|
finally:
|
|
if owns_stream is not None:
|
|
with contextlib.suppress(Exception):
|
|
cuda_call(cudart.cudaStreamDestroy(owns_stream))
|
|
if owns_ctx:
|
|
del shape_ctx
|
|
|
|
return inputs, outputs, bindings
|
|
|
|
|
|
# Frees the resources allocated in allocate_buffers
|
|
def free_buffers(inputs: List[HostDeviceMem], outputs: List[HostDeviceMem]):
|
|
"""
|
|
Explicitly free CUDA memory resources.
|
|
|
|
While __del__ methods provide automatic cleanup, they are not guaranteed to be called.
|
|
This function provides explicit resource management for critical applications.
|
|
"""
|
|
for inp in inputs:
|
|
if hasattr(inp, '_device_mem') and hasattr(inp._device_mem, 'free'):
|
|
inp._device_mem.free()
|
|
if hasattr(inp, '_host_mem') and hasattr(inp._host_mem, 'free'):
|
|
inp._host_mem.free()
|
|
|
|
for out in outputs:
|
|
if hasattr(out, '_device_mem') and hasattr(out._device_mem, 'free'):
|
|
out._device_mem.free()
|
|
if hasattr(out, '_host_mem') and hasattr(out._host_mem, 'free'):
|
|
out._host_mem.free()
|
|
|
|
|
|
# Wrapper for cudaMemcpy which infers copy size and does error checking
|
|
def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray):
|
|
cuda_call(cudart.cudaMemcpy(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice))
|
|
|
|
# Wrapper for cudaMemcpy which infers copy size and does error checking
|
|
def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int):
|
|
cuda_call(cudart.cudaMemcpy(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost))
|
|
|
|
|
|
# Additional CUDA wrapper functions for common operations
|
|
|
|
|
|
def cuda_init():
|
|
"""Initialize CUDA driver API with error checking."""
|
|
cuda_call(cuda.cuInit(0))
|
|
|
|
|
|
def cuda_get_device(device_id: int = 0):
|
|
"""Get CUDA device handle with error checking."""
|
|
return cuda_call(cuda.cuDeviceGet(device_id))
|
|
|
|
|
|
# CUDA Runtime API functions (preferred over driver API when available)
|
|
|
|
|
|
def cuda_memcpy_htod(device_ptr: int, host_data: np.ndarray):
|
|
"""Copy data from host to device using CUDA runtime API with error checking.
|
|
|
|
Note: Consider using HostDeviceMem.host setter for integrated memory management.
|
|
"""
|
|
cuda_call(cudart.cudaMemcpy(device_ptr, host_data, host_data.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice))
|
|
|
|
|
|
def _do_inference_base(inputs, outputs, stream, execute_async_func):
|
|
# Transfer input data to the GPU.
|
|
kind = cudart.cudaMemcpyKind.cudaMemcpyHostToDevice
|
|
[cuda_call(cudart.cudaMemcpyAsync(inp.device_ptr, inp.host.ctypes.data, inp.nbytes, kind, stream)) for inp in inputs]
|
|
# Run inference.
|
|
execute_async_func()
|
|
# Transfer predictions back from the GPU.
|
|
kind = cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost
|
|
[cuda_call(cudart.cudaMemcpyAsync(out.host.ctypes.data, out.device_ptr, out.nbytes, kind, stream)) for out in outputs]
|
|
# Synchronize the stream
|
|
cuda_call(cudart.cudaStreamSynchronize(stream))
|
|
# Return only the host outputs.
|
|
return [out.host.copy() for out in outputs]
|
|
|
|
|
|
# This function is generalized for multiple inputs/outputs.
|
|
# inputs and outputs are expected to be lists of HostDeviceMem objects.
|
|
def do_inference(context, engine, bindings, inputs, outputs, stream):
|
|
"""
|
|
Perform inference using the provided context and stream.
|
|
|
|
Usage with context manager:
|
|
with stream: # Ensures proper stream lifecycle
|
|
outputs = do_inference(context, engine, bindings, inputs, outputs, stream)
|
|
"""
|
|
stream_handle = stream.stream
|
|
def execute_async_func():
|
|
context.execute_async_v3(stream_handle=stream_handle)
|
|
# Setup context tensor address.
|
|
num_io = engine.num_io_tensors
|
|
for i in range(num_io):
|
|
context.set_tensor_address(engine.get_tensor_name(i), bindings[i])
|
|
return _do_inference_base(inputs, outputs, stream_handle, execute_async_func)
|