From 006637cf73fc0393e58546596b7bd5ee782f01cd Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sun, 12 Jul 2026 22:51:19 +0000 Subject: [PATCH] remove unmeasured CIFAR micro-optimizations --- test/unit/test_helpers.py | 18 ++----------- tinygrad/codegen/__init__.py | 42 +++++------------------------ tinygrad/codegen/late/linearizer.py | 6 ++--- tinygrad/codegen/opt/postrange.py | 13 +++------ tinygrad/engine/realize.py | 6 ++--- tinygrad/helpers.py | 1 - tinygrad/mixin/movement.py | 9 ++----- tinygrad/mixin/op.py | 30 ++++++++++----------- tinygrad/renderer/__init__.py | 7 +++-- tinygrad/schedule/indexing.py | 2 +- tinygrad/uop/ops.py | 15 +++++------ tinygrad/uop/spec.py | 10 +------ tinygrad/uop/symbolic.py | 2 +- 13 files changed, 45 insertions(+), 116 deletions(-) diff --git a/test/unit/test_helpers.py b/test/unit/test_helpers.py index 4fee1ce719..994497fd2b 100644 --- a/test/unit/test_helpers.py +++ b/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() diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 07bd9e418a..eefb8148c6 100644 --- a/tinygrad/codegen/__init__.py +++ b/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]) diff --git a/tinygrad/codegen/late/linearizer.py b/tinygrad/codegen/late/linearizer.py index 8a99ef2c38..e9366cb4d5 100644 --- a/tinygrad/codegen/late/linearizer.py +++ b/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] diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 1e6c7255fb..0dfec14aec 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/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] = {} diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 62e4e28f19..7d9c93a99a 100644 --- a/tinygrad/engine/realize.py +++ b/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"): diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 4838ecd70b..92b336279a 100644 --- a/tinygrad/helpers.py +++ b/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() diff --git a/tinygrad/mixin/movement.py b/tinygrad/mixin/movement.py index 37270991d4..dda49540a5 100644 --- a/tinygrad/mixin/movement.py +++ b/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: diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py index 909cf83c13..e895e11bac 100644 --- a/tinygrad/mixin/op.py +++ b/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: """ diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index 900d540708..bb0d484ac2 100644 --- a/tinygrad/renderer/__init__.py +++ b/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") diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 9fc1fe06e9..ca61dfe942 100644 --- a/tinygrad/schedule/indexing.py +++ b/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 diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index c076353c40..b879ebcc35 100644 --- a/tinygrad/uop/ops.py +++ b/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 diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index ec04dbda2b..dadb039c7b 100644 --- a/tinygrad/uop/spec.py +++ b/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 ***** diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 32cf98443f..e1c89b9ebf 100644 --- a/tinygrad/uop/symbolic.py +++ b/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