bump tinygrad usb wrap fix (#39204)

openpilot/starpilot/models/tests/test_modeld_jit_copies.py

@@ -0,0 +1,20 @@
+"""The production modern copy helper must see fresh NumPy storage on every call."""
+import __future__
+import ast
+from pathlib import Path
+
+import numpy as np
+
+
+def test_modern_copy_jit_reloads_mutated_host_inputs():
+  from tinygrad.tensor import Tensor
+  path = Path(__file__).resolve().parents[3] / 'selfdrive/modeld/modeld.py'
+  function = next(n for n in ast.parse(path.read_text()).body if isinstance(n, ast.FunctionDef) and n.name == 'compile_copy')
+  namespace = {}
+  exec(compile(ast.Module(body=[function], type_ignores=[]), str(path), 'exec', flags=__future__.annotations.compiler_flag), namespace)
+  source = np.arange(32, dtype=np.float32)
+  tensor = Tensor(source, device='NPY').to('CPU')
+  copy = namespace['compile_copy'](tensor)
+  for iteration in range(5):
+    source[:] = np.arange(32, dtype=np.float32) + iteration * 0.25
+    np.testing.assert_array_equal(copy().numpy(), source)

openpilot/starpilot/models/tests/test_stock_model_error_state.py

@@ -105,6 +105,8 @@ class TestStockModelErrorState(unittest.TestCase):
 
     model.run_warp, model.run_model = run_warp, run_model
     model.outputs = {'outputs': Output()}
+    model.copy_inputs = lambda: model.input_device.copy_from(model.input_host)
+    model.copy_outputs = lambda: model.outputs['outputs']
     model.output_slices = {'fixture': slice(0, 2)}
     parsed = {'fixture': np.array([[1., 2.]], dtype=np.float32)}
 

tinygrad_repo/.github/actions/setup-tinygrad/action.yml

@@ -74,17 +74,22 @@ runs:
       with:
         enable-cache: 'false' # see below for manual caching
 
-    # **** Caching packages ****
+    # **** Caching ****
+    # NOTE: self-hosted CI (ie. collaborator PRs) should never cache anything, instead, either:
+    #   * bake into the CI image: https://git.tinygrad.win/tinygrad/images, OR
+    #   * handle in CI runner HTTP cache
+    # Because CI for master and non-collaborator PRs are github-hosted, they need caching to workaround the flaky internet.
+    # Additionally, the actions/cache/restore pattern should be used to avoid PRs polluting the limited github cache storage.
 
     - name: Cache Python packages (PR)
-      if: github.event_name == 'pull_request' && env.RUNNER_ENVIRONMENT == 'github-hosted'
+      if: github.event_name == 'pull_request' && runner.environment == 'github-hosted'
       id: restore-venv-pr
       uses: actions/cache/restore@v5
       with:
         path: /tmp/.uv-cache
         key: uv-${{ runner.os }}-${{ runner.arch }}-python-${{ inputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
     - name: Cache Python packages
-      if: github.event_name != 'pull_request'
+      if: github.event_name != 'pull_request' && runner.environment == 'github-hosted'
       id: restore-venv
       uses: actions/cache@v5
       with:
@@ -94,13 +99,13 @@ runs:
     # **** Caching downloads ****
 
     - name: Cache downloads (PR)
-      if: inputs.key != '' && github.event_name == 'pull_request' && env.RUNNER_ENVIRONMENT == 'github-hosted'
+      if: inputs.key != '' && github.event_name == 'pull_request' && runner.environment == 'github-hosted'
       uses: actions/cache/restore@v5
       with:
         path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
         key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
     - name: Cache downloads
-      if: inputs.key != '' && github.event_name != 'pull_request'
+      if: inputs.key != '' && github.event_name != 'pull_request' && runner.environment == 'github-hosted'
       uses: actions/cache@v5
       with:
         path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
@@ -205,13 +210,13 @@ runs:
         echo "installed=$installed" >> "$GITHUB_OUTPUT"
 
     - name: Cache apt (PR)
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false' && env.RUNNER_ENVIRONMENT == 'github-hosted'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false' && runner.environment == 'github-hosted'
       uses: actions/cache/restore@v5
       with:
         path: /var/cache/apt/archives/
         key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
     - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false' && runner.environment == 'github-hosted'
       uses: actions/cache@v5
       with:
         path: /var/cache/apt/archives/

tinygrad_repo/.github/workflows/autogen.yml

@@ -1,7 +1,7 @@
 name: Autogen
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '14'
+  CACHE_VERSION: '15'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
@@ -53,7 +53,7 @@ jobs:
         python3 -c "from tinygrad.runtime.autogen import mesa"
         python3 -c "from tinygrad.runtime.autogen import avcodec"
         python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
-        python3 -c "from tinygrad.runtime.autogen import mlx5, bnxt"
+        python3 -c "from tinygrad.runtime.autogen import mlx5, bnxt, nvme"
         python3 -c "from tinygrad.runtime.autogen import ggml_common"
         REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
     - name: Check for differences

tinygrad_repo/.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '20'
+  CACHE_VERSION: '21'
   CAPTURE_PROCESS_REPLAY: ${{ github.event_name == 'pull_request' && contains(github.event.pull_request.title, '[pr]') && '1' || '0' }}
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}

tinygrad_repo/extra/export_model.py

@@ -66,13 +66,16 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
   inputs = [(name, dtype_map[bufs[name][1]], bufs[name][0]) for name in input_names + list(symbolic_vars.values())]
   outputs = [(name, dtype_map[bufs[name][1]], bufs[name][0]) for name in output_names]
   forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs))
+  def kernel_call(name:str, args:List[str]) -> str:
+    fields = " ".join(f"{dtype_map[bufs[a][1]] if a in symbolic_vars.values() else 'void*'} a{i};" for i,a in enumerate(args))
+    return f"{name}((unsigned char*)&(struct {{ {fields} }}){{{', '.join(args)}}});"
 
   if not wasm:
     for name,cl in bufs_to_save.items():
       weight = ''.join(["\\x%02X"%x for x in cl.as_memoryview()])
       cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
     cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
