Wrapper for CUDA memory buffer that implements __cuda_array_interface__.
| 21 | |
| 22 | # docs_tag: begin_cuda_buffer |
| 23 | class CudaBuffer: |
| 24 | """Wrapper for CUDA memory buffer that implements __cuda_array_interface__.""" |
| 25 | |
| 26 | def __init__(self, shape, dtype, ptr=None): |
| 27 | """Initialize CUDA buffer. |
| 28 | |
| 29 | Args: |
| 30 | shape: tuple of dimensions |
| 31 | dtype: numpy dtype |
| 32 | ptr: CUDA device pointer (if None, allocates new memory) |
| 33 | """ |
| 34 | self.shape = shape |
| 35 | self.dtype = np.dtype(dtype) |
| 36 | self.size = int(np.prod(shape)) * self.dtype.itemsize |
| 37 | |
| 38 | if ptr is None: |
| 39 | err, self.ptr = cudart.cudaMalloc(self.size) |
| 40 | if err != cudart.cudaError_t.cudaSuccess: |
| 41 | raise RuntimeError(f"cudaMalloc failed: {err}") |
| 42 | self.owns_memory = True |
| 43 | else: |
| 44 | self.ptr = ptr |
| 45 | self.owns_memory = False |
| 46 | |
| 47 | @staticmethod |
| 48 | def from_cuda(cuda_obj) -> "CudaBuffer": |
| 49 | """Create from __cuda_array_interface__ or object that implements it.""" |
| 50 | if hasattr(cuda_obj, "__cuda_array_interface__"): |
| 51 | cuda_array_interface = cuda_obj.__cuda_array_interface__ |
| 52 | else: |
| 53 | cuda_array_interface = cuda_obj |
| 54 | return CudaBuffer( |
| 55 | shape=cuda_array_interface["shape"], |
| 56 | dtype=np.dtype(cuda_array_interface["typestr"]), |
| 57 | ptr=cuda_array_interface["data"][0], |
| 58 | ) |
| 59 | |
| 60 | @property |
| 61 | def __cuda_array_interface__(self): |
| 62 | """CUDA Array Interface for zero-copy interop.""" |
| 63 | return { |
| 64 | "version": 3, |
| 65 | "shape": self.shape, |
| 66 | "typestr": self.dtype.str, |
| 67 | "data": (int(self.ptr), False), |
| 68 | "strides": None, |
| 69 | } |
| 70 | |
| 71 | def __del__(self): |
| 72 | """Free CUDA memory if we own it.""" |
| 73 | if self.owns_memory and hasattr(self, "ptr"): |
| 74 | cudart.cudaFree(self.ptr) |
| 75 | |
| 76 | |
| 77 | # docs_tag: end_cuda_buffer |
no outgoing calls
no test coverage detected