more test

maketuple
fix
2026-08-18 14:38:26 +00:00 · 2026-03-16 12:37:16 +08:00 · 2026-03-16 12:04:39 +08:00 · 2026-03-16 11:56:36 +08:00 · 2026-03-16 11:44:15 +08:00 · 2026-03-16 11:33:52 +08:00
102 changed files with 1447 additions and 4170 deletions
@@ -66,14 +66,14 @@ runs:
      uses: actions/cache/restore@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
    - name: Cache Python packages
      if: github.event_name != 'pull_request'
      id: restore-venv
      uses: actions/cache@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}

    # **** Caching downloads ****

@@ -199,13 +199,13 @@ runs:
      uses: actions/cache/restore@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name != 'pull_request'
      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}

    - name: Run apt Update + Install
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
@@ -1282,7 +1282,7 @@ def train_bert():
        previous_step = i

 def train_llama3():
-  from extra.models.llama import Transformer
+  from examples.mlperf.models.llama import Transformer
  from examples.llama3 import MODEL_PARAMS
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
  from examples.mlperf.optim import GradAccClipAdamW
@@ -1343,7 +1343,8 @@ def train_llama3():
  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size

-  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+  model = Transformer(**model_params, max_context=SEQLEN)
+
  params = get_parameters(model)
  # weights are all bfloat16 for now
  assert params and all(p.dtype == dtypes.bfloat16 for p in params)
@@ -1381,16 +1382,20 @@ def train_llama3():

    vocab_mask.shard_(device, axis=2).realize()

-  optim_device = "CPU" if getenv("OFFLOAD_OPTIM") else None
+  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(get_parameters(model), 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)

  # init grads
-  for p in optim.params:
-    p.grad = p.empty_like().realize()
+  if is_offload_optim:
+    for p in optim.params:
+      p.grad = Tensor.zeros(p.shape, dtype=p.dtype, device=optim_device, requires_grad=False).contiguous().realize()
+  else:
+    for p in optim.params:
+      p.grad = p.zeros_like().contiguous().realize()
  grads: list[Tensor] = [p.grad for p in optim.params]
-  for p in optim.params:
-    p.grad.assign(p.grad.zeros_like()).realize()

  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)

@@ -1412,7 +1417,7 @@ def train_llama3():
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
      tokens = tokens.shard(device)
    if DP == 1 and MP == 1: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
+    logits:Tensor = model(tokens[:, :-1])
    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
    loss.backward()
    assert all(p.grad is g for p,g in zip(optim.params, grads))
@@ -1426,7 +1431,7 @@ def train_llama3():
    scheduler.step()

    for g in grads:
-      g.assign(g.zeros_like()).realize()
+      g.assign(g.zeros_like())

    lr_cpu = optim.lr.float().to("CPU")
    grad_norm_cpu = grad_norm.float().to("CPU")
@@ -1444,7 +1449,7 @@ def train_llama3():
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
      tokens = tokens.shard(device)
    if DP == 1 and MP == 1: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
+    logits:Tensor = model(tokens[:, :-1])
    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
    return loss.flatten().float().to("CPU")

@@ -0,0 +1,80 @@
+from tinygrad import Tensor, nn
+from tinygrad.helpers import getenv
+from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+
+class Attention:
+  def __init__(self, dim:int, n_heads:int, n_kv_heads:int|None=None, linear=nn.Linear):
+    self.n_heads = n_heads
+    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
+    self.head_dim = dim // n_heads
+    self.n_rep = self.n_heads // self.n_kv_heads
+
+    if getenv("WQKV"):
+      self.wqkv = linear(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2, bias=False)
+    else:
+      self.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
+      self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+      self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+
+    self.wo = linear(self.n_heads * self.head_dim, dim, bias=False)
+
+  def __call__(self, x:Tensor, freqs_cis:Tensor) -> Tensor:
+    if getenv("WQKV"):
+      xqkv = self.wqkv(x)
+      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
+      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+    else:
+      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+
+    xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
+    xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
+    xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
+
+    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+    bsz, seqlen, _, _ = xq.shape
+
+    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
+    attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
+
+    attn = attn.reshape(bsz, seqlen, -1)
+    return self.wo(attn)
+
+class FeedForward:
+  def __init__(self, dim:int, hidden_dim:int, linear=nn.Linear):
+    self.w1 = linear(dim, hidden_dim, bias=False)
+    self.w2 = linear(hidden_dim, dim, bias=False)
+    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit
+
+  def __call__(self, x:Tensor) -> Tensor:
+    w1 = self.w1(x).silu()
+    w3 = self.w3(x)
+    return self.w2(w1 * w3)
+
+class TransformerBlock:
+  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int|None, norm_eps:float, linear=nn.Linear):
+    self.attention = Attention(dim, n_heads, n_kv_heads, linear)
+    self.feed_forward = FeedForward(dim, hidden_dim, linear)
+    self.attention_norm = nn.RMSNorm(dim, norm_eps)
+    self.ffn_norm = nn.RMSNorm(dim, norm_eps)
+
+  def __call__(self, x:Tensor, freqs_cis:Tensor):
+    h = x + self.attention(self.attention_norm(x), freqs_cis)
+    return h + self.feed_forward(self.ffn_norm(h))
+
+class Transformer:
+  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
+               rope_theta:int=10000, max_context:int=1024, linear=nn.Linear, embedding=nn.Embedding):
+    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, linear) for _ in range(n_layers)]
+    self.norm = nn.RMSNorm(dim, norm_eps)
+    self.tok_embeddings = embedding(vocab_size, dim)
+    self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
+
+  def __call__(self, tokens:Tensor):
+    h = self.tok_embeddings(tokens)
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
+    for layer in self.layers: h = layer(h, freqs_cis)
+    logits = self.output(self.norm(h))
+    return logits
@@ -34,10 +34,10 @@ class GradAccClipAdamW(Optimizer):
      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).realize()
-      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous().realize()
+        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)).realize()
+        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
@@ -45,8 +45,8 @@ class GradAccClipAdamW(Optimizer):
    for i, g in enumerate(grads):
      self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
      self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
-      m_hat = self.m[i] / (1.0 - self.b1_t)
-      v_hat = self.v[i] / (1.0 - self.b2_t)
+      m_hat = (self.m[i] / (1.0 - self.b1_t)).cast(self.m[i].dtype)
+      v_hat = (self.v[i] / (1.0 - self.b2_t)).cast(self.v[i].dtype)
      up = m_hat / (v_hat.sqrt() + self.eps)
      ret.append((self.lr * up).cast(g.dtype))
    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
@@ -12,6 +12,7 @@ export ALL2ALL=${ALL2ALL:-1}
 export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-1} MP=${MP:-8}
@@ -12,6 +12,7 @@ export ALL2ALL=${ALL2ALL:-1}
 export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-1} MP=${MP:-8}
@@ -65,7 +65,7 @@ def get_bar0_size(pcibus):
 class AMSMI(AMDev):
  def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
    self.pcibus = pcibus
-    self.vram, self.doorbell64, self.mmio, self.dma_regions = vram_bar, doorbell_bar, mmio_bar, None
+    self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
    self.pci_state = self.read_pci_state()
    if self.pci_state == "D0": self._init_from_d0()

@@ -7,8 +7,8 @@ class GFXFake:
  def __init__(self): self.xccs = 8

 class AMDFake(AMDev):
-  def __init__(self, pci_dev, dma_regions=None):
-    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
+  def __init__(self, pci_dev):
+    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
    self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
    self._run_discovery()
    self._build_regs()
