realize sched

Merge branch 'master' into outerworld_work
clean up the LLM tokenizer (#12653 )
2025-10-14 14:44:56 +08:00 · 2025-10-14 14:26:52 +08:00 · 2025-10-14 14:22:01 +08:00 · 2025-10-14 13:05:26 +08:00 · 2025-10-14 12:35:20 +08:00 · 2025-10-14 10:49:08 +08:00
84 changed files with 1634 additions and 380 deletions
@@ -310,9 +310,9 @@ jobs:
    - name: Fuzz Test fast idiv
      run: python test/external/fuzz_fast_idiv.py
    - name: Fuzz Test shapetracker
-      run: |
-        python test/external/fuzz_shapetracker.py
-        python test/external/fuzz_shapetracker_math.py
+      run: CNT=50 python test/external/fuzz_shapetracker.py
+    - name: Fuzz Test shapetracker math
+      run: CNT=200 python test/external/fuzz_shapetracker_math.py
    - name: Fuzz Test shape ops
      run: python test/external/fuzz_shape_ops.py

@@ -377,7 +377,7 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=41 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+          ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2081 ALLOWED_GATED_READ_IMAGE=28 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot alt model correctness (float32)
        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot fastvits model correctness (float32)
@@ -30,10 +30,6 @@ persistent=yes
 # Specify a configuration file.
 #rcfile=

-# When enabled, pylint would attempt to guess common misconfiguration and emit
-# user-friendly hints instead of false-positive error messages
-suggestion-mode=yes
-
 # Allow loading of arbitrary C extensions. Extensions are imported into the
 # active Python interpreter and may run arbitrary code.
 unsafe-load-any-extension=no
@@ -435,8 +435,8 @@ generate_sqtt() {
    -o extra/sqtt/rocprof/rocprof.py
  fixup extra/sqtt/rocprof/rocprof.py
  sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py
-  sed -i "s/import ctypes/import ctypes\nfrom tinygrad.helpers import fetch/g" extra/sqtt/rocprof/rocprof.py
-  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(str(fetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so')))|g" extra/sqtt/rocprof/rocprof.py
+  sed -i "s/import ctypes/import ctypes, ctypes.util/g" extra/sqtt/rocprof/rocprof.py
+  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder'))|g" extra/sqtt/rocprof/rocprof.py
 }

 generate_webgpu() {
@@ -1188,7 +1188,9 @@ def train_bert():
      if MLLOGGER and RUNMLPERF:
        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
      if getenv("RESET_STEP"): train_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None:
+        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
+        train_step_bert.captured.free_intermediates()
      eval_lm_losses = []
      eval_clsf_losses = []
      eval_lm_accs = []
@@ -1222,7 +1224,7 @@ def train_bert():
          return

      if getenv("RESET_STEP"): eval_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 1) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()

      del eval_data
      avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
+
+export IGNORE_OOB=1
+
+export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+# export BEAM_LOG_SURPASS_MAX=1
+# export BASEDIR="/raid/datasets/wiki"
+
+export RESET_STEP=1
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,69 @@
+# 1. Problem
+
+This problem uses BERT for NLP.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
+```
+pip install gdown numpy tqdm tensorflow
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download and verify data
+
+### 1. Download raw data
+
+```
+BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
+```
+
+### 2. Preprocess train and validation data
+
+Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
+
+#### Training:
+```
+BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
+```
+
+Generating a specific topic (Between 0 and 499)
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
+```
+
+#### Validation:
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
+```
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+```
+
+### tinybox_red
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+```
+### tinybox_8xMI300X
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+```
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+export BENCHMARK=10 BERT_LAYERS=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,31 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_8xMI300X"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
+
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,69 @@
+# 1. Problem
+
+This problem uses BERT for NLP.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
+```
+pip install gdown numpy tqdm tensorflow
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download and verify data
+
+### 1. Download raw data
+
+```
+BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
+```
+
+### 2. Preprocess train and validation data
+
+Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
+
+#### Training:
+```
+BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
+```
+
+Generating a specific topic (Between 0 and 499)
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
+```
+
+#### Validation:
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
+```
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+```
+
+### tinybox_red
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+```
+### tinybox_8xMI300X
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+```
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BEAM_LOG_SURPASS_MAX=1
+export BASEDIR="/raid/datasets/wiki"
+
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,16 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,28 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." NV=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_green"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_green_${DATETIME}_${SEED}.log"
+
+# init
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,69 @@
+# 1. Problem
+
+This problem uses BERT for NLP.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
+```
+pip install gdown numpy tqdm tensorflow
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download and verify data
+
+### 1. Download raw data
+
+```
+BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
+```
+
+### 2. Preprocess train and validation data
+
+Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
+
+#### Training:
+```
+BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
+```
+
+Generating a specific topic (Between 0 and 499)
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
+```
+
+#### Validation:
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
+```
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+```
+
+### tinybox_red
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+```
+### tinybox_8xMI300X
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+```
@@ -0,0 +1,18 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BEAM_LOG_SURPASS_MAX=1
+export BASEDIR="/raid/datasets/wiki"
+
+export RESET_STEP=1
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,16 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,31 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_red"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_red_${DATETIME}_${SEED}.log"
+
+export HCQDEV_WAIT_TIMEOUT_MS=100000  # prevents hang?
+
+# init
+sleep 5 && sudo rmmod amdgpu || true
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,50 @@
+# 1. Problem
+
+This problem uses the ResNet-50 CNN to do image classification.
+
+## Requirements
+
+Install tinygrad and mlperf-logging from master.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+### tinybox_red
+Disable cwsr
+This is the default on production tinybox red.
+```
+sudo vi /etc/modprobe.d/amdgpu.conf
+cat <<EOF > /etc/modprobe.d/amdgpu.conf
+options amdgpu cwsr_enable=0
+EOF
+sudo update-initramfs -u
+sudo reboot
+
+# validate
+sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
+```
+
+# 2. Directions
+
+## Steps to download and verify data
+
+```
+IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
+```
+
+## Steps for one time setup
+
+### tinybox_red
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
+```
+
+## Steps to run benchmark
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
+```
@@ -0,0 +1,13 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
+
+export BENCHMARK=10 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
+
+export EVAL_START_EPOCH=3 EVAL_FREQ=4
+
+export WANDB=1 PARALLEL=0
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,25 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." NV=1
+export MODEL="resnet"
+export SUBMISSION_PLATFORM="tinybox_green"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=${LOGMLPERF:-1}
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
+
+# init
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,50 @@
+# 1. Problem
+
+This problem uses the ResNet-50 CNN to do image classification.
+
+## Requirements
+
+Install tinygrad and mlperf-logging from master.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+### tinybox_red
+Disable cwsr
+This is the default on production tinybox red.
+```
+sudo vi /etc/modprobe.d/amdgpu.conf
+cat <<EOF > /etc/modprobe.d/amdgpu.conf
+options amdgpu cwsr_enable=0
+EOF
+sudo update-initramfs -u
+sudo reboot
+
+# validate
+sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
+```
+
+# 2. Directions
+
+## Steps to download and verify data
+
+```
+IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
+```
+
+## Steps for one time setup
+
+### tinybox_red
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
+```
+
+## Steps to run benchmark
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
+```
@@ -0,0 +1,13 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=10 DEBUG=${DEBUG:-2}
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export EVAL_START_EPOCH=3 EVAL_FREQ=4
+
+export WANDB=1 PARALLEL=0
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,26 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." AMD=1
+export MODEL="resnet"
+export SUBMISSION_PLATFORM="tinybox_red"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=${LOGMLPERF:-1}
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
+
+# init
+sleep 5 && sudo rmmod amdgpu || true
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,8 @@
+#!/bin/bash
+
+rocm-smi --setprofile compute
+rocm-smi --setmclk 3
+rocm-smi --setperflevel high
+
+# power cap to 350W
+echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
@@ -0,0 +1,38 @@
+# 1. Problem
+
+This problem uses RetinaNet for SSD.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+
+Also install the following dependencies:
+```
+pip install tqdm numpy pycocotools boto3 pandas torch torchvision
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download data
+
+Run the following:
+```
+BASEDIR=/raid/datasets/openimages python3 extra/datasets/openimages.py
+```
+
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh
+```
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=5 DEBUG=2
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export WANDB=1 PARALLEL=0
+export RUNMLPERF=1
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,25 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export SUBMISSION_PLATFORM="tinybox_green"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export TRAIN_BEAM=2 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/openimages"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="retinanet_green_${DATETIME}_${SEED}.log"
+
+# init
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=5 DEBUG=2
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export WANDB=1 PARALLEL=0
+export RUNMLPERF=1
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,38 @@
+{
+    "submitter": "tinycorp",
+    "division": "closed",
+    "status": "Available on-premise",
+    "system_name": "tinybox 8xMI300X",
+    "number_of_nodes": "1",
+    "host_processors_per_node": "2",
+    "host_processor_model_name": "AMD EPYC 9354",
+    "host_processor_core_count": "32",
+    "host_processor_vcpu_count": "64",
+    "host_processor_frequency": "",
+    "host_processor_caches": "",
+    "host_processor_interconnect": "",
+    "host_memory_capacity": "2304GB",
+    "host_storage_type": "NVMe SSD",
+    "host_storage_capacity": "3x 4TB raid array",
+    "host_networking": "",
+    "host_networking_topology": "",
+    "host_memory_configuration": "24x 96GB DDR5",
+    "accelerators_per_node": "8",
+    "accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
+    "accelerator_host_interconnect": "PCIe 5.0 x16",
+    "accelerator_frequency": "",
+    "accelerator_on-chip_memories": "",
+    "accelerator_memory_configuration": "HBM3",
+    "accelerator_memory_capacity": "192GB",
+    "accelerator_interconnect": "",
+    "accelerator_interconnect_topology": "",
+    "cooling": "air",
+    "hw_notes": "",
+    "framework": "tinygrad, branch mlperf_training_v5.0",
+    "other_software_stack": {
+        "python": "3.10.16",
+        "ROCm": "3.0.0+94441cb"
+    },
+    "operating_system": "Ubuntu 24.04.1 LTS",
+    "sw_notes": ""
+  }
@@ -0,0 +1,38 @@
+{
+  "submitter": "tinycorp",
+  "division": "closed",
+  "status": "Available on-premise",
+  "system_name": "tinybox green",
+  "number_of_nodes": "1",
+  "host_processors_per_node": "1",
+  "host_processor_model_name": "AMD EPYC 7532",
+  "host_processor_core_count": "32",
+  "host_processor_vcpu_count": "64",
+  "host_processor_frequency": "",
+  "host_processor_caches": "",
+  "host_processor_interconnect": "",
+  "host_memory_capacity": "128GB",
+  "host_storage_type": "NVMe SSD",
+  "host_storage_capacity": "4 TB raid array + 1 TB boot",
+  "host_networking": "",
+  "host_networking_topology": "",
+  "host_memory_configuration": "8x 16GB DDR4",
+  "accelerators_per_node": "6",
+  "accelerator_model_name": "NVIDIA GeForce RTX 4090",
+  "accelerator_host_interconnect": "PCIe 4.0 x16",
+  "accelerator_frequency": "",
+  "accelerator_on-chip_memories": "",
+  "accelerator_memory_configuration": "GDDR6X",
+  "accelerator_memory_capacity": "24GB",
+  "accelerator_interconnect": "",
+  "accelerator_interconnect_topology": "",
+  "cooling": "air",
+  "hw_notes": "",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "other_software_stack": {
+    "python": "3.10.12",
+    "CUDA": "12.4"
+  },
+  "operating_system": "Ubuntu 22.04.4",
+  "sw_notes": ""
+}
@@ -0,0 +1,37 @@
+{
+  "submitter": "tinycorp",
+  "division": "closed",
+  "status": "Available on-premise",
+  "system_name": "tinybox red",
+  "number_of_nodes": "1",
+  "host_processors_per_node": "1",
+  "host_processor_model_name": "AMD EPYC 7532",
+  "host_processor_core_count": "32",
+  "host_processor_vcpu_count": "64",
+  "host_processor_frequency": "",
+  "host_processor_caches": "",
+  "host_processor_interconnect": "",
+  "host_memory_capacity": "128GB",
+  "host_storage_type": "NVMe SSD",
+  "host_storage_capacity": "4 TB raid array + 1 TB boot",
+  "host_networking": "",
+  "host_networking_topology": "",
+  "host_memory_configuration": "8x 16GB DDR4",
+  "accelerators_per_node": "6",
+  "accelerator_model_name": "AMD Radeon RX 7900 XTX",
+  "accelerator_host_interconnect": "PCIe 4.0 x16",
+  "accelerator_frequency": "",
+  "accelerator_on-chip_memories": "",
+  "accelerator_memory_configuration": "GDDR6",
+  "accelerator_memory_capacity": "24GB",
+  "accelerator_interconnect": "",
+  "accelerator_interconnect_topology": "",
+  "cooling": "air",
+  "hw_notes": "",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "other_software_stack": {
+    "python": "3.10.12"
+  },
+  "operating_system": "Ubuntu 22.04.4",
+  "sw_notes": ""
+}
@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+import os, shutil
+from pathlib import Path
+from tinygrad.helpers import fetch, OSX
+
+DEST = Path("/usr/local/lib")
+DEST.mkdir(exist_ok=True)
+
+if __name__ == "__main__":
+  if OSX:
+    fp = fetch("https://github.com/ROCm/rocprof-trace-decoder/releases/download/0.1.4/rocprof-trace-decoder-macos-arm64-0.1.4-Darwin.sh")
+    lib = fp.parent/"rocprof-trace-decoder-macos-arm64-0.1.4-Darwin"/"lib"/"librocprof-trace-decoder.dylib"
+    os.chmod(fp, 0o755)
+    os.system(f"sudo {fp} --prefix={fp.parent} --include-subdir")
+  else:
+    lib = fetch("https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so", name="librocprof-trace-decoder.so")
+  shutil.copy2(lib, DEST)
+  print(f"Installed {lib.name} to", DEST)
@@ -7,8 +7,7 @@
 # POINTER_SIZE is: 8
 # LONGDOUBLE_SIZE is: 16
 #
-import ctypes
-from tinygrad.helpers import fetch
+import ctypes, ctypes.util


 class AsDictMixin:
@@ -156,7 +155,7 @@ class FunctionFactoryStub:
 # You can either re-run clan2py with -l /path/to/library.so
 # Or manually fix this by comment the ctypes.CDLL loading
 _libraries = {}
-_libraries['FIXME_STUB'] = ctypes.CDLL(str(fetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so'))) #  ctypes.CDLL('FIXME_STUB')
+_libraries['FIXME_STUB'] = ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder')) #  ctypes.CDLL('FIXME_STUB')



@@ -50,4 +50,7 @@ exclude = [
  "E303", "E304", "E501", "E702", "E703", "E731", "W191",
  "W291", "W293", "UP039", "C416", "RET506", "RET507", "A",
  "FURB110", "RUF018", "F541", "F841"
-]
+]
+
+[format]
+exclude = ["*"]
@@ -1,41 +1,50 @@
+import functools, multiprocessing
 from transformers import AutoTokenizer
 from datasets import load_dataset
-from tinygrad.apps.llm import SimpleTokenizer, gpt2_decode_vocab, get_llama_re
+from tinygrad.apps.llm import SimpleTokenizer
 from tinygrad.helpers import tqdm, getenv, partition

+@functools.cache
+def get_tokenizers():
+  print("getting tokenizers")
+  base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
+  special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()), lambda e: e[1] in base_tokenizer.all_special_ids)
+  simple_tokenizer = SimpleTokenizer(dict(normal_tokens), dict(special_tokens))
+  return base_tokenizer, simple_tokenizer
+
+def test_tokenize(samp) -> bool:
+  base_tokenizer, simple_tokenizer = get_tokenizers()
+  idx, txt = samp
+  try: simple_tokens = tuple(simple_tokenizer.encode(txt))
+  except RuntimeError: simple_tokens = ()
+  base_tokens = tuple(base_tokenizer.encode(txt, add_special_tokens=False))
+  if simple_tokens != base_tokens:
+    print(f"tokens mismatch at index: {idx}.\n")
+    color_codes = [91, 92, 94, 93, 95]
+    def color_tokens(tids):
+      return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m"
+    print("simple:  ", color_tokens(simple_tokens))
+    print("official:", color_tokens(base_tokens) + "\n")
+    return False
+  if simple_tokenizer.decode(simple_tokens) != txt:
+    print(f"decode mismatch at {idx}")
+    return False
+  return True
+
 # use ALLOW_FAILED=-1 to go over the entire dataset without printing.
 if __name__ == "__main__":
-  base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
-  special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()),
-    lambda e: e[1] in base_tokenizer.all_special_ids)
-  inv_vocab = { tid: word for word, tid in base_tokenizer.get_vocab().items() }
-  simple_tokenizer = SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens))
-
-  color_codes = [ 91, 92, 94, 93, 95 ]
-  def color_tokens(tids):
-    return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m"
-
+  print("loading datasets")
  ds = load_dataset("OpenAssistant/oasst1")
+  loaded_ds = [(idx, el["text"]) for idx, el in enumerate(ds["train"])]
+  print(f"loaded {len(loaded_ds)}")
+
  allow_failed = getenv("ALLOW_FAILED", 10)
-
  fail_count, total = 0, 0
-
-  for idx, el in enumerate(tqdm(ds["train"])):
-    total += 1
-
-    try: simple_tokens = tuple(simple_tokenizer.encode(el["text"]))
-    except RuntimeError: simple_tokens = ()
-    base_tokens = tuple(base_tokenizer.encode(el["text"], add_special_tokens=False))
-
-    if simple_tokens != base_tokens:
-      fail_count += 1
-      allow_failed -= 1
-
-      if allow_failed >= 0:
-        print(f"tokens mismatch at index: {idx}.\n")
-
-        print("simple:  ", color_tokens(simple_tokens))
-        print("official:", color_tokens(base_tokens) + "\n")
-
-      if allow_failed == 0: break
-  print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.")
+  with multiprocessing.Pool(16) as pool:
+    for good in tqdm(pool.imap_unordered(test_tokenize, loaded_ds), total=len(loaded_ds)):
+      total += 1
+      if not good:
+        fail_count += 1
+        allow_failed -= 1
+        if allow_failed == 0: break
+    print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.")
@@ -42,13 +42,13 @@ class ProcessReplayWarning(Warning): pass

 # *** replay the function and convert return values to string

-def replay_kernelize(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[Any, ...]]:
+def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[Any, ...]]:
  UOp.unique_num = itertools.count(max([u.arg for u in big_sink.toposort() if u.op is Ops.UNIQUE], default=0)+1)
  new_sink = big_sink.substitute(get_rangeify_map(big_sink))
  def to_str(ret:UOp) -> str:
    asts = [repr(u.arg.ast) for u in ret.toposort() if u.op is Ops.KERNEL]
    return "\n".join([f"{len(asts)} kernels", *asts])
-  return to_str(new_sink), to_str(ret[big_sink]), (big_sink,)
+  return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)

 def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
  # NOTE: this always uses the opts_to_apply path
@@ -65,7 +65,7 @@ def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts
  ast_repr = codecs.decode(str(input_ast), "unicode_escape")
  return to_str(p2), to_str(p), (ast_repr, renderer)

-replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {"get_kernelize_map":replay_kernelize, "get_program":replay_get_program}
+replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {"get_rangeify_map":replay_get_rangeify_map, "get_program":replay_get_program}

 # *** run replayers on captured rows and print diffs

@@ -8,9 +8,11 @@ from tinygrad.uop.ops import UOp, Ops, GroupOp
 from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
-from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT
+from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, getenv
 from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.cstyle import CUDARenderer
+MOCKGPU = getenv("MOCKGPU")

 class TestLinearizer(unittest.TestCase):
  def test_arg_dedup(self):
@@ -68,7 +70,8 @@ class TestLinearizer(unittest.TestCase):
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
    assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
-    assert not any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]])
+    # the index of the load doesnt depend on the second range
+    assert any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]])
    assert any(x.op in {*GroupOp.ALU, Ops.LOAD} for x in uops[ranges[1]:])

  def test_range_outer_op_before_phi(self):
@@ -314,7 +317,7 @@ class TestLinearizer(unittest.TestCase):
    a.realize()
    np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexes differently. might be ok?")
+  @unittest.skipIf(MOCKGPU and isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CUDARenderer)), "PTX indexes differently. might be ok?")
  def test_where_fold(self):
    a = Tensor.ones(4, 4).contiguous().realize()
    b = a.shrink(((1, 2), None)).pad(((1, 2), None))
@@ -390,7 +390,6 @@ class TestMultiTensor(unittest.TestCase):

  # NOTE: this is failing on LLVM CI, no idea why. Works locally.
  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "CPU", "AMD"), "slow, and flaky on CPU")
-  @unittest.skip("TODO: pm_rangeify hangs")
  def test_data_parallel_resnet(self):
    from extra.models.resnet import ResNet18

@@ -1,7 +1,63 @@
 import unittest
-from tinygrad import Tensor, UOp
+from tinygrad import Tensor, UOp, Variable, nn
 from tinygrad.uop.ops import AxisType, Ops

+class TestOuterworldTrain(unittest.TestCase):
+  @Tensor.train()
+  def test_train(self):
+    # same example over and over
+    X = Tensor.rand(1, 32).expand(16,32).contiguous()
+    Y = Tensor.rand(1, 1).expand(16,1).contiguous()
+
+    layer = nn.Linear(32, 1, bias=False)
+    opt = nn.optim.SGD(nn.state.get_parameters(layer))
+    Tensor.realize(X, Y, *nn.state.get_parameters(layer))
+
+    print("train")
+
+    # if everything is correct, this should be a 16 step training loop
+    steps = UOp.range(16, -1)
+    opt.zero_grad()
+    loss = (layer(X[steps]) - Y[steps]).square().mean().backward()
+    sched = opt.schedule_step() # TODO: does this need to know anything about steps?
+    # NOTE: this can't work. the inputs to layer are not the assign, need to run twice for the fixed point?
+    all_losses = Tensor.realize(loss.reshape(1).expand(steps).contiguous(), *sched)
+    print(all_losses.numpy())
+
+#@unittest.skip("TODO: understand assign")
+class TestOuterworldAssign(unittest.TestCase):
+  def test_triple_add_inner(self):
+    t = Tensor.zeros(5).contiguous().realize()
+    t2 = Tensor.ones(3).contiguous().realize()
+    a = UOp.range(3, -1)
+    t = t.reshape(1,5).expand(a+1,5)[a].assign(t+t2[a])
+    self.assertListEqual(t.tolist(), [3,3,3,3,3])
+
+  def test_triple_add_outer(self):
+    t = Tensor.zeros(5).contiguous().realize()
+    t2 = Tensor.ones(3).contiguous().realize()
+
+    # OUTER is a loop at the schedule level
+    a = UOp.range(3, -1, AxisType.OUTER)
+    va = Variable("loop", 0, 2).bind(a)
+    t = t.assign(t+t2[va])
+    t = Tensor(UOp(Ops.ENDRANGE, dtype=t.uop.dtype, src=(a, t.uop)))
+
+    self.assertListEqual(t.tolist(), [3,3,3,3,3])
+
+  def test_triple_gemm(self):
+    x = Tensor.rand(1, 16).realize()
+    W = Tensor.rand(3, 16, 16).realize()
+
+    #manual = (x @ W[0] @ W[1] @ W[2]).contiguous().realize()
+
+    a = UOp.range(3, -1)
+
+    out = (x @ W[a]).contiguous()
+    t = Tensor(UOp(Ops.ASSIGN, dtype=out.uop.dtype, src=(x.uop, out.uop, a)))
+    #t = Tensor(UOp(Ops.REDUCE, dtype=out.uop.dtype, src=(out.uop, x.uop, a), arg=Ops.NOOP))
+    t.realize()
+
 class TestOuterworldReduce(unittest.TestCase):
  def test_reduce(self):
    x = Tensor.ones(5, 5).contiguous()
@@ -40,6 +96,7 @@ class TestOuterworld(unittest.TestCase):
    # passthrough ranges
    a = UOp.range(10, -1)
    sel = t[9-a]
+    assert sel.shape == (10,)
    cpy = sel.reshape(1, 10).expand(a, 10).contiguous().realize()

    self.assertTrue((t.flip(0)==cpy).all().item())
@@ -1526,7 +1526,7 @@ class TestSchedule(unittest.TestCase):
    # run_schedule(check_schedule(out, 1))
    run_schedule(check_schedule(out, 4))
    np.testing.assert_allclose(out.numpy(), np.pad(np.log2(np.abs(np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum() + \
-                                                   b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-6)
+                                                   b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-5)

  def test_shrink_pad_safe(self):
    a = Tensor.ones((3, )).contiguous().realize()
@@ -1,19 +1,21 @@
 import unittest, base64, functools, sys
-from tinygrad.apps.llm import SimpleTokenizer, get_llama_re
+from tinygrad.apps.llm import SimpleTokenizer
 from tinygrad.helpers import fetch

@unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
 class TestLLMTokenizer(unittest.TestCase):
-  @functools.cached_property
-  def basic_tok(self): return SimpleTokenizer(".*", { b"a": 0, b"b": 1, b"c": 2, b"ab": 3, b"bc": 4 }, { "<x>": 5, "<y>": 6, "<z>": 7 })
-
  @functools.cached_property
  def llama_tok(self):
    # from https://github.com/tinygrad/tinygrad/blob/e0106b6b257ebc003eb3694144e3e198f7d8cc37/examples/llama3.py#L14
    model_file = fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model")
    with open(model_file, "rt") as fd:
-      str_vocab = [ line.split(maxsplit=1) for line in fd.read().splitlines() if line ]
-      normal_tokens = { base64.b64decode(stok): int(srank) for stok, srank in str_vocab }
+      str_vocab = [line.split(maxsplit=1) for line in fd.read().splitlines() if line]
+
+      # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
+      bs = [*range(33, 127), *range(161, 173), *range(174, 256)]  # bytes that map to themselves
+      _byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
+      _byte_encoder = {v:k for k,v in _byte_decoder.items()}
+      normal_tokens = {''.join([_byte_encoder[x] for x in base64.b64decode(stok)]): int(srank) for stok, srank in str_vocab}

    special_tokens = [
      "<|begin_of_text|>",
@@ -27,22 +29,12 @@ class TestLLMTokenizer(unittest.TestCase):
      "<|reserved_special_token_4|>",
      "<|eot_id|>",
    ] + [ f"<|reserved_special_token_{i}|>" for i in range(5, 256 - 5) ]
-    return SimpleTokenizer(get_llama_re(), normal_tokens, { token: len(normal_tokens) + i for i, token in enumerate(special_tokens) })
+    return SimpleTokenizer(normal_tokens, {token: len(normal_tokens) + i for i, token in enumerate(special_tokens)})

  def _test_coding(self, tok: SimpleTokenizer, text: str, expected_tokens: list[int]):
    self.assertEqual(tok.encode(text), expected_tokens)
    self.assertEqual(tok.decode(expected_tokens), text)

-  def test_abc(self): self._test_coding(self.basic_tok, "abc", [ 3, 2 ])
-  def test_abbc(self): self._test_coding(self.basic_tok, "abbc", [ 3, 4 ])
-  def test_aabbbcc(self): self._test_coding(self.basic_tok, "aabbbcc", [ 0, 3, 1, 4, 2 ])
-  def test_specials1(self): self._test_coding(self.basic_tok, "a<x>a<y>a<z>a", [ 0, 5, 0, 6, 0, 7, 0 ])
-  def test_specials2(self): self._test_coding(self.basic_tok, "<x>a<y>a<z>", [ 5, 0, 6, 0, 7 ])
-  def test_invalid_token(self):
-    with self.assertRaises(RuntimeError): self._test_coding(self.basic_tok, "L", [])
-
-  def test_no_specials(self): self._test_coding(SimpleTokenizer(".*", { bytes([i]): i for i in range(256) }, {}), "abc", [97, 98, 99])
-
  # NOTE: the correct tokenization for this can only be found by looking up the text chunk in the vocab, not by applying merges
  def test_llama_early_tokenize(self): self._test_coding(self.llama_tok, " например", [ 111797 ])

@@ -5,6 +5,7 @@ from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.uop.symbolic import simplify_valid
 from tinygrad.helpers import Context
+from .test_uop_symbolic import check_uop_against_string

 def get_gated_load_uop(valid:UOp, idx:UOp):
  return UOp(Ops.LOAD, dtypes.float, (
@@ -49,8 +50,8 @@ class TestValidIdxSimplification(unittest.TestCase):
    with Context(NOOPT=1):
      load = full_rewrite_to_sink(load.sink()).src[0]
    idx, valid = load.src[0].src[1], load.src[0].src[2]
-    self.assertEqual(idx.render(simplify=False), sidx)
-    self.assertEqual(valid.render(simplify=False), svalid)
+    check_uop_against_string(self, idx, sidx)
+    check_uop_against_string(self, valid, svalid)

  def test_cumsum(self):
    gidx0 = Special("gidx0", 5)
@@ -186,13 +187,13 @@ class TestValidIdxSimplification(unittest.TestCase):
    print("The expressions are not equivalent.")
    print(s.model())

-  @unittest.expectedFailure  # TODO: improve uop_given_valid
  def test_valid_becomes_const2(self):
    ridx0 = Range(0, 4)
    ridx1 = Range(1, 4)
    ridx2 = Range(2, 4)
    ridx3 = Range(3, 4)
-    idx= ((ridx0+ridx1+ridx2+ridx3+28)//30)
+    # TODO: this should also work without the extra nesting
+    idx = (((ridx0+ridx1)+(ridx2+ridx3)+28)//30)
    valid = ((ridx0+ridx1)<1).ne(True) & ((ridx2+ridx3)<1).ne(True)
    load = get_gated_load_uop(valid, idx)
    self.check(load,
@@ -218,10 +219,10 @@ class TestImageSimplification(unittest.TestCase):
    self.assertEqual(idx.op, Ops.VECTORIZE)
    self.assertEqual(len(idx.src), 2)
    idx0, idx1 = idx.src[0], idx.src[1]
-    self.assertEqual(idx0.render(simplify=False), sidx0)
-    self.assertEqual(idx1.render(simplify=False), sidx1)
+    check_uop_against_string(self, idx0, sidx0)
+    check_uop_against_string(self, idx1, sidx1)
    if svalid is not None:
-      self.assertEqual(load.src[0].src[2].render(simplify=False), svalid)
+      check_uop_against_string(self, load.src[0].src[2], svalid)
    else:
      self.assertEqual(len(load.src[0].src), 2, "svalid is None but load still has a valid")

@@ -0,0 +1,22 @@
+import unittest
+from tinygrad import Tensor
+
+class TestLoadStore(unittest.TestCase):
+  def test_load_shape(self):
+    t = Tensor(bytes(16)).load(1024).kernelize()
+    assert t.shape == (1024,), t.shape
+
+  def test_store_shape(self):
+    t = Tensor.zeros(1024).store().kernelize()
+    assert t.shape == (16,), t.shape
+
+  def test_load_large_shape(self):
+    t = Tensor(bytes(16)).load(10_000_000).kernelize()
+    assert t.shape == (10_000_000,), t.shape
+
+  def test_store_large_shape(self):
+    t = Tensor.zeros(10_000_000).store().kernelize()
+    assert t.shape == (16,), t.shape
+
+if __name__ == "__main__":
+  unittest.main()
@@ -5,14 +5,17 @@ import z3
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from tinygrad.codegen import full_rewrite
 from tinygrad.helpers import Context
-from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer, track_rewrites
-from tinygrad.uop.symbolic import sym
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
+from tinygrad.uop.symbolic import sym, commutative
 from tinygrad.uop.spec import uops_to_z3

-@track_rewrites(name="simplify symbolic uop")
-def render(v) -> UOp:
-  v_simplified = graph_rewrite(v, sym)
-  return v_simplified
+def check_uop_against_string(self, v:UOp, s:str):
+  sym_vars = {v.render():v for v in v.toposort() if v.op in (Ops.DEFINE_VAR, Ops.RANGE, Ops.SPECIAL)}
+  s_eval = eval(s, sym_vars)
+  if isinstance(s_eval, int) and v.dtype==dtypes.index: s_eval = UOp.const(dtypes.index, s_eval)
+  elif isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(dtypes.from_py(s_eval), s_eval)
+  s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval")
+  self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v} for {s}")

 def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.index): return UOp.variable(name,min_val,max_val,dtype)
 def uconst(val): return UOp.const(dtypes.index, val)
@@ -33,11 +36,11 @@ class TestSymbolic(unittest.TestCase):
    self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original")

  def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    v_simplified = render(v)
+    v_simplified = graph_rewrite(v, sym, name="simplify symbolic uop")
    if test_z3: self.check_equal_z3(v, v_simplified)
-    rendered, nmin, nmax = v_simplified.render(simplify=False), v_simplified.vmin, v_simplified.vmax
-    if isinstance(s, tuple): self.assertIn(rendered, s)
-    else: self.assertEqual(rendered, s)
+    nmin, nmax = v_simplified.vmin, v_simplified.vmax
+    check_uop_against_string(self, v_simplified, s)
+    # eval the test string and see if we get the same uop
    self.assertEqual(nmin, n)
    self.assertEqual(nmax, m)

@@ -76,7 +79,7 @@ class TestSymbolic(unittest.TestCase):

  def test_lt_factors(self):
    expr = (Variable("idx1", 0, 511)*4 + Variable("FLOAT4_INDEX", 0, 256)) < 512
-    self.helper_test_variable(expr, 0, 1, ("(((idx1*4)+FLOAT4_INDEX)<512)", "((FLOAT4_INDEX+(idx1*4))<512)"))
+    self.helper_test_variable(expr, 0, 1, "(((idx1*4)+FLOAT4_INDEX)<512)")

  def test_div_reduction(self):
    self.helper_test_variable(Variable("a", 2, 3)//2, 1, 1, "1")
@@ -187,7 +190,7 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(Variable("a", 0, 8)%1, 0, 0, "0")

  def test_max_folds(self):
-    self.helper_test_variable(Variable("a", 0, 20).maximum(10).maximum(11), 11, 20, "max(a, 11)")
+    self.helper_test_variable(Variable("a", 0, 20).maximum(10).maximum(11), 11, 20, "a.maximum(11)")

  def test_add_min_max(self):
    self.helper_test_variable(Variable("a", 0, 8) * 2 + 12, 12, 16+12, "((a*2)+12)")
@@ -216,7 +219,7 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(usum([Variable("a", 0, 7)*4, Variable("b", 0, 3)*4]) % 2, 0, 0, "0")

  def test_sum_div_some_factor(self):
-    self.helper_test_variable(usum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*4]) // 2, 0, 23, ("(((a*5)//2)+(b*2))", "((b*2)+((a*5)//2))"))
+    self.helper_test_variable(usum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*4]) // 2, 0, 23, "(((a*5)//2)+(b*2))")

  def test_sum_div_trim_const(self):
    self.helper_test_variable((Variable("a", 0, 7)*4 + Variable("b", 0, 3)*4 + 7) // 16, 0, 2, "(((a+b)+1)//4)")
@@ -279,7 +282,7 @@ class TestSymbolic(unittest.TestCase):
  def test_mod_congruence_multiple_vars(self):
    self.helper_test_variable((9+9*Variable("x",0,3)+9*Variable("y",0,3))%10, 3, 9, "(((x*-1)+(y*-1))+9)")
    self.helper_test_variable((7+9*Variable("x",0,2)+9*Variable("y",0,2)+Variable("z",0,2))%10, 3, 9,
-                              ("(((z+(x*-1))+(y*-1))+7)", "(((y*-1)+(z+(x*-1)))+7)"))
+                              "(((z+(x*-1))+(y*-1))+7)")
    self.helper_test_variable((10+12*Variable("x",0,2)+Variable("y", 0, 4)%3)%13, 8, 12, "(((x*-1)+(y%3))+10)")

  def test_div_congruence(self):
@@ -301,8 +304,7 @@ class TestSymbolic(unittest.TestCase):

  def test_sum_lt_fold(self):
    self.helper_test_variable(usum([Variable("a", 0, 7) * 4, Variable("b", 0, 3)]) < 16, 0, 1, "(a<4)")
-    self.helper_test_variable(usum([Variable("a", 0, 7) * 4, Variable("b", 0, 4)]) < 16, 0, 1,
-                              ("(((a*4)+b)<16)", "((b+(a*4))<16)"))
+    self.helper_test_variable(usum([Variable("a", 0, 7) * 4, Variable("b", 0, 4)]) < 16, 0, 1, "(((a*4)+b)<16)")
    self.helper_test_variable(usum([Variable("uidx", 0, 3), Variable("a", 0, 1529) * 12]) < (4 * 67), 0, 1, "(a<23)")

  def test_mul_mod_large(self):
@@ -364,7 +366,7 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable((1+Variable("a", 0, 3))*(-2)+12, 4, 10, "((a*-2)+10)")

  def test_mod_mul_sum(self):
-    self.helper_test_variable(usum([Variable("b", 0, 2), Variable("a", 0, 5)*10])%9, 0, 7, ("(b+a)", "(a+b)"))
+    self.helper_test_variable(usum([Variable("b", 0, 2), Variable("a", 0, 5)*10])%9, 0, 7, "(b+a)")

  def test_sum_0(self):
    self.helper_test_variable(usum([Variable("a", 0, 7)]), 0, 7, "a")
@@ -395,11 +397,11 @@ class TestSymbolic(unittest.TestCase):

  def test_lt_sum_factor_rhs_partial(self):
    self.helper_test_variable((Variable("a", 0, 6)*6 + Variable("b", 0, 6)*4 + Variable("c", 0, 6)*8) < 4, 0, 1,
-                              ("((((a*3)+(b*2))+(c*4))<2)", "(((b*2)+((a*3)+(c*4)))<2)"))
+                              "((((a*3)+(b*2))+(c*4))<2)")

  def test_lt_sum_factor_rhs_all(self):
    self.helper_test_variable((Variable("a", 0, 6)*6 + Variable("b", 0, 6)*4 + Variable("c", 0, 6)*8) < 2, 0, 1,
-                              ("((((a*3)+(b*2))+(c*4))<1)", "(((b*2)+((a*3)+(c*4)))<1)"))
+                              "((((a*3)+(b*2))+(c*4))<1)")

  def test_and_fold(self):
    self.helper_test_variable(uand([uconst(0), Variable("a", 0, 1)]), 0, 0, "0")
@@ -561,38 +563,35 @@ class TestSymbolic(unittest.TestCase):
    lidx2 = Variable("lidx2", 0, 3)
    alu0 = gidx2*640+gidx1*160+(gidx0//5)*2+lidx0*320+lidx1*10
    self.helper_test_variable((alu0+lidx2*2+1)//20, 0, 8192,
-                              ("((((((gidx0//5)+lidx2)//5)+lidx1)//2)+(((gidx2*32)+(gidx1*8))+(lidx0*16)))",
-                               "(((lidx1+((lidx2+(gidx0//5))//5))//2)+((gidx2*32)+((gidx1*8)+(lidx0*16))))",
-                               "((((gidx1*8)+(gidx2*32))+(lidx0*16))+((lidx1+((lidx2+(gidx0//5))//5))//2))"))
+                              "((((((gidx0//5)+lidx2)//5)+lidx1)//2)+(((gidx2*32)+(gidx1*8))+(lidx0*16)))")

  def test_sum_div_complex2(self):
    gidx0 = Variable("gidx0", 0, 7)
    lidx2 = Variable("lidx2", 0, 1)
    lidx3 = Variable("lidx3", 0, 1)
-    self.helper_test_variable((gidx0*4+lidx2*2+1)//10, 0, 3, ("(((gidx0*2)+lidx2)//5)", "((lidx2+(gidx0*2))//5)"))
-    self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//10, 0, 3, ("(((gidx0*2)+lidx2)//5)", "((lidx2+(gidx0*2))//5)"))
+    self.helper_test_variable((gidx0*4+lidx2*2+1)//10, 0, 3, "(((gidx0*2)+lidx2)//5)")
+    self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//10, 0, 3, "(((gidx0*2)+lidx2)//5)")
    self.helper_test_variable((gidx0*2+lidx2)//10, 0, 1, "(gidx0//5)")

  def test_sum_div_complex3(self):
    gidx0 = Variable("gidx0", 0, 7)
    lidx2 = Variable("lidx2", 0, 12)
    lidx3 = Variable("lidx3", 0, 1)
-    self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//12, 0, 4, ("(((lidx2//2)+gidx0)//3)", "((gidx0+(lidx2//2))//3)"))
-    self.helper_test_variable((lidx2*2+gidx0*4+lidx3)//12, 0, 4, ("(((lidx2//2)+gidx0)//3)", "((gidx0+(lidx2//2))//3)"))
+    self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//12, 0, 4, "(((lidx2//2)+gidx0)//3)")
+    self.helper_test_variable((lidx2*2+gidx0*4+lidx3)//12, 0, 4, "(((lidx2//2)+gidx0)//3)")

  @unittest.expectedFailure  # TODO: improve nest_div_by_smallest_factor
  def test_sum_div_complex4(self):
    gidx0 = Variable("gidx0", 0, 2)
    lidx2 = Variable("lidx2", 0, 12)
    lidx3 = Variable("lidx3", 0, 12)
-    self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, ("((lidx2+(lidx3*2))//3)"))
+    self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, "((lidx2+(lidx3*2))//3)")

  def test_sum_mul_distribute(self):
    gidx0 = Variable("gidx0", 0, 7)
    lidx2 = Variable("lidx2", 0, 12)
    lidx3 = Variable("lidx3", 0, 1)
-    self.helper_test_variable((gidx0+lidx2+lidx3)*4, 0, 80,
-                              ("(((gidx0*4)+(lidx2*4))+(lidx3*4))","((lidx3*4)+((gidx0*4)+(lidx2*4)))"))
+    self.helper_test_variable((gidx0+lidx2+lidx3)*4, 0, 80, "(((gidx0*4)+(lidx2*4))+(lidx3*4))")

  @unittest.expectedFailure
  def test_variable_divmod(self):
@@ -662,7 +661,7 @@ class TestSymbolic(unittest.TestCase):
    idx = Variable("idx", 0, 24)
    self.helper_test_variable(idx//4, 0, 6, "(idx//4)")
    # TODO: simplify the true branch
-    self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "((idx//4) if (idx<4) else -1)")
+    self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "(idx<4).where((idx//4), -1)")

  def test_idiv_lt(self):
    idx = Variable("idx", 0, 24)
@@ -681,8 +680,8 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable((a*3+b*4<1).ne(True), 0, 1, "(((a+b)<1)!=True)")
    self.helper_test_variable((a*(-3)+b*4<1).ne(True), 0, 1, "((((a*-3)+(b*4))<1)!=True)")  # negative coeff, should not be simplified
    self.helper_test_variable((a*3+d*4<1).ne(True), 0, 1, "((((a*3)+(d*4))<1)!=True)")  # var can be negative, should not be simplified
-    self.helper_test_variable((a+b+c*2<1).ne(True), 0, 1, ("((((a+b)+c)<1)!=True)", "(((c+(a+b))<1)!=True)", '(((b+(a+c))<1)!=True)'))
-    self.helper_test_variable((a+b*2+c*4<1).ne(True), 0, 1, ("((((a+b)+c)<1)!=True)", "(((c+(a+b))<1)!=True)", '(((b+(a+c))<1)!=True)'))
+    self.helper_test_variable((a+b+c*2<1).ne(True), 0, 1, "((((a+b)+c)<1)!=True)")
+    self.helper_test_variable((a+b*2+c*4<1).ne(True), 0, 1, "((((a+b)+c)<1)!=True)")

  def test_where_removal(self):
    cond = Variable("a", 0, 3) < 2
@@ -700,30 +699,30 @@ class TestSymbolic(unittest.TestCase):
    c = Variable("c", 0, 3)
    aa = cond.where(a, a.ufix(0))
    bb = cond.where(b, b.ufix(1))
-    self.helper_test_variable(aa, 0, 3, "(a if (x<2) else 0)")
-    self.helper_test_variable(bb, 0, 3, "(b if (x<2) else 1)")
-    self.helper_test_variable(aa+bb, 0, 6, "((a+b) if (x<2) else 1)")
-    self.helper_test_variable(aa.maximum(bb), 0, 3, "(max(a, b) if (x<2) else 1)")
-    self.helper_test_variable((c+aa)+bb, 0, 9, "(c+((a+b) if (x<2) else 1))")
+    self.helper_test_variable(aa, 0, 3, "(x<2).where(a, 0)")
+    self.helper_test_variable(bb, 0, 3, "(x<2).where(b, 1)")
+    self.helper_test_variable(aa+bb, 0, 6, "(x<2).where((a+b), 1)")
+    self.helper_test_variable(aa.maximum(bb), 0, 3, "(x<2).where(a.maximum(b), 1)")
+    self.helper_test_variable((c+aa)+bb, 0, 9, "(c+(x<2).where((a+b), 1))")

    # not combining because it increased total ALU
    cc = cond.where(c, c+1)
-    self.helper_test_variable(bb+cc, 0, 7, "((b if (x<2) else 1)+(c if (x<2) else (c+1)))")
+    self.helper_test_variable(bb+cc, 0, 7, "((x<2).where(b, 1)+(x<2).where(c, (c+1)))")

    # not combining  # TODO: can combine if it can further simplify?
    ab = cond.where(a, b)
    ba = cond.where(b, a)
-    self.helper_test_variable(ab+ba, 0, 6, "((a if (x<2) else b)+(b if (x<2) else a))")
+    self.helper_test_variable(ab+ba, 0, 6, "((x<2).where(a, b)+(x<2).where(b, a))")

    # not combining  # TODO: can combine if one is identity element const
-    self.helper_test_variable(aa+ab, 0, 6, "((a if (x<2) else b)+(a if (x<2) else 0))")
+    self.helper_test_variable(aa+ab, 0, 6, "((x<2).where(a, b)+(x<2).where(a, 0))")

  def test_negation_in_where(self):
    cond = Variable("x", 0, 3) < 2
    a = Variable("a", 0, 3)
    b = Variable("b", 0, 3)
    w = cond.logical_not().where(a, b)
-    self.helper_test_variable(w, 0, 3, "(b if (x<2) else a)")
+    self.helper_test_variable(w, 0, 3, "(x<2).where(b, a)")

  def test_neg_in_comp(self):
    a = Variable("a", 0, 3)
@@ -750,7 +749,7 @@ class TestSymbolic(unittest.TestCase):
    a = Variable("a", 0, 3)
    b = Variable("b", 0, 3)
    expr = cond1.where(cond2.where(a, b), b)
-    self.helper_test_variable(expr, 0, 3, "(a if ((s<6)&(2<s)) else b)")
+    self.helper_test_variable(expr, 0, 3, "((s<6)&(2<s)).where(a, b)")

  def test_where_merge_branches2(self):
    cond1 = Variable("s", 0, 10) < 5
@@ -759,7 +758,7 @@ class TestSymbolic(unittest.TestCase):
    b = Variable("b", 0, 3)
    expr = cond1.where(cond2.where(a, b), b)
    # (a if ((s<5)&(s<6)) else b) -> (a if (s<5) else b)
-    self.helper_test_variable(expr, 0, 3, "(a if (s<5) else b)")
+    self.helper_test_variable(expr, 0, 3, "(s<5).where(a, b)")

  def test_symbolic_div(self):
    # from symbolic arange
@@ -774,7 +773,7 @@ class TestSymbolic(unittest.TestCase):
    a = Variable("a", 1, 10, dtypes.float)
    # TODO: bounds for reciprocal
    # TODO: should z3 work?
-    self.helper_test_variable(2*(2*a).reciprocal(), -math.inf, math.inf, "(1/a)", test_z3=False)
+    self.helper_test_variable(2*(2*a).reciprocal(), -math.inf, math.inf, "a.reciprocal()", test_z3=False)

  def test_trunc_noop(self):
    a = Variable("a", 1, 10, dtypes.int)
@@ -783,8 +782,8 @@ class TestSymbolic(unittest.TestCase):
  def test_do_math_in_int32(self):
    a = Variable("a", 1, 10, dtypes.int)
    b = Variable("b", 1, 10, dtypes.int)
-    self.helper_test_variable(a.cast(dtypes.long)+b.cast(dtypes.long), 2, 20, "(long)((a+b))")
-    self.helper_test_variable(a.cast(dtypes.long)*b.cast(dtypes.long), 1, 100, "(long)((a*b))")
+    self.assertIn((a.cast(dtypes.long)+b.cast(dtypes.long)).render(), "(long)((a+b))")
+    self.assertIn((a.cast(dtypes.long)*b.cast(dtypes.long)).render(), "(long)((a*b))")

 class TestSymbolicNumeric(unittest.TestCase):
  def helper_test_numeric(self, f):
@@ -30,6 +30,7 @@ class BaseTestViz(unittest.TestCase):
    # clear the global context
    for lst in [tracked_keys, tracked_ctxs, active_rewrites, _name_cnt]: lst.clear()
    Buffer.profile_events.clear()
+    cpu_events.clear()
    self.tms = TRACK_MATCH_STATS.value
    self.profile = PROFILE.value
    TRACK_MATCH_STATS.value = 2
@@ -462,5 +463,21 @@ class TestVizMemoryLayout(BaseTestViz):
    self.assertEqual(ret["peak"], 2)
    self.assertEqual(len(ret["events"]), 4)

+  def test_free_last(self):
+    bufs = []
+    for _ in range(3):
+      bufs.append(_alloc(1))
+      profile_marker("alloc")
+    device = bufs[0].device
+    while bufs:
+      b = bufs.pop()
+      del b
+      profile_marker("free")
+    profile = load_profile(cpu_events+Buffer.profile_events)
+    ret = profile["layout"][f"{device} Memory"]
+    self.assertEqual(ret["peak"], 3)
+    self.assertEqual(len(ret["events"]), 6)
+    self.assertEqual(len(profile["markers"]), 6)
+
 if __name__ == "__main__":
  unittest.main()
@@ -42,7 +42,7 @@ class TestWinograd(unittest.TestCase):
    out = Tensor.conv2d(x,w, padding=1)
    out.mean().backward()
    backward_schedule = Tensor.schedule(x.grad, w.grad)
-    self.assertEqual(len(backward_schedule), 4)
+    self.assertEqual(len(backward_schedule), 5)

  def test_counters(self):
    IC, OC, X, Y = 4,4,9,9
@@ -1,63 +1,61 @@
 from __future__ import annotations
-import sys, argparse, typing, re, itertools, unicodedata
+import sys, argparse, typing, re, unicodedata
 from tinygrad import Tensor, nn, UOp, TinyJit, getenv, helpers

-def gpt2_decode_vocab(voc: dict[str, int]): # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
-  c2b = { chr(cp): cp for cp in itertools.chain(range(ord("!"), ord("~")+1), range(ord("¡"), ord("¬")+1), range(ord("®"), ord("ÿ")+1)) }
-  c2b.update({ chr(256+off): cp for off, cp in enumerate(cp for cp in range(256) if chr(cp) not in c2b) })
-  return { bytes(c2b[c] for c in tok): tid for tok, tid in voc.items() }
-
-def get_llama_re():
-  def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre))
-  r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
-  # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
-  return "(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
-    f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+"
-
 class SimpleTokenizer:
-  def __init__(self, pat: str, normal_tokens: dict[bytes, int], special_tokens: dict[str, int]):
-    self._normal_tokens, self._special_tokens, self._pat = normal_tokens, special_tokens, re.compile(pat)
-    self._tok2str = { tid: tok.encode() for tok, tid in special_tokens.items() } | { tid: tok for tok, tid in normal_tokens.items()  }
-    self._special_re = re.compile("|".join(re.escape(tok) for tok in self._special_tokens.keys()) if special_tokens else r"(?!)")
+  def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int]):
+    # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
+    bs = [*range(33, 127), *range(161, 173), *range(174, 256)]  # bytes that map to themselves
+    self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
+
+    # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
+    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre))
+    r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
+    self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
+      f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
+    self._split_to_sentence = re.compile("|".join(re.escape(tok) for tok in special_tokens.keys()) if special_tokens else r"(?!)")
+
+    self._normal_tokens = {bytes(self._byte_decoder[c] for c in tok): tid for tok, tid in normal_tokens.items()}
+    self._special_tokens = special_tokens
+    self._tok2bytes = {tid: tok for tok, tid in self._normal_tokens.items()} | {tid: tok.encode() for tok, tid in self._special_tokens.items()}

  @staticmethod
-  def from_gguf_kv(kv: dict):
+  def from_gguf_kv(kv:dict):
    # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L1818-L1820
    if kv["tokenizer.ggml.pre"] not in ("llama3","llama-v3","llama-bpe"): raise ValueError(f"Invalid tokenizer preset '{kv['tokenizer.ggml.pre']}'")
    vocab: typing.Iterable[tuple[str, int]] = ((tok, idx) for idx, tok in enumerate(kv["tokenizer.ggml.tokens"]))
    normal_tokens, special_tokens = helpers.partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
-    return SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens))
+    return SimpleTokenizer(dict(normal_tokens), dict(special_tokens))

-  def encode(self, text: str):
+  def _encode_word(self, word:bytes) -> list[int]:
+    if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token]
+    parts = [bytes([b]) for b in word]
+    # greedily merge any parts that we can
+    while True:
+      i = min([(sys.maxsize, -1)] + [(self._normal_tokens.get(parts[j]+parts[j+1], sys.maxsize), j) for j in range(len(parts)-1)])[1]
+      if i == -1: break
+      parts[i:i+2] = [parts[i] + parts[i+1]]
+    try: return [self._normal_tokens[p] for p in parts]
+    except KeyError: raise RuntimeError("token not found")
+  def _encode_sentence(self, chunk:str) -> list[int]:
+    return [tok for word in self._split_to_word.findall(chunk) for tok in self._encode_word(word.encode())]
+  def encode(self, text:str) -> list[int]:
    tokens: list[int] = []
    pos = 0
-    for match in self._special_re.finditer(text):
+    for match in self._split_to_sentence.finditer(text):
      tokens.extend(self._encode_sentence(text[pos:match.start(0)]) + [self._special_tokens[text[match.start(0):match.end(0)]]])
      pos = match.end(0)
    return tokens + self._encode_sentence(text[pos:])

-  def decode(self, ids: list[int]) -> str: return b''.join(self._tok2str[tid] for tid in ids).decode()
+  def decode(self, ids:list[int]) -> str: return b''.join(self._tok2bytes[tid] for tid in ids).decode()
  def role(self, role:str): return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")

-  def _encode_sentence(self, chunk: str): return [ tok for word in self._pat.findall(chunk) for tok in self._encode_word(word.encode()) ]
-  def _encode_word(self, word: bytes):
-    if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token]
-    parts = [word[i:i+1] for i in range(len(word))]
-    while True:
-      min_tid, min_idx = 2**32, -1
-      for idx, (p1, p2) in enumerate(zip(parts[:-1], parts[1:])):
-        tid = self._normal_tokens.get(p1 + p2, min_tid)
-        if tid < min_tid: min_tid, min_idx = tid, idx
-      if min_idx == -1: break
-      parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx+1]] + parts[min_idx+2:]
-    try: return [ self._normal_tokens[p] for p in parts ]
-    except KeyError: raise RuntimeError("token not found")
-
 def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor:
  B, H, T, Hd = x.shape
-  assert (Hd & 1) == 0, "RoPE requires an even head dimension"
+  assert isinstance(Hd, int) and (Hd & 1) == 0, "RoPE requires an even head dimension"
  half = Hd // 2
-  angles = (Tensor.arange(T, dtype="float32") + start_pos)[:, None] * (base ** (-(Tensor.arange(half, dtype="float32") / half)))[None, :]
+  t_start_pos = start_pos if isinstance(start_pos, int) else Tensor(start_pos)
+  angles = (Tensor.arange(T, dtype="float32") + t_start_pos)[:, None] * (base ** (-(Tensor.arange(half, dtype="float32") / half)))[None, :]
  # contiguous here allows RoPE to be pruned in the JIT
  cos, sin = angles.cos().reshape(1, 1, T, half).cast(x.dtype).contiguous(), angles.sin().reshape(1, 1, T, half).cast(x.dtype).contiguous()
  x_pairs = x.reshape(B, H, T, half, 2)
@@ -18,8 +18,8 @@ class Opt:

 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
-axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
-               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
+axis_colors = {AxisType.OUTER: "GREEN", AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN",
+               AxisType.LOOP: "WHITE", AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}

 class KernelOptError(Exception): pass
 def check(cond:bool, msg:str=""):
@@ -178,7 +178,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:

  if k.opts.has_threads and k.opts.global_max is not None:
    for threads in [32,16,12,8,6,5,4,3,2]:
-      # Skip is too many threads. Heuristic: use about 128K ops per thread
+      # Skip if too many threads. Heuristic: use about 128K ops per thread
      if threads > k.opts.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue
      for axis in k.axes_of(AxisType.LOOP):
        if k.full_shape[axis] % threads == 0:
@@ -13,8 +13,8 @@ from tinygrad.renderer import Renderer
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])

 # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
-axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
-               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
+axis_to_pos = {AxisType.OUTER: -2, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2,
+               AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}

 class Scheduler:
  def __init__(self, ast:UOp, opts:Renderer):
@@ -165,15 +165,22 @@ def beam_search(lin:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=Tr
          if isinstance(e, RuntimeError): continue
          raise
        timed_lins.append((acted_lins[i], min(tms)))
-        if BEAM_DEBUG > 1: print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(cast(list, p.uops)):5d} uops {time_to_str(compile_et, w=12)} compile/{time_to_str(timed_lins[-1][1], w=12)} run       {len(timed_lins):4d}/{len(acted_lins):4d}         {timed_lins[-1][0].colored_shape()}")  # noqa: E501
-        elif DEBUG >= 2: print(f"\r{time.perf_counter() - st:7.2f}s: {time_to_str(timed_lins[-1][1], w=12)}       {len(timed_lins):4d}/{len(acted_lins):4d}         {timed_lins[-1][0].colored_shape()}\033[K", end="")  # noqa: E501
+        if BEAM_DEBUG > 1:
+          print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(cast(list, p.uops)):5d} uops",
+                f"{time_to_str(compile_et, w=12)} compile/{time_to_str(timed_lins[-1][1], w=12)} run",
+                f"      {len(timed_lins):4d}/{len(acted_lins):4d}         {timed_lins[-1][0].colored_shape()}")
+        elif DEBUG >= 2:
+          print(f"\r{time.perf_counter() - st:7.2f}s: {time_to_str(timed_lins[-1][1], w=12)}",
+                f"      {len(timed_lins):4d}/{len(acted_lins):4d}         {timed_lins[-1][0].colored_shape()}\033[K", end="")

      # done
      opts = sorted(timed_lins, key=lambda x: x[1])
      exiting = len(opts) == 0 or (opts[0][1] < min_progress) or (len(beam) > 0 and ((beam[0][1]-opts[0][1]) < min_progress))
      if not exiting: beam = opts[:amt]
      elif len(opts) > 0 and opts[0][1] < beam[0][1]: beam = opts[:1]
-      if DEBUG >= 2: print(f"\r{time.perf_counter() - st:7.2f}s:", colored(time_to_str(beam[0][1], w=12), "green" if exiting else None), f"from {len(acted_lins):3d} -> {len(opts):3d} actions\033[K", beam[0][0].colored_shape())  # noqa: E501
+      if DEBUG >= 2:
+        print(f"\r{time.perf_counter() - st:7.2f}s:", colored(time_to_str(beam[0][1], w=12), "green" if exiting else None),
+              f"from {len(acted_lins):3d} -> {len(opts):3d} actions\033[K", beam[0][0].colored_shape())
  except KeyboardInterrupt as e:
    if beam_pool is not None: beam_pool.terminate()
    raise e
@@ -23,7 +23,7 @@ class _Device:
  def __getitem__(self, ix:str) -> Compiled: return self.__get_canonicalized_item(self.canonicalize(ix))
  @functools.cache  # this class is a singleton, pylint: disable=method-cache-max-size-none
  def __get_canonicalized_item(self, ix:str) -> Compiled:
-    assert ALLOW_DEVICE_USAGE or ix.split(":")[0] in ["DISK", "NPY", "PYTHON"], f"usage of device {ix} disallowed"
+    assert ALLOW_DEVICE_USAGE or ix.split(":")[0] in ["DISK", "TINYFS", "NPY", "PYTHON"], f"usage of device {ix} disallowed"
    base = (__package__ or __name__).split('.')[0]  # tinygrad
    x = ix.split(":")[0].lower()
    ret = [cls for cname, cls in inspect.getmembers(importlib.import_module(f'{base}.runtime.ops_{x}')) \
@@ -39,7 +39,7 @@ class _Device:
  @functools.cached_property
  def DEFAULT(self) -> str:
    dev = [dev] if (dev:=getenv("DEV", "").upper()) else []
-    from_env = dedup(dev + [d for d in self._devices if d not in ["DISK", "NPY"] and getenv(d) == 1])
+    from_env = dedup(dev + [d for d in self._devices if d not in ["DISK", "TINYFS", "NPY"] and getenv(d) == 1])
    assert len(from_env) < 2, f"multiple devices set in env: {from_env}"
    if len(from_env) == 1: return from_env[0]
    try:
@@ -137,16 +137,14 @@ class Buffer:
    else:
      self._buf = opaque if opaque is not None else self.allocator.alloc(self.nbytes, self.options)
      if not self.device.startswith("DISK"): GlobalCounters.mem_used += self.nbytes
-      if PROFILE:
-        self._prof_num = num = len(Buffer.profile_events)
-        Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", num, {"dtype":self.dtype, "sz":self.size}))
+      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"dtype":self.dtype, "sz":self.size}))
    return self
  def deallocate(self):
    assert hasattr(self, '_buf'), "buffer must be allocated to deallocate"
    if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}")
    if self._base is None and (self.options is None or self.options.external_ptr is None):
      if GlobalCounters is not None and not self.device.startswith("DISK"): GlobalCounters.mem_used -= self.nbytes
-      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self._prof_num))
+      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num))
      self.allocator.free(self._buf, self.nbytes, self.options)
    elif self._base is not None: self._base.allocated_views -= 1
    del self._buf
@@ -160,6 +158,10 @@ class Buffer:
      self.copyout(memoryview(buf))
    return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount)
  @property
+  def trace_num(self) -> int:
+    if not hasattr(self, '_trace_num'): self._trace_num = len(Buffer.profile_events)
+    return self._trace_num
+  @property
  def nbytes(self): return self.size*self.dtype.itemsize
  def __del__(self): (not hasattr(self, '_buf')) or self.deallocate()
  def __repr__(self):
@@ -121,7 +121,7 @@ class BufferCopy(Runner):
      getattr(src.allocator.dev, 'fd', None) is not None and dest.allocator.supports_copy_from_disk
    if src.device.startswith("DISK") and hasattr(dest.allocator, 'copy_from_disk') and disk_supports_fast_copyout and src.nbytes >= 4096:
      dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
-    elif src.device.startswith("DISK") and hasattr(dest.allocator, '_as_buffer'):
+    elif (src.device.startswith("DISK") or src.device.startswith("TINYFS")) and hasattr(dest.allocator, '_as_buffer'):
      # fast(ish) path, uses readinto in diskbuffers
      src.allocator._copyout(dest.allocator._as_buffer(dest._buf), src._buf)
    else:
@@ -165,7 +165,9 @@ class ExecItem:
  def run(self, _var_vals:dict[str, int]|None=None, wait=False, jit=False, do_update_stats=True) -> float|None:
    var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars)
    bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs]
-    if PROFILE: cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, {"metadata":self.metadata, "var_vals":var_vals}))
+    if PROFILE:
+      payload = {"metadata":self.metadata, "var_vals":var_vals, "bufs":[b.trace_num for b in bufs]}
+      cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, payload))
    et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2)
    if do_update_stats:
      GlobalCounters.kernel_count += 1
@@ -39,7 +39,7 @@ pm_gradient = PatternMatcher([
  (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: (ctx.r(Ops.ADD, tuple(i for i,(si,so) in enumerate(zip(ret.src[0].shape, ret.arg)) if si!=so)),)),
  (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
  # there's no gradient for bitcast
-  (UPat(Ops.BITCAST), lambda ctx: (None,)),
+  (UPat(Ops.BITCAST), lambda: (None,)),
 ])

 def _deepwalk(root:UOp, targets:set[UOp]) -> list[UOp]:
@@ -23,10 +23,13 @@ def argfix(*x):
    if len(x) != 1: raise ValueError(f"bad arg {x}")
    return tuple(x[0])
  return x
-def argsort(x): return type(x)(sorted(range(len(x)), key=x.__getitem__)) # https://stackoverflow.com/questions/3382352/equivalent-of-numpy-argsort-in-basic-python
+# https://stackoverflow.com/questions/3382352/equivalent-of-numpy-argsort-in-basic-python
+def argsort(x): return type(x)(sorted(range(len(x)), key=x.__getitem__))
 def all_same(items:tuple[T, ...]|list[T]): return all(x == items[0] for x in items)
 def all_int(t: Sequence[Any]) -> TypeGuard[tuple[int, ...]]: return all(isinstance(s, int) for s in t)
-def colored(st, color:str|None, background=False): return f"\u001b[{10*background+60*(color.upper() == color)+30+['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'].index(color.lower())}m{st}\u001b[0m" if color is not None else st  # replace the termcolor library with one line  # noqa: E501
+def colored(st, color:str|None, background=False): # replace the termcolor library
+  colors = ['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white']
+  return f"\u001b[{10*background+60*(color.upper() == color)+30+colors.index(color.lower())}m{st}\u001b[0m" if color is not None else st
 def colorize_float(x: float): return colored(f"{x:7.2f}x", 'green' if x < 0.75 else 'red' if x > 1.15 else 'yellow')
 def time_to_str(t:float, w=8) -> str: return next((f"{t * d:{w}.2f}{pr}" for d,pr in [(1, "s "),(1e3, "ms")] if t > 10/d), f"{t * 1e6:{w}.2f}us")
 def ansistrip(s:str): return re.sub('\x1b\\[(K|.*?m)', '', s)
@@ -96,7 +99,7 @@ def suppress_finalizing(func):
      if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
  return wrapper

-def unwrap_class_type(cls_t:T): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t
+def unwrap_class_type(cls_t): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t

 def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's')

@@ -150,7 +153,7 @@ CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), Co
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
 FUSE_ATTENTION = ContextVar("FUSE_ATTENTION", 0)
 EMULATE = ContextVar("EMULATE", "")
-CPU_COUNT = ContextVar("CPU_COUNT", max(1, (os.cpu_count() or 1) // (4 if ARCH_X86 else 2))) # take 1/2 of the cores, accounting HT
+CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
 CPU_LLVM, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("AMD_LLVM", 1)
 VIZ = PROFILE = ContextVar("VIZ", 0)
 SPEC = ContextVar("SPEC", 0)
@@ -218,11 +221,12 @@ class TracingKey:
 class ProfileEvent: pass

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

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

 cpu_events:list[ProfileEvent] = []
@contextlib.contextmanager
@@ -281,7 +285,8 @@ def diskcache_put(table:str, key:dict|str|int, val:Any, prepickled=False):
    ltypes = ', '.join(f"{k} {TYPES[type(key[k])]}" for k in key.keys())
    cur.execute(f"CREATE TABLE IF NOT EXISTS '{table}_{VERSION}' ({ltypes}, val blob, PRIMARY KEY ({', '.join(key.keys())}))")
    _db_tables.add(table)
-  cur.execute(f"REPLACE INTO '{table}_{VERSION}' ({', '.join(key.keys())}, val) VALUES ({', '.join(['?']*len(key))}, ?)", tuple(key.values()) + (val if prepickled else pickle.dumps(val), ))  # noqa: E501
+  cur.execute(f"REPLACE INTO '{table}_{VERSION}' ({', '.join(key.keys())}, val) VALUES ({', '.join(['?']*len(key))}, ?)",
+              tuple(key.values()) + (val if prepickled else pickle.dumps(val),))
  conn.commit()
  cur.close()
  return val
@@ -347,10 +352,10 @@ def capstone_flatdump(lib: bytes):
    print(f"{instr.address:#08x}: {instr.mnemonic}\t{instr.op_str}")
  sys.stdout.flush()

-def wait_cond(cb, value=True, timeout_ms=10000, msg="") -> bool:
+def wait_cond(cb, *args, value=True, timeout_ms=10000, msg="") -> bool:
  start_time = int(time.perf_counter() * 1000)
  while int(time.perf_counter() * 1000) - start_time < timeout_ms:
-    if (val:=cb()) == value: return val
+    if (val:=cb(*args)) == value: return val
  raise TimeoutError(f"{msg}. Timed out after {timeout_ms} ms, condition not met: {val} != {value}")

 # *** ctypes helpers
@@ -223,7 +223,7 @@ class InstanceNorm:
  print(t.mean().item(), t.std().item())
  ```
  """
-  def __init__(self, num_features:int, eps=1e-5, affine=True):
+  def __init__(self, num_features:int, eps:float=1e-5, affine:bool=True):
    self.num_features, self.eps = num_features, eps
    self.weight: Tensor|None = Tensor.ones(num_features) if affine else None
    self.bias: Tensor|None = Tensor.zeros(num_features) if affine else None
@@ -249,16 +249,16 @@ class LayerNorm:
  print(t.mean().item(), t.std().item())
  ```
  """
-  def __init__(self, normalized_shape:int|tuple[int, ...], eps=1e-5, elementwise_affine=True):
+  def __init__(self, normalized_shape:int|tuple[int, ...], eps:float=1e-5, elementwise_affine:bool=True):
    self.normalized_shape: tuple[int, ...] = make_tuple(normalized_shape, 1)
-    self.axis, self.eps, self.elementwise_affine = tuple(-1-i for i in range(len(self.normalized_shape))), eps, elementwise_affine
+    self.axis, self.eps = tuple(-1-i for i in range(len(self.normalized_shape))), eps
    self.weight: Tensor|None = Tensor.ones(*self.normalized_shape) if elementwise_affine else None
    self.bias: Tensor|None = Tensor.zeros(*self.normalized_shape) if elementwise_affine else None

  def __call__(self, x:Tensor) -> Tensor:
    assert self.normalized_shape == x.shape[-len(self.normalized_shape):], f"last dimensions of {x.shape} must match {self.normalized_shape}"
    x = x.layernorm(eps=self.eps, axis=self.axis)
-    if not self.elementwise_affine: return x
+    if self.weight is None or self.bias is None: return x
    return x * self.weight + self.bias

 class LayerNorm2d(LayerNorm):
@@ -108,7 +108,9 @@ class CStyleLanguage(Renderer):
  extra_matcher = extra_pm

  def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[DType,bool]]], uops:list[UOp], prefix=None) -> str:
-    tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n" if any(isinstance(dtype, ImageDType) for _,(dtype,_) in bufs) else ""  # noqa: E501
+    tmp = ""
+    if any(isinstance(dtype, ImageDType) for _,(dtype,_) in bufs):
+      tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n"
    buftypes = [(name, self.render_dtype(dtype, mutable)+self.buffer_suffix if isinstance(dtype, (ImageDType, PtrDType)) else
                self.arg_int_prefix if dtype == dtypes.int else None) for name,(dtype,mutable) in bufs]
    local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
@@ -229,10 +231,12 @@ class ClangRenderer(CStyleLanguage):
      # 'static' in C roughly means that function symbol isn't exported. LLVM puts those symbols at the end of object file which allows Clang JIT
      # to just jump at the start of a shellcode without having to deal with symbols or trampolines at all. This is better than having to inline
      # wmma function every time it is called or wasting complexity on a symbol parsing and a memory page on trampoline.
-      prefix += [f"""static {(out := self.render_dtype(dtype_in.vec(N*N)))} __{name}({self.render_dtype(dtype_in.vec(N))} data1, {self.render_dtype(dtype_in.vec(M))} data2, {out} data0){{
+      out, dt1, dt2 = self.render_dtype(dtype_in.vec(N*N)), self.render_dtype(dtype_in.vec(N)), self.render_dtype(dtype_in.vec(M))
+      prefix += [f"""static {out} __{name}({dt1} data1, {dt2} data2, {out} data0){{
  AMX_SET(0);\n  for(int ridx0 = 0; ridx0 < 16; ridx0++){{ AMX(4, (int *)(&data0), 0ull<<62 | (ridx0*4ull)<<56 | ridx0*64ull); }}
  AMX(0, (int *)(&data2), 0ull<<62); AMX(1, (int *)(&data1), 0ull<<62); AMX(12, 0, 0ull);
-  for(int ridx0 = 0; ridx0 < 16; ridx0++){{ AMX(5, (int *)(&data0), 0ull<<62 | (ridx0*4ull)<<56 | ridx0*64ull); }}\n  AMX_SET(1);\n  return data0;\n}}"""] # noqa: E501
+  for(int ridx0 = 0; ridx0 < 16; ridx0++){{ AMX(5, (int *)(&data0), 0ull<<62 | (ridx0*4ull)<<56 | ridx0*64ull); }}
+  AMX_SET(1);\n  return data0;\n}}"""]
    return prefix
  def _render_body(self, function_name, kernel, bufs, uops, pref=None) -> str: return super().render_kernel(function_name, kernel, bufs, uops, pref)
  def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[DType,bool]]]) -> str: return ""
@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H
 from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator
 from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerPairT
-from tinygrad.helpers import getenv, to_mv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32, colored
+from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32, colored
 from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
@@ -670,7 +670,7 @@ class PCIIface(PCIIfaceBase):
  gpus:ClassVar[list[str]] = []

  def __init__(self, dev, dev_id):
-    super().__init__(dev, dev_id, vendor=0x1002, devices=[0x744c, 0x7480, 0x7550], bars=[0, 2, 5], vram_bar=0,
+    super().__init__(dev, dev_id, vendor=0x1002, devices=[0x744c, 0x7480, 0x7550, 0x7590], bars=[0, 2, 5], vram_bar=0,
      va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size)
    self._setup_adev(self.pci_dev.pcibus, self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I'))
    self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
@@ -713,7 +713,7 @@ class PCIIface(PCIIfaceBase):
  def device_fini(self): self.dev_impl.fini()

 class USBIface(PCIIface):
-  def __init__(self, dev, dev_id):
+  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
    self.dev = dev
    self.usb = ASM24Controller()
    self.bars = setup_pci_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30))
@@ -875,13 +875,12 @@ class AMDDevice(HCQCompiled):
  def _at_profile_finalize(self):
    if self.sqtt_enabled:
      wptrs_buf = self.allocator.alloc(round_up(len(self.sqtt_buffers), 0x1000), BufferSpec(cpu_access=True, nolru=True))
-      wptrs = to_mv(wptrs_buf.va_addr, wptrs_buf.size)
      cast(AMDComputeQueue, self.hw_compute_queue_t()).sqtt_stop(len(self.sqtt_buffers), wptrs_buf) \
                                                      .signal(self.timeline_signal, self.next_timeline()).submit(self)
      self.synchronize()
      if DEBUG >= 2: print(f'{self.device}: Saving SQTT in profile...')
      for i,buf0 in enumerate(self.sqtt_buffers):
-        wptr = ((struct.unpack('<I', wptrs[i*4:i*4+4])[0] & 0x1FFFFFFF) - (((buf0.va_addr//32) & 0x1FFFFFFF) if self.target < (12,0,0) else 0)) * 32
+        wptr = ((wptrs_buf.cpu_view().view(fmt='I')[i] & 0x1FFFFFFF) - (((buf0.va_addr//32) & 0x1FFFFFFF) if self.target < (12,0,0) else 0)) * 32
        if DEBUG >= 2: print(f'\t{self.device}: SE {i} blob size {wptr:#x}')
        assert wptr >= 0 and wptr <= buf0.size, f"{wptr} > {buf0.size}, should never happen"
        # When sqtt buffer overflows, wptr stops at the last dword
@@ -23,10 +23,12 @@ class CLCompiler(Compiler):
    build_status: int = cl.clBuildProgram(program, 1, self.dev.device_id, None, cl.clBuildProgram.argtypes[4](), None)
    if build_status != 0:
      cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG, 0, None, log_size := ctypes.c_size_t())
-      cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG, log_size.value, mstr := ctypes.create_string_buffer(log_size.value), None)  # noqa: E501
+      cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG,
+                               log_size.value, mstr := ctypes.create_string_buffer(log_size.value), None)
      raise CompileError(f"OpenCL Compile Error\n\n{mstr.value.decode()}")
    check(cl.clGetProgramInfo(program, cl.CL_PROGRAM_BINARY_SIZES, ctypes.sizeof(ctypes.c_size_t), binary_sizes := (ctypes.c_size_t * 1)(), None))
-    check(cl.clGetProgramInfo(program, cl.CL_PROGRAM_BINARIES, ctypes.sizeof(ctypes.c_void_p), (ctypes.c_void_p * 1)(ctypes.addressof(binary := ctypes.create_string_buffer(binary_sizes[0]))), None))  # noqa: E501
+    check(cl.clGetProgramInfo(program, cl.CL_PROGRAM_BINARIES, ctypes.sizeof(ctypes.c_void_p),
+                              (ctypes.c_void_p * 1)(ctypes.addressof(binary := ctypes.create_string_buffer(binary_sizes[0]))), None))
    check(cl.clReleaseProgram(program))
    return bytes(binary)

@@ -97,16 +99,22 @@ class CLDevice(Compiled):
        err = cl.clGetDeviceIDs(platform_ids[0], device_type, 0, None, num_devices := ctypes.c_uint32())
        if err == 0 and num_devices.value != 0: break
      if DEBUG >= 1: print(f"CLDevice: got {num_platforms.value} platforms and {num_devices.value} devices")
-      CLDevice.device_ids = init_c_var((cl.cl_device_id * num_devices.value)(), lambda x: check(cl.clGetDeviceIDs(platform_ids[0], device_type, num_devices, x, None)))  # noqa: E501
+      CLDevice.device_ids = init_c_var((cl.cl_device_id * num_devices.value)(),
+                                       lambda x: check(cl.clGetDeviceIDs(platform_ids[0], device_type, num_devices, x, None)))

    self.device_id = CLDevice.device_ids[0 if ":" not in device else int(device.split(":")[1])]
-    self.device_name = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_NAME, 256, buf := ctypes.create_string_buffer(256), None), buf.value.decode())[1]  # noqa: E501
-    self.driver_version = (cl.clGetDeviceInfo(self.device_id, cl.CL_DRIVER_VERSION, 256, buf := ctypes.create_string_buffer(256), None), buf.value.decode())[1]  # noqa: E501
+    self.device_name = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_NAME, 256,
+                                           buf:=ctypes.create_string_buffer(256), None), buf.value.decode())[1]
+    self.driver_version = (cl.clGetDeviceInfo(self.device_id, cl.CL_DRIVER_VERSION, 256,
+                                              buf:=ctypes.create_string_buffer(256), None), buf.value.decode())[1]
    if DEBUG >= 1: print(f"CLDevice: opening {self.device_name} with version {self.driver_version}")
    self.context = checked(cl.clCreateContext(None, 1, self.device_id, cl.clCreateContext.argtypes[3](), None, status := ctypes.c_int32()), status)
    self.queue = checked(cl.clCreateCommandQueue(self.context, self.device_id, cl.CL_QUEUE_PROFILING_ENABLE, status), status)
    self.pending_copyin: list[memoryview] = []
-    self.device_exts = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, 4096, ctypes.byref(buf := ctypes.create_string_buffer(4096)), ctypes.byref(total := ctypes.c_size_t())), ctypes.string_at(buf, size=total.value).decode())[1]  # noqa: E501
+    self.device_exts = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, 4096,
+                                           ctypes.byref(buf := ctypes.create_string_buffer(4096)),
+                                           ctypes.byref(total := ctypes.c_size_t())),
+                                           ctypes.string_at(buf, size=total.value).decode())[1]

    compilers = [(IntelRenderer if "cl_intel_subgroup_matrix_multiply_accumulate" in self.device_exts else OpenCLRenderer,
      functools.partial(CLCompiler, self, f"compile_cl_{hashlib.md5(self.device_name.encode() + self.driver_version.encode()).hexdigest()}"))]
@@ -105,8 +105,8 @@ class CPUAllocator(HCQAllocatorBase):
    else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_PRIVATE, mmap.PROT_READ | mmap.PROT_WRITE))
    return HCQBuffer(va:=addr, sz:=size, meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev)
  def _as_buffer(self, src) -> memoryview:
-   self.dev.synchronize()
-   return to_mv(src.va_addr, src.size)
+    self.dev.synchronize()
+    return to_mv(src.va_addr, src.size)
  def _as_dmaref(self, buf):
    self.dev.synchronize()
    return DMACPURef(buf.va_addr, buf.size)
@@ -10,7 +10,9 @@ if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl  # noqa: F401  # pylint:
 if MOCKGPU:=getenv("MOCKGPU"): from test.mockgpu.cuda import cuda # type: ignore # pylint: disable=reimported

 def check(status):
-  if status != 0: raise RuntimeError(f"CUDA Error {status}, {ctypes.string_at(init_c_var(ctypes.POINTER(ctypes.c_char)(), lambda x: cuda.cuGetErrorString(status, ctypes.byref(x)))).decode()}")  # noqa: E501
+  if status != 0:
+    error = ctypes.string_at(init_c_var(ctypes.POINTER(ctypes.c_char)(), lambda x: cuda.cuGetErrorString(status, ctypes.byref(x)))).decode()
+    raise RuntimeError(f"CUDA Error {status}, {error}")

 def encode_args(args, vals) -> tuple[ctypes.Structure, ctypes.Array]:
  c_args = init_c_struct_t(tuple([(f'f{i}', cuda.CUdeviceptr_v2) for i in range(len(args))] +
@@ -424,7 +424,7 @@ class RemoteConnection:
    conns = RemoteConnection.all.keys()
    datas = {conn: conn.req.serialize() for conn in conns}
    reqs, hashes, hash_datas = sum(len(c.req._q) for c in conns), sum(len(c.req._h) for c in conns), sum(len(data) for data in datas.values())
-    resps = []
+    ret, resps = None, []
    with Timing(f"*** send {reqs:-3d} requests {hashes:-3d} hashes with len {hash_datas/1024:.2f} kB in ", enabled=DEBUG>=3):
      for conn,data in datas.items(): conn.conn.request("POST", "/batch", data)
      for conn in datas.keys():
@@ -0,0 +1,137 @@
+import socket, uuid, json, asyncio, threading
+from contextlib import asynccontextmanager
+from tinygrad.device import Compiled, Allocator
+from tinygrad.helpers import DEBUG, getenv
+from tinygrad import Tensor
+
+TINYFS_ENDPOINT = getenv("TINYFS_ENDPOINT", "localhost:6767")
+
+class TinyFSDevice(Compiled):
+  def __init__(self, device:str):
+    self.op = device[len("tinyfs:"):].upper()
+    super().__init__(device, TinyFSAllocator(self), None, None, None)
+
+    self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    self.sock.connect((TINYFS_ENDPOINT.rsplit(":", 1)[0], int(TINYFS_ENDPOINT.rsplit(":", 1)[1])))
+    self.sfile = self.sock.makefile("rwb")
+
+    # fetch node info
+    self.sfile.write(b"INFO\r\n")
+    self.sfile.flush()
+    info = self.sfile.readline()
+    self.node_info = json.loads(info)
+    if DEBUG >= 2: print(f"nodes: {self.node_info}")
+
+    # spawn thread for async copyout
+    self.start_event = threading.Event()
+    self.t = threading.Thread(target=self._start_thread, daemon=True)
+    self.t.start()
+    self.start_event.wait()
+
+    # connection pools
+    self.conn_pools: dict[str, asyncio.Queue] = {}
+    self.conn_pools_lock = asyncio.Lock()
+
+  def finalize(self):
+    self.sfile.close()
+
+    for pool in self.conn_pools.values():
+      while not pool.empty():
+        _, w = pool.get_nowait()
+        w.close()
+        asyncio.run_coroutine_threadsafe(w.wait_closed(), self.loop).result()
+
+    if hasattr(self, "loop"):
+      self.loop.call_soon_threadsafe(self.loop.stop)
+    self.t.join()
+
+  def _start_thread(self):
+    self.loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(self.loop)
+
+    self.start_event.set()
+    self.loop.run_forever()
+    self.loop.close()
+
+  @asynccontextmanager
+  async def connection(self, loc):
+    if loc not in self.conn_pools:
+      await self.conn_pools_lock.acquire()
+      if loc not in self.conn_pools:
+        self.conn_pools[loc] = asyncio.Queue(nw:=getenv("ASYNC_COPY_WORKERS", 4))
+        conn_tasks = [asyncio.open_connection(*self.node_info[loc][-1].rsplit(":", 1)) for _ in range(nw)]
+        connections = await asyncio.gather(*conn_tasks)
+        for reader, writer in connections: self.conn_pools[loc].put_nowait((reader, writer))
+      self.conn_pools_lock.release()
+
+    reader, writer = await self.conn_pools[loc].get()
+    try:
+      yield reader, writer
+    finally:
+      await self.conn_pools[loc].put((reader, writer))
+
+class TinyFSBuffer:
+  def __init__(self, device:TinyFSDevice, size:int, offset=0, request_id=None, copyout_queue=None):
+    self.device, self.size, self.offset = device, size, offset
+    self.request_id: uuid.UUID|None = request_id
+    self.copyout_queue = copyout_queue or []
+  def __repr__(self): return f"<TinyFSBuffer size={self.size} offset={self.offset}>"
+
+class TinyFSAllocator(Allocator[TinyFSDevice]):
+  def _alloc(self, size, options):
+    return TinyFSBuffer(self.dev, size)
+
+  def _copyin(self, dest:TinyFSBuffer, src:memoryview):
+    if DEBUG >= 2: print(f"Copying in {dest.size} bytes to TINYFS:{dest.device.op}")
+    self.dev.sfile.write(f"{dest.device.op}_IN {dest.size}\r\n".encode())
+
+    if dest.device.op == "STORE":
+      self.dev.sfile.flush()
+      dest.request_id = uuid.UUID(bytes=self.dev.sfile.read(16))
+      if DEBUG >= 2: print(f"Request ID: {dest.request_id}")
+
+    self.dev.sfile.write(src)
+    self.dev.sfile.flush()
+
+    if dest.device.op == "LOAD":
+      locs = self.dev.sfile.readline()
+      locs = json.loads(locs)
+
+      dest.copyout_queue = []
+      for i, loc in enumerate(locs):
+        dest.copyout_queue.append((i, loc, src[i*16:(i+1)*16]))
+
+  def _copyout(self, dest:memoryview, src:TinyFSBuffer):
+    if DEBUG >= 2: print(f"Copying out {src.size} bytes from TINYFS:{src.device.op}")
+    if src.device.op == "LOAD":
+      asyncio.run_coroutine_threadsafe(self._copyout_async(dest, src), src.device.loop).result()
+    else:
+      self.dev.sfile.write(f"{src.device.op}_OUT {src.size} {src.request_id}\r\n".encode())
+      self.dev.sfile.flush()
+      src.request_id = uuid.UUID(bytes=self.dev.sfile.read(16))
+      if DEBUG >= 2: print(f"Request ID: {src.request_id}")
+      self.dev.sfile.readinto(dest)
+
+  async def _copyout_async(self, dest:memoryview, src:TinyFSBuffer):
+    async def _worker(item):
+      i, loc, h = item
+      async with self.dev.connection(loc) as (reader, writer):
+        ptr = i * Tensor.CHUNK_SIZE
+        size = min(len(dest[ptr:ptr+Tensor.CHUNK_SIZE]), Tensor.CHUNK_SIZE)
+
+        writer.write(f"CHUNK_OUT {size}\r\n".encode())
+        writer.write(h)
+        await writer.drain()
+
+        chunk = await reader.readexactly(size)
+
+        view = dest[ptr:ptr+len(chunk)]
+        view[:] = chunk
+        del view
+
+    workers = [asyncio.create_task(_worker(item)) for item in src.copyout_queue]
+    await asyncio.gather(*workers)
+    src.copyout_queue.clear()
+
+  def _offset(self, buf:TinyFSBuffer, size:int, offset:int):
+    return TinyFSBuffer(buf.device, size, offset, buf.request_id, buf.copyout_queue)
@@ -113,7 +113,7 @@ class AM_GMC(AM_IP):
    for eng_i in range(18): self.adev.wreg_pair(f"reg{ip}VM_INVALIDATE_ENG{eng_i}_ADDR_RANGE", "_LO32", "_HI32", 0x1fffffffff)
    self.hub_initted[ip] = True

-  @functools.cache
+  @functools.cache  # pylint: disable=method-cache-max-size-none
  def get_pte_flags(self, pte_lv, is_table, frag, uncached, system, snooped, valid, extra=0):
    extra |= (am.AMDGPU_PTE_SYSTEM * system) | (am.AMDGPU_PTE_SNOOPED * snooped) | (am.AMDGPU_PTE_VALID * valid) | am.AMDGPU_PTE_FRAG(frag)
    if not is_table: extra |= (am.AMDGPU_PTE_WRITEABLE | am.AMDGPU_PTE_READABLE | am.AMDGPU_PTE_EXECUTABLE)
@@ -175,7 +175,7 @@ class AM_SMU(AM_IP):

  def _send_msg(self, msg:int, param:int, read_back_arg=False, timeout=10000, debug=False): # default timeout is 10 seconds
    self._smu_cmn_send_msg(msg, param, debug=debug)
-    wait_cond(lambda: (self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).read(), value=1, timeout_ms=timeout,
+    wait_cond((self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).read, value=1, timeout_ms=timeout,
      msg=f"SMU msg {msg:#x} timeout")
    return (self.adev.mmMP1_SMN_C2PMSG_82 if not debug else self.adev.mmMP1_SMN_C2PMSG_53).read() if read_back_arg else None

@@ -60,7 +60,11 @@ def compile_hip(prg:str, arch="gfx1100", asm=False) -> bytes:
    check(comgr.amd_comgr_set_data_name(data_src, b"<null>"))
    check(comgr.amd_comgr_data_set_add(data_set_src, data_src))
    # -include hiprtc_runtime.h was removed
-    check(set_options(action_info, f"-O3 -mcumode --hip-version=6.0.32830 -DHIP_VERSION_MAJOR=6 -DHIP_VERSION_MINOR=0 -DHIP_VERSION_PATCH=32830 -D__HIPCC_RTC__ -std=c++14 -nogpuinc -Wno-gnu-line-marker -Wno-missing-prototypes --offload-arch={arch} -I/opt/rocm/include -Xclang -disable-llvm-passes -Xclang -aux-triple -Xclang x86_64-unknown-linux-gnu".encode())) # noqa: E501
+    options = [
+      "-O3", "-mcumode", "--hip-version=6.0.32830", "-DHIP_VERSION_MAJOR=6", "-DHIP_VERSION_MINOR=0", "-DHIP_VERSION_PATCH=32830",
+      "-D__HIPCC_RTC__", "-std=c++14", "-nogpuinc", "-Wno-gnu-line-marker", "-Wno-missing-prototypes", f"--offload-arch={arch}",
+      "-I/opt/rocm/include", "-Xclang -disable-llvm-passes", "-Xclang -aux-triple", "-Xclang x86_64-unknown-linux-gnu"]
+    check(set_options(action_info, ' '.join(options).encode()))
    status = comgr.amd_comgr_do_action(comgr.AMD_COMGR_ACTION_COMPILE_SOURCE_WITH_DEVICE_LIBS_TO_BC, action_info, data_set_src, data_set_bc)
    if status != 0:
      print(_get_comgr_data(data_set_bc, comgr.AMD_COMGR_DATA_KIND_LOG).decode())
@@ -22,10 +22,12 @@ def jitlink_check(status, ctx=None):

 def pretty_ptx(s):
  # all expressions match `<valid_before><expr><valid_after>` and replace it with `<valid_before>color(<expr>)<valid_after>`
-  s = re.sub(r'([!@<\[\s,\+\-;\n])((?:[_%$][\w%\$_]+(?:\.[xyz])?\:?)|(?:buf\d+))([<>\]\s,\+\-;\n\)])', lambda m:m[1]+colored(m[2], "blue")+m[3], s, flags=re.M) # identifiers  # noqa: E501
+  s = re.sub(r'([!@<\[\s,\+\-;\n])((?:[_%$][\w%\$_]+(?:\.[xyz])?\:?)|(?:buf\d+))([<>\]\s,\+\-;\n\)])',
+             lambda m:m[1]+colored(m[2], "blue")+m[3], s, flags=re.M) # identifiers
  s = re.sub(r'(.)((?:b|s|u|f)(?:8|16|32|64)|pred)([\.\s])', lambda m:m[1]+colored(m[2], "green")+m[3], s, flags=re.M) # types
  s = re.sub(r'^(\s*)([\w]+)(.*?;$)', lambda m:m[1]+colored(m[2], "yellow")+m[3], s, flags=re.M) # instructions
-  s = re.sub(r'([<>\[\]\s,\+\-;])((?:0[fF][0-9a-fA-F]{8})|(?:[0-9]+)|(?:0[xX][0-9a-fA-F]+))([<>\[\]\s,\+\-;])', lambda m:m[1]+colored(m[2], "yellow")+m[3], s, flags=re.M) # numbers  # noqa: E501
+  s = re.sub(r'([<>\[\]\s,\+\-;])((?:0[fF][0-9a-fA-F]{8})|(?:[0-9]+)|(?:0[xX][0-9a-fA-F]+))([<>\[\]\s,\+\-;])',
+             lambda m:m[1]+colored(m[2], "yellow")+m[3], s, flags=re.M) # numbers
  s = re.sub(r'(\.)(param|reg|global)', lambda m:m[1]+colored(m[2], "magenta"), s, flags=re.M) # space
  s = re.sub(r'(\.)(version|target|address_size|visible|entry)', lambda m:m[1]+colored(m[2], "magenta"), s, flags=re.M) # derivatives
  return s
@@ -33,7 +33,7 @@ def elf_loader(blob:bytes, force_section_align:int=1) -> tuple[memoryview, list[
  for sh, trgt_sh_name, c_rels in rel + rela:
    target_image_off = next(tsh for tsh in sections if tsh.name == trgt_sh_name).header.sh_addr
    rels = [(r.r_offset, symtab[libc.ELF64_R_SYM(r.r_info)], libc.ELF64_R_TYPE(r.r_info), getattr(r, "r_addend", 0)) for r in c_rels]
-    for roff, sym, r_type_, r_addend in rels:
+    for _, sym, _, _ in rels:
      if sym.st_shndx == 0: raise RuntimeError(f'Attempting to relocate against an undefined symbol {repr(_strtab(sh_strtab, sym.st_name))}')
    relocs += [(target_image_off + roff, sections[sym.st_shndx].header.sh_addr + sym.st_value, rtype, raddend) for roff, sym, rtype, raddend in rels]

@@ -30,10 +30,10 @@ class TLSFAllocator:
    self.blocks:dict[int, tuple[int, int|None, int|None, bool]] = {0: (size, None, None, True)} # size, next, prev, is_free
    self._insert_block(0, size)

-  @functools.cache
+  @functools.cache # pylint: disable=method-cache-max-size-none
  def lv1(self, size): return size.bit_length()

-  @functools.cache
+  @functools.cache # pylint: disable=method-cache-max-size-none
  def lv2(self, size): return (size - (1 << (size.bit_length() - 1))) // (1 << max(0, size.bit_length() - self.l2_cnt))

  def _insert_block(self, start:int, size:int, prev:int|None=None):
@@ -209,7 +209,7 @@ class MemoryManager:
    if getenv("MM_DEBUG", 0): print(f"mm {self.dev.devfmt}: unmapping {vaddr=:#x} ({size=:#x})")

    ctx = PageTableTraverseContext(self.dev, self.root_page_table, vaddr, free_pts=True)
-    for off, pt, pte_idx, pte_cnt, pte_covers in ctx.next(size):
+    for _, pt, pte_idx, pte_cnt, _ in ctx.next(size):
      for pte_id in range(pte_idx, pte_idx + pte_cnt):
        assert pt.valid(pte_id), f"PTE not mapped: {pt.entry(pte_id):#x}"
        pt.set_entry(pte_id, paddr=0x0, valid=False)
@@ -124,6 +124,7 @@ class NV_FLCN(NV_IP):
    def __patch(cmd_id, cmd):
      patched_image = bytearray(image)

+      dmem_offset = 0
      hdr = nv.FALCON_APPLICATION_INTERFACE_HEADER_V1.from_buffer_copy(image[(app_hdr_off:=self.desc_v3.IMEMLoadSize+self.desc_v3.InterfaceOffset):])
      ents = (nv.FALCON_APPLICATION_INTERFACE_ENTRY_V1 * hdr.entryCount).from_buffer_copy(image[app_hdr_off + ctypes.sizeof(hdr):])
      for i in range(hdr.entryCount):
@@ -334,7 +335,7 @@ class NV_GSP(NV_IP):
    # Fill up arguments
    queue_args = nv.MESSAGE_QUEUE_INIT_ARGUMENTS(sharedMemPhysAddr=queues_sysmem[0], pageTableEntryCount=pte_cnt, cmdQueueOffset=pt_size,
      statQueueOffset=pt_size + queue_size)
-    rm_args, self.rm_args_sysmem = self.nvdev._alloc_boot_struct(nv.GSP_ARGUMENTS_CACHED(bDmemStack=True, messageQueueInitArguments=queue_args))
+    _, self.rm_args_sysmem = self.nvdev._alloc_boot_struct(nv.GSP_ARGUMENTS_CACHED(bDmemStack=True, messageQueueInitArguments=queue_args))

    # Build command queue header
    self.cmd_q_va, self.stat_q_va = queues_va + pt_size, queues_va + pt_size + queue_size
@@ -481,7 +482,7 @@ class NV_GSP(NV_IP):
      params.ramfcMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=ramfc_alloc.paddrs[0][0], size=0x200, addressSpace=2, cacheAttrib=0)
      params.instanceMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=ramfc_alloc.paddrs[0][0], size=0x1000, addressSpace=2, cacheAttrib=0)

-      method_va, method_sysmem = System.alloc_sysmem(0x5000, contiguous=True)
+      _, method_sysmem = System.alloc_sysmem(0x5000, contiguous=True)
      params.mthdbufMem = nv_gpu.NV_MEMORY_DESC_PARAMS(base=method_sysmem[0], size=0x5000, addressSpace=1, cacheAttrib=0)

      if client is not None and client != self.priv_root and params.hObjectError != 0:
@@ -557,7 +558,7 @@ class NV_GSP(NV_IP):
        self.nvdev.wreg(addr, (self.nvdev.rreg(addr) & ~mask) | (val & mask))
      elif op == 0x2: # reg poll
        addr, mask, val, _, _ = next(cmd_iter), next(cmd_iter), next(cmd_iter), next(cmd_iter), next(cmd_iter)
-        wait_cond(lambda: (self.nvdev.rreg(addr) & mask), value=val, msg=f"Register {addr:#x} not equal to {val:#x} after polling")
+        wait_cond(lambda a, m: (self.nvdev.rreg(a) & m), addr, mask, value=val, msg=f"Register {addr:#x} not equal to {val:#x} after polling")
      elif op == 0x3: time.sleep(next(cmd_iter) / 1e6) # delay us
      elif op == 0x4: # save reg
        addr, index = next(cmd_iter), next(cmd_iter)
@@ -152,6 +152,8 @@ class NVDev(PCIDevImplBase):
    return gzip.decompress(struct.pack("<4BL2B", 0x1f, 0x8b, 8, 0, 0, 0, 3) + image) if "COMPRESSION: YES" in info else image

  def include(self, file:str):
+    def _do_eval(s:str): return eval(s) # pylint: disable=eval-used
+
    regs_off = {'NV_PFALCON_FALCON': 0x0, 'NV_PGSP_FALCON': 0x0, 'NV_PSEC_FALCON': 0x0, 'NV_PRISCV_RISCV': 0x1000, 'NV_PGC6_AON': 0x0, 'NV_PFSP': 0x0,
      'NV_PGC6_BSI': 0x0, 'NV_PFALCON_FBIF': 0x600, 'NV_PFALCON2_FALCON': 0x1000, 'NV_PBUS': 0x0, 'NV_PFB': 0x0, 'NV_PMC': 0x0, 'NV_PGSP_QUEUE': 0x0,
      'NV_VIRTUAL_FUNCTION':0xb80000}
@@ -163,13 +165,13 @@ class NVDev(PCIDevImplBase):
        name, hi, lo = m.groups()

        reg = next((r for r in self.reg_names if name.startswith(r+"_")), None)
-        if reg is not None: self.__dict__[reg].add_field(name[len(reg)+1:].lower(), eval(lo), eval(hi))
-        else: self.reg_offsets[name] = (eval(lo), eval(hi))
+        if reg is not None: self.__dict__[reg].add_field(name[len(reg)+1:].lower(), _do_eval(lo), _do_eval(hi))
+        else: self.reg_offsets[name] = (_do_eval(lo), _do_eval(hi))
        continue

      if m:=re.match(r'#define\s+(\w+)\s*\(\s*(\w+)\s*\)\s*(.+)', raw): # reg set
        fn = m.groups()[2].strip().rstrip('\\').split('/*')[0].rstrip()
-        name, value = m.groups()[0], eval(f"lambda {m.groups()[1]}: {fn}")
+        name, value = m.groups()[0], _do_eval(f"lambda {m.groups()[1]}: {fn}")
      elif m:=re.match(r'#define\s+(\w+)\s+([0-9A-Fa-fx]+)(?![^\n]*:)', raw): name, value = m.groups()[0], int(m.groups()[1], 0) # reg value
      else: continue

@@ -10,14 +10,14 @@ MAP_FIXED, MAP_LOCKED, MAP_POPULATE, MAP_NORESERVE = 0x10, 0 if OSX else 0x2000,
 class _System:
  def reserve_hugepages(self, cnt): os.system(f"sudo sh -c 'echo {cnt} > /proc/sys/vm/nr_hugepages'")

-  def memory_barrier(self): lib.atomic_thread_fence(__ATOMIC_SEQ_CST:=5) if (lib:=self.atomic_lib()) is not None else None
+  def memory_barrier(self): lib.atomic_thread_fence(__ATOMIC_SEQ_CST:=5) if (lib:=self.atomic_lib) is not None else None

  def lock_memory(self, addr:int, size:int):
    if libc.mlock(ctypes.c_void_p(addr), size): raise RuntimeError(f"Failed to lock memory at {addr:#x} with size {size:#x}")

  def system_paddrs(self, vaddr:int, size:int) -> list[int]:
-    self.pagemap().seek(vaddr // mmap.PAGESIZE * 8)
-    return [(x & ((1<<55) - 1)) * mmap.PAGESIZE for x in array.array('Q', self.pagemap().read(size//mmap.PAGESIZE*8, binary=True))]
+    self.pagemap.seek(vaddr // mmap.PAGESIZE * 8)
+    return [(x & ((1<<55) - 1)) * mmap.PAGESIZE for x in array.array('Q', self.pagemap.read(size//mmap.PAGESIZE*8, binary=True))]

  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False, data:bytes|None=None) -> tuple[int, list[int]]:
    assert not contiguous or size <= (2 << 20), "Contiguous allocation is only supported for sizes up to 2MB"
@@ -36,17 +36,17 @@ class _System:
      if vendor == target_vendor and device in target_devices: result.append(pcibus)
    return sorted(result)

-  @functools.cache
+  @functools.cached_property
  def atomic_lib(self): return ctypes.CDLL(ctypes.util.find_library('atomic')) if sys.platform == "linux" else None

-  @functools.cache
+  @functools.cached_property
  def pagemap(self) -> FileIOInterface:
    if FileIOInterface(reloc_sysfs:="/proc/sys/vm/compact_unevictable_allowed", os.O_RDONLY).read()[0] != "0":
      os.system(cmd:=f"sudo sh -c 'echo 0 > {reloc_sysfs}'")
      assert FileIOInterface(reloc_sysfs, os.O_RDONLY).read()[0] == "0", f"Failed to disable migration of locked pages. Please run {cmd} manually."
    return FileIOInterface("/proc/self/pagemap", os.O_RDONLY)

-  @functools.cache
+  @functools.cached_property
  def vfio(self) -> FileIOInterface|None:
    try:
      if not FileIOInterface.exists("/sys/module/vfio"): os.system("sudo modprobe vfio-pci disable_idle_d3=1")
@@ -90,7 +90,7 @@ class PCIDevice:
                                " to allow python accessing device or run with sudo") from 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:
+    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")
      FileIOInterface("/sys/bus/pci/drivers_probe", os.O_WRONLY).write(self.pcibus)
      iommu_group = FileIOInterface.readlink(f"/sys/bus/pci/devices/{self.pcibus}/iommu_group").split('/')[-1]
@@ -229,7 +229,7 @@ class ASM24Controller:
    for i in range(0, len(ops), bs:=(4 if OSX else 16)): self.exec_ops(list(itertools.chain.from_iterable(ops[i:i+bs])))

 class USBMMIOInterface(MMIOInterface):
-  def __init__(self, usb, addr, size, fmt, pcimem=True):
+  def __init__(self, usb, addr, size, fmt, pcimem=True): # pylint: disable=super-init-not-called
    self.usb, self.addr, self.nbytes, self.fmt, self.pcimem, self.el_sz = usb, addr, size, fmt, pcimem, struct.calcsize(fmt)

  def __getitem__(self, index): return self._access_items(index)
@@ -256,13 +256,14 @@ class USBMMIOInterface(MMIOInterface):

      acc, acc_size = self._acc_size(sz)
      return bytes(array.array(acc, [self._acc_one(off + i * acc_size, acc_size) for i in range(sz // acc_size)]))
-    else: # write op
-      data = struct.pack(self.fmt, data) if isinstance(data, int) else bytes(data)

-      if not self.pcimem:
-        # Fast path for writing into buffer 0xf000
-        use_cache = 0xa800 <= self.addr <= 0xb000
-        return self.usb.scsi_write(bytes(data)) if self.addr == 0xf000 else self.usb.write(self.addr + off, bytes(data), ignore_cache=not use_cache)
+    # write op
+    data = struct.pack(self.fmt, data) if isinstance(data, int) else bytes(data)

-      _, acc_sz = self._acc_size(len(data) * struct.calcsize(self.fmt))
-      self.usb.pcie_mem_write(self.addr+off, [int.from_bytes(data[i:i+acc_sz], "little") for i in range(0, len(data), acc_sz)], acc_sz)
+    if not self.pcimem:
+      # Fast path for writing into buffer 0xf000
+      use_cache = 0xa800 <= self.addr <= 0xb000
+      return self.usb.scsi_write(bytes(data)) if self.addr == 0xf000 else self.usb.write(self.addr + off, bytes(data), ignore_cache=not use_cache)
+
+    _, acc_sz = self._acc_size(len(data) * struct.calcsize(self.fmt))
+    self.usb.pcie_mem_write(self.addr+off, [int.from_bytes(data[i:i+acc_sz], "little") for i in range(0, len(data), acc_sz)], acc_sz)
@@ -3,7 +3,7 @@ import functools, operator, itertools
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType
-from tinygrad.uop.symbolic import sym, symbolic
+from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, TracingKey

 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
@@ -17,7 +17,7 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
    if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None

 def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
-  if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
+  #if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
  # if it's a kernel, we don't realize it
  if a.src[1].op is not Ops.KERNEL: ctx[a] = None

@@ -25,7 +25,7 @@ pm_generate_realize_map = PatternMatcher([
  # always realize SINK src
  (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
  # always realize COPY/BUFFER_VIEW/CONTIGUOUS
-  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS}, name="tr"), realize),
+  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.ENDRANGE}, name="tr"), realize),
  # realize srcs of COPY, MSELECT, MSTACK
  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
  # realize ASSIGN and input to assign (might be optimized out)
@@ -112,8 +112,10 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
    case Ops.EXPAND:  rngs = tuple(a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, in_shape, arg))
    case Ops.PAD:
      # TODO: why is multiple graph_rewrites faster than one here?
-      rngs = tuple(r if (s == 0 and e == 0) else graph_rewrite(((r >= s) & (r < (sh+s))).where(r-s, UOp.invalid()), sym, name="pad")
-        for r,sh,(s,e) in zip(rngs, in_shape, arg))
+      # TODO: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local`
+      #       wraps the pad with only the newly added valid
+      rngs = tuple(r if (s == 0 and e == 0) else graph_rewrite(((r >= s) & (r < (sh+s))),
+        symbolic+pm_simplify_valid, name="pad").where(r-s, UOp.invalid()) for r,sh,(s,e) in zip(rngs, in_shape, arg))
    case Ops.RESHAPE:
      acc = 1
      axes_in:list[UOp] = []
@@ -126,7 +128,8 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
        axes_out.append(combined_axes % s)
        combined_axes //= s
      # this simplify is doing a lot of heavy lifting. this is the replacement for the reshape view merging code
-      rngs = graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic, name="reshape").src
+      rngs = graph_rewrite(graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic+pm_simplify_valid, name="reshape"),
+        pm_drop_and_clauses, name="reshape drop ands").src
    case _: raise RuntimeError(f"{op} is not a MovementOp")
  return rngs

@@ -110,7 +110,7 @@ pm_mops = PatternMatcher([
 # 3.5 cleanups

 # Ops.NOOP happens when we have a COPY to the device the Tensor is already on. We treat it like COPY here for MSTACK.
-ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.NOOP}
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.NOOP, Ops.ENDRANGE}

 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
@@ -207,7 +207,7 @@ pm_cleanups = pm_mops+PatternMatcher([
 ])

 def late_buffer_view(t:UOp, b:UOp):
-  if isinstance(b.device, str) and b.device.startswith("DISK"):
+  if isinstance(b.device, str) and (b.device.startswith("DISK") or b.device.startswith("TINYFS")):
    rngs = b.src[1:]
    size = prod(shape := [int(r.vmax+1) for r in rngs])

@@ -338,6 +338,7 @@ def handle_assign(ctx:LocalAddBufferContext, assign:UOp):

 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
  if r.tag is not None: return None
+  if r.arg[-1] is AxisType.OUTER: return None
  ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=())
  ctx.range += 1
  return ret
@@ -412,7 +413,7 @@ class Kernel:
    return f"<Kernel {len(list(self.ast.toposort()))} {ast_rep} {self.metadata}>"

 def split_store(ctx:list[UOp], x:UOp):
-  if len(x.ranges): return None
+  if len([r for r in x.ranges if r.arg[-1] != AxisType.OUTER]): return None
  if x.src[0].ptrdtype.addrspace is AddrSpace.LOCAL: return None

  # local kernel rewrite
@@ -424,7 +425,7 @@ def split_store(ctx:list[UOp], x:UOp):

  # NOTE: the hack for COPY is here
  ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts) if lctx.opts is not None else None) \
-    if ret.src[1].op not in {Ops.COPY, Ops.BUFFER_VIEW} else ret.src[1]
+    if ret.src[1].op not in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENDRANGE} else ret.src[1]
  kernel_arg = Kernel(ret,tuple(dedup(flatten([x for x in metadatas if x is not None])))[::-1])
  kernel = UOp(Ops.KERNEL, src=tuple(lctx.map.values())+tuple(lctx.vars.keys()), arg=kernel_arg)
  if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src if x.op is not Ops.BIND]):
@@ -481,6 +482,11 @@ def do_sub_recurse(s:UOp):
  return x.replace(src=tuple([UOp(Ops.SUBSTITUTE, dtype=y.dtype, src=(y,uop_keys,uop_values)) for y in x.src]))
 pm_substitute_recurse = PatternMatcher([(UPat(Ops.SUBSTITUTE, src=(UPat(), UPat(Ops.NOOP), UPat(Ops.NOOP)), name="s"), do_sub_recurse)])

+pm_localize_bufs = PatternMatcher([
+  (UPat(Ops.BUFFERIZE, name="x"), lambda x:
+   x.replace(arg=BufferizeOpts(device=None, addrspace=AddrSpace.LOCAL), tag=None) if len(x.ranges) > 0 else None),
+])
+
@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len([u for u in UOp.sink(*ret.values()).toposort() if u.op is Ops.KERNEL]))}", True)
 def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
  uop_list: list[UOp] = []
@@ -496,7 +502,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
  tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers")
  # TODO: can you substitute and remove costly buffers at the same time?
  tsink = graph_rewrite(tsink, pm_substitute_recurse, bottom_up=True, name="run substitutes")
-  tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
+  tsink = graph_rewrite(tsink, pm_localize_bufs+pm_limit_bufs, ctx=rctx, name="localize/limit buffers")

  # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph
  # MSTACK stacks multiple BUFFERIZEs in one tagged tensor
@@ -9,8 +9,8 @@ from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_u
 from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG, is_numpy_ndarray, FUSE_ATTENTION
 from tinygrad.helpers import suppress_finalizing
 from tinygrad.gradient import compute_gradient
-from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, MathTrait, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop, \
-  srender
+from tinygrad.uop.mathtraits import MathTrait
+from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop, srender
 from tinygrad.uop.spec import tensor_uop_spec, type_verify
 from tinygrad.device import Device, Buffer
 from tinygrad.engine.realize import run_schedule
@@ -411,6 +411,59 @@ class Tensor(MathTrait):
    """
    return self.replace(self.shard(devices, axis))

+  CHUNK_SIZE = 2**20
+  def load(self, size:int) -> Tensor:
+    """
+    Load a tensor from storage.
+
+    self should be a tensor of the hash to load
+    """
+    # TODO: this should work locally as well
+    assert self.dtype == dtypes.uint8, "hash is expected to be uint8"
+    h = self.contiguous().flatten()
+    assert h.shape[0] == 16, "expected hash"
+
+    base_chunks = math.ceil(size / Tensor.CHUNK_SIZE)
+    tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
+    data, level_chunks = h, 0
+    for i in reversed(range(tree_depth + 1)):
+      data = data.to("tinyfs:load")
+
+      # if not last level, its still hashes
+      if i > 0 or tree_depth == 0:
+        level_chunks = max(1, math.ceil(base_chunks / (Tensor.CHUNK_SIZE // 16)**(i-1)))
+        pad_amt = 16 * level_chunks
+      else: pad_amt = Tensor.CHUNK_SIZE * level_chunks
+      if (tsize := data.shape[0]) < pad_amt: data = data.pad((0, pad_amt - tsize))
+      data = data[:pad_amt].contiguous()
+      if i != 0: data = data.to(self.device)
+
+    return data[:size]
+
+  def store(self) -> Tensor:
+    """
+    Store a tensor to storage.
+    """
+    # TODO: this should work locally as well
+    data = self.contiguous().flatten().bitcast(dtypes.uint8)
+
+    # pad to a multiple of 1mb
+    if (tsize := data.shape[0]) % Tensor.CHUNK_SIZE != 0: data = data.pad((0, Tensor.CHUNK_SIZE - tsize % Tensor.CHUNK_SIZE))
+    size = data.shape[0]
+
+    base_chunks = math.ceil(size / Tensor.CHUNK_SIZE)
+    tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))
+
+    to_device = "CPU" if isinstance(self.device, str) and self.device.startswith("DISK") else self.device
+
+    level_chunks = base_chunks
+    for _ in range(tree_depth + 1):
+      data = data.to("tinyfs:store")[:level_chunks * 16].contiguous().to(to_device)
+      if (tsize := data.shape[0]) % Tensor.CHUNK_SIZE != 0: data = data.pad((0, Tensor.CHUNK_SIZE - tsize % Tensor.CHUNK_SIZE))
+      level_chunks = math.ceil(data.shape[0] / Tensor.CHUNK_SIZE)
+
+    return data[:16].contiguous()
+
  @staticmethod
  def from_uop(y:UOp, **kwargs) -> Tensor:
    if y.op is Ops.BIND: return Tensor(y, **kwargs, requires_grad=False)
@@ -1159,12 +1212,13 @@ class Tensor(MathTrait):
      match index:
        case Tensor():
          if not dtypes.is_int(index.dtype): raise IndexError(f"index dtype {index.dtype} is not supported")
+          assert isinstance(size, int), "size must be an int"
          index = (index < 0).where(index+size, index).to(self.device)  # treat negative index values
        case list() | tuple():
          if not dtypes.is_int((ti:=Tensor(index)).dtype): raise IndexError(f"{index=} contains non-int element")
          index = Tensor([i+size if i<0 else i for i in fully_flatten(index)], self.device, requires_grad=False).reshape(ti.shape)
        case int() | UOp(): # sint
-          if index >= size or index < -size: raise IndexError(f"{index=} is out of bounds with {size=}")
+          #if index >= size or index < -size: raise IndexError(f"{index=} is out of bounds with {size=}")
          # TODO: is this right for (negative) symbolic?
          boundary = [index, index+1] if index >= 0 else [index+size, index+size+1]
        case slice():
@@ -2484,17 +2538,20 @@ class Tensor(MathTrait):
    if IMAGE: return self.image_conv2d(weight, bias, groups, stride, dilation, padding, dtype)
    (bs,cin_), (cout,cin), HW = self.shape[:2], weight.shape[:2], weight.shape[2:]
    padding_ = self._resolve_pool_pads(padding, len(HW))
-    assert groups*cin == cin_ and len(self.shape) == len(weight.shape), f"Input Tensor shape {self.shape} does not match the shape of the weights {weight.shape}. ({groups*cin} vs. {cin_})"  # noqa: E501
+    assert groups*cin == cin_ and len(self.shape) == len(weight.shape),\
+        f"Input Tensor shape {self.shape} does not match the shape of the weights {weight.shape}. ({groups*cin} vs. {cin_})"

    # conv2d is a pooling op (with padding)
    x = self.pad(padding_)._pool(HW, stride, dilation)   # (bs, groups*cin, oy, ox, H, W)
    rcout, oyx = cout//groups, x.shape[2:-len(HW)]
    if not all(x == 3 for x in HW) or stride != 1 or dilation != 1 or not WINO:
      # normal conv
-      x = x.reshape(bs, groups, cin, 1, *oyx, *HW).expand(bs, groups, cin, rcout, *oyx, *HW).permute(0,1,3,*[4+i for i in range(len(oyx))],2,*[4+len(oyx)+i for i in range(len(HW))])  # noqa: E501
+      x = x.reshape(bs, groups, cin, 1, *oyx, *HW).expand(bs, groups, cin, rcout, *oyx, *HW)\
+        .permute(0,1,3,*[4+i for i in range(len(oyx))],2,*[4+len(oyx)+i for i in range(len(HW))])

      # conv! broadcasted to (bs, groups, rcout, *oyx, cin, *HW)
-      ret = (x * weight.reshape(1, groups, rcout, *[1] * len(oyx), cin, *HW)).sum([-1-i for i in range(1+len(oyx))], keepdim=True, dtype=dtype).reshape(bs, cout, *oyx)  # noqa: E501
+      ret = (x * weight.reshape(1, groups, rcout, *[1] * len(oyx), cin, *HW))\
+        .sum([-1-i for i in range(1+len(oyx))], keepdim=True, dtype=dtype).reshape(bs, cout, *oyx)
      return ret if bias is None else ret.add(bias.reshape(1, -1, *[1] * len(HW)))

    HWI, HWO = (6,) * len(HW), (4,) * len(HW)  # F(4x4,3x3) winograd tiles
@@ -2505,7 +2562,8 @@ class Tensor(MathTrait):
    # TODO: stride == dilation
    # use padding to round up to 4x4 output tiles
    # (bs, cin_, tyx, HWI)
-    d = self.pad(sum([[padding_[i*2], padding_[i*2+1] + (-(dim + sum(padding_[i * 2:(i + 1) * 2]) - 2) % 4)] for i, dim in enumerate(self.shape[-len(HW):])], []))._pool(HWI, HWO)  # noqa: E501
+    pads = [[padding_[i*2], padding_[i*2+1] + (-(dim + sum(padding_[i * 2:(i + 1) * 2]) - 2) % 4)] for i, dim in enumerate(self.shape[-len(HW):])]
+    d = self.pad(sum(pads, []))._pool(HWI, HWO)
    # move HW to the front: # (HWI, bs, cin_, tyx)
    d = d.permute(*range(len(d.shape)-len(HW),len(d.shape)), *range(len(d.shape)-len(HW)))
    tyx = d.shape[-len(HWI):]  # dim of tiling
@@ -2627,7 +2685,8 @@ class Tensor(MathTrait):
    base = ret[..., -1]._cumalu(-1, op, _include_initial=True)
    base = base.unsqueeze(-1).expand(*base.shape, ret.shape[-1])
    def fix(x: Tensor) -> Tensor: return x.flatten(start_dim=-2)[..., -s:].transpose(axis,-1)
-    return {Ops.ADD: Tensor.__add__, Ops.MAX: Tensor.maximum, Ops.MUL: Tensor.__mul__}[op](fix(ret), fix(base))
+    reduce_fxns: dict[Ops, Callable[[Tensor, Tensor], Tensor]] = {Ops.ADD: Tensor.__add__, Ops.MAX: Tensor.maximum, Ops.MUL: Tensor.__mul__}
+    return reduce_fxns[op](fix(ret), fix(base))

  def cumsum(self, axis:int=0) -> Tensor:
    """
@@ -3666,7 +3725,7 @@ class Tensor(MathTrait):
    if self.dtype != dtypes.bool and not dtypes.is_int(self.dtype): raise RuntimeError(f"{self.dtype} is not supported")
    return self.logical_not() if self.dtype == dtypes.bool else self ^ -1

-  def lshift(self, x:int, reverse=False) -> Tensor:
+  def lshift(self, x:Tensor|int, reverse=False) -> Tensor:
    """
    Computes left arithmetic shift of `self` by `x` bits. `self` must have unsigned dtype.
    Equivalent to `self << x`.
@@ -3678,7 +3737,7 @@ class Tensor(MathTrait):
    assert dtypes.is_unsigned(self.dtype) and isinstance(x, int) and x >= 0 and not reverse, f"not supported {self.dtype=} {x=}"
    return self.mul(2 ** x, reverse)

-  def rshift(self, x:int, reverse=False) -> Tensor:
+  def rshift(self, x:Tensor|int, reverse=False) -> Tensor:
    """
    Computes right arithmetic shift of `self` by `x` bits. `self` must have unsigned dtype.
    Equivalent to `self >> x`.
@@ -3794,18 +3853,20 @@ class Tensor(MathTrait):
  def __rpow__(self, x) -> Tensor: return self.pow(x, True)
  def __rmatmul__(self, x) -> Tensor: return self.matmul(x, True)

-  def __iadd__(self, x) -> Tensor: return self.assign(self.add(x))
-  def __isub__(self, x) -> Tensor: return self.assign(self.sub(x))
-  def __imul__(self, x) -> Tensor: return self.assign(self.mul(x))
-  def __ipow__(self, x) -> Tensor: return self.assign(self.pow(x))
-  def __itruediv__(self, x) -> Tensor: return self.assign(self.div(x))
  def __ifloordiv__(self, x) -> Tensor: return self.assign(self.__floordiv__(x))
+  def __ipow__(self, x) -> Tensor: return self.assign(self.pow(x))
  def __imatmul__(self, x) -> Tensor: return self.assign(self.matmul(x))
-  def __iand__(self, x) -> Tensor: return self.assign(self.bitwise_and(x))
-  def __ior__(self, x) -> Tensor: return self.assign(self.bitwise_or(x))
-  def __ixor__(self, x) -> Tensor: return self.assign(self.bitwise_xor(x))
-  def __ilshift__(self, x) -> Tensor: return self.assign(self.lshift(x))
-  def __irshift__(self, x) -> Tensor: return self.assign(self.rshift(x))
+
+  # unlike Tensors, UOps are immutable, so these don't go in MathTraits
+  def __iadd__(self, x) -> Tensor: return self.assign(self.add(x)) # type: ignore[misc]
+  def __isub__(self, x) -> Tensor: return self.assign(self.sub(x)) # type: ignore[misc]
+  def __imul__(self, x) -> Tensor: return self.assign(self.mul(x)) # type: ignore[misc]
+  def __itruediv__(self, x) -> Tensor: return self.assign(self.div(x)) # type: ignore[misc]
+  def __iand__(self, x) -> Tensor: return self.assign(self.bitwise_and(x)) # type: ignore[misc]
+  def __ior__(self, x) -> Tensor: return self.assign(self.bitwise_or(x)) # type: ignore[misc]
+  def __ixor__(self, x) -> Tensor: return self.assign(self.bitwise_xor(x)) # type: ignore[misc]
+  def __ilshift__(self, x) -> Tensor: return self.assign(self.lshift(x)) # type: ignore[misc]
+  def __irshift__(self, x) -> Tensor: return self.assign(self.rshift(x)) # type: ignore[misc]

  def __lt__(self, x) -> Tensor: return self._apply_broadcasted_uop(UOp.__lt__, x, False)
  def __gt__(self, x) -> Tensor: return self._apply_broadcasted_uop(UOp.__lt__, x, True)
@@ -1,15 +1,17 @@
+from typing import TypeVar
 from tinygrad.uop import Ops
-from tinygrad.helpers import T
-from tinygrad.dtype import dtypes
+from tinygrad.dtype import dtypes, ConstType

+TMT = TypeVar("TMT", bound="MathTrait")
 class MathTrait:
  # required to implement
-  def alu(self:T, op:Ops, *src) -> T: raise NotImplementedError
-  def const_like(self:T, b) -> T: raise NotImplementedError
+  def alu(self:TMT, op:Ops, *src:TMT) -> TMT: raise NotImplementedError
+  def const_like(self:TMT, b:ConstType) -> TMT: raise NotImplementedError

  # great functions you get!
-  def ufix(self, x): return self.const_like(x) if not isinstance(x, MathTrait) else x
-  def _binop(self, op, x, reverse): return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x))
+  def ufix(self:TMT, x:TMT|ConstType) -> TMT: return self.const_like(x) if not isinstance(x, MathTrait) else x
+  def _binop(self:TMT, op:Ops, x:TMT|ConstType, reverse:bool) -> TMT:
+    return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x))
  def logical_not(self): return self.ne(True)
  def neg(self):
    if (dtype:=getattr(self, 'dtype')) is None: raise TypeError(f"MathTraits __neg__ requires a dtype, {self=}")
@@ -18,7 +20,7 @@ class MathTrait:
    if (dtype:=getattr(self, 'dtype')) is not None:
      if isinstance(dtype, tuple): dtype = dtype[0]
      if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)): raise RuntimeError(f"{dtype} is not supported")
-  def add(self, x, reverse=False):
+  def add(self:TMT, x:TMT|ConstType, reverse:bool=False):
    """
    Adds `self` and `x`.
    Equivalent to `self + x`.
@@ -36,7 +38,7 @@ class MathTrait:
    ```
    """
    return self._binop(Ops.ADD, x, reverse)
-  def mul(self, x, reverse=False):
+  def mul(self:TMT, x:TMT|ConstType, reverse:bool=False):
    """
    Multiplies `self` and `x`.
    Equivalent to `self * x`.
@@ -55,7 +57,7 @@ class MathTrait:
    ```
    """
    return self._binop(Ops.MUL, x, reverse)
-  def bitwise_and(self, x, reverse=False):
+  def bitwise_and(self:TMT, x:TMT|ConstType, reverse:bool=False):
    """
    Computes the bitwise AND of `self` and `x`.
    Equivalent to `self & x`.
@@ -69,7 +71,7 @@ class MathTrait:
    """
    self._check_dtype()
    return self._binop(Ops.AND, x, reverse)
-  def bitwise_or(self, x, reverse=False):
+  def bitwise_or(self:TMT, x:TMT|ConstType, reverse:bool=False):
    """
    Computes the bitwise OR of `self` and `x`.
    Equivalent to `self | x`.
@@ -83,7 +85,7 @@ class MathTrait:
    """
    self._check_dtype()
    return self._binop(Ops.OR, x, reverse)
-  def bitwise_xor(self, x, reverse=False):
+  def bitwise_xor(self:TMT, x:TMT|ConstType, reverse:bool=False):
    """
    Computes bitwise xor of `self` and `x`.
    Equivalent to `self ^ x`.
@@ -98,7 +100,7 @@ class MathTrait:
    """
    self._check_dtype()
    return self._binop(Ops.XOR, x, reverse)
-  def idiv(self, x, reverse=False):
+  def idiv(self:TMT, x:TMT|ConstType, reverse:bool=False):
    """
    Divides `self` by `x`.
    Equivalent to `self // x`.
@@ -110,61 +112,61 @@ class MathTrait:
    ```
    """
    return self._binop(Ops.IDIV, x, reverse)
-  def mod(self, x, reverse=False): return self._binop(Ops.MOD, x, reverse)
-  def sub(self, x, reverse=False): return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, self.ufix(-x))
-  def div(self, x, reverse=False): return (self.ufix(x)*self.alu(Ops.RECIP)) if reverse else (self*self.ufix(x).alu(Ops.RECIP))
+  def mod(self:TMT, x:TMT|ConstType, reverse:bool=False): return self._binop(Ops.MOD, x, reverse)
+  def sub(self:TMT, x:TMT|ConstType, reverse:bool=False): return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, self.ufix(-x))
+  def div(self:TMT, x:TMT|ConstType, reverse:bool=False): return (self.ufix(x)*self.alu(Ops.RECIP)) if reverse else (self*self.ufix(x).alu(Ops.RECIP))

  def __neg__(self): return self.neg()

-  def __add__(self, x): return self.add(x)
-  def __sub__(self, x): return self.sub(x)
-  def __mul__(self, x): return self.mul(x)
-  def __truediv__(self, x): return self.div(x)
-  def __floordiv__(self, x): return self.idiv(x)  # TODO: idiv is trunc div, not floordiv
-  def __mod__(self, x): return self.mod(x)
-  def __and__(self, x): return self.bitwise_and(x)
-  def __or__(self, x): return self.bitwise_or(x)
-  def __xor__(self, x): return self.bitwise_xor(x)
+  def __add__(self:TMT, x:TMT|ConstType): return self.add(x)
+  def __sub__(self:TMT, x:TMT|ConstType): return self.sub(x)
+  def __mul__(self:TMT, x:TMT|ConstType): return self.mul(x)
+  def __truediv__(self:TMT, x:TMT|ConstType): return self.div(x)
+  def __floordiv__(self:TMT, x:TMT|ConstType): return self.idiv(x)  # TODO: idiv is trunc div, not floordiv
+  def __mod__(self:TMT, x:TMT|ConstType): return self.mod(x)
+  def __and__(self:TMT, x:TMT|ConstType): return self.bitwise_and(x)
+  def __or__(self:TMT, x:TMT|ConstType): return self.bitwise_or(x)
+  def __xor__(self:TMT, x:TMT|ConstType): return self.bitwise_xor(x)

-  def __radd__(self, x): return self.add(x, True)
-  def __rsub__(self, x): return self.sub(x, True)
-  def __rmul__(self, x): return self.mul(x, True)
-  def __rtruediv__(self, x): return self.div(x, True)
-  def __rfloordiv__(self, x): return self.idiv(x, True)
-  def __rand__(self, x): return self.bitwise_and(x, True)
-  def __ror__(self, x): return self.bitwise_or(x, True)
-  def __rxor__(self, x): return self.bitwise_xor(x, True)
-  def __rmod__(self, x): return self.mod(x, True)
+  def __radd__(self:TMT, x:TMT|ConstType): return self.add(x, True)
+  def __rsub__(self:TMT, x:TMT|ConstType): return self.sub(x, True)
+  def __rmul__(self:TMT, x:TMT|ConstType): return self.mul(x, True)
+  def __rtruediv__(self:TMT, x:TMT|ConstType): return self.div(x, True)
+  def __rfloordiv__(self:TMT, x:TMT|ConstType): return self.idiv(x, True)
+  def __rand__(self:TMT, x:TMT|ConstType): return self.bitwise_and(x, True)
+  def __ror__(self:TMT, x:TMT|ConstType): return self.bitwise_or(x, True)
+  def __rxor__(self:TMT, x:TMT|ConstType): return self.bitwise_xor(x, True)
+  def __rmod__(self:TMT, x:TMT|ConstType): return self.mod(x, True)

-  def __lt__(self, x): return self.alu(Ops.CMPLT, self.ufix(x))
-  def __gt__(self, x): return self.ufix(x).alu(Ops.CMPLT, self)
-  def __ge__(self, x): return (self < x).logical_not()
-  def __le__(self, x): return (self > x).logical_not()
+  def __lt__(self:TMT, x:TMT|ConstType): return self.alu(Ops.CMPLT, self.ufix(x))
+  def __gt__(self:TMT, x:TMT|ConstType): return self.ufix(x).alu(Ops.CMPLT, self)
+  def __ge__(self:TMT, x:TMT|ConstType): return (self < x).logical_not()
+  def __le__(self:TMT, x:TMT|ConstType): return (self > x).logical_not()

-  def ne(self, x): return self.alu(Ops.CMPNE, self.ufix(x))
-  def eq(self, x): return self.ne(x).logical_not()
-  def __ne__(self, x): return self.ne(x)
+  def ne(self:TMT, x:TMT|ConstType): return self.alu(Ops.CMPNE, self.ufix(x))
+  def eq(self:TMT, x:TMT|ConstType): return self.ne(x).logical_not()
+  def __ne__(self:TMT, x:TMT|ConstType): return self.ne(x)  # type: ignore[override]
  # NOTE: __eq__ isn't overridden, and means the same thing as is by default

-  def lshift(self, x, reverse=False): return self._binop(Ops.SHL, x, reverse)
-  def rshift(self, x, reverse=False): return self._binop(Ops.SHR, x, reverse)
-  def __lshift__(self, x): return self.lshift(x)
-  def __rshift__(self, x): return self.rshift(x)
-  def __rlshift__(self, x): return self.lshift(x, True)
-  def __rrshift__(self, x): return self.rshift(x, True)
+  def lshift(self:TMT, x:TMT|int, reverse:bool=False): return self._binop(Ops.SHL, x, reverse)
+  def rshift(self:TMT, x:TMT|int, reverse:bool=False): return self._binop(Ops.SHR, x, reverse)
+  def __lshift__(self:TMT, x:TMT|int): return self.lshift(x)
+  def __rshift__(self:TMT, x:TMT|int): return self.rshift(x)
+  def __rlshift__(self:TMT, x:TMT|int): return self.lshift(x, True)
+  def __rrshift__(self:TMT, x:TMT|int): return self.rshift(x, True)

-  def maximum(self, x): return self.alu(Ops.MAX, self.ufix(x))
-  def minimum(self, x): return -(-self).maximum(-x)
-  def where(self, x, y):
-    if type(self) is type(x): return self.alu(Ops.WHERE, x, x.ufix(y))
-    if type(self) is type(y): return self.alu(Ops.WHERE, y.ufix(x), y)
+  def maximum(self:TMT, x:TMT|ConstType): return self.alu(Ops.MAX, self.ufix(x))
+  def minimum(self:TMT, x:TMT|ConstType): return -(-self).maximum(-x)
+  def where(self:TMT, x:TMT|ConstType, y:TMT|ConstType):
+    if isinstance(x, type(self)): return self.alu(Ops.WHERE, x, x.ufix(y))
+    if isinstance(y, type(self)): return self.alu(Ops.WHERE, y.ufix(x), y)
    raise RuntimeError("where needs at least one UOp arg")
-  def threefry(self, seed): return self.alu(Ops.THREEFRY, seed)
+  def threefry(self:TMT, seed:TMT): 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)
  def exp2(self): return self.alu(Ops.EXP2)
-  def pow(self, x): return self.alu(Ops.POW, self.ufix(x))
-  def __pow__(self, x): return self.pow(x)
+  def pow(self:TMT, x:TMT|ConstType): return self.alu(Ops.POW, self.ufix(x))
+  def __pow__(self:TMT, x:TMT|ConstType): return self.pow(x)
@@ -15,6 +15,7 @@ if TYPE_CHECKING:

 class AxisType(Enum):
  def __repr__(self): return str(self)
+  OUTER = auto()
  GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
  THREAD = auto()

@@ -34,11 +35,11 @@ def resolve(x:UOp|bool, default:bool=True):
 def _suop(lst, uop_fxn, python_fxn):
  uops, nums = partition(lst, lambda x: isinstance(x, UOp))
  return ssimplify(functools.reduce(uop_fxn, uops + ([python_fxn(nums)] if nums else [])))
-def smax(*lst): return _suop(argfix(*lst), UOp.maximum, max)
-def smin(*lst): return _suop(argfix(*lst), UOp.minimum, min)
-def srender(x) -> str: return x.render() if isinstance(x, UOp) else str(x)
+def smax(*lst) -> sint: return _suop(argfix(*lst), UOp.maximum, max)
+def smin(*lst) -> sint: return _suop(argfix(*lst), UOp.minimum, min)
+def srender(x:sint) -> str: return x.render() if isinstance(x, UOp) else str(x)

-def ssimplify(uop): return uop.ssimplify() if isinstance(uop, UOp) else uop
+def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop
 def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop

 def range_str(u:UOp) -> str: return '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])
@@ -240,6 +241,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
        if s in ret: del ret[s]
    else:
      for s in self.src: ret.update(s.ranges)
+    if self.op is Ops.ENDRANGE: del ret[self.src[0]]
    return ret

  @property
@@ -249,11 +251,11 @@ class UOp(MathTrait, metaclass=UOpMetaClass):

  # *** uop evaluation ***

-  def simplify(self, tracked=False):
+  def simplify(self, tracked=False, full_symbolic=True):
    # late import!
-    from tinygrad.uop.symbolic import symbolic
+    from tinygrad.uop.symbolic import symbolic, commutative
    with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value):
-      return graph_rewrite(self, symbolic, name="simplify")
+      return graph_rewrite(self, symbolic if full_symbolic else commutative, name="simplify")
  def ssimplify(self) -> UOp|ConstType: return ret.arg if (ret:=self.simplify()).op is Ops.CONST else ret
  def _eval(self, dtype, expected_type:Type[T]) -> T:
    assert self.dtype in dtype, f"eval with wrong dtype {self}"
@@ -1098,6 +1100,8 @@ pm_lower_index_dtype = PatternMatcher([
  (UPat((Ops.STORE, Ops.LOAD), src=(UPat(), UPat(), UPat().cast(dtypes.index)), allow_any_len=True, name="s"),
    lambda s: s.replace(src=s.src[:2]+tuple(u.src[0] for u in s.src[2:]))),
  (UPat((Ops.SINK, Ops.NOOP), src=UPat().cast(dtypes.index), name="n"), lambda n: n.replace(src=tuple(s.src[0] for s in n.src))),
+  # hack for ENDRANGE
+  (UPat(Ops.ENDRANGE, src=(UPat(Ops.RANGE, name="r").cast(dtypes.index),), allow_any_len=True, name="x"), lambda x,r: x.replace(src=(r,)+x.src[1:])),
 ])
 def _index_to_concrete_int(u:UOp): return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]

@@ -109,6 +109,9 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
  (UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"), UPat(Ops.DEVICE))), lambda red,x: red.dtype == x.dtype and isinstance(red.arg, Ops)),
  (UPat(Ops.MULTI, name="multi"), lambda multi: all(x.dtype == multi.dtype for x in multi.src) and isinstance(multi.arg, int)),

+  # endrange/reduce for outerworld range work
+  (UPat(Ops.ENDRANGE, src=(UPat(Ops.RANGE),), allow_any_len=True), lambda: True),
+
  # REDUCE with an outerworld range
  (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])),
 ])
@@ -177,8 +177,8 @@ def fold_binary_numerator(d: UOp, x: UOp, y: UOp) -> UOp|None:
  x,const = x.pop_const()
  terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in x.split_uop(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
+    y1 = cmod(factors[0]*v.vmin+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmin+const, c)
+    y2 = cmod(factors[0]*v.vmax+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmax+const, c)
    return (y2-y1)*(v-v.vmin) + y1
  return None

@@ -397,7 +397,7 @@ def parse_valid(valid:UOp) -> tuple[UOp, bool, int]:
  if valid.op is Ops.CMPLT and dtypes.is_int(valid.src[0].dtype): return valid.src[0], True, int((valid.src[1]).vmax)-1
  raise ValueError(f"not able to parse {valid=}")

-def uop_given_valid(valid:UOp, uop:UOp) -> UOp:
+def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
  # return simplified uop (might be the same as input)

  # first, parse valid into {expr: (lower_bound, upper_bound)}
@@ -411,26 +411,33 @@ def uop_given_valid(valid:UOp, uop:UOp) -> UOp:
  uop = uop.substitute((load_subs:={u: UOp(Ops.NOOP, arg=u) for u in uop.toposort() if u.op is Ops.INDEX}))

  # simplify uop given that valid is True
-  for expr,v in bounds.items():
+  all_candidates = []
+  for i,(expr,v) in enumerate(bounds.items()):
    v0, v1 = (expr.vmin if v[0] is None else v[0], expr.vmax if v[1] is None else v[1])
    expr = expr.substitute(load_subs)  # make sure expr appears in same form in the uop
-    # some expr has lower bound > upper bound -> valid is an empty set and we return None
-    # every candidate is a set of constrained UOp based on valid, and if every item in a set simplifies the uop into a same output, we rewrite uop
-    candidates = []
-    if expr.op is Ops.ADD and v0 == 1 and all(u.op in GroupOp.Irreducible for u in expr.split_uop(Ops.ADD)):
-      # if the constraint is a simplex: X0 + X1 + ... > 0, we can check if all Xi > 0 simplify into the same output
-      candidates.append([(Xi, UOp.variable("fake", 1, Xi.vmax, Xi.dtype)) for Xi in expr.split_uop(Ops.ADD)])
    # try checking the whole clause
-    candidates.append([(expr, UOp.variable("fake", v0, v1, expr.dtype))])
+    all_candidates.append((expr, UOp.variable(f"fake{i}", v0, v1, expr.dtype)))

-    for candidate in candidates:
-      # if every branch in candidate gives the same simplified uop, we can rewrite the uop
-      newuops = [uop.substitute({X:newX}).simplify().substitute({newX:X}).simplify() for X,newX in candidate]
-      if uop.op is Ops.VECTORIZE and len(uop.src) == 2:
-        if all_same([uops.src[0] for uops in newuops]): uop = uop.replace(src=(newuops[0].src[0], uop.src[1]))
-        if all_same([uops.src[1] for uops in newuops]): uop = uop.replace(src=(uop.src[0], newuops[0].src[1]))
-      elif all_same(newuops): uop = newuops[0]
+    if try_simplex:
+      # every candidate is a set of constrained UOp based on valid, and if every item in a set simplifies the uop into a same output, we rewrite uop
+      candidates = [[all_candidates[-1]]]
+      if expr.op is Ops.ADD and v0 == 1 and all(u.op in GroupOp.Irreducible for u in expr.split_uop(Ops.ADD)):
+        # if the constraint is a simplex: X0 + X1 + ... > 0, we can check if all Xi > 0 simplify into the same output
+        candidates.append([(Xi, UOp.variable(f"fake{i}", 1, Xi.vmax, Xi.dtype)) for Xi in expr.split_uop(Ops.ADD)])

+      for candidate in candidates:
+        # if every branch in candidate gives the same simplified uop, we can rewrite the uop
+        newuops = [uop.substitute({X:newX}) for X,newX in candidate]
+        if any(u is uop for u in newuops): continue  # if any branch doesnt appear in uop, skip
+        newuops = [u.simplify().substitute({newX:X}).simplify(full_symbolic=False) for (X,newX),u in zip(candidate,newuops)]
+        if uop.op is Ops.VECTORIZE and len(uop.src) == 2:
+          if all_same([uops.src[0] for uops in newuops]): uop = uop.replace(src=(newuops[0].src[0], uop.src[1]))
+          if all_same([uops.src[1] for uops in newuops]): uop = uop.replace(src=(uop.src[0], newuops[0].src[1]))
+        elif all_same(newuops): uop = newuops[0]
+
+  # try all the valids together (but only the whole expressions)
+  if (s_uop:=uop.substitute(sub_dict:=dict(all_candidates))) is not uop:
+    uop = s_uop.simplify().substitute({newX:X for X,newX in sub_dict.items()}).simplify(full_symbolic=False)
  # put the loads back in
  uop = uop.substitute({v:k for k,v in load_subs.items()})
  return uop
@@ -463,14 +470,21 @@ def reduce_mul_chain(r:UOp):
  if len(outside) == 0: return None
  return r.replace(src=(prod(inside) if len(inside) else r.src[0].const_like(1),)+r.src[1:])*prod(outside)

+def drop_and_clauses(cond:UOp, x:UOp, i:UOp) -> UOp|None:
+  if not (dropped_clauses:=[c for c in cond.split_uop(Ops.AND) if not any(r in x.ranges for r in c.ranges)]): return None
+  return functools.reduce(operator.and_, [c for c in cond.split_uop(Ops.AND) if c not in dropped_clauses], UOp.const(dtypes.bool, True)).where(x, i)
+pm_drop_and_clauses = PatternMatcher([(UPat.var("cond").where(UPat.var("x", dtype=dtypes.index), invalid_pat), drop_and_clauses)])
+
+pm_simplify_valid = PatternMatcher([
+  # simplify valid
+  (UPat(Ops.AND, name="valid"), simplify_valid),
+  (UPat.var("c").where(UPat.var("x", dtype=dtypes.index), invalid_pat), lambda c,x,i: c.where(uop_given_valid(c, x, try_simplex=False), i)),
+])
+
 # this is symbolic 2.0
 REMOVE_FROM_SINK = {Ops.SINK, Ops.UNROLL, Ops.PTRCAT, Ops.CAT, Ops.NOOP}
 REMOVE_FROM_BARRIER = {Ops.VECTORIZE, Ops.SINK, Ops.CAT, Ops.PTRCAT, Ops.NOOP}
-sym = symbolic_flat+PatternMatcher([
-  # simplify valid
-  (UPat(Ops.AND, name="valid"), simplify_valid),
-  (UPat.var("cond").where(UPat.var("x", dtype=dtypes.index), invalid_pat), lambda cond,x,i: cond.where(newx, i) if
-    (newx:=uop_given_valid(cond, x)) is not x else None),
+sym = symbolic_flat+pm_simplify_valid+PatternMatcher([
  # LOAD/STORE -> NOOP
  (UPat.var('x').store(UPat.var('x').load(), allow_any_len=True), lambda x: None if x.dtype.addrspace != AddrSpace.REG else x.src[0].src[0]),
  (UPat(Ops.LOAD, src=(UPat.cvar('c'))), lambda c: c),
@@ -505,8 +519,8 @@ sym = symbolic_flat+PatternMatcher([
  (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"),
    lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0
  # # Where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer
-  (UPat.var("c1").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("c2").where(UPat(), invalid_pat)).or_casted(),), allow_any_len=True, name="l"), 0),
-    lambda c1,c2,l,i: l.replace(src=(l.src[0],)+l.src[1:]) if any(c in list(c2.split_uop(Ops.AND)) for c in c1.split_uop(Ops.AND)) else None),
+  (UPat.var("c1").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("c2").where(UPat(), invalid_pat)).or_casted(),), name="l"), 0),
+    lambda c1,c2,l,i: l.replace(src=(l.src[0],)+l.src[1:]) if all(c in list(c2.split_uop(Ops.AND)) for c in c1.split_uop(Ops.AND)) else None),
  # remove VECTORIZE from SINK/BARRIER. TODO: SINK/BARRIER are really the same thing at GLOBAL/LOCAL levels
  (UPat(Ops.BARRIER, name="root"),
    lambda root: UOp(Ops.BARRIER, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_BARRIER else (x,) for x in root.src)), root.arg)
@@ -18,6 +18,9 @@ const ANSI_COLORS_LIGHT = ["#d9d9d9","#ff9999","#99cc99","#ffff99","#9999ff","#f
 const parseColors = (name, defaultColor="#ffffff") => Array.from(name.matchAll(/(?:\u001b\[(\d+)m([\s\S]*?)\u001b\[0m)|([^\u001b]+)/g),
  ([_, code, colored_st, st]) => ({ st: colored_st ?? st, color: code != null ? (code>=90 ? ANSI_COLORS_LIGHT : ANSI_COLORS)[(parseInt(code)-30+60)%60] : defaultColor }));

+const colored = n => d3.create("span").call(s => s.selectAll("span").data(typeof n === "string" ? parseColors(n) : n).join("span")
+                       .style("color", d => d.color).text(d => d.st)).node();
+
 const rect = (s) => (typeof s === "string" ? document.querySelector(s) : s).getBoundingClientRect();

 let timeout = null;
@@ -174,7 +177,7 @@ function tabulate(rows) {
 var data, focusedDevice, focusedShape, canvasZoom, zoomLevel = d3.zoomIdentity;
 async function renderProfiler() {
  displayGraph("profiler");
-  d3.select(".metadata").html("");
+  d3.select(".metadata").node().replaceChildren(focusedShape?.html ?? "");
  // layout once!
  if (data != null) return updateProgress({ start:false });
  const profiler = d3.select(".profiler").html("");
@@ -236,8 +239,7 @@ async function renderProfiler() {
          const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
          if (stepIdx !== -1) { ref.step = stepIdx; shapeRef = ref; }
        }
-        const htmlLabel = label.map(({color, st}) => `<span style="color:${color}">${st}</span>`).join('');
-        const arg = { tooltipText:htmlLabel+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...shapeRef };
+        const arg = { tooltipText:colored(e.name).outerHTML+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...shapeRef };
        // offset y by depth
        shapes.push({x:e.st, y:levelHeight*depth, width:e.dur, height:levelHeight, arg, label, fillColor });
      }
@@ -257,8 +259,8 @@ async function renderProfiler() {
          x += 1; y += nbytes; valueMap.set(ts, y);
        } else {
          const free = buf_shapes.get(key);
-          timestamps.push(ts);
-          x += 1; y -= free.nbytes; valueMap.set(ts, y);
+          timestamps.push(ts); valueMap.set(ts, y);
+          x += 1; y -= free.nbytes;
          free.x.push(x);
          free.y.push(free.y.at(-1));
          temp.delete(key);
@@ -333,6 +335,7 @@ async function renderProfiler() {
    const st = visibleX[0], et = visibleX[1];
    xscale.domain(visibleX);
    // draw shapes
+    const paths = [];
    for (const [_, { offsetY, shapes, visible, valueMap }] of data.tracks) {
      visible.length = 0;
      for (const e of shapes) {
@@ -340,18 +343,18 @@ async function renderProfiler() {
        if (e.width == null) {
          if (e.x[0]>et || e.x.at(-1)<st) continue;
          const x = e.x.map(xscale);
-          ctx.beginPath();
-          ctx.moveTo(x[0], offsetY+e.y0[0]);
+          const p = new Path2D();
+          p.moveTo(x[0], offsetY+e.y0[0]);
          for (let i=1; i<x.length; i++) {
-            ctx.lineTo(x[i], offsetY+e.y0[i]);
+            p.lineTo(x[i], offsetY+e.y0[i]);
            let arg = e.arg;
            if (arg == null && valueMap != null) arg = {tooltipText: `Total: ${formatUnit(valueMap.get(e.x[i-1]), 'B')}`}
            visible.push({ x0:x[i-1], x1:x[i], y0:offsetY+e.y1[i-1], y1:offsetY+e.y0[i], arg });
          }
-          for (let i=x.length-1; i>=0; i--) ctx.lineTo(x[i], offsetY+e.y1[i]);
-          ctx.closePath();
-          ctx.fillStyle = e.fillColor; ctx.fill();
-          if (focusedShape && e.arg?.key === focusedShape) { ctx.lineWidth = 1.4; ctx.strokeStyle = "#c9a8ff"; ctx.stroke(); }
+          for (let i=x.length-1; i>=0; i--) p.lineTo(x[i], offsetY+e.y1[i]);
+          p.closePath();
+          ctx.fillStyle = e.fillColor; ctx.fill(p);
+          if (focusedShape && e.arg?.key === focusedShape.key) { paths.push(p); }
          continue;
        }
        // contiguous rect
@@ -400,11 +403,13 @@ async function renderProfiler() {
      }
    }
    // draw markers
+    ctx.textBaseline = "top";
    for (const m of markers) {
      const x = xscale(m.ts);
      drawLine(ctx, [x, x], [0, canvas.clientHeight], { color:m.color });
      ctx.fillText(m.name, x+2, 1);
    }
+    for (const p of paths) { ctx.lineWidth = 1.4; ctx.strokeStyle = "#c9a8ff"; ctx.stroke(p); }
  }

  function resize() {
@@ -445,7 +450,7 @@ async function renderProfiler() {
    e.preventDefault();
    const foundRect = findRectAtPosition(e.clientX, e.clientY);
    if (foundRect?.step != null) return setCtxWithHistory(foundRect.ctx, foundRect.step);
-    if (foundRect?.key != focusedShape) { focusedShape = foundRect?.key; render(zoomLevel); }
+    if (foundRect?.key != focusedShape?.key) { focusedShape = foundRect; render(zoomLevel); }
    return document.querySelector(".metadata").replaceChildren(foundRect?.html ?? "");
  });

@@ -589,7 +594,7 @@ async function main() {
      const ul = ctxList.appendChild(document.createElement("ul"));
      ul.id = `ctx-${i}`;
      const p = ul.appendChild(document.createElement("p"));
-      p.innerHTML = parseColors(name).map(c => `<span style="color: ${c.color}">${c.st}</span>`).join("");
+      p.appendChild(colored(name));
      p.onclick = () => {
        setState(i === state.currentCtx ? { expandSteps:!state.expandSteps } : { expandSteps:true, currentCtx:i, currentStep:0, currentRewrite:0 });
      }
@@ -703,9 +708,7 @@ async function main() {
        metadata.appendChild(codeBlock(upat[1], "python", { loc:upat[0], wrap:true }));
        const diffCode = metadata.appendChild(document.createElement("pre")).appendChild(document.createElement("code"));
        for (const line of diff) {
-          const span = diffCode.appendChild(document.createElement("span"));
-          span.style.color = line.startsWith("+") ? "#3aa56d" : line.startsWith("-") ? "#d14b4b" : "#f0f0f5";
-          span.innerText = line;
+          diffCode.appendChild(colored([{st:line, color:line.startsWith("+") ? "#3aa56d" : line.startsWith("-") ? "#d14b4b" : "#f0f0f5"}]));
          diffCode.appendChild(document.createElement("br"));
        }
        diffCode.className = "wrap";
Author	SHA1	Message	Date
geohot	30ff87eab4	realize sched	2025-10-14 14:44:56 +08:00
George HotzandGitHub	fe683bafa6	Merge branch 'master' into outerworld_work	2025-10-14 14:26:52 +08:00
George HotzandGitHub	b9eb5b5d49	clean up the LLM tokenizer (#12653 ) * clean up the LLM tokenizer * simple tokenizer is actually simple * ugh write good code	2025-10-14 14:22:01 +08:00
qazalandGitHub	a9ef93176f	viz: add colored text helper (#12654 )	2025-10-14 13:05:26 +08:00
George HotzandGitHub	ecdc7539a2	add typing to MathTraits (#12650 ) * add typing to MathTraits * fix assign	2025-10-14 12:35:20 +08:00
qazalandGitHub	9bf032de69	viz: keep focused shape in view (#12648 )	2025-10-14 10:49:08 +08:00
chenyuandGitHub	77b5e6774e	fix bert training config (#12647 ) FREE_INTERMEDIATE=0 REWRITE_STACK_LIMIT=500000	2025-10-13 15:03:47 -04:00
nimlgenandGitHub	f1041dc0ac	pylint 4.0.0 (#12642 ) * cpu: fix spacing * fix pylint * fix pylint * pylint 4.0.0 * lambda * keep eval for now * im so sorry	2025-10-13 23:28:36 +08:00
wozeparrotandGitHub	47e0c43976	feat: Tensor.{load, store} (#12629 )	2025-10-13 08:04:41 -07:00
chenyuandGitHub	0f776c6e46	examples/mlperf/training_submission_v6.0 (#12644 ) copied from v5.1	2025-10-13 09:58:25 -04:00
Sieds LyklesandGitHub	e0139fafc1	UOp symbolic tests use eval to check against string (#12643 )	2025-10-13 14:19:42 +02:00
b1tg GitHub wozeparrot	218225e8d0	pylint error (#12630 ) Co-authored-by: wozeparrot <[email protected]>	2025-10-13 05:05:12 -07:00
nimlgenandGitHub	9096d7cc2e	amd: support for rx9060 (#12640 )	2025-10-13 19:44:15 +08:00
qazalandGitHub	066d25f5fb	refactor to trace_num property in buffers (#12638 )	2025-10-13 18:06:55 +08:00
qazalandGitHub	cd6aeebfee	sqtt: osx decoder installer (#12637 )	2025-10-13 17:26:12 +08:00
Sieds LyklesandGitHub	e537e895b1	drop unused invalid conditions (#12635 ) * drop where conditions if the ranges are not used inside the index * remove allow_any_len	2025-10-13 10:52:21 +02:00
wozeparrotandGitHub	9ab06dffad	hotfix: block from env (#12628 )	2025-10-12 08:07:32 -07:00
wozeparrotandGitHub	12435a2dab	actual tinyfs device (#12620 )	2025-10-12 07:51:17 -07:00
chenyuandGitHub	8f5f57c7d9	smaller CNT fuzz shapetracker (#12626 )	2025-10-12 08:52:30 -04:00
George HotzandGitHub	1ecf403294	cleanup long lines [pr] (#12623 ) * cleanup long lines * more * a few more * all noqa fixed * fix amd + cuda * clean that up	2025-10-12 20:18:05 +08:00
qazalandGitHub	fd51ecf983	process_replay for get_rangeify_map (#12624 )	2025-10-12 15:14:40 +03:00
qazalandGitHub	b5afa3848e	viz: fix memory graph total nbytes (#12622 ) * viz: fix memory graph total nbytes * post increment * simple regression test * loop with markers + slightly off text baseline * cpu events clear	2025-10-12 14:32:46 +03:00
nimlgenandGitHub	822eab057f	cpu: respect taskset + allow all cores (#12619 ) * cpu: account taskset + allow all cores * spaces	2025-10-12 14:31:40 +08:00
chenyuandGitHub	7ac74d1550	remove unused type ignore [pr] (#12618 )	2025-10-11 21:24:04 -04:00
Sieds LyklesandGitHub	772a8dfe31	reshape uses valid when simplifying (#12597 ) * reshape uses valid when simplifying * try with IGNORE_OOB=0 * is it this test? * skipif gpuocelot	2025-10-11 17:02:54 +02:00
nimlgenandGitHub	08e62454b6	amd: use cpu_view() in sqtt (#12610 )	2025-10-11 18:11:25 +08:00
Sieds LyklesandGitHub	a2ae56674a	`uop_given_valid` try multiple clauses (#12615 ) * uop_given_valid uses less simplify * enable test * try all expressions together * enable test	2025-10-11 11:53:42 +02:00
Sieds LyklesandGitHub	dccdd190aa	uop_given_valid uses less simplify (#12612 ) * uop_given_valid uses less simplify * enable test	2025-10-11 10:57:39 +02:00
qazalandGitHub	9205527db0	viz: draw highlights above shapes (#12613 )	2025-10-11 11:39:13 +03:00
George HotzandGitHub	cab034b863	improve typing (#12611 ) * improve typing and bump to 3.11 * no need for Self yet * improve typing * binop also	2025-10-11 16:20:23 +08:00
geohot	ab9064c411	train loop	2025-10-10 20:20:19 +08:00
George HotzandGitHub	8832f08af3	Merge branch 'master' into outerworld_work	2025-10-10 20:07:44 +08:00
geohot	402e1cf48f	work	2025-10-10 19:49:24 +08:00
geohot	b2490b6e31	test assign/reduce	2025-10-10 18:50:16 +08:00
George HotzandGitHub	33e8babdd8	Merge branch 'master' into outerworld_work	2025-10-10 18:25:08 +08:00
geohot	67a409343d	work	2025-10-10 18:10:42 +08:00
geohot	5b24999a36	work on outerworld	2025-10-10 14:48:05 +08:00