github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

missed cleanup from cache_id removal (#2376)

Browse Source
This commit is contained in:
wozeparrot
2023-11-21 01:03:43 -05:00
committed by GitHub
parent 179551a55c
commit abbcc7aefa
1 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
extra/dist/world.py vendored
View File

@@ -60,7 +60,7 @@ def _send_rb(x:RawBuffer, target_rank:int):
# jit support
CacheCollector.add(__send_rb, [x, target_rank, y], {})
setattr(_send_rb, "shared_memory_cache", {})
if DEBUG >= 2: print(f"**** rank {getenv('RANK')} sent {x} to rank {target_rank}")
# receive a rawbuffer from the target rank
def _recv_rb(x:RawBuffer, target_rank:int):
@@ -74,7 +74,6 @@ def _recv_rb(x:RawBuffer, target_rank:int):
y = RawHIPBuffer(x.size, x.dtype, device=str(y_device), buf=ptr, allocator=None)
x._transfer(y)
if DEBUG >= 2: print(f"**** rank {getenv('RANK')} got {x} from rank {target_rank}")
CacheCollector.add(__recv_rb, [x, target_rank, y], {})
else:
shm_name = dist.OOB.recv(target_rank)
@@ -83,10 +82,10 @@ def _recv_rb(x:RawBuffer, target_rank:int):
# fast path when we can directly copyin
if isinstance(x, RawBufferCopyIn): x._copyin(y.toCPU())
else: x.fromCPU(y.toCPU())
if DEBUG >= 2: print(f"**** rank {getenv('RANK')} got {x} from rank {target_rank}")
# jit support
CacheCollector.add(__recv_rb, [x, target_rank, y], {})
if DEBUG >= 2: print(f"**** rank {getenv('RANK')} got {x} from rank {target_rank}")
# sends a lazybuffer from our rank to the target rank
def _send_lb(x:LazyBuffer, target_rank:int) -> None: