casted CONST migration for nir [pr] (#17588)

before rendering, rewrite strong typed const to casted weak const and have renderer adopt the new UOp. starting with PYTHON

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

@@ -42,7 +42,11 @@ inputs:
     required: false
     default: 'false'
   qemu:
-    description: "Install qemu"
+    description: "Install qemu?"
+    required: false
+    default: 'false'
+  ninja:
+    description: "Install ninja?"
     required: false
     default: 'false'
 runs:
@@ -130,7 +134,7 @@ runs:
 
     # ******************* apt *******************
     - name: Setup apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       shell: bash
       run: |
         sudo mkdir -p /var/cache/apt/archives
@@ -158,7 +162,7 @@ runs:
         echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list
 
     - name: Compute Package List + Hash
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       id: apt-pkgs
       shell: bash
       run: |
@@ -183,25 +187,29 @@ runs:
         if [[ "${{ inputs.qemu }}" == "true" ]]; then
           pkgs+=" qemu-user-static"
         fi
+        # **** ninja ****
+        if [[ "${{ inputs.ninja }}" == "true" ]]; then
+          pkgs+=" ninja-build"
+        fi
 
         echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
         echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$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') && github.event_name == 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request'
       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') && github.event_name != 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request'
       uses: actions/cache@v5
       with:
         path: /var/cache/apt/archives/
         key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
 
     - name: Run apt Update + Install
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       shell: bash
       run: |
         sudo apt -qq update || true

.github/workflows/autogen.yml

@@ -35,7 +35,7 @@ jobs:
         key: 'autogen'
         amd: 'true'
         llvm: 'true'
-        pydeps: 'pyyaml mako'
+        deps: 'autogen'
     - name: Install autogen support packages
       run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
     - name: Regenerate autogen files

.github/workflows/docs.yml

@@ -8,7 +8,7 @@ permissions:
   contents: write
 jobs:
   deploy:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     steps:
       - uses: actions/checkout@v6
       - name: Configure Git Credentials

.github/workflows/python-publish.yml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   deploy:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     steps:
     - uses: actions/checkout@v6
     - name: Set up Python

.github/workflows/szdiff.yml

@@ -10,7 +10,7 @@ concurrency:
 jobs:
   checkbranch:
     name: Check PR Branch status
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     outputs:
       branchstat: ${{ steps.brstat.outputs.stat}}
     steps:
@@ -44,7 +44,7 @@ jobs:
     permissions:
       contents: read
       pull-requests: write
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     needs: checkbranch
     if: needs.checkbranch.outputs.branchstat == 'false'
     steps:
@@ -87,7 +87,7 @@ jobs:
     name: Core Library Line Difference
     permissions:
       pull-requests: write
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     needs: checkbranch
     if: needs.checkbranch.outputs.branchstat == 'true'
     steps:

.github/workflows/test.yml

@@ -31,8 +31,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
-        deps: docs
-        pydeps: "capstone torch"
+        deps: "docs testing_minimal"
     - name: Build wheel and show size
       run: |
         uv build --wheel
@@ -73,10 +72,7 @@ jobs:
         deps: testing_unit
         pydeps: "pillow torchvision expecttest"
         llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
+        ninja: 'true'
     - name: Test ResNet-18
       run: DEBUG=2 python3 extra/torch_backend/example.py
     - name: Test one op in torch tests
@@ -98,12 +94,8 @@ jobs:
       with:
         key: torch-backend-pillow-torchvision-et-pt
         deps: testing_unit
-        pydeps: "pillow torchvision expecttest"
         llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
+        ninja: 'true'
     - name: Test beautiful_mnist in torch with TINY_BACKEND
       run: STEPS=20 DEV=CPU TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
 

examples/mlperf/model_train.py

@@ -1742,8 +1742,8 @@ def train_gptoss():
   )
 
   for p in optim.params:
-    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
-    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
+    p.grad = p.zeros_like(dtype=dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype).contiguous()
+    if getattr(p, "_zero2", False): p.grad = optim.optimizers[0]._zero_shard(p.grad)
   grads = [p.grad for p in optim.params]
 
   from extra.gemm.cdna_asm_gemm import _mx_block_scale

examples/mlperf/models/gpt_oss.py

@@ -146,6 +146,7 @@ class GPTOSS:
       return w_q, w_e8.is_param_(False)
     if moe:
       qs = [_one(*shape[1:]) for _ in range(shape[0])]
+      for q in qs: q[0]._zero2 = True  # grad arrives sharded on the expert axis under ZeRO-2 (moe_gemm)
       return [q[0] for q in qs], [q[1] for q in qs]
     return _one(*shape)
 

extra/gemm/moe_gemm.py

@@ -1,10 +1,32 @@
 import functools, pathlib
 from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.helpers import getenv
 from tinygrad.renderer import Estimates
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from extra.gemm.cdna_asm_gemm import quantize_mxfp8, _mx_block_scale, _mx_block_scale_3d
 
+ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
+
+def reduce_scatter_devaxis(out:Tensor, shard_axis:int=0) -> Tensor:
+  # out: sharded on the device axis, shape (ndev, *rest); return the device-axis sum left sharded on shard_axis.
+  u = out.uop
+  devs, rest = u.device, u.shape[1:]
+  assert rest[shard_axis] % len(devs) == 0, f"reduce_scatter needs even shards: {rest[shard_axis]} % {len(devs)}"
+  # reach the raw per-device buffer below the UNSHARD, keeping the AFTERs so reads stay ordered after the kernel writes
+  node, barriers = u, []
+  while node.op is not Ops.UNSHARD:
+    if node.op is Ops.AFTER: barriers += node.src[1:]
+    node = node.src[0]
+  mbuf = node.src[0].after(*barriers) if barriers else node.src[0]
+  sz = rest[shard_axis] // len(devs)
+  shards = []
+  for i in range(len(devs)):
+    bounds = tuple((0,s) if a != shard_axis else (i*sz,(i+1)*sz) for a,s in enumerate(rest))
+    contribs = [mbuf.mselect(j).reshape(rest).shrink(bounds).copy_to_device(devs[i]) for j in range(len(devs))]
+    shards.append(functools.reduce(lambda a,b: a.alu(Ops.ADD, b), contribs))
+  return Tensor(UOp.mstack(*shards).unshard(shard_axis, UOp.range(len(devs), -1, AxisType.DEVICE)), device=devs)
+
 @functools.cache
 def custom_hk_grouped_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:UOp, dname:str, n_experts:int) -> UOp:
   M, K = A.shape
@@ -58,7 +80,8 @@ def grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> T
   out = Tensor(inv.uop.unshard(0), device=g.device) if is_multi else inv
   out = Tensor.custom_kernel(out, gT, xT, g_si, x_si, expert_off,
                              fxn=functools.partial(custom_hk_grouped_mxfp8_wgrad, dname=dname, n_experts=n_experts))[0]
-  out = out.sum(0) if is_multi else out.squeeze(0)
+  if is_multi and ZERO_OPTIM: out = reduce_scatter_devaxis(out, 0)
+  else: out = out.sum(0) if is_multi else out.squeeze(0)
   return out.reshape(n_experts, N, K)
 
 def mx_pack_3d(e8:Tensor) -> Tensor:

extra/hcq2/ops_amd2.py

@@ -179,11 +179,10 @@ class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TR
 
 def sdma_copy(ctx, call):
   sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
-  src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
-  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in (
-    ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
-    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0,
-    *data64_le(src_addr+UOp.const(off, dtypes.uint64)), *data64_le(dst_addr+UOp.const(off, dtypes.uint64)))))
+  hdr = ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR)
+  return call.ins(SDMAOps.COPY, src=tuple(x for off in range(0, sz, ctx.max_copy_size) for x in (
+    *(UOp.const(v, dtypes.uint32) for v in (hdr, ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0)),
+    *(a + UOp.const(off, dtypes.uint64) if off else a for a in (call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs))))))
 
 def sdma_wait(ctx, ins, dst, val):
   op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \

pyproject.toml

@@ -111,6 +111,10 @@ docs = [
   "numpy",
 ]
 mesa = ["tinymesa==25.2.7.2"]
+autogen = [
+  "pyyaml",
+  "mako",
+]
 
 
 [tool.mutmut]

test/backend/test_custom_kernel.py

@@ -4,7 +4,7 @@ import numpy as np
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import assert_kernel_count
+from test.helpers import assert_kernel_count, KernelCountException
 
 # **** kernels ****
 
@@ -474,7 +474,7 @@ class TestCustomKernelInput(unittest.TestCase):
     y.realize()
     kernel_count = GlobalCounters.kernel_count
     self.assertEqual(y.tolist(), x.add(1).tolist())
-    self.assertLessEqual(kernel_count, max_kernels)
+    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
     # same test with @function, input is PARAM
     from tinygrad import function
     x0 = Tensor.arange(32).clone("CPU").realize()
@@ -487,7 +487,7 @@ class TestCustomKernelInput(unittest.TestCase):
     y = run(x0).realize()
     kernel_count = GlobalCounters.kernel_count
     self.assertEqual(y.tolist(), mop_fxn(x0).add(1).tolist())
-    self.assertLessEqual(kernel_count, max_kernels)
+    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
 
   def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=2)
   def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=3)

test/backend/test_linearizer.py

@@ -252,7 +252,7 @@ class TestLinearizer(unittest.TestCase):
     for u in uops:
       if u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.REG:
         if uops.index(u) < begin_range:
-          assert u.src[1].op is Ops.CONST
+          assert u.src[1].op not in GroupOp.ALU
         else:
           assert u.src[1].op in GroupOp.ALU
           assert begin_range < uops.index(u) < end_range
@@ -261,6 +261,7 @@ class TestLinearizer(unittest.TestCase):
         assert end_range < uops.index(u)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipIf(Device[Device.DEFAULT].renderer.casted_consts, "reads a literal, which is casted here. TODO: flip this")
   def test_default_global_reversed(self):
     # shrink so that the dims do not collapse
     t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))

test/backend/test_multitensor.py

@@ -6,7 +6,7 @@ from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import run_linear, compile_linear, pm_beam, pm_compile
 import numpy as np
 from hypothesis import given, strategies as strat, settings
-from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count
+from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count, KernelCountException
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -395,7 +395,7 @@ class TestMultiBufferView(unittest.TestCase):
     linear, var_vals = b_multi.linear_with_vars()
     if all(not d.startswith(("WEBGPU", "CL")) for d in b_multi.device):
       compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
-      self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
+      if len(compiled) != 0: raise KernelCountException(0, len(compiled))
     run_linear(linear, var_vals)
     np.testing.assert_equal(b_multi.numpy(), b_ref.numpy())
 

test/backend/test_pickle.py

@@ -3,6 +3,7 @@ import numpy as np
 from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
 from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
 from tinygrad.uop.ops import PatternMatcher, UPat, UOp, deconstruct_function
+from test.helpers import KernelCountException
 
 class TestPickle(unittest.TestCase):
   def test_pickle_code_object(self):
@@ -41,7 +42,7 @@ class TestPickle(unittest.TestCase):
     t2:Tensor = pickle.loads(st)
     np.testing.assert_equal(t_values, t2.numpy())
     # expect at most one COPY kernel
-    self.assertLessEqual(GlobalCounters.kernel_count, 1)
+    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)
 
   def test_pickle_realized_tensor_alt(self):
     print("** init")

test/null/test_mnist_dataset.py

@@ -1,6 +1,7 @@
 import unittest
 from tinygrad.helpers import GlobalCounters
 from tinygrad.nn.datasets import mnist
+from test.helpers import KernelCountException
 
 class TestDataset(unittest.TestCase):
   def test_dataset_is_realized(self):
@@ -8,7 +9,7 @@ class TestDataset(unittest.TestCase):
     X_train[0].contiguous().realize()
     GlobalCounters.reset()
     X_train[0].contiguous().realize()
-    self.assertLessEqual(GlobalCounters.kernel_count, 1)  # 0 if SLICE (zero-copy), 1 otherwise
+    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)  # 0 if SLICE (zero-copy), 1 otherwise
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_allreduce.py

@@ -2,6 +2,7 @@ import unittest
 from tinygrad import Tensor, UOp, dtypes
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops
+from test.helpers import KernelCountException
 
 class TestRingAllReduce(unittest.TestCase):
   def test_schedule_ring(self):
@@ -13,7 +14,7 @@ class TestRingAllReduce(unittest.TestCase):
       copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
       pairs = [(c.src[1].buffer.device, c.src[2].buffer.device) for c in copies]
       # N*(N-1) scatter reduce, and N*(N-1) allgather
-      self.assertEqual(len(pairs), N*(N-1)*2)
+      if len(pairs) != N*(N-1)*2: raise KernelCountException(N*(N-1)*2, len(pairs))
       # copy topology forms a ring
       self.assertEqual(len(set(pairs)), N)
 
@@ -25,8 +26,8 @@ class TestRingAllReduce(unittest.TestCase):
       linear = t.sum(0).mul(2.0).contiguous().linear_with_vars()[0]
       copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
       sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
-      self.assertEqual(len(copies), 24)
-      self.assertEqual(len(sinks), 26)
+      if len(copies) != 24: raise KernelCountException(24, len(copies))
+      if len(sinks) != 26: raise KernelCountException(26, len(sinks))
 
   @Context(RING=0, ALL2ALL=0)
   def test_schedule_naive(self):
@@ -39,8 +40,8 @@ class TestRingAllReduce(unittest.TestCase):
     sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
     pairs = [(c.src[1].buffer.device, c.src[2].buffer.device) for c in copies]
 
-    self.assertEqual(len(pairs), N*(N-1))
-    self.assertEqual(len(sinks), 2)
+    if len(pairs) != N*(N-1): raise KernelCountException(N*(N-1), len(pairs))
+    if len(sinks) != 2: raise KernelCountException(2, len(sinks))
     self.assertTrue(all(dst != src for dst, src in pairs))
 
   def test_symbolic_shape(self):

test/unit/test_attention.py

@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from tinygrad import Tensor, dtypes
+from tinygrad import Tensor, dtypes, nn
 from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
   apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
@@ -45,10 +45,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)
 
   def _make_config(self, **kwargs):
-    return TransformerConfig(**({"num_blocks":1, "dim":4, "hidden_dim":8, "n_heads":1, "n_kv_heads":1,
-                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":4, "rope_theta":10000.0,
-                                 "rope_dim":4, "v_head_dim":4, "max_context":4, "ssm_layers":(True,),
-                                 "ssm":SSMConfig(conv_kernel=2, state_size=2, group_count=1, time_step_rank=1, inner_size=2)} | kwargs))
+    return TransformerConfig(**({"num_blocks":1, "dim":32, "hidden_dim":64, "n_heads":1, "n_kv_heads":1,
+                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":32, "rope_theta":10000.0,
+                                 "rope_dim":32, "v_head_dim":32, "max_context":4, "ssm_layers":(True,),
+                                 "ssm":SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32)} | kwargs))
 
   def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock:
     block = GatedDeltaNetBlock(config, config.ssm)
@@ -79,6 +79,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
       recurrent_state = cache[:, conv_flat:].reshape(cache.shape[0], block.num_v_heads, block.head_v_dim, block.head_v_dim)
       return conv_state, recurrent_state
 
+  def _reset_state(self, block:GatedDeltaNetBlock):
+    Tensor.realize(block.conv_state.assign(block.conv_state.const_like(0)),
+                   block.recurrent_state.assign(block.recurrent_state.const_like(0)))
+
   def _linear_np(self, x:np.ndarray, weight:np.ndarray) -> np.ndarray:
     return x.astype(np.float32) @ weight.T.astype(np.float32)
 
@@ -86,7 +90,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     x_float = x.astype(np.float32)
     return (x_float / np.sqrt((x_float * x_float).mean(axis=-1, keepdims=True) + eps)) * weight.astype(np.float32)
 
-  def _normalize_np(self, x:np.ndarray, eps:float=1e-12) -> np.ndarray:
+  def _normalize_np(self, x:np.ndarray, eps:float=1e-6) -> np.ndarray:
     return x / np.maximum(np.sqrt((x * x).sum(axis=-1, keepdims=True)), eps)
 
   def _softplus_np(self, x:np.ndarray) -> np.ndarray:
@@ -148,6 +152,12 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     x = Tensor.linspace(-1.0, 1.0, 3 * config.dim, dtype=dtypes.float32).reshape(1, 3, config.dim)
 
     expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, x)
+    out = self._run_attention(block, x, 0)
+    conv_state, recurrent_state = self._cache_views(block)
+    np.testing.assert_allclose(out, np.concatenate(expected_outs, axis=1), rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(conv_state, expected_conv[-1], rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(recurrent_state, expected_recurrent[-1], rtol=1e-3, atol=1e-3)
+    self._reset_state(block)
 
     for step in range(x.shape[1]):
       out = self._run_attention(block, x[:, step:step+1], step)
@@ -163,7 +173,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     prompt = Tensor.linspace(0.75, -0.75, 2 * config.dim, dtype=dtypes.float32).reshape(1, 2, config.dim)
 
     for i in range(warmup.shape[1]): self._run_attention(block, warmup[:, i:i+1], i)
-    Tensor.realize(*block._state_reset_ops())
+    self._reset_state(block)
     expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, prompt)
 
     for step in range(prompt.shape[1]):
@@ -177,18 +187,64 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
                                  err_msg=f"GatedDeltaNet reset recurrent cache mismatch at step {step}")
 
   def test_kda_channel_decay(self):
-    config = self._make_config(n_heads=2, ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True))
-    block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]])
-    # f_b(f_a(x)) = [1, 2, 3, 4]
+    config = self._make_config(dim=4, hidden_dim=8, n_heads=2, head_dim=4, rope_dim=4, v_head_dim=4,
+      ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True))
+    block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.], [2., 1., 0., 0.]]])
     block.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]])
     block.ssm_f_b.weight = Tensor([[1., 0.], [0., 1.], [1., 1.], [2., 1.]])
     block._init_state(x)
     initial_state = Tensor.arange(8, dtype=dtypes.float32).reshape(1, 2, 2, 2)
     block.recurrent_state.assign(initial_state).realize()
     block.ssm_a = Tensor([[-1.], [-1.]])
-    block._attention(x, 0).realize()
-    alpha = np.exp(-self._softplus_np(np.arange(1, 5)).reshape(1, 2, 1, 2))
-    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=1e-5, atol=1e-5)
+    block._attention(x, x.shape[1]).realize()
+    alpha = np.exp(-self._softplus_np(np.array([[1, 2, 3, 4], [2, 1, 3, 5]])).reshape(2, 2, 2)).prod(0)
+    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha[..., None], rtol=1e-5, atol=1e-5)
+
+  def test_kda_prefill_matches_decode(self):
+    config = self._make_config(ssm=SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=True))
+    block = GatedDeltaNetBlock(config, config.ssm)
+    for p in nn.state.get_parameters(block):
+      p.replace(self._tensor_linspace(-0.05, 0.05, p.shape) if len(p.shape) > 1 else self._tensor_linspace(0.05, 0.1, p.shape))
+    x = self._tensor_linspace(-0.5, 0.5, (1, 3, config.dim))
+    prefill = self._run_attention(block, x, 0)
+    prefill_conv, prefill_recurrent = self._cache_views(block)
+    self._reset_state(block)
+    decode = np.concatenate([self._run_attention(block, x[:, i:i+1], i) for i in range(3)], axis=1)
+    decode_conv, decode_recurrent = self._cache_views(block)
+    np.testing.assert_allclose(prefill, decode, rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(prefill_conv, decode_conv, rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(prefill_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3)
+
+  def test_varied_chunk_sizes_match_decode(self):
+    for kda in (False, True):
+      ssm = SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=kda)
+      config = self._make_config(ssm=ssm)
+      if kda:
+        block = GatedDeltaNetBlock(config, config.ssm)
+        for p in nn.state.get_parameters(block):
+          p.replace(self._tensor_linspace(-0.05, 0.05, p.shape) if len(p.shape) > 1 else self._tensor_linspace(0.05, 0.1, p.shape))
+      else: block = self._make_block(config)
+      x = self._tensor_linspace(-0.5, 0.5, (1, 4, config.dim))
+      decode = np.concatenate([self._run_attention(block, x[:, i:i+1], i) for i in range(4)], axis=1)
+      decode_conv, decode_recurrent = self._cache_views(block)
+      for chunking in ([4], [2, 2], [1, 3], [3, 1], [2, 1, 1]):
+        self._reset_state(block)
+        outs, start = [], 0
+        for size in chunking:
+          outs.append(self._run_attention(block, x[:, start:start+size], start))
+          start += size
+        chunked_conv, chunked_recurrent = self._cache_views(block)
+        np.testing.assert_allclose(np.concatenate(outs, axis=1), decode, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")
+        np.testing.assert_allclose(chunked_conv, decode_conv, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")
+        np.testing.assert_allclose(chunked_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")
+
+  def test_start_zero_resets_realized_state(self):
+    config, x = self._make_config(max_context=3), self._tensor_linspace(-1, 1, (1, 3, 32))
+    block = self._make_block(config)
+    self._run_attention(block, x, 0)
+    restarted = self._run_attention(block, x[:, :2], 0)
+    fresh = self._run_attention(self._make_block(config), x[:, :2], 0)
+    np.testing.assert_allclose(restarted, fresh, rtol=1e-3, atol=1e-3)
 
 class TestPairwiseTopk(unittest.TestCase):
   def test_basic_topk(self):

test/unit/test_dtype_weak.py

@@ -3,11 +3,12 @@ import tempfile, unittest, math
 from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import Context
 from tinygrad.dtype import least_upper_float
-from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite
+from tinygrad.uop.ops import UOp, Ops, GroupOp, dtype_from_uop, graph_rewrite
 from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.uop.spec import spec_shared, type_verify
 from tinygrad.engine.jit import JitError
+from test.helpers import full_rewrite
 
 
 class TestWeakPromotion(unittest.TestCase):
@@ -288,5 +289,18 @@ class TestSignedUint64Weakfloat(unittest.TestCase):
     self.assertAlmostEqual((i64 + u64).sin().item(), math.sin(2), places=5)  # Unary lowers before transcendental
 
 
+class TestNoRedundantWide(unittest.TestCase):
+  def wide_alu(self, t:Tensor) -> int:
+    return sum(sum(1 for u in full_rewrite(call.src[0]).toposort() if u.op in GroupOp.ALU and u.dtype in {dtypes.long, dtypes.ulong})
+               for call in t.schedule_linear().src if call.src[0].op is Ops.SINK)
+
+  def test_unbounded_long_stays_long(self):
+    self.assertGreater(self.wide_alu(Tensor.empty(16, dtype=dtypes.long)*3 + 1), 0)
+
+  def test_fancy_index_has_no_wide_alu(self):
+    j, o = Tensor([0, 1, 2]).reshape(3, 1), Tensor([0, 1]).reshape(1, 2)
+    self.assertEqual(self.wide_alu(Tensor.empty(8, 9, 10, 11, 12)[1, j, 2, o, 2]), 0)
+
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_function.py

@@ -4,7 +4,7 @@ from tinygrad.function import function
 from tinygrad import Tensor, GlobalCounters, Device
 from tinygrad.dtype import Invalid
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, ProgramInfo
-from test.helpers import assert_kernel_count
+from test.helpers import assert_kernel_count, KernelCountException
 
 class TestFunction(unittest.TestCase):
   def test_simple(self):
@@ -516,7 +516,7 @@ class TestFunctionTuple(unittest.TestCase):
     Tensor.realize(a)
     c = f(a)
 
-    self.assertEqual(count_kernels(c), 1)
+    if count_kernels(c) != 1: raise KernelCountException(1, count_kernels(c))
 
     c.sum().backward()
     Tensor.realize(a.grad)

tinygrad/codegen/__init__.py

@@ -5,7 +5,7 @@ from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey,
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp, AxisType
 from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak, pm_cast_weak
 from tinygrad.uop.render import pyrender
-from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
+from tinygrad.uop.spec import type_verify, spec_tensor, spec_program, spec_program_casted_consts
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
 from tinygrad.dtype import dtypes, AddrSpace
@@ -153,8 +153,8 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
   # unpack WMMA
   (UPat(Ops.WMMA, name="u"), do_stack_wmma),
   # stacked INDEX is many INDEX
-  (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.STACK, name="s"))),
-   lambda b,s: UOp.stack(*[b.index(u) for u in s.src])),
+  (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.STACK, name="s")), name="x"),
+   lambda b,s,x: UOp.stack(*[x.replace(src=(b,u)) for u in s.src])),
   # INDEX into RESHAPE moves the RESHAPE
   (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.RESHAPE, name="s"))),
    lambda b,s: b.index(s.src[0]).reshape(s.shape)),
@@ -281,6 +281,10 @@ pm_implicit_barriers = PatternMatcher([
   (UPat(Ops.END, name="end"), add_war_barrier),
 ])
 
+pm_casted_consts = PatternMatcher([
+  (UPat(Ops.CONST, dtypes.all, name="c"), lambda c: UOp(Ops.CAST, c.dtype, src=(UOp.const(c.val),), arg=c.dtype)),
+])
+
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
   if DEBUG >= 5: print(pyrender(ast))
@@ -383,8 +387,11 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   num_params = len([x for x in sink.toposort() if x.op is Ops.PARAM and x.arg.slot != -1])
   sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)
 
+  # TODO: delete once migration are done
+  if ren.casted_consts: sink = graph_rewrite(sink, pm_casted_consts, name="casted consts", walk=True)
+
   if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
-  if SPEC: type_verify(sink, spec_program)
+  if SPEC: type_verify(sink, spec_program_casted_consts if ren.casted_consts else spec_program)
 
   # return the rewritten sink
   return sink

tinygrad/codegen/late/coalesce.py

@@ -149,7 +149,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
         grp = full_grp[:length]
         # NOTE: we apply the valid again after we determine the length
         offset = offset.valid(valid) if valid is not None else offset
-        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset)
+        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset, dtype=offsets[grp[0]][0].src[0].dtype)
         if op == Ops.STORE:
           datas = []
           for i,g in enumerate(grp):

tinygrad/device.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
 from typing import Any, Callable, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
-import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal
+import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct
 from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
 from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up
@@ -310,6 +310,14 @@ class Compiler:
       if self.cachekey is not None: diskcache_put(self.cachekey, src, lib)
     return lib
   def disassemble(self, lib:bytes): pass
+  def server(self, cmd:str, arch:str, *args) -> subprocess.Popen:
+    argv = f"{cmd} {pathlib.Path(__file__).parent}/runtime/support/compileserver.py {type(self).__module__}:{type(self).__name__} {arch}"
+    return subprocess.Popen(argv.split() + [str(a) for a in args], stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0)
+  def compile_server(self, src:str, proc:subprocess.Popen) -> bytes:
+    unwrap(proc.stdin).write(struct.pack("I", len(src.encode())) + src.encode())
+    if (lib:=unwrap(proc.stdout).read(struct.unpack("I", unwrap(proc.stdout).read(4))[0])): return lib
+    raise CompileError("Compilation Error")
+
 
 @dataclass
 class TinyELF:

tinygrad/helpers.py

@@ -486,15 +486,11 @@ def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip
       if length and (file_size:=os.stat(fp).st_size) < length: raise RuntimeError(f"fetch size incomplete, {file_size} < {length}")
   return fp
 
-# not all firmware exists at the pinned ref; newer files can be pinned to the commit that introduced them without
-# affecting any other firmware (blob contents are checked by sha256 anyway)
-FW_REF = "1e2c15348485939baf1b6d1f5a7a3b799d80703d"
-FW_REF_OVERRIDES = {"psp_13_0_15_sos.bin": "23e6cdf0409383e29d681c8c14cd6ffd0f394f02"}
 def fetch_fw(path:str, name:str, sha256:str) -> bytes:
   if sys.version_info >= (3,14) and (p:=pathlib.Path(f"/lib/firmware/{path}/{name}.zst")).is_file():
     from compression.zstd import decompress
     if hashlib.sha256(b:=decompress(p.read_bytes())).hexdigest() == sha256: return b
-  return fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/{FW_REF_OVERRIDES.get(name, FW_REF)}/{path}/{name}",
+  return fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/1e2c15348485939baf1b6d1f5a7a3b799d80703d/{path}/{name}",
                subdir="fw", sha256=sha256).read_bytes()
 
 # *** Exec helpers

tinygrad/llm/model.py

@@ -138,8 +138,6 @@ class FFNBlock:
 
   # 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
-  # return writes that reset this block's state after a cache mismatch
-  def _state_reset_ops(self) -> list[Tensor]: return []
   def _init_state(self, x:Tensor): raise NotImplementedError
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: raise NotImplementedError
 
@@ -274,45 +272,65 @@ class GatedDeltaNetBlock(FFNBlock):
 
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     B, T, _ = x.shape
-    assert T == 1, "GatedDeltaNetBlock currently only supports T=1"
+    # bind ints to a variable so the reset flag stays a runtime value (it toggles when generation restarts at position 0)
+    start_pos = start_pos if isinstance(start_pos, UOp) else UOp.variable("start_pos", 0, self.config.max_context-1).bind(start_pos)
+    initial = Tensor(start_pos).eq(0)
     is_kda = hasattr(self, "ssm_g_a")
+    symbolic = isinstance(T, UOp)
+    T_pad = x.max_shape[1]  # symbolic chunks are padded to their max size: one graph serves every size
 
     # 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 = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
-    beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
+    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)
-    alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
-             self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
+    log_alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, T, self.num_v_heads, -1) *
+                 self.ssm_a.reshape(self.num_v_heads, -1))
 
-    # qkv conv
-    conv_window = self.conv_state.cat(self.attn_qkv(x), dim=1)
-    conv_out = (conv_window * self.ssm_conv1d["weight"].T.unsqueeze(0)).sum(1).silu()
+    # qkv conv, conv_state is reset when starting from position 0
+    conv_state = initial.where(0, self.conv_state)
+    # assemble the conv window in a static-size buffer: [conv_state | qkv rows | zero-pad].
+    # padded steps are exact no-ops: beta=0 (delta rule off), log_alpha=0 (decay 1 after exp)
+    win = Tensor.zeros(B, self.ssm_conv_kernel-1 + T_pad, self.conv_channels).uop
+    win = win.after(win[:, :self.ssm_conv_kernel-1].store(conv_state.cast(win.dtype).uop))
+    win = win.after(win[:, self.ssm_conv_kernel-1:self.ssm_conv_kernel-1+T].store(self.attn_qkv(x).cast(win.dtype).uop))
+    conv_window = Tensor(win)
+    # the last conv_kernel-1 columns of the window become the next conv state
+    conv_state_store = self.conv_state.uop.store(conv_window[:, T:T+self.ssm_conv_kernel-1].cast(self.conv_state.dtype).uop)
+
+    conv_out = functools.reduce(lambda a,b: a+b,
+      (conv_window[:, i:i+T_pad] * self.ssm_conv1d["weight"][:, i] for i in range(self.ssm_conv_kernel))).silu()
+    if symbolic:
+      out_gate = out_gate.pad_to((B, T_pad, self.num_v_heads, self.head_v_dim))
+      beta, log_alpha = beta.pad_to((B, T_pad, self.num_v_heads)), log_alpha.pad_to((B, T_pad, *log_alpha.shape[2:]))
     q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-1)
-    q = q.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
-    k = k.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
-    v = v.reshape(B, self.num_v_heads, self.head_v_dim)
-    q, k, v = q.mul(self.head_k_dim**-0.5).unsqueeze(-1), k.unsqueeze(-1), v.unsqueeze(-1)
+    qk_eps = 1e-12 if is_kda else 1e-6
+    q, k = (z.reshape(B, T_pad, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=qk_eps)
+            .repeat(1, 1, self.num_v_heads//self.num_k_heads, 1) for z in (q, k))
+    v = v.reshape(B, T_pad, self.num_v_heads, self.head_v_dim)
+    # layout the per-step operands to broadcast against the (B, H, V, K) state
+    q, k, v, beta = (z.transpose(1, 2).float() for z in (q, k, v, beta))
+    q, k, v, beta = q.unsqueeze(-2) * self.head_k_dim**-0.5, k.unsqueeze(-2), v.unsqueeze(-1), beta.unsqueeze(-1).unsqueeze(-1)
+    alpha = log_alpha.transpose(1, 2).exp().unsqueeze(-1)  # per-channel decay for kda, per-head otherwise (B, H, T, V|1, 1)
 
-    # recurrent
-    recurrent_state = self.recurrent_state * alpha
-    recurrent_state = recurrent_state + ((v - recurrent_state@k) * beta)@k.transpose(-1, -2)
+    # recurrent: scan over the (padded) tokens, updating the recurrent state. collect the per-step outputs
+    state = Tensor(self.recurrent_state.uop.after(conv_state_store)).float()  # carry the conv write into this graph
+    state = initial.where(0, state)
+    outs = []
+    for t in range(T_pad):
+      s1 = state * alpha[:, :, t]  # decay the state
+      delta = (v[:, :, t] - (s1*k[:, :, t]).sum(-1, keepdim=True)) * beta[:, :, t]  # the delta rule update
+      state = s1 + delta * k[:, :, t]
+      outs.append((state * q[:, :, t]).sum(-1))
 
-    # store the updated state
-    conv_state_store = self.conv_state.uop.store(conv_window[:, 1:, :].cast(self.conv_state.dtype).uop)
-    recurrent_state_store = self.recurrent_state.uop.store(recurrent_state.cast(self.recurrent_state.dtype).uop)
-    recurrent_state = Tensor(self.recurrent_state.uop.after(recurrent_state_store, conv_state_store))
+    # store the updated recurrent state in place, then read the stacked outputs after the write
+    core = Tensor(outs[0].stack(*outs[1:], dim=1).contiguous().uop.after(self.recurrent_state.uop.store(state.cast(self.recurrent_state.dtype).uop)))
 
-    # output
-    core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
-    out_gate = out_gate.sigmoid() if is_kda else out_gate.silu()
-    return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
-
-  # recurrent state can't be partially reused after divergence, force a full rebuild
-  def _state_reset_ops(self):
-    return [self.conv_state.assign(self.conv_state.const_like(0)),
-            self.recurrent_state.assign(self.recurrent_state.const_like(0))] if hasattr(self, "conv_state") else []
+    # 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()
+    if symbolic: z = z[:, :T]
+    return self.ssm_out(z.reshape(B, T, -1))
 
   def _init_state(self, x):
     if not hasattr(self, "conv_state"):
@@ -453,7 +471,6 @@ class Transformer:
     t = Tensor(tokens + [0] * (self.max_context - len(tokens)), dtype="int32").reshape(1, self.max_context)
     # recompute start_pos from what's currently valid in the caches
     start_pos = self.get_start_pos(tokens)
-    if start_pos < len(self._cached_tokens) and (resets := [r for b in self.blk for r in b._state_reset_ops()]): Tensor.realize(*resets)
     out, prompt_len = None, len(tokens)
     while len(tokens) < self.max_context:
       n_toks = min(chunk_size, len(tokens) - start_pos)

tinygrad/renderer/__init__.py

@@ -72,6 +72,8 @@ class Renderer:
   tensor_cores: list[TensorCore] = []
   extra_matcher: PatternMatcher|None = None
   code_for_op: dict[Ops, Callable] = {}
+  # migration: this renderer consumes every literal as a casted const CAST(dt, CONST(value))
+  casted_consts: bool = False
 
   compiler: Compiler = Compiler()
 

tinygrad/renderer/llvmir.py

@@ -81,8 +81,8 @@ base_rewrite = PatternMatcher([
   (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat((Ops.BUFFER, Ops.PARAM, Ops.AFTER)),), allow_any_len=True, name="x"), lambda ctx,x:
    f"  {ctx[x]} = getelementptr inbounds {ldt(x.dtype)}, {ldt(x.dtype, ptr=True)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"),
   # register index
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.cvar("idx")), name="x"), lambda ctx,buf,idx,x:
-   f"  {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {idx.val}" if buf.addrspace == AddrSpace.ALU else None),
+  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.cvar("c").cast()), name="x"), lambda ctx,buf,c,x:
+   f"  {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {c.val}" if buf.addrspace == AddrSpace.ALU else None),
 
   # load/store
   (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.var("alt"), UPat.var("mask")), name="x"),
@@ -146,6 +146,7 @@ base_rewrite = PatternMatcher([
 ])
 
 class LLVMRenderer(Renderer):
+  casted_consts = True
   abi: str | None
   string_rewrite: PatternMatcher
   code_for_op = {k:lambda:None for v in lop.values() for k in v.keys()}
@@ -165,7 +166,7 @@ class LLVMRenderer(Renderer):
     local_args: list[str] = []
     name = "test"
     for u in uops:
-      if u.op in {Ops.NOOP, Ops.GROUP}: continue
+      if u.op in {Ops.NOOP, Ops.GROUP, Ops.CONST}: continue
       if u.op is Ops.AFTER:
         r[u] = r[u.src[0]]
         continue
@@ -185,7 +186,7 @@ class LLVMRenderer(Renderer):
           kernel.append(f"  {r[u]} = addrspacecast [{size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{size} x {ldt(u.dtype)}]*")
         else:
           kernel.append(f"  {r[u]} = alloca [{size} x {ldt(u.dtype)}], align 16")
-      elif u.op is Ops.CONST: r[u] = lconst(u.val, u.dtype)
+      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):
         r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast
       else:

tinygrad/renderer/nir.py

@@ -116,6 +116,7 @@ def nidx(b:mesa.nir_builder, buf, off, space, itemsize, gate=None) -> mesa.nir_d
 
 class NIRRenderer(Renderer):
   suffix = "NIR"
+  casted_consts = True
   nir_options: bytes
   global_max, local_max, shared_max = CUDARenderer.global_max, CUDARenderer.local_max, CUDARenderer.shared_max
   code_for_op = {**{k:lambda:None for k in u_aop.keys()}, **{k:lambda:None for k in s_aop.keys()}, **{k:lambda:None for k in f_aop.keys()}}
@@ -145,7 +146,7 @@ class NIRRenderer(Renderer):
   ])
 
   def_rewrite = PatternMatcher([
-    (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.val, x.dtype)),
+    (UPat.cvar("c").cast(name="x"), lambda ctx,x,c: nimm(ctx.b, c.val, x.dtype)),
     (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, x.dtype.itemsize if x.addrspace is AddrSpace.ALU else 8)),
     (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
     (UPat(Ops.STORE, src=(UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val"))),
@@ -186,16 +187,17 @@ class NIRRenderer(Renderer):
 
   def render(self, uops:list[UOp]):
     self.prerender(uops)
-    for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].val
+    for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]:
+      self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].src[0].val
     self.r: dict[UOp, Any] = {}
     self.param_idx = 0
     ranges: list[mesa.nir_def|None] = []
 
     for u in uops:
-      if u.op in {Ops.NOOP, Ops.GROUP} or (u.op is Ops.STACK and len(u.src) == 0): pass
+      if u.op in {Ops.NOOP, Ops.GROUP, Ops.CONST} or (u.op is Ops.STACK and len(u.src) == 0): pass
       elif u.op in {Ops.INDEX, Ops.SHRINK}:
         # INDEX on a register value picks the element, memory INDEX is handled in the LOAD/STORE patterns
-        if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].val)
+        if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].src[0].val)
       elif u.op is Ops.AFTER:
         self.r[u] = self.r[u.src[0]]
       elif u.op == Ops.SINK:

tinygrad/renderer/ptx.py

@@ -79,8 +79,8 @@ def modifier(a: DType, b: DType): return '.rzi' if dtypes.is_int(a) and dtypes.i
   (a.itemsize < b.itemsize or dtypes.is_int(b) or b == dtypes.bool) else ''
 
 string_rewrite = PatternMatcher([
-  (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.val, x.dtype)}, 0;"),
-  (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.val, x.dtype)};"),
+  (UPat.cvar("c").cast(dtypes.bool, name="x"), lambda ctx, x, c: f"setp.ne.s16 {ctx.r[x]}, {render_val(c.val, x.dtype)}, 0;"),
+  (UPat.cvar("c").cast(name="x"), lambda ctx, x, c: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(c.val, x.dtype)};"),
   (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"),
   (UPat(Ops.PARAM, name="x"), lambda ctx, x:
    f"ld.param.{ctx.types[dtypes.ulong] if x.addrspace is AddrSpace.GLOBAL else ctx.mem_types[x.dtype]} {ctx.r[x]}, [data{x.arg.slot}+0];"),
@@ -136,6 +136,7 @@ string_rewrite = PatternMatcher([
 
 class PTXRenderer(Renderer):
   suffix = "PTX"
+  casted_consts = True
   global_max, local_max, shared_max = CUDARenderer.global_max, CUDARenderer.local_max, CUDARenderer.shared_max
   tc_sm80 = [x for x in tc.cuda_sm80 if x.dtype_in in [dtypes.half, dtypes.float]]
   code_for_op = asm_for_op
@@ -186,7 +187,7 @@ class PTXRenderer(Renderer):
 
     name = "test"
     for u in uops:
-      if u.op in {Ops.NOOP, Ops.GROUP}: continue
+      if u.op in {Ops.NOOP, Ops.GROUP, Ops.CONST}: continue
       if u.op is Ops.AFTER:
         self.r[u] = self.r[u.src[0]]
         continue
@@ -201,9 +202,9 @@ class PTXRenderer(Renderer):
         continue
       if u.op in {Ops.INDEX, Ops.SHRINK, Ops.LOAD} and u.src[0].addrspace in (AddrSpace.REG, AddrSpace.ALU):
         # on REG, INDEX/SHRINK pick the register (must be CONST) and LOAD is a noop
-        if u.op is not Ops.LOAD and u.src[1].op is not Ops.CONST:
+        if u.op is not Ops.LOAD and not (u.src[1].op is Ops.CAST and u.src[1].src[0].op is Ops.CONST):
           raise RuntimeError(f"PTX does not support dynamic register indexing: {u}")
-        r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].val]
+        r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].src[0].val]
         continue
       if u.op is Ops.SPECIAL: r[u] = "%" + u.arg
       elif u.op is Ops.LOAD:
@@ -216,7 +217,7 @@ class PTXRenderer(Renderer):
                        [ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.dtype.itemsize)]]
         r[u] = [ssa("wmma", dtype=self.types[u.dtype]) for _ in range(u.max_numel())]
       prefix, dtype = {Ops.CAST: ("cast", None), Ops.BITCAST: ("cast", None), Ops.END: ("pred", "pred"), Ops.RANGE: ("ridx", None),
-        Ops.CONST: ("const", None), Ops.BUFFER: ("local", "u64"), Ops.INDEX: ("bidx", "u64"), Ops.SHRINK: ("bidx", "u64"),
+        Ops.BUFFER: ("local", "u64"), Ops.INDEX: ("bidx", "u64"), Ops.SHRINK: ("bidx", "u64"),
         Ops.PARAM: ("dat", "u64" if u.addrspace is AddrSpace.GLOBAL else None), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None))
       if u.op is Ops.RANGE and u.dtype == dtypes.void: prefix = None  # loop headers don't have a register
       if prefix: r[u] = ssa(prefix, u, dtype)

tinygrad/runtime/autogen/am/fw.py

@@ -3,7 +3,6 @@ hashes = {
   'psp_13_0_10_sos.bin': '0bcaaad9cd8578d3841ae69155a6bd4fc3ceae8f4fb5a6ba4f576e7ace94d1d9',
   'psp_13_0_12_sos.bin': '89da90bf4286b38678b1fd175c78462a426afa3d258d15872cd14072d7098b9b',
   'psp_13_0_14_sos.bin': 'a4f0d5f76d27b77409ec0b71d7cc6a848ddfd29f8c84f3003edf74ad3999fb7d',
-  'psp_13_0_15_sos.bin': '3b28d53e75a88131155e3931378ac8434eca4880ada9211d3b4e8915b6289583',
   'psp_13_0_6_sos.bin': '27657daa0f91ad8095d3610224a7de748b8b348a4cb211ecb5fccabe47369716',
   'psp_13_0_7_sos.bin': 'ef1af0ecea38abbac6f85cce71789f19848c498d0cb8ef13748dab2d65b23c31',
   'psp_14_0_2_sos.bin': '7b538448b57d4f9dd06b2eea90d4f86a16e65e3027cdecee8db71c2c5f1fa243',

tinygrad/runtime/ops_amd.py

@@ -842,7 +842,7 @@ class KFDIface:
 
 class PCIIface(PCIIfaceBase):
   def __init__(self, dev, dev_id):
-    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0,0x75a8)),), vram_bar=0,
+    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0)),), vram_bar=0,
       va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev)
     self._compute_props()
 
@@ -880,11 +880,6 @@ class PCIIface(PCIIfaceBase):
       doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring.va_addr, ring.size, gart.va_addr+rptr, gart.va_addr+wptr,
         eop_buffer.va_addr, eop_buffer.size, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))
 
-    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA and self.dev_impl.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
-      # aqua (NBIO 7.9): kernel submits SDMA queues by writing the RB_WPTR register directly (SDMA 4.4.4 doorbell regs are firmware-managed)
-      doorbell = self.dev_impl.mmio.view(self.dev_impl.reg('regSDMA_GFX_RB_WPTR').addr[idx] * 4, 8, fmt='Q')
-      return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=doorbell, put_value=0, params=rcvr_params,
-        read_ptr=gart.cpu_view().view(offset=rptr, size=8, fmt='Q'), write_ptr=gart.cpu_view().view(offset=wptr, size=8, fmt='Q'))
     return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'), put_value=0,
       read_ptr=gart.cpu_view().view(offset=rptr, size=8, fmt='Q'), write_ptr=gart.cpu_view().view(offset=wptr, size=8, fmt='Q'), params=rcvr_params)
 

tinygrad/runtime/ops_disk.py

@@ -49,8 +49,7 @@ class DiskDevice(Compiled):
     DiskDevice._tried_io_uring_init = True
 
     if sys.platform == 'linux' and not hasattr(sys, "getandroidapilevel"):
-      p = io_uring.struct_io_uring_params(flags=io_uring.IORING_SETUP_SQPOLL, sq_thread_idle=0xffffffff)
-      fd = libc.syscall(io_uring.NR_io_uring_setup, 4096, ctypes.byref(p))
+      fd = libc.syscall(io_uring.NR_io_uring_setup, 4096, ctypes.byref(p:=io_uring.struct_io_uring_params()))
       if fd < 0: return
 
       sq_ptr = libc.mmap(0, p.sq_off.array + p.sq_entries * 4, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_POPULATE, fd, 0)
@@ -68,7 +67,6 @@ class DiskDevice(Compiled):
         kring_mask=u32ptr(sq_ptr+p.cq_off.ring_mask), cqes=ctypes.cast(cq_ptr+p.cq_off.cqes, ctypes.POINTER(io_uring.struct_io_uring_cqe)))
 
       DiskDevice.io_uring = io_uring.struct_io_uring(ring_fd=fd, sq=sqdesc, cq=cqdesc) # type: ignore
-      libc.syscall(io_uring.NR_io_uring_enter, fd, 0, 0, io_uring.IORING_ENTER_SQ_WAKEUP)
 
 class DiskBuffer:
   def __init__(self, device:DiskDevice, size:int, offset=0):
@@ -126,6 +124,7 @@ class DiskAllocator(Allocator):
         # Send sqe
         DiskDevice.io_uring.sq.array[sqe_index] = sqe_index
         DiskDevice.io_uring.sq.ktail[0] = tail + 1
+        libc.syscall(io_uring.NR_io_uring_enter, DiskDevice.io_uring.ring_fd, 1, 1, io_uring.IORING_ENTER_GETEVENTS)
 
         reqs.append((copy_batch, copied_in, minor_offset, real_copy_size:=min(sqe.len - minor_offset, size - copied_in)))
         next_read_offset += sqe.len

tinygrad/runtime/ops_python.py

@@ -23,7 +23,9 @@ def load(inp, j, dtype: DType):
 
 def _store(m, i, v, dtype: DType):
   if i < 0 or i >= len(m): raise IndexError(f"store out of bounds, size is {len(m)}, access is {i}, value is {v}")
-  m[i] = to_storage_scalar(v, dtype)
+  if (w:=m.nbytes // len(m)) >= dtype.itemsize: m[i] = to_storage_scalar(v, dtype)
+  else:
+    for k in range(dtype.itemsize // w): m[i+k] = (v >> 8*w*k) & ((1 << 8*w) - 1)
 
 # here are the models for the WMMA instruction on the different hardware
 def generic_wmma_helper(inp, warp_size, WARP_THREADS, K, NUM_A, NUM_B, NUM_C, a_elem, b_elem, c_map):
@@ -212,6 +214,7 @@ class PythonCompiler(Compiler):
 class PythonRenderer(Renderer):
   code_for_op = python_alu
   compiler = PythonCompiler()
+  casted_consts: bool = True
 
   def __init__(self, target:Target):
     assert (emu:=getenv("EMULATE", "")) == "", ("EMULATE is deprecated, use DEV=PYTHON::" +

tinygrad/runtime/support/am/amdev.py

@@ -152,9 +152,7 @@ class AMDev:
     self._run_discovery()
     self._build_regs()
 
-    # on NBIO 7.9 the HDP flush doorbell remap register must be programmed before any flush (nbio_v7_9_remap_hdp_registers).
-    # the silicon default is bogus and flushing without this hangs the chip. flush_hdp is used before SOC init (by the PSP).
-    if self.ip_ver[am.NBIO_HWIP][:2] == (7,9): self.reg("regBIF_BX0_REMAP_HDP_MEM_FLUSH_CNTL").write(0x1A000)
+    # AM boot Process:
     # The GPU being passed can be in one of several states: 1. Not initialized. 2. Initialized by amdgpu. 3. Initialized by AM.
     # The 1st and 2nd states require a full GPU setup since their states are unknown. The 2nd state also requires a mode1 reset to
     # reinitialize all components.
@@ -174,21 +172,14 @@ class AMDev:
 
     # Init hw for IP blocks where it is needed
     if not self.partial_boot:
-      fw_is_ours = False
       if self.psp.is_sos_alive() and self.smu.is_smu_alive():
         self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2)
         if self.is_hive():
           if reset_mode: return # in reset mode, do not raise
           raise RuntimeError("Malformed state. Use extra/amdpci/hive_reset.py to reset the hive")
-        # mode1 reset is only needed for foreign/unknown firmware state. If the running firmware was set up by AM itself
-        # (SCRATCH_REG7 is ours), a full AM re-init can run on top of it; SMU mode1 leaves the PSP/BL dead on some chips.
-        fw_is_ours = self.reg("regSCRATCH_REG7").read() == AMDev.Version
-        if not fw_is_ours: self.smu.mode1_reset()
+        self.smu.mode1_reset()
       self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
-      # when the firmware is AM's own and still running, skip the PSP stage: PSP ring commands over a live sOS are not
-      # serviced between sessions, and its firmware is already loaded.
-      self.init_hw(*([self.soc, self.gmc, self.ih, self.smu] if (self.psp.is_sos_alive() and self.smu.is_smu_alive() and fw_is_ours) else
-                      [self.soc, self.gmc, self.ih, self.psp, self.smu]))
+      self.init_hw(self.soc, self.gmc, self.ih, self.psp, self.smu)
 
     # Booting done
     self.is_booting = False
@@ -196,9 +187,7 @@ class AMDev:
     # Re-initialize main blocks
     self.init_hw(self.gfx, self.sdma)
 
-    # TODO: MP0 13.0.15 PMFW doesn't answer DPM clock msgs without the full amdgpu pptable/DPM setup, skip clock programming
-    if self.ip_ver[am.MP0_HWIP] == (13,0,15) and (max_power:=0.0) == 0.0: pass
-    elif (max_power:=getenv("AM_POWER_LIMIT", 0.0)) > 0:
+    if (max_power:=getenv("AM_POWER_LIMIT", 0.0)) > 0:
       self.smu.set_power_limit(max_power)
       self.smu.set_clocks(level=None)
     else: self.smu.set_clocks(level=-1) # last level, max perf.
@@ -249,8 +238,7 @@ class AMDev:
     if DEBUG >= 3: print(f"am {self.devfmt}: Recovery complete")
     return True
 
-  # a hive has multiple XGMI regions; single-node parts (like MI350P) may still program LFB_SIZE with region 0 only
-  def is_hive(self) -> bool: return self.gmc.xgmi_seg_sz > 0 and self.gmc.xgmi_max_region > 0
+  def is_hive(self) -> bool: return self.gmc.xgmi_seg_sz > 0
 
   def paddr2mc(self, paddr:int) -> int: return self.gmc.mc_base + paddr
   def paddr2xgmi(self, paddr:int) -> int: return self.gmc.paddr_base + paddr
@@ -327,18 +315,6 @@ class AMDev:
 
         ip_offset += 8 + (8 if ihdr.base_addr_64_bit else 4) * ip.num_base_address
 
-    # HARV(EST) table: harvested instances must be excluded (like amdgpu_discovery_harvest_ip)
-    self.harvested:dict[int, set[int]] = collections.defaultdict(set)
-    if (harv_off:=self.bhdr.table_list[am.HARVEST_INFO].offset) != 0:
-      hv = ctypes.c_uint32.from_address(ctypes.addressof(self.bhdr) + harv_off).value
-      if hv == am.HARVEST_TABLE_SIGNATURE:
-        for i in range(32):
-          hw_id = ctypes.c_uint16.from_address(ctypes.addressof(self.bhdr) + harv_off + 8 + i*4).value
-          if hw_id == 0: continue
-          inst = ctypes.c_uint8.from_address(ctypes.addressof(self.bhdr) + harv_off + 8 + i*4 + 2).value
-          for hw_ip in am.hw_id_map:
-            if am.hw_id_map[hw_ip] == hw_id: self.harvested[hw_ip].add(inst)
-
     gc_info = am.struct_gc_info_v1_0.from_address(gc_addr:=ctypes.addressof(self.bhdr) + self.bhdr.table_list[am.GC].offset)
     self.gc_info = getattr(am, f"struct_gc_info_v{gc_info.header.version_major}_{gc_info.header.version_minor}").from_address(gc_addr)
     self.reserved_vram_size = (384 << 20) if self.ip_ver[am.GC_HWIP][:2] in {(9,4), (9,5)} else (64 << 20)

tinygrad/runtime/support/am/ip.py

@@ -29,8 +29,7 @@ class AM_SOC(AM_IP):
 
   def init_hw(self):
     if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
-      # kernel programs regXCC_DOORBELL_FENCE = 0xff & ~xcc_mask; on this PF 4 of 8 XCCs are harvested, so fence xcc4-7
-      self.adev.regXCC_DOORBELL_FENCE.write(0xF0)
+      self.adev.regXCC_DOORBELL_FENCE.write(0x0)
       for aid in range(1, self.adev.gmc.vmhubs):
         self.adev.indirect_wreg_pcie(self.adev.regXCC_DOORBELL_FENCE.addr[0], self.adev.regXCC_DOORBELL_FENCE.encode(shub_slv_mode=1), aid=aid)
       self.adev.regBIFC_GFX_INT_MONITOR_MASK.write(0x7ff)
@@ -51,19 +50,15 @@ class AM_SOC(AM_IP):
 class AM_GMC(AM_IP):
   def init_sw(self):
     self.vmhubs = len(self.adev.regs_offset[am.MMHUB_HWIP])
-    # aqua (NBIO 7.9): only the first 2 mmhubs exist in the host window, instances 2+ are phantom layouts (like amdgpu's aid_mask)
-    if self.adev.ip_ver[am.NBIO_HWIP][:2] == (7,9): self.vmhubs = min(self.vmhubs, 2)
 
     # XGMI (for supported systems)
-    xgmi_lfb_cntl = self.adev.regMMMC_VM_XGMI_LFB_CNTL.read_bitfields() if hasattr(self.adev, 'regMMMC_VM_XGMI_LFB_CNTL') else {}
-    self.xgmi_phys_id, self.xgmi_max_region = xgmi_lfb_cntl.get('pf_lfb_region', 0), xgmi_lfb_cntl.get('pf_max_region', 0)
+    self.xgmi_phys_id = self.adev.regMMMC_VM_XGMI_LFB_CNTL.read_bitfields()['pf_lfb_region'] if hasattr(self.adev, 'regMMMC_VM_XGMI_LFB_CNTL') else 0
     self.xgmi_seg_sz = self.adev.regMMMC_VM_XGMI_LFB_SIZE.read_bitfields()['pf_lfb_size']<<24 if hasattr(self.adev, 'regMMMC_VM_XGMI_LFB_SIZE') else 0
 
     self.paddr_base = self.xgmi_phys_id * self.xgmi_seg_sz
 
-    # compute fb_end like the kernel does (vram_start + vram_size), MMMC_VM_FB_LOCATION_TOP is not reliable on all SKUs
     self.fb_base = (self.adev.regMMMC_VM_FB_LOCATION_BASE.read() & 0xFFFFFF) << 24
-    self.fb_end = self.fb_base + self.adev.vram_size
+    self.fb_end = (self.adev.regMMMC_VM_FB_LOCATION_TOP.read() & 0xFFFFFF) << 24
 
     # Memory controller aperture
     self.mc_base = self.fb_base + self.paddr_base
@@ -181,22 +176,10 @@ class AM_SMU(AM_IP):
     self.smu_mod = self.adev._ip_module("smu", am.MP1_HWIP)
     self.driver_table_paddr = self.adev.mm.palloc(0x4000, zero=False, boot=True)
 
-  def wait_alive(self):
-    # poll until the SMU mailbox starts ACKing GetSmuVersion (single-shot attempts, mirroring amdgpu which issues one check)
-    t0 = time.time()
-    while time.time() - t0 < 60:
-      if self.is_smu_alive(): return
-      time.sleep(0.5)
-    raise TimeoutError("SMU not alive")
-
   def init_hw(self):
-    self.wait_alive()
-    # MP0 13.0.15 PMFW answers the dram addr msgs with an error response (as seen in amdgpu logs), tolerate any nonzero resp
-    dram_tolerant = self.adev.ip_ver[am.MP0_HWIP] == (13,0,15)
-    self._send_msg(self.smu_mod.PPSMC_MSG_SetDriverDramAddrHigh, hi32(self.adev.paddr2mc(self.driver_table_paddr)), any_resp=dram_tolerant)
-    self._send_msg(self.smu_mod.PPSMC_MSG_SetDriverDramAddrLow, lo32(self.adev.paddr2mc(self.driver_table_paddr)), any_resp=dram_tolerant)
-    # not valid on smu_v13_0_12-family pmfw
-    if self.adev.ip_ver[am.MP0_HWIP] != (13,0,15): self._send_msg(self.smu_mod.PPSMC_MSG_EnableAllSmuFeatures, 0, any_resp=dram_tolerant)
+    self._send_msg(self.smu_mod.PPSMC_MSG_SetDriverDramAddrHigh, hi32(self.adev.paddr2mc(self.driver_table_paddr)))
+    self._send_msg(self.smu_mod.PPSMC_MSG_SetDriverDramAddrLow, lo32(self.adev.paddr2mc(self.driver_table_paddr)))
+    self._send_msg(self.smu_mod.PPSMC_MSG_EnableAllSmuFeatures, 0)
 
   def is_smu_alive(self):
     with contextlib.suppress(TimeoutError): self._send_msg(self.smu_mod.PPSMC_MSG_GetSmuVersion, 0, timeout=100)
@@ -206,13 +189,13 @@ class AM_SMU(AM_IP):
     if DEBUG >= 2: print(f"am {self.adev.devfmt}: mode1 reset")
     if self.adev.ip_ver[am.MP0_HWIP] >= (14,0,0) or self.adev.ip_ver[am.MP0_HWIP] in {(13,0,0), (13,0,7), (13,0,10)}:
       self._send_msg(__DEBUGSMC_MSG_Mode1Reset:=2, 0, debug=True)
-    elif self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6), (13,0,12), (13,0,15)}: self._send_msg(self.smu_mod.PPSMC_MSG_GfxDriverReset, 1)
+    elif self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6), (13,0,12)}: self._send_msg(self.smu_mod.PPSMC_MSG_GfxDriverReset, 1)
     else: self._send_msg(self.smu_mod.PPSMC_MSG_Mode1Reset, 0)
 
     if not self.adev.is_hive(): time.sleep(0.5) # 500ms
 
   def read_table(self, table_t, arg):
-    if self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6),(13,0,12),(13,0,15)}: self._send_msg(self.smu_mod.PPSMC_MSG_GetMetricsTable, arg)
+    if self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6),(13,0,12)}: self._send_msg(self.smu_mod.PPSMC_MSG_GetMetricsTable, arg)
     else: self._send_msg(self.smu_mod.PPSMC_MSG_TransferTableSmu2Dram, arg)
     return table_t.from_buffer(bytearray(self.adev.vram.view(self.driver_table_paddr, ctypes.sizeof(table_t))[:]))
 
@@ -223,7 +206,7 @@ class AM_SMU(AM_IP):
 
   def set_clocks(self, level:int|None):
     clks = tuple([self.smu_mod.PPCLK_UCLK, self.smu_mod.PPCLK_FCLK, self.smu_mod.PPCLK_SOCCLK])
-    if self.adev.ip_ver[am.MP0_HWIP] not in {(13,0,6), (13,0,12), (13,0,15)}: clks += (self.smu_mod.PPCLK_GFXCLK,)
+    if self.adev.ip_ver[am.MP0_HWIP] not in {(13,0,6), (13,0,12)}: clks += (self.smu_mod.PPCLK_GFXCLK,)
 
     if level is None:
       for clck in clks:
@@ -255,27 +238,22 @@ class AM_SMU(AM_IP):
     (self.adev.mmMP1_SMN_C2PMSG_82 if not debug else self.adev.mmMP1_SMN_C2PMSG_53).write(param)
     (self.adev.mmMP1_SMN_C2PMSG_66 if not debug else self.adev.mmMP1_SMN_C2PMSG_75).write(msg)
 
-  def _send_msg(self, msg:int, param:int, read_back_arg=False, timeout=10000, debug=False, any_resp=False): # default timeout is 10 seconds
+  def _send_msg(self, msg:int, param:int, read_back_arg=False, timeout=10000, debug=False): # default timeout is 10 seconds
     self._smu_cmn_send_msg(msg, param, debug=debug)
-    rc = self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54
-    # amdgpu tolerates any nonzero resp
-    cond, val = (lambda: rc.read() != 0, True) if any_resp else (rc.read, 1)
-    wait_cond(cond, value=val, timeout_ms=timeout, msg=f"SMU msg {msg:#x} timeout")
+    wait_cond((self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).read, value=1, timeout_ms=timeout,
+      msg=f"SMU msg {msg:#x} timeout")
     return (self.adev.mmMP1_SMN_C2PMSG_82 if not debug else self.adev.mmMP1_SMN_C2PMSG_53).read() if read_back_arg else None
 
 class AM_GFX(AM_IP):
   def init_sw(self):
-    self.xccs = sum(1 for i in self.adev.regs_offset[am.GC_HWIP] if i not in self.adev.harvested[am.GC_HWIP])
+    self.xccs = len(self.adev.regs_offset[am.GC_HWIP])
     self.mqd_paddr = [self.adev.mm.palloc(0x1000 * self.xccs, zero=False, boot=True) for i in range(2)]
     self.mqd_mc = [self.adev.paddr2mc(mqd_paddr) for mqd_paddr in self.mqd_paddr]
 
   def init_hw(self):
     # Wait for RLC autoload to complete
-    # regRLC_RLCS_BOOTLOAD_STATUS doesn't exist on gc 9.4.3 (used for gfx942/gfx950), gate on it only if present
-    def bootload_done():
-      return getattr(self.adev, 'regRLC_RLCS_BOOTLOAD_STATUS', None) is None or \
-        self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0
-    wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or bootload_done(), value=True, msg="RLC autoload timeout")
+    wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0,
+              value=True, msg="RLC autoload timeout")
 
     self.adev.gmc.init_hub("GC", inst_cnt=self.xccs)
     if self.adev.partial_boot: return self.reset_mec()
@@ -319,8 +297,8 @@ class AM_GFX(AM_IP):
 
     self._enable_mec()
 
-    # Set 1 partition (skip on MP0 13.0.15 (MI350P): the XCP transition is firmware-owned there)
-    if self.xccs > 1 and self.adev.ip_ver[am.MP0_HWIP] != (13,0,15): self.adev.psp._spatial_partition_cmd(1)
+    # Set 1 partition
+    if self.xccs > 1: self.adev.psp._spatial_partition_cmd(1)
 
   def fini_hw(self): self._dequeue_hqds()
 
@@ -336,9 +314,7 @@ class AM_GFX(AM_IP):
     self._enable_mec()
 
   def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, idx:int, aql:bool) -> int:
-    # aqua (NBIO 7.9) uses DOORBELL_LAYOUT1 (see aqua_vanjaram_doorbell_index_init): its mec ring0 starts at 8, not 3
-    pipe, queue, doorbell = idx // 4, idx % 4, (am.AMDGPU_DOORBELL_LAYOUT1_MEC_RING_START if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}
-      else am.AMDGPU_NAVI10_DOORBELL_MEC_RING0)
+    pipe, queue, doorbell = idx // 4, idx % 4, am.AMDGPU_NAVI10_DOORBELL_MEC_RING0
 
     for xcc in range(self.xccs if aql else 1):
       self._grbm_select(me=1, pipe=pipe, queue=queue, inst=xcc)
@@ -436,23 +412,15 @@ class AM_IH(AM_IP):
     def _alloc_ring(size): return (self.adev.mm.palloc(size, zero=False, boot=True), self.adev.mm.palloc(0x1000, zero=False, boot=True))
     self.rings = [(*_alloc_ring(self.ring_size), "", 0), (*_alloc_ring(self.ring_size), "_RING1", 1)]
     self.ring_view = self.adev.vram.view(offset=self.rings[0][0], size=self.ring_size, fmt='I')
-    # on gfx950 (aqua), the IH rings must live in host system memory (like use_bus_addr=true in amdgpu)
-    self.rings_in_sysmem = self.adev.ip_ver[am.GC_HWIP][:2] == (9,5) # scoped to gfx950 for now (validated symptom there)
-    if self.rings_in_sysmem:
-      # OSSSYS 4.4.2 (aqua): only one IH ring, the second one is skipped in amdgpu too
-      self.sysmem_rings = [self.adev.pci_dev.alloc_sysmem(self.ring_size + 0x1000) for _ in range(1)]
-      self.rings = [(sr[1][0], sr[1][self.ring_size // 0x1000], s, i) for sr, (_, _, s, i) in zip(self.sysmem_rings, self.rings)]
-      self.ring_view = self.sysmem_rings[0][0].view(0, self.ring_size, fmt='I')
 
   def init_hw(self):
     for ring_vm, rwptr_vm, suf, ring_id in self.rings:
-      self.adev.wreg_pair("regIH_RB_BASE", suf, f"_HI{suf}", ring_vm >> 8)
+      self.adev.wreg_pair("regIH_RB_BASE", suf, f"_HI{suf}", self.adev.paddr2mc(ring_vm) >> 8)
 
-      mc_space = 1 if self.rings_in_sysmem else 4
-      self.adev.reg(f"regIH_RB_CNTL{suf}").write(mc_space=mc_space, wptr_overflow_clear=1, rb_size=((self.ring_size//4)-1).bit_length(),
+      self.adev.reg(f"regIH_RB_CNTL{suf}").write(mc_space=4, wptr_overflow_clear=1, rb_size=((self.ring_size//4)-1).bit_length(),
         mc_snoop=1, mc_ro=0, mc_vmid=0, **({'wptr_overflow_enable': 1, 'rptr_rearm': 1} if ring_id == 0 else {'rb_full_drain_enable': 1}))
 
-      if ring_id == 0: self.adev.wreg_pair("regIH_RB_WPTR_ADDR", "_LO", "_HI", (rwptr_vm if self.rings_in_sysmem else self.adev.paddr2mc(rwptr_vm)))
+      if ring_id == 0: self.adev.wreg_pair("regIH_RB_WPTR_ADDR", "_LO", "_HI", self.adev.paddr2mc(rwptr_vm))
 
       self.adev.reg(f"regIH_RB_WPTR{suf}").write(0)
       self.adev.reg(f"regIH_RB_RPTR{suf}").write(0)
@@ -464,12 +432,6 @@ class AM_IH(AM_IP):
       self.adev.regIH_INT_FLOOD_CNTL.update(flood_cntl_enable=1)
       self.adev.regIH_MSI_STORM_CTRL.update(delay=3)
 
-    # aqua (OSSSYS 4.4.2): IH_CHICKEN.MC_SPACE_GPA_ENABLE + retry-int-cam must be set before RB_ENABLE (as in vega20_ih)
-    if self.rings_in_sysmem and hasattr(self.adev, 'regIH_CHICKEN'):
-      self.adev.regIH_CHICKEN.update(mc_space_gpa_enable=1)
-      oss_base = self.adev.regs_offset[am.OSSSYS_HWIP][0][0]
-      self.adev.wreg(oss_base + 0xEA, self.adev.rreg(oss_base + 0xEA) | 0x10000) # IH_RETRY_INT_CAM_CNTL_ALDEBARAN
-
     # toggle interrupts
     for _, rwptr_vm, suf, ring_id in self.rings:
       self.adev.reg(f"regIH_RB_CNTL{suf}").update(rb_enable=1, **({'enable_intr': 1} if ring_id == 0 else {}))
@@ -536,9 +498,6 @@ class AM_IH(AM_IP):
 class AM_SDMA(AM_IP):
   def init_sw(self): self.sdma_reginst, self.sdma_name = [], "F32" if self.adev.ip_ver[am.SDMA0_HWIP] < (7,0,0) else "MCU"
   def init_hw(self):
-    # aqua (NBIO 7.9): SDMA doorbell routing/trap config is firmware/RLC-managed; host programming here tears the fabric
-    # (~40ms later: RAS_ATHUB_ERR_EVENT and host BAR0 access to VRAM dies until the next cold boot).
-    if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}: return
     for pipe_id in range(16 if self.adev.ip_ver[am.SDMA0_HWIP] < (5,0,0) else 1):
       pipe, inst = ("", pipe_id) if self.adev.ip_ver[am.SDMA0_HWIP] < (5,0,0) else (str(pipe_id), 0)
 
@@ -589,14 +548,11 @@ class AM_SDMA(AM_IP):
     self.adev.wreg_pair(f"{reg}_RB_BASE", "", "_HI", ring_addr >> 8, inst=inst)
     self.adev.wreg_pair(f"{reg}_RB_RPTR_ADDR", "_LO", "_HI", rptr_addr, inst=inst)
     self.adev.wreg_pair(f"{reg}_RB_WPTR_POLL_ADDR", "_LO", "_HI", wptr_addr, inst=inst)
-    # aqua (NBIO 7.9): kernel leaves SDMA doorbell regs 0 and submits via the WPTR register
-    if self.adev.ip_ver[am.NBIO_HWIP] not in {(7,9,0), (7,9,1)}:
-      self.adev.reg(f"{reg}_DOORBELL_OFFSET").update(offset=doorbell * 2, inst=inst)
-      self.adev.reg(f"{reg}_DOORBELL").update(enable=1, inst=inst)
+    self.adev.reg(f"{reg}_DOORBELL_OFFSET").update(offset=doorbell * 2, inst=inst)
+    self.adev.reg(f"{reg}_DOORBELL").update(enable=1, inst=inst)
     self.adev.reg(f"{reg}_MINOR_PTR_UPDATE").write(0x0, inst=inst)
     self.adev.reg(f"{reg}_RB_CNTL").write(**({f'{self.sdma_name.lower()}_wptr_poll_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP][:2]!=(4,4) else {}),
-      rb_vmid=0, rptr_writeback_enable=1, rptr_writeback_timer=4, rb_enable=1,
-      rb_priv=1 if self.adev.ip_ver[am.NBIO_HWIP] not in {(7,9,0), (7,9,1)} else 0, rb_size=(ring_size//4).bit_length()-1, inst=inst)
+      rb_vmid=0, rptr_writeback_enable=1, rptr_writeback_timer=4, rb_enable=1, rb_priv=1, rb_size=(ring_size//4).bit_length()-1, inst=inst)
     self.adev.reg(f"{reg}_IB_CNTL").update(ib_enable=1, inst=inst)
     return doorbell
 
@@ -619,15 +575,13 @@ class AM_PSP(AM_IP):
     self.ring_paddr = self.adev.mm.palloc(self.ring_size, zero=False, boot=True)
 
     self.max_tmr_size, self.tmr_size = 0x1300000, 0
-    self.boot_time_tmr = self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6), (13,0,14), (13,0,15), (14,0,2), (14,0,3)}
-    self.autoload_tmr = self.adev.ip_ver[am.MP0_HWIP] not in {(13,0,6), (13,0,14), (13,0,15)}
+    self.boot_time_tmr = self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6), (13,0,14), (14,0,2), (14,0,3)}
+    self.autoload_tmr = self.adev.ip_ver[am.MP0_HWIP] not in {(13,0,6), (13,0,14)}
     self.tmr_paddr = self.adev.mm.palloc(self.max_tmr_size, align=am.PSP_TMR_ALIGNMENT, zero=False, boot=True) if not self.boot_time_tmr else 0
 
   def init_hw(self):
     spl_key = am.PSP_FW_TYPE_PSP_SPL if self.adev.ip_ver[am.MP0_HWIP] >= (14,0,0) else am.PSP_FW_TYPE_PSP_KDB
-    # SPL is preloaded on MP0 13.0.15
-    sos_components = [] if self.adev.ip_ver[am.MP0_HWIP] == (13,0,15) else [(spl_key, am.PSP_BL__LOAD_TOS_SPL_TABLE)]
-    sos_components += [(am.PSP_FW_TYPE_PSP_KDB, am.PSP_BL__LOAD_KEY_DATABASE),
+    sos_components = [(am.PSP_FW_TYPE_PSP_KDB, am.PSP_BL__LOAD_KEY_DATABASE), (spl_key, am.PSP_BL__LOAD_TOS_SPL_TABLE),
       (am.PSP_FW_TYPE_PSP_SYS_DRV, am.PSP_BL__LOAD_SYSDRV), (am.PSP_FW_TYPE_PSP_SOC_DRV, am.PSP_BL__LOAD_SOCDRV),
       (am.PSP_FW_TYPE_PSP_INTF_DRV, am.PSP_BL__LOAD_INTFDRV), (am.PSP_FW_TYPE_PSP_DBG_DRV, am.PSP_BL__LOAD_DBGDRV),
       (am.PSP_FW_TYPE_PSP_RAS_DRV, am.PSP_BL__LOAD_RASDRV), (am.PSP_FW_TYPE_PSP_SOS, am.PSP_BL__LOAD_SOSDRV)]

tinygrad/runtime/support/compiler_cuda.py

@@ -1,10 +1,12 @@
 import hashlib, tempfile, ctypes, re, pathlib
-from tinygrad.helpers import to_char_p_p, colored, getenv, system
+from tinygrad.helpers import to_char_p_p, colored, getenv, system, OSX
 from tinygrad.runtime.support.c import init_c_var
 from tinygrad.runtime.autogen import nvrtc, nvjitlink as jitlink
 from tinygrad.device import Compiler, CompileError
 
 CUDA_PATH = getenv("CUDA_PATH", "")
+root = pathlib.Path(__file__).parents[3]
+osx_docker_cmd = f"docker run --rm -i -v {root}:{root} -e PYTHONPATH={root} ghcr.io/tinygrad/cuda-arm64:v2.3"
 
 def _get_bytes(arg, get_str, get_sz, check) -> bytes:
   x = ctypes.create_string_buffer(init_c_var(ctypes.c_size_t, lambda x: check(get_sz(arg, ctypes.byref(x)))).value)
@@ -44,11 +46,14 @@ def cuda_disassemble(lib:bytes, arch:str, ptx=False):
 class NVRTCCompiler(Compiler):
   def __init__(self, arch:str, ptx=True, cache_key:str="cuda"):
     self.ptx, self.arch, self.compile_options = ptx, arch, [f'--gpu-architecture={arch}']
-    self.compile_options += [f"-I{CUDA_PATH}/include"] if CUDA_PATH else ["-I/usr/local/cuda/include", "-I/usr/include", "-I/opt/cuda/include"]
-    nvrtc_check(nvrtc.nvrtcVersion((nvrtcMajor := ctypes.c_int()), (nvrtcMinor := ctypes.c_int())))
-    if (nvrtcMajor.value, nvrtcMinor.value) >= (12, 4): self.compile_options.append("--minimal")
+    if OSX: self.compiler_process = self.server(osx_docker_cmd, arch, ptx)
+    else:
+      self.compile_options += [f"-I{CUDA_PATH}/include"] if CUDA_PATH else ["-I/usr/local/cuda/include", "-I/usr/include", "-I/opt/cuda/include"]
+      nvrtc_check(nvrtc.nvrtcVersion((nvrtcMajor := ctypes.c_int()), (nvrtcMinor := ctypes.c_int())))
+      if (nvrtcMajor.value, nvrtcMinor.value) >= (12, 4): self.compile_options.append("--minimal")
     super().__init__(f"compile_{cache_key}_{self.arch}")
   def compile(self, src:str) -> bytes:
+    if OSX: return self.compile_server(src, self.compiler_process)
     nvrtc_check(nvrtc.nvrtcCreateProgram(ctypes.byref(prog := nvrtc.nvrtcProgram()), src.encode(), "<null>".encode(), 0, None, None))
     nvrtc_check(nvrtc.nvrtcCompileProgram(prog, len(self.compile_options), to_char_p_p([o.encode() for o in self.compile_options])), prog)
     data = _get_bytes(prog, nvrtc.nvrtcGetPTX if self.ptx else nvrtc.nvrtcGetCUBIN,
@@ -80,9 +85,11 @@ class PTXCompiler(Compiler):
 
 class NVPTXCompiler(PTXCompiler):
   def __init__(self, arch:str):
-    jitlink_check(jitlink.nvJitLinkVersion(ctypes.byref(ctypes.c_uint()), ctypes.byref(ctypes.c_uint())))
+    if OSX: self.compiler_process = self.server(osx_docker_cmd, arch)
+    else: jitlink_check(jitlink.nvJitLinkVersion(ctypes.byref(ctypes.c_uint()), ctypes.byref(ctypes.c_uint())))
     super().__init__(arch, cache_key="nv_ptx")
   def compile(self, src:str) -> bytes:
+    if OSX: return self.compile_server(src, self.compiler_process)
     jitlink_check(jitlink.nvJitLinkCreate(handle := jitlink.nvJitLinkHandle(), 1, to_char_p_p([f'-arch={self.arch}'.encode()])), handle)
     jitlink_check(jitlink.nvJitLinkAddData(handle, jitlink.NVJITLINK_INPUT_PTX, ptxsrc:=super().compile(src), len(ptxsrc), "<null>".encode()), handle)
     jitlink_check(jitlink.nvJitLinkComplete(handle), handle)

tinygrad/runtime/support/compiler_qcom.py

@@ -1,6 +1,6 @@
-import ctypes, struct, platform, pathlib, shutil, subprocess, sys, tarfile, tempfile
+import ctypes, struct, platform, pathlib, shutil, tarfile, tempfile
 from tinygrad.device import Compiler
-from tinygrad.helpers import DEBUG, system, fetch, unwrap
+from tinygrad.helpers import DEBUG, system, fetch
 from tinygrad.runtime.support.compiler_mesa import disas_adreno
 # see https://github.com/sirhcm/tinydreno
 from tinygrad.runtime.autogen import llvm_qcom
@@ -12,12 +12,11 @@ class QCOMCompiler(Compiler):
     assert arch.split(',')[0] == "a630", "only a630 supported"
     if platform.machine() == "aarch64": self.arch, self.chip_id, self.llvm_inst = arch, 0x6030001, llvm_qcom.cl_compiler_create_llvm_instance()
     else:
-      self.arch, self.chip_id, self.fs = arch, 0x6030001, tempfile.TemporaryDirectory()
+      self.arch, self.chip_id, self.fs, root = arch, 0x6030001, tempfile.TemporaryDirectory(), pathlib.Path(__file__).parents[3]
       with tarfile.open(fetch('https://git.tinygrad.win/sirhcm/images/releases/download/v2/qcomcl.tar.gz')) as t: t.extractall(fs:=self.fs.name)
-      if (qemu:=shutil.which("qemu-aarch64-static")): argv = f"{qemu} -cpu max,pauth=off -L {fs} {fs}/usr/bin/python3 {__file__} {arch}"
-      else: argv = (f"docker run --rm -i --platform linux/aarch64 -v {fs}/usr:/usr -v {pathlib.Path(__file__).parents[2]}:/tinygrad "
-                    f"-e PYTHONPATH=/ -e QEMU_CPU=max,pauth=off gcr.io/distroless/static python3 /tinygrad/runtime/support/compiler_qcom.py {arch}")
-      self.compiler_process = subprocess.Popen(argv.split(), stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0)
+      self.compiler_process = self.server(f"{qemu} -cpu max,pauth=off -L {fs} {fs}/usr/bin/python3" if (qemu:=shutil.which("qemu-aarch64-static"))
+                                          else (f"docker run --rm -i --platform linux/aarch64 -v {fs}/usr:/usr -v {root}:{root} "
+                                                f"-e PYTHONPATH={root} -e QEMU_CPU=max,pauth=off gcr.io/distroless/static python3"), arch)
     super().__init__(f"compile_qcomcl_{arch}")
 
   def __del__(self): llvm_qcom.cl_compiler_destroy_llvm_instance(self.llvm_inst) if platform.machine() == "aarch64" else self.compiler_process.kill()
@@ -32,10 +31,7 @@ class QCOMCompiler(Compiler):
     return handle
 
   def compile(self, src) -> bytes:
-    if platform.machine() != "aarch64":
-      unwrap(self.compiler_process.stdin).write(struct.pack("I", len(src.encode())) + src.encode())
-      if (lib:=unwrap(self.compiler_process.stdout).read(struct.unpack("I", unwrap(self.compiler_process.stdout).read(4))[0])): return lib
-      raise RuntimeError("QCOM Compilation Error")
+    if platform.machine() != "aarch64": return self.compile_server(src, self.compiler_process)
     ch = self.checked(llvm_qcom.cl_compiler_compile_source(self.llvm_inst, self.chip_id, llvm_qcom.CL_MODE_64BIT, b"", 0, 0, 0, src.encode(), 0,
                                                            llvm_qcom.CL_SRC_STR, None))
     if DEBUG >= 8: print(system("llvm-dis", input=ctypes.string_at((comp:=ch.contents.compiled.contents).llvm_bitcode, comp.llvm_bitcode_size)))
@@ -48,12 +44,3 @@ class QCOMCompiler(Compiler):
 
   def disassemble(self, lib: bytes): disas_adreno(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)], self.chip_id)
 
-if __name__ == "__main__":
-  compiler = QCOMCompiler(sys.argv[1])
-  while (amt:=sys.stdin.buffer.read(4)):
-    try: lib = compiler.compile(sys.stdin.buffer.read(struct.unpack("I", amt)[0]).decode())
-    except Exception as e:
-      lib = b""
-      print(e, file=sys.stderr, flush=True)
-    sys.stdout.buffer.write(struct.pack("I", len(lib)) + lib)
-    sys.stdout.buffer.flush()

tinygrad/runtime/support/compileserver.py

@@ -0,0 +1,13 @@
+import ast, struct, sys
+from tinygrad.helpers import fromimport
+
+if __name__ == "__main__":
+  assert len(sys.argv) >= 3, f"usage: {sys.argv[0]} <compiler> <arch> [<args>]"
+  compiler = fromimport(*sys.argv[1].split(':'))(sys.argv[2], *(ast.literal_eval(arg) for arg in sys.argv[3:]))
+  while (amt:=sys.stdin.buffer.read(4)):
+    try: lib = compiler.compile(sys.stdin.buffer.read(struct.unpack("I", amt)[0]).decode())
+    except Exception as e:
+      lib = b""
+      print(e, file=sys.stderr, flush=True)
+    sys.stdout.buffer.write(struct.pack("I", len(lib)) + lib)
+    sys.stdout.buffer.flush()

tinygrad/runtime/support/hcq2.py

@@ -3,7 +3,7 @@ from typing import cast, TypeVar, Generic, Any, Sequence, Iterable
 import struct, functools, time, collections, itertools, decimal, statistics
 from dataclasses import replace, dataclass
 from tinygrad.helpers import suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap, PROFILE
-from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar, perf_counter_us, Context
+from tinygrad.helpers import to_tuple, round_up, partition, panic, ContextVar, perf_counter_us, Context
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer, DepsTracker
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, rewrite_group, GroupOp
@@ -48,9 +48,16 @@ def is_value_known_at_link(val:UOp) -> bool:
   return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)
 
 def make_patches(buf:UOp, patches:Sequence[tuple[sint, UOp]]) -> tuple[UOp, ...]:
-  return tuple(buf.index(UOp(Ops.STACK, dtypes.int, tuple(UOp.const(off // buf.dtype.itemsize, dtypes.int) for off,_ in ps)))
-               .store(UOp(Ops.STACK, buf.dtype, tuple(val.cast(buf.dtype) for _,val in ps))).rtag(tag)
-               for ps, tag in zip(partition(patches, lambda p: is_value_known_at_link(p[1])), ("link", None)) if ps)
+  def _mk_store(ps:list[tuple[sint, UOp]], tag:str|None) -> UOp:
+    offs = UOp(Ops.STACK, dtypes.int, tuple(UOp.const(off // buf.dtype.itemsize, dtypes.int) for off,_ in ps))
+    vals = UOp(Ops.STACK, ps[0][1].dtype, tuple(val for _,val in ps))
+    return buf.index(offs, dtype=vals.dtype).store(vals).rtag(tag)
+
+  patches = [(off, val.cast(buf.dtype) if val.dtype.itemsize == buf.dtype.itemsize else val) for off, val in patches]
+  link, runtime = partition(patches, lambda p: is_value_known_at_link(p[1]))
+  inputs, runtime = partition(runtime, lambda p: p[1].op is Ops.GETADDR)
+  return tuple(_mk_store(list(ps), tag) for cls, tag in ((link, "link"), (inputs, "inputs"), (runtime, None))
+               for _, ps in itertools.groupby(sorted(cls, key=lambda p: p[1].dtype), key=lambda p: p[1].dtype))
 
 def make_binary_patch(buf:UOp, blob:bytes) -> UOp:
   data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)
@@ -77,8 +84,8 @@ def make_call(name:str, body:UOp, info:HCQInfo) -> UOp: return UOp.custom_functi
 def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp:
   data, info = prg.arg
   buf = UOp.placeholder((data.kernargs_alloc_size // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("kernargs")
-  words = [w for gi in info.globals for w in data64_le(get_call_arg_uops(call)[gi].getaddr(devs))] + list(info.vars)
-  return buf.after(*make_patches(buf, [(i * 4, w) for i, w in enumerate(words)]))
+  words = [get_call_arg_uops(call)[gi].getaddr(devs) for gi in info.globals] + list(info.vars)
+  return buf.after(*make_patches(buf, list(zip(itertools.accumulate((w.dtype.itemsize for w in words), initial=0), words))))
 
 # *****************
 # 0.1. prep: replace buffers with params
@@ -307,27 +314,15 @@ def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[UOp, dict[UOp
   fills = (table.after(*make_patches(table, [(i*table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else ()
   return table, reads, fills, {g:slots[bare[g]] for g in gaddrs}
 
-def is_bare_addr(val:UOp) -> bool: return val.op is Ops.CAST and val.src[0].op in (Ops.AND, Ops.SHR) and val.src[0].src[0].op is Ops.GETADDR
+def make_gather_loop(patches:list[UOp], table:UOp, slots:dict[UOp, int], lt_patches:list[UOp]) -> dict[UOp, UOp]:
+  (dst,), words = dedup(p.buf_uop for p in patches), [(off.val, slots[val]) for p in patches for off, val in zip(p.src[0].src[1].src, p.src[1].src)]
 
-def make_scatter_loops(patches:list[UOp], inputs_table:tuple, lt_patches:list[UOp]) -> dict[UOp, UOp]:
-  table, _, _, slots = inputs_table
-  subs, by_dst = {}, collections.defaultdict(list)
-  for p in patches: by_dst[p.buf_uop].append(p)
-  for dst, patches in by_dst.items():
-    data = []
-    for p in patches:
-      words = [(off, val, get_getaddrs(val)) for off,val in zip(p.src[0].src[1].src, p.src[1].src)]
-      data += [(off.val, slots[gaddrs[0]]) for off,_,gaddrs in words if gaddrs][::2]
-      scalars = [(off.val*dst.dtype.itemsize, val) for off,val,gaddrs in words if not gaddrs]
-      subs[p] = UOp.group(*make_patches(dst, scalars)) if scalars else UOp(Ops.NOOP)
-
-    word_table, slot_table = (UOp.placeholder((len(data),), dtypes.uint32, next(UOp.unique_num), device=dst.device).rtag("systems") for _ in range(2))
-    ridx = UOp.range(len(data), next(UOp.unique_num), dtype=dtypes.int, src=(word_table, slot_table, dst))
-    widx, slot = ((p.index(ridx).load() % bound).cast(dtypes.int) for p,bound in ((word_table, dst.max_numel()-1), (slot_table, table.max_numel())))
-    loop = UOp.group(*[dst.index(widx+i).store((table.index(slot).load() >> 32*i).cast(dtypes.uint32)) for i in range(2)]).end(ridx)
-    lt_patches += [make_binary_patch(buf, struct.pack(f'<{len(data)}I', *vals)) for buf,vals in zip((word_table, slot_table), zip(*data))]
-    subs[patches[0]] = UOp.group(loop, subs[patches[0]])
-  return subs
+  # build a runtime loop that writes every input address
+  pairs = UOp.placeholder((2*len(words),), dtypes.uint32, next(UOp.unique_num), device=dst.device).rtag("systems")
+  lt_patches.append(make_binary_patch(pairs, struct.pack(f'<{2*len(words)}I', *itertools.chain(*words))))
+  r = UOp.range(len(words), next(UOp.unique_num), dtype=dtypes.int, src=(pairs, dst))
+  off, slot = ((pairs.index(2*r+i).load() % bound).cast(dtypes.int) for i, bound in ((0, dst.max_numel()-1), (1, table.max_numel())))
+  return {p: UOp(Ops.NOOP) for p in patches} | {patches[0]: dst.index(off, dtype=table.dtype).store(table.index(slot).load()).end(r)}
 
 def is_input_addr(g:UOp) -> bool: return all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop))
 
@@ -341,10 +336,8 @@ def split_patches(call:UOp) -> UOp|None:
   runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop)))
   tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))]
   reads, fills = {k:v for _,r,_,_ in tables for k,v in r.items()}, [f for t in tables[1:] for f in t[2]] # inputs table is filled by exec
-  input_patches = [p for p in rt_patches if (gs:=get_getaddrs(p)) and all(map(is_input_addr, gs))
-    and all(is_bare_addr(v) for v in p.src[1].src if get_getaddrs(v))]
-  scatter = make_scatter_loops(input_patches, tables[0], lt_patches)
-  body = body.substitute({p:p.substitute(scatter | reads) for p in rt_patches})
+  gathers = make_gather_loop(ipathces, tables[0][0], tables[0][3], lt_patches) if (ipathces:=[p for p in rt_patches if p.tag == "inputs"]) else {}
+  body = body.substitute({p:p.substitute(gathers | reads) for p in rt_patches})
 
   lt_srcs = collections.defaultdict(list)
   for p in lt_patches: lt_srcs[p.buf_uop].append(p)

tinygrad/uop/render.py

@@ -81,7 +81,7 @@ sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.SQRT, Ops.INDEX, Ops.REDUCE
          Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH}
 pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.val}, {x.dtype})"),
-  (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"),
+  (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})" if x.dtype != x.src[0].dtype else None),
   (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].val}, {repr(x.arg)}, dtype={x.dtype})"),
   (UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x:
     f"UOp.new_buffer({repr(x.arg.device)}, {x.max_numel()}, {x.dtype}, {x.arg.slot})"

tinygrad/uop/spec.py

@@ -223,6 +223,12 @@ spec_program = PatternMatcher([
   (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.int32),), name="s"), lambda s,x: matches_dtype(x, s.dtype) and isinstance(s.arg, str)),
 ])+spec_shared
 
+# migration: on a casted_consts renderer every literal is CAST(dt, CONST(value)) with a weak inner CONST
+spec_program_casted_consts = PatternMatcher([
+  (UPat(Ops.CONST, dtype=dtypes.weaks, name="x"), lambda x: x.dtype is dtypes.from_py(x.val)),
+  (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CAST, src=(UPat(Ops.CONST),)))), lambda: True),
+])+spec_program
+
 spec_hcq = PatternMatcher([
   (UPat(Ops.GETADDR, dtypes.uint64, src=(UPat((Ops.BUFFER, Ops.PARAM)).or_after(),), name="x"), lambda x: is_device(x.arg)),
   (UPat(Ops.PROGRAM, dtypes.void, src=(UPat((Ops.BUFFER, Ops.PARAM)).or_after(),)), lambda: True),

tinygrad/uop/symbolic.py

@@ -287,10 +287,10 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   (UPat.var('x', dtypes.ints+(dtypes.weakint,)).cast(dtypes.ints+(dtypes.weakint,), name="a").cast(name="b"),
     lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None),
   # try to do math in int instead of long, keep weak const weak
-  (UPat(GroupOp.Binary, src=(UPat.var("x", dtypes.long), UPat.var("y", dtypes.long)), name="u"), lambda u,x,y:
+  (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 not any(v.overflows(dtypes.int) for v in (u,x,y)) else None),
+    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.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+cast.const_like(c.val)),
   # only RANGE/IF/STORE/KERNEL have side effects
   (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+

tinygrad/uop/validate.py

@@ -52,6 +52,8 @@ z3_renderer = PatternMatcher([
     create_bounded(f"cast{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])),
   # A comparison between floats introduces a new bool variable
   (UPat(GroupOp.Comparison, src=UPat(dtype=dtypes.floats)), lambda ctx: (z3.Bool(f"float_cmp{len(ctx[1])}", ctx=ctx[0]), None)),
+  # a same-dtype cast states a width, which z3 does not model: identity. must precede the rules below (bool->bool)
+  (UPat(Ops.CAST, name="x"), lambda x,ctx: (ctx[1][x.src[0]], None) if x.dtype == x.src[0].dtype else None),
   # casts from bool/int to int/bool
   (UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,),src=(UPat.var("x", dtypes.bool),)), lambda x,ctx: (z3.If(ctx[1][x], 1, 0), None)),
   (UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,), src=(UPat.var("x", dtypes.ints+(dtypes.weakint,)),)), lambda x,ctx: (ctx[1][x], None)),

tinygrad/uop/weak.py

@@ -7,6 +7,35 @@ def default_dtype(u:UOp):
   if u.dtype is dtypes.weakfloat: return dtypes.default_float
   return dtypes.long if u.overflows(dtypes.int32) else dtypes.int
 
+def commit_weak(s:UOp, dt:DType) -> UOp:
+  # a CONST commits directly at dt (the value stays mathematical, emission truncates), a non-const src takes the cast
+  return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt)
+
+def commit_weak_srcs(u:UOp) -> UOp|None:
+  if not any(s.dtype in dtypes.weaks for s in u.src): return None
+  if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None
+  # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too
+  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src))
+
+# runs in index lowering and in the decomps: a rule that mints a weak const commits it in the same rewrite, so none reaches the renderer
+pm_commit_weak = PatternMatcher([
+  (UPat(GroupOp.Broadcastable, name="u"), commit_weak_srcs),
+  # demand from the destination: a STORE's weak value commits at the destination's dtype
+  (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.weaks)), allow_any_len=True, name="u"),
+   lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))),
+])
+
+# a concrete CAST over a weak node states the width the value will live at. that width is a floor, never a narrowing
+def cast_weak_srcs(c:UOp, u:UOp) -> UOp|None:
+  if c.dtype in dtypes.weaks or weak_dtype(c.dtype) is not u.dtype: return None
+  dt = least_upper_dtype(c.dtype, default_dtype(u))
+  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype)
+
+pm_cast_weak = PatternMatcher([
+  (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs),
+  (UPat(Ops.CAST, name="c", src=(UPat(Ops.CONST, dtype=dtypes.weaks, name="u"),)), lambda c,u: commit_weak(u, c.dtype)),
+])
+
 def lower_weak_node(u:UOp) -> UOp|None:
   start, src = (1 if u.op is Ops.WHERE else 0), tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src)
   if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None
@@ -39,35 +68,6 @@ def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None:
   ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype in dtypes.weaks else s for s in u.src))
   return None if ret is u else ret
 
-def commit_weak(s:UOp, dt:DType) -> UOp:
-  # a CONST commits directly at dt (the value stays mathematical, emission truncates), a non-const src takes the cast
-  return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt)
-
-def commit_weak_srcs(u:UOp) -> UOp|None:
-  if not any(s.dtype in dtypes.weaks for s in u.src): return None
-  if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None
-  # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too
-  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src))
-
-# runs in index lowering and in the decomps: a rule that mints a weak const commits it in the same rewrite, so none reaches the renderer
-pm_commit_weak = PatternMatcher([
-  (UPat(GroupOp.Broadcastable, name="u"), commit_weak_srcs),
-  # demand from the destination: a STORE's weak value commits at the destination's dtype
-  (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.weaks)), allow_any_len=True, name="u"),
-   lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))),
-])
-
-# a concrete CAST over a weak node states the width the value will live at. that width is a floor, never a narrowing
-def cast_weak_srcs(c:UOp, u:UOp) -> UOp|None:
-  if c.dtype in dtypes.weaks or weak_dtype(c.dtype) is not u.dtype: return None
-  dt = least_upper_dtype(c.dtype, default_dtype(u))
-  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype)
-
-pm_cast_weak = PatternMatcher([
-  (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs),
-  (UPat(Ops.CAST, name="c", src=(UPat(Ops.CONST, dtype=dtypes.weaks, name="u"),)), lambda c,u: commit_weak(u, c.dtype)),
-])
-
 pm_lower_index_dtype = pm_commit_weak+pm_cast_weak+PatternMatcher([
   # a CAST between two concrete dtypes over a CONST is a value conversion: evaluate it once, at the width the CAST states
   # TODO: delete this once CONST has no dtype

tinygrad/viz/js/index.js

@@ -886,7 +886,7 @@ const evtSources = [];
 // context: collection of steps
 const state = {currentCtx:-1, currentStep:0, currentRewrite:0, expandSteps:false, callSrcMask:new Set(), expandedNodes:new Set()};
 function setState(ns) {
-  saveToHistory(state);
+  if (["currentCtx", "currentStep", "currentRewrite"].some(k => k in ns && state[k] !== ns[k])) saveToHistory(state);
   const { ctx:prevCtx, step:prevStep } = select(state.currentCtx, state.currentStep);
   const prevRewrite = state.currentRewrite;
   Object.assign(state, ns);