remove on_stack from graph_rewrite

* remove restrictions on range ending in indexing

* early simplify

* Revert "early simplify"

This reverts commit 657d9972c2.

* disable const folding tests

test/mockgpu/nv/nvgpu.py

@@ -1,4 +1,4 @@
-import ctypes, ctypes.util, time
+import ctypes, time
 import tinygrad.runtime.autogen.nv_gpu as nv_gpu
 from enum import Enum, auto
 from test.mockgpu.gpu import VirtGPU

tinygrad/dtype.py

@@ -183,7 +183,7 @@ class dtypes:
   uints = (uint8, uint16, uint32, uint64)
   sints = (int8, int16, int32, int64)
   ints = uints + sints
-  all = floats + ints + (bool, index)
+  all = floats + ints + (bool, index) # noqa: A003
 
 if (env_default_float := getenv("DEFAULT_FLOAT", "")):
   dtypes.default_float = getattr(dtypes, env_default_float.lower())

tinygrad/runtime/ops_amd.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 from typing import cast, ClassVar
-import os, ctypes, ctypes.util, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools
+import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools
 assert sys.platform != 'win32'
 from dataclasses import dataclass
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, FileIOInterface

tinygrad/runtime/ops_cuda.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import ctypes, ctypes.util, functools
+import ctypes, functools
 from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t, suppress_finalizing
 from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerPairT
 from tinygrad.renderer.cstyle import CUDARenderer

tinygrad/runtime/support/compiler_cuda.py

@@ -1,4 +1,4 @@
-import subprocess, hashlib, tempfile, ctypes, ctypes.util, re, pathlib
+import subprocess, hashlib, tempfile, ctypes, re, pathlib
 from typing import Callable
 from tinygrad.helpers import to_char_p_p, colored, init_c_var, getenv
 import tinygrad.runtime.autogen.nvrtc as nvrtc

tinygrad/runtime/support/llvm.py

@@ -1,4 +1,4 @@
-import ctypes, ctypes.util, os, sys, subprocess
+import ctypes.util, os, sys, subprocess
 from tinygrad.helpers import DEBUG, OSX, getenv
 
 if sys.platform == 'win32':

tinygrad/runtime/support/webgpu.py

@@ -1,4 +1,4 @@
-import ctypes, ctypes.util, os, subprocess, platform, sysconfig
+import ctypes.util, os, subprocess, platform, sysconfig
 from tinygrad.helpers import OSX
 
 WEBGPU_PATH: str | None

tinygrad/uop/ops.py

@@ -869,11 +869,11 @@ match_stats:dict[UPat, list[int|float]] = dict()
 
 @dataclass(frozen=True)
 class TrackedGraphRewrite:
-  loc:tuple[str, int]                    # location that called graph_rewrite
-  sink:int                               # the sink input to graph_rewrite
-  matches:list[tuple[int, int, tuple]]   # before/after UOp, UPat location
-  name:str|None                          # optional name of the rewrite
-  depth:int                              # depth if it's a subrewrite
+  loc:tuple[str, int]                           # location that called graph_rewrite
+  sink:int                                      # the sink input to graph_rewrite
+  matches:list[tuple[int, int, tuple, float]]   # before/after UOp, UPat location and time
+  name:str|None                                 # optional name of the rewrite
+  depth:int                                     # depth if it's a subrewrite
   bottom_up:bool
 
 tracked_keys:list[TracingKey] = []
@@ -945,14 +945,14 @@ class TrackedPatternMatcher(PatternMatcher):
       try: ret = match(uop, ctx)
       except Exception:
         if TRACK_MATCH_STATS >= 2 and active_rewrites:
-          active_rewrites[-1].matches.append((track_uop(uop), track_uop(UOp(Ops.REWRITE_ERROR, src=uop.src, arg=str(sys.exc_info()[1]))), p.location))
+          active_rewrites[-1].matches.append((track_uop(uop), track_uop(UOp(Ops.REWRITE_ERROR,src=uop.src,arg=str(sys.exc_info()[1]))),p.location,0))
         raise
       if ret is not None and ret is not uop:
         match_stats[p][0] += 1
         match_stats[p][3] += (et:=time.perf_counter()-st)
         if TRACK_MATCH_STATS >= 3: print(f"{et*1e6:7.2f} us -- ", printable(p.location))
         if TRACK_MATCH_STATS >= 2 and isinstance(ret, UOp) and active_rewrites:
-          active_rewrites[-1].matches.append((track_uop(uop), track_uop(ret), p.location))
+          active_rewrites[-1].matches.append((track_uop(uop), track_uop(ret), p.location, et))
         return ret
       match_stats[p][2] += time.perf_counter()-st
     return None
@@ -1009,7 +1009,6 @@ 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)")
@@ -1031,39 +1030,23 @@ class RewriteContext:
             self.replace[n] = new_n
             continue
         stack.append((n, 1, new_n))
-        for x in reversed(new_n.src):
-          if x in on_stack: continue
-          stack.append((x, 0, x))
-          on_stack.add(x)
+        for x in reversed(new_n.src): stack.append((x, 0, x))
       elif stage == 1:
-        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)
+        # 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
         else:
-          # 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))
+          # 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
-        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
+        self.replace[n] = self.replace[new_n]
     return self.replace[root]
 
 @track_matches

tinygrad/viz/serve.py

@@ -97,12 +97,13 @@ def _reconstruct(a:int, i:int):
 def get_details(ctx:TrackedGraphRewrite, i:int=0) -> Generator[GraphRewriteDetails, None, None]:
   yield {"graph":uop_to_json(next_sink:=_reconstruct(ctx.sink, i)), "uop":str(next_sink), "changed_nodes":None, "diff":None, "upat":None}
   replaces: dict[UOp, UOp] = {}
-  for u0_num,u1_num,upat_loc in tqdm(ctx.matches):
+  for u0_num,u1_num,upat_loc,dur in tqdm(ctx.matches):
     replaces[u0:=_reconstruct(u0_num, i)] = u1 = _reconstruct(u1_num, i)
     try: new_sink = next_sink.substitute(replaces)
     except RuntimeError as e: new_sink = UOp(Ops.NOOP, arg=str(e))
+    match_repr = f"# {dur*1e6:.2f} us\n"+printable(upat_loc)
     yield {"graph":(sink_json:=uop_to_json(new_sink)), "uop":str(new_sink), "changed_nodes":[id(x) for x in u1.toposort() if id(x) in sink_json],
-           "diff":list(difflib.unified_diff(str(u0).splitlines(), str(u1).splitlines())), "upat":(upat_loc, printable(upat_loc))}
+           "diff":list(difflib.unified_diff(str(u0).splitlines(),str(u1).splitlines())), "upat":(upat_loc, match_repr)}
     if not ctx.bottom_up: next_sink = new_sink
 
 # encoder helpers