From fab262fb98c196244e2f41dfae2bdfe2da7afc80 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 10 Jul 2026 02:00:42 +0000 Subject: [PATCH] fix codegen correctness and lifetime regressions --- tinygrad/codegen/__init__.py | 6 +++++- tinygrad/codegen/late/coalese.py | 4 ++-- tinygrad/engine/jit.py | 3 +-- tinygrad/uop/ops.py | 17 ++++++++++------- 4 files changed, 18 insertions(+), 12 deletions(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index b8200fa50b..721d2993ba 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -455,6 +455,11 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp: full_sink = graph_rewrite(full_sink, renderer.isel_matcher, ctx=IselContext(full_sink), name="instruction selection", bottom_up=True) prg = UOp(Ops.PROGRAM, src=(full_sink,)) else: raise RuntimeError(f"can't call to_program on {ast.op}") + if VIZ: + if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0])) + prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render") + graph_rewrite(prg, PatternMatcher([]), name="View Program") + return prg # 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]) @@ -463,7 +468,6 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp: 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]) if len(prg.src) == 3 and (compiled:=do_compile(renderer, prg, prg.src[2])) is not None: prg = compiled - if VIZ: graph_rewrite(prg, PatternMatcher([]), name="View Program") return prg to_program_cache: dict[tuple, UOp] = {} diff --git a/tinygrad/codegen/late/coalese.py b/tinygrad/codegen/late/coalese.py index eeac316af4..0a70daa435 100644 --- a/tinygrad/codegen/late/coalese.py +++ b/tinygrad/codegen/late/coalese.py @@ -42,7 +42,7 @@ def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|N if not is_image_shape(buf._shape): return None if idx_x.dtype != idx_y.dtype: idx_x, idx_y = idx_x.cast(dtypes.int), idx_y.cast(dtypes.int) start_idx = idx_x.stack(idx_y) - idx = uop_given_valid(valid, start_idx) + idx = uop_given_valid(valid, start_idx, try_simplex=True) drop_stmt = _drop_valid_stmts(valid, idx, buf._shape[0], buf._shape[1]) if not drop_stmt and idx is start_idx: return None @@ -74,7 +74,7 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None: # search for dims that drop the most valid statements best_drop, cands = -1, [] for ch, cw in [shapes[buf.arg.slot]] if buf.arg.slot in shapes else image_valid_dims(buf.dtype, buf.max_numel(), ren.target.arch): - cidx = uop_given_valid(valid, ((x//4)%cw).stack(x//(4*cw))) + cidx = uop_given_valid(valid, ((x//4)%cw).stack(x//(4*cw)), try_simplex=True) dropped = len(_drop_valid_stmts(valid, cidx, ch, cw)) if dropped > best_drop: best_drop, cands = dropped, [(ch, cw, cidx)] elif dropped == best_drop: cands.append((ch, cw, cidx)) diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index a485364114..a733524b46 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -1,7 +1,7 @@ from typing import TypeVar, Generic, Callable, Any import functools, collections from tinygrad.tensor import Tensor, all_tensors -from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ, disable_gc +from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ from tinygrad.device import Buffer, Compiled, Device, MultiBuffer from tinygrad.dtype import DType from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite @@ -269,7 +269,6 @@ class TinyJit(Generic[ReturnType]): def __get__(self, obj, objtype): return functools.partial(self.__call__, obj) # add support for instance methods - @disable_gc() def __call__(self, *args, **kwargs) -> ReturnType: input_buf_uops, var_vals, names, expected_input_info = _prepare_jit_inputs(args, kwargs) if not JIT or self.cnt == 0: diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index e22a974898..18687803f0 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -221,7 +221,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): arg:Any = None tag:Any = None def __del__(self): - if sys.is_finalizing(): return + if getattr(sys, "is_finalizing", lambda: True)(): return if Ops is not None and self.op is Ops.BUFFER and (buffer:=buffers.get(self)) is not None: buffer.ref(-1) try: del UOpMetaClass.ucache[(self.op, self.dtype, self.src, self.arg, self.tag)] except AttributeError: pass @@ -504,7 +504,15 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def substitute(self, dvars:dict[UOp, UOp], name:str|None=None, extra_pm:PatternMatcher|None=None, walk:bool=False, enter_calls:bool=False): dvars = {k:v for k,v in dvars.items() if k is not v} if len(dvars) == 0: return self - if name is None: return _cached_substitute(self, tuple(dvars.items()), extra_pm, walk, enter_calls) + if name is None: + key = (tuple(dvars.items()), extra_pm, walk, enter_calls) + cache = self.__dict__.setdefault('_substitute_cache', {}) + if (cached:=cache.get(key, SENTINEL)) is not SENTINEL: return cached + with Context(TRACK_MATCH_STATS=0): + ret = graph_rewrite(self, (extra_pm+_substitute) if extra_pm is not None else _substitute, dvars, + bottom_up=True, walk=walk, enter_calls=enter_calls) + cache[key] = ret + return ret with Context(TRACK_MATCH_STATS=(0 if name is None else TRACK_MATCH_STATS.value)): return graph_rewrite(self, (extra_pm+_substitute) if extra_pm is not None else _substitute, dvars, bottom_up=True, walk=walk, enter_calls=enter_calls, name=name) @@ -1739,11 +1747,6 @@ pm_lower_index_dtype = PatternMatcher([ def _index_to_concrete_int(u:UOp) -> UOp: return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0] _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))]) -@functools.lru_cache(maxsize=32768) -def _cached_substitute(uop:UOp, dvars:tuple[tuple[UOp, UOp], ...], extra_pm:PatternMatcher|None, walk:bool, enter_calls:bool) -> UOp: - with Context(TRACK_MATCH_STATS=0): - return graph_rewrite(uop, (extra_pm+_substitute) if extra_pm is not None else _substitute, dict(dvars), - bottom_up=True, walk=walk, enter_calls=enter_calls) _pm_resolve_params = PatternMatcher([(UPat(Ops.PARAM, name="p"), lambda ctx,p: ctx[p.arg.slot])]) remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])