Revert "Revert "fix the misused cast in amd llvm tc (#11711 )" (#11715 )"

This reverts commit ca28db5a97.
improve syntax of UPats using f [pr] (#11717 )
2025-08-18 19:44:00 -07:00 · 2025-08-18 20:49:45 -04:00 · 2025-08-18 20:49:35 -04:00 · 2025-08-19 03:31:16 +03:00 · 2025-08-18 13:28:53 -07:00 · 2025-08-18 13:13:11 -07:00
62 changed files with 2406 additions and 1872 deletions
@@ -62,8 +62,6 @@ jobs:
      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run model inference benchmark
      run: METAL=1 python3.11 test/external/external_model_benchmark.py
-    - name: Run huggingface_onnx test
-      run: METAL=1 python3.11 extra/huggingface_onnx/run_models.py test --debug FacebookAI/xlm-roberta-large
    - name: Test speed vs torch
      run: BIG=2 MPS=1 python3.11 test/test_speed_v_torch.py | tee torch_speed.txt
    - name: Test tensor cores
@@ -605,12 +603,12 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: validate openpilot 0.9.7
-      run: PYTHONPATH=. FLOAT16=0 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
-    - name: benchmark openpilot 0.9.7
-      run: BENCHMARK_LOG=openpilot_0_9_7 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_7.txt
-    - name: benchmark openpilot w IMAGE=2 0.9.7
-      run: BENCHMARK_LOG=openpilot_0_9_7_image PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
+    - name: benchmark openpilot 0.9.9 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    - name: benchmark openpilot 0.9.9 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    - name: benchmark openpilot 0.9.9 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 0.9.9 driving_vision
      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.9.9 driving_policy
@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  DOWNLOAD_CACHE_VERSION: '10'
+  DOWNLOAD_CACHE_VERSION: '11'
  PYTHON_CACHE_VERSION: '2'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -46,6 +46,11 @@ jobs:
      with:
        deps: docs
        pydeps: "capstone"
+    - name: Build wheel and show size
+      run: |
+        pip install build
+        python -m build --wheel --outdir dist
+        ls -lh dist/*.whl
    - name: Use as an external package
      run: |
        mkdir $HOME/test_external_dir
@@ -329,7 +334,6 @@ jobs:
      run: |
        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check .
-        python3 -m ruff check extra/onnx.py
        python3 -m ruff check examples/mlperf/ --ignore E501
    - name: Lint tinygrad with pylint
      run: python -m pylint tinygrad/
@@ -337,7 +341,6 @@ jobs:
      run: |
        python -m mypy --strict-equality --lineprecision-report .
        cat lineprecision.txt
-        python -m mypy --strict-equality extra/onnx.py

  unittest:
    name: Unit Tests
@@ -375,8 +378,8 @@ jobs:
        PYTHONPATH=. python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
        DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
-    - name: Repo line count < 16000 lines
-      run: MAX_LINE_COUNT=16000 python sz.py
+    - name: Repo line count < 17000 lines
+      run: MAX_LINE_COUNT=17000 python sz.py

  fuzzing:
    name: Fuzzing
@@ -21,7 +21,7 @@ if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))

  model = Model()
-  opt = nn.optim.Adam(nn.state.get_parameters(model))
+  opt = (nn.optim.Adam if not getenv("MUON") else nn.optim.Muon)(nn.state.get_parameters(model))

  @TinyJit
  @Tensor.train()
@@ -29,8 +29,7 @@ if __name__ == "__main__":
    opt.zero_grad()
    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
    loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
-    opt.step()
-    return loss
+    return loss.realize(*opt.schedule_step())

  @TinyJit
  def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
@@ -1318,6 +1318,10 @@ def train_llama3():
  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)

+  if getenv("FAKEDATA"):
+    for v in get_parameters(model):
+      v = v.assign(Tensor.empty(v.shape))
+
  if (DP := getenv("DP", 1)) > 1:
    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
    for v in get_parameters(model):
@@ -1339,6 +1343,8 @@ def train_llama3():
      else:
        # attention_norm, ffn_norm, norm
        v.shard_(device, axis=None)
+      # prevents memory spike on device 0
+      v.realize()

  optim = AdamW(get_parameters(model), lr=0.0,
                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
@@ -1,6 +1,16 @@
-import re, ctypes, sys
+import re, ctypes, sys, importlib

-from tinygrad.runtime.autogen.am import am, mp_11_0, mp_13_0_0, nbio_4_3_0, mmhub_3_0_0, gc_11_0_0, osssys_6_0_0
+from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
+class AMDFake(AMDev):
+  def __init__(self, devfmt, vram, doorbell, mmio, dma_regions=None):
+    self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
+    self._run_discovery()
+    self._build_regs()
+
+amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
+amdev.AMDev = AMDFake
+
+from tinygrad.runtime.ops_amd import PCIIface

 def parse_amdgpu_logs(log_content, register_names=None):
  register_map = register_names
@@ -23,16 +33,11 @@ def parse_amdgpu_logs(log_content, register_names=None):
  return processed_log

 def main():
-  regs_offset = {13: {0: [3072, 37784576]}, 28: {0: [93184, 37754880], 1: [201327616, 201461760], 2: [209716224, 209850368], 3: [218104832, 218238976], 4: [226493440, 226627584], 5: [234882048, 235016192], 6: [243270656, 243404800]}, 21: {0: [28672, 12582912, 37795840, 130023424, 306184192], 1: [201326592, 201463808, 201465856, 204210176, 204472320], 2: [209715200, 209852416, 209854464, 212598784, 212860928], 3: [218103808, 218241024, 218243072, 220987392, 221249536], 4: [226492416, 226629632, 226631680, 229376000, 229638144], 5: [234881024, 235018240, 235020288, 237764608, 238026752], 6: [243269632, 243406848, 243408896, 246153216, 246415360]}, 22: {0: [18, 192, 13504, 36864, 37764096]}, 1: {0: [4704, 40960, 114688, 37760000]}, 2: {0: [3872, 37790720]}, 11: {0: [70656, 38103040]}, 12: {0: [106496, 37783552]}, 15: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 16: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 14: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 26: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 23: {0: [4256, 37789696]}, 33: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 25: {0: []}, 3: {0: [4704, 40960, 114688, 37760000]}, 4: {0: [4704, 40960, 114688, 37760000]}, 24: {0: [92160, 92672, 37752832, 54788096]}, 27: {0: [91648, 37751808], 1: [201339904, 201458176], 2: [209728512, 209846784], 3: [218117120, 218235392], 4: [226505728, 226624000], 5: [234894336, 235012608], 6: [243282944, 243401216]}, 29: {0: [201342976, 201344000, 205520896, 205537280], 1: [209731584, 209732608, 213909504, 213925888], 2: [218120192, 218121216, 222298112, 222314496], 3: [226508800, 226509824, 230686720, 230703104], 4: [234897408, 234898432, 239075328, 239091712], 5: [243286016, 243287040, 247463936, 247480320]}, 17: {0: [30720, 32256], 1: [31488, 73728]}}
-
  reg_names = {}
-  def _prepare_registers(modules):
-    for base, m in modules:
-      for k, regval in m.__dict__.items():
-        if k.startswith("reg") and not k.endswith("_BASE_IDX") and (base_idx:=getattr(m, f"{k}_BASE_IDX", None)) is not None:
-          reg_names[regs_offset[am.__dict__.get(f"{base}_HWIP")][0][base_idx] + regval] = k
-
-  _prepare_registers([("MP0", mp_13_0_0), ("NBIO", nbio_4_3_0), ("MMHUB", mmhub_3_0_0), ("GC", gc_11_0_0), ("OSSSYS", osssys_6_0_0)])
+  dev = PCIIface(None, 0)
+  for x, y in dev.dev_impl.__dict__.items():
+    if isinstance(y, AMRegister):
+      for inst, addr in y.addr.keys(): reg_names[addr] = f"{x}, xcc={inst}"

  with open(sys.argv[1], 'r') as f:
    log_content = log_content_them = f.read()
@@ -0,0 +1,61 @@
+# HuggingFace ONNX
+
+Tool for discovering, downloading, and validating ONNX models from HuggingFace.
+
+## Extra Dependencies
+
+```bash
+pip install huggingface_hub pyyaml requests onnx onnxruntime numpy
+```
+
+## Huggingface Manager (discovering and downloading)
+
+The `huggingface_manager.py` script discovers top ONNX models from HuggingFace, collects metadata, and optionally downloads them.
+
+```bash
+# Download top 50 models sorted by downloads
+python huggingface_manager.py --limit 50 --download
+
+# Just collect metadata (no download)
+python huggingface_manager.py --limit 100
+
+# Sort by likes instead of downloads
+python huggingface_manager.py --limit 20 --sort likes --download
+
+# Custom output file
+python huggingface_manager.py --limit 10 --output my_models.yaml
+```
+
+### Output Format
+
+The tool generates a YAML file with the following structure:
+
+```yaml
+repositories:
+  "model-name":
+    url: "https://huggingface.co/model-name"
+    download_path: "/path/to/models/..."  # when --download used
+    files:
+      - file: "model.onnx"
+        size: "90.91MB"
+total_size: "2.45GB"
+created_at: "2024-01-15T10:30:00Z"
+```
+
+## Run Models (validation)
+
+The `run_models.py` script validates ONNX models against ONNX Runtime for correctness.
+
+```bash
+# Validate models from a YAML configuration file
+python run_models.py --validate huggingface_repos.yaml
+
+# Debug specific repository (downloads and validates all ONNX models)
+python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2
+
+# Debug specific model file
+python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2/onnx/model.onnx
+
+# Debug with model truncation for debugging and validating intermediate results
+DEBUGONNX=1 python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2/onnx/model.onnx --truncate 10
+```
@@ -1,85 +0,0 @@
-import yaml, time, requests, argparse
-from pathlib import Path
-from huggingface_hub import list_models, HfApi
-from tinygrad.helpers import tqdm
-
-HUGGINGFACE_URL = "https://huggingface.co"
-SKIPPED_FILES = [
-  "fp16", "int8", "uint8", "quantized",      # numerical accuracy issues
-  "avx2", "arm64", "avx512", "avx512_vnni",  # numerical accuracy issues
-  "q4", "q4f16", "bnb4",                     # unimplemented quantization
-  "model_O4",                                # requires non cpu ort runner and MemcpyFromHost op
-  "merged",                                  # TODO implement attribute with graph type and Loop op
-]
-SKIPPED_REPO_PATHS = [
-  # Invalid model-index
-  "AdamCodd/vit-base-nsfw-detector",
-  # TODO: implement attribute with graph type and Loop op
-  "minishlab/potion-base-8M", "minishlab/M2V_base_output", "minishlab/potion-retrieval-32M",
-  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, GroupQueryAttention
-  "HuggingFaceTB/SmolLM2-360M-Instruct",
-  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, RotaryEmbedding, MultiHeadAttention
-  "HuggingFaceTB/SmolLM2-1.7B-Instruct",
-  # TODO: implmement RandomNormalLike
-  "stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", 'SimianLuo/LCM_Dreamshaper_v7',
-  # TODO: implement NonZero
-  "mangoapps/fb_zeroshot_mnli_onnx",
-  # TODO huge Concat in here with 1024 (1, 3, 32, 32) Tensors, and maybe a MOD bug with const folding
-  "briaai/RMBG-2.0",
-]
-
-def get_top_repos(n: int, sort: str) -> list[str]: # list["FacebookAI/xlm-roberta-large", ...]
-  print(f"** Getting top {n} models sorted by {sort} **")
-  repos = []
-  i = 0
-  for model in list_models(filter="onnx", sort=sort):
-    if model.id in SKIPPED_REPO_PATHS: continue
-    print(f"{i+1}/{n}: {model.id} ({getattr(model, sort)})")
-    repos.append(model.id)
-    i += 1
-    if i == n: break
-  return repos
-
-def get_metadata(repos:list[str]) -> dict:
-  api = HfApi()
-  repos_metadata = {"repositories": {}}
-  total_size = 0
-
-  # TODO: speed head requests up with async?
-  for repo in tqdm(repos, desc="Getting metadata"):
-    files_metadata = []
-    model_info = api.model_info(repo)
-
-    for file in model_info.siblings:
-      filename = file.rfilename
-      if not (filename.endswith('.onnx') or filename.endswith('.onnx_data')): continue
-      if any(skip_str in filename for skip_str in SKIPPED_FILES): continue
-      head = requests.head(f"{HUGGINGFACE_URL}/{repo}/resolve/main/{filename}", allow_redirects=True)
-      file_size = file.size or int(head.headers.get('Content-Length', 0))
-      files_metadata.append({"file": filename, "size": f"{file_size/1e6:.2f}MB"})
-      total_size += file_size
-
-    repos_metadata["repositories"][repo] = {
-      "url": f"{HUGGINGFACE_URL}/{repo}",
-      "download_path": None,
-      "files": files_metadata,
-    }
-  repos_metadata['total_size'] = f"{total_size/1e9:.2f}GB"
-  repos_metadata['created_at'] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
-  return repos_metadata
-
-if __name__ == "__main__":
-  sort = "downloads" # recent 30 days downloads
-  huggingface_onnx_dir = Path(__file__).parent
-
-  parser = argparse.ArgumentParser(description="Produces a YAML file with metadata of top huggingface onnx models")
-  parser.add_argument("--limit", type=int, required=True, help="Number of top repositories to process (e.g., 100)")
-  parser.add_argument("--output", type=str, default="huggingface_repos.yaml", help="Output YAML file name to save the report")
-  args = parser.parse_args()
-
-  top_repos = get_top_repos(args.limit, sort)
-  metadata = get_metadata(top_repos)
-  yaml_path = huggingface_onnx_dir / args.output
-  with open(yaml_path, 'w') as f:
-    yaml.dump(metadata, f, sort_keys=False)
-    print(f"YAML saved to: {str(yaml_path)}")
@@ -1,29 +0,0 @@
-import yaml, argparse
-from pathlib import Path
-from huggingface_hub import snapshot_download
-
-def download_models(yaml_file: str, download_dir: str) -> None:
-  with open(yaml_file, 'r') as f: metadata = yaml.safe_load(f)
-  n = len(metadata["repositories"])
-
-  for i, (model_id, model_data) in enumerate(metadata["repositories"].items()):
-    print(f"Downloading {i+1}/{n}: {model_id}...")
-    allow_patterns = [file_info["file"] for file_info in model_data["files"]]
-    root_path = Path(snapshot_download(repo_id=model_id, allow_patterns=allow_patterns, cache_dir=download_dir))
-    # download configs too (the sizes are small)
-    snapshot_download(repo_id=model_id, allow_patterns=["*config.json"], cache_dir=download_dir)
-    print(f"Downloaded model files to: {root_path}")
-    model_data["download_path"] = str(root_path)
-
-  # Save the updated metadata back to the YAML file
-  with open(yaml_file, 'w') as f: yaml.dump(metadata, f, sort_keys=False)
-  print("Download completed according to YAML file.")
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Download models from Huggingface Hub based on a YAML configuration file.")
-  parser.add_argument("input", type=str, help="Path to the input YAML configuration file containing model information.")
-  args = parser.parse_args()
-
-  models_folder = Path(__file__).parent / "models"
-  models_folder.mkdir(parents=True, exist_ok=True)
-  download_models(args.input, str(models_folder))
@@ -0,0 +1,230 @@
+import yaml
+import time
+import requests
+import argparse
+from pathlib import Path
+from huggingface_hub import list_models, HfApi, snapshot_download
+from tinygrad.helpers import _ensure_downloads_dir
+DOWNLOADS_DIR = _ensure_downloads_dir() / "models"
+from tinygrad.helpers import tqdm
+
+def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, cache_dir: str|Path|None=None,
+                                 tries: int=2, **kwargs) -> Path:
+  for attempt in range(tries):
+    try:
+      return Path(snapshot_download(
+        repo_id=repo_id,
+        allow_patterns=allow_patterns,
+        cache_dir=str(cache_dir) if cache_dir is not None else None,
+        **kwargs
+      ))
+    except Exception as e:
+      if attempt == tries-1: raise
+      time.sleep(1)
+
+# Constants for filtering models
+HUGGINGFACE_URL = "https://huggingface.co"
+SKIPPED_FILES = [
+  "fp16", "int8", "uint8", "quantized",      # numerical accuracy issues
+  "avx2", "arm64", "avx512", "avx512_vnni",  # numerical accuracy issues
+  "q4", "q4f16", "bnb4",                     # unimplemented quantization
+  "model_O4",                                # requires non cpu ort runner and MemcpyFromHost op
+  "merged",                                  # TODO implement attribute with graph type and Loop op
+]
+
+SKIPPED_REPO_PATHS = [
+  # Invalid model-index
+  "AdamCodd/vit-base-nsfw-detector",
+  # TODO: implement attribute with graph type and Loop op
+  "minishlab/potion-base-8M", "minishlab/M2V_base_output", "minishlab/potion-retrieval-32M",
+  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, GroupQueryAttention
+  "HuggingFaceTB/SmolLM2-360M-Instruct",
+  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, RotaryEmbedding, MultiHeadAttention
+  "HuggingFaceTB/SmolLM2-1.7B-Instruct",
+  # TODO: implement RandomNormalLike
+  "stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", 'SimianLuo/LCM_Dreamshaper_v7',
+  # TODO: implement NonZero
+  "mangoapps/fb_zeroshot_mnli_onnx",
+  # TODO huge Concat in here with 1024 (1, 3, 32, 32) Tensors, and maybe a MOD bug with const folding
+  "briaai/RMBG-2.0",
+]
+
+
+class HuggingFaceONNXManager:
+  def __init__(self):
+    self.base_dir = Path(__file__).parent
+    self.models_dir = DOWNLOADS_DIR
+    self.api = HfApi()
+
+  def discover_models(self, limit: int, sort: str = "downloads") -> list[str]:
+    print(f"Discovering top {limit} ONNX models sorted by {sort}...")
+    repos = []
+    i = 0
+
+    for model in list_models(filter="onnx", sort=sort):
+      if model.id in SKIPPED_REPO_PATHS:
+        continue
+
+      print(f"  {i+1}/{limit}: {model.id} ({getattr(model, sort)})")
+      repos.append(model.id)
+      i += 1
+      if i == limit:
+        break
+
+    print(f"Found {len(repos)} suitable ONNX models")
+    return repos
+
+  def collect_metadata(self, repos: list[str]) -> dict:
+    print(f"Collecting metadata for {len(repos)} repositories...")
+    metadata = {"repositories": {}}
+    total_size = 0
+
+    for repo in tqdm(repos, desc="Collecting metadata"):
+      try:
+        files_metadata = []
+        model_info = self.api.model_info(repo)
+
+        for file in model_info.siblings:
+          filename = file.rfilename
+          if not (filename.endswith('.onnx') or filename.endswith('.onnx_data')):
+            continue
+          if any(skip_str in filename for skip_str in SKIPPED_FILES):
+            continue
+
+          # Get file size from API or HEAD request
+          try:
+            head = requests.head(
+              f"{HUGGINGFACE_URL}/{repo}/resolve/main/{filename}",
+              allow_redirects=True,
+              timeout=10
+            )
+            file_size = file.size or int(head.headers.get('Content-Length', 0))
+          except requests.RequestException:
+            file_size = file.size or 0
+
+          files_metadata.append({
+            "file": filename,
+            "size": f"{file_size/1e6:.2f}MB"
+          })
+          total_size += file_size
+
+        if files_metadata:  # Only add repos with valid ONNX files
+          metadata["repositories"][repo] = {
+            "url": f"{HUGGINGFACE_URL}/{repo}",
+            "download_path": None,
+            "files": files_metadata,
+          }
+
+      except Exception as e:
+        print(f"WARNING: Failed to collect metadata for {repo}: {e}")
+        continue
+
+    metadata['total_size'] = f"{total_size/1e9:.2f}GB"
+    metadata['created_at'] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+
+    print(f"Collected metadata for {len(metadata['repositories'])} repositories")
+    print(f"Total estimated download size: {metadata['total_size']}")
+
+    return metadata
+
+  def download_models(self, metadata: dict) -> dict:
+    self.models_dir.mkdir(parents=True, exist_ok=True)
+
+    repos = metadata["repositories"]
+    n = len(repos)
+
+    print(f"Downloading {n} repositories to {self.models_dir}...")
+
+    for i, (model_id, model_data) in enumerate(repos.items()):
+      print(f"  Downloading {i+1}/{n}: {model_id}...")
+
+      try:
+        # Download ONNX model files
+        allow_patterns = [file_info["file"] for file_info in model_data["files"]]
+        root_path = snapshot_download_with_retry(
+          repo_id=model_id,
+          allow_patterns=allow_patterns,
+          cache_dir=str(self.models_dir)
+        )
+
+        # Download config files (usually small)
+        snapshot_download_with_retry(
+          repo_id=model_id,
+          allow_patterns=["*config.json"],
+          cache_dir=str(self.models_dir)
+        )
+
+        model_data["download_path"] = str(root_path)
+        print(f"    Downloaded to: {root_path}")
+
+      except Exception as e:
+        print(f"    ERROR: Failed to download {model_id}: {e}")
+        model_data["download_path"] = None
+        continue
+
+    successful_downloads = sum(1 for repo in repos.values() if repo["download_path"] is not None)
+    print(f"Successfully downloaded {successful_downloads}/{n} repositories")
+    print(f"All models saved to: {self.models_dir}")
+
+    return metadata
+
+  def save_metadata(self, metadata: dict, output_file: str):
+    yaml_path = self.base_dir / output_file
+    with open(yaml_path, 'w') as f:
+      yaml.dump(metadata, f, sort_keys=False)
+    print(f"Metadata saved to: {yaml_path}")
+
+  def discover_and_download(self, limit: int, output_file: str = "huggingface_repos.yaml",
+                          sort: str = "downloads", download: bool = True):
+    print(f"Starting HuggingFace ONNX workflow...")
+    print(f"   Limit: {limit} models")
+    print(f"   Sort by: {sort}")
+    print(f"   Download: {'Yes' if download else 'No'}")
+    print(f"   Output: {output_file}")
+    print("-" * 50)
+
+    repos = self.discover_models(limit, sort)
+
+    metadata = self.collect_metadata(repos)
+
+    if download:
+      metadata = self.download_models(metadata)
+
+    self.save_metadata(metadata, output_file)
+
+    print("-" * 50)
+    print("Workflow completed successfully!")
+    if download:
+      successful = sum(1 for repo in metadata["repositories"].values()
+                     if repo["download_path"] is not None)
+      print(f"{successful}/{len(metadata['repositories'])} models downloaded")
+
+    return metadata
+
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(
+    description="HuggingFace ONNX Model Manager - Discover, collect metadata, and download ONNX models",
+  )
+
+  parser.add_argument("--limit", type=int, help="Number of top repositories to process")
+  parser.add_argument("--output", type=str, default="huggingface_repos.yaml",
+                     help="Output YAML file name (default: huggingface_repos.yaml)")
+  parser.add_argument("--sort", type=str, default="downloads",
+                     choices=["downloads", "likes", "created", "modified"],
+                     help="Sort criteria for model discovery (default: downloads)")
+
+  parser.add_argument("--download", action="store_true", default=False,
+                     help="Download models after collecting metadata")
+
+  args = parser.parse_args()
+
+  if not args.limit: parser.error("--limit is required")
+
+  manager = HuggingFaceONNXManager()
+  manager.discover_and_download(
+    limit=args.limit,
+    output_file=args.output,
+    sort=args.sort,
+    download=args.download
+  )
@@ -1,10 +1,11 @@
-import onnx, yaml, tempfile, time, collections, pprint, argparse, json
+import onnx, yaml, tempfile, time, argparse, json
 from pathlib import Path
+from typing import Any
 from tinygrad.frontend.onnx import OnnxRunner
-from extra.onnx import get_onnx_ops
 from extra.onnx_helpers import validate, get_example_inputs
+from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry

-def get_config(root_path: Path):
+def get_config(root_path: Path) -> dict[str, Any]:
  ret = {}
  for path in root_path.rglob("*config.json"):
    config = json.load(path.open())
@@ -12,19 +13,19 @@ def get_config(root_path: Path):
      ret.update(config)
  return ret

-def run_huggingface_validate(onnx_model_path, config, rtol, atol):
-  onnx_runner = OnnxRunner(onnx_model_path)
-  inputs = get_example_inputs(onnx_runner.graph_inputs, config)
-  validate(onnx_model_path, inputs, rtol=rtol, atol=atol)
-
-def get_tolerances(file_name): # -> rtol, atol
+def get_tolerances(file_name: str) -> tuple[float, float]:
  # TODO very high rtol atol
  if "fp16" in file_name: return 9e-2, 9e-2
  if any(q in file_name for q in ["int8", "uint8", "quantized"]): return 4, 4
  return 4e-3, 3e-2

+def run_huggingface_validate(onnx_model_path: str | Path, config: dict[str, Any], rtol: float, atol: float):
+  onnx_runner = OnnxRunner(onnx_model_path)
+  inputs = get_example_inputs(onnx_runner.graph_inputs, config)
+  validate(onnx_model_path, inputs, rtol=rtol, atol=atol)
+
 def validate_repos(models:dict[str, tuple[Path, Path]]):
-  print(f"** Validating {len(model_paths)} models **")
+  print(f"** Validating {len(models)} models **")
  for model_id, (root_path, relative_path) in models.items():
    print(f"validating model {model_id}")
    model_path = root_path / relative_path
@@ -36,25 +37,6 @@ def validate_repos(models:dict[str, tuple[Path, Path]]):
    et = time.time() - st
    print(f"passed, took {et:.2f}s")

-def retrieve_op_stats(models:dict[str, tuple[Path, Path]]) -> dict:
-  ret = {}
-  op_counter = collections.Counter()
-  unsupported_ops = collections.defaultdict(set)
-  supported_ops = get_onnx_ops()
-  print(f"** Retrieving stats from {len(model_paths)} models **")
-  for model_id, (root_path, relative_path) in models.items():
-    print(f"examining {model_id}")
-    model_path = root_path / relative_path
-    onnx_runner = OnnxRunner(model_path)
-    for node in onnx_runner.graph_nodes:
-      op_counter[node.op] += 1
-      if node.op not in supported_ops:
-        unsupported_ops[node.op].add(model_id)
-    del onnx_runner
-  ret["unsupported_ops"] = {k:list(v) for k, v in unsupported_ops.items()}
-  ret["op_counter"] = op_counter.most_common()
-  return ret
-
 def debug_run(model_path, truncate, config, rtol, atol):
  if truncate != -1:
    model = onnx.load(model_path)
@@ -71,12 +53,9 @@ def debug_run(model_path, truncate, config, rtol, atol):
    run_huggingface_validate(model_path, config, rtol, atol)

 if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Huggingface ONNX Model Validator and Ops Checker")
-  parser.add_argument("input", type=str, help="Path to the input YAML configuration file containing model information.")
-  parser.add_argument("--check_ops", action="store_true", default=False,
-                      help="Check support for ONNX operations in models from the YAML file")
-  parser.add_argument("--validate", action="store_true", default=False,
-                      help="Validate correctness of models from the YAML file")
+  parser = argparse.ArgumentParser(description="Huggingface ONNX Model Validator")
+  parser.add_argument("--validate", type=str, default="",
+                      help="Validate correctness of models from the specified YAML configuration file")
  parser.add_argument("--debug", type=str, default="",
                      help="""Validates without explicitly needing a YAML or models pre-installed.
                      provide repo id (e.g. "minishlab/potion-base-8M") to validate all onnx models inside the repo
@@ -85,13 +64,13 @@ if __name__ == "__main__":
  parser.add_argument("--truncate", type=int, default=-1, help="Truncate the ONNX model so intermediate results can be validated")
  args = parser.parse_args()

-  if not (args.check_ops or args.validate or args.debug):
-    parser.error("Please provide either --validate, --check_ops, or --debug.")
+  if not (args.validate or args.debug):
+    parser.error("Please provide either --validate <yaml_file> or --debug <repo_id>.")
  if args.truncate != -1 and not args.debug:
    parser.error("--truncate and --debug should be used together for debugging")

-  if args.check_ops or args.validate:
-    with open(args.input, 'r') as f:
+  if args.validate:
+    with open(args.validate, 'r') as f:
      data = yaml.safe_load(f)
      assert all(repo["download_path"] is not None for repo in data["repositories"].values()), "please run `download_models.py` for this yaml"
      model_paths = {
@@ -101,22 +80,16 @@ if __name__ == "__main__":
        if model["file"].endswith(".onnx")
      }

-    if args.check_ops:
-      pprint.pprint(retrieve_op_stats(model_paths))
-
-    if args.validate:
-      validate_repos(model_paths)
+    validate_repos(model_paths)

  if args.debug:
-    from huggingface_hub import snapshot_download
-    download_dir = Path(__file__).parent / "models"
    path:list[str] = args.debug.split("/")
    if len(path) == 2:
      # repo id
      # validates all onnx models inside repo
      repo_id = "/".join(path)
-      root_path = Path(snapshot_download(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=download_dir))
-      snapshot_download(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=download_dir)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=DOWNLOADS_DIR)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
      config = get_config(root_path)
      for onnx_model in root_path.rglob("*.onnx"):
        rtol, atol = get_tolerances(onnx_model.name)
@@ -128,8 +101,8 @@ if __name__ == "__main__":
      onnx_model = path[-1]
      assert path[-1].endswith(".onnx")
      repo_id, relative_path = "/".join(path[:2]), "/".join(path[2:])
-      root_path = Path(snapshot_download(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=download_dir))
-      snapshot_download(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=download_dir)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=DOWNLOADS_DIR)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
      config = get_config(root_path)
      rtol, atol = get_tolerances(onnx_model)
      print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")
@@ -1,7 +1,6 @@
 from tinygrad import Tensor
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner
-from extra.onnx import OnnxValue
+from tinygrad.frontend.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort

@@ -0,0 +1,75 @@
+import torch
+
+#credit to KellerJordan at https://github.com/KellerJordan/Muon/tree/master
+#some changes: classic momentum instead of weighting gradient
+#added ns_steps, ns_params, nesterov as hyperparams
+def zeropower_via_newtonschulz5(G:torch.tensor, steps:int, params:tuple[int, ...]):
+  """
+  Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
+  quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
+  of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
+  zero even beyond the point where the iteration no longer converges all the way to one everywhere
+  on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
+  where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
+  performance at all relative to UV^T, where USV^T = G is the SVD.
+  """
+  assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
+
+  a, b, c = params
+  X = G
+  if G.size(-2) > G.size(-1):
+    X = X.mT
+
+  # Ensure spectral norm is at most 1
+  X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
+  # Perform the NS iterations
+  for _ in range(steps):
+    A = X @ X.mT
+    B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
+    X = a * X + B @ X
+
+  if G.size(-2) > G.size(-1):
+    X = X.mT
+
+  return X
+
+def muon_update(grad, momentum, beta=0.95, ns_steps=5, ns_params=(3.4445, -4.7750,  2.0315), nesterov=True):
+  if beta:
+    momentum.mul_(beta).add_(grad)
+    update = grad.add(momentum,alpha=beta) if nesterov else momentum
+  else: update = grad
+  if update.ndim == 4: # for the case of conv filters
+    update = update.view(len(update), -1)
+  update = zeropower_via_newtonschulz5(update, steps=ns_steps, params=ns_params)
+  return update
+
+class SingleDeviceMuon(torch.optim.Optimizer):
+  """
+  Muon variant for usage in non-distributed settings.
+  """
+  def __init__(self, params, lr=0.02, weight_decay=0.0, momentum=0.95, ns_steps=5, ns_params=(3.4445, -4.7750,  2.0315), nesterov=True):
+    defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, ns_steps=ns_steps, ns_params=ns_params, nesterov=nesterov)
+    super().__init__(params, defaults)
+
+  @torch.no_grad()
+  def step(self, closure=None):
+
+    loss = None
+    if closure is not None:
+      with torch.enable_grad():
+        loss = closure()
+
+    for group in self.param_groups:
+      for p in group["params"]:
+        if p.grad is None:
+          p.grad = torch.zeros_like(p)  # Force synchronization
+        state = self.state[p]
+        if len(state) == 0:
+          state["momentum_buffer"] = torch.zeros_like(p)
+        update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"], ns_steps=group["ns_steps"],
+                             ns_params=group["ns_params"], nesterov=group["nesterov"])
+        p.mul_(1.0 - group["lr"] * group["weight_decay"])
+
+        p.add_(update.reshape(p.shape), alpha=-group["lr"])
+
+    return loss
@@ -24,10 +24,28 @@ setup(name='tinygrad',
      license='MIT',
      long_description=long_description,
      long_description_content_type='text/markdown',
-      packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.runtime.autogen.am', 'tinygrad.codegen', 'tinygrad.nn',
-                  'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.schedule',
-                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.codegen.opt',
-                  'tinygrad.runtime.support.nv', 'tinygrad.apps'],
+      packages = [
+        'tinygrad',
+        'tinygrad.apps',
+        'tinygrad.codegen',
+        'tinygrad.codegen.opt',
+        'tinygrad.engine',
+        'tinygrad.frontend',
+        'tinygrad.nn',
+        'tinygrad.renderer',
+        'tinygrad.runtime',
+        'tinygrad.runtime.autogen',
+        'tinygrad.runtime.autogen.am',
+        'tinygrad.runtime.autogen.nv',
+        'tinygrad.runtime.graph',
+        'tinygrad.runtime.support',
+        'tinygrad.runtime.support.am',
+        'tinygrad.runtime.support.nv',
+        'tinygrad.schedule',
+        'tinygrad.shape',
+        'tinygrad.uop',
+        'tinygrad.viz',
+      ],
      package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
      classifiers=[
        "Programming Language :: Python :: 3",
@@ -294,7 +294,6 @@ class TestTrainingOnnxOps(TestOnnxOps):
        outputs = ["X_out", "V_out"]
        self._validate_training("Momentum", onnx_fxn, inputs, attributes, outputs)

-  @unittest.expectedFailure  # TODO: regression from removing StrEnum in Domain
  def test_adam_t_greater_than_zero(self):
    from onnx.backend.test.case.node.adam import apply_adam
    for t in [1, 3, 100]:
@@ -3,8 +3,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.uop.ops import Ops
 from tinygrad.device import is_dtype_supported
-from extra.onnx import OnnxDataType
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner, OnnxDataType
 from hypothesis import given, strategies as st

 # copied from test_const_folding.py
@@ -3,7 +3,7 @@ import z3
 from tinygrad import dtypes
 from tinygrad.uop.spec import z3_renderer, z3_cdiv
 from tinygrad.uop.ops import UOp, graph_rewrite
-from tinygrad.uop.transcendental import fast_idiv
+from tinygrad.uop.decompositions import fast_idiv
 random.seed(42)

 powers_of_two = [2**i for i in range(64)]
@@ -87,16 +87,19 @@ class AMDDriver(VirtDriver):
        functools.partial(TextFileDesc, text=gpu_props.format(drm_render_minor=gpu_id))),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0',
               functools.partial(DirFileDesc, child_names=[str(am.GC_HWID), str(am.SDMA0_HWID), str(am.NBIF_HWID)])),
+      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/major', functools.partial(TextFileDesc, text='11')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/minor', functools.partial(TextFileDesc, text='0')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.GC_HWID}/0/base_addr',
               functools.partial(TextFileDesc, text='0x00001260\n0x0000A000\n0x0001C000\n0x02402C00')),
+      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/major', functools.partial(TextFileDesc, text='6')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/minor', functools.partial(TextFileDesc, text='0')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.SDMA0_HWID}/0/base_addr',
               functools.partial(TextFileDesc, text='0x00001260\n0x0000A000\n0x0001C000\n0x02402C00')),
+      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}', functools.partial(DirFileDesc, child_names=['0'])),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/major', functools.partial(TextFileDesc, text='4')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/minor', functools.partial(TextFileDesc, text='3')),
      VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{am.NBIF_HWID}/0/revision', functools.partial(TextFileDesc, text='0')),
@@ -9,7 +9,15 @@ except ModuleNotFoundError:
  raise unittest.SkipTest("onnx not installed, skipping onnx test")
 from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import CI, fetch, temp
+from tinygrad.device import Device
+from tinygrad.helpers import CI, fetch, temp, Context
+
+try:
+  from extra.onnx_helpers import validate
+  from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry
+  HUGGINGFACE_AVAILABLE = True
+except ModuleNotFoundError:
+  HUGGINGFACE_AVAILABLE = False

 def run_onnx_torch(onnx_model, inputs):
  import torch
@@ -137,5 +145,36 @@ class TestOnnxModel(unittest.TestCase):
    print(cls, _LABELS[cls])
    assert "car" in _LABELS[cls] or _LABELS[cls] == "convertible"

+@unittest.skipUnless(HUGGINGFACE_AVAILABLE and Device.DEFAULT == "METAL", "only run on METAL")
+class TestHuggingFaceOnnxModels(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls._ctx = Context(MAX_BUFFER_SIZE=0)
+    cls._ctx.__enter__()
+
+  @classmethod
+  def tearDownClass(cls):
+    cls._ctx.__exit__()
+
+  def _validate(self, repo_id, model_file, custom_inputs, rtol=1e-4, atol=1e-4):
+    onnx_model_path = snapshot_download_with_retry(
+      repo_id=repo_id,
+      allow_patterns=["*.onnx", "*.onnx_data"],
+      cache_dir=str(DOWNLOADS_DIR)
+    )
+    onnx_model_path = onnx_model_path / model_file
+    file_size = onnx_model_path.stat().st_size
+    print(f"Validating model: {repo_id}/{model_file} ({file_size/1e6:.2f}M)")
+    validate(onnx_model_path, custom_inputs, rtol=rtol, atol=atol)
+
+  def test_xlm_roberta_large(self):
+    repo_id = "FacebookAI/xlm-roberta-large"
+    model_file = "onnx/model.onnx"
+    custom_inputs = {
+      "input_ids": np.random.randint(0, 250002, (1, 11), dtype=np.int64),
+      "attention_mask": np.ones((1, 11), dtype=np.int64),
+    }
+    self._validate(repo_id, model_file, custom_inputs)
+
 if __name__ == "__main__":
  unittest.main()
@@ -143,13 +143,12 @@ class TestIndexingConstFolding(unittest.TestCase):
    _check_ast_count(1, t[:,:,Tensor(1)+2,:])
    _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])

-  @unittest.expectedFailure
  def test_const_tensor_index(self):
-    # TODO: implement const tensor folded indexing
+    # TODO: these can be 0, implement const tensor folded indexing
    t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
-    _check_ast_count(0, t[:,:,Tensor.ones(2,1),:])
-    _check_ast_count(0, t[:,:,Tensor.ones(1,2)+2,:])
-    _check_ast_count(0, t[:,:,Tensor.ones(1,1),Tensor.zeros(2,1,2)])
+    _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
+    _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
+    _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])

 class TestMovedConstFolding(unittest.TestCase):
  def test_add_shrunk_zero(self):
@@ -62,5 +62,15 @@ class TestLinAlg(unittest.TestCase):
      orthogonality_helper(Q)
      reconstruction_helper([Q,R],a)

+  def test_newton_schulz(self):
+    coefficients = [(2, -1.5, 0.5), (2.0, -1.4, 0.2, 0.2)]#these params map to the sign function
+    sizes = [(2,2), (3,2), (2,3), (2,2,2)]
+    for coefs in coefficients:
+      for size in sizes:
+        a = Tensor.randn(size)
+        b = Tensor.newton_schulz(a, steps=20, params=coefs, eps=0.0)
+        # ns(A) = U @ Vt -> (U @ Vt) @ (U @ Vt)t = I
+        orthogonality_helper(b if size[-1] > size[-2] else b.transpose(-2, -1), tolerance=1e-1)
+
 if __name__ == "__main__":
  unittest.main()
@@ -393,7 +393,7 @@ class TestOps(unittest.TestCase):
  def test_trunc(self):
    helper_test_op([()], lambda x: x.trunc(), forward_only=True)
    helper_test_op([(45,35)], lambda x: x.trunc(), forward_only=True)
-    helper_test_op(None, lambda x: x.trunc(), vals=[[1.499, 1.5, 1.501, 1.0, 2.1, 0.0, -5.0, -2.499, -2.5, -2.501]], forward_only=True)
+    helper_test_op(None, lambda x: x.trunc(), vals=[[1.499, 1.5, 1.501, 1.0, 2.1, 0.0, -5.0, -2.499, -2.5, -2.501, 1e12, -1e12]], forward_only=True)
  def test_floor(self):
    helper_test_op([()], lambda x: x.floor(), forward_only=True)
    helper_test_op([(45,35)], lambda x: x.floor(), forward_only=True)
@@ -2,9 +2,10 @@ import numpy as np
 import torch
 import unittest
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.nn.optim import Adam, SGD, AdamW
+from tinygrad.nn.optim import Adam, SGD, AdamW, Muon
 from tinygrad.helpers import CI
 from tinygrad.device import is_dtype_supported
+from extra.torch_muon import SingleDeviceMuon as TorchMuon

 np.random.seed(1337)
 x_init = np.random.randn(1,4).astype(np.float32)
@@ -57,9 +58,12 @@ class TestOptim(unittest.TestCase):
  def _test_sgd(self, steps, opts, atol, rtol): self._test_optim(SGD, torch.optim.SGD, steps, opts, atol, rtol)
  def _test_adam(self, steps, opts, atol, rtol): self._test_optim(Adam, torch.optim.Adam, steps, opts, atol, rtol)
  def _test_adamw(self, steps, opts, atol, rtol): self._test_optim(AdamW, torch.optim.AdamW, steps, opts, atol, rtol)
+  #TODO: use torch.muon when it comes out
+  def _test_muon(self, steps, opts, atol, rtol): self._test_optim(Muon, TorchMuon, steps, opts, atol, rtol)

  def test_multistep_sgd_high_lr_teeny(self): self._test_sgd(2, {'lr': 1.1, 'teeny': True}, 1e-6, 1e-5)
  def test_multistep_adam_high_lr_teeny(self): self._test_adam(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)
+  def test_multistep_muon_high_lr_teeny(self): self._test_muon(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)

  def test_sgd(self): self._test_sgd(1, {'lr': 0.001}, 1e-6, 0)
  def test_sgd_high_lr(self): self._test_sgd(1, {'lr': 10}, 1e-6, 1e-5)
@@ -83,6 +87,28 @@ class TestOptim(unittest.TestCase):
  def test_multistep_sgd_high_lr_nesterov_momentum_wd(self):
    self._test_sgd(10, {'lr': 9, 'momentum': 0.9, 'nesterov': True, 'weight_decay': 0.1}, 1e-5, 3e-4)

+  def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-6, 0)
+  def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
+  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-6, 0)
+  def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 3e-4)
+
+  # NOTE: momentum set to 0.95 by default, nesterov set to True by default
+  def test_multistep_muon_momentum_wd(self): self._test_muon(10, {'lr': 0.001, 'weight_decay': 0.01}, 1e-5, 0)
+  # ns defaults are numerically unstable, but it is tolerable in real training (see nsteps/nparam tests)
+  def test_multistep_muon_high_lr_momentum_wd(self): self._test_muon(10, {'lr': 10, 'weight_decay': 0.01}, 1e-1, 3e-4)
+  def test_multistep_muon_no_nesterov_momentum(self): self._test_muon(10, {'lr': 0.001, 'nesterov': False}, 1e-5, 0)
+  def test_multistep_muon_high_lr_no_nesterov_momentum(self): self._test_muon(10, {'lr': 10, 'nesterov': False}, 0.5e-1, 1e-1)
+
+  def test_muon_ns_steps(self): self._test_muon(1, {'lr': 0.001, 'ns_steps': 3}, 1e-6, 0)
+  def test_muon_high_lr_ns_steps(self): self._test_muon(1, {'lr': 10, 'ns_steps': 3}, 1e-5, 3e-4)
+  def test_muon_ns_params(self): self._test_muon(1, {'lr': 0.001,'ns_params': (2.0,-1.5,0.5)}, 1e-6, 0)
+  def test_muon_high_lr_ns_params(self): self._test_muon(1, {'lr': 10,'ns_params': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+
+  def test_muon_momentum_wd_ns_steps_ns_params(self):
+    self._test_muon(10, {'lr': 0.001, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_params': (2.0,-1.5,0.5)}, 1e-5, 0)
+  def test_multistep_muon_high_lr_momentum_wd_ns_steps_ns_params(self):
+    self._test_muon(10, {'lr': 10, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_params': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
+
  def test_adam(self): self._test_adam(1, {'lr': 0.001}, 1e-5, 0)
  def test_adam_high_lr(self): self._test_adam(1, {'lr': 10}, 1e-4, 1e-4)
  def test_adamw(self): self._test_adamw(1, {'lr': 0.001}, 1e-5, 0)
@@ -117,6 +117,14 @@ class TestFuse(unittest.TestCase):
    c = (a.sum(axis=1) + b.sum(axis=1)).fuse()
    self.assertListEqual(c.tolist(), [30]*16)

+  @unittest.skipUnless(Device.DEFAULT == "METAL", "METAL TC")
+  def test_fuse_and_tc_opt(self):
+    A = Tensor.randn(8, 8).realize()
+    B = Tensor.randn(8, 8).realize()
+    C = Tensor.ones(1, 8, 8).pad(((1,1), None, None),).sum(0)
+    out = (C + (A @ B)).fuse()
+    out.realize()
+
 class TestSoftmaxFusion(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
@@ -106,6 +106,24 @@ class TestTiny(unittest.TestCase):
    probs = Tensor.rand(1, 1, 28, 28).sequential(layers).tolist()
    self.assertEqual(len(probs[0]), 10)

+  # TODO: this is failing because of how swizzling rewrites the ShapeTracker of the final STORE
+  @unittest.skipIf(IMAGE>0 or (CI and Device.DEFAULT == "DSP"), "failing because of make things that can't be images not images")
+  def test_mnist_backward(self):
+    # NOTE: we don't have the whole model here for speed
+    layers = [
+      nn.Conv2d(1, 32, 5), Tensor.relu,
+      nn.Conv2d(32, 32, 5), Tensor.relu]
+
+    # replace random weights with ones
+    # TODO: there's a bug here where it's tying two of the biases together. we need UNIQUE const
+    #Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])
+    for p in nn.state.get_parameters(layers): p.replace(Tensor.empty(p.shape))
+
+    # realize gradients
+    for x in nn.state.get_parameters(layers): x.requires_grad_()
+    Tensor.empty(4, 1, 28, 28).sequential(layers).sum().backward()
+    Tensor.realize(*[x.grad for x in nn.state.get_parameters(layers) if x.grad is not None])
+
  # *** image ***

  @unittest.skipIf(Device.DEFAULT != "GPU", "image only supported on GPU")
@@ -53,5 +53,13 @@ class TestCastConvenienceMethod(unittest.TestCase):
      self.assertEqual(t.float().dtype, dtypes.float)
      self.assertEqual(t.double().dtype, dtypes.double)

+class TestDtypeTolist(unittest.TestCase):
+  def test_bfloat16(self):
+    self.assertEqual(Tensor([-60000, 1.5, 3.1, 60000], device="PYTHON", dtype=dtypes.bfloat16).tolist(), [-59904.0, 1.5, 3.09375, 59904.0])
+    # 448
+    self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e4m3).tolist(), [-448.0, 1.5, 3.0, 448.0])
+    # 57344
+    self.assertEqual(Tensor([-30000, 1.5, 3.1, 30000], device="PYTHON", dtype=dtypes.fp8e5m2).tolist(), [-28672.0, 1.5, 3.0, 28672.0])
+
 if __name__ == "__main__":
  unittest.main()
@@ -303,8 +303,8 @@ class TestRecurse(unittest.TestCase):
  def test_inf_loop(self):
    a = UOp.variable('a', 0, 10)
    pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.DEFINE_REG)),
-      (UPat(Ops.DEFINE_REG, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
+      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
    ])
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm)
@@ -312,8 +312,8 @@ class TestRecurse(unittest.TestCase):
  def test_inf_loop_bottom_up(self):
    a = UOp.variable('a', 0, 10)
    pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.DEFINE_REG)),
-      (UPat(Ops.DEFINE_REG, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
+      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
    ])
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm, bottom_up=True)
@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor
 from tinygrad.uop.ops import PatternMatcher, Ops, UPat, graph_rewrite, RewriteContext, UOp
-from tinygrad.schedule.kernelize import sym, merge_views
+from tinygrad.schedule.kernelize import kernelize_sym, merge_views

 class TestRewriteTrackedChildren(unittest.TestCase):
  @unittest.skip("track_children no longer supported")
@@ -57,7 +57,7 @@ class TestRewriteTrackedChildren(unittest.TestCase):
    extra = PatternMatcher([(UPat(Ops.REDUCE_AXIS, name="r"), print_children)])
    a = Tensor.empty(3, 3)
    r = (a+0).sum()
-    graph_rewrite(r.uop, merge_views+sym+extra, track_children=True)
+    graph_rewrite(r.uop, merge_views+kernelize_sym+extra, track_children=True)

 if __name__ == '__main__':
  unittest.main()
@@ -2,8 +2,8 @@ import unittest, math
 import numpy as np
 from tinygrad import dtypes
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.uop.transcendental import TRANSCENDENTAL_SUPPORTED_DTYPES, payne_hanek_reduction, cody_waite_reduction
-from tinygrad.uop.transcendental import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
+from tinygrad.uop.decompositions import TRANSCENDENTAL_SUPPORTED_DTYPES, payne_hanek_reduction, cody_waite_reduction
+from tinygrad.uop.decompositions import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
 from test.helpers import eval_uop

 class TestTranscendentalFunctions(unittest.TestCase):
@@ -30,16 +30,16 @@ class TestSymbolicPickle(unittest.TestCase):

 class TestSymbolic(unittest.TestCase):
  def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    rendered, nmin, nmax = render(v)
-    if isinstance(s, tuple): self.assertIn(rendered, s)
-    else: self.assertEqual(rendered, s)
-    self.assertEqual(nmin, n)
-    self.assertEqual(nmax, m)
    if test_z3:
      solver = z3.Solver()
      z3_sink = graph_rewrite(v.sink(v.simplify()), z3_renderer, ctx=(solver, {}))
      expr, epxr_simplified = z3_sink.src[0].arg, z3_sink.src[1].arg
      self.assertEqual(solver.check(expr != epxr_simplified), z3.unsat, "simplified expression not equal to original")
+    rendered, nmin, nmax = render(v)
+    if isinstance(s, tuple): self.assertIn(rendered, s)
+    else: self.assertEqual(rendered, s)
+    self.assertEqual(nmin, n)
+    self.assertEqual(nmax, m)

  def test_cmp_simple(self):
    self.helper_test_variable(Variable("a", 3, 8) < 4, 0, 1, "(a<4)")
@@ -266,6 +266,16 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(((5*Variable("a", 0, 31)) % 12) % 5, 0, 4, "(((a*5)%12)%5)")
    self.helper_test_variable((Variable("a", 0, 31) % 4) % 12, 0, 3, "(a%4)")

+  def test_mod_mod_wrong_sign(self):
+    v1=Variable("v1", 0, 128)
+    v3=Variable("v3", 0, 7)
+    self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), -4, 4, "(((((v1%2)*2)+((v3+-1)%5))+-2)%5)")
+
+  def test_mod_mod_wrong_sign2(self):
+    v2=Variable("v2", 0, 8)
+    v3=Variable("v3", 0, 4)
+    self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), -6, 6, "(((v2+((v3+3)%7))+-2)%7)")
+
  def test_mul_mul(self):
    self.helper_test_variable((Variable("a", 0, 5)*10)*9, 0, 5*10*9, "(a*90)")

@@ -375,6 +385,17 @@ class TestSymbolic(unittest.TestCase):
  def test_mul_div(self):
    self.helper_test_variable((Variable("a", 0, 10)*4)//4, 0, 10, "a")

+  def test_div_drop_small_terms(self):
+    # from openpilot, shouldnt simplify
+    gidx0 = UOp.variable("gidx0", 0, 10)
+    gidx1 = UOp.variable("gidx1", 0, 10)
+    lidx0 = UOp.variable("lidx0", 0, 1)
+    lidx1 = UOp.variable("lidx1", 0, 1)
+    ridx1005 = UOp.variable("ridx1005", 0, 2)
+    ridx1006 = UOp.variable("ridx1006", 0, 2)
+    self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -2, 20,
+      "(((((lidx1+(((gidx1*18)+(ridx1005*18))+(lidx0*162)))+(gidx0*2))+(ridx1006*2))+-40)//18)")
+
  def test_add_div(self):
    # careful about the lower bounds and upper bounds
    self.helper_test_variable((Variable("a", 0, 5)-2)//4, 0, 0, "0")
@@ -421,6 +442,11 @@ class TestSymbolic(unittest.TestCase):
  def test_div_numerator_negative(self):
    self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -8, 0, "(((idx*10)//11)*-1)")

+  def test_nest_div_negative_factor(self):
+    ridx0=UOp.variable("ridx0", 0, 9)
+    ridx1=UOp.variable("ridx1", 0, 6)
+    self.helper_test_variable(((((ridx0*-7)+ridx1)+63)//35), 0, 1, "(((ridx0//5)*-1)+1)")
+
  def test_div_into_mod(self):
    self.helper_test_variable((Variable("idx", 0, 16)*4)%8//4, 0, 1, "(idx%2)")

@@ -679,6 +705,10 @@ class TestSymbolic(unittest.TestCase):
    # TODO: should z3 work?
    self.helper_test_variable(2*(2*a).reciprocal(), -math.inf, math.inf, "(1/a)", test_z3=False)

+  def test_trunc_noop(self):
+    a = Variable("a", 1, 10, dtypes.int)
+    self.helper_test_variable(a.trunc(), 1, 10, "a", test_z3=False)
+
 class TestSymbolicNumeric(unittest.TestCase):
  def helper_test_numeric(self, f):
    MIN, MAX = 0, 10
@@ -76,7 +76,7 @@ class TestViz(BaseTestViz):
    self.assertEqual(lineno, inner.__code__.co_firstlineno)

  def test_exceptions(self):
-    # VIZ tracks rewrites up to the error
+    # VIZ tracks rewrites up to and including the error
    def count_3(x:UOp):
      assert x.arg <= 3
      return x.replace(arg=x.arg+1)
@@ -85,7 +85,7 @@ class TestViz(BaseTestViz):
    with self.assertRaises(AssertionError): exec_rewrite(a, [err_pm])
    lst = get_viz_list()
    err_step = lst[0]["steps"][0]
-    self.assertEqual(err_step["match_count"], 3)
+    self.assertEqual(err_step["match_count"], 4) # 3 successful rewrites + 1 err

  def test_default_name(self):
    a = UOp.variable("a", 1, 10)
@@ -124,10 +124,10 @@ class TestViz(BaseTestViz):

  def test_inf_loop(self):
    a = UOp.variable('a', 0, 10)
-    b = a.replace(op=Ops.DEFINE_REG)
+    b = a.replace(op=Ops.CONST)
    pm = PatternMatcher([
-      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.DEFINE_REG)),
-      (UPat(Ops.DEFINE_REG, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (UPat(Ops.DEFINE_VAR, name="x"), lambda x: x.replace(op=Ops.CONST)),
+      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
    ])
    with self.assertRaises(RuntimeError): exec_rewrite(a, [pm])
    graphs = flatten(x["graph"].values() for x in get_details(tracked_ctxs[0][0]))
@@ -250,7 +250,7 @@ class TestVizProfiler(unittest.TestCase):

    j = json.loads(get_profile(prof))

-    dev_events = j['layout']['NV']['timeline']['shapes']
+    dev_events = j['layout']['NV']['shapes']
    self.assertEqual(len(dev_events), 1)
    event = dev_events[0]
    self.assertEqual(event['name'], 'E_2')
@@ -263,7 +263,7 @@ class TestVizProfiler(unittest.TestCase):

    j = json.loads(get_profile(prof))

-    event = j['layout']['NV']['timeline']['shapes'][0]
+    event = j['layout']['NV']['shapes'][0]
    self.assertEqual(event['name'], 'COPYxx')
    self.assertEqual(event['st'], 900) # diff clock
    self.assertEqual(event['dur'], 10)
@@ -278,23 +278,23 @@ class TestVizProfiler(unittest.TestCase):

    j = json.loads(get_profile(prof))

-    devices = list(j['layout'])
-    self.assertEqual(devices[0], 'NV Graph')
-    self.assertEqual(devices[1], 'NV')
-    self.assertEqual(devices[2], 'NV:1')
+    tracks = list(j['layout'])
+    self.assertEqual(tracks[0], 'NV Graph')
+    self.assertEqual(tracks[2], 'NV')
+    self.assertEqual(tracks[4], 'NV:1')

-    nv_events = j['layout']['NV']['timeline']['shapes']
+    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']['timeline']['shapes']
+    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']['timeline']['shapes']
+    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'])

@@ -308,7 +308,7 @@ class TestVizMemoryLayout(BaseTestViz):
    a = _alloc(1)
    _b = _alloc(1)
    profile_ret = json.loads(get_profile(Buffer.profile_events))
-    ret = profile_ret["layout"][a.device]["mem"]
+    ret = profile_ret["layout"][f"{a.device} Memory"]
    self.assertEqual(ret["peak"], 2)
    self.assertEqual(ret["shapes"][0]["x"], [0, 2])
    self.assertEqual(ret["shapes"][1]["x"], [1, 2])
@@ -318,7 +318,7 @@ class TestVizMemoryLayout(BaseTestViz):
    del a
    b = _alloc(1)
    profile_ret = json.loads(get_profile(Buffer.profile_events))
-    ret = profile_ret["layout"][b.device]["mem"]
+    ret = profile_ret["layout"][f"{b.device} Memory"]
    self.assertEqual(ret["peak"], 1)
    self.assertEqual(ret["shapes"][0]["x"], [0, 2])
    self.assertEqual(ret["shapes"][1]["x"], [2, 3])
@@ -331,7 +331,7 @@ class TestVizMemoryLayout(BaseTestViz):
    del a
    c = _alloc(1)
    profile_ret = json.loads(get_profile(Buffer.profile_events))
-    ret = profile_ret["layout"][c.device]["mem"]
+    ret = profile_ret["layout"][f"{c.device} Memory"]
    self.assertEqual(ret["peak"], 2)
    self.assertEqual(ret["shapes"][0]["x"], [0, 3])
    self.assertEqual(ret["shapes"][1]["x"], [1, 3, 3, 4])
@@ -11,7 +11,7 @@ from tinygrad.codegen.lowerer import pm_lowerer, get_index
 from tinygrad.codegen.quantize import pm_quant
 from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
-from tinygrad.uop.optional import get_late_rewrite_patterns
+from tinygrad.uop.decompositions import get_late_rewrite_patterns
 from tinygrad.codegen.expander import migrate_indexing, expander
 from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
  ReduceContext, correct_load_store, pm_render
@@ -85,8 +85,12 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
  # optional pre matcher
  if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher"))

+  # decompositions
+  pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, _TRANSCENDENTAL>=2)
+  ret.append(RewriteStep(pm_decomp, lambda _: opts.device, name="decompositions"))
+
  # final rules for the renderer (without sym)
-  pm_final_rewrite = symbolic_simple+get_late_rewrite_patterns(supported_ops, _TRANSCENDENTAL>=2)+pm_render+extra_matcher
+  pm_final_rewrite = pm_decomp+pm_render+extra_matcher
  ret.append(RewriteStep(pm_final_rewrite, lambda _: opts.device, name="final rewrite"))

  # return the list (with optional linearizer)
@@ -285,7 +285,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
    topo = inp.toposort()
    stored_ranges = flatten([x.src[2:] for x in topo if x.op is Ops.STORE])
    input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in stored_ranges])
-    identity = red.const_like(identity_element(red.arg, red.dtype.scalar()))
+    identity = red.const(red.dtype, identity_element(red.arg, red.dtype.scalar()))
    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=(ctx.acc_num,)).index(UOp.const(dtypes.int, 0))
    do_store = acc.store(identity, UOp(Ops.NOOP, src=input_ranges)) if len(input_ranges) else acc.store(identity)
    lst = [acc.load(do_store, *reduce_range)] + lst  # put acc as the first element
@@ -227,7 +227,7 @@ block_merge = PatternMatcher([

 def finalize(sink:UOp) -> UOp:
  if sink.op is not Ops.BLOCK or not all(x.op in DONT_PLACE_IN_BLOCK for x in sink.src):
-    raise RuntimeError("linearize failure")
+    raise RuntimeError(f"linearize failure {sink.op} {[x.op for x in sink.src if x.op not in DONT_PLACE_IN_BLOCK]}")

  # place the early things
  lst = sorted(dedup(sink.src), key=lambda x: x.tuplize) + list(sink.arg.lst)
@@ -378,9 +378,9 @@ class Kernel:
      tensor_cores = self.opts.tensor_cores if tc_select == -1 else [self.opts.tensor_cores[tc_select]]
      for tc in tensor_cores:
        tensor_core_opts = [self._create_tc_opts(reduceop, tc, axis, opt_level) for reduceop in self.reduceops]
+        if tensor_core_opts[0] is None: continue
        # can only fuse reduces with the same tc options
        assert all_same(tensor_core_opts)
-        if tensor_core_opts[0] is None: continue
        self.tensor_core_opts = tc_opts = tensor_core_opts[0]

        # attempt to pad the tensor axes that require it
@@ -3,8 +3,8 @@ 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, 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, cpu_events, ProfileEvent, ProfilePointEvent, dedup
+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
 from tinygrad.renderer import Renderer

@@ -124,7 +124,7 @@ class Buffer:
  def allocate(self, opaque=None, external_ptr=None) -> Buffer:
    assert not self.is_initialized(), "can't allocate already allocated buffer"
    if DEBUG >= 7: print(f"buffer: allocate {self.nbytes} bytes on {self.device}")
-    if (mbs:=getenv("MAX_BUFFER_SIZE", 0)) > 0 and self.size > mbs: raise RuntimeError(f"buffer of size {self.size/1e6:.2f}M is too large")
+    if MAX_BUFFER_SIZE > 0 and self.size > MAX_BUFFER_SIZE: raise RuntimeError(f"buffer of size {self.size/1e6:.2f}M is too large")
    self.allocator:Allocator = Device[self.device].allocator
    if external_ptr is not None:
      self.options = replace(self.options, external_ptr=external_ptr) if self.options else BufferSpec(external_ptr=external_ptr)
@@ -139,7 +139,7 @@ DISABLE_COMPILER_CACHE = ContextVar("DISABLE_COMPILER_CACHE", 0)
 DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES = ContextVar("DONT_REALIZE_EXPAND", 0), ContextVar("DONT_GROUP_REDUCES", 0)
 QUANTIZE, VALIDATE_WITH_CPU = ContextVar("QUANTIZE", 0), ContextVar("VALIDATE_WITH_CPU", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
-ALLOW_DEVICE_USAGE, AMD_LLVM = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("AMD_LLVM", 1)
+ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, AMD_LLVM = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0), ContextVar("AMD_LLVM", 1)

@dataclass(frozen=True)
 class Metadata:
@@ -77,7 +77,19 @@ def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov

  `classic` is a boolean flag that determines whether to use the popular momentum update rule or the classic momentum update rule.
  """
-  return LARS(params, lr, momentum, weight_decay, nesterov, classic, tcoef=0.0, fused=fused)
+  return LARS(params, lr, momentum, weight_decay, 0, None, nesterov, classic=classic, pre_wd=True, tcoef=0.0, fused=fused)
+
+# Muon applies the newton schulz algorithm on gradient. also can include momentum, nesterov, and weight decay
+def Muon(params: list[Tensor], lr=0.02, momentum=0.95, weight_decay=0.0, ns_steps=5, ns_params=(3.4445, -4.775, 2.0315),
+         nesterov=True, fused=FUSE_OPTIM):
+  """
+  SGD with newton-schulz iteration and post momentum weight decay.
+
+  - Described: https://kellerjordan.github.io/posts/muon/
+  - Paper: https://arxiv.org/pdf/2502.16982
+  """
+  assert not fused, "FUSE_OPTIM not allowed for Muon optimizer"
+  return LARS(params, lr, momentum, weight_decay, ns_steps, ns_params, nesterov, classic=False, pre_wd=False, tcoef=0.0, fused=fused)

 class LARS(Optimizer):
  """
@@ -85,9 +97,11 @@ class LARS(Optimizer):

  - Paper: https://arxiv.org/abs/1708.03888v3
  """
-  def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, nesterov=False, classic=True, tcoef=0.001, fused=FUSE_OPTIM):
+  def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, ns_steps=0, ns_params=None,
+               nesterov=False, classic=True, pre_wd=True, tcoef=0.001, fused=FUSE_OPTIM):
    super().__init__(params, lr, fused)
-    self.momentum, self.wd, self.nesterov, self.classic, self.tcoef = momentum, weight_decay, nesterov, classic, tcoef
+    self.momentum, self.wd, self.ns_steps, self.ns_params  = momentum, weight_decay, ns_steps, ns_params
+    self.nesterov, self.classic, self.pre_wd, self.tcoef = nesterov, classic, pre_wd, tcoef
    self.b = self._new_optim_param() if self.momentum else []

  def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
@@ -98,7 +112,7 @@ class LARS(Optimizer):
        r2 = g.square().sum().sqrt()
        r:Tensor|float = (r1 > 0).where((r2 > 0).where(self.tcoef * r1 / (r2 + self.wd * r1), 1.0), 1.0)
      else: r = 1.0
-      if self.wd > 0: g = g + self.wd * t.detach()
+      if self.pre_wd and self.wd > 0: g = g + self.wd * t.detach()
      # classic momentum does post learning rate update
      if self.classic: g = g * r * self.lr
      if self.momentum:
@@ -106,6 +120,9 @@ class LARS(Optimizer):
        # the scheduler should detect this and just insert contiguous
        self.b[i].assign(self.momentum * self.b[i].contiguous() + g)  # NOTE: self.b[i] is zero on the first run, no if required
        g = (g + self.momentum * self.b[i]) if self.nesterov else self.b[i]
+      if self.ns_params: g = g.reshape(g.shape[0], -1).newton_schulz(self.ns_steps, self.ns_params).reshape(g.shape)
+      # muon does post momentum weight decay
+      if not self.pre_wd and self.wd > 0: t = t.detach() * (1.0 - self.wd * self.lr)
      # popular momentum does pre learning rate update
      if not self.classic: g = g * r * self.lr
      ret.append((t.detach() - g).cast(t.dtype))
@@ -99,6 +99,7 @@ class CStyleLanguage(Renderer):
  code_for_op: dict = {
    Ops.SQRT: lambda x,dtype: f"sqrt({x})", Ops.RECIP: lambda x,dtype: f"(1/{x})", Ops.NEG: lambda x,dtype: f"-{x}",
    Ops.EXP2: lambda x,dtype: f"exp2({x})", Ops.LOG2: lambda x,dtype: f"log2({x})", Ops.SIN: lambda x,dtype: f"sin({x})",
+    Ops.TRUNC: lambda x,dtype: f"trunc({x})",
    Ops.AND: lambda a,b,dtype: f"({a}&{b})", Ops.XOR: lambda a,b,dtype: f"({a}^{b})", Ops.OR: lambda a,b,dtype: f"({a}|{b})",
    Ops.ADD: lambda a,b,dtype: f"({a}+{b})", Ops.SUB: lambda a,b,dtype: f"({a}-{b})", Ops.MUL: lambda a,b,dtype: f"({a}*{b})",
    Ops.MOD: lambda a,b,dtype: f"({a}%{b})", Ops.IDIV: lambda a,b,dtype: f"({a}/{b})", Ops.CMPNE: lambda a,b,dtype: f"({a}!={b})",
@@ -146,7 +147,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}" 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

@@ -200,11 +201,12 @@ class ClangRenderer(CStyleLanguage):
  # language options
  buffer_suffix = " restrict"
  type_map = {dtypes.bool:"_Bool", dtypes.half:"__fp16"}
-  code_for_op = {**({k:v for k,v in CStyleLanguage.code_for_op.items() if k not in [Ops.EXP2, Ops.SIN, Ops.LOG2]}),
-                 Ops.SQRT: lambda x,dtype: f"__builtin_sqrt({x})" if dtype == dtypes.float64 else f"__builtin_sqrtf({x})"}
+  code_for_op = {**({k:v for k,v in CStyleLanguage.code_for_op.items() if k not in [Ops.EXP2, Ops.SIN, Ops.LOG2, Ops.TRUNC]}),
+                 Ops.SQRT: lambda x,dtype: f"__builtin_sqrt({x})" if dtype == dtypes.float64 else f"__builtin_sqrtf({x})",
+                 Ops.TRUNC: lambda x,dtype: f"__builtin_trunc({x})" if dtype == dtypes.float64 else f"__builtin_truncf({x})"}
  # LLVM legalizes double => half cast on systems that don't support it natively (like x86 cpus without AVX512-FP16) into a compiler-rt libcall.
  extra_matcher = PatternMatcher([(UPat.var("x", dtypes.float64).cast(dtypes.float16), lambda x: x.cast(dtypes.float32).cast(dtypes.float16)),
-    (UPat(Ops.SQRT, name="alu"), no_vectorized_alu),]) + CStyleLanguage.extra_matcher
+    (UPat((Ops.SQRT, Ops.TRUNC), name="alu"), no_vectorized_alu),]) + CStyleLanguage.extra_matcher

  if sys.platform == 'win32':
    kernel_typedef = "__attribute__((ms_abi)) void"
@@ -414,7 +416,7 @@ class AMDRenderer(CStyleLanguage):
  ockl = [(f"__ockl_get_{name}", "unsigned int", "size_t", "const") for name in ["local_id", "group_id", "local_size"]]
  ocml = [(f"__ocml_{name}_f{n}", f"{dt}, {dt}" if "fmax" == name else dt, dt, atr)
            for dt, n in [(dtype.name, dtype.itemsize * 8) for dtype in [dtypes.float, dtypes.double, dtypes.half]]
-            for name, atr in [("fmax", "const"), ("exp2", "pure"), ("log2", "pure"), ("sqrt", "const"), ("sin", "")]]
+            for name, atr in [("fmax", "const"), ("exp2", "pure"), ("log2", "pure"), ("sqrt", "const"), ("sin", ""), ("trunc", "")]]

  kernel_typedef = "\n".join(f'extern "C" __attribute__((device{f", {atr}" if atr else ""})) {dto} {meth}({dti});' for meth,dti,dto,atr in ockl+ocml)
  # https://clang.llvm.org/docs/AttributeReference.html#amdgpu-flat-work-group-size
@@ -423,6 +425,7 @@ class AMDRenderer(CStyleLanguage):
  code_for_workitem = {"g": lambda x: f"__ockl_get_group_id({x})", "l": lambda x: f"__ockl_get_local_id({x})",
                       "i": lambda x: f"(__ockl_get_group_id({x})*__ockl_get_local_size({x})+__ockl_get_local_id({x}))"}
  code_for_op = { **CStyleLanguage.code_for_op,
+    Ops.TRUNC: lambda x,dtype: f"__ocml_trunc_f{ {dtypes.half:16, dtypes.double:64}.get(dtype, 32)}({x})",
    Ops.SIN: lambda x,dtype: f"__ocml_sin_f{ {dtypes.half:16, dtypes.double:64}.get(dtype, 32)}({x})",
    Ops.LOG2: lambda x,dtype: f"__ocml_log2_f{ {dtypes.half:16, dtypes.double:64}.get(dtype, 32)}({x})",
    Ops.EXP2: lambda x,dtype: f"__ocml_exp2_f{ {dtypes.half:16, dtypes.double:64}.get(dtype, 32)}({x})",
@@ -92,6 +92,8 @@ base_rewrite = PatternMatcher([
  # unary/binary/ternary ops
  (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"  {ctx[x]} = bitcast {ldt(x.src[0].dtype)} {ctx[x.src[0]]} to {ldt(x.dtype)}"),
  (UPat(Ops.CAST, name="x"), lambda ctx,x: f"  {ctx[x]} = {lcast(x.src[0].dtype, x.dtype)} {ldt(x.src[0].dtype)} {ctx[x.src[0]]} to {ldt(x.dtype)}"),
+  (UPat(Ops.TRUNC, name="x"),
+   lambda ctx,x: f"  {ctx[x]} = call {ldt(x.dtype)} @llvm.trunc.{ldt(x.dtype.scalar())}({ldt(x.src[0].dtype)} {ctx[x.src[0]]})"),
  (UPat(GroupOp.Binary, name="x"), lambda ctx,x:
   f"  {ctx[x]} = {lop[x.src[0].dtype.scalar()][x.op]} {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ctx[x.src[1]]}"),
  (UPat(Ops.WHERE, name="x"), lambda ctx,x:
@@ -204,12 +206,13 @@ class AMDLLVMRenderer(LLVMRenderer):
  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(Ops.BARRIER), lambda ctx: barrier),
-    (UPat(Ops.CAST, name="x", dtype=dtypes.half.vec(16), src=UPat.var("y", dtypes.half.vec(8))), lambda ctx, x, y: f"  {ctx[x]} = shufflevector "\
-      f"<8 x half> {ctx[y]}, <8 x half> zeroinitializer, <16 x i32> <{', '.join([f'i32 {i}, i32 {j}' for i, j in zip(range(0, 8), range(8, 16))])}>"),
-    (UPat(Ops.CAST, name="x", dtype=dtypes.half.vec(8), src=UPat.var("y", dtypes.half.vec(16))), lambda ctx, x, y:
-      f"  {ctx[x]}= shufflevector <16 x half> {ctx[y]}, <16 x half> undef, <8 x i32> <{', '.join([f'i32 {x}' for x in range(0, 16, 2)])}>"),
  ]) + base_rewrite
-  extra_matcher = LLVMRenderer.extra_matcher
+  extra_matcher = LLVMRenderer.extra_matcher + PatternMatcher([
+    (UPat(Ops.CAST, name="x", dtype=dtypes.half.vec(16), src=UPat.var("y", dtypes.half.vec(8))),
+      lambda x, y: UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(y.gep(i // 2) if i % 2 == 0 else UOp.const(dtypes.half, 0.0) for i in range(16)))),
+    (UPat(Ops.CAST, name="x", dtype=dtypes.half.vec(8), src=UPat.var("y", dtypes.half.vec(16))),
+      lambda x, y: UOp(Ops.VECTORIZE, dtypes.half.vec(8), tuple(y.gep(i * 2) for i in range(8)))),
+  ])
  def _render_footer(self, uops: list[UOp]) -> str:
    # TODO: this is copied from cstyle
    requiredMaxThreadsPerBlock = prod(u.arg[1] for u in uops if u.op is Ops.SPECIAL and u.arg[0][0] == "l")
@@ -6,7 +6,7 @@ from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp
 from tinygrad.dtype import dtypes, DType, PtrDType, AddrSpace
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import CUDARenderer
-from tinygrad.helpers import flatten, get_single_element
+from tinygrad.helpers import flatten, get_single_element, prod

 def render_val(x, dtype):
  if dtypes.is_float(dtype):
@@ -19,6 +19,7 @@ asm_for_op: dict[Ops, Callable] = {
  Ops.RECIP: lambda d,a,dt,name: f"rcp{'.approx' if dtypes.is_float(dt) else ''}.{name} {d}, {a};",
  Ops.EXP2: lambda d,a,dt,name: f"ex2.approx.{name} {d}, {a};", Ops.LOG2: lambda d,a,dt,name: f"lg2.approx.{name} {d}, {a};",
  Ops.SIN: lambda d,a,dt,name: f"sin.approx.{name} {d}, {a};", Ops.SQRT: lambda d,a,dt,name: f"sqrt.approx.{name} {d}, {a};",
+  Ops.TRUNC: lambda d,a,dt,name: f"cvt.rzi.{name}.{name} {d}, {a};",
  Ops.SHR: lambda d,a,b,dt,name: f"shr.{name} {d}, {a}, {b};", Ops.SHL: lambda d,a,b,dt,name: f"shl.b{name[1:]} {d}, {a}, {b};",
  Ops.ADD: lambda d,a,b,dt,name: f"{'or' if dt == dtypes.bool else 'add'}.{name} {d}, {a}, {b};",
  Ops.MUL: lambda d,a,b,dt,name: f"{'and' if dt == dtypes.bool else 'mul'}{'.lo' if dtypes.is_int(dt) else ''}.{name} {d}, {a}, {b};",
@@ -33,11 +34,12 @@ asm_for_op: dict[Ops, Callable] = {
    f"selp.{'b16' if name == 'f16' else name} {d}, {b}, {c}, {a};"
 }

-supports_half = (Ops.EXP2, Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPLT, Ops.WHERE)
+supports_half = (Ops.EXP2, Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPLT, Ops.WHERE, Ops.TRUNC)
 doesnt_support_half: tuple[Ops, ...] = tuple(op for op in asm_for_op.keys() if op not in supports_half)
 ptx_matcher = PatternMatcher([
  # bool CMPNE is XOR, bool CMPLT is XOR+AND (universal makes this slow, this is for renderer only)
  (UPat.var('x', dtype=dtypes.bool).ne(UPat.var('y')), lambda x,y: x^y),
+  (UPat.var('x', dtype=dtypes.bool).alu(Ops.CMPEQ, UPat.var('y')), lambda x,y: (x^y)^True),
  (UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
  # upcast to float32 all the ops that don't support half
  (UPat(doesnt_support_half, dtype=dtypes.half, name="x"),
@@ -150,11 +152,12 @@ class PTXRenderer(Renderer):

  mem_types: dict[DType, str] = {**types, dtypes.int8: "s8", dtypes.uint8: "u8", dtypes.bool: "u8", dtypes.float16: "b16"}

-  def render_kernel(self, kernel, function_name, bufs, regs) -> str:
+  def render_kernel(self, kernel, function_name, bufs, regs, uops) -> str:
    def fmt(line): return line if line[0]=="$" else "\t" + line.replace(" ", "\t" if len(line.split(" ")[0]) > 7 else "\t\t", 1)
    kernel = '\n'.join(map(fmt, [f".reg .{reg.split('_')[-2]} %{reg}<{cnt}>;" for reg,cnt in regs] + kernel + ["ret;"]))
+    launch_bounds = prod(u.arg[1] for u in uops if u.op is Ops.SPECIAL and u.arg[0][0] == "l")
    params = ',\n\t'.join([f".param .{'u64' if dtype.__class__ == PtrDType else self.types[dtype]} {name}" for name,dtype in bufs])
-    return f"{self.kernel_prefix} {function_name}(\n\t{params}\n)\n{{\n{kernel}\n}}"
+    return f"{self.kernel_prefix.format(launch_bounds=launch_bounds)} {function_name} (\n\t{params}\n)\n.maxntid {launch_bounds}\n{{\n{kernel}\n}}"

  def render(self, uops:list[UOp]) -> str:
    kernel:list[str] = []
@@ -221,4 +224,4 @@ class PTXRenderer(Renderer):
      kernel.extend([l] if isinstance(l, str) else l)

      if u.op is Ops.SPECIAL: kernel = [f".reg .u32 %{u.arg[0]};"] + kernel
-    return self.render_kernel(kernel, name, bufs, c.items())
+    return self.render_kernel(kernel, name, bufs, c.items(), uops)
@@ -45,12 +45,12 @@ class AMDComputeQueue(HWQueue):

  def wreg(self, reg:AMDReg, *args:sint, **kwargs:int):
    if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified')
-    if self.pm4.PACKET3_SET_SH_REG_START <= reg.addr < self.pm4.PACKET3_SET_SH_REG_END:
+    if self.pm4.PACKET3_SET_SH_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_SH_REG_END:
      set_packet, set_packet_start = self.pm4.PACKET3_SET_SH_REG, self.pm4.PACKET3_SET_SH_REG_START
-    elif self.pm4.PACKET3_SET_UCONFIG_REG_START <= reg.addr < self.pm4.PACKET3_SET_UCONFIG_REG_START + 2**16-1:
+    elif self.pm4.PACKET3_SET_UCONFIG_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_UCONFIG_REG_START + 2**16-1:
      set_packet, set_packet_start = self.pm4.PACKET3_SET_UCONFIG_REG, self.pm4.PACKET3_SET_UCONFIG_REG_START
-    else: raise RuntimeError(f'Cannot set {reg.name} ({reg.addr}) via pm4 packet')
-    self.pkt3(set_packet, reg.addr - set_packet_start, *(args or (reg.encode(**kwargs),)))
+    else: raise RuntimeError(f'Cannot set {reg.name} ({reg.addr[0]}) via pm4 packet')
+    self.pkt3(set_packet, reg.addr[0] - set_packet_start, *(args or (reg.encode(**kwargs),)))

  @contextlib.contextmanager
  def pred_exec(self, xcc_mask:int):
@@ -119,8 +119,8 @@ class AMDComputeQueue(HWQueue):

  def memory_barrier(self):
    pf = '' if self.nbio.version[0] == 2 else '0' if self.nbio.version[:2] != (7, 11) else '1'
-    self.wait_reg_mem(reg_req=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr,
-                      reg_done=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr, value=0xffffffff)
+    self.wait_reg_mem(reg_req=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
+                      reg_done=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff)
    self.acquire_mem()
    return self

@@ -190,17 +190,17 @@ class AMDComputeQueue(HWQueue):
      self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, instance_broadcast_writes=1)
      # Wait for FINISH_PENDING==0
      self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, self.pm4.WAIT_REG_MEM_FUNCTION(WAIT_REG_MEM_FUNCTION_EQ),
-                self.gc.regSQ_THREAD_TRACE_STATUS.addr, 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), 4)
+                self.gc.regSQ_THREAD_TRACE_STATUS.addr[0], 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), 4)
      # Wait for FINISH_DONE!=0
      self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, self.pm4.WAIT_REG_MEM_FUNCTION(WAIT_REG_MEM_FUNCTION_NEQ),
-                self.gc.regSQ_THREAD_TRACE_STATUS.addr, 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_done'), 4)
+                self.gc.regSQ_THREAD_TRACE_STATUS.addr[0], 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_done'), 4)
      # Disable SQTT
      self.sqtt_config(tracing=False)
      # Wait for BUSY==0
      self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, self.pm4.WAIT_REG_MEM_FUNCTION(WAIT_REG_MEM_FUNCTION_EQ),
-                self.gc.regSQ_THREAD_TRACE_STATUS.addr, 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('busy'), 4)
+                self.gc.regSQ_THREAD_TRACE_STATUS.addr[0], 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('busy'), 4)
      # Copy WPTR to memory (src_sel = perf, dst_sel = tc_l2, wr_confirm = True)
-      self.pkt3(self.pm4.PACKET3_COPY_DATA, 1 << 20 | 2 << 8 | 4, self.gc.regSQ_THREAD_TRACE_WPTR.addr, 0, *data64_le(wptrs.va_addr+(se*4)))
+      self.pkt3(self.pm4.PACKET3_COPY_DATA, 1 << 20 | 2 << 8 | 4, self.gc.regSQ_THREAD_TRACE_WPTR.addr[0], 0, *data64_le(wptrs.va_addr+(se*4)))
    # Restore global broadcasting
    self.wreg(self.gc.regGRBM_GFX_INDEX, se_broadcast_writes=1, sa_broadcast_writes=1, instance_broadcast_writes=1)
    self.spi_config(tracing=False)
@@ -539,10 +539,10 @@ class KFDIface:
    self.props = {(p:=l.split())[0]: int(p[1]) for l in FileIOInterface(f"{kfd_topo_path}/{KFDIface.gpus[device_id]}/properties").read().splitlines()}
    ip_base = f"/sys/class/drm/renderD{self.props['drm_render_minor']}/device/ip_discovery/die/0"
    id2ip = {am.GC_HWID: am.GC_HWIP, am.SDMA0_HWID: am.SDMA0_HWIP, am.NBIF_HWID: am.NBIF_HWIP}
-    self.ip_versions = {id2ip[int(hwid)]:tuple(int(FileIOInterface(f'{ip_base}/{hwid}/0/{part}').read()) for part in ['major', 'minor', 'revision'])
-                        for hwid in FileIOInterface(ip_base).listdir() if hwid.isnumeric() and int(hwid) in id2ip}
-    self.ip_offsets = {id2ip[int(hwid)]:tuple(int(x, 16) for x in FileIOInterface(f'{ip_base}/{hwid}/0/base_addr').read().splitlines())
-                        for hwid in FileIOInterface(ip_base).listdir() if hwid.isnumeric() and int(hwid) in id2ip}
+    ip_hw = [(id2ip[int(hwid)], int(hwid)) for hwid in FileIOInterface(ip_base).listdir() if hwid.isnumeric() and int(hwid) in id2ip]
+    self.ip_versions = {ip:tuple(int(FileIOInterface(f'{ip_base}/{hw}/0/{part}').read()) for part in ['major','minor','revision']) for ip,hw in ip_hw}
+    self.ip_offsets = {ip:{int(i):tuple(int(x, 16) for x in FileIOInterface(f'{ip_base}/{hw}/{i}/base_addr').read().splitlines())
+      for i in FileIOInterface(f'{ip_base}/{hw}').listdir()} for ip,hw in ip_hw }
    self.drm_fd = FileIOInterface(f"/dev/dri/renderD{self.props['drm_render_minor']}", os.O_RDWR)

    kfd.AMDKFD_IOC_ACQUIRE_VM(KFDIface.kfd, drm_fd=self.drm_fd.fd, gpu_id=self.gpu_id)
@@ -653,8 +653,7 @@ class PCIIface(PCIIfaceBase):

  def _setup_adev(self, name, vram:MMIOInterface, doorbell:MMIOInterface, mmio:MMIOInterface, dma_regions:list[tuple[int, MMIOInterface]]|None=None):
    self.dev_impl:AMDev = AMDev(name, vram, doorbell, mmio, dma_regions)
-    self.ip_versions = self.dev_impl.ip_ver
-    self.ip_offsets = {hwip: tuple(instances[0]) for hwip,instances in self.dev_impl.regs_offset.items()}
+    self.ip_offsets, self.ip_versions = self.dev_impl.regs_offset, self.dev_impl.ip_ver

    gfxver = int(f"{self.dev_impl.ip_ver[am.GC_HWIP][0]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][1]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][2]:02d}")
    array_count = self.dev_impl.gc_info.gc_num_sa_per_se * self.dev_impl.gc_info.gc_num_se
@@ -762,7 +761,7 @@ class AMDDevice(HCQCompiled):

    nbio_name = 'nbio' if self.target[0] < 12 else 'nbif'
    nbio_pad = (0,) if self.target[0] == 9 else ()
-    self.nbio = AMDIP(nbio_name, self.iface.ip_versions[am.NBIF_HWIP], nbio_pad+self.iface.ip_offsets[am.NBIF_HWIP])
+    self.nbio = AMDIP(nbio_name, self.iface.ip_versions[am.NBIF_HWIP], {i:nbio_pad+x for i,x in self.iface.ip_offsets[am.NBIF_HWIP].items()})

    self.compute_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE, 0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size, debug_memory_size=debug_memory_size)
@@ -1,12 +1,11 @@
-from tinygrad.device import Compiled, Compiler, Renderer, Allocator
-from tinygrad.uop.ops import Ops
+from tinygrad.device import Compiled, Compiler, Allocator
 from tinygrad.engine.jit import MultiGraphRunner
+from tinygrad.renderer.cstyle import CStyleLanguage

-class NullRenderer(Renderer):
+class NullRenderer(CStyleLanguage):
  device = "NULL"
-  code_for_op = {k:lambda:None for k in [Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT]}
  has_local = False
-  def render(self, uops:list) -> str: return ""
+  float4 = "float4"

 class NullProgram:
  def __init__(self, name:str, lib:bytes): pass
@@ -61,7 +61,8 @@ class PythonProgram:
          i += 1
          continue
        if uop in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
-          assert dtype.fmt is not None and isinstance(dtype, PtrDType)
+          assert isinstance(dtype, PtrDType), dtype
+          if dtype.fmt is None: raise RuntimeError(f"{dtype=} is not supported")
          if TYPE_CHECKING or sys.version_info < (3, 12): assert dtype.fmt != "e"
          if uop is Ops.DEFINE_REG:
            # REGs are per thread
@@ -10,16 +10,16 @@ from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU

 AM_DEBUG = getenv("AM_DEBUG", 0)

-@dataclasses.dataclass(frozen=True)
+@dataclasses.dataclass
 class AMRegister(AMDReg):
  adev:AMDev

-  def read(self): return self.adev.rreg(self.addr)
-  def read_bitfields(self) -> dict[str, int]: return self.decode(self.read())
+  def read(self, inst=0): return self.adev.rreg(self.addr[inst])
+  def read_bitfields(self, inst=0) -> dict[str, int]: return self.decode(self.read(inst=inst))

-  def write(self, _am_val:int=0, **kwargs): self.adev.wreg(self.addr, _am_val | self.encode(**kwargs))
+  def write(self, _am_val:int=0, inst=0, **kwargs): self.adev.wreg(self.addr[inst], _am_val | self.encode(**kwargs))

-  def update(self, **kwargs): self.write(self.read() & ~self.fields_mask(*kwargs.keys()), **kwargs)
+  def update(self, inst=0, **kwargs): self.write(self.read(inst=inst) & ~self.fields_mask(*kwargs.keys()), inst=inst, **kwargs)

 class AMFirmware:
  def __init__(self, adev):
@@ -29,7 +29,7 @@ class AMFirmware:
    # Load SOS firmware
    self.sos_fw = {}

-    blob, sos_hdr = self.load_fw(f"psp_{fmt_ver(am.MP0_HWIP)}_sos.bin", am.struct_psp_firmware_header_v2_0)
+    blob, sos_hdr = self.load_fw(f"psp_{fmt_ver(am.MP0_HWIP)}_sos.bin", versioned_header='struct_psp_firmware_header')
    fw_bin = sos_hdr.psp_fw_bin

    for fw_i in range(sos_hdr.psp_fw_bin_count):
@@ -45,11 +45,11 @@ class AMFirmware:
    self.smu_psp_desc = self.desc(blob, hdr.header.ucode_array_offset_bytes, hdr.header.ucode_size_bytes, am.GFX_FW_TYPE_SMU)

    # SDMA firmware
-    blob, hdr, hdr_v3 = self.load_fw(f"sdma_{fmt_ver(am.SDMA0_HWIP)}.bin", am.struct_sdma_firmware_header_v2_0, am.struct_sdma_firmware_header_v3_0)
+    blob, hdr = self.load_fw(f"sdma_{fmt_ver(am.SDMA0_HWIP)}.bin", versioned_header='struct_sdma_firmware_header')
    if hdr.header.header_version_major < 3:
      self.descs += [self.desc(blob, hdr.ctl_ucode_offset, hdr.ctl_ucode_size_bytes, am.GFX_FW_TYPE_SDMA_UCODE_TH1)]
      self.descs += [self.desc(blob, hdr.header.ucode_array_offset_bytes, hdr.ctx_ucode_size_bytes, am.GFX_FW_TYPE_SDMA_UCODE_TH0)]
-    else: self.descs += [self.desc(blob, hdr_v3.header.ucode_array_offset_bytes, hdr_v3.ucode_size_bytes, am.GFX_FW_TYPE_SDMA_UCODE_TH0)]
+    else: self.descs += [self.desc(blob, hdr.header.ucode_array_offset_bytes, hdr.ucode_size_bytes, am.GFX_FW_TYPE_SDMA_UCODE_TH0)]

    # PFP, ME, MEC firmware
    for (fw_name, fw_cnt) in ([('PFP', 1), ('ME', 1)] if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0) else []) + [('MEC', 1)]:
@@ -83,10 +83,13 @@ class AMFirmware:

    self.descs += [self.desc(blob, hdr0.header.ucode_array_offset_bytes, hdr0.header.ucode_size_bytes, am.GFX_FW_TYPE_RLC_G)]

-  def load_fw(self, fname:str, *headers):
+  def load_fw(self, fname:str, *headers, versioned_header:str|None=None):
    fpath = fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/45f59212aebd226c7630aff4b58598967c0c8c91/amdgpu/{fname}", subdir="fw")
    blob = memoryview(bytearray(fpath.read_bytes()))
    if AM_DEBUG >= 1: print(f"am {self.adev.devfmt}: loading firmware {fname}: {hashlib.sha256(blob).hexdigest()}")
+    if versioned_header:
+      chdr = am.struct_common_firmware_header.from_address(mv_address(blob))
+      headers += (getattr(am, versioned_header + f"_v{chdr.header_version_major}_{chdr.header_version_minor}"),)
    return tuple([blob] + [hdr.from_address(mv_address(blob)) for hdr in headers])

  def desc(self, blob:memoryview, offset:int, size:int, *types:int) -> tuple[list[int], memoryview]: return (list(types), blob[offset:offset+size])
@@ -223,7 +226,7 @@ class AMDev(PCIDevImplBase):

    self.bhdr = am.struct_binary_header.from_buffer(bytearray(self.vram.view(tmr_offset, tmr_size)[:]))
    ihdr = am.struct_ip_discovery_header.from_address(ctypes.addressof(self.bhdr) + self.bhdr.table_list[am.IP_DISCOVERY].offset)
-    assert ihdr.signature == am.DISCOVERY_TABLE_SIGNATURE and not ihdr.base_addr_64_bit, f"0x{ihdr.signature:X} != 0x{am.DISCOVERY_TABLE_SIGNATURE:X}"
+    assert self.bhdr.binary_signature == am.BINARY_SIGNATURE and ihdr.signature == am.DISCOVERY_TABLE_SIGNATURE, "discovery signatures mismatch"

    # Mapping of HW IP to Discovery HW IP
    hw_id_map = {am.__dict__[x]: int(y) for x,y in am.hw_id_map}
@@ -254,6 +257,5 @@ class AMDev(PCIDevImplBase):
      ("nbio" if self.ip_ver[am.GC_HWIP] < (12,0,0) else "nbif", am.NBIO_HWIP)]

    for prefix, hwip in mods:
-      self.__dict__.update(import_asic_regs(prefix, self.ip_ver[hwip], cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[hwip][0])))
-    self.__dict__.update(import_asic_regs('mp', (11, 0), cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[am.MP1_HWIP][0])))
-
+      self.__dict__.update(import_asic_regs(prefix, self.ip_ver[hwip], cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[hwip])))
+    self.__dict__.update(import_asic_regs('mp', (11, 0), cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[am.MP1_HWIP])))
@@ -249,7 +249,7 @@ class AM_GFX(AM_IP):
    self._grbm_select(me=1, pipe=pipe, queue=queue)

    mqd_st_mv = to_mv(ctypes.addressof(mqd_struct), ctypes.sizeof(mqd_struct)).cast('I')
-    for i, reg in enumerate(range(self.adev.regCP_MQD_BASE_ADDR.addr, self.adev.regCP_HQD_PQ_WPTR_HI.addr + 1)):
+    for i, reg in enumerate(range(self.adev.regCP_MQD_BASE_ADDR.addr[0], self.adev.regCP_HQD_PQ_WPTR_HI.addr[0] + 1)):
      self.adev.wreg(reg, mqd_st_mv[0x80 + i])
    self.adev.regCP_HQD_ACTIVE.write(0x1)

@@ -459,7 +459,7 @@ class AM_PSP(AM_IP):
    wait_cond(lambda: self.adev.reg(f"{self.reg_pref}_64").read() & 0x8000FFFF, value=0x80000000, msg="sOS ring not created")

  def _ring_submit(self, cmd:am.struct_psp_gfx_cmd_resp) -> am.struct_psp_gfx_cmd_resp:
-    msg = am.struct_psp_gfx_rb_frame(fence_value=(prev_wptr:=self.adev.reg(f"{self.reg_pref}_67").read()),
+    msg = am.struct_psp_gfx_rb_frame(fence_value=(prev_wptr:=self.adev.reg(f"{self.reg_pref}_67").read()) + 1,
      cmd_buf_addr_lo=lo32(self.adev.paddr2mc(self.cmd_paddr)), cmd_buf_addr_hi=hi32(self.adev.paddr2mc(self.cmd_paddr)),
      fence_addr_lo=lo32(self.adev.paddr2mc(self.fence_paddr)), fence_addr_hi=hi32(self.adev.paddr2mc(self.fence_paddr)))

@@ -469,8 +469,7 @@ class AM_PSP(AM_IP):
    # Move the wptr
    self.adev.reg(f"{self.reg_pref}_67").write(prev_wptr + ctypes.sizeof(am.struct_psp_gfx_rb_frame) // 4)

-    while self.adev.vram.view(self.fence_paddr, 4, 'I')[0] != prev_wptr: pass
-    time.sleep(0.005)
+    wait_cond(lambda: self.adev.vram.view(self.fence_paddr, 4, 'I')[0], value=msg.fence_value, msg="sOS ring not responding")

    resp = type(cmd).from_buffer(bytearray(self.adev.vram.view(self.cmd_paddr, ctypes.sizeof(cmd))[:]))
    if resp.resp.status != 0: raise RuntimeError(f"PSP command failed {resp.cmd_id} {resp.resp.status}")
@@ -5,9 +5,10 @@ from tinygrad.helpers import getbits, round_up, fetch
 from tinygrad.runtime.autogen import pci
 from tinygrad.runtime.support.usb import ASM24Controller

-@dataclass(frozen=True)
+@dataclass
 class AMDReg:
-  name:str; offset:int; segment:int; fields:dict[str, tuple[int, int]]; bases:tuple[int, ...] # noqa: E702
+  name:str; offset:int; segment:int; fields:dict[str, tuple[int, int]]; bases:dict[int, tuple[int, ...]] # noqa: E702
+  def __post_init__(self): self.addr:dict[int, int] = { inst: bases[self.segment] + self.offset for inst, bases in self.bases.items() }

  def encode(self, **kwargs) -> int: return functools.reduce(int.__or__, (value << self.fields[name][0] for name,value in kwargs.items()), 0)
  def decode(self, val: int) -> dict: return {name:getbits(val, start, end) for name,(start,end) in self.fields.items()}
@@ -15,12 +16,9 @@ class AMDReg:
  def fields_mask(self, *names) -> int:
    return functools.reduce(int.__or__, ((((1 << (self.fields[nm][1]-self.fields[nm][0]+1)) - 1) << self.fields[nm][0]) for nm in names), 0)

-  @property
-  def addr(self): return self.bases[self.segment] + self.offset
-
@dataclass
 class AMDIP:
-  name:str; version:tuple[int, ...]; bases:tuple[int, ...] # noqa: E702
+  name:str; version:tuple[int, ...]; bases:dict[int, tuple[int, ...]] # noqa: E702
  def __post_init__(self): self.version = fixup_ip_version(self.name, self.version)[0]

  @functools.cached_property
@@ -82,9 +82,9 @@ class PCIDevice:
      FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver/unbind", os.O_WRONLY).write(self.pcibus)

    for i in resize_bars or []:
-      supported_sizes = int(FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource{i}_resize", os.O_RDONLY).read(), 16)
-      try: FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/resource{i}_resize", os.O_RDWR).write(str(supported_sizes.bit_length() - 1))
-      except OSError as e: raise RuntimeError(f"Cannot resize BAR {i}: {e}. Ensure the resizable BAR option is enabled on your system.") from e
+      if FileIOInterface.exists(rpath:=f"/sys/bus/pci/devices/{self.pcibus}/resource{i}_resize"):
+        try: FileIOInterface(rpath, os.O_RDWR).write(str(int(FileIOInterface(rpath, os.O_RDONLY).read(), 16).bit_length() - 1))
+        except OSError as e: raise RuntimeError(f"Cannot resize BAR {i}: {e}. Ensure the resizable BAR option is enabled on your system.") from e

    if getenv("VFIO", 0) and (vfio_fd:=System.vfio()) is not None:
      FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver_override", os.O_WRONLY).write("vfio-pci")
@@ -48,7 +48,7 @@ def copy_reorder_view(copy:UOp, view:UOp, base:UOp):
  if prod(view.shape) < prod(base.shape): return view.contiguous().copy_to_device(copy.device)
  return base.copy_to_device(copy.device).view(view.arg)

-sym = symbolic_simple+PatternMatcher([
+kernelize_sym = symbolic_simple+PatternMatcher([
  # UOp with size 0 is zero
  (UPat(GroupOp.All-{Ops.SINK}, name="root"), lambda root: root.const_like(0) if root.base.st is not None and root.size == 0 else None),
  # DETACH and CONTIGUOUS_BACKWARD are NOOPs here
@@ -327,7 +327,7 @@ def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]:
  """

  # multi + merge_views + simplify
-  tensor_map = graph_rewrite_map(sink, multi_pm+do_fuse+merge_views+sym+replace_contiguous, ctx={}, name="merge_views")
+  tensor_map = graph_rewrite_map(sink, multi_pm+do_fuse+merge_views+kernelize_sym+replace_contiguous, ctx={}, name="merge_views")

  # display the cleaned up tensor graph
  if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Tensor Graph")
@@ -13,6 +13,10 @@ def handle_allreduce_multirank(buf:UOp, red:UOp) -> UOp|None:
  for i,dev in enumerate(buf.device):
    groups.setdefault(Device[dev].group_id, []).append(buf.mselect(i))

+  # Put reduce leader of each group first
+  reduce_leaders = set(getenv("REDUCE_LEADERS", "").split(","))
+  groups = {gid: sorted(bufs, key=lambda x: (x.device not in reduce_leaders, x.device)) for gid,bufs in groups.items()}
+
  # Skip if only one group or if every group has only one buffer
  if len(groups) <= 1 or not any(len(g) > 1 for g in groups.values()): return None

@@ -252,7 +252,7 @@ class Tensor(MathTrait):
    # create the schedule
    schedule, var_vals = create_schedule_with_vars(sink)
    schedule = memory_planner(schedule)
-    if DEBUG >= 1 and len(schedule) >= 10: print(f"scheduled {len(schedule)} kernels in {(time.perf_counter()-st)*1000:.2f} ms")
+    if DEBUG >= 1 and len(schedule) > 1: print(f"scheduled {len(schedule)} kernels in {(time.perf_counter()-st)*1000:.2f} ms")
    return schedule, var_vals

  def schedule(self, *lst:Tensor) -> list[ScheduleItem]:
@@ -345,6 +345,7 @@ class Tensor(MathTrait):
    print(t.tolist())
    ```
    """
+    if self.dtype in (dtypes.bfloat16, *dtypes.fp8s): return self.cast(dtypes.float32).tolist()
    return self.data().tolist()

  def numpy(self) -> 'np.ndarray':  # type: ignore [name-defined] # noqa: F821
@@ -2933,11 +2934,11 @@ class Tensor(MathTrait):
    """
    return self*-1 if self.dtype != dtypes.bool else self.logical_not()

-  def contiguous(self) -> Tensor:
+  def contiguous(self, **kwargs) -> Tensor:
    """
    Returns a contiguous tensor.
    """
-    return self._apply_uop(UOp.contiguous)
+    return self._apply_uop(UOp.contiguous, **kwargs)

  def fuse(self) -> Tensor:
    """
@@ -3143,7 +3144,7 @@ class Tensor(MathTrait):
    print(Tensor([-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]).trunc().numpy())
    ```
    """
-    return self.cast(dtypes.int32).cast(self.dtype)
+    return self._apply_uop(UOp.trunc)

  def ceil(self: Tensor) -> Tensor:
    """
@@ -3173,7 +3174,7 @@ class Tensor(MathTrait):
    print(Tensor([-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]).round().numpy())
    ```
    """
-    return ((self > 0) == ((b := self.cast(dtypes.int32) / 2.0).cast(dtypes.int32) == b)).where((self - 0.5).ceil(), (self + 0.5).floor())
+    return ((self > 0) == ((b := self.trunc() / 2.0).trunc() == b)).where((self - 0.5).ceil(), (self + 0.5).floor())

  def isinf(self:Tensor, detect_positive:bool=True, detect_negative:bool=True) -> Tensor:
    """
@@ -4033,6 +4034,21 @@ class Tensor(MathTrait):
    nll = -self.gather(1, Y.unsqueeze(1)).squeeze(1) * masked_weight
    return nll.sum() / masked_weight.sum() if reduction == "mean" else nll._do_reduction(reduction)

+  def newton_schulz(self, steps:int, params:tuple[int, ...], eps:float=1.0e-7) -> Tensor:
+    """
+    Performs the newton-schulz algorithm for odd polynomials. The degree of the odd polynomial depends on the number of params.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor.randn(4, 4)
+    print(t.newton_schulz(steps=5, params=(2,-1.5,0.5)).numpy())
+    ```
+    """
+    assert self.ndim > 1, "NS only works for two or more dims"
+    G = self / (self.square().sum(axis=(-2, -1), keepdim=True).sqrt() + eps)
+    G = G.transpose(-2, -1) if self.shape[-2] > self.shape[-1] else G
+    for _ in range(steps): G = sum(p * functools.reduce(lambda x, y: (y @ y.transpose(-2, -1)) @ x, [G]*i, G) for i,p in enumerate(params))
+    return G.transpose(-2, -1) if self.shape[-2] > self.shape[-1] else G
+
  def qr(self) -> tuple[Tensor, Tensor]:
    assert self.ndim > 1, f"expected two or more dimensions, got {self.ndim}"
    R = self.clone()
@@ -9,14 +9,17 @@ class FastEnum(IntEnum):
 # the order of these Ops controls the order of the toposort
 class Ops(FastEnum):
  # uops that aren't rendered
-  NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto()  # noqa: E702
+  NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto(); REWRITE_ERROR = auto()  # noqa: E702

  # track children
-  CHILD = auto()
+  CHILD = auto(); CHILDREN = auto() # noqa: E702

  # buffer ops
  COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702

+  # create buffer
+  BUFFERIZE = auto()
+
  # ops that adjust the behavior of the scheduler
  CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); FUSE = auto() # noqa: E702

@@ -49,7 +52,7 @@ class Ops(FastEnum):
  UNROLL = auto(); CONTRACT = auto(); GEP = auto(); VECTORIZE = auto(); CAT = auto(); PTRCAT = auto() # noqa: E702

  # UnaryOps
-  CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto(); SQRT = auto(); RECIP = auto(); NEG = auto() # noqa: E702
+  CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto(); SQRT = auto(); RECIP = auto(); NEG = auto(); TRUNC = auto() # noqa: E702

  # load/store before math
  LOAD = auto(); STORE = auto() # noqa: E702
@@ -80,12 +83,15 @@ class Ops(FastEnum):
  CUSTOM = auto(); CUSTOMI = auto() # noqa: E702

 class GroupOp:
-  Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIP, Ops.NEG}
+  Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIP, Ops.NEG, Ops.TRUNC}
  Binary = {Ops.ADD, Ops.MUL, Ops.IDIV, Ops.MAX, Ops.MOD, Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ,
            Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW}
  Ternary = {Ops.WHERE, Ops.MULACC}
  ALU = set.union(Unary, Binary, Ternary)

+  # TODO: is BITCAST always Elementwise if it's shape changing?
+  Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST})
+
  Defines = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}

  Irreducible = {Ops.CONST, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.RANGE}
@@ -1,8 +1,9 @@
+from typing import Callable
 import math, functools
 from tinygrad.dtype import dtypes, DType, promo_lattice
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import polyN
-from tinygrad.uop.ops import UOp
+from tinygrad.helpers import polyN, getenv
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher

 TRANSCENDENTAL_SUPPORTED_DTYPES = (dtypes.float16, dtypes.float32, dtypes.float64)

@@ -79,10 +80,10 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
  intermediate_dtype = dtypes.float32.vec(d.dtype.count) if d.dtype.base.scalar() == dtypes.float16 else d.dtype

  f, e = frexp(d)
-  ia = (f.cast(intermediate_dtype) * 4.294967296e9).cast_vec(dtypes.uint64)
+  ia = (f.cast(intermediate_dtype) * 4.294967296e9).cast(dtypes.uint64)
  # extract 96 relevant bits of 2/pi based on magnitude of argument
-  i = shr(e.cast_vec(dtypes.uint64), 5)
-  e = e.cast_vec(dtypes.int32) & 31
+  i = shr(e.cast(dtypes.uint64), 5)
+  e = e.cast(dtypes.int32) & 31
  offset = 32 - e

  def _take(an:UOp, offset:int, count:int=0) -> UOp:
@@ -90,8 +91,8 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
    if count+offset < len(two_over_pi_f) - 1:
      an = i.ne(count).where(_take(an, offset, count=count+1), an.const_like(two_over_pi_f[count+offset]))
    return an
-  def _shl_lazy(x, y): return (x.cast_vec(dtypes.uint64) * pow2if(y, d.dtype).cast_vec(dtypes.uint64)).cast_vec(dtypes.uint32)
-  def _shr_lazy(x, y): return (x.cast_vec(dtypes.uint64) // pow2if(y, d.dtype).cast_vec(dtypes.uint64)).cast_vec(dtypes.uint32)
+  def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) * pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
+  def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) // pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)

  a = [_take(UOp.const(dtypes.uint32.vec(d.dtype.count), 0), i) for i in range(4)]
  #  (two_over_pi_f[Int(i) + n] << e) | (two_over_pi_f[Int(i) + n+1] >> (nbits - e))
@@ -100,12 +101,12 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
  mi = _shl_lazy(a[1], e) | _shr_lazy(a[2], offset)
  lo = _shl_lazy(a[2], e) | _shr_lazy(a[3], offset)

-  def _hp_mul(x:UOp, y:UOp) -> UOp: return x.cast_vec(dtypes.uint64) * y.cast_vec(dtypes.uint64)
+  def _hp_mul(x:UOp, y:UOp) -> UOp: return x.cast(dtypes.uint64) * y.cast(dtypes.uint64)
  # compute x * 2/pi
  p = shl(_hp_mul(ia, hi), 32) + _hp_mul(ia, mi) + shr(_hp_mul(ia, lo), 32)

  # round quotient to nearest
-  q = shr(p, 62).cast_vec(dtypes.int32)
+  q = shr(p, 62).cast(dtypes.int32)
  p = p & 0x3fffffffffffffff
  r = (p.cast(intermediate_dtype) * (3.4061215800865545e-19)).cast(d.dtype)

@@ -132,7 +133,7 @@ def cody_waite_reduction(d:UOp) -> tuple[UOp, UOp]:
      d = (qdh + q) * -PI_D + d
    elif x.dtype.scalar() == dtypes.float16:
      # [FIXME] when reducing `d`, FP16 needs FP32 precision to achieve 1.0 ULP precision.
-      d = _reduce_d(x.cast_vec(dtypes.float32), q.cast_vec(dtypes.float32)).cast_vec(dtypes.float16)
+      d = _reduce_d(x.cast(dtypes.float32), q.cast(dtypes.float32)).cast(dtypes.float16)
    else:
      # https://github.com/shibatch/sleef/blob/4e08851f59fc2b545f9c393c6a23dfd311a26308/src/libm/sleefsp.c#L464-L503
      d = q * -3.1414794921875 + x
@@ -142,9 +143,9 @@ def cody_waite_reduction(d:UOp) -> tuple[UOp, UOp]:
    return d

  m_1_pi = 0.318309886183790671537767526745028724
-  qdh = (d * (m_1_pi / 2.0**24)).cast_vec(dtypes.int64).cast(d.dtype) * (2.0**24)
+  qdh = (d * (m_1_pi / 2.0**24)).cast(dtypes.int64).cast(d.dtype) * (2.0**24)
  quadrant = rintk(d * m_1_pi -qdh) if d.dtype.base.scalar() == dtypes.float64 else rintk(d * m_1_pi)
-  return _reduce_d(d, quadrant.cast(d.dtype)), quadrant.cast_vec(dtypes.int32)
+  return _reduce_d(d, quadrant.cast(d.dtype)), quadrant.cast(dtypes.int32)

 # *** approximate sine on small angle. ***
 def trig_poly(d:UOp, coeff32, coeff64): return d * (polyN(d*d, coeff64) if d.dtype.scalar() == dtypes.float64 else polyN(d*d, coeff32))
@@ -223,7 +224,7 @@ def xlog2(d:UOp) -> UOp:
  """
  assert d.dtype.scalar() in TRANSCENDENTAL_SUPPORTED_DTYPES
  # TODO: float16 denormal need float32 to achieve precision
-  if d.dtype.scalar() == dtypes.float16: return xlog2(d.cast_vec(dtypes.float32)).cast_vec(dtypes.float16)
+  if d.dtype.scalar() == dtypes.float16: return xlog2(d.cast(dtypes.float32)).cast(dtypes.float16)
  FLT_MIN = d.const_like(1e-6 if d.dtype.scalar() == dtypes.float16 else 1e-4)
  is_denormal = d<FLT_MIN
  a = is_denormal.where(d * (2 ** 64), d)
@@ -260,9 +261,9 @@ def xpow(base:UOp, exponent:UOp) -> UOp:
  # start with b ** e = exp2(e * log2(b))
  ret = (base < 0).where(-base, base).log2().mul(exponent).exp2()
  # negative base adjustment: nan for non-integer exponent and -1 for odd exponent
-  non_int = exponent != exponent.cast_vec(dtypes.int32).cast(exponent.dtype)
+  non_int = exponent != exponent.cast(dtypes.int32).cast(exponent.dtype)
  adj = non_int.where(ret.const_like(math.nan),
-    (exponent < 0).where(-exponent, exponent).cast_vec(dtypes.int32).mod(2).cast_vec(dtypes.bool).where(ret.const_like(-1), ret.const_like(1)))
+    (exponent < 0).where(-exponent, exponent).cast(dtypes.int32).mod(2).cast(dtypes.bool).where(ret.const_like(-1), ret.const_like(1)))
  # fix 0 ** 0 = 1
  return (base.eq(0) & exponent.eq(0)).where(ret.const_like(1), ret * (base < 0).where(adj, ret.const_like(1)))

@@ -292,3 +293,59 @@ def fast_idiv(device: str, x: UOp, d: int) -> UOp|None:
    if m*vmin >= dtypes.min(next_dtype) and m*vmax <= dtypes.max(next_dtype):
      return ((x.cast(next_dtype)*m) >> s).cast(x.dtype) if is_unsigned else ((x.cast(next_dtype)*m) >> s).cast(x.dtype) + (x<0).where(x.ufix(1), 0)
  return None
+
+# ***** threefry *****
+
+def threefry2x32(x: UOp, key: UOp):
+  # split x and key from uint64 to two uint32
+  x0, x1 = (x & 0xffffffff).cast(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast(dtypes.uint32)
+  key0, key1 = (key & 0xffffffff).cast(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast(dtypes.uint32)
+
+  rotations = [[13, 15, 26, 6], [17, 29, 16, 24]]
+  ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0]
+  xr:list[UOp] = [x0 + ks[-1], x1 + ks[0]]
+  for i in range(5):
+    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r)))
+    xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)]
+
+  return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64)
+
+# ***** decomposition patterns *****
+
+powers_of_two = {2**i:i for i in range(64)}
+@functools.cache
+def get_late_rewrite_patterns(ops:tuple[Ops, ...], force_transcendental=False):
+  pat: list[tuple[UPat, Callable]] = [(UPat(op, dtype=TRANSCENDENTAL_SUPPORTED_DTYPES, src=(UPat.var("d"),)), f) for op,f in \
+           ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)) if op not in ops or force_transcendental]
+  # no real hardware supports THREEFRY
+  pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32))
+  # rewrite SQRT to xpow 0.5
+  if Ops.SQRT not in ops: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5))))
+  # rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1)
+  if Ops.AND in ops: pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)]
+  # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
+  if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
+  if Ops.SHR in ops:
+    # no reason to check x<0 for uints
+    pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
+    pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(
+      c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]  # (x+(x<0).where(c-1, 0)) >> v
+    if not getenv("DISABLE_FAST_IDIV"):
+      pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d"), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))]
+      pat += [(UPat.var("x", dtypes.ints)%UPat.var("d"), lambda x, d: x-d*(x//d))]
+  if Ops.NEG in ops:
+    pat += [(UPat.var('x')*-1, lambda x: x.alu(Ops.NEG))]
+    if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda x,y: x.alu(Ops.SUB, y))]
+  if Ops.CMPLT in ops:
+    # These are late rewrites because simplex expects equalities to be a certain format
+    pat += [
+      ((UPat.var("x", dtypes.sints) < UPat.cvar("c", dtypes.sints)).logical_not(), lambda x,c: c-1<x),
+      ((UPat.cvar("c", dtypes.sints) < UPat.var("x", dtypes.sints)).logical_not(), lambda x,c: x<c+1),
+      (UPat.var("x", dtypes.sints)*-1 < UPat.var("y", dtypes.sints)*UPat.cvar("c"), lambda x,y,c: y*(-c)<x),
+      (UPat.var("x", dtypes.sints)*-1 < UPat.cvar("c"), lambda x,c:-c<x),
+      ((UPat.cvar("c1",vec=False)<UPat.var("x", dtypes.sints)) & (UPat.var("x", dtypes.sints)<UPat.cvar("c2",vec=False)),
+        lambda x,c1,c2: x.eq(c1+1) if c1.arg+1==c2.arg-1 else None),  # (c-1)<x & x<(c+1) -> x==c
+    ]
+  if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))]
+  if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
+  return PatternMatcher(pat)
@@ -161,6 +161,7 @@ class MathTrait:
    raise RuntimeError("where needs at least one UOp arg")
  def threefry(self, seed): return self.alu(Ops.THREEFRY, seed)
  def reciprocal(self): return self.alu(Ops.RECIP)
+  def trunc(self): return self.alu(Ops.TRUNC)
  def sqrt(self): return self.alu(Ops.SQRT)
  def sin(self): return self.alu(Ops.SIN)
  def log2(self): return self.alu(Ops.LOG2)
@@ -182,6 +182,19 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
  @property
  def size(self) -> int: return self.arg[0] if self.op is Ops.BUFFER_VIEW else self.arg if self.op is Ops.BUFFER else unwrap(self.st).size

+  # determine what ranges this is in
+  @functools.cached_property
+  def ranges(self) -> dict[UOp, None]:
+    if self.op is Ops.RANGE: return {self:None}
+    if self.op in {Ops.CONTIGUOUS, Ops.REDUCE, Ops.STORE}:
+      ret = self.src[0].ranges.copy()
+      for s in self.src[1:]:
+        if s in ret: del ret[s]
+    else:
+      ret = {}
+      for s in self.src: ret.update(s.ranges)
+    return ret
+
  # *** uop evaluation ***

  def simplify(self):
@@ -219,7 +232,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    return ret
  def sink(self, *srcs:UOp|None, **kwargs): return UOp(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
  def detach(self): return UOp(Ops.DETACH, self.dtype, (self,))
-  def index(self, idx:UOp, valid:UOp|None=None): return UOp(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx))
+  def index(self, *srcs:UOp|None, **kwargs):
+    return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
  def __getitem__(self, idx): return self.index(idx)
  def const_like(self, b:ConstLike):
    # constants can optionally have a DEVICE source
@@ -229,9 +243,10 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    if count == 1: return self
    return UOp(Ops.VECTORIZE, self.dtype.vec(count), (self,)*count)
  def cast(self, dtype:DType):
+    # TODO: we shouldn't have to check for dtype.count == 1 here, but CAST is misused in AMD LLVM
+    if dtype.count == 1 and dtype.count != self.dtype.count: dtype = dtype.vec(self.dtype.count)
    if self.dtype == dtype: return self
    return UOp(Ops.CAST, dtype, (self,))
-  def cast_vec(self, dtype:DType): return UOp(Ops.CAST, dtype.vec(self.dtype.count), (self,))
  def bitcast(self, dtype:DType): return UOp(Ops.BITCAST, dtype, (self,))
  def gep(self, i:tuple[int, ...]|int):
    if isinstance(i, tuple) and len(i) == 1: return self.gep(i[0])
@@ -259,7 +274,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
      from tinygrad.shape.shapetracker import ShapeTracker
      ret = ret.replace(src=(UOp(Ops.VIEW, dtypes.void, (), ShapeTracker.from_shape(shape, (0,)*len(shape))),))
    if device is not None:
-      ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device).view(unwrap(ret.st)),))
+      if shape is not None: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device).view(unwrap(ret.st)),))
+      else: ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device),))
    return ret
  @staticmethod
  def range(dtype:DType, end:sint, idx:int): return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end),), arg=idx)
@@ -275,7 +291,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
    ret = UOp(Ops.REDUCE_AXIS, self.dtype, (ret,), (op, new_axis))
    return ret.reshape(tuple([x if i not in axis else 1 for i,x in enumerate(self.shape)]))
  def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
-  def contiguous(self): return self.alu(Ops.CONTIGUOUS)
+  def contiguous(self, *args, **kwargs): return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
  def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD)
  def fuse(self): return self.alu(Ops.FUSE)
  def allreduce(self, op, device:str|tuple[str, ...]|UOp):
@@ -468,6 +484,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
      if (d0:=self.src[0].divides(v)) is not None: return d0 * self.src[1]
      if (d1:=self.src[1].divides(v)) is not None: return self.src[0] * d1
    return None # generic None if we aren't sure
+  def pop_const(self) -> tuple[UOp, int]: return (self.src[0], self.src[1].arg) if self.op is Ops.ADD and self.src[1].op is Ops.CONST else (self, 0)
  @property
  def vmin(self) -> ConstType: return self._min_max[0]
  @property
@@ -559,7 +576,7 @@ def safe_pow(x, y):
 python_alu: dict[Ops, Callable]  = {
  Ops.LOG2: lambda x: math.log2(x) if x > 0 else -math.inf if x == 0 else math.nan, Ops.EXP2: safe_exp2,
  Ops.SQRT: lambda x: math.sqrt(x) if x >= 0 else math.nan, Ops.RECIP: lambda x: 1/x if x != 0 else math.copysign(math.inf, x),
-  Ops.SIN: lambda x: math.sin(x) if not math.isinf(x) else math.nan, Ops.POW: safe_pow,
+  Ops.SIN: lambda x: math.sin(x) if not math.isinf(x) else math.nan, Ops.POW: safe_pow, Ops.TRUNC: math.trunc,
  Ops.NEG: operator.neg, Ops.ADD: operator.add, Ops.SUB: operator.sub, Ops.MUL: operator.mul, Ops.CMPNE: operator.ne, Ops.CMPLT: operator.lt,
  Ops.XOR: operator.xor, Ops.OR: operator.or_, Ops.AND: operator.and_, Ops.SHR: operator.rshift, Ops.SHL: operator.lshift, Ops.MAX: max,
  Ops.MOD: cmod, Ops.IDIV: cdiv, Ops.MULACC: lambda x,y,z: (x*y)+z, Ops.WHERE: lambda x,y,z: y if x else z, Ops.CMPEQ: operator.eq}
@@ -641,6 +658,9 @@ class UPat(MathTrait):
  @staticmethod
  def const(dtype:DType|tuple[DType, ...]|None, b:ConstType): return UPat(Ops.CONST, dtype=dtype, arg=b)

+  # lil helper
+  def f(self, op, **kwargs): return UPat(op, src=(self,), **kwargs)
+
  # copied from UOp
  def sink(self, *srcs:UPat|None, **kwargs): return UPat(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
  def index(self, idx:UPat, valid:UPat|None=None): return UPat(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx))
@@ -654,6 +674,7 @@ class UPat(MathTrait):
  def reduce(self, *src:UPat, **kwargs): return UPat(Ops.REDUCE, self.dtype, src=(self,)+src, **kwargs)
  def fuse(self): return self.alu(Ops.FUSE)
  def or_broadcasted(self, **kwargs): return UPat.any(self, UPat(Ops.VECTORIZE, self.dtype, src=self, **kwargs))
+  def contiguous(self, *args, **kwargs): return UPat(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)

  def const_like(self, b:ConstLike): return UPat.const(self.dtype, cast(ConstType, b))
  def alu(self, op:Ops, *src:UPat):
@@ -739,16 +760,6 @@ class PatternMatcher:
      if (ret:=match(uop, ctx)) is not None and ret is not uop: return ret
    return None

-  def fixed_point_rewrite(self, uop:UOp, ctx=None) -> UOp:
-    # apply rewrite rules until a fixed point is reached. may return `uop` itself if PatternMatcher doesn't match
-    new_n: UOp|None = uop
-    seen = set()
-    while new_n is not None:
-      if new_n in seen: raise RuntimeError("infinite loop in fixed_point_rewrite")
-      seen.add(new_n)
-      last_n, new_n = new_n, self.rewrite(new_n, ctx)
-    return last_n
-
 # *** non-blocking UOp tracker ***

 ucount = itertools.count()
@@ -840,7 +851,12 @@ class TrackedPatternMatcher(PatternMatcher):
        match_stats[p][2] += time.perf_counter()-st
        continue
      match_stats[p][1] += 1
-      if (ret:=match(uop, ctx)) is not None and ret is not uop:
+      try: ret = match(uop, ctx)
+      except Exception as e:
+        if TRACK_MATCH_STATS >= 2 and active_rewrites and not isinstance(e, RewriteNotReady):
+          active_rewrites[-1].matches.append((track_uop(uop), track_uop(UOp(Ops.REWRITE_ERROR, src=uop.src, arg=str(sys.exc_info()[1]))), p.location))
+        raise
+      if ret is not None and ret is not uop:
        match_stats[p][0] += 1
        match_stats[p][3] += (et:=time.perf_counter()-st)
        if TRACK_MATCH_STATS >= 3: print(f"{et*1e6:7.2f} us -- ", printable(p.location))
@@ -884,10 +900,21 @@ class RewriteNotReady(Exception): pass
 class RewriteContext:
  def __init__(self, pm, bpm, ctx=None):
    self.pm: PatternMatcher|None = pm
+    self.pm_cache: dict[UOp, UOp|None] = {}
    self.bpm: PatternMatcher|None = bpm
+    self.bpm_cache: dict[UOp, UOp|None] = {}
    self.ctx = ctx
    self.replace: dict[UOp, UOp] = {}
-    self.skip_0: dict[UOp, None] = {}  # NOTE: this is needed for RewriteNotReady. it also detects some infinite loops
+
+  def cached_pm_rewrite(self, x:UOp):
+    if (ret:=self.pm_cache.get(x,False)) is not False: return ret
+    ret = self.pm_cache[x] = cast(PatternMatcher, self.pm).rewrite(x, self.ctx)
+    return ret
+
+  def cached_bpm_rewrite(self, x:UOp):
+    if (ret:=self.bpm_cache.get(x,False)) is not False: return ret
+    ret = self.bpm_cache[x] = cast(PatternMatcher, self.bpm).rewrite(x, self.ctx)
+    return ret

  def unified_rewrite(self, root:UOp) -> UOp:
    stack: list[tuple[UOp, int, UOp]] = [(root, 0, root)]
@@ -897,18 +924,23 @@ class RewriteContext:
      if n in self.replace: continue  # skip any nodes we have seen
      try:
        if stage == 0:
-          if n in self.skip_0: continue
          # if bottom up, we rewrite this node early. in both cases, we add its parents to the stack
-          if self.bpm is not None: new_n = self.bpm.fixed_point_rewrite(new_n, self.ctx)
+          if self.bpm is not None:
+            # apply rewrite rules until a fixed point is reached. may return `uop` itself if PatternMatcher doesn't match
+            test_n: UOp|None = n
+            seen = set()
+            while test_n is not None:
+              if test_n in seen: raise RuntimeError("infinite loop in fixed_point_rewrite")
+              seen.add(test_n)
+              new_n, test_n = test_n, self.cached_bpm_rewrite(test_n)
          stack.append((n, 1, new_n))
          for x in reversed(new_n.src): stack.append((x, 0, x))
-          self.skip_0[n] = None
        elif stage == 1:
          try: new_src = tuple([self.replace[x] for x in new_n.src])
          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.pm.rewrite(new_n, self.ctx)) is None:
+            if self.pm is None or (new_src_n:=self.cached_pm_rewrite(new_n)) is None:
              self.replace[n] = new_n
              continue
          else:
@@ -920,7 +952,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 RuntimeError("infinite loop in graph_rewrite (explicit)")  # pylint: disable=raise-missing-from
+          except KeyError: raise RewriteNotReady  # pylint: disable=raise-missing-from
      except RewriteNotReady:
        # retry this later
        stack.insert(0, (n, stage, new_n))
@@ -936,7 +968,7 @@ def graph_rewrite_map(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, na
                      input_map:dict[UOp, UOp]|None=None, ) -> dict[UOp, UOp]:
  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx)
  new_map: dict[UOp, UOp] = {}
-  for k in sink.toposort():
+  for k in (list(sink.toposort())[::-1] if bottom_up else sink.toposort()):
    new_map[k] = v = rewrite_ctx.unified_rewrite(k)
    if k is not v and k.metadata is not None: all_metadata[v] = tuple(dedup(all_metadata.get(v, ())))+k.metadata
  if input_map is not None:
@@ -1,44 +0,0 @@
-from typing import Callable
-import functools
-from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import Ops, UPat, PatternMatcher
-from tinygrad.helpers import getenv
-from tinygrad.uop.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES, fast_idiv
-
-# ***** optional patterns *****
-
-powers_of_two = {2**i:i for i in range(64)}
-@functools.cache
-def get_late_rewrite_patterns(ops, force_transcendental=False):
-  pat: list[tuple[UPat, Callable]] = [(UPat(op, dtype=TRANSCENDENTAL_SUPPORTED_DTYPES, src=(UPat.var("d"),)), f) for op,f in \
-           ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)) if op not in ops or force_transcendental]
-  # rewrite SQRT to xpow 0.5
-  if Ops.SQRT not in ops: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5))))
-  # rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1)
-  if Ops.AND in ops: pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)]
-  # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
-  if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
-  if Ops.SHR in ops:
-    # no reason to check x<0 for uints
-    pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
-    pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(
-      c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]  # (x+(x<0).where(c-1, 0)) >> v
-    if not getenv("DISABLE_FAST_IDIV"):
-      pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d"), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))]
-      pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("d"), lambda ctx, x, d: x - d*f if (f:=fast_idiv(ctx, x, d.arg)) is not None else None)]
-  if Ops.NEG in ops:
-    pat += [(UPat.var('x')*-1, lambda x: x.alu(Ops.NEG))]
-    if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda x,y: x.alu(Ops.SUB, y))]
-  if Ops.CMPLT in ops:
-    # These are late rewrites because simplex expects equalities to be a certain format
-    pat += [
-      ((UPat.var("x", dtypes.sints) < UPat.cvar("c", dtypes.sints)).logical_not(), lambda x,c: c-1<x),
-      ((UPat.cvar("c", dtypes.sints) < UPat.var("x", dtypes.sints)).logical_not(), lambda x,c: x<c+1),
-      (UPat.var("x", dtypes.sints)*-1 < UPat.var("y", dtypes.sints)*UPat.cvar("c"), lambda x,y,c: y*(-c)<x),
-      (UPat.var("x", dtypes.sints)*-1 < UPat.cvar("c"), lambda x,c:-c<x),
-      ((UPat.cvar("c1",vec=False)<UPat.var("x", dtypes.sints)) & (UPat.var("x", dtypes.sints)<UPat.cvar("c2",vec=False)),
-        lambda x,c1,c2: x.eq(c1+1) if c1.arg+1==c2.arg-1 else None),  # (c-1)<x & x<(c+1) -> x==c
-    ]
-  if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))]
-  if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
-  return PatternMatcher(pat)
@@ -225,4 +225,4 @@ def type_verify(uops:list[UOp], extra_spec:PatternMatcher|None=None):
    with Context(TRACK_MATCH_STATS=0): ret = check_spec.rewrite(u)
    if cast(bool|None, ret) is not True:
      if DEBUG >= 3: print_uops(uops)
-      raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[x.op for x in u.src]} {u.arg}")
+      raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[(x.op, x.dtype, x.arg) for x in u.src]} {u.arg}")
@@ -1,11 +1,11 @@
 # all of symbolic lives here now
-from typing import Any, Literal, cast
+from typing import Any, cast
 import math, operator, struct, functools
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
 from tinygrad.dtype import ConstType, dtypes, PtrDType, AddrSpace, can_safe_cast
 from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, cdiv, cmod, CORRECT_DIVMOD_FOLDING
-from tinygrad.uop.transcendental import xpow
+from tinygrad.uop.decompositions import xpow

 # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ********

@@ -41,6 +41,7 @@ symbolic_simple = PatternMatcher([
  (UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x),
  (UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x),
  (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, True), UPat.const(dtypes.bool, False)), lambda x: x),
+  (UPat.var("x", dtype=dtypes.ints+(dtypes.bool,)).trunc(), lambda x: x),
  # ** zero folding **
  (UPat.var("x") < UPat.var("x"), lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x < x -> False
  (UPat.var("x") % UPat.var("x"), lambda x: x.const_like(0)), # x%x -> 0
@@ -71,6 +72,19 @@ symbolic_simple = PatternMatcher([
  (UPat.var("x").alu(Ops.POW, UPat.cvar("c", vec=False)), simplify_pow),
  # positive const ** x
  (UPat.cvar("c", vec=False).alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.arg == 1 else (x*math.log2(c.arg)).exp2() if c.arg > 0 else None),
+  # rules for threefry
+  ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)&0xFFFFFFFF), # TODO: why is the and needed?
+  (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y',  dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
+  (((UPat.var('x',  dtypes.uint64)*(1<<32)) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x),
+  # hacks for threefry long removal when padded (TODO: genericize)
+  (UPat.var('x', dtypes.uint32).cast(dtypes.uint64) * UPat.var('y').where(UPat.const(dtypes.uint64, 1<<32), UPat.const(dtypes.uint64, 0)),
+   lambda x,y: y.where(x, 0).cast(dtypes.uint64) * (1<<32)),
+  ((UPat.var('x', dtypes.uint64)&(UPat.var('y').where(UPat.const(dtypes.uint64, 0xFFFFFFFF), UPat.const(dtypes.uint64, 0)))).cast(dtypes.uint32),
+   lambda x,y: y.where(x.cast(dtypes.uint32), 0)),
+  # new decomp rules for threefry
+  (((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y',  dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
+  (((UPat.var('x',  dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32, lambda x: x),
+  (UPat.var('b').where(UPat.var('x', dtypes.uint32).cast(dtypes.uint64), UPat.const(dtypes.uint64, 0)).cast(dtypes.uint32), lambda b,x: b.where(x,0))
 ])

 # ******** phase 2 builds on phase 1, it includes the old "symbolic", rules that match deeper ********
@@ -125,65 +139,92 @@ def canonicalize_simplex(X:UOp) -> UOp|None:
    ret.append(u)
  return functools.reduce(operator.add, ret) if changed else None

-def div_and_mod_folding(x: UOp, y: UOp, which: Literal[Ops.MOD, Ops.IDIV], split_rem: bool=False) -> UOp|None:
-  # simplify x // y or x % y, None means no change
-  # simple cancel div/mod case
+def cancel_divmod(d: UOp, x: UOp, y: UOp) -> UOp|None:
+  # simple cancel div/mod case when the range of the numerator lies within a single denominator interval
  x_min, x_max, y_min, y_max = x.vmin, x.vmax, y.vmin, y.vmax
  assert isinstance(x_min, int) and isinstance(x_max, int) and isinstance(y_min, int) and isinstance(y_max, int)
-
+  if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.IDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
  if y_min*y_max > 0 and (q:=cdiv(x_min,y_min)) == cdiv(x_min,y_max) == cdiv(x_max,y_min) == cdiv(x_max,y_max):
-    return x - q*y if which is Ops.MOD else x.const_like(q)
+    return x - q*y if d.op is Ops.MOD else d.const_like(q)
+  return None

-  if (y.op is not Ops.CONST) or ((c := y.arg) < 0) or (x.dtype.count > 1): return None
-  if y.arg == 0: raise ZeroDivisionError(f"{'Division' if which is Ops.IDIV else 'Mod'} by zero trying to rewrite {x.alu(which, y)}")
-
-  svars, factors, quotients, remainders, gcd, div, const, something_changed = [], [], [], [], c, 1, 0, False
+def remove_nested_mod(m: UOp, x: UOp, y: UOp) -> UOp|None:
+  # remove nested mod in case the inner mod is a multiple of the outer mod
+  # example: (a%4 + b)%2 -> (a+b)%2
+  if ((c := y.arg) < 0) or x.vmin<0: return None
+  new_xs = []
+  something_changed = False
  for u in split_uop(x, Ops.ADD):
-    if u.op is Ops.MOD and which is Ops.MOD and u.src[1].op is Ops.CONST and u.src[1].arg%c == 0:
-      u = u.src[0]
-      something_changed = True
-    v: UOp = u.divides(f:=u.const_factor())
-    q, r = divmod(f, c)
-    if r==0 or ((which is Ops.MOD or split_rem or u.op is Ops.CONST) and r!=f): something_changed = True
-    if u.op is Ops.CONST: const += f
-    else:  # div is the smallest common divisor of all terms
-      if f > 1 and c % f == 0 and (div == 1 or div > f): div = f
-      gcd = math.gcd(r, gcd)
-      factors.append(f); svars.append(v); quotients.append(q); remainders.append(r)  # noqa: E702
+    if u.op is Ops.MOD:
+      if u.src[1].divides(c) is not None:
+        something_changed = True
+        u = u.src[0]
+    new_xs.append(u)
+  new_x: UOp = functools.reduce(operator.add, new_xs)
+  if something_changed and new_x.vmin>=0: return new_x % y
+  return None

+def fold_binary_numerator(d: UOp, x: UOp, y: UOp) -> UOp|None:
  # we can fold if the expression has only one non-constant term and this term can only take on two values
-  if len(svars)==1 and (v:=svars[0]).vmax-v.vmin == 1:
-    y1 = cmod(factors[0]*v.vmin+const, c) if which is Ops.MOD else cdiv(factors[0]*v.vmin+const, c)
-    y2 = cmod(factors[0]*v.vmax+const, c) if which is Ops.MOD else cdiv(factors[0]*v.vmax+const, c)
+  if ((c := y.arg) < 0) or (x.dtype.count > 1): return None
+  x,const = x.pop_const()
+  terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in split_uop(x, Ops.ADD)])
+  if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1:
+    y1 = cmod(factors[0]*v.vmin+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmin+const, c)  # type: ignore
+    y2 = cmod(factors[0]*v.vmax+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmax+const, c)  # type: ignore
    return (y2-y1)*(v-v.vmin) + y1
+  return None

-  if not CORRECT_DIVMOD_FOLDING or x_min>=0:
-    # a//c = (a-a%c)/c, if we can fold a%c, we can fold a//c
-    # within a mod we can freely subtract multiples of c, we use this to see if a is congruent to an expression whose vmin/vmax are between 0 and c
-    rems = [min(r, r-c, key=abs) for r in remainders]
-    if (rem:=sum(r*v for r,v in zip(rems,svars))+const%c).vmin//c==rem.vmax//c and all(f > 0 for f in factors):
-      if which is Ops.MOD: return rem - rem.vmin//c*c
-      return sum((f-r)//c * v for f,r,v in zip(factors,rems,svars)) + (const-const%c+rem.vmin//c*c)//c
+def fold_divmod_congruence(d: UOp, x: UOp, y: UOp) -> UOp|None:
+  # within a mod we can freely subtract multiples of c, we use this to see if a is congruent to an expression whose vmin/vmax are between 0 and c
+  if (x.vmin<0 and CORRECT_DIVMOD_FOLDING) or ((c := y.arg) < 0) or (x.dtype.count > 1): return None
+  x,const = x.pop_const()
+  terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in split_uop(x, Ops.ADD)])
+  # a//c = (a-a%c)/c, if we can fold a%c, we can fold a//c
+  rems = [min((r:=f%c), r-c, key=abs) for f in factors]
+  if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c and all(f > 0 for f in factors):
+    if d.op is Ops.MOD: return rem - rem.vmin//c*c
+    return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + (const-const%c+rem.vmin//c*c)//c
+  return None

-  if (g:=math.gcd(gcd, const))!=1:
-    ret = UOp(which, x.dtype, src=(sum(f//g * v for f,v in zip(factors, svars)) + const//g, x.const_like(c//g)))
-    return ret*g if which is Ops.MOD else ret
+def divide_by_gcd(d: UOp, x: UOp, y: UOp) -> UOp|None:
+  # x//y -> (x//gcd)//(y//gcd) or x%y -> gcd*(x//gcd)%(y//gcd)
+  terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in split_uop(x, Ops.ADD)])
+  if (gcd := math.gcd(y.arg, *factors)) == 1: return None
+  ret = sum(f//gcd * v for f,v in zip(factors, terms)).alu(d.op, y.const_like(y.arg//gcd))
+  return ret*gcd if d.op is Ops.MOD else ret
+
+def nest_div_by_smallest_factor(d: UOp, x: UOp, y: UOp) -> UOp|None:
+  # we try and nest the div and see if it allows the numerator to be simplified
+  if ((c := y.arg) < 0) or (x.dtype.count > 1): return None
+  factors = [u.const_factor() for u in split_uop(x.pop_const()[0], Ops.ADD)]
+  # div is the smallest factor of the denominator (greater than 1) out of all "factors"
+  # TODO: there are better ways to pick `div`, this sometimes adds extra divisions
+  # TODO: add same optimization for mod
+  div = min([y.arg]+[abs(f) for f in factors if abs(f) > 1 and (c%f)==0])
+  if (1 < div < c) and (newxs:=(newx:=(x//div)).simplify()) is not newx and x.vmin>=0 and newx.vmin>=0: return newxs//(c//div)
+  return None
+
+def simplify_remainder(d: UOp, x: UOp, y: UOp) -> UOp|None:
+  # we try and take out the quotient and see if it allows the numerator to be simplified
+  if ((c := y.arg) < 0) or (x.dtype.count > 1): return None
+  x_no_const,const = x.pop_const()
+  terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in split_uop(x_no_const, Ops.ADD)])
+  quotients, remainders = zip(*[divmod(f, c) for f in factors])
+  gcd = math.gcd(c, *remainders)  # gcd without const!
+  if const%c==const and gcd==1 and not any(r==0 or (r!=f and d.op is Ops.MOD) for r,f in zip(remainders, factors)): return None

-  if gcd != 1: something_changed = True
-  if not something_changed:
-    if which is Ops.IDIV and (1 < div < c) and (newx:=div_and_mod_folding(x, x.const_like(div), Ops.IDIV)) is not None: return newx//(c//div)
-    return None
  quo, rem = x.const_like(const//c), x.const_like((const%c)//gcd)
-  for q,r,f,v in zip(quotients, remainders, factors, svars):
-    if which is Ops.IDIV and (not split_rem) and r!=0:
+  for q,r,f,v in zip(quotients, remainders, factors, terms):
+    if d.op is Ops.IDIV and r!=0:
      rem += f//gcd * v
    else:
      rem += r//gcd * v
      quo += q * v

  # if numerator before/after is negative, and it has remainder, don't simplify because C divmod is different from python divmod.
-  if (x_min < 0 or rem.vmin < 0) and remainders: return None
-  if which is Ops.MOD: return gcd*(rem % (c//gcd)) + const%gcd
+  if (x.vmin < 0 or rem.vmin < 0) and remainders: return None
+  if d.op is Ops.MOD: return gcd*(rem % (c//gcd)) + const%gcd
  return rem//(c//gcd)+quo

 def gep_through_wmma(gep:UOp, wmma:UOp):
@@ -201,18 +242,18 @@ def gep_through_wmma(gep:UOp, wmma:UOp):

 gep_pushing = PatternMatcher([
  # GEP/VECTORIZE, GEP/GEP, GEP/CONST, GEP/VCONST
-  (UPat(Ops.GEP, src=(UPat(Ops.GEP, name='g2'),), name='g1'),
+  (UPat(Ops.GEP, name='g2').f(Ops.GEP, name='g1'),
   lambda g1, g2: g2.src[0].gep(tuple(g2.arg[g1.arg[i]] for i in range(len(g1.arg))))),
-  (UPat(Ops.GEP, src=(UPat(Ops.VECTORIZE, name="vec"),), name="gep"),
+  (UPat(Ops.VECTORIZE, name='vec').f(Ops.GEP, name='gep'),
   lambda gep, vec: UOp(Ops.VECTORIZE, gep.dtype, tuple(vec.src[i] for i in gep.arg)) if len(gep.arg) > 1 else vec.src[gep.arg[0]]),
-  (UPat(Ops.GEP, src=(UPat.cvar("c", vec=False),), name="gep"), lambda gep, c: gep.const_like(c.arg)),
-  (UPat(Ops.GEP, src=(UPat(Ops.VCONST, name="c"),), name="gep"), lambda gep, c: gep.const_like(tuple(c.arg[x] for x in gep.arg))),
+  (UPat.cvar("c", vec=False).f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(c.arg)),
+  (UPat(Ops.VCONST, name="c").f(Ops.GEP, name="gep"), lambda gep, c: gep.const_like(tuple(c.arg[x] for x in gep.arg))),
  # GEP on void is skipped
  (UPat(Ops.GEP, src=(UPat(dtype=dtypes.void, name="x"),)), lambda x: x),
  # GEP in order is removed
  (UPat(Ops.GEP, name="g"), lambda g: g.src[0] if not isinstance(g.dtype, PtrDType) and g.arg == tuple(range(g.src[0].dtype.count)) else None),
  # push all GEPs through ALUs (fix arange stuff)
-  (UPat(Ops.GEP, src=(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu'),), name='gep'),
+  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu').f(Ops.GEP, name='gep'),
   lambda gep,alu: UOp(alu.op, alu.dtype.scalar().vec(gep.dtype.count), tuple(x.gep(gep.arg) for x in alu.src), alu.arg) \
     if not isinstance(gep.dtype, PtrDType) else None),
  # CAT can't be rendered. it's a VECTORIZE on vectors, we expand to a single VECTORIZEs with GEPs (TODO: move this later)
@@ -221,7 +262,7 @@ gep_pushing = PatternMatcher([
  # VECTORIZE on same GEP
  (UPat(Ops.VECTORIZE, name="v", src=UPat(Ops.GEP, src=(UPat.var("x"),))), lambda v,x: x.gep(tuple(get_single_element(i.arg) for i in v.src))),
  # push some GEPs through WMMAs
-  (UPat(Ops.GEP, src=(UPat(Ops.WMMA, name="wmma"),), name="gep"), gep_through_wmma),
+  (UPat(Ops.WMMA, name="wmma").f(Ops.GEP, name="gep"), gep_through_wmma),
 ])

 commutative = PatternMatcher([
@@ -288,15 +329,20 @@ 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)
-  (UPat.var("x", dtypes.sints) // UPat.var("y"), lambda x,y: div_and_mod_folding(x,y,Ops.IDIV)),
+  (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),
+  (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), divide_by_gcd),
+  (UPat(Ops.MOD, dtypes.sints, name="m", src=(UPat.var("x"), UPat.cvar("y", vec=False))), remove_nested_mod),
+  (UPat((Ops.IDIV), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), nest_div_by_smallest_factor),
+  (UPat((Ops.IDIV, Ops.MOD), dtypes.sints, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), simplify_remainder),
  (UPat.var("x") // UPat.var("d"), lambda x,d: -(x//(-d)) if d.vmax < 0 else None),
  (UPat.var("x") // UPat.var("d"), lambda x,d: -((-x)//d) if x.vmax <=0 else None),
  ((UPat.var("x", dtypes.sints)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
    lambda x,c,n,d: (-(-(c.arg%d.arg + x - (d.arg-1))//d) + c.arg//d.arg) if x.vmax<=0 and n.vmin>=0 and d.arg>0 else None),
  # ** mod **
  # mod folding
-  (UPat.var("x") % UPat.var("y"), lambda x,y: div_and_mod_folding(x,y,Ops.MOD)),
-  (UPat.var("x") % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <=0 else None),
+  (UPat.var("x") % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None),
  (UPat.var("x") % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax <  0 else None),
 ])+gep_pushing

@@ -378,22 +424,6 @@ def simplify_valid(valid:UOp) -> UOp|None:
    if ret[-1] is not stmt: something_changed = True
  return functools.reduce(operator.and_, ret) if something_changed else None

-# ***** threefry *****
-
-def threefry2x32(x: UOp, key: UOp):
-  # split x and key from uint64 to two uint32
-  x0, x1 = (x & 0xffffffff).cast(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast(dtypes.uint32)
-  key0, key1 = (key & 0xffffffff).cast(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast(dtypes.uint32)
-
-  rotations = [[13, 15, 26, 6], [17, 29, 16, 24]]
-  ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0]
-  xr = [x0 + ks[-1], x1 + ks[0]]
-  for i in range(5):
-    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r)))
-    xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)]
-
-  return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64)
-
 # ******** phase 3 is the complete symbolic, and deals with very complex things like loop rewriting and threefry transform ********

 def reduce_mul_chain(r:UOp):
@@ -428,16 +458,6 @@ sym = symbolic_flat+PatternMatcher([
  # tensor core with a 0 input is acc
  (UPat(Ops.WMMA, src=(UPat.const(None, 0.0), UPat.var(), UPat.var("acc"))), lambda acc: acc),
  (UPat(Ops.WMMA, src=(UPat.var(), UPat.const(None, 0.0), UPat.var("acc"))), lambda acc: acc),
-  # threefry + remove longs
-  (UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32),
-  ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)),  # cast does truncation
-  (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y',  dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
-  (((UPat.var('x',  dtypes.uint64)*(1<<32)) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x),
-  # hacks for threefry long removal when padded (TODO: genericize)
-  (UPat.var('x', dtypes.uint32).cast(dtypes.uint64) * UPat.var('y').where(UPat.const(dtypes.uint64, 1<<32), UPat.const(dtypes.uint64, 0)),
-   lambda x,y: y.where(x, UOp.const(dtypes.uint32, 0)).cast(dtypes.uint64) * (1<<32)),
-  ((UPat.var('x', dtypes.uint64)&(UPat.var('y').where(UPat.const(dtypes.uint64, 0xFFFFFFFF), UPat.const(dtypes.uint64, 0)))).cast(dtypes.uint32),
-   lambda x,y: y.where(x.cast(dtypes.uint32), UOp.const(dtypes.uint32, 0))),
  # ** self folding **
  # x!=0 -> (bool)x
  (UPat.var("x")!=0, lambda x: x.cast(dtypes.bool.vec(x.dtype.count))),
@@ -73,7 +73,6 @@
    user-select: auto;
  }
  g.tag circle {
-    r: 5;
    fill: #FFD700;
    stroke: #B8860B;
    stroke-width: 0.8;
@@ -81,10 +80,9 @@
  g.tag text {
    text-anchor: middle;
    font-size: 6px;
-    fill: black;
+    fill: #08090e;
  }
  .label :is(text, p) {
-    color: #08090e;
    font-weight: 350;
  }
  .edgePath {
@@ -230,7 +228,7 @@
  }
  #device-list > div {
    min-height: 32px;
-    max-width: 100px;
+    max-width: 132px;
    overflow-x: auto;
    overflow-y: hidden;
    white-space: nowrap;
@@ -32,6 +32,11 @@ function intersectRect(r1, r2) {
  return {x:r1.x+dx*scale, y:r1.y+dy*scale};
 }

+function addTags(root) {
+  root.selectAll("circle").data(d => [d]).join("circle").attr("r", 5);
+  root.selectAll("text").data(d => [d]).join("text").text(d => d).attr("dy", "0.35em");
+}
+
 let [workerUrl, worker] = [null, null];
 async function renderDag(graph, additions, recenter=false) {
  // start calculating the new layout (non-blocking)
@@ -71,10 +76,8 @@ async function renderDag(graph, additions, recenter=false) {
      return [ret];
    }).join("text").selectAll("tspan").data(d => d).join("tspan").attr("x", "0").attr("dy", 14).selectAll("tspan").data(d => d).join("tspan")
      .attr("fill", d => d.color).text(d => d.st).attr("xml:space", "preserve");
-    const tags = nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
-      .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`);
-    tags.selectAll("circle").data(d => [d]).join("circle");
-    tags.selectAll("text").data(d => [d.tag]).join("text").text(d => d).attr("dy", "0.35em");
+    addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
+      .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => e.tag));
    // draw edges
    const line = d3.line().x(d => d.x).y(d => d.y).curve(d3.curveBasis);
    d3.select("#edges").selectAll("path.edgePath").data(g.edges()).join("path").attr("class", "edgePath").attr("d", (e) => {
@@ -84,7 +87,7 @@ async function renderDag(graph, additions, recenter=false) {
      points.push(intersectRect(g.node(e.w), points[points.length-1]));
      return line(points);
    }).attr("marker-end", "url(#arrowhead)");
-    const edgeLabels = d3.select("#edge-labels").selectAll("g").data(g.edges().filter(e => g.edge(e).label != null)).join("g").attr("transform", (e) => {
+    addTags(d3.select("#edge-labels").selectAll("g").data(g.edges().filter(e => g.edge(e).label != null)).join("g").attr("transform", (e) => {
      // get a point near the end
      const [p1, p2] = g.edge(e).points.slice(-2);
      const dx = p2.x-p1.x;
@@ -98,9 +101,7 @@ async function renderDag(graph, additions, recenter=false) {
      const x = p2.x - ux * offset;
      const y = p2.y - uy * offset;
      return `translate(${x}, ${y})`
-    }).attr("class", "tag");
-    edgeLabels.selectAll("circle").data(e => [g.edge(e).label]).join("circle");
-    edgeLabels.selectAll("text").data(e => [g.edge(e).label]).join("text").text(d => d).attr("dy", "0.35em");
+    }).attr("class", "tag").datum(e => g.edge(e).label));
    if (recenter) document.getElementById("zoom-to-fit-btn").click();
  };

@@ -121,6 +122,33 @@ const colorScheme = {TINY:["#1b5745", "#354f52", "#354f52", "#1d2e62", "#63b0cd"
  CATEGORICAL:["#ff8080", "#F4A261", "#C8F9D4", "#8D99AE", "#F4A261", "#ffffa2", "#ffffc0", "#87CEEB"],}
 const cycleColors = (lst, i) => lst[i%lst.length];

+const createPolygons = (source, height) => {
+  const shapes = [];
+  const yscale = d3.scaleLinear().domain([0, source.peak]).range([height, 0]);
+  for (const [i,e] of source.shapes.entries()) {
+    const x = e.x.map((i,_) => source.timestamps[i]-data.st);
+    const y0 = e.y.map(yscale);
+    const y1 = e.y.map(y => yscale(y+e.arg.nbytes));
+    const arg = { tooltipText:`${e.arg.dtype} len:${formatUnit(e.arg.sz)}\n${formatUnit(e.arg.nbytes, "B")}` };
+    shapes.push({ x, y0, y1, arg, fillColor:cycleColors(colorScheme.BUFFER, i) });
+  }
+  return shapes;
+}
+
+const rescaleTrack = (source, tid, k) => {
+  for (const e of source.shapes) {
+    for (let i=0; i<e.y0.length; i++) {
+      e.y0[i] = e.y0[i]*k;
+      e.y1[i] = e.y1[i]*k;
+    }
+  }
+  const change = (source.height*k)-source.height;
+  const div = document.getElementById(tid);
+  div.style.height = rect(div).height+change+"px";
+  source.height = source.height*k;
+  return change;
+}
+
 const drawLine = (ctx, x, y) => {
  ctx.beginPath();
  ctx.moveTo(x[0], y[0]);
@@ -129,78 +157,68 @@ const drawLine = (ctx, x, y) => {
  ctx.stroke();
 }

-var profileRet, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity;
+var data, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity;
 async function renderProfiler() {
  displayGraph("profiler");
  d3.select(".metadata").html("");
+  // layout once!
+  if (data != null) return;
  const profiler = d3.select(".profiler").html("");
-  const deviceList = profiler.append("div").attr("id", "device-list").node();
+  const { layout, st, et } = await (await fetch("/get_profile")).json();
+  // place devices on the y axis and set vertical positions
+  const [tickSize, padding] = [10, 8];
+  const deviceList = profiler.append("div").attr("id", "device-list").style("padding-top", tickSize+padding+"px");
  const canvas = profiler.append("canvas").attr("id", "timeline").node();
  // NOTE: scrolling via mouse can only zoom the graph
  canvas.addEventListener("wheel", e => (e.stopPropagation(), e.preventDefault()), { passive:false });
-  if (profileRet == null) profileRet = await (await fetch("/get_profile")).json()
-  const { layout, st, et } = profileRet;
-  // place devices on the y axis and set vertical positions
-  const [tickSize, padding] = [10, 8];
-  deviceList.style.paddingTop = `${tickSize+padding}px`;
  const ctx = canvas.getContext("2d");
  const canvasTop = rect(canvas).top;
  // color by key (name/category/device)
  const colorMap = new Map();
-  const data = {shapes:[], axes:{}};
-  const areaScale = d3.scaleLinear().domain([0, Object.entries(layout).reduce((peak, [_,d]) => Math.max(peak, d.mem.peak), 0)]).range([4,maxArea=100]);
-  for (const [k, { timeline, mem }] of Object.entries(layout)) {
-    if (timeline.shapes.length === 0 && mem.shapes.length == 0) continue;
-    const div = deviceList.appendChild(document.createElement("div"));
-    div.innerText = k;
-    div.style.padding = `${padding}px`;
-    div.onclick = () => { // TODO: make this feature more visible
-      focusedDevice = k === focusedDevice ? null : k;
-      const prevScroll = profiler.node().scrollTop;
-      renderProfiler();
-      if (prevScroll) profiler.node().scrollTop = prevScroll;
-    }
-    const { y:baseY, height:baseHeight } = rect(div);
-    const levelHeight = baseHeight-padding;
+  data = {tracks:new Map(), axes:{}, st, et};
+  const heightScale = d3.scaleLinear().domain([0, Object.entries(layout).reduce((peak, [_,d]) => Math.max(peak, d.peak||0), 0)]).range([4,maxheight=100]);
+  for (const [k, v] of Object.entries(layout)) {
+    if (v.shapes.length === 0) continue;
+    const div = deviceList.append("div").attr("id", k).text(k).style("padding", padding+"px");
+    const { y:baseY, height:baseHeight } = rect(div.node());
    const offsetY = baseY-canvasTop+padding/2;
-    let colorKey, ref;
-    for (const e of timeline.shapes) {
-      if (e.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(e.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};
-      else if (ref != null) {
-        const start = ref.step>0 ? ref.step+1 : 0;
-        const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
-        ref = stepIdx === -1 ? null : {ctx:ref.ctx, step:stepIdx};
+    if (v.shapes[0].dur != null) {
+      const levelHeight = baseHeight-padding;
+      const shapes = [];
+      data.tracks.set(k, { shapes, offsetY });
+      let colorKey, ref;
+      for (const e of v.shapes) {
+        if (e.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(e.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};
+        else if (ref != null) {
+          const start = ref.step>0 ? ref.step+1 : 0;
+          const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
+          ref = stepIdx === -1 ? null : {ctx:ref.ctx, step:stepIdx};
+        }
+        const arg = { tooltipText:formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...ref };
+        // offset y by depth
+        shapes.push({x:e.st-st, y:levelHeight*e.depth, width:e.dur, height:levelHeight, arg, label, fillColor });
      }
-      const arg = { tooltipText:formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...ref };
-      // offset y by depth
-      data.shapes.push({x:e.st-st, y:offsetY+levelHeight*e.depth, width:e.dur, height:levelHeight, arg, label, fillColor });
+      div.style("height", levelHeight*v.maxDepth+padding+"px").style("pointerEvents", "none");
+    } else {
+      const height = heightScale(v.peak);
+      data.tracks.set(k, { shapes:createPolygons(v, height), offsetY, height, peak:v.peak, scaleFactor:maxheight*4/height });
+      div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => {
+        const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id;
+        let offset = 0;
+        for (const [tid, track] of data.tracks) {
+          track.offsetY += offset;
+          if (tid === newFocus) offset += rescaleTrack(track, tid, track.scaleFactor);
+          else if (tid === focusedDevice) offset += rescaleTrack(track, tid, 1/track.scaleFactor);
+        }
+        data.axes.y = newFocus != null ? { domain:[0, (t=data.tracks.get(newFocus)).peak], range:[t.offsetY+t.height, t.offsetY], fmt:"B" } : null;
+        focusedDevice = newFocus;
+        return resize();
+      });
    }
-    // position shapes on the canvas and scale to fit fixed area
-    let area = mem.shapes.length === 0 ? 0 : areaScale(mem.peak);
-    if (area === 0) div.style.pointerEvents = "none";
-    else {
-      const startY = offsetY+(levelHeight*timeline.maxDepth)+padding/2;
-      div.style.cursor = "pointer";
-      if (k === focusedDevice) {
-        // expand memory graph for the focused device
-        area = maxArea*4;
-        data.axes.y = { domain:[0, mem.peak], range:[startY+area, startY], fmt:"B" };
-      }
-      const yscale = d3.scaleLinear().domain([0, mem.peak]).range([startY+area, startY]);
-      for (const [i,e] of mem.shapes.entries()) {
-        const x = e.x.map((i,_) => (mem.timestamps[i] ?? et)-st);
-        const y0 = e.y.map(yscale);
-        const y1 = e.y.map(y => yscale(y+e.arg.nbytes));
-        const arg = { tooltipText:`${e.arg.dtype} len:${formatUnit(e.arg.sz)}\n${formatUnit(e.arg.nbytes, "B")}` };
-        data.shapes.push({ x, y0, y1, arg, fillColor:cycleColors(colorScheme.BUFFER, i) });
-      }
-    }
-    // lastly, adjust device rect by number of levels
-    div.style.height = `${Math.max(levelHeight*timeline.maxDepth, baseHeight)+area+padding}px`;
  }
  updateProgress({ "show":false });
  // draw events on a timeline
@@ -221,63 +239,63 @@ async function renderProfiler() {
      yscale = d3.scaleLinear().domain(data.axes.y.domain).range(data.axes.y.range);
    }
    // draw shapes
-    for (const e of data.shapes) {
-      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) {
-        const x = e.x.map(xscale);
-        ctx.beginPath();
-        ctx.moveTo(x[0], e.y0[0]);
-        for (let i=1; i<x.length; i++) ctx.lineTo(x[i], e.y0[i]);
-        for (let i=x.length-1; i>=0; i--) ctx.lineTo(x[i], e.y1[i]);
-        ctx.closePath();
-        ctx.fill();
-        // NOTE: y coordinates are in reverse order
-        for (let i = 0; i < x.length - 1; i++) {
-          let tooltipText = e.arg.tooltipText;
-          if (yscale != null && ((yaxisVal=yscale.invert(e.y1[i]))>0)) {
-            tooltipText += `\nTotal: ${formatUnit(yaxisVal, data.axes.y.fmt)}`;
+    for (const [_, { offsetY, shapes }] of data.tracks) {
+      for (const e of shapes) {
+        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) {
+          const x = e.x.map(xscale);
+          ctx.beginPath();
+          ctx.moveTo(x[0], offsetY+e.y0[0]);
+          for (let i=1; i<x.length; i++) ctx.lineTo(x[i], offsetY+e.y0[i]);
+          for (let i=x.length-1; i>=0; i--) ctx.lineTo(x[i], offsetY+e.y1[i]);
+          ctx.closePath();
+          ctx.fill();
+          // NOTE: y coordinates are in reverse order
+          for (let i = 0; i < x.length - 1; i++) {
+            let tooltipText = e.arg.tooltipText;
+            if (yscale != null && ((yaxisVal=yscale.invert(offsetY+e.y1[i]))>0)) {
+              tooltipText += `\nTotal: ${formatUnit(yaxisVal, data.axes.y.fmt)}`;
+            }
+            rectLst.push({ x0:x[i], x1:x[i+1], y0:offsetY+e.y1[i], y1:offsetY+e.y0[i], arg:{...e.arg, tooltipText} });
          }
-          rectLst.push({ x0:x[i], x1:x[i+1], y0:e.y1[i], y1:e.y0[i], arg:{...e.arg, tooltipText} });
+          continue;
        }
-        continue;
-      }
-      // contiguous rect
-      const x = xscale(start);
-      const width = xscale(end)-x;
-      ctx.fillRect(x, e.y, width, e.height);
-      rectLst.push({ y0:e.y, y1:e.y+e.height, x0:x, x1:x+width, arg:e.arg });
-      // add label
-      if (e.label == null) continue;
-      ctx.textAlign = "left";
-      ctx.textBaseline = "middle";
-      let [labelX, labelWidth] = [x+2, 0];
-      const labelY = e.y+e.height/2;
-      for (const [i,l] of e.label.entries()) {
-        if (labelWidth+l.width+(i===e.label.length-1 ? 0 : ellipsisWidth)+2 > width) {
-          if (labelWidth !== 0) ctx.fillText("...", labelX, labelY);
-          break;
+        // contiguous rect
+        const x = xscale(start);
+        const width = xscale(end)-x;
+        ctx.fillRect(x, offsetY+e.y, width, e.height);
+        rectLst.push({ y0:offsetY+e.y, y1:offsetY+e.y+e.height, x0:x, x1:x+width, arg:e.arg });
+        // add label
+        if (e.label == null) continue;
+        ctx.textAlign = "left";
+        ctx.textBaseline = "middle";
+        let [labelX, labelWidth] = [x+2, 0];
+        const labelY = offsetY+e.y+e.height/2;
+        for (const [i,l] of e.label.entries()) {
+          if (labelWidth+l.width+(i===e.label.length-1 ? 0 : ellipsisWidth)+2 > width) {
+            if (labelWidth !== 0) ctx.fillText("...", labelX, labelY);
+            break;
+          }
+          ctx.fillStyle = l.color;
+          ctx.fillText(l.st, labelX, labelY);
+          labelWidth += l.width;
+          labelX += l.width;
        }
-        ctx.fillStyle = l.color;
-        ctx.fillText(l.st, labelX, labelY);
-        labelWidth += l.width;
-        labelX += l.width;
      }
    }
    // draw axes
    drawLine(ctx, xscale.range(), [0, 0]);
-    const ticks = xscale.ticks();
-    for (const [i, tick] of ticks.entries()) {
+    for (const tick of xscale.ticks()) {
      // tick line
      const x = xscale(tick);
      drawLine(ctx, [x, x], [0, tickSize])
      // tick label
      ctx.textBaseline = "top";
-      ctx.textAlign = i === ticks.length-1 ? "right" : "left";
-      const padding = i === ticks.length-1 ? -1 : 1;
-      ctx.fillText(formatTime(tick, et-st), x+(ctx.lineWidth+2)*padding, tickSize);
+      ctx.textAlign = "left";
+      ctx.fillText(formatTime(tick, et-st), x+ctx.lineWidth+2, tickSize);
    }
    if (yscale != null) {
      drawLine(ctx, [0, 0], yscale.range());
@@ -19,7 +19,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
               **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF",
               Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500",
               Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
-               Ops.CHILD: "#80fff0"}
+               Ops.CHILDREN: "#80ffc0", Ops.CHILD: "#80fff0", Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e"}

 # VIZ API

@@ -73,13 +73,15 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
    if u.dtype != dtypes.void: label += f"\n{u.dtype}"
    for idx,x in enumerate(u.src):
      if x in excluded:
-        if x.op is Ops.CONST and dtypes.is_float(u.dtype): label += f"\nCONST{idx} {x.arg:g}"
-        else: label += f"\n{x.op.name}{idx} {x.arg}"
+        arg = f"{x.arg:g}" if x.op is Ops.CONST and dtypes.is_float(u.dtype) else f"{x.arg}"
+        label += f"\n{x.op.name}{idx} {arg}" + (f" {x.src[0].op}" if len(x.src) else "")
    try:
      if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None:
        label += f"\n{shape_to_str(u.shape)}"
+      elif len(rngs:=u.ranges):
+        label += f"\n{str(sorted([x.arg for x in rngs]))}"
    except Exception:
-      label += "\n<ISSUE GETTING SHAPE>"
+      label += "\n<ISSUE GETTING LABEL>"
    if (ref:=ref_map.get(u.arg.ast) if u.op is Ops.KERNEL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}"
    # NOTE: kernel already has metadata in arg
    if TRACEMETA >= 2 and u.metadata is not None and u.op is not Ops.KERNEL: label += "\n"+repr(u.metadata)
@@ -142,7 +144,7 @@ def timeline_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
    shapes.append({"name":name, "ref":ref, "st":st, "dur":dur, "depth":depth, "cat":cat, "info":info})
  return {"shapes":shapes, "maxDepth":len(levels)}

-def mem_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
+def mem_layout(events:list[tuple[int, int, float, DevEvent]], max_ts:int) -> dict:
  step, peak, mem = 0, 0, 0
  shps:dict[int, dict] = {}
  temp:dict[int, dict] = {}
@@ -168,9 +170,10 @@ def mem_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
  for v in temp.values():
    v["x"].append(step)
    v["y"].append(v["y"][-1])
+  timestamps.append(max_ts)
  return {"shapes":list(shps.values()), "peak":peak, "timestamps":timestamps}

-def get_profile(profile:list[ProfileEvent]):
+def get_profile(profile:list[ProfileEvent]) -> bytes|None:
  # start by getting the time diffs
  for ev in profile:
    if isinstance(ev,ProfileDeviceEvent): device_ts_diffs[ev.device] = (ev.comp_tdiff, ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff)
@@ -182,10 +185,14 @@ def get_profile(profile:list[ProfileEvent]):
    dev_events.setdefault(e.device,[]).append((st:=int(ts), et:=int(en), float(en-ts), e))
    if min_ts is None or st < min_ts: min_ts = st
    if max_ts is None or et > max_ts: max_ts = et
+  if min_ts is None: return None
  # return layout of per device events
-  for events in dev_events.values(): events.sort(key=lambda v:v[0])
-  dev_layout = {k:{"timeline":timeline_layout(v), "mem":mem_layout(v)} for k,v in dev_events.items()}
-  return json.dumps({"layout":dev_layout, "st":min_ts, "et":max_ts}).encode("utf-8")
+  layout:dict[str, dict] = {}
+  for k,v in dev_events.items():
+    v.sort(key=lambda e:e[0])
+    layout[k] = timeline_layout(v)
+    layout[f"{k} Memory"] = mem_layout(v, unwrap(max_ts))
+  return json.dumps({"layout":layout, "st":min_ts, "et":max_ts}).encode("utf-8")

 def get_runtime_stats(key) -> list[dict]:
  ret:list[dict] = []
Author	SHA1	Message	Date
George HotzandGitHub	2036361d4d	Revert "Revert "fix the misused cast in amd llvm tc (#11711 )" (#11715 )" This reverts commit `ca28db5a97`.	2025-08-18 19:44:00 -07:00
George Hotz GitHub chenyu	2ea54d7337	improve syntax of UPats using f [pr] (#11717 ) Co-authored-by: chenyu <[email protected]>	2025-08-18 20:49:45 -04:00
chenyuandGitHub	b67345caa3	use truncate in onnx read_int64 [pr] (#11720 )	2025-08-18 20:49:35 -04:00
qazalandGitHub	50e789e290	hotfix: add device to decompositions ctx (#11721 ) fast_idiv requires it for checking if a dtype is supported. Without this, codegen creates non reproducible output without a complete os.environ. since `is_dtype_supported` will open devices based on the env var unless the device is specified by the caller.	2025-08-19 03:31:16 +03:00
George HotzandGitHub	4b3fcb4064	Revert "REDUCE_AXIS keepdim=False (#11311 )" (#11718 ) This reverts commit `b518a7378a`.	2025-08-18 13:28:53 -07:00
George HotzandGitHub	67d0ba5bd8	new ops from rangeify (#11716 )	2025-08-18 13:13:11 -07:00
geohot	4afa0b86bb	hotfix: ls -lh on wheel size	2025-08-18 11:52:59 -07:00
George HotzandGitHub	ca28db5a97	Revert "fix the misused cast in amd llvm tc (#11711 )" (#11715 ) This reverts commit `799a637b03`.	2025-08-18 11:51:28 -07:00
chenyuandGitHub	c10e4c4e20	print wheel build size (#11714 )	2025-08-18 14:29:47 -04:00
b1tg GitHub b1tg George Hotz nimlgen	b518a7378a	REDUCE_AXIS keepdim=False (#11311 ) * progress * fix tests * fix tests * remove hack for test_symfold * fix test_conv.py on llvm * hack test_cache_speed * lint * remove hack for helper_linearizer_opt * tests * fix DSP * clean up * remove hack for kernelize.py * hack for test/test_multitensor.py TestMultiTensor.test_matmul_shard_none * clean * uop.r need reshape? * lower_store cause fail * fix lower? * avoid contiguous hack * 2134 * conv2d count * remove unused * hack lower * reduced and clean up * fix TestMultiTensor.test_matmul_shard_none * src sync + fix TestMultiTensor.test_matmul_shard_none * remove excluded in mop --------- Co-authored-by: b1tg <[email protected]> Co-authored-by: George Hotz <[email protected]> Co-authored-by: nimlgen <[email protected]>	2025-08-18 10:09:17 -07:00
b1tg GitHub b1tg	61884f2057	add cstyle renderer to the NULL device (#11709 ) Co-authored-by: b1tg <[email protected]>	2025-08-18 09:52:22 -07:00
uuuvn GitHub wozeparrot	18db8fa311	Allow choosing leaders in multinode reduce (#11506 ) Co-authored-by: wozeparrot <[email protected]>	2025-08-18 12:43:20 -04:00
b1tg GitHub b1tg	799a637b03	fix the misused cast in amd llvm tc (#11711 ) Co-authored-by: b1tg <[email protected]>	2025-08-18 09:15:34 -07:00
qazalandGitHub	fef97547f9	viz: preset the final timestamp (#11712 )	2025-08-18 17:51:21 +03:00
chenyuandGitHub	c30a113b2a	support bf16 and fp8 in Tensor.tolist (#11704 ) memoryview does not support it, but casting works fine so cast is fine	2025-08-17 15:11:13 -04:00
nimlgenandGitHub	1c62a3833b	am: add versioned_header to load_fw (#11702 ) * am: add versioned_header to load_fw * fix mypy	2025-08-17 20:11:57 +03:00
qazalandGitHub	eb3c918c5b	viz: s/area/height (#11703 )	2025-08-17 19:20:01 +03:00
qazalandGitHub	d762edd694	viz: define tracks in python (#11701 ) * viz: defines tracks in python * update unittests * figuring it out * works * diff cleanup * math * y axis is back	2025-08-17 18:19:13 +03:00
qazalandGitHub	eeeea29171	viz: device list refactor (#11700 ) * viz: device list refactor * paddingTop/padding-top	2025-08-17 15:08:54 +03:00
George HotzandGitHub	9366a23eb0	test backward in test_tiny (#11697 ) * test backward in test_tiny * empty	2025-08-16 20:29:39 -07:00
chenyuandGitHub	4666df71c1	fix test_fuse_and_tc_opt (#11699 )	2025-08-16 21:10:53 -04:00
geohotstanandGitHub	3d7c35d615	add fuse and tc opt bug repro (#11695 ) * FINALLY HAVE A SMALL REPRO OH BOY * show failure in CI * cleaner? * 1 possible fix * Revert "1 possible fix" This reverts commit `9e0fd215dd`.	2025-08-16 18:24:49 -04:00
nimlgenandGitHub	d1224a7c4a	am: check both signatures (#11694 ) * am: check both signatures * fix	2025-08-16 20:01:07 +03:00
qazalandGitHub	58c8991fa4	add Ops.REWRITE_ERROR (#11689 )	2025-08-16 00:56:53 +03:00
qazalandGitHub	ec4fccb1da	viz: pass through RewriteNotReady (#11690 )	2025-08-16 00:33:59 +03:00
qazalandGitHub	e954decb44	viz: pass UOp.st errors (#11688 )	2025-08-16 00:07:56 +03:00
nimlgenandGitHub	bf0c45fd16	system: resource_resize might be unavail (#11680 )	2025-08-15 22:03:23 +03:00
George HotzandGitHub	4ab9fb2edd	explicit fixed point rewrite (#11685 ) * explicit fixed point rewrite * local cache * fix that	2025-08-15 11:08:41 -07:00
chenyuandGitHub	5d6963c968	RuntimeError for unsupported dtype in PYTHON (#11686 )	2025-08-15 13:59:27 -04:00
nimlgenandGitHub	b970cd6895	am: fix psp ring completion (#11679 ) * am: psp ring timeout + fix 0 fence_value * no sleep	2025-08-15 20:15:49 +03:00
qazalandGitHub	c8ba48b223	show rewrite errors in viz (#11684 )	2025-08-15 19:09:47 +03:00
George HotzandGitHub	560984fd8d	small changes from rangeify (#11682 ) * small changes from rangeify * const like thing * ksym	2025-08-15 08:45:52 -07:00
chenyuandGitHub	d0d39885c3	onnx in tinygrad (#11675 )	2025-08-14 19:57:21 -04:00
wozeparrotandGitHub	71260a5ea4	feat: only bench openpilot 0.9.9 models (#11664 )	2025-08-14 19:27:18 -04:00
chenyuandGitHub	4ddefbccb4	update setup packages (#11674 ) sorted, and added missing 'tinygrad.frontend' and 'tinygrad.runtime.autogen.nv'	2025-08-14 19:24:57 -04:00
chenyuandGitHub	48c4033ae1	fix pylint for onnx (#11673 ) * fix pylint for onnx * too long	2025-08-14 18:48:02 -04:00
chenyuandGitHub	e9d0027591	llama MP realize weight after shard (#11672 ) * llama MP realize weight after shard prevents memory spike on device 0 * empty weight for FAKEDATA	2025-08-14 16:17:46 -04:00
nimlgenandGitHub	4176b24264	amd: support xcc in regs (#11670 ) * amd: support xcc in regs * mockamd * typong	2025-08-14 21:20:11 +03:00
Sieds LyklesandGitHub	f399d0d75d	Render mod in terms of idiv (#11668 ) * Render mod in terms of idiv * cvar -> var	2025-08-14 19:59:39 +02:00
nimlgenandGitHub	d747eeed32	amd logs parser based on device (#11669 )	2025-08-14 19:49:33 +03:00
geohotstanandGitHub	1e904155e3	Add Onnx Huggingface to test/models/test_onnx.py (#11468 ) * BOOM * cache extra/huggingface/models/ * why max buffer size is not 0 * override MAX_BUFFER_SIZE * less models * remove more models and change cache dir to already cached dir * only metal * less is more? * remove check ops * why is this not setting the ENVVAR * ughhhhh just test in models * only cpu and gpu * only cpu actually * just override it idk * final * move extra dependencies up top * simplification * fix print * make README better * revert ops_disk fix for now * clean up test_onnx * remove testing fashion clip model cuz sloooowwwwww * actually let METAL run this * fix comment mistake * fix download path in run_models * does this work? * cleanup setup and teardown * contextvar like this? * prove model is cached * do I need to increment DOWNLOAD_CACHE_VERSION? * see if cached with incremented DOWNLOAD_CACHE_VERSION * use warnings to see if the model exists * revert DOWNLOAD_CACHE_VERSION stuff and clean up * add retry to download * nit	2025-08-14 11:16:41 -04:00
Sieds LyklesandGitHub	06beeb6e13	Nest div even if factor is negative (#11666 )	2025-08-14 13:58:59 +02:00
Sieds LyklesandGitHub	661e9a2d5d	div_and_mod_folding refactor (#11585 ) * divmod const folding is its own function * split nested mod optimization out of div and mod folding * make `fold_binary_numerator` its own function * factor out `fold_divmod_congruence` * check sign of numerator * add tests * assert int on vmin and vmax * add type: ignore * factor out more rules * remove div_and_mod_folding * cached_property to property * remove import * add returns * restore old order * check sign of x.vmin and newx.vmin * check more signs * add some test that would have caught bugs * better test if the div simplified * shorten line * replace terms_factors_const with pop_const * move that back * minor cleanup * remove comments * some cleanup	2025-08-14 11:52:42 +02:00
chenyuandGitHub	0fc43c2e54	fix test_const_tensor_index index (#11660 ) index should be ints	2025-08-13 19:50:16 -04:00
chenyuandGitHub	4fe19eec72	Ops.TRUNC (#11659 )	2025-08-13 18:40:48 -04:00
qazalandGitHub	eb10a9c76a	viz: always left align timeline values (#11658 )	2025-08-13 23:55:28 +03:00
George HotzandGitHub	22bdf48cdd	render ranges in viz, name gbufs with sizes. changes from rangeify (#11656 ) * render ranges in viz, name gbufs with sizes. changes from rangeify * fix unit test dtypes	2025-08-13 12:46:16 -07:00
George HotzandGitHub	9b4da590bb	remove need for cast_vec (#11653 ) * remove need for cast_vec * fix amdllvm	2025-08-13 12:09:47 -07:00
kevvz GitHub George Hotz	e2873a3a41	[bounty] Muon optim (#11414 ) * newton schulz * add muon + move newton schulz to tensor * compact newton schulz * better tests * cleanup * add comments for muon * cleanup * add export with tests * match muon optim with test optim * cleanup * unsed import * correct comment * whitespace * move export * muon test fix * match reference impl + tests * remove export by moving muon device * add credit * cleanup * remove print * spacing * spacing * comma * cleanup * removal * fix tests + optim momentum * consistent is not/ not * more consistency * fix test * cleanup * fix the nones * remove comment * cast * comment * comment * muon teeny test * muon flag beautiful mnist * set steps * steps as hyperparam * match default test steps * name * large cleanup * dont care about steps * nesterov false default * match each other impl * steps * switch nest * swap defaults * update docstring * add no nesterov test * ban fuse_optim * prints * classical momentum * alternative condition * recon * pre + post wd * false default * detach * signature changes * context * swap order * big cleanup * 0 step instead * parity * remove fuse * remove fused * better paper * assert message * correct shape check + eps * multidim * add eps * cleanup * correct assert message * lint * better tests * naming * ns_steps,ns_params * update docstring * docstring * match sgd and muon together * sandwich * add back fused * parity --------- Co-authored-by: George Hotz <[email protected]>	2025-08-13 14:27:55 -04:00
chenyuandGitHub	94e6d84e32	rewrite Tensor.round to not use cast int (#11654 )	2025-08-13 13:51:08 -04:00
George HotzandGitHub	d2521d828a	transcendental+idiv+threefry are uop decompositions (#11636 ) * transcendental+idiv+threefry are uop decompositions [pr] * threefry decomp * fix randomness tests * fix webgpu * unneeded now * fix * move prematcher * all cast should probably be cast_vec	2025-08-13 09:37:12 -07:00
geohotstanandGitHub	cf7224ce3e	fully lint onnx.py (#11647 ) * mypy * ruff ruff ruff	2025-08-13 08:22:06 -07:00
geohotstanandGitHub	925555b62a	Fix onnx Domain bug (#11650 )	2025-08-13 08:20:50 -07:00
Sieds LyklesandGitHub	67df617fe1	add launch bounds to ptx (#11646 )	2025-08-13 13:05:39 +02:00
qazalandGitHub	88f95e9f59	viz: minor fixups for firefox (#11645 ) * fix circle attr * set fill color	2025-08-13 12:59:28 +03:00
qazalandGitHub	6f88eac0fc	viz: refactor node and edge tagging (#11644 )	2025-08-13 12:41:01 +03:00
qazalandGitHub	8140bf9778	viz: create layout once (#11643 ) * start * work * works * diff cleanup	2025-08-13 09:24:58 +03:00
chenyuandGitHub	3fb79bb43a	minor onnx cleanups (#11642 )	2025-08-13 01:05:19 -04:00
chenyuandGitHub	e9e5a08a04	simplify onnx cubic (#11641 ) we can drop the double where and abs since we know which ranges the inputs map into	2025-08-12 19:57:31 -04:00
George HotzandGitHub	18cdbec447	split decompositions pass (#11638 ) * split decompositions pass * fix ptx * pack load store early * restore that	2025-08-12 12:56:05 -07:00