github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-25 06:48:22 -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

2
tinygrad/device.py
View File

@@ -130,7 +130,7 @@ class BufferXfer(BufferCopy):
def copy(self, dest, src):
if hasattr(dest.allocator.device, "track_cross_buffer") and hasattr(src.allocator, "track_cross_device"):
dest.allocator.device.track_cross_buffer.append(src)
src.allocator.track_cross_device.append(dest.allocator.device)
src.allocator.track_cross_device.add(dest.allocator.device)
dest.allocator.transfer(dest._buf, src._buf, dest.nbytes)
# TODO: size, dest, src are the same type. can we enforce this?

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):