Merge branch 'master' into call_inline

better cl compiler name (#14660 )
cl_compiler instead of compiler because overriding Compiled.compiler seems more confusing
2026-02-10 14:12:08 +08:00 · 2026-02-10 01:03:46 -05:00 · 2026-02-10 14:02:43 +08:00 · 2026-02-10 13:35:17 +08:00 · 2026-02-10 13:25:29 +08:00 · 2026-02-09 23:45:24 -05:00
90 changed files with 4231 additions and 1804 deletions
@@ -525,6 +525,8 @@ jobs:
      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training steps w 6 GPUS
      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Test full tinyfs load
+      run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '15'
+  CACHE_VERSION: '16'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
@@ -581,7 +581,7 @@ jobs:
        load: true
        tags: qemu-hexagon:latest
        cache-from: type=gha
-        cache-to: type=gha,mode=min
+        cache-to: ${{ github.event_name != 'pull_request' && 'type=gha,mode=min' || '' }}
    - name: Set MOCKDSP env
      run: printf "MOCKDSP=1" >> $GITHUB_ENV
    - name: Run test_tiny on DSP
@@ -739,7 +739,7 @@ jobs:
          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
      - name: Run pytest (cuda)
        # skip multitensor because it's slow
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --ignore test/test_multitensor.py --durations=20
      - name: Run TestOps.test_add with PMA
        run: VIZ=-1 PMA=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
      - name: Run process replay tests
@@ -2,6 +2,7 @@

 export PYTHONPATH="."
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000

@@ -12,7 +13,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=8 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=1
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

 export MODEL="llama3"
@@ -20,13 +21,14 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="2.5e-4" END_LR="2.5e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}

-export SEED=5760
+export SEED=${SEED:-5760}

-export JITBEAM=3
+export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

 export FAKEDATA=1 BENCHMARK=10 LLAMA_LAYERS=2
@@ -13,7 +13,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

 export MODEL="llama3"
@@ -21,9 +21,10 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="2.5e-4" END_LR="2.5e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}

 export SEED=${SEED:-5760}

@@ -0,0 +1,6 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+export VIZ=${VIZ:--1}
+examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
@@ -92,7 +92,7 @@ class SMICtx:
    self.prev_terminal_width = 0
    self.prev_terminal_height = 0

-    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
+    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
    lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
    self.lspci = {l.split()[0]: l.split(" ", 1)[1] for l in lspci}
    for k,v in self.lspci.items():
@@ -153,7 +153,8 @@ class SMICtx:
    tables = {}
    for dev in self.devs:
      match dev.ip_ver[am.MP1_HWIP]:
-        case (13,0,6)|(13,0,12): table_t = dev.smu.smu_mod.MetricsTableX_t
+        case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableV0_t
+        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTableV2_t
        case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
      tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
    return tables
@@ -279,7 +280,7 @@ class SMICtx:
      device_line = [f"{bold(dev.pcibus)} {trim(self.lspci[dev.pcibus[5:]], col_size - 20)}"] + [pad("", col_size)]
      activity_line = [f"GFX Activity {draw_bar(self.get_gfx_activity(dev, metrics) / 100, activity_line_width)}"] \
                    + [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
-                    + [f"MEM Usage    {draw_bar((mem_used / mem_total) / 100, activity_line_width, opt_text=mem_fmt)}"] \
+                    + [f"MEM Usage    {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \

      temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
      temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]
@@ -1,12 +1,18 @@
 #!/usr/bin/env python3

+import os
 from tinygrad.helpers import Context
 from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
+from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev

 if __name__ == "__main__":
  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
-  pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
+  for gpu in gpus:
+    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, bars=[0, 2, 5]) for gpu in gpus]
  amdevs = []
  with Context(DEBUG=2):
    for pcidev in pcidevs:
@@ -28,7 +28,7 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
  return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))

 def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
-  _, input_uop, weight_uop = kernel.src
+  _, input_uop, weight_uop = kernel.src[1:]
  input_tensor = Tensor(input_uop, device=input_uop.device)
  grad_tensor = Tensor(gradient, device=gradient.device)
  weight_tensor = Tensor(weight_uop, device=weight_uop.device)
@@ -14,7 +14,12 @@ GEMM_ARGS = {
  (8192, 8192, 8192): (256, 128, 131072),
  (4096, 4096, 4096): (256, 64, 16384),
  (4096, 14336, 4096): (256, 64, 57344),
+  (4096, 14336, 8192): (256, 128, 114688),
  (4096, 4096, 14336): (256, 224, 57344),
+  (14336, 4096, 8192): (256, 128, 114688),
+  (4096, 8192, 14336): (256, 224, 114688),
+  (4096, 4096, 8192): (256, 128, 32768),
+  (4096, 8192, 4096): (256, 64, 32768),
 }
 ITERS_ARGS = {64: (67108864, 0), 128: (33554432, 0), 224: (613566757, 2147483656)}

@@ -30,13 +30,13 @@ atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters
 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}")
-  # only sharding on the batch is tested, others might work too
-  if isinstance(a.device, tuple) and not (a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None):
-    return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
  N = b.shape[1]
+  # only sharding on the batch or K is tested, others might work too
  if isinstance(a.device, tuple):
-    batch //= len(a.device)
+    if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= 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", "")
@@ -62,9 +62,11 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
 # ** backward gemm, might use the asm gemm

 def custom_gemm_bw(gradient:UOp, kernel:UOp):
-  out, a, b = kernel.src
+  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)
@@ -80,9 +82,10 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
  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 is_multi:
-    out = Tensor(Tensor.empty(batch//len(a.device), M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+    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)

@@ -93,4 +96,5 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=numWG, arch=arch), 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)
  return out.squeeze(0) if squeeze else out
@@ -89,13 +89,13 @@ class Attention:
      assert start_pos == 0
      keys, values = xk, xv

-    if Tensor.training:
-      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
-      attn = xq.scaled_dot_product_attention(keys, values, is_causal=True, enable_gqa=True).transpose(1, 2)
-    else:
+    if self.max_context:
      keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
      attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
+    else:
+      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(keys, values, is_causal=True, enable_gqa=True).transpose(1, 2)
    if getenv("STUB_ATTENTION"):
      from tinygrad.uop.ops import UOp, KernelInfo
      def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
@@ -103,9 +103,9 @@ class Attention:
      def fa_custom_backward(out_q:UOp, out_k:UOp, out_v:UOp, grad:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
        return UOp.sink(arg=KernelInfo(name="fa_custom_backward"))
      def fa_backward(grad:UOp, kernel:UOp) -> tuple[None, UOp, UOp, UOp]:
-        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[1]))
-        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[2]))
-        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[3]))
+        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[2]))
+        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[3]))
+        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[4]))
        ck = Tensor.custom_kernel(grad_q, grad_k, grad_v, Tensor(grad), q, k, v, fxn=fa_custom_backward)[:3]
        return (None, ck[0].uop, ck[1].uop, ck[2].uop)
      attn = Tensor.empty_like(attn).custom_kernel(xq, keys, values, fxn=fa_custom_forward, grad_fxn=fa_backward)[0]
@@ -203,7 +203,7 @@ class Transformer:
    h = self.tok_embeddings(tokens)
    freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :]

-    if not Tensor.training and seqlen > 1:
+    if self.max_context != 0 and seqlen > 1:
      mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1)
    else: mask = None
    for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
@@ -2,27 +2,26 @@ import sys, pickle, decimal, json
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.helpers import tqdm, temp, ProfileEvent, ProfileRangeEvent, TracingKey

-devices:dict[str, tuple[decimal.Decimal, decimal.Decimal, int]] = {}
-def prep_ts(device:str, ts:decimal.Decimal, is_copy): return int(decimal.Decimal(ts) + devices[device][is_copy])
-def dev_to_pid(device:str, is_copy=False): return {"pid": devices[device][2], "tid": int(is_copy)}
+devices:dict[str, tuple[decimal.Decimal, int]] = {}
+def prep_ts(device:str, ts:decimal.Decimal): return int(decimal.Decimal(ts) + devices[device][0])
+def dev_to_pid(device:str): return {"pid": devices[device][1], "tid": 0}
 def dev_ev_to_perfetto_json(ev:ProfileDeviceEvent):
-  devices[ev.device] = (ev.comp_tdiff, ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff, len(devices))
+  devices[ev.device] = (ev.tdiff, len(devices))
  return [{"name": "process_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "args": {"name": ev.device}},
-          {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 0, "args": {"name": "COMPUTE"}},
-          {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 1, "args": {"name": "COPY"}}]
+          {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 0, "args": {"name": ev.device}}]
 def range_ev_to_perfetto_json(ev:ProfileRangeEvent):
  name = ev.name.display_name if isinstance(ev.name, TracingKey) else ev.name
-  return [{"name": name, "ph": "X", "ts": prep_ts(ev.device, ev.st, ev.is_copy), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device, ev.is_copy)}]
+  return [{"name": name, "ph": "X", "ts": prep_ts(ev.device, ev.st), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device)}]
 def graph_ev_to_perfetto_json(ev:ProfileGraphEvent, reccnt):
  ret = []
  for i,e in enumerate(ev.ents):
    st, en = ev.sigs[e.st_id], ev.sigs[e.en_id]
    name = e.name.display_name if isinstance(e.name, TracingKey) else e.name
-    ret += [{"name": name, "ph": "X", "ts": prep_ts(e.device, st, e.is_copy), "dur": float(en-st), **dev_to_pid(e.device, e.is_copy)}]
+    ret += [{"name": name, "ph": "X", "ts": prep_ts(e.device, st), "dur": float(en-st), **dev_to_pid(e.device)}]
    for dep in ev.deps[i]:
      d = ev.ents[dep]
-      ret += [{"ph": "s", **dev_to_pid(d.device, d.is_copy), "id": reccnt+len(ret), "ts": prep_ts(d.device, ev.sigs[d.en_id], d.is_copy), "bp": "e"}]
-      ret += [{"ph": "f", **dev_to_pid(e.device, e.is_copy), "id": reccnt+len(ret)-1, "ts": prep_ts(e.device, st, e.is_copy), "bp": "e"}]
+      ret += [{"ph": "s", **dev_to_pid(d.device), "id": reccnt+len(ret), "ts": prep_ts(d.device, ev.sigs[d.en_id]), "bp": "e"}]
+      ret += [{"ph": "f", **dev_to_pid(e.device), "id": reccnt+len(ret)-1, "ts": prep_ts(e.device, st), "bp": "e"}]
  return ret
 def to_perfetto(profile:list[ProfileEvent]):
  # Start json with devices.
@@ -145,7 +145,8 @@ class RGP:
  @staticmethod
  def from_profile(profile_pickled, device:str|None=None):
    profile: list[ProfileEvent] = pickle.loads(profile_pickled)
-    device_events = {x.device:x for x in profile if isinstance(x, ProfileDeviceEvent) and x.device.startswith('AMD')}
+    def _is_base_dev(d): return all(p.isdigit() for p in d.split(":")[1:])
+    device_events = {x.device:x for x in profile if isinstance(x, ProfileDeviceEvent) and x.device.startswith('AMD') and _is_base_dev(x.device)}
    if device is None:
      if len(device_events) == 0: raise RuntimeError('No supported devices found in profile')
      if len(device_events) > 1: raise RuntimeError(f"More than one supported device found, select which one to export: {', '.join(device_events.keys())}")
@@ -1,11 +1,36 @@
 from tinygrad.tensor import Tensor
-import argparse
+import argparse, math, hashlib
+
+def _python_hash_1mb(data:bytes|bytearray):
+  chunks = [data[i:i+4096] for i in range(0, len(data), 4096)]
+  chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
+  return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
+
+def hash_file(data: bytes|bytearray):
+  if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)
+  base_chunks = math.ceil(len(data) / Tensor.CHUNK_SIZE)
+  tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
+
+  for _ in range(tree_depth + 1):
+    data_chunks = [data[i:i+Tensor.CHUNK_SIZE] for i in range(0, len(data), Tensor.CHUNK_SIZE)]
+    data_chunk_hashes = [_python_hash_1mb(chunk) for chunk in data_chunks]
+    data = b''.join(data_chunk_hashes)
+    if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)
+
+  return data[:16]

 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument("--hash", type=str, required=True, help="file hash to fetch")
  parser.add_argument("--len", type=int, required=True, help="file length to fetch")
  parser.add_argument("--dest", type=str, required=True, help="destination path to save the file")
+  parser.add_argument("--check", action="store_true", help="verify the file hash after fetching")
  args = parser.parse_args()

  Tensor(bytes.fromhex(args.hash), device="CPU").fs_load(args.len).to(f"disk:{args.dest}").realize()
+
+  if args.check:
+    with open(args.dest, "rb") as f:
+      data = f.read()
+      assert hash_file(data) == bytes.fromhex(args.hash), "Hash mismatch after fetching file"
+    print("File hash verified successfully!")
@@ -0,0 +1,18 @@
+A command line tool for exploring the VIZ trace.
+
+After running with VIZ=-1, use `PYTHONPATH=. extra/viz/cli.py` to explore the saved trace files.
+
+## Inspect runtime profiling
+
+Use `PYTHONPATH=. extra/viz/cli.py --profile` to list all traced devices.
+
+List top slowest kernels on a device: `--profile --device "AMD"`
+List samples of a kernel on a device: `--profile --device "AMD" --kernel E_3`
+
+## Inspect codegen and PatternMatcher
+
+Use `PYTHONPATH=. extra/viz/cli.py --rewrites` to list all traced kernels.
+
+List all codegen steps for a kernel: `--rewrites --kernel E_3`
+Get source code: `--rewrites --kernel E_3 --select "View Program"`
+Inspect a graph rewrite: `--rewrites --kernel E_3 --select "initial symbolic"`
@@ -26,16 +26,25 @@ def print_data(data:dict) -> None:

 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
-  parser.add_argument('--kernel', type=str, default=None, metavar="NAME", help='Select a kernel by name (optional name, default: only list names)')
-  parser.add_argument('--select', type=str, default=None, metavar="NAME",
-                      help='Rewrites: Select an item within the chosen kernel (optional name, default: only list names)')
-  parser.add_argument('--profile', action="store_true", help="View profiling trace (default: views rewrites)")
-  parser.add_argument('--device', type=str, default=None, metavar="NAME", help="Profile only: Select a device (default: prints all devices)")
-  parser.add_argument('--profile-path', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
-                      default=pathlib.Path(temp("profile.pkl", append_user=True)))
-  parser.add_argument('--rewrites-path', type=pathlib.Path, metavar="PATH", help='Path to rewrites (optional file, default: latest rewrites)',
-                      default=pathlib.Path(temp("rewrites.pkl", append_user=True)))
+  g_mode = parser.add_argument_group("mode")
+  g_mode.add_argument("--profile", action="store_true", help="View profile trace")
+  g_mode.add_argument("--rewrites", action="store_true", help="View rewrites trace")
+  g_profile = parser.add_argument_group("profile options")
+  g_profile.add_argument("--device", type=str, default=None, metavar="NAME", help="Select a device (optional name, default: only list names)")
+  g_profile.add_argument("--top", type=int, default=10, metavar="N", help="Number of top kernels to show (-1 for all, default: 10)")
+  g_rewrites = parser.add_argument_group("rewrites options")
+  g_rewrites.add_argument("--select", type=str, default=None, metavar="NAME",
+                          help="Select an item within the chosen kernel (optional name, default: only list names)")
+  g_common = parser.add_argument_group("common options")
+  g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
+  parser.add_argument("--profile-path", type=pathlib.Path, metavar="PATH", help="Path to profile (optional file, default: latest profile)",
+                        default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  parser.add_argument("--rewrites-path", type=pathlib.Path, metavar="PATH", help="Path to rewrites (optional file, default: latest rewrites)",
+                        default=pathlib.Path(temp("rewrites.pkl", append_user=True)))
  args = parser.parse_args()
+  if not args.profile and not args.rewrites:
+    parser.print_help()
+    exit(0)

  viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
  viz.ctxs = viz.get_rewrites(viz.trace)
@@ -44,9 +53,10 @@ if __name__ == "__main__":
    from tabulate import tabulate
    profile = load_profile(viz.load_pickle(args.profile_path, default=[]))
    agg, total, n = {}, 0, 0
+    if args.device is None: print("Select a device:")
    for k,v in profile["layout"].items():
      if not optional_eq({"name":k}, args.device): continue
-      print(k)
+      print(f"  {k}")
      if args.device is None: continue
      for e in v.get("events", []):
        et = e["dur"]*1e-6
@@ -61,11 +71,14 @@ if __name__ == "__main__":
          a[0] += et
          a[1] += 1
          total += et
-
-    if agg:
-      rows = [[n, t, time_to_str(t, w=9), t / c if c else 0.0, c, (t / total * 100.0) if total else 0.0] for n, (t, c) in agg.items()]
-      rows.sort(key=lambda r: r[1], reverse=True)
-      print(tabulate([[r[0], r[2], r[4], f"{r[5]:.2f}%"] for r in rows[:30]], headers=["name", "total", "count", "pct"], tablefmt="github"))
+    if agg and total > 0:
+      items = sorted(agg.items(), key=lambda kv:kv[1][0], reverse=True)
+      sel = items if args.top == -1 else items[:args.top]
+      table = [[name, time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in sel]
+      if args.top != -1 and (other:=items[len(sel):]):
+        other_t = total-sum(t for _, (t, _) in sel)
+        table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"])
+      print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
    exit(0)

  for k in viz.ctxs:
@@ -67,7 +67,7 @@ testing_minimal = [
  "pytest-timeout",
  "pytest-split",
  "hypothesis>=6.148.9",
-  "z3-solver",
+  "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
 ]
 testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "ggml-python"]
 testing = [
@@ -1,5 +1,5 @@
 import unittest
-from tinygrad.device import CompileError, Device, Compiler
+from tinygrad.device import CompileError, Device
 if Device.DEFAULT=="METAL":
  from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
@unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
@@ -48,28 +48,4 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr
 """)
    with self.assertRaises(RuntimeError):
      compiled = compiled[:40] # corrupt the compiled program
-      MetalProgram(device, "r_5", compiled)
-
-  def test_program_w_empty_compiler(self):
-    device = MetalDevice("metal")
-    compiler = Compiler(device)
-    compiled = compiler.compile("""
-#include <metal_stdlib>
-kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]){
-  data0[0] = 0;
-}
-""")
-    MetalProgram(device, "r_5", compiled)
-
-  def test_bad_program_w_empty_compiler(self):
-    device = MetalDevice("metal")
-    compiler = Compiler(device)
-    # this does not raise
-    compiled = compiler.compile("""
-#include <metal_stdlib>
-kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]){
-  invalid codes;
-}
-""")
-    with self.assertRaises(RuntimeError):
      MetalProgram(device, "r_5", compiled)
@@ -1,9 +1,19 @@
 import unittest
+from unittest.mock import patch
 from tinygrad import Device
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes
 from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler, CLProgram

+@unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
+class TestCLCompileCache(unittest.TestCase):
+  def test_compile_cached(self):
+    device = Device[Device.DEFAULT]
+    src = "__kernel void cached_test(__global int* a) { a[0] = 1; }"
+    CLProgram(device, name="cached_test", lib=src.encode())
+    with patch.object(CLCompiler, 'compile', side_effect=RuntimeError("compile should not be called on cache hit")):
+      CLProgram(device, name="cached_test", lib=src.encode())
+
@unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
 class TestCLError(unittest.TestCase):
  @unittest.skip("allocates tons of memory")
@@ -17,7 +27,7 @@ class TestCLError(unittest.TestCase):
  def test_invalid_kernel_name(self):
    device = Device[Device.DEFAULT]
    with self.assertRaises(RuntimeError) as err:
-      CLProgram(device, name="", lib=CLCompiler(device, "test").compile("__kernel void test(__global int* a) { a[0] = 1; }"))
+      CLProgram(device, name="", lib="__kernel void test(__global int* a) { a[0] = 1; }".encode())
    assert str(err.exception) == "OpenCL Error -46: CL_INVALID_KERNEL_NAME"

  def test_unaligned_copy(self):
@@ -1,12 +1,16 @@
-import random, operator
+# NOTE: z3-solver 4.15.4 segfaults (exit code 139) when creating many z3.Context() with complex expressions.
+# Reproduces consistently with seed=74 around iteration 1767. Versions <=4.15.3 are fine.
+# Workaround: reuse a single z3.Context, or pin z3-solver<4.15.4 (see pyproject.toml).
+# To repro: pip install z3-solver==4.15.4.0 && python test/external/fuzz_symbolic.py 74
+import random, operator, sys
 import z3
 from tinygrad import Variable, dtypes
 from tinygrad.uop.ops import UOp
 from tinygrad.uop.validate import uops_to_z3
 from tinygrad.helpers import DEBUG, Context

-seed = random.randint(0, 100)
-print(f"Seed: {seed}")
+seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
+print(f"Seed: {seed}", flush=True)
 random.seed(seed)

 unary_ops = [lambda a:a+random.randint(-4, 4), lambda a: a*random.randint(-4, 4),
@@ -1,11 +1,11 @@
-import random
+import random, sys
 import z3
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.uop.validate import uops_to_z3
 from tinygrad.helpers import DEBUG, Context, colored

-seed = random.randint(0, 100)
-print(f"Seed: {seed}")
+seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
+print(f"Seed: {seed}", flush=True)
 random.seed(seed)

 def get_random_term(ranges, factors):
@@ -47,13 +47,15 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str,
  UOp.unique_num = itertools.count(max([u.arg for u in big_sink.toposort() if u.op is Ops.UNIQUE], default=0)+1)
  new_sink = big_sink.substitute(get_rangeify_map(big_sink))
  def to_str(ret:UOp) -> str:
-    asts = [repr(u.arg.ast) for u in ret.toposort() if u.op is Ops.KERNEL]
+    asts = [repr(u.arg.ast) for u in ret.toposort() if u.op is Ops.CALL]
    return "\n".join([f"{len(asts)} kernels", *asts])
  return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)

 def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
  # the ast.arg is non None if we are inside of search.py
-  sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p.applied_opts if BEAM>=1 else None)
+  sink_arg = ast.arg or KernelInfo()
+  if opts is not None: sink_arg = replace(sink_arg, opts_to_apply=tuple(opts))
+  elif BEAM >= 1 and sink_arg.opts_to_apply is None: sink_arg = replace(sink_arg, opts_to_apply=p.applied_opts)
  input_ast = ast if ast.op is Ops.PROGRAM else ast.replace(arg=replace(sink_arg, name=p.name))
  p2 = get_program(input_ast, renderer=renderer)
  def to_str(ret:ProgramSpec) -> str:
@@ -130,15 +130,18 @@ class NVDriver(VirtDriver):
      struct.hObjectNew = self._alloc_handle()
      self.object_by_handle[struct.hObjectNew] = NVContextShare(self.object_by_handle[struct.hObjectParent])
    elif struct.hClass == nv_gpu.AMPERE_CHANNEL_GPFIFO_A:
-      assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVChannelGroup)
+      parent = self.object_by_handle.get(struct.hObjectParent)
+      assert parent is not None and isinstance(parent, (NVChannelGroup, NVGPU))
      struct.hObjectNew = self._alloc_handle()
      params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS.from_address(params_ptr)
-      gpu = self.object_by_handle[struct.hObjectParent].device
+      gpu = parent.device if isinstance(parent, NVChannelGroup) else parent
      gpfifo_token = gpu.add_gpfifo(params.gpFifoOffset, params.gpFifoEntries)
      self.object_by_handle[struct.hObjectNew] = NVGPFIFO(gpu, gpfifo_token)
-    elif struct.hClass == nv_gpu.AMPERE_DMA_COPY_B or struct.hClass == nv_gpu.ADA_COMPUTE_A:
+    elif struct.hClass in (nv_gpu.AMPERE_DMA_COPY_B, nv_gpu.ADA_COMPUTE_A, nv_gpu.NVC9B0_VIDEO_DECODER, nv_gpu.NVCFB0_VIDEO_DECODER):
      assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPFIFO)
      struct.hObjectNew = self._alloc_handle()
+      gpfifo = self.object_by_handle[struct.hObjectParent]
+      gpfifo.device.queues[gpfifo.token].bound_engines.add(struct.hClass)
    elif struct.hClass == nv_gpu.GT200_DEBUGGER:
      struct.hObjectNew = self._alloc_handle()
    elif struct.hClass == nv_gpu.MAXWELL_PROFILER_DEVICE:
@@ -194,7 +197,7 @@ class NVDriver(VirtDriver):
      params = nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS.from_address(params_ptr)
      gpu_fifo = self.object_by_handle[struct.hObject]
      params.workSubmitToken = gpu_fifo.token
-    elif struct.cmd == nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE: pass
+    elif struct.cmd in (nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE, nv_gpu.NVA06F_CTRL_CMD_BIND, nv_gpu.NVA06F_CTRL_CMD_GPFIFO_SCHEDULE): pass
    elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_PERF_BOOST: pass
    elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_FB_FLUSH_GPU_CACHE: pass
    elif struct.cmd == nv_gpu.NV83DE_CTRL_CMD_DEBUG_READ_ALL_SM_ERROR_STATES:
@@ -26,6 +26,7 @@ class GPFIFO:
    self.gpfifo = to_mv(self.base, self.entries_cnt * 8).cast("Q")
    self.ctrl = nv_gpu.AmpereAControlGPFifo.from_address(self.base + self.entries_cnt * 8)
    self.state = {}
+    self.bound_engines: set[int] = set()

    # Buf exec state
    self.buf = None
@@ -115,9 +116,11 @@ class GPFIFO:
  def execute_cmd(self, cmd) -> SchedResult:
    if cmd == nv_gpu.NVC56F_SEM_EXECUTE: return self._exec_signal()
    elif cmd == nv_gpu.NVC6C0_LAUNCH_DMA: return self._exec_nvc6c0_dma()
-    elif cmd == nv_gpu.NVC6B5_LAUNCH_DMA: return self._exec_nvc6b5_dma()
+    elif cmd == nv_gpu.NVC6B5_LAUNCH_DMA: # NOTE: NVC6B5_LAUNCH_DMA == NVC9B0_EXECUTE == 0x300
+      return self._exec_vid_decode() if self.bound_engines & {nv_gpu.NVC9B0_VIDEO_DECODER, nv_gpu.NVCFB0_VIDEO_DECODER} else self._exec_nvc6b5_dma()
    elif cmd == nv_gpu.NVC6C0_SEND_SIGNALING_PCAS2_B: return self._exec_pcas2()
    elif cmd == 0x0320: return self._exec_load_inline_qmd() # NVC6C0_LOAD_INLINE_QMD_DATA
+    elif cmd == nv_gpu.NVC9B0_SEMAPHORE_D: return self._exec_vid_semaphore()
    else: self.state[cmd] = self._next_dword() # just state update
    return SchedResult.CONT

@@ -136,6 +139,27 @@ class GPFIFO:
    else: raise RuntimeError(f"Unsupported type={typ} in exec wait/signal")
    return SchedResult.CONT

+  def _exec_vid_decode(self) -> SchedResult:
+    self._next_dword() # consume execute flags
+    # validate that all required decode state was set up correctly
+    assert self._state(nv_gpu.NVC9B0_SET_APPLICATION_ID) == nv_gpu.NVC9B0_SET_APPLICATION_ID_ID_HEVC
+    pic_desc_addr = self._state(nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET) << 8
+    pic = nv_gpu.nvdec_hevc_pic_s.from_address(pic_desc_addr)
+    assert pic.stream_len > 0 and pic.pic_width_in_luma_samples > 0 and pic.pic_height_in_luma_samples > 0
+    assert self._state(nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET) != 0
+    assert self._state(nv_gpu.NVC9B0_SET_COLOC_DATA_OFFSET) != 0
+    assert self._state(nv_gpu.NVC9B0_SET_NVDEC_STATUS_OFFSET) != 0
+    assert self._state(nv_gpu.NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET) != 0
+    return SchedResult.CONT
+
+  def _exec_vid_semaphore(self) -> SchedResult:
+    signal = self._state64(nv_gpu.NVC9B0_SEMAPHORE_A)
+    val = self._state(nv_gpu.NVC9B0_SEMAPHORE_C)
+    self._next_dword() # flags
+    to_mv(signal, 8).cast('Q')[0] = val
+    to_mv(signal + 8, 8).cast('Q')[0] = int(time.perf_counter() * 1e9)
+    return SchedResult.CONT
+
  def _exec_load_inline_qmd(self):
    qmd_addr = self._state64(nv_gpu.NVC6C0_SET_INLINE_QMD_ADDRESS_A) << 8
    assert qmd_addr != 0x0, f"invalid qmd address {qmd_addr}"
@@ -102,8 +102,8 @@ class TestCompiler(unittest.TestCase):

 class TestRunAsModule(unittest.TestCase):
  def test_module_runs(self):
-    out = '\n'.join(enumerate_devices_str())
-    self.assertIn("CPU", out) # for sanity check
+    cpu_line = [l for l in enumerate_devices_str() if "CPU" in l][0]
+    self.assertIn("PASS", cpu_line, f"expected CPU to PASS, got: {cpu_line}")

 if __name__ == "__main__":
  unittest.main()
@@ -72,7 +72,6 @@ class TestGC(unittest.TestCase):
    ys = y.schedule()
    del x
    run_schedule(ys)
-    np.testing.assert_equal(y.numpy(), np.full((256,), 2))
    self.assertEqual(bufs_allocated()-init, 1)
    del y
    self.assertEqual(bufs_allocated()-init, 0)
@@ -0,0 +1,34 @@
+import unittest
+from tinygrad import Tensor, Device
+from tinygrad.engine.realize import get_program
+from tinygrad.codegen.opt import Opt, OptOps
+from test.external.process_replay.process_replay import replay_get_program
+
+N = 16
+class TestProcessReplay(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1].ast
+    cls.renderer = Device[Device.DEFAULT].renderer
+
+  def test_replay_no_opts(self):
+    # opts=None means use default heuristic path
+    p = get_program(self.ast, self.renderer)
+    good, compare, _ = replay_get_program(p, self.ast, self.renderer)
+    self.assertEqual(good, compare)
+
+  def test_replay_empty_opts(self):
+    # opts=[] means explicitly apply zero opts (unoptimized)
+    p = get_program(self.ast, self.renderer, opts=[])
+    good, compare, _ = replay_get_program(p, self.ast, self.renderer, opts=[])
+    self.assertEqual(good, compare)
+
+  def test_replay_with_opt(self):
+    # opts=[Opt(...)] means apply a specific opt
+    opts = [Opt(OptOps.UPCAST, 0, 4)]
+    p = get_program(self.ast, self.renderer, opts=opts)
+    good, compare, _ = replay_get_program(p, self.ast, self.renderer, opts=opts)
+    self.assertEqual(good, compare)
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)
@@ -365,8 +365,8 @@ def load_profile(lst:list[ProfileEvent]) -> dict:

 class TestVizProfiler(BaseTestViz):
  def test_node(self):
-    prof = [ProfileRangeEvent(device='NV', name='E_2', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=False),
-            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100))]
+    prof = [ProfileRangeEvent(device='NV', name='E_2', st=decimal.Decimal(1000), en=decimal.Decimal(1010)),
+            ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000))]

    j = load_profile(prof)

@@ -379,28 +379,28 @@ class TestVizProfiler(BaseTestViz):
    assert event['ref'] is None

  def test_copy_node(self):
-    prof = [ProfileRangeEvent(device='NV', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
-            ProfileRangeEvent(device='NV:2', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010), is_copy=True),
-            ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
-            ProfileDeviceEvent(device='NV:2', comp_tdiff=decimal.Decimal(-800), copy_tdiff=decimal.Decimal(-80))]
+    prof = [ProfileRangeEvent(device='NV:SDMA:0', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010)),
+            ProfileRangeEvent(device='NV:2:SDMA:0', name='COPYxx', st=decimal.Decimal(1000), en=decimal.Decimal(1010)),
+            ProfileDeviceEvent(device='NV:SDMA:0', tdiff=decimal.Decimal(-100)),
+            ProfileDeviceEvent(device='NV:2:SDMA:0', tdiff=decimal.Decimal(-80))]

    j = load_profile(prof)

-    event = j['layout']['NV']['events'][0]
+    event = j['layout']['NV:SDMA:0']['events'][0]
    self.assertEqual(event['name'], 'COPYxx')
    self.assertEqual(event['st'], 0)   # first event
    self.assertEqual(event['dur'], 10)

-    event2 = j['layout']['NV:2']['events'][0]
+    event2 = j['layout']['NV:2:SDMA:0']['events'][0]
    self.assertEqual(event2['st'], 20) # second event, diff clock

    self.assertEqual(j["dur"], (event2["st"]+event2["dur"])-event["st"])

  def test_graph(self):
-    prof = [ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
-            ProfileDeviceEvent(device='NV:1', comp_tdiff=decimal.Decimal(-500), copy_tdiff=decimal.Decimal(-50)),
-            ProfileGraphEvent(ents=[ProfileGraphEntry(device='NV', name='E_25_4n2', st_id=0, en_id=1, is_copy=False),
-                                    ProfileGraphEntry(device='NV:1', name='NV -> NV:1', st_id=2, en_id=3, is_copy=True)],
+    prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)),
+            ProfileDeviceEvent(device='NV:1:SDMA:0', tdiff=decimal.Decimal(-50)),
+            ProfileGraphEvent(ents=[ProfileGraphEntry(device='NV', name='E_25_4n2', st_id=0, en_id=1),
+                                    ProfileGraphEntry(device='NV:1:SDMA:0', name='NV -> NV:1', st_id=2, en_id=3)],
                              deps=[[], [0]],
                              sigs=[decimal.Decimal(1000), decimal.Decimal(1002), decimal.Decimal(1004), decimal.Decimal(1008)])]

@@ -409,22 +409,20 @@ class TestVizProfiler(BaseTestViz):
    tracks = list(j['layout'])
    self.assertEqual(tracks[0], 'NV')
    self.assertEqual(tracks[1], 'NV Graph')
-    self.assertEqual(tracks[2], 'NV:1')
+    self.assertEqual(tracks[2], 'NV:1:SDMA:0')

    nv_events = j['layout']['NV']['events']
    self.assertEqual(nv_events[0]['name'], 'E_25_4n2')
    self.assertEqual(nv_events[0]['st'], 0)
    self.assertEqual(nv_events[0]['dur'], 2)
-    #self.assertEqual(j['devEvents'][6]['pid'], j['devEvents'][0]['pid'])

-    nv1_events = j['layout']['NV:1']['events']
-    self.assertEqual(nv1_events[0]['name'], 'NV -> NV:1')
-    self.assertEqual(nv1_events[0]['st'], 954)
-    #self.assertEqual(j['devEvents'][7]['pid'], j['devEvents'][3]['pid'])
+    sdma_events = j['layout']['NV:1:SDMA:0']['events']
+    self.assertEqual(sdma_events[0]['name'], 'NV -> NV:1')
+    self.assertEqual(sdma_events[0]['st'], 954)

    graph_events = j['layout']['NV Graph']['events']
    self.assertEqual(graph_events[0]['st'], nv_events[0]['st'])
-    self.assertEqual(graph_events[0]['st']+graph_events[0]['dur'], nv1_events[0]['st']+nv1_events[0]['dur'])
+    self.assertEqual(graph_events[0]['st']+graph_events[0]['dur'], sdma_events[0]['st']+sdma_events[0]['dur'])

  def test_bytes_per_kernel(self):
    step = 10
@@ -465,11 +463,11 @@ class TestVizProfiler(BaseTestViz):

  def test_layout_order(self):
    def fn(): return
-    for dname in ["TINY", "USER", "TEST:1 N1", "TEST:2 N1", "TEST:1 N2"]:
+    for dname in ["TINY", "USER", "TEST:1 N1", "TEST:2 N1", "TEST:1 N2", "TEST:1:ENGINE:0", "TEST:1"]:
      with cpu_profile("fn", dname): fn()
    layout = list(load_profile(cpu_events)["layout"])
    self.assertListEqual(layout[:2], ["USER","TINY"])
-    self.assertListEqual(layout[2:], ["TEST:1 N1","TEST:1 N2", "TEST:2 N1"])
+    self.assertListEqual(layout[2:], ["TEST:1", "TEST:1:ENGINE:0", "TEST:1 N1","TEST:1 N2", "TEST:2 N1"])

 def _alloc(b:int):
  a = Tensor.empty(b, device="NULL", dtype=dtypes.char)
@@ -52,7 +52,7 @@ def flip_contract_kernel(dest:UOp, src:UOp):
  j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
  vec = src[i, j].contract(j)
  store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
-  return store.end(i).sink(arg=KernelInfo(name=f"flip_contract_{dest.size}", opts_to_apply=()))
+  return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.size}", opts_to_apply=()))

 def slice_sum_kernel(dest:UOp, src:UOp):
  G = UOp.range(src.shape[0], 0)
@@ -88,13 +88,13 @@ def simple_qkv_kernel(O:UOp, Q:UOp, K:UOp, V:UOp) -> UOp:
 # **** backward callbacks ****

 def backward_gemm(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
-  out, a, b = kernel.src
+  out, a, b = kernel.src[1:]
  grad_a = (Tensor(gradient) @ Tensor(b).T).uop
  grad_b = (Tensor(a).T @ Tensor(gradient)).uop
  return (None, grad_a, grad_b)

 def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
-  out, a, b = kernel.src
+  out, a, b = kernel.src[1:]
  grad_a = Tensor.empty_like(Tensor(a)).custom_kernel(Tensor(gradient), Tensor(b).T, fxn=custom_gemm)[0].uop
  grad_b = Tensor.empty_like(Tensor(b)).custom_kernel(Tensor(a).T, Tensor(gradient), fxn=custom_gemm)[0].uop
  return (None, grad_a, grad_b)
@@ -128,6 +128,14 @@ class TestCustomKernel(unittest.TestCase):
    out = c.flatten().tolist()
    assert all(x == 2 for x in out), "all 2"

+  def test_sharded_add_one(self):
+    # PYTHON backend explicitly checks for OOB access for wrong multi shape regression
+    devs = ("PYTHON:0", "PYTHON:1")
+    a = Tensor.ones(4, 4).contiguous().shard(devs, axis=0)
+    c = Tensor(Tensor.empty(2, 4, device=devs).uop.multi(0), device=devs)
+    c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
+    assert (c == 2).all().item()
+
  def test_multioutput(self):
    a = Tensor.full((16, 16), 3.).contiguous()
    b = Tensor.full((16, 16), 3.).contiguous()
@@ -380,9 +380,40 @@ class TestBoolDType(TestDType): DTYPE = dtypes.bool

 class TestBFloat16Type(TestDType): DTYPE = dtypes.bfloat16

+class TestEmulatedBFloat16Type(TestBFloat16Type):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="bfloat16"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
 class TestFp8e4m3(TestDType): DTYPE = dtypes.fp8e4m3
+
+class TestEmulatedFp8e4m3(TestFp8e4m3):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="fp8e4m3"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
 class TestFp8e5m2(TestDType): DTYPE = dtypes.fp8e5m2

+class TestEmulatedFp8e5m2(TestFp8e5m2):
+  @classmethod
+  def setUpClass(cls):
+    cls.stack = contextlib.ExitStack()
+    cls.stack.enter_context(Context(EMULATED_DTYPES="fp8e5m2"))
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+
+  @classmethod
+  def tearDownClass(cls): cls.stack.close()
+
 class TestPtrDType(unittest.TestCase):
  def test_vec_double(self):
    dt1 = dtypes.float.vec(4).ptr().vec(4)
@@ -1,6 +1,6 @@
 import unittest, operator, math
 from tinygrad import Context, Tensor, dtypes, Device
-from tinygrad.dtype import DType, truncate
+from tinygrad.dtype import DType, truncate, fp8_to_float
 from tinygrad.helpers import CI, EMULATED_DTYPES, getenv
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -59,9 +59,10 @@ def universal_test(a, b, dtype, op):
  # lt and max with nan is undefined in tinygrad
  if op[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
  ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
-  tensor_value = (op[0](ta, tb)).numpy()
-  numpy_value = op[1](ta.numpy(), tb.numpy())
-  if dtype in dtypes.fp8s: numpy_value = truncate[dtype](numpy_value.item())
+  if dtype in dtypes.fp8s and op[0] not in (operator.lt, operator.eq):
+    tensor_value = fp8_to_float((op[0](ta.realize(), tb.realize())).bitcast(dtypes.uint8).item(), dtype)
+    numpy_value = truncate[dtype](op[1](ta.numpy(), tb.numpy()).item())
+  else: tensor_value, numpy_value = (op[0](ta, tb)).numpy(), op[1](ta.numpy(), tb.numpy())
  if dtype in dtypes.floats:
    if not is_dtype_supported(dtype) or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
      fe, fm = dtypes.finfo(dtype)
@@ -76,13 +77,14 @@ def universal_test_unary(a, dtype, op):
  # TODO: cos does not match for large input
  if op[0] == Tensor.cos and abs(a) > 30: return
  if op[0] == Tensor.log and a <= 0: return
-  out: Tensor = op[0](ta)
-  tensor_value = out.numpy()
-  numpy_value = op[1](ta.numpy())
  if dtype in dtypes.fp8s:
+    # normals are zero
+    if dtype in EMULATED_DTYPES.tolist(dtypes) and abs(ta.numpy().item()) < 0.015625: return
+    tensor_value = fp8_to_float(op[0](ta.realize()).bitcast(dtypes.uint8).item(), dtype)
+    numpy_value = truncate[dtype](v:=op[1](ta.numpy()).item())
    # cuda cast f32 inf to f8 MAX, amd cast it to nan(E4M3)/inf(E5M2)
-    if math.isinf(numpy_value.item()): return
-    numpy_value = truncate[dtype](numpy_value.item())
+    if math.isinf(v): return
+  else: tensor_value, numpy_value = op[0](ta).numpy(), op[1](ta.numpy())
  if dtype in dtypes.floats:
    atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1)}.get(dtype, (1e-6, 1e-5))
@@ -128,16 +130,31 @@ class TestDTypeALU(unittest.TestCase):
  def test_bfloat16(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)

+  @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="bfloat16")
+  def test_emulated_bfloat16(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
+
  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), f"no fp8e4m3 on {Device.DEFAULT}")
  @given(ht.fp8e4m3, ht.fp8e4m3, strat.sampled_from(binary_operations))
  def test_fp8e4m3(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op)

+  @given(ht.fp8e4m3, ht.fp8e4m3, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="fp8e4m3")
+  def test_emulated_fp8e4m3(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op)
+
  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2), f"no fp8e5m2 on {Device.DEFAULT}")
  @given(ht.fp8e5m2, ht.fp8e5m2, strat.sampled_from(binary_operations))
  def test_fp8e5m2(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)

+  @given(ht.fp8e5m2, ht.fp8e5m2, strat.sampled_from(binary_operations))
+  @Context(EMULATED_DTYPES="fp8e5m2")
+  def test_emulated_fp8e5m2(self, a, b, op):
+    universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)
+
  @given(ht.float32, strat.sampled_from(unary_operations))
  def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)

@@ -153,18 +170,34 @@ class TestDTypeALU(unittest.TestCase):
  @given(ht.bfloat16, strat.sampled_from(unary_operations))
  def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)

+  @given(ht.bfloat16, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="bfloat16")
+  def test_emulated_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
+
  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), f"no fp8e4m3 on {Device.DEFAULT}")
  @given(ht.fp8e4m3, strat.sampled_from(unary_operations))
  def test_fp8e4m3_unary(self, a, op):
    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0)
    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3), dtypes.fp8e4m3, op)

+  @given(ht.fp8e4m3, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="fp8e4m3")
+  def test_emulated_fp8e4m3_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3), dtypes.fp8e4m3, op)
+
  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2), f"no fp8e5m2 on {Device.DEFAULT}")
  @given(ht.fp8e5m2, strat.sampled_from(unary_operations))
  def test_fp8e5m2_unary(self, a, op):
    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)

+  @given(ht.fp8e5m2, strat.sampled_from(unary_operations))
+  @Context(EMULATED_DTYPES="fp8e5m2")
+  def test_emulated_fp8e5m2_unary(self, a, op):
+    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
+    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)
+
  @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
  def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)

@@ -175,7 +175,6 @@ class TestZeroFolding(unittest.TestCase):
 class TestAssignIssues(unittest.TestCase):
  # these are good failures. i'm not sure we need more, but we need to fix these.

-  @unittest.expectedFailure
  def test_assign_permuted_view_constant(self):
    # assigning to a permuted view should modify the underlying tensor
    arr = np.arange(6).reshape(2, 3).astype(np.float32)
@@ -185,7 +184,6 @@ class TestAssignIssues(unittest.TestCase):
    t.permute(1, 0).assign(Tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]))
    np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())

-  @unittest.expectedFailure
  def test_assign_shrink_view_constant(self):
    # assigning to a shrunk view should update the base tensor
    arr = np.arange(9).reshape(3, 3).astype(np.float32)
@@ -1,68 +0,0 @@
-import unittest
-import time
-import numpy as np
-from tinygrad import Tensor, dtypes
-from tinygrad.engine.realize import run_schedule
-
-class TestFusionOp(unittest.TestCase):
-  def test_contiguous_add(self):
-    def test(contig=False):
-      bt = Tensor(np.arange(16), dtype=dtypes.float32).reshape(4,4)
-      x = bt.permute(1,0)
-      if contig: x = x.contiguous()
-      return (x.permute(1,0) + bt).data()
-    assert test() == test(True)
-
-  def test_expand_fuse(self):
-    bt = Tensor(np.ones((10, 1)), dtype=dtypes.float32)
-    out = (bt*2).expand(10,10).sum(1)
-    sched = out.schedule()
-    run_schedule(sched)
-    outd = out.tolist()
-    assert all(x == 20.0 for x in outd)
-
-  def test_recursive_add(self):
-    st = time.perf_counter()
-    a = Tensor([1,2,3,4])
-    for _ in range(24): a = a + a
-    sched = a.schedule()
-    sched[-1].lower()
-    self.assertLess(time.perf_counter()-st, 2.0)
-    assert len(sched[-1].prg.p.src.splitlines()) < 250
-
-  def test_recursive_add_cmp(self):
-    st = time.perf_counter()
-    a = Tensor([1,2,3,4])
-    for _ in range(24): a = a + a
-    sched1 = a.schedule()
-    b = Tensor([1,2,3,4])
-    for _ in range(24): b = b + b
-    sched2 = b.schedule()
-    c = Tensor([1,2,3,4])
-    for _ in range(23): c = c + c
-    sched3 = c.schedule()
-    self.assertEqual(sched1[-1].ast, sched2[-1].ast)
-    with self.assertRaises(AssertionError): self.assertEqual(sched1[-1].ast, sched3[-1].ast)
-    self.assertLess(time.perf_counter()-st, 2.0)
-
-  def test_recursive_pad(self):
-    st = time.perf_counter()
-    val = 1.0
-    a = Tensor(val)
-    for _ in range(24): a = Tensor.stack(a, a)[0]
-    sched = a.schedule()
-    self.assertEqual(len(sched), 0)
-    self.assertLess(time.perf_counter()-st, 2.0)
-
-  def test_recursive_reshape(self):
-    st = time.perf_counter()
-    a = Tensor.empty(32, 32).realize()
-    b = Tensor.empty(16, 2).realize()
-    r = a.sum(1)
-    for _ in range(24): r = r.reshape(16, 2) + b
-    sched = r.schedule()
-    self.assertEqual(len(sched), 1)
-    self.assertLess(time.perf_counter()-st, 2.0)
-
-if __name__ == '__main__':
-  unittest.main(verbosity=2)
@@ -111,30 +111,18 @@ class TestJitFootguns(unittest.TestCase):
      self.assertEqual(first.numpy().item(), expected_first)
      buf = new_buf

-  def test_slice_assign_requires_realize(self):
-    """Slice assign then read from same buffer - assign isn't connected to read without explicit realize()."""
+  def test_slice_assign_works_without_realize(self):
+    """Slice assign then read from same buffer - pending assigns are side-realized."""
    from tinygrad import Variable
    v_pos = Variable("pos", 0, 3)
-
-    # without .realize() after assign, the read doesn't see the assigned values
    cache = Tensor.zeros(4, 4).contiguous().realize()
    @TinyJit
-    def f_broken(pos):
+    def f(pos):
      cache[pos:pos+1, :].assign(Tensor.ones(1, 4))
      return cache.sum().realize()
    for i in range(4):
      cache.assign(Tensor.zeros(4, 4)).realize()
-      self.assertEqual(f_broken(v_pos.bind(i)).item(), 0.0)  # should be 4.0!
-
-    # workaround: add .realize() after assign
-    cache2 = Tensor.zeros(4, 4).contiguous().realize()
-    @TinyJit
-    def f_fixed(pos):
-      cache2[pos:pos+1, :].assign(Tensor.ones(1, 4)).realize()
-      return cache2.sum().realize()
-    for i in range(4):
-      cache2.assign(Tensor.zeros(4, 4)).realize()
-      self.assertEqual(f_fixed(v_pos.bind(i)).item(), 4.0)
+      self.assertEqual(f(v_pos.bind(i)).item(), 4.0)

  def test_symbolic_pad_view_frozen(self):
    """Symbolic pad view has BIND values baked in at capture time. TODO: pad should be captured in jit."""
@@ -57,7 +57,7 @@ class TestOuterRange(unittest.TestCase):
    out.realize()

    # TODO: testing allclose
-    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+    assert Tensor.allclose(ref, out, atol=1e-6), f"max diff {(ref-out).abs().max().item()}"

 class TestOuterScan(unittest.TestCase):
  def _test_scan(self):
@@ -85,7 +85,7 @@ class TestOuterScan(unittest.TestCase):
    out.realize()

    # TODO: testing allclose
-    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+    assert Tensor.allclose(ref, out, atol=1e-5), f"max diff {(ref-out).abs().max().item()}"

 class TestOuterworld(unittest.TestCase):
  def test_range_plus_1(self):
@@ -166,7 +166,7 @@ class TestVmap(unittest.TestCase):
    out.realize()

    # TODO: testing allclose
-    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+    assert Tensor.allclose(ref, out, atol=1e-6), f"max diff {(ref-out).abs().max().item()}"
  def test_vmap_inner_fuse(self): self.test_vmap_inner(fuse=True)
  def test_vmap_outer(self): self.test_vmap_inner(AxisType.OUTER)
  def test_vmap_outer_fuse(self): self.test_vmap_inner(AxisType.OUTER, fuse=True)
@@ -6,6 +6,9 @@ from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.engine.realize import get_runner

 MOCKGPU = getenv("MOCKGPU")
+def _dev_base(d):
+  p = d.split(":")
+  return p[0] if len(p) < 2 or not p[1].isdigit() else f"{p[0]}:{p[1]}"

@contextlib.contextmanager
 def helper_collect_profile(*devs):
@@ -55,7 +58,7 @@ class TestProfiler(unittest.TestCase):
    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent)]
    assert len(kernel_runs) == 1, "one kernel run is expected"
    assert kernel_runs[0].name == runner_name, "kernel name is not correct"
-    assert not kernel_runs[0].is_copy, "kernel should not be copy"
+    assert _dev_base(kernel_runs[0].device) == kernel_runs[0].device, "kernel should not be on a sub-device"

  def test_profile_copyin(self):
    buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
@@ -63,10 +66,8 @@ class TestProfiler(unittest.TestCase):
    with helper_collect_profile(TestProfiler.d0) as profile:
      buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))

-    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
-    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent)]
+    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
    assert len(kernel_runs) == 1, "one kernel run is expected"
-    assert kernel_runs[0].is_copy, "kernel should be copy"

  def test_profile_multiops(self):
    runner_name = TestProfiler.runner._prg.name
@@ -77,16 +78,12 @@ class TestProfiler(unittest.TestCase):
      TestProfiler.runner([buf1, TestProfiler.a.uop.buffer], var_vals={})
      buf1.copyout(memoryview(bytearray(buf1.nbytes)))

-    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
-    evs = [x for x in profile if isinstance(x, ProfileRangeEvent)]
+    evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]

    assert len(evs) == 3, "3 kernel runs are expected"
    # NOTE: order of events does not matter, the tool is responsible for sorting them
-    copy_events = [e for e in evs if e.is_copy]
-    self.assertEqual(len(copy_events), 2)
-
-    prg_events = [e for e in evs if not e.is_copy]
-    assert prg_events[0].name == runner_name, "kernel name is not correct"
+    prg_events = [e for e in evs if e.device == TestProfiler.d0.device]
+    assert any(e.name == runner_name for e in prg_events), "kernel name is not correct"

    #for i in range(1, 3):
    #  assert evs[i].st > evs[i-1].en, "timestamp not aranged"
@@ -102,13 +99,9 @@ class TestProfiler(unittest.TestCase):
      buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
      buf2.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))

-    profile0, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
-    profile1, _ = helper_profile_filter_device(profile, d1.device)
-
-    for p in [profile0, profile1]:
-      evs = [x for x in p if isinstance(x, ProfileRangeEvent)]
+    for dev in [TestProfiler.d0.device, d1.device]:
+      evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and _dev_base(x.device) == dev]
      assert len(evs) == 1, "one kernel runs are expected"
-      assert evs[0].is_copy, "kernel should be copy"

  def test_profile_multidev_transfer(self):
    try: d1 = Device[f"{Device.DEFAULT}:1"]
@@ -118,10 +111,8 @@ class TestProfiler(unittest.TestCase):
    with helper_collect_profile(TestProfiler.d0, d1) as profile:
      buf1.to(f"{Device.DEFAULT}:1").realize()

-    profile0, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
-    kernel_runs = [x for x in profile0 if isinstance(x, ProfileRangeEvent)]
+    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
    assert len(kernel_runs) == 1, "one kernel run is expected"
-    assert kernel_runs[0].is_copy, "kernel should be copy"

  @unittest.skipIf(Device.DEFAULT in "METAL" or (MOCKGPU and Device.DEFAULT == "AMD"), "AMD mockgpu does not support queue wait interrupts")
  def test_profile_graph(self):
@@ -167,17 +158,18 @@ class TestProfiler(unittest.TestCase):
      return d2.timeline_signal.timestamp - d1.timeline_signal.timestamp

    # then test it by timing the GPU to GPU times
+    dev_evs = {x.device:x for x in Compiled.profile_events if isinstance(x, ProfileDeviceEvent)}
    jitter_matrix = [[float('nan')] * len(devs) for _ in range(len(devs))]
    pairs = [(p1, p2) for p1 in enumerate(devs) for p2 in enumerate(devs) if p1 != p2]
    for (i1, d1), (i2, d2) in pairs:
-      cpu_diff = d1.gpu2cpu_compute_time_diff - d2.gpu2cpu_compute_time_diff
+      cpu_diff = dev_evs[d1.device].tdiff - dev_evs[d2.device].tdiff
      jitter_matrix[i1][i2] = statistics.median(_sync_d2d(d1, d2) - _sync_d2d(d2, d1) for _ in range(20)) / 2 - cpu_diff

+    print("pairwise clock jitter matrix (us):\n" + '\n'.join([''.join([f'{float(item):8.3f}' for item in row]) for row in jitter_matrix]))
+
    for (i1, d1), (i2, d2) in pairs:
      assert abs(jitter_matrix[i1][i2]) < 0.5, "jitter should be less than 0.5us"

-    print("pairwise clock jitter matrix (us):\n" + '\n'.join([''.join([f'{float(item):8.3f}' for item in row]) for row in jitter_matrix]))
-
  @unittest.skip("this test is flaky")
  def test_cpu_profile(self):
    def test_fxn(err=False):
@@ -220,6 +212,7 @@ class TestProfiler(unittest.TestCase):
    for ge in graphs:
      self.assertEqual(len(ge.ents), len(graphs))

+  @unittest.skip("this test is flaky")
  def test_trace_metadata(self):
    with Context(TRACEMETA=1):
      a = Tensor.empty(1)+2
@@ -228,7 +221,7 @@ class TestProfiler(unittest.TestCase):
        Tensor.realize(a, b)
    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
    exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
-    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent) and not e.is_copy]
+    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent) and _dev_base(e.device) == e.device]
    self.assertEqual(len(exec_points), len(range_events), 2)
    self.assertEqual(len(dedup(e.arg['name'] for e in exec_points)), 1)
    self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)
@@ -1,8 +1,6 @@
-import unittest
-import random
-from os import getenv
+import unittest, random
 from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
-from tinygrad.helpers import Context
+from tinygrad.helpers import Context, getenv
 import numpy as np

 class TestSetitem(unittest.TestCase):
@@ -13,7 +11,7 @@ class TestSetitem(unittest.TestCase):
      ((6,6), (slice(2,4), slice(3,5)), 1.0),
      ((6,6), (3, 4), 1.0),
      ((6,6), (3, None, 4, None), 1.0),
-      ((4,4,4,4), (Ellipsis, slice(1,3), slice(None)), Tensor(4)),
+      ((4,4,4,4), (Ellipsis, slice(1,3), slice(None)), Tensor(4.0)),
      ((4,4,4,4), (Ellipsis, slice(1,3)), 4),
      ((4,4,4,4), (2, slice(1,3), None, 1), 4),
      ((4,4,4,4), (slice(1,3), slice(None), slice(0,4,2)), 4),
@@ -50,6 +48,10 @@ class TestSetitem(unittest.TestCase):
        t[1] = v
        self.assertEqual(t.dtype, dt)

+  def test_setitem_dtype_mismatch(self):
+    t = Tensor.zeros(6, dtype=dtypes.float).contiguous().realize()
+    with self.assertRaises(RuntimeError): t[2:4] = Tensor([1, 2], dtype=dtypes.int)
+
  def test_setitem_into_noncontiguous(self):
    t = Tensor.ones(4)
    with self.assertRaises(RuntimeError): t[1] = 5
@@ -109,8 +111,6 @@ class TestSetitem(unittest.TestCase):
  def test_setitem_consecutive_inplace_operator(self):
    t = Tensor.arange(4).reshape(2, 2).contiguous()
    t[1] += 2
-    t = t.contiguous()
-    # TODO: RuntimeError: can't double realize in one schedule
    t[1] -= 1
    np.testing.assert_allclose(t.numpy(), [[0, 1], [3, 4]])

@@ -182,7 +182,7 @@ class TestSetitem(unittest.TestCase):
    t[:-1] = t[1:]
    self.assertEqual(t.tolist(), [[2.0], [1.0], [1.0]])

-  # TODO: WEBGPU pipeline validation error
+  # TODO: WEBGPU pipeline validation error. this generates (1==gidx0)|(2==gidx0)|(3==gidx0)|(4==gidx0)|(5==gidx0) ...
  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU pipeline validation error")
  def test_setitem_big(self):
    idx_size, val = 256, 4
@@ -193,23 +193,23 @@ class TestSetitem(unittest.TestCase):

  def test_setitem_advanced_indexing(self):
    # Example from https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing
-    t = Tensor.zeros(10,20,30,40,50).contiguous()
+    t = Tensor.zeros(10,20,30,40,50, dtype=dtypes.int).contiguous()
    ind_1 = Tensor([5,3,7,8])
    ind_2 = Tensor([[[0],[1],[2]],[[3],[4],[5]]])
    v = Tensor.arange(2*3*4*10*30*50).reshape(2,3,4,10,30,50)
    t[:, ind_1, :, ind_2, :] = v
-    n = np.zeros((10,20,30,40,50))
+    n = np.zeros((10,20,30,40,50), dtype=np.int32)
    n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy()
-    np.testing.assert_allclose(t.numpy(), n)
+    np.testing.assert_equal(t.numpy(), n)

  def test_setitem_2d_tensor_indexing(self):
-    t = Tensor.zeros(2).contiguous()
+    t = Tensor.zeros(2, dtype=dtypes.int).contiguous()
    index = Tensor([[0, 1], [1,0]])
    v = Tensor.arange(2*2).reshape(2, 2).contiguous()
    t[index] = v
-    n = np.zeros((2,))
+    n = np.zeros((2,), dtype=np.int32)
    n[index.numpy()] = v.numpy()
-    np.testing.assert_allclose(t.numpy(), n)
+    np.testing.assert_equal(t.numpy(), n)

  @unittest.skip("slow")
  def test_setitem_tensor_indexing_fuzz(self):
@@ -1,51 +1,76 @@
 import unittest
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import getenv, CI
+from tinygrad.helpers import getenv
 from extra.gemm.asm.cdna.gemm import asm_gemm
 from test.helpers import needs_second_gpu

-def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
+# On non CDNA4 it will only validate the Tensor.custom_kernel integration
+# Use NULL=1 EMULATE=AMD_CDNA4 to also test the assembly
+def is_cdna4(): return getattr(Device[Device.DEFAULT].renderer, "arch", "").startswith("gfx950")
+
+def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=None, gpus:int=1) -> None:
  Tensor.manual_seed(0)
-  a_rand = Tensor.randn((batch, M, K), dtype=dtypes.float).sub(0.5).cast(dtype)
-  b_rand = Tensor.randn((K, N), dtype=dtypes.float).sub(0.5).cast(dtype)
+  a_rand = Tensor.randn(a_shape, dtype=dtypes.float).sub(0.5).cast(dtype)
+  b_rand = Tensor.randn(b_shape, dtype=dtypes.float).sub(0.5).cast(dtype)
  with Context(DEBUG=0):
    Tensor.realize(a_rand, b_rand)

  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(gpus)) if (multi:=gpus>1) else None

  a, b = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
-  if multi: a, b = a.shard(devs, axis=0), b.shard(devs, axis=None)
-  tst = asm_gemm(a, b)
-  tst.sum().backward()
+  if multi: a, b = a.shard(devs, axis=a_shard), b.shard(devs, axis=b_shard)
+  with Context(ASM_GEMM=1):
+    tst = asm_gemm(a, b)
+    tst.sum().backward()
  Tensor.realize(tst, a.grad, b.grad)

  a_ref, b_ref = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
-  if multi: a_ref, b_ref = a_ref.shard(devs, axis=0), b_ref.shard(devs, axis=None)
-  with Context(ASM_GEMM=0): ref = a_ref @ b_ref
-  ref.sum().backward()
+  if multi: a_ref, b_ref = a_ref.shard(devs, axis=a_shard), b_ref.shard(devs, axis=b_shard)
+  with Context(ASM_GEMM=0):
+    ref = asm_gemm(a_ref, b_ref)
+    ref.sum().backward()
  Tensor.realize(ref, a_ref.grad, b_ref.grad)

+  # no validation on the NULL device
+  if a_rand.device.startswith("NULL"): return None
+  atol, rtol = (1e-2, 1e-3)
  with Context(DEBUG=0):
-    assert (tst - ref).square().max().float().item() < 1e-6, "forward mismatch"
-    assert (a.grad - a_ref.grad).square().max().float().item() < 1e-3, "grad_a mismatch"
-    assert (b.grad - b_ref.grad).square().max().float().item() < 1e-3, "grad_b mismatch"
+    assert tst.allclose(ref, atol=atol, rtol=rtol), "forward mismatch"
+    assert a.grad.allclose(a_ref.grad, atol=atol, rtol=rtol), "grad_a mismatch"
+    assert b.grad.allclose(b_ref.grad, atol=atol, rtol=rtol), "grad_b mismatch"

-SCALE = 128 if CI else 1

+def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
+  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
+
+def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=8) -> None:
+  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=1, b_shard=0, gpus=gpus)
+
+# 128x smaller than usual
+# uses the UOp GEMM, runs on non CDNA4 and CI
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
 class TestGemm(unittest.TestCase):
-  def test_simple(self): verify_asm_gemm(1, N:=(getenv("N", 4096)//SCALE), N, N, dtype=dtypes.half)
-  def test_gemm(self): verify_asm_gemm(1, 8192//SCALE, 4096//SCALE, 14336//SCALE)
-  def test_gemm_batched(self): verify_asm_gemm(2, 8192//SCALE, 4096//SCALE, 4096//SCALE)
+  def setUp(self):
+    if is_cdna4(): self.skipTest("shapes are too small for the assembly GEMM")
+  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 32), N, N, dtype=dtypes.half)
+  def test_gemm(self): verify_asm_gemm(1, 64, 32, 112)
+  def test_gemm_batched(self): verify_asm_gemm(2, 64, 32, 32)
  @needs_second_gpu
-  def test_gemm_multi(self): verify_asm_gemm(2, 8192//SCALE, 4096//SCALE, 4096//SCALE, gpus=2)
+  def test_gemm_multi(self): verify_asm_gemm(2, 64, 32, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_k_sharded(self): verify_asm_gemm_k_sharded(64, 64, 2*64, gpus=2)

+# uses the Asm GEMM on CDNA4 only for speed reasons
 class TestGemmLarge(unittest.TestCase):
  def setUp(self):
-    if getattr(Device[Device.DEFAULT].renderer, "arch", "") != "gfx950":
+    if not is_cdna4():
      self.skipTest("very slow on non mi350x")

+  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)
+  def test_gemm_batched(self): verify_asm_gemm(2, 8192, 4096, 4096)
+
  def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
  @unittest.skip("disabled, asm in this shape is slower than tinygrad")
  def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, dtype=dtypes.bfloat16, gpus=8)
@@ -55,6 +80,12 @@ class TestGemmLarge(unittest.TestCase):
  def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
  @unittest.skip("disabled, asm in this shape is slower than tinygrad")
  def test_gemm7(self): verify_asm_gemm(1, 8192, 128256, 4096)
+  def test_gemm8(self): verify_asm_gemm(1, 4096, 14336, 8192)
+  def test_gemm9(self): verify_asm_gemm(8, 4096, 14336, 8192, dtype=dtypes.bfloat16, gpus=8)
+  def test_gemm10(self): verify_asm_gemm(1, 4096, 8192, 4096)
+  def test_k_sharded_1(self): verify_asm_gemm_k_sharded(14336, 4096, 8*8192, gpus=8)
+  def test_k_sharded_2(self): verify_asm_gemm_k_sharded(4096, 14336, 8*8192, gpus=8)
+  def test_k_sharded_3(self): verify_asm_gemm_k_sharded(4096, 4096, 8*8192, gpus=8)
  def test_gemm_unsupported(self):
    with self.assertRaisesRegex(AssertionError, "shape not supported"):
      verify_asm_gemm(8, 1024, 1024, 4096, gpus=8)
@@ -1,8 +1,9 @@
 import unittest

-from tinygrad import Device
-from tinygrad.helpers import fetch
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.helpers import fetch, round_up
 from extra.hevc.hevc import parse_hevc_file_headers, nv_gpu
+from extra.hevc.decode import hevc_decode

 class TestHevc(unittest.TestCase):
  def test_hevc_parser(self):
@@ -61,5 +62,26 @@ class TestHevc(unittest.TestCase):
    self.assertEqual(list(frame3.initreflistidxl0), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
    self.assertEqual(list(frame3.initreflistidxl1), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
    self.assertEqual(list(frame3.RefDiffPicOrderCnts), [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+
+  @unittest.skipUnless(Device.DEFAULT == "NV", "NV only")
+  def test_hevc_decode(self):
+    url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
+    dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
+
+    opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat)
+    frame_info = frame_info[:4]
+    out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
+
+    hevc_tensor = Tensor(dat, device="NV")
+    opaque_nv = opaque.to("NV").contiguous().realize()
+
+    frames = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w))
+    Device.default.synchronize()
+    self.assertEqual(len(frames), 4)
+    for f in frames:
+      self.assertEqual(f.shape, out_image_size)
+      self.assertEqual(f.dtype, dtypes.uint8)
+      self.assertEqual(f.device, "NV")
+
 if __name__ == "__main__":
  unittest.main()
@@ -560,30 +560,16 @@ class TestAssignOrdering(unittest.TestCase):

  def test_overlapping_slice_assigns(self):
    """Overlapping slice assigns - later write should win for overlapping elements."""
-    # without .realize(): assigns not executed, buffer stays zeros
    buf = Tensor.zeros(8).contiguous().realize()
    buf[0:4].assign(Tensor.ones(4))
    buf[2:6].assign(Tensor.ones(4) * 2)
-    np.testing.assert_equal(buf.numpy(), [0,0,0,0,0,0,0,0])  # TODO: wrong! should be [1,1,2,2,2,2,0,0]
-
-    # with .realize(): assigns execute in order
-    buf = Tensor.zeros(8).contiguous().realize()
-    buf[0:4].assign(Tensor.ones(4)).realize()
-    buf[2:6].assign(Tensor.ones(4) * 2).realize()
    np.testing.assert_equal(buf.numpy(), [1,1,2,2,2,2,0,0])

  def test_overlapping_slice_assigns_reverse(self):
    """Overlapping slice assigns in reverse order."""
-    # without .realize(): assigns not executed
    buf = Tensor.zeros(8).contiguous().realize()
    buf[2:6].assign(Tensor.ones(4) * 2)
    buf[0:4].assign(Tensor.ones(4))
-    np.testing.assert_equal(buf.numpy(), [0,0,0,0,0,0,0,0])  # TODO: wrong! should be [1,1,1,1,2,2,0,0]
-
-    # with .realize(): assigns execute in order
-    buf = Tensor.zeros(8).contiguous().realize()
-    buf[2:6].assign(Tensor.ones(4) * 2).realize()
-    buf[0:4].assign(Tensor.ones(4)).realize()
    np.testing.assert_equal(buf.numpy(), [1,1,1,1,2,2,0,0])

  def test_read_between_writes(self):
@@ -619,26 +605,14 @@ class TestAssignOrdering(unittest.TestCase):

  def test_slice_write_then_full_read(self):
    """Write to slice, then read full buffer."""
-    # without .realize(): orphan slice assign not triggered by .numpy()
    buf = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
    buf[1:3].assign(Tensor([5, 6]))
-    np.testing.assert_equal(buf.numpy(), [0, 0, 0, 0])  # TODO: wrong! should be [0, 5, 6, 0]
-
-    # with .realize(): assign executes
-    buf = Tensor.zeros(4, dtype=dtypes.int32).contiguous().realize()
-    buf[1:3].assign(Tensor([5, 6])).realize()
    np.testing.assert_equal(buf.numpy(), [0, 5, 6, 0])

  def test_chained_slice_copies(self):
    """Copy from one slice to another within same buffer."""
-    # without .realize(): orphan slice assign not triggered
    buf = Tensor([1, 2, 3, 4, 5, 6, 7, 8]).contiguous().realize()
    buf[4:8].assign(buf[0:4].contiguous())
-    np.testing.assert_equal(buf.numpy(), [1, 2, 3, 4, 5, 6, 7, 8])  # TODO: wrong! should be [1,2,3,4,1,2,3,4]
-
-    # with .realize(): assign executes
-    buf = Tensor([1, 2, 3, 4, 5, 6, 7, 8]).contiguous().realize()
-    buf[4:8].assign(buf[0:4].contiguous()).realize()
    np.testing.assert_equal(buf.numpy(), [1, 2, 3, 4, 1, 2, 3, 4])

  def test_swap_slices(self):
@@ -661,16 +635,9 @@ class TestAssignOrdering(unittest.TestCase):

  def test_reduction_after_partial_assign(self):
    """Reduction over buffer after partial assign - must see the assigned values."""
-    # without .realize(): orphan slice assign not triggered by reduction
    buf = Tensor.zeros(4, 4).contiguous().realize()
    buf[0:2, :].assign(Tensor.ones(2, 4))  # top half = 1
    total = buf.sum()
-    self.assertEqual(total.item(), 0)  # TODO: wrong! should be 8 (2*4 ones)
-
-    # with .realize(): assign executes before reduction
-    buf = Tensor.zeros(4, 4).contiguous().realize()
-    buf[0:2, :].assign(Tensor.ones(2, 4)).realize()
-    total = buf.sum()
    self.assertEqual(total.item(), 8)

  def test_multiple_reductions_different_views(self):
@@ -734,34 +701,18 @@ class TestAssignOrdering(unittest.TestCase):
  def test_variable_slice_ordering(self):
    """Variable-indexed slices - tests symbolic dependency tracking."""
    v_i = Variable("i", 0, 3)
-
-    # without .realize(): orphan slice assigns not triggered
    buf = Tensor.zeros(4, 4).contiguous().realize()
    buf[v_i.bind(0):v_i.bind(0)+1, :].assign(Tensor.ones(1, 4))
-    row0_sum = buf[0:1, :].sum()
-    self.assertEqual(row0_sum.item(), 0)  # TODO: wrong! should be 4
-
-    # with .realize(): assigns execute
-    buf = Tensor.zeros(4, 4).contiguous().realize()
-    buf[v_i.bind(0):v_i.bind(0)+1, :].assign(Tensor.ones(1, 4)).realize()
-    row0_sum = buf[0:1, :].sum()
-    buf[v_i.bind(1):v_i.bind(1)+1, :].assign(Tensor.ones(1, 4) * 2).realize()
-    row1_sum = buf[1:2, :].sum()
-    self.assertEqual(row0_sum.item(), 4)
-    self.assertEqual(row1_sum.item(), 8)
+    buf[v_i.bind(1):v_i.bind(1)+1, :].assign(Tensor.ones(1, 4) * 2)
+    self.assertEqual(buf[0:1, :].sum().item(), 4)
+    self.assertEqual(buf[1:2, :].sum().item(), 8)

  def test_multiple_slice_assigns_then_read(self):
-    """Multiple non-overlapping slice assigns then read - RAW dependencies must ensure all writes complete before read."""
+    """Multiple non-overlapping slice assigns then read."""
    buf = Tensor.zeros(4).contiguous().realize()
    buf[0:1].assign(Tensor.ones(1))
    buf[1:2].assign(Tensor.full((1,), 2.0))
    buf[2:3].assign(Tensor.full((1,), 3.0))
-    self.assertEqual(buf.sum().realize().item(), 0.0)  # TODO: wrong! should be 1 + 2 + 3 + 0 = 6
-
-    buf = Tensor.zeros(4).contiguous().realize()
-    buf[0:1].assign(Tensor.ones(1)).realize()
-    buf[1:2].assign(Tensor.full((1,), 2.0)).realize()
-    buf[2:3].assign(Tensor.full((1,), 3.0)).realize()
    self.assertEqual(buf.sum().realize().item(), 6.0)

 if __name__ == "__main__":
@@ -13,7 +13,7 @@ class TestCall(unittest.TestCase):
    # we define a plus function
    plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))

-    c = Tensor.call(a, b, fxn=plus_fxn)
+    c = Tensor.call(a, b, fxn=plus_fxn, inline=True)
    np.testing.assert_equal(c.numpy(), (a+b).numpy())

  def test_call_plus_backward(self):
@@ -30,7 +30,7 @@ class TestCall(unittest.TestCase):

    # we define a plus function
    plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
-    c = Tensor.call(a, b, fxn=plus_fxn, grad_fxn=grad_fxn)
+    c = Tensor.call(a, b, fxn=plus_fxn, grad_fxn=grad_fxn, inline=True)
    c.mean().backward()

    np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
@@ -46,7 +46,7 @@ class TestCall(unittest.TestCase):
    a.grad, b.grad = None, None

    plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
-    c = Tensor.call(a, b, fxn=plus_fxn)
+    c = Tensor.call(a, b, fxn=plus_fxn, inline=True)
    c.mean().backward()

    np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
@@ -57,7 +57,7 @@ class TestCall(unittest.TestCase):
    a = Tensor.randn(M, K)
    b = Tensor.randn(K, N)
    Tensor.realize(a, b)
-    c = Tensor.call(a, b, fxn=a.as_param(0) @ b.as_param(1))
+    c = Tensor.call(a, b, fxn=a.as_param(0) @ b.as_param(1), inline=True)
    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)

  @unittest.skip("needs GEMM on mixins")
@@ -70,7 +70,7 @@ class TestCall(unittest.TestCase):
    # we define a gemm function
    x = UOp.param(0, dtypes.float, shape=(M, K))
    y = UOp.param(1, dtypes.float, shape=(K, N))
-    c = Tensor.call(a, b, fxn=x@y)
+    c = Tensor.call(a, b, fxn=x@y, inline=True)

    np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)

@@ -86,7 +86,7 @@ class TestCall(unittest.TestCase):

    p0, p1 = UOp.param(0, dtypes.float, (10,10)), UOp.param(1, dtypes.float, (10,10))
    complex_fxn = (p0*p1 + p0).exp2() * p1.reciprocal()
-    c = Tensor.call(a, b, fxn=complex_fxn)
+    c = Tensor.call(a, b, fxn=complex_fxn, inline=True)
    c.mean().backward()

    np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
@@ -4,7 +4,7 @@ from tinygrad import Tensor, Device, dtypes
 from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType, DTYPES_DICT
 from tinygrad.nn.state import safe_load, safe_save, get_state_dict, torch_load
-from tinygrad.helpers import Timing, fetch, OSX
+from tinygrad.helpers import Timing, fetch, OSX, dedup
 from test.helpers import slow

 class TempDirTestCase(unittest.TestCase):
@@ -172,7 +172,7 @@ class TestSafetensors(TempDirTestCase):
    assert json.loads(dat[8:8+sz])['__metadata__']['hello'] == 'world'

  def test_save_all_dtypes(self):
-    for dtype in DTYPES_DICT.values():
+    for dtype in dedup(DTYPES_DICT.values()):
      if dtype in [dtypes.bfloat16]: continue # not supported in numpy
      if not is_dtype_supported(dtype): continue
      path = self.tmp(f"ones.{dtype}.safetensors")
@@ -301,6 +301,11 @@ class TestDiskTensor(TempDirTestCase):
    # self.assertEqual(dt.tolist(), [10, 2, 20, 4, 30, 6])
    self.assertEqual(dt.tolist(), [10, 20, 30, 4, 5, 6])  # wrong!

+  def test_advanced_setitem_not_supported(self):
+    dt = Tensor.arange(12).reshape(3, 4).to(f"disk:{self.tmp('dt_advanced_setitem')}")
+    with self.assertRaises(RuntimeError, msg="advanced setitem is not supported for DISK tensors"):
+      dt[Tensor([0, 2]), Tensor([1, 3])] = 99
+
  def test_assign_const_to_disk(self):
    # assign from CONST (Tensor.full) to disk - source has no buffer, needs contiguous first
    dt = Tensor.empty(4, device=f"disk:{self.tmp('dt_assign_const')}", dtype=dtypes.int32)
@@ -40,8 +40,6 @@ def ggml_tensor_to_numpy(tensor: ggml.ggml_tensor_p):
  return np.lib.stride_tricks.as_strided(output, shape=shape, strides=strides), ctx

@unittest.skipIf(any(not is_dtype_supported(t) for t in [ dtypes.uint8, dtypes.half ]), "Backend must support uint8 and half")
-# TODO: WEBGPU GGUF dequantization produces incorrect values
-@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU GGUF dequantization issue")
 class TestGGUF(unittest.TestCase):
  def setUp(self) -> None:
    params = ggml.ggml_init_params(mem_size=0, mem_buffer=None, no_alloc=False)
@@ -1,36 +1,31 @@
 from typing import cast
 from dataclasses import replace
 import itertools
-from tinygrad.helpers import DISABLE_FAST_IDIV, EMULATED_DTYPES, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, getenv, TracingKey, Context
+from tinygrad.helpers import DISABLE_FAST_IDIV, EMULATED_DTYPES, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, TracingKey, Context
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, pyrender
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
 from tinygrad.renderer import Renderer, ProgramSpec
-from tinygrad.dtype import dtypes, PtrDType
+from tinygrad.dtype import dtypes, promo_lattice
+from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import panic
 from tinygrad.codegen.opt import Opt

 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load
-from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_unsupported_dtypes_patterns, get_transcendental_patterns
+from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcendental_patterns, pm_float_decomp, pm_long_decomp
 from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
  ReduceContext, correct_load_store, pm_render, pm_add_loads
 from tinygrad.codegen.opt.postrange import apply_opts, 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
+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

-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),
-])
-
 def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp:
  if ren is None: ren = Renderer()

-  if getenv("VIZ"): graph_rewrite(sink, PatternMatcher([]), name="View Base AST")
+  if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Base AST")
  if DEBUG >= 5: print(pyrender(sink))
  if SPEC: type_verify(sink, kernel_spec)

@@ -94,10 +89,13 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
  # decompositions
  supported_ops = tuple(ren.code_for_op.keys())
  pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, ren.device, bool(DISABLE_FAST_IDIV))
-  pm_unsupported = get_unsupported_dtypes_patterns(ren.device, tuple(EMULATED_DTYPES.tolist(dtypes)))
  pm_transcendental = symbolic_simple+get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
  sink = graph_rewrite(sink, pm_decomp, ctx=ren.device, name="decompositions")
-  sink = graph_rewrite(sink, pm_unsupported, ctx=ren.device, name="unsupported dtypes", bottom_up=True)
+  if not is_dtype_supported(dtypes.long, ren.device) or dtypes.long in EMULATED_DTYPES.tolist(dtypes):
+    sink = graph_rewrite(sink, pm_long_decomp, name="decomp long -> int", bottom_up=True)
+  for fr, to in [(fr, next((to for to in promo_lattice[fr] if is_dtype_supported(to, ren.device)), dtypes.float))
+                 for fr in EMULATED_DTYPES.tolist(dtypes) if fr in dtypes.floats]:
+    sink = graph_rewrite(sink, pm_float_decomp, ctx=(fr, to), name=f"decomp {fr} -> {to}", bottom_up=True)
  sink = graph_rewrite(sink, pm_transcendental, ctx=ren.device, name="transcendental")

  # final rules for the renderer (without sym)
@@ -114,7 +112,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
 # inject IF/ENDIF. only needed if device doesn't support gated stores
 pm_linearize_cleanups = PatternMatcher([
  # if statements are not allowed in the graph
-  (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError("if not allowed in graph"))),
+  (UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError, "if not allowed in graph")),
  # gated INDEX becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
  (UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat())),
   lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))]))
@@ -141,7 +139,6 @@ def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
  return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),), arg=ctx.aux(list(lin.src)) if ctx.has_aux else prg.arg)

 def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None:
-  if ctx.compiler is None: return None
  lib = ctx.compiler.compile_cached(source.arg)
  return prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),))

@@ -7,17 +7,13 @@ from tinygrad.helpers import prod, getenv, TUPLE_ORDER
 def linearize(sink:UOp) -> list[UOp]:
  # this is a toposort with priority
  lst = list(sink.toposort())
-  consumers: defaultdict[UOp, list[UOp]] = defaultdict(list)
-  in_degree:dict[UOp, int] = {}
-  out_degree:dict[UOp, int] = {}
+  out_degree:defaultdict[UOp, int] = defaultdict(int)
  priorities:dict[UOp, tuple[int, int, Any]] = {}

  # get consumers and assign priorities
  # NOTE: this requires the lst be locally toposorted
  for u in reversed(lst):
-    for s in u.src: consumers[s].append(u)
-    in_degree[u] = len(u.src)
-    out_degree[u] = len(consumers[u])
+    for s in u.src: out_degree[s] += 1

    # we place UOps with higher run_counts later
    run_count = prod([int(r.vmax)+1 for r in u.ranges])
@@ -1,14 +1,14 @@
 from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
-from typing import Any, Generic, TypeVar, Iterator, Generator
+from typing import Any, Generic, TypeVar, Iterator, Generator, TYPE_CHECKING
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup, ContextVar
 from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ, CPU_LLVM, CPU_LVP, NV_PTX, CUDA_PTX, NV_NAK
 from tinygrad.helpers import EMULATED_DTYPES
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
-from tinygrad.renderer import Renderer
+if TYPE_CHECKING: from tinygrad.renderer import Renderer

 # **************** Device ****************

@@ -56,14 +56,13 @@ atexit.register(lambda: [Device[dn].finalize() for dn in Device._opened_devices]
 # **************** Profile ****************

@dataclass(frozen=True)
-class ProfileDeviceEvent(ProfileEvent):
-  device:str; comp_tdiff:decimal.Decimal=decimal.Decimal(0); copy_tdiff:decimal.Decimal=decimal.Decimal(0); props:dict[str,Any]|None=None # noqa: E702
+class ProfileDeviceEvent(ProfileEvent): device:str; tdiff:decimal.Decimal=decimal.Decimal(0); props:dict[str,Any]|None=None # noqa: E702

@dataclass(frozen=True)
 class ProfileProgramEvent(ProfileEvent): device:str; name:str; lib:bytes|None; base:int|None; tag:int|None=None # noqa: E702

@dataclass(frozen=True)
-class ProfileGraphEntry: device:str; name:str; st_id:int; en_id:int; is_copy:bool # noqa: E702
+class ProfileGraphEntry: device:str; name:str; st_id:int; en_id:int # noqa: E702

@dataclass(frozen=True)
 class ProfileGraphEvent(ProfileEvent): ents:list[ProfileGraphEntry]; deps:list[list[int]]; sigs:list[decimal.Decimal] # noqa: E702
@@ -279,37 +278,34 @@ class Compiler:
  def disassemble(self, lib:bytes): pass

@dataclass(frozen=True)
-class CompilerPair: renderer:type[Renderer]|functools.partial; compiler:type[Compiler]|functools.partial|None; ctrl_var:ContextVar|None = None # noqa: E702
-
-@dataclass(frozen=True)
-class CompilerSet: cset:list[CompilerPair]; ctrl_var:ContextVar|None = None # noqa: E702
+class CompilerSet: cset:list[tuple[type[Renderer]|functools.partial, ContextVar|None]]; ctrl_var:ContextVar|None = None # noqa: E702

 class Compiled:
  profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.

  def __init__(self, device:str, allocator:Allocator, compilers:CompilerSet|None, runtime, graph=None, group_id=None):
+    from tinygrad.renderer import Renderer
+
    self.device, self.allocator, self.runtime, self.graph, self.group_id = device, allocator, runtime, graph, group_id

    self.comps_ctrl_var = compilers.ctrl_var if compilers is not None else None
-    self.comp_sets:dict[Any, tuple[ContextVar|None, tuple[type[Renderer]|functools.partial, type[Compiler]|functools.partial|None]]] = {}
-    self.cached_pair:dict[Any, tuple[Renderer, Compiler|None]] = {}
-    for cpair in (compilers.cset if compilers is not None else [CompilerPair(Renderer, Compiler)]):
-      self.comp_sets[self._compiler_name(cpair.renderer, cpair.compiler)] = (cpair.ctrl_var, (cpair.renderer, cpair.compiler))
+    self.comp_sets:dict[str, tuple[ContextVar|None, type[Renderer]|functools.partial]] = {}
+    self.cached_pair:dict[Any, Renderer] = {}
+    for ren, var in (compilers.cset if compilers is not None else [(Renderer, None)]):
+      self.comp_sets[var.key.split('_', 1)[-1] if var is not None else self._compiler_name(ren)] = (var, ren)

  @property
-  def renderer(self) -> Renderer: return self._select_compiler_pair()[0]
+  def renderer(self) -> Renderer: return self._select_compiler_pair()

  @property
  def compiler(self) -> Compiler:
-    if (ret:=self.renderer.compiler or self._select_compiler_pair()[1]) is None: raise RuntimeError(f"no compiler for {self.device}")
+    if (ret:=self.renderer.compiler) is None: raise RuntimeError(f"no compiler for {self.device}")
    return ret

-  def _compiler_name(self, r:type[Renderer]|functools.partial, c:type[Compiler]|functools.partial|None) -> str:
-    devname = self.device.split(':')[0].upper()
-    if c is None: return unwrap_class_type(r).__name__.upper().removesuffix("RENDERER").removeprefix(devname) or devname
-    return unwrap_class_type(c).__name__.upper().removesuffix("COMPILER").removeprefix(devname) or devname
+  def _compiler_name(self, r:type[Renderer]|functools.partial) -> str:
+    return unwrap_class_type(r).__name__.upper().removesuffix("RENDERER").removeprefix(devname:=self.device.split(':')[0].upper()) or devname

-  def _select_compiler_pair(self) -> tuple[Renderer, Compiler|None]:
+  def _select_compiler_pair(self) -> Renderer:
    # select forced compiler from global env var.
    forced_comps = set([self.comp_sets[val][1]] if self.comps_ctrl_var is not None and (val:=self.comps_ctrl_var.value) else [])

@@ -398,18 +394,18 @@ def enumerate_devices_str() -> Generator[str, None, None]:
      d = Device[device]
      default_comp_pairs, default_compiler, cc_ctrl_var = d.comp_sets, d.compiler, d.comps_ctrl_var
      try:
-        for k,(en,(r,c)) in default_comp_pairs.items():
-          d.comp_sets = {k:(None,(r,c))} # env var set to None, so it doesn't interfere
+        for k,(en,r) in default_comp_pairs.items():
+          d.comp_sets = {k:(None,r)} # env var set to None, so it doesn't interfere
          d.comps_ctrl_var = None
          try:
            # d.renderer, d.compiler = r(), c()
            with Context(CACHELEVEL=0): test = (Tensor([1,2,3], device=device) * 2).tolist()
            if test != [2,4,6]: raise ValueError(f"got {test} instead of [2, 4, 6]")
-            set_text = f'({cc_ctrl_var.key}={d._compiler_name(r, c)} to make default)' if cc_ctrl_var is not None else ''
+            set_text = f'({cc_ctrl_var.key}={d._compiler_name(r)} to make default)' if cc_ctrl_var is not None else ''
            default_text = '(default)' if type(default_compiler) is type(d.compiler) else set_text
-            compilers_results.append(f"{colored('+', 'green')} {unwrap_class_type(c).__name__} {default_text}")
+            compilers_results.append(f"{colored('+', 'green')} {d._compiler_name(r)} {default_text}")
            any_works = True
-          except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {unwrap_class_type(c).__name__}: {e}")
+          except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {d._compiler_name(r)}: {e}")
      finally:
        # put the defaults back!
        d.comp_sets, d.comps_ctrl_var = default_comp_pairs, cc_ctrl_var
@@ -211,6 +211,10 @@ class dtypes:

  fp8s = (fp8e4m3, fp8e5m2)
  floats = fp8s + (float16, bfloat16, float32, float64)
+  int8s = (uint8, int8)
+  int16s = (uint16, int16)
+  int32s = (uint32, int32)
+  int64s = (uint64, int64)
  uints = (uint8, uint16, uint32, uint64)
  sints = (int8, int16, int32, int64)
  ints = uints + sints
@@ -1,7 +1,7 @@
 import time
 from typing import cast
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, Kernel
+from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, gate_kernel_sink
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Buffer, MultiBuffer
 from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize, SCACHE, Metadata
@@ -22,14 +22,14 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
    # build kernel dependency graph: edges from producer kernel to consumer kernels
    children: dict[UOp, list[UOp]] = {}
    in_degree: dict[UOp, int] = {}
-    for u in sched_sink.toposort():
+    for u in sched_sink.toposort(gate_kernel_sink):
      if u.op is Ops.RANGE: in_degree.setdefault(u, 0)
      if u.op is not Ops.AFTER: continue
      if (k:=u.src[1]).op is Ops.RANGE: continue  # RANGEs are scheduled directly, not through dependency graph
-      assert k.op in {Ops.KERNEL, Ops.END}, f"AFTER src[1] should be KERNEL or END, not {k.op}"
+      assert k.op in {Ops.CALL, Ops.END}, f"AFTER src[1] should be KERNEL or END, not {k.op}"
      in_degree.setdefault(k, 0)
-      if k.op is Ops.END: assert k.src[0].op is Ops.KERNEL, f"END src[0] should be KERNEL, not {k.src[0].op}"
-      for s in k.src[0].src if k.op is Ops.END else k.src:
+      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}"
+      for s in k.src[0].src[1:] if k.op is Ops.END else k.src[1:]:
        match (s := _unwrap_src(s)).op:
          case Ops.AFTER:
            children.setdefault(s.src[1], []).append(k)
@@ -54,13 +54,13 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
    while len(queue):
      k = rk = queue.popleft()
      if k.op is Ops.END: k = k.src[0]
-      assert k.op in {Ops.RANGE, Ops.KERNEL}, f"unexpected op in queue: {k.op}"
+      assert k.op in {Ops.RANGE, Ops.CALL}, f"unexpected op in queue: {k.op}"
      if k.op is Ops.RANGE: schedule.append(k)
-      elif k.op is Ops.KERNEL:
-        ast = (kernel:=cast(Kernel, k.arg)).ast
-        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src if s.op is not Ops.BIND)
-        bound_ranges = tuple(s for s in k.src if s.op is Ops.BIND and len(s.src) > 1 and s.src[1].op is Ops.RANGE)
-        sched_item[k] = (ast, buf_uops, kernel.metadata, bound_ranges)
+      elif k.op is Ops.CALL:
+        ast = k.src[0]
+        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND)
+        bound_ranges = tuple(s for s in k.src[1:] if s.op is Ops.BIND and len(s.src) > 1 and s.src[1].op is Ops.RANGE)
+        sched_item[k] = (ast, buf_uops, k.arg.metadata, bound_ranges)
        schedule.append(k)
        if rk.op is Ops.END: schedule.append(rk)
      for x in children.get(rk, []):
@@ -86,7 +86,7 @@ def unroll_outer_ranges(schedule:list[UOp], sched_item:dict[UOp, ScheduleItem])
        sched_ptr = range_ptrs[si.src[1]]
        continue
    else:
-      assert si.op is Ops.KERNEL, f"unexpected op in schedule: {si.op}"
+      assert si.op is Ops.CALL, f"unexpected op in schedule: {si.op}"
      ast, buf_uops, metadata, bound_ranges = sched_item[si]
      fixedvars = {s.src[0].arg[0]:in_ranges[s.src[1]] for s in bound_ranges}
      pre_schedule.append(ExecItem(ast, [], metadata, fixedvars))
@@ -52,7 +52,6 @@ pm_gradient = PatternMatcher([
  (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
  # NOTE: this is only correct when the KERNEL has a single output
  (UPat(Ops.AFTER), lambda ctx: (ctx, ctx)),
-  (UPat(Ops.CUSTOM_KERNEL, name="k"), lambda ctx, k: k.arg.grad_fxn(ctx, k)),
  # gradient on CALL: use provided grad_fxn or auto-differentiate
  (UPat(Ops.CALL, name="k"), call_gradient),
  # there's no gradient for bitcast
@@ -86,7 +86,9 @@ def word_wrap(x, wrap=80):
  while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1
  return x[:i] + "\n" + word_wrap(x[i:], wrap)
 def pad_bytes(b:bytes, align:int) -> bytes: return b + b'\x00' * ((align - (len(b) % align)) % align)
-def panic(e:Exception|None=None): raise e if e is not None else RuntimeError("PANIC!")
+
+# NOTE: you must create the exception inside the function where it's raised or you will get a GC cycle!
+def panic(e:type[Exception]|None=None, *arg): raise e(*arg) if e is not None else RuntimeError("PANIC!")

@functools.cache
 def canonicalize_strides(shape:tuple[T, ...], strides:tuple[T, ...]) -> tuple[T, ...]:
@@ -183,11 +185,12 @@ ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), Conte
 EMULATE, EMULATED_DTYPES = ContextVar("EMULATE", ""), ContextVar("EMULATED_DTYPES", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
 # Compilers
-CPU_LLVM, CPU_LVP, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0), ContextVar("AMD_LLVM", 0)
-NV_PTX, CUDA_PTX, NV_NAK, QCOM_IR3 = ContextVar("NV_PTX", 0), ContextVar("CUDA_PTX", 0), ContextVar("NV_NAK", 0), ContextVar("QCOM_IR3", 0)
+CPU_CC, CPU_LLVM, CPU_LVP = ContextVar("CPU_CC", ""), ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0)
+NV_CC, NV_PTX, NV_NAK = ContextVar("NV_CC", ""), ContextVar("NV_PTX", 0), ContextVar("NV_NAK", 0)
+CUDA_CC, CUDA_PTX, CUDA_NVCC = ContextVar("CUDA_CC", ""), ContextVar("CUDA_PTX", 0), ContextVar("CUDA_NVCC", 0)
 NULL_IR3, NULL_NAK, NULL_ALLOW_COPYOUT = ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0), ContextVar("NULL_ALLOW_COPYOUT", 0)
-AMD_CC, CPU_CC, NV_CC, CUDA_CC = ContextVar("AMD_CC", ""), ContextVar("CPU_CC", ""), ContextVar("NV_CC", ""), ContextVar("CUDA_CC", "")
-QCOM_CC = ContextVar("QCOM_CC", "")
+AMD_CC, AMD_LLVM, AMD_HIPCC  = ContextVar("AMD_CC", ""), ContextVar("AMD_LLVM", 0), ContextVar("AMD_HIPCC", 0)
+QCOM_CC, QCOM_IR3 = ContextVar("QCOM_CC", ""), ContextVar("QCOM_IR3", 0)
 # VIZ implies PROFILE, but you can run PROFILE without VIZ
 VIZ = ContextVar("VIZ", 0)
 PROFILE = ContextVar("PROFILE", abs(VIZ.value))
@@ -287,8 +290,7 @@ class TracingKey:
 class ProfileEvent: pass

@dataclass
-class ProfileRangeEvent(ProfileEvent):
-  device:str; name:str|TracingKey; st:decimal.Decimal; en:decimal.Decimal|None=None; is_copy:bool=False # noqa: E702
+class ProfileRangeEvent(ProfileEvent): device:str; name:str|TracingKey; st:decimal.Decimal; en:decimal.Decimal|None=None # noqa: E702

@dataclass(frozen=True)
 class ProfilePointEvent(ProfileEvent):
@@ -296,8 +298,8 @@ class ProfilePointEvent(ProfileEvent):

 cpu_events:list[ProfileEvent] = []
@contextlib.contextmanager
-def cpu_profile(name:str|TracingKey, device="TINY", is_copy=False, display=True) -> Generator[ProfileRangeEvent, None, None]:
-  res = ProfileRangeEvent(device, name, perf_counter_us(), is_copy=is_copy)
+def cpu_profile(name:str|TracingKey, device="TINY", display=True) -> Generator[ProfileRangeEvent, None, None]:
+  res = ProfileRangeEvent(device, name, perf_counter_us())
  try: yield res
  finally:
    res.en = perf_counter_us()
@@ -360,7 +360,8 @@ class Embedding:

  def __call__(self, idx:Tensor) -> Tensor:
    if not dtypes.is_int(idx.dtype): raise TypeError(f"Expected integer dtype for index in embedding, got {idx.dtype}")
-    if USE_ATOMICS: return Tensor.call(self.weight, idx, fxn=_embedding_fwd(self.weight.as_param(0), idx.as_param(1)), grad_fxn=_embedding_bwd)
+    if USE_ATOMICS:
+      return Tensor.call(self.weight, idx, fxn=_embedding_fwd(self.weight.as_param(0), idx.as_param(1)), grad_fxn=_embedding_bwd, inline=True)
    return _embedding_fwd(self.weight, idx)

 class LSTMCell:
@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Callable, cast, TYPE_CHECKING
+from typing import Callable, cast
 import functools
 from dataclasses import dataclass, field
 from tinygrad.helpers import to_function_name, dedup, prod, DEBUG
@@ -7,7 +7,7 @@ from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, Gro
 from tinygrad.dtype import AddrSpace, PtrDType
 from tinygrad.codegen.opt.tc import TensorCore
 from tinygrad.codegen.opt import Opt
-if TYPE_CHECKING: from tinygrad.device import Compiler
+from tinygrad.device import Compiler

@dataclass(frozen=True)
 class Estimates:
@@ -150,7 +150,8 @@ class Renderer:
  pre_matcher: PatternMatcher|None = None
  extra_matcher: PatternMatcher|None = None
  code_for_op: dict[Ops, Callable] = {}
-  compiler: Compiler|None = None
+
+  compiler: Compiler = Compiler()

  def __reduce__(self): return self.__class__, ()
  def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer")
@@ -340,7 +340,9 @@ class IntelRenderer(OpenCLRenderer):
 class MetalRenderer(CStyleLanguage):
  device = "METAL"
  shared_max = 32768
-  def __init__(self): self.tensor_cores = tc.metal if hasattr(os, 'uname') and os.uname().machine == "arm64" else []
+  def __init__(self):
+    from tinygrad.runtime.ops_metal import MetalCompiler
+    self.compiler, self.tensor_cores = MetalCompiler(), tc.metal if hasattr(os, 'uname') and os.uname().machine == "arm64" else []

  # language options
  kernel_typedef = "kernel void"
@@ -382,15 +384,17 @@ class MetalRenderer(CStyleLanguage):
 _nms = list("xyzwabcdefghijkl") + [f'v{i}' for i in range(16, 32)]

 class CUDARenderer(CStyleLanguage):
-  device = "CUDA"
  global_max = (2147483647, 65535, 65535)
  local_max = (1024, 1024, 64)
  shared_max = 49152

-  def __init__(self, arch:str):
-    self.arch, arch_ver = arch, int(arch[3:])
-    self.tensor_cores = tc.cuda_sm89 if arch_ver >= 89 else tc.cuda_sm80 if arch_ver >= 80 else tc.cuda_sm75 if arch_ver >= 75 else []
-  def __reduce__(self): return self.__class__, (self.arch,)
+  def __init__(self, arch:str, device:str="NV", use_nvcc=False):
+    from tinygrad.runtime.support.compiler_cuda import NVRTCCompiler, NVCCCompiler
+    from tinygrad.runtime.support.hcq import MOCKGPU
+    self.device, self.arch, self.use_nvcc = device, arch, use_nvcc
+    self.compiler = (NVCCCompiler if use_nvcc else NVRTCCompiler)(arch, ptx=bool(MOCKGPU) or device == "CUDA", cache_key=device.lower())
+    self.tensor_cores = tc.cuda_sm89 if (ver:=int(arch[3:])) >= 89 else tc.cuda_sm80 if ver >= 80 else tc.cuda_sm75 if ver >= 75 else []
+  def __reduce__(self): return self.__class__, (self.arch, self.device, self.use_nvcc)

  # language options
  # https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html
@@ -547,7 +551,6 @@ class AMDHIPRenderer(CStyleLanguage):
  for (int n = 0; n < 8; n++) { d[n] = c_frag[n*2]; } return d;\n}""")
    return super().render_kernel(function_name, kernel, bufs, uops, prefix)

-class NVRenderer(CUDARenderer): device = "NV"
 class HIPRenderer(AMDHIPRenderer): device = "HIP"
 class AMDHIPCCRenderer(AMDHIPRenderer):
  def __init__(self, arch:str):
@@ -205,6 +205,9 @@ class CPULLVMRenderer(LLVMRenderer):
  string_rewrite = base_rewrite + PatternMatcher([(UPat(Ops.WMMA, name="wmma"), render_wmma_amx)])
  def render(self, uops: list[UOp]) -> str: return "\n".join((k:=self._render_kernel(uops))[0] + (k[1], self._render_footer(uops)))
  def _render_footer(self, uops: list[UOp]) -> str: return 'attributes #0 = { alwaysinline nounwind "no-builtins" "no-trapping-math"="true" }'
+  def __init__(self):
+    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
+    self.compiler = CPULLVMCompiler()

 barrier = 'fence syncscope("workgroup") release\ntail call void @llvm.amdgcn.s.barrier()\nfence syncscope("workgroup") acquire\n'
 code_for_workitem = {"g": lambda x: f"tail call i32 @llvm.amdgcn.workgroup.id.{chr(120+int(x))}()",
@@ -144,9 +144,11 @@ class PTXRenderer(Renderer):
  tc_sm80 = [x for x in tc.cuda_sm80 if x.dtype_in in [dtypes.half, dtypes.float]]
  code_for_op = asm_for_op
  extra_matcher = ptx_matcher
-  def __init__(self, arch:str, device="CUDA"):
-    self.device, self.arch, arch_ver = device, arch, int(arch[3:])
-    self.tensor_cores = PTXRenderer.tc_sm80 if arch_ver >= 80 else tc.cuda_sm75 if arch_ver >= 75 else []
+  def __init__(self, arch:str, device="NV"):
+    from tinygrad.runtime.support.compiler_cuda import NVPTXCompiler, PTXCompiler
+    from tinygrad.runtime.support.hcq import MOCKGPU
+    self.compiler, self.device, self.arch = (PTXCompiler if bool(MOCKGPU) or device == "CUDA" else NVPTXCompiler)(arch), device, arch
+    self.tensor_cores = PTXRenderer.tc_sm80 if (ver:=int(arch[3:])) >= 80 else tc.cuda_sm75 if ver >= 75 else []
  def __reduce__(self): return self.__class__, (self.arch, self.device)

  # language options
@@ -69,9 +69,12 @@ class WGSLRenderer(CStyleLanguage):
    (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"var {ctx[x]}: array<{ctx.buf_map(x.dtype)},{_packed_size(x.dtype)}>;"),
    (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)),
     lambda ctx,x: f"bitcast<vec2<f16>>({ctx[x.src[0]]})[0]"),
-    (UPat(Ops.BITCAST, dtype=(dtypes.char, dtypes.uchar), name="x"), lambda ctx,x: f"bitcast<{ctx.type_map[x.dtype]}>({ctx[x.src[0]]}&0xFF)"),
-    (UPat(Ops.BITCAST, dtype=(dtypes.short, dtypes.ushort), name="x"),lambda ctx,x:f"bitcast<{ctx.type_map[x.dtype]}>(vec2<f16>({ctx[x.src[0]]},0))" \
-     if x.src[0].dtype == dtypes.half else f"bitcast<{ctx.type_map[x.dtype]}>({ctx[x.src[0]]}&0xFFFF)"),
+    (UPat(Ops.BITCAST, dtype=dtypes.uchar, name="x"), lambda ctx,x: f"bitcast<u32>({ctx[x.src[0]]}&0xFF)"),
+    (UPat(Ops.BITCAST, dtype=dtypes.char, name="x"), lambda ctx,x: f"((i32({ctx[x.src[0]]}&0xFF)<<24)>>24)"),
+    (UPat(Ops.BITCAST, dtype=dtypes.ushort, name="x"), lambda ctx,x: f"bitcast<u32>(vec2<f16>({ctx[x.src[0]]},0))" \
+     if x.src[0].dtype == dtypes.half else f"bitcast<u32>({ctx[x.src[0]]}&0xFFFF)"),
+    (UPat(Ops.BITCAST, dtype=dtypes.short, name="x"), lambda ctx,x: f"bitcast<i32>(vec2<f16>({ctx[x.src[0]]},0))" \
+     if x.src[0].dtype == dtypes.half else f"((i32({ctx[x.src[0]]}&0xFFFF)<<16)>>16)"),
    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast<{ctx.type_map[x.dtype]}>({ctx[x.src[0]]})"),
    # TODO: load alt value doesnt have to be a const
    (UPat.load(UPat.var("b"), UPat.cvar("v"), allow_any_len=True),
@@ -57,7 +57,7 @@ def __getattr__(nm):
      ], args=[
        "-include", "{}/src/common/sdk/nvidia/inc/nvtypes.h", "-I{}/src/common/inc", "-I{}/kernel-open/nvidia-uvm", "-I{}/kernel-open/common/inc",
        "-I{}/src/common/sdk/nvidia/inc", "-I{}/src/nvidia/arch/nvalloc/unix/include", "-I{}/src/common/sdk/nvidia/inc/ctrl"
-      ], rules=[(r'MW\(([^:]+):(.+)\)',r'(\1, \2)')], tarball=nv_src[nm], anon_names={"{}/kernel-open/common/inc/nvstatus.h:37":"nv_status_codes"})
+      ], rules=[(r'MW\(([^:]+):(.+)\)',r'(\1, \2)'), (r'(\d+):(\d+)', r'(\1, \2)')], tarball=nv_src[nm], anon_names={"{}/kernel-open/common/inc/nvstatus.h:37":"nv_status_codes"})
    case "nv": return load("nv", None, [
      *[f"{{}}/src/nvidia/inc/kernel/gpu/{s}.h" for s in ["fsp/kern_fsp_cot_payload", "gsp/gsp_init_args"]],
      *[f"{{}}/src/nvidia/arch/nvalloc/common/inc/{s}.h" for s in ["gsp/gspifpub", "gsp/gsp_fw_wpr_meta", "gsp/gsp_fw_sr_meta", "rmRiscvUcode",
@@ -115,7 +115,7 @@ def __getattr__(nm):
      return load("rocprof", "['rocprof-trace-decoder', p:='/usr/local/lib/rocprof-trace-decoder.so', p.replace('so','dylib')]",
                  [f"{{}}/include/{s}.h" for s in ["rocprof_trace_decoder", "trace_decoder_instrument", "trace_decoder_types"]],
                  tarball="https://github.com/ROCm/rocprof-trace-decoder/archive/dd0485100971522cc4cd8ae136bdda431061a04d.tar.gz")
-    case "mesa": return load("mesa", "['tinymesa_cpu', 'tinymesa']", [
+    case "mesa": return load("mesa", "([] if CPU_CC.value == 'LVP' or bool(CPU_LVP) else ['tinymesa']) + ['tinymesa_cpu']", [
        *[f"{{}}/src/compiler/nir/{s}.h" for s in ["nir", "nir_builder", "nir_shader_compiler_options", "nir_serialize"]], "{}/gen/nir_intrinsics.h",
        *[f"{{}}/src/nouveau/{s}.h" for s in ["headers/nv_device_info", "compiler/nak"]],
        *[f"{{}}/src/gallium/auxiliary/gallivm/lp_bld{s}.h" for s in ["", "_passmgr", "_misc", "_type", "_init", "_nir", "_struct", "_jit_types",
@@ -134,7 +134,8 @@ def __getattr__(nm):
          *[f"python3 src/compiler/{s}_h.py > gen/{s.split('/')[-1]}.h" for s in ["nir/nir_opcodes", "nir/nir_builder_opcodes"]],
          *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
  tarball="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
-  prolog=["import gzip, base64"], epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
+  prolog=["from tinygrad.helpers import CPU_CC, CPU_LVP", "import gzip, base64"],
+  epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
    case "libclang":
      return load("libclang", clang_lib,
                  lambda: [f"{system('llvm-config-20 --includedir')}/clang-c/{s}.h" for s in ["Index", "CXString", "CXSourceLocation", "CXFile"]],
@@ -4,8 +4,9 @@ import ctypes
 from typing import Annotated, Literal, TypeAlias
 from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
 from tinygrad.runtime.support import c
+from tinygrad.helpers import CPU_CC, CPU_LVP
 import gzip, base64
-dll = c.DLL('mesa', ['tinymesa_cpu', 'tinymesa'])
+dll = c.DLL('mesa', ([] if CPU_CC.value == 'LVP' or bool(CPU_LVP) else ['tinymesa']) + ['tinymesa_cpu'])
 class struct_u_printf_info(ctypes.Structure): pass
 u_printf_info: TypeAlias = struct_u_printf_info
 uint32_t: TypeAlias = Annotated[int, ctypes.c_uint32]
@@ -131,7 +131,8 @@ class HCQGraph(MultiGraphRunner):
        # Description based on the command.
        prof_ji_desc = ji.prg._prg.name if is_exec_prg else f"{ji.bufs[1].device} -> {ji.bufs[0].device}" # type: ignore

-        self.prof_graph_entries.append(ProfileGraphEntry(enqueue_dev.device, prof_ji_desc, sig_st, j * 2 + 1, is_copy=not is_exec_prg))
+        prof_name = f"{enqueue_dev.device}:SDMA:{queue_idx}" if not is_exec_prg else enqueue_dev.device
+        self.prof_graph_entries.append(ProfileGraphEntry(prof_name, prof_ji_desc, sig_st, j * 2 + 1))
        self.prof_graph_deps.append([d - 1 for _, d in rdeps])

      last_j[enqueue_queue] = j
@@ -101,7 +101,7 @@ class MetalGraph(GraphRunner):
  def collect_timestamps(self):
    # create a graph event and evenly space each program
    st, en = decimal.Decimal(self.command_buffer.GPUStartTime()) * 1000000, decimal.Decimal(self.command_buffer.GPUEndTime()) * 1000000
-    ents = [ProfileGraphEntry(self.device, cast(CompiledRunner, ji.prg)._prg.name, i, i+1, is_copy=False) for i,ji in enumerate(self.jit_cache)]
+    ents = [ProfileGraphEntry(self.device, cast(CompiledRunner, ji.prg)._prg.name, i, i+1) for i,ji in enumerate(self.jit_cache)]
    step = (en-st)/len(ents)
    self.dev.profile_events += [ProfileGraphEvent(ents, [], [st+step*i for i in range(len(ents)+1)])]

@@ -6,9 +6,9 @@ from dataclasses import dataclass
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, FileIOInterface
 from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filter_visible_devices
 from tinygrad.uop.ops import sint
-from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerSet
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar
-from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, ceildiv
+from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, AMD_HIPCC, ceildiv, unwrap
 from tinygrad.renderer.cstyle import AMDHIPRenderer, AMDHIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt, amdgpu_kd, amdgpu_drm
@@ -703,13 +703,15 @@ class KFDIface:
    # Event to wait for queues completion
    self.dev.queue_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_SIGNAL, auto_reset=1)
    self.dev.queue_event_mailbox_ptr = KFDIface.event_page.va_addr + self.dev.queue_event.event_slot_index * 8
-    self.queue_event_arr = (kfd.struct_kfd_event_data)(event_id=self.dev.queue_event.event_id)
-    self.queue_event_arr_ptr = ctypes.addressof(self.queue_event_arr)

    # OS events to collect memory and hardware faults
    self.mem_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_MEMORY)
    self.hw_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_HW_EXCEPTION)

+    self.queue_event_arr = (kfd.struct_kfd_event_data * 3)(kfd.struct_kfd_event_data(event_id=self.dev.queue_event.event_id),
+      kfd.struct_kfd_event_data(event_id=self.mem_fault_event.event_id), kfd.struct_kfd_event_data(event_id=self.hw_fault_event.event_id))
+    self.queue_event_arr_ptr = ctypes.addressof(self.queue_event_arr)
+
  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> HCQBuffer:
    flags = kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE

@@ -778,18 +780,20 @@ class KFDIface:
                        doorbell=MMIOInterface(self.doorbells + queue.doorbell_offset - self.doorbells_base, 8, fmt='Q'))

  def sleep(self, tm:int):
-    kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=self.queue_event_arr_ptr, num_events=1, wait_for_all=1, timeout=tm)
+    kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=self.queue_event_arr_ptr, num_events=3, wait_for_all=0, timeout=tm)
+    if self.queue_event_arr[1].memory_exception_data.gpu_id or self.queue_event_arr[2].hw_exception_data.gpu_id: raise RuntimeError("Device fault")

  def on_device_hang(self):
-    def _collect_str(st): return ' '.join(f'{k[0]}={getattr(st, k[0])}' for k in st._real_fields_)
+    def _str(st): return ' '.join(f'{k[0]}={getattr(st, k[0])}' for k in st._real_fields_)
+
+    # try to collect fault info if not already set from sleep().
+    if not self.queue_event_arr[1].memory_exception_data.gpu_id and not self.queue_event_arr[2].hw_exception_data.gpu_id:
+      with contextlib.suppress(RuntimeError): self.sleep(tm=1)

    report = []
-    for evnt in [self.mem_fault_event, self.hw_fault_event]:
-      ev = (kfd.struct_kfd_event_data)(event_id=evnt.event_id)
-      kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=ctypes.addressof(ev), num_events=1, wait_for_all=1)
-      if evnt == self.mem_fault_event and ev.memory_exception_data.gpu_id:
-        report += [f"MMU fault: 0x{ev.memory_exception_data.va:X} | {_collect_str(ev.memory_exception_data.failure)}"]
-      if evnt == self.hw_fault_event and ev.hw_exception_data.gpu_id: report += [f"HW fault: {_collect_str(ev.hw_exception_data)}"]
+    if self.queue_event_arr[1].memory_exception_data.gpu_id:
+      report += [f"MMU fault: 0x{self.queue_event_arr[1].memory_exception_data.va:X} | {_str(self.queue_event_arr[1].memory_exception_data.failure)}"]
+    if self.queue_event_arr[2].hw_exception_data.gpu_id: report += [f"HW fault: {_str(self.queue_event_arr[2].hw_exception_data)}"]

    raise RuntimeError("\n".join(report))

@@ -955,11 +959,12 @@ class AMDDevice(HCQCompiled):
      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size, debug_memory_size=debug_memory_size)

    self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
+    self.sdma_queues:dict = {}
    self.has_sdma_queue = self.sdma_queue(0) is not None

-    compilers = CompilerSet([CompilerPair(functools.partial(AMDHIPRenderer, self.arch), None),
-                             CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), None, AMD_LLVM),
-                             CompilerPair(functools.partial(AMDHIPCCRenderer, self.arch), None)], ctrl_var=AMD_CC)
+    compilers = CompilerSet([(functools.partial(AMDHIPRenderer, self.arch), None),
+                             (functools.partial(AMDLLVMRenderer, self.arch), AMD_LLVM),
+                             (functools.partial(AMDHIPCCRenderer, self.arch), AMD_HIPCC)], ctrl_var=AMD_CC)

    super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
                     functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),
@@ -1018,11 +1023,12 @@ class AMDDevice(HCQCompiled):
            wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
            ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx))

-  @functools.lru_cache(None)
  def sdma_queue(self, idx:int):
    if getenv("AMD_DISABLE_SDMA"): return None
-    with contextlib.suppress(OSError): return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
-    return None
+    if idx in self.sdma_queues: return self.sdma_queues[idx]
+    with contextlib.suppress(OSError):
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
+    return self.sdma_queues.get(idx, None)

  def _ensure_has_local_memory(self, private_segment_size):
    if self.max_private_segment_size >= private_segment_size: return
@@ -1063,3 +1069,5 @@ class AMDDevice(HCQCompiled):
  def on_device_hang(self): self.iface.on_device_hang()

  def device_props(self): return self.iface.props
+
+  def hw_copy_queues(self): return [(f"SDMA:{i}", functools.partial(unwrap(self.hw_copy_queue_t), queue_idx=i)) for i in self.sdma_queues]
@@ -5,7 +5,7 @@ from tinygrad.runtime.autogen import opencl as cl
 from tinygrad.runtime.support import c
 from tinygrad.helpers import to_char_p_p, from_mv, OSX, DEBUG, mv_address, suppress_finalizing
 from tinygrad.renderer.cstyle import OpenCLRenderer, IntelRenderer
-from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError, CompilerPair, CompilerSet
+from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError, CompilerSet
 from tinygrad.dtype import ImageDType

 CC_CB = c.CFUNCTYPE[None, [c.POINTER[ctypes.c_char], c.POINTER[None], cl.size_t, c.POINTER[None]]]
@@ -39,9 +39,9 @@ class CLCompiler(Compiler):

 class CLProgram:
  def __init__(self, device:CLDevice, name:str, lib:bytes, buf_dtypes=[], **kwargs):
-    self.dev, self.name, self.lib, self.buf_dtypes = device, name, lib, buf_dtypes
-    self.program = checked(cl.clCreateProgramWithBinary(device.context, 1, device.device_id, (ctypes.c_size_t * 1)(len(lib)),
-                                                        to_char_p_p([lib], ctypes.c_ubyte), binary_status := ctypes.c_int32(),
+    self.dev, self.name, self.lib, self.buf_dtypes = device, name, device.cl_compiler.compile_cached(lib.decode()), buf_dtypes
+    self.program = checked(cl.clCreateProgramWithBinary(device.context, 1, device.device_id, (ctypes.c_size_t * 1)(len(self.lib)),
+                                                        to_char_p_p([self.lib], ctypes.c_ubyte), binary_status := ctypes.c_int32(),
                                                        errcode_ret := ctypes.c_int32()), errcode_ret)
    check(binary_status.value)
    check(cl.clBuildProgram(self.program, 1, device.device_id, None, BP_CB(), None)) # NOTE: OSX requires this
@@ -125,8 +125,9 @@ class CLDevice(Compiled):
                                           ctypes.string_at(buf, size=total.value).decode())[1]

    renderer = IntelRenderer if "cl_intel_subgroup_matrix_multiply_accumulate" in self.device_exts else OpenCLRenderer
-    compiler = functools.partial(CLCompiler, self, f"{hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest()}")
-    super().__init__(device, CLAllocator(self), CompilerSet([CompilerPair(renderer, compiler)]), functools.partial(CLProgram, self))
+    self.cl_compiler = CLCompiler(self, f"{hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest()}")
+    super().__init__(device, CLAllocator(self), CompilerSet([(renderer, None)]), functools.partial(CLProgram, self))
+
  def synchronize(self):
    check(cl.clFinish(self.queue))
    self.pending_copyin.clear()
@@ -2,13 +2,12 @@ from __future__ import annotations
 import platform, sys, ctypes, functools, time, mmap, threading, queue
 from tinygrad.helpers import to_mv, OSX, WIN, mv_address, wait_cond, suppress_finalizing, unwrap, data64_le
 from tinygrad.helpers import CPU_CC, CPU_LVP, CPU_LLVM
-from tinygrad.device import BufferSpec, DMACPURef, CompilerSet, CompilerPair
+from tinygrad.device import BufferSpec, DMACPURef, CompilerSet
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.runtime.support.hcq import CLikeArgsState
 from tinygrad.renderer.cstyle import ClangJITRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
 from tinygrad.renderer.nir import LVPRenderer
-from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
 from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.uop.ops import sint

@@ -134,6 +133,5 @@ class CPUDevice(HCQCompiled):
  def __init__(self, device:str=""):
    self.tasks:queue.Queue = queue.Queue()
    CPUWorker(self, self.tasks, thread_id=0).start()
-    compilers = CompilerSet([CompilerPair(ClangJITRenderer, None), CompilerPair(CPULLVMRenderer, CPULLVMCompiler, ctrl_var=CPU_LLVM),
-                             CompilerPair(LVPRenderer, None, ctrl_var=CPU_LVP)], ctrl_var=CPU_CC)
+    compilers = CompilerSet([(ClangJITRenderer, None), (CPULLVMRenderer, CPU_LLVM), (LVPRenderer, CPU_LVP)], ctrl_var=CPU_CC)
    super().__init__(device, CPUAllocator(self), compilers, functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue)
@@ -1,11 +1,11 @@
 from __future__ import annotations
 import ctypes, functools
-from tinygrad.helpers import DEBUG, getenv, mv_address, suppress_finalizing, CUDA_CC, CUDA_PTX
-from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerPair, CompilerSet
+from tinygrad.helpers import DEBUG, getenv, mv_address, suppress_finalizing, CUDA_CC, CUDA_PTX, CUDA_NVCC
+from tinygrad.device import Compiled, BufferSpec, LRUAllocator, CompilerSet
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.runtime.autogen import cuda
-from tinygrad.runtime.support.compiler_cuda import pretty_ptx, CUDACompiler, PTXCompiler, NVCCCompiler
+from tinygrad.runtime.support.compiler_cuda import pretty_ptx
 from tinygrad.runtime.support.c import init_c_struct_t, init_c_var
 if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl  # noqa: F401  # pylint: disable=unused-import
 if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore # pylint: disable=reimported
@@ -118,9 +118,9 @@ class CUDADevice(Compiled):
    CUDADevice.devices.append(self)

    from tinygrad.runtime.graph.cuda import CUDAGraph
-    compilers = CompilerSet([CompilerPair(functools.partial(CUDARenderer, self.arch), functools.partial(CUDACompiler, self.arch)),
-                             CompilerPair(functools.partial(PTXRenderer, self.arch), functools.partial(PTXCompiler, self.arch), CUDA_PTX),
-                             CompilerPair(functools.partial(CUDARenderer, self.arch), functools.partial(NVCCCompiler, self.arch))], ctrl_var=CUDA_CC)
+    compilers = CompilerSet([(functools.partial(CUDARenderer, self.arch, device="CUDA"), None),
+                             (functools.partial(PTXRenderer, self.arch, device="CUDA"), CUDA_PTX),
+                             (functools.partial(CUDARenderer, self.arch, device="CUDA", use_nvcc=True), CUDA_NVCC)], ctrl_var=CUDA_CC)
    super().__init__(device, CUDAAllocator(self), compilers, functools.partial(CUDAProgram, self), None if MOCKGPU else CUDAGraph)

  def synchronize(self):
@@ -1,7 +1,7 @@
 from __future__ import annotations
 import ctypes, os, mmap, tempfile, pathlib, array, functools, threading, contextlib, sys, subprocess, struct
 assert sys.platform != 'win32'
-from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler, CompilerSet, CompilerPair
+from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler, CompilerSet
 from tinygrad.dtype import dtypes, DType, PtrDType
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, system, DEBUG, suppress_finalizing
@@ -45,6 +45,8 @@ class DSPRenderer(ClangRenderer):
  type_map = { **ClangRenderer.type_map, dtypes.uint64: "unsigned long long", dtypes.int64: "long long" }
  code_for_op = {k:v for k,v in ClangRenderer.code_for_op.items() if k != Ops.SQRT}

+  def __init__(self): self.compiler = DSPCompiler()
+
  def _render_defines(self, uops) -> list[str]:
    return ['''/* DSP boilerplate */ struct dcvs_v2_req { int type; int _pad; _Bool dcvs_enable; char dcvs_option; _Bool set_latency; int latency;
      _Bool set_dcvs_params; short _pad2; char target_corner; char min_corner; char max_corner; int _pad3[3];};''','int HAP_power_set(void*, void*);',
@@ -116,28 +118,11 @@ class DSPAllocator(Allocator['DSPDevice']):
  def _copyout(self, dest:memoryview, src:DSPBuffer): ctypes.memmove(mv_address(dest), src.va_addr, dest.nbytes)
  def _offset(self, buf, size:int, offset:int): return DSPBuffer(buf.va_addr+offset, size, buf.share_info, buf.offset+offset)

-class ClangCompiler(Compiler):
-  def __init__(self, cachekey="compile_clang", args:list[str]|None=None, objdump_tool='objdump'):
-    self.args = ['-shared', '-march=native'] if args is None else args
-    self.objdump_tool = objdump_tool
-    super().__init__(cachekey)
-
-  def compile(self, src:str) -> bytes:
-    # TODO: remove file write. sadly clang doesn't like the use of /dev/stdout here
-    with tempfile.NamedTemporaryFile(delete=True) as f:
-      system(f"{getenv('CC','clang')} {' '.join(self.args)} -O2 -Wall -Werror -x c -fPIC -ffreestanding -nostdlib - -o {f.name}", input=src.encode())
-      return pathlib.Path(f.name).read_bytes()
-
-  def disassemble(self, lib:bytes): return cpu_objdump(lib, self.objdump_tool)
-
-class DSPDevice(Compiled):
-  def __init__(self, device:str=""):
-    compiler_args = ["--target=hexagon", "-mcpu=hexagonv65", "-fuse-ld=lld", "-nostdlib",  "-mhvx=v65", "-mhvx-length=128b"]
-    if getenv("MOCKDSP"):
-      mock_compilers = CompilerSet([CompilerPair(MockDSPRenderer, functools.partial(ClangCompiler, None, ["-static"]+compiler_args, 'llvm-objdump'))])
-      super().__init__(device, DSPAllocator(self), mock_compilers, MockDSPProgram)
+class DSPCompiler(Compiler):
+  def __init__(self, mock:bool=False):
+    compiler_args = "--target=hexagon -mcpu=hexagonv65 -fuse-ld=lld -nostdlib -mhvx=v65 -mhvx-length=128b"
+    if mock: self.args = f"-static {compiler_args}"
    else:
-      self.ion_fd = os.open('/dev/ion', os.O_RDONLY)
      # Generate link script to pass into clang. Aligning all used sections to 4k fixes invoke problem.
      sections = ['text', 'rela.plt', 'rela.dyn', 'plt', 'data', 'bss', 'hash', 'dynamic',
                  'got', 'got.plt', 'dynsym', 'dynstr', 'symtab', 'shstrtab', 'strtab']
@@ -146,8 +131,25 @@ class DSPDevice(Compiled):
        self.link_ld.write(f"SECTIONS {{ . = 0x0; {sections_link}\n /DISCARD/ : {{ *(.note .note.* .gnu.hash .comment) }} }}".encode())
        self.link_ld.flush()

-      compiler = functools.partial(ClangCompiler, "compile_dsp", ["-shared"] + compiler_args + [f"-T{self.link_ld.name}"], 'llvm-objdump')
-      super().__init__(device, DSPAllocator(self), CompilerSet([CompilerPair(DSPRenderer, compiler)]), functools.partial(DSPProgram, self))
+      self.args = f"-shared {compiler_args} -T{self.link_ld.name}"
+
+    super().__init__(None if mock else "compile_dsp")
+
+  def compile(self, src:str) -> bytes:
+    # TODO: remove file write. sadly clang doesn't like the use of /dev/stdout here
+    with tempfile.NamedTemporaryFile(delete=True) as f:
+      system(f"{getenv('CC','clang')} {self.args} -O2 -Wall -Werror -x c -fPIC -ffreestanding -nostdlib - -o {f.name}", input=src.encode())
+      return pathlib.Path(f.name).read_bytes()
+
+  def disassemble(self, lib:bytes): return cpu_objdump(lib, "llvm-objdump")
+
+
+class DSPDevice(Compiled):
+  def __init__(self, device:str=""):
+    if getenv("MOCKDSP"): super().__init__(device, DSPAllocator(self), CompilerSet([(MockDSPRenderer, None)]), MockDSPProgram)
+    else:
+      self.ion_fd = os.open('/dev/ion', os.O_RDONLY)
+      super().__init__(device, DSPAllocator(self), CompilerSet([(DSPRenderer, None)]), functools.partial(DSPProgram, self))
      fastrpc_shell = memoryview(bytearray(pathlib.Path('/dsp/cdsp/fastrpc_shell_3').read_bytes()))
      self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True))
      ctypes.memmove(self.shell_buf.va_addr, mv_address(fastrpc_shell), fastrpc_shell.nbytes)
@@ -268,6 +270,7 @@ static void *mmap2(void *addr, unsigned int length, int prot, int flags, int fd,
 return (void*)syscall((long)addr, length, prot, flags, fd, offset, 222); }}'''

 class MockDSPRenderer(DSPRenderer):
+  def __init__(self): self.compiler = DSPCompiler(mock=True)
  def _render_defines(self, uops) -> list[str]: return ClangRenderer._render_defines(self, uops)
  def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[DType,bool]]]) -> str:
    # https://gpages.juszkiewicz.com.pl/syscalls-table/syscalls.html
@@ -1,6 +1,6 @@
 import ctypes, functools
 from tinygrad.helpers import mv_address, getenv, suppress_finalizing
-from tinygrad.device import Compiled, LRUAllocator, BufferSpec, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, LRUAllocator, BufferSpec, CompilerSet
 from tinygrad.runtime.autogen import hip
 from tinygrad.renderer.cstyle import HIPRenderer
 from tinygrad.runtime.support.c import init_c_var, init_c_struct_t
@@ -15,7 +15,7 @@ class HIPDevice(Compiled):
    self.arch = init_c_var(hip.hipDeviceProp_t, lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode()
    self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t, lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)]

-    compilers = CompilerSet([CompilerPair(functools.partial(HIPRenderer, self.arch), None)])
+    compilers = CompilerSet([(functools.partial(HIPRenderer, self.arch), None)])
    super().__init__(device, HIPAllocator(self), compilers, functools.partial(HIPProgram, self))
  def synchronize(self):
    check(hip.hipSetDevice(self.device_id))
@@ -1,7 +1,7 @@
-import subprocess, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform
-from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, PROFILE, ProfileRangeEvent, cpu_profile, unwrap, suppress_finalizing
+import subprocess, pathlib, struct, ctypes, tempfile, functools, decimal, platform
+from tinygrad.helpers import prod, to_mv, round_up, cache_dir, PROFILE, ProfileRangeEvent, cpu_profile, unwrap, suppress_finalizing
 import tinygrad.runtime.support.objc as objc
-from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent, CompilerSet
 from tinygrad.renderer.cstyle import MetalRenderer
 from tinygrad.runtime.autogen import metal
 from tinygrad.runtime.support.c import DLL
@@ -42,7 +42,7 @@ class MetalDevice(Compiled):
    from tinygrad.runtime.graph.metal import MetalGraph
    # NOTE: GitHub CI macOS runners use paravirtualized metal which is broken with graph.
    # This can be reproduced locally with any virtualization software (like utm) that can create macOS VMs with apple's own virtualization framework.
-    super().__init__(device, MetalAllocator(self), CompilerSet([CompilerPair(MetalRenderer, MetalCompiler), CompilerPair(MetalRenderer, Compiler)]),
+    super().__init__(device, MetalAllocator(self), CompilerSet([(MetalRenderer, None)]),
      functools.partial(MetalProgram, self), MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None)

  def synchronize(self):
@@ -51,23 +51,15 @@ class MetalDevice(Compiled):
      st, en = decimal.Decimal(cbuf.GPUStartTime()) * 1000000, decimal.Decimal(cbuf.GPUEndTime()) * 1000000
      # NOTE: command buffers from MetalGraph are not profiled here
      if PROFILE and (lb:=cmdbuf_label(cbuf)) is not None and not lb.startswith("batched"):
-        Compiled.profile_events += [ProfileRangeEvent(self.device, lb, st, en, is_copy=lb.startswith("COPY"))]
+        Compiled.profile_events += [ProfileRangeEvent(self.device, lb, st, en)]
    self.mtl_buffers_in_flight.clear()

-def metal_src_to_library(device:MetalDevice, src:str) -> metal.MTLLibrary:
-  options = metal.MTLCompileOptions.new()
-  options.setFastMathEnabled(getenv("METAL_FAST_MATH"))
-  library = device.sysdevice.newLibraryWithSource_options_error(to_ns_str(src), options, ctypes.byref(compileError:=metal.NSError().retained()))
-  error_check(compileError, CompileError)
-  return library
-
 class MetalCompiler(Compiler):
  # Opening METAL after LLVM doesn't fail because ctypes.CDLL opens with RTLD_LOCAL but MTLCompiler opens it's own llvm with RTLD_GLOBAL
  # This means that MTLCompiler's llvm will create it's own instances of global state because RTLD_LOCAL doesn't export symbols, but if RTLD_GLOBAL
  # library is loaded first then RTLD_LOCAL library will just use it's symbols. On linux there is RTLD_DEEPBIND to prevent that, but on macos there
  # doesn't seem to be anything we can do.
-  with contextlib.suppress(FileNotFoundError, ModuleNotFoundError):
-    import tinygrad.runtime.autogen.llvm # noqa: F401
+  import tinygrad.runtime.autogen.llvm as _
  support = DLL("MTLCompiler", "MTLCompiler")
  support.MTLCodeGenServiceCreate.restype = ctypes.c_void_p

@@ -118,15 +110,9 @@ class MetalCompiler(Compiler):
 class MetalProgram:
  def __init__(self, dev:MetalDevice, name:str, lib:bytes, **kwargs):
    self.dev, self.name, self.lib = dev, name, lib
-    if lib[:4] == b"MTLB":
-      # binary metal library
-      data = objc.dispatch_data_create(lib, len(lib), None, None)
-      self.library = self.dev.sysdevice.newLibraryWithData_error(data, ctypes.byref(error_lib:=metal.NSError().retained())).retained()
-      error_check(error_lib)
-    else:
-      # metal source. rely on OS caching
-      try: self.library = metal_src_to_library(self.dev, lib.decode())
-      except CompileError as e: raise RuntimeError from e
+    data = objc.dispatch_data_create(lib, len(lib), None, None)
+    self.library = self.dev.sysdevice.newLibraryWithData_error(data, ctypes.byref(error_lib:=metal.NSError().retained())).retained()
+    error_check(error_lib)
    self.fxn = self.library.newFunctionWithName(to_ns_str(name)).retained()
    descriptor = metal.MTLComputePipelineDescriptor.new()
    descriptor.setComputeFunction(self.fxn)
@@ -191,7 +177,7 @@ class MetalAllocator(LRUAllocator[MetalDevice]):
    # There is no real metal multidevice support for now, so transfer is used only for tests.
    src_dev.synchronize()
  def _cp_mv(self, dst, src, prof_desc):
-    with cpu_profile(prof_desc, self.dev.device, is_copy=True): dst[:] = src
+    with cpu_profile(prof_desc, self.dev.device): dst[:] = src
  def _as_buffer(self, src:MetalBuffer) -> memoryview:
    self.dev.synchronize()
    return to_mv(src.buf.contents(), src.size + src.offset)[src.offset:]
@@ -1,5 +1,5 @@
 import functools
-from tinygrad.device import Compiled, Compiler, Allocator, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, Allocator, CompilerSet
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.renderer.cstyle import Renderer, CStyleLanguage, AMDHIPRenderer
 from tinygrad.uop.ops import Ops
@@ -39,6 +39,6 @@ class NullDevice(Compiled):
      case "AMD_CDNA4": renderer = functools.partial(AMDHIPRenderer, "gfx950")
      case "": renderer = NullRenderer
      case _: raise RuntimeError(f"can't EMULATE device: {EMULATE.value}")
-    compilers = CompilerSet([CompilerPair(renderer, Compiler), CompilerPair(functools.partial(IR3Renderer, 0x6030001), None, NULL_IR3), # adreno 630
-                             CompilerPair(functools.partial(NAKRenderer, "sm_120", 48), None, NULL_NAK)]) # 5090
+    compilers = CompilerSet([(renderer, None), (functools.partial(IR3Renderer, 0x6030001), NULL_IR3), # adreno 630
+                             (functools.partial(NAKRenderer, "sm_120", 48), NULL_NAK)]) # 5090
    super().__init__(device, NullAllocator(self), compilers, functools.partial(NullProgram, device), NullGraph)
@@ -6,12 +6,11 @@ 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
 from tinygrad.uop.ops import sint
-from tinygrad.device import Compiled, BufferSpec, CompilerPair, CompilerSet
+from tinygrad.device import Compiled, BufferSpec, CompilerSet
 from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32, NV_CC, NV_PTX, NV_NAK, PROFILE
 from tinygrad.helpers import ContextVar, VIZ, ProfileEvent
 from tinygrad.renderer.ptx import PTXRenderer
-from tinygrad.renderer.cstyle import NVRenderer
-from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, NVPTXCompiler, NVCompiler
+from tinygrad.renderer.cstyle import CUDARenderer
 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
@@ -36,6 +35,9 @@ def get_error_str(status): return f"{status}: {nv_gpu.nv_status_codes.get(status
 NV_PFAULT_FAULT_TYPE = {dt:name for name,dt in nv_gpu.__dict__.items() if name.startswith("NV_PFAULT_FAULT_TYPE_")}
 NV_PFAULT_ACCESS_TYPE = {dt:name.split("_")[-1] for name,dt in nv_gpu.__dict__.items() if name.startswith("NV_PFAULT_ACCESS_TYPE_")}

+def nv_flags(reg, **kwargs): return functools.reduce(int.__or__, ((getattr(nv_gpu, f"{reg}_{k}_{v}".upper()) if isinstance(v, str) else v) <<
+  getattr(nv_gpu, f"{reg}_{k}".upper())[1] for k, v in kwargs.items()), 0)
+
 def nv_iowr(fd:FileIOInterface, nr, args, cmd=None):
  ret = fd.ioctl(cmd or ((3 << 30) | (ctypes.sizeof(args) & 0x1FFF) << 16 | (ord('F') & 0xFF) << 8 | (nr & 0xFF)), args)
  if ret != 0: raise RuntimeError(f"ioctl returned {ret}")
@@ -94,7 +96,8 @@ class NVCommandQueue(HWQueue[HCQSignal, 'NVDevice', 'NVProgram', 'NVArgsState'])
    return self

  def wait(self, signal:HCQSignal, value:sint=0):
-    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(signal.value_addr), *data64_le(value), (3 << 0) | (1 << 24)) # ACQUIRE | PAYLOAD_SIZE_64BIT
+    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(signal.value_addr), *data64_le(value),
+             nv_flags("NVC56F_SEM_EXECUTE", operation="acq_circ_geq", payload_size="64bit"))
    self.active_qmd = None
    return self

@@ -125,7 +128,8 @@ class NVCommandQueue(HWQueue[HCQSignal, 'NVDevice', 'NVProgram', 'NVArgsState'])

 class NVComputeQueue(NVCommandQueue):
  def memory_barrier(self):
-    self.nvm(1, nv_gpu.NVC6C0_INVALIDATE_SHADER_CACHES_NO_WFI, (1 << 12) | (1 << 4) | (1 << 0))
+    self.nvm(1, nv_gpu.NVC6C0_INVALIDATE_SHADER_CACHES_NO_WFI,
+             nv_flags("NVC6C0_INVALIDATE_SHADER_CACHES_NO_WFI", instruction="true", global_data="true", constant="true"))
    self.active_qmd:QMD|None = None
    return self

@@ -169,7 +173,7 @@ class NVComputeQueue(NVCommandQueue):
          return self

    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(signal.value_addr), *data64_le(value),
-             (1 << 0) | (1 << 20) | (1 << 24) | (1 << 25)) # RELEASE | RELEASE_WFI | PAYLOAD_SIZE_64BIT | RELEASE_TIMESTAMP
+             nv_flags("NVC56F_SEM_EXECUTE", operation="release", release_wfi="en", payload_size="64bit", release_timestamp="en"))
    self.nvm(0, nv_gpu.NVC56F_NON_STALL_INTERRUPT, 0x0)
    self.active_qmd = None
    return self
@@ -185,12 +189,13 @@ class NVCopyQueue(NVCommandQueue):
    for off in range(0, copy_size, step:=(1 << 31)):
      self.nvm(4, nv_gpu.NVC6B5_OFFSET_IN_UPPER, *data64(src+off), *data64(dest+off))
      self.nvm(4, nv_gpu.NVC6B5_LINE_LENGTH_IN, min(copy_size-off, step))
-      self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA, 0x182) # TRANSFER_TYPE_NON_PIPELINED | DST_MEMORY_LAYOUT_PITCH | SRC_MEMORY_LAYOUT_PITCH
+      self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA,
+               nv_flags("NVC6B5_LAUNCH_DMA", data_transfer_type="non_pipelined", src_memory_layout="pitch", dst_memory_layout="pitch"))
    return self

  def signal(self, signal:HCQSignal, value:sint=0):
    self.nvm(4, nv_gpu.NVC6B5_SET_SEMAPHORE_A, *data64(signal.value_addr), value)
-    self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA, 0x14)
+    self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA, nv_flags("NVC6B5_LAUNCH_DMA", flush_enable="true", semaphore_type="release_four_word_semaphore"))
    return self

  def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.dma_gpfifo)
@@ -199,7 +204,8 @@ class NVVideoQueue(NVCommandQueue):
  def decode_hevc_chunk(self, pic_desc:HCQBuffer, in_buf:HCQBuffer, out_buf:HCQBuffer, out_buf_pos:int, hist_bufs:list[HCQBuffer], hist_pos:list[int],
                        chroma_off:int, coloc_buf:HCQBuffer, filter_buf:HCQBuffer, intra_top_off:int, intra_unk_off:int|None, status_buf:HCQBuffer):
    self.nvm(4, nv_gpu.NVC9B0_SET_APPLICATION_ID, nv_gpu.NVC9B0_SET_APPLICATION_ID_ID_HEVC)
-    self.nvm(4, nv_gpu.NVC9B0_SET_CONTROL_PARAMS, 0x52057)
+    self.nvm(4, nv_gpu.NVC9B0_SET_CONTROL_PARAMS, nv_flags("NVC9B0_SET_CONTROL_PARAMS", codec_type="hevc", testrun_env="prod_run", gptimer_on=1,
+             err_conceal_on=1, mbtimer_on=1, event_trace_logging_on=1))
    self.nvm(4, nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET, pic_desc.va_addr >> 8)
    self.nvm(4, nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET, in_buf.va_addr >> 8)
    for pos, buf in zip(hist_pos + [out_buf_pos], hist_bufs + [out_buf]):
@@ -216,7 +222,7 @@ class NVVideoQueue(NVCommandQueue):

  def signal(self, signal:HCQSignal, value:sint=0):
    self.nvm(4, nv_gpu.NVC9B0_SEMAPHORE_A, *data64(signal.value_addr), value)
-    self.nvm(4, nv_gpu.NVC9B0_SEMAPHORE_D, (1 << 24) | (1 << 0))
+    self.nvm(4, nv_gpu.NVC9B0_SEMAPHORE_D, nv_flags("NVC9B0_SEMAPHORE_D", structure_size="four", payload_size="64bit"))
    return self

  def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.vid_gpfifo)
@@ -333,7 +339,7 @@ class NVAllocator(HCQAllocator['NVDevice']):
    self.dev._ensure_has_vid_hw(w, h)

    q = NVVideoQueue().wait(self.dev.timeline_signal, self.dev.timeline_value - 1)
-    with hcq_profile(self.dev, queue=q, desc="NVDEC", enabled=PROFILE):
+    with hcq_profile(self.dev, queue=q, desc="HEVC Decode", enabled=PROFILE, dev_suff="NVDEC"):
      q.decode_hevc_chunk(desc_buf, bufin, bufout, frame_pos, hist, [(frame_pos-x) % (len(hist) + 1) for x in range(len(hist), 0, -1)],
                          round_up(w, 64)*round_up(h, 64), self.dev.vid_coloc_buf, self.dev.vid_filter_buf, self.dev.intra_top_off,
                          self.dev.intra_unk_off, self.dev.vid_stat_buf)
@@ -613,10 +619,9 @@ class NVDevice(HCQCompiled[NVSignal]):
    self.arch: str = "sm_120" if self.sm_version==0xa04 else f"sm_{(self.sm_version>>8)&0xff}{(val>>4) if (val:=self.sm_version&0xff) > 0xf else val}"
    self.sass_version = ((self.sm_version & 0xf00) >> 4) | (self.sm_version & 0xf)

-    cucc, ptxcc = (CUDACompiler, PTXCompiler) if MOCKGPU else (NVCompiler, NVPTXCompiler)
-    compilers = CompilerSet(ctrl_var=NV_CC, cset=[CompilerPair(functools.partial(NVRenderer, self.arch),functools.partial(cucc, self.arch)),
-       CompilerPair(functools.partial(PTXRenderer, self.arch, device="NV"), functools.partial(ptxcc, self.arch), NV_PTX),
-       CompilerPair(functools.partial(NAKRenderer, self.arch, self.max_warps_per_sm), None, NV_NAK)])
+    compilers = CompilerSet(ctrl_var=NV_CC, cset=[(functools.partial(CUDARenderer, self.arch), None),
+       (functools.partial(PTXRenderer, self.arch, device="NV"), NV_PTX),
+       (functools.partial(NAKRenderer, self.arch, self.max_warps_per_sm), NV_NAK)])
    super().__init__(device, NVAllocator(self), compilers, functools.partial(NVProgram, self), NVSignal, NVComputeQueue, NVCopyQueue)

    self.pma_enabled = PMA.value > 0 and PROFILE >= 1
@@ -6,7 +6,7 @@ from typing import Any, TYPE_CHECKING
 import pickle, base64, itertools, time, struct, sys, functools
 from tinygrad.dtype import DType, dtypes, ImageDType, PtrDType, truncate, float_to_fp16, float_to_bf16, float_to_fp8, fp8_to_float
 from tinygrad.helpers import all_same, getenv, flatten, get_single_element, EMULATE
-from tinygrad.device import Compiled, Compiler, Allocator, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, Compiler, Allocator, CompilerSet
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
 from tinygrad.renderer import Renderer
@@ -214,9 +214,14 @@ class PythonProgram:
        i += 1
    return time.perf_counter() - st

+class PythonCompiler(Compiler):
+  def compile(self, src:str) -> bytes: return base64.b64decode(src)
+
 class PythonRenderer(Renderer):
  device = "PYTHON"
  code_for_op = python_alu
+  compiler = PythonCompiler()
+
  def __init__(self):
    match EMULATE.value:
      case "METAL": self.device, self.tensor_cores = "METAL", tc.metal
@@ -236,9 +241,6 @@ class PythonRenderer(Renderer):
    lops = [(u.op, u.dtype, [uops.index(v) for v in u.src if u.op is not Ops.SPECIAL], u.arg) for u in uops]
    return base64.b64encode(pickle.dumps(lops)).decode()

-class PythonCompiler(Compiler):
-  def compile(self, src:str) -> bytes: return base64.b64decode(src)
-
 class PythonAllocator(Allocator['PythonDevice']):
  def _alloc(self, size, options): return memoryview(bytearray(size))
  def _copyin(self, dest, src:memoryview): dest[:] = src
@@ -246,4 +248,4 @@ class PythonAllocator(Allocator['PythonDevice']):

 class PythonDevice(Compiled):
  def __init__(self, device:str):
-    super().__init__(device, PythonAllocator(self), CompilerSet([CompilerPair(PythonRenderer, PythonCompiler)]), PythonProgram)
+    super().__init__(device, PythonAllocator(self), CompilerSet([(PythonRenderer, None)]), PythonProgram)
@@ -2,15 +2,15 @@ from __future__ import annotations
 import os, ctypes, functools, mmap, struct, array, math, sys, weakref, contextlib
 assert sys.platform != 'win32'
 from typing import Any
-from tinygrad.device import BufferSpec, CompilerSet, CompilerPair, Device
+from tinygrad.device import BufferSpec, CompilerSet, Device
 from tinygrad.runtime.support.hcq import HCQBuffer, HWQueue, HCQProgram, HCQCompiled, HCQAllocatorBase, HCQSignal, HCQArgsState, BumpAllocator
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
 from tinygrad.runtime.autogen import kgsl, mesa
-from tinygrad.runtime.ops_cl import CLCompiler, CLDevice
+from tinygrad.runtime.ops_cl import CLDevice
 from tinygrad.renderer.cstyle import QCOMRenderer
 from tinygrad.renderer.nir import IR3Renderer
 from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, ceildiv, prod, fromimport, cpu_profile, lo32, suppress_finalizing
-from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC, PROFILE
+from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC, PROFILE, DEBUG
 from tinygrad.dtype import ImageDType, dtypes
 from tinygrad.runtime.support.system import System
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import
@@ -49,10 +49,6 @@ def pkt7_hdr(opcode: int, cnt: int): return mesa.CP_TYPE7_PKT | cnt & 0x3FFF | p
 def pkt4_hdr(reg: int, cnt: int): return mesa.CP_TYPE4_PKT | cnt & 0x7F | parity(cnt) << 7 | (reg & 0x3FFFF) << 8 | parity(reg) << 27

 def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
-class QCOMCompiler(CLCompiler):
-  def __init__(self, device:str=""): super().__init__(CLDevice(device), 'compile_qcom')
-  def disassemble(self, lib:bytes):
-    fromimport('tinygrad.runtime.support.compiler_mesa', 'disas_adreno')(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)])

 class QCOMSignal(HCQSignal):
  def __init__(self, *args, **kwargs): super().__init__(*args, **{**kwargs, 'timestamp_divider': 19.2})
@@ -231,7 +227,7 @@ class QCOMArgsState(HCQArgsState):
 class QCOMProgram(HCQProgram):
  def __init__(self, dev: QCOMDevice, name: str, lib: bytes, buf_dtypes=[], **kwargs):
    self.dev: QCOMDevice = dev
-    self.buf_dtypes, self.name, self.lib, self.NIR = buf_dtypes, name, lib, isinstance(dev.renderer, IR3Renderer)
+    self.buf_dtypes, self.name, self.NIR = buf_dtypes, name, isinstance(dev.renderer, IR3Renderer)

    if self.NIR:
      from tinygrad.runtime.support.compiler_mesa import IR3Compiler
@@ -252,7 +248,9 @@ class QCOMProgram(HCQProgram):
      self.tex_off, self.ibo_off, self.samp_off = 2048, 2048 + 0x40 * self.tex_cnt, 2048 + 0x40 * (self.tex_cnt + self.ibo_cnt)
      self.fregs, self.hregs = v.info.max_reg + 1, v.info.max_half_reg + 1
      self.consts_info:list[tuple] = []
-    else: self._parse_lib()
+    else:
+      self._parse_lib(lib:=self.dev.cl_dev.cl_compiler.compile_cached(lib.decode()))
+      if DEBUG >= 7: fromimport('tinygrad.runtime.support.compiler_mesa', 'disas_adreno')(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)])

    self.lib_gpu: HCQBuffer = self.dev.allocator.alloc(self.image_size, buf_spec:=BufferSpec(cpu_access=True, nolru=True))
    to_mv(self.lib_gpu.va_addr, self.image_size)[:] = self.image
@@ -274,21 +272,21 @@ class QCOMProgram(HCQProgram):
      raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
    return super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait)

-  def _parse_lib(self):
+  def _parse_lib(self, lib):
    # Extract image binary
-    self.image_size = _read_lib(self.lib, 0x100)
-    self.image = bytearray(self.lib[(image_offset:=_read_lib(self.lib, 0xc0)):image_offset+self.image_size])
+    self.image_size = _read_lib(lib, 0x100)
+    self.image = bytearray(lib[(image_offset:=_read_lib(lib, 0xc0)):image_offset+self.image_size])

    # Parse image descriptors
-    image_desc_off = _read_lib(self.lib, 0x110)
-    self.prg_offset, self.brnchstck = _read_lib(self.lib, image_desc_off+0xc4), _read_lib(self.lib, image_desc_off+0x108) // 2
-    self.pvtmem, self.shmem = _read_lib(self.lib, image_desc_off+0xc8), _read_lib(self.lib, image_desc_off+0xd8)
+    image_desc_off = _read_lib(lib, 0x110)
+    self.prg_offset, self.brnchstck = _read_lib(lib, image_desc_off+0xc4), _read_lib(lib, image_desc_off+0x108) // 2
+    self.pvtmem, self.shmem = _read_lib(lib, image_desc_off+0xc8), _read_lib(lib, image_desc_off+0xd8)

    # Fill up constants and buffers info
    self.consts_info = []

    # Collect sampler info.
-    self.samp_cnt = samp_cnt_in_file = _read_lib(self.lib, image_desc_off + 0xdc)
+    self.samp_cnt = samp_cnt_in_file = _read_lib(lib, image_desc_off + 0xdc)
    assert self.samp_cnt <= 1, "Up to one sampler supported"
    if self.samp_cnt:
      self.samp_cnt += 1
@@ -298,8 +296,8 @@ class QCOMProgram(HCQProgram):

    # Collect kernel arguments (buffers) info.
    bdoff, binfos = round_up(image_desc_off + 0x158 + len(self.name), 4) + 8 * samp_cnt_in_file, []
-    while bdoff + 32 <= len(self.lib):
-      length, _, _, offset_words, _, _, _, typ = struct.unpack("8I", self.lib[bdoff:bdoff+32])
+    while bdoff + 32 <= len(lib):
+      length, _, _, offset_words, _, _, _, typ = struct.unpack("8I", lib[bdoff:bdoff+32])
      if length == 0: break
      binfos.append((offset_words * 4, typ))
      bdoff += length
@@ -309,16 +307,16 @@ class QCOMProgram(HCQProgram):
    self.tex_cnt, self.ibo_cnt = sum(typ is BUFTYPE_TEX for _,typ in binfos), sum(typ is BUFTYPE_IBO for _,typ in binfos)
    self.ibo_off, self.tex_off, self.samp_off = 2048, 2048 + 0x40 * self.ibo_cnt, 2048 + 0x40 * self.tex_cnt + 0x40 * self.ibo_cnt

-    if _read_lib(self.lib, 0xb0) != 0: # check if we have constants.
-      cdoff = _read_lib(self.lib, 0xac)
+    if _read_lib(lib, 0xb0) != 0: # check if we have constants.
+      cdoff = _read_lib(lib, 0xac)
      while cdoff + 40 <= image_offset:
-        cnst, offset_words, _, is32 = struct.unpack("I", self.lib[cdoff:cdoff+4])[0], *struct.unpack("III", self.lib[cdoff+16:cdoff+28])
+        cnst, offset_words, _, is32 = struct.unpack("I", lib[cdoff:cdoff+4])[0], *struct.unpack("III", lib[cdoff+16:cdoff+28])
        self.consts_info.append((cnst, offset_words * (sz_bytes:=(2 << is32)), sz_bytes))
        cdoff += 40

    # Registers info
-    reg_desc_off = _read_lib(self.lib, 0x34)
-    self.fregs, self.hregs = _read_lib(self.lib, reg_desc_off + 0x14), _read_lib(self.lib, reg_desc_off + 0x18)
+    reg_desc_off = _read_lib(lib, 0x34)
+    self.fregs, self.hregs = _read_lib(lib, reg_desc_off + 0x14), _read_lib(lib, reg_desc_off + 0x18)

 class QCOMTextureInfo:
  def __init__(self, pitch:int, real_stride:int, desc:list[int], ibo:list[int]):
@@ -331,7 +329,7 @@ class QCOMAllocator(HCQAllocatorBase):
    return self.dev._gpu_map(opts.external_ptr, size, image=opts.image) if opts.external_ptr else self.dev._gpu_alloc(size, image=opts.image)

  def _do_copy(self, src_addr, dest_addr, src_size, real_size, src_stride, dest_stride, prof_text, dest_off=0, src_off=0):
-    with cpu_profile(prof_text, self.dev.device, is_copy=True):
+    with cpu_profile(prof_text, self.dev.device):
      while src_off < src_size:
        ctypes.memmove(dest_addr+dest_off, src_addr+src_off, real_size)
        src_off, dest_off = src_off+src_stride, dest_off+dest_stride
@@ -385,8 +383,8 @@ class QCOMDevice(HCQCompiled):
    if PROFILE and self.gpu_id[:2] < (7, 3):
      System.write_sysfs("/sys/class/kgsl/kgsl-3d0/idle_timer", value="4000000000", msg="Failed to disable suspend mode", expected="4294967276")

-    compilers = CompilerSet(ctrl_var=QCOM_CC, cset=[CompilerPair(QCOMRenderer, functools.partial(QCOMCompiler, device)),
-                                                    CompilerPair(functools.partial(IR3Renderer, info.chip_id), None, QCOM_IR3)])
+    self.cl_dev = CLDevice(device)
+    compilers = CompilerSet(ctrl_var=QCOM_CC, cset=[(QCOMRenderer, None), (functools.partial(IR3Renderer, info.chip_id), QCOM_IR3)])
    super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal,
                     functools.partial(QCOMComputeQueue, self), None)

@@ -1,5 +1,5 @@
 import functools, struct
-from tinygrad.device import  Compiled, Allocator, Compiler, BufferSpec, CompilerSet, CompilerPair
+from tinygrad.device import Compiled, Allocator, BufferSpec, CompilerSet
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.helpers import round_up, suppress_finalizing
 from tinygrad.runtime.autogen import webgpu
@@ -217,7 +217,7 @@ class WebGpuDevice(Compiled):
    self.device_res = _run(webgpu.wgpuAdapterRequestDeviceF, webgpu.WGPURequestDeviceCallbackInfo, webgpu.WGPURequestDeviceCallback,
      webgpu.WGPURequestDeviceStatus, 1, 2, adapter_res, dev_desc)

-    super().__init__(device, WebGpuAllocator(self), CompilerSet([CompilerPair(WGSLRenderer, Compiler)]),
+    super().__init__(device, WebGpuAllocator(self), CompilerSet([(WGSLRenderer, None)]),
      functools.partial(WebGPUProgram, (self.device_res, webgpu.WGPUFeatureName_TimestampQuery in supported)))

  def synchronize(self):
@@ -264,7 +264,8 @@ class AM_GFX(AM_IP):
      self.adev.regGRBM_CNTL.update(read_timeout=0xff, inst=xcc)
      for i in range(0, 16):
        self._grbm_select(vmid=i, inst=xcc)
-        self.adev.regSH_MEM_CONFIG.write(**({'initial_inst_prefetch':3} if self.adev.ip_ver[am.GC_HWIP][0]>=10 else {'retry_disable':1, 'f8_mode':1}),
+        self.adev.regSH_MEM_CONFIG.write(**({'initial_inst_prefetch':3} if self.adev.ip_ver[am.GC_HWIP][0]>=10 else {'retry_disable':1}),
+          **({'f8_mode':1} if self.adev.ip_ver[am.GC_HWIP][:2]==(9,4) else {}),
          address_mode=self.adev.soc.module.SH_MEM_ADDRESS_MODE_64, alignment_mode=self.adev.soc.module.SH_MEM_ALIGNMENT_MODE_UNALIGNED, inst=xcc)

        # Configure apertures:
@@ -1,5 +1,4 @@
-import subprocess, hashlib, tempfile, ctypes, re, pathlib
-from typing import Callable
+import hashlib, tempfile, ctypes, re, pathlib
 from tinygrad.helpers import to_char_p_p, colored, getenv, system
 from tinygrad.runtime.support.c import init_c_var
 from tinygrad.runtime.autogen import nvrtc, nvjitlink as jitlink
@@ -42,36 +41,32 @@ def cuda_disassemble(lib:bytes, arch:str, ptx=False):
    print(system(f'nvdisasm {fn}'))
  except Exception as e: print("Failed to generate SASS", str(e), "Make sure your PATH contains ptxas/nvdisasm binary of compatible version.")

-class CUDACompiler(Compiler):
-  def __init__(self, arch:str, cache_key:str="cuda"):
-    self.arch, self.compile_options = arch, [f'--gpu-architecture={arch}']
+class NVRTCCompiler(Compiler):
+  def __init__(self, arch:str, ptx=True, cache_key:str="cuda"):
+    self.ptx, self.arch, self.compile_options = ptx, arch, [f'--gpu-architecture={arch}']
    self.compile_options += [f"-I{CUDA_PATH}/include"] if CUDA_PATH else ["-I/usr/local/cuda/include", "-I/usr/include", "-I/opt/cuda/include"]
    nvrtc_check(nvrtc.nvrtcVersion((nvrtcMajor := ctypes.c_int()), (nvrtcMinor := ctypes.c_int())))
    if (nvrtcMajor.value, nvrtcMinor.value) >= (12, 4): self.compile_options.append("--minimal")
    super().__init__(f"compile_{cache_key}_{self.arch}")
-  def _compile_program(self, src:str, nvrtc_get_content:Callable, nvrtc_get_size:Callable) -> bytes:
+  def compile(self, src:str) -> bytes:
    nvrtc_check(nvrtc.nvrtcCreateProgram(ctypes.byref(prog := nvrtc.nvrtcProgram()), src.encode(), "<null>".encode(), 0, None, None))
    nvrtc_check(nvrtc.nvrtcCompileProgram(prog, len(self.compile_options), to_char_p_p([o.encode() for o in self.compile_options])), prog)
-    data = _get_bytes(prog, nvrtc_get_content, nvrtc_get_size, nvrtc_check)
+    data = _get_bytes(prog, nvrtc.nvrtcGetPTX if self.ptx else nvrtc.nvrtcGetCUBIN,
+                      nvrtc.nvrtcGetPTXSize if self.ptx else nvrtc.nvrtcGetCUBINSize, nvrtc_check)
    nvrtc_check(nvrtc.nvrtcDestroyProgram(ctypes.byref(prog)))
    return data
-  def compile(self, src:str) -> bytes: return self._compile_program(src, nvrtc.nvrtcGetPTX, nvrtc.nvrtcGetPTXSize)
-  def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch, ptx=True)
-
-class NVCompiler(CUDACompiler):
-  def __init__(self, arch:str): super().__init__(arch, cache_key="nv")
-  def compile(self, src:str) -> bytes: return self._compile_program(src, nvrtc.nvrtcGetCUBIN, nvrtc.nvrtcGetCUBINSize)
-  def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch)
+  def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch, ptx=self.ptx)

 class NVCCCompiler(Compiler):
-  def __init__(self, arch:str, extra_options:list[str]=[]):
+  def __init__(self, arch:str, ptx:bool=True, cache_key:str="cuda", extra_options:list[str]=[]):
+    assert ptx, "NVCCCompiler cubin support unimplemented"
    self.arch, self.extra_options = arch, extra_options
-    super().__init__(f"compile_nvcc_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}")
+    super().__init__(f"compile_nvcc_{cache_key}_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}")
  def compile(self, src:str) -> bytes:
    with tempfile.NamedTemporaryFile(suffix=".cu") as srcf, tempfile.NamedTemporaryFile(suffix=".ptx") as libf:
      srcf.write(src.encode())
      srcf.flush()
-      subprocess.run(["nvcc", f"-arch={self.arch}", "-ptx", "-o", libf.name, srcf.name] + self.extra_options, check=True)
+      system(f"nvcc -arch={self.arch} -ptx -o {libf.name} {srcf.name}" + ' '.join(self.extra_options))
      return libf.read()
  def disassemble(self, lib:bytes): cuda_disassemble(lib, self.arch, ptx=True)

@@ -257,15 +257,14 @@ class HCQSignal(Generic[HCQDeviceType]):
      value: The value to wait for.
      timeout: Maximum time to wait in milliseconds. Defaults to 30s.
    """
-    start_time = last_sleep_time = int(time.perf_counter() * 1000)
+    start_time = int(time.perf_counter() * 1000)
    while (not_passed:=(prev_value:=self.value) < value) and (cur_time:=int(time.perf_counter() * 1000)) - start_time < timeout:
-      self._sleep(cur_time - last_sleep_time)
-      last_sleep_time = int(time.perf_counter() * 1000)
-      if self.value != prev_value: start_time = last_sleep_time # progress was made, reset timer
+      self._sleep(cur_time - start_time)
+      if self.value != prev_value: start_time = int(time.perf_counter() * 1000) # progress was made, reset timer
    if not_passed and self.value < value: raise RuntimeError(f"Wait timeout: {timeout} ms! (the signal is not set to {value}, but {self.value})")

@contextlib.contextmanager
-def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]|None=None, queue:HWQueue|None=None):
+def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]|None=None, queue:HWQueue|None=None, dev_suff:str|None=None):
  st, en = (dev.new_signal(), dev.new_signal()) if enabled else (None, None)
  assert queue is not None or queue_type is not None, "Either queue or queue_type must be provided"

@@ -279,7 +278,7 @@ def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]
    elif enabled and queue_type is not None:
      queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(en).signal(dev.timeline_signal, dev.next_timeline()).submit(dev)

-    if enabled and PROFILE: dev.sig_prof_records.append((unwrap(st), unwrap(en), desc, (queue_type or type(queue)) is dev.hw_copy_queue_t))
+    if enabled and PROFILE: dev.sig_prof_records.append((unwrap(st), unwrap(en), desc, f"{dev.device}:{dev_suff}" if dev_suff else dev.device))

 class HCQArgsState(Generic[ProgramType]):
  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint|None, ...]=()):
@@ -376,7 +375,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
    self.signal_t, self.hw_compute_queue_t, self.hw_copy_queue_t = signal_t, comp_queue_t, copy_queue_t
    self.timeline_value:int = 1
    self.timeline_signal, self._shadow_timeline_signal = self.new_signal(value=0, is_timeline=True), self.new_signal(value=0, is_timeline=True)
-    self.sig_prof_records:list[tuple[HCQSignal, HCQSignal, str, bool]] = []
+    self.sig_prof_records:list[tuple[HCQSignal, HCQSignal, str, str]] = []
    self.prof_exec_counter:int = 0
    self.prof_prg_counter:int = 0

@@ -402,7 +401,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):

    if self.timeline_value > (1 << 31): self._wrap_timeline_signal()
    if PROFILE:
-      Compiled.profile_events += [ProfileRangeEvent(self.device, name, st.timestamp, en.timestamp, cp) for st,en,name,cp in self.sig_prof_records]
+      Compiled.profile_events += [ProfileRangeEvent(dev, name, st.timestamp, en.timestamp) for st,en,name,dev in self.sig_prof_records]
      self.sig_prof_records = []

  def next_timeline(self):
@@ -418,6 +417,10 @@ class HCQCompiled(Compiled, Generic[SignalType]):

  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.

+  def hw_compute_queues(self) -> list[tuple[str|None, Callable[[], HWQueue]]]: return [(None, self.hw_compute_queue_t)]
+  def hw_copy_queues(self) -> list[tuple[str, Callable[[], HWQueue]]]:
+    return [("SDMA:0", self.hw_copy_queue_t)] if self.hw_copy_queue_t is not None else []
+
  def _at_profile_finalize(self):
    self.synchronize() # Expect device to be synchronizes

@@ -428,10 +431,9 @@ class HCQCompiled(Compiled, Generic[SignalType]):
      et = time.perf_counter_ns()
      return (decimal.Decimal(et+st) / 2000) - d.timeline_signal.timestamp

-    self.gpu2cpu_compute_time_diff = statistics.median([_sync(self, self.hw_compute_queue_t) for _ in range(40)])
-    if self.hw_copy_queue_t is None: gpu2cpu_copy_time_diff = decimal.Decimal(0)
-    else: gpu2cpu_copy_time_diff = statistics.median([_sync(self, self.hw_copy_queue_t) for _ in range(40)])
-    Compiled.profile_events += [ProfileDeviceEvent(self.device, self.gpu2cpu_compute_time_diff, gpu2cpu_copy_time_diff, props=self.device_props())]
+    for prefix, q_t in self.hw_compute_queues() + self.hw_copy_queues():
+      devname = f"{self.device}:{prefix}" if prefix else self.device
+      Compiled.profile_events += [ProfileDeviceEvent(devname, statistics.median([_sync(self, q_t) for _ in range(40)]), props=self.device_props())]

  def _wrap_timeline_signal(self):
    self.timeline_signal, self._shadow_timeline_signal, self.timeline_value = self._shadow_timeline_signal, self.timeline_signal, 1
@@ -514,10 +516,10 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
  def _copyin(self, dest:HCQBuffer, src:memoryview):
    if self.dev.hw_copy_queue_t is None:
      self.dev.synchronize()
-      with cpu_profile(f'TINY -> {self.dev.device}', self.dev.device, is_copy=True): ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
+      with cpu_profile(f'TINY -> {self.dev.device}', self.dev.device): ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
      return

-    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"TINY -> {self.dev.device}", enabled=PROFILE):
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"TINY -> {self.dev.device}", enabled=PROFILE, dev_suff="SDMA:0"):
      for i in range(0, src.nbytes, self.b[0].size):
        self.b_next = (self.b_next + 1) % len(self.b)
        self.dev.timeline_signal.wait(self.b_timeline[self.b_next])
@@ -538,7 +540,7 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
      return None

    assert self.dev.hw_copy_queue_t is not None
-    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"DISK -> {self.dev.device}", enabled=PROFILE):
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"DISK -> {self.dev.device}", enabled=PROFILE, dev_suff="SDMA:0"):
      for (batch_info, dst_off, src_off, copy_size) in src.device.allocator._copyout_sharded(src, size, _get_temp_buf, seg_len=self.b[0].size):
        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
                                  .copy(dest.va_addr + dst_off, batch_info[0] + src_off, copy_size) \
@@ -548,10 +550,10 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
  def _copyout(self, dest:memoryview, src:HCQBuffer):
    self.dev.synchronize()
    if self.dev.hw_copy_queue_t is None:
-      with cpu_profile(f'{self.dev.device} -> TINY', self.dev.device, is_copy=True): ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
+      with cpu_profile(f'{self.dev.device} -> TINY', self.dev.device): ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
      return

-    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"{self.dev.device} -> TINY", enabled=PROFILE):
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"{self.dev.device} -> TINY", enabled=PROFILE, dev_suff="SDMA:0"):
      for i in range(0, dest.nbytes, cp_size:=(self.max_copyout_size or self.b[0].size)):
        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
                                  .copy(self.b[0].va_addr, src.va_addr+i, lsize:=min(cp_size, dest.nbytes-i)) \
@@ -563,7 +565,7 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
    cast(HCQAllocator, src_dev.allocator).map(dest)

    assert src_dev.hw_copy_queue_t is not None
-    with hcq_profile(src_dev, queue_type=src_dev.hw_copy_queue_t, desc=f"{src_dev.device} -> {dest_dev.device}", enabled=PROFILE):
+    with hcq_profile(src_dev, queue_type=src_dev.hw_copy_queue_t, desc=f"{src_dev.device} -> {dest_dev.device}", enabled=PROFILE, dev_suff="SDMA:0"):
      src_dev.hw_copy_queue_t().wait(src_dev.timeline_signal, src_dev.timeline_value - 1) \
                               .wait(dest_dev.timeline_signal, dest_dev.timeline_value - 1) \
                               .copy(dest.va_addr, src.va_addr, sz) \
@@ -3,13 +3,13 @@ import functools, itertools
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches
-from tinygrad.uop.ops import consumer_map_from_toposort
+from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink, pm_gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored

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

 def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None

@@ -17,21 +17,15 @@ 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_assign(ctx:dict[UOp, None], a:UOp) -> None:
-  if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
-  ctx[a] = None
-
-pm_generate_realize_map = PatternMatcher([
+pm_generate_realize_map = pm_gate_kernel_sink+PatternMatcher([
  # always realize SINK src
  (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
-  # always realize COPY/BUFFER_VIEW/CONTIGUOUS/STORE/ENCDEC
-  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE, Ops.ENCDEC}, name="tr"), realize),
+  # always realize
+  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE, Ops.ASSIGN, Ops.ENCDEC}, name="tr"), realize),
  # always realize REDUCE on outer ranges
  (UPat(Ops.REDUCE, name="r"), lambda ctx,r: realize(ctx, r) if any(tr.arg[-1] == AxisType.OUTER for tr in r.src[1:]) else None),
-  # realize srcs of COPY, MSELECT, MSTACK, ENCDEC
-  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK, Ops.ENCDEC), name="rb"), realize_srcs),
-  # realize ASSIGN and input to assign (might be optimized out)
-  (UPat(Ops.ASSIGN, name="a"), realize_assign),
+  # realize srcs of these
+  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK, Ops.ASSIGN, Ops.ENCDEC), name="rb"), realize_srcs),
 ])

@dataclass(frozen=True)
@@ -97,20 +91,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 add_third_op_to_assign_to_track_shape(ctx:IndexingContext, assign:UOp):
-  if assign.src[1].op is Ops.KERNEL: return None
-  to_mop = graph_rewrite(assign.src[0], PatternMatcher([(UPat(GroupOp.Movement, name="x"), lambda x: x.replace(tag=()))]))
-  ret = assign.replace(src=assign.src+(to_mop,))
-  ctx.range_map[ret] = ctx.range_map[assign]
-  return ret
+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),
-  # add third op to assign
-  (UPat(Ops.ASSIGN, src=(UPat(), UPat()), name="assign"), add_third_op_to_assign_to_track_shape),
+  # 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
@@ -161,11 +160,11 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
  rctx = IndexingContext()

  # get ops to realize
-  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="get realize")
+  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, bottom_up=True, name="get realize")

  # get the consumer map
  with cpu_profile("consumer map in rangeify", "TINY"):
-    consumer_map = consumer_map_from_toposort(tsink_toposort:=tsink.toposort())
+    consumer_map = consumer_map_from_toposort(tsink_toposort:=tsink.toposort(gate_kernel_sink))

  # explicit rangeify
  ending_ranges: dict[UOp, list[UOp]] = {}
@@ -173,7 +172,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
    if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue

    # no ranges on kernels, they are internal
-    if x.op is Ops.KERNEL: continue
+    if x.op is Ops.CALL: continue

    if x.dtype.scalar() == dtypes.index: continue  # TODO: why do I need this?
    ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
@@ -1,7 +1,8 @@
 from typing import cast
 import functools, itertools
-from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, ALL2ALL, getenv
+from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, ALL2ALL, VIZ, getenv
 from tinygrad.uop.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, graph_rewrite_map, graph_rewrite
+from tinygrad.dtype import dtypes
 from tinygrad.device import Device

 # *** allreduce implementation ***
@@ -214,17 +215,18 @@ multi_pm = PatternMatcher([
  (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)),
  (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)),
  (UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD),
        src=(UPat(Ops.MULTI, name="multi"), ), name="root"), passthrough_multi),
-  # multi supports custom kernels with CUSTOM_KERNEL + AFTER
-  (UPat(Ops.CUSTOM_KERNEL, src=UPat((Ops.MULTI, Ops.CONTIGUOUS)), name="ck"),
-    lambda ck: ck.replace(src=tuple(m.src[0] if m.op is Ops.MULTI else m for m in ck.src))),
-  (UPat(Ops.AFTER, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.CUSTOM_KERNEL)), name="a"),
-    lambda multi,a: a.replace(src=(multi.src[0],)+a.src[1:]).multi(multi.axis))
+  # 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

 def get_multi_map(big_sink:UOp) -> dict[UOp, UOp]:
-  if getenv("VIZ"): graph_rewrite(big_sink, PatternMatcher([]), name="View Multi AST")
+  if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Multi AST")
  ret = graph_rewrite_map(big_sink, multi_pm, name="multi_pm")
-  if getenv("VIZ"): graph_rewrite(ret[big_sink], PatternMatcher([]), name="View Post Multi AST")
+  if VIZ: graph_rewrite(ret[big_sink], PatternMatcher([]), name="View Post Multi AST")
  return ret
@@ -1,10 +1,10 @@
 from dataclasses import dataclass, field, replace
 import itertools
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo
-from tinygrad.uop.ops import graph_rewrite, identity_element, sint, AxisType, BottomUpGate, Kernel, _remove_all_tags, range_str
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo, pm_gate_kernel_sink
+from tinygrad.uop.ops import graph_rewrite, identity_element, sint, AxisType, BottomUpGate, _remove_all_tags, range_str
 from tinygrad.uop.symbolic import symbolic
-from tinygrad.helpers import argsort, prod, all_same, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY
+from tinygrad.helpers import argsort, prod, all_same, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ
 from tinygrad.helpers import PCONTIG, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
@@ -14,6 +14,12 @@ from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTI
 import sys
 sys.setrecursionlimit(10000)

+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),
+])
+
 # movement op on INDEX as a PatternMatcher
 pm_mops = PatternMatcher([
  (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
@@ -27,13 +33,21 @@ pm_mops = PatternMatcher([
 # *****************
 # 0. do some cleanup rewrites, mostly copied from the old stuff

-def fix_assign_hazard(dest:UOp, src:UOp, assign:UOp):
+def collapse_nested_assign(assign:UOp, target:UOp, src:UOp):
+  """nested ASSIGN to the same buffer (e.g. __iadd__ in __setitem__): collapse the redundant outer ASSIGN"""
+  if src.src[0].base is target.base: return src if src.src[0] is target else assign.replace(src=(target, src.src[1]))
+
+def assign_to_contiguous(assign:UOp, target:UOp, src:UOp):
+  if (t := target.base).op is Ops.BUFFER or (t.op is Ops.MSTACK and all(s.op is Ops.BUFFER for s in t.src)): return None
+  return src.f(Ops.CONTIGUOUS, tag=assign.tag)
+
+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 dest.op_in_backward_slice_with_self(Ops.SHRINK) else set())
+  unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set())
  if not (hazards:=[s for s in src.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS) if s.op in unsafe]): return
  for h in hazards:
-    if any(s is dest.base for s in h.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS-{Ops.BUFFER})):
-      return assign.replace(src=(dest, src.contiguous()))
+    if any(s is target.base for s in h.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS-{Ops.BUFFER})):
+      return assign.replace(src=(target, src.contiguous()))

 def split_reduceop(reduce:UOp, x:UOp):
  if prod(reduce.shape) == 0: return None
@@ -64,11 +78,9 @@ mop_cleanup = PatternMatcher([
   lambda x,x2: x.replace(src=(x2.src[0], x.src[1])) if x.tag is None and x2.tag is None else None),
 ])

-def resolve_custom_kernel(ck:UOp) -> UOp:
-  placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(ck.src)]
-  return UOp(Ops.KERNEL, src=ck.src, arg=Kernel(ck.arg.fxn(*placeholders)))
-
-def resolve_call(c:UOp) -> UOp:
+def resolve_call(c:UOp) -> UOp|None:
+  # we only resolve here if the call is inlined
+  if not c.arg.inline: return None
  params = sorted([x for x in c.src[0].toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
  args = c.src[1:]
  # TODO: this check belongs in spec, not here
@@ -86,11 +98,8 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
  # resolve calls
  (UPat(Ops.CALL, name="c"), resolve_call),

-  # resolve custom kernels
-  (UPat(Ops.CUSTOM_KERNEL, name="ck"), resolve_custom_kernel),
-
-  # remove CONTIGUOUS if the BUFFER is already contiguous
-  (UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER), UPat()), name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
+  # remove CONTIGUOUS if the source is already contiguous
+  (UPat(Ops.RESHAPE, src=(UPat((Ops.BUFFER, Ops.CONTIGUOUS)), UPat()), name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),

  # split_reduceop
  (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop),
@@ -126,17 +135,18 @@ earliest_rewrites = mop_cleanup+PatternMatcher([

  # ** assign rules **

+  # collapse nested ASSIGN to the same buffer (e.g. __iadd__ in __setitem__)
+  (UPat(Ops.ASSIGN, src=(UPat(name="target"), UPat(Ops.ASSIGN, name="src")), name="assign"), collapse_nested_assign),
+
  # 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")), name="assign"),
-   lambda target, src, assign: target.assign(src.bitcast(target.dtype)).replace(tag=assign.tag)),
+   lambda assign, target, src: target.assign(src.bitcast(target.dtype)).replace(tag=assign.tag)),

  # assign only to buffer, otherwise make it a CONTIGUOUS
-  (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.BUFFER}, name="target"), UPat(name="x")), name="assign"),
-   lambda x,target,assign: x.f(Ops.CONTIGUOUS, tag=assign.tag) if ((t:=target.base).op is not Ops.BUFFER and \
-       not (t.op is Ops.MSTACK and all(s.op is Ops.BUFFER for s in t.src))) else None),
+  (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.BUFFER}, name="target"), UPat(name="src")), name="assign"), assign_to_contiguous),

   # make source contiguous if it has hazardous movement ops on the dest buffer
-   (UPat(Ops.ASSIGN, src=(UPat.var("dest"), UPat.var("src")), name="assign"), fix_assign_hazard),
+   (UPat(Ops.ASSIGN, src=(UPat.var("target"), UPat.var("src")), name="assign"), fix_assign_hazard),
 ])

 # *****************
@@ -325,19 +335,13 @@ 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)
-  if x.src[0].op is Ops.ASSIGN:
-    assign_target, assign_src, assign_mops = x.src[0].src
+  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"
    # in assign, this is the buffer size, not the bufferize size
-    # TODO: assign_mops here
    do_store = assign_target.replace(dtype=sdtype).store(assign_src, tag=x.tag).end(*rngs)
    ret = assign_target.src[0].after(do_store)
-    mops = []
-    walk = assign_mops
-    while walk is not assign_mops.base:
-      mops.append((walk.op, walk.marg))
-      walk = walk.src[0]
-    for m in mops[::-1]: ret = ret._mop(*m)
+    for op, marg in reversed(assign.arg or ()): ret = ret._mop(op, marg)
    return ret

  # lower outerworld reduce here
@@ -384,7 +388,7 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
   lambda m: m.replace(src=tuple([x.src[0].base for x in m.src]), tag=None).reshape(m.shape).rtag(m.tag)),

  # remove any RESHAPEs on KERNEL
-  (UPat(Ops.KERNEL, name="k"), lambda k: k.replace(src=tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src))),
+  (UPat(Ops.CALL, name="k"), lambda k: k.replace(src=tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src))),
 ])

 pm_add_buffers_local = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
@@ -515,10 +519,10 @@ def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
  if stored.op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENCDEC}: ret = stored
  else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))

-  kernel_arg = Kernel(ret,tuple(dedup(flatten([x for x in metadatas if x is not None])))[::-1])
-  kernel = UOp(Ops.KERNEL, src=tuple(lctx.map.values())+tuple(lctx.vars.keys()), arg=kernel_arg)
-  if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src if x.op is not Ops.BIND]):
-    raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop for b in kernel.src)}")
+  metadata = tuple(dedup(flatten([x for x in metadatas if x is not None])))[::-1]
+  kernel = ret.call(*lctx.map.values(), *lctx.vars.keys(), metadata=metadata)
+  if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src[1:] if x.op is not Ops.BIND]):
+    raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop for b in kernel.src[1:])}")
  return kernel

 split_kernels = PatternMatcher([
@@ -533,9 +537,9 @@ def tag_uop(ctx:tuple[list[UOp], set[UOp]], x:UOp):
  if x.dtype.scalar() == dtypes.index: return None
  ctx[0].append(x)
  return x.replace(tag=(len(ctx[0])-1,))
-add_tags = PatternMatcher([
+add_tags = pm_gate_kernel_sink+PatternMatcher([
  # don't tag BUFFERs, they are global
-  (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.LUNIQUE, Ops.DEFINE_VAR, Ops.BIND, Ops.KERNEL, Ops.END,
+  (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.LUNIQUE, Ops.DEFINE_VAR, Ops.BIND, Ops.CALL, Ops.END,
                     Ops.MSTACK, Ops.MSELECT, Ops.RANGE}.union(GroupOp.Movement), name="x"), tag_uop),
  (UPat({Ops.MSTACK, Ops.MSELECT}, name="x"), lambda ctx,x: None if all(s.op is Ops.BUFFER for s in x.src) else tag_uop(ctx, x)),
 ])
@@ -556,11 +560,11 @@ replace_contiguous = PatternMatcher([
 ])

 def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
-  if getenv("VIZ"): graph_rewrite(sink, PatternMatcher([]), name="View Input Graph")
+  if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Input Graph")
  uop_list: list[UOp] = []
  tsink = graph_rewrite(sink, add_tags, ctx=(uop_list, set()), bottom_up=True, name="number the uops")

-  tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")
+  tsink = graph_rewrite(tsink, pm_syntactic_sugar+pm_mops+earliest_rewrites+replace_contiguous, ctx={}, bottom_up=True, name="earliest rewrites")

  # convert movement ops to ranges
  tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
@@ -574,12 +578,13 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
  tsink = UOp.sink(*[x for x in tsink.backward_slice if x.base.op in {Ops.BUFFERIZE, Ops.MSTACK, Ops.CONST, Ops.BUFFER, Ops.AFTER} and \
                     x.tag is not None and len(x.tag)])

-  if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")
+  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")

  # bufferize -> store
  lunique_start: int = max([-1]+[x.arg for x in tsink.toposort() if x.op is Ops.LUNIQUE]) + 1
-  tsink = graph_rewrite(tsink, pm_add_buffers+pm_add_range_tags, ctx=itertools.count(lunique_start), bottom_up=True, name="bufferize to store")
-  tsink = graph_rewrite(tsink, split_kernels, ctx=uop_list, name="split kernels")
+  tsink = graph_rewrite(tsink, pm_gate_kernel_sink+pm_add_buffers+pm_add_range_tags, ctx=itertools.count(lunique_start), bottom_up=True,
+                        name="bufferize to store")
+  tsink = graph_rewrite(tsink, pm_gate_kernel_sink+split_kernels, ctx=uop_list, bottom_up=True, name="split kernels")

  # if a kernel depends on a buffer, and that buffer is later assigned to, make the assign depend on the kernel's assign
  kernel_assign: dict[UOp, UOp] = {}
@@ -599,7 +604,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
  sink_tags = [s.tag for s in tsink.src]
  tsink = graph_rewrite(tsink, _remove_all_tags, name="remove all tags")

-  if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
+  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")

  becomes_map: dict[UOp, UOp] = {}
  for tag, s in zip(sink_tags, tsink.src):
@@ -24,6 +24,7 @@ def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
 # *** all in scope Tensors are here. this gets relevant UOps ***

 all_tensors: dict[weakref.ref[Tensor], None] = {}
+_pending_assigns: dict[UOp, list[UOp]] = {}  # buffer_uop -> [assign_uops in insertion order]
 def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
  with cpu_profile(TracingKey(name), "TINY"):
    # get tensors in scope
@@ -234,13 +235,12 @@ class Tensor(OpMixin):

  def as_param(self, slot:int):
    if self.uop.axis is not None:
-      multi_shape = tuple([s//len(self.device) if i==self.uop.axis else s for i,s in enumerate(self.shape)])
-      param = UOp.param(slot, self.dtype, multi_shape, self.device).multi(self.uop.axis)
+      param = UOp.param(slot, self.dtype, self.uop.shard_shape, self.device).multi(self.uop.axis)
    else:
      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:
-    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 call(self, *lst:Tensor, fxn:Tensor|UOp, **kwargs) -> Tensor:
+    return Tensor((fxn.uop if isinstance(fxn, Tensor) else fxn).call(*[t.uop for t in (self,)+lst], **kwargs), device=self.device)

  def custom_kernel(self, *lst:Tensor, fxn:Callable, grad_fxn:Callable|None=None) -> list[Tensor]:
    """
@@ -272,6 +272,13 @@ class Tensor(OpMixin):
  @disable_gc()
  def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
    """Triggers the computation needed to create these Tensor(s)."""
+    # side-realize pending assigns for buffers referenced by these tensors
+    if _pending_assigns:
+      for buf in {u for t in (self,)+lst for u in t.uop.toposort() if u.op is Ops.BUFFER}:
+        for assign_uop in _pending_assigns.pop(buf, []):
+          becomes_map, schedule, var_vals = complete_create_schedule_with_vars(UOp.sink(assign_uop))
+          _apply_map_to_tensors(becomes_map, name="Apply Pending Assign")
+          run_schedule(schedule, var_vals, do_update_stats=do_update_stats)
    if len(to_realize:=[x for x in (self,)+lst if not x.uop.has_buffer_identity()]):
      run_schedule(*Tensor.schedule_with_vars(*to_realize), do_update_stats=do_update_stats)
    return self
@@ -300,7 +307,11 @@ class Tensor(OpMixin):
    if is_disk:
      self._buffer().copyin(x._data())
      return self
-    return self.replace(self._apply_uop(UOp.assign, x))
+    result = self._apply_uop(UOp.assign, x)
+    # track view assigns (not full-buffer or assign-chain) so they can be side-realized when the buffer is read
+    if (buf_uop:=self.uop.base).op is Ops.BUFFER and self.uop.op is not Ops.ASSIGN and not self.uop.has_buffer_identity():
+      _pending_assigns.setdefault(buf_uop, []).append(result.uop)
+    return self.replace(result)

  def detach(self) -> Tensor:
    """
@@ -756,8 +767,7 @@ class Tensor(OpMixin):
    dtype = kwargs.pop("dtype", self.dtype)
    if kwargs.get("device") is not None: raise RuntimeError("cannot specify `device` on `*_like` of a multi device tensor")
    if self.uop.axis is None: return fxn(self.shape, *args, dtype=dtype, **kwargs).shard(self.device)
-    sharded_shape = tuple(s//len(self.device) if a==self.uop.axis else s for a,s in enumerate(self.shape))
-    stacked = UOp(Ops.MSTACK, dtype=dtype, src=tuple([fxn(sharded_shape, *args, device=d, dtype=dtype, **kwargs).uop for d in self.device]))
+    stacked = UOp(Ops.MSTACK, dtype=dtype, src=tuple([fxn(self.uop.shard_shape, *args, device=d, dtype=dtype, **kwargs).uop for d in self.device]))
    return Tensor(UOp.multi(stacked, axis=self.uop.axis), device=self.device, dtype=dtype)

  def full_like(self, fill_value:PyConst, **kwargs) -> Tensor:
@@ -1281,21 +1291,20 @@ class Tensor(OpMixin):
    return self._getitem(indices)

  def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None:
-    if isinstance(self.device, str) and self.device.startswith("DISK"):
-      self.realize()._getitem(indices).assign(v)
-      return
-    # NOTE: check that setitem target is valid first
-    if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
-    if self.requires_grad or v.requires_grad: raise NotImplementedError("setitem with requires_grad is not supported")
-    self.realize()
-    if not self.uop.is_writable_view(): raise RuntimeError("setitem target must be a writable view backed by a buffer")
-    res = self._getitem(indices, v)
-    # if shapes match and data is not shared it's a copy and we assign to self
-    if res.shape == self.shape and res.uop is not self.uop:
-      self.assign(res).realize()
-    else: # no copy, basic setitem
-      v = v.cast(res.dtype)._broadcast_to(_broadcast_shape(res.shape, v.shape)).contiguous()
-      res.assign(v).realize()
+    if isinstance(v, Tensor) and v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
+    if self.requires_grad or (isinstance(v, Tensor) and v.requires_grad): raise NotImplementedError("setitem with requires_grad is not supported")
+    idx = [indices] if (isinstance(indices, list) and all_int(indices)) or not isinstance(indices, (tuple, list)) else list(indices)
+    is_disk = isinstance(self.device, str) and self.device.startswith("DISK")
+    if any(isinstance(i, (Tensor, list, tuple)) for i in idx): # advanced setitem
+      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))
+    else: # basic setitem
+      if is_disk: self[indices].assign(v)
+      else:
+        self.realize()
+        if not self.uop.is_writable_view(): raise RuntimeError("setitem target must be a writable view backed by a buffer")
+        self[indices].assign(v).realize()

  def __delitem__(self, indices) -> None:
    raise TypeError("Tensor does not support deleting items")
@@ -3786,7 +3795,7 @@ class Tensor(OpMixin):
    S, indices = U.square().sum(-2).sqrt().sort(dim = -1, descending=True)
    new_indices = indices.reshape(b_shape + (1, num)).expand(b_shape + (num, num))
    U = U.gather(-1, new_indices) / (S != 0).where(S, 1).unsqueeze(-2)
-    V = V.gather(-1, new_indices).realize()
+    V = V.gather(-1, new_indices)

    padded_u = Tensor.eye(q_num, dtype=U.dtype).reshape((1,) * len(b_shape) + (q_num, q_num)).expand(b_shape + (q_num, q_num)).contiguous()
    padded_u[..., 0:num, 0:num] = U
@@ -79,8 +79,7 @@ class Ops(FastEnum):
  # ** 6 -- ops that don't exist in programs **

  # tensor graph ops
-  UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); ASSIGN = auto()
-  CUSTOM_KERNEL = auto()
+  UNIQUE = auto(); DEVICE = auto(); ASSIGN = auto()

  # local unique
  LUNIQUE = auto()
@@ -14,8 +14,8 @@ def _lazy_map_numbers(x:UOp, inf:UOp, _inf:UOp, nan:UOp, ratio:UOp):

 # *** helper functions for bit manipulation ***
 def mantissa_bits(d:DType) -> int: return dtypes.finfo(d.scalar())[1]
-def exponent_bias(d:DType) -> int: return {dtypes.float64: 1023, dtypes.float32: 127, dtypes.float16: 15}[d.scalar()]
-def exponent_mask(d:DType) -> int: return {dtypes.float64: 2047, dtypes.float32: 255, dtypes.float16: 31}[d.scalar()]
+def exponent_bias(d:DType) -> int: return (1 << (dtypes.finfo(d.scalar())[0] - 1)) - 1
+def exponent_mask(d:DType) -> int: return (1 << dtypes.finfo(d.scalar())[0]) - 1

 # **** utils ****
 def shr(x:UOp|int, y:UOp|int) -> UOp: return x // (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y)
@@ -378,33 +378,46 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
    case _: raise NotImplementedError(f"long decomposition of {op} unsupported")

 # ***** floats *****
-f2f_dt = { dtypes.half: dtypes.ushort, dtypes.float: dtypes.uint }
+f2f_dt = { f:getattr(dtypes, f"uint{f.bitsize}") for f in dtypes.floats }

 def rne(v: UOp, s) -> UOp: return shr(v, s) + ((shr(v, s - 1) & 1) & ((v & ((1 << (s - 1)) - 1)).ne(0).cast(v.dtype) | (shr(v, s) & 1)))

 def f2f(v, fr:DType, to:DType):
  fs, fb, (fe, fm), ts, tb, (te, tm) = fr.bitsize, exponent_bias(fr), dtypes.finfo(fr), to.bitsize, exponent_bias(to), dtypes.finfo(to)
  # NB: denormals are zero!
-  if fe < te and fm < tm:
+  if fe <= te and fm < tm:
    sign, nosign = shl((v & shl(1, fs-1)).cast(f2f_dt[to]), ts - fs), (v & (shl(1, fs-1) - 1)).cast(f2f_dt[to])
    exp, norm = shr(nosign, fm), shl(nosign, tm - fm) + shl(tb - fb, tm)
-    inf_or_nan = shl(nosign, tm - fm) | shl((shl(1, te) - 1), tm)
-    return (sign | exp.eq(0).where(0, exp.eq(shl(1, fe) - 1).where(inf_or_nan, norm))).bitcast(to)
-  elif fe > te and fm > tm:
-    sign, nosign, exp = shr(v, fs - ts) & shl(1, ts - 1), v & (shl(1, fs - 1) - 1), shr(v, fm) & (shl(1, fe) - 1)
+    nan = shl(nosign, tm - fm) | shl((shl(1, te) - 1), tm)
+    # fp8e4m3 has only one nan
+    is_nan = (nosign.eq(shl(1, fm + fe) - 1) if fr == dtypes.fp8e4m3 else exp.eq(shl(1, fe) - 1))
+    return (sign | exp.eq(0).where(0, is_nan.where(nan, norm))).bitcast(to)
+  elif fe >= te and fm > tm:
+    v = f2f_clamp(v.bitcast(fr), to).bitcast(f2f_dt[fr])
+    sign, nosign = shr(v, fs - ts) & shl(1, ts - 1), v & (shl(1, fs - 1) - 1)
    norm = (rne(nosign, fm - tm) - shl(fb - tb, tm)).cast(f2f_dt[to])
-    infnan = (sign | (shr(nosign, fm - tm) & (shl(1, tm) - 1)) | shl(shl(1, te) - 1, tm)).cast(f2f_dt[to])
-    underflow, overflow = exp < (1 + fb - tb), exp > (shl(1, te) - 2 + (fb - tb))
-    return exp.eq(shl(1, fe) - 1).where(infnan, sign.cast(f2f_dt[to]) | underflow.where(0, overflow.where(shl(shl(1, te) - 1, tm), norm)))
+    underflow = (shr(v, fm) & (shl(1, fe) - 1)) < (1 + fb - tb)
+    nan_mantissa = (shl(1, tm) - 1) if to == dtypes.fp8e4m3 else (shr(nosign, fm - tm) & (shl(1, tm) - 1))
+    nan = (sign | nan_mantissa | shl(shl(1, te) - 1, tm)).cast(f2f_dt[to])
+    is_nan = (shr(v, fm) & (shl(1, fe) - 1)).eq(shl(1, fe) - 1)
+    return is_nan.where(nan, sign.cast(f2f_dt[to]) | underflow.where(0, norm))
  else: raise NotImplementedError(f"unsupported decomp {fr} -> {to}")

-def f2f_load(x: UOp) -> UOp:
-  if (n:=x.dtype.count) == 1: return f2f(x.replace(dtype=dtypes.ushort), dtypes.half, dtypes.float)
-  return UOp.vectorize(*(f2f(x.replace(dtype=dtypes.ushort, src=(reindex(x.src[0].src[0], i, 1),)), dtypes.half, dtypes.float) for i in range(n)))
+def f2f_clamp(val:UOp, dt:DType) -> UOp:
+  e, m = dtypes.finfo(dt)
+  max_exp, max_man = ((1 << e) - 1, (1 << m) - 2) if dt == dtypes.fp8e4m3 else ((1 << e) - 2, (1 << m) - 1)
+  mx = val.const_like(2.0**(max_exp - exponent_bias(dt)) * (1.0 + max_man / (1 << m)))
+  sat = mx if dt in dtypes.fp8s else val.const_like(float('inf'))
+  # FIXME: CMPLT of nan is undefined
+  return val.ne(val).where(val, (val < -mx).where(-sat, (mx < val).where(sat, val)))

-def f2f_store(st, idx, val):
-  if (n:=val.dtype.count) == 1: return st.replace(src=(idx, f2f(val.bitcast(dtypes.uint), dtypes.float, dtypes.half)))
-  return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.gep(i).bitcast(dtypes.uint), dtypes.float, dtypes.half))) for i in range(n)))
+def f2f_load(x: UOp, fr:DType, to:DType) -> UOp:
+  if (n:=x.dtype.count) == 1: return f2f(x.replace(dtype=f2f_dt[fr]), fr, to)
+  return UOp.vectorize(*(f2f(x.replace(dtype=f2f_dt[fr], src=(reindex(x.src[0].src[0], i, 1),)), fr, to) for i in range(n)))
+
+def f2f_store(st, idx, val, fr:DType, to:DType):
+  if (n:=val.dtype.count) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr)))
+  return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.gep(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))

 # ***** decomposition patterns *****

@@ -463,40 +476,44 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], device:str, disable_fast_idiv
    pat += [(UPat.var("a", dtypes.floats) * UPat.const(dtypes.floats, 1).alu(Ops.FDIV, UPat.var("b")), lambda a,b: a.alu(Ops.FDIV, b))]
  return PatternMatcher(pat)

-@functools.cache
-def get_unsupported_dtypes_patterns(device:str, emulated_dtypes:tuple[DType, ...]) -> PatternMatcher:
-  pat: list[tuple[UPat, Callable]] = []
-  if not is_dtype_supported(dtypes.long, device) or dtypes.long in emulated_dtypes:
-    pat += [(UPat((*GroupOp.Defines, Ops.INDEX), name="x"), lambda x:
-             x.replace(dtype=l2i_dt[x.dtype.base].ptr(x.dtype.size * 2)) if hasattr(x.dtype, 'size') and x.dtype.base in l2i_dt else None)]
-    pat += [(UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x: reindex(x, x.tag).replace(dtype=l2i_dt[x.dtype]))]
-    pat += [(UPat(Ops.STORE, src=(UPat.var('idx'), UPat.var('val', tuple(l2i_dt.keys()))), name='st'), lambda st,idx,val:
-             st.replace(src=(reindex(idx, 0), val.rtag(0))).group(st.replace(src=(reindex(idx, 1), val.rtag(1)))) if val.tag is None else None)]
-    pat += [(UPat(GroupOp.Comparison, src=(UPat.var('a', tuple(l2i_dt.keys())), UPat.var('b', tuple(l2i_dt.keys()))), name="x"), lambda a,b,x:
-             l2i(x.op, dt:=l2i_dt[a.dtype], a.rtag(0).cast(dt), a.rtag(1).cast(dt), b.rtag(0).cast(dt), b.rtag(1).cast(dt)))]
-    pat += [(UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda a,x:
-             l2i(x.op, x.dtype, a)[x.tag] if x.tag is not None and a.dtype not in l2i_dt else None)]
-    pat += [(UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
-             (a.rtag(0).cast(dt:=l2i_dt[a.dtype]).bitcast(xdt:=l2i_dt[x.dtype]), a.rtag(1).cast(dt).bitcast(xdt))[x.tag])]
-    pat += [(UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
-             l2i(x.op, x.dtype, a.rtag(0).cast(dt:=l2i_dt[a.dtype]), a.rtag(1).cast(dt)) if x.dtype not in l2i_dt and a.tag is None else None)]
-    pat += [(UPat((*(GroupOp.ALU - GroupOp.Comparison), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda x:
-             l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag(0).cast(dt:=l2i_dt[x.src[-1].dtype]), a.rtag(1).cast(dt))
-                                                                            if a.dtype in l2i_dt else (a,) for a in x.src))[x.tag])]
-    pat += [(UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx:
-             x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag),)))]
-    pat += [(UPat(Ops.CONST, tuple(l2i_dt.keys()), name='x'), lambda x:
-             UOp.const(dt:=l2i_dt[x.dtype], truncate[dt]((x.arg >> 32) if x.tag == 1 else (x.arg & 0xFFFFFFFF))))]
-  if dtypes.half in emulated_dtypes:
-    pat += [(UPat((*GroupOp.Defines, Ops.INDEX), name="x"), lambda x:
-             x.replace(dtype=dtypes.uint16.ptr(x.dtype.size), tag=dtypes.half) if x.dtype.base == dtypes.half else None)]
-    pat += [(UPat(Ops.LOAD, dtypes.half, name="x"), f2f_load)]
-    pat += [(UPat(Ops.BITCAST, src=(UPat(Ops.LOAD, dtypes.half, name="ld"),), name="bc"), lambda bc,ld:
-             ld.replace(dtype=dtypes.ushort).bitcast(bc.dtype))]
-    pat += [(UPat(Ops.BITCAST, (dtypes.ushort, dtypes.short, dtypes.bfloat16), src=(UPat.var("x", dtypes.float),), name="bc"), lambda bc,x:
-             bc.replace(src=(f2f(x.bitcast(dtypes.uint), dtypes.float, dtypes.half),)))]
-    pat += [(UPat(GroupOp.All, dtypes.half, name="x"), lambda x:
-             x.replace(dtype=dtypes.float.vec(x.dtype.count), src=tuple(s.cast(dtypes.float) if s.dtype == dtypes.half else s for s in x.src)))]
-    pat += [(UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val", dtypes.float)), name='st'), lambda st,idx,val:
-             f2f_store(st, idx, val) if (idx:=idx.src[0] if idx.op == Ops.CAST else idx).tag == dtypes.half else None)]
-  return PatternMatcher(pat)
+pm_long_decomp = PatternMatcher([
+  (UPat((*GroupOp.Defines, Ops.INDEX), name="x"), lambda x:
+   x.replace(dtype=l2i_dt[x.dtype.base].ptr(x.dtype.size * 2)) if hasattr(x.dtype, 'size') and x.dtype.base in l2i_dt else None),
+  (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x: reindex(x, x.tag).replace(dtype=l2i_dt[x.dtype])),
+  (UPat(Ops.STORE, src=(UPat.var('idx'), UPat.var('val', tuple(l2i_dt.keys()))), name='st'), lambda st,idx,val:
+   st.replace(src=(reindex(idx, 0), val.rtag(0))).group(st.replace(src=(reindex(idx, 1), val.rtag(1)))) if val.tag is None else None),
+  (UPat(GroupOp.Comparison, src=(UPat.var('a', tuple(l2i_dt.keys())), UPat.var('b', tuple(l2i_dt.keys()))), name="x"), lambda a,b,x:
+   l2i(x.op, dt:=l2i_dt[a.dtype], a.rtag(0).cast(dt), a.rtag(1).cast(dt), b.rtag(0).cast(dt), b.rtag(1).cast(dt))),
+  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda a,x:
+   l2i(x.op, x.dtype, a)[x.tag] if x.tag is not None and a.dtype not in l2i_dt else None),
+  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
+   (a.rtag(0).cast(dt:=l2i_dt[a.dtype]).bitcast(xdt:=l2i_dt[x.dtype]), a.rtag(1).cast(dt).bitcast(xdt))[x.tag]),
+  (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
+   l2i(x.op, x.dtype, a.rtag(0).cast(dt:=l2i_dt[a.dtype]), a.rtag(1).cast(dt)) if x.dtype not in l2i_dt and a.tag is None else None),
+  (UPat((*(GroupOp.ALU - GroupOp.Comparison), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda x:
+   l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag(0).cast(dt:=l2i_dt[x.src[-1].dtype]), a.rtag(1).cast(dt))
+                                       if a.dtype in l2i_dt else (a,) for a in x.src))[x.tag]),
+  (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx: x.replace(dtype=l2i_dt[x.dtype],src=(reindex(idx, x.tag),))),
+  (UPat(Ops.CONST, tuple(l2i_dt.keys()), name='x'), lambda x:
+   UOp.const(dt:=l2i_dt[x.dtype], truncate[dt]((x.arg >> 32) if x.tag == 1 else (x.arg & 0xFFFFFFFF))))
+])
+
+# float decomposition patterns - ctx is (fr, to) tuple
+pm_float_decomp = PatternMatcher([
+  (UPat((*GroupOp.Defines, Ops.INDEX), name="x"), lambda ctx,x:
+   x.replace(dtype=f2f_dt[ctx[0]].ptr(x.dtype.size), tag=ctx[0]) if x.dtype.base == ctx[0] else None),
+  (UPat(Ops.LOAD, dtypes.floats, name="x"), lambda ctx,x: f2f_load(x, *ctx) if x.dtype.scalar() == ctx[0] else None),
+  (UPat(Ops.BITCAST, src=(UPat(Ops.LOAD, name="ld"),), name="bc"), lambda ctx,bc,ld:
+   ld.replace(dtype=f2f_dt[ctx[0]]).bitcast(bc.dtype) if ld.dtype.bitsize == ctx[0].bitsize else None),
+  (UPat(Ops.BITCAST, src=(UPat.var("x", dtypes.floats),), name="bc"), lambda ctx,bc,x:
+   bc.replace(src=(f2f(x.bitcast(f2f_dt[ctx[1]]), ctx[1], ctx[0]),)) if x.dtype == ctx[1] and bc.dtype.bitsize == ctx[0].bitsize else None),
+  (UPat(Ops.CAST, dtypes.floats, src=(UPat.var("val"),), name="x"), lambda ctx,x,val:
+   f2f_clamp(val.cast(ctx[1]), ctx[0]) if x.dtype.scalar() == ctx[0] else None),
+  (UPat(GroupOp.All-{Ops.BITCAST}, dtypes.floats, name="x"), lambda ctx,x:
+   x.replace(dtype=ctx[1].vec(x.dtype.count), src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src))
+   if x.dtype.scalar() == ctx[0] else None),
+  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat(Ops.BITCAST, dtypes.floats, name="val")), name='st'), lambda ctx,st,idx,val:
+   st.replace(src=(idx, val.replace(dtype=f2f_dt[ctx[0]]))) if val.dtype == ctx[0] and idx.tag == ctx[0] else None),
+  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val", dtypes.floats)), name='st'), lambda ctx,st,idx,val:
+   f2f_store(st, idx, val, *ctx) if val.dtype.scalar() == ctx[1] and (idx:=idx.src[0] if idx.op == Ops.CAST else idx).tag == ctx[0] else None),
+])
@@ -67,7 +67,8 @@ def consumer_map_from_toposort(lst:Iterable[UOp]):
  ret: dict[UOp, dict[UOp, None]] = {}
  for u in lst:
    ret[u] = {}
-    for s in u.src: ret[s][u] = None
+    for s in u.src:
+      if s in ret: ret[s][u] = None
  return ret

 def pretty_print(x:UOp, cache=None, d=0)->str:
@@ -206,7 +207,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    match self.op:
      # late ops don't have shape
      case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
-           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.CUSTOM_KERNEL | Ops.SINK | \
+           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.SINK | \
           Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY:
        return None

@@ -235,9 +236,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
      case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END | Ops.CALL:
        return self.src[0]._shape

-      # ops with custom handling
-      case Ops.KERNEL: return self.arg.ast._shape
-
      # TODO: disallow shape changing bitcast
      case Ops.BITCAST:
        ps = self.src[0]._shape
@@ -287,10 +285,11 @@ 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.ASSIGN: return self.src[1]._shape
+
    # elementwise ops keep the shape the same. all inputs with shape must match
-    if self.op in GroupOp.ALU.union({Ops.CAST, Ops.COPY, Ops.ASSIGN, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE, Ops.STORE}):
-      # TODO: remove this hack for 3 op assign
-      input_shapes = [x._shape for x in (self.src[:2] if self.op is Ops.ASSIGN else self.src) if x._shape is not None]
+    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}")
      return input_shapes[0]
@@ -310,6 +309,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
  def ended_ranges(self):
    if self.op in range_start: return self.src[range_start[self.op]:]
    if self.op is Ops.AFTER: return tuple(flatten([x.ended_ranges for x in self.src[1:]]))
+    # TODO: copy isn't using range properly and isn't ending the range it uses, remove this
+    if self.op in {Ops.COPY, Ops.BUFFER_VIEW}: return self.src[0].ranges
    return ()

  # determine what ranges this is in
@@ -364,9 +365,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
  @recursive_property
  def trace_num(self):
    num = next(ucount)
-    # KERNEL also has a UOp in the arg
-    arg = type(self.arg)(self.arg.ast.trace_num, self.arg.metadata) if self.op is Ops.KERNEL else self.arg
-    uop_fields[num] = (self.op, self.dtype, tuple(s.trace_num for s in self.src), arg, self.tag)+((self.metadata,) if TRACEMETA>=2 else ())
+    uop_fields[num] = (self.op, self.dtype, tuple(s.trace_num for s in self.src), self.arg, self.tag)+((self.metadata,) if TRACEMETA>=2 else ())
    return num

  # *** uop syntactic sugar ***
@@ -539,6 +538,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
  def base(self) -> UOp:
    if self.op in GroupOp.Movement: return self.src[0].base
    if self.op is Ops.MULTI: return self.src[0].base  # MULTI is really a VIEW
+    if self.op is Ops.DETACH: return self.src[0].base  # DETACH can't change base
    return self

  # like gep, but might return an integer
@@ -798,6 +798,11 @@ class UOp(OpMixin, metaclass=UOpMetaClass):

  # *** uop high level syntactic sugar ***

+  @property
+  def shard_shape(self):
+    if self.axis is None: return self.shape
+    return tuple(x//len(self.device) if i == self.axis else x for i,x in enumerate(self.shape))
+
  @staticmethod
  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL):
    lookup = {AddrSpace.GLOBAL: Ops.PARAM, AddrSpace.LOCAL: Ops.DEFINE_LOCAL, AddrSpace.REG: Ops.DEFINE_REG}
