llm kernels: adapt to Ops.BIND removal

Variables are 0-d ALU BUFFERs in the tensor graph and take the ALU PARAM form
inside kernels (UOp.variable(param=True)). Add kernel_var helper for the
conversion, and keep start_pos in bound form at the graph level so function
implicit-input collection and the schedule's binds rename-back line up.

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

@@ -42,11 +42,7 @@ inputs:
     required: false
     default: 'false'
   qemu:
-    description: "Install qemu?"
-    required: false
-    default: 'false'
-  ninja:
-    description: "Install ninja?"
+    description: "Install qemu"
     required: false
     default: 'false'
 runs:
@@ -134,7 +130,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' || inputs.ninja == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
       shell: bash
       run: |
         sudo mkdir -p /var/cache/apt/archives
@@ -162,7 +158,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' || inputs.ninja == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
       id: apt-pkgs
       shell: bash
       run: |
@@ -187,29 +183,25 @@ 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' || inputs.ninja == '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') && 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' || inputs.ninja == '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') && 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' || inputs.ninja == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
       shell: bash
       run: |
         sudo apt -qq update || true

.github/workflows/autogen.yml

@@ -35,7 +35,7 @@ jobs:
         key: 'autogen'
         amd: 'true'
         llvm: 'true'
-        deps: 'autogen'
+        pydeps: 'pyyaml mako'
     - 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/benchmark.yml

@@ -108,6 +108,10 @@ jobs:
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
       run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -125,6 +129,10 @@ jobs:
       # just metal for now
       if: ${{ matrix.dev == 'METAL' }}
       run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
+    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
+      # only run on machines with multiple gpus
+      if: ${{ matrix.dev != 'METAL' }}
+      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -174,6 +182,10 @@ jobs:
       # slow on metal
       if: ${{ matrix.dev != 'METAL' }}
       run: time BENCHMARK_LOG=cifar DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run full CIFAR training steps w 6 GPUS
+      # only run on machines with multiple gpus
+      if: ${{ matrix.dev != 'METAL' }}
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -219,6 +231,11 @@ jobs:
       run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
     - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
       run: BENCHMARK_LOG=resnet_10steps DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf Bert training steps (6 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -268,58 +285,6 @@ jobs:
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
-  multigpubenchmark:
-    name: Multi-GPU Benchmarks (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
-    strategy:
-      fail-fast: false
-      matrix:
-        dev: ['AMD', 'NV']
-    timeout-minutes: 60
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    env:
-      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup (AMD)
-      if: ${{ matrix.dev == 'AMD' }}
-      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
-    - name: Setup (NV)
-      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        mkdir -p extra/datasets
-        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: python3 test/external/process_replay/reset.py
-    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf Bert training steps (6 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
   tests:
     name: Tests (DEV=${{ matrix.dev }})
     runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
@@ -576,7 +541,7 @@ jobs:
     - name: openpilot run_pickle big_driving_supercombo
       run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
     - name: Test copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
 
   driverbenchmarks:
     name: PCI Driver Benchmark (DEV=${{ matrix.dev }})

.github/workflows/docs.yml

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

.github/workflows/python-publish.yml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   deploy:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
     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-24.04
+    runs-on: ubuntu-latest
     outputs:
       branchstat: ${{ steps.brstat.outputs.stat}}
     steps:
@@ -44,7 +44,7 @@ jobs:
     permissions:
       contents: read
       pull-requests: write
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
     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-24.04
+    runs-on: ubuntu-latest
     needs: checkbranch
     if: needs.checkbranch.outputs.branchstat == 'true'
     steps:

.github/workflows/test.yml

@@ -31,7 +31,8 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
-        deps: "docs testing_minimal"
+        deps: docs
+        pydeps: "capstone torch"
     - name: Build wheel and show size
       run: |
         uv build --wheel
@@ -72,7 +73,10 @@ jobs:
         deps: testing_unit
         pydeps: "pillow torchvision expecttest"
         llvm: 'true'
-        ninja: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
     - name: Test ResNet-18
       run: DEBUG=2 python3 extra/torch_backend/example.py
     - name: Test one op in torch tests
@@ -82,23 +86,6 @@ jobs:
     - name: Custom tests
       run: DEV=CPU:LLVM GPUS=4 TINY_BACKEND=1 python3 -m pytest -nauto extra/torch_backend/test.py extra/torch_backend/test_inplace.py extra/torch_backend/test_multigpu.py extra/torch_backend/test_kernel_fusion.py --durations=20
 
-  torchbackendtrain:
-    name: Torch Backend Training
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_unit
-        llvm: 'true'
-        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
-
   bepython:
     name: Python Backend
     runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}

README.md

@@ -140,7 +140,7 @@ Documentation along with a quick start guide can be found on the [docs website](
 ```python
 from tinygrad import Tensor
 
-x = Tensor.eye(3).clone()  # clone to make it a buffer
+x = Tensor.eye(3)
 y = Tensor([[2.0,0,-2.0]])
 z = y.matmul(x).sum()
 z.backward()

examples/mlperf/model_train.py

@@ -1742,8 +1742,8 @@ def train_gptoss():
   )
 
   for p in optim.params:
-    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)
+    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
+    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
   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,7 +146,6 @@ 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)
 
@@ -183,12 +182,10 @@ class GPTOSS:
     xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
     xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)  # (B,N,H,D)/(B,N,KV,D)
 
-    fa_saves = []
     if getenv("HK_FLASH_ATTENTION"):
       from extra.thunder.amd.fa import flash_attention
-      attn, _, l_vec = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
+      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
       attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
-      fa_saves = [xq, xk, xv, l_vec]
     elif sliding:
       attn = self._sliding_attention(xq, xk, xv, sinks)
     else:
@@ -202,7 +199,7 @@ class GPTOSS:
       attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
 
     out = matmul_mx(attn, wo, wo_scale) + wo_bias
-    return out, [x_normed, rrms, attn] + fa_saves
+    return out, [x_normed, rrms, attn]
 
   def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
                    w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
@@ -223,7 +220,6 @@ class GPTOSS:
       z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
           + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
       out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
-      return out, [x_normed, rrms, xg, h, y, z, r.weights, r.dest_row, r.off]
     else:
       thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
       weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)

extra/gemm/moe_gemm.py

@@ -1,32 +1,10 @@
 import functools, pathlib
 from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 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
@@ -80,8 +58,7 @@ 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]
-  if is_multi and ZERO_OPTIM: out = reduce_scatter_devaxis(out, 0)
-  else: out = out.sum(0) if is_multi else out.squeeze(0)
+  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,10 +179,11 @@ 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
-  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))))))
+  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)))))
 
 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) \
@@ -560,7 +561,9 @@ class AMDDevice(HCQ2Compiled):
   def is_usb(self) -> bool: return False
 
   def __init__(self, device:str=""):
-    self.iface = self._select_iface(device)
+    self.device_id = int(device.split(":")[1]) if ":" in device else 0
+
+    self.iface = self._select_iface()
 
     self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
     self.arch = "gfx%d%x%x" % self.target

extra/thunder/amd/fa.py

@@ -19,33 +19,16 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
 @functools.cache
 def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp,
                                   device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
-  group_size = H // H_KV
-  q, k, v = q.reshape(B, N, H, D), k.reshape(B, N, H_KV, D), v.reshape(B, N, H_KV, D)
-  xqkv = xqkv.reshape(B, N, H_KV, group_size + 2, D)
-  b, n = UOp.range(B, 0), UOp.range(N, 1)
-  pair = UOp.range(D // 2, 2)
-  even = pair * 2
-  c = freqs_cis[0, n, 0, pair, 0].cast(dtypes.float)
-  s = freqs_cis[0, n, 0, pair, 1].cast(dtypes.float)
-  ordered:UOp|None = None
-  for kvh in range(H_KV):
-    q_out, k_out, v_out = (x.after(ordered) if ordered is not None else x for x in (q, k, v))
-    x_in = xqkv.after(ordered) if ordered is not None else xqkv
-    stores:list[UOp] = []
-    for rep in range(group_size):
-      a = x_in[b, n, kvh, rep, even].cast(dtypes.float)
-      bb = x_in[b, n, kvh, rep, even + 1].cast(dtypes.float)
-      h = kvh * group_size + rep
-      stores += [q_out[b, n, h, even].store((a * c - bb * s).cast(q.dtype)), q_out[b, n, h, even + 1].store((a * s + bb * c).cast(q.dtype))]
-    a = x_in[b, n, kvh, group_size, even].cast(dtypes.float)
-    bb = x_in[b, n, kvh, group_size, even + 1].cast(dtypes.float)
-    stores += [k_out[b, n, kvh, even].store((a * c - bb * s).cast(k.dtype)),
-               k_out[b, n, kvh, even + 1].store((a * s + bb * c).cast(k.dtype)),
-               v_out[b, n, kvh, even].store(x_in[b, n, kvh, group_size + 1, even]),
-               v_out[b, n, kvh, even + 1].store(x_in[b, n, kvh, group_size + 1, even + 1])]
-    ordered = UOp.group(*stores)
-  assert ordered is not None
-  return ordered.end(pair, n, b).sink(arg=KernelInfo(name="fused_qkv_rope_forward"))
+  code = (pathlib.Path(__file__).parent / "fused_qkv_rope.cpp").read_text()
+  threads = 256
+  thread_idx = UOp.special(threads, "lidx0")
+  block_idx_x, block_idx_y = UOp.special(B, "gidx0"), UOp.special(N, "gidx1")
+  sink = UOp.sink(q.base, k.base, v.base, xqkv.base, freqs_cis.base, thread_idx, block_idx_x, block_idx_y,
+                  arg=KernelInfo(name="fused_qkv_rope_forward"))
+  compile_args = ["-std=c++20", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}",
+                  f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DTHREADS_PER_BLOCK={threads}"]
+  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
 
 @functools.cache
 def custom_fused_qkv_rope_backward(dxqkv:UOp, dq:UOp, dk:UOp, dv:UOp, freqs_cis:UOp,

extra/thunder/amd/fused_qkv_rope.cpp

@@ -0,0 +1,69 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+#ifndef ATTN_B
+#define ATTN_B 2
+#endif
+#ifndef ATTN_N
+#define ATTN_N 8192
+#endif
+#ifndef ATTN_H
+#define ATTN_H 32
+#endif
+#ifndef ATTN_H_KV
+#define ATTN_H_KV 8
+#endif
+#ifndef ATTN_D
+#define ATTN_D 128
+#endif
+#ifndef THREADS_PER_BLOCK
+#define THREADS_PER_BLOCK 256
+#endif
+
+constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
+constexpr int HALF_D = ATTN_D / 2;
+constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D;
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
+fused_qkv_rope_forward(
+    __hip_bfloat16*       __restrict__ q,
+    __hip_bfloat16*       __restrict__ k,
+    __hip_bfloat16*       __restrict__ v,
+    const __hip_bfloat16* __restrict__ xqkv,
+    const __hip_bfloat16* __restrict__ freqs_cis) {
+  const int b = blockIdx.x;
+  const int n = blockIdx.y;
+  const int bn = b * ATTN_N + n;
+  const int packed_bn = bn * ATTN_H_KV * PACKED_D;
+  const int q_bn = bn * ATTN_H * ATTN_D;
+  const int kv_bn = bn * ATTN_H_KV * ATTN_D;
+
+  if (threadIdx.x < HALF_D) {
+    const int pair = threadIdx.x;
+    const int even = pair << 1;
+    const float c = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 0]);
+    const float s = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 1]);
+
+    for (int kvh = 0; kvh < ATTN_H_KV; kvh++) {
+      const int base = packed_bn + kvh * PACKED_D;
+
+      for (int rep = 0; rep < GROUP_SIZE; rep++) {
+        const int qbase = base + rep * ATTN_D;
+        const int h = kvh * GROUP_SIZE + rep;
+        const float a = static_cast<float>(xqkv[qbase + even]);
+        const float bb = static_cast<float>(xqkv[qbase + even + 1]);
+        const int out = q_bn + h * ATTN_D + even;
+        q[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
+        q[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
+      }
+
+      const float a = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even]);
+      const float bb = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even + 1]);
+      const int out = kv_bn + kvh * ATTN_D + even;
+      k[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
+      k[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
+      v[out] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even];
+      v[out + 1] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even + 1];
+    }
+  }
+}

extra/thunder/tiny/tk/tiles.py

@@ -209,7 +209,7 @@ class ST:
     return cls(uop, rows, cols, layout, base_shape, ker)
 
   def swizzle(self, row, col):
-    swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype)
+    swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype.scalar())
 
     row = swizzled_offset // self.base_shape.cols
     col = swizzled_offset % self.base_shape.cols

extra/torch_backend/backend.py

@@ -4,7 +4,7 @@
 # A006 Lambda argument `input` is shadowing a Python builtin
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.uop.ops import Ops, GroupOp
-from tinygrad.helpers import getenv, prod, strides_for_shape
+from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
 import torch.lib
 TORCH_DEBUG = getenv("TORCH_DEBUG")
 import torch, pathlib, operator, functools, weakref
@@ -73,12 +73,6 @@ def wrap_view_op(fn):
     return wrap(ret)
   return _wrap
 
-# NOTE: list assignment raises IndexError on an out of range dim, and the index must be a tuple: a list of all ints is one advanced index
-def _index_dim(self, dim, idx):
-  idxs = [slice(None)] * self.ndim
-  idxs[dim] = idx
-  return self[tuple(idxs)]
-
 view_ops = {
   "aten.view": Tensor.reshape,
   "aten._unsafe_view": Tensor.reshape,  # when are views unsafe, and do we care?
@@ -88,13 +82,15 @@ view_ops = {
   "aten.transpose.int": Tensor.transpose,
   "aten.squeeze.dim": Tensor.squeeze,
   "aten.unsqueeze": Tensor.unsqueeze,
-  "aten.select.int": _index_dim,
+  "aten.select.int": lambda self, dim, idx: self[(slice(None),) * (dim%self.ndim) + (idx,)],
   "aten.permute": Tensor.permute,
   "aten.alias": lambda self: self,
   "aten.diagonal": Tensor.diagonal,
-  "aten.slice.Tensor": lambda self, dim=0, start=None, end=None, step=1: _index_dim(self, dim, slice(start, end, step)),
   }
 
+# torch 2.10 handles this natively
+if tuple(map(int, torch.__version__.split('.')[:2])) < (2, 10): view_ops.update({"aten.detach": Tensor.detach})
+
 for k,v in view_ops.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_view_op(v))
 
 def _get_view_ops(view): return getattr(view, "_view_ops", [])
@@ -103,21 +99,46 @@ def _apply_view_ops(target, ops):
   for fn, args, kwargs in ops: target = fn(target, *args, **kwargs)
   return target
 
-# a chain of reshapes is undone by reshaping the value back to the base
+# similar to https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/InferSize.h
+def _reshape_target_shape(shape:tuple[int, ...], args) -> tuple[int, ...]|None:
+  if not (req := argfix(*args)): return None
+  new_shape, infer_idx = [], -1
+  for i, s in enumerate(req):
+    if s is None: s = shape[i] if i < len(shape) else None
+    if not isinstance(s, int): return None
+    if s == -1:
+      if infer_idx != -1: return None
+      infer_idx = len(new_shape)
+    new_shape.append(s)
+  total = prod(shape)
+  if infer_idx != -1:
+    known = prod(x for x in new_shape if x != -1)
+    if known == 0:
+      if total != 0: return None
+      new_shape[infer_idx] = 0
+    else: new_shape[infer_idx] = total // known
+  return tuple(new_shape) if prod(new_shape) == total else None
+
+# TODO: can we get rid of this? only for test_flatten_reshape_add
 def _try_simple_reshape_view_write(base: Tensor, view: Tensor, val: Tensor) -> bool:
   if not (ops := _get_view_ops(view)): return False
-  if any(fn is not Tensor.reshape for fn, _, _ in ops): return False
-  base.assign(val.reshape(base.shape))
+  shapes = [base.shape]
+  for fn, args, _ in ops:
+    if fn is Tensor.reshape:
+      if not (next_shape := _reshape_target_shape(shapes[-1], args)): return False
+      shapes.append(next_shape)
+  if shapes[-1] != view.shape: return False
+  for s in reversed(shapes[:-1]): val = val.reshape(s)
+  base.assign(val)
   return True
 
 def _view_write(base: Tensor, view: Tensor, value: Tensor) -> None:
   val = value if value.dtype == base.dtype else value.cast(base.dtype)
+  if view.shape == base.shape: return base.assign(val)
   if _try_simple_reshape_view_write(base, view, val): return
   idx_base = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape)
   idx_view = _apply_view_ops(idx_base, _get_view_ops(view)).reshape(-1)
-  # clone, not contiguous: contiguous() on a base that already owns its buffer returns the base itself, and scattering
-  # into that is an in-place write to a buffer other tensors still hold, which setitem refuses
-  flat_base = base.reshape(base.numel()).clone()
+  flat_base = base.reshape(base.numel()).contiguous()
   flat_base[idx_view] = val.reshape(-1)
   base.assign(flat_base.reshape(base.shape))
 
@@ -145,6 +166,11 @@ def _index_put_impl_(self, indices, values, accumulate=False, unsafe=False):
 def index_put(self, indices, values, accumulate=False):
   return aten.index_put(self.cpu(), [z.cpu() if isinstance(z, torch.Tensor) else None for z in indices], values.clone().cpu(), accumulate).tiny()
 
+@torch.library.impl("aten::isin.Tensor_Tensor_out", "privateuseone")
+def isin_tensor_tensor_out(x, y, *, assume_unique=False, invert=False, out=None):
+  result = (unwrap(x).unsqueeze(-1) == unwrap(y).flatten()).any(-1)
+  return out.copy_(wrap(~result if invert else result))
+
 @torch.library.impl("aten::randperm.generator_out", "privateuseone")
 def randperm_generator(n, generator=None, out=None):
   if generator is not None: raise NotImplementedError("tinygrad torch backend does not support torch.Generator for randperm")
@@ -205,6 +231,49 @@ def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
 def _reshape_alias(tensor:torch.Tensor, size, stride):
   return _as_strided(tensor, size, stride)
 
+@torch.library.impl("aten::empty_strided", "privateuseone")
+def empty_strided(size, stride, dtype=None, layout=None, device=None, pin_memory=False):
+  if TORCH_DEBUG: print(f"empty_strided {size=} {stride=} {dtype=} {layout=} {device=} {pin_memory=}")
+  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
+  # TODO: should return with requested strides
+  return wrap(ret)
+
+@torch.library.impl("aten::empty.memory_format", "privateuseone")
+def empty_memory_format(size, dtype=None, layout=None, device=None, pin_memory=False, memory_format=None):
+  if TORCH_DEBUG: print(f"empty.memory_format {size=} {dtype=} {layout=} {device=} {pin_memory=} {memory_format=}")
+  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
+  return wrap(ret)
+
+@torch.library.impl("aten::max_pool2d_with_indices", "privateuseone")
+def max_pool2d_with_indices(self:torch.Tensor, kernel_size:tuple[int, ...], stride=None, padding=0, dilation=1, ceil_mode=False):
+  # TODO: supprt stride [] in tinygrad?
+  if stride is not None and len(stride) == 0: stride = None
+  ret, idx = unwrap(self).max_pool2d(kernel_size, stride, dilation, padding, ceil_mode, return_indices=True)
+  return (wrap(ret), wrap(idx.cast(dtypes.int64)))
+
+@torch.library.impl("aten::max_pool2d_with_indices_backward", "privateuseone")
+def max_pool2d_with_indices_backward(grad_out:torch.Tensor, self:torch.Tensor, kernel_size:tuple[int, ...], stride=None, padding=0, dilation=1, ceil_mode=False, indices=None):
+  return wrap(Tensor.max_unpool2d(unwrap(grad_out), unwrap(indices), output_size=unwrap(self).shape))
+
+@torch.library.impl("aten::max_unpool2d", "privateuseone")
+def max_unpool2d(self:torch.Tensor, indices:torch.Tensor, output_size):
+  return wrap(unwrap(self).max_unpool2d(unwrap(indices), output_size=output_size))
+
+@torch.library.impl("aten::arange", "privateuseone")
+def arange(end, dtype=None, device=None, pin_memory=None):
+  has_float = isinstance(end, float)
+  return wrap(Tensor.arange(0, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
+
+@torch.library.impl("aten::arange.start", "privateuseone")
+def arange_start(start, end, dtype=None, device=None, pin_memory=None):
+  has_float = any(isinstance(x, float) for x in (start, end))
+  return wrap(Tensor.arange(start, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
+
+@torch.library.impl("aten::arange.start_step", "privateuseone")
+def arange_start_step(start, end, step, dtype=None, device=None, pin_memory=None):
+  has_float = any(isinstance(x, float) for x in (start, end, step))
+  return wrap(Tensor.arange(start, end, step, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
+
 @torch.library.impl("aten::convolution_overrideable", "privateuseone")
 def convolution_overrideable(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups):
   if TORCH_DEBUG >= 1:
@@ -225,27 +294,12 @@ def convolution_backward_overrideable(grad_out, input, weight, stride, padding,
   grads = out.gradient(*[t for t,m in zip([input, weight, bias], output_mask) if m], gradient=grad_out)
   return tuple([wrap(grads.pop(0)) if m else None for m in output_mask])
 
-# the functional scatters. without an impl aten falls back to a path that assumes a real storage: "self.has_storage() INTERNAL ASSERT FAILED"
-def _scatter_into(self, src, dim, index):
-  out = unwrap(self).clone()
-  slices = [slice(None)] * out.ndim
-  slices[dim] = index
-  out[slices] = unwrap(src).cast(out.dtype)  # torch casts src to self's dtype, tinygrad setitem demands they already match
-  return wrap(out)
-
-@torch.library.impl("aten::slice_scatter", "privateuseone")
-def slice_scatter(self, src, dim=0, start=None, end=None, step=1): return _scatter_into(self, src, dim, slice(start, end, step))
-
-@torch.library.impl("aten::select_scatter", "privateuseone")
-def select_scatter(self, src, dim, index): return _scatter_into(self, src, dim, index)
-
-@torch.library.impl("aten::diagonal_scatter", "privateuseone")
-def diagonal_scatter(self, src, offset=0, dim1=0, dim2=1):
-  # a diagonal is not one axis, so scatter through the flat indices it picks out
-  base, out = unwrap(self), unwrap(self).clone().reshape(-1)
-  idx = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape).diagonal(offset, dim1, dim2).reshape(-1)
-  out[idx] = unwrap(src).cast(base.dtype).reshape(-1)
-  return wrap(out.reshape(base.shape))
+@torch.library.impl("aten::slice.Tensor", "privateuseone")
+@wrap_view_op
+def slice_tensor(self, dim=0, start=None, end=None, step=1):
+  slices = [slice(None)] * self.ndim
+  slices[dim] = slice(start, end, step)
+  return self[slices]
 
 @torch.library.impl("aten::slice_backward", "privateuseone")
 def slice_backward(grad_out, input_sizes, dim, start, end, step):
@@ -287,14 +341,19 @@ for dim in [1, 2, 3]:
     torch.library.impl(f"aten::{pad_type}_pad{dim}d", "privateuseone")(functools.partial(pad_forward, mode=mode))
     torch.library.impl(f"aten::{pad_type}_pad{dim}d_backward", "privateuseone")(functools.partial(pad_backward, mode=mode))
 
-# the schemas are all positional: (self, output_size, align_corners, *scales) for linear, (self, output_size, *scales) for nearest.
-def upsample(self, size, *args, mode=None):
-  return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=args[0] if mode == "linear" else False))
+def upsample(self, size, align_corners=False, mode=None): return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=align_corners))
 for i,pre in enumerate(["", "bi", "tri"]):
   torch.library.impl(f"aten::upsample_{pre}linear{i+1}d", "privateuseone")(functools.partial(upsample, mode="linear"))
   torch.library.impl(f"aten::upsample_nearest{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest"))
   torch.library.impl(f"aten::_upsample_nearest_exact{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest-exact"))
 
+@torch.library.impl("aten::scatter_add.out", "privateuseone")
+def scatter_add(self, dim, index, src, out):
+  self, index, src, out_unwrapped = unwrap(self), unwrap(index), unwrap(src), unwrap(out)
+  if self.shape == (): _apply_inplace(out_unwrapped, src)
+  else: _apply_inplace(out_unwrapped, Tensor.scatter_reduce(self, dim, index, src, reduce='sum'))
+  return out
+
 def _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking=False):
   if src.is_tiny and dest.is_tiny:
     src_t, dest_t = unwrap(src), unwrap(dest)
@@ -345,11 +404,15 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic
   _apply_inplace(unwrap(indices), out_indices.cast(dtypes.int64))
   return values, indices
 
+@torch.library.impl("aten::_linalg_svd", "privateuseone")
+def _linalg_svd(self, full_matrices=False):
+  U, S, Vh = unwrap(self).svd(full_matrices)
+  return wrap(U), wrap(S), wrap(Vh)
+
 # register some decompositions
 from torch._decomp import get_decompositions
 decomps = [
   aten.native_layer_norm_backward,
-  aten.native_group_norm_backward,
   aten.linalg_cross,
   aten.addmm,
   aten.addcmul,
@@ -384,20 +447,12 @@ decomps = [
   aten._softmax_backward_data, aten.embedding_dense_backward,
   aten.linalg_vector_norm,
   aten.binary_cross_entropy, aten.binary_cross_entropy_backward,
-  # the C++ mse/smooth_l1 kernels resize their out tensor, and a tiny tensor has no storage to resize
-  aten.mse_loss, aten.mse_loss_backward,
-  aten.smooth_l1_loss, aten.smooth_l1_loss_backward,
   aten.upsample_nearest2d.out,
-  # NOTE: only the "out" overload, the "vec" one is CompositeImplicitAutograd and overriding it loses the autograd kernel
-  aten.upsample_bicubic2d.out,
-  aten._adaptive_avg_pool2d,
   # activations
   aten.hardswish, aten.hardswish_backward,
   aten.hardtanh, aten.hardtanh_backward,
   aten.gelu, aten.gelu_backward,
-  # NOTE: no aten.logical_or here, its decomposition reaches aten.bitwise_or through a path that checks aliasing by
-  # reading storage, which a tiny tensor has none of. it gets a direct impl below instead
-  aten.logical_and, aten.logical_xor,
+  aten.logical_and,
   aten.randint,
   aten.eye,
   aten.hardsigmoid_backward,
@@ -440,7 +495,7 @@ simple_tensor_methods = [
   # reduce
   "all", "any", "argmax", "argmin", "cumsum", "cumprod",
   # complex
-  "linspace"]
+  "avg_pool2d", "linspace"]
 
 tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_methods}, **{
   "aten.add.out": lambda input,other,alpha=1: input+alpha*other,
@@ -485,8 +540,6 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
   "aten.where.self_out": Tensor.where,
   "aten.prod.int_out": Tensor.prod,
   "aten.scatter.src_out": Tensor.scatter,
-  "aten.scatter_add.out": lambda self,dim,index,src: src if self.shape == () else Tensor.scatter_reduce(self, dim, index, src, reduce="sum"),
-  "aten.isin.Tensor_Tensor_out": lambda x,y,assume_unique=False,invert=False: (x.unsqueeze(-1)==y.flatten()).any(-1) != invert,
   # NOTE: axis=[] in torch means all, change tinygrad?
   "aten.sum.IntList_out": lambda self,axis,keepdim=False,dtype=None:
     self.sum(axis if axis is None or len(axis) else None, keepdim,
@@ -502,9 +555,10 @@ def wrap_out(f):
     assert out.shape == assigned.shape, f"shape mismatch: {assigned.shape} -> {out.shape}"
     assert out.device == assigned.device or out.device is None or assigned.device is None, f"device mismatch: {assigned.device} -> {out.device}"
     assert out.dtype == assigned.dtype, f"dtype mismatch: {assigned.dtype} -> {out.dtype}"
-    # writing out= is an in-place write like any other: through the base if it is a view, refreshing any derived views
-    _apply_inplace(out, assigned)
-    return out
+    # an out= that is a view has to be written through its base, and _apply_inplace gives a deviceless base its buffer first
+    if canonical_base(out) is not out: return _apply_inplace(out, assigned) or out
+    if out.device is None and assigned.device is not None: out.replace(out.empty_like(device=assigned.device))
+    return out.assign(assigned)
   return _wrap_out
 
 def _inplace_op(t, new_value):
@@ -512,14 +566,7 @@ def _inplace_op(t, new_value):
   else: _apply_inplace(t, new_value)
   return t
 
-# the three arange overloads are one function at different arity, and dtype/layout/device/pin_memory are keyword only in all of them
-def _arange(*args, dtype=None, **_):
-  return Tensor.arange(*args, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if any(isinstance(x, float) for x in args) else torch.int64)))
-
-def _empty(size, dtype=None, device=None, **_):
-  return Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
-
-tiny_backend = {**tiny_backend_out, **{
+tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.remainder.Scalar_Tensor": lambda x,y: x%y,
   "aten.floor_divide": lambda x,y: x//y,
   "aten.floor_divide_.Tensor": lambda x,y: x//y,
@@ -532,8 +579,8 @@ tiny_backend = {**tiny_backend_out, **{
   # inplace ops using replace for fusion
   "aten.zero_": lambda x: x.const_like(0),
   "aten.fill_.Scalar": lambda x, y: x.const_like(y),
-  "aten.add_.Tensor": lambda self, other, alpha=1: self + other * alpha,
-  "aten.add_.Scalar": lambda self, other, alpha=1: self + other * alpha,
+  "aten.add_.Tensor": lambda self, other, alpha=1.0: self + other * alpha,
+  "aten.add_.Scalar": lambda self, other, alpha=1.0: self + other * alpha,
   "aten.mul_.Tensor": lambda self, other: self * other,
   "aten.mul_.Scalar": lambda self, other: self * other,
   # relu doesn't have an out form?
@@ -566,9 +613,7 @@ tiny_backend = {**tiny_backend_out, **{
   # these don't work in out form, they have size 0
   "aten.abs": Tensor.abs,
   "aten.logical_not": Tensor.logical_not,
-  # compare against zero first: logical_* is bool-valued for any input dtype, while | is bitwise
-  "aten.logical_or": lambda x, y: (x != 0) | (y != 0),
-  "aten.logical_or_": lambda x, y: (x != 0) | (y != 0),
+  "aten.logical_or_": lambda x, y: x | y,
   "aten.multinomial": Tensor.multinomial,
   "aten.masked_fill_.Scalar": lambda self, mask, value: self.masked_fill(mask, value),
   "aten.masked_fill_.Tensor": lambda self, mask, value: self.masked_fill(mask, value),
@@ -577,7 +622,14 @@ tiny_backend = {**tiny_backend_out, **{
   "aten.masked_select": Tensor.masked_select,
   "aten.all": Tensor.all,
   "aten.sgn": Tensor.sign,
+  "aten.acos": Tensor.acos,
   "aten.any": Tensor.any,
+  "aten.bitwise_not": Tensor.bitwise_not,
+  "aten.argmax": Tensor.argmax,
+  "aten.argmin": Tensor.argmin,
+  "aten.asinh": Tensor.asinh,
+  "aten.mul": Tensor.mul,
+  "aten.atanh": Tensor.atanh,
   "aten.fill_.Tensor": lambda self, value: self.const_like(value.reshape(()).item()),
   "aten.flip": Tensor.flip,
   "aten.scatter_reduce.two": Tensor.scatter_reduce,
@@ -588,22 +640,10 @@ tiny_backend = {**tiny_backend_out, **{
   "aten.add.Tensor": lambda input,other,alpha=1: input+alpha*other,
   "aten.linspace": lambda start, stop, steps, dtype=None, **kwargs:
     Tensor.linspace(start, stop, steps, **({"dtype": _from_torch_dtype(dtype)} if dtype is not None else {})),
-  # the functional copy_. without an impl the fallback segfaults on a tensor with no storage
-  "aten.copy": lambda self,src,non_blocking=False: src.cast(self.dtype).to(self.device).expand(self.shape),
-  "aten.arange": lambda end, **kwargs: _arange(0, end, **kwargs),
-  "aten.arange.start": _arange,
-  "aten.arange.start_step": _arange,
-  # empty_strided takes the strides and drops them: we always allocate contiguous
-  "aten.empty_strided": lambda size, stride, **kwargs: _empty(size, **kwargs),
-  "aten.empty.memory_format": _empty,
-  # TODO: supprt stride [] in tinygrad?
-  "aten.max_pool2d_with_indices": lambda self,kernel_size,stride=None,padding=0,dilation=1,ceil_mode=False: ((r:=Tensor.max_pool2d(self, kernel_size, stride or None, dilation, padding, ceil_mode, return_indices=True))[0], r[1].cast(dtypes.int64)),
-  "aten.max_pool2d_with_indices_backward": lambda grad_out,self,kernel_size,stride=None,padding=0,dilation=1,ceil_mode=False,indices=None: Tensor.max_unpool2d(grad_out, indices, output_size=self.shape),
-  "aten.max_unpool2d": lambda self,indices,output_size: Tensor.max_unpool2d(self, indices, output_size=output_size),
-  "aten._linalg_svd": lambda self,full_matrices=False: Tensor.svd(self, full_matrices),
   "aten.topk": Tensor.topk,
   "aten.constant_pad_nd": lambda self, padding, value=0.0: self.pad(padding, mode="constant", value=value).contiguous(),
-  "aten.cumsum": lambda self, dim: self.cumsum(dim),
+  # TODO: input contiguous is needed to prevent CFGContext circular dependency assertion for shapes >512 (see test_cumsum_arange_large)
+  "aten.cumsum": lambda self, dim: self.contiguous().cumsum(dim),
   "aten.logsumexp": lambda self, axis, keepdim=False: self.logsumexp(axis[0], keepdim=keepdim),
   "aten.roll": Tensor.roll,
   "aten.logcumsumexp": Tensor.logcumsumexp,
@@ -612,7 +652,6 @@ tiny_backend = {**tiny_backend_out, **{
     self.ones_like(**{k: v for k, v in {"dtype": _from_torch_dtype(dtype) if dtype else None,
                                         "device": _from_torch_device(device) if device else None}.items() if v is not None}),
   "aten.max.dim": lambda self, dim, keepdim=False: (self.max(dim, keepdim), self.argmax(dim, keepdim).cast(dtype=dtypes.int64)),
-  "aten.min.dim": lambda self, dim, keepdim=False: (self.min(dim, keepdim), self.argmin(dim, keepdim).cast(dtype=dtypes.int64)),
   "aten.cummax": lambda self, dim: ((r := self.cummax(dim))[0], r[1].cast(dtypes.int64)),
   "aten.cummin": lambda self, dim: ((r := self.cummin(dim))[0], r[1].cast(dtypes.int64)),
   "aten.nonzero": Tensor.nonzero,
@@ -674,16 +713,15 @@ def wrap_inplace_view_op(f):
   return nf
 
 # the aten schema says how an op is called: an inplace view retargets the view, a writable first arg is inplace,
-# and a writable out arg gets wrap_out's dtype cast, shape assert, and view write-through
+# and a writable out arg must have come from tiny_backend_out so that wrap_out was applied
 for k,v in tiny_backend.items():
   name, _, overload = k.removeprefix("aten.").partition(".")
   op = getattr(getattr(aten, name), overload or "default")
   writes = [a.name for a in op._schema.arguments if a.alias_info is not None and a.alias_info.is_write]
   if torch.Tag.inplace_view in op.tags: fxn = wrap_inplace_view_op(v)
   elif writes == [op._schema.arguments[0].name] and op._schema.returns: fxn = wrap_inplace(v)
-  elif not writes: fxn = wrap_fxn(k, v)
-  elif writes == ["out"]: fxn = wrap_fxn(k, wrap_out(v))
-  else: raise RuntimeError(f"{k} writes {writes}: unhandled writable arg in schema")
+  elif not writes or (writes == ["out"] and k in tiny_backend_out): fxn = wrap_fxn(k, v)
+  else: raise RuntimeError(f"{k} writes {writes}: expected an inplace first arg, or an out arg with {k} in tiny_backend_out")
   torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(fxn)
 
 @torch.library.impl("aten::equal", "privateuseone")

extra/torch_backend/test.py

@@ -83,12 +83,6 @@ class TestTorchBackend(unittest.TestCase):
     torch.add(torch.ones(5, device=device), torch.ones(5, device=device), out=a)
     self.assertEqual(a.detach().storage_offset(), 3)
 
-  def test_out_refreshes_views_of_base(self):
-    a = torch.zeros(4, device=device)
-    v = a[2:]
-    torch.add(torch.ones(4, device=device), torch.ones(4, device=device), out=a)
-    np.testing.assert_equal(v.cpu().numpy(), [2., 2.])
-
   @unittest.expectedFailure  # TODO: storage offset assumes a contiguous source, use UOp.contiguous_view_offset
   def test_storage_offset_non_contiguous_source(self):
     a = torch.arange(12., device=device).reshape(3,4)
@@ -172,15 +166,6 @@ class TestTorchBackend(unittest.TestCase):
     expected = np.array([[1.5, 5.2, 9.0], [13.2, 17.1, 18.4]], dtype=np.float32)
     np.testing.assert_equal(y3.cpu().numpy(), expected)
 
-  def test_argmax_argmin(self):
-    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
-    c = a.cpu()
-    for got, want in [(a.argmax(), c.argmax()), (a.argmin(0), c.argmin(0)), (a.argmax(1, keepdim=True), c.argmax(1, keepdim=True)),
-                      (torch.min(a, 1).indices, torch.min(c, 1).indices), (torch.max(a, 1).indices, torch.max(c, 1).indices),
-                      (torch.min(a, 1).values, torch.min(c, 1).values), (torch.min(a, 1, keepdim=True).indices, torch.min(c, 1, keepdim=True).indices)]:
-      self.assertEqual(got.dtype, want.dtype)  # torch's arg reduces are int64, tinygrad's are int32
-      np.testing.assert_equal(got.cpu().numpy(), want.numpy())
-
   def test_isfinite(self):
     a = torch.ones(4, device=device)
     np.testing.assert_equal(torch.isfinite(a).cpu().numpy(), [True, True, True, True])
@@ -388,22 +373,6 @@ class TestTorchBackend(unittest.TestCase):
     for bwd_eps in [1e-5, 0.3]:
       for got, want in zip(run(device, bwd_eps), run("cpu", bwd_eps)): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
 
-  def test_groupnorm_backward(self):
-    def run(dev):
-      x = torch.arange(24., device=dev).reshape(2, 4, 3).requires_grad_()
-      w = torch.linspace(0.5, 2.0, 4).to(dev).requires_grad_()
-      torch.nn.functional.group_norm(x, 2, w, torch.zeros(4, device=dev)).square().sum().backward()
-      return x.grad.cpu().numpy(), w.grad.cpu().numpy()
-    for got, want in zip(run(device), run("cpu")): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
-
-  def test_mse_smooth_l1_loss_backward(self):
-    def run(dev, loss):
-      x = torch.arange(4., device=dev).requires_grad_()
-      loss(x, torch.ones(4, device=dev)).backward()
-      return x.grad.cpu().numpy()
-    for loss in [torch.nn.functional.mse_loss, torch.nn.functional.smooth_l1_loss]:
-      np.testing.assert_allclose(run(device, loss), run("cpu", loss), atol=1e-6)
-
   def test_batchnorm_unsqueeze(self):
     bn = torch.nn.BatchNorm2d(4).to(device)
     x = torch.randn(8, 4, 3, 3, device=device)
@@ -547,15 +516,6 @@ class TestTorchBackend(unittest.TestCase):
     cpu_res = torch.arange(20, dtype=torch.float32)[::2][1:4].numpy()
     np.testing.assert_equal(torch_res, cpu_res)
 
-  def test_select_out_of_range_dim(self):
-    a = torch.arange(12, dtype=torch.int32, device=device).reshape(3, 4)
-    with self.assertRaises(IndexError): a.select(5, 0)
-
-  def test_select_collapses_the_only_dim(self):
-    a = torch.arange(3, dtype=torch.int32, device=device)
-    self.assertEqual(a.select(0, 1).shape, ())
-    np.testing.assert_equal(a.select(0, 1).cpu().numpy(), 1)
-
   def test_slice_negative_dim(self):
     a = torch.arange(13, dtype=torch.int32, device=device).repeat(8, 1)
     torch_chunks = a.chunk(3, -1)
@@ -836,86 +796,6 @@ class TestTorchBackend(unittest.TestCase):
     np.testing.assert_allclose(w_tiny.grad.cpu().numpy(), w_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
     np.testing.assert_allclose(b_tiny.grad.cpu().numpy(), b_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
 
-  def test_write_through_detach_of_unrealized(self):
-    a = torch.empty(4, device=device)
-    a.detach().fill_(3)
-    np.testing.assert_equal(a.cpu().numpy(), [3, 3, 3, 3])
-
-  def test_square_transpose_inplace(self):
-    # a same-shape transpose is not a reshape: writing the transposed values straight back would scramble the base
-    a = torch.tensor([[0., 1., 2.], [3., 4., 5.], [6., 7., 8.]], device=device)
-    a.transpose(0, 1).add_(100)
-    np.testing.assert_equal(a.cpu().numpy(), [[100., 101., 102.], [103., 104., 105.], [106., 107., 108.]])
-
-  def test_interpolate(self):
-    a = torch.arange(4, dtype=torch.float32, device=device).reshape(1, 1, 2, 2)
-    nearest = torch.nn.functional.interpolate(a, scale_factor=2.0)
-    np.testing.assert_equal(nearest.cpu().numpy()[0, 0], [[0, 0, 1, 1], [0, 0, 1, 1], [2, 2, 3, 3], [2, 2, 3, 3]])
-    linear = torch.nn.functional.interpolate(a, size=(4, 4), mode="bilinear", align_corners=False)
-    ref = torch.nn.functional.interpolate(a.cpu(), size=(4, 4), mode="bilinear", align_corners=False)
-    np.testing.assert_allclose(linear.cpu().numpy(), ref.numpy(), rtol=1e-5)
-
-  def test_interpolate_bicubic_area(self):
-    a = torch.arange(32, dtype=torch.float32, device=device).reshape(1, 2, 4, 4)
-    for mode, scale in [("bicubic", 2.0), ("area", 0.5)]:
-      ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=scale, mode=mode)
-      np.testing.assert_allclose(torch.nn.functional.interpolate(a, scale_factor=scale, mode=mode).cpu().numpy(), ref.numpy(), atol=1e-4)
-
-  @unittest.expectedFailure
-  def test_interpolate_bicubic_backward(self):
-    # the forward comes from a decomposition, but aten::upsample_bicubic2d_backward has none (nor does
-    # aten::_adaptive_avg_pool2d_backward, for area), so training through these modes needs a real kernel
-    x = torch.arange(32., dtype=torch.float32, device=device).reshape(1, 2, 4, 4).requires_grad_()
-    torch.nn.functional.interpolate(x, scale_factor=2.0, mode="bicubic").sum().backward()
-
-  @unittest.expectedFailure
-  def test_interpolate_inexact_scale(self):
-    # torch forwards the raw scale_factor, Tensor.interpolate recomputes it from output_size, and they disagree here
-    a = torch.arange(6, dtype=torch.float32, device=device).reshape(1, 1, 2, 3)
-    tiny = torch.nn.functional.interpolate(a, scale_factor=2.5, mode="bilinear")
-    ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=2.5, mode="bilinear")
-    np.testing.assert_allclose(tiny.cpu().numpy(), ref.numpy(), rtol=1e-5)
-
-  def test_logical_or_xor(self):
-    a = torch.tensor([True, True, False, False], device=device)
-    b = torch.tensor([True, False, True, False], device=device)
-    np.testing.assert_equal(torch.logical_or(a, b).cpu().numpy(), [True, True, True, False])
-    np.testing.assert_equal(torch.logical_xor(a, b).cpu().numpy(), [False, True, True, False])
-    # bool-valued whatever the input dtype, so this is not | and ^
-    i, j = torch.tensor([2, 0, 5, 0], device=device), torch.tensor([0, 0, 1, 1], device=device)
-    np.testing.assert_equal(torch.logical_or(i, j).cpu().numpy(), [True, False, True, True])
-    np.testing.assert_equal(torch.logical_xor(i, j).cpu().numpy(), [True, False, False, True])
-
-  def test_slice_scatter(self):
-    # the scatters are functional: they return a new tensor and must leave the one they were given alone
-    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
-    out = torch.slice_scatter(a, torch.ones(1, 4, device=device), 0, 0, 1)
-    np.testing.assert_equal(out.cpu().numpy(), [[1, 1, 1, 1], [4, 5, 6, 7], [8, 9, 10, 11]])
-    np.testing.assert_equal(a.cpu().numpy(), np.arange(12, dtype=np.float32).reshape(3, 4))
-
-  def test_slice_scatter_casts_src(self):
-    a = torch.zeros(3, 4, device=device)
-    out = torch.slice_scatter(a, torch.ones(1, 4, dtype=torch.int32, device=device), 0, 0, 1)
-    self.assertEqual(out.dtype, torch.float32)
-    np.testing.assert_equal(out.cpu().numpy()[0], np.ones(4, dtype=np.float32))
-
-  def test_select_scatter(self):
-    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
-    out = torch.select_scatter(a, torch.ones(4, device=device), 0, 1)
-    np.testing.assert_equal(out.cpu().numpy(), [[0, 1, 2, 3], [1, 1, 1, 1], [8, 9, 10, 11]])
-
-  def test_diagonal_scatter(self):
-    a = torch.zeros(3, 3, device=device)
-    out = torch.diagonal_scatter(a, torch.arange(3, dtype=torch.float32, device=device))
-    np.testing.assert_equal(out.cpu().numpy(), np.diag([0., 1., 2.]))
-    np.testing.assert_equal(a.cpu().numpy(), np.zeros((3, 3), dtype=np.float32))
-
-  def test_copy_functional(self):
-    # without an impl this segfaults rather than fails: a regression here takes the whole run down
-    a = torch.arange(4, dtype=torch.float32, device=device)
-    out = torch.ops.aten.copy(a, torch.zeros(4, device=device))
-    np.testing.assert_equal(out.cpu().numpy(), [0., 0., 0., 0.])
-    np.testing.assert_equal(a.cpu().numpy(), [0., 1., 2., 3.])
 
 from tinygrad import Tensor
 class TestBackendHelpers(unittest.TestCase):

pyproject.toml

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

test/amd/test_emu2_pcode.py

@@ -67,26 +67,32 @@ class TestParseExpr(unittest.TestCase):
 
   def test_integer_literals(self):
     """Test parsing integer literals."""
-    self.assertIs(parse_expr('0', {}), UOp.const(0, dtypes.uint32))
-    self.assertIs(parse_expr('42', {}), UOp.const(42, dtypes.uint32))
-    self.assertIs(parse_expr('42U', {}), UOp.const(42, dtypes.uint32))
+    self.assertEqual(parse_expr('0', {}).val, 0)
+    self.assertEqual(parse_expr('42', {}).val, 42)
+    self.assertEqual(parse_expr('42U', {}).val, 42)
 
   def test_negative_integers(self):
     """Test parsing negative integer literals."""
-    self.assertIs(parse_expr('-1', {}), UOp.const(-1, dtypes.int))
+    result = parse_expr('-1', {})
+    self.assertEqual(result.val, -1)
+    self.assertEqual(result.dtype, dtypes.int)
 
   def test_float_literals(self):
     """Test parsing float literals."""
-    self.assertIs(parse_expr('1.0F', {}), UOp.const(1.0, dtypes.float32))
+    result = parse_expr('1.0F', {})
+    self.assertEqual(result.val, 1.0)
+    self.assertEqual(result.dtype, dtypes.float32)
 
   def test_hex_literals(self):
     """Test parsing hex literals."""
-    self.assertIs(parse_expr('0xFF', {}), UOp.const(255, dtypes.uint32))
+    result = parse_expr('0xFF', {})
+    self.assertEqual(result.val, 255)
 
   def test_variable_lookup(self):
     """Test variable lookup in parse_expr."""
     vrs = {'x': UOp.const(42, dtypes.uint32)}
-    self.assertIs(parse_expr('x', vrs), vrs['x'])
+    result = parse_expr('x', vrs)
+    self.assertEqual(result.val, 42)
 
   def test_binary_ops(self):
     """Test parsing binary operations."""
@@ -97,7 +103,9 @@ class TestParseExpr(unittest.TestCase):
     self.assertEqual(result.op, Ops.ADD)
 
     # Subtraction with constant folding
-    self.assertIs(parse_expr('10 - 5', {}), UOp.const(5, dtypes.uint32))
+    result = parse_expr('10 - 5', {})
+    self.assertEqual(result.op, Ops.CONST)
+    self.assertEqual(result.val, 5)
 
   def test_ternary(self):
     """Test parsing ternary expressions."""
@@ -134,8 +142,15 @@ class TestForLoopParsing(unittest.TestCase):
     S0 = UOp.const(0, dtypes.uint32)
     _vrs, assigns = parse_pcode(pcode, {'S0': S0})
 
-    # every cond folds (S0 is a const), leaving the default branch: -1 in the destination dtype
-    self.assertIs(assigns[0][1].simplify(), UOp.const(-1, dtypes.uint32))
+    # Check that the innermost value (default) is -1 (may be wrapped in CAST)
+    val = assigns[0][1]
+    # Traverse to innermost WHERE
+    while val.op == Ops.WHERE:
+      val = val.src[2]  # false branch
+    # Unwrap CAST if present
+    while val.op == Ops.CAST:
+      val = val.src[0]
+    self.assertEqual(val.val, -1)
 
   def test_ctz_parsing(self):
     """Test CTZ pcode parsing."""
@@ -247,8 +262,8 @@ class TestDSPcodePatterns(unittest.TestCase):
     _, assigns = parse_pcode(pcode, srcs)
     # Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
     # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
-    self.assertIs(assigns[0][1][0].simplify(), UOp.const(108, dtypes.uint32))  # type: ignore[index]
-    self.assertIs(assigns[1][1][0].simplify(), UOp.const(120, dtypes.uint32))  # type: ignore[index]
+    self.assertEqual(assigns[0][1][0].simplify().val, 108)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][0].simplify().val, 120)  # type: ignore[index]
 
   def test_ds_store_data_values(self):
     """Test DS_STORE_2ADDR_B32 uses correct data values."""
@@ -265,8 +280,8 @@ class TestDSPcodePatterns(unittest.TestCase):
     _, assigns = parse_pcode(pcode, srcs)
     # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
     # DATA[31:0] should preserve the value
-    self.assertIs(assigns[0][1][1].simplify(), UOp.const(0xAAAAAAAA, dtypes.uint32))  # type: ignore[index]
-    self.assertIs(assigns[1][1][1].simplify(), UOp.const(0xBBBBBBBB, dtypes.uint32))  # type: ignore[index]
+    self.assertEqual(assigns[0][1][1].simplify().val, 0xAAAAAAAA)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][1].simplify().val, 0xBBBBBBBB)  # type: ignore[index]
 
 class TestConditionalParsing(unittest.TestCase):
   """Test conditional (if/elsif/else) pcode parsing."""
@@ -291,12 +306,12 @@ class TestConcatWidthParsing(unittest.TestCase):
   def test_permlanex16_altrow_concat(self):
     for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]:
       parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(row, dtypes.uint32)})
-      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))
+      self.assertEqual(parsed.simplify().val, expected)
 
   def test_permlane64_altlane_concat(self):
     for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]:
       parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(lane, dtypes.uint32)})
-      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))
+      self.assertEqual(parsed.simplify().val, expected)
 
   def test_permlane64_wave64_pcode_indices(self):
     vgpr = UOp.param(0, dtypes.uint32, (256,))
@@ -312,17 +327,19 @@ class TestConcatWidthParsing(unittest.TestCase):
       'S2': UOp.const(0, dtypes.uint32),
     }
 
-    def check_load_idx(v: UOp, expected: int):
+    def load_idx(v: UOp) -> int:
       simp = v.simplify()
       self.assertEqual(simp.op, Ops.LOAD)
       self.assertEqual(simp.src[0].op, Ops.INDEX)
-      self.assertIs(simp.src[0].src[1].simplify(), UOp.const(expected, dtypes.uint32))
+      idx = simp.src[0].src[1].simplify()
+      self.assertEqual(idx.op, Ops.CONST)
+      return idx.val
 
     _, assigns = parse_pcode(PCODE[VOP1Op.V_PERMLANE64_B32_E32], srcs)
     self.assertEqual(len(assigns), 64)
     for lane, (dst_idx, src_idx) in {0: (64, 32), 31: (95, 63), 32: (96, 0), 63: (127, 31)}.items():
-      self.assertIs(assigns[lane][1][0].simplify(), UOp.const(dst_idx, dtypes.uint32))  # type: ignore[index]
-      check_load_idx(assigns[lane][1][1], src_idx)  # type: ignore[index]
+      self.assertEqual(assigns[lane][1][0].simplify().val, dst_idx)  # type: ignore[index]
+      self.assertEqual(load_idx(assigns[lane][1][1]), src_idx)  # type: ignore[index]
 
 class TestAllPcode(unittest.TestCase):
   """Test that all pcode from all architectures can be parsed."""

test/backend/test_asm_gemm.py

@@ -1,5 +1,4 @@
 import unittest
-import functools
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.helpers import getenv, system, DEV
 from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm
@@ -10,7 +9,6 @@ from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
 # Use DEV=NULL:HIP:gfx950 to also test the assembly
 def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950")
 
-@functools.cache
 def has_hipcc():
   try: system("hipcc --version")
   except Exception: return False

test/backend/test_const_folding.py

@@ -1,7 +1,7 @@
 import unittest, math
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DTYPES_DICT
-from tinygrad.uop.ops import Ops, UOp, GroupOp
+from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen.decomp.op import threefry2x32
 import numpy as np
 from test.helpers import not_support_multi_device
@@ -17,7 +17,7 @@ def _check_ast_count(desired_count:int, t:Tensor):
 class TestMovedConstFolding(unittest.TestCase):
   def test_contiguous_deviceless_const(self):
     t = Tensor(UOp.const(2.0, dtypes.float)).contiguous()
-    self.assertIs(t.uop, UOp.const(2.0, dtypes.float))
+    self.assertIs(t.uop.op, Ops.CONST)
     self.assertIsNone(t.uop.device)
 
   def test_add_shrunk_zero(self):
@@ -169,8 +169,8 @@ class TestMultiConstFolding(unittest.TestCase):
 class TestThreefryConstFolding(unittest.TestCase):
   def test_threefry(self):
     # THREEFRY(const,const) folds to a const once decomposed
-    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64)).simplify()
-    self.assertEqual([u.op for u in x.toposort() if u.op in GroupOp.ALU], [])
+    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64))
+    self.assertIs(x.simplify().op, Ops.CONST)
 
 class TestTautologicalCompare(unittest.TestCase):
   # without const folding, these would have triggered -Wtautological-compare in clang
@@ -188,6 +188,7 @@ class TestTautologicalCompare(unittest.TestCase):
     np.testing.assert_equal((Tensor(True) < Tensor(False)).numpy(), False)
     np.testing.assert_equal((Tensor(True) < Tensor(True)).numpy(), False)
 
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly")
   def test_a_eq_a(self):
     # self eq is always true for int or bool
     a = Tensor([1, 2, 3])

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, KernelCountException
+from test.helpers import assert_kernel_count
 
 # **** kernels ****
 
@@ -474,7 +474,7 @@ class TestCustomKernelInput(unittest.TestCase):
     y.realize()
     kernel_count = GlobalCounters.kernel_count
     self.assertEqual(y.tolist(), x.add(1).tolist())
-    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
+    self.assertLessEqual(kernel_count, max_kernels)
     # 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())
-    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
+    self.assertLessEqual(kernel_count, max_kernels)
 
   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_dtype.py

@@ -340,9 +340,6 @@ class TestUint64DType(TestDType):
   DTYPE = dtypes.uint64
   def test_uint64_load(self):
     assert Tensor(2**64 - 1, dtype=dtypes.uint64).numpy() == 2**64 - 1
-  @unittest.skipIf(dtypes.double not in supported_dtypes, "needs float64")
-  def test_uint64_cast_double(self):
-    assert Tensor([2**32 + 1], dtype=dtypes.uint64).cast(dtypes.double).numpy() == 2**32 + 1
 
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
 class TestEmulatedUInt64DType(TestUint64DType):

test/backend/test_isel.py

@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext
 
 # INDEX on a register value with a constant index extracts a single element (the old GEP)
-def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int), dtype=y.dtype)
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype.scalar())
 
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):
@@ -46,10 +46,10 @@ class TestIselX86(unittest.TestCase):
   # complex address is [base + index*scale + displacement]
   def test_complex_address(self):
     a = UOp.variable("a", 0, 0, dtypes.int32)
-    load = UOp.param(0, dtypes.int32, (16,)).index(a + UOp.cconst(1, dtypes.int32)).load()
+    load = UOp.param(0, dtypes.int32, (16,)).index(a + 1).load()
     n = self.isel_rewrite(load)
     # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
-    self.assertTrue(n.src[2].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4)
+    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].val == 4)
 
 if __name__ == "__main__":
   unittest.main()

test/backend/test_linearizer.py

@@ -16,6 +16,8 @@ from test.helpers import replace_opts, check_schedule
 from test.backend.test_softmax_fusion import single_kernel_softmax
 MOCKGPU = DEV.interface.startswith("MOCK")
 
+from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
+
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, ISARenderer), "isa backends don't preserve the op spec when lowering")
 class TestLinearizer(unittest.TestCase):
   def test_arg_dedup(self):
@@ -246,10 +248,11 @@ class TestLinearizer(unittest.TestCase):
     uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
     begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
     end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
+    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
     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 not in GroupOp.ALU
+          assert u.src[1].op is Ops.CONST
         else:
           assert u.src[1].op in GroupOp.ALU
           assert begin_range < uops.index(u) < end_range
@@ -265,9 +268,9 @@ class TestLinearizer(unittest.TestCase):
     uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
     idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
     idxs = sorted(idxs, key=lambda uop: uop.arg)
-    assert (idxs[0].arg, idxs[0].src[0].src[0].val) == ('gidx0', 6), idxs[0]
-    assert (idxs[1].arg, idxs[1].src[0].src[0].val) == ('gidx1', 5), idxs[1].arg
-    assert (idxs[2].arg, idxs[2].src[0].src[0].val) == ('gidx2', 4), idxs[2].arg
+    assert (idxs[0].arg, idxs[0].src[0].val) == ('gidx0', 6), idxs[0]
+    assert (idxs[1].arg, idxs[1].src[0].val) == ('gidx1', 5), idxs[1].arg
+    assert (idxs[2].arg, idxs[2].src[0].val) == ('gidx2', 4), idxs[2].arg
 
   def test_sum_collapse(self):
     t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()

test/backend/test_llama_kernels.py

@@ -99,20 +99,22 @@ class TestLocalAmax(unittest.TestCase):
     assert_kernel_count(2)
     self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])
 
+@unittest.skipUnless(has_hipcc() and Device.DEFAULT == "AMD", "requires hipcc to compile and amd device to run")
 class TestFusedQKVRoPE(unittest.TestCase):
   SHAPE = (2, 8192, 32, 8, 128)
 
-  def setUp(self):
-    if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("test uses bf16 inputs")
-
   def rand_bf16(self, *shape:int) -> Tensor:
     return (Tensor.randn(*shape) * 0.1).cast(dtypes.bfloat16).contiguous().realize()
 
-  def test_forward(self):
+  def freqs_cis(self) -> Tensor:
+    _, N, _, _, D = self.SHAPE
+    return precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
+
+  def test_llama31_8b_forward(self):
     Tensor.manual_seed(0)
-    B, N, H, H_KV, D = 1, 32, 8, 2, 16
+    B, N, H, H_KV, D = self.SHAPE
     GROUP = H // H_KV
-    freqs_cis = (Tensor.randn(1, N * 2, 1, D // 2, 2) * 0.1).cast(dtypes.bfloat16).contiguous().realize()
+    freqs_cis = self.freqs_cis()
 
     x = self.rand_bf16(B, N, H_KV * (GROUP + 2) * D)
     q, k, v = fused_qkv_rope(x, freqs_cis, H, H_KV, D)
@@ -129,13 +131,12 @@ class TestFusedQKVRoPE(unittest.TestCase):
       self.assertTrue(k.allclose(k_ref, atol=2e-2, rtol=0).item(), "K forward mismatch")
       self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")
 
-  @unittest.skipUnless(has_hipcc(), "backward kernel requires hipcc to compile")
-  def test_llama31_8b(self):
+  def test_llama31_8b_backward(self):
     Tensor.manual_seed(1)
     B, N, H, H_KV, D = self.SHAPE
     PARTIALS = 2
     GROUP = H // H_KV
-    freqs_cis = precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
+    freqs_cis = self.freqs_cis()
     dq = self.rand_bf16(B, N, H, D)
     dk_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
     dv_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)

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, KernelCountException
+from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count
 
 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]
-      if len(compiled) != 0: raise KernelCountException(0, len(compiled))
+      self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
     run_linear(linear, var_vals)
     np.testing.assert_equal(b_multi.numpy(), b_ref.numpy())
 

test/backend/test_ops.py

@@ -720,11 +720,10 @@ class TestOps(unittest.TestCase):
       return torch.autograd.grad(t ** c, t)[0].item()
     for x in [-math.inf, 0, 1, math.inf]:
       for c in [-1, 0, 0.3, 1, 2]:
-        torch_out = get_torch_gradient(x, c)
-        # the pow backward routes through exp2/log2, whose 0/inf behavior is undefined on WEBGPU
-        if Device.DEFAULT == "WEBGPU" and not math.isfinite(torch_out): continue
         tiny_out = get_tiny_gradient(x, c)
+        torch_out = get_torch_gradient(x, c)
         if math.isnan(tiny_out):
+          if Device.DEFAULT == "WEBGPU": continue # TODO: WEBGPU issue with nan
           assert math.isnan(torch_out)
         else:
           self.assertAlmostEqual(tiny_out, torch_out, msg=f"{x}, {c}")
@@ -750,7 +749,6 @@ class TestOps(unittest.TestCase):
   def test_exp2_log2_zero_times_negative(self):
     # gallivm's exp2/log2 have "undefined behavior with infs, 0s and nans", so exp2(log2(0)*y) returns 0 instead of inf
     helper_test_op(None, lambda x,y: (x.log2()*y).exp2(), lambda x,y: (x.log2()*y).exp2(), vals=[[0.0], [-0.7]], forward_only=True)
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "pow at 0 routes through exp2/log2, whose 0/inf behavior is undefined on WEBGPU")
   def test_pow_zero_const(self):
     helper_test_op(None, lambda x: x**0.3, vals=[[0.0]])
     helper_test_op(None, lambda x: x**0.0, vals=[[0.0]])
@@ -2164,10 +2162,6 @@ class TestOps(unittest.TestCase):
   def test_roll(self):
     helper_test_op([(2, 4)], lambda x: x.roll(1))
     helper_test_op([(2, 4)], lambda x: x.roll((1,)))
-    helper_test_op([(0,)], lambda x: x.roll(1, 0))
-    helper_test_op([(2, 0, 3)], lambda x: x.roll(1, 0))
-    helper_test_op([(2, 0, 3)], lambda x: x.roll(1, 1))
-    helper_test_op([(2, 0, 3)], lambda x: x.roll(1))
     self.helper_test_exception([(2, 4)], lambda x: x.roll((1, 2)), expected=RuntimeError)
     helper_test_op([(2, 4)], lambda x: x.roll(1, 0))
     helper_test_op([(2, 4)], lambda x: x.roll(-1, 0))

test/backend/test_pickle.py

@@ -3,7 +3,6 @@ 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):
@@ -42,7 +41,7 @@ class TestPickle(unittest.TestCase):
     t2:Tensor = pickle.loads(st)
     np.testing.assert_equal(t_values, t2.numpy())
     # expect at most one COPY kernel
-    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)
+    self.assertLessEqual(GlobalCounters.kernel_count, 1)
 
   def test_pickle_realized_tensor_alt(self):
     print("** init")

test/backend/test_quantize_onnx.py

@@ -82,7 +82,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
       linear = run_onnx({"input":inp})["output"].schedule_linear()
       prg = to_program(linear.src[-2].src[0], renderer=Device[Device.DEFAULT].renderer)
       daccs = [u for u in tuple(prg.src[1].src) if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
-      assert all(u.dtype is dtypes.int for u in daccs)
+      assert all(u.dtype.scalar() is dtypes.int for u in daccs)
 
 @unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")
 class TestQuantizeOnnx(unittest.TestCase):

test/backend/test_tensor.py

@@ -653,19 +653,9 @@ class TestZeroShapeTensor(unittest.TestCase):
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3).numpy(), [[1, 2, 0], [0, 0, 0]])
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(1, 3).numpy(), [[1, 2, 0]])
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(None, 3).numpy(), [[1, 2, 0]])
-    np.testing.assert_equal(Tensor([1, 2]).pad_to(4, value=2).numpy(), [1, 2, 2, 2])
-    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3, value=-1).numpy(), [[1, 2, -1], [-1, -1, -1]])
-    np.testing.assert_equal(Tensor([1, 2]).pad_to(None, value=5).numpy(), [1, 2])  # no-op pad ignores the fill
     with self.assertRaises(ValueError): Tensor([1, 2]).pad_to(2, 3)
     with self.assertRaises(ValueError): Tensor([[1, 2]]).pad_to(3)
 
-  def test_max_shape(self):
-    from tinygrad import UOp
-    t = Tensor.empty(2, UOp.variable('v', 1, 32), 4)
-    self.assertEqual(t.max_shape, (2, 32, 4))
-    self.assertEqual(t.max_numel(), 2*32*4)
-    self.assertEqual(Tensor.empty(2, 3).max_shape, (2, 3))
-
   def test_shrink_into_zero(self):
     t = Tensor.rand(3, 4).realize()
     assert t.shrink((None, (2, 2))).realize().shape == (3, 0)

test/device/test_hcq2.py

@@ -10,13 +10,5 @@ class TestHCQ2(unittest.TestCase):
     with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):
       np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61))
 
-  def test_overlapping_device_tuples(self):
-    # an op on a wide device tuple followed by an op on an overlapping smaller tuple used to MMU-fault the smaller one
-    d4, d2 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4)), tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
-    ref = Tensor.arange(16).contiguous().realize()
-    Tensor(ref.uop.copy_to_device(d4)).realize()
-    out = Tensor.ones(8).shard(d2, axis=0).contiguous().realize()
-    np.testing.assert_equal(out.numpy(), np.ones(8))
-
 if __name__ == "__main__":
   unittest.main()

test/external/external_test_usb_asm24.py

@@ -6,7 +6,7 @@ import numpy as np
 class TestDevCopySpeeds(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
-    cls.sz = getenv("SIZE", 2000000)
+    cls.sz = getenv("SIZE", 2e6)
     cls.dev = Device["AMD"]
     if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
 

test/external/external_uop_gc.py

@@ -44,7 +44,7 @@ def realized_matmul():
   z = y.matmul(x)
   Tensor.realize(z)
 def realized_gradient():
-  x = Tensor.eye(3).clone()
+  x = Tensor.eye(3)
   y = Tensor([[2.0,0,-2.0]])
   z = y.matmul(x).sum()
   z.backward()

test/helpers.py

@@ -86,9 +86,7 @@ def assert_jit_cache_len(fxn, expected_len):
   if linear is None or not linear.src:
     if expected_len != 0: raise KernelCountException(expected_len, 0)
     return
-  if expected_len and all(call_is_hcq(call) for call in linear.src): # HCQ2: one batch submitter, or fence + reset + merged calls + finalizer
-    from tinygrad.runtime.support.hcq2 import HCQ_RUNTIME_DEV
-    expected_len = 1 if HCQ_RUNTIME_DEV.value == "CPU" else 4
+  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 4 # HCQ2: fence + reset + merged same-queue calls + finalizer
   if call_is_graph(linear.src[0]):
     if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
     inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION

test/null/test_const_folding.py

@@ -1,10 +1,39 @@
 import unittest, itertools, math
-from tinygrad import dtypes, Context
+from tinygrad import Tensor, dtypes, Context
 from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
 
+def _check_ast_count(desired_count:int, t:Tensor):
+  # NOTE: this has side effect because everything can be scheduled only once
+  linear = t.schedule_linear()
+  asts = [s for s in linear.src if s.src[0].op is Ops.SINK]
+  len(asts)
+  # NOT SUPPORTED ANYMORE
+  #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
+
+class TestUnaryOpsConstFolding(unittest.TestCase):
+  def test_all_consts_ops(self):
+    _check_ast_count(0, Tensor.ones(4).exp())
+    _check_ast_count(0, Tensor.ones(4).sqrt())
+    _check_ast_count(0, Tensor.ones(4) + Tensor.ones(4))
+    _check_ast_count(0, Tensor.ones(4) / Tensor.ones(4))
+
+  def test_cast(self):
+    _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
+    _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
+
+  def test_neg_folding(self):
+    _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
+    _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
+    _check_ast_count(0, Tensor([1, 2, 3]).neg().neg())
+
+  def test_neg_realized_no_fold(self):
+    x = Tensor.randn(32, 32)
+    x = x.clip(0, 1).realize()
+    _check_ast_count(1, x.neg())
+
 class TestWeakConstFolding(unittest.TestCase):
   def test_weakint_math(self):
     out = (UOp.const(2**40) + UOp.const(2**40)).simplify()
@@ -22,18 +51,84 @@ class TestWeakConstFolding(unittest.TestCase):
   def test_invalid_poison(self):
     self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
 
-class TestBitcastConstFolding(unittest.TestCase):
-  def test_out_of_range_source_value(self):
-    for val, src_dt, dst_dt, bits in ((3000000000, dtypes.int32, dtypes.uint32, 3000000000),
-                                      (70000, dtypes.int16, dtypes.uint16, 4464),
-                                      (-5, dtypes.uint32, dtypes.int32, -5)):
-      self.assertEqual(UOp.const(val, src_dt).bitcast(dst_dt).simplify().val, bits)
+  def test_single_rounding_log10_backward(self):
+    # log10 backward folds log10(2)/log(2) = 1/log(10) in one rounding, not the double-rounded 1/float32(log(10))
+    x = Tensor([1.0, 2.0, 3.0])
+    ast = next(s.src[0] for s in x.log10().sum().gradient(x)[0].schedule_linear().src if s.src[0].op is Ops.SINK)
+    const = next(u.arg for u in full_rewrite(ast).toposort() if u.op is Ops.CONST and u.dtype is dtypes.float32)
+    # correctly rounded: within half a float32 ulp of the exact value (folding at float32 lands 0.66 ulp off)
+    self.assertLess(abs(const - 1/math.log(10)), 2**-26)
 
+class TestBinaryOpsConstFolding(unittest.TestCase):
+  def test_add_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
+  def test_add_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(4))
+  def test_literal_zero_add(self):
+    _check_ast_count(0, 0 + Tensor([1.0, 2, 3, 4]))
+  def test_tensor_zero_add(self):
+    _check_ast_count(0, Tensor.zeros(4) + Tensor([1.0, 2, 3, 4]))
+
+  def test_sub_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - 0)
+  def test_sub_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - Tensor.zeros(4))
+
+  def test_mul_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 0)
+  def test_mul_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.zeros(4))
+  def test_literal_zero_mul(self):
+    _check_ast_count(0, 0 * Tensor([1.0, 2, 3, 4]) * 0)
+  def test_tensor_zero_mul(self):
+    _check_ast_count(0, Tensor.zeros(4) * Tensor([1.0, 2, 3, 4]))
+
+  def test_mul_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 1)
+  def test_mul_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(4))
+  def test_literal_one_mul(self):
+    _check_ast_count(0, 1 * Tensor([1.0, 2, 3, 4]))
+  def test_tensor_one_mul(self):
+    _check_ast_count(0, Tensor.ones(4) * Tensor([1.0, 2, 3, 4]))
+
+  def test_bool_tensor_mul_bool(self):
+    _check_ast_count(0, Tensor([True, False]) * True)
+    _check_ast_count(0, Tensor([True, False]) * False)
+  def test_bool_mul_bool_tensor(self):
+    _check_ast_count(0, True * Tensor([True, False]))
+    _check_ast_count(0, False * Tensor([True, False]))
+
+  def test_div_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / 1)
+  def test_div_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / Tensor.ones(4))
+
+  def test_floordiv_literal_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
+  def test_floordiv_tensor_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) // Tensor.ones(4, dtype=dtypes.int32))
+
+  def test_pow_literal_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 0)
+  def test_pow_tensor_zero(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.zeros(4))
+
+  def test_pow_literal_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 1)
+  def test_pow_tensor_one(self):
+    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.ones(4))
+  def test_literal_one_pow(self):
+    _check_ast_count(0, 1 ** Tensor([1.0, 2, 3, 4]))
+  def test_tensor_one_pow(self):
+    _check_ast_count(0, Tensor.ones(4) ** Tensor([1.0, 2, 3, 4]))
+
+class TestBitcastConstFolding(unittest.TestCase):
   def test_scalar_bitcast(self):
     def t(cases: dict[DType, ConstType]):
       for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
         if not math.isnan(from_v):
-          r = UOp.const(from_v, from_dt).bitcast(to_dt).simplify()
+          r = full_rewrite(UOp.const(from_v, from_dt).bitcast(to_dt).sink()).src[0]
           self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
           self.assertEqual(r.dtype, to_dt, msg)
           np.testing.assert_equal(r.val, to_v, msg)
@@ -57,9 +152,24 @@ class TestBitcastConstFolding(unittest.TestCase):
 
   def test_vec_bitcast(self):
     with Context(SPEC=0):
-      result = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink())
-      expected = full_rewrite(UOp.const((2**32-1, 2**31, 75), dtypes.uint32).sink())
-    self.assertEqual(result.src, expected.src)
+      srcs = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink()).src
+    self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs))
+    self.assertEqual(tuple(x.val for x in srcs), (2**32-1, 2**31, 75))
+
+# folds advance indexing into basic indexing
+class TestIndexingConstFolding(unittest.TestCase):
+  def test_scalar_index(self):
+    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
+    _check_ast_count(1, t[:,:,Tensor(1),:])
+    _check_ast_count(1, t[:,:,Tensor(1)+2,:])
+    _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
+
+  def test_const_tensor_index(self):
+    # TODO: these can be 0, implement const tensor folded indexing
+    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
+    _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
+    _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
+    _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_dtype_spec.py

@@ -51,6 +51,10 @@ class TestHelpers(unittest.TestCase):
     assert dtypes.is_float(dtypes.fp8e4m3)
     assert dtypes.is_float(dtypes.fp8e5m2)
 
+  @given(strat.sampled_from([d for d in DTYPES_DICT.values() if dtypes.is_float(d) or dtypes.is_int(d)]))
+  def test_scalar(self, dtype):
+    assert dtype.scalar() == dtype
+
   def test_from_py(self):
     assert dtypes.from_py(True) == dtypes.bool
     assert dtypes.from_py(Invalid) == dtypes.bool

test/null/test_elf.py

@@ -1,7 +1,6 @@
 import unittest, subprocess, platform
 from tinygrad.runtime.support.compiler_cpu import ClangCompiler
 from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.runtime.support.c import DLL
 
 class TestElfLoader(unittest.TestCase):
   def test_load_clang_jit_strtab(self):
@@ -24,7 +23,7 @@ class TestElfLoader(unittest.TestCase):
       }
     '''
     with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
-      elf_loader(ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src))
+      ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src)
   def test_link(self):
     src = '''
       float powf(float, float); // from libm
@@ -33,7 +32,7 @@ class TestElfLoader(unittest.TestCase):
     args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
     obj = subprocess.check_output(('clang',) + args + ('-', '-o', '-'), input=src.encode())
     with self.assertRaisesRegex(RuntimeError, 'powf'): elf_loader(obj)
-    elf_loader(obj, link_libs=[DLL('m', 'm')])
+    elf_loader(obj, link_libs=['m'])
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_graph_rewrite.py

@@ -1,7 +1,8 @@
 import unittest, math
 from tinygrad import dtypes
+from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import all_same, Context
-from tinygrad.uop.ops import GroupOp, UOp, Ops, PatternMatcher, TrackedPatternMatcher, UPat
+from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, TrackedPatternMatcher, UPat
 from test.helpers import full_rewrite
 from hypothesis import given, strategies as strat
 
@@ -10,14 +11,125 @@ from hypothesis import given, strategies as strat
 def apply_rewrite(expr):
   return full_rewrite(expr.sink()).src[0]
 
-def const_value(uop:UOp):
-  if uop.op is Ops.CAST: uop = uop.src[0]
-  assert uop.op is Ops.CONST
-  return uop.val
+@Context(SPEC=0)
+def apply_rewrite_values(expr):
+  srcs = full_rewrite(expr.sink()).src
+  if len(srcs) == 1:
+    if srcs[0].op is Ops.CONST: return (srcs[0].val,)
+    if srcs[0].op is Ops.STACK: return tuple(s.val for s in srcs[0].src)
+  return tuple(s.val for s in srcs)
+
+def evaluate_uop(uop, variables):
+  if uop.op == Ops.CONST:
+    return uop.val
+  elif uop.op == Ops.PARAM and uop.arg.addrspace is AddrSpace.ALU:
+    return variables[uop.expr]
+  elif uop.op in GroupOp.ALU:
+    src_values = [evaluate_uop(src, variables) for src in uop.src]
+    return exec_alu(uop.op, uop.dtype, src_values)
+  else:
+    raise NotImplementedError(f"Unsupported UOp {uop.op}")
+
+class TestArithmeticSimplifications(unittest.TestCase):
+  def test_full_graph_rewrite_division_by_zero(self):
+    optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0))
+    self.assertEqual(optimized_div_uop.op, Ops.CONST)
+    self.assertTrue(math.isinf(optimized_div_uop.val) or math.isnan(optimized_div_uop.val))
+
+  def test_full_graph_rewrite_redundant_operations(self):
+    optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0))
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.val, 10.0)
+
+  def test_full_graph_rewrite_large_graph(self):
+    prev_uop = UOp.const(0)
+    for i in range(1, 101):
+      prev_uop += UOp.const(i)
+    optimized_uop = apply_rewrite(prev_uop)
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.val, sum(range(1, 101)))
+
+  def test_full_graph_rewrite_division_by_one(self):
+    optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0))
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.val, 42.0)
+
+  def test_full_graph_rewrite_modulo_by_one(self):
+    optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1))
+    self.assertEqual(optimized_uop.op, Ops.CONST)
+    self.assertEqual(optimized_uop.val, 0)
+
+
+class TestFoldingAndReduction(unittest.TestCase):
+  @unittest.skip("reduce is removed now")
+  def test_full_graph_rewrite_constant_reduction_folding(self):
+    const1 = UOp.const(5)
+    const2 = UOp.const(10)
+    const3 = UOp.const(20)
+    optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD))
+    expected_sum = 5 + 10 + 20
+    self.assertEqual(optimized_sink.val, expected_sum)
+
+  @unittest.skip("reduce is removed now")
+  def test_full_graph_rewrite_reduction_with_unused_range(self):
+    const1 = UOp.const(15)
+    const2 = UOp.const(25)
+    rng = UOp.range(10, idx=0)
+    optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
+    expected_sum = 10 * (15 + 25)
+    self.assertEqual(optimized_sink.val, expected_sum)
+
+  @unittest.skip("currently failing")
+  def test_full_graph_rewrite_range_reduction(self):
+    simple_range = UOp.range(5, idx=0)
+    optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
+    expected_sum = sum(range(5))
+    self.assertEqual(optimized_sink.val, expected_sum)
+
+  @unittest.skip("currently failing")
+  def test_full_graph_rewrite_simple_reduction_folding(self):
+    simple_range = UOp.range(4, idx=0)
+    add_uop = simple_range + UOp.const(1)
+    optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
+    expected_sum = sum(i + 1 for i in range(4))
+    self.assertEqual(optimized_sink.val, expected_sum)
+
+  @unittest.skip("currently failing")
+  def test_full_graph_rewrite_nested_loop_collapse(self):
+    outer_range = UOp.range(8, 0)
+    inner_range = UOp.range(4, 1)
+    expr = (outer_range * 10) + inner_range
+    optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range))
+    self.assertEqual(optimized_reduce_uop.op, Ops.CONST)
+    self.assertEqual(optimized_reduce_uop.val, sum((i * 10) + j for i in range(8) for j in range(4)))
+
 
 class TestModuloAndDivisionFolding(unittest.TestCase):
+  def test_full_graph_rewrite_modulo_folding_with_define_var(self):
+    # index dtype because div-mod rules only work on index
+    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
+    optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
+    self.assertEqual(optimized_mod_uop.op, Ops.CONST)
+    self.assertEqual(optimized_mod_uop.val, 2)
+
+  def test_full_graph_rewrite_division_folding_with_define_var(self):
+    # index dtype because div-mod rules only work on index
+    n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
+    optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
+    self.assertEqual(optimized_div_uop.op, Ops.MUL)
+    self.assertEqual(optimized_div_uop.src[1].val, 2)
+
+  def test_full_graph_rewrite_complex_mod_div_folding(self):
+    # index dtype because div-mod rules only work on index
+    k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
+    optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
+    self.assertEqual(optimized_div_uop.op, Ops.CONST)
+    self.assertEqual(optimized_div_uop.val, 1)
+
   def test_graph_rewrite_div_folding_bug(self):
-    lhs = UOp.stack(*(UOp.special(32, 'lidx0'),)*4) + UOp.const((0, 256, 512, 768))
+    lhs = UOp(Ops.ADD, src=(
+      UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32),), arg='lidx0'),)*4),
+      UOp.const((0, 256, 512, 768))))
     rhs = UOp.const((2,)*4)
     unopt = lhs<rhs
     opt = apply_rewrite(unopt)
@@ -25,31 +137,74 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
     print(opt)
     if opt.op is Ops.STACK: self.assertFalse(all_same(opt.src))
 
+  def test_full_graph_rewrite_modulo_large_divisor(self):
+    # index dtype because div-mod rules only work on index
+    x_var_uop = UOp.variable('x', 1, 5)
+    self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.weakint) % 10).render(simplify=False), x_var_uop.render(simplify=False))
+
+  def test_full_graph_rewrite_division_with_remainder(self):
+    x_var_uop = UOp.variable('x', 7, 9, param=True)
+    optimized_sink = apply_rewrite(x_var_uop // 2)
+    for x_value in range(7, 10):
+      self.assertEqual(x_value // 2, evaluate_uop(optimized_sink, {'x': x_value}))
+
+  def test_full_graph_rewrite_complex_mod_div_expression(self):
+    x_var_uop = UOp.variable('x', 1, 10, param=True)
+    optimized_sink = apply_rewrite(((x_var_uop * 5) % 3) // 2)
+    for x_value in range(1, 11):
+      original_result = ((x_value * 5) % 3) // 2
+      optimized_result = evaluate_uop(optimized_sink, {'x': x_value})
+      self.assertEqual(original_result, optimized_result)
+
+
 class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
   def test_full_graph_rewrite_transcendental_edge_cases(self):
     optimized_sink = full_rewrite(UOp.const(-1.0).log2().sink(UOp.const(0.0).reciprocal()))
     optimized_log2_neg, optimized_recip_zero = optimized_sink.src
-    log2_neg, recip_zero = const_value(optimized_log2_neg), const_value(optimized_recip_zero)
-    self.assertTrue(math.isnan(log2_neg), f"Expected NaN for log2(-1.0), got {log2_neg}")
-    self.assertTrue(math.isinf(recip_zero) and recip_zero > 0, f"Expected +inf for reciprocal(0.0), got {recip_zero}")
+    self.assertTrue(math.isnan(optimized_log2_neg.val), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.val}")
+    self.assertTrue(math.isinf(optimized_recip_zero.val) and optimized_recip_zero.val > 0,
+                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.val}")
+
+  @unittest.skip("broken")
+  def test_full_graph_rewrite_modulo_negative_dividend(self):
+    x_var_uop = UOp.variable('x', -5, -1)
+    optimized_sink = full_rewrite((x_var_uop % 3).sink())
+    for x_value in range(-5, 0):
+      self.assertEqual(x_value % 3, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
+
+  @unittest.skip("broken")
+  def test_full_graph_rewrite_division_negative_divisor(self):
+    x_var_uop = UOp.variable('x', 1, 5)
+    optimized_sink = full_rewrite((x_var_uop // -2).sink())
+    for x_value in range(1, 6):
+      self.assertEqual(x_value // -2, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
 
 class TestGEPAndVectorizeRewrite(unittest.TestCase):
   def test_gep_single_element_extraction(self):
     # GEP on a vector dtype to extract a single element
     base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertIs(apply_rewrite(base_vector.index(2)), apply_rewrite(base_vector.src[2]))
+    self.assertEqual(apply_rewrite(base_vector.index(2)).val, 3.0)
 
   def test_gep_tuple_extraction(self):
     # GEP on a vector dtype to extract multiple elements as a vector
     base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertIs(apply_rewrite(UOp.stack(*[base_vector.index(i) for i in (2, 3)])),
-                  apply_rewrite(UOp.stack(base_vector.src[2], base_vector.src[3])))
+    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])
+
+  def test_gep_on_const_stack(self):
+    # GEP on a const STACK to extract a single element
+    const_stack = UOp.const((1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(const_stack.index(2)).val, 3.0)
+
+  def test_gep_tuple_on_const_stack(self):
+    # GEP on a const STACK using a tuple to extract multiple elements
+    const_stack = UOp.const((7.0, 8.0, 9.0, 10.0))
+    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])
 
   def test_vectorize_multiple_elements(self):
     # Vectorizing multiple elements using GEP
     base_vector = UOp.const((5.0, 10.0, 15.0, 20.0))
-    vectorized_uop = UOp.stack(*(base_vector.index(i) for i in range(4)))
-    self.assertIs(apply_rewrite(vectorized_uop), apply_rewrite(base_vector))
+    vectorized_uop = UOp(Ops.STACK, src=tuple(base_vector.index(i) for i in range(4)))
+    self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])
 
 
 import inspect
@@ -101,6 +256,16 @@ class TestSubstitute(unittest.TestCase):
     ret = substitute(ret, {a.sin():b})
     self.assertIs(ret, b.sin())
 
+  # broken due to infinite recursion
+  # NOTE: VIZ hangs and doesn't recover if you click this one
+  @unittest.skip("recursion error no longer raised")
+  def test_assert_inf_recurse(self):
+    a = UOp.variable('a', 0, 10)
+    n1 = a.sin()
+    ret = n1
+    with self.assertRaises(RecursionError):
+      ret = substitute(ret, {n1:n1.sqrt()})
+
   def test_sin_to_sqrt(self):
     a = UOp.variable('a', 0, 10, dtype=dtypes.float)
     n1 = a.sin()

test/null/test_mnist_dataset.py

@@ -1,7 +1,6 @@
 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):
@@ -9,7 +8,7 @@ class TestDataset(unittest.TestCase):
     X_train[0].contiguous().realize()
     GlobalCounters.reset()
     X_train[0].contiguous().realize()
-    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)  # 0 if SLICE (zero-copy), 1 otherwise
+    self.assertLessEqual(GlobalCounters.kernel_count, 1)  # 0 if SLICE (zero-copy), 1 otherwise
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_simplify_valid_idx.py

@@ -15,12 +15,16 @@ def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move
 def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")
 
 def get_gated_load_uop(valid:UOp, idx:UOp):
-  return UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)).load()
+  return UOp(Ops.LOAD, src=(
+    UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)),
+  ))
 
 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
-  return UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)).load()
+  return UOp(Ops.LOAD, src=(
+    UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)),
+  ))
 
-def Special(expr, nmax): return UOp.special(nmax, expr)
+def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(nmax),), arg=expr)
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax, param=True)
 def Range(n, nmax): return UOp.range(nmax, n)
 
@@ -508,7 +512,7 @@ class TestDropTrueGate(unittest.TestCase):
     buf = UOp.param(0, dtypes.int, (1,))
     idx = UOp.const(0)
     true_gate = UOp.const(True)
-    index_with_gate = buf.index(idx.valid(true_gate))
+    index_with_gate = UOp(Ops.INDEX, src=(buf, idx.valid(true_gate)))
     # apply the optimization
     result = graph_rewrite(index_with_gate, sym+indexing_simplify)
     # the True valid should be dropped (INDEX should only have 2 sources)
@@ -520,17 +524,13 @@ class TestRangeShrink(unittest.TestCase):
       result = full_rewrite(sink)
     return [u for u in result.toposort() if u.op is Ops.RANGE]
 
-  def assert_range_end(self, ranges:list[UOp], end:int):
-    self.assertEqual(len(ranges), 1)
-    with Context(NOOPT=1, SPEC=0): expected = full_rewrite(UOp.const(end, dtypes.int).sink()).src[0]
-    self.assertIs(ranges[0].src[0], expected)
-
   def test_range_shrink_single_guard(self):
     # range 0..203 guarded by r < 4 everywhere -> shrink to 0..3
     r = Range(0, 204)
     load = get_gated_load_uop(r < UOp.const(4), r)
     ranges = self.get_ranges(load.sink())
-    self.assert_range_end(ranges, 4)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 4)
 
   def test_range_shrink_picks_max_guard(self):
     # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
@@ -538,22 +538,25 @@ class TestRangeShrink(unittest.TestCase):
     load1 = get_gated_load_uop(r < UOp.const(4), r)
     load2 = get_gated_load_uop(r < UOp.const(8), r)
     ranges = self.get_ranges(UOp.sink(load1, load2))
-    self.assert_range_end(ranges, 8)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 8)
 
   def test_range_no_shrink_guard_ge_max(self):
     # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
     r = Range(0, 204)
     load = get_gated_load_uop(r < UOp.const(300), r)
     ranges = self.get_ranges(load.sink())
-    self.assert_range_end(ranges, 204)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 204)
 
   def test_range_no_shrink_when_unguarded_elsewhere(self):
     # one load guards r < 4, but another load uses r without a gate -> no shrink
     r = Range(0, 204)
     load1 = get_gated_load_uop(r < UOp.const(4), r)
-    load2 = UOp.param(1, dtypes.float, (204,)).index(r).load()
+    load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),))
     ranges = self.get_ranges(UOp.sink(load1, load2))
-    self.assert_range_end(ranges, 204)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 204)
 
   def test_range_no_shrink_when_used_in_reduce(self):
     # range used in both a gated load AND directly in the reduce expression -> no shrink
@@ -561,7 +564,8 @@ class TestRangeShrink(unittest.TestCase):
     gated_load = get_gated_load_uop(r < UOp.const(4), r)
     red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
     ranges = self.get_ranges(red.sink())
-    self.assert_range_end(ranges, 204)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 204)
 
   def test_range_shrink_to_single_iteration(self):
     # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
@@ -576,7 +580,8 @@ class TestRangeShrink(unittest.TestCase):
     r = Range(0, 204)
     x = (r < 4).where(UOp.const(1.0), Invalid)
     ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
-    self.assert_range_end(ranges, 4)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 4)
 
   def test_range_shrink_store_where_invalid_flipped(self):
     # above, but flipped
@@ -584,7 +589,8 @@ class TestRangeShrink(unittest.TestCase):
     r = Range(0, 204)
     x = (r < 4).where(UOp.const(1.0), Invalid)
     ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
-    self.assert_range_end(ranges, 4)
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].val, 4)
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_tensor.py

@@ -69,7 +69,7 @@ class TestIdxUpcast(unittest.TestCase):
     if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
       assert idx.op is Ops.INDEX
       idx_val = idx.src[1]
-      self.assertFalse(idx_val.overflows(idx_val.dtype))
+      self.assertFalse(idx_val.overflows(idx_val.dtype.scalar()))
 
   # use expand to generate kernel that uses large idx
   def do_op_then_assert(self, dtype: DType, dim1, dim2, dim3):
@@ -171,11 +171,6 @@ class TestTensorConstLike(unittest.TestCase):
     t = Tensor.ones(8, 4).shard(("NULL:0", "NULL:1"), axis=0)
     with self.assertRaises(RuntimeError): t.full_like(5, device="NULL")
 
-class TestTensorShape(unittest.TestCase):
-  def test_float_shape_raises(self):
-    for dim in (2.0, 2.5):
-      with self.subTest(dim=dim), self.assertRaisesRegex(RuntimeError, "shape must be int"): Tensor.ones(dim)
-
 class TestTensorDevice(unittest.TestCase):
   def test_create_from_single_device_tuple(self):
     (Tensor([1.0], device=(Device.DEFAULT,)) + Tensor([2.0])).realize()

test/null/test_uop_graph.py

@@ -1,9 +1,10 @@
 import unittest, pytest
 from tinygrad import dtypes, Variable, Device
 from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import DEBUG, Context
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes, KernelInfo
 from tinygrad.uop.symbolic import sym
-from test.helpers import full_rewrite, to_uops_list
+from test.helpers import to_uops_list
 from tinygrad.codegen import full_rewrite_to_sink
 
 simple_pm = PatternMatcher([
@@ -13,27 +14,43 @@ simple_pm = PatternMatcher([
   ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.val+c2.val)),
 ])
 
+def const_values(u:UOp):
+  if u.op is Ops.CONST: return (u.val,)
+  if u.op is Ops.STACK: return tuple(x.val for x in u.src)
+  raise AssertionError(f"expected const-like UOp, got {u.op}")
+
 class TestGraphRewriteConst(unittest.TestCase):
   def test_gep_const(self):
     v1 = UOp.const((0,1,2), dtypes.int)
     v2 = v1.index(1)
     ret = graph_rewrite(v2, sym)
-    self.assertIs(ret, UOp.const(1, dtypes.int))
+    self.assertEqual(ret.dtype, dtypes.int)
+    self.assertEqual(ret.val, 1)
 
   def test_add_const(self):
     v1 = UOp.const((0,1,2))
     v2 = UOp.const((5,6,7))
-    self.assertIs(graph_rewrite(v1+v2, sym), UOp.const((5,7,9)))
+    ret = graph_rewrite(v1+v2, sym)
+    self.assertEqual(ret.op, Ops.STACK)
+    self.assertEqual(const_values(ret), (5,7,9))
+
+  def test_add_const_lose_v(self):
+    v1 = UOp.const((0,1,2))
+    v2 = UOp.const((2,1,0))
+    ret = graph_rewrite(v1+v2, sym)
+    self.assertEqual(ret.op, Ops.STACK)
+    self.assertEqual(const_values(ret), (2,2,2))
 
 def xfail_broken_const_wraparound(fn):
   fn = pytest.mark.xfail(reason="const folding does not properly implement modular arithmetic")(fn)
   return unittest.expectedFailure(fn)
 class TestModularWraparound(unittest.TestCase):
   def _test(self, uop:UOp, expected:int):
-    result = uop.simplify()
-    self.assertEqual(result.op, Ops.CONST)
-    self.assertEqual(result.dtype, uop.dtype)
-    self.assertEqual(result.val, expected)
+    results = to_uops_list([uop])
+    self.assertEqual(len(results), 2)  # +1 for SINK
+    self.assertEqual(results[0].op, Ops.CONST)
+    self.assertEqual(results[0].dtype, uop.dtype)
+    self.assertEqual(results[0].val, expected)
 
   @xfail_broken_const_wraparound
   def test_cast(self):
@@ -174,25 +191,63 @@ class TestGraphRewrite(unittest.TestCase):
       self.assertEqual(len([x for x in sink.toposort() if x.op is Ops.CONST]), 1)
 
 class TestUOpGraph(unittest.TestCase):
+  def test_add_constant_fold(self):
+    c1 = UOp.const(1.0, dtypes.float)
+    c2 = UOp.const(2.0, dtypes.float)
+    out = c1+c2
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.val, 3.0)
+
   def test_where_same_fold(self):
     v = UOp.variable('tmp', 0, 1)
     c0 = UOp.const(0)
     vc = v != c0
     c1 = UOp.const(1.0, dtypes.float)
     out = vc.where(c1, c1)
-    self.assertIs(out.simplify(), c1)
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.val, 1.0)
 
   def test_where_const_fold(self):
     bf = UOp.const(False)
     c1 = UOp.const(1.0, dtypes.float)
     c2 = UOp.const(2.0, dtypes.float)
     out = bf.where(c1, c2)
-    self.assertIs(out.simplify(), c2)
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.val, 2.0)
 
   def test_const_cast(self):
     bf = UOp.const(False)
     out = bf.cast(dtypes.int)
-    self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.val, 0)
+
+  def test_const_bitcast(self):
+    bf = UOp.const(1.0, dtypes.float)
+    out = bf.bitcast(dtypes.uint32)
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
+    out = uops[-2]
+    self.assertEqual(out.op, Ops.CONST)
+    self.assertEqual(out.val, 0x3F800000)
+
+  @unittest.expectedFailure
+  def test_const_shape_change_bitcast(self):
+    bf = UOp.const(0x3F).cast(dtypes.uint8)
+    out = bf.bitcast(dtypes.half)
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 2)  # +1 for SINK
 
   def test_devectorize_derives_lane_dtype(self):
     from tinygrad.codegen import do_devectorize
@@ -202,11 +257,66 @@ class TestUOpGraph(unittest.TestCase):
     invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL)
     self.assertIs(invalid_lane_mul.dtype, dtypes.bool)
 
+  @unittest.skip("this test isn't valid uops")
+  def test_noop_vectorize_fold(self):
+    d0 = UOp.param(0, dtypes.float, (1,))
+    idx = UOp.const(0)
+    ld = d0.load(idx, dtype=dtypes.float)
+    vec = UOp(Ops.STACK, dtypes.float, (ld,))
+    x = vec.index(0)
+    alu = UOp(Ops.SQRT, src=(x, ))
+    out = UOp(Ops.STORE, src=(d0, idx, alu))
+    uops = to_uops_list([out])
+    self.assertEqual(len([x for x in uops if x.op is Ops.STACK]), 0)
+
+  @unittest.skip("this test isn't valid uops")
+  def test_gep_vec_fold(self):
+    d0 = UOp.param(0, dtypes.float, (1,))
+    d1 = UOp.param(1, dtypes.float, (1,))
+    d2 = UOp.param(2, dtypes.float, (1,))
+    idx = UOp.const(0)
+    def _test_vec(geps, count=4):
+      vec = UOp(Ops.STACK, dtypes.float, geps)
+      out = d0.index(idx).store(vec)
+      uops = to_uops_list([out])
+      if DEBUG >= 4:
+        from tinygrad import Device
+        print(Device[Device.DEFAULT].renderer.render(uops))
+      return uops[-2].src[-1]  # -2 to skip SINK
+
+    # possible
+    val = d1.index(idx).load(dtype=dtypes.float)
+    xyzw = tuple(val.index(i) for i in range(4))
+    self.assertIs(_test_vec(xyzw).op, Ops.LOAD)
+
+    # unaligned
+    val = d1.index(idx).load(dtype=dtypes.float)
+    wzyx = tuple(val.index(i) for i in reversed(range(4)))
+    self.assertIs(_test_vec(wzyx).op, Ops.STACK)
+
+    # different_size
+    val = d1.index(idx).load(dtype=dtypes.float)
+    xy = tuple(val.index(i) for i in range(2))
+    self.assertIs(_test_vec(xy+xy).op, Ops.STACK)
+    val = d1.index(idx).load(dtype=dtypes.float)
+    xy = tuple(val.index(i) for i in range(2))
+    self.assertIs(_test_vec(xy, count=2).op, Ops.STACK)
+
+    # different vals
+    val1 = d1.index(idx).load(dtype=dtypes.float)
+    val2 = d2.index(idx).load(dtype=dtypes.float)
+    xy1 = tuple(val1.index(i) for i in range(2))
+    xy2 = tuple(val2.index(i) for i in range(2))
+    self.assertIs(_test_vec(xy1+xy2).op, Ops.STACK)
+
   def test_gep_vec_const_fold(self):
     for vec_size in [2, 4, 8]:
       consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
-      vec = UOp.stack(*consts)
-      for i, const in enumerate(consts): self.assertIs(vec.index(i), const)
+      vec = UOp(Ops.STACK, src=tuple(consts))
+      with Context(SPEC=0):
+        uops = to_uops_list([vec.index(i) for i in range(vec_size)])
+        for uop, const in zip(uops, consts):
+          self.assertEqual(uop, const)
 
   def test_cast_alu_fold(self):
     d0 = UOp.param(0, dtypes.bool, (1,))
@@ -216,7 +326,7 @@ class TestUOpGraph(unittest.TestCase):
     alu = (ld<1).cast(dtypes.bool)
     out = d0.index(idx).store(alu)
     uops = to_uops_list([out])
-    self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 0)
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
 
   def test_double_cast_fold(self):
     d0 = UOp.param(0, dtypes.float, (1,))
@@ -226,7 +336,7 @@ class TestUOpGraph(unittest.TestCase):
     alu = ld.cast(dtypes.float).cast(dtypes.float)
     out = d0.index(idx).store(alu)
     uops = to_uops_list([out])
-    self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 1)
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
 
   def test_depth_2_const_fold(self):
     v = UOp.variable("tmp", 0, 1, dtypes.int, param=True)
@@ -234,7 +344,12 @@ class TestUOpGraph(unittest.TestCase):
     c4 = UOp.const(4, dtypes.int)
     vc = v+c2
     out = vc+c4
-    self.assertIs(out.simplify(), (v+UOp.const(6, dtypes.int)).simplify())
+    uops = to_uops_list([out])
+    self.assertEqual(len(uops), 5)  # +1 for SINK, +1 for the PARAM shape STACK
+    out = uops[-2]  # -2 to skip SINK
+    self.assertEqual(out.op, Ops.ADD)
+    self.assertEqual(out.src[1].op, Ops.CONST)
+    self.assertEqual(out.src[1].val, 6)
 
   def test_bitcast_to_same_dtype_fold(self):
     for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
@@ -245,8 +360,9 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_sub_with_cast_folds(self):
     a = Variable("a", 0, 5)
-    out = a.cast(dtypes.int)+(-a).cast(dtypes.int)
-    self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])
+    uops = to_uops_list([a.cast(dtypes.int)+(-a).cast(dtypes.int)])
+    assert uops[0] == UOp.const(0, dtypes.int)
+    assert uops[-1].op == Ops.SINK
 
   def test_where_on_gated_load_fold(self):
     ridx0 = UOp.range(100, 0)
@@ -255,10 +371,9 @@ class TestUOpGraph(unittest.TestCase):
     w = (ridx0<50).where(ld, 5)
     out = UOp.param(1, dtypes.long, (100,))
     uops = to_uops_list([out.index(ridx0).store(w)])
-    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: self.assertIs(u.src[1], expected)
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val==5
 
   def test_where_on_gated_load_folds_swapped_branches(self):
     ridx0 = UOp.range(100, 0)
@@ -266,10 +381,9 @@ class TestUOpGraph(unittest.TestCase):
     ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
     w = (ridx0<50).where(5, ld)
     uops = to_uops_list([w])
-    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD: self.assertIs(u.src[1], expected)
+      if u.op is Ops.LOAD: assert u.src[1].val==5
 
   def test_where_on_gated_load_with_cast(self):
     ridx0 = UOp.range(100, 0)
@@ -279,10 +393,9 @@ class TestUOpGraph(unittest.TestCase):
     w = (ridx0<50).where(ld, 5.0)
     out = UOp.param(1, dtypes.float, (100,))
     uops = to_uops_list([out.index(ridx0).store(w)])
-    expected = full_rewrite(UOp.const(5, dtypes.int).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: self.assertIs(u.src[1], expected)
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val == 5
 
   def test_where_on_casted_gated_load_extra_cond(self):
     ridx0 = UOp.range(100, 0)
@@ -312,10 +425,9 @@ class TestUOpGraph(unittest.TestCase):
     val = (ridx0<50).where(5, ld)
     st = idx.store(val).end(ridx0)
     uops = to_uops_list([st])
-    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.STORE: self.assertIs(u.src[1], expected)
+      if u.op is Ops.STORE: assert u.src[1].val==5
 
   def test_load_idx_becomes_int(self):
     # mnist indexing with split reduceop
@@ -389,6 +501,13 @@ class TestUOpGraph(unittest.TestCase):
     # only the second store happens
     self.assertEqual(len([u for u in uops if u.op is Ops.STORE]), 1)
 
+  @unittest.skip("this is a uop type error")
+  def test_asserts_bad_gate(self):
+    glbl0 = UOp.param(0, dtypes.int, (1,))
+    idx = UOp.const(0)
+    bad_gate = UOp.const(1)
+    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(42), bad_gate))])
+
   def test_after_end(self):
     r = UOp.range(10, 0)
 
@@ -456,7 +575,7 @@ class TestConstBufferize(unittest.TestCase):
     from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
     c = UOp.const(42.0)
     r1 = UOp.range(3, 0)
-    bufferize_with_range = c.bufferize(r1, arg=BufferizeOpts(device="CPU"))
+    bufferize_with_range = UOp(Ops.STAGE, src=(c, r1), arg=BufferizeOpts(device="CPU"))
     self.assertEqual(len(bufferize_with_range.src), 2)  # const + 1 range
 
     result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test')
@@ -471,7 +590,7 @@ class TestConstBufferize(unittest.TestCase):
     c = UOp.const(3.14)
     r1 = UOp.range(3, 0)
     r2 = UOp.range(4, 1)
-    bufferize_with_ranges = c.bufferize(r1, r2, arg=BufferizeOpts(device="CPU"))
+    bufferize_with_ranges = UOp(Ops.STAGE, src=(c, r1, r2), arg=BufferizeOpts(device="CPU"))
     self.assertEqual(len(bufferize_with_ranges.src), 3)  # const + 2 ranges
 
     result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test')

test/null/test_uop_symbolic.py

@@ -5,8 +5,7 @@ import z3
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
-from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
-from tinygrad.uop.weak import pm_cast_weak
+from tinygrad.uop.symbolic import sym, pm_fold_cast_const, commutative, pm_simplify_valid, pm_move_where_on_load
 from tinygrad.uop.validate import uops_to_z3
 
 def check_uop_against_string(self, v:UOp, s:str):
@@ -36,7 +35,7 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original")
 
   def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    v_simplified = graph_rewrite(v, sym+pm_cast_weak, name="simplify symbolic uop")
+    v_simplified = graph_rewrite(v, sym+pm_fold_cast_const, name="simplify symbolic uop")
     if test_z3: self.check_equal_z3(v, v_simplified)
     nmin, nmax = v_simplified.vmin, v_simplified.vmax
     check_uop_against_string(self, v_simplified, s)
@@ -1013,13 +1012,6 @@ class TestSymbolic(unittest.TestCase):
     b = Variable("b", 0, 3)
     self.helper_test_variable(-a<-b, False, True, "(b<a)")
 
-  def test_where_cast(self):
-    cond = Variable("s", 0, 3, dtypes.int) < 2
-    a = Variable("a", 0, 3, dtypes.int)
-    self.assertIs(graph_rewrite(cond.where(a, a+1).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), (a+1).cast(dtypes.half)))
-    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.const(2, dtypes.half)))
-    self.assertIs(graph_rewrite(cond.where(a, UOp.invalid()).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.invalid()))
-
   def test_where_merge_branches(self):
     cond1 = Variable("s", 0, 10) < 6
     cond2 = Variable("s", 0, 10) > 2

test/null/test_uops.py

@@ -41,6 +41,11 @@ class TestDTypeFromUOp(unittest.TestCase):
     # an explicit (strong) const dtype is legal until the field is removed
     self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)
 
+  def test_weak_dtype_rejected_by_program_spec(self):
+    for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)):
+      with self.assertRaises(RuntimeError): type_verify(UOp.const(value, weak).sink(), spec_program)
+      type_verify(UOp.const(value, concrete).sink(), spec_program)
+
   def test_invalid_stated_dtype(self):
     # UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not,
     # and the spec is what rejects a non-bool Invalid
@@ -129,7 +134,7 @@ class TestConstFloatEq(unittest.TestCase):
     self.assertFalse(Invalid != HoldsInvalid())
 
   def test_matchers_agree_on_nan(self):
-    n = UOp.const(math.nan)
+    n = UOp.const(math.nan, dtypes.float32)
     for compiled in (False, True):
       pm = PatternMatcher([(UPat(Ops.CONST, arg=math.nan), lambda: True)], compiled=compiled)
       self.assertTrue(pm.rewrite(n), f"{compiled=}")
@@ -343,9 +348,10 @@ class TestFastIdiv(unittest.TestCase):
   def test_fast_idiv_remove_powers_of_two(self):
     ridx = UOp.range(2**20, 0)
     uops = to_uops_list([ridx//(7*64)], ren=Device[Device.DEFAULT].renderer)
+    ops = [x.op for x in uops]
     # this requires shifting out the powers of two before doing fast_idiv
     # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
-    self.assertNotIn(dtypes.long, [x.dtype for x in uops])
+    self.assertNotIn(Ops.CAST, ops)
 
   @unittest.expectedFailure
   def test_fast_idiv_overflow(self):

test/null/test_viz.py

@@ -236,8 +236,8 @@ class TestViz(unittest.TestCase):
   def test_const_node_visibility(self):
     with save_viz() as viz:
       a = UOp.variable("a", 0, 10, dtype=dtypes.int)
-      z = UOp.const(0)
-      y = UOp.const(math.pi)
+      z = UOp.const(0, a.dtype)
+      y = UOp.const(math.pi, dtypes.float)
       alu = a*z
       ret = exec_rewrite(sink:=UOp.sink(alu, y), [sym])
     lst = viz.list_items()
@@ -249,7 +249,7 @@ class TestViz(unittest.TestCase):
     self.assertTrue(graphs[0][id(y)]["exclude"])
     self.assertFalse(graphs[0][id(alu)]["exclude"])
     self.assertEqual(graphs[0][id(y)]["label"].split("\n")[:2], ["CONST", "3.14159"])
-    self.assertEqual(list(graphs[1]), [id(u) for u in ret.toposort()])  # rewrite graph keys follow the rewritten sink's toposort
+    self.assertEqual(list(graphs[1]), [id(z), id(y), id(ret)])
 
   def test_const_reshape_expand_folded(self):
     # CONST->EXPAND should be folded into the ALU node, not shown as separate EXPAND nodes

test/opt/test_gen_float4.py

@@ -10,12 +10,12 @@ from test.helpers import replace_opts
 class TestFloat4(unittest.TestCase):
   @staticmethod
   def count_float4(uops: list[UOp], n=4):
-    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.float and uop.shape == (4,)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.float and uop.shape == (4,)]))
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.float and uop.shape == (4,)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.float and uop.shape == (4,)]))
   @staticmethod
   def count_half4(uops: list[UOp]):
-    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.half and uop.shape == (4,)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.half and uop.shape == (4,)]))
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.half and uop.shape == (4,)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.half and uop.shape == (4,)]))
 
   def test_float4_basic(self):
     a = Tensor.empty(2, 8).realize()

test/unit/test_allreduce.py

@@ -2,7 +2,6 @@ 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):
@@ -14,7 +13,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
-      if len(pairs) != N*(N-1)*2: raise KernelCountException(N*(N-1)*2, len(pairs))
+      self.assertEqual(len(pairs), N*(N-1)*2)
       # copy topology forms a ring
       self.assertEqual(len(set(pairs)), N)
 
@@ -26,8 +25,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]
-      if len(copies) != 24: raise KernelCountException(24, len(copies))
-      if len(sinks) != 26: raise KernelCountException(26, len(sinks))
+      self.assertEqual(len(copies), 24)
+      self.assertEqual(len(sinks), 26)
 
   @Context(RING=0, ALL2ALL=0)
   def test_schedule_naive(self):
@@ -40,8 +39,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]
 
-    if len(pairs) != N*(N-1): raise KernelCountException(N*(N-1), len(pairs))
-    if len(sinks) != 2: raise KernelCountException(2, len(sinks))
+    self.assertEqual(len(pairs), N*(N-1))
+    self.assertEqual(len(sinks), 2)
     self.assertTrue(all(dst != src for dst, src in pairs))
 
   def test_symbolic_shape(self):
@@ -65,7 +64,7 @@ class TestAllreduceCast(unittest.TestCase):
     with Context(ALLREDUCE_CAST=allreduce_cast, RING=0, SCACHE=0):
       t = Tensor.empty(4, 4, dtype=dtype).shard(ds, axis=0)
       linear = t.sum(0).linear_with_vars()[0]
-      return {si.src[1].buffer.dtype for si in linear.src if si.src[0].op is Ops.COPY}
+      return {si.src[1].buffer.dtype.scalar() for si in linear.src if si.src[0].op is Ops.COPY}
 
   def test_allreduce_cast_bf16(self):
     # with ALLREDUCE_CAST, allreduce copies stay in bfloat16 instead of promoting to float32

test/unit/test_attention.py

@@ -5,6 +5,8 @@ from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
   apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
 )
+from tinygrad.llm.kernels import Linear, gated_delta_prefill
+from tinygrad.llm.gguf import ggml_data_to_tensor
 
 def apply_rope(x:Tensor, start_pos:int):
   B, H, T, Hd = x.shape
@@ -12,6 +14,15 @@ def apply_rope(x:Tensor, start_pos:int):
   freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
   return apply_rope_new(x, freqs_cis)
 
+class TestLinear(unittest.TestCase):
+  def test_recovers_packed_ggml_weight(self):
+    for ggml_type,packed_size,words in ((13, 176, 44), (14, 210, 210), (23, 136, 34)):
+      packed = Tensor.empty(packed_size+4, dtype=dtypes.uint8, device="CPU")[4:]
+      decoded = ggml_data_to_tensor(packed, 256, ggml_type).reshape(1, 256)
+      linear = Linear(256, 1, bias=False)
+      linear.set_quantized(decoded)
+      self.assertEqual((linear.ggml_type, linear.weight.numel()), (ggml_type, words))
+
 class TestAttention(unittest.TestCase):
   def test_apply_rope(self):
     x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
@@ -41,6 +52,22 @@ class TestAttention(unittest.TestCase):
     np.testing.assert_allclose(block.cache_kv[0, :, :, :seqlen, :].numpy(), expected.numpy(), rtol=1e-5, atol=1e-5)
 
 class TestGatedDeltaNetBlock(unittest.TestCase):
+  def test_gated_delta_rectangular_state_and_row_decay(self):
+    rng = np.random.default_rng(42)
+    q, k = (rng.normal(size=(1, 1, 3, 32)).astype(np.float32) for _ in range(2))
+    v, beta = rng.normal(size=(1, 1, 3, 4)).astype(np.float32), rng.uniform(size=(1, 1, 3)).astype(np.float32)
+    alpha, initial = rng.uniform(0.8, 1, size=(1, 1, 3, 4)).astype(np.float32), rng.normal(size=(1, 1, 4, 32)).astype(np.float32)
+    expected_state, expected_out = initial.copy(), np.empty_like(v)
+    for t in range(3):
+      previous, av = expected_state.copy(), alpha[:, :, t, :, None]
+      delta = (v[:, :, t] - (previous*k[:, :, t, None]).sum(-1)*alpha[:, :, t]) * beta[:, :, t, None]
+      expected_state = previous*av + delta[..., None]*k[:, :, t, None, :]
+      expected_out[:, :, t] = (previous*q[:, :, t, None]).sum(-1)*alpha[:, :, t] + delta*(q[:, :, t]*k[:, :, t]).sum(-1)
+    state = Tensor(initial).contiguous().realize()
+    out = gated_delta_prefill(Tensor(q), Tensor(k), Tensor(v), Tensor(beta), Tensor(alpha), state).realize()
+    np.testing.assert_allclose(out.numpy(), expected_out, rtol=1e-4, atol=1e-4)
+    np.testing.assert_allclose(state.numpy(), expected_state, rtol=1e-4, atol=1e-4)
+
   def _tensor_linspace(self, start:float, stop:float, shape:tuple[int, ...]) -> Tensor:
     return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)
 
@@ -215,29 +242,6 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     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)

test/unit/test_dtype_weak.py

@@ -3,12 +3,11 @@ 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, GroupOp, dtype_from_uop, graph_rewrite
+from tinygrad.uop.ops import UOp, Ops, 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):
@@ -77,11 +76,6 @@ class TestWeakPromotion(unittest.TestCase):
       committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={})
     self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32])
 
-  def test_div_sub_operand_kept_weak(self):
-    a = Tensor.empty(4, dtype=dtypes.float32)
-    for t in (a / 1, a - 0):
-      self.assertEqual(t.uop.src[1].dtype, dtypes.weakfloat)
-
   def test_cast_weak_expression_commits_at_cast_floor(self):
     # the floor never narrows: a cast BELOW the default does not pull the compute width down with it
     with Context(DEFAULT_FLOAT=dtypes.float32):
@@ -94,13 +88,6 @@ class TestWeakPromotion(unittest.TestCase):
       out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom
       self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10)
 
-  def test_stacked_weak_casts_convert_each_kind(self):
-    # each weak cast is a kind conversion: weakint truncates before weakfloat re-lifts (neither is only a marker)
-    x = Tensor([2.5, -3.7], dtype=dtypes.float32, device="CPU")
-    stacked = x.cast(dtypes.weakint).cast(dtypes.weakfloat)
-    self.assertIs(stacked.dtype, dtypes.weakfloat)
-    self.assertEqual(stacked.tolist(), [2.0, -3.0])
-
   def test_uop_scalar_const_lifts_kind(self):
     for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
                                                  (dtypes.int32, 1, dtypes.int32, dtypes.weakint),
@@ -289,18 +276,5 @@ 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, KernelCountException
+from test.helpers import assert_kernel_count
 
 class TestFunction(unittest.TestCase):
   def test_simple(self):
@@ -516,7 +516,7 @@ class TestFunctionTuple(unittest.TestCase):
     Tensor.realize(a)
     c = f(a)
 
-    if count_kernels(c) != 1: raise KernelCountException(1, count_kernels(c))
+    self.assertEqual(count_kernels(c), 1)
 
     c.sum().backward()
     Tensor.realize(a.grad)

test/unit/test_gradient.py

@@ -51,10 +51,6 @@ class TestTensorGradient(unittest.TestCase):
     with self.assertRaises(RuntimeError): x.sum().gradient(x)
     with self.assertRaises(RuntimeError): x.float().sum().gradient(x)
 
-  def test_const_target_raise(self):
-    t = Tensor(2.0)
-    with self.assertRaises(RuntimeError): (t * 2.0).gradient(t)
-
   def test_copy_to_device_gradient(self):
     t = Tensor([1.0, 2, 3]).realize()
     t.to("CPU:1").square().sum().backward()
@@ -104,7 +100,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_where_gradient(self):
     # WHERE with a bare ()-shape branch: the scalar's gradient counts the positions where it is selected
-    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0, dtype=dtypes.float32)
+    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0)
     dw = Tensor(cond.uop.alu(Ops.WHERE, x.uop, w.uop)).sum().gradient(w)[0]
     self.assertEqual(dw.shape, ())
     self.assertEqual(dw.item(), 1.0)
@@ -113,7 +109,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_alu_gradient(self):
     # MUL with a bare ()-shape src, no EXPAND in the graph
-    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0, dtype=dtypes.float32)
+    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0)
     m = x.uop.alu(Ops.MUL, w.uop)
     self.assertIs(m.src[1], w.uop)
     dw = Tensor(m).sum().gradient(w)[0]
@@ -122,7 +118,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_intermediate_accumulation(self):
     # s is used directly and through an implicit broadcast edge, each edge's gradient reduces to s's shape before they sum
-    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5, dtype=dtypes.float32)
+    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5)
     s = p.sin()
     z = Tensor(x.uop.alu(Ops.MUL, s.uop)).sum() + s
     dp = z.gradient(p)[0]

test/unit/test_llm_amd.py

@@ -0,0 +1,91 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor, UOp, dtypes, nn
+from tinygrad.llm.kernels import Linear, amd_custom_kernels_supported
+from tinygrad.llm.kernels.amd import q8_quantize, quantized_attention
+from tinygrad.llm.gguf import ggml_data_to_tensor
+
+class TestQ8Quantize(unittest.TestCase):
+  def test_word_quant_weights_use_typed_buffer_view(self):
+    for ggml_type, type_size in ((13, 176), (23, 136)):
+      with self.subTest(ggml_type=ggml_type):
+        raw = Tensor(np.zeros(type_size + 4, dtype=np.uint8), device="CPU").contiguous().realize()[4:]
+        decoded = ggml_data_to_tensor(raw, 256, ggml_type).reshape(1, 256)
+        linear = Linear(256, 1, bias=False)
+        linear.set_quantized(decoded)
+        self.assertEqual(linear.ggml_type, ggml_type)
+        self.assertEqual(linear.weight.dtype, dtypes.uint32)
+        self.assertEqual(linear.weight.nbytes(), type_size)
+        self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)
+
+  def test_values_and_scales(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    x = np.linspace(-3.1, 2.7, 64, dtype=np.float32).reshape(2, 32)
+    quant, scale = q8_quantize(Tensor(x), 2, 32)
+    scale_np = np.maximum(np.max(np.abs(x), axis=-1, keepdims=True) / 127, 1e-8)
+    expected = np.clip(np.rint(x / scale_np), -127, 127).astype(np.int8)
+    np.testing.assert_array_equal(quant.bitcast(dtypes.int8).reshape(2, 32).numpy(), expected)
+    np.testing.assert_allclose(scale.numpy(), scale_np, rtol=1e-6)
+
+  def test_q6_linear_compiles(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    packed = rng.integers(0, 256, 210, dtype=np.uint8)
+    packed[-2:] = np.array([0.01], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, 256, 14).reshape(1, 256)
+    linear = Linear(256, 1, bias=False)
+    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
+    self.assertTrue(np.isfinite(linear(Tensor.randn(1, 256)).realize().item()))
+    self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)
+
+  def test_q6_linear_multiple_tokens(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    in_features, blocks = 2048, 16*2048//256
+    packed = rng.integers(0, 256, blocks*210, dtype=np.uint8)
+    for i in range(blocks): packed[i*210+208:i*210+210] = np.array([0.01], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, 16*in_features, 14).reshape(16, in_features)
+    weight = decoded.numpy()
+    linear = Linear(in_features, 16, bias=False)
+    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
+    x = rng.normal(size=(3, in_features)).astype(np.float32)
+    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
+    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
+    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
+    self.assertEqual(linear.ggml_type, 14)
+
+    generic = Linear(in_features, 16, bias=False)
+    nn.state.load_state_dict(generic, {"weight":decoded}, verbose=False, realize=False)
+    generic(Tensor.randn(4, in_features)[:UOp.variable("tokens", 1, 4).bind(2)])
+    self.assertFalse(generic.use_custom_quant)
+    self.assertIsNone(generic.ggml_type)
+
+  def test_attention_uses_physical_cache_length(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    q, k, v = Tensor.zeros(1, 2, 1, 32), Tensor.randn(1, 1, 1, 32), Tensor.randn(1, 1, 1, 32)
+    cache = Tensor.empty(2, 1, 1, 256, 32, dtype=dtypes.int8).contiguous()
+    scale = Tensor.empty(2, 1, 1, 256, dtype=dtypes.float16).contiguous()
+    out = quantized_attention(q, Tensor.stack(k, v), cache, scale, 0).realize()
+    np.testing.assert_allclose(out.numpy(), v.expand(1, 2, 1, 32).numpy(), rtol=2e-2, atol=2e-2)
+
+  def test_prefill_attention_unaligned_start(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    start_pos = 1718
+    q = Tensor.zeros(1, 8, 32, 128)
+    old_kv = rng.normal(size=(2, 1, 1, start_pos, 128)).astype(np.float32)
+    new_kv = rng.normal(size=(2, 1, 1, 32, 128)).astype(np.float32)
+    cache = Tensor.empty(2, 1, 1, 2048, 128, dtype=dtypes.int8).contiguous()
+    scale = Tensor.zeros(2, 1, 1, 2048, dtype=dtypes.float16).contiguous()
+    old_scale = np.maximum(np.max(np.abs(old_kv), axis=-1, keepdims=True) / 127, 1e-8).astype(np.float16)
+    Tensor.realize(cache[:, :, :, :start_pos].assign(Tensor(np.rint(old_kv / old_scale).astype(np.int8))),
+                   scale[:, :, :, :start_pos].assign(Tensor(old_scale.squeeze(-1))))
+    out = quantized_attention(q, Tensor(new_kv), cache, scale, UOp.variable("start_pos", 0, 2047).bind(start_pos)).realize()
+    values = cache[1, 0, 0, :start_pos+32].numpy().astype(np.float32) * \
+      scale[1, 0, 0, :start_pos+32].numpy().astype(np.float32)[:, None]
+    expected = np.stack([values[:start_pos+i+1].mean(0) for i in range(32)])[None, None].repeat(8, axis=1)
+    np.testing.assert_allclose(out.numpy(), expected, rtol=2e-3, atol=2e-3)
+
+if __name__ == "__main__": unittest.main()

test/unit/test_llm_moe.py

@@ -2,7 +2,7 @@ import unittest
 import numpy as np
 from dataclasses import replace
 from tinygrad import Tensor
-from tinygrad.llm.model import ExpertGating, TransformerBlock, TransformerConfig
+from tinygrad.llm.model import TransformerBlock, TransformerConfig
 
 def _moe_config(dim=8, hidden=16, n_heads=2, num_experts=4, num_experts_per_tok=2):
   return TransformerConfig(
@@ -96,32 +96,5 @@ class TestMoEFeedForward(unittest.TestCase):
     expected = moe_expected + shared_expected
     np.testing.assert_allclose(out.numpy(), expected, rtol=1e-2)
 
-  def test_moe_feed_forward_gating_funcs(self):
-    dim, hidden, n_heads = 8, 16, 2
-    num_experts, k = 4, 2
-    logits = np.array([4.0, 3.0, 0.0, -1.0], dtype=np.float32)
-    def softmax(x):
-      probs = np.exp(x - x.max())
-      return probs / probs.sum()
-    for gating_func in ExpertGating:
-      for norm_topk_prob in (False, True):
-        block = TransformerBlock(replace(_moe_config(dim, hidden, n_heads, num_experts, k),
-                                         expert_gating_func=gating_func, norm_topk_prob=norm_topk_prob))
-        block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
-        block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
-        block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
-        block.ffn_gate_inp.weight = Tensor((logits / dim)[None, :].repeat(dim, 0).T)
-        out = block._feed_forward(Tensor.ones(1, 1, dim)).numpy()[0, 0, 0]
-
-        if gating_func == ExpertGating.SOFTMAX: selection_scores = softmax(logits)
-        elif gating_func == ExpertGating.SIGMOID: selection_scores = 1 / (1 + np.exp(-logits))
-        elif gating_func == ExpertGating.SOFTMAX_WEIGHT: selection_scores = logits
-        else: selection_scores = np.sqrt(np.logaddexp(0, logits))
-        sel = np.argsort(selection_scores)[-k:]
-        weights = softmax(logits[sel]) if gating_func == ExpertGating.SOFTMAX_WEIGHT else selection_scores[sel]
-        if norm_topk_prob: weights /= weights.sum()
-        expected = (weights * (sel + 1)).sum() / (1 + np.exp(-1))
-        np.testing.assert_allclose(out, expected, rtol=1e-3)
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_multitensor.py

@@ -390,12 +390,6 @@ class TestMultiTensor(unittest.TestCase):
     self.assertEqual(out.shape, (rows, 8))
     np.testing.assert_equal(out[:3].to(Device.DEFAULT).numpy(), np.ones((3, 8)))
 
-  def test_symbolic_broadcast_consumed(self):
-    rows = Variable("rows", 1, 4).bind(3)
-    out = (Tensor.ones(rows).to(devices_2) + 1).realize()
-    self.assertEqual(out.shape, (rows,))
-    np.testing.assert_equal(out[:3].to(Device.DEFAULT).numpy(), np.full(3, 2))
-
   def test_multitensor_jit_in_list(self):
     # test MULTI tensor inside a list container - exercises the container unpacking + MULTI unpacking
     @TinyJit

test/unit/test_schedule_cache.py

@@ -4,13 +4,8 @@ from tinygrad import Tensor, Variable, UOp, function
 from tinygrad.uop.ops import KernelInfo
 from tinygrad.schedule import schedule_cache
 
-def custom_add_kernel(A:UOp, B:UOp, num:int=0) -> UOp:
-  return A[0].set(B[0] + num).sink(arg=KernelInfo(f"custom_add_{num}"))
-
-def custom_add_backward(grad_output:UOp, _) -> tuple[None, UOp]:
-  grad = Tensor.invalids(*grad_output.shape, dtype=grad_output.dtype, device=grad_output.device)
-  grad = Tensor.custom_kernel(grad, Tensor(grad_output, device=grad_output.device), fxn=functools.partial(custom_add_kernel, num=0))[0]
-  return None, grad.uop
+def custom_set0_kernel(A:UOp, num:int) -> UOp:
+  return A[0].set(num).sink(arg=KernelInfo(f"custom_set0_{num}"))
 
 class TestScheduleCache(unittest.TestCase):
   def test_bound_variable_reuses_cache(self):
@@ -30,27 +25,27 @@ class TestScheduleCache(unittest.TestCase):
 
   def test_custom_kernel(self):
     for i in range(4):
-      a, b = Tensor.empty(1), Tensor.ones(1)
-      a = Tensor.custom_kernel(a, b, fxn=functools.partial(custom_add_kernel, num=i))[0]
+      a = Tensor.empty(1)
+      a = Tensor.custom_kernel(a, fxn=functools.partial(custom_set0_kernel, num=i))[0]
       a.realize()
-      self.assertEqual(a.item(), i+1)
+      self.assertEqual(a.item(), i)
 
   def test_same_custom_function_reuses_cache(self):
     schedule_cache.clear()
-    fxn = functools.partial(custom_add_kernel, num=10)
+    fxn = functools.partial(custom_set0_kernel, num=10)
 
     # first run
-    a, x = Tensor.empty(1), Tensor.ones(1)
-    a = Tensor.custom_kernel(a, x, fxn=fxn)[0]
+    a = Tensor.empty(1)
+    a = Tensor.custom_kernel(a, fxn=fxn)[0]
     a.realize()
-    self.assertEqual(a.item(), 11)
+    self.assertEqual(a.item(), 10)
     cache_size_after_first = len(schedule_cache)
 
     # second run with same function should reuse cache
-    b, x = Tensor.empty(1), Tensor.ones(1)
-    b = Tensor.custom_kernel(b, x, fxn=fxn)[0]
+    b = Tensor.empty(1)
+    b = Tensor.custom_kernel(b, fxn=fxn)[0]
     b.realize()
-    self.assertEqual(b.item(), 11)
+    self.assertEqual(b.item(), 10)
     self.assertEqual(len(schedule_cache), cache_size_after_first)
 
   def test_simple(self):
@@ -70,28 +65,23 @@ class TestScheduleCache(unittest.TestCase):
       print(num)
     self.assertEqual(len(schedule_cache), start_len_schedule_cache)
 
+  @unittest.expectedFailure
   def test_simple_precompile(self):
-    @function(precompile=True, precompile_backward=True)
+    @function(precompile=True)
     def f(x:Tensor) -> Tensor:
       out = Tensor.invalids(*x.shape, dtype=x.dtype, device=x.device)
-      out = Tensor.custom_kernel(out, x, fxn=functools.partial(custom_add_kernel, num=10), grad_fxn=custom_add_backward)[0]
+      out = Tensor.custom_kernel(out, fxn=functools.partial(custom_set0_kernel, num=10))[0]
       return out + x
 
     # warmup
     x = Tensor.ones(1).realize()
-    out = f(x)
-    out.backward(x)
-    self.assertEqual(out.item(), 12)
-    self.assertEqual(x.grad.item(), 2)
+    _ = f(x).realize()
 
     # use the cache next time function is called
     start_len_schedule_cache = len(schedule_cache)
     for _ in range(3):
-      x = Tensor.ones(1).realize()
-      out = f(x)
-      out.backward(x)
-      self.assertEqual(out.item(), 12)
-      self.assertEqual(x.grad.item(), 2)
+      num = f(x).realize()
+      self.assertEqual(num.item(), 11)
     self.assertEqual(len(schedule_cache), start_len_schedule_cache)
 
 if __name__ == "__main__":

tinygrad/codegen/__init__.py

@@ -12,7 +12,7 @@ from tinygrad.dtype import dtypes, AddrSpace
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
+from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_fold_cast_const, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.codegen.decomp.dtype import pm_dtype_decomps
 from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns
@@ -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")), name="x"),
-   lambda b,s,x: UOp.stack(*[x.replace(src=(b,u)) for u in s.src])),
+  (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])),
   # 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,10 +281,6 @@ 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.cconst(c.val, 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))
@@ -350,7 +346,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # lower index dtype
   # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
-  sink = graph_rewrite(sink, symbolic_simple+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
+  sink = graph_rewrite(sink, symbolic_simple+pm_fold_cast_const+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
 
   # final symbolic before decomp
   sink = graph_rewrite(sink, symbolic, name="final symbolic")
@@ -387,10 +383,6 @@ 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)
 
-  # spell every literal as a casted const CAST(dt, CONST(value))
-  # TODO: remove once consts are always weak
-  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)
 

tinygrad/codegen/decomp/dtype.py

@@ -33,8 +33,7 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
       return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0))
     case Ops.CAST if dt in dtypes.floats:
       small = (a1.eq(0) & (a0 >= 0)) | (a1.eq(-1) & (a0 < 0))
-      cdt = dt if dt == dtypes.float64 else dtypes.float32
-      return small.where(a0.cast(dt), ((a1.cast(cdt) * (2**32)) + a0.bitcast(dtypes.uint).cast(cdt)).cast(dt))
+      return small.where(a0.cast(dt), ((a1.cast(dtypes.float32) * (2**32)) + a0.bitcast(dtypes.uint).cast(dtypes.float32)).cast(dt))
     case Ops.CAST: return a0.bitcast(dtypes.uint).cast(dt)
     case Ops.BITCAST: return a0.bitcast(dt), a1.bitcast(dt)
     case Ops.SHL:

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, dtype=offsets[grp[0]][0].src[0].dtype)
+        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset)
         if op == Ops.STORE:
           datas = []
           for i,g in enumerate(grp):

tinygrad/codegen/late/regalloc.py

@@ -4,7 +4,7 @@ from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
 from tinygrad.renderer.isa import ISARenderer, Register, greg
 from tinygrad.dtype import dtypes
 
-PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
+PSEUDO_OPS = {Ops.CONST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
 
 class LinearScanRegallocContext:
   # returns the uop that defines the virtual register
@@ -52,7 +52,7 @@ class LinearScanRegallocContext:
         # the value of a BUFFER is its 64bit address, XMM registers need 16 bytes
         sz = 16 if v.cons[0].size == 16 else (8 if self.vdef(v).op is Ops.BUFFER else self.vdef(v).dtype.itemsize)
         offset = self.stack_size + (sz - self.stack_size % sz) % sz
-        self.spills[v] = UOp.cconst(offset, dtypes.int32)
+        self.spills[v] = UOp.const(offset, dtypes.int32)
         self.stack_size = offset + sz
       r = alloc(cons if cons is not None else v.cons, i)
       self.insert_before.setdefault(i, []).append((v, r))
@@ -84,7 +84,7 @@ class LinearScanRegallocContext:
 
       # allocate stack array
       if u.op is Ops.BUFFER:
-        self.locals[u] = UOp.cconst(self.stack_size, dtypes.int32)
+        self.locals[u] = UOp.const(self.stack_size, dtypes.int32)
         self.stack_size += u.max_numel() * u.dtype.itemsize
 
       # loop prologue, avoid loading inside the loop
@@ -125,7 +125,7 @@ def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
   # alloc/dealloc stack
   if ctx.stack_size > 0:
     sp = ctx.ren.stack_pointer()
-    offset = UOp.cconst(ctx.stack_size, sp.dtype)
+    offset = UOp.const(ctx.stack_size, sp.dtype)
     if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, src=(sp, offset), tag=sp.tag))] + before
     elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, src=(sp, offset), tag=sp.tag))]
 

tinygrad/device.py

@@ -1,8 +1,8 @@
 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, subprocess, struct
+from typing import Any, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
+import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal
 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
@@ -103,7 +103,7 @@ class Buffer:
   def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None,
                initial_value:bytes|pickle.PickleBuffer|None=None, uop_refcount=0, base:Buffer|None=None, offset:int=0, preallocate=False):
     assert isinstance(dtype, DType)
-    self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = Device.canonicalize(device), size, dtype, options, offset, 0
+    self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = device, size, dtype, options, offset, 0
     self._bufs: dict[str, Any] = {}
     if base is None:
       assert offset == 0, "base buffers can't have offset"
@@ -116,7 +116,7 @@ class Buffer:
         if isinstance(initial_value, pickle.PickleBuffer): initial_value.release()
     else:
       assert base._base is None, "base can't have a base"
-      assert self.device == base.device, "base must have the same device"
+      assert device == base.device, "base must have the same device"
       self._base = base
     if preallocate: self.allocate()
   @property
@@ -133,7 +133,7 @@ class Buffer:
   # check if the underlying buffer is allocated, possibly from the base object
   def is_allocated(self) -> bool: return self.base.is_allocated() if self._base is not None else self.device in self._bufs
   def get_buf(self, device: str) -> Any:
-    if device not in self._bufs and (device:=Device.canonicalize(device)) not in self._bufs:
+    if device not in self._bufs:
       allocator = Device[device].allocator
       if device == self.device: self.ensure_allocated()
       elif self._base is not None: self._bufs[device] = allocator._offset(self._base.get_buf(device), self.nbytes, self.offset)
@@ -310,14 +310,6 @@ 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:
@@ -339,18 +331,17 @@ class Program(Generic[DeviceType]):
                wait=False) -> float|None: pass
 
 class Compiled:
-  ifaces:list[Callable] = []
   profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.
 
-  has_copy_queue:bool = True
-
   pm_lower:Any = None
   pm_bufferize:Any = None
 
+  has_copy_queue:bool = True
+
   def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):
     from tinygrad.renderer import Renderer
     self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]
-    self.device_id, self.arch = (int(idx) if ":" in device and (idx:=device.split(":")[1]).isdigit() else 0), arch
+    self.arch = arch
     self.cached_renderer:dict[Any, Renderer] = {}
 
   @property
@@ -373,21 +364,11 @@ class Compiled:
     return select_first_inited(select_by_name(self.renderers, self._renderer_name, t.renderer, f"{self.device} has no renderer {t.renderer!r}"),
                                f"No renderer for {self.device} is available", self.cached_renderer, t)
 
-  def _select_iface(self, device:str):
-    self.device_id = int(device.split(":")[1]) if ":" in device else 0
-    assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "",  \
-      f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
-    t = DEV.target(dev:=type(self).__name__[:-6])
-    filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
-    filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fallback to mock ifaces
-    return select_first_inited([functools.partial(iface, self, self.device_id) for iface in filtered],
-                               f"No interface for {dev}:{self.device_id} is available")
-
   def count(self) -> int:
     """
     Returns the number of physical accelerators available to the runtime.
     """
-    return self.iface.count if hasattr(self, 'iface') else 1
+    return 1
 
   def synchronize(self):
     """
@@ -405,7 +386,7 @@ class Compiled:
     """
     Called at the end of process lifetime to allow the device to finalize.
     """
-    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
+    # override this in your device implementation
 
 if PROFILE:
   @atexit.register
@@ -427,7 +408,7 @@ def enumerate_devices_str() -> Generator[str, None, None]:
     ren_results, iface_results = [], []
     try:
       d = Device[device]
-      for iface in [i for i in d.ifaces if not i.__name__.startswith("MOCK")]:
+      for iface in [i for i in getattr(d, 'ifaces', []) if not i.__name__.startswith("MOCK")]:
         try:
           name = iface.__name__[:-5]
           default_text, count = ("(default)", d.count()) if type(d.iface) is iface else (f"(DEV={name}+{device} to make default)", iface(d, 0).count) # type: ignore

tinygrad/dtype.py

@@ -66,6 +66,7 @@ class DType(metaclass=DTypeMetaClass):
   def __reduce__(self): return type(self), tuple(getattr(self, f.name) for f in fields(self))
   def __repr__(self): return f"dtypes.{INVERSE_DTYPES_DICT[self.name]}"
   def __lt__(self, o:DType): return (self.priority, self.bitsize, self.name, self.fmt) < (o.priority, o.bitsize, o.name, o.fmt)
+  def scalar(self) -> DType: return self
   @functools.cached_property
   def min(self):
     if dtypes.is_int(self): return 0 if dtypes.is_unsigned(self) else -2**(self.bitsize-1)

tinygrad/engine/jit.py

@@ -269,14 +269,10 @@ class _TinyJit(Generic[ReturnType]):
         big_linear, onetime_linear = prune_linear(big_linear, set(input_buf_uops))
         if DEBUG >= 1: print(f"pruned from {len(big_linear.src) + len(onetime_linear.src)} -> {len(big_linear.src)} kernels")
         run_linear(onetime_linear, var_vals)
-        del onetime_linear
 
       # hold all buffers reachable from live Tensors (e.g. lazy .grad created during capture), the memory planner can't suballocate those
       held_bufs = set(buffers) | {u for tref in list(all_tensors) if (t:=tref()) is not None for u in t.uop.toposort() if u.op is Ops.BUFFER}
       linear = jit_lower(big_linear, held_bufs, input_buf_uops)
-      # drop the pre-planning graph: it keeps the whole capture-time working set allocated (big_linear) or referenced (held_bufs).
-      # the planned linear only uses the arena/held buffers, so the intermediates must be freed before linking and first exec
-      del big_linear, held_bufs
       self.captured = CapturedJit(ret, linear, names, expected_input_info)
       ret = self.captured(input_buf_uops, var_vals)
     elif self.cnt >= 2:

tinygrad/engine/realize.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
 from typing import cast, Iterator, Any, Sequence
-import random, itertools, math, weakref, array, decimal
+import time, random, itertools, math, contextlib, weakref, array
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
-from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
+from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
-from tinygrad.dtype import dtypes
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import args_from_ast
@@ -14,9 +13,6 @@ from tinygrad.codegen.opt.postrange import args_from_ast
 # **************** Helpers ****************
 
 def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call.src[1:] if not s.is_bound_var)
-def get_call_var_uops(call:UOp, prg:UOp) -> list[UOp]:
-  bound = {s.src[0].expr: s.src[1].src[1] for s in call.src[1:] if s.is_bound_var}
-  return [bound.get(v.expr, v) for v in prg.arg.vars]
 
 def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
   ast = call.src[0]
@@ -25,12 +21,6 @@ def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
   return (), ()
 
-def get_call_kernels(call:UOp) -> list[tuple[str, UOp]]:
-  if (ast:=call.src[0]).op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return [(d, k) for devs, k, _ in call.arg.aux.kernels for d in devs]
-  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return [(to_tuple(ast.device)[0], call)]
-  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "validate": return []
-  return [(d, call) for d in to_tuple(call.src[1].device)]
-
 def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|None=None) -> str:
   def _uop_sz_to_str(uop:UOp) -> str: return size_to_str(sym_infer(prod(uop.shape) * uop.dtype.itemsize, var_vals or {}))
   def _dev_str(buf:Buffer|UOp) -> str: return ', '.join(d[:7] for d in to_tuple(buf.device))
@@ -46,52 +36,49 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N
 # **************** Stat ****************
 
 def estimate_uop(call:UOp) -> Estimates:
-  if (ast:=call.src[0]).op is Ops.PROGRAM: return ast.src[0].arg.estimates or Estimates()
+  ast = call.src[0]
+  if ast.op is Ops.PROGRAM: return ast.src[0].arg.estimates or Estimates()
   if ast.op is Ops.COPY or (ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec"):
-    return Estimates(lds=(nbytes:=prod(call.src[1].shape) * call.src[1].dtype.itemsize), mem=nbytes)
+    nbytes = prod(call.src[1].shape) * call.src[1].dtype.itemsize
+    return Estimates(lds=nbytes, mem=nbytes)
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return get_graph_runtime(ast).estimates
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return call.arg.aux.estimates
   return Estimates()
 
 first_run_cache:set[bytes] = set()
-def track_stats(ctx:ExecContext, call:UOp, st:decimal.Decimal, ets:list[float|None]):
-  if ctx.update_stats:
-    is_hcq = (ast:=call.src[0]).op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq"
-    estimates, n = estimate_uop(call), 1 if is_hcq else len(get_call_kernels(call))
-    GlobalCounters.kernel_count += len(call.arg.aux.kernels) if is_hcq else n
-    GlobalCounters.global_ops += n*sym_infer(estimates.ops, ctx.var_vals)
-    GlobalCounters.global_mem += n*sym_infer(estimates.mem, ctx.var_vals)
-    GlobalCounters.time_sum_s += sum(et for et in ets if et is not None)
-  if DEBUG < 2 and not PROFILE: return
+@contextlib.contextmanager
+def track_stats(ctx:ExecContext, call:UOp, device:str, bufs:list[Buffer], var_vals:dict[str, int]):
+  if PROFILE:
+    outputs, inputs = get_call_outs_ins(call)
+    cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"var_vals": var_vals,
+      "bufs": [b.trace_num for b in bufs], "name": get_call_name(call, bufs, var_vals), "outputs": outputs, "inputs": inputs}))
+  et: list[float|None] = [None]
+  if DEBUG >= 2: st = time.perf_counter()
+  yield et
+  if not ctx.update_stats: return
 
-  kernels = get_call_kernels(call) # everything below is the per kernel display: exec events for the profiler and DEBUG=2 lines
-  args = resolve_params(call, ctx.input_uops) if kernels and kernels[0][1] is call else []
-  lanes = list(unwrap_multi(call, [args[g] for g in call.src[0].arg.globals] if call.src[0].op is Ops.PROGRAM else args)) if args else []
-  for i, (device, kcall) in enumerate(kernels):
-    et, bufs = ets[i] if i < len(ets) else None, lanes[i][0] if i < len(lanes) else []
-    if PROFILE: # backdate the event to the start of the call, the viz matches a device range with the exec event before it
-      outputs, inputs = get_call_outs_ins(kcall)
-      cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"var_vals": ctx.var_vals,
-        "bufs": [b.trace_num for b in bufs], "name": get_call_name(kcall, bufs, ctx.var_vals), "outputs": outputs, "inputs": inputs}, ts=st))
-    if DEBUG < 2 or not ctx.update_stats: continue
-    if et is None:
-      Device[device].synchronize()
-      et, st = float(perf_counter_us() - st)*1e-6, perf_counter_us()
-      GlobalCounters.time_sum_s += et
+  if DEBUG >= 2 and et[0] is None:
+    Device[device].synchronize()
+    et[0] = time.perf_counter() - st
 
-    estimates = estimate_uop(kcall)
-    display_name = get_call_name(kcall, bufs, ctx.var_vals)
-    op_est, mem_est, lds_est = (sym_infer(x, ctx.var_vals) for x in (estimates.ops, estimates.mem, estimates.lds))
-    header_color = 'magenta' if ctx.jit else ('green' if kcall.src[0].key not in first_run_cache else None)
-    ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
-    flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
+  estimates = estimate_uop(call)
+  GlobalCounters.kernel_count += 1
+  GlobalCounters.global_ops += (op_est:=sym_infer(estimates.ops, var_vals))
+  GlobalCounters.global_mem += (mem_est:=sym_infer(estimates.mem, var_vals))
+  if et[0] is not None: GlobalCounters.time_sum_s += et[0]
+  if DEBUG >= 2:
+    display_name = get_call_name(call, bufs, var_vals)
+    lds_est = sym_infer(estimates.lds, var_vals)
+    header_color = 'magenta' if ctx.jit else ('green' if call.src[0].key not in first_run_cache else None)
+    ptm = colored(time_to_str(et[0], w=9), "yellow" if et[0] > 0.01 else None) if et[0] is not None else ""
+    flops, membw, ldsbw = op_est/(et[0] or 1e-20), mem_est/(et[0] or 1e-20), lds_est/(et[0] or 1e-20)
     flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
     mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \
       colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green')
     print(f"{colored(f'*** {device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
       f" {display_name+' '*(46-ansilen(display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
-      ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})"))
-    first_run_cache.add(kcall.src[0].key)
+      ("" if et[0] is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})"))
+    first_run_cache.add(call.src[0].key)
 
 local_size_cache: dict[bytes, tuple[int, ...]] = {}
 def optimize_local_size(call:UOp, prg:UOp) -> UOp|None:
@@ -164,31 +151,32 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d
                    for x in call.src[0].toposort())
     for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {"_device_num": j} if has_dnum else {}
 
-def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
+def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
     dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated()
-    if hasattr(dest.allocator,'_transfer') and dest.allocator.supports_transfer and dest.device.split(":")[0] == src.device.split(":")[0]:
-      dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev)
-    elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \
-         and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
-      dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
-    elif hasattr(dest.allocator, '_as_buffer'): src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf)
-    else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True))
-  return []
+    with track_stats(ctx, call, dest.device, [dest, src], ctx.var_vals):
+      if hasattr(dest.allocator,'_transfer') and dest.allocator.supports_transfer and dest.device.split(":")[0] == src.device.split(":")[0]:
+        dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev)
+      elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \
+           and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
+        dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
+      elif hasattr(dest.allocator, '_as_buffer'): src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf)
+      else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True))
+  return None
 
-def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
-  ets:list[float|None] = []
-  resolved = resolve_params(call, ctx.input_uops)
-  for device, (bufs, device_vars) in zip(to_tuple(call.src[1].device), unwrap_multi(call, [resolved[i] for i in ast.arg.globals])):
+def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  et = None
+  for device, (bufs, device_vars) in zip(to_tuple(call.src[1].device), unwrap_multi(call, resolve_params(call, ctx.input_uops))):
     var_vals = {**ctx.var_vals, **device_vars}
-    prg_bufs = [b.ensure_allocated() for b in bufs]
+    prg_bufs = [bufs[i].ensure_allocated() for i in ast.arg.globals]
     rt = get_runtime(device, ast, cache=ctx.cache)
     global_size, local_size = ast.arg.launch_dims(var_vals)
-    ets.append(rt(*[b.get_buf(device) for b in prg_bufs], global_size=global_size, local_size=local_size, vals=ast.arg.vals(var_vals),
-                  wait=ctx.wait, timeout=ctx.timeout))
-  return ets
+    with track_stats(ctx, call, device, prg_bufs, var_vals) as tm:
+      et = tm[0] = rt(*[b.get_buf(device) for b in prg_bufs], global_size=global_size, local_size=local_size, vals=ast.arg.vals(var_vals),
+                       wait=ctx.wait, timeout=ctx.timeout)
+  return et
 
-def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
+def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   import numpy as np
   for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
     bufs, dev_bufs = bufs[:len(bufs)//2], bufs[len(bufs)//2:]
@@ -197,36 +185,45 @@ def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
     global_size, local_size = prg.arg.launch_dims(var_vals)
     cpu_rt(*[bufs[i].ensure_allocated()._buf for i in prg.arg.globals], global_size=global_size, local_size=local_size, vals=prg.arg.vals(var_vals))
     for i in prg.arg.outs: np.testing.assert_allclose(dev_bufs[i].ensure_allocated().numpy(), bufs[i].numpy(), rtol=1e-3, atol=1e-3)
-  return []
+  return None
 
-def exec_encdec(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
+def exec_encdec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   bufs = [cast(Buffer, b.buffer).ensure_allocated() for b in resolve_params(call, ctx.input_uops)]
   shape, pos_var = tuple(s.val for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr
-  bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var])
-  return []
+  with track_stats(ctx, call, bufs[0].device, bufs, ctx.var_vals):
+    bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var])
+  return None
 
-def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
-  return [get_graph_runtime(ast, ctx.input_uops)(ctx.input_uops, ctx.var_vals, wait=ctx.wait)]
+def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  rt = get_graph_runtime(ast, ctx.input_uops)
+  with track_stats(ctx, call, rt.device, [], ctx.var_vals) as t: t[0] = rt(ctx.input_uops, ctx.var_vals, wait=ctx.wait)
+  return t[0]
 
-def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
-  dev = cast(Any, Device[(info:= call.arg.aux).device[0]])
-  addrs = [(b.bufs[j] if isinstance(b:=_resolve(ctx.input_uops[k], ctx.input_uops).buffer, MultiBuffer) else b).get_buf(dev_name).va_addr
-           for devs, idxs in info.input_idxs for j, dev_name in enumerate(devs) for k in idxs]
-  dev.rt_buffer._buf.cpu_view().view(offset=(base:=dev.rt_allocator.alloc(len(addrs) * 8)), fmt='Q')[:len(addrs)] = array.array('Q', addrs)
+def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  if (info:=call.arg.aux).inputs is not None:
+    bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
+    table = call.src[1+info.inputs].buffer
+    for j,dev in enumerate(call.arg.aux.device):
+      addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs])
+      mv = (table.bufs[j] if isinstance(table, MultiBuffer) else table).ensure_allocated()._buf.cpu_view().view(fmt='Q')
+      wait_cond(lambda: mv[0], value=0, timeout_ms=ctx.timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000), msg=f"{dev} hang detected")
+      mv[:len(addrs)] = addrs
 
-  if info.inputs is not None:
-    tables = [UOp.from_buffer(dev.rt_buffer.view(len(idxs), dtypes.uint64, base + j*len(idxs)*8), HCQ_RUNTIME_DEV.value)
-              for devs, idxs in info.input_idxs for j in range(len(devs))]
-    call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*tables)})
-  exec_kernel(replace(ctx, var_vals={**ctx.var_vals, "hcq_inputs_ptr": dev.rt_buffer._buf.va_addr + base}), call, ast)
+  exec_kernel(replace(ctx, update_stats=DEBUG>=3), call, ast)
 
-  def _prof_tm(device:str, stat_call:UOp, prof:tuple[int, ...]) -> float|None:
-    (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, stat_call.arg.name, *prof)
-    if not ctx.wait: return None
-    d.synchronize(timeout=ctx.timeout)
-    st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
-    return float(en-st)/d.timestamp_divider/1e6
-  return [_prof_tm(device, k, prof) for devices, k, prof in info.kernels if prof for device in devices] if PROFILE or ctx.wait else []
+  tms = []
+  for devices,name,estimates,prof in info.kernels:
+    for device in devices:
+      tm = None
+      if prof:
+        (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, name, *prof)
+        if ctx.wait:
+          d.synchronize(timeout=ctx.timeout)
+          st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
+          tms.append(tm:=float(en-st)/d.timestamp_divider/1e6)
+      stat_call = call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates, kernels=())))
+      with track_stats(ctx, stat_call, device, [], ctx.var_vals) as et: et[0] = tm
+  return max(tms) if tms else None
 
 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -265,15 +262,14 @@ pm_exec = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
 ])
 
-if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV # noqa: E402 # down here, hcq2 imports realize
+if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link # noqa: E402 # down here, hcq2 imports the helpers above
 
 def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
   if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
   linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
-  linear = graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
   if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
-  return linear
+  return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
 
 def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear
 
@@ -281,7 +277,7 @@ def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequenc
   inputs = list(input_uops)
   if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs))
   ctx = ExecContext(var_vals or {}, tuple(inputs), update_stats, jit, wait or DEBUG>=2)
-  for call in linear.src: track_stats(ctx, call, perf_counter_us(), pm_exec.rewrite(call, ctx))
+  for call in linear.src: pm_exec.rewrite(call, ctx)
 
 def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> float:
   if clear_l2:
@@ -291,4 +287,4 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None
       with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
   ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
   linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), cache=ctx.cache)
-  return max(et for c in linear.src for et in pm_exec.rewrite(c, ctx) or [0.0])
+  return max(pm_exec.rewrite(c, ctx) or 0.0 for c in linear.src)

tinygrad/function.py

@@ -1,5 +1,4 @@
 import functools, time
-from dataclasses import replace
 from typing import Generic, TypeVar, Callable, cast, overload
 from tinygrad.helpers import Context, dedup, getenv, DEBUG
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
@@ -24,10 +23,6 @@ def invalid_outputs(uret:UOp) -> set[UOp]:
   return {u.src[0].buf_uop for u in uret.backward_slice_with_self
           if u.op is Ops.STORE and u.src[1].base.is_invalid and not u.src[0].buf_uop.is_realized}
 
-def renumber_invalid_outputs(uret:UOp) -> UOp:
-  return uret.substitute({b:b.replace(arg=replace(b.arg, slot=i))
-                          for i,b in enumerate(x for x in uret.toposort(enter_calls=False) if x in invalid_outputs(uret))})
-
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):
   depth = 0
@@ -70,7 +65,6 @@ class _function(Generic[ReturnType]):
     # the BUFFERs that are left are the implicit inputs
     num_explicit = len(call_uops)
     uret = graph_rewrite(uret, pm_ctx, (call_uops, invalid_outputs(uret)), bottom_up=True, name="get_implicit_inputs")
-    uret = renumber_invalid_outputs(uret)
     name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
     if not self.allow_implicit:
       implicit_buffers = [x for x in call_uops[num_explicit:] if x.op is Ops.BUFFER]

tinygrad/helpers.py

@@ -490,7 +490,7 @@ 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/0a6871b19abf5d6e024b5d208b101ae53e7fa0de/{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/kernels/__init__.py

@@ -0,0 +1,92 @@
+import functools
+from typing import cast
+from tinygrad import Tensor, UOp, nn, dtypes, Device, Context
+from tinygrad.device import Buffer
+from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import prod
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+
+def kernel_var(x:UOp) -> UOp:
+  # a Variable is a 0-d ALU BUFFER in the tensor graph; inside kernels it takes the ALU PARAM form (same name keeps the value binding)
+  return x.substitute({v: UOp.variable(v.expr, v.vmin, v.vmax, dtype=v.dtype, multiple_of=v.arg.multiple_of, param=True)
+                       for v in x.toposort() if v.is_variable})
+
+def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
+  # the custom kernels are tuned for RDNA3 (gfx11): the WMMA register layouts don't match gfx12 (RDNA4)
+  # or CDNA (MFMA-only, wave64), and the dp4a builtins and 32-lane wave ops aren't portable either.
+  if isinstance(device, tuple): device = device[0]
+  if device is None or device.split(":")[0] != "AMD": return False
+  # Device[...] trips ALLOW_DEVICE_USAGE=0 in function contexts, the device is always open here anyway
+  with Context(ALLOW_DEVICE_USAGE=1):
+    return (t:=getattr(Device[device], "target", None)) is not None and t[0] == 11
+
+class Linear(nn.Linear):
+  ggml_type:int|None = None
+  def __init__(self, in_features:int, out_features:int, bias=True):
+    super().__init__(in_features, out_features, bias)
+    self.in_features, self.out_features = in_features, out_features
+    self.use_custom_quant = True
+  def set_quantized(self, decoded:Tensor):
+    packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in ((13, 176), (14, 210), (23, 136))}
+    raw = next((u for u in decoded.uop.toposort() if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
+    if raw is None: return
+    raw_offset = raw.contiguous_view_offset()
+    assert raw_offset is not None and raw_offset % 4 == 0 and raw.buf_uop.dtype == dtypes.uint8
+    self.ggml_type = packed_sizes[prod(raw.shape)]
+    # Q5_K and IQ4_XS kernels consume words. Store a typed buffer view directly: a lazy BITCAST is decomposed into
+    # byte-combining ALU before custom-kernel scheduling and would copy the entire packed weight on every JIT graph.
+    packed_dtype = dtypes.uint8 if self.ggml_type == 14 else dtypes.uint32
+    self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer).view(raw.max_numel() * raw.dtype.itemsize // packed_dtype.itemsize,
+                                                                              packed_dtype, raw_offset)))
+  def __call__(self, x:Tensor) -> Tensor:
+    static = isinstance(x.numel(), int)
+    supported = self.use_custom_quant and amd_custom_kernels_supported(self.weight.device)
+    if self.ggml_type is None and not static: supported = self.use_custom_quant = False
+    if self.ggml_type is None and supported: self.set_quantized(self.weight)
+    if self.ggml_type in (13, 14, 23) and supported:
+      from tinygrad.llm.kernels.amd import q8_linear
+      return q8_linear(self, x)
+    return super().__call__(x)
+
+@functools.cache
+def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp, start_pos:UOp|None=None) -> UOp:
+  batch, heads, tokens, value_dim = cast(tuple[int, int, int, int], core.shape)
+  key_dim, alpha_dim = cast(int, q.shape[-1]), cast(int, alpha.shape[-1]) if len(alpha.shape) == 4 else 1
+  core, v = (x.reshape(batch*heads, tokens, value_dim) for x in (core, v))
+  q, k = (x.reshape(batch*heads, tokens, key_dim) for x in (q, k))
+  beta, kq = (x.reshape(batch*heads, tokens) for x in (beta, kq))
+  alpha, state = alpha.reshape(batch*heads, tokens, alpha_dim), state.reshape(batch*heads, value_dim, key_dim)
+  bh, row, cols = UOp.range(batch*heads, 0, AxisType.GLOBAL), UOp.range(value_dim, 2), tuple(range(key_dim))
+  current = UOp.placeholder((key_dim,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  initial = None if start_pos is None else start_pos.eq(0)
+  current = current.after(UOp.group(*(current[col].store(state[bh, row, col].float() if initial is None else
+    initial.where(0, state[bh, row, col].float())) for col in cols)))
+  token = UOp.range(tokens, 1, AxisType.REDUCE)
+  previous = tuple(current.after(token)[col].load() for col in cols)
+  keys, queries = (tuple(x[bh, token, col].load() for col in cols) for x in (k, q))
+  av, bv = alpha[bh, token, row if alpha_dim > 1 else 0].load(), beta[bh, token].load()
+  state_k = sum((x*y for x,y in zip(previous, keys)), UOp.const(0, dtypes.float32))
+  state_q = sum((x*y for x,y in zip(previous, queries)), UOp.const(0, dtypes.float32))
+  delta = (v[bh, token, row].load() - state_k*av) * bv
+  step = UOp.group(core[bh, token, row].store(state_q*av + delta*kq[bh, token]),
+                   *(current[col].store(x*av + delta*y) for col,x,y in zip(cols, previous, keys))).end(token)
+  stores = (state[bh, row, col].store(current.after(step)[col].load().cast(state.dtype)) for col in cols)
+  return UOp.group(*stores).end(row, bh).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
+
+def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, state:Tensor, start_pos:Tensor|None=None) -> Tensor:
+  batch, heads, tokens, key_dim = q.shape
+  value_dim = v.shape[-1]
+  assert q.shape == k.shape and v.shape[:3] == q.shape[:3] and beta.shape == (batch, heads, tokens)
+  assert alpha.shape in ((batch, heads, tokens), (batch, heads, tokens, value_dim))
+  assert state.shape == (batch, heads, value_dim, key_dim)
+  kernel = _gated_delta_prefill_kernel
+  if amd_custom_kernels_supported(q.device) and key_dim % 32 == 0 and value_dim % 4 == 0:
+    from tinygrad.llm.kernels.amd import _gated_delta_prefill_kernel as kernel
+  core, kq = Tensor.empty_like(v), (q*k).sum(-1).contiguous()
+  srcs = (core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq)
+  if start_pos is None: return Tensor.custom_kernel(*srcs, fxn=kernel)[0]
+  contig = tuple(x.uop if x.uop.op is Ops.AFTER else x.uop.contiguous() for x in srcs)
+  params = tuple(UOp.placeholder_like(x, slot=i) for i,x in enumerate(contig))
+  assert start_pos.uop.is_bound_var
+  call = kernel(*params, kernel_var(start_pos.uop.src[0])).call(*contig, start_pos.uop)
+  return Tensor(contig[0].after(call))

tinygrad/llm/kernels/amd.py

@@ -0,0 +1,451 @@
+from __future__ import annotations
+import functools, math
+from typing import Callable, cast
+from tinygrad import Tensor, UOp
+from tinygrad.llm.kernels import Linear, kernel_var
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve
+from tinygrad.dtype import AddrSpace, dtypes
+
+BLOCK_M, BLOCK_N, DECODE_HEAD_TILE, WARP_SIZE = 32, 32, 8, 32
+WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
+WAVES_M, WAVES_N, LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 2, 2, 16
+WMMA_ACC, THREADS_PER_BLOCK = WMMA_M // LANES_PER_WAVE_M, WARP_SIZE * WAVES_M * WAVES_N
+LDS_PAD, WMMA_ARG, LOG2E = 4, ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32), math.log2(math.e)
+Q5_K, Q6_K, IQ4_XS, GGML_BLOCK_SIZE, Q8_GROUP_SIZE, Q5_WORDS, Q6_BYTES, IQ4_WORDS = 13, 14, 23, 256, 32, 44, 210, 34
+
+def warp_reduce(val:UOp, maximum:bool=False, full_wave:bool=False) -> UOp:
+  for offset in ((16, 8, 4, 2, 1) if full_wave else (8, 4, 2, 1)):
+    if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
+    other = UOp(Ops.CUSTOM, dtypes.float, (val,), arg=
+      f"__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {{0}}), {0x1f | offset<<10}))")
+    val = val.maximum(other) if maximum else val + other
+  return val
+
+def _reg(shape:tuple[int, ...], slot:int, value:float, dep:UOp|None=None) -> UOp:
+  ret = UOp.placeholder(shape, dtypes.float, slot=slot, addrspace=AddrSpace.REG)
+  return ret.after((ret if dep is None else ret.after(dep)).store(ret.const_like(value)))
+
+@functools.cache
+def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, cache_scale, valid_kv_len, max_kv_len, block_n):
+  if isinstance(valid_kv_len, UOp): valid_kv_len = kernel_var(valid_kv_len.unbind_all()[0])
+  _, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
+  _, H, M, _ = cast(tuple[int, int, int, int], q.shape)
+  assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % block_n == 0
+  G, CHUNK, DV, heads_per_wave = H // H_KV, block_n, D // WARP_SIZE, 2
+  head_tile = min(DECODE_HEAD_TILE, G)  # share each KV stream across two GQA heads per wave
+  assert G % head_tile == 0 and head_tile % heads_per_wave == 0
+  decode_waves, decode_group = head_tile // heads_per_wave, 4
+  block_bhkv = UOp.range(B*H_KV*(G//head_tile), 0, AxisType.GLOBAL)
+  valid_chunks = (valid_kv_len+CHUNK-1)//CHUNK
+  group_count = min(valid_chunks, out.shape[2]) if isinstance(valid_chunks, int) else valid_chunks.minimum(out.shape[2])
+  block_n = UOp.range(group_count, 1, AxisType.GLOBAL)
+  lane, wave = UOp.range(WARP_SIZE, 2, AxisType.LOCAL), UOp.range(decode_waves, 3, AxisType.LOCAL)
+  head_group, bhkv = block_bhkv % (G//head_tile), block_bhkv // (G//head_tile)
+  b, kv_head = bhkv // H_KV, bhkv % H_KV
+  dims = tuple(lane + i*WARP_SIZE for i in range(DV))
+  acc, row_max, row_sum = _reg((heads_per_wave, DV), 0, 0), _reg((heads_per_wave,), 1, -math.inf), _reg((heads_per_wave,), 2, 0)
+  groups_per_chunk, offset = CHUNK // decode_group, UOp.range(((valid_chunks+group_count-1)//group_count)*(CHUNK//decode_group), 100, AxisType.REDUCE)
+  chunk = block_n + (offset // groups_per_chunk) * group_count
+  keys = tuple(chunk*CHUNK + (offset % groups_per_chunk)*decode_group + i for i in range(decode_group))
+  valid = tuple(key < valid_kv_len for key in keys)
+  kvals, vvals = (tuple(tuple(cache_kv[kv, b, kv_head, key, d].float() *
+    is_valid.where(cache_scale[kv, b, kv_head, key].float(), UOp.const(0, dtypes.float)) for d in dims)
+    for key,is_valid in zip(keys, valid)) for kv in range(2))
+  q_heads = tuple(kv_head*G + head_group*head_tile + wave*heads_per_wave + head for head in range(heads_per_wave))
+  updates:list[UOp] = []
+  for head,q_head in enumerate(q_heads):
+    scores = tuple(warp_reduce(sum((q[b, q_head, 0, d].float()*k for d,k in zip(dims, key_kvals)),
+                                   UOp.const(0, dtypes.float)), full_wave=True) / math.sqrt(D) for key_kvals in kvals)
+    prev_acc, prev_max, prev_sum = acc.after(offset)[head], row_max.after(offset)[head], row_sum.after(offset)[head]
+    new_max = functools.reduce(lambda a,vs:a.maximum(vs[0].where(vs[1], UOp.const(-math.inf, dtypes.float))), zip(valid, scores), prev_max)
+    alpha = ((prev_max-new_max)*LOG2E).exp2()
+    betas = tuple(is_valid.where(((score-new_max)*LOG2E).exp2(), UOp.const(0, dtypes.float)) for is_valid,score in zip(valid, scores))
+    updates += [acc[head].store(prev_acc*alpha + sum((UOp.stack(*value)*beta for value,beta in zip(vvals, betas)), acc[head].const_like(0))),
+                row_sum[head].store(prev_sum*alpha + sum(betas, UOp.const(0, dtypes.float))), row_max[head].store(new_max)]
+  update = UOp.group(*updates).end(offset)
+  acc, row_max, row_sum = acc.after(update), row_max.after(update), row_sum.after(update)
+  stores = [out[b, q_head, block_n, d].store(acc[head, i]) for head,q_head in enumerate(q_heads) for i,d in enumerate(dims)] + \
+    [stats[b, q_head.valid(lane.eq(0)), block_n, i].store(x[head])
+     for head,q_head in enumerate(q_heads) for i,x in enumerate((row_max, row_sum))]
+  return UOp.group(*stores).end(lane, wave, block_n, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
+
+def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, cache_scale:Tensor, max_kv_len:int) -> Tensor:
+  B, H, D = cache_kv.shape[1], q.shape[1], cache_kv.shape[4]
+  block_n = 128
+  chunks = min(64, max_kv_len // block_n)
+  partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
+  stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
+  decode_partial = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=block_n)
+  partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, cache_scale, fxn=decode_partial)[:2]
+  live_chunks = (valid_kv_len+block_n-1)//block_n
+  live_chunks = min(live_chunks, chunks) if isinstance(live_chunks, int) else live_chunks.minimum(chunks)
+  partial, stats = partial[:, :, :live_chunks], stats[:, :, :live_chunks]
+  weights = ((stats[..., 0]-stats[..., 0].max(2, keepdim=True))*LOG2E).exp2()
+  return ((partial*weights.unsqueeze(-1)).sum(2) / (stats[..., 1]*weights).sum(2, keepdim=True)).unsqueeze(2)
+
+@functools.cache
+def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, kv_scale:UOp, valid_kv_len:int|UOp) -> UOp:
+  if isinstance(valid_kv_len, UOp): valid_kv_len = kernel_var(valid_kv_len.unbind_all()[0])
+  BH, M, D = q.shape
+  _, B, H_KV, physical_n, cache_dim = cache.shape
+  k, v = cache[0].reshape(B*H_KV, physical_n, cache_dim), cache[1].reshape(B*H_KV, physical_n, cache_dim)
+  kv_scale = kv_scale.reshape(2, B*H_KV, physical_n)
+  assert k.shape == v.shape and BH % k.shape[0] == 0 and k.shape[2] == D
+  gqa_group = BH // k.shape[0]
+  if isinstance(M, int) and isinstance(valid_kv_len, int):
+    assert M % BLOCK_M == 0 and valid_kv_len % BLOCK_N == 0
+  assert isinstance(D, int) and D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0
+  TM, TN, TD, SCALE = BLOCK_M//(WAVES_M*LANES_PER_WAVE_M), BLOCK_N//LANES_PER_WAVE_N, D//(WAVES_N*LANES_PER_WAVE_N), 1/math.sqrt(D)
+  block_bh, block_m = UOp.range(BH, 0, AxisType.GLOBAL), UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL)
+  q = q.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
+  kv_head = block_bh // gqa_group
+  k, v = k[kv_head], v[kv_head]
+  o = o.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
+  wave_m, wave_n, lane = UOp.range(WAVES_M, 2, AxisType.LOCAL), UOp.range(WAVES_N, 3, AxisType.LOCAL), UOp.range(WARP_SIZE, -1, AxisType.WARP)
+  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
+  lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
+  Q_ELEMS_PER_THREAD, KV_ELEMS_PER_THREAD = BLOCK_M * D // THREADS_PER_BLOCK, BLOCK_N * D // THREADS_PER_BLOCK
+  QP_lds = UOp.placeholder((BLOCK_M, D + LDS_PAD), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
+  KV_lds = UOp.placeholder((BLOCK_N, D + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :D]
+  acc, m_i, l_i = _reg((TM, TD), 2, 0), _reg((TM,), 3, -math.inf), _reg((TM,), 4, 0)
+  n_tiles = (valid_kv_len - M + (block_m + 1) * BLOCK_M + BLOCK_N - 1) // BLOCK_N
+  n_tile = UOp.range(n_tiles, 100, AxisType.REDUCE)
+  Q_lds = QP_lds[:, :D]
+  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid].store(q.reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid])
+  load_k = UOp.range(KV_ELEMS_PER_THREAD, 90, AxisType.WEAK)
+  kidx = n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_k
+  kval = k.reshape(physical_n*D)[kidx].float() * kv_scale[0, kv_head, kidx // D].float()
+  K_store = KV_lds.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_k].store(kval).end(load_k)
+  qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
+  Q_lds, KV_lds_k = Q_lds.after(qk_load_barrier), KV_lds.after(qk_load_barrier)
+  S_reg = _reg((TM, TN), 6, 0, n_tile)
+  k_qk, tm1, tn1 = UOp.range(D//WMMA_K, 101, AxisType.REDUCE), UOp.range(TM//WMMA_ACC, 200), UOp.range(TN, 201)
+  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
+  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
+  k_frag = KV_lds_k.reshape(TN, WMMA_N, D // WMMA_K, WMMA_K)[tn1, lane_n, k_qk]
+  qk_done = S_frag.store(UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)).end(tm1, tn1).end(k_qk)
+  S_reg = S_reg.after(qk_done)
+  S_reg = S_reg.after(S_reg.store(S_reg * SCALE))
+  rm, rn = UOp.range(TM, 250, AxisType.WEAK), UOp.range(TN, 251, AxisType.WEAK)
+  q_idx = valid_kv_len - M + block_m * BLOCK_M + wave_m * WMMA_M + rm * LANES_PER_WAVE_M + lane_m
+  k_idx = n_tile * BLOCK_N + rn * LANES_PER_WAVE_N + lane_n
+  valid = k_idx <= q_idx
+  S_reg = S_reg.after(S_reg[rm, rn].store(valid.where(S_reg[rm, rn], S_reg[rm, rn].const_like(-math.inf))).end(rm, rn))
+  m_ij, rm2 = _reg((TM,), 7, -math.inf, n_tile), UOp.range(TN, 261, AxisType.REDUCE)
+  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
+  ri_w = UOp.range(TM, 270)
+  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce(m_ij[ri_w], maximum=True)).end(ri_w))
+  tile_max = m_ij.reshape(TM, 1).expand(TM, TN).maximum(-1e30)
+  S_reg = S_reg.after(S_reg.store(((S_reg - tile_max) * LOG2E).exp2()))
+  p_local, ri_ws = _reg((TM,), 8, 0, n_tile), UOp.range(TM, 295, AxisType.WEAK)
+  p_sum = p_local.after(p_local[ri_ws].store(sum((warp_reduce(S_reg[ri_ws, rn]) for rn in range(TN)), S_reg.const_like(0))).end(ri_ws))
+  P_lds = QP_lds.flatten()[:WAVES_N * BLOCK_M * BLOCK_N].reshape(WAVES_N, BLOCK_M, BLOCK_N)
+  P_write = P_lds.reshape(WAVES_N, WAVES_M, TM, LANES_PER_WAVE_M, 1, TN, LANES_PER_WAVE_N, 1)
+  P_write = P_write.permute((1, 0, 3, 6, 2, 4, 5, 7)).reshape(THREADS_PER_BLOCK, TM, TN)
+  P_store = P_write[tid].store(S_reg.cast(dtypes.half))
+  beta_i, ri4 = UOp.placeholder((TM,), dtypes.float, slot=9, addrspace=AddrSpace.REG), UOp.range(TM, 330, AxisType.WEAK)
+  m_new_val = m_i[ri4].maximum(m_ij[ri4])
+  alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2()
+  beta_val = ((m_ij[ri4] - m_new_val) * LOG2E).exp2()
+  rj4 = UOp.range(TD, 331)
+  correction = UOp.group(acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
+                         l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
+                         m_i[ri4].store(m_new_val), beta_i[ri4].store(beta_val)).end(ri4)
+  acc, l_i, m_i, beta_i = acc.after(correction), l_i.after(correction), m_i.after(correction), beta_i.after(correction)
+  V_lds = UOp.placeholder((D, BLOCK_N + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :BLOCK_N]
+  V_copy, load_v = V_lds.after(qk_done).permute(1, 0), UOp.range(KV_ELEMS_PER_THREAD, 390, AxisType.WEAK)
+  vidx = n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v
+  vval = v.reshape(physical_n*D)[vidx].float() * kv_scale[1, kv_head, vidx // D].float()
+  V_store = V_copy.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_v].store(vval).end(load_v)
+  pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
+  P_lds, V_lds = P_lds.after(pv_barrier), V_lds.after(pv_barrier)
+  pv_acc = _reg((TM, TD), 10, 0, n_tile).after(pv_barrier)
+  k_pv, tm2, tn2 = UOp.range(BLOCK_N//WMMA_K, 400, AxisType.REDUCE), UOp.range(TM//WMMA_ACC, 401, AxisType.WEAK), UOp.range(TD, 402, AxisType.WEAK)
+  pv_frag = pv_acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
+  p_frag = P_lds[wave_n].reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
+  v_frag = V_lds.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
+  pv_done = pv_frag.store(UOp.wmma(p_frag, v_frag, pv_frag.after(k_pv), *WMMA_ARG)).end(tm2, tn2).end(k_pv)
+  pv_acc = pv_acc.after(pv_done)
+  ri5, rj5 = UOp.range(TM, 410, AxisType.WEAK), UOp.range(TD, 411, AxisType.WEAK)
+  accumulate = acc[ri5, rj5].store(acc[ri5, rj5] + beta_i[ri5] * pv_acc[ri5, rj5]).end(ri5, rj5)
+  n_tile_end = accumulate.barrier().end(n_tile)
+  acc, l_i, m_i = acc.after(n_tile_end), l_i.after(n_tile_end), m_i.after(n_tile_end)
+  acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD)))
+  o = o.reshape(WAVES_M, TM, LANES_PER_WAVE_M, 1, WAVES_N, TD, LANES_PER_WAVE_N, 1)
+  o = o.permute((0, 4, 2, 6, 1, 3, 5, 7)).reshape(THREADS_PER_BLOCK, TM, TD)
+  return o[tid].store(acc).end(wave_m, wave_n, lane).end(block_m, block_bh).sink(arg=KernelInfo(opts_to_apply=()))
+
+def flash_attention_causal_cached(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, cache_scale:Tensor) -> Tensor:
+  B, H, T, D = cast(tuple[int, int, int, int], q.shape)
+  out = Tensor.empty(B*H, T, D, dtype="float32", device=q.device)
+  flash_cached = functools.partial(_amd_flash_attention, valid_kv_len=valid_kv_len)
+  return Tensor.custom_kernel(out, q.reshape(B*H, T, D), cache_kv, cache_scale, fxn=flash_cached)[0].reshape(B, H, T, D)
+
+def quantized_attention(q:Tensor, stacked_kv:Tensor, cache_kv:Tensor, cache_scale:Tensor, start_pos:int|UOp) -> Tensor:
+  T = q.shape[2]
+  scale = (stacked_kv.float().abs().max(axis=-1, keepdim=True) / 127).maximum(1e-8).half()
+  packed_kv = (stacked_kv.float() / scale).round().clip(-127, 127).cast(dtypes.int8)
+  store_kv = cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(packed_kv.uop)
+  store_scale = cache_scale[:, :, :, start_pos:start_pos+T].uop.store(scale.squeeze(-1).uop)
+  # each store goes on its own buffer's AFTER: sharing both stores across both AFTERs leaves
+  # un-ended stores with open ranges in the kernel graph
+  assigned_kv, assigned_scale = Tensor(cache_kv.uop.after(store_kv)), Tensor(cache_scale.uop.after(store_scale))
+  # keep start_pos in its bound form at the graph level, the kernel builders unbind it to the kernel-side PARAM form
+  valid_end = start_pos+T
+  return amd_flash_attention_decode(q.half(), assigned_kv, valid_end, assigned_scale, cast(int, cache_kv.shape[3])) if resolve(T == 1) else \
+    flash_attention_causal_cached(q.half(), assigned_kv, valid_end, assigned_scale)
+
+def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp:
+  return UOp(Ops.CUSTOMI, dtypes.int32, (a.int(), b.int(), c), arg="__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)")
+
+def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp:
+  return UOp(Ops.CUSTOMI, dtypes.uint32, tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg="__builtin_amdgcn_perm({}, {}, {})")
+
+def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
+  assert ptr.op is Ops.INDEX
+  if lanes is None: return UOp(Ops.CUSTOMI, ptr.dtype, (ptr,), arg="__builtin_nontemporal_load({0})")
+  buf, coords = ptr.src[0], ptr.src[1:]
+  idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0, dtypes.weakint))
+  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes, dtypes.weakint))).load(dtype=ptr.dtype)
+
+def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
+def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()
+
+def _iq4_bytes(packed:UOp, shift:int) -> UOp:
+  selectors = (packed >> shift) & 0x0f0f0f0f
+  low = _amd_byte_perm(UOp.const(0xf6eaddcf, dtypes.uint32), UOp.const(0xbfad9881, dtypes.uint32), selectors)
+  high = _amd_byte_perm(UOp.const(0x71594535, dtypes.uint32), UOp.const(0x26190d01, dtypes.uint32), selectors & 0x07070707)
+  return _amd_byte_perm(high, low, 0x03020100 | ((selectors & 0x08080808) >> 1))
+
+@functools.cache
+def _q8_quantize_kernel(q:UOp, scale:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
+  groups = in_features//Q8_GROUP_SIZE
+  token_group, lane = UOp.range(tokens*groups, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, group = token_group//groups, token_group%groups
+  x = x.reshape(tokens, groups, 32)
+  group_scale = warp_reduce(x[token, group, lane].float().abs(), maximum=True, full_wave=True) / 127
+  group_scale = group_scale.maximum(1e-8)
+  word_lane = lane.minimum(7)
+  xs = tuple(x[token, group, word_lane*4+i].float() for i in range(4))
+  word = sum(((v/group_scale).round().clip(-127, 127).cast(dtypes.int8).cast(dtypes.uint8).cast(dtypes.uint32) << (i*8)
+              for i,v in enumerate(xs)), UOp.const(0, dtypes.uint32))
+  stores = (q[token, group, lane.valid(lane < 8)].store(word), scale[token, group.valid(lane.eq(0))].store(group_scale))
+  return UOp.group(*stores).end(token_group, lane).sink(arg=KernelInfo(name="q8_quantize", opts_to_apply=()))
+
+def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor]:
+  groups = in_features//Q8_GROUP_SIZE
+  q, scale = Tensor.empty(tokens, groups, 8, dtype=dtypes.uint32, device=x.device), \
+    Tensor.empty(tokens, groups, dtype=dtypes.float32, device=x.device)
+  q, scale = Tensor.custom_kernel(q, scale, x, fxn=functools.partial(_q8_quantize_kernel, tokens=tokens, in_features=in_features))[:2]
+  return q, scale
+
+@functools.cache
+def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp, start_pos:UOp|None=None) -> UOp:
+  batch, heads, tokens, value_dim, row_tile = *core.shape, 4
+  key_dim, alpha_dim = q.shape[-1], alpha.shape[-1] if len(alpha.shape) == 4 else 1
+  assert all(isinstance(x, int) for x in (batch, heads, tokens, value_dim, key_dim)) and key_dim % 32 == 0 and value_dim % row_tile == 0
+  batch, heads, tokens, value_dim, key_dim = cast(tuple[int, int, int, int, int], (batch, heads, tokens, value_dim, key_dim))
+  core, v = (x.reshape(batch*heads, tokens, value_dim) for x in (core, v))
+  q, k = (x.reshape(batch*heads, tokens, key_dim) for x in (q, k))
+  beta, kq = (x.reshape(batch*heads, tokens) for x in (beta, kq))
+  alpha, state = alpha.reshape(batch*heads, tokens, alpha_dim), state.reshape(batch*heads, value_dim, key_dim)
+  bh_row, lane = UOp.range(batch*heads*value_dim//row_tile, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  bh, row_base = bh_row // (value_dim//row_tile), (bh_row % (value_dim//row_tile))*row_tile
+  rows = tuple(row_base+i for i in range(row_tile))
+  cols = tuple(lane + i*32 for i in range(key_dim//32))
+  current = UOp.placeholder((row_tile*key_dim//32,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  initial = None if start_pos is None else start_pos.eq(0)
+  current = current.after(current.store(UOp.stack(*(state[bh, row, col].float() if initial is None else
+    initial.where(0, state[bh, row, col].float()) for row in rows for col in cols))))
+  token = UOp.range(tokens, 2, AxisType.REDUCE)
+  keys = tuple(k[bh, token, col].load() for col in cols)
+  queries = tuple(q[bh, token, col].load() for col in cols)
+  updates:list[UOp] = []
+  stores:list[UOp] = []
+  for row_idx,row in enumerate(rows):
+    previous = tuple(current.after(token)[row_idx*key_dim//32+i].load() for i in range(key_dim//32))
+    av, bv = alpha[bh, token, row if alpha_dim > 1 else 0].load(), beta[bh, token].load()
+    state_k = warp_reduce(sum((x*y for x,y in zip(previous, keys)), UOp.const(0, dtypes.float32)), full_wave=True)
+    state_q = warp_reduce(sum((x*y for x,y in zip(previous, queries)), UOp.const(0, dtypes.float32)), full_wave=True)
+    delta = (v[bh, token, row].load() - state_k*av) * bv
+    updates += [x*av + delta*y for x,y in zip(previous, keys)]
+    stores.append(core[bh, token, row.valid(lane.eq(0))].store(state_q*av + delta*kq[bh, token]))
+  step = UOp.group(*stores, current.store(UOp.stack(*updates))).end(token)
+  state_stores = (state[bh, row, col].store(current.after(step)[row_idx*key_dim//32+i].load().cast(state.dtype))
+                  for row_idx,row in enumerate(rows) for i,col in enumerate(cols))
+  return UOp.group(*state_stores).end(lane, bh_row).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
+
+def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int):
+  output_waves = 2 if out_features % (32*output_tiles) == 0 else 1
+  token_block, output_block = UOp.range(out.shape[0]//token_tile, 0), UOp.range(out_features//(16*output_tiles*output_waves), 1)
+  lane, wave = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL), UOp.range(output_waves, 3, axis_type=AxisType.LOCAL)
+  hw_lane = UOp(Ops.CUSTOM, dtypes.int32, (lane.int(),), arg="__builtin_amdgcn_mbcnt_lo(-1, 0)").cast(dtypes.weakint)
+  col, half = hw_lane % 16, hw_lane // 16
+  outputs = tuple((output_block*output_waves+wave)*(16*output_tiles) + tile*16 + col for tile in range(output_tiles))
+  inputs = tuple(token_block*token_tile + tile*16 + col for tile in range(token_tile//16))
+  tokens = tuple(tuple(token_block*token_tile + tile*16 + half*8 + i for i in range(8)) for tile in range(token_tile//16))
+  return output_waves, token_block, output_block, lane, wave, half, outputs, inputs, tokens
+
+def _wmma_stores(out, outputs, tokens, accs, update, half):
+  def values(acc:UOp) -> tuple[UOp, ...]:
+    vals = tuple(acc.after(update)[i].load() for i in range(8))
+    swapped = tuple(UOp(Ops.CUSTOM, dtypes.float32, (value,),
+      arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {0}), 50688))") for value in vals)
+    low = half.eq(0)
+    return tuple(low.where(vals[i], swapped[i+4]) if j == 0 else low.where(swapped[i], vals[i+4]) for i in range(4) for j in range(2))
+  return [out[token, output].store(value) for output,output_accs in zip(outputs, accs)
+          for tile_tokens,acc in zip(tokens, output_accs) for token,value in zip(tile_tokens, values(acc))]
+
+def _decode_linear(out:UOp, out_features:int, group_count:int, group_dot, name:str) -> UOp:
+  chunks = (group_count+31)//32
+  token_output_chunk, lane = UOp.range(out.shape[0]*out_features*chunks, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, output, chunk = token_output_chunk // (out_features*chunks), (token_output_chunk//chunks) % out_features, token_output_chunk % chunks
+  group = lane+chunk*32
+  value = group_dot(token, output, group) if group_count % 32 == 0 else \
+    (group < group_count).where(group_dot(token, output, group.minimum(group_count-1)), UOp.const(0, dtypes.float32))
+  total = warp_reduce(value, full_wave=True)
+  return out[token, output, chunk.valid(lane.eq(0))].store(total.cast(out.dtype)).end(token_output_chunk, lane).sink(
+    arg=KernelInfo(name=name, opts_to_apply=()))
+
+def _q5_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp, UOp, UOp]:
+  scale = (subgroup < 4).where(_load_byte(raw, base, 4 + subgroup) & 63,
+    (_load_byte(raw, base, 8 + subgroup) & 15) | ((_load_byte(raw, base, subgroup) >> 6) << 4))
+  minimum = (subgroup < 4).where(_load_byte(raw, base, 8 + subgroup) & 63,
+    (_load_byte(raw, base, 8 + subgroup) >> 4) | ((_load_byte(raw, base, 4 + subgroup) >> 6) << 4))
+  d, dmin = (raw[base] & 0xffff).cast(dtypes.uint16), (raw[base] >> 16).cast(dtypes.uint16)
+  return _half(d), _half(dmin), scale.float(), minimum.float()
+
+@functools.cache
+def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
+  group_count = in_features // Q8_GROUP_SIZE
+  def group_dot(token:UOp, output:UOp, group:UOp) -> UOp:
+    block, subgroup = group // 8, group % 8
+    xwords = _amd_load(xq[token, group, 0], 8)
+    if ggml_type == Q5_K:
+      base = (output * in_features//GGML_BLOCK_SIZE + block) * Q5_WORDS
+      qs_base, dot, qsum = base + 12 + (subgroup//2)*8, UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)
+      for word_idx in range(8):
+        word = (raw[qs_base+word_idx] >> ((subgroup&1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f
+        word |= ((raw[base+4+word_idx] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4
+        dot, qsum = _amd_dp4a(word, xwords[word_idx], dot), _amd_dp4a(UOp.const(0x01010101, dtypes.uint32), xwords[word_idx], qsum)
+      d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
+      return (dot.float()*d*scale - qsum.float()*dmin*minimum) * xd[token, group]
+    if ggml_type == IQ4_XS:
+      base = (output * in_features//GGML_BLOCK_SIZE + block) * IQ4_WORDS
+      dot = UOp.const(0, dtypes.int32)
+      for word_idx in range(8):
+        packed = _amd_load(raw[base + 2 + subgroup*4 + word_idx%4])
+        dot = _amd_dp4a(_iq4_bytes(packed, 4*(word_idx//4)), xwords[word_idx], dot)
+      d, scale = _iq4_scales(raw, base, subgroup)
+      return dot.float() * xd[token, group] * d * scale
+    base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_BYTES
+    dots = [UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)]
+    for word_idx in range(8):
+      pos, within = subgroup*32 + word_idx*4, (subgroup*32 + word_idx*4)%128
+      low = _amd_load(raw[base + (pos//128)*64 + within%64], 4) >> ((within//64)*4).cast(dtypes.uint8)
+      high = _amd_load(raw[base + 128 + (pos//128)*32 + within%32], 4) >> ((within//32)*2).cast(dtypes.uint8)
+      quant = ((low & 15) | ((high & 3) << 4)).bitcast(dtypes.int8) - 32
+      word = sum((quant[i].cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i in range(4)), UOp.const(0, dtypes.uint32))
+      dots[word_idx//4] = _amd_dp4a(word, xwords[word_idx], dots[word_idx//4])
+    scales = [raw[base + 192 + subgroup*2+i].cast(dtypes.uint8).bitcast(dtypes.int8).float() for i in range(2)]
+    dbits = raw[base+208].cast(dtypes.uint16) | (raw[base+209].cast(dtypes.uint16) << 8)
+    return (dots[0].float()*scales[0] + dots[1].float()*scales[1]) * xd[token, group] * _half(dbits)
+  return _decode_linear(out, out_features, group_count, group_dot, {Q5_K:"linear_q5_k", IQ4_XS:"linear_iq4_xs", Q6_K:"linear_q6"}[ggml_type])
+
+def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, dequant, name):
+  x = x.reshape(out.shape[0], in_features)
+  _, token_block, output_block, lane, wave, physical_half, outputs, input_tokens, tokens = layout
+  token_tile, output_tiles = len(tokens)*16, len(outputs)
+  output_words = in_features // GGML_BLOCK_SIZE * type_words
+  accs = tuple(tuple(UOp.placeholder((8,), dtypes.float32, slot=ot*(token_tile//16)+tile, addrspace=AddrSpace.REG)
+                     for tile in range(token_tile // 16)) for ot in range(output_tiles))
+  accs = tuple(tuple(acc.after(acc.store(acc.const_like(0))) for acc in output_accs) for output_accs in accs)
+  group = UOp.range(in_features // Q8_GROUP_SIZE, 4, AxisType.REDUCE)
+  block, subgroup = group // 8, group % 8
+  wmma_accs = [list(output_accs) for output_accs in accs]
+  for half in range(2):
+    afrags = tuple(UOp.stack(*(x[input_token, group*32 + half*16 + i].cast(dtypes.float16) for i in range(16)))
+                   for input_token in input_tokens)
+    for output_tile,output in enumerate(outputs):
+      bfrag = UOp.stack(*dequant(output*output_words + block*type_words, subgroup, half))
+      for tile,afrag in enumerate(afrags):
+        previous = accs[output_tile][tile].after(group) if half == 0 else wmma_accs[output_tile][tile]
+        wmma_accs[output_tile][tile] = UOp.wmma(afrag, bfrag, previous, *WMMA_ARG)
+  update = UOp.group(*(acc.store(value) for output_accs,output_values in zip(accs, wmma_accs)
+                       for acc,value in zip(output_accs, output_values))).end(group)
+  return UOp.group(*_wmma_stores(out, outputs, tokens, accs, update, physical_half)).end(token_block, output_block, lane, wave).sink(
+    arg=KernelInfo(name=name, opts_to_apply=()))
+
+@functools.cache
+def _q5_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, out_features:int, in_features:int) -> UOp:
+  token_tile, output_tiles = (64, 1) if out_features <= 1024 and out.shape[0] % 64 == 0 else \
+    (64, 2) if out.shape[0] % 64 == 0 else (32 if out.shape[0] % 32 == 0 else 16, 2)
+  def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
+    d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
+    qs_base = base + 12 + (subgroup // 2)*8 + half*4
+    words = tuple((raw[qs_base+i] >> ((subgroup&1)*4).cast(dtypes.uint32) & 0x0f0f0f0f) |
+      ((raw[base+4+half*4+i] >> subgroup.cast(dtypes.uint32) & 0x01010101) << 4) for i in range(4))
+    return tuple(((word >> (byte*8) & 255).float()*d*scale-dmin*minimum).cast(dtypes.float16) for word in words for byte in range(4))
+  return _quant_linear_wmma(out, x, out_features, in_features, Q5_WORDS,
+                            _wmma_layout(out, out_features, token_tile, output_tiles), dequant, "linear_q5_k_f16_wmma")
+
+def _iq4_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp]:
+  low = _load_byte(raw, base, 4 + subgroup//2)
+  scale = ((low >> (4*(subgroup%2)).cast(dtypes.uint32)) & 15) | ((((raw[base] >> 16) >> (2*subgroup).cast(dtypes.uint32)) & 3) << 4)
+  return _half(raw[base] & 0xffff), (scale.cast(dtypes.uint8).bitcast(dtypes.int8)-32).float()
+
+@functools.cache
+def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, out_features:int, in_features:int) -> UOp:
+  token_tile = 32 if out_features <= 1024 and out.shape[0] % 32 == 0 else 64 if out.shape[0] % 64 == 0 and \
+    (out_features <= 6144 or out_features == 5120 and in_features > 8192) else 128 if out.shape[0] % 128 == 0 else \
+    32 if out.shape[0] % 32 == 0 else 16
+  output_tiles = 1 if out_features <= 1024 else 2 if out_features <= 6144 else 1 if out_features < 8192 else 2
+  layout = _wmma_layout(out, out_features, token_tile, output_tiles)
+  output_waves, _, _, lane, wave, _, _, _, _ = layout
+  local_lut = UOp.placeholder((256,), dtypes.uint32, slot=32, addrspace=AddrSpace.LOCAL)
+  tid, lut_items = wave*32+lane, 256//(32*output_waves)
+  lut = local_lut.after(UOp.group(*(local_lut[tid*lut_items+i].store(lut[tid*lut_items+i]) for i in range(lut_items))).barrier())
+  def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
+    d, scale = _iq4_scales(raw, base, subgroup)
+    scale = scale * d
+    if out_features <= 6144:
+      pairs = tuple(lut[((raw[base + 2 + subgroup*4 + word] >> (byte*8)) & 255).cast(dtypes.weakint)]
+                    for word in range(4) for byte in range(4))
+      return tuple((_half((pair >> (half*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in pairs)
+    def nibble(packed:UOp, index:int): return (packed >> (8*index+4*half)) & 15
+    lut_pairs = (lut[(nibble(packed, i) | nibble(packed, i+1)<<4).cast(dtypes.weakint)]
+                 for packed in (raw[base+2+subgroup*4+i] for i in range(4)) for i in (0, 2))
+    return tuple((_half((pair >> (i*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in lut_pairs for i in range(2))
+  return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma")
+
+def q8_linear(layer:Linear, x:Tensor) -> Tensor:
+  assert layer.ggml_type in (Q5_K, Q6_K, IQ4_XS)
+  tokens = int(x.numel()) // layer.in_features
+  raw = layer.weight.uop.buf_uop
+  out_features, in_features = layer.out_features, layer.in_features
+  use_wmma = tokens % 16 == 0 and layer.out_features % 16 == 0
+  def run(fxn:Callable[..., UOp], out:UOp, *srcs:UOp) -> Tensor:
+    all_srcs = (out,)+srcs
+    params = tuple(UOp.placeholder_like(src, slot=i) for i,src in enumerate(all_srcs))
+    kernel = fxn(*params, out_features=out_features, in_features=in_features).call(*all_srcs)
+    result = Tensor(out.after(kernel))
+    if len(result.shape) == 3: result = result.sum(-1)
+    result = result.reshape(*x.shape[:-1], layer.out_features)
+    return result if layer.bias is None else result + layer.bias
+
+  out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device).uop
+  if layer.ggml_type == Q5_K and use_wmma:
+    return run(_q5_linear_f16_wmma_kernel, out, raw, x.cast(dtypes.float16).contiguous().uop)
+  if layer.ggml_type == IQ4_XS and use_wmma:
+    return run(_iq4_linear_f16_wmma_kernel, out, raw, x.cast(dtypes.float16).contiguous().uop,
+               iq4_half_lut(str(x.device)).uop)
+  xq, xd = q8_quantize(x, tokens, layer.in_features)
+  decode = functools.partial(_quant_decode_kernel, ggml_type=layer.ggml_type)
+  out = Tensor.empty(tokens, layer.out_features, (layer.in_features+1023)//1024, dtype=dtypes.float32, device=x.device).uop
+  return run(decode, out, raw, xq.uop, xd.uop)
+
+@functools.cache
+def iq4_half_lut(device:str) -> Tensor:
+  from tinygrad.runtime.autogen.ggml_common import kvalues_iq4nl
+  return Tensor([x for j in range(16) for i in range(16) for x in (kvalues_iq4nl[i], kvalues_iq4nl[j])],
+                dtype=dtypes.float16, device=device).bitcast(dtypes.uint32).contiguous()

tinygrad/llm/model.py

@@ -1,17 +1,11 @@
 from __future__ import annotations
-import enum, functools, itertools, pathlib
+import functools, itertools, pathlib
 from dataclasses import dataclass, replace
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
-from tinygrad.nn import Linear
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, Context, dtypes
+from tinygrad.llm.kernels import Linear, gated_delta_prefill, amd_custom_kernels_supported
 from tinygrad.llm.gguf import gguf_load
 from tinygrad.uop.ops import resolve
 
-class ExpertGating(enum.IntEnum):
-  SOFTMAX = 1
-  SIGMOID = 2
-  SOFTMAX_WEIGHT = 3  # softmax over the top-k selected logits
-  SQRT_SOFTPLUS = 4
-
 @functools.cache
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str|None=None) -> Tensor:
   freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
@@ -67,7 +61,6 @@ class TransformerConfig:
   num_experts: int = 0
   num_experts_per_tok: int = 0
   norm_topk_prob: bool = False
-  expert_gating_func: ExpertGating = ExpertGating.SOFTMAX
   q_lora_rank: int = 0
   kv_lora_rank: int = 0
   shared_expert_dim: int = 0
@@ -110,21 +103,14 @@ class FFNBlock:
     if hasattr(self, 'ffn_gate_exps'):
       h = x.unsqueeze(2)  # (B, T, 1, D) - add expert dim for broadcasting
       logits = self.ffn_gate_inp(x)
-      bias = self.exp_probs_b["bias"] if hasattr(self, 'exp_probs_b') else None
-      gating, normalize_topk = self.config.expert_gating_func, self.config.norm_topk_prob
-      # fast path: without selection bias, normalized SOFTMAX is equivalent to SOFTMAX_WEIGHT
-      if gating == ExpertGating.SOFTMAX and bias is None and normalize_topk:
-        gating, normalize_topk = ExpertGating.SOFTMAX_WEIGHT, False
-      if   gating == ExpertGating.SOFTMAX_WEIGHT: scores = logits
-      elif gating == ExpertGating.SOFTMAX:        scores = logits.softmax(-1)
-      elif gating == ExpertGating.SIGMOID:        scores = logits.sigmoid()
-      elif gating == ExpertGating.SQRT_SOFTPLUS:  scores = logits.softplus().sqrt()
-
-      _, sel = pairwise_topk(scores if bias is None else scores + bias, self.config.num_experts_per_tok)
-      probs = scores.gather(-1, sel)
-      # SOFTMAX_WEIGHT applies softmax after top-k selection
-      if gating == ExpertGating.SOFTMAX_WEIGHT: probs = probs.softmax(-1)
-      if normalize_topk: probs = probs / probs.sum(axis=-1, keepdim=True)
+      if hasattr(self, 'exp_probs_b'):
+        probs = logits.sigmoid()
+        _, sel = pairwise_topk(probs + self.exp_probs_b["bias"], self.config.num_experts_per_tok)
+        probs = probs.gather(-1, sel)
+        if self.config.norm_topk_prob: probs = probs / probs.sum(axis=-1, keepdim=True)
+      else:
+        vals, sel = pairwise_topk(logits, self.config.num_experts_per_tok)
+        probs = vals.softmax(-1) if self.config.norm_topk_prob else logits.softmax(-1).gather(-1, sel)
       probs = probs * self.config.routed_scaling_factor
       x_down = self.ffn_down_exps(sel, (self.ffn_gate_exps(sel, h).silu() * self.ffn_up_exps(sel, h)).contiguous())  # (B, T, k, D)
       out = (x_down * probs.unsqueeze(-1)).sum(axis=2)  # (B, T, D)
@@ -181,26 +167,38 @@ class TransformerBlock(FFNBlock):
     k = apply_rope(k[..., :self.config.rope_dim], self.freqs_cis[start_pos:start_pos+T]).cat(k[..., self.config.rope_dim:], dim=-1)
 
     # NOTE: we don't want to change self.cache_kv, the function API doesn't support this well
-    assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).uop)))
-    k = assigned_kv[0, :, :, 0:start_pos+T, :]
-    v = assigned_kv[1, :, :, 0:start_pos+T, :]
+    stacked_kv = Tensor.stack(k, v)
+    if hasattr(self, "cache_kv_scale"):
+      from tinygrad.llm.kernels.amd import quantized_attention
+      attn = quantized_attention(q, stacked_kv, self.cache_kv, self.cache_kv_scale, start_pos)
+    else:
+      assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(stacked_kv.uop)))
+      k = assigned_kv[0, :, :, 0:start_pos+T, :]
+      v = assigned_kv[1, :, :, 0:start_pos+T, :]
 
-    #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
-    #k = self.cache_kv[0, :, :, 0:start_pos+T, :]
-    #v = self.cache_kv[1, :, :, 0:start_pos+T, :]
+      #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(stacked_kv)
+      #k = self.cache_kv[0, :, :, 0:start_pos+T, :]
+      #v = self.cache_kv[1, :, :, 0:start_pos+T, :]
 
-    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
-    # TODO: this if statement should be removed and it shouldn't generate extra kernels
-    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
-      if resolve(T != 1) else None
-    attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)     # (B,H,T,Hd)
+      # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
+      # TODO: this if statement should be removed and it shouldn't generate extra kernels
+      mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
+        if resolve(T != 1) else None
+      attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)     # (B,H,T,Hd)
     attn = attn.transpose(1, 2).reshape(B, T, -1)                                    # back to (B,T,D)
     return self.attn_output(attn if not self.config.attn_output_gate else (attn * gate.sigmoid()))
 
   def _init_state(self, x:Tensor):
     if not hasattr(self, "cache_kv"):
+      # hybrid models use a quantized KV cache on AMD, sized in flash decode blocks of 256
+      quantize = amd_custom_kernels_supported(x.device) and self.config.ssm is not None
+      assert not quantize or self.config.max_context % 256 == 0, \
+        f"quantized KV cache needs max_context to be a multiple of 256, got {self.config.max_context}"
       self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
-                                   dtype=dtypes.default_float, device=x.device)
+                                   dtype=dtypes.int8 if quantize else dtypes.default_float, device=x.device)
+      if quantize:
+        self.cache_kv_scale = Tensor.zeros(2, x.shape[0], self.config.n_kv_heads, self.config.max_context,
+                                           dtype=dtypes.float16, device=x.device)
       self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
 class MLATransformerBlock(FFNBlock):
@@ -272,12 +270,8 @@ class GatedDeltaNetBlock(FFNBlock):
 
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     B, T, _ = x.shape
-    # 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()
@@ -285,57 +279,36 @@ class GatedDeltaNetBlock(FFNBlock):
     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)
-    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))
+    log_alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, T, self.num_v_heads, -1) * self.ssm_a).squeeze(-1) \
+      if is_kda else ((alpha.float() + self.ssm_dt["bias"]).softplus() * self.ssm_a).reshape(B, T, self.num_v_heads)
 
     # 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:]))
+    conv_state = Tensor(start_pos).eq(0).where(0, self.conv_state)
+    conv_window = conv_state.cat(self.attn_qkv(x), dim=1)
+    conv_out = ((conv_window * self.ssm_conv1d["weight"].T.unsqueeze(0)).sum(1) if is_kda and resolve(T == 1) else functools.reduce(lambda a,b: a+b,
+      (conv_window[:, i:i+T] * self.ssm_conv1d["weight"][:, i] for i in range(self.ssm_conv_kernel)))).silu()
     q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-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)
+    q = q.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-12 if is_kda else 1e-6)
+    k = k.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-12 if is_kda else 1e-6)
+    q, k = q.repeat(1, 1, self.num_v_heads//self.num_k_heads, 1), k.repeat(1, 1, self.num_v_heads//self.num_k_heads, 1)
+    v = v.reshape(B, T, self.num_v_heads, self.head_v_dim)
+    q, k, v, beta = [z.transpose(1, 2).float() for z in (q, k, v, beta)]
+    alpha = log_alpha.transpose(1, 2).exp()
 
-    # 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))
+    # recurrent, the conv and recurrent states are updated in place
+    conv_state = conv_window[:, T:T+self.ssm_conv_kernel-1].cast(self.conv_state.dtype).contiguous()
+    state = Tensor(self.recurrent_state.uop.after(self.conv_state.uop.store(conv_state.uop)))
+    core = gated_delta_prefill(q * self.head_k_dim**-0.5, k, v, beta, alpha, state, Tensor(start_pos)).transpose(1, 2)
 
-    # 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; 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))
+    # output
+    core_attn_out = self.ssm_norm(core)
+    out_gate = out_gate.sigmoid() if is_kda else out_gate.silu()
+    return self.ssm_out((core_attn_out * out_gate).reshape(B, T, -1).cast(x.dtype)).contiguous()
 
   def _init_state(self, x):
     if not hasattr(self, "conv_state"):
-      self.conv_state = Tensor.zeros(x.shape[0], self.ssm_conv_kernel-1, self.conv_channels, device=x.device).clone()
-      self.recurrent_state = Tensor.zeros(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device).clone()
+      self.conv_state = Tensor.empty(x.shape[0], self.ssm_conv_kernel-1, self.conv_channels, device=x.device)
+      self.recurrent_state = Tensor.empty(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device)
 
 class Transformer:
   def __init__(self, config:TransformerConfig):
@@ -430,7 +403,6 @@ class Transformer:
       qk_norm=int(state_dict['blk.0.attn_q_norm.weight'].shape[0]) if 'blk.0.attn_q_norm.weight' in state_dict else 0,
       num_experts=kv.get(f'{arch}.expert_count', 0), num_experts_per_tok=kv.get(f'{arch}.expert_used_count', 0),
       norm_topk_prob=kv.get(f'{arch}.expert_weights_norm', arch in ('qwen3moe', 'qwen35moe', 'kimi-linear')),
-      expert_gating_func=ExpertGating(kv.get(f'{arch}.expert_gating_func', ExpertGating.SOFTMAX)),
       kv_lora_rank=kv_lora_rank, q_lora_rank=kv.get(f'{arch}.attention.q_lora_rank', 0),
       leading_dense_blocks=kv.get(f'{arch}.leading_dense_block_count', 0),
       shared_expert_dim=kv.get(
@@ -450,9 +422,6 @@ class Transformer:
       Tensor.realize(*params)
     return model, kv
 
-  def warmup(self):
-    for _ in range(2): list(zip(range(2), self.generate([0])))
-
   def get_start_pos(self, tokens:list[int]) -> int:
     # recurrent state can't be partially reused after divergence: reuse it only when tokens extend the cached prefix
     if self.has_recurrent_block:
@@ -461,8 +430,19 @@ class Transformer:
     prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
     return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
 
+  def warmup(self, chunk_size:int=32):
+    prompt = [0] * (min(chunk_size, 256, self.max_context-1) if self.has_recurrent_block else 1)
+    if self.has_recurrent_block:
+      x = Tensor.empty(1, 1, self.blk[0].config.dim, device=self.token_embd.weight.device)
+      for block in self.blk: block._init_state(x)
+    for _ in range(2):
+      # NOTE: chunk_size must match what generate uses at serve time, otherwise the captured JIT rejects the new toks range
+      warm = self.generate(prompt, chunk_size=chunk_size)
+      with Context(JIT_BATCH_SIZE=getenv("PREFILL_JIT_BATCH_SIZE", 512) if self.has_recurrent_block else 0): next(warm)
+      with Context(JIT_BATCH_SIZE=0): next(warm)
+      self._cached_tokens = []
+
   def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0):
-    if self.has_recurrent_block: chunk_size = 1
     v_start_pos = UOp.variable("start_pos", 0, self.max_context-1)
     v_toks = UOp.variable("toks", 1, chunk_size)
     # TODO: use UOp.variable for temperature once float variables are supported
@@ -473,8 +453,10 @@ class Transformer:
     start_pos = self.get_start_pos(tokens)
     out, prompt_len = None, len(tokens)
     while len(tokens) < self.max_context:
-      n_toks = min(chunk_size, len(tokens) - start_pos)
-      sp, nt = v_start_pos.bind(start_pos), v_toks.bind(n_toks)
+      # recurrent blocks prefill full chunks with a static shape, the tail of the prompt goes through the decode graph
+      remaining = len(tokens)-start_pos
+      n_toks = 1 if self.has_recurrent_block and remaining < chunk_size else min(chunk_size, remaining)
+      sp, nt = v_start_pos.bind(start_pos), n_toks if self.has_recurrent_block else v_toks.bind(n_toks)
       out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize()
       start_pos += n_toks
       # chunked prefill: keep processing until all prompt tokens are consumed

tinygrad/mixin/elementwise.py

@@ -115,8 +115,7 @@ class ElementwiseMixin(CreationMixin):
     ```
     """
     a, b = self._broadcasted(x, reverse)
-    # alu, not +: _broadcasted already promoted these, and a second promote would cast -b (only a bare weak CONST is kept weak)
-    return a.alu(Ops.ADD, -b)
+    return a + (-b)
 
   def mul(self, x: Self | ConstType, reverse: bool = False) -> Self:
     """
@@ -246,9 +245,8 @@ class ElementwiseMixin(CreationMixin):
     if dtypes.is_int(a.dtype) and dtypes.is_int(b.dtype):
       if rounding_mode == "trunc": return a.alu(Ops.CDIV, b)
       if rounding_mode == "floor": return a.alu(Ops.FLOORDIV, b)
-    if dtypes.is_int(a.dtype) or a.dtype == dtypes.bool: a = a.cast(dtypes.default_float)
-    # alu, not *: _broadcasted already promoted these, and a second promote would cast 1/b (only a bare weak CONST is kept weak)
-    d = a.alu(Ops.MUL, b.reciprocal())
+      a = a.cast(dtypes.default_float)
+    d = a * b.reciprocal()
     if rounding_mode is None: return d
     if rounding_mode == "trunc": return d.trunc()
     if rounding_mode == "floor": return d.floor()

tinygrad/mixin/gradient.py

@@ -3,7 +3,6 @@ import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes
 from tinygrad.helpers import argsort
 from tinygrad.dtype import sum_acc_dtype
-from tinygrad.function import renumber_invalid_outputs
 
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
   if op == Ops.ADD: return (ctx._broadcast_to(ret.src[0].shape),)
@@ -33,7 +32,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   params = {x.arg.slot:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
   grad_args = ctx.src
   root_grad = UOp(Ops.TUPLE, src=tuple(UOp(Ops.NOOP) if g.op is Ops.NOOP else
-    g if g.device is None else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
+    g if g.base.op is Ops.CONST else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
   grads = compute_gradient(fxn, root_grad, set(params.values()))
   # for precompiled calls, substitute forward outputs with params so intermediates aren't recomputed
   fwd_subs = {src: src.param_like(len(args)+len(grad_args)+i) for i, src in enumerate(fxn.src)} if k.arg.precompile else {}
@@ -41,7 +40,6 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   # collect needed gradient bodies, compact unused params, create a single backward CALL
   grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
   bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
-  bwd_body = renumber_invalid_outputs(bwd_body)
   bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))
   bwd_call = bwd_body.call(*compact_args, name=(k.arg.name or "")+"_backward", precompile=k.arg.precompile_backward)
   gb_map = {i: idx for idx, (i, _) in enumerate(grad_bodies)}

tinygrad/mixin/movement.py

@@ -46,16 +46,6 @@ class MovementMixin:
     """
     return prod(self.shape)
 
-  @property
-  def max_shape(self) -> tuple[int, ...]:
-    """The shape with every symbolic dimension replaced by its maximum."""
-    from tinygrad.uop.ops import to_max_shape  # deferred: ops.py imports the mixins
-    return to_max_shape(self.shape)
-
-  def max_numel(self) -> int:
-    """The number of elements in `max_shape`."""
-    return prod(self.max_shape)
-
   def size(self, dim:int|None=None) -> sint|tuple[sint, ...]:
     """
     Returns the size of the tensor. If `dim` is specified, return the length along dimension `dim`. Otherwise return the shape of the tensor.
@@ -550,7 +540,6 @@ class MovementMixin:
     if dims is None: return self.flatten().roll(shifts, 0).reshape(self.shape)
     dims, shifts = tuple(self._resolve_dim(d) for d in make_tuple(dims, 1)), make_tuple(shifts, 1)
     if len(dims) != len(shifts): raise RuntimeError(f"{len(dims)=} != {len(shifts)=}")
-    if 0 in self.shape: return self
     shrink_arg: list[tuple[sint, sint]|None] = [None] * self.ndim
     for d, s in zip(dims, shifts): shrink_arg[d] = (delta:=self.shape[d]-s%self.shape[d], delta+self.shape[d])
     return self.repeat(*tuple(2 if i in dims else 1 for i in range(self.ndim))).shrink(tuple(shrink_arg))

tinygrad/mixin/op.py

@@ -289,12 +289,6 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     if value == 0: return base
     return MovementMixin.pad(X.const_like(True, dtypes.bool), pads).where(base, value)
 
-  def pad_to(self, shape, *args, value:ConstType=0) -> Self:
-    # same mask trick as _pad_constant so the fill survives backends that realize PAD as 0-fill
-    ret = MovementMixin.pad_to(self, shape, *args)
-    if value == 0 or ret is self: return ret
-    return MovementMixin.pad_to(self.const_like(True, dtypes.bool), shape, *args).where(ret, value)
-
   def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self:
     # shrink first for negative pads, then wrap the non-negative remainder
     X = self.shrink(tuple((-smin(pB,0), smin(pA+sh,sh)) for (pB,pA),sh in zip(pX, self.shape)))
@@ -466,7 +460,6 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     """
     assert gradient is not None or self.shape == tuple(), "when no gradient is provided, backward must be called on a scalar tensor"
     if not (self.is_floating_point() and all(t.is_floating_point() for t in targets)): raise RuntimeError("only float Tensors have gradient")
-    if any(t.dtype in dtypes.weaks for t in targets): raise RuntimeError("cannot take gradient wrt a weak Tensor")
     from tinygrad.mixin.gradient import compute_gradient
     if gradient is None: gradient = self.const_like(1.0)
     target_uops = [t._uop for t in targets]

tinygrad/renderer/__init__.py

@@ -35,8 +35,8 @@ class Estimates:
         while len(buf.src) and buf.op is not Ops.PARAM: buf = buf.src[0]
         if buf.op is Ops.PARAM:
           # u.src[0] is INDEX, cap at buffer size for re-reads (e.g. matmul)
-          accessed = mem.get((buf, u.op), 0) + u.src[0].max_numel() * u.src[0].dtype.itemsize * mults
-          mem[(buf, u.op)] = smin(accessed, buf.max_numel() * buf.dtype.itemsize)
+          accessed = mem.get((buf, u.op), 0) + u.src[0].max_numel() * u.src[0].dtype.scalar().itemsize * mults
+          mem[(buf, u.op)] = smin(accessed, buf.max_numel() * buf.dtype.scalar().itemsize)
       if u.op is Ops.RANGE:
         mult_stack.append(mults)
         if u.dtype is not dtypes.void:  # unbounded loop, unknown trip count
@@ -47,9 +47,9 @@ class Estimates:
       elif u.op is Ops.SPECIAL: mults *= cast(sint, u.src[0].ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these
       elif u.op is Ops.PARAM and u.arg.addrspace == AddrSpace.ALU and u.expr == 'core_id': mults *= int(u.vmax) + 1
       elif u.op is Ops.LOAD and u.src[0].addrspace != AddrSpace.REG:
-        lds += u.max_numel() * u.dtype.itemsize * mults
+        lds += u.max_numel() * u.dtype.scalar().itemsize * mults
       elif u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG:
-        lds += u.max_numel() * u.src[1].dtype.itemsize * mults
+        lds += u.max_numel() * u.src[1].dtype.scalar().itemsize * mults
       elif u.op in GroupOp.ALU and u not in excluded:
         flops += (mults * (2 if u.op is Ops.MULACC else 1)) * u.max_numel()
       elif u.op is Ops.WMMA and u not in excluded:

tinygrad/renderer/cstyle.py

@@ -7,6 +7,7 @@ from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, NUM_CPU_
 from tinygrad.dtype import dtypes, DType, AddrSpace, truncate, float_to_bf16
 from tinygrad.renderer import Renderer
 
+
 base_rewrite = PatternMatcher([
   # local/reg buffers
   (UPat(Ops.BUFFER, name="x"), lambda ctx,x: ctx.render_buffer(x)),
@@ -19,21 +20,6 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"),
   (UPat((Ops.ENDIF, Ops.END)), lambda ctx: "}"),
 
-  # const
-  (UPat.cvar("c").cast(dtypes.floats, name="x"), lambda ctx,x,c: None if math.isfinite(v:=c.val) else \
-    f"({ctx.render_cast(x, ctx.nan if math.isnan(v) else ctx.infinity if v > 0 else f'-{ctx.infinity}')})"),
-  (UPat.cvar("c").cast(dtypes.float), lambda ctx,c: f"{c.val}f"),
-  (UPat.cvar("c").cast(dtypes.int64), lambda ctx,c: f"{c.val}l"),
-  (UPat.cvar("c").cast(dtypes.uint64, name="x"), lambda ctx,x,c: f"{truncate[x.dtype](c.val)}ul"),
-  (UPat.cvar("c").cast(dtypes.uint32, name="x"), lambda ctx,x,c: f"{truncate[x.dtype](c.val)}u"),
-  (UPat.cvar("c").cast(dtypes.bool), lambda ctx,c: "1" if c.val else "0"),
-  # consts are rendered to larger type and casted
-  (UPat.cvar("c").cast((*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x,c: f"({ctx.render_cast(x, f'{c.val}f')})"),
-  (UPat.cvar("c").cast((dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x,c: f"({ctx.render_cast(x, f'{c.val}u')})"),
-  (UPat.cvar("c").cast((dtypes.int8, dtypes.int16), name="x"), lambda ctx,x,c: f"({ctx.render_cast(x, str(c.val))})"),
-  # default const render
-  (UPat.cvar("c").cast(), lambda ctx,c: str(c.val)),
-
   # casting
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})" \
     if x.max_numel() > 1 and x.addrspace is AddrSpace.REG else None),
@@ -45,9 +31,25 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.BARRIER), lambda ctx: ctx.barrier),
   (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0]](x.arg[-1])}; /* {(x.src[0]).render()} */"),
 
+  # const
+  (UPat(Ops.CONST, arg=math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x, ctx.infinity)})"),
+  (UPat(Ops.CONST, arg=-math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x, f'-{ctx.infinity}')})"),
+  (UPat(Ops.CONST, dtype=dtypes.floats, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx.nan)})" if math.isnan(x.val) else None),
+  (UPat(Ops.CONST, dtype=dtypes.float, name="x"), lambda ctx,x: f"{x.val}f"),
+  (UPat(Ops.CONST, dtype=dtypes.int64, name="x"), lambda ctx,x: f"{x.val}l"),
+  (UPat(Ops.CONST, dtype=dtypes.uint64, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}ul"),
+  (UPat(Ops.CONST, dtype=dtypes.uint32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}u"),
+  (UPat(Ops.CONST, dtype=dtypes.bool, name="x"), lambda ctx,x: "1" if x.val else "0"),
+  # consts are rendered to larger type and casted
+  (UPat(Ops.CONST, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.val}f')})"),
+  (UPat(Ops.CONST, (dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.val}u')})"),
+  (UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, str(x.val))})"),
+  # default const render
+  (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.val)),
+
   # SHRINK/INDEX
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)),
-  (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var('idx'), UPat.cvar().cast()), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)),
+  (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var('idx'), UPat.cvar()), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)),
   (UPat(Ops.STACK, name="x"),
    lambda ctx,x: f"{ctx.float4.replace('float4', ctx.render_type(x))}" + \
                  f"{ctx.float4_style[0]}{','.join([ctx[y] for y in x.src])}{ctx.float4_style[1]}"),
@@ -105,11 +107,11 @@ def uops_to_dtypes(uops:list[UOp]) -> list[tuple[DType, int]]:
 
 def _wmma_name(u:UOp) -> str:
   # sanitize spaces in DType.name (int8 = "signed char")
-  return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.name}".replace(" ", "_")
+  return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.scalar().name}".replace(" ", "_")
 
 # (name, dims, dtype_in, dtype_out, device, threads, upcast_sizes)
 def wmma_args(uops:list[UOp]):
-  return dedup((_wmma_name(uop), uop.arg[0], uop.arg[1], uop.dtype, *(uop.arg[2:4]),
+  return dedup((_wmma_name(uop), uop.arg[0], uop.arg[1], uop.dtype.scalar(), *(uop.arg[2:4]),
                tuple(uop.src[i].shape[-1] for i in range(3)))
               for uop in uops if uop.op is Ops.WMMA)
 
@@ -161,8 +163,8 @@ class CStyleLanguage(Renderer):
   def render_index(self, x:UOp, buf:UOp, idx:UOp):
     if buf.addrspace == AddrSpace.ALU:
       # this is lane access in C
-      if not (idx.op is Ops.CAST and idx.src[0].op is Ops.CONST): return f"({self[buf]})[{self[idx]}]"
-      return self[buf]+(f"[{idx.src[0].val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.src[0].val]}")
+      if idx.op is not Ops.CONST: return f"({self[buf]})[{self[idx]}]"
+      return self[buf]+(f"[{idx.val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.val]}")
     return f"({self[buf]}+{strip_parens(self[idx]) if idx.arg == Ops.ADD else self[idx]})"
 
   def render_buffer(self, x:UOp):
@@ -180,8 +182,8 @@ class CStyleLanguage(Renderer):
     if addrspace in (AddrSpace.LOCAL, AddrSpace.GLOBAL) or override_ptr:
       suffix = "*"
     if sz > 1:
-      return prefix + self.type_map.get(dtype, dtype.name).replace(" ", "_") + str(sz) + suffix
-    return prefix + self.type_map.get(dtype, dtype.name) + suffix
+      return prefix + self.type_map.get(scalar:=dtype.scalar(), scalar.name).replace(" ", "_") + str(sz) + suffix
+    return prefix + self.type_map.get(scalar:=dtype.scalar(), scalar.name) + suffix
 
   def render_type(self, u:UOp): return self._render_dtype(u.dtype, u.max_numel(), u.addrspace, shape=u._shape)
   def render_access(self, u:UOp):
@@ -208,7 +210,7 @@ class CStyleLanguage(Renderer):
     c: defaultdict[str, int] = defaultdict(int)
     name = "test"
     for u in uops:
-      if u.op in {Ops.NOOP, Ops.GROUP, Ops.CONST}: continue
+      if u.op in {Ops.NOOP, Ops.GROUP}: continue
       if u.op == Ops.STACK and len(u.src) == 0: continue
       if u.op is Ops.AFTER:
         r[u] = r[u.src[0]]
@@ -226,7 +228,7 @@ class CStyleLanguage(Renderer):
       if u.op is Ops.SPECIAL: r[u] = u.arg
       elif u.op is Ops.RANGE: r[u] = f"{axis_letters[u.arg[-1]]}idx"+range_str(u)
       else:
-        prefix = {Ops.WMMA: "wmma", Ops.BUFFER: "buf", Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.STACK: "cast",
+        prefix = {Ops.WMMA: "wmma", Ops.CONST: "const", Ops.BUFFER: "buf", Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.STACK: "cast",
                   Ops.INDEX: "bidx", Ops.LOAD: "val"}.get(u.op, "alu")
         r[u] = f"{prefix}{c[prefix]}"
 
@@ -234,8 +236,7 @@ class CStyleLanguage(Renderer):
       assert l is not None, f"failed to render {u.op} {u.dtype} {[(x.op,x.dtype) for x in u.src]} {u.arg}"
 
       if u.op in {Ops.ENDIF, Ops.END}: depth -= 1
-      if (u.op is not Ops.CAST or u.max_numel() == 1) and ((u.op is Ops.CAST and u.src[0].op is Ops.CONST) or \
-        u.op in {Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
+      if (u.op is not Ops.CAST or u.max_numel() == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
         (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \
         (u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \
         (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
@@ -319,7 +320,8 @@ class OpenCLRenderer(CStyleLanguage):
   string_rewrite = PatternMatcher([
     (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
     # bfloat16 constants need to be rendered as their bit pattern since bf16 is stored as ushort
-    (UPat.cvar("c").cast(dtypes.bfloat16), lambda ctx,c: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(c.val)))[0] >> 16)}u"),
+    (UPat(Ops.CONST, dtypes.bfloat16, name="x"),
+      lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.val)))[0] >> 16)}u"),
     # load/store image (OpenCL)
     (UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), lambda ctx,buf,idx_y,idx_x: f"IMAGE<{ctx[buf]}, {ctx[idx_y]}, {ctx[idx_x]}>"),
     (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))),
@@ -470,7 +472,7 @@ class CUDARenderer(CStyleLanguage):
 class NVCCRenderer(CUDARenderer):
   def __init__(self, target:Target): super().__init__(target, use_nvcc=True)
 
-def fp8_index(dtype: DType): return (dtypes.fp8e4m3, dtypes.fp8e5m2).index(dtype)
+def fp8_index(dtype: DType): return (dtypes.fp8e4m3, dtypes.fp8e5m2).index(dtype.scalar())
 def _ocml(op): return lambda x,dtype: f"__ocml_{op}_f{ {dtypes.half:16, dtypes.double:64}.get(dtype, 32)}({x})"
 
 class HIPRenderer(CStyleLanguage):
@@ -493,9 +495,10 @@ class HIPRenderer(CStyleLanguage):
         (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]},"
           f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[0][2] == 128 else None),
         (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"),
-        (UPat.cvar("c").cast(dtypes.fp8s, name="x"), lambda ctx,x,c:
-          f"f32_to_fp8({ctx.nan if math.isnan(v:=c.val) else ctx.infinity if v == math.inf else f'-{ctx.infinity}' if v == -math.inf else f'{v}f'},"
-          f" {fp8_index(x.dtype)})"),
+        (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(x.val) else None),
+        (UPat(Ops.CONST, dtypes.fp8s, arg=math.inf, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.infinity}, {fp8_index(x.dtype)})"),
+        (UPat(Ops.CONST, dtypes.fp8s, arg=-math.inf, name="x"), lambda ctx,x: f"f32_to_fp8(-{ctx.infinity}, {fp8_index(x.dtype)})"),
+        (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({x.val}f, {fp8_index(x.dtype)})"),
         (UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",),
           lambda ctx,x: f"f32_to_fp8({ctx[x.src[0]]}, {fp8_index(x.dtype)})"),
         (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",),
@@ -536,21 +539,21 @@ class HIPRenderer(CStyleLanguage):
     prefix, ockl = [], []
     type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16", dtypes.fp8e4m3: "_fp8_fp8", dtypes.fp8e5m2: "_bf8_bf8" }
     used_dtypes = uops_to_dtypes(uops)
-    if any(u.op is Ops.CAST and u.src[0].op is Ops.CONST and not math.isfinite(u.src[0].val) for u in uops):
+    if any(u.op is Ops.CONST and not math.isfinite(u.val) for u in uops):
       prefix += ["#define INFINITY (__builtin_inff())", "#define NAN (__builtin_nanf(\"\"))"]
     if any(u.op is Ops.SPECIAL for u in uops):
       prefix.append("typedef long unsigned int size_t;")
       ockl = [(f"__ockl_get_{name}", "unsigned int", "size_t", "const") for name in ["local_id", "group_id", "local_size"]]
     ocml_ops = {Ops.EXP2: ("exp2", "pure"), Ops.LOG2: ("log2", "pure"), Ops.SQRT: ("sqrt", "const"), Ops.SIN: ("sin", ""), Ops.TRUNC: ("trunc", "")}
     ocml = [(f"__ocml_{ocml_ops[op][0]}_f{dt.bitsize}", dt.name, dt.name, ocml_ops[op][1])
-      for op, dt in dedup((u.op, u.dtype) for u in uops) if op in ocml_ops and dt in (dtypes.half, dtypes.float, dtypes.double)]
+      for op, dt in dedup((u.op, u.dtype.scalar()) for u in uops) if op in ocml_ops and dt in (dtypes.half, dtypes.float, dtypes.double)]
     if any(dt == dtypes.bfloat16 for dt, _ in used_dtypes):
       prefix.append(f"typedef {'__bf16' if self.is_cdna4(self.target.arch) else 'unsigned short'} hip_bfloat16;")
     if any(dt == dtypes.half for dt, _ in used_dtypes): prefix.append("#define half _Float16")
     if any(dt in dtypes.fp8s for dt, _ in used_dtypes):
       prefix += ["typedef unsigned char hip_bf8;", "typedef unsigned char hip_fp8;"]
     if any((u.op is Ops.CAST and u.dtype in dtypes.fp8s and u.src[0].dtype == dtypes.float) or
-           (u.op is Ops.CAST and u.src[0].op is Ops.CONST and u.dtype in dtypes.fp8s) for u in uops):
+           (u.op is Ops.CONST and u.dtype in dtypes.fp8s) for u in uops):
       prefix.append("""static inline __attribute__((device)) unsigned char f32_to_fp8(float v, int is_bf8) {
   v = (((*(unsigned*)&v)&0x7F800000)!=0x7F800000)?__builtin_amdgcn_fmed3f(v,is_bf8?57344.0f:448.0f,is_bf8?-57344.0f:-448.0f) : v;
   return (unsigned char)(is_bf8?__builtin_amdgcn_cvt_pk_bf8_f32(v,v,0,false):__builtin_amdgcn_cvt_pk_fp8_f32(v,v,0,false));\n}""")

tinygrad/renderer/isa/x86.py

@@ -165,16 +165,16 @@ def scratch_buffer(elem_dt:DType, count:int, slot:int) -> UOp:
   return UOp.placeholder((count,), elem_dt, slot, AddrSpace.LOCAL)
 
 def gated_load(ctx, addr:UOp, alt:UOp, gate:UOp, x:UOp):
-  local = scratch_buffer(addr.src[0].dtype, x.max_numel(), next(ctx))
-  local_idx = local.index(UOp.cconst(0, dtypes.int32), dtype=dtypes.uint64)
+  local = scratch_buffer(addr.src[0].dtype.scalar(), x.max_numel(), next(ctx))
+  local_idx = local.index(UOp.const(0, dtypes.int32), dtype=dtypes.uint64)
   # the selected address is a 64bit value, the AFTER orders the load after the scratch store and carries the element dtype for the encoder
   sel = gate.where(addr.replace(dtype=dtypes.uint64), local_idx)
   ptr = UOp(Ops.AFTER, addr.dtype, (sel, (local_idx if x.max_numel() == 1 else local).store(alt)))
   return ptr.load(dtype=x.dtype)
 
 def gated_store(addr:UOp, gate:UOp, val:UOp):
-  local = scratch_buffer(addr.src[0].dtype, val.max_numel(), -1)
-  sel = gate.where(addr.replace(dtype=dtypes.uint64), local.index(UOp.cconst(0, dtypes.int32), dtype=dtypes.uint64))
+  local = scratch_buffer(addr.src[0].dtype.scalar(), val.max_numel(), -1)
+  sel = gate.where(addr.replace(dtype=dtypes.uint64), local.index(UOp.const(0, dtypes.int32), dtype=dtypes.uint64))
   return UOp(Ops.AFTER, addr.dtype, (sel,)).store(val)
 
 # legalize the new style graph for isel. NOTE: this runs after the spec is verified, some of these rewrites violate it
@@ -195,7 +195,7 @@ pre_isel_matcher = PatternMatcher([
   # if gate in scalar int cmove is not a comparison need to add one to set the flag
   # NOTE: the 0 is int so the bool gate zero-extends and compares as int (a byte compare renders different kernels)
   (UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")),
-   lambda m,a,b: m.ne(UOp.cconst(0, dtypes.int)).where(a,b) if m.op not in GroupOp.Comparison else None),
+   lambda m,a,b: m.ne(UOp.const(0, dtypes.int)).where(a,b) if m.op not in GroupOp.Comparison else None),
 ])
 
 # ***** X86 registers *****
@@ -221,15 +221,15 @@ reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"},
 
 # ***** X86 instruction selection *****
 def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.INDEX else s
-def lane(x:UOp, i:int) -> int: return s.src[1].src[0].val if (s:=x.src[i]).op is Ops.INDEX else 0
+def lane(x:UOp, i:int) -> int: return s.src[1].val if (s:=x.src[i]).op is Ops.INDEX else 0
 def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt]
 def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, dt, arg=X86Ops.DEFINE, tag=None if reg is None else (reg,))
-def imm(dt:DType, v:int) -> UOp: return UOp.cconst(truncate[dt](v), dt).rtag()
+def imm(dt:DType, v:int) -> UOp: return UOp.const(truncate[dt](v), dt).rtag()
 def to_imm(c:UOp) -> UOp|None:
-  if not (c.op is Ops.CAST and (v:=c.src[0]).op is Ops.CONST): return None
-  if c.dtype is dtypes.int64: return imm(dtypes.int32, v.val) if not v.overflows(dtypes.int32) else None
-  if c.dtype is dtypes.uint64: return imm(dtypes.uint32, v.val) if not v.overflows(dtypes.uint32) else None
-  if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, v.val)
+  if c.op is not Ops.CONST: return None
+  if c.dtype is dtypes.int64: return imm(dtypes.int32, c.val) if not c.overflows(dtypes.int32) else None
+  if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.val) if not c.overflows(dtypes.uint32) else None
+  if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, c.val)
   return None
 def cmp(x:UOp) -> UOp:
   if x.src[0].dtype is dtypes.float32: return x.ins(X86Ops.VUCOMISS, dtype=dtypes.void)
@@ -237,7 +237,7 @@ def cmp(x:UOp) -> UOp:
   return x.ins(X86Ops.CMP, dtype=dtypes.void) if (i:=to_imm(x.src[1])) is None else x.ins(X86Ops.CMPi, dtype=dtypes.void, src=(x.src[0], i))
 def vcmp(x:UOp) -> UOp:
   v = imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op])
-  if x.dtype is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.max_numel() == 1 else X86Ops.VCMPPS, src=x.src + (v,))
+  if x.dtype.scalar() is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.max_numel() == 1 else X86Ops.VCMPPS, src=x.src + (v,))
   return x.ins(X86Ops.VCMPSD if x.max_numel() == 1 else X86Ops.VCMPPD, src=x.src + (v,))
 
 # vinsertps xmm2, xmm0, xmm1, imm
@@ -252,7 +252,7 @@ def vinsertps(x:UOp) -> UOp:
 # vpinsq xmm2, xmm0, rax, imm
 # inserts element in rax into any position in xmm0, result is written to xmm2 according to imm
 def vpins(x:UOp) -> UOp:
-  op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.itemsize]
+  op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.scalar().itemsize]
   return functools.reduce(lambda ret,i: x.ins(op, src=(ret, x.src[i], imm(dtypes.uint8, i))), range(len(x.src)), def_reg(x.dtype))
 
 # we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg
@@ -289,9 +289,8 @@ def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]:
   # buffers are indexed by element, everything else (the stack pointer) by byte
   scale = base.dtype.itemsize if base.op in {Ops.PARAM, Ops.BUFFER, Ops.AFTER} else 1
   sz = imm(dtypes.uint8, base.dtype.itemsize)
-  if idx.op is Ops.ADD and (c:=idx.src[1]).op is Ops.CAST and c.src[0].op is Ops.CONST:
-    return (base, _cast(idx.src[0]), _disp(c.src[0].val * scale), sz)
-  if idx.op is Ops.CAST and idx.src[0].op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.src[0].val * scale), sz)
+  if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return (base, _cast(idx.src[0]), _disp(idx.src[1].val * scale), sz)
+  if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.val * scale), sz)
   return (base, _cast(idx), _disp(0), sz)
 
 def abi(ctx:IselContext, x:UOp) -> UOp|None:
@@ -354,7 +353,7 @@ isel_matcher = PatternMatcher([
   # cast of void is a noop
   (UPat.var("y").cast(name="x"), lambda y,x: y if y.dtype == dtypes.void else None),
   # range is lowered to acc, cmp, jmp after regalloc
-  (UPat(Ops.RANGE, src=(UPat.cvar("c").cast(),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(x.dtype, c.val),) + x.src[1:])),
+  (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.val),) + x.src[1:])),
   (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None),
   # really all a backedge END is is an IF with a tag referencing the RANGE start label
   (UPat(Ops.END, src=(UPat(), UPat(), UPat(GroupOp.Comparison, name="cond")), name="x"),
@@ -368,10 +367,10 @@ isel_matcher = PatternMatcher([
   # function abi constraints
   (UPat((Ops.PARAM, Ops.SPECIAL), name="x"), abi),
   # constants that can't be immediates, move them to registers
-  (UPat.cvar("c").cast(dtypes.int64s, name="x"), lambda c,x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, c.val),)) if not x.tag else None),
-  (UPat.cvar("c").cast(dtypes.ints+(dtypes.bool,), name="x"), lambda c,x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, c.val),)) if not x.tag else None),
-  (UPat.cvar("c").cast(dtypes.floats, name="x"), lambda c,x:
-   UOp.cconst(struct.unpack((dt:=to_int(x.dtype)).fmt, struct.pack(x.dtype.fmt, c.val))[0], dt).bitcast(x.dtype) if not x.tag else None),
+  (UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.val),)) if not x.tag else None),
+  (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.val),)) if not x.tag else None),
+  (UPat.cvar("x", dtypes.floats), lambda x:
+   UOp.const(struct.unpack((dt:=to_int(x.dtype)).fmt, struct.pack(x.dtype.fmt, x.val))[0], dt).bitcast(x.dtype) if not x.tag else None),
   # conditional moves that use masks NOTE: these currently assume a mask producing cmp exists
   (UPat.var("m").where(UPat.var("a", dtypes.int8s+dtypes.int16s+dtypes.int32s+(dtypes.int64,)), UPat.var("b")), lambda m,a,b:
    a.ins(X86Ops.VPBLENDVB, src=(b, a, m.replace(dtype=m.src[0].dtype))) if a.max_numel() > 1 else None),
@@ -381,7 +380,7 @@ isel_matcher = PatternMatcher([
    a.ins(X86Ops.VBLENDVPD, src=(b, a, m.replace(dtype=m.src[0].dtype)))),
   # in this case we have a mask producing comparison whose user expects a bool, so we convert to bool
   (UPat(GroupOp.Comparison, dtypes.bool, (UPat.var("y", (dtypes.float32, dtypes.float64)), UPat()), name="x"), lambda y,x:
-   UOp(Ops.AND, src=(x.replace(dtype=y.dtype).bitcast(dt:=to_int(y.dtype)), UOp.cconst(1, dt))).f(Ops.NOOP, dtype=dtypes.bool)),
+   UOp(Ops.AND, src=(x.replace(dtype=y.dtype).bitcast(dt:=to_int(y.dtype)), UOp.const(1, dt))).f(Ops.NOOP, dtype=dtypes.bool)),
   # conditional moves that use flags
   (UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.sints), UPat()), name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b:
    a.ins(X86Ops.CMOVL, src=(b, a, cmp(m)))),
@@ -421,15 +420,15 @@ isel_matcher = PatternMatcher([
   (UPat(Ops.STACK, dtypes.float32, name="x"), vinsertps),
   (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), vpins),
   # INDEX on a vector register value extracts a single element
-  (UPat.var("y", dtypes.int8s+(dtypes.bool,)).index(UPat.cvar("c").cast(), name="x"),
+  (UPat.var("y", dtypes.int8s+(dtypes.bool,)).index(UPat.cvar("c"), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.int16s).index(UPat.cvar("c").cast(), name="x"),
+  (UPat.var("y", dtypes.int16s).index(UPat.cvar("c"), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.int32s).index(UPat.cvar("c").cast(), name="x"),
+  (UPat.var("y", dtypes.int32s).index(UPat.cvar("c"), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.int64s).index(UPat.cvar("c").cast(), name="x"),
+  (UPat.var("y", dtypes.int64s).index(UPat.cvar("c"), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.floats).index(UPat.cvar("c").cast(), name="x"),
+  (UPat.var("y", dtypes.floats).index(UPat.cvar("c"), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, c.val * x.dtype.itemsize))) if _is_vec_xmm(y) else None),
   # packed bitwise
   ((UPat() & UPat()).named("x"), lambda x: x.ins(X86Ops.VPAND) if x.max_numel() > 1 else None),
@@ -454,19 +453,15 @@ isel_matcher = PatternMatcher([
   # scalar int binary
   ((UPat(dtype=dtypes.ints).alu(Ops.CDIV, UPat())).named("x"), idiv),
   # scalar int binary with immediate
-  (UPat.var("a", dtypes.ints) << UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.val)))),
-  (UPat.var("a", dtypes.uints) >> UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.val)))),
-  (UPat.var("a", dtypes.sints) >> UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.val)))),
-  (UPat.var("a", dtypes.ints) + UPat.cvar().cast(name="c"), lambda a,c: a.ins(X86Ops.ADDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  (UPat.var("a", dtypes.ints) * UPat.cvar().cast(name="c"), lambda a,c: a.ins(X86Ops.IMULi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.cvar().cast(name="c"),
-   lambda a,c: a.ins(X86Ops.ANDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  (UPat.var("a", dtypes.ints+(dtypes.bool,)) | UPat.cvar().cast(name="c"),
-   lambda a,c: a.ins(X86Ops.ORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.cvar().cast(name="c"),
-   lambda a,c: a.ins(X86Ops.XORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.cvar().cast(name="c"))),
-   lambda a,c: a.ins(X86Ops.SUBi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints) << UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.val)))),
+  (UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.val)))),
+  (UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.val)))),
+  (UPat.var("a", dtypes.ints) + UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ADDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints) * UPat.cvar("c"), lambda a,c: a.ins(X86Ops.IMULi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ANDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) | UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.cvar("c"), lambda a,c: a.ins(X86Ops.XORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
+  (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.cvar("c"))), lambda a,c: a.ins(X86Ops.SUBi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
   # scalar int binary with register
   ((UPat(dtype=dtypes.ints) << UPat()).named("x"), lambda x: shift(x, X86Ops.SHL)),
   ((UPat(dtype=dtypes.uints) >> UPat()).named("x"), lambda x: shift(x, X86Ops.SHR)),
@@ -577,7 +572,7 @@ def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
   if x.dtype is dtypes.void: return (label, [label])
   else:
     acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + x.src[1:])
-    cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CAST else X86Ops.CMP, src=(acc, x.src[0]))
+    cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CONST else X86Ops.CMP, src=(acc, x.src[0]))
     jump_out = UOp(Ops.INS, arg=X86Ops.JGE, src=(cmp,), tag=f".LOOP_OUT_{loop_label}")
     ctx.loop_label[acc] = loop_label
     return (acc, [acc, label, cmp, jump_out])
@@ -596,7 +591,7 @@ def lower_loop(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
 # final rewrite to match the isa spec
 post_regalloc_matcher = PatternMatcher([
   # rewrite FRAME_INDEX to IMM now that the stack size is known
-  (UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=UOp.cconst(ctx.stack_size + x.tag, x.dtype), [nx])),
+  (UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=x.const_like(ctx.stack_size + x.tag), [nx])),
   # expand the cmp here so we can preserve rng src edge to get label from ctx
   (UPat(Ops.INS, arg=X86Ops.LOOP_CMP, name="x"), lower_loop),
   # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound
@@ -619,7 +614,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     rm = cast(Register, greg(rm_uop)).index
     idx = cast(Register, greg(idx_uop)).index if idx_uop is not None and greg(idx_uop) is not None else 4
     # for a memory operand the rm size is the element size from the address, otherwise it's the size of the value in the register
-    rm_sz = sz_uop.src[0].val if sz_uop is not None else rm_uop.dtype.itemsize
+    rm_sz = sz_uop.val if sz_uop is not None else rm_uop.dtype.itemsize
     reg_sz = reg_uop.dtype.itemsize if reg_uop is not None else 0
     sz = reg_sz or rm_sz
 
@@ -652,10 +647,10 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     # 0b10 -- signals memory access with 32bit displacement
     # 0b11 -- signals no memory access
     if disp_uop is not None:
-      assert disp_uop.op is Ops.CAST, "displacement must be a literal"
+      assert disp_uop.op is Ops.CONST, "displacement must be a constant"
       assert disp_uop.dtype in (dtypes.int8, dtypes.int32), "displacement can only be 1 or 4 byte signed int"
       # rbp/r13 always require a displacement
-      if disp_uop.src[0].val != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10
+      if disp_uop.val != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10
       else: mod = 0b00
     else: mod = 0b11
     # x 0b0 and idx 0b100 means rsp which means no index exists
@@ -669,10 +664,10 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     # DISP byte
     if mod == 0b01 or mod == 0b10:
       assert disp_uop is not None
-      inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.src[0].val)
+      inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.val)
     # IMM byte
     if imm_uop is not None:
-      if imm_uop.op is Ops.CAST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.src[0].val)
+      if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.val)
       elif isinstance(greg(imm_uop), Register): inst += bytes([(greg(imm_uop).index & 0b1111) << 4 | 0b0000])
     return inst
 
@@ -682,13 +677,13 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
   if x.arg in X86GroupOp.WriteMem:
     if len(x.src) > 4: address, rest = x.src[:4], x.src[4:]
     else: address, rest = (x, None, None, None), x.src
-    imm_uop = rest[:1] if rest and rest[0].op is Ops.CAST else (None,)
+    imm_uop = rest[:1] if rest and rest[0].op is Ops.CONST else (None,)
     return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *imm_uop))
 
   if x.arg in X86GroupOp.Rm1st:
     if len(x.src) > 3: address, rest = x.src[:4], x.src[4:]
     else: address, rest = (x.src[0], None, None, None), x.src[1:]
-    imm_uop = rest[:1] if rest and rest[0].op is Ops.CAST else (None,)
+    imm_uop = rest[:1] if rest and rest[0].op is Ops.CONST else (None,)
     return _encode(x, *address, *(None, *imm_uop)) if reg is None else _encode(None, *address, *(x if sel else None, *imm_uop))
 
   if x.arg in X86GroupOp.Rm2nd:
@@ -706,7 +701,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
 encodings = {
   # moves
   X86Ops.MOVABS: lambda x:
-   bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].src[0].val),
+   bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].val),
   X86Ops.MOV: lambda x: encode(x, 0x8B), X86Ops.MOVi: lambda x: encode(x, 0xC7, reg=0),
   X86Ops.MOVm: lambda x: encode(x, 0x89), X86Ops.LEA: lambda x: encode(x, 0x8D),
   X86Ops.VMOVSS: lambda x: encode(x, 0x10, pp=2, sel=1), X86Ops.VMOVSSm: lambda x: encode(x, 0x11, pp=2, sel=1),
@@ -729,8 +724,8 @@ encodings = {
   X86Ops.VCVTPS2PD: lambda x: encode(x, 0x5A, pp=0, sel=1), X86Ops.VCVTPD2PS: lambda x: encode(x, 0x5A, pp=1, sel=1),
   X86Ops.VCVTTPS2DQ: lambda x: encode(x, 0x5B, pp=2, sel=1), X86Ops.VCVTTPD2DQ: lambda x: encode(x, 0xE6, pp=1, sel=1),
   # the int src is the 2nd src (the rm field), if it was folded into a memory operand its width is the element size of the address
-  X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=(x.src[4].src[0].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
-  X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=(x.src[4].src[0].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
+  X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=(x.src[4].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
+  X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=(x.src[4].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
   X86Ops.VCVTTSS2SI: lambda x: encode(x, 0x2C, pp=2, sel=1, we=x.dtype.itemsize == 8),
   X86Ops.VCVTTSD2SI: lambda x: encode(x, 0x2C, pp=3, sel=1, we=x.dtype.itemsize == 8),
   # int division
@@ -845,10 +840,10 @@ class X86Renderer(ISARenderer):
     def _format_op(x:UOp) -> str: return f"    {(o[7:-1] if (o:=str(x.arg))[-1] in ('i', 'm') else o[7:]).lower():7s}"
     def _format_operands(x:UOp) -> str:
       def _format(src:tuple[UOp, ...]) -> list[str]:
-        return [str(s.src[0].val) if s.op is Ops.CAST else reg_strs[o].get(s.dtype.itemsize, o) if \
+        return [str(s.val) if s.op is Ops.CONST else reg_strs[o].get(s.dtype.itemsize, o) if \
                 (o:=str(greg(s))) in reg_strs else o for s in src if greg(s) is not None]
       def _mem_adress(base:UOp, idx:UOp, disp:UOp, sz:UOp) -> list[str]:
-        return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.src[0].val}" if greg(idx) else "") + (f" + {d}" if (d:=disp.src[0].val) else "") + "]"]
+        return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.val}" if greg(idx) else "") + (f" + {disp.val}" if disp.val else "") + "]"]
 
       if len(x.src) > 4 and x.arg in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:4]) + _format(x.src[4:])
       elif len(x.src) > 3 and x.arg in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:4]) + _format(x.src[4:])

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("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),
+  (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),
 
   # load/store
   (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.var("alt"), UPat.var("mask")), name="x"),
@@ -165,7 +165,7 @@ class LLVMRenderer(Renderer):
     local_args: list[str] = []
     name = "test"
     for u in uops:
-      if u.op in {Ops.NOOP, Ops.GROUP, Ops.CONST}: continue
+      if u.op in {Ops.NOOP, Ops.GROUP}: continue
       if u.op is Ops.AFTER:
         r[u] = r[u.src[0]]
         continue
@@ -185,7 +185,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.CAST and u.src[0].op is Ops.CONST: r[u] = lconst(u.src[0].val, u.dtype)
+      elif u.op is Ops.CONST: r[u] = lconst(u.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

@@ -145,7 +145,7 @@ class NIRRenderer(Renderer):
   ])
 
   def_rewrite = PatternMatcher([
-    (UPat.cvar("c").cast(name="x"), lambda ctx,x,c: nimm(ctx.b, c.val, x.dtype)),
+    (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.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,17 +186,16 @@ 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].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].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, Ops.CONST} or (u.op is Ops.STACK and len(u.src) == 0): pass
+      if u.op in {Ops.NOOP, Ops.GROUP} 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].src[0].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].val)
       elif u.op is Ops.AFTER:
         self.r[u] = self.r[u.src[0]]
       elif u.op == Ops.SINK:

tinygrad/renderer/ptx.py

@@ -64,7 +64,7 @@ def render_wmma(ctx: "PTXRenderer", wmma: UOp):
 
   for src, regs in zip(wmma.src, ctx.wmma_r):
     for i, reg in enumerate(regs): # pack input and acc registers
-      if (elems_per_reg := 4 // src.dtype.itemsize) == 1: yield f"mov.b32 {reg}, {ctx.r[src][i]};"
+      if (elems_per_reg := 4 // src.dtype.scalar().itemsize) == 1: yield f"mov.b32 {reg}, {ctx.r[src][i]};"
       else: yield f"mov.b32 {reg}, {{{', '.join(ctx.r[src][i * elems_per_reg : (i+1) * elems_per_reg])}}};"
 
   dt_map_in, dt_map_out = {dtypes.float: "tf32", dtypes.half: "f16"}, {dtypes.float: "f32", dtypes.half: "f16"}
@@ -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("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.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(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];"),
@@ -101,17 +101,17 @@ string_rewrite = PatternMatcher([
      if loc.addrspace == AddrSpace.REG else None),
   (UPat(Ops.STORE, src=(UPat((Ops.INDEX, Ops.SHRINK), name="loc"), UPat.var("var"))),
    lambda ctx, loc, var: f"st.{mem_type(loc)}" + \
-    f"{f'.v{cnt}' if ((cnt:=var.max_numel())>1) else ''}.{ctx.mem_types[var.dtype]} " + \
+    f"{f'.v{cnt}' if ((cnt:=var.max_numel())>1) else ''}.{ctx.mem_types[var.dtype.scalar()]} " + \
     f"[{ctx.r[loc]}+0], {('{' + ', '.join(ctx.r[var]) + '}') if var.max_numel() > 1 else ctx.r[var]};"),
   (UPat(Ops.LOAD, name="x", src=(UPat((Ops.INDEX, Ops.SHRINK), name="loc"), UPat.var("alt"), UPat.var("gate"))),
     lambda ctx, x, loc, alt, gate: flatten([
-    [f"mov.{ctx.mem_types[x.dtype]} {v}, {render_val(0, x.dtype)};" for v in ctx.r[x]],
-    [f"@{ctx.r[gate]} ld.{mem_type(loc)}.v{x.max_numel()}.{ctx.mem_types[x.dtype]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];"]
+    [f"mov.{ctx.mem_types[x.dtype.scalar()]} {v}, {render_val(0, x.dtype.scalar())};" for v in ctx.r[x]],
+    [f"@{ctx.r[gate]} ld.{mem_type(loc)}.v{x.max_numel()}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];"]
   ]) if alt.max_numel() > 1 else [
-    f"@{ctx.r[gate]} ld.{mem_type(loc)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];",
-    f"@!{ctx.r[gate]} mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {ctx.r[alt]};"]),
+    f"@{ctx.r[gate]} ld.{mem_type(loc)}.{ctx.mem_types[x.dtype.scalar()]} {ctx.r[x]}, [{ctx.r[loc]}+0];",
+    f"@!{ctx.r[gate]} mov.b{ctx.types[x.dtype.scalar()][1:]} {ctx.r[x]}, {ctx.r[alt]};"]),
   (UPat(Ops.LOAD, name="x", src=(UPat((Ops.INDEX, Ops.SHRINK), name="loc"),)),
-    lambda ctx, x, loc: f"ld.{mem_type(loc)}.v{x.max_numel()}.{ctx.mem_types[x.dtype]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \
+    lambda ctx, x, loc: f"ld.{mem_type(loc)}.v{x.max_numel()}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \
      if x.max_numel() > 1 else f"ld.{mem_type(loc)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"),
   # simple
   (UPat(Ops.BUFFER, name="x"), lambda ctx, x: [] if x.addrspace == AddrSpace.REG else [
@@ -186,7 +186,7 @@ class PTXRenderer(Renderer):
 
     name = "test"
     for u in uops:
-      if u.op in {Ops.NOOP, Ops.GROUP, Ops.CONST}: continue
+      if u.op in {Ops.NOOP, Ops.GROUP}: continue
       if u.op is Ops.AFTER:
         self.r[u] = self.r[u.src[0]]
         continue
@@ -197,26 +197,26 @@ class PTXRenderer(Renderer):
         r[u] = [cast(str,r[x]) for x in u.src]
         continue
       if u.op is Ops.BUFFER and u.addrspace == AddrSpace.REG:
-        r[u] = [ssa("reg", u, self.types[u.dtype]) for _ in range(u.max_numel())]
+        r[u] = [ssa("reg", u, self.types[u.dtype.scalar()]) for _ in range(u.max_numel())]
         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 not (u.src[1].op is Ops.CAST and u.src[1].src[0].op is Ops.CONST):
+        if u.op is not Ops.LOAD and u.src[1].op is not 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].src[0].val]
+        r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].val]
         continue
       if u.op is Ops.SPECIAL: r[u] = "%" + u.arg
       elif u.op is Ops.LOAD:
-        r[u] = [ssa('val', dtype=self.types[u.dtype]) for _ in range(u.max_numel())] if u.max_numel() > 1 else ssa('val', u)
+        r[u] = [ssa('val', dtype=self.types[u.dtype.scalar()]) for _ in range(u.max_numel())] if u.max_numel() > 1 else ssa('val', u)
       elif u.op is Ops.PARAM: bufs.append((f"data{u.arg.slot}", u))
       elif u.op is Ops.WMMA:
         # registers for packing/unpacking input and acc
-        self.wmma_r = [[ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[0]]), 4 // u.src[0].dtype.itemsize)],
-                       [ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[1]]), 4 // u.src[0].dtype.itemsize)],
-                       [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())]
+        self.wmma_r = [[ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[0]]), 4 // u.src[0].dtype.scalar().itemsize)],
+                       [ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[1]]), 4 // u.src[0].dtype.scalar().itemsize)],
+                       [ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.dtype.scalar().itemsize)]]
+        r[u] = [ssa("wmma", dtype=self.types[u.dtype.scalar()]) 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.BUFFER: ("local", "u64"), Ops.INDEX: ("bidx", "u64"), Ops.SHRINK: ("bidx", "u64"),
+        Ops.CONST: ("const", None), 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/renderer/wgsl.py

@@ -50,9 +50,8 @@ wgsl_matcher = PatternMatcher([
   (UPat.store(UPat.var("b"), UPat.var("var"), name="s"), lambda b,var,s: packed_store(b,var) if is_packed(s) else None),
   (UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<<b.cast(dtypes.uint32)).bitcast(a.dtype) if b.dtype!=dtypes.uint32 else None),
   (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
-  # fix nan check: 'a != a -> is_nan()'. the decomp rewrites (a != a).logical_not() to CMPEQ, so match both forms
-  (UPat.var("a", dtypes.floats) != UPat.var("a"), is_nan),
-  (UPat.var("a", dtypes.floats).alu(Ops.CMPEQ, UPat.var("a")), lambda a: is_nan(a).ne(True)),
+  # fix nan check: 'a != a -> is_nan()'
+  (UPat.var("a") != UPat.var("a"), is_nan),
   ])
 
 class WGSLRenderer(CStyleLanguage):
@@ -69,10 +68,10 @@ class WGSLRenderer(CStyleLanguage):
 
   string_rewrite = PatternMatcher([
     (UPat(Ops.NEG, dtypes.uints, src=(UPat.var('x'))), lambda ctx,x: f"(0-{ctx[x]})"),
-    (UPat.cvar("c").cast(dtypes.bool), lambda c: "true" if c.val else "false"),
-    (UPat.cvar("c").cast((dtypes.uchar, dtypes.ushort, dtypes.uint32)),
-     lambda c: f"bitcast<u32>({c.val})" if c.val < 0 else f"{c.val&0xFFFFFFFF}u"),
-    (UPat.cvar("c").cast(dtypes.int32, name="x"), lambda ctx,x,c: f"{truncate[x.dtype](c.val)}"),
+    (UPat.cvar("x", dtype=dtypes.bool), lambda x: "true" if x.val else "false"),
+    (UPat(Ops.CONST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), name="x"),
+     lambda x: f"bitcast<u32>({x.val})" if x.val < 0 else f"{x.val&0xFFFFFFFF}u"),
+    (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}"),
     (UPat(Ops.BUFFER, name="x"), lambda ctx,x:
      f"var{'<workgroup>' if x.addrspace == AddrSpace.LOCAL else ''} {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"),
     (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)),

tinygrad/runtime/autogen/am/__init__.py

@@ -8,7 +8,7 @@ am_src="https://github.com/ROCm/ROCK-Kernel-Driver/archive/33970e1351f5e51102960
 rocm_src="https://github.com/ROCm/rocm-systems/archive/cccc350dc620e61ae2554978b62ab3532dc10bd9.tar.gz"
 AMD, AMDINC = "{}/drivers/gpu/drm/amd", "{}/drivers/gpu/drm/amd/include"
 inc, kern_rules = ["-include", "stdint.h"], [(r'le32_to_cpu', ''),]
-fw_src="https://gitlab.com/kernel-firmware/linux-firmware/-/archive/0a6871b19abf5d6e024b5d208b101ae53e7fa0de/0a6871b19abf5d6e024b5d208b101ae53e7fa0de.tar.gz"
+fw_src="https://gitlab.com/kernel-firmware/linux-firmware/-/archive/1e2c15348485939baf1b6d1f5a7a3b799d80703d/1e2c15348485939baf1b6d1f5a7a3b799d80703d.tar.gz"
 pmc_src="https://raw.githubusercontent.com/ROCm/rocm-systems/cccc350dc620e61ae2554978b62ab3532dc10bd9/projects/rocprofiler-compute/src/rocprof_compute_soc/profile_configs/counter_defs.yaml"
 
 reg_files = {

tinygrad/runtime/autogen/am/fw.py

@@ -1,82 +1,81 @@
 hashes = {
-  'psp_13_0_0_sos.bin': '4a51299f6d0a15bbba9694419f7891e6accc01dbd2dd67c06add7bfd75a45ac6',
+  'psp_13_0_0_sos.bin': 'b5592f46885585b935e013f46c949db8ff2f15c0b346caf70e7fcd2776623d13',
   'psp_13_0_10_sos.bin': '0bcaaad9cd8578d3841ae69155a6bd4fc3ceae8f4fb5a6ba4f576e7ace94d1d9',
-  'psp_13_0_12_sos.bin': '7113a165c75c232d4cb7193a920b503e0bf082689adde3b45fdc38f58bfd18b3',
-  'psp_13_0_14_sos.bin': 'db863768cb25e806b68033e9237e0869f9f3603119df4d369ff4d80418d585d0',
-  'psp_13_0_15_sos.bin': '3b28d53e75a88131155e3931378ac8434eca4880ada9211d3b4e8915b6289583',
-  'psp_13_0_6_sos.bin': '36cce3a9441a0dcde81badd8fcf0416de8e4c39a7707865eff4d9d75e6bb0466',
-  'psp_13_0_7_sos.bin': '94db505fa6482f258c33a0a8d412050f6d843ab4ada368252e988f82f8a26fa8',
+  'psp_13_0_12_sos.bin': '89da90bf4286b38678b1fd175c78462a426afa3d258d15872cd14072d7098b9b',
+  'psp_13_0_14_sos.bin': 'a4f0d5f76d27b77409ec0b71d7cc6a848ddfd29f8c84f3003edf74ad3999fb7d',
+  'psp_13_0_6_sos.bin': '27657daa0f91ad8095d3610224a7de748b8b348a4cb211ecb5fccabe47369716',
+  'psp_13_0_7_sos.bin': 'ef1af0ecea38abbac6f85cce71789f19848c498d0cb8ef13748dab2d65b23c31',
   'psp_14_0_2_sos.bin': '7b538448b57d4f9dd06b2eea90d4f86a16e65e3027cdecee8db71c2c5f1fa243',
-  'psp_14_0_3_sos.bin': '28469a0857c813c54a0492423cdf0b0caf757428400036377e19c47e5af62478',
-  'smu_13_0_0.bin': '93e46a5526f19dcc3d13bfd9e23f88bc8eee52138bfe9caf0951b4eef5e49914',
-  'smu_13_0_0_kicker.bin': 'd0ef51d9ed06d0c17e06667302be21e7aedd86ed7a72be6e2f55b102214131cc',
-  'smu_13_0_10.bin': '9376ae64149e6b0b684898ffbc12c2230f8c50a2e9447dc7dafc95c0c16b5227',
+  'psp_14_0_3_sos.bin': '23bea01a0c6f36d00759d0765d46cb4cb4aa87398b2fbccacbf547a890c0bf51',
+  'smu_13_0_0.bin': '2ffac37fd8534965eeba19755db0e5ec80278213487dc4af0fbc8453befb64b1',
+  'smu_13_0_0_kicker.bin': '7f83656a2a89b7fce1c8a85e96d91cd8265a91fe883a7027f1a0ed18ced501de',
+  'smu_13_0_10.bin': 'daedb9cbdf48942be7ffe00d31b7c16bb36e11ff5a9d7495f218e95c07717b71',
   'smu_13_0_14.bin': 'a4f36de75fdcecd8000246762e027b4be489b6787afea57675225b0b39d35625',
   'smu_13_0_6.bin': 'ad7232264e8c57c2094244fbdd5a55d7a4575ffe9b44d229884bc0b6a44fb0b1',
-  'smu_13_0_7.bin': '68ec18bd605e680085c927ff72c609f8c771aff0718d0cfab58a3744dff8e5b7',
-  'smu_14_0_2.bin': '1b2054e3f710d1ab8dbdf6ff35914ad376b51caa6337831260c955add874b2ee',
-  'smu_14_0_3.bin': '4e1522d3c96c1028be2961dfcfc5f1ff783fb77724b260a99c4c8b4a901ef3fe',
-  'smu_14_0_3_kicker.bin': '9ff142656ae5f57be1b5ecc134e9da8f76650e793fbc5c499acd75094ff24453',
-  'sdma_4_4_2.bin': 'ff885711d2d5d75ceed51cf239e93c882584b918cd5d5d1ff58ee5aecc0c50ae',
-  'sdma_4_4_4.bin': '06a9d4d02c187844313a78469321d6091e59a334f0ce3b61d770d810c984d70b',
-  'sdma_4_4_5.bin': 'e2a30faa9403933fbfba7ce8e9feba460fff6ecdb15304818d24c9f3eeaad0a6',
+  'smu_13_0_7.bin': 'ccecc0fd0196b9613c920a51c2fd9436e739ff19dda5bdf74d97562387231732',
+  'smu_14_0_2.bin': '6951995d1d606f4dc60c895f19d34ed18aa40e62129f83d8510c45e8aa9ae2fc',
+  'smu_14_0_3.bin': 'df230947ddb7bbfd6e77d1280001db886e69adf2b2a448b47fa668a48bc0009f',
+  'smu_14_0_3_kicker.bin': '8ddc1da5b4e1619796c2cc81f19f388a35bf7d78bfe476cee559625589cb4dc7',
+  'sdma_4_4_2.bin': '456061b814268425843537da6f2191c8861d4e1a18d4c5d90c44ea6be18c78ff',
+  'sdma_4_4_4.bin': 'af47a2940e72b932d3e3a7e8f34f7a182624e5e433f7c56dff939ca5549cd33a',
+  'sdma_4_4_5.bin': '6127baabea3de7b18db3868c983b02c0fbf2cd75997f7f11241a5b1be27e5134',
   'sdma_5_2_6.bin': '3a163db00eb7e4752be8adbd61cf7dd8f08d924e59a6f798ced7dfcd89f340ed',
   'sdma_5_2_7.bin': '16fe80dc866b323e15a06f51646ef0f036878ad34da66921fcdb8167207d6b2b',
-  'sdma_6_0_0.bin': '82cd01a24171af12de6c7ac4ee7471aa2dfcf51f9677e7bae9cd4c75e07761ee',
-  'sdma_6_0_1.bin': '708c2c2c45262c98ebe8e34e32c3f1ce8eb5b007bab560c9ea9b576a8e4d6768',
-  'sdma_6_0_2.bin': '16c374344d2894da751f7028f9ec1f7520035fa9548d8c750d99a00a4afa86c7',
-  'sdma_6_0_3.bin': 'd47ac4db523aa35d77b27d57c35d4c04f431229dec6c0d667c36d98b985a6933',
-  'sdma_6_1_0.bin': '85f9f3714de68eee74cdf4852d709bc36a5c73a00e943b707bc2ce10d6b7bab4',
-  'sdma_6_1_1.bin': 'e7b7a23923ab691665e6ad16bbc8431a92f7c049da4b0b19a82c45fba03d4979',
-  'sdma_6_1_2.bin': '5947d78eb308a3f6a62d772c5a6493b21439c73eac139f9c22f080f660b4f4f3',
-  'sdma_6_1_3.bin': '8c651f32cbf030b6239ecc44f01bc9f5d5a193f333e21f2103736aff33227361',
+  'sdma_6_0_0.bin': '0f3da6b211f376356335b41be07149f650c10cfa4e23f7e25d53836006ed11f5',
+  'sdma_6_0_1.bin': 'ff565d3c215a30737560d4e3df6fc2c637738407e91d212fb200fdfb185b6744',
+  'sdma_6_0_2.bin': '398380184bb69113ef4c8964a3b55f6184deb0c1ffd96c9683490a3eec3ba8f3',
+  'sdma_6_0_3.bin': '0e8a83513087db865ba926f8b65cfb003fd41098f707e178d7a7ae2941fed0b1',
+  'sdma_6_1_0.bin': '22e55d0ad5f0247a7f0fffc67cfd3161b39f24ad6062ff3c91ec7ff38bd7e1e1',
+  'sdma_6_1_1.bin': '74533a581b8e3e2743b3c9c803d0666405e80898c4a630acefed82cb6b516ba2',
+  'sdma_6_1_2.bin': '4fe04b0286ec739b0414e8aee17e62e85e691f0246d1d9b56bc18a1219072314',
+  'sdma_6_1_3.bin': '35c9ed7e3a237c0d4a83b4975c63b62488f72aeafbb648342f384618e103f66b',
   'sdma_7_0_0.bin': 'beaafb53993a106edd392392d5896245ae2a957c6d0f495d0002eec72ad8ad38',
   'sdma_7_0_1.bin': '73c29e1c1714ebc95d2221ba56e187910902891593010653bf9518937e414a59',
-  'gc_10_3_6_pfp.bin': '042f5d2d223aac6a62b500a47d0d0bf33984200110da0ffca4fe5df9a96571c0',
+  'gc_10_3_6_pfp.bin': '793d678427887a0e724c79e356440aec33e6d1301f2a4e63543500249ebec064',
   'gc_10_3_7_pfp.bin': '3ae29aac3f424f7de97f82ce7158beba69509afb2dcbf1a428dc315df474a524',
-  'gc_11_0_0_pfp.bin': 'b360393c8629144b194f69a3cd961ed509331feff7a5cc1e4eb21c901da2710a',
-  'gc_11_0_1_pfp.bin': 'fb1ee527c05c55679c80a8bcf60fbb533724891baeb0eabc2917fc44e63a45dc',
-  'gc_11_0_2_pfp.bin': '9020f53788ad881fa01aa656fc082f9f8d3cdfc81f70aaac0bed6e6001491128',
-  'gc_11_0_3_pfp.bin': '362db904fa16c1fea2af7ad1295532434df7f85662b4a69332f51ae6c7290b61',
-  'gc_11_0_4_pfp.bin': 'aad22ca342c47d857bc1107a9aa9127e5e4ba7f7fd42d432213b1850bda1f4e1',
-  'gc_11_5_0_pfp.bin': '82ccf0265d841351183b011a79422799431f0c11f6d11165d64d7dfe404bda31',
-  'gc_11_5_1_pfp.bin': '633404d8db1dc03fe997f7d0d0e15ef908069727abaf9de55841be3f3c97348b',
-  'gc_11_5_2_pfp.bin': 'baee1456dd1800cdaedd4998c2dd7d76cdc0cf0ec928679fe67b02485905ea2c',
-  'gc_11_5_3_pfp.bin': 'fee840b049b5e082215df72a93fad80a64f07ef6f638408a2d56fae97449a2cb',
-  'gc_12_0_0_pfp.bin': 'd1b043c60920e509e5c8f9677221fb78ff7985f68b605e8f39a04a57333a9366',
-  'gc_12_0_1_pfp.bin': '9d8d6188efeca5ef05482d9299c4f102fab7db3dae51a23e59de9baa34997123',
-  'gc_10_3_6_me.bin': '776d2299bc4f3abffd4a7999f5a21a4e38aced8b6b4c199a83610dbabf08176d',
+  'gc_11_0_0_pfp.bin': 'e175cb0f580a38c961a6f7366142c08e413995f57f78f39795368b15442df8a3',
+  'gc_11_0_1_pfp.bin': 'f5bf21dfbd9e72a30b4caf4704282c27854710c1b7c4affbb2a19530466b12a8',
+  'gc_11_0_2_pfp.bin': '001c4dec1119e29314d725cc1280fc4f0cd9cabdf61ea5ee2260cfd4e62ec141',
+  'gc_11_0_3_pfp.bin': '0488034c85be97125e39e860308d33c3f76a01df8250092a32d4d55acb2526fd',
+  'gc_11_0_4_pfp.bin': '5ae8b7bb6316f87ae8b978354c088e3bd8c890959382d72886377cda25b1ffd1',
+  'gc_11_5_0_pfp.bin': '0124f540871a7759fa8aaae046d458dfb34aeea12a1183ff962c3f1a33067d5a',
+  'gc_11_5_1_pfp.bin': '7794ea46d0d3cf9cb3f7938affbdf09dd7a9970340da5cd02b774cb393436d24',
+  'gc_11_5_2_pfp.bin': '55e64741de28c506524959f7f696713a72aafe46f49ccd827781d67a9475b386',
+  'gc_11_5_3_pfp.bin': 'ce805040fb347fddbc89b2715e66b446865dda9e2056a9b233269b72bc09c387',
+  'gc_12_0_0_pfp.bin': '16bfd64c10fe73b5e760055069a60e5841dba16c0ed4edb56c20d675e23901f6',
+  'gc_12_0_1_pfp.bin': '49efb319305c5fffd90ac1eef7d7a0bdec72998ecb5cf4526996311788a53dc3',
+  'gc_10_3_6_me.bin': '141b59faad3f2f1be16a2178833b7ca8e97519e1e844c8fda6689572c3767902',
   'gc_10_3_7_me.bin': '9eb0b56e9bcc9dad5d53437b162226fcb37e5df102832260f1232832f3658edf',
-  'gc_11_0_0_me.bin': 'f2f5a793d811c6abad1a18af0fcf7694c443478f224176da86650c22aa71ca7a',
-  'gc_11_0_1_me.bin': '476db2ec7e33d1e126b1736649208443e3ccc68aa60e4978574cdbced2b26543',
-  'gc_11_0_2_me.bin': 'f5fe48f97acbd3ce13b35929290bfbac01ce522631cc91dcef1fdeb3ff35c8ed',
-  'gc_11_0_3_me.bin': 'd02c25070e5bdf0ec0146f5c9d6d2f8b86de43bd2a318a0b67eb5201963bafdc',
-  'gc_11_0_4_me.bin': 'f075220f75ffe43eacc5986ff8448946c27405e632764fda83323e7ec8d55566',
-  'gc_11_5_0_me.bin': '338019a1fcdab39729e3f492ffc9f5970c2c81b12c8a4f431494ca28cfdadedf',
-  'gc_11_5_1_me.bin': '4c4dd30c22d4f7f2c5d3a19c645f505e30cdac115a91c65791e2651b22932175',
-  'gc_11_5_2_me.bin': 'cab2999186d26c0e9a3d46b5a43d2854d88be880cb764c096ad2b43038566384',
-  'gc_11_5_3_me.bin': '94e2d74e834725b3d51e03e830160e95c56f3a31e93f5d61c852afd8fe8cc779',
-  'gc_12_0_0_me.bin': 'fb10cb3535ae4a6a8fb3e78166cf30c5b717341b1f20cde73065c62b642adfed',
-  'gc_12_0_1_me.bin': '56a1ae0031aa938f6b61348a56404ab2cee92f1f45630fc82a801aa4d908f98a',
-  'gc_10_3_6_mec.bin': '7003c4a77537e9edaf67064104cd9371fac38a84f71f948349140b28d3c210e8',
+  'gc_11_0_0_me.bin': 'f8fba8a63dd4293b8fc1e4aab78b6fac630e575d1d62838c7996d9210f82aea1',
+  'gc_11_0_1_me.bin': '5030040b00955de94876341ec64ea43b96640413d7a03dc460a83c8386bf76e0',
+  'gc_11_0_2_me.bin': '0f21fd43f1dfbc6ccced9a2b3774de25c993c61a689aabab8b45333937b7945e',
+  'gc_11_0_3_me.bin': '3acb5061dba342ade81d329d1932f19ec01f0c5bf44e6e3568008a951a351bac',
+  'gc_11_0_4_me.bin': 'e4f1f6abcd213d54ad9e885d9f550083b0e2f67d983566015e8a53981e1cb155',
+  'gc_11_5_0_me.bin': '8f906b64d0a29503daa662c93ec44d076fcac11b78f70cd50ce0af2b500a05a6',
+  'gc_11_5_1_me.bin': '7e42602bcbaf1e511f8b4f6ed2246844ad1f6e351ce2b663d89062a7be263663',
+  'gc_11_5_2_me.bin': 'aae26255d8efff81e0e3bbcb727efb8b837d8e25fe85c708545f5328f1077b50',
+  'gc_11_5_3_me.bin': '93cd588348b16fe432609fe8da6e6b5da0a52da5c5884882aecf7b1001f72700',
+  'gc_12_0_0_me.bin': 'd7eba5197f2580f32b8256b1d9cb68e723e9e644293a34446a7913e3c093cba5',
+  'gc_12_0_1_me.bin': '365e7f193b39cbb10d3af44905fefaca0e9844721801755276baebac7b19c1ea',
+  'gc_10_3_6_mec.bin': '247943415658159704a21f670dd7b3e7cb2d2fc0c17b000a5098715979c8d95e',
   'gc_10_3_7_mec.bin': 'ee58a523375bcf5b89400b32b801f95e182b632a26bce4f2bed5c07928d486dc',
-  'gc_11_0_0_mec.bin': '1dd1de8ecf5455ea4719c502b64b32ac18763d5601128c01b4a4a36211a122c2',
-  'gc_11_0_1_mec.bin': '505ae64eccb2e4b4751fe18ec1b584e1f6b4c81d0f5ec089afbcf378cad59711',
-  'gc_11_0_2_mec.bin': '19bf080d6e672de5ed3fb86e3fdbdda4d700d8e3bda2dbdcc923101484ad645b',
-  'gc_11_0_3_mec.bin': 'a37bc1a4e245300a5c3e26da34ea213842447d7df6c5c81e9fc78887a2fde26f',
-  'gc_11_0_4_mec.bin': '850d5302b4fee6022f42f706c2de103531b45b7794a45f2d6dce6015767a1ad6',
-  'gc_11_5_0_mec.bin': '5e022bae6638967d82e2b1077e3024f52bc83b3cb850aa31fba51469c7517c4c',
-  'gc_11_5_1_mec.bin': 'e49964d5e58686c53e66d98d4e3b9fab70e98fad3b28379c6e60aed03c83ee80',
-  'gc_11_5_2_mec.bin': '9691d7bff5d2c933d8eecb7d171635612a76a2dd1441cffcd65a8a02bdb5a2c5',
-  'gc_11_5_3_mec.bin': 'd368f3886b9245dd0d21d57fccfd8aa7e872c2564e23f292abe735348121277e',
-  'gc_12_0_0_mec.bin': '9c7602d6ebf1f7e6ec7a5d1ceefded18f35fa1c08fbea1e3e1a0d78d519db8e8',
-  'gc_12_0_1_mec.bin': 'caf1dbaf72b0ef0c4c973947414033aeec002994f63967bb53e9165195a3c2c3',
-  'gc_9_4_3_mec.bin': '99bc12230f00b930cf286105a35cc6110d87461cd48cb4fdf3cb6caff73ac1e7',
-  'gc_9_4_3_sjt_mec.bin': '2945dbd098c4158870df7dc4ccb33d40031fd1cce37cdbe5df291d8941d03567',
-  'gc_9_4_4_mec.bin': '7f14258f8301d2717e0a707ccfad7b3091af478b0df6d5134adfd56caa7429d8',
-  'gc_9_4_4_sjt_mec.bin': '0bbef279bbc07c502098b80765b876f69fcda9834e5ed269a7d8236c85e89e19',
-  'gc_9_5_0_mec.bin': '0c39078c53e10e99538901df5fc14e7f1b1f3639ea825b1b3126ae87a28b2464',
-  'gc_9_5_0_sjt_mec.bin': 'a769745367567fc6f389695aa5f48c154c07560e21a93052185e19f950205240',
+  'gc_11_0_0_mec.bin': '801a09c9bf06188260db9b51ad8f978f15d84c72ca91b90643a2ef8af4074776',
+  'gc_11_0_1_mec.bin': '6afadcb7504bb11bcc9d4a205cdf73f7934a615e28f178fcf7285971df2ccd05',
+  'gc_11_0_2_mec.bin': '0da0edee28c73a6fa1191f77853d380ec2503cbf43e0aaae4617f32f1f8a48fa',
+  'gc_11_0_3_mec.bin': '323cfa6658b6b5169830f852e2ff0552acae8dfb9e44b42c63de7b2900d3fd9e',
+  'gc_11_0_4_mec.bin': '5d89cf6b60354f3746c2cbd1ff0cb1a741556ca20d72745242cb69b553d0985c',
+  'gc_11_5_0_mec.bin': 'a01c324ab14ec89792449a621a541829b9af26865019027a411a14b910145dfa',
+  'gc_11_5_1_mec.bin': 'eab05719371caa68df09d4f7574e3958a3c4f5044ab3c7b0d2b214add0c6d1c4',
+  'gc_11_5_2_mec.bin': 'a374b2335802e24f8b9a3ce40000a1d37a52a14eb87099bebcc6680c27cc93e5',
+  'gc_11_5_3_mec.bin': '165025437cba80dd32c19ebbc83b756fa7adac7053ff7780ba4aa2f8089c6a3f',
+  'gc_12_0_0_mec.bin': '1931593440b8f9423580d9e2cdc5b34e7c682cdffe1ca4b74b0c2f6a0420236d',
+  'gc_12_0_1_mec.bin': 'f57541688a5108730bf210663f1137ffc2121f3acfe614a6de09ec1982c69a2f',
+  'gc_9_4_3_mec.bin': '3159176e72301fb88dc416721fb3d0ab82ece484cf93a43c3f37430c7e6673a1',
+  'gc_9_4_3_sjt_mec.bin': 'd19468dbb47849640bd0e6cdc8d7e25a3c8442c7ca2ca81357702e0d6baab50f',
+  'gc_9_4_4_mec.bin': '5004f73e43db2dd45e77d65942e33d4a69e7157618cfd23944c30f801c77a0f3',
+  'gc_9_4_4_sjt_mec.bin': '627a9e98102e70fe3bf0947eb764187f29f5e775d1130c7310e0ba5fc0502dbe',
+  'gc_9_5_0_mec.bin': 'c5eca4311a6f6e8f81cf41c2c46941d5dcf90789ee8326901da2dfc86ac14c31',
+  'gc_9_5_0_sjt_mec.bin': 'f162e509379288e3f3b1eead541b315c2262d625d433287ecd34ca185614d312',
   'gc_11_0_0_imu.bin': 'b4f8fc056b45709a6abf48e7885fb1b4ab8d3cc092cbfa2c554a78564a6403bc',
   'gc_11_0_1_imu.bin': 'ac71f4eec713fc35b4a1fe27531e3eb04edd81eeac2cef64df01ac50d8510805',
   'gc_11_0_2_imu.bin': '9befca62b0b0cfd252c3df4a9edca295526f4d43821cd99a6326454995a6ca2d',
@@ -91,17 +90,17 @@ hashes = {
   'gc_10_3_6_rlc.bin': 'acfbac75c0dcfbfe40e222640ef17eb3dc8d206d30bc3863f275f2dd1cb132a5',
   'gc_10_3_7_rlc.bin': 'a02585ebe3b36d942e883057119572d9497600c52fc65b8a523487eb65d874f2',
   'gc_11_0_0_rlc.bin': 'dabd49039772d02f5fd5e48dc21d35ad52a6b1283b470dabca86ca159c4c7c8e',
-  'gc_11_0_1_rlc.bin': '5f07dc1f0a75ecd9cb56d805ea869184a50ed9e43d811ebf833b8906534650ef',
+  'gc_11_0_1_rlc.bin': '86145719a58e9428562930c6b5ee3b6ced4701d34a80d0b4d84d6026c93134f2',
   'gc_11_0_2_rlc.bin': 'b43eb2fd0600f50a1a5796bc9983d6b39b5c20960234920f5e89cb362193e0b8',
-  'gc_11_0_3_rlc.bin': '890d8e0123efb40c0179dd8ac3e9af073a0b87cbbccfec1db54e5ed2315a8d39',
-  'gc_11_0_4_rlc.bin': '257ced82d7bec41249b06592ee0c44fb8f9262de2c6af9c52dc6f6a8a702063e',
-  'gc_11_5_0_rlc.bin': '0dc8b6ef5530a4a53938c8baa0d49cd458607d95233237859fa98d44feb3e985',
+  'gc_11_0_3_rlc.bin': '29b0b456f5b53076ddffa6f09de3bb697219e8e7b33504bf6c197e8b858426dc',
+  'gc_11_0_4_rlc.bin': '823573078b608108fbe4dd8176c396ec582632913db9c59a512d82b068f8eba0',
+  'gc_11_5_0_rlc.bin': '68cd85567f4f2f8d6b80db294988806d956bf826979c3597daccb71c7ee6aadd',
   'gc_11_5_1_rlc.bin': '92731ecabbeb77865fb71787b4268dc738a58779f1190bdc2056482cb88a08f6',
-  'gc_11_5_2_rlc.bin': 'c9ad70b8ac309257cb8929bb6b4efa6b551ec1e5229d7a419332a9797f31fc9e',
+  'gc_11_5_2_rlc.bin': 'ef3a9209d3eccfbe18fce9e972c146ac283719798bb788096c176b796dc9aee5',
   'gc_11_5_3_rlc.bin': '10a68940c6258d5818d9c05fd98eb0ccc8d5aee99b2769fbad30e5abd0d9327e',
   'gc_12_0_0_rlc.bin': '6436b582734a413456fff3d3c7195e71cc9e78a7ed31ee21c83ffd6fae1ad186',
   'gc_12_0_1_rlc.bin': '6ba4459532246a5c415d3cb33c9b1248294e48f67b827e2accb292a8d1a5c0ec',
-  'gc_9_4_3_rlc.bin': '54cbd0de3a0ec35d2e58e992babeee2a237f870ccdf37e734652e4daeeba59d5',
+  'gc_9_4_3_rlc.bin': '5345d388712d547b0ae16f199ad5ccadb65643584b3efa7817049ddeb3fdcd12',
   'gc_9_4_4_rlc.bin': 'e0c3585c72f8136670ca63e607fba32c1ae4948f493f13e33fc4d466bd6318a8',
   'gc_9_5_0_rlc.bin': '9b1268f5751153fe57f527c9acb417bfa53ed42c9bc083c9d3da2ba61fe5fdc4',
 }

tinygrad/runtime/ops_amd.py

@@ -944,7 +944,9 @@ class AMDDevice(HCQCompiled):
   def is_usb(self) -> bool: return isinstance(self.iface, USBIface)
 
   def __init__(self, device:str=""):
-    self.iface = self._select_iface(device)
+    self.device_id = int(device.split(":")[1]) if ":" in device else 0
+
+    self.iface = self._select_iface()
 
     self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
     self.arch = "gfx%d%x%x" % self.target

tinygrad/runtime/ops_cl.py

@@ -24,10 +24,10 @@ class CLCompiler(Compiler):
     super().__init__(f"compile_cl_{compile_key}")
   def compile(self, src:str) -> bytes:
     program = checked(cl.clCreateProgramWithSource(self.dev.context, 1, to_char_p_p([src.encode()]), None, status := ctypes.c_int32()), status)
-    build_status: int = cl.clBuildProgram(program, 1, self.dev.cl_dev, None, BP_CB(), None)
+    build_status: int = cl.clBuildProgram(program, 1, self.dev.device_id, None, BP_CB(), None)
     if build_status != 0:
-      cl.clGetProgramBuildInfo(program, self.dev.cl_dev, cl.CL_PROGRAM_BUILD_LOG, 0, None, log_size := ctypes.c_size_t())
-      cl.clGetProgramBuildInfo(program, self.dev.cl_dev, cl.CL_PROGRAM_BUILD_LOG,
+      cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG, 0, None, log_size := ctypes.c_size_t())
+      cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG,
                                log_size.value, mstr := ctypes.create_string_buffer(log_size.value), None)
       raise CompileError(f"OpenCL Compile Error\n\n{mstr.value.decode()}")
     check(cl.clGetProgramInfo(program, cl.CL_PROGRAM_BINARY_SIZES, ctypes.sizeof(ctypes.c_size_t), binary_sizes := (ctypes.c_size_t * 1)(), None))
@@ -39,11 +39,11 @@ class CLCompiler(Compiler):
 class CLProgram(Program['CLDevice']):
   def __init__(self, device:CLDevice, obj:TinyELF):
     self.dev, self.lib, self.signature = device, device.cl_compiler.compile_cached(obj.lib.decode()), obj.signature
-    self.program = checked(cl.clCreateProgramWithBinary(device.context, 1, device.cl_dev, (ctypes.c_size_t * 1)(len(self.lib)),
+    self.program = checked(cl.clCreateProgramWithBinary(device.context, 1, device.device_id, (ctypes.c_size_t * 1)(len(self.lib)),
                                                         to_char_p_p([self.lib], ctypes.c_ubyte), binary_status := ctypes.c_int32(),
                                                         errcode_ret := ctypes.c_int32()), errcode_ret)
     check(binary_status.value)
-    check(cl.clBuildProgram(self.program, 1, device.cl_dev, None, BP_CB(), None)) # NOTE: OSX requires this
+    check(cl.clBuildProgram(self.program, 1, device.device_id, None, BP_CB(), None)) # NOTE: OSX requires this
     self.kernel = checked(cl.clCreateKernel(self.program, obj.name.encode(), status := ctypes.c_int32()), status)
 
   def __del__(self):
@@ -101,17 +101,17 @@ class CLDevice(Compiled):
       CLDevice.device_ids = c.init_c_var((cl.cl_device_id * num_devices.value),
                                          lambda x: check(cl.clGetDeviceIDs(platform_ids[0], device_type, num_devices, x, None)))
 
-    self.cl_dev = CLDevice.device_ids[0 if ":" not in device else int(device.split(":")[1])]
-    self.device_name = (cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_NAME, 256,
+    self.device_id = CLDevice.device_ids[0 if ":" not in device else int(device.split(":")[1])]
+    self.device_name = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_NAME, 256,
                                            buf:=ctypes.create_string_buffer(256), None), buf.value.decode())[1]
-    self.driver_version = (cl.clGetDeviceInfo(self.cl_dev, cl.CL_DRIVER_VERSION, 256,
+    self.driver_version = (cl.clGetDeviceInfo(self.device_id, cl.CL_DRIVER_VERSION, 256,
                                               buf:=ctypes.create_string_buffer(256), None), buf.value.decode())[1]
     if DEBUG >= 1: print(f"CLDevice: opening {self.device_name} with version {self.driver_version}")
-    self.context = checked(cl.clCreateContext(None, 1, self.cl_dev, CC_CB(), None, status := ctypes.c_int32()), status)
-    self.queue = checked(cl.clCreateCommandQueue(self.context, self.cl_dev, cl.CL_QUEUE_PROFILING_ENABLE, status), status)
+    self.context = checked(cl.clCreateContext(None, 1, self.device_id, CC_CB(), None, status := ctypes.c_int32()), status)
+    self.queue = checked(cl.clCreateCommandQueue(self.context, self.device_id, cl.CL_QUEUE_PROFILING_ENABLE, status), status)
     self.pending_copyin: list[memoryview] = []
-    check(cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_EXTENSIONS, 0, None, ctypes.byref(exts_len:=ctypes.c_size_t())))
-    self.device_exts = (cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_EXTENSIONS, exts_len.value,
+    check(cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, 0, None, ctypes.byref(exts_len:=ctypes.c_size_t())))
+    self.device_exts = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, exts_len.value,
                                            ctypes.byref(buf := ctypes.create_string_buffer(exts_len.value)), None),
                                            ctypes.string_at(buf).decode().split())[1]
 
@@ -119,7 +119,7 @@ class CLDevice(Compiled):
 
     arch = ",".join(self.device_exts)
     if "cl_khr_image2d_from_buffer" in self.device_exts:
-      check(cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_IMAGE_PITCH_ALIGNMENT, 4, ctypes.byref(ipa := ctypes.c_uint32()), None))
+      check(cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_IMAGE_PITCH_ALIGNMENT, 4, ctypes.byref(ipa := ctypes.c_uint32()), None))
       arch += f",IMAGE_PITCH_ALIGNMENT={ipa.value}"
     super().__init__(device, CLAllocator(self), [OpenCLRenderer], CLProgram, arch=arch)
 

tinygrad/runtime/ops_cpu.py

@@ -6,7 +6,6 @@ from tinygrad.helpers import to_mv, from_mv, OSX, WIN, Context, mv_address, supp
 from tinygrad.device import Buffer, BufferSpec, TinyELF, Program, Device
 from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface
 from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, make_cmdbuf, make_signal
-from tinygrad.runtime.support.c import DLL
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
 from tinygrad.renderer.nir import LVPRenderer
@@ -14,7 +13,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer
 from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.runtime.autogen import libc
 from tinygrad.codegen import do_to_program
-from tinygrad.engine.realize import pm_flatten_linear, get_call_arg_uops, get_call_var_uops, get_runtime
+from tinygrad.engine.realize import pm_flatten_linear, get_call_arg_uops, get_runtime
 from tinygrad import UOp, dtypes
 from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import KernelInfo, Ops, UPat, PatternMatcher, graph_rewrite
@@ -65,7 +64,7 @@ def cpu_cmd(devs:tuple[str, ...], prog, *args:UOp) -> UOp:
   return UOp(Ops.INS, dtypes.void, words + (UOp.const(0, dtypes.uint64),) * (CMD_SIZE - len(words)), arg="cmd")
 
 def cpu_exec(ctx:tuple[str, ...], call:UOp, prg:UOp) -> UOp:
-  args = [get_call_arg_uops(call)[i].getaddr(ctx) for i in prg.arg.globals] + [v.cast(dtypes.uint64) for v in get_call_var_uops(call, prg)]
+  args = [get_call_arg_uops(call)[i].getaddr(ctx) for i in prg.arg.globals] + [v.cast(dtypes.uint64) for v in prg.arg.vars]
   if (core:=prg.arg.runtimevars.get('core_id')) is None: return cpu_cmd(ctx, prg, *args)
 
   la = [cpu_cmd(ctx,prg,*args[:(cid:=(len(prg.arg.globals)+core))],UOp.const(t, dtypes.uint64),*args[cid+1:]) for t in range(prg.arg.global_size[0])]
@@ -100,11 +99,10 @@ def encode_queue(q:UOp) -> UOp:
   e = UOp.range(cnt, next(UOp.unique_num), dtype=dtypes.int, src=(cmdbuf, ring))
   copy = UOp.group(*[ring.index((base + e*CMD_SIZE + w) % ring_words).store(cmdbuf.index(e*CMD_SIZE + w).load()) for w in range(CMD_SIZE)])
 
-  bumped = put.after(copy.end(e)).index(0).store(put.index(0).load() + cnt)
-  if WIN: return sysbuf.after(bumped).index(0).store(put.after(bumped).index(0).load())
-
-  e = UOp.range(cnt, next(UOp.unique_num), dtype=dtypes.int, src=(bumped,))
-  return make_signal(devs, tag="func:sem_post").after(e).index(0).load().call(sem.after(e).index(0), ret_dtype=dtypes.void).end(e)
+  # wake the worker after each entry, keeping the post with the stores stops it from hoisting out of the loop
+  wake = copy.end(e) if WIN else make_signal(devs, tag="func:sem_post").after(copy).index(0).load().call(sem.index(0), ret_dtype=dtypes.void).end(e)
+  bumped = put.after(wake).index(0).store(put.index(0).load() + cnt)
+  return sysbuf.after(bumped).index(0).store(put.index(0).load() + cnt) if WIN else bumped
 
 # *****************
 
@@ -112,9 +110,9 @@ def encode_queue(q:UOp) -> UOp:
 MAP_JIT = 0x0800
 
 class CPUProgram(Program['CPUDevice']):
-  rt_lib, libm = DLL('rt', 'System' if OSX else 'kernel' if WIN else 'gcc_s'), DLL('m', 'm')
-
-  def _load(self, lib, base=0): return lib if lib[:4] != libc.ELFMAG.encode() else jit_loader(lib, base=base, link_libs=[self.libm, self.rt_lib])
+  rt_lib = None
+  try: rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or WIN else 'libgcc_s.so.1')
+  except OSError: pass
 
   def __init__(self, dev:CPUDevice, obj:TinyELF):
     self.dev, self.name, self.signature = dev, obj.name, obj.signature
@@ -126,10 +124,10 @@ class CPUProgram(Program['CPUDevice']):
       ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_void_p
       self.addr = ctypes.windll.kernel32.VirtualAlloc(ctypes.c_void_p(0), ctypes.c_size_t(len(obj.lib)), MEM_COMMIT | MEM_RESERVE,
                                                       PAGE_EXECUTE_READWRITE)
-      ctypes.memmove(self.addr, (loaded:=self._load(obj.lib, self.addr)), len(loaded))
+      ctypes.memmove(self.addr, obj.lib, len(obj.lib))
       ctypes.windll.kernel32.GetCurrentProcess.restype = ctypes.c_void_p
       proc = ctypes.windll.kernel32.GetCurrentProcess()
-      ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.addr), ctypes.c_size_t(len(loaded)))
+      ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.addr), ctypes.c_size_t(len(obj.lib)))
       self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr) if self.lvp else ctypes.CFUNCTYPE(None)(self.addr)
     else:
       # On apple silicon with SPRR enabled (it always is in macos) RWX pages are unrepresentable: https://blog.svenpeter.dev/posts/m1_sprr_gxf/
@@ -138,17 +136,18 @@ class CPUProgram(Program['CPUDevice']):
       self.addr = mv_address(self.mem)
 
       if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(False)
-      self.mem.write(loaded:=self._load(obj.lib, mv_address(self.mem)))
+      lib = jit_loader(obj.lib, base=ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)), link_libs=['m']) if self.lvp else obj.lib
+      self.mem.write(lib)
       if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(True)
 
       # __clear_cache isn't a normal libc function, but a compiler support routine found in libgcc_s for gcc and compiler-rt for clang.
       # libgcc_s comes as shared library but compiler-rt is only a bunch of static library archives which we can't directly load, but fortunately
       # it somehow found its way into libSystem on macos (likely because it used __builtin_clear_cache) and libgcc_s is ~always present on linux
       # Using ["name"] instead of .name because otherwise name is getting mangled: https://docs.python.org/3.12/reference/expressions.html#index-5
-      if 'rt' in DLL._loaded_: CPUProgram.rt_lib["__clear_cache"](ctypes.c_void_p(self.addr), ctypes.c_void_p(self.addr + len(loaded)))
+      if CPUProgram.rt_lib is not None: CPUProgram.rt_lib["__clear_cache"](ctypes.c_void_p(self.addr), ctypes.c_void_p(self.addr + len(lib)))
       else:
         # msync should be a universal POSIX way to do this
-        libc.msync(ctypes.c_void_p(self.addr), len(loaded), libc.MS_SYNC | libc.MS_INVALIDATE)
+        libc.msync(ctypes.c_void_p(self.addr), len(lib), libc.MS_SYNC | libc.MS_INVALIDATE)
 
       self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr) if self.lvp else ctypes.CFUNCTYPE(None)(self.addr)
 
@@ -197,13 +196,11 @@ class CPUDevice(HCQ2Compiled):
   pm_lower = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue)])
 
   def __init__(self, device:str=""):
-    self.workers:list[CPUWorker] = []
     super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram,
       arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native")
 
     self.pm_bufferize = PatternMatcher(
-      [(UPat(Ops.PARAM, tag=f"{q}_{n}"), lambda ctx, q=q,n=n: getattr(ctx[0].worker(q), n))
-       for q in ("COMPUTE:0", "SUBMIT:0") for n in ("ring", "put", "sem", "sys", "done")] +
+      [(UPat(Ops.PARAM, tag=f"COMPUTE:0_{n}"), lambda ctx, n=n: getattr(ctx[0].worker, n)) for n in ("ring", "put", "sem", "sys", "done")] +
       [(UPat(Ops.PARAM, tag=f"func:{f}"), lambda ctx, f=f: ctx[0].func_ptr(f)) for f in FUNCS]) + self.pm_bufferize
 
     with Context(EMULATED_DTYPES="", TRACK_MATCH_STATS=0):
@@ -213,12 +210,6 @@ class CPUDevice(HCQ2Compiled):
 
   def func_ptr(self, name:str) -> Buffer: return self.func_table.view(1, dtypes.uint64, FUNCS.index(name)*8).ensure_allocated()
 
-  def synchronize(self, timeout:int|None=None):
-    for worker in self.workers:
-      put, done = (getattr(worker, x)._buf.cpu_view().view(fmt='Q') for x in ("put", "done"))
-      while done[0] < put[0]: self._wait_signal(done, put[0], timeout)
-    super().synchronize(timeout)
-
   @functools.cached_property
   def func_table(self) -> Buffer:
     lib = ctypes.windll.kernel32 if sys.platform == "win32" else libc.dll # type: ignore[attr-defined]
@@ -226,8 +217,8 @@ class CPUDevice(HCQ2Compiled):
       array.array('Q', [unwrap(ctypes.cast(getattr(lib, f), ctypes.c_void_p).value) for f in FUNCS])
     return ft
 
-  @functools.cache
-  def worker(self, queue:str) -> CPUWorker:
+  @functools.cached_property
+  def worker(self) -> CPUWorker:
     ring, put, sysbuf, done = (Buffer(self.device, sz, dtypes.uint64, preallocate=True) for sz in (RING_SLOTS*CMD_SIZE, 1, 1, 1))
     addr, hsem = 0, None
 
@@ -239,6 +230,5 @@ class CPUDevice(HCQ2Compiled):
     sem = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(external_ptr=addr), preallocate=True)
 
     worker_args = [ring._buf.va_addr, sysbuf._buf.va_addr if WIN else self.func_ptr('sem_wait')._buf.va_addr, done._buf.va_addr, addr]
-    (thread:=threading.Thread(target=self.prgs[worker_prog].fxn, daemon=True, args=[ctypes.c_uint64(x) for x in worker_args])).start()
-    self.workers.append(worker:=CPUWorker(ring, put, sem, sysbuf, done, thread))
-    return worker
+    (worker:=threading.Thread(target=self.prgs[worker_prog].fxn, daemon=True, args=[ctypes.c_uint64(x) for x in worker_args])).start()
+    return CPUWorker(ring, put, sem, sysbuf, done, worker)

tinygrad/runtime/ops_nv.py

@@ -588,7 +588,8 @@ class NVDevice(HCQCompiled[NVSignal]):
   def is_nvd(self) -> bool: return isinstance(self.iface, PCIIface)
 
   def __init__(self, device:str=""):
-    self.iface = self._select_iface(device)
+    self.device_id = int(device.split(":")[1]) if ":" in device else 0
+    self.iface = self._select_iface()
 
     device_params = nv_gpu.NV0080_ALLOC_PARAMETERS(deviceId=self.iface.gpu_instance, hClientShare=self.iface.root,
                                                    vaMode=nv_gpu.NV_DEVICE_ALLOCATION_VAMODE_OPTIONAL_MULTIPLE_VASPACES)

tinygrad/runtime/ops_python.py

@@ -23,9 +23,7 @@ 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}")
-  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)
+  m[i] = to_storage_scalar(v, dtype)
 
 # 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):

tinygrad/runtime/ops_qcom.py

@@ -23,8 +23,8 @@ def dcache_flush():
   buf, n = UOp.param(0, dtypes.uint8, shape=(1,)), UOp.param(1, dtypes.int, shape=(), name="n", addrspace=AddrSpace.ALU)
   i = UOp.range(n, 0, dtype=dtypes.int)
   flush = UOp(Ops.CUSTOM, src=(buf.index(i * 64),), arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");')
-  sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush"), tag=1)
-  prg = to_program(sink, Device["CPU"].renderer)
+  sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush"))
+  prg = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())))), Device["CPU"].renderer)
   return Device["CPU"].runtime(prg.to_elf())
 
 #Parse C-style defines: <regname>_<field_x>__SHIFT and <regname>_<field_y>__MASK from the adreno module into the following format:

tinygrad/runtime/ops_rdma.py

@@ -101,5 +101,6 @@ class RDMAAllocator(HCQAllocatorBase):
 
 class RDMADevice(HCQCompiled):
   def __init__(self, device:str=""):
-    self.iface = MLXIface(self, int(device.split(":")[1]) if ":" in device else 0)
+    self.device_id = int(device.split(":")[1]) if ":" in device else 0
+    self.iface = MLXIface(self, self.device_id)
     super().__init__(device, RDMAAllocator(self), [], None, signal_t=None)

tinygrad/runtime/support/c.py

@@ -91,7 +91,7 @@ class DLL(ctypes.CDLL):
 
   @staticmethod
   def findlib(nm:str, paths:list[str], extra_paths=[]):
-    if nm in ('libc', 'm') and OSX: return f'/usr/lib/lib{nm.removeprefix("lib")}.dylib'
+    if nm == 'libc' and OSX: return '/usr/lib/libc.dylib'
     if pathlib.Path(path:=getenv(nm.replace('-', '_').upper()+"_PATH", '')).is_file(): return path
     for p in paths:
       libpaths = {"posix": [d for d in os.environ.get('LD_LIBRARY_PATH', '').split(os.pathsep) if d] + ["/usr/lib64", "/usr/lib", "/usr/local/lib"],

tinygrad/runtime/support/compiler_cpu.py

@@ -1,9 +1,10 @@
 import subprocess
 from tinygrad.device import Compiler
-from tinygrad.helpers import getenv, capstone_flatdump, cpu_objdump
+from tinygrad.helpers import getenv, capstone_flatdump
+from tinygrad.runtime.support.elf import jit_loader
 
 class ClangCompiler(Compiler):
-  def __init__(self, arch:list[str], cachekey="compile_clang_obj"):
+  def __init__(self, arch:list[str], cachekey="compile_clang_jit"):
     assert len(arch) >= 2, f"invalid arch string: {','.join(arch)!r}, expected '<arch>,<cpu>,[<feats>]' (eg. 'x86_64,znver2')"
     self.arch, cpu, *feats = arch
     match self.arch:
@@ -15,13 +16,15 @@ class ClangCompiler(Compiler):
       case _: raise RuntimeError(f"unsupported arch: {self.arch!r}")
     super().__init__(f"{cachekey}_{'_'.join(arch)}")
 
-  def compile(self, src:str) -> bytes:
+  def compile_to_obj(self, src:str) -> bytes:
     """Compile C source to ELF object file (before linking)."""
     # -fno-math-errno is required for __builtin_sqrt to become an instruction instead of a function call
     return subprocess.check_output([getenv("CC", 'clang'), '-c', '-x', 'c', '-O2', '-fPIC', '-ffreestanding', '-fno-math-errno', '-nostdlib',
                                     '-fno-ident', f'--target={self.arch}-none-unknown-elf', *self.args, '-', '-o', '-'], input=src.encode('utf-8'))
 
-  def disassemble(self, lib: bytes): cpu_objdump(lib)
+  def compile(self, src:str) -> bytes: return jit_loader(self.compile_to_obj(src))
+
+  def disassemble(self, lib:bytes): return capstone_flatdump(lib, self.arch)
 
 
 class X86Compiler(Compiler):

tinygrad/runtime/support/compiler_cuda.py

@@ -1,12 +1,10 @@
 import hashlib, tempfile, ctypes, re, pathlib
-from tinygrad.helpers import to_char_p_p, colored, getenv, system, OSX
+from tinygrad.helpers import to_char_p_p, colored, getenv, system
 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)
@@ -46,14 +44,11 @@ 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}']
-    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")
+    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,
@@ -85,11 +80,9 @@ class PTXCompiler(Compiler):
 
 class NVPTXCompiler(PTXCompiler):
   def __init__(self, arch:str):
-    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())))
+    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_llvm.py

@@ -1,6 +1,7 @@
 import ctypes
 from tinygrad.device import Compiler, CompileError
-from tinygrad.helpers import getenv, cpu_objdump, amdgpu_disassemble, unwrap, DEBUG
+from tinygrad.helpers import getenv, capstone_flatdump, amdgpu_disassemble, unwrap, DEBUG
+from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.runtime.autogen import llvm
 
 def cerr(): return ctypes.pointer(ctypes.pointer(ctypes.c_char()))
@@ -10,6 +11,7 @@ def expect(x, err, ret=None):
   return ret
 
 class LLVMCompiler(Compiler):
+  jit = True
   def __init__(self, arch:str, processor:str, feats:str, cache_key=None):
     for component in ['Target', 'TargetInfo', 'TargetMC', 'AsmParser', 'AsmPrinter']:
       getattr(llvm, "LLVMInitialize" + {'arm64': 'AArch64', 'x86_64': 'X86', 'riscv64': 'riscv64'}.get(arch, "AMDGPU") + component)()
@@ -41,13 +43,13 @@ class LLVMCompiler(Compiler):
         self.diag_msgs.append(msg)
     self.handle_diag = handle_diag
     llvm.LLVMContextSetDiagnosticHandler(self.context, handle_diag, None)
-    super().__init__(cache_key or f"compile_llvm_{processor}_{feats}{'_opt' if opt else ''}")
+    super().__init__(cache_key or f"compile_llvm_{processor}_{feats}{'_jit' if self.jit else ''}{'_opt' if opt else ''}")
 
   def __del__(self):
     if hasattr(self, 'pbo'): llvm.LLVMDisposePassBuilderOptions(self.pbo)
     if hasattr(self, 'context'): llvm.LLVMContextDispose(self.context)
 
-  def compile(self, src:str) -> bytes:
+  def compile_to_obj(self, src:str) -> bytes:
     self.diag_msgs.clear()
     src_buf = llvm.LLVMCreateMemoryBufferWithMemoryRangeCopy(ctypes.create_string_buffer(src_bytes:=src.encode()), len(src_bytes), b'src')
     mod = expect(llvm.LLVMParseIRInContext(self.context, src_buf, ctypes.pointer(m:=llvm.LLVMModuleRef()), err:=cerr()), err, m)
@@ -62,6 +64,9 @@ class LLVMCompiler(Compiler):
     if self.diag_msgs: raise RuntimeError("llvm diagnostic: " + "\n".join(self.diag_msgs))
     return obj
 
+  def compile(self, src:str) -> bytes: return jit_loader(self.compile_to_obj(src)) if self.jit else self.compile_to_obj(src)
+
+
 class CPULLVMCompiler(LLVMCompiler):
   def __init__(self, arch:list[str], cache_key=None):
     assert len(arch) >= 2, f"invalid arch string: {','.join(arch)!r}, expected '<arch>,<cpu>,[<feats>]' (eg. 'x86_64,znver2')"
@@ -73,9 +78,10 @@ class CPULLVMCompiler(LLVMCompiler):
     # +reserve-x18 here does the same thing as -ffixed-x18 in ClangCompiler, see comments there for why it's needed on arm osx
     super().__init__(self.arch, cpu, ('+reserve-x18,' if self.arch == "arm64" else '') + featstr, cache_key)
 
-  def disassemble(self, lib: bytes): cpu_objdump(lib)
+  def disassemble(self, lib:bytes): capstone_flatdump(lib, self.arch)
 
 class AMDLLVMCompiler(LLVMCompiler):
+  jit = False
   def __init__(self, arch: str):
     self.arch = arch
     super().__init__("AMDGPU", self.arch, "+cumode")

tinygrad/runtime/support/compiler_qcom.py

@@ -1,6 +1,6 @@
-import ctypes, struct, platform, pathlib, shutil, tarfile, tempfile
+import ctypes, struct, platform, pathlib, shutil, subprocess, sys, tarfile, tempfile
 from tinygrad.device import Compiler
-from tinygrad.helpers import DEBUG, system, fetch
+from tinygrad.helpers import DEBUG, system, fetch, unwrap
 from tinygrad.runtime.support.compiler_mesa import disas_adreno
 # see https://github.com/sirhcm/tinydreno
 from tinygrad.runtime.autogen import llvm_qcom
@@ -12,11 +12,12 @@ 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, root = arch, 0x6030001, tempfile.TemporaryDirectory(), pathlib.Path(__file__).parents[3]
+      self.arch, self.chip_id, self.fs = arch, 0x6030001, tempfile.TemporaryDirectory()
       with tarfile.open(fetch('https://git.tinygrad.win/sirhcm/images/releases/download/v2/qcomcl.tar.gz')) as t: t.extractall(fs:=self.fs.name)
-      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)
+      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)
     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()
@@ -31,7 +32,10 @@ class QCOMCompiler(Compiler):
     return handle
 
   def compile(self, src) -> bytes:
-    if platform.machine() != "aarch64": return self.compile_server(src, self.compiler_process)
+    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")
     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)))
@@ -44,3 +48,12 @@ 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

@@ -1,13 +0,0 @@
-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/elf.py

@@ -1,4 +1,4 @@
-import struct, ctypes
+import struct, ctypes, ctypes.util
 from dataclasses import dataclass
 from tinygrad.helpers import getbits, i2u, unwrap
 from tinygrad.runtime.autogen import libc
@@ -6,13 +6,13 @@ from tinygrad.runtime.autogen import libc
 @dataclass(frozen=True)
 class ElfSection: name:str; header:libc.Elf64_Shdr|libc.Elf32_Shdr; content:bytes # noqa: E702
 
-def link_sym(sym:str, libs:list[ctypes.CDLL]) -> int:
+def link_sym(sym:str, libs:list[str]) -> int:
   for lib in libs:
-    try: return unwrap(ctypes.cast(getattr(lib, sym), ctypes.c_void_p).value)
+    try: return unwrap(ctypes.cast(getattr(ctypes.CDLL(ctypes.util.find_library(lib)), sym), ctypes.c_void_p).value)
     except (OSError, AttributeError): pass
   raise RuntimeError(f'Attempting to relocate against an undefined symbol {sym}')
 
-def elf_loader(blob:bytes, force_section_align:int=1, link_libs:list[ctypes.CDLL]|None=None) -> tuple[memoryview, list[ElfSection], list[tuple]]:
+def elf_loader(blob:bytes, force_section_align:int=1, link_libs:list[str]|None=None) -> tuple[memoryview, list[ElfSection], list[tuple]]:
   assert blob[:4] == libc.ELFMAG.encode(), "blob is not an ELF, missing magic bytes"
   ecls = {libc.ELFCLASS32: "Elf32", libc.ELFCLASS64: "Elf64"}[blob[libc.EI_CLASS]]
 
@@ -49,7 +49,7 @@ def elf_loader(blob:bytes, force_section_align:int=1, link_libs:list[ctypes.CDLL
 
   return memoryview(image), sections, relocs
 
-def jit_loader(obj: bytes, base:int=0, link_libs:list[ctypes.CDLL]|None=None) -> bytes:
+def jit_loader(obj: bytes, base:int=0, link_libs:list[str]|None=None) -> bytes:
   image_, _, relocs = elf_loader(obj, link_libs=link_libs)
   image = bytearray(image_)
 

tinygrad/runtime/support/hcq.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 from typing import cast, Callable, Type, TypeVar, Generic, Any
-import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, itertools
+import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, functools, itertools
+from dataclasses import replace
 try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
-from tinygrad.helpers import DEV, PROFILE, getenv, to_mv, from_mv, cpu_profile, ProfileRangeEvent, unwrap
+from tinygrad.helpers import DEV, PROFILE, getenv, to_mv, from_mv, cpu_profile, ProfileRangeEvent, select_first_inited, select_by_name, unwrap
 from tinygrad.helpers import suppress_finalizing, pluralize, TracingKey
 from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, Program, TinyELF
 from tinygrad.uop.ops import sym_infer, sint, UOp
@@ -392,6 +393,8 @@ class HCQCompiled(Compiled, Generic[SignalType]):
   def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:list[type[Renderer]], runtime:type[Program]|None,
                signal_t:Type[SignalType]|None=None, comp_queue_t:Callable[..., HWQueue]|None=None, copy_queue_t:Callable[..., HWQueue]|None=None,
                kernargs_size=(16 << 20), sigalloc_size=0x1000, can_recover:bool=False, arch=None):
+    self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
+
     from tinygrad.runtime.graph.hcq import HCQGraph
     super().__init__(device, allocator, compilers, runtime, HCQGraph, arch=arch)
 
@@ -421,6 +424,8 @@ class HCQCompiled(Compiled, Generic[SignalType]):
 
     if self._is_cpu(): HCQCompiled.cpu_devices.append(self)
 
+  def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1
+
   def synchronize(self, timeout:int|None=None):
     if self.error_state is not None: raise self.error_state
     if not hasattr(self, 'timeline_signal'): return
@@ -486,6 +491,16 @@ class HCQCompiled(Compiled, Generic[SignalType]):
       buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False
     return buf, realloced
 
+  def _select_iface(self):
+    assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "",  \
+      f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
+    assert hasattr(self, "ifaces"), "must have ifaces to select an iface"
+    t = DEV.target(dev:=type(self).__name__[:-6])
+    filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
+    filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fallback to mock ifaces
+    return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in filtered],
+                               f"No interface for {dev}:{self.device_id} is available")
+
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
 
   def rdma_dev(self):
@@ -497,7 +512,9 @@ class HCQCompiled(Compiled, Generic[SignalType]):
   def finalize(self):
     try: self.synchronize() # Try to finalize device in any case.
     except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
-    super().finalize()
+
+    # If the device has an interface, call its device_fini method to clean up resources.
+    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
 
 class HCQBuffer:
   def __init__(self, va_addr:sint, size:int, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None, owner:Any=None):

tinygrad/runtime/support/hcq2.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
-from typing import cast, TypeVar, Generic, Any, Sequence, Iterable
+from typing import cast, Callable, 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, panic, ContextVar, perf_counter_us, Context
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, 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.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
@@ -23,16 +23,17 @@ HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
 HCQ_RUNTIME_DEV = ContextVar("HCQ_RUNTIME_DEV", "CPU")
 
 HCQ_DEVS = frozenset(("AMD", "CPU"))
-HCQ_CACHE_TAGS = frozenset(("program", "systems"))
+HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))
+HCQ_CACHE_TAGS = frozenset(("program", "systems", "template"))
 
 @dataclass(frozen=True)
 class HCQInfo:
   device:tuple[str, ...]
   estimates:Estimates = Estimates()
 
-  input_idxs:tuple[tuple[tuple[str, ...], tuple[int, ...]], ...] = () # per inputs table: (devices, indexes into input_uops)
-  inputs:int|None = None # index of the inputs table in call.src
-  kernels:tuple[tuple[tuple[str, ...], UOp, tuple[int, ...]], ...] = () # per kernel: (devices, a call carrying its name and estimates, timestamps)
+  input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call
+  inputs:int|None = None
+  kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...]], ...] = ()
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
 
@@ -48,16 +49,9 @@ 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, ...]:
-  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))
+  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 make_binary_patch(buf:UOp, blob:bytes) -> UOp:
   data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)
@@ -84,8 +78,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 = [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))))
+  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)]))
 
 # *****************
 # 0.1. prep: replace buffers with params
@@ -100,11 +94,10 @@ pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_
 # *****************
 # 1.1. prep: staging copies
 
-def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not all_devices_in(b.device, HCQ_DEVS)
+def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not all_devices_in(b.device, HCQ_P2P_DEVS)
 
 def _get_enqueue_devs(call:UOp) -> Any|None:
-  if not (bufs:=call.src[1:]) or not all(all_devices_in(b.device, HCQ_DEVS) for b in bufs): return None
-  if call.src[0].op is Ops.COPY: bufs = bufs[::-1] # copies push from the src device: p2p writes are faster than reads
+  if not (bufs:=call.src[1:]) or not all(all_devices_in(b.device, HCQ_P2P_DEVS) for b in bufs): return None
   devs = min(bufs, key=lambda b: to_tuple(b.device)[0].startswith("CPU")).device # prio to enqueue on not CPU device
   return devs if all_devices_in(devs, HCQ_DEVS) else None
 
@@ -223,7 +216,7 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
     # and make hcq call
     name, info = get_call_name(call, get_call_arg_uops(call)), HCQInfo(devices, estimate_uop(call))
     ts_ids = [next(UOp.unique_num) for _ in range(2)] if profile else []
-    kerns.append((devices, make_call(name, call.src[0], info), tuple(ts_ids)))
+    kerns.append((devices, name, info.estimates, tuple(ts_ids)))
 
     ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_signal(devices, s),)) for s in ts_ids]
     q += ts_ins[:1] + [call.replace(arg=replace(call.arg, aux=info))] + ts_ins[1:]
@@ -252,7 +245,7 @@ def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
   devs, queue = get_submit(calls[0]).src[0].arg
   body = make_submit(*[cmd for c in calls for cmd in get_submit(c).src[0].src], devs=devs, queue=queue).sink()
   return make_call(f"submit {queue} ({len(calls)})", body,
-    replace(calls[0].arg.aux, estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()).simplify()))
+    replace(calls[0].arg.aux, estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates())))
 
 def merge_queues(linear:UOp) -> UOp:
   new_src:list[UOp] = []
@@ -314,15 +307,27 @@ 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 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 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
 
-  # 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 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
 
 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))
 
@@ -336,14 +341,19 @@ 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
-  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})
+  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})
+
+  if inputs: # fence inputs
+    fills.append((t:=tables[0][0]).after(make_binary_patch(t, bytes(t.max_numel() * 8)))) # zeroed at link, slot 0 is the host fence
+    body = body.replace(src=(UOp.sink(*body.src[0].src, t.after(*body.src[0].src).index(0).store(0)),)) # open it once consumed
 
   lt_srcs = collections.defaultdict(list)
   for p in lt_patches: lt_srcs[p.buf_uop].append(p)
   return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()], *fills),
-    arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=((call.arg.aux.device,
-      tuple(sorted(dedup(b.arg.slot for g in inputs for b in unwrap_mstack(g.buf_uop))))),) if inputs else call.arg.aux.input_idxs)))
+    arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(sorted(dedup(b.arg.slot for g in inputs for b in unwrap_mstack(g.buf_uop)))))))
 pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
 
 # *****************
@@ -362,13 +372,14 @@ def replace_params(call:UOp) -> UOp|None:
 
   # keep buffers whose addresses become link-time constants alive and mapped
   held = args + [r.without_after for r in refhold]
-  addrs = dedup([g.src[0].without_after for g in call.toposort() if g.op is Ops.GETADDR])
+  addrs = dedup([g.src[0].without_after for x in call.src for g in x.toposort() if g.op is Ops.GETADDR])
   refhold += [a for a in addrs if a not in held and all(b.op is not Ops.PARAM or b.tag is not None for b in unwrap_mstack(a))]
 
   sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=HCQ_RUNTIME_DEV.value, volatile=b.op is Ops.PARAM and b.arg.volatile)
          for i,u in enumerate(c_args)} | {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM} | _rank_ranges(tops)
-  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args + refhold) if u.without_after.tag == "inputs"), None))
-  return call.replace(src=(body.substitute(sub).replace(arg="hcq_args"), *c_args, *refhold), arg=replace(call.arg, aux=info))
+  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.without_after.tag == "inputs"), None))
+  return call.replace(src=(body.substitute(sub).replace(arg="hcq_args"), *c_args, *refhold),
+                      arg=replace(call.arg, aux=info)) # TODO: call.after(*refhold)?
 pm_replace_params = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), replace_params)])
 
@@ -414,44 +425,8 @@ def callify_hcq(call:UOp, cf:UOp) -> UOp:
 pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, src=(
   UPat(Ops.CUSTOM_FUNCTION, arg="hcq_args", src=(UPat(Ops.SINK),), name="cf"),), name="call", allow_any_len=True), callify_hcq)])
 
-# *****************
-# 9. merge submitters
-
-def _lane_arg(a:UOp, lane:int, table:UOp) -> UOp: return table if a.tag == "inputs" else a.mselect(lane) if len(to_tuple(a.device)) > 1 else a
-
-def merge_batch(batch:list[UOp]) -> UOp:
-  tables = UOp.variable("hcq_inputs_ptr", 0, 2**64-1, dtypes.uint64, param=True)
-  lanes = [(c, j, sum(len(idxs) * 8 for _, idxs in c.arg.aux.input_idxs)) for c in batch for j in range(len(c.arg.aux.device))] # (call, lane, bytes)
-  offs = itertools.accumulate((table_bytes for _, _, table_bytes in lanes), initial=0) # every lane owns the next table of the region
-  cmds = [c.src[0].src[0].call(*[_lane_arg(a.without_after, j, tables + off) for a in c.src[1:]], UOp.variable("_device_num", 0, 1 << 30).bind(j))
-          for (c, j, _), off in zip(lanes, offs)]
-
-  info = HCQInfo((HCQ_RUNTIME_DEV.value,), sum((c.arg.aux.estimates for c in batch), start=Estimates()).simplify(),
-                 input_idxs=tuple(x for c in batch for x in c.arg.aux.input_idxs), kernels=tuple(k for c in batch for k in c.arg.aux.kernels))
-  body = UOp.custom_function("hcq", make_submit(*cmds, devs=HCQ_RUNTIME_DEV.value, queue="SUBMIT:0").sink())
-  return body.call(*[s for c in batch for s in c.src[1:] if s.without_after.tag != "inputs"], name=f"hcq_submitter ({len(batch)})", aux=info)
-
-def merge_submitters(linear:UOp) -> UOp:
-  batches = [(k, list(g)) for k, g in itertools.groupby(linear.src, key=lambda c: isinstance(c.arg.aux, HCQInfo))]
-  return linear.replace(src=tuple(c for is_hcq, b in batches for c in ([merge_batch(b)] if is_hcq else b)))
-
-# *****************
-# hcq schedule
-
 hcq_compile_cache:dict[tuple[bytes, bool], UOp] = {}
 
-def hcq_lower(linear:UOp, pm_encode:PatternMatcher) -> UOp:
-  # lowering to hcq ir
-  linear = graph_rewrite(linear, pm_encode, walk=True, name="encode and pack", enter_calls=True)
-
-  # patches and runtime uops
-  linear = graph_rewrite(linear, pm_early_simplify+symbolic, bottom_up=False, name="simplify patches", enter_calls=True)
-  linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
-
-  # and compile it
-  linear = graph_rewrite(linear, pm_replace_params, name="replace params")
-  return graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True)
-
 @rewrite_group(lambda linear,input_uops,profile,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
 def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
   if input_uops is not None:
@@ -466,9 +441,16 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
     # schedule
     linear = graph_rewrite(linear, pm_schedule_and_merge, ctx=({s:p for p,s in back_map.items()}, profile), walk=True, name="schedule and merge hcq")
 
-    # lower to hcq programs, then pack the programs of every batch into one C submitter (needs a C runtime device for the program addresses)
-    linear = hcq_lower(linear, pm_encode_cmdbufs+pm_pack_placeholders)
-    final_linear = hcq_compile_cache[cache_key] = hcq_lower(merge_submitters(linear), pm_encode_cmdbufs) if HCQ_RUNTIME_DEV.value == "CPU" else linear
+    # lowering to hcq ir
+    linear = graph_rewrite(linear, pm_encode_cmdbufs+pm_pack_placeholders, walk=True, name="encode and pack", enter_calls=True)
+
+    # patches and runtime uops
+    linear = graph_rewrite(linear, pm_early_simplify+symbolic, bottom_up=False, name="simplify patches", enter_calls=True)
+    linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
+
+    # and compile it
+    linear = graph_rewrite(linear, pm_replace_params, name="replace params")
+    final_linear = hcq_compile_cache[cache_key] = graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True)
 
   return final_linear
 
@@ -485,7 +467,7 @@ pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, name="buf"), bufferize_buf)])
 # 7. resolve patches
 
 def push_stack(op, s): return UOp(Ops.STACK,
-  src=tuple(op.replace(dtype=op.dtype, src=tuple(x if y is s else y for y in op.src)) for x in s.src))
+  src=tuple(op.replace(dtype=op.dtype.scalar(), src=tuple(x if y is s else y for y in op.src)) for x in s.src))
 
 def fold_binary(buf:UOp, blob:UOp) -> UOp:
   for b in (m.bufs if isinstance(m:=buf.buffer, MultiBuffer) else (m,)):
@@ -550,6 +532,7 @@ class HCQ2Compiled(Compiled):
   wait_timeout_ms: float = 30000.0
 
   def __init__(self, device:str, allocator:HCQAllocator, compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
+    self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
     self.can_recover = can_recover
 
     self.pm_bufferize = PatternMatcher([
@@ -562,6 +545,7 @@ class HCQ2Compiled(Compiled):
 
     super().__init__(device, allocator, compilers, runtime, None, arch=arch)
 
+    self.rt_buffer = Buffer(self.device, 64 << 20, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True))
     self.rt_allocator = BumpAllocator(64 << 20)
     self.prof_ents:dict[int, ProfileGraphEntry] = {}
 
@@ -585,13 +569,9 @@ class HCQ2Compiled(Compiled):
       tdiffs.append((st+perf_counter_us())/2 - gpu)
     Compiled.profile_events.append(ProfileDeviceEvent(self.device, statistics.median(tdiffs), self.device_props()))
 
-  @functools.cached_property
-  def rt_buffer(self) -> Buffer:
-    return Buffer(self.device, self.rt_allocator.size, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True), preallocate=True)
-
   def new_buffer(self, b:UOp, cache:bool) -> Buffer:
     if cache or b.tag in HCQ_CACHE_TAGS:
-      return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(uncached=b.tag not in ("program","kernargs"), cpu_access=True,nolru=True))
+      return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(uncached=True, cpu_access=True, nolru=True))
     return self.rt_buffer.view(b.max_numel(), b.dtype, self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))
 
   @functools.cache
@@ -600,31 +580,42 @@ class HCQ2Compiled(Compiled):
     buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
     return buf
 
-  def _wait_signal(self, sig:memoryview, value:int, timeout:int|None=None):
-    timeout = timeout if timeout is not None and self.can_recover else None
-    st, done = time.perf_counter(), sig[0]
-    while done < value:
-      if done != (done:=sig[0]): st = time.perf_counter()
-      elif time.perf_counter() - st > (timeout or self.wait_timeout_ms) / 1000: self.on_device_hang()
-
   def synchronize(self, timeout:int|None=None):
     if HCQ_RUNTIME_DEV.value != self.device: Device[HCQ_RUNTIME_DEV.value].synchronize()
 
     sig = self.signal("timeline").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
     tl = self.signal("value", 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
-    self._wait_signal(sig, tl[0] - 1, timeout)
+    timeout = timeout if timeout is not None and self.can_recover else None
+    st, done = time.perf_counter(), sig[0]
+    while done < tl[0] - 1:
+      if done != (done:=sig[0]): st = time.perf_counter()
+      elif time.perf_counter() - st > (timeout or self.wait_timeout_ms) / 1000: self.on_device_hang()
     if self.prof_ents: self.collect_prof()
 
   def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected")
 
   def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
 
+  def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1
+
+  def _select_iface(self):
+    assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "",  \
+      f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
+    assert hasattr(self, "ifaces"), "must have ifaces to select an iface"
+    t = DEV.target(dev:=type(self).__name__[:-6])
+    filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
+    filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fall back to mock ifaces
+    return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in filtered],
+                               f"No interface for {dev}:{self.device_id} is available")
+
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
 
   def finalize(self):
     try: self.synchronize() # try to finalize the device in any case
     except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
-    super().finalize()
+
+    # if the device has an interface, call device_fini to clean up resources
+    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
 
 @dataclass
 class HCQ2Buffer:

tinygrad/schedule/multi.py

@@ -22,14 +22,14 @@ def mstack_early_shrink(ms:UOp, shrink:UOp):
 
 def lower_broadcast_copy(c:UOp, x:UOp):
   if not (isinstance(c.device, tuple) and isinstance(x.device, str)): return None
-  if (sx:=x.simplify()).device is None: return UOp(Ops.MSTACK, src=(sx,)*len(c.device))
+  if (sx:=x.simplify()).device is None and sx.base.op is Ops.CONST: return UOp(Ops.MSTACK, src=(sx,)*len(c.device))
   return UOp(Ops.MSTACK, src=tuple(x.copy_to_device(d) for d in c.device))
 
 replace_allreduce = PatternMatcher([
   # BROADCAST: explicitly expand broadcast copies and combine with MSTACK
-  (UPat(Ops.COPY, name="c", src=(UPat(name="x"),)), lower_broadcast_copy),
+  (UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lower_broadcast_copy),
   # COPY_TO_ONE: if copying from multidevice to one, MSELECT the first (TODO: a little from each?)
-  (UPat(Ops.COPY, name="c", src=(UPat(name="x"),)), lambda c,x:
+  (UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lambda c,x:
     x.mselect(0).copy_to_device(c.device) if isinstance(c.device, str) and isinstance(x.device, tuple) else None),
   # MSELECT on MSTACK is replaced with nothing
   (UPat(Ops.MSELECT, src=(UPat(Ops.MSTACK, name="mstack"),), name="ms"), lambda mstack, ms: mstack.src[ms.arg]),

tinygrad/schedule/rangeify.py

@@ -4,7 +4,7 @@ import itertools
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype, strong_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element
-from tinygrad.uop.symbolic import symbolic
+from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
@@ -313,9 +313,9 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([
    lambda idx,after: idx.const_like(Invalid) if after_all_invalid(after) else None),
   # hack if a noop turned to a const
   (UPat(Ops.NOOP, src=(UPat.cvar("c"),)), lambda c: c),
-  # a deviceless MSTACK src is the same value on every device, so indexing the stack is just indexing that value
-  (UPat(Ops.MSTACK, src=(UPat.var("s"),), allow_any_len=True).f(Ops.INDEX, allow_any_len=True, name="idx"),
-   lambda s,idx: idx.replace(src=(s,)+idx.src[1:]) if s.device is None else None),
+  # mstack on CONST is CONST
+  (UPat(Ops.MSTACK, src=(UPat.var("s"),), allow_any_len=True).f(Ops.INDEX, allow_any_len=True),
+   lambda s: c if (c:=s.base).op is Ops.CONST else None),
 ])
 
 pm_remove_bufferize = PatternMatcher([
@@ -327,9 +327,6 @@ pm_remove_bufferize = PatternMatcher([
   (UPat(Ops.END, src=(UPat(Ops.NOOP, name="x"),), allow_any_len=True), lambda x: x),
 ])
 
-def strip_zero_offset_shrink(x:UOp) -> UOp:
-  return x.src[0] if x.op is Ops.SHRINK and all(resolve(start == 0, False) for start,_ in x.marg) else x
-
 def no_indexing_calls(u:UOp):
   new_srcs = []
   for x in u.src:
@@ -339,9 +336,8 @@ def no_indexing_calls(u:UOp):
       new_srcs.append(x.src[0])
     elif x.op is Ops.SHRINK:
       # SHRINK with offset 0 is fine
-      new_srcs.append(strip_zero_offset_shrink(x))
-    elif x.op is Ops.MSTACK:
-      new_srcs.append(x.replace(src=tuple(strip_zero_offset_shrink(s) for s in x.src)))
+      # TODO: check offset
+      new_srcs.append(x.src[0])
     else:
       # everything else we pass through
       new_srcs.append(x)
@@ -588,7 +584,7 @@ def get_kernel_graph(sink:UOp) -> UOp:
   tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
 
   tsink = graph_rewrite(tsink,
-                        symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize,
+                        symbolic+pm_fold_cast_const+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize+pm_no_indexing_calls,
                         name="symbolic+reduce_collapse+debuf")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 
@@ -599,7 +595,6 @@ def get_kernel_graph(sink:UOp) -> UOp:
   paramarg_start: int = max([-1]+slots) + 1
   tsink = graph_rewrite(tsink, pm_add_buffers+pm_add_param_range_tags, ctx=itertools.count(paramarg_start), bottom_up=True, name="stage to store")
   tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
-  tsink = graph_rewrite(tsink, pm_no_indexing_calls, name="remove indexing from call args")
 
   if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
   if SPEC: