stuff

.github/workflows/benchmark.yml

@@ -617,10 +617,6 @@ jobs:
       run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/e8bea2c78ffa92685ece511e9b554122aaf1a79d/selfdrive/modeld/models/supercombo.onnx
     - name: openpilot dmonitoring compile3 0.9.7
       run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: openpilot compile3 Space Lab policy + vision
-      run: |
-        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
-        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
     - name: benchmark MobileNetV2 on DSP
       run: |
         # generate quantized weights
@@ -653,10 +649,10 @@ jobs:
     steps:
     - name: Checkout Code
       uses: actions/checkout@v4
-    - name: Remove amd modules
-      run: ./extra/hcq/hcq_smi.py amd rmmod
-    - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Remove amdgpu
+      run: sudo rmmod amdgpu || true
+    - name: Cleanup running AM processes
+      run: python extra/amdpci/am_smi.py --pids --kill
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -720,9 +716,7 @@ jobs:
     - name: Checkout Code
       uses: actions/checkout@v4
     - name: Remove nv modules
-      run: ./extra/hcq/hcq_smi.py nv rmmod
-    - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py nv kill_pids
+      run: ./extra/nvpci/nv_smi.py rmmod
     - name: Symlink models and datasets
       run: |
         mkdir -p weights

.github/workflows/mlperf.yml

@@ -12,7 +12,7 @@ jobs:
   run_script_job:
     runs-on: [self-hosted, Linux, tinybox]
     if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 360
+    timeout-minutes: 240
 
     steps:
     - name: Checkout Code
@@ -27,4 +27,4 @@ jobs:
       run: |
         rm "~/.cache/tinygrad/cache_mlperf.db" || true
         BENCHMARK_LOG=mlpert_train_resnet LOGMLPERF=0 CACHEDB="~/.cache/tinygrad/cache_mlperf.db" examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-        rm "~/.cache/tinygrad/cache_mlperf.db"
+        rm "~/.cache/tinygrad/cache_mlperf.db"

docs/ramp.py

@@ -126,7 +126,7 @@ print(t_log_grad.uop)
 """
 void E_(float* restrict data0, float* restrict data1) {
   float val0 = *(data1+0);
-  *(data0+0) = (1/val0);
+  *(data0+0) = (0.6931471805599453f*(1/(val0*0.6931471805599453f)));
 }
 """
 # the derivative is close to 1/3

examples/mlperf/dataloader.py

@@ -1,6 +1,4 @@
-import functools
-import hashlib
-import os, random, pickle, queue, struct, math
+import os, random, pickle, queue
 from typing import List
 from pathlib import Path
 from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
@@ -8,7 +6,6 @@ from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu
 import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
-from tinygrad.nn.state import TensorIO
 
 ### ResNet
 
@@ -513,202 +510,6 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
       # happens with BENCHMARK set
       pass
 
-# llama3
-
-class BinIdxDataset:
-  def __init__(self, base_path:Path):
-    self.idx_t = Tensor(base_path.with_name(f"{base_path.name}.idx"))
-    self.idx = TensorIO(self.idx_t)
-
-    # parse idx file
-    magic = self.idx.read(9)
-    assert magic == b"MMIDIDX\x00\x00", "invalid index file format"
-    version, = struct.unpack("<Q", self.idx.read(8))
-    assert version == 1, "unsupported index version"
-    dtype_code, = struct.unpack("<B", self.idx.read(1))
-    self.dtype = {1:dtypes.uint8, 2:dtypes.int8, 3:dtypes.int16, 4:dtypes.int32, 5:dtypes.int64, 6:dtypes.float64, 7:dtypes.double, 8:dtypes.uint16}[dtype_code]
-    self.count, = struct.unpack("<Q", self.idx.read(8))
-    doc_count, = struct.unpack("<Q", self.idx.read(8))
-
-    start = self.idx.tell()
-    end = start + self.count * dtypes.int32.itemsize
-    self.sizes = self.idx_t[start:end].bitcast(dtypes.int32)
-
-    start = end
-    end = start + self.count * dtypes.int64.itemsize
-    self.pointers = self.idx_t[start:end].bitcast(dtypes.int64)
-
-    start = end
-    end = start + doc_count * dtypes.int64.itemsize
-    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64)
-
-    # bin file
-    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
-
-  def _index(self, idx) -> tuple[int, int]:
-    return self.pointers[idx].item(), self.sizes[idx].item()
-
-  def get(self, idx, offset:int=0, length:int|None=None):
-    ptr, size = self._index(idx)
-    if length is None: length = size - offset
-    ptr += offset * self.dtype.itemsize
-    return self.bin_t[ptr:ptr+length*self.dtype.itemsize].bitcast(self.dtype).to(None)
-
-# https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
-class GPTDataset:
-  def __init__(self, base_path:Path, samples:int, seqlen:int, seed:int, shuffle:bool):
-    self.samples, self.seqlen = samples, seqlen
-    self.shuffle = shuffle
-    self.rng = np.random.RandomState(seed)
-
-    self.indexed_dataset = BinIdxDataset(base_path)
-
-    # check for cache
-    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
-    cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
-    if cache_path.exists():
-      with open(cache_path, "rb") as f:
-        self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
-    else:
-      self.doc_idx = self._build_doc_idx()
-      self.sample_idx = self._build_sample_idx()
-      self.shuffle_idx = self._build_shuffle_idx()
-      # save cache
-      with open(cache_path, "wb") as f:
-        pickle.dump((self.doc_idx, self.sample_idx, self.shuffle_idx), f)
-
-  def __getitem__(self, idx):
-    if idx is None:
-      text = self._get(0)
-    else:
-      text = self._get(idx)
-
-    return text
-
-  def _get(self, idx):
-    idx = self.shuffle_idx[idx]
-
-    doc_idx_beg, doc_idx_beg_offset = self.sample_idx[idx]
-    doc_idx_end, doc_idx_end_offset = self.sample_idx[idx + 1]
-
-    doc_ids, sample_parts = [], []
-
-    if doc_idx_beg == doc_idx_end:
-      doc_ids.append(self.doc_idx[doc_idx_beg])
-
-      sample_parts.append(
-          self.indexed_dataset.get(
-            int(self.doc_idx[doc_idx_beg]), offset=int(doc_idx_beg_offset), length=int(doc_idx_end_offset - doc_idx_beg_offset + 1)))
-    else:
-      for i in range(doc_idx_beg, doc_idx_end + 1):
-        doc_ids.append(self.doc_idx[i])
-
-        offset = 0 if i > doc_idx_beg else doc_idx_beg_offset
-        length = None if i < doc_idx_end else int(doc_idx_end_offset + 1)
-        sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
-
-    # concat all parts
-    text = Tensor.cat(*sample_parts)
-
-    return text
-
-  @functools.cached_property
-  def tokens_per_epoch(self) -> int:
-    return sum(self.indexed_dataset.sizes.tolist())
-
-  @functools.cached_property
-  def num_epochs(self) -> int:
-    # we need enough epochs to cover the requested amount of tokens
-    num_epochs = 1
-    num_tokens = self.tokens_per_epoch
-    while num_tokens < self.samples * self.seqlen:
-      num_epochs += 1
-      num_tokens += self.tokens_per_epoch
-    return num_epochs
-
-  # https://github.com/NVIDIA/Megatron-LM/blob/94bd476bd840c2fd4c3ebfc7448c2af220f4832b/megatron/core/datasets/gpt_dataset.py#L558
-  def _build_doc_idx(self):
-    doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
-    doc_idx = doc_idx.reshape(-1)
-    doc_idx = doc_idx.astype(np.int32)
-    if self.shuffle: self.rng.shuffle(doc_idx)
-    return doc_idx
-
-  def _build_sample_idx(self):
-    sample_idx = np.empty((self.samples + 1, 2), dtype=np.int32)
-
-    sample_idx_idx, doc_idx_idx, doc_offset = 0, 0, 0
-    sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
-    sample_idx_idx += 1
-
-    for _ in tqdm(range(1, self.samples + 1)):
-      remaining_seqlen = self.seqlen + 1
-      while remaining_seqlen > 0:
-        doc_idx = int(self.doc_idx[doc_idx_idx])
-        doc_len = self.indexed_dataset.sizes[doc_idx].item() - doc_offset
-        remaining_seqlen -= doc_len
-        if remaining_seqlen <= 0:
-          doc_offset += remaining_seqlen + doc_len - 1
-          remaining_seqlen = 0
-        else:
-          if doc_idx_idx == len(self.doc_idx) - 1:
-            assert sample_idx_idx == self.samples
-            doc_idx = int(self.doc_idx[doc_idx_idx])
-            doc_offset = self.indexed_dataset.sizes[doc_idx].item() - 1
-            break
-          doc_idx_idx += 1
-          doc_offset = 0
-
-      sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
-      sample_idx_idx += 1
-
-    return sample_idx
-
-  def _build_shuffle_idx(self):
-    shuffle_idx = np.arange(self.samples, dtype=np.int32)
-    if self.shuffle: self.rng.shuffle(shuffle_idx)
-    return shuffle_idx
-
-class BlendedGPTDataset:
-  def __init__(self, paths:list[Path], weights:list[float], samples:int, seqlen:int, seed:int, shuffle:bool):
-    self.seed = seed
-
-    # normalize weights
-    total_weight = sum(weights)
-    self.weights = [w / total_weight for w in weights]
-
-    self.samples = samples
-    surplus = 0.005
-    samples_per_blend = [math.ceil(math.ceil(self.samples * w) * (1 + surplus)) for w in self.weights]
-
-    self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
-
-  def get(self, idx:int):
-    tokens = self.datasets[0][idx]
-    return tokens
-
-def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
-  if val:
-    dataset = BlendedGPTDataset([
-      base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
-    ], [
-      1.0
-    ], samples, seqlen, seed, False)
-  else:
-    dataset = BlendedGPTDataset([
-      base_dir / "c4-train.en_6_text_document",
-      base_dir / "c4-train.en_7_text_document",
-    ], [
-      1.0, 1.0
-    ], samples, seqlen, seed, True)
-
-  for b in range(math.ceil(samples / bs)):
-    batch = []
-    for i in range(bs):
-      tokens = dataset.get(b * bs + i)
-      batch.append(tokens)
-    yield Tensor.stack(batch, dim=0)
-
 if __name__ == "__main__":
   def load_unet3d(val):
     assert not val, "validation set is not supported due to different sizes on inputs"
@@ -737,18 +538,6 @@ if __name__ == "__main__":
       for x in batch_load_retinanet(dataset, val, base_dir):
         pbar.update(x[0].shape[0])
 
