tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 16:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Revert "remove the on_stack stuff and any retries"

Browse Source 
This reverts commit 49a2b328b9.
This commit is contained in:
geohot
2025-10-08 21:11:36 +08:00
parent a4edf78351
commit 1b1752b6f8
1 changed files with 30 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/uop/ops.py
+30 -13
View File
@@ -1019,6 +1019,7 @@ class RewriteContext:
def unified_rewrite(self, root:UOp) -> UOp:
stack: collections.deque[tuple[UOp, int, UOp]] = collections.deque([(root, 0, root)])
on_stack = {root} # all UOps either on the stack or in self.replace, i.e. dont have to be placed again
REWRITE_STACK_LIMIT = getenv("REWRITE_STACK_LIMIT", 250000)
while stack:
if len(stack) > REWRITE_STACK_LIMIT: raise RuntimeError("infinite loop in graph_rewrite (stack too big)")
@@ -1040,23 +1041,39 @@ class RewriteContext:
self.replace[n] = new_n
continue
stack.append((n, 1, new_n))
for x in reversed(new_n.src): stack.append((x, 0, x))
for x in reversed(new_n.src):
if x in on_stack: continue
stack.append((x, 0, x))
on_stack.add(x)
elif stage == 1:
# in stage 1, once all srcs are rewritten, rebuild (if changed) or run top-down rewrite
if (new_src:=tuple([self.replace[x] for x in new_n.src])) == new_n.src:
# if top down, do the rewrite. if no rewrite or bottom up, we are done rewriting this node so we add it to the dict
if self.pm is None or (new_src_n:=self.cached_pm_rewrite(new_n)) is None:
self.replace[n] = new_n
continue
tmp = []
for x in new_n.src:
if (rx:=self.replace.get(x, SENTINEL)) is SENTINEL:
# if some new sources aren't ready, we try this again later. happens with on_stack, maybe should remove?
stack.appendleft((n, 1, new_n))
break
tmp.append(rx)
else:
# if srcs changed from rewrites, construct a new UOp with the new srcs
new_src_n = UOp(new_n.op, new_n.dtype, new_src, new_n.arg, new_n.tag)
# trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n
stack.append((n, 2, new_src_n))
stack.append((new_src_n, 0, new_src_n))
# in stage 1, once all srcs are rewritten, rebuild (if changed) or run top-down rewrite
if (new_src:=tuple(tmp)) == new_n.src:
# if top down, do the rewrite. if no rewrite or bottom up, we are done rewriting this node so we add it to the dict
if self.pm is None or (new_src_n:=self.cached_pm_rewrite(new_n)) is None:
self.replace[n] = new_n
continue
else:
# if srcs changed from rewrites, construct a new UOp with the new srcs
new_src_n = UOp(new_n.op, new_n.dtype, new_src, new_n.arg, new_n.tag)
# trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n
stack.append((n, 2, new_src_n))
stack.append((new_src_n, 0, new_src_n))
else:
# in stage 2, we link the result of new_n to the result of n
self.replace[n] = self.replace[new_n]
if (replaced_new_n:=self.replace.get(new_n, SENTINEL)) is SENTINEL:
# not ready, try the link later
stack.appendleft((n, 2, new_n))
else:
# otherwise we are done
self.replace[n] = replaced_new_n
return self.replace[root]
@track_matches