remove unmeasured CIFAR micro-optimizations

test/unit/test_helpers.py

@@ -1,6 +1,6 @@
-import unittest, gc
+import unittest
 import numpy as np
-from tinygrad.helpers import polyN, is_numpy_ndarray, disable_gc
+from tinygrad.helpers import polyN, is_numpy_ndarray
 from tinygrad.tensor import Tensor
 
 class TestPolyN(unittest.TestCase):
@@ -11,19 +11,5 @@ class TestIsNumpyNdarray(unittest.TestCase):
   def test_tensor_numpy(self):
     self.assertTrue(is_numpy_ndarray(Tensor([1, 2, 3]).numpy()))
 
-class TestDisableGC(unittest.TestCase):
-  def test_recursive_decorator(self):
-    was_enabled = gc.isenabled()
-    @disable_gc()
-    def recurse(depth:int):
-      self.assertFalse(gc.isenabled())
-      if depth: recurse(depth-1)
-      self.assertFalse(gc.isenabled())
-    try:
-      recurse(2)
-      self.assertEqual(gc.isenabled(), was_enabled)
-    finally:
-      (gc.enable if was_enabled else gc.disable)()
-
 if __name__ == '__main__':
   unittest.main()

tinygrad/codegen/__init__.py

@@ -156,21 +156,6 @@ class DevectorizeContext:
   @property
   def rewrite_cache_key(self): return (type(self.ren), self.ren.target)
 
-devectorize_caches: list[dict[tuple[UOp, tuple[UOp, ...]], UOp]] = []
-class scoped_devectorize_cache:
-  def __enter__(self): devectorize_caches.append({})
-  def __exit__(self, *args): devectorize_caches.pop()
-
-apply_opts_caches: list[dict[tuple[UOp, type[Renderer], object, int], UOp]] = []
-class scoped_apply_opts_cache:
-  def __enter__(self): apply_opts_caches.append({})
-  def __exit__(self, *args): apply_opts_caches.pop()
-
-postopt_codegen_caches: list[dict[tuple[UOp, type[Renderer], object], UOp]] = []
-class scoped_postopt_codegen_cache:
-  def __enter__(self): postopt_codegen_caches.append({})
-  def __exit__(self, *args): postopt_codegen_caches.pop()
-
 def index_lane(ctx:DevectorizeContext, x:UOp, idxs:tuple[UOp, ...]) -> UOp:
   key = (x, idxs)
   if (ret:=ctx.lanes.get(key)) is not None: return ret
@@ -341,15 +326,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
     sink = graph_rewrite(sink, pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges")
 
     # do postrange optimization, BEAM or hand_coded_optimizations
-    opt_key = (sink, type(ren), ren.target, ast.arg.beam)
-    if apply_opts_caches and (optimized:=apply_opts_caches[-1].get(opt_key)) is not None: sink = optimized
-    else:
-      sink = apply_opts(sink, ren, beam=ast.arg.beam)
-      if apply_opts_caches: apply_opts_caches[-1][opt_key] = sink
-
-  postopt_key = (sink, type(ren), ren.target)
-  use_postopt_cache = bool(postopt_codegen_caches) and not (VIZ or PROFILE or TRACK_MATCH_STATS)
-  if use_postopt_cache and (cached_sink:=postopt_codegen_caches[-1].get(postopt_key)) is not None: return cached_sink
+    sink = apply_opts(sink, ren, beam=ast.arg.beam)
 
   # ** expander (expand_rewrite) **
   sink = graph_rewrite(sink, sym+pm_move_where_on_load+pm_flatten_range, name="postopt symbolic")
@@ -372,8 +349,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, symbolic_simple+unbroadcast+pm_add_loads, name="*** unbroadcast / add loads")
 
   # devectorize
-  sink = graph_rewrite(sink, symbolic_simple+devectorizer2,
-                       ctx=DevectorizeContext(ren, devectorize_caches[-1] if devectorize_caches else {}), name="devectorize2")
+  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=DevectorizeContext(ren, {}), name="devectorize2")
 
   # some coalesing misses without this
   sink = graph_rewrite(sink, sym, name="early symbolic")
