something

just the new devectorizer
viz/cli: support passing only a start marker (#16802 )
2026-06-30 11:59:32 -07:00 · 2026-06-30 11:19:52 -07:00 · 2026-06-30 17:06:32 +09:00 · 2026-06-29 19:12:10 -07:00 · 2026-06-29 18:37:49 -07:00 · 2026-06-29 20:11:16 -04:00
136 changed files with 2309 additions and 2408 deletions
@@ -81,85 +81,181 @@ jobs:
  #      source /tmp/tinygrad_pytest_ci/bin/activate
  #      pytest -nauto --durations=20

-  testmacbenchmark:
-    name: Mac Benchmark
-    runs-on: [self-hosted, macOS]
+  sharedbenchmarks:
+    name: Benchmark (DEV=${{ matrix.dev }})
+    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    strategy:
+      fail-fast: false
+      matrix:
+        dev: ['METAL', 'AMD', 'NV']
    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
+    env:
+      DEV: ${{ matrix.dev }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
-    - name: Symlink models and datasets
+    - name: Setup (AMD)
+      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        mkdir -p weights
-        mkdir -p extra/disassemblers
-        ln -s ~/tinygrad/extra/disassemblers/applegpu extra/disassemblers/applegpu
-        ln -s ~/tinygrad/weights/sd-v1-4.ckpt weights/sd-v1-4.ckpt
-        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
-        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
-      run: python3.11 test/external/process_replay/reset.py
-    - name: Print macOS version
-      run: sw_vers
-    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
-    - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing
-    - name: Run Stable Diffusion v2
-      # TODO: very slow step time
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing
-    # process replay can't capture this, the graph is too large
-    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing
-    - name: Run model inference benchmark
-      run: DEV=METAL NOCLANG=1 python3.11 test/external/external_model_benchmark.py
-    - name: Test speed vs torch
-      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py
+      run: python3 test/external/process_replay/reset.py
    - name: Test tensor cores
-      run: DEV=METAL python3.11 test/opt/test_tensor_cores.py
-    - name: Run Tensor Core GEMM (float)
-      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py
-    - name: Run Tensor Core GEMM (half)
-      run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py
-    - name: Run Tensor Core GEMM (bfloat16)
-      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py
-    - name: Fuzz Padded Tensor Core GEMM
-      run: DEV=METAL M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3.11 ./extra/gemm/fuzz_matmul.py
+      run: |
+        if [[ "${{ matrix.dev }}" == "METAL" ]]; then
+          python3 test/opt/test_tensor_cores.py
+          DEBUG=2 SHOULD_USE_TC=1 python3 extra/gemm/simple_matmul.py
+          DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3 extra/gemm/simple_matmul.py
+          DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3 extra/gemm/simple_matmul.py
+          M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3 ./extra/gemm/fuzz_matmul.py
+        elif [[ "${{ matrix.dev }}" == "NV" ]]; then
+          ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+          DEV=NV:PTX ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+          DEV=CUDA SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+          DEV=CUDA SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+          DEV=CUDA SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+          DEV=CUDA SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+          DEV=NV:PTX SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+          SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+          # TODO: too slow
+          # M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+          # DEV=NV:PTX M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+        else
+          python3 test/opt/test_tensor_cores.py
+          # TODO: this is flaky
+          # DEV=AMD:LLVM python3 test/opt/test_tensor_cores.py
+          SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+          SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+          # TODO: AMD compiler bug causes this to fail
+          # HSA=1 M_START=12 M_STOP=20 M_STEP=1 N_START=12 N_STOP=20 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 DEBUG=2 python3 ./extra/gemm/fuzz_matmul.py
+        fi
+    - name: Run model inference benchmark
+      # TODO: unstable on AMD
+      if: ${{ matrix.dev != 'AMD' }}
+      run: CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
+    - name: Test speed vs torch
+      # TODO: unstable on AMD
+      if: ${{ matrix.dev != 'AMD' }}
+      env:
+        HALF: ${{ matrix.dev == 'NV' && '1' || '0' }}
+      run: CAPTURE_PROCESS_REPLAY=0 BIG=2 ${{ matrix.dev == 'METAL' && 'MPS=1' || 'TORCHCUDA=1' }} python3 test/speed/external_test_speed_v_torch.py
+    - name: Run Stable Diffusion
+      env:
+        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'METAL' && '720' || matrix.dev == 'AMD' && '550' || '0' }}
+      run: BENCHMARK_LOG=stable_diffusion python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
+    - name: Run SDXL
+      if: ${{ matrix.dev != 'NV' }}
+      env:
+        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'METAL' && '5000' || matrix.dev == 'AMD' && '3200' || '2000' }}
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
    - name: Run llama3.2
-      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
+      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
+    - name: Run qwen3.5
+      # qwen3.5:35b-a3b doesn't fit on mac
+      if: ${{ matrix.dev != 'METAL' }}
+      run: BENCHMARK_LOG=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.5:35b-a3b --benchmark --warmup
    - name: Run olmoe
-      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 -m tinygrad.llm -m olmoe --benchmark --warmup
+      # just metal for now
+      if: ${{ matrix.dev == 'METAL' }}
+      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
    - name: Train MNIST
-      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py
-
-    # NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
-    # the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
-    #- name: Run 10 CIFAR training steps
-    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py
-    #- name: Run 10 CIFAR training steps w HALF
-    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py
-
-    #- name: Run 10 CIFAR training steps w BF16
-    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py
-    # TODO: too slow
-    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py
+      run: time TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
+    - name: Run 10 CIFAR training steps
+      env:
+        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'NV' && '130' || matrix.dev == 'AMD' && '200' || '3000' }}
+      run: BENCHMARK_LOG=cifar_10steps STEPS=10 python3 examples/hlb_cifar10.py
+    - name: Run 10 CIFAR training steps w HALF
+      env:
+        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'NV' && '120' || matrix.dev == 'AMD' && '230' || '3000' }}
+      run: BENCHMARK_LOG=cifar_10steps_half STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    - uses: actions/upload-artifact@v7
+      if: ${{ matrix.dev != 'AMD' }}
      with:
-        name: Speed (Mac)
+        name: Speed (${{ matrix.dev }})
        path: |
          onnx_inference_speed.csv
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

+  tinyboxbenchmark:
+    name: Tinybox Benchmark (${{ matrix.dev }})
+    runs-on: [self-hosted, Linux, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    strategy:
+      fail-fast: false
+      matrix:
+        dev: ['AMD', 'NV']
+    timeout-minutes: 60
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    env:
+      DEV: ${{ matrix.dev }}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup (AMD)
+      if: ${{ matrix.dev == 'AMD' }}
+      run: |
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        mkdir -p extra/datasets
+        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
+        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
+    - name: Test tiny
+      run: |
+        DEBUG=2 python -m pytest -rA test/test_tiny.py
+        if [[ "${{ matrix.dev }}" == "NV" ]]; then
+          DEBUG=2 DEV=CUDA python -m pytest -rA test/test_tiny.py
+        fi
+    - name: Test speed vs theoretical
+      run: IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+    - name: Test benchmark allreduce
+      if: ${{ matrix.dev == 'NV' }}
+      run: python test/external/external_benchmark_multitensor_allreduce.py
+    - name: HEVC Decode Benchmark
+      if: ${{ matrix.dev == 'NV' }}
+      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
+    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
+      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run full CIFAR training steps w 6 GPUS
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run MLPerf resnet eval on training data
+      run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf Bert training steps (6 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
  testusbgpu:
    name: UsbGPU Benchmark
    runs-on: [self-hosted, macOS]
@@ -196,337 +292,6 @@ jobs:
    - name: UsbGPU (USB4/TB) tiny tests
      run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py

-  testnvidiabenchmark:
-    name: tinybox green Benchmark
-    runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 60
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Print nvidia-smi
-      run: nvidia-smi
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
-        mkdir -p extra/datasets
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
-    - name: Run model inference benchmark
-      run: DEV=NV CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
-    - name: Test speed vs torch
-      run: DEV=NV CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
-    - name: Test speed vs theoretical
-      run: DEV=NV IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
-    - name: Test benchmark allreduce
-      run: DEV=NV python test/external/external_benchmark_multitensor_allreduce.py
-    - name: Test tensor cores
-      run: |
-        DEV=NV ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
-        DEV=NV:PTX ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
-    - name: Run Tensor Core GEMM (CUDA)
-      run: |
-        DEV=CUDA SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-        DEV=CUDA SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-        DEV=CUDA SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
-        DEV=CUDA SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-    - name: Run Tensor Core GEMM (PTX)
-      run: DEV=NV:PTX SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-    - name: Run Tensor Core GEMM (NV)
-      run: DEV=NV SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-    - name: Test DEV=NV
-      run: DEBUG=2 DEV=NV python -m pytest -rA test/test_tiny.py
-    - name: Test DEV=CUDA
-      run: DEBUG=2 DEV=CUDA python -m pytest -rA test/test_tiny.py
-    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion DEV=NV python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
-    # TODO: too slow
-    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
-    - name: Run llama3.2
-      run: DEV=NV BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.5
-      run: DEV=NV BENCHMARK_LOG=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 -m tinygrad.llm -m qwen3.5:35b-a3b --benchmark --warmup
-    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu DEV=NV JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
-    # - name: Run LLaMA-3 8B on 6 GPUs
-    #   run: DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
-    # - name: Run LLaMA-2 70B
-    #   run: DEV=NV CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
-    - uses: actions/upload-artifact@v7
-      with:
-        name: Speed (NVIDIA)
-        path: |
-          onnx_inference_speed.csv
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testmorenvidiabenchmark:
-    name: tinybox green Training Benchmark
-    runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 60
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
-        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
-        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
-        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
-        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
-        mkdir -p extra/datasets
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
-    # TODO: too slow
-    # - name: Fuzz Padded Tensor Core GEMM (NV)
-    #   run: DEV=NV M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
-    # TODO: too slow
-    # - name: Fuzz Padded Tensor Core GEMM (PTX)
-    #   run: DEV=NV:PTX M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
-    - name: HEVC Decode Benchmark
-      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 DEV=NV PYTHONPATH=. python3 extra/hevc/decode.py
-    - name: Train MNIST
-      run: time PYTHONPATH=. DEV=NV TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
-    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=130 DEV=NV STEPS=10 python3 examples/hlb_cifar10.py
-    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
-    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
-    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 DEV=NV WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar DEV=NV DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEV=NV DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    - name: Run MLPerf resnet eval on training data
-      run: time BENCHMARK_LOG=resnet_eval DEV=NV MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps DEV=NV DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf Bert training steps (6 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testamdbenchmark:
-    name: tinybox red Benchmark
-    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setcap to python
-      run: ./extra/amdpci/setup_python_cap.sh
-    - name: Remove amd modules
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
-    - name: Kill stale pids
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
-    #- name: Insert amdgpu
-    #  run: sudo modprobe amdgpu
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
-        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
-        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
-        mkdir -p extra/datasets
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
-    #- name: setup perflevel
-    #  run: |
-    #    examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
-    #    rocm-smi
-    #- name: Show off tinybox
-    #  run: /opt/rocm/bin/rocm-bandwidth-test
-    # TODO: unstable on AMD
-    #- name: Run model inference benchmark
-    #  run: LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 NOCLANG=1 python3 test/external/external_model_benchmark.py
-    # TODO: unstable on AMD
-    #- name: Test speed vs torch
-    #  run: |
-    #    python3 -c "import torch; print(torch.__version__)"
-    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
-    - name: Test speed vs theoretical
-      run: DEV=AMD IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
-    - name: Test tensor cores (no LLVM)
-      run: DEV=AMD python3 test/opt/test_tensor_cores.py
-    # TODO: this is flaky
-    # - name: Test tensor cores AMD:LLVM
-    #   run: DEV=AMD:LLVM python3 test/opt/test_tensor_cores.py
-    - name: Run Tensor Core GEMM (AMD)
-      run: |
-        DEV=AMD SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-        DEV=AMD SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
-    - name: Test DEV=AMD
-      run: DEBUG=2 DEV=AMD python -m pytest -rA test/test_tiny.py
-    #- name: Test HIP=1
-    #  run: DEBUG=2 HIP=1 python -m pytest -rA test/test_tiny.py
-    # TODO: AMD compiler bug causes this to fail
-    #- name: Fuzz Padded Tensor Core GEMM
-    #  run: HSA=1 M_START=12 M_STOP=20 M_STEP=1 N_START=12 N_STOP=20 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 DEBUG=2 python3 ./extra/gemm/fuzz_matmul.py
-    #- name: Remove amdgpu
-    #  run: sleep 10 && sudo rmmod amdgpu # sleep a bit to let the driver unload the prev pid.
-    - name: Test AM cold start time
-      run: time DEV=AMD AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
-    - name: Test AM warm start time
-      run: time DEV=AMD python3 test/test_tiny.py TestTiny.test_plus
-    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 DEV=AMD python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
-    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 DEV=AMD python3 examples/sdxl.py --seed 0 --noshow --timing
-    - name: Run llama3.2
-      run: DEV=AMD BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.5
-      run: DEV=AMD BENCHMARK_LOG=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 -m tinygrad.llm -m qwen3.5:35b-a3b --benchmark --warmup
-    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu DEV=AMD JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
-    # - name: Run LLaMA-3 8B on 6 GPUs
-    #   run: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
-    #- name: Restore amdgpu
-    #  run: sudo modprobe amdgpu
-    # - name: Run LLaMA-2 70B
-    #   run: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testmoreamdbenchmark:
-    name: tinybox red Training Benchmark
-    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setcap to python
-      run: ./extra/amdpci/setup_python_cap.sh
-    - name: Remove amd modules
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
-    - name: Kill stale pids
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
-        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
-        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
-        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
-        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
-        mkdir -p extra/datasets
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
-    - name: Test GPU crash recovery
-      run: DEV=AMD python3 -m pytest -rA test/external/external_test_gpu_crash.py
-    - name: Train MNIST
-      run: time PYTHONPATH=. DEV=AMD TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
-    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 DEV=AMD STEPS=10 python3 examples/hlb_cifar10.py
-    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=230 DEV=AMD STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
-    # - name: Run 10 CIFAR training steps w BF16
-    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 DEV=AMD STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
-    # TODO: too slow
-    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 DEV=AMD WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar DEV=AMD DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu DEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    # TODO: broken on some of the machines
-    #- name: Test full tinyfs load
-    #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testmlperfamdbenchmark:
-    name: tinybox red MLPerf Benchmark
-    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setcap to python
-      run: ./extra/amdpci/setup_python_cap.sh
-    - name: Remove amd modules
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
-    - name: Kill stale pids
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
-        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
-        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
-        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
-        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
-        mkdir -p extra/datasets
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
-    - name: Run MLPerf resnet eval
-      run: time BENCHMARK_LOG=resnet_eval DEV=AMD MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps DEV=AMD DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf Bert training steps (6 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
  testcommalatest:
    name: comma Benchmark (0.11.0)
    runs-on: [self-hosted, Linux, comma]
@@ -550,7 +315,7 @@ jobs:
    - name: openpilot compile3 0.11.0 driving_vision (from pickle)
      run: BENCHMARK_LOG=openpilot_0_11_0_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
    - name: IR3 openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.11.0 driving_policy
      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.11.0 dmonitoring
@@ -648,31 +413,28 @@ jobs:
    - name: openpilot run_pickle 0.10.1 driving_vision
      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py

-  testreddriverbenchmark:
-    name: AM Benchmark
+  driverbenchmarks:
+    name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
    runs-on: [self-hosted, Linux, tinyboxrandom]
+    strategy:
+      fail-fast: false
+      matrix:
+        dev: ['AMD', 'NV']
    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
+    env:
+      DEV: ${{ matrix.dev }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
-    - name: Setcap to python
-      run: ./extra/amdpci/setup_python_cap.sh
-    - name: Remove amd modules
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
-    - name: Kill stale pids
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
-    - name: Symlink models and datasets
+    - name: Setup
      run: |
-        mkdir -p weights
-        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
-        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
-        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
-        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
-        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} rmmod
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} kill_pids
        mkdir -p extra/datasets
        ln -s /raid/datasets/imagenet extra/datasets/imagenet
    - name: setup staging db
@@ -683,104 +445,47 @@ jobs:
    - name: reset process replay
      run: test/external/process_replay/reset.py
    - name: Test driver cold start time
-      run: time DEBUG=3 DEV=AMD AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
+      run: time DEBUG=3 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test driver warm start time
-      run: time DEBUG=3 DEV=AMD python3 test/test_tiny.py TestTiny.test_plus
+      if: ${{ matrix.dev == 'AMD' }}
+      run: time DEBUG=3 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test GPU crash recovery
-      run: DEV=AMD python3 -m pytest -rA test/external/external_test_gpu_crash.py
-    # Fails on 9070
-    # - name: Test tensor cores
-    #   run: |
-    #     DEV=AMD python3 test/test_linearizer.py test/opt/test_tensor_cores.py
-    #     DEV=AMD:LLVM python3 test/test_linearizer.py test/opt/test_tensor_cores.py
-    #     DEV=AMD SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-    - name: Run Tensor Core GEMM (AMD)
-      run: DEV=AMD SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
-    - name: Test DEV=AMD
-      run: DEBUG=2 DEV=AMD python -m pytest -rA test/test_tiny.py
+      if: ${{ matrix.dev == 'AMD' }}
+      run: python3 -m pytest -rA test/external/external_test_gpu_crash.py
+    - name: Test tensor cores
+      run: |
+        if [[ "${{ matrix.dev }}" == "AMD" ]]; then
+          # Fails on 9070
+          # python3 test/test_linearizer.py test/opt/test_tensor_cores.py
+          # DEV=AMD:LLVM python3 test/test_linearizer.py test/opt/test_tensor_cores.py
+          # SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+          SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+        else
+          ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+        fi
    - name: Test DISK copy time
-      run: DEV=AMD TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
+      run: TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
    - name: Test CPU copy time
      run: |
-        DEV=AMD GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
-        DEV=AMD GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
+        GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
+        GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar DEV=AMD DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #   run: BENCHMARK_LOG=resnet_10steps DEV=AMD MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
+      run: time BENCHMARK_LOG=cifar DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      if: ${{ matrix.dev == 'NV' }}
+      run: BENCHMARK_LOG=resnet_10steps MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: BENCHMARK_LOG=bert_10steps CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Remote
      run: |
        pkill -f 'extra/remote/serve.py' || true
        PYTHONPATH=. python3 extra/remote/serve.py 6482 &
        sleep 1
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 DEV=PCI+AMD python3 test/test_tiny.py
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 DEV=PCI+AMD AMD_AQL=1 python3 test/test_tiny.py
-        pkill -f 'extra/remote/serve.py' || true
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testgreendriverbenchmark:
-    name: NV Benchmark
-    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setcap to python
-      run: ./extra/amdpci/setup_python_cap.sh
-    - name: Remove nv modules
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv rmmod
-    - name: Kill stale pids
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
-        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
-        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
-        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
-        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
-        mkdir -p extra/datasets
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
-    - name: Test driver start time
-      run: time DEBUG=3 DEV=NV python3 test/test_tiny.py TestTiny.test_plus
-    - name: Test tensor cores
-      run: DEV=NV ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
-    - name: Test DISK copy time
-      run: DEV=NV TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
-    - name: Test CPU copy time
-      run: |
-        DEV=NV GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
-        DEV=NV GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
-    - name: Test LLAMA-3
-      run: BENCHMARK_LOG=llama3_beam DEV=NV JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0
-    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar DEV=NV DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps DEV=NV MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf Bert training steps (1 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Remote
-      run: |
-        pkill -f 'extra/remote/serve.py' || true
-        PYTHONPATH=. python3 extra/remote/serve.py 6483 &
-        sleep 1
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6483 DEV=NV python3 test/test_tiny.py
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 python3 test/test_tiny.py
+        if [[ "${{ matrix.dev }}" == "AMD" ]]; then
+          DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD_AQL=1 python3 test/test_tiny.py
+        fi
        pkill -f 'extra/remote/serve.py' || true
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
@@ -1,8 +1,8 @@
 name: Run MLPerf Training

 on:
-  schedule:
-    - cron: '5 8 * * *'  # Runs at 08:05 UTC (12:05 AM Pacific Time)
+  #schedule:
+  #  - cron: '5 8 * * *'  # Runs at 08:05 UTC (12:05 AM Pacific Time)
  push:
    branches:
      - update_mlperf
@@ -77,45 +77,14 @@ jobs:
      run: |
        sudo apt update || true
        sudo apt install -y --no-install-recommends ninja-build
-    - name: Test one op
-      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test ResNet-18
      run: DEBUG=2 python3 extra/torch_backend/example.py
-    - name: custom tests
-      run: python3 -m pytest -n auto extra/torch_backend/test.py --durations=20
    - name: Test one op in torch tests
      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
    - name: Test Ops with TINY_BACKEND
      run: DEV=CPU:LLVM LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/backend/test_ops.py --durations=20
-    - name: Test in-place operations on views
-      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
-    - name: Test multi-gpu
-      run: DEV=CPU:LLVM GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
-    - name: Test kernel fusion
-      run: python3 extra/torch_backend/test_kernel_fusion.py
-
-
-  torchbackendmore:
-    name: Torch Backend Tests More
-    runs-on: *linux
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_unit
-        llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
-    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: STEPS=20 DEV=CPU TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
-    - name: Test some torch tests (expect failure)
-      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+    - name: Custom tests
+      run: DEV=CPU:LLVM GPUS=4 TINY_BACKEND=1 python3 -m pytest -nauto extra/torch_backend/test.py extra/torch_backend/test_inplace.py extra/torch_backend/test_multigpu.py extra/torch_backend/test_kernel_fusion.py --durations=20

  bepython:
    name: Python Backend
@@ -327,7 +296,7 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1468 ALLOWED_GATED_READ_IMAGE=10 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=55 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp32 (test correctness)
        run: |
          DEV=CL IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
@@ -370,7 +339,6 @@ jobs:
        with:
          key: optim
          deps: testing
-          pydeps: "tensorflow==2.19"
          opencl: 'true'
      #- name: Test Optimization Helpers
      #  run: DEBUG=1 python3 extra/optimization/test_helpers.py
@@ -379,7 +347,7 @@ jobs:
      - name: Test Beam Search
        run: DEV=CL IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
      - name: Test MLPerf stuff
-        run: DEV=CL python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
+        run: DEV=CL python -m pytest -n=auto test/external/external_test_lr_schedule.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
      - name: DEV=NULL beautiful_mnist_multigpu
        run: DEV=NULL NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py
      - name: Test Bert training
@@ -685,6 +653,33 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

+  testmetal:
+    strategy:
+      fail-fast: false
+      matrix:
+        group: [1, 2]
+    name: MacOS (DEV=METAL) (${{ matrix.group }})
+    runs-on: *macos
+    timeout-minutes: 20
+    env:
+      DEV: METAL
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: macos-metal
+        deps: testing_unit
+    - name: Check Device.DEFAULT and print some source
+      run: |
+        python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL'"
+        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+    - name: Run backend tests
+      run: python -m pytest -n=auto test/backend --durations=20 --splits 2 --group ${{ matrix.group }}
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
  testmacos:
    strategy:
      fail-fast: false
@@ -693,7 +688,6 @@ jobs:
          - 'CPU:CLANG'
          - 'CPU:LLVM'
          - 'CPU:LVP'
-          - 'METAL'
          - 'WEBGPU'

    name: MacOS (DEV=${{ matrix.dev }})
@@ -716,8 +710,8 @@ jobs:
      run: |
        python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
-    - name: Run backend tests
-      run: python -m pytest -n=auto test/backend --durations=20
+    - name: Run test_tiny
+      run: python -m pytest -n=auto test/test_tiny.py --durations=20
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -11,7 +11,7 @@ X_train -= X_train.mean()
 # *****
 # 1. Define an MNIST model.

-from tinygrad import Tensor
+from tinygrad import Tensor, Context

 l1 = Tensor.kaiming_uniform(128, 784)
 l2 = Tensor.kaiming_uniform(10, 128)
@@ -24,11 +24,11 @@ l1n, l2n = l1.numpy(), l2.numpy()
 from tinygrad.nn.optim import SGD
 optim = SGD([l1, l2])

-Tensor.training = True
-X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
-optim.zero_grad()
-model(X).sparse_categorical_crossentropy(Y).backward()
-optim.schedule_step()   # this will step the optimizer without running realize
+with Context(TRAINING=1):
+  X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
+  optim.zero_grad()
+  model(X).sparse_categorical_crossentropy(Y).backward()
+  optim.schedule_step()   # this will step the optimizer without running realize

 # *****
 # 3. Create a schedule (linear uop).
@@ -67,8 +67,7 @@ def example_2_hip(a:Tensor, correct):
    # the sink specifies the GLOBAL and LOCAL sizes, along with the input buffers and name
    sink = UOp.sink(UOp.special(GLOBALS, 'gidx0'), UOp.special(THREADS, 'lidx0'), out, buf,
                    arg=KernelInfo(name="hip_reduce_sum_kernel"))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT),
-                UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
  eval_harness("HIP kernel", a, lambda x: Tensor.empty(GLOBALS).custom_kernel(x, fxn=hip_reduce_sum)[0].sum(), check=correct)

 def example_3_custom_uop(a:Tensor, correct):
@@ -123,8 +122,7 @@ def example_5_custom_assembly(a:Tensor, correct):
        offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
        if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
        inst.simm16 = offset_dwords
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT),
-                                   UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))

  CU_COUNT = 32
  LANES = 64
@@ -24,7 +24,7 @@ You will see `CUDA` here on a GPU instance, or `CPU` here on a CPU instance.
 We'll use the model from [the Keras tutorial](https://keras.io/examples/vision/mnist_convnet/).

 ```python
-from tinygrad import Tensor, nn
+from tinygrad import Tensor, nn, Context

 class Model:
  def __init__(self):
@@ -74,8 +74,8 @@ We'll use the Adam optimizer. The `nn.state.get_parameters` will walk the model
 ```python
 optim = nn.optim.Adam(nn.state.get_parameters(model))
 batch_size = 128
+@Context(TRAINING=1)
 def step():
-  Tensor.training = True  # makes dropout work
  samples = Tensor.randint(batch_size, high=X_train.shape[0])
  X, Y = X_train[samples], Y_train[samples]
  optim.zero_grad()
@@ -143,7 +143,6 @@ Since we are just randomly sampling from the dataset, there's no real concept of
 for step in range(7000):
  loss = jit_step()
  if step%100 == 0:
-    Tensor.training = False
    acc = (model(X_test).argmax(axis=1) == Y_test).mean().item()
    print(f"step {step:4d}, loss {loss.item():.2f}, acc {acc*100.:.2f}%")
 ```
@@ -165,7 +165,7 @@ from extra.datasets import fetch_mnist
 Now we have everything we need to start training our neural network.
 We will be training for 1000 steps with a batch size of 64.

-We use `with Context(TRAINING=1)` to set the internal flag `Tensor.training` to `True` during training.
+We use `with Context(TRAINING=1)` to enable training mode.
 Upon exit, the flag is restored to its previous value by the context manager.

 ```python
@@ -182,7 +182,7 @@ if __name__ == "__main__":
  # *** run on Machine and compare ***

  # NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
-  src = eval(prg.src[3].arg)
+  src = eval(prg.src[2].arg)
  max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
  print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else "       <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
  machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)
@@ -10,7 +10,7 @@ from extra.lr_scheduler import OneCycleLR
 from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit, Variable
 from tinygrad.nn.state import get_state_dict
 from tinygrad.nn import optim
-from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod
+from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod, TRAINING
 from extra.bench_log import BenchEvent, WallTimeEvent

 cifar_mean = [0.4913997551666284, 0.48215855929893703, 0.4465309133731618]
@@ -44,7 +44,7 @@ class UnsyncedBatchNorm:
    return ret.reshape(x.shape).cast(x.dtype)

  def calc_stats(self, x:Tensor):
-    if Tensor.training:
+    if TRAINING:
      # This requires two full memory accesses to x
      # https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh
      # There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm
@@ -363,7 +363,7 @@ def train_cifar():
    st = time.monotonic()
    while i <= STEPS:
      if i % getenv("EVAL_STEPS", STEPS) == 0 and i > 1 and not getenv("DISABLE_BACKWARD"):
-        # Use Tensor.training = False here actually bricks batchnorm, even with track_running_stats=True
+        # Using Context(TRAINING=0) here actually bricks batchnorm, even with track_running_stats=True
        corrects = []
        corrects_ema = []
        losses = []
@@ -3,7 +3,7 @@ import os
 if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
 from tinygrad import Device, nn, Tensor, dtypes
 from train_gpt2 import GPT, GPTConfig
-from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name
+from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name, Context
 from tinygrad.engine.realize import get_kernel
 from tinygrad.schedule.memory import memory_planner
 from tinygrad.uop.ops import Ops
@@ -23,23 +23,23 @@ if __name__ == "__main__":
  #B, T = Variable("B", 1, 128).bind(4), 64 #Variable("T", 1, 1024).bind(64)
  B, T = 4, 64

-  Tensor.training = True
  optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=1e-4)
  warmup_count = getenv("WARMUP", 3)
-  for i in range(warmup_count):  # TODO: why does it take three and not two to stabilize
-    GlobalCounters.reset()
-    X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
-    Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
-    _, loss = model(X, Y)
-    optimizer.zero_grad()
-    if getenv("BACKWARD", 1):
-      loss.backward()
-      tensors = optimizer.schedule_step()
-    else:
-      tensors = []
-    sched = loss.schedule(*tensors)
-    print(f"calls {i}:", len(sched))
-    #run_schedule(sched[:])
+  with Context(TRAINING=1):
+    for i in range(warmup_count):  # TODO: why does it take three and not two to stabilize
+      GlobalCounters.reset()
+      X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
+      Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
+      _, loss = model(X, Y)
+      optimizer.zero_grad()
+      if getenv("BACKWARD", 1):
+        loss.backward()
+        tensors = optimizer.schedule_step()
+      else:
+        tensors = []
+      sched = loss.schedule(*tensors)
+      print(f"calls {i}:", len(sched))
+      #run_schedule(sched[:])
  sched = memory_planner(sched)
  ast_dedup = dedup([si.ast for si in sched if si.ast.op is Ops.SINK])
  srcs = {}
@@ -2,7 +2,7 @@ import math
 from typing import Union

 from tinygrad import Tensor, nn, dtypes
-from tinygrad.helpers import prod, argfix, Context
+from tinygrad.helpers import prod, argfix, Context, TRAINING
 from tinygrad.nn.state import get_parameters
 from extra.models.unet import UNetModel

@@ -85,7 +85,7 @@ class FrozenBatchNorm2dRetinaNet(nn.BatchNorm2d):

  def __call__(self, x:Tensor) -> Tensor:
    batch_mean, batch_var = super().calc_stats(x.cast(dtypes.float32))
-    if self.track_running_stats and Tensor.training:
+    if self.track_running_stats and TRAINING:
      self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
      self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * x.numel()/(x.numel()-x.shape[1]) * batch_var.detach().cast(self.running_var.dtype))
      self.num_batches_tracked += 1
@@ -498,11 +498,10 @@ def eval_stable_diffusion():

 if __name__ == "__main__":
  # inference only
-  Tensor.training = False
-
  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
-  for m in models:
-    nm = f"eval_{m}"
-    if nm in globals():
-      print(f"eval {m}")
-      globals()[nm]()
+  with Context(TRAINING=0):
+    for m in models:
+      nm = f"eval_{m}"
+      if nm in globals():
+        print(f"eval {m}")
+        globals()[nm]()
@@ -1,6 +1,6 @@
 # load each model here, quick benchmark
 from tinygrad import Tensor, GlobalCounters
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, Context
 import numpy as np

 def test_model(model, *inputs):
@@ -59,11 +59,10 @@ def spec_mrcnn():

 if __name__ == "__main__":
  # inference only for now
-  Tensor.training = False
-
-  for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
-    nm = f"spec_{m}"
-    if nm in globals():
-      print(f"testing {m}")
-      globals()[nm]()
+  with Context(TRAINING=0):
+    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
+      nm = f"spec_{m}"
+      if nm in globals():
+        print(f"testing {m}")
+        globals()[nm]()

@@ -157,6 +157,7 @@ def train_resnet():
  # input_std = Tensor([0.229, 0.224, 0.225], device=GPUS, dtype=dtypes.float32).reshape(1, -1, 1, 1)
  def normalize(x): return (x.permute([0, 3, 1, 2]) - input_mean).cast(dtypes.default_float)
  @TinyJit
+  @Context(TRAINING=1)
  def train_step(X, Y):
    optimizer_group.zero_grad()
    X = normalize(X)
@@ -170,6 +171,7 @@ def train_resnet():
    return loss.realize(), top_1.realize()

  @TinyJit
+  @Context(TRAINING=0)
  def eval_step(X, Y):
    X = normalize(X)
    out = model.forward(X)
@@ -192,7 +194,6 @@ def train_resnet():
    # ** train loop **
    if MLLOGGER and RUNMLPERF:
      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=e+1, metadata=dict(epoch_num=e+1))
-    Tensor.training = True
    BEAM.value = TRAIN_BEAM

    if INITMLPERF:
@@ -271,7 +272,6 @@ def train_resnet():
      eval_loss = 0.0
      eval_top_1 = 0
      eval_num_samples = 0
-      Tensor.training = False
      BEAM.value = EVAL_BEAM

      if INITMLPERF:
@@ -919,6 +919,7 @@ def train_rnnt():
  pass

@TinyJit
+@Context(TRAINING=0)
 def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
                   masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
  for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
@@ -1106,6 +1107,7 @@ def train_bert():
      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})

  @TinyJit
+  @Context(TRAINING=1)
  def train_step_bert(input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
                      masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
@@ -1133,7 +1135,6 @@ def train_bert():

  while train_data is not None and i < train_steps and not achieved:
    if getenv("TRAIN", 1):
-      Tensor.training = True
      BEAM.value = TRAIN_BEAM
      st = time.perf_counter()
      GlobalCounters.reset()
@@ -1186,7 +1187,6 @@ def train_bert():
      eval_lm_accs = []
      eval_clsf_accs = []
      eval_times = []
-      Tensor.training = False
      BEAM.value = EVAL_BEAM

      for j in tqdm(range(max_eval_steps), desc="Evaluating", total=max_eval_steps, disable=BENCHMARK):
@@ -38,7 +38,7 @@ def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):

 def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
           x_fp8:Tensor|None=None, x_new_amax:Tensor|None=None,
-           grad_amax_state:Tensor|None=None) -> tuple[Tensor,...]:
+           grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None) -> tuple[Tensor,...]:
  if not fp8:
    if ASM_GEMM:
      from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
@@ -47,12 +47,14 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
  assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
  if MXFP8:
    from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
-    x_q, x_e8, x_si = quantize_mxfp8(x.reshape(-1, x.shape[-1]))
+    if x_prequant_mx is not None: x_q, x_e8, x_si = x_prequant_mx       # fused producer already quantized (2d)
+    else: x_q, x_e8, x_si = quantize_mxfp8(x.reshape(-1, x.shape[-1]))
+    l_shape = x.shape[:-1] if x is not None else x_q.shape[:-1]
    if can_use_asm_gemm(x_q, w.T):
      out = asm_gemm(x_q, w.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(w_inv_scale), w_inv_scale),
-                     mx_w_stored=True).reshape(*x.shape[:-1], w.shape[0])
+                     mx_w_stored=True).reshape(*l_shape, w.shape[0])
    else:
-      x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*x.shape[:-1], x.shape[-1])
+      x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*l_shape, x_q.shape[-1])
      out = x_phys @ (w.cast(dtypes.bfloat16) * _mx_block_scale(w_inv_scale)).T
    return out, (amax_x.detach() if amax_x is not None else None), x_q
  if x_fp8 is None:
@@ -214,8 +216,15 @@ class FlatTransformer:
      x_w3, new_amax, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"], grad_amax_state=kwargs["grad_amax_xw3"])
      amaxs.append(new_amax)
      saves.extend([*s, x_w3])
-      out, new_amax, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
-                                 grad_amax_state=kwargs["grad_amax_xout"])
+      if FUSED_SILU_W13 and MXFP8:
+        from extra.llama_kernels.fused_silu_mul_quantize_mxfp8 import fused_silu_mul_quantize_mxfp8
+        aq, ae8, asi = fused_silu_mul_quantize_mxfp8(x_w1.reshape(-1, x_w1.shape[-1]), x_w3.reshape(-1, x_w3.shape[-1]))
+        out, new_amax, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
+                                   w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"])
+        out = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
+      else:
+        out, new_amax, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
+                                   grad_amax_state=kwargs["grad_amax_xout"])
      amaxs.append(new_amax)
      saves.extend([*s, out])
    else:
@@ -1,6 +1,6 @@
 #!/bin/bash
-export BENCHMARK=5
+export BENCHMARK=${BENCHMARK:-5}
 export EVAL_BS=0
 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
 SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
-python -m tinygrad.viz.cli -s "$SRC" -t --interval "train @ 2" "train @ 3"
+[ "$BENCHMARK" -le 3 ] || python -m tinygrad.viz.cli -s "$SRC" -t --interval "train @ 2" "train @ 3"
@@ -42,7 +42,7 @@ def compile(onnx_file):
  kernel_calls = [u for u in run_onnx_jit.captured.linear.toposort(gate=lambda x: x.op not in kernel_asts)
                  if u.op is Ops.CALL and u.src[0].op in kernel_asts]
  print(f"captured {len(kernel_calls)} kernels")
-  np.testing.assert_equal(test_val, ret, "JIT run failed")
+  if getenv("TEST", 1): np.testing.assert_equal(test_val, ret, "JIT run failed")
  print("jit run validated")

  # check gated read_image usage
@@ -50,7 +50,7 @@ def compile(onnx_file):
  read_image_count = 0
  gated_read_image_count = 0
  for call in kernel_calls:
-    _, _, _, source, _ = call.src[0].src
+    _, _, source, _ = call.src[0].src
    src = source.arg
    kernel_count += 1
    read_image_count += src.count("read_image")
@@ -38,7 +38,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
    prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
    info = prg.arg
-    functions[info.function_name] = prg.src[3].arg
+    functions[info.function_name] = prg.src[2].arg
    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + list(info.vars)
    statements.append((info.function_name, cargs, info.global_size, info.local_size))

@@ -462,7 +462,7 @@ def test_matmul():
    lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
                                                                                  estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
  linear = c.schedule_linear()

@@ -17,7 +17,7 @@ def make_matmul_kernel(name:str, src:str, local_size:int):
    wg_y = UOp.special(N//128, "gidx1")
    sink = UOp.sink(a.base, b.base, c.base, threads, wg_x, wg_y, arg=KernelInfo(name, estimates=Estimates(ops=2*N**3, mem=3*N*N*4)))
    lib = Device[Device.DEFAULT].compiler.compile_cached(src)
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
                                 UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
  return fxn

@@ -2622,8 +2622,7 @@ def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
  lds = UOp.placeholder((133_120,), dtypes.uint8, 0, AddrSpace.LOCAL)
  sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
                  arg=KernelInfo(name=f"gemm_{batch}_{M}_{N}_{K}", estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname),
-                                UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

 # ** FP8 GEMM custom kernel

@@ -2646,7 +2645,7 @@ def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str, scale_mode:int
  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
                                 f"-DSCALE_MODE={scale_mode}"]).compile_cached(src)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                               UOp(Ops.BINARY, arg=lib)))

 # ** MXFP8 GEMM custom kernel
@@ -2669,7 +2668,7 @@ def custom_hk_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:U
  src = (kittens_path/"gemm_mxfp8.cpp").read_text()
  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                               UOp(Ops.BINARY, arg=lib)))

 def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
@@ -2759,7 +2758,7 @@ def custom_hk_bf16_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str) -> UOp:
  src = (kittens_path/"gemm_bf16.cpp").read_text()
  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                                UOp(Ops.BINARY, arg=lib)))

@functools.cache
@@ -2777,7 +2776,7 @@ def custom_hk_bf16_atb_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
  src = (kittens_path/"gemm_bf16_atb.cpp").read_text()
  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                                UOp(Ops.BINARY, arg=lib)))

 def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
@@ -2851,11 +2850,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=
    grad_a = asm_gemm(g_fp8, b_t, x_scale=s_x_t, w_scale=s_w_t, g_scale=g_scale) if has_w else asm_gemm(g_fp8, b_t, x_scale=s_x_t, w_scale=g_scale)
    # wgrad: no w_scale
    g_fp8_2d = g_fp8.reshape(-1, g_fp8.shape[-1])
-    if getenv("FAST_FP8_TRANSPOSE", 0) and g_fp8_2d.shape[0] % 64 == 0 and g_fp8_2d.shape[1] % 64 == 0:
-      from extra.llama_kernels.fp8_transpose import fast_fp8_transpose
-      g_fp8_T = fast_fp8_transpose(g_fp8_2d)
-    else:
-      g_fp8_T = g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1)
+    g_fp8_T = g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1)
    grad_b = asm_gemm(g_fp8_T, a_t.reshape(-1, a_t.shape[-1]), x_scale=s_x_t, w_scale=g_scale)
    # wgrad: rescale if not scalar
    if w_post_t is not None:
@@ -2899,7 +2894,7 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F

  g = Tensor(gradient, device=aq.device)[:aq.shape[0]].reshape(aq.shape[0]*aq.shape[1], bq.shape[0]).cast(dtypes.bfloat16)
  grad_a = asm_gemm(g, b_phys, mx=True)
-  grad_b = asm_gemm(g.T, a_phys, mx=True)
+  grad_b = asm_gemm(g.T, a_phys, mx=True, a_pretranspose=g)

  grad_a = (grad_a * _mx_block_scale(ae8)).reshape(aq.shape)
  if not w_stored: grad_b = grad_b * _mx_block_scale(be8)
@@ -2909,7 +2904,8 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F
 # ** main gemm function

 def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None,
-             w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_scale:Tensor|None=None) -> Tensor:
+             w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_scale:Tensor|None=None,
+             a_pretranspose:Tensor|None=None) -> Tensor:
  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
  counters["used"] += 1
  unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
@@ -2946,6 +2942,11 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
      if mx_scales is not None:
        a_si, a_e8, b_si, b_e8 = mx_scales
        a_q, b_q = a.reshape(-1, a.shape[-1]), b.T
+      elif (a_pretranspose is not None and getenv("FUSED_GRAD_QUANTIZE", 0) and a_pretranspose.dtype == dtypes.bfloat16
+            and a_pretranspose.shape[0] % 32 == 0 and a_pretranspose.shape[1] % 256 == 0):
+        from extra.llama_kernels.transpose_quantize_mxfp8 import transpose_quantize_mxfp8
+        a_q, a_e8, a_si = transpose_quantize_mxfp8(a_pretranspose)
+        b_q, b_e8, b_si = quantize_mxfp8(b.T)
      else:
        a_q, a_e8, a_si = quantize_mxfp8(a.reshape(-1, a.shape[-1]))
        b_q, b_e8, b_si = quantize_mxfp8(b.T)
@@ -223,7 +223,7 @@ def test_matmul():
    lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs,
                    arg=KernelInfo(name=colored("kernel","cyan"), estimates=Estimates(ops=N*N*N*2, mem=N*N*2*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
  linear = c.schedule_linear()
@@ -93,7 +93,7 @@ if __name__ == "__main__":
  info = ProgramInfo(name="matmul_kernel",
                     global_size=(M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1), local_size=(32*compiled.metadata.num_warps, 1, 1))
  sink = UOp.sink(arg=KernelInfo(name="matmul_kernel"))
-  prg_uop = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info),
+  prg_uop = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info),
                       Device.default.renderer)
  rt = get_runtime(Device.DEFAULT, prg_uop)
  all_bufs = [x.ensure_allocated() for x in bufs]
@@ -27,7 +27,7 @@ class HCQ2Compiled(Compiled):
      (UPat(Ops.BUFFER, tag="timeline_value"), lambda ctx: ctx.timeline_value()),
      (UPat(Ops.BUFFER, tag="sentinel_signal"), lambda ctx: ctx.timeline_signal("sentinel", (1 << 64) - 1)),
      (UPat(Ops.BUFFER, name="b"), lambda ctx, b:
-        Buffer(ctx.device, b.arg, b.dtype, options=BufferSpec(host=False, uncached=True, cpu_access=True, nolru=True))), # TODO: remove nolru
+        Buffer(ctx.device, b.max_numel(), b.dtype, options=BufferSpec(host=False, uncached=True, cpu_access=True, nolru=True))), # TODO: remove nolru
    ])

    super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch)
@@ -137,20 +137,23 @@ def unwrap_after(uop):
  return uop

 def make_getaddr(u, device=None):
-  if unwrap_after(u).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT): return u
+  if unwrap_after(u).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM): return u
  return UOp(Ops.GETADDR, dtypes.uint64, src=(u, UOp(Ops.DEVICE, arg=device or to_tuple(u.device)[0])))

 def make_ins(op, *srcs):
  return UOp(Ops.INS, dtypes.void, tuple(UOp.const(dtypes.uint32, s) if isinstance(s, int) else s.cast(dtypes.uint32) for s in srcs), op)

+def make_patch(buf:UOp, off:sint, val:UOp, dtype=None) -> UOp:
+  dt = dtype or val.dtype
+  return UOp(Ops.SHRINK, buf.dtype.base, (buf, UOp.const(dtypes.int, off), UOp.const(dtypes.int, dt.itemsize))).bitcast(dt).store(val.cast(dt))
+
 def make_cmdbuf(lin, devs, tag):
  blob, patches = b'', []
  for s in (s for ins in lin.src for s in ins.src):
    if s.op is not Ops.CONST: patches.append((len(blob), s))
    blob += struct.pack(f'<{s.dtype.fmt}', s.arg if s.op is Ops.CONST else 0x0)
  buf = UOp.new_buffer(devs, len(blob), dtypes.uint8).rtag(tag)
-  stores = [buf.index(UOp.const(dtypes.int, off), dtype=buf.dtype.ptr()).cast(s.dtype.ptr()).store(s) for off, s in patches]
-  return buf.after(buf.store(UOp(Ops.BINARY, dtypes.void, src=(), arg=blob)), *stores)
+  return buf.after(buf.store(UOp(Ops.BINARY, dtypes.void, src=(), arg=blob)), *[make_patch(buf, off, s) for off, s in patches])

 def make_mstack(uops): return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, tuple(uops))

@@ -180,19 +183,19 @@ class HCQInfo:
  def from_call(call:UOp) -> HCQInfo: return HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), get_call_outs_ins(call)[0])

 # *****************
-# 1.1. prep runtimes: staging copies
+# 1.1. prep: staging copies

 def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS)

 def stage_copy(dst:UOp, src:UOp) -> UOp|None:
  if not (_need_staging(src, dst) or _need_staging(dst, src)): return None

-  stage = UOp.new_buffer("CPU", src.buffer.nbytes, dtypes.uint8)
+  stage = UOp.new_buffer("CPU", src.nbytes(), dtypes.uint8)
  return UOp(Ops.LINEAR, dtypes.void, (src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
 pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])

 # *****************
-# 1.2. prep runtimes: programs/kernargs
+# 2.1. hcq lowering: programs/kernargs

@functools.cache
 def get_pm_prep_program(name:str) -> PatternMatcher|None:
@@ -211,19 +214,15 @@ def prep_program(call:UOp, prg:UOp) -> UOp|None:
  return prg.replace(src=(buf.after(buf.store(blob)),), arg=(data, prg.arg)).call(*call.src[1:], aux=HCQInfo.from_call(call))

 def prep_kernargs(call:UOp, prg:UOp) -> UOp:
-  data, info = prg.arg
-  patches = [(i*dtypes.uint64.itemsize, UOp(Ops.GETADDR, dtypes.uint64, src=(call.src[1+gi], UOp(Ops.DEVICE, arg=call.src[1+gi].device))),
-              dtypes.uint64) for i,gi in enumerate(info.globals)] \
-          + [(len(info.globals)*dtypes.uint64.itemsize + i*dtypes.uint32.itemsize, v, dtypes.uint32) for i,v in enumerate(info.vars)]
-
-  buf = UOp.new_buffer(call.src[1].device, data.kernargs_alloc_size, dtypes.uint8).rtag("kernargs")
-  kernargs = buf.after(*tuple(buf.index(UOp.const(dtypes.int, o), dtype=buf.dtype.ptr()).cast(dt.ptr()).store(val.cast(dt)) for o, val, dt in patches))
-
-  return call.replace(src=(prg.replace(src=prg.src + (kernargs,), arg=(data, info)),) + call.src[1:])
+  (data, info), dev_uop = prg.arg, UOp(Ops.DEVICE, arg=call.src[1].device)
+  buf = UOp.new_buffer(dev_uop.arg, data.kernargs_alloc_size, dtypes.uint8).rtag("kernargs")
+  patches = [make_patch(buf, i*8, make_getaddr(call.src[1+gi], dev_uop.arg)) for i,gi in enumerate(info.globals)] \
+          + [make_patch(buf, len(info.globals)*8 + i*4, v, dtypes.uint32) for i,v in enumerate(info.vars)]
+  return call.replace(src=(prg.replace(src=prg.src + (buf.after(*patches),), arg=(data, info)),) + call.src[1:])

 pm_prep_runtime = PatternMatcher([
  # bind generic PROGRAM device to the call's actual dev(s), then run device-specific lowering
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"),),
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"),),
    name="call", allow_any_len=True), prep_program),

  # lower kernargs (PROGRAM.src[0] is now AFTER(BUFFER, COPY) — the lowered program image)
@@ -231,7 +230,7 @@ pm_prep_runtime = PatternMatcher([
 ])

 # *****************
-# 2. lowering to hcq ir
+# 2.2. hcq lowering: ops to ir

 def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
  devs:tuple[str, ...] = to_tuple(devs)
@@ -471,15 +470,16 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None:
  off_per_buf:dict[UOp, int] = {}
  size_per_tag:dict[str, int] = {}
  for b in bufs:
-    if b.tag in maxtags: size_per_tag[b.tag] = max(size_per_tag.get(b.tag, 0), b.arg)
+    bsz = b.max_numel()
+    if b.tag in maxtags: size_per_tag[b.tag] = max(size_per_tag.get(b.tag, 0), bsz)
    elif b.tag in sumtags:
      off_per_buf[b] = round_up(size_per_tag.get(b.tag, 0), {"program": 0x1000}.get(b.tag, 128))
-      size_per_tag[b.tag] = off_per_buf[b] + b.arg
+      size_per_tag[b.tag] = off_per_buf[b] + bsz

  count_per_tag = collections.Counter(b.tag for b in bufs)
  ref_bufs = {b.tag:b for b in bufs if count_per_tag[b.tag] > 1}
-  bases = {tag:UOp.new_buffer(b.src[1].arg, size_per_tag[tag], b.dtype).rtag(tag) for tag,b in ref_bufs.items()}
-  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.weakint, off_per_buf.get(b, 0))), b.arg) for b in bufs if b.tag in bases}
+  bases = {tag:UOp.new_buffer(b.device, size_per_tag[tag], b.dtype).rtag(tag) for tag,b in ref_bufs.items()}
+  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.weakint, off_per_buf.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
  return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
 pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), pack_hcq_placeholders)])

@@ -496,7 +496,7 @@ pm_hold_call_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"),

 def bufferize_buf(buf:UOp) -> UOp|None:
  if buf.tag is None: return None
-  uops = tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=dv), "CPU") for dev in to_tuple(buf.src[1].arg))
+  uops = tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=dv), "CPU") for dev in to_tuple(buf.device))
  return make_mstack(uops)
 pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, name="buf"), bufferize_buf)])

@@ -532,9 +532,9 @@ pm_resolve_patches = PatternMatcher([
  (UPat(GroupOp.ALU, src=[UPat(Ops.STACK, name="s"), UPat(Ops.CONST)], name="op"), push_stack),
  (UPat(Ops.CAST, src=(UPat(Ops.STACK, name="s"),), name="op"), push_stack),

-  # index on slice is index
-  (UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"), UPat()), name="idx", allow_any_len=True),
-    lambda idx, bv: idx.replace(src=(bv.src[0], idx.src[1] + bv.src[1].cast(idx.src[1].dtype), *idx.src[2:]))),
+  # shrink on slice is shrink on base at offset
+  (UPat(Ops.SHRINK, src=(UPat(Ops.SLICE, name="bv"), UPat(), UPat()), name="shr"),
+    lambda shr, bv: shr.replace(src=(bv.src[0], shr.src[1] + bv.src[1].cast(shr.src[1].dtype), shr.src[2]))),

  # getaddr
  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"), UPat(Ops.DEVICE, name="dev"))), resolve_getaddr_slice), # getaddr(slice(x)) -> offset+getaddr(x)
@@ -542,8 +542,8 @@ pm_resolve_patches = PatternMatcher([

  # folders
  (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store),
-  (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat.cvar("off")).or_casted().store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))),
-    fold_const_store),
+  (UPat(Ops.SHRINK, src=(UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf"), UPat.cvar("off"), UPat(Ops.CONST))).bitcast()
+    .store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))), fold_const_store),
 ]) + symbolic_simple

 # *****************
@@ -570,8 +570,8 @@ def callify_hcq(call:UOp) -> UOp:
  return UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(prg,), arg="hcq").call(*call.src[1:], aux=call.arg.aux)
 pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), callify_hcq)])

-@track_rewrites(lambda _,ret: f"HCQ Schedule {pluralize('Kernel', len(ret.src))}")
-def hcq_schedule(linear:UOp) -> UOp:
+@track_rewrites(lambda _,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
+def hcq_compile(linear:UOp) -> UOp:
  linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
  linear = graph_rewrite(linear, pm_prep_runtime, name="prepare runtime")

@@ -587,12 +587,11 @@ def hcq_schedule(linear:UOp) -> UOp:
  linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True)
  linear = graph_rewrite(linear, pm_lift_patches_to_cmdbuf, name="lift patches to cmdbuf", enter_calls=True)
  linear = graph_rewrite(linear, pm_pack_placeholders, walk=True, name="pack placeholders")
-  linear = graph_rewrite(linear, pm_hold_call_buffers, walk=True, name="hold call buffers")
+  return graph_rewrite(linear, pm_hold_call_buffers, walk=True, name="hold call buffers")

-  # realize starts from here
+@track_rewrites(lambda _,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
+def hcq_link(linear:UOp) -> UOp:
  linear = graph_rewrite(linear, pm_bufferize, bottom_up=True, walk=True, name="bufferize placeholders", enter_calls=True)
  linear = graph_rewrite(linear, pm_resolve_patches, bottom_up=False, name="simplify patches", enter_calls=True)
  linear = graph_rewrite(linear, pm_parametrize_host_buffers, walk=True, name="parametrize host buffers")
-  linear = graph_rewrite(linear, pm_callify_hcq, name="callify hcq")
-
-  return linear
+  return graph_rewrite(linear, pm_callify_hcq, name="callify hcq")
@@ -142,7 +142,7 @@ pm_pm4_opsel = PatternMatcher([
 ])

 def pm4_submit(cmdbuf, devs):
-  size, zero = UOp.const(dtypes.uint32, cmdbuf.src[0].arg // dtypes.uint32.itemsize), UOp.const(dtypes.int, 0)
+  size, zero = UOp.const(dtypes.uint32, cmdbuf.nbytes() // dtypes.uint32.itemsize), UOp.const(dtypes.int, 0)

  # the compute queue's ring and its host-side ring/write/put pointers (placeholders, resolved in pm_bufferize)
  for d in devs: q = Device[d].compute_queue
@@ -205,7 +205,7 @@ pm_sdma_opsel = PatternMatcher([

 def sdma_submit(cmdbuf, devs):
  # the cmdbuf to submit + the patch writes that fill it
-  size_dw, zero = cmdbuf.src[0].arg // dtypes.uint32.itemsize, UOp.const(dtypes.int, 0)
+  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(dtypes.int, 0)

  # the sdma queue's ring and its host-side ring/write/put pointers
  for d in devs: q = Device[d].sdma_queue(0)
@@ -247,8 +247,8 @@ class AMDProgramData:
 _amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}

 def amd_build_program(prg:UOp) -> UOp:
-  dev = Device[prg.src[1].arg] # TODO: rm this
-  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[4].arg, dev.device))) is None:
+  dev = Device[to_tuple(prg.device)[0]] # TODO: rm this
+  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, dev.device))) is None:
    image, sections, relocs = elf_loader(lib)
    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
    for off, sym, typ, addent in relocs:
@@ -271,7 +271,7 @@ def amd_build_program(prg:UOp) -> UOp:
  return cached

 pm_prep_program = PatternMatcher([
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE, arg="AMD"), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
+  (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
 ])

 class AMDAllocator(HCQAllocator['AMDDevice']):
@@ -579,7 +579,7 @@ class AMDDevice(HCQ2Compiled):

    # Scratch setup
    self.max_private_segment_size = 0
-    self.pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.arg))]) + self.pm_bufferize
+    self.pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.max_numel()))]) + self.pm_bufferize

    self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
    if self.pmc_enabled:
@@ -21,7 +21,7 @@ def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp,
                  xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups,
                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem)))
  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))

@functools.cache
@@ -34,7 +34,7 @@ def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:
  sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups,
                  arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_fwd_w13.cpp", n_elems, hidden)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))

 def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
@@ -1,41 +0,0 @@
-from __future__ import annotations
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from extra.llama_kernels import THREADS_PER_WG, alloc_like, dname_of, compile_hip
-
-TILE = 64
-
-@functools.cache
-def _custom_fp8_transpose(out:UOp, inp:UOp, dname:str) -> UOp:
-  M, N = inp.shape
-  num_wg = (M // TILE) * (N // TILE)
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0")
-  mem = M * N * 2  # one byte read + one byte write per element
-  sink = UOp.sink(out.base, inp.base, threads, workgroups,
-                  arg=KernelInfo(f"fp8_transpose_{M}_{N}",
-                                 estimates=Estimates(ops=M*N, mem=mem)))
-  src = (pathlib.Path(__file__).parent/"fp8_transpose.cpp").read_text()
-  defines = [f"-DM_DIM={M}", f"-DN_DIM={N}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
-
-def fast_fp8_transpose(t:Tensor) -> Tensor:
-  assert t.ndim == 2, f"fast_fp8_transpose needs 2D input, got shape {t.shape}"
-  assert t.dtype in dtypes.fp8s, f"fast_fp8_transpose needs fp8 dtype, got {t.dtype}"
-  M, N = t.shape
-  assert M % TILE == 0 and N % TILE == 0, f"M={M}, N={N} must be multiples of {TILE}"
-
-  device = t.device
-  axis = t.uop.axis if isinstance(device, tuple) else None
-  out_axis = None
-  if axis == 0: out_axis = 1
-  elif axis == 1: out_axis = 0
-  elif axis is not None:
-    raise ValueError(f"fast_fp8_transpose: unsupported axis {axis}")
-
-  out = alloc_like((N, M), t.dtype, device, out_axis)
-  fxn = functools.partial(_custom_fp8_transpose, dname=dname_of(device))
-  out, _ = Tensor.custom_kernel(out, t, fxn=fxn)
-  return out
@@ -1,74 +0,0 @@
-#include <hip/hip_runtime.h>
-
-// LDS-staged 64x64 fp8 transpose.
-//   in : (M_DIM, N_DIM) fp8 contiguous
-//   out: (N_DIM, M_DIM) fp8 contiguous, out[c][r] = in[r][c]
-//
-// One WG processes one 64x64 output tile.  Each thread reads one uint4 (16 fp8) coalesced
-// from input rows, stages into LDS, then writes one uint4 coalesced to the output (whose
-// 16 fp8 come from 16 different input rows via in-LDS gather).
-//
-// LDS layout: lds[64][LDS_STRIDE] with LDS_STRIDE=65 (1 byte pad) to mitigate bank conflicts
-// during the column-direction read of the write phase.
-
-#ifndef M_DIM
-#define M_DIM 16384
-#endif
-#ifndef N_DIM
-#define N_DIM 28672
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-
-constexpr int TILE = 64;
-constexpr int VEC = 16;                    // fp8 per uint4 (128-bit) load/store
-constexpr int LDS_PAD = 1;
-constexpr int LDS_STRIDE = TILE + LDS_PAD; // 65 fp8 per row
-
-static_assert(THREADS_PER_WG * VEC == TILE * TILE, "256 threads * 16 fp8 = 64*64");
-static_assert(M_DIM % TILE == 0, "M_DIM must be a multiple of 64");
-static_assert(N_DIM % TILE == 0, "N_DIM must be a multiple of 64");
-
-constexpr int N_TILES_N = N_DIM / TILE;
-
-struct alignas(16) fp8x16 { uint8_t v[16]; };
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-fp8_transpose(uint8_t* __restrict__ out,         // (N_DIM, M_DIM)
-              const uint8_t* __restrict__ in)    // (M_DIM, N_DIM)
-{
-  __shared__ uint8_t lds[TILE * LDS_STRIDE];
-
-  const int tid     = threadIdx.x;
-  const int wg_id   = blockIdx.x;
-  const int tile_r  = wg_id / N_TILES_N;       // tile index along M dim of input
-  const int tile_c  = wg_id % N_TILES_N;       // tile index along N dim of input
-
-  const int a = tid / (TILE / VEC);            // 0..63 (row within tile during read; col within tile during write)
-  const int b = tid % (TILE / VEC);            // 0..3
-  const int b16 = b * VEC;                     // 0,16,32,48
-
-  // ---- Read phase: input rows -> LDS rows
-  {
-    const long long src = (long long)(tile_r * TILE + a) * (long long)N_DIM
-                        + (long long)(tile_c * TILE + b16);
-    fp8x16 v = *reinterpret_cast<const fp8x16*>(&in[src]);
-    *reinterpret_cast<fp8x16*>(&lds[a * LDS_STRIDE + b16]) = v;
-  }
-  __syncthreads();
-
-  // ---- Write phase: LDS columns (gathered) -> output rows
-  // out[(tile_c*TILE + a)][(tile_r*TILE + b16 + i)] = in[(tile_r*TILE + b16 + i)][(tile_c*TILE + a)]
-  //                                                  = lds[b16 + i][a]
-  {
-    fp8x16 v;
-    #pragma unroll
-    for (int i = 0; i < VEC; ++i) {
-      v.v[i] = lds[(b16 + i) * LDS_STRIDE + a];
-    }
-    const long long dst = (long long)(tile_c * TILE + a) * (long long)M_DIM
-                        + (long long)(tile_r * TILE + b16);
-    *reinterpret_cast<fp8x16*>(&out[dst]) = v;
-  }
-}
@@ -22,7 +22,7 @@ def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
             f"-DEPS_LITERAL={eps_val}f"]
  src = _src()
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))

@functools.cache
@@ -39,7 +39,7 @@ def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax
  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
             f"-DEPS_LITERAL={eps_val}f", f"-DHAS_RESIDUAL=1"]
  src = _src()
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))

@functools.cache
@@ -55,7 +55,7 @@ def _custom_bwd(grad_x:UOp, grad_weight_partial:UOp,
                                 estimates=Estimates(ops=8*n_elems, mem=mem)))
  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
  src = _src_bwd()
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))

 def _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel:UOp):
@@ -0,0 +1,104 @@
+import functools
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from extra.llama_kernels import FP8_MAX, THREADS_PER_WG, alloc_like
+
+BLK = 32
+PACK = 4
+LOG2E = 1.4426950408889634
+
+@functools.cache
+def _custom_silu_mul_quantize_mxfp8(fp8_out:UOp, e8_out:UOp, si_out:UOp, x_w1:UOp, x_w3:UOp) -> UOp:
+  rows, K = x_w1.shape
+  scale_K = K // BLK
+  n_elems = rows * K
+  n_super = n_elems // (BLK * PACK)
+  sk4 = scale_K // PACK
+  assert n_super % THREADS_PER_WG == 0, f"{n_super=} must divide over {THREADS_PER_WG=}"
+  nwg = n_super // THREADS_PER_WG
+
+  x_w1, x_w3 = x_w1.reshape(n_elems), x_w3.reshape(n_elems)
+  fp8_out = fp8_out.reshape(n_elems)
+  e8_out = e8_out.reshape(rows * scale_K)
+  si_out = si_out.reshape(sk4 * rows)
+
+  wg = UOp.range(nwg, 0, AxisType.GLOBAL)
+  tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
+  sb = UOp.range(PACK, 2, AxisType.UNROLL)
+  lane = UOp.range(BLK, 3, AxisType.UNROLL)
+
+  super_idx = wg * THREADS_PER_WG + tid
+  idx = super_idx * (BLK * PACK) + sb * BLK + lane
+
+  w1 = x_w1[idx].cast(dtypes.float)
+  w3 = x_w3[idx].cast(dtypes.float)
+  sig = (1.0 + (w1 * -LOG2E).exp2()).reciprocal()
+  act = w1 * sig * w3
+  abs_a = (act < 0.0).where(-act, act)
+  blk_max = abs_a.reduce(lane, arg=Ops.MAX)
+  e8f = (blk_max.maximum(1e-38).log2().floor() + 127.0).maximum(0.0).minimum(254.0)
+  qscale = (127.0 - e8f).exp2()
+  scaled = (act * qscale).maximum(-FP8_MAX).minimum(FP8_MAX)
+  e8u8 = e8f.cast(dtypes.uint8)
+
+  fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype.base)).end(lane)
+  e8_store = e8_out.after(fp8_store)[super_idx * PACK + sb].store(e8u8)
+  packed = (e8u8.cast(dtypes.uint32) << (sb.cast(dtypes.uint32) * 8)).reduce(sb, arg=Ops.ADD)
+  row, col4 = super_idx // sk4, super_idx % sk4
+  si_store = si_out.after(e8_store.end(sb))[col4 * rows + row].store(packed)
+  return si_store.end(tid, wg).sink(arg=KernelInfo(f"silu_mul_quantize_mxfp8_{n_elems}", opts_to_apply=()))
+
+@functools.cache
+def _custom_silu_mul_bwd_mxfp8(gx1_out:UOp, gx3_out:UOp, x_w1:UOp, x_w3:UOp, grad_aq:UOp, e8:UOp) -> UOp:
+  rows, K = x_w1.shape
+  scale_K = K // BLK
+  n_elems = rows * K
+  VEC = 8
+  assert n_elems % (THREADS_PER_WG * VEC) == 0, f"{n_elems=} must divide {THREADS_PER_WG*VEC=}"
+  nwg = n_elems // (THREADS_PER_WG * VEC)
+  x_w1, x_w3, grad_aq = x_w1.reshape(n_elems), x_w3.reshape(n_elems), grad_aq.reshape(n_elems)
+  gx1_out, gx3_out, e8 = gx1_out.reshape(n_elems), gx3_out.reshape(n_elems), e8.reshape(rows * scale_K)
+
+  wg = UOp.range(nwg, 0, AxisType.GLOBAL)
+  tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
+  lane = UOp.range(VEC, 2, AxisType.UNROLL)
+  idx = (wg * THREADS_PER_WG + tid) * VEC + lane
+
+  e8v = e8[idx // BLK].cast(dtypes.float)
+  qscale = (127.0 - e8v).exp2()
+  ga = grad_aq[idx].cast(dtypes.float) * qscale
+  w1 = x_w1[idx].cast(dtypes.float)
+  w3 = x_w3[idx].cast(dtypes.float)
+  sig = (1.0 + (w1 * -LOG2E).exp2()).reciprocal()
+  s = w1 * sig
+  sprime = sig * (1.0 + w1 * (1.0 - sig))
+  gx1 = gx1_out[idx].store((ga * sprime * w3).cast(gx1_out.dtype.base))
+  gx3 = gx3_out.after(gx1)[idx].store((ga * s).cast(gx3_out.dtype.base))
+  return gx3.end(lane, tid, wg).sink(arg=KernelInfo(f"silu_mul_bwd_mxfp8_{n_elems}", opts_to_apply=()))
+
+def _silu_mul_quantize_mxfp8_bwd(gradient:UOp, kernel:UOp):
+  _, e8_out, _, x_w1, x_w3 = kernel.src[1:]
+  device = x_w1.device
+  rows, K = x_w1.shape
+  axis = x_w1.axis if isinstance(device, tuple) else None
+  gx1 = alloc_like((rows, K), dtypes.bfloat16, device, axis)
+  gx3 = alloc_like((rows, K), dtypes.bfloat16, device, axis)
+  gx1, gx3, *_ = Tensor.custom_kernel(gx1, gx3, Tensor(x_w1, device=device), Tensor(x_w3, device=device),
+                                      Tensor(gradient, device=device).cast(dtypes.bfloat16), Tensor(e8_out.after(kernel), device=device),
+                                      fxn=_custom_silu_mul_bwd_mxfp8)
+  return (None, None, None, gx1.uop, gx3.uop)
+
+def fused_silu_mul_quantize_mxfp8(x_w1:Tensor, x_w3:Tensor) -> tuple[Tensor, Tensor, Tensor]:
+  assert x_w1.shape == x_w3.shape, f"{x_w1.shape} != {x_w3.shape}"
+  assert x_w1.dtype == dtypes.bfloat16 and x_w3.dtype == dtypes.bfloat16
+  assert x_w1.ndim == 2, f"expected 2d, got {x_w1.shape}"
+  from extra.gemm.cdna_asm_gemm import FP8_DTYPE
+  rows, K = x_w1.shape
+  scale_K = K // BLK
+  axis = x_w1.uop.axis if isinstance(x_w1.device, tuple) else None
+  fp8_out = alloc_like((rows, K), FP8_DTYPE, x_w1.device, axis)
+  e8_out = alloc_like((rows, scale_K), dtypes.uint8, x_w1.device, axis)
+  si_out = alloc_like((scale_K // PACK, rows), dtypes.uint32, x_w1.device, None if axis is None else (1 if axis == 0 else 0))
+  fp8_out, e8_out, si_out, *_ = Tensor.custom_kernel(fp8_out, e8_out, si_out, x_w1, x_w3,
+                                                     fxn=_custom_silu_mul_quantize_mxfp8, grad_fxn=_silu_mul_quantize_mxfp8_bwd)
+  return fp8_out, e8_out, si_out
@@ -42,8 +42,8 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_st
  step = THREADS_PER_WG // 2
  while step:
    active = tid < step
-    other = lds[tid + step].load(UOp.const(dtypes.float, 0.0), active)
-    lds = lds.after(lds[tid].store(lds[tid].maximum(other), gate=active).barrier())
+    other = lds[(tid + step).valid(active)].load()
+    lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)).barrier())
    step //= 2

  amax_store = amax_partial[tid.eq(0).where(wg, UOp.invalid())].store(lds[0])
@@ -0,0 +1,37 @@
+from __future__ import annotations
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import THREADS_PER_WG, alloc_like, dname_of, compile_hip
+
+TILE_N = THREADS_PER_WG   # 256
+BLK = 32
+
+@functools.cache
+def _custom_transpose_quantize_mxfp8(q:UOp, e8:UOp, g:UOp, dname:str) -> UOp:
+  M, N = g.shape
+  num_wg = (M // BLK) * (N // TILE_N)
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0")
+  mem = M * N * 2 + M * N + (M // BLK) * N   # read bf16, write fp8 + e8
+  sink = UOp.sink(q.base, e8.base, g.base, threads, workgroups,
+                  arg=KernelInfo(f"transpose_quantize_mxfp8_{M}_{N}", estimates=Estimates(ops=M*N, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"transpose_quantize_mxfp8.cpp").read_text()
+  defines = [f"-DM_DIM={M}", f"-DN_DIM={N}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+def transpose_quantize_mxfp8(g:Tensor) -> tuple[Tensor, Tensor, Tensor]:
+  # fused g.T quantize: returns (q, e8, si) == quantize_mxfp8(g.T) — q (N,M) fp8, e8 (N, M/32), si packed (M/128, N)
+  assert g.ndim == 2 and g.dtype == dtypes.bfloat16, f"{g.shape} {g.dtype}"
+  from extra.gemm.cdna_asm_gemm import FP8_DTYPE, mx_pack
+  M, N = g.shape
+  assert M % BLK == 0 and N % TILE_N == 0, f"M={M} must%{BLK}, N={N} must%{TILE_N}"
+  device = g.device
+  axis = g.uop.axis if isinstance(device, tuple) else None
+  out_axis = None if axis is None else (1 if axis == 0 else 0)
+  q = alloc_like((N, M), FP8_DTYPE, device, out_axis)
+  e8 = alloc_like((N, M // BLK), dtypes.uint8, device, out_axis)
+  fxn = functools.partial(_custom_transpose_quantize_mxfp8, dname=dname_of(device))
+  q, e8, *_ = Tensor.custom_kernel(q, e8, g, fxn=fxn)
+  return q, e8, mx_pack(e8)
@@ -0,0 +1,62 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_fp8.h>
+#include <hip/hip_bf16.h>
+
+#ifndef M_DIM
+#define M_DIM 8192
+#endif
+#ifndef N_DIM
+#define N_DIM 14336
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int BLK = 32;
+constexpr int TILE_M = BLK;                 // one mxfp8 block along M per tile
+constexpr int TILE_N = THREADS_PER_WG;      // 256, one output column per thread
+constexpr int LDS_STRIDE = TILE_N + 1;      // +1 pad: stride 257 ≡ 1 (mod 32) -> conflict-free column reads
+constexpr int N_TILES_N = N_DIM / TILE_N;
+constexpr float FP8_MAX = 448.0f;
+
+static_assert(M_DIM % TILE_M == 0, "M_DIM must be a multiple of 32");
+static_assert(N_DIM % TILE_N == 0, "N_DIM must be a multiple of 256");
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+transpose_quantize_mxfp8(__hip_fp8_storage_t* __restrict__ q,        // (N_DIM, M_DIM)
+                         uint8_t* __restrict__ e8_out,               // (N_DIM, M_DIM/32)
+                         const __hip_bfloat16* __restrict__ g)       // (M_DIM, N_DIM)
+{
+  __shared__ __hip_bfloat16 lds[TILE_M * LDS_STRIDE];
+  const int tid    = threadIdx.x;
+  const int tile_m = blockIdx.x / N_TILES_N;     // which 32-block along M
+  const int tile_n = blockIdx.x % N_TILES_N;
+
+  #pragma unroll
+  for (int mm = 0; mm < TILE_M; mm++)
+    lds[mm * LDS_STRIDE + tid] = g[(long long)(tile_m * TILE_M + mm) * N_DIM + (tile_n * TILE_N + tid)];
+  __syncthreads();
+
+  float vals[TILE_M];
+  float amax = 0.0f;
+  #pragma unroll
+  for (int mm = 0; mm < TILE_M; mm++) {
+    float v = (float)lds[mm * LDS_STRIDE + tid];
+    vals[mm] = v;
+    amax = fmaxf(amax, fabsf(v));
+  }
+  int e8 = (int)floorf(log2f(fmaxf(amax, 1e-38f))) + 127;
+  e8 = max(0, min(254, e8));
+  float qscale = exp2f((float)(127 - e8));
+
+  const long long n = tile_n * TILE_N + tid;
+  __hip_fp8_storage_t out[TILE_M];
+  #pragma unroll
+  for (int mm = 0; mm < TILE_M; mm++)
+    out[mm] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, vals[mm] * qscale)), __HIP_SATFINITE, __HIP_E4M3);
+  // 32 contiguous fp8 along M -> two 16-byte vector stores
+  long long obase = n * M_DIM + (long long)(tile_m * TILE_M);
+  *reinterpret_cast<uint4*>(&q[obase])      = *reinterpret_cast<uint4*>(&out[0]);
+  *reinterpret_cast<uint4*>(&q[obase + 16]) = *reinterpret_cast<uint4*>(&out[16]);
+  e8_out[n * (M_DIM / BLK) + tile_m] = (uint8_t)e8;
+}
@@ -37,7 +37,7 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs)
    gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
    FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
    sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
  dummy = Tensor.zeros(1).contiguous().realize()
  out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
  linear = out.schedule_linear()
@@ -1,6 +1,6 @@
 from tinygrad.tensor import Tensor
 from tinygrad.nn import Conv2d, LayerNorm, LayerNorm2d, Linear
-from tinygrad.helpers import fetch, get_child
+from tinygrad.helpers import fetch, get_child, Context

 class Block:
  def __init__(self, dim):
@@ -58,7 +58,6 @@ if __name__ == "__main__":
  from test.models.test_efficientnet import chicken_img, preprocess, _LABELS
  img = Tensor(preprocess(chicken_img))

-  Tensor.training = False
-
-  out = model(img).numpy()
+  with Context(TRAINING=0):
+    out = model(img).numpy()
  print(_LABELS[out.argmax()])
@@ -5,7 +5,7 @@ import numpy as np
 from pathlib import Path
 from tinygrad import nn, Tensor, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.helpers import get_child, fetch
+from tinygrad.helpers import get_child, fetch, TRAINING
 from tinygrad.nn.state import torch_load
 from examples.mlperf.helpers import BoxCoder
 from extra.models.resnet import ResNet
@@ -1069,7 +1069,7 @@ class RoIBoxHead:
  def __call__(self, features, proposals, targets=None):
    x = self.feature_extractor(features, proposals)
    class_logits, box_regression = self.predictor(x)
-    if not Tensor.training:
+    if not TRAINING:
      result = self.post_processor((class_logits, box_regression), proposals)
      return x, result, {}

@@ -1111,7 +1111,7 @@ class Mask:
    x = self.feature_extractor(features, proposals)
    if x is not None:
      mask_logits = self.predictor(x)
-      if not Tensor.training:
+      if not TRAINING:
        result = self.post_processor(mask_logits, proposals)
        return x, result, {}
    return x, [], {}
@@ -1,6 +1,6 @@
 import math
 from tinygrad import Tensor, dtypes
-from tinygrad.helpers import flatten, get_child
+from tinygrad.helpers import flatten, get_child, TRAINING
 from examples.mlperf.helpers import generate_anchors, BoxCoder
 from examples.mlperf.losses import sigmoid_focal_loss, l1_loss
 from extra.models.resnet import ResNet
@@ -141,7 +141,7 @@ class ClassificationHead:
    out = [self.cls_logits(feat.sequential(self.conv)).permute(0, 2, 3, 1).reshape(feat.shape[0], -1, self.num_classes) for feat in x]
    out = out[0].cat(*out[1:], dim=1)

-    if Tensor.training:
+    if TRAINING:
      assert labels is not None and matches is not None, "labels and matches should be passed in when training"
      return self._compute_loss(out.cast(dtypes.float32), labels, matches)

@@ -167,7 +167,7 @@ class RegressionHead:
    out = [self.bbox_reg(feat.sequential(self.conv)).permute(0, 2, 3, 1).reshape(feat.shape[0], -1, 4) for feat in x]
    out = out[0].cat(*out[1:], dim=1)

-    if Tensor.training:
+    if TRAINING:
      assert bboxes is not None and matches is not None and anchors is not None, "bboxes, matches, and anchors should be passed in when training"
      return self._compute_loss(out, bboxes, matches, anchors)

@@ -187,7 +187,7 @@ class RetinaHead:
    self.regression_head = RegressionHead(in_channels, num_anchors)

  def __call__(self, x:Tensor, **kwargs) -> Tensor|dict[str, Tensor]:
-    if Tensor.training:
+    if TRAINING:
      return {
        "classification_loss": self.classification_head(x, labels=kwargs["labels"], matches=kwargs["matches"]),
        "regression_loss": self.regression_head(x, bboxes=kwargs["bboxes"], matches=kwargs["matches"], anchors=kwargs["anchors"])
@@ -111,7 +111,7 @@ def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, device:str, arch:st
  lib = bytes(lib)

  return UOp(Ops.PROGRAM,
-             src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))

@functools.cache
 def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
@@ -141,7 +141,7 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
  lib = bytes(lib)

  return UOp(Ops.PROGRAM,
-             src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))

@functools.cache
 def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
@@ -171,7 +171,7 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
  lib = bytes(lib)

  return UOp(Ops.PROGRAM,
-             src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))

@functools.cache
 def custom_fa_backward_post(dq_out:UOp, dq_in:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
@@ -201,4 +201,4 @@ def custom_fa_backward_post(dq_out:UOp, dq_in:UOp, device:str, arch:str, B:int,
  lib = bytes(lib)

  return UOp(Ops.PROGRAM,
-             src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
@@ -1,10 +1,14 @@
 import unittest
 import torch
 import tinygrad.nn.torch
-torch.set_default_device("tiny")
 import numpy as np

 class TestTorchBackendInplace(unittest.TestCase):
+  def setUp(self):
+    self._prev_device = torch.get_default_device()
+    torch.set_default_device("tiny")
+  def tearDown(self): torch.set_default_device(self._prev_device)
+
  def test_zero(self):
    a = torch.ones(4)
    a.zero_()
@@ -2,11 +2,15 @@ import unittest
 from tinygrad.helpers import getenv
 import torch
 import tinygrad.nn.torch
-torch.set_default_device("tiny")
 import numpy as np

@unittest.skipIf(getenv("GPUS",1)<=1, "only single GPU")
 class TestTorchBackendMultiGPU(unittest.TestCase):
+  def setUp(self):
+    self._prev_device = torch.get_default_device()
+    torch.set_default_device("tiny")
+  def tearDown(self): torch.set_default_device(self._prev_device)
+
  def test_transfer(self):
    a = torch.Tensor([[1,2],[3,4]]).to("tiny:0")
    b = torch.Tensor([[3,2],[1,0]]).to("tiny:1")
@@ -40,7 +40,7 @@ def train(model, X_train, Y_train, optim, steps, BS=128, lossfn=lambda out,y: ou

 def evaluate(model, X_test, Y_test, num_classes=None, BS=128, return_predict=False, transform=lambda x: x,
             target_transform=lambda y: y):
-  Tensor.training = False
+  @Context(TRAINING=0)
  def numpy_eval(Y_test, num_classes):
    Y_test_preds_out = np.zeros(list(Y_test.shape)+[num_classes])
    for i in trange((len(Y_test)-1)//BS+1, disable=None):
@@ -20,6 +20,7 @@ include-package-data = true
 packages = [
  'tinygrad',
  'tinygrad.codegen',
+  'tinygrad.codegen.decomp',
  'tinygrad.codegen.opt',
  'tinygrad.codegen.late',
  'tinygrad.engine',
@@ -30,7 +30,7 @@ def custom_add_one(A:UOp) -> UOp:
    s_endpgm(),
  ]
  sink = UOp.sink(A.base, threads, arg=KernelInfo(f"custom_add_one_{A.numel()}", estimates=Estimates(ops=A.numel(), mem=A.numel()*4*2)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

 def custom_add_var(A:UOp, B:UOp) -> UOp:
  A,B = A.flatten(), B.flatten()
@@ -49,7 +49,7 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
    s_endpgm(),
  ]
  sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.numel()}"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

 def custom_wave_sync(A:UOp, arch:str) -> UOp:
  # 4 waves across 1024 WG — enough to saturate a SIMD with many concurrent WGs
@@ -63,7 +63,7 @@ def custom_wave_sync(A:UOp, arch:str) -> UOp:
    insts += [s_nop(0)]*4
  insts.append(s_endpgm())
  sink = UOp.sink(A.base, threads, wg, arg=KernelInfo("custom_wave_sync"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

 def custom_lds_sync(A:UOp, arch:str) -> UOp:
  A = A.flatten()
@@ -97,7 +97,7 @@ def custom_lds_sync(A:UOp, arch:str) -> UOp:
    isa.s_endpgm(),
  ]
  sink = UOp.sink(A.base, lds, threads, wg, arg=KernelInfo("custom_lds_sync"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

 def custom_handwritten(A:UOp) -> UOp:
  A = A.flatten()
@@ -143,7 +143,7 @@ def custom_handwritten(A:UOp) -> UOp:
  k.emit(r4.s_endpgm())
  insts = k.finalize()
  sink = UOp.sink(A.base, threads, wg, lds, arg=KernelInfo("custom_handwritten"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

 def custom_data_deps(A:UOp) -> UOp:
  A = A.flatten()
@@ -159,7 +159,7 @@ def custom_data_deps(A:UOp) -> UOp:
  k.emit(s_endpgm())
  insts = k.finalize()
  sink = UOp.sink(A.base, threads, arg=KernelInfo("custom_data_deps"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestAsmKernel(unittest.TestCase):
@@ -89,7 +89,7 @@ class TestTinygradIntegration(unittest.TestCase):
    assert len(sink_items) > 0, "No SINK in schedule"
    renderer = AMDLLVMRenderer(Target("AMD", arch='gfx1100'))
    prg = to_program(sink_items[0].src[0], renderer)
-    lib = renderer.compiler.compile(prg.src[3].arg)
+    lib = renderer.compiler.compile(prg.src[2].arg)
    return next(s.content for s in elf_loader(lib)[1] if s.name == ".text")

  def test_simple_add_kernel(self):
@@ -15,7 +15,7 @@ dev = Device["AMD"]
 a = Tensor([1.0]).realize()
 b = a + 1
 linear = compile_linear(b.schedule_linear())
-lib = bytearray(linear.src[-1].src[0].src[4].arg)
+lib = bytearray(linear.src[-1].src[0].src[3].arg)

 # Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
 # VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400
@@ -83,8 +83,8 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, in
            buf_data[dst_id] = src_data
      elif ast.op is Ops.PROGRAM:
        info = ast.arg
-        if len(ast.src) > 4 and ast.src[4].op is Ops.BINARY:
-          lib = bytes(ast.src[4].arg)
+        if len(ast.src) > 3 and ast.src[3].op is Ops.BINARY:
+          lib = bytes(ast.src[3].arg)
          _, sections, _ = elf_loader(lib)
          for sec in sections:
            if sec.name == '.text':
@@ -98,7 +98,7 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, in
                buf_sizes.append(b.nbytes)
              kernels.append(KernelSnapshot(
                code=bytes(sec.content),
-                src=ast.src[3].arg,
+                src=ast.src[2].arg,
                global_size=tuple(info.global_size),
                local_size=tuple(info.local_size),
                buf_idxs=buf_idxs,
@@ -430,7 +430,7 @@ class TestCustomKernel(unittest.TestCase):
    binary = Device[a.device].renderer.compiler.compile(src)
    def custom_src_kernel(A:UOp) -> UOp:
      sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="CPU"), UOp(Ops.LINEAR, src=tuple(sink.toposort())),
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())),
                                   UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))

    a = Tensor.custom_kernel(a, fxn=custom_src_kernel)[0]
@@ -47,7 +47,7 @@ class TestLinearizer(unittest.TestCase):
    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
    out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
    ast = helper_linearizer_opt(out)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)

  @unittest.expectedFailure
@@ -55,7 +55,7 @@ class TestLinearizer(unittest.TestCase):
    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
    out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
    ast = helper_linearizer_opt(out)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)

  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
@@ -65,7 +65,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.empty(16)
    out = img.conv2d(w, b)
    ast = helper_linearizer_opt(out)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    # slice at the last loop end
    uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
    # only valid test if outermost range is the reduce
@@ -86,7 +86,7 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor.randn(2, ).realize()
    out = a.reshape(2, 1).expand(2, 3).sum()
    ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

@@ -94,7 +94,7 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor.randn(2, ).realize()
    out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
    ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

@@ -102,7 +102,7 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor([2, 2]).realize()
    out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
    ast = helper_linearizer_opt(out, wanna_output=[24])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    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]])
@@ -115,7 +115,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.randn(1, 1).realize()
    out = (a + b[0]).sum() + b[0]
    ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    # LOAD -> RANGE -> LOAD -> STORE
    assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
@@ -125,7 +125,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.randn(1, 1).realize()
    out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
    ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

@@ -137,10 +137,10 @@ class TestLinearizer(unittest.TestCase):
    r = a[:-1] + a[1:]

    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+                       renderer=Device[Device.DEFAULT].renderer).src[1].src)
    num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
    assert num_loads <= 4, "more load uops than needed"
-    assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
+    assert num_loads >= 1, "expected at least one load uop"

  @unittest.skip("this is handled at higher level now")
  def test_upcast_cse(self):
@@ -150,7 +150,7 @@ class TestLinearizer(unittest.TestCase):
    r = a.expand([2]) + b.expand([2])

    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+                       renderer=Device[Device.DEFAULT].renderer).src[1].src)
    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops <= 1, "more alu uops than needed"

@@ -160,7 +160,7 @@ class TestLinearizer(unittest.TestCase):
    r = Tensor.conv2d(x,w,padding=1).relu()

    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0],
-      [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+      [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    accs = [u for u in uops if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
    stores = [u for u in uops if u.op is Ops.STORE]
    assert len(accs) == 0  # it's removed now
@@ -173,7 +173,7 @@ class TestLinearizer(unittest.TestCase):
    out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
    prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.LOCAL, axis=0, arg=4)]),
                      renderer=Device[Device.DEFAULT].renderer)
-    self.assertEqual(len(prg.src[3].arg.split("for")), 5)
+    self.assertEqual(len(prg.src[2].arg.split("for")), 5)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -185,7 +185,7 @@ class TestLinearizer(unittest.TestCase):
    opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
    program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer)

-    stores = [u for u in tuple(program.src[2].src) if u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG]
+    stores = [u for u in tuple(program.src[1].src) if u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG]

    # the first store is to lds and can be upcasted
    assert stores[0].src[1].max_numel() == 4
@@ -199,7 +199,7 @@ class TestLinearizer(unittest.TestCase):
    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
    r = Tensor.stack(a, b)
    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+                       renderer=Device[Device.DEFAULT].renderer).src[1].src)
    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops == 0, "more alu uops than needed"

@@ -210,14 +210,14 @@ class TestLinearizer(unittest.TestCase):
        a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
        realized_ast = a.schedule_linear().src[-1].src[0]
        program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
-        local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
+        local = [uop for uop in tuple(program.src[1].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
        assert local[0].dtype.base == acc_dtype

  def test_arg_acc_dtype(self):
    def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
      realized_ast = c.schedule_linear().src[-1].src[0]
      program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
-      local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
+      local = [uop for uop in tuple(program.src[1].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
      self.assertEqual(local[0].dtype.base, expected_dtype)

    tests = (
@@ -244,7 +244,7 @@ class TestLinearizer(unittest.TestCase):
    opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
    ast = helper_linearizer_opt(r, [opt])
    # the uops graph is reg BUFFER -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
-    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
    end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
@@ -264,7 +264,7 @@ class TestLinearizer(unittest.TestCase):
    # shrink so that the dims do not collapse
    t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
    ast = helper_linearizer_opt(t+1)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
    idxs = sorted(idxs, key=lambda uop: uop.arg)
    assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
@@ -297,12 +297,12 @@ class TestLinearizer(unittest.TestCase):
    run_linear(linear, var_vals)
    np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
    program = to_program(replace_opts(linear.src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer)
-    assert not any(u.op == Ops.WHERE for u in tuple(program.src[2].src)), "found where where where should be folded"
+    assert not any(u.op == Ops.WHERE for u in tuple(program.src[1].src)), "found where where where should be folded"

  def test_phi_simplification(self):
    def helper(t, max_ops=0):
      ast = helper_linearizer_opt(t)
-      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
+      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src)
      # ignore kernel optimized IF statements for now
      if if_op:=next((u for u in uops if u.op is Ops.IF), None):
        uops = uops[:uops.index(if_op)]
@@ -335,7 +335,7 @@ class TestLinearizer(unittest.TestCase):
    out = x.matmul(y)
    with Context(TC=0):
      ast = helper_linearizer_opt(out)
-      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
+      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src)
    # check that the float4 cast collapses
    store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
    for val in store_vals:
@@ -347,7 +347,7 @@ class TestLinearizer(unittest.TestCase):
    x = Tensor.randn((4,3,6,6)).realize()
    out = x.flip((0,1)).contiguous()
    ast = helper_linearizer_opt(out)
-    store_val = [u.src[1] for u in tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src) if u.op is Ops.STORE][0]
+    store_val = [u.src[1] for u in tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src) if u.op is Ops.STORE][0]
    assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.STACK

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -360,15 +360,14 @@ class TestLinearizer(unittest.TestCase):
            Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
    ast = helper_linearizer_opt(out, opts=[opt])
    def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
-    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
-    local_stores = [u for u in uops if u.op is Ops.STORE and any(
-      x.op is Ops.BUFFER and x.addrspace is AddrSpace.LOCAL for x in get_recursive(u.src[0]))]
-    global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.PARAM for x in get_recursive(u.src[0]))]
+    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    local_stores = [u for u in uops if u.op is Ops.STORE and u.addrspace == AddrSpace.LOCAL]
+    global_stores = [u for u in uops if u.op is Ops.STORE and u.addrspace == AddrSpace.GLOBAL]
    barrier = [u for u in uops if u.op is Ops.BARRIER]
    assert len(barrier) == 1
    # check that the float4 cast collapses for all stores
    for store in local_stores+global_stores:
-      assert store.src[1].max_numel() > 1 # and store.src[2].op is not Ops.VECTORIZE
+      assert store.src[1].max_numel() > 1, f"store shape {store.src[1].shape} on {store.addrspace}"
    # # check the children's vins
    # TODO: src ALU are not the same, should it?
    # assert barrier.src == tuple(local_stores)
@@ -383,7 +382,7 @@ class TestLinearizer(unittest.TestCase):
    x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
    r = (x@y).relu()
    ast = helper_linearizer_opt(r)
-    uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src)
    stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]

    # the float4 value stores directly in lds and we skip upcast
@@ -3,6 +3,7 @@ import torch
 import unittest
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.nn.optim import Adam, SGD, AdamW, Muon, LAMB
+from tinygrad.helpers import Context
 from test.helpers import needs_second_gpu, slow

 np.random.seed(1337)
@@ -45,11 +46,7 @@ def step(tensor, optim, steps=1, teeny=False, **kwargs):

@slow
 class TestOptim(unittest.TestCase):
-  def setUp(self):
-    self.old_training = Tensor.training
-    Tensor.training = True
-  def tearDown(self):
-    Tensor.training = self.old_training
+  def setUp(self): self.enterContext(Context(TRAINING=1))

  def _test_optim(self, tinygrad_optim, torch_optim, steps, opts, atol, rtol):
    for x,y in zip(step(Tensor, tinygrad_optim, steps, **opts),
@@ -157,13 +154,11 @@ class TestOptim(unittest.TestCase):
    t = Tensor.ones((1,1))
    optimizer = Adam([t])
    optimizer.zero_grad()
-    old_state = Tensor.training
    t.sum().backward()
-    Tensor.training = False
-    self.assertRaises(RuntimeError, optimizer.step)
-    Tensor.training = True
-    optimizer.step()
-    Tensor.training = old_state
+    with Context(TRAINING=0):
+      self.assertRaises(RuntimeError, optimizer.step)
+    with Context(TRAINING=1):
+      optimizer.step()

  def test_lamb_cpu_offload(self):
    # test that LAMB works when optimizer params (m, v, b1_t, b2_t) are moved to CPU
@@ -43,10 +43,10 @@ def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
  call = linear.src[-1]
  prg = to_program(replace_opts(call.src[0], opts), renderer=Device[Device.DEFAULT].renderer)
  if replace_src is not None:
-    old_name = prg.src[3].arg.split("__attribute__((noinline)) void ")[1].split("(")[0]
-    new_src = replace_src + "/* DSP boilerplate */" + prg.src[3].arg.split("/* DSP boilerplate */")[1].replace(old_name, "fxn")
+    old_name = prg.src[2].arg.split("__attribute__((noinline)) void ")[1].split("(")[0]
+    new_src = replace_src + "/* DSP boilerplate */" + prg.src[2].arg.split("/* DSP boilerplate */")[1].replace(old_name, "fxn")
    # drop BINARY and replace SOURCE so run_linear recompiles
-    prg = prg.replace(src=prg.src[:3] + (UOp(Ops.SOURCE, arg=new_src),))
+    prg = prg.replace(src=prg.src[:2] + (UOp(Ops.SOURCE, arg=new_src),))
  linear = linear.replace(src=linear.src[:-1] + (call.replace(src=(prg, *call.src[1:])),))
  for _ in range(run_count): run_linear(linear)

@@ -81,7 +81,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
    with Context(QUANTIZE=1):
      linear = run_onnx({"input":inp})["output"].schedule_linear()
      prg = to_program(linear.src[-2].src[0], renderer=Device[Device.DEFAULT].renderer)
-      daccs = [u for u in tuple(prg.src[2].src) if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
+      daccs = [u for u in tuple(prg.src[1].src) if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
      assert all(u.dtype.scalar() is dtypes.int for u in daccs)

@unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")
@@ -77,7 +77,7 @@ class TestRandomness(unittest.TestCase):
      ast = call.src[0]
      if ast.op is Ops.SINK:
        prg = to_program(ast, renderer=Device[Device.DEFAULT].renderer)
-        for u in tuple(prg.src[2].src):
+        for u in tuple(prg.src[1].src):
          self.assertNotIn(u.dtype, {dtypes.long, dtypes.ulong}, msg=f"long found in {prg.arg.name}")

  def test_threefry_against_reference_full(self):
@@ -63,7 +63,7 @@ class TestCStyleFailures(unittest.TestCase):
    for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype))
    linear = ret.schedule_linear()
    assert len(linear.src) == 1
-    src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[3].arg
+    src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[2].arg
    self.assertEqual("("*5 not in src, should_strip_paren)

  def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)
@@ -166,6 +166,13 @@ class TestSchedule(unittest.TestCase):
    self.assertEqual(t.device, dev)
    np.testing.assert_equal(t[:64].numpy(), np.arange(64).cumsum())

+  def test_copy_multi_scalar(self):
+    devs = ("CPU:0", "CPU:1")
+    x = Tensor.ones(2, device="CPU").shard(devs, axis=0).realize()
+    out = (x.sum()*2).reshape(1).to("CPU")
+    run_linear(*check_schedule(out, 5))
+    np.testing.assert_equal(out.numpy(), [4.])
+
 class TestLimitBufs(unittest.TestCase):
  @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV", "crashes in ocelot")
  def test_limit_bufs_with_var(self):
@@ -448,6 +455,14 @@ class TestCopyFolding(unittest.TestCase):
    b.realize()
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])

+  def test_permute_copy_to_device(self):
+    b = Tensor([[0, 1, 2, 3], [4, 5, 6, 7]], device="CPU").permute(1, 0).to("PYTHON")
+    self.assertListEqual(b.tolist(), [[0, 4], [1, 5], [2, 6], [3, 7]])
+
+  def test_flip_copy_to_device(self):
+    b = Tensor([0, 1, 2, 3], device="CPU").flip(0).to("PYTHON")
+    self.assertListEqual(b.tolist(), [3, 2, 1, 0])
+
 class TestUOpBecome(unittest.TestCase):
  def test_setitem_offset(self):
    a = Tensor.full((16,), 0.).contiguous().realize()
@@ -249,7 +249,7 @@ class TestAssembly(unittest.TestCase):
    opts_to_apply = [Opt(OptOps.UNROLL, 0, 4)]
    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
    program = to_program(ast, Device[Device.DEFAULT].renderer)
-    uops = tuple(program.src[2].src)
+    uops = tuple(program.src[1].src)
    self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)

  def test_mulacc_shl(self):
@@ -0,0 +1,68 @@
+from tinygrad import Tensor, dtypes, Device, TinyJit
+from tinygrad.helpers import getenv, Timing, round_up
+
+# llama 70b shapes
+DIM, HIDDEN, N_HEADS, N_KV_HEADS, HEAD_DIM = 8192, 28672, 64, 8, 128
+QKV_OUT = N_HEADS*HEAD_DIM + 2*N_KV_HEADS*HEAD_DIM
+LAYER_SHAPES = {
+  "wqkv": ((QKV_OUT, DIM), 0),
+  "wo": ((DIM, N_HEADS*HEAD_DIM), 1),
+  "w1": ((HIDDEN, DIM), 0),
+  "w3": ((HIDDEN, DIM), 0),
+  "w2": ((DIM, HIDDEN), 1),
+}
+
+def build_params(layers:int, devs:tuple[str,...], dtype):
+  ts, nbytes = [], 0
+  ndev = len(devs)
+  for _ in range(layers):
+    for (a,b), ax in LAYER_SHAPES.values():
+      shape = (round_up(a, ndev) if ax==0 else a, round_up(b, ndev) if ax==1 else b)
+      t = Tensor.empty(*shape, dtype=dtype).shard(devs, axis=ax).realize()
+      ts.append(t)
+      nbytes += t.numel()*dtype.itemsize
+  return ts, nbytes
+
+def sync():
+  Device["CPU"].synchronize()
+  Device[Device.DEFAULT].synchronize()
+
+if __name__ == "__main__":
+  ngpus = getenv("GPUS", 8)
+  layers = getenv("LAYERS", 8)
+  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(ngpus))
+
+  # to cpu
+  grads, nbytes = build_params(layers, devs, dtypes.bfloat16)
+  cpu_dst = [Tensor.rand(*g.shape, dtype=dtypes.bfloat16, device="CPU").contiguous().realize() for g in grads]
+  sync()
+  def to_cpu():
+    for d,g in zip(cpu_dst, grads):
+      d.assign(g.to("CPU"))
+    return Tensor.realize(*cpu_dst)
+  to_cpu_jit = TinyJit(to_cpu)
+  for _ in range(3):
+    to_cpu_jit()
+    sync()
+  for _ in range(5):
+    with Timing("to cpu ", on_exit=lambda ns: f" @ {nbytes/ns:.2f} GB/s"):
+      to_cpu_jit()
+      sync()
+
+  # to gpu
+  from extra.gemm.cdna_asm_gemm import FP8_DTYPE
+  w_gpu, nbytes = build_params(layers, devs, FP8_DTYPE)
+  cpu_src = [Tensor.rand(*w.shape, dtype=FP8_DTYPE, device="CPU").contiguous().realize() for w in w_gpu]
+  sync()
+  def to_gpu():
+    for w,c in zip(w_gpu, cpu_src):
+      w.assign(c.shard_like(w))
+    return Tensor.realize(*w_gpu)
+  to_gpu_jit = TinyJit(to_gpu)
+  for _ in range(3):
+    to_gpu_jit()
+    sync()
+  for _ in range(5):
+    with Timing("to gpu ", on_exit=lambda ns: f" @ {nbytes/ns:.2f} GB/s"):
+      to_gpu_jit()
+      sync()
@@ -1,16 +1,15 @@
 from tinygrad import nn, Tensor, dtypes
-from tinygrad.helpers import DEV, Timing
+from tinygrad.helpers import DEV, Timing, Context

 from extra.models.llama import Transformer
 from examples.llama3 import MODEL_PARAMS

 if __name__ == "__main__":
  DEV.value = "NULL"
-  Tensor.training = True
  #model_size = "8B"
  model_size = "405B"

-  with Timing("total "):
+  with Context(TRAINING=1), Timing("total "):
    with Timing("***** create model in    "):
      model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=nn.Linear, embedding=nn.Embedding,
                          max_context=1024, jit=True, disable_kv_cache=True)
@@ -63,8 +63,8 @@ class BenchmarkResnetTrain(unittest.TestCase):
    with Context(TRACK_MATCH_STATS=0): Tensor.realize(*[t.assign(t.detach().contiguous()) for t in get_parameters(optim)])

    JITCNT = getenv("JITCNT", 1)
-    Tensor.training = True
    @TinyJit
+    @Context(TRAINING=1)
    def step(x):
      optim.zero_grad()
      x.grad = None
@@ -0,0 +1,67 @@
+import unittest, math
+import numpy as np
+from tinygrad.tensor import Tensor
+from tinygrad.nn.optim import AdamW
+from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup, LambdaLR, LambdaLinearScheduler
+
+np.random.seed(1337)
+x_init = np.random.randn(1,4).astype(np.float32)
+W_init = np.random.randn(4,4).astype(np.float32)
+m_init = np.random.randn(1,4).astype(np.float32)
+
+class TinyNet:
+  def __init__(self):
+    self.x = Tensor(x_init.copy())
+    self.W = Tensor(W_init.copy())
+    self.m = Tensor(m_init.copy())
+
+  def forward(self):
+    out = self.x.matmul(self.W).relu()
+    out = out.log_softmax(1)
+    out = out.mul(self.m).add(self.m).sum()
+    return out
+
+class TestCosineAnnealingLRWithWarmup(unittest.TestCase):
+  # only tests the lr
+  def _test_lr(self, base_lr, end_lr, warmup_steps, decay_steps):
+    net = TinyNet()
+    optim = AdamW([net.W], lr=0.0)
+    tiny_lr = CosineAnnealingLRWithWarmup(optim, base_lr, end_lr, warmup_steps, decay_steps)
+    lr = []
+    for _ in range(warmup_steps+decay_steps):
+      lr.append(optim.lr.item())
+      tiny_lr.step()
+    # reimplemented in python
+    expected = []
+    for i in range(warmup_steps): expected.append((i+1)/warmup_steps*base_lr)
+    for i in range(decay_steps): expected.append(end_lr+(base_lr-end_lr)*(1+math.cos((i+1)/decay_steps*math.pi))/2)
+    np.testing.assert_allclose(lr, expected, rtol=1e-5)
+
+  def test_lr_0(self): self._test_lr(3e-4, 8e-5, 3, 5)
+  def test_lr_1(self): self._test_lr(3e-4, 8e-5, 10, 20)
+  def test_lr_llama3(self): self._test_lr(8e-5, 8e-7, 20, 100)
+
+class TestLambdaLRLinearWarmup(unittest.TestCase):
+  def test_linear_lr_warmup(self):
+    BS, BASE_LR = 304, 2.5e-7
+    lr = BS * BASE_LR
+    # Use a dummy Tensor parameter for optimizer because the lr_scheduler only needs the optimizer's device and lr, the params aren't touched.
+    optimizer = AdamW([Tensor([1.])])
+    lambda_lr_callback = LambdaLinearScheduler(1000, 1.0, 1.0, 1e-06, 10000000000000).schedule
+    lr_scheduler = LambdaLR(optimizer, Tensor(lr, device=optimizer.device), lambda_lr_callback)
+    lrs = {}
+
+    # with above settings, optimizer.lr should warm up to lr over 1000 steps linearly
+    for i in range(1200):
+      lr_scheduler.step()
+      if i in {0, 499, 998, 999, 1000, 1199}:
+        lrs[i] = optimizer.lr.item()
+
+    np.testing.assert_allclose(lr, lrs[999], rtol=0, atol=1e-11)
+    np.testing.assert_equal(lrs[999], lrs[1000])
+    np.testing.assert_equal(lrs[999], lrs[1199])
+    np.testing.assert_allclose(lrs[999] / lrs[0], 1000, rtol=0, atol=1)
+    np.testing.assert_allclose(lrs[999] / lrs[499], 2, rtol=0, atol=1e-5)
+
+if __name__ == '__main__':
+  unittest.main()
@@ -4,7 +4,7 @@ import numpy as np
 import torch

 from tinygrad import GlobalCounters, Tensor, Device
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, Context
 from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import capturing, run_linear
 from tinygrad.tensor import _to_np_dtype
@@ -38,11 +38,7 @@ from extra.models.vit import ViT

@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
 class TestInferenceMinKernels(unittest.TestCase):
-  def setUp(self):
-    self.training_old = Tensor.training
-    Tensor.training = False
-  def tearDown(self):
-    Tensor.training = self.training_old
+  def setUp(self): self.enterContext(Context(TRAINING=0))

  def test_convnext(self):
    model = ConvNeXt()
@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-import unittest, math
+import unittest
 import numpy as np
 import tensorflow as tf
 from tensorflow.keras.optimizers import Lamb
@@ -7,11 +7,12 @@ from tensorflow.python.ops import math_ops
 from extra.lr_scheduler import LRSchedulerGroup

 from tinygrad.tensor import Tensor
-from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, AdamW
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup
+from tinygrad.helpers import Context

 from test.external.mlperf_resnet.lars_optimizer import LARSOptimizer

-from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup, CosineAnnealingLRWithWarmup, LambdaLR, LambdaLinearScheduler
+from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup
 from test.external.mlperf_resnet.lars_util import PolynomialDecayWithWarmup as PolynomialDecayWithWarmup_tf

 np.random.seed(1337)
@@ -103,11 +104,7 @@ def create_tf_polylr(initial_lr, end_lr, train_steps, warmup, power=2, skip_list
                                      initial_learning_rate=initial_lr, end_learning_rate=end_lr, warmup_epochs=warmup)

 class ExternalTestOptim(unittest.TestCase):
-  def setUp(self):
-    self.old_training = Tensor.training
-    Tensor.training = True
-  def tearDown(self):
-    Tensor.training = self.old_training
+  def setUp(self): self.enterContext(Context(TRAINING=1))

  def _test_optim(self, tinygrad_optim, tensorflow_optim, steps, opts, atol, rtol, tiny_sched=None, tf_sched=None, schedopts=None, do_optim=True):
    for x,y in zip(step(tinygrad_optim, steps=steps, kwargs=opts, scheduler=tiny_sched, schedopts=schedopts, do_optim=do_optim),
@@ -173,48 +170,5 @@ class ExternalTestOptim(unittest.TestCase):
      'warmup': steps_per_epoch * warmup_epochs,
    }, 1e-5, 1e-5, do_optim=False)

-
-class TestCosineAnnealingLRWithWarmup(unittest.TestCase):
-  # only tests the lr
-  def _test_lr(self, base_lr, end_lr, warmup_steps, decay_steps):
-    net = TinyNet()
-    optim = AdamW([net.W], lr=0.0)
-    tiny_lr = CosineAnnealingLRWithWarmup(optim, base_lr, end_lr, warmup_steps, decay_steps)
-    lr = []
-    for _ in range(warmup_steps+decay_steps):
-      lr.append(optim.lr.item())
-      tiny_lr.step()
-    # reimplemented in python
-    expected = []
-    for i in range(warmup_steps): expected.append((i+1)/warmup_steps*base_lr)
-    for i in range(decay_steps): expected.append(end_lr+(base_lr-end_lr)*(1+math.cos((i+1)/decay_steps*math.pi))/2)
-    np.testing.assert_allclose(lr, expected, rtol=1e-5)
-
-  def test_lr_0(self): self._test_lr(3e-4, 8e-5, 3, 5)
-  def test_lr_1(self): self._test_lr(3e-4, 8e-5, 10, 20)
-  def test_lr_llama3(self): self._test_lr(8e-5, 8e-7, 20, 100)
-
-class TestLambdaLRLinearWarmup(unittest.TestCase):
-  def test_linear_lr_warmup(self):
-    BS, BASE_LR = 304, 2.5e-7
-    lr = BS * BASE_LR
-    # Use a dummy Tensor parameter for optimizer because the lr_scheduler only needs the optimizer's device and lr, the params aren't touched.
-    optimizer = AdamW([Tensor([1.])])
-    lambda_lr_callback = LambdaLinearScheduler(1000, 1.0, 1.0, 1e-06, 10000000000000).schedule
-    lr_scheduler = LambdaLR(optimizer, Tensor(lr, device=optimizer.device), lambda_lr_callback)
-    lrs = {}
-
-    # with above settings, optimizer.lr should warm up to lr over 1000 steps linearly
-    for i in range(1200):
-      lr_scheduler.step()
-      if i in {0, 499, 998, 999, 1000, 1199}:
-        lrs[i] = optimizer.lr.item()
-
-    np.testing.assert_allclose(lr, lrs[999], rtol=0, atol=1e-11)
-    np.testing.assert_equal(lrs[999], lrs[1000])
-    np.testing.assert_equal(lrs[999], lrs[1199])
-    np.testing.assert_allclose(lrs[999] / lrs[0], 1000, rtol=0, atol=1)
-    np.testing.assert_allclose(lrs[999] / lrs[499], 2, rtol=0, atol=1e-5)
-
 if __name__ == '__main__':
  unittest.main()
@@ -3,7 +3,7 @@ from tinygrad import Tensor, UOp, Device, nn
 from tinygrad.schedule import schedule_cache
 from tinygrad.codegen import to_program, to_program_cache
 from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
-from tinygrad.uop.divandmod import fold_divmod_general
+from tinygrad.codegen.decomp.divandmod import fold_divmod_general
 from test.test_tiny import TestTiny

 def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
@@ -3,7 +3,7 @@ import z3
 from tinygrad import dtypes, Device
 from tinygrad.uop.validate import uops_to_z3, z3_cdiv
 from tinygrad.uop.ops import UOp
-from tinygrad.uop.decompositions import fast_idiv
+from tinygrad.codegen.decomp.op import fast_idiv
 random.seed(42)

 powers_of_two = [2**i for i in range(64)]
@@ -11,7 +11,7 @@ if __name__ == "__main__":
  for i in range(10_000):
    if i % 1000 == 0:
      print(f"Progress: {i}")
-    dt = random.choice(dtypes.ints + tuple(dt.vec(4) for dt in dtypes.ints))
+    dt = random.choice(dtypes.ints)
    u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt)
    d = random.randint(1, max(1, u.vmax)*2)
    if d in powers_of_two: continue
@@ -41,8 +41,8 @@ class BenchmarkBertTrain(unittest.TestCase):
    with Context(TRACK_MATCH_STATS=0): Tensor.realize(*[t.assign(t.detach().contiguous()) for t in get_parameters(optim)])

    JITCNT = getenv("JITCNT", 1)
-    Tensor.training = True
    @TinyJit
+    @Context(TRAINING=1)
    def step(inputs):
      optim.zero_grad()
      for i in inputs: i.grad = None
@@ -31,8 +31,6 @@ if __name__ == "__main__":
  assert all(os.path.exists(os.path.join(INIT_CKPT_DIR, f)) for f in required_files), \
    f"Missing checkpoint files in INIT_CKPT_DIR: {required_files}"

-  Tensor.training = False
-
  model = get_mlperf_bert_model(INIT_CKPT_DIR)

  for _, x in get_state_dict(model).items():
@@ -48,9 +48,9 @@ def replay_to_program(p:UOp, ast:UOp, renderer:Renderer) -> tuple[str, str, tupl
    if sink_arg.beam: sink_arg = replace(sink_arg, opts_to_apply=p.src[0].arg.applied_opts)
    input_ast = ast.replace(arg=replace(sink_arg, name=p.src[0].arg.name))
  p2 = to_program(input_ast, renderer=renderer)
-  device = p.src[1].arg
+  device = renderer.target.device
  def to_str(ret:UOp) -> str:
-    src = ret.src[3].arg
+    src = next(x.arg for x in ret.src if x.op is Ops.SOURCE)
    # PYTHON renderer pickles UOps, first unpickle and decode here
    if device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(src))])
    return src
@@ -84,7 +84,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
  g = UOp.param(0, uop.dtype.ptr(1))
  prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0), ptr=True), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
-  prog = PythonProgram("run", PythonCompiler().compile(prg.src[3].arg))
+  prog = PythonProgram("run", PythonCompiler().compile(prg.src[2].arg))
  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]

@@ -375,7 +375,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32,
  """Conditional memory store with sub-word support. Returns list of store UOps."""
  adt = dtypes.uint64 if addr_bits == 64 else dtypes.uint32
  word_addr = addr >> UOp.const(adt, 2)
-  idx = mem.index(word_addr.cast(dtypes.int).valid(active))
+  idx = mem.index(word_addr.valid(active))
  if data_bits == 32: return [idx.store(active.where(_to_u32(val), idx))]
  # Sub-word store: read-modify-write with mask
  byte_pos = addr.cast(dtypes.uint32) & _c(3)
@@ -388,7 +388,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32,
  is_cross = byte_pos.eq(_c(3))
  cross_word0 = (idx & _c(0x00FFFFFF)) | ((val_u32 & _c(0xFF)) << _c(24))
  store0 = idx.store(active.where(is_cross.where(cross_word0, new_word), idx))
-  next_idx = mem.index((word_addr + UOp.const(adt, 1)).cast(dtypes.int).valid(active & is_cross))
+  next_idx = mem.index((word_addr + UOp.const(adt, 1)).valid(active & is_cross))
  cross_word1 = (next_idx & _c(0xFFFFFF00)) | ((val_u32 >> _c(8)) & _c(0xFF))
  return [store0, next_idx.store((active & is_cross).where(cross_word1, next_idx))]

@@ -398,7 +398,7 @@ def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int
  val_u32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
  for i in range(data_bits // 8):
    byte_val = (val_u32 >> UOp.const(dtypes.uint32, i * 8)) & UOp.const(dtypes.uint32, 0xFF)
-    stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).cast(dtypes.int).valid(active)).store(byte_val.cast(dtypes.uint8)))
+    stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).valid(active)).store(byte_val.cast(dtypes.uint8)))
  return stores

 def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode_vars: dict | None = None, op_name: str = "") -> dict[int, UOp]:
@@ -463,7 +463,7 @@ class _Ctx:
    """Read instruction dword from vmem at PC + dword_idx*4."""
    pc = self.rpc()
    addr = pc if dword_idx == 0 else pc + UOp.const(dtypes.uint64, dword_idx * 4)
-    return self.vmem.index((addr >> UOp.const(dtypes.uint64, 2)).cast(dtypes.int), ptr=True).load()
+    return self.vmem.index(addr >> UOp.const(dtypes.uint64, 2), ptr=True).load()

  def inst_field(self, field) -> UOp:
    """Extract field bits from instruction encoding. Tracks field for canonical key computation."""
@@ -516,14 +516,14 @@ class _Ctx:
  # Dynamic register access (takes UOp index instead of int)
  def rsgpr_dyn(self, reg: UOp, valid: UOp | None = None) -> UOp:
    """Read SGPR with dynamic register index."""
-    if valid is not None: return self.sgpr.index(reg.cast(dtypes.int).valid(valid), ptr=True).load()
-    return self.sgpr.index(reg.cast(dtypes.int), ptr=True).load()
+    if valid is not None: return self.sgpr.index(reg.valid(valid), ptr=True).load()
+    return self.sgpr.index(reg, ptr=True).load()

  def wsgpr_dyn(self, reg: UOp, val: UOp) -> UOp:
    """Write SGPR with dynamic register index. On RDNA, index 124 = NULL (writes discarded). On CDNA, index 124 = M0 (read/write)."""
    # RDNA: NULL (124) discards writes. CDNA: M0 (124) is writable.
    valid = None if self.wave_size == 64 else reg.ne(_c(124))
-    return self.sgpr.index(reg.cast(dtypes.int).valid(valid) if valid is not None else reg.cast(dtypes.int), ptr=True).store(val.cast(dtypes.uint32))
+    return self.sgpr.index(reg.valid(valid) if valid is not None else reg, ptr=True).store(val.cast(dtypes.uint32))

  def wmask(self, reg: UOp, val: UOp) -> list[UOp]:
    """Write a lane mask (VCC/EXEC). Splits into lo/hi for wave64."""
@@ -656,7 +656,7 @@ class _Ctx:
    stores = []
    for dest, val in assigns:
      if dest.startswith('D0'): stores.append(self.wsgpr_dyn(vdst_off, val.cast(dtypes.uint32)))
-      elif dest.startswith('VGPR['): stores.append(self.vgpr.index(val[0].cast(dtypes.int)).store(val[1].cast(dtypes.uint32)))
+      elif dest.startswith('VGPR['): stores.append(self.vgpr.index(val[0]).store(val[1].cast(dtypes.uint32)))
    return UOp.sink(*stores, *self.inc_pc())

  def compile_vop_pcode(self, op, srcs: dict[str, UOp | int], lane: UOp, vdst_reg: UOp, exec_mask: UOp,
@@ -711,11 +711,11 @@ class _Ctx:
        # VGPR bit-slice: (vgpr_idx, rhs_val, hi_bit, lo_bit) - hi/lo are UOp constants
        hi_bit, lo_bit = int(val[2].arg), int(val[3].arg)
        width = hi_bit - lo_bit + 1
-        old = self.vgpr.index(val[0].cast(dtypes.int), ptr=True).load()
+        old = self.vgpr.index(val[0], ptr=True).load()
        new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32)
        active = _lane_active(exec_mask, lane)
        if len(val) > 4: active = active & _to_bool(val[4])
-        raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].cast(dtypes.int).valid(active)).store(new_val)))
+        raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].valid(active)).store(new_val)))
        continue
      if 'D0' in dest and '[laneId]' in dest:
        old_vcc = self.rmask(_c(VCC_LO.offset))
@@ -834,7 +834,7 @@ def _compile_smem(inst: ir3.SMEM | ir4.SMEM, ctx: _Ctx) -> UOp:
  nval = int(part.removeprefix('DWORD').removeprefix('X') or '1') if 'DWORD' in part else int(part[1:]) / 32 * (-1 if part[0] == 'I' else 1)
  ndwords = max(1, int(abs(nval)))
  dword_base = addr >> UOp.const(dtypes.uint64, 2)
-  vals = [ctx.vmem.index((dword_base + UOp.const(dtypes.uint64, i)).cast(dtypes.int)) for i in range(ndwords)]
+  vals = [ctx.vmem.index(dword_base + UOp.const(dtypes.uint64, i)) for i in range(ndwords)]
  if abs(nval) < 1:
    nbits = int(abs(nval) * 32)
    byte_off = (addr & UOp.const(dtypes.uint64, 3)).cast(dtypes.uint32) * UOp.const(dtypes.uint32, 8)
@@ -1819,7 +1819,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
    srcs = {'ADDR': addr_reg, 'DATA0': vdata_reg, 'VDST': vdst_reg, 'OFFSET': offset,
            'EXEC': exec_mask.cast(dtypes.uint64), '_vgpr': ctx.vgpr, '_wave_size': ctx.wave_size}
    _, assigns = parse_pcode(pcode, srcs)
-    stores = [ctx.vgpr.index(val[0].cast(dtypes.int)).store(val[1].cast(dtypes.uint32)) for dest, val in assigns if dest.startswith('VGPR[')]
+    stores = [ctx.vgpr.index(val[0]).store(val[1].cast(dtypes.uint32)) for dest, val in assigns if dest.startswith('VGPR[')]
    return UOp.sink(*stores, *ctx.inc_pc())

  def make_addr(lane: UOp) -> UOp:
@@ -1853,7 +1853,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
  def wmem(addr: UOp, val: UOp, active: UOp, data_bits: int = 32) -> UOp:
    if data_bits < 32:
      # Sub-dword LDS write: read-modify-write within the uint32 slot
-      word_addr = (addr >> addr_shift).cast(dtypes.int)
+      word_addr = addr >> addr_shift
      idx = mem.index(word_addr.valid(active))
      byte_pos = addr.cast(dtypes.uint32) & _c(3)
      byte_shift = byte_pos * _c(8)
@@ -1861,7 +1861,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
      mask = size_mask << byte_shift
      new_word = (idx & (mask ^ _c(0xFFFFFFFF))) | ((val.cast(dtypes.uint32) & size_mask) << byte_shift)
      return idx.store(active.where(new_word, idx))
-    idx = mem.index((addr >> addr_shift).cast(dtypes.int))
+    idx = mem.index(addr >> addr_shift)
    return idx.store(active.where(val, idx.load()))

  def make_srcs(lane: UOp) -> dict:
@@ -2011,7 +2011,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
    for i in range(n_dwords):
      word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
      val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64), ptr=True).load(), _c(0))
-      lds_idx = ((lds_addr + _c(i * 4)) >> _c(2)).cast(dtypes.int)
+      lds_idx = (lds_addr + _c(i * 4)) >> _c(2)
      lds_slot = ctx.lds.index(lds_idx.valid(active))
      stores.append(lds_slot.store(active.where(val, lds_slot)))
  elif is_store:
@@ -2,7 +2,7 @@
 from typing import Any, Callable
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.uop.decompositions import f2f
+from tinygrad.codegen.decomp.dtype import f2f

 # Type alias for vars dict: stores UOps and tuples for lambda definitions
 VarVal = UOp | tuple[str, list[str], str]
@@ -831,7 +831,7 @@ class Parser:
    def mindex(idx:UOp, ptr=False): return mem.index(idx.valid(active) if active is not None else idx, ptr=ptr)
    byte_mem = mem.dtype.base == dtypes.uint8
    if byte_mem:
-      idx = addr.cast(dtypes.int)
+      idx = addr
      if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
        val = _u32(0).cast(dtypes.uint64)
        for i in range(8): val = val | (mindex(idx + _const(dtypes.int, i), ptr=True).load().cast(dtypes.uint64) << _u64(i * 8))
@@ -845,10 +845,10 @@ class Parser:
        val = _u32(0)
        for i in range(4): val = val | (mindex(idx + _const(dtypes.int, i), ptr=True).load().cast(dtypes.uint32) << _u32(i * 8))
    else:
-      idx = (addr >> _const(addr.dtype, 2)).cast(dtypes.int)
+      idx = addr >> _const(addr.dtype, 2)
      val = mindex(idx)
      if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
-        idx2 = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int)
+        idx2 = (addr + _const(adt, 4)) >> _const(adt, 2)
        val = val.cast(dtypes.uint64) | (mindex(idx2).cast(dtypes.uint64) << _u64(32))
      elif dt in (dtypes.uint8, dtypes.int8): val = (val >> ((addr & _const(adt, 3)).cast(dtypes.uint32) * _u32(8))) & _u32(0xFF)
      elif dt in (dtypes.uint16, dtypes.int16):
@@ -21,7 +21,7 @@ class TestDisassembly(unittest.TestCase):
    c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
    s = c.schedule_linear().src[-1]
    p = to_program(s.src[0], Device[Device.DEFAULT].renderer)
-    lib = Device[Device.DEFAULT].compiler.compile(p.src[3].arg)
+    lib = Device[Device.DEFAULT].compiler.compile(p.src[2].arg)
    out = io.StringIO()
    with redirect_stdout(out): Device[Device.DEFAULT].compiler.disassemble(lib)
    assert "fcvt" not in out.getvalue()
@@ -15,7 +15,7 @@ class TestLinearizerRewrite(unittest.TestCase):
      opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4))
      ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
      prg = to_program(ast, Device["CPU"].renderer)
-      print(prg.src[3].arg)
+      print(prg.src[2].arg)

  def test_arange(self):
    out = Tensor.arange(32).clone("NULL")
@@ -25,7 +25,7 @@ class TestLinearizerRewrite(unittest.TestCase):
      opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
      ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
      prg = to_program(ast, Device["CPU"].renderer)
-      print(prg.src[3].arg)
+      print(prg.src[2].arg)

  def test_kernel_info(self):
    out = Tensor.arange(4).clone("NULL")
@@ -1842,7 +1842,7 @@ class TestFusionOp(unittest.TestCase):
    linear = a.schedule_linear()
    prg = to_program(linear.src[-1].src[0], renderer=Device[Device.DEFAULT].renderer)
    self.assertLess(time.perf_counter()-st, 2.0)
-    assert len(prg.src[3].arg.splitlines()) < 250
+    assert len(prg.src[2].arg.splitlines()) < 250

  def test_recursive_add_cmp(self):
    st = time.perf_counter()
@@ -1,6 +1,6 @@
 import unittest, itertools

-from tinygrad.codegen.late.devectorizer import load_store_indexing
+from tinygrad.codegen.late.devectorizer import indexing_simplify
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
@@ -11,46 +11,22 @@ from test.null.test_uop_symbolic import check_uop_against_string
 # symbolic-only idx + valid simplification (no late lowering of FLOORDIV/FLOORMOD)
 def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load, name="simplify_valid_idx")
 # image-aware idx + valid simplification: adds the codegen-layer matcher that drops provably in-bounds gates
-def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+load_store_indexing, name="simplify_image_idx")
+def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")

 def get_gated_load_uop(valid:UOp, idx:UOp):
  return UOp(Ops.LOAD, dtypes.float, (
    UOp.param(0, dtypes.float.ptr()).index(idx.valid(valid), ptr=True),
-    UOp.const(dtypes.float, 0.0)
  ))

 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
  return UOp(Ops.LOAD, dtypes.float.vec(4), (
    UOp.param(0, dtypes.imagef(image_shape)).index(idx[1].valid(valid), idx[0].valid(valid), ptr=True),
-    UOp(Ops.STACK, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
  ))

 def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.weakint, (UOp.const(dtypes.weakint, nmax),), expr)
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
 def Range(n, nmax): return UOp.range(nmax, n)