@@ -0,0 +1,84 @@
+from tinygrad import UOp, getenv
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+from tinygrad.dtype import AddrSpace, dtypes
+
+N = getenv("N", 4096)
+M = getenv("M", N)
+K = getenv("K", N)
+
+WARP_SIZE = 32
+BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 8
+TM, TN = 4, 4
+LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
+assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
+
+is_kernel5 = getenv("K5", 0)
+THREADS_PER_BLOCK = 128 if is_kernel5 else 256
+WAVES_PER_BLOCK_N = 1 if is_kernel5 else 2
+WAVES_PER_BLOCK_M = THREADS_PER_BLOCK // WARP_SIZE // WAVES_PER_BLOCK_N
+REG_TILES_PER_WAVE_N = BLOCK_N // (WAVES_PER_BLOCK_N * LANES_PER_WAVE_N * TN)
+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"
+
+consts = {"wpb_m":WAVES_PER_BLOCK_M, "lpw_m":LANES_PER_WAVE_M, "rt_m":REG_TILES_PER_WAVE_M, "t_m": TM,
+          "wpb_n":WAVES_PER_BLOCK_N, "lpw_n":LANES_PER_WAVE_N, "rt_n":REG_TILES_PER_WAVE_N, "t_n": TN}
+
+# 128x128 out, kx128, kx128 in
+def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
+  tid = UOp.range(THREADS_PER_BLOCK, 2, AxisType.LOCAL)
+
+  #tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
+  warp, lane = tid // WARP_SIZE, tid % WARP_SIZE
+  wave_n, wave_m = warp % WAVES_PER_BLOCK_N, warp // WAVES_PER_BLOCK_N
+  lane_n, lane_m = lane % LANES_PER_WAVE_N, lane // LANES_PER_WAVE_N
+
+  # define locals
+  A_local = UOp.placeholder((BLOCK_K, BLOCK_M), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
+  B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+
+  # open the main reduction range and copy in GLOBAL -> LOCAL
+  a = a.reshape(K // BLOCK_K, BLOCK_K, BLOCK_M)
+  b = b.reshape(K // BLOCK_K, BLOCK_K, BLOCK_N)
+  k_tile_range = UOp.range(K // BLOCK_K, 3, AxisType.REDUCE)
+  A_store = A_local.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile_range].reshape(-1, THREADS_PER_BLOCK)[:, tid])
+  B_store = B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile_range].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)
+
+  # define accumulator (128x128), but broadcast across tid
+  c_regs = UOp.placeholder((REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  c_regs = c_regs.after(c_regs.store(UOp.const(dtypes.float, 0).reshape((1,)*len(c_regs.shape)).expand(c_regs.shape)))
+
+  # define registers (NOTE: the thread count is the device count for this multi, it's sharded across the THREADS_PER_BLOCK)
+  A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+
+  # LOCAL -> REGS
+  k = UOp.range(BLOCK_K, 4, AxisType.REDUCE)
+  A_col = A_col.after(A_col.store(A_local[k].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[wave_m, :, lane_m, :]))
+  B_row = B_row.after(B_row.store(B_local[k].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[wave_n, :, lane_n, :]))
+
+  # do FMA
+  A_col = A_col.reshape(REG_TILES_PER_WAVE_M*TM, 1).expand(REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN)
+  B_row = B_row.reshape(1, REG_TILES_PER_WAVE_N*TN).expand(REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN)
+  c_regs = c_regs.after(c_regs.store(c_regs.after(k) + (A_col * B_row)).end(k).barrier().end(k_tile_range))
+
+  # store back to c
+  c_store = c.rearrange("(wpb_m rt_m lpw_m t_m) (wpb_n rt_n lpw_n t_n) -> (wpb_m wpb_n lpw_m lpw_n) (rt_m t_m) (rt_n t_n)", **consts)
+  return c_store[tid].store(c_regs).end(tid)
+
+def amd_copy_matmul(c:UOp, a:UOp, b:UOp) -> UOp:
+  block_id_n = UOp.range(N // BLOCK_N, 0, AxisType.GLOBAL)
+  block_id_m = UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL)
+
+  # index the output with the globals
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
+  a = a.T.reshape(K, M // BLOCK_M, BLOCK_M)[:, block_id_m, :]
+  b = b.reshape(K, N // BLOCK_N, BLOCK_N)[:, block_id_n, :]
+  return block_128x128_gemm(c, a, b).end(block_id_n, block_id_m).sink(arg=KernelInfo(opts_to_apply=()))
+
+if __name__ == "__main__":
+  from amd_uop_matmul import eval_custom_matmul
+  eval_custom_matmul(amd_copy_matmul)
@@ -1,98 +1,74 @@
-import numpy as np
-from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
+from tinygrad import Tensor, Context, GlobalCounters, dtypes
 from tinygrad.uop.ops import UOp, KernelInfo, sint, AxisType
-from tinygrad.engine.realize import ExecItem, get_runner
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import getenv
+from tinygrad.helpers import DEBUG, getenv

 N = getenv("N", 4096)
-M = K = N
-run_count = getenv("CNT", 5)
+M = getenv("M", N)
+K = getenv("K", N)
+NUM_RUNS = getenv("CNT", 5)

 # ---------------------------
 # launch/config constants
 # ---------------------------

 WARP_SIZE = 32
-
-# Threadblock tile sizes (block-level tile of C that a block computes)
-BLOCK_N = 128   # columns of C (N-dim) per block
-BLOCK_M = 128   # rows of C (M-dim) per block
-BLOCK_K = 8     # K-slice per block iteration
-
-# Register tile sizes (per-thread accumulator tile of C)
-TN = 4     # columns per thread
-TM = 4     # rows per thread
+BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 8
+TM, TN = 4, 4
+LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
+assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0

 is_kernel5 = getenv("K5", 0)
 THREADS_PER_BLOCK = 128 if is_kernel5 else 256
-assert THREADS_PER_BLOCK % BLOCK_N == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_N"
-assert THREADS_PER_BLOCK % BLOCK_K == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_K"
-assert (BLOCK_N * BLOCK_K) % THREADS_PER_BLOCK == 0
-assert (BLOCK_M * BLOCK_K) % THREADS_PER_BLOCK == 0
+WAVES_PER_BLOCK_N = 1 if is_kernel5 else 2
+WAVES_PER_BLOCK_M = THREADS_PER_BLOCK // WARP_SIZE // WAVES_PER_BLOCK_N
+REG_TILES_PER_WAVE_N = BLOCK_N // (WAVES_PER_BLOCK_N * LANES_PER_WAVE_N * TN)
+REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)

-WARPS_PER_BLOCK = THREADS_PER_BLOCK // WARP_SIZE
-WAVE_TILE_N = 128 if is_kernel5 else 64
-WAVE_TILE_M = BLOCK_N * BLOCK_M // WARPS_PER_BLOCK // WAVE_TILE_N
-assert BLOCK_N % WAVE_TILE_N == 0, "BN must be a multiple of WN"
-assert BLOCK_M % WAVE_TILE_M == 0, "BM must be a multiple of WM"
-WAVES_IN_BLOCK_X = BLOCK_N // WAVE_TILE_N
-WAVES_IN_BLOCK_Y = BLOCK_M // WAVE_TILE_M
-assert WAVES_IN_BLOCK_X * WAVES_IN_BLOCK_Y == WARPS_PER_BLOCK, "wave grid must match warps/block"
-
-LANES_PER_WAVE_X = 8
-LANES_PER_WAVE_Y = 4
-ITERS_PER_WAVE_N = WAVE_TILE_N // (LANES_PER_WAVE_X * TN)
-ITERS_PER_WAVE_M = WAVE_TILE_M // (LANES_PER_WAVE_Y * TM)
-assert WAVE_TILE_N % (LANES_PER_WAVE_X * TN) == 0, "WAVE_TILE_N must be divisible by LANES_PER_WAVE_X*TN"
-assert WAVE_TILE_M % (LANES_PER_WAVE_Y * TM) == 0, "WAVE_TILE_M must be divisible by LANES_PER_WAVE_Y*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 copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
+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)
-  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
-  return dest.after(copy) if set else copy
+  return dest[*rngs].store(src[*rngs]).end(*rngs)

-def hand_spec_kernel3():
+def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
  # ---------------------------
-  # block indices & placeholders
+  # block indices
  # ---------------------------
-  blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
-  blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")
-
-  a = UOp.placeholder((N, N), dtypes.float, slot=1)
-  b = UOp.placeholder((N, N), dtypes.float, slot=2)
-  c = UOp.placeholder((N, N), dtypes.float, slot=0)
+  block_id_n = UOp.special(N // BLOCK_N, "gidx0")
+  block_id_m = UOp.special(M // BLOCK_M, "gidx1")

  # index the output with the globals
-  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]

  # open the main reduction range
-  k_tile_range = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)
-  a = a.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_K, BLOCK_K)[blockIdx_y, :, k_tile_range, :]
-  b = b.reshape(N // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, blockIdx_x, :]
+  k_tile_range = UOp.range(K // BLOCK_K, 0, AxisType.REDUCE)
+  a = a.reshape(M // BLOCK_M, BLOCK_M, K // BLOCK_K, BLOCK_K)[block_id_m, :, k_tile_range, :]
+  b = b.reshape(K // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, block_id_n, :]

  # globals are no longer used, they are already in the indexes
-  del blockIdx_y, blockIdx_x
+  del block_id_m, block_id_n

  # ---------------------------
-  # GLOBAL -> LOCAL (As, Bs)
+  # GLOBAL -> LOCAL (A_local, B_local)
  # ---------------------------
  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")

  # A: read BM x BK tiles (permute on store into locals)
-  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
-  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
-  As_store = copy(As.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
+  BM_A_local_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
+  A_local = UOp.placeholder((BLOCK_K, BM_A_local_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
+  A_local_store = copy(A_local.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)

  # B: read BK x BN tiles
-  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-  Bs_store = copy(Bs.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
+  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(As_store, Bs_store)
-  As, Bs = As.after(barrier), Bs.after(barrier)
+  barrier = UOp.barrier(A_local_store, B_local_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)

  # open inner k range
  k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -100,31 +76,30 @@ def hand_spec_kernel3():
  # ---------------------------
  # LOCAL -> REG (per-wave tiles)
  # ---------------------------
-  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
-  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
-  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
+  warp, lane = tid // WARP_SIZE, tid % WARP_SIZE
+  waveIdx, waveIdy = warp % WAVES_PER_BLOCK_N, warp // WAVES_PER_BLOCK_N
+  laneIdx, laneIdy = lane % LANES_PER_WAVE_N, lane // LANES_PER_WAVE_N
+  assert waveIdy.vmax+1 == WAVES_PER_BLOCK_M and laneIdy.vmax+1 == LANES_PER_WAVE_M

-  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
-  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
-  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
+  A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  A_local_slice = A_local[k, :].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[waveIdy, :, laneIdy, :]
+  A_col = A_col.after(copy(A_col, A_local_slice, 300, upcast=True))

-  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  A_col = copy(A_col, As[k, :].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :], 300, set=True, upcast=True)
-
-  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  B_row = copy(B_row, Bs[k, :].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :], 400, set=True, upcast=True)
+  B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  B_local_slice = B_local[k, :].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[waveIdx, :, laneIdx, :]
+  B_row = B_row.after(copy(B_row, B_local_slice, 400, upcast=True))

  # ---------------------------
  # FMA: c_regs += A_col * B_row
  # ---------------------------
-  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  c_regs = UOp.placeholder((REG_TILES_PER_WAVE_M, TM, REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
  i = UOp.range(c_regs.size, 16)
  c_regs = c_regs.after(c_regs.flatten()[i].store(0.0).end(i))

  # TODO: why don't these work as upcast?
  # why if the ranges merge is it slow?!? (if you change the order on end, they will merge. big slowdown on METAL)
-  iterWaveM, yt, iterWaveN, xt = rngs = rngs_for_shape(c_regs.shape, 500)
-  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iterWaveM, yt] * B_row[iterWaveN, xt]).end(iterWaveM, iterWaveN, yt, xt)
+  iter_m, t_m, iter_n, t_n = rngs = rngs_for_shape(c_regs.shape, 500)
+  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)
@@ -132,38 +107,37 @@ def hand_spec_kernel3():
  # ---------------------------
  # REG -> GLOBAL (epilogue)
  # ---------------------------
-  c = c.reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
-                WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
+  c = c.reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM,
+                WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)
  c = c[waveIdy, :, laneIdy, :,
        waveIdx, :, laneIdx, :]
  sink = copy(c, c_regs.after(sink), rng=600)

  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()

-def test_matmul(sink:UOp, dtype=dtypes.float32, N=N):
-  rng = np.random.default_rng()
-  a = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
-  b = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
-  hc = Tensor.empty(N, N, dtype=dtype)
-  Tensor.realize(a, b, hc)
-
-  ei = ExecItem(sink, [t.uop.buffer for t in [hc, a, b]], prg=get_runner(Device.DEFAULT, sink))
+def eval_custom_matmul(fxn):
+  a = Tensor.randn(M, K, dtype=dtypes.float)
+  b = Tensor.randn(K, N, dtype=dtypes.float)
+  c = Tensor.empty(M, N, dtype=dtypes.float)
+  with Context(DEBUG=0): Tensor.realize(a, b)

  ets = []
-  with Context(DEBUG=2):
-    for _ in range(run_count):
-      ets.append(ei.run(wait=True))
-  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
+  with Context(DEBUG=max(2, DEBUG.value)):
+    for _ in range(NUM_RUNS):
+      GlobalCounters.reset()
+      tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize()
+      ets.append(GlobalCounters.time_sum_s)
+  print(f"REAL TFLOPS {M * N * K * 2 / min(ets) * 1e-12:.2f}")

  if getenv("VERIFY", 1):
    GlobalCounters.reset()
    with Context(DEBUG=2):
      tc = (a @ b).realize()
    with Context(DEBUG=0):
-      err = (hc - tc).square().mean().item()
+      err = (tc - tst).square().mean().item()
    print(f"mean squared error {err}")
    if err > 1e-06:
      raise RuntimeError("matmul is wrong!")

 if __name__ == "__main__":
-  test_matmul(hand_spec_kernel3(), N=N)
+  eval_custom_matmul(hand_spec_kernel3)
@@ -1,122 +0,0 @@
-import atexit, functools
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, DEBUG
-from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG
-
-# ** CDNA4 assembly gemm
-
-WORKGROUP_SIZE = 256
-
-@functools.cache
-def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
-  batch, M, K = A.shape
-  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
-  assert K == K2
-  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
-  gidx = UOp.special(NUM_WG, "gidx0")
-  insts = build_kernel(batch, M, N, K, A.dtype.base)
-  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
-  sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
-                  arg=KernelInfo(name=f"gemm_{batch}_{M}_{N}_{K}", estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname),
-                                UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-
-counters = {"used":0, "todos":[]}
-def todo(msg:str) -> bool: counters["todos"].append(msg); return False
-def _asm_gemm_report():
-  print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used')
-  if DEBUG >= 2 and counters["todos"]:
-    from collections import Counter
-    for msg, cnt in Counter(counters["todos"]).most_common(): print(f'  {cnt:3d}x {msg}')
-atexit.register(_asm_gemm_report)
-
-def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
-  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
-  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
-  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
-  N = b.shape[1]
-  if isinstance(a.device, tuple):
-    if a.ndim == 2 and a.uop.axis == 0 and b.uop.axis is None: M //= len(a.device)
-    elif a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
-    elif a.ndim == 2 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis == 2 and b.uop.axis == 0: K //= len(a.device)
-    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
-    dname = a.device[0]
-  else: dname = a.device
-  arch = getattr(Device[dname].renderer, "arch", "")
-  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
-  if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
-    return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
-  return True
-
-# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
-# note: this can be removed after we have GEMM on mixins
-
-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)
-  k = UOp.range(K, 0, AxisType.REDUCE)
-  mul = (A.index((m*UOp.const(dtypes.index, K)+k))*B.index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
-  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
-  store = C.index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
-  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
-
-# ** backward gemm, might use the asm gemm
-
-def custom_gemm_bw(gradient:UOp, kernel:UOp):
-  out, a, b = kernel.src[1:]
-  assert all_same([gradient.device, a.device, b.device, out.device])
-  a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
-  # TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
-  g_t = g_t[:a.shape[0]]
-  grad_a = (g_t @ b_t.T).uop
-  grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
-  return (None, grad_a, grad_b)
-
-# ** main gemm function
-
-def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
-  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
-  counters["used"] += 1
-  unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
-  if unfold_batch:
-    orig_batch = a.shape[0]
-    a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
-  squeeze = a.ndim == 2
-  if squeeze: a = a.unsqueeze(0)
-
-  batch, M, K = a.shape
-  N = b.shape[1]
-  is_multi = isinstance(a.device, tuple)
-  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
-  if (m_sharded:=is_multi and a.uop.axis == 1): M //= len(a.device)
-  n_sharded = is_multi and b.uop.axis == 1
-
-  if is_multi:
-    if n_sharded:
-      out = Tensor(Tensor.empty(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
-    elif m_sharded:
-      out = Tensor(Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
-    else:
-      out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
-  else:
-    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
-
-  renderer = Device[a.device[0] if is_multi else a.device].renderer
-  dname, arch = renderer.device, getattr(renderer, "arch", "")
-  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
-  else:
-    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
-  if k_sharded: out = out.sum(0)
-  out = out.squeeze(0) if squeeze else out
-  if unfold_batch: out = out.reshape(orig_batch, -1, out.shape[-1])
-  return out
@@ -1,76 +0,0 @@
-.text
-.section	.text.
-.global	gemm
-.p2align	8
-.type	gemm,@function
-
-gemm:
-INSTRUCTIONS
-
-.section .rodata,"a",@progbits
-.p2align 6, 0x0
-.amdhsa_kernel gemm
-  # basic memory requirements
-  .amdhsa_group_segment_fixed_size 30336
-  .amdhsa_private_segment_fixed_size 0
-  .amdhsa_kernarg_size 32
-  # register usage (RSRC1)
-  .amdhsa_next_free_vgpr 256
-  .amdhsa_next_free_sgpr 100
-  # workgroup / workitem IDs (RSRC2)
-  .amdhsa_system_sgpr_workgroup_id_x 1
-  .amdhsa_system_sgpr_workgroup_id_y 1
-  .amdhsa_system_sgpr_workgroup_id_z 1
-  # user SGPRs: kernarg ptr in s[0:1]
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_user_sgpr_count 2
-  # gfx10+ / gfx11 specifics (RSRC1[29..31])
-  .amdhsa_wavefront_size32 1
-  .amdhsa_workgroup_processor_mode 1
-  .amdhsa_memory_ordered 1
-  .amdhsa_forward_progress 1
-  # misc for gfx11
-  .amdhsa_dx10_clamp 1
-  .amdhsa_ieee_mode 1
-  .amdhsa_uses_dynamic_stack 0
-.end_amdhsa_kernel
-
-.amdgpu_metadata
---
-amdhsa.kernels:
-  - .args:
-      - .address_space:  generic
-        .name:           C
-        .offset:         0
-        .size:           8
-        .value_kind:     global_buffer
-        .value_type:     f16
-      - .address_space:  generic
-        .name:           A
-        .offset:         8
-        .size:           8
-        .value_kind:     global_buffer
-        .value_type:     f16
-      - .address_space:  generic
-        .name:           B
-        .offset:         16
-        .size:           8
-        .value_kind:     global_buffer
-        .value_type:     f16
-    .group_segment_fixed_size: 30336
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 32
-    .max_flat_workgroup_size: 128
-    .name:           gemm
-    .private_segment_fixed_size: 0
-    .sgpr_count:     70
-    .sgpr_spill_count: 0
-    .symbol:         gemm.kd
-    .vgpr_count:     256
-    .vgpr_spill_count: 0
-    .wavefront_size: 32
-amdhsa.version:
-  - 1
-  - 1
-...
-.end_amdgpu_metadata
@@ -1,30 +0,0 @@
-import math, pathlib
-
-from tinygrad import Device, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-
-from extra.gemm.amd_uop_matmul import test_matmul
-
-N = 4096
-TN = 96
-THREADS_PER_WG = 128
-NUM_WG = math.ceil(N / TN) * math.ceil(N / TN)
-
-dname:str = Device.DEFAULT
-template:str = (pathlib.Path(__file__).parent/"template.s").read_text()
-
-def asm_kernel() -> UOp:
-  lidx = UOp.special(THREADS_PER_WG, "lidx0")
-  gidx = UOp.special(NUM_WG, "gidx0")
-
-  a = UOp.placeholder((N*N,), dtypes.half, slot=1)
-  b = UOp.placeholder((N*N,), dtypes.half, slot=2)
-  c = UOp.placeholder((N*N,), dtypes.half, slot=0)
-
-  src = template.replace("INSTRUCTIONS", (pathlib.Path(__file__).parent/"gemm.s").read_text())
-
-  sink = UOp.sink(a, b, c, lidx, gidx, arg=KernelInfo(name="gemm"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
-
-if __name__ == "__main__":
-  test_matmul(asm_kernel(), dtype=dtypes.half, N=N)
@@ -1,179 +0,0 @@
-# unpack the complete kernel descriptor of an amdgpu ELF
-# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
-import struct, pathlib, sys
-from tinygrad.runtime.support.elf import elf_loader
-
-def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
-def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
-
-with open(sys.argv[1], "rb") as f:
-  lib = f.read()
-
-image, sections, relocs = elf_loader(lib)
-rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"))
-
-# rodata is exactly 64 bytes
-kd = image[rodata_entry:rodata_entry+64]
-desc = int.from_bytes(kd, byteorder="little")
-
-group_segment_fixed_size = bits(desc, 0, 31)
-private_segment_fixed_size = bits(desc, 32, 63)
-kernarg_size = bits(desc, 64, 95)
-reserved_127_96 = bits(desc, 96, 127)
-assert reserved_127_96 == 0
-
-print("GROUP_SEGMENT_FIXED_SIZE:", group_segment_fixed_size)
-print("PRIVATE_SEGMENT_FIXED_SIZE:", private_segment_fixed_size)
-print("KERNARG_SIZE:", kernarg_size)
-print("RESERVED 127:96:", reserved_127_96)
-
-entry_off = bits(desc, 128, 191)
-
-# sign-extend manually if needed
-if entry_off & (1 << 63):
-  entry_off -= 1 << 64
-
-print("KERNEL_CODE_ENTRY_BYTE_OFFSET:", entry_off)
-
-kd_addr = 0x1840
-entry_addr = kd_addr + entry_off
-
-print("Computed entry address: 0x%016x" % entry_addr)
-print("256B aligned:", entry_addr % 256 == 0)
-
-pgm_rsrc3 = bits(desc, 352, 383)
-pgm_rsrc1 = bits(desc, 384, 415)
-pgm_rsrc2 = bits(desc, 416, 447)
-
-print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
-print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
-print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
-
-# rsrc 3 (gfx950)
-
-accum_offset_raw = bits(pgm_rsrc3, 0, 5)
-assert_zero(pgm_rsrc3, 6, 15)
-tg_split = bits(pgm_rsrc3, 16, 16)
-accum_offset_vgprs = (accum_offset_raw + 1) * 4
-print("RSRC3.ACCUM_OFFSET (AccVGPR index):", accum_offset_vgprs)
-print("RSRC3.TG_SPLIT:", tg_split)
-
-# rsrc 1
-
-vgpr_gran = bits(pgm_rsrc1, 0, 5)
-sgpr_gran = bits(pgm_rsrc1, 6, 9)
-assert_zero(pgm_rsrc1, 27, 28)
-
-# NOTE: this is vgprs + agprs
-vgprs_used = (vgpr_gran + 1) * 8
-assert 0 <= vgprs_used <= 512
-
-k = sgpr_gran // 2
-sgprs_used = (k + 1) * 16
-
-print("RSRC1.VGPRS:", vgprs_used)
-print("RSRC1.SGPRS:", sgprs_used)
-
-assert_zero(pgm_rsrc1, 10, 11)
-
-float_round_mode_32 = bits(pgm_rsrc1, 12, 13)
-float_round_mode_16_64 = bits(pgm_rsrc1, 15, 14)
-float_denorm_mode_32 = bits(pgm_rsrc1, 16, 17)
-float_denorm_mode_16_64 = bits(pgm_rsrc1, 18, 19)
-
-priv = bits(pgm_rsrc1, 20, 20)
-assert priv == 0
-enable_dx10_clamp_wg_rr_en = bits(pgm_rsrc1, 21, 21)
-debug_mode = bits(pgm_rsrc1, 22, 22)
-enable_ieee_mode = bits(pgm_rsrc1, 23, 23)
-bulky = bits(pgm_rsrc1, 24, 24)
-assert bulky == 0
-cdbg_user = bits(pgm_rsrc1, 25, 25)
-assert cdbg_user == 0
-fp16_ovfl = bits(pgm_rsrc1, 26, 26)
-assert_zero(pgm_rsrc1, 27, 28)  # reserved
-assert_zero(pgm_rsrc1, 29, 29)  # WGP_MODE (reserved on gfx9)
-assert_zero(pgm_rsrc1, 30, 30)  # MEM_ORDERED (reserved on gfx9)
-assert_zero(pgm_rsrc1, 31, 31)  # FWD_PROGRESS (reserved on gfx9)
-
-# rsrc 2
-
-enable_private_segment = bits(pgm_rsrc2, 0, 0)  # SCRATCH_EN
-user_sgpr_count = bits(pgm_rsrc2, 1, 5)         # USER_SGPR
-enable_trap_handler = bits(pgm_rsrc2, 6, 6)     # TRAP_PRESENT (must be 0 here)
-assert enable_trap_handler == 0
-
-enable_sgpr_workgroup_id_x = bits(pgm_rsrc2, 7, 7)
-enable_sgpr_workgroup_id_y = bits(pgm_rsrc2, 8, 8)
-enable_sgpr_workgroup_id_z = bits(pgm_rsrc2, 9, 9)
-enable_sgpr_workgroup_info = bits(pgm_rsrc2, 10, 10)
-
-enable_vgpr_workitem_id = bits(pgm_rsrc2, 11, 12)  # TIDIG_CMP_CNT enum (0..3)
-
-enable_exception_address_watch = bits(pgm_rsrc2, 13, 13)
-assert enable_exception_address_watch == 0
-enable_exception_memory = bits(pgm_rsrc2, 14, 14)
-assert enable_exception_memory == 0
-
-granulated_lds_size = bits(pgm_rsrc2, 15, 23)
-assert granulated_lds_size == 0  # spec: must be 0; CP uses dispatch packet rounding
-
-enable_exception_fp_invalid = bits(pgm_rsrc2, 24, 24)
-enable_exception_fp_denorm_src = bits(pgm_rsrc2, 25, 25)
-enable_exception_fp_div0 = bits(pgm_rsrc2, 26, 26)
-enable_exception_fp_overflow = bits(pgm_rsrc2, 27, 27)
-enable_exception_fp_underflow = bits(pgm_rsrc2, 28, 28)
-enable_exception_fp_inexact = bits(pgm_rsrc2, 29, 29)
-enable_exception_int_div0 = bits(pgm_rsrc2, 30, 30)
-
-assert_zero(pgm_rsrc2, 31, 31)
-
-print("RSRC2.ENABLE_PRIVATE_SEGMENT:", enable_private_segment)
-print("RSRC2.USER_SGPR_COUNT:", user_sgpr_count)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_X:", enable_sgpr_workgroup_id_x)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Y:", enable_sgpr_workgroup_id_y)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Z:", enable_sgpr_workgroup_id_z)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_INFO:", enable_sgpr_workgroup_info)
-print("RSRC2.ENABLE_VGPR_WORKITEM_ID (enum):", enable_vgpr_workitem_id)
-
-print("RSRC2.EXC_FP_INVALID:", enable_exception_fp_invalid)
-print("RSRC2.EXC_FP_DENORM_SRC:", enable_exception_fp_denorm_src)
-print("RSRC2.EXC_FP_DIV0:", enable_exception_fp_div0)
-print("RSRC2.EXC_FP_OVERFLOW:", enable_exception_fp_overflow)
-print("RSRC2.EXC_FP_UNDERFLOW:", enable_exception_fp_underflow)
-print("RSRC2.EXC_FP_INEXACT:", enable_exception_fp_inexact)
-print("RSRC2.EXC_INT_DIV0:", enable_exception_int_div0)
-
-# user sgprs
-
-enable_sgpr_private_segment_buffer = bits(desc, 448, 448)
-enable_sgpr_dispatch_ptr = bits(desc, 449, 449)
-enable_sgpr_queue_ptr = bits(desc, 450, 450)
-enable_sgpr_kernarg_segment_ptr = bits(desc, 451, 451)
-enable_sgpr_dispatch_id = bits(desc, 452, 452)
-enable_sgpr_flat_scratch_init = bits(desc, 453, 453)
-enable_sgpr_private_segment_size = bits(desc, 454, 454)
-
-assert_zero(desc, 455, 457)
-
-print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER:", enable_sgpr_private_segment_buffer)
-print("DESC.ENABLE_SGPR_DISPATCH_PTR:", enable_sgpr_dispatch_ptr)
-print("DESC.ENABLE_SGPR_QUEUE_PTR:", enable_sgpr_queue_ptr)
-print("DESC.ENABLE_SGPR_KERNARG_SEGMENT_PTR:", enable_sgpr_kernarg_segment_ptr)
-print("DESC.ENABLE_SGPR_DISPATCH_ID:", enable_sgpr_dispatch_id)
-print("DESC.ENABLE_SGPR_FLAT_SCRATCH_INIT:", enable_sgpr_flat_scratch_init)
-print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_SIZE:", enable_sgpr_private_segment_size)
-
-assert_zero(desc, 458, 459)
-
-uses_dynamic_stack = bits(desc, 459, 460)
-print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
-
-# gfx950 only
-assert_zero(desc, 460, 463)
-kernarg_preload_spec_length = bits(desc, 464, 470)
-print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
-kernarg_preload_spec_offset = bits(desc, 471, 479)
-print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
-
-assert_zero(desc, 480, 511)
@@ -1,5 +1,14 @@
+import atexit, functools
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.renderer import Estimates
+from tinygrad.helpers import getenv, all_same, DEBUG
 from tinygrad.runtime.autogen.amd.cdna.ins import *
-from tinygrad.dtype import dtypes
+
+# ** CDNA4 assembly gemm
+
+WORKGROUP_SIZE = 256

 # M0 is encoded with 124 (NULL in RDNA) in CDNA
 M0 = NULL
@@ -2597,3 +2606,117 @@ def build_kernel(batch, M, N, K, dtype):
  k.label('KernelEnd')
  k.emit(s_endpgm())
  return k.finalize()
+
+# ** ASM_GEMM custom kernel
+
+@functools.cache
+def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
+  batch, M, K = A.shape
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
+  gidx = UOp.special(NUM_WG, "gidx0")
+  insts = build_kernel(batch, M, N, K, A.dtype.base)
+  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
+  sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
+                  arg=KernelInfo(name=f"gemm_{batch}_{M}_{N}_{K}", estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname),
+                                UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+
+counters = {"used":0, "todos":[]}
+def todo(msg:str) -> bool: counters["todos"].append(msg); return False
+def _asm_gemm_report():
+  print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used')
+  if DEBUG >= 2 and counters["todos"]:
+    from collections import Counter
+    for msg, cnt in Counter(counters["todos"]).most_common(): print(f'  {cnt:3d}x {msg}')
+atexit.register(_asm_gemm_report)
+
+def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
+  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
+  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
+  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
+  N = b.shape[1]
+  if isinstance(a.device, tuple):
+    if a.ndim == 2 and a.uop.axis == 0 and b.uop.axis is None: M //= len(a.device)
+    elif a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
+    elif a.ndim == 2 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis == 2 and b.uop.axis == 0: K //= len(a.device)
+    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
+    dname = a.device[0]
+  else: dname = a.device
+  arch = getattr(Device[dname].renderer, "arch", "")
+  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
+  if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
+    return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
+  return True
+
+# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
+# note: this can be removed after we have GEMM on mixins
+
+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)
+  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)
+  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
+  store = C.flatten().index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
+  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
+
+# ** backward gemm, might use the asm gemm
+
+def custom_gemm_bw(gradient:UOp, kernel:UOp):
+  out, a, b = kernel.src[1:]
+  assert all_same([gradient.device, a.device, b.device, out.device])
+  a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
+  # TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
+  g_t = g_t[:a.shape[0]]
+  grad_a = (g_t @ b_t.T).uop
+  grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
+  return (None, grad_a, grad_b)
+
+# ** main gemm function
+
+def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
+  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
+  counters["used"] += 1
+  unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
+  if unfold_batch:
+    orig_batch = a.shape[0]
+    a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
+  squeeze = a.ndim == 2
+  if squeeze: a = a.unsqueeze(0)
+
+  batch, M, K = a.shape
+  N = b.shape[1]
+  is_multi = isinstance(a.device, tuple)
+  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
+  if (m_sharded:=is_multi and a.uop.axis == 1): M //= len(a.device)
+  n_sharded = is_multi and b.uop.axis == 1
+
+  if is_multi:
+    if n_sharded:
+      out = Tensor(Tensor.empty(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
+    elif m_sharded:
+      out = Tensor(Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
+    else:
+      out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+  else:
+    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
+
+  renderer = Device[a.device[0] if is_multi else a.device].renderer
+  dname, arch = renderer.device, getattr(renderer, "arch", "")
+  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
+    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
+  else:
+    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
+  if k_sharded: out = out.sum(0)
+  out = out.squeeze(0) if squeeze else out
+  if unfold_batch: out = out.reshape(orig_batch, -1, out.shape[-1])
+  return out
@@ -1,148 +0,0 @@
-#!/usr/bin/env python3
-# Run all ALU and memory instructions in the ISA
-import functools, inspect
-from enum import Enum
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AddrSpace
-from tinygrad.renderer.amd.dsl import Inst, Reg, OPERANDS, SrcField, VGPRField, SGPRField, SSrcField, SBaseField, AlignedSGPRField, BitField
-from tinygrad.renderer.amd.dsl import FixedBitField, EnumBitField, s, v, NULL, VCC_LO
-from extra.gemm.amd_asm_matmul import Kernel
-
-# skip instructions that mutate wave state (PC, EXEC, allocations, signals)
-SKIP = {"S_SETPC_B64", "S_SWAPPC_B64", "S_RFE_B64", "S_BARRIER_SIGNAL_ISFIRST", "S_GET_BARRIER_STATE", "S_ALLOC_VGPR", "S_SLEEP_VAR", "S_GETPC_B64",
-        "S_SENDMSG_RTN_B32", "S_SENDMSG_RTN_B64"}
-# skip barriers, s_waits, wrap level atomics, and ray tracing (bvh)
-SKIP_SUBSTR = ["SAVEEXEC", "CMPX", "WREXEC", "MOVREL", "ATOMIC", "S_BUFFER_", "S_ATC_PROBE", "BARRIER", "S_WAITCNT", "BVH",
-               "DS_CMPSTORE_RTN", "DS_WRAP_RTN_B32", "DS_ORDERED_COUNT", "DS_GWS", "GS_REG", "GLOBAL_LOAD_LDS", "GLOBAL_STORE_BLOCK"]
-
-ALU_FORMATS = {"VOP1", "VOP1_LIT", "VOP1_SDST", "VOP2", "VOP2_LIT", "VOP3", "VOP3_SDST", "VOP3SD", "VOP3P", "VOP3P_MFMA", "VOP3PX2",
-               "VOPC", "SOP1", "SOP1_LIT", "SOP2", "SOP2_LIT", "SOPC", "SOPC_LIT", "SOPK", "SOPK_LIT", "VINTERP"}
-# intentionally not testing scratch memory ops
-MEM_FORMATS = {"VGLOBAL", "GLOBAL", "SMEM", "DS"}
-
-def should_skip(op:Enum) -> bool: return (name:=op.name) in SKIP or any(sub in name for sub in SKIP_SUBSTR)
-
-# ** named register assignments
-
-# ALU operands
-ALU_VGPR_STRIDE = 16          # v[0], v[16], v[32], ... per ALU operand slot
-ALU_SGPR_STRIDE = 4           # s[0], s[4], s[8], ... per ALU operand slot
-
-# memory address registers
-S_KERNARG_PTR = (0, 1)
-S_BUF_PTR = (2, 3)
-V_VADDR = (0, 1)
-V_DS_ADDR = 0
-
-# memory data registers
-MEM_VGPR_BASE = 32            # v[32], v[48], ... for vdst/vdata/vsrc
-MEM_VGPR_STRIDE = 16          # spacing between memory data vgpr slots
-MEM_SGPR_BASE = 8             # s[8], s[10], ... for SMEM sdata
-MEM_SGPR_STRIDE = 2           # spacing between memory data sgpr slots
-
-# ** create an ALU instruction based on the operands
-
-def create_alu_inst(op:Enum, builder:functools.partial[Inst]) -> Inst:
-  inst_cls, operands, slot = builder.func, OPERANDS[op], 0
-  kwargs:dict[str, Reg|int] = {}
-  for name, field in inst_cls._fields:
-    if isinstance(field, (FixedBitField, EnumBitField)): continue
-    nregs = max(1, operands[name][1] // 32) if name in operands else 1
-    is_sreg = name in operands and "SREG" in str(operands[name][2])
-    base_v, base_s = slot * ALU_VGPR_STRIDE, slot * ALU_SGPR_STRIDE
-    if name == "sdst" and isinstance(field, SGPRField): reg = VCC_LO
-    elif is_sreg and not isinstance(field, VGPRField): reg = VCC_LO
-    elif isinstance(field, VGPRField): reg = v[base_v:base_v+nregs-1] if nregs > 1 else v[base_v]
-    elif isinstance(field, SSrcField): reg = VCC_LO if nregs <= 2 else s[base_s:base_s+nregs-1] if nregs > 1 else s[base_s]
-    elif isinstance(field, SGPRField): reg = s[base_s:base_s+nregs-1] if nregs > 1 else s[base_s]
-    elif isinstance(field, SrcField): reg = v[base_v:base_v+nregs-1] if nregs > 1 else v[base_v]
-    else: reg = None
-    if reg is not None: kwargs[name] = reg; slot += 1
-    elif isinstance(field, BitField): kwargs[name] = field.default
-  return builder(**kwargs)
-
-# ** create a memory instruction with pre set address registers
-
-MEM_PRESET_REGS:dict[str, dict[str, Reg]] = {
-  "VGLOBAL":{"saddr":s[S_BUF_PTR[0]:S_BUF_PTR[1]], "vaddr":v[V_VADDR[0]:V_VADDR[1]]},
-  "GLOBAL":{"saddr":s[S_BUF_PTR[0]:S_BUF_PTR[1]], "addr":v[V_DS_ADDR]},  # addr is 32-bit offset when saddr is valid SGPR
-  "DS":{"addr":v[V_DS_ADDR]},
-  "SMEM":{"sbase":s[S_KERNARG_PTR[0]:S_KERNARG_PTR[1]], "soffset":NULL},
-}
-
-def create_mem_inst(op:Enum, builder:functools.partial[Inst]) -> Inst:
-  inst_cls, operands, field_map = builder.func, OPERANDS.get(op, {}), MEM_PRESET_REGS.get(builder.func.__name__, {})
-  kwargs:dict[str, Reg|int] = {}
-  vslot, sslot = 0, 0
-  for name, field in inst_cls._fields:
-    if isinstance(field, (FixedBitField, EnumBitField)): continue
-    if name in field_map:
-      kwargs[name] = field_map[name]
-      continue
-    nregs = max(1, operands[name][1] // 32) if name in operands else 1
-    if isinstance(field, VGPRField):
-      vi = MEM_VGPR_BASE + vslot * MEM_VGPR_STRIDE
-      kwargs[name] = v[vi:vi+nregs-1] if nregs > 1 else v[vi]
-      vslot += 1
-    elif isinstance(field, (SGPRField, AlignedSGPRField, SBaseField)):
-      si = MEM_SGPR_BASE + sslot * MEM_SGPR_STRIDE
-      kwargs[name] = s[si:si+nregs-1] if nregs > 1 else s[si]
-      sslot += 1
-    elif isinstance(field, BitField): kwargs[name] = field.default
-  return builder(**kwargs)
-
-# ** collect all memory and ALU instructions from the ISA autogen
-
-def collect_instructions() -> tuple[list[Inst], list[Inst], list[str]]:
-  op_map:dict[Enum, functools.partial[Inst]] = {}
-  for name, obj in inspect.getmembers(all_insts):
-    if isinstance(obj, functools.partial) and len(obj.args) == 1: op_map[obj.args[0]] = obj
-  alu_insts:list[Inst] = []
-  mem_insts:list[Inst] = []
-  skipped:list[str] = []
-  for op_enum, builder in op_map.items():
-    if should_skip(op_enum) or op_enum not in OPERANDS: skipped.append(op_enum.name); continue
-    fmt = builder.func.__name__
-    if fmt in ALU_FORMATS: alu_insts.append(create_alu_inst(op_enum, builder))
-    elif fmt in MEM_FORMATS: mem_insts.append(create_mem_inst(op_enum, builder))
-  return alu_insts, mem_insts, skipped
-
-def exec_insts(insts:list):
-  k = Kernel(arch)
-  # ** prologue for global memory
-  k.emit(s_load_b64(sdata=s[S_BUF_PTR[0]:S_BUF_PTR[1]], sbase=s[S_KERNARG_PTR[0]:S_KERNARG_PTR[1]], soffset=NULL))
-  k.waitcnt(lgkm=0)
-  k.emit(v_mov_b32_e32(v[V_VADDR[0]], 0))
-  k.emit(v_mov_b32_e32(v[V_VADDR[1]], 0))
-  # ** emit
-  for inst in insts: k.emit(inst)
-  k.emit(s_endpgm())
-  # ** run
-  NUM_THREADS, NUM_GRIDS, BUF_SIZE = 32, 1, 1024*1024
-  def fxn(A:UOp, B:UOp, C:UOp) -> UOp:
-    lidx, gidx = UOp.special(NUM_THREADS, "lidx0"), UOp.special(NUM_GRIDS, "gidx0")
-    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=BUF_SIZE, addrspace=AddrSpace.LOCAL), (), "lds")
-    sink = UOp.sink(A.base, B.base, C.base, lds, lidx, gidx, arg=KernelInfo(name="discover_ops"))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in k.finalize()))))
-  A = Tensor.empty(BUF_SIZE, dtype=dtypes.uint8)
-  B = Tensor.empty(1, dtype=dtypes.uint8)
-  C = Tensor.empty(1, dtype=dtypes.uint8)
-  Tensor.custom_kernel(A, B, C, fxn=fxn)[0].realize()
-
-if __name__ == "__main__":
-  import sys
-  arch = Device[Device.DEFAULT].renderer.arch
-  if arch.startswith("gfx12"):
-    from tinygrad.runtime.autogen.amd.rdna4.ins import *
-    import tinygrad.runtime.autogen.amd.rdna4.ins as all_insts
-  elif arch.startswith("gfx11"):
-    from tinygrad.runtime.autogen.amd.rdna3.ins import *
-    import tinygrad.runtime.autogen.amd.rdna3.ins as all_insts
-    # these don"t exist in RDNA3, only RDNA3.5 and above
-    SKIP.update(["S_FMAAK_F32", "S_FMAMK_F32"])
-  else:
-    print(f"{arch} not supported yet")
-    sys.exit(0)
-  alu_insts, mem_insts, skipped = collect_instructions()
-  print(f"collected {len(alu_insts)} ALU + {len(mem_insts)} memory instructions ({len(skipped)} skipped)")
-  exec_insts(mem_insts+alu_insts)
@@ -8,8 +8,8 @@ PROFILE_PATH = Path(temp("profile.pkl", append_user=True))
 EXAMPLES = {
  "empty":"test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
  "plus":"test/test_tiny.py TestTiny.test_plus",
-  "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=64, N)@Tensor.empty(N, N)).realize()\"",
-  "ops":"extra/sqtt/examples/discover_ops.py"
+  "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=32, N)@Tensor.empty(N, N)).realize()\"",
+  "sync":"test/amd/test_custom_kernel.py TestCustomKernel.test_wave_sync",
 }

 if __name__ == "__main__":
@@ -590,7 +590,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  "aten.repeat": lambda x,*repeats: Tensor.repeat(x,*repeats).contiguous(), # not a view
  "aten._softmax": lambda self,dim,half_to_float: self.softmax(dim),
  "aten._log_softmax": lambda self,dim,half_to_float: self.log_softmax(dim),
-  "aten.random_": lambda self: Tensor.randint(*self.shape, low=dtypes.min(self.dtype), high=dtypes.max(self.dtype), device=self.device, dtype=self.dtype),
+  "aten.random_": lambda self: Tensor.randint(*self.shape, low=self.dtype.min, high=self.dtype.max, device=self.device, dtype=self.dtype),
  "aten.random_.from": lambda self, from_, to: Tensor.randint(*self.shape, low=from_, high=to, device=self.device, dtype=self.dtype),
  "aten.uniform_": lambda self, low=0, high=1: Tensor.uniform(*self.shape, low=low, high=high, dtype=self.dtype),
  "aten.normal_": lambda self, mean=0, std=1: Tensor.normal(*self.shape, mean=mean, std=std, dtype=self.dtype),
@@ -48,11 +48,16 @@ def decode_profile(data:bytes) -> dict:
        name, ref, key, st, dur, fmt = u("<IIIIfI")
        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":strings[fmt]})
    else:
+      v["linear"] = u("<B")[0]
      v["peak"] = u("<Q")[0]
      for _ in range(event_count):
-        alloc, ts, key = u("<BII")
-        if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
+        if v["linear"]:
+          ts, value = u("<IQ")
+          v["events"].append({"event":"freq", "ts":ts, "value":value})
+        else:
+          alloc, ts, key = u("<BII")
+          if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
+          else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}

 if __name__ == "__main__":
@@ -6,7 +6,7 @@ from tinygrad.runtime.support.elf import elf_loader

 ARCH_TO_TARGET:dict[str, list[str]] = {
  "rdna3":["gfx1100"],
-  "rdna4":["gfx1200"],
+  "rdna4":["gfx1200", "gfx1201"],
  "cdna":["gfx950", "gfx942"],
 }

@@ -1,10 +1,12 @@
 import unittest
+import functools
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
+from tinygrad.runtime.autogen.amd.rdna4.ins import s_barrier_wait, s_barrier_signal
 from tinygrad.renderer.amd.dsl import s, v
+from test.amd.helpers import TARGET_TO_ARCH

 def custom_add_one(A:UOp) -> UOp:
  A = A.flatten()
@@ -43,9 +45,26 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
  sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.size}"))
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

+def custom_wave_sync(A:UOp, arch:str) -> UOp:
+  # 4 waves across 1024 WG — enough to saturate a SIMD with many concurrent WGs
+  # s_sleep yields the SIMD so waves from different WGs interleave, causing barrier packet reordering
+  threads = UOp.special(128, "lidx0")
+  wg = UOp.special(1024, "gidx0")
+  insts = []
+  for _ in range(4):
+    insts.append(s_sleep(4))
+    insts += [s_barrier()] if arch == "rdna3" else [s_barrier_signal(), s_barrier_wait()]
+    insts += [s_nop(0)]*4
+  insts.append(s_endpgm())
+  sink = UOp.sink(A.base, threads, wg, arg=KernelInfo("custom_wave_sync"))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+
@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestCustomKernel(unittest.TestCase):
+  def setUp(self): self.arch = TARGET_TO_ARCH[Device["AMD"].arch]
+
  def test_simple(self):
+    if self.arch != "rdna3": self.skipTest("only rdna3")
    a = Tensor.full((16, 16), 1.).contiguous().realize()
    a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
    ei = a.schedule()[-1].lower()
@@ -55,6 +74,7 @@ class TestCustomKernel(unittest.TestCase):
    self.assertTrue((a.numpy() == 2.).all())

  def test_variable(self):
+    if self.arch != "rdna3": self.skipTest("only rdna3")
    b = Tensor.full((16, 16), 1, dtype=dtypes.uint32).contiguous().realize()
    a = Tensor.zeros_like(b).contiguous().realize()
    a = Tensor.custom_kernel(a, b, fxn=custom_add_var)[0]
@@ -63,5 +83,9 @@ class TestCustomKernel(unittest.TestCase):
      ei.run({"var":i})
      self.assertTrue((a.numpy() == 1+i).all())

+  def test_wave_sync(self):
+    if self.arch not in {"rdna3", "rdna4"}: self.skipTest("only rdna3 or rdna4")
+    Tensor.empty(1).custom_kernel(fxn=functools.partial(custom_wave_sync, arch=self.arch))[0].realize()
+
 if __name__ == "__main__":
  unittest.main()
@@ -10,7 +10,7 @@ from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
 from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
 from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART,
-                                     InstOp, InstOpRDNA4, print_packets, CDNA_WAVEEND)
+                                     InstOp, InstOpRDNA4, print_packets, CDNA_WAVEEND, CDNA_INST)
 from test.amd.helpers import TARGET_TO_ARCH

 import tinygrad
@@ -21,7 +21,7 @@ OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LD
                  InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
                  InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
                  InstOp.OTHER_GLOBAL_STORE_VADDR_128}
-OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_5}
+OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_5, InstOpRDNA4.OTHER_LDS_1, InstOpRDNA4.OTHER_LDS_2}

 # ═══════════════════════════════════════════════════════════════════════════════
 # ROCPROF DECODER
@@ -153,9 +153,10 @@ class SQTTExamplesTestBase(unittest.TestCase):
      if "gemm" not in name: continue
      with self.subTest(example=name):
        all_packets = [p for e in events for p in decode(e.blob)]
-        inst_names = [p.op.name for p in all_packets if isinstance(p, (INST, INST_RDNA4))]
-        self.assertGreater(len(inst_names), 0, f"no INST packets in {name}")
-        self.assertGreater(len([n for n in inst_names if n.startswith("JUMP")]), 0, f"no JUMP packets in {name}")
+        inst_packets = [p for p in all_packets if isinstance(p, (INST, INST_RDNA4, CDNA_INST))]
+        self.assertGreater(len(inst_packets), 0, f"no INST packets in {name}")
+        if isinstance(inst_packets[0], (INST, INST_RDNA4)):
+          self.assertGreater(len([p for p in inst_packets if p.op.name.startswith("JUMP")]), 0, f"no JUMP packets in {name}")

  expected: dict[str, list[int]] = {}  # override in subclasses
  def test_packet_counts(self):
@@ -210,23 +211,22 @@ class SQTTExamplesTestBase(unittest.TestCase):
 class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
  target = "gfx1100"
  expected = {
-    "profile_empty_run_0": [1974, 1961, 2014, 2065, 2092, 1998],
-    "profile_empty_run_1": [1979, 1972, 2019, 2070, 2097, 2003],
-    "profile_gemm_run_0": [2038, 11076, 2324, 2129, 2156, 2062],
-    "profile_gemm_run_1": [2038, 11037, 2318, 2129, 2156, 2062],
-    "profile_ops_run_0": [2038, 5070, 2078, 2129, 2156, 2062],
-    "profile_ops_run_1": [2038, 5007, 2078, 2129, 2156, 2062],
-    "profile_plus_run_0": [1979, 1979, 2030, 2070, 2097, 2003],
-    "profile_plus_run_1": [1979, 2043, 2030, 2070, 2097, 2003],
+    "profile_empty_run_0": [1880, 1867, 1920, 1971, 1998, 1904],
+    "profile_empty_run_1": [1880, 1867, 1920, 1971, 1998, 1904],
+    "profile_gemm_run_0": [3275, 3278, 2426, 2475, 2511, 2431],
+    "profile_gemm_run_1": [3264, 3268, 2420, 2469, 2504, 2401],
+    "profile_ops_run_0": [1944, 4903, 1984, 2035, 2062, 1968],
+    "profile_ops_run_1": [1944, 4918, 1984, 2035, 2062, 1968],
+    "profile_plus_run_0": [1938, 1932, 1978, 2029, 2056, 1962],
+    "profile_plus_run_1": [1891, 1874, 1931, 1982, 2009, 1915],
  }

 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"

 class TestSQTTExamplesCDNA(SQTTExamplesTestBase):
  target = "gfx950"
-  def test_decode_all_examples(self): self.skipTest("TODO: correct deltas in the timestamp packet types, first packet is REGCS_CDNA")
-  def test_gemm_has_instructions(self): self.skipTest("TODO: decode CDNA inst packets")
  def test_rocprof_wave_times_match(self): self.skipTest("TODO: requires timestamp patching")
+  def test_rocprof_inst_times_match(self): self.skipTest("TODO: requires timestamp patching")

 if __name__ == "__main__":
  unittest.main()
@@ -2,7 +2,7 @@
 import unittest, pickle
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG, OSX, getenv, temp
+from tinygrad.helpers import DEBUG, getenv, temp
 from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
 from tinygrad.viz.serve import sqtt_timeline
@@ -11,7 +11,7 @@ from test.amd.disasm import disasm
 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"

-def rocprof_inst_traces_match(sqtt, prg, target, pass_rocprof_err=False):
+def rocprof_inst_traces_match(sqtt, prg, target):
  from tinygrad.viz.serve import amd_decode
  from extra.sqtt.roc import decode as roc_decode, InstExec
  addr_table = amd_decode(prg.lib, target)
@@ -31,7 +31,7 @@ def rocprof_inst_traces_match(sqtt, prg, target, pass_rocprof_err=False):
    rocprof_inst = next(rwaves_iter[info.wave][0])
    ref_pc = rocprof_inst.pc-prg.base
    # always check pc matches
-    assert ref_pc == info.pc or pass_rocprof_err, f"pc mismatch {ref_pc}:{disasm_map[rocprof_inst.pc]} != {info.pc}:{disasm(info.inst)}"
+    assert ref_pc == info.pc, f"pc mismatch {ref_pc}:{disasm_map[rocprof_inst.pc]} != {info.pc}:{disasm(info.inst)}"
    # special handling for s_endpgm, it marks the wave completion.
    if info.inst == s_endpgm():
      completed_wave = list(rwaves_iter[info.wave].pop(0))
@@ -64,13 +64,13 @@ class TestSQTTMapBase(unittest.TestCase):

  def test_rocprof_inst_traces_match(self):
    for name, (events, kern_events, target) in self.examples.items():
+      if "sync" in name and self.target.startswith("gfx12"):
+        self.skipTest("our timestamps are off by a few cycles because rocprof patches timestamps for rdna4 barriers")
      for event in events:
        if not event.itrace: continue
        if event.kern not in kern_events: continue
        with self.subTest(example=name, kern=event.kern):
-          # rocprof OSX has a bug for sopk decoding, linux rocprof works
-          pass_rocprof_err = OSX and target == "gfx1200" and name.startswith("profile_ops")
-          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target, pass_rocprof_err)
+          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target)
          if n_waves: print(f"{name}: passed for {passed_insts} instructions across {n_waves} waves scheduled on {n_units} wave units")

  def test_sqtt_timeline(self):
@@ -78,17 +78,34 @@ class TestSQTTMapBase(unittest.TestCase):
      for event in events:
        if (p:=kern_events.get(event.kern)) is None: continue
        with self.subTest(example=name, kern=event.kern):
-          events = [e for e in sqtt_timeline(event.blob, p.lib, target) if type(e).__name__ == "ProfileRangeEvent"]
+          if not (timeline:=sqtt_timeline(event.blob, p.lib, target)): continue
+          frequency = [e.key for e in timeline if type(e).__name__ == "ProfilePointEvent" and e.name == "freq_hz"]
+          mean = sum(frequency) / len(frequency)
+          variance = sum((v - mean) ** 2 for v in frequency) / len(frequency)
+          self.assertGreater(mean, 0)
+          if DEBUG >= 2: print(f"{name:20s} SE:{event.se} {mean/1e9:.2f} GHz mean, {variance/1e18:.2f} GHz^2 variance")
+          events = [e for e in timeline if type(e).__name__ == "ProfileRangeEvent"]
          insts, execs = 0, 0
          for e in events:
            if "EXEC" in e.device:
              if "ALT" not in e.name.display_name: execs += 1
            elif "WAVE" in e.device:
              # sopk/immediates don't get ALU/MEM EXEC
-              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE"}: insts += 1
+              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL"}: insts += 1
            else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
          self.assertEqual(execs, insts)

+  def test_wave_sync(self):
+    for name, (events, kern_events, target) in self.examples.items():
+      for event in events:
+        wave_barriers = {}
+        for e in sqtt_timeline(event.blob, kern_events[event.kern].lib, target):
+          if type(e).__name__ == "ProfileRangeEvent" and e.name.display_name == "BARRIER": wave_barriers.setdefault(e.device, []).append(e)
+        if not wave_barriers: continue
+        for row, events in wave_barriers.items():
+          for e in events:
+            assert e.en-e.st > 1, f"all barriers must have a duration greater than 1, got {e}"
+
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"

 class TestSQTTMapRDNA4(TestSQTTMapBase): target = "gfx1200"
@@ -2,7 +2,7 @@ import unittest
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv
-from extra.gemm.asm.cdna.gemm import asm_gemm
+from extra.gemm.cdna_asm_gemm import asm_gemm
 from test.helpers import needs_second_gpu

 # On non CDNA4 it will only validate the Tensor.custom_kernel integration
@@ -157,7 +157,7 @@ class TestGemmLarge(unittest.TestCase):

 class TestMagicGu(unittest.TestCase):
  def test_magicgu_matches_old(self):
-    from extra.gemm.asm.cdna.asm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
+    from extra.gemm.cdna_asm_gemm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
    old_iters_args = {64: (67108864, 0), 128: (33554432, 0), 224: (613566757, 2147483656)}
    old_gemm_shapes = [
      (8192, 4096, 4096), (8192, 14336, 4096), (8192, 4096, 14336),
@@ -10,7 +10,7 @@ from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Context, Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype
-from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
+from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")

@@ -101,14 +101,14 @@ class TestDType(unittest.TestCase):
  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "skip for now")
  def test_uint_overflow(self):
    if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
-    v = dtypes.max(self.DTYPE)
+    v = self.DTYPE.max
    _test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2)
    _test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2)

  def test_dtypes_DTYPES_DICT(self):
    self.assertIn("float", DTYPES_DICT)
    self.assertIn("float32", DTYPES_DICT)
-    self.assertEqual(len(DTYPES_DICT), 26)
+    self.assertEqual(len(DTYPES_DICT), 28)
    self.assertTrue(all(isinstance(value, DType) for value in DTYPES_DICT.values()))
    self.assertTrue(all(issubclass(_to_np_dtype(value), np.generic) for value in DTYPES_DICT.values() if _to_np_dtype(value) is not None))

@@ -143,6 +143,8 @@ def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
 class TestFp8s(unittest.TestCase):
  def test_fp8e4m3_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3).dtype == dtypes.fp8e4m3
  def test_fp8e5m2_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2).dtype == dtypes.fp8e5m2
+  def test_fp8e4m3fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3fnuz).dtype == dtypes.fp8e4m3fnuz
+  def test_fp8e5m2fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2fnuz).dtype == dtypes.fp8e5m2fnuz

 class TestFp8sConversions(unittest.TestCase):
  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
@@ -169,6 +171,30 @@ class TestFp8sConversions(unittest.TestCase):
  def test_fp8e5m2_to_float(self, x):
    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2).float().item())

+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3FNUZ_MAX, max_value=FP8E4M3FNUZ_MAX))
+  def test_float_to_fp8e4m3fnuz(self, x):
+    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
+
+  def test_float_to_fp8e4m3fnuz_extreme_values(self):
+    for x in [FP8E4M3FNUZ_MAX, FP8E4M3FNUZ_MAX*1.01, -FP8E4M3FNUZ_MAX, -FP8E4M3FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
+      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
+
+  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2FNUZ_MAX, max_value=FP8E5M2FNUZ_MAX))
+  def test_float_to_fp8e5m2fnuz(self, x):
+    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
+
+  def test_float_to_fp8e5m2fnuz_extreme_values(self):
+    for x in [FP8E5M2FNUZ_MAX, FP8E5M2FNUZ_MAX*1.01, -FP8E5M2FNUZ_MAX, -FP8E5M2FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
+      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
+
+  @given(strat.integers(min_value=0, max_value=255))
+  def test_fp8e4m3fnuz_to_float(self, x):
+    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fnuz).float().item())
+
+  @given(strat.integers(min_value=0, max_value=255))
+  def test_fp8e5m2fnuz_to_float(self, x):
+    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())
+
@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
 class TestBFloat16(unittest.TestCase):
  def test_bf16_creation_numpy(self):
@@ -490,4 +516,3 @@ class TestOpsBFloat16(unittest.TestCase):

 if __name__ == '__main__':
  unittest.main()
-
@@ -52,6 +52,8 @@ class ht:
 ht.bfloat16 = ht.uint16.filter(lambda x: ((x >> 7) & 0xFF) != 0)  # filter subnormal bfloat16
 ht.fp8e4m3 = ht.uint8
 ht.fp8e5m2 = ht.uint8
+ht.fp8e4m3fnuz = ht.uint8
+ht.fp8e5m2fnuz = ht.uint8

 def universal_test(a, b, dtype, op):
  if not isinstance(op, tuple): op = (op, op)
@@ -67,7 +69,8 @@ def universal_test(a, b, dtype, op):
    if not is_dtype_supported(dtype) or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
      fe, fm = dtypes.finfo(dtype)
      atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm)
-    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype, (1e-10, 1e-7))
+    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1),
+                        dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz:(5e-1, 5e-1)}.get(dtype, (1e-10, 1e-7))
    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
  else: np.testing.assert_equal(tensor_value, numpy_value)

@@ -87,7 +90,8 @@ def universal_test_unary(a, dtype, op):
  else: tensor_value, numpy_value = op[0](ta).numpy(), op[1](ta.numpy())
  if dtype in dtypes.floats:
    atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
-      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1)}.get(dtype, (1e-6, 1e-5))
+      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1),
+      dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz: (5e-1, 5e-1)}.get(dtype, (1e-6, 1e-5))
    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
  else: np.testing.assert_equal(tensor_value, numpy_value)

@@ -155,6 +159,26 @@ class TestDTypeALU(unittest.TestCase):
  def test_emulated_fp8e5m2(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)

+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
+  def test_fp8e4m3fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
+  def test_fp8e5m2fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
+  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
+  def test_emulated_fp8e4m3fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
+  def test_emulated_fp8e5m2fnuz(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
  @given(ht.float32, strat.sampled_from(unary_operations))
  def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)

@@ -198,6 +222,30 @@ class TestDTypeALU(unittest.TestCase):
    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)

+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
+  def test_fp8e4m3fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
+  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
+  def test_fp8e5m2fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
+  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
+  def test_emulated_fp8e4m3fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
+
+  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
+  def test_emulated_fp8e5m2fnuz_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
+
  @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
  def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)

@@ -318,7 +366,7 @@ class TestDTypeALU(unittest.TestCase):

  @unittest.expectedFailure
  def test_unsafe_cast_float_to_int_failure(self):
-    val = float(dtypes.max(dtypes.int32) - 1)
+    val = float(dtypes.int32.max - 1)
    t1 = Tensor([val], dtype=dtypes.float32).cast(dtypes.int32)
    t2 = Tensor(val, dtype=dtypes.float32).cast(dtypes.int32)
    np.testing.assert_equal(t1.item(), t2.item())
@@ -479,9 +479,9 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 0., 3., -4.], 3.])
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 0., 3., -4.], [-1., -2., 3., 0.]])
    helper_test_op(None, torch.maximum, Tensor.maximum,
-                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.max(dtypes.int)], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.max], forward_only=True)
    helper_test_op(None, torch.maximum, Tensor.maximum,
-                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.min(dtypes.int)], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.min], forward_only=True)
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[True, False, False], True], forward_only=True)
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[True, False, False], [True, True, False]], forward_only=True)

@@ -496,9 +496,9 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[1., 0., 3., -4.], 3.])
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[1., 0., 3., -4.], [-1., -2., 3., 0.]])
    helper_test_op(None, torch.minimum, Tensor.minimum,
-                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.max(dtypes.int)], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.max], forward_only=True)
    helper_test_op(None, torch.minimum, Tensor.minimum,
-                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.min(dtypes.int)], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.min], forward_only=True)
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[True, False, False], True], forward_only=True)
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[True, False, False], [True, True, False]], forward_only=True)

@@ -204,13 +204,13 @@ class TestQuantizeOnnx(unittest.TestCase):
    W = Tensor(m2:=(np.random.uniform(0, 255, size=(N,N)).astype(wi))).realize()
    tg_dtype = dtypes.int8 if xi == np.int8 else dtypes.uint8
    out = (X.int().matmul(W.int())//1000)
-    if clip: out = out.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
+    if clip: out = out.clip(tg_dtype.min, tg_dtype.max)
    out = out.cast(tg_dtype)
    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
    sexec(out, opts, replace_src, run_count=1)
    tout = out.numpy()
    mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
-    if clip: mout = mout.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
+    if clip: mout = mout.clip(tg_dtype.min, tg_dtype.max)
    mout = mout.astype(xi)
    print(tout)
    print(mout)
@@ -62,7 +62,7 @@ class TestRendererFailures(unittest.TestCase):
 class TestCStyleFailures(unittest.TestCase):
  def test_inline_const_alu(self):
    # CPU doesn't use the max function
-    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.min(dtypes.int)+1))
+    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.int.min+1))
    self.assertEqual(ret[0], 1)

  def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):
@@ -75,20 +75,20 @@ class TestHelpers(unittest.TestCase):
  def test_dtype_range(self):
    for dt in core_dtypes:
      if dtypes.is_float(dt):
-        np.testing.assert_equal(dtypes.min(dt), -math.inf)
-        np.testing.assert_equal(dtypes.max(dt), math.inf)
+        np.testing.assert_equal(dt.min, -math.inf)
+        np.testing.assert_equal(dt.max, math.inf)
        np.testing.assert_equal(dt.min, -math.inf)
        np.testing.assert_equal(dt.max, math.inf)
      elif dtypes.is_int(dt):
        info = np.iinfo(_to_np_dtype(dt))
-        np.testing.assert_equal(dtypes.min(dt), info.min)
-        np.testing.assert_equal(dtypes.max(dt), info.max)
+        np.testing.assert_equal(dt.min, info.min)
+        np.testing.assert_equal(dt.max, info.max)
        np.testing.assert_equal(dt.min, info.min)
        np.testing.assert_equal(dt.max, info.max)
      else:
        assert dt == dtypes.bool, dt
-        np.testing.assert_equal(dtypes.min(dt), False)
-        np.testing.assert_equal(dtypes.max(dt), True)
+        np.testing.assert_equal(dt.min, False)
+        np.testing.assert_equal(dt.max, True)
        np.testing.assert_equal(dt.min, False)
        np.testing.assert_equal(dt.max, True)

@@ -46,9 +46,9 @@ class TestHelpers(unittest.TestCase):
    self.assertTrue((rng+2).is_increasing())

 class TestValidIdxSimplification(unittest.TestCase):
-  def check(self, load, sidx, svalid):
+  def check(self, load, sidx, svalid, extra=()):
    with Context(NOOPT=1, SPEC=0):
-      load = full_rewrite_to_sink(load.sink()).src[0]
+      load = full_rewrite_to_sink(UOp.sink(load, *extra)).src[0]
    idx, valid = load.src[0].src[1], load.src[0].src[2]
    check_uop_against_string(self, idx, sidx)
    check_uop_against_string(self, valid, svalid)
@@ -156,9 +156,12 @@ class TestValidIdxSimplification(unittest.TestCase):
    idx = (alu15*-31)+(((((alu11+218)//224)+ridx0)%30)*1568)
    valid = (ridx2<1)&(ridx1<6)
    load = get_gated_load_uop(valid, idx)
+    # prevent ridx1 and ridx2 from being shrunk
+    red = UOp(Ops.REDUCE, dtypes.float, (load, ridx1, ridx2), Ops.ADD)
    self.check(load,
      "(r0*1568)",
-      "((r2<1)&(r1<6))")
+      "((r2<1)&(r1<6))",
+      extra=(red,))

  def test_valid_becomes_const1_z3(self):
    from z3 import Ints, Solver, And, If, Not, unsat
@@ -483,5 +486,79 @@ class TestDropTrueGate(unittest.TestCase):
    # the True gate should be dropped (INDEX should only have 2 sources)
    self.assertEqual(len(result.src), 2, "True gate should be dropped from INDEX")

+class TestRangeShrink(unittest.TestCase):
+  def get_ranges(self, sink):
+    with Context(NOOPT=1, SPEC=0):
+      result = full_rewrite_to_sink(sink)
+    return [u for u in result.toposort() if u.op is Ops.RANGE]
+
+  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)
+    ranges = self.get_ranges(load.sink())
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 4)
+
+  def test_range_shrink_picks_max_guard(self):
+    # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
+    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)
+    ranges = self.get_ranges(UOp.sink(load1, load2))
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 8)
+
+  def test_range_no_shrink_guard_ge_max(self):
+    # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
+    r = Range(0, 204)
+    load = get_gated_load_uop(r < UOp.const(dtypes.index, 300), r)
+    ranges = self.get_ranges(load.sink())
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 204)
+
+  def test_range_no_shrink_when_unguarded_elsewhere(self):
+    # one load guards r < 4, but another load uses r without a gate -> no shrink
+    r = Range(0, 204)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
+    load2 = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.float.ptr(), arg=1).index(r, ptr=True),))
+    ranges = self.get_ranges(UOp.sink(load1, load2))
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 204)
+
+  def test_range_no_shrink_when_used_in_reduce(self):
+    # range used in both a gated load AND directly in the reduce expression -> no shrink
+    r = Range(0, 204)
+    gated_load = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
+    red = UOp(Ops.REDUCE, dtypes.float, (r.cast(dtypes.float) + gated_load, r), Ops.ADD)
+    ranges = self.get_ranges(red.sink())
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 204)
+
+  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)
+    ranges = self.get_ranges(load.sink())
+    self.assertEqual(len(ranges), 0)
+
+  def test_range_shrink_store_where_invalid(self):
+    # emulates mask.where(x.pad_to(mask.shape), Invalid): range should shrink accordingly
+    from tinygrad.dtype import Invalid
+    r = Range(0, 204)
+    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
+    ranges = self.get_ranges(UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(r).store((r < 4).where(x, 0)).sink())
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 4)
+
+  def test_range_shrink_store_where_invalid_flipped(self):
+    # above, but flipped
+    from tinygrad.dtype import Invalid
+    r = Range(0, 204)
+    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
+    ranges = self.get_ranges(UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(r).store((r < 4).where(0, x)).sink())
+    self.assertEqual(len(ranges), 1)
+    self.assertEqual(ranges[0].src[0].arg, 4)
+
 if __name__ == '__main__':
  unittest.main()
@@ -423,10 +423,12 @@ class TestUOpGraph(unittest.TestCase):
    d0 = UOp(Ops.PARAM, dtypes.long.ptr(), (), 0)
    ld = d0.index(ridx0.valid(ridx0<50))
    w = (ridx0<50).where(ld, 5)
-    uops = to_uops_list([w])
+    # prevent ridx0 from being shrunk
+    red = UOp(Ops.REDUCE, dtypes.long, (ridx0.cast(dtypes.long), ridx0), Ops.ADD)
+    uops = to_uops_list([w, red])
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD: assert u.src[1].arg==5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg==5

  def test_where_on_gated_load_folds_swapped_branches(self):
    ridx0 = UOp.range(100, 0)
@@ -444,10 +446,12 @@ class TestUOpGraph(unittest.TestCase):
    gate_idx = ridx0.valid((ridx0<50))
    ld = d0.index(gate_idx).cast(dtypes.float)
    w = (ridx0<50).where(ld, 5.0)
-    uops = to_uops_list([w])
+    # prevent ridx0 from being shrunk
+    red = UOp(Ops.REDUCE, dtypes.long, (ridx0.cast(dtypes.long), ridx0), Ops.ADD)
+    uops = to_uops_list([w, red])
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD: assert u.src[1].arg == 5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg == 5

  def test_where_in_store_becomes_gate(self):
    ridx0 = UOp.range(100, 0)
@@ -786,6 +786,14 @@ class TestSymbolic(unittest.TestCase):
    # div nests: y//12 -> a//2, mod nests: y%12 -> (a%2)*6+b, recombine
    self.helper_test_variable((y//12)*12 + y%12, 0, 43, "(b+a*6)")

+  def test_div_mod_recombine_after_asymmetric_fold(self):
+    a = Variable("a", 0, 7)
+    b = Variable("b", 0, 14)
+    x = a*15+b
+    # TODO: expected "(b+a*15)"
+    self.helper_test_variable((x//10)*10 + x%10, 0, 119, "(a*10+(a+b//5)//2*10+(b+a*5)%10)")
+    self.helper_test_variable((x//10)*2 + (x//5)%2, 0, 23, "(a*3+b//5)")
+
  def test_div_mod_recombine_in_additive_sum(self):
    x = Variable("x", 0, 31)
    y = Variable("y", 0, 5)
@@ -64,8 +64,8 @@ class TestVminVmaxProperties(unittest.TestCase):

    # negative mask: x & -1 could be anything since -1 has all bits set
    uop = x & -1
-    self.assertEqual(uop.vmin, dtypes.min(dtypes.int32))
-    self.assertEqual(uop.vmax, dtypes.max(dtypes.int32))
+    self.assertEqual(uop.vmin, dtypes.int32.min)
+    self.assertEqual(uop.vmax, dtypes.int32.max)

  def test_vmin_vmax_multiplication_with_variable(self):
    # vmin and vmax for multiplication with a variable
@@ -136,8 +136,8 @@ class TestVminVmaxProperties(unittest.TestCase):
    self.assertEqual(x_bool.vmin, False)
    self.assertEqual(x_bool.vmax, True)
    x_uint = x.cast(dtypes.uint)
-    self.assertEqual(x_uint.vmin, dtypes.min(dtypes.uint))
-    self.assertEqual(x_uint.vmax, dtypes.max(dtypes.uint))
+    self.assertEqual(x_uint.vmin, dtypes.uint.min)
+    self.assertEqual(x_uint.vmax, dtypes.uint.max)

  def test_vmin_vmax_invalid(self):
    i = UOp.invalid()
@@ -58,7 +58,8 @@ class TestCfg(unittest.TestCase):
    k.emit(s_endpgm())
    k.emit(s_code_end())
    ei = run_asm("diamond", k)
-    cfg = amdgpu_cfg(ei.prg.p.lib, self.arch)["data"]
+    ret = amdgpu_cfg(ei.prg.p.lib, self.arch)
+    cfg = ret["data"]
    self.assertEqual(len(cfg["blocks"]), 5)
    edge_count = sum(len(v) for v in cfg["paths"].values())
    self.assertEqual(edge_count, 5)
@@ -69,8 +70,11 @@ class TestCfg(unittest.TestCase):
    self.assertEqual(len(references["r0"]), 2)
    insts = [cfg["pc_tokens"][pc][0]["st"] for pc in references["r0"]]
    self.assertEqual(insts, ['s_mov_b32', 's_cmp_eq_u64'])
-    end_block_content = "\n".join(" ".join(t["st"] for t in cfg["pc_tokens"][pc]) for pc in list(cfg["blocks"].values())[-1])
-    self.assertEqual(end_block_content, "s_endpgm\ns_code_end (217x)")
+    end_block = [" ".join(t["st"] for t in cfg["pc_tokens"][pc]) for pc in list(cfg["blocks"].values())[-1]]
+    code_line = ret["src"].splitlines()[-1]
+    self.assertEqual(len(end_block), 2)
+    for st in [end_block[-1], code_line]:
+      assert st.startswith("s_code_end") and st.endswith("x)"), st

  def test_loop(self):
    k = Kernel(arch=Device["AMD"].arch)
@@ -280,6 +280,11 @@ class TestAssign(unittest.TestCase):
    t.uop = t.uop.after(t[:5].uop.assign(Tensor.ones(5).uop))
    np.testing.assert_allclose(t.numpy(), [1.,1.,1.,1.,1.,0.,0.,0.,0.,0.])

+  def test_assign_after_target_chain(self):
+    t = Tensor.arange(16).reshape(4, 4).permute(1, 0).contiguous()
+    t.assign(t + 100)
+    np.testing.assert_equal(t.numpy(), [[100, 104, 108, 112], [101, 105, 109, 113], [102, 106, 110, 114], [103, 107, 111, 115]])
+
  def test_assign_contiguous(self):
    b = Tensor.arange(16).reshape(4,4).contiguous().realize()
    a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1)
@@ -237,5 +237,13 @@ class TestCallSchedule(unittest.TestCase):
    self.assertIsInstance(out.shape[0], UOp)
    np.testing.assert_allclose(out[:5].numpy(), (np.arange(16*4).reshape(16, 4)[:5] * 2 + 1).astype(np.float32))

+  def test_precompile_multi_sharded(self):
+    @function(precompile=True)
+    def f(x:Tensor) -> Tensor: return x + 1
+    devs = ("CPU:0", "CPU:1")
+    a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
+    out = f(a) + 2
+    np.testing.assert_allclose(out.numpy(), np.arange(8, dtype=np.float32).reshape(4, 2) + 3)
+
 if __name__ == '__main__':
  unittest.main()
@@ -17,13 +17,15 @@ dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and is_dtype_sup

 FP8E4M3_MAX = 448.0
 FP8E5M2_MAX = 57344.0
+FP8E4M3FNUZ_MAX = 240.0
+FP8E5M2FNUZ_MAX = 57344.0

 def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float=1e-7):
  if DEBUG >= 2: print(tensor.numpy())
  try:
    assert tensor.dtype == target_dtype
    np.testing.assert_allclose(tensor.numpy(), target, rtol={dtypes.float16:1e-3, dtypes.bfloat16:1e-2,
-      dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1}.get(target_dtype, tol_target_dtype))
+      dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1, dtypes.fp8e4m3fnuz:1e-1, dtypes.fp8e5m2fnuz:5e-1}.get(target_dtype, tol_target_dtype))

  except AssertionError as e:
    raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e
@@ -3,6 +3,7 @@ import unittest
 from tinygrad.function import function
 from tinygrad import Tensor
 from tinygrad.uop.ops import UOp
+from tinygrad.helpers import GlobalCounters

 class TestFunction(unittest.TestCase):
  def test_simple(self):
@@ -335,5 +336,100 @@ class TestFunctionMulti(unittest.TestCase):
      f(x).sum().backward()
      np.testing.assert_allclose(x.grad.numpy(), expected)

+class TestFunctionTuple(unittest.TestCase):
+  def test_tuple(self, precompile=False):
+    x = Tensor.ones(3).contiguous()
+    @function(precompile=precompile)
+    def f(t:Tensor): return (t+1, t+2)
+    t1, t2 = f(x)
+    t1.realize(t2)
+    print(t1.tolist(), t2.tolist())
+    assert t1.tolist() == [2,2,2]
+    assert t2.tolist() == [3,3,3]
+  def test_tuple_precompile(self): self.test_tuple(True)
+
+class TestFunctionBackward(unittest.TestCase):
+  def test_backward_has_grad(self):
+    N = 16
+    x = Tensor.empty(N, N)
+    w1 = Tensor.empty(N, N, requires_grad=True)
+    @function
+    def f(t:Tensor, w1:Tensor): return (t@w1)
+    f(x, w1).sum().backward()
+    assert w1.grad is not None
+
+  def test_double_matmul_backward(self):
+    N = 16
+    x = Tensor.empty(N, N)
+    w1 = Tensor.empty(N, N, requires_grad=True)
+    w2 = Tensor.empty(N, N, requires_grad=True)
+    ref = Tensor.empty(N)
+
+    @function(precompile=True)
+    def f(t:Tensor, w1:Tensor, w2:Tensor): return (t@w1)@w2
+    loss = (f(x, w1, w2)-ref).square().mean().backward()
+    loss.realize(w1.grad, w2.grad)
+
+  def test_backward_single_call(self):
+    N = 4
+    x = Tensor.arange(N*N).reshape(N, N).float()
+    w1 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_()
+    w2 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_()
+    xn, w1n, w2n = x.numpy(), w1.numpy(), w2.numpy()
+    @function
+    def f(t:Tensor, w1:Tensor, w2:Tensor): return (t@w1)@w2
+    f(x, w1, w2).sum().backward()
+    assert w1.grad is not None and w2.grad is not None
+    np.testing.assert_allclose(w1.grad.numpy(), xn.T @ np.ones((N,N)) @ w2n.T, atol=1e-3)
+    np.testing.assert_allclose(w2.grad.numpy(), (xn @ w1n).T @ np.ones((N,N)), atol=1e-3)
+
+  def test_backward_precompile_backward(self):
+    N = 4
+    x = Tensor.arange(N*N).reshape(N, N).float().contiguous()
+    w1 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_().contiguous()
+    w2 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_().contiguous()
+    Tensor.realize(x, w1, w2)
+    xn, w1n, w2n = x.numpy(), w1.numpy(), w2.numpy()
+    @function(precompile=True, precompile_backward=True)
+    def f(t:Tensor, w1:Tensor, w2:Tensor): return (t@w1)@w2
+    loss = f(x, w1, w2).sum().backward()
+    assert w1.grad is not None and w2.grad is not None
+    GlobalCounters.reset()
+    Tensor.realize(loss, w1.grad, w2.grad)
+    np.testing.assert_allclose(w1.grad.numpy(), xn.T @ np.ones((N,N)) @ w2n.T, atol=1e-3)
+    np.testing.assert_allclose(w2.grad.numpy(), (xn @ w1n).T @ np.ones((N,N)), atol=1e-3)
+
+  def test_backward_precompile_backward_tuple(self):
+    N = 4
+    x = Tensor.arange(N*N).reshape(N, N).float().contiguous()
+    w1 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_().contiguous()
+    w2 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_().contiguous()
+    Tensor.realize(x, w1, w2)
+    xn, w1n, w2n = x.numpy(), w1.numpy(), w2.numpy()
+    # non-tuple reference
+    ref_w1 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_().contiguous()
+    ref_w2 = Tensor.arange(N*N).reshape(N, N).float().requires_grad_().contiguous()
+    Tensor.realize(ref_w1, ref_w2)
+    @function(precompile=True, precompile_backward=True)
+    def g(t:Tensor, w1:Tensor, w2:Tensor): return (t@w1)@w2
+    g(x, ref_w1, ref_w2).sum().backward()
+    GlobalCounters.reset()
+    Tensor.realize(ref_w1.grad, ref_w2.grad)
+    ref_ops = GlobalCounters.global_ops
+    # tuple version with intermediate — should not redo forward compute
+    @function(precompile=True, precompile_backward=True)
+    def f(t:Tensor, w1:Tensor, w2:Tensor):
+      h = t@w1
+      return (h, h@w2)
+    h, out = f(x, w1, w2)
+    loss = out.sum().backward()
+    assert w1.grad is not None and w2.grad is not None
+    GlobalCounters.reset()
+    Tensor.realize(loss, w1.grad, w2.grad)
+    np.testing.assert_allclose(w1.grad.numpy(), xn.T @ np.ones((N,N)) @ w2n.T, atol=1e-3)
+    np.testing.assert_allclose(w2.grad.numpy(), (xn @ w1n).T @ np.ones((N,N)), atol=1e-3)
+    # tuple version should have fewer ops than non-tuple (saves recomputing h=t@w1 in backward)
+    self.assertLessEqual(GlobalCounters.global_ops, ref_ops)
+
 if __name__ == '__main__':
  unittest.main()
@@ -1,27 +1,20 @@
-import math, unittest
-from unittest.mock import patch
+import unittest
 from tinygrad import Tensor
-from tinygrad.device import Device
 from tinygrad.dtype import Invalid, dtypes
-from tinygrad.helpers import unwrap_class_type
+from tinygrad.engine.realize import run_schedule

 class TestInvalidTensor(unittest.TestCase):
  def _invalid_test_helper(self, out, expected):
-    before = None
-    original_call = (runtime_cls:=unwrap_class_type(Device[Device.DEFAULT].runtime)).__call__
+    sched = out.schedule()
+    buf = out.uop.buffer
+    buf.allocate()
+    sentinel = memoryview(bytearray(b'\x42' * buf.nbytes))
+    buf.copyin(sentinel)
+    before = buf.as_memoryview().cast(out.dtype.fmt).tolist()
+    run_schedule(sched)
+    ret = buf.as_memoryview().cast(out.dtype.fmt).tolist()

-    def patched_call(self_prg, *bufs, **kwargs):
-      nonlocal before
-      before = Device[Device.DEFAULT].allocator._as_buffer(bufs[0]).cast(out.dtype.fmt).tolist()
-      return original_call(self_prg, *bufs, **kwargs)
-
-    with patch.object(runtime_cls, '__call__', patched_call): ret = out.tolist()
-
-    for i,v in enumerate(expected):
-      if v is None: assert before[i] == ret[i] or (math.isnan(before[i]) and math.isnan(ret[i]))
-      else: assert ret[i] == v
-
-    return before, ret
+    for i,v in enumerate(expected): self.assertEqual(ret[i], before[i] if v is None else v)

  def test_where_x_invalid(self):
    mask = Tensor.arange(4) < 2
@@ -37,11 +30,7 @@ class TestInvalidTensor(unittest.TestCase):
    mask = Tensor.arange(6).reshape(2, 3) < 3
    vals = Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
    out = mask.where(vals, Invalid)
-    before, ret = self._invalid_test_helper(out, [])
-    assert ret[0] == [1.0, 2.0, 3.0]
-    assert before[3] == ret[1][0] or (math.isnan(before[3]) and math.isnan(ret[1][0]))
-    assert before[4] == ret[1][1] or (math.isnan(before[4]) and math.isnan(ret[1][1]))
-    assert before[5] == ret[1][2] or (math.isnan(before[5]) and math.isnan(ret[1][2]))
+    self._invalid_test_helper(out, [1.0, 2.0, 3.0, None, None, None])

  def test_where_invalid_int(self):
    mask = Tensor.arange(3) < 2
@@ -89,8 +78,7 @@ class TestInvalidTensor(unittest.TestCase):
  def test_where_reduce_always_true(self):
    mask = Tensor.arange(4) < 9
    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).sum()
-    before, ret = self._invalid_test_helper(out, [])
-    assert ret == 10.0
+    self._invalid_test_helper(out, [10.0])

  def test_invalid_unary(self):
    mask = Tensor.arange(4) < 2
@@ -110,10 +98,7 @@ class TestInvalidTensor(unittest.TestCase):
  def test_invalid_reshape(self):
    mask = Tensor.arange(4) < 2
    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).reshape(2,2)
-    before, ret = self._invalid_test_helper(out, [])
-    assert ret[0] == [1.0, 2.0]
-    assert ret[1][0] == before[2] or (math.isnan(ret[1][0]) and math.isnan(before[2]))
-    assert ret[1][1] == before[3] or (math.isnan(ret[1][1]) and math.isnan(before[3]))
+    self._invalid_test_helper(out, [1.0, 2.0, None, None])

  def test_invalid_cast(self):
    mask = Tensor.arange(4) < 2
@@ -36,8 +36,7 @@ class TestSetitemInto(unittest.TestCase):
    self.assertEqual(GlobalCounters.kernel_count, 0)
    t.realize()
    self.assertEqual(GlobalCounters.kernel_count, 1)
-    # TODO: this can be just 4 if empty goes through is_realized setitem path
-    self.assertEqual(GlobalCounters.global_mem, 4*(3*2+1)) # 3 elements had +1, 1 is assigned directly
+    self.assertEqual(GlobalCounters.global_mem, 4)
    t[1].realize()
    t.realize()
    self.assertEqual(GlobalCounters.kernel_count, 1)
@@ -19,7 +19,7 @@ from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_in
  ReduceContext, correct_load_store, pm_render, pm_add_loads
 from tinygrad.codegen.opt.postrange import apply_opts, pm_make_images
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
-from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar
+from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar, pm_store_ranges
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.renderer.amd.elf import do_assemble_amd

@@ -31,7 +31,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
  if SPEC: type_verify(sink, kernel_spec)

  # preprocess
-  sink = graph_rewrite(sink, pm_mops+pm_syntactic_sugar, name="early movement ops", bottom_up=True)
+  sink = graph_rewrite(sink, pm_mops+pm_syntactic_sugar+pm_store_ranges, ctx=itertools.count(1000), name="early movement ops", bottom_up=True)

  # first we optimize
  if optimize:
@@ -48,7 +48,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
    sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")

    # optimize (schedule) the AST
-    sink = graph_rewrite(sink, pm_simplify_ranges, name="simplify ranges")
+    sink = graph_rewrite(sink, pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges")

    # do postrange optimization, BEAM or hand_coded_optimizations
    sink = apply_opts(sink, ren)
@@ -190,9 +190,9 @@ def _do_image_fixup(dt:ImageDType, idx:UOp) -> tuple[UOp, UOp, int, int]:
  if IMAGE == 1 and valid is not None:
    h, w = max(ImageDType.valid_dims(dt), key=lambda hw:
                # maximize number of valids removed
-               (len(_drop_valid_stmts(valid, idx:=uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw)),
+               (len(_drop_valid_stmts(valid, idx:=uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1])).simplify()), *hw)),
                # and minimize idx complexity (number of nodes)
-                -len(idx.backward_slice)))
+                -len(idx.gep(1).backward_slice)))
    buf = buf.replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4), w * 4 * dt.itemsize))
  oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % w, (x // (4*w))))
  return x, idx.replace(src=(buf, oidx.valid(valid))), w, h
@@ -1,4 +1,5 @@
 import itertools
+from typing import Callable
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import partition
@@ -36,22 +37,32 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
          u = nidx
  return u

+def mark_gated(ctx, idx):
+  if idx.src[1].op is Ops.WHERE:
+    x, cond = idx.src[1].get_idx(), idx.src[1].get_valid()
+    # get all ranges r with guards "r < c" for some const c
+    guards = {r:c for v in cond.split_uop(Ops.AND) if v.op is Ops.CMPLT and (r:=v.src[0]).op is Ops.RANGE and (c:=v.src[1]).op is Ops.CONST}
+  else: x, guards = idx, {}
+  # ensure that we choose max(c_i) for all i where r < c_i
+  ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].arg < c.arg)}
+  # but if a range is ever ungated, we cannot shrink it
+  ctx |= {r:r.src[0] for r in x.ranges if r not in guards}
+
 pm_simplify_ranges = PatternMatcher([
  (UPat((Ops.END, Ops.REDUCE), name="u"), simplify_merge_adjacent),
+  (UPat(Ops.INDEX, name="idx"), mark_gated),
+  # reduce ranges can't be shrunk
+  (UPat(Ops.REDUCE, name="red"), lambda ctx, red: ctx.update({r:r.src[0] for r in red.src[1:]})),
+  (UPat(Ops.SINK, name="x"), lambda ctx, x: do_substitute(ctx, x, lambda r,c: r.replace(src=(c,)))),
 ])

 def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None:
  if r not in ctx and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c

-def do_substitute(ctx:dict[UOp, UOp|None], x: UOp) -> UOp|None:
-  subs = {}
-  for k,v in ctx.items():
-    if v is not None:
-      subs[k] = k.replace(src=(k.src[0]//v,), arg=k.arg[0:-1]+(0,k.arg[-1]))*v + k.replace(src=(v,), arg=k.arg[0:-1]+(1,k.arg[-1]))
-  if not len(subs): return None
-  ret = x.substitute(subs).simplify()
+def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None:
+  ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None})
  ctx.clear()
-  return ret
+  return None if ret is x else ret.simplify()

 def dont_sub_ranges_for_image(ctx:dict[UOp, UOp|None], x:UOp) -> None:
  if isinstance(x.src[0].src[0].dtype, ImageDType):
@@ -60,7 +71,8 @@ def dont_sub_ranges_for_image(ctx:dict[UOp, UOp|None], x:UOp) -> None:
 pm_split_ranges = PatternMatcher([
  (UPat(Ops.RANGE, name="r")%UPat.cvar("c"), mark_range_mod),
  (UPat(Ops.STORE, name="x"), dont_sub_ranges_for_image),
-  (UPat(Ops.SINK, name="x"), do_substitute),
+  (UPat(Ops.SINK, name="x"), lambda ctx, x: do_substitute(ctx, x,
+    lambda k,v: k.replace(src=(k.src[0]//v,), arg=k.arg[0:-1]+(0,k.arg[-1]))*v + k.replace(src=(v,), arg=k.arg[0:-1]+(1,k.arg[-1])))),
 ])

 # **** reduce simplification ****
@@ -353,11 +353,12 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
    if device == "NV": return not CI and not NV_PTX and not NV_NAK
    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not CPU_LVP
    return device in {"AMD", "CL", "PYTHON", "NULL"}
-  if dtype in dtypes.fp8s:
+  if dtype in dtypes.fp8_ocp:
    if device == "CUDA": return not CI and not CUDA_PTX
    if device == "NV": return not CI and not NV_PTX and not NV_NAK
-    if device == "AMD": return not CI and getattr(Device["AMD"], "target") in {(9,4,2), (9,5,0)}
+    if device == "AMD": return not CI and getattr(Device["AMD"], "target") == (9,5,0)
    return device in {"PYTHON", "NULL"}
+  if dtype in dtypes.fp8_fnuz: return device in {"PYTHON", "NULL"}
  if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
                                          dtypes.ushort, dtypes.float, dtypes.int32, dtypes.uint32, dtypes.half]
  # for CI GPU and OSX, cl_khr_fp16 isn't supported
@@ -79,10 +79,14 @@ class DType(metaclass=DTypeMetaClass):
    return PtrDType(self.priority, self.bitsize, self.name, self.fmt, self.count, None, self, addrspace, 1, size)
  def scalar(self) -> DType: return self._scalar if self._scalar is not None else self
  def nbytes(self) -> int: raise RuntimeError("only ptr types have nbytes")
-  @property
-  def min(self): return dtypes.min(self)
-  @property
-  def max(self): return dtypes.max(self)
+  @functools.cached_property
+  def min(self):
+    if dtypes.is_int(self): return 0 if dtypes.is_unsigned(self) else -2**(self.scalar().bitsize-1)
+    return -float("inf") if dtypes.is_float(self) else False
+  @functools.cached_property
+  def max(self):
+    if dtypes.is_int(self): return 2**(self.scalar().bitsize)-1+self.min
+    return float("inf") if dtypes.is_float(self) else True

@dataclass(frozen=True, eq=False)
 class PtrDType(DType):
@@ -172,21 +176,11 @@ class dtypes:
    # int is the default. wrap floats in ConstFloat to distinguish -0.0 from 0.0 in cache
    return ConstFloat(float(val)) if dtypes.is_float(dtype) else bool(val) if dtypes.is_bool(dtype) else int(val)
  @staticmethod
-  @functools.cache
-  def min(dtype:DType):
-    if dtypes.is_int(dtype): return 0 if dtypes.is_unsigned(dtype) else -2**(dtype.scalar().bitsize-1)
-    return -float("inf") if dtypes.is_float(dtype) else False
-  @staticmethod
-  @functools.cache
-  def max(dtype:DType):
-    if dtypes.is_int(dtype): return 2**(dtype.scalar().bitsize)-1+dtypes.min(dtype)
-    return float("inf") if dtypes.is_float(dtype) else True
-  @staticmethod
  def finfo(dtype:DType) -> tuple[int, int]:
    """(exponent, mantissa)"""
    if not dtypes.is_float(dtype): raise ValueError(f"{dtype} is not a floating point type")
    return {dtypes.float16: (5, 10), dtypes.bfloat16: (8, 7), dtypes.float32: (8, 23), dtypes.float64: (11, 52),
-            dtypes.fp8e5m2: (5, 2), dtypes.fp8e4m3: (4, 3)}[dtype]
+            dtypes.fp8e4m3: (4, 3), dtypes.fp8e5m2: (5, 2), dtypes.fp8e4m3fnuz: (4, 3), dtypes.fp8e5m2fnuz: (5, 2)}[dtype]
  void: Final[DType] = DType.new(-1, 0, "void", None)
  index: Final[DType] = DType.new(-1, 800, "index", None)
  bool: Final[DType] = DType.new(0, 1, "bool", '?')
@@ -202,6 +196,8 @@ class dtypes:
  _uint256: Final[DType] = DType.new(8, 256, "uint256", None)
  fp8e4m3: Final[DType] = DType.new(9, 8, "float8_e4m3", None)
  fp8e5m2: Final[DType] = DType.new(10, 8, "float8_e5m2", None)
+  fp8e4m3fnuz: Final[DType] = DType.new(9, 8, "float8_e4m3fnuz", None)
+  fp8e5m2fnuz: Final[DType] = DType.new(10, 8, "float8_e5m2fnuz", None)
  float16: Final[DType] = DType.new(11, 16, "half", 'e')
  # bfloat16 has higher priority than float16, so least_upper_dtype(dtypes.int64, dtypes.uint64) = dtypes.float16
  bfloat16: Final[DType] = DType.new(12, 16, "__bf16", None)
@@ -222,7 +218,9 @@ class dtypes:
  default_float: ClassVar[DType] = float32
  default_int: ClassVar[DType] = int32

-  fp8s = (fp8e4m3, fp8e5m2)
+  fp8_ocp = (fp8e4m3, fp8e5m2)
+  fp8_fnuz = (fp8e4m3fnuz, fp8e5m2fnuz)
+  fp8s = fp8_ocp + fp8_fnuz
  floats = fp8s + (float16, bfloat16, float32, float64)
  int8s = (uint8, int8)
  int16s = (uint16, int16)
@@ -244,8 +242,9 @@ def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType)
 # we don't support weak type and complex type
 promo_lattice = { dtypes.bool: [dtypes.int8, dtypes.uint8], dtypes.int8: [dtypes.int16], dtypes.int16: [dtypes.int32], dtypes.int32: [dtypes.int64],
  dtypes.int64: [dtypes.uint64], dtypes.uint8: [dtypes.int16, dtypes.uint16], dtypes.uint16: [dtypes.int32, dtypes.uint32],
-  dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.fp8e4m3, dtypes.fp8e5m2],
-  dtypes.fp8e5m2: [dtypes.float16, dtypes.bfloat16], dtypes.fp8e4m3: [dtypes.float16, dtypes.bfloat16],
+  dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.fp8e4m3, dtypes.fp8e5m2, dtypes.fp8e4m3fnuz, dtypes.fp8e5m2fnuz],
+  dtypes.fp8e4m3: [dtypes.float16, dtypes.bfloat16], dtypes.fp8e5m2: [dtypes.float16, dtypes.bfloat16],
+  dtypes.fp8e4m3fnuz: [dtypes.float16, dtypes.bfloat16], dtypes.fp8e5m2fnuz: [dtypes.float16, dtypes.bfloat16],
  dtypes.float16: [dtypes.float32], dtypes.bfloat16: [dtypes.float32], dtypes.float32: [dtypes.float64], }

@functools.cache
@@ -300,10 +299,14 @@ def float_to_bf16(x):
 _fp8_cfg = {
  dtypes.fp8e4m3: (7, 4, 0x7, 0x3F50000000000000, 0x407D000000000000, 0x7E, 0x3F90000000000000),
  dtypes.fp8e5m2: (15, 3, 0x3, 0x3EE0000000000000, 0x40EE000000000000-1, 0x7B, 0x3F10000000000000),
+  dtypes.fp8e4m3fnuz: (8, 4, 0x7, 0x3F40000000000000, 0x406F000000000000-1, 0x7F, 0x3F80000000000000),
+  dtypes.fp8e5m2fnuz: (16, 3, 0x3, 0x3ED0000000000000, 0x40EE000000000000-1, 0x7F, 0x3F00000000000000),
 }

 def float_to_fp8(x: float, dtype: DType) -> int:
  assert dtype in dtypes.fp8s, "Only for fp8s"
+  if dtype in dtypes.fp8_fnuz and not math.isfinite(x): return 0x80
+  if dtype in dtypes.fp8_fnuz and x == 0.0: return 0x00
  # e4m3 don't support inf, return 0x7f(+NaN) and 0xff(-NaN) to match jax
  # NaN is unordered, can't compare with zero, use math.copysign to get sign
  if dtype == dtypes.fp8e4m3 and not math.isfinite(x): return 0x7f if math.copysign(1, x) > 0 else 0xff
@@ -323,16 +326,17 @@ def float_to_fp8(x: float, dtype: DType) -> int:
    res, half = mantissa >> shift, half_ulp << shift
    round_bits = (xbits | (1 << 52)) & ((half << 1) - 1)
    if round_bits > half or (round_bits == half and res & 1): res += 1
-  return int(res | sign)
+  return 0 if dtype in dtypes.fp8_fnuz and res == 0 else int(res | sign)  # fnuz has no negative zero

 def fp8_to_float(x: int, dtype: DType) -> float:
  assert dtype in dtypes.fp8s, "Only for fp8s"
+  if dtype in dtypes.fp8_fnuz and x == 0x80: return math.nan
  if (x & 0x7F) == 0: return -0.0 if x & 0x80 else 0.0
  bias, sig_bits, *_ = _fp8_cfg[dtype]
  mant_bits, exp_bits = sig_bits - 1, 8 - sig_bits
  exp_max, mant_max = (1 << exp_bits) - 1, (1 << mant_bits) - 1
  sign, exp, mantissa = (x >> 7) & 1, (x >> mant_bits) & exp_max, x & mant_max
-  if exp == exp_max:
+  if dtype not in dtypes.fp8_fnuz and exp == exp_max:
    if dtype == dtypes.fp8e5m2: return math.copysign(math.nan if mantissa else math.inf, -1 if sign else 1)
    if mantissa == mant_max: return math.nan
  val = (mantissa / (mant_max + 1)) * 2 ** (1 - bias) if exp == 0 else (1 + mantissa / (mant_max + 1)) * 2 ** (exp - bias)
@@ -46,18 +46,19 @@ def _buffer_like(u:UOp) -> UOp:
    if prod(dtype.shape) != prod(u.max_shard_shape) or ([x for x in u.max_shard_shape if x != 1] or [1])[-1] % 4 != 0:
      if DEBUG >= 1: print(f"demoting Image {dtype} with shape {u.max_shard_shape}")
      dtype = dtype.base
-  buffer = UOp.new_buffer(u.device, u.shard_size, dtype).reshape(u.max_shard_shape)
+  buffer = UOp.new_buffer(u.device, u.shard_size, dtype).reshape(u.max_shard_shape).shrink_to(u.shard_shape)
  if isinstance(u.device, tuple) and u.axis is not None: buffer = buffer.multi(u.axis)
  return buffer

-def replace_contig_with_assign(u:UOp):
+def replace_contig_with_store_after(u:UOp):
  # can't allocate a buffer without a device (e.g., inside a CALL function body with only PARAMs)
  if u._device is None: return None
  # if size is 0, remove the contig
  if u.size == 0: return u.src[0]
  # no real contig for DISK/TINYFS tensors, they are left alone
  if isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS")): return u.rtag(None)
-  return _buffer_like(u).assign(u.src[0]).rtag(u.tag)
+  buf = _buffer_like(u)
+  return buf.after(buf.store(u.src[0])).rtag(u.tag)

 def replace_assign_with_contig(u:UOp):
  assigned_to = u
@@ -93,9 +94,23 @@ def contiguous_mops_to_view(c:UOp):
 def transform_precompiled_call(c:UOp) -> UOp|None:
  if not c.arg.precompile: return None
  if c.src[0].op is Ops.SINK: return None
-  out = _buffer_like(c)
  input_buffers = tuple(x.contiguous() if x.op not in {Ops.AFTER, Ops.BIND} else x for x in c.src[1:])
-  fxn = out.param_like(len(c.src)-1).assign(c.src[0]).sink()
+  # multi-output (TUPLE) precompiled calls: allocate a buffer per output and return a TUPLE of buffers
+  if c.src[0].op is Ops.TUPLE:
+    tuple_body = c.src[0]
+    out_bufs = []
+    sink_srcs = []
+    for i, elem in enumerate(tuple_body.src):
+      buf = UOp.new_buffer(c.device, prod(elem.max_shape), elem.dtype).reshape(elem.max_shape).shrink_to(elem.shape)
+      out_bufs.append(buf)
+      target = buf.param_like(len(c.src) - 1 + i).shrink_to(elem.shape)
+      sink_srcs.append(target.after(target.store(elem)))
+    fxn = UOp.sink(*sink_srcs)
+    new_call = c.replace(src=(fxn, *input_buffers, *out_bufs), dtype=dtypes.void, tag=None)
+    return UOp.maketuple(*[buf.after(new_call) for buf in out_bufs])
+  out = _buffer_like(c)
+  target = out.param_like(len(c.src)-1).shrink_to(c.shape)
+  fxn = target.after(target.store(c.src[0])).sink()
  ret = out.after(c.replace(src=(fxn, *input_buffers, out), dtype=dtypes.void, tag=None))
  # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
  if any(isinstance(s, UOp) for s in c.shape): ret = ret.shrink(tuple((0, s) for s in c.shape))
@@ -110,14 +125,15 @@ pm_early_transform_tensor_graph = PatternMatcher([
  (UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Movement),), name="c"), contiguous_mops_to_view),

  # add CONTIGUOUS to tagged UOps
-  (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.ASSIGN}, name="x"), lambda x: x.rtag(None).contiguous(tag=x.tag) if x.tag else x.replace(tag=None)),
-  # remove extra CONTIGUOUS on ASSIGN (only when assign target is contiguous)
-  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.ASSIGN, name="a"),), name="c"),
+  (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.ASSIGN, Ops.AFTER, Ops.STORE}, name="x"),
+   lambda x: x.rtag(None).contiguous(tag=x.tag) if x.tag else x.replace(tag=None)),
+  # remove extra CONTIGUOUS on ASSIGN/AFTER (only when target is contiguous)
+  (UPat(Ops.CONTIGUOUS, src=(UPat({Ops.ASSIGN, Ops.AFTER}, name="a"),), name="c"),
   lambda a,c: a.replace(tag=(a.tag or ())+(c.tag or ())) if a.src[0].has_buffer_identity() else None),
  # replace ASSIGN with CONTIGUOUS
  (UPat(Ops.ASSIGN, name="u"), replace_assign_with_contig),
-  # replace CONTIGUOUS with ASSIGNs
-  (UPat(Ops.CONTIGUOUS, name="u"), replace_contig_with_assign),
+  # replace CONTIGUOUS with STORE+AFTER
+  (UPat(Ops.CONTIGUOUS, name="u"), replace_contig_with_store_after),
  # remove DETACH/CONTIGUOUS_BACKWARD (allows more contiguous removal)
  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
 ])
@@ -128,9 +144,9 @@ def untag_and_append(ctx:AllocCtx, x:UOp):
  for t in x.tag:
    original_uop: UOp = ctx.uop_list[t]
    replace_uop = ret
-    while replace_uop.op is Ops.ASSIGN: replace_uop = replace_uop.src[0]
+    while replace_uop.op in {Ops.ASSIGN, Ops.AFTER}: replace_uop = replace_uop.src[0]
    ctx.buffer_map[original_uop] = replace_uop.shrink_to(original_uop.shape)
-  ctx.assigns.append(ret)
+  if ret.op is not Ops.AFTER: ctx.assigns.append(ret)  # AFTER gets appended by append_after
  return ret

 def append_after(ctx:AllocCtx, x:UOp):
@@ -142,7 +158,7 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp):
                   b._min_max if b.op is Ops.BIND else None, b.src[0].arg[0] if b.op is Ops.BIND else None)

 pm_finalize_call = PatternMatcher([
-  (UPat(Ops.ASSIGN, name="x"), untag_and_append),
+  (UPat({Ops.ASSIGN, Ops.AFTER}, name="x"), untag_and_append),
  (UPat(Ops.AFTER, name="x"), append_after),
  (UPat(Ops.COPY, name="x"), lambda ctx,x: append_after(ctx,x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
  # remove unique from const. TODO: this is copied in function.py
@@ -2,6 +2,7 @@ import functools
 from typing import Generic, TypeVar, Callable, cast, overload
 from tinygrad.helpers import Context, dedup, getenv
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
+from tinygrad.gradient import compute_gradient
 from tinygrad.tensor import Tensor

 def add_to_ctx(ctx, x:UOp):
@@ -19,9 +20,10 @@ pm_ctx = PatternMatcher([

 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):
-  def __init__(self, fxn:Callable[..., ReturnType], *, precompile:bool=False):
+  def __init__(self, fxn:Callable[..., ReturnType], *, precompile:bool=False, precompile_backward:bool=False):
    self.fxn = fxn
    self.precompile = precompile
+    self.precompile_backward = precompile_backward

  def __get__(self, obj, objtype=None): return functools.partial(self.__call__, obj) if obj is not None else self

@@ -39,12 +41,17 @@ class _function(Generic[ReturnType]):
    # run it and do surgery later
    with Context(ALLOW_DEVICE_USAGE=getenv("DEVICE_IN_FUNCTION_BUG", 0)):
      ret = self.fxn(*args, **kwargs)
-    assert isinstance(ret, Tensor), "only supports one tensor return for now"
+    if isinstance(ret, Tensor):
+      uret = ret.uop
+    elif isinstance(ret, tuple) and all(isinstance(x, Tensor) for x in ret):
+      uret = UOp.maketuple(*[x.uop for x in ret])
+    else:
+      raise RuntimeError(f"function return type {type(ret)} not supported")

    # replace the known inputs with params (using deduplicated slots)
    subs = {}
    for i,x in enumerate(call_uops): subs[x] = x.param_like(i)
-    uret = ret.uop.substitute(subs)
+    uret = uret.substitute(subs)

    # add contiguous to call_uops
    #call_uops = [x.contiguous() for x in call_uops]
@@ -60,14 +67,77 @@ class _function(Generic[ReturnType]):
    #call = assigned.call(*call_uops, buffer, name=name)
    #ret = buffer.after(call)

-    ret = uret.call(*call_uops, name=name, precompile=self.precompile)
-    return cast(ReturnType, Tensor(ret, device=ret.device))
+    # precompute the backward: determine which inputs need gradients and build the backward CALL body
+    grad_fxn = None
+    inputs = [t for t in list(args)+[kwargs[k] for k in sorted(kwargs)] if isinstance(t, (Tensor, UOp))]
+    grad_params = {x.arg:x for x in uret.toposort(enter_calls=False) if x.op == Ops.PARAM}
+    # find which param slots correspond to requires_grad inputs
+    need_grad = {i for i, t in enumerate(inputs) if isinstance(t, Tensor) and t.requires_grad}
+    target_params = {grad_params[i] for i in need_grad if i in grad_params}
+    if target_params:
+      grad_fxn = self._make_grad_fxn(uret, len(call_uops), target_params, need_grad, name, self.precompile_backward)
+
+    fret = uret.call(*call_uops, name=name, precompile=self.precompile, precompile_backward=self.precompile_backward, grad_fxn=grad_fxn)
+    if isinstance(ret, tuple):
+      return cast(ReturnType, tuple(Tensor(fret.gettuple(i), device=fret.device) for i in range(len(ret))))
+    else:
+      return cast(ReturnType, Tensor(fret, device=fret.device))
+
+  @staticmethod
+  def _make_grad_fxn(uret:UOp, num_args:int, target_params:set[UOp], need_grad:set[int], name:str, precompile_backward:bool):
+    def grad_fxn(ctx, k):
+      fxn, args = k.src[0], k.src[1:]
+      params = {x.arg:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
+      # compute gradients only for needed params
+      if isinstance(ctx, dict):
+        all_grads: dict[UOp, UOp] = {}
+        for idx, grad_out in ctx.items():
+          elem_grads = compute_gradient(fxn.src[idx], grad_out.param_like(len(args) + idx), target_params)
+          for p, g in elem_grads.items():
+            if p in all_grads: all_grads[p] = all_grads[p] + g
+            else: all_grads[p] = g
+        grads = all_grads
+        grad_ctx_inputs = tuple(ctx.get(i, fxn.src[i].const_like(0)) for i in range(len(fxn.src)))
+      else:
+        grads = compute_gradient(fxn, ctx.param_like(len(args)), target_params)
+        grad_ctx_inputs = (ctx,)
+      # collect gradients for needed params only
+      grad_indices: list[int] = []
+      grad_uops: list[UOp] = []
+      for i in range(len(args)):
+        if i in need_grad and (p:=params.get(i)) is not None and p in grads:
+          grad_indices.append(i)
+          grad_uops.append(grads[p])
+      if len(grad_uops) == 0: return (None,) * len(args)
+      # replace forward output references with PARAMs to avoid recomputation
+      if fxn.op is Ops.TUPLE:
+        fwd_subs = {elem: elem.param_like(len(args) + len(grad_ctx_inputs) + i) for i, elem in enumerate(fxn.src)}
+        fwd_inputs = tuple(k.gettuple(i) for i in range(len(fxn.src)))
+      else:
+        fwd_subs = {fxn: fxn.param_like(len(args) + 1)}
+        fwd_inputs = (k,)
+      grad_uops = [g.substitute(fwd_subs) for g in grad_uops]
+      # build a single backward CALL returning a TUPLE of all gradients
+      bwd_body = UOp.maketuple(*grad_uops)
+      bwd_call = bwd_body.call(*args, *grad_ctx_inputs, *fwd_inputs, name=name+"_backward", precompile=precompile_backward)
+      # extract each gradient via GETTUPLE
+      ret: list[UOp|None] = []
+      gi = 0
+      for i in range(len(args)):
+        if gi < len(grad_indices) and grad_indices[gi] == i:
+          ret.append(bwd_call.gettuple(gi))
+          gi += 1
+        else:
+          ret.append(None)
+      return tuple(ret)
+    return grad_fxn

 # overload signatures support both @function and @function(precompile=True) syntax
@overload
-def function(fxn:Callable[..., ReturnType], *, precompile:bool=False) -> _function[ReturnType]: ...
+def function(fxn:Callable[..., ReturnType], *, precompile:bool=False, precompile_backward:bool=False) -> _function[ReturnType]: ...
@overload
-def function(fxn:None=None, *, precompile:bool=False) -> Callable[[Callable[..., ReturnType]], _function[ReturnType]]: ...
-def function(fxn=None, *, precompile:bool=False):
-  if fxn is None: return lambda f: _function(f, precompile=precompile)
-  return _function(fxn, precompile=precompile)
+def function(fxn:None=None, *, precompile:bool=False, precompile_backward:bool=False) -> \
+  Callable[[Callable[..., ReturnType]], _function[ReturnType]]: ...
+def function(fxn=None, *, precompile:bool=False, precompile_backward:bool=False):
+  if fxn is None: return lambda f: _function(f, precompile=precompile, precompile_backward=precompile_backward)
+  return _function(fxn, precompile=precompile, precompile_backward=precompile_backward)
@@ -13,18 +13,53 @@ def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
    return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
  if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)

-def call_gradient(ctx:UOp, k:UOp) -> tuple[UOp|None, ...]:
-  if k.arg.grad_fxn is not None: return (None,) + k.arg.grad_fxn(ctx, k)
+def call_gradient(ctx:UOp|dict[int, UOp], k:UOp) -> tuple[UOp|None, ...]:
+  if k.arg.grad_fxn is not None:
+    if isinstance(ctx, dict): return (None,) + k.arg.grad_fxn(ctx, k)
+    return (None,) + k.arg.grad_fxn(ctx, k)
  # auto-differentiate the function
  fxn, args = k.src[0], k.src[1:]
  params = {x.arg:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
-  grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params.values()))
-  ret: list[UOp|None] = [None]
+  # for tuple-returning CALLs, ctx is a dict mapping output index to gradient; differentiate each output separately and sum
+  if isinstance(ctx, dict):
+    all_grads: dict[UOp, UOp] = {}
+    for idx, grad_out in ctx.items():
+      elem_grads = compute_gradient(fxn.src[idx], grad_out.param_like(len(args) + idx), set(params.values()))
+      for p, g in elem_grads.items():
+        if p in all_grads: all_grads[p] = all_grads[p] + g
+        else: all_grads[p] = g
+    grads = all_grads
+    grad_ctx_inputs = tuple(ctx.get(i, fxn.src[i].const_like(0)) for i in range(len(fxn.src)))
+  else:
+    grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params.values()))
+    grad_ctx_inputs = (ctx,)
+  # collect which args have gradients
+  grad_indices: list[int] = []
+  grad_uops: list[UOp] = []
  for i in range(len(args)):
    if (p:=params.get(i, None)) is not None and p in grads:
-      # TODO: compact the args and remove unused ones
      assert not grads[p].op_in_backward_slice_with_self(Ops.BUFFER), "BUG: BUFFER in backward slice of grad"
-      ret.append(grads[p].call(*args, ctx, name=(k.arg.name or "")+f"_backward_{i}"))
+      grad_indices.append(i)
+      grad_uops.append(grads[p])
+  if len(grad_uops) == 0: return (None,) * (len(args) + 1)
+  # replace forward output references with PARAMs, passing the forward CALL output(s) as inputs to avoid recomputation
+  if fxn.op is Ops.TUPLE:
+    fwd_subs = {elem: elem.param_like(len(args) + len(grad_ctx_inputs) + i) for i, elem in enumerate(fxn.src)}
+    fwd_inputs = tuple(k.gettuple(i) for i in range(len(fxn.src)))
+  else:
+    fwd_subs = {fxn: fxn.param_like(len(args) + 1)}
+    fwd_inputs = (k,)
+  grad_uops = [g.substitute(fwd_subs) for g in grad_uops]
+  # build a single backward CALL returning a TUPLE of all gradients
+  bwd_body = UOp.maketuple(*grad_uops)
+  bwd_call = bwd_body.call(*args, *grad_ctx_inputs, *fwd_inputs, name=(k.arg.name or "")+"_backward", precompile=k.arg.precompile_backward)
+  # extract each gradient via GETTUPLE
+  ret: list[UOp|None] = [None]
+  gi = 0
+  for i in range(len(args)):
+    if gi < len(grad_indices) and grad_indices[gi] == i:
+      ret.append(bwd_call.gettuple(gi))
+      gi += 1
    else:
      ret.append(None)
  return tuple(ret)
@@ -72,11 +107,26 @@ def _deepwalk(root:UOp, targets:set[UOp]) -> list[UOp]:
  return list(root.toposort(lambda node: node.op not in {Ops.DETACH, Ops.ASSIGN} and in_target_path[node]))

 def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp]:
-  grads = {root: root_grad}
+  grads: dict[UOp, UOp] = {root: root_grad}
+  # for GETTUPLE nodes on tuple-returning CALLs, collect per-output gradients
+  tuple_call_grads: dict[UOp, dict[int, UOp]] = {}
  for t0 in reversed(_deepwalk(root, targets)):
-    if t0 not in grads: continue
-    lgrads: tuple[UOp|None, ...]|None = cast(tuple[UOp|None, ...]|None, pm_gradient.rewrite(t0, ctx=grads[t0]))
-    if lgrads is None: raise RuntimeError(f"failed to compute gradient for {t0.op}\n\nin {str(t0)[0:1000]}...")
+    if t0 not in grads and t0 not in tuple_call_grads: continue
+    # for tuple-returning CALLs, use accumulated per-output gradients
+    if t0.op is Ops.CALL and t0 in tuple_call_grads:
+      lgrads = cast(tuple[UOp|None, ...], call_gradient(tuple_call_grads[t0], t0))
+    elif t0 not in grads:
+      continue
+    # for GETTUPLE on a CALL, accumulate gradient per output index instead of propagating to CALL directly
+    elif t0.op is Ops.GETTUPLE and t0.src[0].op is Ops.CALL:
+      call = t0.src[0]
+      if call not in tuple_call_grads: tuple_call_grads[call] = {}
+      if t0.arg in tuple_call_grads[call]: tuple_call_grads[call][t0.arg] = tuple_call_grads[call][t0.arg] + grads[t0]
+      else: tuple_call_grads[call][t0.arg] = grads[t0]
+      continue
+    else:
+      lgrads = cast(tuple[UOp|None, ...]|None, pm_gradient.rewrite(t0, ctx=grads[t0]))
+      if lgrads is None: raise RuntimeError(f"failed to compute gradient for {t0.op}\n\nin {str(t0)[0:1000]}...")
    assert len(lgrads) == len(t0.src), f"got {len(lgrads)} gradient, expected {len(t0.src)}"
    for k,v in zip(t0.src, lgrads):
      if v is None: continue
@@ -424,7 +424,9 @@ def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip

 def system(cmd:str, **kwargs) -> str:
  st = time.perf_counter()
-  ret = subprocess.check_output(cmd.split(), **kwargs).decode().strip()
+  try: ret = subprocess.check_output(cmd.split(), stderr=subprocess.STDOUT, **kwargs).decode().strip()
+  except subprocess.CalledProcessError as e:
+    raise RuntimeError(f"system: '{cmd}' failed with exit code {e.returncode}\n{(e.output or b'').decode().strip()}") from e
  if DEBUG >= 1: print(f"system: '{cmd}' returned {len(ret)} bytes in {(time.perf_counter() - st)*1e3:.2f} ms")
  return ret

@@ -303,7 +303,7 @@ class RMSNorm:
    x = self._norm(x.float()).cast(x.dtype)
    return x if self.weight is None else x * self.weight

-from tinygrad.uop.ops import UOp, KernelInfo, Ops
+from tinygrad.uop.ops import UOp, KernelInfo, Ops, AxisType
 def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
  weight, idx = call.src[1:]
  # for multi-device: unshard inputs to one device
@@ -326,15 +326,29 @@ def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
  # this is the real atomic kernel
  def _embedding_bwd_kernel(grad_weight:UOp, grad_emb:UOp, idx:UOp) -> UOp:
    idx_flat, grad_emb_flat = idx.flatten(), grad_emb.reshape((idx.size, grad_weight.shape[-1]))
+
    i = UOp.range(grad_emb_flat.shape[0], 0)  # batch_size * sequence_length
    j = UOp.range(grad_emb_flat.shape[1], 1)  # embed_size
+
+    embed_size = grad_weight.shape[-1]
+    BLOCK_J = min(256, embed_size)
+    assert embed_size % BLOCK_J == 0, f"embed_size {embed_size} must be divisible by {BLOCK_J}"
+
+    n_j_blocks = embed_size // BLOCK_J
+    i = UOp.range(grad_emb_flat.shape[0], 0)         # batch_size * sequence_length -> GLOBAL
+    j_inner = UOp.range(BLOCK_J, 2, AxisType.LOOP if device in ("CPU", "NULL") else AxisType.LOCAL)  # BLOCK_J threads per workgroup
+    j_outer = UOp.range(n_j_blocks, 1)
+    j = j_outer * BLOCK_J + j_inner
+
    token_id = idx_flat[i].clip(0, grad_weight.shape[0]-1).cast(dtypes.index)
+
    # atomic scatter-add: grad_weight[token_id, j] += grad_emb_flat[i, j]
    if device in ("CPU", "NULL"): atomic_arg = "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
    elif device == "AMD": atomic_arg = "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
    else: raise NotImplementedError(f"no atomics for device {device}")
    atomic = UOp(Ops.CUSTOM, dtypes.void, (grad_weight.index(token_id, j, ptr=True), grad_emb_flat[i, j].cast(dtypes.float)), arg = atomic_arg)
-    return atomic.end(i, j).sink(arg=KernelInfo(name="embedding_bwd", opts_to_apply=()))
+    return atomic.end(i, j_outer, j_inner).sink(arg=KernelInfo(name="embedding_bwd", opts_to_apply=()))
+
  grad_weight_uop = grad_weight_uop.custom_kernel(grad_emb, idx, fxn=_embedding_bwd_kernel)[0]

  return (grad_weight_uop.cast(weight.dtype), None)
@@ -497,7 +497,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    o_ = [((i - 1) // s + 1) for i,s in zip(i_, s_)]
    return _onnx_pads_to_tiny_pads(_auto_pad([(o-1)*s+k-i for o,i,k,s in zip(o_, i_, k_, s_)], auto_pad))

-  def _clamp_cast(x:Tensor, dtype:DType): return x.clamp(dtypes.min(dtype), dtypes.max(dtype)).cast(dtype)
+  def _clamp_cast(x:Tensor, dtype:DType): return x.clamp(dtype.min, dtype.max).cast(dtype)

  def _prepare_quantize(x:Tensor, scale:Tensor, zero_point:Tensor|int, axis=1, block_size=0):
    if axis < 0: axis += x.ndim
@@ -1209,7 +1209,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT

  def DynamicQuantizeLinear(x: Tensor):
    # only support uint8
-    qmin, qmax = dtypes.min(dtypes.uint8), dtypes.max(dtypes.uint8)
+    qmin, qmax = dtypes.uint8.min, dtypes.uint8.max
    scale = (x.max().maximum(0) + ((-x).max()).maximum(0)) / (qmax - qmin)
    zero_point = _clamp_cast((qmin - x.min() / scale).round(), dtypes.uint8)
    y = _clamp_cast((x / scale).round() + zero_point, dtypes.uint8)
@@ -126,6 +126,9 @@ class InstOpRDNA4(Enum):
  LDS_WR_3 = 0x2c
  LDS_WR_4 = 0x2d
  LDS_WR_5 = 0x2e
+  OTHER_LDS_1 = 0x50
+  OTHER_LDS_2 = 0x51
+  BARRIER_SIGNAL = 0x7a
  WMMA_8 = 0x8c
  WMMA_16 = 0x8d
  VALU_DPFP = 0x92
@@ -381,25 +384,27 @@ PACKET_TYPES_RDNA4: dict[int, type[PacketType]] = {
 # CDNA PACKET TYPE DEFINITIONS
 # ═══════════════════════════════════════════════════════════════════════════════

-class CDNA_DELTA(PacketType):
-  """pkt_fmt=0: 16-bit timestamp delta packet"""
+class CDNA_MISC(PacketType):
+  """pkt_fmt=0: 16-bit (Misc)"""
  encoding = bits[3:0] == 0
-  delta = bits[11:4]      # (data >> 4) & 0xff
-  unk_0 = bits[12:12]     # (data >> 0xc) & 1
-  unk_1 = bits[15:13]     # (data >> 0xd)
+  delta = bits[11:4]
+  sh = bits[12:12]
+  misc_type = bits[15:13]

 class CDNA_TIMESTAMP(PacketType):
  """pkt_fmt=1: 64-bit timestamp packet (case 0x0)"""
  encoding = bits[3:0] == 1
-  unk_0 = bits[15:4]
+  _reserved = bits[15:4]
  timestamp = bits[63:16]   # stored as (data_word >> 0x10) in low 46 bits of local_58

-class CDNA_PKT_2(PacketType):
-  """pkt_fmt=2: 64-bit packet (case 0x4)"""
+class CDNA_REG(PacketType):
+  """pkt_fmt=2: 64-bit (Reg)"""
  encoding = bits[3:0] == 2
-  unk_0 = bits[6:5]       # (data >> 5) & 3
-  unk_1 = bits[7:7]       # (data >> 7) + 1 & 1
-  unk_padding = bits[63:8]
+  pipe = bits[6:5]
+  _me_raw = bits[8:7]
+  _reserved = bits[15:9]
+  regaddr = bits[31:16]
+  regdata = bits[63:32]

 class CDNA_WAVESTART(PacketType):
  """type 3: 32-bit wave start (Wave/group_id)"""
@@ -410,19 +415,19 @@ class CDNA_WAVESTART(PacketType):
  simd = bits[15:14]
  pipe = bits[17:16]
  me = bits[19:18]
-  _gap = bits[21:20]
+  _reserved = bits[21:20]
  count = bits[28:22]
  _padding = bits[31:29]

-class CDNA_PKT_4(PacketType):
-  """pkt_fmt=4: 16-bit packet (case 0xc, same as 0x8/0x14)"""
+class CDNA_WAVEALLOC(PacketType):
+  """pkt_fmt=4: 16-bit (Wave)"""
  encoding = bits[3:0] == 4
-  unk_0 = bits[5:5]       # (data_word >> 5) & 1
-  unk_1 = bits[9:6]       # (data_word >> 6) & 0xf
-  unk_2 = bits[13:10]     # (data_word >> 10) & 0xf
-  unk_3 = bits[15:14]     # (data_word >> 0xe)
+  sh = bits[5:5]
+  cu = bits[9:6]
+  wave = bits[13:10]
+  simd = bits[15:14]

-class REGCS_CDNA(PacketType):
+class CDNA_REG_CS(PacketType):
  """type 5: 48-bit register CS write (RegCs)"""
  encoding = bits[3:0] == 5
  pipe = bits[6:5]
@@ -438,80 +443,86 @@ class CDNA_WAVEEND(PacketType):
  wave = bits[13:10]
  simd = bits[15:14]

-class CDNA_EXEC(PacketType):
-  """pkt_fmt=10: 16-bit EXEC packet (case 0x24)"""
-  encoding = bits[3:0] == 10
-  unk_0 = bits[8:5]       # (data_word >> 5) & 0xf
-  unk_1 = bits[10:9]      # (data_word >> 9) & 3
-  unk_2 = bits[15:11]     # (data_word >> 0xb)
-
-class CDNA_PKT_11(PacketType):
-  """pkt_fmt=11: 64-bit packet (case 0x28)"""
-  encoding = bits[3:0] == 11
-  unk_0 = bits[8:5]       # (data_word >> 5) & 0xf
-  unk_1 = bits[10:9]      # (data_word >> 9) & 3
-  unk_2 = bits[15:15]     # (data_word >> 0xf) & 1
-  unk_padding = bits[63:16]
-
 class CDNA_INST(PacketType):
-  """pkt_fmt=13: 32-bit INST packet (case 0x30)"""
+  """pkt_fmt=10: 16-bit (MsgInst)"""
+  encoding = bits[3:0] == 10
+  wave = bits[8:5]
+  simd = bits[10:9]
+  inst_type = bits[15:11]
+
+class CDNA_INST_PC(PacketType):
+  """pkt_fmt=11: 64-bit (MsgInstPc)"""
+  encoding = bits[3:0] == 11
+  wave = bits[8:5]
+  simd = bits[10:9]
+  _reserved = bits[14:11]
+  err = bits[15:15]
+  pc = bits[63:16]
+
+class CDNA_ISSUE(PacketType):
+  """pkt_fmt=13: 32-bit (Issue)"""
  encoding = bits[3:0] == 13
-  unk_0 = bits[6:5]       # (data >> 5) & 3
-  unk_1 = bits[9:8]       # (data >> 8) & 3
-  unk_2 = bits[11:10]     # (data >> 10) & 3
-  unk_3 = bits[13:12]     # (data >> 0xc) & 3
-  unk_4 = bits[15:14]     # (data >> 0xe) & 3
-  unk_5 = bits[19:18]     # (data >> 0x12) & 3
-  unk_6 = bits[21:20]     # (data >> 0x14) & 3
-  unk_7 = bits[23:22]     # (data >> 0x16) & 3
-  unk_8 = bits[25:24]     # (data >> 0x18) & 3
-  unk_9 = bits[27:26]     # (data >> 0x1a) & 3
-  unk_padding = bits[31:28]
+  simd = bits[6:5]
+  _gap = bits[7:7]
+  inst0 = bits[9:8]
+  inst1 = bits[11:10]
+  inst2 = bits[13:12]
+  inst3 = bits[15:14]
+  inst4 = bits[17:16]
+  inst5 = bits[19:18]
+  inst6 = bits[21:20]
+  inst7 = bits[23:22]
+  inst8 = bits[25:24]
+  inst9 = bits[27:26]
+  _padding = bits[31:28]

-class CDNA_PKT_14(PacketType):
-  """pkt_fmt=14: 64-bit packet (case 0x34)"""
+class CDNA_PERF(PacketType):
+  """pkt_fmt=14: 64-bit (MsgPerf)"""
  encoding = bits[3:0] == 14
-  unk_0 = bits[5:5]       # (data >> 5) & 1
-  unk_1 = bits[9:6]       # (data >> 6) & 0xf
-  unk_2 = bits[11:10]     # (data >> 10) & 3
-  unk_3 = bits[24:12]     # (data >> 0xc) & 0x1fff
-  unk_4 = bits[37:25]     # (data >> 0x19) & 0x1fff
-  unk_5 = bits[50:38]     # (data >> 0x26) & 0x1fff
-  unk_6 = bits[51:51]     # (data >> 0x33) & 1
-  unk_padding = bits[63:52]
+  sh = bits[5:5]
+  cu = bits[9:6]
+  cntr_bank = bits[11:10]
+  cntr0 = bits[24:12]
+  cntr1 = bits[37:25]
+  cntr2 = bits[50:38]
+  cntr3 = bits[63:51]

-class CDNA_PKT_7(PacketType):
-  """pkt_fmt=7: 16-bit packet"""
+class CDNA_EVENT(PacketType):
+  """pkt_fmt=7: 16-bit"""
  encoding = bits[3:0] == 7
-  unk_padding = bits[15:4]
+  _reserved = bits[15:4]

-class CDNA_PKT_8(PacketType):
-  """pkt_fmt=8: 16-bit packet"""
+class CDNA_EVENT_CS(PacketType):
+  """pkt_fmt=8: 16-bit"""
  encoding = bits[3:0] == 8
-  unk_padding = bits[15:4]
+  _reserved = bits[15:4]

-class CDNA_PKT_9(PacketType):
-  """pkt_fmt=9: 16-bit packet"""
+class CDNA_EVENT_GFX1(PacketType):
+  """pkt_fmt=9: 16-bit"""
  encoding = bits[3:0] == 9
-  unk_padding = bits[15:4]
+  _reserved = bits[15:4]

-class CDNA_PKT_12(PacketType):
-  """pkt_fmt=12: 48-bit packet"""
+class CDNA_USERDATA(PacketType):
+  """pkt_fmt=12: 48-bit (UserData)"""
  encoding = bits[3:0] == 12
-  unk_padding = bits[47:4]
+  sh = bits[5:5]
+  cu = bits[9:6]
+  wave = bits[13:10]
+  simd = bits[15:14]
+  data = bits[47:16]

-class CDNA_PKT_15(PacketType):
-  """pkt_fmt=15: 48-bit packet (case 0x38, same as 0x10)"""
+class CDNA_REG_CS_PRIV(PacketType):
+  """pkt_fmt=15: 48-bit (RegCs)"""
  encoding = bits[3:0] == 15
-  unk_0 = bits[6:5]       # (data >> 5) & 3
-  unk_1 = bits[7:7]       # (data >> 7) + 1 & 1
-  unk_2 = bits[15:9]      # (data >> 9) & 0x7f
-  unk_padding = bits[47:16]
+  pipe = bits[6:5]
+  _me_raw = bits[8:7]
+  regaddr = bits[15:9]
+  regdata = bits[47:16]

 PACKET_TYPES_CDNA: dict[int, type[PacketType]] = {
-  0: CDNA_DELTA, 1: CDNA_TIMESTAMP, 2: CDNA_PKT_2, 3: CDNA_WAVESTART, 4: CDNA_PKT_4, 5: REGCS_CDNA, 6: CDNA_WAVEEND,
-  7: CDNA_PKT_7, 8: CDNA_PKT_8, 9: CDNA_PKT_9, 10: CDNA_EXEC, 11: CDNA_PKT_11, 12: CDNA_PKT_12,
-  13: CDNA_INST, 14: CDNA_PKT_14, 15: CDNA_PKT_15,
+  0: CDNA_MISC, 1: CDNA_TIMESTAMP, 2: CDNA_REG, 3: CDNA_WAVESTART, 4: CDNA_WAVEALLOC, 5: CDNA_REG_CS, 6: CDNA_WAVEEND,
+  7: CDNA_EVENT, 8: CDNA_EVENT_CS, 9: CDNA_EVENT_GFX1, 10: CDNA_INST, 11: CDNA_INST_PC, 12: CDNA_USERDATA,
+  13: CDNA_ISSUE, 14: CDNA_PERF, 15: CDNA_REG_CS_PRIV, 16: LAYOUT_HEADER,
 }

 # ═══════════════════════════════════════════════════════════════════════════════
@@ -523,8 +534,8 @@ def _build_decode_tables(packet_types: dict[int, type[PacketType]]) -> tuple[dic
  sorted_types = sorted(packet_types.items(), key=lambda x: (-bin(x[1].encoding.mask).count('1'), x[0] == 16))
  state_table = bytes(next((op for op, cls in sorted_types if (b & cls.encoding.mask) == cls.encoding.default), 16) for b in range(256))
  # Build decode info: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
-  # special_case: 0=none, 1=TS_DELTA_OR_MARK (check is_marker), 2=TS_DELTA_SHORT (add 8), 3=CDNA_DELTA (*4), 4=CDNA_TIMESTAMP (absolute)
-  _special = {TS_DELTA_OR_MARK: 1, TS_DELTA_OR_MARK_RDNA4: 1, TS_DELTA_SHORT: 2, CDNA_DELTA: 3, CDNA_TIMESTAMP: 4}
+  # special_case: 0=none, 1=TS_DELTA_OR_MARK (check is_marker), 2=TS_DELTA_SHORT (add 8), 3=CDNA_MISC (*4), 4=CDNA_TIMESTAMP (absolute)
+  _special = {TS_DELTA_OR_MARK: 1, TS_DELTA_OR_MARK_RDNA4: 1, TS_DELTA_SHORT: 2, CDNA_MISC: 3, CDNA_TIMESTAMP: 4}
  decode_info = {}
  for opcode, pkt_cls in packet_types.items():
    delta_field = getattr(pkt_cls, 'delta', None)
@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import cast, ClassVar
+from typing import cast
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
@@ -11,7 +11,7 @@ from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, Profil
 from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, AMD_HIPCC, ceildiv, unwrap
 from tinygrad.renderer.cstyle import AMDHIPRenderer, AMDHIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
-from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt, amdgpu_kd, amdgpu_drm
+from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
@@ -824,19 +824,16 @@ class KFDIface:
  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return ((self.drm_dev_info.cu_bitmap[se % 4][sa + (se // 4) * 2] >> (2 * wgp)) & 0x3) == 0x3

 class PCIIface(PCIIfaceBase):
-  gpus:ClassVar[list[str]] = []
-
  def __init__(self, dev, dev_id):
    super().__init__(dev, dev_id, vendor=0x1002, devices=[(0xffff, [0x74a1, 0x744c, 0x7480, 0x7550, 0x7590, 0x75a0])], bars=[0, 2, 5], vram_bar=0,
      va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size)
    self._setup_adev(self.pci_dev)
-    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)

  def require_profile_mode(self): return True
  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.

-  def _setup_adev(self, pci_dev:PCIDevice, dma_regions:list[tuple[int, MMIOInterface]]|None=None):
-    self.dev_impl:AMDev = AMDev(pci_dev, dma_regions)
+  def _setup_adev(self, pci_dev:PCIDevice):
+    self.dev_impl:AMDev = AMDev(pci_dev)
    self.ip_versions = self.dev_impl.ip_ver

    gfxver = int(f"{self.dev_impl.ip_ver[am.GC_HWIP][0]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][1]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][2]:02d}")
@@ -894,7 +891,7 @@ class PCIIface(PCIIfaceBase):
 class USBIface(PCIIface):
  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
    self.dev, self.pci_dev = dev, USBPCIDevice(dev.__class__.__name__[:2], f"usb:{dev_id}", bars=[0, 2, 5])
-    self._setup_adev(self.pci_dev, dma_regions=[(0x200000, self.pci_dev.dma_view(0xf000, 0x80000))])
+    self._setup_adev(self.pci_dev)
    self.pci_dev.usb._pci_cacheable += [(self.pci_dev.bar_info[2].addr, self.pci_dev.bar_info[2].size)] # doorbell region is cacheable

    # special regions
@@ -905,7 +902,7 @@ class USBIface(PCIIface):
    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
    return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)

-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
    if (host or (uncached and cpu_access)) and self.sys_next_off + size < self.sys_buf.size:
      self.sys_next_off += size
      return self.sys_buf.offset(self.sys_next_off - size, size)
@@ -1,7 +1,7 @@
 from __future__ import annotations
 import os, ctypes, contextlib, re, functools, mmap, struct, array, sys, weakref
 assert sys.platform != 'win32'
-from typing import cast, ClassVar
+from typing import cast
 from dataclasses import dataclass
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQProgram, HCQSignal, BumpAllocator
 from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU, hcq_filter_visible_devices, hcq_profile
@@ -11,7 +11,7 @@ from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, pr
 from tinygrad.helpers import ContextVar, VIZ, ProfileEvent
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
-from tinygrad.runtime.autogen import nv_570, nv_580, pci, mesa
+from tinygrad.runtime.autogen import nv_570, nv_580, mesa
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.nv.nvdev import NVDev, NVMemoryManager
 from tinygrad.runtime.support.system import System, PCIIfaceBase, MAP_FIXED
@@ -535,8 +535,6 @@ class NVKIface:
  def sleep(self, tm:int): pass

 class PCIIface(PCIIfaceBase):
-  gpus:ClassVar[list[str]] = []
-
  def __init__(self, dev, dev_id):
    # PCIIface's MAP_FIXED mmap will overwrite UVM allocations made by NVKIface, so don't try PCIIface if kernel driver was already used.
    if NVKIface.root is not None: raise RuntimeError("Cannot use PCIIface after NVKIface has been initialized (would corrupt UVM memory)")
@@ -544,7 +542,6 @@ class PCIIface(PCIIfaceBase):
      base_class=0x03, bars=[0, 1, 3], vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
    if not OSX: System.reserve_hugepages(64)

-    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
    self.dev_impl:NVDev = NVDev(self.pci_dev)
    self.root, self.gpu_instance = 0xc1000000, 0
    self.rm_alloc(0, nv_gpu.NV01_ROOT, nv_gpu.NV0000_ALLOC_PARAMETERS())
@@ -553,11 +550,11 @@ class PCIIface(PCIIfaceBase):
    self.gpfifo_class, self.compute_class, self.dma_class = (gsp:=self.dev_impl.gsp).gpfifo_class, gsp.compute_class, gsp.dma_class
    self.viddec_class = None

-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, **kwargs) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
    # Force use of huge pages for large allocations. NVDev will attempt to use huge pages in any case,
    # but if the size is not aligned, the tail will be allocated with 4KB pages, increasing TLB pressure.
-    page_size = mmap.PAGESIZE if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << 10))
-    return super().alloc(round_up(size, page_size), host=host, uncached=uncached, cpu_access=cpu_access, contiguous=contiguous, **kwargs)
+    return super().alloc(round_up(size, mmap.PAGESIZE if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << 10))),
+      host=host, uncached=uncached, cpu_access=cpu_access, contiguous=contiguous, force_devmem=force_devmem, **kwargs)

  def setup_usermode(self): return 0xce000000, self.pci_dev.map_bar(bar=0, fmt='I', off=0xbb0000, size=0x10000)
  def setup_vm(self, vaspace): pass
@@ -1,8 +1,8 @@
 from __future__ import annotations
 import ctypes, collections, dataclasses, functools, hashlib, array
 from tinygrad.helpers import mv_address, getenv, DEBUG, fetch, lo32, hi32
+from tinygrad.runtime.autogen import pci
 from tinygrad.runtime.autogen.am import am
-from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.amd import AMDReg, import_module, import_asic_regs
 from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace
 from tinygrad.runtime.support.system import PCIDevice, PCIDevImplBase
@@ -146,8 +146,8 @@ class AMMemoryManager(MemoryManager):
 class AMDev(PCIDevImplBase):
  Version = 0xA0000008

-  def __init__(self, pci_dev:PCIDevice, dma_regions:list[tuple[int, MMIOInterface]]|None=None, reset_mode=False):
-    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
+  def __init__(self, pci_dev:PCIDevice, reset_mode=False):
+    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
    self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')

    self._run_discovery()
@@ -185,6 +185,7 @@ class AMDev(PCIDevImplBase):

    # Re-initialize main blocks
    self.init_hw(self.gfx, self.sdma)
+    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)

    self.smu.set_clocks(level=-1) # last level, max perf.
    for ip in [self.soc, self.gfx]: ip.set_clockgating_state()
@@ -547,10 +547,10 @@ class AM_PSP(AM_IP):
  def init_sw(self):
    self.reg_pref = "regMP0_SMN_C2PMSG" if self.adev.ip_ver[am.MP0_HWIP] < (14,0,0) else "regMPASP_SMN_C2PMSG"

-    msg1_region = next((reg for reg in self.adev.dma_regions or [] if reg[1].nbytes >= (512 << 10)), None)
-    if msg1_region is not None:
-      self.msg1_addr, self.msg1_view = self.adev.mm.alloc_vaddr(size=msg1_region[1].nbytes, align=am.PSP_1_MEG), msg1_region[1]
-      self.adev.mm.map_range(self.msg1_addr, msg1_region[1].nbytes, [(msg1_region[0],msg1_region[1].nbytes)], AddrSpace.SYS, uncached=True, boot=True)
+    if self.adev.devfmt.startswith("usb:"):
+      self.msg1_view, paddrs = self.adev.pci_dev.alloc_sysmem(512 << 10)
+      self.msg1_addr = self.adev.mm.alloc_vaddr(size=self.msg1_view.nbytes, align=am.PSP_1_MEG)
+      self.adev.mm.map_range(self.msg1_addr, self.msg1_view.nbytes, [(paddrs[0], self.msg1_view.nbytes)], AddrSpace.SYS, uncached=True, boot=True)
    else:
      self.msg1_paddr = self.adev.mm.palloc(am.PSP_1_MEG, align=am.PSP_1_MEG, zero=False, boot=True)
      self.msg1_addr, self.msg1_view = self.adev.paddr2mc(self.msg1_paddr), self.adev.vram.view(self.msg1_paddr, am.PSP_1_MEG, 'B')
@@ -8,6 +8,7 @@ def _do_ioctl(__idir, __base, __nr, __struct, __fd, *args, __payload=None, **kwa
  assert not WIN, "ioctl not supported"
  import tinygrad.runtime.support.hcq as hcq, fcntl
  ioctl = __fd.ioctl if isinstance(__fd, hcq.FileIOInterface) else functools.partial(fcntl.ioctl, __fd)
+  if __struct is None: return ioctl((__base<<8)|__nr, __payload or (args[0] if args else 0))
  if (rc:=ioctl((__idir<<30)|(ctypes.sizeof(out:=(__payload or __struct(*args, **kwargs)))<<16)|(__base<<8)|__nr, out)):
    raise RuntimeError(f"ioctl returned {rc}")
  return out
@@ -1,6 +1,7 @@
 from __future__ import annotations
 import ctypes, time, functools, re, gzip, struct
 from tinygrad.helpers import getenv, DEBUG, fetch, getbits
+from tinygrad.runtime.autogen import pci
 from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace
 from tinygrad.runtime.support.nv.ip import NV_FLCN, NV_FLCN_COT, NV_GSP
 from tinygrad.runtime.support.system import PCIDevice, PCIDevImplBase, MMIOInterface
@@ -72,6 +73,7 @@ class NVMemoryManager(MemoryManager):
 class NVDev(PCIDevImplBase):
  def __init__(self, pci_dev:PCIDevice):
    self.pci_dev, self.devfmt, self.mmio = pci_dev, pci_dev.pcibus, pci_dev.map_bar(0, fmt='I')
+    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)

    self.smi_dev, self.is_booting, self.is_err_state = False, True, False
    self._early_ip_init()
@@ -1,10 +1,9 @@
 from __future__ import annotations
 import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys, itertools, struct, socket, subprocess, time, enum
-from typing import ClassVar
-from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv, unwrap, fetch, system
+from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv, unwrap, fetch, system, _ensure_downloads_dir
 from tinygrad.runtime.autogen import libc, pci, vfio, iokit, corefoundation
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer, hcq_filter_visible_devices
-from tinygrad.runtime.support.memory import MemoryManager, VirtMapping, AddrSpace
+from tinygrad.runtime.support.memory import MemoryManager, VirtMapping, AddrSpace, BumpAllocator
 from tinygrad.runtime.support.usb import ASM24Controller, USBMMIOInterface

 MAP_FIXED, MAP_FIXED_NOREPLACE = 0x10, 0x100000
@@ -44,6 +43,11 @@ class _System:

  def reserve_hugepages(self, cnt): os.system(f"sudo sh -c 'echo {cnt} > /proc/sys/vm/nr_hugepages'")

+  @functools.cache
+  def reserve_va(self, va_start, va_size):
+    # cached, runs only once per range. used to not collide with other mappings.
+    FileIOInterface.anon_mmap(va_start, va_size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE | MAP_FIXED_NOREPLACE, 0)
+
  def memory_barrier(self): lib.atomic_thread_fence(__ATOMIC_SEQ_CST:=5) if (lib:=self.libsys if OSX else self.atomic_lib) is not None else None

  def lock_memory(self, addr:int, size:int):
@@ -72,6 +76,13 @@ class _System:

    return sorted([val for vndr, device, val in all_devs if vndr == vendor and any((device & mask) in devlist for mask, devlist in devices)])

+  def pci_probe_device(self, devpref:str, dev_id:int, vendor:int, devices:list[tuple[int, list[int]]], bars:list[int],
+                       resize_bars:list[int]|None=None, base_class:int|None=None):
+    gpus = hcq_filter_visible_devices(System.pci_scan_bus(vendor, devices, base_class))
+    if not gpus: raise RuntimeError("No supported GPUs found")
+    if OSX: return APLRemotePCIDevice(devpref, f'usb4:{dev_id}', bars)
+    return PCIDevice(devpref, gpus[dev_id], bars=bars, resize_bars=resize_bars)
+
  def pci_setup_usb_bars(self, usb:ASM24Controller, gpu_bus:int, mem_base:int, pref_mem_base:int) -> dict[int, PCIBarInfo]:
    for bus in range(gpu_bus):
      # All 3 values must be written at the same time.
@@ -172,8 +183,8 @@ class PCIDevice:
      self.irq_poller.register(self.irq_fd.fd, select.POLLIN)

      irqs = vfio.struct_vfio_irq_set(index=vfio.VFIO_PCI_MSI_IRQ_INDEX, flags=vfio.VFIO_IRQ_SET_DATA_EVENTFD|vfio.VFIO_IRQ_SET_ACTION_TRIGGER,
-        argsz=ctypes.sizeof(vfio.struct_vfio_irq_set), count=1, data=(ctypes.c_int * 1)(self.irq_fd.fd))
-      vfio.VFIO_DEVICE_SET_IRQS(self.vfio_dev, irqs)
+        argsz=ctypes.sizeof(vfio.struct_vfio_irq_set) + ctypes.sizeof(ctypes.c_int), count=1)
+      vfio.VFIO_DEVICE_SET_IRQS(self.vfio_dev, (ctypes.c_byte * irqs.argsz).from_buffer(bytearray(bytes(irqs)) + struct.pack('i', self.irq_fd.fd)))
    else: FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/enable", os.O_RDWR).write("1")

    self.cfg_fd = FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/config", os.O_RDWR | os.O_SYNC | os.O_CLOEXEC)
