fix beam

work
beam in RKernel
2026-08-14 17:18:28 +00:00 · 2025-08-29 19:05:49 -07:00 · 2025-08-29 18:42:50 -07:00 · 2025-08-29 18:42:50 -07:00 · 2025-08-29 12:37:01 -07:00 · 2025-08-29 11:23:17 -07:00
45 changed files with 380 additions and 419 deletions
@@ -226,15 +226,10 @@ runs:
      if: inputs.ocelot == 'true' && runner.os == 'macOS'
      shell: bash
      run: |
-        pkgs=(cmake ninja llvm@15 zlib glew flex bison boost@1.85 zstd ncurses)
+        pkgs=(cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses)
        for f in "${pkgs[@]}"; do
          brew ls --versions "$f" >/dev/null 2>&1 || brew install --quiet "$f"
        done
-
-        # Fix boost 1.85 for gpuocelot
-        ln -s /opt/homebrew/opt/[email protected] /opt/homebrew/opt/boost || true
-        ln -s /opt/homebrew/opt/boost/lib/libboost_atomic-mt.dylib /opt/homebrew/opt/boost/lib/libboost_atomic.dylib || true
-        ln -s /opt/homebrew/opt/boost/lib/libboost_thread-mt.dylib /opt/homebrew/opt/boost/lib/libboost_thread.dylib || true
    - name: Cache gpuocelot
      if: inputs.ocelot == 'true'
      id: cache-build
@@ -253,8 +248,7 @@ runs:
        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
        mkdir build
        cd build
-        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF \
-          -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib -DCMAKE_POLICY_VERSION_MINIMUM=3.5
+        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5
        ninja
    - name: Install gpuocelot
      if: inputs.ocelot == 'true'
@@ -688,10 +688,6 @@ jobs:
      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
    - name: Test DISK copy time
      run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
-    - name: Test CPU copy time
-      run: |
-        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
-        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Run full CIFAR training w 1 GPU
      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
    # TODO: enable
@@ -749,10 +745,6 @@ jobs:
      run: NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test DISK copy time
      run: NV=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
-    - name: Test CPU copy time
-      run: |
-        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
-        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Test LLAMA-3
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
    - name: Run full CIFAR training w 1 GPU
@@ -54,12 +54,11 @@ confidence=
 # --enable=similarities". If you want to run only the classes checker, but have
 # no Warning level messages displayed, use"--disable=all --enable=classes
 # --disable=W"
-disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method,W0707
+disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method
 # E1101 for function binding
 # W0221 for Function class
 # W0105 for comment strings
 # E0401 for missing imports
-# W0707 for not reraising

 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option
@@ -2,6 +2,7 @@ import time
 start_tm = time.perf_counter()
 import math
 from typing import Tuple, cast
+import numpy as np
 from tinygrad import Tensor, nn, GlobalCounters, TinyJit, dtypes, Device
 from tinygrad.helpers import partition, trange, getenv, Context
 from extra.lr_scheduler import OneCycleLR
@@ -149,12 +150,13 @@ if __name__ == "__main__":
      acc.append((out.argmax(-1) == Y).sum() / eval_batchsize)
    return Tensor.stack(*loss).mean() / (batchsize*loss_batchsize_scaler), Tensor.stack(*acc).mean()

-  Tensor.manual_seed(1337)
-  num_train_samples = X_train.shape[0]
-
+  np.random.seed(1337)
  for epoch in range(math.ceil(hyp['misc']['train_epochs'])):
+    # TODO: move to tinygrad
    gst = time.perf_counter()
-    tidxs = Tensor.randperm(num_train_samples, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)
+    idxs = np.arange(X_train.shape[0])
+    np.random.shuffle(idxs)
+    tidxs = Tensor(idxs, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)  # NOTE: long doesn't fold
    train_loss:float = 0
    for epoch_step in (t:=trange(num_steps_per_epoch)):
      st = time.perf_counter()
@@ -758,27 +758,6 @@ def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0
      batch.append(tokens)
    yield Tensor.stack(batch, dim=0)

-def batch_load_llama3_small(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
-  if val:
-    dataset = BlendedGPTDataset([
-      base_dir / "c4-validation-91205-samples.en_text_document",
-    ], [
-      1.0
-    ], samples, seqlen, seed, False)
-  else:
-    dataset = BlendedGPTDataset([
-      base_dir / "c4-train.en_6_text_document",
-    ], [
-      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"
@@ -243,49 +243,31 @@ def eval_mrcnn():

 def eval_llama3():
  from extra.models.llama import Transformer
-  from examples.llama3 import MODEL_PARAMS, load, convert_from_huggingface
+  from examples.llama3 import MODEL_PARAMS
  from tinygrad.helpers import tqdm

-  BASEDIR = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
-  BS = getenv("BS", 4)
-  SMALL = getenv("SMALL", 0)
-  SEQLEN = getenv("SEQLEN", 8192)
-  MODEL_PATH = Path(getenv("MODEL_PATH", "/raid/weights/llama31_8b/"))
+  bs = 4
+  sequence_length = 512

-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
-  params = params | {"vocab_size": 32000} if not SMALL else params
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
-  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
-
-  # load weights
-  weights = load(str(MODEL_PATH / "model.safetensors.index.json"))
-  if "model.embed_tokens.weight" in weights:
-    print("converting from huggingface format")
-    weights = convert_from_huggingface(weights, params["n_layers"], params["n_heads"], params["n_kv_heads"])
-
-  load_state_dict(model, weights, strict=False, consume=True)
+  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().float()
+    return loss.flatten()

-  if SMALL:
-    from examples.mlperf.dataloader import batch_load_llama3_small
-    iter = batch_load_llama3_small(BS, 5760, SEQLEN, BASEDIR, val=True)
-  else:
-    from examples.mlperf.dataloader import batch_load_llama3
-    iter = batch_load_llama3(BS, 5760, SEQLEN, BASEDIR, val=True)
+  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):
+  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 = np.mean(losses)
-  print(f"Log Perplexity: {log_perplexity}")
+  log_perplexity = Tensor(losses).mean()
+  print(f"Log Perplexity: {log_perplexity.item()}")

 if __name__ == "__main__":
  # inference only
@@ -1290,14 +1290,12 @@ def train_llama3():
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup

  config = {}
-  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
  BS                 = config["BS"]                     = getenv("BS", 16)
  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)
-  SMALL              = config["SMALL"]                  = getenv("SMALL", 0)
  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
  EVAL_FREQ          = config["EVAL_FREQ"]              = getenv("EVAL_FREQ", 46080)
  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
@@ -1313,14 +1311,13 @@ def train_llama3():

  opt_gradient_clip_norm = 1.0
  opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
-  opt_learning_rate_decay_steps = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS)) - opt_learning_rate_warmup_steps
+  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
-  opt_end_learning_rate = getenv("END_LR", 8e-7)
+  opt_end_learning_rate = 8e-7

  # TODO: confirm weights are in bf16
  # vocab_size from the mixtral tokenizer
-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
-  params = params | {"vocab_size": 32000} if not SMALL else params
+  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)

@@ -1406,29 +1403,21 @@ def train_llama3():
  # ** data iters **
  def fake_data(bs, samples):
    for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)

  def get_train_iter():
    if getenv("FAKEDATA", 0):
      return fake_data(GBS, SAMPLES)
    else:
-      if SMALL:
-        from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
-      else:
-        from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+      from examples.mlperf.dataloader import batch_load_llama3
+      return batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))

  def get_eval_iter():
    if getenv("FAKEDATA", 0):
      return fake_data(EVAL_BS, 5760)
    else:
-      if SMALL:
-        from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
-      else:
-        from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
+      from examples.mlperf.dataloader import batch_load_llama3
+      return batch_load_llama3(EVAL_BS, 5760, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=True)

  iter = get_train_iter()
  i, sequences_seen = 0, 0
@@ -1437,7 +1426,7 @@ def train_llama3():
    GlobalCounters.reset()
    loss, lr = train_step(model, tokens, grad_acc)
    loss = loss.float().item()