-class TestHelpers(unittest.TestCase):
-  def test_is_increasing(self):
-    idx1 = Special("idx1", 32)
-    idx2 = Special("idx2", 64)
-    ridx0 = Variable("ridx0", 0, 5)
-    ridx1 = Variable("ridx1", 0, 2)
-    ridx2 = Variable("ridx2", 0, 2)
-    # (ridx0+(idx1*48)+(ridx2*6)+(-6)),((idx2*2)+ridx1+(-1)))
-    f0 = ((idx1*24)+(ridx2*3)+ridx0+765)%768
-    f1 = ridx0+(idx1*48)+(ridx2*6)+(-6)
-    f2 = (idx2*2)+ridx1+((idx1+((ridx2+7)//8)+31)//32)+(-2)
-    f3 = (idx2*2)+ridx1+(-1)
-
-    self.assertFalse(f0.is_increasing())
-    self.assertTrue(f1.is_increasing())
-    self.assertTrue(f2.is_increasing())
-    self.assertTrue(f3.is_increasing())
-
-    rng = UOp.range(5, 2)
-    self.assertTrue(rng.is_increasing())
-    self.assertTrue((rng+2).is_increasing())
-
 class TestValidIdxSimplification(unittest.TestCase):
  def check(self, load, sidx, svalid, extra=()):
    load = simplify_valid_idx(UOp.sink(load, *extra)).src[0]
@@ -506,10 +482,20 @@ class TestImageSimplification(unittest.TestCase):
      self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<11))",
                       "(lidx2+gidx0*4+lidx1*256+(lidx0*1024+r0*1024)+-3264)", "0")

+  def test_drop_non_monotonic_window(self):
+    # two-sided window valid (645 <= gidx0 < 653) on a non-monotonic index (lane split via %4 and //4):
+    # gidx0 outside the window pushes idx_x out of the (1, 48) image, so the gate is dropped
+    gidx0 = Special("gidx0", 1064)
+    r12 = Range(12, 3)
+    valid = ((gidx0 < 645).ne(True)) & (gidx0 < 653)
+    idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(dtypes.weakint, 0))
+    load = get_load_image_uop((1, 48, 4), valid, idx)
+    self.check(load, None, "(r12*4+(gidx0+3)%4+(gidx0+3)//4*24+-3888)", "0")
+
 class TestDropTrueGate(unittest.TestCase):
  def test_drop_true_gate_on_index(self):
    # test that INDEX with a constant True valid gets simplified to drop the valid
-    from tinygrad.codegen.late.devectorizer import load_store_indexing
+    from tinygrad.codegen.late.devectorizer import indexing_simplify
    from tinygrad.uop.ops import graph_rewrite
    from tinygrad.uop.symbolic import sym
    buf = UOp.param(0, dtypes.int.ptr())
@@ -517,7 +503,7 @@ class TestDropTrueGate(unittest.TestCase):
    true_gate = UOp.const(dtypes.bool, True)
    index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx.valid(true_gate)))
    # apply the optimization
-    result = graph_rewrite(index_with_gate, sym+load_store_indexing)
+    result = graph_rewrite(index_with_gate, sym+indexing_simplify)
    # the True valid should be dropped (INDEX should only have 2 sources)
    self.assertEqual(len(result.src), 2, "True valid should be dropped from INDEX")

@@ -1,7 +1,7 @@
 # tensor tests that pass on NULL backend (no copyout needed)
 import numpy as np
 import unittest
-from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
@@ -12,15 +12,6 @@ x_init = np.random.randn(1,3).astype(np.float32)
 W_init = np.random.randn(3,3).astype(np.float32)
 m_init = np.random.randn(1,3).astype(np.float32)

-class TestTrainMode(unittest.TestCase):
-  def test_train_mode(self):
-    assert not Tensor.training
-    @Context(TRAINING=1)
-    def f():
-      assert Tensor.training
-    f()
-    assert not Tensor.training
-
 class TestInferenceMode(unittest.TestCase):
  def test_inference(self):
    x = Tensor(x_init)
@@ -66,7 +57,7 @@ class TestIdxUpcast(unittest.TestCase):
      if ast.op is Ops.SINK:
        renderer = Device[si.src[1].buffer.device].renderer
        prg = to_program(ast, renderer)
-        return tuple(prg.src[2].src)
+        return tuple(prg.src[1].src)

  def _assert(self, dtype: DType, a: Tensor):
    uops = self._schedule_render(a)
@@ -1,9 +1,10 @@
 import math, unittest
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite
+from dataclasses import replace
+from tinygrad import Tensor, dtypes, Context
+from tinygrad.uop.ops import ParamArg, UOp, UPat, Ops, PatternMatcher, graph_rewrite

 _strip_unique_pm = PatternMatcher([
-  (UPat((Ops.UNIQUE, Ops.LUNIQUE), name="u"), lambda u: u.replace(arg=0) if u.arg != 0 else None),
+  (UPat(Ops.BUFFER, name="b"), lambda b: b.replace(arg=replace(b.arg, slot=0)) if isinstance(b.arg, ParamArg) and b.arg.slot != 0 else None),
 ])
 def _strip_unique(u: UOp) -> UOp: return graph_rewrite(u, _strip_unique_pm)

@@ -394,6 +395,24 @@ class TestTensorUOpConv2d(unittest.TestCase):
  def test_conv2d_3d(self):
    w = _t(1, 1, 2, 2, 2).float()
    _check(self, _t(1, 1, 3, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
+  def test_conv2d_winograd(self):
+    w, a = _t(2, 2, 3, 3).float(), _t(1, 2, 6, 6).float()
+    with Context(WINO=0): direct = a.conv2d(w).uop
+    with Context(WINO=1):
+      self.assertIsNot(a.conv2d(w).uop, direct)
+      _check(self, a, lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
+  def test_conv2d_image(self):
+    w, a = _t(4, 4, 3, 3).float(), _t(1, 4, 8, 8).float()
+    with Context(IMAGE=0): direct = a.conv2d(w).uop
+    with Context(IMAGE=1):
+      self.assertIsNot(a.conv2d(w).uop, direct)
+      _check(self, a, lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
+  def test_dot_image(self):
+    y, a = _t(4, 3).float(), _t(2, 4).float()
+    with Context(IMAGE=0): direct = a.dot(y).uop
+    with Context(IMAGE=1):
+      self.assertIsNot(a.dot(y).uop, direct)
+      _check(self, a, lambda x: x.dot(y if isinstance(x, Tensor) else y.uop))
  def test_conv_transpose2d_basic(self):
    w = _t(1, 1, 2, 2).float()
    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv_transpose2d(w if isinstance(x, Tensor) else w.uop))
@@ -2,8 +2,7 @@ import unittest, math
 import numpy as np
 from tinygrad import dtypes
 from tinygrad.uop.ops import UOp
-from tinygrad.uop.decompositions import TRANSCENDENTAL_DTYPES, payne_hanek_reduction, cody_waite_reduction
-from tinygrad.uop.decompositions import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
+from tinygrad.codegen.decomp.transcendental import payne_hanek_reduction, cody_waite_reduction, frexp, rintk, pow2if
 from test.helpers import eval_uop

 class TestTranscendentalFunctions(unittest.TestCase):
@@ -71,39 +70,5 @@ class TestTranscendentalFunctions(unittest.TestCase):
    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -10), dtypes.float)), 2**-10)
    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -63), dtypes.float)), 2**-63)

-class TestTranscendentalVectorizedFunctions(unittest.TestCase):
-  # given a scalar and vectorized input, check that the fxn outputs have the same
-  # scalar_dtypes, args, ops, and vcount (only for vectorized input)
-
-  def _check_uop_vcount(self, u:tuple|UOp, vcount:int):
-    # check all UOps in u are vectorized with vcount
-    if isinstance(u, UOp):
-      assert u.dtype.vcount == vcount, f'expected {vcount=} but got {u.dtype.vcount=} for UOp\n{u=}'
-    [self._check_uop_vcount(x, vcount) for x in (u if isinstance(u, tuple) else u.src)]
-
-  def _check_uops_match(self, u1:tuple|UOp, u2:tuple|UOp):
-    # check all UOps in u1, u2 have the same scalar_dtype, args, ops
-    if isinstance(u1, UOp) and isinstance(u2, UOp):
-      assert u1.dtype.scalar() == u2.dtype.scalar(), f'expected {u1.dtype.scalar()=} but got {u2.dtype.scalar()=} for UOps\n{u1=}\n{u2}'
-      assert u1.arg == u2.arg or (math.isnan(u1.arg) and math.isnan(u2.arg)), f'expected {u1.arg=} but got {u2.arg=} for UOps\n{u1=}\n{u2}'
-      assert u1.op == u2.op, f'expected {u1.op=} but got {u2.op=} for UOps\n{u1=}\n{u2}'
-    [self._check_uops_match(x1, x2) for x1, x2 in zip((u1 if isinstance(u1, tuple) else u1.src), (u2 if isinstance(u2, tuple) else u2.src))]
-
-  def _test_vectorized(self, fxn, scalar_dtypes=TRANSCENDENTAL_DTYPES, vals=[-2,1.3,194], vcounts=[1,4,19]):
-    for scalar_dtype in scalar_dtypes:
-      for val in vals:
-        for vcount in vcounts:
-          in_scalar, in_vec = UOp.const(scalar_dtype, val), UOp.const(scalar_dtype.vec(vcount), val)
-          out_scalar, out_vec = fxn(in_scalar), fxn(in_vec)
-          self._check_uops_match(out_scalar, out_vec)
-          self._check_uop_vcount(out_vec, vcount)
-
-  def test_xpow(self): return self._test_vectorized(lambda x: xpow(x, x))
-  def test_xexp2(self): return self._test_vectorized(xexp2)
-  def test_xlog2(self): return self._test_vectorized(xlog2)
-  def test_payne_hanek_reduction(self): return self._test_vectorized(payne_hanek_reduction)
-  def test_cody_waite_reduction(self): return self._test_vectorized(cody_waite_reduction)
-  def test_trig_poly(self): return self._test_vectorized(lambda x: trig_poly(x, [0.0], [1.0]))
-
 if __name__ == '__main__':
  unittest.main()
@@ -35,19 +35,17 @@ class TestGraphRewriteConst(unittest.TestCase):
    self.assertEqual(ret.arg, 4)

  def test_add_const(self):
-    v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
-    v2 = UOp.const(dtypes.int.vec(3), (5,6,7))
+    v1 = UOp.const(dtypes.int, (0,1,2))
+    v2 = UOp.const(dtypes.int, (5,6,7))
    ret = graph_rewrite(v1+v2, sym)
    self.assertEqual(ret.op, Ops.STACK)
-    self.assertEqual(ret.dtype, dtypes.int.vec(3))
    self.assertEqual(const_values(ret), (5,7,9))

  def test_add_const_lose_v(self):
-    v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
-    v2 = UOp.const(dtypes.int.vec(3), (2,1,0))
+    v1 = UOp.const(dtypes.int, (0,1,2))
+    v2 = UOp.const(dtypes.int, (2,1,0))
    ret = graph_rewrite(v1+v2, sym)
    self.assertEqual(ret.op, Ops.STACK)
-    self.assertEqual(ret.dtype, dtypes.int.vec(3))
    self.assertEqual(const_values(ret), (2,2,2))

 def xfail_broken_const_wraparound(fn):
@@ -771,45 +769,6 @@ class TestMovementOps(unittest.TestCase):
    self.assertEqual(result.op, Ops.INDEX)
    self.assertEqual(result.src[0].op, Ops.RESHAPE)

-class TestLoadStoreFolding(unittest.TestCase):
-  def test_gated_load_gep_preserves_alt(self):
-    """Test that LOAD(GEP, alt) preserves alt value after rewrite"""
-    from tinygrad.codegen.late.devectorizer import load_store_folding
-    buf = UOp.param(0, dtypes.float.vec(4).ptr())
-    idx = UOp.const(dtypes.int, 0)
-    gate = UOp.const(dtypes.bool, True)
-    gated_index = buf.index(idx.valid(gate))
-    gep = gated_index.gep(0)
-    alt = UOp.const(dtypes.float, 42.0)
-    gated_load = gep.load(alt)
-    self.assertEqual(len(gated_load.src), 2)  # GEP + alt
-    result = graph_rewrite(gated_load, load_store_folding, name='test')
-    # After rewrite, should still have alt value preserved
-    self.assertEqual(result.op, Ops.GEP)
-    inner_load = result.src[0]
-    self.assertEqual(inner_load.op, Ops.LOAD)
-    self.assertEqual(len(inner_load.src), 2)  # INDEX + alt
-
-  def test_gated_load_ptrcat_preserves_alt(self):
-    """Test that LOAD(PTRCAT, alt) preserves alt value after rewrite"""
-    from tinygrad.codegen.late.devectorizer import load_store_folding
-    buf1 = UOp.param(0, dtypes.float.ptr())
-    buf2 = UOp.param(1, dtypes.float.ptr())
-    idx = UOp.const(dtypes.int, 0)
-    idx1 = buf1.index(idx)
-    idx2 = buf2.index(idx)
-    ptrcat = UOp(Ops.PTRCAT, dtypes.float.ptr().vec(2), (idx1, idx2))
-    alt = UOp.const(dtypes.float.vec(2), 42.0)
-    gated_load = ptrcat.load(alt)
-    self.assertEqual(len(gated_load.src), 2)  # PTRCAT + alt
-    result = graph_rewrite(gated_load, load_store_folding, name='test')
-    # After rewrite, should be CAT of LOADs, each preserving alt
-    self.assertEqual(result.op, Ops.VCAT)
-    for inner_load in result.src:
-      self.assertEqual(inner_load.op, Ops.LOAD)
-      self.assertEqual(len(inner_load.src), 2)  # INDEX + alt
-    self.assertEqual(inner_load.src[1].arg, 42.0)  # alt value preserved
-
 class TestConstBufferize(unittest.TestCase):
  def test_const_bufferize_with_ranges(self):
    """Test that CONST.BUFFERIZE with ranges is folded correctly.
@@ -5,7 +5,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, ConstFloat  # noqa: F401
 from tinygrad.device import Device
-from tinygrad.uop.ops import Ops, UOp, UPat, exec_alu
+from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, exec_alu  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
 from tinygrad.uop.spec import spec_shared
 from tinygrad.uop.symbolic import sym
 from test.helpers import eval_uop, to_uops_list
@@ -185,7 +185,7 @@ class TestStatsOptimized(unittest.TestCase):
                      renderer=Device[Device.DEFAULT].renderer)
    except KernelOptError:
      raise unittest.SkipTest("no tensor cores")
-    print(p.src[3].arg)
+    print(p.src[2].arg)
    self.check_gemm(p)

  # this is a good lesson about why UPCASTing is a good idea
@@ -7,7 +7,7 @@ from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatch
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.helpers import colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
-from tinygrad.helpers import cpu_profile, ProfilePointEvent, unwrap, VIZ
+from tinygrad.helpers import cpu_profile, ProfilePointEvent, unwrap, VIZ, BEAM
 from tinygrad.device import Buffer

 from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace
@@ -215,6 +215,23 @@ class TestViz(unittest.TestCase):
    nop = UOp(Ops.NOOP, arg="infinite loop in fixed_point_rewrite")
    self.assertEqual(graphs[2], uop_to_json(VizData(), nop)[id(nop)])

+  def test_walk_rewrite(self):
+    from tinygrad.uop.ops import _substitute
+    with save_viz() as viz:
+      a = UOp.variable("a", 0, 10)
+      graph_rewrite(a + 4, TrackedPatternMatcher(_substitute.patterns), {a:a+1}, walk=True)
+    list(viz.get_details(0, 0))
+
+  def test_enter_calls_rewrite(self):
+    pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4))])
+    with save_viz() as viz:
+      inner = UOp.const(dtypes.int, 3)
+      func = UOp(Ops.FUNCTION, src=(UOp(Ops.SINK, src=(inner,)),))
+      call = UOp(Ops.CALL, src=(func,))
+      graph_rewrite(call, TrackedPatternMatcher(pm.patterns), enter_calls=True)
+    details = list(viz.get_details(0, 0))
+    self.assertTrue(details[-1]["change"], "viz replay should detect change inside CALL")
+
  def test_const_node_visibility(self):
    with save_viz() as viz:
      a = UOp.variable("a", 0, 10, dtype=dtypes.int)
@@ -341,41 +358,28 @@ class TestVizGC(unittest.TestCase):
 from tinygrad import Tensor, Device, TinyJit, Variable, function

 class TestVizIntegration(unittest.TestCase):
-  # codegen supports rendering of code blocks
-  def test_codegen_tracing(self):
-    with save_viz() as viz:
-      ast = (Tensor.empty(4)+Tensor.empty(4)).schedule_linear().src[0].src[0]
-      prg = do_to_program(ast, Device[Device.DEFAULT].renderer)
-    lst = viz.list_items()
-    self.assertEqual(len(lst), 3)
-    self.assertEqual(lst[0]["name"], "Callify 1 Buffer n1")
-    self.assertEqual(lst[1]["name"], "Schedule 1 Kernel n1")
-    self.assertEqual(lst[2]["name"], prg.arg.name)
-    input_ast = next(viz.get_details(2, 0))["graph"].values()
-    for u in input_ast:
-      if u["label"].startswith("PARAM\n"): self.assertEqual(u["addrspace"], addrspace_colors[AddrSpace.GLOBAL])
-
-  # schedule graph CALL nodes have a link to jump to codegen
  def test_link_sched_codegen(self):
+    c1 = Tensor.empty(4, device="NULL")
+    c2 = Tensor.empty(8, device="NULL")
+    # uniquely named A = B + 1 kernel
+    kernel_name = f"custom_add1_link_sched_codegen_{BEAM.value}"
+    def custom_add1(A:UOp, B:UOp): return A[0].store(B[0]+1).sink(arg=KernelInfo(kernel_name))
    with save_viz() as viz:
-      c1 = Tensor.empty(4, device="NULL").add(1)
-      c2 = Tensor.empty(8, device="NULL").add(1)
-      with Context(SCACHE=0):
-        sched = c1.schedule_linear(c2)
-      from tinygrad.engine.realize import compile_linear
-      sched = compile_linear(sched)
-      with Context(NO_COLOR=0):
-        prgs = [do_to_program(si.src[0], Device[c1.device].renderer).arg.name for si in sched.src]
+      c1 = Tensor.custom_kernel(c1, c2, fxn=custom_add1)[0]
+      c1.realize()
    lst = viz.list_items()
+    # schedule graph CALL nodes have a link to jump to codegen
    sched_idx = next(i for i,l in enumerate(lst) if l["name"].startswith("Schedule"))
    viz_kernel = next(i for i,s in enumerate(lst[sched_idx]["steps"]) if s["name"] == "View Kernel Graph")
-    with Context(NO_COLOR=1):
-      graph = next(viz.get_details(sched_idx, viz_kernel))["graph"]
+    graph = next(viz.get_details(sched_idx, viz_kernel))["graph"]
    call_nodes = [n for n in graph.values() if n["label"].startswith("CALL")]
    for i,n in enumerate(call_nodes):
      assert n["ref"] is not None
-      self.assertEqual(lst[n["ref"]]["name"], prgs[i])
-      assert ansistrip(prgs[i]) in n["label"], f"CALL must contain kernel name, got {n['label']}"
+      self.assertEqual(lst[n["ref"]]["name"], kernel_name)
+      assert kernel_name[i] in n["label"], f"CALL must contain kernel name, got {n['label']}"
+    # UOp addrspace is colored
+    for u in graph.values():
+      if u["label"].startswith("PARAM\n"): self.assertEqual(u["addrspace"], addrspace_colors[AddrSpace.GLOBAL])

  def test_link_sched_codegen_beam(self):
    with Context(BEAM=2):
@@ -493,6 +497,24 @@ class TestVizIntegration(unittest.TestCase):
    bin_render = get_render(viz.data, steps[bin_idx]["query"])["src"]
    self.assertIn(type(e.exception).__name__, bin_render)

+  def test_view_source_alt(self):
+    src = "void E_3(float* data0_3) {}"
+    binary = Device["CPU"].renderer.compiler.compile(src)
+    def custom_binary(X:UOp):
+      sink = UOp.sink(X, arg=KernelInfo("custom_binary"))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=sink.src+(sink,)), UOp(Ops.SOURCE, arg=src),
+                                   UOp(Ops.BINARY, arg=binary)))
+    x = Tensor.custom_kernel(Tensor.empty(1, device="CPU"), fxn=custom_binary)[0]
+    with save_viz() as viz:
+      x.realize()
+    lst = viz.list_items()
+    codegen_idx = len(lst)-1
+    steps = lst[codegen_idx]["steps"]
+    src_idx = next((i for i,s in enumerate(steps) if s["name"] == "View Source"), None)
+    assert src_idx is not None, "must have source rendering in list"
+    src_render = get_render(viz.data, steps[src_idx]["query"])["src"]
+    self.assertEqual(src, src_render)
+
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
 from tinygrad.viz.cli import decode_profile
@@ -804,7 +826,7 @@ class TestCfg(unittest.TestCase):
      lidx = UOp.special(1, "lidx0")
      gidx = UOp.special(1, "gidx0")
      sink = UOp.sink(out.base, lidx, gidx, arg=KernelInfo(name=name))
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="NULL"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
    with save_viz() as viz:
      with Context(DEV=f"NULL::{self.arch}"):
        out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
@@ -1089,8 +1111,10 @@ class TestCLI(unittest.TestCase):
    with write_files(viz) as files, Context(NO_COLOR=1):
      flat = run_cli(*files, "-s", "NULL", "--interval", "interval_start", "interval_end")
      aggregate = run_cli(*files, "-s", "NULL", "--interval", "interval_start", "interval_end", "-t")
+      final = run_cli(*files, "-s", "NULL", "--interval", "interval_end", "-t")
    self.assertEqual([s["name"] for s in flat], ["interval_start", "target_1", "target_2", "interval_end"])
    self.assertEqual(sorted(s["name"] for s in aggregate), ["target_1", "target_2"])
+    assert all(s["name"].startswith("post_") for s in final), f"post_* kernels must be present in final, got {final}"

 if __name__ == "__main__":
  unittest.main()
@@ -27,7 +27,7 @@ class TestFloat4(unittest.TestCase):
    opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
    program = to_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)

-    assert TestFloat4.count_float4(tuple(program.src[2].src)) == (2, 1)
+    assert TestFloat4.count_float4(tuple(program.src[1].src)) == (2, 1)

  def test_float4_multidim(self):
    a = Tensor.empty(2, 8).realize()
@@ -36,7 +36,7 @@ class TestFloat4(unittest.TestCase):

    s = c.schedule_linear().src[0]
    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+                       renderer=Device[Device.DEFAULT].renderer).src[1].src)
    assert TestFloat4.count_float4(uops) == (4, 2)

  def test_float4_unaligned_load(self):
@@ -49,7 +49,7 @@ class TestFloat4(unittest.TestCase):
    opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
    program = to_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)

-    assert TestFloat4.count_float4(tuple(program.src[2].src)) == (0, 1)
+    assert TestFloat4.count_float4(tuple(program.src[1].src)) == (0, 1)

  def test_float4_multidim_unaligned_load(self):
    a = Tensor.empty(2, 9).realize().shrink(((0, 2), (1, 9),))
@@ -58,7 +58,7 @@ class TestFloat4(unittest.TestCase):

    s = c.schedule_linear().src[0]
    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+                       renderer=Device[Device.DEFAULT].renderer).src[1].src)

    assert TestFloat4.count_float4(uops) == (0, 2)

@@ -70,7 +70,7 @@ class TestFloat4(unittest.TestCase):
    # float4 should be emitted (the reduce axis of size 4 is the float4 axis here)

    s = c.schedule_linear().src[0]
-    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UNROLL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UNROLL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)

    assert TestFloat4.count_float4(uops) == (0, 0)

@@ -85,7 +85,7 @@ class TestFloat4(unittest.TestCase):

    s = c.schedule_linear().src[0]
    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+                       renderer=Device[Device.DEFAULT].renderer).src[1].src)

    assert TestFloat4.count_float4(uops) in {(0,1), (1,1)}

@@ -98,7 +98,7 @@ class TestFloat4(unittest.TestCase):
    # since the top axis is not contiguous.

    s = c.schedule_linear().src[0]
-    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)

    assert TestFloat4.count_float4(uops) == (0, 1)

@@ -110,7 +110,7 @@ class TestFloat4(unittest.TestCase):
    # should float4 b but not a

    s = c.schedule_linear().src[0]
-    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)

    assert TestFloat4.count_float4(uops) == (1, 1)

@@ -33,7 +33,7 @@ def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, d

  if ensure_triggered:
    program = to_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer)
-    wmmas = len([uop for uop in tuple(program.src[2].src) if uop.op is Ops.WMMA])
+    wmmas = len([uop for uop in tuple(program.src[1].src) if uop.op is Ops.WMMA])
    tcs = len([x for x in program.src[0].arg.applied_opts if x.op is OptOps.TC])
    assert wmmas > 0, "tensor core not triggered"
    assert tcs == 1, "tensor core opt not included"
@@ -52,7 +52,7 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
  opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))]
  ast = replace_opts(realized_ast, opts)
  pu = to_program(ast, Device[Device.DEFAULT].renderer)
-  if use_tensor_cores == 1: assert len([uop for uop in pu.src[2].src if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
+  if use_tensor_cores == 1: assert len([uop for uop in pu.src[1].src if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
  assert len([x for x in pu.src[0].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
  run_program(ast, bufs)
  if dtype_in == dtypes.half: tc_atol, tc_rtol = 1e-2, 1e-3
@@ -80,13 +80,13 @@ class TestTensorCores(unittest.TestCase):
      prg = to_program(replace_opts(r.schedule_linear().src[-1].src[0],
                        [Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))]), Device[Device.DEFAULT].renderer)
      if Device.DEFAULT == "CPU" and DEV.renderer == "LLVM":
-        assert "0x201000" in prg.src[3].arg
+        assert "0x201000" in prg.src[2].arg
      elif Device.DEFAULT == "AMD" and DEV.renderer == "LLVM":
-        assert "@llvm.amdgcn.wmma" in prg.src[3].arg
+        assert "@llvm.amdgcn.wmma" in prg.src[2].arg
      elif Device[Device.DEFAULT].renderer.suffix == "PTX":
-        assert "mma.sync.aligned" in prg.src[3].arg
+        assert "mma.sync.aligned" in prg.src[2].arg
      else:
-        assert "__WMMA_" in prg.src[3].arg
+        assert "__WMMA_" in prg.src[2].arg

  @Context(ALLOW_TF32=1)
  @unittest.skipIf((Device.DEFAULT == "AMD") or (Device.DEFAULT == "PYTHON" and Device.default.renderer.target.device == "AMD"), "broken for AMD")
@@ -144,7 +144,7 @@ class TestTensorCores(unittest.TestCase):

        ast = replace_opts(realized_ast, [Opt(OptOps.TC, axis, (-1, 2, 1))])
        program = to_program(ast, Device[Device.DEFAULT].renderer)
-        assert len([uop for uop in tuple(program.src[2].src) if uop.op is Ops.WMMA]) > 0, "tensor core not triggered"
+        assert len([uop for uop in tuple(program.src[1].src) if uop.op is Ops.WMMA]) > 0, "tensor core not triggered"
        assert len([x for x in program.src[0].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"

        # TODO: support this even if numpy doesn't
@@ -166,7 +166,7 @@ class TestTensorCores(unittest.TestCase):
    r = x.matmul(y, dtype=tc.dtype_out)
    opts = [Opt(OptOps.UNROLL, 0, 4)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[2].src):
+    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
      if u.op is Ops.WMMA:
        assert u.src[-1].src[0].op != Ops.STORE

@@ -180,7 +180,7 @@ class TestTensorCores(unittest.TestCase):
    r = x.matmul(y, dtype=tc.dtype_out)
    opts = [Opt(OptOps.UNROLL, 0, 4)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[2].src):
+    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
      if u.op is Ops.WMMA:
        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
        assert u.src[-1].src[0].op != Ops.STORE
@@ -196,7 +196,7 @@ class TestTensorCores(unittest.TestCase):
    r = x.matmul(y, dtype=tc.dtype_out).relu()
    opts = [Opt(OptOps.UNROLL, 0, 4)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[2].src):
+    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
      if u.op is Ops.WMMA:
        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
        assert u.src[-1].src[0].op != Ops.STORE
@@ -1,7 +1,7 @@
 # basic self-contained tests of the external functionality of tinygrad
 import unittest, random
 from tinygrad import Tensor, Context, Variable, TinyJit, dtypes, Device, nn
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, OSX

 class TestTiny(unittest.TestCase):

@@ -89,6 +89,7 @@ class TestTiny(unittest.TestCase):

  # *** BEAM (for Kernel speed) ***

+  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and OSX, "WEBGPU's timestamp-query is unreliable on dawn's metal backend")
  def test_beam(self):
    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): self.test_plus()

@@ -4,7 +4,7 @@ import unittest
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import get_parameters
 from tinygrad.nn.optim import Adam, SGD
-from tinygrad.helpers import DEBUG
+from tinygrad.helpers import DEBUG, Context
 from extra.lr_scheduler import MultiStepLR, ReduceLROnPlateau, CosineAnnealingLR, OneCycleLR
 from extra.training import train, evaluate
 from extra.datasets import fetch_mnist
@@ -53,11 +53,7 @@ def get_lrs(optim, sched, epochs, steps=1, accs=None):
  return lrs

 class TestLrScheduler(unittest.TestCase):
-  def setUp(self):
-    self.old_training = Tensor.training
-    Tensor.training = True
-  def tearDown(self):
-    Tensor.training = self.old_training
+  def setUp(self): self.enterContext(Context(TRAINING=1))

  def _test_lr_scheduler(self, tinygrad_sched, torch_sched, epochs, opts, atol, rtol, adam=True):
    accs = opts.pop('accs', None)
@@ -17,7 +17,7 @@ class TestCPU(unittest.TestCase):
          if ren is LVPRenderer: continue # LVP does not play nice with cross compilation
          r = ren(Target(device="CPU", arch=arch))
          p = to_program(ast, r)
-          lib = r.compiler.compile(p.src[3].arg)
+          lib = r.compiler.compile(p.src[2].arg)
          out = io.StringIO()
          with redirect_stdout(out): r.compiler.disassemble(lib)
          self.assertEqual("vmov" in out.getvalue(), expect_vmov, out.getvalue())
@@ -557,7 +557,7 @@ class TestFunctionTuple(unittest.TestCase):
    lib = Device["CPU"].compiler.compile(src)
    def prog(C:UOp, A:UOp) -> UOp:
      sink = UOp.sink(C.base, A.base, arg=KernelInfo(name="k"))
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="CPU"), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
                                   UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)),
                 arg=ProgramInfo(name="k", global_size=(1, 1, 1), local_size=(1, 1, 1), globals=(0, 1)))

@@ -33,11 +33,11 @@ class TestKeccak(unittest.TestCase):
      out_shape = Tensor.randint(*s[i:], high=255, dtype=dtypes.uint8).keccak().shape
      self.assertTupleEqual(s[i:-1], out_shape[:-1])

-  @unittest.skipUnless(Device.DEFAULT=="METAL", "slow")
+  @slow
  def test_sha3_224(self): self._test_preset("sha3_224", [143, 144])
-  @unittest.skipUnless(Device.DEFAULT=="METAL", "slow")
+  @slow
  def test_sha3_256(self): self._test_preset("sha3_256", [135, 136])
-  @unittest.skipUnless(Device.DEFAULT=="METAL", "slow")
+  @slow
  def test_shake_128(self): self._test_preset("shake_128", [167, 168], lambda d: hashlib.shake_128(d).digest(16))

  def _test_preset(self, name: str, special_sizes: list[int], hasher: Callable[[bytes], bytes] | None = None):
@@ -21,8 +21,8 @@ class TestHCQUnit(unittest.TestCase):
    for _ in range(5): f(inp, inp_cpu)

    # construct minimal CALL UOps for supports_uop (graphs only see PROGRAMs after compile_linear)
-    gpu_call = UOp(Ops.PROGRAM, src=(UOp.sink(), UOp(Ops.DEVICE, arg=Device.DEFAULT))).call(UOp.new_buffer(Device.DEFAULT, 1, dtypes.float))
-    cpu_call = UOp(Ops.PROGRAM, src=(UOp.sink(), UOp(Ops.DEVICE, arg="CPU"))).call(UOp.new_buffer("CPU", 1, dtypes.float))
+    gpu_call = UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(UOp.new_buffer(Device.DEFAULT, 1, dtypes.float))
+    cpu_call = UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(UOp.new_buffer("CPU", 1, dtypes.float))
    gpu_devs = [d0]

    # local MMIO: GPU works alone and with CPU in batch (cpu_support=True)
@@ -1,6 +1,6 @@
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, AddrSpace, PtrDType, ImageDType
-from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, track_rewrites
+from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, ParamArg, graph_rewrite, track_rewrites
 from tinygrad.helpers import VIZ, pluralize, all_int

@dataclass
@@ -193,7 +193,8 @@ pm_finalize_call = PatternMatcher([

 pm_replace_buf = PatternMatcher([
  # replace BUFFER with PARAM for cache key normalization
-  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
+  (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
+   replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
  # replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input
  (UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer),
  # strip value from BIND for cache key normalization, so different values hit same cache
@@ -1,37 +1,30 @@
-from typing import cast
 from dataclasses import replace
-import itertools
+import itertools, functools
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
-from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic
+from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic, all_same, flatten
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
-from tinygrad.dtype import dtypes, PtrDType, ImageDType
+from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace

 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
-from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcendental_patterns, pm_dtype_decomps
+from tinygrad.codegen.decomp.dtype import pm_dtype_decomps
+from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns
+from tinygrad.codegen.decomp.transcendental import get_transcendental_patterns
 from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
-from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize_buf_and_index, devectorize_alu, pm_reduce, \
-  ReduceContext, correct_load_store, pm_render, pm_add_loads, pm_make_images
+from tinygrad.codegen.late.devectorizer import indexing_simplify, ReduceContext, pm_render, merge_reduce_ends
 from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
-from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar, pm_store_ranges
+from tinygrad.schedule.rangeify import pm_mops, pm_syntactic_sugar, pm_store_ranges
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
-
-pm_index_is_shrink = PatternMatcher([
-  # rewrite non-image INDEX to SHRINK
-  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"))).cast(name="x"), lambda buf,idx,x:
-    UOp(Ops.SHRINK, dtype=x.dtype.base, src=(buf, idx, UOp.const(dtypes.int, x.dtype.count))) \
-      if isinstance(buf.dtype, PtrDType) and x.dtype.count > 1 else None),
-  # rewrite GEP to INDEX
-  (UPat(Ops.GEP, name="x"), lambda x: x.replace(op=Ops.INDEX, src=x.src+(UOp.const(dtypes.int, x.arg if len(x.arg) > 1 else x.arg[0]),), arg=None)),
-])
+from tinygrad.codegen.late.coalese import memory_coalesing, pm_simplify_add_image
+from tinygrad.uop.ops import identity_element

 pm_remove_vec_dtypes = PatternMatcher([
  # rewrite PARAM to non pointer
@@ -41,6 +34,11 @@ pm_remove_vec_dtypes = PatternMatcher([
  # remove all vec dtypes
  (UPat(GroupOp.All-{Ops.PARAM, Ops.BUFFER}, name="x"),
   lambda x: x.replace(dtype=x.dtype.base.scalar().base)),
+  # rewrite GEP to INDEX
+  (UPat(Ops.GEP, name="x"), lambda x: x.replace(op=Ops.INDEX, src=x.src+(UOp.const(dtypes.int, x.arg if len(x.arg) > 1 else x.arg[0]),), arg=None)),
+  # don't stack PARAMs/BUFFERs on INDEX. TODO: this is an expander bug
+  (UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat((Ops.PARAM, Ops.BUFFER), name="b")),), allow_any_len=True, name="idx"),
+   lambda b,idx: idx.replace(src=(b,)+idx.src[1:])),
 ])+pm_clean_up_group_sink

 def do_number_param(ctx:list[int], x:UOp):
@@ -52,6 +50,84 @@ pm_number_params = PatternMatcher([
  (UPat(Ops.PARAM, name="x"), do_number_param),
 ])

+pm_no_weakints = PatternMatcher([
+  (UPat(GroupOp.All, dtype=dtypes.weakint, name="x"), lambda x: x.replace(dtype=dtypes.int))
+])
+
+def maybe_load(u:UOp): return u.load() if u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL, AddrSpace.REG) else u
+pm_move_regs = PatternMatcher([
+  # BITCAST?
+  (UPat(GroupOp.Elementwise|{Ops.REDUCE, Ops.STACK}, name="x"), lambda x: x.replace(src=tuple([maybe_load(u) for u in x.src]))),
+  (UPat(Ops.STORE, name="x"), lambda x: x.replace(src=(x.src[0], maybe_load(x.src[1]))+x.src[2:])),
+])
+
+def do_devectorize(b:UOp):
+  if b.shape == (): return None
+  # broadcasting needs to be already unpacked
+  if not all_same([x.shape for x in b.src]): return None
+  src = []
+  for idx in itertools.product(*[range(x) for x in b.shape]):
+    idx_c = [UOp.const(dtypes.weakint, i) for i in idx]
+    src.append(b.replace(src=tuple([x.index(*idx_c) for x in b.src])))
+  return UOp.vectorize(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
+
+devectorizer2 = pm_mops+PatternMatcher([
+  # unpack broadcasting
+  (UPat(GroupOp.Elementwise|{Ops.LOAD,Ops.STORE}, name="b"), do_devectorize),
+  # const INDEX into STACK is src
+  (UPat(Ops.INDEX, src=(UPat(Ops.STACK, name="a"), UPat.cvar("i"))), lambda a,i: a.src[i.arg]),
+  # stacked INDEX is many INDEX
+  (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.STACK, name="s"))),
+   lambda b,s: UOp.vectorize(*[b.index(u) for u in s.src])),
+  # INDEX into RESHAPE moves the RESHAPE
+  (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.RESHAPE, name="s"))),
+   lambda b,s: b.index(s.src[0]).reshape(s.shape)),
+  # RESHAPE a void is removed (hack for AFTER)
+  (UPat(Ops.RESHAPE, dtype=dtypes.void, name="x"), lambda x: x.src[0]),
+  # reshape of a single element shaped value to scalar is an index
+  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(UOp.const(dtypes.weakint, 0)) if x.marg == () and x.src[0].shape == (1,) else None),
+  # INDEX without src is nothing
+  (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
+  # RESHAPE+EXPAND -> STACK
+  (UPat(Ops.EXPAND, src=(UPat(Ops.RESHAPE, src=(UPat.var("x"), UPat())), UPat()), name="out"),
+   lambda x,out: UOp.vectorize(*([x]*out.max_numel())) if out.shape == (out.max_numel(),) else None),
+  # INDEX on INDEX is INDEX
+  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"),
+   lambda idx1, idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:])),
+])
+
+def reduce_ranges_to_acc(ctx:ReduceContext, r:UOp):
+  acc = UOp.placeholder_like(r, ctx.acc_num, AddrSpace.REG)
+  ctx.acc_num += 1
+  topo = r.src[0].toposort()
+  ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END])
+  input_ranges = tuple(x for x in topo if x.op is Ops.RANGE and x not in r.src[1:] and x not in ended_ranges)
+  acc_init = acc.after(*input_ranges).store(identity_element(r.arg[0], r.dtype.scalar()))
+  acc_initted = acc.after(acc_init, *r.src[1:])
+  inp = r.src[0].reduce(arg=r.arg) if r.arg[1] else r.src[0]
+  acc_out = acc_initted.store(acc_initted.alu(r.arg[0], inp)).end(*r.src[1:]).rtag("mergeable")
+  return acc.after(acc_out)
+
+def expand_horizontal_reduce(r:UOp):
+  axes = r.arg[1]
+  vals = [r.src[0].shrink(tuple((idx[axes.index(i)], idx[axes.index(i)]+1) if i in axes else None for i in range(r.src[0].ndim)))
+          for idx in itertools.product(*[range(r.src[0].max_shape[a]) for a in axes])]
+  return functools.reduce(lambda x,y: x.alu(r.arg[0], y), vals)
+
+pm_reduce_local = PatternMatcher([
+  (UPat(Ops.REDUCE, src=(UPat(), UPat()), allow_any_len=True, name="r"), reduce_ranges_to_acc),
+  (UPat(Ops.REDUCE, src=(UPat(),), name="r"), expand_horizontal_reduce),
+  (UPat(Ops.SINK, name="sink"), merge_reduce_ends),
+])+pm_clean_up_group_sink
+
+def add_local_buffer(ctx, x:UOp):
+  buf = UOp.placeholder(x.max_shape, x.dtype, slot=next(ctx), addrspace=x.arg.addrspace)
+  return buf.after(buf.index(*x.src[1:]).store(x.src[0]).end(*x.src[1:]).barrier())
+
+pm_add_local_buffers = PatternMatcher([
+  (UPat(Ops.STAGE, name="x"), add_local_buffer),
+])+pm_mops
+
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
  if DEBUG >= 5: print(pyrender(ast))
@@ -82,13 +158,24 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:

  # expand
  sink = graph_rewrite(sink, sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander")
+  #sink = graph_rewrite(sink, gep_pushing, name="gep pushing")
+
+  # this is new style (TODO: this should all be removed)
+  sink = graph_rewrite(sink, pm_render, name="pm_render gep/stack")
+  sink = graph_rewrite(sink, pm_remove_vec_dtypes, name="transform to new style")
+
+  # ****** new style ******

  # add locals
-  sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, ctx=itertools.count(0), name="add local buffers")
+  sink = graph_rewrite(sink, pm_add_local_buffers, ctx=itertools.count(0), name="add local buffers")
+  sink = graph_rewrite(sink, pm_reduce_local, ctx=ReduceContext(), name="remove_reduce")
+
+  # add locals
+  #sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, ctx=itertools.count(0), name="add local buffers")

  # ** devectorizer (full_graph_rewrite) **
  # remove reduce
-  sink = graph_rewrite(sink, pm_reduce+gep_pushing, ctx=ReduceContext(), name="remove_reduce")
+  #sink = graph_rewrite(sink, pm_reduce+gep_pushing, ctx=ReduceContext(), name="remove_reduce")

  # add gpu dims (late). this works after devectorize, but it's faster here
  sink = graph_rewrite(sink, pm_add_gpudims, ctx=ren, name="add gpudims")
@@ -96,44 +183,55 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  # **** optimizations are done, now we lower to actual code ****

  # add loads and remove invalids
-  sink = graph_rewrite(sink, pm_add_loads+pm_remove_invalid, name="** add loads (code)")
-
-  # create image buffers
-  if IMAGE and ren.target.device in {"QCOM", "CL", "PYTHON", "NULL"}:
-    sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True, ctx=ren.target.arch)
+  #sink = graph_rewrite(sink, pm_add_loads+pm_remove_invalid, name="** add loads (code)")

  # devectorize
-  sink = graph_rewrite(sink, sym+devectorize_alu+devectorize_buf_and_index+load_store_folding+correct_load_store+load_store_indexing,
-                       ctx=ren, name="devectorize")
+  #sink = graph_rewrite(sink, sym+devectorize_alu+devectorize_buf_and_index+load_store_folding, ctx=ren, name="devectorize")

-  # lower the index dtype to a concrete int
-  sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing+gep_pushing, name="lower all index dtypes")
-  sink = graph_rewrite(sink, symbolic, name="post index symbolic")
+  # add loads
+  sink = graph_rewrite(sink, pm_move_regs, name="** add loads")
+
+  # devectorize
+  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
+
+  # simplify indexing
+  sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing")
+
+  # do memory coalesing (late)
+  sink = memory_coalesing(sink, ren)
+  sink = graph_rewrite(sink, pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
+
+  # extra symbolic before decomp. crashes without this?
+  sink = graph_rewrite(sink, sym, name="extra symbolic")
+
+  # lower index dtype
+  # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
+  sink = graph_rewrite(sink, pm_lower_index_dtype+indexing_simplify, name="lower all index dtypes")
+
+  # final symbolic before decomp
+  sink = graph_rewrite(sink, symbolic, name="final symbolic")
+
+  # **** decomps ****

  # optional pre matcher
  if ren.pre_matcher is not None: sink = graph_rewrite(sink, ren.pre_matcher, name="pre_matcher")

-  # decompositions
+  # floordiv+mod / dtype decomp (early)
  supported_ops = tuple(ren.code_for_op.keys())
-  pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))
-  pm_transcendental = symbolic_simple+get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
-  sink = graph_rewrite(sink, pm_decomp, ctx=ren, name="decompositions")
+  pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops)
+  sink = graph_rewrite(sink, pm_decomp, name="early decompositions")
+
+  # late decomps + move gates from unrenderable INVALID where
  sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes")
-  sink = graph_rewrite(sink, pm_transcendental, name="transcendental")
-
-  # GEP/STACK stuff
-  sink = graph_rewrite(sink, pm_render, name="pm_render gep/stack")
-
-  # this is new style
-  sink = graph_rewrite(sink, pm_index_is_shrink, name="index is shrink")
-  sink = graph_rewrite(sink, pm_remove_vec_dtypes, name="transform to new style")
-
-  # move gates from unrenderable INVALID where
+  pm_decomp = pm_decomp+\
+    get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\
+    get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
+  sink = graph_rewrite(sink, pm_decomp, ctx=ren, name="late decompositions")
  sink = graph_rewrite(sink, pm_move_gates_from_index, name="move gates from index")

  # final rules for the renderer (without sym)
  extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
-  pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends
+  pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends+pm_no_weakints
  sink = graph_rewrite(sink, pm_final_rewrite, ctx=ren, name="final rewrite")

  # this was the linearizer
@@ -164,7 +262,7 @@ def line_rewrite(lst:list[UOp], pm:PatternMatcher, ctx=None) -> list[UOp]:
  replaced: dict[UOp, UOp] = {}
  for u in lst:
    nu = u.replace(src=tuple([replaced.get(x, x) for x in u.src]))
-    ret: tuple[UOp, list[UOp]] = cast(tuple[UOp, list[UOp]]|None, pm.rewrite(nu, ctx)) or (nu, [nu])
+    ret: tuple[UOp, list[UOp]] = pm.rewrite(nu, ctx) or (nu, [nu])
    replaced[u] = ret[0]
    newlst.extend(ret[1])
  return newlst
@@ -191,7 +289,7 @@ def do_assemble(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
  src = "\n".join(str(u.arg) for u in lin.src)
  if DEBUG >= 4: print(src)
  binary = ctx.asm(prg, lin)
-  return prg.replace(src=prg.src[:3]+(UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
+  return prg.replace(src=prg.src[:2]+(UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))

 def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
  src = ctx.render(list(lin.src))
@@ -205,11 +303,11 @@ def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None:
  return prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),))

 pm_to_program = PatternMatcher([
-  (UPat(Ops.PROGRAM, src=(UPat(Ops.SINK, name="sink"), UPat(Ops.DEVICE)), name="prg"), do_linearize),
-  (UPat(Ops.PROGRAM, src=(UPat(Ops.SINK, name="sink"), UPat(Ops.DEVICE), UPat(Ops.LINEAR, name="lin")), name="prg"), do_estimates),
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR, src=UPat(Ops.INS), name="lin")), name="prg"), do_assemble),
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR, name="lin")), name="prg"), do_render),
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
+  (UPat(Ops.PROGRAM, src=(UPat(Ops.SINK, name="sink"),), name="prg"), do_linearize),
+  (UPat(Ops.PROGRAM, src=(UPat(Ops.SINK, name="sink"), UPat(Ops.LINEAR, name="lin")), name="prg"), do_estimates),
+  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR, src=UPat(Ops.INS), name="lin")), name="prg"), do_assemble),
+  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR, name="lin")), name="prg"), do_render),
+  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
 ])

@track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
@@ -223,7 +321,7 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
    renderer: The renderer used to generate the code

  Returns:
-    The Ops.PROGRAM with SINK/DEVICE/LINEAR/SOURCE/BINARY.
+    The Ops.PROGRAM with SINK/LINEAR/SOURCE/BINARY.
  """
  if ast.op is Ops.PROGRAM: prg = ast
  elif ast.op is Ops.SINK:
@@ -234,7 +332,7 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
    if isinstance(renderer, ISARenderer):
      full_sink = graph_rewrite(full_sink, renderer.pre_isel_matcher, ctx=itertools.count(-1, -1), name="pre instruction selection", bottom_up=True)
      full_sink = graph_rewrite(full_sink, renderer.isel_matcher, ctx=IselContext(full_sink), name="instruction selection", bottom_up=True)
-    prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.target.device)), arg=prog_info)
+    prg = UOp(Ops.PROGRAM, src=(full_sink,), arg=prog_info)
  else: raise RuntimeError(f"can't call to_program on {ast.op}")
  if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0]))
  prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
@@ -0,0 +1,185 @@
+from tinygrad.dtype import dtypes, DType, truncate
+from tinygrad.helpers import flatten, DEBUG, EMULATED_DTYPES
+from tinygrad.uop import GroupOp
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite
+from tinygrad.renderer import Renderer
+from tinygrad.codegen.decomp.transcendental import exponent_bias, shl, shr
+
+# ***** long as 2 ints *****
+
+l2i_dt = {dtypes.long: dtypes.int, dtypes.ulong: dtypes.uint}
+def unpack32(v:UOp) -> tuple[UOp, UOp]: return v.bitcast(dtypes.uint) & 0xFFFF, shr(v.bitcast(dtypes.uint), 16)
+def reindex(idx:UOp, off:int, mul=2) -> UOp:
+  if idx.op is Ops.SHRINK:
+    assert mul == 1, "can't reindex SHRINK with mul != 1"
+    return idx.replace(op=Ops.INDEX, src=(idx.src[0], idx.src[1]+off))
+  return idx.replace(src=(idx.src[0], idx.src[1]*mul+off, *idx.src[2:]))
+
+# 4.3.1 is the relevant section in TAOCP
+def l2i(op: Ops, dt: DType, *uops:UOp):
+  zero = UOp.const(dt, 0)
+  if len(uops) == 2: a0, a1 = uops
+  elif len(uops) == 4: a0, a1, b0, b1 = uops
+  match op:
+    case Ops.NEG: return l2i(Ops.SUB, dt, zero, zero, *uops)
+    case Ops.CAST if dt in (dtypes.long, dtypes.ulong) and uops[0].dtype not in dtypes.floats:
+      return uops[0].cast(l2i_dt[dt]), (uops[0] < 0).where(UOp.const(l2i_dt[dt], -1), UOp.const(l2i_dt[dt], 0))
+    case Ops.CAST if dt in (dtypes.long, dtypes.ulong):
+      return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0)).cast(l2i_dt[dt])
+    case Ops.CAST if dt in dtypes.floats:
+      small = (a1.eq(0) & (a0 >= 0)) | (a1.eq(-1) & (a0 < 0))
+      return small.where(a0.cast(dt), ((a1.cast(dtypes.float32) * (2**32)) + a0.bitcast(dtypes.uint).cast(dtypes.float32)).cast(dt))
+    case Ops.CAST: return a0.bitcast(dtypes.uint).cast(dt)
+    case Ops.BITCAST: return a0.bitcast(dt), a1.bitcast(dt)
+    case Ops.SHL:
+      lo, hi = shl(a0, b0_mod:=b0 & 31), shl(a1, b0_mod) | shr(shr(a0, 1), 31 - b0_mod)
+      return (b0 >= 32).where(zero, lo), (b0 >= 32).where(lo, hi)
+    case Ops.SHR:
+      lo, hi = shr(a0, b0_mod:=b0 & 31) | shl(shl(a1, 1), 31 - b0_mod), shr(a1, b0_mod)
+      return (b0 >= 32).where(hi, lo), (b0 >= 32).where(zero, hi)
+    case Ops.ADD: return (low:=a0+b0), (a1 + b1).replace(dtype=dt) + (low.bitcast(dtypes.uint) < a0.bitcast(dtypes.uint)).cast(dt)
+    case Ops.SUB: return a0 - b0, a1 - b1 - (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint)).cast(dt)
+    case Ops.MUL:
+      (a00, a01), (b00, b01) = unpack32(a0), unpack32(b0)
+      mid = l2i(Ops.ADD, dt, shl(a00*b01, 16).bitcast(dt), shr(a00*b01, 16).bitcast(dt), shl(a01*b00, 16).bitcast(dt), shr(a01*b00, 16).bitcast(dt))
+      return l2i(Ops.ADD, dt, *mid, (a00*b00).bitcast(dt), (a01*b01).bitcast(dt) + a0*b1 + a1*b0)
+    case Ops.CDIV | Ops.CMOD:
+      # TAOCP Algorithm 4.3.1D could be faster here, but must be parameterized over the width of b
+      if dt == dtypes.int:
+        ua0, ua1, ub0, ub1 = a0.bitcast(dtypes.uint), a1.bitcast(dtypes.uint), b0.bitcast(dtypes.uint), b1.bitcast(dtypes.uint)
+        a0, a1 = (a_neg:=a1 < zero).where((n:=l2i(Ops.NEG, dtypes.uint, ua0, ua1))[0], ua0), a_neg.where(n[1], ua1)
+        b0, b1 = (b_neg:=b1 < zero).where((n:=l2i(Ops.NEG, dtypes.uint, ub0, ub1))[0], ub0), b_neg.where(n[1], ub1)
+      q, r = (z:=UOp.const(dtypes.uint, 0), z), (z, z)
+      for i in range(63, -1, -1):
+        r = l2i(Ops.SHL, dtypes.uint, *r, UOp.const(dtypes.uint, 1), z)
+        r = (r[0] | l2i(Ops.SHR, dtypes.uint, a0, a1, UOp.const(dtypes.uint, i), z)[0] & 1), r[1]
+        cond = l2i(Ops.CMPLT, dtypes.uint, *r, b0, b1).logical_not()
+        diff = l2i(Ops.SUB, dtypes.uint, *r, b0, b1)
+        q = ((q[0] | shl(cond.cast(dtypes.uint), i % 32), q[1]) if i < 32 else (q[0], q[1] | shl(cond.cast(dtypes.uint), i % 32)))
+        r = l2i(Ops.WHERE, dtypes.uint, cond, *diff, *r)
+      if dt == dtypes.int:
+        (nq0, nq1), (nr0, nr1) = l2i(Ops.BITCAST, dt, *l2i(Ops.NEG, dtypes.uint, *q)), l2i(Ops.BITCAST, dt, *l2i(Ops.NEG, dtypes.uint, *r))
+        (q0, q1), (r0, r1) = l2i(Ops.BITCAST, dt, *q), l2i(Ops.BITCAST, dt, *r)
+        return (a_neg.where(nr0, r0), a_neg.where(nr1, r1)) if op == Ops.CMOD else ((a_neg^b_neg).where(nq0, q0), (a_neg^b_neg).where(nq1, q1))
+      return (r[0].bitcast(dt), r[1].bitcast(dt)) if op == Ops.CMOD else (q[0].bitcast(dt), q[1].bitcast(dt))
+    case Ops.CMPLT: return (a1 < b1) | ((a1.eq(b1)) & (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint)))
+    case Ops.CMPEQ: return a0.eq(b0) & a1.eq(b1)
+    case Ops.CMPNE: return a0.ne(b0) | a1.ne(b1)
+    case Ops.XOR | Ops.OR | Ops.AND: return UOp(op, dt, src=(a0, b0)), UOp(op, dt, src=(a1, b1))
+    case Ops.WHERE: return uops[0].where(uops[1], uops[3]), uops[0].where(uops[2], uops[4])
+    case Ops.MAX: return l2i(Ops.WHERE, dt, l2i(Ops.CMPLT, dt, *uops), b0, b1, a0, a1)
+    case _: raise NotImplementedError(f"long decomposition of {op} unsupported")
+
+# ***** floats *****
+f2f_dt = { f:getattr(dtypes, f"uint{f.bitsize}") for f in dtypes.floats }
+
+def rne(v: UOp, s) -> UOp: return shr(v, s) + ((shr(v, s - 1) & 1) & ((v & ((1 << (s - 1)) - 1)).ne(0).cast(v.dtype) | (shr(v, s) & 1)))
+
+def f2f(v, fr:DType, to:DType, sat=True):
+  fs, fb, (fe, fm), ts, tb, (te, tm) = fr.bitsize, exponent_bias(fr), dtypes.finfo(fr), to.bitsize, exponent_bias(to), dtypes.finfo(to)
+  # NB: denormals are zero!
+  if fe <= te and fm < tm:
+    sign, nosign = shl((v & shl(1, fs-1)).cast(f2f_dt[to]), ts - fs), (v & (shl(1, fs-1) - 1)).cast(f2f_dt[to])
+    exp, norm = shr(nosign, fm), shl(nosign, tm - fm) + shl(tb - fb, tm)
+    nan = shl(nosign, tm - fm) | shl((shl(1, te) - 1), tm)
+    if fr in dtypes.fp8_fnuz:
+      fnuz_nan = sign.ne(0) & nosign.eq(0)
+      qnan = shl(shl(1, te) - 1, tm) | shl(1, tm - 1)
+      return fnuz_nan.where(qnan, sign | exp.eq(0).where(0, norm)).bitcast(to)
+    # fp8e4m3 has only one nan
+    is_nan = (nosign.eq(shl(1, fm + fe) - 1) if fr == dtypes.fp8e4m3 else exp.eq(shl(1, fe) - 1))
+    return (sign | exp.eq(0).where(0, is_nan.where(nan, norm))).bitcast(to)
+  elif fe >= te and fm > tm:
+    v = f2f_clamp(v.bitcast(fr), to, sat).bitcast(f2f_dt[fr])
+    sign, nosign = shr(v, fs - ts) & shl(1, ts - 1), v & (shl(1, fs - 1) - 1)
+    norm = (rne(nosign, fm - tm) - shl(fb - tb, tm)).cast(f2f_dt[to])
+    underflow = (shr(v, fm) & (shl(1, fe) - 1)) < (1 + fb - tb)
+    nan_mantissa = (shl(1, tm) - 1) if to == dtypes.fp8e4m3 else (shr(nosign, fm - tm) & (shl(1, tm) - 1))
+    nan = (sign | nan_mantissa | shl(shl(1, te) - 1, tm)).cast(f2f_dt[to])
+    is_nan = (shr(v, fm) & (shl(1, fe) - 1)).eq(shl(1, fe) - 1)
+    if to in dtypes.fp8_fnuz: return is_nan.where(shl(1, ts - 1), underflow.where(0, sign.cast(f2f_dt[to]) | norm))
+    return is_nan.where(nan, sign.cast(f2f_dt[to]) | underflow.where(0, norm))
+  else: raise NotImplementedError(f"unsupported decomp {fr} -> {to}")
+
+def f2f_clamp(val:UOp, dt:DType, sat=True) -> UOp:
+  e, m = dtypes.finfo(dt)
+  if dt in dtypes.fp8_fnuz: max_exp, max_man = (1 << e) - 1, (1 << m) - 1
+  else: max_exp, max_man = ((1 << e) - 1, (1 << m) - 2) if dt == dtypes.fp8e4m3 else ((1 << e) - 2, (1 << m) - 1)
+  mx = val.const_like(2.0**(max_exp - exponent_bias(dt)) * (1.0 + max_man / (1 << m)))
+  sat = mx if dt in dtypes.fp8s and sat else val.const_like(float('inf'))
+  # FIXME: CMPLT of nan is undefined
+  return val.ne(val).where(val, (val < -mx).where(-sat, (mx < val).where(sat, val)))
+
+def f2f_load(x: UOp, fr:DType, to:DType) -> UOp:
+  if (n:=x.max_numel()) == 1: return f2f(x.replace(dtype=f2f_dt[fr]), fr, to)
+  return UOp(Ops.STACK, to, tuple(f2f(x.replace(dtype=f2f_dt[fr], src=(reindex(x.src[0], i, 1),)), fr, to) for i in range(n)))
+
+def f2f_store(st, idx, val, fr:DType, to:DType):
+  if (n:=val.max_numel()) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr)))
+  return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.gep(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))
+
+pm_long_decomp = PatternMatcher([
+  (UPat(GroupOp.Defines, src=(UPat.var("sz"),), name="x"), lambda x,sz:
+   x.replace(dtype=l2i_dt[x.dtype], src=(sz*2,)) if x.dtype in l2i_dt else None),
+  (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x: reindex(x, x.tag).replace(dtype=l2i_dt[x.dtype]) if x.tag is not None else None),
+  (UPat(Ops.STORE, src=(UPat.var('idx'), UPat.var('val', tuple(l2i_dt.keys()))), name='st'), lambda st,idx,val:
+   st.replace(src=(idx.rtag(0), val.rtag(0))).group(st.replace(src=(idx.rtag(1), val.rtag(1)))) if val.tag is None else None),
+  (UPat(GroupOp.Comparison, src=(UPat.var('a', tuple(l2i_dt.keys())), UPat.var('b', tuple(l2i_dt.keys()))), name="x"), lambda a,b,x:
+   l2i(x.op, dt:=l2i_dt[a.dtype], a.rtag(0).cast(dt), a.rtag(1).cast(dt), b.rtag(0).cast(dt), b.rtag(1).cast(dt))),
+  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda a,x:
+   l2i(x.op, x.dtype, a)[x.tag] if x.tag is not None and a.dtype not in l2i_dt else None),
+  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
+   (a.rtag(0).cast(dt:=l2i_dt[a.dtype]).bitcast(xdt:=l2i_dt[x.dtype]), a.rtag(1).cast(dt).bitcast(xdt))[x.tag]),
+  (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
+   l2i(x.op, x.dtype, a.rtag(0).cast(dt:=l2i_dt[a.dtype]), a.rtag(1).cast(dt)) if x.dtype not in l2i_dt and a.tag is None else None),
+  (UPat((*(GroupOp.ALU - GroupOp.Comparison), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda x:
+   l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag(0).cast(dt:=l2i_dt[x.src[-1].dtype]), a.rtag(1).cast(dt))
+                                       if a.dtype in l2i_dt else (a,) for a in x.src))[x.tag] if x.tag is not None else None),
+  (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx:
+   x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag).replace(dtype=l2i_dt[x.dtype]),))),
+  (UPat(Ops.CONST, tuple(l2i_dt.keys()), name='x'), lambda x:
+   UOp.const(dt:=l2i_dt[x.dtype], truncate[dt]((x.arg >> 32) if x.tag == 1 else (x.arg & 0xFFFFFFFF))))
+])
+
+# float decomposition patterns - ctx is (fr, to) tuple
+pm_float_decomp = PatternMatcher([
+  (UPat((*GroupOp.Defines, Ops.INDEX, Ops.SHRINK), name="x"), lambda ctx,x:
+   x.replace(dtype=f2f_dt[ctx[0]], tag=ctx[0])
+   if x.dtype == ctx[0] and (x.op is not Ops.INDEX or x.src[0].op not in {Ops.LOAD, Ops.STACK}) else None),
+  (UPat(Ops.LOAD, dtypes.floats, name="x"), lambda ctx,x: f2f_load(x, *ctx) if x.dtype == ctx[0] else None),
+  # bitcasted load should just replace load
+  (UPat(Ops.BITCAST, src=(UPat(Ops.LOAD, name="ld"),), name="bc"), lambda ctx,bc,ld:
+   ld.replace(dtype=f2f_dt[ctx[0]]).bitcast(bc.dtype) if ld.dtype == ctx[0] else None),
+  # bitcast from
+  (UPat(Ops.BITCAST, src=(UPat.var("x", dtypes.floats),), name="bc"), lambda ctx,bc,x:
+   bc.replace(src=(f2f(x.bitcast(f2f_dt[ctx[1]]), ctx[1], ctx[0]),)) if x.dtype == ctx[1] and bc.dtype.bitsize == ctx[0].bitsize else None),
+  # bitcast to
+  (UPat(Ops.BITCAST, src=(UPat.var("x"),), name="bc"), lambda ctx,bc,x:
+   f2f(x.bitcast(f2f_dt[ctx[0]]), ctx[0], ctx[1]) if bc.dtype == ctx[0] else None),
+  (UPat(Ops.CAST, dtypes.floats, src=(UPat.var("val"),), name="x"), lambda ctx,x,val:
+   f2f_clamp(val.cast(ctx[1]), ctx[0]) if x.dtype == ctx[0] else None),
+  (UPat(GroupOp.All-{Ops.BITCAST}, dtypes.floats, name="x"), lambda ctx,x:
+   x.replace(dtype=ctx[1], src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src))
+   if x.dtype == ctx[0] else None),
+  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat(Ops.BITCAST, dtypes.floats, name="val")), name='st'), lambda ctx,st,idx,val:
+   st.replace(src=(idx, val.replace(dtype=f2f_dt[ctx[0]]))) if val.dtype == ctx[0] and idx.tag == ctx[0] else None),
+  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val", dtypes.floats)), name='st'), lambda ctx,st,idx,val:
+   f2f_store(st, idx, val, *ctx) if val.dtype == ctx[1] and (idx:=idx.src[0] if idx.op == Ops.CAST else idx).tag == ctx[0] else None),
+])
+
+def do_dtype_decomps(sink:UOp, ctx:tuple[set[DType], Renderer]) -> UOp:
+  def _should_emulate(dt): return dt in EMULATED_DTYPES.tolist(dtypes) or dt not in ctx[1].supported_dtypes()
+  for fr in sorted(filter(_should_emulate, ctx[0])):
+    to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float
+    if DEBUG >= 2: print(f"emulating {fr} as {to}")
+    sink = graph_rewrite(sink, pm_float_decomp if fr in dtypes.floats else pm_long_decomp, name=f"decomp {fr} -> {to}", ctx=(fr, to), bottom_up=True)
+  ctx[0].clear()
+  return sink
+
+pm_dtype_decomps = PatternMatcher([
+  # detect dtypes to decompose
+  (UPat(GroupOp.All, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half, dtypes.long, dtypes.ulong), name="x"), lambda x,ctx:
+   ctx[0].add({dtypes.ulong:dtypes.long}.get(dt:=x.dtype, dt))),
+  # do the rewrites
+  (UPat(Ops.SINK, name="sink"), do_dtype_decomps),
+])
@@ -0,0 +1,131 @@
+from typing import Callable
+import functools
+from tinygrad.dtype import dtypes, promo_lattice
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher
+from tinygrad.renderer import Renderer
+
+# *** integer division ***
+
+@functools.lru_cache(None)
+def magicgu(vmax:int, d:int) -> tuple[int,int]:
+  # calculate m,s such that x//d == (x*m) >> s for all 0 <= x <= vmax, d>0; adapted from Hacker's Delight, Chapter 10
+  nc = (vmax+1)//(d) * d - 1
+  nbits = vmax.bit_length()
+  for s in range(0, 2*nbits + 1):
+    if 2**s > nc*(d - 1 - (2**s - 1) % d):
+      m = (2**s + d - 1 - (2**s - 1) % d)//d
+      return m, s
+  assert False
+
+def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None:
+  from tinygrad.renderer.cstyle import MetalRenderer
+  # NOTE: disable for METAL due to compiler bug. keccak with -O0 works but not with optimization
+  if isinstance(ren, MetalRenderer): return None
+  # If d is a power of two this is not valid for signed ints!
+  is_unsigned = x.vmin>=0 or x.dtype in dtypes.uints
+  assert d>0, "Sign should have been taken out of divisor"
+  vmin,vmax = max(x.vmin, x.dtype.min), min(x.vmax, x.dtype.max)
+  if vmin > -d and vmax < d: return x.const_like(0)
+  m,s = magicgu(max(vmax, abs(vmin)), d)
+  if m*vmin >= x.dtype.min and m*vmax <= x.dtype.max:
+    return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)
+  # before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
+  # use explicit Ops.CDIV (trunc) since the recursion assumes trunc semantics throughout
+  if (largest_factor_of_two_in_d := (d & -d)) > 1:
+    if (ret:=fast_idiv(ren, x.alu(Ops.CDIV, x.const_like(largest_factor_of_two_in_d)),
+                       d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
+  if dont_cast: return None
+  # promo_lattice needs to return an unsigned type if the type is unsigned
+  if dtypes.is_int(next_dtype := promo_lattice[x.dtype][-1]) and next_dtype in ren.supported_dtypes():
+    if m*vmin >= next_dtype.min and m*vmax <= next_dtype.max:
+      return ((x.cast(next_dtype)*m) >> s).cast(x.dtype) if is_unsigned else ((x.cast(next_dtype)*m) >> s).cast(x.dtype) + (x<0).where(x.ufix(1), 0)
+  return None
+
+# ***** threefry *****
+
+def threefry2x32(x: UOp, key: UOp):
+  # split x and key from uint64 to two uint32
+  x0, x1 = (x & 0xffffffff).cast(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast(dtypes.uint32)
+  key0, key1 = (key & 0xffffffff).cast(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast(dtypes.uint32)
+
+  rotations = [[13, 15, 26, 6], [17, 29, 16, 24]]
+  ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0]
+  xr:list[UOp] = [x0 + ks[-1], x1 + ks[0]]
+  for i in range(5):
+    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r)))
+    xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)]
+
+  return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64)
+
+# ***** decomposition patterns *****
+
+def floordiv_to_idiv(a:UOp, b:UOp) -> UOp:
+  if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.CDIV, b)
+  return a.alu(Ops.CDIV, b) - (a.alu(Ops.CMOD, b).ne(0) & (a<0).ne(b<0)).cast(a.dtype)
+
+def floormod_to_mod(a:UOp, b:UOp) -> UOp:
+  if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.CMOD, b)
+  r = a.alu(Ops.CMOD, b)
+  # use where instead of mul to avoid being fused into MULACC (which int64 long-decomp doesn't handle)
+  return r + (r.ne(0) & (a<0).ne(b<0)).where(b, b.const_like(0))
+
+powers_of_two: dict[int, int] = {2**i:i for i in range(64)}
+@functools.cache
+def get_simplifying_rewrite_patterns(ops:tuple[Ops, ...]) -> PatternMatcher:
+  # these are rewrites that make things simpler
+  pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)]
+  # FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod
+  if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None))
+  pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod))
+  # no real hardware supports THREEFRY, but NullRenderer does
+  if Ops.THREEFRY not in ops: pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32))
+  # MAX can be rewritten as CMPLT + WHERE (max function is annoying on many cstyle backends)
+  if Ops.MAX not in ops and Ops.CMPLT in ops: pat.append((UPat(Ops.MAX, name="m"), lambda m: (m.src[0] < m.src[1]).where(m.src[1], m.src[0])))
+  return PatternMatcher(pat)
+
+@functools.cache
+def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> PatternMatcher:
+  pat: list[tuple[UPat, Callable]] = []
+  if Ops.OR in ops: pat += [(UPat.var("x", dtypes.bool).logical_not()&UPat.var("y", dtypes.bool).logical_not(),
+    lambda x,y: (x | y).logical_not())]
+  # rewrite MUL/CDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
+  if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
+  if Ops.SHR in ops:
+    # uint CDIV by 2**v -> x >> v (FLOORDIV is lowered to CDIV by the rule above before reaching here)
+    pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.uints), UPat.cvar("c"))),
+      lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
+    # signed CDIV (trunc) by 2**v -> (x + (x<0 ? c-1 : 0)) >> v
+    pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("c"))),
+      lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v
+        if (v:=powers_of_two.get(c.arg, 0)) else None)]
+    if not disable_fast_idiv:
+      # fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0)
+      pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))),
+        lambda ctx, x, d: fast_idiv(ctx, x, d.arg) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
+      # rewrite raw CMOD -> x - d*CDIV(x,d) so fast_idiv can pick up the CDIV. only on non-negative inputs;
+      # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail)
+      pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.var("d"))),
+        lambda x, d: x - d * x.alu(Ops.CDIV, d) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
+  if Ops.NEG in ops:
+    pat += [(UPat.var('x')*-1, lambda ctx,x: x.alu(Ops.NEG))]
+    if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda ctx,x,y: x.alu(Ops.SUB, y))]
+  if Ops.CMPLT in ops:
+    # These are late rewrites because simplex expects equalities to be a certain format
+    pat += [
+      ((UPat.var("x", dtypes.sints) < UPat.cvar("c", dtypes.sints)).logical_not(), lambda x,c: c-1<x),
+      ((UPat.cvar("c", dtypes.sints) < UPat.var("x", dtypes.sints)).logical_not(), lambda x,c: x<c+1),
+      (UPat.var("x", dtypes.sints)*-1 < UPat.var("y", dtypes.sints)*UPat.cvar("c"), lambda x,y,c: y*(-c)<x),
+      (UPat.var("x", dtypes.sints)*-1 < UPat.cvar("c"), lambda x,c:-c<x),
+      ((UPat.cvar("c1")<UPat.var("x", dtypes.sints)) & (UPat.var("x", dtypes.sints)<UPat.cvar("c2")),
+        lambda x,c1,c2: x.eq(c1+1) if c1.arg+1==c2.arg-1 else None),  # (c-1)<x & x<(c+1) -> x==c
+    ]
+  if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))]
+  if Ops.MULACC in ops:
+    pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
+    # also fuse (x << n) + c → MULACC(x, 2^n, c) since MUL→SHL may run first
+    if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1<<n.arg), c))]
+  # some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b
+  if Ops.FDIV in ops:
+    pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))]
+    pat += [(UPat.var("a", dtypes.floats) * UPat.const(dtypes.floats, 1).alu(Ops.FDIV, UPat.var("b")), lambda a,b: a.alu(Ops.FDIV, b))]
+  return PatternMatcher(pat)
@@ -0,0 +1,277 @@
+from typing import Callable
+import math, functools
+from tinygrad.dtype import dtypes, DType
+from tinygrad.helpers import polyN
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher
+
+TRANSCENDENTAL_DTYPES = (dtypes.float16, dtypes.float32, dtypes.float64)
+
+def _lazy_map_numbers(x:UOp, inf:UOp, _inf:UOp, nan:UOp, ratio:UOp):
+  """replace inf -> inf, -inf -> _inf, nan -> nan, otherwise -> ratio"""
+  return x.ne(math.inf).where(x.ne(x).where(nan, x.ne(-math.inf).where(ratio, _inf)), inf)
+
+# *** helper functions for bit manipulation ***
+def mantissa_bits(d:DType) -> int: return dtypes.finfo(d)[1]
+def exponent_bias(d:DType) -> int: return (1 << (dtypes.finfo(d)[0] - 1)) - (0 if d in dtypes.fp8_fnuz else 1)
+def exponent_mask(d:DType) -> int: return (1 << dtypes.finfo(d)[0]) - 1
+
+# **** utils ****
+def shr(x:UOp|int, y:UOp|int) -> UOp: return x // (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y)
+def shl(x:UOp|int, y:UOp|int) -> UOp: return x * (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y)
+
+def rintk(d:UOp) -> UOp:
+  """round d:float to int away from 0"""
+  out_dtype = {dtypes.float64: dtypes.int64, dtypes.float32: dtypes.int32, dtypes.float16: dtypes.int16}[d.dtype]
+  return (d + (d<0.0).where(d.const_like(-0.5), d.const_like(0.5))).cast(out_dtype)
+
+def pow2if(q:UOp, float_dtype:DType):
+  """cast(2^q, float_dtype) where q is any integer in the range of [-126, 127]"""
+  out_dtype = {dtypes.int64: dtypes.float64, dtypes.int32: dtypes.float32, dtypes.int16: float_dtype}[q.dtype]
+  return shl(q + exponent_bias(out_dtype), mantissa_bits(out_dtype)).bitcast(out_dtype)
+
+def ilogb2k(d:UOp) -> UOp:
+  """calculate the integer part of log2(d), where d is normalized fp value in the range of [0, +inf)."""
+  assert d.dtype in TRANSCENDENTAL_DTYPES
+  dint = d.bitcast({dtypes.float64: dtypes.int64, dtypes.float32: dtypes.int32, dtypes.float16: dtypes.int16}[d.dtype])
+  # -1 <= ilog2bk(d) <= 128
+  return (shr(dint, mantissa_bits(d.dtype)) & exponent_mask(d.dtype)) - exponent_bias(d.dtype)
+
+def ldexp3k(d:UOp, e:UOp) -> UOp:
+  """d*2^e. e is a number obtained by casting an integer in the range [-127, 127] to a float. d is any float number."""
+  assert d.dtype in TRANSCENDENTAL_DTYPES and e.dtype in TRANSCENDENTAL_DTYPES
+  dtype = {dtypes.float64: dtypes.int64, dtypes.float32: dtypes.int32, dtypes.float16: dtypes.int16}[d.dtype]
+  m1 = d.bitcast(dtype)
+  m2 = shl(e.cast(dtype), mantissa_bits(d.dtype))
+  return (m1 + m2).bitcast(d.dtype)
+
+def ldexp2k(d:UOp, e:UOp) -> UOp:
+  """d*2^e. much faster than ldexp3k but risky. d > 0 and d is not denormal."""
+  assert d.dtype in TRANSCENDENTAL_DTYPES and e.dtype in (dtypes.int16, dtypes.int32, dtypes.int64)
+  return (d * pow2if(shr(e, 1), d.dtype)) * pow2if(e - shr(e, 1), d.dtype)
+
+def frexp(v:UOp) -> tuple[UOp, UOp]:
+  """frexp(v) -> (mantissa, exponent) assuming v != 0"""
+  assert v.dtype in TRANSCENDENTAL_DTYPES
+  # m1 = masks for mantissa, m2 = masks to normalize the mantissa.
+  m1 = {dtypes.float64: 0x000FFFFFFFFFFFFF, dtypes.float32: 0x807FFFFF, dtypes.float16: 0x83FF}[v.dtype]
+  m2 = {dtypes.float64: 0x3FE0000000000000, dtypes.float32: 0x3F000000, dtypes.float16: 0x3800}[v.dtype]
+  bits = v.bitcast({dtypes.float64: dtypes.uint64, dtypes.float32: dtypes.uint32, dtypes.float16: dtypes.uint16}[v.dtype])
+  exponent = shr(bits, mantissa_bits(v.dtype)) & exponent_mask(v.dtype)
+  # Set the exponent bits appropriately to normalize the mantissa into the range of [0.5, 1.0).
+  mantissa = ((bits & m1) | m2).bitcast(v.dtype)
+  exp = exponent - exponent_bias(v.dtype) + 1
+  return mantissa, exp
+
+# *** reduction algorithms for sine ***
+def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
+  """
+  Performs Payne-Hanek Reduction: computes the remainder of `d` modulo pi/2 for the values `d` where
+    39800.0 <= d <= +Inf
+  Returns a tuple of `(r, q)`:
+  - `r`[d.dtype] is the reminder value corresponding to `round_to_nearest(x % pi/2)`.
+  - `q`[int32] is an integer, and q % 4 is corresponding to the quadrant of the original angle `d`.
+  """
+  assert d.dtype in TRANSCENDENTAL_DTYPES
+  # https://stackoverflow.com/questions/30463616/payne-hanek-algorithm-implementation-in-c/30465751#30465751
+  # 190 bits of 2/pi for Payne-Hanek style argument reduction
+  two_over_pi_f = [0x00000000, 0x28be60db, 0x9391054a, 0x7f09d5f4, 0x7d4d3770, 0x36d8a566, 0x4f10e410]
+
+  intermediate_dtype = dtypes.float32 if d.dtype == dtypes.float16 else d.dtype
+
+  f, e = frexp(d)
+  ia = (f.cast(intermediate_dtype) * 4.294967296e9).cast(dtypes.uint64)
+  # extract 96 relevant bits of 2/pi based on magnitude of argument
+  i = shr(e.cast(dtypes.uint64), 5)
+  e = e.cast(dtypes.int32) & 31
+  offset = 32 - e
+
+  def _take(an:UOp, offset:int, count:int=0) -> UOp:
+    """an = two_over_pi_f[i+offset]"""
+    if count+offset < len(two_over_pi_f) - 1:
+      an = i.ne(count).where(_take(an, offset, count=count+1), an.const_like(two_over_pi_f[count+offset]))
+    return an
+  def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) * pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
+  def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) // pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
+
+  a = [_take(UOp.const(dtypes.uint32, 0), i) for i in range(4)]
+  #  (two_over_pi_f[Int(i) + n] << e) | (two_over_pi_f[Int(i) + n+1] >> (nbits - e))
+  # Note: e >= 1 for all numbers d >= 1.0. assume e != 0
+  hi = _shl_lazy(a[0], e) | _shr_lazy(a[1], offset)
+  mi = _shl_lazy(a[1], e) | _shr_lazy(a[2], offset)
+  lo = _shl_lazy(a[2], e) | _shr_lazy(a[3], offset)
+
+  def _hp_mul(x:UOp, y:UOp) -> UOp: return x.cast(dtypes.uint64) * y.cast(dtypes.uint64)
+  # compute x * 2/pi
+  p = shl(_hp_mul(ia, hi), 32) + _hp_mul(ia, mi) + shr(_hp_mul(ia, lo), 32)
+
+  # round quotient to nearest
+  q = shr(p, 62).cast(dtypes.int32)
+  p = p & 0x3fffffffffffffff
+  r = (p.cast(intermediate_dtype) * (3.4061215800865545e-19)).cast(d.dtype)
+
+  # if fraction >= 0.5, r -= pi/2, q += 1
+  return (f<0.5).where(r, r - math.pi/2), (f<0.5).where(q, q + 1)
+
+def cody_waite_reduction(d:UOp) -> tuple[UOp, UOp]:
+  """
+  Performs Cody-Waite Reduction: computes the reminder of `d` modulo pi/2 for the values `d` where
+      0 <= abs(d) <= 39800.0
+  Returns a tuple of `(r, q)`, where the output format is the same as that of `payne_hanek_reduction`.
+  """
+  def _reduce_d(x:UOp, q:UOp):
+    # https://github.com/shibatch/sleef/blob/4e08851f59fc2b545f9c393c6a23dfd311a26308/src/libm/sleefdp.c#L789-L823
+    if x.dtype == dtypes.float64:
+      # https://github.com/shibatch/sleef/blob/f6d8a841fbfddd26ce712834d4da220cd76048fb/src/common/misc.h#L77
+      PI_A, PI_B, PI_C, PI_D = 3.1415926218032836914, 3.1786509424591713469e-08, 1.2246467864107188502e-16, 1.2736634327021899816e-24
+      d = qdh * -PI_A + x
+      d = q * -PI_A + d
+      d = qdh * -PI_B + d
+      d = q * -PI_B + d
+      d = qdh * -PI_C + d
+      d = q * -PI_C + d
+      d = (qdh + q) * -PI_D + d
+    elif x.dtype == dtypes.float16:
+      # [FIXME] when reducing `d`, FP16 needs FP32 precision to achieve 1.0 ULP precision.
+      d = _reduce_d(x.cast(dtypes.float32), q.cast(dtypes.float32)).cast(dtypes.float16)
+    else:
+      # https://github.com/shibatch/sleef/blob/4e08851f59fc2b545f9c393c6a23dfd311a26308/src/libm/sleefsp.c#L464-L503
+      d = q * -3.1414794921875 + x
+      d = q * -0.00011315941810607910156 + d
+      d = q * -1.9841872589410058936e-09 + d
+      d = q * -1.2154201256553420762e-10 + d
+    return d
+
+  m_1_pi = 0.318309886183790671537767526745028724
+  qdh = (d * (m_1_pi / 2.0**24)).cast(dtypes.int64).cast(d.dtype) * (2.0**24)
+  quadrant = rintk(d * m_1_pi -qdh) if d.dtype == dtypes.float64 else rintk(d * m_1_pi)
+  return _reduce_d(d, quadrant.cast(d.dtype)), quadrant.cast(dtypes.int32)
+
+# *** approximate sine on small angle. ***
+def trig_poly(d:UOp, coeff32, coeff64): return d * (polyN(d*d, coeff64) if d.dtype == dtypes.float64 else polyN(d*d, coeff32))
+# approximate sine on [-pi/2, pi/2]
+def sin_poly(d:UOp) -> UOp:
+  return trig_poly(d, [2.6083159809786593541503e-06, -0.0001981069071916863322258, 0.00833307858556509017944336, -0.166666597127914428710938, 1.0],
+                      [-7.97255955009037868891952e-18, 2.81009972710863200091251e-15, -7.64712219118158833288484e-13, 1.60590430605664501629054e-10,
+                       -2.50521083763502045810755e-08, 2.75573192239198747630416e-06, -0.000198412698412696162806809, 0.00833333333333332974823815,
+                       -0.166666666666666657414808,    1.0])
+
+def _ifand(q:UOp, n:int): return (q & n).ne(0)
+
+def sin_poly_small(d:UOp, q:UOp) -> UOp:
+  r = sin_poly(d)
+  return r * _ifand(q, 1).where(r.const_like(-1), r.const_like(1))
+
+def sin_poly_large(d:UOp, q:UOp) -> UOp:
+  r = sin_poly(d + _ifand(q, 1).where(d.const_like(math.pi / 2), d.const_like(0)))
+  return r * _ifand(q, 2).where(r.const_like(-1), r.const_like(1))
+
+# *** toplevel functions for xsin/xlog2/xexp2 ***
+
+def xsin(d:UOp, fast:bool=False, switch_over:float=30.0) -> UOp:
+  """
+  Implements a 1.0 ULP approximation for Ops.SIN.
+  - fast=True assumes x <= switch_over.
+  - switch_over is the threshold for switching to payne_hanek_reduction.
+  """
+  assert d.dtype in TRANSCENDENTAL_DTYPES
+  # mask +-inf/nan as zero
+  x = _lazy_map_numbers(d, d.const_like(0.0), d.const_like(0.0), d.const_like(0.0), d)
+  # x_sign = sign(x)
+  x_sign = x.ne(0).where((x<0).where(x.const_like(-1), x.const_like(1)), x.const_like(0))
+  x_abs = x * x_sign
+  r, q = (cody_waite_reduction if fast else payne_hanek_reduction)(x_abs)
+  if fast: result = sin_poly_small(r, q)
+  else:
+    # Payne Hanek Reduction assumes abs(x) >= pi/4, so for smaller values, use cody_waite_reduction.
+    r_small, q_small = cody_waite_reduction(x_abs)
+    result = (x_abs<switch_over).where(sin_poly_small(r_small, q_small), sin_poly_large(r, q))
+  # adjusts the sign for abs(x)
+  result = result * x_sign
+  # sin(Inf) = NaN, sin(-Inf) = NaN, sin(NaN) = NaN
+  return _lazy_map_numbers(d, d.const_like(math.nan), d.const_like(math.nan), d.const_like(math.nan), result)
+
+def xexp2(d:UOp) -> UOp:
+  """
+  Implements a 1.0 ULP approximation for Ops.EXP2
+  - Paper: https://arxiv.org/pdf/2001.09258
+  """
+  assert d.dtype in TRANSCENDENTAL_DTYPES
+  # mask +=inf/nan as zero.
+  x = _lazy_map_numbers(d, d.const_like(0.0), d.const_like(0.0), d.const_like(0.0), d)
+  q = rintk(x)
+  # s = d - round(d)
+  s = x - q.cast(x.dtype)
+  # a polynomial approximation with 13 non-zero terms in the range of [−(log 2)/2,(log 2)/2].
+  if d.dtype == dtypes.float64:
+    u = polyN(s, [0.4434359082926529454e-9, 0.7073164598085707425e-8, 0.1017819260921760451e-6, 0.1321543872511327615e-5, 0.1525273353517584730e-4,
+                  0.1540353045101147808e-3, 0.1333355814670499073e-2, 0.9618129107597600536e-2, 0.5550410866482046596e-1, 0.2402265069591012214e+0,
+                  0.6931471805599452862e+0, 0.1000000000000000000e+1])
+  else: u = polyN(s, [0.1535920892e-3, 0.1339262701e-2, 0.9618384764e-2, 0.5550347269e-1, 0.2402264476e+0, 0.6931471825e+0, 1.0])
+  u = ldexp2k(u, q) # u*2^q
+  upper, lower = {dtypes.float64: (1024, -2000), dtypes.float32: (128, -150), dtypes.float16: (23, -22)}[d.dtype]
+  # Replace x >= upper with +inf
+  u = (d >= upper).where(d.const_like(math.inf), u)
+  # Replace x < lower with zero.
+  u = (d<lower).where(d.const_like(0.0), u)
+  # exp2(NaN) = NaN
+  return d.ne(d).where(d.const_like(math.nan), u)
+
+def xlog2(d:UOp) -> UOp:
+  """
+  Implements a 1.0 ULP approximation for Ops.LOG2
+  Paper: https://arxiv.org/pdf/2001.09258 5.5
+  """
+  assert d.dtype in TRANSCENDENTAL_DTYPES
+  # float16 uses 2^10 for denormal scaling (2^64 overflows), float32/64 use 2^64
+  denormal_exp = 10 if d.dtype == dtypes.float16 else 64
+  FLT_MIN = d.const_like({dtypes.float16: 6.1e-5, dtypes.float32: 1e-4, dtypes.float64: 1e-4}[d.dtype])
+  is_denormal = d<FLT_MIN
+  a = is_denormal.where(d * (2 ** denormal_exp), d)
+
+  e = ilogb2k(a * (1.0 / 0.75)).cast(a.dtype)
+  m = ldexp3k(a, -e)
+  e = is_denormal.where(e - denormal_exp, e)
+
+  x = (m - 1.0) / (m + 1.0)
+  x2 = x * x
+  if d.dtype == dtypes.float64:
+    t = polyN(x2, [0.2211941750456081490e+0, 0.2200768693152277689e+0, 0.2623708057488514656e+0, 0.3205977477944495502e+0,
+                   0.4121985945485324709e+0, 0.5770780162997058982e+0, 0.96179669392608091449])
+    r = t * (x * x2) + e + x * 2.885390081777926774
+  else:
+    t = polyN(x2, [0.4374550283e+0, 0.5764790177e+0, 0.9618012905120])
+    # s_lo term (x*3.27e-08) only for float32 - underflows in float16
+    r = t * (x * x2) + e + x * 2.8853900432586669922 + (x * 3.2734474483568488616e-08 if d.dtype == dtypes.float32 else 0)
+
+  # log2(Inf) = Inf
+  r = d.ne(math.inf).where(r, r.const_like(math.inf))
+  # log2(0) = -Inf (handle both +0.0 and -0.0)
+  r = d.ne(0.0).where(r, r.const_like(-math.inf))
+  # log2(x) = NaN for x < 0
+  r = (d<-0.0).where(r.const_like(math.nan), r)
+  # log2(NaN) = NaN
+  r = d.ne(d).where(r.const_like(math.nan), r)
+  # log2(-0.0) = -Inf. In certain devices like PTX, x == -0.0 won't be true. so making reciprocal.
+  return d.reciprocal().ne(-math.inf).where(r, r.const_like(-math.inf))
+
+def xpow(base:UOp, exponent:UOp) -> UOp:
+  # start with b ** e = exp2(e * log2(b))
+  ret = (base < 0).where(-base, base).log2().mul(exponent).exp2()
+  # negative base: nan for non-integer exponent, negate for odd integer exponent
+  non_int = exponent != exponent.cast(dtypes.int32).cast(exponent.dtype)
+  is_odd = (exponent < 0).where(-exponent, exponent).cast(dtypes.int32).mod(2).cast(dtypes.bool)
+  neg_base = non_int.where(ret.const_like(math.nan), is_odd.where(-ret, ret))
+  # fix 0 ** 0 = 1
+  return (base.eq(0) & exponent.eq(0)).where(ret.const_like(1), (base < 0).where(neg_base, ret))
+
+@functools.cache
+def get_transcendental_patterns(ops:tuple[Ops, ...], force_transcendental:bool) -> PatternMatcher:
+  pat: list[tuple[UPat, Callable]] = []
+  for op,f in ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)):
+    if op not in ops or force_transcendental:
+      pat += [(UPat(op, dtype=TRANSCENDENTAL_DTYPES, src=(UPat.var("d"),)), f),
+              (UPat(op, dtype=tuple(dt for dt in dtypes.floats if dt not in TRANSCENDENTAL_DTYPES), src=(UPat.var("d"),), name="x"),
+                lambda x,d: d.cast(dtypes.float32).alu(x.op).cast(x.dtype))]
+  # rewrite SQRT to xpow 0.5
+  if Ops.SQRT not in ops or force_transcendental: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5))))
+  return PatternMatcher(pat)
@@ -96,7 +96,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
      if len(missing_locals):
        assert len(idx.src) == 2, "index has 2 sources"
        mask: UOp = UOp.uprod(*[x.eq(0) for x in missing_locals])
-        subs[idx] = idx.replace(src=(idx.src[0], mask.broadcast(idx.src[1].dtype.count).where(idx.src[1], Invalid)))
+        subs[idx] = idx.replace(src=(idx.src[0], mask.broadcast(idx.src[1].max_numel()).where(idx.src[1], Invalid)))
    if r.op is not Ops.RANGE: continue
    try:
      ii = (global_dims+local_dims).index(r.arg[0:-1])
@@ -0,0 +1,123 @@
+from typing import Any
+import itertools
+from collections import defaultdict
+from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp
+from tinygrad.helpers import getenv, IMAGE, all_same
+from tinygrad.renderer import Renderer
+from tinygrad.uop.symbolic import symbolic_simple
+from tinygrad.codegen.late.devectorizer import image_valid_dims, _drop_valid_stmts, uop_given_valid
+
+def do_devectorize(b:UOp):
+  if b.shape == (): return None
+  # broadcasting needs to be already unpacked
+  if not all_same([x.shape for x in b.src]): return None
+  src = []
+  for idx in itertools.product(*[range(x) for x in b.shape]):
+    idx_c = [UOp.const(dtypes.weakint, i) for i in idx]
+    src.append(b.replace(src=tuple([x.index(*idx_c) for x in b.src])))
+  return UOp._stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
+
+devectorizer2 = PatternMatcher([
+  # unpack broadcasting
+  (UPat(GroupOp.Elementwise, name="b"), do_devectorize),
+])
+
+def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
+  shapes, ren = ctx
+  if not IMAGE or ren.target.device not in {"QCOM", "CL", "PYTHON", "NULL"}: return None
+  valid = UOp.const(dtypes.bool, True)
+  if x.op == Ops.WHERE and x.src[2].op == Ops.CONST and x.src[2].arg == Invalid: valid,x,_= x.src
+  # search for dims that drop the most valid statements
+  best_drop, cands = -1, []
+  for ch, cw in [shapes[buf.arg.slot]] if buf.arg.slot in shapes else image_valid_dims(buf.dtype, buf.max_numel(), ren.target.arch):
+    cidx = uop_given_valid(valid, UOp.vectorize((x//4)%cw, x//(4*cw)))
+    dropped = len(_drop_valid_stmts(valid, cidx, ch, cw))
+    if dropped > best_drop: best_drop, cands = dropped, [(ch, cw, cidx)]
+    elif dropped == best_drop: cands.append((ch, cw, cidx))
+  # if no candidates, we don't rewrite
+  if len(cands) == 0: return None
+  # and tiebreak with indexing complexity (ie. number of nodes)
+  h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].gep(1).simplify().backward_slice))
+  buf = buf.replace(dtype=(dtypes.imageh if buf.dtype.itemsize == 2 else dtypes.imagef)((h, w, 4)))
+  shapes[buf.arg.slot] = (h, w)
+  if valid.op is not Ops.CONST or valid.arg is not True:
+    return buf.index(valid.where(cidx.src[1], cidx.src[1].const_like(Invalid)),
+                     valid.where(cidx.src[0], cidx.src[0].const_like(Invalid)))
+  else:
+    return buf.index(cidx.src[1], cidx.src[0])
+
+pm_simplify_add_image = PatternMatcher([
+  (UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))), transform_to_image),
+  # image load/store is always float
+  (UPat(Ops.INDEX, dtype=dtypes.float, name="x").load(dtype=dtypes.half), lambda x: x.load().cast(dtypes.half)),
+  (UPat(Ops.INDEX, dtype=dtypes.float, name="x").store(UPat(name="d", dtype=dtypes.half)), lambda x,d: x.store(d.cast(dtypes.float))),
+  (UPat.var("x", dtype=dtypes.float).cast(dtypes.half).cast(dtypes.float), lambda x: x),
+])+devectorizer2+symbolic_simple
+
+def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
+  if getenv("DMC"): return sink
+
+  # collect
+  memory: defaultdict[tuple[Ops, UOp, Any, Any], dict[int, list[UOp]]]  = defaultdict(dict)
+  for u in sink.toposort():
+    # TODO: this should handle images too, it's just memory coalesing
+    if u.op in {Ops.LOAD, Ops.STORE}:
+      assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalesing does not support gated loads/stores"
+      assert u.src[0].op is Ops.INDEX, f"src must be INDEX, not {u.src[0].op} on {u.op}"
+      buf, idx_u = u.src[0].src
+      if buf.addrspace == AddrSpace.REG: continue
+      idx: Any = idx_u.src[1] if idx_u.op is Ops.WHERE and idx_u.src[2].arg is Invalid else idx_u
+      valid: Any = idx_u.src[0] if idx_u.op is Ops.WHERE and idx_u.src[2].arg is Invalid else None
+      if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg
+      elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].arg
+      elif idx.op is Ops.CONST and idx.arg is Invalid: root_src, arg = "INVALID", 0
+      elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.arg
+      else: root_src, arg = idx, 0
+      memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u)
+
+  # build replacements
+  replacements = {}
+  for (op,buf,base,valid),offsets in memory.items():
+    # allowed lengths (copied in)
+    lengths = []
+    must_divide = True
+    if ctx is not None and ctx.target.device == "DSP":
+      lengths = [128,64,32,16,8,4]
+      must_divide = False
+    elif buf.dtype not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not isinstance(buf.dtype, ImageDType):
+      pass
+    elif buf.addrspace == AddrSpace.REG:
+      pass
+    elif isinstance(buf.dtype, ImageDType):
+      lengths = [4]
+    elif ctx is not None and ctx.supports_float4:
+      # TODO: a better way to get this than ctx
+      lengths = [8,4,2] if buf.dtype == dtypes.half and getenv("ALLOW_HALF8") else [4,2]
+    lengths.append(1)  # worst case, it's not folded
+    # do the grouping
+    grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
+    for full_grp in grouped_offsets:
+      while len(full_grp):
+        offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(dtypes.weakint, full_grp[0])
+        length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.divides(l) is not None)][0]
+        grp = full_grp[:length]
+        # NOTE: we apply the valid again after we determine the length
+        offset = valid.where(offset, UOp(Ops.CONST, offset.dtype, arg=Invalid)) if valid is not None else offset
+        idx = UOp(Ops.SHRINK, dtype=buf.dtype, src=(buf, offset, UOp.const(dtypes.weakint, len(grp)))) if len(grp) > 1 else buf.index(offset)
+        if op == Ops.STORE:
+          datas = []
+          for i,g in enumerate(grp):
+            assert len(offsets[g]) == 1, f"attempting multiple stores: {len(offsets[g])}"
+            datas.append(offsets[g][0].src[1])
+          store = idx.store(UOp._stack(*datas) if len(datas) > 1 else datas[0])
+          for i,g in enumerate(grp): replacements[offsets[g][0]] = store
+        else:
+          ld = idx.load()
+          for i,g in enumerate(grp):
+            for oo in offsets[g]:
+              replacements[oo] = ld.index(UOp.const(dtypes.weakint, i)) if len(grp) > 1 else ld
+        full_grp = full_grp[length:]
+
+  # apply
+  return sink.substitute(replacements, name="memory coalesing")
@@ -1,12 +1,9 @@
-from typing import Any, cast
 import functools, itertools
-from collections import defaultdict
 from dataclasses import dataclass
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
-from tinygrad.helpers import getenv, flatten, prod
-from tinygrad.renderer import Renderer
+from tinygrad.helpers import getenv, flatten, prod, OSX, ceildiv

 # ***** image load valid simplification *****

@@ -14,7 +11,7 @@ from tinygrad.renderer import Renderer
 def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
  # can drop valid if idx is out of bound when valid is False
  drop_stmt = []
-  for stmt in valid.split_uop(Ops.AND):
+  for i,stmt in enumerate(valid.split_uop(Ops.AND)):
    if (res:=parse_valid(stmt)) is None: continue
    X, is_upper_bound, c = res

@@ -25,12 +22,12 @@ def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
        drop_stmt.append(stmt)
        continue

-    # if X <= c, check if it's out of bound when X = c+1
-    # if X >= c, check if it's out of bound when X = c-1
-    test_value = c + 1 if is_upper_bound else c - 1
-    for i,b in zip(idx.src, (width, height)):
-      if i.is_increasing():
-        rw = i.substitute({X:X.const_like(test_value)})
+    # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1]
+    lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1)
+    if lo <= hi:
+      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype)
+      for coord,b in zip(idx.src, (width, height)):
+        rw = coord.substitute({X:fake}).simplify()
        if rw.vmin >= b or rw.vmax < 0:
          drop_stmt.append(stmt)
          break
@@ -51,7 +48,7 @@ def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|N
  idx_y, idx_x = idx.gep(1), idx.gep(0)
  return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), ptr=True) if new_valid is not None else buf.index(idx_y, idx_x, ptr=True)

-load_store_indexing = PatternMatcher([
+indexing_simplify = PatternMatcher([
  # image load valid idx simplification
  (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)),
@@ -60,168 +57,42 @@ load_store_indexing = PatternMatcher([

 # ***** load/store grouping *****

+# get list of (height, width) that do not require pitch padding
+def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]:
+  if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return []
+  MAXW, pxls = 16384, size // 4
+  if base not in (dtypes.half, dtypes.float) or size > 4*MAXW*MAXW: return []
+  # height=1 images just need to abide by alignment requirements in bytes, not pixels!
+  if size % (ALIGN * 4) != 0: return [] if (base.itemsize * size) % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)]
+  return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]
+
 def expand_index(ctx, buf:UOp, vec:UOp):
-  # determine optimal image shapes
-  if isinstance(dt:=buf.dtype, ImageDType):
-    idxs, valids = vec.get_idx(), vec.get_valid()
-    lane = next((i for i in range(vec.dtype.count) if valids.gep(i).vmax != 0), 0)
-    x, valid = idxs.gep(lane), valids.gep(lane)
-    # search for dims that drop the most valid statements
-    best_drop, cands = -1, []
-    for ch, cw in ImageDType.valid_dims(dt, ctx.target.arch):
-      if (dropped:=len(_drop_valid_stmts(valid, cidx:=uop_given_valid(valid, UOp.vectorize((x//4)%cw, x//(4*cw))), ch, cw))) > best_drop:
-        best_drop, cands = dropped, [(ch, cw, cidx)]
-      elif dropped == best_drop: cands.append((ch, cw, cidx))
-    # and tiebreak with indexing complexity (ie. number of nodes)
-    h, w, _ = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].gep(1).simplify().backward_slice))
-    assert buf.op is Ops.RESHAPE
-    buf = buf.src[0].replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4))).flatten()
  if getenv("UNSAFE_DISABLE_MASK", 0): vec = vec.get_idx()
  # generate the individual indexes
  return UOp(Ops.STACK, buf.dtype, tuple(buf.index(vec.gep(i), ptr=True) for i in range(vec.dtype.count)))

-def fold_expanded_index(midx:UOp):
-  buf = midx.src[0].src[0]
-  if not all(s.src[0] is buf for s in midx.src): return None
-  if not all(isinstance(s.dtype, PtrDType) for s in midx.src): return None
+def load_stack(stack:UOp, ld:UOp):
+  offset, ret = 0, []
+  for x in stack.src:
+    src = [x]
+    for s in ld.src[1:]:
+      src.append(s.gep(tuple(range(offset, offset+x.dtype.count))) if s.dtype.vcount > 1 else s)
+    ret.append(ld.replace(dtype=x.dtype.base, src=tuple(src)))
+    offset += x.dtype.count
+  return UOp(Ops.STACK, stack.dtype.base.vec(len(stack.src)), tuple(ret))

-  # extract all the relevant offsets
-  offsets_rootsrc: defaultdict[Any, dict[int, list[int]]] = defaultdict(dict)
-  for i in range(len(midx.src)):
-    idx: Any = midx.src[i].src[1].get_idx()
-    if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg
-    elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].arg
-    elif idx.op is Ops.CONST and idx.arg is Invalid: root_src, arg = "INVALID", 0
-    elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.arg
-    else: root_src, arg = idx, 0
-    root_src = (midx.src[i].src[1].get_valid(), root_src)
-    offsets_rootsrc[root_src].setdefault(arg, []).append(i)
-
-  # then rewrite everything we can into groups
-  ret = []
-  idxs: list[int|None] = [None]*len(midx.src)
-  global_offset = 0
-  for offsets in offsets_rootsrc.values():
-    grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
-    for grp in grouped_offsets:
-      # get the index offset for this element. using [0] is okay, because they are the same
-      lidx = midx.src[offsets[grp[0]][0]]
-      if len(grp) > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(len(grp)).ptr(size=buf.max_numel(), addrspace=buf.addrspace))
-      # set the idxs of the output
-      for i,g in enumerate(grp):
-        for oo in offsets[g]: idxs[oo] = global_offset+i
-      # add this lidx to the CAT
-      ret.append(lidx)
-      global_offset += len(grp)
-  assert None not in idxs, f"some idxs are missing {idxs}"
-  # this base thing is for image, we want the CAT to be a normal pointer
-  post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.max_numel(), addrspace=buf.addrspace).vec(global_offset), tuple(ret))
-  return post_cat.gep(tuple(cast(list[int], idxs)))
-
-def cat_after_store(cat:UOp, data:UOp):
-  # TODO: this is written in many places
-  offset = 0
-  ret: list[UOp] = []
-  for s in cat.src:
-    ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count)))))
-    offset += s.dtype.count
+def store_stack(stack:UOp, data:UOp):
+  offset, ret = 0, []
+  for x in stack.src:
+    ret.append(x.store(data.gep(tuple(range(offset, offset+x.dtype.count)))))
+    offset += x.dtype.count
  return UOp.group(*ret)

-def gep_on_store(gep:UOp, st:UOp):
-  # NOTE: we need to invert the gep here, but it may be an expanding gep
-  # fake argsort. TODO: handle duplicates
-  a = {}
-  for i,x in enumerate(gep.arg): a[x] = i
-  new_arg = tuple(x[1] for x in sorted(a.items()))
-  return gep.src[0].store(st.gep(new_arg))
-
 load_store_folding = PatternMatcher([
  (UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat(name="buf")), UPat.var("vec"))), expand_index),
-  (UPat(Ops.STACK, src=UPat(Ops.INDEX), name="midx"), fold_expanded_index),
-  # GEP after LOAD
-  (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True),
-   lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)),
-  # GEP on data of STORE
-  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st"))), gep_on_store),
-  # put PTRCAT after LOAD
-  (UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True),
-   lambda cat,ld: UOp(Ops.VCAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
-  # put PTRCAT after STORE
-  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data"))), cat_after_store),
-])
-
-# *** correct load/store ***
-
-def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
-  # this splits loads and stores into multiple chunks
-
-  # if there's only one element to load/store, no splitting needed
-  if (sz:=ls.src[0].dtype.count) == 1: return None
-  buf = idx.src[0]
-
-  # determine fold lengths
-  lengths = []
-  must_divide = True
-  if ctx is not None and ctx.target.device == "DSP":
-    lengths = [128,64,32,16,8,4]
-    must_divide = False
-  elif buf.dtype.base not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not isinstance(buf.dtype, ImageDType):
-    pass
-  elif buf.addrspace == AddrSpace.REG:
-    pass
-  elif isinstance(buf.dtype, ImageDType):
-    lengths = [4]
-  elif ctx is not None and ctx.supports_float4:
-    # TODO: a better way to get this than ctx
-    lengths = [8,4,2] if buf.dtype.base == dtypes.half and getenv("ALLOW_HALF8") else [4,2]
-  lengths.append(1)  # worst case, it's not folded
-
-  # filter fold lengths that don't divide
-  offset, mask = idx.src[1].get_idx(), idx.src[1].get_valid()
-  if must_divide: lengths = [x for x in lengths if offset.divides(x) is not None]
-
-  # split based on the fold lengths
-  global_offset = 0
-  ret = []
-  while global_offset < sz:
-    # with 1 at the end of the lengths list, this will always hit
-    for fold_length in lengths:
-      if global_offset+fold_length > sz: continue
-      lidx = buf.index((offset + global_offset).valid(mask), ptr=True)
-      if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.max_numel(), addrspace=buf.addrspace))
-      if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))))
-      else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length)))
-      global_offset += fold_length
-      break
-
-  # if it wasn't split, we return None. otherwise we CAT them
-  if len(ret) <= 1: return None
-  return UOp(Ops.VCAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
-
-def get_image_idx(idx:UOp, width:int):
-  x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
-  idx_x, idx_y = (x // 4) % width, x // (4*width)
-  assert idx.src[0].op is Ops.RESHAPE, "image idx must be on reshape"
-  return idx.replace(src=(idx.src[0].src[0], idx_y.valid(valid), idx_x.valid(valid)))
-
-def image_fixup(ls:UOp):
-  # normal image load or store, with the CAST from expand_index
-  if isinstance(dt:=ls.src[0].src[0].dtype, ImageDType) and ls.src[0].op is Ops.CAST:
-    assert ls.src[0].dtype.count == 4, "image must be casted to 4"
-    return ls.replace(src=(get_image_idx(ls.src[0].src[0], dt.shape[1]),)+ls.src[1:])
-
-  # this is an unprocessed image without a cast, we should just make it a buffer
-  if isinstance(dt, ImageDType) and len(ls.src[0].src) == 2:
-    off = ls.src[0].src[1]
-    assert ls.src[0].src[0].op is Ops.RESHAPE, "image idx must be on reshape"
-    idx = ls.src[0].src[0].src[0].replace(dtype=(new_dt:=dtypes.half if dt.itemsize == 2 else dtypes.float).ptr(dt.size)).index(off)
-    return ls.replace(src=(idx,), dtype=new_dt).cast(dtypes.float) if ls.op is Ops.LOAD else ls.replace(src=(idx, ls.src[1].cast(new_dt)))
-
-correct_load_store = PatternMatcher([
-  # split LOAD/STORE
-  (UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.INDEX, name="idx").cast(),), name="ls", allow_any_len=True), split_load_store),
-  # image indexing, including unfoldable images
-  (UPat((Ops.LOAD, Ops.STORE), name="ls"), image_fixup),
+  # put STACK of indexes after LOAD/STORE
+  (UPat(Ops.LOAD, src=(UPat(Ops.STACK, src=UPat(Ops.INDEX), name="stack"),), name="ld", allow_any_len=True), load_stack),
+  (UPat(Ops.STORE, src=(UPat(Ops.STACK, src=UPat(Ops.INDEX), name="stack"), UPat(name="data"))), store_stack),
 ])

 # *** uop expander ***
@@ -240,7 +111,6 @@ def no_vectorized_wmma(wmma:UOp):

 def no_vectorized_alu(alu:UOp):
  if alu.dtype.vcount == 1: return None
-  if alu.op is Ops.WHERE and alu.src[2].arg is Invalid: return None  # image load/store has cond.where(idx.vec(2), Invalid) as the index
  alus = tuple(UOp(alu.op, alu.dtype.scalar(), tuple(s.gep(i) for s in alu.src), alu.arg) for i in range(alu.dtype.vcount))
  return UOp(Ops.STACK, alu.dtype, alus)

@@ -370,28 +240,3 @@ pm_add_loads = PatternMatcher([
  (UPat(Ops.STORE, src=(UPat(Ops.LOAD),), allow_any_len=True, name="s"), lambda s: s.replace(src=(s.src[0].src[0],)+s.src[1:])),
  (UPat(Ops.LOAD, src=(UPat(Ops.LOAD),), allow_any_len=True, name="l"), lambda l: l.replace(src=(l.src[0].src[0],)+l.src[1:])),
 ])
-
-# make images
-
-pm_imageh_store = PatternMatcher([
-  # store<imageh>(idx, x) is actually store(idx, x.cast(half)) so we can pull the cast into the store
-  (UPat.var("x", dtypes.float).cast(dtypes.half), lambda x: x),
-  # store(imageh, a.where(b.half(), c).float()) -> store(imageh, a.where(b, c.float()))
-  (UPat(Ops.WHERE, src=(UPat.var("a"), UPat.var("b", dtypes.float).cast(dtypes.half), UPat.var("c"))), lambda a,b,c: a.where(b,c.cast(dtypes.float))),
-  # otherwise, we cast to float
-  (UPat(GroupOp.All, name="x"), lambda x: x.cast(dtypes.float))
-])
-
-def make_image(ctx, ls, buf, off):
-  if (vcount:=buf.dtype.vcount) != 1: buf = buf.src[0]
-  if buf.op == Ops.PARAM and not isinstance(dt:=buf.dtype, ImageDType) and (dims:=ImageDType.valid_dims(dt, ctx)):
-    buf = buf.replace(dtype=(dtypes.imageh if dt.base == dtypes.half else dtypes.imagef)((*dims[0], 4))).flatten()
-    if vcount != 1: buf = UOp.vectorize(*([buf] * vcount))
-    if ls.op is Ops.LOAD: return ls.replace(src=(buf.index(off, ptr=True),), dtype=dtypes.float.vec(ls.dtype.vcount)).cast(dt.base)
-    return buf.index(off, ptr=True).store(pm_imageh_store.rewrite(ls.src[1]) if dt.base == dtypes.half else ls.src[1])
-
-pm_make_images = PatternMatcher([
-  (UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))),), allow_any_len=True, name="ls"), make_image),
-  # load<imageh> is actually load<half>.cast(float), so load<imageh>.half().float() -> load<half>.float().half().float() -> load<half>.float()
-  (UPat(Ops.LOAD, name="li").cast(dtypes.half).cast(dtypes.float), lambda li: li if isinstance(li.src[0].dtype, ImageDType) else None),
-])
@@ -51,7 +51,7 @@ def do_expand(root:UOp):
        new_srcs.append(src)
      elif src.dtype.count > 1:
        # put any input dtype > 1 grouped together
-        new_srcs.append(UOp(Ops.VCAT, src.dtype.scalar().vec(expand_sz*src.dtype.count), (src,)*expand_sz))
+        new_srcs.append(src.gep(tuple(i for _ in range(expand_sz) for i in range(src.dtype.count))))
      else:
        # repeat the arg
        new_srcs.append(src.broadcast(expand_sz))
@@ -51,7 +51,7 @@ class LinearScanRegallocContext:
      if v not in self.spills:
        # the value of a BUFFER is its 64bit address
        dt = self.vdef(v).dtype
-        sz = 8 if self.vdef(v).op is Ops.BUFFER else dt.scalar().itemsize * dt.count
+        sz = 8 if self.vdef(v).op is Ops.BUFFER else dt.itemsize
        offset = self.stack_size + (sz - self.stack_size % sz) % sz
        self.spills[v] = UOp.const(dtypes.int32, offset)
        self.stack_size = offset + sz
@@ -1,8 +1,9 @@
 import itertools
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, IMAGE
-from tinygrad.dtype import PtrDType, ImageDType
+from tinygrad.dtype import PtrDType
 from tinygrad.uop.ops import Ops, resolve, AxisType
+from tinygrad.codegen.late.devectorizer import image_valid_dims
 from tinygrad.codegen.opt.postrange import Scheduler

 def hand_coded_optimizations(k:Scheduler) -> Scheduler:
@@ -50,7 +51,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
  # upcast float4 images, this must be early so we don't accidentally add locals before the upcast
  if IMAGE:
    for buf_index,buf in enumerate(k.bufs):
-      if isinstance(buf.src[0].dtype, PtrDType) and ImageDType.valid_dims(buf.src[0].dtype, k.ren.target.arch):
+      if isinstance(buf.src[0].dtype, PtrDType) and image_valid_dims(buf.src[0].dtype.base, buf.src[0].dtype.size, k.ren.target.arch):
        # part of is_expanded
        unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].get_idx().split_uop(Ops.ADD) if
          c.op is Ops.RANGE and (c.vmax+1)%4 == 0]
@@ -101,6 +102,12 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
    # for Schedule, we check if the range is used in INDEX gates or WHERE gates
    is_masked = k.rngs[axis] in where_gate_rngs
    if k.full_shape[axis] <= 7 and is_masked and prod(k.full_shape[j] for j in to_upcast) * k.full_shape[axis] <= 7 * 7:
+      # upcasting a masked global axis moves that range out of the launch grid into each work-item
+      # under IMAGE, skip the upcast unless enough global work-items remain after it to hide memory latency
+      if IMAGE and k.axis_types[axis] is AxisType.GLOBAL:
+        global_upcast = prod(k.full_shape[i] for i in to_upcast if k.axis_types[i] is AxisType.GLOBAL) * k.full_shape[axis]
+        global_items_after = prod(k.full_shape[i] for i in k.axes_of(AxisType.GLOBAL)) // global_upcast
+        if resolve(global_items_after < getenv("OCCUPANCY_FLOOR", 4096), False): continue
      if DEBUG >= 4: print(f"upcasting masked axis : {axis}")
      to_upcast.append(axis)
  for axis in to_upcast[::-1]: k.apply_opt(Opt(OptOps.UPCAST, axis, 0))
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	6a05336468	something	2026-06-30 11:59:32 -07:00
geohot	f7b5502ec2	just the new devectorizer	2026-06-30 11:19:52 -07:00
qazalandGitHub	795559b3d4	viz/cli: support passing only a start marker (#16802 ) * failing test * support end interval	2026-06-30 17:06:32 +09:00
George HotzandGitHub	b683ff9836	test updates + addrspace work from new codegen (#16800 )	2026-06-29 19:12:10 -07:00
geohot	af0e339114	test fix	2026-06-29 18:37:49 -07:00
sirhcmandGitHub	395eec1866	disable mlperf training benchmark (#16798 )	2026-06-29 20:11:16 -04:00
chenyuandGitHub	3a9b87e48f	remove Ops.DEVICE from Ops.PROGRAM [PR] (#16797 )	2026-06-29 17:41:31 -04:00
chenyuandGitHub	379b8d37fc	read PROGRAM device from ParamArg [PR] (#16796 ) nothing reads from its Ops.DEVICE after this	2026-06-29 16:06:37 -04:00
nimlgenandGitHub	724f58cdbe	Revert "do not assert index in mem coalesing (#16723 )" (#16795 ) This reverts commit `b99b9e1875`.	2026-06-29 18:05:56 +03:00
nimlgenandGitHub	2ea8439d4d	viz: fix enter_calls (#16793 )	2026-06-29 17:10:26 +03:00
chenyuandGitHub	e836899336	remove Tensor conv2d and dot override [PR] (#16792 )	2026-06-29 09:58:21 -04:00
nimlgenandGitHub	d9c65bd843	viz: fix walk rewrite (#16791 )	2026-06-29 16:32:32 +03:00
qazalandGitHub	905b405820	fix contiguous_view_offset for scalar index (#16789 ) * simple failing test * scalar passes	2026-06-29 12:50:21 +09:00
chenyuandGitHub	59de2fa3d3	combine pm_reduce_collapse rules [PR] (#16787 )	2026-06-28 13:10:13 -04:00
George HotzandGitHub	0bee0e6e39	do full sym in extra symbolic [pr] (#16786 ) * do full sym in extra symbolic [pr] * put that back to 55	2026-06-28 08:31:47 -07:00
nimlgenandGitHub	4c3789661d	hcq2: split hcq_schedule into hcq_compile/hcq_link for jit (#16784 )	2026-06-28 10:19:48 +03:00
George HotzandGitHub	a6ea5373fa	shape arg as const (#16783 ) * shape arg can be single const * shape arg can be single const * fix just the const for a shape	2026-06-27 23:38:13 -07:00
chenyuandGitHub	9db8cbe1a0	minor geitem cleanup [PR] (#16777 )	2026-06-27 15:09:37 -04:00
chenyuandGitHub	5ff82838f0	Ops.BUFFER arg is always ParamArg [PR] (#16776 )	2026-06-27 14:02:35 -04:00
chenyuandGitHub	36c0e81979	deprecate Tensor.training [PR] (#16774 )	2026-06-27 11:02:28 -04:00
George HotzandGitHub	afe9fcaec8	remove dtype.vec from transcendental (glm) [pr] (#16773 ) * remove dtype.vec from transcendental (glm) [pr] * more cleanups * fix fuzzer	2026-06-26 21:41:31 -07:00
George HotzandGitHub	59aec35a7d	more devec cleanups (glm) (#16772 )	2026-06-26 19:27:04 -07:00
George HotzandGitHub	a6ed4237e5	remove Ops.VCAT and Ops.PTRCAT (glm) (#16771 ) * remove Ops.VCAT and Ops.PTRCAT (glm) * fix stack, and remove useless casts	2026-06-26 18:35:38 -07:00
chenyuandGitHub	5d1867a6c3	more Tensor.training -> TRAINING [PR] (#16770 ) code and test and doc	2026-06-26 20:11:59 -04:00
sirhcmandGitHub	f8ddca9a00	qcom: ir3 needs SP_MODE_CNTL.CONSTANT_DEMOTION_ENABLE (#16769 )	2026-06-26 19:48:29 -04:00
chenyuandGitHub	f5612ebbc8	remove Ops.UNIQUE and Ops.LUNIQUE (#16768 ) 88 ops	2026-06-26 16:56:10 -04:00
chenyuandGitHub	9a864701b8	no Ops.DEVICE in Ops.COPY [PR] (#16765 ) moved to arg	2026-06-26 13:34:28 -04:00
nimlgenandGitHub	fe9cd76e8b	hcq2: cleanup (#16767 )	2026-06-26 20:26:35 +03:00
chenyuandGitHub	8c9e8d38a9	fix permute copy (#16766 ) flip and permute copy was wrong	2026-06-26 11:42:50 -04:00
chenyuandGitHub	f6e7426028	no Ops.DEVICE in Ops.ALLREDUCE [PR] (#16764 ) moved to arg[1]	2026-06-26 09:51:40 -04:00
chenyuandGitHub	434e233640	LUNIQUE -> ParamArg [PR] (#16756 )	2026-06-26 09:04:49 -04:00
qazalandGitHub	b8224e19a7	llama: remove unused fp8 transpose kernel (#16763 )	2026-06-26 16:31:13 +09:00
wozeparrotandGitHub	5857297e64	llama: benchmark offload copy script (#16762 )	2026-06-26 00:26:32 -07:00
George HotzandGitHub	cd80c87fdb	remove dead code from devectorizer (#16761 )	2026-06-26 00:00:52 -07:00
George HotzandGitHub	a94a32ff71	move image to post coalese (#16749 ) * move image to post coalese * fix types * dv2 * work * load_store_indexing * lsi * simplify indexing * always simplify * more * whitespace * no simplify * no early opt * fix linter * exit early without valid * always simplify * relax ir3 * no gep * Revert "no gep" This reverts commit `3fa4dd6dfa`. * shapes in ctx * disable IR3	2026-06-25 23:33:07 -07:00
George HotzandGitHub	5fcb21c1b3	remove vectorize/const symbolic rule [pr] (#16759 ) * remove vectorize/const symbolic rule * post index symbolic is late * simpler * less gates * that's fine * this * work * don't rerun * this order * late removal of index dtype * skip that test * it's fine with extra symbolic	2026-06-25 17:05:14 -07:00
sirhcmandGitHub	d5a852b9ee	ci: cleanup macos tests (#16760 )	2026-06-25 19:47:29 -04:00
George HotzandGitHub	b34a0017b8	GEP -> INDEX stays (#16757 )	2026-06-25 14:48:11 -07:00
George HotzandGitHub	ca52ba7ec1	move floordiv late [pr] (#16755 ) * move floordiv late * index/on/stack	2026-06-25 13:59:30 -07:00
chenyuandGitHub	d953eb95d5	don't use Ops.UNIQUE in Buffer [PR] (#16754 ) build new_buffer with ParamArg	2026-06-25 16:00:59 -04:00
George HotzandGitHub	7f4ceeb2b9	move memory coalesing after new style (#16753 )	2026-06-25 11:53:02 -07:00
chenyuandGitHub	93a82d42bd	Buffer arg -> max_numel [PR] (#16752 ) prep removing UNIQUE and use ParamArg for arg	2026-06-25 14:50:20 -04:00
sirhcmandGitHub	c759301865	ci: cleanup torch backend tests (#16750 )	2026-06-25 14:39:36 -04:00
George HotzandGitHub	dc57efc3a3	move down dtype decomps (try 2) (#16746 ) * move down dtype decomps (try 2) * gpt fixes * fixes * that did need reindex * valid.where * upd to master * reindex	2026-06-25 09:49:43 -07:00
George HotzandGitHub	b5e0bca764	move the load/store gate back to gater (#16747 ) * move the load/store gate back to gater * order	2026-06-25 09:29:32 -07:00
geohot	fa60fa96a6	hotfix: decomp in pyproject	2026-06-25 08:56:34 -07:00
George HotzandGitHub	e72ecd101f	split decomp into files (#16745 ) * split decomp into files * that's in op	2026-06-25 08:40:02 -07:00
George HotzandGitHub	54887cbbdf	move decomp files into codegen (#16744 )	2026-06-25 08:19:02 -07:00
chenyuandGitHub	96c9a68205	imrpove some type casts (#16743 )	2026-06-25 10:48:28 -04:00
chenyuandGitHub	db2257ad09	don't run tinyboxbenchmark on fork (#16741 )	2026-06-25 09:12:56 -04:00
qazalandGitHub	daf8c8f509	viz: simpler test_link_sched_codegen, merge with the other codegen test (#16739 ) * simpler test_link_sched_codegen * merge test_codegen_tracing in the other test * include BEAM.value	2026-06-25 15:02:00 +09:00
qazalandGitHub	695ffb3a08	viz: fix rendering custom Ops.PROGRAM source (#16738 ) * failing test * fix * simplify test	2026-06-25 14:21:31 +09:00
George HotzandGitHub	2f91a401a1	move late decomp after new style (#16734 )	2026-06-24 22:04:59 -07:00
qazalandGitHub	cb9df8bd33	fix profile.sh BENCHMARK=3 exit status (#16736 )	2026-06-25 13:41:03 +09:00
sirhcmandGitHub	63f0e08593	start benchmarks cleanup (#16729 )	2026-06-24 23:33:46 -04:00
wozeparrotandGitHub	e0c69d7a12	llama: fused grad quantize (#16731 )	2026-06-24 23:25:39 -04:00
chenyuandGitHub	65dd099b63	support Invalid in _pad_constant [PR] (#16733 ) and simplified image_conv2d	2026-06-24 21:03:36 -04:00
George HotzandGitHub	9d5d2253b3	refactor to image_valid_dims (#16730 )	2026-06-24 16:53:09 -07:00
nimlgenandGitHub	b99b9e1875	do not assert index in mem coalesing (#16723 )	2026-06-24 14:28:49 +03:00
Nikolay Murzin GitHub Claude Opus 4.8	f38836d14d	cuda: pass kernel dynamic shared memory in the graph runner (#16699 ) The CUDA graph runner hardcoded sharedMemBytes=0 in the kernel node params, so a custom_kernel using dynamic shared memory (extern __shared__) failed with CUDA error 700 on graph replay while the eager path worked. Pass runtime.smem (0 for normal kernels, which use static shared). Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>	2026-06-24 14:20:10 +03:00
qazalandGitHub	535c806c9f	llama: profile.sh BENCHMARK=3 (#16722 )	2026-06-24 19:20:51 +09:00
chenyuandGitHub	857a435af2	remove dead added_ox (#16721 )	2026-06-23 23:59:17 -04:00
chenyuandGitHub	687ade119e	IMAGE hand_coded_optimizations update (#16720 )	2026-06-23 21:55:28 -04:00
George HotzandGitHub	0a8e61d0c5	switch to the new memory coaleser [pr] (#16716 ) * switch to the new memory coalese * move that stuff * copy in allowed length logic * mulitple buffers * new coalese is better * fine * earlier * fixes * work * work * valid * stack on index const	2026-06-23 18:03:48 -07:00
wozeparrotandGitHub	dfea9e7994	llama: fused silu mul quantize mxfp8 (#16704 )	2026-06-23 16:59:50 -07:00
chenyuandGitHub	ce87d80911	better _drop_valid_stmts [pr] (#16719 ) also dropped the unused is_increasing	2026-06-23 19:35:01 -04:00
George HotzandGitHub	5a2b3b7b06	early dtype decomp (#16718 ) * early dtype decomp * simplify * cleanup * that goes there * doing too much * stupid symbolic rules	2026-06-23 16:07:20 -07:00
sirhcmandGitHub	116045cc8e	ci: remove tensorflow from testoptim (#16717 )	2026-06-23 18:11:48 -04:00
nimlgenandGitHub	7c1d0b6d9a	hcq2: use shrink(bitcast) (#16713 ) * hcq2: use shrink(bitcast) * x	2026-06-23 18:11:39 +03:00