Merge branch 'master' into simpler_postrange

* start

* tiny clean up

* whoops, didn't mean to accidentally fix this

* fix .to(device), kinda hacky and this fix makes it slower?

* merge properly

* FINALLY figured out slowness, also hack pylint for now

* add DEBUGONNX print for subgraph

* oops

* WOOOOOOOO SHAPE CACHE 50% SPEED INCREASE

* small fix, but maybe all deterministic Tensor creation in fp should be cached

* cache condition

* sliiiightly cleaner

* better abstraction?

* remove sam from model_benchmark

* remove shape cache speed up for now

* less lines

* isinstance fix

---------

Co-authored-by: chenyu <[email protected]>

test/external/external_test_onnx_backend.py

@@ -134,7 +134,6 @@ backend_test.exclude('test_simple_rnn_*')
 
 # no control flow
 # control flow uses AttributeProto.GRAPH
-backend_test.exclude('test_if_*')
 backend_test.exclude('test_loop*')
 backend_test.exclude('test_range_float_type_positive_delta_expanded_cpu') # requires loop
 backend_test.exclude('test_affine_grid_2d_align_corners_expanded_cpu')
@@ -183,6 +182,8 @@ backend_test.exclude('test_resize_downsample_scales_cubic_antialias_cpu') # anti
 backend_test.exclude('test_resize_downsample_sizes_cubic_antialias_cpu') # antialias not implemented
 backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_value_only_mapping_cpu') # bad data type string
 backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad data type string
+backend_test.exclude('test_if_opt_cpu') # ValueError: 13 is not a valid AttributeType
+backend_test.exclude('test_if_seq_cpu') # NotImplementedError: op='SequenceConstruct' is not supported
 
 backend_test.exclude('test_scatternd_min_cpu') # min not yet supported
 backend_test.exclude('test_scatternd_max_cpu') # max not yet supported

test/external/external_test_onnx_ops.py

@@ -100,6 +100,25 @@ class TestMainOnnxOps(TestOnnxOps):
     self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=1)
     self._test_resize_scales([0.01, 0.25, 0.5, 0.51, 0.6, 1.0, 1.5, 2.0, 3.5, 20.0], mode="cubic", exclude_outside=0)
 
+  def _test_if(self, then_value, else_value):
+    then_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, then_value.shape)
+    else_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, else_value.shape)
+
+    then_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(then_value))
+    else_const_node = onnx.helper.make_node("Constant", inputs=[], outputs=["res"], value=onnx.numpy_helper.from_array(else_value))
+
+    then_body = onnx.helper.make_graph([then_const_node], "then_body", [], [then_out])
+    else_body = onnx.helper.make_graph([else_const_node], "else_body", [], [else_out])
+
+    self.helper_test_single_op("If", {"cond": np.array(False).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
+    self.helper_test_single_op("If", {"cond": np.array(True).astype(bool)}, {"then_branch": then_body, "else_branch": else_body}, ["res"])
+
+  def test_if_different_shapes_broadcastable(self):
+    self._test_if(np.array([[1], [2]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
+
+  def test_if_different_shapes_not_broadcastable(self):
+    self._test_if(np.array([[1, 2, 3], [4, 5, 6]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
+
   def test_resize_downsample_scales_linear_align_corners(self):
     # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-131
     X = np.array([[[[1, 2, 3, 4], [5, 6, 7, 8]]]], dtype=np.float32)

test/external/fuzz_fast_idiv.py

@@ -1,8 +1,8 @@
 import random
 import z3
 from tinygrad import dtypes
-from tinygrad.uop.spec import z3_renderer, z3_cdiv
-from tinygrad.uop.ops import UOp, graph_rewrite
+from tinygrad.uop.spec import uops_to_z3, z3_cdiv
+from tinygrad.uop.ops import UOp
 from tinygrad.uop.decompositions import fast_idiv
 random.seed(42)
 
@@ -19,8 +19,7 @@ if __name__ == "__main__":
     if expr is None: continue
 
     solver = z3.Solver()
-    z3_sink = graph_rewrite(expr.sink(u), z3_renderer, ctx=(solver, {}))
-    z3_expr, x = z3_sink.src[0].arg, z3_sink.src[1].arg
+    z3_expr, x =uops_to_z3(solver, expr, u)
 
     if solver.check(z3_expr != z3_cdiv(x, d)) == z3.sat:
       assert False, f"Failed: {expr.render()} != x//{d} at x={solver.model()}\nx={u}\nd={d}\n{z3_expr=}\n{x/d=}"

test/external/fuzz_symbolic.py

@@ -1,8 +1,8 @@
 import random, operator
 import z3
 from tinygrad import Variable, dtypes
-from tinygrad.uop.ops import UOp, graph_rewrite
-from tinygrad.uop.spec import z3_renderer
+from tinygrad.uop.ops import UOp
+from tinygrad.uop.spec import uops_to_z3
 from tinygrad.helpers import DEBUG, Context
 
 seed = random.randint(0, 100)
@@ -57,8 +57,7 @@ if __name__ == "__main__":
 
     solver = z3.Solver()
     solver.set(timeout=5000)  # some expressions take very long verify, but its very unlikely they actually return sat
-    z3_sink = graph_rewrite(expr.sink(simplified_expr, u1, u2, u3), z3_renderer, ctx=(solver, {}))
-    z3_expr, z3_simplified_expr = z3_sink.src[0].arg, z3_sink.src[1].arg
+    z3_expr, z3_simplified_expr, v1, v2, v3 = uops_to_z3(solver, expr, simplified_expr, u1, u2, u3)
     check = solver.check(z3_simplified_expr != z3_expr)
     if check == z3.unknown and DEBUG>=1:
       skipped += 1
@@ -69,7 +68,6 @@ if __name__ == "__main__":
             f"expr = {expr.render(simplify=False)}\n")
     elif check == z3.sat:
       m = solver.model()
-      v1, v2, v3 = z3_sink.src[2].arg, z3_sink.src[3].arg, z3_sink.src[4].arg
       n1, n2, n3 = m[v1], m[v2], m[v3]
       u1_val, u2_val, u3_val = u1.const_like(n1.as_long()), u2.const_like(n2.as_long()), u3.const_like(n3.as_long())
       with Context(CORRECT_DIVMOD_FOLDING=1):

test/test_const_folding.py

@@ -1,11 +1,10 @@
 import unittest, itertools, math
-from typing import Any
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import DType
+from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen import full_rewrite_to_sink
-import numpy as np
 from tinygrad.device import is_dtype_supported
+import numpy as np
 from test.helpers import not_support_multi_device
 
 def _check_ast_count(desired_count:int, t:Tensor):
@@ -25,7 +24,7 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
     _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
     _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
 
-  @unittest.expectedFailure  # no two level fold at lazybuffer
+  @unittest.expectedFailure  # no two level fold
   def test_neg_folding(self):
     _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
     _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
@@ -104,7 +103,7 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
 
 class TestBitcastConstFolding(unittest.TestCase):
   def test_scalar_bitcast(self):
-    def t(cases: dict[DType, Any]):
+    def t(cases: dict[DType, ConstType]):
       for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
         if not math.isnan(from_v):
           r = full_rewrite_to_sink(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0]
@@ -165,7 +164,6 @@ class TestMovedConstFolding(unittest.TestCase):
     _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
 
   def test_cast_padded(self):
-    # NOTE: this is folded due to CAST_BEFORE_VIEW
     if is_dtype_supported(dtypes.int16):
       _check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
       np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])

test/test_memory_planner.py

@@ -120,5 +120,19 @@ class TestMemoryPlanner(unittest.TestCase):
     ]
     check_assign(bs)
 
+  def test_very_small_buffers(self):
+    bs = [
+      [b(0, pin=True), b(1, size=32)],
+      [b(3, size=4), b(4, size=6)],
+    ]
+    check_assign(bs)
+
+  def test_very_big_buffers(self):
+    bs = [
+      [b(0, pin=True), b(1, size=34359738368000)],
+      [b(3, size=1 << 128), b(4, size=1 << 64)],
+    ]
+    check_assign(bs)
+
 if __name__ == "__main__":
   unittest.main()

test/test_rangeify.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor
-from tinygrad.helpers import RANGEIFY
+from tinygrad.helpers import RANGEIFY, Context, GlobalCounters
 
 N = 256
 
@@ -11,6 +11,26 @@ class TestRangeify(unittest.TestCase):
     ba = A.expand(N, N)
     ((ba+1).sum(axis=1) + (ba+2).sum(axis=0)).realize()
 
+  def test_partial_contig(self):
+    A = Tensor.empty(64, 64, 64)
+    ret = A.sum(axis=2).contiguous(arg=(1,)).sum(axis=1)
+    ret.realize()
+
+  def test_double_gemm_real(self):
+    def go():
+      with Context(DEBUG=0):
+        Tensor.manual_seed(1337)
+        A,B,C = [Tensor.randn(N, N) for _ in range(3)]
+        Tensor.realize(A, B, C)
+      GlobalCounters.reset()
+      return (A@B@C).realize()
+    rng = go()
+    with Context(RANGEIFY=0, DEBUG=2):
+      ref = go()
+      mse = ((rng-ref)**2).sum().item()
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-2)
+
   def test_double_gemm(self):
     A = Tensor.empty(N, N)
     B = Tensor.empty(N, N)
@@ -96,14 +116,30 @@ class TestRangeify(unittest.TestCase):
     out.realize()
 
   def test_flash_attention(self):
-    BS = 4
-    HEADS = 2
-    MATDIM = 16
-    EMB = 8
-    q = Tensor.empty(BS, HEADS, MATDIM, EMB)
-    k = Tensor.empty(BS, HEADS, MATDIM, EMB)
-    v = Tensor.empty(BS, HEADS, MATDIM, EMB)
-    q.scaled_dot_product_attention(k, v).realize()
+    #BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
+
+    # bigger
+    BS, HEADS, SEQLEN, EMB = 4, 32, 1024, 64
+
+    # llama 8B
+    #BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
+
+    def fa():
+      Tensor.manual_seed(1337)
+      with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
+      return q.scaled_dot_product_attention(k, v).realize()
+
+    with Context(DEBUG=4):
+      GlobalCounters.reset()
+      ret = fa()
+    with Context(RANGEIFY=0):
+      with Context(DEBUG=2):
+        GlobalCounters.reset()
+        cmp = fa()
+      with Context(DEBUG=0):
+        mse = ((cmp-ret)**2).sum().item()
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
 
 from tinygrad import dtypes
 from tinygrad.uop.ops import UOp

test/test_schedule.py

@@ -1050,6 +1050,14 @@ class TestSchedule(unittest.TestCase):
       compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
       np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
 
+    with Context(FUSE_ATTENTION=1):
+      out = Tensor.scaled_dot_product_attention(q,k,v)
+      run_schedule(check_schedule(out, 1))
+      if getenv("CHECK", 1):
+        import torch
+        compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
+        np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
+
   def test_ugly_reduceop_pairing(self):
     Tensor.manual_seed(0)
     a = Tensor.randn(4, 32).realize()

test/test_tiny.py

@@ -30,7 +30,10 @@ class TestTiny(unittest.TestCase):
   def test_gemm(self, N=64, out_dtype=dtypes.float):
     a = Tensor.ones(N,N).contiguous()
     b = Tensor.eye(N).contiguous()
-    self.assertListEqual((out:=a@b).flatten().tolist(), [1.0]*(N*N))
+    lst = (out:=a@b).tolist()
+    for y in range(N):
+      for x in range(N):
+        self.assertEqual(lst[y][x], 1.0, msg=f"mismatch at ({y},{x})")
     if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
 
   # *** randomness ***

test/unit/test_gguf.py

@@ -53,11 +53,37 @@ class TestGGUF(unittest.TestCase):
   def test_load_tinyllama_q4_0(self): self._test_gguf_load("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q4_0.gguf?download=true")
   def test_load_gpt2_q4_1(self): self._test_gguf_load("https://huggingface.co/PrunaAI/gpt2-GGUF-smashed/resolve/main/gpt2.Q4_1.gguf?download=true")
   def test_load_sample_q6_k(self): self._test_gguf_load("https://huggingface.co/Isotr0py/test-gguf-sample/resolve/main/Quant_Q6_K_1024.gguf?download=true")
+  def test_load_sample_mxfp4(self): self._test_gguf_load("https://huggingface.co/ngxson/boring-testing-tiny/resolve/main/stories260K-mxfp4.gguf?download=true")
 
   def test_dequantization_q4_0(self): self._test_dequantization(ggml.GGML_TYPE_Q4_0)
   def test_dequantization_q4_1(self): self._test_dequantization(ggml.GGML_TYPE_Q4_1)
   def test_dequantization_q8_0(self): self._test_dequantization(ggml.GGML_TYPE_Q8_0)
   def test_dequantization_q6_k(self): self._test_dequantization(ggml.GGML_TYPE_Q6_K)
+  def test_dequantization_mxfp4(self):
+    MXFP4 = 39
+
+    def encode(nibbles, E):
+      packed = [(low & 0xF) | ((high & 0xF) << 4) for low, high in zip(nibbles[:16], nibbles[16:])]
+      return np.array([E] + packed, dtype=np.uint8)
+
+    def decode(code, E):
+      sign = -1.0 if code * 0b1000 else 1.0
+      exp = (code >> 1) & 0b11
+      mant = code & 0b1
+      val = (1.0 + 0.5 * mant) * np.exp2(exp - 1) if exp else 0.5 * mant
+      scale = np.exp2(E - 128) if E >= 2 else np.exp2(-127 if E == 1 else -128)
+      return sign * val * scale
+
+    blocks, expected = [], []
+    rng = np.random.default_rng(42)
+    for _ in range(4):
+      E = rng.integers(0, 256)
+      codes = rng.integers(0, 16, size=32, dtype=np.uint8)
+      blocks.append(encode(codes, E))
+      expected.extend(decode(c, E) for c in codes)
+    tensor = Tensor(np.concatenate(blocks))
+    out = ggml_data_to_tensor(tensor, len(expected), MXFP4)
+    self.assertListEqual(out.numpy().tolist(), np.array(expected, dtype=np.float32).tolist())
 
   def test_expected_failure_unknown_type(self):
     with self.assertRaises(ValueError):

test/unit/test_uop_spec.py

@@ -81,5 +81,16 @@ class TestUOpSpec(unittest.TestCase):
     with self.assertRaisesRegex(RuntimeError, "UOp verification failed"):
       type_verify([a], tensor_uop_spec)
 
+class TestUOpSink(unittest.TestCase):
+  def test_0(self):
+    s = UOp.sink()
+    self.assertEqual(len(s.src), 0)
+
+  def test_1(self):
+    a = UOp.const(dtypes.int, 0)
+    s1 = UOp.sink(a)
+    s2 = a.sink()
+    self.assertIs(s1, s2)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_uop_symbolic.py

@@ -8,7 +8,7 @@ from tinygrad.codegen.late.devectorizer import sym
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad import Variable
-from tinygrad.uop.spec import z3_renderer
+from tinygrad.uop.spec import uops_to_z3
 
 def render(self) -> tuple[str, ConstType, ConstType]:
   # NOTE: we need STORE so the ALU op has children
@@ -32,9 +32,8 @@ class TestSymbolic(unittest.TestCase):
   def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
     if test_z3:
       solver = z3.Solver()
-      z3_sink = graph_rewrite(v.sink(v.simplify()), z3_renderer, ctx=(solver, {}))
-      expr, epxr_simplified = z3_sink.src[0].arg, z3_sink.src[1].arg
-      self.assertEqual(solver.check(expr != epxr_simplified), z3.unsat, "simplified expression not equal to original")
+      expr, expr_simplified = uops_to_z3(solver, v, v.simplify())
+      self.assertEqual(solver.check(expr != expr_simplified), z3.unsat, "simplified expression not equal to original")
     rendered, nmin, nmax = render(v)
     if isinstance(s, tuple): self.assertIn(rendered, s)
     else: self.assertEqual(rendered, s)

tinygrad/codegen/__init__.py

@@ -1,7 +1,7 @@
 from typing import Any, Callable
 import functools
 from dataclasses import dataclass
-from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL
+from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, RANGEIFY, POSTOPT
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp
 from tinygrad.uop.spec import type_verify
 from tinygrad.renderer import Renderer
@@ -9,16 +9,18 @@ from tinygrad.renderer import Renderer
 # import all pattern matchers here
 from tinygrad.codegen.lowerer import pm_lowerer, get_index
 from tinygrad.codegen.quantize import pm_quant
-from tinygrad.codegen.gpudims import pm_add_gpudims
+from tinygrad.codegen.gpudims import pm_add_gpudims, pm_tensor_cores, pm_group_for_reduce, pm_fix_locals, pm_bufferize_loop
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
 from tinygrad.uop.decompositions import get_late_rewrite_patterns
-from tinygrad.codegen.late.expander import migrate_indexing, expander
+from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render
 from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
-from tinygrad.codegen.opt import pm_get_optimization, pm_do_optimize
+from tinygrad.codegen.opt import pm_get_optimization, pm_do_optimize, pm_postrange_opt
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
-from tinygrad.schedule.rangeify import pm_add_buffers, rangeify_codegen
+from tinygrad.codegen.opt.postrange import pm_postrange_opt
+from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen
+from tinygrad.codegen.opt.postrange import pm_flatten_range
 
 @dataclass
 class RewriteStep:
@@ -45,10 +47,10 @@ rewrites_for_linearizer = [
 
 def get_rewrites_for_renderer(opts:Renderer, linearizer:bool=True) -> list[RewriteStep]:
   # cache with the values of the context vars
-  return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)
+  return _get_rewrites_for_renderer(opts, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value, RANGEIFY.value, POSTOPT.value)
 
 @functools.cache
-def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
+def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL, _RANGEIFY, _POSTOPT) -> list[RewriteStep]:
   # ** lowerer (rewrite_shapetracker_with_index) **
   ret: list[RewriteStep] = []
 
@@ -56,28 +58,41 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   ret.extend(rewrites_for_views)
 
   # this is kernel.py
-  ret.append(RewriteStep(pm_get_optimization, ctx=lambda _: opts, name="get optimization"))
-  ret.append(RewriteStep(pm_do_optimize, ctx=lambda _: opts, name="optimize ast"))
+  if not _RANGEIFY: ret.append(RewriteStep(pm_get_optimization, ctx=lambda _: opts, name="get optimization"))
+
+  if not _POSTOPT and not _RANGEIFY: ret.append(RewriteStep(pm_do_optimize, ctx=lambda _: opts, name="optimize ast"))
 
   if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
   ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
+  
+  # add tensor cores
+  if _RANGEIFY:
+    ret.append(RewriteStep(pm_bufferize_loop, name="bufferize loop"))
+    ret.append(RewriteStep(pm_tensor_cores, lambda _: ({}, opts), name="tensor cores", bottom_up=True))
+
+  if _POSTOPT or _RANGEIFY: ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
 
   # ** expander (expand_rewrite) **
   ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
 
+  ret.append(RewriteStep(pm_fix_locals, name="fix locals"))
+
+  # add locals
+  ret.append(RewriteStep(pm_flatten_range+pm_add_buffers_local+rangeify_codegen+pm_group_for_reduce, name="add local buffers"))
+
   # add gpu dims (late). this also handles UNROLL range
   ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
 
   # expand
-  ret.append(RewriteStep(sym+expander, name="expander"))
-
-  # add locals
-  ret.append(RewriteStep(pm_add_buffers+rangeify_codegen, name="add local buffers"))
+  ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))
 
   # ** devectorizer (full_graph_rewrite) **
   # remove reduce
   ret.append(RewriteStep(pm_reduce+gep_pushing, lambda _: ReduceContext(), name="remove_reduce"))
 