@@ -806,7 +811,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    return ret
  def placeholder_like(self, slot:int):
    assert all_int(self.shape), "no placeholder-like on symbolic shape"
-    return UOp.placeholder(self.shape, self.dtype, slot)
+    return UOp.placeholder(self.shard_shape, self.dtype, slot)

  # set is store+end+after
  def set(self:UOp, val:UOp|ConstType, end:UOp|tuple[UOp, ...]|list[UOp]=()) -> UOp:
@@ -818,11 +823,14 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
    src = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),) + (() if device is None else (UOp(Ops.DEVICE, arg=device),))
    return UOp(Ops.PARAM, dtype, src, arg=slot)

-  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=()) -> UOp:
-    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata))
+  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), inline=False) -> UOp:
+    # TODO: reenable this after ENCDEC is fixed
+    #assert len(self.ranges) == 0, f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}"
+    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata, inline))
  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)
-    kernel = UOp(Ops.CUSTOM_KERNEL, src=contig_srcs, arg=CustomKernel(fxn=fxn, grad_fxn=grad_fxn))
+    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
+    kernel = fxn(*placeholders).call(*contig_srcs, grad_fxn=grad_fxn)
    return [s.after(kernel) for s in contig_srcs]

@dataclass(frozen=True)
@@ -836,27 +844,14 @@ class KernelInfo:
  @property
  def function_name(self): return to_function_name(self.name)

-@dataclass(frozen=True)
-class CustomKernel:
-  fxn: Callable
-  grad_fxn: Callable|None = None
-  # sadly CustomKernel can't be pickled or reconstructed as a str
-  def __reduce__(self): return (CustomKernel, (panic,))
-  def __repr__(self): return f"CustomKernel({id(self.fxn)})"
-
@dataclass(frozen=True)
 class CallInfo:
  grad_fxn: Callable|None = None
  metadata: tuple[Metadata, ...] = ()
+  inline: bool = False
  # grad_fxn can't be pickled, but metadata can
-  def __reduce__(self): return (CallInfo, (None, self.metadata))
-  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata})"
-
-@dataclass(frozen=True)
-class Kernel:
-  ast: UOp
-  metadata: tuple[Metadata, ...] = ()
-  grad_fxn: Callable|None = None
+  def __reduce__(self): return (CallInfo, (None, self.metadata, self.inline))
+  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata}, {self.inline})"

 # ******** ops in python ********

@@ -1317,6 +1312,9 @@ def _index_to_concrete_int(u:UOp) -> UOp: return graph_rewrite(u.sink(), pm_lowe
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
 _remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])

+def gate_kernel_sink(x:UOp) -> bool: return not (x.op is Ops.SINK and isinstance(x.arg, KernelInfo))
+pm_gate_kernel_sink = PatternMatcher([(UPat(Ops.SINK, name="sink"), lambda sink: None if gate_kernel_sink(sink) else panic(BottomUpGate))])
+
 def do_unbind(ctx:dict[Variable, int], x:UOp):
  v,i = x.unbind()
  ctx[v] = i
@@ -1414,7 +1412,6 @@ pm_pyrender_extra = PatternMatcher([

 # NOTE: you can remove pm_pyrender_extra and it'll still be correct
 pm_pyrender = pm_pyrender_extra+PatternMatcher([
-  (UPat(Ops.KERNEL, name="u"), lambda ctx,u: f"UOp(Ops.KERNEL, src={srcs(ctx,u.src)}, arg=Kernel({ctx[u.arg.ast]}(), {u.arg.metadata}))"),
  (UPat(GroupOp.All, name="u"), lambda ctx,u: f"UOp({u.op}, {u.dtype}, {srcs(ctx,u.src)}"+(f", {repr(u.arg)})" if u.arg is not None else ")")),
 ])

@@ -1424,7 +1421,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.KERNEL, Ops.WHERE, Ops.END, Ops.ASSIGN}
+                     Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.WHERE, Ops.END, Ops.ASSIGN}

  to_render: set[UOp] = {ast}
  for u in lst:
@@ -1432,7 +1429,7 @@ def pyrender(ast:UOp) -> str:
      for s in u.src: to_render.add(s)
    if u.op is Ops.STORE: to_render.add(u.src[1])
    if u.op in {Ops.REDUCE, Ops.REDUCE_AXIS}: to_render.add(u.src[0])
-    if u.op in {Ops.CUSTOM_KERNEL, Ops.CALL}: raise NotImplementedError("custom_kernel / call can't be pyrendered")
+    if u.op is Ops.CALL: raise NotImplementedError("call can't be pyrendered")
    if u.op in not_rendered: continue
    # checking the consumers is not enough, you have to make sure it's not used twice by the one consumer
    if len(cmap[u]) == 1 and len([x for x in list(cmap[u].keys())[0].src if x is u]) == 1 and u.op not in always_rendered: continue
@@ -1447,11 +1444,6 @@ def pyrender(ast:UOp) -> str:
    op_depth = 1 + max([depth[s] for s in u.src], default=0)
    if op_depth > 100: to_render.add(u)
    depth[u] = 0 if u in to_render else op_depth
-    # do the rendering
-    if u.op is Ops.KERNEL:
-      if u.arg.ast not in kernels:
-        kernels[u.arg.ast] = (f"k{len(kernels)}", f"def k{len(kernels)}():\n  " + pyrender(u.arg.ast).replace('\n', '\n  ') + "\n  return ast\n\n")
-      r[u.arg.ast] = kernels[u.arg.ast][0]
    ren = cast(str, pm_pyrender.rewrite(u, ctx=r))
    assert isinstance(ren, str)
    if u.tag is not None: ren += f".rtag({repr(u.tag)})"
@@ -1,6 +1,6 @@
 import math
 from typing import cast, Any
-from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender, Kernel, CustomKernel
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid, ConstFloat
 from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata, panic, CHECK_OOB

@@ -73,9 +73,6 @@ movement_ops = PatternMatcher([

  # AFTER on Movement Op
  (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.MULTI, Ops.CONTIGUOUS})),), allow_any_len=True), lambda: True),
-
-  # custom kernels allowed here
-  (UPat(Ops.CUSTOM_KERNEL), lambda: True),
 ])

 _tensor_spec = PatternMatcher([
@@ -88,7 +85,7 @@ _tensor_spec = PatternMatcher([
   lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, (DType, ImageDType))),

  # KERNEL can attach to an AFTER to describe the compute required to realize a BUFFER
-  (UPat(Ops.KERNEL, src=UPat((Ops.BUFFER, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
+  (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:])),
@@ -139,12 +136,19 @@ _tensor_spec = PatternMatcher([
  # allow CALL/PARAM
  (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),
-])+movement_ops+shared_spec

-tensor_spec = PatternMatcher([
-  # no tags allowed in tensor graph
-  (UPat(GroupOp.All, name="x"), lambda x: None if x.tag is None else False),
-])+_tensor_spec
+  # ** for custom kernels **
+
+  # codegen: PROGRAM with progressive sources through the pipeline (SINK, DEVICE, LINEAR?, SOURCE?, BINARY?)
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE))), lambda: True),
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR))), lambda: True),
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE))), lambda: True),
+  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE), UPat(Ops.BINARY))), lambda: True),
+  # codegen: standalone LINEAR/SOURCE/BINARY
+  (UPat(Ops.LINEAR, dtypes.void), lambda: True),
+  (UPat(Ops.SOURCE, dtypes.void, src=()), lambda: True),
+  (UPat(Ops.BINARY, dtypes.void, src=()), lambda: True),
+])+movement_ops+shared_spec

 # ***** UOp spec in codegen shared between kernel and program *****

@@ -204,6 +208,11 @@ kernel_spec = PatternMatcher([
  (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype in (dtypes.index, dtypes.int) for y in x.src[1:])),
 ])+movement_ops+shared_codegen_spec+shared_spec

+tensor_spec = PatternMatcher([
+  # no tags allowed in tensor graph
+  (UPat(GroupOp.All, name="x"), lambda x: None if x.tag is None else False),
+])+_tensor_spec+kernel_spec
+
 # ***** UOp spec in linearized programs *****

 program_spec = PatternMatcher([
@@ -235,8 +244,6 @@ full_spec = PatternMatcher([

  # rangeify: buffer view with index or load is okay
  (UPat(Ops.BUFFER_VIEW, src=(UPat((Ops.INDEX, Ops.LOAD)),)), lambda: True),
-  # assign on index. the third op is the shape
-  (UPat(Ops.ASSIGN, src=(UPat(), UPat(), UPat())), lambda: True),

  # expander: unroll/contract/gep/ptrcat/cat
  (UPat((Ops.UNROLL, Ops.CONTRACT), src=(UPat(),)), lambda: True),
@@ -250,7 +257,7 @@ full_spec = PatternMatcher([
  (UPat(Ops.INDEX, src=(UPat((Ops.VECTORIZE, Ops.CAST)), UPat())), lambda: True),

  # linearizer: outputs + intermediate KERNELs
-  (UPat(Ops.KERNEL, dtype=dtypes.void), lambda: True),
+  (UPat(Ops.CALL, dtype=dtypes.void), lambda: True),

  # Invalid must have type Index
  (UPat(Ops.CONST, arg=Invalid, name="x"), lambda x: x.dtype.scalar() == dtypes.index),
@@ -266,16 +273,6 @@ full_spec = PatternMatcher([
  # in progress MSTACK may lose device
  (UPat((Ops.MSELECT, Ops.MSTACK), name="x"), lambda x: True),

-  # codegen: PROGRAM with progressive sources through the pipeline (SINK, DEVICE, LINEAR?, SOURCE?, BINARY?)
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE))), lambda: True),
-  (UPat(Ops.PROGRAM, dtypes.void, src=(UPat(Ops.SINK), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE), UPat(Ops.BINARY))), lambda: True),
-  # codegen: standalone LINEAR/SOURCE/BINARY
-  (UPat(Ops.LINEAR, dtypes.void), lambda: True),
-  (UPat(Ops.SOURCE, dtypes.void, src=()), lambda: True),
-  (UPat(Ops.BINARY, dtypes.void, src=()), lambda: True),
-
  # temp VECTORIZE/INDEX during rewrite have the wrong dtype
  (UPat(Ops.VECTORIZE), lambda: True),
  (UPat(Ops.INDEX), lambda: True),
@@ -303,8 +300,8 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher):
 # late imports to avoid circular import
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.schedule.rangeify import BufferizeOpts
-glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Kernel": Kernel, "Metadata": Metadata,
-                        "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid, "CustomKernel": CustomKernel,
+glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Metadata": Metadata,
+                        "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid,
                        "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace, "panic": panic,
                        "ConstFloat": ConstFloat}
 def eval_pyrender(code:str) -> UOp:
@@ -251,7 +251,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
  ((UPat.var("x", dtypes.index) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
  # only RANGE/IF/STORE/KERNEL have side effects
  (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
-    tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.KERNEL, Ops.BARRIER, Ops.END, Ops.UNROLL} else y.src for y in x.src[1:]])))),
+    tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.BARRIER, Ops.END, Ops.UNROLL} else y.src for y in x.src[1:]])))),
  # after with 1 src is just src[0]
  (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s),
  # VECTORIZE/CONST
@@ -25,9 +25,10 @@ z3_renderer = PatternMatcher([
  (UPat(Ops.SPECIAL, name="x"), lambda x,ctx: create_bounded(x.arg, 0, ctx[1][x.src[0]]-1, ctx[0])),
  (UPat(Ops.DEFINE_VAR, name="x"), lambda x,ctx: create_bounded(x.arg[0], x.arg[1], x.arg[2], ctx[0])),
  (UPat(Ops.RANGE, name="x"), lambda x,ctx: create_bounded(x.render(simplify=False), 0, ctx[1][x.src[0]]-1, ctx[0])),
-  # loads are variables bounded by the min/max of the dtype
-  (UPat(Ops.LOAD, dtypes.ints+(dtypes.index,), name="x"), lambda x,ctx: create_bounded(f"load{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])),
-  (UPat(Ops.LOAD, dtypes.bool, name="x"), lambda x,ctx: (z3.Bool(f"load{len(ctx[1])}", ctx=ctx[0]), None)),
+  # loads are variables bounded by the min/max of the dtype. non-pointer INDEX is also a LOAD
+  (UPat((Ops.LOAD, Ops.INDEX), dtypes.ints+(dtypes.index,), name="x"), lambda x,ctx:
+    create_bounded(f"load{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])),
+  (UPat((Ops.LOAD, Ops.INDEX), dtypes.bool, name="x"), lambda x,ctx: (z3.Bool(f"load{len(ctx[1])}", ctx=ctx[0]), None)),
  # constants
  (UPat(Ops.CONST, arg=Invalid, name="x"), lambda x,ctx: (z3.Int("Invalid", ctx=ctx[0]), None)),
  (UPat(Ops.CONST, dtypes.ints+(dtypes.index,), name="x"), lambda x,ctx: (z3.IntVal(x.arg, ctx=ctx[0]), None)),
@@ -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",
               Ops.PARAM:"#cb9037", **{x:"#f2cb91" for x in {Ops.DEFINE_LOCAL, Ops.DEFINE_REG}}, Ops.REDUCE_AXIS: "#FF6B6B",
               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.KERNEL: "#3e7f55", Ops.CUSTOM_KERNEL: "#3ebf55",
+               Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff",
               **{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.ENCDEC: "#bf71b6",
               Ops.CALL: "#00B7C8", Ops.PARAM: "#14686F",
@@ -106,9 +106,6 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
    if u in excluded: continue
    argst = codecs.decode(str(u.arg), "unicode_escape")
    if u.op in GroupOp.Movement: argst = (mask_to_str if u.op in {Ops.SHRINK, Ops.PAD} else shape_to_str)(u.marg)
-    if u.op is Ops.KERNEL:
-      ast_str = f"SINK{tuple(s.op for s in u.arg.ast.src)}" if u.arg.ast.op is Ops.SINK else repr(u.arg.ast.op)
-      argst = f"<Kernel {len(list(u.arg.ast.toposort()))} {ast_str} {[str(m) for m in u.arg.metadata]}>"
    if u.op is Ops.BINARY: argst = f"<{len(u.arg)} bytes>"
    label = f"{str(u.op).split('.')[1]}{(chr(10)+word_wrap(argst.replace(':', ''))) if u.arg is not None else ''}"
    if u.dtype != dtypes.void: label += f"\n{u.dtype}"
@@ -130,9 +127,9 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
        label += "\n"+' '.join([f"{range_str(s, color=True)}({s.vmax+1})" for s in trngs])
    except Exception:
      label += "\n<ISSUE GETTING LABEL>"
-    if (ref:=ref_map.get(u.arg.ast) if u.op is Ops.KERNEL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}"
+    if (ref:=ref_map.get(u.src[0]) if u.op is Ops.CALL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}"
    # NOTE: kernel already has metadata in arg
-    if TRACEMETA >= 2 and u.metadata is not None and u.op is not Ops.KERNEL: label += "\n"+str(u.metadata)
+    if TRACEMETA >= 2 and u.metadata is not None and u.op is not Ops.CALL: label += "\n"+str(u.metadata)
    graph[id(u)] = {"label":label, "src":[(i,id(x)) for i,x in enumerate(u.src) if x not in excluded], "color":uops_colors.get(u.op, "#ffffff"),
                    "ref":ref, "tag":repr(u.tag) if u.tag is not None else None}
  return graph
@@ -140,12 +137,10 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
@functools.cache
 def _reconstruct(a:int):
  op, dtype, src, arg, *rest = trace.uop_fields[a]
-  arg = type(arg)(_reconstruct(arg.ast), arg.metadata) if op is Ops.KERNEL else arg
  return UOp(op, dtype, tuple(_reconstruct(s) for s in src), arg, *rest)

 def get_full_rewrite(ctx:TrackedGraphRewrite) -> Generator[GraphRewriteDetails, None, None]:
  next_sink = _reconstruct(ctx.sink)
-  # in the schedule graph we don't show indexing ops (unless it's in a kernel AST or rewriting dtypes.index sink)
  yield {"graph":uop_to_json(next_sink), "uop":pystr(next_sink), "change":None, "diff":None, "upat":None}
  replaces: dict[UOp, UOp] = {}
  for u0_num,u1_num,upat_loc,dur in tqdm(ctx.matches):
@@ -168,19 +163,19 @@ def option(s:int|None) -> int: return 0 if s is None else s+1

 # Profiler API

-device_ts_diffs:dict[str, tuple[Decimal, Decimal]] = {}
-def cpu_ts_diff(device:str, thread=0) -> Decimal: return device_ts_diffs.get(device, (Decimal(0),))[thread]
+device_ts_diffs:dict[str, Decimal] = {}
+def cpu_ts_diff(device:str) -> Decimal: return device_ts_diffs.get(device, Decimal(0))

 amdgpu_targets:dict[str, int] = {}

 DevEvent = ProfileRangeEvent|ProfileGraphEntry|ProfilePointEvent
 def flatten_events(profile:list[ProfileEvent]) -> Generator[tuple[Decimal, Decimal, DevEvent], None, None]:
  for e in profile:
-    if isinstance(e, ProfileRangeEvent): yield (e.st+(diff:=cpu_ts_diff(e.device, e.is_copy)), (e.en if e.en is not None else e.st)+diff, e)
+    if isinstance(e, ProfileRangeEvent): yield (e.st+(diff:=cpu_ts_diff(e.device)), (e.en if e.en is not None else e.st)+diff, e)
    elif isinstance(e, ProfilePointEvent): yield (e.ts, e.ts, e)
    elif isinstance(e, ProfileGraphEvent):
      cpu_ts = []
-      for ent in e.ents: cpu_ts += [e.sigs[ent.st_id]+(diff:=cpu_ts_diff(ent.device, ent.is_copy)), e.sigs[ent.en_id]+diff]
+      for ent in e.ents: cpu_ts += [e.sigs[ent.st_id]+(diff:=cpu_ts_diff(ent.device)), e.sigs[ent.en_id]+diff]
      yield (st:=min(cpu_ts)), (et:=max(cpu_ts)), ProfileRangeEvent(f"{e.ents[0].device.split(':')[0]} Graph", f"batched {len(e.ents)}", st, et)
      for i,ent in enumerate(e.ents): yield (cpu_ts[i*2], cpu_ts[i*2+1], ent)

@@ -375,16 +370,19 @@ def unpack_sqtt(key:tuple[str, int], data:list, p:ProfileProgramEvent) -> tuple[

 def device_sort_fn(k:str) -> tuple[int, str, int]:
  order = {"GC": 0, "USER": 1, "TINY": 2, "DISK": 999}
-  dname = k.split()[0]
+  dname, *rest = k.split()
  dev_rank = next((v for k,v in order.items() if dname.startswith(k)), len(order))
-  return (dev_rank, dname, len(k))
+  if len(parts:=dname.split(":")) < 2 or not parts[1].isdigit(): parts.insert(1, "0")
+  eng_rank = 2 if rest else 1 if len(parts) > 2 else 0
+  # 3 levels of hierarchy: device class, index in multi device, engine within device
+  return (dev_rank, parts[1], eng_rank)

 def get_profile(profile:list[ProfileEvent], sort_fn:Callable[[str], Any]=device_sort_fn) -> bytes|None:
  # start by getting the time diffs
  device_decoders:dict[str, Callable[[list[ProfileEvent]], None]] = {}
  for ev in profile:
    if isinstance(ev, ProfileDeviceEvent):
-      device_ts_diffs[ev.device] = (ev.comp_tdiff,ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff)
+      device_ts_diffs[ev.device] = ev.tdiff
      if (d:=ev.device.split(":")[0]) == "AMD":
        device_decoders[d] = load_counters
        amdgpu_targets[d] = unwrap(ev.props)["gfx_target_version"]
Author	SHA1	Message	Date
George HotzandGitHub	2fe45b0660	Merge branch 'master' into call_inline	2026-02-10 14:12:08 +08:00
sirhcmandGitHub	cdb78954cb	better cl compiler name (#14660 ) cl_compiler instead of compiler because overriding Compiled.compiler seems more confusing	2026-02-10 01:03:46 -05:00
George HotzandGitHub	1d88723aa0	Merge branch 'master' into call_inline	2026-02-10 14:02:43 +08:00
George HotzandGitHub	cc9bf8ccbc	move more to null/unit tests (#14658 ) * move more to null tests * move test_gc * no test fusion op	2026-02-10 13:35:17 +08:00
geohot	b0dd3af093	inline all these calls	2026-02-10 13:25:29 +08:00
chenyuandGitHub	83f6d28579	two less realize in setitem (#14655 )	2026-02-09 23:45:24 -05:00
geohot	e89221e9aa	add inline flag for call	2026-02-10 12:19:51 +08:00
wozeparrotandGitHub	69574542ab	fix: use correct fa implementation in eval (#14651 )	2026-02-09 18:20:44 -08:00
chenyuandGitHub	0dedf4063c	minor test_setitem cleanup (#14654 )	2026-02-09 20:40:29 -05:00
sirhcmandGitHub	b36b62eb59	don't push docker cache for PRs (#14652 )	2026-02-09 19:55:55 -05:00
sirhcmandGitHub	e6562a5061	remove CompilerPair (#14638 )	2026-02-09 19:51:18 -05:00
sirhcmandGitHub	396e1320fb	bump cache version for z3 (#14650 )	2026-02-09 19:32:07 -05:00
chenyuandGitHub	9e3f24db9f	assign realize fix (#14649 ) fix the need for explicit assign. track pending assigns for each buffer, and run those before the main realize in order	2026-02-09 17:46:46 -05:00
chenyuandGitHub	0913c068ea	clean up setitem disk path (#14648 )	2026-02-09 15:58:04 -05:00
chenyuandGitHub	205a1212b7	delegate non Tensor src setitem to assign (#14647 ) cannot do this for DISK in the unified path	2026-02-09 13:53:20 -05:00
chenyuandGitHub	e9f40f49d4	explicitly check advanced setitem (#14644 ) advanced setitem DISK would failed in rangeify with bad error, now it's checked directly in setitem. eventully DISK can use regular setitem path	2026-02-09 13:36:46 -05:00
chenyuandGitHub	20a132b1c4	relax atol for test_uop_scan_matmul (#14646 ) flaky, also log max diff	2026-02-09 13:25:19 -05:00
qazalandGitHub	50d3f6cea5	EVAL_BS=0 in llama profile (#14643 )	2026-02-10 00:49:43 +09:00
chenyuandGitHub	8a2c23d3dc	raise RuntimeError for setitem dtype mismatch (#14642 )	2026-02-09 10:37:08 -05:00
qazalandGitHub	80b0119cef	llama: add new asm gemm shape (#14611 ) * llama: add new asm gemm shape * work * cleanup * half dtype * more comment	2026-02-10 00:34:29 +09:00
chenyuandGitHub	a49e038c0c	dont manually broadcast in setitem (#14641 ) handled by assign	2026-02-09 09:34:09 -05:00
chenyuandGitHub	2c3e3559eb	remove a contiguous in basic setitem (#14640 ) handled in rangeify	2026-02-09 09:19:46 -05:00
chenyuandGitHub	6c0c8e2ac3	setitem push a realize to basic setitem (#14637 ) advanced setitem does not need it	2026-02-09 08:54:07 -05:00
nimlgenandGitHub	e087c58ae0	print tables in llama/profile.sh (#14639 )	2026-02-09 12:32:54 +03:00
sirhcmandGitHub	27f7ea478b	new style DSP renderer (#14636 ) * new style DSP renderer * cleanup	2026-02-09 00:39:03 -05:00
sirhcmandGitHub	efac5b9ef6	new style NV/CUDA renderers, try 2 (#14634 ) * new style NV/CUDA renderers, try 2 * fix diskcache	2026-02-08 22:58:48 -05:00
sirhcmandGitHub	0ebb508b85	new style metal compiler (#14632 )	2026-02-08 21:58:25 -05:00
sirhcmandGitHub	9eef9f38ad	new style python renderer (#14631 )	2026-02-08 21:45:07 -05:00
sirhcmandGitHub	5f2f2cc956	Revert "new style NV/CUDA renderers (#14627 )" (#14633 ) This reverts commit `0e505951b0`.	2026-02-08 21:16:03 -05:00
sirhcmandGitHub	4ad787ece2	new style CPULLVMRenderer (#14629 )	2026-02-08 21:05:01 -05:00
sirhcmandGitHub	0e505951b0	new style NV/CUDA renderers (#14627 ) * new style NV/CUDA renderers * fix pickle * oops * fix CUDA_CC=NVCC * mockgpu uses PTXCompiler * oops * ruff * dont discard stderr * ugh	2026-02-08 21:04:51 -05:00
Filip BrzekandGitHub	1667669c46	fix: python3 -m tinygrad.device reporting on AMD/CPU (#14622 ) * test: device module expects PASS in -m tinygrad.device for CPU * fix: use device._compiler_name instead of unwrap_class_type(compiler).__name__ in enumerate_devices_str	2026-02-08 20:22:35 +03:00
nimlgenandGitHub	01a4ee4d66	do not hive_reset when amdgpu (#14624 )	2026-02-08 19:14:13 +03:00
nimlgenandGitHub	a615b9d781	am: f8_mode for gfx94x only (#14620 )	2026-02-08 17:38:48 +03:00
chenyuandGitHub	c28f7d0167	remove realize in Tensor.svd (#14623 )	2026-02-08 09:36:31 -05:00
qazalandGitHub	087dab4c3b	gemm/asm: split out cdna tests from CI (#14619 ) * gemm/asm: split out cdna tests from CI * reorder * work	2026-02-08 21:33:42 +09:00
George HotzandGitHub	183d38b128	remove CUSTOM_KERNEL / directly construct it (#14604 ) * remove CUSTOM_KERNEL / directly construct it * clean that up * simpler multi * custom kernel spec * remove Kernel * fix multi * use sharded shape * explicit regression test	2026-02-08 18:43:33 +08:00
nimlgenandGitHub	e29a88ca09	hive_reset respects lock (#14618 )	2026-02-08 10:47:25 +03:00
qazalandGitHub	b10802eb53	use existing VIZ ContextVar instead of getenv (#14610 )	2026-02-08 15:37:55 +09:00
chenyuandGitHub	510b65489e	style change rangeify assign [pr] (#14616 ) consistent naming, also a standalone fucntion to replace complicated lambda	2026-02-07 15:47:32 -05:00
chenyuandGitHub	b7afd4471c	use arg instead of 3rd op for ASSIGN [pr] (#14613 )	2026-02-07 14:17:10 -05:00
nimlgenandGitHub	88c3022223	amd: kfd iface early exit (#14612 ) * amd: kfd iface early exit * l * revert	2026-02-07 18:57:10 +03:00
nimlgenandGitHub	ce7bfc6ce8	nv: use nv_flags for all fields (#14607 )	2026-02-07 15:01:38 +03:00
qazalandGitHub	c2544e2252	viz: remove outdated comment (#14608 )	2026-02-07 20:05:24 +09:00
nimlgenandGitHub	6838b35cff	mockgpu: hevc (#14606 ) * mockgpu: hevc * eng	2026-02-07 12:27:55 +03:00
chenyuandGitHub	884592f6c8	pin z3-solver version (#14605 ) found exact input that crashes z3 4.15.4	2026-02-06 22:49:31 -05:00
George HotzandGitHub	7a2a3b5c71	Remove Ops.KERNEL, it's all Ops.CALL now (#14603 )	2026-02-07 10:21:54 +08:00
George HotzandGitHub	ca6604eae2	kernel is call (#14577 ) * call is kernel * closer * fix bugs * dedup * pm_gate_kernel_sink * better * Revert "better" This reverts commit `b4c799b810`. * Reapply "better" This reverts commit `e53f094ce7`. * cleanups * work * remove junk * subtle fix * index * viz cleanups * disable assert for now	2026-02-07 10:10:14 +08:00
wozeparrotandGitHub	d87ae1c84c	feat: tinyfs load test in benchmark (#14602 )	2026-02-06 18:00:00 -08:00
ttomsaandGitHub	462b455562	cleanup linearize (#14523 )	2026-02-07 08:54:02 +08:00
ttomsaandGitHub	d5652e4da2	new dtype aliases (#14596 )	2026-02-07 08:53:35 +08:00
sirhcmandGitHub	ad9e2f0de7	decompose bf16 (#14601 )	2026-02-06 19:24:09 -05:00
sirhcmandGitHub	7bb45e7df0	decompose fp8 to bigger floats [skip_process_replay] (#14554 ) * decompose fp8 also * it works * cleanup * no shift required * default to float * cleanup * fixes * fp8e5m2 * don't rely on behavior comparing nans * cleanup	2026-02-06 19:05:40 -05:00
chenyuandGitHub	81f6cdb4ab	delete realize_assign [pr] (#14575 ) use realize and realize_srcs like COPY and STORE. src[0] always has BUFFER for base	2026-02-06 17:12:33 -05:00
chenyuandGitHub	7d193a6e26	fix wgsl bitcast (#14600 ) was wrong for signed int	2026-02-06 16:57:36 -05:00
chenyuandGitHub	b9fe8b7591	fix opt in process replay [pr] (#14599 )	2026-02-06 16:49:56 -05:00
chenyuandGitHub	197ebcbbbc	log seed with flush=True in fuzz_symbolic (#14597 ) * log seed with flush=True in fuzz_symbolic i think z3 can crash. added reading seed from argv to see if we repro later * fuzz_symbolic_symbolic_div	2026-02-06 15:03:57 -05:00
nimlgenandGitHub	fbb67a3f95	am_smi: fix after regen (#14594 )	2026-02-06 20:57:41 +03:00
qazalandGitHub	a80fb4e641	viz: better ordering of device engines in profiler (#14590 )	2026-02-06 23:08:09 +09:00
qazalandGitHub	b7e3fbe07e	llama: add VIZ=-1 to dev_run (#14583 ) * llama: add VIZ=-1 to dev_run * readme * cleaner * add profile.sh script * better grouping of options * add other row * readme edits * work	2026-02-06 22:59:22 +09:00
nimlgenandGitHub	fbeb978170	diff devices for sdma (#14589 ) * start * x * fix * sdma * c * clean * x * hm * cleaer	2026-02-06 16:39:12 +03:00
George HotzandGitHub	7cb996e153	bottom up earliest rewrites (#14587 ) * better * bottom up earliest rewrites * fix	2026-02-06 18:13:07 +08:00
George HotzandGitHub	03af2404e2	small changes and test fixes from kernel is call (#14586 )	2026-02-06 17:08:33 +08:00