-
+    # above as tqdm.write f-string
    tqdm.write(f"{loss:.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:
@@ -673,7 +673,6 @@ impl<'a> Thread<'a> {
                                        39 => f32::log2(s0),
                                        42 => 1.0 / s0,
                                        43 => 1.0 / s0,
-                                        46 => 1.0 / f32::sqrt(s0),
                                        51 => f32::sqrt(s0),
                                        _ => todo_instr!(instruction)?,
                                    }
@@ -1247,7 +1246,7 @@ impl<'a> Thread<'a> {
                            }

                            let ret = match op {
-                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 430 | 531 | 537 | 540 | 551 | 567 | 796 => {
+                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 531 | 537 | 540 | 551 | 567 | 796 => {
                                    let s0 = f32::from_bits(s0).negate(0, neg).absolute(0, abs);
                                    let s1 = f32::from_bits(s1).negate(1, neg).absolute(1, abs);
                                    let s2 = f32::from_bits(s2).negate(2, neg).absolute(2, abs);
@@ -1259,7 +1258,6 @@ impl<'a> Thread<'a> {
                                        272 => f32::max(s0, s1),
                                        299 => f32::mul_add(s0, s1, f32::from_bits(self.vec_reg[vdst])),
                                        426 => s0.recip(),
-                                        430 => 1.0 / f32::sqrt(s0),
                                        531 => f32::mul_add(s0, s1, s2),
                                        537 => f32::min(f32::min(s0, s1), s2),
                                        540 => f32::max(f32::max(s0, s1), s2),
@@ -2627,14 +2625,6 @@ mod test_vop1 {
        assert_eq!(thread.vec_reg[3], 1071644672);
    }

-    #[test]
-    fn test_v_rsq_f32() {
-        let mut thread = _helper_test_thread();
-        thread.vec_reg[0] = f32::to_bits(4.0);
-        r(&vec![0x7E005D00, END_PRG], &mut thread);
-        assert_eq!(f32::from_bits(thread.vec_reg[0]), 0.5);
-    }
-
    #[test]
    fn test_v_frexp_exp_i32_f64() {
        [(3573412790272.0, 42), (69.0, 7), (2.0, 2), (f64::NEG_INFINITY, 0)]
@@ -128,10 +128,7 @@ def cuModuleUnload(hmod) -> int:
 def cuLaunchKernel(f, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, sharedMemBytes: int,
                   hStream: Any, kernelParams: Any, extra: Any) -> int:
  cargs = [ctypes.cast(getattr(extra, field[0]), ctypes.c_void_p) for field in extra._fields_]
-  try: gpuocelot_lib.ptx_run(ctypes.cast(f.value, ctypes.c_char_p), len(cargs), (ctypes.c_void_p*len(cargs))(*cargs), lx, ly, lz, gx, gy, gz, 0)
-  except Exception as e:
-    print("Error in cuLaunchKernel:", e)
-    return orig_cuda.CUDA_ERROR_LAUNCH_FAILED
+  gpuocelot_lib.ptx_run(ctypes.cast(f.value, ctypes.c_char_p), len(cargs), (ctypes.c_void_p*len(cargs))(*cargs), lx, ly, lz, gx, gy, gz, 0)
  return orig_cuda.CUDA_SUCCESS

 def cuDeviceComputeCapability(major, minor, dev: int) -> int:
@@ -97,10 +97,7 @@ class GPFIFO:
    cargs = [ctypes.cast(args[i], ctypes.c_void_p) for i in range(args_cnt)] + [ctypes.cast(vals[i], ctypes.c_void_p) for i in range(vals_cnt)]
    gx, gy, gz = qmd.cta_raster_width, qmd.cta_raster_height, qmd.cta_raster_depth
    lx, ly, lz = qmd.cta_thread_dimension0, qmd.cta_thread_dimension1, qmd.cta_thread_dimension2
-    try:
-      gpuocelot_lib.ptx_run(ctypes.cast(prg_addr, ctypes.c_char_p), args_cnt+vals_cnt,
-        (ctypes.c_void_p*len(cargs))(*cargs), lx, ly, lz, gx, gy, gz, 0)
-    except Exception as e: print("failed to execute:", e)
+    gpuocelot_lib.ptx_run(ctypes.cast(prg_addr, ctypes.c_char_p), args_cnt+vals_cnt, (ctypes.c_void_p*len(cargs))(*cargs), lx, ly, lz, gx, gy, gz, 0)
    if qmd.release0_enable:
      rel0 = to_mv(qmd.release0_address_lower + (qmd.release0_address_upper << 32), 0x10).cast('Q')
      rel0[0] = qmd.release0_payload_lower + (qmd.release0_payload_upper << 32)
@@ -1,9 +1,9 @@
 import unittest, numpy as np
 from tinygrad import Tensor, Device, TinyJit
-from tinygrad.helpers import Timing, CI, OSX, getenv
+from tinygrad.helpers import Timing, CI, OSX
 import multiprocessing.shared_memory as shared_memory

-N = getenv("NSZ", 256)
+N = 256
 class TestCopySpeed(unittest.TestCase):
  @classmethod
  def setUpClass(cls): Device[Device.DEFAULT].synchronize()
@@ -54,24 +54,22 @@ class TestCopySpeed(unittest.TestCase):
    @TinyJit
    def _do_copy(t): return t.to('CPU').realize()

-    t = Tensor.randn(N, N).contiguous().realize()
-    Device[Device.DEFAULT].synchronize()
+    t = Tensor.randn(N, N, 4).contiguous().realize()
    for _ in range(5):
-      with Timing(f"copy {Device.DEFAULT} -> CPU {t.nbytes()/(1024**2)}M:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
+      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
        x = _do_copy(t)
        Device[Device.DEFAULT].synchronize()
      np.testing.assert_equal(t.numpy(), x.numpy())

-  def testCopyCPUtoDefaultJit(self):
+  def testCopytoCPUtoDefaultJit(self):
    if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")

    @TinyJit
-    def _do_copy(x): return x.to(Device.DEFAULT).realize()
+    def _do_copy(x): return t.to(Device.DEFAULT).realize()

    for _ in range(5):
-      t = Tensor.randn(N, N, device="CPU").contiguous().realize()
-      Device["CPU"].synchronize()
-      with Timing(f"copy CPU -> {Device.DEFAULT} {t.nbytes()/(1024**2)}M:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
+      t = Tensor.randn(N, N, 4, device="CPU").contiguous().realize()
+      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
        x = _do_copy(t)
        Device[Device.DEFAULT].synchronize()
      np.testing.assert_equal(t.numpy(), x.numpy())
@@ -428,10 +428,5 @@ class TestOpsBFloat16(unittest.TestCase):
    data = [60000.0, 70000.0, 80000.0]
    np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())

-  def test_no_approximation(self):
-    data = [326.0, 339.0, 10603200512.0]
-    expected = torch.tensor(data, dtype=torch.bfloat16).sqrt().float().numpy()
-    np.testing.assert_allclose(Tensor(data, dtype=dtypes.bfloat16).sqrt().numpy(), expected)
-
 if __name__ == '__main__':
  unittest.main()
@@ -1,10 +1,9 @@
 import unittest, operator, math
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.dtype import DType
-from tinygrad.helpers import CI, getenv, AMD_LLVM
+from tinygrad.helpers import CI, getenv
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
-from tinygrad.runtime.ops_python import from_storage_scalar
 import numpy as np
 import pytest
 from hypothesis import given, strategies as strat, settings, HealthCheck
@@ -21,13 +20,13 @@ dtypes_bool = (dtypes.bool,)
 binary_operations = [operator.add, operator.sub, operator.mul, operator.lt, operator.eq]

 # TODO: LLVM comparing with nan is incorrect
-if (Device.DEFAULT == "LLVM") or (Device.DEFAULT == "AMD" and AMD_LLVM):
+if Device.DEFAULT == "LLVM" or getenv("AMD_LLVM", 0):
  binary_operations.remove(operator.lt)

 integer_binary_operations = binary_operations + [(Tensor.bitwise_xor, np.bitwise_xor), (Tensor.bitwise_and, np.bitwise_and),
                                                 (Tensor.bitwise_or, np.bitwise_or), operator.mod]
 unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin),
-                    (Tensor.sqrt, np.sqrt), (Tensor.reciprocal, np.reciprocal), (Tensor.cos, np.cos)]
+                    (Tensor.sqrt, np.sqrt), (Tensor.reciprocal, np.reciprocal)]

 # TODO: enable this (this is a dtype issue)
 #binary_operations.append(operator.truediv)
@@ -36,14 +35,13 @@ unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.
 #binary_operations += [(Tensor.maximum, np.maximum)]

 # TODO: CI CUDA segfaults on sin, WEBGPU sin is not precise enough for large numbers
-if (getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}) or Device.DEFAULT == "WEBGPU":
-  unary_operations.remove((Tensor.sin, np.sin))
-  unary_operations.remove((Tensor.cos, np.cos))
+if (getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}) or Device.DEFAULT == "WEBGPU": unary_operations.remove((Tensor.sin, np.sin))

 class ht:
  float64 = strat.floats(width=64, allow_subnormal=False)
  float32 = strat.floats(width=32, allow_subnormal=False)
  float16 = strat.floats(width=16, allow_subnormal=False)
+  bfloat16 = strat.floats(width=16, allow_subnormal=False)
  uint8 = strat.integers(0, 255)
  uint16 = strat.integers(0, 65535)
  uint32 = strat.integers(0, 2**32-1)
@@ -53,7 +51,6 @@ class ht:
  int32 = strat.integers(-2147483648, 2147483647)
  int64 = strat.integers(-9223372036854775808, 9223372036854775807)
  bool = strat.booleans()
-ht.bfloat16 = ht.uint16

 def universal_test(a, b, dtype, op):
  # The 'nan' cases only fail with Vulkan WebGPU backend (CI)
@@ -71,13 +68,11 @@ def universal_test(a, b, dtype, op):
 def universal_test_unary(a, dtype, op):
  if not isinstance(op, tuple): op = (op, op)
  ta = Tensor([a], dtype=dtype)
-  # TODO: cos does not match for large input
-  if op[0] == Tensor.cos and abs(a) > 100: return
  out: Tensor = op[0](ta)
  tensor_value = out.numpy()
  numpy_value = op[1](ta.numpy())
  if dtype in dtypes.floats:
-    atol, rtol = {dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2)}.get(dtype, (1e-6, 1e-5))
+    atol, rtol = {dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 1e-2)}.get(dtype, (1e-6, 1e-5))
    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
  else: np.testing.assert_equal(tensor_value, numpy_value)

@@ -110,8 +105,7 @@ class TestDTypeALU(unittest.TestCase):

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

  @given(ht.float32, strat.sampled_from(unary_operations))
  def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)
@@ -122,7 +116,7 @@ class TestDTypeALU(unittest.TestCase):

  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
  @given(ht.bfloat16, strat.sampled_from(unary_operations))
-  def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
+  def test_bfloat16_unary(self, a, op): universal_test_unary(a, dtypes.bfloat16, op)

  @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
  def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)
@@ -1,22 +0,0 @@
-import unittest
-from tinygrad import Tensor, Device
-from tinygrad.helpers import RANGEIFY
-from tinygrad.codegen.opt.kernel import Opt, OptOps
-from tinygrad.engine.realize import get_program
-
-@unittest.skipIf(RANGEIFY>0, "arg is partial contig in rangeify")
-class TestOpts(unittest.TestCase):
-  def test_opt_upcast(self):
-    opts = (Opt(OptOps.UPCAST, 0, 4),)
-    a = Tensor.empty(16)
-    b = Tensor.empty(16)
-    out = (a+b).contiguous(arg=opts)
-    s = out.schedule()
-    self.assertEqual(s[-1].ast.arg.opts_to_apply, opts)
-    if Device.DEFAULT in {"CPU", "GPU", "METAL"}:
-      prg = get_program(s[-1].ast)
-      self.assertIn('float4', prg.src)
-
-if __name__ == '__main__':
-  unittest.main()
-
@@ -192,26 +192,5 @@ class TestOuterworld(unittest.TestCase):

    self.assertTrue((manual==out).all().item())

-  def test_setitem_pyrange(self):
-    with Context(DEBUG=0):
-      t = Tensor.rand(10).realize()
-      o = Tensor.empty(10)
-    GlobalCounters.reset()
-    for i in range(10):
-      o[i] = t[i]
-    o.realize()
-    self.assertTrue((t==o).all().item())
-
-  @unittest.skip("TODO: fix this")
-  def test_setitem(self):
-    with Context(DEBUG=0):
-      t = Tensor.rand(10).realize()
-      o = Tensor.empty(10)
-    GlobalCounters.reset()
-    i = UOp.range(10, -1)
-    o[i] = t[i]
-    o.contiguous(i).realize()
-    self.assertTrue((t==o).all().item())
-
 if __name__ == '__main__':
  unittest.main()
@@ -1,9 +1,10 @@
 import unittest
-from tinygrad import Tensor, Variable, GlobalCounters
+from tinygrad import Tensor, Variable
 from tinygrad.shape.shapetracker import View
+from tinygrad.helpers import GlobalCounters
 from tinygrad.uop.ops import sym_infer
 from tinygrad.dtype import dtypes
-from tinygrad.device import is_dtype_supported
+from tinygrad.device import Device
 from examples.gpt2 import Attention
 import numpy as np

@@ -262,7 +263,7 @@ class TestSymbolicOps(unittest.TestCase):
      symbolic = a[:vi].bitcast(dtypes.uint8).reshape(expected.shape).numpy()
      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)

-  @unittest.skipUnless(is_dtype_supported(dtypes.uint64), "no uint64")
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "no uint64")
  def test_bitcast_up(self):
    a = Tensor.rand(10, 4)
    for i in range(1, 5):
@@ -1,7 +1,7 @@
 # basic self-contained tests of the external functionality of tinygrad
 import unittest, random
 from tinygrad import Tensor, Context, Variable, TinyJit, dtypes, Device, nn
-from tinygrad.helpers import IMAGE, CI, getenv
+from tinygrad.helpers import IMAGE, CI

 class TestTiny(unittest.TestCase):

@@ -27,7 +27,7 @@ class TestTiny(unittest.TestCase):
    out = Tensor.ones(256).contiguous().sum()
    self.assertEqual(out.item(), 256)

-  def test_gemm(self, N=getenv("GEMM_N", 64), out_dtype=dtypes.float):
+  def test_gemm(self, N=64, out_dtype=dtypes.float):
    a = Tensor.ones(N,N).contiguous()
    b = Tensor.eye(N).contiguous()
    lst = (out:=a@b).tolist()