@@ -201,9 +212,14 @@ class USBPCIDevice(PCIDevice):
    self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
    self.usb = ASM24Controller()
    self.pcibus, self.bar_info = pcibus, System.pci_setup_usb_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30))
+    self.sram = BumpAllocator(size=0x80000, wrap=False) # asm24 controller sram
+  def read_config(self, offset:int, size:int): return self.usb.pcie_cfg_req(offset, bus=4, dev=0, fn=0, size=size)
+  def write_config(self, offset:int, value:int, size:int): self.usb.pcie_cfg_req(offset, bus=4, dev=0, fn=0, value=value, size=size)
  def map_bar(self, bar, off=0, addr=0, size=None, fmt='B'):
    return USBMMIOInterface(self.usb, self.bar_info[bar].addr + off, size or self.bar_info[bar].size, fmt)
  def dma_view(self, ctrl_addr, size): return USBMMIOInterface(self.usb, ctrl_addr, size, fmt='B', pcimem=False)
+  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
+    return self.dma_view(0xf000 + (off:=self.sram.alloc(size)), size), [0x200000 + off]

 class PCIDevImplBase:
  mm: MemoryManager
@@ -211,22 +227,20 @@ class PCIDevImplBase:
@dataclasses.dataclass
 class PCIAllocationMeta: mapping:VirtMapping; has_cpu_mapping:bool; hMemory:int=0 # noqa: E702

-class LNXPCIIfaceBase:
+class PCIIfaceBase:
  dev_impl:PCIDevImplBase
-  gpus:ClassVar[list[str]] = []
+
+  def is_local(self) -> bool: return not isinstance(self.pci_dev, RemotePCIDevice)
+  def is_bar_small(self) -> bool: return self.pci_dev.bar_info[self.vram_bar].size == (256 << 20)

  def __init__(self, dev, dev_id, vendor, devices:list[tuple[int, list[int]]], bars, vram_bar, va_start, va_size, base_class:int|None=None):
-    if len((cls:=type(self)).gpus) == 0:
-      cls.gpus = hcq_filter_visible_devices(System.pci_scan_bus(vendor, devices, base_class))
-
-      # Acquire va range to avoid collisions.
-      FileIOInterface.anon_mmap(va_start, va_size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE | MAP_FIXED_NOREPLACE, 0)
-    self.pci_dev, self.dev, self.vram_bar = PCIDevice(dev.__class__.__name__[:2], cls.gpus[dev_id], bars=bars, resize_bars=[vram_bar]), dev, vram_bar
+    self.pci_dev = System.pci_probe_device(dev.__class__.__name__[:2], dev_id, vendor, devices, bars, resize_bars=[vram_bar], base_class=base_class)
+    if self.is_local(): System.reserve_va(va_start, va_size)
+    self.dev, self.vram_bar = dev, vram_bar
    self.p2p_base_addr = self.pci_dev.bar_info[vram_bar].addr

  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
-    # NOTE: logic on macos is different, since bar is small
-    should_use_sysmem = host or ((cpu_access if OSX else (uncached and cpu_access)) and not force_devmem)
+    should_use_sysmem = host or ((cpu_access if self.is_bar_small() else (uncached and cpu_access)) and not force_devmem)
    if should_use_sysmem:
      vaddr = self.dev_impl.mm.alloc_vaddr(size:=round_up(size, mmap.PAGESIZE), align=mmap.PAGESIZE)
      memview, paddrs = self.pci_dev.alloc_sysmem(size, vaddr=vaddr, contiguous=contiguous)
@@ -240,14 +254,16 @@ class LNXPCIIfaceBase:
  def free(self, b:HCQBuffer):
    for dev in b.mapped_devs[1:]: dev.iface.dev_impl.mm.unmap_range(b.va_addr, b.size)
    if b.meta.mapping.aspace is AddrSpace.PHYS: self.dev_impl.mm.vfree(b.meta.mapping)
-    if b.owner == self.dev and b.meta.has_cpu_mapping and not OSX: FileIOInterface.munmap(b.va_addr, b.size)
+    if self.is_local() and b.owner == self.dev and b.meta.has_cpu_mapping: FileIOInterface.munmap(b.va_addr, b.size)

  def map(self, b:HCQBuffer):
+    if not self.is_local(): raise RuntimeError(f"P2P mapping not supported for remote devices: {b.owner} -> {self.dev}")
+
    if b.owner is not None and b.owner._is_cpu():
      System.lock_memory(int(b.va_addr), b.size)
      paddrs, aspace = [(x, 0x1000) for x in System.system_paddrs(int(b.va_addr), round_up(b.size, 0x1000))], AddrSpace.SYS
      snooped, uncached = True, True
-    elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, LNXPCIIfaceBase):
+    elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, PCIIfaceBase):
      snooped, uncached = True, b.meta.mapping.uncached
      if b.meta.mapping.aspace is AddrSpace.SYS: paddrs, aspace = b.meta.mapping.paddrs, AddrSpace.SYS
      elif hasattr(ifa.dev_impl, 'paddr2xgmi') and ifa.dev_impl.gmc.xgmi_seg_sz > 0:
@@ -323,13 +339,16 @@ class RemotePCIDevice(PCIDevice):
 class APLRemotePCIDevice(RemotePCIDevice):
  APP_PATH = "/Applications/TinyGPU.app/Contents/MacOS/TinyGPU"

-  @staticmethod
-  def install_tinygpu():
+  @classmethod
+  def ensure_app(cls):
+    commit = "8120b5508b43149d27bf22f9a4e6d7c5a4b401e9"
+    if os.path.exists(cls.APP_PATH) and (_ensure_downloads_dir() / (app_name:=f"TinyGPU_{commit}.zip")).is_file(): return
    print("Downloading TinyGPU.app...")