+  # add gpu dims (late). this works after devectorize, but it's faster here
+  ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))
+
   # devectorize (TODO: does this need opts?)
   if _DEVECTORIZE >= 2: pm_devectorize = sym+load_store_folding+load_store_indexing
   elif _DEVECTORIZE: pm_devectorize = sym+devectorize+load_store_folding+correct_load_store+load_store_indexing

tinygrad/codegen/gpudims.py

@@ -1,9 +1,10 @@
 import math, functools, operator
-from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType
-from tinygrad.helpers import all_int, partition, flatten, prod, dedup
+from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, graph_rewrite
+from tinygrad.helpers import all_int, partition, flatten, prod, dedup, USE_TC, DEBUG
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.shape.view import get_contraction
 from tinygrad.renderer import Renderer
+from tinygrad.codegen.opt.tc import TensorCore
 
 def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
   # TODO: symbolic shape
@@ -56,17 +57,17 @@ def add_gpudims(ctx:Renderer, s:UOp):
   if any(x.op is Ops.SPECIAL for x in s_topo): return None
 
   # get ranges
-  all_ranges = {x.arg[0]%1000:x for x in s_topo if x.op is Ops.RANGE}
+  all_ranges = {x.arg[0:-1]:x for x in s_topo if x.op is Ops.RANGE}
 
   # extract global/local dims
-  global_dims = sorted(dedup([x.arg[0]%1000 for x in all_ranges.values() if x.arg[1] is AxisType.GLOBAL]))
-  local_dims = sorted(dedup([x.arg[0]%1000 for x in all_ranges.values() if x.arg[1] in (AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
+  global_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] is AxisType.GLOBAL]))
+  local_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
   if not global_dims and not local_dims: return None
 
   # get global and local shape
   ranges = [all_ranges[r] for r in global_dims+local_dims if r in all_ranges]
-  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0]%1000 in global_dims])
-  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0]%1000 in local_dims])
+  global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in global_dims])
+  local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in local_dims])
 
   # get the idxs
   ki: KernelInfo = s.arg
@@ -82,7 +83,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
   for r in s_topo:
     if r.op is not Ops.RANGE: continue
     try:
-      ii = (global_dims+local_dims).index(r.arg[0]%1000)
+      ii = (global_dims+local_dims).index(r.arg[0:-1])
       if r.arg[1] == AxisType.REDUCE: continue
       subs[r] = idxs[ii]
     except ValueError: continue
@@ -114,12 +115,17 @@ def fix_group_for_reduce(x:UOp):
   # do only the non grouped reduces early
   ret = x.replace(src=(x.src[0],)+tuple(reduce_r))
   reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr]
-  buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=(AddrSpace.LOCAL, reduce_gfr[0].arg[0])).index(*upstream_locals, *reduce_loop)
+  buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=AddrSpace.LOCAL).index(*upstream_locals, *reduce_loop)
 
   # gate with an if on the store + do the final reduce
   buf = UOp(Ops.IF, dtype=buf.dtype, src=(functools.reduce(operator.and_, [x.eq(0) for x in reduce_gfr]), buf))
   return buf.reduce(*reduce_loop, arg=x.arg)
 
+pm_group_for_reduce = PatternMatcher([
+  # fix group for reduce
+  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
+])
+
 pm_add_gpudims = PatternMatcher([
   # add gpudims must be last
   (UPat(Ops.SINK, name="s"), add_gpudims),
@@ -130,6 +136,134 @@ pm_add_gpudims = PatternMatcher([
   # fix REDUCEs with UNROLLs
   (UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
   (UPat(Ops.STORE, name="x"), fix_store_unroll),
-  # fix group for reduce
-  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
 ])
+
+def apply_tensor_cores(ctx:tuple[dict, Renderer], in0:UOp, in1:UOp, r_range:UOp, reduceop:UOp):
+  if not USE_TC: return None
+  # tensor cores have three ranges. X, Y, and REDUCE
+  in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0])
+  in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0])
+  #print(len(in0_ranges), len(in1_ranges))
+  if not len(in0_ranges) or not len(in1_ranges): return None
+  in0_range, in1_range = in0_ranges[0], in1_ranges[0]
+  if DEBUG >= 2: print('TC', in0_range.arg, in1_range.arg, r_range.arg)
+
+  # confirm the dtype and size is good
+  tc_opts: list[TensorCore] = []
+  for tc in ctx[1].tensor_cores:
+    if reduceop.dtype == tc.dtype_out and in0.dtype == tc.dtype_in and in1.dtype == tc.dtype_in:
+      if all(i <= j for i,j in zip(tc.dims, [in0_range.vmax+1, in1_range.vmax+1, r_range.vmax+1])):
+        tc_opts.append(tc)
+  if len(tc_opts) == 0: return None
+  tc = tc_opts[0]
+
+  # create the new ranges as speced by the tensor core
+  old_range = [in0_range, in1_range, r_range]
+  new_range = [r.replace(src=(r.src[0]//tc.dims[i],)) for i,r in enumerate(old_range)]
+  new_reduce_range = new_range[2]
+  tc_range = 9050 #+ r_range.arg[0]*100
+  red_ranges = []
+
+  ne: list[UOp] = []
+  for o in tc.opts:
+    lrange = UOp.range(dtypes.int, 2, tc_range, AxisType.UPCAST if o[0] == "u" else AxisType.LOCAL)
+    ne.append(lrange)
+    tc_range += 1
+    new_range[1-int(o[1])] = (2 * new_range[1-int(o[1])]) + lrange
+  for _, amt in tc.get_reduce_axes():
+    lrange = UOp.range(dtypes.int, amt, tc_range, AxisType.UNROLL)
+    ne.append(lrange)
+    red_ranges.append(lrange)
+    tc_range += 1
+    new_range[2] = (amt * new_range[2]) + lrange
+  tne = [x.replace(tag=1) for x in ne]
+
+  # replace ranges in other parts of the graph
+  for x,y in zip(old_range, new_range): ctx[0][x] = y
+
+  # apply the swizzled ranges to the srcs
+  srcs = [s.substitute(dict(zip(old_range, new_range))).substitute(dict(zip(ne, tne))) for s in (in0, in1)]
+  srcs = [x.substitute(dict(zip(tne, [ne[i] for i in p]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
+
+  ned = dict(zip(tc.base_shape_str(), ne))
+  tc_reduce_axes = tuple([ned[f"r{i}"].arg[0] for i in range(len(tc.get_reduce_axes()))])
+  base_upcast_axes = tuple([(ned[s].arg[0], 2) for s in tc.base_upcast_axes()])
+  tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
+
+  # construct the op
+  # TODO: remove tc_upcast_axes from the arg
+  wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, ctx[1].device, tc.threads, tc_upcast_axes, tc_reduce_axes)
+  wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
+    UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0]),
+    UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1]),
+    UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg)
+  tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2])
+  ret = tc_uop.reduce(new_reduce_range, arg=Ops.ADD)
+  # confirm the UNROLLs aren't actually used, these need to be broadcast MUL
+  assert all(u not in red_ranges for u in ret.toposort()), "UNROLLs in TC"
+  return ret
+
+from tinygrad.codegen.opt.postrange import pm_flatten_range
+
+pm_tensor_cores = PatternMatcher([
+  ((UPat.var("in0")*UPat.var("in1")).reduce(UPat(Ops.RANGE, name="r_range"), name="reduceop", arg=Ops.ADD), apply_tensor_cores),
+
+  # replace range
+  #(UPat(Ops.RANGE, name="r"), lambda ctx,r: ctx[0].get(r, None)),
+  (UPat(Ops.SINK, name="s"), lambda ctx,s: graph_rewrite(s.substitute(ctx[0]), pm_flatten_range, name="flatten")),
+])
+
+def fix_bufferize(x:UOp):
+  if x.arg != AddrSpace.LOCAL: return None
+  locals_left = [r for r in x.ranges if r.arg[1] == AxisType.LOCAL]
+  if not len(locals_left): return None
+  acc = x.size
+  st = []
+  for l in locals_left:
+    st.append(l*acc)
+    acc *= l.vmax+1
+  return x.replace(src=(x.src[0],) + tuple(locals_left[::-1]) + x.src[1:]).index(sum(st))
+
+
+pm_double_index = PatternMatcher([
+  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("x"))), UPat.var("y"))), lambda b,x,y: b.index(x+y)),
+])
+
+pm_fix_locals = pm_double_index+PatternMatcher([
+  (UPat(Ops.BUFFERIZE, name="x"), fix_bufferize),
+])
+
+B = 8
+
+def loop_store(x:UOp):
+  r_maybe = [r for r in x.src[0].ranges if r.arg[0] == 2 and r.tag is None]
+  ur_maybe = [r for r in x.ranges if r.arg[0] < 0]
+  #print("store", len(r_maybe), len(ur_maybe))
+  if not len(r_maybe) or not len(ur_maybe): return None
+
+  r = r_maybe[0]
+  ur = ur_maybe[0]
+  rr = r.replace(src=(r.src[0]//B,), tag=1)
+  return x.substitute({r:rr*B+ur})
+
+def loop_bufferize(x:UOp):
+  if x.arg != AddrSpace.LOCAL: return None
+  r_maybe = [r for r in x.ranges if r.arg[0] == 2 and r.tag is None]
+  ur_maybe = [r for r in x.ranges if r.arg[0] < 0]
+
+  if len(ur_maybe):
+    ur = ur_maybe[0]
+  else:
+    if len(r_maybe) == 0: return None
+    r = r_maybe[0]
+    ur = UOp.range(dtypes.int, B, -2010)
+    rr = r.replace(src=(r.src[0]//B,), tag=1)
+    x = x.substitute({r:rr*B + ur})
+
+  ur1 = UOp.range(dtypes.int, B, ur.arg[0]+1)
+  return x.replace(src=(x.src[0],ur)+x.src[1:]).index(x.size*ur1)
+
+pm_bufferize_loop = pm_double_index+PatternMatcher([
+  (UPat(Ops.BUFFERIZE, name="x"), loop_bufferize),
+  (UPat(Ops.STORE, name="x"), loop_store),
+])

tinygrad/codegen/late/devectorizer.py

@@ -232,17 +232,21 @@ def no_vectorized_alu(alu:UOp):
   alus = tuple(UOp(alu.op, alu.dtype.scalar(), tuple(s.gep(i) for s in alu.src), alu.arg) for i in range(alu.dtype.vcount))
   return UOp(Ops.VECTORIZE, alu.dtype, alus)
 
-def no_vectorized_buf(buf:UOp, idx:UOp):
-  # NOTE: this should work for define reg too
-  if (cnt:=buf.dtype.count) == 1: return None
-  new_buf = buf.replace(dtype=buf.dtype.base.scalar().ptr(cast(PtrDType, buf.dtype).size*cnt, cast(PtrDType, buf.dtype).addrspace))
-  return new_buf.broadcast(cnt).index(idx.broadcast(cnt)*cnt+UOp.const(dtypes.int.vec(cnt), tuple(range(cnt))))
+def no_vectorized_buf(buf:UOp):
+  dtype = cast(PtrDType, buf.dtype)
+  return buf.replace(dtype=dtype.base.scalar().ptr(dtype.size*dtype.count, dtype.addrspace)).cast(dtype)
+
+def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp):
+  cnt = cast.dtype.count
+  assert idx.dtype.count == 1, f"idx dtype must be 1 {idx.dtype}"
+  return buf.broadcast(cnt).index(idx.broadcast(cnt)*cnt+UOp.const(dtypes.int.vec(cnt), tuple(range(cnt))))
 
 devectorize = PatternMatcher([
   # no ALU on vectorized dtypes
   (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name="alu"), no_vectorized_alu),
   (UPat(Ops.WMMA, name="wmma"), no_vectorized_wmma),
-  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf").index(UPat.var("idx")), no_vectorized_buf),
+  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf"), no_vectorized_buf),
+  (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf").cast(name="cast").index(UPat.var("idx")), no_vectorized_index),
 ])
 
 pm_render = PatternMatcher([

tinygrad/codegen/late/expander.py

@@ -1,9 +1,9 @@
 # this converts a lowerer program into a vectorized program
 
 import functools, itertools, operator
-from tinygrad.dtype import dtypes, PtrDType
-from tinygrad.helpers import AMX, dedup, flatten, all_same, prod
-from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp
+from tinygrad.dtype import dtypes, PtrDType, AddrSpace
+from tinygrad.helpers import AMX, dedup, flatten, all_same, prod, partition
+from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, AxisType
 
 def _expand_arg_to_idx(args:tuple[tuple[int, int], ...], rpk:dict[int, int]) -> int:
   idx, mul = 0, 1
@@ -50,7 +50,7 @@ def do_expand(root:UOp):
       if root.op is Ops.IF or src.op is Ops.IF:
         # for the first arg of IF, just pass them through ignoring UNROLLS
         new_srcs.append(src)
-      elif (root.op is Ops.STORE and i >= 2) or (root.op in {Ops.REDUCE, Ops.BUFFERIZE} and i >= 1):
+      elif (root.op is Ops.STORE and i >= 2) or (root.op in {Ops.REDUCE, Ops.BUFFERIZE} and i >= 1) or (root.op is Ops.WMMA and i >= 3):
         # for any range args of STORE/REDUCE, pass them through
         new_srcs.append(src)
       elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType):
@@ -114,3 +114,49 @@ migrate_indexing = PatternMatcher([
   # create gate MUST BE BEFORE expander
   (UPat(Ops.STORE, name="root"), create_gate),
 ])
+
+# ****
+
+def fix_reduce_unroll(x:UOp):
+  reduce_range, reduce_expand = partition(x.src[1:], lambda y: y.op is Ops.RANGE)
+  if len(reduce_expand) == 0: return None
+  reduce_expand = [x for x in reduce_expand if x.op is not Ops.CONST]
+  assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand}"
+  ret = x.src[0]
+  if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
+    ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis), tag=1)
+  # REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
+  return x.replace(src=(ret,)+tuple(reduce_range))
+
+def fix_store_unroll(x:UOp):
+  store_expand, store_range = partition(x.src[2:], lambda y: y.op is Ops.UNROLL)
+  if len(store_expand) == 0: return None
+  return UOp(Ops.CONTRACT, dtypes.void, (x.replace(src=x.src[:2]+tuple(store_range)),), tuple(flatten(x.arg for x in store_expand)), tag=1)
+
+def fix_group_for_reduce(x:UOp):
+  reduce_gfr, reduce_r = partition(x.src[1:], lambda u: u.op is Ops.RANGE and u.arg[1] == AxisType.GROUP_REDUCE)
+  if len(reduce_gfr) == 0: return None
+
+  # NOTE: if there's other locals here, we need them in the buffer too
+  upstream_locals = [u for u in x.toposort() if u.op is Ops.RANGE and u.arg[1] == AxisType.LOCAL]
+
+  # do only the non grouped reduces early
+  ret = x.replace(src=(x.src[0],)+tuple(reduce_r))
+  reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr]
+  buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=(AddrSpace.LOCAL, reduce_gfr[0].arg[0])).index(*upstream_locals, *reduce_loop)
+
+  # gate with an if on the store + do the final reduce
+  buf = UOp(Ops.IF, dtype=buf.dtype, src=(functools.reduce(operator.and_, [x.eq(0) for x in reduce_gfr]), buf))
+  return buf.reduce(*reduce_loop, arg=x.arg)
+
+pm_pre_expander = PatternMatcher([
+  # rewrite UPCAST/UNROLL range to something to be expanded
+  (UPat(Ops.RANGE, name="r"),
+   lambda r: UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
+    if r.arg[1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
+  # fix REDUCEs with UNROLLs
+  (UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
+  (UPat(Ops.STORE, name="x"), fix_store_unroll),
+  # fix group for reduce
+  (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
+])

tinygrad/codegen/lowerer.py

@@ -1,7 +1,7 @@
 # the job of the lowerer is to do indexing
 from dataclasses import dataclass
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite
+from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite, resolve
 
 # ***** indexing *****
 
@@ -12,12 +12,12 @@ class IndexContext:
   start: int = 0
 
 def shape_to_idx(s, axis_types, start=0):
-  return [UOp.range(dtypes.int, sint_to_uop(s), start+i, axistype=at) for i, (s, at) in enumerate(zip(s, axis_types))]
+  return [UOp.range(dtypes.int, sint_to_uop(s), start+i, at) for i, (s, at) in enumerate(zip(s, axis_types))]
 
 def get_index(ast:UOp) -> IndexContext:
   axis_types = ast.arg.axis_types if isinstance(ast.arg, KernelInfo) else ()
   if len(ast.full_shape) != len(axis_types):
-    axis_types = tuple([AxisType.REDUCE if s is not fs else AxisType.LOOP for s,fs in zip(ast.shape, ast.full_shape)])
+    axis_types = tuple([AxisType.REDUCE if resolve(s != fs) else AxisType.LOOP for s,fs in zip(ast.shape, ast.full_shape)])
   return IndexContext(axis_types, [], 0)
 
 # ***** lowering (given index) *****

tinygrad/codegen/opt/__init__.py

@@ -1,6 +1,7 @@
 # opt opinionatedly transforms an ast into an optimized ast using either heuristics or beam search
 
 from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.postrange import RKernel, pm_flatten_range
 from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, KernelInfo
 from tinygrad.helpers import NOOPT, BEAM, USE_TC, getenv
@@ -44,3 +45,15 @@ def apply_opt(ast:UOp, renderer:Renderer):
 pm_do_optimize = PatternMatcher([
   (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_opt(ast, ctx) if ast.arg is not None and ast.arg.opts_to_apply is not None else None),
 ])
+
+# ** postrange **
+
+def apply_ropt(ast:UOp, renderer:Renderer):
+  k = RKernel(ast, opts=renderer)
+  if ast.arg is not None: k.apply_opts(ast.arg.opts_to_apply)
+  return k.get_optimized_ast()
+
+pm_postrange_opt = pm_flatten_range+PatternMatcher([
+  (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_ropt(ast, ctx) if ast.arg is None or \
+   (ast.arg is not None and ast.arg.opts_to_apply is not None) else None),
+])

tinygrad/codegen/opt/kernel.py

@@ -92,10 +92,6 @@ class Kernel:
     global_loops = AxisType.GLOBAL if self.opts.has_local else AxisType.LOOP
     self.axis_types: list[AxisType] = [AxisType.REDUCE if resolve(x!=y) else global_loops for x,y in zip(self.output_shape, self.full_shape)]
 
-    # confirm all reduce axes are at the end
-    if (final_reduces := [x for x in self.axis_types if x == AxisType.REDUCE]) and final_reduces != self.axis_types[-len(final_reduces):]:
-      raise RuntimeError(f"reduces are not at the end of the shape {self.full_shape} -> {self.output_shape}")
-
   def copy(self):
     ret = type(self).__new__(type(self))
 
@@ -245,7 +241,7 @@ class Kernel:
       if axis is None: return -1
       if op is OptOps.UNROLL: return self.unrollable_dims[axis]
       if op in {OptOps.GROUP, OptOps.GROUPTOP}: return self.axes_of(AxisType.REDUCE)[axis]
-      check(axis < self.shape_len, "invalid axis")
+      check(axis < self.shape_len, f"invalid axis on {axis=} {op=} {self.shape_len=}")
       return axis
     except IndexError as e: raise KernelOptError from e
 

tinygrad/codegen/opt/postrange.py

@@ -0,0 +1,183 @@
+import math
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import UOp, Ops, sint, ssimplify, AxisType, KernelInfo, PatternMatcher, UPat, graph_rewrite, _substitute
+from tinygrad.uop.symbolic import symbolic
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.renderer import Renderer
+from tinygrad.dtype import dtypes
+
+def flatten_range(r:UOp):
+  off = 2 if r.op is Ops.STORE else 1
+  rngs = r.src[off:]
+  if not len(rngs): return None
+  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
+  return r.replace(src=r.src[:off]+tuple(new_rngs))
+
+pm_flatten_range = PatternMatcher([
+  # real ranges only
+  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
+])
+
+def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
+
+class RKernel(Kernel):
+  def __init__(self, ast:UOp, opts:Renderer|None=None):
+    self.rng = sorted([u for u in ast.toposort() if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: x.arg)
+    super().__init__(ast, opts)
+    self.sts.clear()
+
+    # convert LOOP to GLOBAL
+    self.replaces = {}
+    if self.opts.has_local:
+      store_rngs = self.ast.src[0].src[2:]
+
+      # filter any not in local stores
+      local_store_rngs = [x.ranges for x in self.ast.toposort() if (x.op is Ops.STORE and x.src[0].dtype.addrspace == AddrSpace.LOCAL) \
+                          or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
+      for ls in local_store_rngs: store_rngs = [x for x in store_rngs if x in ls]
+
+      store_rng = [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE] if store_rngs else []
+      rng = [x.replace(arg=(x.arg[0], AxisType.GLOBAL)) if x.arg[1] == AxisType.LOOP and x in store_rng else x for x in self.rng]
+      self.replaces.update(dict(zip(self.rng, rng)))
+      self.rng = rng
+
+  def simplify_merge_adjacent(self):
+    return
+    # NOTE: this one is better than the one in kernel.py, which is kind of a problem
+    terminators = [u for u in self.ast.toposort() if u.op in {Ops.REDUCE, Ops.STORE}]
+    termination = {}
+    for t in terminators:
+      for u in t.src[1 if t.op is Ops.REDUCE else 2:]: termination[u] = t
+
+    replaces = {}
+    i = 0
+    while i < len(self.rng)-1:
+      r0, r1 = self.rng[i], self.rng[i+1]
+      # same axistype and same termination
+      if r0.arg[1] == r1.arg[1] and termination[r0] == termination[r1]:
+        s0, s1 = r0.src[0], r1.src[0]
+        new_range = r0.replace(src=(s0*s1,)).simplify()
+        # this checks the legality of a merge
+        oidx = self.ast.simplify()
+        nidx = graph_rewrite(oidx, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1}, name=f"check_merge_{i}_{i+1}")
+        # it simplifies
+        if count_divmod(nidx) <= count_divmod(oidx):
+          # it is correct
+          midx = graph_rewrite(nidx, _substitute+symbolic+pm_flatten_range, ctx={new_range:r0*s1+r1}, name=f"correct_merge_{i}_{i+1}")
+          if oidx is midx:
+            termination[new_range] = termination[r0]
+            replaces[r0] = new_range//s1
+            replaces[r1] = new_range%s1
+            self.rng[i] = new_range
+            del self.rng[i+1]
+            continue
+      i += 1
+    self.ast = self.ast.substitute(replaces, name="simplify_merge_adjacent")
+
+  def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None):
+    old_sz = self.rng[axis].src[0].arg // amount
+    assert old_sz > 0, f"bad old_sz on {axis} {amount} {self.rng[axis]}"
+
+    maxarg = max([x.arg[0] for x in self.rng])
+    new_rng = UOp.range(dtypes.int, amount, maxarg+1, new_type)
+
+    if old_sz == 1:
+      self.replaces[self.rng[axis]] = new_rng
+      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
+      del self.rng[axis]
+    else:
+      replaced_rng = self.rng[axis].replace(src=(UOp.const(dtypes.int, old_sz),))
+      self.replaces[self.rng[axis]] = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
+      self.rng[axis] = replaced_rng
+      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
+    return new_rng
+
+  @property
+  def axis_types(self) -> list[AxisType]: return [x.arg[1] for x in self.rng]
+  @property
+  def shape_len(self): return len(self.rng)
+
+  @property
+  def full_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.rng])
+  @property
+  def output_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.ast.src[0].src[2:]])
+
+  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
+    ret = self.ast
+    kernel_name = ret.arg.name if ret.arg is not None and ret.arg.name != "test" else self.name if name_override is None else name_override
+    rarg = KernelInfo(kernel_name, tuple(self.axis_types), self.dont_use_locals, tuple(self.applied_opts))
+    return ret.substitute(self.replaces).replace(arg=rarg)
+
+  # does nothing
+  @axis_types.setter
+  def axis_types(self, value): pass
+
+  def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> bool:
+    reduceop = [x for x in self.ast.toposort() if x.op is Ops.REDUCE][0]
+    if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD:
+      tensor_cores = self.opts.tensor_cores if tc_select == -1 else [self.opts.tensor_cores[tc_select]]
+      for tc in tensor_cores:
+        if tc.dtype_in == dtypes.float and tc.dtype_out == dtypes.float:
+          axes = [1,0]
+
+          # do optimizations and save the ranges
+          ne: list[UOp] = []
+          for opt in tc.opts:
+            ne.append(self.apply_opt(Opt({"u":OptOps.UPCAST, "l":OptOps.LOCAL}[opt[0]], axes[int(opt[1])], 2), append_opt=False))
+          for _, amt in tc.get_reduce_axes():
+            ne.append(self.apply_opt(Opt(OptOps.UNROLL, 0, amt), append_opt=False)) # TODO: this should be the reduce, not 0
+
+          # early realize for TC
+          self.ast = self.ast.substitute(self.replaces)
+          self.replaces = {}
+
+          # fix the srcs
+          reduceop = [x for x in self.ast.toposort() if x.op is Ops.REDUCE][0]
+          tne = [x.replace(tag=1) for x in ne]
+          ret = reduceop.substitute(dict(zip(ne, tne)))
+          srcs = list((ret.src[0] if ret.src[0].op is not Ops.CAST else ret.src[0].src[0]).src)
+          srcs = [x.substitute(dict(zip(tne, [ne[i] for i in p]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
+
+          # get reduce/upcast axes for the tensor cores
+          tc_reduce_axes = self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))])
+          base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis(tc.base_upcast_axes())])
+          tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
+
+          # axes to range number (was done in lowerer)
+          tc_upcast_axes = tuple([tuple([(self.rng[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes])
+          tc_reduce_axes = tuple([self.rng[a].arg[0] for a in tc_reduce_axes])
+
+          # construct the op
+          # TODO: remove tc_upcast_axes from the arg
+          wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.opts.device, tc.threads, tc_upcast_axes, tc_reduce_axes)
+          wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
+            UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0]),
+            UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1]),
+            UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg)
+          tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2])
+
+          # preserve extra reduces
+          reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
+          if len(reduce_ranges): tc_uop = UOp(Ops.REDUCE, tc_uop.dtype, (tc_uop,)+tuple(reduce_ranges), Ops.ADD)
+          self.ast = self.ast.substitute({reduceop: tc_uop})
+          return True
+    return False
+
+from dataclasses import replace
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo
+from tinygrad.helpers import colored
+from tinygrad.codegen.opt.kernel import axis_colors
+
+def rename_sink(s:UOp):
+  if s.arg is not None and s.arg.name != "test": return None
+
+  # get all ranges (sorted)
+  rngs = sorted([u for u in s.parents if u.op is Ops.RANGE], key=lambda x: x.arg[0:-1])
+
+  # add name to kernel
+  name = "k" + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), axis_colors[x.arg[-1]]) for x in rngs])
+  return s.replace(arg=KernelInfo(name=name) if s.arg is None else replace(s.arg, name=name))
+
+pm_postrange_opt = PatternMatcher([
+  (UPat(Ops.SINK, name="s"), rename_sink),
+])