-    cprog += [f"void net({forward_args}) {{"] + [f"{name}(&(struct {name}_args){{{', '.join(args)}}});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    cprog += [f"void net({forward_args}) {{"] + [kernel_call(name, args) for (name, args, _global_size, _local_size) in statements] + ["}"]
     return '\n'.join(headers + cprog)
   else:
     if bufs_to_save:
@@ -88,7 +91,7 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
 
     cprog += [f"void net({forward_args})"] + ["{"]
     get_weight_ptr = lambda x: f"({dtype_map[bufs_to_save[x][1]]} *)bufs[{buf_to_name[x]['idx']}]" if x in bufs_to_save else x
-    cprog += [f"  {name}(&(struct {name}_args){{{', '.join(map(get_weight_ptr, args))}}});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    cprog += ["  " + kernel_call(name, list(map(get_weight_ptr, args))) for (name, args, _global_size, _local_size) in statements] + ["}"]
     weightMapping = "" if not bufs_to_save else f"""\nconst weightNames = [{", ".join([f'"{weight_name}"' for weight_name in [v["name"] for v in buf_to_name.values()]])}];
 const {model_name}_name_to_id = Object.fromEntries(weightNames.map((name, index) => [name, index]));\n"""
     top = f"""import {model_name}Module from './{model_name}.js'{weightMapping}"""

tinygrad_repo/extra/nvme_driver/bench_nvme.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+# GPUS=4 DRIVES=4 GB=2 python3 extra/nvme_driver/bench_nvme.py: each gpu reads GB GiB from its DRIVES drives at once. never writes
+import time
+from tinygrad import Tensor, Device, TinyJit, dtypes
+from tinygrad.helpers import getenv
+
+GPUS, DRIVES, N = getenv("GPUS", 4), getenv("DRIVES", 4), getenv("GB", 2) << 30
+gpus = [f"AMD:{g}" for g in range(GPUS)]
+srcs = [Tensor.empty(sz:=Device[f"NVME:{d}"].nvme.size, dtype=dtypes.uint8, device=f"NVME:{d}")[sz // 2:sz // 2 + N] for d in range(GPUS * DRIVES)]
+dsts = [Tensor.empty(N, dtype=dtypes.uint8, device=gpus[d // DRIVES]).realize() for d in range(GPUS * DRIVES)]
+@TinyJit
+def read() -> list[Tensor]: return [dst.assign(src.to(dst.device)).realize() for dst, src in zip(dsts, srcs)]
+def run(n:int) -> float:
+  for g in gpus: Device[g].synchronize()
+  st = time.perf_counter()
+  for _ in range(n): read()
+  for g in gpus: Device[g].synchronize()
+  return time.perf_counter() - st
+run(3)
+for n in (1, getenv("ITERS", 8)):
+  dt = run(n)
+  print(f"{len(srcs)} drives -> {GPUS} gpus, {n} x {len(srcs) * N / 2**30:.0f} GiB in {dt * 1e3:.0f} ms: "
+        f"{n * len(srcs) * N / dt / 1e9:.1f} GB/s, {n * DRIVES * N / dt / 1e9:.1f} GB/s a gpu")

tinygrad_repo/extra/nvme_driver/test_nvme.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+# NVME=0 GPU=AMD python3 extra/nvme_driver/test_nvme.py: reads of the last 256 MiB of a drive (which must hold data) into gpu memory. never writes
+import unittest
+import numpy as np
+from tinygrad import Tensor, Device, TinyJit, dtypes
+from tinygrad.helpers import getenv
+
+NVME, GPU, N = f"NVME:{getenv('NVME', 0)}", getenv("GPU", "AMD"), 256 << 20
+
+class TestNVME(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.region = Tensor.empty(size:=Device[NVME].nvme.size, dtype=dtypes.uint8, device=NVME)[size - N:]
+    cls.ref = cls.region.to(GPU).numpy()
+    assert np.count_nonzero(np.bincount(cls.ref)) > 64, "the last 256 MiB of the drive hold no data to compare"
+  def check(self, reads:list[tuple[int, int]]):
+    Tensor.realize(*(ts:=[self.region[off:off + n].to(GPU) for off, n in reads]))
+    for t, (off, n) in zip(ts, reads): np.testing.assert_equal(t.numpy(), self.ref[off:off + n], err_msg=f"{n} bytes at {off}")
+
+  def test_reads(self):
+    for read in [(0, 4096), (4096, 8192), (0, 3 * 4096), (0, 1 << 20), (12288, (1 << 20) + 4096), (4096, N - 8192)]: self.check([read])
+  def test_reads_one_submit(self): self.check([(0, 1 << 20), (8 << 20, 3 << 20)])
+  def test_jit(self):
+    @TinyJit
+    def f(t:Tensor) -> Tensor: return (t.to(GPU) + 1).realize()
+    for _ in range(5): np.testing.assert_equal(f(self.region[1 << 20:3 << 20]).numpy(), self.ref[1 << 20:3 << 20] + 1)
+
+if __name__ == "__main__": unittest.main()

tinygrad_repo/test/amd/hw/helpers.py

@@ -136,6 +136,7 @@ def get_prologue_epilogue(n_lanes: int) -> tuple[list, list]:
   for i in range(N_SGPRS):
     prologue.append(s_mov_b32(s[i], 0))
   prologue.append(s_mov_b32(VCC_LO, 0))
+  prologue.append(s_cmp_eq_u32(0, 1))  # initialize SCC independently of the emulator/runtime's entry state
 
   epilogue = [
     s_mov_b32(s[90], VCC_LO),

tinygrad_repo/test/amd/hw/test_cdna_sdwa.py

@@ -1,11 +1,12 @@
 #!/usr/bin/env python3
 import unittest
 import tinygrad.runtime.autogen.amd.cdna.ins as cdna
-from test.amd.hw.test_cdna_vop3 import run_cdna
+from test.amd.hw.test_cdna_vop3 import run_cdna, run_cdna_lanes
 
 class TestCDNASDWA(unittest.TestCase):
   def test_v_add_co_u32_e32_writes_vcc(self):
     out = run_cdna([
+      cdna.s_mov_b64(cdna.VCC, 0),
       cdna.s_mov_b32(cdna.s[0], 0xffffffff),
       cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]),
       cdna.v_mov_b32_e32(cdna.v[13], 1),
@@ -17,5 +18,17 @@ class TestCDNASDWA(unittest.TestCase):
     ])
     self.assertEqual(out, 0x80000000)
 
+  def test_v_add_co_u32_e32_vcc_multiple_lanes(self):
+    out = run_cdna_lanes([
+      cdna.s_mov_b32(cdna.VCC_LO, 0),
+      cdna.s_mov_b32(cdna.s[0], 0xffffffff),
+      cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]),
+      cdna.v_mov_b32_e32(cdna.v[13], 1),
+      cdna.v_add_co_u32_e32(cdna.v[0], cdna.SDWA, cdna.v[13], vsrc0=cdna.v[0], dst_sel=6, src0_sel=6),
+      cdna.s_nop(1),
+      cdna.v_mov_b32_e32(cdna.v[2], cdna.VCC_LO),
+    ], lanes=2)
+    self.assertEqual(out, [3, 3])
+
 if __name__ == "__main__":
   unittest.main()

tinygrad_repo/test/amd/hw/test_sop.py

@@ -267,6 +267,12 @@ class Test64BitLiterals(unittest.TestCase):
 class TestSCCBehavior(unittest.TestCase):
   """Tests for SCC condition code behavior."""
 
+  def test_scc_initial(self):
+    """The shared prologue clears SCC before the test instructions run."""
+    st = run_program([s_cselect_b32(s[0], 1, 0)], n_lanes=1)
+    self.assertEqual(st.sgpr[0], 0)
+    self.assertEqual(st.scc, 0)
+
   def test_scc_from_s_cmp(self):
     """SCC should be set by scalar compare."""
     instructions = [

tinygrad_repo/test/amd/hw/test_vop1.py

@@ -1258,7 +1258,7 @@ class TestF16EdgeCases(unittest.TestCase):
 
 
 class TestExpEdgeCases(unittest.TestCase):
-  """Additional edge cases for V_EXP_F32."""
+  """Tests for V_EXP_F32 and V_EXP_F16."""
 
   def test_v_exp_f32_zero(self):
     """V_EXP_F32 of 0.0 returns 1.0 (2^0 = 1)."""
@@ -1287,6 +1287,19 @@ class TestExpEdgeCases(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     self.assertAlmostEqual(i2f(st.vgpr[0][1]), 0.5, places=5)
 
+  def test_v_exp_f16(self):
+    """V_EXP_F16 computes 2^x in half precision: 2^3 = 8, 2^-1 = 0.5."""
+    instructions = [
+      s_mov_b32(s[0], f32_to_f16(3.0)),
+      s_mov_b32(s[1], f32_to_f16(-1.0)),
+      v_mov_b32_e32(v[0], s[0]),
+      v_mov_b32_e32(v[2], s[1]),
+      v_exp_f16_e32(v[1], v[0]),
+      v_exp_f16_e32(v[3], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual((f16(st.vgpr[0][1]), f16(st.vgpr[0][3])), (8.0, 0.5))
+
   def test_v_exp_f32_inf(self):
     """V_EXP_F32 of +inf returns +inf."""
     import math

tinygrad_repo/test/amd/hw/test_vop3.py

@@ -4,7 +4,6 @@ Includes: v_fma_f32, v_div_scale_f32, v_div_fmas_f32, v_div_fixup_f32,
           v_alignbit_b32, v_bfe_i32, v_mad_u64_u32, v_readlane_b32, v_writelane_b32
 """
 import unittest
-from tinygrad.helpers import OSX
 from test.amd.hw.helpers import *
 
 class TestFMA(unittest.TestCase):
@@ -3721,7 +3720,6 @@ class TestClampLdExpRegressions(unittest.TestCase):
 class TestNaNPropagationRegressions(unittest.TestCase):
   """Regression tests: float arithmetic propagates a NaN from the FIRST NaN operand, quieted with its own sign/payload."""
 
-  @unittest.skipIf(OSX, "broken on mac, TODO: why?")
   def test_mul_nan_priority(self):
     # first NaN operand wins (sign+payload), not x86's second-source propagation
     for a, b, want in [(0x7FC00001, 0x7F800003, 0x7FC00001), (0xFFC00005, 0x7F800003, 0xFFC00005),

tinygrad_repo/test/device/dsp/test_quantize_onnx.py

@@ -115,11 +115,9 @@ class TestQuantizeOnnx(unittest.TestCase):
         unsigned_char128 hi128;
       };
     };
-    struct fxn_args { unsigned char* data0; unsigned char* data1; signed char* data2; };
-    __attribute__((noinline)) void fxn(const struct fxn_args* args) {
-      unsigned char* restrict __attribute__((align_value(128))) data0 = args->data0;
-      unsigned char* restrict __attribute__((align_value(128))) data1 = args->data1;
-      signed char* restrict __attribute__((align_value(128))) data2 = args->data2;
+    __attribute__((noinline)) void fxn(unsigned char* restrict __attribute__((align_value(128))) data0,
+                                       unsigned char* restrict __attribute__((align_value(128))) data1,
+                                       signed char* restrict __attribute__((align_value(128))) data2) {
       for (int ridx0 = 0; ridx0 < 512; ridx0++) {
         int alu0 = (ridx0<<9);
         for (int ridx1 = 0; ridx1 < 4; ridx1++) {

tinygrad_repo/test/external/external_test_usb_asm24.py

@@ -1,7 +1,13 @@
 import unittest
-from tinygrad.helpers import Timing, getenv
-from tinygrad import Tensor, Device, TinyJit
+from tinygrad.helpers import Timing, getenv, unwrap, wait_cond
+from tinygrad import Tensor, Device, TinyJit, dtypes
+from tinygrad.device import Buffer
+from tinygrad.engine.realize import run_linear
+from tinygrad.uop.ops import UOp
+from tinygrad.runtime.ops_amd import AMDSDMAQueue
+from tinygrad.runtime.support.hcq2 import make_submit
 from tinygrad.runtime.support.usb import HALF, CHUNK, SLOT
+from test.null.test_hcq2 import lower_hcq
 import numpy as np
 
 class USBTestCase(unittest.TestCase):
@@ -81,10 +87,26 @@ class TestUSBIntegrity(USBTestCase):
       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
+  def testRingWrap(self):
+    self.dev.synchronize()
+    q = unwrap(self.dev.sdma_queue(0))
+    put, read = [b.host.view(fmt='Q') for b in (q.put_value, q.read_ptr)]
+    ring_size = q.ring.nbytes
+    # advance with NOPs, leave one dword so the next eager copy forces wrapping
+    target = (put[0] // ring_size + 1) * ring_size - 4
+    if nbytes := target - put[0]:
+      nops = UOp.from_buffer(Buffer("CPU", nbytes, initial_value=bytes(nbytes)), dtypes.uint8)
+      hq = AMDSDMAQueue(make_submit(dev=self.dev.device, queue="COPY:0", deps=(nops,)))
+      run_linear(lower_hcq(hq.submit(nops)))
+    wait_cond(lambda: read[0], value=target)
+    self.assertEqual(put[0], target)
+
     a = self.rng.integers(0, 256, 64 << 20, dtype=np.uint8)
     t = Tensor(a, device="NPY").to(Device.DEFAULT).realize()
+    self.assertGreater(put[0] // ring_size, target // ring_size)
     np.testing.assert_array_equal(a, t.numpy())
+    self.dev.synchronize()
+    wait_cond(lambda: read[0], value=put[0])
 
 if __name__ == "__main__":
   unittest.main()

tinygrad_repo/test/mockgpu/amd/amdgpu.py

@@ -193,12 +193,6 @@ class PM4Executor(AMDQueue):
     args_addr = self.gpu.regs[regCOMPUTE_USER_DATA_0] + (self.gpu.regs[regCOMPUTE_USER_DATA_0 + 1] << 32)
     lc = [self.gpu.regs[i] for i in range(regCOMPUTE_NUM_THREAD_X, regCOMPUTE_NUM_THREAD_X+3)]
     rsrc2 = self.gpu.regs[regCOMPUTE_PGM_RSRC2]
-    # Read all user data registers (hardware loads these directly into s[0:N])
-    user_sgpr_count = (rsrc2 >> 1) & 0x1F  # USER_SGPR_COUNT is bits 1:5
-    user_data = []
-    for i in range(user_sgpr_count):
-      try: user_data.append(self.gpu.regs[regCOMPUTE_USER_DATA_0 + i])
-      except KeyError: user_data.append(0)
 
     prg_sz = 0
     for st,sz in sorted(self.gpu.mapped_ranges):
@@ -216,7 +210,7 @@ class PM4Executor(AMDQueue):
     assert 0 < self.kernel_text_size <= prg_sz, "Invalid prg ptr or text size"
     from test.mockgpu.amd.emu import run_asm
     err = run_asm(prg_addr, self.kernel_text_size, gl[0], gl[1], gl[2], lc[0], lc[1], lc[2], args_addr,
-                  rsrc2, scratch_size, self.gpu.arch, user_data, self.kernel_name)
+                  rsrc2, scratch_size, self.gpu.arch, self.kernel_name)
     if err != 0: raise RuntimeError("remu does not support the new instruction introduced in this kernel")
 
   def _exec_indirect_buffer(self, n):

tinygrad_repo/test/mockgpu/amd/emu.py

@@ -1,11 +1,11 @@
 # RDNA3 emulator v2 - compiles pcode to UOps executed via tinygrad CPU backend
 # Each instruction is compiled to a kernel that operates on buffers:
-#   arg=0: sgpr - sgpr[0-127], inline constants[128-255], PC_LO=256, PC_HI=257, SCC=258, SCRATCH_STRIDE=259
+#   arg=0: sgpr - sgpr[0-127], inline constants[128-255], SCC=258
 #   arg=1: vgpr - vgpr[reg * 32 + lane]
 #   arg=2: vmem - base address 0, INDEX offsets directly to host memory
 #   arg=3: lds - local data share
 #   arg=4: scratch - per-lane scratch memory
-import ctypes, functools, itertools, re, struct, platform, subprocess, tempfile, traceback
+import ctypes, functools, itertools, re, platform, subprocess, tempfile, traceback
 from typing import Callable
 
 # Set/restore DAZ+FTZ (denormals-are-zero + flush-to-zero) to match RDNA3 default float mode
@@ -50,7 +50,7 @@ class _MXCSRContext:
     if lib is None or not hasattr(self, '_saved'): return
     lib.set_fpcr(self._saved)
 
-from tinygrad.uop.ops import UOp, sint, Ops, KernelInfo, uopfunc, AxisType, rewrite_group
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, uopfunc, AxisType, rewrite_group
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.device import Buffer, BufferSpec, Device
 from tinygrad.engine.realize import run_linear
@@ -75,12 +75,12 @@ def _c(val, dtype=dtypes.uint32): return UOp.const(val, dtype)
 
 def _u64(lo: UOp, hi: UOp) -> UOp:
   """Combine two 32-bit UOps into a 64-bit UOp."""
-  return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
+  return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << _c(32, dtypes.uint64))
 
 def _split64(val: UOp) -> tuple[UOp, UOp]:
   """Split a 64-bit value into (lo, hi) 32-bit values."""
-  v64 = val.bitcast(dtypes.uint64) if val.dtype == dtypes.float64 else val.cast(dtypes.uint64) if val.dtype != dtypes.uint64 else val
-  return v64.cast(dtypes.uint32), (v64 >> UOp.const(32, dtypes.uint64)).cast(dtypes.uint32)
+  v64 = val.bitcast(dtypes.uint64) if val.dtype == dtypes.float64 else val.cast(dtypes.uint64)
+  return v64.cast(dtypes.uint32), (v64 >> _c(32, dtypes.uint64)).cast(dtypes.uint32)
 
 _SRC_MOD_TYPES = {16: (dtypes.uint16, dtypes.half, 0x7FFF), 32: (dtypes.uint32, dtypes.float32, 0x7FFFFFFF),
                   64: (dtypes.uint64, dtypes.float64, 0x7FFFFFFFFFFFFFFF)}
@@ -89,9 +89,9 @@ def _apply_src_mods(val: UOp, mod_bit: int, abs_bits: int, neg_bits: int, bits:
   if not (abs_bits & (1 << mod_bit)) and not (neg_bits & (1 << mod_bit)): return val
   ut, ft, mask = _SRC_MOD_TYPES[bits]
   fv = val.cast(ut).bitcast(ft) if bits == 16 else val.bitcast(ft) if val.dtype == ut else val
-  if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & UOp.const(mask, ut)).bitcast(ft)
+  if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & _c(mask, ut)).bitcast(ft)
   # neg modifier is a pure sign-bit toggle (preserves NaN payloads), not an arithmetic negate
-  if neg_bits & (1 << mod_bit): fv = (fv.bitcast(ut) ^ UOp.const((mask + 1) & (1 << (bits - 1)), ut)).bitcast(ft)
+  if neg_bits & (1 << mod_bit): fv = (fv.bitcast(ut) ^ _c((mask + 1) & (1 << (bits - 1)), ut)).bitcast(ft)
   return fv.bitcast(ut).cast(dtypes.uint32) if bits == 16 else fv.bitcast(ut)
 
 # Map VOPD ops to VOP2/VOP1 ops for pcode lookup (both RDNA3 and RDNA4 share these targets)
@@ -113,17 +113,15 @@ def _iattr(inst, name: str, default: int = 0) -> int:
   """Optional integer attribute of a decoded instruction (None/missing -> default)."""
   v = getattr(inst, name, default)
   return default if v is None else v
-# Special registers stored after inline constants (256-259)
-SCRATCH_STRIDE_IDX = 259
-# SGPR buffer: 0-127 = SGPRs, 128-255 = inline constants, 256-259 = special registers
-SGPR_COUNT = 260
+# SGPR buffer: 0-127 = SGPRs, 128-255 = inline constants
+SGPR_COUNT = 256
 
 def _op_name(inst) -> str:
   if hasattr(inst, 'opx'): return f"{inst.opx.name}_{inst.opy.name}"  # VOPD has opx/opy not op
-  return inst.op.name if hasattr(inst.op, 'name') else str(inst.op)
+  return inst.op.name
 
 def _lane_active(exec_mask: UOp, lane: UOp) -> UOp:
-  if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & UOp.const(1, dtypes.uint64)).ne(UOp.const(0, dtypes.uint64))
+  if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & _c(1, dtypes.uint64)).ne(_c(0, dtypes.uint64))
   return ((exec_mask >> lane.cast(dtypes.uint32)) & _c(1)).ne(_c(0))
 def _hi16(v: UOp) -> UOp: return (v >> _c(16)) & _c(0xFFFF)
 def _cond(cond, if_true, if_false):
@@ -134,7 +132,6 @@ def _apply_opsel(val: UOp, sel_bit: int, opsel: int) -> UOp: return _hi16(val) i
 
 def _val_to_u32(val: UOp) -> UOp:
   """Convert any value to uint32 for storage (bitcast floats, cast ints)."""
-  if val.dtype == dtypes.uint32: return val
   if val.dtype == dtypes.float32: return val.bitcast(dtypes.uint32)
   if val.dtype == dtypes.half: return val.bitcast(dtypes.uint16).cast(dtypes.uint32)
   return val.cast(dtypes.uint32)
@@ -168,6 +165,7 @@ _pcode_fixes = {
   'V_CMP_CLASS_F32': ('64\'F(abs(S0.f32)) > 0.0', '(64\'U(S0.u32 & 0x7FFFFFFF) != 0)'),
   'V_CMP_CLASS_F16': ('64\'F(abs(S0.f16)) > 0.0', '(64\'U(S0.u32 & 0x7FFF) != 0)'),
   'V_CMP_CLASS_F64': ('64\'F(abs(S0.f64)) > 0.0', '(64\'U(S0.u64 & 0x7FFFFFFFFFFFFFFF) != 0)'),
+  'V_MUL_F32': ('D0.f32 = S0.f32 * S1.f32', 'D0.f32 = isNAN(S0.f32) ? cvtToQuietNAN(S0.f32) : isNAN(S1.f32) ? cvtToQuietNAN(S1.f32) : S0.f32*S1.f32'),
 }
 
 def _get_pcode_dict(op) -> dict:
@@ -223,7 +221,7 @@ def _write_64bit(val: UOp, wfn, reg_or_addr, is_mem: bool, *args) -> list[UOp]:
   """Write a 64-bit value as two 32-bit writes. args passed to wfn after reg/addr and lo/hi value."""
   lo, hi = _split64(val)
   incr = 4 if is_mem else 1  # 4 bytes for memory addresses, 1 for register indices
-  return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (UOp.const(incr, reg_or_addr.dtype) if isinstance(reg_or_addr, UOp) else incr), hi, *args)]
+  return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (_c(incr, reg_or_addr.dtype) if isinstance(reg_or_addr, UOp) else incr), hi, *args)]
 
 def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = False) -> list[UOp]:
   """Write value, splitting 64-bit if needed. bits=64 for 64-bit writes, otherwise 32-bit."""
@@ -232,10 +230,10 @@ def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = Fals
 def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int = 32) -> list[UOp]:
   """Store to byte-addressable memory. addr is byte offset, mem is uint8 buffer."""
   stores = []
-  val_u32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
+  val_u32 = val.cast(dtypes.uint32)
   for i in range(data_bits // 8):
-    byte_val = (val_u32 >> UOp.const(i * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
-    stores.append(mem.index((addr + UOp.const(i, dtypes.uint64)).valid(active)).store(byte_val.cast(dtypes.uint8)))
+    byte_val = (val_u32 >> _c(i * 8)) & _c(0xFF)
+    stores.append(mem.index((addr + _c(i, dtypes.uint64)).valid(active)).store(byte_val.cast(dtypes.uint8)))
   return stores
 
 def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode_vars: dict | None = None, op_name: str = "") -> dict[int, UOp]:
@@ -291,11 +289,12 @@ def _int_clamp(op_name: str, srcs: dict) -> UOp | None:
 
 class _Ctx:
   """Context for instruction compilation - holds buffers and helpers."""
-  __slots__ = ('dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr', 'branch_cond')
+  __slots__ = ('dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr', 'branch_cond', 'sgpr_reads')
   sgpr = UOp.param(0, dtypes.uint32, SGPR_COUNT, name="sgpr")
   vmem = UOp.param(2, dtypes.uint32, 1 << 46, name="vmem")
   lds = UOp.param(3, dtypes.uint32, 16384, name="lds")
   scratch = UOp.param(4, dtypes.uint8, 1 << 30, name="scratch")
+  scratch_stride = UOp.param(9, dtypes.uint32, name="scratch_stride", addrspace=AddrSpace.ALU)
 
   def __init__(self, wave_size: int = 32):
     self._axis_id, self.wave_size = 0, wave_size
@@ -303,6 +302,7 @@ class _Ctx:
     self.vgpr = UOp.param(1, dtypes.uint32, 256 * wave_size, name="vgpr")
     self.accvgpr = UOp.param(5, dtypes.uint32, 256 * wave_size, name="accvgpr") if wave_size == 64 else self.vgpr
     self.branch_cond: UOp | None = None
+    self.sgpr_reads: dict[UOp, UOp] = {}
 
   def range(self, n: int | None = None) -> UOp:
     """Create a lane range UOp with unique axis ID."""
@@ -331,15 +331,15 @@ class _Ctx:
     word = self.inst_word(dword_idx)
     if lo // 32 == hi // 32:  # Same dword
       mask = (1 << (hi - lo + 1)) - 1
-      shifted = word if lo_in_dword == 0 else word >> UOp.const(lo_in_dword, dtypes.uint32)
-      return shifted & UOp.const(mask, dtypes.uint32)
+      shifted = word if lo_in_dword == 0 else word >> _c(lo_in_dword)
+      return shifted & _c(mask)
     else:  # Spans two dwords
       lo_bits = 32 - lo_in_dword
       lo_mask = (1 << lo_bits) - 1
       hi_mask = (1 << (hi_in_dword + 1)) - 1
-      lo_part = (word >> UOp.const(lo_in_dword, dtypes.uint32)) & UOp.const(lo_mask, dtypes.uint32)
-      hi_part = self.inst_word(dword_idx + 1) & UOp.const(hi_mask, dtypes.uint32)
-      return lo_part | (hi_part << UOp.const(lo_bits, dtypes.uint32))
+      lo_part = (word >> _c(lo_in_dword)) & _c(lo_mask)
+      hi_part = self.inst_word(dword_idx + 1) & _c(hi_mask)
+      return lo_part | (hi_part << _c(lo_bits))
 
   def optional_field(self, inst: Inst, name: str) -> UOp | None:
     """Extract a field that only exists on some instruction formats (e.g. 'literal', 'saddr'). None if absent."""
@@ -370,15 +370,14 @@ class _Ctx:
 
   # Dynamic register access (takes UOp index instead of int)
   def rsgpr_dyn(self, reg: UOp, valid: UOp | None = None) -> UOp:
-    """Read SGPR with dynamic register index."""
-    if valid is not None: return self.sgpr.index(reg.valid(valid)).load()
-    return self.sgpr.index(reg).load()
+    index = reg.valid(valid) if valid is not None else reg
+    if index not in self.sgpr_reads: self.sgpr_reads[index] = UOp.param(SGPR_COUNT+len(self.sgpr_reads), dtypes.uint32, addrspace=AddrSpace.ALU)
+    return self.sgpr_reads[index]
 
   def wsgpr_dyn(self, reg: UOp, val: UOp) -> UOp:
     """Write SGPR with dynamic register index. On RDNA, index 124 = NULL (writes discarded). On CDNA, index 124 = M0 (read/write)."""
     # RDNA: NULL (124) discards writes. CDNA: M0 (124) is writable.
-    valid = None if self.wave_size == 64 else reg.ne(_c(124))
-    return self.sgpr.index(reg.valid(valid) if valid is not None else reg).store(val.cast(dtypes.uint32))
+    return self.sgpr.index(reg if self.wave_size == 64 else reg.valid(reg.ne(_c(124)))).store(val.cast(dtypes.uint32))
 
   def wmask(self, reg: UOp, val: UOp) -> list[UOp]:
     """Write a lane mask (VCC/EXEC). Splits into lo/hi for wave64."""
@@ -451,7 +450,7 @@ class _Ctx:
       inline = is_float_const.where(float_inline.bitcast(dtypes.uint64), int_inline.bitcast(dtypes.uint64))
       # Literal handling: F64 VOP puts literal in high 32 bits; B64/I64/U64 VOP and SOP zero-extend
       if literal is not None:
-        lit_val = literal.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64) if is_f64 else literal.cast(dtypes.uint64)
+        lit_val = literal.cast(dtypes.uint64) << _c(32, dtypes.uint64) if is_f64 else literal.cast(dtypes.uint64)
         inline = off.eq(_c(255)).where(lit_val, inline)
       scalar_val = (off < _c(128)).where(sgpr_val, inline)
     else:
@@ -492,16 +491,16 @@ class _Ctx:
   def compile_lane_pcode(self, op, inst) -> UOp:
     """Compile cross-lane ops (READLANE/WRITELANE/PERMLANE) using pcode parser."""
     pcode = get_pcode(op)
-    op_name = op.name if hasattr(op, 'name') else str(op)
+    op_name = op.name
     src0_off, vdst_off = self.inst_field(type(inst).src0), self.inst_field(type(inst).vdst)
     src0_reg = (src0_off >= _c(256)).where(src0_off - _c(256), _c(0))  # VGPR index or 0
     src1_off, src2_off = self.optional_field(inst, 'src1'), self.optional_field(inst, 'src2')
     src1_reg = (src1_off >= _c(256)).where(src1_off - _c(256), src1_off) if src1_off is not None else _c(0)
     src2_reg = (src2_off >= _c(256)).where(src2_off - _c(256), src2_off) if src2_off is not None else _c(0)
     exec_val = self.rexec()
-    exec_lo = exec_val.cast(dtypes.uint32) if exec_val.dtype == dtypes.uint64 else exec_val
+    exec_lo = exec_val.cast(dtypes.uint32)
     srcs = {
-      'SRC0': src0_reg, 'VDST': vdst_off, 'EXEC_LO': exec_lo, 'EXEC': exec_val if exec_val.dtype == dtypes.uint64 else exec_val.cast(dtypes.uint64),
+      'SRC0': src0_reg, 'VDST': vdst_off, 'EXEC_LO': exec_lo, 'EXEC': exec_val.cast(dtypes.uint64),
       '_vgpr': self.vgpr, '_wave_size': self.wave_size, 'SRC1': src1_reg, 'SRC2': src2_reg,
       'S0': self.rsrc_dyn(src0_off, _c(0, dtypes.int)) if 'WRITELANE' in op_name else src0_reg,
       'S1': self.rsrc_dyn(src1_off, _c(0, dtypes.int)) if src1_off is not None else _c(0),
@@ -531,7 +530,7 @@ class _Ctx:
     pcode = get_pcode(op)
     vcc_reg = sdst_reg if sdst_reg is not None else VCC_LO.offset
     if 'VCC' not in srcs: srcs['VCC'] = self.rmask(_c(vcc_reg))
-    srcs.update(self.base_srcs(exec_mask, lane), VDST=vdst_reg, MAX_FLOAT_F32=UOp.const(3.4028234663852886e38, dtypes.float32))
+    srcs.update(self.base_srcs(exec_mask, lane), VDST=vdst_reg, MAX_FLOAT_F32=_c(3.4028234663852886e38, dtypes.float32))
     # f32 min/max/median ops flush denormal inputs to signed zero (select-style ops: results propagate inputs bitwise)
     # (RDNA4 calls them _NUM_: V_MIN_NUM_F32 etc.)
     if any(p in op.name for p in ('MIN_F32', 'MAX_F32', 'MIN3_F32', 'MAX3_F32', 'MED3_F32', 'MIN_NUM_F32', 'MAX_NUM_F32')):
@@ -554,23 +553,22 @@ class _Ctx:
       elif dest.startswith('D0'):
         if (dest_suffix := re.match(r'D0\.(\w+)', dest)) is not None:
           target_dt = {'u16': dtypes.uint16, 'i16': dtypes.int16, 'f16': dtypes.half}.get(dest_suffix.group(1))
-          if target_dt is not None and val.dtype != target_dt: val = val.cast(target_dt)
+          if target_dt is not None: val = val.cast(target_dt)
         if (slice_match := re.match(r'D0\[(\d+)\s*:\s*(\d+)\]', dest)) and (int(slice_match.group(1)), int(slice_match.group(2))) != (31, 0):
           slice_stores.append((int(slice_match.group(2)), int(slice_match.group(1)) - int(slice_match.group(2)) + 1, _val_to_bits(val)))
           continue
         if int_saturate is not None: val = int_saturate
         elif clmp and val.dtype in (dtypes.float32, dtypes.half, dtypes.float64):
           # hardware clamp: -0 becomes +0 and NaN becomes 0 (hardware verified)
-          val = (val > UOp.const(0.0, val.dtype)).where(val.minimum(UOp.const(1.0, val.dtype)), UOp.const(0.0, val.dtype))
+          val = (val > _c(0.0, val.dtype)).where(val.minimum(_c(1.0, val.dtype)), _c(0.0, val.dtype))
         if val.dtype in (dtypes.uint64, dtypes.int64, dtypes.float64):
           lo, hi = _split64(val)
           lane_stores.extend([self.wvgpr_dyn(vdst_reg, lane, lo, exec_mask), self.wvgpr_dyn(vdst_reg + _c(1), lane, hi, exec_mask)])
         elif val.dtype in (dtypes.half, dtypes.uint16, dtypes.int16):
           result, old_val = _val_to_u32(val), self.rvgpr_dyn(vdst_reg, lane)
-          hi_result = (old_val & UOp.const(0xFFFF, dtypes.uint32)) | (result << UOp.const(16, dtypes.uint32))
+          hi_result = (old_val & _c(0xFFFF)) | (result << _c(16, dtypes.uint32))
           # GFX9/CDNA zeroes upper 16 bits on lo-half write; RDNA preserves them
-          lo_result = (result & UOp.const(0xFFFF, dtypes.uint32)) if self.wave_size == 64 else \
-                      (old_val & UOp.const(0xFFFF0000, dtypes.uint32)) | (result & UOp.const(0xFFFF, dtypes.uint32))
+          lo_result = (result & _c(0xFFFF)) if self.wave_size == 64 else (old_val & _c(0xFFFF0000)) | (result & _c(0xFFFF))
           result = opsel_dst_hi.where(hi_result, lo_result) if isinstance(opsel_dst_hi, UOp) else hi_result if opsel_dst_hi else lo_result
           lane_stores.append(self.wvgpr_dyn(vdst_reg, lane, result, exec_mask))
         else: lane_stores.append(self.wvgpr_dyn(vdst_reg, lane, _val_to_u32(val), exec_mask))
@@ -615,7 +613,7 @@ def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
     pcode = get_pcode(inst.op)
     vcc, exec_val = ctx.rmask(_c(VCC_LO.offset)), ctx.rexec()
     srcs = {'PC': _c(0, dtypes.uint32), 'SIMM16': simm16, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)), 'VCC': vcc,
-            'VCCZ': vcc.eq(UOp.const(0, vcc.dtype)).cast(dtypes.uint32), 'EXECZ': exec_val.eq(UOp.const(0, exec_val.dtype)).cast(dtypes.uint32)}
+            'VCCZ': vcc.eq(_c(0, vcc.dtype)).cast(dtypes.uint32), 'EXECZ': exec_val.eq(_c(0, exec_val.dtype)).cast(dtypes.uint32)}
     for dest, val in parse_pcode(pcode, srcs)[1]:
       if dest.startswith('PC') and val.op is Ops.WHERE: ctx.branch_cond = val.src[0]
   return UOp.sink()
@@ -761,7 +759,7 @@ def _load_dpp16_src0(ctx: _Ctx, inst, lane: UOp, fallback: UOp) -> UOp:
   src_lane, enabled, valid = _dpp16_ctrl(lane, _iattr(inst, 'dpp'), _iattr(inst, 'row_mask', 0xf), _iattr(inst, 'bank_mask', 0xf), ctx.wave_size)
   safe_src_lane = (enabled & valid).where(src_lane, _c(0, dtypes.int))
   swizzled = ctx.rvgpr_dyn(ctx.inst_field(type(inst).vsrc0), safe_src_lane)
-  invalid = UOp.const(0, fallback.dtype) if _iattr(inst, 'bc') else fallback
+  invalid = _c(0, fallback.dtype) if _iattr(inst, 'bc') else fallback
   return enabled.where(valid.where(swizzled, invalid), fallback)
 
 def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc.VOPC_SDWA_SDST, ctx: _Ctx) -> UOp:
@@ -785,17 +783,16 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc
 
   # For VOPC: use unroll_lanes to build the bitmask from scratch (no read-modify-write on stale data)
   if is_vopc:
-    def get_cmp_bit(lane) -> UOp:
-      lc = lane.cast(dtypes.int) if isinstance(lane, UOp) else _c(lane, dtypes.int)
+    def get_cmp_bit(lane: UOp) -> UOp:
+      lc = lane.cast(dtypes.int)
       s0_raw = ctx.rsgpr_dyn(vsrc0_reg) if inst.s0 else ctx.rvgpr_dyn(vsrc0_reg, lc)
       s0 = _sdwa_select(s0_raw, src0_sel, src0_sext)
       s1_raw = ctx.rsgpr_dyn(vsrc1_reg) if inst.s1 else ctx.rvgpr_dyn(vsrc1_reg, lc)
       s1 = _sdwa_select(s1_raw, src1_sel, src1_sext)
-      srcs = {'S0': s0, 'S1': s1, 'laneId': lc}
-      for dest, val in parse_pcode(pcode, srcs)[1]:
+      for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc})[1]:
         if '[laneId]' in dest and ('D0' in dest or 'EXEC' in dest): return val.cast(dtypes.uint32)
       return _c(0)
-    new_result = ctx.unroll_lanes(get_cmp_bit, exec_mask, apply_exec=False) & exec_mask
+    new_result = ctx.unroll_lanes(get_cmp_bit, exec_mask)
     stores = ctx.wmask(sdst_off, new_result)
     return UOp.sink(*stores)
 
@@ -828,15 +825,10 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc
         old = ctx.rvgpr_dyn(vdst_reg, lane)
         result = _sdwa_write(old, result, dst_sel, dst_unused)
       stores.append(ctx.wvgpr_dyn(vdst_reg, lane, result, exec_mask))
-    elif dest.startswith('VCC'):
-      stores.extend(ctx.wmask_lane_bit(_c(VCC_LO.offset), lane, val, exec_mask))
+    elif dest.startswith('VCC'): vcc_val = val
   if vcc_val is not None:
-    # Initialize sdst to 0 before lane loop (old value may be unrelated data), then set lane bits in loop
-    init_stores = [ctx.wsgpr_dyn(sdst_off, _c(0)), ctx.wsgpr_dyn(sdst_off + _c(1), _c(0))]
-    stores.extend(ctx.wmask_lane_bit(sdst_off, lane, vcc_val, exec_mask))
-    if stores:
-      return UOp.sink(*init_stores, UOp.sink(*stores).end(lane))
-    return UOp.sink(*init_stores)
+    mask = ctx.unroll_lanes(lambda l: (_to_u32(vcc_val.substitute({lane:l})) & _c(1)).cast(dtypes.uint32), exec_mask)
+    return UOp.sink(*ctx.wmask(sdst_off, mask), UOp.sink(*stores).end(lane))
   if stores:
     return UOp.sink(UOp.sink(*stores).end(lane))
   return UOp.sink()
@@ -878,7 +870,6 @@ def _compile_vop12(inst: ir3.VOP1 | ir3.VOP1_SDST | ir3.VOP1_DPP16 | ir3.VOP2 |
   vdst_reg = ctx.inst_field(type(inst).vdst)
   write_hi_half = bits['d'] == 16 and not isinstance(inst, (irc.VOP1, irc.VOP2)) and (vdst_reg >= _c(128))
   if isinstance(write_hi_half, UOp): vdst_reg = write_hi_half.where(vdst_reg - _c(128), vdst_reg)
-  elif write_hi_half: vdst_reg -= 128
   if isinstance(inst, (ir3.VOP1, ir4.VOP1, irc.VOP1)):
     d0 = _cond_hi16(write_hi_half, ctx.rvgpr_dyn(vdst_reg, lane))
     s0, src0_off = _load_vsrc0(ctx, inst, lane, bits, literal, is_f64, is_float, d0)
@@ -909,8 +900,8 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16
   is_dpp16 = hasattr(type(inst), 'dpp') and hasattr(type(inst), 'vsrc0')
 
   # Handle both VOPC (vsrc1) and VOP3 (src1) instruction formats - read operands dynamically
+  src0_off = ctx.inst_field(type(inst).src0)
   if is_vopc:
-    src0_off = ctx.inst_field(type(inst).src0)
     vsrc1_off = ctx.inst_field(type(inst).vsrc1)  # type: ignore[union-attr]
     # For 16-bit ops, vsrc1 >= 128 means hi-half of v[vsrc1-128]
     if bits['s0'] == 16 and not isinstance(inst, irc.VOPC):
@@ -920,15 +911,14 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16
       vsrc1_hi = False
       src1_off = _c(256) + vsrc1_off
   else:
-    src0_off = ctx.inst_field(type(inst).src0)
     src1_off = ctx.inst_field(type(inst).src1)  # type: ignore[union-attr]
     dst_off = ctx.inst_field(type(inst).vdst)  # type: ignore[union-attr]
     vsrc1_hi = False
   literal = ctx.optional_field(inst, 'literal')
 
   is_float, is_f64, pcode = any(x in op_name for x in ('_F32', '_F64', '_F16')), '_F64' in op_name, get_pcode(inst.op)
-  def get_cmp_bit(lane) -> UOp:
-    lc = lane.cast(dtypes.int) if isinstance(lane, UOp) else _c(lane, dtypes.int)
+  def get_cmp_bit(lane: UOp) -> UOp:
+    lc = lane.cast(dtypes.int)
     s0 = _load_dpp16_src0(ctx, inst, lc, _c(0)) if is_dpp16 else ctx.rsrc_dyn(src0_off, lc, bits['s0'], literal, is_f64)
     if is_vopc and not isinstance(inst, irc.VOPC) and bits['s0'] == 16 and not is_dpp16:
       src0_hi = src0_off >= _c(384)
@@ -936,22 +926,17 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16
     s1 = _cond_hi16(vsrc1_hi, ctx.rsrc_dyn(src1_off, lc, bits['s1'], literal, is_f64)) if bits['s0'] == 16 \
       else ctx.rsrc_dyn(src1_off, lc, bits['s1'], literal, is_f64)
     if bits['s0'] == 16 and opsel: s0, s1 = _apply_opsel(s0, 0, opsel), _apply_opsel(s1, 1, opsel)
-    if is_float:
-      if is_dpp16:
-        s0 = _apply_src_mods(s0, 0, 1 if _iattr(inst, 'src0_abs') else 0, 1 if _iattr(inst, 'src0_neg') else 0, bits['s0'])
-        s1 = _apply_src_mods(s1, 0, 1 if _iattr(inst, 'src1_abs') else 0, 1 if _iattr(inst, 'src1_neg') else 0, bits['s1'])
-      s0 = _apply_src_mods(s0, 0, abs_bits, neg_bits, bits['s0'])
-      s1 = _apply_src_mods(s1, 1, abs_bits, neg_bits, bits['s1'])
-    elif abs_bits or neg_bits:  # int compares also honor abs/neg, as bit-level sign clear/flip (not integer abs/negate)
-      s0 = _apply_src_mods(s0, 0, abs_bits, neg_bits, bits['s0'])
-      s1 = _apply_src_mods(s1, 1, abs_bits, neg_bits, bits['s1'])
-    for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(0, dtypes.uint64)})[1]:
+    if is_float and is_dpp16:
+      s0 = _apply_src_mods(s0, 0, 1 if _iattr(inst, 'src0_abs') else 0, 1 if _iattr(inst, 'src0_neg') else 0, bits['s0'])
+      s1 = _apply_src_mods(s1, 0, 1 if _iattr(inst, 'src1_abs') else 0, 1 if _iattr(inst, 'src1_neg') else 0, bits['s1'])
+    s0 = _apply_src_mods(s0, 0, abs_bits, neg_bits, bits['s0'])
+    s1 = _apply_src_mods(s1, 1, abs_bits, neg_bits, bits['s1'])
+    for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': _c(0, dtypes.uint64)})[1]:
       if '[laneId]' in dest and ('D0' in dest or 'EXEC' in dest): return val.cast(dtypes.uint32)
     return _c(0)
 
-  new_bits = ctx.unroll_lanes(get_cmp_bit, exec_mask, apply_exec=False)
   # Both VOPC and VOP3 clear inactive lane bits (hardware verified)
-  new_result = new_bits & exec_mask
+  new_result = ctx.unroll_lanes(get_cmp_bit, exec_mask)
 
   # CMPX writes EXEC only (hardware verified: e64 CMPX does not write SDST); non-CMPX writes SDST/VCC
   if is_cmpx: stores = ctx.wmask(_c(EXEC_LO.offset), new_result)
@@ -972,8 +957,8 @@ def _compile_bitop3(inst, ctx: _Ctx, exec_mask: UOp, bits: dict, op_name: str) -
   is_16 = 'B16' in op_name
   dt, mask = (dtypes.uint16, 0xFFFF) if is_16 else (dtypes.uint32, 0xFFFFFFFF)
   s0, s1, s2 = src0.cast(dt), src1.cast(dt), src2.cast(dt)
-  def bnot(v): return v ^ UOp.const(mask, dt)
-  result = UOp.const(0, dt)
+  def bnot(v): return v ^ _c(mask, dt)
+  result = _c(0, dt)
   for i in range(8):
     if not (ttbl & (1 << i)): continue
     result = result | ((s0 if i & 4 else bnot(s0)) & (s1 if i & 2 else bnot(s1)) & (s2 if i & 1 else bnot(s2)))
@@ -993,7 +978,7 @@ def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3 | irc.VOP3, ctx: _Ctx) -> UOp:
     return ctx.compile_lane_pcode(inst.op, inst)
 
    # VOP3 VOPC (v_cmp_*_e64) - delegate to unified VOPC handler
-  if 'V_CMP' in op_name or 'V_CMPX' in op_name:
+  if 'V_CMP' in op_name:
     return _compile_vopc(inst, ctx, opsel=opsel, abs_bits=_iattr(inst, 'abs'), neg_bits=_iattr(inst, 'neg'))
 
   # BITOP3: abs/neg/omod encode truth table, not source modifiers
@@ -1025,8 +1010,8 @@ def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3 | irc.VOP3, ctx: _Ctx) -> UOp:
   src0 = _apply_src_mods(src0, 0, abs_bits, neg_bits, bits['s0'])
   src1 = _apply_src_mods(src1, 1, abs_bits, neg_bits, bits['s1'])
   src2 = _apply_src_mods(src2, 2, abs_bits, neg_bits, bits['s2'])
-  srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': UOp.const(opsel, dtypes.uint32)}
-  if 'CNDMASK' in op_name and src2 is not None: srcs['VCC'] = src2
+  srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': _c(opsel, dtypes.uint32)}
+  if 'CNDMASK' in op_name: srcs['VCC'] = src2
   # FMAC instructions need D0 (accumulator) from destination register
   if 'FMAC' in op_name: srcs['D0'] = ctx.rvgpr_dyn(vdst_reg, lane)
   opsel_dst_hi = bool(opsel & 0b1000) and bits['d'] == 16
@@ -1127,8 +1112,8 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:
     scale0 = ctx.inst_field(type(inst).scale_src0), _iattr(inst, 'opsel') & 3
     scale1 = ctx.inst_field(type(inst).scale_src1), _iattr(inst, 'opsel_hi') & 3
     def _scale_exp(off_sel: tuple[UOp, int], lane: UOp) -> UOp:
-      byte = (ctx.rsrc_dyn(off_sel[0], lane, 32) >> UOp.const(off_sel[1] * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
-      return byte.cast(dtypes.int32) - UOp.const(127, dtypes.int32)
+      byte = (ctx.rsrc_dyn(off_sel[0], lane, 32) >> _c(off_sel[1] * 8, dtypes.uint32)) & _c(0xFF, dtypes.uint32)
+      return byte.cast(dtypes.int32) - _c(127, dtypes.int32)
     def scale_factor(lane: UOp) -> UOp:  # combined A*B scale for this lane: 2^(ea-127) * 2^(eb-127)
       return UOp.exp2((_scale_exp(scale0, lane) + _scale_exp(scale1, lane)).cast(dtypes.float32))
 
@@ -1163,26 +1148,26 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:
 
   def cvt_elem(raw: UOp, sub_idx: int, fp8_fmt: str) -> UOp:
     if is_i8:  # extract i8, sign-extend to i32
-      byte = (raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
-      return (byte.cast(dtypes.int32) ^ UOp.const(0x80, dtypes.int32)) - UOp.const(0x80, dtypes.int32)
+      byte = (raw >> _c(sub_idx * 8)) & _c(0xFF, dtypes.uint32)
+      return (byte.cast(dtypes.int32) ^ _c(0x80, dtypes.int32)) - _c(0x80, dtypes.int32)
     if is_f32_src: return raw  # already uint32 (f32 bit pattern)
-    if is_fp8: return _FUNCS[f"{fp8_fmt}_to_f32"](raw >> UOp.const(sub_idx * 8, dtypes.uint32)).bitcast(dtypes.uint32)
-    h = (raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)
-    if is_bf16: return h << UOp.const(16, dtypes.uint32)  # bf16 is the upper 16 bits of f32
+    if is_fp8: return _FUNCS[f"{fp8_fmt}_to_f32"](raw >> _c(sub_idx * 8, dtypes.uint32)).bitcast(dtypes.uint32)
+    h = (raw >> _c(sub_idx * 16)) & _c(0xFFFF)
+    if is_bf16: return h << _c(16)  # bf16 is the upper 16 bits of f32
     return _FUNCS['f16_to_f32'](h).bitcast(dtypes.uint32)
 
   def mn_idx(lane: UOp) -> UOp:  # M/N matrix index held by a lane
     if M == 32:  # (lane%32)/16 selects the 16-wide block, (lane%32)%16 the index within it
-      return ((lane // UOp.const(16, dtypes.int)) % UOp.const(2, dtypes.int)) * UOp.const(16, dtypes.int) + lane % UOp.const(16, dtypes.int)
-    return lane % UOp.const(grp_size, dtypes.int)
+      return ((lane // _c(16, dtypes.int)) % _c(2, dtypes.int)) * _c(16, dtypes.int) + lane % _c(16, dtypes.int)
+    return lane % _c(grp_size, dtypes.int)
 
-  def grp_idx(lane: UOp) -> UOp: return lane // UOp.const(grp_size, dtypes.int)
+  def grp_idx(lane: UOp) -> UOp: return lane // _c(grp_size, dtypes.int)
 
   read_lane = ctx.range()
   mn, grp = mn_idx(read_lane), grp_idx(read_lane)
   def mat_idx(base: int, dim: int, kl: int) -> UOp:  # tmp index of element (mn, kl) of this lane's group
-    if M == 4: return UOp.const(base + kl, dtypes.int) + grp * UOp.const(dim * K, dtypes.int) + mn * UOp.const(K, dtypes.int)
-    return UOp.const(base + kl, dtypes.int) + mn * UOp.const(K, dtypes.int) + grp * UOp.const(k_per_grp, dtypes.int)
+    if M == 4: return _c(base + kl, dtypes.int) + grp * _c(dim * K, dtypes.int) + mn * _c(K, dtypes.int)
+    return _c(base + kl, dtypes.int) + mn * _c(K, dtypes.int) + grp * _c(k_per_grp, dtypes.int)
 
   read_stores = []
   for kl in range(k_per_grp):
@@ -1203,7 +1188,7 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:
   def _dot_accum(acc: UOp, a_row: UOp, b_row: UOp, lane: UOp) -> UOp:
     """acc += sum_k A[a_row+k] * B[b_row+k] in order (FP-associativity matters). For scaled MFMA only the dot is scaled: D = dot*scale + C."""
     def prod(k: int) -> UOp:
-      return tmp2.index(a_row + UOp.const(k, dtypes.int)).bitcast(acc_dt) * tmp2.index(b_row + UOp.const(k, dtypes.int)).bitcast(acc_dt)
+      return tmp2.index(a_row + _c(k, dtypes.int)).bitcast(acc_dt) * tmp2.index(b_row + _c(k, dtypes.int)).bitcast(acc_dt)
     if not scaled:
       for k in range(K): acc = acc + prod(k)
       return acc
@@ -1213,16 +1198,16 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:
 
   compute_lane = ctx.range()
   c_mn, c_grp = mn_idx(compute_lane), grp_idx(compute_lane)
-  b_off = UOp.const(n_a_elems, dtypes.int)
+  b_off = _c(n_a_elems, dtypes.int)
   def out_ab(out_reg: int) -> tuple[UOp, UOp]:  # A/B tmp base indices for one output element of this lane
     if M == 32:  # 16 outputs per lane: rows (lane//32)*16 + blocks of 4 within the wave
-      m_base = c_grp * UOp.const(16, dtypes.int) + UOp.const((out_reg // 4) * 4 + out_reg % 4, dtypes.int)
-      return m_base * UOp.const(K, dtypes.int), b_off + c_mn * UOp.const(K, dtypes.int)
+      m_base = c_grp * _c(16, dtypes.int) + _c(out_reg, dtypes.int)
+      return m_base * _c(K, dtypes.int), b_off + c_mn * _c(K, dtypes.int)
     if M == 4:  # each group of 4 lanes computes an independent 4x4 block
-      a_base = c_grp * UOp.const(M * K, dtypes.int) + UOp.const(out_reg * K, dtypes.int)
-      return a_base, b_off + (c_grp * UOp.const(N, dtypes.int) + c_mn) * UOp.const(K, dtypes.int)
-    m_base = c_grp * UOp.const(out_per_lane, dtypes.int) + UOp.const(out_reg, dtypes.int)  # 16x16: 4 outputs, one row each
-    return m_base * UOp.const(K, dtypes.int), b_off + c_mn * UOp.const(K, dtypes.int)
+      a_base = c_grp * _c(M * K, dtypes.int) + _c(out_reg * K, dtypes.int)
+      return a_base, b_off + (c_grp * _c(N, dtypes.int) + c_mn) * _c(K, dtypes.int)
+    m_base = c_grp * _c(out_per_lane, dtypes.int) + _c(out_reg, dtypes.int)  # 16x16: 4 outputs, one row each
+    return m_base * _c(K, dtypes.int), b_off + c_mn * _c(K, dtypes.int)
 
   compute_stores = []
   for out_reg in range(out_per_lane):
@@ -1250,9 +1235,8 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
 
   # read a source element from VGPRs: (src, lane, vgpr, element-in-vgpr) -> f32/i32
   def gval(src, lane, vgpr, ridx, *, fmt):
-    v = ctx.rvgpr_dyn(src + _c(vgpr), UOp.const(lane, dtypes.int))
-    pkd = v >> UOp.const(ridx * sz, dtypes.uint32)
-    pkd = pkd & UOp.const((1 << sz) - 1, dtypes.uint32)
+    v = ctx.rvgpr_dyn(src + _c(vgpr), _c(lane, dtypes.int))
+    pkd = (v >> _c(ridx * sz)) & _c((1 << sz) - 1)
     if not fmt.startswith('IU'): return _FUNCS[f'{fmt.lower()}_to_f32'](pkd)
     return (pkd << _c(24, dtypes.uint)).bitcast(dtypes.int32) >> _c(24, dtypes.int32)  # sign extend
 
@@ -1274,8 +1258,7 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
   # src2 may be a VGPR or an inline/scalar constant (128 = int 0, the usual ", 0" C form); the runner must handle both dynamically
   out_dt = dtypes.int32 if output_type == "I32" else dtypes.float32
   cbits = ctx.rsrc_dyn(src2_r, None, 32)
-  cval_const = _FUNCS[f'{output_type.lower()}_to_f32'](cbits & UOp.const(0xFFFF, dtypes.uint32)) \
-    if output_type in ("F16", "BF16") else cbits.bitcast(out_dt)
+  cval_const = _FUNCS[f'{output_type.lower()}_to_f32'](cbits & _c(0xFFFF)) if output_type in ("F16", "BF16") else cbits.bitcast(out_dt)
   elem = ctx.range(256)
   row, col = elem // 16, elem % 16
   lane, vgpr = d_map(row, col)
@@ -1291,13 +1274,13 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
 
   def w_store(m: int, n: UOp, val: UOp, vgpr_off: int) -> UOp:  # store one output element to its (lane, vgpr) slot
     lane_i, _ = d_map(m, n)
-    return ctx.wvgpr_dyn(vdst_reg + _c(vgpr_off), UOp.const(lane_i, dtypes.int), val, exec_mask)
+    return ctx.wvgpr_dyn(vdst_reg + _c(vgpr_off), _c(lane_i, dtypes.int), val, exec_mask)
   if output_type in ("F16", "BF16"):
     def to_bits(v: UOp) -> UOp:  # f32 result -> 16 output bits
-      return ((v.bitcast(dtypes.uint32) >> UOp.const(16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)) if output_type == "BF16" \
+      return ((v.bitcast(dtypes.uint32) >> _c(16, dtypes.uint32)) & _c(0xFFFF)) if output_type == "BF16" \
         else v.cast(dtypes.half).bitcast(dtypes.uint16).cast(dtypes.uint32)
     if is_rdna4:  # pack 2 outputs per VGPR (adjacent m values share a VGPR)
-      stores = [w_store(m, n, to_bits(mat_d[m]) | (to_bits(mat_d[m+1]) << UOp.const(16, dtypes.uint32)), d_map(m, n)[1] // 2)
+      stores = [w_store(m, n, to_bits(mat_d[m]) | (to_bits(mat_d[m+1]) << _c(16, dtypes.uint32)), d_map(m, n)[1] // 2)
                 for m in range(0, 16, 2)]
     else:  # one output per VGPR (lo half)
       stores = [w_store(m, n, to_bits(mat_d[m]), d_map(m, n)[1]) for m in range(16)]
@@ -1323,11 +1306,11 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _
     exec_mask = ctx.rexec()
     vdst_reg = ctx.inst_field(type(inst).vdst)
     src0_off = ctx.inst_field(type(inst).src0)
-    if src0_is_acc and not vdst_is_acc:
+    if not vdst_is_acc:
       # v_accvgpr_read: VGPR[vdst] = ACCVGPR[src0]
       val = ctx.raccvgpr_dyn(src0_off - _c(256), lane)
       return UOp.sink(ctx.wvgpr_dyn(vdst_reg, lane, val, exec_mask).end(lane))
-    elif vdst_is_acc and not src0_is_acc:
+    elif not src0_is_acc:
       # v_accvgpr_write: ACCVGPR[vdst] = src0 (src0 can be VGPR or SGPR/const)
       src0 = ctx.rsrc_dyn(src0_off, lane, 32)
       return UOp.sink(ctx.waccvgpr_dyn(vdst_reg, lane, src0, exec_mask).end(lane))
@@ -1362,8 +1345,7 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _
     # v_pk_mov_b32: D[lo] = src0[opsel_bit0 ? hi : lo], D[hi] = src1[opsel_bit1 ? hi : lo]
     lo_val = _pk_sel(src0, ctx.inst_field(type(inst).src0), opsel & 1)
     hi_val = _pk_sel(src1, ctx.inst_field(type(inst).src1), opsel & 2)
-    lo_out, hi_out = _split64(_u64(lo_val, hi_val))
-    stores = [ctx.wvgpr_dyn(vdst_reg, lane, lo_out, exec_mask), ctx.wvgpr_dyn(vdst_reg + _c(1), lane, hi_out, exec_mask)]
+    stores = [ctx.wvgpr_dyn(vdst_reg, lane, lo_val, exec_mask), ctx.wvgpr_dyn(vdst_reg + _c(1), lane, hi_val, exec_mask)]
     return UOp.sink(UOp.group(*stores).end(lane))
 
   srcs: dict[str, UOp | int] = {}
@@ -1373,8 +1355,8 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _
     hi_bits = (opsel_hi & 1, opsel_hi & 2, 1 if opsel_hi2 else 0)
     for i, (bit, s0) in enumerate(zip((1, 2, 4), (src0, src1, src2))):
       lo, hi = _pk_sel(s0, src_offs[i], opsel & bit), _pk_sel(s0, src_offs[i], hi_bits[i])
-      if neg & bit: lo = lo ^ UOp.const(0x80000000, dtypes.uint32)
-      if neg_hi & bit: hi = hi ^ UOp.const(0x80000000, dtypes.uint32)
+      if neg & bit: lo = lo ^ _c(0x80000000, dtypes.uint32)
+      if neg_hi & bit: hi = hi ^ _c(0x80000000, dtypes.uint32)
       srcs[f'S{i}'] = _u64(lo, hi)
   elif 'FMA_MIX' in op_name or 'MAD_MIX' in op_name:
     combined_opsel_hi = (opsel_hi & 0x3) | ((opsel_hi2 & 0x1) << 2)
@@ -1382,35 +1364,35 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _
     def apply_abs(v, bit, opsel_hi_bit, opsel_bit):
       if not (neg_hi & bit): return v
       # Apply abs based on whether source is f32 or f16
-      if not (combined_opsel_hi & opsel_hi_bit): return v & UOp.const(0x7FFFFFFF, dtypes.uint32)  # f32 abs
-      if opsel & opsel_bit: return v & UOp.const(0x7FFF0000, dtypes.uint32)  # f16 hi abs (preserve lo)
-      return v & UOp.const(0xFFFF7FFF, dtypes.uint32)  # f16 lo abs (preserve hi)
+      if not (combined_opsel_hi & opsel_hi_bit): return v & _c(0x7FFFFFFF, dtypes.uint32)  # f32 abs
+      if opsel & opsel_bit: return v & _c(0x7FFF0000, dtypes.uint32)  # f16 hi abs (preserve lo)
+      return v & _c(0xFFFF7FFF, dtypes.uint32)  # f16 lo abs (preserve hi)
     def apply_neg_mix(v, bit, opsel_hi_bit, opsel_bit):
       if not (neg & bit): return v
-      if not (combined_opsel_hi & opsel_hi_bit): return v ^ UOp.const(0x80000000, dtypes.uint32)  # f32 neg
-      if opsel & opsel_bit: return v ^ UOp.const(0x80000000, dtypes.uint32)  # f16 hi neg
-      return v ^ UOp.const(0x00008000, dtypes.uint32)  # f16 lo neg
+      if not (combined_opsel_hi & opsel_hi_bit): return v ^ _c(0x80000000, dtypes.uint32)  # f32 neg
+      if opsel & opsel_bit: return v ^ _c(0x80000000, dtypes.uint32)  # f16 hi neg
+      return v ^ _c(0x00008000, dtypes.uint32)  # f16 lo neg
     s0_mod = apply_neg_mix(apply_abs(src0, 1, 1, 1), 1, 1, 1)
     s1_mod = apply_neg_mix(apply_abs(src1, 2, 2, 2), 2, 2, 2)
     s2_mod = apply_neg_mix(apply_abs(src2, 4, 4, 4), 4, 4, 4)
     srcs = {'S@0': s0_mod, 'S@1': s1_mod, 'S@2': s2_mod,
-            'OPSEL_HI': UOp.const(combined_opsel_hi, dtypes.uint32), 'OPSEL': UOp.const(opsel, dtypes.uint32)}
+            'OPSEL_HI': _c(combined_opsel_hi, dtypes.uint32), 'OPSEL': _c(opsel, dtypes.uint32)}
   else:
     def get_half_bits(val: UOp, use_hi: bool, apply_neg: bool = False) -> UOp:
-      bits = ((val >> UOp.const(16, dtypes.uint32)) if use_hi else val) & UOp.const(0xFFFF, dtypes.uint32)
+      bits = ((val >> _c(16, dtypes.uint32)) if use_hi else val) & _c(0xFFFF)
       if apply_neg: bits = bits.cast(dtypes.uint16).bitcast(dtypes.half).neg().bitcast(dtypes.uint16).cast(dtypes.uint32)
       return bits
     def build_remapped_src(src: UOp, opsel_lo_bit: int, opsel_hi_bit: int, neg_lo_bit: int, neg_hi_bit: int) -> UOp:
       lo = get_half_bits(src, bool(opsel_lo_bit), bool(neg_lo_bit))
       hi = get_half_bits(src, bool(opsel_hi_bit), bool(neg_hi_bit))
-      return lo | (hi << UOp.const(16, dtypes.uint32))
+      return lo | (hi << _c(16, dtypes.uint32))
     # DOT IU instructions use NEG bits for signed/unsigned selection, not fp16 negation
     is_dot_iu = 'DOT' in op_name and 'IU' in op_name
     n0, n1, n2, nh0, nh1, nh2 = (0, 0, 0, 0, 0, 0) if is_dot_iu else (neg & 1, neg & 2, neg & 4, neg_hi & 1, neg_hi & 2, neg_hi & 4)
     srcs = {'S0': build_remapped_src(src0, opsel & 1, opsel_hi & 1, n0, nh0),
             'S1': build_remapped_src(src1, opsel & 2, opsel_hi & 2, n1, nh1),
             'S2': build_remapped_src(src2, opsel & 4, 1 if opsel_hi2 else 0, n2, nh2)}
-    if is_dot_iu: srcs['NEG'] = UOp.const(neg, dtypes.uint32)
+    if is_dot_iu: srcs['NEG'] = _c(neg, dtypes.uint32)
   return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask)
 
 def _compile_vopd(inst: ir3.VOPD | ir4.VOPD, ctx: _Ctx) -> UOp:
@@ -1441,7 +1423,6 @@ def _compile_vopd(inst: ir3.VOPD | ir4.VOPD, ctx: _Ctx) -> UOp:
     if vop in (ir3.VOP2Op.V_FMAAK_F32_E32, ir3.VOP2Op.V_FMAMK_F32_E32):
       assert literal is not None
       srcs['SIMM32'] = literal
-    if op in (ir3.VOPDOp.V_DUAL_CNDMASK_B32, ir4.VOPDOp.V_DUAL_CNDMASK_B32): srcs['VCC'] = ctx.rmask(_c(VCC_LO.offset))
     pcode = get_pcode(vop)
     srcs.update(ctx.base_srcs(exec_mask, lane), VCC=ctx.rmask(_c(VCC_LO.offset)))
     for dest, val in parse_pcode(pcode, srcs)[1]:
@@ -1463,7 +1444,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
   # CDNA acc bit: when set, VGPR operands (vdst/vdata) target ACCVGPR file instead of VGPR
   use_acc = bool(_iattr(inst, 'acc'))
   mem = ctx.lds if is_lds else ctx.scratch if is_scratch else ctx.vmem
-  addr_shift = UOp.const(2, dtypes.uint32 if is_lds else dtypes.uint64)
+  addr_shift = _c(2, dtypes.uint32 if is_lds else dtypes.uint64)
 
   # Field names differ per format: DS (addr/data0/offset0+offset1), RDNA4 V* (vaddr/vsrc/ioffset), RDNA3+CDNA (addr/data/offset)
   if is_lds: addr_field, data_field = 'addr', 'data0'
@@ -1483,7 +1464,6 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
   # Data width from canonical_op_bits (32/64/96/128), default to 32 for untyped ops
   data_bits_mem = inst.canonical_op_bits.get('data', 32)
   is_atomic, glc = 'ATOMIC' in op_name, _iattr(inst, 'glc')
-  has_data1 = is_lds and hasattr(inst, 'data1') and inst.data1 is not None
   data1_reg = ctx.inst_field(type(inst).data1) if is_lds else _c(0)  # type: ignore[union-attr]
 
   if is_lds and op_name == 'DS_SWIZZLE_B32':
@@ -1516,18 +1496,17 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
     # Dynamic saddr check: saddr < 124 means valid SGPR, otherwise use VGPR pair for address
     use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False)
     if is_scratch:
-      scratch_stride = ctx.rsgpr_dyn(_c(SCRATCH_STRIDE_IDX)).cast(dtypes.uint64)
-      base = lane.cast(dtypes.uint64) * scratch_stride
+      base = lane.cast(dtypes.uint64) * ctx.scratch_stride.cast(dtypes.uint64)
       # SVE (Scratch VGPR Enable): when SVE=1, VADDR is used as offset; when SVE=0, VADDR is ignored
       sve = _iattr(inst, 'sve')
       vaddr = ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64)
-      addr_offset = vaddr if sve == 1 else UOp.const(0, dtypes.uint64)
+      addr_offset = vaddr if sve == 1 else _c(0, dtypes.uint64)
       # Add saddr value only if use_saddr is true (saddr < 124)
-      saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), UOp.const(0, dtypes.uint64)) \
-        if saddr_reg is not None else UOp.const(0, dtypes.uint64)
+      saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), _c(0, dtypes.uint64)) \
+        if saddr_reg is not None else _c(0, dtypes.uint64)
       return base + addr_offset + saddr_contrib + offset64
     # GLOBAL can use an SGPR base plus a VGPR offset; FLAT always uses the full VGPR pair.
-    saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64)
+    saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else _c(0, dtypes.uint64)
     vaddr_base = _u64(ctx.rvgpr_dyn(addr_reg, lane), ctx.rvgpr_dyn(addr_reg + _c(1), lane))
     # When saddr is valid: base = saddr pair, vaddr is 32-bit offset; otherwise: base = 0, vaddr is 64-bit address
     base_addr = use_saddr.where(saddr_base + ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64), vaddr_base)
@@ -1553,10 +1532,10 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
     _rvdata = ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn
     if is_lds:
       if data_bits_mem <= 32:
-        data = {'DATA': _rvdata(vdata_reg, lane), 'DATA2': _rvdata(data1_reg, lane) if has_data1 else UOp.const(0, dtypes.uint32)}
+        data = {'DATA': _rvdata(vdata_reg, lane), 'DATA2': _rvdata(data1_reg, lane)}
       elif data_bits_mem == 64:  # DATA/DATA2 are the 64-bit input registers, formed from register pairs
         data = {'DATA': _u64(_rvdata(vdata_reg, lane), _rvdata(vdata_reg + _c(1), lane)),
-                'DATA2': _u64(_rvdata(data1_reg, lane), _rvdata(data1_reg + _c(1), lane)) if has_data1 else UOp.const(0, dtypes.uint64)}
+                'DATA2': _u64(_rvdata(data1_reg, lane), _rvdata(data1_reg + _c(1), lane))}
       else:  # 96/128-bit: one register per dword
         data = {'DATA': _rvdata(vdata_reg, lane), **{f'DATA{i}': _rvdata(vdata_reg + _c(i), lane) for i in range(1, data_bits_mem // 32)}}
       if 'D16' in op_name: data['RETURN_DATA'] = _rvdata(vdst_reg, lane)
@@ -1567,8 +1546,8 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
     active = _lane_active(exec_mask, lane)
     # saddr < 124 means valid SGPR pair, otherwise use 0 (NULL means no saddr contribution)
     use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False)
-    saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64)
-    saddr_base = use_saddr.where(saddr_raw, UOp.const(0, dtypes.uint64))
+    saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else _c(0, dtypes.uint64)
+    saddr_base = use_saddr.where(saddr_raw, _c(0, dtypes.uint64))
     # Sign-extend offset to 64-bit for the final address calculation
     ioffset64 = offset.cast(dtypes.int64).cast(dtypes.uint64)
     # v_addr for CalcGlobalAddr: when saddr valid, use low 32 bits as offset; otherwise full 64-bit address. Include ioffset.
@@ -1580,13 +1559,13 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
         if data_bits_mem == 64 else ctx.rvgpr_dyn(vdata_reg, lane)
       return {'ADDR': addr, 'DATA': atomic_data, '_vmem': mem, '_active': active, 'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base}
     vdata = _rvdata(vdata_reg, lane).cast(dtypes.uint64) if 'STORE' in op_name \
-      else _rvdata(vdst_reg, lane) if 'D16' in op_name else UOp.const(0, dtypes.uint32)
+      else _rvdata(vdst_reg, lane) if 'D16' in op_name else _c(0, dtypes.uint32)
     if 'STORE' in op_name and data_bits_mem >= 64:
-      vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
+      vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << _c(32, dtypes.uint64))
     srcs = {'ADDR': addr, 'VDATA': vdata, '_vmem': mem, '_active': active,
             'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base, 'SADDR': saddr_base, 'OFFSET': offset}
     for i in range(data_bits_mem // 32):
-      srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(0, dtypes.uint32)
+      srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else _c(0, dtypes.uint32)
     return srcs
 
   def make_stores(dest: str, val: UOp, lane: UOp, active: UOp, writes_return_data: bool) -> list[UOp]:
@@ -1664,7 +1643,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
   offset, offen, idxen = ctx.inst_field(type(inst).offset), ctx.inst_field(type(inst).offen), ctx.inst_field(type(inst).idxen)
 
   # V# descriptor: base[0:1], num_records[2], stride=word3[13:0]
-  base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & UOp.const(0xFFFFFFFFFFFF, dtypes.uint64)
+  base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & _c(0xFFFFFFFFFFFF, dtypes.uint64)
   num_records = ctx.rsgpr_dyn(srsrc + _c(2))
   stride = (ctx.rsgpr_dyn(srsrc + _c(3)) & _c(0x3FFF)).cast(dtypes.uint64)
 
@@ -1681,7 +1660,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
   buffer_offset = (stride * index + voff + offset.cast(dtypes.uint64)).cast(dtypes.uint32)
   in_bounds = active & buffer_offset.__lt__(num_records)
   addr = base + soff + buffer_offset.cast(dtypes.uint64)
-  addr = in_bounds.where(addr, UOp.const(0, dtypes.uint64))  # safe address when OOB
+  addr = in_bounds.where(addr, _c(0, dtypes.uint64))  # safe address when OOB
   mem = ctx.vmem
 
   stores: list[UOp] = []
@@ -1690,20 +1669,20 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
     lds_base = ctx.rsgpr_dyn(_c(124)) & _c(0x3FFFF)
     lds_addr = lds_base + lane.cast(dtypes.uint32) * _c(n_dwords * 4)
     for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + _c(i * 4, dtypes.uint64)) >> _c(2, dtypes.uint64)
       val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64)).load(), _c(0))
       lds_idx = (lds_addr + _c(i * 4)) >> _c(2)
       lds_slot = ctx.lds.index(lds_idx.valid(active))
       stores.append(lds_slot.store(active.where(val, lds_slot)))
   elif is_store:
     for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + _c(i * 4, dtypes.uint64)) >> _c(2, dtypes.uint64)
       idx = mem.index(word_addr.cast(dtypes.int64).valid(in_bounds))
       val = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(vdata + _c(i), lane)
       stores.append(idx.store(in_bounds.where(_to_u32(val), idx)))
   else:
     for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + _c(i * 4, dtypes.uint64)) >> _c(2, dtypes.uint64)
       val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64).valid(in_bounds)).load(), _c(0))
       stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(vdata + _c(i), lane, val, exec_mask))
   return UOp.sink(UOp.group(*stores).end(lane))
@@ -1741,49 +1720,19 @@ def _get_handler(inst: Inst) -> Callable[..., UOp]:
 # PROGRAM DECODE AND COMPILATION
 # ═══════════════════════════════════════════════════════════════════════════════
 
-_BARRIER_OPS = {ir3.SOPPOp.S_BARRIER, irc.SOPPOp.S_BARRIER}
-if hasattr(ir4.SOPPOp, 'S_BARRIER_WAIT'): _BARRIER_OPS.add(ir4.SOPPOp.S_BARRIER_WAIT)
-_BARRIER_SOP1_OPS: set = set()
-if hasattr(ir4.SOP1Op, 'S_BARRIER_SIGNAL'): _BARRIER_SOP1_OPS.add(ir4.SOP1Op.S_BARRIER_SIGNAL)
+_BARRIER_OPS = {ir3.SOPPOp.S_BARRIER, irc.SOPPOp.S_BARRIER, ir4.SOPPOp.S_BARRIER_WAIT}
+_BARRIER_SOP1_OPS = {ir4.SOP1Op.S_BARRIER_SIGNAL}
 
 # Inline float constants (as bit patterns) for GPU instructions
 F32_INLINE = {240: 0x3f000000, 241: 0xbf000000, 242: 0x3f800000, 243: 0xbf800000,  # 0.5, -0.5, 1.0, -1.0
               244: 0x40000000, 245: 0xc0000000, 246: 0x40800000, 247: 0xc0800000, 248: 0x3e22f983}  # 2.0, -2.0, 4.0, -4.0, 1/(2*pi)
 
-instruction_cache:dict[tuple, list[tuple[int, int, tuple[UOp|None, UOp|None]]]] = {}
-
-# this is meant to replace the old emulator
 @uopfunc
-def init_wave(wg:UOp, wave:UOp, sgpr:UOp, vgpr:UOp, args_ptr:UOp, gx:int, gy:int, lx:int, ly:int, total_threads:int, wave_size:int,
-              scratch_size:int, rsrc2:int, arch:str="rdna3", user_data_count:int=0):
+def store_consts(sgpr:UOp, wave_size:int):
   ci = UOp.range(81, 2, AxisType.LOOP, dtype=dtypes.int)
   init_ints = sgpr.index(ci+128).store((ci<65).where(ci, 64-ci).cast(dtypes.uint32)).end(ci)
-  fi = UOp.range(len(F32_INLINE), 3, AxisType.LOOP, dtype=dtypes.int)
-  floats_buf = UOp(Ops.BINARY, arg=struct.pack(f"<{len(F32_INLINE)}I", *F32_INLINE.values())).bitcast(dtypes.uint32)
-  init_floats = sgpr.index(fi+240).store(floats_buf[fi]).end(fi)
-  # set RANGE registers
-  gidx, gidy, gidz = wg%gx, (wg//gx)%gy, wg//(gx*gy)
-  n_lanes = (total_threads-wave*wave_size).minimum(wave_size)
-  initial:list[tuple[int, sint]] = [(124, 0)] if wave_size == 32 else []
-  if user_data_count: initial += [(i, UOp.variable(f"user_data_{i}", 0, dtypes.uint32.max, dtypes.uint32)) for i in range(user_data_count)]
-  else: initial += [(0, args_ptr.cast(dtypes.uint32)), (1, (args_ptr>>32).cast(dtypes.uint32))]
-  if arch == "rdna4": initial += [(ttmp[7].offset, (gidy&0xFFFF)|((gidz&0xFFFF)<<16)), (ttmp[9].offset, gidx)]
-  else:
-    sgpr_id = (rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT_SHIFT
-    for enabled, gid in [(hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X, gidx),
-                         (hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Y, gidy),
-                         (hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Z, gidz)]:
-      if rsrc2 & enabled:
-        initial.append((sgpr_id, gid))
-        sgpr_id += 1
-  initial += [(EXEC_LO.offset, ((UOp.const(1, dtypes.uint64)<<n_lanes.minimum(32).cast(dtypes.uint64))-1).cast(dtypes.uint32)),
-              (SCRATCH_STRIDE_IDX, scratch_size), (SGPR_COUNT-16+4, (wave&15)|((wave&3)<<4))]
-  if wave_size == 64: initial.append((EXEC_LO.offset+1,
-                                     ((UOp.const(1, dtypes.uint64)<<(n_lanes-32).maximum(0).cast(dtypes.uint64))-1).cast(dtypes.uint32)))
-  lane = UOp.range(wave_size, 5, AxisType.LOOP, dtype=dtypes.int)
-  tid = wave*wave_size+lane
-  init_vgpr = vgpr.index(lane.valid(tid<total_threads)).store((((tid//(lx*ly))<<20)|(((tid//lx)%ly)<<10)|(tid%lx)).cast(dtypes.uint32)).end(lane)
-  return UOp.sink(init_vgpr, *(sgpr.after(init_ints, init_floats).index(i).store(UOp.const(v, dtypes.uint32)) for i,v in dict(initial).items()))
+  init_floats = UOp.group(*[sgpr.index(i).store(_c(v)) for i,v in F32_INLINE.items()])
+  return UOp.sink(init_ints, init_floats, *([sgpr.index(124).store(_c(0))] if wave_size == 32 else []))
 
 @uopfunc
 def loop_active(routes:UOp, active:UOp, header:UOp, waves:UOp):
@@ -1795,6 +1744,8 @@ def lower_body(sink:UOp, name:str) -> UOp:
   body = sink.substitute({b:b.param_like(i, name=b.arg.name) for i,b in enumerate(bufs)}, walk=True)
   return body.call(*bufs, name=name)
 
+instruction_cache:dict[tuple, list[tuple[int, int, tuple[UOp|None, UOp|None, dict[UOp, UOp]]]]] = {}
+
 def instruction_call(inst:Inst, arch:str, renderer:Renderer) -> tuple[UOp|None, UOp|None]:
   entries = instruction_cache.setdefault((arch, type(inst), inst.size(), type(renderer), renderer.target, *(x.value for x in to_program_config)), [])
   inst_bytes = inst.to_bytes()
@@ -1806,26 +1757,32 @@ def instruction_call(inst:Inst, arch:str, renderer:Renderer) -> tuple[UOp|None,
     sink = _get_handler(inst)(inst, ctx)
     base, mask = ctx.canonical_mask(inst_bytes)
     name = f"{_op_name(inst).lower()}_{base.to_bytes(len(inst_bytes), 'little').hex()}"
-    ret = lower_body(sink, name) if sink.src else None, ctx.branch_cond
+    call = lower_body(sink, name) if sink.src else None
+    used = UOp.sink(call, ctx.branch_cond).toposort()
+    ret = call, ctx.branch_cond, {p:index for index,p in ctx.sgpr_reads.items() if p in used}
     entries.append((base, mask, ret))
-  call, cond = ret
-  words = {_Ctx.inst_word(i):_c((bits >> (32*i)) & 0xffffffff) for i in range(inst.size()//4)}
-  if call is not None: call = call.replace(src=(call.body, *(words.get(b, b) for b in call.src[1:])))
-  return call, cond.substitute(words) if cond is not None else None
+  call, cond, reads = ret
+  words = {f"inst_w{i}":(bits >> (32*i)) & 0xffffffff for i in range(inst.size()//4)}
+  bindings = {_Ctx.inst_word(i):_c(words[f"inst_w{i}"]) for i in range(inst.size()//4)}
+  bindings.update({p:_Ctx.sgpr.index(_c(index.get_idx().sym_infer(words))).load() if index.get_valid().sym_infer(words) else _c(0)
+                   for p,index in reads.items()})
+  if call is not None: call = call.replace(src=(call.body, *(bindings.get(b, b) for b in call.src[1:])))
+  return call, cond.substitute(bindings) if cond is not None else None
 
 def instruction_block(instructions:list[UOp]) -> tuple[UOp, list[UOp]]:
-  bufs = sorted({b for call in instructions for b in call.src[1:] if b.op is Ops.PARAM}, key=lambda b:b.arg.slot)
+  bufs = sorted({buf for call in instructions for b in call.src[1:] if (buf:=b.buf_uop).op is Ops.PARAM}, key=lambda b: b.arg.slot)
   params = {b:b.param_like(i, name=b.arg.name) for i,b in enumerate(bufs)}
   linear = list(UOp.sink(*params.values()).toposort())[:-1]
-  seen = set(linear)
+  seen: set[UOp] = set()
   for call in instructions:
     args = [params.get(b, b) for b in call.src[1:]]
     inputs = UOp.sink(*args)
-    nodes = inputs.toposort(lambda u:u not in seen)
+    nodes = inputs.toposort(lambda u:u.op is not Ops.PARAM and u not in seen)
     nodes.pop(inputs)
     seen.update(nodes)
     linear.extend(nodes)
     linear.append(call.replace(src=(call.body, *args)))
+    if _Ctx.sgpr in call.src[1:]: seen.clear()
   return UOp(Ops.PROGRAM, src=(sink:=UOp.sink(*linear, arg=KernelInfo("asm_block", optimized=True)), UOp(Ops.LINEAR, src=(*linear, sink)))), bufs
 
 def thread_branches(insts, blocks:dict[int, list[int]], paths:dict[int, dict[int, int]]):
@@ -1923,8 +1880,8 @@ def lift(code:bytes, arch:str, renderer:Renderer, dispatch:tuple, name:str="") -
       possible = (possible-{pc}) | (set(cfg["paths"][pc]) or {-1})
   routed = len(unconditional) != len(regions)
   axes = itertools.count(sum(inst.size() for inst in insts.values()))
-  gx, gy, gz, lx, ly, lz, rsrc2, scratch_size, user_data_count = dispatch
-  # construct CALL graph
+  gx, gy, gz, lx, ly, lz, rsrc2, scratch_size = dispatch
+  # init ranges
   wave_size, total_threads = _wave_size(arch), lx*ly*lz
   n_waves = (total_threads+wave_size-1)//wave_size
   wg = UOp.range(gx*gy*gz, 0, AxisType.LOOP, dtype=dtypes.uint64)
@@ -1946,11 +1903,30 @@ def lift(code:bytes, arch:str, renderer:Renderer, dispatch:tuple, name:str="") -
   route = UOp.param(6, dtypes.int, (1,), name="route")
   taken, fallthrough = [UOp.param(i, dtypes.int, name=name, addrspace=AddrSpace.ALU) for i,name in enumerate(("taken", "fallthrough"), 7)]
   buffers[route] = (routes, 1)
-  init = init_wave(wg, wave, sgpr.shrink(((wave*SGPR_COUNT, (wave+1)*SGPR_COUNT),)), vgpr.shrink(((wave*256*wave_size, (wave+1)*256*wave_size),)),
-                   args_ptr, gx, gy, lx, ly, total_threads, wave_size, scratch_size, rsrc2, arch, user_data_count)
-  user_vars = [u for u in init.body.toposort() if u.is_variable]
-  init = init.replace(src=(init.body.substitute({u:UOp.param(len(init.src)-1+i, u.dtype, name=u.arg.name, addrspace=AddrSpace.ALU)
-                                                 for i,u in enumerate(user_vars)}, walk=True), *init.src[1:], *user_vars))
+  # set per wave registers
+  wave_sgpr = sgpr.shrink(((wave*SGPR_COUNT, (wave+1)*SGPR_COUNT),)).after(wg)
+  wave_vgpr = vgpr.shrink(((wave*256*wave_size, (wave+1)*256*wave_size),)).after(wg)
+  consts = store_consts(wave_sgpr, wave_size)
+  gidx, gidy, gidz = wg%gx, (wg//gx)%gy, wg//(gx*gy)
+  n_lanes = (total_threads-wave*wave_size).minimum(wave_size)
+  runtime_sgprs = [(0, args_ptr.cast(dtypes.uint32)), (1, (args_ptr>>32).cast(dtypes.uint32))]
+  if arch == "rdna4": runtime_sgprs += [(ttmp[7].offset, (gidy&0xFFFF)|((gidz&0xFFFF)<<16)), (ttmp[9].offset, gidx)]
+  else:
+    sgpr_id = (rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT_SHIFT
+    for enabled, gid in [(hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X, gidx),
+                         (hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Y, gidy),
+                         (hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Z, gidz)]:
+      if rsrc2 & enabled:
+        runtime_sgprs.append((sgpr_id, gid))
+        sgpr_id += 1
+  runtime_sgprs += [(EXEC_LO.offset, ((_c(1, dtypes.uint64)<<n_lanes.minimum(32).cast(dtypes.uint64))-1).cast(dtypes.uint32))]
+  if wave_size == 64:
+    runtime_sgprs.append((EXEC_LO.offset+1, ((_c(1, dtypes.uint64)<<(n_lanes-32).maximum(0).cast(dtypes.uint64))-1).cast(dtypes.uint32)))
+  lane = UOp.range(wave_size, 5, AxisType.LOOP, dtype=dtypes.int)
+  tid = wave*wave_size+lane
+  init_vgpr = wave_vgpr.index(lane.valid(tid<total_threads)).store((((tid//(lx*ly))<<20)|(((tid//lx)%ly)<<10)|(tid%lx)).cast(dtypes.uint32)).end(lane)
+  init = UOp.group(init_vgpr, *(wave_sgpr.after(consts).index(i).store(_c(v)) for i,v in dict(runtime_sgprs).items()))
+  # construct CALL graph
   initialized = (routes.after(init).index(wave).store(0) if routed else UOp.group(init, wave)).end(wave)
   def emit_block(block_pc:int, end:UOp, scopes:tuple[UOp, ...]) -> UOp:
     cond = UOp.const(True)
@@ -1962,19 +1938,19 @@ def lift(code:bytes, arch:str, renderer:Renderer, dispatch:tuple, name:str="") -
       call, branch_cond = instruction_call(inst, arch, renderer)
       if branch_cond is not None: cond = branch_cond
       if call is not None: instructions.append(call)
-    values = {}
+    values = {ctx.scratch_stride:_c(scratch_size)}
     if routed:
       dsts = {kind:dst for dst,kind in cfg["paths"][block_pc].items()}
       target = cond.where(taken, fallthrough) if len(dsts) == 2 else taken
-      values = {taken:UOp.const(dsts[COND_TAKEN], dtypes.int), fallthrough:UOp.const(dsts[COND_NOT_TAKEN], dtypes.int)} if len(dsts) == 2 \
-          else {taken:UOp.const(next(iter(dsts.values()), -1), dtypes.int)}
+      values.update({taken:_c(dsts[COND_TAKEN], dtypes.int), fallthrough:_c(dsts[COND_NOT_TAKEN], dtypes.int)} if len(dsts) == 2
+                    else {taken:_c(next(iter(dsts.values()), -1), dtypes.int)})
       instructions.append(lower_body(route[0].store(target).sink(), "asm_branch"))
     if not instructions: return end
     body, bufs = instruction_block(instructions)
     args = [values[b] if b.addrspace is AddrSpace.ALU else buffers[b][0] for b in bufs]
-    args += [UOp.const(block_pc, dtypes.int).after(end, *scopes)]
+    args += [_c(block_pc, dtypes.int).after(end, *scopes)]
     params = [UOp.param(i, a.dtype, a.shape, addrspace=AddrSpace.ALU if i >= len(bufs) or bufs[i].addrspace is AddrSpace.ALU
-                        else AddrSpace.GLOBAL) for i,a in enumerate(args)]
+                        else AddrSpace.GLOBAL, name=bufs[i].arg.name if i < len(bufs) else None) for i,a in enumerate(args)]
     # lower the wave loop with the block, so repeated blocks share its indexing and route check too.
     w = UOp.range(n_waves, 0, AxisType.LOOP, dtype=dtypes.int)
     gates = () if block_pc in unconditional else (UOp.range(params[bufs.index(route)][w].load().eq(params[-1]).cast(dtypes.int),
@@ -1989,7 +1965,7 @@ def lift(code:bytes, arch:str, renderer:Renderer, dispatch:tuple, name:str="") -
         loop = UOp.loop(next(axes)).replace(src=(UOp(Ops.NOOP).after(end),))
         end = emit_dispatch(body, scopes+(loop,), loop)
         # waves that exited keep their route while the remaining waves repeat the loop.
-        check = loop_active(routes.after(end, *scopes, loop), active, UOp.const(pc, dtypes.int), UOp.const(n_waves, dtypes.int))
+        check = loop_active(routes.after(end, *scopes, loop), active, _c(pc, dtypes.int), _c(n_waves, dtypes.int))
         end = end.backedge(loop, active.after(check, *scopes, loop)[0].load())
       else:
         # complete the block for every participating wave before entering its successors, including barrier successors.
@@ -2004,12 +1980,12 @@ RELIFT = ContextVar("RELIFT", 0) # lift and recompile to assembly N times
 REPLAY_EMU = ContextVar("REPLAY_EMU", 0)
 
 def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int, rsrc2: int = 0x19c,
-            scratch_size: int = 0, arch: str = "rdna3", user_data: list[int]|None = None, name: str = "asm") -> int:
+            scratch_size: int = 0, arch: str = "rdna3", name: str = "asm") -> int:
   """Execute AMD assembly program. scratch_size is private_segment_fixed_size from kernel descriptor (per-lane)."""
   with Context(DEV=DEV.target('CPU', renderer='LLVM'), NOOPT=1, CHECK_OOB=0, TUPLE_ORDER=0, EMULATED_DTYPES="",
                CAPTURE_PROCESS_REPLAY=REPLAY_EMU.value, DEFAULT_FLOAT="float32", DEFAULT_INT="int32", TRANSCENDENTAL=1):
     code = bytes(to_mv(lib, lib_sz))
-    dispatch = (gx, gy, gz, lx, ly, lz, rsrc2, scratch_size, len(user_data or []))
+    dispatch = (gx, gy, gz, lx, ly, lz, rsrc2, scratch_size)
     renderer = Device['CPU'].renderer
     if (key:=(code, arch, dispatch, name, type(renderer), renderer.target, *(x.value for x in to_program_config))) not in dispatch_cache:
       dispatch_cache[key] = to_program(lift(code, arch, renderer, dispatch, f"emu {name}"), renderer)
@@ -2022,11 +1998,10 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
           break
         relift_lib = get_elf_section(ret.src[3].arg, ".text").content
     prg = dispatch_cache[key]
-  vals = {"args_ptr":args_ptr, **{f"user_data_{i}":v for i,v in enumerate(user_data or [])}}
+  vmem_buf = UOp.from_buffer(Buffer('CPU', 1 << 48, options=BufferSpec(external_ptr=0), opaque=0), dtypes.uint32)
+  args = [vmem_buf] * (max(p.arg.slot for p in prg.kernel_params) + 1)
+  for p in prg.kernel_params:
+    if p.addrspace is AddrSpace.ALU: args[p.arg.slot] = p.bind(args_ptr)
   with _MXCSRContext():
-    vmem_buf = UOp.from_buffer(Buffer('CPU', 1 << 48, options=BufferSpec(external_ptr=0), opaque=0), dtypes.uint32)
-    args = [vmem_buf] * (max(p.arg.slot for p in prg.kernel_params) + 1)
-    for p in prg.kernel_params:
-      if p.addrspace is AddrSpace.ALU: args[p.arg.slot] = p.bind(vals[p.arg.name])
     run_linear(UOp(Ops.LINEAR, src=(prg.call(*args),)), update_stats=False)
   return 0

tinygrad_repo/test/mockgpu/amd/pcode.py

@@ -364,7 +364,8 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
   'isEven': lambda a: (UOp(Ops.TRUNC, src=(a,)).cast(dtypes.int) & _const(dtypes.int, 1)).eq(_const(dtypes.int, 0)),
   'max': lambda a, b: UOp(Ops.MAX, src=(a, b)),
   'min': lambda a, b: UOp(Ops.MAX, src=(a.neg(), b.neg())).neg(),
-  'pow': lambda a, b: UOp(Ops.EXP2, src=(b.bitcast(dtypes.float32),)),
+  # pow(2.0, x): interpret the operand's float bits at its original width.
+  'pow': lambda a, b: UOp(Ops.EXP2, src=(b.bitcast({2: dtypes.half, 4: dtypes.float32, 8: dtypes.float64}[b.dtype.itemsize]),)),
   'fma': lambda a, b, c: a * b + c,
   'i32_to_f32': lambda a: a.cast(dtypes.int).cast(dtypes.float32),
   'u32_to_f32': lambda a: a.cast(dtypes.uint32).cast(dtypes.float32),

tinygrad_repo/test/null/test_hcq2.py

@@ -93,7 +93,7 @@ class TestHCQ2Deps(unittest.TestCase):
   def test_dependencies_through_selected_slices(self):
     b = UOp.param(0, dtypes.float32, 64, device=("NULL", "NULL:1"))
     for view in [b.mselect(0).shrink(((8, 16),)), b.shrink(((8, 16),)).mselect(0), b.shrink(((4, 32),)).mselect(0).shrink(((4, 12),))]:
-      tracker = hcq2.DepsTracker()
+      tracker = hcq2.DepsTracker(range(5))
       tracker.access_resources([(view, True)], 0)
       self.assertEqual(tracker.access_resources([(b.mselect(1), False)], 1), [])
       self.assertEqual(tracker.access_resources([(b.mselect(0).shrink(((16, 24),)), False)], 2), [])
@@ -102,7 +102,7 @@ class TestHCQ2Deps(unittest.TestCase):
   def test_disjoint_write_preserves_dependencies(self):
     b = UOp.param(0, dtypes.uint8, 16, device="NULL")
     for write in (False, True):
-      tracker = hcq2.DepsTracker()
+      tracker = hcq2.DepsTracker(range(5))
       tracker.access_resources([(b.shrink(((0, 4),)), write)], 0)
       self.assertEqual(tracker.access_resources([(b.shrink(((4, 8),)), True)], 1), [])
       self.assertEqual(tracker.access_resources([(b.shrink(((0, 4),)), True)], 2), [0])
@@ -110,7 +110,7 @@ class TestHCQ2Deps(unittest.TestCase):
   def test_partial_write_preserves_dependencies(self):
     b = UOp.param(0, dtypes.uint8, 16, device="NULL")
     for write in (False, True):
-      tracker = hcq2.DepsTracker()
+      tracker = hcq2.DepsTracker(range(5))
       tracker.access_resources([(b, write)], 0)
       self.assertEqual(tracker.access_resources([(b.shrink(((4, 12),)), True)], 1), [0])
       self.assertEqual(tracker.access_resources([(b.shrink(((0, 4),)), True)], 2), [0])
@@ -119,16 +119,21 @@ class TestHCQ2Deps(unittest.TestCase):
 
   def test_write_waits_for_all_readers(self):
     b = UOp.param(0, dtypes.uint8, 16, device="NULL")
-    tracker = hcq2.DepsTracker()
+    tracker = hcq2.DepsTracker(range(5))
     self.assertEqual(tracker.access_resources([(b, False)], 0), [])
     self.assertEqual(tracker.access_resources([(b, False)], 1), [])
     self.assertEqual(tracker.access_resources([(b, True)], 2), [0, 1])
     self.assertEqual(tracker.access_resources([(b, False)], 3), [2])
 
+  def test_write_waits_for_latest_reader_per_queue(self):
+    b, tracker = UOp.param(0, dtypes.uint8, 16, device="NULL"), hcq2.DepsTracker("AABAA")
+    for tag, view in enumerate([b, b, b, b.shrink(((0, 8),))]): self.assertEqual(tracker.access_resources([(view, False)], tag), [])
+    self.assertEqual(tracker.access_resources([(b, True)], 4), [1, 2, 3])
+
   def test_aliases_do_not_wait_on_themselves(self):
     b = UOp.param(0, dtypes.uint8, 16, device="NULL")
     for whole, part in ((True, False), (False, True), (True, True)):
-      tracker = hcq2.DepsTracker()
+      tracker = hcq2.DepsTracker(range(5))
       tracker.access_resources([(b, True)], 0)
       self.assertEqual(tracker.access_resources([(b, whole), (b.shrink(((4, 12),)), part)], 1), [0])
       self.assertEqual(tracker.access_resources([(b.shrink(((4, 12),)), True)], 2), [1])

tinygrad_repo/test/null/test_llm_server.py

@@ -44,6 +44,7 @@ class TestLLMServer(unittest.TestCase):
     cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
     cls.mock_tok.decode = Mock(return_value="Hello")
     cls.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: "Hello" if tid is not None else "")
+    cls.mock_tok._special_tokens = {}
     cls.mock_tok.preset = "llama3"
     cls.mock_tok.bos_id = 1
     cls.mock_tok.eos_id = 999
@@ -245,6 +246,7 @@ class TestLLMToolCalls(unittest.TestCase):
     cls.mock_tok = Mock()
     cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
     cls.mock_tok.decode = Mock(return_value="")
+    cls.mock_tok._special_tokens = {}
     cls.mock_tok.preset = "qwen2"
     cls.mock_tok.bos_id, cls.mock_tok.eos_id, cls.mock_tok.eot_id = None, 999, None
     cls.mock_tok.is_end = Mock(return_value=False)
@@ -347,6 +349,20 @@ class TestLLMToolCalls(unittest.TestCase):
     self.assertEqual(second.choices[0].message.content, "done")
     self.assertEqual(second.choices[0].finish_reason, "stop")
 
+  def test_kimi_tool_call(self):
+    self.set_output('a </think> b<|close|>think<|sep|><|open|>response<|sep|>c <tool_call>{}</tool_call> d<|close|>response<|sep|>'
+      '<|open|>tools<|sep|><|open|>call tool="read" index="1"<|sep|><|open|>argument key="x&quot;y" type="string"<|sep|>\n</tool_call>\n'
+      '<|close|>argument<|sep|><|open|>argument key="n" type="number"<|sep|>3<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|>')
+    with (patch.object(self.mock_tok, "_special_tokens", {"<|close|>": 0}),
+          patch.object(self.server.template, "render", return_value='<|open|>message role="assistant"<|sep|><|open|>think<|sep|>')):
+      chunks = list(self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"hi"}], tools=self.tools(),
+                                                       stream=True))
+    self.assertEqual(''.join(getattr(c.choices[0].delta, 'reasoning_content', '') for c in chunks if c.choices), 'a </think> b')
+    self.assertEqual(''.join(c.choices[0].delta.content or '' for c in chunks if c.choices), 'c <tool_call>{}</tool_call> d')
+    calls = [tc for c in chunks if c.choices for tc in c.choices[0].delta.tool_calls or []]
+    self.assertEqual([(tc.function.name, json.loads(tc.function.arguments)) for tc in calls], [('read', {'x"y':'\n</tool_call>\n', 'n':3})])
+    self.assertEqual(chunks[-1].choices[0].finish_reason, 'tool_calls')
+
 class TestTransformerGenerate(unittest.TestCase):
   def test_warmup(self):
     model, calls = Transformer(TEST_CONFIG), []
@@ -357,7 +373,7 @@ class TestTransformerGenerate(unittest.TestCase):
     self.assertEqual(calls, [[0], [0]])
 
   def test_template_starts_reasoning(self):
-    router = StreamRouter(reasoning=True)
+    router = StreamRouter({}, reasoning=True)
     self.assertEqual(list(router.route("reasoning</think>answer")),
                      [("reasoning_content", "reasoning"), ("content", "answer")])
 

tinygrad_repo/test/null/test_llm_tokenizer.py

@@ -162,6 +162,15 @@ class TestLLMTokenizer(unittest.TestCase):
       self.assertEqual(split["kimi-k2"](s), go, f"kimi-k2 {s!r}")
     self.assertEqual(split["kimi-k2"]("a.\n/b"), ["a", ".\n", "/b"])
 
+  def test_joyai_llm_split(self):
+    tok = SimpleTokenizer({}, {}, "joyai-llm")
+    tok._encode_word = lambda word: [word.decode()]
+    self.assertEqual(tok.encode("é café"), ["é", " café"])
+    self.assertEqual(tok.encode("  12345"), ["  ", "123", "45"])
+    self.assertEqual(tok.encode("it's I'M don'T"), ["it", "'s", " I", "'M", " don", "'T"])
+    self.assertEqual(tok.encode("  你好hello"), ["  ", "你好", "hello"])
+    self.assertEqual(tok.encode("a\u200b"), ["a", "\u200b"])
+
   def test_stream_decoder(self):
     """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""
     bs = [*range(33, 127), *range(161, 173), *range(174, 256)]

tinygrad_repo/test/null/test_tensor.py

@@ -7,7 +7,8 @@ from tinygrad.dtype import DTYPES_DICT, DType
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.codegen import to_program
+from tinygrad.codegen import full_rewrite_to_sink
+from tinygrad.codegen.late.linearizer import linearize
 
 x_init = np.random.randn(1,3).astype(np.float32)
 W_init = np.random.randn(3,3).astype(np.float32)
@@ -56,9 +57,7 @@ class TestIdxUpcast(unittest.TestCase):
     for si in linear.src:
       ast = si.src[0]
       if ast.op is Ops.SINK:
-        renderer = Device[si.src[1].buffer.device].renderer
-        prg = to_program(ast, renderer)
-        return tuple(prg.src[1].src)
+        return tuple(linearize(full_rewrite_to_sink(ast, Device[si.src[1].buffer.device].renderer)))
 
   def _assert(self, dtype: DType, a: Tensor):
     uops = self._schedule_render(a)

tinygrad_repo/test/null/test_uops.py

@@ -73,6 +73,26 @@ class TestDTypeFromUOp(unittest.TestCase):
     type_verify(out.sink(), spec_program)
 
 class TestMemoryCoalescing(unittest.TestCase):
+  def test_store_deps_ignore_reg_stores(self):
+    buf, reg = UOp.param(0, dtypes.int, 2), UOp.placeholder((2,), dtypes.int, 1, addrspace=AddrSpace.REG)
+    lanes = [buf.index(i).store(UOp.const(i, dtypes.int).after(reg.index(i).store(i))) for i in range(2)]
+    sink = memory_coalescing(UOp.sink(*lanes), Renderer(Target()))
+    self.assertEqual(sum(u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.GLOBAL for u in sink.toposort()), 1)
+
+  def test_store_deps_ignore_local_stores(self):
+    buf, local = UOp.param(0, dtypes.int, 4), UOp.placeholder((4,), dtypes.int, 1, addrspace=AddrSpace.LOCAL)
+    lanes = [buf.index(0).store(0)]+[buf.index(i).store(UOp.const(i, dtypes.int).after(local.index(i).store(i))) for i in range(1, 4)]
+    sink = memory_coalescing(UOp.sink(*lanes), Renderer(Target()))
+    self.assertEqual(sum(u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.GLOBAL for u in sink.toposort()), 1)
+
+  def test_store_deps_local_waits_on_global_stores(self):
+    buf, local = UOp.param(0, dtypes.int, 2), UOp.placeholder((2,), dtypes.int, 1, addrspace=AddrSpace.LOCAL)
+    # the GLOBAL stores coalesce, so the LOCAL stores must not: t1 waits on s1 and s0 waits on t0
+    s1, t0 = buf.index(1).store(5), local.index(0).store(7)
+    s0, t1 = buf.index(0).store(local.after(t0).index(0).load()), local.index(1).store(buf.after(s1).index(1).load())
+    sink = memory_coalescing(UOp.sink(s0, s1, t0, t1), Renderer(Target()))
+    self.assertEqual(sum(u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.LOCAL for u in sink.toposort()), 2)
+
   def test_volatile_view_not_coalesced(self):
     buf = UOp.param(0, dtypes.uint32, 4, volatile=True).bitcast(dtypes.int32)
     sink = memory_coalescing(UOp.sink(*(buf.index(i).load() for i in range(4))), Renderer(Target()))
@@ -107,6 +127,20 @@ class TestSafeCast(unittest.TestCase):
     self.assertEqual(a.cast(dtypes.int64).cast(dtypes.uint8).simplify(), a)
     self.assertEqual(a.cast(dtypes.uint32).cast(dtypes.uint8).simplify(), a)
 
+  def test_const_float_to_int(self):
+    for dtype in dtypes.ints:
+      for value, expected in ((math.nan, 0), (math.inf, dtype.max), (-math.inf, dtype.min), (1.9, 1)):
+        with self.subTest(dtype=dtype, value=value):
+          self.assertEqual(UOp.const(value).cast(dtype).ssimplify(), expected)
+
+  def test_const_float_to_int_wraps(self):
+    for value, dtype, expected in ((128.0, dtypes.int8, -128), (256.0, dtypes.uint8, 0), (-1.0, dtypes.uint8, 255)):
+      with self.subTest(dtype=dtype, value=value):
+        self.assertEqual(UOp.const(value).cast(dtype).ssimplify(), expected)
+
+  def test_const_nan_to_int_tensor(self):
+    self.assertEqual(Tensor(math.nan).cast(dtypes.int32).item(), 0)
+
   def test_remove_intermediate_cast(self):
     a = UOp.variable("a", 0., 100., dtype=dtypes.half)
     self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float))

tinygrad_repo/test/null/test_uops_stats.py

@@ -107,6 +107,13 @@ class TestUOpsStats(unittest.TestCase):
     self.assertEqual(prg.src[0].arg.estimates.mem, 4)
     self.assertEqual(prg.src[0].arg.estimates.lds, 4)
 
+  def test_triangular_loop_estimate(self):
+    buf = UOp.param(0, dtypes.int32, 6)
+    i = UOp.range(3, 0)
+    j = UOp.range(i+1, 1)
+    prg = to_program(buf.index(i*(i+1)//2+j).store(j).end(j, i).sink(arg=KernelInfo()), X86Renderer(Target("CPU", arch="x86_64")))
+    self.assertEqual(prg.src[0].arg.estimates.lds, 3*3*4)
+
   def test_simple_add(self):
     a = Tensor.empty(100,100)
     b = Tensor.empty(100,100)

tinygrad_repo/test/runtime/test_arange.py

@@ -39,6 +39,10 @@ class TestArange(unittest.TestCase):
 DSET, DDIM = 2048, 32
 
 class TestIndexing(unittest.TestCase):
+  def test_index_compared_with_offset(self):
+    a, b = Tensor([1, 2, 3, 4]), Tensor([0, 2, 2, 4])
+    self.assertEqual((a == b + 1).float()[Tensor([0, 2, 1])].tolist(), [1.0, 1.0, 0.0])
+
   def test_arange_2_reduce(self):
     needle = Tensor.zeros(16384, dtype=dtypes.int).contiguous()
     needle[1337] = 1
@@ -121,6 +125,18 @@ class TestIndexing(unittest.TestCase):
     self.assertEqual(((x*x) == 2.).where(data, 0.).tolist(), [0.]*4)
     self.assertEqual(((x*x) == 2.).reshape(2, 2).where(data.reshape(2, 2), 0.).sum(axis=1).tolist(), [0.]*2)
 
+  def test_index_collapse_float_idx(self):
+    x = Tensor.arange(4).cast(dtypes.float32)
+    self.assertEqual((x == 1.5).where(x+1, 0.).sum().item(), 0.)
+    self.assertEqual(((Tensor.arange(4)+1).cast(dtypes.float32) == 2.5).where(x+1, 0.).sum().item(), 0.)
+    self.assertEqual(((Tensor.arange(8).cast(dtypes.float32)+1e8) < 100000008.).where(1, 0).sum().item(), 5)
+
+  def test_index_collapse_cast_wraps(self):
+    r, y = Tensor.arange(8, dtype=dtypes.int32), Tensor([254], dtype=dtypes.int32).realize()
+    self.assertEqual(((r+y).cast(dtypes.uint8) == Tensor([1], dtype=dtypes.uint8).realize()).where(r+1, 0).sum().item(), 4)
+    self.assertEqual(((r+y).cast(dtypes.int64) == Tensor([2**32+254], dtype=dtypes.int64).realize()).where(r+1, 0).sum().item(), 0)
+    self.assertEqual((Tensor.arange(260).cast(dtypes.uint8) == 1).where(Tensor.arange(260)+1, 0).sum().item(), 260)
+
   def test_index_mnist(self, noopt=1, op_limit=512*784*13, split_reduceop=0):
     # WEBGPU generates more ops due to bitpacking of < 4-byte dtypes
     if Device.DEFAULT == "WEBGPU": op_limit *= 15

tinygrad_repo/test/runtime/test_assign.py

@@ -257,6 +257,17 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_equal(b0.numpy(), 128)
     np.testing.assert_equal(b1.numpy(), 608)
 
+  def test_assign_reads_behind_other_after(self):
+    a = Tensor([1, 2]).contiguous().realize()
+    b = Tensor([0, 0]).contiguous().realize()
+    b.assign(a + 1)
+    a.assign(b.sum().expand(2))
+    GlobalCounters.reset()
+    Tensor.realize(a, b)
+    assert_kernel_count(2)
+    self.assertEqual(a.tolist(), [5, 5])
+    self.assertEqual(b.tolist(), [2, 3])
+
   def test_crossunder_assign(self):
     a = Tensor.full((4,), 2).contiguous().realize()
     b = Tensor.full((4,), 3).contiguous().realize()

tinygrad_repo/test/runtime/test_attention.py

@@ -3,7 +3,7 @@ import numpy as np
 from tinygrad import Tensor, dtypes, nn
 from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig, MLATransformerBlock, Indexer, IndexerConfig,
-  HyperConnection, HCConfig,
+  HyperConnection, HCConfig, AttentionResidual,
   apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
 )
 from test.helpers import slow
@@ -83,6 +83,14 @@ class TestAttention(unittest.TestCase):
         self.assertIs(returned_mask, mask)
         np.testing.assert_allclose(actual.numpy(), expected, rtol=1e-6, atol=1e-7)
 
+  def test_attention_residual(self):
+    # rmsnorm makes the streams [1, -3] into scores [log 3, -log 3], softmax [9/10, 1/10] mixes them to 0.6, the sublayer squares it
+    for new_block, expected in ((True, [1, -3, 0.36]), (False, [1, -3 + 0.36])):
+      res = AttentionResidual(self._make_config(dim=1), new_block)
+      res.weight = Tensor([np.log(3)])
+      actual, _ = res(Tensor([[[[1.], [-3.]]]]), lambda x: (x.square(), None))
+      np.testing.assert_allclose(actual.numpy().flatten(), expected, rtol=1e-5)
+
   def test_indexer(self):
     indexer = Indexer(self._make_config(dim=2, q_lora_rank=2, indexer=IndexerConfig(n_heads=1, head_dim=2, top_k=2, kpool=2, norm_eps=1e-6)))
     indexer.attn_q_b.weight, indexer.attn_k.weight = Tensor.eye(2), Tensor.eye(2)
@@ -131,7 +139,8 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     return TransformerConfig(**({"num_blocks":1, "dim":8, "hidden_dim":16, "n_heads":1, "n_kv_heads":1,
                                  "norm_eps":1e-5, "vocab_size":32, "head_dim":8, "rope_theta":10000.0,
                                  "rope_dim":8, "v_head_dim":8, "max_context":4, "ssm_layers":(True,),
-                                 "ssm":SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4)} | kwargs))
+                                 "ssm":SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4,
+                                                 silu_gate=True)} | kwargs))
 
   def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock:
     block = GatedDeltaNetBlock(config, config.ssm)
@@ -319,7 +328,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     # full prefill is proven equivalent to decode by test_gatedeltanet_reference_and_reset (delta rule) and
     # test_kda_prefill_matches_decode (kda), so use it as the baseline and only exercise multi-chunk handoffs here
     for kda in (False, True):
-      ssm = SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4, kda=kda)
+      ssm = SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4, kda=kda, silu_gate=not kda)
       config = self._make_config(ssm=ssm)
       if kda:
         block = GatedDeltaNetBlock(config, config.ssm)

tinygrad_repo/test/runtime/test_custom_kernel.py

@@ -4,7 +4,7 @@ from tinygrad.engine.realize import run_linear, get_runtime
 from tinygrad.codegen import to_program
 import numpy as np
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
-from tinygrad.helpers import getenv, DEV
+from tinygrad.helpers import getenv
 from tinygrad.schedule.rangeify import BufferizeOpts
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops, uopfunc
 from tinygrad.codegen.opt import Opt, OptOps
@@ -289,6 +289,69 @@ class TestCustomKernel(unittest.TestCase):
     tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
     self.assertTrue(tst.allclose(a@b, atol=1e-3).item())
 
+  def test_unused_reduce_dependency(self):
+    for value in (UOp.const(2), UOp.const((1, 2, 3, 4)).reshape((2, 2))):
+      with self.subTest(shape=value.shape):
+        def kernel(C:UOp) -> UOp:
+          i = UOp.range(3, 0, AxisType.LOOP)
+          cell = C.after(i)[0]
+          v = value.after(cell.store(cell + 1))
+          return C.after(v.reduce(i, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=()))
+        self.assertEqual(Tensor.zeros(1, dtype=dtypes.int).contiguous().custom_kernel(fxn=kernel)[0].tolist(), [3])
+
+  def test_range_merge_keeps_after_domain(self):
+    # C is read once per i or once per o, so it counts to 4, and merging o and i must not make that 16
+    o, i = UOp.range(4, 0, AxisType.LOOP), UOp.range(4, 1, AxisType.LOOP)
+    stores = {"after_i": lambda C,A: (c:=C.after(i)[0]).store(c+1),
+              "after_o": lambda C,A: UOp.group((c:=C.after(o)[0]).store(c+1), A.after(o, i)[0].store(1)),
+              "load_after_o": lambda C,A: C.after(o, i)[0].store(C.after(o)[0].load()+1)}
+    for name, store in stores.items():
+      with self.subTest(name):
+        def kernel(C:UOp, A:UOp, store=store) -> UOp: return store(C, A).end(o, i).sink(arg=KernelInfo(opts_to_apply=()))
+        self.assertEqual(Tensor.zeros(1).contiguous().custom_kernel(Tensor.empty(1), fxn=kernel)[0].item(), 4)
+
+  def test_range_merge_keeps_index_domain(self):
+    # C[c] is indexed by o alone, so each cell counts to 2 while the B store merges o and i
+    o, i = UOp.range(4, 0, AxisType.LOOP), UOp.range(4, 1, AxisType.LOOP)
+    for name, c in {"div": o//2, "cmp": (o<2).cast(dtypes.int)}.items():
+      with self.subTest(name):
+        def kernel(C:UOp, B:UOp, c=c) -> UOp:
+          return UOp.group(C[c].store(C[c]+1), B[o*4+i].store(1)).end(o, i).sink(arg=KernelInfo(opts_to_apply=()))
+        self.assertEqual(Tensor.custom_kernel(Tensor.zeros(2).contiguous(), Tensor.zeros(16).contiguous(), fxn=kernel)[0].tolist(), [2, 2])
+
+  def test_range_merge_keeps_inner_reduce(self):
+    # the inner reduce ends r1 alone, so r0 and r1 don't merge. the inner sum is 4, the outer adds it 4*4*4 times
+    def kernel(C:UOp, A:UOp) -> UOp:
+      r0, r1, r2 = [UOp.range(4, n, AxisType.UPCAST) for n in range(3)]
+      return C[0].store(A[r2].reduce(r1, arg=Ops.ADD).reduce(r0, r1, r2, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=()))
+    self.assertEqual(Tensor.custom_kernel(Tensor.empty(1), Tensor.ones(4).contiguous(), fxn=kernel)[0].item(), 256)
+
+  def test_range_merge_keeps_reduce_order(self):
+    # the reduce body stores, so it loops a, k, b in order, and the digits of S are the A[a*2+b] it visits
+    def kernel(O:UOp, S:UOp, A:UOp) -> UOp:
+      a, k, b = [UOp.range(2, n, AxisType.LOOP) for n in range(3)]
+      v = A[a*2+b].after((cell:=S.after(a, k, b)[0]).store(cell*10+A[a*2+b]))
+      return O[0].store(v.reduce(a, k, b, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=()))
+    o, s = Tensor.custom_kernel(Tensor.empty(1, dtype=dtypes.int), Tensor.zeros(1, dtype=dtypes.int).contiguous(),
+                                Tensor([1, 2, 3, 4], dtype=dtypes.int), fxn=kernel)[:2]
+    Tensor.realize(o, s)
+    self.assertEqual(s.item(), 12123434)
+
+  def test_range_merge_keeps_loop_order(self):
+    # loops nest by axis id whatever order END lists them, so x is outer and the digits of S are A[0], A[2], A[1], A[3]
+    def kernel(S:UOp, A:UOp) -> UOp:
+      x, y = UOp.range(2, 0, AxisType.LOOP), UOp.range(2, 1, AxisType.LOOP)
+      return (cell:=S.after(x, y)[0]).store(cell*10+A[y*2+x]).end(y, x).sink(arg=KernelInfo(opts_to_apply=()))
+    s = Tensor.custom_kernel(Tensor.zeros(1, dtype=dtypes.int).contiguous(), Tensor([1, 2, 3, 4], dtype=dtypes.int), fxn=kernel)[0]
+    self.assertEqual(s.item(), 1324)
+
+  def test_range_merge_then_shrink(self):
+    # r0<2 shrinks the merged range of 8 to 4, which is the same UOp as r0, so it isn't shrunk again to 2
+    def kernel(C:UOp) -> UOp:
+      r0, r1 = UOp.range(4, 0, AxisType.LOOP), UOp.range(2, 1, AxisType.LOOP)
+      return (cell:=C[(r0*2+r1).valid(r0<2)]).store(cell+1).end(r0, r1).sink(arg=KernelInfo(opts_to_apply=()))
+    self.assertEqual(Tensor.zeros(8, dtype=dtypes.int).contiguous().custom_kernel(fxn=kernel)[0].tolist(), [1, 1, 1, 1, 0, 0, 0, 0])
+
   @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "test requires locals")
   def test_group_reduce_split_range(self):
     # j%2 splits j into two ranges, both are still LOCAL
@@ -458,6 +521,14 @@ class TestCustomKernel(unittest.TestCase):
       out = Tensor.custom_kernel(Tensor.zeros(2, dtype=dtypes.int).contiguous().realize(), fxn=functools.partial(kernel, loop=loop))[0]
       self.assertEqual(out.tolist(), [7, 9])
 
+  def test_store_value_behind_coalesced_store(self):
+    def c(v): return UOp.const(v, dtypes.int)
+    for stores, want in ((lambda A: (A[0].store(7), A[1].store(A.after(A[0].store(7))[0]+5)), [7, 12, 5, 0]),
+                         (lambda A: (A[1].store(c(1).after(A[0].store(7)) + A[2].load()),), [7, 6, 5, 0]),
+                         (lambda A: (A[1].store(c(1).after(A[2].store(2))), A[3].store(c(3).after(A[0].store(10)))), [10, 1, 2, 3])):
+      out = Tensor([0, 0, 5, 0], dtype=dtypes.int).contiguous().custom_kernel(fxn=lambda A: UOp.sink(*stores(A), arg=KernelInfo(opts_to_apply=())))[0]
+      self.assertEqual(out.tolist(), want)
+
   def test_deleted_store_keeps_deps(self):
     for f in (lambda a, s: s[1].store(s[1].load()), lambda a, s: s[UOp.invalid()].store(5),
               lambda a, s: s[1].store(UOp.invalid()), lambda a, s: a[UOp.invalid()].store(s[1].load())):
@@ -499,6 +570,42 @@ class TestCustomKernel(unittest.TestCase):
     a = Tensor.arange(16).reshape(4, 4).float().contiguous().realize()
     self.assertEqual(Tensor.custom_kernel(Tensor.empty(4, 4), a, fxn=kernel)[0].tolist(), a[:, :1].pad(((0, 0), (0, 3))).tolist())
 
+  def test_shrink_range_keeps_ungated_index(self):
+    # q is ungated in C[q], so q keeps its 4 iterations after r shrinks to its gate r<2
+    q, r = UOp.range(4, 0, AxisType.LOOP), UOp.range(4, 1, AxisType.LOOP)
+    def kernel(C:UOp, D:UOp) -> UOp:
+      return UOp.group(C[q].store(1), D[(q*4+r).valid((q<2)&(r<2))].store(1)).end(q, r).sink(arg=KernelInfo(opts_to_apply=()))
+    c, d = Tensor.custom_kernel(Tensor.zeros(4, dtype=dtypes.int).contiguous(), Tensor.zeros(16, dtype=dtypes.int).contiguous(), fxn=kernel)[:2]
+    Tensor.realize(c, d)
+    self.assertEqual(c.tolist(), [1, 1, 1, 1])
+    self.assertEqual(d.tolist(), [1, 1, 0, 0, 1, 1, 0, 0] + [0]*8)
+
+  def test_shrink_range_keeps_ungated_store(self):
+    # C[0] stores r every iteration, so r keeps its 4 iterations though D gates it to r<2
+    r = UOp.range(4, 0, AxisType.LOOP)
+    def kernel(C:UOp, D:UOp) -> UOp:
+      return UOp.group(C[0].store(r.cast(dtypes.int)), D[r.valid(r<2)].store(1)).end(r).sink(arg=KernelInfo(opts_to_apply=()))
+    c, d = Tensor.custom_kernel(Tensor.zeros(1, dtype=dtypes.int).contiguous(), Tensor.zeros(4, dtype=dtypes.int).contiguous(), fxn=kernel)[:2]
+    Tensor.realize(c, d)
+    self.assertEqual((c.tolist(), d.tolist()), ([3], [1, 1, 0, 0]))
+
+  def test_shrink_range_keeps_store_gated_otherwise(self):
+    # C[0] is gated by r>=2, not r<c, so r keeps its 4 iterations though D gates it to r<2
+    r = UOp.range(4, 0, AxisType.LOOP)
+    def kernel(C:UOp, D:UOp) -> UOp:
+      return UOp.group(C[UOp.const(0).valid(r>=2)].store(r.cast(dtypes.int)), D[r.valid(r<2)].store(1)).end(r).sink(arg=KernelInfo(opts_to_apply=()))
+    c, d = Tensor.custom_kernel(Tensor.zeros(1, dtype=dtypes.int).contiguous(), Tensor.zeros(4, dtype=dtypes.int).contiguous(), fxn=kernel)[:2]
+    Tensor.realize(c, d)
+    self.assertEqual((c.tolist(), d.tolist()), ([3], [1, 1, 0, 0]))
+
+  def test_shrink_range_keeps_symbolic_end(self):
+    # D gates r to r<4, but r only runs to n=2
+    def kernel(D:UOp, N:UOp) -> UOp:
+      r = UOp.range(N, 0, AxisType.LOOP)
+      return D[r.valid(r<4)].store(1).end(r).sink(arg=KernelInfo(opts_to_apply=()))
+    d = Tensor.custom_kernel(Tensor.zeros(8, dtype=dtypes.int).contiguous(), Tensor(Variable("n", 1, 8, dtypes.int).bind(2)), fxn=kernel)[0]
+    self.assertEqual(d.tolist(), [1, 1, 0, 0, 0, 0, 0, 0])
+
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "kernel timing not supported")
   def test_invalids_into_custom_kernel_with_beam(self):
     a = Tensor.full((4, 4), 3.).contiguous()
@@ -592,16 +699,17 @@ class TestCustomKernel(unittest.TestCase):
       self.assertEqual(a.flatten().tolist(), [2, 2, 3, 3])
     self.assertEqual(a.shape, (2, 2))
 
-  @unittest.skipUnless((isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage) or Device.DEFAULT == "PYTHON") and
+  @unittest.skipUnless((isinstance(Device[Device.DEFAULT].renderer, (CStyleLanguage, LLVMRenderer)) or Device.DEFAULT == "PYTHON") and
                        Device.DEFAULT != "WEBGPU", "binary not supported on this backend")
-  @unittest.skipIf(Device.DEFAULT == "AMD" and DEV.interface.startswith("MOCK"), "AMD emulator doesn't support direct pc read and write")
   def test_binary(self):
     payload = bytes(range(256))
+    def view(x): return x.reshape((16, 16)).permute((1, 0)).reshape((256,))
     def kernel(out:UOp):
       i = UOp.range(len(payload), 0)
-      data = UOp(Ops.BINARY, arg=payload)
+      data = view(UOp(Ops.BINARY, arg=payload))
       return out[i].store(data[i]).end(i).sink(arg=KernelInfo(name="binary", opts_to_apply=()))
-    self.assertEqual(Tensor.empty(len(payload), dtype=dtypes.uint8).custom_kernel(fxn=kernel)[0].tolist(), list(payload))
+    self.assertEqual(Tensor.empty(len(payload), dtype=dtypes.uint8).custom_kernel(fxn=kernel)[0].tolist(),
+                     view(Tensor(payload, dtype=dtypes.uint8)).tolist())
 
   def test_one_lane_stack_runtime_index(self):
     def kernel(A:UOp): return A[0].store(UOp.stack(UOp.const(42, dtypes.int)).index(A[1].load())).sink(arg=KernelInfo(opts_to_apply=()))
@@ -638,6 +746,18 @@ class TestCallInKernel(unittest.TestCase):
     out2 = Tensor.custom_kernel(Tensor.empty(1, dtype=x.dtype), x, fxn=main)[0]
     self.assertEqual(out2.tolist(), [N*(N+3)//2])
 
+  def test_reduce_over_call_keeps_loop_order(self):
+    # a call can hide stores, so the reduce loops a, k, b in order, and the digits of S are the A[a*2+b] it visits
+    @uopfunc
+    def update(S:UOp, A:UOp) -> UOp: return S[0].store(S[0]*10+A[0]).sink()
+    def kernel(O:UOp, S:UOp, A:UOp) -> UOp:
+      a, k, b = [UOp.range(2, n, AxisType.LOOP) for n in range(3)]
+      return O[0].store(A[a*2+b].after(update(S.after(k), A[a*2+b])).reduce(a, k, b, arg=Ops.ADD)).sink(arg=KernelInfo(opts_to_apply=()))
+    o, s = Tensor.custom_kernel(Tensor.empty(1, dtype=dtypes.int), Tensor.zeros(1, dtype=dtypes.int).contiguous(),
+                                Tensor([1, 2, 3, 4], dtype=dtypes.int), fxn=kernel)[:2]
+    Tensor.realize(o, s)
+    self.assertEqual(s.item(), 12123434)
+
   def test_nested_call(self):
     @uopfunc
     def incr(out:UOp, A:UOp):

tinygrad_repo/test/runtime/test_dtype.py

@@ -338,6 +338,9 @@ class TestUint8DType(TestDType):
   @unittest.skipIf(Device.DEFAULT == "CUDA" or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently")
   def test_uint8_to_int8_overflow(self):
     _test_op(lambda: Tensor([255, 254, 253, 252], dtype=dtypes.uint8).cast(dtypes.int8), dtypes.int8, [-1, -2, -3, -4])
+  @unittest.skipIf(Device.DEFAULT == "CUDA" or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently")
+  def test_uint8_to_int8_overflow_compare(self):
+    self.assertEqual((Tensor([128, 5], dtype=dtypes.uint8).cast(dtypes.int8) < 0).tolist(), [True, False])
 
 class TestBitCast(unittest.TestCase):
   def test_subword_slice(self):

tinygrad_repo/test/runtime/test_dtype_alu.py

@@ -5,6 +5,7 @@ from tinygrad.helpers import EMULATED_DTYPES, DEV, getenv
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.nir import NIRRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
 from tinygrad.renderer.x86 import X86Renderer
@@ -334,6 +335,28 @@ class TestDTypeALU(unittest.TestCase):
     self.assertEqual(((t+t*(2**31-1))//2).tolist(), [-2**30, 0])
     self.assertEqual((((u+t*(2**31-1))+t)//2).tolist(), [-2**30, 0])
 
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, CUDARenderer), "nvrtc 11.5 reassociates (x+c)+y<d and folds it to True")
+  def test_reassociate_no_overflow(self):
+    t, c = Tensor.arange(3, dtype=dtypes.int32), Tensor([False, True])
+    a, b, x, y = (Tensor(v, dtype=dtypes.int32) for v in ([1, 0], [1, 1], [2, 0], [2**30, 1]))
+    self.assertEqual((((t*2**29-2**30)+t*2**29) < 2**29).tolist(), [True, True, False])
+    self.assertEqual((((t-2**30)+t*(2**30-1)) < 2**29).tolist(), [True, True, False])
+    self.assertEqual((((a*2**30-2**30)+b*2**30) < 2**29).tolist(), [False, True])
+    self.assertEqual((((x*-1)*y) < -1).tolist(), [True, False])
+    self.assertEqual((((c.where(0, b*2**30)-2**30)+c.where(0, b*2**30)) < 2**29).tolist(), [False, True])
+
+  def test_float_reassociate(self):
+    x, y, c = Tensor([1e8]), Tensor([3.]), Tensor([True])
+    self.assertEqual(((x-1e8)+Tensor([1.])).tolist(), [1.])
+    self.assertEqual(((Tensor([1e30])*1e10)*Tensor([1e-30])).tolist(), [math.inf])
+    self.assertEqual(((x+y)+y).tolist(), [1e8])
+    self.assertEqual(((Tensor([1.])+(z:=Tensor([1e38]))*4)+z*-3).tolist(), [math.inf])
+    self.assertEqual(((Tensor([1.])+z*-3.5)+z).tolist(), [-math.inf])
+    self.assertEqual(((Tensor([1.])+c.where(1e8, x))+c.where(-1e8, x)).tolist(), [0.])
+
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer), "backend compiler reassociates floats")
+  def test_float_reassociate_consts(self): self.assertEqual(((Tensor([1.])+1e8)-1e8).tolist(), [0.])
+
   def test_fused_div_matches_unfused(self):
     x, y, z = Tensor([3e38, 1e30, 1.7]), Tensor([1e20, 1e30, 1.3]), Tensor([1e20, 1e10, 3.1])
     for f, g in ((lambda: x*y, lambda p: p/y), (lambda: x/y, lambda p: p/z), (lambda: z*z, lambda p: 1/p), (lambda: z*z*z, lambda p: 1/p)):
@@ -409,6 +432,11 @@ class TestDTypeALU(unittest.TestCase):
 
   def test_shl(self): self._test_shl()
   def test_shr(self): self._test_shr()
+  def test_shr_long_of_small_value(self):
+    self.assertEqual((Tensor([-1, 5], dtype=dtypes.int8).cast(dtypes.int64) >> 32).tolist(), [-1, 0])
+  def test_mod_long_of_int_min(self):
+    a, b = Tensor([-2**31, 5], dtype=dtypes.int32).cast(dtypes.int64), Tensor([-1, -2], dtype=dtypes.int8).minimum(-1).cast(dtypes.int64)
+    self.assertEqual((a % b).tolist(), [0, -1])
 
   @Context(EMULATED_DTYPES="long")
   def test_emulated_shl(self): self._test_shl()

tinygrad_repo/test/runtime/test_linearizer.py

@@ -109,7 +109,7 @@ class TestLinearizer(unittest.TestCase):
     uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     # LOAD -> RANGE -> LOAD -> STORE
-    assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
+    assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD and x.src[0].buf_uop.arg.name != "args"]) == 1
 
   def test_range_outer_op_before_phi_nested_range(self):
     a = Tensor.randn(2, ).realize()

tinygrad_repo/test/runtime/test_llm_tp.py

@@ -105,14 +105,16 @@ class TestTensorParallel(unittest.TestCase):
     typ = Q.Q4_K
     data = random_data(shape, typ).reshape(-1)
     w = ggml_data_to_tensor(Tensor(data, device='CPU'), E*O*I, typ).reshape(shape).float().numpy()
-    sel = np.array([[[3, 1]]], dtype=np.int32)
-    for axis in (1, 2):
+    sel, T = np.array([[[3, 1], [0, 2]]], dtype=np.int32), Variable('T', 1, 2).bind(1)
+    for axis in (0, 1, 2):
       with self.subTest(axis=axis):
         layer = ExpertWeights(E, I, O)
         layer.weight = gguf_shard({'w': (Tensor(data, device='CPU'), shape, typ)}, DEVICES, {'w': axis})['w'].half()
-        x = np.random.default_rng(1).normal(size=(1, 1, 1 if axis == 1 else 2, I)).astype(np.float16)
+        x = np.random.default_rng(1).normal(size=(1, 2, 1 if axis == 1 else 2, I)).astype(np.float16)
         ref = np.stack([w[e] @ x[0, 0, j % x.shape[2]].astype(np.float32) for j, e in enumerate(sel[0, 0])])
-        got = layer(Tensor(sel).shard(DEVICES), Tensor(x).shard(DEVICES, 3 if axis == 2 else None)).float().to(Device.DEFAULT).numpy()
+        got = layer(Tensor(sel).shard(DEVICES)[:, :T], Tensor(x).shard(DEVICES, 3 if axis == 2 else None)[:, :T])
+        got = got.pad_to(got.max_shape).float().to(Device.DEFAULT).numpy()
+        if axis == 0: got = got.sum(0)
         self.assertEqual(layer.ggml_type, typ)
         np.testing.assert_allclose(got[0, 0], ref, rtol=2e-2, atol=2e-2*np.abs(ref).max())
 
@@ -121,7 +123,7 @@ class TestTensorParallel(unittest.TestCase):
     rng = np.random.default_rng(42)
     # Keep GQA, partial RoPE, gated attention and replicated SSM state; custom kernels have separate tests.
     config = TransformerConfig(num_blocks=2, dim=8, hidden_dim=8, n_heads=4, n_kv_heads=2, norm_eps=1e-5, vocab_size=16, head_dim=4,
-      v_head_dim=4, rope_theta=10000, rope_dim=2, max_context=4, qk_norm=4, attn_output_gate=True, ssm=SSMConfig(2, 2, 2, 4, 8),
+      v_head_dim=4, rope_theta=10000, rope_dim=2, max_context=4, qk_norm=4, attn_output_gate=True, ssm=SSMConfig(2, 2, 2, 4, 8, silu_gate=True),
       ssm_layers=(True, False))
     with tempfile.TemporaryDirectory() as folder:
       writer = GGUFWriter(path:=pathlib.Path(folder)/'model.gguf', 'qwen35')

tinygrad_repo/tinygrad/codegen/__init__.py

@@ -27,6 +27,7 @@ from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_c
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
 from tinygrad.codegen.late.coalesce import memory_coalescing, pm_simplify_add_image
 from tinygrad.helpers import all_same, argsort, partition, to_function_name
+from tinygrad.device import TinyELF
 from tinygrad.uop.ops import _broadcast_shape, identity_element
 from tinygrad.schedule.rangeify import BufferizeOpts
 
@@ -406,7 +407,7 @@ class LowerCallsCtx:
 def lower_call(ctx:LowerCallsCtx, call:UOp, b:UOp) -> UOp|None:
   if b.op is Ops.PROGRAM and b.src[-1].op is Ops.BINARY: return None
   bodies = ctx.names.setdefault(name:=to_function_name(call.arg.name), {})
-  ki = KernelInfo(name + (f"_{i}" if (i:=bodies.setdefault(b, len(bodies))) else ""), optimized=True)
+  ki = KernelInfo(name + (f"_{i}" if (i:=bodies.setdefault(b, len(bodies))) else ""), optimized=True, kernel=False)
   if ctx.renderer.inline_prefix:
     b = graph_rewrite(b, pm_lower_calls, ctx=ctx, name="lower calls", walk=True)
     b = UOp(Ops.LINEAR, src=tuple(line_rewrite(list(b.src[1].src) if b.op is Ops.PROGRAM \
@@ -418,6 +419,13 @@ pm_lower_calls = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM), name="b"),), allow_any_len=True, name="call"), lower_call),
 ])
 