-    system(f"ditto -xk {fetch('https://github.com/nimlgen/tinygpu_releases/raw/8120b5508b43149d27bf22f9a4e6d7c5a4b401e9/TinyGPU.zip')} /Applications")
-    print(system(f"{APLRemotePCIDevice.APP_PATH} install"))
+    system(f"ditto -xk {fetch(f'https://github.com/nimlgen/tinygpu_releases/raw/{commit}/TinyGPU.zip', name=app_name)} /Applications")
+    print(system(f"{cls.APP_PATH} install"))

  def __init__(self, devpref:str, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
+    self.ensure_app()
    sock_path, sock = getenv("APL_REMOTE_SOCK", temp("tinygpu.sock")), socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    for i in range(100):
      with contextlib.suppress(ConnectionRefusedError, FileNotFoundError):
@@ -339,20 +358,3 @@ class APLRemotePCIDevice(RemotePCIDevice):
      time.sleep(0.05)
    else: raise RuntimeError(f"Failed to connect to TinyGPU server at {sock_path}.")
    super().__init__(devpref, pcibus, bars, sock)
-
-class APLRemoteIfaceBase(LNXPCIIfaceBase):
-  def __init__(self, dev, dev_id, vendor, devices:list[tuple[int, list[int]]], bars, vram_bar, va_start, va_size, base_class:int|None=None):
-    if not (cls:=type(self)).gpus:
-      cls.gpus = System.pci_scan_bus(vendor, devices, base_class)
-      if not cls.gpus: raise RuntimeError("No supported GPUs found")
-      if not os.path.exists(APLRemotePCIDevice.APP_PATH): APLRemotePCIDevice.install_tinygpu()
-    if dev_id >= len(cls.gpus): raise RuntimeError(f"No device found for {dev_id}. Requesting more devices than the system has ({cls.gpus})?")
-    self.pci_dev = APLRemotePCIDevice(dev.__class__.__name__[:2], f'remote:{dev_id}', bars)
-    self.dev, self.vram_bar = dev, vram_bar
-
-  def free(self, b:HCQBuffer):
-    for dev in b.mapped_devs[1:]: dev.iface.dev_impl.mm.unmap_range(b.va_addr, b.size)
-
-  def map(self, b:HCQBuffer): raise RuntimeError(f"P2P mapping not supported for remote devices: {b.owner} -> {self.dev}")
-
-PCIIfaceBase:type = APLRemoteIfaceBase if OSX else LNXPCIIfaceBase
@@ -107,7 +107,7 @@ class USB3:
      if self.use_bot:
        dir_in = rlen > 0
        data_len = rlen if dir_in else (len(send_data) if send_data is not None else 0)
-        assert (data_len == 0) if dir_in else (rlen == 0), "BOT mode only supports either read or write per command"
+        assert not (rlen > 0 and send_data is not None), "BOT mode only supports either read or write per command"

        # CBW
        self._tag += 1
@@ -7,7 +7,7 @@ from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored

-ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
+ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.CALL}

@@ -25,15 +25,21 @@ def realize_assign_src(ctx:dict[UOp, None], buf:UOp, x:UOp):
  # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
  if buf.base in x.backward_slice_with_self: ctx[x] = None

+def unrealize_store_src(ctx:dict[UOp, None], x:UOp):
+  """Don't realize COPY/BUFFER_VIEW consumed by STORE inside AFTER — bufferize_to_store handles them."""
+  if x in ctx: del ctx[x]
+
 pm_generate_realize_map = PatternMatcher([
-  # always realize SINK src
-  (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
  # always realize
-  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE, Ops.ASSIGN}, name="tr"), realize),
+  (UPat({Ops.COPY, Ops.CONTIGUOUS, Ops.ASSIGN}, name="tr"), realize),
+  # realize AFTER of STORE+AFTER
+  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE)), allow_any_len=True, name="tr"), realize),
  # realize srcs of these
  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
  # sometimes realize src of assign
  (UPat(Ops.ASSIGN, src=(UPat.var("buf"), UPat.var("x"))), realize_assign_src),
+  # don't realize COPY/BUFFER_VIEW consumed by STORE inside AFTER (like realize_assign_src for ASSIGN)
+  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat({Ops.COPY, Ops.BUFFER_VIEW}, name="x"))))), unrealize_store_src),
 ])

@dataclass(frozen=True)
@@ -60,7 +66,8 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
  new_srcs = []
  for s in x.src:
    new_src = s
-    if s.op in {Ops.PARAM, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
+    if s.op in {Ops.PARAM, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or \
+       (s.op is Ops.AFTER and not any(c.op in {Ops.STORE, Ops.END} for c in s.src[1:])):
      if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
    elif s in ctx.realize_map:
      realized_ranges = ctx.realize_map[s]
@@ -99,25 +106,11 @@ def convert_reduce_axis_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
 def remove_movement_op_after_rangeify(ctx:IndexingContext, x:UOp):
  if x in ctx.range_map or x.src[0].op is Ops.INDEX: return x.src[0]

-def handle_assign_mops(ctx:IndexingContext, assign:UOp, target:UOp, src:UOp):
-  if target.op in GroupOp.Movement and src.op is not Ops.CALL:
-    mops = []
-    while target.op in GroupOp.Movement:
-      mops.append((target.op, target.marg))
-      target = target.src[0]
-    if mops and assign in ctx.range_map:
-      ret = assign.replace(arg=tuple(mops))
-      ctx.range_map[ret] = ctx.range_map[assign]
-      return ret
-  return None
-
 pm_apply_rangeify = PatternMatcher([
  # REDUCE_AXIS -> REDUCE
  (UPat(Ops.REDUCE_AXIS, name="x"), convert_reduce_axis_to_reduce_with_ranges),
  # PAD -> WHERE
  (UPat(Ops.PAD, name="x"), convert_pad_to_where_to_keep_behavior_local),
-  # store movement ops in ASSIGN arg
-  (UPat(Ops.ASSIGN, src=(UPat(name="target"), UPat(name="src")), name="assign"), handle_assign_mops),
  # finally, apply_rangeify
  (UPat(GroupOp.All, name="x"), create_bufferize_and_index_based_on_ranges),
  # remove movement op
@@ -179,8 +172,8 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
    # no ranges on kernels, they are internal
    if x.op in {Ops.CALL, Ops.LINEAR}: continue

-    # no range on after
-    if x.op is Ops.AFTER: continue
+    # only STORE+AFTER has range
+    if x.op is Ops.AFTER and all(s.op is not Ops.STORE for s in x.src[1:]): continue

    # treat MSTACK/MSELECT like SINK
    if x.op in {Ops.MSTACK, Ops.MSELECT}: continue
@@ -141,13 +141,10 @@ multi_pm = PatternMatcher([
    lambda multi,device,red: multi.src[0].allreduce(red.arg, device).multi(axis=multi.axis)),
  # rewrite into calls explicitly for MULTI
  (UPat(Ops.CALL, name="call"), rewrite_into_call),
-  (UPat(Ops.CALL, src=(UPat(Ops.MULTI, name="multi"), ), name="root", allow_any_len=True), passthrough_multi),
+  (UPat((Ops.CALL, Ops.AFTER, Ops.STORE), src=(UPat(Ops.MULTI, name="multi"), ), name="root", allow_any_len=True), passthrough_multi),
  # we just remove the MULTI from CALLs with dtypes.void and assume they are handled by the user for custom kernels
  (UPat(Ops.CALL, dtype=dtypes.void, name="root", custom_early_reject=set([Ops.MULTI])), lambda root:
    UOp(root.op, root.dtype, tuple(x.src[0] if x.op is Ops.MULTI else x for x in root.src), root.arg)),
  (UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD),
        src=(UPat(Ops.MULTI, name="multi"), ), name="root"), passthrough_multi),
-  # after CALL
-  (UPat(Ops.AFTER, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.CALL)), name="a"),
-    lambda multi,a: a.replace(src=(multi.src[0],)+a.src[1:]).multi(multi.axis)),
 ])+replace_allreduce
@@ -16,24 +16,39 @@ from tinygrad.schedule.allreduce import create_allreduce_function
 import sys
 sys.setrecursionlimit(10000)

+def add_ranges_to_store(ctx, x):
+  if x.src[0]._shape is None or x.src[1]._shape is None or x.src[0].shape == (): return None
+  assert x.src[0].shape == x.src[1].shape, "bad store shape"
+  idxs = [UOp.range(r, next(ctx), AxisType.LOOP) for r in x.src[0].shape]
+  return UOp.store(x.src[0].index(*idxs), x.src[1].index(*idxs)).end(*idxs)
+
+pm_store_ranges = PatternMatcher([
+  (UPat(Ops.STORE, name="x"), add_ranges_to_store),
+])
+
 pm_syntactic_sugar = PatternMatcher([
  # INDEX on ptr INDEX concats them
  (UPat(Ops.INDEX, name="i1").f(Ops.INDEX, name="i2", allow_any_len=True),
   lambda i1,i2: i2.replace(src=i1.src+i2.src[1:]) if isinstance(i1.dtype, PtrDType) and not isinstance(i2.dtype, PtrDType) else None),
+  # early rangeify
+  (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise | {Ops.CONST}, name="x"),), allow_any_len=True, name="idx"),
+   lambda idx,x: x.replace(src=tuple([s.index(*idx.src[1:]) for s in x.src]))),
 ])

 def found_assign(ctx:dict[UOp, UOp], assign:UOp, src:UOp):
  if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, assign = x.src[0], assign.cast(dtypes.float)
-  while x is not x.base:
-    if x.op is Ops.PERMUTE: assign = assign.permute(argsort(x.marg))
-    elif x.op is Ops.RESHAPE: assign = assign.reshape(x.src[0].shape)
-    else: return None
-    x = x.src[0]
+  while True:
+    if x.op is Ops.PERMUTE: x, assign = x.src[0], assign.permute(argsort(x.marg))
+    elif x.op is Ops.RESHAPE: x, assign = x.src[0], assign.reshape(x.src[0].shape)
+    elif x.op is Ops.WHERE and x.src[2].base.arg == Invalid and x.src[1].op is Ops.PAD:
+      x, assign = x.src[1].src[0], assign.shrink(tuple((l, s-r) for (l,r),s in zip(x.src[1].marg, x.shape)))
+    else: break
  ctx[x] = assign

-# *** fold moved ASSIGNs (hack for openpilot) ***
+# *** fold moved ASSIGNs/AFTERs (hack for openpilot) ***
 pm_fold_moved_assign = PatternMatcher([
  (UPat(Ops.ASSIGN, src=(UPat(), UPat((*GroupOp.Movement, Ops.CAST), name="src")), name="assign"), found_assign),
+  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat((*GroupOp.Movement, Ops.CAST), name="src")))), name="assign"), found_assign),
  # replace ALU sources with assign versions found above
  (UPat(GroupOp.ALU, name="alu"), lambda ctx,alu: alu.replace(src=new_src) if (new_src:=tuple(ctx.get(s, s) for s in alu.src)) != alu.src else None),
 ])
@@ -41,10 +56,12 @@ pm_fold_moved_assign = PatternMatcher([
 # movement op on INDEX as a PatternMatcher
 pm_mops = PatternMatcher([
  (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
-   lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)),
-  # move movement ops after AFTER
+   lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)
+     if len(idx.src[1:]) == len(r.shape) else None),
+  # move movement ops after AFTER (but not when AFTER has a raw STORE with shaped children — from replace_contig_with_store_after)
  (UPat(GroupOp.Movement, name="r").after(name="a", allow_any_len=True),
-   lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], r.arg)),
+   lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], r.arg)
+     if not any(s.op is Ops.STORE and s.src[0]._shape is not None for s in a.src[1:]) else None),
  (UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])),
 ])

@@ -54,12 +71,12 @@ pm_mops = PatternMatcher([
 def fix_assign_hazard(assign:UOp, target:UOp, src:UOp):
  # PERMUTE and FLIP reorder indices, SHRINK can have overlapping regions when dest is also shrunk
  unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set())
-  if any(s.op in unsafe and target.base in s.backward_slice for s in src.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS)):
+  if any(s.op in unsafe and target.base in s.backward_slice for s in src.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS or s.op is Ops.AFTER)):
    return assign.replace(src=(target, src.contiguous()))

 def normalize_assign_target_chain(assign:UOp, target:UOp, src:UOp):
  root_target = target
-  while root_target.op is Ops.ASSIGN: root_target = root_target.src[0]
+  while root_target.op in {Ops.ASSIGN, Ops.AFTER}: root_target = root_target.src[0]
  # when RHS depends on the previous assign result, break with contiguous
  if target in src.toposort(): src = src.contiguous()
  return assign.replace(src=(root_target, src))
@@ -120,8 +137,10 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
  # resolve calls
  (UPat(Ops.CALL, name="c"), resolve_call),