@@ -36,14 +36,6 @@ class TestTiny(unittest.TestCase):
        self.assertEqual(lst[y][x], 1.0, msg=f"mismatch at ({y},{x})")
    if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)

-  def test_gemv(self, N=getenv("GEMV_N", 64), out_dtype=dtypes.float):
-    a = Tensor.ones(1,N).contiguous()
-    b = Tensor.eye(N).contiguous()
-    lst = (out:=a@b).tolist()
-    for x in range(N):
-      self.assertEqual(lst[0][x], 1.0, msg=f"mismatch at {x}")
-    if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
-
  # *** randomness ***

  def test_random(self):
@@ -445,10 +445,10 @@ class TestUOpGraph(unittest.TestCase):
    with Context(IGNORE_OOB=0):
      glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), src=(), arg=0)
      v = Variable("v", 0, 20)
-      st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v), UOp.const(dtypes.int, 0), UOp(Ops.IF, src=(v<16,))))
+      st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v, v<16), UOp.const(dtypes.int, 0)))
      to_uops_list([st0])

-      st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v), v, v<20))
+      st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v, v<20), v))
      with self.assertRaises(RuntimeError): to_uops_list([st1])

  def test_in_bounds_access_gated_local(self):
@@ -463,7 +463,7 @@ class TestUOpGraph(unittest.TestCase):

      gate = (gidx<400) & (lidx<8)

-      local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx), UOp.const(dtypes.uint, 1), UOp(Ops.IF, src=(lidx<8,))))
+      local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx, lidx<8), UOp.const(dtypes.uint, 1)))

      barrier = UOp(Ops.BARRIER, dtypes.void, (local_store,))
      if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))
@@ -177,31 +177,6 @@ class TestBoolUOps(TestUOps):
  def test_cmplt_bool(self): self._test_bop_bool_fxn(Ops.CMPLT, lambda a,b: a < b)
  def test_where_bool(self): self._test_top_bool_fxn(Ops.WHERE, lambda a,b,c: b if a else c)

-class TestSafeCast(TestUOps):
-  def test_cast_folds(self):
-    a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
-    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int32).simplify(), a)
-    self.assertEqual(a.cast(dtypes.double).cast(dtypes.int32).simplify(), a)
-    a = UOp.variable("a", 1, 10, dtype=dtypes.uint8)
-    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.uint8).simplify(), a)
-    self.assertEqual(a.cast(dtypes.uint32).cast(dtypes.uint8).simplify(), a)
-
-  def test_remove_intermediate_cast(self):
-    a = UOp.variable("a", 0., 100., dtype=dtypes.half)
-    self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float))
-    a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
-    # TODO: double preserves certain int dtypes
-    self.assertEqual(a.cast(dtypes.double).cast(dtypes.float).simplify(), a.cast(dtypes.float))
-    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int16).simplify(), a.cast(dtypes.int16))
-    a = UOp.variable("a", 1, 10, dtype=dtypes.uint8)
-    self.assertEqual(a.cast(dtypes.int64).cast(dtypes.int32).simplify(), a.cast(dtypes.int32))
-
-  def test_safe_cast_using_bounds(self):
-    a = UOp.variable("a", 1, 10, dtype=dtypes.uint64)
-    self.assertEqual(a.cast(dtypes.int16).cast(dtypes.int).simplify(), a.cast(dtypes.int))
-    a = UOp.variable("a", -10, 10, dtype=dtypes.int32)
-    self.assertEqual(a.cast(dtypes.int8).cast(dtypes.int64).simplify(), a.cast(dtypes.int64))
-
 class TestExecALU(TestUOps):
  def test_sqrt(self):
    self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)
@@ -100,11 +100,6 @@ class TestHelpers(unittest.TestCase):
        np.testing.assert_equal(dt.min, False)
        np.testing.assert_equal(dt.max, True)

-  def test_dtype_range_vec(self):
-    for dt in core_dtypes:
-      self.assertEqual(dt.min, dt.vec(4).min)
-      self.assertEqual(dt.max, dt.vec(4).max)
-
  def test_truncate_fp16(self):
    self.assertEqual(truncate_fp16(1), 1)
    self.assertEqual(truncate_fp16(65504), 65504)
@@ -618,4 +613,4 @@ class TestAutoCastType(unittest.TestCase):
    np.testing.assert_allclose(out.numpy(), tt.log_softmax(0).numpy(), rtol=1e-3)
    out = t.log_softmax(0, dtype=dtypes.float)
    self.assertEqual(out.dtype, dtypes.float)
-    np.testing.assert_allclose(out.numpy(), tt.log_softmax(0, dtype=torch.float).numpy(), rtol=1e-3)
+    np.testing.assert_allclose(out.numpy(), tt.log_softmax(0, dtype=torch.float).numpy(), rtol=1e-3)
@@ -208,16 +208,6 @@ class TestSymbolic(unittest.TestCase):
    self.assertEqual((Variable("x", -10, 0)%Variable("y", -10, -1))._min_max, (-9, 0))
    self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (-9, 0))

-  def test_range_div_its_symbolic_bound(self):
-    a = Variable("a", 1, 10)
-    ridx0 = UOp.range(a+2, 0)
-    self.helper_test_variable(ridx0//(a+2), 0, 0, "0")
-
-  def test_range_mod_its_symbolic_bound(self):
-    a = Variable("a", 1, 10)
-    ridx = UOp.range(a+2, 0)
-    self.helper_test_variable(ridx%(a+2), 0, 11, "ridx0")
-
  def test_div_min_max(self):
    self.helper_test_variable(Variable("a", 2, 7) // 2, 1, 3, "(a//2)")
    self.helper_test_variable(Variable("a", 0, 6) // 2, 0, 3, "(a//2)")
@@ -60,7 +60,7 @@ class TestVminVmaxProperties(unittest.TestCase):
  def test_vmin_vmax_variable_inside_special(self):
    uop = UOp(Ops.SPECIAL, dtypes.int, arg=('gidx0', UOp(Ops.DEFINE_VAR, dtypes.int, arg=('i', 1, 10))))
    self.assertEqual(uop.vmin, 0)
-    self.assertEqual(uop.vmax, 9)
+    self.assertEqual(uop.vmax, 10)

  def test_vmin_vmax_multiplication_0_inf(self):
    # vmin and vmax for multiplication with a variable
@@ -16,9 +16,10 @@ def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=Non

 # real VIZ=1 pickles these tracked values
 from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, _name_cnt
-traces = [(tracked_keys, tracked_ctxs, uop_fields)]
+from tinygrad.viz import serve
+serve.contexts = (tracked_keys, tracked_ctxs, uop_fields)
 from tinygrad.viz.serve import get_metadata, uop_to_json, get_details
-def get_viz_list(): return get_metadata(traces)
+def get_viz_list(): return get_metadata(tracked_keys, tracked_ctxs)

 class BaseTestViz(unittest.TestCase):
  def setUp(self):
@@ -141,8 +142,6 @@ class TestViz(BaseTestViz):
    z = UOp.const(dtypes.int, 0)
    alu = a*z
    exec_rewrite(alu, [sym])
-    lst = get_viz_list()
-    self.assertEqual(len(lst), 1)
    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)])
@@ -266,7 +265,7 @@ def option(i:int) -> int|None: return None if i == 0 else i-1
 def load_profile(lst:list[ProfileEvent]) -> dict:
  ret = get_profile(lst)
  u = TinyUnpacker(ret)
-  total_dur, global_peak, index_len, layout_len = u("<IQII")
+  dur, global_peak, index_len, layout_len = u("<IQII")
  strings, dtypes = json.loads(ret[u.offset:u.offset+index_len]).values()
  u.offset += index_len
  layout:dict[str, dict] = {}
@@ -274,19 +273,19 @@ def load_profile(lst:list[ProfileEvent]) -> dict:
    klen = u("<B")[0]
    k = ret[u.offset:u.offset+klen].decode()
    u.offset += klen
-    layout[k] = v = {"events":[]}
+    layout[k] = v = {"shapes":[]}
    event_type, event_count = u("<BI")
    if event_type == 0:
      for _ in range(event_count):
-        name, ref, st, dur, _ = u("<IIIfI")
-        v["events"].append({"name":strings[name], "ref":option(ref), "st":st, "dur":dur})
+        name, ref, st, dur, cat, _ = u("<IIIfBI")
+        v["shapes"].append({"name":strings[name], "ref":option(ref), "st":st, "dur":dur, "cat":option(cat)})
    else:
      v["peak"] = u("<Q")[0]
      for _ in range(event_count):
        alloc, ts, key = u("<BII")
-        if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-        else: v["events"].append({"event":"free", "ts":ts, "key":key})
-  return {"dur":total_dur, "peak":global_peak, "layout":layout}
+        if alloc: v["shapes"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
+        else: v["shapes"].append({"event":"free", "ts":ts, "key":key})
+  return {"dur":dur, "peak":global_peak, "layout":layout}

 class TestVizProfiler(unittest.TestCase):
  def test_perfetto_node(self):
@@ -295,7 +294,7 @@ class TestVizProfiler(unittest.TestCase):

    j = load_profile(prof)

-    dev_events = j['layout']['NV']['events']
+    dev_events = j['layout']['NV']['shapes']
    self.assertEqual(len(dev_events), 1)
    event = dev_events[0]
    self.assertEqual(event['name'], 'E_2')
@@ -311,16 +310,14 @@ class TestVizProfiler(unittest.TestCase):

    j = load_profile(prof)

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

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

-    self.assertEqual(j["dur"], (event2["st"]+event2["dur"])-event["st"])
-
  def test_perfetto_graph(self):
    prof = [ProfileDeviceEvent(device='NV', comp_tdiff=decimal.Decimal(-1000), copy_tdiff=decimal.Decimal(-100)),
            ProfileDeviceEvent(device='NV:1', comp_tdiff=decimal.Decimal(-500), copy_tdiff=decimal.Decimal(-50)),
@@ -336,18 +333,18 @@ class TestVizProfiler(unittest.TestCase):
    self.assertEqual(tracks[1], 'NV')
    self.assertEqual(tracks[2], 'NV:1')

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

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

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

@@ -379,7 +376,7 @@ class TestVizMemoryLayout(BaseTestViz):
    profile_ret = load_profile(Buffer.profile_events)
    ret = profile_ret["layout"][f"{a.device} Memory"]
    self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["events"]), 2)
+    self.assertEqual(len(ret["shapes"]), 2)

  def test_del_once(self):
    a = _alloc(1)
@@ -388,7 +385,7 @@ class TestVizMemoryLayout(BaseTestViz):
    profile_ret = load_profile(Buffer.profile_events)
    ret = profile_ret["layout"][f"{b.device} Memory"]
    self.assertEqual(ret["peak"], 1)
-    self.assertEqual(len(ret["events"]), 3)
+    self.assertEqual(len(ret["shapes"]), 3)

  def test_alloc_free(self):
    a = _alloc(1)
@@ -398,7 +395,7 @@ class TestVizMemoryLayout(BaseTestViz):
    profile_ret = load_profile(Buffer.profile_events)
    ret = profile_ret["layout"][f"{c.device} Memory"]
    self.assertEqual(ret["peak"], 2)
-    self.assertEqual(len(ret["events"]), 4)
+    self.assertEqual(len(ret["shapes"]), 4)

 if __name__ == "__main__":
  unittest.main()
@@ -2,7 +2,7 @@ from typing import Any, cast
 import functools, operator, itertools
 from collections import defaultdict
 from dataclasses import dataclass
-from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace
+from tinygrad.dtype import dtypes, ImageDType, PtrDType, DType, AddrSpace
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, graph_rewrite, GroupOp, identity_element
 from tinygrad.uop.symbolic import split_uop, uop_given_valid, parse_valid, simplify_valid, sym, symbolic_flat
 from tinygrad.helpers import getenv, flatten, AMX, prod, partition
@@ -80,6 +80,9 @@ def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None):
    if len(midx.src[i].src) == 3: root_src = (midx.src[i].src[2], root_src)
    offsets_rootsrc[root_src].setdefault(arg, []).append(i)

+  # the buf.dtype is always a pointer
+  ptrdtype = cast(PtrDType, buf.dtype)
+
  # then rewrite everything we can into groups
  ret = []
  idxs: list[int|None] = [None]*vec.dtype.count
@@ -89,7 +92,7 @@ def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None):
    for grp in grouped_offsets:
      # get the index offset for this element. using [0] is okay, because they are the same
      lidx = midx.src[offsets[grp[0]][0]]
-      if len(grp) > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(len(grp)).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
+      if len(grp) > 1: lidx = lidx.cast(ptrdtype.base.vec(len(grp)).ptr(size=ptrdtype.size, addrspace=ptrdtype.addrspace))
      # set the idxs of the output
      for i,g in enumerate(grp):
        for oo in offsets[g]: idxs[oo] = global_offset+i
@@ -98,7 +101,7 @@ def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None):
      global_offset += len(grp)
  assert None not in idxs, f"some idxs are missing {idxs}"
  # this base thing is for image, we want the CAT to be a normal pointer
-  post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(vec.dtype.count), tuple(ret))
+  post_cat = UOp(Ops.PTRCAT, ptrdtype.base.ptr(size=ptrdtype.size, addrspace=ptrdtype.addrspace).vec(vec.dtype.count), tuple(ret))
  return post_cat.gep(tuple(cast(list[int], idxs)))

 def cat_after_store(cat:UOp, data:UOp, sto:UOp):
@@ -151,7 +154,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
    must_divide = False
  elif buf.dtype.base != dtypes.float and buf.dtype.base != dtypes.half and not isinstance(buf.dtype, ImageDType):
    pass
-  elif buf.ptrdtype.addrspace == AddrSpace.REG:
+  elif cast(PtrDType, buf.dtype).addrspace == AddrSpace.REG:
    pass
  elif isinstance(buf.dtype, ImageDType):
    lengths = [4]
@@ -166,12 +169,13 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
  # split based on the fold lengths
  global_offset = 0
  ret = []
+  ptrdtype = cast(PtrDType, buf.dtype)
  while global_offset < sz:
    # with 1 at the end of the lengths list, this will always hit
    for fold_length in lengths:
      if global_offset+fold_length > sz: continue
      lidx = buf.index(idx.src[1] + global_offset, idx.src[2] if len(idx.src) > 2 else None)
-      if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
+      if fold_length > 1: lidx = lidx.cast(ptrdtype.base.vec(fold_length).ptr(size=ptrdtype.size, addrspace=ptrdtype.addrspace))
      if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))+ls.src[2:]))
      else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length)))
      global_offset += fold_length
@@ -229,7 +233,8 @@ def no_vectorized_alu(alu:UOp):
  return UOp(Ops.VECTORIZE, alu.dtype, alus)

 def no_vectorized_buf(buf:UOp):
-  return buf.replace(dtype=buf.ptrdtype.base.scalar().ptr(buf.ptrdtype.size*buf.ptrdtype.count, buf.ptrdtype.addrspace)).cast(buf.dtype)
+  dtype = cast(PtrDType, buf.dtype)
+  return buf.replace(dtype=dtype.base.scalar().ptr(dtype.size*dtype.count, dtype.addrspace)).cast(dtype)

 def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp):
  cnt = cast.dtype.count
@@ -22,17 +22,11 @@ def get_optimized_ast(ast:UOp, renderer:Renderer) -> UOp|None:
  # no shape, no opt
  if ast.src[0].st is None: return None
  new_arg = ast.arg
-  if new_arg is None:
+  if new_arg is None and not NOOPT and not BEAM:
    k = Kernel(ast, opts=renderer)
-    if not NOOPT:
-      if not k.apply_tensor_cores(USE_TC.value): k.apply_opts(hand_coded_optimizations(k))
-      if BEAM >= 1:
-        from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
-        kb = Kernel(ast, opts=renderer)
-        rawbufs = bufs_from_lin(kb, allocate=False)
-        k = beam_search(kb, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
+    if not k.apply_tensor_cores(USE_TC.value): k.apply_opts(hand_coded_optimizations(k))
    new_arg = KernelInfo(opts_to_apply=tuple(k.applied_opts))
-  elif len(new_arg.applied_opts): return None
+  elif new_arg is not None and len(new_arg.applied_opts): return None
  return Kernel(ast.replace(arg=None), opts=renderer).get_optimized_ast().replace(arg=new_arg)

 pm_get_optimization = PatternMatcher([
@@ -92,10 +92,6 @@ class Kernel:
    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)]

-    # 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):]:
-      raise RuntimeError(f"reduces are not at the end of the shape {self.full_shape} -> {self.output_shape}")
-
  def copy(self):
    ret = type(self).__new__(type(self))

@@ -449,6 +445,7 @@ class Kernel:
  def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms])

  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
+    if self.applied_opts: raise RuntimeError("not supported")
    @functools.cache
    def fixup_ast(op:UOp) -> UOp:
      ret = op.replace(src=tuple(fixup_ast(x) for x in op.src)) # noqa: F821
@@ -1,18 +1,197 @@
-from dataclasses import replace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo
-from tinygrad.helpers import colored
-from tinygrad.codegen.opt.kernel import axis_colors
+import math, itertools
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import UOp, Ops, sint, ssimplify, AxisType, KernelInfo, PatternMatcher, UPat, graph_rewrite
+from tinygrad.helpers import DEBUG, BEAM, getenv
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps, KernelOptError
+from tinygrad.renderer import Renderer
+from tinygrad.dtype import dtypes
+from tinygrad.device import Buffer

-def rename_sink(s:UOp):
-  if s.arg is not None and s.arg.name != "test": return None
+def flatten_range(r:UOp):
+  off = 2 if r.op is Ops.STORE else 1
+  rngs = r.src[off:]
+  if not len(rngs): return None
+  new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
+  return r.replace(src=r.src[:off]+tuple(new_rngs))

-  # get all ranges (sorted)
-  rngs = sorted([u for u in s.parents if u.op is Ops.RANGE], key=lambda x: x.arg[0:-1])
+pm_flatten_range = PatternMatcher([
+  # real ranges only
+  (UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
+])

-  # add name to kernel
-  name = "k" + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), axis_colors[x.arg[-1]]) for x in rngs])
-  return s.replace(arg=KernelInfo(name=name) if s.arg is None else replace(s.arg, name=name))
+class RKernel(Kernel):
+  def __init__(self, ast:UOp, opts:Renderer|None=None):
+    self.rng = sorted([u for u in ast.toposort() if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: x.arg)
+    super().__init__(ast, opts)
+    self.sts.clear()

-pm_postrange_opt = PatternMatcher([
-  (UPat(Ops.SINK, name="s"), rename_sink),
+    # convert LOOP to GLOBAL
+    self.replaces = {}
+    if self.opts.has_local:
+      store_rngs = self.ast.src[0].src[2:]
+
+      # filter any not in local stores
+      local_store_rngs = [x.ranges for x in self.ast.toposort() if (x.op is Ops.STORE and x.src[0].dtype.addrspace == AddrSpace.LOCAL) \
+                          or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
+      for ls in local_store_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])
+
+      store_rng = [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE] if store_rngs else []
+      rng = [x.replace(arg=(x.arg[0], AxisType.GLOBAL)) if x.arg[1] == AxisType.LOOP and x in store_rng else x for x in self.rng]
+      self.replaces.update(dict(zip(self.rng, rng)))
+      self.rng = rng
+
+      # NOTE: needed for tensor cores
+      self.substitute()
+
+    self.maxarg = max([x.arg[0] for x in self.rng]) if len(self.rng) else 0
+
+  def substitute(self) -> UOp:
+    self.ast = graph_rewrite(self.ast.substitute(self.replaces), pm_flatten_range)
+    self.replaces = {}
+    return self.ast
+
+  def copy(self):
+    self.substitute()
+    return RKernel(self.ast, self.opts)
+
+  # must be done earlier
+  def simplify_merge_adjacent(self): return
+
+  def apply_opt(self, opt:Opt, append_opt:bool=True) -> UOp|None:
+    if opt.op == OptOps.PADTO: raise KernelOptError("PAD is not supported yet. needs INVALID")
+    if opt.op == OptOps.SWAP: raise KernelOptError("SWAP is not supported yet")
+    return super().apply_opt(opt, append_opt)
+
+  def shift_to(self, axis:int, amount:int, new_type:AxisType, top:bool=False, insert_at:int|None=None):
+    old_sz = self.rng[axis].src[0].arg // amount
+    assert old_sz > 0, f"bad old_sz on {axis} {amount} {self.rng[axis]}"
+
+    self.maxarg += 1
+    new_rng = UOp.range(amount, self.maxarg, new_type)
+
+    if old_sz == 1:
+      self.replaces[self.rng[axis]] = new_rng
+      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
+      del self.rng[axis]
+    else:
+      replaced_rng = self.rng[axis].replace(src=(UOp.const(dtypes.int, old_sz),))
+      self.replaces[self.rng[axis]] = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
+      self.rng[axis] = replaced_rng
+      self.rng.insert(insert_at if insert_at is not None else len(self.rng), new_rng)
+    return new_rng
+
+  @property
+  def axis_types(self) -> list[AxisType]: return [x.arg[-1] for x in self.rng]
+  @property
+  def shape_len(self): return len(self.rng)
+
+  @property
+  def full_shape(self) -> tuple[sint, ...]: return tuple([ssimplify(x.src[0]) for x in self.rng])
+  @property
+  def output_shape(self) -> tuple[sint, ...]:
+    if self.ast.src[0].op is not Ops.STORE: return ()
+    return tuple([ssimplify(x.src[0]) for x in self.ast.src[0].src[2:]])
+
+  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
+    ret = self.substitute()
+    kernel_name = ret.arg.name if ret.arg is not None and ret.arg.name != "test" else self.name if name_override is None else name_override
+    rarg = KernelInfo(kernel_name, tuple(self.axis_types), self.dont_use_locals, tuple(self.applied_opts))
+    return ret.replace(arg=rarg)
+
+  # does nothing
+  @axis_types.setter
+  def axis_types(self, value): pass
+
+  def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> bool:
+    reduceops = [x for x in self.ast.toposort() if x.op is Ops.REDUCE]
+    if not len(reduceops): raise KernelOptError("no reduce ops for TensorCore")
+    reduceop = reduceops[0]
+    if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD:
+      mul = reduceop.src[0] if reduceop.src[0].op is not Ops.CAST else reduceop.src[0].src[0]
+      if mul.op is not Ops.MUL: return False
+      in0, in1 = mul.src
+      tensor_cores = self.opts.tensor_cores if tc_select == -1 else [self.opts.tensor_cores[tc_select]]
+      for tc in tensor_cores:
+        if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
+          # early realize for TC
+          self.substitute()
+
+          # tensor cores have three ranges. X, Y, and REDUCE
+          in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0])
+          in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0])
+          red_ranges = sorted(reduceop.src[1:], key=lambda x: x.arg[0])
+          if DEBUG >= 3:
+            print(f"TC({axis}): {[(x.arg[0],x.vmax+1) for x in in0_ranges]}",
+                              f"{[(x.arg[0],x.vmax+1) for x in in1_ranges]} {[(x.arg[0],x.vmax+1) for x in red_ranges]}")
+          if not len(in0_ranges) or not len(in1_ranges) or not len(red_ranges): return None
+
+          # pick ranges
+          # NOTE: why are in1 and in0 switched?
+          axis_choices = list(itertools.product(in1_ranges, in0_ranges, red_ranges))
+          if not (axis < len(axis_choices)): return None
+          axes = axis_choices[axis]
+
+          # do optimizations and save the ranges
+          try:
+            for i,a in enumerate(axes):
+              if a.src[0].divides(tc.dims[i]) is None:
+                self.apply_opt(Opt(OptOps.PADTO, self.rng.index(a), tc.dims[i]), append_opt=False) # PADTO might fail
+          except KernelOptError: continue
+          ne: list[UOp] = []
+          for opt in tc.opts:
+            ne.append(self.apply_opt(Opt({"u":OptOps.UPCAST, "l":OptOps.LOCAL}[opt[0]], axes[int(opt[1])].arg[0], 2), append_opt=False))
+          reduce_axis = [self.rng[i] for i in self.axes_of(AxisType.REDUCE)].index(axes[2])
+          for _, amt in tc.get_reduce_axes():
+            ne.append(self.apply_opt(Opt(OptOps.UNROLL, reduce_axis, amt), append_opt=False))
+
+          if use_tensor_cores != 2:
+            # fix the srcs
+            reduceop = [x for x in self.ast.toposort() if x.op is Ops.REDUCE][0]
+            tne = [x.replace(tag=1) for x in ne]
+            ret = reduceop.substitute(dict(zip(ne, tne)))
+            srcs = list((ret.src[0] if ret.src[0].op is not Ops.CAST else ret.src[0].src[0]).src)
+            srcs = [x.substitute(dict(zip(tne, [ne[i] for i in p]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
+
+            # get reduce/upcast axes for the tensor cores
+            tc_reduce_axes = self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))])
+            base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis(tc.base_upcast_axes())])
+            tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
+
+            # axes to range number (was done in lowerer)
+            tc_upcast_axes = tuple([tuple([(self.rng[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes])
+            tc_reduce_axes = tuple([self.rng[a].arg[0] for a in tc_reduce_axes])
+
+            # construct the op
+            # TODO: remove tc_upcast_axes from the arg
+            wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.opts.device, tc.threads, tc_upcast_axes, tc_reduce_axes)
+            wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
+              UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0], tag=1),
+              UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1], tag=1),
+              UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg, tag=1)
+            tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2], tag=1)
+
+            # preserve extra reduces
+            reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
+            if len(reduce_ranges): tc_uop = UOp(Ops.REDUCE, tc_uop.dtype, (tc_uop,)+tuple(reduce_ranges), Ops.ADD)
+            self.ast = self.ast.substitute({reduceop: tc_uop})
+          return True
+    return False
+
+def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
+  glbls = sorted([x for x in ast.parents if x.op is Ops.DEFINE_GLOBAL], key=lambda x: x.arg)
+  return [Buffer(dname, x.dtype.size, x.dtype.base) for x in glbls]
+
+def apply_ropt(ast:UOp, renderer:Renderer):
+  k = RKernel(ast, opts=renderer)
+  if BEAM >= 1:
+    from tinygrad.codegen.opt.search import beam_search
+    kb = RKernel(ast, opts=renderer)
+    rawbufs = bufs_from_ast(ast, renderer.device)
+    k = beam_search(kb, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
+  elif ast.arg is not None: k.apply_opts(ast.arg.opts_to_apply)
+  return k.get_optimized_ast()
+
+pm_postrange_opt = pm_flatten_range+PatternMatcher([
+  (UPat(Ops.SINK, name="ast"), lambda ctx,ast: apply_ropt(ast, ctx) if ast.arg is None or \
+   (ast.arg is not None and ast.arg.opts_to_apply is not None) else None),
 ])
@@ -2,7 +2,7 @@ from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
 from typing import Any, Generic, TypeVar, Iterator
-import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
+import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal, time
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, \
                             Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
@@ -138,14 +138,15 @@ class Buffer:
      if not self.device.startswith("DISK"): GlobalCounters.mem_used += self.nbytes
      if PROFILE:
        self._prof_num = num = len(Buffer.profile_events)
-        Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", num, {"dtype":self.dtype, "sz":self.size}))
+        ts = decimal.Decimal(time.perf_counter_ns())/1000
+        Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", ts, num, {"dtype":self.dtype, "sz":self.size}))
    return self
  def deallocate(self):
    assert hasattr(self, '_buf'), "buffer must be allocated to deallocate"
    if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}")
    if self._base is None and (self.options is None or self.options.external_ptr is None):
      if GlobalCounters is not None and not self.device.startswith("DISK"): GlobalCounters.mem_used -= self.nbytes
-      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self._prof_num))
+      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", decimal.Decimal(time.perf_counter_ns())/1000, self._prof_num))
      self.allocator.free(self._buf, self.nbytes, self.options)
    elif self._base is not None: self._base.allocated_views -= 1
    del self._buf
@@ -112,12 +112,12 @@ class dtypes:
  @staticmethod
  @functools.cache
  def min(dtype:DType):
-    if dtypes.is_int(dtype): return 0 if dtypes.is_unsigned(dtype) else -2**(dtype.scalar().itemsize*8-1)
+    if dtypes.is_int(dtype): return 0 if dtypes.is_unsigned(dtype) else -2**(dtype.itemsize*8-1)
    return -float("inf") if dtypes.is_float(dtype) else False
  @staticmethod
  @functools.cache
  def max(dtype:DType):
-    if dtypes.is_int(dtype): return 2**(dtype.scalar().itemsize*8)-1+dtypes.min(dtype)
+    if dtypes.is_int(dtype): return 2**(dtype.itemsize*8)-1+dtypes.min(dtype)
    return float("inf") if dtypes.is_float(dtype) else True
  @staticmethod
  def finfo(dtype:DType) -> tuple[int, int]:
@@ -198,9 +198,8 @@ def can_safe_cast(dt0:DType, dt1:DType) -> bool:
  # 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,
-      dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
-    case dtypes.float: return dt0 in (dtypes.half, dtypes.bfloat16, dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8)
+    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)
@@ -315,4 +314,4 @@ def _to_torch_dtype(dtype:DType) -> 'torch.dtype'|None:  # type: ignore [name-de
  except TypeError: return None
@functools.cache
 def _from_torch_dtype(torchdtype:'torch.dtype') -> DType: # type: ignore [name-defined] # noqa: F821
-  return {v:k for k in dtypes.all if (v:=_to_torch_dtype(k)) is not None}[torchdtype]
+  return {v:k for k in dtypes.all if (v:=_to_torch_dtype(k)) is not None}[torchdtype]
@@ -1,5 +1,5 @@
 from typing import cast, Generator, Callable
-import time, pprint, random, itertools, math
+import time, pprint, decimal, random, itertools, math
 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, PROFILE, ProfilePointEvent, cpu_events, prod
@@ -34,9 +34,8 @@ def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None)
    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
  try:
    uops = full_rewrite(ast, renderer)
-  except RuntimeError as e:
+  except RuntimeError:
    print("***** LINEARIZE FAILURE *****")
-    print(e)
    print(f"ast = {ast}")
    raise
  assert uops[-1].op is Ops.SINK, "last uop must be sink"
@@ -80,7 +79,7 @@ class CompiledRunner(Runner):
    self.p:ProgramSpec = p
    if precompiled is not None: self.lib = precompiled
    else:
-      with cpu_profile(TracingKey(f"compile {p.name}", (p.function_name,)), "TINY"):
+      with cpu_profile(TracingKey(f"compile {p.name}", (p.function_name,), cat="compiler"), "TINY"):
        self.lib = 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
@@ -162,7 +161,8 @@ class ExecItem:
  def run(self, _var_vals:dict[Variable, int]|None=None, wait=False, jit=False, do_update_stats=True) -> float|None:
    var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars)
    bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs]
-    if PROFILE: cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, {"metadata":self.metadata, "var_vals":var_vals}))
+    if PROFILE: cpu_events.append(ProfilePointEvent(self.prg.device, "exec", decimal.Decimal(time.perf_counter_ns())/1000, self.prg.display_name,
+                                                    {"metadata":self.metadata, "var_vals":var_vals}))
    et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2)
    if do_update_stats:
      GlobalCounters.kernel_count += 1
@@ -140,7 +140,7 @@ DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES = ContextVar("DONT_REALIZE_EXPAND", 0),
 QUANTIZE, VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("QUANTIZE", 0), ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, AMD_LLVM = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0), ContextVar("AMD_LLVM", 1)
-RANGEIFY, POSTOPT, FUSE_ATTENTION = ContextVar("RANGEIFY", 0), ContextVar("POSTOPT", 0), ContextVar("FUSE_ATTENTION", 0)
+RANGEIFY, POSTOPT, FUSE_ATTENTION = ContextVar("RANGEIFY", 0), ContextVar("POSTOPT", 1), ContextVar("FUSE_ATTENTION", 0)

@dataclass(frozen=True)
 class Metadata:
@@ -192,12 +192,12 @@ class Profiling(contextlib.ContextDecorator):
              colored(_format_fcn(fcn).ljust(50), "yellow"),
              colored(f"<- {(scallers[0][1][2]/tottime)*100:3.0f}% {_format_fcn(scallers[0][0])}", "BLACK") if scallers else '')

-def perf_counter_us() -> decimal.Decimal: return decimal.Decimal(time.perf_counter_ns())/1000

@dataclass(frozen=True)
 class TracingKey:
  display_name:str                       # display name of this trace event
  keys:tuple[Any, ...]=()                # optional keys to search for related traces
+  cat:str|None=None                      # optional category to color this by
  ret:Any=None

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

@dataclass(frozen=True)
-class ProfilePointEvent(ProfileEvent): device:str; name:str; key:Any; arg:dict=field(default_factory=dict); \
-    ts:decimal.Decimal=field(default_factory=perf_counter_us) # noqa: E702
+class ProfilePointEvent(ProfileEvent): device:str; name:str; ts:decimal.Decimal; key:Any; arg:dict=field(default_factory=dict) # noqa: E702

 cpu_events:list[ProfileEvent] = []
@contextlib.contextmanager
 def cpu_profile(name:str|TracingKey, device="CPU", is_copy=False, display=True) -> Generator[ProfileRangeEvent, None, None]:
-  res = ProfileRangeEvent(device, name, perf_counter_us(), is_copy=is_copy)
+  res = ProfileRangeEvent(device, name, decimal.Decimal(time.perf_counter_ns()) / 1000, is_copy=is_copy)
  try: yield res
  finally:
-    res.en = perf_counter_us()
+    res.en = decimal.Decimal(time.perf_counter_ns()) / 1000
    if PROFILE and display: cpu_events.append(res)

 # *** universal database cache ***