+def pack_args(sink:UOp, params:tuple[UOp, ...]) -> UOp:
+  sig = TinyELF.iter_sig(tuple((None, p.addrspace, p.dtype, ()) for p in params))
+  args = UOp.param(0, dtypes.uint8, len(params) * 8, name="args")
+  loads = [args.bitcast(dt).index(UOp.cconst(off // dt.itemsize, dtypes.int)).load() for off, dt in sig]
+  reads = {p: v if p.addrspace is AddrSpace.ALU else UOp(Ops.REFPARAM, src=(p.src[0], v), arg=p.arg) for p, v in zip(params, loads)}
+  return sink.substitute(reads, name="pack args")
+
 def do_linearize(ctx:Renderer, prg:UOp, sink:UOp) -> UOp:
   if DEBUG >= 3 and sink.arg.applied_opts: print(f"{sink.arg.function_name:<25} opts: {sink.arg.applied_opts}")
   lst = line_rewrite(linearize(sink), pm_linearize_cleanups+pm_alloc_to_buf)
@@ -484,11 +492,11 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
     assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program"
     full_sink = full_rewrite_to_sink(ast, renderer, optimize=not ast.arg.optimized)
     prog_info = ProgramInfo.from_sink(full_sink, renderer.target)
+    if (isinstance(renderer, ISARenderer) or renderer.packed_args) and full_sink.arg.estimates is None:
+      full_sink = full_sink.replace(arg=replace(full_sink.arg, estimates=Estimates.from_uops(tuple(linearize(full_sink)), ignore_indexing=True)))
+    if renderer.packed_args and full_sink.arg.kernel: full_sink = pack_args(full_sink, prog_info.params)
     # instruction selection
     if isinstance(renderer, ISARenderer):
-      # Instruction selection replaces LOAD/STORE/ALU with INS, so estimate while their meaning is still available.
-      if full_sink.arg.estimates is None:
-        full_sink = full_sink.replace(arg=replace(full_sink.arg, estimates=Estimates.from_uops(tuple(linearize(full_sink)), ignore_indexing=True)))
       full_sink = graph_rewrite(full_sink, renderer.pre_isel_matcher, ctx=itertools.count(-1, -1), name="pre instruction selection", bottom_up=True)
       full_sink = graph_rewrite(full_sink, renderer.isel_matcher, ctx=IselContext(full_sink), name="instruction selection", bottom_up=True)
     prg = UOp(Ops.PROGRAM, src=(full_sink,), arg=prog_info)

tinygrad_repo/tinygrad/codegen/gpudims.py

@@ -92,9 +92,6 @@ pm_device_to_var = PatternMatcher([
   # the DEVICE axis is not a program axis, it's bound per device at launch. lower it to the _device_num variable (like SPECIAL for devices)
   (UPat(Ops.RANGE, name="r"),
    lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype) if r.axis_type is AxisType.DEVICE else None),
-  # ENDs that closed a DEVICE range no longer close it
-  (UPat(Ops.END, name="e"), lambda e: e.replace(src=(e.src[0],)+tuple(s for s in e.src[1:] if s.op is not Ops.PARAM))
-   if any(s.op is Ops.PARAM and s.arg.name == '_device_num' for s in e.src[1:]) else None),
 ])
 
 # Run once: grouping creates new GLOBAL/LOCAL ranges, which must not be grouped again.

tinygrad_repo/tinygrad/codegen/late/coalesce.py

@@ -107,9 +107,12 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
   if getenv("DMC"): return sink
 
   # collect
-  memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp, object], dict[int, list[UOp]]] = defaultdict(dict)
+  memory: defaultdict[tuple[Ops, UOp, UOp|None, UOp, object, int|tuple[int, int]], dict[int, list[UOp]]] = defaultdict(dict)
+  # depth[u] is the lexicographically largest (GLOBAL stores, LOCAL stores) count on any chain u waits on, u included
+  depth: dict[UOp, tuple[int, int]] = {}
   for u in sink.toposort():
-    # TODO: this should handle images too, it's just memory coalescing
+    depth[u] = waits = max([depth[s] for s in u.src_without_body], default=(0, 0))
+    if u.op is Ops.STORE: depth[u] = (waits[0]+(u.src[0].addrspace is AddrSpace.GLOBAL), waits[1]+(u.src[0].addrspace is AddrSpace.LOCAL))
     if u.op in {Ops.LOAD, Ops.STORE}:
       assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores"
       assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}"
@@ -118,48 +121,47 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
       if buf.buf_uop.op is Ops.PARAM and buf.buf_uop.arg.volatile: continue # volatile accesses never merge
       idx, valid = idx_u.get_idx(), idx_u.get_valid()
       if idx.is_invalid: continue
-      root_src: UOp|str
+      root_src: UOp|None
       if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].val
       elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].val
-      elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.val
+      elif idx.op is Ops.CONST: root_src, arg = None, idx.val
       else: root_src, arg = idx, 0
       # loads/stores only coalesce with others carrying the same arg (e.g. the nontemporal flag)
-      memory[(u.op, buf, root_src, valid, u.arg)].setdefault(arg, []).append(u)
+      # and the same depth (GLOBAL count only for GLOBAL accesses), so no merge closes a dependency cycle
+      memory[(u.op, buf, root_src, valid, u.arg, waits[0] if buf.addrspace is AddrSpace.GLOBAL else waits)].setdefault(arg, []).append(u)
 
   # build replacements
   replacements = {}
-  for (op,buf,_,_,ld_arg),offsets in memory.items():
+  for (op,buf,_,_,ld_arg,_),offsets in memory.items():
     # allowed lengths (copied in)
     lengths = []
     must_divide = True
     if ctx.target.device == "DSP":
       lengths = [128,64,32,16,8,4]
       must_divide = False
-    elif is_image_shape(buf._shape):
-      lengths = [4]
     elif buf.dtype in (dtypes.float, dtypes.half, dtypes.int, dtypes.uint, *dtypes.fp8s) and ctx.supports_float4:
       # TODO: a better way to get this than ctx
       lengths = [8,4,2] if buf.dtype == dtypes.half and getenv("ALLOW_HALF8") else [4,2]
     lengths.append(1)  # worst case, it's not folded
     # do the grouping
-    grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
-    for full_grp in grouped_offsets:
+    for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0]):
+      full_grp = [offsets[x] for _,x in group]
       while len(full_grp):
         # the group starts at its first access's index, valid included
-        offset = offsets[full_grp[0]][0].src[0].src[1]
+        offset = full_grp[0][0].src[0].src[1]
         length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.get_idx().divides(l) is not None)][0]
         grp, full_grp = full_grp[:length], full_grp[length:]
-        assert op is not Ops.STORE or all(len(offsets[g]) == 1 for g in grp), "attempting multiple stores"
+        assert op is not Ops.STORE or all(len(g) == 1 for g in grp), "attempting multiple stores"
         if length == 1: continue
         idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(length)))
         if op == Ops.STORE:
-          datas = [offsets[g][0].src[1] for g in grp]
+          datas = [g[0].src[1] for g in grp]
           store = idx.store(UOp.stack(*datas))
-          for g in grp: replacements[offsets[g][0]] = store
+          for g in grp: replacements[g[0]] = store
         else:
           ld = idx.load(arg=ld_arg)
           for i,g in enumerate(grp):
-            for oo in offsets[g]: replacements[oo] = ld.index(i)
+            for oo in g: replacements[oo] = ld.index(i)
 
   # apply
   return sink.substitute(replacements, name="memory coalescing")

tinygrad_repo/tinygrad/codegen/simplify.py

@@ -15,48 +15,50 @@ pm_flatten_range = PatternMatcher([
   (UPat((Ops.REDUCE, Ops.END), name="r"), flatten_range),
 ])
 
-# index/range arithmetic uses DIV/MOD prior to late rewrite
-def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.DIV, Ops.MOD} for u in x.backward_slice)
-def simplify_merge_adjacent(u:UOp) -> UOp|None:
-  reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
-  # on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations
-  for r0, r1 in (zip(u.ended_ranges, u.ended_ranges[1:]) if u.op is Ops.END else itertools.permutations(u.ended_ranges, 2)):
-    # check same type and neither bound depends on the other range
-    if r0.axis_type == r1.axis_type and r0 not in r1.ranges and r1 not in r0.ranges:
-      # check if the ranges to merge are in the same reduces
-      if all((r0 in rngs) == (r1 in rngs) for rngs in reduce_ranges):
-        s0, s1 = r0.src[0], r1.src[0]
-        # do the merge
-        new_range = r0.replace(src=(s0*s1,))
-        nidx = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
-                             name=f"check_merge_{r0.axis_id}_{r1.axis_id}")
+# TODO: something in GroupOp
+pure_ops = GroupOp.Elementwise|GroupOp.Irreducible|{Ops.INDEX, Ops.LOAD, Ops.STACK}
 
-        # check if it simplifies. return after one merge so the next rewrite uses the new ranges,
-        # rather than continuing with stale pairs from the original ended_ranges.
-        if count_divmod(nidx) <= count_divmod(u): return nidx
+def mergeable_pairs(x:UOp) -> list[tuple[UOp, UOp]]:
+  # an END or a REDUCE with effects runs its ranges as loops nested by axis id, so only neighbors merge
+  if x.op is Ops.REDUCE and all(y.op in pure_ops|{Ops.REDUCE} for y in x.backward_slice): return list(itertools.permutations(x.ended_ranges, 2))
+  return list(itertools.pairwise(sorted(x.ended_ranges, key=lambda r: r.arg))) if x.op in {Ops.END, Ops.REDUCE} else []
+
+def merge_ranges(u:UOp) -> UOp|None:
+  enders = {x:mergeable_pairs(x) for x in u.backward_slice if x.op in range_start}
+  impure = [x for x in u.backward_slice if x.op not in pure_ops]
+  for r0, r1 in dedup(p for pairs in enders.values() for p in pairs):
+    if r0.axis_type is not r1.axis_type or r0 in r1.ranges or r1 in r0.ranges: continue
+    # every op that ends r0 or r1 must merge them, and an op that isn't pure must run as often as before
+    if any((r0, r1) not in pairs for x, pairs in enders.items() if {r0, r1} & set(x.ended_ranges)): continue
+    if any((r0 in x.ranges) != (r1 in x.ranges) for x in impure): continue
+    s1 = r1.src[0]
+    n = r0.replace(src=((r0.src[0]*s1).simplify(),))
+    sub = {r0:n//s1, r1:n%s1}
+    ret = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx=sub, name=f"check_merge_{r0.axis_id}_{r1.axis_id}")
+    if not any(v in ret.backward_slice for v in sub.values()): return ret
   return None
 
-def mark_gated(ctx, idx):
+def mark_gated(ctx, idx, store=None):
   if len(idx.src) > 1 and idx.src[1].op is Ops.WHERE:
     x, cond = idx.src[1].get_idx(), idx.src[1].get_valid()
     # get all ranges r with guards "r < c" for some const c
     guards = {r:c for v in cond.split_uop(Ops.AND) if v.op is Ops.CMPLT and (r:=v.src[0]).op is Ops.RANGE and (c:=v.src[1]).op is Ops.CONST}
   else: x, guards = idx, {}
   # ensure that we choose max(c_i) for all i where r < c_i
-  ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].vmax < c.val)}
-  # but if a range is ever ungated, we cannot shrink it
-  ctx |= {r:r.src[0] for r in x.ranges if r not in guards}
+  ctx |= {r:c.minimum(r.src[0]) for r,c in guards.items() if (r not in ctx or ctx[r].vmax < c.val)}
+  # but if a range is ever ungated, we cannot shrink it, and a STORE runs every iteration of a range its index doesn't gate
+  ctx |= {r:r.src[0] for r in (x if store is None else store).ranges if r not in guards}
 
 def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None:
   # Only the kernel root: rewriting a nested SINK would leave its enclosing END's binders unchanged.
   if x.arg is None: return None
-  ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None})
-  ctx.clear()
+  ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None and k in x.backward_slice})
   return None if ret is x else ret.simplify()
 
 pm_simplify_ranges = PatternMatcher([
-  (UPat((Ops.END, Ops.REDUCE), name="u"), simplify_merge_adjacent),
+  (UPat(Ops.SINK, name="u"), merge_ranges),
   (UPat(Ops.INDEX, name="idx"), mark_gated),
+  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, name="idx"), UPat()), name="store"), mark_gated),
   # reduce ranges can't be shrunk
   (UPat(Ops.REDUCE, name="red"), lambda ctx, red: ctx.update({r:r.src[0] for r in red.src[1:]})),
   (UPat(Ops.SINK, name="x"), lambda ctx, x: do_substitute(ctx, x, lambda r,c: r.replace(src=(c,)))),
@@ -94,7 +96,7 @@ pm_reduce_unparented = PatternMatcher([
 
 pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
   # lift x+y out of reduce on lt
-  (UPat(Ops.ADD, dtypes.sints+dtypes.floats+dtypes.weaks, [UPat.var("x"), UPat.var("y")]) < UPat.var("c"),
+  (UPat(Ops.ADD, dtypes.sints+(dtypes.weakint,), [UPat.var("x"), UPat.var("y")]) < UPat.var("c"),
    lambda x,y,c: (x < (c-y)) if no_range(y) and no_range(c) else None),
   # lift x*y out of reduce
   (UPat(Ops.MUL, dtypes.sints+(dtypes.weakint,), [UPat.var("x"), UPat.var("y")]) < UPat.var("c"),
@@ -118,10 +120,14 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
 
 pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([
   # lift x+y out of reduce on eq
-  ((UPat.var("x")+UPat.var("y")).or_casted().eq(UPat.var("c")), lambda x,y,c: x.eq(c.cast(y.dtype)-y) if no_range(y) and no_range(c) else None),
+  ((UPat.var("x", dtypes.wints)+UPat.var("y")).or_casted("a").eq(UPat.var("c", dtypes.wints)),
+   lambda x,y,a,c: x.eq(c.cast(y.dtype)-y) if no_range(y) and no_range(c) and not no_range(x) and
+   (a.op is not Ops.CAST or not a.src[0].overflows(a.dtype)) and not c.overflows(y.dtype) else None),
   # reduce on gated load becomes can substitute the range and remove the reduce
-  (UPat.var("idx").eq(UPat(Ops.RANGE, name="r").or_casted()).where(UPat.var("expr"), 0).reduce(UPat.var("r"), arg=Ops.ADD),
-   lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),
+  (UPat.var("idx", dtypes.wints).eq(UPat(Ops.RANGE, name="r").or_casted("rc")).where(UPat.var("expr"), 0)
+   .reduce(UPat.var("r"), arg=Ops.ADD),
+   lambda r,rc,idx,expr: None if r.overflows(rc.dtype) else
+   (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),
 ])
 
 def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp|None:

tinygrad_repo/tinygrad/device.py

@@ -69,7 +69,7 @@ def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
   return canonical[0] if len(canonical:=tuple(Device.canonicalize(d) for d in device)) == 1 else canonical
 
 def is_disk_device(device:str|tuple[str, ...]) -> bool:
-  return any(d.split(":", 1)[0].upper() == "DISK" for d in to_tuple(device))
+  return any(d.split(":", 1)[0].upper() in ("DISK", "NVME") for d in to_tuple(device))
 
 # **************** Profile ****************
 
@@ -182,7 +182,7 @@ class Buffer:
     else: storage = self.allocator.alloc(self.nbytes, self.options)
     storage = replace(storage, host=storage.host.view(self.offset, self.nbytes, fmt='B') if storage.host is not None else None)
     if self._base is None:
-      if not self.device.startswith("DISK") and self.options.external_ptr is None:
+      if not is_disk_device(self.device) and self.options.external_ptr is None:
         GlobalCounters.mem_used += self.nbytes
         GlobalCounters.mem_used_per_device[self.device] += self.nbytes
       if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"nbytes":self.nbytes}))
@@ -193,7 +193,7 @@ class Buffer:
     assert self._storage is not None, "buffer must be allocated to deallocate"
     if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}")
     if self._base is None:
-      if GlobalCounters is not None and not self.device.startswith("DISK") and self.options.external_ptr is None:
+      if GlobalCounters is not None and not is_disk_device(self.device) and self.options.external_ptr is None:
         GlobalCounters.mem_used -= self.nbytes
         GlobalCounters.mem_used_per_device[self.device] -= self.nbytes
       if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num))

tinygrad_repo/tinygrad/dtype.py

@@ -78,6 +78,8 @@ class DType(metaclass=DTypeMetaClass):
     return float("inf") if dtypes.is_float(self) else True
   def const(self, val: ConstType):
     if isinstance(val, InvalidType): return val
+    if self in dtypes.ints and isinstance(val, float) and not math.isfinite(val):
+      val = 0 if math.isnan(val) else self.max if val > 0 else self.min
     if self in dtypes.ints: return truncate[self](int(val))
     # NOTE: float('nan') != float('nan'), so we canonicalize here
     if isinstance(val, float) and math.isnan(val): val = math.nan
@@ -91,7 +93,7 @@ class DTypes:
   def is_float(x: DType) -> bool: return x in (dtypes.floats + (dtypes.weakfloat,))
   @staticmethod # static methods on top, or bool in the type info will refer to dtypes.bool
   @functools.cache
-  def is_int(x: DType) -> bool: return x in (dtypes.ints + (dtypes.weakint,))
+  def is_int(x: DType) -> bool: return x in dtypes.wints
   @staticmethod
   @functools.cache
   def is_unsigned(x: DType) -> bool: return x in dtypes.uints
@@ -158,6 +160,7 @@ class DTypes:
   uints = (uint8, uint16, uint32, uint64)
   sints = (int8, int16, int32, int64)
   ints = uints + sints
+  wints = ints + (weakint,)
   weaks = (weakint, weakfloat)
   all = floats + ints + (bool,) # noqa: A003
 
@@ -313,12 +316,9 @@ def _from_np_dtype(npdtype:'np.dtype') -> DType: # type: ignore [name-defined] #
 def _to_torch_dtype(dtype:DType) -> 'torch.dtype'|None:  # type: ignore [name-defined] # noqa: F821
   import numpy as np, torch
   dtype = strong_dtype(dtype)
-  if dtype == dtypes.uint64: return torch.uint64
   if dtype == dtypes.bfloat16: return torch.bfloat16
   if dtype in dtypes.fp8s: return torch.uint8
-  # NOTE: torch doesn't expose this mapping with a stable API
-  try: return torch.from_numpy(np.array([], dtype=_to_np_dtype(dtype))).dtype
-  except TypeError: return None
+  return getattr(torch, np.dtype(dtype.fmt).name, None) if dtype.fmt is not None else None
 @functools.cache
 def _from_torch_dtype(torchdtype:'torch.dtype') -> DType: # type: ignore [name-defined] # noqa: F821
   return {v:k for k in DTYPES_DICT.values() if (v:=_to_torch_dtype(k)) is not None}[torchdtype]

tinygrad_repo/tinygrad/llm/chat.html