+  # resolve TUPLE+GETTUPLE
+  (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
+
  # resolve allreduce (must be bottom up)
-  (UPat(Ops.ASSIGN, src=(UPat.var("output"), UPat(Ops.ALLREDUCE, src=(UPat.var("buf"), UPat()), name="red"))), create_allreduce_function),
  (UPat(Ops.ALLREDUCE, src=(UPat.var("buf"), UPat()), name="red"), create_allreduce_function),

  # split_reduceop
@@ -156,8 +175,9 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
  (UPat(Ops.ASSIGN, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src"))),
   lambda target, src: target.assign(src.bitcast(target.dtype))),

-  # if assign target is itself an ASSIGN chain, canonicalize to the original buffer target
-  (UPat(Ops.ASSIGN, src=(UPat(Ops.ASSIGN, name="target"), UPat(name="src")), allow_any_len=True, name="assign"), normalize_assign_target_chain),
+  # if assign target is itself an ASSIGN/AFTER chain, canonicalize to the original buffer target
+  (UPat(Ops.ASSIGN, src=(UPat({Ops.ASSIGN, Ops.AFTER}, name="target"), UPat(name="src")), allow_any_len=True, name="assign"),
+   normalize_assign_target_chain),

  # make source contiguous if it has hazardous movement ops on the dest buffer
  (UPat(Ops.ASSIGN, src=(UPat.var("target"), UPat.var("src")), name="assign"), fix_assign_hazard),
@@ -178,8 +198,8 @@ ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.NOOP}

 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
-  # don't optimize ALWAYS_RUN_OPS
-  if b.src[0].op in ALWAYS_RUN_OPS: return None
+  # don't optimize ALWAYS_RUN_OPS or AFTER (AFTER is a buffer identity — ranges define consumer access, not computation)
+  if b.src[0].op in ALWAYS_RUN_OPS or b.src[0].op is Ops.AFTER: return None

  new_rng = []
  hit = False
@@ -210,56 +230,54 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
  # if it's user contiguous, we never remove it
  if src.op in ALWAYS_RUN_OPS or not buf.arg.removable: return None

-  # we don't want to bufferize threefry, also causes problems because not all platforms support long
-  if src.op is not Ops.THREEFRY:
-    # *** here is where we compute the cost ***
-    # if we return None, the bufferize is kept
+  # *** here is where we compute the cost ***
+  # if we return None, the bufferize is kept

-    accessed_buffers: list[UOp] = []
-    indexes: list[UOp] = []
-    reduces: list[UOp] = []
-    def red_gate(x:UOp):
-      if (x.op is Ops.BUFFERIZE and x.arg.addrspace == AddrSpace.GLOBAL) or x.op is Ops.MSTACK:
-        accessed_buffers.append(x)
-        return False
-      if x.op is Ops.PARAM:
-        accessed_buffers.append(x)
-      if x.op is Ops.INDEX:
-        indexes.append(x)
-      if x.op is Ops.REDUCE: reduces.append(x)
-      return True
-    src.toposort(gate=red_gate)
-    del red_gate
-    accessed_buffers = dedup(accessed_buffers)
+  accessed_buffers: list[UOp] = []
+  indexes: list[UOp] = []
+  reduces: list[UOp] = []
+  def red_gate(x:UOp):
+    if (x.op is Ops.BUFFERIZE and x.arg.addrspace == AddrSpace.GLOBAL) or x.op is Ops.MSTACK:
+      accessed_buffers.append(x)
+      return False
+    if x.op is Ops.PARAM:
+      accessed_buffers.append(x)
+    if x.op is Ops.INDEX:
+      indexes.append(x)
+    if x.op is Ops.REDUCE: reduces.append(x)
+    return True
+  src.toposort(gate=red_gate)
+  del red_gate
+  accessed_buffers = dedup(accessed_buffers)

-    # if this is generated from multiple buffers, don't remove this buffer
-    if len(accessed_buffers) > 3 and not (PCONTIG > 2): return None
+  # if this is generated from multiple buffers, don't remove this buffer
+  if len(accessed_buffers) > 3 and not (PCONTIG > 2): return None

-    # if any reduces access a buffer, don't remove this buffer
-    buffer_in_reduce = False
-    def buf_gate(x:UOp):
-      nonlocal buffer_in_reduce
-      if x.op in {Ops.PARAM, Ops.BUFFERIZE}: buffer_in_reduce = True
-      return not buffer_in_reduce
-    UOp.sink(*[x.src[0] for x in reduces]).toposort(gate=buf_gate)
-    del buf_gate
-    if buffer_in_reduce:
-      if PCONTIG > 2:
-        out_in_ratio = (prod(buf.shape)+1) / (sum([x.size for x in accessed_buffers])+1)
-        if out_in_ratio < 10: return None
-        # here we have to check the indexes, we might do a partial contig here
-        local_indexes = [x for x in indexes if x.src[0].op is Ops.BUFFERIZE and x.src[0].arg.addrspace == AddrSpace.LOCAL]
-        exclude_ranges = UOp.group(*[UOp.group(*x.src[1:]) for x in local_indexes]).ranges
-        subs = [(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST]
-        # if it's bufferized or a reduce, it's pcontig
-        is_pcontig, is_subs = partition(subs, lambda x: x[0] in exclude_ranges or any([r.arg[-1] == AxisType.REDUCE for r in x[1].ranges]))
-        if not len(is_subs):
-          return None
-        if len(is_pcontig):
-          ret = src.substitute(dict(is_subs), extra_pm=pm_gate_substitute)
-          return ret.bufferize(*[x[0] for x in is_pcontig], arg=BufferizeOpts(None, AddrSpace.LOCAL)).index(*[x[1] for x in is_pcontig])
-      else:
+  # if any reduces access a buffer, don't remove this buffer
+  buffer_in_reduce = False
+  def buf_gate(x:UOp):
+    nonlocal buffer_in_reduce
+    if x.op in {Ops.PARAM, Ops.BUFFERIZE}: buffer_in_reduce = True
+    return not buffer_in_reduce
+  UOp.sink(*[x.src[0] for x in reduces]).toposort(gate=buf_gate)
+  del buf_gate
+  if buffer_in_reduce:
+    if PCONTIG > 2:
+      out_in_ratio = (prod(buf.shape)+1) / (sum([x.size for x in accessed_buffers])+1)
+      if out_in_ratio < 10: return None
+      # here we have to check the indexes, we might do a partial contig here
+      local_indexes = [x for x in indexes if x.src[0].op is Ops.BUFFERIZE and x.src[0].arg.addrspace == AddrSpace.LOCAL]
+      exclude_ranges = UOp.group(*[UOp.group(*x.src[1:]) for x in local_indexes]).ranges
+      subs = [(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST]
+      # if it's bufferized or a reduce, it's pcontig
+      is_pcontig, is_subs = partition(subs, lambda x: x[0] in exclude_ranges or any([r.arg[-1] == AxisType.REDUCE for r in x[1].ranges]))
+      if not len(is_subs):
        return None
+      if len(is_pcontig):
+        ret = src.substitute(dict(is_subs), extra_pm=pm_gate_substitute)
+        return ret.bufferize(*[x[0] for x in is_pcontig], arg=BufferizeOpts(None, AddrSpace.LOCAL)).index(*[x[1] for x in is_pcontig])
+    else:
+      return None

  # if it makes it here, the bufferize is removed
  # this is the ranges replaced
@@ -355,20 +373,24 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
  assert size > 0 and isinstance(size, int), f"no zero sized or symbolic sized buffers {size}"

  sdtype = x.dtype.ptr(size=size, addrspace=x.arg.addrspace)
+  # AFTER: add END to the existing STORE, return buffer with kernel dependency
+  if x.src[0].op is Ops.AFTER:
+    buf = x.src[0].src[0].buf_uop.base
+    stores = [s for s in x.src[0].src[1:] if s.op is Ops.STORE]
+    return buf.after(*[s.end(*rngs) for s in stores]) if stores else buf
  if (assign := x.src[0]).op is Ops.ASSIGN:
    assign_target, assign_src = assign.src[0], assign.src[1]
    assert assign_target.op is Ops.INDEX, f"{assign_target.op} is not index"
    while assign_src.op is Ops.NOOP: assign_src = assign_src.src[0]

    store_target = assign_target
-    if assign.arg and assign_target.src[0].op is Ops.BUFFERIZE and assign_target.src[0].src[0].op is Ops.INDEX:
+    if assign_target.src[0].op is Ops.BUFFERIZE and assign_target.src[0].src[0].op is Ops.INDEX:
      # BUFFERIZE(INDEX(...)); store through the underlying global index instead.
      store_target = assign_target.src[0].src[0]

    end_rngs = sorted(dedup(tuple(store_target.ranges) + tuple(rngs)), key=lambda x: x.arg)
    ret = store_target.buf_uop.base
-    if assign_src is not store_target: ret = ret.after(store_target.replace(dtype=sdtype).store(assign_src).end(*end_rngs))
-    for op, marg in reversed(assign.arg or ()): ret = ret._mop(op, marg)
+    if assign_src is not assign_target: ret = ret.after(store_target.replace(dtype=sdtype).store(assign_src).end(*end_rngs))
    return ret

  # NOTE: the DEFINE_LOCAL needs to be disambiguated here
@@ -518,6 +540,8 @@ pm_add_range_tags = PatternMatcher([
 def split_store(x:UOp) -> UOp|None:
  # if we have any open ranges here, we don't split
  if x.ranges: return None
+  # raw STORE (not from bufferize_to_store) should be processed through its END wrapper, not independently
+  if x.op is Ops.STORE and x.src[0]._shape is not None: return None

  # local kernel rewrite
  lctx = LocalAddBufferContext()
@@ -313,7 +313,14 @@ class Tensor(OpMixin):
    assign_uop = self.uop.assign(x.uop)
    base = self.uop.base
    if base.op in {Ops.BUFFER, Ops.AFTER} and not self.uop.has_buffer_identity():
-      _apply_map_to_tensors({base: base.after(assign_uop)}, name="Embed View Assign", walk=True)
+      original_uop = self.uop
+      assigned_base = base.after(assign_uop)
+      _apply_map_to_tensors({base: assigned_base}, name="Embed View Assign", walk=True)
+      def replace_view_base(u:UOp) -> UOp:
+        return u.replace(src=((assigned_base if u.src[0] is base else replace_view_base(u.src[0])),)+u.src[1:])
+      ret = Tensor(replace_view_base(original_uop), device=self.device, requires_grad=self.requires_grad)
+      self.replace(self._apply_uop(lambda *_: assign_uop, x))
+      return ret
    return self.replace(self._apply_uop(lambda *_: assign_uop, x))

  def detach(self) -> Tensor:
@@ -1064,7 +1071,7 @@ class Tensor(OpMixin):
    for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient, materialize_grads=True)):
      assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
      if t.grad is None: t.grad = g
-      else: t.grad.assign(t.grad + g)
+      else: t.grad.assign(t.grad + g.to(t.grad.device))
    return self

  # ***** movement low level ops *****
@@ -1340,7 +1347,7 @@ class Tensor(OpMixin):
      if is_disk: raise RuntimeError("advanced setitem is not supported for DISK tensors")
      if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
      self.assign(self._getitem(indices, v))
-    elif is_disk or self.uop.is_realized or self.uop.base.op is Ops.AFTER: # basic setitem, self is realized
+    elif is_disk or self.uop.is_realized or self.uop.base.op in (Ops.AFTER, Ops.BUFFER): # basic setitem, self is realized
      view = self[indices]
      if isinstance(v, Tensor) and v.uop.op is Ops.ASSIGN and v.uop in view.uop.base.src: return
      view.assign(v)
@@ -2109,7 +2116,7 @@ class Tensor(OpMixin):
    x_unsqueezed = x.unsqueeze(-2).expand((None,)*(self.ndim-1)+(last_dim_size, None))
    x_cummax, _ = x.cummax(-1)
    mask = Tensor.ones(last_dim_size, last_dim_size, requires_grad=False, device=self.device).tril()
-    ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), dtypes.min(self.dtype)).exp().sum(-1).log() + x_cummax
+    ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax
    return ret.transpose(-1, axis)

  def argmax(self, axis=None, keepdim=False) -> Tensor:
@@ -2306,12 +2313,12 @@ class Tensor(OpMixin):
    axis = tuple(range(-len(k_ := make_tuple(kernel_size, 2)), 0))
    pads = self._resolve_pool_pads(padding, len(k_))
    if ceil_mode: pads = self._apply_ceil_mode(pads, k_, stride if stride is not None else k_, dilation)
-    pooled = self.pad(pads, value=dtypes.min(self.dtype))._pool(k_, stride if stride is not None else k_, dilation)
+    pooled = self.pad(pads, value=self.dtype.min)._pool(k_, stride if stride is not None else k_, dilation)
    if not return_indices: return pooled.max(axis)
    spatial_sz = int(math.prod(spatial_shape := self.shape[-len(k_):]))
    idx = Tensor.arange(spatial_sz,0,-1, requires_grad=False, device=self.device).reshape(spatial_shape)
    m = pooled == pooled.max(axis, keepdim=True)
-    idx = m * idx.pad(pads, value=dtypes.min(idx.dtype))._pool(k_, stride if stride is not None else k_, dilation)
+    idx = m * idx.pad(pads, value=idx.dtype.min)._pool(k_, stride if stride is not None else k_, dilation)
    return pooled.max(axis), spatial_sz - idx.max(axis)

  def max_unpool2d(self, indices:Tensor, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0, output_size=None):
@@ -2484,7 +2491,7 @@ class Tensor(OpMixin):
    """
    if IMAGE: return self.image_dot(w, dtype)
    if ASM_GEMM:
-      from extra.gemm.asm.cdna.gemm import can_use_asm_gemm, asm_gemm
+      from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
      if can_use_asm_gemm(self, w): return asm_gemm(self, w)
    x, dx, dw = self, self.ndim, w.ndim
    if not (dx > 0 and dw > 0): raise RuntimeError(f"both tensors need to be at least 1D, got {dx}D and {dw}D")
@@ -2752,8 +2759,8 @@ class Tensor(OpMixin):
    def _inv_mask(a:Tensor|PyConst, b:Tensor|PyConst) -> Tensor: return mask.any(-1).logical_not().where(a, b)
    if reduce == "sum": return mask.where(src, 0).sum(-1).add(self if include_self else _inv_mask(self, 0))
    if reduce == "prod": return mask.where(src, 1).prod(-1).mul(self if include_self else _inv_mask(self, 1))
-    if reduce == "amax": return mask.where(src, m := dtypes.min(src.dtype)).max(-1).maximum(self if include_self else _inv_mask(self, m))
-    if reduce == "amin": return mask.where(src, m := dtypes.max(src.dtype)).min(-1).minimum(self if include_self else _inv_mask(self, m))
+    if reduce == "amax": return mask.where(src, m := src.dtype.min).max(-1).maximum(self if include_self else _inv_mask(self, m))
+    if reduce == "amin": return mask.where(src, m := src.dtype.max).min(-1).minimum(self if include_self else _inv_mask(self, m))
    if reduce == "mean":
      count = mask.where(1, 0).sum(-1).add(1 if include_self else _inv_mask(1, 0))
      return mask.where(src, 0).sum(-1).add(self if include_self else _inv_mask(self, 0)).div(count)
@@ -2782,7 +2789,7 @@ class Tensor(OpMixin):
    # pad to power of 2
    n_stages = (orig_len-1).bit_length()
    pads = tuple((0, 2**n_stages - orig_len) if i == dim else None for i in range(x.ndim))
-    x = x.pad(pads, value=dtypes.min(x.dtype) if descending else dtypes.max(x.dtype)).unflatten(dim, (2,)*n_stages)
+    x = x.pad(pads, value=x.dtype.min if descending else x.dtype.max).unflatten(dim, (2,)*n_stages)
    # https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort1.svg
    for stage in range(1, n_stages+1):
      if stage != n_stages:
@@ -3655,21 +3662,27 @@ class Tensor(OpMixin):

    # contiguous creates the image, and early realize static weights (TODO: test for the static weight)
    if IMAGE == 1:
-      # pad with Invalid
-      def _invalid_pad_to(t, shape):
-        if all(p is None or p == s for p,s in zip(shape, t.shape)): return t
-        return Tensor(True, device=t.device).expand(t.shape).pad_to(shape).where(t.pad_to(shape), Invalid)
-      # hacks for pitch alignment
-      assert isinstance(ix, int) and isinstance(H, int)
-      ALIGN = 64 // dtsz
-      x = _invalid_pad_to(x, (None, None, round_up(ix, ALIGN // math.gcd(groups * cin, ALIGN)), None))
-      w = _invalid_pad_to(w, (None, round_up(H, ALIGN // math.gcd(W * cin * 4, ALIGN))) + (None,) * (w.ndim - 2))
+      def is_pow2(v): return v > 0 and v & (v - 1) == 0
+      # pad dimension i to amt with invalids
+      def ipad(t, i, amt):
+        shape = (None,)*i + (amt,) + (None,)*(t.ndim-i-1)
+        return Tensor(True, device=t.device).expand(t.shape).pad_to(shape).where(t.pad_to(shape), Invalid) if amt != t.shape[i] else t
+      # align a dimension to 64 bytes
+      def pad_align(t, dim):
+        return ipad(t, dim, round_up(t.shape[dim], (64 // dtsz) // math.gcd(prod(t.shape) // t.shape[dim], (64 // dtsz))))
+
+      # bank conflicts
+      if cin >= 8 and is_pow2(cin // 4): x, w = ipad(x.reshape(bs, iy, ix, groups, cin // 4, 4), 4, cin // 4 + 1), ipad(w, 2, cin // 4 + 1)
+
+      # 64-byte pitch alignment
+      x, w = pad_align(x, 2), pad_align(w, 1)

      if FLOAT16: x, w = x.cast(dtypes.half).contiguous().cast(dtypes.float), w.cast(dtypes.half).contiguous().cast(dtypes.float)
      else: x, w = x.contiguous(), w.contiguous()

      # undo alignment hacks
-      x, w = x[:, :, :ix, :], w[:, :H, ...]
+      if cin >= 8 and is_pow2(cin // 4): x, w = x[:, :, :ix, :, :cin // 4, :], w[:, :H, :cin // 4, ...]
+      else: x, w = x[:, :, :ix, :], w[:, :H, ...]

    elif IMAGE: x, w = x.cast(base_image_type((bs*iy, ix*groups*cin//4, 4))).contiguous(), w.cast(base_image_type((cout//4, H*W*cin, 4))).contiguous()
    else: x, w = x.contiguous(), w.contiguous()
@@ -39,6 +39,9 @@ class Ops(FastEnum):
  # vector creation / item selection
  GEP = auto(); VECTORIZE = auto()

+  # tuple/gettuple for function with multiple returns
+  TUPLE = auto(); GETTUPLE = auto()
+
  # ** 3 -- load/store **

  # INDEX is a BinaryOp similar to ADD, but it operates on pointers
@@ -14,7 +14,7 @@ def _lazy_map_numbers(x:UOp, inf:UOp, _inf:UOp, nan:UOp, ratio:UOp):

 # *** helper functions for bit manipulation ***
 def mantissa_bits(d:DType) -> int: return dtypes.finfo(d.scalar())[1]
-def exponent_bias(d:DType) -> int: return (1 << (dtypes.finfo(d.scalar())[0] - 1)) - 1
+def exponent_bias(d:DType) -> int: return (1 << (dtypes.finfo(d.scalar())[0] - 1)) - (0 if d.scalar() in dtypes.fp8_fnuz else 1)
 def exponent_mask(d:DType) -> int: return (1 << dtypes.finfo(d.scalar())[0]) - 1

 # **** utils ****
@@ -287,7 +287,7 @@ def fast_idiv(device: str, x: UOp, d: int, dont_cast=False) -> UOp|None:
  assert d>0, "Sign should have been taken out of divisor"
  vmin,vmax = max(x.vmin, x.dtype.min), min(x.vmax, x.dtype.max)
  m,s = magicgu(max(vmax, abs(vmin)), d)
-  if m*vmin >= dtypes.min(x.dtype) and m*vmax <= dtypes.max(x.dtype):
+  if m*vmin >= x.dtype.min and m*vmax <= x.dtype.max:
    return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)
  # before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
  if (largest_factor_of_two_in_d := (d & -d)) > 1:
@@ -295,7 +295,7 @@ def fast_idiv(device: str, x: UOp, d: int, dont_cast=False) -> UOp|None:
  if dont_cast: return None
  # promo_lattice needs to return an unsigned type if the type is unsigned
  if dtypes.is_int(next_dtype := promo_lattice[x.dtype.scalar()][-1]) and is_dtype_supported(next_dtype, device):
-    if m*vmin >= dtypes.min(next_dtype) and m*vmax <= dtypes.max(next_dtype):
+    if m*vmin >= next_dtype.min and m*vmax <= next_dtype.max:
      return ((x.cast(next_dtype)*m) >> s).cast(x.dtype) if is_unsigned else ((x.cast(next_dtype)*m) >> s).cast(x.dtype) + (x<0).where(x.ufix(1), 0)
  return None

@@ -388,6 +388,10 @@ def f2f(v, fr:DType, to:DType):
    sign, nosign = shl((v & shl(1, fs-1)).cast(f2f_dt[to]), ts - fs), (v & (shl(1, fs-1) - 1)).cast(f2f_dt[to])
    exp, norm = shr(nosign, fm), shl(nosign, tm - fm) + shl(tb - fb, tm)
    nan = shl(nosign, tm - fm) | shl((shl(1, te) - 1), tm)
+    if fr in dtypes.fp8_fnuz:
+      fnuz_nan = sign.ne(0) & nosign.eq(0)
+      qnan = shl(shl(1, te) - 1, tm) | shl(1, tm - 1)
+      return fnuz_nan.where(qnan, sign | exp.eq(0).where(0, norm)).bitcast(to)
    # fp8e4m3 has only one nan
    is_nan = (nosign.eq(shl(1, fm + fe) - 1) if fr == dtypes.fp8e4m3 else exp.eq(shl(1, fe) - 1))
    return (sign | exp.eq(0).where(0, is_nan.where(nan, norm))).bitcast(to)
@@ -399,12 +403,14 @@ def f2f(v, fr:DType, to:DType):
    nan_mantissa = (shl(1, tm) - 1) if to == dtypes.fp8e4m3 else (shr(nosign, fm - tm) & (shl(1, tm) - 1))
    nan = (sign | nan_mantissa | shl(shl(1, te) - 1, tm)).cast(f2f_dt[to])
    is_nan = (shr(v, fm) & (shl(1, fe) - 1)).eq(shl(1, fe) - 1)
+    if to in dtypes.fp8_fnuz: return is_nan.where(shl(1, ts - 1), underflow.where(0, sign.cast(f2f_dt[to]) | norm))
    return is_nan.where(nan, sign.cast(f2f_dt[to]) | underflow.where(0, norm))
  else: raise NotImplementedError(f"unsupported decomp {fr} -> {to}")

 def f2f_clamp(val:UOp, dt:DType) -> UOp:
  e, m = dtypes.finfo(dt)
-  max_exp, max_man = ((1 << e) - 1, (1 << m) - 2) if dt == dtypes.fp8e4m3 else ((1 << e) - 2, (1 << m) - 1)
+  if dt in dtypes.fp8_fnuz: max_exp, max_man = (1 << e) - 1, (1 << m) - 1
+  else: max_exp, max_man = ((1 << e) - 1, (1 << m) - 2) if dt == dtypes.fp8e4m3 else ((1 << e) - 2, (1 << m) - 1)
  mx = val.const_like(2.0**(max_exp - exponent_bias(dt)) * (1.0 + max_man / (1 << m)))
  sat = mx if dt in dtypes.fp8s else val.const_like(float('inf'))
  # FIXME: CMPLT of nan is undefined
@@ -29,7 +29,7 @@ axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisTy
 range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.COPY: 2, Ops.BUFFER_VIEW: 1}

 # https://en.wikipedia.org/wiki/Identity_element
-def identity_element(op:Ops, dt:DType) -> PyConst: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt)
+def identity_element(op:Ops, dt:DType) -> PyConst: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dt.min}[op], dt)

 # With True as the default, this matches the old symbolic behavior
 def resolve(x:UOp|bool, default:bool=True):
@@ -207,11 +207,17 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
  def _shape(self) -> tuple[sint, ...]|None:
    match self.op:
      # late ops don't have shape
-      case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
+      case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | Ops.STORE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
           Ops.VECTORIZE | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.SINK | \
-           Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS:
+           Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS | Ops.TUPLE:
        return None

+      case Ops.GETTUPLE:
+        # GETTUPLE extracts from a TUPLE
+        in_tuple = self.src[0].src[0] if self.src[0].op is Ops.CALL else self.src[0]
+        assert in_tuple.op is Ops.TUPLE
+        return in_tuple.src[self.arg]._shape
+
      case Ops.CAST:
        # when PTX casts from ptr to non ptr, remove the shape
        if isinstance(self.src[0].dtype, PtrDType) and not isinstance(self.src[0].dtype, ImageDType) and not isinstance(self.dtype, PtrDType):
@@ -304,7 +310,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    if self.op is Ops.ASSIGN: return self.src[1]._shape

    # elementwise ops keep the shape the same. all inputs with shape must match
-    if self.op in GroupOp.ALU.union({Ops.CAST, Ops.COPY, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE, Ops.STORE}):
+    if self.op in GroupOp.ALU.union({Ops.CAST, Ops.COPY, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE}):
      input_shapes = [x._shape for x in self.src if x._shape is not None]
      if len(input_shapes) == 0: return None
      if not all_same(input_shapes): raise RuntimeError(f"shape mismatch at {self.op}: {input_shapes}")
@@ -404,6 +410,12 @@ class UOp(OpMixin, metaclass=UOpMetaClass):

  def sink(*srcs:UOp|None, **kwargs):  # pylint: disable=no-self-argument
    return UOp(Ops.SINK, dtypes.void, tuple([x for x in srcs if x is not None]), **kwargs)
+  def maketuple(*srcs:UOp):  # pylint: disable=no-self-argument
+    return UOp(Ops.TUPLE, dtypes.void, srcs)
+  def gettuple(self, idx:int) -> UOp:
+    in_tuple = self.src[0] if self.op is Ops.CALL else self
+    assert in_tuple.op is Ops.TUPLE, f"gettuple requires CALL or TUPLE source, got {self.op}"
+    return UOp(Ops.GETTUPLE, in_tuple.src[idx].dtype, (self,), idx)
  def group(*srcs:UOp|None):  # pylint: disable=no-self-argument
    if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
    return UOp(Ops.GROUP, dtypes.void, tuple([x for x in srcs if x is not None]))
@@ -414,6 +426,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
  def __getitem__(self, idx):
    idx = argfix(idx)
+    # add : to the end
+    if len(idx) < len(self.shape): idx += tuple([slice(None)]*(len(self.shape)-len(idx)))
    assert len(idx) == len(self.shape), f"__getitem__ shape mismatch, indexing {self.shape} with {len(idx)} args"
    if len(slice_idx:=[i for i,x in enumerate(idx) if isinstance(x, slice)]):
      perm = self.permute(tuple([i for i in range(self.ndim) if i not in slice_idx] + slice_idx))
@@ -843,8 +857,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    if self.op is Ops.GEP: return self.src[0]._min_max
    # TODO: CAST to bool/unsigned is not monotone, still some case can be simplified
    if self.op is Ops.CAST and self.dtype in dtypes.floats+dtypes.sints+(dtypes.index,):
-      return max(dtypes.min(self.dtype), self.src[0].vmin), min(self.src[0].vmax, dtypes.max(self.dtype))
-    return dtypes.min(self.dtype), dtypes.max(self.dtype)
+      return max(self.dtype.min, self.src[0].vmin), min(self.src[0].vmax, self.dtype.max)
+    return self.dtype.min, self.dtype.max

  @functools.cached_property
  def _sym_fxn(self):
@@ -901,9 +915,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    if self.axis is not None: p = p.replace(src=p.src + (UOp(Ops.MULTI, arg=self.axis),))
    return p

-  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), name:str|None=None, precompile:bool=False) -> UOp:
+  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(),
+           name:str|None=None, precompile:bool=False, precompile_backward:bool=False) -> UOp:
    assert len(self.ranges) == 0, f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}"
-    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata, name, precompile))
+    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata, name, precompile, precompile_backward))
  def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
    contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
@@ -927,9 +942,12 @@ class CallInfo:
  metadata: tuple[Metadata, ...] = ()
  name: str|None = None
  precompile: bool = False
+  precompile_backward: bool = False
  # grad_fxn can't be pickled, but metadata can
-  def __reduce__(self): return (CallInfo, (None, self.metadata, self.name, self.precompile))
-  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata}, {repr(self.name)}, {self.precompile})"
+  def __reduce__(self): return (CallInfo, (None, self.metadata, self.name, self.precompile, self.precompile_backward))
+  def __repr__(self):
+    gf = id(self.grad_fxn) if self.grad_fxn else None
+    return f"CallInfo({gf}, {self.metadata}, {repr(self.name)}, {self.precompile}, {self.precompile_backward})"

 def should_resolve_call(c:UOp) -> bool:
  # don't resolve real kernel calls, sink or program
@@ -59,6 +59,9 @@ shared_spec = PatternMatcher([
  # RANGE/SPECIAL define loops, END closes them
  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE))), lambda: True),

+  # STORE in tensor graph: store a value into a target
+  (UPat(Ops.STORE, dtypes.void, (UPat(), UPat())), lambda: True),
+
  # NOOP
  (UPat(Ops.NOOP), lambda: True)
 ])
@@ -75,7 +78,7 @@ movement_ops = PatternMatcher([
  (UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.index), lambda: True),

  # AFTER on Movement Op or ASSIGN
-  (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.MULTI, Ops.CONTIGUOUS, Ops.ASSIGN})),), allow_any_len=True), lambda: True),
+  (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.MULTI, Ops.CONTIGUOUS, Ops.ASSIGN, Ops.BUFFER})),), allow_any_len=True), lambda: True),
 ])

 _tensor_spec = PatternMatcher([
@@ -136,6 +139,10 @@ _tensor_spec = PatternMatcher([
  (UPat(Ops.PARAM), lambda: True),
  (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda x: isinstance(x.arg, str)),

+  # TUPLE must have void dtype, GETTUPLE can only appear on CALL or TUPLE
+  (UPat(Ops.TUPLE, dtypes.void), lambda: True),
+  (UPat(Ops.GETTUPLE, src=(UPat((Ops.CALL, Ops.TUPLE)),), name="g"), lambda g: isinstance(g.arg, int)),
+
  # ** for custom kernels **

  # codegen: PROGRAM with progressive sources through the pipeline (SINK, DEVICE, LINEAR?, SOURCE?, BINARY?)
@@ -230,8 +237,8 @@ program_spec = PatternMatcher([
  # END closes ranges
  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE)), dtype=dtypes.void), lambda: True),

-  # make sure all index dtypes have been lowered
-  (UPat(GroupOp.All, dtype=dtypes.index), lambda: False),
+  # make sure all index dtypes have been lowered (except CONST/RANGE/DEFINE_VAR which are valid index-typed)
+  (UPat(GroupOp.All-{Ops.CONST, Ops.RANGE, Ops.DEFINE_VAR, Ops.VCONST, Ops.VECTORIZE}, dtype=dtypes.index), lambda: False),
  (UPat(Ops.CONST, arg=Invalid), lambda: False),
  (UPat(Ops.VCONST, name="x"), lambda x: all(v is not Invalid for v in x.arg) and len(x.arg)==x.dtype.vcount>1 and
    type(x.arg) is type(dtypes.as_const(x.arg, x.dtype))),
@@ -65,6 +65,8 @@ propagate_invalid = PatternMatcher([
  (UPat(GroupOp.Binary-GroupOp.Comparison, src=[invalid_pat, UPat()]), lambda i: i),
  # normalize where(cond, Invalid, val) -> where(~cond, val, Invalid)
  (UPat.var("cond").where(invalid_pat, UPat.var("val")), lambda cond, i, val: cond.logical_not().where(val, i) if val.arg != Invalid else i),
+  (UPat.var("a").where(invalid_gate, UPat.var("c")), lambda cond,i,x,a,c: cond.where(a.where(x, c), i) if c.arg != Invalid else None),
+  (UPat.var("a").where(UPat.var("b"), invalid_gate), lambda cond,i,x,a,b: cond.where(a.where(b, x), i) if b.arg != Invalid else None),
  (UPat(Ops.BITCAST, src=(invalid_pat,), name="bc"), lambda bc,i: i.cast(bc.dtype)),
  (UPat(Ops.BITCAST, src=(invalid_gate,), name="bc"), lambda bc,cond,x,i: cond.where(x.bitcast(bc.dtype), i.bitcast(bc.dtype))),
 ])
@@ -138,12 +140,6 @@ symbolic_simple = propagate_invalid + PatternMatcher([
  (UPat.cvar("gate", vec=False).where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
  # a.where(b.where(c, d), d) -> (a & b).where(c, d)
  (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
-  # where over invalid -> invalid
-  (invalid_gate.where(UPat.var("a"), UPat.var("b")), lambda a,b,cond,x,i: i.cast(a.dtype)),
-  (invalid_pat.where(UPat.var("a"), UPat.var("b")), lambda a,b,i: i.cast(a.dtype)),
-  # reduce with invalid -> invalid
-  (UPat(Ops.REDUCE, src=(invalid_gate,), allow_any_len=True, name="r"), lambda r,cond,x,i: i.cast(r.dtype)),
-  (UPat(Ops.REDUCE, src=(invalid_pat,), allow_any_len=True, name="r"), lambda r,i: i.cast(r.dtype)),
 ])

 # ******** phase 2 builds on phase 1, it includes the old "symbolic", rules that match deeper ********
@@ -157,10 +157,6 @@
  #insts .left.highlight {
    background-color: rgba(0, 199, 47, 0.2);
  }
-  #insts .n {
-    color: #787fa1;
-    min-width: 5ch;
-  }
  #insts .wave {
    color: #7aa2f7;
    min-width: 2ch;
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	e4d8e03954	more test	2026-03-16 12:37:16 +08:00
geohot	09574a096a	maketuple	2026-03-16 12:04:39 +08:00
geohot	b89c233917	fix	2026-03-16 11:56:36 +08:00
geohot	00d847afdf	single backward gradient	2026-03-16 11:44:15 +08:00
geohot	8d75eed0a4	fix precompile	2026-03-16 11:33:52 +08:00
geohot	e59cbf78bd	Add TUPLE and GETTUPLE	2026-03-16 11:22:03 +08:00
qazalandGitHub	4445f50356	viz: variable duration rdna barriers (#15277 ) * viz: variable length rdna barriers * work * tiny changes * simple wave simd test * small wave sync test * good multi barrier bug find * simple fix * wave_sync asserts * rdna4 work * more rdna4 * find more bugs in my model * it's so much simpler * wave_sync tests duration * r4 * should just call this rdna4	2026-03-16 06:06:19 +09:00
qazalandGitHub	5cd1daa3bc	cdna asm_gemm in one file, remove old rdna3 asm (#15281 )	2026-03-16 04:32:30 +09:00
chenyuandGitHub	cd14e8e64b	allocations contiguous is store+after (#15280 )	2026-03-15 11:58:40 -04:00
qazalandGitHub	7b6211fdd7	sqtt: remove discover_ops script (#15279 )	2026-03-15 22:17:06 +09:00
wozeparrotandGitHub	473e5e4368	feat: make USE_ATOMICS embedding bwd faster (#15151 )	2026-03-14 21:21:10 -07:00
qazalandGitHub	3858bfc83d	sqtt: CDNA inst decodes (#15274 ) * sqtt: CDNA inst decodes * JUMP packets other way * cdna insts * r3 * r4 * lds from simd1 and simd2	2026-03-14 21:03:46 +09:00
sirhcmandGitHub	d753c5d7e5	IMAGE=1 image_conv2d pads for bank conflicts (#15252 )	2026-03-14 07:59:16 -04:00
sirhcmandGitHub	9047249a7c	m.where(x.pad_to(m.shape), Invalid) ranges shrink (#15275 )	2026-03-14 07:26:36 -04:00
nimlgenandGitHub	f392c53c66	system: merge remote into pciiface (#15273 ) * system: merge remote into pciiface * clenaer * move * mypy * fix	2026-03-14 18:44:20 +08:00
chenyuandGitHub	13eec8fbe8	remove unused assign rules [pr] (#15268 )	2026-03-14 05:37:49 -04:00
sirhcmandGitHub	dabdc986df	shrink guarded ranges, try 2 (#15272 )	2026-03-14 04:24:05 -04:00
sirhcmandGitHub	7cf4b16c91	Revert "shrink guarded ranges" (#15271 )	2026-03-14 03:44:38 -04:00
sirhcmandGitHub	d9951e2f8e	shrink guarded ranges (#15263 )	2026-03-14 03:38:48 -04:00
qazalandGitHub	43ffd66fda	viz: oneline inst list (#15269 ) * viz: oneline inst list * save 5 chars * gradual padding	2026-03-14 15:37:18 +09:00
George HotzandGitHub	86f17468ed	store in spec + USB BOT fix (#15265 ) * move spec to store * usb bot flag * Revert "usb bot flag" This reverts commit `7b8b7824f0`. * fix assert	2026-03-14 13:25:05 +08:00
George HotzandGitHub	06d7cddb33	amd_copy_matmul is cleaner (#15248 ) * amd_copy_matmul is cleaner * it runs * replicated stuff * add tid there * it runs * cleanup * x.src[1] * flatten * move that * keep that assert	2026-03-14 12:56:09 +08:00
chenyuandGitHub	b3600e4774	don't emit assign in transform_precompiled_call [pr] (#15262 )	2026-03-13 22:42:35 -04:00
qazalandGitHub	4d60312f7f	viz: asm python dsl syntax highlighting (#15259 )	2026-03-14 06:37:43 +09:00
qazalandGitHub	6209ddfc90	viz: improve disasm of s_code_end (#15258 ) * viz: improve amd disasm of s_code_end * better tests * order was good	2026-03-14 03:31:14 +09:00
wozeparrotandGitHub	a191ac0566	llama: use mlperf model (#15257 )	2026-03-13 08:08:32 -07:00
Sieds LyklesandGitHub	4b59083d7c	assign into empty works (#15256 )	2026-03-13 10:24:29 -04:00
qazalandGitHub	60b1b908c6	sqtt: CDNA layout header packet is the same size (#15255 )	2026-03-13 22:28:24 +09:00
nimlgenandGitHub	4e21735f31	system: update tinygpu app (#15247 )	2026-03-13 20:36:57 +08:00
nimlgenandGitHub	1fbe1fef2c	move write_configs to drivers (#15253 )	2026-03-13 19:02:34 +08:00
chenyuandGitHub	018c01508d	test case for call precompile multi (#15254 )	2026-03-13 06:28:43 -04:00
nimlgenandGitHub	bc16f80b50	am: remove dma_regions param (#15251 ) * am: remove dma_regions param * linter	2026-03-13 18:12:48 +08:00
chenyuandGitHub	576e7f985f	remove handle_assign_mops [pr] (#15249 )	2026-03-13 01:53:21 -04:00
sirhcmandGitHub	c251fc67c5	ci: consider arch in venv and apt caches and go back to 3.12 (#15250 )	2026-03-13 00:36:49 -04:00
sirhcmandGitHub	d4b947ea9a	ci: explicitly request python 3.12.10 instead of 3.12 (#15246 ) 3.12.10 is the most recent 3.12 version that has toolcache builds for linux, macos, and windows	2026-03-12 23:00:46 -04:00
George HotzandGitHub	a7d2429c21	amd_uop_matmul more cleanups (#15240 )	2026-03-13 10:24:43 +08:00
qazalandGitHub	d893b14193	sqtt: update cdna packet names (#15243 ) * sqtt: update cdna packet names * change * order	2026-03-13 08:49:09 +09:00
wozeparrotandGitHub	749162bd2f	llama memory tweaks (#15223 )	2026-03-12 12:36:23 -07:00
qazalandGitHub	9a7173b7a0	viz: visualize full range of shader clock frequency, auto zoom to kernel range (#15225 ) * start this * work * rm those * relative to start cycle * cleanup * cover the full range of packets * correct event type * start the ui change * fit=true * better * always the zoom identity * diff cleanup * shader engine itrace can be turned off	2026-03-13 00:07:31 +09:00
chenyuandGitHub	d9c09397c0	Ops.STORE is shapeless [pr] (#15239 )	2026-03-12 09:05:30 -04:00
nimlgenandGitHub	d746ccb791	system: fix vfio (#15235 )	2026-03-12 18:31:00 +08:00
nimlgenandGitHub	d104a903f8	system: print output when err (#15230 )	2026-03-12 18:30:49 +08:00
George HotzandGitHub	e560a46f59	update amd_uop_matmul (#15236 ) * update amd_uop_matmul * use custom kernel * simpler * ignore	2026-03-12 17:33:12 +08:00
chenyuandGitHub	90b7f4341d	failed two level divmod recombine case (#15233 )	2026-03-12 04:04:36 -04:00
chenyuandGitHub	8b8d9a443c	remove unused invalid rules [pr] (#15231 )	2026-03-12 03:10:34 -04:00
George HotzandGitHub	bdd62fd484	remove unneeded realize map entries (#15229 ) * remove unneeded realize map entries * not that	2026-03-12 14:23:19 +08:00
chenyuandGitHub	842c978df3	remove staticmethod dtypes.max/min (#15227 ) always use x.dtype.max/min	2026-03-11 23:11:24 -04:00
b1tg GitHub b1tg chenyu	18dc77ccab	add fp8 fnuz dtypes with PYTHON backend support (#14945 ) * add fp8 fnuz dtypes with PYTHON backend support * rm emu related change * clarify fp8 fnuz zero handling * Revert "rm emu related change" This reverts commit `efa4763c22`. --------- Co-authored-by: b1tg <[email protected]> Co-authored-by: chenyu <[email protected]>	2026-03-11 22:30:18 -04:00
George HotzandGitHub	4f3f55328b	do not patch on invalid tensor tests (#15226 ) * do not patch on invalid tensor tests * cleanup	2026-03-12 09:35:20 +08:00
wozeparrotandGitHub	4fab320abe	llama: clean (#15224 )	2026-03-11 13:33:59 -07:00
wozeparrotandGitHub	05d6d9120a	llama offload null (#15222 )	2026-03-11 10:04:31 -07:00
qazalandGitHub	d3eef70162	viz: render shader clock frequency graph (#15197 )	2026-03-12 01:32:49 +09:00
chenyuandGitHub	39b0f4bcc1	remove Ops.THREEFRY in remove_bufferize [pr] (#15220 )	2026-03-11 05:30:33 -04:00
chenyuandGitHub	6489a6f212	Revert "remove mop_cleanup [pr] (#15217 )" (#15218 ) This reverts commit `6b50df940a`.	2026-03-11 04:17:56 -04:00
chenyuandGitHub	6b50df940a	remove mop_cleanup [pr] (#15217 ) no kernel diff, i think this was needed due to force_reshape? test/external/external_benchmark_schedule.py is about the same speed	2026-03-11 03:54:42 -04:00