changes

* no replacement multinomial

Efraimidis–Spirakis

* num_samples == 1 can use fast path

.github/workflows/autogen.yml

@@ -48,7 +48,7 @@ jobs:
         python3 -c "from tinygrad.runtime.autogen import opencl"
         python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
         python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
-        python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v13_0_12, smu_v14_0_2"
+        python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v13_0_12, smu_v14_0_2, fw"
         python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
         python3 -c "from tinygrad.runtime.autogen import llvm"
         python3 -c "from tinygrad.runtime.autogen import webgpu"

test/amd/test_custom_kernel.py

@@ -3,7 +3,7 @@ import functools
 import numpy as np
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.engine.realize import run_linear, estimate_uop
+from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer import Estimates
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv
@@ -169,7 +169,7 @@ class TestCustomKernel(unittest.TestCase):
     if self.arch != "rdna3": self.skipTest("only rdna3")
     a = Tensor.full((16, 16), 1.).contiguous().realize()
     a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
-    linear = a.schedule_linear()
+    linear = compile_linear(a.schedule_linear())
     est = estimate_uop(linear.src[-1])
     self.assertEqual(est.ops, a.numel())
     self.assertEqual(est.mem, a.nbytes()*2)

test/backend/test_arange.py

@@ -2,14 +2,14 @@ import unittest
 import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
-from tinygrad.engine.realize import run_linear, estimate_uop
+from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer.ptx import PTXRenderer
 from test.helpers import needs_second_gpu
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, tensor, desired):
     GlobalCounters.reset()
-    linear = tensor.schedule_linear()
+    linear = compile_linear(tensor.schedule_linear())
     self.assertEqual(len(linear.src), 1)
     run_linear(linear)
     np.testing.assert_equal(tensor.numpy(), desired)
@@ -36,7 +36,7 @@ class TestArange(unittest.TestCase):
   def test_tri_complexity(self):
     with Context(NOOPT=1):
       t = Tensor.ones(256, 256).contiguous().realize()
-      linear = t.triu().schedule_linear()
+      linear = compile_linear(t.triu().schedule_linear())
       self.assertLessEqual(estimate_uop(linear.src[-1]).ops, 4 * 256 * 256)
 
 DSET, DDIM = 2048, 32
@@ -229,7 +229,7 @@ class TestIndexing(unittest.TestCase):
     xq = xq.reshape(bs, seqlen, n_heads, head_dim)
     xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis)
     xq_rope.sum().backward()
-    linear = wq.grad.schedule_linear()
+    linear = compile_linear(wq.grad.schedule_linear())
     assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}"
     bwd_ops = estimate_uop(linear.src[0]).ops
     # bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation

test/backend/test_custom_kernel.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor, UOp
+from tinygrad import Tensor, UOp, GlobalCounters
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.uop.ops import KernelInfo, AxisType
 
@@ -308,6 +308,22 @@ class TestCustomKernel(unittest.TestCase):
     expected = (3+2)*2+2
     assert all(x == expected for x in result), f"expected all {expected}, got {result}"
 
+  def test_custom_kernel_sched(self, use_custom=False):
+    x = Tensor.arange(32).reshape(8, 4).realize()
+    y = Tensor.empty_like(x)
+    y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
+    if use_custom:
+      z = Tensor.empty_like(x)
+      z = Tensor.custom_kernel(y, y.T.T, fxn=custom_add_one_kernel)[0]
+    else: z = y.T.T+1
+    GlobalCounters.reset()
+    z.realize()
+    self.assertEqual(GlobalCounters.kernel_count, 2)
+    self.assertEqual(z.tolist(), x.add(2).tolist())
+
+  @unittest.expectedFailure
+  def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custom=True)
+
 class TestUOpReduce(unittest.TestCase):
   def test_uop_sum(self):
     a = Tensor([1.0, 2, 3, 4, 5])

test/backend/test_randomness.py

@@ -307,17 +307,26 @@ class TestRandomness(unittest.TestCase):
     with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3.5)
     with self.assertRaises(TypeError): Tensor.randint((3, 4), low=1, high=3, dtype="float")
     with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3, dtype=dtypes.float32)
+    # check low < high
+    with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=5)
+    with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=10)
+    np.testing.assert_array_equal(Tensor.randint(16, low=5, high=6).numpy(), 5)
 
   def test_normal(self):
     self.assertTrue(normal_test(Tensor.normal))
     self.assertTrue(equal_distribution(Tensor.normal, lambda x: torch.nn.init.normal_(torch.empty(x), mean=0, std=1),
                                                       lambda x: np.random.normal(loc=0, scale=1, size=x)))
+    # check std >= 0
+    with self.assertRaises(ValueError): Tensor.normal((3, 4), mean=0, std=-1)
 
   def test_uniform(self):
     self.assertFalse(normal_test(Tensor.uniform))
     self.assertTrue(equal_distribution(Tensor.uniform, lambda x: torch.nn.init.uniform_(torch.empty(x)), lambda x: np.random.uniform(size=x)))
     self.assertTrue(equal_distribution(partial(Tensor.uniform, low=-100, high=100, dtype=dtypes.int32),
                                        numpy_func=lambda x: np.random.randint(low=-100, high=100, size=x)))
+    # check low < high
+    with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=5.0, high=3.0)
+    with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=1.0, high=1.0)
 
   def test_scaled_uniform(self):
     self.assertFalse(normal_test(Tensor.scaled_uniform))
@@ -352,7 +361,7 @@ class TestRandomness(unittest.TestCase):
     _check_with_torch(w=[0.231, 0., 1., 0.5], num_samples=300, replacement=True)
     _check_with_torch(w=[[0.2, 0.8]], num_samples=300, replacement=True)  # 2D but only 1 row
     _check_with_torch(w=[[0.453, 0., 1., 0.81], [0.1, 0.8, 0., 0.1]], num_samples=300, replacement=True)
-    # no-replacement isn't supported, unless taking only one sample
+    # no-replacement
     w = [0.1, 0.9]
     self.assertRaises(AssertionError, lambda: Tensor(w).multinomial(100, replacement=False))
 
@@ -363,6 +372,23 @@ class TestRandomness(unittest.TestCase):
     torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(1000)]
     self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples)))
 
+    w = list(range(32))
+    s1 = Tensor(w).multinomial(5, replacement=False).numpy()
+    self.assertEqual(len(set(s1.tolist())), 5)
+    s2 = Tensor(w).multinomial(5, replacement=False).numpy()
+    self.assertFalse(np.array_equal(s1, s2))
+    full = Tensor(w).multinomial(len(w), replacement=False).numpy()
+    self.assertEqual(sorted(full.tolist()), w)
+
+    w = [0.1, 0.2, 0.3, 0.4]
+    @TinyJit
+    def sample_three(): return Tensor(w).multinomial(3, replacement=False).realize()
+
+    tiny_draws = np.array([sample_three().numpy() for _ in range(1000)])
+    torch_draws = np.array([torch.tensor(w).multinomial(3, replacement=False).numpy() for _ in range(1000)])
+    for pos in range(3):
+      self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_draws[:, pos]), lambda _: torch.tensor(torch_draws[:, pos])))
+
   @unittest.skip("this test is flaky")
   def test_multinomial_counterexample(self):
     tiny_res = Tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)

test/null/test_simplify_valid_idx.py

@@ -15,7 +15,7 @@ def get_gated_load_uop(valid:UOp, idx:UOp):
 
 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
   return UOp(Ops.LOAD, dtypes.float.vec(4), (
-    UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.STACK, dtypes.weakint.vec(2), idx).valid(valid), ptr=True),
+    UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(idx[0], idx[1], valid, ptr=True),
     UOp(Ops.STACK, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
   ))
 
@@ -218,16 +218,13 @@ class TestImageSimplification(unittest.TestCase):
   def check(self, load, svalid, sidx0, sidx1):
     with Context(NOOPT=1, SPEC=0):
       load = full_rewrite_to_sink(load.sink()).src[0]
-    idx = load.src[0].src[1]
-    self.assertEqual(idx.op, Ops.STACK)
-    self.assertEqual(len(idx.src), 2)
-    idx0, idx1 = idx.src[0], idx.src[1]
+    idx0, idx1 = load.src[0].src[1], load.src[0].src[2]
     check_uop_against_string(self, idx0, sidx0)
     check_uop_against_string(self, idx1, sidx1)
     if svalid is not None:
-      check_uop_against_string(self, load.src[0].src[2], svalid)
+      check_uop_against_string(self, load.src[0].src[3], svalid)
     else:
-      self.assertEqual(len(load.src[0].src), 2, "svalid is None but load still has a valid")
+      self.assertEqual(len(load.src[0].src), 3, "svalid is None but load still has a valid")
 
   def test_idx_gt_c(self):
     # (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid
@@ -353,7 +350,7 @@ class TestImageSimplification(unittest.TestCase):
     load = get_load_image_uop(shape, valid, idx)
 
     self.check(load,
-               "((((idx2*2)+r0)<11)&((((idx1*8)+r1)<3)!=True))",
+               "(((idx2*2)+r0)<11)",
                "(idx0+(idx1*512+r1*64)+-192)",
                "((((idx2*2)+r0)+(((idx1+((r1+5)//8))+1)//2))+-4)")
 
@@ -481,7 +478,7 @@ class TestImageSimplification(unittest.TestCase):
       self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0")
     except AssertionError:
       # TODO: fold valid
-      self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<19))",
+      self.check(load, "(((lidx1<1)!=True)&((lidx0+r0)<19))",
                        "(gidx0+lidx1*128+(lidx0*1024+r0*1024)+-3168)", "0")
 
   def test_simplify10(self):
@@ -500,7 +497,7 @@ class TestImageSimplification(unittest.TestCase):
       self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0")
     except AssertionError:
       # TODO: fold valid
-      self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<11))",
+      self.check(load, "(((lidx1<1)!=True)&((lidx0+r0)<11))",
                        "(lidx2+gidx0*4+lidx1*256+(lidx0*1024+r0*1024)+-3264)", "0")
 
 class TestUnfoldableImage(unittest.TestCase):

test/null/test_uop_vmin_vmax.py

@@ -297,7 +297,7 @@ class TestVminVmaxVConst(unittest.TestCase):
     # vmin and vmax for a vector constant of bool values
     d1 = UOp(Ops.PARAM, dtypes.int.ptr(), (), 1)
     idx = UOp.const(dtypes.int, 0)
-    val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx),))
+    val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx).cast(dtypes.int.vec(2).ptr()),))
     uop = (val // 32).gep(0)
     self.assertEqual(uop.vmin, -67108864)
     self.assertEqual(uop.vmax, 67108863)

tinygrad/codegen/__init__.py

@@ -109,6 +109,8 @@ pm_linearize_cleanups = PatternMatcher([
   (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError, "if not allowed in graph")),
   # gated INDEX becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
   (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat())),
+   lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))])),
+  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat())),
    lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))]))
 ])
 
@@ -180,8 +182,6 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
 
 to_program_cache: dict[tuple, UOp] = {}
 def to_program(ast:UOp, renderer:Renderer) -> UOp:
-  if ast.op is Ops.PROGRAM and len(ast.src) >= 5 and ast.src[4].op is Ops.BINARY:
-    return ast if isinstance(ast.arg, ProgramInfo) else ast.replace(arg=ProgramInfo.from_sink(ast.src[0]))
   config = (NOOPT, DEVECTORIZE, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32)
   key = (ast.key, type(renderer), renderer.target, *[x.value for x in config])
   if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)

tinygrad/codegen/late/devectorizer.py

@@ -47,7 +47,10 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
 
   if not drop_stmt and idx is start_idx: return None
   new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
-  return buf.index(idx.valid(new_valid) if new_valid is not None else idx, ptr=True)
+  return buf.index(idx.gep(0), idx.gep(1), new_valid, ptr=True) if new_valid is not None else buf.index(idx.gep(0), idx.gep(1), ptr=True)
+
+def simplify_image_valid(buf:UOp, idx_x:UOp, idx_y:UOp, c:UOp) -> UOp|None:
+  return simplify_valid_load(buf, UOp(Ops.STACK, idx_x.dtype.vec(2), (idx_x, idx_y)), c) if isinstance(buf.dtype, ImageDType) else None
 
 
 load_store_indexing = PatternMatcher([
@@ -55,8 +58,11 @@ load_store_indexing = PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
   # simplify away long after index has been lowered
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("x", dtypes.long), UPat.var("c", dtypes.bool))), lambda buf,x,c: simplify_valid_load(buf, x, c)),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx_x"), UPat.var("idx_y"), UPat.var("c", dtypes.bool))), simplify_image_valid),
   # drop true gate
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("x"), UPat.const(dtypes.bool, True)),), lambda buf,x: buf.index(x, ptr=True)),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx_x"), UPat.var("idx_y"), UPat.const(dtypes.bool, True)),),
+   lambda buf,idx_x,idx_y: buf.index(idx_x, idx_y, ptr=True) if isinstance(buf.dtype, ImageDType) else None),
 ])
 
 # ***** load/store grouping *****
@@ -197,8 +203,9 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
   return UOp(Ops.VCAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
 
 def get_image_idx(idx:UOp, width:int):
-  oidx = UOp(Ops.STACK, dtypes.weakint.vec(2), (((x:=idx.src[1].get_idx()) // 4) % width, (x // (4*width))))
-  return idx.replace(src=(idx.src[0], oidx.valid(idx.src[1].get_valid())))
+  x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
+  idx_x, idx_y = (x // 4) % width, x // (4*width)
+  return idx.replace(src=(idx.src[0], idx_x, idx_y) + (() if valid.op is Ops.CONST and valid.arg is True else (valid,)))
 
 def image_fixup(ls:UOp):
   # normal image load or store, with the CAST from expand_index
@@ -207,7 +214,7 @@ def image_fixup(ls:UOp):
     return ls.replace(src=(get_image_idx(ls.src[0].src[0], dt.shape[1]),)+ls.src[1:])
 
   # this is an unprocessed image without a cast, we should just make it a buffer
-  if isinstance(dt, ImageDType) and (off:=ls.src[0].src[1]).get_idx().dtype != dtypes.weakint.vec(2):
+  if isinstance(dt, ImageDType) and len(ls.src[0].src) == 2 and (off:=ls.src[0].src[1]).get_idx().dtype.count != 2:
     idx = ls.src[0].src[0].replace(dtype=(new_dt:=dtypes.half if dt.itemsize == 2 else dtypes.float).ptr(dt.size)).index(off)
     return ls.replace(src=(idx,), dtype=new_dt).cast(dtypes.float) if ls.op is Ops.LOAD else ls.replace(src=(idx, ls.src[1].cast(new_dt)))
 
@@ -234,13 +241,26 @@ def no_vectorized_wmma(wmma:UOp):
 
 def no_vectorized_alu(alu:UOp):
   if alu.dtype.vcount == 1: return None
-  if alu.op is Ops.WHERE and alu.src[2].arg is Invalid: return None  # image load/store has cond.where(idx.vec(2), Invalid) as the index
+  if alu.op is Ops.WHERE and alu.src[2].arg is Invalid: return None  # masked indexes use cond.where(idx, Invalid)
   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.STACK, alu.dtype, alus)
 
 def no_vectorized_buf(buf:UOp):
   return buf.replace(dtype=buf.ptrdtype.base.scalar().ptr(buf.ptrdtype.size*buf.ptrdtype.count, buf.ptrdtype.addrspace)).cast(buf.dtype)
 
+def drop_load_alt(x:UOp):
+  idx = x.src[0].src[0] if x.src[0].op is Ops.CAST else x.src[0]
+  if idx.op is not Ops.INDEX: return None
+  if len(idx.src) == 2 or (isinstance(idx.src[0].dtype, ImageDType) and len(idx.src) == 3): return x.replace(src=(x.src[0],)+x.src[2:])
+  return None
+
+def add_load_alt(x:UOp):
+  idx = x.src[0].src[0] if x.src[0].op is Ops.CAST else x.src[0]
+  if idx.op is not Ops.INDEX: return None
+  if not (len(idx.src) == 3 and idx.src[2].dtype == dtypes.bool) and \
+     not (isinstance(idx.src[0].dtype, ImageDType) and len(idx.src) == 4 and idx.src[3].dtype == dtypes.bool): return None
+  return x.replace(src=(x.src[0], x.const_like(0))+x.src[1:]) if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None
+
 def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp, bcast:UOp|None=None):
   cnt = cast.dtype.count
   if bcast is not None and bcast.op is Ops.GEP:
@@ -282,8 +302,8 @@ pm_render = PatternMatcher([
   (UPat(Ops.STACK, src=(UPat(name='x'),)), lambda x: x),
   # give any loads that are masked an alt value
   (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"),
-    lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:])
-      if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None),
+    add_load_alt),
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX).or_casted(), UPat()), allow_any_len=True, name="x"), drop_load_alt),
   # Where after gated load becomes alt value
   # NOTE: if a is CAST and a.src[0].dtype == l.dtype, use a.src[0] to avoid roundtrip cast (e.g. uint->float->uint)
   (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat(), UPat.var("c")).or_casted(),), allow_any_len=True, name="l").or_casted(),
@@ -358,10 +378,14 @@ pm_reduce = PatternMatcher([
 
 # add loads
 
+def add_load(idx:UOp):
+  if isinstance(idx.dtype, PtrDType): return None
+  assert isinstance(idx.src[0].dtype, PtrDType), f"param is not PtrDType {idx.src[0].dtype}"
+  return idx.replace(dtype=idx.src[0].dtype).load(dtype=idx.dtype.base)
+
 pm_add_loads = PatternMatcher([
   # add loads to non ptr index
-  (UPat(Ops.INDEX, name="idx"), lambda idx: None if isinstance(idx.dtype, PtrDType) else
-    idx.replace(dtype=idx.src[0].dtype).load(dtype=idx.dtype.base)),
+  (UPat(Ops.INDEX, name="idx"), add_load),
   # remove loads from stores
   (UPat(Ops.STORE, src=(UPat(Ops.LOAD), UPat(name="val")), name="s"), lambda s,val: s.replace(src=(s.src[0].src[0], val))),
 ])

tinygrad/engine/jit.py

@@ -1,12 +1,12 @@
 from typing import TypeVar, Generic, Callable, Any
 import functools, collections
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, colored, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ
+from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ
 from tinygrad.device import Buffer, Compiled, Device, MultiBuffer
 from tinygrad.dtype import DType, dtypes
 from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite
-from tinygrad.engine.realize import capturing, Runner, Estimates, compile_linear, run_linear, graph_cache, estimate_uop, get_runtime
-from tinygrad.engine.realize import unwrap_multi, resolve_params
+from tinygrad.engine.realize import capturing, Estimates, compile_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
 from tinygrad.schedule.memory import memory_plan_rewrite, _collect_bufs
 from tinygrad.nn.state import get_parameters
 from tinygrad.schedule.rangeify import mop_cleanup
@@ -59,14 +59,6 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
   if current_batch: flush_batch()
   return linear.replace(src=tuple(new_src))
 
-def _call_outs_ins(call:UOp) -> tuple[set[int], set[int]]:
-  non_bind = [s for s in call.src[1:] if s.op is not Ops.BIND]
-  ast = call.src[0]
-  if ast.op is Ops.PROGRAM: return set(ast.arg.outs), set(ast.arg.ins)
-  if ast.op in (Ops.COPY, Ops.BUFFER_VIEW): return {0}, {1}
-  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return {0}, set(range(1, len(non_bind)))
-  return set(), set()
-
 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.arg, u.dtype), u, metadata=()),)))
   return new
@@ -95,14 +87,14 @@ def _check_no_non_tensor_return(ret):
 
 def graph_class(dev): return dev.graph.func if isinstance(dev.graph, functools.partial) else dev.graph
 
-class GraphRunner(Runner):
+class GraphRunner:
   def __init__(self, linear:UOp, input_uops:tuple[UOp, ...]=()):
     self.linear = linear.src[0]
     self.calls: list[tuple[int, UOp, list[Buffer], dict[str, int]]] = []
     self.runtimes: list[Any|None] = []
     self.uop_replace: list[list[tuple[int, int]]] = []
     for call in self.linear.src:
-      replace = [(p, b.arg) for p, b in enumerate(b for b in call.src[1:] if b.op is not Ops.BIND) if b.op is Ops.PARAM]
+      replace = [(p, b.arg) for p, b in enumerate(get_call_arg_uops(call)) if b.op is Ops.PARAM]
       for dev_idx, (bufs, device_vars) in enumerate(unwrap_multi(call, resolve_params(call, input_uops))):
         self.calls.append((dev_idx, call.src[0], [b.ensure_allocated() for b in bufs], device_vars))
         self.runtimes.append(get_runtime(bufs[0].device, call.src[0]) if call.src[0].op is Ops.PROGRAM else None)
@@ -135,7 +127,9 @@ class GraphRunner(Runner):
     self.w_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
     self.r_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
 
-    super().__init__(colored(f"<batched {len(self.calls)}>", "cyan"), self.calls[0][2][0].device.split(":")[0], estimates.simplify())
+    self.device, self.estimates = self.calls[0][2][0].device.split(":")[0], estimates.simplify()
+
+  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None: raise NotImplementedError("override this")
 
   def updated_vars(self, var_vals: dict[str, int]):
     vals = [var_vals[v] for v in self.vars]
@@ -168,7 +162,7 @@ class GraphRunner(Runner):
 
   @staticmethod
   def _all_devs(batch_devs:list[Compiled], new_call:UOp) -> list[Compiled]:
-    return dedup(batch_devs + [Device[x] for b in new_call.src[1:] if b.op is not Ops.BIND
+    return dedup(batch_devs + [Device[x] for b in get_call_arg_uops(new_call)
                  for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
 
   @staticmethod
@@ -197,9 +191,9 @@ class CapturedJit(Generic[ReturnType]):
     out: set[UOp] = set()
     for call in self.linear.toposort():
       if call.op is not Ops.CALL: continue
-      non_bind = [s for s in call.src[1:] if s.op is not Ops.BIND]
-      outs, ins = _call_outs_ins(call)
-      out |= {non_bind[k] for k in outs - ins if non_bind[k].op in (Ops.BUFFER, Ops.BUFFER_VIEW)}
+      arg_uops = get_call_arg_uops(call)
+      outs, ins = get_call_outs_ins(call)
+      out |= {arg_uops[k] for k in set(outs) - set(ins) if arg_uops[k].op in (Ops.BUFFER, Ops.BUFFER_VIEW)}
     return out
 
   def __call__(self, input_uops:list[UOp], var_vals:dict[str, int]) -> ReturnType:

tinygrad/engine/realize.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
 from typing import cast, Iterator, Any
 import time, random, itertools, math, contextlib, weakref
 from dataclasses import dataclass, replace, field
-from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, Metadata, TRACEMETA, prod, flatten
+from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, TRACEMETA, prod, flatten
 from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite, ProgramInfo
@@ -10,66 +11,74 @@ from tinygrad.renderer import Estimates
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import bufs_from_ast
 
+# **************** Helpers ****************
+
+def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call.src[1:] if s.op is not Ops.BIND)
+
+def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
+  ast = call.src[0]
+  if ast.op is Ops.PROGRAM: return tuple(ast.arg.outs), tuple(ast.arg.ins)
+  if ast.op in (Ops.COPY, Ops.BUFFER_VIEW): return (0,), (1,)
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
+  return (), ()
+
+def get_call_name(call:UOp, bufs:list[Buffer], var_vals:dict[str, int]|None=None) -> str:
+  def _uop_sz_to_str(uop:UOp) -> str: return size_to_str(sym_infer(prod(uop.shape) * uop.dtype.itemsize, var_vals or {}))
+
+  ast, arg_uops = call.src[0], get_call_arg_uops(call)
+  if ast.op is Ops.PROGRAM: return ast.arg.name
+  if ast.op is Ops.BUFFER_VIEW: return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {ast.arg[1] * arg_uops[1].dtype.itemsize:<10d}", "yellow")
+  if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {bufs[0].device[:7]:>7s} <- {bufs[1].device[:7]:7s}", "yellow")
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow")
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return colored(f"batched {len(ast.src[0].src)}", "cyan")
+  raise NotImplementedError("get_call_name is not implemented")
+
 # **************** Stat ****************
 
 def estimate_uop(call:UOp) -> Estimates:
-  if call.src[0].op is Ops.SINK: call = pm_compile.rewrite(call)
-
   ast = call.src[0]
   if ast.op is Ops.PROGRAM: return ast.src[0].arg.estimates or Estimates()
   if ast.op is Ops.COPY or (ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec"):
     nbytes = prod(call.src[1].shape) * call.src[1].dtype.itemsize
     return Estimates(lds=nbytes, mem=nbytes)
-  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph":
-    return runner.estimates if (runner:=graph_cache.get(ast)) is not None else Estimates()
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return get_graph_runtime(ast).estimates
   return Estimates()
 
-def update_stats(display_name:str, device:str, estimates:Estimates, var_vals:dict[str, int], et:float|None, buf_count:int,
-                 jit=False, metadata:tuple[Metadata, ...]=(), first_run=False):
+first_run_cache:set[bytes] = set()
+@contextlib.contextmanager
+def track_stats(ctx:ExecContext, call:UOp, device:str, bufs:list[Buffer], var_vals:dict[str, int]):
+  if PROFILE:
+    outputs, inputs = get_call_outs_ins(call)
+    cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"metadata": call.arg.metadata, "var_vals": var_vals,
+      "bufs": [b.trace_num for b in bufs], "name": get_call_name(call, bufs, var_vals), "outputs": outputs, "inputs": inputs}))
+  et: list[float|None] = [None]
+  if DEBUG >= 2: st = time.perf_counter()
+  yield et
+  if not ctx.do_update_stats: return
+
+  if DEBUG >= 2 and et[0] is None:
+    Device[device].synchronize()
+    et[0] = time.perf_counter() - st
+
+  estimates = estimate_uop(call)
   GlobalCounters.kernel_count += 1
   GlobalCounters.global_ops += (op_est:=sym_infer(estimates.ops, var_vals))
   GlobalCounters.global_mem += (mem_est:=sym_infer(estimates.mem, var_vals))
-  if et is not None: GlobalCounters.time_sum_s += et
+  if et[0] is not None: GlobalCounters.time_sum_s += et[0]
   if DEBUG >= 2:
+    display_name = get_call_name(call, bufs, var_vals)
     lds_est = sym_infer(estimates.lds, var_vals)
-    header_color = 'magenta' if jit else ('green' if 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 ""
-    flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
+    header_color = 'magenta' if ctx.jit else ('green' if call.src[0].key not in first_run_cache else None)
+    ptm = colored(time_to_str(et[0], w=9), "yellow" if et[0] > 0.01 else None) if et[0] is not None else ""
+    flops, membw, ldsbw = op_est/(et[0] or 1e-20), mem_est/(et[0] or 1e-20), lds_est/(et[0] or 1e-20)
     flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
     mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \
       colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green')
     print(f"{colored(f'*** {device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
-      f" {display_name+' '*(46-ansilen(display_name))} arg {buf_count:2d} mem {GlobalCounters.mem_used/1e9:6.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 metadata] if metadata else ''}")
-
-first_run_cache:set[bytes] = set()
-@contextlib.contextmanager
-def track_stats(ctx:"ExecContext", call:UOp, device:str, display_name:str, bufs:list[Buffer], var_vals:dict[str, int], outputs=(0,), inputs=(1,)):
-  if PROFILE: cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"metadata": call.arg.metadata, "var_vals": var_vals,
-                                                  "bufs": [b.trace_num for b in bufs], "name": display_name, "outputs": outputs, "inputs": inputs}))
-  timing: list[float|None] = [None]
-  if DEBUG >= 2: st = time.perf_counter()
-  yield timing
-  if not ctx.do_update_stats: return
-  if DEBUG >= 2 and timing[0] is None:
-    Device[device].synchronize()
-    timing[0] = time.perf_counter() - st
-  update_stats(display_name, device, estimate_uop(call), var_vals, timing[0], len(bufs), jit=ctx.jit, metadata=call.arg.metadata,
-               first_run=call.src[0].key not in first_run_cache)
-  first_run_cache.add(call.src[0].key)
-
-# **************** Runners ****************
-
-class Runner:
-  def __init__(self, display_name:str, device:str, estimates=Estimates()):
-    self.first_run, self.display_name, self.device, self.estimates = True, display_name, device, estimates
-  @property
-  def dev(self): return Device[self.device]
-  def exec(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None) -> float|None:
-    return self(rawbufs, {} if var_vals is None else var_vals)
-  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False) -> float|None:
-    raise NotImplementedError("override this")
+      f" {display_name+' '*(46-ansilen(display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
+      ("" if et[0] 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 call.arg.metadata] if call.arg.metadata else ''}")
+    first_run_cache.add(call.src[0].key)
 
 local_size_cache: dict[bytes, tuple[int, ...]] = {}
 def optimize_local_size(call:UOp, prg:UOp) -> UOp|None:
@@ -93,7 +102,7 @@ def optimize_local_size(call:UOp, prg:UOp) -> UOp|None:
   new_global = tuple(g//l if g%l == 0 else g/l for g,l in zip(prg.arg.global_size, local_size))
   return call.replace(src=(prg.replace(arg=replace(prg.arg, global_size=new_global, local_size=local_size)), *call.src[1:]))
 
-# **************** method cache ****************
+# **************** runtime cache ****************
 
 runtime_cache: dict[tuple[bytes, str], Any] = {}
 def get_runtime(device:str, ast:UOp):
@@ -105,6 +114,13 @@ def get_runtime(device:str, ast:UOp):
     runtime = runtime_cache[key] = Device[device].runtime(ast.arg.function_name, ast.src[4].arg, *ast.arg.aux, runtimevars=ast.arg.runtimevars)
   return runtime
 
+graph_cache:weakref.WeakKeyDictionary[UOp, Any] = weakref.WeakKeyDictionary()
+def get_graph_runtime(ast:UOp, input_uops:tuple[UOp, ...]|None=None):
+  assert ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph", "get_graph_runtime should only be called with a graph ast"
+  if (runtime:=graph_cache.get(ast)) is None and input_uops is not None:
+    graph_cache[ast] = runtime = Device[ast.device if isinstance(ast.device, str) else ast.device[0]].graph(ast, input_uops=input_uops)
+  return runtime
+
 # **************** run linear ****************
 
 capturing: list = []  # put classes with an add_linear method in here
@@ -119,7 +135,7 @@ class ExecContext:
 def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
   if b.op in (Ops.BUFFER_VIEW, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg], *b.src[1:]))
   return inputs[b.arg] if b.op is Ops.PARAM else b
-def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in call.src[1:] if b.op is not Ops.BIND]
+def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]
 
 def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], dict[str, int]]]:
   bufs = [b.buffer for b in resolved]