@@ -416,22 +392,16 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")
 
   # this was the linearizer
-  final_topo = sink.backward_slice_with_self
-  if any(x.op is Ops.RANGE for x in final_topo):
-    sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink, itertools.chain(sink.backward_slice, (sink,))),
-                         name="add control flow", bottom_up=True)
-    final_topo = sink.backward_slice_with_self
+  sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)
 
   # put unnumbered variable PARAMs in slots
-  params = [x for x in final_topo if x.op is Ops.PARAM]
-  if any(x.arg.slot == -1 for x in params):
-    sink = graph_rewrite(sink, pm_number_params, ctx=[sum(x.arg.slot != -1 for x in params)], name="number params with -1", walk=True)
+  num_params = len([x for x in sink.toposort() if x.op is Ops.PARAM and x.arg.slot != -1])
+  sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)
 
   if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
   if SPEC: type_verify(sink, spec_program)
 
   # return the rewritten sink
-  if use_postopt_cache: postopt_codegen_caches[-1][postopt_key] = sink
   return sink
 
 # inject IF/ENDIF. only needed if device doesn't support gated stores
@@ -542,7 +512,7 @@ def do_to_program(ast:UOp, renderer:Renderer, compile_binary=True) -> UOp:
   # PROGRAM lowering is a linear root-only pipeline. Driving it through graph_rewrite
   # needlessly walks the full SINK and LINEAR graphs between each stage.
   if len(prg.src) == 1: prg = do_linearize(renderer, prg, prg.src[0])
-  if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0], uops=prg.src[1].src))
+  if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0]))
   if prg.src[0].arg.estimates is None and (estimated:=do_estimates(prg, prg.src[0], prg.src[1])) is not None: prg = estimated
   if len(prg.src) == 2:
     prg = do_assemble(renderer, prg, prg.src[1]) if isinstance(renderer, ISARenderer) else do_render(renderer, prg, prg.src[1])

tinygrad/codegen/late/linearizer.py

@@ -1,5 +1,5 @@
 import heapq
-from typing import Any, Iterable
+from typing import Any
 from collections import defaultdict
 from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
 from tinygrad.dtype import AddrSpace
@@ -51,7 +51,7 @@ def linearize(sink:UOp) -> list[UOp]:
   return newlst
 
 class CFGContext:
-  def __init__(self, sink:UOp, topo:Iterable[UOp]|None=None):
+  def __init__(self, sink:UOp):
     # there are 3 relationships between ranges:
     # nested, meaning endrange y is a dependency of endrange x and range x is a dependency of endrange y
     # dependent, meaning endrange y is a dependency of endrange x and range x is not a dependency of endrange y
@@ -59,7 +59,7 @@ class CFGContext:
     # everything is nested inside the sink
     deps: dict[UOp, dict[UOp, None]] = {}
     nesting: dict[UOp, UOp] = {}
-    for u in sink.toposort() if topo is None else topo:
+    for u in sink.toposort():
       # get the deps from the src
       deps[u] = {}
       for s in u.src: deps[u] |= deps[s]

tinygrad/codegen/opt/postrange.py

@@ -55,14 +55,9 @@ class Scheduler:
       special_uops = sorted([x for x in self.ast.toposort() if x.op is Ops.SPECIAL], key=lambda x: x.arg)
       special_ops = [colored(str(x.vmax+1), "blue" if x.arg[0] == "g" else "cyan") for x in special_uops]
       name = k_type + colored('_', 'BLACK').join(['']+special_ops+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())])
-      function_name = to_function_name(name)
-      cache_key = (function_name, self.ast)
-      if (cached_name:=kernel_name_cache.get(cache_key)) is not None: name = cached_name
-      else:
-        Scheduler.kernel_cnt[function_name] += 1
-        num = f"n{Scheduler.kernel_cnt[function_name]-1}" if Scheduler.kernel_cnt[function_name] > 1 else ""
-        name += colored(num, 'BLACK')
-        kernel_name_cache[cache_key] = name
+      Scheduler.kernel_cnt[(function_name := to_function_name(name))] += 1
+      num = f"n{Scheduler.kernel_cnt[function_name]-1}" if Scheduler.kernel_cnt[function_name] > 1 else ""
+      name += colored(num, 'BLACK')
     self.ast = graph_rewrite(self.ast, pm_flatten_range, name="flatten range")
     return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts), dont_use_locals=self.dont_use_locals), tag=1)
 
@@ -357,5 +352,3 @@ def apply_opts(ast:UOp, ren:Renderer, beam:int=0) -> UOp:
     if not any(u.op is Ops.STAGE for u in ast.backward_slice):
       k = hand_coded_optimizations(k)
   return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None)
-
-kernel_name_cache:dict[tuple[str, UOp], str] = {}

tinygrad/engine/realize.py

@@ -7,8 +7,7 @@ from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU,
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite, ProgramInfo, scoped_rewrite_cache
 from tinygrad.device import Device, Buffer, MultiBuffer, Compiler
 from tinygrad.renderer import Estimates
-from tinygrad.codegen import to_program, scoped_apply_opts_cache, scoped_devectorize_cache, scoped_postopt_codegen_cache
-from tinygrad.uop.spec import scoped_type_verify_cache
+from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import bufs_from_ast
 
 # **************** Helpers ****************
@@ -299,8 +298,7 @@ pm_exec = PatternMatcher([
 def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, jit=False) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
   if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
-  with scoped_rewrite_cache(), scoped_apply_opts_cache(), scoped_postopt_codegen_cache(), scoped_devectorize_cache(), \
-       scoped_type_verify_cache():
+  with scoped_rewrite_cache():
     linear = graph_rewrite(linear, pm_compile, ctx=True, name="render kernels", walk=True)
   linear = batch_compile(linear)
   if getenv("HCQ2"):

tinygrad/helpers.py

@@ -572,7 +572,6 @@ class tqdm(Generic[T]):
 def trange(n:int, **kwargs) -> tqdm[int]: return tqdm(range(n), total=n, **kwargs)
 
 class disable_gc(contextlib.ContextDecorator):
-  def _recreate_cm(self): return type(self)()
   def __enter__(self):
     self._was_enabled = gc.isenabled()
     if self._was_enabled: gc.disable()

tinygrad/mixin/movement.py

@@ -85,8 +85,6 @@ class MovementMixin:
         b = index if resolve(index >= 0, False) else index + size
         return {"size":size, "boundary":(b, b+1), "stride":1, "collapse_dim":True}
       case slice():
-        if index.start is None and index.stop is None and index.step is None:
-          return {"size":size, "boundary":(0,size), "stride":1, "collapse_dim":False}
         if not all(s is None or isinstance(s, sint) for s in (index.start, index.stop, index.step)):
           raise TypeError(f"slice {index=} is not supported")
         if resolve(index.step == 0, False): raise ValueError(f"{index=} cannot have 0 as step")
@@ -109,8 +107,7 @@ class MovementMixin:
   def _apply_view_ops(self, mops:list) -> Self:
     # applies shrink + flip + stride from a list of parsed view indices
     # flip negative strides
-    x = self.shrink(tuple(m["boundary"] for m in mops))
-    if flip_dims := tuple(i for i, m in enumerate(mops) if m["stride"] < 0): x = x.flip(flip_dims)
+    x = self.shrink(tuple(m["boundary"] for m in mops)).flip(tuple(i for i, m in enumerate(mops) if m["stride"] < 0))
     strides = tuple(abs(m["stride"]) for m in mops)
     # apply stride
     if any(st != 1 for st in strides):
@@ -203,9 +200,7 @@ class MovementMixin:
     """
     if self.ndim != len(arg):
       raise ValueError(f"{self.ndim=} != {len(arg)=}")
-    mop_arg = [(x[0], x[1]-x[0]) if x is not None else (0, s) for x, s in zip(arg, self.shape)]
-    ret = self._mop(Ops.SHRINK, arg=mop_arg)
-    if any(isinstance(ns, int) and isinstance(s, int) and ns != s for (_,ns),s in zip(mop_arg, self.shape)): return ret
+    ret = self._mop(Ops.SHRINK, arg=[(x[0], x[1]-x[0]) if x is not None else (0, s) for x, s in zip(arg, self.shape)])
     return self if ret.shape == self.shape else ret
 
   def permute(self, order, *args) -> Self:

tinygrad/mixin/op.py

@@ -99,7 +99,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     # apply view ops then dim injection (None) and collapse (int)
     x = self._apply_view_ops(mops := [p for p in indices_parsed if p["index"] is not None])
     x_dims = [p for p in indices_parsed if not p["collapse_dim"]]
-    if any(p["index"] is None or p["collapse_dim"] for p in indices_parsed): x = x.reshape(tuple(p["size"] for p in x_dims))
+    x = x.reshape(tuple(p["size"] for p in x_dims))
 
     # tensor indexing
     if tops := [(d, p) for d, p in enumerate(x_dims) if is_adv(p['index'])]:
@@ -918,35 +918,35 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     if (orig_len := int(x.shape[dim])) <= 1: return x, x.const_like(0).cast(dtypes.default_int)
     # pad to power of 2
     n_stages = (orig_len-1).bit_length()
-    padded_len = 2**n_stages
-    pads = tuple((0, padded_len - orig_len) if i == dim else None for i in range(x.ndim))
-    x = x._pad_constant(pads, x.dtype.min if descending else x.dtype.max)
-    idx = (x.const_like(1).cast(dtypes.default_int)._cumalu(dim, Ops.ADD) - 1).unflatten(dim, (2,)*n_stages)
-    x = x.unflatten(dim, (2,)*n_stages)
+    pads = tuple((0, 2**n_stages - orig_len) if i == dim else None for i in range(x.ndim))
+    x = x._pad_constant(pads, x.dtype.min if descending else x.dtype.max).unflatten(dim, (2,)*n_stages)
     # https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort1.svg
     for stage in range(1, n_stages+1):
       if stage != n_stages:
         # flip so arrows of green boxes point the same way as blue boxes
         crossover_dim = dim + n_stages - stage - 1
         blue_box, green_box = x.split(1, crossover_dim)
-        idx_blue_box, idx_green_box = idx.split(1, crossover_dim)
         flip_dims = tuple(-i for i in range(1, stage+1+(self.ndim-dim)))
         x = (blue_box.cat(green_box.flip(flip_dims), dim=crossover_dim)).contiguous()
-        idx = (idx_blue_box.cat(idx_green_box.flip(flip_dims), dim=crossover_dim)).contiguous()
       for substage in range(stage-1, -1, -1):
         partner_dim = dim + n_stages - substage - 1
         x_top, x_bottom = x.split(1, partner_dim)
-        idx_top, idx_bottom = idx.split(1, partner_dim)
-        top_first = ((x_top > x_bottom) if descending else (x_top < x_bottom)) | (x_top.eq(x_bottom) & (idx_top < idx_bottom))
-        x = top_first.where(x_top, x_bottom).cat(top_first.where(x_bottom, x_top), dim=partner_dim).contiguous()
-        idx = top_first.where(idx_top, idx_bottom).cat(top_first.where(idx_bottom, idx_top), dim=partner_dim).contiguous()
+        x_larger, x_smaller = x_top.maximum(x_bottom), x_top.minimum(x_bottom)
+        x = (x_larger.cat(x_smaller, dim=partner_dim) if descending else x_smaller.cat(x_larger, dim=partner_dim)).contiguous()
       if stage != n_stages:
         # flip wires back to undo the crossover
         blue_box, flipped_green_box = x.split(1, crossover_dim)
-        idx_blue_box, idx_flipped_green_box = idx.split(1, crossover_dim)
         x = blue_box.cat(flipped_green_box.flip(flip_dims), dim=crossover_dim)
-        idx = idx_blue_box.cat(idx_flipped_green_box.flip(flip_dims), dim=crossover_dim)
-    return x.flatten(dim, dim+n_stages-1).shrink_to(self.shape), idx.flatten(dim, dim+n_stages-1).shrink_to(self.shape)
+    x = x.flatten(dim, dim+n_stages-1).shrink_to(self.shape)
+    # compute indices for sorted values
+    mask = type(self).ones(orig_len, orig_len, dtype=dtypes.bool, buffer=False).tril()
+    mask = mask.reshape((None, None) + (1,)*(self.ndim-dim-1))
+    def compute_counts(t:Self): return (mask & t.unsqueeze(dim).eq(t.unsqueeze(dim+1))).sum(dim+1)
+    count_orig, count_sorted = compute_counts(self), compute_counts(x)
+    cond = self.unsqueeze(dim+1).eq(x.unsqueeze(dim)) & count_orig.unsqueeze(dim+1).eq(count_sorted.unsqueeze(dim))
+    idx = type(self).arange(orig_len).reshape(tuple(orig_len if i == dim else 1 for i in range(x.ndim)))
+    idx = (cond * idx.unsqueeze(dim+1)).sum(dim)
+    return x, idx
 
   def argsort(self, dim:int=-1, descending:bool=False) -> Self:
     """

tinygrad/renderer/__init__.py

@@ -26,8 +26,9 @@ class Estimates:
     mult_stack: list[sint] = []
     excluded: set[UOp] = set()
     if ignore_indexing:
-      indexing_srcs = [s for u in uops if u.op in {Ops.INDEX, Ops.SHRINK} for s in u.src[1:]]
-      if indexing_srcs: excluded.update(UOp.sink(*indexing_srcs).toposort(lambda x: x.op is not Ops.END))
+      for u in uops:
+        if u.op in {Ops.INDEX, Ops.SHRINK}:
+          excluded = excluded.union(set(UOp.sink(*u.src[1:]).toposort(lambda x: x.op is not Ops.END)))
     for u in uops:
       if u.op in {Ops.LOAD, Ops.STORE}:
         buf = u
@@ -75,8 +76,6 @@ class Renderer:
   compiler: Compiler = Compiler()
 
   def __init__(self, target:Target): self.target = target
-  @property
-  def rewrite_cache_key(self): return (type(self), self.target)
   def __reduce__(self): return self.__class__, (self.target,)
   def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer")
   def asm(self, prg:UOp, lin:UOp) -> bytes: raise NotImplementedError("needs an assembler")

tinygrad/schedule/indexing.py

@@ -241,7 +241,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     if x.op in {Ops.MSTACK, Ops.MSELECT}: continue
 
     if x.dtype == dtypes.index: continue  # TODO: why do I need this?
-    ending_ranges[x] = [r for u in consumer_map[x] for r in ending_ranges.get(u, ())]
+    ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
 
     # *** the ranges on the output are
     #  1. new if this op is realized

tinygrad/uop/ops.py

@@ -168,8 +168,7 @@ class UOpMetaClass(type):
   def __call__(cls, op:Ops, dtype:DType|None=None, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
                metadata:tuple[Metadata,...]|None=None, _buffer:Buffer|None=None):
     if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void
-    spec = SPEC.value
-    if spec == 2 and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype:
+    if SPEC == 2 and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype:
       raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}")
     if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret
     UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(*key))
@@ -178,12 +177,12 @@ class UOpMetaClass(type):
     if _buffer is not None:
       assert op is Ops.BUFFER, f"trying to set Buffer {_buffer} for {op}"
       buffers[created] = _buffer
-    if spec > 1:
+    if SPEC > 1:
       from tinygrad.uop.spec import spec_full, test_pyrender
-      if spec > 2:
+      if SPEC > 2:
         # SPEC=3 checks the shape
         _ = created._shape
-        if spec > 3:
+        if SPEC > 3:
           test_pyrender(created)
       with Context(CHECK_OOB=0): fret = cast(bool|None, spec_full.rewrite(created))
       if fret is not True: raise RuntimeError(f"SPEC ISSUE {fret}: {created}")
@@ -1209,14 +1208,14 @@ class ProgramInfo:
     except KeyError as e: raise RuntimeError(f"unbound Variable {e} used by {self.function_name}") from None
 
   @staticmethod
-  def from_sink(sink:UOp, aux:tuple=(), uops:Iterable[UOp]|None=None) -> ProgramInfo:
+  def from_sink(sink:UOp, aux:tuple=()) -> ProgramInfo:
     _vars: list[UOp] = []
     _globals: list[int] = []
     outs: list[int] = []
     ins: list[int] = []
     global_size: list[int] = [1, 1, 1]
     local_size: list[int]|None = [1, 1, 1]
