Merge branch 'master' into kernel_is_call

.github/workflows/benchmark.yml

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

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh

@@ -1,4 +0,0 @@
-#!/bin/bash
-export BENCHMARK=5
-export VIZ=${VIZ:--1}
-examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

extra/amdpci/am_smi.py

@@ -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:", "Processing accelerators:"]
+    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
     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,8 +153,7 @@ class SMICtx:
     tables = {}
     for dev in self.devs:
       match dev.ip_ver[am.MP1_HWIP]:
-        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 (13,0,6)|(13,0,12): table_t = dev.smu.smu_mod.MetricsTableX_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
@@ -280,7 +279,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, activity_line_width, opt_text=mem_fmt)}"] \
+                    + [f"MEM Usage    {draw_bar((mem_used / mem_total) / 100, 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()]

extra/amdpci/hive_reset.py

@@ -1,18 +1,12 @@
 #!/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])])
-  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]
+  pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
   amdevs = []
   with Context(DEBUG=2):
     for pcidev in pcidevs:

extra/fp8/fp8_linear.py

@@ -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[1:]
+  _, input_uop, weight_uop = kernel.src
   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)

extra/gemm/asm/cdna/asm.py

@@ -9,8 +9,7 @@ GEMM_ARGS = {
   (8192, 4096, 4096): (256, 64, 32768),
   (8192, 14336, 4096): (256, 64, 114688),
   (8192, 4096, 14336): (256, 224, 114688),
-  # TODO: get a fast gemm for this shape
-  #(8192, 128256, 4096): (16032, 64, 1026048),
+  (8192, 128256, 4096): (16032, 64, 1026048),
   (8192, 8192, 8192): (256, 128, 131072),
   (4096, 4096, 4096): (256, 64, 16384),
   (4096, 14336, 4096): (256, 64, 57344),

extra/gemm/asm/cdna/gemm.py

@@ -62,7 +62,7 @@ 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[1:]
+  out, a, b = kernel.src
   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)
   grad_a = (g_t @ b_t.T).uop

extra/models/llama.py

@@ -55,7 +55,7 @@ class Attention:
       xqkv = x @ self.wqkv.T
       xq, xk, xv = xqkv.split([self.wq.weight.shape[0], self.wk.weight.shape[0], self.wv.weight.shape[0]], dim=2)
     else:
-      xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)
+      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
 
     if self.q_norm is not None and self.k_norm is not None:
       xq = self.q_norm(xq)
@@ -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[2]))
-        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[3]))
-        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[4]))
+        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]))
         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]

extra/perfetto/to_perfetto.py

@@ -2,26 +2,27 @@ 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, 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}
+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)}
 def dev_ev_to_perfetto_json(ev:ProfileDeviceEvent):
-  devices[ev.device] = (ev.tdiff, len(devices))
+  devices[ev.device] = (ev.comp_tdiff, ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_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": 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"}}]
 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), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device)}]
+  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)}]
 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), "dur": float(en-st), **dev_to_pid(e.device)}]
+    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)}]
     for dep in ev.deps[i]:
       d = ev.ents[dep]
-      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"}]
+      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"}]
   return ret
 def to_perfetto(profile:list[ProfileEvent]):
   # Start json with devices.

extra/sqtt/rgptool.py

@@ -145,8 +145,7 @@ class RGP:
   @staticmethod
   def from_profile(profile_pickled, device:str|None=None):
     profile: list[ProfileEvent] = pickle.loads(profile_pickled)
-    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)}
+    device_events = {x.device:x for x in profile if isinstance(x, ProfileDeviceEvent) and x.device.startswith('AMD')}
     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())}")

extra/tinyfs/fetch_file.py

@@ -1,36 +1,11 @@
 from tinygrad.tensor import Tensor
-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]
+import argparse
 
 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!")

extra/viz/README

@@ -1,18 +0,0 @@
-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"`

extra/viz/cli.py

@@ -26,24 +26,16 @@ def print_data(data:dict) -> None:
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  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_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)))
+  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)))
   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)
@@ -52,10 +44,9 @@ 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(f"  {k}")
+      print(k)
       if args.device is None: continue
       for e in v.get("events", []):
         et = e["dur"]*1e-6
@@ -70,14 +61,11 @@ if __name__ == "__main__":
           a[0] += et
           a[1] += 1
           total += et
-    if agg and total > 0:
-      items = sorted(agg.items(), key=lambda kv:kv[1][0], reverse=True)
-      sel = items[:10]
-      table = [[name, time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in sel]
-      if (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"))
+
+    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"))
     exit(0)
 
   for k in viz.ctxs:

pyproject.toml

@@ -67,7 +67,7 @@ testing_minimal = [
   "pytest-timeout",
   "pytest-split",
   "hypothesis>=6.148.9",
-  "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
+  "z3-solver",
 ]
 testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "ggml-python"]
 testing = [

test/external/fuzz_symbolic.py

@@ -1,16 +1,12 @@
-# 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 random, operator
 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 = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
-print(f"Seed: {seed}", flush=True)
+seed = random.randint(0, 100)
+print(f"Seed: {seed}")
 random.seed(seed)
 
 unary_ops = [lambda a:a+random.randint(-4, 4), lambda a: a*random.randint(-4, 4),

test/external/fuzz_symbolic_symbolic_div.py

@@ -1,11 +1,11 @@
-import random, sys
+import random
 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 = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
-print(f"Seed: {seed}", flush=True)
+seed = random.randint(0, 100)
+print(f"Seed: {seed}")
 random.seed(seed)
 
 def get_random_term(ranges, factors):

test/external/process_replay/process_replay.py

@@ -47,15 +47,13 @@ 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.CALL]
+    asts = [repr(u.arg.ast) for u in ret.toposort() if u.op is Ops.KERNEL]
     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()
-  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)
+  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)
   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:

test/helpers.py

@@ -62,7 +62,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
   g = UOp(Ops.PARAM, uop.dtype.ptr(), arg=0, src=())
-  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer())
+  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(), PythonRenderer())
   prog = PythonProgram("run", PythonCompiler().compile(prg.src))
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
   return out_buf.cast(uop.dtype.fmt or "").tolist()[0]

test/mockgpu/nv/nvdriver.py

@@ -130,18 +130,15 @@ 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:
-      parent = self.object_by_handle.get(struct.hObjectParent)
-      assert parent is not None and isinstance(parent, (NVChannelGroup, NVGPU))
+      assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVChannelGroup)
       struct.hObjectNew = self._alloc_handle()
       params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS.from_address(params_ptr)
-      gpu = parent.device if isinstance(parent, NVChannelGroup) else parent
+      gpu = self.object_by_handle[struct.hObjectParent].device
       gpfifo_token = gpu.add_gpfifo(params.gpFifoOffset, params.gpFifoEntries)
       self.object_by_handle[struct.hObjectNew] = NVGPFIFO(gpu, gpfifo_token)
-    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):
+    elif struct.hClass == nv_gpu.AMPERE_DMA_COPY_B or struct.hClass == nv_gpu.ADA_COMPUTE_A:
       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:
@@ -197,7 +194,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 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.NVA06C_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:

test/mockgpu/nv/nvgpu.py

@@ -26,7 +26,6 @@ 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
@@ -116,11 +115,9 @@ 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: # 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.NVC6B5_LAUNCH_DMA: return 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
 
@@ -139,27 +136,6 @@ 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}"

test/null/test_device.py

@@ -102,8 +102,8 @@ class TestCompiler(unittest.TestCase):
 
 class TestRunAsModule(unittest.TestCase):
   def test_module_runs(self):
-    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}")
+    out = '\n'.join(enumerate_devices_str())
+    self.assertIn("CPU", out) # for sanity check
 
 if __name__ == "__main__":
   unittest.main()

test/null/test_egraph.py

@@ -1,502 +0,0 @@
-import unittest
-from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import Ops, UOp, GroupOp, PatternMatcher, UPat, graph_rewrite
-from tinygrad.uop.egraph import uf_find, uf_union, rewrite_all, EGraph, egraph_saturate, egraph_extract, node_cost, _rebuild_tree
-
-# *** test union-find ***
-
-class TestUnionFind(unittest.TestCase):
-  def test_find_self(self):
-    a, b = UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)
-    parent = {a: a, b: b}
-    self.assertIs(uf_find(parent, a), a)
-    self.assertIs(uf_find(parent, b), b)
-
-  def test_union_basic(self):
-    a, b = UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)
-    parent = {a: a, b: b}
-    size = {a: 1, b: 1}
-    root = uf_union(parent, size, a, b)
-    self.assertIs(uf_find(parent, a), uf_find(parent, b))
-    self.assertIs(root, uf_find(parent, a))
-
-  def test_union_chain(self):
-    a, b, c = UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2), UOp.const(dtypes.int, 3)
-    parent = {a: a, b: b, c: c}
-    size = {a: 1, b: 1, c: 1}
-    uf_union(parent, size, a, b)
-    uf_union(parent, size, b, c)
-    self.assertIs(uf_find(parent, a), uf_find(parent, c))
-
-  def test_union_idempotent(self):
-    a, b = UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)
-    parent = {a: a, b: b}
-    size = {a: 1, b: 1}
-    r1 = uf_union(parent, size, a, b)
-    r2 = uf_union(parent, size, a, b)
-    self.assertIs(r1, r2)
-
-# *** test rewrite_all ***
-
-class TestRewriteAll(unittest.TestCase):
-  def test_single_match(self):
-    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
-    a = UOp.variable("a", 0, 10)
-    results = rewrite_all(pm, a + 0)
-    self.assertEqual(len(results), 1)
-    self.assertIs(results[0], a)
-
-  def test_no_match(self):
-    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    results = rewrite_all(pm, a + b)
-    self.assertEqual(len(results), 0)
-
-  def test_multiple_matches(self):
-    pm = PatternMatcher([
-      (UPat.var("x") + 0, lambda x: x),
-      (UPat.var("x") * 1, lambda x: x),
-    ])
-    a = UOp.variable("a", 0, 10)
-    results = rewrite_all(pm, a + 0)
-    self.assertEqual(len(results), 1)
-    self.assertIs(results[0], a)
-
-  def test_both_rules_fire(self):
-    pm = PatternMatcher([
-      (UPat.var("x") + UPat.var("x"), lambda x: x * 2),
-      (UPat.var("x") + UPat.var("x"), lambda x: UOp(Ops.SHL, x.dtype, (x, x.const_like(1)))),
-    ])
-    a = UOp.variable("a", 0, 10)
-    results = rewrite_all(pm, a + a)
-    self.assertEqual(len(results), 2)
-
-  def test_const_folding(self):
-    pm = PatternMatcher([
-      (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name="a"),
-       lambda a: a.const_like(a.src[0].arg + a.src[1].arg) if a.op is Ops.ADD else None),
-    ])
-    results = rewrite_all(pm, UOp.const(dtypes.int, 3) + UOp.const(dtypes.int, 4))
-    self.assertEqual(len(results), 1)
-    self.assertEqual(results[0].arg, 7)
-
-# *** test EGraph class ***
-
-class TestEGraphClass(unittest.TestCase):
-  def test_init(self):
-    a = UOp.variable("a", 0, 10)
-    expr = a + 0
-    eg = EGraph(expr)
-    self.assertEqual(len(eg.eclass), len(list(expr.toposort())))
-    self.assertIn(expr, eg.all_nodes)
-
-  def test_add_node(self):
-    a = UOp.variable("a", 0, 10)
-    eg = EGraph(a)
-    b = UOp.variable("b", 0, 10)
-    eg._add_node(b)
-    self.assertIn(b, eg.all_nodes)
-
-  def test_merge(self):
-    a = UOp.variable("a", 0, 10)
-    expr = a + 0
-    eg = EGraph(expr)
-    result = eg._merge(expr, a)
-    self.assertIsNotNone(result)
-    self.assertIs(uf_find(eg.parent, expr), uf_find(eg.parent, a))
-
-  def test_merge_idempotent(self):
-    a = UOp.variable("a", 0, 10)
-    eg = EGraph(a)
-    result = eg._merge(a, a)
-    self.assertIsNone(result)
-
-# *** test egraph_saturate ***
-
-class TestEGraphSaturate(unittest.TestCase):
-  def test_identity_rules(self):
-    pm = PatternMatcher([
-      (UPat.var("x") + 0, lambda x: x),
-      (UPat.var("x") * 1, lambda x: x),
-    ])
-    a = UOp.variable("a", 0, 10)
-    expr = a + 0
-    eclass = egraph_saturate(expr, pm)
-    # a+0 and a should be in the same e-class
-    a_class = expr_class = None
-    for canon, members in eclass.items():
-      if a in members: a_class = canon
-      if expr in members: expr_class = canon
-    self.assertIsNotNone(a_class)
-    self.assertIsNotNone(expr_class)
-    self.assertIs(a_class, expr_class)
-
-  def test_const_fold_saturation(self):
-    from tinygrad.uop.symbolic import symbolic_simple
-    c2, c3 = UOp.const(dtypes.int, 2), UOp.const(dtypes.int, 3)
-    expr = c2 + c3
-    eclass = egraph_saturate(expr, symbolic_simple)
-    c5 = UOp.const(dtypes.int, 5)
-    for canon, members in eclass.items():
-      if expr in members:
-        self.assertIn(c5, members, f"expected CONST(5) in eclass of 2+3, got {members}")
-        return
-    self.fail("expr not found in any eclass")
-
-  def test_no_rules_match(self):
-    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    eclass = egraph_saturate(a + b, pm)
-    for canon, members in eclass.items():
-      self.assertEqual(len(members), 1)
-
-  def test_max_iters_respected(self):
-    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
-    a = UOp.variable("a", 0, 10)
-    expr = a + 0
-    eclass = egraph_saturate(expr, pm, max_iters=1)
-    a_class = expr_class = None
-    for canon, members in eclass.items():
-      if a in members: a_class = canon
-      if expr in members: expr_class = canon
-    self.assertIs(a_class, expr_class)
-
-  def test_rebuilding_propagates(self):
-    """After a*0 merges with 0, rebuilding should create (0+a) which then matches x+0 -> x."""
-    pm = PatternMatcher([
-      (UPat.var("x") * 0, lambda x: x.const_like(0)),
-      (UPat.var("x") + 0, lambda x: x),
-    ])
-    a = UOp.variable("a", 0, 10)
-    expr = (a * 0) + a
-    eclass = egraph_saturate(expr, pm)
-    expr_cls = a_cls = None
-    for canon, members in eclass.items():
-      if expr in members: expr_cls = canon
-      if a in members: a_cls = canon
-    self.assertIsNotNone(expr_cls)
-    self.assertIsNotNone(a_cls)
-    self.assertIs(expr_cls, a_cls)
-
-  def test_rebuilding_chain(self):
-    """((a*0)+0)+b should simplify to b through multiple rebuild steps."""
-    pm = PatternMatcher([
-      (UPat.var("x") * 0, lambda x: x.const_like(0)),
-      (UPat.var("x") + 0, lambda x: x),
-    ])
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    expr = ((a * 0) + 0) + b
-    eclass = egraph_saturate(expr, pm)
-    expr_cls = b_cls = None
-    for canon, members in eclass.items():
-      if expr in members: expr_cls = canon
-      if b in members: b_cls = canon
-    self.assertIsNotNone(expr_cls)
-    self.assertIsNotNone(b_cls)
-    self.assertIs(expr_cls, b_cls)
-
-# *** test egraph_extract ***
-
-class TestEGraphExtract(unittest.TestCase):
-  def test_extract_identity(self):
-    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
-    a = UOp.variable("a", 0, 10)
-    self.assertIs(egraph_extract(a + 0, pm), a)
-
-  def test_extract_mul_identity(self):
-    pm = PatternMatcher([(UPat.var("x") * 1, lambda x: x)])
-    a = UOp.variable("a", 0, 10)
-    self.assertIs(egraph_extract(a * 1, pm), a)
-
-  def test_extract_const_fold(self):
-    from tinygrad.uop.symbolic import symbolic_simple
-    result = egraph_extract(UOp.const(dtypes.int, 2) + UOp.const(dtypes.int, 3), symbolic_simple)
-    self.assertEqual(result.op, Ops.CONST)
-    self.assertEqual(result.arg, 5)
-
-  def test_extract_chain(self):
-    pm = PatternMatcher([
-      (UPat.var("x") + 0, lambda x: x),
-      (UPat.var("x") * 1, lambda x: x),
-    ])
-    a = UOp.variable("a", 0, 10)
-    self.assertIs(egraph_extract((a + 0) * 1, pm), a)
-
-  def test_extract_no_change(self):
-    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    self.assertIs(egraph_extract(a + b, pm), a + b)
-
-  def test_extract_prefers_cheaper(self):
-    pm = PatternMatcher([(UPat.var("x") + UPat.var("x"), lambda x: x * 2)])
-    a = UOp.variable("a", 0, 10)
-    result = egraph_extract(a + a, pm)
-    self.assertEqual(result.op, Ops.ADD)  # ADD cost 1 < MUL cost 2
-
-  def test_extract_with_symbolic_simple(self):
-    from tinygrad.uop.symbolic import symbolic_simple
-    a = UOp.variable("a", 0, 10)
-    self.assertIs(egraph_extract((a + 0) * 1, symbolic_simple), a)
-
-  def test_combine_terms(self):
-    from tinygrad.uop.symbolic import symbolic
-    a = UOp.variable("a", 0, 10)
-    result = egraph_extract(a * 3 + a * 4, symbolic)
-    self.assertEqual(result.op, Ops.MUL)
-    self.assertEqual(result.src[1].arg, 7)
-
-  # *** tests that REQUIRE rebuilding ***
-
-  def test_rebuild_mul_zero_plus(self):
-    pm = PatternMatcher([
-      (UPat.var("x") * 0, lambda x: x.const_like(0)),
-      (UPat.var("x") + 0, lambda x: x),
-    ])
-    a = UOp.variable("a", 0, 10)
-    self.assertIs(egraph_extract((a * 0) + a, pm), a)
-
-  def test_rebuild_nested_zero(self):
-    pm = PatternMatcher([
-      (UPat.var("x") * 0, lambda x: x.const_like(0)),
-      (UPat.var("x") + 0, lambda x: x),
-    ])
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    self.assertIs(egraph_extract(((a * 0) + 0) + b, pm), b)
-
-  def test_rebuild_distribute_then_fold(self):
-    pm = PatternMatcher([(UPat.var("x") * 0, lambda x: x.const_like(0))])
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    result = egraph_extract((a + b) * 0, pm)
-    self.assertEqual(result.op, Ops.CONST)
-    self.assertEqual(result.arg, 0)
-
-  def test_rebuild_symmetric(self):
-    pm = PatternMatcher([
-      (UPat.var("x") * 0, lambda x: x.const_like(0)),
-      (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name="a"),
-       lambda a: a.const_like(a.src[0].arg + a.src[1].arg) if a.op is Ops.ADD else None),
-    ])
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    result = egraph_extract((a * 0) + (b * 0), pm)
-    self.assertEqual(result.op, Ops.CONST)
-    self.assertEqual(result.arg, 0)
-
-  def test_rebuild_with_real_rules(self):
-    from tinygrad.uop.symbolic import symbolic_simple
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    self.assertIs(egraph_extract((a * 0) + (b * 1), symbolic_simple), b)
-
-  def test_rebuild_deep_chain(self):
-    pm = PatternMatcher([
-      (UPat.var("x") * 0, lambda x: x.const_like(0)),
-      (UPat.var("x") + 0, lambda x: x),
-      (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name="a"),
-       lambda a: a.const_like(a.src[0].arg + a.src[1].arg) if a.op is Ops.ADD else None),
-    ])
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    c = UOp.variable("c", 0, 10)
-    self.assertIs(egraph_extract(((a * 0) + (b * 0)) + c, pm), c)
-
-# *** test cost model ***
-
-class TestCostModel(unittest.TestCase):
-  def test_const_is_free(self):
-    self.assertEqual(node_cost(UOp.const(dtypes.int, 0)), 0)
-
-  def test_add_is_cheap(self):
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    self.assertEqual(node_cost(a + b), 1)
-
-  def test_div_is_expensive(self):
-    a = UOp.variable("a", 0, 10).cast(dtypes.index)
-    b = UOp.variable("b", 1, 10).cast(dtypes.index)
-    self.assertEqual(node_cost(a // b), 5)
-
-  def test_mul_more_than_add(self):
-    a = UOp.variable("a", 0, 10)
-    b = UOp.variable("b", 0, 10)
-    self.assertGreater(node_cost(a * b), node_cost(a + b))
-
-# *** test e-graph matches greedy rewrite ***
-
-class TestEGraphVsGreedy(unittest.TestCase):
-  def test_matches_greedy_identity(self):
-    from tinygrad.uop.ops import graph_rewrite
-    from tinygrad.uop.symbolic import symbolic_simple
-    a = UOp.variable("a", 0, 10)
-    greedy = graph_rewrite(a + 0, symbolic_simple)
-    egraph = egraph_extract(a + 0, symbolic_simple)
-    self.assertIs(greedy, egraph)
-
-  def test_matches_greedy_const_fold(self):
-    from tinygrad.uop.ops import graph_rewrite
-    from tinygrad.uop.symbolic import symbolic_simple
-    expr = UOp.const(dtypes.int, 10) + UOp.const(dtypes.int, 20)
-    greedy = graph_rewrite(expr, symbolic_simple)
-    egraph = egraph_extract(expr, symbolic_simple)
-    self.assertEqual(greedy.op, Ops.CONST)
-    self.assertEqual(egraph.op, Ops.CONST)
-    self.assertEqual(greedy.arg, egraph.arg)
-
-  def test_matches_greedy_double_identity(self):
-    from tinygrad.uop.ops import graph_rewrite
-    from tinygrad.uop.symbolic import symbolic_simple
-    a = UOp.variable("a", 0, 10)
-    expr = (a + 0) * 1
-    self.assertIs(graph_rewrite(expr, symbolic_simple), a)
-    self.assertIs(egraph_extract(expr, symbolic_simple), a)
-
-# *** test e-graph beats greedy (phase-ordering problems) ***
-
-# helper PMs that create phase-ordering traps
-_pm_strength_reduce = PatternMatcher([
-  # strength reduction x*2 -> x+x fires FIRST and destroys the x*c form needed by combine-terms
-  (UPat.var('x') * UPat.cvar('c', vec=False), lambda x,c: x+x if c.arg == 2 else None),
-  # combine terms: x*c0 + x*c1 -> x*(c0+c1) can only match if both sides are x*c
-  (UPat.var('x') * UPat.cvar('c0') + UPat.var('x') * UPat.cvar('c1'), lambda x,c0,c1: x*(c0+c1)),
-  # constant folding
-  (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name='a'),
-   lambda a: a.const_like(a.src[0].arg + a.src[1].arg) if a.op is Ops.ADD else
-             a.const_like(a.src[0].arg * a.src[1].arg) if a.op is Ops.MUL else None),
-  (UPat.var('x') + 0, lambda x: x),
-  (UPat.var('x') * 1, lambda x: x),
-])
-
-_pm_shift_reduce = PatternMatcher([
-  # strength reduction x*2 -> x<<1 fires FIRST and destroys the x*c form
-  (UPat.var('x') * UPat.cvar('c', vec=False),
-   lambda x,c: UOp(Ops.SHL, x.dtype, (x, x.const_like(1))) if c.arg == 2 else None),
-  (UPat.var('x') * UPat.cvar('c0') + UPat.var('x') * UPat.cvar('c1'), lambda x,c0,c1: x*(c0+c1)),
-  (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name='a'),
-   lambda a: a.const_like(a.src[0].arg + a.src[1].arg) if a.op is Ops.ADD else
-             a.const_like(a.src[0].arg * a.src[1].arg) if a.op is Ops.MUL else None),
-  (UPat.var('x') + 0, lambda x: x),
-  (UPat.var('x') * 1, lambda x: x),
-])
-
-_pm_strength_fold = PatternMatcher([
-  # strength reduction x*2 -> x+x blocks two-stage folding (x*c1)*c2 -> x*(c1*c2)
-  (UPat.var('x') * UPat.cvar('c', vec=False), lambda x,c: x+x if c.arg == 2 else None),
-  ((UPat.var('x') * UPat.cvar('c1')) * UPat.cvar('c2'), lambda x,c1,c2: x*(c1*c2)),
-  (UPat(GroupOp.Binary, src=(UPat((Ops.CONST, Ops.VCONST)),)*2, name='a'),
-   lambda a: a.const_like(a.src[0].arg * a.src[1].arg) if a.op is Ops.MUL else None),
-])
-
-def _total_cost(u:UOp) -> int:
-  return sum(node_cost(n) for n in u.toposort())
-
-class TestEGraphBeatsGreedy(unittest.TestCase):
-  """Tests where the e-graph finds a cheaper result than the greedy rewriter due to phase-ordering.
-
-  The core problem: when Rule A fires first and transforms a node, it can destroy the pattern
-  that Rule B needs to match. Rule B would have led to a cheaper result, but the greedy rewriter
-  never tries it. The e-graph explores BOTH paths and picks the cheapest.
-  """
-  def test_strength_reduce_blocks_combine(self):
-    """a*2 + a*3: strength reduction x*2->x+x destroys the x*c form needed by combine-terms x*c0+x*c1->x*(c0+c1)."""
-    a = UOp.variable("a", 0, 10)
-    expr = a * 2 + a * 3
-    greedy = graph_rewrite(expr, _pm_strength_reduce)
-    egraph = egraph_extract(expr, _pm_strength_reduce)
-    # greedy: (a+a) + a*3 (cost 4) — strength reduction destroyed the a*2 pattern
-    self.assertEqual(greedy.op, Ops.ADD)
-    self.assertGreater(_total_cost(greedy), _total_cost(egraph))
-    # egraph: a*5 (cost 2) — combine-terms wins because the e-graph explored both paths
-    self.assertEqual(egraph.op, Ops.MUL)
-    self.assertEqual(egraph.src[1].arg, 5)
-
-  def test_shift_reduce_blocks_combine(self):
-    """a*2 + a*3: shift reduction x*2->x<<1 also destroys the combine-terms pattern."""
-    a = UOp.variable("a", 0, 10)
-    expr = a * 2 + a * 3
-    greedy = graph_rewrite(expr, _pm_shift_reduce)
-    egraph = egraph_extract(expr, _pm_shift_reduce)
-    self.assertEqual(greedy.op, Ops.ADD)
-    self.assertGreater(_total_cost(greedy), _total_cost(egraph))
-    self.assertEqual(egraph.op, Ops.MUL)
-    self.assertEqual(egraph.src[1].arg, 5)
-
-  def test_strength_reduce_chain(self):
-    """a*2 + a*3 + a*4: strength reduction causes greedy to miss the combined a*9."""
-    a = UOp.variable("a", 0, 10)
-    expr = a * 2 + a * 3 + a * 4
-    greedy = graph_rewrite(expr, _pm_strength_reduce)
-    egraph = egraph_extract(expr, _pm_strength_reduce)
-    self.assertGreater(_total_cost(greedy), _total_cost(egraph))
-
-  def test_strength_reduce_blocks_two_stage_fold(self):
-    """(a*2)*3: strength reduction x*2->x+x blocks two-stage constant folding (x*c1)*c2->x*(c1*c2)."""
-    a = UOp.variable("a", 0, 10)
-    expr = (a * 2) * 3
-    greedy = graph_rewrite(expr, _pm_strength_fold)
-    egraph = egraph_extract(expr, _pm_strength_fold)
-    # greedy: (a+a)*3 (cost 3) — can't fold constants because *2 was rewritten to +
-    self.assertGreater(_total_cost(greedy), _total_cost(egraph))
-    # egraph: a*6 (cost 2) — two-stage folding path was explored
-    self.assertEqual(egraph.op, Ops.MUL)
-    self.assertEqual(egraph.src[1].arg, 6)
-
-  def test_both_sides_strength_reduced(self):
-    """a*2 + a*2: both sides get strength-reduced, blocking combine-terms."""
-    a = UOp.variable("a", 0, 10)
-    expr = a * 2 + a * 2
-    greedy = graph_rewrite(expr, _pm_strength_reduce)
-    egraph = egraph_extract(expr, _pm_strength_reduce)
-    # greedy: (a+a)+(a+a) — both a*2 were rewritten before combine could fire
-    # egraph: a*4 — combine-terms path was found
-    self.assertEqual(egraph.op, Ops.MUL)
-    self.assertEqual(egraph.src[1].arg, 4)
-    # both have cost 2 here (shared subexpression), but egraph result is canonical
-    self.assertLessEqual(_total_cost(egraph), _total_cost(greedy))
-
-# *** test cycle-breaking in extraction ***
-
-class TestExtractionCycles(unittest.TestCase):
-  def test_self_referencing_eclass(self):
-    """x+0 -> x merges x+0 into x's eclass. Extraction must not recurse on the self-reference."""
-    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
-    a = UOp.variable("a", 0, 10)
-    self.assertIs(egraph_extract(a + 0, pm), a)
-
-  def test_nested_self_referencing_eclass(self):
-    """((a+0)+0)+0 — all merge into a's eclass. Deep self-reference chain."""
-    pm = PatternMatcher([(UPat.var("x") + 0, lambda x: x)])
-    a = UOp.variable("a", 0, 10)
-    self.assertIs(egraph_extract(((a + 0) + 0) + 0, pm), a)
-
-  def test_mutual_eclass_cycle(self):
-    """Two eclasses whose best nodes reference each other — extraction must terminate via cycle-breaking cache."""
-    x = UOp.variable("x", 0, 10)
-    y = UOp.variable("y", 0, 10)
-    one = UOp.const(dtypes.index, 1)
-    two = UOp.const(dtypes.index, 2)
-    node1 = x + one  # E1's best, child x is in E2
-    node2 = y + two  # E2's best, child y is in E1
-    eclass_of = {node1: node1, x: node2, node2: node2, y: node1, one: one, two: two}
-    cost_of = {node1: (2, node1), node2: (2, node2), one: (0, one), two: (0, two)}
-    # without cycle-breaking cache, this would recurse: E1->E2->E1->...
-    result = _rebuild_tree(node1, eclass_of, cost_of)
-    self.assertIsNotNone(result)  # just verify it terminates
-
-  def test_mutual_rewrite_cycle(self):
-    """x+x <-> x*2 mutual rewrite. Both forms in same eclass, extraction picks cheaper (ADD)."""
-    pm = PatternMatcher([
-      (UPat.var("x") + UPat.var("x"), lambda x: x * 2),
-      (UPat.var("x") * UPat.cvar("c", vec=False), lambda x,c: x+x if c.arg == 2 else None),
-    ])
-    a = UOp.variable("a", 0, 10)
-    result = egraph_extract(a + a, pm)
-    self.assertEqual(result.op, Ops.ADD)  # ADD cost 1 < MUL cost 2
-
-if __name__ == '__main__':
-  unittest.main(verbosity=2)

test/null/test_linearizer_failures.py

@@ -1,6 +1,6 @@
 # ruff: noqa: E501
 import unittest
-from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
+from tinygrad.uop.ops import UOp, Ops, AxisType
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import get_program
 from tinygrad.device import Device
@@ -18,7 +18,7 @@ class TestLinearizerFailures(unittest.TestCase):
     c8 = c7.index(c3)
     c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
     c10 = c0.index(c3).store(c9).end(c1, c2)
-    ast = c10.sink(arg=KernelInfo())
+    ast = c10.sink()
     get_program(ast, renderer=Device[Device.DEFAULT].renderer)
 
 if __name__ == '__main__':

test/null/test_linearizer_rewrite.py

@@ -35,7 +35,7 @@ class TestLinearizerRewrite(unittest.TestCase):
     prg = get_program(ast, Device["CPU"].renderer)
     assert prg.applied_opts == (), f"expected no opts, got {prg}"
 
-    prg = get_program(ast.replace(arg=KernelInfo()), Device["CPU"].renderer)
+    prg = get_program(ast.replace(arg=None), Device["CPU"].renderer)
     assert prg.applied_opts != (), f"expected opts to apply, got {prg.applied_opts}"
 
     prg = get_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer)

test/null/test_process_replay.py

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

test/null/test_tinyfs.py

@@ -5,22 +5,18 @@ class TestLoadStore(unittest.TestCase):
   def test_load_shape(self):
     t = Tensor(bytes(16)).fs_load(1024)
     assert t.shape == (1024,), t.shape
-    t.schedule()
 
   def test_store_shape(self):
     t = Tensor.zeros(1024).fs_store()
     assert t.shape == (16,), t.shape
-    t.schedule()
 
   def test_load_large_shape(self):
     t = Tensor(bytes(16)).fs_load(10_000_000)
     assert t.shape == (10_000_000,), t.shape
-    t.schedule()
 
   def test_store_large_shape(self):
     t = Tensor.zeros(10_000_000).fs_store()
     assert t.shape == (16,), t.shape
-    t.schedule()
 
 if __name__ == "__main__":
   unittest.main()

test/null/test_viz.py

@@ -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)),
-            ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000))]
+    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))]
 
     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: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))]
+    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))]
 
     j = load_profile(prof)
 
-    event = j['layout']['NV:SDMA:0']['events'][0]
+    event = j['layout']['NV']['events'][0]
     self.assertEqual(event['name'], 'COPYxx')
     self.assertEqual(event['st'], 0)   # first event
     self.assertEqual(event['dur'], 10)
 
-    event2 = j['layout']['NV:2:SDMA:0']['events'][0]
+    event2 = j['layout']['NV:2']['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', 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)],
+    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)],
                               deps=[[], [0]],
                               sigs=[decimal.Decimal(1000), decimal.Decimal(1002), decimal.Decimal(1004), decimal.Decimal(1008)])]
 
@@ -409,20 +409,22 @@ class TestVizProfiler(BaseTestViz):
     tracks = list(j['layout'])
     self.assertEqual(tracks[0], 'NV')
     self.assertEqual(tracks[1], 'NV Graph')
-    self.assertEqual(tracks[2], 'NV:1:SDMA:0')
+    self.assertEqual(tracks[2], 'NV:1')
 
     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'])
 
-    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)
+    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'])
 
     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'], sdma_events[0]['st']+sdma_events[0]['dur'])
+    self.assertEqual(graph_events[0]['st']+graph_events[0]['dur'], nv1_events[0]['st']+nv1_events[0]['dur'])
 
   def test_bytes_per_kernel(self):
     step = 10
@@ -463,11 +465,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", "TEST:1:ENGINE:0", "TEST:1"]:
+    for dname in ["TINY", "USER", "TEST:1 N1", "TEST:2 N1", "TEST:1 N2"]:
       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", "TEST:1:ENGINE:0", "TEST:1 N1","TEST:1 N2", "TEST:2 N1"])
+    self.assertListEqual(layout[2:], ["TEST:1 N1","TEST:1 N2", "TEST:2 N1"])
 
 def _alloc(b:int):
   a = Tensor.empty(b, device="NULL", dtype=dtypes.char)

test/test_custom_kernel.py

@@ -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, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.size}", opts_to_apply=()))
+  return store.end(i).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[1:]
+  out, a, b = kernel.src
   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[1:]
+  out, a, b = kernel.src
   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)
@@ -104,7 +104,7 @@ def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
 class TestCustomKernel(unittest.TestCase):
   def test_empty(self):
     a = Tensor.empty(1)
-    a = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink(arg=KernelInfo()))[0]
+    a = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink())[0]
     a.realize()
 
   def test_simple(self):
@@ -128,14 +128,6 @@ 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()

test/test_dtype.py

@@ -380,40 +380,9 @@ 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)

test/test_dtype_alu.py

@@ -1,6 +1,6 @@
 import unittest, operator, math
 from tinygrad import Context, Tensor, dtypes, Device
-from tinygrad.dtype import DType, truncate, fp8_to_float
+from tinygrad.dtype import DType, truncate
 from tinygrad.helpers import CI, EMULATED_DTYPES, getenv
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -59,10 +59,9 @@ 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)
-  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())
+  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.floats:
     if not is_dtype_supported(dtype) or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
       fe, fm = dtypes.finfo(dtype)
@@ -77,14 +76,13 @@ 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(v): return
-  else: tensor_value, numpy_value = op[0](ta).numpy(), op[1](ta.numpy())
+    if math.isinf(numpy_value.item()): return
+    numpy_value = truncate[dtype](numpy_value.item())
   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))
@@ -130,31 +128,16 @@ 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)
 
@@ -170,34 +153,18 @@ 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)
 

test/test_profiler.py

@@ -6,9 +6,6 @@ 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):
@@ -58,7 +55,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 _dev_base(kernel_runs[0].device) == kernel_runs[0].device, "kernel should not be on a sub-device"
+    assert not kernel_runs[0].is_copy, "kernel should not be copy"
 
   def test_profile_copyin(self):
     buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
@@ -66,8 +63,10 @@ class TestProfiler(unittest.TestCase):
     with helper_collect_profile(TestProfiler.d0) as profile:
       buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
 
-    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
+    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
+    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].is_copy, "kernel should be copy"
 
   def test_profile_multiops(self):
     runner_name = TestProfiler.runner._prg.name
@@ -78,12 +77,16 @@ class TestProfiler(unittest.TestCase):
       TestProfiler.runner([buf1, TestProfiler.a.uop.buffer], var_vals={})
       buf1.copyout(memoryview(bytearray(buf1.nbytes)))
 
-    evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
+    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
+    evs = [x for x in profile if isinstance(x, ProfileRangeEvent)]
 
     assert len(evs) == 3, "3 kernel runs are expected"
     # NOTE: order of events does not matter, the tool is responsible for sorting them
-    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"
+    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"
 
     #for i in range(1, 3):
     #  assert evs[i].st > evs[i-1].en, "timestamp not aranged"
@@ -99,9 +102,13 @@ class TestProfiler(unittest.TestCase):
       buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
       buf2.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
 
-    for dev in [TestProfiler.d0.device, d1.device]:
-      evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and _dev_base(x.device) == dev]
+    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)]
       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"]
@@ -111,8 +118,10 @@ class TestProfiler(unittest.TestCase):
     with helper_collect_profile(TestProfiler.d0, d1) as profile:
       buf1.to(f"{Device.DEFAULT}:1").realize()
 
-    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
+    profile0, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
+    kernel_runs = [x for x in profile0 if isinstance(x, ProfileRangeEvent)]
     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):
@@ -158,19 +167,17 @@ 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 = dev_evs[d1.device].tdiff - dev_evs[d2.device].tdiff
+      cpu_diff = d1.gpu2cpu_compute_time_diff - d2.gpu2cpu_compute_time_diff
       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"
 
-  @unittest.skip("this test is flaky")
+    print("pairwise clock jitter matrix (us):\n" + '\n'.join([''.join([f'{float(item):8.3f}' for item in row]) for row in jitter_matrix]))
+
   def test_cpu_profile(self):
     def test_fxn(err=False):
       time.sleep(0.1)
@@ -212,7 +219,6 @@ 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
@@ -221,7 +227,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 _dev_base(e.device) == e.device]
+    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent) and not e.is_copy]
     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)

test/test_renderer_failures.py

@@ -9,7 +9,7 @@ from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.runtime.ops_python import PythonRenderer
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu
+from tinygrad.uop.ops import UOp, Ops, python_alu
 from tinygrad.tensor import Tensor, _to_np_dtype
 
 def _test_uop_result(inputs:list[Tensor], prg, local_size=None):
@@ -32,7 +32,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
   ld = b.index(idx)
   alu = ld.alu(alu_op, *alu_src_uops)
   store = UOp.store(a.index(idx), alu)
-  sink = UOp(Ops.SINK, dtypes.void, (store,), arg=KernelInfo())
+  sink = UOp(Ops.SINK, dtypes.void, (store,))
   prg = get_program(sink, Device[Device.DEFAULT].renderer)
   return _test_uop_result([Tensor([input_val])], prg)[0]
 
@@ -42,7 +42,7 @@ class TestRendererFailures(unittest.TestCase):
     a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
     gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
-    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
+    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
     prg = get_program(sink, Device[Device.DEFAULT].renderer)
     ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])
@@ -53,7 +53,7 @@ class TestRendererFailures(unittest.TestCase):
     gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
     gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
-    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
+    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
     prg = get_program(sink, Device[Device.DEFAULT].renderer)
     ret = _test_uop_result([], prg, local_size=[4, 2, 1])[0]
     np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
@@ -99,7 +99,7 @@ class TestPTXFailures(unittest.TestCase):
     val = UOp.const(dtypes.int, 1)
     if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))
     gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, if_uop), val))
-    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
+    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,))
     prg = get_program(sink, Device[Device.DEFAULT].renderer)
     ret = _test_uop_result([], prg, local_size=[4, 1, 1])[0]
     np.testing.assert_equal(ret, [0, 1, 1, 1])

test/test_schedule.py

@@ -12,6 +12,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType, ImageDType
 from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat
 from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
+from tinygrad.schedule.rangeify import Kernel
 from tinygrad.engine.realize import CompiledRunner, run_schedule
 
 class KernelCountException(Exception): pass
@@ -677,6 +678,17 @@ class TestSchedule(unittest.TestCase):
     c = (a.sum(2).contiguous() + b).contiguous()
     check_schedule(c, 2)
 
+  # TODO: this requires supporting multiple stores in the AST
+  @unittest.expectedFailure
+  def test_multioutput_ast(self):
+    a = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().uop
+    b = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().uop
+    c = Tensor.arange(4).realize().uop
+    kernel = UOp(Ops.KERNEL, src=(a.base, b.base, c.base), arg=Kernel(UOp.sink(c.r(Ops.ADD, (0,))+1, c.r(Ops.ADD, (0,))*2)))
+    run_schedule(check_schedule(UOp.sink(a.assign(kernel), b.assign(kernel)), 1))
+    self.assertEqual(a.buffer.numpy(), [7])
+    self.assertEqual(b.buffer.numpy(), [12])
+
   @unittest.skip("no longer supported")
   def test_double_from(self):
     x = Tensor([1,2,3,4])

test/test_uops.py

@@ -23,7 +23,7 @@ def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
   return ret
 
 def _uops_to_prg(uops_list):
-  prg = get_program(UOp.sink(*uops_list, arg=KernelInfo()), Device[Device.DEFAULT].renderer)
+  prg = get_program(UOp.sink(*uops_list), Device[Device.DEFAULT].renderer)
   return CompiledRunner(replace(prg, device=Device.DEFAULT))
 
 def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
@@ -281,7 +281,7 @@ class TestUOpPrograms(unittest.TestCase):
     ptr = UOp.placeholder(out.shape, out.dtype, slot=0)
     i, j = UOp.range(10, axis_id=0), UOp.range(10, axis_id=1)
     prog = ptr[i,j].set(42).end(i,j)
-    self._run(prog.sink(arg=KernelInfo()), out)
+    self._run(prog.sink(), out)
 
     with Context(DEBUG=0): self.assertTrue((out == 42).all().item())
 

test/testextra/test_asm_gemm.py

@@ -1,13 +1,8 @@
 import unittest
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, CI
 from extra.gemm.asm.cdna.gemm import asm_gemm
-from test.helpers import needs_second_gpu
-
-# On non CDNA4 it will only validate the Tensor.custom_kernel integration
-# Use NULL=1 EMULATE=AMD_CDNA4 to also test the assembly
-def is_cdna4(): return getattr(Device[Device.DEFAULT].renderer, "arch", "").startswith("gfx950")
 
 def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
   Tensor.manual_seed(0)
@@ -20,55 +15,41 @@ def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:i
 
   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)
-  with Context(ASM_GEMM=1):
-    tst = asm_gemm(a, b)
-    tst.sum().backward()
+  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 = asm_gemm(a_ref, b_ref)
-    ref.sum().backward()
+  with Context(ASM_GEMM=0): ref = 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
   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"
 
-# 128x smaller than usual
-# uses the UOp GEMM, runs on non CDNA4 and CI
+SCALE = 128 if CI else 1
+
 @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
 class TestGemm(unittest.TestCase):
-  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, 64, 32, 32, gpus=2)
+  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 test_gemm_multi(self): verify_asm_gemm(2, 8192//SCALE, 4096//SCALE, 4096//SCALE, gpus=2)
 
-# uses the Asm GEMM on CDNA4 only for speed reasons
 class TestGemmLarge(unittest.TestCase):
   def setUp(self):
-    if not is_cdna4():
+    if getattr(Device[Device.DEFAULT].renderer, "arch", "") != "gfx950":
       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)
   def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
   def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
   def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
   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_gemm_unsupported(self):
     with self.assertRaisesRegex(AssertionError, "shape not supported"):

test/testextra/test_hevc.py

@@ -1,9 +1,8 @@
 import unittest
 
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.helpers import fetch, round_up
+from tinygrad import Device
+from tinygrad.helpers import fetch
 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):
@@ -62,26 +61,5 @@ 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()

test/unit/test_assign.py

@@ -259,15 +259,15 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_allclose(out, [1.,1.,1.,1.,1.,1.,0.,0.,1.,1.,1.,1.,1.,1.,0.,0.])
 
   def test_assign_contiguous(self):
-    b = Tensor.arange(16).reshape(4,4).contiguous().realize()
-    a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1)
+    b = Tensor.rand(4,4).realize()
+    a = (Tensor.rand(4,4).realize() + 1)
     kc = GlobalCounters.kernel_count
     b.assign(a.contiguous()).realize()
     assert GlobalCounters.kernel_count - kc == 2
 
   def test_assign_contiguous_permute(self):
-    b = Tensor.arange(16).reshape(4,4).contiguous().realize()
-    a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1).permute((1,0))
+    b = Tensor.rand(4,4).realize()
+    a = (Tensor.rand(4,4).realize() + 1).permute((1,0))
     kc = GlobalCounters.kernel_count
     b.assign(a.contiguous()).realize()
     assert GlobalCounters.kernel_count - kc == 2
@@ -333,7 +333,7 @@ class TestAssign(unittest.TestCase):
 
   @unittest.skip("multi output not supported anymore")
   def test_simple_assignment_multioutput(self):
-    a = Tensor.arange(32*32).reshape(32, 32).contiguous().realize()
+    a = Tensor.randn(32, 32).realize()
     b = Tensor.full((32, ), 1.).contiguous().realize()
     c = Tensor.full((32, ), 2.).contiguous().realize()
     d = Tensor.full((32, ), 3.).contiguous().realize()
@@ -361,16 +361,16 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_equal(a.numpy(), np.arange(4 * 4).reshape(4, 4).transpose(1, 0) + np.arange(4 * 4).reshape(4, 4))
 
   def test_permuted_reduceop_child_dual_use(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
-    b = Tensor.ones(32, 32, dtype=dtypes.int).contiguous().realize()
+    a = Tensor.randn(32, 32, 32).realize()
+    b = Tensor.full((32, 32), 1.).contiguous().realize()
     r = a.sum(axis=1)
     b.assign(r + b.permute(1, 0))
     b.realize()
-    np.testing.assert_equal(b.numpy(), a.numpy().sum(axis=1)+np.ones((32, 32), dtype=np.int32).transpose(1, 0))
+    np.testing.assert_allclose(b.numpy(), a.numpy().sum(axis=1)+np.ones((32, 32)).transpose(1, 0), atol=1e-6, rtol=1e-3)
 
   @unittest.skip("multi output not supported anymore")
   def test_permuted_reduceop_multioutput_dual_use(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
+    a = Tensor.randn(32, 32, 32).realize()
     b = Tensor.full((32, 32), 1.).contiguous().realize()
     c = Tensor.full((32, 32), 2.).contiguous().realize()
 
@@ -383,7 +383,7 @@ class TestAssign(unittest.TestCase):
 
   @unittest.skip("multi output not supported anymore")
   def test_permuted_reduceop_multioutput_dual_use_possible(self):
-    a = Tensor.arange(32*32*32).reshape(32, 32, 32).contiguous().realize()
+    a = Tensor.randn(32, 32, 32, dtype=dtypes.int).realize()
     b = Tensor.arange(32 * 32).reshape(32, 32).realize()
     c = Tensor.arange(32 * 32).reshape(32, 32).realize()
 
@@ -529,7 +529,6 @@ class TestAssign(unittest.TestCase):
     a = Tensor.empty(5, device=f"disk:{temp('disk_assignment')}").assign(Tensor.ones(5)).numpy()
     np.testing.assert_equal(a, np.ones(5))
 
-  @unittest.skip("this test is crashing!")
   def test_assign_slice_then_read(self):
     """Assign to slice then read from buffer - read should see the assigned values.
     This is the KV cache pattern from llm.py.

test/unit/test_disk_tensor.py

@@ -4,17 +4,9 @@ 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, dedup
+from tinygrad.helpers import Timing, fetch, temp, OSX
 from test.helpers import slow
 
-class TempDirTestCase(unittest.TestCase):
-  def setUp(self):
-    self.temp_dir = tempfile.TemporaryDirectory()
-  def tearDown(self):
-    self.temp_dir.cleanup()
-  def tmp(self, name:str) -> str:
-    return (pathlib.Path(self.temp_dir.name) / name).as_posix()
-
 def compare_weights_both(url):
   import torch
   fn = fetch(url)
@@ -92,7 +84,7 @@ class TestRawDiskBuffer(unittest.TestCase):
     pathlib.Path(tmp).unlink()
 
 @unittest.skipUnless(is_dtype_supported(dtypes.uint8), "need uint8")
-class TestSafetensors(TempDirTestCase):
+class TestSafetensors(unittest.TestCase):
   def test_real_safetensors(self):
     import torch
     from safetensors.torch import save_file
@@ -103,19 +95,19 @@ class TestSafetensors(TempDirTestCase):
       "weight3": torch.arange(0, 17, dtype=torch.int32).reshape(17,1,1),
       "weight4": torch.arange(0, 2, dtype=torch.uint8),
     }
-    save_file(tensors, self.tmp("real.safetensors"))
+    save_file(tensors, temp("real.safetensors"))
 
-    ret = safe_load(self.tmp("real.safetensors"))
+    ret = safe_load(temp("real.safetensors"))
     for k,v in tensors.items(): np.testing.assert_array_equal(ret[k].numpy(), v.numpy())
-    safe_save(ret, self.tmp("real.safetensors_alt"))
-    with open(self.tmp("real.safetensors"), "rb") as f:
-      with open(self.tmp("real.safetensors_alt"), "rb") as g:
+    safe_save(ret, temp("real.safetensors_alt"))
+    with open(temp("real.safetensors"), "rb") as f:
+      with open(temp("real.safetensors_alt"), "rb") as g:
         assert f.read() == g.read()
-    ret2 = safe_load(self.tmp("real.safetensors_alt"))
+    ret2 = safe_load(temp("real.safetensors_alt"))
     for k,v in tensors.items(): np.testing.assert_array_equal(ret2[k].numpy(), v.numpy())
 
   def test_real_safetensors_open(self):
-    fn = self.tmp("real_safe")
+    fn = temp("real_safe")
     state_dict = {"tmp": Tensor.rand(10,10)}
     safe_save(state_dict, fn)
     import os
@@ -131,15 +123,15 @@ class TestSafetensors(TempDirTestCase):
     from extra.models.efficientnet import EfficientNet
     model = EfficientNet(0)
     state_dict = get_state_dict(model)
-    safe_save(state_dict, self.tmp("eff0"))
-    state_dict_loaded = safe_load(self.tmp("eff0"))
+    safe_save(state_dict, temp("eff0"))
+    state_dict_loaded = safe_load(temp("eff0"))
     assert sorted(state_dict_loaded.keys()) == sorted(state_dict.keys())
     for k,v in state_dict.items():
       np.testing.assert_array_equal(v.numpy(), state_dict_loaded[k].numpy())
 
     # load with the real safetensors
     from safetensors import safe_open
-    with safe_open(self.tmp("eff0"), framework="pt", device="cpu") as f:
+    with safe_open(temp("eff0"), framework="pt", device="cpu") as f:
       assert sorted(f.keys()) == sorted(state_dict.keys())
       for k in f.keys():
         np.testing.assert_array_equal(f.get_tensor(k).numpy(), state_dict[k].numpy())
@@ -163,19 +155,19 @@ class TestSafetensors(TempDirTestCase):
 
   def test_metadata(self):
     metadata = {"hello": "world"}
-    safe_save({}, self.tmp('metadata.safetensors'), metadata)
+    safe_save({}, temp('metadata.safetensors'), metadata)
     import struct
-    with open(self.tmp('metadata.safetensors'), 'rb') as f:
+    with open(temp('metadata.safetensors'), 'rb') as f:
       dat = f.read()
     sz = struct.unpack(">Q", dat[0:8])[0]
     import json
     assert json.loads(dat[8:8+sz])['__metadata__']['hello'] == 'world'
 
   def test_save_all_dtypes(self):
-    for dtype in dedup(DTYPES_DICT.values()):
+    for dtype in 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")
+      path = temp(f"ones.{dtype}.safetensors")
       ones = Tensor(np.random.rand(10,10), dtype=dtype)
       safe_save(get_state_dict(ones), path)
       np.testing.assert_equal(ones.numpy(), list(safe_load(path).values())[0].numpy())
@@ -197,9 +189,9 @@ class TestSafetensors(TempDirTestCase):
       "weight_I16": torch.tensor([127, 64], dtype=torch.short),
       "weight_BF16": torch.randn((2, 2), dtype=torch.bfloat16),
     }
-    save_file(tensors, self.tmp("dtypes.safetensors"))
+    save_file(tensors, temp("dtypes.safetensors"))
 
-    loaded = safe_load(self.tmp("dtypes.safetensors"))
+    loaded = safe_load(temp("dtypes.safetensors"))
     for k,v in loaded.items():
       if v.dtype != dtypes.bfloat16:
         assert v.numpy().dtype == tensors[k].numpy().dtype
@@ -211,52 +203,57 @@ class TestSafetensors(TempDirTestCase):
       "weight_U32": np.array([1, 2, 3], dtype=np.uint32),
       "weight_U64": np.array([1, 2, 3], dtype=np.uint64),
     }
-    np_save_file(tensors, self.tmp("dtypes.safetensors"))
+    np_save_file(tensors, temp("dtypes.safetensors"))
 
-    loaded = safe_load(self.tmp("dtypes.safetensors"))
+    loaded = safe_load(temp("dtypes.safetensors"))
     for k,v in loaded.items():
       assert v.numpy().dtype == tensors[k].dtype
       np.testing.assert_allclose(v.numpy(), tensors[k])
 
-def helper_test_disk_tensor(tmp, fn, data, np_fxn, tinygrad_fxn=None):
+def helper_test_disk_tensor(fn, data, np_fxn, tinygrad_fxn=None):
   if tinygrad_fxn is None: tinygrad_fxn = np_fxn
-  pathlib.Path(tmp(fn)).unlink(missing_ok=True)
-  tinygrad_tensor = Tensor(data, device="CPU").to(f"disk:{tmp(fn)}")
+  pathlib.Path(temp(fn)).unlink(missing_ok=True)
+  tinygrad_tensor = Tensor(data, device="CPU").to(f"disk:{temp(fn)}")
   numpy_arr = np.array(data)
   tinygrad_fxn(tinygrad_tensor)
   np_fxn(numpy_arr)
   np.testing.assert_allclose(tinygrad_tensor.numpy(), numpy_arr)
 
-class TestDiskTensor(TempDirTestCase):
+class TestDiskTensor(unittest.TestCase):
   def test_empty(self):
-    Tensor.empty(100, 100, device=f"disk:{self.tmp('dt_empty')}")
+    pathlib.Path(temp("dt_empty")).unlink(missing_ok=True)
+    Tensor.empty(100, 100, device=f"disk:{temp('dt_empty')}")
 
   def test_simple_read(self):
-    fn = pathlib.Path(self.tmp("dt_simple_read"))
+    fn = pathlib.Path(temp("dt_simple_read"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256)))
-    t = Tensor.empty(16, 16, device=f"disk:{self.tmp('dt_simple_read')}", dtype=dtypes.uint8)
+    t = Tensor.empty(16, 16, device=f"disk:{temp('dt_simple_read')}", dtype=dtypes.uint8)
     out = t[1].to(Device.DEFAULT).tolist()
     assert out == list(range(16, 32))
 
   def test_simple_read_bitcast(self):
-    fn = pathlib.Path(self.tmp("dt_simple_read_bitcast"))
+    fn = pathlib.Path(temp("dt_simple_read_bitcast"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256))*2)
-    t = Tensor.empty(16, 16*2, device=f"disk:{self.tmp('dt_simple_read_bitcast')}", dtype=dtypes.uint8)
+    t = Tensor.empty(16, 16*2, device=f"disk:{temp('dt_simple_read_bitcast')}", dtype=dtypes.uint8)
     out = t[1].bitcast(dtypes.uint16).to(Device.DEFAULT).tolist()
     tout = [(x//256, x%256) for x in out]
     assert tout == list([(x+1,x) for x in range(32,64,2)])
 
   def test_simple_read_bitcast_alt(self):
-    fn = pathlib.Path(self.tmp("dt_simple_read_bitcast_alt"))
+    fn = pathlib.Path(temp("dt_simple_read_bitcast_alt"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256))*2)
-    t = Tensor.empty(16, 16*2, device=f"disk:{self.tmp('dt_simple_read_bitcast_alt')}", dtype=dtypes.uint8)
+    t = Tensor.empty(16, 16*2, device=f"disk:{temp('dt_simple_read_bitcast_alt')}", dtype=dtypes.uint8)
     out = t.bitcast(dtypes.uint16)[1].to(Device.DEFAULT).tolist()
     tout = [(x//256, x%256) for x in out]
     assert tout == list([(x+1,x) for x in range(32,64,2)])
 
   def test_strided_read(self):
     # test non-contiguous (strided) read - should read elements at indices 0, 2, 4
-    dt = Tensor([0, 1, 2, 3, 4, 5]).to(f"disk:{self.tmp('dt_strided_read')}")
+    pathlib.Path(temp(fn:="dt_strided_read")).unlink(missing_ok=True)
+    dt = Tensor([0, 1, 2, 3, 4, 5]).to(f"disk:{temp(fn)}")
     result = dt[::2].tolist()
     # TODO: dt[::2] selects indices 0, 2, 4, so result should be [0, 2, 4]
     # self.assertEqual(result, [0, 2, 4])
@@ -264,38 +261,43 @@ class TestDiskTensor(TempDirTestCase):
 
   def test_permuted_read(self):
     # test non-contiguous (permuted) read - should read transposed
-    dt = Tensor([[0, 1, 2], [3, 4, 5]]).to(f"disk:{self.tmp('dt_permuted_read')}")
+    pathlib.Path(temp(fn:="dt_permuted_read")).unlink(missing_ok=True)
+    dt = Tensor([[0, 1, 2], [3, 4, 5]]).to(f"disk:{temp(fn)}")
     result = dt.T.tolist()
     # TODO: transpose should give [[0, 3], [1, 4], [2, 5]]
     # self.assertEqual(result, [[0, 3], [1, 4], [2, 5]])
     self.assertEqual(result, [[0, 1], [2, 3], [4, 5]])  # wrong!
 
   def test_write_ones(self):
+    pathlib.Path(temp("dt_write_ones")).unlink(missing_ok=True)
+
     out = Tensor.ones(10, 10, device="CPU").contiguous()
-    outdisk = out.to(f"disk:{self.tmp('dt_write_ones')}")
+    outdisk = out.to(f"disk:{temp('dt_write_ones')}")
     print(outdisk)
     outdisk.realize()
     del out, outdisk
 
     import struct
     # test file
-    with open(self.tmp("dt_write_ones"), "rb") as f:
+    with open(temp("dt_write_ones"), "rb") as f:
       assert f.read() == struct.pack('<f', 1.0) * 100 == b"\x00\x00\x80\x3F" * 100
 
     # test load alt
-    reloaded = Tensor.empty(10, 10, device=f"disk:{self.tmp('dt_write_ones')}")
+    reloaded = Tensor.empty(10, 10, device=f"disk:{temp('dt_write_ones')}")
     np.testing.assert_almost_equal(reloaded.numpy(), np.ones((10, 10)))
 
   def test_simple_setitem(self):
+    pathlib.Path(temp(fn:="dt_simple_setitem")).unlink(missing_ok=True)
     data = [[1],[2]]
     src = Tensor(data)
-    dt = src.to(f"disk:{self.tmp('dt_simple_setitem')}")
+    dt = src.to(f"disk:{temp(fn)}")
     dt[1] = [3]
     self.assertEqual(dt.tolist(), [[1], [3]])
 
   def test_strided_setitem(self):
     # test non-contiguous (strided) setitem - should set elements at indices 0, 2, 4
-    dt = Tensor([1, 2, 3, 4, 5, 6]).to(f"disk:{self.tmp('dt_strided_setitem')}")
+    pathlib.Path(temp(fn:="dt_strided_setitem")).unlink(missing_ok=True)
+    dt = Tensor([1, 2, 3, 4, 5, 6]).to(f"disk:{temp(fn)}")
     dt[::2] = Tensor([10, 20, 30])
     # TODO: dt[::2] selects indices 0, 2, 4, so result should be [10, 2, 20, 4, 30, 6]
     # self.assertEqual(dt.tolist(), [10, 2, 20, 4, 30, 6])
@@ -303,35 +305,39 @@ class TestDiskTensor(TempDirTestCase):
 
   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)
+    pathlib.Path(temp(fn:="dt_assign_const")).unlink(missing_ok=True)
+    dt = Tensor.empty(4, device=f"disk:{temp(fn)}", dtype=dtypes.int32)
     dt.assign(Tensor.full((4,), 42, dtype=dtypes.int32)).realize()
     np.testing.assert_array_equal(dt.numpy(), [42, 42, 42, 42])
 
   def test_assign_slice_from_const(self):
     # slice assign from CONST to disk - tests size calculation when no RANGE ops
-    dt = Tensor([0, 1, 2, 3], dtype=dtypes.int32).to(f"disk:{self.tmp('dt_slice_const')}")
+    pathlib.Path(temp(fn:="dt_slice_const")).unlink(missing_ok=True)
+    dt = Tensor([0, 1, 2, 3], dtype=dtypes.int32).to(f"disk:{temp(fn)}")
     dt[1:3].assign(Tensor.full((2,), 99, dtype=dtypes.int32)).realize()
     np.testing.assert_array_equal(dt.numpy(), [0, 99, 99, 3])
 
   def test_disk_to_disk_copy(self):
     # disk-to-disk copy needs to go through CPU
-    src = Tensor([1, 2, 3, 4], dtype=dtypes.int32).to(f"disk:{self.tmp('dt_d2d_src')}")
-    dst = Tensor.empty(4, device=f"disk:{self.tmp('dt_d2d_dst')}", dtype=dtypes.int32)
+    pathlib.Path(temp(fn1:="dt_d2d_src")).unlink(missing_ok=True)
+    pathlib.Path(temp(fn2:="dt_d2d_dst")).unlink(missing_ok=True)
+    src = Tensor([1, 2, 3, 4], dtype=dtypes.int32).to(f"disk:{temp(fn1)}")
+    dst = Tensor.empty(4, device=f"disk:{temp(fn2)}", dtype=dtypes.int32)
     dst.assign(src.to("CPU")).realize()
     np.testing.assert_array_equal(dst.numpy(), [1, 2, 3, 4])
 
   def test_assign_slice(self):
     def assign(x,s,y): x[s] = y
-    helper_test_disk_tensor(self.tmp, "dt_assign_slice_1", [0,1,2,3], lambda x: assign(x, slice(0,2), [13, 12]))
-    helper_test_disk_tensor(self.tmp, "dt_assign_slice_2", [[0,1,2,3],[4,5,6,7]], lambda x: assign(x, slice(0,1), [[13, 12, 11, 10]]))
+    helper_test_disk_tensor("dt_assign_slice_1", [0,1,2,3], lambda x: assign(x, slice(0,2), [13, 12]))
+    helper_test_disk_tensor("dt_assign_slice_2", [[0,1,2,3],[4,5,6,7]], lambda x: assign(x, slice(0,1), [[13, 12, 11, 10]]))
 
   def test_reshape(self):
-    helper_test_disk_tensor(self.tmp, "dt_reshape_1", [1,2,3,4,5], lambda x: x.reshape((1,5)))
-    helper_test_disk_tensor(self.tmp, "dt_reshape_2", [1,2,3,4], lambda x: x.reshape((2,2)))
+    helper_test_disk_tensor("dt_reshape_1", [1,2,3,4,5], lambda x: x.reshape((1,5)))
+    helper_test_disk_tensor("dt_reshape_2", [1,2,3,4], lambda x: x.reshape((2,2)))
 
   def test_assign_to_different_dtype(self):
     # NOTE: this is similar to Y_train in fetch_cifar
-    t = Tensor.empty(10, device=f'disk:{self.tmp("dt_assign_to_different_dtype")}', dtype=dtypes.int64)
+    t = Tensor.empty(10, device=f'disk:{temp("dt_assign_to_different_dtype")}', dtype=dtypes.int64)
 
     for i in range(5):
       data = np.array([3, 3])
@@ -343,7 +349,8 @@ class TestDiskTensor(TempDirTestCase):
   def test_assign_with_bitcast(self):
     # bitcast assign is used in safe_save for writing header length
     # bitcast on source side works, bitcast on target side raises
-    t = Tensor.empty(16, device=f"disk:{self.tmp('dt_assign_bitcast')}", dtype=dtypes.uint8)
+    pathlib.Path(temp(fn:="dt_assign_bitcast")).unlink(missing_ok=True)
+    t = Tensor.empty(16, device=f"disk:{temp(fn)}", dtype=dtypes.uint8)
     # correct way: bitcast the source to match target dtype
     t[0:8].assign(Tensor([12345], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
     val = int.from_bytes(t[0:8].data(), 'little')
@@ -354,7 +361,8 @@ class TestDiskTensor(TempDirTestCase):
 
   def test_assign_to_bitcast_view(self):
     # assign float values to a float32 view of a uint8 disk buffer (used by safe_save)
-    t = Tensor.empty(32, device=f"disk:{self.tmp('dt_bitcast_view_assign')}", dtype=dtypes.uint8)
+    pathlib.Path(temp(fn:="dt_bitcast_view_assign")).unlink(missing_ok=True)
+    t = Tensor.empty(32, device=f"disk:{temp(fn)}", dtype=dtypes.uint8)
     # create float32 view of bytes 8-24 (4 floats)
     float_view = t[8:24].bitcast(dtypes.float32)
     float_view.assign(Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32, device="CPU"))
@@ -362,20 +370,21 @@ class TestDiskTensor(TempDirTestCase):
 
   def test_assign_cross_device(self):
     # disk assign allows cross-device (source on GPU/CPU, target on disk)
-    t = Tensor.empty(4, device=f"disk:{self.tmp('dt_assign_cross')}", dtype=dtypes.float32)
+    pathlib.Path(temp(fn:="dt_assign_cross")).unlink(missing_ok=True)
+    t = Tensor.empty(4, device=f"disk:{temp(fn)}", dtype=dtypes.float32)
     src = Tensor([1.0, 2.0, 3.0, 4.0])  # on default device
     t.assign(src)
     np.testing.assert_array_equal(t.numpy(), [1.0, 2.0, 3.0, 4.0])
 
   def test_bitcast(self):
-    with open(self.tmp('dt_bitcast'), "wb") as f: f.write(bytes(range(10,20)))
-    t = Tensor.empty(5, dtype=dtypes.int16, device=f"disk:{self.tmp('dt_bitcast')}")
+    with open(temp('dt_bitcast'), "wb") as f: f.write(bytes(range(10,20)))
+    t = Tensor.empty(5, dtype=dtypes.int16, device=f"disk:{temp('dt_bitcast')}")
     ret = t.to("CPU").bitcast(dtypes.uint16) + 1
     assert ret.tolist() == [2827, 3341, 3855, 4369, 4883]
 
   def test_bitcast_view(self):
-    with open(self.tmp('dt_bitcast_view'), "wb") as f: f.write(bytes(range(10, 24)))
-    t = Tensor.empty(3, dtype=dtypes.uint, device=f"disk:{self.tmp('dt_bitcast_view')}").shrink([(0, 2)])
+    with open(temp('dt_bitcast_view'), "wb") as f: f.write(bytes(range(10, 24)))
+    t = Tensor.empty(3, dtype=dtypes.uint, device=f"disk:{temp('dt_bitcast_view')}").shrink([(0, 2)])
     ret = t.bitcast(dtypes.uint16).to("CPU") + 1
     assert ret.tolist() == [2827, 3341, 3855, 4369]
 
@@ -383,55 +392,59 @@ class TestDiskTensor(TempDirTestCase):
   @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
   def test_bf16_disk_write_read(self):
     t = Tensor([10000, -1, -1000, -10000, 20], dtype=dtypes.float32)
-    t.to(f"disk:{self.tmp('dt_bf16_disk_write_read_f32')}").realize()
+    t.to(f"disk:{temp('dt_bf16_disk_write_read_f32')}").realize()
 
     # hack to "cast" f32 -> bf16
-    with open(self.tmp('dt_bf16_disk_write_read_f32'), "rb") as f: dat = f.read()
+    with open(temp('dt_bf16_disk_write_read_f32'), "rb") as f: dat = f.read()
     adat = b''.join([dat[i+2:i+4] for i in range(0, len(dat), 4)])
-    with open(self.tmp('dt_bf16_disk_write_read_bf16'), "wb") as f: f.write(adat)
+    with open(temp('dt_bf16_disk_write_read_bf16'), "wb") as f: f.write(adat)
 
-    t = Tensor.empty(5, dtype=dtypes.bfloat16, device=f"disk:{self.tmp('dt_bf16_disk_write_read_bf16')}")
+    t = Tensor.empty(5, dtype=dtypes.bfloat16, device=f"disk:{temp('dt_bf16_disk_write_read_bf16')}")
     ct = t.to(Device.DEFAULT).cast(dtypes.float)
     assert ct.numpy().tolist() == [9984., -1, -1000, -9984, 20]
 
   def test_copy_from_disk(self):
-    fn = pathlib.Path(self.tmp("dt_copy_from_disk"))
+    fn = pathlib.Path(temp("dt_copy_from_disk"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256))*1024)
 
-    t = Tensor.empty(256*1024, device=f"disk:{self.tmp('dt_copy_from_disk')}", dtype=dtypes.uint8)
+    t = Tensor.empty(256*1024, device=f"disk:{temp('dt_copy_from_disk')}", dtype=dtypes.uint8)
     on_dev = t.to(Device.DEFAULT).realize()
     np.testing.assert_equal(on_dev.numpy(), t.numpy())
 
   def test_copy_from_disk_offset(self):
-    fn = pathlib.Path(self.tmp("dt_copy_from_disk_offset"))
+    fn = pathlib.Path(temp("dt_copy_from_disk_offset"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256))*1024)
 
     for off in [314, 991, 2048, 4096]:
-      t = Tensor.empty(256*1024, device=f"disk:{self.tmp('dt_copy_from_disk_offset')}", dtype=dtypes.uint8)[off:]
+      t = Tensor.empty(256*1024, device=f"disk:{temp('dt_copy_from_disk_offset')}", dtype=dtypes.uint8)[off:]
       on_dev = t.to(Device.DEFAULT).realize()
       np.testing.assert_equal(on_dev.numpy(), t.numpy())
 
   @slow
   def test_copy_from_disk_huge(self):
-    fn = pathlib.Path(self.tmp("dt_copy_from_disk_huge"))
+
+    fn = pathlib.Path(temp("dt_copy_from_disk_huge"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256))*1024*256)
 
     for off in [0, 551]:
-      t = Tensor.empty(256*1024*256, device=f"disk:{self.tmp('dt_copy_from_disk_huge')}", dtype=dtypes.uint8)[off:]
+      t = Tensor.empty(256*1024*256, device=f"disk:{temp('dt_copy_from_disk_huge')}", dtype=dtypes.uint8)[off:]
       on_dev = t.to(Device.DEFAULT).realize()
       np.testing.assert_equal(on_dev.numpy(), t.numpy())
 
   @unittest.skip("this allocates a lot of RAM")
   @unittest.skipUnless(OSX, "seems to only be an issue on macOS with file size >2 GiB")
   def test_copy_to_cpu_not_truncated(self):
-    fn = self.tmp("dt_copy_to_cpu_not_truncated")
-    with open(fn, "wb") as f: f.write(b'\x01' * (size := int(2 * 1024**3)) + (test := b"test"))
+    with open((fn:=temp("dt_copy_to_cpu_not_truncated")), "wb") as f: f.write(b'\x01' * (size := int(2 * 1024**3)) + (test := b"test"))
     x = Tensor.empty(size + len(test), dtype=dtypes.uint8, device=f"disk:{fn}").to("CPU").realize()
     assert x[size:].data().tobytes() == test
 
   def test_disk_device_reuse(self):
     from tinygrad.runtime.ops_disk import DiskDevice
-    fn = pathlib.Path(self.tmp("dt_device_reuse"))
+    fn = pathlib.Path(temp("dt_device_reuse"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256)))
     # create first tensor and realize it
     t1 = Tensor.empty(128, device=f"disk:{fn}", dtype=dtypes.uint8)
@@ -453,7 +466,8 @@ class TestDiskTensor(TempDirTestCase):
 
   def test_disk_open_failure_state(self):
     from tinygrad.runtime.ops_disk import DiskDevice
-    fn = pathlib.Path(self.tmp("dt_open_failure"))
+    fn = pathlib.Path(temp("dt_open_failure"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256)))
     os.chmod(fn, 0o000)
     try:
@@ -472,7 +486,8 @@ class TestDiskTensor(TempDirTestCase):
     assert disk_device.size == 200
 
   def test_disk_permission_error(self):
-    fn = pathlib.Path(self.tmp("dt_permission"))
+    fn = pathlib.Path(temp("dt_permission"))
+    fn.unlink(missing_ok=True)
     fn.write_bytes(bytes(range(256)))
     os.chmod(fn, 0o000)
     try:
@@ -481,14 +496,17 @@ class TestDiskTensor(TempDirTestCase):
     finally:
       os.chmod(fn, 0o644)
 
-class TestPathTensor(TempDirTestCase):
+class TestPathTensor(unittest.TestCase):
   def setUp(self):
-    super().setUp()
+    self.temp_dir = tempfile.TemporaryDirectory()
     self.test_file = pathlib.Path(self.temp_dir.name) / "test_file.bin"
     self.test_data = np.arange(100, dtype=np.uint8).tobytes()
     with open(self.test_file, "wb") as f:
       f.write(self.test_data)
 
+  def tearDown(self):
+    self.temp_dir.cleanup()
+
   def test_path_tensor_no_device(self):
     t = Tensor(self.test_file)
     self.assertEqual(t.shape, (100,))
@@ -539,10 +557,10 @@ class TestPathTensor(TempDirTestCase):
     os.chmod(test_file, 0o644)
     assert Tensor(pathlib.Path(test_file)).tolist(), list(range(10))
 
-class TestDiskTensorMovement(TempDirTestCase):
+class TestDiskTensorMovement(unittest.TestCase):
   def setUp(self):
-    super().setUp()
-    self.fn = pathlib.Path(self.tmp("custom_disk_range"))
+    self.fn = pathlib.Path(temp("custom_disk_range"))
+    self.fn.unlink(missing_ok=True)
     Tensor.arange(100, dtype=dtypes.uint8).to(f"disk:{str(self.fn)}").realize()
 
   def test_simple_read(self):

test/unit/test_gguf.py

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

tinygrad/codegen/__init__.py

@@ -1,38 +1,35 @@
 from typing import cast
-from dataclasses import replace
 import itertools
-from tinygrad.helpers import DISABLE_FAST_IDIV, EMULATED_DTYPES, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, EGRAPH, TracingKey, Context
+from tinygrad.helpers import DISABLE_FAST_IDIV, EMULATED_DTYPES, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, getenv, 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, promo_lattice
-from tinygrad.device import is_dtype_supported
+from tinygrad.dtype import dtypes, PtrDType
 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_transcendental_patterns, pm_float_decomp, pm_long_decomp
+from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_unsupported_dtypes_patterns, get_transcendental_patterns
 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, pm_syntactic_sugar
+from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 
-def _sym_rewrite(sink:UOp, sym_pm:PatternMatcher, extra_pm:PatternMatcher|None=None, ctx=None, name:str|None=None) -> UOp:
-  """Symbolic rewrite: uses e-graph extraction when EGRAPH is set, otherwise greedy graph_rewrite."""
-  if EGRAPH:
-    from tinygrad.uop.egraph import egraph_rewrite
-    return egraph_rewrite(sink, sym_pm, extra_pm, ctx=ctx, name=name)
-  return graph_rewrite(sink, sym_pm+extra_pm if extra_pm is not None else sym_pm, ctx=ctx, name=name)
+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 VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Base AST")
+  if getenv("VIZ"): graph_rewrite(sink, PatternMatcher([]), name="View Base AST")
   if DEBUG >= 5: print(pyrender(sink))
   if SPEC: type_verify(sink, kernel_spec)
 
@@ -48,7 +45,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
     sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")
 
     # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
-    sink = _sym_rewrite(sink, sym, pm_flatten_range, name="initial symbolic")
+    sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")
 
     # optimize (schedule) the AST
     sink = graph_rewrite(sink, pm_simplify_ranges, name="simplify ranges")
@@ -60,10 +57,10 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
     sink = apply_opts(sink, ren)
 
   # ** expander (expand_rewrite) **
-  sink = _sym_rewrite(sink, sym, pm_move_where_on_load, name="postopt symbolic")
+  sink = graph_rewrite(sink, sym+pm_move_where_on_load, name="postopt symbolic")
 
   # expand
-  sink = _sym_rewrite(sink, sym, pm_pre_expander+pm_group_for_reduce+expander, name="expander")
+  sink = graph_rewrite(sink, sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander")
 
   # add locals
   sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, ctx=itertools.count(0), name="add local buffers")
@@ -81,33 +78,29 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
   sink = graph_rewrite(sink, pm_add_loads, name="** add loads (code)")
 
   # devectorize (TODO: does this need opts?)
-  if DEVECTORIZE >= 2: pm_devec_extra = load_store_folding+load_store_indexing
-  elif DEVECTORIZE: pm_devec_extra = devectorize+load_store_folding+correct_load_store+load_store_indexing
-  else: pm_devec_extra = load_store_folding+correct_load_store+load_store_indexing
-  if DEVECTORIZE >= 0: sink = _sym_rewrite(sink, sym, pm_devec_extra, ctx=ren, name="devectorize")
+  if DEVECTORIZE >= 2: pm_devectorize = sym+load_store_folding+load_store_indexing
+  elif DEVECTORIZE: pm_devectorize = sym+devectorize+load_store_folding+correct_load_store+load_store_indexing
+  else: pm_devectorize = sym+load_store_folding+correct_load_store+load_store_indexing
+  if DEVECTORIZE >= 0: sink = graph_rewrite(sink, pm_devectorize, ctx=ren, name="devectorize")
 
   # lower the index dtype to a concrete int
   sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing+gep_pushing, ctx=ren.device, name="lower all index dtypes")
-  sink = _sym_rewrite(sink, symbolic, name="post index symbolic")
+  sink = graph_rewrite(sink, symbolic, name="post index symbolic")
 
   # optional pre matcher
   if ren.pre_matcher is not None: sink = graph_rewrite(sink, ren.pre_matcher, name="pre_matcher")
 
   # decompositions
   supported_ops = tuple(ren.code_for_op.keys())
-  pm_decomp_extra = get_late_rewrite_patterns(supported_ops, ren.device, bool(DISABLE_FAST_IDIV))
-  pm_transcend_extra = get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
-  sink = _sym_rewrite(sink, symbolic_simple, pm_decomp_extra, ctx=ren.device, name="decompositions")
-  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 = _sym_rewrite(sink, symbolic_simple, pm_transcend_extra, ctx=ren.device, name="transcendental")
+  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)
+  sink = graph_rewrite(sink, pm_transcendental, ctx=ren.device, name="transcendental")
 
   # final rules for the renderer (without sym)
   extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
-  pm_decomp = symbolic_simple+pm_decomp_extra
   pm_final_rewrite = pm_decomp+pm_render+extra_matcher+pm_split_ends
   sink = graph_rewrite(sink, pm_final_rewrite, ctx=ren.device, name="final rewrite")
 
@@ -120,7 +113,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,))]))
@@ -171,19 +164,17 @@ def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> Program
     The ProgramSpec of the program.
   """
 
+  # fix up KernelInfo
+  if opts is not None:
+    assert ast.arg is None, "can't apply opts if sink has an arg"
+    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
+  if ast.arg is None and ast.op is Ops.SINK: ast = ast.replace(arg=KernelInfo())
+
+  # rewrite to prg
   if ast.op is Ops.PROGRAM: prg = ast
-  elif ast.op is Ops.SINK:
-    # rewrite to prg
-    assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to get_program"
-    if opts is not None:
-      # TODO: should this be here?
-      assert ast.arg.opts_to_apply is None, "can't apply opts if there's already opts to apply"
-      ast = ast.replace(arg=replace(ast.arg, opts_to_apply=tuple(opts)))
+  else:
     full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
     prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.device)))
-  else:
-    raise RuntimeError(f"can't call get_program on {ast.op}")
-
   prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
 
   # create the ProgramSpec

tinygrad/codegen/late/linearizer.py

@@ -7,13 +7,17 @@ from tinygrad.helpers import prod, getenv, TUPLE_ORDER
 def linearize(sink:UOp) -> list[UOp]:
   # this is a toposort with priority
   lst = list(sink.toposort())
-  out_degree:defaultdict[UOp, int] = defaultdict(int)
+  consumers: defaultdict[UOp, list[UOp]] = defaultdict(list)
+  in_degree:dict[UOp, int] = {}
+  out_degree:dict[UOp, 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: out_degree[s] += 1
+    for s in u.src: consumers[s].append(u)
+    in_degree[u] = len(u.src)
+    out_degree[u] = len(consumers[u])
 
     # we place UOps with higher run_counts later
     run_count = prod([int(r.vmax)+1 for r in u.ranges])

tinygrad/device.py

@@ -56,13 +56,14 @@ atexit.register(lambda: [Device[dn].finalize() for dn in Device._opened_devices]
 # **************** Profile ****************
 
 @dataclass(frozen=True)
-class ProfileDeviceEvent(ProfileEvent): device:str; tdiff:decimal.Decimal=decimal.Decimal(0); props:dict[str,Any]|None=None # noqa: E702
+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
 
 @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 # noqa: E702
+class ProfileGraphEntry: device:str; name:str; st_id:int; en_id:int; is_copy:bool # noqa: E702
 
 @dataclass(frozen=True)
 class ProfileGraphEvent(ProfileEvent): ents:list[ProfileGraphEntry]; deps:list[list[int]]; sigs:list[decimal.Decimal] # noqa: E702
@@ -406,9 +407,9 @@ def enumerate_devices_str() -> Generator[str, None, None]:
             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 ''
             default_text = '(default)' if type(default_compiler) is type(d.compiler) else set_text
-            compilers_results.append(f"{colored('+', 'green')} {d._compiler_name(r, c)} {default_text}")
+            compilers_results.append(f"{colored('+', 'green')} {unwrap_class_type(c).__name__} {default_text}")
             any_works = True
-          except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {d._compiler_name(r, c)}: {e}")
+          except Exception as e: compilers_results.append(f"{colored('-', 'yellow')} {unwrap_class_type(c).__name__}: {e}")
       finally:
         # put the defaults back!
         d.comp_sets, d.comps_ctrl_var = default_comp_pairs, cc_ctrl_var

tinygrad/dtype.py

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

tinygrad/engine/schedule.py

@@ -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, gate_kernel_sink
+from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, CallInfo
 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,13 @@ 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(gate_kernel_sink):
+    for u in sched_sink.toposort():
       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.CALL, Ops.END}, f"AFTER src[1] should be KERNEL or END, not {k.op}"
+      assert k.op in {Ops.KERNEL, 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.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:]:
+      for s in k.src[0].src 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)
@@ -60,8 +59,7 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
         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)
+        schedule.append((ast, buf_uops, cast(CallInfo, k.arg).metadata, {}, bound_ranges))
         if rk.op is Ops.END: schedule.append(rk)
       for x in children.get(rk, []):
         in_degree[x] -= 1
@@ -86,7 +84,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.CALL, f"unexpected op in schedule: {si.op}"
+      assert si.op is Ops.KERNEL, 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))

tinygrad/gradient.py

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

tinygrad/helpers.py

@@ -86,9 +86,7 @@ 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)
-
-# 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!")
+def panic(e:Exception|None=None): raise e if e is not None else RuntimeError("PANIC!")
 
 @functools.cache
 def canonicalize_strides(shape:tuple[T, ...], strides:tuple[T, ...]) -> tuple[T, ...]:
@@ -180,7 +178,7 @@ SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), Contex
 RING, ALL2ALL = ContextVar("RING", 1), ContextVar("ALL2ALL", 0)
 CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
 VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
-CORRECT_DIVMOD_FOLDING, FUSE_OPTIM, EGRAPH = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0), ContextVar("EGRAPH", 0)
+CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
 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)))
@@ -289,7 +287,8 @@ class TracingKey:
 class ProfileEvent: pass
 
 @dataclass
-class ProfileRangeEvent(ProfileEvent): device:str; name:str|TracingKey; st:decimal.Decimal; en:decimal.Decimal|None=None # noqa: E702
+class ProfileRangeEvent(ProfileEvent):
+  device:str; name:str|TracingKey; st:decimal.Decimal; en:decimal.Decimal|None=None; is_copy:bool=False # noqa: E702
 
 @dataclass(frozen=True)
 class ProfilePointEvent(ProfileEvent):
@@ -297,8 +296,8 @@ class ProfilePointEvent(ProfileEvent):
 
 cpu_events:list[ProfileEvent] = []
 @contextlib.contextmanager
-def cpu_profile(name:str|TracingKey, device="TINY", display=True) -> Generator[ProfileRangeEvent, None, None]:
-  res = ProfileRangeEvent(device, name, perf_counter_us())
+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)
   try: yield res
   finally:
     res.en = perf_counter_us()

tinygrad/renderer/wgsl.py

@@ -69,12 +69,9 @@ 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.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, 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, 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),

tinygrad/runtime/autogen/__init__.py

@@ -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)'), (r'(\d+):(\d+)', 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)')], 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",

tinygrad/runtime/graph/hcq.py

@@ -131,8 +131,7 @@ 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
 
-        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_entries.append(ProfileGraphEntry(enqueue_dev.device, prof_ji_desc, sig_st, j * 2 + 1, is_copy=not is_exec_prg))
         self.prof_graph_deps.append([d - 1 for _, d in rdeps])
 
       last_j[enqueue_queue] = j

tinygrad/runtime/graph/metal.py

@@ -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) for i,ji in enumerate(self.jit_cache)]
+    ents = [ProfileGraphEntry(self.device, cast(CompiledRunner, ji.prg)._prg.name, i, i+1, is_copy=False) 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)])]
 

tinygrad/runtime/ops_amd.py

@@ -8,7 +8,7 @@ from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filte
 from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerSet, CompilerPair
 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, unwrap
+from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, ceildiv
 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,15 +703,13 @@ 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
 
@@ -780,20 +778,18 @@ 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=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")
+    kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=self.queue_event_arr_ptr, num_events=1, wait_for_all=1, timeout=tm)
 
   def on_device_hang(self):
-    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)
+    def _collect_str(st): return ' '.join(f'{k[0]}={getattr(st, k[0])}' for k in st._real_fields_)
 
     report = []
-    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)}"]
+    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)}"]
 
     raise RuntimeError("\n".join(report))
 
@@ -959,7 +955,6 @@ 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),
@@ -1023,12 +1018,11 @@ 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
-    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)
+    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
 
   def _ensure_has_local_memory(self, private_segment_size):
     if self.max_private_segment_size >= private_segment_size: return
@@ -1069,5 +1063,3 @@ 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]

tinygrad/runtime/ops_metal.py

@@ -51,7 +51,7 @@ 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)]
+        Compiled.profile_events += [ProfileRangeEvent(self.device, lb, st, en, is_copy=lb.startswith("COPY"))]
     self.mtl_buffers_in_flight.clear()
 
 def metal_src_to_library(device:MetalDevice, src:str) -> metal.MTLLibrary:
@@ -191,7 +191,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): dst[:] = src
+    with cpu_profile(prof_desc, self.dev.device, is_copy=True): dst[:] = src
   def _as_buffer(self, src:MetalBuffer) -> memoryview:
     self.dev.synchronize()
     return to_mv(src.buf.contents(), src.size + src.offset)[src.offset:]

tinygrad/runtime/ops_nv.py

@@ -36,9 +36,6 @@ 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}")
@@ -97,8 +94,7 @@ 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),
-             nv_flags("NVC56F_SEM_EXECUTE", operation="acq_circ_geq", payload_size="64bit"))
+    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.active_qmd = None
     return self
 
@@ -129,8 +125,7 @@ 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,
-             nv_flags("NVC6C0_INVALIDATE_SHADER_CACHES_NO_WFI", instruction="true", global_data="true", constant="true"))
+    self.nvm(1, nv_gpu.NVC6C0_INVALIDATE_SHADER_CACHES_NO_WFI, (1 << 12) | (1 << 4) | (1 << 0))
     self.active_qmd:QMD|None = None
     return self
 
@@ -174,7 +169,7 @@ class NVComputeQueue(NVCommandQueue):
           return self
 
     self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(signal.value_addr), *data64_le(value),
-             nv_flags("NVC56F_SEM_EXECUTE", operation="release", release_wfi="en", payload_size="64bit", release_timestamp="en"))
+             (1 << 0) | (1 << 20) | (1 << 24) | (1 << 25)) # RELEASE | RELEASE_WFI | PAYLOAD_SIZE_64BIT | RELEASE_TIMESTAMP
     self.nvm(0, nv_gpu.NVC56F_NON_STALL_INTERRUPT, 0x0)
     self.active_qmd = None
     return self
@@ -190,13 +185,12 @@ 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,
-               nv_flags("NVC6B5_LAUNCH_DMA", data_transfer_type="non_pipelined", src_memory_layout="pitch", dst_memory_layout="pitch"))
+      self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA, 0x182) # TRANSFER_TYPE_NON_PIPELINED | DST_MEMORY_LAYOUT_PITCH | SRC_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, nv_flags("NVC6B5_LAUNCH_DMA", flush_enable="true", semaphore_type="release_four_word_semaphore"))
+    self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA, 0x14)
     return self
 
   def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.dma_gpfifo)
@@ -205,8 +199,7 @@ 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, 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_CONTROL_PARAMS, 0x52057)
     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]):
@@ -223,7 +216,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, nv_flags("NVC9B0_SEMAPHORE_D", structure_size="four", payload_size="64bit"))
+    self.nvm(4, nv_gpu.NVC9B0_SEMAPHORE_D, (1 << 24) | (1 << 0))
     return self
 
   def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.vid_gpfifo)
@@ -340,7 +333,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="HEVC Decode", enabled=PROFILE, dev_suff="NVDEC"):
+    with hcq_profile(self.dev, queue=q, desc="NVDEC", enabled=PROFILE):
       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)

tinygrad/runtime/ops_qcom.py

@@ -331,7 +331,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):
+    with cpu_profile(prof_text, self.dev.device, is_copy=True):
       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

tinygrad/runtime/support/am/ip.py

@@ -264,8 +264,7 @@ 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} if self.adev.ip_ver[am.GC_HWIP][:2]==(9,4) else {}),
+        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}),
           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:

tinygrad/runtime/support/hcq.py

@@ -257,14 +257,15 @@ class HCQSignal(Generic[HCQDeviceType]):
       value: The value to wait for.
       timeout: Maximum time to wait in milliseconds. Defaults to 30s.
     """
-    start_time = int(time.perf_counter() * 1000)
+    start_time = last_sleep_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 - start_time)
-      if self.value != prev_value: start_time = int(time.perf_counter() * 1000) # progress was made, reset timer
+      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
     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, dev_suff:str|None=None):
+def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]|None=None, queue:HWQueue|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"
 
@@ -278,7 +279,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, f"{dev.device}:{dev_suff}" if dev_suff else dev.device))
+    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))
 
 class HCQArgsState(Generic[ProgramType]):
   def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint|None, ...]=()):
@@ -375,7 +376,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, str]] = []
+    self.sig_prof_records:list[tuple[HCQSignal, HCQSignal, str, bool]] = []
     self.prof_exec_counter:int = 0
     self.prof_prg_counter:int = 0
 
@@ -401,7 +402,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
 
     if self.timeline_value > (1 << 31): self._wrap_timeline_signal()
     if PROFILE:
-      Compiled.profile_events += [ProfileRangeEvent(dev, name, st.timestamp, en.timestamp) for st,en,name,dev in self.sig_prof_records]
+      Compiled.profile_events += [ProfileRangeEvent(self.device, name, st.timestamp, en.timestamp, cp) for st,en,name,cp in self.sig_prof_records]
       self.sig_prof_records = []
 
   def next_timeline(self):
@@ -417,10 +418,6 @@ 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
 
@@ -431,9 +428,10 @@ class HCQCompiled(Compiled, Generic[SignalType]):
       et = time.perf_counter_ns()
       return (decimal.Decimal(et+st) / 2000) - d.timeline_signal.timestamp
 
-    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())]
+    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())]
 
   def _wrap_timeline_signal(self):
     self.timeline_signal, self._shadow_timeline_signal, self.timeline_value = self._shadow_timeline_signal, self.timeline_signal, 1
@@ -516,10 +514,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): ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
+      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))
       return
 
-    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"):
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"TINY -> {self.dev.device}", enabled=PROFILE):
       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])
@@ -540,7 +538,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, dev_suff="SDMA:0"):
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"DISK -> {self.dev.device}", enabled=PROFILE):
       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) \
@@ -550,10 +548,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): ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
+      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))
       return
 
-    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"):
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"{self.dev.device} -> TINY", enabled=PROFILE):
       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)) \
@@ -565,7 +563,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, dev_suff="SDMA:0"):
+    with hcq_profile(src_dev, queue_type=src_dev.hw_copy_queue_t, desc=f"{src_dev.device} -> {dest_dev.device}", enabled=PROFILE):
       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) \

tinygrad/schedule/indexing.py

@@ -3,7 +3,7 @@ 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, gate_kernel_sink, pm_gate_kernel_sink
+from tinygrad.uop.ops import consumer_map_from_toposort
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored
 
@@ -17,15 +17,22 @@ 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
 
-pm_generate_realize_map = pm_gate_kernel_sink+PatternMatcher([
+def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
+  if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
+  # if it's a kernel, we don't realize it
+  if a.src[1].op is not Ops.CALL: ctx[a] = None
+
+pm_generate_realize_map = PatternMatcher([
   # always realize SINK src
   (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
-  # always realize
-  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE, Ops.ASSIGN, Ops.ENCDEC}, name="tr"), realize),
+  # 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 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 these
-  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK, Ops.ASSIGN, Ops.ENCDEC), name="rb"), realize_srcs),
+  # 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),
 ])
 
 @dataclass(frozen=True)
@@ -91,25 +98,20 @@ def convert_reduce_axis_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
 def remove_movement_op_after_rangeify(ctx:IndexingContext, x:UOp):
   if x in ctx.range_map or x.src[0].op is Ops.INDEX: return x.src[0]
 
-def handle_assign_mops(ctx:IndexingContext, assign:UOp, target:UOp, src:UOp):
-  if target.op in GroupOp.Movement and src.op is not Ops.CALL:
-    mops = []
-    while target.op in GroupOp.Movement:
-      mops.append((target.op, target.marg))
-      target = target.src[0]
-    if mops and assign in ctx.range_map:
-      ret = assign.replace(arg=tuple(mops))
-      ctx.range_map[ret] = ctx.range_map[assign]
-      return ret
-  return None
+def add_third_op_to_assign_to_track_shape(ctx:IndexingContext, assign:UOp):
+  if assign.src[1].op is Ops.CALL: 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
 
 pm_apply_rangeify = PatternMatcher([
   # REDUCE_AXIS -> REDUCE
   (UPat(Ops.REDUCE_AXIS, name="x"), convert_reduce_axis_to_reduce_with_ranges),
   # PAD -> WHERE
   (UPat(Ops.PAD, name="x"), convert_pad_to_where_to_keep_behavior_local),
-  # store movement ops in ASSIGN arg
-  (UPat(Ops.ASSIGN, src=(UPat(name="target"), UPat(name="src")), name="assign"), handle_assign_mops),
+  # add third op to assign
+  (UPat(Ops.ASSIGN, src=(UPat(), UPat()), name="assign"), add_third_op_to_assign_to_track_shape),
   # finally, apply_rangeify
   (UPat(GroupOp.All, name="x"), create_bufferize_and_index_based_on_ranges),
   # remove movement op
@@ -160,11 +162,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, bottom_up=True, name="get realize")
+  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, 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(gate_kernel_sink))
+    consumer_map = consumer_map_from_toposort(tsink_toposort:=tsink.toposort())
 
   # explicit rangeify
   ending_ranges: dict[UOp, list[UOp]] = {}

tinygrad/schedule/multi.py

@@ -1,8 +1,7 @@
 from typing import cast
 import functools, itertools
-from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, ALL2ALL, VIZ, getenv
+from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, ALL2ALL, 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 ***
@@ -215,18 +214,17 @@ 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),
-  # 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)),
+  # 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))
 ])+replace_allreduce
 
 def get_multi_map(big_sink:UOp) -> dict[UOp, UOp]:
-  if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Multi AST")
+  if getenv("VIZ"): graph_rewrite(big_sink, PatternMatcher([]), name="View Multi AST")
   ret = graph_rewrite_map(big_sink, multi_pm, name="multi_pm")
-  if VIZ: graph_rewrite(ret[big_sink], PatternMatcher([]), name="View Post Multi AST")
+  if getenv("VIZ"): graph_rewrite(ret[big_sink], PatternMatcher([]), name="View Post Multi AST")
   return ret

tinygrad/schedule/rangeify.py

@@ -1,11 +1,11 @@
 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, pm_gate_kernel_sink
+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, _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, VIZ
-from tinygrad.helpers import PCONTIG, partition, get_single_element
+from tinygrad.helpers import argsort, prod, all_same, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY
+from tinygrad.helpers import PCONTIG, partition, get_single_element, panic
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
@@ -14,12 +14,6 @@ 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"),
@@ -33,17 +27,13 @@ pm_mops = PatternMatcher([
 # *****************
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
-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):
+def fix_assign_hazard(dest:UOp, src:UOp, assign:UOp):
   # PERMUTE and FLIP reorder indices, SHRINK can have overlapping regions when dest is also shrunk
-  unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set())
+  unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if dest.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 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()))
+    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()))
 
 def split_reduceop(reduce:UOp, x:UOp):
   if prod(reduce.shape) == 0: return None
@@ -74,10 +64,13 @@ 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 ck.arg.fxn(*placeholders).call(*ck.src)
+
 def resolve_call(c:UOp) -> UOp|None:
-  # don't resolve real kernel calls, sink or program
-  if c.src[0].op is Ops.SINK and isinstance(c.src[0].arg, KernelInfo): return None
-  if c.src[0].op is Ops.PROGRAM: return None
+  # don't resolve CALLs with SINK - those are kernel calls from split_store/custom_kernel
+  if c.src[0].op is Ops.SINK: 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
@@ -92,9 +85,12 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   # just removing it works...
   (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
 
-  # resolve calls
+  # resolve calls (but not CALLs with SINK - those are kernel calls from split_store)
   (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)),
 
@@ -134,13 +130,15 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
 
   # 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 assign, target, src: target.assign(src.bitcast(target.dtype)).replace(tag=assign.tag)),
+   lambda target, src, assign: 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="src")), name="assign"), assign_to_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),
 
    # make source contiguous if it has hazardous movement ops on the dest buffer
-   (UPat(Ops.ASSIGN, src=(UPat.var("target"), UPat.var("src")), name="assign"), fix_assign_hazard),
+   (UPat(Ops.ASSIGN, src=(UPat.var("dest"), UPat.var("src")), name="assign"), fix_assign_hazard),
 ])
 
 # *****************
@@ -275,14 +273,15 @@ def late_buffer_view(t:UOp, b:UOp):
   size = prod(shape)
 
   # walk up for the INDEX
+  # NOTE: even though we allow RESHAPE and SHRINK, they can combine to form non-contiguous access patterns (e.g. t[::2])
   x = t
-  while not any(u.op is Ops.INDEX for u in x.src):
-    assert x.op not in GroupOp.Elementwise, "can't buffer view elementwise"
+  while x.op is not Ops.INDEX:
+    assert x.op in {Ops.BITCAST, Ops.CONTIGUOUS, Ops.SHRINK, Ops.RESHAPE}, f"unexpected op {x.op} in buffer view walk"
     x = x.src[0]
-  x = next(u for u in x.src if u.op is Ops.INDEX)
 
   if len(shape) == 0: offset = x.src[1].arg
-  else: offset = max(sum(idx.vmin for idx in x.src[1:]), 0)
+  else: offset = sum(idx.vmin for idx in x.src[1:])
+  if offset < 0: raise RuntimeError(f"negative offset {offset} in buffer view")
 
   return b.replace(src=(UOp(Ops.BUFFER_VIEW, t.dtype, (x.base,), (size, offset), tag=t.tag), b.src[1]))
 
@@ -329,13 +328,19 @@ 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 (assign := x.src[0]).op is Ops.ASSIGN:
-    assign_target, assign_src = assign.src[0], assign.src[1]
+  if x.src[0].op is Ops.ASSIGN:
+    assign_target, assign_src, assign_mops = x.src[0].src
     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)
-    for op, marg in reversed(assign.arg or ()): ret = ret._mop(op, marg)
+    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)
     return ret
 
   # lower outerworld reduce here
@@ -381,8 +386,8 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
   (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
    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.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))),
+  # remove any RESHAPEs on CALL
+  (UPat(Ops.CALL, name="k"), lambda k: k.replace(src=(k.src[0],)+tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src[1:]))),
 ])
 
 pm_add_buffers_local = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
@@ -511,8 +516,8 @@ def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
   elif ret.op is Ops.END and ret.src[0].op is Ops.STORE: stored = ret.src[0].src[1]
   else: raise RuntimeError(f"unknown kernel type {ret.op}")
   if stored.op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENCDEC}: ret = stored
-  else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))
-
+  else:
+    ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts) if lctx.opts is not None else KernelInfo())
   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]):
@@ -521,6 +526,8 @@ def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
 
 split_kernels = PatternMatcher([
   (UPat((Ops.STORE, Ops.END), name="x"), split_store),
+  # if it's a Kernel, stop
+  (UPat(Ops.SINK, name="sink"), lambda sink: panic(BottomUpGate()) if isinstance(sink.arg, KernelInfo) else None),
 ])
 
 def tag_uop(ctx:tuple[list[UOp], set[UOp]], x:UOp):
@@ -531,7 +538,7 @@ 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 = pm_gate_kernel_sink+PatternMatcher([
+add_tags = 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.CALL, Ops.END,
                      Ops.MSTACK, Ops.MSELECT, Ops.RANGE}.union(GroupOp.Movement), name="x"), tag_uop),
@@ -554,11 +561,11 @@ replace_contiguous = PatternMatcher([
 ])
 
 def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
-  if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Input Graph")
+  if getenv("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_syntactic_sugar+pm_mops+earliest_rewrites+replace_contiguous, ctx={}, bottom_up=True, name="earliest rewrites")
+  tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")
 
   # convert movement ops to ranges
   tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
@@ -572,13 +579,12 @@ 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 VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")
+  if getenv("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_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")
+  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", bottom_up=True)
 
   # 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] = {}
@@ -598,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 VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
+  if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
 
   becomes_map: dict[UOp, UOp] = {}
   for tag, s in zip(sink_tags, tsink.src):

tinygrad/tensor.py

@@ -234,7 +234,8 @@ class Tensor(OpMixin):
 
   def as_param(self, slot:int):
     if self.uop.axis is not None:
-      param = UOp.param(slot, self.dtype, self.uop.shard_shape, self.device).multi(self.uop.axis)
+      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)
     else:
       param = UOp.param(slot, self.dtype, self.shape, self.device)
     return Tensor(param, device=self.device)
@@ -755,7 +756,8 @@ 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)
-    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]))
+    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]))
     return Tensor(UOp.multi(stacked, axis=self.uop.axis), device=self.device, dtype=dtype)
 
   def full_like(self, fill_value:PyConst, **kwargs) -> Tensor:
@@ -3784,7 +3786,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)
+    V = V.gather(-1, new_indices).realize()
 
     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

tinygrad/uop/__init__.py

@@ -79,7 +79,8 @@ class Ops(FastEnum):
   # ** 6 -- ops that don't exist in programs **
 
   # tensor graph ops
-  UNIQUE = auto(); DEVICE = auto(); ASSIGN = auto()
+  UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); ASSIGN = auto()
+  CUSTOM_KERNEL = auto()
 
   # local unique
   LUNIQUE = auto()

tinygrad/uop/decompositions.py

@@ -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 (1 << (dtypes.finfo(d.scalar())[0] - 1)) - 1
-def exponent_mask(d:DType) -> int: return (1 << dtypes.finfo(d.scalar())[0]) - 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()]
 
 # **** utils ****
 def shr(x:UOp|int, y:UOp|int) -> UOp: return x // (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y)
@@ -378,46 +378,33 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
     case _: raise NotImplementedError(f"long decomposition of {op} unsupported")
 
 # ***** floats *****
-f2f_dt = { f:getattr(dtypes, f"uint{f.bitsize}") for f in dtypes.floats }
+f2f_dt = { dtypes.half: dtypes.ushort, dtypes.float: dtypes.uint }
 
 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)
-    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)
+    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)
     norm = (rne(nosign, fm - tm) - shl(fb - tb, tm)).cast(f2f_dt[to])
-    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))
+    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)))
   else: raise NotImplementedError(f"unsupported decomp {fr} -> {to}")
 
-def f2f_clamp(val:UOp, dt:DType) -> UOp:
-  e, m = dtypes.finfo(dt)
-  max_exp, max_man = ((1 << e) - 1, (1 << m) - 2) if dt == dtypes.fp8e4m3 else ((1 << e) - 2, (1 << m) - 1)
-  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_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_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)))
+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)))
 
 # ***** decomposition patterns *****
 
@@ -476,44 +463,40 @@ 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)
 
-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),
-])
+@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)

tinygrad/uop/egraph.py

@@ -1,227 +0,0 @@
-# e-graph (equality saturation) for UOp rewriting
-# instead of greedy first-match rewriting, we explore ALL equivalent forms and extract the cheapest
-from __future__ import annotations
-from tinygrad.uop.ops import UOp, Ops, PatternMatcher, graph_rewrite
-
-# *** union-find (keyed by UOp identity) ***
-
-def uf_find(parent:dict[UOp, UOp], x:UOp) -> UOp:
-  while parent[x] is not x:
-    parent[x] = parent[parent[x]]
-    x = parent[x]  # path compression
-  return x
-
-def uf_union(parent:dict[UOp, UOp], size:dict[UOp, int], a:UOp, b:UOp) -> UOp:
-  a, b = uf_find(parent, a), uf_find(parent, b)
-  if a is b: return a
-  if size[a] < size[b]: a, b = b, a  # merge smaller into larger
-  parent[b] = a
-  size[a] += size[b]
-  return a
-
-# *** e-graph core ***
-
-def rewrite_all(pm:PatternMatcher, uop:UOp, ctx=None) -> list[UOp]:
-  """Apply ALL matching rewrite rules to uop, returning every distinct result."""
-  results: list[UOp] = []
-  seen: dict[UOp, None] = {}
-  for _, match, early_reject in pm.pdict.get(uop.op, []):
-    if not early_reject.issubset({u.op for u in uop.src}): continue
-    try: ret = match(uop, ctx)
-    except Exception: continue  # skip rules that crash on this node (e.g. division by zero in divmod folding)
-    if ret is not None and ret is not uop and ret not in seen:
-      results.append(ret)
-      seen[ret] = None
-  return results
-
-class EGraph:
-  """E-graph with full equality saturation (including rebuilding)."""
-  __slots__ = ("parent", "size", "eclass", "eclass_uses", "all_nodes")
-  def __init__(self, root:UOp):
-    nodes = list(root.toposort())
-    self.parent: dict[UOp, UOp] = {u: u for u in nodes}
-    self.size: dict[UOp, int] = {u: 1 for u in nodes}
-    self.eclass: dict[UOp, dict[UOp, None]] = {u: {u: None} for u in nodes}  # canonical -> members
-    # canonical eclass representative -> dict of nodes that USE this eclass as a child
-    self.eclass_uses: dict[UOp, dict[UOp, None]] = {u: {} for u in nodes}
-    self.all_nodes: dict[UOp, None] = dict.fromkeys(nodes)
-    # build initial parent-child uses
-    for u in nodes:
-      for s in u.src:
-        canon = uf_find(self.parent, s)
-        self.eclass_uses.setdefault(canon, {})[u] = None
-
-  def _add_node(self, u:UOp):
-    """Register a new UOp (and its subtree) in the e-graph."""
-    for sub in u.toposort():
-      if sub in self.parent: continue
-      self.parent[sub] = sub
-      self.size[sub] = 1
-      self.eclass[sub] = {sub: None}
-      self.all_nodes[sub] = None
-      self.eclass_uses[sub] = {}
-      for s in sub.src:
-        canon = uf_find(self.parent, s)
-        self.eclass_uses.setdefault(canon, {})[sub] = None
-
-  def _merge(self, a:UOp, b:UOp) -> UOp|None:
-    """Merge two e-classes. Returns the winner, or None if already merged."""
-    ra, rb = uf_find(self.parent, a), uf_find(self.parent, b)
-    if ra is rb: return None
-    winner = uf_union(self.parent, self.size, ra, rb)
-    loser = rb if winner is ra else ra
-    self.eclass[winner] = {**self.eclass[winner], **self.eclass[loser]}
-    # merge uses
-    winner_uses = self.eclass_uses.setdefault(winner, {})
-    winner_uses.update(self.eclass_uses.pop(loser, {}))
-    del self.eclass[loser]
-    return winner
-
-  def _canonical(self, u:UOp) -> UOp:
-    """Rebuild node with canonical representative for each child's eclass."""
-    if not u.src: return u
-    new_src = []
-    for s in u.src:
-      canon = uf_find(self.parent, s)
-      members = self.eclass.get(canon)
-      if members is not None:
-        best = min(members, key=lambda m: (len(m.src), m.op.value, m.arg if isinstance(m.arg, (int, float, str)) else 0))
-        new_src.append(best)
-      else:
-        new_src.append(s)
-    new_src_tuple = tuple(new_src)
-    if new_src_tuple == u.src: return u
-    return UOp(u.op, u.dtype, new_src_tuple, u.arg, u.tag)
-
-  def _rebuild(self, dirty:dict[UOp, None]) -> list[tuple[UOp, UOp]]:
-    """Rebuild parents of dirty eclasses, creating canonical versions."""
-    new_equalities: list[tuple[UOp, UOp]] = []
-    affected: dict[UOp, None] = {}
-    for d in dirty:
-      canon = uf_find(self.parent, d)
-      affected.update(self.eclass_uses.get(canon, {}))
-    for u in affected:
-      rebuilt = self._canonical(u)
-      if rebuilt is not u:
-        if rebuilt in self.parent and uf_find(self.parent, rebuilt) is uf_find(self.parent, u): continue
-        self._add_node(rebuilt)
-        new_equalities.append((u, rebuilt))
-    return new_equalities
-
-def egraph_saturate(root:UOp, pm:PatternMatcher, max_iters:int=10, ctx=None) -> dict[UOp, dict[UOp, None]]:
-  """Build an e-graph with full equality saturation (with rebuilding). Returns eclass map."""
-  eg = EGraph(root)
-  node_limit = len(eg.all_nodes) * 3  # stop growing at 3x initial size to prevent combinatorial blowup
-  worklist: dict[UOp, None] = dict(eg.all_nodes)  # nodes to match rules on
-  for _ in range(max_iters):
-    # phase 1: match rules only on worklist nodes
-    new_equalities: list[tuple[UOp, UOp]] = []
-    next_worklist: dict[UOp, None] = {}
-    prev_nodes = dict(eg.all_nodes)
-    for u in list(worklist):
-      if len(eg.all_nodes) >= node_limit: break
-      for new in rewrite_all(pm, u, ctx):
-        if new in eg.parent and uf_find(eg.parent, new) is uf_find(eg.parent, u): continue
-        eg._add_node(new)
-        new_equalities.append((u, new))
-    # all newly added nodes (including sub-nodes of rewrite results) go on next worklist
-    for u in eg.all_nodes:
-      if u not in prev_nodes: next_worklist[u] = None
-    if not new_equalities: break
-
-    # phase 2: merge eclasses, then rebuild canonical forms (no rule matching in rebuild)
-    while new_equalities:
-      dirty: dict[UOp, None] = {}
-      for a, b in new_equalities:
-        merged = eg._merge(a, b)
-        if merged is not None: dirty[merged] = None
-      if not dirty: break
-      new_equalities = eg._rebuild(dirty)
-      for _, b in new_equalities: next_worklist[b] = None
-    worklist = next_worklist
-
-  return eg.eclass
-
-# *** cost model ***
-
-OP_COST: dict[Ops, int] = {
-  Ops.CONST: 0, Ops.VCONST: 0, Ops.DEFINE_VAR: 0,
-  Ops.ADD: 1, Ops.MUL: 2, Ops.SUB: 1, Ops.NEG: 1,
-  Ops.IDIV: 5, Ops.MOD: 5, Ops.FDIV: 3,
-  Ops.SHL: 1, Ops.SHR: 1,
-  Ops.AND: 1, Ops.OR: 1, Ops.XOR: 1,
-  Ops.MAX: 1, Ops.CMPLT: 1, Ops.CMPNE: 1, Ops.CMPEQ: 1,
-  Ops.CAST: 1, Ops.BITCAST: 1,
-  Ops.WHERE: 2, Ops.MULACC: 2,
-  Ops.EXP2: 8, Ops.LOG2: 8, Ops.SIN: 8, Ops.SQRT: 4, Ops.RECIPROCAL: 3,
-  Ops.POW: 10, Ops.TRUNC: 1,
-}
-
-def node_cost(u:UOp) -> int:
-  c = OP_COST.get(u.op, 3)
-  # tiebreaker: penalize non-canonical operand order (consts should be on the right for commutative ops)
-  if len(u.src) == 2 and u.src[0].op is Ops.CONST and u.src[1].op is not Ops.CONST: c += 1
-  return c
-
-# *** extraction ***
-
-def egraph_extract(root:UOp, pm:PatternMatcher, max_iters:int=10, ctx=None) -> UOp:
-  """Run equality saturation on root, then extract the cheapest equivalent expression."""
-  eclass = egraph_saturate(root, pm, max_iters, ctx)
-
-  # build eclass lookup: node -> canonical eclass representative
-  eclass_of: dict[UOp, UOp] = {}
-  for canon, members in eclass.items():
-    for u in members: eclass_of[u] = canon
-
-  all_nodes: list[UOp] = [u for members in eclass.values() for u in members]
-
-  # bottom-up DP: for each eclass, find the cheapest representative
-  cost_of: dict[UOp, tuple[int, UOp]] = {}  # eclass_canon -> (cost, best_uop)
-
-  depth_cache: dict[UOp, int] = {}
-  def _depth(u:UOp) -> int:
-    if u in depth_cache: return depth_cache[u]
-    depth_cache[u] = 0  # break cycles
-    depth_cache[u] = (1 + max((_depth(s) for s in u.src), default=0)) if u.src else 0
-    return depth_cache[u]
-
-  for u in sorted(all_nodes, key=_depth):
-    canon = eclass_of[u]
-    child_cost = 0
-    for s in u.src:
-      if (s_canon := eclass_of.get(s)) is not None and s_canon in cost_of: child_cost += cost_of[s_canon][0]
-      else: child_cost += node_cost(s)
-    total = node_cost(u) + child_cost
-    if canon not in cost_of or total < cost_of[canon][0]:
-      cost_of[canon] = (total, u)
-
-  root_canon = eclass_of.get(root)
-  if root_canon is not None and root_canon in cost_of: return _rebuild_tree(cost_of[root_canon][1], eclass_of, cost_of)
-  return root
-
-def _rebuild_tree(u:UOp, eclass_of:dict[UOp, UOp], cost_of:dict[UOp, tuple[int, UOp]], cache:dict[UOp, UOp]|None=None) -> UOp:
-  """Recursively rebuild a UOp tree, picking the cheapest representative for each child's eclass."""
-  if not u.src: return u
-  if cache is None: cache = {}
-  new_src = []
-  for s in u.src:
-    s_canon = eclass_of.get(s)
-    if s_canon is not None and s_canon in cost_of:
-      if s_canon in cache: new_src.append(cache[s_canon])
-      else:
-        cache[s_canon] = s  # placeholder breaks cycles
-        cache[s_canon] = _rebuild_tree(cost_of[s_canon][1], eclass_of, cost_of, cache)
-        new_src.append(cache[s_canon])
-    else:
-      new_src.append(_rebuild_tree(s, eclass_of, cost_of, cache))
-  new_src_tuple = tuple(new_src)
-  return u if new_src_tuple == u.src else UOp(u.op, u.dtype, new_src_tuple, u.arg, u.tag)
-
-# *** graph-level rewrite: drop-in replacement for graph_rewrite when EGRAPH is set ***
-
-def egraph_rewrite(sink:UOp, sym_pm:PatternMatcher, extra_pm:PatternMatcher|None=None, ctx=None, name:str|None=None) -> UOp:
-  """Replace graph_rewrite(sink, sym+extra, ctx) with e-graph extraction for sym, then greedy for the rest."""
-  combined = sym_pm+extra_pm if extra_pm is not None else sym_pm
-  sink = egraph_extract(sink, combined, ctx=ctx)
-  return graph_rewrite(sink, combined, ctx=ctx, name=name)

tinygrad/uop/ops.py

@@ -67,8 +67,7 @@ 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:
-      if s in ret: ret[s][u] = None
+    for s in u.src: ret[s][u] = None
   return ret
 
 def pretty_print(x:UOp, cache=None, d=0)->str:
@@ -207,7 +206,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.SINK | \
+           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.CUSTOM_KERNEL | Ops.SINK | \
            Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY:
         return None
 
@@ -236,6 +235,9 @@ 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
@@ -285,11 +287,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
             raise ValueError(f"invalid type for axis: {axis_arg}")
           return tuple(1 if i in axis_arg else s for i,s in enumerate(ps))
 
-    if self.op is Ops.ASSIGN: return self.src[1]._shape
-
     # elementwise ops keep the shape the same. all inputs with shape must match
-    if self.op in GroupOp.ALU.union({Ops.CAST, Ops.COPY, Ops.NOOP, Ops.GROUP, Ops.SINK, Ops.ALLREDUCE, Ops.STORE}):
-      input_shapes = [x._shape for x in self.src if x._shape is not None]
+    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 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]
@@ -309,8 +310,6 @@ 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
@@ -365,7 +364,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   @recursive_property
   def trace_num(self):
     num = next(ucount)
-    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 ())
+    # KERNEL has a UOp in the arg, CALL has it in src[0] so no special handling needed
+    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 ())
     return num
 
   # *** uop syntactic sugar ***
@@ -538,7 +539,6 @@ 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,11 +798,6 @@ 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}
@@ -811,7 +806,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.shard_shape, self.dtype, slot)
+    return UOp.placeholder(self.shape, self.dtype, slot)
 
   # set is store+end+after
   def set(self:UOp, val:UOp|ConstType, end:UOp|tuple[UOp, ...]|list[UOp]=()) -> UOp:
@@ -824,13 +819,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return UOp(Ops.PARAM, dtype, src, arg=slot)
 
   def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=()) -> 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))
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
     contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
-    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
-    kernel = fxn(*placeholders).call(*contig_srcs, grad_fxn=grad_fxn)
+    kernel = UOp(Ops.CUSTOM_KERNEL, src=contig_srcs, arg=CustomKernel(fxn=fxn, grad_fxn=grad_fxn))
     return [s.after(kernel) for s in contig_srcs]
 
 @dataclass(frozen=True)
@@ -844,6 +836,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
@@ -852,6 +852,12 @@ class CallInfo:
   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
+
 # ******** ops in python ********
 
 def safe_exp2(x):
@@ -1311,9 +1317,6 @@ 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
@@ -1411,6 +1414,7 @@ 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.CALL, name="u"), lambda ctx,u: f"{ctx[u.src[0]]}.call({', '.join(ctx[s] for s in u.src[1:])}, metadata={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 ")")),
 ])
 
@@ -1428,7 +1432,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 is Ops.CALL: raise NotImplementedError("call can't be pyrendered")
+    if u.op in {Ops.CUSTOM_KERNEL, Ops.CALL}: raise NotImplementedError("custom_kernel / 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
@@ -1443,6 +1447,11 @@ 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)})"

tinygrad/uop/spec.py

@@ -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
+from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, AxisType, KernelInfo, pyrender, Kernel, CustomKernel
 from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid, ConstFloat
 from tinygrad.helpers import DEBUG, Context, prod, SPEC, Metadata, panic, CHECK_OOB
 
@@ -73,6 +73,9 @@ 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([
@@ -84,8 +87,9 @@ _tensor_spec = PatternMatcher([
   (UPat(Ops.BUFFER, src=(UPat((Ops.LUNIQUE, Ops.UNIQUE)), UPat(Ops.DEVICE)), name="buf"),
    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.CALL, src=UPat((Ops.BUFFER, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
+  # CALL can attach to an AFTER to describe the compute required to realize a BUFFER
+  # src[0] is the function (SINK), src[1:] are buffers/bindings
+  (UPat(Ops.CALL, src=(UPat(Ops.SINK), UPat((Ops.BUFFER, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), allow_any_len=True), 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:])),
@@ -136,20 +140,13 @@ _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),
-
-  # ** 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
 
+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
+
 # ***** UOp spec in codegen shared between kernel and program *****
 
 shared_codegen_spec = PatternMatcher([
@@ -208,11 +205,6 @@ 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([
@@ -244,6 +236,8 @@ 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),
@@ -256,7 +250,7 @@ full_spec = PatternMatcher([
   # vectorized index
   (UPat(Ops.INDEX, src=(UPat((Ops.VECTORIZE, Ops.CAST)), UPat())), lambda: True),
 
-  # linearizer: outputs + intermediate KERNELs
+  # linearizer: outputs + intermediate CALLs
   (UPat(Ops.CALL, dtype=dtypes.void), lambda: True),
 
   # Invalid must have type Index
@@ -273,6 +267,16 @@ 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),
@@ -300,8 +304,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, "Metadata": Metadata,
-                        "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid,
+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,
                         "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace, "panic": panic,
                         "ConstFloat": ConstFloat}
 def eval_pyrender(code:str) -> UOp:

tinygrad/uop/validate.py

@@ -25,10 +25,9 @@ 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. 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)),
+  # 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)),
   # 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)),

tinygrad/viz/serve.py

@@ -45,7 +45,7 @@ from tinygrad.dtype import dtypes
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
                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.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", Ops.CUSTOM_KERNEL: "#3ebf55",
                **{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,6 +106,9 @@ 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}"
@@ -127,9 +130,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.src[0]) if u.op is Ops.CALL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}"
+    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']}"
     # NOTE: kernel already has metadata in arg
-    if TRACEMETA >= 2 and u.metadata is not None and u.op is not Ops.CALL: label += "\n"+str(u.metadata)
+    if TRACEMETA >= 2 and u.metadata is not None and u.op is not Ops.KERNEL: 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
@@ -137,10 +140,12 @@ 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):
@@ -163,19 +168,19 @@ def option(s:int|None) -> int: return 0 if s is None else s+1
 
 # Profiler API
 
-device_ts_diffs:dict[str, Decimal] = {}
-def cpu_ts_diff(device:str) -> Decimal: return device_ts_diffs.get(device, Decimal(0))
+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]
 
-amdgpu_targets:dict[str, int] = {}
+device_props:dict[str, dict] = {}
 
 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.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.is_copy)), (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)), e.sigs[ent.en_id]+diff]
+      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]
       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)
 
@@ -303,7 +308,7 @@ def load_counters(profile:list[ProfileEvent]) -> None:
     if (sqtt:=v.get(ProfileSQTTEvent)):
       for e in sqtt:
         if e.itrace: steps.append(create_step(f"PKTS SE:{e.se}", (f"/prg-pkts-{e.se}", len(ctxs), len(steps)),
-                                              data=(e.blob, prg_events[k].lib, amdgpu_targets[e.device])))
+                                              data=(e.blob, prg_events[k].lib, device_props[e.device]["gfx_target_version"])))
       steps.append(create_step("SQTT", ("/prg-sqtt", len(ctxs), len(steps)), ((k, tag), sqtt, prg_events[k])))
     ctxs.append({"name":f"Exec {name}"+(f" n{run_number[k]}" if run_number[k] > 1 else ""), "steps":steps})
 
@@ -343,7 +348,7 @@ def unpack_sqtt(key:tuple[str, int], data:list, p:ProfileProgramEvent) -> tuple[
   # * init decoder
   from extra.sqtt.roc import decode
   base = unwrap(p.base)
-  addr_table = amd_decode(unwrap(p.lib), amdgpu_targets[p.device])
+  addr_table = amd_decode(unwrap(p.lib), device_props[p.device]["gfx_target_version"])
   disasm:dict[int, tuple[str, int]] = {addr+base:(inst.disasm(), inst.size()) for addr, inst in addr_table.items()}
   rctx = decode(data, {p.tag:disasm})
   cu_events:dict[str, list[ProfileEvent]] = {}
@@ -370,22 +375,18 @@ 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, *rest = k.split()
+  dname = k.split()[0]
   dev_rank = next((v for k,v in order.items() if dname.startswith(k)), len(order))
-  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)
+  return (dev_rank, dname, len(k))
 
 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.tdiff
-      if (d:=ev.device.split(":")[0]) == "AMD":
-        device_decoders[d] = load_counters
-        amdgpu_targets[d] = unwrap(ev.props)["gfx_target_version"]
+      device_ts_diffs[ev.device] = (ev.comp_tdiff,ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff)
+      if ev.props is not None: device_props[ev.device] = ev.props
+      if (d:=ev.device.split(":")[0]) == "AMD": device_decoders[d] = load_counters
   # load device specific counters
   for fxn in device_decoders.values(): fxn(profile)
   # map events per device
@@ -509,7 +510,7 @@ def get_render(query:str) -> dict:
   if fmt == "asm":
     ret:dict = {"metadata":[]}
     if data.device.startswith("AMD") and data.lib is not None:
-      with soft_err(lambda err: ret.update(err)): ret.update(amdgpu_cfg(data.lib, amdgpu_targets[data.device]))
+      with soft_err(lambda err: ret.update(err)): ret.update(amdgpu_cfg(data.lib, device_props[data.device]["gfx_target_version"]))
       with soft_err(lambda err: ret["metadata"].append(err)): ret["metadata"].append(amd_readelf(data.lib))
     else: ret["src"] = get_stdout(lambda: (compiler:=Device[data.device].compiler).disassemble(compiler.compile(data.src)))
     return ret