From bf05534c6ebc2cf201797f7648e190cb93e7ece3 Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Fri, 1 Sep 2023 09:46:13 -0400 Subject: [PATCH] hip multidevice (#1728) * feat: hip multidevice support + p2p * feat: default device --- extra/hip_wrapper.py | 13 +++++++++++++ tinygrad/runtime/ops_hip.py | 32 ++++++++++++++++++++++---------- 2 files changed, 35 insertions(+), 10 deletions(-) diff --git a/extra/hip_wrapper.py b/extra/hip_wrapper.py index de442154ac..969351d5e5 100644 --- a/extra/hip_wrapper.py +++ b/extra/hip_wrapper.py @@ -115,6 +115,9 @@ hipMemcpyDefault = 4 _libhip.hipMemcpy.restype = int _libhip.hipMemcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] +def hipMemcpy(dst, src, count, direction): + status = _libhip.hipMemcpy(dst, src, ctypes.c_size_t(count), direction) + hipCheckStatus(status) def hipMemcpy_htod(dst, src, count): status = _libhip.hipMemcpy(dst, src, ctypes.c_size_t(count), hipMemcpyHostToDevice) @@ -178,6 +181,16 @@ def hipGetDevice(): return dev.value +_libhip.hipGetDeviceCount.restype = int +_libhip.hipGetDeviceCount.argtypes = [ctypes.POINTER(ctypes.c_int)] + +def hipGetDeviceCount(): + count = ctypes.c_int() + status = _libhip.hipGetDeviceCount(ctypes.byref(count)) + hipCheckStatus(status) + return count.value + + class hipDeviceArch(ctypes.Structure): _fields_ = [ # *32-bit Atomics* diff --git a/tinygrad/runtime/ops_hip.py b/tinygrad/runtime/ops_hip.py index a663928d44..a5b7f45897 100644 --- a/tinygrad/runtime/ops_hip.py +++ b/tinygrad/runtime/ops_hip.py @@ -1,9 +1,9 @@ import numpy as np import ctypes, functools import extra.hip_wrapper as hip -from tinygrad.helpers import DEBUG +from tinygrad.helpers import DEBUG, getenv from tinygrad.ops import Compiled -from tinygrad.runtime.lib import RawBufferCopyInOut, LRUAllocator +from tinygrad.runtime.lib import RawBufferCopyInOut, LRUAllocator, RawBufferTransfer from tinygrad.codegen.kernel import LinearizerOptions from tinygrad.renderer.cstyle import uops_to_cstyle, CStyleLanguage @@ -15,22 +15,31 @@ if DEBUG >= 5: # The default HIP stream is used for everything. class HIPAllocator(LRUAllocator): - def _do_alloc(self, size, dtype, device, **kwargs): return hip.hipMalloc(size * dtype.itemsize) + def _do_alloc(self, size, dtype, device, **kwargs): + hip.hipSetDevice(device) + return hip.hipMalloc(size * dtype.itemsize) def _do_free(self, buf): hip.hipFree(buf) def _cached_bufkey(self, size, dtype, device): return (device, size*dtype.itemsize) # Buffers of the same length could be reused, no matter what dtype. -HIPAlloc = HIPAllocator(hip.hipGetDeviceProperties(hip.hipGetDevice()).totalGlobalMem) -class RawHIPBuffer(RawBufferCopyInOut): - def __init__(self, size, dtype): super().__init__(size, dtype, allocator=HIPAlloc) +class _HIP: + def __init__(self): + self.device_count = hip.hipGetDeviceCount() + self.default_device = getenv("HIP_DEFAULT_DEVICE") + self.allocator = HIPAllocator(hip.hipGetDeviceProperties(self.default_device).totalGlobalMem) +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 _transfer(self, x): hip.hipMemcpyAsync(self._buf, x._buf, self.size * self.dtype.itemsize, hip.hipMemcpyDeviceToDevice, 0) class HIPProgram: def __init__(self, name:str, prg:str, binary=False): try: if not binary: prog = hip.hiprtcCreateProgram(prg, name, [], []) - device_properties = hip.hipGetDeviceProperties(hip.hipGetDevice()) + device_properties = hip.hipGetDeviceProperties(HIP.default_device) hip.hiprtcCompileProgram(prog, [f'--offload-arch={device_properties.gcnArchName}']) prg = hip.hiprtcGetCode(prog) except Exception as e: @@ -40,17 +49,20 @@ class HIPProgram: asm = early_exec((["/opt/rocm/llvm/bin/llvm-objdump", '-d', '-'], prg)) print('\n'.join([x for x in asm.decode('utf-8').split("\n") if 's_code_end' not in x])) - module = hip.hipModuleLoadData(prg) - self.prg = hip.hipModuleGetFunction(module, name) + self.prgs = [] + for i in range(HIP.device_count): + hip.hipSetDevice(i) + self.prgs.append(hip.hipModuleGetFunction(hip.hipModuleLoadData(prg), name)) def __call__(self, global_size, local_size, *args, wait=False): + hip.hipSetDevice(args[0]._device) if wait: start, end = hip.hipEventCreate(), hip.hipEventCreate() hip.hipEventRecord(start) class PackageStruct(ctypes.Structure): _fields_ = [(f'field{idx}', ctypes.c_void_p) for idx in range(len(args))] struct = PackageStruct(*[data._buf for data in args]) - hip.hipModuleLaunchKernel(self.prg, global_size[0], global_size[1], global_size[2], local_size[0], local_size[1], local_size[2], 0, 0, struct) + hip.hipModuleLaunchKernel(self.prgs[args[0]._device], global_size[0], global_size[1], global_size[2], local_size[0], local_size[1], local_size[2], 0, 0, struct) if wait: hip.hipEventRecord(end) hip.hipEventSynchronize(end)