deepwalk without recursion [pr] (#10002)

* deepwalk without recursion [pr]

* comment and remove that test

tinygrad/gradient.py

@@ -1,5 +1,5 @@
-from typing import cast, Iterator
-import math, functools, dataclasses
+from typing import cast
+import math, dataclasses
 from tinygrad.dtype import dtypes, sum_acc_dtype
 from tinygrad.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
 from tinygrad.helpers import argsort
@@ -45,20 +45,12 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.BITCAST), lambda ctx: (None,)),
 ])
 
-# copied from tensor.py, get relevant toposort of gradients
 def _deepwalk(root:UOp, targets:set[UOp]) -> list[UOp]:
-  @functools.cache
-  def is_in_target_path(x:UOp) -> bool: return any(u in targets or is_in_target_path(u) for u in x.src) # noqa: F821
-  def _walk(node:UOp, visited:set[UOp]) -> Iterator[UOp]:
-    visited.add(node)
-    if node.op in {Ops.DETACH, Ops.ASSIGN}: return
-    if is_in_target_path(node): # noqa: F821
-      for i in node.src:
-        if i not in visited: yield from _walk(i, visited) # noqa: F821
-      yield node
-  ret = list(_walk(root, set()))
-  del is_in_target_path, _walk
-  return ret
+  # compute the target path (top down)
+  in_target_path: dict[UOp, bool] = {}
+  for u in root.toposort: in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
+  # don't flow through DETACH/ASSIGN or anything not in target path
+  return list(root.toposort_gated(lambda node: node.op not in {Ops.DETACH, Ops.ASSIGN} and in_target_path[node]))
 
 def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp]:
   grads = {root: root_grad}

tinygrad/ops.py

@@ -279,6 +279,20 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       else: ret[node] = None # second time i'm seeing this node, add it to returned toposort
     return ret
 
+  # TODO: refactor with above
+  def toposort_gated(self, gate:Callable|None) -> dict[UOp, None]:
+    ret: dict[UOp, None] = {}
+    stack: list[tuple[UOp, bool]] = [(self, False)] # each stack entry is (node, visited_flag)
+    while stack:
+      node, visited = stack.pop()
+      if node in ret: continue
+      if not visited:
+        if gate is None or gate(node):
+          stack.append((node, True))  # push node back on stack to process after its parents
+          for parent in reversed(node.src): stack.append((parent, False)) # push parents on the stack
+      else: ret[node] = None # second time i'm seeing this node, add it to returned toposort
+    return ret
+
   # returns map of UOps to their children in the graph rooted by self
   def get_children_map(self) -> dict[UOp, dict[UOp, None]]:
     ret: dict[UOp, dict[UOp, None]] = {}