found_assign

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

@@ -66,14 +66,14 @@ runs:
       uses: actions/cache/restore@v4
       with:
         path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-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 }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-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 }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-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 }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-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')

examples/mlperf/model_train.py

@@ -1282,7 +1282,7 @@ def train_bert():
         previous_step = i
 
 def train_llama3():
-  from examples.mlperf.models.llama import Transformer
+  from extra.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,7 @@ 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)
+  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
 
   params = get_parameters(model)
   # weights are all bfloat16 for now
@@ -1417,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])
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
     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))
@@ -1449,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])
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
     loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
     return loss.flatten().float().to("CPU")
 

examples/mlperf/models/llama.py

@@ -1,80 +0,0 @@
-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

examples/mlperf/optim.py

@@ -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)
-      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+        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()
       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))
+        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)).realize()
 
     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)).cast(self.m[i].dtype)
-      v_hat = (self.v[i] / (1.0 - self.b2_t)).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)
       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]

extra/amdpci/am_smi.py

@@ -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 = vram_bar, doorbell_bar, mmio_bar
+    self.vram, self.doorbell64, self.mmio, self.dma_regions = vram_bar, doorbell_bar, mmio_bar, None
     self.pci_state = self.read_pci_state()
     if self.pci_state == "D0": self._init_from_d0()
 

extra/amdpci/hive_reset.py

@@ -2,7 +2,7 @@
 
 import os
 from tinygrad.helpers import Context
-from tinygrad.runtime.support.system import System, PCIDevice
+from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
 from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev
 
@@ -12,7 +12,7 @@ if __name__ == "__main__":
     drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
     if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
       raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
-  pcidevs = [PCIDevice("AM", gpu) for gpu in gpus]
+  pcidevs = [PCIDevice("AM", gpu, bars=[0, 2, 5]) for gpu in gpus]
   amdevs = []
   with Context(DEBUG=2):
     for pcidev in pcidevs:

extra/amdpci/proclogs.py

@@ -7,8 +7,8 @@ class GFXFake:
   def __init__(self): self.xccs = 8
 
 class AMDFake(AMDev):
-  def __init__(self, pci_dev):
-    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
+  def __init__(self, pci_dev, dma_regions=None):
+    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
     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()

extra/gemm/amd_copy_matmul.py

@@ -1,84 +0,0 @@
-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)

extra/gemm/amd_uop_matmul.py

@@ -1,74 +1,98 @@
-from tinygrad import Tensor, Context, GlobalCounters, dtypes
+import numpy as np
+from tinygrad import Tensor, Device, 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 DEBUG, getenv
+from tinygrad.helpers import getenv
 
 N = getenv("N", 4096)
-M = getenv("M", N)
-K = getenv("K", N)
-NUM_RUNS = getenv("CNT", 5)
+M = K = N
+run_count = getenv("CNT", 5)
 
 # ---------------------------
 # launch/config constants
 # ---------------------------
 
 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
+
+# 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
 
 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 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
 
-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"
+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"
 
 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, upcast=False):
+def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
   assert dest.shape == src.shape
   rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
-  return dest[*rngs].store(src[*rngs]).end(*rngs)
+  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
+  return dest.after(copy) if set else copy
 
-def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
+def hand_spec_kernel3():
   # ---------------------------
-  # block indices
+  # block indices & placeholders
   # ---------------------------
-  block_id_n = UOp.special(N // BLOCK_N, "gidx0")
-  block_id_m = UOp.special(M // BLOCK_M, "gidx1")
+  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)
 
   # 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, :]
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]
 
   # open the main reduction range
-  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, :]
+  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, :]
 
   # globals are no longer used, they are already in the indexes
-  del block_id_m, block_id_n
+  del blockIdx_y, blockIdx_x
 
   # ---------------------------
-  # GLOBAL -> LOCAL (A_local, B_local)
+  # GLOBAL -> LOCAL (As, Bs)
   # ---------------------------
   tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
 
   # A: read BM x BK tiles (permute on store into locals)
-  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)
+  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)
 
   # B: read BK x BN tiles
-  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)
+  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)
 
   # TODO: can we automate barrier?
-  barrier = UOp.barrier(A_local_store, B_local_store)
-  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+  barrier = UOp.barrier(As_store, Bs_store)
+  As, Bs = As.after(barrier), Bs.after(barrier)
 
   # open inner k range
   k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -76,30 +100,31 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   # ---------------------------
   # LOCAL -> REG (per-wave tiles)
   # ---------------------------
-  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
+  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
+  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
+  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_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))
+  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
+  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
+  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
 
-  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))
+  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)
 
   # ---------------------------
   # FMA: c_regs += A_col * B_row
   # ---------------------------
-  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 = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_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)
-  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)
+  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)
 
   # Close k, sync, and close K tiles
   sink = sink.end(k).barrier().end(k_tile_range)
@@ -107,37 +132,38 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   # ---------------------------
   # REG -> GLOBAL (epilogue)
   # ---------------------------
-  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.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[waveIdy, :, laneIdy, :,
         waveIdx, :, laneIdx, :]
   sink = copy(c, c_regs.after(sink), rng=600)
 
   return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
 
-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)
+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))
 
   ets = []
-  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}")
+  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}")
 
   if getenv("VERIFY", 1):
     GlobalCounters.reset()
     with Context(DEBUG=2):
       tc = (a @ b).realize()
     with Context(DEBUG=0):
-      err = (tc - tst).square().mean().item()
+      err = (hc - tc).square().mean().item()
     print(f"mean squared error {err}")
     if err > 1e-06:
       raise RuntimeError("matmul is wrong!")
 
 if __name__ == "__main__":
-  eval_custom_matmul(hand_spec_kernel3)
+  test_matmul(hand_spec_kernel3(), N=N)

extra/gemm/asm/cdna/asm.py

@@ -1,14 +1,5 @@
-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 *
-
-# ** CDNA4 assembly gemm
-
-WORKGROUP_SIZE = 256
+from tinygrad.dtype import dtypes
 
 # M0 is encoded with 124 (NULL in RDNA) in CDNA
 M0 = NULL
@@ -2606,117 +2597,3 @@ 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

extra/gemm/asm/cdna/gemm.py

@@ -0,0 +1,122 @@
+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

extra/gemm/asm/rdna3/template.s

@@ -0,0 +1,76 @@
+.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

extra/gemm/asm/rdna3/test.py

@@ -0,0 +1,30 @@
+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)

extra/gemm/asm/unpack_kd.py

@@ -0,0 +1,179 @@
+# 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)

extra/sqtt/examples/discover_ops.py

@@ -0,0 +1,148 @@
+#!/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)

extra/sqtt/examples/generate_examples.py

@@ -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:=32, N)@Tensor.empty(N, N)).realize()\"",
-  "sync":"test/amd/test_custom_kernel.py TestCustomKernel.test_wave_sync",
+  "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=64, N)@Tensor.empty(N, N)).realize()\"",
+  "ops":"extra/sqtt/examples/discover_ops.py"
 }
 
 if __name__ == "__main__":

extra/usbgpu/tbgpu/installer/Shared/server.c

@@ -17,7 +17,6 @@
 // Protocol
 
 enum {
-  CMD_PROBE = 0,          // probe devices, returns count
   CMD_MAP_BAR = 1,        // map PCI BAR, returns size
   CMD_MAP_SYSMEM_FD = 2,  // alloc DMA memory, returns fd via SCM_RIGHTS
   CMD_CFG_READ = 3,       // read PCI config space
@@ -25,14 +24,11 @@ enum {
   CMD_RESET = 5,          // reset device
   CMD_MMIO_READ = 6,      // bulk read from BAR
   CMD_MMIO_WRITE = 7,     // bulk write to BAR
-  CMD_MAP_SYSMEM = 8,     // map system memory
-  CMD_SYSMEM_READ = 9,    // bulk read from system memory
-  CMD_SYSMEM_WRITE = 10,  // bulk write to system memory
   RESP_OK = 0, RESP_ERR = 1,
 };
 
-typedef struct { uint8_t cmd; uint32_t dev_id, bar; uint64_t arg0, arg1, arg2; } __attribute__((packed)) request_t;
-typedef struct { uint8_t status; uint64_t resp0, resp1; } __attribute__((packed)) response_t;
+typedef struct { uint8_t cmd, bar; uint64_t offset, size, value; } __attribute__((packed)) request_t;
+typedef struct { uint8_t status; uint64_t value, addr; } __attribute__((packed)) response_t;
 
 // Constants and state
 
@@ -81,7 +77,7 @@ static int send_response(int fd, response_t *resp, int send_fd) {
 }
 
 static void send_error(int fd, const char *msg) {
-  response_t resp = {.status = RESP_ERR, .resp0 = strlen(msg)};
+  response_t resp = {.status = RESP_ERR, .value = strlen(msg)};
   send_response(fd, &resp, -1);
   send(fd, msg, strlen(msg), 0);
 }
@@ -120,8 +116,8 @@ static int dext_rpc(uint32_t sel, uint64_t *in, uint32_t in_cnt, uint64_t *out_v
 static int map_bar(uint32_t bar, response_t *resp) {
   if (bar >= MAX_BARS) return -1;
   if (!g_bars[bar].addr && IOConnectMapMemory64(g_conn, bar, mach_task_self(), &g_bars[bar].addr, &g_bars[bar].size, kIOMapAnywhere)) return -1;
-  resp->resp0 = g_bars[bar].addr;
-  resp->resp1 = g_bars[bar].size;
+  resp->addr = g_bars[bar].addr;
+  resp->value = g_bars[bar].size;
   return 0;
 }
 
@@ -152,7 +148,7 @@ static int map_sysmem_fd(uint64_t size, response_t *resp, int *out_fd) {
   strncpy(g_sysmem[idx].shm_name, shm_name, sizeof(g_sysmem[idx].shm_name));
   g_sysmem_count++;
 
-  *resp = (response_t){.resp0 = alloc_sz, .resp1 = idx};
+  *resp = (response_t){.addr = idx, .value = alloc_sz};
   *out_fd = fd;
   return 0;
 
@@ -207,19 +203,19 @@ static void handle_client(int fd) {
 
     case CMD_MAP_SYSMEM_FD: {
       int shm_fd = -1;
-      resp.status = map_sysmem_fd(req.arg0, &resp, &shm_fd) ? 1 : 0;
+      resp.status = map_sysmem_fd(req.size, &resp, &shm_fd) ? 1 : 0;
       send_response(fd, &resp, shm_fd);
       continue;
     }
 
     case CMD_CFG_READ: {
-      uint64_t in[2] = {req.arg0, req.arg1};
-      resp.status = dext_rpc(0, in, 2, &resp.resp0) ? 1 : 0;
+      uint64_t in[2] = {req.offset, req.size};
+      resp.status = dext_rpc(0, in, 2, &resp.value) ? 1 : 0;
       break;
     }
 
     case CMD_CFG_WRITE: {
-      uint64_t in[3] = {req.arg0, req.arg1, req.arg2};
+      uint64_t in[3] = {req.offset, req.size, req.value};
       resp.status = dext_rpc(1, in, 3, NULL) ? 1 : 0;
       break;
     }
@@ -229,17 +225,17 @@ static void handle_client(int fd) {
       break;
 
     case CMD_MMIO_READ:
-      if (validate_bar(req.bar, req.arg0, req.arg1)) { resp.status = 1; break; }
-      mmio_copy(g_bulk_buf, (void*)(g_bars[req.bar].addr + req.arg0), req.arg1);
-      resp.resp0 = req.arg1;
+      if (validate_bar(req.bar, req.offset, req.size)) { resp.status = 1; break; }
+      mmio_copy(g_bulk_buf, (void*)(g_bars[req.bar].addr + req.offset), req.size);
+      resp.value = req.size;
       send_response(fd, &resp, -1);
-      send(fd, g_bulk_buf, req.arg1, 0);
+      send(fd, g_bulk_buf, req.size, 0);
       continue;
 
     case CMD_MMIO_WRITE:
-      recvall(fd, g_bulk_buf, req.arg1);
-      if (!validate_bar(req.bar, req.arg0, req.arg1))
-        mmio_copy((void*)(g_bars[req.bar].addr + req.arg0), g_bulk_buf, req.arg1);
+      recvall(fd, g_bulk_buf, req.size);
+      if (!validate_bar(req.bar, req.offset, req.size))
+        mmio_copy((void*)(g_bars[req.bar].addr + req.offset), g_bulk_buf, req.size);
       continue;
 
     default:

test/amd/helpers.py

@@ -5,8 +5,8 @@ from tinygrad.runtime.autogen import llvm
 from tinygrad.runtime.support.elf import elf_loader
 
 ARCH_TO_TARGET:dict[str, list[str]] = {
-  "rdna3":["gfx1100", "gfx1151"],
-  "rdna4":["gfx1200", "gfx1201"],
+  "rdna3":["gfx1100"],
+  "rdna4":["gfx1200"],
   "cdna":["gfx950", "gfx942"],
 }
 

test/amd/test_custom_kernel.py

@@ -1,12 +1,10 @@
 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()
@@ -45,26 +43,9 @@ 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()
@@ -74,7 +55,6 @@ 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]
@@ -83,9 +63,5 @@ 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()

test/amd/test_sqtt_examples.py

@@ -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, CDNA_INST)
+                                     InstOp, InstOpRDNA4, print_packets, CDNA_WAVEEND)
 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, InstOpRDNA4.OTHER_LDS_1, InstOpRDNA4.OTHER_LDS_2}
+OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_5}
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # ROCPROF DECODER
@@ -153,10 +153,9 @@ 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_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}")
+        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}")
 
   expected: dict[str, list[int]] = {}  # override in subclasses
   def test_packet_counts(self):
@@ -211,22 +210,23 @@ class SQTTExamplesTestBase(unittest.TestCase):
 class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
   target = "gfx1100"
   expected = {
-    "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],
+    "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],
   }
 
 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()

test/amd/test_sqttmap.py

@@ -2,7 +2,7 @@
 import unittest, pickle
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG, getenv, temp
+from tinygrad.helpers import DEBUG, OSX, 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):
+def rocprof_inst_traces_match(sqtt, prg, target, pass_rocprof_err=False):
   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):
     rocprof_inst = next(rwaves_iter[info.wave][0])
     ref_pc = rocprof_inst.pc-prg.base
     # always check pc matches
-    assert ref_pc == info.pc, f"pc mismatch {ref_pc}:{disasm_map[rocprof_inst.pc]} != {info.pc}:{disasm(info.inst)}"
+    assert ref_pc == info.pc or pass_rocprof_err, 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):
-          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target)
+          # 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)
           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,12 +78,12 @@ 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):
-          # skip if there's no SQTT frequency data
           if not (timeline:=sqtt_timeline(event.blob, p.lib, target)): continue
-          if not (frequency:=[e.key for e in timeline if type(e).__name__ == "ProfilePointEvent" and e.name == "freq_hz"]): 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)
+          self.assertGreater(variance, 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
@@ -92,28 +92,13 @@ class TestSQTTMapBase(unittest.TestCase):
               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", "BARRIER", "BARRIER_SIGNAL"}: insts += 1
+              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE"}: 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"
 
-class TestSQTTMapCDNA(TestSQTTMapBase):
-  target = "gfx950"
-  def test_rocprof_inst_traces_match(self): self.skipTest("requires timestamp patching to match rocprof, currently it's off by a few cycles")
-
 if __name__ == "__main__":
   unittest.main()

test/backend/test_arange.py

@@ -8,7 +8,6 @@ from tinygrad.engine.schedule import ExecItem
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer import Estimates
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import needs_second_gpu
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, tensor, desired):
@@ -189,33 +188,6 @@ class TestIndexing(unittest.TestCase):
     for i in idx.flatten().numpy(): expected_grad[i] += 2
     np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
 
-  @needs_second_gpu
-  @unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
-  @Context(USE_ATOMICS=1, SPEC=1)
-  def test_embedding_backward_vocab_sharded(self):
-    from tinygrad.renderer.cstyle import CStyleLanguage
-    if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
-    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
-    vocab_size, embed_size = 1000, 128
-    bs, seqlen = 4, 256
-    idx = Tensor.randint(bs, seqlen, high=vocab_size)
-    emb = nn.Embedding(vocab_size, embed_size)
-    emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True)
-    gt = Tensor.zeros(bs, seqlen, embed_size)
-    Tensor.realize(idx, emb.weight, gt)
-    # compute expected grad on single device
-    expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
-    for i in idx.flatten().numpy(): expected_grad[i] += 2
-    # now shard the embedding weight on vocab axis and recompute
-    emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True)
-    emb.weight.shard_(devices, axis=0)
-    idx = idx.shard(devices, axis=None)
-    gt = gt.shard(devices, axis=None)
-    Tensor.realize(idx, emb.weight, gt)
-    loss = (emb(idx)-gt).square().sum()
-    loss.backward()
-    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
-
   @unittest.skipUnless(Device.DEFAULT == "AMD" or (Device.DEFAULT == "NULL" and EMULATE.value.startswith("AMD")), "tests AMD bf16 cast overhead")
   def base_test_llama_8b_rope_backward(self, dtype):
     from extra.models.llama import precompute_freqs_cis, apply_rotary_emb

test/backend/test_asm_gemm.py

@@ -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.cdna_asm_gemm import asm_gemm
+from extra.gemm.asm.cdna.gemm import asm_gemm
 from test.helpers import needs_second_gpu
 
 # On non CDNA4 it will only validate the Tensor.custom_kernel integration
@@ -87,9 +87,6 @@ class TestGemmLarge(unittest.TestCase):
     if not is_cdna4():
       self.skipTest("very slow on non mi350x")
 
-  @Context(ASM_GEMM=1)
-  def test_empty(self): (Tensor.empty(N:=getenv("N", 4096), N, dtype=dtypes.half)@Tensor.empty(N, N, dtype=dtypes.half)).realize()
-
   def test_tiny(self): verify_asm_gemm(1, 256, 256, 64)
   def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=dtypes.half)
   def test_gemm(self): verify_asm_gemm(1, 8192, 4096, 14336)
@@ -160,7 +157,7 @@ class TestGemmLarge(unittest.TestCase):
 
 class TestMagicGu(unittest.TestCase):
   def test_magicgu_matches_old(self):
-    from extra.gemm.cdna_asm_gemm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
+    from extra.gemm.asm.cdna.asm 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),

test/null/test_simplify_valid_idx.py

@@ -46,9 +46,9 @@ class TestHelpers(unittest.TestCase):
     self.assertTrue((rng+2).is_increasing())
 
 class TestValidIdxSimplification(unittest.TestCase):
-  def check(self, load, sidx, svalid, extra=()):
+  def check(self, load, sidx, svalid):
     with Context(NOOPT=1, SPEC=0):
-      load = full_rewrite_to_sink(UOp.sink(load, *extra)).src[0]
+      load = full_rewrite_to_sink(load.sink()).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,12 +156,9 @@ 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))",
-      extra=(red,))
+      "((r2<1)&(r1<6))")
 
   def test_valid_becomes_const1_z3(self):
     from z3 import Ints, Solver, And, If, Not, unsat
@@ -486,79 +483,5 @@ 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()

test/null/test_uop_graph.py

@@ -423,12 +423,10 @@ 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)
-    # prevent ridx0 from being shrunk
-    red = UOp(Ops.REDUCE, dtypes.long, (ridx0.cast(dtypes.long), ridx0), Ops.ADD)
-    uops = to_uops_list([w, red])
+    uops = to_uops_list([w])
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg==5
+      if u.op is Ops.LOAD: assert u.src[1].arg==5
 
   def test_where_on_gated_load_folds_swapped_branches(self):
     ridx0 = UOp.range(100, 0)
@@ -446,12 +444,10 @@ 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)
-    # prevent ridx0 from being shrunk
-    red = UOp(Ops.REDUCE, dtypes.long, (ridx0.cast(dtypes.long), ridx0), Ops.ADD)
-    uops = to_uops_list([w, red])
+    uops = to_uops_list([w])
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg == 5
+      if u.op is Ops.LOAD: assert u.src[1].arg == 5
 
   def test_where_in_store_becomes_gate(self):
     ridx0 = UOp.range(100, 0)

test/null/test_uop_symbolic.py

@@ -786,14 +786,6 @@ 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)

test/testextra/test_cfg_viz.py

@@ -58,8 +58,7 @@ class TestCfg(unittest.TestCase):
     k.emit(s_endpgm())
     k.emit(s_code_end())
     ei = run_asm("diamond", k)
-    ret = amdgpu_cfg(ei.prg.p.lib, self.arch)
-    cfg = ret["data"]
+    cfg = amdgpu_cfg(ei.prg.p.lib, self.arch)["data"]
     self.assertEqual(len(cfg["blocks"]), 5)
     edge_count = sum(len(v) for v in cfg["paths"].values())
     self.assertEqual(edge_count, 5)
@@ -70,11 +69,8 @@ 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 = [" ".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
+    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)")
 
   def test_loop(self):
     k = Kernel(arch=Device["AMD"].arch)

test/unit/test_assign.py

@@ -272,19 +272,14 @@ class TestAssign(unittest.TestCase):
 
   def test_assign_after(self):
     t = Tensor.zeros(10).contiguous().realize()
-    t.uop = t.uop.after(t.uop.store((t+1).uop))
+    t.uop = t.uop.after(t.uop.assign((t+1).uop))
     np.testing.assert_allclose(t.numpy(), [1.,1.,1.,1.,1.,1.,1.,1.,1.,1.])
 
   def test_assign_after_partial(self):
     t = Tensor.zeros(10).contiguous().realize()
-    t.uop = t.uop.after(t[:5].uop.after(t[:5].uop.store(Tensor.ones(5).uop)))
+    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)

test/unit/test_call.py

@@ -237,116 +237,5 @@ 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)
-
-class TestCallMultiSharded(unittest.TestCase):
-  # TODO: multi-output + sharded needs per-device CALL execution, which requires reworking how MULTI propagates through TUPLE bodies
-  def test_tuple_sharded(self):
-    """multi-output function with sharded input"""
-    devs = ("CPU:0", "CPU:1")
-    @function
-    def f(x:Tensor): return (x + 1, x * 2)
-    a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
-    t1, t2 = f(a)
-    ref = np.arange(8, dtype=np.float32).reshape(4, 2)
-    np.testing.assert_allclose(t1.numpy(), ref + 1)
-    np.testing.assert_allclose(t2.numpy(), ref * 2)
-
-  def test_tuple_sharded_precompile(self):
-    """multi-output precompiled function with sharded input"""
-    devs = ("CPU:0", "CPU:1")
-    @function(precompile=True)
-    def f(x:Tensor): return (x + 1, x * 2)
-    a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
-    t1, t2 = f(a)
-    ref = np.arange(8, dtype=np.float32).reshape(4, 2)
-    np.testing.assert_allclose(t1.numpy(), ref + 1)
-    np.testing.assert_allclose(t2.numpy(), ref * 2)
-
-  def test_tuple_sharded_different_axis(self):
-    """multi-output function where outputs have different sharding: one reduces on sharded axis, one doesn't"""
-    devs = ("CPU:0", "CPU:1")
-    @function
-    def f(x:Tensor): return (x.sum(axis=0), x.sum(axis=1))
-    a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
-    t1, t2 = f(a)
-    ref = np.arange(8, dtype=np.float32).reshape(4, 2)
-    np.testing.assert_allclose(t1.numpy(), ref.sum(axis=0))
-    np.testing.assert_allclose(t2.numpy(), ref.sum(axis=1))
-
-  def test_tuple_sharded_different_ops(self):
-    """multi-output function with different operations per output"""
-    devs = ("CPU:0", "CPU:1")
-    @function
-    def f(x:Tensor, y:Tensor): return (x + y, x * y)
-    a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
-    b = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0) + 1
-    t1, t2 = f(a, b)
-    ref_a = np.arange(8, dtype=np.float32).reshape(4, 2)
-    ref_b = ref_a + 1
-    np.testing.assert_allclose(t1.numpy(), ref_a + ref_b)
-    np.testing.assert_allclose(t2.numpy(), ref_a * ref_b)
-
-  def test_tuple_sharded_mixed_use(self):
-    """multi-output sharded results used in further computation"""
-    devs = ("CPU:0", "CPU:1")
-    @function
-    def f(x:Tensor): return (x + 1, x * 2)
-    a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
-    t1, t2 = f(a)
-    out = (t1 + t2).sum()
-    ref = np.arange(8, dtype=np.float32).reshape(4, 2)
-    np.testing.assert_allclose(out.numpy(), ((ref + 1) + (ref * 2)).sum())
-
-  def test_tuple_sharded_outputs_different_axis(self):
-    """multi-output function where the two outputs are sharded on different axes"""
-    devs = ("CPU:0", "CPU:1")
-    @function
-    def f(x:Tensor, y:Tensor): return (x + 1, y + 2)
-    a = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=0)
-    b = Tensor.arange(8).reshape(4, 2).float().shard(devs, axis=1)
-    t1, t2 = f(a, b)
-    ref_a = np.arange(8, dtype=np.float32).reshape(4, 2)
-    ref_b = np.arange(8, dtype=np.float32).reshape(4, 2)
-    np.testing.assert_allclose(t1.numpy(), ref_a + 1)
-    np.testing.assert_allclose(t2.numpy(), ref_b + 2)
-
-  def test_call_reduce_sharded(self):
-    devs = ("CPU:0", "CPU:1")
-    a = Tensor.ones(10, 10).shard(devs, axis=0)
-    Tensor.realize(a)
-    c = Tensor.call(a, fxn=a.as_param(0).sum(axis=0))
-    np.testing.assert_equal(c.numpy(), 10 * np.ones(10))
-
-  def test_call_reduce_sharded_mixed_args(self):
-    devs = ("CPU:0", "CPU:1")
-    a = Tensor.ones(10, 10).shard(devs, axis=0)
-    b = Tensor.ones(10).shard(devs, axis=None)
-    Tensor.realize(a, b)
-    c = Tensor.call(a, b, fxn=a.as_param(0).sum(axis=0) + b.as_param(1))
-    np.testing.assert_equal(c.numpy(), 11 * np.ones(10))
-
-  def test_call_reduce_sharded_backward(self):
-    devs = ("CPU:0", "CPU:1")
-    a = Tensor.randn(10, 10, requires_grad=True).shard(devs, axis=0)
-    b = Tensor.randn(10, 10, requires_grad=True).shard(devs, axis=0)
-    Tensor.realize(a, b)
-
-    def grad_fxn(grad, call):
-      a_arg, b_arg = call.src[1], call.src[2]
-      return (grad.expand(a_arg.shape) * b_arg, grad.expand(b_arg.shape) * a_arg)
-
-    body = (a.as_param(0) * b.as_param(1)).sum(axis=0)
-    c = Tensor.call(a, b, fxn=body, grad_fxn=grad_fxn)
-    c.sum().backward()
-    np.testing.assert_allclose(a.grad.numpy(), b.numpy(), rtol=1e-5)
-    np.testing.assert_allclose(b.grad.numpy(), a.numpy(), rtol=1e-5)
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_function.py

@@ -1,7 +1,7 @@
 import numpy as np
 import unittest
 from tinygrad.function import function
-from tinygrad import Tensor, GlobalCounters
+from tinygrad import Tensor
 from tinygrad.uop.ops import UOp
 
 class TestFunction(unittest.TestCase):
@@ -165,13 +165,13 @@ class TestFunction(unittest.TestCase):
   def test_name(self):
     @function
     def f(a:Tensor) -> Tensor: return a + 1
-    assert f(Tensor([1])).uop.src[0].arg.name.endswith("f")
+    assert f(Tensor([1])).uop.arg.name.endswith("f")
 
   def test_method_name(self):
     class Foo:
       @function
       def __call__(self, x:Tensor) -> Tensor: return x + 1
-    assert Foo()(Tensor([1])).uop.src[0].arg.name.endswith("Foo.__call__")
+    assert Foo()(Tensor([1])).uop.arg.name.endswith("Foo.__call__")
 
   def test_callable_instance(self):
     class Foo:
@@ -180,7 +180,7 @@ class TestFunction(unittest.TestCase):
     foo = Foo()
     f = function(foo)
     np.testing.assert_equal(f(Tensor([1,2,3])).numpy(), [11,22,33])
-    assert f(Tensor([1,2,3])).uop.src[0].arg.name.endswith("Foo")
+    assert f(Tensor([1,2,3])).uop.arg.name.endswith("Foo")
 
   def test_iadd(self):
     @function
@@ -213,17 +213,6 @@ class TestFunction(unittest.TestCase):
     np.testing.assert_equal(a.numpy(), [11,21,31])  # TODO: should be [1,2,3]
     np.testing.assert_equal(b.numpy(), [10,20,30])
 
-  def test_view_assign_explicit_buffer(self):
-    """view assign on an explicit param's buffer should not create implicit inputs."""
-    class State:
-      def __init__(self): self.buf = Tensor.zeros(2, 4).contiguous().realize()
-      @function(allow_implicit=False)
-      def __call__(self, x:Tensor) -> Tensor:
-        self.buf[:, 0:2].assign(x)
-        return self.buf[:, 0:2]
-    s = State()
-    np.testing.assert_equal(s(Tensor([[5., 6.], [7., 8.]])).numpy(), [[5., 6.], [7., 8.]])
-
   @unittest.expectedFailure
   def test_assign_slice(self):
     @function
@@ -346,69 +335,5 @@ 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)
-
-  def test_grad_tuple(self, precompile=False):
-    x = Tensor.ones(3, requires_grad=True).contiguous()
-    y = Tensor.ones(3, requires_grad=True).contiguous()
-    @function(precompile=precompile)
-    def f(u1:Tensor, u2:Tensor): return (u1+1, u2+2)
-    t1, t2 = f(x,y)
-    (t1+t2).sum().backward()
-    x.grad.realize(y.grad)
-  def test_grad_tuple_precompile(self): self.test_grad_tuple(True)
-
-  def test_grad_fxn_tuple(self):
-    # grad_fxn for tuple: receives one gradient per output as positional args
-    def grad_fxn(d_out0:UOp, d_out1:UOp, call:UOp):
-      # f(u1, u2) = (u1+1, u2+2)
-      # df/du1 = d_out0, df/du2 = d_out1
-      return (d_out0, d_out1)
-
-    x = Tensor.ones(3, requires_grad=True).contiguous()
-    y = Tensor.ones(3, requires_grad=True).contiguous()
-    @function(grad_fxn=grad_fxn)
-    def f(u1:Tensor, u2:Tensor): return (u1+1, u2+2)
-    t1, t2 = f(x, y)
-    (t1+t2).sum().backward()
-    np.testing.assert_allclose(x.grad.numpy(), [1., 1., 1.])
-    np.testing.assert_allclose(y.grad.numpy(), [1., 1., 1.])
-
-class TestFunctionGrad(unittest.TestCase):
-  def test_function_grad_ops(self, precompile=False, precompile_backward=False):
-    N = 64
-    x = Tensor.ones(N,N).contiguous()
-    w1 = Tensor.ones(N,N, requires_grad=True).contiguous()
-    w2 = Tensor.ones(N,N, requires_grad=True).contiguous()
-    w3 = Tensor.ones(N,N, requires_grad=True).contiguous()
-    ref = Tensor.ones(N,N).contiguous()
-    Tensor.realize(x, w1, w2, w3, ref)
-    @function(precompile=precompile, precompile_backward=precompile_backward)
-    def f(x, w1, w2, w3) -> tuple[Tensor, ...]:
-      p1 = x@w1
-      p2 = p1@w2
-      p3 = p2@w3
-      return p1, p2, p3, p3.contiguous()
-    ret = f(x, w1, w2, w3)[-1]
-    loss = (ret-ref).square().mean().backward()
-    print("RESET")
-    GlobalCounters.reset()
-    loss.realize(w1.grad, w2.grad, w3.grad)
-    print(GlobalCounters.global_ops, GlobalCounters.global_mem)
-    self.assertLessEqual(GlobalCounters.global_ops, 4739073)
-  def test_function_grad_ops_precompile(self): self.test_function_grad_ops(precompile=True)
-  def test_function_grad_ops_precompile_backward(self):
-    self.test_function_grad_ops(precompile=True, precompile_backward=True)
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_llm_moe.py

@@ -49,26 +49,5 @@ class TestMoEFeedForward(unittest.TestCase):
     expected = 1 + (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
     np.testing.assert_allclose(out.numpy(), expected, rtol=1e-2)
 
-  def test_moe_feed_forward_norm_topk_prob(self):
-    from tinygrad.apps.llm import TransformerBlock
-    dim, hidden, n_heads = 8, 16, 2
-    num_experts, k = 4, 2
-
-    block = TransformerBlock(dim, hidden, n_heads, n_heads, norm_eps=1e-5, head_dim=dim//n_heads,
-                             rope_theta=10000, max_context=16, num_experts=num_experts, num_experts_per_tok=k)
-    block.norm_topk_prob = True
-
-    block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
-    block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
-    block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
-    block.ffn_gate_inp.weight = Tensor([[0.1, 0, 0.1, 0]] * dim).T  # equal top-2 experts, but only ~69% mass before renorm
-    block.ffn_norm.weight = Tensor.ones(dim)
-
-    h = Tensor.ones(1, 1, dim)
-    out = block._feed_forward(h)
-
-    expected = 1 + (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
-    np.testing.assert_allclose(out.numpy()[0, 0, 0], expected, rtol=1e-2)
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_setitem_schedule.py

@@ -36,7 +36,8 @@ class TestSetitemInto(unittest.TestCase):
     self.assertEqual(GlobalCounters.kernel_count, 0)
     t.realize()
     self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 4)
+    # 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
     t[1].realize()
     t.realize()
     self.assertEqual(GlobalCounters.kernel_count, 1)

tinygrad/apps/llm.py

@@ -84,7 +84,7 @@ def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
 
 class TransformerBlock:
   def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, head_dim:int, rope_theta:float,
-               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0, norm_topk_prob:bool=False):
+               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0):
     self.n_heads      = n_heads
     self.n_kv_heads   = n_kv_heads
     self.head_dim     = head_dim
@@ -107,7 +107,6 @@ class TransformerBlock:
 
     # --- feed-forward (MoE or dense) -------------------------------------
     if num_experts > 0:
-      self.norm_topk_prob = norm_topk_prob
       self.num_experts_per_tok = num_experts_per_tok
       self.ffn_gate_inp = nn.Linear(dim, num_experts, bias=False)  # router
       self.ffn_gate_exps = ExpertWeights(num_experts, dim, hidden_dim)
@@ -118,7 +117,7 @@ class TransformerBlock:
       self.ffn_up      = nn.Linear(dim, hidden_dim, bias=False)
       self.ffn_down    = nn.Linear(hidden_dim, dim, bias=False)
 
-  @function(precompile=bool(getenv("PRECOMPILE", 0)), allow_implicit=False)
+  @function(precompile=bool(getenv("PRECOMPILE", 0)))
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     x_norm = self.attn_norm(x)                       # (B,T,D)
     q, k, v = self.attn_q(x_norm), self.attn_k(x_norm), self.attn_v(x_norm)
@@ -130,12 +129,19 @@ class TransformerBlock:
     v = v.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
     if self.qk_norm == self.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
 
-    q = apply_rope(q, self.freqs_cis[start_pos:start_pos+T])
-    k = apply_rope(k, self.freqs_cis[start_pos:start_pos+T])
+    freqs_cis = precompute_freqs_cis(self.head_dim, self.max_context, self.rope_theta)[start_pos:start_pos+T]
+    q = apply_rope(q, freqs_cis)
+    k = apply_rope(k, freqs_cis)
 
-    self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
-    k = self.cache_kv[0, :, :, 0:start_pos+T, :]
-    v = self.cache_kv[1, :, :, 0:start_pos+T, :]
+    # TODO: fix assign to behave like this
+    assigned_kv = self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.assign(Tensor.stack(k, v).contiguous().uop))
+    tensor_assigned_kv = Tensor(assigned_kv, device=assigned_kv.device)
+    k = tensor_assigned_kv[0, :, :, 0:start_pos+T, :]
+    v = tensor_assigned_kv[1, :, :, 0:start_pos+T, :]
+
+    #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
+    #k = self.cache_kv[0, :, :, 0:start_pos+T, :]
+    #v = self.cache_kv[1, :, :, 0:start_pos+T, :]
 
     # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
     # TODO: this if statement should be removed and it shouldn't generate extra kernels
@@ -145,13 +151,12 @@ class TransformerBlock:
     attn = self.attn_output(attn)
     return x + attn
 
-  @function(precompile=bool(getenv("PRECOMPILE", 0)), allow_implicit=False)
+  @function(precompile=bool(getenv("PRECOMPILE", 0)))
   def _feed_forward(self, h: Tensor) -> Tensor:
     h_norm = self.ffn_norm(h)
     if hasattr(self, 'ffn_gate_exps'):
       x = h_norm.unsqueeze(2)  # (B, T, 1, D) - add expert dim for broadcasting
       probs, sel = self.ffn_gate_inp(h_norm).softmax(-1).topk(self.num_experts_per_tok)  # (B, T, k) each
-      if self.norm_topk_prob: probs = probs / probs.sum(axis=-1, keepdim=True)
       x_down = self.ffn_down_exps(sel, self.ffn_gate_exps(sel, x).silu() * self.ffn_up_exps(sel, x))  # (B, T, k, D)
       return h + (x_down * probs.unsqueeze(-1)).sum(axis=2)  # (B, T, D)
     # TODO: remove the need for this contiguous
@@ -163,14 +168,13 @@ class TransformerBlock:
       # TODO: how is the dtype of this determined?
       # NOTE: clone is used to promise the creation of a specific buffer
       self.cache_kv = Tensor.zeros(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device).clone()
-      self.freqs_cis = precompute_freqs_cis(self.head_dim, self.max_context, self.rope_theta)
     return self._feed_forward(self._attention(x, start_pos)).contiguous()
 
 class Transformer:
   def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, head_dim:int, rope_theta:float,
-               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0, norm_topk_prob:bool=False):
+               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0):
     self.blk = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, head_dim, rope_theta, max_context, qk_norm,
-                                 num_experts, num_experts_per_tok, norm_topk_prob) for _ in range(num_blocks)]
+                                 num_experts, num_experts_per_tok) for _ in range(num_blocks)]
     self.token_embd  = nn.Embedding(vocab_size, dim)
     self.output_norm = nn.RMSNorm(dim, norm_eps)
     self.output = nn.Linear(dim, vocab_size, bias=False)
@@ -217,8 +221,7 @@ class Transformer:
                         head_dim=kv.get(f'{arch}.attention.key_length', kv[f'{arch}.embedding_length'] // n_heads),
                         rope_theta=kv[f'{arch}.rope.freq_base'], max_context=max_context,
                         qk_norm=int(state_dict['blk.0.attn_q_norm.weight'].shape[0]) if 'blk.0.attn_q_norm.weight' in state_dict else 0,
-                        num_experts=kv.get(f'{arch}.expert_count', 0), num_experts_per_tok=kv.get(f'{arch}.expert_used_count', 0),
-                        norm_topk_prob=True if arch=='qwen3moe' else False)
+                        num_experts=kv.get(f'{arch}.expert_count', 0), num_experts_per_tok=kv.get(f'{arch}.expert_used_count', 0))
     nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
     # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
     if realize:

tinygrad/codegen/__init__.py

@@ -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, pm_store_ranges
+from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar
 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+pm_store_ranges, ctx=itertools.count(1000), name="early movement ops", bottom_up=True)
+  sink = graph_rewrite(sink, pm_mops+pm_syntactic_sugar, 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_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges")
+    sink = graph_rewrite(sink, pm_simplify_ranges, name="simplify ranges")
 
     # do postrange optimization, BEAM or hand_coded_optimizations
     sink = apply_opts(sink, ren)

tinygrad/codegen/late/devectorizer.py

@@ -10,7 +10,6 @@ from tinygrad.renderer import Renderer
 
 # ***** image load valid simplification *****
 
-@functools.cache
 def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
   # can drop valid if idx is out of bound when valid is False
   drop_stmt = []
@@ -62,18 +61,6 @@ load_store_indexing = PatternMatcher([
 # ***** load/store grouping *****
 
 def expand_index(buf:UOp, vec:UOp):
-  # determine optimal image shapes
-  if IMAGE == 1 and isinstance(dt:=buf.dtype, ImageDType):
-    x, valid = vec.get_idx().gep(0), vec.get_valid().gep(0)
-    # search for dims that drop the most valid statements
-    best_drop, cands = -1, []
-    for ch, cw in ImageDType.valid_dims(dt):
-      if (dropped:=len(_drop_valid_stmts(valid, cidx:=uop_given_valid(valid, UOp.vectorize((x//4)%cw, x//(4*cw))), ch, cw))) > best_drop:
-        best_drop, cands = dropped, [(ch, cw, cidx)]
-      elif dropped == best_drop: cands.append((ch, cw, cidx))
-    # and tiebreak with indexing complexity (ie. number of nodes)
-    h, w, _ = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].gep(1).simplify().backward_slice))
-    buf = buf.replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4), w * 4 * dt.itemsize))
   if getenv("UNSAFE_DISABLE_MASK", 0): vec = vec.get_idx()
   # generate the individual indexes
   return UOp(Ops.VECTORIZE, buf.dtype, tuple(buf.index(vec.gep(i), ptr=True) for i in range(vec.dtype.count)))
@@ -196,21 +183,31 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
   if len(ret) <= 1: return None
   return UOp(Ops.VCAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
 
-def get_image_idx(idx:UOp, width:int):
-  oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), (((x:=idx.src[1].get_idx()) // 4) % width, (x // (4*width))))
-  return idx.replace(src=(idx.src[0], oidx.valid(idx.src[1].get_valid())))
+def _do_image_fixup(dt:ImageDType, idx:UOp) -> tuple[UOp, UOp, int, int]:
+  buf = idx.src[0]
+  x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
+  h, w = dt.shape[0], dt.shape[1]
+  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)),
+                # and minimize idx complexity (number of nodes)
+                -len(idx.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
 
 def image_fixup(ls:UOp):
   # normal image load or store, with the CAST from expand_index
   if ls.src[0].op is Ops.CAST and isinstance(image_dtype:=ls.src[0].src[0].dtype, ImageDType):
     assert ls.src[0].dtype.count == 4, "image must be casted to 4"
-    idx = get_image_idx(ls.src[0].src[0], image_dtype.shape[1])
+    _, idx, _, _ = _do_image_fixup(image_dtype, ls.src[0].src[0])
     return ls.replace(src=(idx,)+ls.src[1:])
 
   # this is an unprocessed image without a cast, aka unfoldable image load. this doesn't work for stores
   if isinstance(image_dtype:=ls.src[0].dtype, ImageDType) and ls.src[0].src[1].get_idx().dtype != dtypes.index.vec(2):
     assert ls.op is Ops.LOAD, "if an image store isn't upcasted to 4, we can't store it"
-    x, idx = ls.src[0].src[1].get_idx(), get_image_idx(ls.src[0], image_dtype.shape[1])
+    x, idx, width, height = _do_image_fixup(image_dtype, ls.src[0])
     vec_load = ls.replace(dtype=ls.dtype.vec(4), src=(idx,)+ls.src[1:])
     # image pixels have 4 channels (.xyzw), select channel based on x % 4
     x_mod_4 = x % 4
@@ -251,27 +248,32 @@ def no_vectorized_alu(alu:UOp):
 def no_vectorized_buf(buf:UOp):
   return buf.replace(dtype=buf.ptrdtype.base.scalar().ptr(buf.ptrdtype.size*buf.ptrdtype.count, buf.ptrdtype.addrspace)).cast(buf.dtype)
 
-def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp, bcast:UOp|None=None):
+def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp):
   cnt = cast.dtype.count
-  if bcast is not None and bcast.op is Ops.GEP:
-    # GEP selects specific lanes; bcast.arg[k] is the offset for lane k, iterate groups × selected lanes
-    pairs = [(k, g + bcast.arg[k]) for g, k in itertools.product(range(cast.dtype.vcount), range(len(bcast.arg)))]
-  elif bcast is not None:
-    # BROADCAST: cross product of components × lanes
-    pairs = [(j, c) for c, j in itertools.product(range(cnt), range(bcast.dtype.vcount))]
+  assert idx.dtype.count == 1, f"idx dtype must be 1 {idx.dtype}"
+  return buf.broadcast(cnt).index(idx.broadcast(cnt)*cnt+UOp.const(dtypes.index.vec(cnt), tuple(range(cnt))), ptr=True)
+
+def no_vectorized_index_broadcast(buf:UOp, cast:UOp, bcast:UOp, idx:UOp):
+  cnt = cast.dtype.count
+  vcnt = cast.dtype.vcount
+  precnt = bcast.dtype.vcount
+  # TODO: I have no idea *why* this is. I just change things until the tests pass. No AI, old school.
+  if bcast.op is Ops.GEP:
+    gep_arg = tuple(flatten([range(precnt) for _ in range(vcnt)]))
+    sum_arg = tuple(flatten([[i+y for y in bcast.arg] for i in range(vcnt)]))
   else:
-    # simple scalar index: one lane, all components
-    pairs = [(0, c) for c in range(cnt)]
-  idx_lanes, offsets = (tuple(x) for x in zip(*pairs))
-  return buf.broadcast(len(pairs)).index(idx.gep(idx_lanes)*cnt + UOp.const(dtypes.index.vec(len(pairs)), offsets), ptr=True)
+    gep_arg = tuple(flatten([range(precnt) for _ in range(cnt)]))
+    sum_arg = tuple(flatten([[i]*precnt for i in range(cnt)]))
+  new_idx = idx.gep(gep_arg)*cnt + UOp.const(dtypes.index.vec(len(sum_arg)), sum_arg)
+  return buf.broadcast(cnt*precnt).index(new_idx, ptr=True)
 
 devectorize_buf_and_index = PatternMatcher([
   (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf"), no_vectorized_buf),
   (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").index(UPat.var("idx")), no_vectorized_index),
   (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").broadcast(name="bcast").index(UPat.var("idx")),
-   no_vectorized_index),
+   no_vectorized_index_broadcast),
   (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").gep(name="bcast").index(UPat.var("idx")),
-   no_vectorized_index),
+   no_vectorized_index_broadcast),
 ])
 
 devectorize = PatternMatcher([

tinygrad/codegen/simplify.py

@@ -1,5 +1,4 @@
 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
@@ -37,32 +36,22 @@ 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, 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})
+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()
   ctx.clear()
-  return None if ret is x else ret.simplify()
+  return ret
 
 def dont_sub_ranges_for_image(ctx:dict[UOp, UOp|None], x:UOp) -> None:
   if isinstance(x.src[0].src[0].dtype, ImageDType):
@@ -71,8 +60,7 @@ 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"), 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])))),
+  (UPat(Ops.SINK, name="x"), do_substitute),
 ])
 
 # **** reduce simplification ****

tinygrad/engine/allocations.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
 from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, track_rewrites
-from tinygrad.dtype import ImageDType
-from tinygrad.helpers import prod, DEBUG, VIZ, pluralize, all_int
+from tinygrad.dtype import dtypes, ImageDType
+from tinygrad.helpers import prod, DEBUG, VIZ, pluralize
 
 @dataclass
 class AllocCtx:
@@ -29,15 +29,14 @@ def apply_after(ctx:AllocCtx, u:UOp):
   while base.op is Ops.AFTER: base = base.src[0]
   ctx.buffer_map[u] = base
 
-# CONTIGUOUS and AFTER+STORE + parents are the only nodes that get updated
+# CONTIGUOUS and ASSIGN + parents are the only nodes that get updated
 add_tags = PatternMatcher([
   (UPat(Ops.COPY, name="u"), disk_copy_is_buffer),
-  # no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
-   lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None),
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE)), name="x"), tag_uop),
+  # no tag on copies that are assigned
+  (UPat(Ops.ASSIGN, src=(UPat(), UPat(Ops.COPY, name="c")), name="a"),
+   lambda a,c: a.replace(src=(a.src[0], c.rtag(())), tag=a.tag+c.tag) if a.tag and c.tag else None),
   (UPat(Ops.AFTER, name="u"), apply_after),
-  (UPat(Ops.CONTIGUOUS, name="x"), tag_uop),
+  (UPat({Ops.CONTIGUOUS, Ops.ASSIGN}, name="x"), tag_uop),
   (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(ctx,x) if x in ctx.bases else None),
 ])
 
@@ -47,33 +46,34 @@ 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).shrink_to(u.shard_shape)
+  buffer = UOp.new_buffer(u.device, u.shard_size, dtype).reshape(u.max_shard_shape)
   if isinstance(u.device, tuple) and u.axis is not None: buffer = buffer.multi(u.axis)
   return buffer
 
-def replace_contig_with_store_after(u:UOp):
+def replace_contig_with_assign(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)
-  buf = _buffer_like(u)
-  return buf.after(buf.store(u.src[0])).rtag(u.tag)
+  return _buffer_like(u).assign(u.src[0]).rtag(u.tag)
 
-def replace_store_after_with_contig(u:UOp, src:UOp):
+def replace_assign_with_contig(u:UOp):
   assigned_to = u
-  while assigned_to.op in {Ops.BITCAST, Ops.AFTER}: assigned_to = assigned_to.src[0].base
-  if assigned_to.op is not Ops.BUFFER: return src.contiguous(tag=u.tag)
+  while assigned_to.op in {Ops.ASSIGN, Ops.BITCAST, Ops.AFTER}: assigned_to = assigned_to.src[0].base
+  if assigned_to.op is not Ops.BUFFER:
+    return u.src[1].contiguous(tag=u.tag)
 
-def contiguous_mops_to_view(c:UOp, src:UOp):
+def contiguous_mops_to_view(c:UOp):
   """CONTIGUOUS(MOPS(BUFFER)) → CONTIGUOUS(BUFFER_VIEW) when movement ops collapse to a contiguous range."""
+  src = c.src[0]
   buf = src.base
   if buf.op not in {Ops.BUFFER, Ops.BUFFER_VIEW}: return None
   if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.BUFFER_VIEW}: return None
 
   # no symbolic shape
-  if not all_int(c.shape): return None
+  if not all(isinstance(x, int) for x in c.shape): return None
 
   # check if view is supported
   if not isinstance(c.device, str): return None
@@ -81,7 +81,8 @@ def contiguous_mops_to_view(c:UOp, src:UOp):
   if not hasattr(Device[c.device].allocator, "_offset"): return None
 
   # see if this can be a view
-  if (offset := src.contiguous_view_offset()) is None: return None
+  offset = src.contiguous_view_offset()
+  if offset is None: return None
 
   # merge BUFFER_VIEWs
   if buf.op is Ops.BUFFER_VIEW: offset, buf = offset + buf.arg[1], buf.src[0]
@@ -92,47 +93,32 @@ def contiguous_mops_to_view(c:UOp, src: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
-  assert c.src[0].op is Ops.TUPLE, f"expected TUPLE body for precompiled call, got {c.src[0].op}"
+  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:])
-
-  # add the outputs to the call
-  srcs = c.src[0].src
-  resolved = [c.gettuple(i) for i in range(len(srcs))]
-  outs = tuple(_buffer_like(r) for r in resolved)
-  targets = [o.param_like(len(c.src)-1+i).shrink_to(s.shape) for i,(o,s) in enumerate(zip(outs, srcs))]
-  fxn = UOp.sink(*[t.after(t.store(s)) for t,s in zip(targets, srcs)])
-
-  # create the new thing for the big graph
-  new_call = c.replace(src=(fxn, *input_buffers, *outs), tag=None)
-  rets = tuple(o.after(new_call) for o in outs)
-
+  fxn = out.param_like(len(c.src)-1).assign(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
-  # NOTE: must use resolved shapes from the CALL (which substitutes PARAMs with external args), not raw body shapes
-  rets = tuple(r.shrink_to(rs.shape) for r,rs in zip(rets, resolved))
-
-  return UOp.maketuple(*rets)
+  if any(isinstance(s, UOp) for s in c.shape): ret = ret.shrink(tuple((0, s) for s in c.shape))
+  return ret
 
 # NOTE: adding rules to here is bad. these all need to run before the schedule cache
 pm_early_transform_tensor_graph = PatternMatcher([
   # transform precompiled CALLs
   (UPat(Ops.CALL, name="c"), transform_precompiled_call),
 
-  # resolve TUPLE+GETTUPLE (for precompiled calls)
-  (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
-
   # CONTIGUOUS(MOPS(BUFFER/BUFFER_VIEW)) → CONTIGUOUS(BUFFER_VIEW) when movement ops collapse to contiguous range
-  (UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Movement, name="src"),), name="c"), contiguous_mops_to_view),
+  (UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Movement),), name="c"), contiguous_mops_to_view),
 
   # add CONTIGUOUS to tagged UOps
-  (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.AFTER, Ops.STORE}, name="x"),
+  (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.ASSIGN, Ops.AFTER}, name="x"),
    lambda x: x.rtag(None).contiguous(tag=x.tag) if x.tag else x.replace(tag=None)),
-  # remove extra CONTIGUOUS on AFTER (only when target is contiguous)
+  # remove extra CONTIGUOUS on AFTER (only when after target is contiguous)
   (UPat(Ops.CONTIGUOUS, src=(UPat(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 AFTER+STORE with CONTIGUOUS when target is not a buffer
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat(name="src")))), name="u"), replace_store_after_with_contig),
-  # replace CONTIGUOUS with STORE+AFTER
-  (UPat(Ops.CONTIGUOUS, name="u"), replace_contig_with_store_after),
+  # 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),
   # remove DETACH/CONTIGUOUS_BACKWARD (allows more contiguous removal)
   (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
 ])
@@ -145,7 +131,7 @@ def untag_and_append(ctx:AllocCtx, x:UOp):
     replace_uop = ret
     while replace_uop.op is Ops.AFTER: replace_uop = replace_uop.src[0]
     ctx.buffer_map[original_uop] = replace_uop.shrink_to(original_uop.shape)
-  if ret.op is not Ops.AFTER: ctx.assigns.append(ret)  # AFTER gets appended by append_after
+  ctx.assigns.append(ret)
   return ret
 
 def append_after(ctx:AllocCtx, x:UOp):

tinygrad/engine/schedule.py

@@ -22,7 +22,6 @@ def create_schedule(sched_sink:UOp) -> UOp:
     for u in sched_sink.toposort(gate_kernel_sink):
       if u.op is not Ops.AFTER: continue
       k = u.src[1]
-      if k.op is Ops.STORE: continue  # skip unprocessed STORE+AFTER inside precompiled CALL bodies
       assert k.op in {Ops.CALL, Ops.END}, f"AFTER src[1] should be CALL or END, not {k.op}"
       in_degree.setdefault(k, 0)
       if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"

tinygrad/function.py

@@ -3,7 +3,6 @@ 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.tensor import Tensor
-from tinygrad.nn.state import get_state_dict
 
 def add_to_ctx(ctx, x:UOp):
   ret = x.param_like(len(ctx))
@@ -12,7 +11,7 @@ def add_to_ctx(ctx, x:UOp):
 
 pm_ctx = PatternMatcher([
   (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
-  (UPat((Ops.AFTER, Ops.CONTIGUOUS), name="x"),
+  (UPat((Ops.ASSIGN, Ops.CONTIGUOUS), name="x"),
    lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) else None),
   # strip UNIQUE from unique consts — they don't need buffer identity inside function bodies
   (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="x"), lambda ctx,x: x.replace(src=(x.src[1],))),
@@ -20,50 +19,39 @@ pm_ctx = PatternMatcher([
 
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):
-  def __init__(self, fxn:Callable[..., ReturnType], *, precompile:bool=False, precompile_backward:bool=False,
-               allow_implicit:bool=True, grad_fxn:Callable|None=None):
+  def __init__(self, fxn:Callable[..., ReturnType], *, precompile:bool=False):
     self.fxn = fxn
     self.precompile = precompile
-    self.precompile_backward = precompile_backward
-    self.allow_implicit = allow_implicit
-    self.grad_fxn = grad_fxn
 
   def __get__(self, obj, objtype=None): return functools.partial(self.__call__, obj) if obj is not None else self
 
   def __call__(self, *args, **kwargs) -> ReturnType:
-    params = get_state_dict((args, kwargs), tensor_type=(Tensor, UOp)).values()
+    input_uops: list[UOp] = [(t.uop if isinstance(t, Tensor) else t)
+                             for name,t in list(enumerate(args))+sorted(kwargs.items()) if isinstance(t, (Tensor, UOp))]
+
+    # use the base
+    #input_uops = [x.multibase for x in input_uops]
 
     # deduplicate input_uops, keeping the first occurrence index for each unique uop
-    call_uops: list[UOp] = dedup([(t.uop if isinstance(t, Tensor) else t) for t in params])
+    call_uops: list[UOp] = dedup(input_uops)
 
     # disable realize/schedule while this is running
     # run it and do surgery later
     with Context(ALLOW_DEVICE_USAGE=getenv("DEVICE_IN_FUNCTION_BUG", 0)):
       ret = self.fxn(*args, **kwargs)
-    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")
+    assert isinstance(ret, Tensor), "only supports one tensor return for now"
 
     # replace the known inputs with params (using deduplicated slots)
     subs = {}
     for i,x in enumerate(call_uops): subs[x] = x.param_like(i)
-    uret = uret.substitute(subs)
+    uret = ret.uop.substitute(subs)
 
     # add contiguous to call_uops
     #call_uops = [x.contiguous() for x in call_uops]
 
     # the BUFFERs that are left are the implicit inputs
-    num_explicit = len(call_uops)
     uret = graph_rewrite(uret, pm_ctx, call_uops, bottom_up=True, name="get_implicit_inputs")
     name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
-    if not self.allow_implicit:
-      implicit_buffers = [x for x in call_uops[num_explicit:] if x.op is Ops.BUFFER]
-      if implicit_buffers:
-        buf_strs = '\n  '.join(f"{i}: dtype={b.dtype}, size={b.size}, device={b.device}" for i,b in enumerate(implicit_buffers))
-        raise RuntimeError(f"function {name} has {len(implicit_buffers)} implicit buffer(s), but allow_implicit=False\n  {buf_strs}")
 
     # assign output
     #pbuffer = uret.param_like(len(call_uops))
@@ -72,23 +60,14 @@ class _function(Generic[ReturnType]):
     #call = assigned.call(*call_uops, buffer, name=name)
     #ret = buffer.after(call)
 
-    fret = uret.call(*call_uops, grad_fxn=self.grad_fxn, name=name, precompile=self.precompile,
-                     precompile_backward=self.precompile_backward)
-    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.gettuple(0), device=fret.device))
+    ret = uret.call(*call_uops, name=name, precompile=self.precompile)
+    return cast(ReturnType, Tensor(ret, device=ret.device))
 
 # overload signatures support both @function and @function(precompile=True) syntax
 @overload
-def function(fxn:Callable[..., ReturnType], *, precompile:bool=False, precompile_backward:bool=False,
-             allow_implicit:bool=True, grad_fxn:Callable|None=None) -> _function[ReturnType]: ...
+def function(fxn:Callable[..., ReturnType], *, precompile:bool=False) -> _function[ReturnType]: ...
 @overload
-def function(fxn:None=None, *, precompile:bool=False, precompile_backward:bool=False,
-             allow_implicit:bool=True, grad_fxn:Callable|None=None) -> Callable[[Callable[..., ReturnType]], _function[ReturnType]]: ...
-def function(fxn=None, *, precompile:bool=False, precompile_backward:bool=False, allow_implicit:bool=True, grad_fxn:Callable|None=None):
-  if fxn is None:
-    return lambda f: _function(f, precompile=precompile, precompile_backward=precompile_backward,
-                               allow_implicit=allow_implicit, grad_fxn=grad_fxn)
-  return _function(fxn, precompile=precompile, precompile_backward=precompile_backward,
-                   allow_implicit=allow_implicit, grad_fxn=grad_fxn)
+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)

tinygrad/gradient.py

@@ -13,33 +13,21 @@ 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 _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp, ...]]:
-  """Remove unused PARAMs from body and return compacted (body, args)."""
-  used = sorted({p.arg: p for p in body.toposort() if p.op is Ops.PARAM}.items())
-  return body.substitute({p: p.replace(arg=j) for j,(_, p) in enumerate(used)}, walk=True), tuple(all_args[i] for i,_ in used)
-
-def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
+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)
+  # auto-differentiate the function
   fxn, args = k.src[0], k.src[1:]
-  if k.arg.grad_fxn is not None:
-    if ctx.op is Ops.TUPLE:
-      real = [g for g in ctx.src if g.op is not Ops.NOOP]
-      return (None,) + (k.arg.grad_fxn(*real, call=k) if len(real) > 1 else k.arg.grad_fxn(real[0], k))
-    return (None,) + k.arg.grad_fxn(ctx, k)
-  assert fxn.op is Ops.TUPLE, f"expected TUPLE body for gradient, got {fxn.op}"
   params = {x.arg:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
-  grad_args = ctx.src
-  root_grad = UOp(Ops.TUPLE, src=tuple(fxn.src[i].const_like(0) if g.op is Ops.NOOP else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
-  grads = compute_gradient(fxn, root_grad, set(params.values()))
-  # for precompiled calls, substitute forward outputs with params so intermediates aren't recomputed
-  fwd_subs = {src: src.param_like(len(args)+len(grad_args)+i) for i, src in enumerate(fxn.src)} if k.arg.precompile else {}
-  fwd_outs = tuple(k.gettuple(i) for i in range(len(fxn.src))) if k.arg.precompile else ()
-  # collect needed gradient bodies, compact unused params, create a single backward CALL
-  grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
-  bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
-  bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))
-  bwd_call = bwd_body.call(*compact_args, name=(k.arg.name or "")+"_backward", precompile=k.arg.precompile_backward)
-  gb_map = {i: idx for idx, (i, _) in enumerate(grad_bodies)}
-  return (None,) + tuple(bwd_call.gettuple(gb_map[i]) if i in gb_map else None for i in range(len(args)))
+  grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params.values()))
+  ret: list[UOp|None] = [None]
+  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}"))
+    else:
+      ret.append(None)
+  return tuple(ret)
 
 # ctx is grad_output
 pm_gradient = PatternMatcher([
@@ -68,40 +56,26 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip([i for i,x in enumerate(ret.marg) if x]),)),
   (UPat(Ops.COPY, name="ret"), lambda ctx, ret: (ctx.copy_to_device(ret.src[0].device), None)),
   (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
-  (UPat(Ops.TUPLE), lambda ctx: ctx.src),
   # NOTE: this is only correct when the KERNEL has a single output
   (UPat(Ops.AFTER), lambda ctx: (ctx, ctx)),
+  # gradient on CALL: use provided grad_fxn or auto-differentiate
+  (UPat(Ops.CALL, name="k"), call_gradient),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),
 ])
 
-def _deepwalk(root:UOp, targets:set[UOp]) -> tuple[list[UOp], dict[UOp, bool]]:
+def _deepwalk(root:UOp, targets:set[UOp]) -> list[UOp]:
   # compute the target path (top down)
   in_target_path: dict[UOp, bool] = {}
   for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
-  # don't flow through DETACH or anything not in target path
-  return list(root.toposort(lambda node: node.op is not Ops.DETACH and in_target_path[node])), in_target_path
+  # don't flow through DETACH/ASSIGN or anything not in target path
+  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]:
-  walk, in_target_path = _deepwalk(root, targets)
-  grads: dict[UOp, UOp] = {root: root_grad}
-  for t0 in reversed(walk):
+  grads = {root: root_grad}
+  for t0 in reversed(_deepwalk(root, targets)):
     if t0 not in grads: continue
-    # GETTUPLE: accumulate gradient into a TUPLE UOp on the CALL, process when we hit the CALL
-    if t0.op is Ops.GETTUPLE:
-      k = t0.src[0]  # the CALL
-      assert k.op is Ops.CALL and k.src[0].op is Ops.TUPLE
-      n_outputs = len(k.src[0].src)
-      prev = grads[k].src if k in grads else tuple(UOp(Ops.NOOP) for _ in range(n_outputs))
-      grads[k] = UOp.maketuple(*(prev[i] + grads[t0] if i == t0.arg and prev[i].op is not Ops.NOOP else
-                                 grads[t0] if i == t0.arg else prev[i] for i in range(n_outputs)))
-      continue
-    # CALL: pass needed param set so backward only computes required gradients
-    if t0.op is Ops.CALL:
-      needed = {i for i, arg in enumerate(t0.src[1:]) if arg in targets or in_target_path.get(arg, False)}
-      lgrads:tuple[UOp|None, ...]|None = call_gradient(grads[t0], t0, needed)
-    else:
-      lgrads = cast(tuple[UOp|None, ...]|None, pm_gradient.rewrite(t0, ctx=grads[t0]))
+    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]}...")
     assert len(lgrads) == len(t0.src), f"got {len(lgrads)} gradient, expected {len(t0.src)}"
     for k,v in zip(t0.src, lgrads):
@@ -113,4 +87,8 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
         # we add the backward metadata to everything new in the graph
         for bw_uop in v.toposort(lambda x: x not in (t0, *t0.src, grads[t0])):
           all_metadata[bw_uop] = all_metadata.get(bw_uop, ())+backward_metadata
+  # end any ranges on grads with a reduce sum
+  for k,v in grads.items():
+    if len(v.ranges):
+      grads[k] = v.reduce(*v.ranges, arg=Ops.ADD)
   return grads

tinygrad/helpers.py

@@ -424,9 +424,7 @@ 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()
-  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
+  ret = subprocess.check_output(cmd.split(), **kwargs).decode().strip()
   if DEBUG >= 1: print(f"system: '{cmd}' returned {len(ret)} bytes in {(time.perf_counter() - st)*1e3:.2f} ms")
   return ret
 

tinygrad/nn/__init__.py

@@ -303,22 +303,16 @@ 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, AxisType
+from tinygrad.uop.ops import UOp, KernelInfo, Ops
 def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
   weight, idx = call.src[1:]
-  is_vocab_sharded = isinstance(weight.device, tuple) and weight.axis == 0
-  # for multi-device: replicate grad_emb and idx on all devices
+  # for multi-device: unshard inputs to one device
   if isinstance(weight.device, tuple):
-    assert weight.axis is None or weight.axis == 0, "only vocab (axis=0) sharding supported on Embedding with USE_ATOMICS"
+    assert weight.axis is None, "sharded weights on Embedding not supported with USE_ATOMICS"
     grad_emb = grad_emb.copy_to_device(weight.device)
     idx = idx.copy_to_device(weight.device)
-  if is_vocab_sharded:
-    ndev = len(weight.device)
-    local_vocab_size = weight.shape[0] // ndev
-    grad_weight_uop = Tensor.empty(local_vocab_size, weight.shape[1], dtype=dtypes.float, device=weight.device).uop.multi(axis=0)
-  else:
-    # weight is replicated (or single device), grad_weight should match
-    grad_weight_uop = Tensor.empty(weight.shape, dtype=dtypes.float, device=weight.device).uop
+  # weight is replicated, grad_weight should match
+  grad_weight_uop = Tensor.empty(weight.shape, dtype=dtypes.float, device=weight.device).uop
 
   # TODO: how do we remove this dumb kernel and use Tensor.zeros?
   def _zero_kernel(out:UOp) -> UOp:
@@ -332,35 +326,15 @@ 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]))
-
-    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
-
-    if is_vocab_sharded:
-      # each device owns [offset, offset+local_vocab_size) of the global vocabulary
-      dnum = UOp.variable("_device_num", 0, ndev-1)
-      offset = dnum * local_vocab_size
-      global_token_id = idx_flat[i].cast(dtypes.index)
-      local_token_id = (global_token_id - offset).clip(0, grad_weight.shape[0]-1)
-      in_range = (global_token_id >= offset) & (global_token_id < (offset + local_vocab_size))
-      grad_val = in_range.where(grad_emb_flat[i, j].cast(dtypes.float), 0.0)
-    else:
-      local_token_id = idx_flat[i].clip(0, grad_weight.shape[0]-1).cast(dtypes.index)
-      grad_val = grad_emb_flat[i, j].cast(dtypes.float)
+    i = UOp.range(grad_emb_flat.shape[0], 0)  # batch_size * sequence_length
+    j = UOp.range(grad_emb_flat.shape[1], 1)  # embed_size
+    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(local_token_id, j, ptr=True), grad_val), arg = atomic_arg)
-    return atomic.end(i, j_outer, j_inner).sink(arg=KernelInfo(name="embedding_bwd", opts_to_apply=()))
-
+    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=()))
   grad_weight_uop = grad_weight_uop.custom_kernel(grad_emb, idx, fxn=_embedding_bwd_kernel)[0]
 
   return (grad_weight_uop.cast(weight.dtype), None)

tinygrad/renderer/amd/sqtt.py

@@ -126,9 +126,6 @@ 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
@@ -139,24 +136,6 @@ class InstOpRDNA4(Enum):
   OTHER_VMEM = 0xbd
   OTHER_VMEM_5 = 0xc1
 
-class InstOpCDNA(Enum):
-  SMEM_RD = 0
-  SALU_32 = 1
-  VMEM_RD = 2
-  VMEM_WR = 3
-  FLAT_WR = 4
-  VALU_32 = 5
-  LDS = 6
-  PC = 7
-  JUMP = 12
-  NEXT = 13
-  FLAT_RD = 14
-  OTHER_MSG = 15
-  SMEM_WR = 16
-  SALU_64 = 17
-  VALU_64 = 18
-  VALU_MAI = 28
-
 # ═══════════════════════════════════════════════════════════════════════════════
 # PACKET TYPE BASE CLASS
 # ═══════════════════════════════════════════════════════════════════════════════
@@ -402,27 +381,25 @@ PACKET_TYPES_RDNA4: dict[int, type[PacketType]] = {
 # CDNA PACKET TYPE DEFINITIONS
 # ═══════════════════════════════════════════════════════════════════════════════
 
-class CDNA_MISC(PacketType):
-  """pkt_fmt=0: 16-bit (Misc)"""
+class CDNA_DELTA(PacketType):
+  """pkt_fmt=0: 16-bit timestamp delta packet"""
   encoding = bits[3:0] == 0
-  delta = bits[11:4]
-  sh = bits[12:12]
-  misc_type = bits[15:13]
+  delta = bits[11:4]      # (data >> 4) & 0xff
+  unk_0 = bits[12:12]     # (data >> 0xc) & 1
+  unk_1 = bits[15:13]     # (data >> 0xd)
 
 class CDNA_TIMESTAMP(PacketType):
   """pkt_fmt=1: 64-bit timestamp packet (case 0x0)"""
   encoding = bits[3:0] == 1
-  _reserved = bits[15:4]
+  unk_0 = bits[15:4]
   timestamp = bits[63:16]   # stored as (data_word >> 0x10) in low 46 bits of local_58
 
-class CDNA_REG(PacketType):
-  """pkt_fmt=2: 64-bit (Reg)"""
+class CDNA_PKT_2(PacketType):
+  """pkt_fmt=2: 64-bit packet (case 0x4)"""
   encoding = bits[3:0] == 2
-  pipe = bits[6:5]
-  _me_raw = bits[8:7]
-  _reserved = bits[15:9]
-  regaddr = bits[31:16]
-  regdata = bits[63:32]
+  unk_0 = bits[6:5]       # (data >> 5) & 3
+  unk_1 = bits[7:7]       # (data >> 7) + 1 & 1
+  unk_padding = bits[63:8]
 
 class CDNA_WAVESTART(PacketType):
   """type 3: 32-bit wave start (Wave/group_id)"""
@@ -433,19 +410,19 @@ class CDNA_WAVESTART(PacketType):
   simd = bits[15:14]
   pipe = bits[17:16]
   me = bits[19:18]
-  _reserved = bits[21:20]
+  _gap = bits[21:20]
   count = bits[28:22]
   _padding = bits[31:29]
 
-class CDNA_WAVEALLOC(PacketType):
-  """pkt_fmt=4: 16-bit (Wave)"""
+class CDNA_PKT_4(PacketType):
+  """pkt_fmt=4: 16-bit packet (case 0xc, same as 0x8/0x14)"""
   encoding = bits[3:0] == 4
-  sh = bits[5:5]
-  cu = bits[9:6]
-  wave = bits[13:10]
-  simd = bits[15:14]
+  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)
 
-class CDNA_REG_CS(PacketType):
+class REGCS_CDNA(PacketType):
   """type 5: 48-bit register CS write (RegCs)"""
   encoding = bits[3:0] == 5
   pipe = bits[6:5]
@@ -461,86 +438,80 @@ class CDNA_WAVEEND(PacketType):
   wave = bits[13:10]
   simd = bits[15:14]
 
-class CDNA_INST(PacketType):
-  """pkt_fmt=10: 16-bit (MsgInst)"""
+class CDNA_EXEC(PacketType):
+  """pkt_fmt=10: 16-bit EXEC packet (case 0x24)"""
   encoding = bits[3:0] == 10
-  wave = bits[8:5]
-  simd = bits[10:9]
-  op = bits[15:11].enum(InstOpCDNA)
+  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_INST_PC(PacketType):
-  """pkt_fmt=11: 64-bit (MsgInstPc)"""
+class CDNA_PKT_11(PacketType):
+  """pkt_fmt=11: 64-bit packet (case 0x28)"""
   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]
+  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_ISSUE(PacketType):
-  """pkt_fmt=13: 32-bit (Issue)"""
+class CDNA_INST(PacketType):
+  """pkt_fmt=13: 32-bit INST packet (case 0x30)"""
   encoding = bits[3:0] == 13
-  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]
+  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]
 
-class CDNA_PERF(PacketType):
-  """pkt_fmt=14: 64-bit (MsgPerf)"""
+class CDNA_PKT_14(PacketType):
+  """pkt_fmt=14: 64-bit packet (case 0x34)"""
   encoding = bits[3:0] == 14
-  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]
+  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]
 
-class CDNA_EVENT(PacketType):
-  """pkt_fmt=7: 16-bit"""
+class CDNA_PKT_7(PacketType):
+  """pkt_fmt=7: 16-bit packet"""
   encoding = bits[3:0] == 7
-  _reserved = bits[15:4]
+  unk_padding = bits[15:4]
 
-class CDNA_EVENT_CS(PacketType):
-  """pkt_fmt=8: 16-bit"""
+class CDNA_PKT_8(PacketType):
+  """pkt_fmt=8: 16-bit packet"""
   encoding = bits[3:0] == 8
-  _reserved = bits[15:4]
+  unk_padding = bits[15:4]
 
-class CDNA_EVENT_GFX1(PacketType):
-  """pkt_fmt=9: 16-bit"""
+class CDNA_PKT_9(PacketType):
+  """pkt_fmt=9: 16-bit packet"""
   encoding = bits[3:0] == 9
-  _reserved = bits[15:4]
+  unk_padding = bits[15:4]
 
-class CDNA_USERDATA(PacketType):
-  """pkt_fmt=12: 48-bit (UserData)"""
+class CDNA_PKT_12(PacketType):
+  """pkt_fmt=12: 48-bit packet"""
   encoding = bits[3:0] == 12
-  sh = bits[5:5]
-  cu = bits[9:6]
-  wave = bits[13:10]
-  simd = bits[15:14]
-  data = bits[47:16]
+  unk_padding = bits[47:4]
 
-class CDNA_REG_CS_PRIV(PacketType):
-  """pkt_fmt=15: 48-bit (RegCs)"""
+class CDNA_PKT_15(PacketType):
+  """pkt_fmt=15: 48-bit packet (case 0x38, same as 0x10)"""
   encoding = bits[3:0] == 15
-  pipe = bits[6:5]
-  _me_raw = bits[8:7]
-  regaddr = bits[15:9]
-  regdata = bits[47:16]
+  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]
 
 PACKET_TYPES_CDNA: dict[int, type[PacketType]] = {
-  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,
+  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,
 }
 
 # ═══════════════════════════════════════════════════════════════════════════════
@@ -552,8 +523,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_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}
+  # 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}
   decode_info = {}
   for opcode, pkt_cls in packet_types.items():
     delta_field = getattr(pkt_cls, 'delta', None)
@@ -630,7 +601,7 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
   def simd_select(p) -> bool: return getattr(p, "cu", 0) == 0 and getattr(p, "simd", 0) == 0
   for p in decode(data):
     if not simd_select(p): continue
-    if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART)):
+    if isinstance(p, (WAVESTART, WAVESTART_RDNA4)):
       assert p.wave not in wave_pc, "only one inflight wave per unit"
       wave_pc[p.wave] = next(iter(pc_map))
     elif isinstance(p, WAVEEND):

tinygrad/runtime/ops_amd.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import cast
+from typing import cast, ClassVar
 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,12 +11,12 @@ 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, sqtt, amdgpu_kd, amdgpu_drm
+from tinygrad.runtime.autogen import kfd, hsa, pci, 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
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_ip_offsets, import_pmc
-from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
+from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, PCIDevice, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
 from tinygrad.runtime.support.memory import AddrSpace
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
@@ -824,18 +824,19 @@ 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):
-  def __init__(self, dev, dev_id):
-    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7590,0x75a0)),), vram_bar=0,
-      va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev)
-    self._compute_props()
+  gpus:ClassVar[list[str]] = []
 
-  def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]:
-    return ([(self.dev_impl.paddr2xgmi(p), sz) for p, sz in paddrs], AddrSpace.PEER) if self.dev_impl.is_hive() else super().p2p_paddrs(paddrs)
+  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 _compute_props(self):
+  def _setup_adev(self, pci_dev:PCIDevice, dma_regions:list[tuple[int, MMIOInterface]]|None=None):
+    self.dev_impl:AMDev = AMDev(pci_dev, dma_regions)
     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}")
@@ -892,10 +893,9 @@ 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}")
-    self.dev_impl = AMDev(self.pci_dev)
-    self._compute_props()
-    self.pci_dev.usb._pci_cacheable += [self.pci_dev.bar_info(2)] # doorbell region is cacheable
+    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.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
     self.copy_bufs = [self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)]
@@ -905,7 +905,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, force_devmem=False, **kwargs) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=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)

tinygrad/runtime/ops_nv.py

@@ -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
+from typing import cast, ClassVar
 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, mesa
+from tinygrad.runtime.autogen import nv_570, nv_580, pci, 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,13 +535,17 @@ 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)")
+    super().__init__(dev, dev_id, vendor=0x10de, devices=[(0xff00, [0x2200, 0x2400, 0x2500, 0x2600, 0x2700, 0x2800, 0x2b00, 0x2c00, 0x2d00, 0x2f00])],
+      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)
-    super().__init__(dev, dev_id, vendor=0x10de, devices=((0xff00, (0x2200,0x2400,0x2500,0x2600,0x2700,0x2800,0x2b00,0x2c00,0x2d00,0x2f00)),),
-      base_class=0x03, vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size, dev_impl_t=NVDev)
 
+    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())
 
@@ -549,11 +553,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, force_devmem=False, **kwargs) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=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.
-    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)
+    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)
 
   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

tinygrad/runtime/support/am/amdev.py

@@ -1,11 +1,11 @@
 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
+from tinygrad.runtime.support.system import PCIDevice, PCIDevImplBase
 from tinygrad.runtime.support.am.ip import AM_IP, AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
 
 AM_DEBUG = getenv("AM_DEBUG", 0)
@@ -143,11 +143,11 @@ class AMMemoryManager(MemoryManager):
     self.dev.gmc.flush_tlb(ip='GC', vmid=0)
     self.dev.gmc.flush_tlb(ip='MM', vmid=0)
 
-class AMDev:
+class AMDev(PCIDevImplBase):
   Version = 0xA0000008
 
-  def __init__(self, pci_dev:PCIDevice, reset_mode=False):
-    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
+  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
     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,7 +185,6 @@ class AMDev:
 
     # 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()

tinygrad/runtime/support/am/ip.py

@@ -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"
 
-    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)
+    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)
     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')

tinygrad/runtime/support/c.py

@@ -8,7 +8,6 @@ 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

tinygrad/runtime/support/nv/ip.py

@@ -575,7 +575,7 @@ class NV_GSP(NV_IP):
     def bdf_as_int(s): return 0x000 if s.startswith("usb") or s.startswith("remote") else (int(s[5:7],16)<<8) | (int(s[8:10],16)<<3) | int(s[-1],16)
 
     pcidev = self.nvdev.pci_dev
-    data = nv.GspSystemInfo(gpuPhysAddr=pcidev.bar_info(0)[0], gpuPhysFbAddr=pcidev.bar_info(1)[0], gpuPhysInstAddr=pcidev.bar_info(3)[0],
+    data = nv.GspSystemInfo(gpuPhysAddr=pcidev.bar_info[0].addr, gpuPhysFbAddr=pcidev.bar_info[1].addr, gpuPhysInstAddr=pcidev.bar_info[3].addr,
       pciConfigMirrorBase=[0x88000, 0x92000][self.nvdev.fmc_boot], pciConfigMirrorSize=0x1000, nvDomainBusDeviceFunc=bdf_as_int(self.nvdev.devfmt),
       bIsPassthru=1, PCIDeviceID=pcidev.read_config(pci.PCI_VENDOR_ID, 4), PCISubDeviceID=pcidev.read_config(pci.PCI_SUBSYSTEM_VENDOR_ID, 4),
       PCIRevisionID=pcidev.read_config(pci.PCI_REVISION_ID, 1), maxUserVa=0x7ffffffff000)

tinygrad/runtime/support/nv/nvdev.py

@@ -1,10 +1,9 @@
 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, MMIOInterface
+from tinygrad.runtime.support.system import PCIDevice, PCIDevImplBase, MMIOInterface
 
 NV_DEBUG = getenv("NV_DEBUG", 0)
 
@@ -70,10 +69,9 @@ class NVMemoryManager(MemoryManager):
 
   def on_range_mapped(self): self.dev.NV_VIRTUAL_FUNCTION_PRIV_MMU_INVALIDATE.write((1 << 0) | (1 << 1) | (1 << 6) | (1 << 31))
 
-class NVDev:
+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()

tinygrad/runtime/support/system.py

@@ -1,14 +1,18 @@
 from __future__ import annotations
 import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys, itertools, struct, socket, subprocess, time, enum
-from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv, unwrap, fetch, system, _ensure_downloads_dir
+from typing import ClassVar
+from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv, unwrap, fetch, system
 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 VirtMapping, AddrSpace, BumpAllocator
+from tinygrad.runtime.support.memory import MemoryManager, VirtMapping, AddrSpace
 from tinygrad.runtime.support.usb import ASM24Controller, USBMMIOInterface
 
 MAP_FIXED, MAP_FIXED_NOREPLACE = 0x10, 0x100000
 MAP_LOCKED, MAP_POPULATE, MAP_NORESERVE = 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000), 0x400
 
+@dataclasses.dataclass(frozen=True)
+class PCIBarInfo: addr:int; size:int # noqa: E702
+
 class _System:
   def write_sysfs(self, path:str, value:str, msg:str, expected:str|None=None):
     if FileIOInterface(path, os.O_RDONLY).read().splitlines()[0] != (expected or value):
@@ -40,11 +44,6 @@ 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):
@@ -54,7 +53,7 @@ class _System:
     self.pagemap.seek(vaddr // mmap.PAGESIZE * 8)
     return [(x & ((1<<55) - 1)) * mmap.PAGESIZE for x in array.array('Q', self.pagemap.read(size//mmap.PAGESIZE*8, binary=True))]
 
-  def pci_scan_bus(self, vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None=None) -> list[str]:
+  def pci_scan_bus(self, vendor:int, devices:list[tuple[int, list[int]]], base_class:int|None=None) -> list[str]:
     all_devs = []
     if OSX:
       def read_prop(svc, key) -> int:
@@ -73,15 +72,7 @@ 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)])
 
-  @functools.cache
-  def list_devices(self, vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None=None) -> list[tuple[type, str]]:
-    return [(APLRemotePCIDevice if OSX else PCIDevice, x) for x in System.pci_scan_bus(vendor, devices, base_class)]
-
-  def pci_probe_device(self, devpref:str, dev_id:int, vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None=None):
-    cl, pcibus = hcq_filter_visible_devices(self.list_devices(vendor, devices, base_class))[dev_id]
-    return cl(devpref, pcibus)
-
-  def pci_setup_usb_bars(self, usb:ASM24Controller, gpu_bus:int, mem_base:int, pref_mem_base:int) -> dict[int, tuple[int, int]]:
+  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.
       buses = (0 << 0) | ((bus+1) << 8) | ((gpu_bus) << 16)
@@ -123,7 +114,7 @@ class _System:
         usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off, bus=gpu_bus, dev=0, fn=0, value=mem_space_addr[bar_mem] & 0xffffffff, size=4)
         if bar_64: usb.pcie_cfg_req(pci.PCI_BASE_ADDRESS_0 + bar_off + 4, bus=gpu_bus, dev=0, fn=0, value=mem_space_addr[bar_mem] >> 32, size=4)
 
-        bars[bar_off // 4] = (mem_space_addr[bar_mem], bar_size)
+        bars[bar_off // 4] = PCIBarInfo(mem_space_addr[bar_mem], bar_size)
         mem_space_addr[bar_mem] += round_up(bar_size, 2 << 20)
 
       bar_off += 8 if bar_64 else 4
@@ -150,7 +141,7 @@ System = _System()
 # *** PCI Devices
 
 class PCIDevice:
-  def __init__(self, devpref:str, pcibus:str):
+  def __init__(self, devpref:str, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
     self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
     self.pcibus, self.irq_poller = pcibus, None
 
@@ -160,6 +151,11 @@ class PCIDevice:
     if FileIOInterface.exists(f"/sys/bus/pci/devices/{self.pcibus}/driver"):
       FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver/unbind", os.O_WRONLY).write(self.pcibus)
 
+    for i in resize_bars or []:
+      if FileIOInterface.exists(rpath:=f"/sys/bus/pci/devices/{self.pcibus}/resource{i}_resize"):
+        try: FileIOInterface(rpath, os.O_RDWR).write(str(int(FileIOInterface(rpath, os.O_RDONLY).read(), 16).bit_length() - 1))
+        except OSError as e: raise RuntimeError(f"Cannot resize BAR {i}: {e}. Ensure the resizable BAR option is enabled.") from e
+
     if getenv("VFIO", 0) and (vfio_fd:=System.vfio) is not None:
       FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver_override", os.O_WRONLY).write("vfio-pci")
       FileIOInterface("/sys/bus/pci/drivers_probe", os.O_WRONLY).write(self.pcibus)
@@ -176,11 +172,15 @@ 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) + 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)))
+        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)
     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)
+    self.bar_fds = {b: FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource{b}", os.O_RDWR | os.O_SYNC | os.O_CLOEXEC) for b in bars}
+
+    res = FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource", os.O_RDONLY).read().splitlines()
+    self.bar_info = {j:PCIBarInfo(int(s,16), int(e,16)-int(s,16)+1) for j,(s,e,_) in enumerate(l.split() for l in res)}
 
   def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
     assert not contiguous or size <= (2 << 20), "Contiguous allocation is only supported for sizes up to 2MB"
@@ -188,63 +188,45 @@ class PCIDevice:
     va = FileIOInterface.anon_mmap(vaddr, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS|MAP_POPULATE|MAP_LOCKED|flags, 0)
     sysmem_view, paddrs = MMIOInterface(va, size), [(x, mmap.PAGESIZE) for x in System.system_paddrs(va, size)]
     return sysmem_view, [p + i for p, sz in paddrs for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]
-
-  def reset(self): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{self.pcibus}/reset'")
   def read_config(self, offset:int, size:int): return int.from_bytes(self.cfg_fd.read(size, binary=True, offset=offset), byteorder='little')
   def write_config(self, offset:int, value:int, size:int): self.cfg_fd.write(value.to_bytes(size, byteorder='little'), binary=True, offset=offset)
-
-  @functools.cache
-  def bar_fd(self, bar_idx:int) -> FileIOInterface:
-    return FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource{bar_idx}", os.O_RDWR | os.O_SYNC | os.O_CLOEXEC)
-  @functools.cache
-  def bar_info(self, bar_idx:int) -> tuple[int, int]:
-    s, e, _ = FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource", os.O_RDONLY).read().splitlines()[bar_idx].split()
-    return (int(s, 16), int(e, 16) - int(s, 16) + 1)
   def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface:
-    fd, sz = self.bar_fd(bar), size or (self.bar_info(bar)[1] - off)
+    fd, sz = self.bar_fds[bar], size or (self.bar_info[bar].size - off)
     libc.madvise(loc:=fd.mmap(addr, sz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if addr else 0), off), sz, libc.MADV_DONTFORK)
     return MMIOInterface(loc, sz, fmt=fmt)
-  def resize_bar(self, bar_idx:int):
-    rpath = f"/sys/bus/pci/devices/{self.pcibus}/resource{bar_idx}_resize"
-    try: FileIOInterface(rpath, os.O_RDWR).write(str(int(FileIOInterface(rpath, os.O_RDONLY).read(), 16).bit_length() - 1))
-    except OSError as e: raise RuntimeError(f"Cannot resize BAR {bar_idx}: {e}. Ensure the resizable BAR option is enabled.") from e
+  def reset(self): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{self.pcibus}/reset'")
 
 class USBPCIDevice(PCIDevice):
-  def __init__(self, devpref:str, pcibus:str):
+  def __init__(self, devpref:str, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
     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 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]
-
-  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 bar_info(self, bar_idx:int) -> tuple[int, int]: return self._bar_info[bar_idx]  # type: ignore[override]
+    self.pcibus, self.bar_info = pcibus, System.pci_setup_usb_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30))
   def map_bar(self, bar, off=0, addr=0, size=None, fmt='B'):
-    return USBMMIOInterface(self.usb, self.bar_info(bar)[0] + off, size or self.bar_info(bar)[1], fmt)
-  def resize_bar(self, bar_idx:int): pass # already resized
+    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)
+
+class PCIDevImplBase:
+  mm: MemoryManager
 
 @dataclasses.dataclass
 class PCIAllocationMeta: mapping:VirtMapping; has_cpu_mapping:bool; hMemory:int=0 # noqa: E702
 
-class PCIIfaceBase:
-  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)[1] == (256 << 20)
+class LNXPCIIfaceBase:
+  dev_impl:PCIDevImplBase
+  gpus:ClassVar[list[str]] = []
 
-  def __init__(self, dev, dev_id, vendor, devices:tuple[tuple[int, tuple[int, ...]], ...], vram_bar, va_start, va_size,
-               dev_impl_t, base_class:int|None=None):
-    self.pci_dev = System.pci_probe_device(dev.__class__.__name__[:2], dev_id, vendor, devices, base_class=base_class)
-    if self.is_local(): System.reserve_va(va_start, va_size)
-    with contextlib.suppress(Exception): self.pci_dev.resize_bar(vram_bar)
-    self.dev_impl = dev_impl_t(self.pci_dev)
-    self.dev, self.vram_bar = dev, vram_bar
+  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.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:
-    should_use_sysmem = host or ((cpu_access if self.is_bar_small() else (uncached and cpu_access)) and not force_devmem)
+    # 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)
     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)
@@ -258,108 +240,96 @@ class PCIIfaceBase:
   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 self.is_local() and b.owner == self.dev and b.meta.has_cpu_mapping: FileIOInterface.munmap(b.va_addr, b.size)
