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

do not block gc in UOp.toposort (#9623)

Browse Source
This commit is contained in:
qazal
2025-03-30 11:54:17 +08:00
committed by GitHub
parent 54e1e59b44
commit f1a35bbb54
1 changed files with 9 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tinygrad/ops.py
View File

@@ -241,6 +241,15 @@ class UOpMetaClass(type):
buffers:weakref.WeakKeyDictionary[UOp, Buffer] = weakref.WeakKeyDictionary() # this maps BUFFER uops to their device Buffers
all_metadata:weakref.WeakKeyDictionary[UOp, Metadata] = weakref.WeakKeyDictionary()
def _toposort(u:UOp, cache:set[UOp]):
if u in cache: return {}
nodes: dict[UOp, None] = {}
# NOTE: this is a lot faster than the comprehension in parents
for parent in u.src: nodes.update(_toposort(parent, cache))
nodes[u] = None
cache.add(u)
return nodes
# NOTE: this should be frozen, but frozen is slower
@dataclass(eq=False, slots=True)
class UOp(MathTrait, metaclass=UOpMetaClass):
@@ -271,14 +280,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
@property
def toposort(self) -> dict[UOp, None]:
def _toposort(u:UOp, cache:set[UOp]):
if u in cache: return {}
nodes: dict[UOp, None] = {}
# NOTE: this is a lot faster than the comprehension in parents
for parent in u.src: nodes.update(_toposort(parent, cache))
nodes[u] = None
cache.add(u)
return nodes
return _toposort(self, cache=set())
# returns map of UOps to their children in the graph rooted by self