@@ -132,16 +148,13 @@ def exec_view(ctx:ExecContext, call, ast):
   resolved = resolve_params(call, ctx.input_uops)
   bufs = [cast(Buffer, b.buffer) for b in resolved]
   bv = bufs[1].view(resolved[0].arg, ast.dtype, ast.arg[1]*bufs[1].dtype.itemsize)
-  with track_stats(ctx, call, bv.device, colored(f"view {bv.nbytes:8d} @ {bv.offset:<10d}", "yellow"), [bv, bufs[1]], ctx.var_vals):
-    buffers[resolved[0]] = bv
+  with track_stats(ctx, call, bv.device, [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv
 
 def exec_copy(ctx:ExecContext, call, ast):
   for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
     dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated()
-    xfer = hasattr(dest.allocator,'_transfer') and dest.allocator.supports_transfer and dest.device.split(":")[0] == src.device.split(":")[0]
-    name = colored(f"{'xfer' if xfer else 'copy'} {size_to_str(bufs[0].nbytes):>10}, {dest.device[:7]:>7s} <- {src.device[:7]:7s}", "yellow")
-    with track_stats(ctx, call, dest.device, name, [dest, src], ctx.var_vals):
-      if xfer:
+    with track_stats(ctx, call, dest.device, [dest, src], ctx.var_vals):
+      if hasattr(dest.allocator,'_transfer') and dest.allocator.supports_transfer and dest.device.split(":")[0] == src.device.split(":")[0]:
         dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev) # type:ignore[attr-defined]
       elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \
            and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
@@ -156,7 +169,7 @@ def exec_kernel(ctx:ExecContext, call, ast):
     prg_bufs = [bufs[i].ensure_allocated() for i in ast.arg.globals]
     rt = get_runtime(device:=bufs[0].device, ast)
     global_size, local_size = ast.arg.launch_dims(var_vals)
-    with track_stats(ctx, call, device, ast.arg.name, prg_bufs, var_vals, outputs=ast.arg.outs, inputs=ast.arg.ins) as tm:
+    with track_stats(ctx, call, device, prg_bufs, var_vals) as tm:
       tm[0] = rt(*[b._buf for b in prg_bufs], global_size=global_size, local_size=local_size, vals=ast.arg.vals(var_vals), wait=DEBUG>=2)
 
 def exec_validate(ctx:ExecContext, call, ast):
@@ -172,16 +185,12 @@ def exec_validate(ctx:ExecContext, call, ast):
 def exec_encdec(ctx:ExecContext, call, ast):
   bufs = [cast(Buffer, b.buffer).ensure_allocated() for b in resolve_params(call, ctx.input_uops)]
   shape, pos_var = tuple(s.arg for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr
-  with track_stats(ctx, call, bufs[0].device, colored(f"enc/dec {size_to_str(bufs[0].nbytes)}", "yellow"), bufs, ctx.var_vals):
+  with track_stats(ctx, call, bufs[0].device, bufs, ctx.var_vals):
     bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var])
 
-graph_cache:weakref.WeakKeyDictionary[UOp, Runner] = weakref.WeakKeyDictionary()
-def exec_graph(ctx:ExecContext, call, cf):
-  bufs = flatten([b.bufs if isinstance(b, MultiBuffer) else [b] for b in (u.buffer for u in resolve_params(call, ctx.input_uops))])
-  if (runner:=graph_cache.get(cf)) is None:
-    graph_cache[cf] = runner = Device[cf.device if isinstance(cf.device, str) else cf.device[0]].graph(cf, input_uops=ctx.input_uops)
-  with track_stats(ctx, call, runner.device, runner.display_name, bufs, ctx.var_vals) as t:
-    t[0] = runner(bufs, ctx.var_vals, wait=DEBUG >= 2, input_uops=ctx.input_uops) # type: ignore[call-arg]
+def exec_graph(ctx:ExecContext, call, ast):
+  rt = get_graph_runtime(ast, ctx.input_uops)
+  with track_stats(ctx, call, rt.device, [], ctx.var_vals) as t: t[0] = rt(ctx.input_uops, ctx.var_vals, wait=DEBUG>=2) # type: ignore[call-arg]
 
 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -190,7 +199,7 @@ pm_flatten_linear = PatternMatcher([
 ])
 
 def _validate(call:UOp, sink:UOp) -> UOp:
-  params = tuple(p for p in call.src[1:] if p.op is not Ops.BIND)
+  params = get_call_arg_uops(call)
   shadows = tuple(UOp.new_buffer(("CPU",)*len(p.device) if isinstance(p.device, tuple) else "CPU", prod(p.max_shape), p.dtype.base) for p in params)
   copies = tuple(p.copy_to_device(s.device).call(s, p) for s, p in zip(shadows, params))
   return UOp(Ops.LINEAR, src=copies + (call, UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(sink,), arg="validate").call(*shadows, *params)))
@@ -216,7 +225,7 @@ pm_exec = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="ast"),), name="call", allow_any_len=True), exec_copy),
   (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="ast"),), name="call", allow_any_len=True), exec_kernel),
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="ast"),), name="call", allow_any_len=True), exec_encdec),
-  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="graph", name="cf"),), name="call", allow_any_len=True), exec_graph),
+  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="graph", name="ast"),), name="call", allow_any_len=True), exec_graph),
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
 ])
 

tinygrad/renderer/cstyle.py

@@ -45,7 +45,7 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)),
   # new load/store
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), allow_any_len=True),
-   lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
+   lambda ctx,buf,idx: None if isinstance(buf.dtype, ImageDType) else f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
   (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat.var("gate"))).or_casted("bidx"), UPat.var("var"))),
    lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
   (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"),
@@ -296,19 +296,27 @@ class OpenCLRenderer(CStyleLanguage):
               dtypes.bfloat16: "ushort" }
   extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast + extra_pm
 
+  @staticmethod
+  def image_coord(ctx, idx_x, idx_y): return f"({ctx.render_dtype(dtypes.int.vec(2))})({ctx[idx_x]},{ctx[idx_y]})"
+
   string_rewrite = PatternMatcher([
     (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
     # bfloat16 constants need to be rendered as their bit pattern since bf16 is stored as ushort
     (UPat(Ops.CONST, dtypes.bfloat16, name="x"),
       lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.arg)))[0] >> 16)}u"),
     # load/store image (OpenCL)
-    (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2)), UPat.var("gate")), UPat.var("var"))),
-      lambda ctx,buf,idx,var,gate: f"({ctx[gate]}?read_imagef({ctx[buf]}, smp, {ctx[idx]}):{ctx[var]})"),
-    (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2))),)),
-      lambda ctx,buf,idx: f"read_imagef({ctx[buf]}, smp, {ctx[idx]})"),
-    (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx', dtypes.int.vec(2)), allow_any_len=True),
-                          UPat.var("var", dtypes.float.vec(4))), allow_any_len=True),
-      lambda ctx,buf,idx,var: f"write_imagef({ctx[buf]}, {ctx[idx]}, {ctx[var]});"),
+    (UPat(Ops.INDEX, src=(UPat.var('buf'), UPat.var('idx_x', dtypes.int), UPat.var('idx_y', dtypes.int)), allow_any_len=True),
+      lambda ctx,buf,idx_x,idx_y: ctx.image_coord(ctx, idx_x, idx_y) if isinstance(buf.dtype, ImageDType) else None),
+    (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(
+      UPat.var('buf').index(UPat.var('idx_x', dtypes.int), UPat.var('idx_y', dtypes.int), UPat.var("gate")), UPat.var("var"))),
+      lambda ctx,buf,idx_x,idx_y,var,gate: f"({ctx[gate]}?read_imagef({ctx[buf]}, smp, {ctx.image_coord(ctx, idx_x, idx_y)}):{ctx[var]})"),
+    (UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(
+      UPat.var('buf').index(UPat.var('idx_x', dtypes.int), UPat.var('idx_y', dtypes.int)),)),
+      lambda ctx,buf,idx_x,idx_y: f"read_imagef({ctx[buf]}, smp, {ctx.image_coord(ctx, idx_x, idx_y)})"),
+    (UPat(Ops.STORE, src=(
+      UPat.var('buf').index(UPat.var('idx_x', dtypes.int), UPat.var('idx_y', dtypes.int), allow_any_len=True),
+      UPat.var("var", dtypes.float.vec(4))), allow_any_len=True),
+      lambda ctx,buf,idx_x,idx_y,var: f"write_imagef({ctx[buf]}, {ctx.image_coord(ctx, idx_x, idx_y)}, {ctx[var]});"),
   ]) + base_rewrite
 
   def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:

tinygrad/renderer/nir.py

@@ -136,8 +136,9 @@ class NIRRenderer(Renderer):
     # ref: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpConvertFToU
     (UPat(Ops.CAST, (dtypes.uchar, dtypes.ushort), src=(UPat.var("x", dtypes.floats),), name="c"), lambda x,c: x.cast(dtypes.int32).cast(c.dtype)),
     # load/store use pointer arithmetic, and the cast does nothing
-    (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace(
-      src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.dtype.addrspace != AddrSpace.REG and off.op not in (Ops.CAST, Ops.STACK) else None),
+    (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off:
+      x.replace(src=(buf,off.cast(dtypes.long))+x.src[2:]) if not isinstance(buf.dtype, ImageDType) and \
+        buf.dtype.addrspace != AddrSpace.REG and off.op not in (Ops.CAST, Ops.STACK) else None),
     (UPat(Ops.CAST, name="x"), lambda x: x.src[0] if isinstance(x.dtype, PtrDType) or x.src[0].dtype == dtypes.void else None),
   ])
 
@@ -249,30 +250,31 @@ class LVPRenderer(NIRRenderer):
     self.param_sz = sum([8 if u.op == Ops.PARAM else u.dtype.itemsize for u in uops if u.op in (Ops.PARAM, Ops.DEFINE_VAR)])
 
 # FIXME: this should be a rewrite rule
-def tovec(b, coord): return nalu(b, "vec4", nchannel(b, coord, 0), nchannel(b, coord, 1), nundef(b, dtypes.int), nundef(b, dtypes.int))
+def tovec(b, coord_x, coord_y): return nalu(b, "vec4", coord_x, coord_y, nundef(b, dtypes.int), nundef(b, dtypes.int))
 def nfloat(dtype): return mesa.nir_type_float16 if dtype == dtypes.half else mesa.nir_type_float32
 nstore_img = nir_instr(has_def=False, df=lambda img:img, num_components=lambda val:val.num_components,
   intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2D, 'ACCESS':mesa.ACCESS_CAN_REORDER, 'SRC_TYPE':nfloat(dtype)},
-  srcs=lambda b,img,coord,val:[nsrc(x) for x in [img, tovec(b, coord), nundef(b, dtypes.int), val, nimm(b, 0, dtypes.int)]])(
-    lambda b,img,coord,val,dtype:mesa.nir_intrinsic_instr_create(b.shader,g("nir_intrinsic_image_store")))
+  srcs=lambda b,img,coord_x,coord_y,val:[nsrc(x) for x in [img, tovec(b, coord_x, coord_y), nundef(b, dtypes.int), val, nimm(b, 0, dtypes.int)]])(
+    lambda b,img,coord_x,coord_y,val,dtype:mesa.nir_intrinsic_instr_create(b.shader,g("nir_intrinsic_image_store")))
 
 _nload_img = nir_instr(intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2D, 'ACCESS':mesa.ACCESS_CAN_REORDER, 'DEST_TYPE':nfloat(dtype)},
-  nc=4, bs=32, num_components=4, srcs=lambda b,img,coord:[nsrc(x) for x in [img, tovec(b, coord), nundef(b, dtypes.int), nimm(b, 0, dtypes.int)]])(
-    lambda b,img,coord,dtype: mesa.nir_intrinsic_instr_create(b.shader, g("nir_intrinsic_image_load")))
+  nc=4, bs=32, num_components=4,
+  srcs=lambda b,img,coord_x,coord_y:[nsrc(x) for x in [img, tovec(b, coord_x, coord_y), nundef(b, dtypes.int), nimm(b, 0, dtypes.int)]])(
+    lambda b,img,coord_x,coord_y,dtype: mesa.nir_intrinsic_instr_create(b.shader, g("nir_intrinsic_image_load")))
 
 class IR3Renderer(NIRRenderer, OpenCLRenderer):
   has_aux = True
 
-  def nload_img(ctx,img,coord):
+  def nload_img(ctx,img,idx_x,idx_y):
     ctx.texs.add(img)
-    return _nload_img(ctx.b, ctx.r[img], ctx.r[coord], img.dtype)
+    return _nload_img(ctx.b, ctx.r[img], ctx.r[idx_x], ctx.r[idx_y], img.dtype)
 
   def_rewrite = PatternMatcher([
-    (UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), allow_any_len=True), UPat.var("val")),
-          allow_any_len=True), lambda ctx,img,coord,val: nstore_img(ctx.b, ctx.r[img], ctx.r[coord], ctx.r[val], val.dtype)),
-    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2)), UPat.var("gate")), UPat.var("alt"))),
-     lambda ctx,img,coord,alt,gate: if_phi(ctx.b, ctx.r[gate], lambda: ctx.nload_img(img, coord), lambda: ctx.r[alt])),
-    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('coord', dtypes.int.vec(2))),)), nload_img),
+    (UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('idx_x', dtypes.int), UPat.var('idx_y', dtypes.int), allow_any_len=True), UPat.var("val")),
+          allow_any_len=True), lambda ctx,img,idx_x,idx_y,val: nstore_img(ctx.b, ctx.r[img], ctx.r[idx_x], ctx.r[idx_y], ctx.r[val], val.dtype)),
+    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('idx_x', dtypes.int), UPat.var('idx_y', dtypes.int), UPat.var("gate")), UPat.var("alt"))),
+     lambda ctx,img,idx_x,idx_y,alt,gate: if_phi(ctx.b, ctx.r[gate], lambda: ctx.nload_img(img, idx_x, idx_y), lambda: ctx.r[alt])),
+    (UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('idx_x', dtypes.int), UPat.var('idx_y', dtypes.int)),)), nload_img),
   ]) + NIRRenderer.def_rewrite
 
   _param = LVPRenderer.param

tinygrad/runtime/autogen/__init__.py

@@ -22,7 +22,7 @@ webgpu_lib = "os.path.join(sysconfig.get_paths()['purelib'], 'pydawn', 'lib', 'l
 nv_lib_path = ("[f'/{pre}/cuda/targets/{tgt}/lib' for pre in ['opt', 'usr/local'] for tgt in "
                "[sysconfig.get_config_vars().get(\"MULTIARCH\", \"\").rsplit(\"-\", 1)[0], 'sbsa-linux']]")
 
-def load(name, dll, files, **kwargs):
+def load(name, files, **kwargs):
   if not (f:=(root/(path:=kwargs.pop("path", __name__)).replace('.','/')/f"{name}.py")).exists() or getenv('REGEN'):
     files, kwargs['args'] = files() if callable(files) else files, args() if callable(args:=kwargs.get('args', [])) else args
     if (srcs:=kwargs.pop('srcs', None)):
@@ -39,23 +39,24 @@ def load(name, dll, files, **kwargs):
       if (preprocess:=kwargs.pop('preprocess', None)): preprocess(srcpath)
     files = flatten(sorted(glob.glob(p, recursive=True)) if isinstance(p, str) and '*' in p else [p] for p in files)
     kwargs['epilog'] = (epi(srcpath) if srcs else epi()) if callable(epi:=kwargs.get('epilog', [])) else epi
-    f.write_text(importlib.import_module("tinygrad.runtime.support.autogen").gen(name, dll, files, **kwargs))
+    try: f.write_text(kwargs.pop("gen", importlib.import_module("tinygrad.runtime.support.autogen").gen)(name, files, **kwargs))
+    except Exception as e: raise RuntimeError(f"error while generating {name}") from e
     if srcs: td.cleanup()
   return importlib.import_module(f"{path}.{name.replace('/', '.')}")
 
 def __getattr__(nm):
   match nm:
-    case "libc": return load("libc", "'c'", lambda: (
+    case "libc": return load("libc", lambda: (
       [i for i in system("dpkg -L libc6-dev").split() if 'sys/mman.h' in i or 'sys/syscall.h' in i] +
-      ["/usr/include/string.h", "/usr/include/elf.h", "/usr/include/unistd.h", "/usr/include/asm-generic/mman-common.h"]), errno=True)
-    case "avcodec": return load("avcodec", None, ["{}/libavcodec/hevc/hevc.h", "{}/libavcodec/cbs_h265.h"], srcs=ffmpeg_src)
-    case "opencl": return load("opencl", "'OpenCL'", ["/usr/include/CL/cl.h"])
-    case "cuda": return load("cuda", "'cuda'", ["/usr/include/cuda.h"], args=["-D__CUDA_API_VERSION_INTERNAL"], parse_macros=False)
-    case "nvrtc": return load("nvrtc", "'nvrtc'", ["/usr/include/nvrtc.h"], paths=nv_lib_path, prolog=["import sysconfig"])
-    case "nvjitlink": load("nvjitlink", "'nvJitLink'", [root/"extra/nvJitLink.h"], paths=nv_lib_path, prolog=["import sysconfig"])
-    case "kfd": return load("kfd", None, [root/"extra/hip_gpu_driver/kfd_ioctl.h"])
+      ["/usr/include/string.h", "/usr/include/elf.h", "/usr/include/unistd.h", "/usr/include/asm-generic/mman-common.h"]), dll="'c'", errno=True)
+    case "avcodec": return load("avcodec", ["{}/libavcodec/hevc/hevc.h", "{}/libavcodec/cbs_h265.h"], srcs=ffmpeg_src)
+    case "opencl": return load("opencl", ["/usr/include/CL/cl.h"], dll="'OpenCL'")
+    case "cuda": return load("cuda", ["/usr/include/cuda.h"], dll="'cuda'", args=["-D__CUDA_API_VERSION_INTERNAL"], macros=False)
+    case "nvrtc": return load("nvrtc", ["/usr/include/nvrtc.h"], dll="'nvrtc'", paths=nv_lib_path, prolog=["import sysconfig"])
+    case "nvjitlink": load("nvjitlink", [root/"extra/nvJitLink.h"], dll="'nvJitLink'", paths=nv_lib_path, prolog=["import sysconfig"])
+    case "kfd": return load("kfd", [root/"extra/hip_gpu_driver/kfd_ioctl.h"])
     case "nv_570" | "nv_580":