@@ -39,7 +39,7 @@ class Estimates:
        buf = u
        while len(buf.src): buf = buf.src[0]
        if buf.op is Ops.DEFINE_GLOBAL: # assume all DEFINE_GLOBAL memory is accessed
-          mem[(buf, u.op)] = buf.ptrdtype.size * buf.dtype.itemsize
+          mem[(buf, u.op)] = cast(PtrDType, buf.dtype).size * buf.dtype.itemsize
      if u.op is Ops.RANGE:
        mult_stack.append(mults)
        mults *= cast(sint, u.src[0].ssimplify())
@@ -2,7 +2,7 @@ from typing import Literal, Callable, cast
 import os, math, sys
 from collections import defaultdict, Counter
 from tinygrad.codegen.opt import tc
-from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, sint_to_uop
+from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat
 from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX
 from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate
 from tinygrad.renderer import Renderer
@@ -26,7 +26,7 @@ base_rewrite = PatternMatcher([
  (UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"{ctx.smem_align}{ctx.smem_prefix}{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
  (UPat(Ops.BARRIER), lambda ctx: ctx.barrier),
  (UPat(Ops.PRECAST, name="x"), lambda ctx,x: ctx[x.src[0]]),
-  (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0][0]](x.arg[0][-1])}; /* {sint_to_uop(x.arg[1]).render()} */"),
+  (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0][0]](x.arg[0][-1])}; /* {x.arg[1]} */"),
  # const
  (UPat(Ops.CONST, arg=math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x.dtype, ctx.infinity)})"),
  (UPat(Ops.CONST, arg=-math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x.dtype, f'-{ctx.infinity}')})"),
@@ -145,7 +145,7 @@ class CStyleLanguage(Renderer):
        if u.arg is not None: name = u.arg.function_name
        continue
      if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR):
-        r[u] = (f"data{u.arg}_{sz}" if (sz:=u.ptrdtype.size) > 0 else f"data{u.arg}") if u.op is Ops.DEFINE_GLOBAL else u.arg[0]
+        r[u] = (f"data{u.arg}_{sz}" if (sz:=cast(PtrDType, u.dtype).size) > 0 else f"data{u.arg}") if u.op is Ops.DEFINE_GLOBAL else u.arg[0]
        bufs[u] = (r[u], (u.dtype, False))
        continue

@@ -169,7 +169,7 @@ class CStyleLanguage(Renderer):

      if u.op in {Ops.ENDIF, Ops.ENDRANGE}: depth -= 1
      if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.INDEX, Ops.CUSTOMI} or \
