might end if

replace linearizer with toposort
2026-08-16 02:38:26 +00:00 · 2025-10-10 16:40:30 +08:00 · 2025-10-10 16:02:28 +08:00
106 changed files with 857 additions and 1814 deletions
@@ -625,11 +625,11 @@ jobs:
    - name: benchmark openpilot 0.9.9 dmonitoring
      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 0.9.9 driving_vision
-      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=22 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.9.9 driving_policy
      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.9.9 dmonitoring
-      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=15 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 Space Lab policy + vision
      run: |
        PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
@@ -310,9 +310,9 @@ jobs:
    - name: Fuzz Test fast idiv
      run: python test/external/fuzz_fast_idiv.py
    - name: Fuzz Test shapetracker
-      run: CNT=50 python test/external/fuzz_shapetracker.py
-    - name: Fuzz Test shapetracker math
-      run: CNT=200 python test/external/fuzz_shapetracker_math.py
+      run: |
+        python test/external/fuzz_shapetracker.py
+        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=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
+          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
      - 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,6 +30,10 @@ 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, 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
+  sed -i "s/import ctypes/import ctypes, tinygrad.helpers.fetch as tgfetch/g" extra/sqtt/rocprof/rocprof.py
+  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(str(tgfetch('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
 }

 generate_webgpu() {
@@ -1188,9 +1188,7 @@ 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", 0) and train_step_bert.captured is not None:
-        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
-        train_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
      eval_lm_losses = []
      eval_clsf_losses = []
      eval_lm_accs = []
@@ -1224,7 +1222,7 @@ def train_bert():
          return

      if getenv("RESET_STEP"): eval_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE", 1) 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)
@@ -1,17 +0,0 @@
-#!/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
@@ -1,69 +0,0 @@
-# 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
-```
@@ -1,17 +0,0 @@
-#!/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
@@ -1,20 +0,0 @@
-#!/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
@@ -1,31 +0,0 @@
-#!/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
@@ -1,69 +0,0 @@
-# 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
-```
@@ -1,17 +0,0 @@
-#!/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
@@ -1,16 +0,0 @@
-#!/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
@@ -1,28 +0,0 @@
-#!/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
@@ -1,69 +0,0 @@
-# 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
-```
@@ -1,18 +0,0 @@
-#!/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
@@ -1,16 +0,0 @@
-#!/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
@@ -1,31 +0,0 @@
-#!/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
@@ -1,50 +0,0 @@
-# 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
-```
@@ -1,13 +0,0 @@
-#!/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
@@ -1,15 +0,0 @@
-#!/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
@@ -1,25 +0,0 @@
-#!/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
@@ -1,50 +0,0 @@
-# 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
-```
@@ -1,13 +0,0 @@
-#!/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
@@ -1,15 +0,0 @@
-#!/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
@@ -1,26 +0,0 @@
-#!/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
@@ -1,8 +0,0 @@
-#!/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
@@ -1,38 +0,0 @@
-# 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
-```
@@ -1,14 +0,0 @@
-#!/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
@@ -1,15 +0,0 @@
-#!/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
@@ -1,25 +0,0 @@
-#!/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
@@ -1,14 +0,0 @@
-#!/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
@@ -1,15 +0,0 @@
-#!/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
@@ -1,38 +0,0 @@
-{
-    "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": ""
-  }
@@ -1,38 +0,0 @@
-{
-  "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": ""
-}
@@ -1,37 +0,0 @@
-{
-  "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": ""
-}
@@ -156,9 +156,6 @@ class RGP:
    sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
    if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
    device_props = sqtt_events[0].props
-    gfx_ver = device_props['gfx_target_version'] // 10000
-    gfx_iplvl = getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}_{(device_props['gfx_target_version']//100)%100}",
-                        getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}", None))
    sqtt_itrace_enabled = any([event.itrace for event in sqtt_events])
    sqtt_itrace_masked = not all_same([event.itrace for event in sqtt_events])
    sqtt_itrace_se_mask = functools.reduce(lambda a,b: a|b, [int(event.itrace) << event.se for event in sqtt_events], 0) if sqtt_itrace_masked else 0
@@ -196,7 +193,7 @@ class RGP:
        flags=0,
        trace_shader_core_clock=0x93f05080,
        trace_memory_clock=0x4a723a40,
-        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550}[device_props['gfx_target_version']],
+        device_id={110000: 0x744c, 110003: 0x7480}[device_props['gfx_target_version']],
        device_revision_id=0xc8,
        vgprs_per_simd=1536,
        sgprs_per_simd=128*16,
@@ -210,7 +207,7 @@ class RGP:
        sgpr_alloc_granularity=128,
        hardware_contexts=8,
        gpu_type=sqtt.SQTT_GPU_TYPE_DISCRETE,
-        gfxip_level=gfx_iplvl,
+        gfxip_level=sqtt.SQTT_GFXIP_LEVEL_GFXIP_11_0,
        gpu_index=0,
        gds_size=0,
        gds_per_shader_engine=0,
@@ -261,7 +258,7 @@ class RGP:
            major_version=0, minor_version=2,
          ),
          shader_engine_index=sqtt_event.se,
-          sqtt_version={11: sqtt.SQTT_VERSION_3_2, 12: sqtt.SQTT_VERSION_3_3}.get(gfx_ver),
+          sqtt_version=sqtt.SQTT_VERSION_3_2,
          _0=sqtt.union_sqtt_file_chunk_sqtt_desc_0(
            v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_0_v1(
              instrumentation_spec_version=1,
@@ -20,15 +20,14 @@ class InstInfo:
 class _ROCParseCtx:
  def __init__(self, sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
    self.sqtt_evs, self.prog_evs = iter(sqtt_evs), prog_evs
-    self.wave_events, self.disasms, self.addr2prg = {}, {}, {}
-
-    for prog in prog_evs:
-      for addr, info in comgr_get_address_table(prog.lib).items():
-        self.disasms[prog.base + addr] = info
-        self.addr2prg[prog.base + addr] = prog
+    self.wave_events = {}

  def next_sqtt(self): return next(self.sqtt_evs, None)
-  def find_program(self, addr): return self.addr2prg[addr]
+  def find_program(self, idx): return self.prog_evs[idx]
+  def get_instr_info(self, idx, exec_addr): return self.disasm_program(idx)[exec_addr - self.find_program(idx).base]
+
+  @functools.lru_cache(None)
+  def disasm_program(self, idx): return comgr_get_address_table(self.find_program(idx).lib)

  def on_occupancy_ev(self, ev):
    if DEBUG >= 4: print("OCC", ev.time, ev.cu, ev.simd, ev.wave_id, ev.start)
@@ -40,10 +39,10 @@ class _ROCParseCtx:
    for j in range(ev.instructions_size):
      inst_ev = ev.instructions_array[j]
      inst_typ = rocprof.rocprofiler_thread_trace_decoder_inst_category_t__enumvalues[inst_ev.category]
-      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=self.disasms[inst_ev.pc.address][0]))
+      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=self.get_instr_info(inst_ev.pc.code_object_id, inst_ev.pc.address)[0]))
      asm[inst_ev.pc.address].on_ev(inst_ev)

-    self.wave_events[(self.find_program(ev.instructions_array[0].pc.address).name, ev.wave_id, ev.cu, ev.simd)] = asm
+    self.wave_events[(self.find_program(ev.instructions_array[0].pc.code_object_id).name, ev.wave_id, ev.cu, ev.simd)] = asm

 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
@@ -79,7 +78,7 @@ if __name__ == "__main__":

  @rocprof.rocprof_trace_decoder_isa_callback_t
  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
-    instr, mem_size_ptr[0] = ROCParseCtx.disasms[pc.address]
+    instr, mem_size_ptr[0] = ROCParseCtx.get_instr_info(pc.code_object_id, pc.address)

    # this is the number of bytes to next instruction, set to 0 for end_pgm
    if instr == "s_endpgm": mem_size_ptr[0] = 0
@@ -1,18 +0,0 @@
-#!/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,7 +7,7 @@
 # POINTER_SIZE is: 8
 # LONGDOUBLE_SIZE is: 16
 #
-import ctypes, ctypes.util
+import ctypes, tinygrad.helpers.fetch as tgfetch


 class AsDictMixin:
@@ -155,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(ctypes.util.find_library('rocprof-trace-decoder')) #  ctypes.CDLL('FIXME_STUB')
+_libraries['FIXME_STUB'] = ctypes.CDLL(str(tgfetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so'))) #  ctypes.CDLL('FIXME_STUB')



@@ -43,7 +43,6 @@ enum sqtt_version
   SQTT_VERSION_2_3 = 0x6, /* GFX9 */
   SQTT_VERSION_2_4 = 0x7, /* GFX10+ */
   SQTT_VERSION_3_2 = 0xb, /* GFX11+ */
-   SQTT_VERSION_3_3 = 0xc, /* GFX12+ */
 };

 enum sqtt_file_chunk_type
@@ -145,8 +144,6 @@ enum sqtt_gfxip_level
   SQTT_GFXIP_LEVEL_GFXIP_10_1 = 0x7,
   SQTT_GFXIP_LEVEL_GFXIP_10_3 = 0x9,
   SQTT_GFXIP_LEVEL_GFXIP_11_0 = 0xc,
-   SQTT_GFXIP_LEVEL_GFXIP_11_5 = 0xd,
-   SQTT_GFXIP_LEVEL_GFXIP_12 = 0x10,
 };

 enum sqtt_memory_type
@@ -430,8 +427,6 @@ enum elf_gfxip_level
   EF_AMDGPU_MACH_AMDGCN_GFX1010 = 0x033,
   EF_AMDGPU_MACH_AMDGCN_GFX1030 = 0x036,
   EF_AMDGPU_MACH_AMDGCN_GFX1100 = 0x041,
-   EF_AMDGPU_MACH_AMDGCN_GFX1150 = 0x043,
-   EF_AMDGPU_MACH_AMDGCN_GFX1200 = 0x04e,
 };

 struct sqtt_file_chunk_spm_db {
@@ -50,7 +50,4 @@ exclude = [
  "E303", "E304", "E501", "E702", "E703", "E731", "W191",
  "W291", "W293", "UP039", "C416", "RET506", "RET507", "A",
  "FURB110", "RUF018", "F541", "F841"
-]
-
-[format]
-exclude = ["*"]
+]
@@ -1,50 +1,41 @@
-import functools, multiprocessing
 from transformers import AutoTokenizer
 from datasets import load_dataset
-from tinygrad.apps.llm import SimpleTokenizer
+from tinygrad.apps.llm import SimpleTokenizer, gpt2_decode_vocab, get_llama_re
 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__":
-  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)}")
+  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"
+
+  ds = load_dataset("OpenAssistant/oasst1")
  allow_failed = getenv("ALLOW_FAILED", 10)
+
  fail_count, total = 0, 0
-  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.")
+
+  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.")
@@ -2,7 +2,6 @@ import gc
 from tinygrad import Tensor, UOp, Device, nn
 from tinygrad.shape.shapetracker import views_to_valid_uop
 from tinygrad.engine.realize import method_cache, get_program
-from tinygrad.schedule.indexing import apply_movement_op
 from test.test_tiny import TestTiny

 def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
@@ -70,7 +69,6 @@ if __name__ == "__main__":
    # these caches will keep uops alive
    method_cache.clear()
    views_to_valid_uop.cache_clear()
-    apply_movement_op.cache_clear()
    Tensor._device_seeds.clear()
    Tensor._device_rng_counters.clear()

@@ -42,13 +42,13 @@ class ProcessReplayWarning(Warning): pass

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

-def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[Any, ...]]:
+def replay_kernelize(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(big_sink.substitute(ret)), (big_sink,)
+  return to_str(new_sink), to_str(ret[big_sink]), (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_rangeify_map":replay_get_rangeify_map, "get_program":replay_get_program}
+replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {"get_kernelize_map":replay_kernelize, "get_program":replay_get_program}

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

@@ -8,11 +8,9 @@ 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, getenv
+from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT
 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):
@@ -70,8 +68,7 @@ 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]])
-    # 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 not 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):
@@ -317,7 +314,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(MOCKGPU and isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CUDARenderer)), "PTX indexes differently. might be ok?")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "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,6 +390,7 @@ class TestMultiTensor(unittest.TestCase):

  # NOTE: this is failing on LLVM CI, no idea why. Works locally.
  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "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

@@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, AMD_LLVM
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, CPU_LLVM, AMD_LLVM
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -901,6 +901,7 @@ class TestOps(unittest.TestCase):
  def test_abs_exact(self):
    helper_test_op(None, torch.abs, Tensor.abs, vals=[[-1.,0,1]])

+  @unittest.skipIf(TRANSCENDENTAL and Device.DEFAULT=="AMD", "TODO: remu crashes")
  def test_log(self):
    helper_test_op([(45,65)], torch.log, Tensor.log)
    helper_test_op(None, torch.log, Tensor.log, vals=[[math.inf, -math.inf, math.nan]])
@@ -910,6 +911,7 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, torch.log2, Tensor.log2, vals=[[math.inf, -math.inf, math.nan]])
    helper_test_op([()], torch.log2, Tensor.log2)

+  @unittest.skipIf(TRANSCENDENTAL and Device.DEFAULT=="AMD", "TODO: remu crashes")
  def test_exp(self):
    helper_test_op([(45,65)], torch.exp, Tensor.exp)
    helper_test_op(None, torch.exp, Tensor.exp, vals=[[math.inf, -math.inf, math.nan]])
@@ -1547,6 +1549,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([(3,4,5,6)], lambda x: torch.stack(torch.std_mean(x, axis=(1,2))),
                                lambda x: Tensor.stack(*x.std_mean(axis=(1,2))))

+  @unittest.skip("TODO: this fails because of loaded nan in mul folding")
  def test_std_mean_loaded_nan(self):
    helper_test_op([(1,0,3,0,5)], lambda x: torch.stack(torch.std_mean(x, axis=(1,3))),
                                  lambda x: Tensor.stack(*x.std_mean(axis=(1,3))))
@@ -1,71 +1,5 @@
 import unittest
-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()
-    a = UOp.range(5, -1, AxisType.REDUCE)
-    out = x[a]
-    # TODO: syntax for this
-    t = Tensor(UOp(Ops.REDUCE, dtype=out.uop.dtype, src=(out.uop, a), arg=Ops.ADD))
-    self.assertListEqual(t.tolist(), [5.,5.,5.,5.,5.])
+from tinygrad import Tensor, UOp, GlobalCounters, Context

 class TestOuterworld(unittest.TestCase):
  def test_range_plus_1(self):
@@ -79,24 +13,12 @@ class TestOuterworld(unittest.TestCase):

    self.assertTrue((t+1==cpy).all().item())

-  def test_range_plus_1_transpose(self):
-    t = Tensor.arange(100).reshape(10,10).realize()
-
-    # passthrough ranges
-    a = UOp.range(10, -1)
-    sel = t[a] + 1
-    assert sel.shape == (10,)
-    cpy = sel.reshape(10, 1).expand(10, a).contiguous().realize()
-
-    self.assertTrue(((t+1).T==cpy).all().item())
-
  def test_flip_range(self):
    t = Tensor.rand(10, 10).realize()

    # 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())
@@ -115,17 +37,39 @@ class TestOuterworld(unittest.TestCase):
    out.realize()
    self.assertTrue((out==20).all().item())

-  def test_fancy_vmap(self):
-    def f(x,y): return x+y
+  @unittest.skip("opts don't work")
+  def test_triple_gemm(self):
+    x = Tensor.rand(1, 16).realize()
+    W = Tensor.rand(3, 16, 16).realize()

-    x = Tensor.arange(9).reshape(3,3).contiguous()
-    y = Tensor.arange(9).reshape(3,3).contiguous()
+    manual = (x @ W[0] @ W[1] @ W[2]).contiguous().realize()

    a = UOp.range(3, -1)
-    out = f(x[:,a], y[a,:])
-    # TODO: this should support flatten
-    out = out.reshape(1, 3).expand(a, 3).contiguous().realize()
-    self.assertListEqual([[0,4,8],[4,8,12],[8,12,16]], out.tolist())
+    x = x.assign(x @ W[a])
+    out = x.contiguous(a)[-1].contiguous().realize()
+
+    self.assertTrue((manual==out).all().item())
+
+  def test_setitem_pyrange(self):
+    with Context(DEBUG=0):
+      t = Tensor.rand(10).realize()
+      o = Tensor.empty(10)
+    GlobalCounters.reset()
+    for i in range(10):
+      o[i] = t[i]
+    o.realize()
+    self.assertTrue((t==o).all().item())
+
+  @unittest.skip("TODO: fix this")
+  def test_setitem(self):
+    with Context(DEBUG=0):
+      t = Tensor.rand(10).realize()
+      o = Tensor.empty(10)
+    GlobalCounters.reset()
+    i = UOp.range(10, -1)
+    o[i] = t[i]
+    o.contiguous(i).realize()
+    self.assertTrue((t==o).all().item())

 if __name__ == '__main__':
  unittest.main()
@@ -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-5)
+                                                   b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-6)

  def test_shrink_pad_safe(self):
    a = Tensor.ones((3, )).contiguous().realize()
@@ -1,8 +1,9 @@
 import ctypes, gzip, unittest, timeit
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI, mv_address, get_contraction
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI, mv_address
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
 from tinygrad.tensor import Tensor, get_shape
+from tinygrad.shape.view import get_contraction
 import numpy as np

 VARIABLE = ContextVar("VARIABLE", 0)
@@ -1,21 +1,19 @@
 import unittest, base64, functools, sys
-from tinygrad.apps.llm import SimpleTokenizer
+from tinygrad.apps.llm import SimpleTokenizer, get_llama_re
 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]
-
-      # 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}
+      str_vocab = [ line.split(maxsplit=1) for line in fd.read().splitlines() if line ]
+      normal_tokens = { base64.b64decode(stok): int(srank) for stok, srank in str_vocab }

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

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

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

@@ -154,6 +154,21 @@ class TestRealStrides(unittest.TestCase):
    ))
    self.assertEqual(st.is_expanded(), (False, False, False, True, False))

+class TestRealSimplifies(unittest.TestCase):
+  def tearDown(self):
+    self.st = self.st.simplify()
+    assert len(self.st.views) == 1
+
+  def test_1(self):
+    self.st = ShapeTracker((
+      View.create((1, 3, 2, 11, 4, 28), (0, 308, 0, 28, 0, 1), 0, None),
+      View.create((1, 3, 2, 11, 26, 1, 1, 3), (0, 2464, 0, 112, 1, 0, 0, 29), 0, None)))
+
+  def test_2(self):
+    self.st = ShapeTracker((
+      View.create((8, 3, 3, 11, 2, 28), (924, 308, 0, 28, 0, 1), 0, None),
+      View.create((8, 1, 6, 10, 28, 3, 2, 1), (5544, 0, 0, 56, 1, 1848, 672, 0), 0, None)))
+
 class TestIndexExpressions2d(unittest.TestCase):
  def setUp(self):
    shapes = [(30, 5), (15, 10), (15, 1), (5, 10), (5, 1)] # Make sure dim0 is a multiple of 5, one of the tests divides this dimension by 5
@@ -62,7 +62,6 @@ class TestShapeTrackerAdd(unittest.TestCase):
    b = ShapeTracker.from_shape((100,))
    assert a+b == b

-  @unittest.skip("no longer simplifies")
  def test_simple_add_permute(self):
    a = ShapeTracker.from_shape((10, 10))
    a = a.permute((1,0))
@@ -5,7 +5,6 @@ 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, (
@@ -50,8 +49,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]
-    check_uop_against_string(self, idx, sidx)
-    check_uop_against_string(self, valid, svalid)
+    self.assertEqual(idx.render(simplify=False), sidx)
+    self.assertEqual(valid.render(simplify=False), svalid)

  def test_cumsum(self):
    gidx0 = Special("gidx0", 5)
@@ -187,13 +186,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)
-    # TODO: this should also work without the extra nesting
-    idx = (((ridx0+ridx1)+(ridx2+ridx3)+28)//30)
+    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,
@@ -219,10 +218,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]
-    check_uop_against_string(self, idx0, sidx0)
-    check_uop_against_string(self, idx1, sidx1)
+    self.assertEqual(idx0.render(simplify=False), sidx0)
+    self.assertEqual(idx1.render(simplify=False), sidx1)
    if svalid is not None:
-      check_uop_against_string(self, load.src[0].src[2], svalid)
+      self.assertEqual(load.src[0].src[2].render(simplify=False), svalid)
    else:
      self.assertEqual(len(load.src[0].src), 2, "svalid is None but load still has a valid")

@@ -28,7 +28,7 @@ class TestSymbolic(unittest.TestCase):
  def test_merge_view_recursion_err(self):
    vm2 = View(shape=(Variable('j', 1, 10),), strides=(0,), offset=0, mask=None, contiguous=False)
    vm1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
-    self.assertEqual(vm2+vm1, None)
+    self.assertEqual(vm2+vm1, vm1)

  def test_merge_view_recursion_err2(self):
    vm2 = View(shape=(Variable('a', 1, 10).bind(4),), strides=(0,), offset=0, mask=None, contiguous=False)
@@ -1,22 +0,0 @@
-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,17 +5,14 @@ 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
-from tinygrad.uop.symbolic import sym, commutative
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer, track_rewrites
+from tinygrad.uop.symbolic import sym
 from tinygrad.uop.spec import uops_to_z3

-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}")
+@track_rewrites(name="simplify symbolic uop")
+def render(v) -> UOp:
+  v_simplified = graph_rewrite(v, sym)
+  return v_simplified

 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)
@@ -36,11 +33,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 = graph_rewrite(v, sym, name="simplify symbolic uop")
+    v_simplified = render(v)
    if test_z3: self.check_equal_z3(v, v_simplified)
-    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
+    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)
    self.assertEqual(nmin, n)
    self.assertEqual(nmax, m)

@@ -79,7 +76,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)")
+    self.helper_test_variable(expr, 0, 1, ("(((idx1*4)+FLOAT4_INDEX)<512)", "((FLOAT4_INDEX+(idx1*4))<512)"))

  def test_div_reduction(self):
    self.helper_test_variable(Variable("a", 2, 3)//2, 1, 1, "1")
@@ -190,7 +187,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, "a.maximum(11)")
+    self.helper_test_variable(Variable("a", 0, 20).maximum(10).maximum(11), 11, 20, "max(a, 11)")

  def test_add_min_max(self):
    self.helper_test_variable(Variable("a", 0, 8) * 2 + 12, 12, 16+12, "((a*2)+12)")
@@ -219,7 +216,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))")
+    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))"))

  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)")
@@ -282,7 +279,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)")
+                              ("(((z+(x*-1))+(y*-1))+7)", "(((y*-1)+(z+(x*-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):
@@ -304,7 +301,8 @@ 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)")
+    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("uidx", 0, 3), Variable("a", 0, 1529) * 12]) < (4 * 67), 0, 1, "(a<23)")

  def test_mul_mod_large(self):
@@ -366,7 +364,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)")
+    self.helper_test_variable(usum([Variable("b", 0, 2), Variable("a", 0, 5)*10])%9, 0, 7, ("(b+a)", "(a+b)"))

  def test_sum_0(self):
    self.helper_test_variable(usum([Variable("a", 0, 7)]), 0, 7, "a")
@@ -397,11 +395,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)")
+                              ("((((a*3)+(b*2))+(c*4))<2)", "(((b*2)+((a*3)+(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)")
+                              ("((((a*3)+(b*2))+(c*4))<1)", "(((b*2)+((a*3)+(c*4)))<1)"))

  def test_and_fold(self):
    self.helper_test_variable(uand([uconst(0), Variable("a", 0, 1)]), 0, 0, "0")
@@ -563,35 +561,38 @@ 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)))")
+                              ("((((((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))"))

  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)")
-    self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//10, 0, 3, "(((gidx0*2)+lidx2)//5)")
+    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*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)")
-    self.helper_test_variable((lidx2*2+gidx0*4+lidx3)//12, 0, 4, "(((lidx2//2)+gidx0)//3)")
+    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)"))

  @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))")
+    self.helper_test_variable((gidx0+lidx2+lidx3)*4, 0, 80,
+                              ("(((gidx0*4)+(lidx2*4))+(lidx3*4))","((lidx3*4)+((gidx0*4)+(lidx2*4)))"))

  @unittest.expectedFailure
  def test_variable_divmod(self):
@@ -661,7 +662,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).where((idx//4), -1)")
+    self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "((idx//4) if (idx<4) else -1)")

  def test_idiv_lt(self):
    idx = Variable("idx", 0, 24)
@@ -680,8 +681,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)")
-    self.helper_test_variable((a+b*2+c*4<1).ne(True), 0, 1, "((((a+b)+c)<1)!=True)")
+    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)'))

  def test_where_removal(self):
    cond = Variable("a", 0, 3) < 2
@@ -699,30 +700,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, "(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))")
+    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))")

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

    # 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, "((x<2).where(a, b)+(x<2).where(b, a))")
+    self.helper_test_variable(ab+ba, 0, 6, "((a if (x<2) else b)+(b if (x<2) else a))")

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

  def test_neg_in_comp(self):
    a = Variable("a", 0, 3)
@@ -749,7 +750,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, "((s<6)&(2<s)).where(a, b)")
+    self.helper_test_variable(expr, 0, 3, "(a if ((s<6)&(2<s)) else b)")

  def test_where_merge_branches2(self):
    cond1 = Variable("s", 0, 10) < 5
@@ -758,7 +759,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, "(s<5).where(a, b)")
+    self.helper_test_variable(expr, 0, 3, "(a if (s<5) else b)")

  def test_symbolic_div(self):
    # from symbolic arange
@@ -773,7 +774,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, "a.reciprocal()", test_z3=False)
+    self.helper_test_variable(2*(2*a).reciprocal(), -math.inf, math.inf, "(1/a)", test_z3=False)

  def test_trunc_noop(self):
    a = Variable("a", 1, 10, dtypes.int)
@@ -782,8 +783,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.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))")
+    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))")

 class TestSymbolicNumeric(unittest.TestCase):
  def helper_test_numeric(self, f):
@@ -69,5 +69,161 @@ class TestMergeDims(unittest.TestCase):
    # print(f"{ShapeTracker.from_shape((2, 1, 1)).pad(((0, 0), (0, 1), (0, 1))).views[-1]}")
    self.assertEqual(merge_dims((2, 2, 2), (1, 0, 0), ((0, 2), (0, 2), (0, 1))), ((2, 1, 2), (4, 0, 4)))

+class TestMergeViews(unittest.TestCase):
+  def test_with_mask_0(self):
+    # from test/test_ops.py::TestOps::test_pad_reflect_mode
+    v0 = View(shape=(1, 1, 5, 8), strides=(0, 0, 5, 1), offset=-3, mask=((0, 1), (0, 1), (0, 5), (3, 8)), contiguous=False)
+    v1 = View(shape=(1, 1, 2, 2), strides=(0, 0, 8, 1), offset=3, mask=None, contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, View(shape=(1, 1, 2, 2), strides=(0, 0, 5, 1), offset=0, mask=None, contiguous=False))
+
+  def test_with_mask_1(self):
+    # from test/test_ops.py::TestOps::test_pad_reflect_mode
+    v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, 3, 1), offset=-6, mask=((0, 3), (0, 3), (2, 4), (1, 3)), contiguous=False)
+    v1 = View(shape=(3, 3, 2, 2), strides=(45, 15, 3, 1), offset=7, mask=None, contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, View(shape=(3, 3, 2, 2), strides=(27, 9, 3, 1), offset=1, mask=None, contiguous=False))
+
+  def test_with_mask_2(self):
+    # from test/test_ops.py::TestOps::test_pad_reflect_mode
+    v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, -3, 1), offset=6, mask=((0, 3), (0, 3), (0, 2), (0, 2)), contiguous=False)
+    v1 = View(shape=(3, 3, 2, 2), strides=(45, 15, -3, 1), offset=3, mask=None, contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, View(shape=(3, 3, 2, 2), strides=(27, 9, 3, 1), offset=3, mask=None, contiguous=False))
+
+  def test_with_mask_3(self):
+    # from test/test_ops.py::TestOps::test_pad_reflect_mode
+    # has a mask in the final view
+    v0 = View(shape=(3, 3, 4, 4), strides=(27, 9, 3, 1), offset=-5, mask=((0, 3), (0, 3), (2, 4), (0, 2)), contiguous=False)
+    v1 = View(shape=(3, 3, 4, 2), strides=(48, 16, 4, 1), offset=0, mask=None, contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, View(shape=(3, 3, 4, 2), strides=(27, 9, 3, 1), offset=-5, mask=((0, 3), (0, 3), (2, 4), (0, 2)), contiguous=False))
+
+  def test_with_mask_4(self):
+    # from test/test_ops.py::TestOps::test_pad_reflect_mode
+    # has a mask in the final view
+    v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, -3, 1), offset=6, mask=((0, 3), (0, 3), (0, 2), (1, 3)), contiguous=False)
+    v1 = View(shape=(3, 3, 3, 3), strides=(45, 15, 3, 1), offset=6, mask=None, contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, View(shape=(3, 3, 3, 3), strides=(0, 0, 0, 0), offset=0, mask=((0, 0), (0, 0), (0, 0), (0, 0)), contiguous=False))
+
+  def test_with_mask_5(self):
+    # from test/test_ops.py::TestOps::test_pad_reflect_mode
+    # has a mask in the final view
+    v0 = View(shape=(1, 1, 6, 5), strides=(0, 0, 5, 1), offset=-5, mask=((0, 1), (0, 1), (1, 6), (0, 5)), contiguous=False)
+    v1 = View(shape=(1, 1, 6, 3), strides=(0, 0, 5, -1), offset=3, mask=None, contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, View(shape=(1, 1, 6, 3), strides=(0, 0, 5, -1), offset=-2, mask=((0, 1), (0, 1), (1, 6), (0, 3)), contiguous=False))
+
+  @unittest.expectedFailure  # TODO: fix these
+  def test_merges_from_fuzzer1(self):
+    v0 = View(shape=(2, 4), strides=(2, 1), offset=-2, mask=((0, 2), (2, 4)), contiguous=False)
+    v1 = View(shape=(2, 4, 2, 2), strides=(4, 0, -2, -1), offset=3, mask=None, contiguous=False)
+    target = View(shape=(2, 4, 2, 2), strides=(2, 0, 0, -1), offset=1, mask=((0, 2), (0, 4), (0, 1), (0, 2)), contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, target)
+
+  @unittest.expectedFailure  # TODO: fix these
+  def test_merges_from_fuzzer2(self):
+    v0 = View(shape=(5, 10, 12), strides=(100, 1, 10), offset=-20, mask=((0, 5), (0, 10), (2, 12)), contiguous=False)
+    v1 = View(shape=(10, 6, 5, 2, 2), strides=(12, 2, 120, 1, 0), offset=0, mask=None, contiguous=False)
+    target = View(shape=(10, 6, 5, 2, 2), strides=(1, 20, 100, 10, 0), offset=-20, mask=((0, 10), (1, 6), (0, 5), (0, 2), (0, 2)), contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, target)
+
+  @unittest.expectedFailure  # TODO: fix these
+  def test_merges_from_fuzzer3(self):
+    v0 = View(shape=(8, 7, 3), strides=(1, 12, -4), offset=6, mask=((2, 6), (0, 7), (0, 3)), contiguous=False)
+    v1 = View(shape=(4, 2, 6, 2, 1), strides=(42, 21, 3, 1, 0), offset=4, mask=None, contiguous=False)
+    target = View(shape=(4, 2, 6, 2, 1), strides=(2, 1, 12, -4, 0), offset=14, mask=((1, 3), (0, 2), (0, 6), (0, 2), (0, 1)), contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, target)
+
+  @unittest.expectedFailure  # TODO: fix these
+  def test_merges_from_fuzzer4(self):
+    v0 = View(shape=(7, 21, 3), strides=(54, 3, 1), offset=-9, mask=((0, 6), (3, 21), (0, 3)), contiguous=False)
+    v1 = View(shape=(5, 3, 3, 7), strides=(63, 1, 3, 9), offset=63, mask=None, contiguous=False)
+    target = View(shape=(5, 3, 3, 7), strides=(54, 1, 3, 9), offset=45, mask=((0, 5), (0, 3), (0, 3), (1, 7)), contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, target)
+
+  @unittest.expectedFailure  # TODO: fix these
+  def test_merges_from_fuzzer5(self):
+    v0 = View(shape=(5, 1, 24), strides=(20, 0, 1), offset=-2, mask=((0, 5), (0, 1), (2, 22)), contiguous=False)
+    v1 = View(shape=(12, 2, 5, 2, 1), strides=(2, 1, 24, 0, 0), offset=0, mask=None, contiguous=False)
+    target = View(shape=(12, 2, 5, 2, 1), strides=(2, 1, 20, 0, 0), offset=-2, mask=((1, 11), (0, 2), (0, 5), (0, 2), (0, 1)), contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, target)
+
+  def test_merge_views_variable(self):
+    from tinygrad import Variable
+    N = 100
+    start_pos = Variable("start_pos", 1, N-1)
+    v0 = View(shape=(N, 32, 2), strides=(32, 1, 0), offset=0, mask=((0, N), (0, 32), (0, 1)), contiguous=False)
+    v1 = View(shape=(1, 8, 1, 32), strides=(0, 0, 0, 2), offset=start_pos*64, mask=None, contiguous=False)
+    target = View(shape=(1, 8, 1, 32), strides=(0,0,0,1), offset=start_pos*32, mask=None, contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, target)
+
+  def test_view_padded_area1(self):
+    # test_multinomial
+    v0 = View(shape=(2,), strides=(0,), offset=0, mask=((1, 2),), contiguous=False)
+    v1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, View(shape=(1,), strides=(0,), offset=0, mask=((0, 0),), contiguous=False))
+
+  def test_view_padded_area2(self):
+    # test_pad_reflect_mode
+    v0 = View(shape=(1, 1, 10, 7), strides=(0, 0, 5, 1), offset=-15, mask=((0, 1), (0, 1), (3, 8), (0, 5)), contiguous=False)
+    v1 = View(shape=(0, 0, 0, 0), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=True)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, View(shape=(0, 0, 0, 0), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=True))
+
+  def test_view_padded_area3(self):
+    # test_roll
+    v0 = View(shape=(2, 4), strides=(0, 1), offset=4, mask=((0, 1), (0, 4)), contiguous=False)
+    v1 = View(shape=(1, 4), strides=(0, 1), offset=4, mask=None, contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, View(shape=(1, 4), strides=(0, 0), offset=0, mask=((0, 0), (0, 0)), contiguous=False))
+
+  def test_view_padded_area4(self):
+    # test_std_mean
+    v0 = View(shape=(2,), strides=(0,), offset=0, mask=((0, 1),), contiguous=False)
+    v1 = View(shape=(1, 1, 1), strides=(0, 0, 0), offset=1, mask=None, contiguous=False)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, View(shape=(1, 1, 1), strides=(0, 0, 0), offset=0, mask=((0, 0), (0, 0), (0, 0)), contiguous=False))
+
+  def test_empty_shape_view1(self):
+    # test_stack_slice
+    v0 = View(shape=(3, 5), strides=(0, 1), offset=0, mask=((0, 1), (0, 5)), contiguous=False)
+    v1 = View(shape=(), strides=(), offset=0, mask=None, contiguous=True)
+    v = v0 + v1
+    self.assertIsNotNone(v)
+    self.assertEqual(v, View(shape=(), strides=(), offset=0, mask=None, contiguous=True))
+
+  def test_empty_shape_view2(self):
+    # test_std_mean
+    v0 = View(shape=(2,), strides=(0,), offset=0, mask=((1, 2),), contiguous=False)
+    v1 = View(shape=(), strides=(), offset=0, mask=None, contiguous=True)
+    v = v0 + v1
+    # TODO: why is this different?
+    self.assertIsNone(v)
+
 if __name__ == '__main__':
  unittest.main()
@@ -30,7 +30,6 @@ 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
@@ -463,21 +462,5 @@ 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), 5)
+    self.assertEqual(len(backward_schedule), 4)

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

+def gpt2_decode_vocab(voc: dict[str, int]): # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
+  c2b = { chr(cp): cp for cp in itertools.chain(range(ord("!"), ord("~")+1), range(ord("¡"), ord("¬")+1), range(ord("®"), ord("ÿ")+1)) }
+  c2b.update({ chr(256+off): cp for off, cp in enumerate(cp for cp in range(256) if chr(cp) not in c2b) })
+  return { bytes(c2b[c] for c in tok): tid for tok, tid in voc.items() }
+
+def get_llama_re():
+  def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre))
+  r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
+  # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
+  return "(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
+    f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+"
+
 class SimpleTokenizer:
-  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()}
+  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"(?!)")

  @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(dict(normal_tokens), dict(special_tokens))
+    return SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens))

-  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]:
+  def encode(self, text: str):
    tokens: list[int] = []
    pos = 0
-    for match in self._split_to_sentence.finditer(text):
+    for match in self._special_re.finditer(text):
      tokens.extend(self._encode_sentence(text[pos:match.start(0)]) + [self._special_tokens[text[match.start(0):match.end(0)]]])
      pos = match.end(0)
    return tokens + self._encode_sentence(text[pos:])

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

+  def _encode_sentence(self, chunk: str): return [ tok for word in self._pat.findall(chunk) for tok in self._encode_word(word.encode()) ]
+  def _encode_word(self, word: bytes):
+    if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token]
+    parts = [word[i:i+1] for i in range(len(word))]
+    while True:
+      min_tid, min_idx = 2**32, -1
+      for idx, (p1, p2) in enumerate(zip(parts[:-1], parts[1:])):
+        tid = self._normal_tokens.get(p1 + p2, min_tid)
+        if tid < min_tid: min_tid, min_idx = tid, idx
+      if min_idx == -1: break
+      parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx+1]] + parts[min_idx+2:]
+    try: return [ self._normal_tokens[p] for p in parts ]
+    except KeyError: raise RuntimeError("token not found")
+
 def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor:
  B, H, T, Hd = x.shape
-  assert isinstance(Hd, int) and (Hd & 1) == 0, "RoPE requires an even head dimension"
+  assert (Hd & 1) == 0, "RoPE requires an even head dimension"
  half = Hd // 2
-  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, :]
+  angles = (Tensor.arange(T, dtype="float32") + 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)
@@ -14,7 +14,7 @@ from tinygrad.uop.decompositions import get_late_rewrite_patterns
 from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_expander, pm_group_for_reduce
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
  ReduceContext, correct_load_store, pm_render
-from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
+from tinygrad.codegen.late.control_flow import pm_endranges, linearize, CFGContext, pm_control_flow
 from tinygrad.codegen.opt.postrange import pm_postrange_opt
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_reduce_simplify, pm_flatten_range, pm_split_ranges
 from tinygrad.schedule.rangeify import pm_add_buffers, rangeify_codegen
@@ -30,12 +30,6 @@ class RewriteStep:

 def apply_rewrites(sink:UOp, rewrites:list[RewriteStep]): return functools.reduce(lambda x,f: f(x), rewrites, sink)

-rewrites_for_linearizer = [
-  RewriteStep(block_create, ctx=BlockContext.from_sink, name="Linearizer: Create Blocks", bottom_up=True),
-  RewriteStep(pm_blockend_merge, name="Linearizer: Merge Blockends"),
-  RewriteStep(block_merge, name="Linearizer: Merge Blocks"),
-  RewriteStep(pm_finalize, name="Linearizer: Finalize")]
-
 def get_rewrites_for_renderer(opts:Renderer, optimize:bool=True, linearizer:bool=True) -> list[RewriteStep]:
  # cache with the values of the context vars
  return _get_rewrites_for_renderer(opts, optimize, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value)
@@ -77,6 +71,9 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
  # add gpu dims (late). this works after devectorize, but it's faster here
  ret.append(RewriteStep(pm_add_gpudims, lambda _: opts, name="add gpudims"))

+  # add end ranges
+  ret.append(RewriteStep(pm_endranges, name="add end ranges"))
+
  # devectorize (TODO: does this need opts?)
  if _DEVECTORIZE >= 2: pm_devectorize = sym+load_store_folding+load_store_indexing
  elif _DEVECTORIZE: pm_devectorize = sym+devectorize+load_store_folding+correct_load_store+load_store_indexing
@@ -101,11 +98,14 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
  pm_final_rewrite = pm_decomp+pm_render+extra_matcher
  ret.append(RewriteStep(pm_final_rewrite, lambda _: opts.device, name="final rewrite"))

-  # return the list (with optional linearizer)
-  return ret + (rewrites_for_linearizer if linearizer else [])
+  # build CFG
+  ret.append(RewriteStep(pm_control_flow, lambda sink: CFGContext(sink), name="add control flow"))

-def full_rewrite_to_sink(sink:UOp, opts:Renderer|None=None, optimize:bool=True, linearizer:bool=False) -> UOp:
-  return apply_rewrites(sink, get_rewrites_for_renderer(opts if opts is not None else Renderer(), optimize, linearizer))
+  # return the list (with optional linearizer)
+  return ret
+
+def full_rewrite_to_sink(sink:UOp, opts:Renderer|None=None, optimize:bool=True) -> UOp:
+  return apply_rewrites(sink, get_rewrites_for_renderer(opts if opts is not None else Renderer(), optimize))

 def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]:
  """
@@ -119,6 +119,6 @@ def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]:
    Linear program in UOps.
  """

-  lst = list(full_rewrite_to_sink(sink, opts, optimize=sink.tag is None, linearizer=True).arg.lst)
+  lst = linearize(full_rewrite_to_sink(sink, opts, optimize=sink.tag is None))
  if __debug__: type_verify(lst)
  return lst
@@ -1,7 +1,8 @@
 import math
 from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, sint_to_uop
-from tinygrad.helpers import all_int, dedup, get_contraction
+from tinygrad.helpers import all_int, dedup
 from tinygrad.dtype import dtypes
+from tinygrad.shape.view import get_contraction
 from tinygrad.renderer import Renderer

 def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
@@ -0,0 +1,106 @@
+import heapq
+from collections import defaultdict
+from tinygrad.helpers import partition
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat
+
+def end_store_ranges(x:UOp):
+  ranges_to_end, others = partition(x.src[2:], lambda x: x.op is Ops.RANGE)
+  if not len(ranges_to_end): return None
+  ret = x.replace(src=x.src[:2]+tuple(others))
+  for r in ranges_to_end: ret = UOp(Ops.ENDRANGE, src=(ret,r))
+  return ret
+
+def might_end_if(x:UOp):
+  ifs_to_end = []
+  ended_ifs = []
+  def find_if(y:UOp):
+    if y.op is Ops.BARRIER: return False
+    if y.op is Ops.IF: ifs_to_end.append(y)
+    if y.op is Ops.ENDIF: ended_ifs.append(y.src[1])
+    return True
+  x.toposort(find_if)
+  del find_if
+  ifs_to_end = [x for x in ifs_to_end if x not in ended_ifs]
+  if not len(ifs_to_end): return None
+  ret = x.src[0] if len(x.src) == 1 else UOp(Ops.NOOP, src=x.src)
+  for r in ifs_to_end: ret = UOp(Ops.ENDIF, src=(ret,r))
+  return x.replace(src=(ret,))
+
+pm_endranges = PatternMatcher([
+  # all ranges are ended by STORE
+  (UPat(Ops.STORE, name="x"), end_store_ranges),
+  # all ENDIF either goes on BARRIER or SINK
+  (UPat((Ops.BARRIER, Ops.SINK), name="x"), might_end_if),
+])
+
+class CFGContext:
+  def __init__(self, sink:UOp):
+    # there are 3 relationships between ranges:
+    # nested, meaning endrange y is a dependency of endrange x and range x is a dependency of endrange y
+    # dependent, meaning endrange y is a dependency of endrange x and range x is not a dependency of endrange y
+    # independent, endrange y is not a dependency of endrange x
+    deps: dict[UOp, set[UOp]] = {}
+    nesting: dict[UOp, UOp] = {}
+    for u in sink.toposort():
+      deps[u] = set().union(*(deps[s] for s in u.src))
+      if u.op in (Ops.ENDRANGE, Ops.ENDIF):
+        for n in [x for x in deps[u] if x.op in (Ops.ENDRANGE, Ops.ENDIF) and u.src[1] in deps[x] and x not in nesting]: nesting[n] = u
+      if u.op is Ops.SINK:
+        for n in [x for x in deps[u] if x.op in (Ops.ENDRANGE, Ops.ENDIF) and x not in nesting]: nesting[n] = u
+      if u.op in (Ops.RANGE, Ops.ENDRANGE, Ops.IF, Ops.ENDIF): deps[u] |= {u}
+
+    self.edges: dict[UOp, UOp] = {}
+    siblings: dict[UOp, list[UOp]] = {}
+    for k,vv in nesting.items(): siblings.setdefault(vv, []).append(k)
+    for k,v in siblings.items():
+      # range/if that have dependencies on other siblings need to run after them
+      order = sorted(v, key=lambda x: len([y for y in v if y in deps[x]]))
+      zipped = zip(order, order[1:]) if k.op is Ops.SINK else zip([k.src[1]] + order, order)
+      for x,y in zipped: self.edges[y.src[1]] = x
+
+pm_control_flow = PatternMatcher([
+  (UPat(Ops.RANGE, src=(UPat(),), name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None),
+])
+
+def linearize(x:UOp) -> list[UOp]:
+  lst = x.toposort()
+  in_this_block = set(lst)
+  local_children: defaultdict[UOp, list[UOp]] = defaultdict(list)
+  in_degree:dict[UOp, int] = {}
+  priorities:dict[UOp, int] = {}
+
+  # get local children and assign priorities
+  # NOTE: this requires the lst be locally toposorted
+  for u in reversed(lst):
+    in_degree[u] = 0
+    for s in u.src:
+      if s in in_this_block:
+        local_children[s].append(u)
+        in_degree[u] += 1
+    # put loads in the beginning of the block and prevent priority inversion. hack for BARRIER grouping too
+    priority = [0] + [priorities[x] for x in local_children[u]]
+    # if needs to be the first thing
+    if u.op is Ops.IF: priority.append(-10000)
+    if u.op is Ops.LOAD: priority.append(-1000)
+    if u.op is Ops.BARRIER: priority.append(-1500)
+    # ranges are scheduled as late as possible so anything that can be outside is
+    #if u.op is Ops.RANGE: priority = [2000]
+    # move defines and consts to the top
+    if u.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}: priority.append(-2000)
+    priorities[u] = min(priority)
+
+  # number the uops in "ideal" order
+  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: (priorities[x],)+x.tuplize))}
+
+  # then force then to be toposorted in as close to the ideal order as possible
+  heapq.heapify(heap:=[(nkey[u],u) for u in lst if in_degree[u] == 0])
+  newlst = []
+  while heap:
+    newlst.append(u:=heapq.heappop(heap)[1])
+    for v in local_children[u]:
+      in_degree[v] -= 1
+      if in_degree[v] == 0: heapq.heappush(heap, (nkey[v],v))
+
+  assert len(newlst) == len(lst), f"len mismatch {len(newlst)} != {len(lst)}"
+  return newlst
+
@@ -11,7 +11,7 @@ from tinygrad.renderer import Renderer
 # ***** image load valid simplification *****

 def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
-  idx = uop_given_valid(valid, start_idx)
+  if (idx:=uop_given_valid(valid, start_idx)) is None: return buf.index(UOp.invalid())
  if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx.valid(valid))

  # wait for it to be image indexed before running simplification
@@ -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.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"}
+axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
+               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "green", 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 if too many threads. Heuristic: use about 128K ops per thread
+      # Skip is 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.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}
+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}

 class Scheduler:
  def __init__(self, ast:UOp, opts:Renderer):
@@ -165,22 +165,15 @@ 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",
-                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="")
+        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

      # 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())
+      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
  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", "TINYFS", "NPY", "PYTHON"], f"usage of device {ix} disallowed"
+    assert ALLOW_DEVICE_USAGE or ix.split(":")[0] in ["DISK", "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", "TINYFS", "NPY"] and getenv(d) == 1])
+    from_env = dedup(dev + [d for d in self._devices if d not in ["DISK", "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,14 +137,16 @@ 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: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"dtype":self.dtype, "sz":self.size}))
+      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}))
    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.trace_num))
+      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self._prof_num))
      self.allocator.free(self._buf, self.nbytes, self.options)
    elif self._base is not None: self._base.allocated_views -= 1
    del self._buf
@@ -158,10 +160,6 @@ 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") or src.device.startswith("TINYFS")) and hasattr(dest.allocator, '_as_buffer'):
+    elif src.device.startswith("DISK") 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,9 +165,7 @@ 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:
-      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))
+    if PROFILE: cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, {"metadata":self.metadata, "var_vals":var_vals}))
    et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2)
    if do_update_stats:
      GlobalCounters.kernel_count += 1
@@ -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: (None,)),
+  (UPat(Ops.BITCAST), lambda ctx: (None,)),
 ])

 def _deepwalk(root:UOp, targets:set[UOp]) -> list[UOp]:
@@ -1,6 +1,6 @@
 from __future__ import annotations
 import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass
-import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
+import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal
 from dataclasses import dataclass, field
 from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator

@@ -23,13 +23,10 @@ def argfix(*x):
    if len(x) != 1: raise ValueError(f"bad arg {x}")
    return tuple(x[0])
  return x
-# 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 argsort(x): return type(x)(sorted(range(len(x)), key=x.__getitem__)) # https://stackoverflow.com/questions/3382352/equivalent-of-numpy-argsort-in-basic-python
 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): # 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 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 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)
@@ -85,13 +82,6 @@ def word_wrap(x, wrap=80):
  while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1
  return x[:i] + "\n" + word_wrap(x[i:], wrap)

-# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
-def get_contraction(old_shape:tuple[T, ...], new_shape:tuple[T, ...]) -> list[list[int]]|None: # T is sint
-  acc_old, acc_new = list(itertools.accumulate(old_shape, operator.mul)), list(itertools.accumulate(new_shape, operator.mul))
-  try: split = [acc_old.index(acc)+1 if acc != 1 else 0 for acc in acc_new]
-  except ValueError: return None
-  return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])]
-
 def suppress_finalizing(func):
  def wrapper(*args, **kwargs):
    try: return func(*args, **kwargs)
@@ -99,7 +89,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): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t
+def unwrap_class_type(cls_t: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')

@@ -153,7 +143,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, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
+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_LLVM, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("AMD_LLVM", 1)
 VIZ = PROFILE = ContextVar("VIZ", 0)
 SPEC = ContextVar("SPEC", 0)
@@ -221,12 +211,11 @@ 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
@@ -285,8 +274,7 @@ 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),))
+  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
  conn.commit()
  cur.close()
  return val
@@ -352,10 +340,10 @@ def capstone_flatdump(lib: bytes):
    print(f"{instr.address:#08x}: {instr.mnemonic}\t{instr.op_str}")
  sys.stdout.flush()

-def wait_cond(cb, *args, value=True, timeout_ms=10000, msg="") -> bool:
+def wait_cond(cb, 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(*args)) == value: return val
+    if (val:=cb()) == 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:float=1e-5, affine:bool=True):
+  def __init__(self, num_features:int, eps=1e-5, affine=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:float=1e-5, elementwise_affine:bool=True):
+  def __init__(self, normalized_shape:int|tuple[int, ...], eps=1e-5, elementwise_affine=True):
    self.normalized_shape: tuple[int, ...] = make_tuple(normalized_shape, 1)
-    self.axis, self.eps = tuple(-1-i for i in range(len(self.normalized_shape))), eps
+    self.axis, self.eps, self.elementwise_affine = tuple(-1-i for i in range(len(self.normalized_shape))), eps, elementwise_affine
    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 self.weight is None or self.bias is None: return x
+    if not self.elementwise_affine: return x
    return x * self.weight + self.bias

 class LayerNorm2d(LayerNorm):
@@ -2,7 +2,7 @@ from typing import Literal, Callable, cast
 import os, math, sys
 from collections import defaultdict, Counter
 from tinygrad.codegen.opt import tc
-from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
+from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, sint_to_uop, range_str
 from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX, CPU_COUNT
 from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate
 from tinygrad.renderer import Renderer
@@ -108,13 +108,11 @@ 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 = ""
-    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"
+    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
    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"]
-    launch_bounds = prod([d.vmax for d in local_dims])
+    launch_bounds = sint_to_uop(prod(local_dims)).vmax
    prg = ''.join([f"{self.kernel_typedef.format(launch_bounds=launch_bounds)} {function_name}(",] +
    [', '.join([f'{t} {name}' for name,t in buftypes] + self.extra_args)] +
    [") {\n" + tmp] + ['\n'.join(kernel), "\n}"])
@@ -231,12 +229,10 @@ 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.
-      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){{
+      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){{
  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); }}
-  AMX_SET(1);\n  return data0;\n}}"""]
+  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
    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 ""
@@ -4,7 +4,7 @@ from tinygrad.codegen.opt import tc
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.uop.decompositions import xexp2, xlog2
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, sint_to_uop, range_str
 from tinygrad.dtype import dtypes, DType, PtrDType, truncate
 from tinygrad.helpers import prod, AMX

@@ -226,7 +226,7 @@ class AMDLLVMRenderer(LLVMRenderer):
  def _render_footer(self, uops: list[UOp]) -> str:
    # TODO: this is copied from cstyle
    local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
-    requiredMaxThreadsPerBlock = prod([d.vmax for d in local_dims])
+    requiredMaxThreadsPerBlock = sint_to_uop(prod(local_dims)).vmax
    attributes = ["alwaysinline", "nounwind", '"no-builtins"',
                  f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"']
    return 'attributes #0 = { ' + ' '.join(attributes) + ' }'
@@ -2,7 +2,7 @@ from typing import cast, Callable
 import struct
 from collections import defaultdict
 from tinygrad.codegen.opt import tc
-from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp
+from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, sint_to_uop
 from tinygrad.dtype import dtypes, DType, PtrDType, AddrSpace
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import CUDARenderer
@@ -157,7 +157,7 @@ class PTXRenderer(Renderer):
    def fmt(line): return line if line[0]=="$" else "\t" + line.replace(" ", "\t" if len(line.split(" ")[0]) > 7 else "\t\t", 1)
    kernel = '\n'.join(map(fmt, [f".reg .{reg.split('_')[-2]} %{reg}<{cnt}>;" for reg,cnt in regs] + kernel + ["ret;"]))
    local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
-    launch_bounds = prod([d.vmax for d in local_dims])
+    launch_bounds = sint_to_uop(prod(local_dims)).vmax
    params = ',\n\t'.join([f".param .{'u64' if dtype.__class__ == PtrDType else self.types[dtype]} {name}" for name,dtype in bufs])
    return f"{self.kernel_prefix.format(launch_bounds=launch_bounds)} {function_name} (\n\t{params}\n)\n.maxntid {launch_bounds}\n{{\n{kernel}\n}}"

@@ -174,14 +174,12 @@ sqtt_version__enumvalues = {
    6: 'SQTT_VERSION_2_3',
    7: 'SQTT_VERSION_2_4',
    11: 'SQTT_VERSION_3_2',
-    12: 'SQTT_VERSION_3_3',
 }
 SQTT_VERSION_NONE = 0
 SQTT_VERSION_2_2 = 5
 SQTT_VERSION_2_3 = 6
 SQTT_VERSION_2_4 = 7
 SQTT_VERSION_3_2 = 11
-SQTT_VERSION_3_3 = 12
 sqtt_version = ctypes.c_uint32 # enum

 # values for enumeration 'sqtt_file_chunk_type'
@@ -338,8 +336,6 @@ sqtt_gfxip_level__enumvalues = {
    7: 'SQTT_GFXIP_LEVEL_GFXIP_10_1',
    9: 'SQTT_GFXIP_LEVEL_GFXIP_10_3',
    12: 'SQTT_GFXIP_LEVEL_GFXIP_11_0',
-    13: 'SQTT_GFXIP_LEVEL_GFXIP_11_5',
-    16: 'SQTT_GFXIP_LEVEL_GFXIP_12',
 }
 SQTT_GFXIP_LEVEL_NONE = 0
 SQTT_GFXIP_LEVEL_GFXIP_6 = 1
@@ -350,8 +346,6 @@ SQTT_GFXIP_LEVEL_GFXIP_9 = 5
 SQTT_GFXIP_LEVEL_GFXIP_10_1 = 7
 SQTT_GFXIP_LEVEL_GFXIP_10_3 = 9
 SQTT_GFXIP_LEVEL_GFXIP_11_0 = 12
-SQTT_GFXIP_LEVEL_GFXIP_11_5 = 13
-SQTT_GFXIP_LEVEL_GFXIP_12 = 16
 sqtt_gfxip_level = ctypes.c_uint32 # enum

 # values for enumeration 'sqtt_memory_type'
@@ -812,16 +806,12 @@ elf_gfxip_level__enumvalues = {
    51: 'EF_AMDGPU_MACH_AMDGCN_GFX1010',
    54: 'EF_AMDGPU_MACH_AMDGCN_GFX1030',
    65: 'EF_AMDGPU_MACH_AMDGCN_GFX1100',
-    67: 'EF_AMDGPU_MACH_AMDGCN_GFX1150',
-    78: 'EF_AMDGPU_MACH_AMDGCN_GFX1200',
 }
 EF_AMDGPU_MACH_AMDGCN_GFX801 = 40
 EF_AMDGPU_MACH_AMDGCN_GFX900 = 44
 EF_AMDGPU_MACH_AMDGCN_GFX1010 = 51
 EF_AMDGPU_MACH_AMDGCN_GFX1030 = 54
 EF_AMDGPU_MACH_AMDGCN_GFX1100 = 65
-EF_AMDGPU_MACH_AMDGCN_GFX1150 = 67
-EF_AMDGPU_MACH_AMDGCN_GFX1200 = 78
 elf_gfxip_level = ctypes.c_uint32 # enum
 class struct_sqtt_file_chunk_spm_db(Structure):
    pass
@@ -1617,8 +1607,7 @@ __all__ = \
    'ApiCmdUpdateBuffer', 'ApiCmdWaitEvents', 'ApiCmdWriteTimestamp',
    'ApiInvalid', 'ApiRayTracingSeparateCompiled',
    'EF_AMDGPU_MACH_AMDGCN_GFX1010', 'EF_AMDGPU_MACH_AMDGCN_GFX1030',
-    'EF_AMDGPU_MACH_AMDGCN_GFX1100', 'EF_AMDGPU_MACH_AMDGCN_GFX1150',
-    'EF_AMDGPU_MACH_AMDGCN_GFX1200', 'EF_AMDGPU_MACH_AMDGCN_GFX801',
+    'EF_AMDGPU_MACH_AMDGCN_GFX1100', 'EF_AMDGPU_MACH_AMDGCN_GFX801',
    'EF_AMDGPU_MACH_AMDGCN_GFX900', 'EventCmdBlitImage',
    'EventCmdBuildAccelerationStructuresIndirectKHR',
    'EventCmdBuildAccelerationStructuresKHR',
@@ -1682,8 +1671,7 @@ __all__ = \
    'SQTT_FILE_CHUNK_TYPE_SQTT_DESC', 'SQTT_FILE_MAGIC_NUMBER',
    'SQTT_FILE_VERSION_MAJOR', 'SQTT_FILE_VERSION_MINOR',
    'SQTT_GFXIP_LEVEL_GFXIP_10_1', 'SQTT_GFXIP_LEVEL_GFXIP_10_3',
-    'SQTT_GFXIP_LEVEL_GFXIP_11_0', 'SQTT_GFXIP_LEVEL_GFXIP_11_5',
-    'SQTT_GFXIP_LEVEL_GFXIP_12', 'SQTT_GFXIP_LEVEL_GFXIP_6',
+    'SQTT_GFXIP_LEVEL_GFXIP_11_0', 'SQTT_GFXIP_LEVEL_GFXIP_6',
    'SQTT_GFXIP_LEVEL_GFXIP_7', 'SQTT_GFXIP_LEVEL_GFXIP_8',
    'SQTT_GFXIP_LEVEL_GFXIP_8_1', 'SQTT_GFXIP_LEVEL_GFXIP_9',
    'SQTT_GFXIP_LEVEL_NONE', 'SQTT_GPU_NAME_MAX_SIZE',
@@ -1709,9 +1697,9 @@ __all__ = \
    'SQTT_QUEUE_TYPE_COMPUTE', 'SQTT_QUEUE_TYPE_DMA',
    'SQTT_QUEUE_TYPE_UNIVERSAL', 'SQTT_QUEUE_TYPE_UNKNOWN',
    'SQTT_SA_PER_SE', 'SQTT_VERSION_2_2', 'SQTT_VERSION_2_3',
-    'SQTT_VERSION_2_4', 'SQTT_VERSION_3_2', 'SQTT_VERSION_3_3',
-    'SQTT_VERSION_NONE', 'UserEventObjectName', 'UserEventPop',
-    'UserEventPush', 'UserEventTrigger', 'elf_gfxip_level',
+    'SQTT_VERSION_2_4', 'SQTT_VERSION_3_2', 'SQTT_VERSION_NONE',
+    'UserEventObjectName', 'UserEventPop', 'UserEventPush',
+    'UserEventTrigger', 'elf_gfxip_level',
    'rgp_sqtt_marker_event_type', 'rgp_sqtt_marker_general_api_type',
    'rgp_sqtt_marker_identifier', 'rgp_sqtt_marker_user_event_type',
    'sqtt_api_type', 'sqtt_engine_type',
@@ -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, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32, colored
+from tinygrad.helpers import getenv, to_mv, 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
@@ -130,9 +130,8 @@ class AMDComputeQueue(HWQueue):
      self.wreg(self.gc.regSQ_THREAD_TRACE_USERDATA_2, *data_ints[i:i+2])

  def sqtt_config(self, tracing:bool):
-    trace_ctrl = {'rt_freq': self.soc.SQ_TT_RT_FREQ_4096_CLK} if self.dev.target < (12,0,0) else {}
-    self.wreg(self.gc.regSQ_THREAD_TRACE_CTRL, draw_event_en=1, spi_stall_en=1, sq_stall_en=1, reg_at_hwm=2, hiwater=1, util_timer=1,
-      mode=int(tracing), **trace_ctrl)
+    self.wreg(self.gc.regSQ_THREAD_TRACE_CTRL, draw_event_en=1, spi_stall_en=1, sq_stall_en=1, reg_at_hwm=2, hiwater=1,
+              rt_freq=self.soc.SQ_TT_RT_FREQ_4096_CLK, util_timer=self.soc.SQ_TT_UTIL_TIMER_250_CLK, mode=int(tracing))

  # Magic values from mesa/src/amd/vulkan/radv_sqtt.c:radv_emit_spi_config_cntl and src/amd/common/ac_sqtt.c:ac_sqtt_emit_start
  def sqtt_start(self, buf0s:list[HCQBuffer], se_mask:int):
@@ -141,35 +140,24 @@ class AMDComputeQueue(HWQueue):
    # One buffer for one SE, mesa does it with a single buffer and ac_sqtt_get_data_offset, but this is simpler and should work just as well
    for se in range(len(buf0s)):
      self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, instance_broadcast_writes=1)
-      buf0_lo, buf0_hi = data64_le(buf0s[se].va_addr >> 12)
-      if self.dev.target >= (12,0,0):
-        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, size=buf0s[se].size >> 12)
-        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_LO, base_lo=buf0_lo)
-        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_HI, base_hi=buf0_hi)
-      else:
-        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, base_hi=buf0_hi, size=buf0s[se].size >> 12)
-        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE, base_lo=buf0_lo)
+      buf0_lo, buf0_hi = data64_le(buf0s[se].va_addr>>12)
+      self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, base_hi=buf0_hi, size=buf0s[se].size>>12)
+      self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE, base_lo=buf0_lo)
      # NOTE: SQTT can only trace instructions on one simd per se, this selects first simd in first wgp in first sa.
      # For RGP to display instruction trace it has to see it on first SE. Howerver ACE/MEC/whatever does the dispatching starting with second se,
      # and on amdgpu/non-AM it also does weird things with dispatch order inside se: around 7 times out of 10 it starts from the last cu, but
      # sometimes not, especially if the kernel has more than one wavefront which means that kernels with small global size might get unlucky and
      # be dispatched on something else and not be seen in instruction tracing tab. You can force the wavefronts of a kernel to be dispatched on the
      # CUs you want to by disabling other CUs via bits in regCOMPUTE_STATIC_THREAD_MGMT_SE<x> and trace even kernels that only have one wavefront.
-      cs_wtype = (1 << 6) if self.dev.target >= (12,0,0) else self.soc.SQ_TT_WTYPE_INCLUDE_CS_BIT
-      self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, wtype_include=cs_wtype, simd_sel=0, wgp_sel=0, sa_sel=0)
+      self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, wtype_include=self.soc.SQ_TT_WTYPE_INCLUDE_CS_BIT, simd_sel=0, wgp_sel=0, sa_sel=0)
      reg_include = self.soc.SQ_TT_TOKEN_MASK_SQDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_SHDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_GFXUDEC_BIT | \
-                    self.soc.SQ_TT_TOKEN_MASK_COMP_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT
-      token_exclude = (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT) if self.dev.target < (12,0,0) else 0
-
-      # disable tracing
+                    self.soc.SQ_TT_TOKEN_MASK_COMP_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT
+      token_exclude = 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT
      if not (se_mask >> se) & 0b1:
-        # gfx12 doesn't have enums with all fields, so it's hardcoded, but it's the same as gfx11.
-        token_exclude |= (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT | \
+        token_exclude |= 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT | \
                         1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT | \
-                         1 << self.soc.SQ_TT_TOKEN_EXCLUDE_INST_SHIFT) if self.dev.target < (12,0,0) else 0x927
-
-      token_mask = {} if self.dev.target < (12,0,0) else {'exclude_barrier_wait': 1}
-      self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_include=reg_include, token_exclude=token_exclude, bop_events_token_include=1, **token_mask)
+                         1 << self.soc.SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
+      self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_include=reg_include, token_exclude=token_exclude, bop_events_token_include=1)
      # Enable SQTT
      self.sqtt_config(tracing=True)
    # Restore global broadcasting
@@ -190,6 +178,9 @@ class AMDComputeQueue(HWQueue):
      # Wait for FINISH_PENDING==0
      self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, self.pm4.WAIT_REG_MEM_FUNCTION(WAIT_REG_MEM_FUNCTION_EQ),
                self.gc.regSQ_THREAD_TRACE_STATUS.addr[0], 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), 4)
+      # Wait for FINISH_DONE!=0
+      self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, self.pm4.WAIT_REG_MEM_FUNCTION(WAIT_REG_MEM_FUNCTION_NEQ),
+                self.gc.regSQ_THREAD_TRACE_STATUS.addr[0], 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_done'), 4)
      # Disable SQTT
      self.sqtt_config(tracing=False)
      # Wait for BUSY==0
@@ -670,7 +661,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, 0x7590], bars=[0, 2, 5], vram_bar=0,
+    super().__init__(dev, dev_id, vendor=0x1002, devices=[0x744c, 0x7480, 0x7550], 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 +704,7 @@ class PCIIface(PCIIfaceBase):
  def device_fini(self): self.dev_impl.fini()

 class USBIface(PCIIface):
-  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
+  def __init__(self, dev, dev_id):
    self.dev = dev
    self.usb = ASM24Controller()
    self.bars = setup_pci_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30))
@@ -813,7 +804,7 @@ class AMDDevice(HCQCompiled):
    # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
    self.sqtt_enabled = PROFILE and bool(getenv("SQTT", 0))
    if self.sqtt_enabled:
-      if self.target[0] < 11: raise RuntimeError(f'SQ Thread Tracing is not supported on gc:{self.target}')
+      if self.target[0] != 11: raise RuntimeError(f'SQ Thread Tracing is not supported on gc:{self.target}')
      if not self.is_am() and (ppfeaturemask:=int(FileIOInterface('/sys/module/amdgpu/parameters/ppfeaturemask', os.O_RDONLY).read(), 16))&0x8000:
        raise RuntimeError("SQTT can't be enabled because of hardware bug, to workaround either use AMD_IFACE=PCI or add "
                           f"ppfeaturemask={(ppfeaturemask&~0x8000):#x} (current {ppfeaturemask=:#x} & ~PP_GFXOFF_MASK) to amdgpu module parameters\n"
@@ -875,12 +866,13 @@ 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 = ((wptrs_buf.cpu_view().view(fmt='I')[i] & 0x1FFFFFFF) - (((buf0.va_addr//32) & 0x1FFFFFFF) if self.target < (12,0,0) else 0)) * 32
+        wptr = ((struct.unpack('<I', wptrs[i*4:i*4+4])[0] & 0x1FFFFFFF) - ((buf0.va_addr//32) & 0x1FFFFFFF)) * 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,12 +23,10 @@ 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)
+      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
      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))
+    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.clReleaseProgram(program))
    return bytes(binary)

@@ -99,22 +97,16 @@ 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)))
+      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

    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]
-    self.driver_version = (cl.clGetDeviceInfo(self.device_id, cl.CL_DRIVER_VERSION, 256,
-                                              buf:=ctypes.create_string_buffer(256), None), buf.value.decode())[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
    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]
+    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

    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,9 +10,7 @@ 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:
-    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}")
+  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

 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())
-    ret, resps = None, []
+    resps = []
    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():
@@ -1,137 +0,0 @@
-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  # pylint: disable=method-cache-max-size-none
+  @functools.cache
  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((self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).read, value=1, timeout_ms=timeout,
+    wait_cond(lambda: (self.adev.mmMP1_SMN_C2PMSG_90 if not debug else self.adev.mmMP1_SMN_C2PMSG_54).read(), value=1, timeout_ms=timeout,
      msg=f"SMU msg {msg:#x} timeout")
    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,11 +60,7 @@ 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
-    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()))
+    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
    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,12 +22,10 @@ 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
+  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'(.)((?: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
+  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'(\.)(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 _, sym, _, _ in rels:
+    for roff, sym, r_type_, r_addend 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]

@@ -16,7 +16,7 @@ elif OSX:
 else:
  LLVM_PATH = ctypes.util.find_library('LLVM')
  # use newer LLVM if possible
-  for ver in reversed(range(14, 21+1)):
+  for ver in reversed(range(14, 20+1)):
    if LLVM_PATH is not None: break
    LLVM_PATH = ctypes.util.find_library(f'LLVM-{ver}')
  if LLVM_PATH is None:
@@ -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 # pylint: disable=method-cache-max-size-none
+  @functools.cache
  def lv1(self, size): return size.bit_length()

-  @functools.cache # pylint: disable=method-cache-max-size-none
+  @functools.cache
  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 _, pt, pte_idx, pte_cnt, _ in ctx.next(size):
+    for off, pt, pte_idx, pte_cnt, pte_covers 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,7 +124,6 @@ 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):
@@ -335,7 +334,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)
-    _, self.rm_args_sysmem = self.nvdev._alloc_boot_struct(nv.GSP_ARGUMENTS_CACHED(bDmemStack=True, messageQueueInitArguments=queue_args))
+    rm_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
@@ -482,7 +481,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_sysmem = System.alloc_sysmem(0x5000, contiguous=True)
+      method_va, 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:
@@ -558,7 +557,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 a, m: (self.nvdev.rreg(a) & m), addr, mask, value=val, msg=f"Register {addr:#x} not equal to {val:#x} after polling")
+        wait_cond(lambda: (self.nvdev.rreg(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,8 +152,6 @@ 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}
@@ -165,13 +163,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(), _do_eval(lo), _do_eval(hi))
-        else: self.reg_offsets[name] = (_do_eval(lo), _do_eval(hi))
+        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))
        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], _do_eval(f"lambda {m.groups()[1]}: {fn}")
+        name, value = m.groups()[0], 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.cached_property
+  @functools.cache
  def atomic_lib(self): return ctypes.CDLL(ctypes.util.find_library('atomic')) if sys.platform == "linux" else None

-  @functools.cached_property
+  @functools.cache
  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.cached_property
+  @functools.cache
  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): # pylint: disable=super-init-not-called
+  def __init__(self, usb, addr, size, fmt, pcimem=True):
    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,14 +256,13 @@ 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)

-    # 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)

-    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)
+      _, 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)
@@ -1,9 +1,9 @@
-from typing import Iterator
+from typing import Iterator, Sequence
 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 symbolic, pm_simplify_valid, pm_drop_and_clauses
+from tinygrad.uop.symbolic import sym, symbolic
 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, Ops.ENDRANGE}, name="tr"), realize),
+  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS}, 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)
@@ -41,7 +41,7 @@ class BufferizeOpts:
@dataclass
 class IndexingContext:
  realize_map: dict[UOp, None] = field(default_factory=dict)
-  range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict)
+  range_map: dict[UOp, tuple[list[UOp], list[UOp]]] = field(default_factory=dict)

  # create ranges
  range_idx: Iterator[int] = field(default_factory=itertools.count)
@@ -103,34 +103,30 @@ pm_apply_rangeify = PatternMatcher([
 ])

 # this is the definition of the movement ops
-@functools.cache
-def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
-  match op:
-    case Ops.SHRINK:  rngs = tuple(a if ss == 0 else a+ss for a,(ss,_) in zip(rngs, arg))
-    case Ops.PERMUTE: rngs = tuple(rngs[p] for p in argsort(arg))
-    case Ops.FLIP:    rngs = tuple(((s-1)-a) if f else a for a,s,f in zip(rngs, in_shape, arg))
-    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))
+def apply_movement_op(x:UOp, rngs:Sequence[UOp]) -> list[UOp]:
+  match x.op:
+    case Ops.SHRINK:  rngs = [a if ss == 0 else a+ss for a,(ss,_) in zip(rngs, x.arg)]
+    case Ops.PERMUTE: rngs = [rngs[p] for p in argsort(x.arg)]
+    case Ops.FLIP:    rngs = [((s-1)-a) if f else a for a,s,f in zip(rngs, x.shape, x.arg)]
+    case Ops.EXPAND:  rngs = [a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, x.src[0].shape, x.shape)]
    case Ops.PAD:
      # TODO: why is multiple graph_rewrites faster than one here?
-      # 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))
+      rngs = [r if (s == 0 and e == 0) else graph_rewrite(((r >= s) & (r < (sh-e))).where(r-s, UOp.invalid()), sym, name="pad")
+              for r,sh,(s,e) in zip(rngs, x.shape, x.arg)]
    case Ops.RESHAPE:
      acc = 1
      axes_in:list[UOp] = []
-      for s,src in list(zip(arg, rngs))[::-1]:
+      for s,src in list(zip(x.shape, rngs))[::-1]:
        axes_in.append(acc*src)
        acc *= s
      combined_axes = sum(axes_in, start=UOp.const(dtypes.index, 0))
      axes_out:list[UOp] = []
-      for s in in_shape[::-1]:
+      for s in x.src[0].shape[::-1]:
        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(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")
+      rngs = list(graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic, name="reshape").src)
+    case _: raise RuntimeError(f"{x.op} is not a MovementOp")
  return rngs

@cpu_profile(TracingKey("run_rangeify"), "TINY")
@@ -161,7 +157,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
    consumer_rngs = [rctx.range_map[c][0] for c in consumer_map[x] if c in rctx.range_map]
    if x in rctx.realize_map:
      # if this is in the realize_map, we create new ranges (at the output)
-      out_rngs = tuple(rctx.new_range(s) if not isinstance(s, UOp) or s.op is not Ops.RANGE else s for s in x.shape)
+      out_rngs = [rctx.new_range(s) if not isinstance(s, UOp) or s.op is not Ops.RANGE else s for s in x.shape]
      # all ranges are ended now
      ending_ranges[x] = False
    elif x.op in {Ops.MSTACK, Ops.MSELECT}:
@@ -185,16 +181,15 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:

      # TODO: in RANGEIFY > 1 all_all_same isn't required
      all_all_same = all(same_rngs for _,_,same_rngs in rngs_valids)
-      _out_rngs = []
+      out_rngs = []
      for i,(local_rngs,valids,same_rngs) in enumerate(rngs_valids):
        # we compare the ranges without their valids
        if all_all_same:
          # the new valid is the OR of all the children valids
          minimum_valid = functools.reduce(operator.or_, valids, UOp.const(dtypes.bool, False))
-          _out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid"))
+          out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid"))
        else:
-          _out_rngs.append(rctx.new_range(x.shape[i]))
-      out_rngs = tuple(_out_rngs)
+          out_rngs.append(rctx.new_range(x.shape[i]))

      # we have to realize here if there's new ranges
      if not all_all_same: rctx.realize_map[x] = None
@@ -208,16 +203,18 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
    #  2. newly created for REDUCE_AXIS
    #  3. passed through for everything else

-    rngs = out_rngs  # rngs is the input ranges  # pylint: disable=possibly-used-before-assignment
+    rngs = out_rngs  # rngs is the input ranges

    # apply movement ops
-    if x.op in GroupOp.Movement: rngs = apply_movement_op(x.op, x.src[0].shape, x.arg, rngs)
+    if x.op in GroupOp.Movement: rngs = apply_movement_op(x, rngs)
    # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do.
    if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape): ending_ranges[x] = True

    # REDUCE_AXIS creates ranges for the axes it is reducing
    if x.op is Ops.REDUCE_AXIS:
-      rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) if i in x.arg[1] else r for i,(r,s) in enumerate(zip(rngs, x.src[0].shape)))
+      rngs = rngs[:]
+      for i,s in enumerate(x.src[0].shape):
+        if i in x.arg[1]: rngs[i] = rctx.new_range(s, axistype=AxisType.REDUCE)

    if debug:
      print("***" if x in rctx.realize_map else "   ", len(consumer_map[x]), f"{str(x.op):20s}",
@@ -103,14 +103,14 @@ earliest_rewrites = PatternMatcher([
 # movement op on INDEX as a PatternMatcher
 pm_mops = PatternMatcher([
  (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
-   lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.arg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)),  # type: ignore
+   lambda r,idx: r.src[0].index(*apply_movement_op(r, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)),
 ])

 # *****************
 # 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, Ops.ENDRANGE}
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.NOOP}

 # 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") or b.device.startswith("TINYFS")):
+  if isinstance(b.device, str) and b.device.startswith("DISK"):
    rngs = b.src[1:]
    size = prod(shape := [int(r.vmax+1) for r in rngs])

@@ -338,7 +338,6 @@ 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
@@ -413,7 +412,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([r for r in x.ranges if r.arg[-1] != AxisType.OUTER]): return None
+  if len(x.ranges): return None
  if x.src[0].ptrdtype.addrspace is AddrSpace.LOCAL: return None

  # local kernel rewrite
@@ -425,7 +424,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, Ops.ENDRANGE} else ret.src[1]
+    if ret.src[1].op not in {Ops.COPY, Ops.BUFFER_VIEW} 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]):
@@ -443,7 +442,7 @@ def tag_uop(ctx:list[UOp], x:UOp):
 add_tags = PatternMatcher([
  # don't tag BUFFERs, they are global
  (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND,
-                     Ops.MSTACK, Ops.MSELECT, Ops.RANGE}.union(GroupOp.Movement), name="x"), tag_uop),
+                     Ops.MSTACK, Ops.MSELECT}.union(GroupOp.Movement), name="x"), tag_uop),
  (UPat({Ops.MSTACK, Ops.MSELECT}, name="x"), lambda ctx,x: None if all(s.op is Ops.BUFFER for s in x.src) else tag_uop(ctx, x)),
 ])

@@ -482,11 +481,6 @@ 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] = []
@@ -502,7 +496,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_localize_bufs+pm_limit_bufs, ctx=rctx, name="localize/limit buffers")
+  tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="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
@@ -4,7 +4,14 @@ from dataclasses import dataclass
 from typing import cast, Sequence
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import resolve, UOp, Variable, sint, smax, smin, sint_to_uop, Ops, ssimplify
-from tinygrad.helpers import prod, all_int, flatten
+from tinygrad.helpers import prod, all_int, flatten, ceildiv
+
+# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
+def get_contraction(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) -> list[list[int]]|None:
+  acc_old, acc_new = list(itertools.accumulate(old_shape, operator.mul)), list(itertools.accumulate(new_shape, operator.mul))
+  try: split = [acc_old.index(acc)+1 if acc != 1 else 0 for acc in acc_new]
+  except ValueError: return None
+  return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])]

@functools.cache
 def canonicalize_strides(shape:tuple[sint, ...], strides:tuple[sint, ...]) -> tuple[sint, ...]:
@@ -164,6 +171,7 @@ class View:
    if not all_int(vm1.shape):
      # if all strides are 0 and vm2 is unmasked, return vm1
      if all(x == 0 for x in vm2.strides+vm1.strides) and vm2.mask is None: return vm1
+      # TODO: handle more cases
      return None

    # Project vm1's offset and strides on to vm2.
@@ -176,7 +184,47 @@ class View:
        if not resolve((s1 := s1 - o)!=0): continue  # if s1 can possibly be 0
        terms[d2].append((d1, s1))
        strides[d1] += ssimplify(s1 * vm2.strides[d2])
-    return None
+
+    # Merge dimensions in vm2 if required.
+    # NB: Merging too many dimensions can make it difficult to project vm2's mask, hence only combining when required.
+    idxs: list[UOp] = [UOp.variable(f"idx{i}", 0, s-1, dtypes.index) for i,s in enumerate(vm1.shape)]
+    merged_size, merged_term = 1, UOp.const(dtypes.index, 0)
+    extents: list[tuple[sint, UOp]] = []
+    for term, s, o in zip(reversed(terms), reversed(vm2.shape), reversed(origin)):
+      merged_term += (sum([idxs[d1] * s1 for d1, s1 in term]) + o) * merged_size
+      merged_size *= s
+      if resolve(merged_term < merged_size, False) and resolve(0 <= merged_term, False):
+        extents.append((merged_size, merged_term))
+        merged_size, merged_term = 1, UOp.const(dtypes.index, 0)
+    if resolve(merged_term != 0): return None
+    if (vm2_shape := tuple(s for s,_ in reversed(extents))) != vm2.shape:
+      if (reshaped_vm2 := vm2.reshape(vm2_shape)) is None: return None
+      # NOTE: this != to prevent infinite loop
+      if reshaped_vm2.shape != vm2.shape: return reshaped_vm2 + vm1
+
+    if vm2.mask:
+      # Try to project vm2's mask on to vm1.
+      newb, newe, bad = [0] * len(vm1.shape), list(vm1.shape), False
+      for (b, e), o, term, (_, t) in zip(vm2.mask, origin, terms, reversed(extents)):
+        if resolve(b <= (t := t.simplify()).vmin and t.vmax < e, False): continue
+        if len(term) != 1:
+          if not term and newe:
+            # t should be a constant if no terms contribute to this dimension, but it might not be simplified
+            if t.vmin != t.vmax: return None
+            newe[0] = 0
+          else: bad = True
+          continue
+        d1, s1 = term[0]
+        newb[d1] = smax(newb[d1], ceildiv(b - o if s1 > 0 else e - o - 1, s1))
+        newe[d1] = smin(newe[d1], (b - o if s1 < 0 else e - o - 1) // s1 + 1)
+
+      # If any of vm1 was masked off, try again with that mask in place.
+      if any((b, e) != (0, s) for b, e, s in zip(newb, newe, vm1.shape)):
+        return vm2 + View.create(vm1.shape, vm1.strides, vm1.offset, tuple(zip(newb, newe)))
+      # Otherwise if vm2's mask was violated, then cannot merge.
+      if bad: return None
+
+    return View.create(vm1.shape, tuple(strides), ssimplify(sum(o * s for o, s in zip(origin, vm2.strides)) + vm2.offset))

  def __unsafe_resize(self, arg: tuple[tuple[sint, sint], ...], mask=None) -> View:
    offset = sum([s * x[0] for s, x in zip(self.strides,arg)])
@@ -244,6 +292,7 @@ class View:

    r_strides, r_new_shape = [], reversed(new_shape)
    for merged_size, new_stride, real_size in reversed(merge_dims(self.shape, self.strides, self.mask)):
+      # TODO: write with get_contraction
      acc = 1
      # TODO: third resolve shouldn't be needed
      while resolve(acc <= merged_size) and resolve(acc != merged_size) and resolve((new_dim := next(r_new_shape, 0)) > 0):
@@ -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.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.ops import smax, smin, resolve, UOp, Ops, sint, MathTrait, 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,59 +411,6 @@ 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)
@@ -1212,13 +1159,12 @@ 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():
@@ -2538,20 +2484,17 @@ 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_})"
+    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

    # 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))])
+      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

      # 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)
+      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
      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
@@ -2562,8 +2505,7 @@ class Tensor(MathTrait):
    # TODO: stride == dilation
    # use padding to round up to 4x4 output tiles
    # (bs, cin_, tyx, HWI)
-    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)
+    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
    # 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
@@ -2685,8 +2627,7 @@ 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)
-    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))
+    return {Ops.ADD: Tensor.__add__, Ops.MAX: Tensor.maximum, Ops.MUL: Tensor.__mul__}[op](fix(ret), fix(base))

  def cumsum(self, axis:int=0) -> Tensor:
    """
@@ -3725,7 +3666,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:Tensor|int, reverse=False) -> Tensor:
+  def lshift(self, x:int, reverse=False) -> Tensor:
    """
    Computes left arithmetic shift of `self` by `x` bits. `self` must have unsigned dtype.
    Equivalent to `self << x`.
@@ -3737,7 +3678,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:Tensor|int, reverse=False) -> Tensor:
+  def rshift(self, x:int, reverse=False) -> Tensor:
    """
    Computes right arithmetic shift of `self` by `x` bits. `self` must have unsigned dtype.
    Equivalent to `self >> x`.
@@ -3853,20 +3794,18 @@ class Tensor(MathTrait):
  def __rpow__(self, x) -> Tensor: return self.pow(x, True)
  def __rmatmul__(self, x) -> Tensor: return self.matmul(x, True)

-  def __ifloordiv__(self, x) -> Tensor: return self.assign(self.__floordiv__(x))
+  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 __imatmul__(self, x) -> Tensor: return self.assign(self.matmul(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 __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))

  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,17 +1,15 @@
-from typing import TypeVar
 from tinygrad.uop import Ops
-from tinygrad.dtype import dtypes, ConstType
+from tinygrad.helpers import T
+from tinygrad.dtype import dtypes

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

  # great functions you get!
-  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 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 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=}")
@@ -20,7 +18,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:TMT, x:TMT|ConstType, reverse:bool=False):
+  def add(self, x, reverse=False):
    """
    Adds `self` and `x`.
    Equivalent to `self + x`.
@@ -38,7 +36,7 @@ class MathTrait:
    ```
    """
    return self._binop(Ops.ADD, x, reverse)
-  def mul(self:TMT, x:TMT|ConstType, reverse:bool=False):
+  def mul(self, x, reverse=False):
    """
    Multiplies `self` and `x`.
    Equivalent to `self * x`.
@@ -57,7 +55,7 @@ class MathTrait:
    ```
    """
    return self._binop(Ops.MUL, x, reverse)
-  def bitwise_and(self:TMT, x:TMT|ConstType, reverse:bool=False):
+  def bitwise_and(self, x, reverse=False):
    """
    Computes the bitwise AND of `self` and `x`.
    Equivalent to `self & x`.
@@ -71,7 +69,7 @@ class MathTrait:
    """
    self._check_dtype()
    return self._binop(Ops.AND, x, reverse)
-  def bitwise_or(self:TMT, x:TMT|ConstType, reverse:bool=False):
+  def bitwise_or(self, x, reverse=False):
    """
    Computes the bitwise OR of `self` and `x`.
    Equivalent to `self | x`.
@@ -85,7 +83,7 @@ class MathTrait:
    """
    self._check_dtype()
    return self._binop(Ops.OR, x, reverse)
-  def bitwise_xor(self:TMT, x:TMT|ConstType, reverse:bool=False):
+  def bitwise_xor(self, x, reverse=False):
    """
    Computes bitwise xor of `self` and `x`.
    Equivalent to `self ^ x`.
@@ -100,7 +98,7 @@ class MathTrait:
    """
    self._check_dtype()
    return self._binop(Ops.XOR, x, reverse)
-  def idiv(self:TMT, x:TMT|ConstType, reverse:bool=False):
+  def idiv(self, x, reverse=False):
    """
    Divides `self` by `x`.
    Equivalent to `self // x`.
@@ -112,61 +110,61 @@ class MathTrait:
    ```
    """
    return self._binop(Ops.IDIV, x, reverse)
-  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 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 __neg__(self): return self.neg()

-  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 __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 __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 __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 __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 __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 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]
+  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)
  # NOTE: __eq__ isn't overridden, and means the same thing as is by default

-  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 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 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)
+  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)
    raise RuntimeError("where needs at least one UOp arg")
-  def threefry(self:TMT, seed:TMT): return self.alu(Ops.THREEFRY, seed)
+  def threefry(self, seed): return self.alu(Ops.THREEFRY, seed)
  def reciprocal(self): return self.alu(Ops.RECIP)
  def trunc(self): return self.alu(Ops.TRUNC)
  def sqrt(self): return self.alu(Ops.SQRT)
  def sin(self): return self.alu(Ops.SIN)
  def log2(self): return self.alu(Ops.LOG2)
  def exp2(self): return self.alu(Ops.EXP2)
-  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)
+  def pow(self, x): return self.alu(Ops.POW, self.ufix(x))
+  def __pow__(self, x): return self.pow(x)
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	88f1d82bed	might end if	2025-10-10 16:40:30 +08:00
geohot	b9019b566f	replace linearizer with toposort	2025-10-10 16:02:28 +08:00