-    for u in sink.toposort() if uops is None else uops:
+    for u in sink.toposort():
       if u.op is Ops.PARAM and u.addrspace == AddrSpace.ALU: _vars.append(u)
       if u.op is Ops.PARAM and u.addrspace != AddrSpace.ALU: _globals.append(u.arg.slot)
       if u.op in (Ops.STORE, Ops.LOAD):
@@ -1429,7 +1428,6 @@ def upat_deferred_compile(p:UPat, fxn:Callable, entry:list) -> Callable:
   return lazy_compile
 
 class PatternMatcher:
-  __slots__ = ("patterns", "pdict")
   def __init__(self, patterns:Sequence[tuple[UPat, Callable|tuple]], compiled=bool(getenv("UPAT_COMPILE", 1))):
     # if this comes from a pickle, we reconstruct the lambda functions here
     self.patterns:list[tuple[UPat, Callable]] = [(p,types.FunctionType(*fxn) if isinstance(fxn, tuple) else fxn) for p,fxn in patterns]
@@ -1624,7 +1622,6 @@ class scoped_rewrite_cache:
   def __exit__(self, *args): rewrite_caches.pop()
 
 class RewriteContext:
-  __slots__ = ("pm", "bpm", "bpm_cache", "ctx", "replace", "enter_calls", "shared_cache")
   def __init__(self, pm, bpm, ctx=None, enter_calls=False, shared_cache=None):
     self.pm: PatternMatcher|None = pm
     self.bpm: PatternMatcher|None = bpm

tinygrad/uop/spec.py

@@ -33,23 +33,15 @@ def validate_index(uidx:UOp, gate:UOp|None=None):
   return validate_index_with_z3(sz, idx, gate)
 
 def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher):
-  lst = [*ast.backward_slice, ast] if isinstance(ast, UOp) else ast
+  lst = list(ast.toposort()) if isinstance(ast, UOp) else ast
   if SPEC > 1: test_pyrender(lst[-1])  # assume this is the sink
 
-  verified = type_verify_caches[-1].setdefault(check_spec, set()) if type_verify_caches else None
   with Context(TRACK_MATCH_STATS=0):
     for i,u in enumerate(lst):
-      if verified is not None and u in verified: continue
       ret: bool|None = check_spec.rewrite(u)
       if ret is not True:
         if DEBUG >= 3: print_uops(lst)
         raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[(x.op, x.dtype, x.arg) for x in u.src]} {u.arg}")
-      if verified is not None: verified.add(u)
-
-type_verify_caches: list[dict[PatternMatcher, set[UOp]]] = []
-class scoped_type_verify_cache:
-  def __enter__(self): type_verify_caches.append({})
-  def __exit__(self, *args): type_verify_caches.pop()
 
 # ***** new specs *****
 

tinygrad/uop/symbolic.py

@@ -359,7 +359,7 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=False) -> UOp:
 
 def _valid_priority(v: UOp, valids:list[UOp]) -> int:
   # we want valid that's in other valids' parents to be first, so it's more likely the other valids get simplified
-  return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.backward_slice_with_self else 0 for other in valids)
+  return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.toposort() else 0 for other in valids)
 
 def simplify_valid(valid:UOp) -> UOp|None:
   if valid.op_in_backward_slice_with_self(Ops.INDEX): return None  # this should only be for indexing, skip if there's a INDEX