Pair of host and device memory, where the host memory is wrapped in a numpy array
| 137 | |
| 138 | |
| 139 | class HostDeviceMem: |
| 140 | """Pair of host and device memory, where the host memory is wrapped in a numpy array""" |
| 141 | def __init__(self, size: int, dtype: np.dtype): |
| 142 | nbytes = size * dtype.itemsize |
| 143 | host_mem = cuda_call(cudart.cudaMallocHost(nbytes)) |
| 144 | pointer_type = ctypes.POINTER(np.ctypeslib.as_ctypes_type(dtype)) |
| 145 | |
| 146 | self._host = np.ctypeslib.as_array(ctypes.cast(host_mem, pointer_type), (size,)) |
| 147 | self._device = cuda_call(cudart.cudaMalloc(nbytes)) |
| 148 | self._nbytes = nbytes |
| 149 | |
| 150 | @property |
| 151 | def host(self) -> np.ndarray: |
| 152 | return self._host |
| 153 | |
| 154 | @host.setter |
| 155 | def host(self, arr: np.ndarray): |
| 156 | if arr.size > self.host.size: |
| 157 | raise ValueError( |
| 158 | f"Tried to fit an array of size {arr.size} into host memory of size {self.host.size}" |
| 159 | ) |
| 160 | np.copyto(self.host[:arr.size], arr.flat, casting='safe') |
| 161 | |
| 162 | @property |
| 163 | def device(self) -> int: |
| 164 | return self._device |
| 165 | |
| 166 | @property |
| 167 | def nbytes(self) -> int: |
| 168 | return self._nbytes |
| 169 | |
| 170 | def __str__(self): |
| 171 | return f"Host:\n{self.host}\nDevice:\n{self.device}\nSize:\n{self.nbytes}\n" |
| 172 | |
| 173 | def __repr__(self): |
| 174 | return self.__str__() |
| 175 | |
| 176 | def free(self): |
| 177 | cuda_call(cudart.cudaFree(self.device)) |
| 178 | cuda_call(cudart.cudaFreeHost(self.host.ctypes.data)) |
| 179 | |
| 180 | |
| 181 | # Allocates all buffers required for an engine, i.e. host/device inputs/outputs. |