tinygrad/engine/memory.py

@@ -23,12 +23,13 @@ def _internal_memory_planner(buffers:list[list[Buffer]], noopt_buffers=None, ign
   # Sort buffer operations in timeline order. Two events: buffer is allocated or buffer is freed.
   buffer_requests = sorted([((first_appearance[buf], True), buf) for buf in first_appearance.keys()] + \
                            [((last_appearance[buf] + 1, False), buf) for buf in first_appearance.keys()], key=lambda x: x[0])
+  total_memory = sum(round_up(buf.nbytes, min_block_size:=0x1000) for buf in first_appearance.keys()) * 2 # *2 for fragmentation (which is about 15%)
 
   # Try to suballocate from a shared buffer managed by global_planner using TLSFAllocator.
   # Also track buffer replacements for buffers that do not support suballocation.
   buffer_replace:dict[Buffer, tuple[Buffer|None, int|None]] = {}
   reuse_buffers:dict[tuple, list[Buffer]] = defaultdict(list)
-  global_planner:dict[str, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(1 << 44, block_size=0x1000, lv2_cnt=32)))
+  global_planner:dict[str, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(total_memory, block_size=min_block_size, lv2_cnt=32)))
   for (_, is_open_ev), buf in buffer_requests:
     # Check if suballocation is possible for the given buffer and device.
     if hasattr(Device[buf.device].allocator, "_offset") and not isinstance(buf.dtype, ImageDType):

tinygrad/engine/realize.py

@@ -160,10 +160,15 @@ class ExecItem:
       if DEBUG >= 2:
         lds_est = sym_infer(self.prg.estimates.lds, var_vals)
         mem_est = min(mem_est, lds_est)   # there can't be more memory accessed than loads/stores. remove this when symbolic is fixed
+        header_color = 'magenta' if jit else ('green' if self.prg.first_run else None)
         ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
-        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', 'magenta' if jit else ('green' if self.prg.first_run else None))} {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB " +  # noqa: E501
-              (str() if et is None else f"tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({op_est/((et or 1e-20)*1e9):9.2f} GFLOPS {mem_est/((et or 1e-20)*1e9):6.1f}|{lds_est/((et or 1e-20)*1e9):<7.1f} GB/s)" +  # noqa: E501
-               f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}"))
+        flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
+        flops_str = f"{flops*1e-9:9.2f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:9.2f} TFLOPS", 'green')
+        mem_str = f"{membw*1e-9:6.1f}|{ldsbw*1e-9:<7.1f} GB/s" if membw < 1e13 else colored(f"{membw*1e-12:6.1f}|{ldsbw*1e-12:<7.1f} TB/s", 'green')
+        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
+          f" {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB"+
+          ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})")+
+          f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}")
       self.prg.first_run = False
     return et
 

tinygrad/frontend/onnx.py

@@ -21,9 +21,9 @@ class AttributeType(enum.IntEnum):
   ONNX attribute type identifiers.
   Reference: https://github.com/onnx/onnx/blob/rel-1.18.0/onnx/onnx.proto3#L128-L145
   """
-  FLOAT = 1; INT = 2; STRING = 3; TENSOR = 4; FLOATS = 6; INTS = 7; STRINGS = 8 # noqa: E702
+  FLOAT = 1; INT = 2; STRING = 3; TENSOR = 4; GRAPH = 5; FLOATS = 6; INTS = 7; STRINGS = 8 # noqa: E702
 
-  def to_field_name(self) -> str: return {1: "f", 2: "i", 3: "s", 4: "t", 6: "floats", 7: "ints", 8: "strings"}[self.value]
+  def to_field_name(self) -> str: return {1: "f", 2: "i", 3: "s", 4: "t", 5: "g", 6: "floats", 7: "ints", 8: "strings"}[self.value]
 
 class OnnxDataType(enum.IntEnum):
   """
@@ -266,6 +266,7 @@ class OnnxPBParser:
         case 3: obj["i"] = self.reader.read_int64()
         case 4: obj["s"] = self.reader.read_bytes().data().tobytes().decode("utf8")
         case 5: obj["t"] = self._parse_TensorProto()['parsed_tensor']
+        case 6: obj["g"] = OnnxRunner._from_subgraph(self._parse_GraphProto())
         case 7: obj["floats"].append(self.reader.read_float())
         case 8: obj["ints"].append(self.reader.read_int64())
         case 9: obj["strings"].append(self.reader.read_bytes().data().tobytes().decode("utf8"))
@@ -401,8 +402,11 @@ class OnnxRunner:
   """
   def __init__(self, model_path: Tensor | str | pathlib.Path):
     model = OnnxPBParser(model_path, load_external_data=True).parse()
-    graph = model["graph"]
+    self._init_from_graph(model["graph"])
+
+  def _init_from_graph(self, graph: dict, is_subgraph: bool = False):
     self.is_training = any(n['parsed_node'].opset_id.domain in {Domain.AI_ONNX_TRAINING, Domain.AI_ONNX_PREVIEW_TRAINING} for n in graph["node"])
+    self.graph_name = graph["name"] if is_subgraph else ""
     self.graph_values = {"": None, **{i["name"]: i["parsed_tensor"] for i in graph["initializer"]}}
     self.graph_inputs = {i["name"]: i["parsed_type"] for i in graph["input"] if i["name"] not in self.graph_values}
     self.graph_outputs = tuple(o["name"] for o in graph["output"])
@@ -414,6 +418,12 @@ class OnnxRunner:
     self.variable_dims: dict[str, int] = {}
     self.onnx_ops = onnx_ops
 
+  @classmethod
+  def _from_subgraph(cls, graph: dict) -> "OnnxRunner":
+    subgraph = cls.__new__(cls)
+    subgraph._init_from_graph(graph, is_subgraph=True)
+    return subgraph
+
   def _parse_input(self, name: str, value: Any, spec: OnnxValue):
     if spec.is_optional and value is None: return None
     if spec.is_sequence:
@@ -445,9 +455,10 @@ class OnnxRunner:
     return {name:Tensor.empty(*spec.shape, device=device, dtype=dtype or spec.dtype) for name, spec in self.graph_inputs.items()}
 
   def to(self, device:str|None):
-    self.graph_values = {k:v.to(device) if isinstance(v, Tensor) else v for k,v in self.graph_values.items()}
+    self.graph_values = {k: (v.to(device) if isinstance(v, Tensor) else v) for k,v in self.graph_values.items()}
     self.graph_nodes = tuple(OnnxNode(n.op, n.opset_id, tuple(n.inputs), tuple(n.outputs),
-                                      {k:v.to(device) if isinstance(v, Tensor) else v for k,v in n.opts.items()}) for n in self.graph_nodes)
+                                      {k: (v.to(device) if isinstance(v, (Tensor, OnnxRunner)) else v) for k,v in n.opts.items()})
+                                      for n in self.graph_nodes)
     return self
 
   def __call__(self, inputs:dict[str, Any], debug=debug):
@@ -461,9 +472,9 @@ class OnnxRunner:
 
       # provide additional opts
       if node.op == "Split" and 'num_outputs' not in opts: opts['num_outputs'] = len(node.outputs)
-      if node.op == "Gradient": opts['intermediate_tensors'] = self.graph_values
+      if node.op in {"Gradient", "If"}: opts['intermediate_tensors'] = self.graph_values
 
-      if debug >= 1: print(f"{num}: op '{node.op}' opt {opts}")
+      if debug >= 1: print((f"[{self.graph_name}] " if self.graph_name else "") + f"{num}: op '{node.op}' opt {opts}")
       if debug >= 2 and node.inputs: print("\tinputs:\n" + "\n".join(f"\t\t{x} - {i!r}" for x,i in zip(node.inputs, inps)))
       ret = self._select_op(node.op, node.opset_id)(*inps, **opts)
       ret = ret if isinstance(ret, tuple) else (ret,)
@@ -543,6 +554,23 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     return __decorator
 
   # ***** Property/Graph Ops *****
+  def If(condition:Tensor, else_branch:OnnxRunner, then_branch:OnnxRunner, intermediate_tensors:dict[str, Tensor]):
+    def run_branch(branch:OnnxRunner):
+      branch.graph_values.update(intermediate_tensors)
+      out = branch({k:intermediate_tensors[k] for k in branch.graph_inputs.keys()})
+      # dereference intermediate tensors so Buffer can be deallocated
+      for k in intermediate_tensors: del branch.graph_values[k]
+      return out
+    # both branch must be ran before the condition can be evaluated
+    else_out, then_out = run_branch(else_branch), run_branch(then_branch)
+    assert len(else_out) == len(then_out), f"else_out and then_out must have the same number of outputs: {len(else_out)} != {len(then_out)}"
+    # can use where op when output shape is the same
+    if all(t.shape == e.shape for t,e in zip(then_out.values(), else_out.values())):
+      return tuple(condition.where(t,e) for t,e in zip(then_out.values(), else_out.values()))
+    # otherwise, use condition to select the output in python
+    cond = _resolve_const(_cached_to_python_const(condition))
+    return tuple(t if cond else e for t,e in zip(then_out.values(), else_out.values()))
+
   def Identity(x:Tensor): return x
   def Constant(sparse_value:Tensor|None=None, value:Tensor|None=None, value_float:float|None=None, value_floats:list[float]|None=None,
               value_int:int|None=None, value_ints:list[int]|None=None, value_string:str|None=None, value_strings:list[str]|None=None):

tinygrad/gradient.py

@@ -22,11 +22,10 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.SQRT, name="ret"), lambda ctx, ret: (ctx / (ret*2),)),
   (UPat((Ops.CMPLT, Ops.CMPNE)), lambda: (None, None)),
   (UPat(Ops.ADD), lambda ctx: (ctx, ctx)),
-  (UPat(Ops.POW, name="ret"), lambda ctx, ret:
-   (ctx*(ret.src[0].eq(0) & ret.src[1].eq(0)).where(ret.src[1], ret.src[1]*ret.src[0].pow(ret.src[1]-1)),
-    ctx*ret.src[0].eq(0).where((ret.src[1]<0).where(ret.const_like(-math.inf), ret.const_like(0)), ret*ret.src[0].log2()*math.log(2.0)))),
-  (UPat(Ops.MAX, name="ret"), lambda ctx, ret: ((ret.src[0]>ret.src[1]).where(ctx, (ret.src[0]!=ret.src[1]).where(ctx.const_like(0), ctx * 0.5)),
-                                                (ret.src[0]<ret.src[1]).where(ctx, (ret.src[0]!=ret.src[1]).where(ctx.const_like(0), ctx * 0.5)))),
+  (UPat(Ops.POW, name="ret", src=(UPat.var("b"), UPat.var("e"))), lambda ctx, ret, b, e:
+    (ctx * (b.eq(0)&e.eq(0)).where(e, e*b.pow(e-1)), ctx * b.eq(0).where((e<0).where(ret.const_like(-math.inf), 0), ret*b.log2()*math.log(2.0)))),
+  (UPat(Ops.MAX, name="ret", src=(UPat.var("x"), UPat.var("y"))), lambda ctx, ret, x, y:
+    ((x>y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)), (x<y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)))),
   (UPat(Ops.MUL, name="ret"), lambda ctx, ret: (ret.src[1]*ctx, ret.src[0]*ctx)),
   (UPat(Ops.WHERE, name="ret"), lambda ctx, ret: (None, ret.src[0].where(ctx, ctx.const_like(0)), ret.src[0].where(ctx.const_like(0), ctx))),
   (UPat(Ops.REDUCE_AXIS, name="ret"), reduce_gradient),

tinygrad/helpers.py

@@ -140,7 +140,7 @@ DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES = ContextVar("DONT_REALIZE_EXPAND", 0),
 QUANTIZE, VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("QUANTIZE", 0), ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, AMD_LLVM = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0), ContextVar("AMD_LLVM", 1)
-RANGEIFY = ContextVar("RANGEIFY", 0)
+RANGEIFY, POSTOPT, FUSE_ATTENTION = ContextVar("RANGEIFY", 0), ContextVar("POSTOPT", 0), ContextVar("FUSE_ATTENTION", 0)
 
 @dataclass(frozen=True)
 class Metadata:
@@ -196,7 +196,7 @@ class Profiling(contextlib.ContextDecorator):
 @dataclass(frozen=True)
 class TracingKey:
   display_name:str                       # display name of this trace event
-  keys:tuple[str, ...]=()                # optional keys to search for related traces
+  keys:tuple[Any, ...]=()                # optional keys to search for related traces
   cat:str|None=None                      # optional category to color this by
   ret:Any=None
 

tinygrad/nn/state.py

@@ -274,9 +274,9 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
   Converts ggml tensor data to a tinygrad tensor.
 
   Supported native types: float32 (id: 0), float16 (id: 1), int8 (id: 16), int16 (id: 17), int32 (id: 18)
-  Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q8_0 (id: 8), Q6_K (id: 14)
+  Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q8_0 (id: 8), Q6_K (id: 14), MXFP4 (id: 39)
   """
-  # https://github.com/ggerganov/ggml/blob/6dccc647264f5429df2624f36138f601e7ce23e5/include/ggml.h#L356
+  # https://github.com/ggerganov/ggml/blob/323951f1bdcdfbd5b5ff3a9a7c3770e63b1a560e/include/ggml.h#L356
 
   # native types
   if (dtype := { 0: dtypes.float32, 1: dtypes.float16, 16: dtypes.int8, 17: dtypes.int16, 18: dtypes.int32 }.get(ggml_type)) is not None:
@@ -288,7 +288,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
     return t.unsqueeze(-1).expand((*t.shape,8//b)).idiv(shift_tensor).bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
 
   # map to (number of elements, number of bytes)
-  if (nelements_nbytes := { 2: (32, 18), 3: (32, 20), 14: (256, 210), 8: (32, 34) }.get(ggml_type)) is not None:
+  if (nelements_nbytes := { 2: (32, 18), 3: (32, 20), 14: (256, 210), 8: (32, 34), 39: (32, 17) }.get(ggml_type)) is not None:
     blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1]))
     if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
     if ggml_type == 3:
@@ -300,6 +300,17 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
       scales = blocks[:,192:208].bitcast(dtypes.int8).unsqueeze(-1).expand((-1, 16, 16)).reshape((-1, 256))
       d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32).expand((-1, 256))
       return d * (xl.bitwise_or(xh).bitcast(dtypes.int8) - 32).flatten(-2) * scales
+    if ggml_type == 39:
+      e_int = blocks[:, 0].cast(dtypes.int32)
+      d = ((e_int >= 2).cast(dtypes.float32) * (e_int.cast(dtypes.float32) - 128).exp2() +
+           (e_int == 1).cast(dtypes.float32) * 2.0**(-127) +
+           (e_int == 0).cast(dtypes.float32) * 2.0**(-128)).unsqueeze(-1)
+      codes = q_to_uint8(blocks[:, 1:17], 4)
+      sign = 1.0 - codes.rshift(3).cast(dtypes.float32) * 2.0
+      exp, mant = codes.rshift(1).bitwise_and(0x3).cast(dtypes.float32), codes.bitwise_and(0x1).cast(dtypes.float32)
+      fp4_val = sign * ((exp != 0).cast(dtypes.float32) * (1.0 + 0.5 * mant) * (exp - 1.0).exp2() +
+                        (exp == 0).cast(dtypes.float32) * 0.5 * mant)
+      return (fp4_val * d).flatten(-2)[:n]
   raise ValueError(f"GGML type '{ggml_type}' is not supported!")
 
 @accept_filename

tinygrad/renderer/cstyle.py

@@ -157,7 +157,7 @@ class CStyleLanguage(Renderer):
       # naming
       prefix = None
       if u.op is Ops.SPECIAL: r[u] = u.arg[0]
-      elif u.op is Ops.RANGE: r[u] = f"ridx{u.arg[0]}" if u.arg[0] >= 0 else f"ridxm{-u.arg[0]}"
+      elif u.op is Ops.RANGE: r[u] = "ridx"+'_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])
       else:
         prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const",
                   Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.VECTORIZE: "cast", Ops.PRECAST: "precast",
@@ -316,7 +316,9 @@ class MetalRenderer(CStyleLanguage):
 
   def render_kernel(self, function_name, kernel, bufs, uops, prefix=None):
     prefix = ["#include <metal_stdlib>","using namespace metal;"]