@@ -1,38 +1,63 @@
-<!DOCTYPE html><html><head><meta charset="utf-8"><title>tinygrad chat</title><style>
+<!DOCTYPE html><html><head><meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, interactive-widget=resizes-content">
+<title>tinygrad chat</title><style>
   * { margin: 0 }
   body { background: #212121; color: #e3e3e3; font-family: system-ui;
-         height: 100vh; display: flex; flex-direction: column }
-  #chat { flex: 1; overflow-y: auto; padding: 20px }
-  .msg { padding: 10px 16px; margin: 8px 0; white-space: pre-wrap; border-radius: 18px }
+         height: 100dvh; overflow: hidden; display: flex; flex-direction: column }
+  #chat { flex: 1; min-height: 0; overflow-y: auto; padding: 20px }
+  .msg { padding: 10px 16px; margin: 8px 0; white-space: pre-wrap; overflow-wrap: anywhere; border-radius: 18px }
   table { border-collapse: collapse; table-layout: fixed; width: 100%; overflow-wrap: anywhere }
   th, td { border: 1px solid #555; padding: 6px 10px; text-align: left }
   a { color: #8ab4f8 } hr { border: 0; border-top: 1px solid #555 }
   .answer { white-space: normal; line-height: 1.65 } .answer > * { margin: 12px 0 }
   pre, blockquote { background: #2f2f2f; padding: 12px 16px; border-radius: 8px } pre { white-space: pre-wrap }
   .user { background: #2f2f2f; margin-left: auto; width: fit-content; max-width: 70% }
-  #input { max-width: 768px; width: 100%; margin: 20px auto; padding: 14px 20px;
-           background: #2f2f2f; color: inherit; font: inherit;
-           border: none; outline: none; resize: none; border-radius: 24px; field-sizing: content }
+  #composer { display: flex; flex-shrink: 0; align-items: center; gap: 12px; max-width: 808px; width: calc(100% - 24px);
+              box-sizing: border-box; margin: 12px auto max(12px, env(safe-area-inset-bottom));
+              padding: 12px 16px 12px 20px; background: #2f2f2f; border-radius: 24px }
+  #input { flex: 1; min-width: 0; padding: 0; background: transparent; color: inherit; font: inherit;
+           border: none; outline: none; resize: none; field-sizing: content }
+  #model { min-width: 0; max-width: 45%; padding: 6px 10px; background: #383838; color: inherit; font: inherit;
+           border: 1px solid #484848; border-radius: 16px; cursor: pointer }
 </style></head><body><div id="chat"></div>
+<div id="composer">
 <textarea id="input" rows="1" placeholder="Ask anything" autofocus></textarea>
-<script src="/assets/cdn.jsdelivr.net/npm/markdown-it@15.0.1/dist/browser/markdown-it.umd.min.js"></script>
+<select id="model" aria-label="Model"></select></div>
+<script src="./assets/cdn.jsdelivr.net/npm/markdown-it@15.0.1/dist/browser/markdown-it.umd.min.js"></script>
 <script>
+  const params = new URLSearchParams(location.hash.slice(1)), server = (params.get('server') || '').replace(/\/+$/, '');
+  const headers = {'Content-Type': 'application/json'};
+  if (params.get('api_key')) headers.Authorization = 'Bearer ' + params.get('api_key');
+  function showError(e) {
+    const d = document.createElement('div'); d.className = 'msg'; d.style.color = '#f88'; d.setAttribute('role', 'alert');
+    d.textContent = e.message; chat.appendChild(d); chat.scrollTop = chat.scrollHeight;
+  }
+  async function request(path, options = {}) {
+    const r = await fetch(server + path, {headers, ...options});
+    if (!r.ok) throw new Error(`HTTP ${r.status}: ${(await r.text()).slice(0, 500)}`);
+    return r;
+  }
+  request('/v1/models').then(r => r.json()).then(r => {
+    if (!r.data?.length) throw new Error('No models returned by the server.');
+    model.replaceChildren(...r.data.map(m => new Option(m.id, m.id, false, m.id === params.get('model'))));
+  }).catch(showError);
   let generating = false;
   input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) {
-    e.preventDefault(); if (generating) return;
-    generating = true; send().finally(() => generating = false);
+    e.preventDefault(); if (generating || !model.value) return;
+    generating = true; send().catch(showError).finally(() => generating = false);
   } };
   const msgs = [];
   const md = markdownit();
   async function send() {
     if (!input.value.trim()) return;
     msgs.push({role: 'user', content: input.value.trim()});
-    chat.innerHTML += '<div class="msg user">' + input.value.trim().replace(/</g, '&lt;') + '</div>';
+    const user = document.createElement('div'); user.className = 'msg user';
+    user.textContent = input.value.trim(); chat.appendChild(user);
     input.value = '';
     const d = document.createElement('div'); d.className = 'msg'; chat.appendChild(d);
     d.innerHTML = '<span style="color:#888"></span><div class="answer"></div>'; const [thinking, answer] = d.children;
-    const r = await fetch('/v1/chat/completions', {method: 'POST', headers: {'Content-Type': 'application/json'},
-      body: JSON.stringify({model: 'llama', messages: msgs, stream: true, temperature: 0.7})});
+    const r = await request('/v1/chat/completions', {method: 'POST',
+      body: JSON.stringify({model: model.value, messages: msgs, stream: true, temperature: 0.7})});
     let buf = '', txt = '', rsn = '';
     for (const rd = r.body.getReader(), dec = new TextDecoder();;) {
       const {done, value} = await rd.read();
@@ -41,8 +66,10 @@
       const lines = buf.split('\n');
       buf = lines.pop();
       for (const ln of lines)
-        if (ln.startsWith('data: ') && !ln.includes('[DONE]')) {
-          const dl = JSON.parse(ln.slice(6)).choices[0]?.delta;
+        if (ln.startsWith('data: ') && ln.slice(6).trim() !== '[DONE]') {
+          const data = JSON.parse(ln.slice(6));
+          if (data.error) throw new Error(data.error.message || 'API error');
+          const dl = data.choices?.[0]?.delta;
           if (dl?.reasoning_content) { rsn += dl.reasoning_content; thinking.textContent = rsn }
           if (dl?.content) { txt += dl.content; answer.innerHTML = md.render(txt) }
         }

tinygrad_repo/tinygrad/llm/cli.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import sys, argparse, codecs, itertools, typing, re, unicodedata, json, time, pathlib
+import sys, argparse, codecs, itertools, typing, re, string, unicodedata, json, time, pathlib
 from typing import TYPE_CHECKING
 from tinygrad import nn
 from tinygrad.uop.ops import UOp, Ops
@@ -12,17 +12,17 @@ class SimpleTokenizer:
   def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3",
                bos_id:int|None=None, eos_id:int=0, eot_id:int|None=None, eom_id:int|None=None, chat_template:str|None=None):
     preset = {"qwen35":"qwen2","qwen35moe":"qwen2"}.get(preset, preset)
-    if preset not in ("llama3","llama-v3","llama-bpe","qwen2","olmo","kimi-k2","tekken","glm4","gpt-4o"):
+    if preset not in ("llama3","llama-v3","llama-bpe","qwen2","olmo","kimi-k2","tekken","glm4","gpt-4o","joyai-llm"):
       raise ValueError(f"Invalid tokenizer preset '{preset}'")
     # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
     bs = [*range(33, 127), *range(161, 173), *range(174, 256)]  # bytes that map to themselves
     self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
 
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
-    # each limit is one past the category's max codepoint (Z→U+3000, N→U+1FBF9, L→U+323AF, M→U+E01EF)
+    # each limit is one past the category's max codepoint (Z→U+3000, N→U+1FBF9, L→U+323AF, M→U+E01EF, P→U+1E95F, S→U+1FBCA)
     # compact adjacent codepoints into ranges: listing them all makes re spend seconds on large prompts
     def ucat_range(pre:str|tuple[str, ...]) -> str:
-      limits = {"Z": 0x3001, "N": 0x1fbfa, "L": 0x323b0, "M": 0xe01f0}
+      limits = {"Z": 0x3001, "N": 0x1fbfa, "L": 0x323b0, "M": 0xe01f0, "P": 0x1e960, "S": 0x1fbcb}
       limit = max(limits[p if p in limits else p[0]] for p in (pre if isinstance(pre, tuple) else (pre,)))
       cps = enumerate(cp for cp in range(limit) if unicodedata.category(chr(cp)).startswith(pre))
       runs = [list(g) for _, g in itertools.groupby(cps, lambda e: e[1]-e[0])]
@@ -35,7 +35,13 @@ class SimpleTokenizer:
       sfx = "" if preset == "tekken" else f"{contr}?"
       r_w = f"{r_l}[{r_up}]*[{r_lo}]+{sfx}|{r_l}[{r_up}]+[{r_lo}]*{sfx}"
       if preset != "kimi-k2": r_p = f" ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n/]*"
-    self._split_to_word = re.compile(f"{r_w}|{r_n}|{r_p}|{r_t}")
+    self._pre_split = None
+    if preset == "joyai-llm":
+      r_ps = ucat_range(("P", "S"))
+      r_p, r_w = f" ?[{r_ps}]+[\\r\\n]*", f"[{re.escape(string.punctuation)}][A-Za-z]+|[^\\r\\n{r_p_L}{r_ps}]?[{r_p_L}{ucat_range('M')}]+"
+      # numbers and CJK runs are split off before the words
+      self._pre_split = re.compile(f"({r_n}|[\u4e00-\u9fa5\u3040-\u309f\u30a0-\u30ff]+)")
+    self._split_to_word = re.compile(f"({r_w}|{r_n}|{r_p}|{r_t})")
     self._split_to_sentence = re.compile("|".join(re.escape(tok) for tok in special_tokens.keys()) if special_tokens else r"(?!)")
 
     self._normal_tokens = {bytes(self._byte_decoder[c] for c in tok): tid for tok, tid in normal_tokens.items()}
@@ -80,7 +86,9 @@ class SimpleTokenizer:
     try: return [self._normal_tokens[p] for p in parts]
     except KeyError: raise RuntimeError("token not found")
   def _encode_sentence(self, chunk:str) -> list[int]:
-    return [tok for word in self._split_to_word.findall(chunk) for tok in self._encode_word(word.encode())]
+    if self._pre_split is None: words = self._split_to_word.findall(chunk)
+    else: words = [w for p in self._pre_split.split(chunk) for w in self._split_to_word.split(p) if w]
+    return [tok for word in words for tok in self._encode_word(word.encode())]
   def encode(self, text:str) -> list[int]:
     tokens: list[int] = []
     pos = 0

tinygrad_repo/tinygrad/llm/hf.py

@@ -1,5 +1,6 @@
 import json, pathlib
 from tinygrad import Tensor, dtypes, nn, Device
+from tinygrad.helpers import prod
 
 def e4m3fn_to_fnuz(weight:Tensor, scale:Tensor) -> tuple[Tensor, Tensor]:
   bits = weight.bitcast(dtypes.uint8)
@@ -7,25 +8,52 @@ def e4m3fn_to_fnuz(weight:Tensor, scale:Tensor) -> tuple[Tensor, Tensor]:
 
 def hf_load(path:pathlib.Path, devices:tuple[str, ...], names:dict[str, str|list[str]], shard_map:dict[str, int]) -> dict[str, Tensor]:
   index = path / "model.safetensors.index.json"
-  weights:dict[str, Tensor] = {}
   if index.exists(): weight_map = json.loads(index.read_text())["weight_map"]
   else:
-    weights = nn.state.safe_load(path / "model.safetensors")
-    weight_map = dict.fromkeys(weights, "model.safetensors")
+    _, _, metadata = nn.state.safe_load_metadata(path / "model.safetensors")
+    weight_map = dict.fromkeys(metadata, "model.safetensors")
+  needed = {s for source in names.values() for s in (source if isinstance(source, list) else [source])}
+  needed |= {s + "_scale_inv" for s in needed}
+  weights:dict[str, tuple[Tensor, dict]] = {}
 
   def pop_tensor(source:str|list[str], axis:int|None) -> Tensor:
     sources = source if isinstance(source, list) else [source]
     stack = isinstance(source, list)
     parts = []
     for s in sources:
-      if s not in weights: weights.update(nn.state.safe_load(path / weight_map[s]))
+      if s not in weights:
+        data, data_start, metadata = nn.state.safe_load_metadata(path / weight_map[s])
+        data = data[data_start:].to(devices[0]).realize()
+        weights.update({k: (data, m) for k, m in metadata.items() if k in needed})
       parts.append(weights.pop(s))
-    dtype, shape = parts[0].dtype, parts[0].shape
-    if stack: shape = (len(parts), *shape)
+
+    m0 = parts[0][1]
+    shape, dtype = m0['shape'], nn.state.safe_dtypes[m0['dtype']]
+    nbytes = prod(shape) * dtype.itemsize
+    assert all(m['shape'] == shape and m['dtype'] == m0['dtype'] and
+               0 <= m['data_offsets'][0] <= m['data_offsets'][1] <= d.numel() and
+               m['data_offsets'][1] - m['data_offsets'][0] == nbytes for d, m in parts), f"mismatched parts for {sources[0]}"
+    slices = [(d, m['data_offsets'][0]) for d, m in parts]
+    byte_shape = (*shape[:-1], shape[-1] * dtype.itemsize)
+    if stack: shape, byte_shape = (len(parts), *shape), (len(parts), *byte_shape)
     assert axis is None or shape[axis] % len(devices) == 0, f"cannot split {sources[0]} on axis {axis}"
-    if stack: w = Tensor.stack(*[t.bitcast(dtypes.uint8).to(devices[0]) for t in parts]).contiguous().realize()
-    else: w = parts[0].bitcast(dtypes.uint8)
-    return w.shard(devices, axis).contiguous().realize().bitcast(dtype)
+
+    if stack:
+      if nbytes % 4 == 0 and all(d.numel() % 4 == 0 and o % 4 == 0 for d, o in slices):
+        words = Tensor.arange(nbytes // 4, dtype=dtypes.int64).to(devices[0])
+        chunks = []
+        for data in {id(d): d for d, _ in slices}.values():
+          offsets = Tensor([-1 if d is not data else o // 4 for d, o in slices], device=devices[0], dtype=dtypes.int64)
+          gathered = data.bitcast(dtypes.uint32)[offsets.clamp(min_=0)[:, None] + words[None, :]]
+          chunks.append((offsets >= 0)[:, None].where(gathered, 0))
+        w = Tensor.usum(*chunks).bitcast(dtypes.uint8)
+      else:
+        w = Tensor.stack(*[d[o:o+nbytes] for d, o in slices])
+      w = w.reshape(byte_shape).clone().realize().shard(devices, axis).contiguous()
+    else:
+      data, o = slices[0]
+      w = data[o:o+nbytes].reshape(byte_shape).shard(devices, axis).clone()
+    return w.realize().bitcast(dtype)
 
   use_fnuz = dtypes.fp8e4m3fnuz in Device[devices[0]].renderer.supported_dtypes()
   state_dict:dict[str, Tensor] = {}

tinygrad_repo/tinygrad/llm/kernels/amd.py

@@ -146,7 +146,7 @@ class ExpertWeights:
     self.weight = Tensor.zeros(num_experts, out_features, in_features)
     if bias: self.bias = Tensor.zeros(num_experts, out_features)
   def set_quantized(self):
-    if self.weight.uop.axis == 0 or (found:=find_packed(self.weight, self.in_features)) is None: return
+    if (found:=find_packed(self.weight, self.in_features)) is None: return
     self.ggml_type, self.packed = found
   def __call__(self, sel:Tensor, x:Tensor) -> Tensor:
     # sel: (B, T, k), x: (B, T, 1, in) or (B, T, k, in) -> output: (B, T, k, out)
@@ -154,9 +154,10 @@ class ExpertWeights:
     if self.ggml_type is None and supported:
       self.set_quantized()
       if self.ggml_type is None: self.use_custom_quant = supported = False  # not a supported quant format
+    assert self.weight.uop.axis != 0 or (self.ggml_type is not None and supported and not hasattr(self, 'bias')), "split experts need the moe kernel"
     if self.weight_scale_inv is not None:
       ret = fp8_matmul(x.unsqueeze(-2), self.weight[sel], self.weight_scale_inv[sel], self.weight_block_size).squeeze(-2)
-    elif self.ggml_type is not None and supported and isinstance(x.shape[1], int) and x.shape[1] == 1: ret = moe_linear(self, sel, x)
+    elif self.ggml_type is not None and supported: ret = moe_linear(self, sel, x)
     else: ret = (x.unsqueeze(-2) @ self.weight[sel].transpose(-1, -2)).contiguous().squeeze(-2)
     return ret + self.bias[sel] if hasattr(self, 'bias') else ret
 
@@ -367,27 +368,32 @@ def _quant_decode_body(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, grids:tuple[UOp
 
 @functools.cache
 def _moe_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, sel:UOp, *grids:UOp,
-                       out_features:int, in_features:int, ggml_type:int, x_share:int, cdna:bool=False) -> UOp:
+                       out_features:int, in_features:int, ggml_type:int, x_share:int, cdna:bool=False, ndev:int=1) -> UOp:
   # out rows are (token, slot): the slot's expert reads its weight rows, x_share slots share an activation row
   expert_words = out_features * in_features * QUANT_SIZES[ggml_type] // 256 // raw.dtype.itemsize
-  def gate(row:UOp, base:UOp) -> UOp: return base + sel[row].cast(dtypes.weakint).maximum(0).minimum(raw.shape[0]//expert_words-1) * expert_words
+  local_experts = raw.shape[0]//expert_words
+  def expert(row:UOp) -> UOp: return sel[row].cast(dtypes.weakint) - UOp.range(ndev, 4, AxisType.DEVICE)*local_experts
+  def gate(row:UOp, base:UOp) -> UOp:
+    base = base + expert(row).maximum(0).minimum(local_experts-1) * expert_words
+    return base.valid((expert(row) >= 0) & (expert(row) < local_experts)) if ndev > 1 else base
   def xtoken(row:UOp) -> UOp: return row // x_share
   return _quant_decode_body(out, raw, xq, xd, xs, grids, out_features, in_features, ggml_type, cdna, gate, xtoken, "moe_")
 
 def moe_linear(layer:ExpertWeights, sel:Tensor, x:Tensor) -> Tensor:
+  if not isinstance(T:=sel.shape[1], int): return moe_linear(layer, sel.pad_to(sel.max_shape, value=-1), x.pad_to(x.max_shape))[..., :T, :, :]
   B, T, slots = cast(tuple[int, int, int], sel.shape[-3:]) # sel: (B, T, k), weight: (E, O, I)
   in_features, x_slots = int(x.uop.shard_shape[-1]), int(x.shape[-2]) # gate/up: (B, T, 1, D), down: (B, T, k, H)
   assert x_slots in (1, slots), f"x has {x_slots} rows per token, expected 1 or {slots}"
   local_out, global_out, ggml_type = int(layer.weight.uop.max_shard_shape[1]), int(layer.weight.shape[1]), unwrap(layer.ggml_type)
-  splits = layer.in_features // in_features
+  splits, ndev = layer.in_features // in_features, int(layer.weight.shape[0]) // int(layer.weight.uop.shard_shape[0])
   assert (layer.weight.uop.axis == 2) == (splits > 1), f"input features split over {splits} devices, weight split on axis {layer.weight.uop.axis}"
   xq, xd, xs = q8_quantize(x.contiguous(), B*T*x_slots, in_features)
-  out = Tensor.empty(B*T*slots, global_out, splits*((in_features+1023)//1024), dtype=dtypes.float32, device=x.device, axis=layer.weight.uop.axis)
+  out = Tensor.empty(ndev*B*T*slots, global_out, splits*((in_features+1023)//1024), dtype=dtypes.float32, device=x.device, axis=layer.weight.uop.axis)
   fxn = functools.partial(_moe_decode_kernel, out_features=local_out, in_features=in_features, ggml_type=ggml_type, x_share=slots//x_slots,
-                          cdna=_is_cdna(x.device))
+                          cdna=_is_cdna(x.device), ndev=ndev)
   extra = (_iq_grid(x.device, ggml_type),) if ggml_type in (IQ2_XS, IQ3_XXS, IQ3_S, IQ2_S) else ()
   result = Tensor.custom_kernel(out, layer.packed, xq, xd, xs, sel.cast(dtypes.int32).flatten(), *extra, fxn=fxn)[0].sum(-1)
-  return result.reshape(B, T, slots, global_out)
+  return result.reshape(ndev, B, T, slots, global_out) if ndev > 1 else result.reshape(B, T, slots, global_out)
 
 # kgroups: groups of 16 lanes that split k in a 16x16x16 fragment, RDNA3 copies a and b in both halves
 def _wmma_words(lane_group:UOp, kgroups:int) -> tuple[UOp|int, ...]:

tinygrad_repo/tinygrad/llm/model.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass, replace
 from typing import Callable
 from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes, Device
 from tinygrad.llm.kernels.amd import Linear, ExpertWeights, fp8_matmul, gated_delta_prefill, flash_attention, amd_custom_kernels_supported
-from tinygrad.llm.gguf import gguf_parse, gguf_shard
+from tinygrad.llm.gguf import gguf_parse, gguf_shard, _GGML_QUANT
 from tinygrad.llm.hf import hf_load
 from tinygrad.uop.ops import resolve
 from tinygrad.helpers import unwrap, round_up, get_child
@@ -36,7 +36,9 @@ def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str|
   freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
   return (freqs.cos() * concentration).cat(freqs.sin() * concentration, dim=-1).clone(device)
 
-def gated_activation(gate:Tensor, up:Tensor, *, alpha:float=1.0, limit:float|None=None, up_bias:float=0.0) -> Tensor:
+def gated_activation(gate:Tensor, up:Tensor, *, alpha:float=1.0, limit:float|None=None,
+                     up_bias:float=0.0, situ:tuple[float, float]|None=None) -> Tensor:
+  if situ is not None: return (gate / situ[0]).tanh() * situ[0] * gate.sigmoid() * (up / situ[1]).tanh() * situ[1]
   if limit is not None and limit > 0: gate, up = gate.clamp(max_=limit), up.clamp(-limit, limit)
   # TODO: remove the need for this contiguous
   return gate * (gate * alpha).sigmoid().contiguous() * (up + up_bias)
@@ -70,6 +72,8 @@ class SSMConfig:
   inner_size: int
   kda: bool = False
   gate_lower_bound: float|None = None  # kda: bounded decay gate
+  full_rank_gate: bool = False
+  silu_gate: bool = False
 
 @dataclass(frozen=True)
 class HCConfig:  # manifold-constrained hyper-connections: count residual streams, mixed by a doubly stochastic matrix
@@ -131,6 +135,23 @@ class TransformerConfig:
   indexer: IndexerConfig|None = None
   indexer_shared: tuple[bool, ...] = ()
   rope_interleave: bool = False
+  expert_latent_dim: int = 0
+  situ: tuple[float, float]|None = None
+  attn_res_block_size: int = 0
+  attn_res_new_block: bool = False
+
+class AttentionResidual:
+  def __init__(self, config:TransformerConfig, new_block:bool=False):
+    self.weight, self.norm_eps, self.new_block = Tensor.zeros(config.dim), config.norm_eps, new_block
+
+  def mix(self, xs:Tensor) -> Tensor:
+    scores = (xs.rmsnorm(eps=self.norm_eps) * self.weight).sum(-1, keepdim=True)
+    return (scores.softmax(2) * xs).sum(2)  # xs: (B, T, N, D) streams: (B, T, D)
+
+  def __call__(self, xs:Tensor, sublayer:Callable[[Tensor], tuple[Tensor, Tensor|None]]) -> tuple[Tensor, Tensor|None]:
+    out, aux = sublayer(self.mix(xs))
+    # (B, T, N+1, D) if new_block else (B, T, N, D)
+    return (xs.cat(out.unsqueeze(2), dim=2) if self.new_block else xs[:, :, :-1].cat(xs[:, :, -1:] + out.unsqueeze(2), dim=2)), aux
 
 class HyperConnection:
   def __init__(self, config:TransformerConfig):
@@ -157,19 +178,25 @@ class FFNBlock:
     self.attn_norm   = nn.RMSNorm(config.dim, config.norm_eps)
     self.ffn_norm    = nn.RMSNorm(config.dim, config.norm_eps)
     if config.hc is not None: self.hc_attn, self.hc_ffn = HyperConnection(config), HyperConnection(config)
+    if config.attn_res_block_size:
+      self.attn_res_score, self.ffn_res_score = AttentionResidual(config, config.attn_res_new_block), AttentionResidual(config)
 
     # --- feed-forward (MoE or dense) -------------------------------------
     if config.num_experts > 0:
       self.ffn_gate_inp = Linear(config.dim, config.num_experts, bias=config.expert_proj_bias)  # router
       if config.expert_bias: self.exp_probs_b = {"bias": Tensor.zeros(config.num_experts)}
-      self.ffn_gate_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim, bias=config.expert_proj_bias)
-      self.ffn_up_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim, bias=config.expert_proj_bias)
-      self.ffn_down_exps = ExpertWeights(config.num_experts, config.hidden_dim, config.dim, bias=config.expert_proj_bias)
+      expert_dim = config.expert_latent_dim or config.dim
+      self.ffn_gate_exps = ExpertWeights(config.num_experts, expert_dim, config.hidden_dim, bias=config.expert_proj_bias)
+      self.ffn_up_exps = ExpertWeights(config.num_experts, expert_dim, config.hidden_dim, bias=config.expert_proj_bias)
+      self.ffn_down_exps = ExpertWeights(config.num_experts, config.hidden_dim, expert_dim, bias=config.expert_proj_bias)
       if config.shared_expert_dim > 0:
         self.ffn_gate_shexp = Linear(config.dim, config.shared_expert_dim, bias=False)
         self.ffn_up_shexp = Linear(config.dim, config.shared_expert_dim, bias=False)
         self.ffn_down_shexp = Linear(config.shared_expert_dim, config.dim, bias=False)
         if config.shared_expert_gate: self.ffn_gate_inp_shexp = {"weight": Tensor.zeros(config.dim)}
+      if config.expert_latent_dim:
+        self.ffn_routed_down, self.ffn_routed_up = Linear(config.dim, expert_dim, bias=False), Linear(expert_dim, config.dim, bias=False)
+        self.ffn_routed_norm = nn.RMSNorm(expert_dim, config.norm_eps)
     else:
       self.ffn_gate    = Linear(config.dim, config.hidden_dim, bias=False)
       self.ffn_up      = Linear(config.dim, config.hidden_dim, bias=False)
@@ -177,7 +204,7 @@ class FFNBlock:
 
   def _feed_forward(self, x:Tensor) -> Tensor:
     if hasattr(self, 'ffn_gate_exps'):
-      h = x.unsqueeze(2)  # (B, T, 1, D) - add expert dim for broadcasting
+      h = (self.ffn_routed_down(x) if self.config.expert_latent_dim else x).unsqueeze(2)  # (B, T, 1, D) - add expert dim for broadcasting
       logits = self.ffn_gate_inp(x)
       bias = self.exp_probs_b["bias"] if hasattr(self, 'exp_probs_b') else None
       gating, normalize_topk = self.config.expert_gating_func, self.config.norm_topk_prob
@@ -196,15 +223,19 @@ class FFNBlock:
       if normalize_topk: probs = probs / probs.sum(axis=-1, keepdim=True)
       probs = probs * self.config.routed_scaling_factor
       gate, up = self.ffn_gate_exps(sel, h), self.ffn_up_exps(sel, h)
-      act = gated_activation(gate, up, alpha=self.config.swiglu_alpha, limit=self.config.swiglu_limit, up_bias=self.config.swiglu_up_bias)
-      x_down = self.ffn_down_exps(sel, act.contiguous())  # (B, T, k, D)
-      out = (x_down * probs.unsqueeze(-1)).sum(axis=2)  # (B, T, D)
+      act = gated_activation(gate, up, alpha=self.config.swiglu_alpha, limit=self.config.swiglu_limit,
+                             up_bias=self.config.swiglu_up_bias, situ=self.config.situ)
+      x_down = self.ffn_down_exps(sel, act.contiguous())  # ([ndev,] B, T, k, D)
+      out = (x_down * probs.unsqueeze(-1)).sum(axis=-2)  # ([ndev,] B, T, D)
+      if out.ndim == x.ndim + 1: out = out.sum(0)
+      if self.config.expert_latent_dim: out = self.ffn_routed_up(self.ffn_routed_norm(out))
       if hasattr(self, 'ffn_gate_shexp'):
-        shexp = self.ffn_down_shexp(gated_activation(self.ffn_gate_shexp(x), self.ffn_up_shexp(x), limit=self.config.swiglu_limit).contiguous())
+        shexp = self.ffn_down_shexp(gated_activation(self.ffn_gate_shexp(x), self.ffn_up_shexp(x),
+                                                        limit=self.config.swiglu_limit, situ=self.config.situ).contiguous())
         if hasattr(self, 'ffn_gate_inp_shexp'): shexp = shexp * (x * self.ffn_gate_inp_shexp["weight"]).sum(axis=-1, keepdim=True).sigmoid()
         out = out + shexp
       return out
-    return self.ffn_down(gated_activation(self.ffn_gate(x), self.ffn_up(x), limit=self.config.swiglu_limit))
+    return self.ffn_down(gated_activation(self.ffn_gate(x), self.ffn_up(x), limit=self.config.swiglu_limit, situ=self.config.situ))
 
   # given the token-prefix match, return how much cached state this block can still reuse
   def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return prefix_len
@@ -225,6 +256,7 @@ class FFNBlock:
         return self._attention(h, start_pos, mask), mask
       res_attn = res_ffn = lambda xs, sublayer: (xs + (ret:=sublayer(xs))[0], ret[1])
       if self.config.hc is not None: res_attn, res_ffn = self.hc_attn, self.hc_ffn
+      elif self.config.attn_res_block_size: res_attn, res_ffn = self.attn_res_score, self.ffn_res_score
       x, idx_mask = res_attn(x, attend)
       out = res_ffn(x, lambda h: (self._feed_forward(self.ffn_norm(h)), None))[0].contiguous()
       return (out, idx_mask) if idx_mask is not None else out
@@ -355,6 +387,7 @@ class MLATransformerBlock(FFNBlock):
     self.attn_v_b = {"weight": Tensor.zeros(config.n_heads, config.v_head_dim, config.kv_lora_rank)}
     self.attn_output = Linear(config.n_heads * config.v_head_dim, config.dim, bias=False)
     self.indexer = Indexer(config) if config.indexer else None
+    if config.attn_output_gate: self.attn_gate = Linear(config.dim, config.n_heads * config.v_head_dim, bias=False)
 
   def _attention(self, x:Tensor, start_pos:int|UOp, idx_mask:Tensor|None=None) -> Tensor:
     B, T, _ = x.shape
@@ -385,7 +418,7 @@ class MLATransformerBlock(FFNBlock):
       attn = fp8_matmul(attn, self.attn_v_b['weight'], self.attn_v_b['weight_scale_inv'], self.kv_b_block_size)
     else: attn = attn @ self.attn_v_b['weight'].transpose(-1, -2)
     attn = attn.transpose(1, 2).reshape(B, T, -1)
-    return self.attn_output(attn)
+    return self.attn_output(attn * self.attn_gate(x).sigmoid() if self.config.attn_output_gate else attn)
 
   def _attn_mask(self, x:Tensor, start_pos:int|UOp, idx_mask:Tensor|None=None) -> Tensor|None:  # (1, 1, T, start_pos+T), None: everything is visible
     if idx_mask is not None: return idx_mask
@@ -416,7 +449,8 @@ class GatedDeltaNetBlock(FFNBlock):
       self.qkv_dims = (self.q_dim, self.q_dim, ssm.inner_size)
       self.attn_q, self.attn_k, self.attn_v = (Linear(config.dim, d, bias=False) for d in self.qkv_dims)
       self.ssm_conv1d_q, self.ssm_conv1d_k, self.ssm_conv1d_v = ({"weight": Tensor.zeros(d, 1, ssm.conv_kernel)} for d in self.qkv_dims)
-      self.ssm_g_a, self.ssm_g_b = Linear(config.dim, self.head_v_dim, bias=False), Linear(self.head_v_dim, ssm.inner_size, bias=False)
+      if ssm.full_rank_gate: self.ssm_g = Linear(config.dim, ssm.inner_size, bias=False)
+      else: self.ssm_g_a, self.ssm_g_b = Linear(config.dim, self.head_v_dim, bias=False), Linear(self.head_v_dim, ssm.inner_size, bias=False)
       self.ssm_f_a, self.ssm_f_b = Linear(config.dim, self.head_k_dim, bias=False), Linear(self.head_k_dim, ssm.inner_size, bias=False)
     else:
       self.attn_qkv = Linear(config.dim, self.conv_channels, bias=False)
@@ -439,7 +473,7 @@ class GatedDeltaNetBlock(FFNBlock):
 
     # input processing
     x = x.half()
-    out_gate = self.ssm_g_b(self.ssm_g_a(x)) if is_kda else self.attn_gate(x)
+    out_gate = (self.ssm_g(x) if hasattr(self, "ssm_g") else self.ssm_g_b(self.ssm_g_a(x))) if is_kda else self.attn_gate(x)
     out_gate = out_gate.reshape(B, T, self.num_v_heads, self.head_v_dim)
     beta = self.ssm_beta(x).sigmoid().reshape(B, T, self.num_v_heads)
     alpha = self.ssm_f_b(self.ssm_f_a(x)) if is_kda else self.ssm_alpha(x)
@@ -495,7 +529,7 @@ class GatedDeltaNetBlock(FFNBlock):
       core = Tensor(outs[0].stack(*outs[1:], dim=1).contiguous().uop.after(state_store, *conv_writes))
 
     # output; undo the padding before the output projection
-    z = (self.ssm_norm(core) * (out_gate.sigmoid() if is_kda else out_gate.silu())).cast(x.dtype).contiguous()
+    z = (self.ssm_norm(core) * (out_gate.silu() if unwrap(self.config.ssm).silu_gate else out_gate.sigmoid())).cast(x.dtype).contiguous()
     if symbolic: z = z[:, :T]
     return self.ssm_out(z.reshape(B, T, -1))
 
@@ -518,9 +552,11 @@ class Transformer:
       c = dense_config if i < config.leading_dense_blocks else config
       if config.sliding_window_pattern != 0 and (i+1) % config.sliding_window_pattern == 0: c = replace(c, sliding_window=0)
       if config.indexer_shared and config.indexer_shared[i]: c = replace(c, indexer=None)
+      if config.attn_res_block_size and i % config.attn_res_block_size == 0: c = replace(c, attn_res_new_block=True)
       self.blk.append(GatedDeltaNetBlock(c, config.ssm) if config.ssm and config.ssm_layers[i] else block_cls(c))
     self.token_embd  = nn.Embedding(config.vocab_size, config.dim)
     self.output_norm = nn.RMSNorm(config.dim, config.norm_eps)
+    if config.attn_res_block_size: self.output_res_score = AttentionResidual(config)
     self.output = Linear(config.dim, config.vocab_size, bias=False)
     self.max_context = config.max_context
     self.has_recurrent_block = any(isinstance(b, GatedDeltaNetBlock) for b in self.blk)
@@ -532,9 +568,11 @@ class Transformer:
   def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
     x = self.token_embd(tokens.to(self.token_embd.weight.device)).float()  # (B, T, D)
     if (hc:=self.config.hc) is not None: x = x.unsqueeze(2).expand(*x.shape[:2], hc.count, x.shape[-1]).contiguous()
+    elif self.config.attn_res_block_size: x = x.unsqueeze(2)
     idx_mask = None
     for block in self.blk: x, idx_mask = block(x, start_pos, idx_mask)
     if hc is not None: x = x.mean(2)
+    elif self.config.attn_res_block_size: x = self.output_res_score.mix(x)
     # only run the output projection on the last token
     logits = self.output(self.output_norm(x[:, -1:]))[:, -1, :].to(tokens.device)
     # Gumbel-max trick: argmax(logits/temp - log(-log(uniform))) is equivalent to sampling from softmax(logits/temp)
@@ -559,11 +597,12 @@ class Transformer:
       rules = {**{w: 0 for w in ('token_embd.weight', 'output.weight', 'attn_q.weight', 'attn_k.weight', 'attn_v.weight', 'ffn_gate.weight',
         'ffn_up.weight', 'attn_q_b.weight', 'attn_k_b.weight', 'attn_v_b.weight')},
         **{w: 1 for w in ('attn_output.weight', 'ffn_down.weight', 'ffn_gate_exps.weight', 'ffn_up_exps.weight')}, 'ffn_down_exps.weight':2}
-      if arch in ('kimi-linear', 'glm5-next'): rules.update({w: 0 for w in ('ssm_conv1d_q.weight', 'ssm_conv1d_k.weight', 'ssm_conv1d_v.weight',
-        'ssm_dt.bias', 'ssm_a', 'ssm_beta.weight', 'ssm_f_b.weight', 'ssm_g_b.weight')})
+      if arch in ('kimi-linear', 'glm5-next', 'kimi-k3'): rules.update({w: 0 for w in ('ssm_conv1d_q.weight', 'ssm_conv1d_k.weight',
+        'ssm_conv1d_v.weight', 'ssm_dt.bias', 'ssm_a', 'ssm_beta.weight', 'ssm_f_b.weight', 'ssm_g.weight', 'ssm_g_b.weight', 'attn_gate.weight')})
       shard_map = {name: rules[k] for name in entries if (k:=re.sub(r"^blk\.\d+\.", "", name)) in rules}
     devices = tuple(Device.canonicalize(f'{Device.DEFAULT}:{i}') for i in range(shard))
-    state_dict = gguf_shard(entries, devices, shard_map)
+    ep = any(shape[-1] % (shard*_GGML_QUANT.get(typ, (1,))[0]) for k,(_,shape,typ) in entries.items() if 'ffn_down_exps.' in k)
+    state_dict = gguf_shard(entries, devices, {k: 0 if ep and '_exps.' in k else v for k,v in shard_map.items()})
 
     # all state items should be float16, not float32
     state_dict = {k:v.cast('float16') if getenv("HALF", 1) else v for k,v in state_dict.items()}
@@ -571,8 +610,8 @@ class Transformer:
     # some models like Llama 3.2 don't have an output.weight, they just tie to the token_embd.weight
     if 'output.weight' not in state_dict: state_dict['output.weight'] = state_dict['token_embd.weight']
 
-    hc = HCConfig(kv[f'{arch}.hyper_connection.count'], kv[f'{arch}.hyper_connection.sinkhorn_iterations'],
-                  kv[f'{arch}.hyper_connection.epsilon']) if f'{arch}.hyper_connection.count' in kv else None
+    hc = HCConfig(kv[f'{arch}.hyper_connection.count'], kv.get(f'{arch}.hyper_connection.sinkhorn_iterations', 0),
+                  kv.get(f'{arch}.hyper_connection.epsilon', 0.0)) if f'{arch}.hyper_connection.count' in kv else None
     if hc is not None: state_dict = {k.replace('.hc_attn_', '.hc_attn.').replace('.hc_ffn_', '.hc_ffn.'):v for k,v in state_dict.items()}
 
     indexer = IndexerConfig(kv[f'{arch}.attention.indexer.head_count'], kv[f'{arch}.attention.indexer.key_length'],
@@ -586,13 +625,13 @@ class Transformer:
     ssm = None
     ssm_layers: tuple[bool, ...] = ()
     if arch in ('qwen35', 'qwen35moe'):
-      ssm = SSMConfig(**{k: kv[f'{arch}.ssm.{k}'] for k in ('conv_kernel','state_size','group_count','time_step_rank','inner_size')})
+      ssm = SSMConfig(**{k: kv[f'{arch}.ssm.{k}'] for k in ('conv_kernel','state_size','group_count','time_step_rank','inner_size')}, silu_gate=True)
       ssm_layers = tuple((i+1) % kv[f'{arch}.full_attention_interval'] != 0 for i in range(num_blocks))
-    elif arch in ('kimi-linear', 'glm5-next'):
+    elif arch in ('kimi-linear', 'glm5-next', 'kimi-k3'):
       ssm_layers = tuple(x == 0 for x in n_kv_heads)[:num_blocks]
       n_kv_heads = max(n_kv_heads)
       ssm = SSMConfig(kv[f'{arch}.ssm.conv_kernel'], kv[f'{arch}.kda.head_dim'], n_heads, n_heads, n_heads*kv[f'{arch}.kda.head_dim'], kda=True,
-                      gate_lower_bound=kv.get(f'{arch}.kda.gate_lower_bound'))
+                      gate_lower_bound=kv.get(f'{arch}.kda.gate_lower_bound'), full_rank_gate=arch == 'kimi-k3')
       for i, is_ssm in enumerate(ssm_layers):
         if not is_ssm: continue
         state_dict[f"blk.{i}.ssm_a"] = state_dict[f"blk.{i}.ssm_a"].reshape(n_heads, 1)
@@ -610,7 +649,7 @@ class Transformer:
 
     # Permute RoPE weights from interleaved to half-split layout.
     for name in state_dict:
-      if arch in ('kimi-linear', 'glm5-next'): continue
+      if arch in ('kimi-linear', 'glm5-next', 'kimi-k3'): continue
       if ('attn_q.weight' in name or 'attn_q_b.weight' in name) and (arch == 'llama' or kv_lora_rank):
         state_dict[name] = permute_rope(state_dict[name], n_heads, rope_dim)
       elif arch == 'llama' and 'attn_k.weight' in name:
@@ -638,7 +677,7 @@ class Transformer:
         kv.get(f'{arch}.expert_shared_count', 0) * kv.get(f'{arch}.expert_feed_forward_length', 0)),
       shared_expert_gate=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.ffn_gate_inp_shexp.weight" in state_dict,
       dense_hidden_dim=kv.get(f'{arch}.feed_forward_length', 0) if kv.get(f'{arch}.leading_dense_block_count', 0) else 0,
-      routed_scaling_factor=kv.get(f'{arch}.expert_weights_scale', 1.0), attn_output_gate=arch in ('qwen35', 'qwen35moe'), ssm=ssm,
+      routed_scaling_factor=kv.get(f'{arch}.expert_weights_scale', 1.0), attn_output_gate=arch in ('qwen35', 'qwen35moe', 'kimi-k3'), ssm=ssm,
       ssm_layers=ssm_layers,
       qkv_bias='blk.0.attn_q.bias' in state_dict,
       expert_bias=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.exp_probs_b.bias" in state_dict,
@@ -646,7 +685,9 @@ class Transformer:
       swiglu_alpha=1.702 if arch == 'gpt-oss' else 1.0, swiglu_limit=7.0 if arch == 'gpt-oss' else kv.get(f'{arch}.swiglu_clamp_exp', [None])[0],
       swiglu_up_bias=1.0 if arch == 'gpt-oss' else 0.0, attn_sinks='blk.0.attn_sinks.weight' in state_dict,
       sliding_window=kv.get(f'{arch}.attention.sliding_window', 0),
-      sliding_window_pattern=kv.get(f'{arch}.attention.sliding_window_pattern', 2 if arch == 'gpt-oss' else 0), hc=hc, indexer=indexer)
+      sliding_window_pattern=kv.get(f'{arch}.attention.sliding_window_pattern', 2 if arch == 'gpt-oss' else 0), hc=hc, indexer=indexer,
+      expert_latent_dim=kv.get(f'{arch}.expert_latent_length', 0), attn_res_block_size=kv.get(f'{arch}.attn_res.block_size', 0),
+      situ=(kv[f'{arch}.activation.situ_beta'], kv[f'{arch}.activation.situ_linear_beta']) if f'{arch}.activation.situ_beta' in kv else None)
     model = Transformer(config)
     for p in (nn.state.get_parameters(model) if shard > 1 else []): p.to_(devices)
     nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused

tinygrad_repo/tinygrad/llm/serve.py

@@ -9,6 +9,13 @@ if TYPE_CHECKING:
 
 def parse_tool_call(s:str) -> tuple[str, typing.Any]|None:
   s = s.strip()
+  # Kimi-K3 format: tool="name" index="1"<|sep|><|open|>argument key="key" type="string"<|sep|>value<|close|>argument<|sep|>...
+  if (m := re.fullmatch(r'tool="([^"]+)"(?: index="\d+")?<\|sep\|>(.*)', s, re.DOTALL)):
+    pattern = r'<\|open\|>argument key="([^"]+)" type="([^"]+)"<\|sep\|>(.*?)<\|close\|>argument<\|sep\|>'
+    if re.sub(pattern, '', m[2], flags=re.DOTALL).strip(): return None
+    def unescape(v:str) -> str: return v.replace('&quot;', '"').replace('&amp;', '&')
+    try: return unescape(m[1]), {unescape(a[1]): a[3] if a[2] == "string" else json.loads(a[3]) for a in re.finditer(pattern, m[2], re.DOTALL)}
+    except json.JSONDecodeError: return None
   if s.startswith("{"):  # hermes JSON format: {"name": ..., "arguments": {...}}
     try:
       call = json.loads(s)
@@ -42,7 +49,13 @@ def normalize_messages(messages:list[dict]) -> None:
 
 class StreamRouter:
   # routes streamed output text to (field, text) deltas, keeping tool_call regions in .buf for the final parse
-  def __init__(self, reasoning:bool=False):
+  # (start think, end think, end content, tool call)
+  FORMATS = (("<think>", "</think>", "<tool_call>", r"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)"),
+             ("<|open|>think<|sep|>", "<|close|>think<|sep|><|open|>response<|sep|>", "<|close|>response<|sep|>",
+              r"<\|open\|>call \s*(.*?)\s*(?:<\|close\|>call<\|sep\|>|$)"))
+  def __init__(self, special_tokens:dict[str, int], reasoning:bool=False):
+    self.think, self.end_think, self.end_content, self.call = next((f for f in self.FORMATS if any(f[2].startswith(t) for t in special_tokens)),
+                                                                   self.FORMATS[0])
     self.buf = ""
     self.mode = "reasoning" if reasoning else "undecided"  # output inside a think block is sent as reasoning_content
   def split(self, tag:str, final:bool) -> tuple[str, bool]:
@@ -56,17 +69,17 @@ class StreamRouter:
   def route(self, piece:str, final:bool=False) -> typing.Iterator[tuple[str, str]]:
     self.buf += piece
     if self.mode == "undecided":  # decide whether the output starts with a think block
-      if not final and len(self.buf) < len("<think>") and "<think>".startswith(self.buf): return
-      self.mode, self.buf = ("reasoning", self.buf[len("<think>"):]) if self.buf.startswith("<think>") else ("content", self.buf)
+      if not final and len(self.buf) < len(self.think) and self.think.startswith(self.buf): return
+      self.mode, self.buf = ("reasoning", self.buf[len(self.think):]) if self.buf.startswith(self.think) else ("content", self.buf)
     if self.mode == "reasoning":
-      emit, done = self.split("</think>", final)
+      emit, done = self.split(self.end_think, final)
       if emit: yield "reasoning_content", emit
       if not done: return
       self.mode = "content"
     if self.mode == "tool": return
-    emit, found = self.split("<tool_call>", final)
+    emit, found = self.split(self.end_content, final)
     if emit: yield "content", emit
-    if found: self.mode, self.buf = "tool", "<tool_call>" + self.buf
+    if found: self.mode, self.buf = "tool", self.end_content + self.buf
 
 class Handler(VizHandler):
   server: LLMServer
@@ -87,7 +100,7 @@ class Handler(VizHandler):
     finish_reason = "stop"
     st = pt = time.perf_counter()
     dec = tok.stream_decoder()
-    router = StreamRouter(reasoning)
+    router = StreamRouter(tok._special_tokens, reasoning)
     def log_stats(interrupted:bool=False):
       et = time.perf_counter()
       total = f"total:{et-st:6.2f}s"
@@ -107,7 +120,7 @@ class Handler(VizHandler):
           break
       for field, delta in router.route(dec(), final=True): yield chunk({field:delta})
       tool_calls: list[dict] = []
-      for m in re.finditer(r"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)", router.buf, re.DOTALL):
+      for m in re.finditer(router.call, router.buf, re.DOTALL):
         if (parsed := parse_tool_call(m.group(1))) is None:
           stderr_log(f"failed to parse tool call: {m.group(1)[:200]}")
           yield chunk({"content":m.group(0)})  # don't silently drop output the client can't use
@@ -150,7 +163,7 @@ class Handler(VizHandler):
       max_tokens = body.get("max_completion_tokens") or body.get("max_tokens")
       chunks = self.run_model(ids, body["model"], not body.get("stream") or body.get("stream_options",{}).get("include_usage", False),
                               max_tokens=max_tokens, temperature=float(body.get("temperature", 0.0)),
-                              reasoning=rendered.rstrip().endswith("<think>"))
+                              reasoning=rendered.rstrip().endswith(tuple(f[0] for f in StreamRouter.FORMATS)))
       if body.get("stream"): self.stream_json(chunks)
       else:
         out, reasoning, tool_calls, finish_reason = [], [], [], "stop"

tinygrad_repo/tinygrad/mixin/elementwise.py

@@ -785,7 +785,7 @@ class ElementwiseMixin(CreationMixin):
     ```
     """
     if approximate == "tanh":
-      return 0.5 * self * (1 + (math.sqrt(2 / math.pi) * (self + 0.044715 * self ** 3)).tanh())
+      return self * (math.sqrt(8 / math.pi) * (self + 0.044715 * self ** 3)).sigmoid()
     elif approximate == "none":
       return self * 0.5 * (1.0 + (self / math.sqrt(2)).erf())
     else:

tinygrad_repo/tinygrad/mixin/op.py

@@ -1092,8 +1092,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
         low, high, perc = [y.reshape(reshape).expand(expand) for y in (num//den, (num+den-1)//den, (num % den).cast(dtypes.float32)/den)]
         x = x.gather(i, low).lerp(x.gather(i, high), perc)
       else:
-        scale, arr = in_sz / size[i], type(self).arange(size[i], dtype=dtypes.float32)
-        index = (scale*(arr+0.5) if mode=="nearest-exact" else scale*arr).cast(dtypes.int32).reshape(reshape).expand(expand)
+        scale, arr = in_sz / size[i], type(self).arange(1, size[i]+1, dtype=dtypes.float32) - (0.5 if mode=="nearest-exact" else 1)
+        index = (scale*arr).cast(dtypes.int32).reshape(reshape).expand(expand)
         x = x.gather(i, index)
     return x.cast(self.dtype)
 

tinygrad_repo/tinygrad/renderer/__init__.py

@@ -45,7 +45,8 @@ class Estimates:
       if u.op is Ops.RANGE:
         mult_stack.append(mults)
         if u.dtype is not dtypes.void:  # unbounded loop, unknown trip count
-          mults *= cast(sint, u.src[0].without_after.ssimplify())
+          bound = u.src[0].without_after.simplify()
+          mults *= cast(sint, bound.vmax if bound.ranges else bound.ssimplify())
           # SPECIAL are already counted in mults
           mults = mults.substitute({x:x.const_like(0) for x in mults.toposort() if x.op is Ops.SPECIAL}) if isinstance(mults, UOp) else mults
       elif u.op in {Ops.END, Ops.BACKEDGE}: mults = mult_stack.pop(-1)
@@ -75,6 +76,7 @@ class Renderer:
   shared_max: int = 32768
   tensor_cores: list[TensorCore] = []
   extra_matcher: PatternMatcher|None = None
+  packed_args: bool = False
   code_for_op: dict[Ops, Callable] = {}
 
   compiler: Compiler = Compiler()

tinygrad_repo/tinygrad/renderer/cstyle.py

@@ -41,6 +41,7 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"(({ctx._render_dtype(x.dtype, addrspace=x.addrspace)})({ctx[x.src[0]]}))"
    if x.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
   (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"__builtin_bit_cast({ctx.render_type(x)}, ({ctx.render_type(x.src[0])})({ctx[x.src[0]]}))"),
+  (UPat(Ops.REFPARAM, name="x"), lambda ctx,x: f"(({ctx.param_type(x)})({ctx[x.src[1]]}))"),
 
   # GPU stuff
   (UPat(Ops.BARRIER), lambda ctx: ctx.barrier),
@@ -71,8 +72,7 @@ base_rewrite = PatternMatcher([
 
   # call to PROGRAM links BINARY at runtime
   (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="body"),), allow_any_len=True, name="x"), lambda ctx,x,body:
-   f"{body.src[0].arg.function_name}(&(struct {body.src[0].arg.function_name}_args){{"
-   f"{', '.join(f'({ctx.param_type(p)}){ctx[x.src[p.arg.slot+1]]}' for p in body.arg.params)}}});"),
+   f"{body.src[0].arg.function_name}({', '.join(f'({ctx.param_type(p)}){ctx[x.src[p.arg.slot+1]]}' for p in body.arg.params)});"),
 
   # call to LINEAR renders the call in SOURCE
   (UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="body"),), allow_any_len=True, name="x"), lambda ctx,x,body: f"{body.arg}("
@@ -139,7 +139,6 @@ class CStyleLanguage(Renderer):
   var_suffix: str = ""
   barrier: str = ""
   code_for_workitem: dict[Literal["g", "l"], Callable] = {}
-  args_struct_prefix: str|None = None
   extra_args: list[str] = []
   float4: str|None = None
   float4_style: tuple[str, str] = ('(', ')')
@@ -168,12 +167,8 @@ class CStyleLanguage(Renderer):
     if any(is_image_shape(u._shape) for _,(u,_) in bufs):
       tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n"
     buftypes = [(name, ("volatile " if u.arg.volatile else "")+(self.var_prefix if u.addrspace == AddrSpace.ALU else "")+
-                       self._render_dtype(u.dtype, sz=1, addrspace=u.addrspace, mutable=mutable, shape=u._shape)+
+                       self._render_dtype(u.dtype, addrspace=u.addrspace, mutable=mutable, override_ptr=u.addrspace != AddrSpace.ALU, shape=u._shape)+
                        (self.var_suffix if u.addrspace == AddrSpace.ALU else self.buffer_suffix)) for name,(u,mutable) in bufs]
-    if self.args_struct_prefix is not None:
-      prefix = (prefix or []) + [f"struct {function_name}_args {{ " + "".join(f"{self.param_type(u)} arg{u.arg.slot}; " for _,(u,_) in bufs) + "};"]
-      kernel = [f"  {self.param_type(u)} {n} = args->arg{u.arg.slot};" for n,(u,_) in bufs] + kernel
-      buftypes = [("args", f"{self.args_struct_prefix} struct {function_name}_args*")]
     local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
     launch_bounds = prod([d.vmax for d in local_dims])
     prg = ''.join([f"{self.kernel_typedef.format(launch_bounds=launch_bounds)} {function_name}(",] +
@@ -263,7 +258,7 @@ class CStyleLanguage(Renderer):
 
       if u.op in {Ops.ENDIF, Ops.END, Ops.BACKEDGE}: depth -= 1
       if (u.op is not Ops.CAST or u.max_numel() == 1) and ((u.op is Ops.CAST and u.src[0].op is Ops.CONST) or \
-        u.op in {Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
+        u.op in {Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI, Ops.REFPARAM} or \
         (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \
         (u.op in {Ops.CAST, Ops.BITCAST} and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \
         (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
@@ -282,11 +277,9 @@ class CStyleLanguage(Renderer):
   def param_type(self, p:UOp): return "volatile "*p.arg.volatile + self._render_dtype(p.dtype, 1, p.addrspace, True, p.addrspace != AddrSpace.ALU)
   def render(self, uops:list[UOp]) -> str:
     bodies = dedup(u.body for u in UOp.sink(*uops).toposort(enter_calls=bool(self.inline_prefix)) if u.op is Ops.CALL)
-    fns = {b:b.src[0].arg.function_name for b in bodies if b.op is Ops.PROGRAM}
     inlines = {b:self._render(list(b.src)) for b in bodies if b.op is Ops.LINEAR}
     prefix = [f"extern void {b.arg.name}();" if b.op is Ops.CUSTOM_FUNCTION else
-              f"struct {fns[b]}_args {{ " + ''.join(f"{self.param_type(p)} arg{p.arg.slot}; " for p in b.arg.params) +
-              f"}}; extern void {fns[b]}(const struct {fns[b]}_args*);" if b.op is Ops.PROGRAM else
+              f"extern void {b.src[0].arg.function_name}({', '.join(self.param_type(p) for p in b.arg.params)});" if b.op is Ops.PROGRAM else
               f"static inline {self.inline_prefix} void {b.arg}({', '.join(f'{self.param_type(p)} {n}' for n,(p,_) in inlines[b][2])}) {{\n" +
               '\n'.join(inlines[b][1]) + "\n}" for b in bodies]
     return self.render_kernel(*self._render(uops), uops+[u for b in inlines for u in b.src], prefix or None)
@@ -297,7 +290,7 @@ class ClangRenderer(CStyleLanguage):
   gep_arr_threshold = 0
   has_local = False
   global_max = (1, 0, 0)
-  args_struct_prefix = "const"
+  packed_args = True
   infinity = "__builtin_inff()"
   nan = '__builtin_nanf("")'
 
@@ -383,7 +376,7 @@ class MetalRenderer(CStyleLanguage):
   barrier = "threadgroup_barrier(mem_flags::mem_threadgroup);"
   float4 = "float4"
   code_for_workitem = {"g": lambda x: f"gid.{chr(120+int(x))}", "l": lambda x: f"lid.{chr(120+int(x))}"}
-  args_struct_prefix = "constant"
+  packed_args = True
   # uint3 used for gid/lid - TODO: this should probably be `ushort3 lid [[thread_position_in_threadgroup]]`
   extra_args = ['uint3 gid [[threadgroup_position_in_grid]]', 'uint3 lid [[thread_position_in_threadgroup]]']
   type_map = {**CStyleLanguage.type_map, dtypes.u32: "uint", dtypes.bf16: "bfloat"}

tinygrad_repo/tinygrad/renderer/llvmir.py

@@ -7,7 +7,7 @@ from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str
 from tinygrad.dtype import dtypes, float_to_fp8, DType, truncate, AddrSpace
 from tinygrad.helpers import prod, Target
 
-def is_volatile(u:UOp) -> bool: return (buf:=u.buf_uop).op is Ops.PARAM and buf.arg.volatile
+def is_volatile(u:UOp) -> bool: return (buf:=u.buf_uop).op in {Ops.PARAM, Ops.REFPARAM} and buf.arg.volatile
 
 def ldt(dt:DType, count=1, ptr=False):
   if ptr: return ldt(dt, count) + "*"
@@ -112,6 +112,7 @@ base_rewrite = PatternMatcher([
    lambda ctx, x: f"  {ctx[x]} = call {ldt(x.dtype)} @llvm.{llvm_intrinsics[x.op]}.{ldt(x.dtype)}({ldt(x.src[0].dtype)} {ctx[x.src[0]]})"),
   (UPat(Ops.BITCAST, name="x"), lambda ctx,x:
    f"  {ctx[x]} = bitcast {ldt(x.src[0].dtype, x.src[0].max_numel())} {ctx[x.src[0]]} to {ldt(x.dtype, x.max_numel())}"),
+  (UPat(Ops.REFPARAM, name="x"), lambda ctx,x: f"  {ctx[x]} = inttoptr i64 {ctx[x.src[1]]} to {ldt(x.dtype, ptr=True)}"),
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"  {ctx[x]} = {lcast(x.src[0].dtype, x.dtype)} {ldt(x.src[0].dtype)} {ctx[x.src[0]]} to {ldt(x.dtype)}"),
   (UPat(Ops.TRUNC, name="x"),
    lambda ctx,x: f"  {ctx[x]} = call {ldt(x.dtype)} @llvm.trunc.{ldt(x.dtype)}({ldt(x.src[0].dtype)} {ctx[x.src[0]]})"),
@@ -168,20 +169,15 @@ class LLVMRenderer(Renderer):
     super().__init__(target)
     self.extra_matcher = create_non_native_float_pats((dtypes.bfloat16,))+ \
       pm_manual_bf16_cast+pm_widen_uidx if target.arch != "gfx950" else pm_widen_uidx
-  def _render_fn(self, name:str, args:list[tuple[str,UOp]], kernel:list[str], prefix:list[str]|None=None, args_struct=False) -> str:
+  def _render_fn(self, name:str, args:list[tuple[str,UOp]], kernel:list[str], prefix:list[str]|None=None) -> str:
     # Buffer views may start at an offset from the aligned allocation.
     tys = [ldt(u.dtype, ptr=u.addrspace != AddrSpace.ALU) for _,u in args]
     sargs = ", ".join(f"{t}{' noalias' if u.addrspace == AddrSpace.GLOBAL else ''} {name}" for t,(name,u) in zip(tys, args))
-    if args_struct:
-      args_t = "{ " + ", ".join(tys) + " }"
-      kernel = [f"  %args_t = load {args_t}, ptr %args"] + [f"  {n} = extractvalue {args_t} %args_t, {i}" for i,(n,_) in enumerate(args)] + kernel
-      sargs = "ptr %args"
     return "\n".join((prefix or []) + [f"define{' ' + self.abi if self.abi else ''} void @{name}({sargs}) #0", "{"] + kernel + ["  ret void\n}"])
-  def _render_kernel(self, uops: list[UOp], prefix:list[str]|None=None, name="test", args_struct=False) -> tuple[tuple[str, ...], str]:
+  def _render_kernel(self, uops: list[UOp], prefix:list[str]|None=None, name="test") -> tuple[tuple[str, ...], str]:
     r: dict[UOp, str] = {}
     args: list[tuple[str, UOp]] = []
     kernel: list[str] = []
-    call_args: list[str] = []
     vc = -1
 
     local_args: list[str] = []
@@ -209,15 +205,8 @@ class LLVMRenderer(Renderer):
           kernel.append(f"  {r[u]} = alloca [{size} x {ldt(u.dtype)}], align 16")
       elif u.op is Ops.BINARY: r[u] = f"@bin_{u.key.hex()}"
       elif u.op is Ops.CALL and u.body.op is Ops.PROGRAM:
-        vc += 1
-        types = [ldt(p.dtype, ptr=p.addrspace != AddrSpace.ALU) for p in u.body.arg.params]
-        arg_type = "{ " + ", ".join(types) + " }"
-        r[u] = f"%v{vc}"
-        call_args.append(f"  {r[u]} = alloca {arg_type}")
-        for i, (p, ty) in enumerate(zip(u.body.arg.params, types)):
-          kernel.append(f"  {r[u]}_{i} = getelementptr inbounds {arg_type}, ptr {r[u]}, i32 0, i32 {i}")
-          kernel.append(f"  store {ty} {r[u.src[p.arg.slot+1]]}, ptr {r[u]}_{i}")
-        kernel.append(f"  call void @{u.body.src[0].arg.function_name}(ptr {r[u]})")
+        call_args = ', '.join(f'{lparam(p)} {r[u.src[p.arg.slot+1]]}' for p in u.body.arg.params)
+        kernel.append(f"  call void @{u.body.src[0].arg.function_name}({call_args})")
       elif u.op is Ops.CAST and u.src[0].op is Ops.CONST: r[u] = lconst(u.src[0].val, u.dtype)
       elif (u.op is Ops.CAST and ldt(u.dtype) == ldt(u.src[0].dtype)) or (u.op is Ops.BITCAST and u.addrspace not in (None, AddrSpace.ALU)):
         r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast
@@ -230,18 +219,19 @@ class LLVMRenderer(Renderer):
         if l is None:
           raise RuntimeError(f"failed to render {u.op} with {u.dtype} srcs {[x.dtype for x in u.src]}")
         kernel.append(l)
-    return tuple(local_args), self._render_fn(name, args, call_args + kernel, prefix, args_struct)
+    return tuple(local_args), self._render_fn(name, args, kernel, prefix)
 
 class CPULLVMRenderer(LLVMRenderer):
   has_local = False
   global_max = (1, 0, 0)
   abi = None
+  packed_args = True
   string_rewrite = base_rewrite
-  def render(self, uops: list[UOp]) -> str: return "\n".join((k:=self._render_kernel(uops, args_struct=True))[0] + (k[1], self._render_footer(uops)))
+  def render(self, uops: list[UOp]) -> str: return "\n".join((k:=self._render_kernel(uops))[0] + (k[1], self._render_footer(uops)))
   def _render_footer(self, uops: list[UOp]) -> str:
     decls = {x.src[0].arg.name: f"declare {ldt(x.dtype)} @{x.src[0].arg.name}({', '.join(map(lparam, x.src[1:]))})"
              for x in uops if x.op is Ops.CALL and x.src[0].op is Ops.CUSTOM_FUNCTION}
-    decls |= {x.body.src[0].arg.function_name: f"declare void @{x.body.src[0].arg.function_name}(ptr)"
+    decls |= {x.body.src[0].arg.function_name: f"declare void @{x.body.src[0].arg.function_name}({', '.join(map(lparam, x.body.arg.params))})"
               for x in uops if x.op is Ops.CALL and x.body.op is Ops.PROGRAM}
     blobs = {x.key: f"@bin_{x.key.hex()} = private constant [{len(x.arg)} x i8] c\"" + "".join("\\%02X" % b for b in x.arg) + '", align 16'
              for x in uops if x.op is Ops.BINARY}

tinygrad_repo/tinygrad/runtime/autogen/__init__.py

@@ -185,6 +185,9 @@ def __getattr__(nm):
                   args=["-Du8=unsigned char", "-Du16=unsigned short", "-Du32=unsigned int", "-Du64=unsigned long long",
                         "-D__be16=unsigned short", "-D__be32=unsigned int", "-D__be64=unsigned long long", f"-I{kh}"],
                   preprocess=_extract_deb)
+    case "nvme":
+      kh = "{}/usr/src/linux-headers-6.18.9+deb14-common/include"
+      return load("nvme", [f"{kh}/linux/nvme.h"], srcs=linux_headers_kern_deb, args=["-D__KERNEL__", f"-I{kh}"], preprocess=_extract_deb)
     case "bnxt":
       kh = "{}/usr/src/linux-headers-6.18.9+deb14-common/include"
       return load("bnxt", [f"{kh}/linux/bnxt/hsi.h", *[f"{{}}/{s.split('/')[-1].split('?')[0]}" for s in bnxt_src]],

tinygrad_repo/tinygrad/runtime/ops_dsp.py

@@ -11,6 +11,7 @@ if getenv("IOCTL"): import extra.dsp.run # noqa: F401 # pylint: disable=unused-i
 class DSPRenderer(ClangRenderer):
   buffer_suffix = " restrict __attribute__((align_value(128)))"
   kernel_typedef = "__attribute__((noinline)) void"
+  packed_args = False
   type_map = { **ClangRenderer.type_map, dtypes.uint64: "unsigned long long", dtypes.int64: "long long" }
   code_for_op = {k:v for k,v in ClangRenderer.code_for_op.items() if k != Ops.SQRT}
 
@@ -36,7 +37,7 @@ class DSPRenderer(ClangRenderer):
     msrc += [f'void *buf_{i} = HAP_mmap(0,sz_or_val_{i},3,0,pra[{j+3}].dma.fd,0)+off{i};' for j,i in enumerate(gbufs)]
     msrc += ["unsigned long long start = HAP_perf_get_time_us();"]
     fbufs = [(f'buf_{i}' if b[1][0].addrspace == AddrSpace.GLOBAL else f'sz_or_val_{i}') for i,b in enumerate(bufs)]
-    msrc += [f"{function_name}(&(struct {function_name}_args){{{', '.join(fbufs)}}});"]
+    msrc += [f"{function_name}({', '.join(fbufs)});"]
     msrc += ["*(unsigned long long *)(pra[2].buf.pv) = HAP_perf_get_time_us() - start;"]
     msrc += [f'HAP_munmap(buf_{i}, sz_or_val_{i});' for i,b in enumerate(bufs) if b[1][0].addrspace == AddrSpace.GLOBAL]
     msrc += ["return 0; }"]
@@ -266,7 +267,7 @@ class MockDSPRenderer(DSPRenderer):
         msrc.append(f"{self._render_dtype(b[1][0].dtype)} val{i}; read(0, &val{i}, {b[1][0].dtype.itemsize});")
     msrc.append("unsigned int st = inscount();")
     params = [(f'(void*)buf{i}' if b[1][0].addrspace == AddrSpace.GLOBAL else f'val{i}') for i,b in enumerate(bufs)]
-    msrc.append(f"{function_name}(&(struct {function_name}_args){{{', '.join(params)}}});")
+    msrc.append(f"{function_name}({', '.join(params)});")
     msrc.append("unsigned int et = inscount() - st; write(1, &et, sizeof(et));")
     for i,b in enumerate(bufs):
       if b[1][0].addrspace == AddrSpace.GLOBAL: msrc.append(f"write(1, buf{i}, {b[1][0].max_numel()*b[1][0].dtype.itemsize});")

tinygrad_repo/tinygrad/runtime/ops_nvme.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+import functools, os, glob, pathlib, struct
+from typing import Callable, cast
+from tinygrad.device import Allocator, Buffer, BufferSpec, BufferStorage, Compiled, Device
+from tinygrad.dtype import dtypes
+from tinygrad.helpers import ceildiv, round_up, to_tuple, flatten, wait_cond
+from tinygrad.engine.realize import get_call_arg_uops
+from tinygrad.runtime.autogen import nvme, pci
+from tinygrad.runtime.support.hcq2 import to_name, ins, unwrap_view
+from tinygrad.runtime.support.memory import AddrSpace, MemoryManager, MMIOInterface
+from tinygrad.runtime.support.system import PCIDevice, PCIIfaceBase, System, filter_visible_devices
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat
+
+NVME_IDS = (0x15b7, ((0xffff, (0x5030, 0x5050)),), 0x01) # vendor, (mask, device ids), base class
+SQE, CQE = 1 << nvme.NVME_NVM_IOSQES, 1 << nvme.NVME_NVM_IOCQES
+
+class NVMEQueue:
+  def __init__(self, dev:NVMEDev, qid:int):
+    self.dev, self.qid, self.tail, self.phase, self.entries = dev, qid, 0, 1, 0x1000 // SQE # one page, so it's contiguous even without an iommu
+    (self.sq, (self.sq_addr,)), (self.cq, (self.cq_addr,)) = dev.pci_dev.alloc_sysmem(0x1000), dev.pci_dev.alloc_sysmem(0x1000)
+    self.sq_db, self.cq_db = nvme.NVME_REG_DBS + 2 * qid * dev.db_stride, nvme.NVME_REG_DBS + (2 * qid + 1) * dev.db_stride
+
+  def submit(self, opcode:int, nsid:int=0, prp1:int=0, cdw10:int=0, cdw11:int=0):
+    slot, self.tail = self.tail, (self.tail + 1) % self.entries
+    self.sq[slot * SQE:(slot + 1) * SQE] = struct.pack("<IIQQQQII16x", opcode | slot << 16, nsid, 0, 0, prp1, 0, cdw10, cdw11)
+    System.memory_barrier()
+    self.dev.regs[self.sq_db // 4] = self.tail
+    wait_cond(lambda: self.cq.view(slot * CQE + 12, 4, fmt='I')[0] >> 16 & 1, value=self.phase, msg=f"nvme queue {self.qid}")
+    if (dw3:=self.cq.view(slot * CQE + 12, 4, fmt='I')[0]) >> 17:
+      raise RuntimeError(f"nvme {self.dev.pci_dev.pcibus}: queue {self.qid} command {opcode:#x} failed, status {dw3 >> 17:#x}")
+    self.phase ^= int(self.tail == 0)
+    self.dev.regs[self.cq_db // 4] = self.tail
+
+class NVMEDev:
+  def __init__(self, pci_dev:PCIDevice):
+    self.pci_dev = pci_dev
+    pci_dev.write_config(pci.PCI_COMMAND, pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MEMORY | pci.PCI_COMMAND_MASTER, 2)
+    self.regs = pci_dev.map_bar(0, fmt='I')
+    cap = self.regs[nvme.NVME_REG_CAP // 4] | self.regs[nvme.NVME_REG_CAP // 4 + 1] << 32
+    self.db_stride, timeout_ms = 4 << (cap >> 32 & 0xf), (cap >> 24 & 0xff) * 500
+
+    self.regs[nvme.NVME_REG_CC // 4] = 0
+    wait_cond(lambda: self.regs[nvme.NVME_REG_CSTS // 4] & nvme.NVME_CSTS_RDY, value=0, timeout_ms=timeout_ms, msg="nvme disable")
+    self.admin_q = NVMEQueue(self, 0)
+    self.regs[nvme.NVME_REG_AQA // 4] = (self.admin_q.entries - 1) << 16 | (self.admin_q.entries - 1)
+    for reg, addr in ((nvme.NVME_REG_ASQ, self.admin_q.sq_addr), (nvme.NVME_REG_ACQ, self.admin_q.cq_addr)):
+      self.regs[reg // 4], self.regs[reg // 4 + 1] = addr & 0xffffffff, addr >> 32
+    self.regs[nvme.NVME_REG_CC // 4] = nvme.NVME_CC_IOCQES | nvme.NVME_CC_IOSQES | nvme.NVME_CC_CSS_NVM | nvme.NVME_CC_ENABLE
+    wait_cond(lambda: self.regs[nvme.NVME_REG_CSTS // 4] & nvme.NVME_CSTS_RDY, value=1, timeout_ms=timeout_ms, msg="nvme enable")
+
+    ident, (ident_addr,) = pci_dev.alloc_sysmem(0x1000)
+    self.admin_q.submit(nvme.nvme_admin_identify, 0, ident_addr, 1)
+    self.max_xfer = 0x1000 << min(ident[77] or 8, 8) # at most prp1 and a 256 entry list
+    self.admin_q.submit(nvme.nvme_admin_identify, 1, ident_addr, 0)
+    self.lba_shift = ident[128 + 4 * (ident[26] & 0xf) + 2]
+    self.size = struct.unpack_from("<Q", bytes(ident[0:8]))[0] << self.lba_shift
+    self.admin_q.submit(nvme.nvme_admin_set_features, 0, 0, nvme.NVME_FEAT_NUM_QUEUES, 0xfffe << 16 | 0xfffe)
+    self.queues:list[NVMEQueue] = []
+
+  def new_queue(self) -> NVMEQueue:
+    self.queues.append(q:=NVMEQueue(self, len(self.queues) + 1))
+    self.admin_q.submit(nvme.nvme_admin_create_cq, 0, q.cq_addr, (q.entries - 1) << 16 | q.qid, 1)
+    self.admin_q.submit(nvme.nvme_admin_create_sq, 0, q.sq_addr, (q.entries - 1) << 16 | q.qid, q.qid << 16 | 1)
+    return q
+
+def prp_table(dev:NVMEDev, pages:list[int]) -> tuple[MMIOInterface, list[int]]:
+  lists, list_addrs = dev.pci_dev.alloc_sysmem(ceildiv(len(pages) + 1, 256) * 0x1000)
+  for j in range(len(list_addrs)):
+    chunk = pages[256 * j:256 * j + 512] # list pages overlap, so a prp list never crosses a page and needs no chaining
+    lists[j * 0x1000:j * 0x1000 + 8 * len(chunk)] = struct.pack(f"<{len(chunk)}Q", *chunk)
+  table, table_addrs = dev.pci_dev.alloc_sysmem(round_up(16 * len(pages), 0x1000))
+  # each page and a pointer to the prp list of the pages after it
+  entries = flatten((page, list_addrs[(p + 1) // 256] + (p + 1) % 256 * 8) for p, page in enumerate(pages))
+  table[:8 * len(entries)] = struct.pack(f"<{len(entries)}Q", *entries)
+  return table, table_addrs
+
+class NVMEIface(PCIIfaceBase):
+  def __init__(self, dev:NVMEDevice, index:int):
+    # make sure the drive isn't mounted or used by md/lvm
+    pcibus = filter_visible_devices(System.list_devices(*NVME_IDS), "NVME")[index][1]
+    mounted = {line.split()[0] for line in pathlib.Path("/proc/mounts").read_text().splitlines()}
+    for b in [b for b in glob.glob("/sys/class/block/*") if f"/{pcibus}/" in os.path.realpath(b)]:
+      assert f"/dev/{os.path.basename(b)}" not in mounted and not os.listdir(f"{b}/holders"), f"{dev.device} ({pcibus}) is in use by the system: {b}"
+    super().__init__(dev, index, *NVME_IDS[:2], vram_bar=0, va_start=MemoryManager.va_allocator.base, va_size=MemoryManager.va_allocator.size,
+                     dev_impl_t=NVMEDev, base_class=NVME_IDS[2])
+
+class NVMEAllocator(Allocator['NVMEDevice']):
+  lru = False
+  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage: return BufferStorage(0)
+  def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset
+  def _map(self, buf:Buffer) -> BufferStorage:
+    iface = getattr(Device[buf.device], "iface", None)
+    if not isinstance(iface, PCIIfaceBase) or iface.peer_group != self.dev.peer_group: raise RuntimeError(f"NVME can't reach {buf.device}")
+    mapping = buf.meta.mapping
+    paddrs = mapping.paddrs if mapping.aspace is AddrSpace.SYS else PCIIfaceBase.p2p_paddrs(iface, mapping.paddrs)[0]
+    table = self.dev.iface.buffer(*prp_table(self.dev.nvme, [p + off for p, size in paddrs for off in range(0, size, 0x1000)]))
+    return BufferStorage(table.get_buf(buf.device), table)
+
+class NVMEDevice(Compiled):
+  ifaces = [NVMEIface]
+
+  def __init__(self, device:str):
+    self.iface = self._select_iface(device)
+    self.nvme = cast(NVMEDev, self.iface.dev_impl)
+    super().__init__(device, NVMEAllocator(self), [], None)
+  def synchronize(self, timeout:int|None=None): pass
+
+# *****************
+# UOps implementation
+
+@functools.cache
+def nvme_queue(drive:str, gpu:str, queue:str) -> NVMEQueue:
+  dev = cast(NVMEDevice, Device[drive])
+  q, (seq, seq_addrs), bar0 = dev.nvme.new_queue(), dev.nvme.pci_dev.alloc_sysmem(0x1000), dev.iface.pci_dev.bar_info(0)[0]
+  bufs = {"sq": dev.iface.buffer(q.sq, [q.sq_addr]), "cq": dev.iface.buffer(q.cq, [q.cq_addr]), "seq": dev.iface.buffer(seq.view(0, 8), seq_addrs),
+          "regs": dev.iface.buffer(dev.nvme.regs, [bar0 + o for o in range(0, dev.nvme.regs.nbytes, 0x1000)], snooped=False)}
+  Compiled.link_bufs |= {to_name(drive, gpu, queue, n): b for n, b in bufs.items()}
+  return q
+
+def loop(n:int, body:Callable[[UOp], list[UOp]]) -> list[UOp]: # loops in the cmdbuf, so the host program doesn't grow with the read size
+  return [UOp(Ops.LINEAR, src=tuple(body(r:=UOp.range(n, next(UOp.unique_num), dtype=dtypes.uint64)))).end(r)] if n else []
+
+def nvme_copies(queue:UOp, calls:list[UOp]) -> list[list[UOp]]: # builds the commands for all reads from one drive on one gpu queue
+  drive, gpu, qn = to_tuple(get_call_arg_uops(calls[0])[1].device)[0], queue.device, queue.arg.name
+  dev, q = cast(NVMEDevice, Device[drive]).nvme, nvme_queue(drive, gpu, qn)
+  def mem(name:str, size:int, dtype=dtypes.uint8) -> UOp: return UOp.alloc((size,), dtype, 0, device=drive, name=to_name(gpu, qn, name))
+  sq, cq, regs = mem("sq", 0x1000).getaddr(gpu), mem("cq", 0x1000).getaddr(gpu), mem("regs", dev.regs.nbytes).getaddr(gpu)
+
+  # where this submit starts in the queue, bumped once for all its commands
+  cmds = [ceildiv(get_call_arg_uops(c)[1].nbytes(), dev.max_xfer) for c in calls]
+  n = (seq:=mem("seq", 1, dtypes.uint64)).after(seq.index(0).store(seq.index(0).load() + sum(cmds))).index(0).load() - sum(cmds)
+
+  copies = []
+  for call, ncmds in zip(calls, cmds):
+    (dst, src), (full, rem) = get_call_arg_uops(call), divmod(get_call_arg_uops(call)[1].nbytes(), dev.max_xfer)
+    (dbase, doff), soff = unwrap_view(dst), unwrap_view(src)[1]
+    assert doff % 0x1000 == 0 and (soff | src.nbytes()) % (1 << dev.lba_shift) == 0 and soff + src.nbytes() <= dev.size, \
+      f"{drive} reads whole blocks of {1 << dev.lba_shift} bytes into page aligned memory, not {src.nbytes()} bytes at {soff} to {doff}"
+    prps, lba0 = dbase.getaddr(drive) + doff // 0x1000 * 16, src.getaddr(drive) >> dev.lba_shift
+    def wave(k:UOp|int, size:int) -> list[UOp]: # each wave is drained before the next, so cids (slot + 1) stay unique in flight
+      ops = []
+      for j in range(size): # write the sqe, then copy prp1 and prp2 (the next page or a list) from the table
+        c, slot = k * (q.entries - 1) + j, (n + k * (q.entries - 1) + j) % q.entries
+        nb = rem if isinstance(c, int) and c == full else dev.max_xfer
+        entry, lba, page = sq + slot * SQE, lba0 + c * (dev.max_xfer >> dev.lba_shift), c * (dev.max_xfer // 0x1000)
+        ops += [ins("write", entry, ((slot + 1) << 16 | nvme.nvme_cmd_read).cast(dtypes.uint32), 1, *[0] * 8, lba.cast(dtypes.uint32),
+                    (lba >> 32).cast(dtypes.uint32), (nb >> dev.lba_shift) - 1, 0, 0, 0),
+                ins("copy", entry + 24, prps + page * 16, 16 if ceildiv(nb, 0x1000) > 2 else 8)] + \
+               ([ins("copy", entry + 32, prps + (page + 1) * 16, 8)] if ceildiv(nb, 0x1000) == 2 else [])
+      cqe = cq + ((n + k * (q.entries - 1) + size - 1) % q.entries) * CQE + 12 # the wave's last cqe, cleared before ringing the doorbell
+      end = ((n + k * (q.entries - 1) + size) % q.entries).cast(dtypes.uint32)
+      return ops + [ins("write", cqe, 0), ins("write", regs + q.sq_db, end), ins("wait", cqe, 1), ins("write", regs + q.cq_db, end)]
+    waves, last = full // (q.entries - 1), full % (q.entries - 1) + (rem > 0) # loop over the full waves, then do the rest
+    copies.append(loop(waves, lambda k: wave(k, q.entries - 1)) + (wave(waves, last) if last else []) + [ins("memory_barrier")])
+    n = n + ncmds
+  return copies
+
+def nvme_submit(s:UOp) -> UOp|None:
+  lin, queue = s.without_after.body, s.without_after.src[1]
+  if not (drives:={i: d for i, u in enumerate(lin.src) if u.op is Ops.CALL and u.src[0].op is Ops.STORE and
+                   (d:=to_tuple(get_call_arg_uops(u)[1].device)[0]).startswith("NVME")}): return None
+  ops = [[u] for u in lin.src]
+  for d in dict.fromkeys(drives.values()):
+    positions = [i for i in drives if drives[i] == d]
+    for i, copy_ops in zip(positions, nvme_copies(queue, [lin.src[i] for i in positions])): ops[i] = copy_ops
+  return s.without_after.replace(src=(lin.replace(src=tuple(flatten(ops))), queue)).after(*s.src[1:])
+Compiled.pm_encode = PatternMatcher([(UPat.sink().call(UPat(Ops.PARAM)).after(allow_any_len=True, name="s"), nvme_submit)]) + Compiled.pm_encode

tinygrad_repo/tinygrad/runtime/ops_rdma.py

@@ -8,8 +8,8 @@ from tinygrad.engine.realize import get_call_arg_uops
 from tinygrad.runtime.autogen import bnxt
 from tinygrad.runtime.support.rdma.bnxtdev import BNXTDev, BNXTQP, db_value, send_wqe, recv_wqe, WQE_SIZE, RING_ENTRIES, CQ_ENTRIES, MTU
 from tinygrad.runtime.support.hcq2 import unwrap_view, to_name, ins
-from tinygrad.runtime.support.memory import AddrSpace, MMIOInterface, VirtMapping, MemoryManager
-from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, System, filter_visible_devices
+from tinygrad.runtime.support.memory import AddrSpace, MemoryManager
+from tinygrad.runtime.support.system import PCIIfaceBase, System, filter_visible_devices
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat
 
 RDMA_CHUNK = 1 << 30 # a wqe length is 32 bits
@@ -22,13 +22,6 @@ class BNXTIface(PCIIfaceBase):
                      dev_impl_t=BNXTDev, base_class=BNXT_IDS[2])
   def device_fini(self): self.dev_impl.fini()
 
-  def storage(self, mem:MMIOInterface, paddrs:list[int], snooped:bool=True) -> BufferStorage: # nic memory any gpu of the node maps
-    va = MemoryManager.alloc_vaddr(size:=round_up(mem.nbytes, 0x1000), 0x1000)
-    mapping = VirtMapping(va, size, [(p, 0x1000) for p in paddrs], AddrSpace.SYS, uncached=True, snooped=snooped)
-    return BufferStorage(va, PCIAllocationMeta(mapping, True), mem)
-  def buffer(self, mem:MMIOInterface, paddrs:list[int], snooped:bool=True) -> Buffer:
-    return Buffer(self.dev.device, mem.nbytes, opaque=self.storage(mem, paddrs, snooped))
-
   @functools.cached_property
   def doorbell(self) -> Buffer:
     off = self.dev_impl.db_off & ~0xfff

tinygrad_repo/tinygrad/runtime/support/hcq2.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from typing import cast, Any, Sequence
 import itertools, ctypes, struct, time
 from dataclasses import replace, dataclass, field
-from collections import defaultdict
+from collections import defaultdict, Counter
 from tinygrad.helpers import dedup, pluralize, unwrap, to_tuple, Context, panic, partition, to_name, DEBUG, VIZ, ALL2ALL, PROFILE
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, TinyELF, HCQ_RUNTIME_DEV, HCQ_DEVS, HOST_DEVS, UVM_DEVS, ProfileProgramEvent
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo, GroupOp, graph_rewrite, rewrite_group, exec_alu, uopfunc
@@ -184,7 +184,8 @@ pm_insert_copy_staging = PatternMatcher([
 # **** batch: calls become queue programs
 
 class DepsTracker:
-  def __init__(self):
+  def __init__(self, queue_of:Sequence[Any]):
+    self.queue_of = queue_of
     self.ranges:dict[tuple[UOp, int|None], list[tuple[int, int, bool, int]]] = defaultdict(list) # per (base, lane): (start, end, is_write, tag)
 
   def access_resources(self, accesses:Sequence[tuple[UOp, bool]], tag:int) -> list[int]: # (buffer, is_write) accesses -> earlier tags to wait on
@@ -195,10 +196,11 @@ class DepsTracker:
     waits = [t for key, start, end, is_write in spans for lo, hi, was_write, t in self.ranges[key]
              if lo < end and start < hi and (is_write or was_write)]
 
-    # a write owns its span: the earlier accesses keep only the parts outside it
+    # a write owns its span: the earlier accesses keep only the parts outside it.
     for key, start, end, is_write in spans:
       if is_write:
         self.ranges[key] = [(lo, hi, w, t) for a, b, w, t in self.ranges[key] for lo, hi in ((a, min(b, start)), (max(a, end), b)) if lo < hi]
+      else: self.ranges[key] = [r for r in self.ranges[key] if r[:3] != (start, end, False) or self.queue_of[r[3]] != self.queue_of[tag]]
       self.ranges[key].append((start, end, is_write, tag))
     return dedup(waits)
 
@@ -244,7 +246,7 @@ def _build_queues(ctx:BatchCtx) -> dict[tuple[str, str], list[UOp]]:
   queue_of = [(dev, queue) for _, dev, queue in ctx.batch] # the (device, queue) of each call by tag
 
   # calc deps for each tag
-  tracker = DepsTracker()
+  tracker = DepsTracker(queue_of)
   accs = [[(s, k in get_call_outs_ins(c)[0]) for k, s in enumerate(c.src[1:]) if not s.is_variable] for c, _, _ in ctx.batch]
   dep_tags:list[list[int]] = [tracker.access_resources(acc, tag) for tag, acc in enumerate(accs)]
 
@@ -311,7 +313,7 @@ def _finalize_batch(ctx:BatchCtx, skip_wait:bool=False) -> UOp:
   kerns = tuple(zip([d for _, d, _ in ctx.batch], names, estimates, stamps, profile_keys, bufs, [get_call_outs_ins(c) for c, _, _ in ctx.batch]))
   slots = ctx.slots if ctx.profile else {} # profiling reads them
   info = HCQInfo(tuple(ctx.queues), skip_wait=skip_wait, kernels=kerns, slots=tuple((d, i) for i, d in enumerate(slots)),
-                 estimates=Estimates(*[sum(x*t.count(x) for x in dict.fromkeys(t)) for t in zip(*[vars(e).values() for e in estimates])]).simplify())
+                 estimates=Estimates(*[sum(x*n for x, n in Counter(t).items()) for t in zip(*[vars(e).values() for e in estimates])]).simplify())
   return sink.call(*slots.values(), aux=info)
 
 @rewrite_group(new_ctx=False)

tinygrad_repo/tinygrad/runtime/support/system.py

@@ -4,7 +4,7 @@ import fcntl
 from tinygrad.device import BufferStorage, Buffer, Device
 from tinygrad.helpers import round_up, getenv, OSX, temp, DEBUG, pluralize, DEV
 from tinygrad.runtime.autogen import libc, pci, vfio
-from tinygrad.runtime.support.memory import VirtMapping, AddrSpace, BumpAllocator, MMIOInterface, TLSFAllocator
+from tinygrad.runtime.support.memory import VirtMapping, AddrSpace, BumpAllocator, MMIOInterface, TLSFAllocator, MemoryManager
 from tinygrad.runtime.support.usb import USB3, CustomASM24Controller, USBMMIOInterface
 
 def filter_visible_devices(devs, device):
@@ -80,6 +80,16 @@ class _System:
   def lock_memory(self, addr:int, size:int):
     if libc.mlock(ctypes.c_void_p(addr), size): raise RuntimeError(f"Failed to lock memory at {addr:#x} with size {size:#x}")
 
+  @functools.cache
+  def vfio_group(self, grp:str) -> FileIOInterface:
+    for dev in FileIOInterface(f"/sys/kernel/iommu_groups/{grp}/devices").listdir():
+      if FileIOInterface.exists(drv:=f"/sys/bus/pci/devices/{dev}/driver") and not FileIOInterface.readlink(drv).endswith(("/vfio-pci", "/pcieport")):
+        FileIOInterface(f"{drv}/unbind", os.O_WRONLY).write(dev)
+    self.vfio = self.vfio or FileIOInterface("/dev/vfio/vfio", os.O_RDWR)
+    vfio.VFIO_GROUP_SET_CONTAINER(group:=FileIOInterface(f"/dev/vfio/{grp}", os.O_RDWR), ctypes.c_int(self.vfio.fd))
+    with contextlib.suppress(OSError): vfio.VFIO_SET_IOMMU(self.vfio, vfio.VFIO_TYPE1v2_IOMMU) # set iommu works only once for the fd.
+    return group
+
   def vfio_map(self, vaddr:int, size:int, align:int=1) -> int:
     iova = self.iova.alloc(size, align)
     flags = vfio.VFIO_DMA_MAP_FLAG_READ | vfio.VFIO_DMA_MAP_FLAG_WRITE
@@ -217,16 +227,12 @@ class PCIDevice:
     for fn in range(1, 8):
       if FileIOInterface.exists(sib:=f"/sys/bus/pci/devices/{self.pcibus[:-1]}{fn}"): FileIOInterface(f"{sib}/remove", os.O_WRONLY).write("1")
 
-    if FileIOInterface.exists(grp:=f"/sys/bus/pci/devices/{self.pcibus}/iommu_group") and FileIOInterface(f"{grp}/type").read().strip() != "identity":
+    if FileIOInterface.exists(grp:=f"/sys/bus/pci/devices/{self.pcibus}/iommu_group"):
       if not FileIOInterface.exists("/sys/module/vfio_pci"): raise RuntimeError("IOMMU is on: run `sudo modprobe vfio-pci disable_idle_d3=1`")
       FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver_override", os.O_WRONLY).write("vfio-pci")
       FileIOInterface("/sys/bus/pci/drivers_probe", os.O_WRONLY).write(self.pcibus)
-      vfio_fd = System.vfio = System.vfio or FileIOInterface("/dev/vfio/vfio", os.O_RDWR)
-      self.vfio_group = FileIOInterface(f"/dev/vfio/{FileIOInterface.readlink(grp).split('/')[-1]}", os.O_RDWR)
-      vfio.VFIO_GROUP_SET_CONTAINER(self.vfio_group, ctypes.c_int(vfio_fd.fd))
-
-      with contextlib.suppress(OSError): vfio.VFIO_SET_IOMMU(vfio_fd, vfio.VFIO_TYPE1v2_IOMMU) # set iommu works only once for the fd.
-      self.vfio_dev = FileIOInterface(fd=vfio.VFIO_GROUP_GET_DEVICE_FD(self.vfio_group, ctypes.create_string_buffer(self.pcibus.encode())))
+      group = System.vfio_group(FileIOInterface.readlink(grp).split('/')[-1])
+      self.vfio_dev = FileIOInterface(fd=vfio.VFIO_GROUP_GET_DEVICE_FD(group, ctypes.create_string_buffer(self.pcibus.encode())))
 
       self.irq_fd = FileIOInterface.eventfd(0, 0)
       self.irq_poller = select.poll()
@@ -336,6 +342,13 @@ class PCIIfaceBase:
   def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]:
     return [(p + self.pci_dev.bar_info(self.vram_bar)[0], sz) for p, sz in paddrs], AddrSpace.SYS
 
+  def storage(self, mem:MMIOInterface, paddrs:list[int], snooped:bool=True) -> BufferStorage:
+    va = MemoryManager.alloc_vaddr(size:=round_up(mem.nbytes, 0x1000), 0x1000)
+    mapping = VirtMapping(va, size, [(p, 0x1000) for p in paddrs], AddrSpace.SYS, uncached=True, snooped=snooped)
+    return BufferStorage(va, PCIAllocationMeta(mapping, True), mem)
+  def buffer(self, mem:MMIOInterface, paddrs:list[int], snooped:bool=True) -> Buffer:
+    return Buffer(self.dev.device, mem.nbytes, opaque=self.storage(mem, paddrs, snooped))
+
   def map(self, b:Buffer) -> BufferStorage:
     if (host:=b.device.split(":")[0] in {"CPU", "PYTHON", "NPY"}):
       if self.dev.host != Device[b.device].host: raise RuntimeError(f"host memory is not on the node of {self.dev.device}")

tinygrad_repo/tinygrad/runtime/support/usb.py

@@ -418,7 +418,7 @@ Compiled.pm_encode = PatternMatcher([(UPat.sink().call(UPat(Ops.PARAM)).after(al
 Compiled.pm_lower += PatternMatcher([
   # the gpu's memory is reached over usb
   (UPat.custom_function("memcpy", dtypes.uint64).call(UPat.var("dst"), UPat.var("src"), UPat.var("n")), usb_memcpy),
-  (UPat.custom_function("memset", dtypes.uint64).call(UPat.var("dst"), UPat.const(0), UPat.var("n")),
+  (UPat.custom_function("memset", dtypes.uint64).call(UPat.var("dst"), UPat.const(0).or_casted(), UPat.var("n")),
    lambda dst, n: usb_memcpy(dst, usb_zeros(dst.src[0].device).index(0), n, fill=True)),
   (UPat.var("b").index(UPat.var("idx")).store(UPat.var("v")), lambda b, idx, v: usb_store(b, idx, v) if is_remote(b) else None),
   (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),

tinygrad_repo/tinygrad/schedule/indexing.py

@@ -31,7 +31,7 @@ def realize_srcs(ctx:IndexingContext, rb:UOp) -> None:
 
 def realize_store_after_src(ctx:IndexingContext, dest:UOp, src:UOp):
   # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
-  if dest.base in src.toposort(): ctx.realize_map[src] = None
+  if dest.base in src.toposort(gate=lambda u: u is dest.base or u.op is not Ops.AFTER): ctx.realize_map[src] = None
 
 def realize_custom_kernel_srcs(ctx:IndexingContext, c:UOp) -> None:
   for s in c.src[1:]:

tinygrad_repo/tinygrad/schedule/memory.py

@@ -12,7 +12,7 @@ def _collect_bufs(u:UOp) -> list[UOp]:
 def _can_plan(b:UOp, held_bufs:set[UOp]) -> bool:
   if b in held_bufs: return False
   # CL and WEBGPU do not support views, see explanation in contiguous_view_offset
-  return all(not d.startswith(("DISK", "CL", "WEBGPU")) for d in to_tuple(b.device))
+  return all(not d.startswith(("DISK", "NVME", "CL", "WEBGPU")) for d in to_tuple(b.device))
 
 LaneKey = tuple[str, int]
 

tinygrad_repo/tinygrad/schedule/prepare.py

@@ -69,7 +69,6 @@ def _mop_index(r:UOp, idx:UOp):
   if r.op is Ops.RESHAPE:
     src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
     if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
-      if src_prefix == 0: return r.src[0]
       ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), arg=idx.arg)
       return ret if ret.shape == idx.shape else None
 

tinygrad_repo/tinygrad/uop/__init__.py

@@ -24,7 +24,7 @@ class Ops(FastEnum):
   # uops that aren't rendered
   NOOP = auto(); REWRITE_ERROR = auto()
   # CALL is a kernel invocation; value-producing calls write to ALLOC arguments
-  PARAM = auto(); CALL = auto()
+  PARAM = auto(); REFPARAM = auto(); CALL = auto()
 
   # renderer
   # LINEAR is a list of UOps, SOURCE has a str arg that's human readable, BINARY has bytes arg that's compiled
@@ -111,7 +111,7 @@ class GroupOp:
   # TODO: is BITCAST always Elementwise if it's shape changing?
   Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST})
 
-  Defines = {Ops.PARAM, Ops.BUFFER, Ops.ALLOC}
+  Defines = {Ops.PARAM, Ops.REFPARAM, Ops.BUFFER, Ops.ALLOC, Ops.BINARY}
 
   Irreducible = {Ops.CONST, Ops.SPECIAL, Ops.RANGE, Ops.PARAM, Ops.GETADDR}
   Movement = {Ops.RESHAPE, Ops.EXPAND, Ops.PERMUTE, Ops.PAD, Ops.SHRINK}

tinygrad_repo/tinygrad/uop/ops.py

@@ -166,7 +166,7 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType:
       if not all(dtypes.is_int(x.dtype) or x.base.is_invalid for x in src):
         raise RuntimeError(f"shift operands must be int, got {[x.dtype for x in src]}")
       return src[0].dtype
-    case Ops.BUFFER | Ops.ALLOC | Ops.PARAM:
+    case Ops.BUFFER | Ops.ALLOC | Ops.PARAM | Ops.REFPARAM:
       assert isinstance(arg, ParamArg), f"{op} must have ParamArg"
       return arg.dtype
     case Ops.BINARY:
@@ -347,7 +347,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       case Ops.GETADDR: return ()
       case Ops.RANGE | Ops.SPECIAL: return ()
       case Ops.BINARY: return (len(self.arg),)
-      case Ops.BUFFER | Ops.ALLOC | Ops.PARAM:
+      case Ops.BUFFER | Ops.ALLOC | Ops.PARAM | Ops.REFPARAM:
         assert len(self.src[0].as_shape) <= 1
         if (img:=self.arg.image) is not None: return (img[0], img[1], 4)
         return self.src[0].as_shape
@@ -501,7 +501,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     assert isinstance(vmin, expected_type), f"vmin is wrong dtype {type(vmin)} != {expected_type}"
     return vmin
   def __bool__(self): return self._eval((dtypes.bool,), bool)
-  def __int__(self): return self._eval(dtypes.ints+(dtypes.weakint,), int)
+  def __int__(self): return self._eval(dtypes.wints, int)
   def __float__(self): return float(self._eval(dtypes.floats+(dtypes.weakfloat,), float))
   def substitute(self, dvars:dict[UOp, UOp], name:str|None=None, extra_pm:PatternMatcher|None=None, walk:bool=False, enter_calls:bool=False):
     dvars = {k:v for k,v in dvars.items() if k is not v}
@@ -733,7 +733,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def mstack(self, *srcs: UOp) -> UOp: return UOp(Ops.MSTACK, src=(self,)+srcs) if len(srcs) else self
 
   # little helpers
-  def on_disk(self:UOp): return isinstance(self.device, str) and self.device.startswith("DISK")
+  def on_disk(self:UOp): return isinstance(self.device, str) and is_disk_device(self.device)
   def on_creation_device(self:UOp): return isinstance(self.device, str) and self.device.startswith(("DISK", "NPY", "PYTHON"))
   def needs_storage(self:UOp) -> bool: return not self.is_virtual and (self.storage_base.op is Ops.ALLOC or not self.has_buffer_identity())
 
@@ -870,8 +870,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def addrspace(self) -> AddrSpace|None:
     # anything that has dtype void has no memory backing, it lives in the ALU
     if self.dtype is dtypes.void: return AddrSpace.ALU
-    if self.op is Ops.PARAM: return self.arg.addrspace
-    if self.op in {Ops.BUFFER, Ops.ALLOC}: return self.arg.addrspace
+    if self.op in {Ops.PARAM, Ops.REFPARAM, Ops.BUFFER, Ops.ALLOC}: return self.arg.addrspace
     if self.op in {Ops.SPECIAL, Ops.RANGE, Ops.CONST}: return AddrSpace.ALU
     if self.op is Ops.BINARY: return AddrSpace.GLOBAL
     if self.op is Ops.LOAD: return AddrSpace.ALU # LOAD brings things into the ALU
@@ -1112,7 +1111,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       smin, smax = self.src[0]._min_max
       trunc = truncate.get(self.dtype) if dtypes.is_float(self.dtype) else math.trunc if dtypes.is_int(self.dtype) else None
       if trunc is not None: smin, smax = (trunc(v) if math.isfinite(v) else v for v in (smin, smax))
-      if dtypes.is_unsigned(self.dtype) and 0 <= smin and smax <= self.dtype.max: return smin, smax
+      if dtypes.is_int(self.dtype) and self.dtype.min <= smin and smax <= self.dtype.max: return smin, smax
+      if dtypes.is_int(self.src[0].dtype) and dtypes.is_int(self.dtype): return self.dtype.min, self.dtype.max
       # a signed or float destination holds the part of the source that overlaps it: overflow is undefined, a nan bound overlaps nothing
       if self.dtype in dtypes.floats+dtypes.sints+dtypes.weaks and smin <= self.dtype.max and self.dtype.min <= smax:
         return max(self.dtype.min, smin), min(smax, self.dtype.max)
@@ -1267,8 +1267,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def kernel_params(self) -> tuple[UOp, ...]:
-    # the kernel's params in its arg order, which is slot order. the LINEAR keeps two IMAGE params of one slot in order, x86 isel replaces every PARAM
-    params = [u for u in self.src[1].src if u.op is Ops.PARAM]
+    # the kernel's params in slot order. the LINEAR keeps two IMAGE params of one slot in order, isel and packing replace PARAMs
+    params = [u for u in self.src[1].src if u in self.arg.params]
     params += [p for p in self.arg.params if p.arg.slot not in {u.arg.slot for u in params}]
     return tuple(sorted(params, key=lambda p: p.arg.slot))
 
@@ -1299,6 +1299,7 @@ class KernelInfo:
   estimates: Estimates|None = None
   beam: int = 0
   optimized: bool = False       # the AST has been through the optimizer, skip re-optimization
+  kernel: bool = True           # launched by the runtime, not called from another kernel
   @property
   def function_name(self): return to_function_name(self.name)
 

tinygrad_repo/tinygrad/uop/render.py

@@ -37,7 +37,7 @@ def _render_arg(x:UOp) -> str:
       if dt in dtypes.weaks: return repr(v) if dt is dtypes.weakint else repr(float(v))
       if dtypes.is_float(dt): return f"{dt.name}:{float(v).hex()}"   # float.hex() roundtrips exactly, inf/nan included
       return f"{dt.name}:{v}"
-    case Ops.PARAM | Ops.BUFFER | Ops.ALLOC:
+    case Ops.PARAM | Ops.REFPARAM | Ops.BUFFER | Ops.ALLOC:
       a, opts = x.arg, ""
       if a.vmin_vmax is not None: opts += f" bounds=[{a.vmin_vmax[0]},{a.vmin_vmax[1]}]"
       if a.addrspace not in (None, AddrSpace.GLOBAL): opts += f" addrspace={a.addrspace.name}"

tinygrad_repo/tinygrad/uop/spec.py

@@ -90,6 +90,7 @@ spec_shared = PatternMatcher([
 
   # PARAM/BUFFER carry a size CONST or an empty STACK for scalars
   (UPat(Ops.PARAM, src=(p_size,), name="x"), lambda x: isinstance(x.arg, ParamArg)),
+  (UPat(Ops.REFPARAM, src=(p_size, UPat(dtype=dtypes.uint64)), name="x"), lambda x: isinstance(x.arg, ParamArg)),
   (UPat(Ops.BUFFER, src=(p_size,), name="x"), lambda x: isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.REG, AddrSpace.LOCAL)),
 
   (UPat(Ops.BINARY, dtypes.uint8, src=(), name="x"), lambda x: isinstance(x.arg, bytes)),

tinygrad_repo/tinygrad/uop/symbolic.py

@@ -30,6 +30,8 @@ def fold_bitcast(root:UOp, c:UOp) -> UOp|None:
 def fold_const_alu(a:UOp) -> UOp: return a.const_like(exec_alu(a.op, a.dtype, [const_arg(s) for s in a.src]))
 # a fold mints a signed const only if it fits, the mint would wrap it. unsigned ring ops wrap by definition
 def fits(x:UOp, v:PyConst) -> bool: return x.dtype not in dtypes.sints or x.dtype.min <= v <= x.dtype.max
+# a rewrite that introduces a signed intermediate must keep its range in the dtype, overflow is UB
+def no_overflow(u:UOp) -> bool: return fits(u, u.vmin) and fits(u, u.vmax)
 
 def _quotient_base(q:UOp, base:UOp, div:int) -> UOp|None:
   # the B with q == B//div and B%div == base%div, or None. only such congruence is needed to recombine, and canonicalization
@@ -139,7 +141,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
   (UPat.var("x", dtype=dtypes.bool).where(UPat.const(True, dtypes.bool), UPat.const(False, dtypes.bool)), lambda x: x),
   (UPat.var("x", dtype=dtypes.bool).where(UPat.const(False, dtypes.bool), UPat.const(True, dtypes.bool)), lambda x: x.logical_not()),
   # CAST(bool -> int) == const — CAST(True)=1, CAST(False)=0, so fold based on const value
-  (UPat.var("x", dtype=dtypes.bool).cast(dtypes.ints+(dtypes.weakint,)).eq(UPat.cvar("c")),
+  (UPat.var("x", dtype=dtypes.bool).cast(dtypes.wints).eq(UPat.cvar("c")),
    lambda x,c: x.logical_not() if c.val == 0 else x if c.val == 1 else x.const_like(False)),
   (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)).trunc(), lambda x: x),
   # ** zero folding **
@@ -166,7 +168,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
    a.replace(src=tuple(s.ccast(dt) if s.dtype in dtypes.weaks else s for s in a.src))
    if (dt:=promo_dtype(a.src)) not in dtypes.weaks else None),
   # ** two stage ALU folding **
-  *((UPat.var("x").alu(op, UPat.cvar("c1")).alu(op, UPat.cvar("c2")).named("f"),
+  *((UPat.var("x", dtypes.wints if op is Ops.ADD else None).alu(op, UPat.cvar("c1")).alu(op, UPat.cvar("c2")).named("f"),
      lambda f,x,c1,c2: x.alu(f.op,c1.alu(f.op,c2)) if fits(x, exec_alu(f.op, dtypes.weakint, [c1.vmin, c2.vmin])) else None)
     for op in GroupOp.Associative),
   # bool MUL is AND, ADD/MAX is OR. prevents other rules to rewrite bool ADD/MUL incorrectly
@@ -248,14 +250,17 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   (UPat.var("x", dtype=dtypes.bool) | UPat.var("x", dtype=dtypes.bool).logical_not(), lambda x: x.const_like(True)),  # x|!x -> True
   # ** combine terms **
   # (x*c0)+(x*c1) -> x*(c0+c1)
-  (UPat.var("x") * UPat.cvar("c0") + UPat.var("x") * UPat.cvar("c1"), lambda x,c0,c1: x*(c0+c1) if fits(x, c0.vmin+c1.vmin) else None),
-  ((UPat.var("y") + UPat.var("x") * UPat.cvar("c0")) + UPat.var("x") * UPat.cvar("c1"),
-   lambda x,y,c0,c1: y+x*(c0+c1) if fits(x, c0.vmin+c1.vmin) else None),
-  (UPat.var("x") + UPat.var("x") * UPat.cvar("c"), lambda x,c: x*(c+1) if fits(x, c.vmin+1) else None), # (x+x*c)-> x*(c+1)
-  ((UPat.var("y") + UPat.var("x")) + UPat.var("x") * UPat.cvar("c"), lambda x,y,c: y+x*(c+1) if fits(x, c.vmin+1) else None),
-  ((UPat.var("y") + UPat.var("x") * UPat.cvar("c")) + UPat.var("x"), lambda x,y,c: y+x*(c+1) if fits(x, c.vmin+1) else None),
+  (UPat.var("x", dtypes.wints) * UPat.cvar("c0") + UPat.var("x") * UPat.cvar("c1"), lambda x,c0,c1: x*(c0+c1) if fits(x, c0.vmin+c1.vmin) else None),
+  ((UPat.var("y") + UPat.var("x", dtypes.wints) * UPat.cvar("c0")) + UPat.var("x") * UPat.cvar("c1"),
+   lambda x,y,c0,c1: y+m if fits(x, c0.vmin+c1.vmin) and no_overflow(m:=x*(c0+c1)) else None),
+  (UPat.var("x") + UPat.var("x") * UPat.cvar("c"), # (x+x*c)-> x*(c+1)
+   lambda x,c: x*(c+1) if (x.dtype in dtypes.wints or c.arg == -1) and fits(x, c.vmin+1) else None),
+  ((UPat.var("y") + UPat.var("x", dtypes.wints)) + UPat.var("x") * UPat.cvar("c"),
+   lambda x,y,c: y+m if fits(x, c.vmin+1) and no_overflow(m:=x*(c+1)) else None),
+  ((UPat.var("y") + UPat.var("x", dtypes.wints) * UPat.cvar("c")) + UPat.var("x"),
+   lambda x,y,c: y+m if fits(x, c.vmin+1) and no_overflow(m:=x*(c+1)) else None),
   (UPat.var("x") + UPat.var("x"), lambda x: x*2), # (x+x)-> x*2
-  (((UPat.var("y") + UPat.var("x")) + UPat.var("x")).named("a"),
+  (((UPat.var("y") + UPat.var("x", dtypes.wints)) + UPat.var("x")).named("a"),
    lambda a,y,x: y+x*2 if fits(a, 2*x.vmin) and fits(a, 2*x.vmax) else None),
   (UPat.cvar("y") * (UPat.var("x", dtype=dtypes.weakint) + UPat.cvar("c")), lambda x,y,c: (y*x)+(y*c)),  # y*(x+c) -> y*x + y*c
   # ** where folding **
@@ -274,8 +279,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   (UPat(GroupOp.Binary, name="alu", src=(UPat.var("c").where(UPat.var("t"), UPat.var("f")), UPat.var("c").where(UPat.var("tt"), UPat.var("ff")))), \
    lambda alu,c,t,tt,f,ff: c.where(t.alu(alu.op, tt), f.alu(alu.op, ff)) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None),
   # if its a plus we add the associative variation too
-  ((UPat.var("y")+UPat.var("c").where(UPat.var("t"), UPat.var("f"))) + UPat.var("c").where(UPat.var("tt"), UPat.var("ff")), \
-   lambda y,c,t,tt,f,ff: y+c.where(t+tt, f+ff) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None),
+  ((UPat.var("y", dtypes.wints)+UPat.var("c").where(UPat.var("t"), UPat.var("f"))) + UPat.var("c").where(UPat.var("tt"), UPat.var("ff")), \
+   lambda y,c,t,tt,f,ff: y+m if (t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST) and no_overflow(m:=c.where(t+tt, f+ff)) else None),
   # complementary zero branches under the same condition select directly
   (UPat.var("c").where(UPat.var("t"), 0) + UPat.var("c").where(0, UPat.var("f")), lambda c,t,f: c.where(t, f)),
   # ALU/variable min==max -> CONST
@@ -301,8 +306,9 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   ((UPat.var("x", dtype=dtypes.weakint)//UPat.cvar("d"))<UPat.cvar("c"),
    lambda x,d,c: (x<c.val*d.val) if d.val > 0 else None),
   # ** move add/mul consts to end (NOTE: this is still happening before constant folding) **
-  ((UPat.var("x") + UPat.cvar("c1")) + UPat.var("y"), lambda x,c1,y: (x+y)+c1 if y.op is not Ops.CONST else None),
-  ((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: (x*y)*c1 if y.op is not Ops.CONST else None),
+  ((UPat.var("x", dtypes.wints) + UPat.cvar("c1")) + UPat.var("y"), lambda x,c1,y: m+c1 if y.op is not Ops.CONST and no_overflow(m:=x+y) else None),
+  ((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: m*c1 if y.op is not Ops.CONST and
+   (x.dtype in dtypes.wints or c1.arg == -1) and no_overflow(m:=x*y) else None),
   # *** rules from symbolic ***
   # generic lt folding
   (UPat.var("x", dtypes.weakint)<UPat.cvar("c"), lambda x,c: lt_folding(x, c.val)),
@@ -316,18 +322,18 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   # cast/long folding
   # if the intermediate cast doesnt narrow we can do it in one cast
   (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_lossless_cast(x.dtype, a.dtype) else None),
-  (UPat.var('x', dtypes.ints+(dtypes.weakint,)).cast(dtypes.ints+(dtypes.weakint,), name="a").cast(name="b"),
-    lambda x,a,b: x.ccast(b.dtype) if not x.overflows(a.dtype) else None),
-  # try to do math in int instead of long, keep weak const weak
-  (UPat(GroupOp.Binary, src=(UPat.var("x", (dtypes.long, dtypes.weakint)), UPat.var("y", (dtypes.long, dtypes.weakint))), name="u"), lambda u,x,y:
-    (UOp.const(x.val) if x.op is Ops.CONST else x.cast(dtypes.int)).alu(u.op,
-     UOp.const(y.val) if y.op is Ops.CONST else y.cast(dtypes.int)).cast(u.dtype)
-    if dtypes.long in (x.dtype, y.dtype) and not any(v.overflows(dtypes.int) for v in (u,x,y)) else None),
+  (UPat.var('x', dtypes.wints).cast(dtypes.wints, name="a").cast(name="b"), lambda x,a,b: x.ccast(b.dtype) if not x.overflows(a.dtype) else None),
+  # try to do math in int instead of long, keep weak const weak=
+  (UPat(GroupOp.Binary, src=(UPat(dtype=(dtypes.long, dtypes.weakint)),)*2, name="u"), lambda u:
+    UOp(u.op, src=tuple(UOp.const(s.val) if s.op is Ops.CONST else s.cast(dtypes.int) for s in u.src)).cast(u.dtype)
+    if dtypes.long in (s.dtype for s in u.src) and not any(v.overflows(dtypes.int) for v in (u,*u.src))
+    and not (u.op is Ops.SHR and u.src[1].vmax >= 32)
+    and not (u.op in (Ops.MOD, Ops.TRUNCMOD) and u.src[0].vmin <= dtypes.int.min and u.src[1].vmin <= -1 <= u.src[1].vmax) else None),
   ((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+cast.const_like(c.val)),
   # an AFTER waits only on the effect ops listed here, any other dep is replaced by its srcs
   (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
     tuple(dedup(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.BARRIER, Ops.END, Ops.BACKEDGE, Ops.LINEAR, Ops.STAGE}
-                        else y.src for y in x.src[1:]]))))),
+                        else (UOp(Ops.NOOP).after(y.src[0]).end(*y.src[1:]),) if y.op is Ops.REDUCE else y.src for y in x.src[1:]]))))),
   # after/end with 1 src is just src[0]
   (UPat((Ops.AFTER, Ops.END), src=(UPat.var("s"),)), lambda s: s),
   # the rules above key on bare CONSTs, so a redundantly committed const has to be uncast in the same fixpoint

tinygrad_repo/tinygrad/viz/index.html

@@ -136,7 +136,7 @@
     stroke-linejoin: miter;
   }
   .label :is(text, p) {
-    font-weight: 350;
+    font-weight: 300;
   }
   g.node rect {
     stroke-width: 1.4;

tinygrad_repo/tinygrad/viz/js/index.js

@@ -843,6 +843,8 @@ function codeBlock(st, language, { loc, wrap }={}) {
   if (!language || language === "txt") code.appendChild(colored(st));
   else code.innerHTML = hljs.highlight(st, { language }).value;
   code.className = "hljs";
+  code.tabIndex = 0;
+  code.onkeydown = (e) => { if ((e.metaKey || e.ctrlKey) && e.key === "a") { e.preventDefault(); window.getSelection().selectAllChildren(code); }};
   const ret = document.createElement("pre");
   if (wrap) ret.className = "wrap";
   if (loc != null) ret.appendChild(pathLink(loc[0], loc[1]).style("margin-bottom", "4px").node());
@@ -850,7 +852,7 @@ function codeBlock(st, language, { loc, wrap }={}) {
   return ret;
 }
 function traceBlock(trace) {
-  const root = d3.create("pre").append("code").classed("hljs", true);
+  const root = d3.select(codeBlock("", "txt")).select("code");
   for (let i=trace.length-1; i>=0; i--) {
     const [fp, lineno, fn, code] = trace[i];
     root.append("div").style("margin-bottom", "2px").style("display","flex").text(fn+" ").append(() => pathLink(fp, lineno).node());

tinygrad_repo/tinygrad/viz/js/worker.js

@@ -24,7 +24,7 @@ onmessage = (e) => {
 const layoutCfg = (g, { blocks, paths, pc_tokens }) => {
   const lineHeight = 18;
   g.setGraph({ rankdir:"TD", font:"monospace", lh:lineHeight, textSpace:"1ch" });
-  ctx.font = `350 ${lineHeight}px ${g.graph().font}`;
+  ctx.font = `300 ${lineHeight}px ${g.graph().font}`;
   // basic blocks render the assembly in nodes
   const tokenColors = {0:"#7aa2f7", 1:"#9aa5ce"};
   for (const [lead, members] of Object.entries(blocks)) {
@@ -54,7 +54,7 @@ const layoutCfg = (g, { blocks, paths, pc_tokens }) => {
 const layoutUOp = (g, { graph, change }, opts) => {
   const lineHeight = 14;
   g.setGraph({ rankdir: "LR", font:"sans-serif", lh:lineHeight });
-  ctx.font = `350 ${lineHeight}px ${g.graph().font}`;
+  ctx.font = `300 ${lineHeight}px ${g.graph().font}`;
   if (change?.length) g.setNode("overlay", {label:"", labelWidth:0, labelHeight:0, labelX:0, className:"overlay"});
   let callCount = 0;
   for (const [k, {label, src, ref, color, exclude, addrspace}] of Object.entries(graph)) {

upstream-sync.json

@@ -43,8 +43,8 @@
     {
       "path": "tinygrad_repo",
       "url": "https://github.com/tinygrad/tinygrad.git",
-      "commit": "d39227beebd0c3a340dc42b2cdc1e3db1c36aed9",
-      "tree": "b02c108aacae57aea75b417c43b1e1caa22ee609",
+      "commit": "ec91b74b5b6b0fc924b72d3f27689679ca2438b9",
+      "tree": "1cb22a40b0aa242b2e432fc7b88f8b85fc391f81",
       "exclude": [
         "AGENTS.md"
       ]