tensor graph is a single sink

extra/benchmark_llm.py

@@ -1,31 +0,0 @@
-import argparse, time
-from tinygrad.llm.model import Transformer
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  parser.add_argument("--model", required=True, help="path to gguf model")
-  parser.add_argument("--max-context", type=int, default=8192, help="max context length (default: %(default)s)")
-  parser.add_argument("--prompt-tokens", type=int, default=1024, help="number of prompt tokens (default: %(default)s)")
-  parser.add_argument("--decode-tokens", type=int, default=16, help="number of tokens to decode (default: %(default)s)")
-  parser.add_argument("--chunk-size", type=int, default=32, help="chunk size for prefill (default: %(default)s)")
-  args = parser.parse_args()
-
-  st = time.perf_counter()
-  model, _ = Transformer.from_gguf(args.model, args.max_context)
-  print(f"load {time.perf_counter()-st:.3f}s", flush=True)
-
-  st = time.perf_counter()
-  model.warmup()
-  print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
-
-  prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
-  gen = model.generate(prompt, chunk_size=args.chunk_size)
-  st = time.perf_counter()
-  # first token is time-to-first-token; counted as part of prefill
-  output = [next(gen)]
-  pt = time.perf_counter()
-  print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
-
-  for _ in range(args.decode_tokens): output.append(next(gen))
-  et = time.perf_counter()
-  print(f"decode {args.decode_tokens/(et-pt):.3f} tok/s output {output}", flush=True)

test/backend/test_symbolic_ops.py

@@ -6,15 +6,6 @@ from examples.gpt2 import Attention
 import numpy as np
 
 class TestSymbolicOps(unittest.TestCase):
-  def test_negative_slice(self):
-    a = Tensor.rand(3, 10, 4)
-    for i in range(3, 10):
-      vi = Variable("i", 1, 10).bind(i)
-      # negative int bounds against a symbolic dim must resolve against the size, like slice.indices
-      np.testing.assert_allclose(a[:, :vi][:, -3:-1].numpy(), a[:, :i][:, -3:-1].numpy(), atol=1e-6, rtol=1e-6)
-      np.testing.assert_allclose(a[:, :vi][:, -1:].numpy(), a[:, :i][:, -1:].numpy(), atol=1e-6, rtol=1e-6)
-      np.testing.assert_allclose(a[:, :vi][:, -1].numpy(), a[:, :i][:, -1].numpy(), atol=1e-6, rtol=1e-6)
-
   def test_plus1(self):
     def f(a): return (a+1).realize()
     a = Tensor.rand(3, 10)

test/null/test_graph_rewrite.py

@@ -208,7 +208,7 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase):
 
 
 import inspect
-from tinygrad.uop.ops import graph_rewrite, _substitute, rewrite_group
+from tinygrad.uop.ops import graph_rewrite, _substitute, track_rewrites
 from tinygrad.uop.symbolic import symbolic_simple
 
 class TestBottomUpRewrite(unittest.TestCase):
@@ -220,7 +220,7 @@ class TestBottomUpRewrite(unittest.TestCase):
     self.assertIs(gt, ret)
 
 # normally .substitute would be fine, but it's not tracked
-@rewrite_group()
+@track_rewrites()
 def named_substitute(name:str, uop:UOp, rel:dict[UOp, UOp]): return graph_rewrite(uop, _substitute, rel, bottom_up=True)
 def substitute(uop:UOp, rel:dict[UOp, UOp]): return named_substitute(inspect.stack()[1].function, uop, rel)
 

test/null/test_upat_compile.py

@@ -1,11 +1,11 @@
 import unittest
 from tinygrad.helpers import DEBUG, Context
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UPat, rewrite_group, GroupOp, Ops
+from tinygrad.uop.ops import UPat, track_rewrites, GroupOp, Ops
 from tinygrad.uop.upat import _get_code, upat_compile
 import dis
 
-@rewrite_group()
+@track_rewrites()
 def do_compile(up):
   print("\n***** COMPILE", up)
   match_code = _get_code(up, False)

test/null/test_viz.py

@@ -3,7 +3,7 @@ from pathlib import Path
 from dataclasses import dataclass
 from typing import Generator
 
-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, rewrite_group
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, profile_matches
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.helpers import colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
@@ -14,7 +14,7 @@ from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewr
 from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData, get_render, addrspace_colors
 from tinygrad.codegen import do_to_program
 
