feat: way cleaner hip wrapper (#1895)

This commit is contained in:
wozeparrot 2023-09-21 19:20:03 -04:00 committed by GitHub
parent cb5d6576cb
commit 009a99a0b1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 505 additions and 618 deletions

File diff suppressed because it is too large Load Diff

View File

@ -30,8 +30,8 @@ HIP = _HIP()
class RawHIPBuffer(RawBufferCopyInOut, RawBufferTransfer):
def __init__(self, size, dtype, device=str(HIP.default_device)): super().__init__(size, dtype, allocator=HIP.allocator, **{'device': int(device)})
def _copyin(self, x:np.ndarray): hip.hipMemcpyAsync_htod(self._buf, x.ctypes.data, self.size * self.dtype.itemsize, 0)
def _copyout(self, x:np.ndarray): hip.hipMemcpy_dtoh(x.ctypes.data, self._buf, self.size * self.dtype.itemsize)
def _copyin(self, x:np.ndarray): hip.hipMemcpyAsync(self._buf, x.ctypes.data, self.size * self.dtype.itemsize, hip.hipMemcpyHostToDevice, 0)
def _copyout(self, x:np.ndarray): hip.hipMemcpy(x.ctypes.data, self._buf, self.size * self.dtype.itemsize, hip.hipMemcpyDeviceToHost)
def _transfer(self, x): hip.hipMemcpyAsync(self._buf, x._buf, self.size * self.dtype.itemsize, hip.hipMemcpyDeviceToDevice, 0)
class HIPProgram: