27b 935/47

# Conflicts:
#	extra/gemm/amd_flash_attention.py
#	tinygrad/codegen/__init__.py
#	tinygrad/runtime/ops_cpu.py
#	tinygrad/uop/ops.py

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

@@ -137,6 +137,7 @@ runs:
       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
         sudo chown -R $USER:$USER /var/cache/apt/archives
 
         echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
@@ -214,6 +215,7 @@ runs:
           sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
         fi
 
+        sudo mkdir -p /var/cache/apt/archives
         sudo chown -R $USER:$USER /var/cache/apt/archives/
 
     - name: Add clang to PATH (Linux)

.github/workflows/benchmark.yml

@@ -117,10 +117,10 @@ jobs:
       run: python3 test/external/process_replay/reset.py
     - name: Run llama3.2
       run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.5
-      # qwen3.5:35b-a3b doesn't fit on mac
+    - name: Run qwen3.6
+      # qwen3.6:35b-a3b doesn't fit on mac
       if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.5:35b-a3b --benchmark --warmup
+      run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
     - name: Run olmoe
       # just metal for now
       if: ${{ matrix.dev == 'METAL' }}

.github/workflows/test.yml

@@ -167,6 +167,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: unittest-13
+        python-version: '3.11'
         deps: testing_unit
         llvm: 'true'
         amd: 'true'
@@ -176,13 +177,14 @@ jobs:
       run: |
         DEV=NULL python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
         DEV=NULL VIZ=1 python3 -m pytest -n=auto test/null/test_viz.py
+        DEBUG=7 python -m tinygrad.viz.cli --json | jq empty
     # TODO: too slow
     # - name: Run SDXL on NULL backend
     #   run: DEV=NULL DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
     - name: Run Clip tests for SD MLPerf on NULL backend
       run: DEV=NULL python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
     - name: Run AMD emulated BERT training on NULL backend
-      run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     # TODO: support fake weights
     #- name: Run LLaMA 7B on 4 fake devices
     #  run: DEV=NULL python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -199,6 +201,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: unittest-13
+        python-version: '3.11'
         pydeps: "pre-commit"
         deps: testing_unit
         llvm: 'true'
@@ -291,7 +294,7 @@ jobs:
           llvm: 'true'
       - name: Test openpilot model kernel count and gate usage
         run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1391 ALLOWED_GATED_READ_IMAGE=58 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=54 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
           # IMAGE_PITCH_ALIGNMENT=64 matches adreno 630
       - name: Test openpilot CL compile fp32 (test correctness)
         run: |
@@ -478,7 +481,7 @@ jobs:
           wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
           echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
           sudo apt-get update
-          sudo apt-get install llvm-21 llvm-21-tools cloc
+          sudo apt-get install -y llvm-21 llvm-21-tools cloc
       - name: Install rocprof-trace-decoder
         run: sudo PYTHONPATH="." ./extra/sqtt/install_rocprof_decoder.py
       - name: Run AMD renderer tests
@@ -628,7 +631,7 @@ jobs:
 
   unittestmacos:
     name: MacOS (unit)
-    runs-on: &macos macos-26
+    runs-on: macos-26
     timeout-minutes: 20
     steps:
     - name: Checkout Code
@@ -677,7 +680,7 @@ jobs:
       matrix:
         group: [1, 2]
     name: MacOS (DEV=METAL) (${{ matrix.group }})
-    runs-on: *macos
+    runs-on: macos-26
     timeout-minutes: 20
     env:
       DEV: METAL
@@ -709,7 +712,7 @@ jobs:
           - 'WEBGPU'
 
     name: MacOS (DEV=${{ matrix.dev }})
-    runs-on: *macos
+    runs-on: macos-26
     timeout-minutes: 20
     steps:
     - name: Checkout Code

AGENTS.md

@@ -3,3 +3,4 @@
 - Run tests with `-n12` for speed (e.g. `python -m pytest test/null/test_dtype.py -x -q -n12`)
 - Run `python -m mypy tinygrad/` to typecheck
 - Run `python -m ruff check .` to lint
+- Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules

conftest.py

@@ -0,0 +1,10 @@
+import os, pytest, signal, threading
+
+@pytest.hookimpl(wrapper=True)
+def pytest_runtest_call(item):
+  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 300)), os.kill, args=(os.getpid(), signal.SIGABRT))
+  t.start()
+  try: yield
+  finally:
+    t.cancel()
+    t.join()

docs/abstractions4.py

@@ -88,7 +88,7 @@ def example_3_custom_uop(a:Tensor, correct):
 
     # store all the per lane accumulators to LOCAL
     local_accs = UOp.placeholder((LCLS,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-    local_accs = local_accs.after(local_accs[lane].store(acc[0]).barrier())
+    local_accs = local_accs.after(local_accs[lane].store(acc[0]))
 
     # accumulate LOCALs into a single per CU accumulator
     late_reduce_loop = UOp.range(LCLS, 3, AxisType.REDUCE)

docs/dtypes.md

@@ -1,7 +1,8 @@
 ::: tinygrad.dtype.DType
 
-::: tinygrad.dtype.dtypes
+::: tinygrad.dtype.DTypes
     options:
+        heading: dtypes
         members: true
         members_order: source
         show_labels: false

examples/beautiful_cifar.py

@@ -9,8 +9,7 @@ from extra.lr_scheduler import OneCycleLR
 GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
 
 # override tinygrad defaults
-dtypes.default_float = dtypes.half
-Context(FUSE_OPTIM=1).__enter__()
+Context(DEFAULT_FLOAT=dtypes.half, FUSE_OPTIM=1).__enter__()
 
 # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 batchsize = getenv("BS", 1024)

examples/mlperf/dataloader.py

@@ -1,11 +1,11 @@
 import os, random, pickle, queue, struct, math, functools, hashlib, time
 from typing import List
 from pathlib import Path
-from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
+from multiprocessing import Queue, Process, shared_memory, connection, Lock
 
 import numpy as np
 from tinygrad import dtypes, Tensor
-from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
+from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS
 from tinygrad.nn.state import TensorIO
 
 ### ResNet
@@ -131,7 +131,7 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir
     else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}")
     Y = [None] * (batch_size*BATCH_COUNT)
 
-    for _ in range(cpu_count()):
+    for _ in range(NUM_CPU_THREADS.value):
       p = Process(target=loader_process, args=(q_in, q_out, X, seed))
       p.daemon = True
       p.start()
@@ -212,7 +212,7 @@ def batch_load_train_bert(BS:int, seed:int|None=None):
     rng.shuffle(fs)
     train_files.append(fs.pop(0))
 
-  cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
+  cycle_length = min(NUM_CPU_THREADS.value, len(train_files))
   assert cycle_length > 0, "cycle_length must be greater than 0"
 
   dataset = InterleavedDataset(train_files, cycle_length)
@@ -301,7 +301,7 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
     X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}")
     Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}")
 
-    for _ in range(cpu_count()):
+    for _ in range(NUM_CPU_THREADS.value):
       proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y))
       proc.daemon = True
       proc.start()
@@ -437,7 +437,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
   dataset_iter = iter(image_ids)
 
   try:
-    for _ in range(cpu_count()):
+    for _ in range(NUM_CPU_THREADS.value):
       proc = Process(
         target=load_retinanet_data,
         args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),

examples/mlperf/model_train.py

@@ -1285,7 +1285,7 @@ def train_llama3():
   from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW
+  from examples.mlperf.optim import GradAccClipAdamW, clip_grads
 
   INITMLPERF = getenv("INITMLPERF")
   RUNMLPERF = getenv("RUNMLPERF")
@@ -1482,7 +1482,8 @@ def train_llama3():
 
   @TinyJit
   def optim_step():
-    grad_norm = optim.fstep(grads)
+    grad_norm = clip_grads(grads, grad_acc, 1.0)
+    optim.fstep(grads, grad_norm)
     scheduler.step()
 
     for g in grads: g.assign(0)
@@ -1667,7 +1668,7 @@ def train_llama3():
 def train_gptoss():
   from examples.mlperf.models.gpt_oss import GPTOSS, GPT_OSS_20B, apply_grad, FP8_DTYPE
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW
+  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, clip_grads
 
   BENCHMARK = getenv("BENCHMARK")
 
@@ -1733,7 +1734,12 @@ def train_gptoss():
   is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
   is_fake_offload = Device.DEFAULT == "NULL"
   optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
-  optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
+  params_wd = [p for p in params if p.ndim >= 3]
+  params_no_wd = [p for p in params if p.ndim < 3]
+  optim = GradAccClipAdamWGroup(
+    GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device),
+    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=grad_acc, device=optim_device),
+  )
 
   for p in optim.params:
     grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
@@ -1775,7 +1781,8 @@ def train_gptoss():
 
   @TinyJit
   def optim_step():
-    grad_norm = optim.fstep(grads)
+    grad_norm = clip_grads(grads, grad_acc, 1.0)
+    optim.fstep(grads, grad_norm)
     scheduler.step()
 
     for g in grads: g.assign(0)

examples/mlperf/models/gpt_oss.py

@@ -10,7 +10,7 @@ if __name__ == "__main__":
 from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
 from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
-from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
 from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
 
@@ -70,6 +70,11 @@ def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor:
   x_linear = x_linear.clamp(-limit, limit)
   return (x_glu * (alpha * x_glu).sigmoid()) * (x_linear + 1)
 
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
+  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
+  freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
+  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
+
 class GPTOSS:
   def __init__(self, dim:int, n_layers:int, n_heads:int, n_kv_heads:int, head_dim:int, n_experts:int, experts_per_tok:int,
                intermediate_size:int, vocab_size:int, norm_eps:float=1e-5, rope_theta:int=150000, sliding_window:int=128,
@@ -112,14 +117,30 @@ class GPTOSS:
     w_q, w_e8, _ = quantize_mxfp8(w)
     return w_q, w_e8.is_param_(False)
 
-  def _attn_mask(self, seqlen:int, sliding:bool, dtype) -> Tensor:
+  def _attn_mask(self, seqlen:int, dtype) -> Tensor:
     i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
-    allowed = j <= i
-    if sliding: allowed = allowed & (i - j < self.sliding_window)
-    return allowed.where(0.0, -1e30).cast(dtype).contiguous()
+    return (j <= i).where(0.0, -1e30).cast(dtype).contiguous()
 
-  def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wqkv_scale:Tensor,
-                wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
+  def _sliding_attention(self, xq:Tensor, xk:Tensor, xv:Tensor, sinks:Tensor) -> Tensor:
+    bsz, seqlen, H, hd = xq.shape
+    KV, R, W = self.n_kv_heads, self.n_rep, self.sliding_window
+    assert seqlen % W == 0, f"seqlen {seqlen} must be a multiple of sliding_window {W} for banded attention"
+    nb = seqlen // W
+    q = xq.reshape(bsz, seqlen, KV, R, hd).permute(0, 2, 3, 1, 4).reshape(bsz, KV, R, nb, W, hd).float()
+    k, v = (x.permute(0, 2, 1, 3).reshape(bsz, KV, 1, nb, W, hd).float() for x in (xk, xv))
+    kk, vv = (x.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb].cat(x, dim=-2) for x in (k, v))
+    sc = (q @ kk.transpose(-1, -2)) * self.sm_scale   # (B,KV,R,nb,W,2W)
+    i, j, pv = Tensor.arange(W).reshape(W, 1), Tensor.arange(2 * W).reshape(1, 2 * W), Tensor.arange(nb).reshape(nb, 1, 1) >= 1
+    sc = ((j > i) & (j <= i + W) & (pv | (j >= W))).where(sc, -float("inf"))
+    sink = sinks.reshape(1, KV, R, 1, 1, 1).float()
+    m = sc.max(-1, keepdim=True).maximum(sink)
+    e = (sc - m).exp()
+    p = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
+    attn = p @ vv.cast(dtypes.bfloat16)
+    return attn.reshape(bsz, KV, R, seqlen, hd).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, H * hd)
+
+  def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, *, attention_norm:Tensor, wqkv:Tensor,
+                wqkv_scale:Tensor, wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
     bsz, seqlen, _ = x.shape
     x_normed, rrms = rmsnorm(x, self.norm_eps)
     qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
@@ -127,16 +148,23 @@ class GPTOSS:
     xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
     xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
     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)
 
-    xq = xq.cast(dtypes.bfloat16).reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
-    xk = xk.cast(dtypes.bfloat16).permute(0, 2, 1, 3).unsqueeze(2)
-    xv = xv.cast(dtypes.bfloat16).permute(0, 2, 1, 3).unsqueeze(2)
-    scores = (xq @ xk.transpose(-2, -1)).float() * self.sm_scale + mask
-    sink = sinks.reshape(1, self.n_kv_heads, self.n_rep, 1, 1).float()
-    m = scores.max(-1, keepdim=True).maximum(sink)
-    e = (scores - m).exp()
-    w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
-    attn = (w @ xv).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
+    if sliding:
+      attn = self._sliding_attention(xq, xk, xv, sinks)
+    elif getenv("HK_FLASH_ATTENTION"):
+      from extra.thunder.amd.fa import flash_attention
+      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
+      attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
+    else:
+      xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
+      xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
+      scores = (xqm @ xkm.transpose(-2, -1)).float() * self.sm_scale + mask
+      sink = sinks.reshape(1, self.n_kv_heads, self.n_rep, 1, 1).float()
+      m = scores.max(-1, keepdim=True).maximum(sink)
+      e = (scores - m).exp()
+      w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
+      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]
@@ -160,8 +188,8 @@ class GPTOSS:
     return out, [x_normed, rrms]
 
   @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
-    attn, attn_saves = self.attention(x, freqs_cis, mask, **attn_kwargs)
+  def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
+    attn, attn_saves = self.attention(x, freqs_cis, mask, sliding, **attn_kwargs)
     h = x + attn
     ffn, ffn_saves = self.feed_forward(h, **ffn_kwargs)
     h = h + ffn
@@ -178,8 +206,7 @@ class GPTOSS:
     h = self.tok_embeddings(tokens)
     bsz, seqlen = tokens.shape
     freqs_cis = self.freqs_cis.cast(h.dtype)[:, :seqlen, :, :, :]
-    mask_full = self._attn_mask(seqlen, False, dtypes.float32)
-    mask_sliding = self._attn_mask(seqlen, True, dtypes.float32)
+    mask_full = None if getenv("HK_FLASH_ATTENTION") else self._attn_mask(seqlen, dtypes.float32)
     for i in range(self.n_layers):
       attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wqkv_scale=self.wqkv_scale[i],
                          wqkv_bias=self.wqkv_bias[i], wo=self.wo[i], wo_scale=self.wo_scale[i], wo_bias=self.wo_bias[i],
@@ -187,8 +214,7 @@ class GPTOSS:
       ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], gate=self.gate[i], gate_bias=self.gate_bias[i],
                         w_gate_up=self.w_gate_up[i], w_gate_up_scale=self.w_gate_up_scale[i], w_gate_up_bias=self.w_gate_up_bias[i],
                         w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
-      mask = mask_sliding if i % 2 == 0 else mask_full
-      h, *_ = self.run_layer(h, freqs_cis, mask, attn_kwargs, ffn_kwargs, save=save)
+      h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
 
     logits = self.norm(h) @ self.output.T
     return logits

examples/mlperf/optim.py

@@ -1,6 +1,6 @@
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.nn.optim import Optimizer
+from tinygrad.nn.optim import Optimizer, OptimizerGroup
 from tinygrad.helpers import FUSE_OPTIM, getenv
 from tinygrad.uop.ops import UOp, Ops
 
@@ -21,6 +21,12 @@ def stochastic_round_bf16(x:Tensor) -> Tensor:
   noise = (noise * 0xFFFF).cast(dtypes.uint32)
   return ((bits + noise) & 0xFFFF0000).bitcast(dtypes.float32).cast(dtypes.bfloat16)
 
+def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
+  for g in grads: g.assign(g / grad_acc)
+  total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+  for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
+  return total_norm
+
 class GradAccClipAdamW(Optimizer):
   def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
     super().__init__(params, lr, device, fused)
@@ -37,45 +43,28 @@ class GradAccClipAdamW(Optimizer):
 
   def _zero_shard(self, t:Tensor) -> Tensor:
     if not self.zero or (t.shape[0] % len(self.device)) != 0: return t
-    return Tensor(t.uop._shard(0, len(self.device)).multi(0)).clone()
+    return Tensor(t.uop._shard(0, len(self.device)).unshard(0)).clone()
 
   def _zero_gather(self, t:Tensor) -> Tensor:
     if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
     n, sz = len(t.device), t.shape[0] // len(t.device)
     return Tensor.cat(*[t[p*sz:(p+1)*sz] for p in range(n)], dim=0)
 
-  def fstep(self, grads:list[Tensor]):
-    if self.fused:
-      out, extra = self._step([], grads)
-      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
-    else:
-      updates, extra = self._step([], grads)
+  def fschedule_step(self, grads:list[Tensor]) -> list[Tensor]:
+    updates, extra = self._step([], grads)
     for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
-    # collect inv_scale tensors attached to fp8 params (set by _apply_update)
     fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
     fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')]
-    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales+fp8_next_inv_scales
+    return extra + self.params + self.buffers + (self.master_params or []) + fp8_inv_scales + fp8_next_inv_scales
 
-    Tensor.realize(*to_realize)
-    return extra[-1]
+  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
+    Tensor.realize(*([grad_norm] if grad_norm is not None else []), *self.fschedule_step(grads))
 
   def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
     grads = list(grads)
 
     for i in range(len(grads)):
       if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
-
-    if self.fused:
-      grads[0].assign(grads[0] / self.grad_acc)
-      total_norm = grads[0].float().square().sum().sqrt()
-      grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
-    else:
-      for i in range(len(grads)):
-        grads[i].assign(grads[i] / self.grad_acc)
-      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
-      for i in range(len(grads)):
-        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype))
-
     ret = []
     self.b1_t *= self.b1
     self.b2_t *= self.b2
@@ -88,7 +77,7 @@ class GradAccClipAdamW(Optimizer):
       v_hat = v_new / (1.0 - self.b2_t)
       up = m_hat / (v_hat.sqrt() + self.eps)
       ret.append(self.lr * up)
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v
 
   def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor:
     w = master if master is not None else t
@@ -109,7 +98,7 @@ class GradAccClipAdamW(Optimizer):
         if self.zero: w_q, w_e8 = self._zero_gather(w_q), self._zero_gather(w_e8)
         new_e8 = w_e8.reshape(t._inv_scale.shape)
         t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
-        ret = w_q.reshape(new_w.shape)
+        ret = w_q.reshape(t.shape)
         return ret.shard_like(t) if offloaded else ret
       from examples.mlperf.models.flat_llama import FP8_MAX
       if IMMEDIATE_SCALE:
@@ -132,3 +121,21 @@ class GradAccClipAdamW(Optimizer):
       return ret.shard_like(t) if offloaded else ret
     out = new_w.cast(t.dtype)
     return out.shard_like(t) if offloaded else out
+
+class GradAccClipAdamWGroup(OptimizerGroup):
+  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
+    offset = 0
+    to_realize = []
+    for o in self.optimizers:
+      n = len(o.params)
+      to_realize += o.fschedule_step(grads[offset:offset+n])
+      offset += n
+    Tensor.realize(*to_realize, *([grad_norm] if grad_norm is not None else []))
+  @property
+  def lr(self): return self.optimizers[0].lr
+  @property
+  def device(self): return self.optimizers[0].device
+  @property
+  def master_params(self):
+    mp = [mp for o in self.optimizers for mp in (o.master_params or [])]
+    return mp if mp else None

extra/benchmark_llm.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+"""Benchmark tinygrad LLM prefill and decode independently."""
+from __future__ import annotations
+
+import argparse, gc, json, statistics, time
+from dataclasses import asdict, dataclass
+
+from tinygrad import Context, Device, Tensor, UOp
+from tinygrad.helpers import fetch, profile_marker
+from tinygrad.llm.cli import models
+from tinygrad.llm.model import Transformer
+
+@dataclass
+class Result:
+  prompt_tokens: int
+  decode_tokens: int
+  time_to_first_token_s: float
+  prefill_tokens_per_s: float
+  decode_tokens_per_s: float
+  decode_p50_ms: float
+  decode_p95_ms: float
+  output_tokens: list[int]
+
+def percentile(values:list[float], percentile:float) -> float:
+  ordered = sorted(values)
+  return ordered[round((len(ordered) - 1) * percentile)]
+
+def synthetic_prompt(length:int, vocab_size:int, salt:int) -> list[int]:
+  assert length > 0 and vocab_size > 256
+  return [256 + salt % (vocab_size - 256)] + [256 + (i * 7919) % (vocab_size - 256) for i in range(1, length)]
+
+def benchmark(model:Transformer, prompt:list[int], decode_tokens:int, chunk_size:int) -> Result:
+  gen = model.generate(prompt.copy(), chunk_size=chunk_size)
+  profile_marker(f"prefill {len(prompt)} start")
+  begin = time.perf_counter()
+  output_tokens = [next(gen)]
+  ttft = time.perf_counter() - begin
+  profile_marker(f"prefill {len(prompt)} end")
+
+  decode_times: list[float] = []
+  profile_marker(f"decode {len(prompt)} start")
+  for _ in range(decode_tokens):
+    begin = time.perf_counter()
+    output_tokens.append(next(gen))
+    decode_times.append(time.perf_counter() - begin)
+  profile_marker(f"decode {len(prompt)} end")
+
+  return Result(len(prompt), decode_tokens, ttft, len(prompt) / ttft, decode_tokens / sum(decode_times),
+                statistics.median(decode_times) * 1e3, percentile(decode_times, 0.95) * 1e3, output_tokens)
+
+def benchmark_decode_position(model:Transformer, position:int, decode_tokens:int) -> Result:
+  token = Tensor([[0]], dtype="int32", device=Device.DEFAULT).realize()
+  temperature = Tensor([0.0], device=Device.DEFAULT).realize()
+  decode_times, output_tokens = [], []
+  for pos in range(position, position + decode_tokens):
+    begin = time.perf_counter()
+    output_tokens.append(int(model(token, UOp.variable("start_pos", 0, model.max_context-1).bind(pos), temperature).realize().item()))
+    decode_times.append(time.perf_counter() - begin)
+  return Result(position, decode_tokens, 0.0, 0.0, decode_tokens / sum(decode_times),
+                statistics.median(decode_times) * 1e3, percentile(decode_times, 0.95) * 1e3, output_tokens)
+
+def main() -> None:
+  parser = argparse.ArgumentParser(description="Measure LLM prefill and steady-state decode speed")
+  parser.add_argument("--model", default="qwen3:0.6b", help="Model preset or local GGUF path")
+  parser.add_argument("--max-context", type=int, default=32768)
+  parser.add_argument("--prompt-tokens", type=int, nargs="+", default=[128, 2048, 8192])
+  parser.add_argument("--decode-tokens", type=int, default=32)
+  parser.add_argument("--decode-position", type=int, nargs="+")
+  parser.add_argument("--chunk-size", type=int, default=256)
+  parser.add_argument("--beam", type=int, default=2)
+  parser.add_argument("--jit-batch-size", type=int, default=448)
+  parser.add_argument("--parallel-compile", type=int, default=12)
+  parser.add_argument("--realize", action="store_true")
+  parser.add_argument("--json", action="store_true")
+  args = parser.parse_args()
+
+  if args.decode_tokens < 1: parser.error("--decode-tokens must be positive")
+  if args.chunk_size < 1: parser.error("--chunk-size must be positive")
+  if args.decode_position is None and max(args.prompt_tokens) + args.decode_tokens >= args.max_context:
+    parser.error("prompt plus decode tokens must fit within --max-context")
+  if args.decode_position is not None and max(args.decode_position) + args.decode_tokens >= args.max_context:
+    parser.error("decode position plus decode tokens must fit within --max-context")
+
+  begin = time.perf_counter()
+  path = fetch(models.get(args.model, args.model))
+  fetched = time.perf_counter()
+  model, kv = Transformer.from_gguf(path, args.max_context, realize=args.realize)
+  loaded = time.perf_counter()
+  vocab_size = len(kv["tokenizer.ggml.tokens"])
+  print(f"startup: fetch={fetched-begin:.2f}s load={loaded-fetched:.2f}s", flush=True)
+  with Context(BEAM=args.beam, JIT_BATCH_SIZE=args.jit_batch_size, PARALLEL_COMPILE=args.parallel_compile):
+    model.warmup(args.chunk_size)
+  startup = time.perf_counter() - begin
+  print(f"startup: warmup={startup-(loaded-begin):.2f}s total={startup:.2f}s", flush=True)
+  gc.freeze()
+
+  results = [benchmark_decode_position(model, pos, args.decode_tokens) for pos in args.decode_position] if args.decode_position is not None else \
+    [benchmark(model, synthetic_prompt(n, vocab_size, salt=i+1), args.decode_tokens, args.chunk_size)
+     for i, n in enumerate(args.prompt_tokens)]
+  if args.json:
+    print(json.dumps({"model": args.model, "max_context": args.max_context, "chunk_size": args.chunk_size,
+                      "beam": args.beam, "jit_batch_size": args.jit_batch_size, "parallel_compile": args.parallel_compile,
+                      "realize": args.realize, "startup_s": startup, "results": [asdict(x) for x in results]}, indent=2))
+    return
+
+  print(f"model={args.model} max_context={args.max_context} chunk_size={args.chunk_size} beam={args.beam} "
+        f"jit_batch_size={args.jit_batch_size} parallel_compile={args.parallel_compile} realize={args.realize} startup={startup:.2f}s")
+  print(f"{'prompt':>8} {'TTFT':>10} {'prefill':>14} {'decode':>14} {'decode p50':>12} {'decode p95':>12}")
+  for result in results:
+    print(f"{result.prompt_tokens:8d} {result.time_to_first_token_s:9.3f}s {result.prefill_tokens_per_s:11.1f} t/s "
+          f"{result.decode_tokens_per_s:11.1f} t/s {result.decode_p50_ms:9.2f} ms {result.decode_p95_ms:9.2f} ms")
+
+if __name__ == "__main__": main()

extra/export_model.py

@@ -241,8 +241,8 @@ export default {model_name};
 def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
   assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
 
-  # NOTE: CPU_COUNT=1, since export does not support threading
-  with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs)
+  # NOTE: NUM_CPU_THREADS=1, since export does not support threading
+  with Context(JIT=2, NUM_CPU_THREADS=1): linear, output_bufs = jit_model(model, *inputs)
   functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
   state = get_state_dict(model)
   weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}

extra/f16_decompress.py

@@ -5,10 +5,10 @@ def bit_extract(x: Tensor, e: int, s: int) -> Tensor:
   return (x >> s) & mask
 
 def u16_to_f16(x: Tensor) -> Tensor:
-  sign = bit_extract(x, 15, 15).float()
+  sign = bit_extract(x, 15, 15).bool()
   exponent = bit_extract(x, 14, 10).float()
   fraction = bit_extract(x, 9, 0).float()
-  return sign.where(-1, 1) * exponent.where((exponent - 15.0).exp2() * (1 + fraction / 1024.0), 6.103515625e-5 * (fraction / 1024.0))
+  return sign.where(-1, 1) * exponent.bool().where((exponent - 15.0).exp2() * (1 + fraction / 1024.0), 6.103515625e-5 * (fraction / 1024.0))
 
 def u32_to_f16(oo: Tensor) -> Tensor:
   f1 = u16_to_f16(oo>>16)

extra/fp8/fp8_linear.py

@@ -18,9 +18,9 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
   SEQ = inp.shape[1]
   OUT = weight.shape[0]
   IN = weight.shape[-1]
-  seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
-  out_idx = UOp.range(OUT, 3, AxisType.LOOP)
-  batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
+  seq_idx = UOp.range(SEQ, 2)
+  out_idx = UOp.range(OUT, 3)
+  batch_idx = UOp.range(output.size//SEQ//OUT, 1)
   reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
   product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
   reduced = product.reduce(reduce_idx, arg=Ops.ADD)
@@ -53,7 +53,7 @@ class FP8Linear:
     x_fp8, x_scale = quantize_to_fp8(x)
     GPUS = self.weight.device
     if isinstance(GPUS, tuple) and len(GPUS) > 1:
-      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.multi(0), device=GPUS)
+      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.unshard(0), device=GPUS)
     else:
       y = Tensor.empty((batch, seq, self.weight.shape[0]), dtype=dtypes.float)
     y = Tensor.custom_kernel(y, x_fp8, w_fp8, fxn=custom_matmul, grad_fxn=custom_matmul_backward)[0]

extra/gemm/amd_copy_matmul.py

@@ -58,8 +58,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
   B_copy = B_local.permute((1,0)) if use_wmma else B_local
   A_store = A_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
   B_store = B_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
-  barrier = UOp.barrier(A_store, B_store)
-  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  A_local, B_local = A_local.after(A_store, B_store), B_local.after(A_store, B_store)
 
   # -- COMPUTE --
   lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
@@ -70,8 +70,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
 
   if use_wmma:
     k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
-    tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
-    tile_n = UOp.range(TN, 201, AxisType.LOOP)
+    tile_m = UOp.range(TM // WMMA_ACC, 200)
+    tile_n = UOp.range(TN, 201)
 
     acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
     a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
@@ -96,8 +96,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
     b_frag = b_frag.reshape(1, TN).expand(TM, TN)
     acc_store = acc.store(acc.after(k) + (a_frag * b_frag))
 
-  # store accumulator and loop
-  acc = acc.after(acc_store.end(k).barrier().end(k_tile))
+  # store accumulator and loop (the barrier at the end of the loop is implied by the LOCAL buffers stored and loaded in the loop)
+  acc = acc.after(acc_store.end(k).end(k_tile))
 
   # store accumulator to output (unified)
   c = c.reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M,

extra/gemm/amd_flash_attention.py

@@ -1,203 +1,33 @@
-from tinygrad import Tensor, UOp, getenv
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
-from tinygrad.dtype import AddrSpace, dtypes
-from tinygrad.helpers import DEBUG, GlobalCounters, Context
-import math
-
-BLOCK_M, BLOCK_N = 64, 64
-WARP_SIZE = 32
-WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
-WAVES_M, WAVES_N = 4, 1
-LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 16
-WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
-THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
-LDS_PAD = 4  # pad LDS rows to reduce bank conflicts
-
-WMMA_ARG = (WMMA_M, WMMA_N, WMMA_K), 'AMD', 32
-LOG2E = math.log2(math.e)
-
-def warp_shfl_xor(val, offset, lane):
-  """Read val from lane ^ offset using ds_bpermute."""
-  idx = ((lane ^ offset) * 4).cast(dtypes.int)
-  if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
-  return UOp(Ops.CUSTOM, dtypes.float, (idx, val),
-             arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_bpermute({0}, __builtin_bit_cast(int, {1})))")
-
-def warp_reduce_max(val, lane):
-  """Tree reduce MAX across LANES_PER_WAVE_N=16 lanes."""
-  for offset in [8, 4, 2, 1]:
-    val = UOp(Ops.MAX, dtypes.float, (val, warp_shfl_xor(val, offset, lane)))
-  return val
-
-def warp_reduce_sum(val, lane):
-  """Tree reduce SUM across LANES_PER_WAVE_N=16 lanes."""
-  for offset in [8, 4, 2, 1]:
-    val = val + warp_shfl_xor(val, offset, lane)
-  return val
-
-def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
-  # inputs are (B*H, N, D)
-  BH, N, D = q.shape
-  assert N % BLOCK_M == 0 and N % BLOCK_N == 0, f"N={N} must be divisible by BLOCK_M={BLOCK_M} and BLOCK_N={BLOCK_N}"
-  assert D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0, f"D={D} must be divisible by WMMA_K={WMMA_K} and LANES_PER_WAVE_N={LANES_PER_WAVE_N}"
-  assert BLOCK_M % (WAVES_M * WMMA_M) == 0 and BLOCK_N % LANES_PER_WAVE_N == 0
-  TM = BLOCK_M // (WAVES_M * LANES_PER_WAVE_M)
-  TN = BLOCK_N // (WAVES_N * LANES_PER_WAVE_N)
-  TD = D // (WAVES_N * LANES_PER_WAVE_N)
-  SCALE = 1.0 / math.sqrt(D)
-
-  block_bh = UOp.range(BH, 0, AxisType.GLOBAL)
-  block_m = UOp.range(N // BLOCK_M, 1, AxisType.GLOBAL)
-
-  q = q.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
-  k = k.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
-  v = v.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
-  o = o.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
-
-  wave_m = UOp.range(WAVES_M, 2, AxisType.LOCAL)
-  wave_n = UOp.range(WAVES_N, 3, AxisType.LOCAL)
-  lane = UOp.range(WARP_SIZE, -1, AxisType.WARP)
-  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
-  lane_m = lane // LANES_PER_WAVE_N
-  lane_n = lane % LANES_PER_WAVE_N
-
-  # LDS allocation: slot 0 = Q then P (shared), slot 1 = K then V
-  # TODO: the memory planner should be able to find this reuse
-  ELEMS_PER_THREAD = BLOCK_M * 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]
-
-  # register state
-  acc = UOp.placeholder((TM, TD), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  m_i = UOp.placeholder((TM,), dtypes.float, slot=3, addrspace=AddrSpace.REG)
-  l_i = UOp.placeholder((TM,), dtypes.float, slot=4, addrspace=AddrSpace.REG)
-  acc = acc.after(acc.store(acc.const_like(0)))
-  m_i = m_i.after(m_i.store(m_i.const_like(-math.inf)))
-  l_i = l_i.after(l_i.store(l_i.const_like(0)))
-
-  # ====== KV tile loop ======
-  n_tile = UOp.range(N // BLOCK_N, 100, AxisType.REDUCE)
-
-  # load Q + K into LDS (Q reloaded each iteration since P overwrites slot 0)
-  Q_lds = QP_lds[:, :D]
-  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    q.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  K_store = KV_lds.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    k[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
-  Q_lds = Q_lds.after(qk_load_barrier)
-  KV_lds_k = KV_lds.after(qk_load_barrier)
-
-  # -- S = Q @ K^T via WMMA (re-init each n_tile) --
-  S_reg = UOp.placeholder((TM, TN), dtypes.float, slot=6, addrspace=AddrSpace.REG)
-  S_reg = S_reg.after(S_reg.after(n_tile).store(S_reg.const_like(0)))
-  k_qk = UOp.range(D // WMMA_K, 101, AxisType.REDUCE)
-  tm1 = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
-  tn1 = UOp.range(TN, 201, AxisType.LOOP)
-  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(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
-  qk = UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)
-  qk_done = S_frag.store(qk).end(tm1, tn1).end(k_qk)
-  S_reg = S_reg.after(qk_done)
-
-  # -- softmax in registers with warp shuffles --
-  S_reg = S_reg.after(S_reg.store(S_reg * SCALE))
-
-  # per-thread local row max over TN=4 elements, then warp reduce across 16 lanes
-  m_ij = UOp.placeholder((TM,), dtypes.float, slot=7, addrspace=AddrSpace.REG)
-  m_ij = m_ij.after(m_ij.after(n_tile).store(m_ij.const_like(-math.inf)))
-  rm2 = UOp.range(TN, 261, AxisType.REDUCE)
-  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
-  # warp reduce max (in-place)
-  ri_w = UOp.range(TM, 270, AxisType.LOOP)
-  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce_max(m_ij[ri_w], lane)).end(ri_w))
-
-  # compute P = exp(S - m_ij) in S_reg
-  S_reg = S_reg.after(S_reg.store(((S_reg - m_ij.reshape(TM, 1).expand(TM, TN)) * LOG2E).exp2()))
-
-  p_local = UOp.placeholder((TM,), dtypes.float, slot=8, addrspace=AddrSpace.REG)
-  p_local = p_local.after(p_local.after(n_tile).store(p_local.const_like(0)))
-  rp2 = UOp.range(TN, 291, AxisType.REDUCE)
-  p_local = p_local.after(p_local.store(p_local.after(rp2) + S_reg[:, rp2]).end(rp2))
-  ri_ws = UOp.range(TM, 295, AxisType.LOOP)
-  p_sum = p_local.after(p_local[ri_ws].store(warp_reduce_sum(p_local[ri_ws], lane)).end(ri_ws))
-
-  # write P = exp(S - m_ij) to P_lds (reuses slot 0, Q no longer needed)
-  P_lds = QP_lds[:, :BLOCK_N]
-  P_write = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TN, LANES_PER_WAVE_N)
-  P_write = P_write.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TN)
-  P_store = P_write[tid].store(S_reg.cast(dtypes.half))
-
-  # -- online softmax correction --
-  ri4 = UOp.range(TM, 330, AxisType.LOOP)
-  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, AxisType.LOOP)
-  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),
-  ).end(ri4)
-  acc = acc.after(correction)
-  l_i = l_i.after(correction)
-  m_i = m_i.after(correction)
-
-  # load V into KV_lds (must wait for QK WMMA to finish reading K from KV_lds)
-  V_store = KV_lds.after(qk_done).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    v[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
-  P_lds = P_lds.after(pv_barrier)
-  KV_lds_v = KV_lds.after(pv_barrier)
-
-  # -- acc += P @ V via WMMA --
-  k_pv = UOp.range(BLOCK_N // WMMA_K, 400, AxisType.REDUCE)
-  tm2 = UOp.range(TM // WMMA_ACC, 401, AxisType.LOOP)
-  tn2 = UOp.range(TD, 402, AxisType.LOOP)
-  acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
-  p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
-  v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
-  pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), *WMMA_ARG)
-
-  # end KV tile loop
-  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
-  acc = acc.after(n_tile_end)
-  l_i = l_i.after(n_tile_end)
-  m_i = m_i.after(n_tile_end)
-
-  # normalize: acc /= l_i
-  acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD)))
-
-  # store output
-  o = o.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TD, LANES_PER_WAVE_N)
-  o = o.permute((0, 4, 3, 6, 1, 2, 5)).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=()))
+from tinygrad import Tensor, getenv
+from tinygrad.dtype import dtypes
+from tinygrad.helpers import GlobalCounters, Context
+from tinygrad.llm.kernels.amd import amd_flash_attention, amd_flash_attention_causal
 
 if __name__ == "__main__":
   B, H, N, D = getenv("B", 1), getenv("H", 32), getenv("N", 1024), getenv("D", 64)
-  q = Tensor.rand(B, H, N, D).cast(dtypes.half)
+  M, causal = getenv("M", N), getenv("CAUSAL", 0)
+  q = Tensor.rand(B, H, M, D).cast(dtypes.half)
   k = Tensor.rand(B, H, N, D).cast(dtypes.half)
   v = Tensor.rand(B, H, N, D).cast(dtypes.half)
-  o = Tensor.empty(B, H, N, D, dtype=dtypes.float)
+  o = Tensor.empty(B, H, M, D, dtype=dtypes.float)
   with Context(DEBUG=0): Tensor.realize(q, k, v)
 
-  q_flat, k_flat, v_flat, o_flat = q.reshape(B*H, N, D), k.reshape(B*H, N, D), v.reshape(B*H, N, D), o.reshape(B*H, N, D)
-  NUM_RUNS = getenv("CNT", 5)
+  q_flat, k_flat, v_flat, o_flat = q.reshape(B*H, M, D), k.reshape(B*H, N, D), v.reshape(B*H, N, D), o.reshape(B*H, M, D)
   ets = []
   with Context(DEBUG=2):
-    for _ in range(NUM_RUNS):
+    for _ in range(getenv("CNT", 5)):
       GlobalCounters.reset()
-      tst = Tensor.custom_kernel(o_flat, q_flat, k_flat, v_flat, fxn=amd_flash_attention)[0].realize()
+      tst = Tensor.custom_kernel(o_flat, q_flat, k_flat, v_flat,
+                                 fxn=amd_flash_attention_causal if causal else amd_flash_attention)[0].realize()
       ets.append(GlobalCounters.time_sum_s)
   print(f"best time: {min(ets)*1e3:.2f}ms")
 
   if getenv("VERIFY", 1):
     with Context(DEBUG=0):
-      ref = q.float().scaled_dot_product_attention(k.float(), v.float()).reshape(B*H, N, D).realize()
-      err = (ref - tst).square().mean().item()
-    print(f"mean squared error {err}")
-    if err > 1e-2:
-      raise RuntimeError("flash attention is wrong!")
-    else:
-      print("flash attention is correct!")
+      mask = Tensor.full((1, 1, M, N), float("-inf"), buffer=False).triu(N-M+1) if causal else None
+      ref = q.float().scaled_dot_product_attention(k.float(), v.float(), attn_mask=mask).reshape(B*H, M, D).realize()
+      diff = (ref - tst).abs()
+      err, max_err = diff.square().mean().item(), diff.max().item()
+    print(f"mean squared error {err}, max error {max_err}")
+    if err > 1e-2: raise RuntimeError("flash attention is wrong!")
+    print("flash attention is correct!")

extra/gemm/amd_uop_matmul.py

@@ -28,10 +28,10 @@ REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
 assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
 assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
 
-def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
+def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.WEAK): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
 def copy(dest:UOp, src:UOp, rng:int, upcast=False):
   assert dest.shape == src.shape
-  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
+  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.WEAK)
   return dest[*rngs].store(src[*rngs]).end(*rngs)
 
 def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
@@ -66,9 +66,8 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
   B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
 
-  # TODO: can we automate barrier?
-  barrier = UOp.barrier(A_local_store, B_local_store)
-  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  A_local, B_local = A_local.after(A_local_store, B_local_store), B_local.after(A_local_store, B_local_store)
 
   # open inner k range
   k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -102,7 +101,7 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)
 
   # Close k, sync, and close K tiles
-  sink = sink.end(k).barrier().end(k_tile_range)
+  sink = sink.end(k).end(k_tile_range)
 
   # ---------------------------
   # REG -> GLOBAL (epilogue)

extra/gemm/cdna_asm_gemm.py

@@ -72,7 +72,7 @@ def hk_fp8_atb_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, g_amax:Tensor|
     elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=dtypes.bfloat16, device=a.device), 2
     elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=dtypes.bfloat16, device=a.device), 1
     else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device), 0, True
-    out = Tensor(inv.uop.multi(out_axis), device=a.device)
+    out = Tensor(inv.uop.unshard(out_axis), device=a.device)
     dname = a.device[0]
   else:
     out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device)
@@ -171,13 +171,13 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   M, K = A.shape[0]*A.shape[1], A.shape[2]
   K2, N = B.shape[(1 if B.ndim == 3 else 0):]
   assert K == K2
-  m = UOp.range(M, 1, AxisType.LOOP)
-  n = UOp.range(N, 2, AxisType.LOOP)
+  m = UOp.range(M, 1)
+  n = UOp.range(N, 2)
   k = UOp.range(K, 0, AxisType.REDUCE)
-  mul = (A.flatten().index((m*UOp.const(dtypes.index, K)+k))*
-         B.flatten().index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
+  mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
+         B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
   red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
-  store = C.flatten().index((m*UOp.const(dtypes.index, N)+n)).store(red).end(m, n)
+  store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n)).store(red).end(m, n)
   return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
 
 # ** bf16 A @ B.T kernel in C
@@ -234,7 +234,7 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
     elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=a.dtype, device=a.device), 2
     elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=a.dtype, device=a.device), 1
     else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device), 0, True
-    out = Tensor(inv.uop.multi(out_axis), device=a.device)
+    out = Tensor(inv.uop.unshard(out_axis), device=a.device)
     dname = a.device[0]
   else:
     out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device)
@@ -366,11 +366,11 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
 
   if is_multi:
     if n_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.multi(2), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.unshard(2), device=a.device)
     elif m_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.multi(1), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(1), device=a.device)
     else:
-      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.multi(0),
+      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(0),
                    device=a.device)
   else:
     out = Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device)

extra/gemm/mi350x_uop_matmul_2.py

@@ -29,7 +29,7 @@ TID_SIZE = WARPGROUP_SIZE*WARP_SIZE
 
 def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
   assert dest.shape == src.shape
-  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.LOOP) for i,s in enumerate(src.shape)]
+  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.WEAK) for i,s in enumerate(src.shape)]
   copy = dest[*rngs].store(src[*rngs]).end(*rngs)
   return dest.after(copy) if set else copy
 

extra/gemm/moe_gemm.py

@@ -0,0 +1,111 @@
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+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
+
+@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
+  E, N, K2 = B.shape
+  assert K == K2, f"{A.shape} {B.shape}"
+  assert E == n_experts, f"{E} != {n_experts}"
+  threads = UOp.special(64 * 8, "lidx0")
+  workgroups = UOp.special((M // 256) * (N // 256), "gidx0")
+  sink_inputs = (C.base, A.base, B.base, scale_A.base, scale_B.base, extra[0].base, extra[1].base, extra[2].base, threads, workgroups)
+  sink = UOp.sink(*sink_inputs,
+                  arg=KernelInfo(f"hk_grouped_mxfp8_gemm_{E}_{M}_{N}_{K}",
+                                 estimates=Estimates(ops=2*M*N*K, mem=(M*K+E*N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
+  kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
+  src = (kittens_path/"grouped_mxfp8_gemm.cpp").read_text()
+  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
+                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
+                                 f"-DGEMM_E={E}"]).compile_cached(src)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+                               UOp(Ops.BINARY, arg=lib)))
+
+@functools.cache
+def custom_hk_grouped_mxfp8_wgrad(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, expert_off:UOp, *, dname:str, n_experts:int) -> UOp:
+  N, M = A.shape
+  K, M2 = B.shape
+  assert M == M2, f"{A.shape} {B.shape}"
+  E = n_experts
+  threads = UOp.special(64 * 8, "lidx0")
+  workgroups = UOp.special(E * (N // 256) * (K // 256), "gidx0")
+  sink = UOp.sink(C.base, A.base, B.base, scale_A.base, scale_B.base, expert_off.base, threads, workgroups,
+                  arg=KernelInfo(f"hk_grouped_mxfp8_wgrad_{E}_{M}_{N}_{K}",
+                                 estimates=Estimates(ops=2*M*N*K, mem=(N*M+K*M)*A.dtype.itemsize+E*N*K*C.dtype.itemsize)))
+  kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
+  src = (kittens_path/"grouped_mxfp8_wgrad.cpp").read_text()
+  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
+                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DWGRAD_M={M}", f"-DWGRAD_N={N}", f"-DWGRAD_K={K}",
+                                 f"-DWGRAD_E={E}"]).compile_cached(src)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+                               UOp(Ops.BINARY, arg=lib)))
+
+def grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> Tensor:
+  from extra.llama_kernels.transpose_quantize_mxfp8 import transpose_quantize_mxfp8
+  M, N = g.shape
+  M2, K = xg.shape
+  assert M == M2, f"{g.shape} {xg.shape}"
+  assert M % 128 == 0 and N % 256 == 0 and K % 256 == 0, f"wgrad needs M%128,N%256,K%256, got {g.shape} {xg.shape}"
+  gT, _, g_si = transpose_quantize_mxfp8(g.contiguous())
+  xT, _, x_si = transpose_quantize_mxfp8(xg.contiguous())
+  dname = (g.device[0] if isinstance(g.device, tuple) else g.device).split(":")[0]
+  is_multi = isinstance(g.device, tuple)
+  inv = Tensor.invalids(1, n_experts * N, K, dtype=dtypes.bfloat16, device=g.device)
+  out = Tensor(inv.uop.unshard(0), device=g.device) if is_multi else inv
+  out = Tensor.custom_kernel(out, gT, xT, g_si, x_si, expert_off,
+                             fxn=functools.partial(custom_hk_grouped_mxfp8_wgrad, dname=dname, n_experts=n_experts))[0]
+  out = out.sum(0) if is_multi else out.squeeze(0)
+  return out.reshape(n_experts, N, K)
+
+def mx_pack_3d(e8:Tensor) -> Tensor:
+  E, rows, scale_K = e8.shape
+  return e8.reshape(E, rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(E, rows, scale_K // 4).permute(0, 2, 1).contiguous()
+
+@functools.cache
+def custom_grouped_mx_gemm_bw(gradient:UOp, kernel:UOp, w_stored:bool=False) -> tuple:
+  inputs = kernel.src[1:]
+  aq = Tensor(inputs[1], device=inputs[1].device)
+  bq = Tensor(inputs[2], device=inputs[2].device)
+  ae8 = Tensor(inputs[5], device=inputs[5].device)
+  be8 = Tensor(inputs[6], device=inputs[6].device)
+  E, N = bq.shape[0], bq.shape[1]
+  M, K = aq.shape
+  g = Tensor(gradient, device=aq.device).reshape(M, N).cast(dtypes.bfloat16)
+  x_phys = (aq.cast(dtypes.bfloat16) * _mx_block_scale(ae8).cast(dtypes.bfloat16))
+  w_phys = (bq.cast(dtypes.bfloat16) * _mx_block_scale_3d(be8).cast(dtypes.bfloat16))
+  expert_off = Tensor(inputs[7], device=inputs[7].device)
+  grad_x = grouped_mx_gemm(g, w_phys.transpose(1, 2), expert_off)
+  grad_w = grouped_mx_wgrad(g, x_phys, expert_off, E)
+  grad_xq = grad_x * _mx_block_scale(ae8).cast(dtypes.bfloat16)
+  grad_wq = grad_w.contiguous() if w_stored else (grad_w * _mx_block_scale_3d(be8).cast(dtypes.bfloat16)).contiguous()
+  return (None, grad_xq.uop, grad_wq.uop) + tuple(None for _ in inputs[3:])
+
+_grouped_bw_stored = functools.partial(custom_grouped_mx_gemm_bw, w_stored=True)
+
+def grouped_mx_gemm(x:Tensor, w:Tensor|tuple[Tensor, Tensor], expert_off:Tensor) -> Tensor:
+  if (pre_quantized := isinstance(w, tuple)):
+    w_q, w_e8 = w
+    E, N, K2 = w_q.shape
+  else:
+    E, N, K2 = w.shape
+  M, K = x.shape
+  assert K == K2, f"shape mismatch {x.shape} {w.shape}"
+  assert M % 256 == 0 and N % 256 == 0 and K % 128 == 0, f"grouped mxfp8 needs M%256,N%256,K%128, got {x.shape} {w.shape}"
+  dname = (x.device[0] if isinstance(x.device, tuple) else x.device).split(":")[0]
+  x_q, x_e8, x_si = quantize_mxfp8(x)
+  if not pre_quantized: w_q, w_e8, _ = quantize_mxfp8(w)
+  w_si = mx_pack_3d(w_e8)
+  xe_in, out_shape = x_e8.reshape(M, K // 32), (M, N)
+  if isinstance(x.device, tuple) and (row_axis := x.uop.axis) is not None:
+    ndev = len(x.device)
+    out = Tensor(Tensor.invalids(*(s // ndev if i == row_axis else s for i, s in enumerate(out_shape)),
+                                 dtype=dtypes.bfloat16, device=x.device).uop.unshard(row_axis), device=x.device)
+  else:
+    out = Tensor.invalids(*out_shape, dtype=dtypes.bfloat16, device=x.device)
+  return Tensor.custom_kernel(out, x_q, w_q, x_si, w_si, xe_in, w_e8, expert_off,
+                              fxn=functools.partial(custom_hk_grouped_mxfp8_gemm, dname=dname, n_experts=E),
+                              grad_fxn=(_grouped_bw_stored if pre_quantized else custom_grouped_mx_gemm_bw))[0]

extra/gemm_fragment.py

@@ -0,0 +1,195 @@
+"""
+tilelang-style matmul_relu written with tinygrad UOp APIs.
+
+Demonstrates that tilelang's T.alloc_fragment is expressible with existing
+tinygrad primitives: a per-thread REG buffer, wrapped in one Ops.UNSHARD per
+sharded axis over the LOCAL thread-grid ranges to form the full logical tile.
+Here the 64 threads are an 8x8 grid and each thread owns an 8x8 sub-tile --
+the 2-D fragment layout tilelang infers. The kernel is written against the
+full-tile UNSHARD view, and multi_pm (the same pass that lowers multi-device
+UNSHARDs) resolves it into per-thread shard code.
+
+Reference tilelang kernel:
+
+    @tilelang.jit
+    def matmul_relu(A, B, block_M=64, block_N=64, block_K=64,
+                    dtype=T.float16, accum_dtype=T.float32):
+        M, N, K = T.const('M, N, K')
+        C = T.empty([M, N], dtype)
+        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
+            A_shared = T.alloc_shared((block_M, block_K), dtype)
+            B_shared = T.alloc_shared((block_K, block_N), dtype)
+            C_local  = T.alloc_fragment((block_M, block_N), accum_dtype)
+            T.clear(C_local)
+            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
+                T.copy(A[by * block_M, ko * block_K], A_shared)
+                T.copy(B[ko * block_K, bx * block_N], B_shared)
+                T.gemm(A_shared, B_shared, C_local)
+            for i, j in T.Parallel(block_M, block_N):
+                C_local[i, j] = T.max(C_local[i, j], 0)
+            T.copy(C_local, C[by * block_M, bx * block_N])
+        return C
+
+API mapping (tilelang -> tinygrad UOps, idioms from test/backend/test_custom_kernel.py):
+
+    T.Kernel(gx, gy, threads=T)    -> AxisType.GLOBAL ranges (blocks) + AxisType.LOCAL ranges (thread grid)
+    T.alloc_shared(shape, dtype)   -> UOp.placeholder(shape, dtype, slot, AddrSpace.LOCAL)
+    T.alloc_fragment(shape, dt)    -> per-thread REG placeholder, wrapped in one Ops.UNSHARD per sharded axis over
+                                      the AxisType.LOCAL ranges: fragment.unshard((axis_y, axis_x), (ty, tx)).
+                                      The full logical tile is the shard with each sharded axis multiplied by its
+                                      range size, exactly like device sharding, but the sharding axes are thread
+                                      axes carried by the RANGE metadata instead of a device tuple. C_local[i, j]
+                                      with [i, j] in this thread's shard is INDEX on the UNSHARD, which multi_pm
+                                      resolves into INDEX on the per-thread REG shard, axis by axis.
+    T.copy(gmem_slice, smem)       -> smem[thread_idx].set(gmem_slice[thread_idx], end=copy_rng). set returns the
+                                      smem tile AFTER the copy; the implicit-barrier pass turns the store->load
+                                      dependency of the loop that consumes it into a workgroup barrier
+    T.gemm (no WMMA)               -> C_local[..].set(C_local.after(k)[..] + a_shared[..] * b_shared[..], end=k)
+                                      with k a loop-carried LOOP range (codegen builds the register accumulator
+                                      from this self-referential store automatically)
+    T.copy(fragment, gmem)         -> gmem.index(gidx).store(C_local[..]).end(all_ranges)
+    UNSHARD lowering               -> multi_pm in codegen (full_rewrite_to_sink): INDEX/AFTER/STORE ops on the
+                                      full-tile view become per-thread shard ops, no UNSHARD survives into the program.
+"""
+
+from tinygrad.dtype import dtypes, AddrSpace, DType
+from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
+from tinygrad.helpers import cdiv, getenv
+from tinygrad.tensor import Tensor
+
+# ---------------------------------------------------------------------------
+# tilelang builtins, expressed with tinygrad UOp APIs
+# ---------------------------------------------------------------------------
+
+def alloc_shared(shape:tuple[int, ...], dtype:DType) -> UOp:
+  """T.alloc_shared: one LOCAL buffer shared by all threads in the block."""
+  return UOp.placeholder(tuple(shape), dtype, next(UOp.unique_num), AddrSpace.LOCAL)
+
+def alloc_fragment(shape:tuple[int, ...], dtype:DType, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp:
+  """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid.
+
+  Each thread privately owns shape[axis]//threads elements along every sharded
+  axis in a REG buffer. The UNSHARDs over the LOCAL thread ranges present the
+  full logical tile: full_shape = shard_shape with each sharded axis multiplied
+  by its range size. This is exactly how UNSHARD carries a DEVICE axis today,
+  except the sharding axes are thread axes carried by the RANGE metadata.
+  """
+  assert len(axes) == len(rngs)
+  assert all(tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL for tnum in rngs), "fragments shard over LOCAL ranges"
+  assert all(shape[a] % (int(rng.vmax)+1) == 0 for a, rng in zip(axes, rngs))
+  by_axis = dict(zip(axes, rngs))
+  shard_shape = tuple(s // (int(by_axis[i].vmax)+1) if i in by_axis else s for i, s in enumerate(shape))
+  fragment = UOp.placeholder(shard_shape, dtype, next(UOp.unique_num), AddrSpace.REG)
+  return fragment.unshard(axes, rngs)
+
+# ---------------------------------------------------------------------------
+# GEMM kernel: C = relu(A @ B), float inputs (fp16 or fp32), fp32 fragment accumulator, no WMMA
+# ---------------------------------------------------------------------------
+
+# 64x64 output tile per block, 128 threads as an 8x16 grid; each thread owns an 8x4 fragment sub-tile
+# (the 2-D per-thread layout tilelang infers for this GEMM). The 4 contiguous columns (TN=4) are what
+# let codegen vectorize loads/stores to float4, matching tilelang's lowering exactly.
+BLOCK_M = BLOCK_N = BLOCK_K = 64
+TY = 8
+TX = 16
+THREADS = TY * TX
+TM = BLOCK_M // TY   # fragment rows per thread  (8)
+TN = BLOCK_N // TX   # fragment columns per thread (4)
+
+def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp:
+  """C[M, N] = relu(A[M, K] @ B[K, N]) -- one 64x64 tile per block, locals + a 2-D fragment."""
+  M, K = a.shape
+  K2, N = b.shape
+  assert K == K2 and a.dtype == b.dtype == c.dtype and not dtypes.is_int(a.dtype)
+  assert not (K % BLOCK_K or M % BLOCK_M or N % BLOCK_N), "test sizes must be multiples of the block sizes"
+
+  # with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=128) as (bx, by):
+  bx = UOp.range(cdiv(N, BLOCK_N), 0, AxisType.GLOBAL)
+  by = UOp.range(cdiv(M, BLOCK_M), 1, AxisType.GLOBAL)
+  # tx (N, 16) is the fast/inner LOCAL axis so a warp covers 16 cols x 2 rows --
+  # matching tilelang's (tidx>>4, tidx&15) warp composition. This keeps the 8 A_shared
+  # reads in a warp on only 2 row-groups (broadcast across 16 cols) instead of 8 rows
+  # (8-way bank conflict), since A_shared[row*512 + ...] all map to the same bank when 8
+  # distinct rows land in one warp.
+  tx = UOp.range(TX, 2, AxisType.LOCAL)
+  ty = UOp.range(TY, 3, AxisType.LOCAL)
+
+  # A_shared = T.alloc_shared((BLOCK_M, BLOCK_K), dtype)
+  # B_shared = T.alloc_shared((BLOCK_K, BLOCK_N), dtype)
+  A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype)
+  B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype)
+
+  # C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), accum_dtype) -- an 8x4 REG tile per thread of the 8x16 grid
+  C_local = alloc_fragment((BLOCK_M, BLOCK_N), dtypes.float32, (0, 1), (ty, tx))
+
+  # T.clear(C_local) -- each thread zeroes its own fragment sub-tile
+  ic, jc = UOp.range(TM, 4, AxisType.LOOP), UOp.range(TN, 5, AxisType.UPCAST)
+  C_loc = C_local[ic*TM + ty, tx*TN + jc].set(0.0, end=(ic, jc))
+
+  # for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
+  # (num_stages pipelining is async copy + multi-buffering; this is the synchronous single-buffer version)
+  ko = UOp.range(cdiv(K, BLOCK_K), 6, AxisType.LOOP)
+
+  # T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared) -- each thread copies its own 8x4 sub-tile.
+  # Row index is iar*TM + ty (strided by TM across ty), matching tilelang's layout: thread ty owns
+  # rows {ty, ty+8, ..., ty+56} not {ty*8, ..., ty*8+7}.
+  iar, ka = UOp.range(TM, 7, AxisType.LOOP), UOp.range(TN, 8, AxisType.UPCAST)
+  A_store = A_shared[iar*TM + ty, tx*TN + ka].store(a[by*BLOCK_M + iar*TM + ty, ko*BLOCK_K + tx*TN + ka]).end(iar, ka)
+
+  # T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared)
+  kb, ibr = UOp.range(TM, 9, AxisType.LOOP), UOp.range(TN, 10, AxisType.UPCAST)
+  B_store = B_shared[kb*TM + ty, tx*TN + ibr].store(b[ko*BLOCK_K + kb*TM + ty, bx*BLOCK_N + tx*TN + ibr]).end(kb, ibr)
+
+  # get the shared after the stores (single barrier)
+  A_shared = A_shared.after(A_store, B_store)
+  B_shared = B_shared.after(A_store, B_store)
+
+  # T.gemm(A_shared, B_shared, C_local), no WMMA -- per-thread accumulate over its fragment sub-tile.
+  # identical to custom_gemm: a self-referential store over the loop-carried kk range,
+  # which codegen turns into a register accumulator
+  # kk is the outer compute loop (axis 11) so that for each kk we read all 8 A rows and reuse
+  # the B[kk] read across them -- matching tilelang's ko > kk > row > col access order exactly.
+  kk, ir = UOp.range(BLOCK_K, 11, AxisType.LOOP), UOp.range(TM, 12, AxisType.LOOP)
+  jj = UOp.range(TN, 13, AxisType.UPCAST)
+  acc = C_loc.after(kk)[ir*TM + ty, tx*TN + jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32)
+  # closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them
+  C_loc = C_loc[ir*TM + ty, tx*TN + jj].set(acc, end=(kk, ir, jj, ko))
+
+  # for i, j in T.Parallel(BLOCK_M, BLOCK_N): C_local[i, j] = T.max(C_local[i, j], 0)
+  # T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N]) -- per-thread store of the fragment shard (relu fused into it)
+  # LOOP: these loops are the per-thread output layout; convert_loop_to_global must not globalize them
+  ie, je = UOp.range(TM, 14, AxisType.LOOP), UOp.range(TN, 15, AxisType.UPCAST)
+  c_st = c[by*BLOCK_M + ie*TM + ty, bx*BLOCK_N + tx*TN + je].store(C_loc[ie*TM + ty, tx*TN + je].relu().cast(c.dtype))
+
+  # all open ranges are closed at the final store (ko was closed above).
+  # the fragment UNSHARDs go to codegen as is: multi_pm there resolves the full-tile view into per-thread shard code
+  return c_st.end(je, ie, tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=()))
+
+# ---------------------------------------------------------------------------
+# python wrapper: same signature as the tilelang function
+# ---------------------------------------------------------------------------
+
+def matmul_relu(a:Tensor, b:Tensor) -> Tensor:
+  """C = relu(A @ B), fp16 in/out with an fp32 fragment accumulator."""
+  c = Tensor.empty(a.shape[0], b.shape[1], dtype=a.dtype, device=a.device)
+  return c.custom_kernel(a, b, fxn=matmul_relu_kernel)[0]
+
+# ---------------------------------------------------------------------------
+# test
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+  from tinygrad import Device
+  assert Device[Device.DEFAULT].renderer.has_local, "this GPU-style kernel needs a backend with local memory (LOCAL ranges + barriers)"
+  M = K = N = getenv("N", 256)  # 4x4 grid of 64x64 tiles, 4 K chunks
+  dtype_in = dtypes.half if getenv("HALF") else dtypes.float
+
+  a = Tensor.randn(M, K, dtype=dtype_in).contiguous()
+  b = Tensor.randn(K, N, dtype=dtype_in).contiguous()
+  ref = (a @ b).relu().realize()
+
+  out = matmul_relu(a, b).realize()
+
+  import numpy as np
+  np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-1, rtol=1e-2)
+  print("matmul_relu passed!")

extra/hcq2/ops_amd2.py

@@ -3,8 +3,8 @@ from typing import cast, Any, Callable
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_getaddr, make_ins, make_cmdbuf, make_placeholder
-from extra.hcq2.hcq2 import make_binary_patch
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
+from tinygrad.runtime.support.hcq2 import make_binary_patch, make_patches
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
@@ -36,7 +36,9 @@ class PM4Ops(FastEnum):
   SET_SH_REG = auto(); SET_UCONFIG_REG = auto(); WAIT_REG_MEM = auto(); ACQUIRE_MEM = auto()  # noqa: E702
   RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto()  # noqa: E702
 
-def pkt3(ctx, op:PM4Ops, *vals): return make_ins(op, ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)
+def pkt3(ctx, op:PM4Ops, *vals):
+  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(dtypes.uint32, x)
+    for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)))
 
 def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
   if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified')
@@ -90,26 +92,26 @@ def memory_barrier(ctx):
                  reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
     acquire_mem(ctx)))
 
-def pm4_wait(ctx, x, y): return wait_reg_mem(ctx, y, mem=make_getaddr(x.buf_uop, ctx.devs))
+def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val, mem=dst.getaddr(ctx.devs))
 
 def pm4_barrier(ctx): return memory_barrier(ctx)
 
 def pm4_store(ctx, dst, val):
   if val.op is Ops.BINARY: return None
-  return release_mem(ctx, make_getaddr(dst, ctx.devs), val, ctx.pm4.data_sel__mec_release_mem__send_32_bit_low,
+  return release_mem(ctx, dst.getaddr(ctx.devs), val, ctx.pm4.data_sel__mec_release_mem__send_32_bit_low,
                      ctx.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True)
 
 def pm4_timestamp(ctx, dst):
-  return release_mem(ctx, make_getaddr(dst, ctx.devs), 0, ctx.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
+  return release_mem(ctx, dst.getaddr(ctx.devs), 0, ctx.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
                      ctx.pm4.int_sel__mec_release_mem__none)
 
 def pm4_program(ctx, call, prg):
   data, info = prg.arg
   lib_gpu = prg.src[0]
   args = encode_kernargs_clike(call, prg, ctx.devs)
-  prog_addr = make_getaddr(lib_gpu, ctx.devs) + data.entry_point_offset
-  scratch_addr = make_getaddr(make_placeholder(ctx.devs, data.private_segment_size, dtypes.uint8, "scratch", unique=False), ctx.devs)
-  args_addr = make_getaddr(args, ctx.devs)
+  prog_addr = lib_gpu.getaddr(ctx.devs) + data.entry_point_offset
+  scratch_addr = UOp.placeholder((data.private_segment_size,), dtypes.uint8, 0, device=ctx.devs).rtag("scratch").getaddr(ctx.devs)
+  args_addr = args.getaddr(ctx.devs)
 
   user_regs = []
   if data.enable_private_segment_sgpr:
@@ -138,37 +140,43 @@ def pm4_program(ctx, call, prg):
 pm_pm4_opsel = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
 
-  (UPat(Ops.WAIT, src=(UPat.var("x") >= UPat.var("y"),)), pm4_wait),
-  (UPat(Ops.BARRIER), pm4_barrier),
-  (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
-  (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
+  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
+  (UPat(Ops.INS, arg="barrier"), pm4_barrier),
+  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
+  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])
 
-def pm4_submit(cmdbuf, devs):
-  size, zero = UOp.const(dtypes.uint32, cmdbuf.nbytes() // dtypes.uint32.itemsize), UOp.const(dtypes.int, 0)
-
-  # the compute queue's ring and its host-side ring/write/put pointers (placeholders, resolved in pm_bufferize)
-  for d in devs: q = Device[d].compute_queue
-  ring, wptr, doorbell, put_ptr = (make_placeholder(devs, b.size, b.dtype, ("COMPUTE:0", name), unique=False)
+def pm4_submit(ctx, lin):
+  # ensure compute queues are allocated
+  for d in (devs:=ctx.devs): q = Device[d].compute_queue
+  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
     for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
 
-  # place the cmdbuf at the ring's write offset, wrapping the ring
-  put = put_ptr.index(zero)
-  next_put = put + size.cast(put.dtype)
-  i = UOp.range(size, 0, dtype=dtypes.int, src=(cmdbuf,))
-  ring_idx = ((put + i.cast(put.dtype)) % q.ring.size).cast(dtypes.int)
+  # two tail dwords coordinate safe IB reuse: GPU completions and host submits
+  size_dw = sum(len(ins.src) for ins in lin.src) + len(release_mem(ctx, 0, 0).src)
+  assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"
 
-  # copy the cmdbuf into the ring and advance the put/write pointers
-  copy_to_ring = ring.index(ring_idx).store(cmdbuf.index(i).load()).end(i)
-  bump_put_ptr = put_ptr.index(zero).store(next_put)
-  bump_wptr = wptr.index(zero).store(next_put)
+  ib = UOp.placeholder((size_dw + 2,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
+  done_idx, submit_idx = UOp.const(dtypes.int, size_dw + 0), UOp.const(dtypes.int, size_dw + 1)
+  submitted = (counter:=ib.after(make_patches(ib, [((size_dw + i) * 4, UOp.const(dtypes.uint32, 0)) for i in range(2)])).index(submit_idx)).load()
+  completed = ib.after(loop:=UOp.loop(0)).index(done_idx).load()
+  ib_free = completed.end(loop, completed != submitted)
 
-  # ring the doorbell once the copy and pointer bumps have landed
-  flush = UOp.barrier(copy_to_ring, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero).store(next_put)
+  bump_fence = pm4_store(ctx, UOp(Ops.SLICE, dtypes.uint32, (ib, UOp.const(dtypes.weakint, size_dw)), 2), (submitted + 1).cast(dtypes.uint64))
+  cmdbuf = make_cmdbuf(lin.replace(src=lin.src + (bump_fence,)), devs, buf=ib, dep=ib_free)
 
-pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
-  lambda lin: pm4_submit(make_cmdbuf(lin, to_tuple(lin.arg[0])), to_tuple(lin.arg[0])))])
+  # the ring itself only carries a packet pointing at the ib, wrapping the ring
+  put = put_ptr.index(zero:=UOp.const(dtypes.int, 0))
+  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(cmdbuf.getaddr(devs)), size_dw | ctx.pm4.INDIRECT_BUFFER_VALID)
+  write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(dtypes.uint32, x)) for off,x in enumerate(pkt)])
+
+  # advance the put/write pointers past the packet
+  bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
+  bump_wptr = wptr.index(zero).store(put + len(pkt))
+  flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr, counter.store(submitted + 1))
+  return doorbell.after(flush).index(zero).store(put + len(pkt))
+
+pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), pm4_submit)])
 
 # *****************
 # SDMA
