restore that

.github/workflows/benchmark.yml

@@ -62,6 +62,8 @@ 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
@@ -603,12 +605,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: 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: 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: 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

.github/workflows/test.yml

@@ -329,6 +329,7 @@ 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/
@@ -336,6 +337,7 @@ jobs:
       run: |
         python -m mypy --strict-equality --lineprecision-report .
         cat lineprecision.txt
+        python -m mypy --strict-equality extra/onnx.py
 
   unittest:
     name: Unit Tests
@@ -373,8 +375,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 < 17000 lines
-      run: MAX_LINE_COUNT=17000 python sz.py
+    - name: Repo line count < 16000 lines
+      run: MAX_LINE_COUNT=16000 python sz.py
 
   fuzzing:
     name: Fuzzing

examples/beautiful_mnist.py

@@ -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 if not getenv("MUON") else nn.optim.Muon)(nn.state.get_parameters(model))
+  opt = nn.optim.Adam(nn.state.get_parameters(model))
 
   @TinyJit
   @Tensor.train()
@@ -29,7 +29,8 @@ 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()
-    return loss.realize(*opt.schedule_step())
+    opt.step()
+    return loss
 
   @TinyJit
   def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100

examples/mlperf/model_train.py

@@ -1318,10 +1318,6 @@ 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):
@@ -1343,8 +1339,6 @@ 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)

extra/amdpci/proclogs.py

@@ -1,16 +1,6 @@
-import re, ctypes, sys, importlib
+import re, ctypes, sys
 
-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
+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
 
 def parse_amdgpu_logs(log_content, register_names=None):
   register_map = register_names
@@ -33,11 +23,16 @@ 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 = {}
-  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}"
+  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)])
 
   with open(sys.argv[1], 'r') as f:
     log_content = log_content_them = f.read()

extra/huggingface_onnx/README.md

@@ -1,61 +0,0 @@
-# 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
-```

extra/huggingface_onnx/collect_metadata.py

@@ -0,0 +1,85 @@
+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)}")

extra/huggingface_onnx/download_models.py

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

extra/huggingface_onnx/huggingface_manager.py

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

extra/huggingface_onnx/run_models.py

@@ -1,11 +1,10 @@
-import onnx, yaml, tempfile, time, argparse, json
+import onnx, yaml, tempfile, time, collections, pprint, 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) -> dict[str, Any]:
+def get_config(root_path: Path):
   ret = {}
   for path in root_path.rglob("*config.json"):
     config = json.load(path.open())
@@ -13,19 +12,19 @@ def get_config(root_path: Path) -> dict[str, Any]:
       ret.update(config)
   return ret
 
-def get_tolerances(file_name: str) -> tuple[float, float]:
+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
   # 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(models)} models **")
+  print(f"** Validating {len(model_paths)} models **")
   for model_id, (root_path, relative_path) in models.items():
     print(f"validating model {model_id}")
     model_path = root_path / relative_path
@@ -37,6 +36,25 @@ 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)
@@ -53,9 +71,12 @@ 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")
-  parser.add_argument("--validate", type=str, default="",
-                      help="Validate correctness of models from the specified YAML configuration file")
+  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.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
@@ -64,13 +85,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.validate or args.debug):
-    parser.error("Please provide either --validate <yaml_file> or --debug <repo_id>.")
+  if not (args.check_ops or args.validate or args.debug):
+    parser.error("Please provide either --validate, --check_ops, or --debug.")
   if args.truncate != -1 and not args.debug:
     parser.error("--truncate and --debug should be used together for debugging")
 
-  if args.validate:
-    with open(args.validate, 'r') as f:
+  if args.check_ops or args.validate:
+    with open(args.input, '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 = {
@@ -80,16 +101,22 @@ if __name__ == "__main__":
         if model["file"].endswith(".onnx")
       }
 
-    validate_repos(model_paths)
+    if args.check_ops:
+      pprint.pprint(retrieve_op_stats(model_paths))
+
+    if args.validate:
+      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 = 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)
+      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)
       config = get_config(root_path)
       for onnx_model in root_path.rglob("*.onnx"):
         rtol, atol = get_tolerances(onnx_model.name)
@@ -101,8 +128,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 = 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)
+      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)
       config = get_config(root_path)
       rtol, atol = get_tolerances(onnx_model)
       print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")

extra/onnx_helpers.py

@@ -1,6 +1,7 @@
 from tinygrad import Tensor
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner, OnnxValue
+from tinygrad.frontend.onnx import OnnxRunner
+from extra.onnx import OnnxValue
 import numpy as np
 import onnxruntime as ort
 

extra/torch_muon.py

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

setup.py

@@ -24,28 +24,10 @@ setup(name='tinygrad',
       license='MIT',
       long_description=long_description,
       long_description_content_type='text/markdown',
-      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',
-      ],
+      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'],
       package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
       classifiers=[
         "Programming Language :: Python :: 3",

test/external/external_test_onnx_ops.py

@@ -294,6 +294,7 @@ 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]:

test/external/external_test_onnx_runner.py

@@ -3,7 +3,8 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.uop.ops import Ops
 from tinygrad.device import is_dtype_supported
-from tinygrad.frontend.onnx import OnnxRunner, OnnxDataType
+from extra.onnx import OnnxDataType
+from tinygrad.frontend.onnx import OnnxRunner
 from hypothesis import given, strategies as st
 
 # copied from test_const_folding.py

test/external/fuzz_fast_idiv.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.decompositions import fast_idiv
+from tinygrad.uop.transcendental import fast_idiv
 random.seed(42)
 
 powers_of_two = [2**i for i in range(64)]

test/mockgpu/amd/amddriver.py

@@ -87,19 +87,16 @@ 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')),

test/models/test_onnx.py

@@ -9,15 +9,7 @@ except ModuleNotFoundError:
   raise unittest.SkipTest("onnx not installed, skipping onnx test")
 from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.tensor import Tensor
-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
+from tinygrad.helpers import CI, fetch, temp
 
 def run_onnx_torch(onnx_model, inputs):
   import torch
@@ -145,36 +137,5 @@ 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()

test/test_const_folding.py

@@ -143,12 +143,13 @@ 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: these can be 0, implement const tensor folded indexing
+    # TODO: implement const tensor folded indexing
     t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
-    _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)])
+    _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)])
 
 class TestMovedConstFolding(unittest.TestCase):
   def test_add_shrunk_zero(self):

test/test_linalg.py

@@ -62,15 +62,5 @@ 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()

test/test_ops.py

@@ -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, 1e12, -1e12]], 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)
   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)

test/test_optim.py

@@ -2,10 +2,9 @@ import numpy as np
 import torch
 import unittest
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.nn.optim import Adam, SGD, AdamW, Muon
+from tinygrad.nn.optim import Adam, SGD, AdamW
 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)
@@ -58,12 +57,9 @@ 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)
@@ -87,28 +83,6 @@ 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)

test/test_rangeify.py

@@ -1,113 +0,0 @@
-import unittest
-from tinygrad import Tensor
-
-class TestRangeify(unittest.TestCase):
-  def test_add(self):
-    N = 1024
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    (A+B).realize()
-
-  def test_double_gemm(self):
-    N = 1024
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (A@B@C).realize()
-
-  def test_double_gemm_exp(self):
-    N = 1024
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).exp()@C).exp()).realize()
-
-  def test_double_gemm_relu(self):
-    N = 1024
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).relu()@C).relu()).realize()
-
-  def test_double_gemm_relu_half_contig(self):
-    N = 1024
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    (((A@B).relu().contiguous(arg=(1,))@C).relu()).realize()
-
-  def test_double_gemm_half_contig(self):
-    N = 1024
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    ((A@B).contiguous(arg=(1,))@C).realize()
-
-  def test_double_gemm_contig(self):
-    N = 1024
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    ((A@B).contiguous()@C).realize()
-
-  def test_many_gemm(self):
-    N = 1024
-    A = Tensor.empty(N, N)
-    B = Tensor.empty(N, N)
-    C = Tensor.empty(N, N)
-    D = Tensor.empty(N, N)
-    E = Tensor.empty(N, N)
-    F = Tensor.empty(N, N)
-    (A@B@C@D@E@F).realize()
-
-  def test_conv2d(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    x.conv2d(w1).realize()
-
-  def test_conv2d_t(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    (x*2).conv2d(w1).realize()
-
-  def test_double_conv2d(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    x.conv2d(w1).conv2d(w2).realize()
-
-  def test_double_conv2d_half_contig(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    # NOTE: this contiguous doesn't help
-    x.conv2d(w1).contiguous(arg=(1,)).conv2d(w2).permute(0,2,3,1).contiguous().realize()
-
-  def test_double_conv2d_contig(self):
-    x = Tensor.empty(1, 4, 32, 32)
-    w1 = Tensor.empty(8, 4, 3, 3)
-    w2 = Tensor.empty(12, 8, 3, 3)
-    x.conv2d(w1).contiguous().conv2d(w2).realize()
-
-  def test_ffn(self):
-    from tinygrad.apps.llm import TransformerBlock
-    from tinygrad import nn
-    blk = TransformerBlock(1024, 4096, 1, 1, 1e-5)
-    for p in nn.state.get_parameters(blk): p.replace(Tensor.empty(p.shape))
-
-    x = Tensor.empty(128, 1024)
-    out = blk._feed_forward(x)
-    out.realize()
-
-  def test_flash_attention(self):
-    BS = 4
-    HEADS = 2
-    MATDIM = 16
-    EMB = 8
-    q = Tensor.empty(BS, HEADS, MATDIM, EMB)
-    k = Tensor.empty(BS, HEADS, MATDIM, EMB)
-    v = Tensor.empty(BS, HEADS, MATDIM, EMB)
-    q.scaled_dot_product_attention(k, v).realize()
-
-if __name__ == '__main__':
-  unittest.main()

test/test_schedule.py

@@ -1044,7 +1044,7 @@ class TestSchedule(unittest.TestCase):
     k = Tensor.randn(32,8,16,8).realize()
     v = Tensor.randn(32,8,16,8).realize()
     out = Tensor.scaled_dot_product_attention(q,k,v)
-    #run_schedule(check_schedule(out, 5))
+    run_schedule(check_schedule(out, 5))
     if getenv("CHECK", 1):
       import torch
       compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))

test/test_tiny.py

@@ -100,30 +100,12 @@ class TestTiny(unittest.TestCase):
       lambda x: x.flatten(1), nn.Linear(576, 10)]
 
     # replace random weights with ones
-    for p in nn.state.get_parameters(layers): p.replace(Tensor.empty(p.shape))
+    Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])
 
     # run model inference
-    probs = Tensor.empty(1, 1, 28, 28).sequential(layers).tolist()
+    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
-    for p in nn.state.get_parameters(layers): p.replace(Tensor.empty(p.shape))
-    #for p in nn.state.get_parameters(layers): p.replace(Tensor.ones_like(p).contiguous().realize())
-
-    # 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")

test/unit/test_graph_rewrite.py

@@ -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.CONST)),
-      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (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)),
     ])
     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.CONST)),
-      (UPat(Ops.CONST, name="x"), lambda x: x.replace(op=Ops.DEFINE_VAR)),
+      (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)),
     ])
     with self.assertRaises(RuntimeError):
       graph_rewrite(a, pm, bottom_up=True)

test/unit/test_rewrite_tracked_childen.py

@@ -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 kernelize_sym, merge_views
+from tinygrad.schedule.kernelize import 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+kernelize_sym+extra, track_children=True)
+    graph_rewrite(r.uop, merge_views+sym+extra, track_children=True)
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_transcendental_helpers.py

@@ -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.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 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 test.helpers import eval_uop
 
 class TestTranscendentalFunctions(unittest.TestCase):

test/unit/test_uop_symbolic.py

@@ -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):
-    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)
+    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")
 
   def test_cmp_simple(self):
     self.helper_test_variable(Variable("a", 3, 8) < 4, 0, 1, "(a<4)")
@@ -266,16 +266,6 @@ 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)")
 
@@ -385,17 +375,6 @@ 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")
@@ -442,11 +421,6 @@ 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)")
 
@@ -705,10 +679,6 @@ 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

test/unit/test_viz.py

@@ -76,7 +76,7 @@ class TestViz(BaseTestViz):
     self.assertEqual(lineno, inner.__code__.co_firstlineno)
 
   def test_exceptions(self):
-    # VIZ tracks rewrites up to and including the error
+    # VIZ tracks rewrites up to 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"], 4) # 3 successful rewrites + 1 err
+    self.assertEqual(err_step["match_count"], 3)
 
   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.CONST)
+    b = a.replace(op=Ops.DEFINE_REG)
     pm = PatternMatcher([
-      (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)),
+      (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)),
     ])
     with self.assertRaises(RuntimeError): exec_rewrite(a, [pm])
     graphs = flatten(x["graph"].values() for x in get_details(tracked_ctxs[0][0]))

tinygrad/codegen/__init__.py

@@ -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.decompositions import get_late_rewrite_patterns
+from tinygrad.uop.optional 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
@@ -82,13 +82,13 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
   supported_ops = tuple(opts.code_for_op.keys())
   extra_matcher = opts.extra_matcher if opts.extra_matcher is not None else PatternMatcher([])
 
-  # 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, name="decompositions"))
 
+  # optional pre matcher
+  if opts.pre_matcher is not None: ret.append(RewriteStep(opts.pre_matcher, name="pre_matcher"))
+
   # final rules for the renderer (without sym)
   pm_final_rewrite = pm_decomp+pm_render+extra_matcher
   ret.append(RewriteStep(pm_final_rewrite, lambda _: opts.device, name="final rewrite"))

tinygrad/codegen/devectorizer.py

@@ -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(red.dtype, identity_element(red.arg, red.dtype.scalar()))
+    identity = red.const_like(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

tinygrad/codegen/linearize.py

@@ -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(f"linearize failure {sink.op} {[x.op for x in sink.src if x.op not in DONT_PLACE_IN_BLOCK]}")
+    raise RuntimeError("linearize failure")
 
   # place the early things
   lst = sorted(dedup(sink.src), key=lambda x: x.tuplize) + list(sink.arg.lst)

tinygrad/device.py

@@ -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, MAX_BUFFER_SIZE, 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, 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 MAX_BUFFER_SIZE > 0 and self.size > MAX_BUFFER_SIZE: raise RuntimeError(f"buffer of size {self.size/1e6:.2f}M is too large")
+    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")
     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)

tinygrad/helpers.py

@@ -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, MAX_BUFFER_SIZE, AMD_LLVM = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0), ContextVar("AMD_LLVM", 1)
+ALLOW_DEVICE_USAGE, AMD_LLVM = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("AMD_LLVM", 1)
 
 @dataclass(frozen=True)
 class Metadata:

tinygrad/nn/optim.py

@@ -77,19 +77,7 @@ 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, 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)
+  return LARS(params, lr, momentum, weight_decay, nesterov, classic, tcoef=0.0, fused=fused)
 
 class LARS(Optimizer):
   """
@@ -97,11 +85,9 @@ 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, ns_steps=0, ns_params=None,
-               nesterov=False, classic=True, pre_wd=True, tcoef=0.001, fused=FUSE_OPTIM):
+  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):
     super().__init__(params, lr, fused)
-    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.momentum, self.wd, self.nesterov, self.classic, self.tcoef = momentum, weight_decay, nesterov, classic, 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]]:
@@ -112,7 +98,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.pre_wd and self.wd > 0: g = g + self.wd * t.detach()
+      if 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:
@@ -120,9 +106,6 @@ 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))

tinygrad/renderer/cstyle.py

@@ -99,14 +99,11 @@ 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})",
     Ops.SHR: lambda a,b,dtype: f"({a}>>{b})", Ops.SHL: lambda a,b,dtype: f"({a}<<{b})", Ops.CMPLT: lambda a,b,dtype: f"({a}<{b})",
-    Ops.WHERE: lambda a,b,c,dtype: f"({a}?{b}:{c})", Ops.CMPEQ: lambda a,b,dtype: f"({a}=={b})",
-    # NOTE: these don't work, but they are nice for rendering
-    Ops.THREEFRY: lambda a,b,dtype: f"threefry({a},{b})", Ops.MAX: lambda a,b,dtype: f"max({a},{b})"}
+    Ops.WHERE: lambda a,b,c,dtype: f"({a}?{b}:{c})", Ops.CMPEQ: lambda a,b,dtype: f"({a}=={b})"}
 
   string_rewrite = base_rewrite
   extra_matcher = extra_pm
@@ -149,7 +146,7 @@ class CStyleLanguage(Renderer):
         if u.arg is not None: name = u.arg.function_name
         continue
       if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR):
-        r[u] = (f"data{u.arg}_{sz}" if (sz:=cast(PtrDType, u.dtype).size) > 0 else f"data{u.arg}") if u.op is Ops.DEFINE_GLOBAL else u.arg[0]
+        r[u] = f"data{u.arg}" if u.op is Ops.DEFINE_GLOBAL else u.arg[0]
         bufs[u] = (r[u], (u.dtype, False))
         continue
 
@@ -203,12 +200,11 @@ 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.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})"}
+  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})"}
   # 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, Ops.TRUNC), name="alu"), no_vectorized_alu),]) + CStyleLanguage.extra_matcher
+    (UPat(Ops.SQRT, name="alu"), no_vectorized_alu),]) + CStyleLanguage.extra_matcher
 
   if sys.platform == 'win32':
     kernel_typedef = "__attribute__((ms_abi)) void"
@@ -418,7 +414,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", ""), ("trunc", "")]]
+            for name, atr in [("fmax", "const"), ("exp2", "pure"), ("log2", "pure"), ("sqrt", "const"), ("sin", "")]]
 
   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
@@ -427,7 +423,6 @@ 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})",

tinygrad/renderer/llvmir.py

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

tinygrad/renderer/ptx.py

@@ -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, prod
+from tinygrad.helpers import flatten, get_single_element
 
 def render_val(x, dtype):
   if dtypes.is_float(dtype):
@@ -19,7 +19,6 @@ 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};",
@@ -34,7 +33,7 @@ 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, Ops.TRUNC)
+supports_half = (Ops.EXP2, Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPLT, Ops.WHERE)
 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)
@@ -152,12 +151,11 @@ 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, uops) -> str:
+  def render_kernel(self, kernel, function_name, bufs, regs) -> 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.format(launch_bounds=launch_bounds)} {function_name} (\n\t{params}\n)\n.maxntid {launch_bounds}\n{{\n{kernel}\n}}"
+    return f"{self.kernel_prefix} {function_name}(\n\t{params}\n)\n{{\n{kernel}\n}}"
 
   def render(self, uops:list[UOp]) -> str:
     kernel:list[str] = []
@@ -224,4 +222,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(), uops)
+    return self.render_kernel(kernel, name, bufs, c.items())

tinygrad/runtime/ops_amd.py

@@ -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[0] < self.pm4.PACKET3_SET_SH_REG_END:
+    if self.pm4.PACKET3_SET_SH_REG_START <= reg.addr < 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[0] < self.pm4.PACKET3_SET_UCONFIG_REG_START + 2**16-1:
+    elif self.pm4.PACKET3_SET_UCONFIG_REG_START <= reg.addr < 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[0]}) via pm4 packet')
-    self.pkt3(set_packet, reg.addr[0] - set_packet_start, *(args or (reg.encode(**kwargs),)))
+    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),)))
 
   @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[0],
-                      reg_done=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff)
+    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.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, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), 4)
+                self.gc.regSQ_THREAD_TRACE_STATUS.addr, 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, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_done'), 4)
+                self.gc.regSQ_THREAD_TRACE_STATUS.addr, 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, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('busy'), 4)
+                self.gc.regSQ_THREAD_TRACE_STATUS.addr, 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], 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, *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}
-    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.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}
     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,7 +653,8 @@ 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_offsets, self.ip_versions = self.dev_impl.regs_offset, self.dev_impl.ip_ver
+    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()}
 
     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
@@ -761,7 +762,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], {i:nbio_pad+x for i,x in self.iface.ip_offsets[am.NBIF_HWIP].items()})
+    self.nbio = AMDIP(nbio_name, self.iface.ip_versions[am.NBIF_HWIP], nbio_pad+self.iface.ip_offsets[am.NBIF_HWIP])
 
     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)

tinygrad/runtime/ops_python.py

@@ -61,8 +61,7 @@ class PythonProgram:
           i += 1
           continue
         if uop in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}:
-          assert isinstance(dtype, PtrDType), dtype
-          if dtype.fmt is None: raise RuntimeError(f"{dtype=} is not supported")
+          assert dtype.fmt is not None and isinstance(dtype, PtrDType)
           if TYPE_CHECKING or sys.version_info < (3, 12): assert dtype.fmt != "e"
           if uop is Ops.DEFINE_REG:
             # REGs are per thread

tinygrad/runtime/support/am/amdev.py

@@ -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
+@dataclasses.dataclass(frozen=True)
 class AMRegister(AMDReg):
   adev:AMDev
 
-  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 read(self): return self.adev.rreg(self.addr)
+  def read_bitfields(self) -> dict[str, int]: return self.decode(self.read())
 
-  def write(self, _am_val:int=0, inst=0, **kwargs): self.adev.wreg(self.addr[inst], _am_val | self.encode(**kwargs))
+  def write(self, _am_val:int=0, **kwargs): self.adev.wreg(self.addr, _am_val | self.encode(**kwargs))
 
-  def update(self, inst=0, **kwargs): self.write(self.read(inst=inst) & ~self.fields_mask(*kwargs.keys()), inst=inst, **kwargs)
+  def update(self, **kwargs): self.write(self.read() & ~self.fields_mask(*kwargs.keys()), **kwargs)
 
 class AMFirmware:
   def __init__(self, adev):
@@ -254,6 +254,6 @@ 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])))
-    self.__dict__.update(import_asic_regs('mp', (11, 0), cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[am.MP1_HWIP])))
+      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])))
 

tinygrad/runtime/support/am/ip.py

@@ -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[0], self.adev.regCP_HQD_PQ_WPTR_HI.addr[0] + 1)):
+    for i, reg in enumerate(range(self.adev.regCP_MQD_BASE_ADDR.addr, self.adev.regCP_HQD_PQ_WPTR_HI.addr + 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()) + 1,
+    msg = am.struct_psp_gfx_rb_frame(fence_value=(prev_wptr:=self.adev.reg(f"{self.reg_pref}_67").read()),
       cmd_buf_addr_lo=lo32(self.adev.paddr2mc(self.cmd_paddr)), cmd_buf_addr_hi=hi32(self.adev.paddr2mc(self.cmd_paddr)),
       fence_addr_lo=lo32(self.adev.paddr2mc(self.fence_paddr)), fence_addr_hi=hi32(self.adev.paddr2mc(self.fence_paddr)))
 
@@ -469,7 +469,8 @@ 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)
 
-    wait_cond(lambda: self.adev.vram.view(self.fence_paddr, 4, 'I')[0], value=msg.fence_value, msg="sOS ring not responding")
+    while self.adev.vram.view(self.fence_paddr, 4, 'I')[0] != prev_wptr: pass
+    time.sleep(0.005)
 
     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}")

tinygrad/runtime/support/amd.py

@@ -5,10 +5,9 @@ from tinygrad.helpers import getbits, round_up, fetch
 from tinygrad.runtime.autogen import pci
 from tinygrad.runtime.support.usb import ASM24Controller
 
-@dataclass
+@dataclass(frozen=True)
 class AMDReg:
-  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() }
+  name:str; offset:int; segment:int; fields:dict[str, tuple[int, int]]; bases:tuple[int, ...] # noqa: E702
 
   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()}
@@ -16,9 +15,12 @@ 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:dict[int, tuple[int, ...]] # noqa: E702
+  name:str; version:tuple[int, ...]; bases:tuple[int, ...] # noqa: E702
   def __post_init__(self): self.version = fixup_ip_version(self.name, self.version)[0]
 
   @functools.cached_property

tinygrad/runtime/support/system.py

@@ -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 []:
-      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
+      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 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")

tinygrad/schedule/kernelize.py

@@ -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)
 
-kernelize_sym = symbolic_simple+PatternMatcher([
+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+kernelize_sym+replace_contiguous, ctx={}, name="merge_views")
+  tensor_map = graph_rewrite_map(sink, multi_pm+do_fuse+merge_views+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")

tinygrad/schedule/rangeify.py

@@ -1,663 +0,0 @@
-from typing import Any
-from dataclasses import dataclass, field
-from tinygrad.dtype import dtypes, AddrSpace, PtrDType
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, colored, flatten, dedup
-from tinygrad.uop.symbolic import symbolic_simple, sym
-
-from tinygrad.schedule.kernelize import Kernel
-from tinygrad.uop.ops import track_rewrites, graph_rewrite_map, graph_rewrite, KernelInfo, identity_element
-
-imported_rewrites = 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
-  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
-  # reduce of size 0 is the identity element
-  (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)),
-   lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if x.size == 0 and reduce.size != 0 else None),
-])
-
-earliest_rewrites = imported_rewrites+PatternMatcher([
-  # RESHAPE on RESHAPE is the second reshape
-  (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE),), name="x"), lambda x: x.replace(src=(x.src[0].src[0],))),
-  # non shape changing RESHAPE is NOOP
-  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0] if x.src[0].shape == x.arg else None),
-  # RESHAPE after COPY
-  (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).reshape(r.arg)),
-  # TODO: this should be BUFFER_VIEW
-  (UPat(Ops.COPY, src=(UPat(Ops.SHRINK, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).shrink(r.arg)),
-  # const hacks
-  (UPat(Ops.CONST, name="x"), lambda x:
-   x.replace(src=(x.src[0].src[0],)).reshape((1,)*len(x.shape)).expand(x.shape) if \
-    len(x.src) and x.src[0].op is Ops.VIEW and not any(s == 0 for s in x.shape) else None),
-  # assign only to buffer
-  (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.BUFFER}), UPat(name="x"))), lambda x: x if x.src[0].base.op is not Ops.BUFFER else None),
-])
-
-# 1. add contiguous where we have to
-
-ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
-                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
-                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD}
-
-def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None
-
-def realize_parents(ctx:dict[UOp, None], rb:UOp) -> None:
-  for s in rb.src:
-    if s.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
-
-do_realize = PatternMatcher([
-  # always realize SINK parents
-  (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
-  # always realize ASSIGN/CONTIGUOUS/COPY/BUFFER_VIEW
-  (UPat({Ops.ASSIGN, Ops.COPY, Ops.BUFFER_VIEW}, name="tr"), realize),
-  # realize parents of COPY, MSELECT, MSTACK
-  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_parents),
-])
-
-add_contiguous = PatternMatcher([(UPat(GroupOp.All-{Ops.CONTIGUOUS}, name="x"),
-                                  lambda ctx,x: x.replace(tag=1).contiguous() if x in ctx and x.tag is None else None)])
-remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
-early_cleanups = PatternMatcher([(UPat().contiguous(name="c").contiguous(), lambda c: c),])
-
-# 2. mark all children
-
-@dataclass
-class ChildrenContext: children: dict[UOp, list[UOp]]|None = None
-def extract_children(ctx:ChildrenContext, x:UOp):
-  if ctx.children is not None: return
-  children_map = x.get_children_map()
-  ctx.children = {}
-  for k,v in children_map.items():
-    non_sink_children = [u for u in v if u.op is not Ops.SINK]
-    if len(non_sink_children) <= 1: continue
-    if any(x.op is Ops.REDUCE_AXIS for x in k.toposort()):
-      ctx.children[k] = non_sink_children
-
-def mark_children(ctx:ChildrenContext, x:UOp):
-  new_srcs = [(UOp(Ops.CHILD, s.dtype, src=(UOp(Ops.CHILDREN, s.dtype, (s,), arg=len(ctx.children[s])),),
-                   arg=(ctx.children[s].index(x), len(ctx.children[s]))) if s in ctx.children else s) for s in x.src]
-  return x.replace(src=tuple(new_srcs))
-
-pm_children = PatternMatcher([
-  (UPat(Ops.SINK, name="x"), extract_children),
-  (UPat(GroupOp.All-{Ops.CHILD, Ops.CHILDREN}, name="x"), mark_children),
-
-  # hack for one kernel threefry
-  #(UPat(Ops.CHILD, src=(UPat(Ops.THREEFRY, name="x"),)), lambda x: x),
-])
-
-# 3. rangeify
-
-@dataclass
-class RangeifyContext:
-  idx: int = 0
-  regs: int = 0
-  seen_children: dict[UOp, dict[int, UOp]] = field(default_factory=dict)
-  seen_child: dict[UOp, Any] = field(default_factory=dict)
-  progress: int = 0
-  children: dict[UOp, list[UOp]]|None = None
-
-def map_reshape(idx:UOp, r:UOp):
-  acc = 1
-  to_sum = []
-  for s,src in list(zip(idx.shape, idx.src[1:]))[::-1]:
-    to_sum.append(acc*src)
-    acc *= s
-  mish = sum(to_sum)
-  ret = []
-  for s in r.src[0].shape[::-1]:
-    if resolve(s!=1):
-      # this MOD should limit any ranges outside s
-      ret.append(mish % s)
-      mish //= s
-    else:
-      ret.append(UOp.const(dtypes.int, 0))
-  ret = UOp.sink(*ret).simplify().src[::-1] if len(ret) else ()
-  return r.src[0].index(*ret, dtype=idx.dtype, arg=idx.arg)
-
-def map_pad(idx:UOp, r:UOp):
-  ret = list(idx.src[1:])
-  bigwhere = UOp.const(dtypes.bool, True)
-  for i,(sh,(s,e)) in enumerate(zip(r.shape, r.arg)):
-    if s == 0 and e == 0: continue
-    where = UOp.const(dtypes.bool, True)
-    if e > 0: where = where & (ret[i] < (sh-e))
-    if s > 0: where = where & (ret[i] >= s)
-    bigwhere = bigwhere & where
-    # this is safe but dumb
-    ret[i] = (ret[i] - s).maximum(0).minimum(r.src[0].shape[i]-1)
-  # PAD is with 0
-  return bigwhere.simplify().where(r.src[0].index(*ret, dtype=idx.dtype, arg=idx.arg), UOp.const(r.dtype, 0))
-
-def map_expand(r:UOp, idx:UOp):
-  new_rngs = []
-  ending_ranges = []
-  non_ending_ranges = []
-  for a,x,y in zip(idx.src[1:], r.src[0].shape, r.shape):
-    axis_to_range = [u for u in a.toposort() if u.op is Ops.RANGE]
-    if resolve(x!=y, False):
-      ending_ranges.extend(axis_to_range)
-      new_rngs.append(a.const_like(0))
-    else:
-      non_ending_ranges.extend(axis_to_range)
-      new_rngs.append(a)
-  ending_ranges = [x.arg for x in ending_ranges if x not in non_ending_ranges]
-  if idx.arg is not None: ending_ranges.append(idx.arg)
-  return r.src[0].index(*new_rngs, arg=min([x for x in ending_ranges]) if ending_ranges else None)
-
-pm_mops = PatternMatcher([
-  # this is like the definitions of these
-  (UPat(Ops.INDEX, src=(UPat(Ops.SHRINK, name="r"),), allow_any_len=True, name="idx"),
-   lambda r,idx: r.src[0].index(*[a+ss if resolve(ss != 0) else a for a,(ss,_) in zip(idx.src[1:], r.arg)], dtype=idx.dtype, arg=idx.arg)),
-  (UPat(Ops.INDEX, src=(UPat(Ops.PERMUTE, name="r"),), allow_any_len=True, name="idx"),
-   lambda r,idx: r.src[0].index(*[idx.src[1+p] for p in argsort(idx.src[0].arg)], dtype=idx.dtype, arg=idx.arg)),
-  (UPat(Ops.INDEX, src=(UPat(Ops.FLIP, name="r"),), allow_any_len=True, name="idx"),
-   lambda r,idx: r.src[0].index(*[((s-1)-a) if f else a for a,s,f in zip(idx.src[1:], r.shape, r.arg)], dtype=idx.dtype, arg=idx.arg)),
-  # expand needs to end ranges
-  (UPat(Ops.INDEX, src=(UPat(Ops.EXPAND, name="r"),), allow_any_len=True, name="idx"), map_expand),
-  # reshape does a lot of symbolic stuff
-  (UPat(Ops.INDEX, src=(UPat(Ops.RESHAPE, name="r"),), allow_any_len=True, name="idx"), map_reshape),
-  # pad adds min and max
-  (UPat(Ops.INDEX, src=(UPat(Ops.PAD, name="r"),), allow_any_len=True, name="idx"), map_pad),
-])
-
-def map_contiguous(ctx:RangeifyContext, x:UOp, idx:UOp|None=None):
-  # NOTE: partial contig is disabled for now
-  #arg = x.arg
-  arg = None
-  if arg is None and idx is not None: return None
-  if arg is not None and idx is None: return None
-
-  ranges = []
-  new_ranges = []
-  passthrough_idx = []
-  for i,s in enumerate(x.shape):
-    if arg is not None and i not in arg:
-      assert idx is not None, "partial contig requires index"
-      ranges.append(idx.src[1+i])
-      continue
-    if idx is not None: passthrough_idx.append(idx.src[1+i])
-    if resolve(s!=1):
-      ranges.append(UOp.range(dtypes.int, s, ctx.idx))
-      new_ranges.append(ranges[-1])
-      ctx.idx += 1
-    else:
-      ranges.append(UOp.const(dtypes.int, 0))
-  ret = x.src[0].index(*ranges).bufferize(*new_ranges, arg=x.device)
-  ret = ret.index(*passthrough_idx) if len(passthrough_idx) else ret.reshape(x.shape)
-  return ret
-
-def map_reduce(ctx:RangeifyContext, idx:UOp, red:UOp):
-  rngs = list(idx.src[1:])
-  new_ranges = []
-  for i,s in enumerate(red.src[0].shape):
-    if i in red.arg[1]:
-      rngs[i] = UOp.range(dtypes.int, s, ctx.idx)
-      ctx.idx += 1
-      new_ranges.append(rngs[i])
-  return UOp(Ops.REDUCE, red.dtype, src=(red.src[0].index(*rngs),)+tuple(new_ranges), arg=red.arg[0])
-
-def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp):
-  if c not in ctx.seen_children: ctx.seen_children[c] = {}
-  # wait here until we have seen all the children
-  if len(ctx.seen_children[c]) != x.arg[1]:
-    ctx.progress += 1
-    if ctx.progress > 10000: raise RuntimeError("children not making progress")
-    # NOTE: we mark this here
-    ctx.seen_children[c][x.arg[0]] = idx
-    raise RewriteNotReady
-  ctx.progress = 0
-
-  if c not in ctx.seen_child:
-    all_rngs = zip(*[ch.src[1:] for ch in ctx.seen_children[c].values()])
-    out_rngs = []
-    end_ranges = []
-    idx_ranges = []
-    for i,r in enumerate(all_rngs):
-      if all_same(r):
-        out_rngs.append(r[0])
-      else:
-        out_rngs.append(UOp.range(dtypes.int, c.shape[i], ctx.idx))
-        ctx.idx += 1
-        end_ranges.append(out_rngs[-1])
-        idx_ranges.append(i)
-    ctx.seen_child[c] = (idx_ranges, end_ranges)
-  else:
-    out_rngs = list(idx.src[1:])
-    idx_ranges, end_ranges = ctx.seen_child[c]
-    for i,nr in zip(idx_ranges, end_ranges): out_rngs[i] = nr
-  if len(idx_ranges) == 0: return c.index(*out_rngs)
-  # NOTE: partial contigs can still come from here
-  return c.index(*out_rngs).bufferize(*end_ranges, arg=x.device).index(*[idx.src[1+i] for i in idx_ranges])
-
-def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp):
-  if len(ctx.seen_children[c]) != c.arg: raise RuntimeError("all children should have been seen by now")
-  return idx.replace(src=(idx.src[0].src[0],)+idx.src[1:])
-
-def might_end_axis(idx:UOp):
-  if idx.arg is None: return None
-  to_end_axis = []
-  for i,a in enumerate(idx.src[1:]):
-    if any(x.arg > idx.arg for x in a.toposort() if x.op is Ops.RANGE):
-      to_end_axis.append(i)
-  if to_end_axis: return idx.replace(src=(idx.src[0].contiguous(arg=tuple(to_end_axis)),)+idx.src[1:], arg=None)
-  return idx.replace(arg=None)
-
-pm_rangeify = pm_mops+PatternMatcher([
-  # sink contigs to kick it off
-  (UPat(Ops.CONTIGUOUS, src=(UPat(),), name="x"), map_contiguous),
-
-  # if there are new ended children, tag the SINK
-  (UPat(Ops.INDEX, src=(UPat(Ops.CHILD, src=(UPat(name="c"), ), name="x"),), allow_any_len=True, name="idx"), index_child),
-  (UPat(Ops.INDEX, src=(UPat(Ops.CHILDREN, name="c"),), allow_any_len=True, name="idx"), children_gate),
-
-  # if we come across this, remove it. it was a CHILD unused in an INDEX
-  (UPat(Ops.CHILD, src=(UPat(Ops.CHILDREN, src=(UPat.var("x"),)),)), lambda x: x),
-
-  # if there's an INDEX it can support partial contig
-  (UPat(Ops.INDEX, src=(UPat(Ops.CONTIGUOUS, src=(UPat(),), name="x"),), allow_any_len=True, name="idx"), map_contiguous),
-
-  # CONST can't have axes. remove srcs when we idx
-  (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),)), lambda c: c.replace(src=())),
-
-  # handle arg on any op with weight. old endrange stuff
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.REDUCE_AXIS})),), allow_any_len=True, name="idx"), might_end_axis),
-
-  # move MAP through elementwise ALU / reduce. these are the items with cost
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.STORE, Ops.ASSIGN, Ops.COPY, Ops.DEVICE})),), allow_any_len=True, name="x"),
-   lambda x: x.src[0].replace(src=tuple([s.index(*x.src[1:]) for s in x.src[0].src]))),
-  (UPat(Ops.INDEX, src=(UPat(Ops.REDUCE_AXIS, name="red"),), allow_any_len=True, name="idx"), map_reduce),
-])
-
-# 4. remove bufferize
-
-def bufferize_to_store(ctx, x:UOp):
-  rngs = x.src[1:]
-  shape = tuple([r.vmax+1 for r in rngs])
-  assert prod(shape) > 0, f"no zero sized buffers {shape}"
-  store_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
-  if x.src[0].op is Ops.ASSIGN:
-    return x.src[0].src[0].replace(dtype=x.dtype.ptr(size=prod(shape))).store(x.src[0].src[1], *store_rngs)
-  #buf = UOp.new_buffer(x.arg, prod(shape), x.dtype)
-  buf = UOp(Ops.DEFINE_LOCAL, x.dtype.ptr(size=prod(shape)), arg=ctx[0])
-  ctx[0] += 1
-  return buf.reshape(shape).index(*rngs, dtype=x.dtype.ptr(size=prod(shape))).store(x.src[0], *store_rngs)
-
-def add_load_on_buffer(idx:UOp, b:UOp):
-  if isinstance(idx.dtype, PtrDType): return None
-  return idx.replace(dtype=idx.dtype.ptr(b.size), arg=None).load()
-
-def add_load_on_store(x:UOp, st:UOp):
-  if isinstance(x.dtype, PtrDType): return None
-  rngs = x.src[1:]
-  shape = tuple([r.vmax+1 for r in rngs])
-  b = st.src[0].src[0]
-  #assert b.op is Ops.BUFFER
-  return b.shrink(((0,prod(shape)),)).reshape(shape).index(*rngs, dtype=x.dtype.ptr(size=b.size)).load(st)
-
-def shp(shp, rng):
-  acc = 1
-  ss = []
-  for s,r in list(zip(shp,rng))[::-1]:
-    ss.append(r*acc)
-    acc *= s
-  return sum(ss)
-
-pm_add_buffers = pm_mops+PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), bufferize_to_store),
-  (UPat(Ops.INDEX, src=(UPat(Ops.BUFFER, name="b"), UPat()), name="idx"), add_load_on_buffer),
-  (UPat(Ops.INDEX, src=(UPat(Ops.STORE, name="st"),), allow_any_len=True, name="x"), add_load_on_store),
-
-  # HACK
-  #(UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
-
-  (UPat(Ops.INDEX, name="idx").contiguous(),
-   lambda idx: UOp.new_buffer(idx.device, prod(idx.arg), idx.dtype).index(shp(idx.arg, idx.src[1:]),
-                                                                          dtype=idx.dtype.ptr(prod(idx.arg))).store(*idx.src))
-])
-
-# 5 (alt). create pointers
-
-def debuf(ctx, b:UOp):
-  ret = UOp(Ops.DEFINE_GLOBAL, b.dtype.ptr(b.arg), arg=ctx[0])
-  ctx[0] += 1
-  return ret
-
-pm_debuf = PatternMatcher([
-  (UPat(Ops.BUFFER, name="b"), debuf),
-  # HACK: consts shouldn't have srcs by here
-  (UPat(Ops.CONST, name="x"), lambda x: x.replace(src=()) if len(x.src) else None),
-  # no movement ops
-  (UPat(GroupOp.Movement, name="x"), lambda x: x.src[0]),
-  # HACK: no copy
-  (UPat(Ops.COPY, name="x"), lambda x: x.src[0]),
-])
-
-# 5. split into kernels
-
-@dataclass
-class LocalAddBufferContext:
-  dg:int = 0
-  map:dict = field(default_factory=dict)
-  vars:dict = field(default_factory=dict)
-
-def debuf(ctx:LocalAddBufferContext, b:UOp): return UOp(Ops.DEFINE_GLOBAL, b.dtype.ptr(b.arg), arg=ctx.map[b][1])
-def unbind_kernel(ctx:LocalAddBufferContext, b:UOp):
-  ctx.vars[b] = None
-  return b.src[0]
-
-def split_load(ctx:LocalAddBufferContext, s:UOp):
-  b = s.src[0].src[0]
-  if b.op is not Ops.BUFFER: return None
-
-  if len(s.src) == 2 and s.src[1].op is Ops.ASSIGN:
-    assert len(s.src) == 2
-    lb = s.src[1]
-    assert b not in ctx.map or ctx.map[b][0] == lb
-  else:
-    lb = b
-  if b not in ctx.map:
-    ctx.map[b] = (lb, ctx.dg)
-    ctx.dg += 1
-  return s.replace(src=s.src[0:1]) if b is not lb else None
-
-def handle_store(ctx:LocalAddBufferContext, s:UOp):
-  b = s.src[0].src[0]
-  if b.op is not Ops.BUFFER: return None
-  if b not in ctx.map:
-    ctx.map[b] = (b, ctx.dg)
-    ctx.dg += 1
-  if s.src[1].op is Ops.COPY: return s.src[1]
-  return None
-
-to_define_global = PatternMatcher([
-  (UPat(Ops.BUFFER, name="b"), debuf),
-  (UPat(Ops.BIND, name="b"), unbind_kernel),
-  (UPat(Ops.LOAD, name="s"), split_load),
-  (UPat(Ops.STORE, name="s"), handle_store),
-])
-
-def split_store(x:UOp):
-  if len(x.ranges): return None
-  store_rngs = x.src[2:]
-
-  ctx = LocalAddBufferContext()
-  ret = graph_rewrite(x, to_define_global, ctx=ctx, name="kernel split", bottom_up=True)
-  rng = sorted([u for u in ret.toposort() if u.op is Ops.RANGE], key=lambda x: x.arg)
-  name = "k"+colored('_', 'BLACK').join(['']+[colored(str(s.vmax+1), "WHITE") if s in store_rngs else colored(str(s.vmax+1), "red") for s in rng])
-
-  ret = ret.sink(arg=KernelInfo(name=name)) if ret.op is Ops.STORE else ret
-  kernel = UOp(Ops.KERNEL, src=tuple([x[0] for x in ctx.map.values()])+tuple(ctx.vars.keys()), arg=Kernel(ret, ()))
-  return kernel.src[0].assign(kernel)
-
-split_kernels = PatternMatcher([
-  (UPat(Ops.STORE, name="x"), split_store),
-])
-
-pm_children_fixup = PatternMatcher([
-  # clone all movement ops
-  (UPat(Ops.CHILD, src=(UPat(Ops.CHILDREN, src=(UPat(GroupOp.Movement, name="m"),)),), name="c"),
-    lambda c,m: UOp(m.op, m.dtype, (c.replace(src=(c.src[0].replace(src=(m.src[0],)),)),), m.arg)),
-])
-
-
-@dataclass
-class RContext:
-  range_num = 0
-
-def new_range(ctx, s):
-  ret = UOp.range(dtypes.int, s, ctx.range_num)
-  ctx.range_num += 1
-  return ret
-
-def td_reshape(ctx, idx:UOp, r:UOp):
-  acc = 1
-  to_sum = []
-  for s,i in list(zip(idx.arg, idx.src[1:]))[::-1]:
-    to_sum.append(i*acc)
-    acc *= s
-  mish = sum(to_sum)
-  ret = []
-  for s in r.arg[::-1]:
-    if resolve(s!=1):
-      # this MOD should limit any ranges outside s
-      ret.append(mish % s)
-      mish //= s
-    else:
-      ret.append(UOp.const(dtypes.int, 0))
-  ret = UOp.sink(*ret).simplify().src[::-1] if len(ret) else ()
-  ii = idx.src[0]
-  out_rng = ret
-
-  """
-  out_rng = []
-  for i,rr in enumerate(ret):
-    if rr.op not in {Ops.RANGE, Ops.CONST}:
-      out_rng.append(new_range(ctx, r.arg[i]))
-    else:
-      out_rng.append(rr)
-
-  mm = [idx.src[0]]
-  for x,y in zip(ret, out_rng):
-    if x is not y:
-      mm.append(x)
-      mm.append(y)
-  if len(mm) > 1:
-    ii = UOp(Ops.MERGE, idx.dtype, tuple(mm))
-  """
-  return ii.index(*out_rng, dtype=idx.dtype, arg=r.arg)
-
-def td_elementwise(ctx, e:UOp):
-  # if the range is closed by a reduce to the left, we can't reuse it
-  # TODO: handle composite ranges better
-  reduces_left = flatten([x.src[1:] for x in e.toposort() if x.op is Ops.REDUCE])
-  shps = [u.arg for u in e.src]
-  assert all_same(shps)
-  rngs = [u.src[1:] for u in e.src]
-  out_rng = []
-  need_merge = False
-  for i,r in enumerate(zip(*rngs)):
-    r = [x for x in r if x is not UOp.const(dtypes.int, 0)]
-    if len(r) == 0:
-      out_rng.append(UOp.const(dtypes.int, 0))
-    elif all_same(r) and r[0] not in reduces_left:
-      out_rng.append(r[0])
-    else:
-      out_rng.append(new_range(ctx, shps[0][i]))
-      need_merge = True
-  if need_merge:
-    new_src = []
-    for u in e.src:
-      assert u.op is Ops.INDEX
-      out = [u.src[0]]
-      rngs_in_src = [x for x in out[0].toposort() if x.op is Ops.RANGE]
-      for i,idx in list(enumerate(u.src[1:]))[::-1]:
-        rngs_in_idx = [x for x in idx.toposort() if x.op is Ops.RANGE]
-        if all(x not in rngs_in_src for x in rngs_in_idx):
-          # for expands
-          continue
-        if idx is not out_rng[i]:
-          out.append(idx)
-          out.append(out_rng[i])
-          #out = UOp(Ops.MERGE, out.dtype, src=(out, idx, out_rng[i]))
-      if len(out) > 1:
-        new_src.append(UOp(Ops.MERGE, u.dtype, tuple(out)))
-      else:
-        new_src.append(out[0])
-      #mm = []
-      #for i,idx in enumerate(u.src[1:]):
-      #  if idx is not out_rng[i] and idx is not UOp.const(dtypes.int, 0):
-      #    mm.append(UOp(Ops.MERGE, src=(idx, out_rng[i])))
-      #new_src.append(UOp(Ops.MBLOCK, u.dtype, (u.src[0],)+tuple(mm)))
-  else:
-    new_src = list([x.src[0] for x in e.src])
-  return e.replace(src=tuple(new_src)).index(*out_rng, arg=shps[0])
-
-def td_shrink(idx:UOp, r:UOp):
-  ret = []
-  shp = []
-  for u,(s,e),shape in zip(idx.src[1:], r.arg, idx.arg):
-    assert s == 0
-    #if u.vmax >= e: u = (u<e).where(u, UOp(Ops.INVALID, u.dtype))
-    ret.append(u)
-    shp.append(min(shape, e))
-  return idx.src[0].index(*ret, dtype=idx.dtype, arg=tuple(shp))
-
-def td_reduce(ctx, idx:UOp, r:UOp):
-  rngs = idx.src[1:]
-  new_shp = tuple([s if i not in r.arg[1] else 1 for i,s in enumerate(idx.arg)])
-  return UOp(Ops.REDUCE, r.dtype, (idx.src[0],)+tuple([x for i,x in enumerate(rngs) if i in r.arg[1]]),
-                     r.arg[0]).index(*[x if i not in r.arg[1] else UOp.const(dtypes.int, 0) for i,x in enumerate(rngs)], arg=new_shp)
-
-pm_td_rangeify = PatternMatcher([
-  #(UPat(Ops.INDEX, src=(UPat(Ops.MERGE, src=(UPat(Ops.LOAD, name="b"),), allow_any_len=True),), allow_any_len=True, name="idx"),
-  # lambda idx,b: b.src[0].src[0].index(*idx.src[1:], dtype=b.src[0].dtype).load().index(*idx.src[1:], arg=idx.arg)),
-  (UPat(Ops.BUFFER, name="b"), lambda ctx, b:
-    b.replace(tag=1).index(nr:=new_range(ctx, b.size), dtype=b.dtype.ptr(size=b.size)).load().index(nr, arg=(b.size,)) if b.tag is None else None),
-    #b.replace(tag=1).index(new_range(ctx, b.size), arg=(b.size,)) if b.tag is None else None),
-  (UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="c"), lambda c: c.replace(src=()).index(arg=())),
-  (UPat(Ops.RESHAPE, src=(UPat(Ops.INDEX, name="idx"),), name="r"), td_reshape),
-  (UPat(Ops.SHRINK, src=(UPat(Ops.INDEX, name="idx"),), name="r"), td_shrink),
-  (UPat(Ops.PERMUTE, src=(UPat(Ops.INDEX, name="idx"),), name="r"),
-    lambda r,idx: idx.src[0].index(*[idx.src[1+p] for p in r.arg], dtype=idx.dtype, arg=tuple(idx.arg[p] for p in r.arg))),
-  # 0s are already in place for EXPAND
-  #(UPat(Ops.EXPAND, src=(UPat(Ops.INDEX, name="idx"),), name="r"), lambda r,idx: idx.replace(arg=r.arg)),
-  (UPat(Ops.EXPAND, src=(UPat(Ops.INDEX, name="idx"),), name="r"),
-    lambda ctx,r,idx: idx.src[0].index(*[ii if s1==s2 else new_range(ctx, s1) for s1,s2,ii in zip(r.arg, idx.arg, idx.src[1:])], arg=r.arg)),
-  (UPat(GroupOp.Elementwise, src=UPat(Ops.INDEX), name="e"), td_elementwise),
-  (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.INDEX, name="idx"),), name="r"), td_reduce),
-])
-
-def remove_merge(m):
-  tr0, tr1 = [], []
-  for r0,r1 in zip(m.src[1::2], m.src[2::2]):
-    if r0 is r1: continue
-    tr0.append(r0)
-    tr1.append(r1)
-  if m.src[0].op is Ops.LOAD and False:
-    # hack for LOAD
-    reps = {k:v for k,v in zip(tr0, tr1)}
-    return m.src[0].substitute(reps)
-  return UOp(Ops.BUFFERIZE, m.dtype, (m.src[0],)+tuple(tr0), arg=m.device).index(*tr1)
-
-no_merge = PatternMatcher([
-  (UPat(Ops.MERGE, name="m"), remove_merge),
-])
-
-@track_rewrites(name=lambda sink,ret: f"Schedule {pluralize('Kernel',len([u for u in ret[sink].toposort() if u.op is Ops.KERNEL]))}", replay=True)
-def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]:
-  tensor_map = graph_rewrite_map(sink, earliest_rewrites, name="earliest")
-  realize_map = {}
-  graph_rewrite(tensor_map[sink], do_realize, ctx=realize_map, name="Input Graph")
-  tensor_map = graph_rewrite_map(tensor_map[sink], add_contiguous, ctx=realize_map, bottom_up=True, input_map=tensor_map, name="add contiguous")
-  tensor_map = graph_rewrite_map(tensor_map[sink], early_cleanups+remove_tags, input_map=tensor_map, name="cleanup")
-  rsink = tensor_map[sink]
-
-  ctx = RContext()
-  rsink = graph_rewrite(rsink, pm_td_rangeify, ctx=ctx, name="td rangeify")
-  rsink = graph_rewrite(rsink, sym, name="symbolic")
-
-  # find MOD on RANGE to split
-  while 1:
-    #break
-    reps = {}
-    for u in rsink.toposort():
-      if u.op is Ops.MOD and u.src[0].op is Ops.RANGE and u.src[1].op is Ops.CONST:
-        r = u.src[0].vmax+1
-        c = u.src[1].arg
-        if r%c == 0:
-          reps[u.src[0]] = new_range(ctx, r//c)*c + new_range(ctx, c)
-    print(len(reps))
-    if len(reps) == 0: break
-    rsink = rsink.substitute(reps)
-    rsink = graph_rewrite(rsink, sym, name="symbolic")
-
-  for i in range(0):
-    print("loop")
-    real_rngs = rsink.ranges.copy()
-    for u in rsink.toposort():
-      if u.op is Ops.REDUCE:
-        for s in u.src[1:]: real_rngs[s] = None
-    real_rngs = {x:[] for x in real_rngs}
-    print("unmovable", [x.arg for x in real_rngs])
-
-    for u in rsink.toposort():
-      if u.op is not Ops.MERGE: continue
-      assert all(x.op is Ops.RANGE for x in u.src)
-      r0, r1 = [x for x in u.src]
-      if r0 is r1: continue
-      if r0 in real_rngs: real_rngs[r0].append(r1)
-      if r1 in real_rngs: real_rngs[r1].append(r0)
-
-    rew = {}
-    for k,v in real_rngs.items():
-      print(k.arg, [x.arg for x in v])
-      for u in v:
-        rew[u] = k
-    rsink = rsink.substitute(rew)
-
-
-  """
-  rngs = [x for x in rsink.toposort() if x.op is Ops.RANGE]
-  mmap = {16:1000, 2:8, 3:9}
-  rep = {}
-  for x in rngs:
-    if x.arg in mmap:
-      rep[x] = x.replace(arg=mmap[x.arg])
-  rsink = rsink.substitute(rep)
-  """
-
-  rsink = graph_rewrite(rsink, no_merge, name="remove merge")
-
-  """
-  tensor_map = graph_rewrite_map(tensor_map[sink], pm_children, ctx=ChildrenContext(), bottom_up=True, input_map=tensor_map, name="children")
-  tensor_map = graph_rewrite_map(tensor_map[sink], pm_children_fixup, bottom_up=True, input_map=tensor_map, name="fixup children")
-  tensor_map = graph_rewrite_map(tensor_map[sink], pm_rangeify, ctx=RangeifyContext(), bottom_up=True, input_map=tensor_map, name="rangeify")
-  """
-  #tensor_map = graph_rewrite_map(tensor_map[sink], symbolic_simple, input_map=tensor_map, name="symbolic")
-  #tensor_map = graph_rewrite_map(tensor_map[sink], pm_add_buffers, bottom_up=True, input_map=tensor_map, name="add buffers")
-  #if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Rangeify Graph")
-  if getenv("VIZ"): graph_rewrite(rsink, PatternMatcher([]), name="View Rangeify Graph")
-
-  rsink = graph_rewrite(rsink, pm_add_buffers, ctx=[0], bottom_up=True, name="add buffers")
-
-  # render
-  if getenv("SRC") or True:
-    #rsink = tensor_map[sink]
-    from tinygrad.codegen.devectorizer import pm_reduce, ReduceContext
-    rsink = graph_rewrite(rsink, pm_reduce, ctx=ReduceContext(), name="remove reduce")
-    rsink = graph_rewrite(rsink, pm_debuf, ctx=[0], name="debuf", bottom_up=True)
-    rsink = graph_rewrite(rsink, sym, name="symbolic 2")
-
-    # renumber ranges
-    #rngs = dedup([x for x in flatten([x.src[2:] for x in list(rsink.toposort())[::-1] if x.op is Ops.STORE]) if x.op is Ops.RANGE])
-    #rsink = rsink.substitute({x:x.replace(arg=i) for i,x in enumerate(rngs)})
-
-    from tinygrad.codegen import rewrites_for_linearizer, apply_rewrites
-    rsink = apply_rewrites(rsink, rewrites_for_linearizer)
-    from tinygrad.renderer.cstyle import CStyleLanguage
-    src = CStyleLanguage().render(rsink.arg.lst)
-    print(src)
-    return {sink:sink}
-
-  tensor_map = graph_rewrite_map(tensor_map[sink], split_kernels, input_map=tensor_map, name="split kernels")
-
-  # if a kernel depends on a buffer, and that buffer is later assigned to, make the assign depend on the kernel's assign
-  kernel_assign: dict[UOp, UOp] = {}
-  assign_rep: dict[UOp, UOp] = {}
-  for u in tensor_map[sink].toposort():
-    if u.op is not Ops.ASSIGN: continue
-    kernel_assign[u.buf_uop] = u
-    for s in u.src[1].src:
-      # TODO: this is probably broken for MSELECT/MSTACK
-      if s.op is not Ops.BUFFER or s is u.buf_uop or (a:=kernel_assign.get(s)) is None: continue
-      if any(x.op is Ops.ASSIGN and x.buf_uop is s for x in u.toposort()):
-        raise RuntimeError(f"cycle detected in graph, kernel for {u.buf_uop} must either depend on ASSIGN or BUFFER")
-      assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,))
-  if assign_rep:
-    tensor_map = graph_rewrite_map(tensor_map[sink], _substitute, ctx=assign_rep, bottom_up=True, input_map=tensor_map, name="fix_assign")
-
-  if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Kernel Graph")
-  return tensor_map

tinygrad/tensor.py

@@ -14,7 +14,7 @@ from tinygrad.device import Device, Buffer
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.memory import memory_planner
 from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
-from tinygrad.schedule.rangeify import get_kernelize_map
+from tinygrad.schedule.kernelize import get_kernelize_map
 
 # *** all in scope Tensors are here. this gets relevant UOps ***
 
@@ -252,8 +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) or (DEBUG >= 2 and len(schedule) > 1):
-      print(f"scheduled {len(schedule)} kernels in {(time.perf_counter()-st)*1000:.2f} ms")
+    if DEBUG >= 1 and len(schedule) >= 10: 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]:
@@ -2934,11 +2933,11 @@ class Tensor(MathTrait):
     """
     return self*-1 if self.dtype != dtypes.bool else self.logical_not()
 
-  def contiguous(self, **kwargs) -> Tensor:
+  def contiguous(self) -> Tensor:
     """
     Returns a contiguous tensor.
     """
-    return self._apply_uop(UOp.contiguous, **kwargs)
+    return self._apply_uop(UOp.contiguous)
 
   def fuse(self) -> Tensor:
     """
@@ -3144,7 +3143,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._apply_uop(UOp.trunc)
+    return self.cast(dtypes.int32).cast(self.dtype)
 
   def ceil(self: Tensor) -> Tensor:
     """
@@ -3174,7 +3173,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.trunc() / 2.0).trunc() == b)).where((self - 0.5).ceil(), (self + 0.5).floor())
+    return ((self > 0) == ((b := self.cast(dtypes.int32) / 2.0).cast(dtypes.int32) == b)).where((self - 0.5).ceil(), (self + 0.5).floor())
 
   def isinf(self:Tensor, detect_positive:bool=True, detect_negative:bool=True) -> Tensor:
     """
@@ -4034,21 +4033,6 @@ 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()

tinygrad/uop/__init__.py

@@ -12,15 +12,13 @@ class Ops(FastEnum):
   NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto()  # noqa: E702
 
   # track children
-  CHILD = auto(); CHILDREN = auto() # noqa: E702
-  MERGE = auto(); MBLOCK = auto(); INVALID = auto()
+  CHILD = auto()
 
   # buffer ops
   COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702
 
   # ops that adjust the behavior of the scheduler
   CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); FUSE = auto() # noqa: E702
-  BUFFERIZE = auto()
 
   # blocks in linearizer (only used there)
   BLOCK = auto(); BLOCKSTART = auto(); BLOCKEND = auto(); BLOCKFINAL = auto() # noqa: E702
@@ -51,7 +49,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(); TRUNC = auto() # noqa: E702
+  CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto(); SQRT = auto(); RECIP = auto(); NEG = auto() # noqa: E702
 
   # load/store before math
   LOAD = auto(); STORE = auto() # noqa: E702
@@ -82,15 +80,12 @@ class Ops(FastEnum):
   CUSTOM = auto(); CUSTOMI = auto() # noqa: E702
 
 class GroupOp:
-  Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIP, Ops.NEG, Ops.TRUNC}
+  Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIP, Ops.NEG}
   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}

tinygrad/uop/mathtraits.py

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

tinygrad/uop/ops.py

@@ -136,17 +136,11 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def st(self) -> ShapeTracker|None:
-    if self.op is Ops.INDEX and self.src[0].op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.BUFFER, Ops.BUFFERIZE}: return None
-    if self.op is Ops.MBLOCK: return None
-    if self.op in GroupOp.Block: return None
+    if self.op in GroupOp.Block or self.op is Ops.INDEX: return None
     from tinygrad.shape.shapetracker import ShapeTracker
     # VIEW and MovementOps define a new ShapeTracker from the arg
     if self.op is Ops.VIEW: return self.arg
-    if self.op is Ops.BUFFERIZE: return ShapeTracker.from_shape((prod([r.vmax+1 for r in self.src[1:]]),))
-    if self.op is Ops.RESHAPE and self.src[0].st is None: return ShapeTracker.from_shape(self.arg)
-    if self.op in GroupOp.Movement:
-      if self.src[0].st is None: return None
-      return unwrap(self.src[0].st).mop(self.op, self.arg)
+    if self.op in GroupOp.Movement: return unwrap(self.src[0].st).mop(self.op, self.arg)
     # CONST with a DEVICE has a shape of ()
     if self.op is Ops.CONST and len(self.src) and self.src[0].op is Ops.DEVICE: return ShapeTracker.from_shape(())
     # BufferOps and ASSIGN flow ShapeTracker from a direct edge
@@ -164,7 +158,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if self.op is Ops.CAST and self.src[0].op is Ops.DEFINE_GLOBAL: return None
 
     # otherwise we get the shape from sources
-    if not (src_sts := [x.st for x in self.src if x.st is not None and x.op is not Ops.INDEX]): return None
+    if not (src_sts := [x.st for x in self.src if x.st is not None]): return None
     assert all_same([x.shape for x in src_sts]), f"UOp sources must have the same shape {self} {[x.shape for x in src_sts]}"
     match self.op:
       case Ops.MULTI: shape = tuple(self.src[0].shape[a]*len(self.device) if a == self.axis else s for a,s in enumerate(self.src[0].shape))
@@ -188,31 +182,6 @@ 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 is Ops.MERGE:
-      ret = self.src[0].ranges.copy()
-      for s in self.src[1::2]:
-        if s in ret: del ret[s]
-      for s in self.src[2::2]:
-        ret.update(s.ranges)
-      return ret
-    if self.op in {Ops.BUFFERIZE, Ops.REDUCE}:
-      ret = self.src[0].ranges.copy()
-      for s in self.src[1:]:
-        if s in ret: del ret[s]
-    elif self.op in {Ops.STORE}:
-      ret = self.src[0].ranges.copy()
-      ret.update(self.src[1].ranges)
-      for s in self.src[2:]:
-        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):
@@ -250,8 +219,7 @@ 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, *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 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 __getitem__(self, idx): return self.index(idx)
   def const_like(self, b:ConstLike):
     # constants can optionally have a DEVICE source
@@ -261,10 +229,9 @@ 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])
@@ -292,8 +259,7 @@ 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:
-      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),))
+      ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device).view(unwrap(ret.st)),))
     return ret
   @staticmethod
   def range(dtype:DType, end:sint, idx:int): return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end),), arg=idx)
@@ -309,9 +275,8 @@ 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, *args, **kwargs): return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
+  def contiguous(self): return self.alu(Ops.CONTIGUOUS)
   def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD)
-  def bufferize(self, *args, **kwargs): return UOp(Ops.BUFFERIZE, dtype=self.dtype, src=(self,)+args, **kwargs)
   def fuse(self): return self.alu(Ops.FUSE)
   def allreduce(self, op, device:str|tuple[str, ...]|UOp):
     assert isinstance(self.device, tuple), f"allreduce must be on tuple {self.device} isn't"
@@ -503,7 +468,6 @@ 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
@@ -595,7 +559,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.TRUNC: math.trunc,
+  Ops.SIN: lambda x: math.sin(x) if not math.isinf(x) else math.nan, Ops.POW: safe_pow,
   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}
@@ -690,7 +654,6 @@ 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):
@@ -776,6 +739,16 @@ 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()
@@ -867,12 +840,7 @@ class TrackedPatternMatcher(PatternMatcher):
         match_stats[p][2] += time.perf_counter()-st
         continue
       match_stats[p][1] += 1
-      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.NOOP, arg=str(sys.exc_info()[1]))), p.location))
-        raise
-      if ret is not None and ret is not uop:
+      if (ret:=match(uop, ctx)) 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))
@@ -916,21 +884,10 @@ 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] = {}
-
-  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
+    self.skip_0: dict[UOp, None] = {}  # NOTE: this is needed for RewriteNotReady. it also detects some infinite loops
 
   def unified_rewrite(self, root:UOp) -> UOp:
     stack: list[tuple[UOp, int, UOp]] = [(root, 0, root)]
@@ -940,23 +897,18 @@ 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:
-            # 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)
+          if self.bpm is not None: new_n = self.bpm.fixed_point_rewrite(new_n, self.ctx)
           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.cached_pm_rewrite(new_n)) is None:
+            if self.pm is None or (new_src_n:=self.pm.rewrite(new_n, self.ctx)) is None:
               self.replace[n] = new_n
               continue
           else:
@@ -968,7 +920,7 @@ class RewriteContext:
         else:
           # in stage 2, we link the result of new_n to the result of n
           try: self.replace[n] = self.replace[new_n]
-          except KeyError: raise RewriteNotReady  # pylint: disable=raise-missing-from
+          except KeyError: raise RuntimeError("infinite loop in graph_rewrite (explicit)")  # pylint: disable=raise-missing-from
       except RewriteNotReady:
         # retry this later
         stack.insert(0, (n, stage, new_n))
@@ -984,7 +936,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 (list(sink.toposort())[::-1] if bottom_up else sink.toposort()):
+  for k in 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:

tinygrad/uop/optional.py

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

tinygrad/uop/spec.py

@@ -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, x.dtype, x.arg) 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 for x in u.src]} {u.arg}")

tinygrad/uop/symbolic.py

@@ -1,11 +1,11 @@
 # all of symbolic lives here now
-from typing import Any, cast
+from typing import Any, Literal, cast
 import math, operator, struct, functools
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
 from tinygrad.dtype import ConstType, dtypes, PtrDType, AddrSpace, can_safe_cast
 from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, cdiv, cmod, CORRECT_DIVMOD_FOLDING
-from tinygrad.uop.decompositions import xpow
+from tinygrad.uop.transcendental import xpow
 
 # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ********
 
@@ -41,7 +41,6 @@ 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
@@ -72,19 +71,6 @@ 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 ********
@@ -139,92 +125,65 @@ def canonicalize_simplex(X:UOp) -> UOp|None:
     ret.append(u)
   return functools.reduce(operator.add, ret) if changed else None
 
-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
+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
   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 d.op is Ops.MOD else d.const_like(q)
-  return None
+    return x - q*y if which is Ops.MOD else x.const_like(q)
 
-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
+  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
   for u in split_uop(x, Ops.ADD):
-    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
+    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
 
-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 ((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
+  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)
     return (y2-y1)*(v-v.vmin) + y1
-  return None
 
-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 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 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 (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
 
+  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, terms):
-    if d.op is Ops.IDIV and r!=0:
+  for q,r,f,v in zip(quotients, remainders, factors, svars):
+    if which is Ops.IDIV and (not split_rem) 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.vmin < 0 or rem.vmin < 0) and remainders: return None
-  if d.op is Ops.MOD: return gcd*(rem % (c//gcd)) + const%gcd
+  if (x_min < 0 or rem.vmin < 0) and remainders: return None
+  if which is Ops.MOD: return gcd*(rem % (c//gcd)) + const%gcd
   return rem//(c//gcd)+quo
 
 def gep_through_wmma(gep:UOp, wmma:UOp):
@@ -329,20 +288,15 @@ 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((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", dtypes.sints) // UPat.var("y"), lambda x,y: div_and_mod_folding(x,y,Ops.IDIV)),
   (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("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None),
+  (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 d.vmax <  0 else None),
 ])+gep_pushing
 
@@ -424,6 +378,22 @@ 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):
@@ -458,6 +428,16 @@ 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))),

tinygrad/uop/transcendental.py

@@ -1,9 +1,8 @@
-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, getenv
-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher
+from tinygrad.helpers import polyN
+from tinygrad.uop.ops import UOp
 
 TRANSCENDENTAL_SUPPORTED_DTYPES = (dtypes.float16, dtypes.float32, dtypes.float64)
 
@@ -80,10 +79,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(dtypes.uint64)
+  ia = (f.cast(intermediate_dtype) * 4.294967296e9).cast_vec(dtypes.uint64)
   # extract 96 relevant bits of 2/pi based on magnitude of argument
-  i = shr(e.cast(dtypes.uint64), 5)
-  e = e.cast(dtypes.int32) & 31
+  i = shr(e.cast_vec(dtypes.uint64), 5)
+  e = e.cast_vec(dtypes.int32) & 31
   offset = 32 - e
 
   def _take(an:UOp, offset:int, count:int=0) -> UOp:
@@ -91,8 +90,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: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)
+  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)
 
   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))
@@ -101,12 +100,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(dtypes.uint64) * y.cast(dtypes.uint64)
+  def _hp_mul(x:UOp, y:UOp) -> UOp: return x.cast_vec(dtypes.uint64) * y.cast_vec(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(dtypes.int32)
+  q = shr(p, 62).cast_vec(dtypes.int32)
   p = p & 0x3fffffffffffffff
   r = (p.cast(intermediate_dtype) * (3.4061215800865545e-19)).cast(d.dtype)
 
@@ -133,7 +132,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(dtypes.float32), q.cast(dtypes.float32)).cast(dtypes.float16)
+      d = _reduce_d(x.cast_vec(dtypes.float32), q.cast_vec(dtypes.float32)).cast_vec(dtypes.float16)
     else:
       # https://github.com/shibatch/sleef/blob/4e08851f59fc2b545f9c393c6a23dfd311a26308/src/libm/sleefsp.c#L464-L503
       d = q * -3.1414794921875 + x
@@ -143,9 +142,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(dtypes.int64).cast(d.dtype) * (2.0**24)
+  qdh = (d * (m_1_pi / 2.0**24)).cast_vec(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(dtypes.int32)
+  return _reduce_d(d, quadrant.cast(d.dtype)), quadrant.cast_vec(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))
@@ -224,7 +223,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(dtypes.float32)).cast(dtypes.float16)
+  if d.dtype.scalar() == dtypes.float16: return xlog2(d.cast_vec(dtypes.float32)).cast_vec(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)
@@ -261,9 +260,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(dtypes.int32).cast(exponent.dtype)
+  non_int = exponent != exponent.cast_vec(dtypes.int32).cast(exponent.dtype)
   adj = non_int.where(ret.const_like(math.nan),
-    (exponent < 0).where(-exponent, exponent).cast(dtypes.int32).mod(2).cast(dtypes.bool).where(ret.const_like(-1), ret.const_like(1)))
+    (exponent < 0).where(-exponent, exponent).cast_vec(dtypes.int32).mod(2).cast_vec(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)))
 
@@ -293,59 +292,3 @@ 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)

tinygrad/viz/index.html

@@ -73,6 +73,7 @@
     user-select: auto;
   }
   g.tag circle {
+    r: 5;
     fill: #FFD700;
     stroke: #B8860B;
     stroke-width: 0.8;
@@ -80,9 +81,10 @@
   g.tag text {
     text-anchor: middle;
     font-size: 6px;
-    fill: #08090e;
+    fill: black;
   }
   .label :is(text, p) {
+    color: #08090e;
     font-weight: 350;
   }
   .edgePath {

tinygrad/viz/js/index.js

@@ -32,11 +32,6 @@ 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)
@@ -76,8 +71,10 @@ 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");
-    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));
+    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");
     // 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) => {
@@ -87,7 +84,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)");
-    addTags(d3.select("#edge-labels").selectAll("g").data(g.edges().filter(e => g.edge(e).label != null)).join("g").attr("transform", (e) => {
+    const edgeLabels = 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;
@@ -101,7 +98,9 @@ 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").datum(e => g.edge(e).label));
+    }).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");
     if (recenter) document.getElementById("zoom-to-fit-btn").click();
   };
 
@@ -122,19 +121,6 @@ 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, area) => {
-  const shapes = [];
-  const yscale = d3.scaleLinear().domain([0, source.peak]).range([area, 0]);
-  for (const [i,e] of source.shapes.entries()) {
-    const x = e.x.map((i,_) => (source.timestamps[i] ?? data.et)-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 drawLine = (ctx, x, y) => {
   ctx.beginPath();
   ctx.moveTo(x[0], y[0]);
@@ -143,18 +129,16 @@ const drawLine = (ctx, x, y) => {
   ctx.stroke();
 }
 
-var data, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity;
+var profileRet, 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 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 });
-  const profileRet = await (await fetch("/get_profile")).json()
+  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];
@@ -163,7 +147,7 @@ async function renderProfiler() {
   const canvasTop = rect(canvas).top;
   // color by key (name/category/device)
   const colorMap = new Map();
-  data = {tracks:new Map(), axes:{}, st, et};
+  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;
@@ -171,30 +155,14 @@ async function renderProfiler() {
     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;
-      let newOffset = null;
-      for (const [track, v] of data.tracks) {
-        if (track === `${k} memory`)  {
-          // expand the y axis or reset to default size
-          const pick = [areaScale(mem.peak), maxArea*4];
-          const expand = k !== focusedDevice;
-          const [newArea, prevArea] = expand ? pick.reverse() : pick;
-          focusedDevice = expand ? k : null;
-          data.axes.y = expand ? { domain:[0, mem.peak], range:[v.offsetY+newArea, v.offsetY], fmt:"B" } : null;
-          // either way update all offsets
-          v.shapes = createPolygons(mem, newArea);
-          newOffset = newArea-prevArea;
-          v.div.style.height = rect(v.div).height+newOffset+"px";
-        } else if (newOffset != null) v.offsetY += newOffset;
-      }
-      d3.select(canvas).call(canvasZoom.transform, zoomLevel);
+      renderProfiler();
       if (prevScroll) profiler.node().scrollTop = prevScroll;
     }
     const { y:baseY, height:baseHeight } = rect(div);
     const levelHeight = baseHeight-padding;
     const offsetY = baseY-canvasTop+padding/2;
-    const shapes = [];
-    data.tracks.set(k, { shapes, offsetY });
     let colorKey, ref;
     for (const e of timeline.shapes) {
       if (e.depth === 0) colorKey = e.cat ?? e.name;
@@ -209,15 +177,27 @@ async function renderProfiler() {
       }
       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 });
+      data.shapes.push({x:e.st-st, y:offsetY+levelHeight*e.depth, width:e.dur, height:levelHeight, arg, label, fillColor });
     }
     // 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;
-      data.tracks.set(`${k} memory`, { shapes:createPolygons(mem, area), offsetY:startY, div });
       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`;
@@ -241,63 +221,63 @@ async function renderProfiler() {
       yscale = d3.scaleLinear().domain(data.axes.y.domain).range(data.axes.y.range);
     }
     // draw shapes
-    for (const [_, { offsetY, shapes }] of data.tracks) {
-      for (const e of shapes) {
-        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} });
+    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)}`;
           }
-          continue;
+          rectLst.push({ x0:x[i], x1:x[i+1], y0:e.y1[i], y1:e.y0[i], arg:{...e.arg, tooltipText} });
         }
-        // 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;
+        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;
         }
+        ctx.fillStyle = l.color;
+        ctx.fillText(l.st, labelX, labelY);
+        labelWidth += l.width;
+        labelX += l.width;
       }
     }
     // draw axes
     drawLine(ctx, xscale.range(), [0, 0]);
-    for (const tick of xscale.ticks()) {
+    const ticks = xscale.ticks();
+    for (const [i, tick] of ticks.entries()) {
       // tick line
       const x = xscale(tick);
       drawLine(ctx, [x, x], [0, tickSize])
       // tick label
       ctx.textBaseline = "top";
-      ctx.textAlign = "left";
-      ctx.fillText(formatTime(tick, et-st), x+ctx.lineWidth+2, tickSize);
+      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);
     }
     if (yscale != null) {
       drawLine(ctx, [0, 0], yscale.range());

tinygrad/viz/serve.py

@@ -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.CHILDREN: "#80ffc0", Ops.CHILD: "#80fff0", Ops.BUFFERIZE: "#FF991C"}
+               Ops.CHILD: "#80fff0"}
 
 # VIZ API
 
@@ -73,15 +73,13 @@ 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}" + (f" {x.src[0].op}" if len(x.src) else "")
+        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}"
     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 LABEL>"
+      label += "\n<ISSUE GETTING SHAPE>"
     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)