-
-  def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]:
-    return [(p + self.pci_dev.bar_info(self.vram_bar)[0], sz) for p, sz in paddrs], AddrSpace.SYS
+    if b.owner == self.dev and b.meta.has_cpu_mapping and not OSX: 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, PCIIfaceBase):
+    elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, LNXPCIIfaceBase):
       snooped, uncached = True, b.meta.mapping.uncached
       if b.meta.mapping.aspace is AddrSpace.SYS: paddrs, aspace = b.meta.mapping.paddrs, AddrSpace.SYS
-      else: paddrs, aspace = ifa.p2p_paddrs(b.meta.mapping.paddrs)
+      elif hasattr(ifa.dev_impl, 'paddr2xgmi') and ifa.dev_impl.gmc.xgmi_seg_sz > 0:
+        paddrs, aspace = [(ifa.dev_impl.paddr2xgmi(p), sz) for p, sz in b.meta.mapping.paddrs], AddrSpace.PEER
+      else: paddrs, aspace = [(p + ifa.p2p_base_addr, sz) for p, sz in b.meta.mapping.paddrs], AddrSpace.SYS
     else: raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
 
     self.dev_impl.mm.map_range(int(b.va_addr), round_up(b.size, 0x1000), paddrs, aspace=aspace, snooped=snooped, uncached=uncached)
 
 # *** Remote PCI Devices
 
-class RemoteCmd(enum.IntEnum):
-  PROBE, MAP_BAR, MAP_SYSMEM_FD, CFG_READ, CFG_WRITE, RESET, MMIO_READ, MMIO_WRITE, MAP_SYSMEM, SYSMEM_READ, SYSMEM_WRITE = range(11)
+class RemoteCmd(enum.IntEnum): MAP_BAR, MAP_SYSMEM_FD, CFG_READ, CFG_WRITE, RESET, MMIO_READ, MMIO_WRITE = 1, 2, 3, 4, 5, 6, 7
 
 class RemoteMMIOInterface(MMIOInterface):
-  def __init__(self, dev:RemotePCIDevice, residx:int, nbytes:int, fmt='B', off=0, rd_cmd=RemoteCmd.MMIO_READ, wr_cmd=RemoteCmd.MMIO_WRITE):
+  def __init__(self, dev:RemotePCIDevice, residx:int, nbytes:int, fmt='B', off=0):
     self.dev, self.residx, self.nbytes, self.fmt, self.off, self.el_sz = dev, residx, nbytes, fmt, off, struct.calcsize(fmt)
-    self.rd_cmd, self.wr_cmd = rd_cmd, wr_cmd
 
   def __getitem__(self, index):
     sl = index if isinstance(index, slice) else slice(index, index + 1)
     start, stop = (sl.start or 0) * self.el_sz, (sl.stop or len(self)) * self.el_sz
-    data = self.dev._bulk_read(self.rd_cmd, self.residx, self.off + start, stop - start)
+    data = self.dev._bulk_read(RemoteCmd.MMIO_READ, self.residx, self.off + start, stop - start)
     result = data if self.fmt == 'B' else list(struct.unpack(f'<{(stop - start) // self.el_sz}{self.fmt}', data))
     return result if isinstance(index, slice) else result[0]
 
   def __setitem__(self, index, val):
     start = (index.start or 0) * self.el_sz if isinstance(index, slice) else index * self.el_sz
     data = (val if self.fmt == 'B' else struct.pack(f'<{len(val)}{self.fmt}', *val)) if isinstance(index, slice) else struct.pack(f'<{self.fmt}', val)
-    self.dev._bulk_write(self.wr_cmd, self.residx, self.off + start, data)
+    self.dev._bulk_write(RemoteCmd.MMIO_WRITE, self.residx, self.off + start, data)
 
   def view(self, offset:int=0, size:int|None=None, fmt=None):
-    return RemoteMMIOInterface(self.dev, self.residx, size or (self.nbytes - offset), fmt or self.fmt, self.off + offset, self.rd_cmd, self.wr_cmd)
+    return RemoteMMIOInterface(self.dev, self.residx, size or (self.nbytes - offset), fmt or self.fmt, self.off + offset)
 
 class RemotePCIDevice(PCIDevice):
-  @staticmethod
-  def _recvall(sock:socket.socket, n:int) -> bytes:
+  def __init__(self, devpref:str, pcibus:str, bars:list[int], sock:socket.socket):
+    self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
+    self.pcibus, self.sock = pcibus, sock
+    for buft in [socket.SO_SNDBUF, socket.SO_RCVBUF]: self.sock.setsockopt(socket.SOL_SOCKET, buft, 64 << 20)
+    self.bar_info = {b: PCIBarInfo(0, self._rpc(RemoteCmd.MAP_BAR, b)[0]) for b in bars}
+
+  def _recvall(self, n:int) -> bytes:
     data = b''
-    while len(data) < n and (chunk:=sock.recv(n - len(data))): data += chunk
+    while len(data) < n and (chunk:=self.sock.recv(n - len(data))): data += chunk
     if len(data) < n: raise RuntimeError("Connection closed")
     return data
 
-  @staticmethod
-  def _rpc(sock:socket.socket, dev_id:int, cmd:int, *args:int, bar:int=0, readout_size:int=0, payload:bytes=b'', has_fd=False):
-    sock.sendall(struct.pack('<BIIQQQ', cmd, dev_id, bar, *(*args, 0, 0, 0)[:3]) + payload)
-    if has_fd:
-      msg, anc, _, _ = sock.recvmsg(17, socket.CMSG_LEN(4))
-      fd = struct.unpack('<i', anc[0][2][:4])[0]
-    else: msg, fd = RemotePCIDevice._recvall(sock, 17), None
+  def _recv_with_fd(self) -> tuple[bytes, int|None]:
+    msg, anc, _, _ = self.sock.recvmsg(17, socket.CMSG_LEN(4))
+    return msg, struct.unpack('<i', anc[0][2][:4])[0]
+
+  def _rpc(self, cmd:int, *args:int, readout_size:int=0, has_fd=False) -> tuple[int, int, bytes|None, int|None]:
+    self.sock.sendall(struct.pack('<BBQQQ', cmd, *(*args, 0, 0, 0, 0)[:4]))
+    msg, fd = self._recv_with_fd() if has_fd else (self._recvall(17), None)
     if (resp:=struct.unpack('<BQQ', msg))[0] != 0:
-      raise RuntimeError(f"RPC failed: {RemotePCIDevice._recvall(sock, resp[1]).decode('utf-8') if resp[1] > 0 else 'unknown error'}")
-    return (resp[1], resp[2]) + ((RemotePCIDevice._recvall(sock, readout_size) if readout_size > 0 else None),) + (fd,)
+      raise RuntimeError(f"RPC failed: {self._recvall(resp[1]).decode('utf-8') if resp[1] > 0 else 'unknown error'}")
+    return (resp[1], resp[2]) + ((self._recvall(readout_size) if readout_size > 0 else None),) + (fd,)
 
-  def __init__(self, devpref:str, pcibus:str, sock:socket.socket|None=None):
-    self.sock, self.pcibus, self.dev_id = unwrap(sock), pcibus, int(pcibus.split(':')[-1]) if ':' in pcibus else 0
-    for buft in [socket.SO_SNDBUF, socket.SO_RCVBUF]: self.sock.setsockopt(socket.SOL_SOCKET, buft, 64 << 20)
+  def _bulk_read(self, cmd:int, idx:int, offset:int, size:int) -> bytes: return unwrap(self._rpc(cmd, idx, offset, size, readout_size=size)[2])
+  def _bulk_write(self, cmd:int, idx:int, offset:int, data:bytes): self.sock.sendall(struct.pack('<BBQQQ', cmd, idx, offset, len(data), 0) + data)
 
-    self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
+  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
+    mapped_size, _, _, fd = self._rpc(RemoteCmd.MAP_SYSMEM_FD, 0, 0, size, has_fd=True)
+    memview = MMIOInterface(FileIOInterface(fd=fd).mmap(0, mapped_size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, 0), mapped_size, fmt='B')
 
-  def _bulk_read(self, cmd:int, idx:int, offset:int, size:int) -> bytes:
-    return unwrap(self._rpc(self.sock, self.dev_id, cmd, offset, size, bar=idx, readout_size=size)[2])
-  def _bulk_write(self, cmd:int, idx:int, offset:int, data:bytes):
-    self.sock.sendall(struct.pack('<BIIQQQ', cmd, self.dev_id, idx, offset, len(data), 0) + data)
-
-  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]: raise NotImplementedError()
-  def reset(self): self._rpc(self.sock, self.dev_id, RemoteCmd.RESET)
-  def read_config(self, offset:int, size:int): return self._rpc(self.sock, self.dev_id, RemoteCmd.CFG_READ, offset, size)[0]
-  def write_config(self, offset:int, value:int, size:int): self._rpc(self.sock, self.dev_id, RemoteCmd.CFG_WRITE, offset, size, value)
-
-  @functools.cache
-  def bar_info(self, bar_idx:int) -> tuple[int, int]: return self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_BAR, bar=bar_idx)[:2]
+    # paddrs are returned as (paddr, size) pairs until a (paddr=0, size=0) terminator in the beginning of the mapping.
+    paddrs_raw = list(itertools.takewhile(lambda p: p[1] != 0, zip(memview.view(fmt='Q')[0::2], memview.view(fmt='Q')[1::2])))
+    return memview, [p + i for p, sz in paddrs_raw for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]
+  def read_config(self, offset:int, size:int): return self._rpc(RemoteCmd.CFG_READ, 0, offset, size)[0]
+  def write_config(self, offset:int, value:int, size:int): self._rpc(RemoteCmd.CFG_WRITE, 0, offset, size, value)
+  def reset(self): self._rpc(RemoteCmd.RESET, 0, 0, 0)
   def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface:
-    return RemoteMMIOInterface(self, bar, size or self.bar_info(bar)[1], fmt).view(off, size, fmt)
-  def resize_bar(self, bar_idx:int): pass
+    return RemoteMMIOInterface(self, bar, size or self.bar_info[bar].size, fmt).view(off, size, fmt)
 
 class APLRemotePCIDevice(RemotePCIDevice):
   APP_PATH = "/Applications/TinyGPU.app/Contents/MacOS/TinyGPU"
 
-  @classmethod
-  def ensure_app(cls):
-    sip_disabled = "disabled" in subprocess.run(["csrutil", "status"], capture_output=True, text=True).stdout.lower()
-    commit = "c4d5a845763e6694ef53ff1c0c8c6d3e130c3724" if sip_disabled else "81f5bb45ba0d65edefa833475b71cd17a470c2c5"
-    app_name = f"TinyGPU_{commit}.zip"
-    if (_ensure_downloads_dir() / app_name).is_file() and os.path.exists(cls.APP_PATH): return
+  @staticmethod
+  def install_tinygpu():
     print("Downloading TinyGPU.app...")
-    with contextlib.suppress(RuntimeError): system("pkill -f TinyGPU")
-    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"))
+    system(f"ditto -xk {fetch('https://github.com/nimlgen/tinygpu_releases/raw/8120b5508b43149d27bf22f9a4e6d7c5a4b401e9/TinyGPU.zip')} /Applications")
+    print(system(f"{APLRemotePCIDevice.APP_PATH} install"))
 
-  def __init__(self, devpref:str, pcibus:str):
-    self.ensure_app()
+  def __init__(self, devpref:str, pcibus:str, bars:list[int], resize_bars:list[int]|None=None):
     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):
@@ -368,12 +338,21 @@ class APLRemotePCIDevice(RemotePCIDevice):
       if i == 0: subprocess.Popen([self.APP_PATH, "server", sock_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
       time.sleep(0.05)
     else: raise RuntimeError(f"Failed to connect to TinyGPU server at {sock_path}.")
-    super().__init__(devpref, "usb4", sock=sock)
+    super().__init__(devpref, pcibus, bars, sock)
 
-  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
-    mapped_size, _, _, fd = self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_SYSMEM_FD, size, has_fd=True)
-    memview = MMIOInterface(FileIOInterface(fd=fd).mmap(0, mapped_size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, 0), mapped_size, fmt='B')
+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
 
-    # paddrs are returned as (paddr, size) pairs until a (paddr=0, size=0) terminator in the beginning of the mapping.
-    paddrs_raw = list(itertools.takewhile(lambda p: p[1] != 0, zip(memview.view(fmt='Q')[0::2], memview.view(fmt='Q')[1::2])))
-    return memview, [p + i for p, sz in paddrs_raw for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]
+  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

tinygrad/runtime/support/usb.py

@@ -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 not (rlen > 0 and send_data is not None), "BOT mode only supports either read or write per command"
+        assert (data_len == 0) if dir_in else (rlen == 0), "BOT mode only supports either read or write per command"
 
         # CBW
         self._tag += 1

tinygrad/schedule/allreduce.py

@@ -60,4 +60,4 @@ def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|Non
   to = red.param_like(0)
   src = buf.param_like(1)
   red = src.allreduce(red.arg, red.src[1])
-  return output.after(to.after(to.store(handle_allreduce(src, red))).sink().call(output, buf.contiguous(), name="allreduce", precompile=True))
+  return output.after(to.assign(handle_allreduce(src, red)).sink().call(output, buf.contiguous(), name="allreduce", precompile=True))

tinygrad/schedule/indexing.py

@@ -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.AFTER, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
+ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, 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}
 
@@ -17,23 +17,21 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
   for s in rb.src:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
-def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp):
-  # don't realize COPY/BUFFER_VIEW when they are the direct source of STORE+AFTER — the target buffer is the output
-  if src.op in {Ops.COPY, Ops.BUFFER_VIEW} and src in ctx \
-     and not dest.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
-    del ctx[src]
+def realize_assign_src(ctx:dict[UOp, None], buf:UOp, x:UOp):
+  # don't realize COPY/BUFFER_VIEW when they are the direct source of ASSIGN — the ASSIGN target buffer is the output
+  if x.op in {Ops.COPY, Ops.BUFFER_VIEW} and x in ctx \
+     and not buf.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
+    del ctx[x]
   # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
-  if dest.base in src.backward_slice_with_self: ctx[src] = None
+  if buf.base in x.backward_slice_with_self: ctx[x] = None
 
 pm_generate_realize_map = PatternMatcher([
   # always realize
-  (UPat({Ops.COPY, Ops.CONTIGUOUS}, name="tr"), realize),
-  # realize AFTER of STORE+AFTER
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE)), allow_any_len=True, name="tr"), realize),
+  (UPat({Ops.COPY, Ops.CONTIGUOUS, Ops.STORE, Ops.ASSIGN}, name="tr"), realize),
   # realize srcs of these
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
-  # sometimes realize/unrealize src of store+after
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat.var("dest"), UPat.var("src"))))), realize_store_after_src),
+  # sometimes realize src of assign
+  (UPat(Ops.STORE, src=(UPat.var("buf"), UPat.var("x"))), realize_assign_src),
 ])
 
 @dataclass(frozen=True)
@@ -60,8 +58,7 @@ 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} or \
-       (s.op is Ops.AFTER and not any(c.op in {Ops.STORE, Ops.END} for c in s.src[1:])):
+    if s.op in {Ops.PARAM, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
       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]
@@ -100,11 +97,25 @@ 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
@@ -166,8 +177,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
 
-    # 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
+    # no range on after
+    if x.op is Ops.AFTER: continue
 
     # treat MSTACK/MSELECT like SINK
     if x.op in {Ops.MSTACK, Ops.MSELECT}: continue

tinygrad/schedule/multi.py

@@ -107,7 +107,9 @@ def copy_multi(multi:UOp, device:str | tuple[str, ...] | UOp):
   assert multi.axis is not None, "all multi ops have axis"
   return multi.src[0]._unshard(multi.axis).allreduce(Ops.ADD, device)
 
-def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).multi(src.axis)
+def assign_multi(dest:UOp, src:UOp):
+  if dest.axis != src.axis: raise RuntimeError(f"axis must match in assign {dest.axis} != {src.axis}")
+  return dest.src[0].assign(src.src[0]).multi(src.axis)
 
 def passthrough_multi(root:UOp, multi:UOp):
   return UOp(root.op, root.dtype, (multi.src[0],)+tuple(x.src[0] if x.op is Ops.MULTI else x for x in root.src[1:]), root.arg).multi(multi.axis)
@@ -116,11 +118,6 @@ def rewrite_into_call(call:UOp):
   if not should_resolve_call(call): return None
   new_body = graph_rewrite(call.src[0], multi_pm, name="subcall")
   new_args = tuple(a.src[0] if a.op is Ops.MULTI else a for a in call.src[1:])
-  # after multi resolution, TUPLE elements may be MULTI — strip MULTI from body, create per-shard CALL, wrap each GETTUPLE in its own MULTI
-  assert new_body.op is Ops.TUPLE
-  if any(s.op is Ops.MULTI for s in new_body.src):
-    shard_call = call.replace(src=(UOp.maketuple(*[s.src[0] if s.op is Ops.MULTI else s for s in new_body.src]),)+new_args)
-    return UOp.maketuple(*[shard_call.gettuple(i).multi(s.axis) if s.op is Ops.MULTI else shard_call.gettuple(i) for i, s in enumerate(new_body.src)])
   return call.replace(src=(new_body,)+new_args)
 
 def param_to_multi(p:UOp):
@@ -138,24 +135,19 @@ multi_pm = PatternMatcher([
   (UPat(Ops.SHRINK, src=(UPat(Ops.MULTI, name="multi"), UPat(), UPat()), name="root"), shrink_multi),
   (UPat(Ops.PERMUTE, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), permute_multi),
   (UPat(Ops.FLIP, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), flip_multi),
-  (UPat(Ops.AFTER, src=(UPat(Ops.MULTI), UPat(Ops.STORE, src=(UPat(Ops.MULTI, name="dest"), UPat(Ops.MULTI, name="src"))))), store_after_multi),
+  (UPat(Ops.ASSIGN, src=(UPat(Ops.MULTI, name="dest"), UPat(Ops.MULTI, name="src"))), assign_multi),
   (UPat(Ops.COPY, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.DEVICE, name="device"))), copy_multi),
   (UPat(Ops.ALLREDUCE, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.DEVICE, name="device")), name="red"),
     lambda multi,device,red: multi.src[0].allreduce(red.arg, device).multi(axis=multi.axis)),
-
-  # resolve TUPLE+GETTUPLE (needed in multi)
-  (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
-  # GETTUPLE on MULTI: passthrough MULTI (e.g. when CALL was replaced by MULTI(GETTUPLE(...)))
-  (UPat(Ops.GETTUPLE, src=(UPat(Ops.MULTI, name="multi"),), name="g"),
-    lambda g, multi: multi.src[0].gettuple(g.arg).multi(multi.axis) if multi.src[0].op in {Ops.CALL, Ops.TUPLE}
-    else multi),
   # rewrite into calls explicitly for MULTI
   (UPat(Ops.CALL, name="call"), rewrite_into_call),
-  (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 non-value-producing CALLs (custom kernels, etc.) — TUPLE body CALLs are handled by rewrite_into_call
+  (UPat(Ops.CALL, 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)
-    if root.src[0].op is not Ops.TUPLE else None),
+    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

tinygrad/schedule/rangeify.py

@@ -8,7 +8,7 @@ from tinygrad.helpers import prod, all_same, getenv, dedup, all_int, DEBUG, SPLI
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
-from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, IndexingContext, apply_movement_op
+from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
 from tinygrad.schedule.multi import multi_pm
 from tinygrad.schedule.allreduce import create_allreduce_function
 
@@ -16,36 +16,22 @@ 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 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
+  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]
   ctx[x] = assign
 
-# *** fold moved AFTERs (hack for openpilot) ***
+# *** fold moved ASSIGNs (hack for openpilot) ***
 pm_fold_moved_assign = PatternMatcher([
   (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
@@ -53,36 +39,30 @@ pm_fold_moved_assign = PatternMatcher([
 ])
 
 # movement op on INDEX as a PatternMatcher
-# TODO: clean up .src[0]._shape is not None
 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)
-     if r.src[0]._shape is not None and 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)
+   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
   (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)
-     if a.src[0]._shape is not None and not any(s.op is Ops.STORE and s.src[0]._shape is not None for s in a.src[1:]) else None),
+   lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], r.arg)),
   (UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])),
 ])
 
 # *****************
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
-def fix_store_after_hazard(after:UOp, target:UOp, src:UOp):
+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())
-  base = target.base
-  reaches_base: dict[UOp, bool] = {}
-  for s in src.toposort(gate=lambda s: s.op is not Ops.CONTIGUOUS):
-    reaches_base[s] = s is base or any(reaches_base.get(c) for c in s.src)
-    if reaches_base[s] and s.op in unsafe: return after.replace(src=(after.src[0], target.store(src.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)):
+    return assign.replace(src=(target, src.contiguous()))
 
-def normalize_store_after_target_chain(after:UOp, target:UOp, src:UOp):
+def normalize_assign_target_chain(assign:UOp, target:UOp, src:UOp):
   root_target = target
-  while root_target.op is Ops.AFTER: root_target = root_target.src[0]
+  while root_target.op is Ops.ASSIGN: 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 after.replace(src=(root_target, root_target.store(src)))
+  return assign.replace(src=(root_target, src))
 
 def split_reduceop(reduce:UOp, x:UOp):
   if prod(reduce.shape) == 0: return None
@@ -140,10 +120,8 @@ 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
@@ -169,18 +147,20 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   # copy only to different device
   (UPat(Ops.COPY, src=(UPat.var("x"), UPat()), name="copy"), lambda x,copy: x.f(Ops.NOOP) if x.device == copy.device else None),
 
-  # ** assign rules (STORE+AFTER) **
+  # ** assign rules **
 
-  # move bitcast from store+after target to source
-  (UPat(Ops.AFTER, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(Ops.STORE, src=(UPat(Ops.BITCAST), UPat(name="src"))))),
-   lambda target, src: target.after(target.store(src.bitcast(target.dtype)))),
+  # collapse nested ASSIGN to the same buffer (e.g. __iadd__ in __setitem__)
+  (UPat(Ops.ASSIGN, src=(UPat(name="target"), UPat(Ops.ASSIGN, src=(UPat(name="target"), UPat()), name="src"))), lambda target, src: src),
+
+  # move bitcast from assign target to source: a.bitcast(X).assign(src) -> a.assign(src.bitcast(a.dtype))
+  (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),
 
   # make source contiguous if it has hazardous movement ops on the dest buffer
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="target"), UPat(name="src")))), name="after"), fix_store_after_hazard),
-
-  # normalize target chain: walk through AFTERs to root, insert contiguous if needed
-  (UPat(Ops.AFTER, src=(UPat(Ops.AFTER, name="target"), UPat(Ops.STORE, src=(UPat(), UPat(name="src")))), name="after"),
-   lambda after, target, src: normalize_store_after_target_chain(after, target, src)),
+  (UPat(Ops.STORE, src=(UPat.var("target"), UPat.var("src")), name="assign"), fix_assign_hazard),
 
   # ** size 0 **
 
@@ -194,12 +174,12 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
 # *****************
 # 3.5 cleanups
 
-ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.NOOP}
+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 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
+  # don't optimize ALWAYS_RUN_OPS
+  if b.src[0].op in ALWAYS_RUN_OPS: return None
 
   new_rng = []
   hit = False
@@ -373,19 +353,21 @@ 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 (after:=x.src[0]).op is Ops.AFTER:
-    buf = after.src[0].buf_uop.base
-    if not (stores := [s for s in after.src[1:] if s.op is Ops.STORE and s.src[0].op is Ops.INDEX]): return buf
-    # BUFFERIZE(INDEX(...)); store through the underlying global index instead.
-    ended_stores = []
-    store_target = stores[0].src[0]
-    if store_target.src[0].op is Ops.BUFFERIZE and store_target.src[0].src[0].op is Ops.INDEX:
-      store_target = store_target.src[0].src[0]
-    if stores[0].src[1] is not store_target:  # skip self-assign
-      end_rngs = sorted(dedup(tuple(store_target.ranges) + tuple(rngs)), key=lambda x: x.arg)
-      ended_stores.append(store_target.replace(dtype=sdtype).store(stores[0].src[1]).end(*end_rngs))
-    return buf.after(*ended_stores)
+  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:
+      # 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)
+    return ret
 
   # NOTE: the DEFINE_LOCAL needs to be disambiguated here
   if sdtype.addrspace == AddrSpace.GLOBAL:
@@ -534,8 +516,6 @@ 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()

tinygrad/tensor.py

@@ -242,8 +242,7 @@ class Tensor(OpMixin):
       param = UOp.param(slot, self.dtype, self.shape, self.device)
     return Tensor(param, device=self.device)
   def call(self, *lst:Tensor, fxn:Tensor|UOp, grad_fxn:Callable|None=None) -> Tensor:
-    fret = (fxn.uop if isinstance(fxn, Tensor) else fxn).call(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn)
-    return Tensor(fret.gettuple(0), device=self.device)
+    return Tensor((fxn.uop if isinstance(fxn, Tensor) else fxn).call(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn), device=self.device)
 
   def custom_kernel(self, *lst:Tensor, fxn:Callable, grad_fxn:Callable|None=None) -> list[Tensor]:
     """
@@ -309,18 +308,13 @@ class Tensor(OpMixin):
     if is_disk:
       self._buffer().copyin(x._data())
       return self
-    # STORE+AFTER: STORE is the write effect (void), AFTER wraps the view for correct shape/ranging
-    assign = self.uop.after(self.uop.store(x.uop))
-    if (base := self.uop.base).op in {Ops.BUFFER, Ops.AFTER} and self.uop is not base and not self.uop.has_buffer_identity():
-      # view assign: replace at the buffer-identity level (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
-      ib = self.uop
-      while not ib.has_buffer_identity() and ib is not base: ib = ib.src[0]
-      assigned_ib = ib.after(assign)
-      _apply_map_to_tensors({ib: assigned_ib}, name="Embed View Assign", walk=True)
-    else:
-      # simple assign
-      self.uop = assign
-    return self
+    # NOTE: assign_uop is created before AFTER embedding (uses original self.uop),
+    # but AFTER must be embedded before _apply_uop (so subsequent assigns see it)
+    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)
+    return self.replace(self._apply_uop(lambda *_: assign_uop, x))
 
   def detach(self) -> Tensor:
     """
@@ -1337,8 +1331,7 @@ class Tensor(OpMixin):
              if (t:=tref()) is not None and t is not self and t.uop is not v_uop and t.uop not in v_bw):
         raise RuntimeError("can't setitem on a tensor that already has other uses and requires grad")
       if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
-      # __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value
-      if v.uop.op is Ops.AFTER and any(s.op is Ops.STORE for s in v.uop.src[1:]): v = v._apply_uop(lambda x: x.src[1].src[1])
+      if v.uop.op is Ops.ASSIGN: v = v._apply_uop(lambda x: x.src[1])
       self.replace(self._getitem(indices, v))
       return
     idx = [indices] if (isinstance(indices, list) and all_int(indices)) or not isinstance(indices, (tuple, list)) else list(indices)
@@ -1347,14 +1340,14 @@ 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.BUFFER or self.uop._base_buffer_is_realized(): # basic setitem
+    elif is_disk or self.uop.is_realized or self.uop.base.op is Ops.AFTER: # basic setitem, self is realized
       view = self[indices]
-      if isinstance(v, Tensor) and v.uop.op is Ops.AFTER and v.uop in view.uop.base.src: return
+      if isinstance(v, Tensor) and v.uop.op is Ops.ASSIGN and v.uop in view.uop.base.src: return
       view.assign(v)
     else: # basic setitem, self is not realized
       if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
-      # __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value
-      if v.uop.op is Ops.AFTER and any(s.op is Ops.STORE for s in v.uop.src[1:]): v = v._apply_uop(lambda x: x.src[1].src[1])
+      # __iadd__/__isub__ on unrealized views creates a no-op ASSIGN; unwrap to get the computed value
+      if v.uop.op is Ops.ASSIGN: v = v._apply_uop(lambda x: x.src[1])
       self.replace(self._getitem(indices, v))
 
   def __delitem__(self, indices) -> None:
@@ -2491,7 +2484,7 @@ class Tensor(OpMixin):
     """
     if IMAGE: return self.image_dot(w, dtype)
     if ASM_GEMM:
-      from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
+      from extra.gemm.asm.cdna.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")
@@ -3662,29 +3655,21 @@ class Tensor(OpMixin):
 
     # contiguous creates the image, and early realize static weights (TODO: test for the static weight)
     if IMAGE == 1:
-      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, use at to specify the dimension to pad in, defaults to first
-      def pad_align(t, dim, at=None, force=False):
-        # align to 64 pixels when height is real, otherwise 64 bytes is sufficient
-        align = (64 // dtsz) if prod(t.shape[:dim]) == 1 or prod(t.shape) < 16384 * 4 else 256
-        return ipad(t, at:=at or dim, round_up(t.shape[at] + int(force), align // math.gcd(prod(t.shape[dim:]) // t.shape[at], align)))
-
-      # bank conflicts
-      if cin >= 8 and is_pow2(cin // 4):
-        x, w = pad_align(x.reshape(bs, iy, ix, groups, cin // 4, 4), 2, at=4, force=True), pad_align(w, 1, at=2, force=True)
-      else: x, w = pad_align(x, 2), pad_align(w, 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))
 
       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
-      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, ...]
+      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()

tinygrad/uop/__init__.py

@@ -39,9 +39,6 @@ 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
@@ -85,7 +82,7 @@ class Ops(FastEnum):
   # ** 6 -- ops that don't exist in programs **
 
   # tensor graph ops
-  UNIQUE = auto(); DEVICE = auto()
+  UNIQUE = auto(); DEVICE = auto(); ASSIGN = auto()
 
   # local unique
   LUNIQUE = auto()

tinygrad/uop/ops.py

@@ -207,22 +207,11 @@ 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.STORE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
+      case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | 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.TUPLE:
+           Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS:
         return None
 
-      case Ops.GETTUPLE:
-        # GETTUPLE extracts from a TUPLE (possibly through a CALL)
-        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
-        inner_shape = in_tuple.src[self.arg]._shape
-        if inner_shape is None: return None
-        # if through a CALL, substitute internal PARAMs in the shape with corresponding args
-        if self.src[0].op is Ops.CALL:
-          return tuple(graph_rewrite(s, _pm_resolve_params, self.src[0].src[1:], walk=True) if isinstance(s, UOp) else s for s in inner_shape)
-        return inner_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):
@@ -253,7 +242,11 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END:
         return self.src[0]._shape
 
-      case Ops.CALL: return None
+      case Ops.CALL:
+        inner_shape = self.src[0]._shape
+        if inner_shape is None: return None
+        # substitute internal PARAMs in the shape with corresponding args
+        return tuple(graph_rewrite(s, _pm_resolve_params, self.src[1:], walk=True) if isinstance(s, UOp) else s for s in inner_shape)
 
       # TODO: disallow shape changing bitcast
       case Ops.BITCAST:
@@ -308,8 +301,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
             raise ValueError(f"invalid type for axis: {axis_arg}")
           return tuple(1 if i in axis_arg else s for i,s in enumerate(ps))
 
+    if self.op is Ops.STORE: 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}):
+    if self.op in GroupOp.ALU.union({Ops.CAST, Ops.COPY, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE, Ops.STORE}):
       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}")
@@ -409,12 +404,6 @@ 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]))
@@ -425,8 +414,6 @@ 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))
@@ -457,9 +444,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return UOp(Ops.GEP, self.dtype.scalar().vec(len(i)) if len(i) > 1 else self.dtype.scalar(), (self,), i)
   def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
   def store(self, src:UOp|ConstType, **kwargs):
-    return UOp(Ops.STORE, dtypes.void, (self, self.const_like(src) if not isinstance(src, UOp) else src), **kwargs)
+    return UOp(Ops.STORE, kwargs.pop("dtype", dtypes.void), (self, UOp.const(self.dtype, src) if not isinstance(src, UOp) else src), **kwargs)
   def end(self, *src:UOp): return UOp(Ops.END, src=(self,)+src) if len(src) else self
   def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, self.dtype, (self,)+src, **kwargs) if len(src) else self
+  def assign(self, x:UOp): return self.after(self.store(x))
   def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
   def contract(self, *rngs:UOp):
     assert all(x.arg[-1] == AxisType.UPCAST for x in rngs), "all contract ranges must be upcast"
@@ -536,10 +524,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     # COPY removes axis. TODO: add more tests for this, and consider MSELECT/MSTACK
     if self.op is Ops.COPY: return None
     if self.op is Ops.MULTI: return self.arg
-    # GETTUPLE: axis comes from the specific TUPLE element, not src[0]
-    if self.op is Ops.GETTUPLE:
-      in_tuple = self.src[0].src[0] if self.src[0].op is Ops.CALL else self.src[0]
-      return in_tuple.src[self.arg].axis if in_tuple.op is Ops.TUPLE else None
     # PARAM: axis is stored as a MULTI source
     if self.op is Ops.PARAM:
       for s in self.src:
@@ -697,15 +681,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def has_buffer_identity(self):
     """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/MULTI -> BUFFER chain)."""
     if self.op in {Ops.RESHAPE, Ops.MULTI}: return self.src[0].has_buffer_identity()
-    if self.op is Ops.GETTUPLE and self.src[0].op is Ops.TUPLE: return self.src[0].src[self.arg].has_buffer_identity()
     return self.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.PARAM}
 
-  def _base_buffer_is_realized(self) -> bool:
-    """Walk through AFTER chain to find if the underlying buffer is realized (has allocated memory)."""
-    u = self.base
-    while u.op is Ops.AFTER: u = u.src[0]
-    return u.is_realized
-
   @property
   def buffer(self) -> Buffer|MultiBuffer:
     from tinygrad.device import Buffer, MultiBuffer
@@ -924,13 +901,9 @@ 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
 
-  _NO_TUPLE_WRAP = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.BUFFER_VIEW, Ops.CUSTOM_FUNCTION, Ops.TUPLE}
-  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:
+  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), name:str|None=None, precompile:bool=False) -> UOp:
     assert len(self.ranges) == 0, f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}"
-    # value-producing bodies are always wrapped in TUPLE so CALL dtype is always void
-    body = self if self.op in UOp._NO_TUPLE_WRAP else UOp.maketuple(self)
-    return UOp(Ops.CALL, dtypes.void, (body,)+srcs, CallInfo(grad_fxn, metadata, name, precompile, precompile_backward))
+    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata, name, precompile))
   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)]
@@ -954,12 +927,9 @@ 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, 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 __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 should_resolve_call(c:UOp) -> bool:
   # don't resolve real kernel calls, sink or program
@@ -1081,6 +1051,7 @@ class UPat(OpMixin):
   def gep(self, i:int|None=None, **kwargs): return UPat(Ops.GEP, None, (self,), (i,) if i is not None else None, **kwargs)
   def load(self, *src:UPat, **kwargs): return UPat(Ops.LOAD, src=(self,)+src, **kwargs)
   def store(self, *src:UPat, **kwargs): return UPat(Ops.STORE, self.match_dtype, (self,)+src, **kwargs)
+  def assign(self, x:UPat, **kwargs): return UPat(Ops.ASSIGN, self.match_dtype, (self,x), **kwargs)
   def reduce(self, *src:UPat, **kwargs): return UPat(Ops.REDUCE, self.match_dtype, src=(self,)+src, **kwargs)
   def broadcast(self, **kwargs): return UPat(Ops.VECTORIZE, self.match_dtype, src=self, **kwargs)
   def contiguous(self, *args, **kwargs): return UPat(Ops.CONTIGUOUS, dtype=self.match_dtype, src=(self,)+args, **kwargs)
@@ -1527,7 +1498,7 @@ def render_marg(ctx,x:UOp):
   return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)"
 
 sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.UNIQUE, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY,
-         Ops.WHERE, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH}
+         Ops.WHERE, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.ASSIGN, Ops.DETACH}
 pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE, name="u"), UPat(Ops.DEVICE, name="d")), name="x"),
    lambda x,u,d: f"UOp.unique_const({x.dtype}, {x.arg}, device={repr(d.arg)}, unique={u.arg})"),
@@ -1595,7 +1566,7 @@ def pyrender(ast:UOp) -> str:
   cmap = consumer_map_from_toposort(lst)
   not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE}
   always_rendered = {Ops.PARAM, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.VECTORIZE,
-                     Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.WHERE, Ops.END}
+                     Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.WHERE, Ops.END, Ops.ASSIGN}
 
   to_render: set[UOp] = {ast}
   for u in lst:

tinygrad/uop/spec.py

@@ -59,9 +59,6 @@ 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)
 ])
@@ -77,9 +74,8 @@ movement_ops = PatternMatcher([
   (UPat((Ops.VECTORIZE, Ops.VCONST), dtype=dtypes.index), lambda: True),
   (UPat({Ops.ADD, Ops.MUL, Ops.IDIV}, dtype=dtypes.index), lambda: True),
 
-  # AFTER on Movement Op, BUFFER, COPY, or BITCAST
-  (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.MULTI, Ops.CONTIGUOUS, Ops.BUFFER, Ops.BITCAST, Ops.COPY})),), allow_any_len=True),
-   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),
 ])
 
 _tensor_spec = PatternMatcher([
@@ -97,6 +93,12 @@ _tensor_spec = PatternMatcher([
   # KERNEL can attach to an AFTER to describe the compute required to realize a BUFFER
   (UPat(Ops.CALL, src=UPat((Ops.BUFFER, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
 
+  # ASSIGN has a target and a value. It can also optionally depend on other assigns
+  (UPat(Ops.ASSIGN, name="x"), lambda x: len(x.src) >= 2 and all(s.op is Ops.ASSIGN for s in x.src[2:])),
+
+  # STORE in tensor graph: store a value into a target
+  (UPat(Ops.STORE, dtypes.void, (UPat(), UPat())), lambda: True),
+
   # MSELECT chooses one of the multi buffers
   (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
 
@@ -130,17 +132,13 @@ _tensor_spec = PatternMatcher([
   (UPat(Ops.REDUCE_AXIS, name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) >= 2 and x.arg[0] in {Ops.ADD, Ops.MUL, Ops.MAX}),
 
   # AFTER if things were kernelized
-  (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.AFTER)),), allow_any_len=True), lambda: True),
+  (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.AFTER, Ops.PARAM)),), allow_any_len=True), lambda: True),
 
-  # allow CALL/PARAM/CUSTOM_FUNCTION — CALL dtype is always void
-  (UPat(Ops.CALL, dtypes.void), lambda: True),
+  # allow CALL/PARAM/CUSTOM_FUNCTION
+  (UPat(Ops.CALL, src=(UPat(name="f"),), name="c", allow_any_len=True), lambda c,f: c.dtype == f.dtype),
   (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?)
@@ -235,8 +233,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 (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),
+  # make sure all index dtypes have been lowered
+  (UPat(GroupOp.All, 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))),

tinygrad/uop/symbolic.py

@@ -65,8 +65,6 @@ 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))),
 ])
@@ -140,6 +138,12 @@ 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 ********

tinygrad/viz/index.html

@@ -157,6 +157,10 @@
   #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;

tinygrad/viz/js/index.js

@@ -248,37 +248,6 @@ function selectShape(key) {
   return { eventType:track?.eventType, e:track?.shapes[idx] };
 }
 
-// scaling function for time to pixels
-const timelineScale = () => d3.scaleLinear().domain([data.first, data.dur]).range([0, document.getElementById("timeline").clientWidth])
-
-function timeAtCycle(clk) {
-  if (clk < data.instSt || clk > data.instEt || data.tracks.get("Shader Clock") == null) return "-";
-  let cur = data.instSt, ns = 0, freq = null;
-  // walk through all frequency changes and accumulate time in nanoseconds
-  for (const [s, v] of data.tracks.get("Shader Clock").valueMap) {
-    if (freq != null && cur < s) {
-      const et = Math.min(clk, s);
-      ns += (et - cur) * 1e9 / freq;
-      cur = et;
-      if (cur === clk) break;
-    }
-    freq = v;
-  }
-  // ending cycles use the last known frequency
-  if (cur < clk) ns += (clk - cur) * 1e9 / freq;
-  const remNs = Math.round(ns % 1000);
-  return ns/1000>1 ? formatMicroseconds(ns / 1000, true) + (remNs ? ` ${remNs}ns` : "") : Math.round(ns)+"ns";
-}
-
-function getZoomIdentity() {
-  // for packets, set zoom to the full range of instruction events
-  if (data.instSt != null) {
-    const k = (data.dur - data.first) / (data.instEt - data.instSt), xscale = timelineScale();
-    return d3.zoomIdentity.translate(-xscale(data.instSt) * k, 0).scale(k);
-  }
-  return d3.zoomIdentity;
-}
-
 const Modes = {0:'read', 1:'write', 2:'write+read'};
 
 function setFocus(key) {
@@ -287,8 +256,8 @@ function setFocus(key) {
     // adjust zoom if the entire shape is off screen
     const { eventType, e } = selectShape(key);
     if (e != null) {
-      const xscale = timelineScale();
       const [x0, x1] = eventType === EventTypes.EXEC ? [e.x, e.x+e.width] : [e.x[0], e.x.at(-1)];
+      const xscale = d3.scaleLinear().domain([data.first, data.dur]).range([0, document.getElementById("timeline").clientWidth]);
       const [st, et] = xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale);
       if (x1 < st || x0 > et) zoomLevel = d3.zoomIdentity.translate(-xscale((x0+x1)/2-(et-st)/2)*zoomLevel.k, 0).scale(zoomLevel.k);
     }
@@ -299,9 +268,7 @@ function setFocus(key) {
   const html = d3.select(".info").html("");
   if (eventType === EventTypes.EXEC) {
     const [n, _, ...rest] = e.arg.tooltipText.split("\n");
-    const tableData = [["Name", colored(e.arg.label)], ["Duration", formatTime(e.width)], ["Start Time", formatTime(e.x)]];
-    if (data.instSt != null) tableData.push(["Timestamp", timeAtCycle(e.x)]);
-    html.append(() => tabulate(tableData));
+    html.append(() => tabulate([["Name", colored(e.arg.label)], ["Duration", formatTime(e.width)], ["Start Time", formatTime(e.x)]]));
     let group = html.append("div").classed("args", true);
     for (const r of rest) group.append("p").text(r);
     group = html.append("div").classed("args", true);
@@ -337,11 +304,10 @@ function setFocus(key) {
   let instList = document.getElementById("insts");
   if (data.pcToShape.size == 0) return d3.select(instList?.parentElement).html("");
   if (instList == null) {
-    let contents = "";
+    let contents = "", i = 0;
     for (const [k, v] of data.pcToShape) {
-      const pcHex = v.pc.toString(16);
-      contents += `<div class="line" data-k="${k}"><span class="left" id="inst-${k}"><span class="wave">${v.wave}</span>
-        <span class="pc">${"0x"+pcHex.padStart(Math.max(4, Math.ceil(pcHex.length/4)*4), 0)}</span><span class="label">${data.pcMap[v.pc]}</span></div>`;
+      contents += `<div class="line" data-k="${k}"><span class="left" id="inst-${k}"><span class="n">${i++}</span><span class="wave">${v.wave}</span>
+        <span class="pc">${"0x"+v.pc.toString(16).padStart(12, "0")}</span></span><span class="label">${data.pcMap[v.pc]}</span></div>`;
     }
     instList = d3.create("pre").append("code").classed("hljs", true).style("margin-top", "20px").attr("id", "insts").html(contents)
       .on("click", e => { const line = e.target.closest(".line"); line && setFocus(line.dataset.k); }).node();
@@ -379,10 +345,8 @@ async function renderProfiler(path, opts) {
   const textDecoder = new TextDecoder("utf-8");
   const { strings, dtypeSize, markers, ...extData } = JSON.parse(textDecoder.decode(new Uint8Array(buf, offset, indexLen))); offset += indexLen;
   // place devices on the y axis and set vertical positions
-  const [tickSize, padding, baseOffset] = [5, 8, markers.length ? 14 : 0];
-  const secondaryTick = opts.unit == "clk" ? timeAtCycle : null;
-  const axisHeight = secondaryTick != null ? tickSize*2+(padding*2) : tickSize;
-  const deviceList = profiler.append("div").attr("id", "device-list").style("padding-top", axisHeight+padding+baseOffset+"px");
+  const [tickSize, padding, baseOffset] = [10, 8, markers.length ? 14 : 0];
+  const deviceList = profiler.append("div").attr("id", "device-list").style("padding-top", tickSize+padding+baseOffset+"px");
   const canvas = profiler.append("canvas").attr("id", "timeline").node();
   // NOTE: scrolling via mouse can only zoom the graph
   canvas.addEventListener("wheel", e => (e.stopPropagation(), e.preventDefault()), { passive:false });
@@ -552,13 +516,6 @@ async function renderProfiler(path, opts) {
   for (const m of markers) m.label = m.name.split(/(\s+)/).map(st => ({ st, color:m.color, width:ctx.measureText(st).width }));
   data.pcToShape = new Map([...data.pcToShape].sort((a, b) => a[1].st - b[1].st));
   if (extData.pcMap != null) data.pcMap = extData.pcMap; setFocus(focusedShape);
-  // secondary axis mapping
-  let instRange = null;
-  for (const [k, { shapes }] of data.tracks) if (k.startsWith("WAVE")) {
-    const first = shapes[0].x, last = shapes.at(-1).x;
-    instRange = instRange == null ? [first, last] : [Math.min(first, instRange[0]), Math.max(last, instRange[1])];
-  }
-  if (instRange != null) [data.instSt, data.instEt] = instRange;
   updateProgress(Status.COMPLETE);
   // draw events on a timeline
   const dpr = window.devicePixelRatio || 1;
@@ -632,24 +589,19 @@ async function renderProfiler(path, opts) {
     }
     // draw axes
     ctx.translate(0, baseOffset);
-    const y = secondaryTick != null ? tickSize+padding : 0;
-    drawLine(ctx, xscale.range(), [y, y]);
+    drawLine(ctx, xscale.range(), [0, 0]);
     let lastLabelEnd = -Infinity;
     for (const tick of xscale.ticks()) {
       if (!Number.isInteger(tick)) continue;
       const x = xscale(tick);
-      drawLine(ctx, [x, x], [y, y+tickSize]);
+      drawLine(ctx, [x, x], [0, tickSize]);
       const labelX = x+ctx.lineWidth+2;
       if (labelX <= lastLabelEnd) continue;
+
       const label = formatTime(tick, et-st <= 1e3);
       ctx.textBaseline = "top";
-      ctx.fillText(label, labelX, y+tickSize);
+      ctx.fillText(label, labelX, tickSize);
       lastLabelEnd = labelX + ctx.measureText(label).width + 4;
-      if (secondaryTick != null) {
-        drawLine(ctx, [x, x], [y, y-tickSize]);
-        const label = secondaryTick(tick, st, et); ctx.fillText(label, labelX, 0);
-        lastLabelEnd = Math.max(lastLabelEnd, labelX + ctx.measureText(label).width + 4);
-      }
     }
     if (data.axes.y != null) {
       drawLine(ctx, [0, 0], data.axes.y.range);
@@ -686,10 +638,10 @@ async function renderProfiler(path, opts) {
     canvas.style.height = `${height}px`;
     canvas.style.width = `${width}px`;
     ctx.scale(dpr, dpr);
-    zoomLevel = getZoomIdentity();
     d3.select(canvas).call(canvasZoom.transform, zoomLevel);
   }
 
+  zoomLevel = d3.zoomIdentity;
   canvasZoom = d3.zoom().filter(vizZoomFilter).on("zoom", e => render(e.transform));
   d3.select(canvas).call(canvasZoom);
   document.addEventListener("contextmenu", e => e.ctrlKey && e.preventDefault());
@@ -744,7 +696,7 @@ d3.select("#graph-svg").call(svgZoom);
 document.getElementById("zoom-to-fit-btn").addEventListener("click", () => {
   const canvas = d3.select("#timeline");
   if (!canvas.empty() && rect(canvas.node()).width !== 0) {
-    return canvas.call(canvasZoom.transform, getZoomIdentity());
+    return canvas.call(canvasZoom.transform, d3.zoomIdentity);
   }
   const svg = d3.select("#graph-svg");
   svg.call(svgZoom.transform, d3.zoomIdentity);

tinygrad/viz/serve.py

@@ -45,7 +45,7 @@ from tinygrad.dtype import dtypes
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
                **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.REDUCE_AXIS: "#FF6B6B",
-               Ops.RANGE: "#c8a0e0", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
+               Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#909090", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
                Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.INS: "#eec4ff",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
                Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
@@ -339,35 +339,29 @@ def load_amd_counters(ctxs:list[dict], profile:list[ProfileEvent]) -> None:
 
 def sqtt_timeline(data:bytes, lib:bytes, target:str) -> list[ProfileEvent]:
   from tinygrad.renderer.amd.sqtt import map_insts, InstructionInfo, PacketType, INST, InstOp, VALUINST, IMMEDIATE, IMMEDIATE_MASK, VMEMEXEC, ALUEXEC
-  from tinygrad.renderer.amd.sqtt import INST_RDNA4, InstOpRDNA4, TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4, CDNA_INST, InstOpCDNA
+  from tinygrad.renderer.amd.sqtt import INST_RDNA4, InstOpRDNA4, TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4
   ret:list[ProfileEvent] = []
   row_ends:dict[str, Decimal] = {}
-  curr_barrier:dict[str, ProfileRangeEvent] = {}
   NS_PER_TICK = 10  # 100MHz
   prev_pair:tuple[int, int]|None = None # (shader, realtime)
-  is_cdna = target.startswith("gfx9")
-  def add(name:str, p:PacketType, op:str|None=None, wave:int|None=None, info:InstructionInfo|None=None) -> None:
+  def add(name:str, p:PacketType, width=1, op:str|None=None, wave:int|None=None, info:InstructionInfo|None=None) -> None:
     row = f"WAVE:{wave}" if (wave:=getattr(p, "wave", wave)) is not None else f"{p.__class__.__name__}:0 {name}"
-    # barrier on this row extends to fill the time our wave was waiting
-    if (barrier:=curr_barrier.pop(row, None)) is not None: barrier.en = Decimal(p._time)
-    ret.append(e:=ProfileRangeEvent(row, TracingKey(op or name, ret=f"PC:{info.pc}" if info else None), Decimal(p._time), Decimal(p._time+1)))
-    # allow CDNA packets to overlap, NOT allowed on RDNA.
-    if (et:=row_ends.get(row)) is not None and e.st < et and not is_cdna: raise RuntimeError(f"packet {p} overlaps another packet in {row}.")
+    ret.append(e:=ProfileRangeEvent(row, TracingKey(op or name, ret=f"PC:{info.pc}" if info else None), Decimal(p._time), Decimal(p._time+width)))
+    if (et:=row_ends.get(row)) is not None and e.st < et: raise RuntimeError(f"packet {p} overlaps another packet in {row}.")
     row_ends[row] = unwrap(e.en)
-    if name == "BARRIER": curr_barrier[row] = e
   for p, info in map_insts(data, lib, target):
     if len(ret) > getenv("MAX_SQTT_PKTS", 50_000): break
     if isinstance(p, (TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4)) and p.is_marker:
       pair = (p._time, p.delta)
       if prev_pair is None: prev_pair = pair
-      else:
+      elif ret:
         (s0, r0), (s1, r1) = prev_pair, pair
         freq_hz = (s1 - s0) * 1_000_000_000 // ((r1 - r0) * NS_PER_TICK)
         ret.append(ProfilePointEvent("LINE:Shader Clock", "freq_hz", freq_hz, ts=Decimal(p._time)))
         prev_pair = pair
-    if isinstance(p, (INST, INST_RDNA4, CDNA_INST)):
-      name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4, InstOpCDNA)) else f"0x{p.op:02x}"
-      add(name, p, info=info)
+    if isinstance(p, (INST, INST_RDNA4)):
+      name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4)) else f"0x{p.op:02x}"
+      add(name, p, width=10 if "BARRIER" in name else 1, info=info)
     if isinstance(p, (VALUINST, IMMEDIATE)): add(p.__class__.__name__, p, info=info)
     if isinstance(p, IMMEDIATE_MASK): add("IMMEDIATE", p, wave=unwrap(info).wave, info=info)
     if isinstance(p, (VMEMEXEC, ALUEXEC)):
@@ -531,7 +525,11 @@ def amdgpu_cfg(lib:bytes, target:str) -> dict:
   curr:int|None = None
   blocks:dict[int, list[int]] = {}
   paths:dict[int, dict[int, int]] = {}
+  lines:list[str] = []
+  disasm = {pc:str(inst) for pc,inst in pc_table.items()}
+  asm_width = max(len(asm) for asm in disasm.values())
   for pc, inst in pc_table.items():
+    lines.append(f"  {disasm[pc]:<{asm_width}}  // {pc:012X}")
     if pc in leaders:
       paths[curr:=pc] = {}
       blocks[pc] = []
@@ -552,9 +550,6 @@ def amdgpu_cfg(lib:bytes, target:str) -> dict:
       elif name in {"op","opx","opy"}: tokens.append({"st":(op_name:=val.name.lower()), "keys":[op_name], "kind":0})
       elif name != "encoding" and val != field.default: tokens.append({"st":(s:=repr(val)), "keys":[s], "kind":1})
   # show a smaller view for repeated instructions in the graph
-  lines:list[str] = []
-  disasm = {pc:str(inst) for pc,inst in pc_table.items()}
-  asm_width = max(len(asm) for asm in disasm.values())
   for pcs in blocks.values():
     new_pcs:list[int] = []
     i, n = 0, len(pcs)
@@ -565,13 +560,12 @@ def amdgpu_cfg(lib:bytes, target:str) -> dict:
       if j-i>1:
         pc_tokens[pcs[i]].append({"st":f"({j-i}x)", "keys":[], "kind":0})
         for k in range(i+1, j): del pc_tokens[pcs[k]]
-      lines.append(f"{disasm[pcs[i]]:<{asm_width}}  # {pcs[i]:012X}"+(f"...{pcs[j-1]:012X} ({j-i}x)" if j-i>1 else ""))
       i = j
     pcs[:] = new_pcs
   from tinygrad.runtime.autogen import amdgpu_kd
   kd = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytearray(get_elf_section(lib, ".rodata").content))
   vgpr_gran = kd.compute_pgm_rsrc1 & amdgpu_kd.COMPUTE_PGM_RSRC1_GRANULATED_WORKITEM_VGPR_COUNT
-  return {"data":{"blocks":blocks, "paths":paths, "pc_tokens":pc_tokens}, "src":"\n".join(lines), "lang":"python",
+  return {"data":{"blocks":blocks, "paths":paths, "pc_tokens":pc_tokens}, "src":"\n".join(lines),
           "metadata":[[{"label":f"{r} Alloc", "value":v} for r,v in [("VGPR", (vgpr_gran+1)*8-7), ("LDS", kd.group_segment_fixed_size),
                                                                      ("Scratch", kd.private_segment_fixed_size)] if v>0]]}