@@ -176,36 +184,36 @@ pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
 class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TRAP = auto(); TIMESTAMP = auto()  # noqa: E702
 
 def sdma_copy(ctx, call):
-  dst, src = call.src[1], call.src[2]
-  sz = src.max_numel() * src.dtype.itemsize
-  src_addr, dst_addr = make_getaddr(src, ctx.devs), make_getaddr(dst, ctx.devs)
-  return UOp(Ops.LINEAR, dtypes.void, tuple([make_ins(SDMAOps.COPY,
-     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 + off), *data64_le(dst_addr + off)) for off in range(0, sz, ctx.max_copy_size)]))
+  sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
+  src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
+  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(dtypes.uint32, x) 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+off), *data64_le(dst_addr+off))))
 
-def sdma_wait(ctx, x, y):
+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) \
      | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-  return make_ins(SDMAOps.POLL_REGMEM, op, *data64_le(make_getaddr(x.buf_uop, ctx.devs)), y, 0xffffffff,
-    ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))
+  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(dtypes.uint32, x) for x in (
+    op, *data64_le(dst.getaddr(ctx.devs)), val, 0xffffffff,
+    ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))))
 
-def sdma_store(ctx, dst, val):
+def sdma_store(ctx, ins, dst, val):
   op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0)
-  return UOp(Ops.LINEAR, dtypes.void, (
-    make_ins(SDMAOps.FENCE, op, *data64_le(make_getaddr(dst, ctx.devs)), val), make_ins(SDMAOps.TRAP, ctx.sdma.SDMA_OP_TRAP, 0)))
+  return UOp(Ops.LINEAR, src=(
+    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
+    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(dtypes.uint32, x) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))
 
-def sdma_timestamp(ctx, dst):
+def sdma_timestamp(ctx, ins, dst):
   op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL)
-  return make_ins(SDMAOps.TIMESTAMP, op, *data64_le(make_getaddr(dst, ctx.devs)))
+  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
 
 pm_sdma_opsel = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
 
-  (UPat(Ops.BARRIER), lambda: UOp(Ops.NOOP, dtypes.void, ())),
-  (UPat(Ops.WAIT, src=(UPat.var("x") >= UPat.var("y"),)), sdma_wait),
-  (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), sdma_timestamp),
-  (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store),
+  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
+  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val")), name="ins"), sdma_wait),
+  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),), name="ins"), sdma_timestamp),
+  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val")), name="ins"), sdma_store),
 ])
 
 def sdma_submit(cmdbuf, devs):
@@ -214,7 +222,7 @@ def sdma_submit(cmdbuf, devs):
 
   # the sdma queue's ring and its host-side ring/write/put pointers
   for d in devs: q = Device[d].sdma_queue(0)
-  ring, wptr, doorbell, put_ptr = (make_placeholder(devs, b.size, b.dtype, ("COPY:0", name), unique=False)
+  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COPY:0_{name}")
     for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
 
   # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
@@ -240,7 +248,7 @@ def sdma_submit(cmdbuf, devs):
   return doorbell.after(flush).index(zero).store(next_put_b)
 
 pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
-  lambda lin: sdma_submit(make_cmdbuf(lin, to_tuple(lin.arg[0])), to_tuple(lin.arg[0])))])
+  lambda ctx, lin: sdma_submit(make_cmdbuf(lin, ctx.devs), ctx.devs))])
 
 @dataclass(frozen=True)
 class AMDEncodeCtx:  # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
@@ -251,7 +259,7 @@ def encode_queue(q:UOp) -> UOp|None:
   d = Device[(devs:=to_tuple(q.arg[0]))[0]]
   ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
   opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
-  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"))
+  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"), ctx)
 
 @dataclass(frozen=True)
 class AMDProgramData:
@@ -279,7 +287,7 @@ def amd_build_program(prg:UOp) -> UOp:
       wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
       kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
       enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
-    buf = make_placeholder(prg.device, len(image), dtypes.uint8, "program")
+    buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=prg.device).rtag("program")
     cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg))
   return cached
 
@@ -509,7 +517,8 @@ class PCIIface(PCIIfaceBase):
       cq = d.compute_queue
       for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
       d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      d.timeline_signal()._buf.cpu_view().mv.cast('Q')[0] = d.timeline_value().as_memoryview(force_zero_copy=True).cast('Q')[0] - 1
+      d.timeline_signal('COMPUTE:0')._buf.cpu_view().mv.cast('Q')[0] = \
+        d.timeline_value('COMPUTE:0').as_memoryview(force_zero_copy=True).cast('Q')[0] - 1
 
   def sleep(self, timeout):
     if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -628,10 +637,10 @@ class AMDDevice(HCQ2Compiled):
 
     qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
     self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.PARAM, tag={(qname, name)}), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
+      (UPat(Ops.PARAM, tag=f"{qname}_{name}"), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
     ] + [
-      (UPat(Ops.PARAM, tag={(qname, "timeline_signal")}), lambda ctx, q=qname: ctx[0].timeline_signal(q)),
-      (UPat(Ops.PARAM, tag={(qname, "timeline_value")}), lambda ctx, q=qname: ctx[0].timeline_value(q)),
+      (UPat(Ops.PARAM, tag=f"{qname}_timeline_signal"), lambda ctx, q=qname: ctx[0].timeline_signal(q)),
+      (UPat(Ops.PARAM, tag=f"{qname}_timeline_value"), lambda ctx, q=qname: ctx[0].timeline_value(q)),
     ]) + self.pm_bufferize
 
     return queue

extra/llama_kernels/__init__.py

@@ -31,12 +31,12 @@ def dname_of(device) -> str:
 
 def alloc_like(shape, dtype, device, axis=None) -> Tensor:
   if isinstance(device, tuple) and axis is not None:
-    return Tensor(Tensor.invalids(*shard_shape(shape, axis, len(device)), dtype=dtype, device=device).uop.multi(axis), device=device)
+    return Tensor(Tensor.invalids(*shard_shape(shape, axis, len(device)), dtype=dtype, device=device).uop.unshard(axis), device=device)
   return Tensor.invalids(*shape, dtype=dtype, device=device)
 
 def alloc_local(shape, dtype, device, axis=None) -> Tensor:
   if isinstance(device, tuple) and axis is not None:
-    return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.multi(0), device=device)
+    return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.unshard(0), device=device)
   return Tensor.invalids(*shape, dtype=dtype, device=device)
 
 def compile_hip(src:str, defines:list[str]):

extra/llama_kernels/fused_ce/__init__.py

@@ -16,7 +16,7 @@ def _custom_fused_ce_loss_fwd(loss_out:UOp, max_out:UOp, lse_out:UOp, logits:UOp
   row_lse = (logits[b, s, v_lse].cast(dtypes.float) - row_max).exp().reduce(v_lse, arg=Ops.ADD).log() + row_max
 
   v_smooth = UOp.range(vocab, 3, axis_type=AxisType.REDUCE)
-  target = logits[b, s, targets[row].cast(dtypes.index)].cast(dtypes.float)
+  target = logits[b, s, targets[row].cast(dtypes.weakint)].cast(dtypes.float)
   mean_logits = logits[b, s, v_smooth].cast(dtypes.float).reduce(v_smooth, arg=Ops.ADD) / vocab
   loss = row_lse - (1.0 - label_smoothing) * target - label_smoothing * mean_logits
   stores = UOp.group(loss_out[row].store(loss), max_out[row].store(row_max), lse_out[row].store(row_lse))
@@ -32,7 +32,7 @@ def _custom_fused_ce_loss_bwd(d_logits:UOp, logits:UOp, lse:UOp, targets:UOp, sc
   s = row % seq
 
   prob = (logits[b, s, v].cast(dtypes.float) - lse[row]).exp()
-  target = v.eq(targets[row].cast(dtypes.index)).where(1.0 - label_smoothing, 0.0)
+  target = v.eq(targets[row].cast(dtypes.weakint)).where(1.0 - label_smoothing, 0.0)
   smooth = label_smoothing / vocab
   grad = (prob - target - smooth) * scale[0]
 
@@ -48,7 +48,7 @@ def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
     axis = logits_u.axis
     ndev = len(device)
     local_shape = tuple(s//ndev if i == axis else s for i,s in enumerate((MBS, SEQ, VOCAB)))
-    d_logits = Tensor(Tensor.invalids(*local_shape, dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
+    d_logits = Tensor(Tensor.invalids(*local_shape, dtype=dtypes.bfloat16, device=device).uop.unshard(axis), device=device)
     rows_per_dev = local_shape[0] * local_shape[1]
     seq_per_dev = local_shape[1]
   else:
@@ -74,11 +74,11 @@ def fused_ce_loss(logits:Tensor, targets:Tensor, label_smoothing:float=0.1) -> T
     axis = logits.uop.axis
     assert axis in (0, 1), f"unsupported sharding axis={axis} for CE loss"
     ndev = len(logits.device)
-    loss_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+    loss_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
                       device=logits.device)
-    max_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+    max_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
                       device=logits.device)
-    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
                       device=logits.device)
     local_shape = tuple(s//ndev if i == axis else s for i,s in enumerate(logits.shape))
     rows_per_dev = local_shape[0] * local_shape[1]

extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -16,7 +16,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
 
   wg = UOp.range(NUM_WG, 0, AxisType.GLOBAL)
   tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
-  it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.LOOP)
+  it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.WEAK)
   lane = UOp.range(VEC, 3, AxisType.UNROLL)
 
   idx = (((it * NUM_WG + wg) * THREADS_PER_WG + tid) * VEC) + lane
@@ -36,19 +36,19 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
   lmax_val = lmax.after(lmax_store.end(it))[0]
 
   lds = UOp.placeholder((THREADS_PER_WG,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-  lds = lds.after(lds[tid].store(lmax_val).barrier())
+  lds = lds.after(lds[tid].store(lmax_val))
 
   step = THREADS_PER_WG // 2
   while step:
     active = tid < step
     other = lds[(tid + step).valid(active)].load()
-    lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)).barrier())
+    lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)))
     step //= 2
 
   device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0]
   if device in {"AMD", "NULL"}: atomic_arg = "if ({2} > {3}) __hip_atomic_fetch_max((int*){0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
   else: raise NotImplementedError(f"no atomic max for device {device}")
-  amax_idx = amax_out.reshape((1,)).index(UOp.const(dtypes.index, 0))
+  amax_idx = amax_out.reshape((1,)).index(UOp.const(dtypes.weakint, 0))
   max_val = lds[0].load()
   atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
   return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))

extra/qcom_gpu_driver/opencl_ioctl.py

@@ -4,7 +4,7 @@ from hexdump import hexdump
 from copy import deepcopy
 import pathlib, sys
 from tinygrad.helpers import to_mv, getenv
-from tinygrad.runtime.autogen import adreno
+from tinygrad.runtime.autogen import mesa
 sys.path.append(pathlib.Path(__file__).parent.parent.parent.as_posix())
 
 IOCTL = getenv("IOCTL", 0)
@@ -23,7 +23,7 @@ for child in xml.getroot():
 CAPTURED_STATE = {}
 
 REGS = {}
-for k, v in adreno.__dict__.items():
+for k, v in mesa.__dict__.items():
   if k.startswith("REG_") and isinstance(v, int) and v > 1024: REGS[v] = k
 
 from extra.qcom_gpu_driver import msm_kgsl
@@ -42,7 +42,7 @@ def get_struct(argp, stype):
 
 def format_struct(s):
   sdats = []
-  for field_name, *_ in s._real_fields_:
+  for field_name, *_ in s._fields_:
     if field_name in {"__pad", "PADDING_0"}: continue
     dat = getattr(s, field_name)
     if isinstance(dat, int): sdats.append(f"{field_name}:0x{dat:X}")
@@ -96,9 +96,9 @@ def parse_cmd_buf(dat):
         CAPTURED_STATE['LOAD_FRAGS'].append((state_block, state_type, num_unit, dst_off))
 
         if state_block == SB6_CS_SHADER:
-          from extra.disassemblers.adreno import disasm_raw
+          from tinygrad.runtime.support.compiler_mesa import disas_adreno
           if state_type == ST6_SHADER and IOCTL > 3:
-            disasm_raw(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
+            disas_adreno(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
           if state_type == ST6_CONSTANTS:
             x = get_mem(((vals[2] << 32) | vals[1]), num_unit*4)
             CAPTURED_STATE['constants'] = x[:]
@@ -142,7 +142,7 @@ def parse_cmd_buf(dat):
       vals = struct.unpack("I"*size, dat[ptr+4:ptr+4+4*size])
       if IOCTL > 0: print(f"{ptr:3X} -- typ 4: {size=:3d}, {reg_name}", hprint(vals))
       for vi,v in enumerate(vals): CAPTURED_STATE[offset+vi] = v
-      if offset == adreno.REG_A6XX_SP_CS_CONFIG:
+      if offset == mesa.REG_A6XX_SP_CS_CONFIG:
         val = vals[0]
         if IOCTL > 0:
           print(f"\tBINDLESS_TEX={(val >> 0) & 0b1}")
@@ -215,79 +215,3 @@ def install_hook(c_function, python_function):
 
 libc = ctypes.CDLL(ctypes.util.find_library("libc"))
 install_hook(libc.ioctl, ioctl)
-
-def before_launch():
-  global CAPTURED_STATE
-  CAPTURED_STATE.clear()
-def collect_last_launch_state():
-  global CAPTURED_STATE
-  return deepcopy(CAPTURED_STATE)
-def compare_launch_state(state, good_state):
-  cmp = [
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NTEX__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NSAMP__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NIBO__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_ENABLED),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_TEX),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_SAMP),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_IBO),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_UBO),
-
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_HALFREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_FULLREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_BRANCHSTACK__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_FULLREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_THREADMODE__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_EARLYPREAMBLE),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_MERGEDREGS),
-
-    (adreno.REG_A6XX_SP_CS_PVT_MEM_PARAM, adreno.A6XX_SP_CS_PVT_MEM_PARAM_MEMSIZEPERITEM__MASK),
-    (adreno.REG_A6XX_SP_CS_PVT_MEM_PARAM, adreno.A6XX_SP_CS_PVT_MEM_PARAM_HWSTACKSIZEPERTHREAD__MASK),
-
-    (adreno.REG_A6XX_SP_CS_UNKNOWN_A9B1, adreno.A6XX_SP_CS_UNKNOWN_A9B1_UNK5),
-    (adreno.REG_A6XX_SP_CS_UNKNOWN_A9B1, adreno.A6XX_SP_CS_UNKNOWN_A9B1_UNK6),
-
-    (adreno.REG_A6XX_SP_CS_BRANCH_COND, 0xffffffff),
-
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_KERNELDIM__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEX__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEY__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEZ__MASK),
-
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_1, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_2, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_3, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_4, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_5, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_6, 0xffffffff),
-
-    (adreno.REG_A6XX_HLSQ_CS_CNTL_0, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_CNTL_1, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_X, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_Y, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_Z, 0xffffffff),
-  ]
-
-  for x,m in cmp:
-    print(f"Field {REGS[x]}, mask: 0x{m:X} cmp: {state.get(x, 0) & m} vs {good_state.get(x, 0) & m}")
-    if state.get(x, 0) & m != good_state.get(x, 0) & m:
-      return False, f"Field {REGS[x]}, mask: 0x{m:X} mismatch: {state.get(x, 0) & m} vs {good_state.get(x, 0) & m}"
-
-  for n in ['descriptors', 'ibos']:
-    if n not in good_state: continue
-    mv1, mv2 = state.get(n), good_state.get(n)
-
-    if len(mv1) != len(mv2): return False, f"{n}: len mismatch {len(mv1)} != {len(mv2)}"
-    mv1 = memoryview(bytearray(mv1)).cast('I')
-    mv2 = memoryview(bytearray(mv2)).cast('I')
-    for i in range(len(mv2)):
-      if i % 8 == 5 or i % 8 == 4: continue # addresses
-      if mv1[i]!=mv2[i]: return False, f"{n}: content mismatch {i} {mv1[i]} {mv2[i]}"
-
-  for n in ['samplers']:
-    if n not in good_state: continue
-    mv1, mv2 = state.get(n), good_state.get(n)
-    if len(mv1) != len(mv2): return False, f"{n}: len mismatch {len(mv1)} != {len(mv2)}"
-    if any(mv1[i]!=mv2[i] for i in range(len(mv1))): return False, f"{n}: content mismatch"
-
-  return True, "PASS"

extra/thunder/amd/fa.py

@@ -14,7 +14,7 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
   shard_axis = ref.uop.axis if axis is None else axis
   shape = tuple(s // len(ref.device) if i == shard_axis else s for i, s in enumerate(shape))
   axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
+  return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.unshard(axis), dtype=dtype, device=ref.device)
 
 @functools.cache
 def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp,

extra/thunder/amd/grouped_mxfp8_gemm.cpp

@@ -0,0 +1,158 @@
+#include "kittens.cuh"
+
+using namespace kittens;
+
+#ifndef GEMM_M
+constexpr int GEMM_M = 8192;
+#endif
+#ifndef GEMM_N
+constexpr int GEMM_N = 8192;
+#endif
+#ifndef GEMM_K
+constexpr int GEMM_K = 8192;
+#endif
+#ifndef GEMM_E
+constexpr int GEMM_E = 8;
+#endif
+
+// Kernel
+constexpr int NUM_WARPS  = 8;
+constexpr int WARPS_ROW  = 2;
+constexpr int WARPS_COL  = 4;
+constexpr int BLOCK_ROW  = 256;
+constexpr int BLOCK_COL  = 256;
+constexpr int BLOCK_K    = 128;
+constexpr int HALF_ROW   = BLOCK_ROW / 2;
+constexpr int HALF_COL   = BLOCK_COL / 2;
+constexpr int REG_M      = BLOCK_ROW / WARPS_ROW / 2;
+constexpr int REG_N      = BLOCK_COL / WARPS_COL / 2;
+
+using G = kittens::group<NUM_WARPS>;
+
+__global__ __launch_bounds__(512, 2) void grouped_mxfp8_gemm_kernel(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr, fp8e8m0 *scale_A_ptr, fp8e8m0 *scale_B_ptr,
+    const uint8_t *__restrict__ a_e8_unused,
+    const uint8_t *__restrict__ b_e8_unused,
+    const int *__restrict__ expert_off) {
+    constexpr int M = GEMM_M, N = GEMM_N, K = GEMM_K, E = GEMM_E;
+
+    kittens::gl<fp8e4m3, 1, 1, M, K>     A{A_ptr, nullptr, nullptr, nullptr, nullptr};
+    kittens::gl<fp8e4m3, 1, 1, E * N, K> B{B_ptr, nullptr, nullptr, nullptr, nullptr};  // all experts stacked on rows
+    kittens::gl<bf16, 1, 1, M, N>        C{C_ptr, nullptr, nullptr, nullptr, nullptr};
+
+    constexpr int k_iters      = K / BLOCK_K;
+    constexpr int NUM_THREADS  = NUM_WARPS * WARP_THREADS;
+
+    kittens::gl<fp8e8m0, k_iters * (M / BLOCK_ROW), 1, 16, 64>     scale_A_gl{scale_A_ptr, nullptr, nullptr, nullptr, nullptr};
+    kittens::gl<fp8e8m0, E * k_iters * (N / BLOCK_COL), 1, 16, 64> scale_B_gl{scale_B_ptr, nullptr, nullptr, nullptr, nullptr};
+
+    using ST_A     = st_fp8e4m3<HALF_ROW, BLOCK_K, st_16x128_s>;
+    using ST_B     = st_fp8e4m3<HALF_COL, BLOCK_K, st_16x128_s>;
+    using ST_Scale = st<fp8e8m0, 16, 64, st_16x64_s>;
+    using RT_A     = rt_fp8e4m3<REG_M, BLOCK_K>;
+    using RT_B     = rt_fp8e4m3<REG_N, BLOCK_K>;
+    using RT_C     = rt_fl<REG_M, REG_N, col_l, rt_16x16_s>;
+
+    __shared__ ST_A As[2][2];
+    __shared__ ST_B Bs[2][2];
+    __shared__ ST_Scale scale_A_smem[2], scale_B_smem[2];
+
+    RT_A a;
+    RT_B b0, b1;
+    RT_C cA, cB, cC, cD;
+    zero(cA); zero(cB); zero(cC); zero(cD);
+
+    constexpr int tiles_M  = M / BLOCK_ROW;
+    constexpr int tiles_N  = N / BLOCK_COL;
+    const int NUM_XCDS     = 8;
+    const int WGM          = 8;
+    int wgid = chiplet_transform_chunked(blockIdx.x, gridDim.x, NUM_XCDS, WGM * WGM);
+    int num_wgid_in_group = WGM * tiles_N;
+    int group_id     = wgid / num_wgid_in_group;
+    int first_pid_m  = group_id * WGM;
+    int group_size_m = min(tiles_M - first_pid_m, WGM);
+    int block_row    = first_pid_m + ((wgid % num_wgid_in_group) % group_size_m);
+    int block_col    = (wgid % num_wgid_in_group) / group_size_m;
+    int block_m      = block_row * BLOCK_ROW;
+    int block_n      = block_col * BLOCK_COL;
+
+    int e = 0;
+    #pragma unroll
+    for (int i = 1; i < E; i++) e += (expert_off[i] <= block_row * BLOCK_ROW);
+    e = __builtin_amdgcn_readfirstlane(e);
+    const int bcol_base = e * (N / HALF_COL);          // expert base in B row-tile (128-row) units
+    const int sb_base   = e * (k_iters * tiles_N);     // expert base into scale_B batches
+
+    int warp_m = warpid() / WARPS_COL;
+    int warp_n = warpid() % WARPS_COL;
+
+    using T = fp8e4m3;
+    constexpr int bpt      = ST_A::underlying_subtile_bytes_per_thread;
+    constexpr int bpm      = bpt * NUM_THREADS;
+    constexpr int copies_A = HALF_ROW * BLOCK_K * sizeof(T) / bpm;
+    constexpr int copies_B = HALF_COL * BLOCK_K * sizeof(T) / bpm;
+    uint32_t sw_A[copies_A], sw_B[copies_B];
+    G::prefill_swizzled_offsets(As[0][0], A, sw_A);
+    G::prefill_swizzled_offsets(Bs[0][0], B, sw_B);
+
+    const T *a_base = (const T *)&A[{0, 0, 0, 0}];
+    const T *b_base = (const T *)&B[{0, 0, 0, 0}];
+    const int a_row_stride = A.template stride<2>() * sizeof(T);
+    const int b_row_stride = B.template stride<2>() * sizeof(T);
+    i32x4 a_srd = make_srsrc(a_base, (uint32_t)((uint64_t)M * a_row_stride), a_row_stride);
+    i32x4 b_srd = make_srsrc(b_base, (uint32_t)((uint64_t)E * N * b_row_stride), b_row_stride);
+
+    const int wid = warpid() % NUM_WARPS;
+    constexpr int elem_per_warp = (16 / sizeof(T)) * kittens::WARP_THREADS;
+    uint32_t a_lds_00 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[0][0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t a_lds_01 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[0][1].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t a_lds_10 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[1][0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t a_lds_11 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[1][1].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_00 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[0][0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_01 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[0][1].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_10 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[1][0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_11 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[1][1].data[0]) + wid * elem_per_warp * sizeof(T)));
+
+    int a_row_h0 = warp_m * REG_M;
+    int a_row_h1 = HALF_ROW + warp_m * REG_M;
+    int b_row_h0 = warp_n * REG_N;
+    int b_row_h1 = HALF_COL + warp_n * REG_N;
+
+
+    uint32_t a_lds[2][2] = {{a_lds_00, a_lds_01}, {a_lds_10, a_lds_11}};
+    uint32_t b_lds[2][2] = {{b_lds_00, b_lds_01}, {b_lds_10, b_lds_11}};
+
+    #pragma unroll 1
+    for (int kk = 0; kk < k_iters; kk++) {
+        G::load(As[0][0], A, {0, 0, block_row * 2,     kk}, sw_A, a_srd, a_base, __builtin_amdgcn_readfirstlane(a_lds[0][0]));
+        G::load(As[0][1], A, {0, 0, block_row * 2 + 1, kk}, sw_A, a_srd, a_base, __builtin_amdgcn_readfirstlane(a_lds[0][1]));
+        G::load(Bs[0][0], B, {0, 0, bcol_base + block_col * 2,     kk}, sw_B, b_srd, b_base, __builtin_amdgcn_readfirstlane(b_lds[0][0]));
+        G::load(Bs[0][1], B, {0, 0, bcol_base + block_col * 2 + 1, kk}, sw_B, b_srd, b_base, __builtin_amdgcn_readfirstlane(b_lds[0][1]));
+        G::load(scale_A_smem[0], scale_A_gl, {kk * tiles_M + block_row, 0, 0, 0});
+        G::load(scale_B_smem[0], scale_B_gl, {sb_base + kk * tiles_N + block_col, 0, 0, 0});
+        asm volatile("s_waitcnt vmcnt(0)");
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        __builtin_amdgcn_s_barrier();
+
+        fp8e8m0_4 sa_h0 = pack_scales(scale_A_smem[0].data, a_row_h0);
+        fp8e8m0_4 sa_h1 = pack_scales(scale_A_smem[0].data, a_row_h1);
+        fp8e8m0_4 sb_h0 = pack_scales(scale_B_smem[0].data, b_row_h0);
+        fp8e8m0_4 sb_h1 = pack_scales(scale_B_smem[0].data, b_row_h1);
+
+        auto bs0 = subtile_inplace<REG_N, BLOCK_K>(Bs[0][0], {warp_n, 0}); load(b0, bs0);
+        auto bs1 = subtile_inplace<REG_N, BLOCK_K>(Bs[0][1], {warp_n, 0}); load(b1, bs1);
+        auto as0 = subtile_inplace<REG_M, BLOCK_K>(As[0][0], {warp_m, 0}); load(a, as0);
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        mma_ABt_scaled(cA, a, b0, cA, &sa_h0, &sb_h0);
+        mma_ABt_scaled(cB, a, b1, cB, &sa_h0, &sb_h1);
+        auto as1 = subtile_inplace<REG_M, BLOCK_K>(As[0][1], {warp_m, 0}); load(a, as1);
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        mma_ABt_scaled(cC, a, b0, cC, &sa_h1, &sb_h0);
+        mma_ABt_scaled(cD, a, b1, cD, &sa_h1, &sb_h1);
+        __builtin_amdgcn_s_barrier();
+    }
+
+    store(C, cA, {0, 0, block_row * WARPS_ROW * 2 + warp_m, block_col * WARPS_COL * 2 + warp_n});
+    store(C, cB, {0, 0, block_row * WARPS_ROW * 2 + warp_m, block_col * WARPS_COL * 2 + WARPS_COL + warp_n});
+    store(C, cC, {0, 0, block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m, block_col * WARPS_COL * 2 + warp_n});
+    store(C, cD, {0, 0, block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m, block_col * WARPS_COL * 2 + WARPS_COL + warp_n});
+}

extra/thunder/amd/grouped_mxfp8_wgrad.cpp

@@ -0,0 +1,141 @@
+#include "kittens.cuh"
+
+using namespace kittens;
+
+#ifndef WGRAD_M
+constexpr int WGRAD_M = 8192;
+#endif
+#ifndef WGRAD_N
+constexpr int WGRAD_N = 8192;
+#endif
+#ifndef WGRAD_K
+constexpr int WGRAD_K = 8192;
+#endif
+#ifndef WGRAD_E
+constexpr int WGRAD_E = 8;
+#endif
+
+constexpr int NUM_WARPS  = 8;
+constexpr int WARPS_ROW  = 2;
+constexpr int WARPS_COL  = 4;
+constexpr int BLOCK_ROW  = 256;
+constexpr int BLOCK_COL  = 256;
+constexpr int BLOCK_K    = 128;
+constexpr int HALF_ROW   = BLOCK_ROW / 2;
+constexpr int HALF_COL   = BLOCK_COL / 2;
+constexpr int REG_M      = BLOCK_ROW / WARPS_ROW / 2;
+constexpr int REG_N      = BLOCK_COL / WARPS_COL / 2;
+
+using G = kittens::group<NUM_WARPS>;
+
+__global__ __launch_bounds__(512, 2) void grouped_mxfp8_wgrad_kernel(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr,
+    fp8e8m0 *scale_A_ptr, fp8e8m0 *scale_B_ptr,
+    const int *__restrict__ expert_off) {
+    constexpr int M = WGRAD_M, N = WGRAD_N, K = WGRAD_K, E = WGRAD_E;
+
+    kittens::gl<fp8e4m3, 1, 1, N, M>     A{A_ptr, nullptr, nullptr, nullptr, nullptr};  // g^T
+    kittens::gl<fp8e4m3, 1, 1, K, M>     B{B_ptr, nullptr, nullptr, nullptr, nullptr};  // x^T
+    kittens::gl<bf16, 1, 1, E * N, K>    C{C_ptr, nullptr, nullptr, nullptr, nullptr};  // grad_w, experts stacked
+
+    constexpr int m_blocks    = M / BLOCK_K;      // 128-wide blocks along the contraction (token) axis
+    constexpr int tiles_N     = N / BLOCK_ROW;
+    constexpr int tiles_K     = K / BLOCK_COL;
+    constexpr int NUM_THREADS = NUM_WARPS * WARP_THREADS;
+
+    kittens::gl<fp8e8m0, m_blocks * tiles_N, 1, 16, 64> scale_A_gl{scale_A_ptr, nullptr, nullptr, nullptr, nullptr};
+    kittens::gl<fp8e8m0, m_blocks * tiles_K, 1, 16, 64> scale_B_gl{scale_B_ptr, nullptr, nullptr, nullptr, nullptr};
+
+    using ST_A     = st_fp8e4m3<HALF_ROW, BLOCK_K, st_16x128_s>;
+    using ST_B     = st_fp8e4m3<HALF_COL, BLOCK_K, st_16x128_s>;
+    using ST_Scale = st<fp8e8m0, 16, 64, st_16x64_s>;
+    using RT_A     = rt_fp8e4m3<REG_M, BLOCK_K>;
+    using RT_B     = rt_fp8e4m3<REG_N, BLOCK_K>;
+    using RT_C     = rt_fl<REG_M, REG_N, col_l, rt_16x16_s>;
+
+    __shared__ ST_A As[2];
+    __shared__ ST_B Bs[2];
+    __shared__ ST_Scale scale_A_smem, scale_B_smem;
+
+    RT_A a;
+    RT_B b0, b1;
+    RT_C cA, cB, cC, cD;
+    zero(cA); zero(cB); zero(cC); zero(cD);
+
+    const int wg        = blockIdx.x;
+    const int e         = wg / (tiles_N * tiles_K);
+    const int rem       = wg % (tiles_N * tiles_K);
+    const int block_row = rem / tiles_K;   // over grad_w rows (N)
+    const int block_col = rem % tiles_K;   // over grad_w cols (K)
+
+    const int o0  = __builtin_amdgcn_readfirstlane(expert_off[e]);
+    const int kk0 = o0 / BLOCK_K;
+    const int nk  = (__builtin_amdgcn_readfirstlane(expert_off[e + 1]) - o0) / BLOCK_K;
+
+    const int warp_m = warpid() / WARPS_COL;
+    const int warp_n = warpid() % WARPS_COL;
+
+    using T = fp8e4m3;
+    constexpr int bpt      = ST_A::underlying_subtile_bytes_per_thread;
+    constexpr int bpm      = bpt * NUM_THREADS;
+    constexpr int copies_A = HALF_ROW * BLOCK_K * sizeof(T) / bpm;
+    constexpr int copies_B = HALF_COL * BLOCK_K * sizeof(T) / bpm;
+    uint32_t sw_A[copies_A], sw_B[copies_B];
+    G::prefill_swizzled_offsets(As[0], A, sw_A);
+    G::prefill_swizzled_offsets(Bs[0], B, sw_B);
+
+    const T *a_base = (const T *)&A[{0, 0, 0, 0}];
+    const T *b_base = (const T *)&B[{0, 0, 0, 0}];
+    const int a_row_stride = A.template stride<2>() * sizeof(T);
+    const int b_row_stride = B.template stride<2>() * sizeof(T);
+    i32x4 a_srd = make_srsrc(a_base, (uint32_t)((uint64_t)N * a_row_stride), a_row_stride);
+    i32x4 b_srd = make_srsrc(b_base, (uint32_t)((uint64_t)K * b_row_stride), b_row_stride);
+
+    const int wid = warpid() % NUM_WARPS;
+    constexpr int elem_per_warp = (16 / sizeof(T)) * kittens::WARP_THREADS;
+    uint32_t a_lds_0 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t a_lds_1 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&As[1].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_0 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[0].data[0]) + wid * elem_per_warp * sizeof(T)));
+    uint32_t b_lds_1 = __builtin_amdgcn_readfirstlane(static_cast<uint32_t>(reinterpret_cast<uintptr_t>(&Bs[1].data[0]) + wid * elem_per_warp * sizeof(T)));
+
+    const int a_row_h0 = warp_m * REG_M;
+    const int a_row_h1 = HALF_ROW + warp_m * REG_M;
+    const int b_row_h0 = warp_n * REG_N;
+    const int b_row_h1 = HALF_COL + warp_n * REG_N;
+
+    #pragma unroll 1
+    for (int t = 0; t < nk; t++) {
+        const int kk = kk0 + t;
+        G::load(As[0], A, {0, 0, block_row * 2,     kk}, sw_A, a_srd, a_base, __builtin_amdgcn_readfirstlane(a_lds_0));
+        G::load(As[1], A, {0, 0, block_row * 2 + 1, kk}, sw_A, a_srd, a_base, __builtin_amdgcn_readfirstlane(a_lds_1));
+        G::load(Bs[0], B, {0, 0, block_col * 2,     kk}, sw_B, b_srd, b_base, __builtin_amdgcn_readfirstlane(b_lds_0));
+        G::load(Bs[1], B, {0, 0, block_col * 2 + 1, kk}, sw_B, b_srd, b_base, __builtin_amdgcn_readfirstlane(b_lds_1));
+        G::load(scale_A_smem, scale_A_gl, {kk * tiles_N + block_row, 0, 0, 0});
+        G::load(scale_B_smem, scale_B_gl, {kk * tiles_K + block_col, 0, 0, 0});
+        asm volatile("s_waitcnt vmcnt(0)");
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        __builtin_amdgcn_s_barrier();
+
+        fp8e8m0_4 sa_h0 = pack_scales(scale_A_smem.data, a_row_h0);
+        fp8e8m0_4 sa_h1 = pack_scales(scale_A_smem.data, a_row_h1);
+        fp8e8m0_4 sb_h0 = pack_scales(scale_B_smem.data, b_row_h0);
+        fp8e8m0_4 sb_h1 = pack_scales(scale_B_smem.data, b_row_h1);
+
+        auto bs0 = subtile_inplace<REG_N, BLOCK_K>(Bs[0], {warp_n, 0}); load(b0, bs0);
+        auto bs1 = subtile_inplace<REG_N, BLOCK_K>(Bs[1], {warp_n, 0}); load(b1, bs1);
+        auto as0 = subtile_inplace<REG_M, BLOCK_K>(As[0], {warp_m, 0}); load(a, as0);
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        mma_ABt_scaled(cA, a, b0, cA, &sa_h0, &sb_h0);
+        mma_ABt_scaled(cB, a, b1, cB, &sa_h0, &sb_h1);
+        auto as1 = subtile_inplace<REG_M, BLOCK_K>(As[1], {warp_m, 0}); load(a, as1);
+        asm volatile("s_waitcnt lgkmcnt(0)");
+        mma_ABt_scaled(cC, a, b0, cC, &sa_h1, &sb_h0);
+        mma_ABt_scaled(cD, a, b1, cD, &sa_h1, &sb_h1);
+        __builtin_amdgcn_s_barrier();
+    }
+
+    const int crow_base = e * (N / REG_M);   // grad_w rows are experts stacked; store coord is in REG_M units
+    store(C, cA, {0, 0, crow_base + block_row * WARPS_ROW * 2 + warp_m,              block_col * WARPS_COL * 2 + warp_n});
+    store(C, cB, {0, 0, crow_base + block_row * WARPS_ROW * 2 + warp_m,              block_col * WARPS_COL * 2 + WARPS_COL + warp_n});
+    store(C, cC, {0, 0, crow_base + block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m,  block_col * WARPS_COL * 2 + warp_n});
+    store(C, cD, {0, 0, crow_base + block_row * WARPS_ROW * 2 + WARPS_ROW + warp_m,  block_col * WARPS_COL * 2 + WARPS_COL + warp_n});
+}

extra/thunder/tiny/fa.py

@@ -16,7 +16,7 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None) -> Tensor:
   if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=ref.dtype, device=ref.device)
   shape = tuple(s // len(ref.device) if i == ref.uop.axis else s for i, s in enumerate(shape))
   axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.empty(*shape, dtype=ref.dtype, device=ref.device).uop.multi(axis), dtype=ref.dtype, device=ref.device)
+  return Tensor(Tensor.empty(*shape, dtype=ref.dtype, device=ref.device).uop.unshard(axis), dtype=ref.dtype, device=ref.device)
 
 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
   return _sharded_empty(ref.shape, ref, axis)

extra/thunder/tiny/tk/group.py

@@ -224,7 +224,7 @@ class Group:
 
       # store to shared memory
       red_local_store = red_local[self.laneid].store(red_reg[0])
-      red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
+      red_local = red_local.after(red_local_store).reshape(red_local.shape)
 
       # reduce from shared memory
       for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
@@ -258,7 +258,7 @@ class Group:
 
       # store to shared memory
       red_local_store = red_local[self.laneid].store(red_reg[0])
-      red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
+      red_local = red_local.after(red_local_store).reshape(red_local.shape)
 
       # reduce from shared memory
       for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
@@ -342,7 +342,7 @@ class Group:
           if src.dtype != dst.dtype:
             src_load = src_load.cast(dst.dtype)
           dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
-          dst_store = dst_store.end(height, width, outer, inner).barrier()
+          dst_store = dst_store.end(height, width, outer, inner)
     elif dst.addrspace == AddrSpace.REG and src.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
       srcf = src.flatten()
       row_stride = prod(src.shape[axis+1:])

extra/thunder/tiny/tk/kernel.py

@@ -48,14 +48,14 @@ class Kernel(AbstractContextManager):
   @property
   def warpgroup(self): return self.group(4)
 
-  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.LOOP, track:bool=True):
+  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.WEAK, track:bool=True):
     if end == 0: start, end = 0, start
     rng = _tk_range(start, end, step, axis_type, self.range_id)
     self.range_id += 1
     if track: self.range_stack.append(rng)
     return rng
 
-  def raw_range(self, end:int=0, axis_type:AxisType=AxisType.LOOP):
+  def raw_range(self, end:int=0, axis_type:AxisType=AxisType.WEAK):
     rng = UOp.range(end, self.range_id, axis_type=axis_type)
     self.range_id += 1
     return rng

extra/torch_backend/backend.py

@@ -3,7 +3,7 @@
 # A002 Function argument `input` is shadowing a Python builtin
 # A006 Lambda argument `input` is shadowing a Python builtin
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.uop.ops import Ops
+from tinygrad.uop.ops import Ops, GroupOp
 from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
 import torch.lib
 TORCH_DEBUG = getenv("TORCH_DEBUG")
@@ -18,12 +18,16 @@ def _to_torch_device(device: str): return torch.device("tiny", int(device.partit
 
 import torch.utils.cpp_extension
 mod = torch.utils.cpp_extension.load(name="custom_device_extension", sources=[str(pathlib.Path(__file__).parent / "wrapped_tensor.cpp")])
+# TODO: this assumes a contiguous source, so PERMUTE/EXPAND/PAD/FLIP are wrong. UOp.contiguous_view_offset does it
+# properly, but it needs a device (these are deviceless)
+alias_ops = GroupOp.Movement | {Ops.BITCAST, Ops.DETACH, Ops.AFTER}
 def calculate_storage_offset(x: Tensor) -> int:
-  offset = 0
-  for u in x.uop.toposort():
-    if u.op == Ops.SHRINK:
+  offset, u = 0, x.uop
+  while u.op in alias_ops:
+    if u.op is Ops.SHRINK:
       u_strides = strides_for_shape(u.src[0].shape)
       for i, (start, _) in enumerate(u.marg): offset += start * u_strides[i]
+    u = u.src[0]
   return offset
 def wrap(x: Tensor, dev: torch.device|None=None) -> torch.Tensor:
   x._strides = strides_for_shape(x.shape) # always recalculate
@@ -220,7 +224,7 @@ def _as_strided(tensor:Tensor, size, stride, storage_offset=0):
 
 @torch.library.impl("aten::as_strided", "privateuseone")
 def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
-  storage_offset = storage_offset or tensor.storage_offset()
+  if storage_offset is None: storage_offset = tensor.storage_offset()
   return _as_strided(tensor, size, stride, storage_offset)
 
 @torch.library.impl("aten::_reshape_alias", "privateuseone")
@@ -228,16 +232,16 @@ 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, layout=None, device=None, pin_memory=False):
+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), device=_from_torch_device(device)).contiguous()
+  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)).contiguous()
+  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")
@@ -551,6 +555,8 @@ 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}"
+    # 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
@@ -652,35 +658,6 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.unfold": Tensor.unfold,
 }}
 
-# operations that need inplace treatment (use _inplace_op instead of wrap_fxn) AKA return original tensor
-inplace_ops = {
-  "aten.zero_",
-  "aten.fill_.Scalar",
-  "aten.fill_.Tensor",
-  "aten.add_.Tensor",
-  "aten.add_.Scalar",
-  "aten.mul_.Tensor",
-  "aten.mul_.Scalar",
-  "aten.floor_divide_.Tensor",
-  "aten.__ilshift__.Scalar",
-  "aten.__irshift__.Scalar",
-  "aten.relu_",
-  "aten.random_",
-  "aten.random_.from",
-  "aten.uniform_",
-  "aten.normal_",
-  "aten.logical_or_",
-  "aten.masked_fill_.Scalar",
-  "aten.masked_fill_.Tensor",
-}
-
-inplace_view_ops = {
-  "aten.squeeze_.dim",
-  "aten.unsqueeze_",
-  "aten.transpose_",
-  "aten.t_",
-}
-
 def wrap_fxn(k,f):
   def nf(*args, **kwargs):
     if TORCH_DEBUG:
@@ -694,7 +671,7 @@ def wrap_fxn(k,f):
     else: raise RuntimeError(f"unknown output type {type(out)}")
   return nf
 
-def wrap_inplace(k,f):
+def wrap_inplace(f):
   def nf(*args, **kwargs):
     orig = args[0]
     args, kwargs = unwrap_args(args, kwargs)
@@ -702,7 +679,7 @@ def wrap_inplace(k,f):
     return orig
   return nf
 
-def wrap_inplace_view_op(k,f):
+def wrap_inplace_view_op(f):
   def nf(*args, **kwargs):
     orig = args[0]
     args, kwargs = unwrap_args(args, kwargs)
@@ -735,11 +712,17 @@ def wrap_inplace_view_op(k,f):
     return orig
   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 must have come from tiny_backend_out so that wrap_out was applied
 for k,v in tiny_backend.items():
-  if k in inplace_view_ops: wrapper = wrap_inplace_view_op
-  elif k in inplace_ops: wrapper = wrap_inplace
-  else: wrapper = wrap_fxn
-  torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrapper(k,v))
+  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 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")
 def equal(x: torch.Tensor, y: torch.Tensor): return (x==y).all().item()
@@ -775,21 +758,17 @@ def native_batch_norm(input, weight, bias, running_mean, running_var, training,
 @torch.library.impl("aten::native_batch_norm_backward", "privateuseone")
 def native_batch_norm_backward(grad_out, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps, output_mask):
   grad_out_t, input_t = unwrap(grad_out), unwrap(input)
-  weight_t = unwrap(weight) if weight is not None else None
-  save_mean_t = unwrap(save_mean)
-  save_invstd_t = unwrap(save_invstd)
-  out = input_t.batchnorm(weight_t, None, save_mean_t, save_invstd_t)
-  targets = [t for t, m in zip([input_t, weight_t], output_mask[:2]) if t is not None and m]
-  if targets:
-    grads = out.gradient(*targets, gradient=grad_out_t)
-    grad_input = grads.pop(0) if output_mask[0] else None
-    grad_weight = grads.pop(0) if output_mask[1] and weight_t is not None else None
-  else:
-    grad_input, grad_weight = None, None
-  grad_bias = grad_out_t.sum(axis=tuple(x for x in range(grad_out_t.ndim) if x != 1)) if output_mask[2] else None
-  return (wrap(grad_input) if grad_input is not None else None,
-          wrap(grad_weight) if grad_weight is not None else None,
-          wrap(grad_bias) if grad_bias is not None else None)
+  dims, shape = tuple(x for x in range(input_t.ndim) if x != 1), (1, -1) + (1,)*(input_t.ndim-2)
+  # training differentiates the batch stats it was given, eval treats the running stats as constants
+  if train: mean, invstd = unwrap(save_mean), unwrap(save_invstd)
+  else: mean, invstd = unwrap(running_mean), unwrap(running_var).add(eps).rsqrt()
+  xhat = (input_t - mean.reshape(shape)) * invstd.reshape(shape)
+  grad_bias, grad_weight = grad_out_t.sum(axis=dims), (grad_out_t * xhat).sum(axis=dims)
+  grad_input = grad_out_t if not train else \
+    grad_out_t - (grad_bias.reshape(shape) + xhat * grad_weight.reshape(shape)) / (input_t.numel() // input_t.shape[1])
+  grad_input = grad_input * invstd.reshape(shape) * (unwrap(weight).reshape(shape) if weight is not None else 1)
+  return (wrap(grad_input) if output_mask[0] else None, wrap(grad_weight) if output_mask[1] else None,
+          wrap(grad_bias) if output_mask[2] else None)
 
 # _pad_circular is not CompositeImplicitAutograd (unlike reflect/replicate pad)
 # we need torch.autograd.Function with explicit AutogradPrivateUse1 registration

extra/torch_backend/test.py

@@ -71,6 +71,48 @@ class TestTorchBackend(unittest.TestCase):
     a = a.as_strided((1,1,5,5), (50,50,7,1), storage_offset=21)
     np.testing.assert_equal(a.cpu().numpy().sum(-1), [[[115,150,185,220,255]]])
 
+  def test_storage_offset_of_computed_tensor(self):
+    # a computed result owns its storage, so a slice anywhere in its history must not shift the offset
+    a = torch.arange(8., device=device)
+    self.assertEqual((a[3:]+1).storage_offset(), 0)
+
+  def test_storage_offset_through_aliases(self):
+    a = torch.arange(8., device=device)[3:]
+    self.assertEqual(a.detach().storage_offset(), 3)
+    self.assertEqual(a.view(torch.int32).storage_offset(), 3)
+    torch.add(torch.ones(5, device=device), torch.ones(5, device=device), out=a)
+    self.assertEqual(a.detach().storage_offset(), 3)
+
+  @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)
+    self.assertEqual(a.permute(1,0)[1:].storage_offset(), 1)
+    self.assertEqual(a.flatten()[3:].flip(0).storage_offset(), 0)
+
+  def test_as_strided_explicit_zero_offset(self):
+    # storage_offset=0 is a real offset, not "unspecified": it must not fall back to the input's own offset
+    a = torch.arange(6., device=device)
+    np.testing.assert_equal(a[3:].as_strided((2,), (1,), 0).cpu().numpy(), [0,1])
+    np.testing.assert_equal(a[3:].as_strided((2,), (1,)).cpu().numpy(), [3,4])
+
+  def test_empty_strided_default_dtype(self):
+    self.assertEqual(torch.empty_strided((2,3), (1,2), device=device).dtype, torch.get_default_dtype())
+
+  @unittest.expectedFailure  # TODO: empty_strided ignores the requested strides, the backend treats everything as contiguous
+  def test_empty_strided_honors_strides(self):
+    self.assertEqual(tuple(torch.empty_strided((2,3), (1,2), device=device).stride()), (1,2))
+
+  @unittest.expectedFailure  # TODO: torch refuses an out= that overlaps an input, we compute silently
+  def test_out_overlapping_input_is_rejected(self):
+    x = torch.arange(6., device=device)
+    with self.assertRaises(RuntimeError): torch.add(x[:-1], 10, out=x[1:])
+
+  def test_out_disjoint_input_is_allowed(self):
+    # torch permits an out= that shares a base with an input as long as they do not overlap
+    x, xc = torch.arange(6., device=device), torch.arange(6.)
+    torch.add(x[:3], 10, out=x[3:]); torch.add(xc[:3], 10, out=xc[3:])
+    np.testing.assert_equal(x.cpu().numpy(), xc.numpy())
+
   def test_plus_inplace(self):
     a = torch.ones(4, device=device)
     b = torch.ones(4, device=device)
@@ -316,6 +358,21 @@ class TestTorchBackend(unittest.TestCase):
     assert b.shape == (4, 2, 3)
     np.testing.assert_equal(b.cpu().numpy(), a.cpu().numpy().transpose(2, 0, 1))
 
+  def test_batchnorm_backward_realized_stats(self):
+    # the saved stats are a function of input in training, so grad_input must flow through them even when handed in realized.
+    # the backward eps is unused in training: torch differentiates the save_invstd it was given
+    x0, g0 = torch.randn(8, 4, 3, 3), torch.randn(8, 4, 3, 3)
+    def run(dev, bwd_eps):
+      x, go = x0.to(dev), g0.to(dev)
+      w, b = torch.linspace(0.5, 2.0, 4).to(dev), torch.zeros(4, device=dev)
+      rm, rv = torch.zeros(4, device=dev), torch.ones(4, device=dev)
+      out, sm, si = torch.ops.aten.native_batch_norm(x, w, b, rm, rv, True, 0.1, 1e-5)
+      grads = torch.ops.aten.native_batch_norm_backward(go, x, w, rm, rv, sm.clone().detach(), si.clone().detach(),
+                                                        True, bwd_eps, [True,True,True])
+      return [t.cpu().numpy() for t in grads]
+    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_batchnorm_unsqueeze(self):
     bn = torch.nn.BatchNorm2d(4).to(device)
     x = torch.randn(8, 4, 3, 3, device=device)
@@ -742,6 +799,20 @@ class TestTorchBackend(unittest.TestCase):
 
 from tinygrad import Tensor
 class TestBackendHelpers(unittest.TestCase):
+  def test_unwrap_rejects_foreign_tensor(self):
+    # unwrap casts to the tiny impl, so a tensor from another backend must be refused rather than reinterpreted
+    with self.assertRaises(RuntimeError): extra.torch_backend.backend.unwrap(torch.ones(4))
+
+  def test_update_metadata_rejects_foreign_tensor(self):
+    # resizing a tensor we don't own would expose memory past its allocation
+    t = torch.ones(4)
+    with self.assertRaises(RuntimeError): extra.torch_backend.backend.mod.update_metadata(t, [8], [1], 0)
+    self.assertEqual(t.shape, (4,))
+
+  def test_unwrap_parameter_and_detached(self):
+    # nn.Parameter and detach rebuild the base OpaqueTensorImpl, which unwrap still has to accept
+    extra.torch_backend.backend.unwrap(torch.nn.Parameter(torch.ones(4, device="tiny")))
+    extra.torch_backend.backend.unwrap(torch.ones(4, device="tiny").detach())
 
   def test_calculate_storage_offset_no_shrink(self):
     t = Tensor.ones(3, 4)

extra/torch_backend/wrapped_tensor.cpp

@@ -124,16 +124,21 @@ at::Tensor wrap_tensor(py::object &py_obj, c10::ScalarType dtype, c10::DeviceInd
     sizes, strides, storage_offset);
 }
 
+// shallow_copy_and_detach (nn.Parameter, aten.detach) rebuilds the base OpaqueTensorImpl, so that is the type every tiny tensor has
+at::OpaqueTensorImpl<std::shared_ptr<c10::SafePyObject>> *tiny_impl(const at::Tensor &tensor) {
+  auto* impl = dynamic_cast<at::OpaqueTensorImpl<std::shared_ptr<c10::SafePyObject>>*>(tensor.unsafeGetTensorImpl());
+  TORCH_CHECK(impl != nullptr, "expected a tiny tensor, got a ", tensor.device().str(), " one. move it with .to(\"tiny\") first");
+  return impl;
+}
+
 py::object unwrap_tensor(const at::Tensor &tensor) {
-  auto* impl = tensor.unsafeGetTensorImpl();
-  auto* opaque_impl = static_cast<at::TinyOpaqueTensorImpl<std::shared_ptr<c10::SafePyObject>>*>(impl);
-  std::shared_ptr<c10::SafePyObject> tiny = opaque_impl->opaque_handle();
+  std::shared_ptr<c10::SafePyObject> tiny = tiny_impl(tensor)->opaque_handle();
   return py::reinterpret_borrow<py::object>(tiny->ptr(getPyInterpreter()));
 }
 
 void update_metadata(const at::Tensor &tensor, const std::vector<int64_t> &sizes,
                      const std::vector<int64_t> &strides, int64_t storage_offset) {
-  auto* impl = tensor.unsafeGetTensorImpl();
+  auto* impl = tiny_impl(tensor);
   impl->set_allow_tensor_metadata_change(true);
   impl->set_sizes_and_strides(sizes, strides, storage_offset);
 }

pyproject.toml

@@ -74,7 +74,6 @@ testing_minimal = [
   "torch==2.9.1",
   "pytest",
   "pytest-xdist",
-  "pytest-timeout",
   "pytest-split",
   "hypothesis>=6.148.9",
   "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
@@ -160,8 +159,6 @@ norecursedirs = [
   ".hypothesis",
   ".git",
 ]
-timeout = 300
-timeout_func_only = true
 testpaths = ["test"]
 filterwarnings = [
   # Ignore SWIG warnings from importlib

spec/tinyspec.tex

@@ -29,6 +29,7 @@
 \definecolor{axbrred}{HTML}{E53935}      % GROUP_REDUCE
 \definecolor{axyellow}{HTML}{F9A825}     % UPCAST
 \definecolor{axmagenta}{HTML}{7B1FA2}    % UNROLL
+\definecolor{axgreen}{HTML}{2E7D32}      % DEVICE
 
 \title{tinygrad: a single dialect from Tensor programs to Command Buffers}
 \author{tinygrad, Corp. \\ \texttt{[email protected]}}
@@ -79,6 +80,7 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\
 \op{Stack} & $(T_0, T_1, \ldots)$ & --- & Join along a newly created leading axis. All shapes must match. \\
 \op{Bitcast} & $(T,)$ & dtype & Reinterpret storage as target dtype; preserve total bytes. \\
+\op{Unshard} & $(T, R_0, R_1, \ldots)$ & axes $(a_0, a_1, \ldots)$ & Concatenate shards of \op{Range} $R_k$ along axis $a_k$; $R_k$ is outer. \\
 \bottomrule
 \end{tabular}
 
@@ -258,6 +260,7 @@ Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{addrspa
 \op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
 \op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
 Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
+\op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, each $a_k \times n_k$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & remove first $n$ axes & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]
 \op{Cast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & clamped to dtype \\
 \op{Bitcast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
@@ -283,9 +286,9 @@ $[a,A]$, $[b,B]$, $[c,C]$ denote min\_max of $\mathrm{src}[0]$, $\mathrm{src}[1]
 Default \emph{dtype range}: $[\mathrm{dtype\_min},\, \mathrm{dtype\_max}]$.
 
 \medskip
-\textbf{axis} tracks the multi-device sharding dimension. \op{Buffer} with $n$-tuple device: axis $= 0$ (device dim).
-\op{Reshape} remaps axis to preserve the shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts axis right by $|\mathbf{n}|$.
-\op{Reduce} on the shard axis $\to$ \textsc{null} (shard axis is among the first $n$ axes). \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.
+\textbf{sharding} tracks multi-device sharding as a set of (axis, \op{Range}) pairs. \op{Unshard} defines it: arg is the tuple of sharded axes, one \op{Range} in src per axis (positional: the $k$-th \op{Range} shards the $k$-th axis). \op{Buffer} with $n$-tuple device: sharded on axis $0$ (device dim). The single-axis convenience \textbf{axis} is \textsc{null} unless exactly one axis is sharded.
+\op{Reshape} remaps each sharded axis to preserve its shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts all sharded axes right by $|\mathbf{n}|$.
+\op{Reduce} on a sharded axis drops it. \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.
 
 %% ============================================================
 \subsection*{Kernel Optimizations (OptOps) \normalfont\small--- schedule-level transforms on kernel ranges}
@@ -297,6 +300,7 @@ Each kernel's iteration space is a set of \op{Range} axes. Every range has an \t
 \toprule
 \textbf{AxisType} & \textbf{Letter} & \textbf{Split from} & \textbf{Direction} & \textbf{Semantics} \\
 \midrule
+{\color{axgreen}\texttt{DEVICE}} & \texttt{d} & --- & --- & Multi-device sharding dimension. \\
 {\color{axblue}\texttt{GLOBAL}} & \texttt{g} & --- & --- & GPU global workgroup dimension. \\
 {\color{axcyan}\texttt{LOCAL}} & \texttt{l} & g, L & inner & Workgroup local dimension (shared memory). \\
 {\color{axbrcyan}\texttt{WARP}} & \texttt{w} & \multicolumn{2}{l}{(created by \op{TC})} & Warp-level lanes for tensor cores. \\
@@ -378,6 +382,8 @@ def scatter_add(T, idx, val):
 Let $D = (d_0, \ldots, d_{n-1})$ be an $n$-tuple device.
 \op{Copy} to an $n$-tuple device reshards with axis $= 0$. \op{Copy} never changes shape.
 
+\textbf{Sharding} splits a tensor along an axis across $n$ devices. It opens a \op{Range} of type \texttt{DEVICE} (a symbolic per-device index $d$), shrinks each device's view to its piece, then closes the range with \op{Unshard}$(T, R, a)$. The result is a logical tensor whose shape along axis $a$ is the full size; each device holds $1/n$ of it. \op{Unshard} is the inverse of sharding --- it marks the boundary between per-device computation and the logical multi-device tensor. The range need not be \texttt{DEVICE}; e.g.\ a \texttt{WARP} range closes the same way, concatenating per-lane shards along $a$ with the range as the outer factor. A tensor may be sharded along several axes at once: \op{Unshard}$(T, R_0, R_1, \ldots;\; a_0, a_1, \ldots)$ carries one \op{Range} per sharded axis, and every movement op maps each sharded axis independently.
+
 \begin{lstlisting}
 # T has shape (s,) on a single device.
 

test/amd/hw/helpers.py

@@ -169,11 +169,10 @@ def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
   return parse_output(bytes(out_buf), n_lanes)
 
 def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
-  """Run instructions on real AMD hardware via HIPCompiler and AMDProgram."""
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  """Run instructions on real AMD hardware via HIPCompiler and the AMD runtime."""
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
-  from tinygrad.helpers import flat_mv
+  from tinygrad.helpers import Target, flat_mv
 
   dev = Device["AMD"]
   compiler = HIPCompiler(dev.arch)  # type: ignore[attr-defined]
@@ -223,7 +222,7 @@ amdhsa.kernels:
 """
 
   lib = compiler.compile(asm_src)
-  prg = AMDProgram(dev, "test", lib)  # type: ignore[arg-type]
+  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
 
   buf_sz = _out_bytes(n_lanes)
   out_gpu = dev.allocator.alloc(buf_sz)

test/amd/hw/test_cdna_vop3.py

@@ -5,7 +5,7 @@ gfx950 hardware when USE_HW=1.
 """
 import ctypes, struct, unittest
 import tinygrad.runtime.autogen.amd.cdna.ins as cdna
-from tinygrad.helpers import flat_mv
+from tinygrad.helpers import Target, flat_mv
 from tinygrad.renderer.amd.dsl import NULL
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
@@ -42,8 +42,7 @@ def _run_emu(instructions: list, out_reg: int = 2) -> int:
   return out_buf[0]
 
 def _run_hw(instructions: list, out_reg: int = 2) -> int:
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device["AMD"]
@@ -86,7 +85,7 @@ amdhsa.kernels:
 ...
 .end_amdgpu_metadata
 """
-  prg = AMDProgram(dev, "test", HIPCompiler(dev.arch).compile(asm_src))
+  prg = dev.runtime(TinyELF(HIPCompiler(dev.arch).compile(asm_src), "test", Target("AMD", arch=dev.arch), ()))
   prg(global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
   out = bytearray(LANES * 4)
   dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)

test/amd/hw/test_dpp.py

@@ -6,7 +6,7 @@ when USE_HW=1.
 """
 import ctypes, unittest
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.helpers import flat_mv
+from tinygrad.helpers import Target, flat_mv
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
 
@@ -37,8 +37,7 @@ def _run_wave64_emu(instructions: list, out_reg: int = 1) -> list[int]:
   return list(out_buf)
 
 def _run_wave64_hw(instructions: list, out_reg: int = 1) -> list[int]:
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device["AMD"]
@@ -84,7 +83,7 @@ amdhsa.kernels:
 .end_amdgpu_metadata
 """
   lib = compiler.compile(asm_src)
-  prg = AMDProgram(dev, "test", lib)  # type: ignore[arg-type]
+  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
   out_gpu = dev.allocator.alloc(WAVE64 * 4)
   prg(out_gpu, global_size=(1, 1, 1), local_size=(WAVE64, 1, 1), wait=True)
   out = bytearray(WAVE64 * 4)

test/amd/hw/test_rdna4_permlane_var.py

@@ -5,7 +5,7 @@ real RDNA4 hardware when USE_HW=1.
 """
 import ctypes, unittest
 import tinygrad.runtime.autogen.amd.rdna4.ins as r4
-from tinygrad.helpers import flat_mv
+from tinygrad.helpers import Target, flat_mv
 from tinygrad.renderer.amd.dsl import NULL
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
@@ -36,8 +36,7 @@ def _run_emu(instructions: list, out_reg: int = 2) -> list[int]:
   return list(out_buf)
 
 def _run_hw(instructions: list, out_reg: int = 2) -> list[int]:
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device['AMD']
@@ -85,7 +84,7 @@ amdhsa.kernels:
 .end_amdgpu_metadata
 """
   lib = compiler.compile(asm_src)
-  prg = AMDProgram(dev, 'test', lib)
+  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
   out_gpu = dev.allocator.alloc(LANES * 4)
   prg(out_gpu, global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
   out = bytearray(LANES * 4)

test/amd/test_asm_kernel.py

@@ -36,7 +36,7 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
   A,B = A.flatten(), B.flatten()
   assert A.dtype == dtypes.uint32, f"buffer dtype must be uint32, got {A.dtype}"
   threads = UOp.special(A.numel(), "lidx0")
-  var = UOp.param(2, dtypes.index, vmin_vmax=(0, 10), name="var", addrspace=AddrSpace.ALU)
+  var = UOp.param(2, dtypes.int, vmin_vmax=(0, 10), name="var", addrspace=AddrSpace.ALU)
   insts = [
     s_load_b128(s[4:7], s[0:1]),
     s_load_b32(s[8], s[0:1], offset=0x10), # all threads load the same variable

test/amd/test_mockgpu_invalid.py

@@ -7,15 +7,16 @@ class TestMockGPUInvalidInstruction(unittest.TestCase):
     """Test that unsupported instructions raise immediately through the full MOCKGPU stack."""
     test_code = '''
 import struct
+from dataclasses import replace
 from tinygrad import Device, Tensor
 from tinygrad.engine.realize import compile_linear
-from tinygrad.runtime.ops_amd import AMDProgram
 
 dev = Device["AMD"]
 a = Tensor([1.0]).realize()
 b = a + 1
 linear = compile_linear(b.schedule_linear())
-lib = bytearray(linear.src[-1].src[0].src[3].arg)
+compiled_prg = linear.src[-1].src[0]
+lib = bytearray(compiled_prg.src[3].arg)
 
 # Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
 # VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400
@@ -27,7 +28,7 @@ for i in range(0, len(lib) - 4, 4):
     break
 assert found, "s_endpgm not found"
 
-patched_prg = AMDProgram(dev, "patched", bytes(lib))
+patched_prg = dev.runtime(replace(compiled_prg.to_elf(), name="patched", lib=bytes(lib)))
 b.uop.buffer.allocate()
 patched_prg(b.uop.buffer._buf, a.uop.buffer._buf, global_size=(1,1,1), local_size=(1,1,1))
 dev.synchronize()

test/amd/test_roundtrip.py

@@ -121,7 +121,8 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
     arch = self.arch
 
     from tinygrad.runtime.support.elf import elf_loader
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    from tinygrad.runtime.support.compiler_llvm import AMDLLVMCompiler
     from tinygrad.helpers import DEV
 
     kernels, _, _ = get_kernels_from_tinygrad(op_fn)

test/backend/test_amd_flash_attention.py

@@ -0,0 +1,65 @@
+import unittest
+import numpy as np
+
+from tinygrad import Device, Tensor, TinyJit
+from tinygrad.llm.kernels.amd import amd_flash_attention_decode, flash_attention_causal_cached
+
+
+@unittest.skipUnless(Device.DEFAULT.startswith("AMD"), "AMD flash attention required")
+class TestAMDFlashAttention(unittest.TestCase):
+  def _test_decode(self, max_kv_len:int, valid_kv_len:int, n_heads:int=16, n_kv_heads:int=2, quantized:bool=False):
+    rng = np.random.default_rng(1)
+    q_np = rng.standard_normal((1, n_heads, 1, 256)).astype(np.float16)
+    kv_np = rng.standard_normal((2, 1, n_kv_heads, max_kv_len, 256)).astype(np.float16)
+    scale_np = np.maximum(np.max(np.abs(kv_np.astype(np.float32)), axis=-1), 1e-8) / 127
+    if quantized:
+      kv_np = np.clip(np.rint(kv_np.astype(np.float32) / scale_np[..., None]), -127, 127).astype(np.int8)
+    q, kv = Tensor(q_np).realize(), Tensor(kv_np).realize()
+    scale = Tensor(scale_np.astype(np.float16)).realize() if quantized else None
+
+    @TinyJit
+    def decode(q:Tensor, kv:Tensor): return amd_flash_attention_decode(q, kv, valid_kv_len, max_kv_len, scale).realize()
+
+    out = None
+    for _ in range(3): out = decode(q, kv).numpy()
+    assert out is not None
+    q_ref = q_np[0, :, 0].astype(np.float32)
+    kv_ref = kv_np.astype(np.float32) * scale_np[..., None] if quantized else kv_np.astype(np.float32)
+    k_ref, v_ref = kv_ref[:, 0, :, :valid_kv_len]
+    expected = np.empty((n_heads, 256), dtype=np.float32)
+    for head in range(n_heads):
+      scores = q_ref[head] @ k_ref[head // (n_heads // n_kv_heads)].T / np.sqrt(256)
+      probs = np.exp(scores - scores.max())
+      expected[head] = probs @ v_ref[head // (n_heads // n_kv_heads)] / probs.sum()
+
+    self.assertTrue(np.isfinite(out).all())
+    np.testing.assert_allclose(out[0, :, 0], expected, rtol=2e-3, atol=2e-3)
+
+  def test_short_decode_is_finite_and_matches_reference(self): self._test_decode(8192, 25)
+
+  def test_q8_cache_matches_dequantized_reference(self): self._test_decode(8192, 25, quantized=True)
+
+  def test_q8_cached_prefill_matches_dequantized_reference(self):
+    rng = np.random.default_rng(2)
+    heads, kv_heads, tokens, dim = 16, 2, 32, 256
+    q = rng.standard_normal((1, heads, tokens, dim)).astype(np.float16)
+    kv = rng.standard_normal((2, 1, kv_heads, tokens, dim)).astype(np.float16)
+    scale = np.maximum(np.max(np.abs(kv.astype(np.float32)), axis=-1), 1e-8) / 127
+    packed = np.clip(np.rint(kv.astype(np.float32) / scale[..., None]), -127, 127).astype(np.int8)
+    got = flash_attention_causal_cached(Tensor(q).realize(), Tensor(packed).realize(), tokens, tokens,
+                                        Tensor(scale.astype(np.float16)).realize()).numpy()
+    dequant = packed.astype(np.float32) * scale.astype(np.float16).astype(np.float32)[..., None]
+    expected = np.empty_like(got)
+    for head in range(heads):
+      scores = q[0, head].astype(np.float32) @ dequant[0, 0, head // (heads // kv_heads)].T / np.sqrt(dim)
+      scores[np.triu_indices(tokens, 1)] = -np.inf
+      probs = np.exp(scores - scores.max(axis=-1, keepdims=True))
+      expected[0, head] = probs @ dequant[1, 0, head // (heads // kv_heads)] / probs.sum(axis=-1, keepdims=True)
+    np.testing.assert_allclose(got, expected, rtol=2e-3, atol=2e-3)
+
+  def test_six_query_heads_per_kv_head(self): self._test_decode(8192, 25, n_heads=12, n_kv_heads=2)
+
+  def test_hierarchical_decode_matches_reference(self): self._test_decode(16384, 4097)
+
+
+if __name__ == "__main__": unittest.main()

test/backend/test_arange.py

@@ -19,6 +19,12 @@ class TestArange(unittest.TestCase):
     self.assertLess(self._get_flops(Tensor.arange(256).clone(), np.arange(256)), 256*4)
     self.assertLess(self._get_flops(Tensor.arange(2560).clone(), np.arange(2560)), 2560*4)
 
+  def test_cat_complexity(self):
+    x = Tensor.arange(2**10) + Tensor.empty((), dtype=dtypes.uint32)
+    out = x.cat(x).cat(Tensor.empty(1, dtype=dtypes.uint32))
+    linear = compile_linear(out.schedule_linear())
+    self.assertLessEqual(estimate_uop(linear.src[-1]).ops, out.numel()*20)
+
   @unittest.skipIf(Device.DEFAULT == "CL", "flaky in CI")
   def test_arange_cumsum(self):
     np.testing.assert_equal(Tensor.arange(513).cumsum(0).numpy(), np.arange(513).cumsum())

test/backend/test_call.py

@@ -0,0 +1,39 @@
+import unittest, ctypes
+from tinygrad import Tensor, UOp
+from tinygrad.device import Device
+from tinygrad.dtype import dtypes
+from tinygrad.renderer.cstyle import CStyleLanguage
+from tinygrad.uop.ops import KernelInfo
+
+def call_out_kernel(F:UOp, C:UOp) -> UOp:
+  call = F[0].load().call(UOp.const(dtypes.int, 3), C[0], ret_dtype=dtypes.void)
+  return C.after(call)[1].store(C.after(call)[0].load() + 1).sink(arg=KernelInfo(name="call_out"))
+
+def call_ret_kernel(F:UOp, C:UOp) -> UOp:
+  val = F[0].load().call(UOp.const(dtypes.int, 21), ret_dtype=dtypes.int)
+  return C[0].store(val * 2).sink(arg=KernelInfo(name="call_ret"))
+
+@unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "TODO: CALL is rendered in C style only")
+class TestCall(unittest.TestCase):
+  def test_call_out_param(self):
+    called = []
+    @ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.POINTER(ctypes.c_int))
+    def fxn(n, out):
+      called.append(n)
+      out[0] = n * 2
+    f = Tensor([ctypes.cast(fxn, ctypes.c_void_p).value], dtype=dtypes.uint64, device="CPU")
+    c = Tensor.empty(2, dtype=dtypes.int, device="CPU")
+    c = Tensor.custom_kernel(f, c, fxn=call_out_kernel)[1]
+    self.assertEqual(c.tolist(), [6, 7])
+    self.assertEqual(called, [3])
+
+  def test_call_ret(self):
+    @ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int)
+    def fxn(n): return n + 1
+    f = Tensor([ctypes.cast(fxn, ctypes.c_void_p).value], dtype=dtypes.uint64, device="CPU")
+    c = Tensor.empty(1, dtype=dtypes.int, device="CPU")
+    c = Tensor.custom_kernel(f, c, fxn=call_ret_kernel)[1]
+    c.realize()
+    self.assertEqual(c.item(), 44)
+
+if __name__ == "__main__": unittest.main()

test/backend/test_custom_kernel.py

@@ -2,6 +2,7 @@ import unittest
 from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
+from tinygrad.renderer.ptx import PTXRenderer
 
 # **** kernels ****
 
@@ -55,7 +56,7 @@ def flip_contract_kernel(dest:UOp, src:UOp):
   return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))
 
 def slice_sum_kernel(dest:UOp, src:UOp):
-  G = UOp.range(src.shape[0], 0)
+  G = UOp.range(src.shape[0], 0, dtype=dtypes.int)
   slice_src = src[G, :]
   reg = UOp.placeholder((1,), dest.dtype, 0, addrspace=AddrSpace.REG)
   reg = reg.after(G)[0].set(0)
@@ -117,13 +118,20 @@ class TestCustomKernel(unittest.TestCase):
     out = c.flatten().tolist()
     assert all(x == 2 for x in out), "all 2"
 
+  def test_duplicate_call_arg(self):
+    x = Tensor.arange(4).clone().realize()
+    x = Tensor.custom_kernel(x, x, fxn=custom_add_one_kernel)[0]
+    # webgpu silently errors when a kernel has duplicate buffer args, so the list stays the same.
+    # https://gpuweb.github.io/gpuweb/#abstract-opdef-encoder-bind-groups-alias-a-writable-resource
+    self.assertEqual(x.tolist(), [1, 2, 3, 4] if Device.DEFAULT != "WEBGPU" else [0, 1, 2, 3])
+
   def test_simple_sharded(self):
     devs = ("CPU:0", "CPU:1")
 
     a = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
     b = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
     # ugly construction to get a sharded empty tensor
-    c = Tensor(Tensor.empty(8, 16, device=devs).uop.multi(0), device=devs)
+    c = Tensor(Tensor.empty(8, 16, device=devs).uop.unshard(0), device=devs)
     c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
     out = c.flatten().tolist()
     assert all(x == 2 for x in out), "all 2"
@@ -132,7 +140,7 @@ class TestCustomKernel(unittest.TestCase):
     # PYTHON backend explicitly checks for OOB access for wrong multi shape regression
     devs = ("PYTHON:0", "PYTHON:1")
     a = Tensor.ones(4, 4).contiguous().shard(devs, axis=0)
-    c = Tensor(Tensor.empty(2, 4, device=devs).uop.multi(0), device=devs)
+    c = Tensor(Tensor.empty(2, 4, device=devs).uop.unshard(0), device=devs)
     c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
     assert (c == 2).all().item()
 
@@ -199,18 +207,16 @@ class TestCustomKernel(unittest.TestCase):
     c = Tensor.empty(N, N)
 
     tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
-    err = (tst - (a@b)).square().max()
-    self.assertLess(err.item(), 1e-6)
+    self.assertTrue(tst.allclose(a@b, atol=1e-3).item())
 
   def test_gemm_multi(self):
     devs = ("CPU:0", "CPU:1")
     N = 16
     a = Tensor.randn(N, N).shard_(devs, axis=0)
     b = Tensor.randn(N, N).to(devs)
-    c = Tensor(Tensor.empty(N//2, N, device=devs).uop.multi(0), device=devs)
+    c = Tensor(Tensor.empty(N//2, N, device=devs).uop.unshard(0), device=devs)
     tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
-    err = (tst - (a@b)).square().max()
-    self.assertLess(err.item(), 1e-6)
+    self.assertTrue(tst.allclose(a@b, atol=1e-3).item())
 
   def test_gemm_backward_custom(self): self.test_gemm_backward(True)
   # NOTE: grad_fxn doesn't work with pyrender
@@ -233,14 +239,9 @@ class TestCustomKernel(unittest.TestCase):
     real_grad_a, real_grad_b = a.grad, b.grad
     Tensor.realize(ref, real_grad_a, real_grad_b)
 
-    err = (tst - ref).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-    err = (grad_a - real_grad_a).square().max()
-    self.assertLess(err.item(), 1e-6)
-
-    err = (grad_b - real_grad_b).square().max()
-    self.assertLess(err.item(), 1e-6)
+    self.assertTrue(tst.allclose(ref, atol=1e-3).item())
+    self.assertTrue(grad_a.allclose(real_grad_a, atol=1e-3).item())
+    self.assertTrue(grad_b.allclose(real_grad_b, atol=1e-3).item())
 
   def test_simple_qkv(self):
     N, d = 8, 4
@@ -253,8 +254,7 @@ class TestCustomKernel(unittest.TestCase):
     O_ref = ((Q @ K.T) / (d ** 0.5)) @ V
 
     Tensor.realize(O_custom, O_ref)
-    err = (O_custom - O_ref).square().max()
-    self.assertLess(err.item(), 1e-6)
+    self.assertTrue(O_custom.allclose(O_ref, atol=1e-3).item())
 
   def test_gemm_qkv(self):
     B, N, K_DIM, H_KV, REP, D = 2, 7, 6, 2, 2, 6
@@ -331,7 +331,7 @@ class TestCustomKernel(unittest.TestCase):
   def test_multi_invalids_custom_kernel_no_copy(self):
     devs = ("CPU:0", "CPU:1")
     a = Tensor.ones(4, 4).shard(devs, axis=0).realize()
-    c = Tensor(UOp.const(dtypes.float, Invalid, shape=(2, 4)).clone(device=devs).multi(0), device=devs)
+    c = Tensor(Tensor.invalids(2, 4, dtype=dtypes.float, device=devs).uop.unshard(0), device=devs)
     c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
     GlobalCounters.reset()
     c.realize()
@@ -396,7 +396,7 @@ class TestCustomKernel(unittest.TestCase):
     y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
     if use_custom:
       z = Tensor.empty_like(x)
-      z = Tensor.custom_kernel(y, y.T.T, fxn=custom_add_one_kernel)[0]
+      z = Tensor.custom_kernel(z, y.T.T, fxn=custom_add_one_kernel)[0]
     else: z = y.T.T+1
     GlobalCounters.reset()
     z.realize()
@@ -422,19 +422,76 @@ class TestCustomKernel(unittest.TestCase):
 
   @Context(DEV="CPU")
   def test_simple_from_source(self):
-    a = Tensor([0., 1., 2.]).realize()
-
-    src = "void test_src(float* restrict a) { a[0] = 1.0; }"
+    a = Tensor.arange(4).clone().realize()
+    src = "void test_src(int* restrict a) { a[0] = 1; }"
     # TODO: it currently requires a compiler for Ops.BINARY
     from tinygrad.device import Device
     binary = Device[a.device].renderer.compiler.compile(src)
     def custom_src_kernel(A:UOp) -> UOp:
       sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())),
-                                   UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
+    a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
+    self.assertEqual(a.tolist(), [[1, 2], [1, 3]])
 
-    a = Tensor.custom_kernel(a, fxn=custom_src_kernel)[0]
-    self.assertEqual(a.tolist(), [1., 1., 2.])
+class TestUnshardIndex(unittest.TestCase):
+  """Regression tests for INDEX on UNSHARD (fragment) resolution in schedule/multi.py.
+
+  A fragment is a per-thread REG buffer wrapped in UNSHARD over LOCAL thread ranges.
+  index_multi must resolve an INDEX on the UNSHARD view into an INDEX on the per-thread
+  shard. Two ownership patterns must work:
+    contiguous: idx = rng*shard_sz + local   (thread rng owns [rng*shard_sz, ...))
+    strided:    idx = rng + ir*shard_sz      (thread rng owns {rng, rng+shard_sz, ...})
+  """
+  def _run(self, kernel, shape=(8, 8)):
+    c = Tensor.empty(*shape)
+    out = Tensor.custom_kernel(c, fxn=kernel)[0]
+    try: return out.numpy()
+    except RuntimeError as e:
+      if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
+        self.skipTest("PTX does not support dynamic register indexing")
+      raise
+
+  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
+  def test_contiguous_fragment_index(self):
+    # thread ty owns rows [ty*8, ty*8+8) of a 64-row fragment -- contiguous ownership.
+    # This is the pre-existing case that index_multi always handled.
+    def kernel(C:UOp) -> UOp:
+      ty = UOp.range(8, 0, AxisType.LOCAL)
+      ir = UOp.range(8, 1, AxisType.LOOP)
+      j = UOp.range(8, 2, AxisType.LOOP)
+      # 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows [ty*8, ty*8+8).
+      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
+      return C[ty*8 + ir, j].store(frag[ty*8 + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="contig_frag"))
+    out = self._run(kernel, (64, 8))
+    assert out.shape == (64, 8)
+
+  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
+  def test_strided_fragment_index(self):
+    # thread ty owns rows {ty, ty+8, ty+16, ty+24, ..., ty+56} of a 64-row fragment --
+    # strided ownership. idx = ty + ir*8 where shard_sz=8 (8 threads, shard rows=8).
+    # The contiguous check (idx - rng*shard_sz) fails; the strided check
+    # (idx-rng) % shard_sz == 0 must succeed. This is the pattern the index_multi fix adds.
+    def kernel(C:UOp) -> UOp:
+      ty = UOp.range(8, 0, AxisType.LOCAL)
+      ir = UOp.range(8, 1, AxisType.LOOP)
+      j = UOp.range(8, 2, AxisType.LOOP)
+      # 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows {ty, ty+8, ..., ty+56}.
+      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
+      return C[ty + ir*8, j].store(frag[ty + ir*8, j]).end(j, ir, ty).sink(arg=KernelInfo(name="strided_frag"))
+    out = self._run(kernel, (64, 8))
+    assert out.shape == (64, 8)
+
+  def test_fragment_index_cannot_shard(self):
+    # thread ty indexing rows [ty, ty+8) overlaps with other threads' rows -- this matches neither
+    # the contiguous nor the strided ownership pattern, so index_multi must raise.
+    def kernel(C:UOp) -> UOp:
+      ty = UOp.range(8, 0, AxisType.LOCAL)
+      ir = UOp.range(8, 1, AxisType.LOOP)
+      j = UOp.range(8, 2, AxisType.LOOP)
+      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
+      return C[ty + ir, j].store(frag[ty + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="bad_frag"))
+    with self.assertRaisesRegex(RuntimeError, "cannot shard index"):
+      self._run(kernel, (64, 8))
 
 class TestUOpReduce(unittest.TestCase):
   def test_uop_sum(self):

test/backend/test_dtype.py

@@ -124,6 +124,7 @@ class TestFp8sConversions(unittest.TestCase):
   def test_float_to_fp8e4m3(self, x):
     np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
 
+  @unittest.skip("fp8 overflow semantics are inconsistent")
   def test_float_to_fp8e4m3_extreme_values(self):
     for x in [FP8E4M3_MAX, FP8E4M3_MAX*1.01, -FP8E4M3_MAX, -FP8E4M3_MAX*1.01, math.inf, -math.inf, math.nan, -math.nan]:
       np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())

test/backend/test_linearizer.py

@@ -13,6 +13,7 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.isa import ISARenderer
 from test.helpers import replace_opts
+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
@@ -290,7 +291,7 @@ class TestLinearizer(unittest.TestCase):
   @unittest.skipIf(MOCKGPU and isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CUDARenderer)), "PTX indexes differently. might be ok?")
   def test_where_fold(self):
     a = Tensor.ones(4, 4).contiguous().realize()
-    b = a.shrink(((1, 2), None)).pad(((1, 2), None))
+    b = a.shrink(((1, 2), None)).pad(((1, 2), None)).bool()
     a.assign(b.where(2, a))
     linear, var_vals = a.linear_with_vars()
     assert len(linear.src) == 1
@@ -392,6 +393,16 @@ class TestLinearizer(unittest.TestCase):
     # the global store doesn't change
     assert stores[1].src[1].dtype == dtypes.float
 
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  def test_two_grouped_stores_local(self):
+    # GROUP on both reduces puts two LOCAL buffers in one kernel, and the store to each needs its own barrier
+    a = Tensor.rand(32, 32).realize()
+    opts = [Opt(OptOps.GROUP, 1, 4), Opt(OptOps.GROUP, 2, 4)]
+    ast = helper_linearizer_opt(single_kernel_softmax(a), [opts])
+    uops = to_program(replace_opts(ast, opts), renderer=Device[Device.DEFAULT].renderer).src[1].src
+    self.assertEqual(len([u for u in uops if u.op is Ops.BARRIER]), 2)
+
 # *** helpers ***
 
 def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:

test/backend/test_linearizer_dumb.py

@@ -12,18 +12,18 @@ class TestLinearizerFailure(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
   def test_failure_beam_mnist(self):
     c0 = UOp.param(0, dtypes.uchar, (4014080,))
-    c1 = UOp.range(UOp.const(dtypes.index, 512), 0, AxisType.GLOBAL)
-    c2 = UOp.range(UOp.const(dtypes.index, 784), 1, AxisType.GLOBAL)
-    c3 = UOp.range(UOp.const(dtypes.index, 10), 3, AxisType.GLOBAL)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 0, AxisType.GLOBAL)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 784), 1, AxisType.GLOBAL)
+    c3 = UOp.range(UOp.const(dtypes.weakint, 10), 3, AxisType.GLOBAL)
     c4 = UOp.param(1, dtypes.int, (512,))
     c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
-    c6 = UOp.range(UOp.const(dtypes.index, 6000), 1004, AxisType.REDUCE)
-    c7 = UOp.range(UOp.const(dtypes.index, 3750), 2006, AxisType.REDUCE)
-    c8 = UOp.range(UOp.const(dtypes.index, 16), 2007, AxisType.GROUP_REDUCE)
+    c6 = UOp.range(UOp.const(dtypes.weakint, 6000), 1004, AxisType.REDUCE)
+    c7 = UOp.range(UOp.const(dtypes.weakint, 3750), 2006, AxisType.REDUCE)
+    c8 = UOp.range(UOp.const(dtypes.weakint, 16), 2007, AxisType.GROUP_REDUCE)
     c9 = UOp.param(2, dtypes.uchar, (47040000,))
-    c10 = c9.index((((c3*UOp.const(dtypes.index, 4704000))+c2)+(c6*UOp.const(dtypes.index, 784))).valid(UOp.const(dtypes.bool, True)))
-    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.index, 6000))+c6)+((c7*UOp.const(dtypes.index, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.index, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
-    c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
+    c10 = c9.index((((c3*UOp.const(dtypes.weakint, 4704000))+c2)+(c6*UOp.const(dtypes.weakint, 784))).valid(UOp.const(dtypes.bool, True)))
+    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
+    c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
     ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
     _ = to_program(ast, Device["METAL"].renderer)
 

test/backend/test_llama_kernels.py

@@ -80,7 +80,7 @@ class TestQuantizeFP8(unittest.TestCase):
   @needs_second_gpu
   def test_multi(self):
     devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8))
-    x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.multi(0)
+    x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.unshard(0)
     x = Tensor(x, device=devs)
     amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous()
     amax_out = Tensor.zeros((), dtype=dtypes.float32, device=devs).realize()
@@ -131,6 +131,7 @@ class TestFusedQKVRoPE(unittest.TestCase):
       self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")
 
   def test_llama31_8b_backward(self):
+    if not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"): self.skipTest("requires CDNA4")
     Tensor.manual_seed(1)
     B, N, H, H_KV, D = self.SHAPE
     PARTIALS = 2

test/backend/test_multitensor.py

@@ -1,6 +1,6 @@
 import unittest, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
-from tinygrad.uop.ops import Ops, UOp
+from tinygrad.uop.ops import Ops, UOp, AxisType
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import run_linear, compile_linear
@@ -52,8 +52,10 @@ class TestMultiTensor(unittest.TestCase):
   def test_shard(self):
     X = Tensor.ones(256).contiguous().realize()
     X.shard_(devices_2, 0)
-    for lb in X.uop.src:
-      assert lb.shape == (128,)
+    assert X.uop.src[0].shape == (128,)
+    # the MULTI carries and ends the DEVICE range as its second src
+    assert X.uop.src[1].op is Ops.RANGE and X.uop.src[1].arg[-1] is AxisType.DEVICE
+    assert X.uop.ended_ranges == X.uop.src[1:]
     (X + X).realize()
 
   @unittest.expectedFailure # TODO: fix
@@ -423,6 +425,49 @@ class TestMultiBufferView(unittest.TestCase):
     run_linear(linear, var_vals)
     np.testing.assert_equal(out.numpy(), ref[5].numpy())
 
+@unittest.skipIf(not_support_multi_device(), "need multi")
+class Test2DShard(unittest.TestCase):
+  def setUp(self):
+    self.devices_4 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
+    self.rng = UOp.range(4, -1, AxisType.DEVICE)
+    self.rng0, self.rng1 = self.rng // 2, self.rng % 2
+
+  def _shard_2d(self, t:Tensor) -> Tensor:
+    u = t.uop.copy_to_device(self.devices_4)._shard(0, self.rng0)._shard(1, self.rng1).unshard((0, 1), (self.rng0, self.rng1))
+    return Tensor(u)
+
+  def test_2d_shard_basic(self):
+    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = t.contiguous().realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy())
+
+  def test_2d_shard_elementwise(self):
+    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = (t + 1).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy() + 1)
+
+  def test_2d_shard_sum_all(self):
+    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = t.sum().contiguous().realize()
+    np.testing.assert_equal(out.numpy(), np.array(ref.numpy().sum()))
+
+  def test_2d_shard_sum_non_sharded_axis(self):
+    ref = Tensor.arange(4*4*2).reshape(4, 4, 2).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = t.sum(axis=2).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy().sum(axis=2))
+
+  def test_2d_shard_matmul(self):
+    a = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    b = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    a_s = self._shard_2d(a)
+    b_s = self._shard_2d(b)
+    out = (a_s @ b_s).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), a.numpy() @ b.numpy())
+
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiTransformer(unittest.TestCase):
   @needs_second_gpu
@@ -450,7 +495,7 @@ class TestMultiTransformer(unittest.TestCase):
       else: v.shard_(device, axis=None)
 
     last_tok = 0
-    for i in range(10):
+    for i in range(5):
       real_tok = real_model(Tensor([[last_tok]], device=Device.DEFAULT), i).item()
       shard_tok = shard_model(Tensor([[last_tok]], device=device), i).item()
 

test/backend/test_ops.py

@@ -5,7 +5,6 @@ import torch
 from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.renderer.cstyle import QCOMCLRenderer
 from tinygrad.renderer.nir import NIRRenderer
 
 TINY_BACKEND = getenv("TINY_BACKEND")
@@ -244,7 +243,6 @@ class TestOps(unittest.TestCase):
 
     self.helper_test_exception([(8,)], lambda x: x.unfold(0, 9, 3), expected=RuntimeError)
     self.helper_test_exception([(8,)], lambda x: x.unfold(1, 8, 3), expected=IndexError)
-    self.helper_test_exception([(8,)], lambda x: x.unfold(0, 9, 3), expected=RuntimeError)
     self.helper_test_exception([(8,)], lambda x: x.unfold(0, 1, -1), expected=RuntimeError)
 
   def test_meshgrid(self):
@@ -285,6 +283,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: torch.arange(-128, 128, dtype=torch.int8), lambda: Tensor.arange(-128, 128, dtype=dtypes.int8), forward_only=True)
     helper_test_op([], lambda: torch.arange(127, -129, -1, dtype=torch.int8),
                    lambda: Tensor.arange(127, -129, -1, dtype=dtypes.int8), forward_only=True)
+    # an int range too large for default_int picks int64
+    self.assertEqual(Tensor.arange(2**31, 2**31+3).dtype, dtypes.int64)
     # overflow: tinygrad raises (torch silently wraps)
     with self.assertRaises(OverflowError): Tensor.arange(2**33, dtype=dtypes.int)
     with self.assertRaises(OverflowError): Tensor.arange(129, dtype=dtypes.int8)            # last=128 overflows
@@ -339,10 +339,12 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: torch.ones(256,256).max(1)[0], lambda: Tensor.ones(256,256).max(1), forward_only=True)
 
   def test_where(self):
+    helper_test_op([], lambda: torch.where(torch.tensor([True, False]), 1, 3).type(torch.int32),
+                   lambda: Tensor([True, False]).where(1, 3).clone(), forward_only=True)
     helper_test_op(
       [(100,)],
       lambda x: torch.where(x > 0.5, 4, 2).type(torch.int32),
-      lambda x: (x > 0.5).where(4, 2), forward_only=True)
+      lambda x: (x > 0.5).where(4, 2).clone(), forward_only=True)
 
     for shps in [[(8,),(1,),(1,)], [(10,10),(10,),(10,)], [(100,)]*3, [(10,10)]*3]:
       helper_test_op(
@@ -354,7 +356,7 @@ class TestOps(unittest.TestCase):
     helper_test_op(
       [(5, 5)],
       lambda x: torch.where(x > 0.5, 4, 2).type(torch.int32).permute((1, 0)),
-      lambda x: (x > 0.5).where(4, 2).permute((1, 0)), forward_only=True)
+      lambda x: (x > 0.5).where(4, 2).clone().permute((1, 0)), forward_only=True)
 
   def _test_cmp(self, fxn, reverse=True):
     # test different dtypes
@@ -450,7 +452,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,35), (45,35), (45,35)], lambda x,y,z: x.lerp(y,z))
     helper_test_op(None, lambda x,y,z: x.lerp(y,z), vals=[[1.,2.,3.], [4.,5.,6.], 0.5])
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_tril(self):
     helper_test_op([(3,3)], lambda x: x.tril())
     helper_test_op([(3,3)], lambda x: x.tril(1))
@@ -468,7 +469,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(5,3,3)], lambda x: x.tril(1))
     helper_test_op(None, lambda x: x.tril(), vals=[[[True] * 3] * 3], forward_only=True)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_triu(self):
     helper_test_op([(3,3)], lambda x: x.triu())
     helper_test_op([(3,3)], lambda x: x.triu(1))
@@ -527,6 +527,7 @@ class TestOps(unittest.TestCase):
 
   def test_add(self):
     helper_test_op([(45,68), (45,68)], lambda x,y: x+y, Tensor.add)
+    helper_test_op([], lambda: torch.tensor(1)+0.5, lambda: Tensor(1)+0.5, forward_only=True)
     helper_test_op([(45,68), (45,68)], lambda x,y: x+y)
     helper_test_op([(), ()], lambda x,y: x+y)
   def test_add3(self):
@@ -635,9 +636,9 @@ class TestOps(unittest.TestCase):
         helper_test_op(None, lambda x,y: x%y, forward_only=True, vals=[va, vb])
         helper_test_op(None, lambda x: x%2, forward_only=True, vals=[va])
         helper_test_op(None, lambda x: x%3, forward_only=True, vals=[va])
-        helper_test_op(None, lambda x: x%3.5, forward_only=True, vals=[va])
+        helper_test_op(None, lambda x: x%3.5, lambda x: (x%3.5).clone(), forward_only=True, vals=[va])
         helper_test_op(None, lambda x: 100%x, forward_only=True, vals=[va])
-        helper_test_op(None, lambda x: 100.5%x, forward_only=True, vals=[va])
+        helper_test_op(None, lambda x: 100.5%x, lambda x: (100.5%x).clone(), forward_only=True, vals=[va])
 
   def test_fmod(self):
     a = [-4, 7, 5, 4, -7, 8, -9]
@@ -648,7 +649,7 @@ class TestOps(unittest.TestCase):
         vb = [float(bi) for bi in b] if float_b else b
         helper_test_op(None, lambda x,y: x.fmod(y), forward_only=True, vals=[va, vb])
         helper_test_op(None, lambda x: x.fmod(2), forward_only=True, vals=[va])
-        helper_test_op(None, lambda x: x.fmod(3.5), forward_only=True, vals=[va])
+        helper_test_op(None, lambda x: x.fmod(3.5), lambda x: x.fmod(3.5).clone(), forward_only=True, vals=[va])
 
   def test_mul_naninf(self):
     helper_test_op([(45,65)], lambda x: x*math.inf)
@@ -705,7 +706,7 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: 0.7**x, vals=[[-2.,-1,0,1,2,3]])
     helper_test_op(None, lambda x: (-2)**x, vals=[[-2.,-1,0,1,2,3]])
     # float to power of int
-    helper_test_op(None, lambda x: 0.7**x, vals=[[-2,-1,0,1,2,3]], forward_only=True)
+    helper_test_op(None, lambda x: 0.7**x, lambda x: (0.7**x).clone(), vals=[[-2,-1,0,1,2,3]], forward_only=True)
 
   @unittest.skipIf(COMPILE_ONLY, "test requires runtime")
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, NIRRenderer), "TODO: broken in LVP")
@@ -771,6 +772,11 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: torch.tensor([2], dtype=torch.int) ** torch.tensor(-2, dtype=torch.int),
                        lambda: Tensor([2]) ** Tensor(-2), forward_only=True)
 
+  def test_pow_int_base_float_exponent(self):
+    for exponent in (0.5, 1.5, 2.0, -1.0, 0.0):
+      helper_test_op([], lambda: torch.tensor([1, 2, 3, 4], dtype=torch.int) ** exponent,
+                         lambda: (Tensor([1, 2, 3, 4], dtype=dtypes.int32) ** exponent).clone(), forward_only=True)
+
   def test_sqrt(self):
     helper_test_op([(45,65)], lambda x: x.sqrt())
     helper_test_op(None, lambda x: x.sqrt(), vals=[[0.0]])
@@ -789,9 +795,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor^0x1337, lambda: ten^0x1337, forward_only=True)
     helper_test_op([], lambda: 0x1337^tor, lambda: 0x1337^ten, forward_only=True)
 
-    self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_xor(y), expected=RuntimeError)
-
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_and(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -807,9 +810,6 @@ class TestOps(unittest.TestCase):
 
     helper_test_op(None, lambda x: (1 < x) & (x < 2), forward_only=True, vals=[[1.2, 1.2, 1.2, 3.2]])
 
-    self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_and(y), expected=RuntimeError)
-
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_or(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -823,8 +823,6 @@ class TestOps(unittest.TestCase):
     ten0, ten1 = Tensor(data[0], dtype=dtypes.bool), Tensor(data[1], dtype=dtypes.bool)
     helper_test_op([], lambda: tor0|tor1, lambda: ten0|ten1, forward_only=True)
 
-    self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_or(y), expected=RuntimeError)
-
   def test_bitwise_not(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -838,8 +836,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor.bitwise_not(), lambda: ten.bitwise_not(), forward_only=True)
     helper_test_op([], lambda: ~tor, lambda: ~ten, forward_only=True)
 
-    self.helper_test_exception([(4)], lambda x: x.bitwise_not(), expected=RuntimeError)
-
   def test_lshift(self):
     data = [[0,1,2],[1<<8,1<<16,1<<31-1]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -852,6 +848,9 @@ class TestOps(unittest.TestCase):
                    lambda: (ten << Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.__lshift__(2), lambda: ten.__lshift__(2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.bitwise_left_shift(2), lambda: ten.lshift(2).cast(dtypes.int32), forward_only=True)
+    self.helper_test_exception([], lambda: torch.tensor([1.0]) << 2, lambda: Tensor([1.0]) << 2, expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << torch.tensor([1.0]), lambda: ten << Tensor([1.0]), expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << 1.0, lambda: ten << 1.0, expected=RuntimeError)
 
   def test_rshift(self):
     data = [[0,1,2],[1<<8,1<<16,1<<31-1]]
@@ -865,6 +864,8 @@ class TestOps(unittest.TestCase):
                    lambda: (ten >> Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.__rshift__(2), lambda: ten.__rshift__(2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.bitwise_right_shift(2), lambda: ten.rshift(2).cast(dtypes.int32), forward_only=True)
+    self.helper_test_exception([], lambda: torch.tensor([4.0]) >> 1, lambda: Tensor([4.0]) >> 1, expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor >> torch.tensor([1.0]), lambda: ten >> Tensor([1.0]), expected=RuntimeError)
 
   def test_lshift_signed(self):
     data = [[-1, -3, 1, 7], [0, -2147483648, 2147483647, -1]]
@@ -905,6 +906,7 @@ class TestOps(unittest.TestCase):
   @unittest.skipIf(DEV.renderer == "NAK", "MUFU.SIN is not accurate enough")
   def test_cos(self):
     helper_test_op([(45,65)], lambda x: x.cos())
+    helper_test_op([], lambda: torch.tensor(2.0).cos(), lambda: Tensor(2).cos(), forward_only=True)
     helper_test_op([()], lambda x: x.cos())
     if not ((DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
       helper_test_op(None, lambda x: x.cos(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
@@ -983,6 +985,8 @@ class TestOps(unittest.TestCase):
 
   def test_exp(self):
     helper_test_op([(45,65)], torch.exp, Tensor.exp)
+    # int scalar input: the python coefficients in the decomposition must stay float
+    helper_test_op([], lambda: torch.tensor(2.0).exp(), lambda: Tensor(2).exp(), forward_only=True)
     helper_test_op(None, torch.exp, Tensor.exp, vals=[[math.inf, -math.inf, math.nan]])
     helper_test_op([()], torch.exp, Tensor.exp)
   def test_exp2(self):
@@ -1048,8 +1052,8 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], torch.nn.functional.hardsigmoid, Tensor.hardsigmoid)
     helper_test_op([()], torch.nn.functional.hardsigmoid, Tensor.hardsigmoid)
   def test_hardsigmoid_extreme(self):
-    helper_test_op([(45,65)], torch.sigmoid, Tensor.sigmoid, low=300, high=400)
-    helper_test_op([(45,65)], torch.sigmoid, Tensor.sigmoid, low=-400, high=-300)
+    helper_test_op([(45,65)], torch.nn.functional.hardsigmoid, Tensor.hardsigmoid, low=300, high=400)
+    helper_test_op([(45,65)], torch.nn.functional.hardsigmoid, Tensor.hardsigmoid, low=-400, high=-300)
   def test_softplus(self):
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
     helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3), lambda t: Tensor.softplus(t, beta=3), grad_atol=1e-6)
@@ -1229,7 +1233,6 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x.type(torch.int32).argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[False, True]])
     helper_test_op(None, lambda x: x.type(torch.int32).argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[True, False]])
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_argmin(self):
     # check if it returns the first index for multiple occurrences
     helper_test_op(None, lambda x: x.argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True, vals=[[2, 2]])
@@ -1266,23 +1269,20 @@ class TestOps(unittest.TestCase):
                    lambda x: x.sort(descending=True)[1], forward_only=True, vals=[[0, 1] * 9])
 
   def test_argsort(self):
-    for dim in [-1, 0, 1]:
-      for descending in [True, False]:
-        helper_test_op([(8,8,6)], lambda x: torch.argsort(x, dim=dim, descending=descending, stable=True).type(torch.int32),
-                                  lambda x: x.argsort(dim, descending), forward_only=True)
+    helper_test_op([(8,8,6)], lambda x: torch.argsort(x, dim=1, descending=True, stable=True).type(torch.int32),
+                              lambda x: x.argsort(1, True), forward_only=True)
 
   def test_topk(self):
     helper_test_op([(8)], lambda x: x.topk(3).values, lambda x: x.topk(3)[0], forward_only=True)
     helper_test_op([(8)], lambda x: x.topk(3).indices.type(torch.int32), lambda x: x.topk(3)[1], forward_only=True)
-    for dim in [0, 1, -1]:
-      for largest in [True, False]:
-        for sorted_ in [True]: # TODO support False
-          helper_test_op([(5,5,4)],
-                          lambda x: x.topk(4, dim, largest, sorted_).values,
-                          lambda x: x.topk(4, dim, largest, sorted_)[0], forward_only=True)
-          helper_test_op([(5,5,4)],
-                          lambda x: x.topk(4, dim, largest, sorted_).indices.type(torch.int32),
-                          lambda x: x.topk(4, dim, largest, sorted_)[1], forward_only=True)
+    for dim, largest in [(0, True), (1, False)]:
+      for sorted_ in [True]: # TODO support False
+        helper_test_op([(5,5,4)],
+                        lambda x: x.topk(4, dim, largest, sorted_).values,
+                        lambda x: x.topk(4, dim, largest, sorted_)[0], forward_only=True)
+        helper_test_op([(5,5,4)],
+                        lambda x: x.topk(4, dim, largest, sorted_).indices.type(torch.int32),
+                        lambda x: x.topk(4, dim, largest, sorted_)[1], forward_only=True)
     # repeated values
     if not COMPILE_ONLY:
       value, indices = Tensor([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]).topk(3)
@@ -1535,7 +1535,6 @@ class TestOps(unittest.TestCase):
   def test_prod_dtype_arg(self):
     with self.assertRaises(AttributeError): Tensor([1.0, 2.0]).prod(dtype="")
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_min(self):
     helper_test_op([(3,3)], lambda x: x.min())
     helper_test_op([(45,3)], lambda x: x.min())
@@ -1575,7 +1574,6 @@ class TestOps(unittest.TestCase):
   def test_any_zero_axis(self):
     helper_test_op([(1,0,3,0,5)], lambda x: x.any(axis=(1,3)), forward_only=True)
 
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_all(self):
     helper_test_op([(3,4,5,6)], lambda x: x.all(), forward_only=True)
     helper_test_op(None, lambda x: x.all(), vals=[[True, True]], forward_only=True)
@@ -1907,9 +1905,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(3,3,3)], lambda x: x[-2:2])
     helper_test_op([(3,3,3)], lambda x: x[-2:-5])
 
-  def test_slice_empty(self):
-    helper_test_op([(10,10)], lambda x: x[1:1])
-
   def test_slice_zero_in_shape(self):
     helper_test_op([(10,10)], lambda x: x[1:1])  # x.shape = (0, 10)
     helper_test_op([(3,3,3)], lambda x: x[-2:-5])  # x.shape = (0, 3, 3)
@@ -2102,7 +2097,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(4,3,1,6)], lambda x: x.squeeze(1))
     helper_test_op([(4,3,6,6)], lambda x: x.squeeze(3))
     self.helper_test_exception([(4,3,6,6)], lambda x: x.squeeze(50), expected=IndexError)
-    self.helper_test_exception([(4,3,6,6)], lambda x: x.squeeze(50), expected=IndexError)
     helper_test_op([(4,3,6,1)], lambda x: x.squeeze(-1))
     helper_test_op([(4,3,6,6)], lambda x: x.squeeze())
     helper_test_op([(1,3,6,6)], lambda x: x.squeeze())
@@ -2375,9 +2369,10 @@ class TestOps(unittest.TestCase):
               lambda x,w: Tensor.conv2d(x,w,groups=groups), grad_rtol=1e-5)
   def test_conv2d(self): self._test_conv2d(bs=1, cin=3)
   @slow_test
+  @unittest.skip("redundant: bs/cout are loop dims, kernel×cin sweep covered by test_conv2d")
   def test_conv2d_bs_4_cin_3(self): self._test_conv2d(bs=4, cin=3, cout=2)
   def test_conv2d_bs_1_cin_1(self): self._test_conv2d(bs=1, cin=1)
-  @slow_test
+  @unittest.skip("redundant: cin=1 covered by test_conv2d_bs_1_cin_1")
   def test_conv2d_bs_4_cin_1(self): self._test_conv2d(bs=4, cin=1)
 
   def test_conv2d_errors(self):
@@ -2497,9 +2492,6 @@ class TestOps(unittest.TestCase):
             helper_test_op([(1,1,n,n), (1,1,k,k)],
               lambda x,w: torch.nn.functional.conv2d(torch.nn.functional.pad(x, p),w),
               lambda x,w: Tensor.conv2d(x,w,padding=p))
-            helper_test_op([(1,1,n,n), (1,1,k,k)],
-              lambda x,w: torch.nn.functional.conv2d(torch.nn.functional.pad(x, p),w),
-              lambda x,w: Tensor.conv2d(x,w,padding=p))
 
   def test_padded_conv2d_p21(self):
     bs,cin,H,W,padding = 4, 3, 3, 3, (2,1)
@@ -2548,7 +2540,7 @@ class TestOps(unittest.TestCase):
 
   @slow_test
   def test_max_pool2d(self):
-    for ksz in [(2,2), (3,3), 2, 3, (3,2), (5,5), (5,1)]:
+    for ksz in [2, (3,3), (3,2), (5,5), (5,1)]:
       with self.subTest(kernel_size=ksz):
         helper_test_op([(32,2,11,28)],
           lambda x: torch.nn.functional.max_pool2d(x, kernel_size=ksz),
@@ -2556,7 +2548,7 @@ class TestOps(unittest.TestCase):
 
   @slow_test
   def test_max_pool2d_padding(self):
-    for ksz in [(2,2), (3,3), 2, 3, (3,2)]:
+    for ksz in [(3,3), 2, (3,2)]:
       for p in [1, (1,0), (0,1)]:
         with self.subTest(kernel_size=ksz, padding=p):
           helper_test_op([(4,2,11,28)],
@@ -2619,7 +2611,7 @@ class TestOps(unittest.TestCase):
 
   def test_max_pool2d_ceil_mode(self):
     shape = (1,1,6,6)
-    for ksz in [(3,3), 3, (3,2), 4]:
+    for ksz in [(3,3), (3,2), 4]:
       with self.subTest(kernel_size=ksz):
         helper_test_op([shape],
           lambda x: torch.nn.functional.max_pool2d(x, kernel_size=ksz, padding=1, stride=3, ceil_mode=True),
@@ -2699,7 +2691,7 @@ class TestOps(unittest.TestCase):
   @slow_test
   def test_avg_pool2d(self):
     shape = (32,2,11,28)
-    for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]:
+    for ksz in [2, (3,3), (3,2), (5,5), (5,1)]:
       with self.subTest(kernel_size=ksz):
         helper_test_op([shape],
           lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz),
@@ -2713,7 +2705,7 @@ class TestOps(unittest.TestCase):
   @slow_test
   def test_avg_pool2d_padding(self):
     shape = (32,2,11,28)
-    for ksz in [(2,2), (3,3), 2, 3, (3,2)]:
+    for ksz in [2, (3,3), (3,2)]:
       for p in [1, (1,0), (0,1)]:
         with self.subTest(kernel_size=ksz, padding=p):
           helper_test_op([shape],
@@ -2735,7 +2727,7 @@ class TestOps(unittest.TestCase):
   @slow_test
   def test_avg_pool2d_padding_not_counted(self):
     shape = (32,2,11,28)
-    for ksz in [(2,2), (3,3), 2, 3, (3,2)]:
+    for ksz in [(3,3), 2, (3,2)]:
       with self.subTest(kernel_size=ksz):
         helper_test_op([shape],
           lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz, padding=1, count_include_pad=False),
@@ -2743,7 +2735,7 @@ class TestOps(unittest.TestCase):
 
   def test_avg_pool2d_ceil_mode(self):
     shape = (1,1,6,6)
-    for ksz in [(3,3), 3, (3,2), 4]:
+    for ksz in [(3,3), (3,2), 4]:
       with self.subTest(kernel_size=ksz):
         helper_test_op([shape],
           lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz, padding=1, stride=3, ceil_mode=True),
@@ -2751,7 +2743,7 @@ class TestOps(unittest.TestCase):
 
   def test_avg_pool2d_ceil_mode_padding_not_counted(self):
     shape = (1,1,6,6)
-    for ksz in [(3,3), 3, (3,2), 4]:
+    for ksz in [(3,3), (3,2), 4]:
       with self.subTest(kernel_size=ksz):
         helper_test_op([shape],
           lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz, padding=1, stride=3, ceil_mode=True, count_include_pad=False),
@@ -2953,7 +2945,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[...,c,:,e], lambda x: x[...,k,:,p])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_dim_collapse_int(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # dim collapse from int
@@ -2964,7 +2955,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[1,:,3:11:2,d,0:2], lambda x: x[1,:,3:11:2,o,0:2])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_dim_inject_none(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # dim injection from None
@@ -2999,7 +2989,6 @@ class TestOps(unittest.TestCase):
                             lambda x: x[Tensor([[0,1,-1],[-1,-2,0]]), Tensor([2,1,-1])])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_list_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[((0,),)])
@@ -3011,7 +3000,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,(2,1,0),c,(-2,1,0),e], lambda x: x[i,(2,1,0),k,(-2,1,0),p])
 
   @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
   def test_slice_fancy_indexing_tuple_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[(((0,),),)], lambda x: x[(((0,),),)])
@@ -3342,6 +3330,7 @@ class TestOps(unittest.TestCase):
   @unittest.skipIf((DEV.interface.startswith("MOCK") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu runtime issue")
   @unittest.skipIf(Device.DEFAULT == "QCOM", "QCOM fails with: Resource deadlock avoided")
+  @unittest.skipIf(Device.DEFAULT == "CPU" and DEV.renderer == "LVP", "extremely slow with LVP")
   def test_masked_select(self):
     helper_test_op([(32, 10)], lambda x: x.masked_select(x>0.5), lambda x: x.masked_select(x>0.5), forward_only=True)
     helper_test_op([(32, 10)], lambda x: x.masked_select(torch.tensor(True)), lambda x: x.masked_select(Tensor(True)), forward_only=True)

test/backend/test_optim.py

@@ -143,12 +143,11 @@ class TestOptim(unittest.TestCase):
 
   @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
   def test_mixed_precision(self):
-    old_default_float, dtypes.default_float = dtypes.default_float, dtypes.half
+    self.enterContext(Context(DEFAULT_FLOAT=dtypes.half))
     # weight update would overflow without upcasting
     self._test_sgd(10, {'lr': 1e10}, 1e-6, 3e-4)
     self._test_adam(1, {'lr': 1e10}, 1e-4, 1e-4)
     self._test_adamw(1, {'lr': 1e10}, 1e-4, 1e-4)
-    dtypes.default_float = old_default_float
 
   def test_assert_tensor_train(self):
     t = Tensor.ones((1,1))

test/backend/test_pickle.py

@@ -1,7 +1,7 @@
-import unittest, pickle, types
+import unittest, pickle, types, tracemalloc
 import numpy as np
-from tinygrad import Tensor, TinyJit, Variable, dtypes
-from tinygrad.helpers import GlobalCounters, ContextVar, Context
+from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
+from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
 from tinygrad.uop.ops import PatternMatcher, UPat, UOp
 
 class TestPickle(unittest.TestCase):
@@ -78,6 +78,22 @@ class TestPickle(unittest.TestCase):
     a2:UOp = pickle.loads(s)
     self.assertListEqual(a2.base.realized.as_memoryview().cast("I").tolist(), [0, 1, 2, 3])
 
+  @unittest.skipIf(DEV.interface.startswith("MOCK"), "mock device buffers live in host RAM, not VRAM")
+  def test_pickle_oob_ram(self):
+    N, M = 8, 10**6
+    ts = [Tensor.rand(M, dtype='float32').realize() for _ in range(N)]
+    tracemalloc.start()
+    st = pickle.dumps(ts, protocol=5, buffer_callback=lambda pb: pb.release())
+    self.assertLess(tracemalloc.get_traced_memory()[1], N*M*4)
+    tracemalloc.reset_peak()
+    def make_fake_buffers():
+      for _ in range(N):
+        Device[Device.DEFAULT].synchronize()
+        yield pickle.PickleBuffer(bytearray(M*4))
+    pickle.loads(st, buffers=make_fake_buffers())
+    self.assertLess(tracemalloc.get_traced_memory()[1], N*M*4)
+    tracemalloc.stop()
+
   def test_pickle_unrealized_tensor(self):
     t = Tensor.ones(10, 10)
     st = pickle.dumps(t)

test/backend/test_randomness.py

@@ -1,7 +1,7 @@
 import unittest, math
 
 from tinygrad import dtypes, Tensor, Device
-from tinygrad.helpers import getenv, DEV
+from tinygrad.helpers import getenv, DEV, Context
 from tinygrad.codegen import to_program
 
 from tinygrad.uop.ops import Ops
@@ -232,16 +232,14 @@ class TestRandomness(unittest.TestCase):
   @given(strat.sampled_from([dtypes.float, dtypes.float16, dtypes.bfloat16]))
   def test_randn_finite(self, default_float):
     if default_float not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
-    old_default_float = dtypes.default_float
     # low precision can result in inf from randn
-    dtypes.default_float = default_float
+    self.enterContext(Context(DEFAULT_FLOAT=default_float))
     t = Tensor.randn(64, 64)
     mx = t.max().numpy().item()
     mn = t.min().numpy().item()
     print(f"testing with {default_float=}")
     assert math.isfinite(mx), mx
     assert math.isfinite(mn), mn
-    dtypes.default_float = old_default_float
 
   def test_random_counter_overflow(self):
     device = Device.DEFAULT

test/backend/test_rangeify.py

@@ -1,7 +1,7 @@
 import unittest
-from tinygrad import Tensor, nn, Device
+from tinygrad import Tensor, nn, Device, dtypes, Variable
 from tinygrad.helpers import Context, GlobalCounters, getenv, PCONTIG, DEBUG
-from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops
+from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
 from tinygrad.codegen.opt import OptOps, Opt
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
@@ -66,6 +66,18 @@ class TestRangeifyAssign(unittest.TestCase):
     self.assertListEqual(lst2, B.permute(1, 0).tolist())
 
 class TestRangeifyEdgeCase(unittest.TestCase):
+  def test_variable_stack_data(self):
+    # a bound-Variable STACK used as data gets ranges from its graph position
+    v = Variable("v", 0, 10).bind(3)
+    t = Tensor(UOp.stack(v, v+1).cast(dtypes.int)) + Tensor.arange(2)
+    self.assertListEqual(t.tolist(), [3, 5])
+
+  def test_variable_data_and_shape(self):
+    # the same Variable has a data edge through CAST and a structural edge through SHRINK
+    v = Variable("shared_v", 1, 10).bind(3)
+    t = Tensor.ones(10)[:v] * Tensor(v.cast(dtypes.float))
+    self.assertEqual(t.sum().item(), 9)
+
   def test_matmul_relu_cat(self):
     a = Tensor.ones(100, 512).contiguous().realize()
     c = Tensor.ones(1, 512).contiguous().realize()

test/backend/test_schedule.py

@@ -363,9 +363,10 @@ class TestCopyFolding(unittest.TestCase):
 
   def test_one_hot_with_copy(self):
     y = Tensor([1, 2, 3]).to("CPU")
-    x = y.one_hot(10)
+    x = y.one_hot(10).int()
     check_schedule(x, 3, filter_sink=False)
 
+  @unittest.skip("no longer supported")
   def test_late_const_copy_folding(self):
     a = Tensor.arange(3).clone().realize()
     zeros = Tensor.zeros(3, buffer=False).realize()
@@ -406,6 +407,7 @@ class TestCopyFolding(unittest.TestCase):
     assert t.uop.is_realized, f"didn't realize Tensor {t}"
     self.assertListEqual(t.tolist(), [1.,1.,1.,1.])
 
+  @unittest.skip("same-device copies are no-ops")
   def test_self_assign_same_device_copy(self):
     a = Tensor.ones(4, 4).contiguous().realize()
     # use copy_to_device to bypass Tensor.to() shortcircuit and force a real same-device COPY in the graph

test/backend/test_tensor.py

@@ -464,6 +464,7 @@ class TestTinygrad(unittest.TestCase):
     data = data + [-x for x in data]
     np.testing.assert_allclose(Tensor(data, dtype=dtypes.int32).numpy(), np.array(data).astype(np.int32))
 
+  @unittest.skip("list elements are python scalars: numpy values in a list have no inferred dtype, use Tensor(np.array(data))")
   def test_tensor_list_ndarray(self):
     data = [np.array([1, 2, 3]), np.array([1, 2, 3]), np.array([1, 2, 3])]
     np.testing.assert_equal(Tensor(data).numpy(), np.array(data))

test/backend/test_tensor_variable.py

@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from tinygrad import Tensor, Variable, dtypes
+from tinygrad import Device, Tensor, Variable, TinyJit, dtypes
 from tinygrad.helpers import CHECK_OOB
 
 class TestTensorVariable(unittest.TestCase):
@@ -18,13 +18,24 @@ class TestTensorVariable(unittest.TestCase):
       self.assertListEqual((vv * t).tolist(), [2, 2, 2])
     except RuntimeError: pass
 
+  @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support")
   def test_large_range_variable(self):
-    vv = Variable("b", 0, 2**40).bind(2**35)
-    # TODO: pm_lower_index_dtype lowers ALU PARAM to int32 unconditionally
-    try:
-      self.assertEqual(Tensor(vv).item(), 2**35)
-    except AssertionError:
-      pass
+    self.assertEqual(Tensor(Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35)).clone(Device.DEFAULT).item(), 2**35)
+
+  @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support")
+  def test_large_range_variable_jit(self):
+    @TinyJit
+    def f(a,b): return (Tensor(a+b).clone(Device.DEFAULT) * 2).realize()
+    for i in range(3):
+      a = Variable("a", 0, 2**10, dtype=dtypes.int).bind(i)
+      b = Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35)
+      self.assertEqual(f(a,b).item(), (2**35 + i) * 2)
+
+  def test_variable_defers_like_a_literal(self):
+    vv = Variable("a", 1, 10).bind(2)
+    self.assertEqual(Tensor(vv).dtype, dtypes.weakint)
+    self.assertEqual((Tensor(vv) + Tensor([1], dtype=dtypes.int8)).dtype, dtypes.int8)  # takes the concrete side, no widening
+    self.assertEqual(Tensor(vv).item(), 2)                                              # a read commits at default_int
 
   def test_variable_tensor_dtype_arg(self):
     vv = Variable("a", 1, 10).bind(2)

test/backend/test_wait_loop.py

@@ -0,0 +1,127 @@
+import unittest, threading
+from tinygrad import Tensor, UOp
+from tinygrad.device import Device, Buffer, BufferSpec
+from tinygrad.dtype import AddrSpace, dtypes
+from tinygrad.engine.realize import run_linear
+from tinygrad.uop.ops import Ops, KernelInfo
+
+def wait_loop_kernel(C:UOp) -> UOp:
+  N = 10
+
+  # a RANGE with no src is a bound-less loop header: a jump target with no induction variable.
+  # the compare and conditional backedge are expanded by the renderers from the loop RANGE/END
+  l = UOp.loop(0)
+
+  i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
+
+  # i = 0
+  i = i.after(i[0].store(0))
+
+  # i + 1, read loop-carried through after(l)
+  inc = i.after(l)[0].load() + 1
+
+  # i = inc; END(store, l, cond): conditional backedge, loop again while inc < N (do-while)
+  # NOTE: the cond uses the computed value, not a reload of the register
+  st = i[0].store(inc)
+  i = i.after(st.end(l, inc < N))
+
+  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="wait_loop"))
+
+def nested_loop_kernel(C:UOp) -> UOp:
+  r = UOp.range(4, 0)
+  l = UOp.loop(1)
+
+  i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
+  i = i.after(i[0].store(0))
+
+  inc = i.after(l, r)[0].load() + 1
+  st = i[0].store(inc)
+
+  lend = st.end(l, inc < (r.cast(dtypes.int)+1)*3)
+  i = i.after(lend.end(r))
+
+  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=()))
+
+def wait_ext_kernel() -> UOp:
+  sig = UOp.param(0, dtypes.int, (1,), volatile=True)
+  l = UOp.loop(0)
+  v = sig.after(l)[0].load()
+  e = v.end(l, v < 1)
+  return e.sink(arg=KernelInfo(name="wait_ext"))
+
+def two_loops_kernel(C:UOp) -> UOp:
+  # two sequential loops on the same counter: ++ until 10, then ++ until 25
+  l1, l2 = UOp.loop(0), UOp.loop(1)
+
+  i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
+  i = i.after(i[0].store(0))
+
+  inc1 = i.after(l1)[0].load() + 1
+  i = i.after(i[0].store(inc1).end(l1, inc1 < 10))
+
+  inc2 = i.after(l2)[0].load() + 1
+  i = i.after(i[0].store(inc2).end(l2, inc2 < 25))
+
+  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="two_loops", opts_to_apply=()))
+
+def loop_in_loop_kernel(C:UOp) -> UOp:
+  # outer loop while i < 12, inner loop increments until i % 4 == 0 -> 12
+  l1, l2 = UOp.loop(0), UOp.loop(1)
+
+  i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
+  i = i.after(i[0].store(0))
+
+  inc = i.after(l1, l2)[0].load() + 1
+  st = i[0].store(inc)
+
+  # the outer END closes the inner END, and its cond reloads the register after the inner loop (in scope at the outer level)
+  e2 = st.end(l2, inc % 4 != 0)
+  oc = i.after(e2)[0].load()
+  i = i.after(e2.end(l1, oc < 12))
+
+  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="loop_in_loop", opts_to_apply=()))
+
+class TestWaitLoop(unittest.TestCase):
+  def test_wait_loop(self):
+    c = Tensor.empty(1, dtype=dtypes.int)
+    c = Tensor.custom_kernel(c, fxn=wait_loop_kernel)[0]
+    c.realize()
+    self.assertEqual(c.item(), 10)
+
+  def test_nested_loop_in_range(self):
+    c = Tensor.empty(1, dtype=dtypes.int)
+    c = Tensor.custom_kernel(c, fxn=nested_loop_kernel)[0]
+    c.realize()
+    self.assertEqual(c.item(), 12)
+
+  def test_two_sequential_loops(self):
+    c = Tensor.empty(1, dtype=dtypes.int)
+    c = Tensor.custom_kernel(c, fxn=two_loops_kernel)[0]
+    c.realize()
+    self.assertEqual(c.item(), 25)
+
+  def test_loop_in_loop(self):
+    c = Tensor.empty(1, dtype=dtypes.int)
+    c = Tensor.custom_kernel(c, fxn=loop_in_loop_kernel)[0]
+    c.realize()
+    self.assertEqual(c.item(), 12)
+
+@unittest.skipUnless(Device.DEFAULT in ("CPU", "AMD", "NV"), "need proper uncached=True handling")
+class TestVolatileLoops(unittest.TestCase):
+  def test_async_wait_ext(self):
+    sig_buf = Buffer(Device.DEFAULT, 1, dtypes.int, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+    try: sig_view = sig_buf.as_memoryview(force_zero_copy=True).cast('i')
+    except (AssertionError, NotImplementedError): self.skipTest(f"{Device.DEFAULT} does not support host-visible buffers")
+    sig_view[0] = 0
+
+    def set_signal():
+      threading.Event().wait(0.3)
+      sig_view[0] = 1
+
+    sync = threading.Thread(target=set_signal, daemon=True)
+    sync.start()
+    run_linear(UOp(Ops.LINEAR, src=(wait_ext_kernel().call(UOp.from_buffer(sig_buf)),)), wait=True)
+    sync.join(timeout=3)
+
+
+if __name__ == "__main__": unittest.main()

test/device/test_amd_llvm.py

@@ -3,7 +3,7 @@ from tinygrad import Device
 from tinygrad.device import CompileError
 if Device.DEFAULT == "AMD":
   # NOTE: if you don't gate this, LVP fails on Mac
-  from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler
+  from tinygrad.runtime.support.compiler_llvm import AMDLLVMCompiler
 
 @unittest.skipUnless(Device.DEFAULT == "AMD", "Runs only on AMD")
 class TestAMDLLVM(unittest.TestCase):

test/device/test_metal.py

@@ -1,7 +1,8 @@
 import unittest
-from tinygrad.device import CompileError, Device, BufferSpec
+from tinygrad.device import CompileError, Device, BufferSpec, TinyELF
+from tinygrad.helpers import Target
 if Device.DEFAULT=="METAL":
-  from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
+  from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler
 @unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
 class TestMetal(unittest.TestCase):
   def test_alloc_oom(self):
@@ -48,7 +49,7 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr
 """)
     with self.assertRaises(RuntimeError):
       compiled = compiled[:40] # corrupt the compiled program
-      MetalProgram(device, "r_5", compiled)
+      device.runtime(TinyELF(compiled, "r_5", Target("METAL"), ()))
 
   def test_free(self):
     size = 2**16

test/device/test_ocl.py

@@ -1,18 +1,20 @@
 import unittest
 from unittest.mock import patch
 from tinygrad import Device
-from tinygrad.device import Buffer
+from tinygrad.device import Buffer, TinyELF
 from tinygrad.dtype import dtypes
-from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler, CLProgram
+from tinygrad.helpers import Target
+from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler
 
 @unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
 class TestCLCompileCache(unittest.TestCase):
   def test_compile_cached(self):
     device = Device[Device.DEFAULT]
     src = "__kernel void cached_test(__global int* a) { a[0] = 1; }"
-    CLProgram(device, name="cached_test", lib=src.encode())
+    obj = TinyELF(src.encode(), "cached_test", Target("CL"), ())
+    device.runtime(obj)
     with patch.object(CLCompiler, 'compile', side_effect=RuntimeError("compile should not be called on cache hit")):
-      CLProgram(device, name="cached_test", lib=src.encode())
+      device.runtime(obj)
 
 @unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
 class TestCLError(unittest.TestCase):
@@ -27,7 +29,7 @@ class TestCLError(unittest.TestCase):
   def test_invalid_kernel_name(self):
     device = Device[Device.DEFAULT]
     with self.assertRaises(RuntimeError) as err:
-      CLProgram(device, name="", lib="__kernel void test(__global int* a) { a[0] = 1; }".encode())
+      device.runtime(TinyELF(b"__kernel void test(__global int* a) { a[0] = 1; }", "", Target("CL"), ()))
     assert str(err.exception) == "OpenCL Error -46: CL_INVALID_KERNEL_NAME"
 
   def test_unaligned_copy(self):

test/external/external_benchmark_bert_matmuls.py

@@ -1,7 +1,8 @@
 from tinygrad import Tensor, dtypes
-dtypes.default_float = dtypes.float16
 from tinygrad.dtype import to_dtype
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, Context
+
+Context(DEFAULT_FLOAT=dtypes.float16).__enter__()
 
 if __name__ == "__main__":
   # matmuls in bert layers

test/external/external_benchmark_bert_softmax.py

@@ -1,9 +1,10 @@
 from tinygrad import Tensor, dtypes, GlobalCounters
-dtypes.default_float = dtypes.float16
 from tinygrad.dtype import to_dtype
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, Context
 from test.backend.test_softmax_fusion import single_kernel_softmax
 
+Context(DEFAULT_FLOAT=dtypes.float16).__enter__()
+
 if __name__ == "__main__":
   # softmax in bert layers
   BS = getenv("BS", 96//6)

test/external/external_benchmark_op_conv.py

@@ -13,9 +13,9 @@ from tinygrad.dtype import Invalid
 
 def vision_conv_143():
   c0 = UOp.param(0, dtypes.half, shape=(16, 1024, 4))
-  c2 = UOp.range(32, 3, AxisType.LOOP)
-  c5 = UOp.range(128, 4, AxisType.LOOP)
-  c8 = UOp.range(16, 2, AxisType.LOOP)
+  c2 = UOp.range(32, 3)
+  c5 = UOp.range(128, 4)
+  c8 = UOp.range(16, 2)
   c16 = UOp.range(7, 0, AxisType.REDUCE)
   c17 = c8*2+c16
   c24 = ((c17<3)!=True)&(c17<35)
@@ -24,7 +24,7 @@ def vision_conv_143():
   c32 = ((c27<3)!=True)&(c27<67)
   c34 = UOp.param(1, dtypes.half, shape=(32, 1024, 4))
   c38 = c5//2
-  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.index, Invalid))
+  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
   c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4))
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
@@ -39,9 +39,9 @@ def vision_conv_143():
 
 def vision_conv_153():
   c0 = UOp.param(0, dtypes.half, shape=(8, 1024, 4))
-  c2 = UOp.range(16, 3, AxisType.LOOP)
-  c5 = UOp.range(256, 4, AxisType.LOOP)
-  c8 = UOp.range(8, 2, AxisType.LOOP)
+  c2 = UOp.range(16, 3)
+  c5 = UOp.range(256, 4)
+  c8 = UOp.range(8, 2)
   c16 = UOp.range(7, 0, AxisType.REDUCE)
   c17 = c8*2+c16
   c24 = ((c17<3)!=True)&(c17<19)
@@ -50,7 +50,7 @@ def vision_conv_153():
   c32 = ((c27<3)!=True)&(c27<35)
   c34 = UOp.param(1, dtypes.half, shape=(16, 1024, 4))
   c38 = c5//2
-  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.index, Invalid))
+  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
   c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4))
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
@@ -65,7 +65,7 @@ def vision_conv_153():
 
 def dm_conv_172():
   c0 = UOp.param(0, dtypes.half, shape=(1, 240, 4))
-  c2 = UOp.range(960, 4, AxisType.LOOP)
+  c2 = UOp.range(960, 4)
   c5 = UOp.param(1, dtypes.half, shape=(8, 384, 4))
   c7 = UOp.range(32, 0, AxisType.REDUCE)
   c10 = UOp.range(4, 1, AxisType.REDUCE)

test/external/external_test_am_fault_recovery.py

@@ -16,9 +16,9 @@ def _run(code:str, timeout:float=15.0) -> subprocess.CompletedProcess:
   return subprocess.run(["python", "-c", code], env={**os.environ, "AMD": "1"}, capture_output=True, text=True, timeout=timeout)
 
 def _run_asm(asm_src:str) -> subprocess.CompletedProcess:
-  return _run('from tinygrad.device import Device; from tinygrad.runtime.ops_amd import AMDProgram; '
+  return _run('from tinygrad.device import Device, TinyELF; from tinygrad.helpers import Target; '
               'from tinygrad.runtime.support.compiler_amd import HIPCompiler; dev = Device["AMD"]; '
-              f'AMDProgram(dev, "test", HIPCompiler(dev.arch).compile("""{asm_src}"""))('
+              f'dev.runtime(TinyELF(HIPCompiler(dev.arch).compile("""{asm_src}"""), "test", Target("AMD", arch=dev.arch), ()))('
               'dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)')
 
 def _verify_recovery() -> subprocess.CompletedProcess:

test/external/external_test_gpu_crash.py

@@ -5,7 +5,9 @@ These tests intentionally cause GPU faults to verify error handling.
 Run with: DEV=AMD python -m pytest test/external/external_test_gpu_crash.py -v
 """
 import unittest, re, importlib
-from tinygrad.device import Device
+from tinygrad import Device, UOp, dtypes
+from tinygrad.engine.realize import run_linear
+from tinygrad.uop.ops import Ops, KernelInfo
 from tinygrad.renderer.amd.dsl import s, v, Inst, NULL
 
 RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "s_waitcnt": "s_waitcnt", "s_endpgm": "s_endpgm",
@@ -13,20 +15,11 @@ RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "
                    "global_atomic_add_u32": "global_atomic_add", "flat_load_b32": "flat_load_dword",
                    "flat_store_b32": "flat_store_dword", "flat_atomic_add_u32": "flat_atomic_add", "s_load_b32": "s_load_dword"}
 
-def assemble(code:str, name:str="test", is_cdna:bool=False) -> str:
-  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
-  if is_cdna: kd["accum_offset"] = 8
-  else: kd["wavefront_size32"] = 1
-  return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
-         "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"
-
 @unittest.skipIf(Device.DEFAULT != "AMD", "AMD required")
 class TestGPUCrash(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
     cls.dev = Device["AMD"]
-    cls.compiler = HIPCompiler(cls.dev.arch)
     cls.is_cdna = cls.dev.target[0] < 10
     ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
     for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
@@ -41,14 +34,11 @@ class TestGPUCrash(unittest.TestCase):
     except Exception:
       self.fail("Device not working before test")
 
-  def _run(self, code: str):
-    from tinygrad.runtime.ops_amd import AMDProgram
-    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code, is_cdna=self.is_cdna)))
-    prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)
-
   def _run_insts(self, insts: list[Inst]):
-    from test.amd.disasm import disasm
-    self._run("\n".join(disasm(i) for i in insts))
+    buf = UOp.new_buffer("AMD", 64, dtypes.uint8)
+    sink = UOp.sink(UOp.param(0, dtypes.uint8, (64,), device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
+    prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=i) for i in insts))))
+    run_linear(UOp(Ops.LINEAR, src=(prg.call(buf),)), wait=True)
 
   def _assert_gpu_fault(self, func):
     """Assert that func raises a RuntimeError indicating a GPU fault (not a setup error)."""

test/external/external_test_jit_on_models.py

@@ -3,6 +3,7 @@ import unittest
 import numpy as np
 from tinygrad import Tensor, dtypes
 from tinygrad.engine.jit import TinyJit
+from tinygrad.helpers import Context
 from test.helpers import derandomize_model
 
 from examples.llama import Transformer
@@ -14,8 +15,7 @@ def helper_test_jitted_correctness(gen, train, train_jit):
 
 class TestJittedModels(unittest.TestCase):
   def test_jitted_tiny_llama(self):
-    old_float = dtypes.default_float
-    dtypes.default_float = dtypes.float16
+    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
 
     args_tiny = {"dim": 1024, "hidden_dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
     model = Transformer(**args_tiny)
@@ -25,7 +25,6 @@ class TestJittedModels(unittest.TestCase):
     @TinyJit
     def test_jit(t): return model(t, 0).realize()
     helper_test_jitted_correctness(lambda: (Tensor([[1,]]),), test, test_jit)
-    dtypes.default_float = old_float
 
   def test_jitted_stable_diffusion(self):
     from examples.stable_diffusion import UNetModel, unet_params

test/external/external_test_llm_opencode.py

@@ -0,0 +1,90 @@
+"""Real-model OpenCode regression.
+
+Run against an existing server:
+  RUN_LLM_OPENCODE_REGRESSION=1 LLM_BASE_URL=http://127.0.0.1:8000/v1 \
+    python -m pytest test/external/external_test_llm_opencode.py -v
+
+Or set LLM_GGUF and let the test start the tinygrad server.
+"""
+from __future__ import annotations
+
+import json, os, pathlib, re, shutil, socket, subprocess, sys, tempfile, time, unittest, urllib.request
+
+RUN_REGRESSION = os.getenv("RUN_LLM_OPENCODE_REGRESSION") == "1"
+
+def _server_ready(base_url:str) -> bool:
+  try:
+    with urllib.request.urlopen(base_url.rstrip("/") + "/models", timeout=1) as response: return response.status == 200
+  except OSError: return False
+
+@unittest.skipUnless(RUN_REGRESSION, "set RUN_LLM_OPENCODE_REGRESSION=1 to run the OpenCode regression")
+class TestLLMOpenCode(unittest.TestCase):
+  server:subprocess.Popen|None = None
+  server_log:tempfile._TemporaryFileWrapper|None = None
+
+  @classmethod
+  def setUpClass(cls):
+    if shutil.which("opencode") is None: raise unittest.SkipTest("opencode is not installed")
+    if base_url := os.getenv("LLM_BASE_URL"):
+      cls.base_url = base_url.rstrip("/")
+      if not cls.base_url.endswith("/v1"): cls.base_url += "/v1"
+      if not _server_ready(cls.base_url): raise RuntimeError(f"LLM server is not responding at {cls.base_url}")
+      return
+
+    model = pathlib.Path(os.environ["LLM_GGUF"])
+    with socket.socket() as sock:
+      sock.bind(("127.0.0.1", 0))
+      port = sock.getsockname()[1]
+    cls.base_url = f"http://127.0.0.1:{port}/v1"
+    cls.server_log = tempfile.NamedTemporaryFile(mode="w+", prefix="tinygrad-llm-")
+    cls.server = subprocess.Popen(
+      [sys.executable, "-m", "tinygrad.llm", "--model", str(model), "--serve", str(port), "--max_context", "262144"],
+      stdout=cls.server_log, stderr=subprocess.STDOUT, start_new_session=True)
+    deadline = time.monotonic() + 180
+    while time.monotonic() < deadline and cls.server.poll() is None:
+      if _server_ready(cls.base_url): return
+      time.sleep(0.25)
+    cls.server_log.seek(0)
+    raise RuntimeError(f"LLM server failed to start:\n{cls.server_log.read()[-8000:]}")
+
+  @classmethod
+  def tearDownClass(cls):
+    if cls.server is not None:
+      cls.server.terminate()
+      try: cls.server.wait(timeout=10)
+      except subprocess.TimeoutExpired:
+        cls.server.kill()
+        cls.server.wait(timeout=10)
+    if cls.server_log is not None: cls.server_log.close()
+
+  def run_opencode(self, prompt:str, cwd:pathlib.Path) -> str:
+    config = cwd / "opencode.json"
+    config.write_text(json.dumps({
+      "$schema": "https://opencode.ai/config.json", "permission": {"*": "allow"}, "formatter": False, "lsp": False,
+      "provider": {"regression": {"npm": "@ai-sdk/openai-compatible", "options": {"baseURL": self.base_url},
+                                  "models": {"tinygrad": {"name": "tinygrad"}}}},
+    }))
+    env = os.environ | {"OPENCODE_CONFIG": str(config)}
+    result = subprocess.run(
+      ["opencode", "run", "--pure", "--auto", "--dir", str(cwd), "-m", "regression/tinygrad", prompt],
+      cwd=cwd, env=env, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=120)
+    self.assertEqual(result.returncode, 0, result.stdout)
+    return re.sub(r"\x1b\[[0-?]*[ -/]*[@-~]", "", result.stdout)
+
+  def test_read_tool(self):
+    with tempfile.TemporaryDirectory() as directory:
+      cwd, marker = pathlib.Path(directory), "tinygrad-opencode-regression-7f3a91c2"
+      (cwd / "exact.txt").write_text(marker + "\n")
+      output = self.run_opencode("Read exact.txt with a tool and reply with its exact contents, with no other text.", cwd)
+      self.assertRegex(output, r"(?im)^\s*(?:→|>)\s*Read\s+exact\.txt\s*$")
+      self.assertIn(marker, output)
+
+  def test_shell_tool(self):
+    with tempfile.TemporaryDirectory() as directory:
+      cwd = pathlib.Path(directory)
+      output = self.run_opencode(
+        "Use the shell tool to run `printf tinygrad-shell-regression > shell-regression.txt`, then report completion.", cwd)
+      self.assertRegex(output, r"(?im)^\s*(?:\$|→|>)\s*.*printf\s+tinygrad-shell-regression")
+      self.assertEqual((cwd / "shell-regression.txt").read_text(), "tinygrad-shell-regression")
+
+if __name__ == "__main__": unittest.main()

test/external/fuzz_symbolic.py

@@ -40,7 +40,7 @@ def random_int_expr(depth=10):
 def random_bool_expr(depth=10, expr1=None):
   if depth == 0: return True
   if expr1 is None: expr1 = random_int_expr(depth-1)
-  expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(dtypes.index, random.randint(-10, 10))])
+  expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(dtypes.weakint, random.randint(-10, 10))])
   return random.choice(comp_ops)(expr1, expr2)
 
 
@@ -82,8 +82,8 @@ if __name__ == "__main__":
             f"v2=Variable(\"{u2.arg[0]}\", {u2.arg[1]}, {u2.arg[2]})\n" +\
             f"v3=Variable(\"{u3.arg[0]}\", {u3.arg[1]}, {u3.arg[2]})\n" +\
             f"expr = {expr}\n" +\
-            f"v1_val, v2_val, v3_val = UOp.const(dtypes.index, {n1.as_long()}), UOp.const(dtypes.index, {n2.as_long()})," +\
-                f"UOp.const(dtypes.index, {n3.as_long()})\n" +\
+            f"v1_val, v2_val, v3_val = UOp.const(dtypes.weakint, {n1.as_long()}), UOp.const(dtypes.weakint, {n2.as_long()})," +\
+                f"UOp.const(dtypes.weakint, {n3.as_long()})\n" +\
             "num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
             "rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
             "assert num==rn, f\"{num} != {rn}\"\n"

test/external/mlperf_stable_diffusion/external_test_eval.py

@@ -51,7 +51,7 @@ class TestEval(unittest.TestCase):
       vae = AutoencoderKL()
       for p in get_parameters(vae): p.to_(GPUS).realize()
       x = Tensor.zeros(48,4,64,64).contiguous().to(GPUS).realize()
-      x.uop = x.uop.multi(0)
+      x.uop = x.uop.unshard(0)
       for _ in range(2): vae_decode(x, vae)
 
 if __name__=="__main__":

test/helpers.py

@@ -15,7 +15,7 @@ from tinygrad.codegen.late.linearizer import linearize
 
 # decorator to skip slow tests by default, run with RUN_SLOW=1 to include them
 slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run")
-from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler
+from tinygrad.runtime.ops_python import PythonRenderer
 
 def full_rewrite(sink:UOp, ren:Renderer|None=None) -> UOp:
   if ren is None: ren = Renderer(Target())
@@ -83,14 +83,15 @@ def timeit(fxn:Callable[..., T], *args, **kwargs) -> tuple[T, float]:
   return ret, (time.perf_counter_ns()-st)*1e-6
 
 def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple[int, ...]=()):
-  allocator = Device['PYTHON'].allocator
+  dev = Device['PYTHON']
+  allocator = dev.allocator
   bufs = []
   for buf_dt, data in inputs or []:
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
   g = UOp.param(0, uop.dtype, (1,))
   prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
-  prog = PythonProgram("run", PythonCompiler().compile(prg.src[2].arg))
+  prog = dev.runtime(prg.to_elf())
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
   return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
 

test/mockgpu/amd/amdgpu.py

@@ -131,7 +131,8 @@ class PM4Executor(AMDQueue):
     _ = self._next_dword() # ev
 
     ptr = to_mv(self.gpu.translate_addr(addr_lo + (addr_hi << 32)), 8)
-    if mem_data_sel == 1 or mem_data_sel == 2: ptr.cast('Q')[0] = val
+    if mem_data_sel == 1: ptr.cast('I')[0] = val & 0xffffffff
+    elif mem_data_sel == 2: ptr.cast('Q')[0] = val
     elif mem_data_sel == 3:
       if mem_event_type == CACHE_FLUSH_AND_INV_TS_EVENT: ptr.cast('Q')[0] = int(time.perf_counter() * 1e8)
       else: raise RuntimeError(f"Unknown {mem_data_sel=} {mem_event_type=}")

test/mockgpu/amd/emu.py

@@ -51,7 +51,7 @@ class _MXCSRContext:
     if lib is None or not hasattr(self, '_saved'): return
     lib.set_fpcr(self._saved)
 
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.device import Buffer, BufferSpec, Device
 from tinygrad.runtime.autogen import hsa
@@ -446,7 +446,7 @@ class _Ctx:
     """Create a lane range UOp with unique axis ID."""
     if n is None: n = self.wave_size
     self._axis_id += 1
-    return UOp.range(n, self._axis_id, AxisType.LOOP, dtype=dtypes.int)
+    return UOp.range(n, self._axis_id, dtype=dtypes.int)
 
   def unroll_lanes(self, get_lane_bit, exec_mask: UOp, apply_exec: bool = True) -> UOp:
     """Combine lane bits into a mask using RANGE+REDUCE (32-bit for RDNA, 64-bit for CDNA)."""

test/mockgpu/amd/pcode.py

@@ -837,16 +837,16 @@ class Parser:
       idx = addr
       if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
         val = _u32(0).cast(dtypes.uint64)
-        for i in range(8): val = val | (mindex(idx + _const(dtypes.int, i)).load().cast(dtypes.uint64) << _u64(i * 8))
+        for i in range(8): val = val | (mindex(idx + _const(adt, i)).load().cast(dtypes.uint64) << _u64(i * 8))
       elif dt in (dtypes.uint8, dtypes.int8):
         val = mindex(idx).load().cast(dt)
       elif dt in (dtypes.uint16, dtypes.int16, dtypes.short):
         lo = mindex(idx).load().cast(dtypes.uint32)
-        hi = mindex(idx + _const(dtypes.int, 1)).load().cast(dtypes.uint32)
+        hi = mindex(idx + _const(adt, 1)).load().cast(dtypes.uint32)
         val = (lo | (hi << _u32(8))).cast(dt)
       else:
         val = _u32(0)
-        for i in range(4): val = val | (mindex(idx + _const(dtypes.int, i)).load().cast(dtypes.uint32) << _u32(i * 8))
+        for i in range(4): val = val | (mindex(idx + _const(adt, i)).load().cast(dtypes.uint32) << _u32(i * 8))
     else:
       idx = addr >> _const(addr.dtype, 2)
       val = mindex(idx)

test/mockgpu/nv/nvgpu.py

@@ -91,7 +91,7 @@ class GPFIFO:
     args_cnt, vals_cnt = const0[80], const0[81]
     args_addr = qmd.constant_buffer_addr_lower_0 + (qmd.constant_buffer_addr_upper_0 << 32) + 0x160
     args = to_mv(args_addr, args_cnt*8).cast('Q')
-    vals = to_mv(args_addr + args_cnt*8, vals_cnt*4).cast('I')
+    vals = to_mv(args_addr + args_cnt*8, vals_cnt*8).cast('Q')
     cargs = [ctypes.cast(args[i], ctypes.c_void_p) for i in range(args_cnt)] + [ctypes.cast(vals[i], ctypes.c_void_p) for i in range(vals_cnt)]
     gx, gy, gz = qmd.cta_raster_width, qmd.cta_raster_height, qmd.cta_raster_depth
     lx, ly, lz = qmd.cta_thread_dimension0, qmd.cta_thread_dimension1, qmd.cta_thread_dimension2

test/null/test_const_folding.py

@@ -40,7 +40,7 @@ class TestWeakConstFolding(unittest.TestCase):
     self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakint, 2**41))
 
   def test_float_unaries(self):
-    for dtype in (dtypes.weakint, dtypes.weakfloat):
+    for dtype in (dtypes.weakfloat,):
       for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
         out = UOp.const(dtype, 4).alu(op).simplify()
         self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))
@@ -50,7 +50,7 @@ class TestWeakConstFolding(unittest.TestCase):
     self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakfloat, 3.75))
 
   def test_invalid_poison(self):
-    self.assertIs(UOp.const(dtypes.weakint, Invalid).alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid)
+    self.assertIs(UOp.invalid().alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid)
 
 class TestBinaryOpsConstFolding(unittest.TestCase):
   def test_add_literal_zero(self):

test/null/test_device.py

@@ -69,11 +69,13 @@ class TestDevice(unittest.TestCase):
   @unittest.skipIf(WIN, "skipping windows test") # TODO: subprocess causes memory violation?
   def test_env_overwrite_default_compiler(self):
     if Device.DEFAULT == "CPU":
-      from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler
+      from tinygrad.runtime.support.compiler_cpu import ClangCompiler
+      from tinygrad.runtime.support.compiler_llvm import CPULLVMCompiler
       try: _, _ = CPULLVMCompiler(), ClangCompiler()
       except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
 
-      imports = "from tinygrad import Device; from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler"
+      imports = ("from tinygrad import Device; from tinygrad.runtime.support.compiler_cpu import ClangCompiler; "
+                 "from tinygrad.runtime.support.compiler_llvm import CPULLVMCompiler")
       subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, CPULLVMCompiler)"'],
                         shell=True, check=True, env={**os.environ, "DEV": "CPU:LLVM"})
       subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangCompiler)"'],
@@ -81,11 +83,13 @@ class TestDevice(unittest.TestCase):
       subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, ClangCompiler)"'],
                         shell=True, check=True, env={**os.environ, "DEV": "CPU:CLANG"})
     elif Device.DEFAULT == "AMD":
-      from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
+      from tinygrad.runtime.support.compiler_amd import HIPCompiler
+      from tinygrad.runtime.support.compiler_llvm import AMDLLVMCompiler
       try: _, _ = HIPCompiler(Device[Device.DEFAULT].arch), AMDLLVMCompiler(Device[Device.DEFAULT].arch)
       except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
 
-      imports = "from tinygrad import Device; from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler"
+      imports = ("from tinygrad import Device; from tinygrad.runtime.support.compiler_amd import HIPCompiler; "
+                 "from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler")
       subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, AMDLLVMCompiler)"'],
                         shell=True, check=True, env={**os.environ, "DEV": "AMD:LLVM"})
       subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, HIPCompiler)"'],
@@ -96,7 +100,8 @@ class TestDevice(unittest.TestCase):
 
   @unittest.skipIf(WIN, "skipping windows test")
   def test_env_online(self):
-    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangCompiler
+    from tinygrad.runtime.support.compiler_cpu import ClangCompiler
+    from tinygrad.runtime.support.compiler_llvm import CPULLVMCompiler
     try: _, _ = CPULLVMCompiler(), ClangCompiler()
     except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
 
@@ -111,7 +116,7 @@ class TestDevice(unittest.TestCase):
 
   @unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
   def test_compiler_autodetect_fallback(self):
-    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
+    from tinygrad.runtime.support.compiler_llvm import CPULLVMCompiler
 
     try: CPULLVMCompiler()
     except Exception as e: self.skipTest(f"skipping: LLVM not available: {e}")

test/null/test_dtype_spec.py

@@ -2,7 +2,7 @@ import unittest, math, struct, operator
 from tinygrad import Tensor, Device
 from tinygrad.dtype import DTYPES_DICT, dtypes, Invalid, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
 
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, Context
 from hypothesis import given, settings, strategies as strat
 import numpy as np
 import torch
@@ -58,8 +58,8 @@ class TestHelpers(unittest.TestCase):
   def test_from_py(self):
     assert dtypes.from_py(True) == dtypes.bool
     assert dtypes.from_py(Invalid) == dtypes.bool
-    assert dtypes.from_py(2) == dtypes.default_int
-    assert dtypes.from_py(3.0) == dtypes.default_float
+    assert dtypes.from_py(2) == dtypes.weakint
+    assert dtypes.from_py(3.0) == dtypes.weakfloat
     assert dtypes.from_py([]) == dtypes.default_float
     assert dtypes.from_py(()) == dtypes.default_float
     assert dtypes.from_py([True]) == dtypes.bool
@@ -205,7 +205,8 @@ class TestTypePromotion(unittest.TestCase):
     assert least_upper_dtype(dtypes.uint16, dtypes.int32) == dtypes.int32
     assert least_upper_dtype(dtypes.int32, dtypes.uint32) == dtypes.int64
     assert least_upper_dtype(dtypes.uint32, dtypes.int64) == dtypes.int64
-    assert least_upper_dtype(dtypes.int64, dtypes.uint64) == dtypes.uint64
+    # uint64 has no common integer supertype with any signed int (JAX JEP), they all defer up to weakfloat
+    for st in dtypes.sints: assert least_upper_dtype(st, dtypes.uint64) == dtypes.weakfloat
     assert least_upper_dtype(dtypes.float16, dtypes.float32) == dtypes.float32
     assert least_upper_dtype(dtypes.float32, dtypes.float64) == dtypes.float64
 
@@ -224,30 +225,17 @@ class TestTypePromotion(unittest.TestCase):
     assert least_upper_dtype(dtypes.fp8e5m2, dtypes.uint64) == dtypes.fp8e5m2
 
   def test_weakint_promo(self):
-    # weakint with itself is weakint
     assert least_upper_dtype(dtypes.weakint, dtypes.weakint) == dtypes.weakint
-    # weakint is above bool
-    assert least_upper_dtype(dtypes.weakint, dtypes.bool) == dtypes.weakint
-    # weakint defers to any concrete int type
+    assert least_upper_dtype(dtypes.bool, dtypes.weakint) == dtypes.weakint
     assert least_upper_dtype(dtypes.weakint, dtypes.int8) == dtypes.int8
-    assert least_upper_dtype(dtypes.weakint, dtypes.uint8) == dtypes.uint8
-    assert least_upper_dtype(dtypes.weakint, dtypes.int16) == dtypes.int16
-    assert least_upper_dtype(dtypes.weakint, dtypes.int32) == dtypes.int32
-    assert least_upper_dtype(dtypes.weakint, dtypes.int64) == dtypes.int64
-    assert least_upper_dtype(dtypes.weakint, dtypes.uint64) == dtypes.uint64
-    # weakint defers to any float type
-    assert least_upper_dtype(dtypes.weakint, dtypes.float16) == dtypes.float16
-    assert least_upper_dtype(dtypes.weakint, dtypes.float32) == dtypes.float32
-    assert least_upper_dtype(dtypes.weakint, dtypes.float64) == dtypes.float64
 
   def test_weakfloat_promo(self):
-    # weakfloat is a float, but like weakint it is not one of dtypes.floats
+    # weakfloat is a float, but is not one of dtypes.floats
     assert dtypes.is_float(dtypes.weakfloat) and dtypes.weakfloat not in dtypes.floats
     # weakfloat with itself is weakfloat
     assert least_upper_dtype(dtypes.weakfloat, dtypes.weakfloat) == dtypes.weakfloat
-    # weakfloat is above bool, weakint and any concrete int (they defer up to it)
+    # weakfloat is above bool and any concrete int (they defer up to it)
     assert least_upper_dtype(dtypes.weakfloat, dtypes.bool) == dtypes.weakfloat
-    assert least_upper_dtype(dtypes.weakfloat, dtypes.weakint) == dtypes.weakfloat
     assert least_upper_dtype(dtypes.weakfloat, dtypes.int32) == dtypes.weakfloat
     assert least_upper_dtype(dtypes.weakfloat, dtypes.uint64) == dtypes.weakfloat
     # weakfloat defers to any concrete float type
@@ -257,19 +245,14 @@ class TestTypePromotion(unittest.TestCase):
     assert least_upper_dtype(dtypes.weakfloat, dtypes.float64) == dtypes.float64
 
 class TestTypeSpec(unittest.TestCase):
-  def setUp(self):
-    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
-  def tearDown(self):
-    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
-
   def test_set_dtype_default(self):
     for default_int in [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64]:
-      dtypes.default_int = default_int
-      assert dtypes.default_int == default_int
+      with Context(DEFAULT_INT=default_int):
+        assert dtypes.default_int == default_int
 
     for default_float in [*dtypes.fp8s, dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
-      dtypes.default_float = default_float
-      assert dtypes.default_float == default_float
+      with Context(DEFAULT_FLOAT=default_float):
+        assert dtypes.default_float == default_float
 
   @given(strat.sampled_from(core_dtypes), strat.sampled_from([operator.gt, operator.ge, operator.le, operator.lt, operator.eq, operator.ne]))
   def test_bool_ops(self, dtype, op):
@@ -277,7 +260,7 @@ class TestTypeSpec(unittest.TestCase):
 
   @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
   def test_functions_return_index(self, dtype, default_int, default_float):
-    dtypes.default_int, dtypes.default_float = default_int, default_float
+    self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
     assert Tensor([0, 1], dtype=dtype).argmax().dtype == dtypes.int32
     assert Tensor([0, 1], dtype=dtype).argmin().dtype == dtypes.int32
     assert Tensor([0, 1], dtype=dtype).multinomial().dtype == dtypes.int32
@@ -297,7 +280,7 @@ class TestTypeSpec(unittest.TestCase):
 
   @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
   def test_attention_returns_same_dtype(self, data_dtype, default_float):
-    dtypes.default_float = default_float
+    self.enterContext(Context(DEFAULT_FLOAT=default_float))
     query = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
     key = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
     value = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
@@ -308,30 +291,38 @@ class TestTypeSpec(unittest.TestCase):
     assert query.scaled_dot_product_attention(key, value, attn_mask=mask).dtype == data_dtype
 
 class TestAutoCastType(unittest.TestCase):
-  def setUp(self):
-    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
-  def tearDown(self):
-    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
-
   @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
   def test_least_upper_float_input_is_float(self, input_dtype, default_float):
-    dtypes.default_float = default_float
+    self.enterContext(Context(DEFAULT_FLOAT=default_float))
     self.assertEqual(least_upper_float(input_dtype), input_dtype)
 
   @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
   def test_least_upper_float_input_is_int(self, input_dtype, default_float):
-    dtypes.default_float = default_float
+    self.enterContext(Context(DEFAULT_FLOAT=default_float))
     self.assertEqual(least_upper_float(input_dtype), default_float)
 
   @given(strat.sampled_from(core_dtypes))
   def test_broadcast_scalar(self, dt):
-    assert (Tensor.ones(4, 4, dtype=dt) + 2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
-    assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert (Tensor.ones(4, 4, dtype=dt) + 2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
+    assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
     assert (Tensor.ones(4, 4, dtype=dt) + True).dtype == dt
 
+  @given(strat.sampled_from(core_dtypes))
+  def test_pad_scalar(self, dt):
+    t = Tensor.ones(4, dtype=dt)
+    assert t.pad(((1, 1),), value=2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
+    assert t.pad(((1, 1),), value=2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
+    assert t.pad(((1, 1),), value=True).dtype == dt
+
+  @given(strat.sampled_from(core_dtypes))
+  def test_sort(self, dt):
+    # sort pads with dtype.min/max, a scalar of its own dtype
+    assert Tensor([3, 1, 2], dtype=dt).sort()[0].dtype == dt
+    assert Tensor([3, 1, 2], dtype=dt).sort(descending=True)[0].dtype == dt
+
   @given(strat.sampled_from(dtype_floats))
   def test_int_div_int(self, default_float):
-    dtypes.default_float = default_float
+    self.enterContext(Context(DEFAULT_FLOAT=default_float))
     self.assertEqual(Tensor([1]).div(Tensor([2])).dtype, default_float)
 
   def test_sum(self):
@@ -419,16 +410,19 @@ class TestAutoCastType(unittest.TestCase):
   @given(strat.sampled_from(core_dtypes))
   def test_where_one_scalar(self, dt):
     t = Tensor(2, dtype=dt)
-    self.check_where_alternate_input_other(t, 3.2, (dt if dtypes.is_float(dt) else dtypes.default_float))
-    self.check_where_alternate_input_other(t, 3, (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int))
+    self.check_where_alternate_input_other(t, 3.2, (dt if dtypes.is_float(dt) else dtypes.weakfloat))
+    self.check_where_alternate_input_other(t, 3, (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint))
     self.check_where_alternate_input_other(t, True, dt)
 
   def test_where_two_scalars(self):
-    self.check_where_alternate_input_other(3.1, 3.2, dtypes.default_float)
-    self.check_where_alternate_input_other(3.1, 3, dtypes.default_float)
-    self.check_where_alternate_input_other(3.1, True, dtypes.default_float)
-    self.check_where_alternate_input_other(3, 2, dtypes.default_int)
-    self.check_where_alternate_input_other(3, True, dtypes.default_int)
+    self.check_where_alternate_input_other(3.1, 3.2, dtypes.weakfloat)
+    self.check_where_alternate_input_other(3.1, 3, dtypes.weakfloat)
+    self.check_where_alternate_input_other(3.1, True, dtypes.weakfloat)
+    self.check_where_alternate_input_other(3, 2, dtypes.weakint)
+    self.check_where_alternate_input_other(3, True, dtypes.weakint)
+
+  def test_where_non_bool_cond_raises(self):
+    with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0)
     self.check_where_alternate_input_other(False, True, dtypes.bool)
 
   @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
@@ -437,8 +431,8 @@ class TestAutoCastType(unittest.TestCase):
 
   @given(strat.sampled_from(core_dtypes))
   def test_maximum_const(self, dt):
-    assert Tensor([1, 2], dtype=dt).maximum(3.1).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
-    assert Tensor([1, 2], dtype=dt).maximum(3).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert Tensor([1, 2], dtype=dt).maximum(3.1).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
+    assert Tensor([1, 2], dtype=dt).maximum(3).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
     assert Tensor([1, 2], dtype=dt).maximum(True).dtype == dt
 
   def test_div(self):
@@ -449,7 +443,7 @@ class TestAutoCastType(unittest.TestCase):
 
   def test_div_const(self):
     assert (Tensor([1, 2], dtype=dtypes.int32) / 2).dtype == dtypes.default_float
-    assert (Tensor([1, 2], dtype=dtypes.int32) / 2.0).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.int32) / 2.0).dtype == dtypes.weakfloat
     assert (Tensor([1, 2], dtype=dtypes.float16) / 2).dtype == dtypes.float16
     assert (Tensor([1, 2], dtype=dtypes.float16) / 2.0).dtype == dtypes.float16
 

test/null/test_gpudims.py

@@ -24,7 +24,7 @@ class TestGroupedDims(unittest.TestCase):
     total = math.prod(dims)
     specials = sorted(dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs])), key=lambda u: u.arg)
     # build flat index and primed flat (same expression with renamed SPECIALs)
-    flat = UOp.const(dtypes.index, 0)
+    flat = UOp.const(dtypes.weakint, 0)
     for i, idx in enumerate(idxs):
       flat = flat + idx * int(math.prod(dims[i+1:]))
     flat_p = flat.substitute({s: UOp(Ops.SPECIAL, src=s.src, arg=s.arg+"_p") for s in specials})

test/null/test_graph_rewrite.py

@@ -107,21 +107,21 @@ class TestFoldingAndReduction(unittest.TestCase):
 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.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.arg, 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.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].arg, 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.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.arg, 1)
@@ -140,7 +140,7 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
   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.index) % 10).render(simplify=False), x_var_uop.render(simplify=False))
+    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)

test/null/test_hcq_iface.py

@@ -1,9 +1,26 @@
 import unittest, array, time
 from tinygrad.helpers import mv_address
 from tinygrad.runtime.support.hcq import MMIOInterface
+from tinygrad.runtime.support.memory import VirtMapping
+from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.runtime.support.usb import USBMMIOInterface
 from test.mockgpu.usb import MockUSB
 
+class TestPCIIface(unittest.TestCase):
+  def test_sysmem_mapping_respects_uncached(self):
+    class MM:
+      def alloc_vaddr(self, size, align): return 0x10000
+      def map_range(self, vaddr, size, paddrs, aspace, uncached=False, snooped=False):
+        return VirtMapping(vaddr, size, paddrs, aspace, uncached, snooped)
+    class PCI:
+      def bar_info(self, bar): return 0, 256 << 20
+      def alloc_sysmem(self, size, **kwargs): return memoryview(bytearray(size)), [0x20000]
+    iface = PCIIfaceBase.__new__(PCIIfaceBase)
+    iface.dev, iface.vram_bar, iface.pci_dev = None, 0, PCI()
+    iface.dev_impl = type("DevImpl", (), {"mm": MM()})()
+    for uncached in (False, True):
+      with self.subTest(uncached=uncached): self.assertEqual(iface.alloc(4096, host=True, uncached=uncached).meta.mapping.uncached, uncached)
+
 class TestHCQIface(unittest.TestCase):
   def setUp(self):
     self.size = 4 << 10

test/null/test_linearizer_failures.py

@@ -8,12 +8,12 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailures(unittest.TestCase):
   def test_fail_1(self):
     c0 = UOp.param(0, dtypes.float, (64,))
-    c1 = UOp.range(UOp.const(dtypes.index, 2), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.index, 32), 2, AxisType.LOOP)
-    c3 = ((c1*UOp.const(dtypes.index, 32))+c2)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.WEAK)
+    c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2)
     c4 = UOp.param(1, dtypes.float, (163840,))
-    c5 = UOp.range(UOp.const(dtypes.index, 2560), 0, AxisType.REDUCE)
-    c6 = c4.index(((((((c5//UOp.const(dtypes.index, 8))%UOp.const(dtypes.index, 8))*UOp.const(dtypes.index, 8))+(c5%UOp.const(dtypes.index, 8)))+(((c2*UOp.const(dtypes.index, 40))+(c5//UOp.const(dtypes.index, 64)))*UOp.const(dtypes.index, 64)))+(c1*UOp.const(dtypes.index, 81920))))
+    c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE)
+    c6 = c4.index(((((((c5//UOp.const(dtypes.weakint, 8))%UOp.const(dtypes.weakint, 8))*UOp.const(dtypes.weakint, 8))+(c5%UOp.const(dtypes.weakint, 8)))+(((c2*UOp.const(dtypes.weakint, 40))+(c5//UOp.const(dtypes.weakint, 64)))*UOp.const(dtypes.weakint, 64)))+(c1*UOp.const(dtypes.weakint, 81920))))
     c7 = UOp.param(2, dtypes.float, (64,))
     c8 = c7.index(c3)
     c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()

test/null/test_llm_server.py

@@ -1,5 +1,5 @@
 import unittest, threading, time, json
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 class TestLLMServer(unittest.TestCase):
   """Integration tests using the real OpenAI client."""
@@ -66,20 +66,18 @@ class TestLLMServer(unittest.TestCase):
       self.assertEqual(chunk.model, "test-model")
 
   def test_stream_with_usage(self):
-    stream = self.client.chat.completions.create(
-      model="test",
-      messages=[{"role": "user", "content": "Hello"}],
-      stream=True,
-      stream_options={"include_usage": True}
-    )
-
-    chunks = list(stream)
+    def generate(ids, **kwargs):
+      for token in (300, 301, 999):
+        ids.append(token)
+        yield token
+    with patch.object(self.mock_model, "generate", side_effect=generate):
+      chunks = list(self.client.chat.completions.create(
+        model="test", messages=[{"role": "user", "content": "Hello"}], stream=True, stream_options={"include_usage": True}))
     last_chunk = chunks[-1]
 
-    self.assertIsNotNone(last_chunk.usage)
-    self.assertIsNotNone(last_chunk.usage.prompt_tokens)
-    self.assertIsNotNone(last_chunk.usage.completion_tokens)
-    self.assertIsNotNone(last_chunk.usage.total_tokens)
+    self.assertEqual(last_chunk.usage.prompt_tokens, 3)
+    self.assertEqual(last_chunk.usage.completion_tokens, 2)
+    self.assertEqual(last_chunk.usage.total_tokens, 5)
 
   def test_multi_turn_conversation(self):
     stream = self.client.chat.completions.create(
@@ -97,6 +95,17 @@ class TestLLMServer(unittest.TestCase):
     self.assertGreater(len(chunks), 0)
     self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
 
+  def test_chat_template_kwargs(self):
+    import jinja2
+    with patch.object(self.server, "template", jinja2.Template(
+        "{% for m in messages %}{% if m.role == 'assistant' and preserve_thinking %}<think>{{ m.reasoning_content }}</think>"
+        "{% endif %}{{ m.content }}|{% endfor %}")):
+      list(self.client.chat.completions.create(model="test", messages=[
+        {"role":"user", "content":"first"}, {"role":"assistant", "content":"answer", "reasoning_content":"reason"},
+        {"role":"user", "content":"next"}], stream=True,
+        extra_body={"chat_template_kwargs":{"preserve_thinking":True, "messages":[]}}))
+    self.mock_tok.encode.assert_called_with("first|<think>reason</think>answer|next|")
+
   def test_content_is_streamed(self):
     stream = self.client.chat.completions.create(
       model="test",
@@ -111,6 +120,28 @@ class TestLLMServer(unittest.TestCase):
 
     self.assertGreater(len(contents), 0)
 
+  def test_interrupted_stream_logs_tokens(self):
+    with patch.object(self.mock_model, "generate", side_effect=lambda ids, **kwargs: iter([300, 301, 999])), \
+         patch("tinygrad.llm.serve.stderr_log") as log, patch("tinygrad.llm.serve.colored", side_effect=lambda text, color: text) as color:
+      stream = self.server.RequestHandlerClass.run_model(Mock(server=self.server), [200, 201, 202], "test")
+      next(stream)
+      next(stream)
+      stream.close()
+    interrupt = log.call_args.args[0]
+    self.assertFalse(interrupt.startswith("\n"))
+    self.assertTrue(interrupt.endswith("\n"))
+    self.assertIn("gen:", interrupt)
+    self.assertIn("out:    1", interrupt)
+    self.assertTrue(any(args[0].startswith("total:") and args[1] == "red" for args, _ in color.call_args_list))
+
+  def test_stream_disconnect_closes_source(self):
+    from tinygrad.viz.serve import HTTPRequestHandler
+    source, handler = Mock(), Mock()
+    source.__iter__ = Mock(return_value=iter([{}]))
+    handler.wfile.write.side_effect = BrokenPipeError
+    HTTPRequestHandler.stream_json(handler, source)
+    source.close.assert_called_once()
+
   def test_non_streaming(self):
     resp = self.client.chat.completions.create(
       model="test-model",

test/null/test_multitensor.py

@@ -1,5 +1,5 @@
 import gc, unittest
-from tinygrad import Tensor, GlobalCounters, dtypes
+from tinygrad import Tensor, UOp, GlobalCounters, dtypes
 from tinygrad.engine.jit import TinyJit
 from tinygrad.helpers import Context
 
@@ -181,7 +181,7 @@ class TestMultiScalarALU(unittest.TestCase):
     @functools.cache
     def _fxn(x_p, device):
       t = Tensor(x_p, device=device)
-      inner = Tensor(t.uop.src[0]) if t.uop.op is Ops.MULTI else t
+      inner = Tensor(t.uop.src[0]) if t.uop.op is Ops.UNSHARD else t
       return (inner.sum(),)
     param = x.as_param(0)
     fxn = _fxn(param.uop, x.device)
@@ -218,5 +218,10 @@ class TestMultiAxis(unittest.TestCase):
     self.assertEqual(e.uop.axis, 0)
     self.assertTrue(e.uop.has_buffer_identity())
 
+  def test_symbolic_reshape_shard_axis(self):
+    rows = UOp.variable("rows", 1, 4).bind(3)
+    x = Tensor.empty(4, 2).shard(("NULL:1", "NULL:2"), axis=1)[:rows]
+    self.assertEqual(x.reshape(rows, 1, 2).uop.axis, 2)
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_real_world.py

@@ -47,12 +47,8 @@ class TestRealWorld(unittest.TestCase):
     gc.collect()
     global global_mem_used
     global_mem_used = GlobalCounters.mem_used
-    self.old_float = dtypes.default_float
     np.random.seed(2002)
 
-  def tearDown(self):
-    dtypes.default_float = self.old_float
-
   @slow
   @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
   def test_stable_diffusion(self):
@@ -81,7 +77,7 @@ class TestRealWorld(unittest.TestCase):
 
   @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
   def test_llama(self):
-    dtypes.default_float = dtypes.float16
+    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
 
     args_tiny = {"dim": 1024, "hidden_dim": 2048, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
     model = LLaMaTransformer(**args_tiny)
@@ -93,7 +89,7 @@ class TestRealWorld(unittest.TestCase):
 
   @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
   def test_gpt2(self):
-    dtypes.default_float = dtypes.float16
+    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
 
     args_tiny = {"dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-5, "vocab_size": 1000}
     model = GPT2Transformer(**args_tiny)
@@ -150,7 +146,7 @@ class TestRealWorld(unittest.TestCase):
 
   @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
   def test_train_cifar_hyp(self):
-    dtypes.default_float = dtypes.float16
+    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
     with Context(TRAINING=1):
       model = SpeedyResNet(Tensor.ones((12,3,2,2)))
       optimizer = optim.SGD(get_parameters(model), lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay'])

test/null/test_schedule.py

@@ -603,14 +603,12 @@ class TestSchedule(unittest.TestCase):
     check_schedule(p, 4)
 
   def test_conv2d(self, allowed=4, dtype=dtypes.float):
-    old_default_float, dtypes.default_float = dtypes.default_float, dtype
-    dtypes.default_float = dtype
+    self.enterContext(Context(DEFAULT_FLOAT=dtype))
     Tensor.manual_seed(0)
     BS, CIN = 2, 3
     img = Tensor.randn(BS, CIN, 64, 64).realize()
     w = Tensor.uniform(16, CIN, 3, 3).realize()
     ret = Tensor.conv2d(img, w).relu().mean().backward()
-    dtypes.default_float = old_default_float
     linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad)
     cnt = len([call for call in linear.src if call.src[0].op is Ops.SINK])
     assert cnt == allowed, f"expected {allowed} kernels, got {cnt}"

test/null/test_schedule_cache.py

@@ -1,10 +1,12 @@
-import unittest
-from tinygrad import Tensor, Variable, Context
+import os, unittest
+from unittest.mock import patch
+from tinygrad import Tensor, Variable, Context, UOp
+from tinygrad.callify import transform_to_call
 from tinygrad.helpers import cpu_events
-from tinygrad.schedule import schedule_cache
+from tinygrad.schedule import lower_sink_to_linear, schedule_cache
 
 def schedule_one():
-  Tensor([1]).schedule_linear()
+  (Tensor.empty(1) + 1).schedule_linear()
 
 class TestScheduleCache(unittest.TestCase):
   def test_bound_variable_var_vals(self):
@@ -37,5 +39,22 @@ class TestScheduleCache(unittest.TestCase):
       num_events_cache = len(cpu_events)
     self.assertLess(num_events_cache, num_events_no_cache)
 
+  def test_disk_schedule_cache(self):
+    function = transform_to_call(UOp.sink((Tensor.empty(1) + 1).uop))[0].src[0]
+    schedule_cache.clear()
+    with patch.dict(os.environ, {"DISK_SCACHE":"1"}), \
+         patch("tinygrad.schedule.diskcache_get", return_value=None), \
+         patch("tinygrad.schedule.diskcache_put") as cache_put:
+      lower_sink_to_linear(function)
+      cached = cache_put.call_args.args[2]
+
+    schedule_cache.clear()
+    with patch.dict(os.environ, {"DISK_SCACHE":"1"}), \
+         patch("tinygrad.schedule.diskcache_get", return_value=cached) as cache_get, \
+         patch("tinygrad.schedule.diskcache_put") as cache_put:
+      self.assertIs(lower_sink_to_linear(function), cached)
+      cache_get.assert_called_once()
+      cache_put.assert_not_called()
+
 if __name__ == "__main__":
   unittest.main()

test/null/test_simplify_valid_idx.py

@@ -1,6 +1,6 @@
 import unittest, itertools
 
-from tinygrad.codegen.late.coalese import indexing_simplify
+from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
@@ -23,7 +23,7 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO
     UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)),
   ))
 
-def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(dtypes.index, nmax),), arg=expr)
+def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(dtypes.weakint, nmax),), arg=expr)
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
 def Range(n, nmax): return UOp.range(nmax, n)
 
@@ -409,7 +409,7 @@ class TestImageSimplification(unittest.TestCase):
     alu1 = ((idx2*1536)+(ridx4*768)+ridx3+(idx1*24)+(ridx5*3)+-771)//768
     valid = (((idx2+ridx4)<1)!=1)&(((idx1+ridx5)<1)!=1)
     load = get_load_image_uop((128, 768, 4), valid, (alu0, alu1))
-    self.check(load, None, "((((idx1*24)+r3)+(r5*3))+-3)", "(((idx2*2)+r4)+-1)")
+    self.check(load, None, "((((idx1*24)+(r5*3))+r3)+-3)", "(((idx2*2)+r4)+-1)")
 
   def test_simplify7(self):
     # DEBUG=2 ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916 # noqa: E501
@@ -455,7 +455,7 @@ class TestImageSimplification(unittest.TestCase):
     A1 = lidx0*32 + r0*32 + lidx1*4 - 99
     valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 19)
     alu0 = gidx0 + (A1 % 32)*32 + (A1 // 32 % 16)*1024
-    load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(dtypes.index, 0)))
+    load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
     try:
       self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0")
     except AssertionError:
@@ -474,7 +474,7 @@ class TestImageSimplification(unittest.TestCase):
     A1 = lidx0*16 + r0*16 + lidx1*4 - 51
     valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 11)
     alu0 = lidx2 + gidx0*4 + (A1 % 16)*64 + (A1 // 16 % 8)*1024
-    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(dtypes.index, 0)))
+    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
     try:
       self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0")
     except AssertionError:
@@ -488,18 +488,18 @@ class TestImageSimplification(unittest.TestCase):
     gidx0 = Special("gidx0", 1064)
     r12 = Range(12, 3)
     valid = ((gidx0 < 645).ne(True)) & (gidx0 < 653)
-    idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(dtypes.index, 0))
+    idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(dtypes.weakint, 0))
     load = get_load_image_uop((1, 48, 4), valid, idx)
     self.check(load, None, "(r12*4+(gidx0+3)%4+(gidx0+3)//4*24+-3888)", "0")
 
 class TestDropTrueGate(unittest.TestCase):
   def test_drop_true_gate_on_index(self):
     # test that INDEX with a constant True valid gets simplified to drop the valid
-    from tinygrad.codegen.late.coalese import indexing_simplify
+    from tinygrad.codegen.late.coalesce import indexing_simplify
     from tinygrad.uop.ops import graph_rewrite
     from tinygrad.uop.symbolic import sym
     buf = UOp.param(0, dtypes.int, (1,))
-    idx = UOp.const(dtypes.index, 0)
+    idx = UOp.const(dtypes.weakint, 0)
     true_gate = UOp.const(dtypes.bool, True)
     index_with_gate = UOp(Ops.INDEX, src=(buf, idx.valid(true_gate)))
     # apply the optimization
@@ -516,7 +516,7 @@ class TestRangeShrink(unittest.TestCase):
   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(dtypes.index, 4), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 4)
@@ -524,8 +524,8 @@ class TestRangeShrink(unittest.TestCase):
   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
     r = Range(0, 204)
-    load1 = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
-    load2 = get_gated_load_uop(r < UOp.const(dtypes.index, 8), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
+    load2 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 8), r)
     ranges = self.get_ranges(UOp.sink(load1, load2))
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 8)
@@ -533,7 +533,7 @@ class TestRangeShrink(unittest.TestCase):
   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(dtypes.index, 300), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 300), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 204)
@@ -541,7 +541,7 @@ class TestRangeShrink(unittest.TestCase):
   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(dtypes.index, 4), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),))
     ranges = self.get_ranges(UOp.sink(load1, load2))
     self.assertEqual(len(ranges), 1)
@@ -550,7 +550,7 @@ class TestRangeShrink(unittest.TestCase):
   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
     r = Range(0, 204)
-    gated_load = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
+    gated_load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
     ranges = self.get_ranges(red.sink())
     self.assertEqual(len(ranges), 1)
@@ -559,7 +559,7 @@ class TestRangeShrink(unittest.TestCase):
   def test_range_shrink_to_single_iteration(self):
     # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
     r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(dtypes.index, 1), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 1), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 0)
 
@@ -568,7 +568,7 @@ class TestRangeShrink(unittest.TestCase):
     from tinygrad.dtype import Invalid
     r = Range(0, 204)
     x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
-    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, 0)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 4)
 
@@ -577,7 +577,7 @@ class TestRangeShrink(unittest.TestCase):
     from tinygrad.dtype import Invalid
     r = Range(0, 204)
     x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
-    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(0, x)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 4)
 

test/null/test_tensor.py

@@ -181,7 +181,7 @@ class TestTensorPad(unittest.TestCase):
     t = Tensor.arange(9).reshape(1, 1, 3, 3)
     self.assertEqual(t.dtype, dtypes.int)
     r = t.pad((1, 2, 0, -1), value=-float('inf'))
-    self.assertEqual(r.dtype, dtypes.float)
+    self.assertEqual(r.dtype, dtypes.weakfloat)
     self.assertEqual(r.shape, (1, 1, 2, 6))
 
 class TestTensorDeviceMismatch(unittest.TestCase):

test/null/test_tensor_uop_mixin.py

@@ -382,7 +382,7 @@ class TestTensorUOpStack(unittest.TestCase):
     self.assertIs(_t(2, 3).uop.stack(w.uop).dtype, dtypes.float32)
   def test_stack_index_dtype(self):
     # index is outside the promotion lattice, equal dtypes bypass promotion
-    self.assertEqual(UOp.const(dtypes.index, 1).stack(UOp.const(dtypes.index, 2)).shape, (2,))
+    self.assertEqual(UOp.const(dtypes.weakint, 1).stack(UOp.const(dtypes.weakint, 2)).shape, (2,))
 
 class TestTensorUOpConv2d(unittest.TestCase):
   def test_conv2d_basic(self):

test/null/test_uop_graph.py

@@ -2,7 +2,7 @@ import unittest, pytest
 from tinygrad import dtypes, Variable
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import DEBUG, Context
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes
 from tinygrad.uop.symbolic import sym
 from test.helpers import to_uops_list
 
@@ -202,7 +202,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_same_fold(self):
     v = UOp.variable('tmp', 0, 1)
-    c0 = UOp.const(dtypes.index, 0)
+    c0 = UOp.const(dtypes.weakint, 0)
     vc = v != c0
     c1 = UOp.const(dtypes.float, 1.0)
     out = vc.where(c1, c1)
@@ -424,16 +424,16 @@ class TestUOpGraph(unittest.TestCase):
     # mnist indexing with split reduceop
     # Make sure we are not doign math on the loaded index, which would promote it to long
     c0 = UOp.param(0, dtypes.uchar, (128000,))
-    c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
     c3 = UOp.param(1, dtypes.int, (512,))
     c4 = c3.index(c1)
-    c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
     c7 = UOp.param(2, dtypes.uchar, (60000,))
     c8 = c7.index(c6)
     c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
-    c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
     uops = to_uops_list([c10])
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
@@ -441,19 +441,19 @@ class TestUOpGraph(unittest.TestCase):
   def test_load_idx_no_math_on_loaded(self):
     # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
     c0 = UOp.param(0, dtypes.uchar, (128000,))
-    c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
     c3 = UOp.param(1, dtypes.int, (512,))
     c4 = c3.index(c1)  # c4 is a load
-    c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
     c7 = UOp.param(2, dtypes.uchar, (60000,))
     c8 = c7.index(c6)
     # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
-    loaded_idx = c4.cast(dtypes.index)
-    comparison = (loaded_idx + c5) < UOp.const(dtypes.index, 60000)
+    loaded_idx = c4.cast(dtypes.weakint)
+    comparison = (loaded_idx + c5) < UOp.const(dtypes.weakint, 60000)
     c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
-    c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
     uops = to_uops_list([c10])
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
@@ -707,5 +707,16 @@ class TestUOpBroadcast(unittest.TestCase):
     c = a + b
     self.assertEqual(c.op, Ops.ADD)
 
+  def test_broadcast_axes(self):
+    t = Variable("t", 1, 10)
+    self.assertEqual(broadcast_axes((4, 8), (4, 8)), ())
+    self.assertEqual(broadcast_axes((8,), (4, 8)), (0,))
+    self.assertEqual(broadcast_axes((), (4, 8)), (0, 1))
+    self.assertEqual(broadcast_axes((3, 1), (4, 3, 8)), (0, 2))
+    self.assertEqual(broadcast_axes((1, 8), (1, 8)), ())
+    self.assertEqual(broadcast_axes((t, 8), (t, 8)), ())
+    self.assertEqual(broadcast_axes((1, 8), (t, 8)), (0,))
+    with self.assertRaises(RuntimeError): broadcast_axes((4, 8), (8,))
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/null/test_uop_symbolic.py

@@ -5,19 +5,20 @@ import z3
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from test.helpers import get_uops
 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.validate import uops_to_z3
 
 def check_uop_against_string(self, v:UOp, s:str):
   sym_vars = {v.render():v for v in v.toposort() if v.op in (Ops.RANGE, Ops.SPECIAL, Ops.PARAM)}
   s_eval = eval(s, sym_vars)
-  if isinstance(s_eval, int) and v.dtype==dtypes.index: s_eval = UOp.const(dtypes.index, s_eval)
+  if isinstance(s_eval, int) and v.dtype==dtypes.weakint: s_eval = UOp.const(dtypes.weakint, s_eval)
   elif isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(dtypes.from_py(s_eval), s_eval)
   s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval")
   self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}")
 
-def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.index): return UOp.variable(name,min_val,max_val,dtype)
-def uconst(val): return UOp.const(dtypes.index, val)
+def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint): return UOp.variable(name,min_val,max_val,dtype)
+def uconst(val): return UOp.const(dtypes.weakint, val)
 def usum(ops): return functools.reduce(lambda x,y: x+y, ops)
 def uand(ops): return functools.reduce(lambda x,y: x*y, ops)
 
@@ -247,12 +248,12 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (0, 9))
 
   def test_range_div_its_symbolic_bound(self):
-    a = Variable("a", 1, 10, dtypes.index)
+    a = Variable("a", 1, 10, dtypes.weakint)
     ridx0 = UOp.range(a+2, 0)
     self.helper_test_variable(ridx0//(a+2), 0, 0, "0")
 
   def test_range_mod_its_symbolic_bound(self):
-    a = Variable("a", 1, 10, dtypes.index)
+    a = Variable("a", 1, 10, dtypes.weakint)
     ridx = UOp.range(a+2, 0)
     self.helper_test_variable(ridx%(a+2), 0, 11, "r0")
 
@@ -333,7 +334,7 @@ class TestSymbolic(unittest.TestCase):
   def test_mod_mod_wrong_sign(self):
     v1=Variable("v1", 0, 128)
     v3=Variable("v3", 0, 7)
-    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), 0, 4, "((v3+v1%2*2+-3)%5)")
+    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), 0, 4, "((v3+v1%2*2+2)%5)")
 
   def test_mod_mod_wrong_sign2(self):
     v2=Variable("v2", 0, 8)
@@ -346,6 +347,12 @@ class TestSymbolic(unittest.TestCase):
   def test_mul_lt(self):
     self.helper_test_variable(Variable("a", 0, 5)*4 < 13, 0, 1, "(a<4)")
     self.helper_test_variable(Variable("a", 0, 5)*4 < 16, 0, 1, "(a<4)")
+    self.helper_test_variable(Variable("a", -5, 5)*4 < -13, 0, 1, "(a<-3)")
+    self.helper_test_variable(Variable("a", -5, 5)*-4 < 13, 0, 1, "((a*-1)<4)")
+    c0, c1 = 2, 2**54+1
+    self.helper_test_variable(Variable("a", 0, c1)*c0 < c1, 0, 1, f"(a<{2**53+1})")
+    c0, c1 = -2, -(2**54-1)
+    self.helper_test_variable(Variable("a", 0, -c1)*c0 < c1, 0, 1, f"((a*-1)<{-(2**53-1)})")
     self.helper_test_variable(Variable("a", 0, 5)*(-2) < 0, 0, 1, "((a*-1)<0)")
     self.helper_test_variable(Variable("a", 0, 5)*4 >= 12, 0, 1, "((a<3)!=True)")
     self.helper_test_variable(Variable("a", 0, 5)*4 >= 13, 0, 1, "((a<4)!=True)")
@@ -365,7 +372,7 @@ class TestSymbolic(unittest.TestCase):
 
   def test_div_const_div_wrong_sign_divisor(self):
     a = Variable("a", 0, 124)
-    self.helper_test_variable(((a+10)//-2+10)//-4, -2, 14, "(((a+10)//-2+10)//-4)")
+    self.helper_test_variable(((a+10)//-2+10)//-4, -2, 14, "((a//-2+-3)//-4+-2)")
 
   def test_nested_div_negative_divisor(self):
     # (x//c1)//c2 -> x//(c1*c2) only when c2>0
@@ -437,7 +444,7 @@ class TestSymbolic(unittest.TestCase):
 
   def test_masked_shr_fold(self):
     x = UOp.variable('x', 0, 255, dtype=dtypes.uint32)
-    self.helper_test_variable((x & -4) >> 2, 0, 63, "(x>>2)", test_z3=False)
+    self.helper_test_variable((x & -4) >> 2, 0, 63, "(x>>2)")
 
   def test_bool_or_not_tautology(self):
     a = Variable("a", 0, 10)
@@ -450,8 +457,15 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(c & c.logical_not(), False, False, "False")
 
   def test_mod_factor_negative(self):
-    self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 10), Variable("b", 0, 10)*28]) % 28, 0, 27, "((a+b*28+-29)%28)")
-    self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 100), Variable("b", 0, 10)*28]) % 28, 0, 27, "((a+b*28+-29)%28)")
+    self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 10), Variable("b", 0, 10)*28]) % 28, 0, 27, "((a+27)%28)")
+    self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 100), Variable("b", 0, 10)*28]) % 28, 0, 27, "((a+27)%28)")
+
+  def test_mod_const_reduction_negative_offset(self):
+    # (x+c)%d -> (x+c%d)%d holds for any sign of x+c and d
+    x = Variable("x", 0, 100)
+    self.helper_test_variable((x-50)%3, 0, 2, "((x+1)%3)")
+    self.helper_test_variable((x-50)%-3, -2, 0, "((x+-2)%-3)")
+    self.helper_test_variable((x+7)%-13, -12, 0, "((x+-6)%-13)")
 
   def test_sum_combine_num(self):
     self.helper_test_variable(usum([uconst(29), Variable("a", 0, 10), uconst(-23)]), 6, 16, "(a+6)")
@@ -579,9 +593,9 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(x%12//4*4 + x%4 + x//12*12, 0, 23, "x")
 
   def test_div_neg_cancel(self):
-    self.helper_test_variable((-Variable("idx", 0, 100)+199)//-4 + 50, 0, 25, "((idx*-1+199)//-4+50)")
-    self.helper_test_variable((-Variable("idx", 0, 100)+200)//-4 + 50, 0, 25, "((idx*-1+200)//-4+50)")
-    self.helper_test_variable((-Variable("idx", 0, 100)+201)//-4 + 50, -1, 24, "((idx*-1+201)//-4+50)")
+    self.helper_test_variable((-Variable("idx", 0, 100)+199)//-4 + 50, 0, 25, "((idx*-1+-1)//-4)")
+    self.helper_test_variable((-Variable("idx", 0, 100)+200)//-4 + 50, 0, 25, "(idx*-1//-4)")
+    self.helper_test_variable((-Variable("idx", 0, 100)+201)//-4 + 50, -1, 24, "((idx*-1+-3)//-4+-1)")
     self.helper_test_variable((-Variable("idx", 0, 100))//2, -50, 0, "(idx*-1//2)")
     self.helper_test_variable(Variable("idx", 0, 100)//-2, -50, 0, "(idx//-2)")
 
@@ -658,20 +672,20 @@ class TestSymbolic(unittest.TestCase):
   def test_div_neg_all_range(self):
     gidx = Variable("gidx", 0, 124)
     lidx = Variable("lidx", 0, 7)
-    self.helper_test_variable((-gidx*8-lidx+999)//-4 + 250, 0, 250, "((gidx*-8+lidx*-1+999)//-4+250)")
-    self.helper_test_variable((-gidx*8-lidx+1000)//-4 + 250, 0, 249, "((gidx*-8+lidx*-1+1000)//-4+250)")
-    self.helper_test_variable((-gidx*8-lidx+1001)//-4 + 250, -1, 249, "((gidx*-8+lidx*-1+1001)//-4+250)")
-    self.helper_test_variable((-gidx*8-lidx+1002)//-4 + 250, -1, 249, "((gidx*-8+lidx*-1+1002)//-4+250)")
+    self.helper_test_variable((-gidx*8-lidx+999)//-4 + 250, 0, 250, "((lidx*-1+gidx*-8+-1)//-4)")
+    self.helper_test_variable((-gidx*8-lidx+1000)//-4 + 250, 0, 249, "((lidx*-1+gidx*-8)//-4)")
+    self.helper_test_variable((-gidx*8-lidx+1001)//-4 + 250, -1, 249, "((lidx*-1+gidx*-8+-3)//-4+-1)")
+    self.helper_test_variable((-gidx*8-lidx+1002)//-4 + 250, -1, 249, "((lidx*-1+gidx*-8+-2)//-4+-1)")
 
   def test_div_neg_then_neg(self):
     # taken from arange opts
     lidx0 = Variable("lidx0", 0, 7)
     lidx1 = Variable("lidx1", 0, 7)
     alu2 = -lidx0-lidx1
-    self.helper_test_variable((((alu2+14)//(-32))+4), 3, 4, "((lidx0*-1+lidx1*-1+14)//-32+4)")
-    self.helper_test_variable(-(((alu2+14)//(-32))+4), -4, -3, "((lidx0*-1+lidx1*-1+14)//-32*-1+-4)")
-    self.helper_test_variable((((alu2+134)//(-32))+4), -1, 0, "((lidx0*-1+lidx1*-1+134)//-32+4)")
-    self.helper_test_variable((((alu2+142)//(-32))+4), -1, 0, "((lidx0*-1+lidx1*-1+142)//-32+4)")
+    self.helper_test_variable((((alu2+14)//(-32))+4), 3, 4, "((lidx0*-1+lidx1*-1+-18)//-32+3)")
+    self.helper_test_variable(-(((alu2+14)//(-32))+4), -4, -3, "((lidx0*-1+lidx1*-1+-18)//-32*-1+-3)")
+    self.helper_test_variable((((alu2+134)//(-32))+4), -1, 0, "((lidx0*-1+lidx1*-1+-26)//-32+-1)")
+    self.helper_test_variable((((alu2+142)//(-32))+4), -1, 0, "((lidx0*-1+lidx1*-1+-18)//-32+-1)")
     self.helper_test_variable((((alu2+150)//(-32))+4), -1, -1, "-1")
     self.helper_test_variable((((alu2+158)//(-32))+4), -1, -1, "-1")
 
@@ -837,12 +851,12 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(x%(-3) + ((x//(-3))%5)*(-3), -14, 0, "x%-15")
 
   def test_div_mod_recombine_shifted_quotient(self):
-    # when vmin<0 blocks const reduction on the mod side, the quotient is stored const-shifted: (x-50)//3 -> (x+1)//3 - 17.
+    # const reduction stores mod/div const-shifted: (x-50)%3 -> (x+1)%3, (x-50)//3 -> (x+1)//3 - 17.
     # recombine only needs a quotient of some b congruent to base mod div, so the shift folds into the result
     x = Variable("x", 0, 100)
     y = Variable("y", 0, 99)
     self.helper_test_variable((x-50)%3 + ((x-50)//3)*3, -50, 50, "(x+-50)")                # shifted literal quotient
-    self.helper_test_variable((x-50)%3 + (((x-50)//3)%5)*3, 0, 14, "((x+-50)%15)")         # shift inside the partial's mod
+    self.helper_test_variable((x-50)%3 + (((x-50)//3)%5)*3, 0, 14, "((x+10)%15)")         # shift inside the partial's mod
     self.helper_test_variable(((y-50)//5)%4 + ((y-50)//20)*4, -10, 9, "(y//5+-10)")        # merged and shifted
 
   def test_div_mod_recombine_in_additive_sum(self):
@@ -872,12 +886,16 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "(idx<4).where((idx//4), -1)")
 
   def test_floordiv_lt(self):
-    # x//d<c <=> x<c*d for d>0
+    # x//d<c <=> x<c*d for d>0, and <=> c*d<x for d<0
     idx = Variable("idx", 0, 24)
     self.helper_test_variable((idx//4<3), 0, 1, "(idx<12)")
     self.helper_test_variable(((idx-20)//4<-3), 0, 1, "(idx<8)")
     self.helper_test_variable(((idx-10)//4<0), 0, 1, "(idx<10)")
-    self.helper_test_variable((idx//-4<-3), 0, 1, "((idx//-4)<-3)")
+    self.helper_test_variable((idx//-4<-3), 0, 1, "(12<idx)")
+    self.helper_test_variable((idx//-4<-5), 0, 1, "(20<idx)")
+    self.helper_test_variable((idx//-4<-6), 0, 0, "False")
+    self.helper_test_variable(((idx-10)//-4<0), 0, 1, "(8<(idx+-2))")
+    self.helper_test_variable(((idx-20)//-4<2), 0, 1, "(12<idx)")
 
   def test_nested_div_mod_negative_inner_divisor(self):
     # (x % (k*c)) // c -> (x // c) % k requires k>0; (x % (k*c)) % c -> x % c is unconditional for c>0
@@ -919,8 +937,8 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond.cast(dtypes.int).ne(2), 1, 1, "True")
     self.helper_test_variable(cond.cast(dtypes.int).ne(-1), 1, 1, "True")
     # CAST(bool -> index) folds too
-    self.helper_test_variable(cond.cast(dtypes.index).ne(0), 0, 1, "(a<2)")
-    self.helper_test_variable(cond.cast(dtypes.index).ne(1), 0, 1, "((a<2)!=True)")
+    self.helper_test_variable(cond.cast(dtypes.weakint).ne(0), 0, 1, "(a<2)")
+    self.helper_test_variable(cond.cast(dtypes.weakint).ne(1), 0, 1, "((a<2)!=True)")
 
   def test_where_removal(self):
     cond = Variable("a", 0, 3) < 2
@@ -956,6 +974,28 @@ class TestSymbolic(unittest.TestCase):
     # not combining  # TODO: can combine if one is identity element const
     self.helper_test_variable(aa+ab, 0, 6, "((x<2).where(a, b)+(x<2).where(a, 0))")
 
+  def test_where_combine_cross_zero(self):
+    cond = Variable("x", 0, 3) < 2
+    a = Variable("a", 0, 3)
+    b = Variable("b", 0, 3)
+    self.helper_test_variable(cond.where(a, a.ufix(0)) + cond.where(b.ufix(0), b), 0, 3, "(x<2).where(a, b)")
+    self.helper_test_variable(cond.where(a, a.ufix(0)) + cond.where(a.ufix(0), a), 0, 3, "a")
+
+  def test_where_or_dual(self):
+    m1 = Variable("x", 0, 3) < 2
+    m2 = Variable("y", 0, 3) < 2
+    a = Variable("a", 0, 3)
+    b = Variable("b", 0, 3)
+    self.helper_test_variable(m1.where(a, m2.where(a, b)), 0, 3, "((x<2)|(y<2)).where(a, b)")
+
+  def test_bool_ne_false(self):
+    cond = Variable("x", 0, 3) < 2
+    self.helper_test_variable(cond.ne(False), 0, 1, "(x<2)")
+
+  def test_bitcast_chain(self):
+    a = UOp.variable("a", 0, 3, dtype=dtypes.int32)
+    self.assertIs(graph_rewrite(a.bitcast(dtypes.float32).bitcast(a.dtype), sym), a)
+
   def test_negation_in_where(self):
     cond = Variable("x", 0, 3) < 2
     a = Variable("a", 0, 3)
@@ -1001,7 +1041,6 @@ class TestSymbolic(unittest.TestCase):
     # (a if ((s<5)&(s<6)) else b) -> (a if (s<5) else b)
     self.helper_test_variable(expr, 0, 3, "(s<5).where(a, b)")
 
-  @unittest.expectedFailure
   def test_where_closure_folding(self):
     # cond.where(t, f) where f contains cond.where(a, b) should fold the inner where to b in false branch
     x = Variable("x", 0, 10)
@@ -1011,6 +1050,41 @@ class TestSymbolic(unittest.TestCase):
     # the inner where should be folded: true branch gets -x, false branch gets x
     self.helper_test_variable(outer, -20, 11, "(x<5).where((x*-2), (x+1))")
 
+  def test_where_closure_folding_deep(self):
+    x = Variable("x", 0, 10)
+    cond = x < 5
+    w1 = cond.where(-x, x)
+    w2 = cond.where(w1*2, w1+1)
+    self.helper_test_variable(cond.where(w2*3, w2+7), -60, 18, "(x<5).where((x*-6), (x+8))")
+
+  def test_where_closure_folding_different_cond(self):
+    # a nested where on a different condition is not folded
+    x = Variable("x", 0, 10)
+    a = Variable("a", 0, 3)
+    b = Variable("b", 0, 3)
+    expr = (x<5).where((x<7).where(a, b), (x<7).where(b, a))
+    self.helper_test_variable(expr, 0, 3, "(x<5).where((x<7).where(a, b), (x<7).where(b, a))")
+
+  def test_where_closure_folding_derived_cond(self):
+    # cond is a value inside the branch: (!cond).where(a, b) is b in the true branch
+    x = Variable("x", 0, 10)
+    a = Variable("a", 0, 3)
+    b = Variable("b", 0, 3)
+    c = Variable("c", 0, 3)
+    expr = (x<5).where((x<5).logical_not().where(a, b)*2, c)
+    self.helper_test_variable(expr, 0, 6, "(x<5).where((b*2), c)")
+
+  def test_where_closure_folding_valid(self):
+    # a valid gate on the same cond folds in the true branch, the live else value is kept
+    x = Variable("x", 0, 10)
+    a = Variable("a", 0, 3)
+    cond = x < 5
+    expr = cond.where(a.valid(cond), Variable("c", 0, 3))
+    self.assertIs(graph_rewrite(expr, sym), cond.where(a, Variable("c", 0, 3)))
+    # a same-cond valid gate in the false branch is Invalid there
+    expr = cond.where(Variable("t", 0, 3), a.valid(cond))
+    self.assertIs(graph_rewrite(expr, sym), cond.where(Variable("t", 0, 3), UOp.invalid()))
+
   def test_symbolic_div(self):
     # from symbolic arange
     a = Variable("a", 1, 10)
@@ -1021,7 +1095,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((numerator//denominator)<=0, 1, 1, "True")
 
   def test_symbolic_range_doesnt_collapse(self):
-    r0 = UOp.range((Variable("a", 1, 10)<5).cast(dtypes.index), 0)
+    r0 = UOp.range((Variable("a", 1, 10)<5).cast(dtypes.weakint), 0)
     self.helper_test_variable(r0, 0, 0, "r0")
 
   def test_const_reciprocal(self):
@@ -1043,8 +1117,8 @@ class TestSymbolic(unittest.TestCase):
   def test_nested_mod_negative_range(self):
     # (x%(k*c))%c = x%c for positive c
     x = Variable("x", 0, 1575)
-    self.helper_test_variable(((x + (-1064)) % 512) % 4, 0, 3, "((x+-1064)%4)")
-    self.helper_test_variable(((x + (-1064)) % 512) % 128, 0, 127, "((x+-1064)%128)")
+    self.helper_test_variable(((x + (-1064)) % 512) % 4, 0, 3, "(x%4)")
+    self.helper_test_variable(((x + (-1064)) % 512) % 128, 0, 127, "((x+88)%128)")
 
 class TestSymbolicNumeric(unittest.TestCase):
   def helper_test_numeric(self, f):
@@ -1254,23 +1328,22 @@ class TestSymbolicSymbolicOps(unittest.TestCase):
 """
 
 class TestInvalidIndex(unittest.TestCase):
+  def test_invalid_lift_keeps_live_else(self):
+    ridx = Variable("ridx", 0, 10)
+    cond = ridx < 5
+    expr = cond.where(cond.where(ridx, UOp.invalid()), ridx+100)
+    self.assertIs(expr.simplify(), cond.where(ridx, ridx+100))
+
   def test_invalid_times_0(self):
     ridx = Variable("ridx", 0, 10)
     idx = (ridx<5).where(ridx, UOp.invalid())*0
-    self.assertIs(idx.simplify(), (ridx<5).where(0, UOp.invalid()), "multiplying an index by 0 should preserve the invalid")
-
-  def test_invalid_comparison_drops_invalid(self):
-    # comparisons return a bool, and bools can't be invalid
-    ridx = Variable("ridx", 0, 10)
-    idx = (ridx<5).where(ridx, UOp.invalid())<3
-    self.assertIs(idx.simplify(), (ridx<3), "comparison of index should drop the invalid")
-    self.assertIs(idx.where(UOp.const(dtypes.int, 1), 0).simplify(), (ridx<3).where(UOp.const(dtypes.int, 1), 0),
-      "comparison of index should drop the invalid")
+    self.assertIs(idx.simplify(), (ridx<5).where(UOp.const(dtypes.weakint, 0), UOp.invalid()),
+                  "multiplying an index by 0 should preserve the invalid")
 
   def test_alu_moves_inside_invalid(self):
     ridx = Variable("ridx", 0, 10)
-    idx = (ridx<5).where(ridx, UOp.invalid())*10
-    self.assertIs(idx.simplify(), (ridx<5).where(ridx*10, UOp.invalid()), "multiplying an index by 0 should preserve the invalid")
+    self.assertIs((10*(ridx<5).where(ridx, UOp.invalid())).simplify(), (ridx<5).where(ridx*10, UOp.invalid()),
+      "Invalid should poison either binary operand position")
 
   def test_merge_invalid_conditions(self):
     ridx0 = Variable("ridx0", 0, 10)
@@ -1289,16 +1362,16 @@ class TestInvalidIndex(unittest.TestCase):
     self.assertIs((UOp.invalid()<Variable("a",0,10)).simplify().dtype, dtypes.bool)
 
   def test_alu_invalid_vconst(self):
-    c1 = UOp.const(dtypes.index, (1, 1, Invalid, Invalid))
-    c2 = UOp.const(dtypes.index, (1, Invalid, 1, 1))
-    self.assertIs((c1+c2).simplify(), UOp.const(dtypes.index, (2, Invalid, Invalid, Invalid)))
+    c1 = UOp.const(dtypes.weakint, (1, 1, Invalid, Invalid))
+    c2 = UOp.const(dtypes.weakint, (1, Invalid, 1, 1))
+    self.assertIs((c1+c2).simplify(), UOp.const(dtypes.weakint, (2, Invalid, Invalid, Invalid)))
 
 class TestStoreLoadFolding(unittest.TestCase):
   """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
   def test_store_load_folding(self):
     # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
     buf = UOp.param(0, dtypes.int, (1,))
-    index = buf.index(UOp.const(dtypes.index, 0))
+    index = buf.index(UOp.const(dtypes.weakint, 0))
     # Direct: store(idx, load(idx)) -> NOOP
     self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
     # Emergent: store(idx, load(idx) + 0) -> store(idx, load(idx)) -> NOOP
@@ -1319,11 +1392,7 @@ class TestMoveWhereOnLoad(unittest.TestCase):
     idx = buf.index(a.valid(valid))
     expr = cond.where(idx, idx.const_like(0))
     out = graph_rewrite(expr, pm_move_where_on_load)
-    # any WHERE in the rewritten graph must have matched-dtype branches
-    for u in out.toposort():
-      if u.op is Ops.WHERE:
-        self.assertEqual(u.dtype, u.src[1].dtype, f"WHERE branch 1 dtype mismatch: {u}")
-        self.assertEqual(u.dtype, u.src[2].dtype, f"WHERE branch 2 dtype mismatch: {u}")
+    type_verify(out, spec_shared)  # Invalid matches any dtype
 
 class TestSymbolicRealWorld(unittest.TestCase):
   def test_resnet_half(self):

test/null/test_uop_vmin_vmax.py

@@ -167,9 +167,8 @@ class TestVminVmaxProperties(unittest.TestCase):
     self.assertNotEqual(i.vmin, i.vmax)
 
   def test_vmin_vmax_invalid_vconst(self):
-    x = UOp.const(dtypes.index, (0, 4, Invalid, Invalid))
-    self.assertLess(x.vmin, 0)
-    self.assertGreater(x.vmax, 4)
+    x = UOp.const(dtypes.weakint, (0, 4, Invalid, Invalid))
+    self.assertEqual((x.vmin, x.vmax), (0, 4))
 
 class TestVminVmaxDivMod(unittest.TestCase):
   def test_vmin_vmax_division_positive(self):