-@rewrite_group(name=True)
+@track_rewrites(name=True)
 def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp:
   for i,pm in enumerate(pm_lst):
     sink = graph_rewrite(sink, TrackedPatternMatcher(pm.patterns), name=names[i] if names else None)
@@ -109,7 +109,7 @@ class TestViz(unittest.TestCase):
   def test_default_name(self):
     with save_viz() as viz:
       a = UOp.variable("a", 1, 10)
-      @rewrite_group()
+      @track_rewrites()
       def name_default(): return graph_rewrite(a, PatternMatcher([]))
       name_default()
     lst = viz.list_items()
@@ -118,7 +118,7 @@ class TestViz(unittest.TestCase):
   # name can also come from a function that returns a string
   def test_dyn_name_fxn(self):
     with save_viz() as viz:
-      @rewrite_group(name=lambda *args,ret,**kwargs: ret.render())
+      @track_rewrites(name=lambda *args,ret,**kwargs: ret.render())
       def name_from_fxn(s:UOp, arg:list|None=None): return graph_rewrite(s, PatternMatcher([]))
       name_from_fxn(UOp.variable("a", 1, 10)+1, arg=["test"])
     lst = viz.list_items()
@@ -128,18 +128,18 @@ class TestViz(unittest.TestCase):
   # name can also come from a function that returns a TracingKey
   def test_tracing_key(self):
     with save_viz() as viz:
-      @rewrite_group(name=lambda inp,ret: TracingKey("custom_name", (inp,)))
+      @track_rewrites(name=lambda inp,ret: TracingKey("custom_name", (inp,)))
       def test(s:UOp): return graph_rewrite(s, PatternMatcher([]))
       test(UOp.variable("a", 1, 10)+1)
     lst = viz.list_items()
     # NOTE: names from TracingKey do not get deduped
     self.assertEqual(lst[0]["name"], "custom_name")
 
-  def test_nested_rewrite_group(self):
+  def test_nested_track_rewrites(self):
     with save_viz() as viz:
-      @rewrite_group(name=lambda x,ret: TracingKey(f"inner fxn for {x.render()}", (ret,)))
+      @track_rewrites(name=lambda x,ret: TracingKey(f"inner fxn for {x.render()}", (ret,)))
       def inner(x:UOp): return graph_rewrite(x, PatternMatcher([]), name="each")
-      @rewrite_group(name=lambda *args,ret: f"outer rewrite of {len(args)} inputs")
+      @track_rewrites(name=lambda *args,ret: f"outer rewrite of {len(args)} inputs")
       def outer(*xs:tuple[UOp, ...]): return graph_rewrite(UOp.sink(*[inner(x) for x in xs]), PatternMatcher([]), name="all")
       items = ["a", "b", "c"]
       outer(*[UOp.variable(x, 1, 10) for x in items])
@@ -156,13 +156,13 @@ class TestViz(unittest.TestCase):
       self.assertEqual(len(steps), 1)
       self.assertEqual(steps[0]["name"], "each")
 
-  def test_rewrite_group_nested(self):
+  def test_profile_matches(self):
     with save_viz() as viz:
-      @rewrite_group(new_ctx=False)
+      @profile_matches
       def nested_function(u:UOp):
         for i in range(2): graph_rewrite(u, PatternMatcher([]), name=f"step {i+1}")
 
-      @rewrite_group()
+      @track_rewrites()
       def main_rewrite(u:UOp):
         graph_rewrite(u, PatternMatcher([]), name="init")
         nested_function(u)
@@ -173,9 +173,9 @@ class TestViz(unittest.TestCase):
     self.assertEqual(steps[1]["name"], "nested_function")
     self.assertEqual(len(steps), 4)
 
-  def test_rewrite_group_invalid_arg(self):
+  def test_profile_matches_invalid_arg(self):
     with save_viz():
-      @rewrite_group(new_ctx=False)
+      @profile_matches
       def invalid_fxn(arg:str): return graph_rewrite(UOp(Ops.SINK), PatternMatcher([]))
       with self.assertRaisesRegex(AssertionError, "invalid match tracing input"):
         invalid_fxn("test")
@@ -395,7 +395,7 @@ class TestVizIntegration(unittest.TestCase):
     graph = next(viz.get_details(0, 0))["graph"]
     self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1)
 
-  # tracing also works without a rewrite_group context
+  # tracing also works without a track_rewrites context
   # all graph_rewrites get put into the default group
   def test_default_tracing(self):
     with save_viz() as viz:
@@ -407,11 +407,11 @@ class TestVizIntegration(unittest.TestCase):
     self.assertEqual(len(ls), 1)
     self.assertEqual(ls[0]["name"], "default graph_rewrite")
 
-  # using @rewrite_group organizes function calls into groups
+  # using @track_rewrites organizes function calls into groups
   # and nicely counts function calls.
   def test_group_traces(self):
     with save_viz() as viz:
-      @rewrite_group()
+      @track_rewrites()
       def test(root):
         return graph_rewrite(root, sym)
       test(c:=UOp.const(1))
@@ -420,11 +420,11 @@ class TestVizIntegration(unittest.TestCase):
     self.assertEqual(len(ls), 2)
     for i in range(2): self.assertEqual(ls[i]["name"], f"test n{i+1}")
 
-  # @rewrite_group always starts a new group.
+  # @track_rewrites always starts a new group.
   def test_group_combined(self):
     with save_viz() as viz:
       def default_test(root): return graph_rewrite(root, sym)
-      tracked_test = rewrite_group()(default_test)
+      tracked_test = track_rewrites()(default_test)
       c = UOp.const(1)
       default_test(c+1) # goes to the default group
       tracked_test(c)   # all rewrites after this go inside the second group.

test/unit/test_dtype_weak.py

@@ -64,28 +64,14 @@ class TestWeakPromotion(unittest.TestCase):
 
   def test_weak_expression_anchors_at_strong_lub(self):
     # regression test for the HALF bert nan (#17408, reverted in #17409): lub(int32, weakfloat)==weakfloat makes
-    # `loss_mask.sum() + 1e-5` a weakfloat EXPRESSION. Meeting a strong float in a binop must pin it at the lub
+    # `loss_mask.sum() + 1e-5` a weakfloat EXPRESSION. Meeting a strong float in a binop must pin it at the lub, otherwise
+    # nothing owns a width until the bufferize/codegen commits it at default_float: a HALF reciprocal of (sum+1e-5) is inf
     denom = (Tensor.zeros(912, dtype=dtypes.int32) != Tensor.zeros(912, dtype=dtypes.float32)).sum() + 1e-5
     self.assertIs(denom.dtype, dtypes.weakfloat)  # the setup: the denominator expression itself is weak
     x, y = Tensor([2048.0], dtype=dtypes.float32)._broadcasted(denom)
     self.assertIs(y.dtype, dtypes.float32)
     recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIPROCAL]
     self.assertEqual([(u.dtype, u.src[0].dtype) for u in recips], [(dtypes.float32, dtypes.float32)])
-    with Context(DEFAULT_FLOAT=dtypes.float16):
-      committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={})
-    self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32])
-
-  def test_cast_weak_expression_commits_at_cast_floor(self):
-    # the floor never narrows: a cast BELOW the default does not pull the compute width down with it
-    with Context(DEFAULT_FLOAT=dtypes.float32):
-      narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_lower_index_dtype, ctx={})
-    self.assertEqual((narrowed.dtype, narrowed.src[0].dtype), (dtypes.float16, dtypes.float32))
-
-  def test_cast_weak_expression_value_uses_cast_floor(self):
-    with Context(DEFAULT_FLOAT=dtypes.float16):
-      denom = Tensor.ones(1, dtype=dtypes.int32, device="CPU").sum() * 70000 + 1e-5
-      out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom
-      self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10)
 
   def test_uop_scalar_const_lifts_kind(self):
     for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),

tinygrad/codegen/__init__.py

@@ -2,7 +2,7 @@ from dataclasses import replace, dataclass
 import itertools, functools
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
 from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic
-from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp
+from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
 from tinygrad.uop.ops import AxisType, pm_commit_weak, pm_cast_weak
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
@@ -448,7 +448,7 @@ pm_to_program = PatternMatcher([
   (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
 ])
 
-@rewrite_group(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
+@track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
 @Context(ALLOW_DEVICE_USAGE=0)
 def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   """

tinygrad/codegen/decomp/dtype.py

@@ -78,11 +78,9 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
     case Ops.MAX: return l2i(Ops.WHERE, dt, l2i(Ops.CMPLT, dt, *uops), b0, b1, a0, a1)
     case _: raise NotImplementedError(f"long decomposition of {op} unsupported")
 
-def split_l2i(ctx:dict, op: Ops, dt: DType, *uops:UOp):
-  # l2i does arithmetic on its inputs; rules enter here to split them to 32-bit words first, l2i recurses on itself.
-  # both word halves of a node ask for the same split, so ctx memos it for the pass
-  if (key:=(op, dt, uops)) not in ctx: ctx[key] = l2i(op, dt, *graph_rewrite(UOp.sink(*uops), pm_long_decomp, ctx=ctx, bottom_up=True).src)
-  return ctx[key]
+def split_l2i(op: Ops, dt: DType, *uops:UOp):
+  # l2i does arithmetic on its inputs; rules enter here to split them to 32-bit words first, l2i recurses on itself
+  return l2i(op, dt, *graph_rewrite(UOp.sink(*uops), pm_long_decomp, bottom_up=True).src)
 
 # ***** floats *****
 f2f_dt = { f:getattr(dtypes, f"uint{f.bitsize}") for f in dtypes.floats }
@@ -141,21 +139,21 @@ pm_long_decomp = PatternMatcher([
   (UPat(Ops.STORE, src=(UPat.var('idx', tuple(l2i_dt.keys())), UPat.var('val')), name='st'), lambda st,idx,val:
    st.replace(src=(idx.rtag((0, dt:=l2i_dt[idx.dtype])), val.rtag((0, dt)))).group(
      st.replace(src=(idx.rtag((1, dt)), val.rtag((1, dt))))) if val.tag is None else None),
-  (UPat(GroupOp.Comparison, src=[UPat.var('a', tuple(l2i_dt.keys())), UPat()], name="x"), lambda ctx,a,x:
-   split_l2i(ctx, x.op, dt:=l2i_dt[a.dtype], *flatten((s.rtag((0, dt)), s.rtag((1, dt))) for s in x.src))),
-  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x:
-   split_l2i(ctx, Ops.BITCAST, l2i_dt[x.dtype], a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt)))[x.tag[0]]),
-  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda ctx,a,x:
-   split_l2i(ctx, x.op, x.dtype, a)[x.tag[0]] if x.tag is not None else None),
-  (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x:
-   split_l2i(ctx, x.op, x.dtype, a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt))) if x.dtype not in l2i_dt and a.tag is None else None),
-  (UPat((Ops.SHL, Ops.SHR), tuple(l2i_dt.keys()), src=(UPat.var('a'), UPat.var('b')), name="x"), lambda ctx,a,b,x:
-   split_l2i(ctx, x.op, dt:=l2i_dt[x.dtype], a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)))[x.tag[0]] if x.tag is not None else None),
-  (UPat(Ops.WHERE, tuple(l2i_dt.keys()), src=(UPat.var('c'), UPat.var('a'), UPat.var('b')), name="x"), lambda ctx,a,b,c,x:
-   split_l2i(ctx, x.op, dt:=l2i_dt[x.dtype], c, a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)), b.rtag((1, dt)))[x.tag[0]]
+  (UPat(GroupOp.Comparison, src=[UPat.var('a', tuple(l2i_dt.keys())), UPat()], name="x"), lambda a,x:
+   split_l2i(x.op, dt:=l2i_dt[a.dtype], *flatten((s.rtag((0, dt)), s.rtag((1, dt))) for s in x.src))),
+  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
+   split_l2i(Ops.BITCAST, l2i_dt[x.dtype], a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt)))[x.tag[0]]),
+  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda a,x:
+   split_l2i(x.op, x.dtype, a)[x.tag[0]] if x.tag is not None else None),
+  (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
+   split_l2i(x.op, x.dtype, a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt))) if x.dtype not in l2i_dt and a.tag is None else None),
+  (UPat((Ops.SHL, Ops.SHR), tuple(l2i_dt.keys()), src=(UPat.var('a'), UPat.var('b')), name="x"), lambda a,b,x:
+   split_l2i(x.op, dt:=l2i_dt[x.dtype], a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)))[x.tag[0]] if x.tag is not None else None),
+  (UPat(Ops.WHERE, tuple(l2i_dt.keys()), src=(UPat.var('c'), UPat.var('a'), UPat.var('b')), name="x"), lambda a,b,c,x:
+   split_l2i(x.op, dt:=l2i_dt[x.dtype], c, a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)), b.rtag((1, dt)))[x.tag[0]]
    if x.tag is not None else None),
-  (UPat((*(GroupOp.ALU - GroupOp.Comparison - {Ops.SHL, Ops.SHR, Ops.WHERE}), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda ctx,x:
-   split_l2i(ctx, x.op, l2i_dt[x.dtype], *flatten((a.rtag((0, l2i_dt[x.dtype])), a.rtag((1, l2i_dt[x.dtype]))) for a in x.src))[x.tag[0]]
+  (UPat((*(GroupOp.ALU - GroupOp.Comparison - {Ops.SHL, Ops.SHR, Ops.WHERE}), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda x:
+   split_l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag((0, l2i_dt[x.dtype])), a.rtag((1, l2i_dt[x.dtype]))) for a in x.src))[x.tag[0]]
    if x.tag is not None else None),
   (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx:
    x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None),
@@ -199,7 +197,7 @@ def do_dtype_decomps(sink:UOp, ctx:tuple[set[DType], Renderer]) -> UOp:
       to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float
       if DEBUG >= 2: print(f"emulating {fr} as {to}")
       pm = pm_float_decomp if fr in dtypes.floats else pm_long_decomp
-      sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx={} if pm is pm_long_decomp else (fr, to), bottom_up=True)
+      sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx=(fr, to), bottom_up=True)
   ctx[0].clear()
   return sink
 

tinygrad/engine/jit.py

@@ -4,7 +4,7 @@ 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.device import Buffer, Compiled, Device, MultiBuffer, DepsTracker
 from tinygrad.dtype import DType
-from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, rewrite_group, graph_rewrite
+from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite
 from tinygrad.renderer import Estimates
 from tinygrad.engine.realize import capturing, compile_linear, link_linear, run_linear, graph_cache, estimate_uop, get_runtime
 from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_outs_ins
@@ -64,7 +64,7 @@ def _copy_input(u:UOp) -> UOp:
   run_linear(UOp(Ops.LINEAR, src=(u.copy_to_device(u.device).call(new:=UOp.new_buffer(u.device, u.max_numel(), u.dtype), u),)))
   return new
 
-@rewrite_group(lambda linear,held_bufs,input_uops,ret=(): f"JIT {pluralize('call', len(linear.src))}")
+@track_rewrites(lambda linear,held_bufs,input_uops,ret=(): f"JIT {pluralize('call', len(linear.src))}")
 def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
   if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View captured linear")
 

tinygrad/mixin/movement.py

@@ -90,11 +90,8 @@ class MovementMixin:
         if resolve(index.step == 0, False): raise ValueError(f"{index=} cannot have 0 as step")
         start, stop = 0 if index.start is None else index.start, size if index.stop is None else index.stop
         step = 1 if index.step is None else index.step
-        # resolve negative int bounds against the (possibly symbolic) size, like slice.indices
-        if isinstance(start, int) and start < 0: start = start + size
-        if isinstance(stop, int) and stop < 0: stop = stop + size
         if all_int((start, stop, step)):
-          # handle int slicing (clamp, stride)
+          # handle int slicing (resolve negative bounds, clamp, stride)
           *bound, stride = index.indices(int(size.vmax) if isinstance(size, UOp) else size)
           bound = [0, 0] if stride * (bound[1] - bound[0]) < 0 else ([bound[1]+1, bound[0]+1] if stride < 0 else bound)
           return {"size":ceildiv(bound[1]-bound[0], abs(stride)), "boundary":tuple(bound), "stride":stride, "collapse_dim":False}

tinygrad/runtime/support/hcq2.py

@@ -5,7 +5,7 @@ from dataclasses import replace, dataclass
 from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap
 from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer, DepsTracker
-from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, rewrite_group, GroupOp
+from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.dtype import dtypes, truncate
 from tinygrad.runtime.support.hcq import MMIOInterface
@@ -393,7 +393,7 @@ pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, src=(
 
 hcq_compile_cache:dict[bytes, UOp] = {}
 
-@rewrite_group(lambda linear,input_uops,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
+@track_rewrites(lambda linear,input_uops,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
 def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None) -> UOp:
   if input_uops is not None:
     slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))}
@@ -477,7 +477,7 @@ def link_buf_key(a:UOp): return a.key, to_tuple(a.device)
 link_buf_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {}
 link_linear_cache:dict[bytes, UOp] = {}
 
-@rewrite_group(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
+@track_rewrites(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
 def hcq_link(linear:UOp, cache=True) -> UOp:
   if (linked:=link_linear_cache.get(linear_key:=linear.key)) is not None: return linked
 

tinygrad/schedule/__init__.py

@@ -1,6 +1,6 @@
 import time, inspect
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, rewrite_group, graph_rewrite, gate_kernel_sink, KernelInfo
+from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo
 from tinygrad.uop.spec import type_verify, spec_tensor
 from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition, dedup
 
@@ -130,10 +130,6 @@ def lower_sink_to_linear(function:UOp) -> UOp|None:
           f" | {len(UOpMetaClass.ucache):7d} uops in cache"+("" if frm is None else f" | {frm.filename}:{frm.lineno}"))
   return linear
 
-pm_schedule = PatternMatcher([
-  (UPat(Ops.SINK, name="function"), lower_sink_to_linear),
-])
-
 def assert_all_same_devices(ast:UOp):
   devices = dedup([x.device for x in ast.toposort() if x.op is Ops.PARAM and x.device is not None])
   if len(devices) >= 2: raise RuntimeError(f"all buffers must be on the same device: {devices}")
@@ -167,10 +163,10 @@ pm_copy_from_store = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"),), allow_any_len=True), assert_all_same_devices),
 ])
 
-@rewrite_group(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}")
+@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}")
 def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]:
-  # big_sink srcs are all the Tensors
-  linear_call = graph_rewrite(big_sink, pm_schedule, name="schedule to linear", enter_calls=True)
+  assert big_sink.op is Ops.CALL and big_sink.src[0].op is Ops.SINK, "tensor graph is a single sink"
+  linear_call = big_sink.replace(src=(lower_sink_to_linear(big_sink.src[0]),)+big_sink.src[1:])
 
   # this recursively resolves the linear_call and allocates buffers
   linear = graph_rewrite(linear_call, pm_resolve_linear_call, name="resolve linear call")

tinygrad/schedule/indexing.py

@@ -2,7 +2,7 @@ from typing import Iterator
 import functools, itertools
 from dataclasses import dataclass, field, replace
 from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, rewrite_group, broadcast_axes
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches, broadcast_axes
 from tinygrad.uop.ops import gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC
@@ -25,8 +25,6 @@ def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp):
   # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
   if dest.base in src.backward_slice_with_self: ctx[src] = None
 
-BUFFER_STATE_OPS: set[Ops] = {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK, Ops.BIND}
-
 pm_generate_realize_map = PatternMatcher([
   # always realize
   (UPat({Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
@@ -178,7 +176,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
     case _: raise RuntimeError(f"{op} is not a MovementOp")
   return rngs
 
-@rewrite_group(new_ctx=False)
+@profile_matches
 def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
   if debug: print("**************************")
   rctx = IndexingContext()

tinygrad/schedule/rangeify.py

@@ -3,7 +3,7 @@ from typing import cast
 import itertools
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype, strong_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
-from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element
+from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
@@ -551,7 +551,7 @@ pm_copy_to_store = PatternMatcher([
   (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
 ])
 
-@rewrite_group(new_ctx=False)
+@profile_matches
 def get_kernel_graph(sink:UOp) -> UOp:
   tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
   if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")

tinygrad/tensor.py

@@ -8,7 +8,7 @@ from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtyp
   _from_np_dtype, _to_np_dtype, PyConst, AddrSpace
 from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc, VIZ, pluralize
-from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike, UPat, PatternMatcher, GroupOp, ParamArg, graph_rewrite, rewrite_group
+from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike, UPat, PatternMatcher, GroupOp, ParamArg, graph_rewrite, track_rewrites
 from tinygrad.mixin.rand import RandMixin
 from tinygrad.schedule import create_linear_with_vars
 from tinygrad.device import Buffer, canonicalize_device
@@ -211,7 +211,7 @@ pm_replace_buf = PatternMatcher([
   (UPat(Ops.BIND, src=(UPat(Ops.PARAM), UPat(Ops.CONST)), name="b"), replace_input_buffer),
 ])
 
-@rewrite_group(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}")
+@track_rewrites(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}")
 def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
   # uop list is a list in the original_sink graph and we can map to the tags later

tinygrad/uop/ops.py

@@ -1515,52 +1515,55 @@ def add_trace_group(kt:TracingKey) -> None:
   tracked_ctxs.append([])
 
 active_group:list[int] = []
-active_rewrites:list[TrackedGraphRewrite] = []
-def rewrite_group(name:Callable[..., str|TracingKey]|bool=True, replay:bool=False, new_ctx:bool=True):
-  if not new_ctx: assert not callable(name) and not replay, "name fxn and replay are only supported for new_ctx groups"
+def track_rewrites(name:Callable[..., str|TracingKey]|bool=True, replay:bool=False):
   def _decorator(func):
     def __wrapper(*args, **kwargs):
-      # without tracking, we just call the function (unless top-level, which always profiles)
-      if TRACK_MATCH_STATS < 2 and not new_ctx: return func(*args, **kwargs)
       fn = key = func.__name__
       idx = -1
       if TRACK_MATCH_STATS >= 2:
-        if new_ctx:
-          add_trace_group(key:=TracingKey(n:=f"{fn} n{next(_name_cnt.setdefault(fn, itertools.count(1)))}", (n,)))
-          active_group.append(idx:=len(tracked_keys)-1)
-        else:
-          rewrite_name = str(kwargs.get("name", None) or fn)
-          assert args and isinstance(args[0], UOp), f"invalid match tracing inputs for {rewrite_name} with {args}"
-          loc = ((frm:=sys._getframe(1)).f_code.co_filename, frm.f_lineno)
-          depth = len(active_rewrites)
-          if not tracked_ctxs: add_trace_group(TracingKey(f"default {fn}"))
-          dest_group = active_group[-1] if active_group else len(tracked_ctxs)-1
-          tracked_ctxs[dest_group].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], rewrite_name, depth, kwargs.get("bottom_up", False),
-                                                                   kwargs.get("walk", False), kwargs.get("enter_calls", False)))
-          active_rewrites.append(ctx)
-          key = rewrite_name  # profile spans are named after the rewrite step
+        add_trace_group(key:=TracingKey(n:=f"{fn} n{next(_name_cnt.setdefault(fn, itertools.count(1)))}", (n,)))
+        active_group.append(idx:=len(tracked_keys)-1)
       with cpu_profile(key, "TINY") as e:
         ret = func(*args, **kwargs)
-      if TRACK_MATCH_STATS >= 2:
-        if new_ctx: active_group.pop()
-        else: active_rewrites.pop()
-        if callable(name):
-          name_ret = name(*args, **kwargs, ret=ret)
-          assert isinstance(name_ret, (TracingKey, str)), f"name function returned {type(name_ret)}"
-          tracked_keys[idx] = k = TracingKey(n:=tracked_keys[idx].display_name.replace(fn, name_ret), (n,)) if isinstance(name_ret, str) else name_ret
-          e.name = TracingKey(k.display_name if isinstance(name_ret, str) else f"{fn} for {k.display_name}", k.keys)
+      if TRACK_MATCH_STATS >= 2: active_group.pop()
+      if TRACK_MATCH_STATS >= 2 and callable(name):
+        name_ret = name(*args, **kwargs, ret=ret)
+        assert isinstance(name_ret, (TracingKey, str)), f"name function returned {type(name_ret)}"
+        tracked_keys[idx] = k = TracingKey(n:=tracked_keys[idx].display_name.replace(fn, name_ret), (n,)) if isinstance(name_ret, str) else name_ret
+        e.name = TracingKey(k.display_name if isinstance(name_ret, str) else f"{fn} for {k.display_name}", k.keys)
       if CAPTURE_PROCESS_REPLAY and replay:
         # find the unittest frame we're capturing in
         frm = sys._getframe(1)
         while (f_back:=frm.f_back) is not None and "unittest" not in f_back.f_code.co_filename: frm = f_back
-        replay_loc = f"{frm.f_code.co_filename.split('/')[-1]}:{frm.f_lineno} {frm.f_code.co_name}"
+        loc = f"{frm.f_code.co_filename.split('/')[-1]}:{frm.f_lineno} {frm.f_code.co_name}"
         # capture global context vars and all the args passed in
         inputs = (fn, args, kwargs, ContextVar._cache)
-        replay_capture.append(pickle.dumps(inputs+(replay_loc, ret)))
+        replay_capture.append(pickle.dumps(inputs+(loc, ret)))
       return ret
     return __wrapper
   return _decorator
 
+active_rewrites:list[TrackedGraphRewrite] = []
+def profile_matches(fxn:Callable):
+  def wrap_profile_matches(*args, **kwargs):
+    if TRACK_MATCH_STATS >= 2:
+      name = str(kwargs.get("name", None) or fxn.__name__)
+      assert args and isinstance(args[0], UOp), f"invalid match tracing inputs for {name} with {args}"
+      loc = ((frm:=sys._getframe(1)).f_code.co_filename, frm.f_lineno)
+      depth = len(active_rewrites)
+      if not tracked_ctxs: add_trace_group(TracingKey(f"default {fxn.__name__}"))
+      dest_group = active_group[-1] if active_group else len(tracked_ctxs)-1
+      tracked_ctxs[dest_group].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], name, depth, kwargs.get("bottom_up", False),
+                                                               kwargs.get("walk", False), kwargs.get("enter_calls", False)))
+      active_rewrites.append(ctx)
+      with cpu_profile(name, "TINY"):
+        ret = fxn(*args, **kwargs)
+      active_rewrites.pop()
+      return ret
+    # without tracking, we just call the function
+    return fxn(*args, **kwargs)
+  return wrap_profile_matches
+
 class TrackedPatternMatcher(PatternMatcher):
   def rewrite(self, uop:UOp, ctx=None):
     if len(pats:=self.pdict.get(uop.op, [])):
@@ -1739,7 +1742,7 @@ class RewriteContext:
           if n in waitlist: stack.extend(waitlist.pop(n))
     return self.replace[root]
 
-@rewrite_group(new_ctx=False)
+@profile_matches
 def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, walk=False, enter_calls=False) -> UOp:
   rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls)
   return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink)
@@ -1786,11 +1789,10 @@ def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None:
   return None if ret is u else ret
 
 def commit_weak(s:UOp, dt:DType) -> UOp:
-  # a bare weak CONST commits directly (the value stays mathematical, emission truncates), a weak non-const src takes the demand cast
+  # a bare weak CONST commits directly (its number must fit), a weak non-const src takes the demand cast
   return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt)
 
 def commit_weak_srcs(u:UOp) -> UOp|None:
-  if not any(s.dtype in dtypes.weaks for s in u.src): return None
   if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None
   # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too
   return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src))
@@ -1803,17 +1805,14 @@ pm_commit_weak = PatternMatcher([
    lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))),
 ])
 
-# a concrete CAST over a weak node states the width the value will live at. that width is a floor, never a narrowing
-def cast_weak_srcs(c:UOp, u:UOp) -> UOp|None:
-  if c.dtype in dtypes.weaks or weak_dtype(c.dtype) is not u.dtype: return None
-  dt = least_upper_dtype(c.dtype, select_dtype(u))
-  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype)
-
+# push cast to weak src
 pm_cast_weak = PatternMatcher([
-  (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs),
+  (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.Broadcastable, dtype=dtypes.weaks, name="u"),)),
+   lambda c,u: u.replace(dtype=None, src=tuple(commit_weak(s, c.dtype) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype)
+   if c.dtype not in dtypes.weaks else None),
 ])
 
-pm_lower_index_dtype = pm_commit_weak+pm_cast_weak+PatternMatcher([
+pm_lower_index_dtype = pm_commit_weak+PatternMatcher([
   (UPat(GroupOp.All, name="u"),
    lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None),
   # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded)

tinygrad/uop/symbolic.py

@@ -68,7 +68,7 @@ invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i")
 invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
 pm_data_invalid = PatternMatcher([
   (invalid_pat.broadcast(), lambda i: i),
-  (UPat(GroupOp.Unary|{Ops.CAST, Ops.BITCAST}, src=(invalid_pat,)), lambda i: i),
+  (UPat(GroupOp.Unary|{Ops.BITCAST}, src=(invalid_pat,)), lambda i: i),
   (UPat(GroupOp.Unary|{Ops.CAST, Ops.BITCAST}, src=(invalid_gate,), name="op"),
    lambda cond,x,op,i: cond.where(op.replace(src=(x,)), i)),
   # binary ops move inside the gate, with Invalid in the false branch