-      return load(nm, None, [
+      return load(nm, [
         *[root/"extra/nv_gpu_driver"/s for s in ["clc9b0.h", "clc6c0qmd.h","clcec0qmd.h", "nvdec_drv.h"]], "{}/kernel-open/common/inc/nvmisc.h",
         *[f"{{}}/src/common/sdk/nvidia/inc/class/cl{s}.h" for s in ["0000", "0070", "0080", "2080", "2080_notification", "c56f", "c86f", "c96f", "c761",
                                                                     "83de", "b2cc", "c6c0", "cdc0"]],
@@ -70,7 +71,7 @@ def __getattr__(nm):
         "-include", "{}/src/common/sdk/nvidia/inc/nvtypes.h", "-I{}/src/common/inc", "-I{}/kernel-open/nvidia-uvm", "-I{}/kernel-open/common/inc",
         "-I{}/src/common/sdk/nvidia/inc", "-I{}/src/nvidia/arch/nvalloc/unix/include", "-I{}/src/common/sdk/nvidia/inc/ctrl"
       ], rules=[(r'MW\(([^:]+):(.+)\)',r'(\1, \2)'), (r'(\d+):(\d+)', r'(\1, \2)')], srcs=nv_src[nm], anon_names={"{}/kernel-open/common/inc/nvstatus.h:37":"nv_status_codes"})
-    case "nv": return load("nv", None, [
+    case "nv": return load("nv", [
       *[f"{{}}/src/nvidia/inc/kernel/gpu/{s}.h" for s in ["fsp/kern_fsp_cot_payload", "gsp/gsp_init_args"]],
       *[f"{{}}/src/nvidia/arch/nvalloc/common/inc/{s}.h" for s in ["gsp/gspifpub", "gsp/gsp_fw_wpr_meta", "gsp/gsp_fw_sr_meta", "rmRiscvUcode",
                                                                    "fsp/fsp_nvdm_format"]],
@@ -89,49 +90,50 @@ def __getattr__(nm):
     })
     # this defines all syscall numbers. should probably unify linux autogen?
     case "io_uring":
-      return load("io_uring", None, ["{}/liburing.h", "{}/usr/include/linux/io_uring.h", "{}/usr/include/asm-generic/unistd.h"],
+      return load("io_uring", ["{}/liburing.h", "{}/usr/include/linux/io_uring.h", "{}/usr/include/asm-generic/unistd.h"],
                   args=["-I{}/usr/include"], srcs=[linux_headers_deb, liburing_src], rules=[('__NR', 'NR')],
                   preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
-    case "ib": return load("ib", "'ibverbs'", ["/usr/include/infiniband/verbs.h", "/usr/include/infiniband/verbs_api.h",
-                                               "/usr/include/infiniband/ib_user_ioctl_verbs.h","/usr/include/rdma/ib_user_verbs.h"], errno=True)
-    case "llvm": return load("llvm", llvm_lib, lambda: [system("llvm-config-20 --includedir")+"/llvm-c/**/*.h"],
+    case "ib": return load("ib", ["/usr/include/infiniband/verbs.h", "/usr/include/infiniband/verbs_api.h",
+                                  "/usr/include/infiniband/ib_user_ioctl_verbs.h", "/usr/include/rdma/ib_user_verbs.h"], dll="'ibverbs'", errno=True)
+    case "llvm": return load("llvm", lambda: [system("llvm-config-20 --includedir")+"/llvm-c/**/*.h"], dll=llvm_lib,
                              args=lambda: system("llvm-config-20 --cflags").split(), recsym=True, prolog=["from tinygrad.helpers import WIN, OSX"])
-    case "pci": return load("pci", None, ["{}/usr/include/linux/pci_regs.h"], srcs=linux_headers_deb,
+    case "pci": return load("pci", ["{}/usr/include/linux/pci_regs.h"], srcs=linux_headers_deb,
                              preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
-    case "vfio": return load("vfio", None, ["{}/usr/include/linux/vfio.h"], args=["-I{}/usr/include"], srcs=linux_headers_deb,
+    case "vfio": return load("vfio", ["{}/usr/include/linux/vfio.h"], args=["-I{}/usr/include"], srcs=linux_headers_deb,
                              preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
     # could add rule: WGPU_COMMA -> ','
-    case "webgpu": return load("webgpu", webgpu_lib, [root/"extra/webgpu/webgpu.h"],
+    case "webgpu": return load("webgpu", [root/"extra/webgpu/webgpu.h"], dll=webgpu_lib,
                                prolog=["from tinygrad.helpers import WIN, OSX", "import sysconfig, os"])
-    case "libusb": return load("libusb", "'usb-1.0'", ["/usr/include/libusb-1.0/libusb.h"])
-    case "hip": return load("hip", "os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamdhip64.so'", ["/opt/rocm/include/hip/hip_ext.h",
-                            "/opt/rocm/include/hip/hiprtc.h", "/opt/rocm/include/hip/hip_runtime_api.h", "/opt/rocm/include/hip/driver_types.h"],
+    case "libusb": return load("libusb", ["/usr/include/libusb-1.0/libusb.h"], dll="'usb-1.0'")
+    case "hip": return load("hip", ["/opt/rocm/include/hip/hip_ext.h", "/opt/rocm/include/hip/hiprtc.h",
+                                    "/opt/rocm/include/hip/hip_runtime_api.h", "/opt/rocm/include/hip/driver_types.h"],
+                            dll="os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamdhip64.so'",
                             args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"], prolog=["import os"])
     case "comgr" | "comgr_3":
-      return load("comgr_3" if nm == "comgr_3" else "comgr", "[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so', 'amd_comgr']",
-                  ["/opt/rocm/include/amd_comgr/amd_comgr.h"], args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"],
-                  prolog=["import os"])
-    case "hsa": return load("hsa", "[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libhsa-runtime64.so', 'hsa-runtime64']", [
-      *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/core/inc/{s}.h" for s in ["registers"]],
-      *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/inc/{s}.h" for s in ["hsa", "hsa_ext_amd", "amd_hsa_signal", "amd_hsa_queue",
-                                                                              "amd_hsa_kernel_code", "hsa_ext_finalize",
-                                                                              "hsa_ext_image", "hsa_ven_amd_aqlprofile"]]],
+      return load("comgr_3" if nm == "comgr_3" else "comgr", ["/opt/rocm/include/amd_comgr/amd_comgr.h"],
+                  dll= "[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so', 'amd_comgr']",
+                  args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"], prolog=["import os"])
+    case "hsa": return load("hsa", [*[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/core/inc/{s}.h" for s in ["registers"]],
+                                    *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/inc/{s}.h" for s in [
+                                        "hsa", "hsa_ext_amd", "amd_hsa_signal", "amd_hsa_queue", "amd_hsa_kernel_code",
+                                        "hsa_ext_finalize", "hsa_ext_image", "hsa_ven_amd_aqlprofile"]]],
+      dll="[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libhsa-runtime64.so', 'hsa-runtime64']",
       srcs=rocr_src, args=["-DLITTLEENDIAN_CPU"], prolog=["import os"])
-    case "amdgpu_kd": return load("amdgpu_kd", None, lambda: [f"{system('llvm-config-20 --includedir')}/llvm/Support/AMDHSAKernelDescriptor.h"],
-                                  args=lambda: system("llvm-config-20 --cflags").split() + ["-x", "c++"], recsym=True, parse_macros=False)
-    case "amd_gpu": return load("amd_gpu", None, [root/f"extra/hip_gpu_driver/{s}.h" for s in ["sdma_registers", "nvd", "gc_11_0_0_offset",
+    case "amdgpu_kd": return load("amdgpu_kd", lambda: [f"{system('llvm-config-20 --includedir')}/llvm/Support/AMDHSAKernelDescriptor.h"],
+                                  args=lambda: system("llvm-config-20 --cflags").split() + ["-x", "c++"], recsym=True, macros=False)
+    case "amd_gpu": return load("amd_gpu", [root/f"extra/hip_gpu_driver/{s}.h" for s in ["sdma_registers", "nvd", "gc_11_0_0_offset",
                                                                                                "sienna_cichlid_ip_offset"]],
                                 args=["-I/opt/rocm/include", "-x", "c++"])
-    case "amdgpu_drm": return load("amdgpu_drm", None, [ "/usr/include/drm/drm.h", *[root/f"extra/hip_gpu_driver/{s}.h" for s in ["amdgpu_drm"]]])
-    case "kgsl": return load("kgsl", None, [root/"extra/qcom_gpu_driver/msm_kgsl.h"], args=["-D__user="])
+    case "amdgpu_drm": return load("amdgpu_drm", [ "/usr/include/drm/drm.h", *[root/f"extra/hip_gpu_driver/{s}.h" for s in ["amdgpu_drm"]]])
+    case "kgsl": return load("kgsl", [root/"extra/qcom_gpu_driver/msm_kgsl.h"], args=["-D__user="])
     case "qcom_dsp":
-      return load("qcom_dsp", None, [root/f"extra/dsp/include/{s}.h" for s in ["ion", "msm_ion", "adsprpc_shared", "remote_default", "apps_std"]])
-    case "sqtt": return load("sqtt", None, [root/"extra/sqtt/sqtt.h"])
+      return load("qcom_dsp", [root/f"extra/dsp/include/{s}.h" for s in ["ion", "msm_ion", "adsprpc_shared", "remote_default", "apps_std"]])
+    case "sqtt": return load("sqtt", [root/"extra/sqtt/sqtt.h"])
     case "rocprof":
-      return load("rocprof", "['rocprof-trace-decoder', p:='/usr/local/lib/rocprof-trace-decoder.so', p.replace('so','dylib')]",
-                  [f"{{}}/include/{s}.h" for s in ["rocprof_trace_decoder", "trace_decoder_instrument", "trace_decoder_types"]],
+      return load("rocprof", [f"{{}}/include/{s}.h" for s in ["rocprof_trace_decoder", "trace_decoder_instrument", "trace_decoder_types"]],
+                  dll= "['rocprof-trace-decoder', p:='/usr/local/lib/rocprof-trace-decoder.so', p.replace('so','dylib')]",
                   srcs="https://github.com/ROCm/rocprof-trace-decoder/archive/dd0485100971522cc4cd8ae136bdda431061a04d.tar.gz")
-    case "mesa": return load("mesa", "([] if DEV.renderer == 'LVP' else ['tinymesa']) + ['tinymesa_cpu']", [
+    case "mesa": return load("mesa", [
         *[f"{{}}/src/compiler/nir/{s}.h" for s in ["nir", "nir_builder", "nir_shader_compiler_options", "nir_serialize"]], "{}/gen/nir_intrinsics.h",
         *[f"{{}}/src/nouveau/{s}.h" for s in ["headers/nv_device_info", "compiler/nak"]],
         *[f"{{}}/src/gallium/auxiliary/gallivm/lp_bld{s}.h" for s in ["", "_passmgr", "_misc", "_type", "_init", "_nir", "_struct", "_jit_types",
@@ -150,28 +152,28 @@ def __getattr__(nm):
           *[f"python3 src/compiler/{s}_h.py > gen/{s.split('/')[-1]}.h" for s in ["nir/nir_opcodes", "nir/nir_builder_opcodes"]],
           *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
   srcs="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
+  dll="([] if DEV.renderer == 'LVP' else ['tinymesa']) + ['tinymesa_cpu']",
   prolog=["from tinygrad.helpers import DEV", "import gzip, base64"],
   epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
     case "libclang":
-      return load("libclang", clang_lib,
+      return load("libclang",
                   lambda: [f"{system('llvm-config-20 --includedir')}/clang-c/{s}.h" for s in ["Index", "CXString", "CXSourceLocation", "CXFile"]],
-                  prolog=["from tinygrad.helpers import WIN, OSX"], args=lambda: system("llvm-config-20 --cflags").split())
+                  dll=clang_lib, prolog=["from tinygrad.helpers import WIN, OSX"], args=lambda: system("llvm-config-20 --cflags").split())
     case "metal":
-      return load("metal", "'Metal'", [f"{macossdk}/System/Library/Frameworks/Metal.framework/Headers/MTL{s}.h" for s in
+      return load("metal", [f"{macossdk}/System/Library/Frameworks/Metal.framework/Headers/MTL{s}.h" for s in
                   ["ComputeCommandEncoder", "ComputePipeline", "CommandQueue", "Device", "IndirectCommandBuffer", "Resource", "CommandEncoder"]],
-                  args=["-xobjective-c","-isysroot",macossdk], types={"dispatch_data_t":"objc.id_"})
-    case "iokit": return load("iokit", "'IOKit'", [f"{macossdk}/System/Library/Frameworks/IOKit.framework/Headers/IOKitLib.h"],
+                  dll="'Metal'", args=["-xobjective-c","-isysroot",macossdk], types={"dispatch_data_t":"objc.id_"})
+    case "iokit": return load("iokit", [f"{macossdk}/System/Library/Frameworks/IOKit.framework/Headers/IOKitLib.h"], dll="'IOKit'",
                               args=["-isysroot", macossdk])
-    case "corefoundation": return load("corefoundation", "'CoreFoundation'",
+    case "corefoundation": return load("corefoundation",
                                        [f"{macossdk}/System/Library/Frameworks/CoreFoundation.framework/Headers/CF{s}.h" for s in ["String", "Data"]],
-                                       args=["-isysroot", macossdk])
-    case "llvm_qcom": return load("llvm_qcom", "'llvm-qcom'", [root/"extra/tinydreno.h"])
-    case "ggml_common":
-      return load("ggml_common", None, ["{}/ggml-common.h"], srcs=ggml_common_src,
-                  args=["-DGGML_COMMON_DECL_C", "-DGGML_COMMON_IMPL_C"], parse_macros=False)
+                                       dll="'CoreFoundation'",args=["-isysroot", macossdk])
+    case "llvm_qcom": return load("llvm_qcom", [root/"extra/tinydreno.h"], dll="'llvm-qcom'")
+    case "ggml_common": return load("ggml_common", ["{}/ggml-common.h"], srcs=ggml_common_src,
+                                    args=["-DGGML_COMMON_DECL_C", "-DGGML_COMMON_IMPL_C"], macros=False)
     case "mlx5":
       kh = "{}/usr/src/linux-headers-6.18.9+deb14-common/include/linux/mlx5"
-      return load("mlx5", None, [root/"extra/mlx_driver/mlx5.h", f"{kh}/mlx5_ifc.h"], srcs=linux_headers_kern_deb,
+      return load("mlx5", [root/"extra/mlx_driver/mlx5.h", f"{kh}/mlx5_ifc.h"], srcs=linux_headers_kern_deb,
                   args=["-Du8=unsigned char", "-Du16=unsigned short", "-Du32=unsigned int", "-Du64=unsigned long long",
                         "-D__be16=unsigned short", "-D__be32=unsigned int", "-D__be64=unsigned long long", f"-I{kh}"],
                   preprocess=lambda path: subprocess.run(f"ar x {linux_headers_kern_deb.split('/')[-1]} && tar xf data.tar.xz",

tinygrad/runtime/autogen/am/__init__.py

@@ -1,29 +1,37 @@
+import pathlib, hashlib
 from tinygrad.runtime.autogen import load, root
 
 am_src="https://github.com/ROCm/ROCK-Kernel-Driver/archive/33970e1351f5e511029602454979f3de7e22260f.tar.gz"
 AMD, AMDINC = "{}/drivers/gpu/drm/amd", "{}/drivers/gpu/drm/amd/include"
 inc, kern_rules = ["-include", "stdint.h"], [(r'le32_to_cpu', ''),]
+fw_src="https://gitlab.com/kernel-firmware/linux-firmware/-/archive/1e2c15348485939baf1b6d1f5a7a3b799d80703d/1e2c15348485939baf1b6d1f5a7a3b799d80703d.tar.gz"
 
 def __getattr__(nm):
   match nm:
-    case "am": return load("am/am", [], [root/f"extra/amdpci/headers/{s}.h" for s in ["v11_structs", "v12_structs", "amdgpu_vm",
+    case "am": return load("am/am", [root/f"extra/amdpci/headers/{s}.h" for s in ["v11_structs", "v12_structs", "amdgpu_vm",
       "discovery", "amdgpu_ucode", "psp_gfx_if", "amdgpu_psp", "amdgpu_irq", "amdgpu_doorbell"]] + [f"{AMD}/amdkfd/soc15_int.h"] + \
       [f"{AMDINC}/ivsrcid/{s}.h" for s in [f"gfx/irqsrcs_gfx_{x}_0" for x in ('9','11_0','12_0')] + [f"sdma0/irqsrcs_sdma0_{x}_0" for x in (4,5)]] + \
       [f"{AMDINC}/{s}.h" for s in ["v9_structs", "soc15_ih_clientid"]], args=inc, srcs=am_src, rules=kern_rules)
-    case "pm4_soc15": return load("am/pm4_soc15", [], [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/soc15d.h"], srcs=am_src)
-    case "pm4_nv": return load("am/pm4_nv", [], [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/nvd.h"], srcs=am_src)
-    case "sdma_4_0_0": return load("am/sdma_4_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/vega10_sdma_pkt_open.h"],
+    case "pm4_soc15": return load("am/pm4_soc15", [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/soc15d.h"], srcs=am_src)
+    case "pm4_nv": return load("am/pm4_nv", [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/nvd.h"], srcs=am_src)
+    case "sdma_4_0_0": return load("am/sdma_4_0_0", [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/vega10_sdma_pkt_open.h"],
                                    args=["-I/opt/rocm/include", "-x", "c++"], srcs=am_src)
-    case "sdma_5_0_0": return load("am/sdma_5_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/navi10_sdma_pkt_open.h"],
+    case "sdma_5_0_0": return load("am/sdma_5_0_0", [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/navi10_sdma_pkt_open.h"],
                                    args=["-I/opt/rocm/include", "-x", "c++"], srcs=am_src)
-    case "sdma_6_0_0": return load("am/sdma_6_0_0", [], [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/sdma_v6_0_0_pkt_open.h"],
+    case "sdma_6_0_0": return load("am/sdma_6_0_0", [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/sdma_v6_0_0_pkt_open.h"],
                                    args=["-I/opt/rocm/include", "-x", "c++"], srcs=am_src)
-    case "smu_v13_0_0": return load("am/smu_v13_0_0",[],[f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_0_ppsmc","smu13_driver_if_v13_0_0"]]
+    case "smu_v13_0_0": return load("am/smu_v13_0_0", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_0_ppsmc","smu13_driver_if_v13_0_0"]]
                                     +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
-    case "smu_v13_0_6": return load("am/smu_v13_0_6",[],[f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_6_ppsmc","smu_v13_0_6_pmfw", \
+    case "smu_v13_0_6": return load("am/smu_v13_0_6", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_6_ppsmc","smu_v13_0_6_pmfw", \
       "smu13_driver_if_v13_0_6"]] +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
-    case "smu_v13_0_12": return load("am/smu_v13_0_12",[],[f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_12_ppsmc","smu_v13_0_12_pmfw",
+    case "smu_v13_0_12": return load("am/smu_v13_0_12", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_12_ppsmc","smu_v13_0_12_pmfw",
       "smu13_driver_if_v13_0_6"]] +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
-    case "smu_v14_0_2": return load("am/smu_v14_0_2", [], [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v14_0_0_pmfw", "smu_v14_0_2_ppsmc",
+    case "smu_v14_0_2": return load("am/smu_v14_0_2", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v14_0_0_pmfw", "smu_v14_0_2_ppsmc",
                                     "smu14_driver_if_v14_0"]]+[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
+    # firmware hashes
+    case "fw":
+      def genfw(name, files, **kwargs): return "\n".join(["hashes = {"] + [f"  {p.name!r}: {hashlib.sha256(p.read_bytes()).hexdigest()!r},"
+                                                                           for f in files if (p:=pathlib.Path(f)).is_file()] + ["}"])
+      return load("am/fw", ["{}/amdgpu/psp_*_sos.bin", "{}/amdgpu/smu_*.bin", "{}/amdgpu/sdma_*.bin"] +
+                           [f"{{}}/amdgpu/gc_*_{x}.bin" for x in ["pfp", "me", "mec", "imu", "rlc"]], srcs=fw_src, gen=genfw)
     case _: raise AttributeError(f"no such autogen: {nm}")

tinygrad/runtime/autogen/am/fw.py

@@ -0,0 +1,106 @@
+hashes = {
+  'psp_13_0_0_sos.bin': 'b5592f46885585b935e013f46c949db8ff2f15c0b346caf70e7fcd2776623d13',
+  'psp_13_0_10_sos.bin': '0bcaaad9cd8578d3841ae69155a6bd4fc3ceae8f4fb5a6ba4f576e7ace94d1d9',
+  'psp_13_0_12_sos.bin': '89da90bf4286b38678b1fd175c78462a426afa3d258d15872cd14072d7098b9b',
+  'psp_13_0_14_sos.bin': 'a4f0d5f76d27b77409ec0b71d7cc6a848ddfd29f8c84f3003edf74ad3999fb7d',
+  'psp_13_0_6_sos.bin': '27657daa0f91ad8095d3610224a7de748b8b348a4cb211ecb5fccabe47369716',
+  'psp_13_0_7_sos.bin': 'ef1af0ecea38abbac6f85cce71789f19848c498d0cb8ef13748dab2d65b23c31',
+  'psp_14_0_2_sos.bin': '7b538448b57d4f9dd06b2eea90d4f86a16e65e3027cdecee8db71c2c5f1fa243',
+  'psp_14_0_3_sos.bin': '23bea01a0c6f36d00759d0765d46cb4cb4aa87398b2fbccacbf547a890c0bf51',
+  'smu_13_0_0.bin': '2ffac37fd8534965eeba19755db0e5ec80278213487dc4af0fbc8453befb64b1',
+  'smu_13_0_0_kicker.bin': '7f83656a2a89b7fce1c8a85e96d91cd8265a91fe883a7027f1a0ed18ced501de',
+  'smu_13_0_10.bin': 'daedb9cbdf48942be7ffe00d31b7c16bb36e11ff5a9d7495f218e95c07717b71',
+  'smu_13_0_14.bin': 'a4f36de75fdcecd8000246762e027b4be489b6787afea57675225b0b39d35625',
+  'smu_13_0_6.bin': 'ad7232264e8c57c2094244fbdd5a55d7a4575ffe9b44d229884bc0b6a44fb0b1',
+  'smu_13_0_7.bin': 'ccecc0fd0196b9613c920a51c2fd9436e739ff19dda5bdf74d97562387231732',
+  'smu_14_0_2.bin': '6951995d1d606f4dc60c895f19d34ed18aa40e62129f83d8510c45e8aa9ae2fc',
+  'smu_14_0_3.bin': 'df230947ddb7bbfd6e77d1280001db886e69adf2b2a448b47fa668a48bc0009f',
+  'smu_14_0_3_kicker.bin': '8ddc1da5b4e1619796c2cc81f19f388a35bf7d78bfe476cee559625589cb4dc7',
+  'sdma_4_4_2.bin': '456061b814268425843537da6f2191c8861d4e1a18d4c5d90c44ea6be18c78ff',
+  'sdma_4_4_4.bin': 'af47a2940e72b932d3e3a7e8f34f7a182624e5e433f7c56dff939ca5549cd33a',
+  'sdma_4_4_5.bin': '6127baabea3de7b18db3868c983b02c0fbf2cd75997f7f11241a5b1be27e5134',
+  'sdma_5_2_6.bin': '3a163db00eb7e4752be8adbd61cf7dd8f08d924e59a6f798ced7dfcd89f340ed',
+  'sdma_5_2_7.bin': '16fe80dc866b323e15a06f51646ef0f036878ad34da66921fcdb8167207d6b2b',
+  'sdma_6_0_0.bin': '0f3da6b211f376356335b41be07149f650c10cfa4e23f7e25d53836006ed11f5',
+  'sdma_6_0_1.bin': 'ff565d3c215a30737560d4e3df6fc2c637738407e91d212fb200fdfb185b6744',
+  'sdma_6_0_2.bin': '398380184bb69113ef4c8964a3b55f6184deb0c1ffd96c9683490a3eec3ba8f3',
+  'sdma_6_0_3.bin': '0e8a83513087db865ba926f8b65cfb003fd41098f707e178d7a7ae2941fed0b1',
+  'sdma_6_1_0.bin': '22e55d0ad5f0247a7f0fffc67cfd3161b39f24ad6062ff3c91ec7ff38bd7e1e1',
+  'sdma_6_1_1.bin': '74533a581b8e3e2743b3c9c803d0666405e80898c4a630acefed82cb6b516ba2',
+  'sdma_6_1_2.bin': '4fe04b0286ec739b0414e8aee17e62e85e691f0246d1d9b56bc18a1219072314',
+  'sdma_6_1_3.bin': '35c9ed7e3a237c0d4a83b4975c63b62488f72aeafbb648342f384618e103f66b',
+  'sdma_7_0_0.bin': 'beaafb53993a106edd392392d5896245ae2a957c6d0f495d0002eec72ad8ad38',
+  'sdma_7_0_1.bin': '73c29e1c1714ebc95d2221ba56e187910902891593010653bf9518937e414a59',
+  'gc_10_3_6_pfp.bin': '793d678427887a0e724c79e356440aec33e6d1301f2a4e63543500249ebec064',
+  'gc_10_3_7_pfp.bin': '3ae29aac3f424f7de97f82ce7158beba69509afb2dcbf1a428dc315df474a524',
+  'gc_11_0_0_pfp.bin': 'e175cb0f580a38c961a6f7366142c08e413995f57f78f39795368b15442df8a3',
+  'gc_11_0_1_pfp.bin': 'f5bf21dfbd9e72a30b4caf4704282c27854710c1b7c4affbb2a19530466b12a8',
+  'gc_11_0_2_pfp.bin': '001c4dec1119e29314d725cc1280fc4f0cd9cabdf61ea5ee2260cfd4e62ec141',
+  'gc_11_0_3_pfp.bin': '0488034c85be97125e39e860308d33c3f76a01df8250092a32d4d55acb2526fd',
+  'gc_11_0_4_pfp.bin': '5ae8b7bb6316f87ae8b978354c088e3bd8c890959382d72886377cda25b1ffd1',
+  'gc_11_5_0_pfp.bin': '0124f540871a7759fa8aaae046d458dfb34aeea12a1183ff962c3f1a33067d5a',
+  'gc_11_5_1_pfp.bin': '7794ea46d0d3cf9cb3f7938affbdf09dd7a9970340da5cd02b774cb393436d24',
+  'gc_11_5_2_pfp.bin': '55e64741de28c506524959f7f696713a72aafe46f49ccd827781d67a9475b386',
+  'gc_11_5_3_pfp.bin': 'ce805040fb347fddbc89b2715e66b446865dda9e2056a9b233269b72bc09c387',
+  'gc_12_0_0_pfp.bin': '16bfd64c10fe73b5e760055069a60e5841dba16c0ed4edb56c20d675e23901f6',
+  'gc_12_0_1_pfp.bin': '49efb319305c5fffd90ac1eef7d7a0bdec72998ecb5cf4526996311788a53dc3',
+  'gc_10_3_6_me.bin': '141b59faad3f2f1be16a2178833b7ca8e97519e1e844c8fda6689572c3767902',
+  'gc_10_3_7_me.bin': '9eb0b56e9bcc9dad5d53437b162226fcb37e5df102832260f1232832f3658edf',
+  'gc_11_0_0_me.bin': 'f8fba8a63dd4293b8fc1e4aab78b6fac630e575d1d62838c7996d9210f82aea1',
+  'gc_11_0_1_me.bin': '5030040b00955de94876341ec64ea43b96640413d7a03dc460a83c8386bf76e0',
+  'gc_11_0_2_me.bin': '0f21fd43f1dfbc6ccced9a2b3774de25c993c61a689aabab8b45333937b7945e',
+  'gc_11_0_3_me.bin': '3acb5061dba342ade81d329d1932f19ec01f0c5bf44e6e3568008a951a351bac',
+  'gc_11_0_4_me.bin': 'e4f1f6abcd213d54ad9e885d9f550083b0e2f67d983566015e8a53981e1cb155',
+  'gc_11_5_0_me.bin': '8f906b64d0a29503daa662c93ec44d076fcac11b78f70cd50ce0af2b500a05a6',
+  'gc_11_5_1_me.bin': '7e42602bcbaf1e511f8b4f6ed2246844ad1f6e351ce2b663d89062a7be263663',
+  'gc_11_5_2_me.bin': 'aae26255d8efff81e0e3bbcb727efb8b837d8e25fe85c708545f5328f1077b50',
+  'gc_11_5_3_me.bin': '93cd588348b16fe432609fe8da6e6b5da0a52da5c5884882aecf7b1001f72700',
+  'gc_12_0_0_me.bin': 'd7eba5197f2580f32b8256b1d9cb68e723e9e644293a34446a7913e3c093cba5',
+  'gc_12_0_1_me.bin': '365e7f193b39cbb10d3af44905fefaca0e9844721801755276baebac7b19c1ea',
+  'gc_10_3_6_mec.bin': '247943415658159704a21f670dd7b3e7cb2d2fc0c17b000a5098715979c8d95e',
+  'gc_10_3_7_mec.bin': 'ee58a523375bcf5b89400b32b801f95e182b632a26bce4f2bed5c07928d486dc',
+  'gc_11_0_0_mec.bin': '801a09c9bf06188260db9b51ad8f978f15d84c72ca91b90643a2ef8af4074776',
+  'gc_11_0_1_mec.bin': '6afadcb7504bb11bcc9d4a205cdf73f7934a615e28f178fcf7285971df2ccd05',
+  'gc_11_0_2_mec.bin': '0da0edee28c73a6fa1191f77853d380ec2503cbf43e0aaae4617f32f1f8a48fa',
+  'gc_11_0_3_mec.bin': '323cfa6658b6b5169830f852e2ff0552acae8dfb9e44b42c63de7b2900d3fd9e',
+  'gc_11_0_4_mec.bin': '5d89cf6b60354f3746c2cbd1ff0cb1a741556ca20d72745242cb69b553d0985c',
+  'gc_11_5_0_mec.bin': 'a01c324ab14ec89792449a621a541829b9af26865019027a411a14b910145dfa',
+  'gc_11_5_1_mec.bin': 'eab05719371caa68df09d4f7574e3958a3c4f5044ab3c7b0d2b214add0c6d1c4',
+  'gc_11_5_2_mec.bin': 'a374b2335802e24f8b9a3ce40000a1d37a52a14eb87099bebcc6680c27cc93e5',
+  'gc_11_5_3_mec.bin': '165025437cba80dd32c19ebbc83b756fa7adac7053ff7780ba4aa2f8089c6a3f',
+  'gc_12_0_0_mec.bin': '1931593440b8f9423580d9e2cdc5b34e7c682cdffe1ca4b74b0c2f6a0420236d',
+  'gc_12_0_1_mec.bin': 'f57541688a5108730bf210663f1137ffc2121f3acfe614a6de09ec1982c69a2f',
+  'gc_9_4_3_mec.bin': '3159176e72301fb88dc416721fb3d0ab82ece484cf93a43c3f37430c7e6673a1',
+  'gc_9_4_3_sjt_mec.bin': 'd19468dbb47849640bd0e6cdc8d7e25a3c8442c7ca2ca81357702e0d6baab50f',
+  'gc_9_4_4_mec.bin': '5004f73e43db2dd45e77d65942e33d4a69e7157618cfd23944c30f801c77a0f3',
+  'gc_9_4_4_sjt_mec.bin': '627a9e98102e70fe3bf0947eb764187f29f5e775d1130c7310e0ba5fc0502dbe',
+  'gc_9_5_0_mec.bin': 'c5eca4311a6f6e8f81cf41c2c46941d5dcf90789ee8326901da2dfc86ac14c31',
+  'gc_9_5_0_sjt_mec.bin': 'f162e509379288e3f3b1eead541b315c2262d625d433287ecd34ca185614d312',
+  'gc_11_0_0_imu.bin': 'b4f8fc056b45709a6abf48e7885fb1b4ab8d3cc092cbfa2c554a78564a6403bc',
+  'gc_11_0_1_imu.bin': 'ac71f4eec713fc35b4a1fe27531e3eb04edd81eeac2cef64df01ac50d8510805',
+  'gc_11_0_2_imu.bin': '9befca62b0b0cfd252c3df4a9edca295526f4d43821cd99a6326454995a6ca2d',
+  'gc_11_0_3_imu.bin': 'beaf704d5acdf4623456b0d0cbcea8b8e428058340cd922a259a9045f5c457a3',
+  'gc_11_0_4_imu.bin': 'ac71f4eec713fc35b4a1fe27531e3eb04edd81eeac2cef64df01ac50d8510805',
+  'gc_11_5_0_imu.bin': '469add57cafead90ab1953d6039cd8e39bc50dfd78aa5fd78f019ccf66a0af41',
+  'gc_11_5_1_imu.bin': '0aaca8a01b2237fca1b3c0cd082b5e12a271df334ff368bbf8e2be17f192b785',
+  'gc_11_5_2_imu.bin': 'fb684842839c61a0706a19df8e15eb8afc17e20c14267eb71e7d7d824c180acf',
+  'gc_11_5_3_imu.bin': 'fb684842839c61a0706a19df8e15eb8afc17e20c14267eb71e7d7d824c180acf',
+  'gc_12_0_0_imu.bin': 'aa15e5b3156bffc45e0c50bccbcd364fbd3f958531b695b7487a803d780b8328',
+  'gc_12_0_1_imu.bin': 'b3b301fb636efc77b63ce4d2ced0f90c851d03c19681852faa45598e6f5773fd',
+  'gc_10_3_6_rlc.bin': 'acfbac75c0dcfbfe40e222640ef17eb3dc8d206d30bc3863f275f2dd1cb132a5',
+  'gc_10_3_7_rlc.bin': 'a02585ebe3b36d942e883057119572d9497600c52fc65b8a523487eb65d874f2',
+  'gc_11_0_0_rlc.bin': 'dabd49039772d02f5fd5e48dc21d35ad52a6b1283b470dabca86ca159c4c7c8e',
+  'gc_11_0_1_rlc.bin': '86145719a58e9428562930c6b5ee3b6ced4701d34a80d0b4d84d6026c93134f2',
+  'gc_11_0_2_rlc.bin': 'b43eb2fd0600f50a1a5796bc9983d6b39b5c20960234920f5e89cb362193e0b8',
+  'gc_11_0_3_rlc.bin': '29b0b456f5b53076ddffa6f09de3bb697219e8e7b33504bf6c197e8b858426dc',
+  'gc_11_0_4_rlc.bin': '823573078b608108fbe4dd8176c396ec582632913db9c59a512d82b068f8eba0',
+  'gc_11_5_0_rlc.bin': '68cd85567f4f2f8d6b80db294988806d956bf826979c3597daccb71c7ee6aadd',
+  'gc_11_5_1_rlc.bin': '92731ecabbeb77865fb71787b4268dc738a58779f1190bdc2056482cb88a08f6',
+  'gc_11_5_2_rlc.bin': 'ef3a9209d3eccfbe18fce9e972c146ac283719798bb788096c176b796dc9aee5',
+  'gc_11_5_3_rlc.bin': '10a68940c6258d5818d9c05fd98eb0ccc8d5aee99b2769fbad30e5abd0d9327e',
+  'gc_12_0_0_rlc.bin': '6436b582734a413456fff3d3c7195e71cc9e78a7ed31ee21c83ffd6fae1ad186',
+  'gc_12_0_1_rlc.bin': '6ba4459532246a5c415d3cb33c9b1248294e48f67b827e2accb292a8d1a5c0ec',
+  'gc_9_4_3_rlc.bin': '5345d388712d547b0ae16f199ad5ccadb65643584b3efa7817049ddeb3fdcd12',
+  'gc_9_4_4_rlc.bin': 'e0c3585c72f8136670ca63e607fba32c1ae4948f493f13e33fc4d466bd6318a8',
+  'gc_9_5_0_rlc.bin': '9b1268f5751153fe57f527c9acb417bfa53ed42c9bc083c9d3da2ba61fe5fdc4',
+}

tinygrad/runtime/graph/cuda.py

@@ -3,7 +3,7 @@ from typing import Any, cast
 import tinygrad.runtime.autogen.cuda as cuda
 from tinygrad.runtime.support.c import init_c_var
 from tinygrad.device import Device, MultiBuffer
-from tinygrad.uop.ops import Ops
+from tinygrad.uop.ops import UOp, Ops
 from tinygrad.runtime.ops_cuda import CUDADevice, check, encode_args, cu_time_execution
 from tinygrad.engine.jit import MultiGraphRunner
 
@@ -44,7 +44,7 @@ class CUDAGraph(MultiGraphRunner):
     deps = self._access_resources(bufs, write, new_dependency=(node:=cuda.CUgraphNode()))
     return (cuda.CUgraphNode*len(deps))(*deps) if deps else None, node
 
-  def __call__(self, input_buffers, var_vals, wait=False, input_uops=None):
+  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False):
     # Update buffers in the c_args struct.
     for j in self.updatable:
       (_, params, c_args, is_copy), dev_idx = self.nodes[j], self.calls[j][0]

tinygrad/runtime/graph/hcq.py

@@ -260,7 +260,7 @@ class HCQGraph(MultiGraphRunner):
 
   def _dev_copy_queues(self, dev): return [q for (d, _), q in self.copy_queues.items() if d == dev]
 
-  def __call__(self, input_buffers: list[Buffer], var_vals: dict[str, int], wait=False, input_uops=None) -> float|None:
+  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None:
     # Map input buffers
     for dev in self.devices:
       for iidx, dev_idx in self.input_replace_map[dev]:

tinygrad/runtime/graph/metal.py

@@ -1,16 +1,17 @@
-from typing import Any
+from typing import Any, cast
 import ctypes, decimal
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import dedup, getenv, PROFILE
-from tinygrad.device import ProfileGraphEntry, ProfileGraphEvent
+from tinygrad.device import Buffer, Device, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.jit import GraphRunner, GraphException
-from tinygrad.runtime.ops_metal import wait_check, to_ns_str
+from tinygrad.runtime.ops_metal import MetalDevice, MetalAllocator, wait_check, to_ns_str
 from tinygrad.runtime.autogen import metal
 
 class MetalGraph(GraphRunner):
   def __init__(self, linear, input_uops=()):
     super().__init__(linear, input_uops)
+    self.dev = cast(MetalDevice, Device[self.device])
 
     # create metal batch exec
     icb_descriptor = metal.MTLIndirectCommandBufferDescriptor.new()
@@ -44,11 +45,11 @@ class MetalGraph(GraphRunner):
     self.all_resources = dedup(all_resources)
     self.all_pipelines = dedup(all_pipelines)
     self.command_buffer: Any = None
-    if len(self.vars): self.int_buf_view = self.dev.allocator._as_buffer(self.int_buf).cast('i')
+    if len(self.vars): self.int_buf_view = cast(MetalAllocator, self.dev.allocator)._as_buffer(self.int_buf).cast('i')
     self.range = metal.NSRange(0, len(self.calls))
     self.updatable = sorted({j for j,r in enumerate(self.uop_replace) if r} | self.var_vals_replace.keys() | self.launch_dims_replace.keys())
 
-  def __call__(self, input_buffers, var_vals, wait=False, input_uops=None):
+  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False):
     if self.command_buffer is not None and self.command_buffer in self.dev.mtl_buffers_in_flight: wait_check(self.command_buffer)
     # NOTE: old command buffer may not be inflight anymore
     if self.command_buffer is not None and PROFILE: self.collect_timestamps()
@@ -57,7 +58,7 @@ class MetalGraph(GraphRunner):
     for j in self.updatable:
       computeCommand = self.icb.indirectComputeCommandAtIndex(j)
       for pos, iidx in self.uop_replace[j]:
-        buf = input_uops[iidx].buffer
+        buf = cast(Buffer, input_uops[iidx].buffer)
         computeCommand.setKernelBuffer_offset_atIndex(buf._buf.buf, buf._buf.offset, pos)
         updated_bufs.append(buf._buf.buf)
 

tinygrad/runtime/ops_null.py

@@ -31,7 +31,7 @@ class NullAllocator(Allocator['NullDevice']):
   def _offset(self, buf, offset:int, size:int): pass
 
 class NullGraph(MultiGraphRunner):
-  def __call__(self, input_buffers, var_vals, wait=False, input_uops=None) -> float|None: return 1e-1
+  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None: return 1e-1
 
 class NullDevice(Compiled):
   def __init__(self, device:str):

tinygrad/runtime/ops_python.py

@@ -93,12 +93,14 @@ class PythonProgram:
         elif uop is Ops.INDEX:
           ret:list = []
           if isinstance(src_dtypes[0], ImageDType):
-            for m,ox,oy in zip(src_values[0], src_values[1][0], src_values[1][1]):
+            for m,ox,oy in zip(src_values[0], src_values[1], src_values[2]):
               if ox < 0 or ox >= src_dtypes[0].shape[1] or oy < 0 or oy >= src_dtypes[0].shape[0]: ret.append((m, None))
               else: ret.append((m, ox*4 + oy*src_dtypes[0].shape[1]*4))
+            gate = src_values[3] if len(src_values) == 4 else [True]*len(ret)
           else:
             for m,o in zip(src_values[0], src_values[1]): ret.append((m,o))
-          values[i] = [(m,o,g) for (m,o),g in zip(ret, src_values[2] if len(src_values) == 3 else [True]*len(ret))] # set the gate last
+            gate = src_values[2] if len(src_values) == 3 else [True]*len(ret)
+          values[i] = [(m,o,g) for (m,o),g in zip(ret, gate)] # set the gate last
         elif uop is Ops.CAST and isinstance(dtype, PtrDType):
           values[i] = src_values[0]
         elif uop is Ops.RANGE:

tinygrad/runtime/support/am/amdev.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
-import ctypes, collections, dataclasses, functools, hashlib, array
+import ctypes, collections, dataclasses, functools, hashlib, array, pathlib, sys
 from tinygrad.helpers import mv_address, getenv, DEBUG, fetch, lo32, hi32
 from tinygrad.runtime.autogen import pci
-from tinygrad.runtime.autogen.am import am
+from tinygrad.runtime.autogen.am import am, fw
 from tinygrad.runtime.support.amd import AMDReg, import_module, import_asic_regs
 from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace
 from tinygrad.runtime.support.system import PCIDevice
 from tinygrad.runtime.support.am.ip import AM_IP, AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
+if sys.version_info >= (3, 14): from compression import zstd
 
 AM_DEBUG = getenv("AM_DEBUG", 0)
 
@@ -108,8 +109,10 @@ class AMFirmware:
     self.descs += [self.desc(blob, hdr0.header.ucode_array_offset_bytes, hdr0.header.ucode_size_bytes, am.GFX_FW_TYPE_RLC_G)]
 
   def load_fw(self, fname:str, *headers, versioned_header:str|None=None):
-    fpath = fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/1e2c15348485939baf1b6d1f5a7a3b799d80703d/amdgpu/{fname}", subdir="fw")
-    blob = memoryview(bytearray(fpath.read_bytes()))
+    if (sys.version_info >= (3,14) and (p:=pathlib.Path("/lib/firmware/amdgpu")/f"{fname}.zst").is_file() and
+        hashlib.sha256(b:=zstd.decompress(p.read_bytes())).hexdigest() == fw.hashes[fname]): blob = memoryview(bytearray(b))
+    else: blob = memoryview(bytearray(fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/1e2c15348485939baf1b6d1f5a7a3b799d80703d/amdgpu/{fname}",
+                                            subdir="fw").read_bytes()))
     if AM_DEBUG >= 1: print(f"am {self.adev.devfmt}: loading firmware {fname}: {hashlib.sha256(blob).hexdigest()}")
     if versioned_header:
       chdr = am.struct_common_firmware_header.from_address(mv_address(blob))

tinygrad/runtime/support/autogen.py

@@ -99,8 +99,8 @@ arc_families = ['alloc', 'copy', 'mutableCopy', 'new']
 
 def normalize(a): return ("_" + n if keyword.iskeyword(n:=nm(a)) else n)
 
-def gen(name, dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, errno=False, anon_names={}, types={}, parse_macros=True, paths=[]):
-  macros, lines, anoncnt, types, objc, fns = [], [], itertools.count().__next__, {k:(v,True) for k,v in types.items()}, False, set()
+def gen(name, files, dll="", args=[], prolog=[], rules=[], epilog=[], recsym=False, errno=False, anon_names={}, types={}, macros=True, paths=[]):
+  extras, lines, anoncnt, types, objc, fns = [], [], itertools.count().__next__, {k:(v,True) for k,v in types.items()}, False, set()
 
   # ctypes automatically "unboxes" simple types
   def typehint(ty) -> str:
@@ -240,13 +240,13 @@ def gen(name, dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False,
             if clang.CXCursor_NSReturnsRetained in attrs(c): lines.append(f"{nm(c)} = objc.returns_retained({nm(c)})")
           case (clang.CXCursor_StructDecl | clang.CXCursor_UnionDecl | clang.CXCursor_TypedefDecl | clang.CXCursor_EnumDecl
                 | clang.CXCursor_ObjCInterfaceDecl): tname(clang.clang_getCursorType(c))
-          case clang.CXCursor_MacroDefinition if parse_macros and nm(c) and len(toks:=Tokens(c)) > 1:
+          case clang.CXCursor_MacroDefinition if macros and nm(c) and len(toks:=Tokens(c)) > 1:
             if nm(toks[1])=='(' and clang.clang_equalLocations(clang.clang_getRangeEnd(extent(toks[0])), clang.clang_getRangeStart(extent(toks[1]))):
               it = iter(toks[1:])
               _args = [nm(t) for t in itertools.takewhile(lambda t:nm(t)!=')', it) if clang.clang_getTokenKind(t) == clang.CXToken_Identifier]
               if len(body:=list(it)) == 0: continue
-              macros += [f"{nm(c)} = lambda{' ' * bool(_args)}{','.join(_args)}: {readext(f,loc(body[0]),clang.clang_getRangeEnd(extent(toks[-1])))}"]
-            else: macros += [f"{nm(c)} = {readext(f, loc(toks[1]), clang.clang_getRangeEnd(extent(toks[-1])))}"]
+              extras += [f"{nm(c)} = lambda{' ' * bool(_args)}{','.join(_args)}: {readext(f,loc(body[0]),clang.clang_getRangeEnd(extent(toks[-1])))}"]
+            else: extras += [f"{nm(c)} = {readext(f, loc(toks[1]), clang.clang_getRangeEnd(extent(toks[-1])))}"]
           case clang.CXCursor_VarDecl if clang.clang_getCursorLinkage(c) == clang.CXLinkage_Internal:
             ty = clang.clang_getCursorType(c)
             if (ty.kind == clang.CXType_ConstantArray and clang.clang_getCanonicalType(clang.clang_getArrayElementType(ty)).kind in ints and
@@ -254,10 +254,10 @@ def gen(name, dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False,
               cs = children(init)
               if all(re.match(r"\[.*\].*=", readext(f, extent(ch))) for ch in cs):
                 items = ','.join(f'{readext(f, extent(next(it:=iter(children(ch)))))}:{readext(f, extent(next(it)))}' for ch in cs)
-                macros += [f"{nm(c)} = {{{items}}}"]
-              else: macros += [f"{nm(c)} = ({','.join(readext(f, extent(ch)) for ch in cs)},)"]
-            elif clang.clang_getCanonicalType(ty).kind in ints: macros += [f"{nm(c)} = {readext(f, extent(children(c)[-1]))}"]
-            else: macros += [f"{nm(c)} = {tname(ty)}({readext(f, extent(children(c)[-1]))})"]
+                extras += [f"{nm(c)} = {{{items}}}"]
+              else: extras += [f"{nm(c)} = ({','.join(readext(f, extent(ch)) for ch in cs)},)"]
+            elif clang.clang_getCanonicalType(ty).kind in ints: extras += [f"{nm(c)} = {readext(f, extent(children(c)[-1]))}"]
+            else: extras += [f"{nm(c)} = {tname(ty)}({readext(f, extent(children(c)[-1]))})"]
           case clang.CXCursor_VarDecl if clang.clang_getCursorLinkage(c) == clang.CXLinkage_External and dll:
             lines.append(f"try: {nm(c)} = {tname(clang.clang_getCursorType(c))}.in_dll(dll, '{nm(c)}') # type: ignore\n" +
                          "except (ValueError,AttributeError): pass")
@@ -272,16 +272,16 @@ def gen(name, dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False,
                     "from typing import Literal, TypeAlias", "from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR",
                     "from tinygrad.runtime.support import c", *prolog, *(["from tinygrad.runtime.support import objc"]*objc),
                     *([f"dll = c.DLL('{name}', {dll}{f', {paths}'*bool(paths)}{', use_errno=True'*errno})"] if dll else []), *lines]) + '\n'
-  macros = [f"{r} # type: ignore" if "lambda" in r else r for m in macros
-            if (r:=functools.reduce(lambda s,r:re.sub(r[0], r[1], s), rules + base_rules, m))]
+  extras = [f"{r} # type: ignore" if "lambda" in r else r
+            for m in extras if (r:=functools.reduce(lambda s,r:re.sub(r[0], r[1], s), rules + base_rules, m))]
   while True:
     try:
-      exec(main + '\n'.join(macros), {})
+      exec(main + '\n'.join(extras), {})
       break
     except (SyntaxError, NameError, TypeError) as e:
-      macrono = unwrap(e.lineno if isinstance(e, SyntaxError) else unwrap(unwrap(e.__traceback__).tb_next).tb_lineno) - main.count('\n') - 1
-      assert macrono >= 0 and macrono < len(macros), f"error outside macro range: {e}"
-      print(f"skipping {macros[macrono]}: {e}")
-      del macros[macrono]
+      extrano = unwrap(e.lineno if isinstance(e, SyntaxError) else unwrap(unwrap(e.__traceback__).tb_next).tb_lineno) - main.count('\n') - 1
+      assert extrano >= 0 and extrano < len(extras), f"error outside extra range: {e}"
+      print(f"skipping {extras[extrano]}: {e}")
+      del extras[extrano]
     except Exception as e: raise Exception("parsing failed") from e
-  return main + '\n'.join(macros + epilog)
+  return main + '\n'.join(extras + epilog)

tinygrad/schedule/rangeify.py

@@ -198,7 +198,7 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.REDUCE, name="reduce", src=(UPat.var("x"),)),
    lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if 0 in x.shape and 0 not in reduce.shape else None),
   # handle size 0
-  (UPat(GroupOp.All-{Ops.SINK, Ops.STACK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and 0 in x.shape else None),
+  (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and 0 in x.shape else None),
 ])
 
 # *****************
@@ -545,8 +545,6 @@ pm_add_range_tags = PatternMatcher([
 def split_store(x:UOp) -> UOp|None:
   # if we have any open ranges here, we don't split
   if x.ranges: return None
-  # raw STORE (not from bufferize_to_store) should be processed through its END wrapper, not independently
-  #if x.op is Ops.STORE and x.src[0]._shape is not None: return None
 
   # local kernel rewrite
   lctx = LocalAddBufferContext()

tinygrad/tensor.py

@@ -551,6 +551,20 @@ class Tensor(OpMixin):
     """
     Tensor._seed, Tensor._device_seeds, Tensor._device_rng_counters = seed, {}, {}
 
+  @staticmethod
+  def _next_counter(device:str, num:int) -> tuple[Tensor, Tensor]:
+    if device not in Tensor._device_seeds:
+      seed = [int.from_bytes(hashlib.sha256(len(Tensor._device_seeds).to_bytes(4, "big")).digest(), "big"), Tensor._seed]
+      Tensor._device_seeds[device] = Tensor(seed, device=device, dtype=dtypes.uint32, requires_grad=False)
+      Tensor._device_rng_counters[device] = Tensor([0, 0], device=device, dtype=dtypes.uint32, requires_grad=False)
+    counter = Tensor._device_rng_counters[device]
+    new_low = counter[0:1] + (num & 0xffffffff)
+    new_high = counter[1:2] + (num >> 32) + (new_low < counter[0])
+    counter.assign(new_low.cat(new_high))
+    low = counter[0:1] - (num & 0xffffffff)
+    high = counter[1:2] - (num >> 32) - (counter[0] < (num & 0xffffffff))
+    return Tensor._device_seeds[device], low.cat(high)
+
   @staticmethod
   def rand(*shape, device:str|None=None, dtype:DTypeLike|None=None, contiguous:bool=True, **kwargs) -> Tensor:
     """
@@ -574,22 +588,8 @@ class Tensor(OpMixin):
     # if shape has 0, return zero tensor
     if (numel := prod(shape)) == 0: return Tensor.zeros(shape, device=device, dtype=dt, **kwargs)
     num = ceildiv(numel * dt.itemsize, 4)
-
-    # generate per device seeds and rng counter if we haven't seen this device yet
-    if device not in Tensor._device_seeds:
-      Tensor._device_seeds[device] = Tensor(
-        [int.from_bytes(hashlib.sha256(len(Tensor._device_seeds).to_bytes(4, "big")).digest(), "big"), Tensor._seed],
-        device=device, dtype=dtypes.uint32, requires_grad=False)
-      Tensor._device_rng_counters[device] = Tensor([0, 0], device=device, dtype=dtypes.uint32, requires_grad=False).contiguous()
-
-    # increment rng counter for devices
-    new_low = Tensor._device_rng_counters[device][0:1] + (num & 0xffffffff)
-    new_high = Tensor._device_rng_counters[device][1:2] + (num >> 32) + (new_low < Tensor._device_rng_counters[device][0]).cast(dtypes.uint32)
-    Tensor._device_rng_counters[device].assign(new_low.cat(new_high))
-
-    low = Tensor._device_rng_counters[device][0:1] - (num & 0xffffffff)
-    high = Tensor._device_rng_counters[device][1:2] - (num >> 32) - (Tensor._device_rng_counters[device][0] < (num & 0xffffffff)).cast(dtypes.uint32)
-    bits = Tensor.random_bits(Tensor._device_seeds[device], low.cat(high), num)
+    key, counter = Tensor._next_counter(device, num)
+    bits = Tensor.random_bits(key, counter, num)
     out = Tensor._bits_to_rand(bits, shape, dt).requires_grad_(kwargs.get("requires_grad"))
     return out.contiguous() if contiguous else out
 
@@ -692,7 +692,7 @@ class Tensor(OpMixin):
   def randint(*shape, low=0, high=10, dtype=dtypes.int32, **kwargs) -> Tensor:
     """
     Creates a tensor with the given shape, filled with random integer values generated uniformly from the interval `[low, high)`.
-    If `dtype` is not specified, the default type is used.
+    Requires `low < high`. If `dtype` is not specified, the default type is used.
 
     You can pass in the `device` keyword argument to control device of the tensor.
     Additionally, all other keyword arguments are passed to the constructor of the tensor.
@@ -704,12 +704,14 @@ class Tensor(OpMixin):
     """
     if not all_int([low, high]): raise TypeError(f"{low=} and {high=} must be integers")
     if not dtypes.is_int(dtype := to_dtype(dtype)): raise TypeError(f"{dtype=} must be int")
+    if low >= high: raise ValueError(f"Tensor.randint requires low < high, got {low=}, {high=}")
     return Tensor.uniform(*shape, low=low, high=high, dtype=dtype, **kwargs)
 
   @staticmethod
   def normal(*shape, mean=0.0, std=1.0, requires_grad:bool|None=None, **kwargs) -> Tensor:
     """
     Creates a tensor with the given shape, filled with random values from a normal distribution with the given `mean` and standard deviation `std`.
+    Requires `std >= 0`.
 
     You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
     Additionally, all other keyword arguments are passed to the constructor of the tensor.
@@ -719,12 +721,14 @@ class Tensor(OpMixin):
     print(Tensor.normal(2, 3, mean=10, std=2).numpy())
     ```
     """
+    if std < 0: raise ValueError(f"Tensor.normal requires std >= 0, got {std=}")
     return (std * Tensor.randn(*shape, **kwargs) + mean).requires_grad_(requires_grad)
 
   @staticmethod
   def uniform(*shape, low=0.0, high=1.0, dtype:DTypeLike|None=None, requires_grad:bool|None=None, **kwargs) -> Tensor:
     """
     Creates a tensor with the given shape, filled with random values from a uniform distribution over the interval `[low, high)`.
+    Requires `low < high`.
 
     You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
     Additionally, all other keyword arguments are passed to the constructor of the tensor.
@@ -734,6 +738,8 @@ class Tensor(OpMixin):
     print(Tensor.uniform(2, 3, low=2, high=10).numpy())
     ```
     """
+    if not all_int(shape:=argfix(*shape)) or not all(s >= 0 for s in shape): raise ValueError(f"invalid input {shape=}")
+    if low >= high: raise ValueError(f"Tensor.uniform requires low < high, got {low=}, {high=}")
     return (((high-low) * Tensor.rand(*shape, **kwargs)).cast(dtype or dtypes.default_float) + low).requires_grad_(requires_grad)
 
   @staticmethod
@@ -816,19 +822,27 @@ class Tensor(OpMixin):
     """
     Returns a tensor with `num_samples` indices sampled from a multinomial distribution weighted by `self`.
 
-    NOTE: `replacement=False` for `num_samples > 1` is not supported yet.
     ```python exec="true" source="above" session="tensor" result="python"
     Tensor.manual_seed(42)
     t = Tensor([1, 2, 3, 4])
     print(t.multinomial(20, replacement=True).numpy())
     ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    Tensor.manual_seed(42)
+    t = Tensor([1, 2, 3, 4])
+    print(t.multinomial(3, replacement=False).numpy())
+    ```
     """
     assert 1 <= self.ndim <= 2 and num_samples > 0, f"{self.ndim=} must be 1 or 2 dim, {num_samples=} must be positive"
-    assert replacement or num_samples == 1, "no replacement only supports num_samples = 1"
     weight = self.unsqueeze(0) if self.ndim == 1 else self
-    cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1)
-    unif_samples = Tensor.rand(num_samples, cdf.shape[0], 1).to(self.device)
-    indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0))
+    assert replacement or num_samples <= weight.shape[1], "no replacement samples must not exceed population size"
+    if replacement or num_samples == 1:
+      cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1)
+      unif_samples = Tensor.rand(num_samples, cdf.shape[0], 1).to(self.device)
+      indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0))
+    else:
+      # Efraimidis–Spirakis
+      indices = (weight.rand_like(dtype=dtypes.float32).log2() / weight).topk(num_samples, dim=1)[1]
     return (indices.squeeze(0) if self.ndim == 1 else indices).cast(dtypes.int32)
 
   # ***** toposort and backward pass *****

tinygrad/uop/ops.py

@@ -212,7 +212,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     match self.op:
       # late ops don't have shape
       case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
-           Ops.UNROLL | Ops.CONTRACT | Ops.SINK | Ops.END | \
+           Ops.STACK | Ops.GEP | Ops.UNROLL | Ops.CONTRACT | Ops.SINK | Ops.END | Ops.REWRITE_ERROR | \
            Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION:
         return None
 
@@ -228,39 +228,22 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
         return inner_shape
 
       case Ops.CAST:
-        if self.dtype.count > 1:
-          return (self.dtype.count,)
-
         # when PTX casts from ptr to non ptr, remove the shape
         if isinstance(self.src[0].dtype, PtrDType) and not isinstance(self.src[0].dtype, ImageDType) and not isinstance(self.dtype, PtrDType):
           return None
 
-      case Ops.STACK: return (len(self.src),)
-      case Ops.GEP:
-        if len(self.arg) > 1: return (len(self.arg),)
-        #assert len(self.arg) == 1
-        return ()
       case Ops.INDEX:
-        shp = []
-        for s in self.src[1:]: shp.extend(list(s.shape))
-        return tuple(shp)
-        """
         # non pointer index doesn't have a shape
         if not isinstance(self.dtype, PtrDType): return None
         # fully indexed doesn't have a shape. TODO: remove this
         if self.src[0]._shape is None or len(self.src[1:]) == len(self.src[0].shape): return None
         # pointer index
         return self.src[0].shape[len(self.src[1:]):]
-        """
 
       # some ops init the shape
-      case Ops.DEFINE_VAR | Ops.BIND | Ops.RANGE | Ops.SPECIAL: return ()
-      case Ops.CONST:
-        if self.dtype.count > 1: return (self.dtype.count,)
-        return ()
-      case Ops.VCONST: return (len(self.arg),)
+      case Ops.CONST | Ops.DEFINE_VAR | Ops.BIND | Ops.RANGE | Ops.SPECIAL: return ()
       # TODO: VCONST should have the shape of the arg
-      #case Ops.VCONST: return ()
+      case Ops.VCONST: return ()
       case Ops.BUFFER: return (self.arg,)
       case Ops.BUFFER_VIEW: return (self.arg[0],)
       case Ops.CUSTOM_FUNCTION: return None
@@ -272,8 +255,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
         if len(self.src) >= 1: return tuple(self.src[0].sgep(i) for i in range(self.src[0].dtype.count))
         return None
 
-      # SHAPED_WMMA output shape = accumulator shape (src[2])
-      case Ops.SHAPED_WMMA: return self.src[2]._shape
+      # wmma output shape = accumulator shape (src[2])
+      case Ops.WMMA | Ops.SHAPED_WMMA: return self.src[2]._shape
 
       # passthrough ops
       case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.LOAD:
@@ -295,10 +278,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
 
     # movement ops change the shape
     # NOTE: ssimplify is required because the shape needs to be canonical for broadcasting and same shape checking
-    if self.op in GroupOp.Movement.union({Ops.MULTI, Ops.REDUCE, Ops.WMMA}):
+    if self.op in GroupOp.Movement.union({Ops.MULTI, Ops.REDUCE}):
       ps = self.src[0]._shape
-      # TODO: WMMA is used for both axis WMMA and op WMMA. fix this and remove this hack. tested by BERT on AMD LLVM
-      if ps is None and self.op is Ops.WMMA: return None
       if ps is None: raise RuntimeError(f"movement op {self.op} requires shape")
       match self.op:
         case Ops.RESHAPE:
@@ -325,7 +306,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
           if len(ps) != len(self.marg) or not all(isinstance(x, bool) for x in self.marg): raise ValueError(f"bad flip on {ps}, {self.marg}")
           return ps
         case Ops.MULTI: return tuple(s*len(self.device) if a == self.axis else s for a,s in enumerate(ps))
-        case Ops.REDUCE | Ops.WMMA:
+        case Ops.REDUCE:
           axis_arg = self.arg[1] if self.op is Ops.REDUCE else self.arg[7]
           if not isinstance(axis_arg, tuple) or not all(isinstance(x, int) and x>=0 and x<len(ps) for x in axis_arg):
             raise ValueError(f"invalid type for axis: {axis_arg}")
@@ -509,7 +490,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     ret = UOp(Ops.VCONST if isinstance(b, tuple) else Ops.CONST, dtype,
               arg=dtype.const(b),
               src=(UOp(Ops.DEVICE, arg=device),) if device is not None else ())
-    return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None and shape != ret.shape else ret
+    return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None and ret.shape != shape else ret
   @staticmethod
   def unique_const(fill_value:ConstType, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None,  # type: ignore[override]
                    shape:tuple[sint, ...]|None=None, unique=True):
@@ -517,7 +498,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     assert not isinstance(fill_value, (UOp, tuple)), "unique const only works on numbers"
     ret = UOp.const(to_dtype(dtype) if dtype is not None else dtypes.from_py(fill_value), fill_value, canonicalize_device(device))
     ret = ret.replace(src=(UOp.unique(None if unique is True else unique),) + ret.src)
-    return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None else ret
+    return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None and ret.shape != shape else ret
   @staticmethod
   def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.weakint, src=(), **kwargs):
     return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs)
@@ -1155,8 +1136,8 @@ class UPat(OpMixin):
 
   # copied from UOp
   def sink(self, *srcs:UPat|None, **kwargs): return UPat(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
-  def index(self, idx:UPat, valid:UPat|None=None, **kwargs):
-    return UPat(Ops.INDEX, self.match_dtype, (self,idx,valid) if valid is not None else (self,idx), **kwargs)
+  def index(self, *srcs:UPat|None, **kwargs):
+    return UPat(Ops.INDEX, self.match_dtype, (self,)+tuple(x for x in srcs if x is not None), **kwargs)
   def cast(self, dtype=None, **kwargs):
     if dtype is not None and self.match_dtype == (dtype,): return self
     return UPat(Ops.CAST, dtype, (self,), **kwargs)
@@ -1525,6 +1506,14 @@ def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(dtype, x)
 def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape)
 
 def select_dtype(u): return (dtypes.long if u.overflows(dtypes.int32) else dtypes.int).vec(u.dtype.count)
+def strip_index_casts(idx:UOp):
+  new_src = (idx.src[0],) + tuple(s.src[0] if s.op is Ops.CAST and s.dtype.scalar() is dtypes.weakint and
+                                  s.src[0].dtype.scalar() in dtypes.ints else s for s in idx.src[1:])
+  return new_src[0].index(*new_src[1:], ptr=True) if new_src != idx.src else None
+def lower_image_index(idx:UOp):
+  if not isinstance(idx.src[0].dtype, ImageDType) or len(idx.src) not in (3, 4): return None
+  new_src = (idx.src[0],) + tuple(s if s.dtype == dtypes.int else s.cast(dtypes.int) for s in idx.src[1:3]) + idx.src[3:]
+  return idx.replace(src=new_src) if new_src != idx.src else None
 pm_lower_index_dtype = PatternMatcher([
   # There are no Unary ops at this point in symbolic, those are introduced later
   (UPat(GroupOp.Binary, name="u", src=(UPat.var("x").cast(dtypes.weakint), UPat.var("y").cast(dtypes.weakint))), lambda u,x,y:
@@ -1545,9 +1534,9 @@ pm_lower_index_dtype = PatternMatcher([
   # lower Invalid
   (UPat.var("buf").index(UPat.var("cond").where(UPat.var("idx"), UPat(Ops.CONST, arg=Invalid))), lambda buf,idx,cond: buf.index(idx, cond, ptr=True)),
   # remove hanging casts
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast()),), lambda buf,idx: buf.index(idx, ptr=True)),
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast(), UPat.var("valid"))),
-   lambda buf,idx,valid: buf.index(idx, valid, ptr=True)),
+  (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="idx"), strip_index_casts),
+  # images are indexed with separate int x/y coordinates
+  (UPat(Ops.INDEX, src=(UPat(), UPat(), UPat()), allow_any_len=True, name="idx"), lower_image_index),
   (UPat((Ops.SINK, Ops.NOOP, Ops.END), name="n"),
    lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.weakint else s for s in n.src))),
   # vectorized indexes (ie. images) must be int

tinygrad/uop/spec.py

@@ -23,6 +23,11 @@ def validate_index(buf:UOp, idx:UOp, gate:UOp|None=None):
   from tinygrad.uop.validate import validate_index_with_z3
   return validate_index_with_z3(sz, idx, gate)
 
+def validate_image_index(buf:UOp, idx_x:UOp, idx_y:UOp, gate:UOp|None=None):
+  if not isinstance(buf.dtype, ImageDType): return None
+  return idx_x.dtype in dtypes.ints+(dtypes.weakint,) and idx_y.dtype in dtypes.ints+(dtypes.weakint,) and \
+    (gate is None or gate.dtype == dtypes.bool)
+
 # four specs:
 #   shared_spec  -- usable anywhere
 #   tensor_spec  -- usable in tensor graph
@@ -175,6 +180,7 @@ shared_codegen_spec = PatternMatcher([
 
   # LOAD(idx) / STORE(idx, val)
   (UPat().index(UPat()).or_casted().load(), lambda: True),
+  (UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted().load(), lambda: True),
   (UPat().index(UPat(), UPat(dtype=dtypes.bool)).or_casted().load(), lambda: True),  # gated load (alt added in program_spec)
   (UPat(Ops.INDEX).or_casted().store(UPat()), lambda: True),
 
@@ -187,6 +193,9 @@ shared_codegen_spec = PatternMatcher([
   # INDEX (2-arg and 3-arg with bool gate)
   (UPat(GroupOp.Defines|{Ops.AFTER}, name="buf").index(UPat.var("idx")), validate_index),
   (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines|{Ops.AFTER}, name="buf"), UPat.var("idx"), UPat.var("gate", dtype=dtypes.bool))), validate_index),
+  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines|{Ops.AFTER}, name="buf"), UPat.var("idx_x"), UPat.var("idx_y"))), validate_image_index),
+  (UPat(Ops.INDEX, src=(UPat(GroupOp.Defines|{Ops.AFTER}, name="buf"), UPat.var("idx_x"), UPat.var("idx_y"),
+                         UPat.var("gate", dtype=dtypes.bool))), validate_image_index),
 
   # SPECIAL
   (UPat(Ops.SPECIAL, src=(UPat.var("x", (dtypes.weakint, dtypes.int32)),), name="s"), lambda s,x: s.dtype == x.dtype and isinstance(s.arg, str)),
@@ -237,6 +246,7 @@ tensor_spec = PatternMatcher([
 program_spec = PatternMatcher([
   # LOAD (idx, alt_value), LOAD can have an alt value, but only if the index has a gate
   (UPat().index(UPat(), UPat(dtype=dtypes.bool)).or_casted().load(UPat()), lambda: True),
+  (UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(), UPat(dtype=dtypes.bool))).or_casted().load(UPat()), lambda: True),
 
   # END closes ranges
   (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE)), dtype=dtypes.void), lambda: True),

tinygrad/uop/symbolic.py

@@ -2,7 +2,7 @@
 import math, struct
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
-from tinygrad.dtype import ConstType, dtypes, PtrDType, can_lossless_cast, Invalid
+from tinygrad.dtype import ConstType, dtypes, PtrDType, ImageDType, can_lossless_cast, Invalid
 from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, unwrap, IMAGE, dedup
 from tinygrad.uop.decompositions import threefry2x32, xpow
 from tinygrad.uop.divandmod import div_and_mod_symbolic
@@ -398,6 +398,13 @@ def where_on_load(cond:UOp, buf:UOp, idx:UOp, or_cast:UOp) -> UOp|None:
   ret_idx = idx.cast(or_cast.dtype) if or_cast.op is Ops.CAST else idx
   return UOp.const(dtypes.bool, True).uprod(*keep).where(ret_idx, ret_idx.const_like(0))
 
+def index_with_gate(index:UOp, gate:UOp) -> UOp:
+  if isinstance(index.src[0].dtype, ImageDType) and len(index.src) in (3, 4):
+    return index.replace(src=index.src[:3] + (gate if len(index.src) == 3 else index.src[3] & gate,))
+  if len(index.src) == 3 and index.src[2].dtype == dtypes.bool:
+    return index.replace(src=(index.src[0], index.src[1], index.src[2] & gate))
+  return index.replace(src=(index.src[0], gate.where(index.src[1], UOp.invalid())) + index.src[2:])
+
 # where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer
 pm_move_where_on_load = PatternMatcher([
   (UPat.var("cond").where(UPat.var("buf").index(UPat.var("idx")).or_casted("or_cast"), 0), where_on_load),
@@ -443,14 +450,19 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
   (UPat.store(UPat(Ops.INDEX, name="index"), UPat.load(UPat(Ops.INDEX, name="index"))), lambda index: UOp(Ops.NOOP)),
   (UPat.store(UPat(Ops.INDEX, name="index"), UPat.var("gate").where(UPat.var("alt"),
                                                                     UPat.load(UPat(Ops.INDEX, name="index"))), allow_any_len=True, name="store"),
-   lambda index, gate, alt, store: UOp.store(index.src[0].index(gate.where(index.src[1], UOp.invalid())), alt, *store.src[2:])),
+   lambda index, gate, alt, store: UOp.store(index_with_gate(index, gate), alt, *store.src[2:])),
   # fold gated LOAD/STORE
-  (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(),), allow_any_len=True, name="x"),
-    lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0
+  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(), UPat.const(dtypes.bool, False)), name="index"),), allow_any_len=True, name="x"),
+   lambda index,x: UOp(Ops.NOOP) if isinstance(index.src[0].dtype, ImageDType) else None),
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(), UPat.const(dtypes.bool, False)), name="index"),), allow_any_len=True, name="x"),
+   lambda index,x: (x.src[1] if len(x.src) > 1 else x.const_like(0)) if isinstance(index.src[0].dtype, ImageDType) else None),
+  (UPat(Ops.STORE, src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(),), allow_any_len=True, name="x"), lambda x: UOp(Ops.NOOP)),
+  (UPat(Ops.LOAD, src=(UPat().index(UPat.const(dtypes.weakint, Invalid)).or_casted(),), allow_any_len=True, name="x"),
+    lambda x: x.src[1] if len(x.src) > 1 else x.const_like(0)), # invalid load produces 0, or the alt value if we have one
   (UPat(Ops.STORE, src=(UPat(), invalid_pat), allow_any_len=True), lambda i: UOp(Ops.NOOP)),
   # store of where with invalid -> gated store
   (UPat(Ops.STORE, src=(UPat(Ops.INDEX, name="index"), UPat.var("cond").where(UPat.var("val"), invalid_pat)), allow_any_len=True, name="store"),
-   lambda index, cond, val, store, i: UOp.store(index.src[0].index(cond.where(index.src[1], UOp.invalid())), val, *store.src[2:])),
+   lambda index, cond, val, store, i: UOp.store(index_with_gate(index, cond), val, *store.src[2:])),
   ((UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()),  # 1/(x^c) -> (1/x)^c
   ((UPat.var("x") * UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()*x.reciprocal()),
   ((UPat.var("x") * UPat.cvar("c")).reciprocal(), lambda x,c: x.reciprocal()*c.reciprocal()), # 1/(x*c) -> (1/c)*(1/x)

tinygrad/viz/serve.py

@@ -116,7 +116,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]:
     if u.op is Ops.VCONST and u.dtype.scalar() == dtypes.weakint and u is not x: excluded.add(u)
     if u.op is Ops.STACK and len(u.src) == 0: excluded.add(u)
     # exclude RESHAPE/EXPAND that only serve to broadcast a CONST
-    #if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u)
+    if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u)
   for u in toposort:
     if u in excluded: continue
     argst = codecs.decode(str(u.arg), "unicode_escape")