x

* hcq2: profile slots

* x

.github/workflows/benchmark.yml

@@ -536,7 +536,7 @@ jobs:
   testcommausbgpubenchmark:
     name: UsbGPU Benchmark (comma)
     runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 10
+    timeout-minutes: 14
     defaults:
       run:
         shell: bash -e -o pipefail {0}

.github/workflows/platform.yml

@@ -70,9 +70,10 @@ jobs:
     - name: Run pytest (amd)
       env:
         DEV: MOCKKFD+AMD
+        HCQ_RUNTIME_DEV: PYTHON
         FORWARD_ONLY: 1
       run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
+        python3 -m pytest -n=auto test/device/test_hcq2.py test/test_tiny.py --durations=20
     - name: Run pytest (ptx)
       env:
         DEV: "MOCK+NV:PTX"

.github/workflows/test.yml

@@ -253,7 +253,7 @@ jobs:
         deps: testing_unit
         llvm: 'true'
     - name: Test SPEC=2
-      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py -k "not test_setitem_big" -k "not test_conv2d_ceildiv_edge_case" --splits 2 --group ${{ matrix.group }}
+      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --splits 2 --group ${{ matrix.group }}
 
   fuzzing:
     name: Fuzzing
@@ -478,6 +478,7 @@ jobs:
     timeout-minutes: 20
     env:
       DEV: MOCKKFD+AMD
+      HCQ_RUNTIME_DEV: PYTHON
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -504,7 +505,7 @@ jobs:
       - name: Run AMD renderer tests (AMD:LLVM)
         run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
       - name: Run SQTT profiling tests
-        run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: SQTT_BUFFER_SIZE=16 VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
       - name: Run AMD emulated tests on NULL backend
         env:
           AMD: 0