-        (u.op is Ops.LOAD and u.src[0].ptrdtype.addrspace == AddrSpace.REG) or \
+        (u.op is Ops.LOAD and cast(PtrDType, u.src[0].dtype).addrspace == AddrSpace.REG) or \
        (u.op is Ops.CAST and isinstance(u.dtype, PtrDType)) or \
        (u.op in {Ops.VECTORIZE, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
        r[u] = l
@@ -196,20 +196,14 @@ class LLVMRenderer(Renderer):
 barrier = 'fence syncscope("workgroup") release\ntail call void @llvm.amdgcn.s.barrier()\nfence syncscope("workgroup") acquire\n'
 code_for_workitem = {"g": lambda x: f"tail call i32 @llvm.amdgcn.workgroup.id.{chr(120+int(x))}()",
                     "l": lambda x: f"tail call i32 @llvm.amdgcn.workitem.id.{chr(120+int(x))}()"}
-# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#llvm-ir-intrinsics
-# llvm.log2/llvm.exp2 don't support double
-llvm_intrinsics = {Ops.SQRT: "sqrt"}
 class AMDLLVMRenderer(LLVMRenderer):
  device = "AMD"
  has_local = True
  shared_max = AMDRenderer.shared_max
  global_max = AMDRenderer.global_max
  abi = "amdgpu_kernel"
-  code_for_op = {**LLVMRenderer.code_for_op, **{op: lambda: None for op in llvm_intrinsics}}
  string_rewrite = PatternMatcher([
    (UPat(Ops.SPECIAL, name="x"), lambda ctx, x: f"  {ctx[x]} = " + f"{ code_for_workitem[x.arg[0][0]](x.arg[0][-1])}; "),
-    (UPat(tuple(llvm_intrinsics), name="x"),
-    lambda ctx, x: f"  {ctx[x]} = call {ldt(x.dtype)} @llvm.{llvm_intrinsics[x.op]}.{ldt(x.dtype.scalar())}({ldt(x.src[0].dtype)} {ctx[x.src[0]]})"),
    (UPat(Ops.BARRIER), lambda ctx: barrier),
  ]) + base_rewrite
  extra_matcher = LLVMRenderer.extra_matcher + PatternMatcher([
@@ -190,7 +190,7 @@ class PTXRenderer(Renderer):
        r[u] = r[u.src[0]]
        continue
      if u.op is Ops.DEFINE_REG:
-        r[u] = [ssa("reg", u, self.types[u.dtype.base.scalar()]) for _ in range(u.ptrdtype.size)]
+        r[u] = [ssa("reg", u, self.types[u.dtype.base.scalar()]) for _ in range(cast(PtrDType, u.dtype).size)]
        continue
      if u.op in {Ops.INDEX, Ops.LOAD, Ops.STORE} and isinstance(u.src[0].dtype, PtrDType) and u.src[0].dtype.addrspace == AddrSpace.REG:
        if u.op is Ops.INDEX:
@@ -8,7 +8,7 @@ from tinygrad.dtype import DType, dtypes, ImageDType, PtrDType, truncate, float_
 from tinygrad.helpers import all_same, getenv, flatten, get_single_element
 from tinygrad.device import Compiled, Compiler, Allocator
 from tinygrad.codegen.opt import tc
-from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
+from tinygrad.uop.ops import exec_alu, Ops, UOp, GroupOp
 from tinygrad.renderer import Renderer

 def storage_fmt_for_dtype(dtype: DType): return 'H' if dtype == dtypes.bfloat16 else dtype.fmt
@@ -200,7 +200,7 @@ class PythonProgram:
          else: raise NotImplementedError(f"unimplemented tensor core {arg}")
        elif uop in GroupOp.ALU:
          assert all_same([len(x) for x in inp]), f"{[len(x) for x in inp]} doesn't match on {uop}"
-          assert all_same([dtype] + dtp) or uop in {*GroupOp.Comparison, Ops.WHERE}, f"dtype mismatch on {uop}"
+          assert all_same([dtype] + dtp) or uop in {Ops.CMPNE, Ops.CMPLT, Ops.WHERE}, f"dtype mismatch on {uop}"
          ul[i] = [exec_alu(uop, dtype, p) for p in zip(*inp)]
        assert i in ul, (uop, dtype, idp, arg)
        i += 1
@@ -208,7 +208,6 @@ class PythonProgram:

 class PythonRenderer(Renderer):
  device = "PYTHON"
-  code_for_op = python_alu
  def __init__(self):
    if getenv("EMULATE_METAL"): self.device, self.tensor_cores = "METAL", tc.metal
    if getenv("EMULATE_AMD"): self.device, self.tensor_cores = "AMD", tc.amd_rdna3
@@ -383,20 +383,15 @@ 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)

-    self.error_state:Exception|None = None # Exception if error is unrecoverable and sync will always fail
-
    if self._is_cpu(): HCQCompiled.cpu_devices.append(self)

  def synchronize(self):
-    if self.error_state is not None: raise self.error_state
-
    # 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:
-      self.error_state = e
      if hasattr(self, 'on_device_hang'): self.on_device_hang()
      else: raise e

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from tinygrad.uop.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, identity_element, resolve
-from tinygrad.uop.ops import track_rewrites, _substitute, KernelInfo
+from tinygrad.uop.ops import track_rewrites, _substitute
 from tinygrad.uop.spec import type_verify, tensor_uop_spec
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.helpers import Metadata, all_int, all_same, prod, dedup, unwrap, getenv, pluralize, FUSE_ARANGE, DEBUG, SPLIT_REDUCEOP
@@ -8,7 +8,6 @@ from tinygrad.dtype import ImageDType
 from tinygrad.schedule.multi import multi_pm
 from tinygrad.schedule.grouper import group_realizes, ALWAYS_CONTIGUOUS
 from tinygrad.codegen.opt.swizzler import merge_views, apply_swizzle, swizzle_reduceop
-from tinygrad.codegen.opt.kernel import Opt

 # creation can recurse a lot
 import sys
@@ -155,10 +154,6 @@ def unbind_view(x:UOp):
  return None

 replace_buffers = PatternMatcher([
-  # sink on contig creates a KernelInfo
-  (UPat(Ops.CONTIGUOUS, name="c").sink(name="s"),
-   lambda s,c: s.replace(src=(c.replace(arg=None),), arg=KernelInfo(opts_to_apply=c.arg)) \
-     if s.arg is None and c.arg is not None and isinstance(c.arg[0], Opt) else None),
  # replace ASSIGN with the target BUFFER
  (UPat(Ops.ASSIGN, src=(UPat((Ops.BUFFER, Ops.LOAD)), UPat(Ops.KERNEL)), name="assign", allow_any_len=True), lambda assign: assign.src[0]),
  # HACK: select the 0 branch of MSTACK (the device is wrong after this, is that okay?)
@@ -3099,7 +3099,6 @@ class Tensor(MathTrait):
    print(Tensor([0., math.pi/2, math.pi, 3*math.pi/2, 2*math.pi]).cos().numpy())
    ```
    """
-    if self.is_floating_point(): return ((math.pi/2)-self.cast(least_upper_dtype(self.dtype, dtypes.float32))).sin().cast(self.dtype)
    return ((math.pi/2)-self).sin()

  def tan(self) -> Tensor:
@@ -135,11 +135,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
  def tuplize(self:UOp) -> tuple:
    return (self.op.value, self.arg, self.dtype,)+tuple([x.tuplize for x in self.src])

-  @property
-  def ptrdtype(self) -> PtrDType:
-    if not isinstance(self.dtype, PtrDType): raise RuntimeError("ptrdtype called on UOp without PtrDType")
-    return self.dtype
-
  # *** uop shape stuff ***

  @functools.cached_property
@@ -168,7 +163,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    if self.op in {Ops.BUFFER, Ops.BUFFER_VIEW}: return ShapeTracker.from_shape((self.size,))
    if self.op is Ops.KERNEL: return ShapeTracker.from_shape((self.arg.ast.size,))
    if self.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
-      sz = self.ptrdtype.size
+      sz = cast(PtrDType, self.dtype).size
      return ShapeTracker.from_shape((sz,)) if sz > 0 else None

    # CONTIGUOUS with RANGE
@@ -560,7 +555,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    if self.op is Ops.BIND: return self.src[0]._min_max # ignore the bound value
    if self.op in {Ops.UNROLL, Ops.VECTORIZE}: return min(x.vmin for x in self.src), max(x.vmax for x in self.src)
    # TODO: Ops.SPECIAL is Ops.DEFINE_VAR
-    if self.op is Ops.SPECIAL: return 0, self.arg[1]-1 if isinstance(self.arg[1], int) else self.arg[1].vmax-1
+    if self.op is Ops.SPECIAL: return 0, self.arg[1]-1 if isinstance(self.arg[1], int) else self.arg[1].vmax
    if self.op is Ops.CONST: return self.arg, self.arg
    if self.op is Ops.VCONST: return (min(self.arg), max(self.arg))
    # TODO: CAST to bool/unsigned is not monotone, still some case can be simplified
@@ -837,7 +832,7 @@ def track_rewrites(name:Callable[..., str|TracingKey]|bool=True, replay:bool=Fal
    def __wrapper(*args, **kwargs):
      fn = key = func.__name__
      if TRACK_MATCH_STATS >= 2:
-        tracked_keys.append(key:=TracingKey(n:=f"{fn} n{next(_name_cnt.setdefault(fn, itertools.count(1)))}", (n,)))
+        tracked_keys.append(key:=TracingKey(n:=f"{fn} n{next(_name_cnt.setdefault(fn, itertools.count(1)))}", (n,), cat=fn))
        tracked_ctxs.append([])
      with cpu_profile(key, "TINY") as e:
        ret = func(*args, **kwargs)
@@ -845,7 +840,7 @@ def track_rewrites(name:Callable[..., str|TracingKey]|bool=True, replay:bool=Fal
        name_ret = name(*args, **kwargs, ret=ret)
        assert isinstance(name_ret, (TracingKey, str)), f"name function returned {type(name_ret)}"
        tracked_keys[-1] = k = TracingKey(n:=tracked_keys[-1].display_name.replace(fn, name_ret), (n,)) if isinstance(name_ret, str) else name_ret
-        e.name = TracingKey(k.display_name if isinstance(name_ret, str) else f"{fn} for {k.display_name}", k.keys)
+        e.name = TracingKey(k.display_name if isinstance(name_ret, str) else f"{fn} for {k.display_name}", k.keys, cat=fn)
      if getenv("CAPTURE_PROCESS_REPLAY") and replay:
        # find the unittest frame we're capturing in
        frm = sys._getframe(1)
@@ -907,7 +902,7 @@ if TRACK_MATCH_STATS or PROFILE:
    if TRACK_MATCH_STATS >= 2:
      with open(fn:=temp("rewrites.pkl", append_user=True), "wb") as f:
        print(f"rewrote {len(tracked_ctxs)} graphs and matched {sum(len(r.matches) for x in tracked_ctxs for r in x)} times, saved to {fn}")
-        pickle.dump([(tracked_keys, tracked_ctxs, uop_fields)], f)
+        pickle.dump((tracked_keys, tracked_ctxs, uop_fields), f)
    if VIZ: launch_viz(VIZ, temp("rewrites.pkl", append_user=True))
    if getenv("PRINT_MATCH_STATS", TRACK_MATCH_STATS.value):
      ret = [0,0,0.0,0.0]
@@ -970,7 +965,7 @@ class RewriteContext:
          for x in reversed(new_n.src): stack.append((x, 0, x))
        elif stage == 1:
          try: new_src = tuple([self.replace[x] for x in new_n.src])
-          except KeyError: raise RewriteNotReady
+          except KeyError: raise RewriteNotReady  # pylint: disable=raise-missing-from
          if new_src == new_n.src:
            # if top down, do the rewrite. if no rewrite or bottom up, we are done rewriting this node so we add it to the dict
            if self.pm is None or (new_src_n:=self.cached_pm_rewrite(new_n)) is None:
@@ -985,7 +980,7 @@ class RewriteContext:
        else:
          # in stage 2, we link the result of new_n to the result of n
          try: self.replace[n] = self.replace[new_n]
-          except KeyError: raise RewriteNotReady
+          except KeyError: raise RewriteNotReady  # pylint: disable=raise-missing-from
      except RewriteNotReady:
        # retry this later
        stack.insert(0, (n, stage, new_n))
@@ -120,7 +120,7 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
 # ***** uop type spec *****

 def validate_index(idx:UOp, gate:UOp=UOp.const(dtypes.bool, True)):
-  if IGNORE_OOB or isinstance(idx.dtype, ImageDType) or (sz := idx.src[0].ptrdtype.size) == -1: return True
+  if IGNORE_OOB or isinstance(idx.dtype, ImageDType) or (sz := cast(PtrDType, idx.src[0].dtype).size) == -1: return True
  # We can use UOp min/max to do a faster check, but it can give false positive since its not an exact bound and doesn't consider the mask
  if 0<=idx.src[1].vmin and idx.src[1].vmax<sz: return True
  mask = idx.src[2]&gate if len(idx.src)==3 else gate
@@ -68,10 +68,7 @@ symbolic_simple = PatternMatcher([
  (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(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_safe_cast(b.dtype, a.dtype) else None),
-  # if the intermediate cast doesnt narrow we can do it in one cast, we have to be carefull with bfloat16
-  (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_safe_cast(x.dtype, a.dtype) and
-    not (a.dtype==dtypes.float and (b.dtype==dtypes.bfloat16 or x.dtype==dtypes.bfloat16)) else None),
+  (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
@@ -287,8 +284,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
  ((UPat.var("y") + UPat.var("x")) + UPat.var("x"), lambda y,x: y+x*2),
  ((UPat.var("x") / UPat.var("x2")) / UPat.var("x3"), lambda x,x2,x3: x/(x2*x3) if x2 is not x3 else None), # (x/x2)/x3 -> x/(x2*x3)
  (-1 * (UPat.var("x") + UPat.cvar("c")), lambda x,c: (-x)+(-c)),  # -(x+c) -> -x + -c
-  (UPat.var('x', dtypes.ints).cast(dtypes.ints, name="a").cast(dtypes.ints, name="b"),
-    lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None),
  # a conditional with the same results either way is a noop, also fold const conditionals
  (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val),
  (UPat.cvar("gate", vec=False).where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
@@ -338,9 +333,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
  # div folding
  ((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d)
    if c.vmin>0 and d.vmin>0 and ((x.vmin>=0 and a.vmin>=0) or (x.vmax<=0 and a.vmax<=0)) else None),  # (x//c+a)//d -> (x+a*c)//(c*d)
-  # a range mod its own upper bound is just the range
-  (UPat(Ops.RANGE, src=UPat.var("end"), name="r")%UPat.var("end"), lambda r,end: r),
-  (UPat(Ops.RANGE, src=UPat.var("end"), name="r")//UPat.var("end"), lambda r,end: r.const_like(0)),
  (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.var("y"))), cancel_divmod),
  (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), fold_binary_numerator),
  (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), fold_divmod_congruence),
@@ -150,7 +150,7 @@
    inset: 0;
    z-index: 1;
  }
-  .profiler, .disasm {
+  .profiler {
    flex: 1 1 auto;
    min-width: 0;
    width: 100%;
@@ -346,7 +346,6 @@
    <div id="progress-message"></div>
    <div class="container ctx-list-parent"><div class="ctx-list"></div></div>
    <div class="view profiler"></div>
-    <div class="view disasm"></div>
    <div class="view graph">
      <svg id="graph-svg" preserveAspectRatio="xMidYMid meet">
        <g id="render">
@@ -191,7 +191,7 @@ async function renderProfiler() {
  canvas.addEventListener("wheel", e => (e.stopPropagation(), e.preventDefault()), { passive:false });
  const ctx = canvas.getContext("2d");
  const canvasTop = rect(canvas).top;
-  // color by key (name/device)
+  // color by key (name/category/device)
  const colorMap = new Map();
  data = {tracks:new Map(), axes:{}};
  const heightScale = d3.scaleLinear().domain([0, peak]).range([4,maxheight=100]);
@@ -210,7 +210,7 @@ async function renderProfiler() {
      data.tracks.set(k, { shapes, offsetY });
      let colorKey, ref;
      for (let j=0; j<eventsLen; j++) {
-        const e = {name:strings[u32()], ref:optional(u32()), st:u32(), dur:f32(), info:strings[u32()] || null};
+        const e = {name:strings[u32()], ref:optional(u32()), st:u32(), dur:f32(), cat:optional(u8()), info:strings[u32()] || null};
        // find a free level to put the event
        let depth = levels.findIndex(levelEt => e.st >= levelEt);
        const et = e.st+Math.trunc(e.dur);
@@ -218,8 +218,8 @@ async function renderProfiler() {
          depth = levels.length;
          levels.push(et);
        } else levels[depth] = et;
-        if (depth === 0) colorKey = e.name.split(" ")[0];
-        if (!colorMap.has(colorKey)) colorMap.set(colorKey, cycleColors(colorScheme[k.split(":")[0]] ?? colorScheme.DEFAULT, colorMap.size));
+        if (depth === 0) colorKey = e.cat ?? e.name;
+        if (!colorMap.has(colorKey)) colorMap.set(colorKey, cycleColors(colorScheme[k] ?? colorScheme.DEFAULT, colorMap.size));
        const fillColor = d3.color(colorMap.get(colorKey)).brighter(depth).toString();
        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};
@@ -299,15 +299,17 @@ async function renderProfiler() {
    ctx.clearRect(0, 0, canvas.clientWidth, canvas.clientHeight);
    // rescale to match current zoom
    const xscale = d3.scaleLinear().domain([0, dur]).range([0, canvas.clientWidth]);
-    const visibleX = xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale);
-    xscale.domain(visibleX);
-    const yscale = data.axes.y != null ? d3.scaleLinear().domain(data.axes.y.domain).range(data.axes.y.range) : null;
+    xscale.domain(xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale));
+    const zoomDomain = transform != null ? xscale.domain() : null;
+    let yscale = null;
+    if (data.axes.y != null) {
+      yscale = d3.scaleLinear().domain(data.axes.y.domain).range(data.axes.y.range);
+    }
    // draw shapes
    for (const [_, { offsetY, shapes }] of data.tracks) {
      for (const e of shapes) {
-        if (e.width == null) { start = e.x[0]; end = end = e.x[e.x.length-1]; }
-        else { start = e.x; end = e.x+e.width; }
-        if (start>visibleX[1] || end<visibleX[0]) continue;
+        const [start, end] = e.width != null ? [e.x, e.x+e.width] : [e.x[0], e.x[e.x.length-1]];
+        if (zoomDomain != null && (start>zoomDomain[1]|| end<zoomDomain[0])) continue;
        ctx.fillStyle = e.fillColor;
        // generic polygon
        if (e.width == null) {
@@ -587,7 +589,7 @@ async function main() {
  // ** Disassembly view
  if (ckey.startsWith("/disasm")) {
    if (!(ckey in cache)) cache[ckey] = ret = await (await fetch(ckey)).json();
-    displayGraph("disasm");
+    displayGraph("profiler");
    const root = document.createElement("div");
    root.className = "raw-text";
    const metadata = document.querySelector(".metadata");
@@ -629,7 +631,7 @@ async function main() {
        appendTd(tr, s.value);
      }
    } else root.appendChild(codeBlock(ret.src, "x86asm"));
-    return document.querySelector(".disasm").replaceChildren(root);
+    return document.querySelector(".profiler").replaceChildren(root);
  }
  // ** UOp view (default)
  // if we don't have a complete cache yet we start streaming rewrites in this step
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import multiprocessing, pickle, difflib, os, threading, json, time, sys, webbrowser, socket, argparse, socketserver, functools, codecs, io, struct
-import subprocess, ctypes, pathlib
+import subprocess, ctypes
 from contextlib import redirect_stdout
 from decimal import Decimal
 from http.server import BaseHTTPRequestHandler
@@ -27,22 +27,19 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
 # ** Metadata for a track_rewrites scope

 ref_map:dict[Any, int] = {}
-traces:dict[int, tuple] = {}
-def get_metadata(trace_bufs:list[tuple]) -> list[dict]:
+def get_metadata(keys:list[TracingKey], contexts:list[list[TrackedGraphRewrite]]) -> list[dict]:
  ret = []
-  for keys,contexts,uop_fields in trace_bufs:
-    for k,v in zip(keys, contexts):
-      traces[i:=len(traces)] = (k, v, uop_fields)
-      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})
-      # 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
+  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})
+    # 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

 # ** Complete rewrite details for a graph_rewrite call
@@ -94,16 +91,16 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
  return graph

@functools.cache
-def _reconstruct(a:int, i:int):
-  op, dtype, src, arg, *rest = traces[i][2][a]
-  arg = type(arg)(_reconstruct(arg.ast, i), arg.metadata) if op is Ops.KERNEL else arg
-  return UOp(op, dtype, tuple(_reconstruct(s, i) for s in src), arg, *rest)
+def _reconstruct(a:int):
+  op, dtype, src, arg, *rest = contexts[2][a]
+  arg = type(arg)(_reconstruct(arg.ast), arg.metadata) if op is Ops.KERNEL else arg
+  return UOp(op, dtype, tuple(_reconstruct(s) for s in src), arg, *rest)

-def get_details(ctx:TrackedGraphRewrite, i:int=0) -> Generator[GraphRewriteDetails, None, None]:
-  yield {"graph":uop_to_json(next_sink:=_reconstruct(ctx.sink, i)), "uop":str(next_sink), "changed_nodes":None, "diff":None, "upat":None}
+def get_details(ctx:TrackedGraphRewrite) -> Generator[GraphRewriteDetails, None, None]:
+  yield {"graph":uop_to_json(next_sink:=_reconstruct(ctx.sink)), "uop":str(next_sink), "changed_nodes":None, "diff":None, "upat":None}
  replaces: dict[UOp, UOp] = {}
  for u0_num,u1_num,upat_loc in tqdm(ctx.matches):
-    replaces[u0:=_reconstruct(u0_num, i)] = u1 = _reconstruct(u1_num, i)
+    replaces[u0:=_reconstruct(u0_num)] = u1 = _reconstruct(u1_num)
    try: new_sink = next_sink.substitute(replaces)
    except RuntimeError as e: new_sink = UOp(Ops.NOOP, arg=str(e))
    yield {"graph":(sink_json:=uop_to_json(new_sink)), "uop":str(new_sink), "changed_nodes":[id(x) for x in u1.toposort() if id(x) in sink_json],
@@ -139,39 +136,41 @@ def flatten_events(profile:list[ProfileEvent]) -> Generator[tuple[Decimal, Decim
 def timeline_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, scache:dict[str, int]) -> bytes|None:
  events:list[bytes] = []
  exec_points:dict[str, dict] = {}
+  category_enum:dict[str, int] = {}
  for st,et,dur,e in dev_events:
    if isinstance(e, ProfilePointEvent) and e.name == "exec": exec_points[e.key] = e.arg
    if dur == 0: continue
-    name, info = e.name, None
+    name, cat, info = e.name, None, None
    if (ref:=ref_map.get(name)) is not None:
      name = ctxs[ref]["name"]
-      if isinstance(p:=traces[ref][0].ret, ProgramSpec) and (ei:=exec_points.get(p.name)) is not None:
+      if isinstance(p:=contexts[0][ref].ret, ProgramSpec) and (ei:=exec_points.get(p.name)) is not None:
        info = f"{sym_infer(p.estimates.ops, ei['var_vals'])/(t:=dur*1e3):.2f} GFLOPS {sym_infer(p.estimates.mem, ei['var_vals'])/t:4.1f}"+ \
               f"|{sym_infer(p.estimates.lds,ei['var_vals'])/t:.1f} GB/s\n{ei['metadata']}"
    elif isinstance(e.name, TracingKey):
-      name = e.name.display_name
+      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)
-    events.append(struct.pack("<IIIfI", enum_str(name, scache), option(ref), st-start_ts, dur, enum_str(info or "", scache)))
+    events.append(struct.pack("<IIIfBI", enum_str(name, scache), option(ref), st-start_ts, dur,
+                              option(None if cat is None else enum_str(cat, category_enum)), enum_str(info or "", scache)))
  return struct.pack("<BI", 0, len(events))+b"".join(events) if events else None

-def mem_layout(dev_events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
+def mem_layout(events:list[tuple[int, int, float, DevEvent]], start_ts:int, end_ts:int, peaks:list[int], dtype_size:dict[str, int],
               scache:dict[str, int]) -> bytes|None:
  peak, mem = 0, 0
  temp:dict[int, int] = {}
-  events:list[bytes] = []
-  for st,_,_,e in dev_events:
+  bufs:list[bytes] = []
+  for st,_,_,e in events:
    if not isinstance(e, ProfilePointEvent): continue
    if e.name == "alloc":
-      events.append(struct.pack("<BIIIQ", 1, int(e.ts)-start_ts, e.key, enum_str(e.arg["dtype"].name, scache), e.arg["sz"]))
+      bufs.append(struct.pack("<BIIIQ", 1, int(e.ts)-start_ts, e.key, enum_str(e.arg["dtype"].name, scache), e.arg["sz"]))
      dtype_size.setdefault(e.arg["dtype"].name, e.arg["dtype"].itemsize)
      temp[e.key] = nbytes = e.arg["sz"]*e.arg["dtype"].itemsize
      mem += nbytes
      if mem > peak: peak = mem
    if e.name == "free":
-      events.append(struct.pack("<BII", 0, int(e.ts)-start_ts, e.key))
+      bufs.append(struct.pack("<BII", 0, int(e.ts)-start_ts, e.key))
      mem -= temp.pop(e.key)
  peaks.append(peak)
-  return struct.pack("<BIQ", 1, len(events), peak)+b"".join(events) if events else None
+  return struct.pack("<BIQ", 1, len(bufs), peak)+b"".join(bufs) if bufs else None

 def get_profile(profile:list[ProfileEvent]) -> bytes|None:
  # start by getting the time diffs
@@ -229,7 +228,7 @@ def get_llvm_mca(asm:str, mtriple:str, mcpu:str) -> dict:
  return {"rows":rows, "cols":["Opcode", "Latency", {"title":"HW Resources", "labels":resource_labels}], "summary":summary}

 def get_disassembly(ctx:list[str]):
-  if not isinstance(prg:=traces[int(ctx[0])][0].ret, ProgramSpec): return
+  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()
@@ -257,7 +256,7 @@ class Handler(BaseHTTPRequestHandler):
      except FileNotFoundError: status_code = 404
    elif (query:=parse_qs(url.query)):
      if url.path == "/disasm": ret, content_type = get_disassembly(**query), "application/json"
-      else: return self.stream_json(get_details(traces[i:=int(query["ctx"][0])][1][int(query["idx"][0])], i))
+      else: return self.stream_json(get_details(contexts[1][int(query["ctx"][0])][int(query["idx"][0])]))
    elif url.path == "/ctxs": ret, content_type = json.dumps(ctxs).encode(), "application/json"
    elif url.path == "/get_profile" and profile_ret: ret, content_type = profile_ret, "application/octet-stream"
    else: status_code = 404
@@ -292,17 +291,17 @@ def reloader():
      os.execv(sys.executable, [sys.executable] + sys.argv)
    time.sleep(0.1)

-def load_pickle(path:pathlib.Path|None) -> list:
-  if path is None or not path.exists(): return []
-  with path.open("rb") as f: return pickle.load(f)
+def load_pickle(path:str|None) -> list:
+  if path is None or not os.path.exists(path): return []
+  with open(path, "rb") as f: return pickle.load(f)

 # NOTE: using HTTPServer forces a potentially slow socket.getfqdn
 class TCPServerWithReuse(socketserver.TCPServer): allow_reuse_address = True

 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
-  parser.add_argument('--kernels', type=pathlib.Path, help='Path to kernels', default=None)
-  parser.add_argument('--profile', type=pathlib.Path, help='Path profile', default=None)
+  parser.add_argument('--kernels', type=str, help='Path to kernels', default=None)
+  parser.add_argument('--profile', type=str, help='Path profile', default=None)
  args = parser.parse_args()

  with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -313,9 +312,12 @@ if __name__ == "__main__":
  st = time.perf_counter()
  print("*** viz is starting")

-  ctxs = get_metadata(load_pickle(args.kernels))
+  contexts, profile = load_pickle(args.kernels), load_pickle(args.profile)

-  profile_ret = get_profile(profile:=load_pickle(args.profile))
+  # NOTE: this context is a tuple of list[keys] and list[values]
+  ctxs = get_metadata(*contexts[:2]) if contexts else []
+
+  profile_ret = get_profile(profile)

  server = TCPServerWithReuse(('', PORT), Handler)
  reloader_thread = threading.Thread(target=reloader)
Author	SHA1	Message	Date
geohot	f4e1b93225	fix beam	2025-08-29 19:05:49 -07:00
geohot	4a31c319b3	work	2025-08-29 18:42:50 -07:00
geohot	59081645f7	beam in RKernel	2025-08-29 18:42:50 -07:00
George HotzandGitHub	6e57905c6d	Merge branch 'master' into working_postopt	2025-08-29 12:37:01 -07:00
George HotzandGitHub	40606c60b0	Merge branch 'master' into working_postopt	2025-08-29 11:23:17 -07:00
George HotzandGitHub	80986321eb	Merge branch 'master' into working_postopt	2025-08-29 09:43:05 -07:00
George HotzandGitHub	bd263cbcb0	Merge branch 'master' into working_postopt	2025-08-28 15:13:02 -07:00
geohot	2e41472b02	support tc 2, all are pad	2025-08-28 14:40:50 -07:00
geohot	ac641f7b10	support tc padding	2025-08-28 14:34:42 -07:00
geohot	226c59fa5a	bugfix	2025-08-28 14:13:34 -07:00
geohot	3bbfcbccde	work	2025-08-28 13:55:07 -07:00
geohot	78a56b3461	fix some tests	2025-08-28 12:55:12 -07:00
George HotzandGitHub	b5ac4501d4	Merge branch 'master' into working_postopt	2025-08-28 12:35:39 -07:00
geohot	038d3bc295	clean up test	2025-08-28 11:39:02 -07:00
geohot	4b223c820a	fix test	2025-08-28 11:29:02 -07:00
geohot	528e285d81	more tests	2025-08-28 11:15:47 -07:00
geohot	cd3dc67636	tensor cores need to pad	2025-08-28 11:04:30 -07:00
geohot	b19a8963c3	revert	2025-08-28 11:00:20 -07:00
George HotzandGitHub	ec10e00cf5	Merge branch 'master' into working_postopt	2025-08-28 10:50:00 -07:00
geohot	d3aa38ad4a	work	2025-08-28 10:46:18 -07:00
geohot	6e41040e91	get tc ranges	2025-08-28 08:24:33 -07:00
geohot	6b3d9d6663	fix some tests	2025-08-28 08:07:45 -07:00
geohot	0ece218588	fix uop test	2025-08-28 07:34:06 -07:00
geohot	feb12685f2	postopt using RKernel	2025-08-28 07:28:52 -07:00