-  def load_llama3(val):
-    bs = 24
-    samples = 5760 if val else 1_200_000
-    seqlen = 512
-
-    max_, min_ = 0, math.inf
-    for tokens in tqdm(batch_load_llama3(bs, samples, seqlen, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=5760, val=bool(val)), total=samples//bs):
-      max_ = max(max_, tokens.shape[1])
-      min_ = min(min_, tokens.shape[1])
-    print(f"max seq length: {max_}")
-    print(f"min seq length: {min_}")
-
   load_fn_name = f"load_{getenv('MODEL', 'resnet')}"
   if load_fn_name in globals():
     globals()[load_fn_name](getenv("VAL", 1))

examples/mlperf/model_eval.py

@@ -1,4 +1,4 @@
-import time, math
+import time
 start = time.perf_counter()
 from pathlib import Path
 import numpy as np
@@ -241,34 +241,6 @@ def eval_mrcnn():
   evaluate_predictions_on_coco(bbox_output, iou_type='bbox')
   evaluate_predictions_on_coco(mask_output, iou_type='segm')
 
-def eval_llama3():
-  from extra.models.llama import Transformer
-  from examples.llama3 import MODEL_PARAMS
-  from tinygrad.helpers import tqdm
-
-  bs = 4
-  sequence_length = 512
-
-  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
-
-  @TinyJit
-  def eval_step(model, tokens):
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten()
-
-  from examples.mlperf.dataloader import batch_load_llama3
-  iter = batch_load_llama3(bs, 5760, sequence_length, Path(getenv("BASEDIR", "/raid/datasets/c4/")), True)
-
-  losses = []
-  for tokens in tqdm(iter, total=5760//bs):
-    GlobalCounters.reset()
-    losses += eval_step(model, tokens).tolist()
-    tqdm.write(f"loss: {np.mean(losses)}")
-
-  log_perplexity = Tensor(losses).mean()
-  print(f"Log Perplexity: {log_perplexity.item()}")
-
 if __name__ == "__main__":
   # inference only
   Tensor.training = False

examples/mlperf/model_train.py

@@ -1290,16 +1290,9 @@ def train_llama3():
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
 
   config = {}
-  BS                 = config["BS"]                     = getenv("BS", 16)
+  BS                 = config["BS"]                     = getenv("BS", 4)
   grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
   GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
-  SEED               = config["SEED"]                   = getenv("SEED", 5760)
-  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
-  TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
-  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000)
-
-  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
-  # trains to 7
 
   opt_adamw_beta_1 = 0.9
   opt_adamw_beta_2 = 0.95
@@ -1307,6 +1300,7 @@ def train_llama3():
   opt_adamw_weight_decay = 0.1
 
   opt_gradient_clip_norm = 1.0
+  sequence_length = 8192
   opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
   opt_learning_rate_decay_steps = getenv("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
   opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
@@ -1314,33 +1308,7 @@ def train_llama3():
 
   # TODO: confirm weights are in bf16
   # vocab_size from the mixtral tokenizer
-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
-  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
-
-  if (DP := getenv("DP", 1)) > 1:
-    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-    for v in get_parameters(model):
-      v.shard_(device, axis=None)
-
-  # TODO: MP
-  # if (GPUS := getenv("GPUS", 1)) > 1:
-  #   device = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
-  #   for k,v in get_state_dict(model).items():
-  #     if 'scale' in k: v.shard_(device, axis=None)  # from quantized
-  #     # elif '.attention.wq' in k: v.shard_(device, axis=0)
-  #     # elif '.attention.wk' in k: v.shard_(device, axis=0)
-  #     # elif '.attention.wv' in k: v.shard_(device, axis=0)
-  #     # elif '.attention.wo' in k: v.shard_(device, axis=1)
-  #     # elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
-  #     # elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
-  #     # elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
-  #     # elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
-  #     elif 'output.weight' in k: v.shard_(device, axis=0)  # 243.32
-  #     else:
-  #       # print(k)
-  #       # attention_norm, ffn_norm, norm
-  #       v.shard_(device, axis=None)
+  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
 
   optim = AdamW(get_parameters(model), lr=0.0,
                 b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
@@ -1348,17 +1316,12 @@ def train_llama3():
 
   @TinyJit
   @Tensor.train()
-  def train_step(model, tokens:Tensor, grad_acc:int):
+  def train_step(model, x, y):
     optim.zero_grad()
-    # grad acc
-    for batch in tokens.split(tokens.shape[0]//grad_acc):
-      if (DP := getenv("DP", 1)) > 1:
-        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-        batch = batch.shard(device, 0)
-      logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
-      loss = logits.sparse_categorical_crossentropy(batch[:, 1:])
-      loss.backward()
-      Tensor.realize(*[p.grad for p in optim.params])
+    logits:Tensor = model(x, start_pos=0, temperature=math.nan)
+    loss = logits.cross_entropy(y)
+    loss.backward()
+
     # L2 norm grad clip
     # https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
     # https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
@@ -1377,32 +1340,19 @@ def train_llama3():
     loss.realize(lr)
     return loss, lr
 
-  if getenv("FAKEDATA", 0):
-    def fake_data():
-      for _ in range(SAMPLES // GBS):
-        yield Tensor.randint(GBS, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
-    iter = fake_data()
-  else:
-    from examples.mlperf.dataloader import batch_load_llama3
-    iter = batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
+  # overfitting this example should give cross_entropy log(BS)
+  fake_input = Tensor([list(range(getenv("SEQLEN", 10)))], dtype="int16").expand(BS, -1)
+  fake_label = Tensor(list(range(BS)), dtype="int16")
 
-  i = 0
-  for tokens in tqdm(iter, total=SAMPLES//GBS):
-    t = time.perf_counter()
+  for _ in range(100):
     GlobalCounters.reset()
-    loss, lr = train_step(model, tokens, grad_acc)
-    # above as tqdm.write f-string
-    tqdm.write(f"{loss.item():.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
-    if (fname:=getenv("LOSS_FILE", "")):
-      with open(fname, "a") as f:
-        f.write(f"{i} {loss.item():.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
-
-    if getenv("CKPT") and (i % 200 == 0 or i == 10):
-      tqdm.write("saving checkpoint")
-      if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-      fn = f"{ckpt_dir}/{i}.safe"
-      safe_save(get_state_dict(model), fn)
-    i += 1
+    loss, lr = train_step(model, fake_input, fake_label)
+    # BS=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
+    # uses 43% ~= 83GB
+    # 8B bf16 = 16GB. model + grad + optim m and v = 64GB
+    # TODO: this OOM
+    # BS=1 SEQLEN=4000 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
+    print(loss.item(), lr.item(), f"{GlobalCounters.global_mem//10**9=}")
 
 if __name__ == "__main__":
   multiprocessing.set_start_method('spawn')

extra/hcq/hcq_smi.py

@@ -1,122 +0,0 @@
-#!/usr/bin/env python3
-
-from tinygrad.runtime.support.system import System
-import argparse, glob, os, re, time, subprocess, sys
-
-def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
-  target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
-
-  devs = []
-  for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
-    dev_id = dev[8:-5]
-    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
-  return devs
-
-def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
-def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
-
-def cmd_remove_module(args):
-  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
-  to_unload = [m for m in modules if _is_module_loaded(m)]
-  if not to_unload: print("Kernel modules are not loaded")
-  else:
-    print("Removing kernel modules:", ", ".join(to_unload))
-    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
-    except subprocess.CalledProcessError as e:
-      print("Failed to unload all modules — they may be in use.", file=sys.stderr)
-      sys.exit(e.returncode)
-
-def cmd_insert_module(args):
-  cmd_remove_module(args)
-  cmd_reset_devices(args)
-
-  module = "nvidia" if args.backend == "nv" else "amdgpu"
-  if _is_module_loaded(module):
-    print(f"{module} kernel module already loaded")
-    return
-
-  print(f"Inserting kernel module: {module}")
-  if args.backend == "nv":
-    subprocess.run(["nvidia-smi"], check=True)
-  elif args.backend == "amd":
-    subprocess.run(["sudo", "modprobe", "amdgpu"], check=True)
-
-def cmd_reset_devices(args):
-  devs = scan_devs_based_on_lock({"amd":"am", "nv":"nv"}[args.backend], args)
-
-  for dev in devs:
-    print(f"Resetting device {dev}")
-    if args.backend != "amd": _do_reset_device(dev)
-    time.sleep(0.2)
-
-def cmd_show_pids(args):
-  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
-
-  for dev in devs:
-    try:
-      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
-      print(f"{dev}: {pid}")
-    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
-
-def cmd_kill_pids(args):
-  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
-
-  for dev in devs:
-    try:
-      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
-      print(f"{dev}: {pid}")
-    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
-
-def cmd_kill_pids(args):
-  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
-
-  for dev in devs:
-    for i in range(128):
-      if i > 0: time.sleep(0.2)
-
-      try:
-        try: pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
-        except subprocess.CalledProcessError: break
-
-        print(f"Killing process {pid} (which uses {dev})")
-        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
-      except subprocess.CalledProcessError as e:
-        print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)
-
-def add_common_commands(parent_subparsers):
-  p_insmod = parent_subparsers.add_parser("insmod", help="Insert a kernel module")
-  p_insmod.set_defaults(func=cmd_insert_module)
-
-  p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
-  p_rmmod.set_defaults(func=cmd_remove_module)
-
-  p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
-  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device to reset")
-  p_reset.set_defaults(func=cmd_reset_devices)
-
-  p_reset = parent_subparsers.add_parser("pids", help="Show pids of processes using the device")
-  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
-  p_reset.set_defaults(func=cmd_show_pids)
-
-  p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
-  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
-  p_reset.set_defaults(func=cmd_kill_pids)
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  backend_subparsers = parser.add_subparsers(dest="backend", required=True, metavar="{nv,amd}", help="Hardware backend to target")
-
-  nv_parser = backend_subparsers.add_parser("nv", help="NVIDIA GPUs")
-  nv_commands = nv_parser.add_subparsers(dest="command", required=True)
-  add_common_commands(nv_commands)
-
-  amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
-  amd_commands = amd_parser.add_subparsers(dest="command", required=True)
-  add_common_commands(amd_commands)
-
-  args = parser.parse_args()
-  if args.command is None:
-    parser.print_help(sys.stderr)
-    sys.exit(1)
-
-  args.func(args)

extra/models/llama.py

@@ -99,9 +99,7 @@ class FeedForward:
     self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit
 
   def __call__(self, x:Tensor) -> Tensor:
-    w1 = self.w1(x).silu()
-    w3 = self.w3(x.contiguous_backward())  # this fixes a strange fusion that makes tensor cores miss
-    return self.w2(w1 * w3)
+    return self.w2(self.w1(x).silu() * self.w3(x)) # SwiGLU [arxiv/2002.05202, eq (5)]
 
 class TransformerBlock:
   def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, max_context:int, linear=nn.Linear,
@@ -113,7 +111,7 @@ class TransformerBlock:
 
   def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]):
     h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
-    return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward()
+    return (h + self.feed_forward(self.ffn_norm(h))).contiguous()
 
 # standard openai sampling
 def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):
@@ -187,10 +185,10 @@ class Transformer:
 
     mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
     for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
-    logits = self.output(self.norm(h)).float()
+    logits = self.output(self.norm(h)).float()[:, -1, :]
     if math.isnan(temperature): return logits
 
-    return sample(logits[:, -1, :].flatten(), temperature, top_k, top_p, alpha_f, alpha_p)
+    return sample(logits.flatten(), temperature, top_k, top_p, alpha_f, alpha_p)
 
   def __call__(self, tokens:Tensor, start_pos:int, temperature:float=0.0, top_k:int=0, top_p:float=0.8, alpha_f:float=0.0, alpha_p:float=0.0):
     # TODO: better way to handle the first call v.s. the rest?

extra/nvpci/gpu_reset.sh

@@ -0,0 +1,2 @@
+GPU="$1"
+echo 1 | sudo tee /sys/bus/pci/devices/$GPU/reset 2>/dev/null

extra/nvpci/nv_smi.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+
+from tinygrad.runtime.support.system import System
+import argparse, glob, os, re, time, subprocess, sys
+
+def scan_devs_based_on_lock(prefix:str) -> list[str]:
+  devs = []
+  for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
+    dev_id = dev[8:-5]
+    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}"): devs.append(dev_id)
+  return devs
+
+def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
+def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
+
+def cmd_remove_module(args):
+  to_unload = [m for m in ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if _is_module_loaded(m)]
+  if not to_unload:
+    print("NVIDIA kernel modules are not loaded")
+  else:
+    print("Removing NVIDIA kernel modules:", ", ".join(to_unload))
+    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
+    except subprocess.CalledProcessError as e:
+      print("Failed to unload all modules — they may be in use.", file=sys.stderr)
+      sys.exit(e.returncode)
+
+def cmd_insert_module(args):
+  cmd_remove_module(args)
+  cmd_reset_devices(args)
+
+  if not os.path.exists("/sys/module/nvidia"):
+    print("Inserting nvidia kernel module")
+    subprocess.run(["nvidia-smi"], check=True)
+  else: print("Nvidia kernel module already loaded")
+
+def cmd_reset_devices(args):
+  devs = scan_devs_based_on_lock("nv")
+  dev_to_reset = args.pci_bus if 'pci_bus' in args.__dir__() else ""
+
+  for dev in devs:
+    if dev.startswith(dev_to_reset):
+      print(f"Resetting device {dev}")
+      _do_reset_device(dev)
+      time.sleep(0.2)
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  subparsers = parser.add_subparsers(required=True, dest="cmd")
+
+  parser_insmod = subparsers.add_parser('insmod', help='Insert a nvidia kernel module')
+  parser_insmod.set_defaults(func=cmd_insert_module)
+
+  parser_rmmod = subparsers.add_parser('rmmod', help='Remove a nvidia kernel module')
+  parser_rmmod.set_defaults(func=cmd_remove_module)
+
+  parser_reset = subparsers.add_parser('reset', help='Reset a nvidia device')
+  parser_reset.add_argument('--pci_bus', type=str, default="", help='PCI bus ID of the device to reset')
+  parser_reset.set_defaults(func=cmd_reset_devices)
+
+  args = parser.parse_args()
+  if args.cmd is None:
+    parser.print_help(sys.stderr)
+    sys.exit(1)
+
+  args.func(args)

extra/torch_backend/backend.py

@@ -128,12 +128,6 @@ def _linalg_eigh(self, UPLO: str = 'U'):
   w, v = torch.linalg.eigh(self.cpu(), UPLO=UPLO)
   return w.tiny(), v.tiny()
 
-@torch.library.impl("aten::_linalg_det", "privateuseone")
-# TODO: move to tinygrad
-def _linalg_det(self: torch.Tensor):
-  result = aten._linalg_det(self.cpu())
-  return result[0].tiny(), result[1].tiny(), result[2].tiny()
-
 def upsample_backward(grad_out, output_size, input_size, *args, f=None): return f(grad_out.cpu(), output_size, input_size, *args).tiny()
 
 for i in [

extra/torch_backend/test.py

@@ -198,11 +198,6 @@ class TestTorchBackend(unittest.TestCase):
     recon = (v @ torch.diag(w) @ v.T).cpu().numpy()
     np.testing.assert_allclose(recon, a.cpu().numpy(), atol=1e-6)
 
-  def test_linalg_det(self):
-    a = torch.diag(torch.tensor([1,2,3,4,5], dtype = torch.float32, device=device))
-    b = torch.linalg.det(a)
-    np.testing.assert_equal(b.cpu().numpy(), 120.0)
-
   def test_scalar_assign(self):
     a = torch.tensor([1, 2, 3], device=device)
     a[1] = 4

test/external/external_llm_eval.py

@@ -10,7 +10,7 @@ if __name__ == "__main__":
 
   model, kv = Transformer.from_gguf(Tensor.from_url(models["1B"]), max_context=4096)
 
-  tok = SimpleTokenizer.from_gguf_kv(kv)
+  tok = SimpleTokenizer(kv["tokenizer.ggml.tokens"])
   bos_id: int = kv['tokenizer.ggml.bos_token_id']
   eos_id: int = kv['tokenizer.ggml.eos_token_id']
 

test/external/external_test_simple_tokenizer.py

@@ -1,19 +1,17 @@
 from transformers import AutoTokenizer
 from datasets import load_dataset
-from tinygrad.apps.llm import SimpleTokenizer, gpt2_decode_vocab, get_llama_re
-from tinygrad.helpers import tqdm, getenv, partition
+from tinygrad.apps.llm import SimpleTokenizer
+from tinygrad.helpers import tqdm, getenv
 
 # use ALLOW_FAILED=-1 to go over the entire dataset without printing.
 if __name__ == "__main__":
   base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
-  special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()),
-    lambda e: e[1] in base_tokenizer.all_special_ids)
+  vocab_words = [ word for word, _ in sorted(base_tokenizer.get_vocab().items(), key=lambda t: t[1]) ]
   inv_vocab = { tid: word for word, tid in base_tokenizer.get_vocab().items() }
-  simple_tokenizer = SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens))
+  simple_tokenizer = SimpleTokenizer(vocab_words)
 
   color_codes = [ 91, 92, 94, 93, 95 ]
-  def color_tokens(tids):
-    return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m"
+  def color_tokens(tids): return "".join(f"\033[{color_codes[i%len(color_codes)]}m{inv_vocab[t]}" for i, t in enumerate(tids)) + "\033[0m"
 
   ds = load_dataset("OpenAssistant/oasst1")
   allow_failed = getenv("ALLOW_FAILED", 10)

test/external/process_replay/process_replay.py

@@ -74,7 +74,6 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
       warnings.warn(f"detected changes in over {MAX_DIFF_PCT}%. skipping further diff generation.", ProcessReplayWarning)
       early_stop.set()
       break
-    name, loc = "", ""
     try:
       name, args, kwargs, ctx_vals, loc, ret = pickle.loads(row[0])
       ctx_vars = {k:v.value for k,v in ctx_vals.items() if k != "DEBUG" and (var:=ContextVar._cache.get(k)) is not None and var.value != v.value}
@@ -91,7 +90,7 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
         warnings.warn("PROCESS REPLAY DETECTED CHANGE", ProcessReplayWarning)
     except Exception as e:
       changed += 1
-      warnings.warn(f"{name=} {loc=} {e=}", ProcessReplayWarning)
+      warnings.warn(e, ProcessReplayWarning)
   conn.commit()
   cur.close()
 
@@ -124,5 +123,5 @@ if __name__ == "__main__":
   logging.info(f"running process replay with {ASSERT_DIFF=}")
   try: _pmap(replayers)
   except Exception as e:
-    logging.info(f"process replay err: {e}")
+    logging.info("process replay err", e)
     exit(int(ASSERT_DIFF))

test/test_dtype.py

@@ -4,7 +4,7 @@ import torch
 from typing import Any, List
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, DEBUG, CI
-from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8
+from tinygrad.dtype import DType, DTYPES_DICT, ImageDType, PtrDType, least_upper_dtype, to_dtype, fp8_to_float, float_to_fp8
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.tensor import _to_np_dtype
 from hypothesis import assume, given, settings, strategies as strat
@@ -384,6 +384,30 @@ class TestPtrDType(unittest.TestCase):
     self.assertEqual(dt.v, 4)
     self.assertEqual(dt.count, 4)
 
+class TestImageDType(unittest.TestCase):
+  def test_image_scalar(self):
+    assert dtypes.imagef((10,10)).base.scalar() == dtypes.float32
+    assert dtypes.imageh((10,10)).base.scalar() == dtypes.float32
+  def test_image_vec(self):
+    assert dtypes.imagef((10,10)).base.vec(4) == dtypes.float32.vec(4)
+    assert dtypes.imageh((10,10)).base.vec(4) == dtypes.float32.vec(4)
+
+class TestEqStrDType(unittest.TestCase):
+  def test_image_ne(self):
+    if ImageDType is None: raise unittest.SkipTest("no ImageDType support")
+    assert dtypes.float == dtypes.float32, "float doesn't match?"
+    assert dtypes.imagef((1,2,4)) != dtypes.imageh((1,2,4)), "different image dtype doesn't match"
+    assert dtypes.imageh((1,2,4)) != dtypes.imageh((1,4,2)), "different shape doesn't match"
+    assert dtypes.imageh((1,2,4)) == dtypes.imageh((1,2,4)), "same shape matches"
+    assert isinstance(dtypes.imageh((1,2,4)), ImageDType)
+  def test_ptr_eq(self):
+    assert dtypes.float32.ptr() == dtypes.float32.ptr()
+    assert not (dtypes.float32.ptr() != dtypes.float32.ptr())
+  def test_strs(self):
+    if PtrDType is None: raise unittest.SkipTest("no PtrDType support")
+    self.assertEqual(str(dtypes.imagef((1,2,4))), "dtypes.imagef((1, 2, 4))")
+    self.assertEqual(str(dtypes.float32.ptr(16)), "dtypes.float.ptr(16)")
+
 class TestImplicitFunctionTypeChange(unittest.TestCase):
   def test_functions(self):
     result = []
@@ -414,6 +438,19 @@ class TestDtypeUsage(unittest.TestCase):
         t = Tensor([[1, 2], [3, 4]], dtype=d)
         (t*t).max().item()
 
+class TestToDtype(unittest.TestCase):
+  def test_dtype_to_dtype(self):
+    dtype = dtypes.int32
+    res = to_dtype(dtype)
+    self.assertIsInstance(res, DType)
+    self.assertEqual(res, dtypes.int32)
+
+  def test_str_to_dtype(self):
+    dtype = "int32"
+    res = to_dtype(dtype)
+    self.assertIsInstance(res, DType)
+    self.assertEqual(res, dtypes.int32)
+
 @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
 class TestOpsBFloat16(unittest.TestCase):
   def test_cast(self):

test/test_graph.py

@@ -107,9 +107,8 @@ class TestGraph(unittest.TestCase):
     helper_test_graphs(Device[d0].graph, graphs)
 
   def skip_if_not_multigraph(self):
-    graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g
+    graph = g.func if isinstance(g:=Device[Device.DEFAULT].graph, functools.partial) else g
     if not issubclass(graph, MultiGraphRunner): self.skipTest("graph is not supported (not MultiGraphRunner)")
-    if not hasattr(d.allocator, '_transfer'): self.skipTest("device is not supported (no transfers)")
 
   def test_order_copy_writed(self):
     self.skip_if_not_multigraph()

test/test_linalg.py

@@ -16,7 +16,7 @@ def reconstruction_helper(A:List[Tensor],B:Tensor, tolerance=1.0e-5):
 class TestLinAlg(unittest.TestCase):
 
   def test_svd_general(self):
-    sizes = [(2,2),(5,3),(3,5),(3,4,4),(2,2,2,2,3)]
+    sizes = [(2,2),(5,3),(3,5),(2,2,2,2,3)]
     for size in sizes:
       a = Tensor.randn(size).realize()
       U,S,V = Tensor.svd(a)

test/test_multitensor.py

@@ -2,7 +2,7 @@ import unittest, functools, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.helpers import CI, getenv, prod, Context
+from tinygrad.helpers import CI, getenv, prod, Context, OSX
 from tinygrad.nn.state import get_parameters, get_state_dict
 from tinygrad.engine.realize import lower_schedule, BufferCopy, CompiledRunner, run_schedule
 import numpy as np
@@ -374,6 +374,7 @@ class TestMultiTensor(unittest.TestCase):
 
   # NOTE: this is failing on LLVM CI, no idea why. Works locally.
   @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU"), "slow, and flaky on LLVM/CPU")
+  @unittest.skipIf(REAL_DEV == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
   def test_data_parallel_resnet(self):
     from extra.models.resnet import ResNet18
 
@@ -410,6 +411,7 @@ class TestMultiTensor(unittest.TestCase):
     np.testing.assert_allclose(grad, shard_grad, atol=1e-5, rtol=1e-5)
 
   @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "LLVM", "CPU"), "slow, and flaky on LLVM/CPU")
+  @unittest.skipIf(REAL_DEV == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
   def test_data_parallel_resnet_train_step(self):
     from extra.models.resnet import ResNet18
     fake_image = Tensor.rand((2, 3, 224//8, 224//8))
@@ -936,6 +938,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
     np.testing.assert_allclose(output.numpy(), expected)
 
 @unittest.skipIf(not_support_multi_device(), "no multi")
+@unittest.skipIf(REAL_DEV == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
 class TestBatchNorm(unittest.TestCase):
   def test_unsynced_backprop_conv_bn(self):
     with Tensor.train():
@@ -963,6 +966,7 @@ class TestBatchNorm(unittest.TestCase):
       optim.step()
       out.numpy()
 
+  @unittest.skipIf(REAL_DEV == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
   def test_unsynced_backprop_standalone_bn(self):
     from extra.lr_scheduler import OneCycleLR
     GPUS = (d1, d2)

test/test_nn.py

@@ -4,7 +4,7 @@ import numpy as np
 import torch
 from tinygrad import Tensor, Device, TinyJit
 from tinygrad.uop.ops import Ops
-from tinygrad.helpers import GlobalCounters, CI, Context
+from tinygrad.helpers import GlobalCounters, CI, Context, OSX
 from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding
 from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell
 from tinygrad.nn.state import load_state_dict
@@ -284,6 +284,7 @@ class TestNN(unittest.TestCase):
     torch_z = torch_layer(torch_x)
     np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
 
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
   def test_groupnorm(self):
     BS, H, W, C, G = 20, 10, 10, 6, 3
 
@@ -310,6 +311,7 @@ class TestNN(unittest.TestCase):
       np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
       np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
 
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
   def test_layernorm(self):
     N, C, H, W = 20, 5, 10, 10
 
@@ -336,6 +338,7 @@ class TestNN(unittest.TestCase):
       np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
       np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
 
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
   def test_layernorm_2d(self):
     N, C, H, W = 20, 5, 10, 10
 
@@ -362,6 +365,7 @@ class TestNN(unittest.TestCase):
       np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
       np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
 
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
   def test_instancenorm_2d(self):
     N, C, H, W = 20, 10, 10, 10
 
@@ -388,6 +392,7 @@ class TestNN(unittest.TestCase):
       np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
       np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
 
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
   def test_instancenorm_3d(self):
     N, C, D, H, W = 20, 10, 10, 10, 10
 
@@ -414,6 +419,7 @@ class TestNN(unittest.TestCase):
       np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=2e-3, rtol=1e-3)
       np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
 
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
   def test_rmsnorm(self):
     class TorchRMSNorm(torch.nn.Module):
       # https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L34C1-L77C36

test/test_ops.py

@@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, AMD_LLVM
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, OSX, AMD_LLVM
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -2682,6 +2682,7 @@ class TestOps(unittest.TestCase):
     i, j, k, o, p = [Tensor(tor.detach().cpu().numpy().astype(np.int32), requires_grad=False) for tor in [a,b,c,d,e]]
     return a,b,c,d,e,i,j,k,o,p
 
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU can only run kernels with up to 10 buffers")
   def test_slice_fancy_indexing_no_dim_collapse(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     # no dim collapse from int or dim injection from None
@@ -2733,6 +2734,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,3)], lambda x: x[torch.tensor([[0,1,-1],[-1,-2,0]]), torch.tensor([2,1,-1])],
                             lambda x: x[Tensor([[0,1,-1],[-1,-2,0]]), Tensor([2,1,-1])])
 
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU can only run kernels with up to 10 buffers")
   def test_slice_fancy_indexing_list_indices(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[[[0]]], lambda x: x[[[0]]])
@@ -2752,6 +2754,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,((2,),(1,),(0,)),c,(2,1,0)], lambda x: x[i,((2,),(1,),(0,)),k,(2,1,0)])
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[1,(2,1,0),None,c,(2,1,0),e], lambda x: x[1,(2,1,0),None,k,(2,1,0),p])
 
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
   def test_slice_fancy_indexing_list_with_tensors(self):
     a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
     helper_test_op([(2,5,6,5,3,4)], lambda x: x[[a]], lambda x: x[[i]])

test/test_sample.py

@@ -3,7 +3,6 @@ import numpy as np
 from tinygrad import Tensor, Variable, Device
 from tinygrad.helpers import OSX
 
-# TODO: still fails with MAX_KERNEL_BUFFERS
 @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
 class TestSample(unittest.TestCase):
   def test_sample(self):

test/test_search.py

@@ -34,29 +34,25 @@ class TestBEAM(unittest.TestCase):
     capturing.clear()
     self.assertNotEqual(k_beam_0[-1].prg.p.src, k_beam_1[-1].prg.p.src)
 
-  def test_get_kernel_actions_dedup(self):
+  def test_get_kernel_actions(self):
     from test.test_linearizer import helper_realized_ast
-    from tinygrad.opt.search import get_kernel_actions
-    a = Tensor.empty(4, 3)
-    b = Tensor.empty(3)
+    a = Tensor.rand(4, 3)
+    b = Tensor.rand(3)
     realized_ast, _ = helper_realized_ast(a @ b)
-    candidates = [
-      Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=4),
-      Opt(op=OptOps.LOCAL, axis=0, arg=0), Opt(op=OptOps.LOCAL, axis=0, arg=4),
-      Opt(op=OptOps.UNROLL, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=3),
-      Opt(op=OptOps.GROUP, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=3),
-      Opt(op=OptOps.GROUPTOP, axis=0, arg=0), Opt(op=OptOps.GROUPTOP, axis=0, arg=3),
-    ]
-    lins = get_kernel_actions(Kernel(realized_ast), include_0=False, candidates=candidates).values()
+    from tinygrad.opt.search import get_kernel_actions
+    lins = get_kernel_actions(Kernel(realized_ast), False).values()
 
     # ensure amt=0 are not duplicated
-    assert all(len(x.applied_opts) == 1 for x in lins)
-    kernel_actions = [x.applied_opts[0] for x in lins]
-    assert Opt(OptOps.UPCAST, axis=0, arg=4) not in kernel_actions, "did not de-dup UPCAST"
-    assert Opt(OptOps.LOCAL, axis=0, arg=4) not in kernel_actions, "did not de-dup LOCAL"
-    assert Opt(OptOps.UNROLL, axis=0, arg=3) not in kernel_actions, "did not de-dup UNROLL"
-    assert Opt(OptOps.GROUP, axis=0, arg=3) not in kernel_actions, "did not de-dup GROUP"
-    assert Opt(OptOps.GROUPTOP, axis=0, arg=3) not in kernel_actions, "did not de-dup GROUPTOP"
+    if Opt(OptOps.UPCAST, 0, 0) in actions:
+      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.UPCAST, axis=0, arg=4)]) == 0, "did not de-dup UPCAST"
+    if Opt(OptOps.LOCAL, 0, 0) in actions:
+      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.LOCAL, axis=0, arg=4)]) == 0, "did not de-dup LOCAL"
+    if Opt(OptOps.UNROLL, 0, 0) in actions:
+      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.UNROLL, axis=0, arg=3)]) == 0, "did not de-dup UNROLL"
+    if Opt(OptOps.GROUP, 0, 0) in actions:
+      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.GROUP, axis=0, arg=3)]) == 0, "did not de-dup GROUP"
+    if Opt(OptOps.GROUPTOP, 0, 0) in actions:
+      assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.GROUPTOP, axis=0, arg=3)]) == 0, "did not de-dup GROUPTOP"
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
   def test_search_over_shape(self):

test/test_setitem.py

@@ -158,13 +158,6 @@ class TestSetitem(unittest.TestCase):
     t[:-1] = t[1:]
     self.assertEqual(t.tolist(), [[2.0], [1.0], [1.0]])
 
-  def test_setitem_big(self):
-    idx_size, val = 256, 4
-    t = Tensor.arange(0, idx_size+1)
-    idx = Tensor.arange(0, idx_size)
-    t[idx] = val
-    self.assertEqual(t.tolist(), [val]*idx_size+[idx_size])
-
 class TestWithGrad(unittest.TestCase):
   def test_no_requires_grad_works(self):
     z = Tensor.rand(8, 8)

test/test_softmax_fusion.py

@@ -121,7 +121,7 @@ class TestSoftmaxFusion(unittest.TestCase):
       out = (inp / div).reshape(32, 10)
       out.realize()
 
-    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
+    np.testing.assert_allclose(sout.numpy(), out.numpy())
 
   def test_softmax(self):
     # this is the softmax from scaled_dot_product_attention

test/unit/test_dtype.py

@@ -1,57 +0,0 @@
-import unittest
-from tinygrad.tensor import Tensor
-from tinygrad.dtype import dtypes, DType, ImageDType, PtrDType, to_dtype
-
-class TestImageDType(unittest.TestCase):
-  def test_image_scalar(self):
-    assert dtypes.imagef((10,10)).base.scalar() == dtypes.float32
-    assert dtypes.imageh((10,10)).base.scalar() == dtypes.float32
-  def test_image_vec(self):
-    assert dtypes.imagef((10,10)).base.vec(4) == dtypes.float32.vec(4)
-    assert dtypes.imageh((10,10)).base.vec(4) == dtypes.float32.vec(4)
-
-class TestEqStrDType(unittest.TestCase):
-  def test_image_ne(self):
-    if ImageDType is None: raise unittest.SkipTest("no ImageDType support")
-    assert dtypes.float == dtypes.float32, "float doesn't match?"
-    assert dtypes.imagef((1,2,4)) != dtypes.imageh((1,2,4)), "different image dtype doesn't match"
-    assert dtypes.imageh((1,2,4)) != dtypes.imageh((1,4,2)), "different shape doesn't match"
-    assert dtypes.imageh((1,2,4)) == dtypes.imageh((1,2,4)), "same shape matches"
-    assert isinstance(dtypes.imageh((1,2,4)), ImageDType)
-  def test_ptr_eq(self):
-    assert dtypes.float32.ptr() == dtypes.float32.ptr()
-    assert not (dtypes.float32.ptr() != dtypes.float32.ptr())
-  def test_strs(self):
-    if PtrDType is None: raise unittest.SkipTest("no PtrDType support")
-    self.assertEqual(str(dtypes.imagef((1,2,4))), "dtypes.imagef((1, 2, 4))")
-    self.assertEqual(str(dtypes.float32.ptr(16)), "dtypes.float.ptr(16)")
-
-class TestToDtype(unittest.TestCase):
-  def test_dtype_to_dtype(self):
-    dtype = dtypes.int32
-    res = to_dtype(dtype)
-    self.assertIsInstance(res, DType)
-    self.assertEqual(res, dtypes.int32)
-
-  def test_str_to_dtype(self):
-    dtype = "int32"
-    res = to_dtype(dtype)
-    self.assertIsInstance(res, DType)
-    self.assertEqual(res, dtypes.int32)
-
-class TestCastConvenienceMethod(unittest.TestCase):
-  def test_method(self):
-    for input_dtype in (dtypes.float, dtypes.int):
-      t = Tensor([1, 2], dtype=input_dtype)
-      self.assertEqual(t.dtype, input_dtype)
-      self.assertEqual(t.bool().dtype, dtypes.bool)
-      self.assertEqual(t.short().dtype, dtypes.short)
-      self.assertEqual(t.int().dtype, dtypes.int)
-      self.assertEqual(t.long().dtype, dtypes.long)
-      self.assertEqual(t.half().dtype, dtypes.half)
-      self.assertEqual(t.bfloat16().dtype, dtypes.bfloat16)
-      self.assertEqual(t.float().dtype, dtypes.float)
-      self.assertEqual(t.double().dtype, dtypes.double)
-
-if __name__ == "__main__":
-  unittest.main()

test/unit/test_llm_tokenizer.py

@@ -1,57 +0,0 @@
-import unittest, base64, functools
-from tinygrad.apps.llm import SimpleTokenizer, get_llama_re
-from tinygrad.helpers import fetch
-
-class TestLLMTokenizer(unittest.TestCase):
-  @functools.cached_property
-  def basic_tok(self): return SimpleTokenizer(".*", { b"a": 0, b"b": 1, b"c": 2, b"ab": 3, b"bc": 4 }, { "<x>": 5, "<y>": 6, "<z>": 7 })
-
-  @functools.cached_property
-  def llama_tok(self):
-    # from https://github.com/tinygrad/tinygrad/blob/e0106b6b257ebc003eb3694144e3e198f7d8cc37/examples/llama3.py#L14
-    model_file = fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model")
-    with open(model_file, "rt") as fd:
-      str_vocab = [ line.split(maxsplit=1) for line in fd.read().splitlines() if line ]
-      normal_tokens = { base64.b64decode(stok): int(srank) for stok, srank in str_vocab }
-
-    special_tokens = [
-      "<|begin_of_text|>",
-      "<|end_of_text|>",
-      "<|reserved_special_token_0|>",
-      "<|reserved_special_token_1|>",
-      "<|reserved_special_token_2|>",
-      "<|reserved_special_token_3|>",
-      "<|start_header_id|>",
-      "<|end_header_id|>",
-      "<|reserved_special_token_4|>",
-      "<|eot_id|>",
-    ] + [ f"<|reserved_special_token_{i}|>" for i in range(5, 256 - 5) ]
-    return SimpleTokenizer(get_llama_re(), normal_tokens, { token: len(normal_tokens) + i for i, token in enumerate(special_tokens) })
-
-  def _test_coding(self, tok: SimpleTokenizer, text: str, expected_tokens: list[int]):
-    self.assertEqual(tok.encode(text), expected_tokens)
-    self.assertEqual(tok.decode(expected_tokens), text)
-
-  def test_abc(self): self._test_coding(self.basic_tok, "abc", [ 3, 2 ])
-  def test_abbc(self): self._test_coding(self.basic_tok, "abbc", [ 3, 4 ])
-  def test_aabbbcc(self): self._test_coding(self.basic_tok, "aabbbcc", [ 0, 3, 1, 4, 2 ])
-  def test_specials1(self): self._test_coding(self.basic_tok, "a<x>a<y>a<z>a", [ 0, 5, 0, 6, 0, 7, 0 ])
-  def test_specials2(self): self._test_coding(self.basic_tok, "<x>a<y>a<z>", [ 5, 0, 6, 0, 7 ])
-  def test_invalid_token(self):
-    with self.assertRaises(RuntimeError): self._test_coding(self.basic_tok, "L", [])
-
-  def test_no_specials(self): self._test_coding(SimpleTokenizer(".*", { bytes([i]): i for i in range(256) }, {}), "abc", [97, 98, 99])
-
-  # NOTE: the correct tokenization for this can only be found by looking up the text chunk in the vocab, not by applying merges
-  def test_llama_early_tokenize(self): self._test_coding(self.llama_tok, " например", [ 111797 ])
-
-  def test_llama_basic(self): self._test_coding(self.llama_tok, "hello world", [ 15339, 1917 ])
-  def test_llama_control_char(self): self._test_coding(self.llama_tok, " \x850", [ 220, 116360, 15 ])
-  def test_llama_bytes(self): self._test_coding(self.llama_tok, " \xec\x8b\xa4\xed", [ 1717, 105, 116174, 82638, 2483 ])
-  def test_llama_special1(self): self._test_coding(self.llama_tok, "hello <|end_of_text|>", [ 15339, 220, 128001 ])
-  def test_llama_special2(self): self._test_coding(self.llama_tok, "<|start_header_id|>user<|end_header_id|>\n\n", [ 128006, 882, 128007, 271 ])
-  def test_llama_repeat(self): self._test_coding(self.llama_tok, "00000000000000000", [ 931, 931, 931, 931, 931, 410 ])
-  def test_llama_pat(self): self._test_coding(self.llama_tok, "today\n  \n", [ 31213, 14211 ])
-
-if __name__ == '__main__':
-  unittest.main()

test/unit/test_uop_symbolic.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-import unittest, pickle, functools, math
+import unittest, pickle, functools
 import z3
 
 from tinygrad.dtype import dtypes, ConstType
@@ -29,17 +29,16 @@ class TestSymbolicPickle(unittest.TestCase):
   def test_pickle_variable_times_2(self): self._test_pickle_unpickle(Variable("a", 3, 8)*2)
 
 class TestSymbolic(unittest.TestCase):
-  def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
+  def helper_test_variable(self, v, n, m, s):
     rendered, nmin, nmax = render(v)
     if isinstance(s, tuple): self.assertIn(rendered, s)
     else: self.assertEqual(rendered, s)
     self.assertEqual(nmin, n)
     self.assertEqual(nmax, m)
-    if test_z3:
-      solver = z3.Solver()
-      z3_sink = graph_rewrite(v.sink(v.simplify()), z3_renderer, ctx=(solver, {}))
-      expr, epxr_simplified = z3_sink.src[0].arg, z3_sink.src[1].arg
-      self.assertEqual(solver.check(expr != epxr_simplified), z3.unsat, "simplified expression not equal to original")
+    solver = z3.Solver()
+    z3_sink = graph_rewrite(v.sink(v.simplify()), z3_renderer, ctx=(solver, {}))
+    expr, epxr_simplified = z3_sink.src[0].arg, z3_sink.src[1].arg
+    self.assertEqual(solver.check(expr != epxr_simplified), z3.unsat, "simplified expression not equal to original")
 
   def test_cmp_simple(self):
     self.helper_test_variable(Variable("a", 3, 8) < 4, 0, 1, "(a<4)")
@@ -673,12 +672,6 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(numerator, 3, 390, "(a*((a*4)+-1))")
     self.helper_test_variable((numerator//denominator)<=0, 1, 1, "True")
 
-  def test_const_reciprocal(self):
-    a = Variable("a", 1, 10, dtypes.float)
-    # TODO: bounds for reciprocal
-    # TODO: should z3 work?
-    self.helper_test_variable(2*(2*a).reciprocal(), -math.inf, math.inf, "(1/a)", test_z3=False)
-
 class TestSymbolicNumeric(unittest.TestCase):
   def helper_test_numeric(self, f):
     MIN, MAX = 0, 10

test/unit/test_viz.py

@@ -106,12 +106,13 @@ class TestViz(BaseTestViz):
 
   # name can also come from a function that returns a TracingKey
   def test_tracing_key(self):
-    @track_rewrites(name=lambda inp,ret: TracingKey("custom_name", (inp,)))
+    @track_rewrites(name=lambda inp,ret: TracingKey("custom_name", (inp,), fmt=f"input={inp.render()}"))
     def test(s:UOp): return graph_rewrite(s, PatternMatcher([]))
     test(UOp.variable("a", 1, 10)+1)
     lst = get_viz_list()
     # NOTE: names from TracingKey do not get deduped
     self.assertEqual(lst[0]["name"], "custom_name")
+    self.assertEqual(lst[0]["fmt"], "input=(a+1)")
 
   def test_colored_label(self):
     # NOTE: dataclass repr prints literal escape codes instead of unicode chars
@@ -137,16 +138,6 @@ class TestViz(BaseTestViz):
     nop = UOp(Ops.NOOP, arg="infinite loop in fixed_point_rewrite")
     self.assertEqual(graphs[2], uop_to_json(nop)[id(nop)])
 
-  def test_const_node_visibility(self):
-    a = UOp.variable("a", 0, 10)
-    z = UOp.const(dtypes.int, 0)
-    alu = a*z
-    exec_rewrite(alu, [sym])
-    graphs = [x["graph"] for x in get_details(tracked_ctxs[0][0])]
-    # embed const in the parent node when possible
-    self.assertEqual(list(graphs[0]), [id(a), id(alu)])
-    self.assertEqual(list(graphs[1]), [id(z)])
-
 # VIZ displays nested graph_rewrites in a tree view
 
 def leaf_rewrite(x:UOp): return x.rtag(1) if x.tag is None else None

tinygrad/apps/llm.py

@@ -1,57 +1,33 @@
 from __future__ import annotations
-import sys, argparse, typing, re, itertools, unicodedata
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, helpers
-
-def gpt2_decode_vocab(voc: dict[str, int]): # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
-  c2b = { chr(cp): cp for cp in itertools.chain(range(ord("!"), ord("~")+1), range(ord("¡"), ord("¬")+1), range(ord("®"), ord("ÿ")+1)) }
-  c2b.update({ chr(256+off): cp for off, cp in enumerate(cp for cp in range(256) if chr(cp) not in c2b) })
-  return { bytes(c2b[c] for c in tok): tid for tok, tid in voc.items() }
-
-def get_llama_re():
-  def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre))
-  r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
-  # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
-  return "(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
-    f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+"
+import sys, argparse
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv
 
 class SimpleTokenizer:
-  def __init__(self, pat: str, normal_tokens: dict[bytes, int], special_tokens: dict[str, int]):
-    self._normal_tokens, self._special_tokens, self._pat = normal_tokens, special_tokens, re.compile(pat)
-    self._tok2str = { tid: tok.encode() for tok, tid in special_tokens.items() } | { tid: tok for tok, tid in normal_tokens.items()  }
-    self._special_re = re.compile("|".join(re.escape(tok) for tok in self._special_tokens.keys()) if special_tokens else r"(?!)")
+  def __init__(self, vocab: list[str]):
+    self.vocab: list[str] = vocab
+    self.biggest_token: int = max(map(len, vocab))
+    self.token_to_id: dict[str, int] = {tok: i for i, tok in enumerate(vocab)}
+    self.replace_space = "Ġ"
+    self.replace_newline = "Ċ"
 
-  @staticmethod
-  def from_gguf_kv(kv: dict):
-    # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L1818-L1820
-    if kv["tokenizer.ggml.pre"] not in ("llama3","llama-v3","llama-bpe"): raise ValueError(f"Invalid tokenizer preset '{kv['tokenizer.ggml.pre']}'")
-    vocab: typing.Iterable[tuple[str, int]] = ((tok, idx) for idx, tok in enumerate(kv["tokenizer.ggml.tokens"]))
-    normal_tokens, special_tokens = helpers.partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
-    return SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens))
+  def encode(self, text:str) -> list[int]:
+    s = text.replace(" ", self.replace_space).replace("\n", self.replace_newline)
+    out: list[int] = []
+    i = 0
+    while i < len(s):
+      j = min(i+self.biggest_token, len(s))
+      while i < j and (tid:=self.token_to_id.get(s[i:j])) is None: j -= 1
+      if tid is None: raise RuntimeError(f"token not found in {s}")
+      assert tid is not None, f"token not found in {s}"
+      out.append(tid)
+      i = j
+    return out
 
-  def encode(self, text: str):
-    tokens: list[int] = []
-    pos = 0
-    for match in self._special_re.finditer(text):
-      tokens.extend(self._encode_sentence(text[pos:match.start(0)]) + [self._special_tokens[text[match.start(0):match.end(0)]]])
-      pos = match.end(0)
-    return tokens + self._encode_sentence(text[pos:])
+  def decode(self, ids: list[int]) -> str:
+    return ''.join(self.vocab[tid] for tid in ids).replace(self.replace_space, " ").replace(self.replace_newline, "\n")
 
-  def decode(self, ids: list[int]) -> str: return b''.join(self._tok2str[tid] for tid in ids).decode()
-  def role(self, role:str): return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
-
-  def _encode_sentence(self, chunk: str): return [ tok for word in self._pat.findall(chunk) for tok in self._encode_word(word.encode()) ]
-  def _encode_word(self, word: bytes):
-    if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token]
-    parts = [word[i:i+1] for i in range(len(word))]
-    while True:
-      min_tid, min_idx = 2**32, -1
-      for idx, (p1, p2) in enumerate(zip(parts[:-1], parts[1:])):
-        tid = self._normal_tokens.get(p1 + p2, min_tid)
-        if tid < min_tid: min_tid, min_idx = tid, idx
-      if min_idx == -1: break
-      parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx+1]] + parts[min_idx+2:]
-    try: return [ self._normal_tokens[p] for p in parts ]
-    except KeyError: raise RuntimeError("token not found")
+  def role(self, role:str):
+    return [t for x in ["<|start_header_id|>", role, "<|end_header_id|>\n\n"] for t in self.encode(x)]  # llama style
 
 def apply_rope(x:Tensor, start_pos:int|UOp, base:int=10000):
   B, H, T, Hd = x.shape
@@ -189,7 +165,7 @@ if __name__ == "__main__":
   model, kv = Transformer.from_gguf(Tensor.from_url(models[args.size]), args.max_context)
 
   # extract some metadata
-  tok = SimpleTokenizer.from_gguf_kv(kv)
+  tok = SimpleTokenizer(kv["tokenizer.ggml.tokens"])
   bos_id: int = kv['tokenizer.ggml.bos_token_id']
   eos_id: int = kv['tokenizer.ggml.eos_token_id']
 

tinygrad/dtype.py

@@ -193,21 +193,6 @@ def least_upper_float(dt:DType) -> DType: return dt if dtypes.is_float(dt) else
 DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void"))}
 INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void"}
 
-@functools.cache
-def can_safe_cast(dt0:DType, dt1:DType) -> bool:
-  # return if dt1 preserves value of dt0
-  # https://numpy.org/doc/stable/reference/generated/numpy.can_cast.html
-  if dt0 == dt1 or dt0 == dtypes.bool: return True
-  match dt1:
-    case dtypes.double: return dt0 in (dtypes.float, dtypes.half, dtypes.bfloat16)
-    case dtypes.float: return dt0 in (dtypes.half, dtypes.bfloat16)
-    case dtypes.uint64: return dt0 in (dtypes.uint32, dtypes.uint16, dtypes.uint8)
-    case dtypes.uint32: return dt0 in (dtypes.uint16, dtypes.uint8)
-    case dtypes.int64: return dt0 in (dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
-    case dtypes.int32: return dt0 in (dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8)
-    case dtypes.int16: return dt0 in (dtypes.uint8, dtypes.int8)
-    case _: return False
-
 def sum_acc_dtype(dt:DType):
   # default acc dtype for sum
   if dtypes.is_unsigned(dt): return least_upper_dtype(dt, dtypes.uint)

tinygrad/engine/jit.py

@@ -21,24 +21,24 @@ def apply_graph_to_jit(jit_cache: list[ExecItem], input_rawbuffers: list[Buffer]
   # This allows the accelerator to run some batches while subsequent graphs are still being updated.
   graphed_jit_cache: list[ExecItem] = []
   current_batch: list[ExecItem] = []
-  current_batch_devs: list[Compiled] = []
+  current_device: Compiled|None = None
 
   def flush_batch():
-    nonlocal current_batch, current_batch_devs, max_batch_size
+    nonlocal current_batch, current_device, max_batch_size
     try:
-      if len(current_batch_devs) == 0: raise GraphException("no device for graph")
+      if current_device is None: raise GraphException("no device for graph")
       if len(current_batch) <= 1 and not getenv("GRAPH_ONE_KERNEL"): raise GraphException("only one kernel doesn't graph")
-      graph_runner = current_batch_devs[0].graph(current_batch, input_rawbuffers, var_vals)
+      graph_runner = current_device.graph(current_batch, input_rawbuffers, var_vals)
       # clear jit inputs to allow their memory to be freed/reused
       for (j,i) in graph_runner.input_replace.keys(): graph_runner.jit_cache[j].bufs[i] = None
       graphed_jit_cache.append(ExecItem(graph_runner, cast(list[Buffer|None], input_rawbuffers)))
       max_batch_size *= 2
-      if DEBUG >= 2: print(f"JIT GRAPHing batch with {len(current_batch)} kernels on device {current_batch_devs[0]}")
+      if DEBUG >= 2: print(f"JIT GRAPHing batch with {len(current_batch)} kernels on device {current_device}")
     except GraphException as e:
       graphed_jit_cache.extend(current_batch)
-      if DEBUG >= 2: print(f"JIT GRAPHing failed batch with {len(current_batch)} kernels on device {current_batch_devs[0]}: {e}")
+      if DEBUG >= 2: print(f"JIT GRAPHing failed batch with {len(current_batch)} kernels on device {current_device}: {e}")
     current_batch = []
-    current_batch_devs = []
+    current_device = None
 
   for ji in jit_cache:
     match ji.prg:
@@ -48,18 +48,13 @@ def apply_graph_to_jit(jit_cache: list[ExecItem], input_rawbuffers: list[Buffer]
       case ViewOp(): continue # ViewOps are just ignored
       case _: ji_graph_dev = None # Everything else is not graphed and flushes existing graph if it's being constructed
 
-    # Check if this jit item can be graphed at all, so check if a new graph supports the current item.
-    can_be_graphed = ji_graph_dev is not None and ji_graph_dev.graph is not None and graph_class(ji_graph_dev).supports_exec_item([ji_graph_dev], ji)
-
-    # Check if the current batch can be extended with this item.
-    new_batched_devs = dedup(current_batch_devs + [ji_graph_dev])
-    can_share_graph = can_be_graphed and len(current_batch_devs) > 0 and graph_class(current_batch_devs[0]).supports_exec_item(new_batched_devs, ji)
+    can_be_graphed = ji_graph_dev is not None and ji_graph_dev.graph is not None and graph_class(ji_graph_dev).supports_exec_item(ji_graph_dev, ji)
+    is_multigraph = can_be_graphed and issubclass(graph_class(ji_graph_dev), MultiGraphRunner)
+    can_share_graph = can_be_graphed and (type(ji_graph_dev) is type(current_device) if is_multigraph else ji_graph_dev == current_device)
     can_extend_graph_batch = can_share_graph and (max_batch_size == 0 or len(current_batch) < max_batch_size)
-
-    # Flush the current batch if any, since it can't be extended or is full.
     if not can_extend_graph_batch and len(current_batch) > 0: flush_batch()
     (current_batch if can_be_graphed else graphed_jit_cache).append(ji)
-    current_batch_devs = new_batched_devs if can_be_graphed else []
+    current_device = ji_graph_dev if can_be_graphed else None
 
   if len(current_batch) > 0: flush_batch()
   return graphed_jit_cache
@@ -132,14 +127,12 @@ class GraphRunner(Runner):
     return list({id(x):x for x in wait_nodes}.values())
 
   @staticmethod
-  def supports_exec_item(devs:list[Compiled], ei:ExecItem) -> bool: return isinstance(ei.prg, CompiledRunner) and len(dedup(devs)) == 1
+  def supports_exec_item(dev, ei:ExecItem) -> bool: return isinstance(ei.prg, CompiledRunner)
 
 # a marker for your graph supporting multiple devices of the same type
 class MultiGraphRunner(GraphRunner):
   @staticmethod
-  def supports_exec_item(devs:list[Compiled], ei:ExecItem) -> bool:
-    # Devices must be the same type
-    return isinstance(ei.prg, (CompiledRunner, BufferXfer)) and len(dedup([type(Device[b.device]) for b in ei.bufs if b]+[type(d) for d in devs]))==1
+  def supports_exec_item(dev, ei:ExecItem) -> bool: return isinstance(ei.prg, (CompiledRunner, BufferXfer))
 
 def get_out_buffers_for_ei(ei:ExecItem) -> list[Buffer]:
   if isinstance(ei.prg, CompiledRunner): return [cast(Buffer, ei.bufs[out]) for out in ei.prg.p.outs if out not in ei.prg.p.ins]

tinygrad/engine/realize.py

@@ -2,7 +2,7 @@ from typing import cast, Generator
 import time, pprint
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
-from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile
+from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer, graph_rewrite, print_uops, track_rewrites
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
@@ -13,7 +13,7 @@ from tinygrad.uop.spec import type_verify
 
 # **************** Program Creation ****************
 
-@track_rewrites(name=lambda _ast,_renderer,ret: TracingKey(ret.name, (ret.function_name, ret.ast), ret=ret))
+@track_rewrites(name=lambda _ast,_renderer,ret: TracingKey(ret.name, (ret.function_name, ret.ast), ret.src, ret=ret))
 def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec:
   """
   Transform an AST into a ProgramSpec. May trigger BEAM search.
@@ -63,10 +63,7 @@ class CompiledRunner(Runner):
   def __init__(self, p:ProgramSpec, precompiled:bytes|None=None, prg=None):
     if DEBUG >= 4: print(p.src)
     self.p:ProgramSpec = p
-    if precompiled is not None: self.lib = precompiled
-    else:
-      with cpu_profile(TracingKey(f"compile {p.name}", (p.function_name,), cat="compiler"), "TINY"):
-        self.lib = Device[p.device].compiler.compile_cached(p.src)
+    self.lib:bytes = precompiled if precompiled is not None else Device[p.device].compiler.compile_cached(p.src)
     if DEBUG >= 7: Device[p.device].compiler.disassemble(self.lib)
     self._prg = Device[p.device].runtime(p.function_name, self.lib) if prg is None else prg
     super().__init__(p.name, p.device, p.estimates)
@@ -159,7 +156,7 @@ class ExecItem:
         lds_est = sym_infer(self.prg.estimates.lds, var_vals)
         mem_est = min(mem_est, lds_est)   # there can't be more memory accessed than loads/stores. remove this when symbolic is fixed
         ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
-        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', 'magenta' if jit else ('green' if self.prg.first_run else None))} {self.prg.display_name+' '*(44-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB " +  # noqa: E501
+        print(f"{colored(f'*** {self.prg.device[:7]:7s} {GlobalCounters.kernel_count:4d}', 'magenta' if jit else ('green' if self.prg.first_run else None))} {self.prg.display_name+' '*(41-ansilen(self.prg.display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:5.2f} GB " +  # noqa: E501
               (str() if et is None else f"tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({op_est/((et or 1e-20)*1e9):9.2f} GFLOPS {mem_est/((et or 1e-20)*1e9):6.1f}|{lds_est/((et or 1e-20)*1e9):<7.1f} GB/s)" +  # noqa: E501
                f" {[repr(m) if TRACEMETA >= 2 else str(m) for m in self.metadata] if self.metadata else ''}"))
       self.prg.first_run = False

tinygrad/gradient.py

@@ -1,5 +1,6 @@
 from typing import cast
 import math, dataclasses
+from tinygrad.dtype import dtypes, sum_acc_dtype
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
 from tinygrad.helpers import argsort
 
@@ -7,7 +8,7 @@ def reduce_gradient(ctx:UOp, ret:UOp):
   def to_inp_shape(x): return x.reshape(x.shape+(1,)*(len(ret.src[0].shape)-len(x.shape))).expand(ret.src[0].shape)
   if ret.arg[0] == Ops.ADD: return (to_inp_shape(ctx),)
   if ret.arg[0] == Ops.MAX:
-    max_is_1s = ret.src[0].eq(to_inp_shape(ret)).cast(ctx.dtype)
+    max_is_1s = ret.src[0].ne(to_inp_shape(ret)).ne(ret.src[0].const_like(1).cast(dtypes.bool)).cast(ctx.dtype)
     div = to_inp_shape(max_is_1s.r(Ops.ADD, ret.arg[1]))
     return ((max_is_1s/div) * to_inp_shape(ctx),)
   if ret.arg[0] == Ops.MUL: return (to_inp_shape(ctx * ret) / ret.src[0],)
@@ -37,7 +38,9 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.arg)])),)),
   (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[1]) for s,p in zip(ret.src[0].shape, ret.arg)])),)),
   (UPat(Ops.FLIP, name="ret"), lambda ctx, ret: (ctx.flip(ret.arg),)),
-  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: (ctx.r(Ops.ADD, tuple(i for i,(si,so) in enumerate(zip(ret.src[0].shape, ret.arg)) if si!=so)),)),
+  # TODO: this cast can be removed by putting the casts around the EXPAND
+  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret:
+    (ctx.cast(sum_acc_dtype(ctx.dtype)).r(Ops.ADD, tuple(i for i,(si,so) in enumerate(zip(ret.src[0].shape, ret.arg)) if si!=so)).cast(ctx.dtype),)),
   (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda ctx: (None,)),

tinygrad/helpers.py

@@ -81,13 +81,6 @@ 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 suppress_finalizing(func):
-  def wrapper(*args, **kwargs):
-    try: return func(*args, **kwargs)
-    except (AttributeError, TypeError, ImportError):
-      if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
-  return wrapper
-
 def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's')
 
 class LazySeq(Generic[T]): # NOTE: Mapping requires __iter__ and __len__, Sequence requires supporting __len__ and slicing in __getitem__
@@ -196,6 +189,7 @@ class Profiling(contextlib.ContextDecorator):
 class TracingKey:
   display_name:str                       # display name of this trace event
   keys:tuple[str, ...]=()                # optional keys to search for related traces
+  fmt:str|None=None                      # optional detailed formatting
   cat:str|None=None                      # optional category to color this by
   ret:Any=None
 

tinygrad/nn/__init__.py

@@ -10,6 +10,7 @@ class BatchNorm:
   """
   Applies Batch Normalization over a 2D or 3D input.
 
+  - Described: https://paperswithcode.com/method/batch-normalization
   - Paper: https://arxiv.org/abs/1502.03167v3
 
   See: `Tensor.batchnorm`
@@ -181,6 +182,7 @@ class GroupNorm:
   """
   Applies Group Normalization over a mini-batch of inputs.
 
+  - Described: https://paperswithcode.com/method/group-normalization
   - Paper: https://arxiv.org/abs/1803.08494v3
 
   ```python exec="true" source="above" session="tensor" result="python"
@@ -211,6 +213,7 @@ class InstanceNorm:
   """
   Applies Instance Normalization over a mini-batch of inputs.
 
+  - Described: https://paperswithcode.com/method/instance-normalization
   - Paper: https://arxiv.org/abs/1607.08022v3
 
   ```python exec="true" source="above" session="tensor" result="python"
@@ -237,6 +240,7 @@ class LayerNorm:
   """
   Applies Layer Normalization over a mini-batch of inputs.
 
+  - Described: https://paperswithcode.com/method/layer-normalization
   - Paper: https://arxiv.org/abs/1607.06450v1
 
   ```python exec="true" source="above" session="tensor" result="python"
@@ -283,6 +287,7 @@ class RMSNorm:
   """
   Applies Root Mean Square Normalization to input.
 
+  - Described: https://paperswithcode.com/method/rmsnorm
   - Paper: https://arxiv.org/abs/1910.07467
 
   ```python exec="true" source="above" session="tensor" result="python"

tinygrad/nn/optim.py

@@ -76,6 +76,8 @@ def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov
   Stochastic Gradient Descent (SGD) optimizer with optional momentum and weight decay.
 
   `classic` is a boolean flag that determines whether to use the popular momentum update rule or the classic momentum update rule.
+
+  - Described: https://paperswithcode.com/method/sgd
   """
   return LARS(params, lr, momentum, weight_decay, nesterov, classic, tcoef=0.0, fused=fused)
 
@@ -83,6 +85,7 @@ class LARS(Optimizer):
   """
   Layer-wise Adaptive Rate Scaling (LARS) optimizer with optional momentum and weight decay.
 
+  - Described: https://paperswithcode.com/method/lars
   - Paper: https://arxiv.org/abs/1708.03888v3
   """
   def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, nesterov=False, classic=True, tcoef=0.001, fused=FUSE_OPTIM):
@@ -116,6 +119,7 @@ def AdamW(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, weight_dec
   """
   AdamW optimizer with optional weight decay.
 
+  - Described: https://paperswithcode.com/method/adamw
   - Paper: https://arxiv.org/abs/1711.05101v3
   """
   return LAMB(params, lr, b1, b2, eps, weight_decay, adam=True, fused=fused)
@@ -123,6 +127,7 @@ def Adam(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, fused=FUSE_
   """
   Adam optimizer.
 
+  - Described: https://paperswithcode.com/method/adam
   - Paper: https://arxiv.org/abs/1412.6980
   """
   return LAMB(params, lr, b1, b2, eps, 0.0, adam=True, fused=fused)
@@ -131,6 +136,7 @@ class LAMB(Optimizer):
   """
   LAMB optimizer with optional weight decay.
 
+  - Described: https://paperswithcode.com/method/lamb
   - Paper: https://arxiv.org/abs/1904.00962
   """
   def __init__(self, params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, adam=False, fused=FUSE_OPTIM):

tinygrad/opt/kernel.py

@@ -90,10 +90,11 @@ class Kernel:
 
     # axis types
     global_loops = AxisType.GLOBAL if self.opts.has_local else AxisType.LOOP
-    self.axis_types: list[AxisType] = [AxisType.REDUCE if resolve(x!=y) else global_loops for x,y in zip(self.output_shape, self.full_shape)]
+    self.axis_types: list[AxisType] = [AxisType.REDUCE if resolve(x!=y) else global_loops for x,y in zip(self.sts[0].shape, self.sts[-1].shape)]
 
     # confirm all reduce axes are at the end
-    if (final_reduces := [x for x in self.axis_types if x == AxisType.REDUCE]) and final_reduces != self.axis_types[-len(final_reduces):]:
+    final_reduces = [i for i,(s,n) in enumerate(zip(self.full_shape, self.output_shape)) if resolve(s != n)]
+    if final_reduces != list(range(len(self.full_shape)-len(final_reduces), len(self.full_shape))):
       raise RuntimeError(f"reduces are not at the end of the shape {self.full_shape} -> {self.output_shape}")
 
   def copy(self):
@@ -200,7 +201,7 @@ class Kernel:
     if self.shape_len == 0: return
     shapes, strides = [x.shape for x in self.sts], [x.real_strides() for x in self.sts]
     # NOTE: we can't use self.first_reduce yet
-    first_reduce = [resolve(x!=y) for x,y in zip(self.output_shape+(0,), self.full_shape+(1,))].index(True)
+    first_reduce = [resolve(x!=y) for x,y in zip(self.sts[0].shape+(0,), self.full_shape+(1,))].index(True)
 
     # if it's an image, insert fake strides such that this fusion doesn't happen across image axes
     # TODO: remove membufs
@@ -448,7 +449,22 @@ class Kernel:
       if op.op in GroupOp.Buffer and op in self.bufs:
         st = self.sts[self.bufs.index(op)]
         # replace the VIEW source
-        return ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
+        if op.op is Ops.LOAD:
+          global_buf = ret.src[0].src[0]
+          # add locals cache
+          local_shape = [s if self.axis_types[i] not in (AxisType.GLOBAL, AxisType.REDUCE)
+                         and ss != 0 else 1 for i,(s,ss) in enumerate(zip(st.shape, st.real_strides()))]
+          # NOTE: this can have any permutation here
+          lst = lst_store = ShapeTracker.from_shape(tuple(local_shape)).expand(st.shape)
+          lbuf = UOp(Ops.DEFINE_LOCAL, dtype=global_buf.dtype.base.ptr(prod(local_shape), AddrSpace.LOCAL), arg=1000+global_buf.arg)
+          # TODO: permute to place any UPCASTs in 0 stride LOCALs
+          # any permutes of st + lst_store together are fine
+          print(list(zip(self.shape_str(), st.shape, st.real_strides())))
+          ret = ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
+          ret = lbuf.view(lst).load(lbuf.view(lst_store).store(ret))
+        else:
+          ret = ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
+        return ret
       if op.op is Ops.SINK:
         # NOTE: should group_for_reduces be added to the local_dims?
         kernel_name = ret.arg.name if ret.arg is not None else self.name if name_override is None else name_override

tinygrad/opt/search.py

@@ -2,7 +2,7 @@ from typing import cast, Callable
 import itertools, functools, random, math, time, multiprocessing, traceback, signal, atexit
 from collections import defaultdict
 from dataclasses import replace
-from tinygrad.uop.ops import UOp, Ops, Variable, sym_infer, AxisType
+from tinygrad.uop.ops import UOp, Ops, Variable, sym_infer
 from tinygrad.device import Device, Buffer, Compiler
 from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
 from tinygrad.helpers import IGNORE_BEAM_CACHE, TC_SEARCH_OVER_SHAPE
@@ -83,7 +83,7 @@ def _try_compile_linearized_w_idx(x:tuple[int,Kernel], compiler:Compiler) -> tup
 
 # workers should not open devices and should ignore ctrl c and should not launch VIZ
 def _init_worker():
-  Context(ALLOW_DEVICE_USAGE=0, VIZ=0, TRACK_MATCH_STATS=0).__enter__()
+  Context(ALLOW_DEVICE_USAGE=0, VIZ=0).__enter__()
   signal.signal(signal.SIGINT, signal.SIG_IGN)
 
 def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_allocated() if buf is not None else buf for buf in bufs]
@@ -108,9 +108,9 @@ def bufs_from_lin(lin:Kernel, allocate:bool=True) -> list[Buffer]:
   return cast(list[Buffer], rawbufs)
 
 # get dictionary of all possible actions
-def get_kernel_actions(lin:Kernel, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Kernel]:
+def get_kernel_actions(lin:Kernel, include_0=True) -> dict[int, Kernel]:
   acted_lins, max_up, max_lcl = {0:lin} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256), getenv("BEAM_LOCAL_MAX", 1024)
-  kernel_actions = (actions if candidates is None else candidates).copy()
+  kernel_actions = actions.copy()
 
   if TC_SEARCH_OVER_SHAPE and len(lin.applied_opts) == 0: # tensor core opts must be first
     for i, action in enumerate(kernel_actions):
@@ -123,14 +123,14 @@ def get_kernel_actions(lin:Kernel, include_0=True, candidates:list[Opt]|None=Non
     if a.axis is not None and a.op is not OptOps.TC:
       try: ax = lin.real_axis(a.op, a.axis)
       except KernelOptError: continue
-      if (ax >= lin.shape_len) or (lin.full_shape[ax] == a.arg and Opt(a.op, a.axis, 0) in kernel_actions): continue
+      if (ax >= lin.shape_len) or (lin.full_shape[ax] == a.arg and Opt(a.op, ax, 0) in kernel_actions): continue
     lin2 = lin.copy()
     try:
       lin2.apply_opt(a)
       up, lcl, tc_up = 1, 1, prod(tc.dims)//tc.threads if (tc:=lin2.tensor_core) else 1
-      for s,c in zip(lin2.full_shape, lin2.axis_types):
-        if c in (AxisType.UPCAST, AxisType.UNROLL): up *= s
-        elif c in (AxisType.LOCAL, AxisType.GROUP_REDUCE): lcl *= s
+      for s,c in zip(lin2.full_shape, lin2.colors()):
+        if c in {"magenta", "yellow"}: up *= s
+        elif c in {"cyan", "green", "white"}: lcl *= s
       if up//tc_up > max_up or lcl > max_lcl:
         if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many upcast/local. {up//tc_up=}, {max_up=}, {lcl=}, {max_lcl=}")
         continue

tinygrad/opt/tc.py

@@ -37,10 +37,6 @@ class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x
     assert 2**local_axes == self.threads, f"{self.threads} threads construct the warp but found {2**local_axes} in {self.opts}"
     assert 2**upcast_axes == self.elements_per_thread[2], \
       f"{self.elements_per_thread[2]} elements from C are processed per thread but found {2**upcast_axes} in {self.opts}"
-    # check dims match opts
-    assert self.dims[0] == 2**len(gd:=[x for x in self.opts if x[1] == '0']), f"opts wrong on dims[0], {self.dims[0]} vs {gd}"
-    assert self.dims[1] == 2**len(gd:=[x for x in self.opts if x[1] == '1']), f"opts wrong on dims[1], {self.dims[1]} vs {gd}"
-    # NOTE: the K opts is implictly set by the dim
     # check swizzle
     assert len(self.swizzle[0]) == 3 and len(self.swizzle[1]) == 3, "swizzle has wrong part count"
     assert len(self.swizzle[0][0]) == len(self.swizzle[1][0]) == local_axes, "local swizzle size is wrong"
@@ -56,7 +52,6 @@ class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x
       if o[1] == '1': zero_stride_1.append(o[0] + str(un if o[0] == 'u' else ln))
       if o[0] == 'u': un += 1
       if o[0] == 'l': ln += 1
-    # NOTE: all the zero_stride dims can be placed in any order in the swizzle
     upcasted_0 = [x for x in (self.swizzle[0][1] + self.swizzle[0][2]) if x not in zero_stride_0 and x[0] != 'l']
     upcasted_1 = [x for x in (self.swizzle[1][1] + self.swizzle[1][2]) if x not in zero_stride_1 and x[0] != 'l']
     assert 2**len(upcasted_0) == self.elements_per_thread[0], f"mismatch in elements_per_thread[0], {upcasted_0} vs {self.elements_per_thread[0]}"

tinygrad/runtime/graph/hcq.py

@@ -1,7 +1,7 @@
 import collections, time
 from typing import Any, cast
-from tinygrad.helpers import round_up, PROFILE, merge_dicts, getenv, dedup
-from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator, MMIOInterface
+from tinygrad.helpers import round_up, PROFILE, merge_dicts, getenv
+from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator
 from tinygrad.device import Buffer, BufferSpec, Compiled, Device, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Variable
@@ -29,7 +29,7 @@ class HCQGraph(MultiGraphRunner):
     for ji in jit_cache:
       if not isinstance(ji.prg, CompiledRunner): continue
       kernargs_size[ji.prg.dev] += round_up(ji.prg._prg.kernargs_alloc_size, 16)
-    self.kernargs_bufs: dict[Compiled, HCQBuffer] = {d:d.allocator._alloc(max(sz, 1), BufferSpec(cpu_access=True)) for d,sz in kernargs_size.items()}
+    self.kernargs_bufs: dict[Compiled, HCQBuffer] = {dev:dev.allocator._alloc(sz, BufferSpec(cpu_access=True)) for dev,sz in kernargs_size.items()}
 
     # Fill initial arguments.
     self.ji_args: dict[int, HCQArgsState] = {}
@@ -51,8 +51,8 @@ class HCQGraph(MultiGraphRunner):
     self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: dev.hw_compute_queue_t() for dev in self.devices}
     self.copy_queues: dict[HCQCompiled, HWQueue] = {} # lazy allocation
 
-    self.signals: dict[Any, HCQSignal] = {**{dev: dev.new_signal(value=0) for dev in self.devices if not dev._is_cpu()},
-      **{"KICK": self.devices[0].new_signal(value=0)}, **{dev: self.devices[0].new_signal(value=0) for dev in self.devices if dev._is_cpu()}}
+    self.signals: dict[Any, HCQSignal] = {**{dev: dev.new_signal(value=0) for dev in self.devices if dev.device != "CPU"},
+      **{"KICK": self.devices[0].new_signal(value=0)}, **{dev: self.devices[0].new_signal(value=0) for dev in self.devices if dev.device == "CPU"}}
     self.kickoff_value: int = 0
     self.kickoff_var = UOp.variable("kickoff_var", 0, 0xffffffff, dtype=dtypes.uint32)
 
@@ -87,7 +87,7 @@ class HCQGraph(MultiGraphRunner):
         assert (enqueue_dev.hw_copy_queue_t is not None), "device must implement a copy queue"
         enqueue_queue = self.copy_queues.setdefault(enqueue_dev, enqueue_dev.hw_copy_queue_t())
 
-      out_signal = self.signals.setdefault(enqueue_queue, self.devices[0].new_signal(value=0))
+      out_signal = self.signals.setdefault(enqueue_queue, enqueue_dev.new_signal(value=0))
 
       # Get dependencies based on input and output buffers.
       rdeps = self._access_resources(ji.bufs, ji.prg.p.outs if is_exec_prg else [0], (enqueue_queue, j + 1)) #type:ignore
@@ -225,17 +225,7 @@ class HCQGraph(MultiGraphRunner):
     for fdev, buf in self.kernargs_bufs.items(): fdev.allocator._free(buf, BufferSpec(cpu_access=True))
 
   @staticmethod
-  def supports_exec_item(devs:list[Compiled], ei:ExecItem) -> bool:
-    # Check if all devices are HCQ
-    all_devs = cast(list[HCQCompiled], dedup(devs + [Device[b.device] for b in ei.bufs if b]))
-    if not all(issubclass(type(d), HCQCompiled) for d in all_devs): return False
-
-    # If all of devices are mapped into CPU address space, can use CPU inside the peer group.
-    cpu_support = all(isinstance(d.timeline_signal.base_buf.view, MMIOInterface) for d in all_devs)
-
-    # Check if all devices are within the same peer group. If CPU is supported, don't count it as a separate peer group.
-    if len(set(d.peer_group for d in all_devs if cpu_support and not d._is_cpu())) > 1: return False
-
+  def supports_exec_item(dev, ei:ExecItem) -> bool:
     # MOCKGPU is not supported, since it can't execute commands in parallel
-    copy = (isinstance(ei.prg, BufferCopy) and cast(HCQCompiled, devs[0]).hw_copy_queue_t is not None) and not getenv("MOCKGPU")
-    return isinstance(ei.prg, (CompiledRunner, BufferXfer)) or copy
+    copy = (isinstance(ei.prg, BufferCopy) and cast(HCQCompiled, dev).hw_copy_queue_t is not None) and not getenv("MOCKGPU")
+    return all(issubclass(type(Device[b.device]), HCQCompiled) for b in ei.bufs if b) and (isinstance(ei.prg, (CompiledRunner, BufferXfer)) or copy)

tinygrad/runtime/ops_amd.py

@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec
-from tinygrad.helpers import getenv, to_mv, round_up, data64_le, all_same, flatten, DEBUG, AMD_LLVM, PROFILE, ProfileEvent, suppress_finalizing
+from tinygrad.helpers import getenv, to_mv, round_up, data64_le, all_same, flatten, DEBUG, AMD_LLVM, PROFILE, ProfileEvent
 from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
@@ -473,10 +473,11 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
     return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access)
 
-  @suppress_finalizing
   def _free(self, opaque, options:BufferSpec):
-    self.dev.synchronize()
-    self.dev.iface.free(opaque)
+    try:
+      self.dev.synchronize()
+      self.dev.iface.free(opaque)
+    except AttributeError: pass
 
   def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
 
@@ -592,7 +593,7 @@ class KFDIface:
 
   def free(self, mem):
     if len(mem.mapped_devs) > 0:
-      gpus = (ctypes.c_int32 * len(mem.mapped_devs))(*[x.iface.gpu_id for x in mem.mapped_devs])
+      gpus = (ctypes.c_int32 * len(mem.mapped_devs))(*[x.gpu_id for x in mem.mapped_devs])
       stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=len(gpus))
       assert stm.n_success == len(gpus)
     if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
@@ -673,8 +674,8 @@ class PCIIface(PCIIfaceBase):
       self.dev_impl.gfx.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr, wptr_addr=gart.va_addr+0x10,
         eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size, doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_MEC_RING0), pipe=0, queue=0)
 
-    return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbells=[self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q')],
-      read_ptrs=[gart.cpu_view().view(size=8, fmt='Q')], write_ptrs=[gart.cpu_view().view(offset=0x10, size=8, fmt='Q')])
+    return AMDQueueDesc(ring=MMIOInterface(ring.va_addr, ring.size, fmt='I'), read_ptrs=[MMIOInterface(gart.va_addr, 8, fmt='Q')],
+      write_ptrs=[MMIOInterface(gart.va_addr+0x10, 8, fmt='Q')], doorbells=[MMIOInterface(self.doorbell_cpu_addr + doorbell_index * 8, 8, fmt='Q')])
 
   def sleep(self, timeout):
     if self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -717,8 +718,16 @@ class USBIface(PCIIface):
       view=USBMMIOInterface(self.usb, self.bars[0][0] + am_mapping.paddrs[0][0], size, fmt='B') if cpu_access else None, owner=self.dev)
 
   def create_queue(self, queue_type, ring, gart, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
-    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE: self.usb._pci_cacheable += [(ring.cpu_view().addr, ring.size)]
-    return super().create_queue(queue_type, ring, gart, eop_buffer, cwsr_buffer, ctl_stack_size, ctx_save_restore_size, xcc_id)
+    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
+      self.dev_impl.sdma.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr, wptr_addr=gart.va_addr+0x10,
+                                    doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0), pipe=0, queue=0)
+    else:
+      self.dev_impl.gfx.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr, wptr_addr=gart.va_addr+0x10,
+        eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size, doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_MEC_RING0), pipe=0, queue=0)
+      self.usb._pci_cacheable += [(ring.cpu_view().addr, ring.size)]
+
+    return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbells=[self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q')],
+      read_ptrs=[gart.cpu_view().view(size=8, fmt='Q')], write_ptrs=[gart.cpu_view().view(offset=0x10, size=8, fmt='Q')])
 
   def sleep(self, timeout): pass
 

tinygrad/runtime/ops_cpu.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import platform, subprocess, sys, ctypes, functools, time, mmap, threading, queue
+import platform, subprocess, sys, ctypes, functools, time, mmap
 from tinygrad.helpers import capstone_flatdump, getenv, from_mv, to_mv, OSX, mv_address, wait_cond, cpu_profile
 from tinygrad.device import Compiler, BufferSpec, DMACPURef
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
@@ -7,10 +7,6 @@ from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.uop.ops import sint
 
-class CPUSignal(HCQSignal):
-  def _sleep(self, time_spent_waiting_ms:int):
-    if self.is_timeline and self.owner is not None: self.owner.tasks.join()
-
 class ClangJITCompiler(Compiler):
   def __init__(self, cachekey="compile_clang_jit"): super().__init__(cachekey)
 
@@ -25,19 +21,6 @@ class ClangJITCompiler(Compiler):
 
   def disassemble(self, lib:bytes): return capstone_flatdump(lib)
 
-class CPUWorker(threading.Thread):
-  def __init__(self, dev):
-    super().__init__()
-    self.dev, self.tasks, self.daemon = dev, dev.tasks, True
-
-  def run(self):
-    while True:
-      cmd_iter = iter(self.tasks.get())
-      for cmd in cmd_iter:
-        args_cnt = next(cmd_iter)
-        cmd(*[next(cmd_iter) for _ in range(args_cnt)])
-      self.tasks.task_done()
-
 class CPUComputeQueue(HWQueue):
   def _exec(self, prg, bufs, *args):
     prg.fxn(*map(ctypes.c_uint64, args[:bufs]), *map(ctypes.c_int64 if platform.machine() == "arm64" else ctypes.c_int32, args[bufs:]))
@@ -54,7 +37,13 @@ class CPUComputeQueue(HWQueue):
   def wait(self, signal, value=0): return self.cmd(self._wait, signal.value_addr, value)
   def timestamp(self, signal): return self.cmd(self._timestamp, signal.timestamp_addr)
   def signal(self, signal, value:sint=0): return self.cmd(self._signal, signal.value_addr, value)
-  def _submit(self, dev): dev.tasks.put(self._q[:])
+
+  def _submit(self, dev):
+    # Execute the commands in the queue: fn, argc, args...
+    off = 0
+    while off < len(self._q):
+      self._q[off](*self._q[off + 2:off + 2 + self._q[off + 1]])
+      off += self._q[off + 1] + 2
 
 # NOTE: MAP_JIT is added to mmap module in python 3.13
 MAP_JIT = 0x0800
@@ -101,23 +90,16 @@ class CPUAllocator(HCQAllocatorBase):
     elif sys.platform == "win32": addr = mv_address(buf:=mmap.mmap(-1, size, access=mmap.ACCESS_WRITE))
     else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_PRIVATE, mmap.PROT_READ | mmap.PROT_WRITE))
     return HCQBuffer(va:=addr, sz:=size, meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev)
-  def _as_buffer(self, src) -> memoryview:
-   self.dev.synchronize()
-   return to_mv(src.va_addr, src.size)
-  def _as_dmaref(self, buf):
-    self.dev.synchronize()
-    return DMACPURef(buf.va_addr, buf.size)
+  def _as_buffer(self, src) -> memoryview: return to_mv(src.va_addr, src.size)
+  def _as_dmaref(self, buf): return DMACPURef(buf.va_addr, buf.size)
   def _copyin(self, dest, src:memoryview):
-    self.dev.synchronize()
     with cpu_profile('TINY -> CPU', self.dev.device, is_copy=True): ctypes.memmove(dest.va_addr, from_mv(src), len(src))
   def _copyout(self, dest:memoryview, src):
-    self.dev.synchronize()
     with cpu_profile('CPU -> TINY', self.dev.device, is_copy=True): ctypes.memmove(from_mv(dest), src.va_addr, len(dest))
   def _map(self, buf:HCQBuffer):
     if buf.view is None or not isinstance(buf.view, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")
 
 class CPUDevice(HCQCompiled):
   def __init__(self, device:str=""):
-    self.tasks:queue.Queue = queue.Queue()
-    CPUWorker(self).start()
-    super().__init__(device, CPUAllocator(self), ClangRenderer(), ClangJITCompiler(), functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue)
+    super().__init__(device, CPUAllocator(self), ClangRenderer(), ClangJITCompiler(), functools.partial(CPUProgram, self), HCQSignal, CPUComputeQueue,
+                     supports_graph=False)

tinygrad/runtime/ops_cuda.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import ctypes, ctypes.util, functools
-from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t, suppress_finalizing
+from tinygrad.helpers import DEBUG, getenv, mv_address, init_c_var, init_c_struct_t
 from tinygrad.device import Compiled, BufferSpec, LRUAllocator
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.ptx import PTXRenderer
@@ -45,8 +45,9 @@ class CUDAProgram:
     self.prg = prg
     if self.smem > 0: check(cuda.cuFuncSetAttribute(self.prg, cuda.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, self.smem))
 
-  @suppress_finalizing
-  def __del__(self): check(cuda.cuModuleUnload(self.module))
+  def __del__(self):
+    try: check(cuda.cuModuleUnload(self.module))
+    except AttributeError: pass
 
   def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
     check(cuda.cuCtxSetCurrent(self.dev.context))

tinygrad/runtime/ops_gpu.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from typing import cast
 import ctypes, functools, hashlib
 from tinygrad.runtime.autogen import opencl as cl
-from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, getenv, mv_address, suppress_finalizing
+from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, getenv, mv_address
 from tinygrad.renderer.cstyle import OpenCLRenderer, IntelRenderer
 from tinygrad.device import BufferSpec, LRUAllocator, Compiled, Compiler, CompileError
 
@@ -69,8 +69,9 @@ class CLAllocator(LRUAllocator['CLDevice']):
                                         cl.cl_image_format(cl.CL_RGBA, {2: cl.CL_HALF_FLOAT, 4: cl.CL_FLOAT}[options.image.itemsize]),
                                         options.image.shape[1], options.image.shape[0], 0, None, status := ctypes.c_int32()), status), options)
     return (checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status), options)
-  @suppress_finalizing
-  def _free(self, opaque:tuple[ctypes._CData, BufferSpec], options:BufferSpec): check(cl.clReleaseMemObject(opaque[0]))
+  def _free(self, opaque:tuple[ctypes._CData, BufferSpec], options:BufferSpec):
+    try: check(cl.clReleaseMemObject(opaque[0]))
+    except AttributeError: pass
   def _copyin(self, dest:tuple[ctypes._CData, BufferSpec], src:memoryview):
     if dest[1].image is not None:
       check(cl.clEnqueueWriteImage(self.dev.queue, dest[0], False, (ctypes.c_size_t * 3)(0,0,0),

tinygrad/runtime/ops_llvm.py

@@ -1,7 +1,7 @@
-import ctypes, platform, functools, queue
+import ctypes, platform, functools
 from tinygrad.device import Compiler
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQSignal
-from tinygrad.runtime.ops_cpu import CPUAllocator, CPUProgram, CPUComputeQueue, CPUWorker
+from tinygrad.runtime.ops_cpu import CPUAllocator, CPUProgram, CPUComputeQueue
 from tinygrad.helpers import OSX, getenv, capstone_flatdump, DEBUG
 from tinygrad.renderer.llvmir import LLVMRenderer
 import tinygrad.runtime.autogen.llvm as llvm
@@ -73,6 +73,5 @@ class HostLLVMCompiler(LLVMCompiler):
 
 class LLVMDevice(HCQCompiled):
   def __init__(self, device:str=""):
-    self.tasks:queue.Queue = queue.Queue()
-    CPUWorker(self).start()
-    super().__init__(device, CPUAllocator(self), LLVMRenderer(), HostLLVMCompiler(), functools.partial(CPUProgram, self), HCQSignal, CPUComputeQueue)
+    super().__init__(device, CPUAllocator(self), LLVMRenderer(), HostLLVMCompiler(), functools.partial(CPUProgram, self), HCQSignal, CPUComputeQueue,
+                     supports_graph=False)

tinygrad/runtime/ops_nv.py

@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H
 from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU
 from tinygrad.uop.ops import sint
 from tinygrad.device import BufferSpec
-from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32, suppress_finalizing
+from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, to_mv, hi32, lo32
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import NVRenderer
 from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, PTX, NVPTXCompiler, NVCompiler
@@ -276,10 +276,11 @@ class NVAllocator(HCQAllocator['NVDevice']):
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
     return self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host)
 
-  @suppress_finalizing
   def _free(self, opaque:HCQBuffer, options:BufferSpec):
-    self.dev.synchronize()
-    self.dev.iface.free(opaque)
+    try:
+      self.dev.synchronize()
+      self.dev.iface.free(opaque)
+    except AttributeError: pass
 
   def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
 

tinygrad/runtime/ops_remote.py

@@ -225,8 +225,7 @@ class RemoteHandler:
             graph_cls = graph_class(Device[self.base_device])
             rp = RemoteProperties(
               real_device=dev.device, renderer=(cls.__module__, cls.__name__, args), offset_supported=hasattr(dev.allocator, '_offset'),
-              graph_supported=graph_cls is not None,
-              graph_supports_multi=graph_cls is not None and issubclass(graph_cls, MultiGraphRunner) and hasattr(dev.allocator, '_transfer'),
+              graph_supported=graph_cls is not None, graph_supports_multi=graph_cls is not None and issubclass(graph_cls, MultiGraphRunner),
               ib_gid=bytes(self.ib_ctx.gid_attr.raw) if self.ib_ctx is not None else None,
             )
             ret = repr(rp).encode()

tinygrad/runtime/ops_webgpu.py

@@ -1,7 +1,7 @@
 import functools, struct
 from tinygrad.device import  Compiled, Allocator, Compiler, BufferSpec
 from tinygrad.renderer.wgsl import WGSLRenderer
-from tinygrad.helpers import round_up, suppress_finalizing
+from tinygrad.helpers import round_up
 from tinygrad.runtime.autogen import webgpu
 from typing import List, Any, TypeAlias
 import ctypes
@@ -188,8 +188,9 @@ class WebGpuAllocator(Allocator['WGPUDevPtr']):
   def _copyout(self, dest:memoryview, src:WGPUBufPtr):
     buffer_data = read_buffer(self.dev, src)
     dest[:] = buffer_data[:dest.nbytes] if webgpu.wgpuBufferGetSize(src)  > dest.nbytes else buffer_data
-  @suppress_finalizing
-  def _free(self, opaque:WGPUBufPtr, options:BufferSpec): webgpu.wgpuBufferDestroy(opaque)
+  def _free(self, opaque:WGPUBufPtr, options:BufferSpec):
+    try: webgpu.wgpuBufferDestroy(opaque)
+    except AttributeError: pass
 
 class WebGpuDevice(Compiled):
   def __init__(self, device:str):

tinygrad/runtime/support/am/ip.py

@@ -169,12 +169,12 @@ class AM_SMU(AM_IP):
       self._send_msg(self.smu_mod.PPSMC_MSG_SetSoftMinByFreq, clck << 16 | (vals[level]))
       self._send_msg(self.smu_mod.PPSMC_MSG_SetSoftMaxByFreq, clck << 16 | (vals[level]))
 
-  def _smu_cmn_send_msg(self, msg:int, param=0, debug=False):
+  def _smu_cmn_send_msg(self, msg, param=0, debug=False):
     (self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).write(0) # resp reg
     (self.adev.mmMP1_SMN_C2PMSG_82 if not debug else self.adev.mmMP1_SMN_C2PMSG_53).write(param)
     (self.adev.mmMP1_SMN_C2PMSG_66 if not debug else self.adev.mmMP1_SMN_C2PMSG_75).write(msg)
 
-  def _send_msg(self, msg:int, param:int, read_back_arg=False, timeout=10000, debug=False): # default timeout is 10 seconds
+  def _send_msg(self, msg, param, read_back_arg=False, timeout=10000, debug=False): # 10s
     self._smu_cmn_send_msg(msg, param, debug=debug)
     wait_cond(lambda: (self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).read(), value=1, timeout_ms=timeout,
       msg=f"SMU msg {msg:#x} timeout")
@@ -414,12 +414,12 @@ class AM_PSP(AM_IP):
 
   def _wait_for_bootloader(self): wait_cond(lambda: self.adev.reg(f"{self.reg_pref}_35").read() & 0x80000000, value=0x80000000, msg="BL not ready")
 
-  def _prep_msg1(self, data:memoryview):
+  def _prep_msg1(self, data):
     assert len(data) <= self.msg1_view.nbytes, f"msg1 buffer is too small {len(data):#x} > {self.msg1_view.nbytes:#x}"
     self.msg1_view[:len(data)+4] = bytes(data) + b'\x00' * 4
     self.adev.gmc.flush_hdp()
 
-  def _bootloader_load_component(self, fw:int, compid:int):
+  def _bootloader_load_component(self, fw, compid):
     if fw not in self.adev.fw.sos_fw: return 0
 
     self._wait_for_bootloader()
@@ -458,7 +458,7 @@ class AM_PSP(AM_IP):
 
     wait_cond(lambda: self.adev.reg(f"{self.reg_pref}_64").read() & 0x8000FFFF, value=0x80000000, msg="sOS ring not created")
 
-  def _ring_submit(self, cmd:am.struct_psp_gfx_cmd_resp) -> am.struct_psp_gfx_cmd_resp:
+  def _ring_submit(self, cmd):
     msg = am.struct_psp_gfx_rb_frame(fence_value=(prev_wptr:=self.adev.reg(f"{self.reg_pref}_67").read()),
       cmd_buf_addr_lo=lo32(self.adev.paddr2mc(self.cmd_paddr)), cmd_buf_addr_hi=hi32(self.adev.paddr2mc(self.cmd_paddr)),
       fence_addr_lo=lo32(self.adev.paddr2mc(self.fence_paddr)), fence_addr_hi=hi32(self.adev.paddr2mc(self.fence_paddr)))
@@ -477,7 +477,7 @@ class AM_PSP(AM_IP):
 
     return resp
 
-  def _load_ip_fw_cmd(self, fw_types:list[int], fw_bytes:memoryview):
+  def _load_ip_fw_cmd(self, fw_types, fw_bytes):
     self._prep_msg1(fw_bytes)
     for fw_type in fw_types:
       if DEBUG >= 2: print(f"am {self.adev.devfmt}: loading fw: {am.psp_gfx_fw_type__enumvalues[fw_type]}")
@@ -487,7 +487,7 @@ class AM_PSP(AM_IP):
       cmd.cmd.cmd_load_ip_fw.fw_type = fw_type
       self._ring_submit(cmd)
 
-  def _tmr_load_cmd(self) -> am.struct_psp_gfx_cmd_resp:
+  def _tmr_load_cmd(self):
     cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_SETUP_TMR)
     cmd.cmd.cmd_setup_tmr.buf_phy_addr_hi, cmd.cmd.cmd_setup_tmr.buf_phy_addr_lo = data64(self.adev.paddr2mc(self.tmr_paddr))
     cmd.cmd.cmd_setup_tmr.system_phy_addr_hi, cmd.cmd.cmd_setup_tmr.system_phy_addr_lo = data64(self.tmr_paddr)
@@ -495,7 +495,7 @@ class AM_PSP(AM_IP):
     cmd.cmd.cmd_setup_tmr.buf_size = self.tmr_size
     return self._ring_submit(cmd)
 
-  def _load_toc_cmd(self, toc_size:int) -> am.struct_psp_gfx_cmd_resp:
+  def _load_toc_cmd(self, toc_size):
     cmd = am.struct_psp_gfx_cmd_resp(cmd_id=am.GFX_CMD_ID_LOAD_TOC)
     cmd.cmd.cmd_load_toc.toc_phy_addr_hi, cmd.cmd.cmd_load_toc.toc_phy_addr_lo = data64(self.msg1_addr)
     cmd.cmd.cmd_load_toc.toc_size = toc_size

tinygrad/runtime/support/hcq.py

@@ -358,14 +358,14 @@ class HCQCompiled(Compiled, Generic[SignalType]):
   peer_groups: dict[str, list[HCQCompiled]] = collections.defaultdict(list)
   signal_pages: dict[str, list[HCQBuffer]] = collections.defaultdict(list) # per peer group
   signal_pool: dict[str, list[HCQBuffer]] = collections.defaultdict(list) # per peer group
-  cpu_devices: list[HCQCompiled] = []
 
   def __init__(self, device:str, allocator:HCQAllocatorBase, renderer:Renderer, compiler:Compiler, runtime, signal_t:Type[SignalType],
-               comp_queue_t:Callable[[], HWQueue], copy_queue_t:Callable[[], HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000):
+               comp_queue_t:Callable[[], HWQueue], copy_queue_t:Callable[[], HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000,
+               supports_graph=True):
     self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
 
     from tinygrad.runtime.graph.hcq import HCQGraph
-    super().__init__(device, allocator, renderer, compiler, runtime, HCQGraph)
+    super().__init__(device, allocator, renderer, compiler, runtime, HCQGraph if supports_graph else None)
 
     # TODO: peer logic is determined based on device name.
     self.peer_group = device.split(":")[0]
@@ -383,13 +383,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     self.kernargs_buf:HCQBuffer = self.allocator.alloc(kernargs_size, BufferSpec(cpu_access=True))
     self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(self.kernargs_buf.size, wrap=True)
 
-    if self._is_cpu(): HCQCompiled.cpu_devices.append(self)
-
   def synchronize(self):
-    # If we have any work on CPU devices, need to synchronize them. This is just an optimization to release GIL allowing to finish faster.
-    if not self._is_cpu():
-      for dev in HCQCompiled.cpu_devices: dev.synchronize()
-
     try: self.timeline_signal.wait(self.timeline_value - 1)
     except RuntimeError as e:
       if hasattr(self, 'on_device_hang'): self.on_device_hang()

tinygrad/runtime/support/nv/ip.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import ctypes, time, array, struct, itertools, dataclasses
-from typing import cast, Any
+from typing import cast
 from tinygrad.runtime.autogen.nv import nv
 from tinygrad.helpers import to_mv, lo32, hi32, DEBUG, round_up, round_down, mv_address, fetch, wait_cond
 from tinygrad.runtime.support.system import System
@@ -8,7 +8,7 @@ from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.autogen import nv_gpu
 
 @dataclasses.dataclass(frozen=True)
-class GRBufDesc: size:int; virt:bool; phys:bool; local:bool=False # noqa: E702
+class GRBufDesc: size:int; v:int; p:int; lc:int=0 # noqa: E702
 
 class NV_IP:
   def __init__(self, nvdev): self.nvdev = nvdev
@@ -26,13 +26,13 @@ class NVRpcQueue:
     self.gsp, self.va, self.queue_va, self.seq = gsp, va, va + self.tx.entryOff, 0
     self.queue_mv = to_mv(self.queue_va, self.tx.msgSize * self.tx.msgCount)
 
-  def _checksum(self, data:bytes):
+  def _checksum(self, data):
     if (pad_len:=(-len(data)) % 8): data += b'\x00' * pad_len
     checksum = 0
     for offset in range(0, len(data), 8): checksum ^= struct.unpack_from('Q', data, offset)[0]
     return hi32(checksum) ^ lo32(checksum)
 
-  def send_rpc(self, func:int, msg:bytes, wait=False):
+  def send_rpc(self, func, msg, wait=False):
     header = nv.rpc_message_header_v(signature=nv.NV_VGPU_MSG_SIGNATURE_VALID, rpc_result=nv.NV_VGPU_MSG_RESULT_RPC_PENDING,
       rpc_result_private=nv.NV_VGPU_MSG_RESULT_RPC_PENDING, header_version=(3<<24), function=func, length=len(msg) + 0x20)
 
@@ -49,7 +49,7 @@ class NVRpcQueue:
     self.seq += 1
     self.gsp.nvdev.NV_PGSP_QUEUE_HEAD[0].write(0x0)
 
-  def wait_resp(self, cmd:int) -> memoryview:
+  def wait_resp(self, cmd) -> memoryview:
     while True:
       System.memory_barrier()
       if self.rx.readPtr == self.tx.writePtr: continue
@@ -60,8 +60,7 @@ class NVRpcQueue:
 
       # Handling special functions
       if hdr.function == nv.NV_VGPU_MSG_EVENT_GSP_RUN_CPU_SEQUENCER: self.gsp.run_cpu_seq(msg)
-      elif hdr.function == nv.NV_VGPU_MSG_EVENT_OS_ERROR_LOG:
-        print(f"nv {self.gsp.nvdev.devfmt}: GSP LOG: {msg[12:].tobytes().rstrip(bytes([0])).decode('utf-8')}")
+      elif hdr.function == nv.NV_VGPU_MSG_EVENT_OS_ERROR_LOG: print(f"GSP LOG: {msg[12:].tobytes().rstrip(bytes([0])).decode('utf-8')}")
 
       # Update the read pointer
       self.rx.readPtr = (self.rx.readPtr + round_up(hdr.length, self.tx.msgSize) // self.tx.msgSize) % self.tx.msgCount
@@ -178,7 +177,7 @@ class NV_FLCN(NV_IP):
     self.nvdev.NV_PFALCON_FALCON_OS.with_base(self.falcon).write(0x0)
     assert self.nvdev.NV_PRISCV_RISCV_CPUCTL.with_base(self.falcon).read_bitfields()['active_stat'] == 1, "GSP Core is not active"
 
-  def execute_dma(self, base:int, cmd:int, dest:int, mem_off:int, sysmem:int, size:int):
+  def execute_dma(self, base, cmd, dest, mem_off, sysmem, size):
     wait_cond(lambda: self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).read_bitfields()['full'], value=0, msg="DMA does not progress")
 
     self.nvdev.NV_PFALCON_FALCON_DMATRFBASE.with_base(base).write(lo32(sysmem >> 8))
@@ -195,7 +194,7 @@ class NV_FLCN(NV_IP):
 
     wait_cond(lambda: self.nvdev.NV_PFALCON_FALCON_DMATRFCMD.with_base(base).read_bitfields()['idle'], msg="DMA does not complete")
 
-  def start_cpu(self, base:int):
+  def start_cpu(self, base):
     if self.nvdev.NV_PFALCON_FALCON_CPUCTL.with_base(base).read_bitfields()['alias_en'] == 1:
       self.nvdev.wreg(base + self.nvdev.NV_PFALCON_FALCON_CPUCTL_ALIAS, 0x2)
     else: self.nvdev.NV_PFALCON_FALCON_CPUCTL.with_base(base).write(startcpu=1)
@@ -233,11 +232,11 @@ class NV_FLCN(NV_IP):
     if mailbox is not None:
       return self.nvdev.NV_PFALCON_FALCON_MAILBOX0.with_base(base).read(), self.nvdev.NV_PFALCON_FALCON_MAILBOX1.with_base(base).read()
 
-  def disable_ctx_req(self, base:int):
+  def disable_ctx_req(self, base):
     self.nvdev.NV_PFALCON_FBIF_CTL.with_base(base).update(allow_phys_no_ctx=1)
     self.nvdev.NV_PFALCON_FALCON_DMACTL.with_base(base).write(0x0)
 
-  def reset(self, base:int, riscv=False):
+  def reset(self, base, riscv=False):
     engine_reg = self.nvdev.NV_PGSP_FALCON_ENGINE if base == self.falcon else self.nvdev.NV_PSEC_FALCON_ENGINE
     engine_reg.write(reset=1)
     time.sleep(0.1)
@@ -409,10 +408,10 @@ class NV_GSP(NV_IP):
       assert self.nvdev.flcn.frts_offset == m.frtsOffset, f"FRTS mismatch: {self.nvdev.flcn.frts_offset} != {m.frtsOffset}"
     self.wpr_meta, self.wpr_meta_sysmem = self.nvdev._alloc_boot_struct(m)
 
-  def promote_ctx(self, client:int, subdevice:int, obj:int, ctxbufs:dict[int, GRBufDesc], bufs=None, virt=None, phys=None):
+  def promote_ctx(self, client, subdevice, obj, ctxbufs, bufs=None, virt=None, phys=None):
     res, prom = {}, nv_gpu.NV2080_CTRL_GPU_PROMOTE_CTX_PARAMS(entryCount=len(ctxbufs), engineType=0x1, hChanClient=client, hObject=obj)
     for i,(buf,desc) in enumerate(ctxbufs.items()):
-      use_v, use_p = (desc.virt if virt is None else virt), (desc.phys if phys is None else phys)
+      use_v, use_p = (desc.v if virt is None else virt), (desc.p if phys is None else phys)
       x = (bufs or {}).get(buf, self.nvdev.mm.valloc(desc.size, contiguous=True)) # allocate buffers
       prom.promoteEntry[i] = nv_gpu.NV2080_CTRL_GPU_PROMOTE_CTX_BUFFER_ENTRY(bufferId=buf, gpuVirtAddr=x.va_addr if use_v else 0, bInitialize=use_p,
         gpuPhysAddr=x.paddrs[0][0] if use_p else 0, size=desc.size if use_p else 0, physAttr=0x4 if use_p else 0, bNonmapped=(use_p and not use_v))
@@ -450,11 +449,10 @@ class NV_GSP(NV_IP):
     gr_size = _ctx_info(nv_gpu.NV0080_CTRL_FIFO_GET_ENGINE_CONTEXT_PROPERTIES_ENGINE_ID_GRAPHICS, add=0x40000)
     patch_size = _ctx_info(nv_gpu.NV0080_CTRL_FIFO_GET_ENGINE_CONTEXT_PROPERTIES_ENGINE_ID_GRAPHICS_PATCH)
     cfgs_sizes = {x: _ctx_info(x + 14, align=(2 << 20) if x == 5 else None) for x in range(3, 11)} # indices 3–10 are mapped to 17–24
-    self.grctx_bufs = {0: GRBufDesc(gr_size, phys=True, virt=True), 1: GRBufDesc(patch_size, phys=True, virt=True, local=True),
-      2: GRBufDesc(patch_size, phys=True, virt=True), **{x: GRBufDesc(cfgs_sizes[x], phys=False, virt=True) for x in range(3, 7)},
-      9: GRBufDesc(cfgs_sizes[9], phys=True, virt=True), 10: GRBufDesc(cfgs_sizes[10], phys=True, virt=False),
-      11: GRBufDesc(cfgs_sizes[10], phys=True, virt=True)} # NOTE: 11 reuses cfgs_sizes[10]
-    self.promote_ctx(self.priv_root, subdev, ch_gpfifo, {k:v for k, v in self.grctx_bufs.items() if not v.local})
+    self.grctx_bufs = {0: GRBufDesc(gr_size, p=1, v=1), 1: GRBufDesc(patch_size, p=1, v=1, lc=1), 2: GRBufDesc(patch_size, p=1, v=1),
+      **{x: GRBufDesc(cfgs_sizes[x], p=0, v=1) for x in range(3, 7)}, 9: GRBufDesc(cfgs_sizes[9], p=1, v=1),
+      10: GRBufDesc(cfgs_sizes[10], p=1, v=0), 11: GRBufDesc(cfgs_sizes[10], p=1, v=1)} # NOTE: 11 reuses cfgs_sizes[10]
+    self.promote_ctx(self.priv_root, subdev, ch_gpfifo, {k:v for k, v in self.grctx_bufs.items() if v.lc == 0})
 
     self.rpc_rm_alloc(hParent=ch_gpfifo, hClass=self.compute_class, params=None)
     self.rpc_rm_alloc(hParent=ch_gpfifo, hClass=self.dma_class, params=None)
@@ -475,7 +473,7 @@ class NV_GSP(NV_IP):
 
   ### RPCs
 
-  def rpc_rm_alloc(self, hParent:int, hClass:int, params:Any, client=None) -> int:
+  def rpc_rm_alloc(self, hParent, hClass, params, client=None) -> int:
     if hClass == self.gpfifo_class:
       ramfc_alloc = self.nvdev.mm.valloc(0x1000, contiguous=True)
       params.ramfcMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=ramfc_alloc.paddrs[0][0], size=0x200, addressSpace=2, cacheAttrib=0)
@@ -501,7 +499,7 @@ class NV_GSP(NV_IP):
       self.promote_ctx(client, self.subdevice, hParent, {k:v for k,v in self.grctx_bufs.items() if k in [0, 1, 2]}, phys_gr_ctx, phys=False)
     return obj if hClass != nv_gpu.NV1_ROOT else client
 
-  def rpc_rm_control(self, hObject:int, cmd:int, params:Any, client=None):
+  def rpc_rm_control(self, hObject, cmd, params, client=None):
     control_args = nv.rpc_gsp_rm_control_v(hClient=(client:=client or self.priv_root), hObject=hObject, cmd=cmd, flags=0x0,
       paramsSize=ctypes.sizeof(params) if params is not None else 0x0)
     self.cmd_q.send_rpc(nv.NV_VGPU_MSG_FUNCTION_GSP_RM_CONTROL, bytes(control_args) + (bytes(params) if params is not None else b''))
@@ -513,7 +511,7 @@ class NV_GSP(NV_IP):
       cast(nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS, st).workSubmitToken |= (1 << 30)
     return st
 
-  def rpc_set_page_directory(self, device:int, hVASpace:int, pdir_paddr:int, client=None, pasid=0xffffffff):
+  def rpc_set_page_directory(self, device, hVASpace, pdir_paddr, client=None, pasid=0xffffffff):
     params = nv.struct_NV0080_CTRL_DMA_SET_PAGE_DIRECTORY_PARAMS_v1E_05(physAddress=pdir_paddr,
       numEntries=self.nvdev.mm.pte_cnt[0], flags=0x8, hVASpace=hVASpace, pasid=pasid, subDeviceId=1, chId=0) # flags field is all channels.
     alloc_args = nv.rpc_set_page_directory_v(hClient=client or self.priv_root, hDevice=device, pasid=pasid, params=params)
@@ -546,7 +544,7 @@ class NV_GSP(NV_IP):
     header = nv.PACKED_REGISTRY_TABLE(size=hdr_size + len(entries_bytes) + len(data_bytes), numEntries=len(table))
     self.cmd_q.send_rpc(nv.NV_VGPU_MSG_FUNCTION_SET_REGISTRY, bytes(header) + entries_bytes + data_bytes)
 
-  def run_cpu_seq(self, seq_buf:memoryview):
+  def run_cpu_seq(self, seq_buf):
     hdr = nv.rpc_run_cpu_sequencer_v17_00.from_address(mv_address(seq_buf))
     cmd_iter = iter(seq_buf[ctypes.sizeof(nv.rpc_run_cpu_sequencer_v17_00):].cast('I')[:hdr.cmdIndex])
 

tinygrad/runtime/support/nv/nvdev.py

@@ -71,7 +71,7 @@ class NVMemoryManager(MemoryManager):
   def on_range_mapped(self): self.dev.NV_VIRTUAL_FUNCTION_PRIV_MMU_INVALIDATE.write((1 << 0) | (1 << 1) | (1 << 6) | (1 << 31))
 
 class NVDev(PCIDevImplBase):
-  def __init__(self, devfmt:str, mmio:MMIOInterface, vram:MMIOInterface, venid:int, subvenid:int, rev:int, bars:dict):
+  def __init__(self, devfmt, mmio:MMIOInterface, vram:MMIOInterface, venid:int, subvenid:int, rev:int, bars:dict):
     self.devfmt, self.mmio, self.vram, self.venid, self.subvenid, self.rev, self.bars = devfmt, mmio, vram, venid, subvenid, rev, bars
     self.lock_fd = System.flock_acquire(f"nv_{self.devfmt}.lock")
 
@@ -101,10 +101,10 @@ class NVDev(PCIDevImplBase):
     for ip in [self.gsp, self.flcn]: ip.fini_hw()
 
   def reg(self, reg:str) -> NVReg: return self.__dict__[reg]
-  def wreg(self, addr:int, value:int):
+  def wreg(self, addr, value):
     self.mmio[addr // 4] = value
     if NV_DEBUG >= 4: print(f"wreg: {hex(addr)} = {hex(value)}")
-  def rreg(self, addr:int) -> int: return self.mmio[addr // 4]
+  def rreg(self, addr): return self.mmio[addr // 4]
 
   def _early_init(self):
     self.reg_names:set[str] = set()
@@ -134,12 +134,12 @@ class NVDev(PCIDevImplBase):
 
     self.vram_size = self.reg("NV_PGC6_AON_SECURE_SCRATCH_GROUP_42").read() << 20
 
-  def _alloc_boot_struct(self, struct:ctypes.Structure) -> tuple[ctypes.Structure, int]:
+  def _alloc_boot_struct(self, struct):
     va, paddrs = System.alloc_sysmem(sz:=ctypes.sizeof(type(struct)), contiguous=True)
     to_mv(va, sz)[:] = bytes(struct)
     return type(struct).from_address(va), paddrs[0]
 
-  def _download(self, file:str) -> str:
+  def _download(self, file) -> str:
     url = f"https://raw.githubusercontent.com/NVIDIA/open-gpu-kernel-modules/8ec351aeb96a93a4bb69ccc12a542bf8a8df2b6f/{file}"
     return fetch(url, subdir="defines").read_text()
 

tinygrad/schedule/kernelize.py

@@ -344,7 +344,7 @@ pm_fuse = PatternMatcher([
 def do_fusion(x:UOp):
   found_contiguous = {}
   def gate_contiguous(x):
-    if is_contiguous:=(x.op is Ops.CONTIGUOUS): found_contiguous[x] = x.replace(src=(UOp(Ops.VIEW, arg=x.st), UOp.unique()))
+    if is_contiguous:=(x.op is Ops.CONTIGUOUS): found_contiguous[x] = x.replace(src=(UOp(Ops.VIEW, arg=x.st),))
     return not is_contiguous
   x.toposort(gate=gate_contiguous)
   del gate_contiguous

tinygrad/tensor.py

@@ -282,7 +282,7 @@ class Tensor(MathTrait):
     # TODO: this is a hack for writing to DISK. remove with working assign
     if isinstance(self.device, str) and self.device.startswith("DISK"):
       if x.__class__ is not Tensor: x = Tensor(x, device="CPU", dtype=self.dtype)
-      self._buffer().copyin(x._data())
+      cast(Buffer, self.contiguous().realize().uop.base.buffer).ensure_allocated().copyin(x._data())
       return self
     if x.__class__ is not Tensor: x = Tensor(x, device=self.device, dtype=self.dtype)
     if self.uop is x.uop: return self  # a self assign is a NOOP
@@ -299,10 +299,7 @@ class Tensor(MathTrait):
     """
     return Tensor(self.uop.detach(), device=self.device, requires_grad=False)
 
-  def _buffer(self) -> Buffer:
-    x = self.cast(self.dtype.base).contiguous()
-    if isinstance(self.device, tuple): x = x.to("CPU")
-    return cast(Buffer, x.realize().uop.base.buffer).ensure_allocated()
+  def _buffer(self) -> Buffer: return cast(Buffer, self.cast(self.dtype.base).contiguous().to("CPU").realize().uop.base.buffer)
   def _data(self) -> memoryview: return self._buffer().as_buffer()
 
   def data(self) -> memoryview:
@@ -2332,6 +2329,8 @@ class Tensor(MathTrait):
 
     NOTE: unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.
 
+    See: https://paperswithcode.com/method/average-pooling
+
     ```python exec="true" source="above" session="tensor" result="python"
     t = Tensor.arange(25).reshape(1, 1, 5, 5)
     print(t.avg_pool2d().numpy())
@@ -2378,6 +2377,8 @@ class Tensor(MathTrait):
 
     NOTE: unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.
 
+    See: https://paperswithcode.com/method/max-pooling
+
     ```python exec="true" source="above" session="tensor" result="python"
     t = Tensor.arange(25).reshape(1, 1, 5, 5)
     print(t.max_pool2d().numpy())
@@ -3006,6 +3007,8 @@ class Tensor(MathTrait):
     """
     Applies the Rectified Linear Unit (ReLU) function element-wise.
 
+    - Described: https://paperswithcode.com/method/relu
+
     ```python exec="true" source="above" session="tensor" result="python"
     print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).relu().numpy())
     ```
@@ -3042,6 +3045,7 @@ class Tensor(MathTrait):
     Applies the Hardsigmoid function element-wise.
     NOTE: default `alpha` and `beta` values are taken from torch
 
+    - Described: https://paperswithcode.com/method/hard-sigmoid
     - See: https://pytorch.org/docs/stable/generated/torch.nn.functional.hardsigmoid.html
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3284,6 +3288,7 @@ class Tensor(MathTrait):
     """
     Applies the Exponential Linear Unit (ELU) function element-wise.
 
+    - Described: https://paperswithcode.com/method/elu
     - Paper: https://arxiv.org/abs/1511.07289v5
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3296,6 +3301,7 @@ class Tensor(MathTrait):
     """
     Applies the Continuously differentiable Exponential Linear Unit (CELU) function element-wise.
 
+    - Described: https://paperswithcode.com/method/celu
     - Paper: https://arxiv.org/abs/1704.07483
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3308,6 +3314,7 @@ class Tensor(MathTrait):
     """
     Applies the Scaled Exponential Linear Unit (SELU) function element-wise.
 
+    - Described: https://paperswithcode.com/method/selu
     - Paper: https://arxiv.org/abs/1706.02515v5
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3332,6 +3339,7 @@ class Tensor(MathTrait):
     """
     Applies the Sigmoid Linear Unit (SiLU) function element-wise.
 
+    - Described: https://paperswithcode.com/method/silu
     - Paper: https://arxiv.org/abs/1606.08415
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3344,6 +3352,7 @@ class Tensor(MathTrait):
     """
     Applies the ReLU6 function element-wise.
 
+    - Described: https://paperswithcode.com/method/relu6
     - Paper: https://arxiv.org/abs/1704.04861v1
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3356,6 +3365,7 @@ class Tensor(MathTrait):
     """
     Applies the Hardswish function element-wise.
 
+    - Described: https://paperswithcode.com/method/hard-swish
     - Paper: https://arxiv.org/abs/1905.02244v5
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3440,6 +3450,8 @@ class Tensor(MathTrait):
     """
     Applies the Hardtanh function element-wise.
 
+    - Described: https://paperswithcode.com/method/hardtanh-activation
+
     ```python exec="true" source="above" session="tensor" result="python"
     print(Tensor([-1.5, -1.0, -0.5, 0., 0.5, 1.0, 1.5]).hardtanh().numpy())
     ```
@@ -3464,6 +3476,7 @@ class Tensor(MathTrait):
     """
     Applies the Gaussian Error Linear Unit (GELU) function element-wise.
 
+    - Described: https://paperswithcode.com/method/gelu
     - Paper: https://arxiv.org/abs/1606.08415v5
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3476,6 +3489,8 @@ class Tensor(MathTrait):
     """
     Applies the Sigmoid GELU approximation element-wise.
 
+    - Described: https://paperswithcode.com/method/gelu
+
     ```python exec="true" source="above" session="tensor" result="python"
     print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).quick_gelu().numpy())
     ```
@@ -3486,6 +3501,8 @@ class Tensor(MathTrait):
     """
     Applies the Leaky ReLU function element-wise.
 
+    - Described: https://paperswithcode.com/method/leaky-relu
+
     ```python exec="true" source="above" session="tensor" result="python"
     print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).leaky_relu().numpy())
     ```
@@ -3499,6 +3516,7 @@ class Tensor(MathTrait):
     """
     Applies the Mish function element-wise.
 
+    - Described: https://paperswithcode.com/method/mish
     - Paper: https://arxiv.org/abs/1908.08681v3
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3511,6 +3529,8 @@ class Tensor(MathTrait):
     """
     Applies the Softplus function element-wise.
 
+    - Described: https://paperswithcode.com/method/softplus
+
     ```python exec="true" source="above" session="tensor" result="python"
     print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).softplus().numpy())
     ```
@@ -3521,6 +3541,8 @@ class Tensor(MathTrait):
     """
     Applies the Softsign function element-wise.
 
+    - Described: https://paperswithcode.com/method/softsign
+
     ```python exec="true" source="above" session="tensor" result="python"
     print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).softsign().numpy())
     ```
@@ -3536,8 +3558,7 @@ class Tensor(MathTrait):
     # for each dimension, check either dim is 1, or it does not change
     if not all(resolve(s == ns) or resolve(s == 1) for s,ns in zip(shape, new_shape)):
       raise ValueError(f"cannot broadcast {self.shape} to {new_shape=}")
-    # NOTE: this cast is no-op in forward and uses sum_acc_dtype in the backward sum
-    return self.reshape(shape).cast(sum_acc_dtype(self.dtype))._apply_uop(UOp.expand, arg=new_shape).cast(self.dtype)
+    return self.reshape(shape)._apply_uop(UOp.expand, arg=new_shape)
 
   def _broadcasted(self, y:Tensor|ConstType|UOp, reverse:bool=False, match_dtype:bool=True) -> tuple[Tensor, Tensor]:
     x: Tensor = self
@@ -3814,6 +3835,7 @@ class Tensor(MathTrait):
     """
     Applies Layer Normalization over a mini-batch of inputs.
 
+    - Described: https://paperswithcode.com/method/layer-normalization
     - Paper: https://arxiv.org/abs/1607.06450v1
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3832,6 +3854,7 @@ class Tensor(MathTrait):
     """
     Applies Batch Normalization over a mini-batch of inputs.
 
+    - Described: https://paperswithcode.com/method/batch-normalization
     - Paper: https://arxiv.org/abs/1502.03167
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3856,6 +3879,7 @@ class Tensor(MathTrait):
 
     NOTE: dropout is only applied when `Tensor.training` is `True`.
 
+    - Described: https://paperswithcode.com/method/dropout
     - Paper: https://jmlr.org/papers/v15/srivastava14a.html
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -3897,6 +3921,7 @@ class Tensor(MathTrait):
     Computes scaled dot-product attention.
     `self` is the query tensor, `key` is the key tensor, and `value` is the value tensor.
 
+    - Described: https://paperswithcode.com/method/scaled
     - Paper: https://arxiv.org/abs/1706.03762v7
 
     ```python exec="true" source="above" session="tensor" result="python"
@@ -4089,8 +4114,8 @@ class Tensor(MathTrait):
     #extract singular values and sort. construct U from Q
     S, indices = U.square().sum(-2).sqrt().sort(dim = -1, descending=True)
     new_indices = Tensor.arange(num).reshape((1,) * (self.ndim - 1) + (num,)).expand(b_shape + 2 * (num,)).contiguous()
-    new_indices[..., :num] = indices.reshape(b_shape + (1,) + (num,)).expand(b_shape + 2 * (num,))
-    U,V = U.gather(-1, new_indices[...,0:num,0:num]) / S.unsqueeze(-2), V.gather(-1, new_indices[..., 0:num, 0:num]).realize()
+    new_indices[..., :num] = indices.reshape(b_shape + (1,) + (U.shape[0],)).expand(b_shape + 2 * (num,))
+    U,V = U.gather(-1, new_indices[...,0:num,0:num]) / S.unsqueeze(-2), V.gather(-1, new_indices[..., 0:num, 0:num])
 
     padded_u = Tensor.eye(q_num, dtype = U.dtype).reshape((1,) * (self.ndim - 2) + 2 * (q_num,)).expand(b_shape + 2 * (q_num,)).contiguous()
     padded_u[..., 0:num, 0:num] = U
@@ -4288,11 +4313,6 @@ class Tensor(MathTrait):
     """
     return self.cast(dtypes.bool)
 
-  def bfloat16(self) -> Tensor: return self.cast(dtypes.bfloat16)
-  def double(self) -> Tensor: return self.cast(dtypes.double)
-  def long(self) -> Tensor: return self.cast(dtypes.long)
-  def short(self) -> Tensor: return self.cast(dtypes.short)
-
   # *** image Tensor function replacements ***
 
   def image_dot(self, w:Tensor, dtype:DTypeLike|None=None) -> Tensor:

tinygrad/uop/ops.py

@@ -941,7 +941,6 @@ renderer = PatternMatcher([
   (UPat(Ops.BIND, src=UPat(Ops.NOOP), name="x"), lambda x: x.src[0]),
   #(UPat(Ops.BIND, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}[={x.src[1].arg}]")),
   (UPat(Ops.NEG, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"(-{x.src[0].arg})")),
-  (UPat(Ops.RECIP, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"(1/{x.src[0].arg})")),
   (UPat(Ops.MAX, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"max({x.src[0].arg}, {x.src[1].arg})")),
   (UPat(Ops.MULACC, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[0].arg}*{x.src[1].arg}+{x.src[2].arg})")),
   (UPat(Ops.WHERE, src=UPat(Ops.NOOP), name="x"), lambda x: UOp(Ops.NOOP, arg=f"({x.src[1].arg} if {x.src[0].arg} else {x.src[2].arg})")),

tinygrad/uop/symbolic.py

@@ -3,7 +3,7 @@ from typing import Any, Literal, cast
 import math, operator, struct, functools
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
-from tinygrad.dtype import ConstType, dtypes, PtrDType, AddrSpace, can_safe_cast
+from tinygrad.dtype import ConstType, dtypes, PtrDType, AddrSpace
 from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, cdiv, cmod, CORRECT_DIVMOD_FOLDING
 from tinygrad.uop.transcendental import xpow
 
@@ -65,8 +65,6 @@ symbolic_simple = PatternMatcher([
   (UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.arg)),
   (UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None),
   (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast),
-  # b.cast(a).cast(b) -> b if a preserves all values in b
-  (UPat.var('x').cast().named('a').cast().named('b'), lambda x,a,b: x if x.dtype == b.dtype and can_safe_cast(b.dtype, a.dtype) else None),
   # ** pow **
   (UPat.var("x").alu(Ops.POW, UPat.cvar("c", vec=False)), simplify_pow),
   # positive const ** x
@@ -429,6 +427,7 @@ sym = symbolic_flat+PatternMatcher([
   (UPat(Ops.WMMA, src=(UPat.var(), UPat.const(None, 0.0), UPat.var("acc"))), lambda acc: acc),
   # threefry + remove longs
   (UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32),
+  (UPat.var('x', dtypes.uint32).cast(dtypes.uint64).cast(dtypes.uint32), lambda x: x),   # cast there and back is noop (TODO: genericize)
   ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)),  # cast does truncation
   (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y',  dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
   (((UPat.var('x',  dtypes.uint64)*(1<<32)) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x),
@@ -467,7 +466,6 @@ sym = symbolic_flat+PatternMatcher([
       if any(x.op in REMOVE_FROM_SINK for x in root.src) else None),
   ((UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()),  # 1/(x^c) -> (1/x)^c
   ((UPat.var("x") * UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()*x.reciprocal()),
-  ((UPat.var("x") * UPat.cvar("c")).reciprocal(), lambda x,c: x.reciprocal()*c.reciprocal()), # 1/(x*c) -> (1/c)*(1/x)
   (UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")), lambda x,d: 1-d), # x*/(1+x) -> 1-1/(1+x)
   (UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")*UPat.var("y")), lambda x,y,d: y*(1-d)),
   (UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")+UPat.var("y")), lambda x,y,d: (1-d)+x*y),

tinygrad/viz/index.html

@@ -173,7 +173,7 @@
     background-color: #1a1b26;
     border: 1px solid #4a4b56;
     color: #f0f0f5;
-    border-radius: 4px;
+    border-radius: 8px;
     padding: 6px;
     cursor: pointer;
     height: 32px;
@@ -184,6 +184,7 @@
   }
   .btn:hover {
     background-color: #2a2b36;
+    border-color: #5a5b66;
   }
   .collapsed .container {
     display: none;
@@ -202,6 +203,7 @@
   pre code.hljs {
     overflow-y: auto;
     max-height: 30vh;
+    border-radius: 8px;
     padding: 8px;
   }
   .progress-message {
@@ -221,7 +223,6 @@
     pointer-events: none;
     display: none;
     font-size: 10px;
-    white-space: pre;
   }
   #device-list > div {
     min-height: 32px;
@@ -238,8 +239,6 @@
     padding: 0 8px;
     width: 100%;
     height: 100%;
-    max-height: 100vh;
-    overflow-x: auto;
   }
   .raw-text code {
     max-height: none !important;
@@ -248,12 +247,14 @@
     width: 100%;
     border-collapse: separate;
     border-spacing: 0;
+    border-radius: 8px;
+    overflow: hidden;
     background-color: #1a1b26;
     color: #f0f0f5;
     font-size: 0.95em;
   }
   table td {
-    border-bottom: 1px solid #4a4b56;
+    border-bottom: 1px solid #2c2f40;
     vertical-align: top;
   }
   table tr:last-child > td {
@@ -268,40 +269,6 @@
   tr.main-row > td, tr.sub-row > td {
     padding: 8px 12px;
   }
-  tr.code-row > td:first-child {
-    font-family: monospace;
-  }
-  td.pct-row > div {
-    height: 12px;
-    width: 100%;
-    display: flex;
-  }
-  td.pct-row > div > div {
-    height: 100%;
-  }
-  thead {
-    position: sticky;
-    top: 0;
-    z-index: 10;
-    background-color: #20222e;
-  }
-  thead th {
-    text-align: left;
-    padding: 10px 12px;
-    font-weight: 600;
-    border-bottom: 1px solid #4a4b56;
-    font-size: 0.95em;
-    letter-spacing: 0.03em;
-  }
-  .legend {
-    display: flex;
-    align-items: center;
-  }
-  .legend > div {
-    width: 0.95em;
-    height: 0.95em;
-    margin-right: 4px;
-  }
   </style>
 </head>
 <body>

tinygrad/viz/js/index.js

@@ -109,11 +109,11 @@ function formatTime(ts, dur=ts) {
 }
 const formatUnit = (d, unit="") => d3.format(".3~s")(d)+unit;
 
-const colorScheme = {TINY:["#1b5745", "#354f52", "#354f52", "#46acc2", "#1d2e62", "#63b0cd"],
-  DEFAULT:["#2b2e39", "#2c2f3a", "#31343f", "#323544", "#2d303a", "#2e313c", "#343746", "#353847", "#3c4050", "#404459", "#444862", "#4a4e65"],
-  BUFFER:["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"],
-  CATEGORICAL:["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"],}
-const cycleColors = (lst, i) => lst[i%lst.length];
+const devColors = {"TINY":["rgb(27 87 69)", "rgb(53 79 82)", "rgb(53 79 82)", "rgb(70 172 194)", "rgb(29, 46, 98)"],
+                   "DEFAULT":["rgb(29,31,42)","rgb(42,45,61)","rgb(55,59,79)","rgb(68,72,98)","rgb(18,19,26)","rgb(47,50,68)","rgb(59,63,84)","rgb(74,78,101)","rgb(24,26,35)","rgb(35,37,50)","rgb(49,53,72)","rgb(64,68,89)"],}
+const bufColors = ["#3A57B7","#5066C1","#6277CD","#7488D8","#8A9BE3","#A3B4F2"];
+
+const lighten = (rgb, depth, step=0.08) => rgb.replace(/\d+/g, n => Math.round(parseInt(n)+(255-parseInt(n)) * Math.min(1, depth*step)));
 
 var profileRet, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity;
 async function renderProfiler() {
@@ -153,10 +153,10 @@ async function renderProfiler() {
     for (const e of timeline.shapes) {
       if (e.depth === 0) colorKey = e.cat ?? e.name;
       if (!colorMap.has(colorKey)) {
-        const colors = colorScheme[k] ?? colorScheme.DEFAULT;
+        const colors = devColors[k] ?? devColors.DEFAULT;
         colorMap.set(colorKey, colors[colorMap.size%colors.length]);
       }
-      const fillColor = d3.color(colorMap.get(colorKey)).brighter(e.depth).toString();
+      const fillColor = lighten(colorMap.get(colorKey), e.depth);
       const label = parseColors(e.name).map(({ color, st }) => ({ color, st, width:ctx.measureText(st).width }));
       if (e.ref != null) ref = {ctx:e.ref, step:0};
       else if (ref != null) {
@@ -164,7 +164,7 @@ async function renderProfiler() {
         const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
         ref = stepIdx === -1 ? null : {ctx:ref.ctx, step:stepIdx};
       }
-      const arg = { tooltipText:formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...ref };
+      const arg = { tooltipText:formatTime(e.dur), ...ref };
       // offset y by depth
       data.shapes.push({x:e.st-st, y:offsetY+levelHeight*e.depth, width:e.dur, height:levelHeight, arg, label, fillColor });
     }
@@ -184,7 +184,7 @@ async function renderProfiler() {
       const y0 = e.y.map(yscale);
       const y1 = e.y.map(y => yscale(y+e.arg.nbytes));
       const arg = { tooltipText:`${e.arg.dtype} len:${formatUnit(e.arg.sz)}\n${formatUnit(e.arg.nbytes, "B")}` };
-      data.shapes.push({ x, y0, y1, arg, fillColor:cycleColors(colorScheme.BUFFER, i) });
+      data.shapes.push({ x, y0, y1, arg, fillColor:bufColors[i%bufColors.length] });
     }
     // lastly, adjust device rect by number of levels
     div.style.height = `${Math.max(levelHeight*timeline.maxDepth, baseHeight)+area+padding}px`;
@@ -377,16 +377,11 @@ function codeBlock(st, language, { loc, wrap }={}) {
   return ret;
 }
 
-function appendTd(tr, value, unit=null) {
-  const fmt = (typeof value === "number" && !Number.isInteger(value)) ? value.toFixed(2) : value;
-  tr.appendChild(document.createElement("td")).innerText = unit == "us" ? formatTime(value) : fmt+(unit ?? "");
-}
-
-function appendRow(table, name, value, unit=null, cls="main-row") {
+function appendRow(table, name, value, unit, cls) {
   const tr = table.appendChild(document.createElement("tr"));
   tr.className = cls;
   tr.appendChild(document.createElement("td")).innerText = name;
-  appendTd(tr, value, unit);
+  tr.appendChild(document.createElement("td")).innerText = unit === "us" ? formatTime(value) : value.toFixed(2)+(unit != null ? " "+unit : "%");
   return tr;
 }
 
@@ -500,44 +495,10 @@ async function main() {
   if (ckey.startsWith("/disasm")) {
     if (!(ckey in cache)) cache[ckey] = ret = await (await fetch(ckey)).json();
     displayGraph("profiler");
+    document.querySelector(".metadata").innerHTML = "";
     const root = document.createElement("div");
     root.className = "raw-text";
-    const metadata = document.querySelector(".metadata");
-    metadata.innerHTML = "";
-    // detailed assembly view
-    if (ret.cols != null) {
-      const asm = root.appendChild(document.createElement("table"));
-      const thead = asm.appendChild(document.createElement("thead"));
-      const usage = {};
-      for (const c of ret.cols) thead.appendChild(document.createElement("th")).innerText = c;
-      for (const r of ret.rows) {
-        const tr = asm.appendChild(document.createElement("tr"));
-        tr.className = "main-row code-row";
-        for (const d of Object.values(r.data)) appendTd(tr, d);
-        const segmentsTd = tr.appendChild(document.createElement("td"));
-        segmentsTd.className = "pct-row";
-        const usageBar = segmentsTd.appendChild(document.createElement("div"));
-        for (const [k, {width, value}] of Object.entries(r.segs)) {
-          const seg = usageBar.appendChild(document.createElement("div"));
-          seg.style.width = width+"%";
-          seg.title = `${ret.segments[k]} ${value}`;
-          seg.style.background = cycleColors(colorScheme.CATEGORICAL, parseInt(k));
-          if (!(k in usage)) usage[k] = 0;
-          usage[k] += value;
-        }
-      }
-      const summary = metadata.appendChild(document.createElement("table"));
-      for (const [i,s] of ret.segments.entries()) {
-        const tr = summary.appendChild(document.createElement("tr"));
-        tr.className = "main-row";
-        const td = tr.appendChild(document.createElement("td"));
-        const div = td.appendChild(document.createElement("div"));
-        div.className = "legend";
-        div.appendChild(document.createElement("div")).style.background = cycleColors(colorScheme.CATEGORICAL, i);
-        div.appendChild(document.createElement("p")).textContent = s;
-        appendTd(tr, usage[i] ?? 0);
-      }
-    } else root.appendChild(codeBlock(ret.src, "x86asm"));
+    root.appendChild(codeBlock(ret.src, "x86asm"));
     return document.querySelector(".profiler").replaceChildren(root);
   }
   // ** UOp view (default)

tinygrad/viz/serve.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io
-import subprocess, ctypes
 from contextlib import redirect_stdout
 from decimal import Decimal
 from http.server import BaseHTTPRequestHandler
@@ -30,13 +29,10 @@ def get_metadata(keys:list[TracingKey], contexts:list[list[TrackedGraphRewrite]]
   for i,(k,v) in enumerate(zip(keys, contexts)):
     steps = [{"name":s.name, "loc":s.loc, "depth":s.depth, "match_count":len(s.matches), "code_line":printable(s.loc),
               "query":f"/ctxs?ctx={i}&idx={j}"} for j,s in enumerate(v)]
-    ret.append(r:={"name":k.display_name, "steps":steps})
+    if isinstance(k.ret, ProgramSpec): steps.append({"name":"View Disassembly", "query":f"/disasm?ctx={i}"})
+    ret.append(r:={"name":k.display_name, "fmt":k.fmt, "steps":steps})
     # use the first key to get runtime profiling data about this context
     if getenv("PROFILE_VALUE") >= 2 and k.keys: r["runtime_stats"] = get_runtime_stats(k.keys[0])
-    # program spec metadata
-    if isinstance(k.ret, ProgramSpec):
-      steps.append({"name":"View Disassembly", "query":f"/disasm?ctx={i}"})
-      r["fmt"] = k.ret.src
     for key in k.keys: ref_map[key] = i
   return ret
 
@@ -58,7 +54,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
   excluded: set[UOp] = set()
   for u in (toposort:=x.toposort()):
     # always exclude DEVICE/CONST/UNIQUE
-    if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE} and u is not x: excluded.add(u)
+    if u.op in {Ops.DEVICE, Ops.CONST, Ops.UNIQUE}: excluded.add(u)
     # only exclude CONST VIEW source if it has no other children in the graph
     if u.op is Ops.CONST and len(u.src) != 0 and all(cr.op is Ops.CONST for c in u.src[0].children if (cr:=c()) is not None and cr in toposort):
       excluded.update(u.src)
@@ -129,15 +125,12 @@ def timeline_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
     depth = next((i for i,level_et in enumerate(levels) if st>=level_et), len(levels))
     if depth < len(levels): levels[depth] = et
     else: levels.append(et)
-    name, cat, info = e.name, None, None
-    if (ref:=ref_map.get(name)) is not None:
-      name = ctxs[ref]["name"]
-      if isinstance(p:=contexts[0][ref].ret, ProgramSpec):
-        info = f"{p.estimates.ops/(t:=dur*1e3):.2f} GFLOPS {p.estimates.mem/t:4.1f}|{p.estimates.lds/t:.1f} GB/s"
+    name, cat = e.name, None
+    if (ref:=ref_map.get(name)) is not None: name = ctxs[ref]["name"]
     elif isinstance(e.name, TracingKey):
       name, cat = e.name.display_name, e.name.cat
       ref = next((v for k in e.name.keys if (v:=ref_map.get(k)) is not None), None)
-    shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat, "info":info})
+    shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat})
   return {"shapes":shapes, "maxDepth":len(levels)}
 
 def mem_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
@@ -192,36 +185,11 @@ def get_runtime_stats(key) -> list[dict]:
       ret.append({"device":e.device, "data":[{"name":"Duration", "value":float(e.en-e.st), "unit":"us"}]})
   return ret
 
-# ** Assembly analyzers
-
-def get_llvm_mca(asm:str, mtriple:str, mcpu:str) -> dict:
-  target_args = [f"-mtriple={mtriple}", f"-mcpu={mcpu}"]
-  # disassembly output can include headers / metadata, skip if llvm-mca can't parse those lines
-  data = json.loads(subprocess.check_output(["llvm-mca","-skip-unsupported-instructions=parse-failure","--json","-"]+target_args, input=asm.encode()))
-  cr = data["CodeRegions"][0]
-  rows:list = [{"data":[instr], "segs":{}} for instr in cr["Instructions"]]
-  for i,info in enumerate(cr["InstructionInfoView"]["InstructionList"]): rows[i]["data"].append(info["Latency"])
-  for d in cr["ResourcePressureView"]["ResourcePressureInfo"]:
-    i, r = d["InstructionIndex"], d["ResourceIndex"]
-    if i>len(rows)-1: continue
-    rows[i]["segs"][r] = rows[i]["segs"].get(r, 0)+d["ResourceUsage"]
-  # rescale segment width to 0-100
-  max_usage = max([sum(x["segs"].values()) for x in rows], default=0)
-  for x in rows: x["segs"] = {k:{"width":(v/max_usage)*100, "value":v} for k,v in x["segs"].items()}
-  return {"rows":rows, "cols":["Opcode", "Latency", "HW Resources"], "segments":data["TargetInfo"]["Resources"]}
-
 def get_disassembly(ctx:list[str]):
   if not isinstance(prg:=contexts[0][int(ctx[0])].ret, ProgramSpec): return
-  lib = (compiler:=Device[prg.device].compiler).compile(prg.src)
-  with redirect_stdout(buf:=io.StringIO()): compiler.disassemble(lib)
-  disasm_str = buf.getvalue()
-  from tinygrad.runtime.ops_llvm import llvm, LLVMCompiler
-  if isinstance(compiler, LLVMCompiler):
-    mtriple = ctypes.string_at(llvm.LLVMGetTargetMachineTriple(tm:=compiler.target_machine)).decode()
-    mcpu = ctypes.string_at(llvm.LLVMGetTargetMachineCPU(tm)).decode()
-    ret = get_llvm_mca(disasm_str, mtriple, mcpu)
-  else: ret = {"src":disasm_str}
-  return json.dumps(ret).encode()
+  lib = Device[prg.device].compiler.compile(prg.src)
+  with redirect_stdout(buf:=io.StringIO()): Device[prg.device].compiler.disassemble(lib)
+  return json.dumps({"src":buf.getvalue()}).encode()
 
 # ** HTTP server