github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-02-15 17:15:48 -05:00
Code Issues Packages Projects Releases Wiki Activity

full sync to fix HIP memory leak (#3364)

Browse Source
This commit is contained in:
George Hotz
2024-02-10 11:50:27 +01:00
committed by GitHub
parent ddb22a60c8
commit d1fb1e0ba4
2 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
tinygrad/runtime/ops_hip.py
View File

@@ -71,11 +71,14 @@ CHUNK_SIZE, PAGE_SIZE = 256*1024*1024, 0x1000
class HIPAllocator(LRUAllocator):
def __init__(self, device:HIPDevice):
self.device = device
self.track_cross_device: List[HIPDevice] = []
self.track_cross_device: Set[HIPDevice] = set()
super().__init__()
def free_cache(self):
def full_synchronize(self):
self.device.synchronize()
for x in self.track_cross_device: x.synchronize()
self.track_cross_device.clear()
def free_cache(self):
self.full_synchronize()
return super().free_cache()
def _alloc(self, size:int):
hip_set_device(self.device.device)
@@ -120,7 +123,7 @@ class HIPAllocator(LRUAllocator):
ctypes.memmove(host_mem, from_mv(src), len(src))
check(hip.hipMemcpyAsync(dest, host_mem, len(src), hip.hipMemcpyHostToDevice, None))
def copyout(self, dest:memoryview, src:T):
self.device.synchronize()
self.full_synchronize()
hip_set_device(self.device.device)
check(hip.hipMemcpy(from_mv(dest), src, len(dest), hip.hipMemcpyDeviceToHost))
def transfer(self, dest:T, src:T, sz:int):