test/amd/test_asm_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, compile_linear
+from tinygrad.engine.realize import run_linear, estimate_uop, lower_and_compile
 from tinygrad.renderer import Estimates
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv
@@ -169,7 +169,7 @@ class TestAsmKernel(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 = compile_linear(a.schedule_linear())
+    linear = lower_and_compile(a.schedule_linear())
     est = estimate_uop(linear.src[-1])
     self.assertEqual(est.ops, a.numel())
     self.assertEqual(est.mem, a.nbytes()*2)

test/amd/test_mockgpu_invalid.py

@@ -1,51 +1,29 @@
 #!/usr/bin/env python3
 """Test that invalid instructions raise exceptions through the mock GPU stack."""
-import unittest, subprocess, os, sys, time
+import unittest, subprocess, os, sys
 
 class TestMockGPUInvalidInstruction(unittest.TestCase):
   def test_unsupported_instruction_raises(self):
     """Test that unsupported instructions raise immediately through the full MOCKGPU stack."""
     test_code = '''
-import struct
-from dataclasses import replace
-from tinygrad import Device, Tensor
-from tinygrad.engine.realize import compile_linear
+import os, sys
+from tinygrad import Tensor
+from tinygrad.engine.realize import lower_and_compile, run_linear
 
-dev = Device["AMD"]
-a = Tensor([1.0]).realize()
-b = a + 1
-linear = compile_linear(b.schedule_linear())
-compiled_prg = linear.src[-1].src[0]
-lib = bytearray(compiled_prg.src[3].arg)
-
-# Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
-# VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400
-found = False
-for i in range(0, len(lib) - 4, 4):
-  if struct.unpack("<I", lib[i:i+4])[0] == 0xBFB00000:
-    lib[i:i+4] = struct.pack("<I", 0x7E008400)
-    found = True
-    break
-assert found, "s_endpgm not found"
-
-patched_prg = dev.runtime(replace(compiled_prg.to_elf(), name="patched", lib=bytes(lib)))
-b.uop.buffer.allocate()
-patched_prg(b.uop.buffer._buf, a.uop.buffer._buf, global_size=(1,1,1), local_size=(1,1,1))
-dev.synchronize()
+linear = lower_and_compile((Tensor.empty(1) + 1).schedule_linear())
+binary = linear.src[-1].src[0].src[3]
+lib = binary.arg.replace(bytes.fromhex("0000b0bf"), bytes.fromhex("00fe017e"), 1)
+try:
+  run_linear(linear.substitute({binary: binary.replace(arg=lib)}, enter_calls=True))
+except ValueError as error:
+  print(error, file=sys.stderr, flush=True)
+  os._exit(1)
 '''
 
-    env = os.environ.copy()
-    env["DEV"] = "MOCKKFD+AMD"
-    env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
-
-    st = time.perf_counter()
-    result = subprocess.run([sys.executable, "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
-    elapsed = time.perf_counter() - st
-
-    self.assertNotEqual(result.returncode, 0, "should have raised")
-    self.assertTrue("Error" in result.stderr, f"expected an error in stderr, got: {result.stderr[:500]}")
-    # Should exit immediately, not wait for the full timeout
-    self.assertLess(elapsed, 9.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
+    env = {**os.environ, "DEV": "MOCKKFD+AMD", "HCQ_RUNTIME_DEV": "PYTHON"}
+    result = subprocess.run([sys.executable, "-c", test_code], env=env, capture_output=True, text=True, timeout=9)
+    self.assertEqual(result.returncode, 1)
+    self.assertIn("unknown rdna3 format word=0x7e01fe00", result.stderr)
 
 if __name__ == "__main__":
   unittest.main()

test/amd/test_roundtrip.py

@@ -58,11 +58,11 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, in
   """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
   from tinygrad import Tensor
   from tinygrad.uop.ops import Ops
-  from tinygrad.engine.realize import compile_linear, resolve_params, unwrap_multi
+  from tinygrad.engine.realize import lower_and_compile, resolve_params, unwrap_multi
   from tinygrad.runtime.support.elf import elf_loader
 
   out = op_fn(Tensor)
-  linear = compile_linear(out.schedule_linear())
+  linear = lower_and_compile(out.schedule_linear())
   kernels = []
   buf_pool: dict[int, int] = {}  # buffer id -> size
   buf_data: dict[int, bytes] = {}  # buffer id -> initial data from COPY

test/backend/test_ops.py

@@ -268,11 +268,6 @@ class TestOps(unittest.TestCase):
 
     self.helper_test_exception([], lambda: torch.meshgrid(x, indexing="bad"), lambda: xt.meshgrid(indexing="bad"), expected=RuntimeError)
 
-  def test_meshgrid_scalar(self):
-    for indexing in ("ij", "xy"):
-      with self.subTest(indexing=indexing):
-        helper_test_op([()], lambda x: torch.meshgrid(x, indexing=indexing)[0], lambda x: x.meshgrid(indexing=indexing)[0])
-
   def test_arange(self):
     helper_test_op([], lambda: torch.arange(10, dtype=torch.int32), lambda: Tensor.arange(10), forward_only=True)
     helper_test_op([], lambda: torch.arange(36, dtype=torch.int32), lambda: Tensor.arange(36), forward_only=True)
@@ -1134,6 +1129,7 @@ class TestOps(unittest.TestCase):
   def test_relu6(self):
     helper_test_op([(45,65)], torch.nn.functional.relu6, Tensor.relu6)
     helper_test_op([()], torch.nn.functional.relu6, Tensor.relu6)
+    helper_test_op(None, torch.nn.functional.relu6, Tensor.relu6, vals=[[6.71089e7, 2.68435e8, 1e9]])
   def test_hardswish(self):
     helper_test_op([(45,65)], torch.nn.functional.hardswish, Tensor.hardswish, grad_atol=1e-6)
     helper_test_op([()], torch.nn.functional.hardswish, Tensor.hardswish, grad_atol=1e-6)
@@ -1188,20 +1184,6 @@ class TestOps(unittest.TestCase):
   def test_small_cummax(self):
     helper_test_op([(10)], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0)[0])
     helper_test_op([(10)], lambda x: torch.cummax(x, dim=0).indices.int(), lambda x: Tensor.cummax(x, axis=0)[1], forward_only=True)
-
-  def test_cumextrema_ties(self):
-    for op in ("cummax", "cummin"):
-      for axis in (0, 1, -1):
-        for values in ([[2, 2, 1, 3, 3, 0, 0]] * 2, [[0, 0, 0]] * 2):
-          with self.subTest(op=op, axis=axis, values=values):
-            helper_test_op(None, lambda x: getattr(torch, op)(x, dim=axis).indices.int(),
-                           lambda x: getattr(x, op)(axis)[1], vals=[values], forward_only=True)
-
-  def test_cumextrema_ties_split(self):
-    for op in ("cummax", "cummin"):
-      helper_test_op(None, lambda x: getattr(torch, op)(x, dim=-1).indices.int(), lambda x: getattr(x, op)(-1)[1],
-                     vals=[[[2.0, 2.0, 1.0, 3.0, 3.0, 0.0, 0.0] * 100] * 2], forward_only=True)
-
   @slow_test
   def test_simple_cummax(self):
     helper_test_op([(512)], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0)[0])
@@ -1668,10 +1650,6 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x.isclose(torch.tensor(1.0)), lambda x: x.isclose(1.0),
                    vals=[[1.0, 1.0 + 1e-7, 2.0, math.inf, -math.inf, math.nan]], forward_only=True)
 
-  def test_isclose_overflow(self):
-    helper_test_op(None, lambda x,y: x.isclose(y, rtol=3),
-                   vals=[[3e38, -3e38, 3e38, 0.0], [-3e38, 3e38, 3e38, 1.0]], forward_only=True)
-
   def test_mean(self):
     helper_test_op([(3,4,5,6)], lambda x: x.mean())
     helper_test_op([()], lambda x: x.mean())
@@ -1716,16 +1694,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(15, 25, 35)], lambda x: x.var(keepdim=True))
     helper_test_op([(15, 25, 35)], lambda x: x.var(0, keepdim=True, correction=0))
 
-  def test_var_std_integer(self):
-    for op in ("var", "std"):
-      for axis in (None, 0, 1):
-        for correction in (0, 1):
-          for keepdim in (False, True):
-            with self.subTest(op=op, axis=axis, correction=correction, keepdim=keepdim):
-              helper_test_op(None, lambda x: getattr(x.float(), op)(dim=axis, correction=correction, keepdim=keepdim),
-                             lambda x: getattr(x, op)(axis=axis, correction=correction, keepdim=keepdim),
-                             vals=[[[0, 1, 3], [1, 2, 4]]], forward_only=True)
-
   @slow_test
   def test_std(self):
     helper_test_op([(15, 25, 35)], lambda x: x.std())
@@ -1842,21 +1810,6 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), atol=1e-7, grad_atol=1e-7, vals=[[0.0, 100.0]])
     helper_test_op(None, lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), vals=[[-math.inf, 0.0, 1.0]], forward_only=True)
 
-  def test_logcumsumexp_scalar_invalid_axis(self):
-    for axis in (-2, 1):
-      with self.subTest(axis=axis):
-        self.helper_test_exception([()], lambda x: torch.logcumsumexp(x, dim=axis), lambda x: x.logcumsumexp(axis), expected=IndexError)
-
-  def test_logcumsumexp_empty(self):
-    for shape, axis in (((0,), 0), ((2, 0, 3), 1), ((2, 0, 3), -1)):
-      with self.subTest(shape=shape, axis=axis):
-        helper_test_op([shape], lambda x: torch.logcumsumexp(x, dim=axis), lambda x: x.logcumsumexp(axis))
-
-  def test_logcumsumexp_nonfinite(self):
-    for values in ([-math.inf, -math.inf], [0., math.inf, -math.inf], [0., math.nan, 1.]):
-      with self.subTest(values=values):
-        helper_test_op(None, lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), vals=[values], forward_only=True)
-
   def test_sinh(self):
     helper_test_op([(45,65)], lambda x: x.sinh(), grad_atol=1e-6)
     # TODO: backward nan instead of inf
@@ -2237,12 +2190,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(3,5)], lambda x: x.diagonal(offset=2))  # offset on rectangular
     self.helper_test_exception([(3,3)], lambda x: x.diagonal(dim1=0, dim2=0), expected=RuntimeError)
 
-  def test_diagonal_outside_matrix(self):
-    for shape, dims in (((2, 3), (0, 1)), ((2, 3, 4), (-2, -1)), ((2, 3, 4), (2, 0))):
-      for offset in (-10, -4, 4, 10):
-        with self.subTest(shape=shape, dims=dims, offset=offset):
-          helper_test_op([shape], lambda x: x.diagonal(offset=offset, dim1=dims[0], dim2=dims[1]))
-
   def test_roll(self):
     helper_test_op([(2, 4)], lambda x: x.roll(1))
     helper_test_op([(2, 4)], lambda x: x.roll((1,)))
@@ -3380,16 +3327,6 @@ class TestOps(unittest.TestCase):
       helper_test_op([(12,10)], lambda x: torch.nn.CrossEntropyLoss(label_smoothing=s)(x, torch.tensor(classes)),
                                 lambda x: x.sparse_categorical_crossentropy(Tensor(classes), label_smoothing=s))
 
-  def test_sparse_categorical_crossentropy_default_ignore_index(self):
-    classes = [-1, 0, 2, -1]
-    for reduction in ("none", "sum", "mean"):
-      for smoothing in (0.0, 0.3, 1.0):
-        with self.subTest(reduction=reduction, smoothing=smoothing):
-          helper_test_op([(4, 3)],
-                         lambda x: torch.nn.functional.cross_entropy(x, torch.tensor(classes), ignore_index=-1,
-                                                                     reduction=reduction, label_smoothing=smoothing),
-                         lambda x: x.sparse_categorical_crossentropy(Tensor(classes), reduction=reduction, label_smoothing=smoothing))
-
   def test_nll_loss(self):
     target = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
     helper_test_op([(32,10)],
@@ -3503,31 +3440,6 @@ class TestOps(unittest.TestCase):
     if not COMPILE_ONLY: assert t == -1
 
 class TestOpsUint8(unittest.TestCase):
-  def test_lerp_integer_end(self):
-    for dtype in dtypes.ints:
-      with self.subTest(dtype=dtype):
-        actual = Tensor([[10], [100]], dtype=dtypes.uint8).lerp(Tensor([20, 20, 100], dtype=dtype), Tensor([0., 0.5, 1.]))
-        self.assertEqual(actual.dtype, dtypes.uint8)
-        actual.realize()
-        if not COMPILE_ONLY: np.testing.assert_equal(actual.numpy(), [[10, 15, 100], [100, 60, 100]])
-
-  def test_lerp_float_end(self):
-    helper_test_op(None, lambda x,y,w: x.float().lerp(y, w), lambda x,y,w: x.cast(dtypes.uint8).lerp(y, w),
-                   vals=[[[10], [100]], [20.5, 9.5, -5.5], [0., 0.5, 1.]], forward_only=True)
-
-  def test_interpolate_bilinear_full_range(self):
-    for values in ([[0, 255]], [[255, 0]], [[1, 200]], [[0, 255], [255, 0]]):
-      for size in ((1, 3), (5, 10)):
-        for align_corners in (False, True):
-          with self.subTest(values=values, size=size, align_corners=align_corners):
-            image = torch.tensor([[values]], dtype=torch.uint8)
-            expected = torch.nn.functional.interpolate(image, size=size, mode="bilinear", align_corners=align_corners)
-            actual = Tensor(image.numpy()).interpolate(size, align_corners=align_corners)
-            self.assertEqual(actual.dtype, dtypes.uint8)
-            # Midpoints are exact; other weights can differ by one with 7-bit fixed-point coefficients.
-            actual.realize()
-            if not COMPILE_ONLY: np.testing.assert_allclose(actual.numpy(), expected.numpy(), rtol=0, atol=0 if size == (1, 3) else 1)
-
   def test_cast(self):
     helper_test_op([(2,3,64,64)], lambda x: x.type(torch.uint8), lambda x: x.cast('uint8'), forward_only=True, low=0, high=255)
 

test/backend/test_setitem.py

@@ -1,6 +1,5 @@
 import unittest, operator
 from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
-from tinygrad.helpers import Context
 import numpy as np
 
 class TestSetitem(unittest.TestCase):
@@ -163,21 +162,20 @@ class TestSetitem(unittest.TestCase):
       np.testing.assert_allclose(t.numpy(), n)
 
   def test_jit_setitem_variable_offset(self):
-    with Context(CHECK_OOB=0):
-      @TinyJit
-      def f(t:Tensor, a:Tensor, v:Variable):
-        t.shrink(((v,v+1), None)).assign(a).realize()
+    @TinyJit
+    def f(t:Tensor, a:Tensor, v:Variable):
+      t.shrink(((v,v+1), None)).assign(a).realize()
 
-      t = Tensor.zeros(6, 6).contiguous().realize()
-      n = np.zeros((6, 6))
+    t = Tensor.zeros(6, 6).contiguous().realize()
+    n = np.zeros((6, 6))
 
-      for i in range(6):
-        v = Variable("v", 0, 6).bind(i)
-        a = Tensor.full((1, 6), fill_value=i+1, dtype=dtypes.float).contiguous()
-        n[i, :] = i+1
-        f(t, a, v)
-        np.testing.assert_allclose(t.numpy(), n)
-      np.testing.assert_allclose(t.numpy(), [[1,1,1,1,1,1],[2,2,2,2,2,2],[3,3,3,3,3,3],[4,4,4,4,4,4],[5,5,5,5,5,5],[6,6,6,6,6,6]])
+    for i in range(6):
+      v = Variable("v", 0, 6).bind(i)
+      a = Tensor.full((1, 6), fill_value=i+1, dtype=dtypes.float).contiguous()
+      n[i, :] = i+1
+      f(t, a, v)
+      np.testing.assert_allclose(t.numpy(), n)
+    np.testing.assert_allclose(t.numpy(), [[1,1,1,1,1,1],[2,2,2,2,2,2],[3,3,3,3,3,3],[4,4,4,4,4,4],[5,5,5,5,5,5],[6,6,6,6,6,6]])
 
   def test_setitem_overlapping_inplace1(self):
     t = Tensor([[3.0], [2.0], [1.0]]).contiguous()

test/device/test_hcq2.py

@@ -1,11 +1,12 @@
 import unittest, contextlib, ctypes, gc, numpy as np
 from unittest.mock import patch
 from tinygrad import Device, Tensor, TinyJit, Variable, dtypes, GlobalCounters
-from tinygrad.device import Buffer, BufferSpec
+from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import Context, dedup, partition
+from tinygrad.helpers import Context, dedup, partition, unwrap
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo
 from tinygrad.engine.realize import compile_linear, link_linear, lower_and_compile, run_linear
+from tinygrad.codegen import do_to_program
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.c import init_c_struct_t
@@ -28,8 +29,10 @@ def chain(x:Tensor, n:int) -> Tensor:
 @contextlib.contextmanager
 def encoded_batches():
   batches, orig = [], hcq2.lower_and_compile
-  with patch.object(hcq2, "lower_and_compile", lambda l, *a, **kw: (batches.extend(c for c in l.src if call_is_hcq(c)), orig(l, *a, **kw))[1]):
-    yield batches
+  def track(l, *args, **kwargs):
+    batches.extend(c.without_after for c in l.src if call_is_hcq(c))
+    return orig(l, *args, **kwargs)
+  with patch.object(hcq2, "lower_and_compile", track): yield batches
 
 def eager_chain(x:Tensor, n:int=64) -> Tensor: # at hcq_compile's use_rt bound: an eager linear this big bakes its inputs and borrows ring slots
   for _ in range(n): x = (x + 1).contiguous()
@@ -42,6 +45,11 @@ def patch_words(batch:UOp) -> list[UOp]:
 def rt_params(batch:UOp) -> list[str]:
   return dedup([u.arg.name for w in patch_words(batch) for u in w.toposort() if u.op is Ops.PARAM and u.arg.addrspace is AddrSpace.GLOBAL])
 
+def cpu_buf(size:int=1, dtype=dtypes.uint8, **kwargs) -> UOp: return UOp.placeholder((size,), dtype, device="CPU", **kwargs)
+
+def lower_hcq(body:UOp) -> UOp:
+  return unwrap(hcq2.lower_call(UOp.sink(body, arg=KernelInfo("test")).call(aux=hcq2.HCQInfo(("CPU",)))))
+
 class TestHCQ2Deps(unittest.TestCase):
   def test_disjoint_write_preserves_dependencies(self):
     b = UOp.param(0, dtypes.uint8, 16, device="CPU")
@@ -61,8 +69,8 @@ class TestHCQ2Deps(unittest.TestCase):
       self.assertEqual(tracker.access_resources([b.shrink(((12, 16),))], [0], 3), [0])
       self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [], 4), [1])
 
-@unittest.skipUnless(all_devices_in(Device.DEFAULT, HCQ_DEVS - {"CPU"}), "non-CPU hcq2 device required")
-class TestHCQ2Core(unittest.TestCase):
+@unittest.skipUnless(all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
+class TestHCQ2Schedule(unittest.TestCase):
   @staticmethod
   def input(value:int=2) -> Tensor: return Tensor.full((4,), value, dtype=dtypes.int32).contiguous().realize()
 
@@ -93,6 +101,33 @@ class TestHCQ2Core(unittest.TestCase):
     linked = link_linear(compiled, input_uops=inputs)
     self.assertIs(link_linear(compiled, input_uops=inputs), linked)
 
+  def test_profile_slots_survive_indirect_access(self):
+    pm = PatternMatcher([(UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat())),), allow_any_len=True),
+                          lambda buf: hcq2.rt_addr(buf, "CPU") if hcq2.unwrap_view(buf)[0].tag == "slots" else None)])
+    with patch.object(Device[Device.DEFAULT], "pm_lower", pm):
+      compiled = compile_linear(Tensor.ones(4).contiguous().schedule_linear(), profile=True)
+    self.assertFalse(any(param.op is Ops.PARAM and (param.arg.name or "").startswith("slots_")
+                         for param in compiled.src[0].without_after.src[0].toposort()))
+    call = link_linear(compiled).src[0].without_after
+    ((device, index),) = call.arg.aux.slots
+    self.assertEqual(device, Device.DEFAULT)
+    self.assertEqual(call.src[1 + index].buffer.dtype, dtypes.uint64)
+
+  def test_host_copies(self):
+    dev = Device[Device.DEFAULT]
+    if not dev.has_copy_queue: self.skipTest("copy queue required")
+    for host_device in ("CPU", "NPY", "DISK"):
+      for direct in (False, True):
+        for upload in (False, True):
+          with self.subTest(host_device=host_device, direct=direct, upload=upload):
+            host, gpu = UOp.new_buffer(host_device, 4, dtypes.uint8), UOp.new_buffer(dev.device, 4, dtypes.uint8)
+            src, dst = (host, gpu) if upload else (gpu, host)
+            linear = UOp(Ops.LINEAR, src=(src.copy_to_device(dst.device).call(dst, src),))
+            with patch.object(dev, "host_devs", frozenset({"CPU", host_device}) if direct else frozenset({"CPU"})):
+              compiled = compile_linear(linear, profile=False)
+            self.assertEqual(len(compiled.src), 1 if direct or host_device == "CPU" else 2)
+            self.assertEqual(sum(call_is_hcq(call) for call in compiled.src), 1)
+
   def test_large_eager_not_cached(self):
     _, compiled, inputs = self.compiled(65)
     linked = link_linear(compiled, input_uops=inputs)
@@ -221,62 +256,70 @@ class TestHCQ2Core(unittest.TestCase):
     refs = [u.buffer for u in call.src[1:] if u.op is Ops.BUFFER]
     for n in names: self.assertTrue(any(r is getattr(dev, n) for r in refs), f"{n} is not a ref of the call")
 
+  def test_usb_renumbering(self):
+    programs = []
+    with Context(HCQ_RUNTIME_DEV="CPU"), patch("tinygrad.codegen.do_to_program", wraps=do_to_program) as build:
+      for ids in ((0, 1, 2, 3), (2, 0, 3, 1), (1, 0, 2, 3), (0, 1, 3, 2), (100, 101, 102, 103)):
+        with self.subTest(ids=ids):
+          regs = [UOp.placeholder((1,), dtypes.uint32, slot=i, addrspace=AddrSpace.REG) for i in ids[:2]]
+          a, b = [r.after(r.index(0).store(v)) for r, v in zip(regs, (3, 5))]
+          i, j = [UOp.range(UOp(Ops.NOOP), n, dtype=dtypes.void, src=(a, b)) for n in ids[2:]]
+          out = cpu_buf(dtype=dtypes.uint32, tag="out")
+          body = out.index(0).store(a.after(i, j).index(0).load()*10 + b.index(0).load()).end(j, UOp.const(False)).end(i, UOp.const(False))
+          compiled = lower_and_compile(UOp(Ops.LINEAR, src=(lower_hcq(body),)))
+          programs.append(compiled.src[0].without_after.src[0])
+          self.assertIs(programs[-1], programs[0])
+          linear = hcq2.hcq_link(compiled, allow_cache=False)
+          run_linear(linear, jit=True)
+          self.assertEqual(linear.src[0].without_after.src[1].buffer._buf.cpu_view().view(fmt='I')[0], 35)
+      self.assertLessEqual(build.call_count, 1)
+
+  def test_patched_view(self):
+    with Context(HCQ_RUNTIME_DEV="CPU"):
+      ctx = hcq2.EncodeCtx(("CPU",))
+      inner = hcq2.patch(cpu_buf(8, tag="inner"), [(4, UOp.const(42, dtypes.uint32))], bytes(8))
+      inner = unwrap(hcq2.hoist_links(ctx, inner))
+      outer = hcq2.patch(cpu_buf(8, tag="outer"), [(0, inner[4:8].getaddr("CPU"))])
+      with patch.object(hcq2, "EncodeCtx", return_value=ctx): call = lower_hcq(outer.bitcast(dtypes.uint64).index(0).load())
+      self.assertEqual(call.without_after.arg.aux.nargs, 1)
+      self.assertTrue(all(s.op is Ops.STORE for s in call.src[1:]))
+      linked = hcq2.hcq_link(UOp(Ops.LINEAR, src=(call,)), allow_cache=False).src[0]
+      inner_buf, outer_buf = linked.src[1].buffer, linked.without_after.src[1].buffer
+      self.assertEqual(inner_buf._buf.cpu_view().view(fmt='I')[1], 42)
+      self.assertEqual(outer_buf._buf.cpu_view().view(fmt='Q')[0], inner_buf._buf.va_addr + 4)
+
 @unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "CALL is rendered in C style only")
 class TestHCQ2FFI(unittest.TestCase):
   @staticmethod
   def _run(body:UOp) -> list[Buffer]:
-    call = hcq2.lower_call(UOp.sink(body, arg=KernelInfo("test_ffi")).call(aux=hcq2.HCQInfo(("CPU",))))
-    assert call is not None
-    linear = hcq2.hcq_link(lower_and_compile(UOp(Ops.LINEAR, src=(call,))), allow_cache=False)
+    linear = hcq2.hcq_link(lower_and_compile(UOp(Ops.LINEAR, src=(lower_hcq(body),))), allow_cache=False)
     run_linear(linear, jit=True)
     return [u.buffer for u in linear.src[0].without_after.src[1:] if u.op is Ops.BUFFER]
 
   def test_ffi_ccall(self):
     with Context(HCQ_RUNTIME_DEV="CPU"):
-      out = UOp.placeholder((1,), dtypes.int32, slot=1, device="CPU", volatile=True, tag="ffi_result")
+      out = cpu_buf(dtype=dtypes.int32, slot=1, volatile=True, tag="ffi_result")
       bufs = self._run(out.index(0).store(hcq2.ccall(libc.dll.ffs, 0x10)))
     self.assertEqual(next(b for b in bufs if b.dtype is dtypes.int)._buf.cpu_view().view(fmt='i')[0], 5)
 
   def test_ffi_cstruct(self):
     struct_t = init_c_struct_t(16, (("u8", ctypes.c_uint8, 0), ("u16", ctypes.c_uint16, 2),
                                   ("u32", ctypes.c_uint32, 4), ("u64", ctypes.c_uint64, 8)))
-    UOp.placeholder((1,), dtypes.uint8, device="CPU") # reserve slot zero for device-owned placeholders
+    cpu_buf() # reserve slot zero for device-owned placeholders
     with Context(HCQ_RUNTIME_DEV="CPU"):
       s = hcq2.cstruct(struct_t, u8=0x12, u16=UOp.const(0x3456, dtypes.uint16), u32=0x789ABCDE, u64=0xFEDCBA9876543210)
       bufs = self._run(s.index(0).load())
     got = struct_t.from_buffer_copy(bytes(next(b for b in bufs if b.nbytes == ctypes.sizeof(struct_t))._buf.cpu_view()))
     self.assertEqual((got.u8, got.u16, got.u32, got.u64), (0x12, 0x3456, 0x789ABCDE, 0xFEDCBA9876543210))
 
-  def test_device_lower_after_encode(self):
-    with Context(HCQ_RUNTIME_DEV="CPU"):
-      out = UOp.placeholder((1,), dtypes.int32, device="CPU", tag="result")
-      encode = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="test_encode"), lambda: UOp.custom_function("test_lower"))])
-      lower = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="test_lower"), lambda out=out: out.index(0).store(42))])
-      with patch.object(Device["CPU"], "pm_encode", encode), patch.object(Device["CPU"], "pm_lower", lower):
-        bufs = self._run(UOp.custom_function("test_encode"))
-    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.int)._buf.cpu_view().view(fmt='i')[0], 42)
-
   def test_nested_cstruct_patches(self):
     with Context(HCQ_RUNTIME_DEV="CPU"):
       inner = hcq2.cstruct(init_c_struct_t(4, (("value", ctypes.c_uint32, 0),)), value=42)
       outer = hcq2.cstruct(init_c_struct_t(8, (("ptr", ctypes.c_uint64, 0),)), ptr=inner.getaddr("CPU"))
-      out = UOp.placeholder((1,), dtypes.uint32, device="CPU", tag="result")
+      out = cpu_buf(dtype=dtypes.uint32, tag="result")
       copied = hcq2.ccall(libc.memcpy, out.index(0), outer.bitcast(dtypes.uint64).index(0).load(), 4)
       bufs = self._run(out.after(copied).index(0).load())
     self.assertEqual(next(b for b in bufs if b.dtype is dtypes.uint32)._buf.cpu_view().view(fmt='I')[0], 42)
 
-
-class TestHCQ2Timeline(unittest.TestCase):
-  def test_reused_timeline_is_zeroed(self):
-    buf = Buffer("CPU", 2, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
-    addr = buf._buf.va_addr
-    buf._buf.cpu_view().view(fmt='B')[:] = b'\xff' * 16
-    buf.deallocate()
-    dev = HCQ2Compiled.__new__(HCQ2Compiled)
-    dev.device = "CPU"
-    self.assertEqual(dev.timeline._buf.va_addr, addr)
-    self.assertEqual(bytes(dev.timeline._buf.cpu_view()), bytes(16))
-
-
 if __name__ == "__main__":
   unittest.main()

test/external/external_test_usb_asm24.py

@@ -1,6 +1,7 @@
 import unittest
 from tinygrad.helpers import Timing, getenv
-from tinygrad import Tensor, Device
+from tinygrad import Tensor, Device, TinyJit
+from tinygrad.runtime.support.usb import HALF, CHUNK, SLOT
 import numpy as np
 
 class USBTestCase(unittest.TestCase):
@@ -8,7 +9,12 @@ class USBTestCase(unittest.TestCase):
   def setUpClass(cls):
     cls.sz = getenv("SIZE", 2000000)
     cls.dev = Device["AMD"]
-    if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
+    if not cls.dev.is_usb: raise unittest.SkipTest("only test this on USB devices")
+    cls.rng = np.random.default_rng(0)
+
+  def roundtrip(self, a:np.ndarray): # a copy in, a kernel, a copy out: the queue must order them
+    np.testing.assert_array_equal(a, Tensor(a, device="NPY").to(Device.DEFAULT).numpy())
+    np.testing.assert_array_equal(a + 1, (Tensor(a, device="NPY").to(Device.DEFAULT) + 1).numpy())
 
 class TestDevCopySpeeds(USBTestCase):
   def testCopyCPUtoDefault(self):
@@ -30,73 +36,55 @@ class TestUSBIntegrity(USBTestCase):
     t = Tensor.randn(self.sz, device="CPU", dtype='uchar').contiguous().realize()
     x = t.to(Device.DEFAULT).realize()
     Device[Device.DEFAULT].synchronize()
-
     y = x.to('CPU').realize()
-
     np.testing.assert_equal(t.numpy(), y.numpy())
-    del x, y, t
 
-  def testCopyinBoundaries(self):
-    rng, chunk = np.random.default_rng(0), 0x40000 - 4
-    for size in (1, 3, 508, 509, 0x3ffc, 0x3ffd, chunk, chunk+1, 2*chunk+31):
-      with self.subTest(size=size):
-        a = rng.integers(0, 256, size, dtype=np.uint8)
-        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+  def testBoundaries(self): # around the slot, the chunk and the read window
+    for size in (1, 3, 508, 509, SLOT - 513, SLOT - 512, SLOT - 511, CHUNK - 1, CHUNK, CHUNK + 1, 2 * CHUNK - 1, 2 * CHUNK, 2 * CHUNK + 31, HALF,
+                 2 * HALF, 1 << 20):
+      with self.subTest(size=size): self.roundtrip(self.rng.integers(0, 256, size, dtype=np.uint8))
 
-  def testCopyinFenceWrap(self):
-    a = np.arange(2*(0x40000-4)+31, dtype=np.uint8)
-    np.testing.assert_array_equal(a[:31], Tensor(a[:31], device="AMD").numpy())
-    self.dev.synchronize()
-    alloc, usb = self.dev.allocator, self.dev.iface.pci_dev.usb
-    clear = usb.read(0xA808, 1)
-    # Model a completed 256-chunk copy instead of the one-chunk warmup. The next clear tag must still change.
-    alloc._usb_seq += 255
-    usb.write(0xA800, bytes([alloc._usb_seq & 0xff]))
-    np.testing.assert_array_equal(a[:31], Tensor(a[:31], device="AMD").numpy())
-    self.assertNotEqual(clear, usb.read(0xA808, 1))
-    for bits in (8, 24):
-      with self.subTest(bits=bits):
-        alloc._usb_seq = ((alloc._usb_seq >> bits)+2)*(1 << bits)-2
-        usb.write(0xA800, bytes([alloc._usb_seq & 0xff]))
-        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+  def testManyCopiesInABatch(self):
+    for n in (2, 7, 64, 300): # 300 chunks: the fence byte wraps
+      with self.subTest(n=n):
+        arrs = [self.rng.integers(0, 256, int(s), dtype=np.uint8) for s in self.rng.integers(1, 5000, n)]
+        ts = [Tensor(a, device="NPY").to(Device.DEFAULT) for a in arrs]
+        Tensor.realize(*ts)
+        for t, a in zip(ts, arrs): np.testing.assert_array_equal(a, t.numpy())
 
-  def testCopyinRingWrap(self):
-    rng = np.random.default_rng(0)
-    a = rng.integers(0, 256, 1 << 20, dtype=np.uint8)
-    np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
-    ring = self.dev.sdma_queue(0)
-    # A 16 MiB copyin needs more than 4 KiB of SDMA packets, forcing the submission to wrap.
-    target = ring.ring.nbytes - 0x1000
-    padding = target - ring.put_value % ring.ring.nbytes - 16  # four-dword timeline fence
-    self.assertGreaterEqual(padding, 0)
-    q = self.dev.hw_copy_queue_t()
-    q.q(*([0] * (padding // 4)))
-    q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
-    self.dev.synchronize()
-    before = ring.put_value // ring.ring.nbytes
-    a = rng.integers(0, 256, 16 << 20, dtype=np.uint8)
-    t = Tensor(a, device="AMD").realize()
-    self.assertGreater(ring.put_value // ring.ring.nbytes, before)
+  def testMixedBatch(self): # copies out and in, in one batch: runs of both directions
+    arrs = [self.rng.integers(0, 256, s, dtype=np.uint8) for s in (5, CHUNK + 7, 9, 2 * CHUNK + 3, 11)]
+    ts = [Tensor(a, device="NPY").to(Device.DEFAULT).realize() for a in arrs]
+    more = [self.rng.integers(0, 256, s, dtype=np.uint8) for s in (5, CHUNK + 7, 9, 2 * CHUNK + 3, 11)]
+    outs = [t.to("NPY") for t in ts] + [Tensor(a, device="NPY").to(Device.DEFAULT) for a in more]
+    Tensor.realize(*outs)
+    for o, a in zip(outs, arrs + more): np.testing.assert_array_equal(a, o.numpy())
+
+  def testRepeatedBatches(self): # a batch numbers its chunks from 0: the same batch again must not see what the last one left behind
+    a = self.rng.integers(0, 256, 2 * CHUNK + 31, dtype=np.uint8)
+    for _ in range(5): self.roundtrip(a)
+    @TinyJit
+    def step(x:Tensor) -> Tensor: return (x + 1).realize()
+    src = Tensor(a, device="NPY")
+    for i in range(5):
+      x = src.to(Device.DEFAULT)
+      np.testing.assert_array_equal(a + 1, step(x).numpy())
+
+  def testStaleSentinel(self): # payloads full of the tags the queue waits for, in both directions, before and around the real chunks
+    tags = np.array([0x51000000 | k for k in range(8)], dtype=np.uint32)
+    for tag in tags: # every dword of every chunk is the tag of some chunk of the copy
+      with self.subTest(payload=hex(tag)):
+        a = np.full((2 * CHUNK + 31) // 4, tag, dtype=np.uint32).view(np.uint8)
+        self.roundtrip(a)
+    with self.subTest(case="copyout residue"): # a read fills the sram with tags, then small chunks land in both halves
+      a = np.tile(tags, 2 * CHUNK // 32).view(np.uint8)
+      np.testing.assert_array_equal(a, (Tensor(a, device="NPY").to(Device.DEFAULT) * 1).numpy())
+      for size in (31, CHUNK + 31, 2 * CHUNK + 31): self.roundtrip(np.tile(tags, size // 32 + 1).view(np.uint8)[:size])
+
+  def testRingWrap(self): # 64MB of chunks: the sdma ring (1MB on usb) wraps within the copy
+    a = self.rng.integers(0, 256, 64 << 20, dtype=np.uint8)
+    t = Tensor(a, device="NPY").to(Device.DEFAULT).realize()
     np.testing.assert_array_equal(a, t.numpy())
 
-  def testCopyinStaleSentinel(self):
-    a = np.arange(16, dtype=np.uint8)
-    np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
-    chunk = 0x40000 - 4
-    for case in ("copyout", "reuse"):
-      with self.subTest(case=case):
-        if case == "copyout":
-          # A 512 KiB copyin takes three chunks. Copyout then fills both SRAM windows with the next expected tag.
-          tag = 0x51000000 | ((self.dev.allocator._usb_seq + 3) & 0xFFFFFF)
-          a = np.full(0x80000 // 4, tag, dtype=np.uint32)
-          np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
-          a = np.arange(31, dtype=np.uint8)
-        else:
-          # The first full chunk contains the tag expected by the short third chunk in the same window.
-          tag = 0x51000000 | ((self.dev.allocator._usb_seq + 2) & 0xFFFFFF)
-          a = np.arange(2 * chunk + 31, dtype=np.uint8)
-          a[:chunk].view(np.uint32)[:] = tag
-        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
-
 if __name__ == "__main__":
   unittest.main()

test/external/fuzz_fast_idiv.py

@@ -12,7 +12,8 @@ if __name__ == "__main__":
     if i % 1000 == 0:
       print(f"Progress: {i}")
     dt = random.choice(dtypes.ints)
-    u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt)
+    vmax = random.randint(1, 2**random.randint(1, dt.max.bit_length()))
+    u = UOp.variable('x', random.randint(0, vmax-1) if vmax > 1 else 0, vmax, dtype=dt)
     d = random.randint(1, max(1, u.vmax)*2)
     if d in powers_of_two: continue
     expr = fast_idiv(Device[Device.DEFAULT].renderer, u, d)

test/helpers.py

@@ -67,7 +67,7 @@ def assert_kernel_count(expected:int):
 
 def is_hcq2_device() -> bool: # an hcq2 device stages every copy from the host through a pinned buffer: such a copy is two calls, not one
   from tinygrad.runtime.support.hcq2 import HCQ_DEVS
-  return Device.DEFAULT.split(":")[0] in HCQ_DEVS - {"CPU"}
+  return Device.DEFAULT.split(":")[0] in HCQ_DEVS
 
 def call_is_graph(call:UOp) -> bool:
   ast = call.src[0]

test/null/test_helpers.py

@@ -78,6 +78,13 @@ class TestContextVars(unittest.TestCase):
     test()
     self.assertEqual(VARIABLE.value, 0)
 
+  def test_decorator_recursive(self):
+    @Context(VARIABLE=1)
+    def test(n):
+      if n: test(n-1)
+    test(2)
+    self.assertEqual(VARIABLE.value, 0)
+
   def test_context_exit_reverts_updated_values(self):
     D = ContextVar("D", 1)
     D.value = 2

test/null/test_uops.py

@@ -269,7 +269,6 @@ class TestGatedStoreRewrite(unittest.TestCase):
     for x in gated_uops: self.assertIs(x.op, Ops.STORE)
     for x in gated_uops: self.assertEqual(len(x.src), 2)
 
-@unittest.skipIf(Device.DEFAULT == "METAL", "compiler bug")
 @unittest.skipUnless(Ops.SHR in Device[Device.DEFAULT].renderer.code_for_op, "fast_idiv requires SHR")
 class TestFastIdiv(unittest.TestCase):
   def test_division_power_of_two(self):
@@ -310,7 +309,7 @@ class TestFastIdiv(unittest.TestCase):
       self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")
 
   @Context(DISABLE_FAST_IDIV=0)
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
+  @unittest.skipUnless(dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "fast_idiv widens uint32 to uint64")
   def test_fast_idiv_and_mod(self):
     g = UOp.param(0, dtypes.uint32, 4)
     c = UOp.const(3)
@@ -329,6 +328,25 @@ class TestFastIdiv(unittest.TestCase):
     self.assertIn(Ops.SHR, ops)
     self.assertNotIn(Ops.CMOD, ops)
 
+  @Context(DISABLE_FAST_IDIV=0)
+  def test_fast_idiv_nonpositive_divisor(self):
+    ridx = UOp.range(20, 0)
+    for d in (-3, 0):
+      for op in (Ops.CDIV, Ops.CMOD):
+        ops = [x.op for x in to_uops_list([ridx.alu(op, UOp.const(d))], ren=Device[Device.DEFAULT].renderer)]
+        self.assertNotIn(Ops.SHR, ops, f"fast_idiv fired on {op} by {d}")
+
+  @Context(DISABLE_FAST_IDIV=0)
+  @unittest.skipUnless(dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "needs a uint64 buffer")
+  def test_fast_idiv_cmod_kept_when_idiv_declines(self):
+    ren = Device[Device.DEFAULT].renderer
+    d = UOp.param(0, dtypes.int32, 4).index(UOp.const(0))
+    ops = [x.op for x in to_uops_list([UOp.range(30, 0).alu(Ops.CMOD, d)], ren=ren)]
+    self.assertIn(Ops.CMOD, ops, "CMOD by a non-const divisor should be left alone")
+    big = UOp.param(1, dtypes.uint64, 4).index(UOp.const(0))
+    ops = [x.op for x in to_uops_list([big.alu(Ops.CMOD, UOp.const(3, dtypes.uint64))], ren=ren)]
+    self.assertIn(Ops.CMOD, ops, "CMOD should be left alone when fast_idiv declines")
+
   @Context(DISABLE_FAST_IDIV=0)
   def test_fast_idiv_bounded_numerator_zero(self):
     x = UOp.variable("x", 0, 1, dtype=dtypes.int32)
@@ -342,6 +360,7 @@ class TestFastIdiv(unittest.TestCase):
     # this requires shifting out the powers of two before doing fast_idiv
     # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
     self.assertNotIn(dtypes.long, [x.dtype for x in uops])
+    self.assertNotIn(Ops.CDIV, [x.op for x in uops])
 
   @unittest.expectedFailure
   def test_fast_idiv_overflow(self):

test/null/test_validate_oob.py

@@ -186,40 +186,20 @@ class TestValidateOOB(unittest.TestCase):
       with self.assertRaises(RuntimeError):
         to_uops_list([buf_int.index(gidx.valid(ld_bool)).load()])  # gidx 0..15, buf_int size 8
 
-  # skipped tests (moved from test_uop_graph.py)
-  @unittest.skip("if not allowed in graph")
-  def test_in_bounds_access_gated_local(self):
-    with Context(CHECK_OOB=1):
-      # Define buffers
+  # local memory
+  def test_gated_local(self):
+    with Context(CHECK_OOB=1, SPEC=2):
       gbuf = UOp.param(0, dtypes.uint, 400)
       sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL)
-
-      # Define indices, valids and barrier
       gidx = UOp(Ops.SPECIAL, src=(UOp.const(416),), arg="gidx0")
       lidx = UOp(Ops.SPECIAL, src=(UOp.const(10),), arg="lidx0")
-
-      gate = (gidx<400) & (lidx<8)
-
-      local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(1))
-
-      barrier = UOp(Ops.BARRIER, src=(local_store,))
-      if_barrier = UOp(Ops.IF, src=(gate, barrier))
-
-      # Load from local memory (after the IF/barrier)
-      local_load = UOp(Ops.LOAD, src=(sbuf.index(lidx), if_barrier))
-
-      # Store to global memory
-      global_store = UOp(Ops.STORE, src=(gbuf.index(gidx), local_load))
-      to_uops_list([global_store])
-
-  @unittest.skip("Bool load is not supported yet")
-  def test_load_mask(self):
-    with Context(CHECK_OOB=1):
-      glbl0 = UOp.param(0, dtypes.int, 16)
-      mask = UOp.param(0, dtypes.bool, 16)
-      ridx = UOp.range(20, 0)
-      ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx))))
-      to_uops_list([ld0])
+      store = sbuf.index(lidx.valid(lidx < 8)).store(UOp.const(1))
+      load = sbuf.after(store).index(lidx.valid(lidx < 8)).load()
+      to_uops_list([gbuf.index(gidx.valid(gidx < 400)).store(load)])  # valid: local store and load gated to 8, global store gated to 400
+      with self.assertRaises(RuntimeError):
+        to_uops_list([gbuf.index(gidx.valid(gidx < 400)).store(sbuf.after(store).index(lidx).load())])  # lidx 0..9 into 8
+      with self.assertRaises(RuntimeError):
+        to_uops_list([gbuf.index(gidx).store(load)])  # gidx 0..415 into 400
 
 if __name__ == "__main__":
   unittest.main()

test/testextra/test_hevc.py

@@ -98,7 +98,7 @@ class TestHevc(unittest.TestCase):
       Variable("pos", 0, max_hist + 1).bind(frame_pos), out_image_size, opaque[1], history)
 
     compiled = compile_linear(decoded.linear_with_vars()[0])
-    self.assertTrue(any(call.src[0].op is Ops.PROGRAM for call in compiled.src))
+    self.assertTrue(any(call.without_after.src[0].op is Ops.PROGRAM for call in compiled.src))
     encdec_calls = [call for call in compiled.src if call.src[0].op is Ops.CUSTOM_FUNCTION and call.src[0].arg == "encdec"]
     self.assertEqual(len(encdec_calls), 1)
 

test/unit/test_dtype_spec.py

@@ -151,17 +151,6 @@ class TestTypeSpec(unittest.TestCase):
     _assert_eq(Tensor.arange(5.0, 3.0), dtypes.default_float, np.arange(5.0, 3.0))
 
 class TestAutoCastType(unittest.TestCase):
-  @unittest.skipUnless(dtypes.float64 in supported_dtypes, "need float64")
-  def test_linspace_float64_precision(self):
-    for start, stop in ((1., 1.+1e-8), (1.+1e-8, 1.), (1e10, 1e10+1)):
-      with self.subTest(start=start, stop=stop):
-        out = Tensor.linspace(start, stop, 3, dtype=dtypes.float64)
-        self.assertEqual(out.dtype, dtypes.float64)
-        np.testing.assert_allclose(out.numpy(), np.linspace(start, stop, 3), rtol=1e-15, atol=0)
-    with Context(DEFAULT_FLOAT=dtypes.float64):
-      out = Tensor.linspace(10**10, 10**10+2, 3, dtype=dtypes.int64)
-      np.testing.assert_array_equal(out.numpy(), [10**10, 10**10+1, 10**10+2])
-
   def test_int_sqrt(self):
     _assert_eq(Tensor([1, 4, 9, 16]).sqrt(), dtypes.default_float, [1, 2, 3, 4])
 
@@ -233,13 +222,6 @@ class TestAutoCastType(unittest.TestCase):
     t.square().mean().backward()
     np.testing.assert_allclose(t.grad.numpy().flatten(), [60000 * 2 / (N*N)] * N*N)
 
-  def test_var_integer_fractional(self):
-    for dtype in [*dtype_ints, dtypes.bool]:
-      with self.subTest(dtype=dtype):
-        out = Tensor([0, 1], dtype=dtype).var()
-        self.assertEqual(out.dtype, dtypes.float32)
-        np.testing.assert_allclose(out.numpy(), 0.5)
-
   @unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
   def test_var_half_precision_large_n(self):
     # the element count (70000) exceeds half max (65504): the denominator must not be materialized in half

test/unit/test_state.py

@@ -1,82 +0,0 @@
-import unittest
-from collections import OrderedDict, namedtuple
-from types import SimpleNamespace
-from tinygrad import Tensor
-from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict
-
-
-class TestStateDict(unittest.TestCase):
-  def test_container_subclasses(self):
-    class TensorDict(dict): pass
-    class TensorList(list): pass
-    class TensorTuple(tuple): pass
-    weight = Tensor([1., 2.])
-    for container, key in ((TensorDict(weight=weight), "weight"), (TensorList([weight]), "0"), (TensorTuple([weight]), "0")):
-      with self.subTest(container=type(container).__name__):
-        container.description = "model weights"
-        model = SimpleNamespace(layers=container)
-        state = get_state_dict(model)
-        self.assertEqual(list(state), [f"layers.{key}"])
-        self.assertIs(state[f"layers.{key}"], weight)
-        params = get_parameters(model)
-        self.assertEqual(len(params), 1)
-        self.assertIs(params[0], weight)
-
-  def test_namedtuple_and_ordered_dict(self):
-    first, second = Tensor([1.]), Tensor([2.])
-    pair = namedtuple("Pair", ["first", "second"])(first, second)
-    state = get_state_dict(OrderedDict(pair=pair))
-    self.assertEqual(list(state), ["pair.first", "pair.second"])
-    self.assertIs(state["pair.first"], first)
-    self.assertIs(state["pair.second"], second)
-
-  def test_load_container_subclass(self):
-    class TensorDict(dict): pass
-    weight = Tensor([1., 2.])
-    model = TensorDict(weight=weight)
-    loaded = load_state_dict(model, {"weight": Tensor([3., 4.])}, verbose=False)
-    self.assertEqual(len(loaded), 1)
-    self.assertIs(loaded[0], weight)
-    self.assertEqual(weight.tolist(), [3., 4.])
-
-  def test_container_tensor_attributes(self):
-    class TensorDict(dict): pass
-    class TensorList(list): pass
-    class TensorTuple(tuple): pass
-    for container_type in (TensorDict, TensorList, TensorTuple):
-      with self.subTest(container=container_type.__name__):
-        model = container_type()
-        model.weight = Tensor([1., 2.])
-        state = get_state_dict(model)
-        self.assertEqual(list(state), ["weight"])
-        self.assertIs(state["weight"], model.weight)
-        params = get_parameters(model)
-        self.assertEqual(len(params), 1)
-        self.assertIs(params[0], model.weight)
-        loaded = load_state_dict(model, {"weight": Tensor([3., 4.])}, verbose=False)
-        self.assertEqual(len(loaded), 1)
-        self.assertIs(loaded[0], model.weight)
-        self.assertEqual(model.weight.tolist(), [3., 4.])
-
-  def test_container_contents_and_attributes(self):
-    class TensorDict(dict): pass
-    class TensorList(list): pass
-    class TensorTuple(tuple): pass
-    item, weight = Tensor([1.]), Tensor([2.])
-    for model, key in ((TensorDict(item=item), "item"), (TensorList([item]), "0"), (TensorTuple([item]), "0")):
-      with self.subTest(container=type(model).__name__):
-        model.weight = weight
-        state = get_state_dict(model, prefix="model.")
-        self.assertEqual(list(state), [f"model.{key}", "model.weight"])
-        self.assertIs(state[f"model.{key}"], item)
-        self.assertIs(state["model.weight"], weight)
-
-  def test_container_attribute_precedence(self):
-    class TensorDict(dict): pass
-    model = TensorDict(weight=Tensor([1.]))
-    model.weight = Tensor([2.])
-    self.assertIs(get_state_dict(model)["weight"], model.weight)
-
-
-if __name__ == '__main__':
-  unittest.main()

tinygrad/codegen/__init__.py

@@ -7,7 +7,7 @@ from tinygrad.uop.weak import pm_lower_weak, pm_commit_weak, pm_cast_const
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
-from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
+from tinygrad.renderer.isa import ISARenderer, IselContext
 from tinygrad.dtype import dtypes, AddrSpace
 
 # import all pattern matchers here
@@ -439,12 +439,13 @@ def do_linearize(ctx:Renderer, prg:UOp, sink:UOp) -> UOp:
   lst = line_rewrite(linearize(sink), pm_linearize_cleanups)
   # isa renderers need to allocate registers
   if isinstance(ctx, ISARenderer):
-    lst = line_rewrite(lst, ctx.pre_regalloc_matcher, PreRegAllocContext())
+    lin_ctx = ctx.linear_ctx_type(ctx)
+    lst = line_rewrite(lst, ctx.pre_regalloc_matcher, lin_ctx)
     # register definitions (INS without srcs) move to the top so regalloc sees their live ranges span the whole program (callee saved regs)
     lst = sorted(lst, key=lambda u: u.op is not Ops.INS or bool(u.src))
-    regalloc_ctx = LinearScanRegallocContext(lst, ctx)
+    regalloc_ctx = LinearScanRegallocContext(lin_ctx, lst, ctx)
     lst = line_rewrite(lst, pm_regalloc_rewrite, regalloc_ctx)
-    lst = line_rewrite(lst, ctx.post_regalloc_matcher, regalloc_ctx)
+    lst = line_rewrite(lst, ctx.post_regalloc_matcher, lin_ctx)
     if DEBUG >= 4: print(ctx.asm_str(lst, sink.arg.function_name))
   return prg.replace(src=prg.src + (UOp(Ops.LINEAR, src=tuple(lst)),))
 

tinygrad/codegen/decomp/op.py

@@ -18,29 +18,19 @@ def magicgu(vmax:int, d:int) -> tuple[int,int]:
   assert False
 
 def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None:
-  from tinygrad.renderer.cstyle import MetalRenderer
-  # NOTE: disable for METAL due to compiler bug. keccak with -O0 works but not with optimization
-  if isinstance(ren, MetalRenderer): return None
-  # If d is a power of two this is not valid for signed ints!
-  is_unsigned = x.vmin>=0 or x.dtype in dtypes.uints
-  assert d>0, "Sign should have been taken out of divisor"
-  vmin,vmax = max(x.vmin, x.dtype.min), min(x.vmax, x.dtype.max)
-  if vmin > -d and vmax < d: return x.const_like(0)
-  m,s = magicgu(max(vmax, abs(vmin)), d)
-  if m*vmin >= x.dtype.min and m*vmax <= x.dtype.max:
-    return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)
+  if d <= 0 or x.vmin < 0: return None
+  if (vmax:=min(x.vmax, x.dtype.max)) < d: return x.const_like(0)
+  m,s = magicgu(vmax, d)
+  if m*vmax <= x.dtype.max: return (x*m) >> s
   # before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
-  # use explicit Ops.CDIV (trunc) since the recursion assumes trunc semantics throughout
-  if (largest_factor_of_two_in_d := (d & -d)) > 1:
-    if (ret:=fast_idiv(ren, x.alu(Ops.CDIV, x.const_like(largest_factor_of_two_in_d)),
-                       d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
+  if (k := (d & -d).bit_length()-1) > 0:
+    if (ret:=fast_idiv(ren, x >> k, d >> k, dont_cast=True)) is not None: return ret
   if dont_cast: return None
   # the next integer width that holds x*m
   widen = {dtypes.int8:dtypes.int16, dtypes.int16:dtypes.int32, dtypes.int32:dtypes.int64, dtypes.int64:dtypes.uint64,
            dtypes.uint8:dtypes.uint16, dtypes.uint16:dtypes.uint32, dtypes.uint32:dtypes.uint64}
   if (next_dtype := widen.get(x.dtype)) is not None and next_dtype in ren.supported_dtypes():
-    if m*vmin >= next_dtype.min and m*vmax <= next_dtype.max:
-      return ((x.cast(next_dtype)*m) >> s).cast(x.dtype) if is_unsigned else ((x.cast(next_dtype)*m) >> s).cast(x.dtype) + (x<0).where(x.ufix(1), 0)
+    if m*vmax <= next_dtype.max: return ((x.cast(next_dtype)*m) >> s).cast(x.dtype)
   return None
 
 # ***** threefry *****
@@ -105,13 +95,12 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
       lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v
         if (v:=powers_of_two.get(c.val, 0)) else None)]
     if not disable_fast_idiv:
-      # fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0)
-      pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))),
-        lambda ctx, x, d: fast_idiv(ctx, x, d.val) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
-      # rewrite raw CMOD -> x - d*CDIV(x,d) so fast_idiv can pick up the CDIV. only on non-negative inputs;
+      # fast_idiv handles non-pow2 divisors on non-negative inputs
+      pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))), lambda ctx, x, d: fast_idiv(ctx, x, d.val))]
+      # rewrite raw CMOD -> x - d*fast_idiv(x,d), only when fast_idiv can actually divide;
       # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail)
-      pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.var("d"))),
-        lambda x, d: x - d * x.alu(Ops.CDIV, d) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
+      pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))),
+        lambda ctx, x, d: x - d * q if (q:=fast_idiv(ctx, x, d.val)) is not None else None)]
   if Ops.NEG in ops:
     pat += [(UPat.var('x')*-1, lambda ctx,x: x.alu(Ops.NEG))]
     if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda ctx,x,y: x.alu(Ops.SUB, y))]

tinygrad/codegen/late/regalloc.py

@@ -1,20 +1,18 @@
 import itertools
 from tinygrad.helpers import dedup
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
-from tinygrad.renderer.isa import ISARenderer, Register, greg
-from tinygrad.dtype import dtypes
+from tinygrad.renderer.isa import ISARenderer, Register, rdef, LinearContext
+from typing import Any
 
 PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.BITCAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
 
 class LinearScanRegallocContext:
   # returns the uop that defines the virtual register
   def vdef(self, v:Register) -> UOp: return self.uops[self.live_range[v][0]]
-  def __init__(self, uops:list[UOp], ren:ISARenderer):
+  def __init__(self, ctx:LinearContext, uops:list[UOp], ren:ISARenderer):
     self.uops = uops
     self.ren = ren
     self.idx = itertools.count()
-    # the label associated with each loop NOTE: this is only used post regalloc and should be removed
-    self.loop_label: dict[UOp, str] = {}
 
     # compute live ranges
     self.live_range: dict[Register, list[int]] = {}
@@ -23,16 +21,15 @@ class LinearScanRegallocContext:
     for idx,u in reversed(list(enumerate(uops))):
       if u.op in PSEUDO_OPS: continue
       defs = u.tag if isinstance(u.tag, tuple) else ()
-      for v in defs + tuple(greg(s) for s in dedup(u.src)):
+      for v in defs + tuple(rdef(s) for s in dedup(u.src)):
         if isinstance(v, Register): lr.setdefault(v, []).insert(0, idx)
       for v in defs:
         if v in lr and (n:=max((e for s,e in loops.items() if s <= lr[v][-1] < e), default=None)): lr[v].append(n)
       if u.op is Ops.RANGE: loops[idx] = max(j for j,x in enumerate(uops) if u in x.src)
 
     # allocate registers
-    self.stack_size: int = 0
     self.locals: dict[UOp, UOp] = {}
-    self.spills: dict[Register, UOp] = {} # mapping from virtual to stack slot
+    self.spills: dict[Register, Any] = {} # mapping from virtual to arbitrary spill slot
     self.reals: dict[int, dict[Register, Register]] = {} # mapping from virtual to real at each program point
     self.insert_before: dict[int, list[tuple[Register, Register]]] = {} # fills to be inserted at each program point
     live: dict[Register, Register] = {} # mapping from virtual to real that's currently assigned to it
@@ -49,11 +46,7 @@ class LinearScanRegallocContext:
     # assign register to spilled virtual and record load to be emitted before current uop, also assign it a stack slot
     def fill(v:Register, i:int, cons:tuple[Register, ...]|None=None) -> Register:
       if v not in self.spills:
-        # the value of a BUFFER is its 64bit address, XMM registers need 16 bytes
-        sz = 16 if v.cons[0].size == 16 else (8 if self.vdef(v).op is Ops.BUFFER else self.vdef(v).dtype.itemsize)
-        offset = self.stack_size + (sz - self.stack_size % sz) % sz
-        self.spills[v] = UOp.cconst(offset, dtypes.int32)
-        self.stack_size = offset + sz
+        self.spills[v] = ctx.assign_spill_slot(v, self.vdef(v))
       r = alloc(cons if cons is not None else v.cons, i)
       self.insert_before.setdefault(i, []).append((v, r))
       return r
@@ -64,7 +57,7 @@ class LinearScanRegallocContext:
       for s in u.src:
         # HACK: cause of later hacks to lower range
         if u.op is Ops.END: continue
-        if not isinstance(v:=greg(s), Register): continue
+        if not isinstance(v:=rdef(s), Register): continue
         if v not in live: live[v] = fill(v, i)
         self.reals.setdefault(i, {})[v] = live[v]
 
@@ -76,17 +69,12 @@ class LinearScanRegallocContext:
           cons = v.cons
           # two address instructions (src is reused by def) can only coalesce reused src. reused src goes first to get priority in case of a tiebreak
           if ren.is_two_address(u) and j == 0:
-            uses = tuple(live.get(greg(s)) for s in u.src)
+            uses = tuple(live.get(rdef(s)) for s in u.src)
             cons = ((uses[0],) if uses[0] in cons else ()) + tuple(r for r in cons if r not in uses)
           # HACK: cause the range is missing the comparison
           live[v] = alloc(cons, i+1 if u.op is not Ops.RANGE else i)
           self.reals.setdefault(i, {})[v] = live[v]
 
-      # allocate stack array
-      if u.op is Ops.BUFFER:
-        self.locals[u] = UOp.cconst(self.stack_size, dtypes.int32)
-        self.stack_size += u.max_numel() * u.dtype.itemsize
-
       # loop prologue, avoid loading inside the loop
       if u.op is Ops.RANGE:
         # we move to registers vars used in the loop sorted by next use, vars not used in the loop will not be reloaded in the epilogue
@@ -113,22 +101,14 @@ def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
   nsrc = []
   for j,s in enumerate(x.src):
     # v here is the virtual defined by the original s as s is the rewritten version
-    if i in ctx.reals and (v:=greg(ctx.uops[i].src[j])) in ctx.spills: nsrc.append(ctx.ren.fill(ctx.spills[v], ctx.vdef(v), ctx.reals[i][v]))
+    if i in ctx.reals and (v:=rdef(ctx.uops[i].src[j])) in ctx.spills: nsrc.append(ctx.ren.fill(ctx.spills[v], ctx.vdef(v), ctx.reals[i][v]))
     else: nsrc.append(s)
   ndefs = tuple(ctx.reals[i][v] for v in x.tag) if isinstance(x.tag, tuple) else x.tag
-  if x.op is Ops.BUFFER: nx = ctx.ren.isel_matcher.rewrite(ctx.ren.stack_pointer().index(ctx.locals[x], tag=ndefs))
-  else: nx = x.replace(src=tuple(nsrc), tag=ndefs)
+  nx = x.replace(src=tuple(nsrc), tag=ndefs)
 
   before = [ctx.ren.fill(ctx.spills[v], ctx.vdef(v), r) for v,r in ctx.insert_before.get(i, [])]
   after = [ctx.ren.spill(ctx.spills[v], nx) for v in x.tag if v in ctx.spills] if isinstance(x.tag, tuple) else []
 
-  # alloc/dealloc stack
-  if ctx.stack_size > 0:
-    sp = ctx.ren.stack_pointer()
-    offset = UOp.cconst(ctx.stack_size, sp.dtype)
-    if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, src=(sp, offset), tag=sp.tag))] + before
-    elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, src=(sp, offset), tag=sp.tag))]
-
   return nx, before + [nx] + after
 
 pm_regalloc_rewrite = PatternMatcher([

tinygrad/engine/realize.py

@@ -52,6 +52,8 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N
 # **************** Stat ****************
 
 def estimate_uop(call:UOp) -> Estimates:
+  call = call.without_after
+  if isinstance(call.arg.aux, HCQInfo): return call.arg.aux.estimates
   if (ast:=call.src[0]).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"):
     return Estimates(lds=(nbytes:=prod(call.src[1].shape) * call.src[1].dtype.itemsize), mem=nbytes)
@@ -192,20 +194,21 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
 
 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   if (info:=call.arg.aux).inputs:
-    addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(dev).va_addr for u, dev in info.inputs]
+    addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(dev).va_addr + off for u, dev, off in info.inputs]
     cast(Buffer, call.src[1 + info.table].buffer)._buf.cpu_view().view(fmt='Q')[:] = array.array('Q', addrs)
   ctx = replace(ctx, var_vals={**ctx.var_vals, **{k: v for d in info.device for k, v in cast(Any, Device[d]).var_vals.items()}})
   ets = exec_kernel(ctx, call, ast, devices=(HCQ_RUNTIME_DEV.value,))
   if not (ctx.wait or PROFILE): return ets
 
   slots = {d: cast(Buffer, call.src[1 + i].buffer) for d, i in info.slots}
-  def _prof_tm(device:str, name:str, prof:tuple[int, ...], profile_key:bytes) -> float|None:
-    (d:=cast(Any, Device[device])).prof_ents[(slots[device], prof[0])] = ProfileGraphEntry(device, name, prof[0], prof[1], profile_key)
-    if not ctx.wait: return None
-    d.synchronize(timeout=ctx.timeout)
+  for devs, name, _, prof, pkey in info.kernels:
+    for d in (devs if prof else ()): cast(Any, Device[d]).prof_ents[(slots[d], prof[0])] = ProfileGraphEntry(d, name, prof[0], prof[1], pkey)
+  if ctx.wait:
+    for device in info.device: cast(Any, Device[device]).synchronize(timeout=ctx.timeout)
+  def _prof_tm(device:str, prof:tuple[int, ...]) -> float:
     st, en = (slots[device]._buf.cpu_view().view(fmt='Q')[x] for x in prof)
-    return float(en-st) / d.timestamp_divider / 1e6
-  return ets + [_prof_tm(device, name, prof, profile_key) for devices,name,_,prof,profile_key in info.kernels if prof for device in devices]
+    return float(en-st) / cast(Any, Device[device]).timestamp_divider / 1e6
+  return ets + [_prof_tm(device, prof) if ctx.wait else None for devices, _, _, prof, _ in info.kernels if prof for device in devices]
 
 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([

tinygrad/helpers.py

@@ -166,6 +166,8 @@ def stderr_log(msg:str): print(msg, end='', file=sys.stderr, flush=True)
 
 class Context(contextlib.ContextDecorator):
   def __init__(self, **kwargs): self.kwargs = kwargs
+  # ContextDecorator otherwise reuses self, so recursive calls overwrite old_context.
+  def _recreate_cm(self): return Context(**self.kwargs)
   def __enter__(self):
     self.old_context:dict[str, Any] = {k: ContextVar._cache[k].value for k in self.kwargs}
     for k,v in self.kwargs.items(): ContextVar._cache[k].value = v
@@ -239,7 +241,7 @@ TRANSCENDENTAL = ContextVar("TRANSCENDENTAL", 1)
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("LRU", 1)
 RING, ALL2ALL, ALLREDUCE_CAST = ContextVar("RING", 1), ContextVar("ALL2ALL", 0), ContextVar("ALLREDUCE_CAST", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0)
-VALIDATE_WITH_CPU, HCQ2 = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("HCQ2", 0)
+VALIDATE_WITH_CPU, HCQ2 = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("HCQ2", 1)
 # TODO: this is broken for some indexing
 DISABLE_FAST_IDIV = ContextVar("DISABLE_FAST_IDIV", 1)
 FUSE_OPTIM = ContextVar("FUSE_OPTIM", 0)

tinygrad/mixin/elementwise.py

@@ -647,9 +647,10 @@ class ElementwiseMixin(CreationMixin):
     ```
     """
     other = self.ufix(other)
-    error = (self - other).abs()
-    is_finite_close = error.isfinite() & (error <= atol + rtol * other.abs())
-    return self.eq(other) | is_finite_close | (self.isnan() & other.isnan() & equal_nan)
+    is_finite_close = self.isfinite() & other.isfinite() & ((self - other).abs() <= atol + rtol * other.abs())
+    is_infinite_close = (self.isinf() | other.isinf()) & self.eq(other)
+    is_nan_close = (self.isnan() & other.isnan()) & equal_nan
+    return is_finite_close | is_infinite_close | is_nan_close
 
   def ceil(self) -> Self:
     """
@@ -704,7 +705,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([-9., -6., -3., 0., 3., 6., 9.]).relu6().numpy())
     ```
     """
-    return self.relu() - (self-6).relu()
+    return self.relu().minimum(6)
 
   def hardswish(self) -> Self:
     """
@@ -1086,7 +1087,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([1., 2., 3.]).lerp(Tensor([4., 5., 6.]), 0.5).numpy())
     ```
     """
-    if self.dtype == dtypes.uint8 and not end.is_floating_point() and not isinstance(weight, ConstType):
-      weight_int = (weight * 128 + 0.5).cast(dtypes.int32)  # 7 fractional bits
-      return ((self * (128 - weight_int) + end.cast(dtypes.int32) * weight_int + 64) >> 7).cast(dtypes.uint8)
+    if self.dtype == dtypes.uint8 and not isinstance(weight, ConstType):
+      w_i = (weight * (1<<(W_PREC:=7)) + 0.5).cast(dtypes.int16)
+      return (self+(((end - self).cast(dtypes.int8) * w_i + (1<<W_PREC-1)).cast(dtypes.uint16) >> W_PREC)).cast(dtypes.uint8)
     return self + (end - self) * weight

tinygrad/mixin/movement.py

@@ -493,8 +493,9 @@ class MovementMixin:
     ```
     """
     if indexing not in ("ij", "xy"): raise RuntimeError(f'indexing must be in ("ij", "xy"), got {indexing}')
-    basis = tuple(range(len(args)+1)) if indexing == "ij" or not args else (1, 0) + tuple(range(2, len(args)+1))
-    tensors = tuple(t.reshape((-1,) + (1,)*(len(args) - i)) for i,t in zip(basis, (self, *args)))
+    if len(tensors:=(self, *args)) == 1: return tensors
+    basis = tuple(range(len(tensors))) if indexing == "ij" else (1, 0) + tuple(range(2, len(tensors)))
+    tensors = tuple(t.reshape((-1,) + (1,)*(len(args) - i)) for i,t in zip(basis, tensors))
     output_shape = _broadcast_shape(*(t.shape for t in tensors))
     return tuple(t._broadcast_to(output_shape) for t in tensors)
 
@@ -527,8 +528,8 @@ class MovementMixin:
     """
     if (dim1:=self._resolve_dim(dim1)) == (dim2:=self._resolve_dim(dim2)): raise RuntimeError("dim1 and dim2 cannot be the same dimension")
     x = self.permute(*[i for i in range(self.ndim) if i != dim1 and i != dim2], dim1, dim2)
-    if offset >= 0: x = x.shrink((None,)*(x.ndim-1) + ((min(offset, x.shape[-1]), x.shape[-1]),))
-    else: x = x.shrink((None,)*(x.ndim-2) + ((min(-offset, x.shape[-2]), x.shape[-2]), None))
+    if offset >= 0: x = x.shrink(tuple(None for _ in x.shape[:-1]) + ((offset, x.shape[-1]),))
+    else: x = x.shrink(tuple(None for _ in x.shape[:-2]) + ((-offset, x.shape[-2]), None))
     if (d := min(int(x.shape[-2]), int(x.shape[-1]))) <= 0: return x.reshape(*x.shape[:-2], 0)
     nones, x = tuple(None for _ in x.shape[:-2]), x.shrink_to(tuple(None for _ in x.shape[:-2]) + (d, d))
     return x.flatten(-2).pad_to(nones+(d*(d+1),)).unflatten(-1, (d, d+1)).shrink_to(nones+(None, 1)).squeeze(-1)

tinygrad/mixin/op.py

@@ -206,7 +206,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     if steps < 0: raise ValueError("number of steps must be non-negative")
     if (dtype := to_dtype(dtype or dtypes.default_float)) == dtypes.bool: raise ValueError("linspace with bool dtype is not supported")
     if steps == 1: return cls.full((1,), start, dtype=dtype, buffer=False)
-    return (start + cls.arange(steps, dtype=least_upper_dtype(dtype, dtypes.default_float)) * ((stop - start) / (steps - 1))).cast(dtype)
+    return (start + cls.arange(steps, dtype=dtypes.default_float) * ((stop - start) / (steps - 1))).cast(dtype)
 
   @classmethod
   def eye(cls, n:int, m:int|None=None, dtype:DTypeLike|None=None) -> Self:
@@ -542,10 +542,11 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     print(t.var(axis=1).numpy())
     ```
     """
+    output_dtype = self.dtype if dtypes.is_float(self.dtype) else dtypes.float32
     squares = (self - self.mean(axis=axis, keepdim=True)).square()
     n = prod([si for si, so in zip(self.shape, squares.sum(axis=axis, keepdim=True).shape) if resolve(si != so)])
-    numerator = squares.sum(axis=axis, keepdim=keepdim, dtype=sum_acc_dtype(squares.commit_dtype()))
-    return numerator.div(smax(n - correction, 0)).cast(squares.dtype)
+    numerator = squares.cast(sum_acc_dtype(self.commit_dtype())).sum(axis=axis, keepdim=keepdim)
+    return numerator.div(smax(n - correction, 0)).cast(output_dtype)
 
   def var_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Self, Self]:
     """
@@ -806,10 +807,11 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     ```
     """
     if self.ndim == 0: return self._split_cumalu(axis, Ops.MAX), type(self).zeros(self.shape, dtype=dtypes.int32, buffer=False)
-    values = self._split_cumalu(axis, Ops.MAX)
-    # Record the latest index matching the running maximum, then carry it forward.
-    idx = self.eq(values).transpose(axis, -1) * type(self).arange(self.shape[axis], dtype=dtypes.int32)
-    return values, idx._split_cumalu(-1, Ops.MAX).transpose(-1, axis)
+    values, n = self._split_cumalu(axis, Ops.MAX), int(self.shape[axis])
+    x, values_t = self.transpose(axis, -1), values.transpose(axis, -1)
+    match = x.unsqueeze(-1).eq(values_t.unsqueeze(-2)) * self._tri(n, n)
+    idx = (-(match * type(self).arange(n, 0, -1).reshape(n, 1)).max(-2) + n).cast(dtypes.int32)
+    return values, idx.transpose(-1, axis)
 
   def cummin(self, axis:int=0) -> tuple[Self, Self]:
     """
@@ -849,12 +851,14 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     print(t.logcumsumexp(axis=1).numpy())
     ```
     """
-    axis = self._resolve_dim(axis)
     if self.ndim == 0: return self
     x = self.transpose(axis, -1)
-    mask = self._tri(x.shape[-1], x.shape[-1], 1)
-    prefixes = mask.where(-math.inf, x.unsqueeze(-2))
-    return prefixes.logsumexp(-1).transpose(-1, axis)
+    last_dim_size = x.shape[-1]
+    x_unsqueezed = x.unsqueeze(-2)
+    x_cummax = (mx:=x.cummax(-1)[0].detach()).isfinite().where(mx, 0)
+    mask = self._tri(last_dim_size, last_dim_size, 1).logical_not()
+    ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax
+    return ret.transpose(-1, axis)
 
   def argmax(self, axis=None, keepdim=False) -> Self:
     """
@@ -1734,10 +1738,10 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     if Y.device is not None and self.device is not None and Y.device != self.device:
       raise RuntimeError(f"expected Y and self on the same device, {Y.device=}, {self.device=}")
     log_probs = self.log_softmax()
-    loss_mask = Y.ne(ignore_index)
-    y = Y.unsqueeze(-1)._one_hot_along_dim(self.shape[-1], dim=-1)
-    smoothing = label_smoothing * log_probs.mean(-1)
-    unreduced = ((1 - label_smoothing) * (log_probs * y).sum(-1) + smoothing) * loss_mask
+    loss_mask = Y.ne(ignore_index) if ignore_index != -1 else Y.const_like(True, dtypes.bool)
+    y = Y.unsqueeze(-1)._one_hot_along_dim(self.shape[-1], dim=-1) * loss_mask.unsqueeze(-1)
+    smoothing = label_smoothing * (log_probs.mean(-1) * loss_mask)
+    unreduced = ((1 - label_smoothing) * (log_probs * y).sum(-1) + smoothing)
     return -unreduced.sum() / loss_mask.sum() if reduction == "mean" else -unreduced._do_reduction(reduction)
 
   def cross_entropy(self, Y:Self, reduction:ReductionStr="mean", label_smoothing:float=0.0) -> Self:

tinygrad/nn/state.py

@@ -1,4 +1,5 @@
 import json, pathlib, struct, functools, io, zlib
+from collections import OrderedDict
 from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
@@ -101,12 +102,13 @@ def get_state_dict(obj, prefix:str='', tensor_type=Tensor) -> dict[str, Tensor]:
   """
   if isinstance(obj, tensor_type): return {prefix.strip('.'):obj}
   if hasattr(obj, '_asdict'): return get_state_dict(obj._asdict(), prefix, tensor_type)  # namedtuple
+  if isinstance(obj, OrderedDict): return get_state_dict(dict(obj), prefix, tensor_type)
+  if hasattr(obj, '__dict__'): return get_state_dict(obj.__dict__, prefix, tensor_type)
   state_dict = {}
   if isinstance(obj, (list, tuple)):
     for i,x in enumerate(obj): state_dict.update(get_state_dict(x, f"{prefix}{str(i)}.", tensor_type))
   elif isinstance(obj, dict):
     for k,v in obj.items(): state_dict.update(get_state_dict(v, f"{prefix}{str(k)}.", tensor_type))
-  if hasattr(obj, '__dict__'): state_dict.update(get_state_dict(obj.__dict__, prefix, tensor_type))
   return state_dict
 
 def get_parameters(obj) -> list[Tensor]:

tinygrad/renderer/amd/sqtt.py

@@ -670,7 +670,6 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
       for wave in range(10):
         if (p.inst >> (wave * 2)) & 3 == 3:
           inst = pc_map[pc:=wave_pc[(p.simd, wave)]]
-          wave_pc[(p.simd, wave)] += inst.size()
           yield (p, InstructionInfo(pc, wave, inst))
     # map INST events on this SIMD to the program counter, we know the waves
     elif isinstance(p, (VALUINST, INST, INST_RDNA4, IMMEDIATE)) and not (isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("OTHER_")):

tinygrad/renderer/isa/__init__.py

@@ -3,6 +3,7 @@ import itertools
 from dataclasses import dataclass, field
 from tinygrad.renderer import Renderer
 from tinygrad.uop.ops import PatternMatcher, UOp, Ops
+from typing import Any
 
 @dataclass(frozen=True)
 class Register:
@@ -23,23 +24,24 @@ class IselContext:
   def vreg(self, cons:tuple[Register, ...]|Register):
     return Register(f"v{next(self.reg_n)}", 0, _cons=cons if isinstance(cons, tuple) else (cons,))
 
-def greg(u:UOp):
-  if u.op in {Ops.NOOP, Ops.AFTER, Ops.BITCAST} and u.src: return greg(u.src[0])
-  if isinstance(u.tag, tuple): return u.tag[0]
-  return u.tag
+def rdef(u:UOp):
+  if u.op in {Ops.NOOP, Ops.AFTER, Ops.BITCAST} and u.src: return rdef(u.src[0])
+  return u.tag[0] if isinstance(u.tag, tuple) else u.tag
 
-@dataclass
-class PreRegAllocContext:
-  lock: UOp|None = None
+class LinearContext:
+  def __init__(self, ren:ISARenderer):
+    self.ren, self.stack_size = ren, 0
+    self.loop_label: dict[UOp, str] = {}
+  def assign_spill_slot(self, r:Register, u:UOp) -> Any: raise NotImplementedError("arch specific")
 
 class ISARenderer(Renderer):
   pre_isel_matcher: PatternMatcher
   isel_matcher: PatternMatcher
   pre_regalloc_matcher: PatternMatcher
   post_regalloc_matcher: PatternMatcher
+  linear_ctx_type: type = LinearContext
 
   def is_two_address(self, x:UOp) -> bool: return False
-  def stack_pointer(self) -> UOp: raise NotImplementedError("arch specific")
-  def spill(self, disp:UOp, x:UOp) -> UOp: raise NotImplementedError("arch specific")
-  def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
+  def spill(self, spill_slot:Any, x:UOp) -> UOp: raise NotImplementedError("arch specific")
+  def fill(self, spill_slot:Any, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
   def asm_str(self, uops:list[UOp], function_name:str) -> str: raise NotImplementedError("arch specific")

tinygrad/renderer/isa/x86.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 # flake8: noqa: E702
 # allow semicolons to put multiple ops on one line
 import sys, struct, functools
@@ -6,7 +7,7 @@ from dataclasses import replace
 from tinygrad.dtype import dtypes, DType, truncate, AddrSpace
 from tinygrad.uop import FastEnum, auto, Ops, GroupOp
 from tinygrad.uop.ops import UOp, UPat, PatternMatcher, promo_dtype
-from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext, greg
+from tinygrad.renderer.isa import ISARenderer, IselContext, Register, LinearContext, rdef
 from tinygrad.helpers import unwrap, Target
 
 # ***** X86 Ops *****
@@ -158,6 +159,9 @@ pre_isel_matcher = PatternMatcher([
 ])
 
 # ***** X86 registers *****
+def def_reg(dt:DType, reg:Register) -> UOp: return UOp(Ops.INS, arg=(X86Ops.DEFINE, dt), tag=(reg,))
+# undefined operand, used for VEX instructions
+def undef(): return UOp(Ops.NOOP)
 
 RAX = Register("rax", 0)
 RCX = Register("rcx", 1)
@@ -178,11 +182,12 @@ reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"},
         "rsp": {4:"esp", 2:"sp", 1:"spl"}, "rbp": {4:"ebp", 2:"bp", 1:"bpl"}, "rsi": {4:"esi", 2:"si", 1:"sil"}, "rdi": {4:"edi", 2:"di", 1:"dil"},
         **{f"r{i}": {4:f"r{i}d", 2:f"r{i}w", 1:f"r{i}b"} for i in range(8, 16)}}
 
+stack_pointer = def_reg(dtypes.uint64, RSP)
+
 # ***** X86 instruction selection *****
 def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.INDEX else s
 def lane(x:UOp, i:int) -> int: return s.src[1].src[0].val if (s:=x.src[i]).op is Ops.INDEX else 0
 def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt]
-def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, arg=(X86Ops.DEFINE, dt), tag=None if reg is None else (reg,))
 def imm(dt:DType, v:int) -> UOp: return UOp.cconst(truncate[dt](v), dt).rtag()
 def to_imm(c:UOp) -> UOp|None:
   if not (c.op is Ops.CAST and (v:=c.src[0]).op is Ops.CONST): return None
@@ -206,13 +211,13 @@ def vinsertps(x:UOp) -> UOp:
   def _insert(ret:UOp, i:int) -> UOp:
     s, v = base(x, i), lane(x, i)
     return x.ins(X86Ops.VINSERTPS, src=(ret, s, imm(dtypes.uint8, v << 6 | i << 4)))
-  return functools.reduce(_insert, range(len(x.src)), def_reg(x.dtype))
+  return functools.reduce(_insert, range(len(x.src)), undef())
 
 # vpinsrd xmm2, xmm0, eax, imm
 # inserts the element in eax into any position in xmm0, result is written to xmm2 according to imm
 def vpins(x:UOp, srcs:tuple[UOp, ...]) -> UOp:
   op = {2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD}[x.dtype.itemsize]
-  return functools.reduce(lambda ret,i: x.ins(op, src=(ret, srcs[i], imm(dtypes.uint8, i))), range(len(srcs)), def_reg(x.dtype))
+  return functools.reduce(lambda ret,i: x.ins(op, src=(ret, srcs[i], imm(dtypes.uint8, i))), range(len(srcs)), undef())
 
 # we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg
 def idiv(ctx:IselContext, x:UOp) -> UOp:
@@ -265,7 +270,7 @@ def abi(ctx:IselContext, x:UOp) -> UOp|None:
   # the shape srcs of a PARAM are not values, tag them so they aren't materialized into registers
   def _reg_arg(r:Register) -> tuple[UOp, ...]: return (x.replace(arg=arg, src=tuple(s.rtag() for s in x.src), tag=(r,)),)
   def _stack_arg(disp:int):
-    return (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), UOp(Ops.INS, arg=(X86Ops.FRAME_INDEX, dtypes.int32), tag=disp), imm(dtypes.uint8, 8))
+    return (stack_pointer, UOp(Ops.NOOP), UOp(Ops.INS, arg=(X86Ops.FRAME_INDEX, dtypes.int32), src=(imm(dtypes.int32, disp),)), imm(dtypes.uint8, 8))
   if sys.platform == "win32": src = _reg_arg((RCX, RDX, GPR[8], GPR[9])[i]) if i < 4 else _stack_arg((i-3)*8+32)
   else: src = _reg_arg((RDI, RSI, RDX, RCX, GPR[8], GPR[9])[i]) if i < 6 else _stack_arg((i-5)*8)
   # this move "cleanses" the abi register constraint
@@ -320,9 +325,9 @@ isel_matcher = PatternMatcher([
     lambda x,cond: cond.ins(X86Ops.LOOP_CMP, tag=cond.op, src=cond.src + x.src[:2])),
   # **** Op -> X86Op ****
   # add callee saved registers to the RET, these will be scheduled at the top of the kernel and will be saved/restored if they are used in regalloc
-  # so regalloc builds the prologue/epilogue naturally
+  # so regalloc builds the prologue/epilogue naturally. they all share the stack pointer define's dtype so the the stack pointer define is first
   (UPat(Ops.SINK, name="x"), lambda x:
-   x.replace(src=(x.ins(X86Ops.RET, src=x.src + tuple(def_reg(dtypes.uint64 if r in GPR else dtypes.float64, r) for r in CALLEE_SAVED)),)) \
+   x.replace(src=(x.ins(X86Ops.RET, src=x.src + (stack_pointer,) + tuple(def_reg(dtypes.uint64, r) for r in CALLEE_SAVED)),))
     if not x.src or x.src[0].op is not Ops.INS or x.src[0].arg[0] is not X86Ops.RET else None),
   # function abi constraints
   (UPat((Ops.PARAM, Ops.SPECIAL), name="x"), abi),
@@ -417,8 +422,8 @@ isel_matcher = PatternMatcher([
   (UPat(dtype=dtypes.float64).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSD2SI)),
   (UPat.var("y", dtypes.float32).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSS2SD, src=(y, y))),
   (UPat.var("y", dtypes.float64).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSD2SS, src=(y, y))),
-  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(def_reg(x.dtype), y))),
-  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(def_reg(x.dtype), y))),
+  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(undef(), y))),
+  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(undef(), y))),
   (UPat(dtype=(dtypes.uint8, dtypes.uint16, dtypes.bool)).cast(dtypes.ints, name="x"), lambda x:
    x.ins(X86Ops.MOVZX) if x.src[0].dtype.itemsize < x.dtype.itemsize else None),
   (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD)),
@@ -436,7 +441,7 @@ isel_matcher = PatternMatcher([
   # TODO: fuse stores, very few cases -- store cmp becomes setcc, store gep int becomes vpextr, store bitcast to int becomes vmovd/q
   # load, store
   (UPat(Ops.LOAD, dtypes.floats, src=(UPat(name="a"),), name="x"), lambda x,a:
-   x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (imm(dtypes.uint8, 0),)) if x.max_numel() * x.dtype.itemsize == 2 else
+   x.ins(X86Ops.VPINSRW, src=(undef(),) + fold_address(a) + (imm(dtypes.uint8, 0),)) if x.max_numel() * x.dtype.itemsize == 2 else
    x.ins(_xmm_sz(x), src=fold_address(a))),
   (UPat(Ops.LOAD, dtypes.ints+(dtypes.bool,), src=(UPat(name="a"),), name="x"), lambda x,a:
    x.ins(X86Ops.MOV, src=fold_address(a)) if x.max_numel() == 1 else x.ins(_xmm_sz(x), src=fold_address(a))),
@@ -454,14 +459,21 @@ isel_matcher = PatternMatcher([
 # the flags belong to the last instruction that wrote them. x86 has no good way to store/restore them (then regalloc would
 # handle it), so a consumer that no longer owns its compare re-emits it. Unlike a regalloc rematerialization this is not
 # optional, there is no fallback load from stack
-def flag_rematerialize(ctx:PreRegAllocContext, x:UOp):
+def flag_rematerialize(ctx:X86LinearContext, x:UOp):
   if x.op in (Ops.RANGE, Ops.END) or x.arg[0] in X86GroupOp.WriteFlags: ctx.lock = x
   elif x.arg[0] in X86GroupOp.ReadFlags and ctx.lock is not (flag_def:=x.src[-1]):
     ctx.lock = flag_def
     return (x, [flag_def, x])
   return None
 
+# TODO: dont use rewrite
+def alloc_buffer(ctx:X86LinearContext, x:UOp):
+  nx = isel_matcher.rewrite(stack_pointer.index(UOp.cconst(ctx.stack_size, dtypes.uint32), tag=x.tag))
+  ctx.stack_size += x.max_numel() * x.dtype.itemsize
+  return nx, [nx]
+
 pre_regalloc_matcher = PatternMatcher([
+  (UPat(Ops.BUFFER, name="x"), alloc_buffer),
   (UPat((Ops.INS, Ops.RANGE, Ops.END), name="x"), flag_rematerialize),
 ])
 
@@ -492,8 +504,14 @@ def lower_loop(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
 
 # final rewrite to match the isa spec
 post_regalloc_matcher = PatternMatcher([
+  # the frame is allocated after the stack pointer define at the top of the program and freed before RET
+  (UPat(Ops.INS, name="x"), lambda ctx,x: (x, [x, x.ins(X86Ops.SUBi, src=(imm(dtypes.int32, ctx.stack_size),))])
+    if ctx.stack_size and x.arg[0] is X86Ops.DEFINE and rdef(x) == RSP else None),
+  (UPat(Ops.INS, name="x"), lambda ctx,x: (x, [stack_pointer.ins(X86Ops.ADDi, src=(imm(dtypes.int32, ctx.stack_size),)), x])
+    if ctx.stack_size and x.arg[0] is X86Ops.RET else None),
   # rewrite FRAME_INDEX to IMM now that the stack size is known
-  (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=UOp.cconst(ctx.stack_size + x.tag, x.dtype), [nx]) if x.arg[0] is X86Ops.FRAME_INDEX else None),
+  (UPat(Ops.INS, src=(UPat.cvar("disp").cast(),), name="x"), lambda ctx,disp,x:
+    (nx:=UOp.cconst(ctx.stack_size + disp.val, x.dtype), [nx]) if x.arg[0] is X86Ops.FRAME_INDEX else None),
   # expand the cmp here so we can preserve rng src edge to get label from ctx
   (UPat(Ops.INS, name="x"), lambda ctx,x: lower_loop(ctx, x) if x.arg[0] is X86Ops.LOOP_CMP else None),
   # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound
@@ -502,7 +520,7 @@ post_regalloc_matcher = PatternMatcher([
   (UPat(Ops.END, name="x"), lower_end),
   # rewrite two address instructions to two address form, if reused src wasn't coalesced insert a move
   (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=x.replace(src=x.src[1:]),
-   [ctx.ren.copy(x.src[0], greg(x)), nx] if greg(x) != greg(x.src[0]) else [nx]) if x.arg[0] in X86GroupOp.TwoAddress else None),
+   [ctx.ren.copy(x.src[0], rdef(x)), nx] if rdef(x) != rdef(x.src[0]) else [nx]) if x.arg[0] in X86GroupOp.TwoAddress else None),
 ])
 
 # ***** X86 instruction encoding *****
@@ -512,9 +530,9 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
               vvvv_uop:UOp|None=None, imm_uop:UOp|None=None) -> bytes:
     nonlocal reg, opc
     # get the encoding values of the different fields
-    reg = cast(int, cast(Register, greg(reg_uop)).index if reg_uop is not None else reg)
-    rm = cast(Register, greg(rm_uop)).index
-    idx = cast(Register, greg(idx_uop)).index if idx_uop is not None and greg(idx_uop) is not None else 4
+    reg = cast(int, cast(Register, rdef(reg_uop)).index if reg_uop is not None else reg)
+    rm = cast(Register, rdef(rm_uop)).index
+    idx = cast(Register, rdef(idx_uop)).index if idx_uop is not None and rdef(idx_uop) is not None else 4
     # for a memory operand the rm size is the element size from the address, otherwise it's the size of the value in the register
     rm_sz = sz_uop.src[0].val if sz_uop is not None else rm_uop.dtype.itemsize
     reg_sz = reg_uop.dtype.itemsize if reg_uop is not None else 0
@@ -526,7 +544,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     # r extends reg field, x extends index field, b extends rm or base field
     r, _x, b = reg >> 3, idx >> 3, rm >> 3
     if sel: # VEX bytes
-      vvvv = cast(Register, greg(vvvv_uop)).index if vvvv_uop is not None else 0
+      vvvv = (vd.index if isinstance(vd := rdef(vvvv_uop), Register) else reg) if vvvv_uop is not None else 0
       if sel == 1 and _x == b == we == 0: inst += bytes([0xC5, (~r & 0b1) << 7 | (~vvvv & 0b1111) << 3 | pp])
       else: inst += bytes([0xC4, (~r & 0b1) << 7 | (~_x & 0b1) << 6 | (~b & 0b1) << 5 | sel, we << 7 | (~vvvv & 0b1111) << 3 | pp])
     else: # optional PREFIX and REX bytes
@@ -571,7 +589,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     # IMM byte
     if imm_uop is not None:
       if imm_uop.op is Ops.CAST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.src[0].val)
-      elif isinstance(greg(imm_uop), Register): inst += bytes([(greg(imm_uop).index & 0b1111) << 4 | 0b0000])
+      elif isinstance(rdef(imm_uop), Register): inst += bytes([(rdef(imm_uop).index & 0b1111) << 4 | 0b0000])
     return inst
 
   # get the encoding structure of the uop
@@ -604,7 +622,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
 encodings = {
   # moves
   X86Ops.MOVABS: lambda x:
-   bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].src[0].val),
+   bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | rdef(x).index >> 3, 0xB8 + (rdef(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].src[0].val),
   X86Ops.MOV: lambda x: encode(x, 0x8B), X86Ops.MOVi: lambda x: encode(x, 0xC7, reg=0),
   X86Ops.MOVm: lambda x: encode(x, 0x89), X86Ops.LEA: lambda x: encode(x, 0x8D),
   X86Ops.VMOVSS: lambda x: encode(x, 0x10, pp=2, sel=1), X86Ops.VMOVSSm: lambda x: encode(x, 0x11, pp=2, sel=1),
@@ -666,6 +684,16 @@ encodings = {
   X86Ops.RET: lambda x: bytes([0xC3]),
 }
 
+class X86LinearContext(LinearContext):
+  def __init__(self, ren:X86Renderer):
+    super().__init__(ren)
+    self.lock: UOp|None = None
+  def assign_spill_slot(self, r:Register, u:UOp) -> int:
+    sz = r.cons[0].size
+    offset = self.stack_size + (sz - self.stack_size % sz) %sz
+    self.stack_size = offset + sz
+    return offset
+
 class X86Renderer(ISARenderer):
   device = "CPU"
   has_local = False
@@ -676,34 +704,36 @@ class X86Renderer(ISARenderer):
   pre_regalloc_matcher = pre_regalloc_matcher
   post_regalloc_matcher = post_regalloc_matcher
   code_for_op = {x: lambda: None for x in (Ops.SQRT, Ops.AND, Ops.OR, Ops.SHL, Ops.SHR, Ops.NEG, Ops.SUB, Ops.FDIV, Ops.CMPLT, Ops.CMPEQ)}
+  linear_ctx_type = X86LinearContext
   def __init__(self, target:Target):
     if target.arch.split(",")[0] != "x86_64": raise RuntimeError(f"X86Renderer only supports x86_64, got {target.arch}")
     super().__init__(target)
     from tinygrad.runtime.support.compiler_cpu import X86Compiler
     self.compiler = X86Compiler()
   def is_two_address(self, x:UOp) -> bool: return x.op is Ops.INS and x.arg[0] in X86GroupOp.TwoAddress
-  def stack_pointer(self) -> UOp: return def_reg(dtypes.uint64, RSP)
   def copy(self, x:UOp, reg:Register) -> UOp: return x.ins(X86Ops.MOV, src=(x,), tag=reg)
 
-  def spill(self, disp:UOp, x:UOp) -> UOp:
+  def spill(self, spill_slot:int, x:UOp) -> UOp:
     is_xmm = isinstance(x.tag, tuple) and x.tag[0].cons[0].size == 16
     op = X86Ops.VMOVUPSm if is_xmm else X86Ops.MOVm
-    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)) + (x,), arg=(op, dtypes.void), tag=x.tag)
+    disp = UOp.cconst(spill_slot, dtypes.int32)
+    return UOp(Ops.INS, src=fold_address(stack_pointer.index(disp)) + (x,), arg=(op, dtypes.void), tag=x.tag)
 
   # the value of a BUFFER is its address, it moves through registers and the stack as a 64bit int
-  def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp:
+  def fill(self, spill_slot:int, x:UOp, reg:Register) -> UOp:
     is_xmm = reg.cons[0].size == 16
     dt = dtypes.uint64 if x.op is Ops.BUFFER else x.dtype
-    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)), arg=(X86Ops.VMOVUPS if is_xmm else X86Ops.MOV, dt), tag=(reg,))
+    disp = UOp.cconst(spill_slot, dtypes.int32)
+    return UOp(Ops.INS, src=fold_address(stack_pointer.index(disp)), arg=(X86Ops.VMOVUPS if is_xmm else X86Ops.MOV, dt), tag=(reg,))
 
   def asm_str(self, uops:list[UOp], function_name:str) -> str:
     def _format_op(x:UOp) -> str: return f"    {(o[7:-1] if (o:=str(x.arg[0]))[-1] in ('i', 'm') else o[7:]).lower():7s}"
     def _format_operands(x:UOp) -> str:
       def _format(src:tuple[UOp, ...]) -> list[str]:
         return [str(s.src[0].val) if s.op is Ops.CAST else reg_strs[o].get(s.dtype.itemsize, o) if \
-                (o:=str(greg(s))) in reg_strs else o for s in src if greg(s) is not None]
+                (o:=str(rdef(s))) in reg_strs else o for s in src if rdef(s) is not None]
       def _mem_adress(base:UOp, idx:UOp, disp:UOp, sz:UOp) -> list[str]:
-        return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.src[0].val}" if greg(idx) else "") + (f" + {d}" if (d:=disp.src[0].val) else "") + "]"]
+        return [f"[{rdef(base)}" + (f" + {rdef(idx)}*{sz.src[0].val}" if rdef(idx) else "") + (f" + {d}" if (d:=disp.src[0].val) else "") + "]"]
 
       if len(x.src) > 4 and x.arg[0] in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:4]) + _format(x.src[4:])
       elif len(x.src) > 3 and x.arg[0] in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:4]) + _format(x.src[4:])

tinygrad/runtime/ops_nv.py

@@ -95,11 +95,6 @@ class QMD:
 
 class NVQueue(HWQueue):
   dev:NVDevice
-  q_rewrite = PatternMatcher([
-    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
-    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
-    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.signal(dst, val)),
-  ])
 
   def nvm(self, subc:int, mthd:int, *vals, typ=2): self.q(*nvm(subc, mthd, *vals, typ=typ))
 
@@ -127,11 +122,6 @@ class NVQueue(HWQueue):
     return doorbell.after(queued).index(0).store(UOp.const(fifo.token, dtypes.uint32))
 
 class NVComputeQueue(NVQueue):
-  q_rewrite = PatternMatcher([
-    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
-    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
-  ]) + NVQueue.q_rewrite
-
   def __init__(self, ctx, submit):
     super().__init__(ctx, submit)
 
@@ -190,11 +180,6 @@ class NVComputeQueue(NVQueue):
     self.prev_qmd = qmd
 
 class NVCopyQueue(NVQueue):
-  q_rewrite = PatternMatcher([
-    (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), lambda ctx, call: ctx.copy(call)),
-    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ()),
-  ]) + NVQueue.q_rewrite
-
   def copy(self, call:UOp):
     dest, src = (a.getaddr(self.devs) for a in call.src[1:3])
     for off in range(0, sz:=call.src[2].max_numel() * call.src[2].dtype.itemsize, step:=(1 << 31)):

tinygrad/runtime/ops_qcom.py

@@ -53,14 +53,6 @@ def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
 
 class QCOMComputeQueue(HWQueue):
   dev:QCOMDevice
-  q_rewrite = PatternMatcher([
-    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
-    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
-    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
-    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
-    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.signal(dst, val)),
-  ])
-
   def cmd(self, opcode:int, *vals): self.q(pkt7_hdr(opcode, sum(x.dtype.itemsize // 4 if isinstance(x, UOp) else 1 for x in vals)), *vals)
 
   def reg(self, reg:int, *vals): self.q(pkt4_hdr(reg, sum(x.dtype.itemsize // 4 if isinstance(x, UOp) else 1 for x in vals)), *vals)

tinygrad/runtime/support/hcq2.py

@@ -21,30 +21,32 @@ if TYPE_CHECKING: from tinygrad.runtime.support.hcq import HCQBuffer # TODO: rem
 HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
 HCQ_RUNTIME_DEV = ContextVar("HCQ_RUNTIME_DEV", "CPU")
 HCQ_CACHE_THRESH = ContextVar("HCQ_CACHE_THRESH", 64)
-HCQ_DEVS = frozenset(("NV", "QCOM", "CPU")) | (frozenset(("AMD",)) if HCQ2 else frozenset())
+HCQ_DEVS = frozenset(("NV", "QCOM")) | (frozenset(("AMD",)) if HCQ2 else frozenset())
 
 @dataclass(frozen=True)
 class HCQInfo:
   device:tuple[str, ...]
 
   kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...], bytes], ...] = () # (devices, name, estimates, timestamp slots, profile key)
+  estimates:Estimates = Estimates()
 
   nargs:int = 0
   table:int = -1
-  inputs:tuple[tuple[UOp, str], ...] = ()
+  inputs:tuple[tuple[UOp, str, int], ...] = ()
   slots:tuple[tuple[str, int], ...] = () # per device, the position of its batch slots in the args
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
 
 def get_enqueue_devs(call:UOp) -> Any|None:
   if call.src[0].op not in (Ops.PROGRAM, Ops.COPY): return None # only these bodies can be enqueued
-  if not (bufs:=get_call_arg_uops(call)) or not all(all_devices_in(b.device, HCQ_DEVS) for b in bufs): return None
+  if not (bufs:=get_call_arg_uops(call)): return None
   if call.src[0].op is Ops.COPY: bufs = bufs[::-1] # copies push from the src device: p2p writes are faster than reads
-  devs = min(bufs, key=lambda b: to_tuple(b.device)[0].startswith("CPU")).device # prio to enqueue on not CPU device
-  # cpu has no queue (yet)
-  if not all_devices_in(devs, HCQ_DEVS) or to_tuple(devs)[0].startswith("CPU"): return None
+  devs = min(bufs, key=lambda b: not all_devices_in(b.device, HCQ_DEVS)).device
+  if not all_devices_in(devs, HCQ_DEVS): return None
+  dev = cast(HCQ2Compiled, Device[to_tuple(devs)[0]])
+  if not all(all_devices_in(b.device, HCQ_DEVS | dev.host_devs) for b in bufs): return None
   # a device without a copy queue leaves copies to its allocator
-  return devs if call.src[0].op is not Ops.COPY or Device[to_tuple(devs)[0]].has_copy_queue else None
+  return devs if call.src[0].op is not Ops.COPY or dev.has_copy_queue else None
 
 def unwrap_view(v:UOp) -> tuple[UOp, int]: # look through views to (base, byte offset)
   if v.op in (Ops.BITCAST, Ops.AFTER): return unwrap_view(v.src[0])
@@ -59,9 +61,10 @@ def to_name(*parts:str) -> str: return "_".join(parts).replace(":", "_").lower()
 def timeline(devs:tuple[str, ...]) -> UOp: return UOp.placeholder((2,), dtypes.uint64, 0, device=devs, volatile=True, tag="timeline")
 def timeline_value(devs:tuple[str, ...]) -> UOp: return timeline(devs).index(1).load()
 
-def rt_addr(b:UOp, dev) -> UOp:
+def rt_addr(b:UOp, dev="CPU") -> UOp:
   base, off = unwrap_view(b)
-  return patch(UOp.placeholder((1,), dtypes.uint64, device=base.device, tag="addr"), [(0, base.getaddr(dev))]).index(0).load() + off
+  word = UOp.placeholder((1,), dtypes.uint64, device="CPU", tag="addr")
+  return patch(word, [(0, base.bitcast(dtypes.uint8)[off:off + b.nbytes()].getaddr(dev))]).index(0).load()
 
 def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
   fn = to_name("submit", (devs:=to_tuple(devs))[0].split(":")[0], queue.split(":")[0])
@@ -82,12 +85,16 @@ def ccall(fn:Any, *args:UOp|int) -> UOp:
   cargs = [UOp.const(a, dtypes.int) if isinstance(a, int) else a for a in args]
   return UOp.custom_function(fn.__name__, ptr.index(0).load()).call(*cargs, ret_dtype=ret)
 
+CDTYPE = {1: dtypes.uchar, 2: dtypes.ushort, 4: dtypes.uint, 8: dtypes.ulong} # a C field as the unsigned int of its size
+
 def cstruct(struct_t, **fields:UOp|int) -> UOp:
-  flds = {n: (o, {1: dtypes.uchar, 2: dtypes.ushort, 4: dtypes.uint, 8: dtypes.ulong}[ctypes.sizeof(t)]) for n, t, o, *_ in struct_t._real_fields_}
+  flds = {n: (o, CDTYPE[ctypes.sizeof(t)]) for n, t, o, *_ in struct_t._real_fields_ if ctypes.sizeof(t)} # skips zero length arrays
   rows = [(flds[n][0], v.cast(flds[n][1]) if isinstance(v, UOp) else UOp.const(v, flds[n][1])) for n, v in fields.items()]
   buf = UOp.placeholder((ctypes.sizeof(struct_t),), dtypes.uint8, device=HCQ_RUNTIME_DEV.value, volatile=True, tag=struct_t.__name__)
   return patch(buf, rows, bytes(ctypes.sizeof(struct_t)))
 
+def cfield(buf:UOp, struct_t, name:str) -> UOp: return buf[(f:=getattr(struct_t, name)).offset:f.offset + f.size].bitcast(CDTYPE[f.size]).index(0)
+
 # *****************
 # 0.1. prep: eager buffers become tagged params
 
@@ -107,16 +114,13 @@ STAGING_SIZE, STAGING_SLOTS = (4 if DEV.interface.startswith("MOCK") else 128) <
 def _staging() -> Buffer: return Buffer("CPU", STAGING_SIZE, dtypes.uint8, preallocate=True)
 
 def _need_staging(a, b):
-  return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not all_devices_in(b.device, HCQ_DEVS) and Device[to_tuple(a.device)[0]].has_copy_queue
-
-def stage_copy_ext(call:UOp) -> UOp|None:
-  if (d:=next((d for b in call.src[1:] for d in to_tuple(b.device) if not d.startswith("CPU")), None)) is None: return None
-  return pm.rewrite(call) if (pm:=getattr(Device[d], "pm_stage_copy", None)) is not None else None
+  if not all_devices_in(a.device, HCQ_DEVS): return False
+  dev = cast(HCQ2Compiled, Device[to_tuple(a.device)[0]])
+  return not all_devices_in(b.device, HCQ_DEVS | dev.host_devs) and dev.has_copy_queue
 
 def stage_copy(dst:UOp, src:UOp) -> UOp|None:
   if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
 
-  assert src.dtype.itemsize == dst.dtype.itemsize, "staged copies must be dtype-size matched"
   base, it, copies = UOp.from_buffer(_staging()), src.dtype.itemsize, []
   chunk = (STAGING_SIZE // STAGING_SLOTS) // it
   for i, off in enumerate(range(0, src.max_numel(), chunk)):
@@ -125,7 +129,6 @@ def stage_copy(dst:UOp, src:UOp) -> UOp|None:
   return UOp(Ops.LINEAR, src=tuple(copies))
 
 pm_insert_copy_staging = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), stage_copy_ext),
   (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy),
 ])
 
@@ -234,10 +237,10 @@ def _finalize_batch(ctx:BatchCtx) -> UOp:
   merged:list[UOp] = [] # the submits in order, after the fence
   for m in _merge_queues(submits): merged.append(m.after(fence, *merged[-1:]))
   estimates = sum((estimate_uop(call) for call, _, _ in ctx.batch), start=Estimates()).simplify()
-  sink = UOp.sink(*merged, arg=KernelInfo("hcq_submit", estimates=estimates), tag=1)
+  sink = UOp.sink(*merged, arg=KernelInfo("hcq_submit"), tag=1)
   for pm in [Device[d].pm_batch for d in ctx.queues if Device[d].pm_batch is not None]: # a device adds its own work to the batch
     if (r:=pm.rewrite(sink)) is not None: sink = r
-  return sink.call(aux=HCQInfo(tuple(ctx.queues), kernels=tuple(kerns)))
+  return sink.call(*(ctx.slots.values() if ctx.profile else ()), aux=HCQInfo(tuple(ctx.queues), kernels=tuple(kerns), estimates=estimates))
 
 @rewrite_group(new_ctx=False)
 def sched_batches(l:UOp, profile:bool) -> UOp:
@@ -254,12 +257,20 @@ def sched_batches(l:UOp, profile:bool) -> UOp:
 @dataclass
 class EncodeCtx:
   devs:tuple[str, ...]
-  inputs:dict[tuple[UOp, str], int] = field(default_factory=dict)
+  inputs:dict[tuple[UOp, str, int], int] = field(default_factory=dict)
   table:UOp = field(default_factory=lambda: UOp.placeholder((1,), dtypes.uint64, device="CPU", tag="inputs"))
-  lt_patches:dict[UOp, list[UOp]] = field(default_factory=dict) # placeholder -> the stores into it that resolve when the linear links
+  lt_patches:list[UOp] = field(default_factory=list)
 
 class HWQueue:
-  q_rewrite:PatternMatcher
+  q_rewrite = PatternMatcher([ # the ops of a queue: a queue defines the methods it supports
+    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
+    (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), lambda ctx, call: ctx.copy(call)),
+    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
+    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
+    (UPat(Ops.INS, arg=("wait_eq", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val, eq=True)),
+    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
+    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.signal(dst, val)),
+  ])
 
   def __init__(self, ctx:EncodeCtx, submit:UOp):
     self.ctx, self.lin = ctx, submit.src[0]
@@ -280,6 +291,7 @@ class HWQueue:
         self.blob += (v & (1 << 8 * n) - 1).to_bytes(n, 'little')
     return len(self.blob)
 
+  def memory_barrier(self): pass # a copy queue has nothing to flush
   def submit(self, cmdbuf:UOp) -> UOp: raise NotImplementedError("queues need a submit")
 
 # *****************
@@ -324,8 +336,8 @@ def _is_input_addr(g:UOp) -> bool:
 def addrs_to_table(ctx:EncodeCtx, g:UOp) -> UOp|None:
   if not _is_input_addr(g): return None
   base, off = unwrap_view(g.src[0])
-  slot = ctx.inputs.setdefault((base, to_tuple(g.arg)[0]), len(ctx.inputs))
-  return ctx.table.index(slot).load() + UOp.const(off, dtypes.uint64)
+  slot = ctx.inputs.setdefault((base, to_tuple(g.arg)[0], off), len(ctx.inputs))
+  return ctx.table.index(slot).load()
 
 def _is_link_patch(w:UOp) -> bool:
   if w.op is Ops.GETADDR: return not _is_input_addr(w)
@@ -337,10 +349,7 @@ def _is_link_patch(w:UOp) -> bool:
 def hoist_links(ctx:EncodeCtx, a:UOp) -> UOp|None:
   links, rest = partition(a.src[1:], lambda s: s.op is Ops.STORE and _is_link_patch(s))
   if not links: return None
-  # nest the addr placeholders patches under their getaddr
-  ws = UOp.sink(*links)
-  sub = {g: g.replace(src=(g.src[0].after(*ctx.lt_patches[g.src[0]]),)) for g in ws.toposort() if g.op is Ops.GETADDR and g.src[0] in ctx.lt_patches}
-  ctx.lt_patches.setdefault(unwrap_view(a.src[0])[0], []).extend(ws.substitute(sub).src)
+  ctx.lt_patches.extend(links)
   return a.src[0].after(*rest)
 
 pm_patches = PatternMatcher([(UPat(Ops.GETADDR, name="g"), addrs_to_table), (UPat(Ops.AFTER, name="a"), hoist_links)])
@@ -374,6 +383,11 @@ def encode_submit(hq:HWQueue) -> UOp:
 # *****************
 # 4. lower call
 
+pm_renumber = PatternMatcher([
+  (UPat(Ops.RANGE, name="u"), lambda ctx, u: u.replace(arg=(next(ctx),)+u.arg[1:])),
+  (UPat(Ops.BUFFER, name="u"), lambda ctx, u: u.replace(arg=replace(u.arg, slot=next(ctx))) if u.addrspace is AddrSpace.REG else None),
+])
+
 def lower_call(call:UOp) -> UOp|None:
   if not isinstance(call.arg.aux, HCQInfo) or call.arg.aux.nargs: return None # not an hcq call, or lowered already
 
@@ -387,28 +401,23 @@ def lower_call(call:UOp) -> UOp|None:
   body = body.substitute({ctx.table: (table:=UOp.placeholder((len(ctx.inputs),), dtypes.uint64, device="CPU", tag="inputs"))})
 
   # the placeholders become the body's params in visit order, variables bind by name after them, the ranges renumber
-  tops = body.toposort()
-  bufs, alus = partition([u for u in tops if u.op is Ops.PARAM], lambda u: u.tag is not None)
+  bufs, alus = partition([u for u in body.toposort() if u.op is Ops.PARAM], lambda u: u.tag is not None)
+  bufs = dedup([*call.src[1:], *bufs])
   names = dedup([a.arg.name for a in alus])
   # bufs to params
   params = {b: UOp.param(i, b.dtype, b.shape, HCQ_RUNTIME_DEV.value, volatile=b.arg.volatile, name=f"{b.arg.name}_{i}") for i, b in enumerate(bufs)}
   # new slots for vars
   vals = {a: a.replace(arg=replace(a.arg, slot=len(bufs) + names.index(a.arg.name))) for a in alus}
-  # reenum ranges
-  rngs = {r: r.replace(arg=(i,)+r.arg[1:]) for i, r in enumerate(sorted([u for u in tops if u.op is Ops.RANGE], key=lambda r: r.arg))}
-  # and sub all of them
-  sink = body.substitute(params | vals | rngs, enter_calls=True)
+  sink = graph_rewrite(body.substitute(params | vals, enter_calls=True), pm_renumber, ctx=itertools.count(), walk=True, enter_calls=True)
 
-  # move all lt-patches to the args
-  patched = {b: b.after(*dedup(stores)) for b, stores in ctx.lt_patches.items()}
-  args = [patched.get(b, b) for b in bufs]
+  patches = dedup(ctx.lt_patches)
 
-  if VIZ: graph_rewrite(UOp.sink(*args), PatternMatcher([]), name="View Link-Time Patches")
+  if VIZ: graph_rewrite(UOp.sink(*patches), PatternMatcher([]), name="View Link-Time Patches")
   if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Body")
 
-  info = replace(call.arg.aux, nargs=len(args), table=bufs.index(table) if table in bufs else -1, inputs=tuple(ctx.inputs),
+  info = replace(call.arg.aux, nargs=len(bufs), table=bufs.index(table) if table in bufs else -1, inputs=tuple(ctx.inputs),
                  slots=tuple((to_tuple(b.device)[0], i) for i, b in enumerate(bufs) if b.tag == "slots"))
-  return call.replace(src=(sink, *args), arg=replace(call.arg, aux=info))
+  return call.replace(src=(sink, *bufs), arg=replace(call.arg, aux=info)).after(*patches)
 pm_encode = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.SINK),), name="call", allow_any_len=True), lower_call)])
 
 hcq_compile_cache:dict[tuple[UOp, bool], UOp] = {} # eager templates: a buffer-free linear (uops are hash-consed) to its compiled form
@@ -478,6 +487,8 @@ pm_link = PatternMatcher([
   (UPat(name="buf").store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast())), fold_binary),
   (UPat(name="buf").index(UPat(Ops.STACK, src=UPat.cvar().or_casted(), name="offs")).store(UPat(Ops.STACK, src=UPat.cvar().or_casted(), name="ws")),
     fold_words),
+  (UPat(Ops.AFTER, src=(UPat(Ops.CALL),), allow_any_len=True, name="a"),
+    lambda a: a.src[0].after(*(s for s in a.src[1:] if s.op is not Ops.NOOP))),
   (UPat(Ops.AFTER, name="a"), lambda a: None if a.is_bound_var or a.src[0].op is Ops.CALL else
    a.src[0] if all(s.op is Ops.NOOP for s in a.src[1:]) else panic(RuntimeError, f"unresolved link words on {a.src[0].op}")),
 ])
@@ -505,6 +516,7 @@ class HCQ2Compiled(Compiled):
   wait_timeout_ms: float = 30000.0
   sleep_timeout_ms: int|None = None
   rt_nbytes: int = 64 << 20 # the pool every per-linear buffer is carved out of
+  host_devs: frozenset[str] = frozenset({"CPU"})
   pm_encode: PatternMatcher = PatternMatcher([]) # the backend's own encode rules, matched by its submit names
   var_vals: dict[str, int] = {}
 
@@ -565,11 +577,14 @@ class HCQ2Compiled(Compiled):
     st, done = time.perf_counter(), sig[0]
     while done < value:
       if done != (done:=sig[0]): st = time.perf_counter()
-      elif (elapsed:=time.perf_counter() - st) > (timeout or self.wait_timeout_ms) / 1000: self.on_device_hang()
+      elif (elapsed:=time.perf_counter() - st) > (timeout or self.wait_timeout_ms) / 1000: raise RuntimeError(f"{self.device} signal wait timed out")
       elif self.sleep_timeout_ms is not None and elapsed > self.sleep_timeout_ms / 1000: self.on_sleep()
 
   def synchronize(self, timeout:int|None=None):
-    self._wait_signal(tl:=self.timeline._buf.cpu_view().view(fmt='Q'), tl[1], timeout)
+    try: self._wait_signal(tl:=self.timeline._buf.cpu_view().view(fmt='Q'), tl[1], timeout)
+    except RuntimeError:
+      self.on_device_hang()
+      raise
     if self.prof_ents: self.collect_prof()
 
   def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected")

tinygrad/runtime/support/usb.py

@@ -1,11 +1,12 @@
+from typing import cast
 import ctypes, struct, time, functools, itertools
-from typing import Any, cast
-from tinygrad.runtime.autogen import libusb
-from tinygrad.helpers import DEBUG, DEV, to_mv, from_mv, round_up, ceildiv, unwrap, to_tuple
-from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher
-from tinygrad.device import Buffer, BufferSpec, Device
-from tinygrad.runtime.support.hcq2 import HCQInfo, make_submit, HCQ_RUNTIME_DEV
+from tinygrad.runtime.autogen import libusb, libc
+from tinygrad.helpers import DEBUG, DEV, to_mv, from_mv, round_up, ceildiv, to_tuple
+from tinygrad.dtype import dtypes, DType, AddrSpace
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite
+from tinygrad.engine.realize import pm_flatten_linear
+from tinygrad.device import Buffer, BufferSpec
+from tinygrad.runtime.support.hcq2 import HCQ_RUNTIME_DEV, HCQ_DEVS, ccall, cfield, patch, rt_addr, unwrap_view, all_devices_in
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support import c
 
@@ -227,7 +228,6 @@ class USBMMIOInterface(MMIOInterface):
     return (index * self.el_sz, self.el_sz)
 
   def __getitem__(self, index):
-    Device[HCQ_RUNTIME_DEV.value].synchronize() # one driver on the link: drain the compiled submits before python touches it
     off, sz = self._off_from_index(index)
     if self.pcimem:
       assert sz % 4 == 0 and off % 4 == 0, f"pcie_mem_read requires 4-byte aligned access, got off={off}, sz={sz}"
@@ -236,7 +236,6 @@ class USBMMIOInterface(MMIOInterface):
     return data if isinstance(index, slice) else int.from_bytes(data, "little")
 
   def __setitem__(self, index, data):
-    Device[HCQ_RUNTIME_DEV.value].synchronize()
     off, _ = self._off_from_index(index)
     data = struct.pack(self.fmt, data) if isinstance(data, int) else bytes(data)
     if not self.pcimem: self.usb.scsi_write(data) if self.addr == 0xf000 else self.usb.write(self.addr + off, data)
@@ -249,75 +248,249 @@ class USBMMIOInterface(MMIOInterface):
     return USBMMIOInterface(self.usb, self.addr+offset, self.nbytes-offset if size is None else size, fmt=fmt or self.fmt, pcimem=self.pcimem)
 
 # *****************
+# UOps implementation
 
-# TODO: unported to the hcq2 rewrite, keeps the old signal placeholder helper alive
-def make_buf(devs, slot:int=0, tag:str="signal") -> UOp: return UOp.placeholder((1,), dtypes.uint64, slot, device=devs, volatile=True, tag=tag)
+# the sram: two halves, the last block of each holds the sentinel. a write starts on a 16KB slot and ends at the half's end
+HALF, CHUNK, SLOT = 0x40000, 0x40000 - 512, 0x4000
 
-def _libusb(devs, dep:tuple[UOp, ...], fn:str, *args) -> UOp:
-  # the CUSTOM_FUNCTION body holds the callee (the loaded function pointer), the call args are plain dataflow
-  fptr = make_buf(devs, tag=f"func:{fn}").after(*dep).index(0).load()
-  return UOp.custom_function(fn, fptr).call(make_buf(devs, tag="usb_handle").index(0).load(),
-    *[UOp.const(a, dtypes.int) if isinstance(a, int) else a for a in args], ret_dtype=dtypes.void)
+# the host block
+def usb_host(dev) -> UOp: return UOp.placeholder((0x180020,), dtypes.uint8, 0, device=to_tuple(dev)[0], tag="usb_host")
+def usb_link(dev) -> UOp: return usb_host(dev)[:24].bitcast(dtypes.uint64) # [the libusb handle, its context, the chunks of the last batch]
+def usb_stage(dev) -> UOp: return usb_host(dev)[32:32 + 2 * HALF] # what streams into the sram, laid out like it. then 1MB of zeros
 
-def usb_bulk(devs, dep, endpoint:int, data:UOp, length, timeout:int=1000) -> UOp: # NULL actual_length out param
-  return _libusb(devs, dep, "libusb_bulk_transfer", endpoint, data, length, UOp.const(0, dtypes.uint64), timeout)
+# libusb's transfer per half: a bulk out on the handle, the program sets the rest and reaps it
+def usb_xfer(dev, half:int) -> UOp:
+  return UOp.placeholder((ctypes.sizeof(libusb.struct_libusb_transfer),), dtypes.uint8, 0, device=to_tuple(dev)[0], tag=f"usb_xfer{half}")
 
-def usb_stream(devs, dep:tuple[UOp, ...], addr:UOp, data:UOp, nbytes:int, write:bool) -> UOp:
-  hdr = UOp.placeholder((2,), dtypes.uint64, device=devs, tag="usb_scratch").after(*dep)
-  arm = _libusb(devs, (hdr.index(0).store(addr), hdr.index(1).store(UOp.const(nbytes // 4, dtypes.uint64))), "libusb_control_transfer",
-                0x40, 0xF0, (0x60 if write else 0x20) | (0x0F << 8), 1 if write else 2, hdr.index(0), 12, 5000)
-  return usb_bulk(devs, (arm,), 0x02 if write else 0x81, data, nbytes)
+# the words in vram the host pokes
+def usb_vram(dev) -> UOp: return UOp.placeholder((2,), dtypes.uint32, 0, device=to_tuple(dev)[0], tag="usb_vram")
+def usb_go(dev) -> UOp: return usb_vram(dev)[:1] # the go of a read
+def usb_scratch(dev) -> UOp: return usb_vram(dev)[1:] # an empty stream's target
 
-def usb_load(b:UOp, idx:UOp, dt) -> UOp:
-  got = UOp.placeholder((1,), dt, device=(devs:=to_tuple(b.device)), tag="usb_scratch")
-  addr = b.getaddr((HCQ_RUNTIME_DEV.value,)) + (idx*dt.itemsize).cast(dtypes.uint64)
-  return got.after(usb_stream(devs, b.src[1:] if b.op is Ops.AFTER else (), addr, got.index(0), dt.itemsize, False)).index(0).load()
+# the asm2464's memory: its sys page, its cq page, the sram it streams
+def usb_asm24(dev) -> UOp: return UOp.placeholder((0x85000,), dtypes.uint8, 0, device=to_tuple(dev)[0], tag="usb_asm24")
+def usb_fence(dev) -> UOp: return usb_asm24(dev)[0x800:0x804].bitcast(dtypes.uint32) # the chunks the queue is done with, 0xa800 to the host
+def usb_cq(dev) -> UOp: return usb_asm24(dev)[0x100c:0x1010].bitcast(dtypes.uint32) # the completion dword: a store releases an armed read
+def usb_sram(dev) -> UOp: return usb_asm24(dev)[0x5000:0x5000 + 2 * HALF]
 
-def usb_write(b:UOp, idx:UOp, v:UOp) -> UOp:
-  val = (s:=UOp.placeholder((1,), v.dtype, device=(devs:=to_tuple(b.device)), tag="usb_scratch")).after(s.index(0).store(v))
-  addr = b.getaddr((HCQ_RUNTIME_DEV.value,)) + (idx*v.dtype.itemsize).cast(dtypes.uint64)
-  return usb_stream(devs, b.src[1:] if b.op is Ops.AFTER else (), addr, val.index(0), v.dtype.itemsize, True)
+# transfers on the link h: the caller puts the link after the transfer, so the link threads through the program in order
+def usb_reg(dt:DType, *vals:UOp|int) -> UOp: # an array on the program's stack: what a transfer reads or writes
+  r = UOp.placeholder((max(1, len(vals)),), dt, addrspace=AddrSpace.REG)
+  return r.after(*[r.index(i).store(v.cast(dt) if isinstance(v, UOp) else UOp.const(v, dt)) for i, v in enumerate(vals)])
 
-def usb_idle(devs) -> UOp:
-  v = usb_load(make_buf(devs, tag="timeline_signal").after(loop:=UOp.loop(0)), UOp.const(0, dtypes.int), dtypes.uint64)
-  return v.end(loop, v + 1 < make_buf(devs, tag="timeline_value").index(0).load())
+def usb_ctrl(h:UOp, rtype:int, req:int, val:UOp|int, idx:UOp|int, data:UOp, n:UOp|int, timeout:int=1000) -> UOp:
+  return ccall(libusb.libusb_control_transfer, h.index(0).load(), rtype, req, val, idx, data, n, timeout)
 
-def usb_scsi(devs, read:bool, nbytes:int) -> UOp:
-  return _libusb(devs, (usb_idle(devs),), "libusb_control_transfer", 0x40, 0xF2, ceildiv(nbytes, 512) | (0x8000 if read else 0),
-                 (ceildiv(nbytes, 0x4000) & 0xFF) << 8, UOp.const(0, dtypes.uint64), 0, 1000)
+def usb_bulk(h:UOp, ep:int, data:UOp, n:UOp|int, timeout:int=10000) -> UOp: # NULL actual_length
+  return ccall(libusb.libusb_bulk_transfer, h.index(0).load(), ep, data, n, UOp.const(0, dtypes.uint64), timeout)
 
-def usb_stage_copy(dst:UOp, src:UOp) -> UOp|None:
-  if (cin:=to_tuple(src.device)[0].startswith("CPU")) == to_tuple(dst.device)[0].startswith("CPU"): return None
+def usb_poke(h:UOp, addr:UOp, val:UOp) -> UOp: # 0xF0 mode 0: a dword in one control transfer. the header: the address, then a dword
+  return usb_ctrl(h, 0x40, 0xF0, 0x60 | 0x0F00, 0, usb_reg(dtypes.uint64, addr, val.bitcast(dtypes.uint32).cast(dtypes.uint64)).index(0), 12, 5000)
 
-  total, ops, win = dst.nbytes(), [], cast(Any, Device[(devs:=to_tuple((dst if cin else src).device))[0]]).iface.usb_sram
-  for off in range(0, total, win.size): # off and nb are bytes, the two ends of the copy can have different dtypes
-    sram = UOp.from_buffer(win)[0:(nb:=min(win.size, total - off))]
-    s, d = src[off // src.dtype.itemsize:(off + nb) // src.dtype.itemsize], dst[off // dst.dtype.itemsize:(off + nb) // dst.dtype.itemsize]
-    if cin:
-      push = usb_bulk(devs, (usb_scsi(devs, False, nb),), 0x02, s.getaddr((HCQ_RUNTIME_DEV.value,)), round_up(nb, 512), 10000)
-      ops += [UOp.custom_function("hcq", push.sink()).call(sram, s, name="hcq_copyin", aux=HCQInfo(devs)),
-              sram.copy_to_device(d.device).call(d, sram)]
+def usb_stream(h:UOp, addr:UOp, data:UOp, n:UOp|int, write:bool) -> UOp: # 0xF0 mode 1/2: the header, then the payload on the bulk endpoint
+  hdr = usb_ctrl(h, 0x40, 0xF0, (0x60 if write else 0x20) | 0x0F00, 1 if write else 2, usb_reg(dtypes.uint64, addr, n // 4).index(0), 12, 5000)
+  return usb_bulk(h.after(hdr), 0x02 if write else 0x81, data, n)
+
+# a copy between the host and vram goes through the sram in chunks, a batch numbers them from 0 in queue order
+def is_host(b:UOp) -> bool: return b.device is None or not all_devices_in(b.device, HCQ_DEVS - {"CPU"}) # the stack, or memory the host reads
+def usb_wire(size:UOp|int) -> UOp|int: return (size + 512 + SLOT - 1) // SLOT * SLOT # a chunk on the wire: the payload, then the sentinel block
+def usb_sentinel(g:UOp) -> UOp: return ((g & 0xFFFFFF) | 0x51000000).cast(dtypes.uint32)
+def is_staged(call:UOp) -> bool: return call.op is Ops.CALL and call.src[0].op is Ops.COPY and is_host(call.src[1]) != is_host(call.src[2])
+def usb_chunks(call:UOp) -> list[tuple[UOp, int, int]]: # (the host view, byte offset, bytes) per chunk of a copy. a read fills both payload areas
+  host, win = (call.src[2], CHUNK) if is_host(call.src[2]) else (call.src[1], 2 * CHUNK)
+  return [(host, off, min(win, host.nbytes() - off)) for off in range(0, host.nbytes(), win)]
+
+def usb_copy_slicer(ctx:dict[UOp, tuple[int, int]], call:UOp, dst:UOp, src:UOp) -> UOp|None:
+  if (nums:=ctx.get(call)) is None: return None
+
+  vram = (dst if is_host(src) else src).bitcast(dtypes.uint8)
+  sram, ops = usb_sram(vram.device), []
+
+  # reassembly copies with comms
+  for n, (_, off, nb) in enumerate(usb_chunks(call), start=nums[0]):
+    if is_host(src):
+      end = ((n - nums[1]) & 1) * HALF + HALF
+      ops += [UOp(Ops.INS, arg=("wait_eq", dtypes.void), src=(sram[end - 4:end].bitcast(dtypes.uint32), usb_sentinel(UOp.const(n, dtypes.uint32)))),
+              sram.copy_to_device(vram.device).call(vram[off:off + nb], sram[end - usb_wire(nb):end - usb_wire(nb) + nb]),
+              UOp(Ops.INS, arg=("store", dtypes.void), src=(sram[end - 4:end].bitcast(dtypes.uint32), UOp.const(0, dtypes.uint32))),
+              UOp(Ops.INS, arg=("store", dtypes.void), src=(usb_fence(vram.device), UOp.const(n + 1, dtypes.uint32)))]
     else:
-      pad = UOp.new_buffer("CPU", round_up(nb, 512), dtypes.uint8)[0:nb]
-      submit = make_submit(s.copy_to_device(sram.device).call(sram, s), devs=devs, queue="COPY:0")
-      pull = usb_bulk(devs, (submit,), 0x81, pad.getaddr((HCQ_RUNTIME_DEV.value,)), round_up(nb, 512), 10000)
-      ops += [UOp.custom_function("hcq", pull.sink()).call(pad, sram, s, name="hcq_copyout", aux=HCQInfo(devs)),
-              pad.copy_to_device("CPU").call(d, pad)]
+      ops += [UOp(Ops.INS, arg=("wait", dtypes.void), src=(usb_go(vram.device), UOp.const(n + 1, dtypes.uint32))),
+              UOp(Ops.INS, arg=("store", dtypes.void), src=(usb_go(vram.device), UOp.const(0, dtypes.uint32)))]
+      ops += [vram.copy_to_device(vram.device).call(sram[wo:wo + pb], vram[off + po:off + po + pb])
+              for wo, po, pb in ((0, 0, min(nb, CHUNK)), (HALF, CHUNK, nb - CHUNK)) if pb > 0]
+      ops += [UOp(Ops.INS, arg=("store", dtypes.void), src=(usb_cq(vram.device), UOp.const(0, dtypes.uint32))),
+              UOp(Ops.INS, arg=("store", dtypes.void), src=(usb_fence(vram.device), UOp.const(n + 1, dtypes.uint32)))]
   return UOp(Ops.LINEAR, src=tuple(ops))
-pm_usb_stage = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), usb_stage_copy)])
+pm_usb_copy_slicer = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call"), usb_copy_slicer)]) + pm_flatten_linear
 
-USB_HOST_TAGS = {"signal", "timeline_signal"}
-pm_usb_hostio = PatternMatcher([
-  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(Ops.PARAM, tag=USB_HOST_TAGS).or_after(name="b"), UPat(name="idx"))),),
-        name="ld"), lambda b, idx, ld: usb_load(b, idx, ld.dtype)),
-  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat(Ops.PARAM, tag=USB_HOST_TAGS).or_after(name="b"), UPat(name="idx"))), UPat(name="v"))), usb_write)])
+def usb_copy_rewriter(s:UOp) -> UOp|None:
+  lins = [submit.without_after.src[0] for submit in s.src]
+  if not (copies:=[call for lin in lins for call in lin.src if is_staged(call)]): return None
 
-pm_usb_bufferize = PatternMatcher([
-  (UPat(Ops.PARAM, tag={"systems", "runtime", "inputs", "usb_scratch"}, name="b"),
-   lambda ctx, b: Buffer("CPU", b.max_numel(), b.dtype, options=BufferSpec(nolru=True), preallocate=True)),
-  (UPat(Ops.PARAM, tag="usb_handle", name="b"), lambda ctx, b: ctx.signal(b.tag, ctx.iface.usb_handle, device="CPU")),
-  (UPat(Ops.PARAM, name="b"), lambda ctx, b: None if not isinstance(b.tag, str) or not b.tag.startswith("func:") else
-   ctx.signal(b.tag, unwrap(ctypes.cast(getattr(libusb.dll, b.tag[5:]), ctypes.c_void_p).value), device="CPU")),
+  # number the chunks in queue order
+  runs, nums, n = [], {}, 0
+  for cin, grp in itertools.groupby(copies, key=lambda call: is_host(call.src[2])):
+    chunks:list = []
+    for call in grp: nums[call], chunks = (n + len(chunks), n), chunks + usb_chunks(call)
+    runs.append((cin, n, chunks))
+    n += len(chunks)
+
+  # queue side
+  s = graph_rewrite(s, pm_usb_copy_slicer, ctx=nums, name="usb copy slicer")
+
+  # host side: drain the last batch, restart the fence
+  h = usb_link(lins[0].arg[0][0]).after(s.src[-1])
+  h = h.after(usb_ctrl(h.after(usb_drained(h, h.index(2).load() + 1)), 0x40, 0xE5, rt_addr(usb_fence(h.device)), 0, UOp.const(0, dtypes.uint64), 0))
+  for cin, run, chunks in runs: h = (usb_copyin if cin else usb_copyout)(h, chunks, run)
+  return s.replace(src=(*s.src, h.index(2).store(UOp.const(n, dtypes.uint64))))
+pm_usb_batch = PatternMatcher([(UPat(Ops.SINK, name="s"), usb_copy_rewriter)])
+
+# the host side: a run of copies one way is a loop over a table of its chunks
+def view_bytes(b:UOp) -> UOp: # the bytes of a view as a view of its base alone: an address doesn't order, so nothing it's after comes along
+  base, off = unwrap_view(b)
+  return base.bitcast(dtypes.uint8)[off:off + b.nbytes()]
+
+def usb_table(chunks:list[tuple[UOp, int, int]], dev) -> UOp: # [host address, bytes] per chunk, the addresses as the device's host program sees them
+  table = UOp.placeholder((2 * len(chunks),), dtypes.uint64, device=HCQ_RUNTIME_DEV.value, tag="usb_table")
+  return patch(table, [(16 * i, view_bytes(host)[off:off + nb].getaddr(to_tuple(dev)[0])) for i, (host, off, nb) in enumerate(chunks)] +
+                      [(16 * i + 8, UOp.const(nb, dtypes.uint64)) for i, (_, _, nb) in enumerate(chunks)])
+
+def usb_reap(h:UOp, xfer:UOp) -> UOp: # until the async transfer is done
+  loop = UOp.range(UOp(Ops.NOOP), next(UOp.unique_num), dtype=dtypes.void, src=(h,))
+  events = ccall(libusb.libusb_handle_events_timeout, h.after(loop).index(1).load(), usb_reg(dtypes.uint64, 0, 0).index(0)) # a poll
+  status = cfield(xfer.after(events), libusb.struct_libusb_transfer, "status").load()
+  return status.end(loop, status.eq(0xff))
+
+def usb_drained(h:UOp, need:UOp) -> UOp: # until the queue is done with the chunks before need, it may lag one
+  loop, slot = UOp.range(UOp(Ops.NOOP), next(UOp.unique_num), dtype=dtypes.void, src=(h,)), usb_reg(dtypes.uint32)
+  fence = slot.after(usb_ctrl(h.after(loop), 0xC0, 0xE4, rt_addr(usb_fence(h.device)), 0, slot.index(0), 1)).index(0).load() # a byte, mod 256
+  return fence.end(loop, ((need - fence.cast(dtypes.uint64)) & 0xff) > 1)
+
+def usb_chunk(h:UOp, table:UOp, i:UOp, half:int, run:int) -> UOp: # chunk i of the run into a half of the sram, an async bulk
+  addr, size = table.index(2 * i).load(), table.index(2 * i + 1).load().cast(dtypes.int)
+  n, wire, end = (i + run).cast(dtypes.uint64), cast(UOp, usb_wire(size)), (half + 1) * HALF
+  xfer, stage = usb_xfer(h.device, half), usb_stage(h.device)
+
+  # stage it, the sentinel last
+  h = h.after(usb_reap(h, xfer))
+  h = h.after(ccall(libc.memcpy, stage.after(h).index(end - wire), addr, size.cast(dtypes.uint64)))
+  h = h.after(stage.after(h).bitcast(dtypes.uint32).index(end // 4 - 1).store(usb_sentinel(n)))
+
+  # once the queue is done with the half: arm the write from its slot, send it
+  h = h.after(usb_drained(h, n))
+  h = h.after(usb_ctrl(h, 0x40, 0xF2, wire // 512, ((end - wire) // SLOT) | (wire // SLOT << 8), UOp.const(0, dtypes.uint64), 0))
+  field = functools.partial(cfield, xfer:=xfer.after(h), libusb.struct_libusb_transfer)
+  xfer = xfer.after(field("status").store(0xff), field("length").store(wire),
+                    field("buffer").store(rt_addr(stage) + (end - wire).cast(dtypes.uint64)))
+  return ccall(libusb.libusb_submit_transfer, xfer.index(0))
+
+def usb_copyin(h:UOp, chunks:list, run:int) -> UOp: # pairs of chunks on the two halves, two transfers in flight
+  table, n = usb_table(chunks, h.device), len(chunks)
+  h = h.after(usb_drained(h, UOp.const(run + 1, dtypes.uint64))) # the sram is free, the run starts on half 0
+
+  # the pairs, a single pair unrolled: the linearizer misplaces a one trip loop
+  if (pairs:=n // 2 if n // 2 > 1 else 0):
+    j = UOp.range(pairs, next(UOp.unique_num), dtype=dtypes.int)
+    hj = h.after(j, usb_chunk(h.after(j), table, j * 2, 0, run))
+    h = h.after(usb_chunk(hj, table, j * 2 + 1, 1, run).end(j))
+  for i in range(pairs * 2, n): h = h.after(usb_chunk(h, table, UOp.const(i, dtypes.int), i & 1, run))
+  return h
+
+def usb_copyout(h:UOp, chunks:list, run:int) -> UOp: # per chunk: arm a read of the sram, release the queue to fill it, pull it
+  table, stage = usb_table(chunks, h.device), usb_stage(h.device)
+  h = h.after(usb_drained(h, UOp.const(run + 1, dtypes.uint64))) # the sram is free, nothing else writes the asm2464's memory
+
+  i = UOp.range(len(chunks), next(UOp.unique_num), dtype=dtypes.int)
+  addr, size = table.index(2 * i).load(), table.index(2 * i + 1).load().cast(dtypes.int)
+  first, second = size.minimum(CHUNK), (size - CHUNK).maximum(0) # the payload areas of the halves
+  wire = (size + (second > 0).where(UOp.const(512, dtypes.int), UOp.const(0, dtypes.int)) + 511) // 512 * 512
+
+  # arm the read, release the queue, pull it
+  hi = h.after(i, usb_ctrl(h.after(i), 0x40, 0xF2, (wire // 512) | 0x8000, (wire + 0x3fff) // 0x4000 << 8, UOp.const(0, dtypes.uint64), 0))
+  hi = hi.after(usb_poke(hi, rt_addr(usb_go(h.device)), (i + run + 1).cast(dtypes.uint32)))
+  hi = hi.after(usb_bulk(hi, 0x81, stage.index(0), wire))
+  hi = hi.after(ccall(libc.memcpy, addr, stage.after(hi).index(0), first.cast(dtypes.uint64)))
+  hi = hi.after(ccall(libc.memcpy, addr + CHUNK, stage.after(hi).index(HALF), second.cast(dtypes.uint64)))
+  return h.after(hi.end(i))
+
+# the host's accesses to device memory: a load streams the value into a register, a store pokes it, a loop of stores streams the source
+def is_remote(b:UOp) -> bool:
+  return (p:=unwrap_view(b)[0]).op is Ops.PARAM and not is_host(p) and not str(p.tag).startswith(("usb_host", "usb_xfer", "put_value", "cmdbuf_copy"))
+def usb_addr(b:UOp, idx:UOp, dt:DType) -> UOp: return rt_addr(b) + (idx * dt.itemsize).cast(dtypes.uint64) # of an element of a view
+def usb_deps(b:UOp) -> tuple[UOp, ...]: # what a view is after
+  return (b.src[1:] if b.op is Ops.AFTER else ()) + (usb_deps(b.src[0]) if b.op in (Ops.BITCAST, Ops.SHRINK, Ops.AFTER) else ())
+def usb_affine(idx:UOp, r:UOp) -> UOp|None: # the base of idx = base + r, None if idx doesn't walk r with unit stride
+  if idx is r: return UOp.const(0, r.dtype)
+  if idx.op is not Ops.ADD or r not in idx.src: return None
+  base = idx.src[1] if idx.src[0] is r else idx.src[0]
+  return base if r not in base.ranges else None
+
+def usb_load(b:UOp, idx:UOp, ld:UOp) -> UOp:
+  slot = usb_reg(ld.dtype)
+  read = usb_stream(usb_link(b.device).after(*usb_deps(b)), usb_addr(b, idx, ld.dtype), slot.index(0), ld.dtype.itemsize, False)
+  return slot.after(read).index(0).load()
+
+def usb_store(b:UOp, idx:UOp, v:UOp) -> UOp:
+  # a patch: word by word, each after the one before
+  if idx.op is Ops.STACK:
+    h = usb_store(b, idx.src[0], v.src[0])
+    for i, w in zip(idx.src[1:], v.src[1:]): h = usb_store(b.after(h), i, w)
+    return h
+
+  # a dword is one poke, a qword two
+  h, addr = usb_link(b.device).after(*usb_deps(b)), usb_addr(b, idx, v.dtype)
+  if v.dtype.itemsize == 4: return usb_poke(h, addr, v)
+  return usb_poke(h.after(usb_poke(h, addr, v.cast(dtypes.uint32))), addr + 4, (v >> 32).cast(dtypes.uint32))
+
+def usb_copy(dst:UOp, di:UOp, v:UOp, r:UOp) -> UOp|None: # a loop of unit stride stores from host memory: one stream
+  if not is_remote(dst): return None
+
+  # the source: host memory the loop walks, or the zeros
+  if v.op is Ops.LOAD and not is_remote(sb:=v.src[0].src[0]): s0, deps = usb_affine(v.src[0].src[1], r), usb_deps(sb)
+  elif v.vmin == v.vmax == 0: sb, s0, deps = usb_host(dst.device)[32 + 2 * HALF:], UOp.const(0, dtypes.int), ()
+  else: return usb_store(dst, di, v).end(r)
+  if s0 is None or (d0:=usb_affine(di, r)) is None: return usb_store(dst, di, v).end(r)
+
+  # an empty loop streams one element into the scratch word: the firmware can't stream zero bytes
+  h, cnt = usb_link(dst.device).after(*usb_deps(dst), *deps, *r.src[1:]), r.src[0]
+  addr = (cnt > 0).where(usb_addr(dst, d0, v.dtype), rt_addr(usb_scratch(dst.device)))
+  return usb_stream(h, addr, sb.index(s0.minimum(sb.max_numel() - 1)), (cnt * v.dtype.itemsize).maximum(v.dtype.itemsize), True)
+
+pm_usb_lower = PatternMatcher([ # a store inside a loop is left to the loop's END
+  (UPat.var("dst").index(UPat.var("di")).store(UPat.var("v")).end(UPat(Ops.RANGE, name="r")), usb_copy),
+  (UPat.var("b").index(UPat.var("idx")).store(UPat.var("v")), lambda b, idx, v: None if idx.ranges or not is_remote(b) else usb_store(b, idx, v)),
+  (UPat.var("b").index(UPat.var("idx")).load(name="ld"), lambda b, idx, ld: usb_load(b, idx, ld) if is_remote(b) else None),
 ])
 
+# the buffers of the device's usb state, bound once for the life of the device
+# the libusb handles are this process's: written here, not compiled in, so a pickled jit links in another process
+@functools.cache
+def _host_block(dev) -> Buffer: # the link words, then zeros: mmap memory
+  b = Buffer("CPU", 0x180020, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
+  b._buf.cpu_view().view(fmt='B')[:16] = struct.pack('QQ', *[ctypes.addressof(x.contents) for x in (dev.iface.pci_dev.usb.usb.handle, USB3.ctx())])
+  return b
+@functools.cache
+def _xfer(dev, tag:str) -> Buffer: # libusb's transfer, a bulk out on the handle: the program sets the rest
+  t = libusb.libusb_alloc_transfer(0).contents
+  t.dev_handle, t.endpoint, t.type, t.timeout = dev.iface.pci_dev.usb.usb.handle, 0x02, libusb.LIBUSB_TRANSFER_TYPE_BULK, 10000
+  return Buffer("CPU", ctypes.sizeof(t), dtypes.uint8, options=BufferSpec(external_ptr=ctypes.addressof(t), nolru=True), preallocate=True)
+@functools.cache
+def _words(dev) -> Buffer: # the words in vram the host pokes, cleared: no stale go
+  b = Buffer(dev.device, 2, dtypes.uint32, options=BufferSpec(uncached=True, cpu_access=True, nolru=True), preallocate=True)
+  b._buf.cpu_view().view(fmt='B')[:8] = bytes(8)
+  return b
+@functools.cache
+def _asm24(dev) -> Buffer: # the asm2464's memory: not ours to free
+  ctrl = dev.iface.ctrl
+  return Buffer(dev.device, 0x85000, dtypes.uint8, options=BufferSpec(external_ptr=ctrl.va_addr, nolru=True)).allocate(opaque=ctrl)
+pm_usb_bufferize = PatternMatcher([
+  (UPat(Ops.PARAM, tag="usb_host"), lambda ctx: _host_block(ctx)),
+  (UPat(Ops.PARAM, tag={"usb_xfer0", "usb_xfer1"}, name="b"), lambda ctx, b: _xfer(ctx, b.tag)),
+  (UPat(Ops.PARAM, tag="usb_vram"), lambda ctx: _words(ctx)),
+  (UPat(Ops.PARAM, tag="usb_asm24"), lambda ctx: _asm24(ctx)),
+  (UPat(Ops.PARAM, name="b"), lambda b: Buffer("CPU", b.max_numel(), b.dtype, preallocate=True) if str(b.tag).startswith("cmdbuf_copy") else None),
+]) # the sdma cmdbuf streams into the ring from the host
+
 if DEV.interface.startswith("MOCK"): from test.mockgpu.usb import MockUSB3 as USB3  # type: ignore  # noqa: F811

tinygrad/uop/spec.py

@@ -21,11 +21,6 @@ def validate_index(uidx:UOp, gate:UOp|None=None):
   # We can use UOp min/max to do a faster check, but it can give false positive since its not an exact bound and doesn't consider the mask
   if 0<=idx.vmin and idx.vmax<sz: return True
 
-  # TODO: validate STACK, z3 can't model vectors
-  # don't descend into PARAM shape metadata; only the PARAM value participates in index arithmetic
-  for x in idx.toposort(gate=lambda x: x.op is not Ops.PARAM) | gate.toposort(gate=lambda x: x.op is not Ops.PARAM):
-    if x.op is Ops.STACK: return True
-
   # if all is good and CHECK_OOB=1, validate with z3
   from tinygrad.uop.validate import validate_index_with_z3
   return validate_index_with_z3(sz, idx, gate)

tinygrad/uop/validate.py

@@ -1,5 +1,5 @@
 from typing import Callable
-from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, python_alu
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, python_alu, range_str
 from tinygrad.dtype import dtypes, Invalid
 from tinygrad.helpers import cpu_profile
 import z3
@@ -45,7 +45,8 @@ z3_renderer = PatternMatcher([
   # the valid condition is a constraint
   (UPat.var("cond").where(UPat.var("x"), UPat(Ops.CONST, arg=Invalid)), lambda x,cond,ctx: ctx[0].add(ctx[1][cond]) or ctx[1][x]),
   # variables
-  (UPat((Ops.SPECIAL, Ops.RANGE), name="x"), lambda x,ctx: create_bounded(x.render(simplify=False), 0, ctx[1][x.src[0]]-1, ctx[0])),
+  (UPat((Ops.SPECIAL, Ops.RANGE), name="x"), lambda x,ctx:
+   create_bounded(x.arg if x.op is Ops.SPECIAL else f"r{range_str(x)}", 0, ctx[1][x.src[0]]-1, ctx[0])),
   # unknown values are variables bounded by their vmin/vmax: params, loads (non-pointer INDEX is a LOAD) and anything from floats
   (UPat((Ops.PARAM, Ops.BUFFER, Ops.LOAD, Ops.INDEX), name="x"), create_var),
   (UPat((Ops.CAST, Ops.BITCAST)+tuple(GroupOp.Comparison), src=UPat(dtype=dtypes.floats), name="x"), create_var),
@@ -65,8 +66,6 @@ def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]:
                                       (x.dtype in dtypes.ints+(dtypes.bool, dtypes.weakint) or x.op is Ops.SINK)))[:-1]
   z3map: dict[UOp, z3.ExprRef] = {}
   for u in lst:
-    # NOTE: we skip STACK here, it can't actually be accessed
-    if u.op is Ops.STACK: continue
     if (z3_rewritten:=z3_renderer.rewrite(u, ctx=(solver, z3map))) is None: raise NotImplementedError(f"{u.op} is not supported by z3")
     z3map[u] = z3_rewritten
   assert all(u in z3map for u in uops), "UOp failed to rewrite to z3!"

tinygrad/viz/serve.py

@@ -479,10 +479,11 @@ def get_profile(data:VizData, profile:list[ProfileEvent], sort_fn:Callable[[str]
   scache:dict[str, int] = {}
   peaks:list[int] = []
   dtype_size:dict[str, int] = {}
-  for k,v in dev_events.items():
-    v.sort(key=lambda e:e[0])
-    layout[k] = timeline_layout(data, v, start_ts, scache)
-    layout.update([graph_layout(k, v, start_ts, unwrap(end_ts), peaks, dtype_size, scache)])
+  with soft_err():
+    for k,v in dev_events.items():
+      v.sort(key=lambda e:e[0])
+      layout[k] = timeline_layout(data, v, start_ts, scache)
+      layout.update([graph_layout(k, v, start_ts, unwrap(end_ts), peaks, dtype_size, scache)])
   sorted_layout = sorted([k for k,v in layout.items() if v is not None], key=sort_fn)
   ret = [b"".join([struct.pack("<B", len(k)), k.encode(), unwrap(layout[k])]) for k in sorted_layout]
   index = json.dumps({"strings":list(scache), "dtypeSize":dtype_size,