github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 23:18:04 -05:00
Code Issues Packages Projects Releases Wiki Activity

skip UOp del when python is shutting down [pr] (#10847)

Browse Source
This commit is contained in:
qazal
2025-06-19 15:31:40 +03:00
committed by GitHub
parent 58252e3c49
commit ac891b78f8
1 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tinygrad/uop/ops.py
View File

@@ -74,10 +74,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
tag:Any = None
children:set[weakref.ref[UOp]] = field(default_factory=set)
def __del__(self):
if self.op is Ops.BUFFER and (buffer:=buffers.get(self)) is not None: buffer.ref(-1)
if (ref:=UOpMetaClass.ucache.get(k:=(self.op, self.dtype, self.src, self.arg, self.tag))) is not None:
for s in self.src: s.children.discard(ref)
del UOpMetaClass.ucache[k]
if Ops is not None and self.op is Ops.BUFFER and (buffer:=buffers.get(self)) is not None: buffer.ref(-1)
try:
if (ref:=UOpMetaClass.ucache.get(k:=(self.op, self.dtype, self.src, self.arg, self.tag))) is not None:
for s in self.src: s.children.discard(ref)
del UOpMetaClass.ucache[k]
except AttributeError: pass
def __reduce__(self):
args = [self.op, self.dtype, self.src, self.arg, self.tag, self.metadata]
if self.op is Ops.BUFFER and self.realized is not None and PICKLE_BUFFERS: args.append(self.realized)