-    for name, _, dtype_in, dtype_out, _, _, _, _ in wmma_args(uops): prefix.append(
+    wargs = wmma_args(uops)
+    if len(wargs) > 0: wargs = wargs[0:1]
+    for name, _, dtype_in, dtype_out, _, _, _, _ in wargs: prefix.append(
   f"""{(dstr_out:=self.render_dtype(dtype_out.vec(2)))} __{name}({(dstr_in:=self.render_dtype(dtype_in.vec(2)))} a, {dstr_in} b, {dstr_out} c){{
   simdgroup_{self.render_dtype(dtype_in)}8x8 mat_a, mat_b; simdgroup_{self.render_dtype(dtype_out)}8x8 mat_c;
   mat_a.thread_elements()[0] = a[0]; mat_b.thread_elements()[0] = b[0]; mat_c.thread_elements()[0] = c[0];

tinygrad/runtime/ops_amd.py

@@ -464,14 +464,14 @@ class AMDProgram(HCQProgram):
     # TODO; this API needs the type signature of the function and global_size/local_size
     self.dev, self.name, self.lib = dev, name, lib
 
-    image, sections, _ = elf_loader(self.lib)
+    image, sections, relocs = elf_loader(self.lib)
 
     rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
-    text_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".text"), -1)
-    assert rodata_entry >= 0 and text_entry >= 0, ".text or .rodata section not found"
+    assert rodata_entry >= 0, ".rodata section not found"
 
-    # Relo for kernel_code_entry_byte_offset for AMD_LLVM. Comgr doesn't need that, but keep shared code path.
-    image[rodata_entry+0x10:rodata_entry+0x10+8] = struct.pack('<q', text_entry - rodata_entry)
+    for apply_image_offset, rel_sym_offset, typ, addent in relocs:
+      if typ == 5: image[apply_image_offset:apply_image_offset+8] = struct.pack('<q', rel_sym_offset - apply_image_offset + addent) # R_AMDGPU_REL64
+      else: raise RuntimeError(f"unknown AMD reloc {typ}")
 
     self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000), buf_spec:=BufferSpec(cpu_access=True, nolru=True))
     self.dev.allocator._copyin(self.lib_gpu, image)

tinygrad/runtime/ops_null.py

@@ -7,6 +7,7 @@ class NullRenderer(CStyleLanguage):
   device = "NULL"
   has_local = False
   float4 = "float4"
+  barrier = "// BARRIER"
   code_for_op = {**CStyleLanguage.code_for_op, Ops.THREEFRY: lambda a,b,dtype: f"threefry({a},{b})", Ops.MAX: lambda a,b,dtype: f"max({a},{b})"}
 
 class NullProgram:

tinygrad/runtime/support/memory.py

@@ -77,11 +77,10 @@ class TLSFAllocator:
       if self.lv1_entries[l1] == 0: continue
       for l2 in range(self.lv2(size) if l1 == size.bit_length() else 0, (1 << self.l2_cnt)):
         if len(self.storage[l1][l2]) > 0:
-          nsize = self.blocks[self.storage[l1][l2][0]][0]
-          assert nsize >= size, "block must be larger"
-
           # Block start address.
           start = self.storage[l1][l2][0]
+          nsize = self.blocks[start][0]
+          assert nsize >= size, "block must be larger"
 
           # If request contains alignment, split the block into two parts.
           if (new_start:=round_up(start, align)) != start:

tinygrad/schedule/rangeify.py

@@ -2,11 +2,11 @@ from typing import Any
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, colored, RANGEIFY
+from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY
 from tinygrad.schedule.multi import multi_pm
 
 from tinygrad.schedule.kernelize import Kernel
-from tinygrad.uop.ops import track_rewrites, graph_rewrite_map, graph_rewrite, KernelInfo, identity_element, sint, AxisType
+from tinygrad.uop.ops import track_rewrites, graph_rewrite_map, graph_rewrite, identity_element, sint, AxisType
 
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
@@ -189,6 +189,7 @@ def map_partial_contiguous(ctx:RangeifyContext, x:UOp, idx:UOp):
     ranges.append(ctx.new_range(s) if resolve(s!=1) else UOp.const(dtypes.int, 0))
     new_ranges.append(ranges[-1])
   ret = x.src[0].index(*ranges).bufferize(*[x for x in new_ranges if x.op is not Ops.CONST], arg=x.device)
+  if len(ret.ranges): ret = ret.replace(arg=AddrSpace.LOCAL) # if some ranges are still open, this has to be LOCAL
   return ret.index(*passthrough_idx)
 
 def map_contiguous(ctx:RangeifyContext, x:UOp):
@@ -238,7 +239,10 @@ def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp):
   # index based on the shared ranges
   ret = c.index(*out_rngs)
   # if all ranges aren't the same between children, we have to bufferize
-  if len(idx_ranges) > 0: ret = ret.bufferize(*end_ranges, arg=x.device).index(*[idx.src[1+i] for i in idx_ranges])
+  if len(idx_ranges) > 0:
+    ret = ret.bufferize(*end_ranges, arg=x.device)
+    if len(ret.ranges): ret = ret.replace(arg=AddrSpace.LOCAL) # if some ranges are still open, this has to be LOCAL
+    ret = ret.index(*[idx.src[1+i] for i in idx_ranges])
   return ret
 
 def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp):
@@ -329,20 +333,28 @@ pm_cleanups = double_reshape+pm_mops+PatternMatcher([
 # BUFFERIZE returns the BUFFER ready for INDEXing (doing this will make splitting a lot easier)
 # NOTE: this has been fixed up a bit
 
-def bufferize_to_store(x:UOp):
+def bufferize_to_store(x:UOp, locals_allowed=False):
   rngs = x.src[1:]
   shape = tuple([int(r.vmax+1) for r in rngs])
   size = prod(shape)
   assert size > 0, f"no zero sized buffers {shape}"
-  sdtype = x.dtype.ptr(size=size, addrspace=AddrSpace.GLOBAL if not isinstance(x.arg, tuple) else x.arg[0])
+  sdtype = x.dtype.ptr(size=size, addrspace=AddrSpace.GLOBAL if not isinstance(x.arg, AddrSpace) else x.arg)
   if x.src[0].op is Ops.ASSIGN:
     assign_target, assign_src = x.src[0].src
     assert assign_target.op is Ops.INDEX
     return assign_target.replace(dtype=sdtype).store(assign_src, *rngs, dtype=sdtype)
-  if sdtype.addrspace == AddrSpace.GLOBAL: buf = UOp.new_buffer(x.arg, size, x.dtype)
-  else: buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=x.arg[1])
+  # NOTE: the DEFINE_LOCAL needs to be disambiguated here
+  if sdtype.addrspace == AddrSpace.GLOBAL:
+    buf = UOp.new_buffer(x.arg, size, x.dtype)
+  else:
+    if not locals_allowed: return None
+    buf = UOp(Ops.DEFINE_LOCAL, sdtype, arg=0) #UOp.unique().arg)
   return buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=sdtype).forced_reshape(shape, dtype=x.dtype)
 
+pm_add_buffers_local = pm_mops+PatternMatcher([
+  (UPat(Ops.BUFFERIZE, name="x"), lambda x: bufferize_to_store(x, True)),
+])
+
 pm_add_buffers = pm_mops+PatternMatcher([
   (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store),
 
@@ -407,12 +419,8 @@ def split_store(x:UOp):
   ctx = LocalAddBufferContext()
   ret = graph_rewrite(x, to_define_global+rangeify_codegen, ctx=ctx, name="kernel split", bottom_up=True)
 
-  # get name
-  rng = sorted([u for u in ret.toposort() if u.op is Ops.RANGE], key=lambda x: x.arg)
-  name = "k"+colored('_', 'BLACK').join(['']+[colored(s.src[0].render(), "WHITE" if s in ret.src[2:] else "red") for s in rng])
-
   # NOTE: the hack for COPY is here
-  ret = ret.sink(arg=KernelInfo(name=name)) if ret.src[1].op is not Ops.COPY else ret.src[1]
+  ret = ret.sink() if ret.src[1].op is not Ops.COPY else ret.src[1]
   kernel = UOp(Ops.KERNEL, src=tuple(ctx.map.values())+tuple(ctx.vars.keys()), arg=Kernel(ret,()))
   return x.as_buf().assign(kernel)
 

tinygrad/tensor.py

@@ -6,7 +6,7 @@ from typing import Callable, ClassVar, Sequence, cast, get_args, Literal, Suppor
 from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
 from tinygrad.dtype import _from_np_dtype, _to_np_dtype
 from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup
-from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray, RANGEIFY
+from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray, RANGEIFY, FUSE_ATTENTION
 from tinygrad.gradient import compute_gradient
 from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, Variable, MathTrait, identity_element, all_metadata
 from tinygrad.uop.spec import tensor_uop_spec, type_verify
@@ -3930,7 +3930,11 @@ class Tensor(MathTrait):
     if enable_gqa:
       key = key.repeat_interleave(self.shape[-3] // key.shape[-3], dim=-3)
       value = value.repeat_interleave(self.shape[-3] // value.shape[-3], dim=-3)
-    qk = self.matmul(key.transpose(-2,-1), dtype=least_upper_dtype(self.dtype, key.dtype, dtypes.float32)) / math.sqrt(self.shape[-1])
+
+    if FUSE_ATTENTION: q, key, value = self.contiguous(), key.contiguous(), value.contiguous()
+    else: q = self
+
+    qk = q.matmul(key.transpose(-2,-1), dtype=least_upper_dtype(q.dtype, key.dtype, dtypes.float32)) / math.sqrt(q.shape[-1])
     # handle attention mask
     if is_causal:
       if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
@@ -3938,7 +3942,8 @@ class Tensor(MathTrait):
     if attn_mask is not None:
       if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
       qk = qk + attn_mask
-    return qk.cast(self.dtype).softmax(-1).dropout(dropout_p) @ value
+    attn = qk.cast(self.dtype).softmax(-1).dropout(dropout_p) @ value
+    return attn.fuse() if FUSE_ATTENTION else attn
 
   def _do_reduction(self, reduction:ReductionStr="mean") -> Tensor:
     if reduction not in get_args(ReductionStr): raise ValueError(f"{reduction=} must be one of {get_args(ReductionStr)}")

tinygrad/uop/ops.py

@@ -142,6 +142,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if self.op is Ops.INDEX and self.src[0].op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG,
                                                    Ops.BUFFER, Ops.BUFFERIZE, Ops.VECTORIZE, Ops.STORE}:
       return None
+    if self.op is Ops.BARRIER: return None
     if self.op in GroupOp.Block: return None
     from tinygrad.shape.shapetracker import ShapeTracker
     # VIEW and MovementOps define a new ShapeTracker from the arg
@@ -202,17 +203,15 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   @functools.cached_property
   def ranges(self) -> dict[UOp, None]:
     if self.op is Ops.RANGE: return {self:None}
-    if self.op in {Ops.BUFFERIZE, Ops.REDUCE}:
-      ret = self.src[0].ranges.copy()
-      for s in self.src[1:]:
-        if s in ret: del ret[s]
-    elif self.op in {Ops.STORE}:
-      ret = self.src[0].ranges.copy()
-      ret.update(self.src[1].ranges)
-      for s in self.src[2:]:
-        if s in ret: del ret[s]
+    range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3}
+    ret: dict[UOp, None] = {}
+    if self.op in range_start.keys():
+      for s in self.src[:range_start[self.op]]: ret.update(s.ranges)
+      delete_ranges = self.src[range_start[self.op]:]
+      if len(delete_ranges):
+        for s in UOp.sink(*delete_ranges).ranges:
+          if s in ret: del ret[s]
     else:
-      ret = {}
       for s in self.src: ret.update(s.ranges)
     return ret
 
@@ -251,7 +250,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     ret = self.arg[1] if self.op is Ops.REDUCE_AXIS else self.arg[7]
     assert isinstance(ret, tuple) and all(isinstance(x, int) for x in ret), f"axis_arg trying to return {ret}"
     return ret
-  def sink(self, *srcs:UOp|None, **kwargs): return UOp(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
+  def sink(*srcs:UOp|None, **kwargs):  # pylint: disable=no-self-argument
+    return UOp(Ops.SINK, dtypes.void, tuple([x for x in srcs if x is not None]), **kwargs)
   def detach(self): return UOp(Ops.DETACH, self.dtype, (self,))
   def index(self, *srcs:UOp|None, **kwargs):
     return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
@@ -299,8 +299,10 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       else: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),))
     return ret
   @staticmethod
-  def range(dtype:DType, end:sint, idx:int, axistype:AxisType=AxisType.LOOP):
-    return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end),), arg=(idx, axistype))
+  def range(dtype:DType, end:sint, *arg):
+    if len(arg) == 0: raise RuntimeError("range needs an arg")
+    if len(arg) == 1: arg = arg+(AxisType.LOOP,)
+    return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end),), arg=arg)
   def r(self, op:Ops, axis:tuple[int, ...]):
     axis = tuple(sorted([x for x in axis if resolve(self.shape[x] != 1)]))
     if len(axis) == 0: return self

tinygrad/uop/spec.py

@@ -17,33 +17,39 @@ try:
     return s
 
   # ctx is (solver, load_number_dict)
+  # each uop gets rewritten to NOOP(arg=(solver, z3_object)), the arg has the solver first due to UOpMetaClass caching. z3 objects from different
+  # contexts can have the same hash but error on comparison
   z3_renderer = PatternMatcher([
     # Ops.SPECIAL can have symbolic arg but it wont be in the toposort beacuse its not a src, we need to add it manually
     (UPat(Ops.SPECIAL, src=(), name="x"), lambda x: UOp(Ops.SPECIAL, arg=x.arg[0], src=(x.ufix(x.arg[1]),))),
-    (UPat(Ops.SPECIAL, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=create_bounded(x.arg, 0, x.src[0].arg-1, ctx[0]))),
-    (UPat(Ops.DEFINE_VAR, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=create_bounded(x.arg[0], x.arg[1], x.arg[2], ctx[0]))),
-    (UPat(Ops.RANGE, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=create_bounded(f"ridx{x.arg}", 0, x.src[0].arg-1, ctx[0]))),
+    (UPat(Ops.SPECIAL, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(x.arg, 0, x.src[0].arg[1]-1, ctx[0])))),
+    (UPat(Ops.DEFINE_VAR, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(x.arg[0], x.arg[1], x.arg[2], ctx[0])))),
+    (UPat(Ops.RANGE, name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(f"ridx{x.arg}", 0, x.src[0].arg[1]-1, ctx[0])))),
     # float loads only become a variable when they get cast to int/bool
     (UPat(Ops.LOAD, dtypes.ints, name="x"),
-      lambda x,ctx: UOp(Ops.NOOP, arg=create_bounded(f"load{ctx[1].setdefault(x, len(ctx[1]))}", x.dtype.min, x.dtype.max, ctx[0]))),
+      lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],create_bounded(f"load{ctx[1].setdefault(x, len(ctx[1]))}", x.dtype.min, x.dtype.max, ctx[0])))),
     (UPat(Ops.CONST, dtype=dtypes.ints+(dtypes.bool,), name="x"),
-      lambda x,ctx: UOp(Ops.NOOP, arg=(z3.BoolVal if dtypes.is_bool(x.dtype) else z3.IntVal)(x.arg, ctx=ctx[0].ctx))),
+      lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0],(z3.BoolVal if dtypes.is_bool(x.dtype) else z3.IntVal)(x.arg, ctx=ctx[0].ctx)))),
     # z3 can cast from bool to int automatically
     (UPat(Ops.CAST, dtype=dtypes.ints, src=UPat(Ops.NOOP), name="x"), lambda x: x.src[0]),
-    (UPat(Ops.CAST, dtype=dtypes.bool, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=(x.src[0].arg!=0))),
+    (UPat(Ops.CAST, dtype=dtypes.bool, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], x.src[0].arg[1]!=0))),
     # if the source of the cast is not a noop it means that it is a float and so we create a new variable
     (UPat(Ops.CAST, dtype=dtypes.ints, name="x"), lambda x,ctx:
-      UOp(Ops.NOOP, arg=create_bounded(f"cast{ctx[1].setdefault(x, len(ctx[1]))}", x.dtype.min, x.dtype.max, ctx[0]))),
+      UOp(Ops.NOOP, arg=(ctx[0], create_bounded(f"cast{ctx[1].setdefault(x, len(ctx[1]))}", x.dtype.min, x.dtype.max, ctx[0])))),
     (UPat(Ops.CAST, dtype=dtypes.bool, name="x"), lambda x,ctx:
-      UOp(Ops.NOOP, arg=z3.Bool(f"cast{ctx[1].setdefault(x, len(ctx[1]))}",ctx=ctx[0].ctx))),
+      UOp(Ops.NOOP, arg=(ctx[0], z3.Bool(f"cast{ctx[1].setdefault(x, len(ctx[1]))}",ctx=ctx[0].ctx)))),
     (UPat(Ops.XOR, src=UPat(Ops.NOOP), name="x"),
-      lambda x: UOp(Ops.NOOP, arg=z3.BV2Int(z3_alu[x.op](*(z3.Int2BV(s.arg, x.dtype.itemsize*8) for s in x.src))))),
-    (UPat(GroupOp.ALU, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=z3_alu[x.op](*(s.arg for s in x.src)))),
+      lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], z3.BV2Int(z3_alu[x.op](*(z3.Int2BV(s.arg[1], x.dtype.itemsize*8) for s in x.src)))))),
+    (UPat(GroupOp.ALU, src=UPat(Ops.NOOP), name="x"), lambda x,ctx: UOp(Ops.NOOP, arg=(ctx[0], z3_alu[x.op](*(s.arg[1] for s in x.src))))),
     # A comparison between floats introduces a new bool variable
     (UPat(GroupOp.Comparison, src=UPat(dtype=dtypes.floats), name="x"), lambda x,ctx:
-      UOp(Ops.NOOP, arg=z3.Bool(f"float_cmp{ctx[1].setdefault(x, len(ctx[1]))}",ctx=ctx[0].ctx))),
+      UOp(Ops.NOOP, arg=(ctx[0], z3.Bool(f"float_cmp{ctx[1].setdefault(x, len(ctx[1]))}",ctx=ctx[0].ctx)))),
   ])
 
+  def uops_to_z3(solver, *uops: UOp) -> 'list[z3.ExprRef]':
+    with Context(TRACK_MATCH_STATS=0):  # cant pickle z3 objects
+      return [s.arg[1] for s in graph_rewrite(uops[0].sink(*uops[1:]), z3_renderer, ctx=(solver, {})).src]
+
   z3_imported = True
 except (ImportError, AttributeError): z3_imported = False
 
@@ -124,9 +130,8 @@ def validate_index(idx:UOp, gate:UOp=UOp.const(dtypes.bool, True)):
 
   if not z3_imported: raise ImportError("z3 is required for bounds checking, try IGNORE_OOB=0 or \"pip install z3-solver\"")
   solver = z3.Solver(ctx=z3.Context())
-  z3_sink = graph_rewrite(idx.src[1].sink(mask), z3_renderer, ctx=(solver, {}))
-  z3_idx = z3_sink.src[0].arg
-  solver.add(z3_sink.src[1].arg)
+  z3_idx, z3_mask = uops_to_z3(solver, idx.src[1], mask)
+  solver.add(z3_mask)
   if solver.check((z3_idx<0)|(sz<=z3_idx)) == z3.sat:
     print(f"idx={idx.src[1].render(simplify=False)}")
     print(f"mask & gate={mask.render(simplify=False)}")

tinygrad/viz/js/index.js

@@ -4,6 +4,15 @@ const displayGraph = (cls) => {
   for (const e of document.getElementsByClassName("view")) e.style.display = e.classList.contains(cls) ? "flex" : "none";
 }
 
+const darkenHex = (h, p = 0) =>
+  `#${(
+    c = parseInt(h.slice(1), 16),
+    f = 1 - p / 100,
+    ((c >> 16 & 255) * f | 0) << 16 |
+    ((c >>  8 & 255) * f | 0) <<  8 |
+    ((c       & 255) * f | 0)
+  ).toString(16).padStart(6, '0')}`;
+
 const ANSI_COLORS = ["#b3b3b3", "#ff6666", "#66b366", "#ffff66", "#6666ff", "#ff66ff", "#66ffff", "#ffffff"];
 const parseColors = (name, defaultColor="#ffffff") => Array.from(name.matchAll(/(?:\u001b\[(\d+)m([\s\S]*?)\u001b\[0m)|([^\u001b]+)/g),
   ([_, code, colored_st, st]) => ({ st: colored_st ?? st, color: code != null ? ANSI_COLORS[(parseInt(code)-30+60)%60] : defaultColor }));
@@ -87,7 +96,7 @@ async function renderDag(graph, additions, recenter=false) {
       }
       return [ret];
     }).join("text").selectAll("tspan").data(d => d).join("tspan").attr("x", "0").attr("dy", 14).selectAll("tspan").data(d => d).join("tspan")
-      .attr("fill", d => d.color).text(d => d.st).attr("xml:space", "preserve");
+      .attr("fill", d => darkenHex(d.color, 25)).text(d => d.st).attr("xml:space", "preserve");
     addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
       .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => e.tag));
     // draw edges
@@ -381,8 +390,7 @@ async function renderProfiler() {
     d3.select(canvas).call(canvasZoom.transform, zoomLevel);
   }
 
-  canvasZoom = d3.zoom().filter(e => (!e.ctrlKey || e.type === 'wheel' || e.type === 'mousedown') && !e.button)
-    .scaleExtent([1, Infinity]).translateExtent([[0,0], [Infinity,0]]).on("zoom", e => render(e.transform));
+  canvasZoom = d3.zoom().filter(vizZoomFilter).scaleExtent([1, Infinity]).translateExtent([[0,0], [Infinity,0]]).on("zoom", e => render(e.transform));
   d3.select(canvas).call(canvasZoom);
   document.addEventListener("contextmenu", e => e.ctrlKey && e.preventDefault());
 
@@ -419,7 +427,8 @@ async function renderProfiler() {
 
 // ** zoom and recentering
 
-const svgZoom = d3.zoom().on("zoom", (e) => d3.select("#render").attr("transform", e.transform));
+const vizZoomFilter = e => (!e.ctrlKey || e.type === 'wheel' || e.type === 'mousedown') && !e.button && e.type !== 'dblclick';
+const svgZoom = d3.zoom().filter(vizZoomFilter).on("zoom", (e) => d3.select("#render").attr("transform", e.transform));
 d3.select("#graph-svg").call(svgZoom);
 
 // zoom to fit into view

tinygrad/viz/serve.py

@@ -11,6 +11,7 @@ from tinygrad.uop.ops import TrackedGraphRewrite, UOp, Ops, printable, GroupOp,
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry, Device
 from tinygrad.renderer import ProgramSpec
 from tinygrad.dtype import dtypes
+from tinygrad.codegen.opt.kernel import axis_colors
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
                Ops.DEFINE_GLOBAL: "#ffe0b0", Ops.DEFINE_LOCAL: "#ffe0d0", Ops.DEFINE_REG: "#f0ffe0", Ops.REDUCE_AXIS: "#FF6B6B",
@@ -79,7 +80,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
       if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None:
         label += f"\n{shape_to_str(u.shape)}"
       elif len(rngs:=u.ranges):
-        label += f"\n{str(sorted([x.arg[0] for x in rngs]))}"
+        label += f"\n({','.join([colored(str(x.arg[0]), axis_colors[x.arg[-1]]) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})"
     except Exception:
       label += "\n<ISSUE GETTING LABEL>"
     if (ref:=ref_map.get(u.arg.ast) if u.op is Ops.KERNEL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}"
@@ -257,7 +258,7 @@ class Handler(BaseHTTPRequestHandler):
       if url.path == "/disasm": ret, content_type = get_disassembly(**query), "application/json"
       else: return self.stream_json(get_details(contexts[1][int(query["ctx"][0])][int(query["idx"][0])]))
     elif url.path == "/ctxs": ret, content_type = json.dumps(ctxs).encode(), "application/json"
-    elif url.path == "/get_profile" and profile_ret is not None: ret, content_type = profile_ret, "application/octet-stream"
+    elif url.path == "/get_profile" and profile_ret: ret, content_type = profile_ret, "application/octet-stream"
     else: status_code = 404
 
     # send response
@@ -290,8 +291,8 @@ def reloader():
       os.execv(sys.executable, [sys.executable] + sys.argv)
     time.sleep(0.1)
 
-def load_pickle(path:str):
-  if path is None or not os.path.exists(path): return None
+def load_pickle(path:str|None) -> list:
+  if path is None or not os.path.exists(path): return []
   with open(path, "rb") as f: return pickle.load(f)
 
 # NOTE: using HTTPServer forces a potentially slow socket.getfqdn
@@ -314,16 +315,16 @@ if __name__ == "__main__":
   contexts, profile = load_pickle(args.kernels), load_pickle(args.profile)
 
   # NOTE: this context is a tuple of list[keys] and list[values]
-  ctxs = get_metadata(*contexts[:2]) if contexts is not None else []
+  ctxs = get_metadata(*contexts[:2]) if contexts else []
 
-  profile_ret = get_profile(profile) if profile is not None else None
+  profile_ret = get_profile(profile)
 
   server = TCPServerWithReuse(('', PORT), Handler)
   reloader_thread = threading.Thread(target=reloader)
   reloader_thread.start()
   print(f"*** started viz on {HOST}:{PORT}")
   print(colored(f"*** ready in {(time.perf_counter()-st)*1e3:4.2f}ms", "green"), flush=True)
-  if len(getenv("BROWSER", "")) > 0: webbrowser.open(f"{HOST}:{PORT}{'/profiler' if contexts is None else ''}")
+  if len(getenv("BROWSER", "")) > 0: webbrowser.open(f"{HOST}:{PORT}")
   try: server.serve_forever()
   except KeyboardInterrupt:
     print("*** viz is shutting down...")