realize sched

Merge branch 'master' into outerworld_work
train loop
2026-08-14 05:58:28 +00:00 · 2025-10-14 14:44:56 +08:00 · 2025-10-14 14:26:52 +08:00 · 2025-10-10 20:20:19 +08:00 · 2025-10-10 20:07:44 +08:00 · 2025-10-10 19:49:24 +08:00
561 changed files with 221313 additions and 206905 deletions
@@ -1,3 +0,0 @@
-[run]
-source = tinygrad
-branch = True
@@ -41,10 +41,6 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
-  mesa:
-    description: "Install mesa"
-    required: false
-    default: 'false'
 runs:
  using: "composite"
  steps:
@@ -61,7 +57,7 @@ runs:
      uses: actions/cache@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ hashFiles('**/setup.py') }}-${{ env.PYTHON_CACHE_VERSION }}

    # **** Caching downloads ****

@@ -70,13 +66,13 @@ runs:
      uses: actions/cache@v4
      with:
        path: ~/.cache/tinygrad/downloads/
-        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
+        key: downloads-cache-${{ inputs.key }}-${{ env.DOWNLOAD_CACHE_VERSION }}
    - name: Cache downloads (macOS)
      if: inputs.key != '' && runner.os == 'macOS'
      uses: actions/cache@v4
      with:
        path: ~/Library/Caches/tinygrad/downloads/
-        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
+        key: osx-downloads-cache-${{ inputs.key }}-${{ env.DOWNLOAD_CACHE_VERSION }}

    # **** Python deps ****

@@ -187,7 +183,7 @@ runs:
      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.APT_CACHE_VERSION }}

    - name: Run apt Update + Install
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
@@ -221,7 +217,7 @@ runs:
        sudo mkdir -p /usr/local/lib
        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
          jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
-          sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
+          sudo xargs curl -L -o /usr/local/lib/libamd_comgr.dylib
        cargo build --release --manifest-path ./extra/remu/Cargo.toml

    # **** gpuocelot ****
@@ -247,7 +243,7 @@ runs:
        cache-name: cache-gpuocelot-build-1
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
    - name: Clone/compile gpuocelot
      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
      shell: bash
@@ -278,7 +274,7 @@ runs:
      if: inputs.webgpu == 'true' && runner.os == 'Linux'
      shell: bash
      run: |
-        sudo curl -fL https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/local/lib/libwebgpu_dawn.so
+        sudo curl -L https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/local/lib/libwebgpu_dawn.so
        sudo ldconfig
    - name: Install WebGPU dawn (macOS)
      if: inputs.webgpu == 'true' && runner.os == 'macOS'
@@ -293,13 +289,3 @@ runs:
      if: inputs.llvm == 'true' && runner.os == 'macOS'
      shell: bash
      run: brew install llvm@20
-
-    # **** mesa ****
-    - name: Install mesa (linux)
-      if: inputs.mesa == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa_cpu-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
-    - name: Install mesa (macOS)
-      if: inputs.mesa == 'true' && runner.os == 'macOS'
-      shell: bash
-      run: brew install sirhcm/tinymesa/tinymesa_cpu
@@ -1,7 +1,10 @@
 name: Autogen
 env:
  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '13'
+  DOWNLOAD_CACHE_VERSION: '12'
+  PYTHON_CACHE_VERSION: '3'
+  APT_CACHE_VERSION: '1'
+  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
@@ -11,17 +14,15 @@ on:
    branches:
      - master
  pull_request:
-    paths:
+    paths: 
    - 'tinygrad/runtime/autogen/**/*'
-    - 'tinygrad/runtime/support/autogen.py'
  workflow_dispatch:
-    paths:
+    paths: 
    - 'tinygrad/runtime/autogen/**/*'
-    - 'tinygrad/runtime/support/autogen.py'

 jobs:
  autogen:
-    name: In-tree Autogen
+    name: Autogen
    runs-on: ubuntu-24.04
    timeout-minutes: 15
    steps:
@@ -33,144 +34,58 @@ jobs:
        opencl: 'true'
        amd: 'true'
        cuda: 'true'
-        llvm: 'true'
        webgpu: 'true'
-        mesa: 'true'
-        pydeps: 'pyyaml mako'
+        llvm: 'true'
    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
+      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
    - name: Verify OpenCL autogen
      run: |
-        mv tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
-        python3 -c "from tinygrad.runtime.autogen import opencl"
+        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
+        ./autogen_stubs.sh opencl
        diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
    - name: Verify CUDA autogen
      run: |
-        mv tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
-        mv tinygrad/runtime/autogen/nvrtc.py /tmp/nvrtc.py.bak
-        mv tinygrad/runtime/autogen/nvjitlink.py /tmp/nvjitlink.py.bak
-        mv tinygrad/runtime/autogen/nv_570.py /tmp/nv_570.py.bak
-        mv tinygrad/runtime/autogen/nv.py /tmp/nv.py.bak
-        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv"
+        cp tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
+        cp tinygrad/runtime/autogen/nv_gpu.py /tmp/nv_gpu.py.bak
+        ./autogen_stubs.sh cuda
+        ./autogen_stubs.sh nv
        diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
-        diff /tmp/nvrtc.py.bak tinygrad/runtime/autogen/nvrtc.py
-        diff /tmp/nvjitlink.py.bak tinygrad/runtime/autogen/nvjitlink.py
-        diff /tmp/nv_570.py.bak tinygrad/runtime/autogen/nv_570.py
-        diff /tmp/nv.py.bak tinygrad/runtime/autogen/nv.py
+        diff /tmp/nv_gpu.py.bak tinygrad/runtime/autogen/nv_gpu.py
    - name: Verify AMD autogen
      run: |
-        mv tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
-        mv tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
-        mv tinygrad/runtime/autogen/hip.py /tmp/hip.py.bak
-        mv tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
-        mv tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
-        mv tinygrad/runtime/autogen/rocprof.py /tmp/rocprof.py.bak
-        mv tinygrad/runtime/autogen/am/am.py /tmp/am_am.py.bak
-        mv tinygrad/runtime/autogen/am/pm4_soc15.py /tmp/am_pm4_soc15.py.bak
-        mv tinygrad/runtime/autogen/am/pm4_nv.py /tmp/am_pm4_nv.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_4_0_0.py /tmp/am_sdma_4_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_5_0_0.py /tmp/am_sdma_5_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_6_0_0.py /tmp/am_sdma_6_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/smu_v13_0_0.py /tmp/am_smu_v13_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/smu_v14_0_2.py /tmp/am_smu_v14_0_2.py.bak
-        python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof; from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v14_0_2"
-        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
+        cp tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
+        cp tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
+        cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
+        cp tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
+        cp tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
+        ./autogen_stubs.sh hsa
+        ./autogen_stubs.sh kfd
+        ./autogen_stubs.sh comgr
+        ./autogen_stubs.sh amd
+        ./autogen_stubs.sh sqtt
        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
-        diff /tmp/hip.py.bak tinygrad/runtime/autogen/hip.py
+        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
+        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
-        diff /tmp/rocprof.py.bak tinygrad/runtime/autogen/rocprof.py
-        diff /tmp/am_am.py.bak tinygrad/runtime/autogen/am/am.py
-        diff /tmp/am_pm4_soc15.py.bak tinygrad/runtime/autogen/am/pm4_soc15.py
-        diff /tmp/am_pm4_nv.py.bak tinygrad/runtime/autogen/am/pm4_nv.py
-        diff /tmp/am_sdma_4_0_0.py.bak tinygrad/runtime/autogen/am/sdma_4_0_0.py
-        diff /tmp/am_sdma_5_0_0.py.bak tinygrad/runtime/autogen/am/sdma_5_0_0.py
-        diff /tmp/am_sdma_6_0_0.py.bak tinygrad/runtime/autogen/am/sdma_6_0_0.py
-        diff /tmp/am_smu_v13_0_0.py.bak tinygrad/runtime/autogen/am/smu_v13_0_0.py
-        diff /tmp/am_smu_v14_0_2.py.bak tinygrad/runtime/autogen/am/smu_v14_0_2.py
    - name: Verify Linux autogen
      run: |
-        mv tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
-        mv tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
-        mv tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
-        mv tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
-        mv tinygrad/runtime/autogen/pci.py /tmp/pci.py.bak
-        mv tinygrad/runtime/autogen/vfio.py /tmp/vfio.py.bak
-        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
+        cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
+        cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
+        cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
+        ./autogen_stubs.sh libc
+        ./autogen_stubs.sh io_uring
+        ./autogen_stubs.sh ib
        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
-        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
-        diff /tmp/pci.py.bak tinygrad/runtime/autogen/pci.py
-        diff /tmp/vfio.py.bak tinygrad/runtime/autogen/vfio.py
-    - name: Verify LLVM autogen
-      run: |
-        mv tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
-        python3 -c "from tinygrad.runtime.autogen import llvm"
-        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
    - name: Verify WebGPU autogen
      run: |
-        mv tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
-        python3 -c "from tinygrad.runtime.autogen import webgpu"
+        cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
+        ./autogen_stubs.sh webgpu
        diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
-    - name: Verify Qualcomm autogen
+    - name: Verify LLVM autogen
      run: |
-        mv tinygrad/runtime/autogen/kgsl.py /tmp/kgsl.py.bak
-        mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
-        python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
-        diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
-        diff /tmp/qcom_dsp.py.bak tinygrad/runtime/autogen/qcom_dsp.py
-    - name: Verify libusb autogen
-      run: |
-        mv tinygrad/runtime/autogen/libusb.py /tmp/libusb.py.bak
-        python3 -c "from tinygrad.runtime.autogen import libusb"
-        diff /tmp/libusb.py.bak tinygrad/runtime/autogen/libusb.py
-    - name: Verify mesa autogen
-      run: |
-        mv tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
-        python3 -c "from tinygrad.runtime.autogen import mesa"
-        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
-    - name: Verify libclang autogen
-      run: |
-        cp tinygrad/runtime/autogen/libclang.py /tmp/libclang.py.bak
-        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
-        diff /tmp/libclang.py.bak tinygrad/runtime/autogen/libclang.py
-  autogen-mac:
-    name: In-tree Autogen (macos)
-    runs-on: macos-14
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        llvm: 'true'
-    - name: Verify macos autogen
-      run: |
-        mv tinygrad/runtime/autogen/metal.py /tmp/metal.py.bak
-        LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
-        diff /tmp/metal.py.bak tinygrad/runtime/autogen/metal.py
-  autogen-comgr-3:
-    name: In-tree Autogen (comgr 3)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-    - name: Install autogen support packages
-      run: |
-        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
-        sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
-        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.4 $(lsb_release -cs) main
-        EOF
-        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
-        sudo apt -qq update || true
-        sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
-    - name: Verify comgr (3) autogen
-      run: |
-        mv tinygrad/runtime/autogen/comgr_3.py /tmp/comgr_3.py.bak
-        python3 -c "from tinygrad.runtime.autogen import comgr_3"
-        diff /tmp/comgr_3.py.bak tinygrad/runtime/autogen/comgr_3.py
+        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
+        ./autogen_stubs.sh llvm
+        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
@@ -14,6 +14,12 @@ on:
      - update_benchmark
      - update_benchmark_staging
  workflow_dispatch:
+    inputs:
+      run_process_replay:
+        description: "Run process replay tests"
+        required: false
+        default: false
+        type: boolean

 jobs:
  testmacbenchmark:
@@ -33,7 +39,6 @@ jobs:
    - name: Symlink models and datasets
      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
@@ -46,20 +51,19 @@ jobs:
        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 | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - 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 | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=900 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - 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 | tee sdv2.txt
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=10000 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
    # 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 | tee sdxl.txt
+    # TODO: too slow
+    # - 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 | tee sdxl.txt
    - name: Run model inference benchmark
-      run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
+      run: METAL=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 | tee torch_speed.txt
    - name: Test tensor cores
@@ -119,6 +123,14 @@ jobs:
    # 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 | tee train_cifar_wino.txt
+    - name: UsbGPU boot time
+      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
+    - name: UsbGPU tiny tests
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
+    - name: UsbGPU copy speeds
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+    #- name: UsbGPU openpilot test
+    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -152,37 +164,6 @@ jobs:
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3.11 process_replay.py

-  testusbgpu:
-    name: UsbGPU Benchmark
-    env:
-      PYTHONPYCACHEPREFIX: /tmp/tiny_python_pycache
-    runs-on: [self-hosted, macOS]
-    timeout-minutes: 10
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - 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: UsbGPU boot time
-      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
-    - name: UsbGPU tiny tests
-      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
-    - name: UsbGPU copy speeds
-      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
-    - name: UsbGPU (USB4/TB) boot time
-      run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
-    - name: UsbGPU (USB4/TB) tiny tests
-      run: PYTHONPATH=. NV=1 NV_IFACE=PCI NV_NAK=1 python3.11 test/test_tiny.py
-
  testnvidiabenchmark:
    name: tinybox green Benchmark
    runs-on: [self-hosted, Linux, tinyboxgreen]
@@ -217,7 +198,7 @@ jobs:
    - name: Test speed vs torch
      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
-      run: NV=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: NV=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test benchmark allreduce
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
@@ -229,7 +210,6 @@ jobs:
        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
-        CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_fp8.txt
    - name: Run Tensor Core GEMM (PTX)
      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
    - name: Run Tensor Core GEMM (NV)
@@ -257,8 +237,6 @@ jobs:
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 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 | tee llama3_four_gpu.txt
-    - name: Run quantized LLaMA3
-      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8 | tee llama3_fp8.txt
    # - name: Run LLaMA-3 8B on 6 GPUs
    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
    # - name: Run LLaMA-2 70B
@@ -292,7 +270,6 @@ jobs:
          llama3_beam.txt
          llama3_four_gpu.txt
          llama3_six_gpu.txt
-          llama3_fp8.txt
          llama_2_70B.txt
          mixtral.txt
          gpt2_unjitted.txt
@@ -336,16 +313,15 @@ jobs:
    # TODO: too slow
    # - name: Fuzz Padded Tensor Core GEMM (PTX)
    #   run: NV=1 NV_PTX=1 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 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
    - name: Train MNIST
      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
@@ -354,13 +330,13 @@ jobs:
      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    - name: Run MLPerf resnet eval on training data
      run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA Training)
@@ -428,16 +404,14 @@ jobs:
    #    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 | tee torch_speed.txt
    - name: Test speed vs theoretical
-      run: AMD=1 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 AMD_LLVM=0
-      run: AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
-    # TODO: this is flaky
-    # - name: Test tensor cores AMD_LLVM=1
-    #   run: AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
-    - name: Run Tensor Core GEMM (AMD)
+      run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+    - name: Test tensor cores
      run: |
+        AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
+        AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
        AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-        AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
+    - name: Run Tensor Core GEMM (AMD)
+      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
    - name: Test AMD=1
      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
    #- name: Test HIP=1
@@ -453,8 +427,9 @@ jobs:
      run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Run Stable Diffusion
      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
-    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    # TODO: too slow
+    # - name: Run SDXL
+    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA 7B
      run: |
        BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -545,9 +520,9 @@ jobs:
    - name: Train MNIST
      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    # - name: Run 10 CIFAR training steps w BF16
    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    # TODO: too slow
@@ -555,8 +530,10 @@ jobs:
    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    #- name: Run full CIFAR training steps w 6 GPUS
+    #  run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    #- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
+    #  run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD Training)
@@ -568,6 +545,7 @@ jobs:
          train_cifar_wino.txt
          train_cifar_one_gpu.txt
          train_cifar_six_gpu.txt
+          train_cifar_six_gpu_remote.txt
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

@@ -605,13 +583,13 @@ jobs:
      run: test/external/process_replay/reset.py
    - name: Run MLPerf resnet eval
      run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD MLPerf)
@@ -640,20 +618,22 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: openpilot compile3 0.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.10.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
-      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
-      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.10.1 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: benchmark openpilot 0.9.9 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    - name: benchmark openpilot 0.9.9 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    - name: benchmark openpilot 0.9.9 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.9.9 driving_vision
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.9.9 driving_policy
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.9.9 dmonitoring
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 Space Lab policy + vision
+      run: |
+        PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
+        PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -661,9 +641,19 @@ jobs:
        ln -s /data/home/tiny/tinygrad/testsig-*.so .
        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
        # benchmark on DSP with NOOPT=1, the devectorizer has issues
-        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+        PYTHONPATH=. CC=clang-19 DSP=1 DONT_REALIZE_EXPAND=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
+    - uses: actions/upload-artifact@v4
+      with:
+        name: Speed (comma)
+        path: |
+          openpilot_compile_0_9_4.txt
+          openpilot_compile_0_9_7.txt
+          openpilot_0_9_4.txt
+          openpilot_0_9_7.txt
+          openpilot_image_0_9_4.txt
+          openpilot_image_0_9_7.txt

  testreddriverbenchmark:
    name: AM Benchmark
@@ -719,6 +709,7 @@ jobs:
        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. 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 AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+    # TODO: enable
    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
    #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
@@ -781,8 +772,8 @@ jobs:
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
    - name: Run full CIFAR training w 1 GPU
      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
+    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
@@ -22,13 +22,13 @@ jobs:
    - name: Run SDXL with new search
      # TODO: GCVM_L2_PROTECTION_FAULT_STATUS with llvm19
      run: |
-        BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CCACHE=0 python examples/sdxl.py --noshow --timing --seed 0
+        BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 python examples/sdxl.py --noshow --timing --seed 0
    - name: Run SDXL with cached search
      run: |
        BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. AMD=1 JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
    - name: Run winograd cifar with new search
      run: |
-        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 CCACHE=0 BS=1024 STEPS=500 python examples/hlb_cifar10.py
+        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BS=1024 STEPS=500 python examples/hlb_cifar10.py
    - name: Run winograd cifar with cached search
      run: |
        BENCHMARK_LOG=search_wino_cifar_cached WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py
@@ -12,7 +12,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 720
+    timeout-minutes: 360

    steps:
    - name: Checkout Code
@@ -27,4 +27,4 @@ jobs:
      run: |
        rm "~/.cache/tinygrad/cache_mlperf.db" || true
        BENCHMARK_LOG=mlpert_train_resnet LOGMLPERF=0 CACHEDB="~/.cache/tinygrad/cache_mlperf.db" examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-        rm "~/.cache/tinygrad/cache_mlperf.db"
+        rm "~/.cache/tinygrad/cache_mlperf.db"
@@ -20,11 +20,11 @@ jobs:
    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
-        pip install setuptools wheel build twine
+        pip install setuptools wheel twine
    - name: Build and publish
      env:
        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
      run: |
-        python -m build
+        python setup.py sdist bdist_wheel
        twine upload dist/*
@@ -56,15 +56,15 @@ jobs:
        uses: actions/checkout@v4
        with:
          path: base
-      - name: Set up Python 3.12
+      - name: Set up Python 3.10
        uses: actions/setup-python@v5
        with:
-          python-version: '3.12'
+          python-version: '3.10'
      - name: Count Line Diff
        run: |
+          pip install tabulate
          BASE="$GITHUB_WORKSPACE/base"
          PR="$GITHUB_WORKSPACE/pr"
-          pip install tabulate $BASE
          cp "$BASE/sz.py" .
          echo "loc_content<<EOF" >> "$GITHUB_ENV"
          python sz.py "$BASE" "$PR" >> "$GITHUB_ENV"
@@ -1,11 +1,13 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '15'
+  DOWNLOAD_CACHE_VERSION: '12'
+  PYTHON_CACHE_VERSION: '3'
+  APT_CACHE_VERSION: '1'
+  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
-  IGNORE_OOB: 0

 on:
  push:
@@ -37,8 +39,6 @@ jobs:
    name: Docs
    runs-on: ubuntu-22.04
    timeout-minutes: 10
-    env:
-      IGNORE_OOB: 1
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -74,7 +74,9 @@ jobs:
    - name: Test Docs Build
      run: python -m mkdocs build --strict
    - name: Test Docs
-      run: python docs/abstractions3.py
+      run: |
+        python docs/abstractions2.py
+        python docs/abstractions3.py
    - name: Test README
      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  python README.py
    - name: Test Quickstart
@@ -105,11 +107,15 @@ jobs:
      run: |
        sudo apt update || true
        sudo apt install -y --no-install-recommends ninja-build
+    - name: Lint with ruff
+      run: |
+        pip3 install --upgrade --force-reinstall ruff==0.11.0
+        python3 -m ruff check extra/torch_backend/backend.py
    - name: Test one op
      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
    - name: Test ResNet-18
      run: DEBUG=2 python3 extra/torch_backend/example.py
-    - name: custom tests
+    - name: My (custom) tests
      run: python3 extra/torch_backend/test.py
    - name: Test one op in torch tests
      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
@@ -119,9 +125,6 @@ jobs:
      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
    - name: Test multi-gpu
      run: CPU=1 CPU_LLVM=1 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
@@ -141,7 +144,7 @@ jobs:
        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 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+      run: CPU=1 CPU_LLVM=1 TARGET_EVAL_ACC_PCT=96.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

@@ -200,7 +203,7 @@ jobs:
        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
+        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated INTEL OpenCL tensor cores
      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
    - name: Test emulated AMX tensor cores
@@ -226,17 +229,17 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: linting-only
-        python-version: '3.11'
+        python-version: '3.10'
        deps: linting
    - name: Lint bad-indentation and trailing-whitespace with pylint
      run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' --recursive=y .
    - name: Lint with ruff
      run: |
-        pip3 install --upgrade --force-reinstall ruff==0.14.10
+        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check .
        python3 -m ruff check examples/mlperf/ --ignore E501
-        python3 -m ruff check extra/thunder/tiny/ --ignore E501 --ignore F841 --ignore E722
-        python3 -m ruff check extra/torch_backend/backend.py
+    - name: Lint tinygrad with pylint
+      run: python -m pylint tinygrad/
    - name: Run mypy
      run: |
        python -m mypy --strict-equality --lineprecision-report .
@@ -261,9 +264,9 @@ jobs:
    - name: Check Device.DEFAULT
      run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
    - name: Run unit tests
-      run: |
-        CPU=1 python test/unit/test_device.py TestRunAsModule.test_module_runs
-        CPU=1 python -m pytest -n=auto test/unit/ --durations=20 --deselect=test/unit/test_device.py::TestRunAsModule::test_module_runs
+      run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Check SPEC=1
+      run: SPEC=1 python3 test/test_tiny.py
    - name: Run targetted tests on NULL backend
      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
    # TODO: too slow
@@ -271,8 +274,6 @@ jobs:
    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
    - name: Run Clip tests for SD MLPerf on NULL backend
      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
-    - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    # TODO: support fake weights
    #- name: Run LLaMA 7B on 4 fake devices
    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -289,28 +290,8 @@ jobs:
        python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
        #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 20000 lines
-      run: MAX_LINE_COUNT=20000 python sz.py
-
-  spec:
-    strategy:
-      fail-fast: false
-      matrix:
-        group: [1, 2]
-    name: SPEC=2 (${{ matrix.group }})
-    runs-on: ubuntu-latest
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: spec-unit
-        deps: testing_unit
-        python-version: '3.14'
-    - name: Test SPEC=2
-      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+    - name: Repo line count < 18000 lines
+      run: MAX_LINE_COUNT=18000 python sz.py

  fuzzing:
    name: Fuzzing
@@ -326,10 +307,12 @@ jobs:
        deps: testing_unit
    - name: Fuzz Test symbolic
      run: python test/external/fuzz_symbolic.py
-    - name: Fuzz Test symbolic (symbolic divisors)
-      run: python test/external/fuzz_symbolic_symbolic_div.py
    - name: Fuzz Test fast idiv
      run: python test/external/fuzz_fast_idiv.py
+    - name: Fuzz Test shapetracker
+      run: CNT=50 python test/external/fuzz_shapetracker.py
+    - name: Fuzz Test shapetracker math
+      run: CNT=200 python test/external/fuzz_shapetracker_math.py
    - name: Fuzz Test shape ops
      run: python test/external/fuzz_shape_ops.py

@@ -346,11 +329,10 @@ jobs:
          key: gpu-image
          deps: testing_minimal
          opencl: 'true'
-      - name: Test CL IMAGE=2 ops
+      - name: Test CL IMAGE=2 ops + training
        run: |
          CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
-        # TODO: training is broken
-        # CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
+          CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -372,7 +354,7 @@ jobs:
      - name: Run Kernel Count Test
        run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
      - name: Run fused optimizer tests
-        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/test_optim.py -k "not muon"
+        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py
      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
@@ -395,13 +377,17 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
-      - name: Test openpilot CL compile fp16
-        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 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: DEBUGCL=1 CL=1 IMAGE=2 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
-      - name: Test openpilot LLVM compile fp16
-        run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2081 ALLOWED_GATED_READ_IMAGE=28 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+      - name: Test openpilot alt model correctness (float32)
+        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
+      - name: Test openpilot fastvits model correctness (float32)
+        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      # - name: Test openpilot simple_plan vision model correctness (float32)
+      #   run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
+      - name: Test openpilot LLVM compile
+        run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      - name: Test openpilot compile4
+        run: NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -447,7 +433,7 @@ jobs:
        with:
          key: onnxoptl
          deps: testing
-          pydeps: "tensorflow==2.19"
+          pydeps: "tensorflow==2.15.1 tensorflow_addons"
          python-version: '3.11'
          opencl: 'true'
      - name: Test ONNX (CL)
@@ -465,7 +451,7 @@ jobs:
      - name: Test Bert training
        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
      - name: Test llama 3 training
-        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -473,8 +459,6 @@ jobs:
    name: Test LLM
    runs-on: ubuntu-24.04
    timeout-minutes: 15
-    env:
-      IGNORE_OOB: 1
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -545,11 +529,11 @@ jobs:
        pydeps: "pillow"
        llvm: "true"
    - name: Test LLVM=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
    - name: Test LLVM=1 DEVECTORIZE=0 for model
      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
    - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
+      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"

  testdsp:
    name: Linux (DSP)
@@ -643,66 +627,18 @@ jobs:
        if: matrix.backend=='amdllvm'
        run: python test/device/test_amd_llvm.py
      - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py --durations=20
+        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py --durations=20
      - name: Run pytest (amd)
        run: python -m pytest test/external/external_test_am.py --durations=20
      - name: Run TRANSCENDENTAL math
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run TestOps.test_add with SQTT
        run: |
-          VIZ=1 PMC=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
          VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-  testamdasm:
-    name: AMD ASM IDE
-    runs-on: ubuntu-24.04
-    timeout-minutes: 10
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: rdna3-emu
-          deps: testing_minimal
-          amd: 'true'
-      - name: Install LLVM 21
-        run: |
-          wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
-          echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
-          sudo apt-get update
-          sudo apt-get install llvm-21 llvm-21-tools cloc
-      - name: RDNA3 Line Count
-        run: cloc --by-file extra/assembly/amd/*.py
-      - name: Run RDNA3 emulator tests
-        run: python -m pytest -n=auto extra/assembly/amd/ --durations 20
-      - name: Run RDNA3 emulator tests (AMD_LLVM=1)
-        run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
-      - name: Run RDNA3 dtype tests
-        run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
-      - name: Run RDNA3 dtype tests (AMD_LLVM=1)
-        run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
-
-  testamdautogen:
-    name: AMD autogen
-    runs-on: ubuntu-24.04
-    timeout-minutes: 10
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: rdna3-autogen
-          pydeps: "pdfplumber"
-      - name: Verify AMD autogen is up to date
-        run: |
-          python -m extra.assembly.amd.pdf --arch all
-          git diff --exit-code extra/assembly/amd/autogen/
-
  testnvidia:
    strategy:
      fail-fast: false
@@ -741,7 +677,7 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu, opencl, lvp]
+        backend: [llvm, cpu, opencl]

    name: Linux (${{ matrix.backend }})
    runs-on: ubuntu-22.04
@@ -755,10 +691,9 @@ jobs:
          key: ${{ matrix.backend }}-minimal
          deps: testing_minimal
          opencl: ${{ matrix.backend == 'opencl' && 'true' }}
-          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
-          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
@@ -770,6 +705,71 @@ jobs:
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

+  amdremote:
+    name: Linux (remote)
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    env:
+      REMOTE: 1
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: linux-remote
+          deps: testing_minimal
+          amd: 'true'
+          llvm: 'true'
+          opencl: 'true'
+      - name: Start remote server
+        run: |
+          start_server() {
+            systemd-run --user \
+              --unit="$1" \
+              --setenv=REMOTEDEV="$2" \
+              --setenv=MOCKGPU=1 \
+              --setenv=PYTHONPATH=. \
+              --setenv=PORT="$3" \
+              --working-directory="$(pwd)" \
+              python tinygrad/runtime/ops_remote.py
+          }
+
+          start_server "remote-server-amd-1" "AMD" 6667
+          start_server "remote-server-amd-2" "AMD" 6668
+          start_server "remote-server-gpu" "CL" 7667
+          start_server "remote-server-cpu" "CPU" 8667
+      - name: Check Device.DEFAULT and print some source
+        env:
+          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
+          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run REMOTE=1 Test (AMD)
+        env:
+          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py --durations 20
+      - name: Run REMOTE=1 Test (CL)
+        env:
+          HOST: 127.0.0.1:7667*6
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py --durations 20
+          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
+      - name: Run REMOTE=1 Test (CPU)
+        env:
+          HOST: 127.0.0.1:8667*6
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py --durations 20
+      - name: Show remote server logs
+        if: always()
+        run: |
+          journalctl --user -u remote-server-amd-1 --no-pager
+          journalctl --user -u remote-server-amd-2 --no-pager
+          journalctl --user -u remote-server-gpu --no-pager
+          journalctl --user -u remote-server-cpu --no-pager
+
 # ****** OSX Tests ******

  testmetal:
@@ -867,11 +867,35 @@ jobs:
    - name: Test ONNX Runner (WEBGPU)
      run: WEBGPU=1 python3 test/external/external_test_onnx_runner.py

+  osxremote:
+   name: MacOS (remote metal)
+   runs-on: macos-15
+   timeout-minutes: 10
+   env:
+     REMOTE: 1
+     REMOTEDEV: METAL
+   steps:
+     - name: Checkout Code
+       uses: actions/checkout@v4
+     - name: Setup Environment
+       uses: ./.github/actions/setup-tinygrad
+       with:
+         key: macos-remote
+         deps: testing_minimal
+     - name: Check Device.DEFAULT and print some source
+       run: |
+         python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
+         python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
+         DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+     - name: Run REMOTE=1 Test
+       run: |
+         python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
+
  osxtests:
    strategy:
      fail-fast: false
      matrix:
-        backend: [metal, llvm, cpu, lvp]
+        backend: [metal, llvm, cpu]
    name: MacOS (${{ matrix.backend }})
    runs-on: macos-15
    timeout-minutes: 20
@@ -884,13 +908,12 @@ jobs:
          key: macos-${{ matrix.backend }}-minimal
          deps: testing_minimal
          pydeps: "capstone"
-          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
-          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (${{ matrix.backend }})
        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
@@ -932,33 +955,3 @@ jobs:
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
-
-# ****** Compile-only Tests ******
-
-  compiletests:
-    strategy:
-      fail-fast: false
-      matrix:
-        backend: [ir3, nak]
-    name: Compile-only (${{ matrix.backend }})
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: compile-${{ matrix.backend }}
-          deps: testing_minimal
-          mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
-          python-version: '3.14'
-      - name: Set env
-        shell: bash
-        run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
-      - name: Run test_ops
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/test_ops.py --durations=20
@@ -38,7 +38,6 @@ extra/huggingface_onnx/models/*
 extra/huggingface_onnx/*.yaml
 extra/weights
 venv
-venv_sd_mlperf
 examples/**/net.*[js,json]
 examples/**/*.safetensors
 node_modules
@@ -63,5 +62,3 @@ profile_stats
 *.log
 target
 .mypy_cache
-mutants
-.mutmut-cache
@@ -20,15 +20,21 @@ repos:
        language: system
        always_run: true
        pass_filenames: false
+      - id: tests
+        name: subset of tests
+        entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
+        language: system
+        always_run: true
+        pass_filenames: false
      - id: example
-        name: test all devices
+        name: multi device tests
        entry: python3 test/external/external_test_example.py
        language: system
        always_run: true
        pass_filenames: false
-      - id: tests
-        name: comprehensive test suite
-        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
+      - id: pylint
+        name: pylint
+        entry: python3 -m pylint tinygrad/
        language: system
        always_run: true
-        pass_filenames: false
+        pass_filenames: false
@@ -1,216 +0,0 @@
-# Claude Code Guide for tinygrad
-
-## Architecture Overview
-
-tinygrad compiles tensor operations into optimized kernels. The pipeline:
-
-1. **Tensor** (`tensor.py`) - User-facing API, creates UOp graph
-2. **UOp** (`uop/ops.py`) - Unified IR for all operations (both tensor and kernel level)
-3. **Schedule** (`engine/schedule.py`, `schedule/`) - Converts tensor UOps to kernel UOps
-4. **Codegen** (`codegen/`) - Converts kernel UOps to device code
-5. **Runtime** (`runtime/`) - Device-specific execution
-
-## Key Concepts
-
-### UOp (Universal Operation)
-Everything is a UOp - tensors, operations, buffers, kernels. Key properties:
- `op`: The operation type (Ops enum)
- `dtype`: Data type
- `src`: Tuple of source UOps
- `arg`: Operation-specific argument
- `tag`: Optional tag for graph transformations
-
-UOps are **immutable and cached** - creating the same UOp twice returns the same object (ucache).
-
-### PatternMatcher
-Used extensively for graph transformations:
-```python
-pm = PatternMatcher([
-  (UPat(Ops.ADD, src=(UPat.cvar("x"), UPat.cvar("x"))), lambda x: x * 2),
-])
-result = graph_rewrite(uop, pm)
-```
-
-### Schedule Cache
-Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.
-
-## Testing
-
-```bash
-# Run specific test
-python -m pytest test/unit/test_schedule_cache.py -xvs
-
-# Run with timeout
-python -m pytest test/test_symbolic_ops.py -x --timeout=60
-
-# Debug with print
-DEBUG=2 python -m pytest test/test_schedule.py::test_name -xvs
-
-# Visualize UOp graphs
-VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
-```
-
-## Common Environment Variables
-
- `DEBUG=1-7` - Increasing verbosity (7 shows assembly output)
- `VIZ=1` - Enable graph visualization
- `SPEC=1` - Enable UOp spec verification
- `NOOPT=1` - Disable optimizations
- `DEVICE=CPU/CUDA/AMD/METAL` - Set default device
-
-## Debugging Tips
-
-1. **Print UOp graphs**: `print(tensor.uop)` or `print(tensor.uop.sink())`
-2. **Check schedule**: `tensor.schedule()` returns list of ExecItems
-3. **Trace graph rewrites**: Use `VIZ=1` or add print in PatternMatcher callbacks
-4. **Find UOps by type**: `[u for u in uop.toposort() if u.op is Ops.SOMETHING]`
-
-## Workflow Rules
-
- **NEVER commit without explicit user approval** - always show the diff and wait for approval
- **NEVER amend commits** - always create a new commit instead
- Run `pre-commit run --all-files` before committing to catch linting/type errors
- Run tests before proposing commits
- Test with `SPEC=2` when modifying UOp-related code
-
-## Auto-generated Files (DO NOT EDIT)
-
-The following files are auto-generated and should never be edited manually:
- `extra/assembly/amd/autogen/{arch}/__init__.py` - Generated by `python -m extra.assembly.amd.dsl --arch {arch}`
- `extra/assembly/amd/autogen/{arch}/gen_pcode.py` - Generated by `python -m extra.assembly.amd.pcode --arch {arch}`
-
-Where `{arch}` is one of: `rdna3`, `rdna4`, `cdna`
-
-To add missing instruction implementations, add them to `extra/assembly/amd/emu.py` instead.
-
-## Style Notes
-
- 2-space indentation, 150 char line limit
- PatternMatchers should be defined at module level (slow to construct)
- Prefer `graph_rewrite` over manual graph traversal
- UOp methods like `.replace()` preserve tags unless explicitly changed
- Use `.rtag(value)` to add tags to UOps
-
-## Lessons Learned
-
-### UOp ucache Behavior
-UOps are cached by their contents - creating a UOp with identical (op, dtype, src, arg) returns the **same object**. This means:
- `uop.replace(tag=None)` on a tagged UOp returns the original untagged UOp if it exists in cache
- Two UOps with same structure are identical (`is` comparison works)
-
-### Spec Validation
-When adding new UOp patterns, update `tinygrad/uop/spec.py`. Test with:
-```bash
-SPEC=2 python3 test/unit/test_something.py
-```
-Spec issues appear as `RuntimeError: SPEC ISSUE None: UOp(...)`.
-
-### Schedule Cache Key Normalization
-The schedule cache strips values from BIND nodes so different bound values (e.g., KV cache positions) hit the same cache entry:
- `pm_pre_sched_cache`: BIND(DEFINE_VAR, CONST) → BIND(DEFINE_VAR) for cache key
- `pm_post_sched_cache`: restores original BIND from context
- When accessing `bind.src[1]`, check `len(bind.src) > 1` first (might be stripped)
- Extract var_vals from `input_buffers` dict after graph_rewrite (avoids extra toposort)
-
-### Avoiding Extra Work
- Use ctx dict from graph_rewrite to collect info during traversal instead of separate toposort
- Only extract var_vals when schedule is non-empty (no kernels = no vars needed)
- PatternMatchers are slow to construct - define at module level, not in functions
-
-### Readability Over Speed
-Don't add complexity for marginal performance gains. Simpler code that's slightly slower is often better:
-```python
-# BAD: "optimized" with extra complexity
-if has_afters:  # skip toposort if no AFTERs
-  after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
-
-# GOOD: simple, always works
-after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
-```
-The conditional check adds complexity, potential bugs, and often negligible speedup. Only optimize when profiling shows a real bottleneck.
-
-### Testing LLM Changes
-```bash
-# Quick smoke test
-echo "Hello" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b"
-
-# Check cache hits (should see "cache hit" after warmup)
-echo "Hello world" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b" 2>&1 | grep cache
-
-# Test with beam search
-echo "Hello" | BEAM=2 python tinygrad/apps/llm.py --model "llama3.2:1b"
-```
-
-## Common Patterns
-
-### Graph Transformation
-```python
-def my_transform(ctx, x):
-  # Return new UOp or None to skip
-  return x.replace(arg=new_arg)
-
-pm = PatternMatcher([
-  (UPat(Ops.SOMETHING, name="x"), my_transform),
-])
-result = graph_rewrite(input_uop, pm, ctx={})
-```
-
-### Finding Variables
-```python
-# Get all variables in a UOp graph
-variables = uop.variables()
-
-# Get bound variable values
-var, val = bind_uop.unbind()
-```
-
-### Shape Handling
-```python
-# Shapes can be symbolic (contain UOps)
-shape = tensor.shape  # tuple[sint, ...] where sint = int | UOp
-```
-
-## Performance Optimization
-
-When optimizing tinygrad internals:
-
-1. **Measure wall time, not just call counts** - Reducing `graph_rewrite` calls doesn't always improve wall time. The overhead of conditional checks can exceed the cost of the operation being skipped.
-
-2. **Profile each optimization individually** - Run benchmarks with and without each change to measure actual impact. Use `test/external/external_benchmark_schedule.py` for schedule/rewrite timing.
-
-3. **Early exits in hot paths are effective** - Simple checks like `if self.op is Ops.CONST: return self` in `simplify()` can eliminate many unnecessary `graph_rewrite` calls.
-
-4. **`graph_rewrite` is expensive** - Each call has overhead even for small graphs. Avoid calling it when the result is trivially known (e.g., simplifying a CONST returns itself).
-
-5. **Beware iterator overhead** - Checks like `all(x.op is Ops.CONST for x in self.src)` can be slower than just running the operation, especially for small sequences.
-
-6. **Verify cache hit rates before adding/keeping caches** - Measure actual hit rates with real workloads. A cache with 0% hit rate is pure overhead (e.g., `pm_cache` was removed because the algorithm guarantees each UOp is only passed to `pm_rewrite` once).
-
-7. **Use `TRACK_MATCH_STATS=2` to profile pattern matching** - This shows match rates and time per pattern. Look for patterns with 0% match rate that still cost significant time - these are pure overhead for that workload.
-
-8. **Cached properties beat manual traversal** - `backward_slice` uses `@functools.cached_property`. A DFS with early-exit sounds faster but is actually slower because it doesn't benefit from caching. The cache hit benefit often outweighs algorithmic improvements.
-
-9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.
-
-## Pattern Matching Profiling
-
-Use `TRACK_MATCH_STATS=2` to identify expensive patterns:
-
-```bash
-TRACK_MATCH_STATS=2 PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
-```
-
-Output format: `matches / attempts -- match_time / total_time ms -- location`
-
-Key patterns to watch (from ResNet50 benchmark):
- `split_load_store`: ~146ms, 31% match rate - does real work
- `simplify_valid`: ~75ms, 0% match rate in this workload - checks AND ops for INDEX in backward slice
- `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches
-
-Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.
-
-## AMD Performance Counter Profiling
-
-Set VIZ to `-2` to save performance counters traces for the AMD backend.
-
-Use the CLI in `./extra/sqtt/roc.py` to explore the trace.
@@ -21,38 +21,17 @@ tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) an

 ---

-tinygrad is an end-to-end deep learning stack:
+Despite tinygrad's size, it is a fully featured deep learning framework.

- **Tensor library** with autograd
- **IR and compiler** that fuse and lower kernels
- **JIT + graph execution**
- **nn / optim / datasets** for real training
+Due to its extreme simplicity, it is the easiest framework to add new accelerators to, with support for both inference and training. If XLA is CISC, tinygrad is RISC.

-It’s inspired by PyTorch (ergonomics), JAX (functional transforms and IR-based AD), and TVM (scheduling and codegen), but stays intentionally tiny and hackable.
+tinygrad is now beta software, we [raised some money](https://geohot.github.io/blog/jekyll/update/2023/05/24/the-tiny-corp-raised-5M.html) to make it good. Someday, we will tape out chips.

---
+## Features

-## How tinygrad compares
+### LLaMA and Stable Diffusion

-**PyTorch**
-
- ✅ Similar: eager `Tensor` API, autograd, `optim`, basic datasets and layers.
- ✅ You can write familiar training loops.
- 🔁 Unlike PyTorch, the entire compiler and IR are visible and hackable.
-
-**JAX**
-
- ✅ IR-based autodiff over primitives (like JAXPR + XLA).
- ✅ Function-level JIT (`TinyJit`) that captures and replays kernels.
- 🔁 Fewer functional transforms (no full `vmap`/`pmap` yet), but far easier to read.
-
-**TVM**
-
- ✅ Multiple lowering passes, scheduling, and BEAM search over kernels.
- ✅ Device “graphs” for batched execution.
- 🔁 tinygrad also ships the **front-end framework** (tensors, nn, optim), not just the compiler.
-
---
+tinygrad can run [LLaMA](/docs/showcase.md#llama) and [Stable Diffusion](/docs/showcase.md#stable-diffusion)!

 ### Laziness

@@ -0,0 +1,489 @@
+#!/bin/bash -e
+
+# setup instructions for clang2py
+if [[ ! $(clang2py -V) ]]; then
+  pushd .
+  cd /tmp
+  sudo apt-get install -y --no-install-recommends clang
+  pip install --upgrade pip setuptools
+  pip install clang==14.0.6
+  git clone https://github.com/nimlgen/ctypeslib.git
+  cd ctypeslib
+  pip install .
+  clang2py -V
+  popd
+fi
+
+BASE=tinygrad/runtime/autogen/
+
+fixup() {
+  sed -i '1s/^/# mypy: ignore-errors\n/' $1
+  sed -i 's/ *$//' $1
+  grep FIXME_STUB $1 || true
+}
+
+patch_dlopen() {
+  path=$1; shift
+  name=$1; shift
+  cat <<EOF | sed -i "/import ctypes.*/r /dev/stdin" $path
+PATHS_TO_TRY = [
+$(for p in "$@"; do echo "  $p,"; done)
+]
+def _try_dlopen_$name():
+  library = ctypes.util.find_library("$name")
+  if library: return ctypes.CDLL(library)
+  for candidate in PATHS_TO_TRY:
+    try: return ctypes.CDLL(candidate)
+    except OSError: pass
+  return None
+EOF
+}
+
+generate_opencl() {
+  clang2py /usr/include/CL/cl.h -o $BASE/opencl.py -l /usr/lib/x86_64-linux-gnu/libOpenCL.so.1 -k cdefstum
+  fixup $BASE/opencl.py
+  # hot patches
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/opencl.py
+  sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libOpenCL.so.1')\ctypes.CDLL(ctypes.util.find_library('OpenCL'))\g" $BASE/opencl.py
+  python3 -c "import tinygrad.runtime.autogen.opencl"
+}
+
+generate_hip() {
+  clang2py /opt/rocm/include/hip/hip_ext.h /opt/rocm/include/hip/hiprtc.h \
+  /opt/rocm/include/hip/hip_runtime_api.h /opt/rocm/include/hip/driver_types.h \
+  --clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o $BASE/hip.py -l /opt/rocm/lib/libamdhip64.so
+  echo "hipDeviceProp_t = hipDeviceProp_tR0600" >> $BASE/hip.py
+  echo "hipGetDeviceProperties = hipGetDevicePropertiesR0600" >> $BASE/hip.py
+  fixup $BASE/hip.py
+  # we can trust HIP is always at /opt/rocm/lib
+  #sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/hip.py
+  #sed -i "s\ctypes.CDLL('/opt/rocm/lib/libhiprtc.so')\ctypes.CDLL(ctypes.util.find_library('hiprtc'))\g" $BASE/hip.py
+  #sed -i "s\ctypes.CDLL('/opt/rocm/lib/libamdhip64.so')\ctypes.CDLL(ctypes.util.find_library('amdhip64'))\g" $BASE/hip.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/hip.py
+  sed -i "s\'/opt/rocm/\os.getenv('ROCM_PATH', '/opt/rocm/')+'/\g" $BASE/hip.py
+  python3 -c "import tinygrad.runtime.autogen.hip"
+}
+
+generate_comgr() {
+  clang2py /opt/rocm/include/amd_comgr/amd_comgr.h \
+  --clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o $BASE/comgr.py -l /opt/rocm/lib/libamd_comgr.so
+  fixup $BASE/comgr.py
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/comgr.py
+  patch_dlopen $BASE/comgr.py amd_comgr "'/opt/rocm/lib/libamd_comgr.so'" "os.getenv('ROCM_PATH', '')+'/lib/libamd_comgr.so'" "'/usr/local/lib/libamd_comgr.dylib'" "'/opt/homebrew/lib/libamd_comgr.dylib'"
+  sed -i "s\ctypes.CDLL('/opt/rocm/lib/libamd_comgr.so')\_try_dlopen_amd_comgr()\g" $BASE/comgr.py
+  python3 -c "import tinygrad.runtime.autogen.comgr"
+}
+
+generate_kfd() {
+  clang2py /usr/include/linux/kfd_ioctl.h -o $BASE/kfd.py -k cdefstum
+
+  fixup $BASE/kfd.py
+  sed -i "s/import ctypes/import ctypes, os/g" $BASE/kfd.py
+  sed -i "s/import fcntl, functools/import functools/g" $BASE/kfd.py
+  sed -i "/import functools/a from tinygrad.runtime.support.hcq import FileIOInterface" $BASE/kfd.py
+  sed -i "s/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd, \*\*kwargs):/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd:FileIOInterface, \*\*kwargs):/g" $BASE/kfd.py
+  sed -i "s/fcntl.ioctl(__fd, (__idir<<30)/__fd.ioctl((__idir<<30)/g" $BASE/kfd.py
+  sed -i "s/!!/not not /g" $BASE/kfd.py
+  python3 -c "import tinygrad.runtime.autogen.kfd"
+}
+
+generate_cuda() {
+  clang2py /usr/include/cuda.h --clang-args="-D__CUDA_API_VERSION_INTERNAL" -o $BASE/cuda.py -l /usr/lib/x86_64-linux-gnu/libcuda.so
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/cuda.py
+  sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libcuda.so')\ctypes.CDLL(ctypes.util.find_library('cuda'))\g" $BASE/cuda.py
+  fixup $BASE/cuda.py
+  python3 -c "import tinygrad.runtime.autogen.cuda"
+}
+
+generate_nvrtc() {
+  clang2py /usr/local/cuda/include/nvrtc.h /usr/local/cuda/include/nvJitLink.h -o $BASE/nvrtc.py -l /usr/local/cuda/lib64/libnvrtc.so -l /usr/local/cuda/lib64/libnvJitLink.so
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/nvrtc.py
+  sed -i "s\ctypes.CDLL('/usr/local/cuda/lib64/libnvrtc.so')\ctypes.CDLL(ctypes.util.find_library('nvrtc'))\g" $BASE/nvrtc.py
+  sed -i "s\ctypes.CDLL('/usr/local/cuda/lib64/libnvJitLink.so')\ctypes.CDLL(ctypes.util.find_library('nvJitLink'))\g" $BASE/nvrtc.py
+  fixup $BASE/nvrtc.py
+  python3 -c "import tinygrad.runtime.autogen.nvrtc"
+}
+
+generate_nv() {
+  NVKERN_COMMIT_HASH=81fe4fb417c8ac3b9bdcc1d56827d116743892a5
+  NVKERN_SRC=/tmp/open-gpu-kernel-modules-$NVKERN_COMMIT_HASH
+  if [ ! -d "$NVKERN_SRC" ]; then
+    git clone https://github.com/NVIDIA/open-gpu-kernel-modules $NVKERN_SRC
+    pushd .
+    cd $NVKERN_SRC
+    git reset --hard $NVKERN_COMMIT_HASH
+    popd
+  fi
+
+  clang2py -k cdefstum \
+    extra/nv_gpu_driver/clc6c0qmd.h \
+    extra/nv_gpu_driver/clcec0qmd.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0000.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0080.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080_notification.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc56f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc86f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc96f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc761.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl83de.h \
+    $NVKERN_SRC/src/nvidia/generated/g_allclasses.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc6c0.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clcdc0.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/clc6b5.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/clc9b5.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/uvm_ioctl.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/uvm_linux_ioctl.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/hwref/ampere/ga100/dev_fault.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv_escape.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl-numbers.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl-numa.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-unix-nvos-params-wrappers.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/alloc/alloc_channel.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/nvos.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl0000/*.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl0080/*.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl2080/*.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl83de/*.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlc36f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlcb33.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrla06c.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl90f1.h \
+    --clang-args="-include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
+    -o $BASE/nv_gpu.py
+  fixup $BASE/nv_gpu.py
+  sed -i "s\(0000000001)\1\g" $BASE/nv_gpu.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/nv_gpu.py
+  sed -i 's/#\?\s\([A-Za-z0-9_]\+\) = MW ( \([0-9]\+\) : \([0-9]\+\) )/\1 = (\2 , \3)/' $BASE/nv_gpu.py # NVC6C0_QMDV03_00 processing
+  sed -i 's/#\sdef NVC6C0_QMD\([A-Za-z0-9_()]\+\):/def NVC6C0_QMD\1:/' $BASE/nv_gpu.py
+  sed -i 's/#\sdef NVCEC0_QMD\([A-Za-z0-9_()]\+\):/def NVCEC0_QMD\1:/' $BASE/nv_gpu.py
+  sed -E -i -n '/^def (NVCEC0_QMDV05_00_RELEASE)(_ENABLE)\(i\):/{p;s//\1'"0"'\2=\1\2(0)\n\1'"1"'\2=\1\2(1)/;H;b};p;${x;s/^\n//;p}' "$BASE/nv_gpu.py"
+  sed -i 's/#\s*return MW(\([0-9i()*+]\+\):\([0-9i()*+]\+\))/    return (\1 , \2)/' $BASE/nv_gpu.py
+  sed -i 's/#\?\s*\(.*\)\s*=\s*\(NV\)\?BIT\(32\)\?\s*(\s*\([0-9]\+\)\s*)/\1 = (1 << \4)/' $BASE/nv_gpu.py # name = BIT(x) -> name = (1 << x)
+  sed -i "s/UVM_\([A-Za-z0-9_]\+\) = \['i', '(', '\([0-9]\+\)', ')'\]/UVM_\1 = \2/" $BASE/nv_gpu.py # UVM_name = ['i', '(', '<num>', ')'] -> UVM_name = <num>
+
+  # Parse status codes
+  sed -n '1i\
+nv_status_codes = {}
+/^NV_STATUS_CODE/ { s/^NV_STATUS_CODE(\([^,]*\), *\([^,]*\), *"\([^"]*\)") *.*$/\1 = \2\nnv_status_codes[\1] = "\3"/; p }' $NVKERN_SRC/src/common/sdk/nvidia/inc/nvstatuscodes.h >> $BASE/nv_gpu.py
+  python3 -c "import tinygrad.runtime.autogen.nv_gpu"
+
+  clang2py -k cdefstum \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/fsp/kern_fsp_cot_payload.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gspifpub.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_wpr_meta.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_sr_meta.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/common/uproc/os/common/include/libos_init_args.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/rmRiscvUcode.h \
+    $NVKERN_SRC/src/common/shared/msgq/inc/msgq/msgq_priv.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_headers.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_global_enums.h \
+    $NVKERN_SRC/src/nvidia/generated/g_rpc-structures.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/fsp/fsp_nvdm_format.h \
+    extra/nv_gpu_driver/g_rpc-message-header.h \
+    extra/nv_gpu_driver/gsp_static_config.h \
+    extra/nv_gpu_driver/vbios.h \
+    --clang-args="-DRPC_MESSAGE_STRUCTURES -DRPC_STRUCTURES -include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/nvidia/generated -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/src/nvidia/inc -I$NVKERN_SRC/src/nvidia/interface/ -I$NVKERN_SRC/src/nvidia/inc/kernel -I$NVKERN_SRC/src/nvidia/inc/libraries -I$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
+    -o $BASE/nv/nv.py
+
+  fixup $BASE/nv/nv.py
+  python3 -c "import tinygrad.runtime.autogen.nv.nv"
+}
+
+generate_amd() {
+  # clang2py broken when pass -x c++ to prev headers
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    extra/hip_gpu_driver/nvd.h \
+    extra/hip_gpu_driver/gc_11_0_0_offset.h \
+    extra/hip_gpu_driver/sienna_cichlid_ip_offset.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/amd_gpu.py
+
+  fixup $BASE/amd_gpu.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/amd_gpu.py
+  python3 -c "import tinygrad.runtime.autogen.amd_gpu"
+}
+
+generate_hsa() {
+  clang2py \
+    /opt/rocm/include/hsa/hsa.h \
+    /opt/rocm/include/hsa/hsa_ext_amd.h \
+    /opt/rocm/include/hsa/amd_hsa_signal.h \
+    /opt/rocm/include/hsa/amd_hsa_queue.h \
+    /opt/rocm/include/hsa/amd_hsa_kernel_code.h \
+    /opt/rocm/include/hsa/hsa_ext_finalize.h /opt/rocm/include/hsa/hsa_ext_image.h \
+    /opt/rocm/include/hsa/hsa_ven_amd_aqlprofile.h \
+    --clang-args="-I/opt/rocm/include" \
+    -o $BASE/hsa.py -l /opt/rocm/lib/libhsa-runtime64.so
+
+  fixup $BASE/hsa.py
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/hsa.py
+  sed -i "s\ctypes.CDLL('/opt/rocm/lib/libhsa-runtime64.so')\ctypes.CDLL(os.getenv('ROCM_PATH')+'/lib/libhsa-runtime64.so' if os.getenv('ROCM_PATH') else ctypes.util.find_library('hsa-runtime64'))\g" $BASE/hsa.py
+  python3 -c "import tinygrad.runtime.autogen.hsa"
+}
+
+generate_io_uring() {
+  clang2py -k cdefstum \
+    /usr/include/liburing.h \
+    /usr/include/linux/io_uring.h \
+    -o $BASE/io_uring.py
+
+  sed -r '/^#define __NR_io_uring/ s/^#define __(NR_io_uring[^ ]+) (.*)$/\1 = \2/; t; d' /usr/include/asm-generic/unistd.h >> $BASE/io_uring.py # io_uring syscalls numbers
+  fixup $BASE/io_uring.py
+}
+
+generate_ib() {
+  clang2py -k cdefstum \
+    /usr/include/infiniband/verbs.h \
+    /usr/include/infiniband/verbs_api.h \
+    /usr/include/infiniband/ib_user_ioctl_verbs.h \
+    /usr/include/rdma/ib_user_verbs.h \
+    -o $BASE/ib.py
+
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" "$BASE/ib.py"
+  sed -i "s\FIXME_STUB\libibverbs\g" "$BASE/ib.py"
+  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(ctypes.util.find_library('ibverbs'), use_errno=True)\g" "$BASE/ib.py"
+
+  fixup $BASE/ib.py
+}
+
+generate_libc() {
+  clang2py -k cdefstum \
+    $(dpkg -L libc6-dev | grep sys/mman.h) \
+    $(dpkg -L libc6-dev | grep sys/syscall.h) \
+    /usr/include/string.h \
+    /usr/include/elf.h \
+    /usr/include/unistd.h \
+    /usr/include/asm-generic/mman-common.h \
+    -o $BASE/libc.py
+
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libc.py
+  sed -i "s\FIXME_STUB\libc\g" $BASE/libc.py
+  sed -i "s\FunctionFactoryStub()\None if (libc_path := ctypes.util.find_library('c')) is None else ctypes.CDLL(libc_path, use_errno=True)\g" $BASE/libc.py
+
+  fixup $BASE/libc.py
+}
+
+generate_llvm() {
+  INC="$(llvm-config-14 --includedir)"
+  clang2py -k cdefstum \
+    $(find "$INC/llvm-c/" -type f -name '*.h' | sort) \
+    "$INC/llvm/Config/Targets.def" \
+    "$INC/llvm/Config/AsmPrinters.def" \
+    "$INC/llvm/Config/AsmParsers.def" \
+    "$INC/llvm/Config/Disassemblers.def" \
+    --clang-args="$(llvm-config-14 --cflags)" \
+    -o "$BASE/llvm.py"
+
+  sed -i "s\import ctypes\import ctypes, tinygrad.runtime.support.llvm as llvm_support\g" "$BASE/llvm.py"
+  sed -i "s\FIXME_STUB\llvm\g" "$BASE/llvm.py"
+  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(llvm_support.LLVM_PATH)\g" "$BASE/llvm.py"
+
+  fixup "$BASE/llvm.py"
+}
+
+generate_kgsl() {
+  clang2py extra/qcom_gpu_driver/msm_kgsl.h -o $BASE/kgsl.py -k cdefstum
+  fixup $BASE/kgsl.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/kgsl.py
+  sed -nE 's/#define ([A-Za-z0-9_]+)_SHIFT\s*[^\S\r\n]*[0-9]*$/def \1(val): return (val << \1_SHIFT) \& \1_MASK/p' extra/qcom_gpu_driver/msm_kgsl.h >> $BASE/kgsl.py
+  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\__fd.ioctl((__idir<<30)\g" $BASE/kgsl.py
+  python3 -c "import tinygrad.runtime.autogen.kgsl"
+}
+
+generate_adreno() {
+  clang2py extra/qcom_gpu_driver/a6xx.xml.h -o $BASE/adreno.py -k cestum
+  sed -nE 's/#define ([A-Za-z0-9_]+)__SHIFT\s*[^\S\r\n]*[0-9]*$/def \1(val): return (val << \1__SHIFT) \& \1__MASK/p' extra/qcom_gpu_driver/a6xx.xml.h >> $BASE/adreno.py
+  fixup $BASE/adreno.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/adreno.py
+  python3 -c "import tinygrad.runtime.autogen.adreno"
+}
+
+generate_qcom() {
+  clang2py -k cdefstum \
+    extra/dsp/include/ion.h \
+    extra/dsp/include/msm_ion.h \
+    extra/dsp/include/adsprpc_shared.h \
+    extra/dsp/include/remote_default.h \
+    extra/dsp/include/apps_std.h \
+    -o $BASE/qcom_dsp.py
+
+  fixup $BASE/qcom_dsp.py
+  python3 -c "import tinygrad.runtime.autogen.qcom_dsp"
+}
+
+generate_pci() {
+  clang2py -k cdefstum \
+    /usr/include/linux/pci_regs.h \
+    -o $BASE/pci.py
+  fixup $BASE/pci.py
+}
+
+generate_vfio() {
+  clang2py -k cdefstum \
+    /usr/include/linux/vfio.h \
+    -o $BASE/vfio.py
+  fixup $BASE/vfio.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/vfio.py
+  sed -i "s\import fcntl, functools\import functools" $BASE/vfio.py
+  sed -i "s\import ctypes,os\a from tinygrad.runtime.support import FileIOInterface\g" $BASE/vfio.py
+  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\return __fd.ioctl((__idir<<30)\g" $BASE/vfio.py
+}
+
+generate_am() {
+  AMKERN_COMMIT_HASH=ceb12c04e2b5b53ec0779362831f5ee40c4921e4
+  AMKERN_SRC=/tmp/ROCK-Kernel-Driver-$AMKERN_COMMIT_HASH
+  if [ ! -d "$AMKERN_SRC" ]; then
+    git clone https://github.com/ROCm/ROCK-Kernel-Driver $AMKERN_SRC --depth 1
+  fi
+  AMKERN_AMD=$AMKERN_SRC/drivers/gpu/drm/amd/
+  AMKERN_INC=$AMKERN_AMD/include/
+
+  clang2py -k cdefstum \
+    extra/amdpci/headers/v11_structs.h \
+    extra/amdpci/headers/v12_structs.h \
+    extra/amdpci/headers/amdgpu_vm.h \
+    extra/amdpci/headers/discovery.h \
+    extra/amdpci/headers/amdgpu_ucode.h \
+    extra/amdpci/headers/psp_gfx_if.h \
+    extra/amdpci/headers/amdgpu_psp.h \
+    extra/amdpci/headers/amdgpu_irq.h \
+    extra/amdpci/headers/amdgpu_doorbell.h \
+    $AMKERN_INC/soc15_ih_clientid.h \
+    --clang-args="-include stdint.h" \
+    -o $BASE/am/am.py
+  fixup $BASE/am/am.py
+  sed -i "s\(int64_t)\ \g" $BASE/am/am.py
+  sed -i "s\AMDGPU_PTE_MTYPE_VG10(2)\AMDGPU_PTE_MTYPE_VG10(0, 2)\g" $BASE/am/am.py # incorrect parsing (TODO: remove when clang2py is gone).
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
+    $AMKERN_AMD/amdgpu/soc15d.h \
+    -o $BASE/am/pm4_soc15.py
+  fixup $BASE/am/pm4_soc15.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
+    $AMKERN_AMD/amdgpu/nvd.h \
+    -o $BASE/am/pm4_nv.py
+  fixup $BASE/am/pm4_nv.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/vega10_sdma_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_4_0_0.py
+  fixup $BASE/am/sdma_4_0_0.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/navi10_sdma_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_5_0_0.py
+  fixup $BASE/am/sdma_5_0_0.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/sdma_v6_0_0_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_6_0_0.py
+  fixup $BASE/am/sdma_6_0_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v13_0_0_ppsmc.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu13_driver_if_v13_0_0.h \
+    extra/amdpci/headers/amdgpu_smu.h \
+    -o $BASE/am/smu_v13_0_0.py
+  fixup $BASE/am/smu_v13_0_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_0_pmfw.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_2_ppsmc.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu14_driver_if_v14_0.h \
+    extra/amdpci/headers/amdgpu_smu.h \
+    --clang-args="-include stdint.h" \
+    -o $BASE/am/smu_v14_0_2.py
+  fixup $BASE/am/smu_v14_0_2.py
+}
+
+generate_sqtt() {
+  clang2py -k cdefstum \
+    extra/sqtt/sqtt.h \
+    -o $BASE/sqtt.py
+  fixup $BASE/sqtt.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
+  python3 -c "import tinygrad.runtime.autogen.sqtt"
+
+  ROCPROF_COMMIT_HASH=dd0485100971522cc4cd8ae136bdda431061a04d
+  ROCPROF_SRC=/tmp/rocprof-trace-decoder-$ROCPROF_COMMIT_HASH
+  if [ ! -d "$ROCPROF_SRC" ]; then
+    git clone https://github.com/ROCm/rocprof-trace-decoder $ROCPROF_SRC
+    pushd .
+    cd $ROCPROF_SRC
+    git reset --hard $ROCPROF_COMMIT_HASH
+    popd
+  fi
+
+  clang2py -k cdefstum \
+    $ROCPROF_SRC/include/rocprof_trace_decoder.h \
+    $ROCPROF_SRC/include/trace_decoder_instrument.h \
+    $ROCPROF_SRC/include/trace_decoder_types.h \
+    -o extra/sqtt/rocprof/rocprof.py
+  fixup extra/sqtt/rocprof/rocprof.py
+  sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py
+  sed -i "s/import ctypes/import ctypes, ctypes.util/g" extra/sqtt/rocprof/rocprof.py
+  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder'))|g" extra/sqtt/rocprof/rocprof.py
+}
+
+generate_webgpu() {
+  clang2py extra/webgpu/webgpu.h -o $BASE/webgpu.py
+  fixup $BASE/webgpu.py
+  sed -i "s/FIXME_STUB/webgpu/g" "$BASE/webgpu.py"
+  sed -i "s/FunctionFactoryStub()/ctypes.CDLL(webgpu_support.WEBGPU_PATH)/g" "$BASE/webgpu.py"
+  sed -i "s/import ctypes/import ctypes, tinygrad.runtime.support.webgpu as webgpu_support/g" "$BASE/webgpu.py"
+  python3 -c "import tinygrad.runtime.autogen.webgpu"
+}
+
+generate_libusb() {
+  clang2py -k cdefstum \
+    /usr/include/libusb-1.0/libusb.h \
+    -o $BASE/libusb.py
+
+  fixup $BASE/libusb.py
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libusb.py
+  sed -i "s/FIXME_STUB/libusb/g" "$BASE/libusb.py"
+  sed -i "s/libusb_le16_to_cpu = libusb_cpu_to_le16//g" "$BASE/libusb.py"
+  sed -i "s/FunctionFactoryStub()/None if (lib_path:=os.getenv('LIBUSB_PATH', ctypes.util.find_library('usb-1.0'))) is None else ctypes.CDLL(lib_path)/g" "$BASE/libusb.py"
+  python3 -c "import tinygrad.runtime.autogen.libusb"
+}
+
+if [ "$1" == "opencl" ]; then generate_opencl
+elif [ "$1" == "hip" ]; then generate_hip
+elif [ "$1" == "comgr" ]; then generate_comgr
+elif [ "$1" == "cuda" ]; then generate_cuda
+elif [ "$1" == "nvrtc" ]; then generate_nvrtc
+elif [ "$1" == "hsa" ]; then generate_hsa
+elif [ "$1" == "kfd" ]; then generate_kfd
+elif [ "$1" == "nv" ]; then generate_nv
+elif [ "$1" == "amd" ]; then generate_amd
+elif [ "$1" == "am" ]; then generate_am
+elif [ "$1" == "nvdrv" ]; then generate_nvdrv
+elif [ "$1" == "sqtt" ]; then generate_sqtt
+elif [ "$1" == "qcom" ]; then generate_qcom
+elif [ "$1" == "io_uring" ]; then generate_io_uring
+elif [ "$1" == "ib" ]; then generate_ib
+elif [ "$1" == "libc" ]; then generate_libc
+elif [ "$1" == "llvm" ]; then generate_llvm
+elif [ "$1" == "kgsl" ]; then generate_kgsl
+elif [ "$1" == "adreno" ]; then generate_adreno
+elif [ "$1" == "pci" ]; then generate_pci
+elif [ "$1" == "vfio" ]; then generate_vfio
+elif [ "$1" == "webgpu" ]; then generate_webgpu
+elif [ "$1" == "libusb" ]; then generate_libusb
+elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu
+else echo "usage: $0 <type>"
+fi
@@ -0,0 +1,137 @@
+# tinygrad is a tensor library, and as a tensor library it has multiple parts
+# 1. a "runtime". this allows buffer management, compilation, and running programs
+# 2. a "Device" that uses the runtime but specifies compute in an abstract way for all
+# 3. a "UOp" that fuses the compute into kernels, using memory only when needed
+# 4. a "Tensor" that provides an easy to use frontend with autograd ".backward()"
+
+
+print("******** first, the runtime ***********")
+
+from tinygrad.runtime.ops_cpu import ClangJITCompiler, CPUDevice, CPUProgram
+
+cpu = CPUDevice()
+
+# allocate some buffers
+out = cpu.allocator.alloc(4)
+a = cpu.allocator.alloc(4)
+b = cpu.allocator.alloc(4)
+
+# load in some values (little endian)
+cpu.allocator._copyin(a, memoryview(bytearray([2,0,0,0])))
+cpu.allocator._copyin(b, memoryview(bytearray([3,0,0,0])))
+
+# compile a program to a binary
+lib = ClangJITCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")
+
+# create a runtime for the program
+fxn = cpu.runtime("add", lib)
+
+# run the program
+fxn(out, a, b)
+
+# check the data out
+print(val := cpu.allocator._as_buffer(out).cast("I").tolist()[0])
+assert val == 5
+
+
+print("******** second, the Device ***********")
+
+DEVICE = "CPU"   # NOTE: you can change this!
+
+import struct
+from tinygrad.dtype import dtypes
+from tinygrad.device import Buffer, Device
+from tinygrad.uop.ops import UOp, Ops
+
+# allocate some buffers + load in values
+out = Buffer(DEVICE, 1, dtypes.int32).allocate()
+a = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
+b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
+# NOTE: a._buf is the same as the return from cpu.allocator.alloc
+
+# describe the computation
+idx = UOp.const(dtypes.index, 0)
+buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
+buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
+ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.index(idx),))
+ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.index(idx),))
+alu = ld_1 + ld_2
+output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
+st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.index(idx), alu))
+s = UOp(Ops.SINK, dtypes.void, (st_0,))
+
+# convert the computation to a "linearized" format (print the format)
+from tinygrad.engine.realize import get_program, CompiledRunner
+program = get_program(s, Device[DEVICE].renderer)
+
+# compile a program (and print the source)
+fxn = CompiledRunner(program)
+print(fxn.p.src)
+# NOTE: fxn.clprg is the CPUProgram
+
+# run the program
+fxn.exec([out, a, b])
+
+# check the data out
+assert out.as_buffer().cast('I')[0] == 5
+
+
+print("******** third, the UOp ***********")
+
+from tinygrad.engine.realize import run_schedule
+from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.schedule.rangeify import get_rangeify_map
+
+# allocate some values + load in values
+a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
+b = UOp.new_buffer(DEVICE, 1, dtypes.int32)
+a.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
+b.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
+
+# describe the computation
+out = a + b
+s = UOp(Ops.SINK, dtypes.void, (out,))
+
+# group the computation into kernels
+becomes_map = get_rangeify_map(s)
+
+# the compute maps to an assign
+assign = becomes_map[a+b].base
+
+# the first source is the output buffer (data)
+assert assign.src[0].op is Ops.BUFFER
+# the second source is the kernel (compute)
+assert assign.src[1].op is Ops.KERNEL
+
+# schedule the kernel graph in a linear list
+s = UOp(Ops.SINK, dtypes.void, (assign,))
+sched, _ = create_schedule_with_vars(s)
+assert len(sched) == 1
+
+# DEBUGGING: print the compute ast
+print(sched[-1].ast)
+# NOTE: sched[-1].ast is the same as st_0 above
+
+# the output will be stored in a new buffer
+out = assign.buf_uop
+assert out.op is Ops.BUFFER and not out.buffer.is_allocated()
+print(out)
+
+# run that schedule
+run_schedule(sched)
+
+# check the data out
+assert out.is_realized and out.buffer.as_buffer().cast('I')[0] == 5
+
+
+print("******** fourth, the Tensor ***********")
+
+from tinygrad import Tensor
+
+a = Tensor([2], dtype=dtypes.int32, device=DEVICE)
+b = Tensor([3], dtype=dtypes.int32, device=DEVICE)
+out = a + b
+
+# check the data out
+print(val:=out.item())
+assert val == 5
@@ -38,19 +38,25 @@ optim.schedule_step()   # this will step the optimizer without running realize
 # The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
 # l1.uop and l2.uop define a computation graph

-from tinygrad.engine.schedule import ExecItem
-schedule: List[ExecItem] = Tensor.schedule(l1, l2)
+from tinygrad.engine.schedule import ScheduleItem
+schedule: List[ScheduleItem] = Tensor.schedule(l1, l2)

 print(f"The schedule contains {len(schedule)} items.")
 for si in schedule: print(str(si)[:80])

 # *****
-# 4. Lower and run the schedule.
+# 4. Lower a schedule.

-for si in tqdm(schedule): si.run()
+from tinygrad.engine.realize import lower_schedule_item, ExecItem
+lowered: List[ExecItem] = [lower_schedule_item(si) for si in tqdm(schedule)]

 # *****
-# 5. Print the weight change
+# 5. Run the schedule
+
+for ei in tqdm(lowered): ei.run()
+
+# *****
+# 6. Print the weight change

 print("first weight change\n", l1.numpy()-l1n)
 print("second weight change\n", l2.numpy()-l2n)
@@ -13,19 +13,19 @@ There's also a [doc describing speed](../developer/speed.md)

 Everything in [Tensor](../tensor/index.md) is syntactic sugar around constructing a graph of [UOps](../developer/uop.md).

-The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view. Inputs to a base can be either base or view, inputs to a view can only be a single base.
+The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view (specified by a ShapeTracker). Inputs to a base can be either base or view, inputs to a view can only be a single base.

 ## Scheduling

-The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of UOps into a list of `ExecItem`. One `ExecItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
+The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of UOps into a list of `ScheduleItem`. One `ScheduleItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.

-::: tinygrad.engine.schedule.ExecItem
+::: tinygrad.engine.schedule.ScheduleItem

 ## Lowering

-The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers `ExecItem` by populating its `prg` field with
+The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers `ScheduleItem` to `ExecItem` with

-::: tinygrad.engine.realize.run_schedule
+::: tinygrad.engine.realize.lower_schedule

 There's a ton of complexity hidden behind this, see the `codegen/` directory.

@@ -0,0 +1,109 @@
+# Kernel Creation
+
+Tinygrad lazily builds up a graph of Tensor operations. The Tensor graph includes a mix of:
+
+- Buffer and Assignment Ops: `BUFFER`, `BUFFER_VIEW`, `COPY`, `ASSIGN`
+- Movement Ops: `RESHAPE`, `EXPAND`, `PERMUTE`, `PAD`, `SHRINK`, `FLIP`
+- Compute Ops: `ADD`, `MUL`, `REDUCE_AXIS`, ...
+
+`Tensor.kernelize` creates the kernels and buffers needed to realize the output Tensor(s).
+
+## Kernelize flow
+
+Let's see how a multiply add Tensor graph becomes a fused elementwise kernel.
+
+```py
+# initialize 3 input buffers on the device
+a = Tensor([1]).realize()
+b = Tensor([2]).realize()
+c = Tensor([3]).realize()
+
+# create the Tensor graph
+mul = a*b
+out = mul+c
+
+print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
+print(out) # <Tensor <UOp METAL (1,) int (<Ops.ADD: 52>, None)> on METAL with grad None>
+
+out.kernelize()
+
+print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
+print(out) # <Tensor <UOp METAL (1,) int (<Ops.ASSIGN: 66>, None)> on METAL with grad None>
+```
+
+The multiply Tensor stays the same because it is fused. The output Tensor's UOp becomes a new ASSIGN UOp:
+
+```py
+print(out.uop)
+```
+
+The first source is the output BUFFER:
+
+```
+UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
+  UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),))
+```
+
+And the second source is the KERNEL and its 4 buffer edges (output_buffer, a, b, c):
+
+```
+UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 45>,) (__add__, __mul__)>, src=(
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+    x1:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
+    UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=3, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=5, src=()),)),))
+```
+
+KERNEL describes the compute AST, metadata and memory dependencies.
+
+BUFFER holds a reference to the device memory where the output will be stored.
+
+Once a Tensor is kernelized, all children will LOAD its BUFFER, instead of fusing it:
+
+```py
+child = out+2
+child.kernelize()
+print(child.uop.src[1].arg.ast)
+```
+
+```
+UOp(Ops.SINK, dtypes.void, arg=None, src=(
+  UOp(Ops.STORE, dtypes.void, arg=None, src=(
+    UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=0, src=()),
+    x2:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),
+    UOp(Ops.ADD, dtypes.int, arg=None, src=(
+      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
+        UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=()),
+         x2,)),
+      UOp(Ops.CONST, dtypes.int, arg=2, src=(
+         x2,)),)),)),))
+```
+
+`Tensor.realize` will execute the kernels and write outputs to memory:
+
+```py
+Tensor.realize(out)
+print(out)        # <Tensor <UOp METAL (1,) int (<Ops.BUFFER: 23>, <buf real:True device:METAL size:1 dtype:dtypes.int offset:0>)> on METAL with grad None>
+print(out.item()) # 5
+```
+
+<hr />
+
+**Summary**
+
+- The large Tensor graph is built from a mix of data, compute and movement Ops.
+
+- `Tensor.kernelize` splits the Tensor graph into data (BUFFER), compute (KERNEL) and links dependencies with ASSIGN.
+
+- `Tensor.realize` executes KERNELs on device and replaces the Tensor graph with just a BUFFER.
+
+- Kernelize can be called multiple times on a Tensor. This allows for incrementally building the kernel fusion layout of a large Tensor graph, without having to call `realize` or `schedule`.
@@ -26,9 +26,9 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi

 ## tinygrad/codegen

-Transform the optimized ast into a linearized and rendered program.
+Transform the optimized ast into a linearized list of UOps.

-::: tinygrad.codegen.get_program
+::: tinygrad.codegen.full_rewrite
    options:
        members: false
        show_labels: false
@@ -41,7 +41,7 @@ BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
 IMAGE               | [1-2]      | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
-HCQ_VISIBLE_DEVICES | [list[int]]| restricts the HCQ devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
+VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
 VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
 ALLOW_TF32          | [1]        | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
@@ -131,7 +131,7 @@ timeit.repeat(jit_step, repeat=5, number=1)

 1.0 ms is 75x faster! Note that we aren't syncing the GPU, so GPU time may be slower.

-The first two runs of the function execute normally, with the JIT capturing the kernels. Starting from the third run, only the tinygrad operations are replayed, removing the overhead by skipping Python code execution. So be aware that any non-tinygrad Python values affecting the kernels will be "frozen" from the second run. Note that `Tensor` randomness functions work as expected.
+The slowness the first two times is the JIT capturing the kernels. And this JIT will not run any Python in the function, it will just replay the tinygrad kernels that were run, so be aware that non tinygrad Python operations won't work. Randomness functions work as expected.

 Unlike other JITs, we JIT everything, including the optimizer. Think of it as a dumb replay on different data.

@@ -0,0 +1,293 @@
+#!/usr/bin/env python3
+
+# this file is a "ramp" for people new to tinygrad to think about how to approach it
+# it is runnable and editable.
+# whenever you see stuff like DEBUG=2 or CPU=1 discussed, these are environment variables
+# in a unix shell like bash `DEBUG=2 CPU=1 python docs/ramp.py`
+
+# this pip installs tinygrad master for the system
+# the -e allows you to edit the tinygrad folder and update system tinygrad
+# tinygrad is pure Python, so you are encouraged to do this
+# git pull in the tinygrad directory will also get you the latest
+"""
+git clone https://github.com/tinygrad/tinygrad.git
+cd tinygrad
+python3 -m pip install -e .
+"""
+
+# %% ********
+print("******* PART 1 *******")
+
+# we start with a Device.
+# a Device is where Tensors are stored and compute is run
+# tinygrad autodetects the best device on your system and makes it the DEFAULT
+from tinygrad import Device
+print(Device.DEFAULT)  # on Mac, you can see this prints METAL
+
+# now, lets create a Tensor
+from tinygrad import Tensor, dtypes
+t = Tensor([1,2,3,4])
+
+# you can see this Tensor is on the DEFAULT device with int dtype and shape (4,)
+assert t.device == Device.DEFAULT
+assert t.dtype == dtypes.int
+assert t.shape == (4,)
+
+# unlike in torch, if we print it, it doesn't print the contents
+# this is because tinygrad is lazy
+# this Tensor has not been computed yet
+print(t)
+# <Tensor <UOp METAL (4,) int (<Ops.COPY: 7>, None)> on METAL with grad None>
+
+# the ".uop" property on Tensor contains the specification of how to compute it
+print(t.uop)
+"""
+UOp(Ops.COPY, dtypes.int, arg=None, src=(
+  UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+    UOp(Ops.UNIQUE, dtypes.void, arg=0, src=()),
+    UOp(Ops.DEVICE, dtypes.void, arg='PYTHON', src=()),)),
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
+"""
+# as you can see, it's specifying a copy from PYTHON device
+# which is where the [1,2,3,4] array lives
+
+# UOps are the specification language in tinygrad
+# they are immutable and form a DAG
+# they have a "Ops", a "dtype", a tuple of srcs (parents), and an arg
+
+t.realize()
+# if we want to "realize" a tensor, we can with the "realize" method
+# now when we look at the uop, it's changed
+print(t.uop)
+"""
+UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+  UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
+"""
+# the copy was actually run, and now the "uop" of the Tensor is just a BUFFER
+# if you run this script with DEBUG=2 in the environment, you can see the copy happen
+# *** METAL      1 copy       16,   METAL <- PYTHON ...
+
+# now let's do some compute
+# we look at the uop to see the specification of the compute
+t_times_2 = t * 2
+print(t_times_2.uop)
+"""
+UOp(Ops.MUL, dtypes.int, arg=None, src=(
+  UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+    x2:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),)),
+  UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
+    UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
+      UOp(Ops.CONST, dtypes.int, arg=2, src=(
+        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)), src=(
+           x2,)),)),)),)),))
+"""
+# the BUFFER from above is being multiplied by a CONST 2
+# it's RESHAPEd and EXPANDed to broadcast the CONST to the BUFFER
+
+# we can check the result with
+assert t_times_2.tolist() == [2, 4, 6, 8]
+
+# UOps are both immutable and globally unique
+# if i multiply the Tensor by 4 twice, these result Tensors will have the same uop specification
+t_times_4_try_1 = t * 4
+t_times_4_try_2 = t * 4
+assert t_times_4_try_1.uop is t_times_4_try_2.uop
+# the specification isn't just the same, it's the exact same Python object
+assert t_times_4_try_1 is not t_times_4_try_2
+# the Tensor is a different Python object
+
+# if we realize `t_times_4_try_1` ...
+t_times_4_try_1.realize()
+print(t_times_4_try_2.uop)
+"""
+UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+  UOp(Ops.UNIQUE, dtypes.void, arg=4, src=()),
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
+"""
+# ... `t_times_4_try_2` also becomes the same BUFFER
+assert t_times_4_try_1.uop is t_times_4_try_2.uop
+# so this print doesn't require any computation, just a copy back to the CPU so we can print it
+print("** only the copy start")
+print(t_times_4_try_2.tolist())  # [4, 8, 12, 16]
+print("** only the copy end")
+# you can confirm this with DEBUG=2, seeing what's printed in between the "**" prints
+
+# tinygrad has an auto differentiation engine that operates according to these same principles
+# the derivative of "log(x)" is "1/x", and you can see this on line 20 of gradient.py
+t_float = Tensor([3.0])
+t_log = t_float.log()
+t_log_grad, = t_log.sum().gradient(t_float)
+# due to how log is implemented, this gradient contains a lot of UOps
+print(t_log_grad.uop)
+# ...not shown here...
+# but if you run with DEBUG=4 (CPU=1 used here for simpler code), you can see the generated code
+"""
+void E_(float* restrict data0, float* restrict data1) {
+  float val0 = *(data1+0);
+  *(data0+0) = (1/val0);
+}
+"""
+# the derivative is close to 1/3
+assert (t_log_grad.item() - 1/3) < 1e-6
+
+# %% ********
+print("******* PART 2 *******")
+
+# we redefine the same t here so this cell can run on it's own
+from tinygrad import Tensor
+t = Tensor([1,2,3,4])
+
+# what's above gives you enough of an understanding to go use tinygrad as a library
+# however, a lot of the beauty of tinygrad is in how easy it is to interact with the internals
+# NOTE: the APIs here are subject to change
+
+t_plus_3_plus_4 = t + 3 + 4
+print(t_plus_3_plus_4.uop)
+"""
+UOp(Ops.ADD, dtypes.int, arg=None, src=(
+  UOp(Ops.ADD, dtypes.int, arg=None, src=(
+    UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+      UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+      x3:=UOp(Ops.DEVICE, dtypes.void, arg='CPU', src=()),)),
+    UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
+      UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
+        UOp(Ops.CONST, dtypes.int, arg=3, src=(
+          x7:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)), src=(
+             x3,)),)),)),)),)),
+  UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
+    UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
+      UOp(Ops.CONST, dtypes.int, arg=4, src=(
+         x7,)),)),)),))
+"""
+# you can see it's adding both 3 and 4
+
+# but by the time we are actually running the code, it's adding 7
+# `kernelize` will simplify and group the operations in the graph into kernels
+t_plus_3_plus_4.kernelize()
+print(t_plus_3_plus_4.uop)
+"""
+UOp(Ops.ASSIGN, dtypes.int, arg=None, src=(
+  x0:=UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+    UOp(Ops.UNIQUE, dtypes.void, arg=7, src=()),
+    x2:=UOp(Ops.DEVICE, dtypes.void, arg='CPU', src=()),)),
+  UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 48>,) (__add__,)>, src=(
+     x0,
+    UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+      UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+       x2,)),)),))
+"""
+# ASSIGN has two srcs, src[0] is the BUFFER that's assigned to, and src[1] is the thing to assign
+# src[1] is the GPU Kernel that's going to be run
+# we can get the ast of the Kernel as follows
+kernel_ast = t_plus_3_plus_4.uop.src[1].arg.ast
+
+# almost everything in tinygrad functions as a rewrite of the UOps
+# the codegen rewrites the ast to a simplified form ready for "rendering"
+from tinygrad.codegen import full_rewrite_to_sink
+rewritten_ast = full_rewrite_to_sink(kernel_ast)
+print(rewritten_ast)
+"""
+UOp(Ops.SINK, dtypes.void, arg=None, src=(
+  UOp(Ops.STORE, dtypes.void, arg=None, src=(
+    UOp(Ops.INDEX, dtypes.int.ptr(4), arg=None, src=(
+      UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(4), arg=0, src=()),
+      x3:=UOp(Ops.SPECIAL, dtypes.int, arg=('gidx0', 4), src=()),)),
+    UOp(Ops.ADD, dtypes.int, arg=None, src=(
+      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
+        UOp(Ops.INDEX, dtypes.int.ptr(4), arg=None, src=(
+          UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(4), arg=1, src=()),
+           x3,)),)),
+      UOp(Ops.CONST, dtypes.int, arg=7, src=()),)),)),))
+"""
+# you can see at this point we are adding 7, not 3 and 4
+
+# with DEBUG=4, we can see the code.
+# since optimizations are on, it UPCASTed the operation, explicitly writing out all 4 +7s
+t_plus_3_plus_4.realize()
+"""
+void E_4n2(int* restrict data0, int* restrict data1) {
+  int val0 = *(data1+0);
+  int val1 = *(data1+1);
+  int val2 = *(data1+2);
+  int val3 = *(data1+3);
+  *(data0+0) = (val0+7);
+  *(data0+1) = (val1+7);
+  *(data0+2) = (val2+7);
+  *(data0+3) = (val3+7);
+}
+"""
+# the function name E_4n2 is "E" for elementwise op (as opposed to "r" for reduce op)
+# "4" for the size, and "n2" for name deduping (it's the 3rd function with the same E and 4 in this session)
+# when you print the name with DEBUG=2, you'll see the 4 is yellow, meaning that it's upcasted
+# if you run with NOOPT=1 ...
+"""
+void E_4n2(int* restrict data0, int* restrict data1) {
+  for (int ridx0 = 0; ridx0 < 4; ridx0++) {
+    int val0 = *(data1+ridx0);
+    *(data0+ridx0) = (val0+7);
+  }
+}
+"""
+# ... you get this unoptimized code with a loop and the 4 is blue (for global). the color code is in kernel.py
+
+# %% ********
+print("******* PART 3 *******")
+
+# now, we go even lower and understand UOps better and how the graph rewrite engine works.
+# it's much simpler than what's in LLVM or MLIR
+
+from tinygrad import dtypes
+from tinygrad.uop.ops import UOp, Ops
+
+# first, we'll construct some const UOps
+a = UOp(Ops.CONST, dtypes.int, arg=2)
+b = UOp(Ops.CONST, dtypes.int, arg=2)
+
+# if you have been paying attention, you should know these are the same Python object
+assert a is b
+
+# UOps support normal Python math operations, so a_plus_b expresses the spec for 2 + 2
+a_plus_b = a + b
+print(a_plus_b)
+"""
+UOp(Ops.ADD, dtypes.int, arg=None, src=(
+  x0:=UOp(Ops.CONST, dtypes.int, arg=2, src=()),
+   x0,))
+"""
+
+# we could actually render this 2+2 into a language like c and run it
+# or, we can use tinygrad's graph rewrite engine to "constant fold"
+
+from tinygrad.uop.ops import graph_rewrite, UPat, PatternMatcher
+
+# a `PatternMatcher` is a list of tuples. for each element in the list:
+# [0] is the pattern to match, and [1] is the function to run.
+# this function can return either a UOp to replace the pattern with, or None to not replace
+simple_pm = PatternMatcher([
+  (UPat(Ops.ADD, src=(UPat(Ops.CONST, name="c1"), UPat(Ops.CONST, name="c2"))),
+   lambda c1,c2: UOp(Ops.CONST, dtype=c1.dtype, arg=c1.arg+c2.arg)),
+])
+# this pattern matches the addition of two CONST and rewrites it into a single CONST UOp
+
+# to actually apply the pattern to a_plus_b, we use graph_rewrite
+a_plus_b_simplified = graph_rewrite(a_plus_b, simple_pm)
+print(a_plus_b_simplified)
+"""
+UOp(Ops.CONST, dtypes.int, arg=4, src=())
+"""
+# 2+2 is in fact, 4
+
+# we can also use syntactic sugar to write the pattern nicer
+simpler_pm = PatternMatcher([
+  (UPat.cvar("c1")+UPat.cvar("c2"), lambda c1,c2: c1.const_like(c1.arg+c2.arg))
+])
+assert graph_rewrite(a_plus_b, simple_pm) is graph_rewrite(a_plus_b, simpler_pm)
+# note again the use of is, UOps are immutable and globally unique
+
+# %% ********
+
+# that brings you to an understanding of the most core concepts in tinygrad
+# you can run this with VIZ=1 to use the web based graph rewrite explorer
+# hopefully now you understand it. the nodes in the graph are just UOps
@@ -41,7 +41,7 @@ The BMC also has a web interface you can use if you find that easier.
 It is recommended that you change the BMC password after setting up the box, as the password on the screen is only the initial password.

 If you do decide to change the BMC password and no longer want the initial password to be displayed, remove the `/root/.bmc_password` file.
-Reboot after making these changes or restart the `tinybox-display.service` service.
+Reboot after making these changes or restart the `displayservice.service` service.

 ## What do I use it for?

@@ -0,0 +1,9 @@
+import globals from "globals";
+import pluginJs from "@eslint/js";
+import pluginHtml from "eslint-plugin-html";
+
+export default [
+  {files: ["**/*.html"], plugins: {html: pluginHtml}, rules:{"max-len": ["error", {"code": 150}]}},
+  {languageOptions: {globals: globals.browser}},
+  pluginJs.configs.recommended,
+];
@@ -21,7 +21,7 @@ if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))

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

  @TinyJit
  @Tensor.train()
@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+import os, sys, traceback
+sys.path.append(os.getcwd())
+
+from io import StringIO
+from contextlib import redirect_stdout
+from tinygrad import Tensor, nn
+from tinygrad.helpers import Timing, colored, getenv, fetch
+from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
+from sentencepiece import SentencePieceProcessor
+
+def create_fixed_tokenizer(output_file):
+  print("creating fixed tokenizer")
+  import extra.junk.sentencepiece_model_pb2 as spb2
+  mp = spb2.ModelProto()
+  mp.ParseFromString(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/tokenizer.model?download=true").read_bytes())
+  mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_end|>", score=0))
+  mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_start|>", score=0))
+  with open(output_file, "wb") as f:
+    f.write(mp.SerializeToString())
+
+# example:
+# echo -en "write 2+2\nwrite hello world\ny\n" | TEMP=0 python3 examples/coder.py
+
+if __name__ == "__main__":
+  # https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/config.json
+  with Timing("create model: "):
+    model = Transformer(4096, 14336, n_heads=32, n_layers=32, norm_eps=1e-5, vocab_size=32002, n_kv_heads=8, max_context=4096, jit=getenv("JIT", 1))
+
+  with Timing("download weights: "):
+    part1 = nn.state.torch_load(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/pytorch_model-00001-of-00002.bin?download=true"))
+    part2 = nn.state.torch_load(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/pytorch_model-00002-of-00002.bin?download=true"))
+
+  with Timing("weights -> model: "):
+    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part1, 32, 32, 8)), strict=False)
+    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part2, 32, 32, 8)), strict=False)
+
+  if not os.path.isfile("/tmp/tokenizer.model"): create_fixed_tokenizer("/tmp/tokenizer.model")
+  spp = SentencePieceProcessor(model_file="/tmp/tokenizer.model")
+
+  # https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/tokenizer_config.json
+  #   "chat_template": "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
+  IM_END = 32000
+  IM_START = 32001
+  def encode_prompt(k, v): return [IM_START]+spp.encode(f"{k}\n{v}")+[IM_END]+spp.encode("\n")
+  def start_prompt(k): return [IM_START]+spp.encode(f"{k}\n")
+  def output(outputted, toks, color):
+    cur = spp.decode(toks)[len(outputted):]
+    sys.stdout.write(colored(cur, color))
+    sys.stdout.flush()
+    outputted += cur
+    return outputted
+
+  # *** app below this line ***
+
+  toks = [spp.bos_id()] + encode_prompt("system", "You are Quentin. Quentin is a useful assistant who writes Python code to answer questions. He keeps the code as short as possible and doesn't read from user input")
+
+  PROMPT = getenv("PROMPT", 1)
+  temperature = getenv("TEMP", 0.7)
+
+  start_pos = 0
+  outputted = output("", toks, "green")
+  turn = True
+  while 1:
+    if PROMPT:
+      toks += encode_prompt("user", input("Q: ")) + start_prompt("assistant")
+    else:
+      toks += start_prompt("user" if turn else "assistant")
+      turn = not turn
+    old_output_len = len(outputted)
+    while 1:
+      tok = model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
+      start_pos = len(toks)
+      toks.append(tok)
+      outputted = output(outputted, toks, "blue" if not turn else "cyan")
+      if tok == IM_END: break
+      if tok == spp.eos_id(): break
+      new_output = outputted[old_output_len:]
+
+      if new_output.endswith("```") and '```python\n' in new_output:
+        python_code = new_output.split('```python\n')[1].split("```")[0]
+        # AI safety. Warning to user. Do not press y if the AI is trying to do unsafe things.
+        if input(colored(f" <-- PYTHON DETECTED, RUN IT? ", "red")).lower() == 'y':
+          my_stdout = StringIO()
+          try:
+            with redirect_stdout(my_stdout): exec(python_code)
+            result = my_stdout.getvalue()
+          except Exception as e:
+            result = ''.join(traceback.format_exception_only(e))
+          toks += spp.encode(f"\nOutput:\n```\n{result}```")
+          outputted = output(outputted, toks, "yellow")
+          old_output_len = len(outputted)
+    print("")
@@ -0,0 +1,341 @@
+import argparse
+import multiprocessing as mp
+import os
+import re
+import sys
+import time
+from contextlib import contextmanager
+from pathlib import Path
+
+import numpy as np
+import pyaudio
+import yaml
+from llama import LLaMa
+from vits import MODELS as VITS_MODELS
+from vits import Y_LENGTH_ESTIMATE_SCALARS, HParams, Synthesizer, TextMapper, get_hparams_from_file, load_model
+from whisper import init_whisper, transcribe_waveform
+from sentencepiece import SentencePieceProcessor
+
+from tinygrad.helpers import Timing, fetch
+from tinygrad import Tensor, dtypes
+
+# Whisper constants
+RATE = 16000
+CHUNK = 1600
+
+# LLaMa constants
+IM_START = 32001
+IM_END = 32002
+
+
+# Functions for encoding prompts to chatml md
+def encode_prompt(spp, k, v): return [IM_START]+spp.encode(f"{k}\n{v}")+[IM_END]+spp.encode("\n")
+def start_prompt(spp, k): return [IM_START]+spp.encode(f"{k}\n")
+
+def chunks(lst, n):
+  for i in range(0, len(lst), n): yield lst[i:i + n]
+
+def create_fixed_tokenizer():
+  """Function needed for extending tokenizer with additional chat tokens"""
+  import extra.junk.sentencepiece_model_pb2 as spb2
+  tokenizer_path = fetch("https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/resolve/main/tokenizer.model")
+  if SentencePieceProcessor(model_file=str(tokenizer_path)).vocab_size() != 32003:
+    print("creating fixed tokenizer")
+    mp = spb2.ModelProto()
+    mp.ParseFromString(tokenizer_path.read_bytes())
+    # https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/blob/main/added_tokens.json
+    mp.pieces.append(spb2.ModelProto.SentencePiece(piece="[PAD]", score=0))
+    mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_start|>", score=0))
+    mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_end|>", score=0))
+    tokenizer_path.write_bytes(mp.SerializeToString())
+  return tokenizer_path
+
+def llama_prepare(llama: LLaMa, temperature: float, pre_prompt_path: Path) -> tuple[list[int], str, str, str]:
+  """Prepares a llama model from a specified pre-prompt file"""
+  with open(str(pre_prompt_path)) as f:
+    config = yaml.safe_load(f.read())
+  toks = [llama.tokenizer.bos_id()] + encode_prompt(llama.tokenizer, "system", config["pre_prompt"].replace("\n", " "))
+  for i in config["examples"]:
+    toks += encode_prompt(llama.tokenizer, config["user_delim"], i["user_prompt"])
+    toks += encode_prompt(llama.tokenizer, config["resp_delim"], i["resp_prompt"])
+  llama.model(Tensor([toks]), 0, temperature).realize()  # NOTE: outputs are not used
+  return toks, config["user_delim"], config["resp_delim"], len(toks), llama.tokenizer.decode(toks)
+
+def llama_generate(
+  llama: LLaMa,
+  toks: list[int],
+  outputted: str,
+  prompt: str,
+  start_pos: int,
+  user_delim: str,
+  resp_delim: str,
+  temperature=0.7,
+  max_tokens=1000
+):
+  """Generates an output for the specified prompt"""
+  toks += encode_prompt(llama.tokenizer, user_delim, prompt)
+  toks += start_prompt(llama.tokenizer, resp_delim)
+
+  outputted = llama.tokenizer.decode(toks)
+  init_length = len(outputted)
+  for _ in range(max_tokens):
+    token = llama.model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
+    start_pos = len(toks)
+    toks.append(token)
+
+    cur = llama.tokenizer.decode(toks)
+
+    # Print is just for debugging
+    sys.stdout.write(cur[len(outputted):])
+    sys.stdout.flush()
+    outputted = cur
+    if toks[-1] == IM_END: break
+  else:
+    toks.append(IM_END)
+  print() # because the output is flushed
+  return outputted, start_pos, outputted[init_length:].replace("<|im_end|>", "")
+
+def tts(
+  text_to_synthesize: str,
+  synth: Synthesizer,
+  hps: HParams,
+  emotion_embedding: Path,
+  speaker_id: int,
+  model_to_use: str,
+  noise_scale: float,
+  noise_scale_w: float,
+  length_scale: float,
+  estimate_max_y_length: bool,
+  text_mapper: TextMapper,
+  model_has_multiple_speakers: bool,
+  pad_length=600,
+  vits_pad_length=1000
+):
+  if model_to_use == "mmts-tts": text_to_synthesize = text_mapper.filter_oov(text_to_synthesize.lower())
+
+  # Convert the input text to a tensor.
+  stn_tst = text_mapper.get_text(text_to_synthesize, hps.data.add_blank, hps.data.text_cleaners)
+  init_shape = stn_tst.shape
+  assert init_shape[0] < pad_length, "text is too long"
+  x_tst, x_tst_lengths = stn_tst.pad(((0, pad_length - init_shape[0]),), value=1).unsqueeze(0), Tensor([init_shape[0]], dtype=dtypes.int64)
+  sid = Tensor([speaker_id], dtype=dtypes.int64) if model_has_multiple_speakers else None
+
+  # Perform inference.
+  audio_tensor = synth.infer(x_tst, x_tst_lengths, sid, noise_scale, length_scale, noise_scale_w, emotion_embedding=emotion_embedding,
+                             max_y_length_estimate_scale=Y_LENGTH_ESTIMATE_SCALARS[model_to_use] if estimate_max_y_length else None, pad_length=vits_pad_length)[0, 0]
+  # Save the audio output.
+  audio_data = (np.clip(audio_tensor.numpy(), -1.0, 1.0) * 32767).astype(np.int16)
+  return audio_data
+
+def init_vits(
+  model_to_use: str,
+  emotion_path: Path,
+  speaker_id: int,
+  seed: int,
+):
+  model_config = VITS_MODELS[model_to_use]
+
+  # Load the hyperparameters from the config file.
+  hps = get_hparams_from_file(fetch(model_config[0]))
+
+  # If model has multiple speakers, validate speaker id and retrieve name if available.
+  model_has_multiple_speakers = hps.data.n_speakers > 0
+  if model_has_multiple_speakers:
+    if speaker_id >= hps.data.n_speakers: raise ValueError(f"Speaker ID {speaker_id} is invalid for this model.")
+    if hps.__contains__("speakers"): # maps speaker ids to names
+      speakers = hps.speakers
+      if isinstance(speakers, list): speakers = {speaker: i for i, speaker in enumerate(speakers)}
+
+  # Load emotions if any. TODO: find an english model with emotions, this is untested atm.
+  emotion_embedding = None
+  if emotion_path is not None:
+    if emotion_path.endswith(".npy"): emotion_embedding = Tensor(np.load(emotion_path), dtype=dtypes.int64).unsqueeze(0)
+    else: raise ValueError("Emotion path must be a .npy file.")
+
+  # Load symbols, instantiate TextMapper and clean the text.
+  if hps.__contains__("symbols"): symbols = hps.symbols
+  elif model_to_use == "mmts-tts": symbols = [x.replace("\n", "") for x in fetch("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/vocab.txt").open(encoding="utf-8").readlines()]
+  else: symbols = ['_'] + list(';:,.!?¡¿—…"«»“” ') + list('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz') + list("ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ")
+  text_mapper = TextMapper(apply_cleaners=True, symbols=symbols)
+
+  # Load the model.
+  if seed is not None:
+    Tensor.manual_seed(seed)
+    np.random.seed(seed)
+  net_g = load_model(text_mapper.symbols, hps, model_config)
+
+  return net_g, emotion_embedding, text_mapper, hps, model_has_multiple_speakers
+
+@contextmanager
+def output_stream(num_channels: int, sample_rate: int):
+  try:
+    p = pyaudio.PyAudio()
+    stream = p.open(format=pyaudio.paInt16, channels=num_channels, rate=sample_rate, output=True)
+    yield stream
+  except KeyboardInterrupt: pass
+  finally:
+    stream.stop_stream()
+    stream.close()
+    p.terminate()
+
+@contextmanager
+def log_writer():
+  try:
+    logs = []
+    yield logs
+  finally:
+    sep = "="*os.get_terminal_size()[1]
+    print(f"{sep[:-1]}\nCHAT LOG")
+    print(*logs, sep="\n")
+    print(sep)
+
+def listener(q: mp.Queue, event: mp.Event):
+  try:
+    p = pyaudio.PyAudio()
+    stream = p.open(format=pyaudio.paInt16, channels=1, rate=RATE, input=True, frames_per_buffer=CHUNK)
+    did_print = False
+    while True:
+      data = stream.read(CHUNK) # read data to avoid overflow
+      if event.is_set():
+        if not did_print:
+          print("listening")
+          did_print = True
+        q.put(((np.frombuffer(data, np.int16)/32768).astype(np.float32)*3))
+      else:
+        did_print = False
+  finally:
+    stream.stop_stream()
+    stream.close()
+    p.terminate()
+
+def mp_output_stream(q: mp.Queue, counter: mp.Value, num_channels: int, sample_rate: int):
+  with output_stream(num_channels, sample_rate) as stream:
+    while True:
+      try:
+        stream.write(q.get())
+        counter.value += 1
+      except KeyboardInterrupt:
+        break
+
+if __name__ == "__main__":
+  import nltk
+  nltk.download("punkt")
+  # Parse CLI arguments
+  parser = argparse.ArgumentParser("Have a tiny conversation with tinygrad")
+
+  # Whisper args
+  parser.add_argument("--whisper_model_name", type=str, default="tiny.en")
+
+  # LLAMA args
+  parser.add_argument("--llama_pre_prompt_path", type=Path, default=Path(__file__).parent / "conversation_data" / "pre_prompt_stacy.yaml", help="Path to yaml file which contains all pre-prompt data needed. ")
+  parser.add_argument("--llama_count", type=int, default=1000, help="Max number of tokens to generate")
+  parser.add_argument("--llama_temperature", type=float, default=0.7, help="Temperature in the softmax")
+  parser.add_argument("--llama_quantize", type=str, default=None, help="Quantize the weights to int8 or nf4 in memory")
+  parser.add_argument("--llama_model", type=Path, default=None, help="Folder with the original weights to load, or single .index.json, .safetensors or .bin file")
+  parser.add_argument("--llama_gen", type=str, default="tiny", required=False, help="Generation of the model to use")
+  parser.add_argument("--llama_size", type=str, default="1B-Chat", required=False, help="Size of model to use")
+  parser.add_argument("--llama_tokenizer", type=Path, default=None, required=False, help="Path to llama tokenizer.model")
+
+  # vits args
+  parser.add_argument("--vits_model_to_use", default="vctk", help="Specify the model to use. Default is 'vctk'.")
+  parser.add_argument("--vits_speaker_id", type=int, default=12, help="Specify the speaker ID. Default is 6.")
+  parser.add_argument("--vits_noise_scale", type=float, default=0.667, help="Specify the noise scale. Default is 0.667.")
+  parser.add_argument("--vits_noise_scale_w", type=float, default=0.8, help="Specify the noise scale w. Default is 0.8.")
+  parser.add_argument("--vits_length_scale", type=float, default=1, help="Specify the length scale. Default is 1.")
+  parser.add_argument("--vits_seed", type=int, default=None, help="Specify the seed (set to None if no seed). Default is 1337.")
+  parser.add_argument("--vits_num_channels", type=int, default=1, help="Specify the number of audio output channels. Default is 1.")
+  parser.add_argument("--vits_sample_width", type=int, default=2, help="Specify the number of bytes per sample, adjust if necessary. Default is 2.")
+  parser.add_argument("--vits_emotion_path", type=Path, default=None, help="Specify the path to emotion reference.")
+  parser.add_argument("--vits_estimate_max_y_length", type=str, default=False, help="If true, overestimate the output length and then trim it to the correct length, to prevent premature realization, much more performant for larger inputs, for smaller inputs not so much. Default is False.")
+  parser.add_argument("--vits_vocab_path", type=Path, default=None, help="Path to the TTS vocabulary.")
+
+  # conversation args
+  parser.add_argument("--max_sentence_length", type=int, default=20, help="Max words in one sentence to pass to vits")
+
+  args = parser.parse_args()
+
+  # Init models
+  model, enc = init_whisper(args.whisper_model_name)
+  synth, emotion_embedding, text_mapper, hps, model_has_multiple_speakers = init_vits(args.vits_model_to_use, args.vits_emotion_path, args.vits_speaker_id, args.vits_seed)
+
+  # Download tinyllama chat as a default model
+  if args.llama_model is None:
+    args.llama_model = fetch("https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/resolve/main/model.safetensors", "tinyllamachat.safetensors")
+    args.llama_gen = "tiny"
+    args.llama_size = "1B-Chat"
+  # Add 3 more tokens to the tokenizer
+  if args.llama_gen == "tiny" and args.llama_size.endswith("Chat"): args.llama_tokenizer = create_fixed_tokenizer()
+  tokenizer_path = args.llama_tokenizer or args.llama_model.parent / "tokenizer.model"
+  llama = LLaMa.build(args.llama_model, tokenizer_path, args.llama_gen, args.llama_size, args.llama_quantize)
+  toks, user_delim, resp_delim, start_pos, outputted = llama_prepare(llama, args.llama_temperature, args.llama_pre_prompt_path)
+
+  # Start child process for mic input
+  q = mp.Queue()
+  is_listening_event = mp.Event()
+  p = mp.Process(target=listener, args=(q, is_listening_event,))
+  p.daemon = True
+  p.start()
+
+  # Start child process for speaker output
+  out_q = mp.Queue()
+  out_counter = mp.Value("i", 0)
+  out_p = mp.Process(target=mp_output_stream, args=(out_q, out_counter, args.vits_num_channels, hps.data.sampling_rate,))
+  out_p.daemon = True
+  out_p.start()
+
+  # JIT tts
+  for i in ["Hello, I'm a chat bot", "I am capable of doing a lot of things"]:
+    tts(
+      i, synth, hps, emotion_embedding,
+      args.vits_speaker_id, args.vits_model_to_use, args.vits_noise_scale,
+      args.vits_noise_scale_w, args.vits_length_scale,
+      args.vits_estimate_max_y_length, text_mapper, model_has_multiple_speakers
+    )
+
+  # Start the pipeline
+  with log_writer() as log:
+    while True:
+      tokens = [enc._special_tokens["<|startoftranscript|>"], enc._special_tokens["<|notimestamps|>"]]
+      total = np.array([])
+      out_counter.value = 0
+
+      s = time.perf_counter()
+      is_listening_event.set()
+      prev_text = None
+      while True:
+        for _ in range(RATE // CHUNK): total = np.concatenate([total, q.get()])
+        txt = transcribe_waveform(model, enc, [total], truncate=True)
+        print(txt, end="\r")
+        if txt == "[BLANK_AUDIO]" or re.match(r"^\([\w+ ]+\)$", txt.strip()): continue
+        if prev_text is not None and prev_text == txt:
+          is_listening_event.clear()
+          break
+        prev_text = txt
+      print() # to avoid llama printing on the same line
+      log.append(f"{user_delim.capitalize()}: {txt}")
+
+      # Generate with llama
+      with Timing("llama generation: "):
+        outputted, start_pos, response = llama_generate(
+          llama, toks, outputted, txt, start_pos,
+          user_delim=user_delim, resp_delim=resp_delim, temperature=args.llama_temperature,
+          max_tokens=args.llama_count
+        )
+        log.append(f"{resp_delim.capitalize()}: {response}")
+
+      # Convert to voice
+      with Timing("tts: "):
+        sentences = nltk.sent_tokenize(response.replace('"', ""))
+        for i in sentences:
+          total = np.array([], dtype=np.int16)
+          for j in chunks(i.split(), args.max_sentence_length):
+            audio_data = tts(
+              " ".join(j), synth, hps, emotion_embedding,
+              args.vits_speaker_id, args.vits_model_to_use, args.vits_noise_scale,
+              args.vits_noise_scale_w, args.vits_length_scale,
+              args.vits_estimate_max_y_length, text_mapper, model_has_multiple_speakers
+            )
+            total = np.concatenate([total, audio_data])
+          out_q.put(total.tobytes())
+      while out_counter.value < len(sentences): continue
+      log.append(f"Total: {time.perf_counter() - s}")
@@ -0,0 +1,89 @@
+# load weights from
+# https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth
+# a rough copy of
+# https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/model.py
+import sys
+import ast
+import time
+import numpy as np
+from PIL import Image
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import getenv, fetch, Timing
+from tinygrad.engine.jit import TinyJit
+from extra.models.efficientnet import EfficientNet
+np.set_printoptions(suppress=True)
+
+# TODO: you should be able to put these in the jitted function
+bias = Tensor([0.485, 0.456, 0.406])
+scale = Tensor([0.229, 0.224, 0.225])
+
+@TinyJit
+def _infer(model, img):
+  img = img.permute((2,0,1))
+  img = img / 255.0
+  img = img - bias.reshape((1,-1,1,1))
+  img = img / scale.reshape((1,-1,1,1))
+  return model.forward(img).realize()
+
+def infer(model, img):
+  # preprocess image
+  aspect_ratio = img.size[0] / img.size[1]
+  img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0))))
+
+  img = np.array(img)
+  y0,x0=(np.asarray(img.shape)[:2]-224)//2
+  retimg = img = img[y0:y0+224, x0:x0+224]
+
+  # if you want to look at the image
+  """
+  import matplotlib.pyplot as plt
+  plt.imshow(img)
+  plt.show()
+  """
+
+  # run the net
+  out = _infer(model, Tensor(img.astype("float32"))).numpy()
+
+  # if you want to look at the outputs
+  """
+  import matplotlib.pyplot as plt
+  plt.plot(out[0])
+  plt.show()
+  """
+  return out, retimg
+
+if __name__ == "__main__":
+  # instantiate my net
+  model = EfficientNet(getenv("NUM", 0))
+  model.load_from_pretrained()
+
+  # category labels
+  lbls = ast.literal_eval(fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt").read_text())
+
+  # load image and preprocess
+  url = sys.argv[1] if len(sys.argv) >= 2 else "https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/showcase/stable_diffusion_by_tinygrad.jpg"
+  if url == 'webcam':
+    import cv2
+    cap = cv2.VideoCapture(0)
+    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
+    while 1:
+      _ = cap.grab() # discard one frame to circumvent capture buffering
+      ret, frame = cap.read()
+      img = Image.fromarray(frame[:, :, [2,1,0]])
+      lt = time.monotonic_ns()
+      out, retimg = infer(model, img)
+      print(f"{(time.monotonic_ns()-lt)*1e-6:7.2f} ms", np.argmax(out), np.max(out), lbls[np.argmax(out)])
+      SCALE = 3
+      simg = cv2.resize(retimg, (224*SCALE, 224*SCALE))
+      retimg = cv2.cvtColor(simg, cv2.COLOR_RGB2BGR)
+      cv2.imshow('capture', retimg)
+      if cv2.waitKey(1) & 0xFF == ord('q'):
+        break
+    cap.release()
+    cv2.destroyAllWindows()
+  else:
+    img = Image.open(fetch(url))
+    for i in range(getenv("CNT", 1)):
+      with Timing("did inference in "):
+        out, _ = infer(model, img)
+        print(np.argmax(out), np.max(out), lbls[np.argmax(out)])
@@ -0,0 +1,498 @@
+# pip3 install sentencepiece
+
+# This file incorporates code from the following:
+# Github Name                    | License | Link
+# black-forest-labs/flux         | Apache  | https://github.com/black-forest-labs/flux/tree/main/model_licenses
+
+from tinygrad import Tensor, nn, dtypes, TinyJit
+from tinygrad.nn.state import safe_load, load_state_dict
+from tinygrad.helpers import fetch, tqdm, colored
+from sdxl import FirstStage
+from extra.models.clip import FrozenClosedClipEmbedder
+from extra.models.t5 import T5Embedder
+import numpy as np
+
+import math, time, argparse, tempfile
+from typing import List, Dict, Optional, Union, Tuple, Callable
+from dataclasses import dataclass
+from pathlib import Path
+from PIL import Image
+
+urls:dict = {
+  "flux-schnell": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors",
+  "flux-dev": "https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/flux1-dev.sft",
+  "ae": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors",
+  "T5_1_of_2": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/text_encoder_2/model-00001-of-00002.safetensors",
+  "T5_2_of_2": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/text_encoder_2/model-00002-of-00002.safetensors",
+  "T5_tokenizer": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/tokenizer_2/spiece.model",
+  "clip": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/text_encoder/model.safetensors"
+}
+
+def tensor_identity(x:Tensor) -> Tensor: return x
+
+class AutoEncoder:
+  def __init__(self, scale_factor:float, shift_factor:float):
+    self.decoder = FirstStage.Decoder(128, 3, 3, 16, [1, 2, 4, 4], 2, 256)
+    self.scale_factor = scale_factor
+    self.shift_factor = shift_factor
+
+  def decode(self, z:Tensor) -> Tensor:
+    z = z / self.scale_factor + self.shift_factor
+    return self.decoder(z)
+
+# Conditioner
+class ClipEmbedder(FrozenClosedClipEmbedder):
+  def __call__(self, texts:Union[str, List[str], Tensor]) -> Tensor:
+    if isinstance(texts, str): texts = [texts]
+    assert isinstance(texts, (list,tuple)), f"expected list of strings, got {type(texts).__name__}"
+    tokens = Tensor.cat(*[Tensor(self.tokenizer.encode(text)) for text in texts], dim=0)
+    return self.transformer.text_model(tokens.reshape(len(texts),-1))[:, tokens.argmax(-1)]
+
+# https://github.com/black-forest-labs/flux/blob/main/src/flux/math.py
+def attention(q:Tensor, k:Tensor, v:Tensor, pe:Tensor) -> Tensor:
+  q, k = apply_rope(q, k, pe)
+  x = Tensor.scaled_dot_product_attention(q, k, v)
+  return x.rearrange("B H L D -> B L (H D)")
+
+def rope(pos:Tensor, dim:int, theta:int) -> Tensor:
+  assert dim % 2 == 0
+  scale = Tensor.arange(0, dim, 2, dtype=dtypes.float32, device=pos.device) / dim # NOTE: this is torch.float64 in reference implementation
+  omega = 1.0 / (theta**scale)
+  out = Tensor.einsum("...n,d->...nd", pos, omega)
+  out = Tensor.stack(Tensor.cos(out), -Tensor.sin(out), Tensor.sin(out), Tensor.cos(out), dim=-1)
+  out = out.rearrange("b n d (i j) -> b n d i j", i=2, j=2)
+  return out.float()
+
+def apply_rope(xq:Tensor, xk:Tensor, freqs_cis:Tensor) -> Tuple[Tensor, Tensor]:
+  xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+  xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+  xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+  xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+  return xq_out.reshape(*xq.shape).cast(xq.dtype), xk_out.reshape(*xk.shape).cast(xk.dtype)
+
+
+# https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py
+class EmbedND:
+  def __init__(self, dim:int, theta:int, axes_dim:List[int]):
+    self.dim = dim
+    self.theta = theta
+    self.axes_dim = axes_dim
+
+  def __call__(self, ids:Tensor) -> Tensor:
+    n_axes = ids.shape[-1]
+    emb = Tensor.cat(*[rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], dim=-3)
+    return emb.unsqueeze(1)
+
+class MLPEmbedder:
+  def __init__(self, in_dim:int, hidden_dim:int):
+    self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+    self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+  def __call__(self, x:Tensor) -> Tensor:
+    return self.out_layer(self.in_layer(x).silu())
+
+class QKNorm:
+  def __init__(self, dim:int):
+    self.query_norm = nn.RMSNorm(dim)
+    self.key_norm = nn.RMSNorm(dim)
+
+  def __call__(self, q:Tensor, k:Tensor) -> Tuple[Tensor, Tensor]:
+    return self.query_norm(q), self.key_norm(k)
+
+class SelfAttention:
+  def __init__(self, dim:int, num_heads:int = 8, qkv_bias:bool = False):
+    self.num_heads = num_heads
+    head_dim = dim // num_heads
+
+    self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+    self.norm = QKNorm(head_dim)
+    self.proj = nn.Linear(dim, dim)
+
+  def __call__(self, x:Tensor, pe:Tensor) -> Tensor:
+    qkv = self.qkv(x)
+    q, k, v = qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+    q, k = self.norm(q, k)
+    x = attention(q, k, v, pe=pe)
+    return self.proj(x)
+
+@dataclass
+class ModulationOut:
+  shift:Tensor
+  scale:Tensor
+  gate:Tensor
+
+class Modulation:
+  def __init__(self, dim:int, double:bool):
+    self.is_double = double
+    self.multiplier = 6 if double else 3
+    self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
+
+  def __call__(self, vec:Tensor) -> Tuple[ModulationOut, Optional[ModulationOut]]:
+    out = self.lin(vec.silu())[:, None, :].chunk(self.multiplier, dim=-1)
+    return ModulationOut(*out[:3]), ModulationOut(*out[3:]) if self.is_double else None
+
+class DoubleStreamBlock:
+  def __init__(self, hidden_size:int, num_heads:int, mlp_ratio:float, qkv_bias:bool = False):
+    mlp_hidden_dim = int(hidden_size * mlp_ratio)
+    self.num_heads = num_heads
+    self.hidden_size = hidden_size
+    self.img_mod = Modulation(hidden_size, double=True)
+    self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+    self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+
+    self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+    self.img_mlp = [nn.Linear(hidden_size, mlp_hidden_dim, bias=True), Tensor.gelu, nn.Linear(mlp_hidden_dim, hidden_size, bias=True)]
+
+    self.txt_mod = Modulation(hidden_size, double=True)
+    self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+    self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
+
+    self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+    self.txt_mlp = [nn.Linear(hidden_size, mlp_hidden_dim, bias=True), Tensor.gelu, nn.Linear(mlp_hidden_dim, hidden_size, bias=True)]
+
+  def __call__(self, img:Tensor, txt:Tensor, vec:Tensor, pe:Tensor) -> tuple[Tensor, Tensor]:
+    img_mod1, img_mod2 = self.img_mod(vec)
+    txt_mod1, txt_mod2 = self.txt_mod(vec)
+    assert img_mod2 is not None and txt_mod2 is not None
+    # prepare image for attention
+    img_modulated = self.img_norm1(img)
+    img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
+    img_qkv = self.img_attn.qkv(img_modulated)
+    img_q, img_k, img_v = img_qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+    img_q, img_k = self.img_attn.norm(img_q, img_k)
+
+    # prepare txt for attention
+    txt_modulated = self.txt_norm1(txt)
+    txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
+    txt_qkv = self.txt_attn.qkv(txt_modulated)
+    txt_q, txt_k, txt_v = txt_qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+    txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k)
+
+    # run actual attention
+    q = Tensor.cat(txt_q, img_q, dim=2)
+    k = Tensor.cat(txt_k, img_k, dim=2)
+    v = Tensor.cat(txt_v, img_v, dim=2)
+
+    attn = attention(q, k, v, pe=pe)
+    txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+    # calculate the img bloks
+    img = img + img_mod1.gate * self.img_attn.proj(img_attn)
+    img = img + img_mod2.gate * ((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift).sequential(self.img_mlp)
+
+    # calculate the txt bloks
+    txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
+    txt = txt + txt_mod2.gate * ((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift).sequential(self.txt_mlp)
+    return img, txt
+
+
+class SingleStreamBlock:
+  """
+  A DiT block with parallel linear layers as described in
+  https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+  """
+
+  def __init__(self,hidden_size:int, num_heads:int, mlp_ratio:float=4.0, qk_scale:Optional[float]=None):
+    self.hidden_dim = hidden_size
+    self.num_heads = num_heads
+    head_dim = hidden_size // num_heads
+    self.scale = qk_scale or head_dim**-0.5
+
+    self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+    # qkv and mlp_in
+    self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
+    # proj and mlp_out
+    self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
+
+    self.norm = QKNorm(head_dim)
+
+    self.hidden_size = hidden_size
+    self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+    self.mlp_act = Tensor.gelu
+    self.modulation = Modulation(hidden_size, double=False)
+
+  def __call__(self, x:Tensor, vec:Tensor, pe:Tensor) -> Tensor:
+    mod, _ = self.modulation(vec)
+    x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
+    qkv, mlp = Tensor.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
+    q, k, v = qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+    q, k = self.norm(q, k)
+
+    # compute attention
+    attn = attention(q, k, v, pe=pe)
+    # compute activation in mlp stream, cat again and run second linear layer
+    output = self.linear2(Tensor.cat(attn, self.mlp_act(mlp), dim=2))
+    return x + mod.gate * output
+
+
+class LastLayer:
+  def __init__(self, hidden_size:int, patch_size:int, out_channels:int):
+    self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+    self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+    self.adaLN_modulation:List[Callable[[Tensor], Tensor]] = [Tensor.silu, nn.Linear(hidden_size, 2 * hidden_size, bias=True)]
+
+  def __call__(self, x:Tensor, vec:Tensor) -> Tensor:
+    shift, scale = vec.sequential(self.adaLN_modulation).chunk(2, dim=1)
+    x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
+    return self.linear(x)
+
+def timestep_embedding(t:Tensor, dim:int, max_period:int=10000, time_factor:float=1000.0) -> Tensor:
+  """
+  Create sinusoidal timestep embeddings.
+  :param t: a 1-D Tensor of N indices, one per batch element.
+                    These may be fractional.
+  :param dim: the dimension of the output.
+  :param max_period: controls the minimum frequency of the embeddings.
+  :return: an (N, D) Tensor of positional embeddings.
+  """
+  t = time_factor * t
+  half = dim // 2
+  freqs = Tensor.exp(-math.log(max_period) * Tensor.arange(0, stop=half, dtype=dtypes.float32) / half).to(t.device)
+
+  args = t[:, None].float() * freqs[None]
+  embedding = Tensor.cat(Tensor.cos(args), Tensor.sin(args), dim=-1)
+  if dim % 2:  embedding = Tensor.cat(*[embedding, Tensor.zeros_like(embedding[:, :1])], dim=-1)
+  if Tensor.is_floating_point(t):  embedding = embedding.cast(t.dtype)
+  return embedding
+
+# https://github.com/black-forest-labs/flux/blob/main/src/flux/model.py
+class Flux:
+  """
+  Transformer model for flow matching on sequences.
+  """
+
+  def __init__(
+      self,
+      guidance_embed:bool,
+      in_channels:int = 64,
+      vec_in_dim:int = 768,
+      context_in_dim:int = 4096,
+      hidden_size:int = 3072,
+      mlp_ratio:float = 4.0,
+      num_heads:int = 24,
+      depth:int = 19,
+      depth_single_blocks:int = 38,
+      axes_dim:Optional[List[int]] = None,
+      theta:int = 10_000,
+      qkv_bias:bool = True,
+      ):
+
+    axes_dim = axes_dim or [16, 56, 56]
+    self.guidance_embed = guidance_embed
+    self.in_channels = in_channels
+    self.out_channels = self.in_channels
+    if hidden_size % num_heads != 0:
+      raise ValueError(f"Hidden size {hidden_size} must be divisible by num_heads {num_heads}")
+    pe_dim = hidden_size // num_heads
+    if sum(axes_dim) != pe_dim:
+      raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}")
+    self.hidden_size = hidden_size
+    self.num_heads = num_heads
+    self.pe_embedder = EmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim)
+    self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
+    self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+    self.vector_in = MLPEmbedder(vec_in_dim, self.hidden_size)
+    self.guidance_in:Callable[[Tensor], Tensor] = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if guidance_embed else tensor_identity
+    self.txt_in = nn.Linear(context_in_dim, self.hidden_size)
+
+    self.double_blocks = [DoubleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias) for _ in range(depth)]
+    self.single_blocks = [SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio) for _ in range(depth_single_blocks)]
+    self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
+
+  def __call__(self, img:Tensor, img_ids:Tensor, txt:Tensor, txt_ids:Tensor, timesteps:Tensor, y:Tensor, guidance:Optional[Tensor] = None) -> Tensor:
+    if img.ndim != 3 or txt.ndim != 3:
+      raise ValueError("Input img and txt tensors must have 3 dimensions.")
+    # running on sequences img
+    img = self.img_in(img)
+    vec = self.time_in(timestep_embedding(timesteps, 256))
+    if self.guidance_embed:
+      if guidance is None:
+        raise ValueError("Didn't get guidance strength for guidance distilled model.")
+      vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+    vec = vec + self.vector_in(y)
+    txt = self.txt_in(txt)
+    ids = Tensor.cat(txt_ids, img_ids, dim=1)
+    pe = self.pe_embedder(ids)
+    for double_block in self.double_blocks:
+      img, txt = double_block(img=img, txt=txt, vec=vec, pe=pe)
+
+    img = Tensor.cat(txt, img, dim=1)
+    for single_block in self.single_blocks:
+      img = single_block(img, vec=vec, pe=pe)
+
+    img = img[:, txt.shape[1] :, ...]
+
+    return self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+
+# https://github.com/black-forest-labs/flux/blob/main/src/flux/util.py
+def load_flow_model(name:str, model_path:str):
+  # Loading Flux
+  print("Init model")
+  model = Flux(guidance_embed=(name != "flux-schnell"))
+  if not model_path: model_path = fetch(urls[name])
+  state_dict = {k.replace("scale", "weight"): v for k, v in safe_load(model_path).items()}
+  load_state_dict(model, state_dict)
+  return model
+
+def load_T5(max_length:int=512):
+  # max length 64, 128, 256 and 512 should work (if your sequence is short enough)
+  print("Init T5")
+  T5 = T5Embedder(max_length, fetch(urls["T5_tokenizer"]))
+  pt_1 = fetch(urls["T5_1_of_2"])
+  pt_2 = fetch(urls["T5_2_of_2"])
+  load_state_dict(T5.encoder, safe_load(pt_1) | safe_load(pt_2), strict=False)
+  return T5
+
+def load_clip():
+  print("Init Clip")
+  clip = ClipEmbedder()
+  load_state_dict(clip.transformer, safe_load(fetch(urls["clip"])))
+  return clip
+
+def load_ae() -> AutoEncoder:
+  # Loading the autoencoder
+  print("Init AE")
+  ae = AutoEncoder(0.3611, 0.1159)
+  load_state_dict(ae, safe_load(fetch(urls["ae"])))
+  return ae
+
+# https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py
+def prepare(T5:T5Embedder, clip:ClipEmbedder, img:Tensor, prompt:Union[str, List[str]]) -> Dict[str, Tensor]:
+  bs, _, h, w = img.shape
+  if bs == 1 and not isinstance(prompt, str):
+    bs = len(prompt)
+
+  img = img.rearrange("b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+  if img.shape[0] == 1 and bs > 1:
+    img = img.expand((bs, *img.shape[1:]))
+
+  img_ids = Tensor.zeros(h // 2, w // 2, 3).contiguous()
+  img_ids[..., 1] = img_ids[..., 1] + Tensor.arange(h // 2)[:, None]
+  img_ids[..., 2] = img_ids[..., 2] + Tensor.arange(w // 2)[None, :]
+  img_ids = img_ids.rearrange("h w c -> 1 (h w) c")
+  img_ids = img_ids.expand((bs, *img_ids.shape[1:]))
+
+  if isinstance(prompt, str):
+    prompt = [prompt]
+  txt = T5(prompt).realize()
+  if txt.shape[0] == 1 and bs > 1:
+    txt = txt.expand((bs, *txt.shape[1:]))
+  txt_ids = Tensor.zeros(bs, txt.shape[1], 3)
+
+  vec = clip(prompt).realize()
+  if vec.shape[0] == 1 and bs > 1:
+    vec = vec.expand((bs, *vec.shape[1:]))
+
+  return {"img": img, "img_ids": img_ids.to(img.device), "txt": txt.to(img.device), "txt_ids": txt_ids.to(img.device), "vec": vec.to(img.device)}
+
+
+def get_schedule(num_steps:int, image_seq_len:int, base_shift:float=0.5, max_shift:float=1.15, shift:bool=True) -> List[float]:
+  # extra step for zero
+  step_size = -1.0 / num_steps
+  timesteps = Tensor.arange(1, 0 + step_size, step_size)
+
+  # shifting the schedule to favor high timesteps for higher signal images
+  if shift:
+    # estimate mu based on linear estimation between two points
+    mu = 0.5 + (max_shift - base_shift) * (image_seq_len - 256) / (4096 - 256)
+    timesteps = math.exp(mu) / (math.exp(mu) + (1 / timesteps - 1))
+  return timesteps.tolist()
+
+@TinyJit
+def run(model, *args): return model(*args).realize()
+
+def denoise(model, img:Tensor, img_ids:Tensor, txt:Tensor, txt_ids:Tensor, vec:Tensor, timesteps:List[float], guidance:float=4.0) -> Tensor:
+  # this is ignored for schnell
+  guidance_vec = Tensor((guidance,), device=img.device, dtype=img.dtype).expand((img.shape[0],))
+  for t_curr, t_prev in tqdm(list(zip(timesteps[:-1], timesteps[1:])), "Denoising"):
+    t_vec = Tensor((t_curr,), device=img.device, dtype=img.dtype).expand((img.shape[0],))
+    pred = run(model, img, img_ids, txt, txt_ids, t_vec, vec, guidance_vec)
+    img = img + (t_prev - t_curr) * pred
+
+  return img
+
+def unpack(x:Tensor, height:int, width:int) -> Tensor:
+  return x.rearrange("b (h w) (c ph pw) -> b c (h ph) (w pw)", h=math.ceil(height / 16), w=math.ceil(width / 16), ph=2, pw=2)
+
+# https://github.com/black-forest-labs/flux/blob/main/src/flux/cli.py
+if __name__ == "__main__":
+  default_prompt = "bananas and a can of coke"
+  parser = argparse.ArgumentParser(description="Run Flux.1", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+  parser.add_argument("--name",       type=str,   default="flux-schnell", help="Name of the model to load")
+  parser.add_argument("--model_path", type=str,   default="",             help="path of the model file")
+  parser.add_argument("--width",      type=int,   default=512,            help="width of the sample in pixels (should be a multiple of 16)")
+  parser.add_argument("--height",     type=int,   default=512,            help="height of the sample in pixels (should be a multiple of 16)")
+  parser.add_argument("--seed",       type=int,   default=None,           help="Set a seed for sampling")
+  parser.add_argument("--prompt",     type=str,   default=default_prompt, help="Prompt used for sampling")
+  parser.add_argument('--out',        type=str,   default=Path(tempfile.gettempdir()) / "rendered.png", help="Output filename")
+  parser.add_argument("--num_steps",  type=int,   default=None,           help="number of sampling steps (default 4 for schnell, 50 for guidance distilled)") #noqa:E501
+  parser.add_argument("--guidance",   type=float, default=3.5,            help="guidance value used for guidance distillation")
+  parser.add_argument("--output_dir", type=str,   default="output",       help="output directory")
+  args = parser.parse_args()
+
+  if args.name not in ["flux-schnell", "flux-dev"]:
+    raise ValueError(f"Got unknown model name: {args.name}, chose from flux-schnell and flux-dev")
+
+  if args.num_steps is None:
+    args.num_steps = 4 if args.name == "flux-schnell" else 50
+
+  # allow for packing and conversion to latent space
+  height = 16 * (args.height // 16)
+  width = 16 * (args.width // 16)
+
+  if args.seed is None: args.seed = Tensor._seed
+  else: Tensor.manual_seed(args.seed)
+
+  print(f"Generating with seed {args.seed}:\n{args.prompt}")
+  t0 = time.perf_counter()
+
+  # prepare input noise
+  x = Tensor.randn(1, 16, 2 * math.ceil(height / 16), 2 * math.ceil(width / 16), dtype="bfloat16")
+
+  # load text embedders
+  T5 = load_T5(max_length=256 if args.name == "flux-schnell" else 512)
+  clip = load_clip()
+
+  # embed text to get inputs for model
+  inp = prepare(T5, clip, x, prompt=args.prompt)
+  timesteps = get_schedule(args.num_steps, inp["img"].shape[1], shift=(args.name != "flux-schnell"))
+
+  # done with text embedders
+  del T5, clip
+
+  # load model
+  model = load_flow_model(args.name, args.model_path)
+
+  # denoise initial noise
+  x = denoise(model, **inp, timesteps=timesteps, guidance=args.guidance)
+
+  # done with model
+  del model, run
+
+  # load autoencoder
+  ae = load_ae()
+
+  # decode latents to pixel space
+  x = unpack(x.float(), height, width)
+  x = ae.decode(x).realize()
+
+  t1 = time.perf_counter()
+  print(f"Done in {t1 - t0:.1f}s. Saving {args.out}")
+
+  # bring into PIL format and save
+  x = x.clamp(-1, 1)
+  x = x[0].rearrange("c h w -> h w c")
+  x = (127.5 * (x + 1.0)).cast("uint8")
+
+  img = Image.fromarray(x.numpy())
+
+  img.save(args.out)
+
+  # validation!
+  if args.prompt == default_prompt and args.name=="flux-schnell" and args.seed == 0 and args.width == args.height == 512:
+    ref_image = Tensor(np.array(Image.open("examples/flux1_seed0.png")))
+    distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item()
+    assert distance < 4e-3, colored(f"validation failed with {distance=}", "red")
+    print(colored(f"output validated with {distance=}", "green"))
@@ -232,7 +232,7 @@ if __name__ == "__main__":
  gpt2 = GPT2.build_gguf(args.model_size) if args.model_size.startswith("gpt2_gguf_") else GPT2.build(args.model_size)

  if args.benchmark != -1:
-    gpt2.model(Tensor.randint(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
+    gpt2.model(Tensor.rand(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
  else:
    texts = gpt2.generate(args.prompt, args.count, args.temperature, timing=args.timing, batch_size=args.batch_size)
    if not args.noshow:
@@ -1,108 +0,0 @@
-import itertools
-from typing import Callable
-from tinygrad import nn, Tensor, dtypes, Device, TinyJit
-from tinygrad.helpers import getenv, trange, partition
-
-class Model:
-  def __init__(self):
-    self.layers: list[Callable[[Tensor], Tensor]] = [
-      nn.Conv2d(1, 32, 5), Tensor.relu,
-      nn.Conv2d(32, 32, 5), Tensor.relu,
-      nn.BatchNorm(32), Tensor.max_pool2d,
-      nn.Conv2d(32, 64, 3), Tensor.relu,
-      nn.Conv2d(64, 64, 3), Tensor.relu,
-      nn.BatchNorm(64), Tensor.max_pool2d,
-      lambda x: x.flatten(1), nn.Linear(576, 10)]
-
-  def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
-
-# TODO: refactor this into optim/onnx
-def functional_adam(g:Tensor, m:Tensor, v:Tensor, b1_t:Tensor, b2_t:Tensor, lr=0.001, b1=0.9, b2=0.999, eps=1e-6) -> Tensor:
-  b1_t *= b1
-  b2_t *= b2
-  m.assign(b1 * m + (1.0 - b1) * g)
-  v.assign(b2 * v + (1.0 - b2) * (g * g))
-  m_hat = m / (1.0 - b1_t)
-  v_hat = v / (1.0 - b2_t)
-  return lr * (m_hat / (v_hat.sqrt() + eps))
-
-if __name__ == "__main__":
-  BS = getenv("BS", 512)
-  ACC_STEPS = getenv("ACC_STEPS", 8)
-
-  X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
-  model = Model()
-
-  params = nn.state.get_parameters(model)
-
-  # init params, set requires grad on the ones we need gradients of
-  for x in params:
-    if x.requires_grad is None: x.requires_grad_()
-    x.replace(x.contiguous())
-  Tensor.realize(*params)
-
-  # split params (with grads) and buffers (without)
-  params, buffers = partition(params, lambda x: x.requires_grad)
-  print(f"params: {len(params)} buffers: {len(buffers)}")
-
-  # optim params
-  pos_params = list(itertools.accumulate(params, lambda x,y: x+y.numel(), initial=0))
-  adam_m = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
-  adam_v = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
-  adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
-  adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
-  adam_params = [adam_m, adam_v, adam_b1_t, adam_b2_t]
-
-  # create loss and grads. init all state so the JIT works on microbatch
-  for x in params: x.assign(x.detach())
-  loss = Tensor.zeros(tuple()).contiguous()
-  grads = Tensor.zeros(pos_params[-1]).contiguous()
-  Tensor.realize(*params, *buffers, *adam_params, loss, grads)
-
-  @TinyJit
-  @Tensor.train()
-  def microbatch():
-    samples = Tensor.randint(BS // ACC_STEPS, high=X_train.shape[0])
-    for t in params: t.grad = None
-    # divide by ACC_STEPS at the loss
-    uloss = (model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]) / ACC_STEPS).backward()
-    ugrads = Tensor.cat(*[t.grad.contiguous().flatten() for t in params], dim=0)
-    for t in params: t.grad = None
-    # concat the grads and assign them
-    loss.assign(loss + uloss)
-    grads.assign(grads + ugrads)
-    Tensor.realize(*params, *buffers, loss, grads)
-
-  @TinyJit
-  def optimizer():
-    # run optimizer (on CPU, where adam params live)
-    delta = functional_adam(grads.to("CPU"), adam_m, adam_v, adam_b1_t, adam_b2_t)
-
-    # update the params, copying back the delta one at a time to avoid OOM
-    # NOTE: the scheduler is ordering things poorly, all the copies are happening before the adds
-    for j,tt in enumerate(params):
-      tt.assign(tt.detach() - delta[pos_params[j]:pos_params[j+1]].reshape(tt.shape).to(Device.DEFAULT))
-
-    # realize everything, zero out loss and grads
-    loss.assign(Tensor.zeros_like(loss))
-    grads.assign(Tensor.zeros_like(grads))
-    Tensor.realize(*params, *adam_params, loss, grads)
-
-  @TinyJit
-  def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
-
-  test_acc = float('nan')
-  for i in (t:=trange(getenv("STEPS", 70))):
-    # microbatch sets the gradients
-    for _ in range(ACC_STEPS): microbatch()
-
-    # get the loss before the optimizer clears it
-    # this is already realized so this isn't a schedule
-    loss_item = loss.item()
-
-    # run the optimizer
-    optimizer()
-
-    # eval
-    if i%10 == 9: test_acc = get_test_acc().item()
-    t.set_description(f"loss: {loss_item:6.2f} test_accuracy: {test_acc:5.2f}%")
@@ -1,6 +1,8 @@
 from pathlib import Path
 from typing import List
 import json, argparse, random, time, os
+import tiktoken
+from tiktoken.load import load_tiktoken_bpe
 from extra.models.llama import Transformer, convert_from_huggingface, convert_from_gguf, fix_bf16
 from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters, gguf_load
 from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
@@ -10,8 +12,6 @@ from extra.bench_log import BenchEvent, WallTimeEvent
 class Tokenizer:
  pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
  def __init__(self, model_path: str):
-    import tiktoken
-    from tiktoken.load import load_tiktoken_bpe
    mergeable_ranks = load_tiktoken_bpe(model_path)
    self.num_base_tokens = len(mergeable_ranks)
    special_tokens = [
@@ -145,41 +145,6 @@ def NF4Linear(block_size):
      return new_state_dict
  return _NF4Linear

-def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
-  fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
-  fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
-  scale = fp8_max / x.abs().max()
-  x_scl_sat = (x * scale).clamp(fp8_min, fp8_max)
-  return x_scl_sat.cast(dtype), scale.float().reciprocal()
-
-class FP8Linear:
-  def __init__(self, in_features, out_features, bias=True):
-    self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.fp8e4m3)
-    self.bias = Tensor.empty(out_features, dtype=dtypes.float16) if bias else None
-    self.weight_scale = Tensor.empty((), dtype=dtypes.float16)
-
-  def __call__(self, x:Tensor):
-    y = x.dot(self.weight.T.cast(dtypes.float32)) * self.weight_scale
-    if self.bias is not None: y = y + self.bias.cast(y.dtype)
-    return y.cast(x.dtype)
-
-  @staticmethod
-  def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
-    assert not quantize_embeds
-    new_tensors = {}
-    for name,v in tensors.items():
-      if "feed_forward" in name or "attention.w" in name:
-        assert "weight" in name, name
-        fp8_weight, scale = quantize_to_fp8(v)
-        new_tensors[name] = fp8_weight
-        new_tensors[name.replace('weight', 'weight_scale')] = scale.cast(scale_dtype)
-        if isinstance(device, tuple):
-          new_tensors[name].shard_(device, axis=-1)
-          new_tensors[name.replace('weight', 'weight_scale')].shard_(device, axis=None)
-      else:
-        new_tensors[name] = v
-    return new_tensors
-
 MODEL_PARAMS = {
  "1B": {
    "args": {"dim": 2048, "n_heads": 32, "n_kv_heads": 8, "n_layers": 16, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 8192},
@@ -202,7 +167,6 @@ def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dt
  # build model
  if quantize == "int8": linear, embedding, quantize_embeds = Int8Linear, Int8Embedding, True
  elif quantize == "nf4": linear, embedding, quantize_embeds = NF4Linear(64), nn.Embedding, False
-  elif quantize == "fp8": linear, embedding, quantize_embeds = FP8Linear, nn.Embedding, False
  else: linear, embedding, quantize_embeds = nn.Linear, nn.Embedding, False
  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, embedding=embedding, max_context=max_context, jit=True)

@@ -278,7 +242,7 @@ if __name__ == "__main__":
  parser.add_argument("--model", type=Path, help="Model path")
  parser.add_argument("--size", choices=["1B", "8B", "70B", "405B"], default="1B", help="Model size")
  parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
-  parser.add_argument("--quantize", choices=["int8", "nf4", "float16", "fp8"], help="Quantization method")
+  parser.add_argument("--quantize", choices=["int8", "nf4", "float16"], help="Quantization method")
  parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
  parser.add_argument("--host", type=str, default="0.0.0.0", help="Web server bind address")
  parser.add_argument("--port", type=int, default=7776, help="Web server port")
@@ -0,0 +1,299 @@
+from extra.models.mask_rcnn import MaskRCNN
+from extra.models.resnet import ResNet
+from extra.models.mask_rcnn import BoxList
+from torch.nn import functional as F
+from torchvision import transforms as T
+from torchvision.transforms import functional as Ft
+import random
+from tinygrad.tensor import Tensor
+from PIL import Image
+import numpy as np
+import torch
+import argparse
+import cv2
+
+
+class Resize:
+  def __init__(self, min_size, max_size):
+    if not isinstance(min_size, (list, tuple)):
+      min_size = (min_size,)
+    self.min_size = min_size
+    self.max_size = max_size
+
+  # modified from torchvision to add support for max size
+  def get_size(self, image_size):
+    w, h = image_size
+    size = random.choice(self.min_size)
+    max_size = self.max_size
+    if max_size is not None:
+      min_original_size = float(min((w, h)))
+      max_original_size = float(max((w, h)))
+      if max_original_size / min_original_size * size > max_size:
+        size = int(round(max_size * min_original_size / max_original_size))
+
+      if (w <= h and w == size) or (h <= w and h == size):
+        return (h, w)
+
+      if w < h:
+        ow = size
+        oh = int(size * h / w)
+      else:
+        oh = size
+        ow = int(size * w / h)
+
+      return (oh, ow)
+
+  def __call__(self, image):
+    size = self.get_size(image.size)
+    image = Ft.resize(image, size)
+    return image
+
+
+class Normalize:
+  def __init__(self, mean, std, to_bgr255=True):
+    self.mean = mean
+    self.std = std
+    self.to_bgr255 = to_bgr255
+
+  def __call__(self, image):
+    if self.to_bgr255:
+      image = image[[2, 1, 0]] * 255
+    else:
+      image = image[[0, 1, 2]] * 255
+    image = Ft.normalize(image, mean=self.mean, std=self.std)
+    return image
+
+transforms = lambda size_scale: T.Compose(
+  [
+    Resize(int(800*size_scale), int(1333*size_scale)),
+    T.ToTensor(),
+    Normalize(
+      mean=[102.9801, 115.9465, 122.7717], std=[1., 1., 1.], to_bgr255=True
+    ),
+  ]
+)
+
+def expand_boxes(boxes, scale):
+  w_half = (boxes[:, 2] - boxes[:, 0]) * .5
+  h_half = (boxes[:, 3] - boxes[:, 1]) * .5
+  x_c = (boxes[:, 2] + boxes[:, 0]) * .5
+  y_c = (boxes[:, 3] + boxes[:, 1]) * .5
+
+  w_half *= scale
+  h_half *= scale
+
+  boxes_exp = torch.zeros_like(boxes)
+  boxes_exp[:, 0] = x_c - w_half
+  boxes_exp[:, 2] = x_c + w_half
+  boxes_exp[:, 1] = y_c - h_half
+  boxes_exp[:, 3] = y_c + h_half
+  return boxes_exp
+
+
+def expand_masks(mask, padding):
+  N = mask.shape[0]
+  M = mask.shape[-1]
+  pad2 = 2 * padding
+  scale = float(M + pad2) / M
+  padded_mask = mask.new_zeros((N, 1, M + pad2, M + pad2))
+  padded_mask[:, :, padding:-padding, padding:-padding] = mask
+  return padded_mask, scale
+
+
+def paste_mask_in_image(mask, box, im_h, im_w, thresh=0.5, padding=1):
+  # TODO: remove torch
+  mask = torch.tensor(mask.numpy())
+  box = torch.tensor(box.numpy())
+  padded_mask, scale = expand_masks(mask[None], padding=padding)
+  mask = padded_mask[0, 0]
+  box = expand_boxes(box[None], scale)[0]
+  box = box.to(dtype=torch.int32)
+
+  TO_REMOVE = 1
+  w = int(box[2] - box[0] + TO_REMOVE)
+  h = int(box[3] - box[1] + TO_REMOVE)
+  w = max(w, 1)
+  h = max(h, 1)
+
+  mask = mask.expand((1, 1, -1, -1))
+
+  mask = mask.to(torch.float32)
+  mask = F.interpolate(mask, size=(h, w), mode='bilinear', align_corners=False)
+  mask = mask[0][0]
+
+  if thresh >= 0:
+    mask = mask > thresh
+  else:
+    mask = (mask * 255).to(torch.uint8)
+
+  im_mask = torch.zeros((im_h, im_w), dtype=torch.uint8)
+  x_0 = max(box[0], 0)
+  x_1 = min(box[2] + 1, im_w)
+  y_0 = max(box[1], 0)
+  y_1 = min(box[3] + 1, im_h)
+
+  im_mask[y_0:y_1, x_0:x_1] = mask[
+                              (y_0 - box[1]): (y_1 - box[1]), (x_0 - box[0]): (x_1 - box[0])
+                              ]
+  return im_mask
+
+
+class Masker:
+  def __init__(self, threshold=0.5, padding=1):
+    self.threshold = threshold
+    self.padding = padding
+
+  def forward_single_image(self, masks, boxes):
+    boxes = boxes.convert("xyxy")
+    im_w, im_h = boxes.size
+    res = [
+      paste_mask_in_image(mask[0], box, im_h, im_w, self.threshold, self.padding)
+      for mask, box in zip(masks, boxes.bbox)
+    ]
+    if len(res) > 0:
+      res = torch.stack(*res, dim=0)[:, None]
+    else:
+      res = masks.new_empty((0, 1, masks.shape[-2], masks.shape[-1]))
+    return Tensor(res.numpy())
+
+  def __call__(self, masks, boxes):
+    if isinstance(boxes, BoxList):
+      boxes = [boxes]
+
+    results = []
+    for mask, box in zip(masks, boxes):
+      result = self.forward_single_image(mask, box)
+      results.append(result)
+    return results
+
+
+masker = Masker(threshold=0.5, padding=1)
+
+def select_top_predictions(predictions, confidence_threshold=0.9):
+  scores = predictions.get_field("scores").numpy()
+  keep = [idx for idx, score in enumerate(scores) if score > confidence_threshold]
+  return predictions[keep]
+
+def compute_prediction(original_image, model, confidence_threshold, size_scale=1.0):
+  image = transforms(size_scale)(original_image).numpy()
+  image = Tensor(image, requires_grad=False)
+  predictions = model(image)
+  prediction = predictions[0]
+  prediction = select_top_predictions(prediction, confidence_threshold)
+  width, height = original_image.size
+  prediction = prediction.resize((width, height))
+
+  if prediction.has_field("mask"):
+    masks = prediction.get_field("mask")
+    masks = masker([masks], [prediction])[0]
+    prediction.add_field("mask", masks)
+  return prediction
+
+def compute_prediction_batched(batch, model, size_scale=1.0):
+  imgs = []
+  for img in batch:
+    imgs.append(transforms(size_scale)(img).numpy())
+  image = [Tensor(image, requires_grad=False) for image in imgs]
+  predictions = model(image)
+  del image
+  return predictions
+
+palette = np.array([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
+
+def findContours(*args, **kwargs):
+  if cv2.__version__.startswith('4'):
+    contours, hierarchy = cv2.findContours(*args, **kwargs)
+  elif cv2.__version__.startswith('3'):
+    _, contours, hierarchy = cv2.findContours(*args, **kwargs)
+  return contours, hierarchy
+
+def compute_colors_for_labels(labels):
+  l = labels[:, None]
+  colors = l * palette
+  colors = (colors % 255).astype("uint8")
+  return colors
+
+def overlay_mask(image, predictions):
+  image = np.asarray(image)
+  masks = predictions.get_field("mask").numpy()
+  labels = predictions.get_field("labels").numpy()
+
+  colors = compute_colors_for_labels(labels).tolist()
+
+  for mask, color in zip(masks, colors):
+    thresh = mask[0, :, :, None]
+    contours, hierarchy = findContours(
+        thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+    )
+    image = cv2.drawContours(image, contours, -1, color, 3)
+
+  composite = image
+
+  return composite
+
+CATEGORIES = [
+    "__background", "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
+    "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant",
+    "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard",
+    "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle",
+    "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
+    "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table",
+    "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster",
+    "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush",
+]
+
+def overlay_boxes(image, predictions):
+  labels = predictions.get_field("labels").numpy()
+  boxes = predictions.bbox
+  image = np.asarray(image)
+  colors = compute_colors_for_labels(labels).tolist()
+
+  for box, color in zip(boxes, colors):
+    box = torch.tensor(box.numpy())
+    box = box.to(torch.int64)
+    top_left, bottom_right = box[:2].tolist(), box[2:].tolist()
+    image = cv2.rectangle(
+        image, tuple(top_left), tuple(bottom_right), tuple(color), 1
+    )
+
+  return image
+
+def overlay_class_names(image, predictions):
+  scores = predictions.get_field("scores").numpy().tolist()
+  labels = predictions.get_field("labels").numpy().tolist()
+  labels = [CATEGORIES[int(i)] for i in labels]
+  boxes = predictions.bbox.numpy()
+  image = np.asarray(image)
+  template = "{}: {:.2f}"
+  for box, score, label in zip(boxes, scores, labels):
+    x, y = box[:2]
+    s = template.format(label, score)
+    x, y = int(x), int(y)
+    cv2.putText(
+        image, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, .5, (255, 255, 255), 1
+    )
+
+  return image
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser(description='Run MaskRCNN', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument('--image', type=str, help="Path of the image to run")
+  parser.add_argument('--threshold', type=float, default=0.7, help="Detector threshold")
+  parser.add_argument('--size_scale', type=float, default=1.0, help="Image resize multiplier")
+  parser.add_argument('--out', type=str, default="/tmp/rendered.png", help="Output filename")
+  args = parser.parse_args()
+
+  resnet = ResNet(50, num_classes=None, stride_in_1x1=True)
+  model_tiny = MaskRCNN(resnet)
+  model_tiny.load_from_pretrained()
+  img = Image.open(args.image)
+  top_result_tiny = compute_prediction(img, model_tiny, confidence_threshold=args.threshold, size_scale=args.size_scale)
+  bbox_image = overlay_boxes(img, top_result_tiny)
+  mask_image = overlay_mask(bbox_image, top_result_tiny)
+  final_image = overlay_class_names(mask_image, top_result_tiny)
+
+  im = Image.fromarray(final_image)
+  print(f"saving {args.out}")
+  im.save(args.out)
+  im.show()
@@ -763,26 +763,48 @@ class BlendedGPTDataset:

    return dataset_idx, dataset_sample_idx

-def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False) -> BlendedGPTDataset:
-  if small:
-    if val:
-      return BlendedGPTDataset(
-        [base_dir / "c4-validation-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
-    return BlendedGPTDataset(
-      [base_dir / "c4-train.en_6_text_document"], [1.0], samples, seqlen, seed, shuffle=True)
+def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
  if val:
-    return BlendedGPTDataset(
-      [base_dir / "validation" / "c4-validationn-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
-  return BlendedGPTDataset(
-    [base_dir / "c4-train.en_6_text_document", base_dir / "c4-train.en_7_text_document"], [1.0, 1.0], samples, seqlen, seed, shuffle=True)
+    dataset = BlendedGPTDataset([
+      base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
+    ], [
+      1.0
+    ], samples, seqlen, seed, False)
+  else:
+    dataset = BlendedGPTDataset([
+      base_dir / "c4-train.en_6_text_document",
+      base_dir / "c4-train.en_7_text_document",
+    ], [
+      1.0, 1.0
+    ], samples, seqlen, seed, True)

-def iterate_llama3_dataset(dataset:BlendedGPTDataset, bs:int):
-  for b in range(math.ceil(dataset.samples / bs)):
-    batch = [dataset.get(b * bs + i) for i in range(bs)]
+  for b in range(math.ceil(samples / bs)):
+    batch = []
+    for i in range(bs):
+      tokens = dataset.get(b * bs + i)
+      batch.append(tokens)
    yield Tensor.stack(batch, dim=0)

-def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
-  return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)
+def batch_load_llama3_small(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
+  if val:
+    dataset = BlendedGPTDataset([
+      base_dir / "c4-validation-91205-samples.en_text_document",
+    ], [
+      1.0
+    ], samples, seqlen, seed, False)
+  else:
+    dataset = BlendedGPTDataset([
+      base_dir / "c4-train.en_6_text_document",
+    ], [
+      1.0
+    ], samples, seqlen, seed, True)
+
+  for b in range(math.ceil(samples / bs)):
+    batch = []
+    for i in range(bs):
+      tokens = dataset.get(b * bs + i)
+      batch.append(tokens)
+    yield Tensor.stack(batch, dim=0)

 if __name__ == "__main__":
  def load_unet3d(val):
@@ -223,13 +223,13 @@ def get_mlperf_bert_model():

 def get_fake_data_bert(BS:int):
  return {
-    "input_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
-    "input_mask": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
-    "segment_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
-    "masked_lm_positions": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
-    "masked_lm_ids": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
-    "masked_lm_weights": Tensor.zeros((BS, 76), dtype=dtypes.float32, device="CPU").contiguous(),
-    "next_sentence_labels": Tensor.zeros((BS, 1), dtype=dtypes.int32, device="CPU").contiguous(),
+    "input_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
+    "input_mask": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
+    "segment_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
+    "masked_lm_positions": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
+    "masked_lm_ids": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
+    "masked_lm_weights": Tensor.empty((BS, 76), dtype=dtypes.float32, device="CPU"),
+    "next_sentence_labels": Tensor.empty((BS, 1), dtype=dtypes.int32, device="CPU"),
  }

 def find_matches(match_quality_matrix:np.ndarray, high_threshold:float=0.5, low_threshold:float=0.4, allow_low_quality_matches:bool=False) -> np.ndarray:
@@ -59,7 +59,9 @@ class EmbeddingBert(nn.Embedding):
    arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).reshape(arange_shp)
    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
-    return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype)
+    # TODO: contiguous() here because the embedding dropout creates different asts on each device, and search becomes very slow.
+    # Should fix with fixing random ast on multi device, and fuse arange to make embedding fast.
+    return (arange == idx).mul(vals).sum(2, dtype=vals.dtype).contiguous()

 class LayerNormBert:
  def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
@@ -204,6 +204,43 @@ def eval_bert():

    st = time.perf_counter()

+def eval_mrcnn():
+  from tqdm import tqdm
+  from extra.models.mask_rcnn import MaskRCNN
+  from extra.models.resnet import ResNet
+  from extra.datasets.coco import BASEDIR, images, convert_prediction_to_coco_bbox, convert_prediction_to_coco_mask, accumulate_predictions_for_coco, evaluate_predictions_on_coco, iterate
+  from examples.mask_rcnn import compute_prediction_batched, Image
+  mdl = MaskRCNN(ResNet(50, num_classes=None, stride_in_1x1=True))
+  mdl.load_from_pretrained()
+
+  bbox_output = '/tmp/results_bbox.json'
+  mask_output = '/tmp/results_mask.json'
+
+  accumulate_predictions_for_coco([], bbox_output, rm=True)
+  accumulate_predictions_for_coco([], mask_output, rm=True)
+
+  #TODO: bs > 1 not as accurate
+  bs = 1
+
+  for batch in tqdm(iterate(images, bs=bs), total=len(images)//bs):
+    batch_imgs = []
+    for image_row in batch:
+      image_name = image_row['file_name']
+      img = Image.open(BASEDIR/f'val2017/{image_name}').convert("RGB")
+      batch_imgs.append(img)
+    batch_result = compute_prediction_batched(batch_imgs, mdl)
+    for image_row, result in zip(batch, batch_result):
+      image_name = image_row['file_name']
+      box_pred = convert_prediction_to_coco_bbox(image_name, result)
+      mask_pred = convert_prediction_to_coco_mask(image_name, result)
+      accumulate_predictions_for_coco(box_pred, bbox_output)
+      accumulate_predictions_for_coco(mask_pred, mask_output)
+    del batch_imgs
+    del batch_result
+
+  evaluate_predictions_on_coco(bbox_output, iou_type='bbox')
+  evaluate_predictions_on_coco(mask_output, iou_type='segm')
+
 def eval_llama3():
  from extra.models.llama import Transformer
  from examples.llama3 import MODEL_PARAMS, load, convert_from_huggingface
@@ -234,9 +271,12 @@ def eval_llama3():
    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
    return loss.flatten().float()

-  from examples.mlperf.dataloader import get_llama3_dataset, iterate_llama3_dataset
-  eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
-  iter = iterate_llama3_dataset(eval_dataset, BS)
+  if SMALL:
+    from examples.mlperf.dataloader import batch_load_llama3_small
+    iter = batch_load_llama3_small(BS, 5760, SEQLEN, BASEDIR, val=True)
+  else:
+    from examples.mlperf.dataloader import batch_load_llama3
+    iter = batch_load_llama3(BS, 5760, SEQLEN, BASEDIR, val=True)

  losses = []
  for tokens in tqdm(iter, total=5760//BS):
@@ -501,7 +541,7 @@ if __name__ == "__main__":
  # inference only
  Tensor.training = False

-  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
+  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(",")
  for m in models:
    nm = f"eval_{m}"
    if nm in globals():
@@ -918,6 +918,40 @@ def train_rnnt():
  # TODO: RNN-T
  pass

+@TinyJit
+def train_step_bert(model, optimizer, scheduler, loss_scaler:float, GPUS, grad_acc:int, **kwargs):
+  optimizer.zero_grad()
+
+  for i in range(grad_acc):
+    input_ids, segment_ids = kwargs[f"input_ids{i}"], kwargs[f"segment_ids{i}"]
+    # NOTE: these two have different names
+    attention_mask, masked_positions = kwargs[f"input_mask{i}"], kwargs[f"masked_lm_positions{i}"]
+    masked_lm_ids, masked_lm_weights, next_sentence_labels = kwargs[f"masked_lm_ids{i}"], kwargs[f"masked_lm_weights{i}"], kwargs[f"next_sentence_labels{i}"]
+
+    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
+      if len(GPUS) > 1: t.shard_(GPUS, axis=0)
+      else: t.to_(GPUS[0])
+
+    lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
+    loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
+    (loss * loss_scaler).backward()
+    # TODO: OOM without this realize with large grad_acc
+    Tensor.realize(*[p.grad for p in optimizer.params])
+
+  global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer[0].device)
+  for p in optimizer.params:
+    p.grad = p.grad / loss_scaler
+    global_norm += p.grad.float().square().sum()
+  global_norm = global_norm.sqrt().contiguous()
+  for p in optimizer.params:
+    p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)
+
+  optimizer.step()
+  scheduler.step()
+  # TODO: no to("CPU") here because it blocks and messes the python time
+  Tensor.realize(loss, global_norm, optimizer.optimizers[0].lr)
+  return loss, global_norm, optimizer.optimizers[0].lr
+
@TinyJit
 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):
@@ -980,8 +1014,7 @@ def train_bert():
  # ** hyperparameters **
  BS                 = config["BS"]                     = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  # TODO: implement grad accumulation + mlperf logging
-  assert grad_acc == 1
+  # TODO: mlperf logging
  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 1 * len(GPUS))
  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(GBS/96))
@@ -1040,8 +1073,8 @@ def train_bert():

  # ** Optimizer **
  parameters_no_wd = [v for k, v in get_state_dict(model).items() if "bias" in k or "LayerNorm" in k]
-  parameters_wd = [x for x in parameters if x not in set(parameters_no_wd)]
-  optimizer_wd = LAMB(parameters_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
+  parameters = [x for x in parameters if x not in set(parameters_no_wd)]
+  optimizer_wd = LAMB(parameters, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
  optimizer_no_wd = LAMB(parameters_no_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=0.0, adam=False)
  optimizer_group = OptimizerGroup(optimizer_wd, optimizer_no_wd)

@@ -1098,38 +1131,12 @@ def train_bert():
  # ** train loop **
  wc_start = time.perf_counter()

-  i, train_data = start_step, next(train_it)
+  i, train_data = start_step, [next(train_it) for _ in range(grad_acc)]

  if RUNMLPERF:
    if MLLOGGER:
      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})

-  @TinyJit
-  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]:
-      if len(GPUS) > 1: t.shard_(GPUS, axis=0)
-      else: t.to_(GPUS[0])
-    optimizer_group.zero_grad()
-
-    lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
-    loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
-    (loss * loss_scaler).backward()
-
-    global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer_group[0].device)
-    for p in optimizer_group.params:
-      p.grad = p.grad / loss_scaler
-      global_norm += p.grad.float().square().sum()
-    global_norm = global_norm.sqrt().contiguous()
-    for p in optimizer_group.params:
-      p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)
-
-    optimizer_group.step()
-    scheduler_group.step()
-    # TODO: no to("CPU") here because it blocks and messes the python time
-    Tensor.realize(loss, global_norm, optimizer_group.optimizers[0].lr)
-    return loss, global_norm, optimizer_group.optimizers[0].lr
-
  while train_data is not None and i < train_steps and not achieved:
    if getenv("TRAIN", 1):
      Tensor.training = True
@@ -1137,12 +1144,16 @@ def train_bert():
      st = time.perf_counter()
      GlobalCounters.reset()
      with WallTimeEvent(BenchEvent.STEP):
-        loss, global_norm, lr = train_step_bert(
-          train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
-          train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"])
+        data = {f"{k}{i}":v for i,d in enumerate(train_data) for k,v in d.items()}
+        loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler, GPUS, grad_acc, **data)

        pt = time.perf_counter()
-        next_data = next(train_it)
+
+        try:
+          next_data = [next(train_it) for _ in range(grad_acc)]
+        except StopIteration:
+          next_data = None
+
        dt = time.perf_counter()

        device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
@@ -1177,8 +1188,8 @@ def train_bert():
      if MLLOGGER and RUNMLPERF:
        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
      if getenv("RESET_STEP"): train_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE") and train_step_bert.captured is not None:
-        # TODO: this hangs on tiny green after 90 minutes of training
+      elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None:
+        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
        train_step_bert.captured.free_intermediates()
      eval_lm_losses = []
      eval_clsf_losses = []
@@ -1213,7 +1224,7 @@ def train_bert():
          return

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

      del eval_data
      avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
@@ -1289,7 +1300,6 @@ def train_llama3():
  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
  BS                 = config["BS"]                     = getenv("BS", 16)
  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  assert grad_acc == 1, f"{grad_acc=} is not supported"
  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
  SEED               = config["SEED"]                   = getenv("SEED", 5760)
  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1314,21 +1324,12 @@ def train_llama3():
  opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
  opt_end_learning_rate = getenv("END_LR", 8e-7)

-  # ** init wandb **
-  WANDB = getenv("WANDB")
-  if WANDB:
-    import wandb
-    wandb_args = {"id": wandb_id, "resume": "must"} if (wandb_id := getenv("WANDB_RESUME", "")) else {}
-    wandb.init(config=config, **wandb_args, project="MLPerf-LLaMA3")
-
-  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
+  # TODO: confirm weights are in bf16
  # vocab_size from the mixtral tokenizer
-  if not SMALL: model_params |= {"vocab_size": 32000}
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
-  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
-  params = get_parameters(model)
-  # weights are all bfloat16 for now
-  assert params and all(p.dtype == dtypes.bfloat16 for p in params)
+  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
+  params = params | {"vocab_size": 32000} if not SMALL else params
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
+  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)

  if getenv("FAKEDATA"):
    for v in get_parameters(model):
@@ -1373,17 +1374,20 @@ def train_llama3():

  @TinyJit
  @Tensor.train()
-  def train_step(model, tokens:Tensor):
+  def train_step(model, tokens:Tensor, grad_acc:int):
    optim.zero_grad()
-    if (DP := getenv("DP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
-    if (MP := getenv("MP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
-      tokens = tokens.shard(device)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    loss.backward()
+    # grad acc
+    for batch in tokens.split(tokens.shape[0]//grad_acc):
+      if (DP := getenv("DP", 1)) > 1:
+        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+        batch = batch.shard(device, 0)
+      if (MP := getenv("MP", 1)) > 1:
+        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+        batch = batch.shard(device)
+      logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
+      loss = logits.sparse_categorical_crossentropy(batch[:, 1:])
+      loss.backward()
+      Tensor.realize(*[p.grad for p in optim.params])
    # L2 norm grad clip
    # https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
    # https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
@@ -1418,62 +1422,55 @@ def train_llama3():
  # ** data iters **
  def fake_data(bs, samples):
    for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)

  def get_train_iter():
    if getenv("FAKEDATA", 0):
-      return fake_data(BS, SAMPLES)
+      return fake_data(GBS, SAMPLES)
    else:
-      from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
-
-  if getenv("FAKEDATA", 0):
-    eval_dataset = None
-  else:
-    from examples.mlperf.dataloader import get_llama3_dataset
-    eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
+      if SMALL:
+        from examples.mlperf.dataloader import batch_load_llama3_small
+        return batch_load_llama3_small(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+      else:
+        from examples.mlperf.dataloader import batch_load_llama3
+        return batch_load_llama3(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))

  def get_eval_iter():
-    if eval_dataset is None:
+    if getenv("FAKEDATA", 0):
      return fake_data(EVAL_BS, 5760)
-    from examples.mlperf.dataloader import iterate_llama3_dataset
-    return iterate_llama3_dataset(eval_dataset, EVAL_BS)
+    else:
+      if SMALL:
+        from examples.mlperf.dataloader import batch_load_llama3_small
+        return batch_load_llama3_small(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
+      else:
+        from examples.mlperf.dataloader import batch_load_llama3
+        return batch_load_llama3(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)

  iter = get_train_iter()
  i, sequences_seen = resume_ckpt, 0
  for tokens in tqdm(iter, total=SAMPLES//GBS):
+    t = time.perf_counter()
    GlobalCounters.reset()
-    if getenv("TRAIN", 1):
-      t = time.perf_counter()
-      loss, lr = train_step(model, tokens)
-      loss = loss.float().item()
-      lr = lr.item()
+    loss, lr = train_step(model, tokens, grad_acc)
+    loss = loss.float().item()

-      i += 1
-      sequences_seen += tokens.shape[0]
+    i += 1
+    sequences_seen += tokens.shape[0]

-      sec = time.perf_counter()-t
-      mem_gb = GlobalCounters.mem_used / 1e9
-      gflops = GlobalCounters.global_ops / 1e9 / sec
-      tqdm.write(
-        f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS")
+    tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
+    if (fname:=getenv("LOSS_FILE", "")):
+      with open(fname, "a") as f:
+        f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")

-      if (fname:=getenv("LOSS_FILE", "")):
-        with open(fname, "a") as f:
-          f.write(f"{i} {loss:.4f} {lr:.12f} {mem_gb:.2f}\n")
+    if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
+      tqdm.write("saving checkpoint")
+      if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
+      fn = f"{ckpt_dir}/llama3_{i}.safe"
+      safe_save(get_state_dict(model), fn)

-      if WANDB:
-        wandb.log({"lr": lr, "train/loss": loss, "train/step_time": sec, "train/GFLOPS": gflops, "train/sequences_seen": sequences_seen})
-
-      if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
-        tqdm.write("saving checkpoint")
-        if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-        fn = f"{ckpt_dir}/llama3_{i}.safe"
-        safe_save(get_state_dict(model), fn)
-
-        tqdm.write("saving optim checkpoint")
-        fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
-        safe_save(get_state_dict(scheduler), fn)
+      tqdm.write("saving optim checkpoint")
+      fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
+      safe_save(get_state_dict(scheduler), fn)

    if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
      tqdm.write(f"evaluating after {sequences_seen} sequences")
@@ -1489,9 +1486,6 @@ def train_llama3():

      tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")

-      if WANDB:
-        wandb.log({"eval/log_perplexity": log_perplexity, "eval/sequences_seen": sequences_seen})
-
      if log_perplexity < EVAL_TARGET:
        tqdm.write(f"target achieved after {sequences_seen} sequences")
        if getenv("CKPT"):
@@ -1,31 +0,0 @@
-#!/bin/bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export SUBMISSION_PLATFORM="tinybox_8xMI350X"
-export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
-
-# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
-export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
-export TRAIN_STEPS=3900
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=5000000
-
-export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
-export BASEDIR="/raid/datasets/wiki"
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=1
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="bert_8xMI350x_${DATETIME}_${SEED}.log"
-
-BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
-
-# run
-PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -2,7 +2,7 @@

 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -2,7 +2,7 @@

 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -5,7 +5,7 @@ set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." NV=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -2,7 +2,7 @@

 export PYTHONPATH="." AMD=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -2,7 +2,7 @@

 export PYTHONPATH="." AMD=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -5,7 +5,7 @@ set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -15,7 +15,7 @@ export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

 # pip install -e ".[mlperf]"
-export LOGMLPERF=${LOGMLPERF:-1}
+export LOGMLPERF=1

 export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
@@ -0,0 +1,118 @@
+import json, pprint
+from tinygrad import fetch, nn, Tensor
+from tinygrad.helpers import DEBUG
+
+class FeedForward:
+  def __init__(self, model_dim, intermediate_dim):
+    self.proj_1 = nn.Linear(model_dim, 2*intermediate_dim, bias=False)
+    self.proj_2 = nn.Linear(intermediate_dim, model_dim, bias=False)
+
+  def __call__(self, x):
+    y_12 = self.proj_1(x)
+    y_1, y_2 = y_12.chunk(2, dim=-1)
+    return self.proj_2(y_1.silu() * y_2)
+
+# NOTE: this RoPE doesn't match LLaMA's?
+def _rotate_half(x: Tensor) -> Tensor:
+  x1, x2 = x.chunk(2, dim=-1)
+  return Tensor.cat(-x2, x1, dim=-1)
+
+def _apply_rotary_pos_emb(x: Tensor, pos_sin: Tensor, pos_cos: Tensor) -> Tensor:
+  return (x * pos_cos) + (_rotate_half(x) * pos_sin)
+
+class Attention:
+  def __init__(self, model_dim, num_query_heads, num_kv_heads, head_dim):
+    self.qkv_proj = nn.Linear(model_dim, (num_query_heads + num_kv_heads*2) * head_dim, bias=False)
+    self.num_query_heads, self.num_kv_heads = num_query_heads, num_kv_heads
+    self.head_dim = head_dim
+    self.q_norm = nn.RMSNorm(head_dim)
+    self.k_norm = nn.RMSNorm(head_dim)
+    self.out_proj = nn.Linear(num_query_heads * head_dim, model_dim, bias=False)
+
+  def __call__(self, x:Tensor) -> Tensor:
+    batch_size, seq_len, embed_dim = x.shape
+    qkv = self.qkv_proj(x)
+    qkv = qkv.reshape(batch_size, seq_len, self.num_query_heads+self.num_kv_heads*2, self.head_dim).transpose(1, 2)
+    xq,xk,xv = qkv.split([self.num_query_heads, self.num_kv_heads, self.num_kv_heads], dim=1)
+    xq = self.q_norm(xq)
+    xk = self.k_norm(xk)
+
+    # add positional embedding (how many kernels is this?)
+    freq_constant = 10000
+    inv_freq = 1.0 / (freq_constant ** (Tensor.arange(0, self.head_dim, 2) / self.head_dim))
+    pos_index_theta = Tensor.einsum("i,j->ij", Tensor.arange(seq_len), inv_freq)
+    emb = Tensor.cat(pos_index_theta, pos_index_theta, dim=-1)
+    cos_emb, sin_emb = emb.cos()[None, None, :, :], emb.sin()[None, None, :, :]
+    xq = _apply_rotary_pos_emb(xq, sin_emb, cos_emb)
+    xk = _apply_rotary_pos_emb(xk, sin_emb, cos_emb)
+
+    # grouped-query attention
+    num_groups = self.num_query_heads // self.num_kv_heads
+    xk = xk.repeat_interleave(num_groups, dim=1)
+    xv = xv.repeat_interleave(num_groups, dim=1)
+
+    # masked attention
+    #start_pos = 0
+    #mask = Tensor.full((1, 1, seq_len, start_pos+seq_len), float("-inf"), dtype=xq.dtype, device=xq.device).triu(start_pos+1)
+    #attn_output = xq.scaled_dot_product_attention(xk, xv, mask).transpose(1, 2)
+
+    # causal is fine, no mask needed
+    attn_output = xq.scaled_dot_product_attention(xk, xv, is_causal=True).transpose(1, 2)
+    return self.out_proj(attn_output.reshape(batch_size, seq_len, self.num_query_heads * self.head_dim))
+
+class Layer:
+  def __init__(self, model_dim, intermediate_dim, num_query_heads, num_kv_heads, head_dim):
+    self.ffn = FeedForward(model_dim, intermediate_dim)
+    self.attn = Attention(model_dim, num_query_heads, num_kv_heads, head_dim)
+    self.ffn_norm = nn.RMSNorm(model_dim)
+    self.attn_norm = nn.RMSNorm(model_dim)
+
+  def __call__(self, x:Tensor) -> Tensor: # (batch, seq_len, embed_dim)
+    x = x + self.attn(self.attn_norm(x))
+    x = x + self.ffn(self.ffn_norm(x))
+    return x
+
+# stupidly complex
+def make_divisible(v, divisor):
+  new_v = max(divisor, int(v + divisor / 2) // divisor * divisor)
+  if new_v < 0.9 * v: new_v += divisor
+  return new_v
+
+class Transformer:
+  def __init__(self, cfg):
+    if DEBUG >= 3: pprint.pp(cfg)
+    self.layers = [Layer(cfg['model_dim'], make_divisible(int(cfg["model_dim"] * cfg['ffn_multipliers'][i]), cfg['ffn_dim_divisor']),
+                         cfg['num_query_heads'][i], cfg['num_kv_heads'][i], cfg['head_dim']) for i in range(cfg['num_transformer_layers'])]
+    self.norm = nn.RMSNorm(cfg['model_dim'])
+    self.token_embeddings = nn.Embedding(cfg['vocab_size'], cfg['model_dim'])
+
+  def __call__(self, tokens:Tensor):
+    # _bsz, seqlen = tokens.shape
+    x = self.token_embeddings(tokens)
+    for l in self.layers: x = l(x)
+    return self.norm(x) @ self.token_embeddings.weight.T
+
+if __name__ == "__main__":
+  #model_name = "OpenELM-270M-Instruct"
+  model_name = "OpenELM-270M"  # this is fp32
+  model = Transformer(json.loads(fetch(f"https://huggingface.co/apple/{model_name}/resolve/main/config.json?download=true").read_bytes()))
+  weights = nn.state.safe_load(fetch(f"https://huggingface.co/apple/{model_name}/resolve/main/model.safetensors?download=true"))
+  if DEBUG >= 3:
+    for k, v in weights.items(): print(k, v.shape)
+  nn.state.load_state_dict(model, {k.removeprefix("transformer."):v for k,v in weights.items()})
+
+  from sentencepiece import SentencePieceProcessor
+  tokenizer = SentencePieceProcessor(fetch("https://github.com/karpathy/llama2.c/raw/master/tokenizer.model").as_posix())
+  toks = [tokenizer.bos_id()] + tokenizer.encode("Some car brands include")
+  for i in range(100):
+    ttoks = Tensor([toks])
+    out = model(ttoks).realize()
+    t0 = out[0].argmax(axis=-1).tolist()
+    toks.append(t0[-1])
+    # hmmm...passthrough still doesn't match (it shouldn't, it outputs the most likely)
+    print(tokenizer.decode(toks))
+    #print(toks)
+    #print(tokenizer.decode(t0))
+    #print(t0)
+
+
@@ -1,10 +1,15 @@
 import os, sys, pickle, time, re
 import numpy as np
+if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
+if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
+if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.engine.realize import CompiledRunner
+
+import onnx
 from tinygrad.nn.onnx import OnnxRunner

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
@@ -16,14 +21,11 @@ def compile(onnx_file):

  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
-
  # Float inputs and outputs to tinyjits for openpilot are always float32
-  # TODO this seems dumb
  input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
  Tensor.manual_seed(100)
-  inputs = {k:Tensor(Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize().numpy(), device='NPY') for k,shp in sorted(input_shapes.items())}
-  if not getenv("NPY_IMG"):
-    inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
+  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in sorted(input_shapes.items())}
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
  print("created tensors")

  run_onnx_jit = TinyJit(lambda **kwargs:
@@ -31,6 +33,8 @@ def compile(onnx_file):
  for i in range(3):
    GlobalCounters.reset()
    print(f"run {i}")
+    inputs = {**{k:v.clone() for k,v in new_inputs.items() if 'img' in k},
+              **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
      ret = run_onnx_jit(**inputs).numpy()
    # copy i == 1 so use of JITBEAM is okay
@@ -39,7 +43,7 @@ def compile(onnx_file):
  np.testing.assert_equal(test_val, ret, "JIT run failed")
  print("jit run validated")

-  # check gated read_image usage
+  # checks from compile2
  kernel_count = 0
  read_image_count = 0
  gated_read_image_count = 0
@@ -65,9 +69,14 @@ def compile(onnx_file):
  print(f"mdl size is {mdl_sz/1e6:.2f}M")
  print(f"pkl size is {pkl_sz/1e6:.2f}M")
  print("**** compile done ****")
-  return inputs, test_val
+  return test_val

-def test_vs_compile(run, inputs, test_val=None):
+def test_vs_compile(run, new_inputs, test_val=None):
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
+
+  # create fake "from_blob" tensors for the inputs, and wrapped NPY tensors for the numpy inputs (these have the same underlying memory)
+  inputs = {**{k:v for k,v in new_inputs.items() if 'img' in k},
+            **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}

  # run 20 times
  step_times = []
@@ -84,58 +93,68 @@ def test_vs_compile(run, inputs, test_val=None):
    min_time = min(step_times)
    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"

+  print(out, val.shape, val.dtype)
  if test_val is not None: np.testing.assert_equal(test_val, val)
  print("**** test done ****")

  # test that changing the numpy changes the model outputs
-  inputs_2x = {k: Tensor(v.numpy()*2, device=v.device) for k,v in inputs.items()}
-  out = run(**inputs_2x)
-  changed_val = out.numpy()
-  np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
+  if any([x.device == 'NPY' for x in inputs.values()]):
+    for v in new_inputs_numpy.values(): v *= 2
+    out = run(**inputs)
+    changed_val = out.numpy()
+    np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
  return val

-def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
-  import onnx
-  import onnxruntime as ort
-  
-  onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}
+def test_vs_onnx(new_inputs, test_val, onnx_file, ort=False):
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
  onnx_model = onnx.load(onnx_file)

-  ORT_TO_NP_DTYPES: dict[str, np.dtype] = {
-    'tensor(float)': np.dtype('float32'),
-    'tensor(float16)': np.dtype('float16'),
-    'tensor(uint8)': np.dtype('uint8'),
-  }
-
  timings = []
-  onnx_session = ort.InferenceSession(onnx_file)
-  onnx_types = {x.name: ORT_TO_NP_DTYPES[x.type] for x in onnx_session.get_inputs()}
-  onnx_inputs = {k:onnx_inputs[k].astype(onnx_types[k]) for k in onnx_inputs}
+  if ort:
+    # test with onnxruntime
+    import onnxruntime as ort
+    onnx_session = ort.InferenceSession(onnx_file)
+    for _ in range(1 if test_val is not None else 5):
+      st = time.perf_counter()
+      onnx_output = onnx_session.run([onnx_model.graph.output[0].name], {k:v.astype(np.float16) for k,v in new_inputs_numpy.items()})
+      timings.append(time.perf_counter() - st)
+    new_torch_out = onnx_output[0]
+  else:
+    # test with torch
+    import torch
+    from onnx2torch import convert
+    inputs = {k.name:new_inputs_numpy[k.name] for k in onnx_model.graph.input}
+    torch_model = convert(onnx_model).float()
+    with torch.no_grad():
+      for _ in range(1 if test_val is not None else 5):
+        st = time.perf_counter()
+        torch_out = torch_model(*[torch.tensor(x) for x in inputs.values()])
+        timings.append(time.perf_counter() - st)
+      new_torch_out = torch_out.numpy()

-  for _ in range(1 if test_val is not None else 5):
-    st = time.perf_counter()
-    onnx_output = onnx_session.run([onnx_model.graph.output[0].name], onnx_inputs)
-    timings.append(time.perf_counter() - st)
-
-  np.testing.assert_allclose(onnx_output[0].reshape(test_val.shape), test_val, atol=tol, rtol=tol)
-  print("test vs onnx passed")
+  if test_val is not None:
+    np.testing.assert_allclose(new_torch_out.reshape(test_val.shape), test_val, atol=1e-4, rtol=1e-2)
+    print("test vs onnx passed")
  return timings

-def bench(run, inputs):
-  from extra.bench_log import WallTimeEvent, BenchEvent
-  for _ in range(10):
-    with WallTimeEvent(BenchEvent.STEP):
-      run(**inputs).numpy()
-
 if __name__ == "__main__":
  onnx_file = fetch(OPENPILOT_MODEL)
-  inputs, outputs = compile(onnx_file)
+  test_val = compile(onnx_file) if not getenv("RUN") else None

  with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)

-  test_vs_compile(pickle_loaded, inputs, outputs)
-  if getenv("SELFTEST"):
-    test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
+  # same randomness as compile
+  Tensor.manual_seed(100)
+  new_inputs = {nm:Tensor.randn(*st.shape, dtype=dtype).mul(8).realize() for nm, (st, _, dtype, _) in
+                sorted(zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_st_vars_dtype_device))}
+
+  test_val = test_vs_compile(pickle_loaded, new_inputs, test_val)
+  if getenv("BENCHMARK"):
+    for be in ["torch", "ort"]:
+      try:
+        timings = test_vs_onnx(new_inputs, None, onnx_file, be=="ort")
+        print(f"timing {be}: {min(timings)*1000:.2f} ms")
+      except Exception as e:
+        print(f"{be} fail with {e}")
+  if not getenv("FLOAT16"): test_vs_onnx(new_inputs, test_val, onnx_file, getenv("ORT"))

-  if getenv("BENCHMARK_LOG", ""):
-    bench(pickle_loaded, inputs)
@@ -0,0 +1,47 @@
+import sys
+from tinygrad import Tensor, fetch, GlobalCounters, dtypes
+from tinygrad.uop.ops import UOp
+from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.schedule.rangeify import get_rangeify_map
+from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.engine.realize import run_schedule
+
+# NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
+
+OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
+OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
+
+if __name__ == "__main__":
+  onnx_file = fetch(OPENPILOT_MODEL)
+  run_onnx = OnnxRunner(onnx_file)
+
+  inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
+  out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
+  root = out.uop
+  targets = [x.uop for x in inputs.values()]
+  print(targets)
+
+  # TODO: abstract this from gradient?
+
+  # compute the target path (top down)
+  in_target_path: dict[UOp, bool] = {}
+  for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
+  independent_set = {}
+  for u in root.toposort():
+    if in_target_path[u]:
+      for s in u.src:
+        if not in_target_path[s]:
+          independent_set[s] = None
+  independent = UOp.sink(*independent_set.keys())
+  kernelized = get_rangeify_map(independent)
+  independent = independent.substitute(kernelized)
+  schedule, var_vals = create_schedule_with_vars(independent)
+  run_schedule(schedule)
+
+  print("**** real ****")
+  GlobalCounters.reset()
+  out.uop = root.substitute(kernelized)
+  out.kernelize()
+
+  # realize
+  out.realize()
@@ -0,0 +1,55 @@
+from tinygrad.helpers import trange
+from tinygrad.nn.datasets import mnist
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.optimizers as optim
+from functools import partial
+
+class Model(nn.Module):
+  def __init__(self):
+    super().__init__()
+    self.c1 = nn.Conv2d(1, 32, 5)
+    self.c2 = nn.Conv2d(32, 32, 5)
+    self.bn1 = nn.BatchNorm(32)
+    self.m1 = nn.MaxPool2d(2)
+    self.c3 = nn.Conv2d(32, 64, 3)
+    self.c4 = nn.Conv2d(64, 64, 3)
+    self.bn2 = nn.BatchNorm(64)
+    self.m2 = nn.MaxPool2d(2)
+    self.lin = nn.Linear(576, 10)
+  def __call__(self, x):
+    x = mx.maximum(self.c1(x), 0)
+    x = mx.maximum(self.c2(x), 0)
+    x = self.m1(self.bn1(x))
+    x = mx.maximum(self.c3(x), 0)
+    x = mx.maximum(self.c4(x), 0)
+    x = self.m2(self.bn2(x))
+    return self.lin(mx.flatten(x, 1))
+
+if __name__ == "__main__":
+  X_train, Y_train, X_test, Y_test = mnist()
+  X_train = mx.array(X_train.float().permute((0,2,3,1)).numpy())
+  Y_train = mx.array(Y_train.numpy())
+  X_test = mx.array(X_test.float().permute((0,2,3,1)).numpy())
+  Y_test = mx.array(Y_test.numpy())
+
+  model = Model()
+  optimizer = optim.Adam(1e-3)
+  def loss_fn(model, x, y): return nn.losses.cross_entropy(model(x), y).mean()
+
+  state = [model.state, optimizer.state]
+  @partial(mx.compile, inputs=state, outputs=state)
+  def step(samples):
+    # Compiled functions will also treat any inputs not in the parameter list as constants.
+    X,Y = X_train[samples], Y_train[samples]
+    loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
+    loss, grads = loss_and_grad_fn(model, X, Y)
+    optimizer.update(model, grads)
+    return loss
+
+  test_acc = float('nan')
+  for i in (t:=trange(70)):
+    samples = mx.random.randint(0, X_train.shape[0], (512,))  # putting this in JIT didn't work well
+    loss = step(samples)
+    if i%10 == 9: test_acc = ((model(X_test).argmax(axis=-1) == Y_test).sum() * 100 / X_test.shape[0]).item()
+    t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
@@ -0,0 +1,45 @@
+import gymnasium as gym
+import numpy as np
+from gymnasium.envs.registration import register
+
+# a very simple game
+# one of <size> lights will light up
+# take the action of the lit up light
+# in <hard_mode>, you act differently based on the step number and need to track this
+
+class PressTheLightUpButton(gym.Env):
+  metadata = {"render_modes": []}
+  def __init__(self, render_mode=None, size=2, game_length=10, hard_mode=False):
+    self.size, self.game_length = size, game_length
+    self.observation_space = gym.spaces.Box(0, 1, shape=(self.size,), dtype=np.float32)
+    self.action_space = gym.spaces.Discrete(self.size)
+    self.step_num = 0
+    self.done = True
+    self.hard_mode = hard_mode
+
+  def _get_obs(self):
+    obs = [0]*self.size
+    if self.step_num < len(self.state):
+      obs[self.state[self.step_num]] = 1
+    return np.array(obs, dtype=np.float32)
+
+  def reset(self, seed=None, options=None):
+    super().reset(seed=seed)
+    self.state = np.random.randint(0, self.size, size=self.game_length)
+    self.step_num = 0
+    self.done = False
+    return self._get_obs(), {}
+
+  def step(self, action):
+    target = ((action + self.step_num) % self.size) if self.hard_mode else action
+    reward = int(target == self.state[self.step_num])
+    self.step_num += 1
+    if not reward:
+      self.done = True
+    return self._get_obs(), reward, self.done, self.step_num >= self.game_length, {}
+
+register(
+  id="PressTheLightUpButton-v0",
+  entry_point="examples.rl.lightupbutton:PressTheLightUpButton",
+  max_episode_steps=None,
+)
@@ -99,7 +99,6 @@ if __name__ == "__main__":
  parser.add_argument('--timing',      action='store_true', help="Print timing per step")
  parser.add_argument('--noshow',      action='store_true', help="Don't show the image")
  parser.add_argument('--fp16',        action='store_true', help="Cast the weights to float16")
-  parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
  args = parser.parse_args()

  N = 1
@@ -113,22 +112,19 @@ if __name__ == "__main__":

  model = StableDiffusionV2(**params)

-  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    if not args.fakeweights:
-      default_weights_url = 'https://huggingface.co/sd2-community/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
-      weights_fn = args.weights_fn
-      if not weights_fn:
-        weights_url = args.weights_url if args.weights_url else default_weights_url
-        weights_fn  = fetch(weights_url, os.path.basename(str(weights_url)))
+  default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
+  weights_fn = args.weights_fn
+  if not weights_fn:
+    weights_url = args.weights_url if args.weights_url else default_weights_url
+    weights_fn  = fetch(weights_url, os.path.basename(str(weights_url)))

-        load_state_dict(model, safe_load(weights_fn), strict=False)
+  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
+    load_state_dict(model, safe_load(weights_fn), strict=False)

    if args.fp16:
      for k,v in get_state_dict(model).items():
        if k.startswith("model"):
-          v.replace(v.cast(dtypes.float16))
-
-    Tensor.realize(*get_state_dict(model).values())
+          v.replace(v.cast(dtypes.float16).realize())

  c  = { "crossattn": model.cond_stage_model(args.prompt) }
  uc = { "crossattn": model.cond_stage_model("") }
@@ -0,0 +1,48 @@
+import os, pathlib, argparse
+from examples.llama3 import Tokenizer
+from tabulate import tabulate
+from tinygrad import fetch
+from tinygrad.helpers import flatten
+
+# llama 3 tokenizer
+tokenizer = Tokenizer(fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model").as_posix())
+
+def read_code(base_path):
+  ret = []
+  for path, _, files in os.walk(os.path.join(base_path, "tinygrad")):
+    for name in files:
+      if not name.endswith(".py"): continue
+      if 'tinygrad/runtime/autogen' in path.replace('\\', '/'): continue
+      fullpath = os.path.join(path, name)
+      code = pathlib.Path(fullpath).read_text()
+      ret.append(("### " + fullpath.split("tinygrad/", 1)[1], code))
+  return ret
+
+def write_code_to_file(filename, code_list):
+  """Writes the combined code to a specified file."""
+  with open(filename, 'w') as f:
+    f.write('\n'.join(flatten(code_list)))
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(description="Analyze and optionally save tinygrad code.")
+  parser.add_argument("--output", help="Output file to write the combined code to.")
+  args = parser.parse_args()
+
+  ret = read_code(".")
+
+  table = []
+  for name,code in ret:
+    table.append([name, len(tokenizer.encode(name+"\x00"+code))])
+  print(tabulate([["name", "llm tokens"]]+sorted(table, key=lambda x: -x[1]), headers="firstrow"))
+
+  code_str = '\x00'.join(flatten(ret))
+  print(f"code has {len(code_str)} chars")
+  newline_count = code_str.count('\n')
+  print(f"code has {newline_count} newlines")
+
+  encoded = tokenizer.encode(code_str)
+  print(f"code has {len(encoded)} tokens")
+
+  if args.output:
+    write_code_to_file(args.output, ret)
+    print(f"Combined code written to {args.output}")
@@ -0,0 +1,136 @@
+#!/usr/bin/env python
+#inspired by https://github.com/Matuzas77/MNIST-0.17/blob/master/MNIST_final_solution.ipynb
+import sys
+import numpy as np
+from tinygrad.nn.state import get_parameters
+from tinygrad.tensor import Tensor
+from tinygrad.nn import BatchNorm2d, optim
+from tinygrad.helpers import getenv
+from extra.datasets import fetch_mnist
+from extra.augment import augment_img
+from extra.training import train, evaluate
+GPU = getenv("GPU")
+QUICK = getenv("QUICK")
+DEBUG = getenv("DEBUG")
+
+class SqueezeExciteBlock2D:
+  def __init__(self, filters):
+    self.filters = filters
+    self.weight1 = Tensor.scaled_uniform(self.filters, self.filters//32)
+    self.bias1 = Tensor.scaled_uniform(1,self.filters//32)
+    self.weight2 = Tensor.scaled_uniform(self.filters//32, self.filters)
+    self.bias2 = Tensor.scaled_uniform(1, self.filters)
+
+  def __call__(self, input):
+    se = input.avg_pool2d(kernel_size=(input.shape[2], input.shape[3])) #GlobalAveragePool2D
+    se = se.reshape(shape=(-1, self.filters))
+    se = se.dot(self.weight1) + self.bias1
+    se = se.relu()
+    se = se.dot(self.weight2) + self.bias2
+    se = se.sigmoid().reshape(shape=(-1,self.filters,1,1)) #for broadcasting
+    se = input.mul(se)
+    return se
+
+class ConvBlock:
+  def __init__(self, h, w, inp, filters=128, conv=3):
+    self.h, self.w = h, w
+    self.inp = inp
+    #init weights
+    self.cweights = [Tensor.scaled_uniform(filters, inp if i==0 else filters, conv, conv) for i in range(3)]
+    self.cbiases = [Tensor.scaled_uniform(1, filters, 1, 1) for i in range(3)]
+    #init layers
+    self._bn = BatchNorm2d(128)
+    self._seb = SqueezeExciteBlock2D(filters)
+
+  def __call__(self, input):
+    x = input.reshape(shape=(-1, self.inp, self.w, self.h))
+    for cweight, cbias in zip(self.cweights, self.cbiases):
+      x = x.pad(padding=[1,1,1,1]).conv2d(cweight).add(cbias).relu()
+    x = self._bn(x)
+    x = self._seb(x)
+    return x
+
+class BigConvNet:
+  def __init__(self):
+    self.conv = [ConvBlock(28,28,1), ConvBlock(28,28,128), ConvBlock(14,14,128)]
+    self.weight1 = Tensor.scaled_uniform(128,10)
+    self.weight2 = Tensor.scaled_uniform(128,10)
+
+  def parameters(self):
+    if DEBUG: #keeping this for a moment
+      pars = [par for par in get_parameters(self) if par.requires_grad]
+      no_pars = 0
+      for par in pars:
+        print(par.shape)
+        no_pars += np.prod(par.shape)
+      print('no of parameters', no_pars)
+      return pars
+    else:
+      return get_parameters(self)
+
+  def save(self, filename):
+    with open(filename+'.npy', 'wb') as f:
+      for par in get_parameters(self):
+        #if par.requires_grad:
+        np.save(f, par.numpy())
+
+  def load(self, filename):
+    with open(filename+'.npy', 'rb') as f:
+      for par in get_parameters(self):
+        #if par.requires_grad:
+        try:
+          par.numpy()[:] = np.load(f)
+          if GPU:
+            par.gpu()
+        except:
+          print('Could not load parameter')
+
+  def forward(self, x):
+    x = self.conv[0](x)
+    x = self.conv[1](x)
+    x = x.avg_pool2d(kernel_size=(2,2))
+    x = self.conv[2](x)
+    x1 = x.avg_pool2d(kernel_size=(14,14)).reshape(shape=(-1,128)) #global
+    x2 = x.max_pool2d(kernel_size=(14,14)).reshape(shape=(-1,128)) #global
+    xo = x1.dot(self.weight1) + x2.dot(self.weight2)
+    return xo
+
+
+if __name__ == "__main__":
+  lrs = [1e-4, 1e-5] if QUICK else [1e-3, 1e-4, 1e-5, 1e-5]
+  epochss = [2, 1] if QUICK else [13, 3, 3, 1]
+  BS = 32
+
+  lmbd = 0.00025
+  lossfn = lambda out,y: out.sparse_categorical_crossentropy(y) + lmbd*(model.weight1.abs() + model.weight2.abs()).sum()
+  X_train, Y_train, X_test, Y_test = fetch_mnist()
+  X_train = X_train.reshape(-1, 28, 28).astype(np.uint8)
+  X_test = X_test.reshape(-1, 28, 28).astype(np.uint8)
+  steps = len(X_train)//BS
+  np.random.seed(1337)
+  if QUICK:
+    steps = 1
+    X_test, Y_test = X_test[:BS], Y_test[:BS]
+
+  model = BigConvNet()
+
+  if len(sys.argv) > 1:
+    try:
+      model.load(sys.argv[1])
+      print('Loaded weights "'+sys.argv[1]+'", evaluating...')
+      evaluate(model, X_test, Y_test, BS=BS)
+    except:
+      print('could not load weights "'+sys.argv[1]+'".')
+
+  if GPU:
+    params = get_parameters(model)
+    [x.gpu_() for x in params]
+
+  for lr, epochs in zip(lrs, epochss):
+    optimizer = optim.Adam(model.parameters(), lr=lr)
+    for epoch in range(1,epochs+1):
+      #first epoch without augmentation
+      X_aug = X_train if epoch == 1 else augment_img(X_train)
+      train(model, X_aug, Y_train, optimizer, steps=steps, lossfn=lossfn, BS=BS)
+      accuracy = evaluate(model, X_test, Y_test, BS=BS)
+      model.save(f'examples/checkpoint{accuracy * 1e6:.0f}')
@@ -0,0 +1,17 @@
+from tinygrad.tensor import Tensor
+from tinygrad.nn import Conv2d, BatchNorm2d
+from tinygrad.nn.state import get_parameters
+
+if __name__ == "__main__":
+  with Tensor.train():
+
+    BS, C1, H, W = 4, 16, 224, 224
+    C2, K, S, P = 64, 7, 2, 1
+
+    x = Tensor.uniform(BS, C1, H, W)
+    conv = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
+    bn = BatchNorm2d(C2, track_running_stats=False)
+    for t in get_parameters([x, conv, bn]): t.realize()
+
+    print("running network")
+    x.sequential([conv, bn]).numpy()
@@ -0,0 +1,669 @@
+# original implementation: https://github.com/svc-develop-team/so-vits-svc
+from __future__ import annotations
+import sys, logging, time, io, math, argparse, operator, numpy as np
+from functools import partial, reduce
+from pathlib import Path
+from typing import Tuple, Optional, Type
+from tinygrad import nn, dtypes, Tensor
+from tinygrad.helpers import getenv, fetch
+from tinygrad.nn.state import torch_load
+from examples.vits import ResidualCouplingBlock, PosteriorEncoder, Encoder, ResBlock1, ResBlock2, LRELU_SLOPE, sequence_mask, split, get_hparams_from_file, load_checkpoint, weight_norm, HParams
+from examples.sovits_helpers import preprocess
+import soundfile
+
+DEBUG = getenv("DEBUG")
+
+F0_BIN = 256
+F0_MAX = 1100.0
+F0_MIN = 50.0
+F0_MEL_MIN = 1127 * np.log(1 + F0_MIN / 700)
+F0_MEL_MAX = 1127 * np.log(1 + F0_MAX / 700)
+
+class SpeechEncoder:
+  def __init__(self, hidden_dim, model:ContentVec): self.hidden_dim, self.model = hidden_dim, model
+  def encode(self, ): raise NotImplementedError("implement me")
+  @classmethod
+  def load_from_pretrained(cls, checkpoint_path:str, checkpoint_url:str) -> ContentVec:
+    contentvec = ContentVec.load_from_pretrained(checkpoint_path, checkpoint_url)
+    return cls(contentvec)
+
+class ContentVec256L9(SpeechEncoder):
+  def __init__(self, model:ContentVec): super().__init__(hidden_dim=256, model=model)
+  def encode(self, wav: Tensor):
+    feats = wav
+    if len(feats.shape) == 2:  # double channels
+      feats = feats.mean(-1)
+    assert len(feats.shape) == 1, feats.dim()
+    feats = feats.reshape(1, -1)
+    padding_mask = Tensor.zeros_like(feats).cast(dtypes.bool)
+    logits = self.model.extract_features(feats.to(wav.device), padding_mask=padding_mask.to(wav.device), output_layer=9)
+    feats = self.model.final_proj(logits[0])
+    return feats.transpose(1,2)
+
+class ContentVec768L12(SpeechEncoder):
+  def __init__(self, model:ContentVec): super().__init__(hidden_dim=768, model=model)
+  def encode(self, wav: Tensor):
+    feats = wav
+    if len(feats.shape) == 2:  # double channels
+      feats = feats.mean(-1)
+    assert len(feats.shape) == 1, feats.dim()
+    feats = feats.reshape(1, -1)
+    padding_mask = Tensor.zeros_like(feats).cast(dtypes.bool)
+    logits = self.model.extract_features(feats.to(wav.device), padding_mask=padding_mask.to(wav.device), output_layer=12)
+    return logits[0].transpose(1,2)
+
+# original code for contentvec: https://github.com/auspicious3000/contentvec/
+class ContentVec:
+  # self.final_proj dims are hardcoded and depend on fairseq.data.dictionary Dictionary in the checkpoint. This param can't yet be loaded since there is no pickle for it. See with DEBUG=2.
+  # This means that the ContentVec only works with the hubert weights used in all SVC models
+  def __init__(self, cfg: HParams):
+    self.feature_grad_mult, self.untie_final_proj = cfg.feature_grad_mult, cfg.untie_final_proj
+    feature_enc_layers = eval(cfg.conv_feature_layers)
+    self.embed = feature_enc_layers[-1][0]
+    final_dim = cfg.final_dim if cfg.final_dim > 0 else cfg.encoder_embed_dim
+    self.feature_extractor = ConvFeatureExtractionModel(conv_layers=feature_enc_layers, dropout=0.0, mode=cfg.extractor_mode, conv_bias=cfg.conv_bias)
+    self.post_extract_proj = nn.Linear(self.embed, cfg.encoder_embed_dim) if self.embed != cfg.encoder_embed_dim else None
+    self.encoder = TransformerEncoder(cfg)
+    self.layer_norm = nn.LayerNorm(self.embed)
+    self.final_proj = nn.Linear(cfg.encoder_embed_dim, final_dim * 1) if self.untie_final_proj else nn.Linear(cfg.encoder_embed_dim, final_dim)
+    self.mask_emb = Tensor.uniform(cfg.encoder_embed_dim, dtype=dtypes.float32)
+    self.label_embs_concat = Tensor.uniform(504, final_dim, dtype=dtypes.float32)
+  def forward_features(self, source, padding_mask):
+    if self.feature_grad_mult > 0:
+      features = self.feature_extractor(source, padding_mask)
+      if self.feature_grad_mult != 1.0: pass  # training: GradMultiply.forward(features, self.feature_grad_mult)
+    else:
+      features = self.feature_extractor(source, padding_mask)
+    return features
+  def forward_padding_mask(self, features, padding_mask):  # replaces original forward_padding_mask for batch inference
+    lengths_org = tilde(padding_mask.cast(dtypes.bool)).cast(dtypes.int64).sum(1)  # ensure its bool for tilde
+    lengths = (lengths_org - 400).float().div(320).floor().cast(dtypes.int64) + 1  # intermediate float to divide
+    padding_mask = lengths_to_padding_mask(lengths)
+    return padding_mask
+  def extract_features(self, source: Tensor, spk_emb:Tensor=None, padding_mask=None, ret_conv=False, output_layer=None, tap=False):
+    features = self.forward_features(source, padding_mask)
+    if padding_mask is not None:
+      padding_mask = self.forward_padding_mask(features, padding_mask)
+    features = features.transpose(1, 2)
+    features = self.layer_norm(features)
+    if self.post_extract_proj is not None:
+      features = self.post_extract_proj(features)
+    x, _ = self.encoder(features, spk_emb, padding_mask=padding_mask, layer=(None if output_layer is None else output_layer - 1), tap=tap)
+    res = features if ret_conv else x
+    return res, padding_mask
+  @classmethod
+  def load_from_pretrained(cls, checkpoint_path:str, checkpoint_url:str) -> ContentVec:
+    fetch(checkpoint_url, checkpoint_path)
+    cfg = load_fairseq_cfg(checkpoint_path)
+    enc = cls(cfg.model)
+    _ = load_checkpoint_enc(checkpoint_path, enc, None)
+    logging.debug(f"{cls.__name__}: Loaded model with cfg={cfg}")
+    return enc
+
+class TransformerEncoder:
+  def __init__(self, cfg: HParams):
+    def make_conv() -> nn.Conv1d:
+      layer = nn.Conv1d(self.embedding_dim, self.embedding_dim, kernel_size=cfg.conv_pos, padding=cfg.conv_pos // 2, groups=cfg.conv_pos_groups)
+      std = std = math.sqrt(4 / (cfg.conv_pos * self.embedding_dim))
+      layer.weight, layer.bias = (Tensor.normal(*layer.weight.shape, std=std)), (Tensor.zeros(*layer.bias.shape))
+      # for training: layer.weights need to be weight_normed
+      return layer
+    self.dropout, self.embedding_dim, self.layer_norm_first, self.layerdrop, self.num_layers, self.num_layers_1 = cfg.dropout, cfg.encoder_embed_dim, cfg.layer_norm_first, cfg.encoder_layerdrop, cfg.encoder_layers, cfg.encoder_layers_1
+    self.pos_conv, self.pos_conv_remove = [make_conv()], (1 if cfg.conv_pos % 2 == 0 else 0)
+    self.layers = [
+      TransformerEncoderLayer(self.embedding_dim, cfg.encoder_ffn_embed_dim, cfg.encoder_attention_heads, self.dropout, cfg.attention_dropout, cfg.activation_dropout, cfg.activation_fn, self.layer_norm_first, cond_layer_norm=(i >= cfg.encoder_layers))
+      for i in range(cfg.encoder_layers + cfg.encoder_layers_1)
+      ]
+    self.layer_norm = nn.LayerNorm(self.embedding_dim)
+    self.cond_layer_norm = CondLayerNorm(self.embedding_dim) if cfg.encoder_layers_1 > 0 else None
+    # training: apply init_bert_params
+  def __call__(self, x, spk_emb, padding_mask=None, layer=None, tap=False):
+    x, layer_results = self.extract_features(x, spk_emb, padding_mask, layer, tap)
+    if self.layer_norm_first and layer is None:
+      x = self.cond_layer_norm(x, spk_emb) if (self.num_layers_1 > 0) else self.layer_norm(x)
+    return x, layer_results
+  def extract_features(self, x: Tensor, spk_emb: Tensor, padding_mask=None, tgt_layer=None, tap=False):
+    if tgt_layer is not None:  # and not self.training
+      assert tgt_layer >= 0 and tgt_layer < len(self.layers)
+    if padding_mask is not None:
+      # x[padding_mask] = 0
+      assert padding_mask.shape == x.shape[:len(padding_mask.shape)]  # first few dims of x must match padding_mask
+      tmp_mask = padding_mask.unsqueeze(-1).repeat((1, 1, x.shape[-1]))
+      tmp_mask = tilde(tmp_mask.cast(dtypes.bool))
+      x = tmp_mask.where(x, 0)
+    x_conv = self.pos_conv[0](x.transpose(1,2))
+    if self.pos_conv_remove > 0: x_conv = x_conv[:, :, : -self.pos_conv_remove]
+    x_conv = x_conv.gelu().transpose(1, 2)
+    x = (x + x_conv).transpose(0, 1)  # B x T x C -> T x B x C
+    if not self.layer_norm_first: x = self.layer_norm(x)
+    x = x.dropout(p=self.dropout)
+    layer_results = []
+    r = None
+    for i, layer in enumerate(self.layers):
+      if i < self.num_layers:  # if (not self.training or (dropout_probability > self.layerdrop)) and (i < self.num_layers):
+        assert layer.cond_layer_norm == False
+        x = layer(x, self_attn_padding_mask=padding_mask, need_weights=False)
+        if tgt_layer is not None or tap:
+          layer_results.append(x.transpose(0, 1))
+      if i>= self.num_layers:
+        assert layer.cond_layer_norm == True
+        x = layer(x, emb=spk_emb, self_attn_padding_mask=padding_mask, need_weights=False)
+      if i == tgt_layer:
+        r = x
+        break
+    if r is not None:
+      x = r
+    x = x.transpose(0, 1)  # T x B x C -> B x T x C
+    return x, layer_results
+
+class TransformerEncoderLayer:
+  def __init__(self, embedding_dim=768.0, ffn_embedding_dim=3072.0, num_attention_heads=8.0, dropout=0.1, attention_dropout=0.1, activation_dropout=0.1, activation_fn="relu", layer_norm_first=False, cond_layer_norm=False):
+    def get_activation_fn(activation):
+      if activation == "relu": return Tensor.relu
+      if activation == "gelu": return Tensor.gelu
+      else: raise RuntimeError(f"activation function={activation} is not forseen")
+    self.embedding_dim, self.dropout, self.activation_dropout, self.layer_norm_first, self.num_attention_heads, self.cond_layer_norm, self.activation_fn = embedding_dim, dropout, activation_dropout, layer_norm_first, num_attention_heads, cond_layer_norm, get_activation_fn(activation_fn)
+    self.self_attn = MultiHeadAttention(self.embedding_dim, self.num_attention_heads)
+    self.self_attn_layer_norm = nn.LayerNorm(self.embedding_dim) if not cond_layer_norm else CondLayerNorm(self.embedding_dim)
+    self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
+    self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)
+    self.final_layer_norm = nn.LayerNorm(self.embedding_dim) if not cond_layer_norm else CondLayerNorm(self.embedding_dim)
+  def __call__(self, x:Tensor, self_attn_mask:Tensor=None, self_attn_padding_mask:Tensor=None, emb:Tensor=None, need_weights=False):
+    #self_attn_padding_mask = self_attn_padding_mask.reshape(x.shape[0], 1, 1, self_attn_padding_mask.shape[1]).expand(-1, self.num_attention_heads, -1, -1).reshape(x.shape[0] * self.num_attention_heads, 1, self_attn_padding_mask.shape[1]) if self_attn_padding_mask is not None else None
+    assert self_attn_mask is None and self_attn_padding_mask is not None
+    residual = x
+    if self.layer_norm_first:
+      x = self.self_attn_layer_norm(x) if not self.cond_layer_norm else self.self_attn_layer_norm(x, emb)
+      x = self.self_attn(x=x, mask=self_attn_padding_mask)
+      x = x.dropout(self.dropout)
+      x = residual + x
+      x = self.final_layer_norm(x) if not self.cond_layer_norm else self.final_layer_norm(x, emb)
+      x = self.activation_fn(self.fc1(x))
+      x = x.dropout(self.activation_dropout)
+      x = self.fc2(x)
+      x = x.dropout(self.dropout)
+      x = residual + x
+    else:
+      x = self.self_attn(x=x, mask=self_attn_padding_mask)
+      x = x.dropout(self.dropout)
+      x = residual + x
+      x = self.self_attn_layer_norm(x) if not self.cond_layer_norm else self.self_attn_layer_norm(x, emb)
+      residual = x
+      x = self.activation_fn(self.fc1(x))
+      x = x.dropout(self.activation_dropout)
+      x = self.fc2(x)
+      x = x.dropout(self.dropout)
+      x = residual + x
+      x = self.final_layer_norm(x) if not self.cond_layer_norm else self.final_layer_norm(x, emb)
+    return x
+
+class MultiHeadAttention:
+  def __init__(self, n_state, n_head):
+    self.n_state, self.n_head = n_state, n_head
+    self.q_proj, self.k_proj, self.v_proj, self.out_proj = [nn.Linear(n_state, n_state) for _ in range(4)]
+  def __call__(self, x:Tensor, xa:Optional[Tensor]=None, mask:Optional[Tensor]=None):
+    x = x.transpose(0,1)  # TxBxC -> BxTxC
+    q, k, v = self.q_proj(x), self.k_proj(xa or x), self.v_proj(xa or x)
+    q, k, v = [x.reshape(*q.shape[:2], self.n_head, -1) for x in (q, k, v)]
+    wv = Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), None).transpose(1, 2).reshape(*x.shape[:2], -1)
+    ret =  self.out_proj(wv).transpose(0,1)  # BxTxC -> TxBxC
+    return ret
+
+class ConvFeatureExtractionModel:
+  def __init__(self, conv_layers, dropout=.0, mode="default", conv_bias=False):
+    assert mode in {"default", "group_norm_masked", "layer_norm"}
+    def block(n_in, n_out, k, stride, is_layer_norm=False, is_group_norm=False, conv_bias=False):
+      def make_conv():
+        conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias)
+        conv.weight = Tensor.kaiming_normal(*conv.weight.shape)
+        return conv
+      assert (is_layer_norm and is_group_norm) == False, "layer norm and group norm are exclusive"
+      if is_layer_norm:
+        return [make_conv(), partial(Tensor.dropout, p=dropout),[partial(Tensor.transpose, dim0=-2, dim1=-1), nn.LayerNorm(dim, elementwise_affine=True), partial(Tensor.transpose, dim0=-2, dim1=-1)], Tensor.gelu]
+      elif is_group_norm and mode == "default":
+        return [make_conv(), partial(Tensor.dropout, p=dropout), nn.GroupNorm(dim, dim, affine=True), Tensor.gelu]
+      elif is_group_norm and mode == "group_norm_masked":
+        return [make_conv(), partial(Tensor.dropout, p=dropout), GroupNormMasked(dim, dim, affine=True), Tensor.gelu]
+      else:
+        return [make_conv(), partial(Tensor.dropout, p=dropout), Tensor.gelu]
+    in_d, self.conv_layers, self.mode = 1, [], mode
+    for i, cl in enumerate(conv_layers):
+      assert len(cl) == 3, "invalid conv definition: " + str(cl)
+      (dim, k, stride) = cl
+      if i == 0: self.cl = cl
+      self.conv_layers.append(block(in_d, dim, k, stride, is_layer_norm=(mode == "layer_norm"), is_group_norm=((mode == "default" or mode == "group_norm_masked") and i == 0), conv_bias=conv_bias))
+      in_d = dim
+  def __call__(self, x:Tensor, padding_mask:Tensor):
+    x = x.unsqueeze(1)  # BxT -> BxCxT
+    if self.mode == "group_norm_masked":
+      if padding_mask is not None:
+        _, k, stride = self.cl
+        lengths_org = tilde(padding_mask.cast(dtypes.bool)).cast(dtypes.int64).sum(1)  # ensure padding_mask is bool for tilde
+        lengths = (((lengths_org - k) / stride) + 1).floor().cast(dtypes.int64)
+        padding_mask = tilde(lengths_to_padding_mask(lengths)).cast(dtypes.int64)  # lengths_to_padding_mask returns bool tensor
+      x = self.conv_layers[0][0](x)  # padding_mask is numeric
+      x = self.conv_layers[0][1](x)
+      x = self.conv_layers[0][2](x, padding_mask)
+      x = self.conv_layers[0][3](x)
+    else:
+      x = x.sequential(self.conv_layers[0])  # default
+    for _, conv in enumerate(self.conv_layers[1:], start=1):
+      conv = reduce(lambda a,b: operator.iconcat(a,b if isinstance(b, list) else [b]), conv, [])  # flatten
+      x = x.sequential(conv)
+    return x
+
+class CondLayerNorm:  # https://github.com/auspicious3000/contentvec/blob/main/contentvec/modules/cond_layer_norm.py#L10
+  def __init__(self, dim_last, eps=1e-5, dim_spk=256, elementwise_affine=True):
+    self.dim_last, self.eps, self.dim_spk, self.elementwise_affine = dim_last, eps, dim_spk, elementwise_affine
+    if self.elementwise_affine:
+      self.weight_ln = nn.Linear(self.dim_spk, self.dim_last, bias=False)
+      self.bias_ln = nn.Linear(self.dim_spk, self.dim_last, bias=False)
+      self.weight_ln.weight, self.bias_ln.weight = (Tensor.ones(*self.weight_ln.weight.shape)), (Tensor.zeros(*self.bias_ln.weight.shape))
+  def __call__(self, x: Tensor, spk_emb: Tensor):
+    axis = tuple(-1-i for i in range(len(x.shape[1:])))
+    x = x.layernorm(axis=axis, eps=self.eps)
+    if not self.elementwise_affine: return x
+    weights, bias = self.weight_ln(spk_emb), self.bias_ln(spk_emb)
+    return weights * x + bias
+
+class GroupNormMasked:  # https://github.com/auspicious3000/contentvec/blob/d746688a32940f4bee410ed7c87ec9cf8ff04f74/contentvec/modules/fp32_group_norm.py#L16
+  def __init__(self, num_groups, num_channels, eps=1e-5, affine=True):
+    self.num_groups, self.num_channels, self.eps, self.affine = num_groups, num_channels, eps, affine
+    self.weight, self.bias = (Tensor.ones(num_channels)), (Tensor.zeros(num_channels)) if self.affine else (None, None)
+  def __call__(self, x:Tensor, mask:Tensor):
+    bsz, n_c, length = x.shape
+    assert n_c % self.num_groups == 0
+    x = x.reshape(bsz, self.num_groups, n_c // self.num_groups, length)
+    if mask is None: mask = Tensor.ones_like(x)
+    else: mask = mask.reshape(bsz, 1, 1, length)
+    x = x * mask
+    lengths = mask.sum(axis=3, keepdim=True)
+    assert x.shape[2] == 1
+    mean_ = x.mean(dim=3, keepdim=True)
+    mean = mean_ * length / lengths
+    var = (((x.std(axis=3, keepdim=True) ** 2) + mean_**2) * length / lengths - mean**2) + self.eps
+    return x.add(-mean).div(var.sqrt()).reshape(bsz, n_c, length).mul(self.weight.reshape(1,-1,1)).add(self.bias.reshape(1,-1,1))
+
+class Synthesizer:
+  def __init__(self, spec_channels, segment_size, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels, ssl_dim, n_speakers, sampling_rate=44100, vol_embedding=False, n_flow_layer=4, **kwargs):
+    self.spec_channels, self.inter_channels, self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout, self.resblock, self.resblock_kernel_sizes, self.resblock_dilation_sizes, self.upsample_rates, self.upsample_initial_channel, self.upsample_kernel_sizes, self.segment_size, self.n_speakers, self.gin_channels, self.vol_embedding = spec_channels, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, segment_size, n_speakers, gin_channels, vol_embedding
+    self.emb_g = nn.Embedding(n_speakers, gin_channels)
+    if vol_embedding: self.emb_vol = nn.Linear(1, hidden_channels)
+    self.pre = nn.Conv1d(ssl_dim, hidden_channels, kernel_size=5, padding=2)
+    self.enc_p = TextEncoder(inter_channels, hidden_channels, kernel_size, n_layers, filter_channels=filter_channels, n_heads=n_heads, p_dropout=p_dropout)
+    self.dec = Generator(sampling_rate, inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels)
+    self.enc_q = PosteriorEncoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
+    self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, n_flow_layer, gin_channels=gin_channels)
+    self.emb_uv = nn.Embedding(vocab_size=2, embed_size=hidden_channels)
+  def infer(self, c:Tensor, f0:Tensor, uv:Tensor, g:Tensor=None, noise_scale=0.35, seed=52468, vol=None) -> Tuple[Tensor, Tensor]:
+    Tensor.manual_seed(getenv('SEED', seed))
+    c_lengths = (Tensor.ones([c.shape[0]]) * c.shape[-1]).to(c.device)
+    if len(g.shape) == 1: g = g.unsqueeze(0)
+    g = self.emb_g(g).transpose(1, 2)
+    x_mask = sequence_mask(c_lengths, c.shape[2]).unsqueeze(1).cast(c.dtype)
+    vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol is not None and self.vol_embedding else 0
+    x = self.pre(c) * x_mask + self.emb_uv(uv.cast(dtypes.int64)).transpose(1, 2) + vol
+    z_p, _, _, c_mask = self.enc_p.forward(x, x_mask, f0=self._f0_to_coarse(f0), noise_scale=noise_scale)
+    z = self.flow.forward(z_p, c_mask, g=g, reverse=True)
+    o = self.dec.forward(z * c_mask, g=g, f0=f0)
+    return o,f0
+  def _f0_to_coarse(self, f0 : Tensor):
+    f0_mel = 1127 * (1 + f0 / 700).log()
+    a = (F0_BIN - 2) / (F0_MEL_MAX - F0_MEL_MIN)
+    b = F0_MEL_MIN * a - 1.
+    f0_mel = (f0_mel > 0).where(f0_mel * a - b, f0_mel)
+    f0_coarse = f0_mel.ceil().cast(dtype=dtypes.int64)
+    f0_coarse = f0_coarse * (f0_coarse > 0)
+    f0_coarse = f0_coarse + ((f0_coarse < 1) * 1)
+    f0_coarse = f0_coarse * (f0_coarse < F0_BIN)
+    f0_coarse = f0_coarse + ((f0_coarse >= F0_BIN) * (F0_BIN - 1))
+    return f0_coarse
+  @classmethod
+  def load_from_pretrained(cls, config_path:str, config_url:str, weights_path:str, weights_url:str) -> Synthesizer:
+    fetch(config_url, config_path)
+    hps = get_hparams_from_file(config_path)
+    fetch(weights_url, weights_path)
+    net_g = cls(hps.data.filter_length // 2 + 1, hps.train.segment_size // hps.data.hop_length, **hps.model)
+    _ = load_checkpoint(weights_path, net_g, None, skip_list=["f0_decoder"])
+    logging.debug(f"{cls.__name__}:Loaded model with hps: {hps}")
+    return net_g, hps
+
+class TextEncoder:
+  def __init__(self, out_channels, hidden_channels, kernel_size, n_layers, gin_channels=0, filter_channels=None, n_heads=None, p_dropout=None):
+    self.out_channels, self.hidden_channels, self.kernel_size, self.n_layers, self.gin_channels = out_channels, hidden_channels, kernel_size, n_layers, gin_channels
+    self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+    self.f0_emb = nn.Embedding(256, hidden_channels)  # n_vocab = 256
+    self.enc_ = Encoder(hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout)
+  def forward(self, x, x_mask, f0=None, noise_scale=1):
+    x = x + self.f0_emb(f0).transpose(1, 2)
+    x = self.enc_.forward(x * x_mask, x_mask)
+    stats = self.proj(x) * x_mask
+    m, logs = split(stats, self.out_channels, dim=1)
+    z = (m + randn_like(m) * logs.exp() * noise_scale) * x_mask
+    return z, m, logs, x_mask
+
+class Upsample:
+  def __init__(self, scale_factor):
+    assert scale_factor % 1 == 0, "Only integer scale factor allowed."
+    self.scale = int(scale_factor)
+  def forward(self, x:Tensor):
+    repeats = tuple([1] * len(x.shape) + [self.scale])
+    new_shape = (*x.shape[:-1], x.shape[-1] * self.scale)
+    return x.unsqueeze(-1).repeat(repeats).reshape(new_shape)
+
+class SineGen:
+  def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voice_threshold=0, flag_for_pulse=False):
+    self.sine_amp, self.noise_std, self.harmonic_num, self.sampling_rate, self.voiced_threshold, self.flag_for_pulse = sine_amp, noise_std, harmonic_num, samp_rate, voice_threshold, flag_for_pulse
+    self.dim = self.harmonic_num + 1
+  def _f02uv(self, f0): return (f0 > self.voiced_threshold).float()  #generate uv signal
+  def _f02sine(self, f0_values):
+    def padDiff(x : Tensor): return (x.pad((0,0,-1,1)) - x).pad((0,0,0,-1))
+    def mod(x: Tensor, n: int) -> Tensor: return x - n * x.div(n).floor()  # this is what the % operator does in pytorch.
+    rad_values = mod((f0_values / self.sampling_rate) , 1)  # convert to F0 in rad
+    rand_ini = Tensor.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device)  # initial phase noise
+
+    #rand_ini[:, 0] = 0
+    m = Tensor.ones(f0_values.shape[0]).unsqueeze(1).pad((0,f0_values.shape[2]-1,0,0)).cast(dtypes.bool)
+    m = tilde(m)
+    rand_ini = m.where(rand_ini, 0)
+
+    #rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
+    tmp = rad_values[:, 0, :] + rand_ini
+    m = Tensor.ones(tmp.shape).pad((0,0,0,rad_values.shape[1]-1,0)).cast(dtypes.bool)
+    m = tilde(m)
+    tmp = tmp.unsqueeze(1).pad((0,0,0,rad_values.shape[1]-1,0))
+    rad_values = m.where(rad_values, tmp)
+
+    tmp_over_one = mod(rad_values.cumsum(1), 1)
+    tmp_over_one_idx = padDiff(tmp_over_one) < 0
+    cumsum_shift = Tensor.zeros_like(rad_values)
+
+    #cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
+    tmp_over_one_idx = (tmp_over_one_idx * -1.0).pad((0,0,1,0))
+    cumsum_shift = tmp_over_one_idx
+
+    sines = ((rad_values + cumsum_shift).cumsum(1) * 2 * np.pi).sin()
+    return sines
+  def forward(self, f0, upp=None):
+    fn = f0.mul(Tensor([[range(1, self.harmonic_num + 2)]], dtype=dtypes.float32).to(f0.device))
+    sine_waves = self._f02sine(fn) * self.sine_amp  #generate sine waveforms
+    uv = self._f02uv(f0)  # generate uv signal
+    noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+    noise = noise_amp * randn_like(sine_waves)
+    sine_waves = sine_waves * uv + noise
+    return sine_waves, uv, noise
+
+class SourceHnNSF:
+  def __init__(self, sampling_rate, harmonic_num=0, sine_amp=0.1, add_noise_std=0.003, voiced_threshold=0):
+    self.sine_amp, self.noise_std = sine_amp, add_noise_std
+    self.l_sin_gen = SineGen(sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold)
+    self.l_linear = nn.Linear(harmonic_num + 1, 1)
+  def forward(self, x, upp=None):
+    sine_waves, uv, _ = self.l_sin_gen.forward(x, upp)
+    sine_merge = self.l_linear(sine_waves.cast(self.l_linear.weight.dtype)).tanh()
+    noise = randn_like(uv) * self.sine_amp / 3
+    return sine_merge, noise, uv
+
+# most of the hifigan in standard vits is reused here, but need to upsample and construct harmonic source from f0
+class Generator:
+  def __init__(self, sampling_rate, inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels):
+    self.sampling_rate, self.inter_channels, self.resblock, self.resblock_kernel_sizes, self.resblock_dilation_sizes, self.upsample_rates, self.upsample_initial_channel, self.upsample_kernel_sizes, self.gin_channels = sampling_rate, inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels
+    self.num_kernels, self.num_upsamples = len(resblock_kernel_sizes), len(upsample_rates)
+    self.conv_pre = nn.Conv1d(inter_channels, upsample_initial_channel, 7, 1, padding=3)
+    self.f0_upsamp = Upsample(scale_factor=np.prod(upsample_rates))
+    self.m_source = SourceHnNSF(sampling_rate, harmonic_num=8)
+    resblock = ResBlock1 if resblock == '1' else ResBlock2
+    self.ups, self.noise_convs, self.resblocks = [], [], []
+    for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+      c_cur = upsample_initial_channel//(2**(i+1))
+      self.ups.append(nn.ConvTranspose1d(upsample_initial_channel//(2**i), c_cur, k, u, padding=(k-u)//2))
+      stride_f0 = int(np.prod(upsample_rates[i + 1:]))
+      self.noise_convs.append(nn.Conv1d(1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2) if (i + 1 < len(upsample_rates)) else nn.Conv1d(1, c_cur, kernel_size=1))
+    for i in range(len(self.ups)):
+      ch = upsample_initial_channel // (2 ** (i + 1))
+      for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
+        self.resblocks.append(resblock(ch, k, d))
+    self.conv_post = nn.Conv1d(ch, 1, 7, 1, padding=3)
+    if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+    self.upp = np.prod(upsample_rates)
+  def forward(self, x, f0, g=None):
+    f0 = self.f0_upsamp.forward(f0[:, None]).transpose(1, 2)  # bs,n,t
+    har_source, _, _ = self.m_source.forward(f0, self.upp)
+    har_source = har_source.transpose(1, 2)
+    x = self.conv_pre(x)
+    if g is not None:  x = x + self.cond(g)
+    for i in range(self.num_upsamples):
+      x, xs = self.ups[i](x.leaky_relu(LRELU_SLOPE)), None
+      x_source = self.noise_convs[i](har_source)
+      x = x + x_source
+      for j in range(self.num_kernels):
+        if xs is None: xs = self.resblocks[i * self.num_kernels + j].forward(x)
+        else: xs += self.resblocks[i * self.num_kernels + j].forward(x)
+      x = xs / self.num_kernels
+    return self.conv_post(x.leaky_relu()).tanh()
+
+# **** helpers ****
+
+def randn_like(x:Tensor) -> Tensor: return Tensor.randn(*x.shape, dtype=x.dtype).to(device=x.device)
+
+def tilde(x: Tensor) -> Tensor:
+  if x.dtype == dtypes.bool: return (1 - x).cast(dtypes.bool)
+  return (x + 1) * -1  # this seems to be what the ~ operator does in pytorch for non bool
+
+def lengths_to_padding_mask(lens:Tensor) -> Tensor:
+  bsz, max_lens = lens.shape[0], lens.max().numpy().item()
+  mask = Tensor.arange(max_lens).to(lens.device).reshape(1, max_lens)
+  mask = mask.expand(bsz, -1) >= lens.reshape(bsz, 1).expand(-1, max_lens)
+  return mask.cast(dtypes.bool)
+
+def repeat_expand_2d_left(content, target_len): # content : [h, t]
+  src_len = content.shape[-1]
+  temp = np.arange(src_len+1) * target_len / src_len
+  current_pos, cols = 0, []
+  for i in range(target_len):
+    if i >= temp[current_pos+1]:
+      current_pos += 1
+    cols.append(content[:, current_pos])
+  return Tensor.stack(*cols).transpose(0, 1)
+
+def load_fairseq_cfg(checkpoint_path):
+  assert Path(checkpoint_path).is_file()
+  state = torch_load(checkpoint_path)
+  cfg = state["cfg"] if ("cfg" in state and state["cfg"] is not None) else None
+  if cfg is None: raise RuntimeError(f"No cfg exist in state keys = {state.keys()}")
+  return HParams(**cfg)
+
+def load_checkpoint_enc(checkpoint_path, model: ContentVec, optimizer=None, skip_list=[]):
+  assert Path(checkpoint_path).is_file()
+  start_time = time.time()
+  checkpoint_dict = torch_load(checkpoint_path)
+  saved_state_dict = checkpoint_dict['model']
+  weight_g, weight_v, parent = None, None, None
+  for key, v in saved_state_dict.items():
+    if any(layer in key for layer in skip_list): continue
+    try:
+      obj, skip = model, False
+      for k in key.split('.'):
+        if k.isnumeric(): obj = obj[int(k)]
+        elif isinstance(obj, dict): obj = obj[k]
+        else:
+          if k in ["weight_g", "weight_v"]:
+            parent, skip = obj, True
+            if k == "weight_g": weight_g = v
+            else: weight_v = v
+          if not skip:
+            parent = obj
+            obj = getattr(obj, k)
+      if weight_g and weight_v:
+        setattr(obj, "weight_g", weight_g.numpy())
+        setattr(obj, "weight_v", weight_v.numpy())
+        obj, v = getattr(parent, "weight"), weight_norm(weight_v, weight_g, 0)
+        weight_g, weight_v, parent, skip = None, None, None, False
+      if not skip and obj.shape == v.shape:
+        if "feature_extractor" in key and (isinstance(parent, (nn.GroupNorm, nn.LayerNorm))):  # cast
+          obj.assign(v.to(obj.device).float())
+        else:
+          obj.assign(v.to(obj.device))
+      elif not skip: logging.error(f"MISMATCH SHAPE IN {key}, {obj.shape} {v.shape}")
+    except Exception as e: raise e
+  logging.info(f"Loaded checkpoint '{checkpoint_path}' in {time.time() - start_time:.4f}s")
+  return model, optimizer
+
+def pad_array(arr, target_length):
+  current_length = arr.shape[0]
+  if current_length >= target_length: return arr
+  pad_width = target_length - current_length
+  pad_left = pad_width // 2
+  pad_right = pad_width - pad_left
+  padded_arr = np.pad(arr, (pad_left, pad_right), 'constant', constant_values=(0, 0))
+  return padded_arr
+
+def split_list_by_n(list_collection, n, pre=0):
+  for i in range(0, len(list_collection), n):
+    yield list_collection[i-pre if i-pre>=0 else i: i + n]
+
+def get_sid(spk2id:HParams, speaker:str) -> Tensor:
+  speaker_id = spk2id[speaker]
+  if not speaker_id and type(speaker) is int:
+    if len(spk2id.__dict__) >= speaker: speaker_id = speaker
+  if speaker_id is None: raise RuntimeError(f"speaker={speaker} not in the speaker list")
+  return Tensor([int(speaker_id)], dtype=dtypes.int64).unsqueeze(0)
+
+def get_encoder(ssl_dim) -> Type[SpeechEncoder]:
+  if ssl_dim == 256: return ContentVec256L9
+  if ssl_dim == 768: return ContentVec768L12
+
+#########################################################################################
+# CODE: https://github.com/svc-develop-team/so-vits-svc
+#########################################################################################
+# CONTENTVEC:
+#   CODE: https://github.com/auspicious3000/contentvec
+#   PAPER: https://arxiv.org/abs/2204.09224
+#########################################################################################
+# INSTALLATION: dependencies are for preprocessing and loading/saving audio.
+# pip3 install soundfile librosa praat-parselmouth
+#########################################################################################
+# EXAMPLE USAGE:
+# python3 examples/so_vits_svc.py --model tf2spy --file ~/recording.wav
+#########################################################################################
+# DEMO USAGE (uses audio sample from LJ-Speech):
+# python3 examples/so_vits_svc.py --model saul_goodman
+#########################################################################################
+SO_VITS_SVC_PATH = Path(__file__).parents[1] / "weights/So-VITS-SVC"
+VITS_MODELS = { # config_path, weights_path, config_url, weights_url
+  "saul_goodman" : (SO_VITS_SVC_PATH / "config_saul_gman.json", SO_VITS_SVC_PATH / "pretrained_saul_gman.pth", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/Saul_Goodman_80000/config.json", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/Saul_Goodman_80000/G_80000.pth"),
+  "drake" : (SO_VITS_SVC_PATH / "config_drake.json", SO_VITS_SVC_PATH / "pretrained_drake.pth", "https://huggingface.co/jaspa/so-vits-svc/resolve/main/aubrey/config_aubrey.json", "https://huggingface.co/jaspa/so-vits-svc/resolve/main/aubrey/pretrained_aubrey.pth"),
+  "cartman" : (SO_VITS_SVC_PATH / "config_cartman.json", SO_VITS_SVC_PATH / "pretrained_cartman.pth", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/EricCartman/config.json", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/EricCartman/G_10200.pth"),
+  "tf2spy" : (SO_VITS_SVC_PATH / "config_tf2spy.json", SO_VITS_SVC_PATH / "pretrained_tf2spy.pth", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_spy_60k/config.json", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_spy_60k/G_60000.pth"),
+  "tf2heavy" : (SO_VITS_SVC_PATH / "config_tf2heavy.json", SO_VITS_SVC_PATH / "pretrained_tf2heavy.pth", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_heavy_100k/config.json", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_heavy_100k/G_100000.pth"),
+  "lady_gaga" : (SO_VITS_SVC_PATH / "config_gaga.json", SO_VITS_SVC_PATH / "pretrained_gaga.pth", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/LadyGaga/config.json", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/LadyGaga/G_14400.pth")
+}
+ENCODER_MODELS = { # weights_path, weights_url
+  "contentvec": (SO_VITS_SVC_PATH / "contentvec_checkpoint.pt", "https://huggingface.co/lj1995/VoiceConversionWebUI/resolve/main/hubert_base.pt")
+}
+ENCODER_MODEL = "contentvec"
+DEMO_PATH, DEMO_URL = Path(__file__).parents[1] / "temp/LJ037-0171.wav", "https://keithito.com/LJ-Speech-Dataset/LJ037-0171.wav"
+if __name__=="__main__":
+  logging.basicConfig(stream=sys.stdout, level=(logging.INFO if DEBUG < 1 else logging.DEBUG))
+  parser = argparse.ArgumentParser()
+  parser.add_argument("-m", "--model", default=None, help=f"Specify the model to use. All supported models: {VITS_MODELS.keys()}", required=True)
+  parser.add_argument("-f", "--file", default=DEMO_PATH, help=f"Specify the path of the input file")
+  parser.add_argument("--out_dir", default=str(Path(__file__).parents[1] / "temp"), help="Specify the output path.")
+  parser.add_argument("--out_path", default=None, help="Specify the full output path. Overrides the --out_dir and --name parameter.")
+  parser.add_argument("--base_name", default="test", help="Specify the base of the output file name. Default is 'test'.")
+  parser.add_argument("--speaker", default=None, help="If not specified, the first available speaker is chosen. Usually there is only one speaker per model.")
+  parser.add_argument("--noise_scale", default=0.4)
+  parser.add_argument("--tran", default=0.0, help="Pitch shift, supports positive and negative (semitone) values. Default 0.0")
+  parser.add_argument("--pad_seconds", default=0.5)
+  parser.add_argument("--lg_num", default=0.0)
+  parser.add_argument("--clip_seconds", default=0.0)
+  parser.add_argument("--slice_db", default=-40)
+  args = parser.parse_args()
+
+  vits_model = args.model
+  encoder_location, vits_location = ENCODER_MODELS[ENCODER_MODEL], VITS_MODELS[vits_model]
+
+  Tensor.training = False
+  # Get Synthesizer and ContentVec
+  net_g, hps = Synthesizer.load_from_pretrained(vits_location[0], vits_location[2], vits_location[1], vits_location[3])
+  Encoder = get_encoder(hps.model.ssl_dim)
+  encoder = Encoder.load_from_pretrained(encoder_location[0], encoder_location[1])
+
+  # model config args
+  target_sample, spk2id, hop_length, target_sample = hps.data.sampling_rate, hps.spk, hps.data.hop_length, hps.data.sampling_rate
+  vol_embedding = hps.model.vol_embedding if hasattr(hps.data, "vol_embedding") and hps.model.vol_embedding is not None else False
+
+  # args
+  slice_db, clip_seconds, lg_num, pad_seconds, tran, noise_scale, audio_path = args.slice_db, args.clip_seconds, args.lg_num, args.pad_seconds, args.tran, args.noise_scale, args.file
+  speaker = args.speaker if args.speaker is not None else list(hps.spk.__dict__.keys())[0]
+
+  ### Loading audio and slicing ###
+  if audio_path == DEMO_PATH: fetch(DEMO_URL, DEMO_PATH)
+  assert Path(audio_path).is_file() and Path(audio_path).suffix == ".wav"
+  chunks = preprocess.cut(audio_path, db_thresh=slice_db)
+  audio_data, audio_sr = preprocess.chunks2audio(audio_path, chunks)
+
+  per_size = int(clip_seconds * audio_sr)
+  lg_size = int(lg_num * audio_sr)
+
+  ### Infer per slice ###
+  global_frame = 0
+  audio = []
+  for (slice_tag, data) in audio_data:
+    print(f"\n====segment start, {round(len(data) / audio_sr, 3)}s====")
+    length = int(np.ceil(len(data) / audio_sr * target_sample))
+
+    if slice_tag:
+      print("empty segment")
+      _audio = np.zeros(length)
+      audio.extend(list(pad_array(_audio, length)))
+      global_frame += length // hop_length
+      continue
+
+    datas = [data] if per_size == 0 else split_list_by_n(data, per_size, lg_size)
+
+    for k, dat in enumerate(datas):
+      per_length = int(np.ceil(len(dat) / audio_sr * target_sample)) if clip_seconds!=0 else length
+      pad_len = int(audio_sr * pad_seconds)
+      dat = np.concatenate([np.zeros([pad_len]), dat, np.zeros([pad_len])])
+      raw_path = io.BytesIO()
+      soundfile.write(raw_path, dat, audio_sr, format="wav")
+      raw_path.seek(0)
+
+      ### Infer START ###
+      wav, sr = preprocess.load_audiofile(raw_path)
+      wav = preprocess.sinc_interp_resample(wav, sr, target_sample)[0]
+      wav16k, f0, uv = preprocess.get_unit_f0(wav, tran, hop_length, target_sample)
+      sid = get_sid(spk2id, speaker)
+      n_frames = f0.shape[1]
+
+      # ContentVec infer
+      start = time.time()
+      c = encoder.encode(wav16k)
+      c = repeat_expand_2d_left(c.squeeze(0).realize(), f0.shape[1])  # interpolate speech encoding to match f0
+      c = c.unsqueeze(0).realize()
+      enc_time = time.time() - start
+
+      # VITS infer
+      vits_start = time.time()
+      out_audio, f0 = net_g.infer(c, f0=f0, uv=uv, g=sid, noise_scale=noise_scale, vol=None)
+      out_audio = out_audio[0,0].float().realize()
+      vits_time = time.time() - vits_start
+
+      infer_time = time.time() - start
+      logging.info("total infer time:{:.2f}s, speech_enc time:{:.2f}s, vits time:{:.2f}s".format(infer_time, enc_time, vits_time))
+      ### Infer END ###
+
+      out_sr, out_frame = out_audio.shape[-1], n_frames
+      global_frame += out_frame
+      _audio = out_audio.numpy()
+      pad_len = int(target_sample * pad_seconds)
+      _audio = _audio[pad_len:-pad_len]
+      _audio = pad_array(_audio, per_length)
+      audio.extend(list(_audio))
+
+  audio = np.array(audio)
+  out_path = Path(args.out_path or Path(args.out_dir)/f"{args.model}{f'_spk_{speaker}'}_{args.base_name}.wav")
+  out_path.parent.mkdir(parents=True, exist_ok=True)
+  soundfile.write(out_path, audio, target_sample, format="flac")
+  logging.info(f"Saved audio output to {out_path}")
@@ -0,0 +1,204 @@
+import math
+from typing import Optional, Tuple
+from tinygrad import Tensor, dtypes
+import librosa
+import soundfile
+import numpy as np
+import parselmouth
+
+class PMF0Predictor:  # from https://github.com/svc-develop-team/so-vits-svc/
+  def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
+    self.hop_length, self.f0_min, self.f0_max, self.sampling_rate, self.name = hop_length, f0_min, f0_max, sampling_rate, "pm"
+  def interpolate_f0(self,f0):
+    vuv_vector = np.zeros_like(f0, dtype=np.float32)
+    vuv_vector[f0 > 0.0] = 1.0
+    vuv_vector[f0 <= 0.0] = 0.0
+    nzindex = np.nonzero(f0)[0]
+    data = f0[nzindex]
+    nzindex = nzindex.astype(np.float32)
+    time_org = self.hop_length / self.sampling_rate * nzindex
+    time_frame = np.arange(f0.shape[0]) * self.hop_length / self.sampling_rate
+    if data.shape[0] <= 0: return np.zeros(f0.shape[0], dtype=np.float32),vuv_vector
+    if data.shape[0] == 1: return np.ones(f0.shape[0], dtype=np.float32) * f0[0],vuv_vector
+    f0 = np.interp(time_frame, time_org, data, left=data[0], right=data[-1])
+    return f0,vuv_vector
+  def compute_f0(self,wav,p_len=None):
+    x = wav
+    if p_len is None: p_len = x.shape[0]//self.hop_length
+    else: assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
+    time_step = self.hop_length / self.sampling_rate * 1000
+    f0 = parselmouth.Sound(x, self.sampling_rate) \
+                    .to_pitch_ac(time_step=time_step / 1000, voicing_threshold=0.6,pitch_floor=self.f0_min, pitch_ceiling=self.f0_max) \
+                    .selected_array['frequency']
+    pad_size=(p_len - len(f0) + 1) // 2
+    if(pad_size>0 or p_len - len(f0) - pad_size>0):
+      f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
+    f0,uv = self.interpolate_f0(f0)
+    return f0
+  def compute_f0_uv(self,wav,p_len=None):
+    x = wav
+    if p_len is None: p_len = x.shape[0]//self.hop_length
+    else: assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
+    time_step = self.hop_length / self.sampling_rate * 1000
+    f0 = parselmouth.Sound(x, self.sampling_rate).to_pitch_ac(
+      time_step=time_step / 1000, voicing_threshold=0.6,
+      pitch_floor=self.f0_min, pitch_ceiling=self.f0_max).selected_array['frequency']
+    pad_size=(p_len - len(f0) + 1) // 2
+    if(pad_size>0 or p_len - len(f0) - pad_size>0):
+      f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
+    f0,uv = self.interpolate_f0(f0)
+    return f0,uv
+
+class Slicer:  # from https://github.com/svc-develop-team/so-vits-svc/
+  def __init__(self, sr: int, threshold: float = -40., min_length: int = 5000, min_interval: int = 300, hop_size: int = 20, max_sil_kept: int = 5000):
+    if not min_length >= min_interval >= hop_size:
+      raise ValueError('The following condition must be satisfied: min_length >= min_interval >= hop_size')
+    if not max_sil_kept >= hop_size:
+      raise ValueError('The following condition must be satisfied: max_sil_kept >= hop_size')
+    min_interval = sr * min_interval / 1000
+    self.threshold = 10 ** (threshold / 20.)
+    self.hop_size = round(sr * hop_size / 1000)
+    self.win_size = min(round(min_interval), 4 * self.hop_size)
+    self.min_length = round(sr * min_length / 1000 / self.hop_size)
+    self.min_interval = round(min_interval / self.hop_size)
+    self.max_sil_kept = round(sr * max_sil_kept / 1000 / self.hop_size)
+  def _apply_slice(self, waveform, begin, end):
+    if len(waveform.shape) > 1: return waveform[:, begin * self.hop_size: min(waveform.shape[1], end * self.hop_size)]
+    else: return waveform[begin * self.hop_size: min(waveform.shape[0], end * self.hop_size)]
+  def slice(self, waveform):
+    samples = librosa.to_mono(waveform) if len(waveform.shape) > 1 else waveform
+    if samples.shape[0] <= self.min_length: return {"0": {"slice": False, "split_time": f"0,{len(waveform)}"}}
+    rms_list = librosa.feature.rms(y=samples, frame_length=self.win_size, hop_length=self.hop_size).squeeze(0)
+    sil_tags, silence_start, clip_start = [], None, 0
+    for i, rms in enumerate(rms_list):
+      if rms < self.threshold:  # Keep looping while frame is silent.
+        if silence_start is None:  # Record start of silent frames.
+          silence_start = i
+        continue
+      if silence_start is None: continue  # Keep looping while frame is not silent and silence start has not been recorded.
+      # Clear recorded silence start if interval is not enough or clip is too short
+      is_leading_silence = silence_start == 0 and i > self.max_sil_kept
+      need_slice_middle = i - silence_start >= self.min_interval and i - clip_start >= self.min_length
+      if not is_leading_silence and not need_slice_middle:
+        silence_start = None
+        continue
+      if i - silence_start <= self.max_sil_kept:  # Need slicing. Record the range of silent frames to be removed.
+        pos = rms_list[silence_start: i + 1].argmin() + silence_start
+        sil_tags.append((0, pos) if silence_start == 0 else (pos, pos))
+        clip_start = pos
+      elif i - silence_start <= self.max_sil_kept * 2:
+        pos = rms_list[i - self.max_sil_kept: silence_start + self.max_sil_kept + 1].argmin()
+        pos += i - self.max_sil_kept
+        pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
+        pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
+        if silence_start == 0:
+          sil_tags.append((0, pos_r))
+          clip_start = pos_r
+        else:
+          sil_tags.append((min(pos_l, pos), max(pos_r, pos)))
+          clip_start = max(pos_r, pos)
+      else:
+        pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
+        pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
+        sil_tags.append((0, pos_r) if silence_start == 0 else (pos_l, pos_r))
+        clip_start = pos_r
+      silence_start = None
+    total_frames = rms_list.shape[0]
+    if silence_start is not None and total_frames - silence_start >= self.min_interval:  # Deal with trailing silence.
+      silence_end = min(total_frames, silence_start + self.max_sil_kept)
+      pos = rms_list[silence_start: silence_end + 1].argmin() + silence_start
+      sil_tags.append((pos, total_frames + 1))
+    if len(sil_tags) == 0: return {"0": {"slice": False, "split_time": f"0,{len(waveform)}"}}  # Apply and return slices.
+    chunks = []
+    if sil_tags[0][0]:
+      chunks.append({"slice": False, "split_time": f"0,{min(waveform.shape[0], sil_tags[0][0] * self.hop_size)}"})
+    for i in range(0, len(sil_tags)):
+      if i: chunks.append({"slice": False, "split_time": f"{sil_tags[i - 1][1] * self.hop_size},{min(waveform.shape[0], sil_tags[i][0] * self.hop_size)}"})
+      chunks.append({"slice": True, "split_time": f"{sil_tags[i][0] * self.hop_size},{min(waveform.shape[0], sil_tags[i][1] * self.hop_size)}"})
+    if sil_tags[-1][1] * self.hop_size < len(waveform):
+      chunks.append({"slice": False, "split_time": f"{sil_tags[-1][1] * self.hop_size},{len(waveform)}"})
+    chunk_dict = {}
+    for i in range(len(chunks)): chunk_dict[str(i)] = chunks[i]
+    return chunk_dict
+
+# sinc_interp_hann audio resampling
+class Resample:
+  def __init__(self, orig_freq:int=16000, new_freq:int=16000, lowpass_filter_width:int=6, rolloff:float=0.99, beta:Optional[float]=None, dtype:Optional[dtypes]=None):
+    self.orig_freq, self.new_freq, self.lowpass_filter_width, self.rolloff, self.beta = orig_freq, new_freq, lowpass_filter_width, rolloff, beta
+    self.gcd = math.gcd(int(self.orig_freq), int(self.new_freq))
+    self.kernel, self.width = self._get_sinc_resample_kernel(dtype) if self.orig_freq != self.new_freq else (None, None)
+  def __call__(self, waveform:Tensor) -> Tensor:
+    if self.orig_freq == self.new_freq: return waveform
+    return self._apply_sinc_resample_kernel(waveform)
+  def _apply_sinc_resample_kernel(self, waveform:Tensor):
+    if not waveform.is_floating_point(): raise TypeError(f"Waveform tensor expected to be of type float, but received {waveform.dtype}.")
+    orig_freq, new_freq = (int(self.orig_freq) // self.gcd), (int(self.new_freq) // self.gcd)
+    shape = waveform.shape
+    waveform = waveform.reshape(-1, shape[-1])  # pack batch
+    num_wavs, length = waveform.shape
+    target_length = int(math.ceil(new_freq * length / orig_freq))
+    waveform = waveform.pad((self.width, self.width + orig_freq))
+    resampled = waveform[:, None].conv2d(self.kernel, stride=orig_freq)
+    resampled = resampled.transpose(1, 2).reshape(num_wavs, -1)
+    resampled = resampled[..., :target_length]
+    resampled = resampled.reshape(shape[:-1] + resampled.shape[-1:])  # unpack batch
+    return resampled
+  def _get_sinc_resample_kernel(self, dtype=None):
+    orig_freq, new_freq = (int(self.orig_freq) // self.gcd), (int(self.new_freq) // self.gcd)
+    if self.lowpass_filter_width <= 0: raise ValueError("Low pass filter width should be positive.")
+    base_freq = min(orig_freq, new_freq)
+    base_freq *= self.rolloff
+    width = math.ceil(self.lowpass_filter_width * orig_freq / base_freq)
+    idx = Tensor.arange(-width, width + orig_freq, dtype=(dtype if dtype is not None else dtypes.float32))[None, None] / orig_freq
+    t = Tensor.arange(0, -new_freq, -1, dtype=dtype)[:, None, None] / new_freq + idx
+    t *= base_freq
+    t = t.clip(-self.lowpass_filter_width, self.lowpass_filter_width)
+    window = (t * math.pi / self.lowpass_filter_width / 2).cos() ** 2
+    t *= math.pi
+    scale = base_freq / orig_freq
+    kernels = Tensor.where(t == 0, Tensor(1.0, dtype=t.dtype).to(t.device), t.sin() / t)
+    kernels *= window * scale
+    if dtype is None: kernels = kernels.cast(dtype=dtypes.float32)
+    return kernels, width
+
+def sinc_interp_resample(x:Tensor, orig_freq:int=16000, new_freq:int=1600, lowpass_filter_width:int=6, rolloff:float=0.99, beta:Optional[float]=None):
+  resamp = Resample(orig_freq, new_freq, lowpass_filter_width, rolloff, beta, x.dtype)
+  return resamp(x)
+
+def cut(audio_path, db_thresh=-30, min_len=5000):
+  audio, sr = librosa.load(audio_path, sr=None)
+  slicer = Slicer(sr=sr, threshold=db_thresh, min_length=min_len)
+  chunks = slicer.slice(audio)
+  return chunks
+
+def chunks2audio(audio_path, chunks):
+  chunks = dict(chunks)
+  audio, sr = load_audiofile(audio_path)
+  if len(audio.shape) == 2 and audio.shape[1] >= 2:
+    audio = audio.mean(0).unsqueeze(0)
+  audio = audio.numpy()[0]
+  result = []
+  for k, v in chunks.items():
+    tag = v["split_time"].split(",")
+    if tag[0] != tag[1]:
+      result.append((v["slice"], audio[int(tag[0]):int(tag[1])]))
+  return result, sr
+
+def load_audiofile(filepath:str, frame_offset:int=0, num_frames:int=-1, channels_first:bool=True):
+  with soundfile.SoundFile(filepath, "r") as file_:
+    frames = file_._prepare_read(frame_offset, None, num_frames)
+    waveform = file_.read(frames, "float32", always_2d=True)
+    sample_rate = file_.samplerate
+  waveform = Tensor(waveform)
+  if channels_first: waveform = waveform.transpose(0, 1)
+  return waveform, sample_rate
+
+def get_unit_f0(wav:Tensor, tran, hop_length, target_sample, f0_filter=False) -> Tuple[Tensor,Tensor,Tensor]:
+  f0_predictor = PMF0Predictor(hop_length, sampling_rate=target_sample)
+  f0, uv = f0_predictor.compute_f0_uv(wav.numpy())
+  if f0_filter and sum(f0) == 0: raise RuntimeError("No voice detected")
+  f0 = Tensor(f0.astype(np.float32)).float()
+  f0 = (f0 * 2 ** (tran / 12)).unsqueeze(0)
+  uv = Tensor(uv.astype(np.float32)).float().unsqueeze(0)
+  wav16k = sinc_interp_resample(wav[None,:], target_sample, 16000)[0]
+  return wav16k.realize(), f0.realize(), uv.realize()
@@ -9,7 +9,7 @@ from typing import Dict, Any
 from PIL import Image
 import numpy as np
 from tinygrad import Device, GlobalCounters, dtypes, Tensor, TinyJit
-from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten, profile_marker
+from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten
 from tinygrad.nn import Conv2d, GroupNorm
 from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
 from extra.models.clip import Closed, Tokenizer, FrozenOpenClipEmbedder
@@ -263,19 +263,13 @@ if __name__ == "__main__":
  parser.add_argument('--timing', action='store_true', help="Print timing per step")
  parser.add_argument('--seed', type=int, help="Set the random latent seed")
  parser.add_argument('--guidance', type=float, default=7.5, help="Prompt strength")
-  parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
  args = parser.parse_args()

-  profile_marker("create model")
  model = StableDiffusion()

-  profile_marker("load in weights")
+  # load in weights
  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    if not args.fakeweights:
-      model_bin = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
-      state_dict = torch_load(model_bin)['state_dict']
-      profile_marker("state dict loaded")
-      load_state_dict(model, state_dict, verbose=False, strict=False, realize=False)
+    load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], verbose=False, strict=False, realize=False)

    if args.fp16:
      for k,v in get_state_dict(model).items():
@@ -284,13 +278,12 @@ if __name__ == "__main__":

    Tensor.realize(*get_state_dict(model).values())

-  profile_marker("run clip (conditional)")
+  # run through CLIP to get context
  tokenizer = Tokenizer.ClipTokenizer()
  prompt = Tensor([tokenizer.encode(args.prompt)])
  context = model.cond_stage_model.transformer.text_model(prompt).realize()
  print("got CLIP context", context.shape)

-  profile_marker("run clip (unconditional)")
  prompt = Tensor([tokenizer.encode("")])
  unconditional_context = model.cond_stage_model.transformer.text_model(prompt).realize()
  print("got unconditional CLIP context", unconditional_context.shape)
@@ -314,7 +307,6 @@ if __name__ == "__main__":
  step_times = []
  with Context(BEAM=getenv("LATEBEAM")):
    for index, timestep in (t:=tqdm(list(enumerate(timesteps))[::-1])):
-      profile_marker(f"step {len(timesteps)-index-1}")
      GlobalCounters.reset()
      st = time.perf_counter_ns()
      t.set_description("%3d %3d" % (index, timestep))
@@ -324,26 +316,24 @@ if __name__ == "__main__":
          latent = run(model, unconditional_context, context, latent, Tensor([timestep]), alphas[tid], alphas_prev[tid], Tensor([args.guidance]))
          if args.timing: Device[Device.DEFAULT].synchronize()
      step_times.append((time.perf_counter_ns() - st)*1e-6)
-    # done with diffusion model
    del run
-    del model.model

  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
    min_time = min(step_times)
    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
-  profile_marker("run decoder") # upsample latent space to image with autoencoder
-  x = model.decode(latent).realize()
+  # upsample latent space to image with autoencoder
+  x = model.decode(latent)
  print(x.shape)

-  profile_marker("save image")
+  # save image
  im = Image.fromarray(x.numpy())
  print(f"saving {args.out}")
  im.save(args.out)
  # Open image.
  if not args.noshow: im.show()

+  # validation!
  if args.prompt == default_prompt and args.steps == 6 and args.seed == 0 and args.guidance == 7.5:
-    profile_marker("validate")
    ref_image = Tensor(np.array(Image.open(Path(__file__).parent / "stable_diffusion_seed0.png")))
    distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item()
    assert distance < 3e-3, colored(f"validation failed with {distance=}", "red")  # higher distance with WINO
@@ -19,8 +19,8 @@ from tinygrad.helpers import fetch, getenv

 # QUANT=1 python3 examples/test_onnx_imagenet.py
 #   https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
-# python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
-# VIZ=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx
+# DONT_REALIZE_EXPAND=1 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+# VIZ=1 DONT_REALIZE_EXPAND=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx

 def imagenet_dataloader(cnt=0):
  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
@@ -1,34 +0,0 @@
-#!/usr/bin/env python3
-from tinygrad import Tensor, Device, GlobalCounters, Context, dtypes
-from tinygrad.helpers import getenv, colored
-
-SZ = 8_000_000_000
-GPUS = getenv("GPUS", 4) # TODO: expose a way in tinygrad to access this
-
-if __name__ == "__main__":
-  # create tensors
-  tens = [Tensor.ones(SZ, dtype=dtypes.uint8, device=f"{Device.DEFAULT}:{i}").contiguous() for i in range(GPUS)]
-  Tensor.realize(*tens)
-
-  bw = [[0.0]*GPUS for _ in range(GPUS)]
-  for i in range(GPUS):
-    for j in range(GPUS):
-      GlobalCounters.reset()
-      with Context(DEBUG=2):
-        if i == j:
-          # this copy would be optimized out, just add 1
-          (tens[i]+1).realize()
-        else:
-          tens[i].to(f"{Device.DEFAULT}:{j}").realize()
-      t = max(GlobalCounters.time_sum_s, 1e-9)
-      bw[i][j] = SZ / t / 1e9  # GB/s
-
-  def fmt(x):
-    c = "green" if x > 50 else "yellow" if x > 20 else "red"
-    return colored(f"{x:6.1f}", c)
-
-  # header
-  print(" " * 8 + " ".join(f"{'d'+str(j):>6}" for j in range(GPUS)))
-  # rows
-  for i in range(GPUS):
-    print(f"{'s'+str(i):>6} -> " + " ".join(fmt(x) for x in bw[i]))
@@ -1,16 +0,0 @@
-from tinygrad import Tensor, Device, TinyJit, dtypes
-from tinygrad.helpers import getenv
-
-GPUS = getenv("GPUS", 4) # TODO: expose a way in tinygrad to access this
-N = 6144
-
-@TinyJit
-def many_matmul(A, B):
-  out = A
-  for _ in range(8): out = out@B
-  return out
-
-if __name__ == "__main__":
-  A = Tensor.ones(GPUS, N, N, dtype=dtypes.half).shard(devices=tuple([f"{Device.DEFAULT}:{i}" for i in range(GPUS)]), axis=0).contiguous()
-  B = Tensor.ones(GPUS, N, N, dtype=dtypes.half).shard(devices=tuple([f"{Device.DEFAULT}:{i}" for i in range(GPUS)]), axis=0).contiguous()
-  while 1: many_matmul(A, B)
@@ -0,0 +1,104 @@
+import traceback
+import time
+from multiprocessing import Process, Queue
+import numpy as np
+from tinygrad.nn.state import get_parameters
+from tinygrad.nn import optim
+from tinygrad.helpers import getenv, trange
+from tinygrad.tensor import Tensor
+from extra.datasets import fetch_cifar
+from extra.models.efficientnet import EfficientNet
+
+class TinyConvNet:
+  def __init__(self, classes=10):
+    conv = 3
+    inter_chan, out_chan = 8, 16   # for speed
+    self.c1 = Tensor.uniform(inter_chan,3,conv,conv)
+    self.c2 = Tensor.uniform(out_chan,inter_chan,conv,conv)
+    self.l1 = Tensor.uniform(out_chan*6*6, classes)
+
+  def forward(self, x):
+    x = x.conv2d(self.c1).relu().max_pool2d()
+    x = x.conv2d(self.c2).relu().max_pool2d()
+    x = x.reshape(shape=[x.shape[0], -1])
+    return x.dot(self.l1)
+
+if __name__ == "__main__":
+  IMAGENET = getenv("IMAGENET")
+  classes = 1000 if IMAGENET else 10
+
+  TINY = getenv("TINY")
+  TRANSFER = getenv("TRANSFER")
+  if TINY:
+    model = TinyConvNet(classes)
+  elif TRANSFER:
+    model = EfficientNet(getenv("NUM", 0), classes, has_se=True)
+    model.load_from_pretrained()
+  else:
+    model = EfficientNet(getenv("NUM", 0), classes, has_se=False)
+
+  parameters = get_parameters(model)
+  print("parameter count", len(parameters))
+  optimizer = optim.Adam(parameters, lr=0.001)
+
+  BS, steps = getenv("BS", 64 if TINY else 16), getenv("STEPS", 2048)
+  print(f"training with batch size {BS} for {steps} steps")
+
+  if IMAGENET:
+    from extra.datasets.imagenet import fetch_batch
+    def loader(q):
+      while 1:
+        try:
+          q.put(fetch_batch(BS))
+        except Exception:
+          traceback.print_exc()
+    q = Queue(16)
+    for i in range(2):
+      p = Process(target=loader, args=(q,))
+      p.daemon = True
+      p.start()
+  else:
+    X_train, Y_train, _, _ = fetch_cifar()
+    X_train = X_train.reshape((-1, 3, 32, 32))
+    Y_train = Y_train.reshape((-1,))
+
+  with Tensor.train():
+    for i in (t := trange(steps)):
+      if IMAGENET:
+        X, Y = q.get(True)
+      else:
+        samp = np.random.randint(0, X_train.shape[0], size=(BS))
+        X, Y = X_train.numpy()[samp], Y_train.numpy()[samp]
+
+      st = time.time()
+      out = model.forward(Tensor(X.astype(np.float32), requires_grad=False))
+      fp_time = (time.time()-st)*1000.0
+
+      y = np.zeros((BS,classes), np.float32)
+      y[range(y.shape[0]),Y] = -classes
+      y = Tensor(y, requires_grad=False)
+      loss = out.log_softmax().mul(y).mean()
+
+      optimizer.zero_grad()
+
+      st = time.time()
+      loss.backward()
+      bp_time = (time.time()-st)*1000.0
+
+      st = time.time()
+      optimizer.step()
+      opt_time = (time.time()-st)*1000.0
+
+      st = time.time()
+      loss = loss.numpy()
+      cat = out.argmax(axis=1).numpy()
+      accuracy = (cat == Y).mean()
+      finish_time = (time.time()-st)*1000.0
+
+      # printing
+      t.set_description("loss %.2f accuracy %.2f -- %.2f + %.2f + %.2f + %.2f = %.2f" %
+        (loss, accuracy,
+        fp_time, bp_time, opt_time, finish_time,
+        fp_time + bp_time + opt_time + finish_time))
+
+      del out, y, loss
@@ -0,0 +1,46 @@
+import ast
+import numpy as np
+from PIL import Image
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import getenv, fetch
+from extra.models.vit import ViT
+"""
+fn = "gs://vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz"
+import tensorflow as tf
+with tf.io.gfile.GFile(fn, "rb") as f:
+  dat = f.read()
+  with open("cache/"+ fn.rsplit("/", 1)[1], "wb") as g:
+    g.write(dat)
+"""
+
+Tensor.training = False
+if getenv("LARGE", 0) == 1:
+  m = ViT(embed_dim=768, num_heads=12)
+else:
+  # tiny
+  m = ViT(embed_dim=192, num_heads=3)
+m.load_from_pretrained()
+
+# category labels
+lbls = ast.literal_eval(fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt").read_text())
+
+#url = "https://upload.wikimedia.org/wikipedia/commons/4/41/Chicken.jpg"
+url = "https://repository-images.githubusercontent.com/296744635/39ba6700-082d-11eb-98b8-cb29fb7369c0"
+
+# junk
+img = Image.open(fetch(url))
+aspect_ratio = img.size[0] / img.size[1]
+img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0))))
+img = np.array(img)
+y0,x0=(np.asarray(img.shape)[:2]-224)//2
+img = img[y0:y0+224, x0:x0+224]
+img = np.moveaxis(img, [2,0,1], [0,1,2])
+img = img.astype(np.float32)[:3].reshape(1,3,224,224)
+img /= 255.0
+img -= 0.5
+img /= 0.5
+
+out = m.forward(Tensor(img))
+outnp = out.numpy().ravel()
+choice = outnp.argmax()
+print(out.shape, choice, outnp[choice], lbls[choice])
@@ -0,0 +1,740 @@
+import json, logging, math, re, sys, time, wave, argparse, numpy as np
+from phonemizer.phonemize import default_separator, _phonemize
+from phonemizer.backend import EspeakBackend
+from phonemizer.punctuation import Punctuation
+from functools import reduce
+from pathlib import Path
+from typing import List
+from tinygrad import nn, dtypes
+from tinygrad.helpers import fetch
+from tinygrad.nn.state import torch_load
+from tinygrad.tensor import Tensor
+from tinygrad.engine.jit import TinyJit
+from unidecode import unidecode
+
+LRELU_SLOPE = 0.1
+
+class Synthesizer:
+  def __init__(self, n_vocab, spec_channels, segment_size, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, n_speakers=0, gin_channels=0, use_sdp=True, emotion_embedding=False, **kwargs):
+    self.n_vocab, self.spec_channels, self.inter_channels, self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout, self.resblock, self.resblock_kernel_sizes, self.resblock_dilation_sizes, self.upsample_rates, self.upsample_initial_channel, self.upsample_kernel_sizes, self.segment_size, self.n_speakers, self.gin_channels, self.use_sdp = n_vocab, spec_channels, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, segment_size, n_speakers, gin_channels, use_sdp
+    self.enc_p = TextEncoder(n_vocab, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, emotion_embedding)
+    self.dec = Generator(inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=gin_channels)
+    self.enc_q = PosteriorEncoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
+    self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
+    self.dp = StochasticDurationPredictor(hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels) if use_sdp else DurationPredictor(hidden_channels, 256, 3, 0.5, gin_channels=gin_channels)
+    if n_speakers > 1: self.emb_g = nn.Embedding(n_speakers, gin_channels)
+  def infer(self, x, x_lengths, sid=None, noise_scale=1.0, length_scale=1, noise_scale_w=1., max_len=None, emotion_embedding=None, max_y_length_estimate_scale=None, pad_length=-1):
+    x, m_p, logs_p, x_mask = self.enc_p.forward(x.realize(), x_lengths.realize(), emotion_embedding.realize() if emotion_embedding is not None else emotion_embedding)
+    g = self.emb_g(sid.reshape(1, 1)).squeeze(1).unsqueeze(-1) if self.n_speakers > 0 else None
+    logw = self.dp.forward(x, x_mask.realize(), g=g.realize(), reverse=self.use_sdp, noise_scale=noise_scale_w if self.use_sdp else 1.0)
+    w_ceil = Tensor.ceil(logw.exp() * x_mask * length_scale)
+    y_lengths = Tensor.maximum(w_ceil.sum([1, 2]), 1).cast(dtypes.int64)
+    return self.generate(g, logs_p, m_p, max_len, max_y_length_estimate_scale, noise_scale, w_ceil, x, x_mask, y_lengths, pad_length)
+  def generate(self, g, logs_p, m_p, max_len, max_y_length_estimate_scale, noise_scale, w_ceil, x, x_mask, y_lengths, pad_length):
+    max_y_length = y_lengths.max().item() if max_y_length_estimate_scale is None else max(15, x.shape[-1]) * max_y_length_estimate_scale
+    y_mask = sequence_mask(y_lengths, max_y_length).unsqueeze(1).cast(x_mask.dtype)
+    attn_mask = x_mask.unsqueeze(2) * y_mask.unsqueeze(-1)
+    attn = generate_path(w_ceil, attn_mask)
+    m_p_2 = attn.squeeze(1).matmul(m_p.transpose(1, 2)).transpose(1, 2)        # [b, t', t], [b, t, d] -> [b, d, t']
+    logs_p_2 = attn.squeeze(1).matmul(logs_p.transpose(1, 2)).transpose(1, 2)  # [b, t', t], [b, t, d] -> [b, d, t']
+    z_p = m_p_2 + Tensor.randn(*m_p_2.shape, dtype=m_p_2.dtype) * logs_p_2.exp() * noise_scale
+    row_len = y_mask.shape[2]
+    if pad_length > -1:
+      # Pad flow forward inputs to enable JIT
+      assert pad_length > row_len, "pad length is too small"
+      y_mask = y_mask.pad(((0, 0), (0, 0), (0, pad_length - row_len))).cast(z_p.dtype)
+      # New y_mask tensor to remove sts mask
+      y_mask = Tensor(y_mask.numpy(), device=y_mask.device, dtype=y_mask.dtype, requires_grad=y_mask.requires_grad)
+      z_p = z_p.squeeze(0).pad(((0, 0), (0, pad_length - z_p.shape[2])), value=1).unsqueeze(0)
+    z = self.flow.forward(z_p.realize(), y_mask.realize(), g=g.realize(), reverse=True)
+    result_length = reduce(lambda x, y: x * y, self.dec.upsample_rates, row_len)
+    o = self.dec.forward((z * y_mask)[:, :, :max_len], g=g)[:, :, :result_length]
+    if max_y_length_estimate_scale is not None:
+      length_scaler = o.shape[-1] / max_y_length
+      o.realize()
+      real_max_y_length = y_lengths.max().numpy()
+      if real_max_y_length > max_y_length:
+        logging.warning(f"Underestimated max length by {(((real_max_y_length / max_y_length) * 100) - 100):.2f}%, recomputing inference without estimate...")
+        return self.generate(g, logs_p, m_p, max_len, None, noise_scale, w_ceil, x, x_mask, y_lengths)
+      if real_max_y_length < max_y_length:
+        overestimation = ((max_y_length / real_max_y_length) * 100) - 100
+        logging.info(f"Overestimated max length by {overestimation:.2f}%")
+        if overestimation > 10: logging.warning("Warning: max length overestimated by more than 10%")
+      o = o[:, :, :(real_max_y_length * length_scaler).astype(np.int32)]
+    return o
+
+class StochasticDurationPredictor:
+  def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, n_flows=4, gin_channels=0):
+    filter_channels = in_channels # it needs to be removed from future version.
+    self.in_channels, self.filter_channels, self.kernel_size, self.p_dropout, self.n_flows, self.gin_channels = in_channels, filter_channels, kernel_size, p_dropout, n_flows, gin_channels
+    self.log_flow, self.flows = Log(), [ElementwiseAffine(2)]
+    for _ in range(n_flows):
+      self.flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3))
+      self.flows.append(Flip())
+    self.post_pre, self.post_proj = nn.Conv1d(1, filter_channels, 1), nn.Conv1d(filter_channels, filter_channels, 1)
+    self.post_convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
+    self.post_flows = [ElementwiseAffine(2)]
+    for _ in range(4):
+      self.post_flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3))
+      self.post_flows.append(Flip())
+    self.pre, self.proj = nn.Conv1d(in_channels, filter_channels, 1), nn.Conv1d(filter_channels, filter_channels, 1)
+    self.convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
+    if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
+  @TinyJit
+  def forward(self, x: Tensor, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
+    x = self.pre(x.detach())
+    if g is not None: x = x + self.cond(g.detach())
+    x = self.convs.forward(x, x_mask)
+    x = self.proj(x) * x_mask
+    if not reverse:
+      flows = self.flows
+      assert w is not None
+      log_det_tot_q = 0
+      h_w = self.post_proj(self.post_convs.forward(self.post_pre(w), x_mask)) * x_mask
+      e_q = Tensor.randn(w.size(0), 2, w.size(2), dtype=x.dtype).to(device=x.device) * x_mask
+      z_q = e_q
+      for flow in self.post_flows:
+        z_q, log_det_q = flow.forward(z_q, x_mask, g=(x + h_w))
+        log_det_tot_q += log_det_q
+      z_u, z1 = z_q.split([1, 1], 1)
+      u = z_u.sigmoid() * x_mask
+      z0 = (w - u) * x_mask
+      log_det_tot_q += Tensor.sum((z_u.logsigmoid() + (-z_u).logsigmoid()) * x_mask, [1,2])
+      log_q = Tensor.sum(-0.5 * (math.log(2*math.pi) + (e_q**2)) * x_mask, [1,2]) - log_det_tot_q
+      log_det_tot = 0
+      z0, log_det = self.log_flow.forward(z0, x_mask)
+      log_det_tot += log_det
+      z = z0.cat(z1, 1)
+      for flow in flows:
+        z, log_det = flow.forward(z, x_mask, g=x, reverse=reverse)
+        log_det_tot = log_det_tot + log_det
+      nll = Tensor.sum(0.5 * (math.log(2*math.pi) + (z**2)) * x_mask, [1,2]) - log_det_tot
+      return (nll + log_q).realize() # [b]
+    flows = list(reversed(self.flows))
+    flows = flows[:-2] + [flows[-1]] # remove a useless vflow
+    z = Tensor.randn(x.shape[0], 2, x.shape[2], dtype=x.dtype).to(device=x.device) * noise_scale
+    for flow in flows: z = flow.forward(z, x_mask, g=x, reverse=reverse)
+    z0, z1 = z.split([1, 1], 1)
+    return z0.realize()
+
+class DurationPredictor:
+  def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0):
+    self.in_channels, self.filter_channels, self.kernel_size, self.p_dropout, self.gin_channels = in_channels, filter_channels, kernel_size, p_dropout, gin_channels
+    self.conv_1, self.norm_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size//2), LayerNorm(filter_channels)
+    self.conv_2, self.norm_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size//2), LayerNorm(filter_channels)
+    self.proj = nn.Conv1d(filter_channels, 1, 1)
+    if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+  def forward(self, x: Tensor, x_mask, g=None):
+    x = x.detach()
+    if g is not None: x = x + self.cond(g.detach())
+    x = self.conv_1(x * x_mask).relu()
+    x = self.norm_1(x).dropout(self.p_dropout)
+    x = self.conv_2(x * x_mask).relu(x)
+    x = self.norm_2(x).dropout(self.p_dropout)
+    return self.proj(x * x_mask) * x_mask
+
+class TextEncoder:
+  def __init__(self, n_vocab, out_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, emotion_embedding):
+    self.n_vocab, self.out_channels, self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout = n_vocab, out_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout
+    if n_vocab!=0:self.emb = nn.Embedding(n_vocab, hidden_channels)
+    if emotion_embedding: self.emo_proj = nn.Linear(1024, hidden_channels)
+    self.encoder = Encoder(hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout)
+    self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+  @TinyJit
+  def forward(self, x: Tensor, x_lengths: Tensor, emotion_embedding=None):
+    if self.n_vocab!=0: x = (self.emb(x) * math.sqrt(self.hidden_channels))
+    if emotion_embedding: x = x + self.emo_proj(emotion_embedding).unsqueeze(1)
+    x = x.transpose(1, -1)  # [b, t, h] -transpose-> [b, h, t]
+    x_mask = sequence_mask(x_lengths, x.shape[2]).unsqueeze(1).cast(x.dtype)
+    x = self.encoder.forward(x * x_mask, x_mask)
+    m, logs = (self.proj(x) * x_mask).split(self.out_channels, dim=1)
+    return x.realize(), m.realize(), logs.realize(), x_mask.realize()
+
+class ResidualCouplingBlock:
+  def __init__(self, channels, hidden_channels, kernel_size, dilation_rate, n_layers, n_flows=4, gin_channels=0):
+    self.channels, self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.n_flows, self.gin_channels = channels, hidden_channels, kernel_size, dilation_rate, n_layers, n_flows, gin_channels
+    self.flows = []
+    for _ in range(n_flows):
+      self.flows.append(ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels, mean_only=True))
+      self.flows.append(Flip())
+  @TinyJit
+  def forward(self, x, x_mask, g=None, reverse=False):
+    for flow in reversed(self.flows) if reverse else self.flows: x = flow.forward(x, x_mask, g=g, reverse=reverse)
+    return x.realize()
+
+class PosteriorEncoder:
+  def __init__(self, in_channels, out_channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0):
+    self.in_channels, self.out_channels, self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.gin_channels = in_channels, out_channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels
+    self.pre, self.proj = nn.Conv1d(in_channels, hidden_channels, 1), nn.Conv1d(hidden_channels, out_channels * 2, 1)
+    self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels)
+  def forward(self, x, x_lengths, g=None):
+    x_mask = sequence_mask(x_lengths, x.size(2)).unsqueeze(1).cast(x.dtype)
+    stats = self.proj(self.enc.forward(self.pre(x) * x_mask, x_mask, g=g)) * x_mask
+    m, logs = stats.split(self.out_channels, dim=1)
+    z = (m + Tensor.randn(m.shape, m.dtype) * logs.exp()) * x_mask
+    return z, m, logs, x_mask
+
+class Generator:
+  def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=0):
+    self.num_kernels, self.num_upsamples = len(resblock_kernel_sizes), len(upsample_rates)
+    self.conv_pre = nn.Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)
+    resblock = ResBlock1 if resblock == '1' else ResBlock2
+    self.ups = [nn.ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)), k, u, padding=(k-u)//2) for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes))]
+    self.resblocks = []
+    self.upsample_rates = upsample_rates
+    for i in range(len(self.ups)):
+      ch = upsample_initial_channel // (2 ** (i + 1))
+      for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
+        self.resblocks.append(resblock(ch, k, d))
+    self.conv_post = nn.Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+    if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+  @TinyJit
+  def forward(self, x: Tensor, g=None):
+    x = self.conv_pre(x)
+    if g is not None:  x = x + self.cond(g)
+    for i in range(self.num_upsamples):
+      x = self.ups[i](x.leaky_relu(LRELU_SLOPE))
+      xs = sum(self.resblocks[i * self.num_kernels + j].forward(x) for j in range(self.num_kernels))
+      x = (xs / self.num_kernels).realize()
+    res = self.conv_post(x.leaky_relu()).tanh().realize()
+    return res
+
+class LayerNorm(nn.LayerNorm):
+  def __init__(self, channels, eps=1e-5): super().__init__(channels, eps, elementwise_affine=True)
+  def forward(self, x: Tensor): return self.__call__(x.transpose(1, -1)).transpose(1, -1)
+
+class WN:
+  def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
+    assert (kernel_size % 2 == 1)
+    self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.gin_channels, self.p_dropout = hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels, p_dropout
+    self.in_layers, self.res_skip_layers = [], []
+    if gin_channels != 0: self.cond_layer = nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers, 1)
+    for i in range(n_layers):
+      dilation = dilation_rate ** i
+      self.in_layers.append(nn.Conv1d(hidden_channels, 2 * hidden_channels, kernel_size, dilation=dilation, padding=int((kernel_size * dilation - dilation) / 2)))
+      self.res_skip_layers.append(nn.Conv1d(hidden_channels, 2 * hidden_channels if i < n_layers - 1 else hidden_channels, 1))
+  def forward(self, x, x_mask, g=None, **kwargs):
+    output = Tensor.zeros_like(x)
+    if g is not None: g = self.cond_layer(g)
+    for i in range(self.n_layers):
+      x_in = self.in_layers[i](x)
+      if g is not None:
+        cond_offset = i * 2 * self.hidden_channels
+        g_l = g[:, cond_offset:cond_offset + 2 * self.hidden_channels, :]
+      else:
+        g_l = Tensor.zeros_like(x_in)
+      acts = fused_add_tanh_sigmoid_multiply(x_in, g_l, self.hidden_channels)
+      res_skip_acts = self.res_skip_layers[i](acts)
+      if i < self.n_layers - 1:
+        x = (x + res_skip_acts[:, :self.hidden_channels, :]) * x_mask
+        output = output + res_skip_acts[:, self.hidden_channels:, :]
+      else:
+        output = output + res_skip_acts
+    return output * x_mask
+
+class ResBlock1:
+  def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+    self.convs1 = [nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[i], padding=get_padding(kernel_size, dilation[i])) for i in range(3)]
+    self.convs2 = [nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)) for _ in range(3)]
+  def forward(self, x: Tensor, x_mask=None):
+    for c1, c2 in zip(self.convs1, self.convs2):
+      xt = x.leaky_relu(LRELU_SLOPE)
+      xt = c1(xt if x_mask is None else xt * x_mask).leaky_relu(LRELU_SLOPE)
+      x = c2(xt if x_mask is None else xt * x_mask) + x
+    return x if x_mask is None else x * x_mask
+
+class ResBlock2:
+  def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
+    self.convs = [nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[i], padding=get_padding(kernel_size, dilation[i])) for i in range(2)]
+  def forward(self, x, x_mask=None):
+    for c in self.convs:
+      xt = x.leaky_relu(LRELU_SLOPE)
+      xt = c(xt if x_mask is None else xt * x_mask)
+      x = xt + x
+    return x if x_mask is None else x * x_mask
+
+class DDSConv: # Dilated and Depth-Separable Convolution
+  def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):
+    self.channels, self.kernel_size, self.n_layers, self.p_dropout = channels, kernel_size, n_layers, p_dropout
+    self.convs_sep, self.convs_1x1, self.norms_1, self.norms_2 = [], [], [], []
+    for i in range(n_layers):
+      dilation = kernel_size ** i
+      padding = (kernel_size * dilation - dilation) // 2
+      self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size, groups=channels, dilation=dilation, padding=padding))
+      self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+      self.norms_1.append(LayerNorm(channels))
+      self.norms_2.append(LayerNorm(channels))
+  def forward(self, x, x_mask, g=None):
+    if g is not None: x = x + g
+    for i in range(self.n_layers):
+      y = self.convs_sep[i](x * x_mask)
+      y = self.norms_1[i].forward(y).gelu()
+      y = self.convs_1x1[i](y)
+      y = self.norms_2[i].forward(y).gelu()
+      x = x + y.dropout(self.p_dropout)
+    return x * x_mask
+
+class ConvFlow:
+  def __init__(self, in_channels, filter_channels, kernel_size, n_layers, num_bins=10, tail_bound=5.0):
+    self.in_channels, self.filter_channels, self.kernel_size, self.n_layers, self.num_bins, self.tail_bound = in_channels, filter_channels, kernel_size, n_layers, num_bins, tail_bound
+    self.half_channels = in_channels // 2
+    self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
+    self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.)
+    self.proj = nn.Conv1d(filter_channels, self.half_channels * (num_bins * 3 - 1), 1)
+  def forward(self, x, x_mask, g=None, reverse=False):
+    x0, x1 = x.split([self.half_channels] * 2, 1)
+    h = self.proj(self.convs.forward(self.pre(x0), x_mask, g=g)) * x_mask
+    b, c, t = x0.shape
+    h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) # [b, cx?, t] -> [b, c, t, ?]
+    un_normalized_widths = h[..., :self.num_bins] / math.sqrt(self.filter_channels)
+    un_normalized_heights = h[..., self.num_bins:2*self.num_bins] / math.sqrt(self.filter_channels)
+    un_normalized_derivatives = h[..., 2 * self.num_bins:]
+    x1, log_abs_det = piecewise_rational_quadratic_transform(x1, un_normalized_widths, un_normalized_heights, un_normalized_derivatives, inverse=reverse, tails='linear', tail_bound=self.tail_bound)
+    x = x0.cat(x1, dim=1) * x_mask
+    return x if reverse else (x, Tensor.sum(log_abs_det * x_mask, [1,2]))
+
+class ResidualCouplingLayer:
+  def __init__(self, channels, hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=0, gin_channels=0, mean_only=False):
+    assert channels % 2 == 0, "channels should be divisible by 2"
+    self.channels, self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.mean_only = channels, hidden_channels, kernel_size, dilation_rate, n_layers, mean_only
+    self.half_channels = channels // 2
+    self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+    self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
+    self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+  def forward(self, x, x_mask, g=None, reverse=False):
+    x0, x1 = x.split([self.half_channels] * 2, 1)
+    stats = self.post(self.enc.forward(self.pre(x0) * x_mask, x_mask, g=g)) * x_mask
+    if not self.mean_only:
+      m, logs = stats.split([self.half_channels] * 2, 1)
+    else:
+      m = stats
+      logs = Tensor.zeros_like(m)
+    if not reverse: return x0.cat((m + x1 * logs.exp() * x_mask), dim=1)
+    return x0.cat(((x1 - m) * (-logs).exp() * x_mask), dim=1)
+
+class Log:
+  def forward(self, x : Tensor, x_mask, reverse=False):
+    if not reverse:
+      y = x.maximum(1e-5).log() * x_mask
+      return y, (-y).sum([1, 2])
+    return x.exp() * x_mask
+
+class Flip:
+  def forward(self, x: Tensor, *args, reverse=False, **kwargs):
+    return x.flip([1]) if reverse else (x.flip([1]), Tensor.zeros(x.shape[0], dtype=x.dtype).to(device=x.device))
+
+class ElementwiseAffine:
+  def __init__(self, channels): self.m, self.logs = Tensor.zeros(channels, 1), Tensor.zeros(channels, 1)
+  def forward(self, x, x_mask, reverse=False, **kwargs): # x if reverse else y, logdet
+    return (x - self.m) * Tensor.exp(-self.logs) * x_mask if reverse \
+      else ((self.m + Tensor.exp(self.logs) * x) * x_mask, Tensor.sum(self.logs * x_mask, [1, 2]))
+
+class MultiHeadAttention:
+  def __init__(self, channels, out_channels, n_heads, p_dropout=0., window_size=None, heads_share=True, block_length=None, proximal_bias=False, proximal_init=False):
+    assert channels % n_heads == 0
+    self.channels, self.out_channels, self.n_heads, self.p_dropout, self.window_size, self.heads_share, self.block_length, self.proximal_bias, self.proximal_init = channels, out_channels, n_heads, p_dropout, window_size, heads_share, block_length, proximal_bias, proximal_init
+    self.attn, self.k_channels  = None, channels // n_heads
+    self.conv_q, self.conv_k, self.conv_v = [nn.Conv1d(channels, channels, 1) for _ in range(3)]
+    self.conv_o = nn.Conv1d(channels, out_channels, 1)
+    if window_size is not None: self.emb_rel_k, self.emb_rel_v = [Tensor.randn(1 if heads_share else n_heads, window_size * 2 + 1, self.k_channels) * (self.k_channels ** -0.5) for _ in range(2)]
+  def forward(self, x, c, attn_mask=None):
+    q, k, v = self.conv_q(x), self.conv_k(c), self.conv_v(c)
+    x, self.attn = self.attention(q, k, v, mask=attn_mask)
+    return self.conv_o(x)
+  def attention(self, query: Tensor, key: Tensor, value: Tensor, mask=None):# reshape [b, d, t] -> [b, n_h, t, d_k]
+    b, d, t_s, t_t = key.shape[0], key.shape[1], key.shape[2], query.shape[2]
+    query = query.reshape(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+    key = key.reshape(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+    value = value.reshape(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+    scores = (query / math.sqrt(self.k_channels)) @ key.transpose(-2, -1)
+    if self.window_size is not None:
+      assert t_s == t_t, "Relative attention is only available for self-attention."
+      key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+      rel_logits = self._matmul_with_relative_keys(query / math.sqrt(self.k_channels), key_relative_embeddings)
+      scores = scores + self._relative_position_to_absolute_position(rel_logits)
+    if mask is not None:
+      scores = Tensor.where(mask, scores, -1e4)
+      if self.block_length is not None:
+        assert t_s == t_t, "Local attention is only available for self-attention."
+        scores = Tensor.where(Tensor.ones_like(scores).triu(-self.block_length).tril(self.block_length), scores, -1e4)
+    p_attn = scores.softmax(axis=-1)  # [b, n_h, t_t, t_s]
+    output = p_attn.matmul(value)
+    if self.window_size is not None:
+      relative_weights = self._absolute_position_to_relative_position(p_attn)
+      value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s)
+      output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings)
+    output = output.transpose(2, 3).contiguous().reshape(b, d, t_t)  # [b, n_h, t_t, d_k] -> [b, d, t_t]
+    return output, p_attn
+  def _matmul_with_relative_values(self, x, y): return x.matmul(y.unsqueeze(0))                 # x: [b, h, l, m], y: [h or 1, m, d], ret: [b, h, l, d]
+  def _matmul_with_relative_keys(self, x, y): return x.matmul(y.unsqueeze(0).transpose(-2, -1)) # x: [b, h, l, d], y: [h or 1, m, d], re, : [b, h, l, m]
+  def _get_relative_embeddings(self, relative_embeddings, length):
+    pad_length, slice_start_position = max(length - (self.window_size + 1), 0), max((self.window_size + 1) - length, 0)
+    padded_relative_embeddings = relative_embeddings if pad_length <= 0\
+      else relative_embeddings.pad(convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]))
+    return padded_relative_embeddings[:, slice_start_position:(slice_start_position + 2 * length - 1)] #used_relative_embeddings
+  def _relative_position_to_absolute_position(self, x: Tensor): # x: [b, h, l, 2*l-1] -> [b, h, l, l]
+    batch, heads, length, _ = x.shape
+    x = x.pad(convert_pad_shape([[0,0],[0,0],[0,0],[0,1]]))
+    x_flat = x.reshape([batch, heads, length * 2 * length]).pad(convert_pad_shape([[0,0],[0,0],[0,length-1]]))
+    return x_flat.reshape([batch, heads, length+1, 2*length-1])[:, :, :length, length-1:]
+  def _absolute_position_to_relative_position(self, x: Tensor): # x: [b, h, l, l] -> [b, h, l, 2*l-1]
+    batch, heads, length, _ = x.shape
+    x = x.pad(convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length-1]]))
+    x_flat = x.reshape([batch, heads, length**2 + length*(length -1)]).pad(convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
+    return x_flat.reshape([batch, heads, length, 2*length])[:,:,:,1:]
+
+class FFN:
+  def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., activation=None, causal=False):
+    self.in_channels, self.out_channels, self.filter_channels, self.kernel_size, self.p_dropout, self.activation, self.causal = in_channels, out_channels, filter_channels, kernel_size, p_dropout, activation, causal
+    self.padding = self._causal_padding if causal else self._same_padding
+    self.conv_1, self.conv_2 = nn.Conv1d(in_channels, filter_channels, kernel_size), nn.Conv1d(filter_channels, out_channels, kernel_size)
+  def forward(self, x, x_mask):
+    x = self.conv_1(self.padding(x * x_mask))
+    x = x * (1.702 * x).sigmoid() if self.activation == "gelu" else x.relu()
+    return self.conv_2(self.padding(x.dropout(self.p_dropout) * x_mask)) * x_mask
+  def _causal_padding(self, x):return x if self.kernel_size == 1 else x.pad(convert_pad_shape([[0, 0], [0, 0], [self.kernel_size - 1, 0]]))
+  def _same_padding(self, x): return x if self.kernel_size == 1 else x.pad(convert_pad_shape([[0, 0], [0, 0], [(self.kernel_size - 1) // 2, self.kernel_size // 2]]))
+
+class Encoder:
+  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=4, **kwargs):
+    self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout, self.window_size = hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, window_size
+    self.attn_layers, self.norm_layers_1, self.ffn_layers, self.norm_layers_2 = [], [], [], []
+    for _ in range(n_layers):
+      self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, window_size=window_size))
+      self.norm_layers_1.append(LayerNorm(hidden_channels))
+      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout))
+      self.norm_layers_2.append(LayerNorm(hidden_channels))
+  def forward(self, x, x_mask):
+    attn_mask, x = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1), x * x_mask
+    for i in range(self.n_layers):
+      y = self.attn_layers[i].forward(x, x, attn_mask).dropout(self.p_dropout)
+      x = self.norm_layers_1[i].forward(x + y)
+      y = self.ffn_layers[i].forward(x, x_mask).dropout(self.p_dropout)
+      x = self.norm_layers_2[i].forward(x + y)
+    return x * x_mask
+
+DEFAULT_MIN_BIN_WIDTH, DEFAULT_MIN_BIN_HEIGHT, DEFAULT_MIN_DERIVATIVE = 1e-3, 1e-3, 1e-3
+def piecewise_rational_quadratic_transform(inputs, un_normalized_widths, un_normalized_heights, un_normalized_derivatives, inverse=False, tails=None, tail_bound=1., min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE):
+  if tails is None: spline_fn, spline_kwargs = rational_quadratic_spline, {}
+  else: spline_fn, spline_kwargs = unconstrained_rational_quadratic_spline, {'tails': tails, 'tail_bound': tail_bound}
+  return spline_fn(inputs=inputs, un_normalized_widths=un_normalized_widths, un_normalized_heights=un_normalized_heights, un_normalized_derivatives=un_normalized_derivatives, inverse=inverse, min_bin_width=min_bin_width, min_bin_height=min_bin_height, min_derivative=min_derivative, **spline_kwargs)
+def unconstrained_rational_quadratic_spline(inputs, un_normalized_widths, un_normalized_heights, un_normalized_derivatives, inverse=False, tails='linear', tail_bound=1., min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE):
+  if not tails == 'linear': raise RuntimeError('{} tails are not implemented.'.format(tails))
+  constant = np.log(np.exp(1 - min_derivative) - 1).item()
+  un_normalized_derivatives = cat_lr(un_normalized_derivatives, constant, constant)
+  output, log_abs_det = rational_quadratic_spline(inputs=inputs.squeeze(dim=0).squeeze(dim=0), unnormalized_widths=un_normalized_widths.squeeze(dim=0).squeeze(dim=0), unnormalized_heights=un_normalized_heights.squeeze(dim=0).squeeze(dim=0), unnormalized_derivatives=un_normalized_derivatives.squeeze(dim=0).squeeze(dim=0), inverse=inverse, left=-tail_bound, right=tail_bound, bottom=-tail_bound, top=tail_bound, min_bin_width=min_bin_width, min_bin_height=min_bin_height, min_derivative=min_derivative)
+  return output.unsqueeze(dim=0).unsqueeze(dim=0), log_abs_det.unsqueeze(dim=0).unsqueeze(dim=0)
+def rational_quadratic_spline(inputs: Tensor, unnormalized_widths: Tensor, unnormalized_heights: Tensor, unnormalized_derivatives: Tensor, inverse=False, left=0., right=1., bottom=0., top=1., min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE):
+  num_bins = unnormalized_widths.shape[-1]
+  if min_bin_width * num_bins > 1.0: raise ValueError('Minimal bin width too large for the number of bins')
+  if min_bin_height * num_bins > 1.0: raise ValueError('Minimal bin height too large for the number of bins')
+  widths = min_bin_width + (1 - min_bin_width * num_bins) * unnormalized_widths.softmax(axis=-1)
+  cum_widths = cat_lr(((right - left) * widths[..., :-1].cumsum(axis=1) + left), left, right + 1e-6 if not inverse else right)
+  widths = cum_widths[..., 1:] - cum_widths[..., :-1]
+  derivatives = min_derivative + (unnormalized_derivatives.exp()+1).log()
+  heights = min_bin_height + (1 - min_bin_height * num_bins) * unnormalized_heights.softmax(axis=-1)
+  cum_heights = cat_lr(((top - bottom) * heights[..., :-1].cumsum(axis=1) + bottom), bottom, top + 1e-6 if inverse else top)
+  heights = cum_heights[..., 1:] - cum_heights[..., :-1]
+  bin_idx = ((inputs[..., None] >= (cum_heights if inverse else cum_widths)).sum(axis=-1) - 1)[..., None]
+  input_cum_widths = gather(cum_widths, bin_idx, axis=-1)[..., 0]
+  input_bin_widths = gather(widths, bin_idx, axis=-1)[..., 0]
+  input_cum_heights = gather(cum_heights, bin_idx, axis=-1)[..., 0]
+  input_delta = gather(heights / widths, bin_idx, axis=-1)[..., 0]
+  input_derivatives = gather(derivatives, bin_idx, axis=-1)[..., 0]
+  input_derivatives_plus_one = gather(derivatives[..., 1:], bin_idx, axis=-1)[..., 0]
+  input_heights = gather(heights, bin_idx, axis=-1)[..., 0]
+  if inverse:
+    a = ((inputs - input_cum_heights) * (input_derivatives + input_derivatives_plus_one - 2 * input_delta) + input_heights * (input_delta - input_derivatives))
+    b = (input_heights * input_derivatives - (inputs - input_cum_heights) * (input_derivatives + input_derivatives_plus_one - 2 * input_delta))
+    c = - input_delta * (inputs - input_cum_heights)
+    discriminant = b.square() - 4 * a * c
+    # assert (discriminant.numpy() >= 0).all()
+    root = (2 * c) / (-b - discriminant.sqrt())
+    theta_one_minus_theta = root * (1 - root)
+    denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta) * theta_one_minus_theta)
+    derivative_numerator = input_delta.square() * (input_derivatives_plus_one * root.square() + 2 * input_delta * theta_one_minus_theta + input_derivatives * (1 - root).square())
+    return root * input_bin_widths + input_cum_widths, -(derivative_numerator.log() - 2 * denominator.log())
+  theta = (inputs - input_cum_widths) / input_bin_widths
+  theta_one_minus_theta = theta * (1 - theta)
+  numerator = input_heights * (input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta)
+  denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta) * theta_one_minus_theta)
+  derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2) + 2 * input_delta * theta_one_minus_theta + input_derivatives * (1 - theta).pow(2))
+  return input_cum_heights + numerator / denominator, derivative_numerator.log() - 2 * denominator.log()
+
+def sequence_mask(length: Tensor, max_length): return Tensor.arange(max_length, dtype=length.dtype, device=length.device).unsqueeze(0) < length.unsqueeze(1)
+def generate_path(duration: Tensor, mask: Tensor):  # duration: [b, 1, t_x], mask: [b, 1, t_y, t_x]
+  b, _, t_y, t_x = mask.shape
+  path = sequence_mask(duration.cumsum(axis=2).reshape(b * t_x), t_y).cast(mask.dtype).reshape(b, t_x, t_y)
+  path = path - path.pad(convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+  return path.unsqueeze(1).transpose(2, 3) * mask
+def fused_add_tanh_sigmoid_multiply(input_a: Tensor, input_b: Tensor, n_channels: int):
+  n_channels_int, in_act = n_channels, input_a + input_b
+  t_act, s_act = in_act[:, :n_channels_int, :].tanh(), in_act[:, n_channels_int:, :].sigmoid()
+  return t_act * s_act
+
+def cat_lr(t, left, right): return Tensor.full(get_shape(t), left).cat(t, dim=-1).cat(Tensor.full(get_shape(t), right), dim=-1)
+def get_shape(tensor):
+  (shape := list(tensor.shape))[-1] = 1
+  return tuple(shape)
+def convert_pad_shape(pad_shape): return tuple(tuple(x) for x in pad_shape)
+def get_padding(kernel_size, dilation=1): return int((kernel_size*dilation - dilation)/2)
+
+def gather(x, indices, axis):
+  indices = (indices < 0).where(indices + x.shape[axis], indices).transpose(0, axis)
+  permute_args = list(range(x.ndim))
+  permute_args[0], permute_args[axis] = permute_args[axis], permute_args[0]
+  permute_args.append(permute_args.pop(0))
+  x = x.permute(*permute_args)
+  reshape_arg = [1] * x.ndim + [x.shape[-1]]
+  return ((indices.unsqueeze(indices.ndim).expand(*indices.shape, x.shape[-1]) ==
+           Tensor.arange(x.shape[-1]).reshape(*reshape_arg).expand(*indices.shape, x.shape[-1])) * x).sum(indices.ndim).transpose(0, axis)
+
+def norm_except_dim(v, dim):
+  if dim == -1: return np.linalg.norm(v)
+  if dim == 0:
+    (output_shape := [1] * v.ndim)[0] = v.shape[0]
+    return np.linalg.norm(v.reshape(v.shape[0], -1), axis=1).reshape(output_shape)
+  if dim == v.ndim - 1:
+    (output_shape := [1] * v.ndim)[-1] = v.shape[-1]
+    return np.linalg.norm(v.reshape(-1, v.shape[-1]), axis=0).reshape(output_shape)
+  transposed_v = np.transpose(v, (dim,) + tuple(i for i in range(v.ndim) if i != dim))
+  return np.transpose(norm_except_dim(transposed_v, 0), (dim,) + tuple(i for i in range(v.ndim) if i != dim))
+def weight_norm(v: Tensor, g: Tensor, dim):
+  v, g = v.numpy(), g.numpy()
+  return Tensor(v * (g / norm_except_dim(v, dim)))
+
+# HPARAMS LOADING
+def get_hparams_from_file(path):
+  with open(path, "r") as f:
+    data = f.read()
+  return HParams(**json.loads(data))
+class HParams:
+  def __init__(self, **kwargs):
+    for k, v in kwargs.items(): self[k] = v if type(v) != dict else HParams(**v)
+  def keys(self): return self.__dict__.keys()
+  def items(self): return self.__dict__.items()
+  def values(self): return self.__dict__.values()
+  def __len__(self): return len(self.__dict__)
+  def __getitem__(self, key): return getattr(self, key)
+  def __setitem__(self, key, value): return setattr(self, key, value)
+  def __contains__(self, key): return key in self.__dict__
+  def __repr__(self): return self.__dict__.__repr__()
+
+# MODEL LOADING
+def load_model(symbols, hps, model) -> Synthesizer:
+  net_g = Synthesizer(len(symbols), hps.data.filter_length // 2 + 1, hps.train.segment_size // hps.data.hop_length, n_speakers = hps.data.n_speakers, **hps.model)
+  _ = load_checkpoint(fetch(model[1]), net_g, None)
+  return net_g
+def load_checkpoint(checkpoint_path, model: Synthesizer, optimizer=None, skip_list=[]):
+  assert Path(checkpoint_path).is_file()
+  start_time = time.time()
+  checkpoint_dict = torch_load(checkpoint_path)
+  iteration, learning_rate = checkpoint_dict['iteration'], checkpoint_dict['learning_rate']
+  if optimizer: optimizer.load_state_dict(checkpoint_dict['optimizer'])
+  saved_state_dict = checkpoint_dict['model']
+  weight_g, weight_v, parent = None, None, None
+  for key, v in saved_state_dict.items():
+    if any(layer in key for layer in skip_list): continue
+    try:
+      obj, skip = model, False
+      for k in key.split('.'):
+        if k.isnumeric(): obj = obj[int(k)]
+        elif isinstance(obj, dict): obj = obj[k]
+        else:
+          if isinstance(obj, (LayerNorm, nn.LayerNorm)) and k in ["gamma", "beta"]:
+            k = "weight" if k == "gamma" else "bias"
+          elif k in ["weight_g", "weight_v"]:
+            parent, skip = obj, True
+            if k == "weight_g": weight_g = v
+            else: weight_v = v
+          if not skip: obj = getattr(obj, k)
+      if weight_g is not None and weight_v is not None:
+        setattr(obj, "weight_g", weight_g.numpy())
+        setattr(obj, "weight_v", weight_v.numpy())
+        obj, v = getattr(parent, "weight"), weight_norm(weight_v, weight_g, 0)
+        weight_g, weight_v, parent, skip = None, None, None, False
+      if not skip and obj.shape == v.shape: obj.assign(v.to(obj.device))
+      elif not skip: logging.error(f"MISMATCH SHAPE IN {key}, {obj.shape} {v.shape}")
+    except Exception as e: raise e
+  logging.info(f"Loaded checkpoint '{checkpoint_path}' (iteration {iteration}) in {time.time() - start_time:.4f}s")
+  return model, optimizer, learning_rate, iteration
+
+# Used for cleaning input text and mapping to symbols
+class TextMapper: # Based on https://github.com/keithito/tacotron
+  def __init__(self, symbols, apply_cleaners=True):
+    self.apply_cleaners, self.symbols, self._inflect = apply_cleaners, symbols, None
+    self._symbol_to_id, _id_to_symbol = {s: i for i, s in enumerate(symbols)}, {i: s for i, s in enumerate(symbols)}
+    self._whitespace_re, self._abbreviations = re.compile(r'\s+'), [(re.compile('\\b%s\\.' % x[0], re.IGNORECASE), x[1]) for x in [('mrs', 'misess'), ('mr', 'mister'), ('dr', 'doctor'), ('st', 'saint'), ('co', 'company'), ('jr', 'junior'), ('maj', 'major'), ('gen', 'general'), ('drs', 'doctors'), ('rev', 'reverend'), ('lt', 'lieutenant'), ('hon', 'honorable'), ('sgt', 'sergeant'), ('capt', 'captain'), ('esq', 'esquire'), ('ltd', 'limited'), ('col', 'colonel'), ('ft', 'fort'), ]]
+    self.phonemizer = EspeakBackend(
+        language="en-us", punctuation_marks=Punctuation.default_marks(), preserve_punctuation=True, with_stress=True,
+    )
+  def text_to_sequence(self, text, cleaner_names):
+    if self.apply_cleaners:
+      for name in cleaner_names:
+        cleaner = getattr(self, name)
+        if not cleaner: raise ModuleNotFoundError('Unknown cleaner: %s' % name)
+        text = cleaner(text)
+    else: text = text.strip()
+    return [self._symbol_to_id[symbol] for symbol in text]
+  def get_text(self, text, add_blank=False, cleaners=('english_cleaners2',)):
+    text_norm = self.text_to_sequence(text, cleaners)
+    return Tensor(self.intersperse(text_norm, 0) if add_blank else text_norm, dtype=dtypes.int64)
+  def intersperse(self, lst, item):
+    (result := [item] * (len(lst) * 2 + 1))[1::2] = lst
+    return result
+  def phonemize(self, text, strip=True): return _phonemize(self.phonemizer, text, default_separator, strip, 1, False, False)
+  def filter_oov(self, text): return "".join(list(filter(lambda x: x in self._symbol_to_id, text)))
+  def base_english_cleaners(self, text): return self.collapse_whitespace(self.phonemize(self.expand_abbreviations(unidecode(text.lower()))))
+  def english_cleaners2(self, text): return self.base_english_cleaners(text)
+  def transliteration_cleaners(self, text): return self.collapse_whitespace(unidecode(text.lower()))
+  def cjke_cleaners(self, text): return re.sub(r'([^\.,!\?\-…~])$', r'\1.', re.sub(r'\s+$', '', self.english_to_ipa2(text).replace('ɑ', 'a').replace('ɔ', 'o').replace('ɛ', 'e').replace('ɪ', 'i').replace('ʊ', 'u')))
+  def cjke_cleaners2(self, text): return re.sub(r'([^\.,!\?\-…~])$', r'\1.', re.sub(r'\s+$', '', self.english_to_ipa2(text)))
+  def cjks_cleaners(self, text): return re.sub(r'([^\.,!\?\-…~])$', r'\1.', re.sub(r'\s+$', '', self.english_to_lazy_ipa(text)))
+  def english_to_ipa2(self, text):
+    _ipa_to_ipa2 = [(re.compile('%s' % x[0]), x[1]) for x in [ ('r', 'ɹ'), ('ʤ', 'dʒ'), ('ʧ', 'tʃ')]]
+    return reduce(lambda t, rx: re.sub(rx[0], rx[1], t), _ipa_to_ipa2, self.mark_dark_l(self.english_to_ipa(text))).replace('...', '…')
+  def mark_dark_l(self, text): return re.sub(r'l([^aeiouæɑɔəɛɪʊ ]*(?: |$))', lambda x: 'ɫ' + x.group(1), text)
+  def english_to_ipa(self, text):
+    import eng_to_ipa as ipa
+    return self.collapse_whitespace(ipa.convert(self.normalize_numbers(self.expand_abbreviations(unidecode(text).lower()))))
+  def english_to_lazy_ipa(self, text):
+    _lazy_ipa = [(re.compile('%s' % x[0]), x[1]) for x in [('r', 'ɹ'), ('æ', 'e'), ('ɑ', 'a'), ('ɔ', 'o'), ('ð', 'z'), ('θ', 's'), ('ɛ', 'e'), ('ɪ', 'i'), ('ʊ', 'u'), ('ʒ', 'ʥ'), ('ʤ', 'ʥ'), ('ˈ', '↓')]]
+    return reduce(lambda t, rx: re.sub(rx[0], rx[1], t), _lazy_ipa, self.english_to_ipa(text))
+  def expand_abbreviations(self, text): return reduce(lambda t, abbr: re.sub(abbr[0], abbr[1], t), self._abbreviations, text)
+  def collapse_whitespace(self, text): return re.sub(self._whitespace_re, ' ', text)
+  def normalize_numbers(self, text):
+    import inflect
+    self._inflect = inflect.engine()
+    text = re.sub(re.compile(r'([0-9][0-9\,]+[0-9])'), self._remove_commas, text)
+    text = re.sub(re.compile(r'£([0-9\,]*[0-9]+)'), r'\1 pounds', text)
+    text = re.sub(re.compile(r'\$([0-9\.\,]*[0-9]+)'), self._expand_dollars, text)
+    text = re.sub(re.compile(r'([0-9]+\.[0-9]+)'), self._expand_decimal_point, text)
+    text = re.sub(re.compile(r'[0-9]+(st|nd|rd|th)'), self._expand_ordinal, text)
+    text = re.sub(re.compile(r'[0-9]+'), self._expand_number, text)
+    return text
+  def _remove_commas(self, m): return m.group(1).replace(',', '') # george won't like this
+  def _expand_dollars(self, m):
+    match = m.group(1)
+    parts = match.split('.')
+    if len(parts) > 2: return match + ' dollars'  # Unexpected format
+    dollars, cents = int(parts[0]) if parts[0] else 0, int(parts[1]) if len(parts) > 1 and parts[1] else 0
+    if dollars and cents: return '%s %s, %s %s' % (dollars, 'dollar' if dollars == 1 else 'dollars', cents, 'cent' if cents == 1 else 'cents')
+    if dollars: return '%s %s' % (dollars, 'dollar' if dollars == 1 else 'dollars')
+    if cents: return '%s %s' % (cents, 'cent' if cents == 1 else 'cents')
+    return 'zero dollars'
+  def _expand_decimal_point(self, m): return m.group(1).replace('.', ' point ')
+  def _expand_ordinal(self, m): return self._inflect.number_to_words(m.group(0))
+  def _expand_number(self, _inflect, m):
+    num = int(m.group(0))
+    if 1000 < num < 3000:
+      if num == 2000: return 'two thousand'
+      if 2000 < num < 2010: return 'two thousand ' + self._inflect.number_to_words(num % 100)
+      if num % 100 == 0: return self._inflect.number_to_words(num // 100) + ' hundred'
+      return _inflect.number_to_words(num, andword='', zero='oh', group=2).replace(', ', ' ')
+    return self._inflect.number_to_words(num, andword='')
+
+#########################################################################################
+# PAPER: https://arxiv.org/abs/2106.06103
+# CODE: https://github.com/jaywalnut310/vits/tree/main
+#########################################################################################
+# INSTALLATION: this is based on default config, dependencies are for preprocessing.
+# vctk, ljs                      | pip3 install unidecode phonemizer          | phonemizer requires [eSpeak](https://espeak.sourceforge.net) backend to be installed on your system
+# mmts-tts                       | pip3 install unidecode                     |
+# uma_trilingual, cjks, voistock | pip3 install unidecode inflect eng_to_ipa  |
+#########################################################################################
+# Some good speakers to try out, there may be much better ones, I only tried out a few:
+# male vctk 1  | --model_to_use vctk --speaker_id 2
+# male vctk 2  | --model_to_use vctk --speaker_id 6
+# anime lady 1 | --model_to_use uma_trilingual --speaker_id 36
+# anime lady 2 | --model_to_use uma_trilingual --speaker_id 121
+#########################################################################################
+VITS_PATH = Path(__file__).parents[1] / "weights/VITS/"
+MODELS = { # config_url, weights_url
+  "ljs": ("https://raw.githubusercontent.com/jaywalnut310/vits/main/configs/ljs_base.json", "https://drive.google.com/uc?export=download&id=1q86w74Ygw2hNzYP9cWkeClGT5X25PvBT&confirm=t"),
+  "vctk": ("https://huggingface.co/csukuangfj/vits-vctk/resolve/main/vctk_base.json", "https://huggingface.co/csukuangfj/vits-vctk/resolve/main/pretrained_vctk.pth"),
+  "mmts-tts": ("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/config.json", "https://huggingface.co/facebook/mms-tts/resolve/main/full_models/eng/G_100000.pth"),
+  "uma_trilingual": ("https://huggingface.co/spaces/Plachta/VITS-Umamusume-voice-synthesizer/raw/main/configs/uma_trilingual.json", "https://huggingface.co/spaces/Plachta/VITS-Umamusume-voice-synthesizer/resolve/main/pretrained_models/G_trilingual.pth"),
+  "cjks": ("https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/14/config.json", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/14/model.pth"),
+  "voistock": ("https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/15/config.json", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/15/model.pth"),
+}
+Y_LENGTH_ESTIMATE_SCALARS = {"ljs": 2.8, "vctk": 1.74, "mmts-tts": 1.9, "uma_trilingual": 2.3, "cjks": 3.3, "voistock": 3.1}
+if __name__ == '__main__':
+  logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--model_to_use", default="vctk", help="Specify the model to use. Default is 'vctk'.")
+  parser.add_argument("--speaker_id", type=int, default=6, help="Specify the speaker ID. Default is 6.")
+  parser.add_argument("--out_path", default=None, help="Specify the full output path. Overrides the --out_dir and --name parameter.")
+  parser.add_argument("--out_dir", default=str(Path(__file__).parents[1] / "temp"), help="Specify the output path.")
+  parser.add_argument("--base_name", default="test", help="Specify the base of the output file name. Default is 'test'.")
+  parser.add_argument("--text_to_synthesize", default="""Hello person. If the code you are contributing isn't some of the highest quality code you've written in your life, either put in the effort to make it great, or don't bother.""", help="Specify the text to synthesize. Default is a greeting message.")
+  parser.add_argument("--noise_scale", type=float, default=0.667, help="Specify the noise scale. Default is 0.667.")
+  parser.add_argument("--noise_scale_w", type=float, default=0.8, help="Specify the noise scale w. Default is 0.8.")
+  parser.add_argument("--length_scale", type=float, default=1, help="Specify the length scale. Default is 1.")
+  parser.add_argument("--seed", type=int, default=1337, help="Specify the seed (set to None if no seed). Default is 1337.")
+  parser.add_argument("--num_channels", type=int, default=1, help="Specify the number of audio output channels. Default is 1.")
+  parser.add_argument("--sample_width", type=int, default=2, help="Specify the number of bytes per sample, adjust if necessary. Default is 2.")
+  parser.add_argument("--emotion_path", type=str, default=None, help="Specify the path to emotion reference.")
+  parser.add_argument("--estimate_max_y_length", type=str, default=False, help="If true, overestimate the output length and then trim it to the correct length, to prevent premature realization, much more performant for larger inputs, for smaller inputs not so much. Default is False.")
+  args = parser.parse_args()
+
+  model_config = MODELS[args.model_to_use]
+
+  # Load the hyperparameters from the config file.
+  hps = get_hparams_from_file(fetch(model_config[0]))
+
+  # If model has multiple speakers, validate speaker id and retrieve name if available.
+  model_has_multiple_speakers = hps.data.n_speakers > 0
+  if model_has_multiple_speakers:
+    logging.info(f"Model has {hps.data.n_speakers} speakers")
+    if args.speaker_id >= hps.data.n_speakers: raise ValueError(f"Speaker ID {args.speaker_id} is invalid for this model.")
+    speaker_name = "?"
+    if hps.__contains__("speakers"): # maps speaker ids to names
+      speakers = hps.speakers
+      if isinstance(speakers, List): speakers = {speaker: i for i, speaker in enumerate(speakers)}
+      speaker_name = next((key for key, value in speakers.items() if value == args.speaker_id), None)
+    logging.info(f"You selected speaker {args.speaker_id} (name: {speaker_name})")
+
+  # Load emotions if any. TODO: find an english model with emotions, this is untested atm.
+  emotion_embedding = None
+  if args.emotion_path is not None:
+    if args.emotion_path.endswith(".npy"): emotion_embedding = Tensor(np.load(args.emotion_path), dtype=dtypes.int64).unsqueeze(0)
+    else: raise ValueError("Emotion path must be a .npy file.")
+
+  # Load symbols, instantiate TextMapper and clean the text.
+  if hps.__contains__("symbols"): symbols = hps.symbols
+  elif args.model_to_use == "mmts-tts": symbols = [x.replace("\n", "") for x in fetch("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/vocab.txt").open(encoding="utf-8").readlines()]
+  else: symbols = ['_'] + list(';:,.!?¡¿—…"«»“” ') + list('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz') + list("ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ")
+  text_mapper = TextMapper(apply_cleaners=True, symbols=symbols)
+
+  # Load the model.
+  if args.seed is not None:
+    Tensor.manual_seed(args.seed)
+    np.random.seed(args.seed)
+  net_g = load_model(text_mapper.symbols, hps, model_config)
+  logging.debug(f"Loaded model with hps: {hps}")
+
+  # Convert the input text to a tensor.
+  text_to_synthesize = args.text_to_synthesize
+  if args.model_to_use == "mmts-tts": text_to_synthesize = text_mapper.filter_oov(text_to_synthesize.lower())
+  stn_tst = text_mapper.get_text(text_to_synthesize, hps.data.add_blank, hps.data.text_cleaners)
+  logging.debug(f"Converted input text to tensor \"{text_to_synthesize}\" -> Tensor({stn_tst.shape}): {stn_tst.numpy()}")
+  x_tst, x_tst_lengths = stn_tst.unsqueeze(0), Tensor([stn_tst.shape[0]], dtype=dtypes.int64)
+  sid = Tensor([args.speaker_id], dtype=dtypes.int64) if model_has_multiple_speakers else None
+
+  # Perform inference.
+  start_time = time.time()
+  audio_tensor = net_g.infer(x_tst, x_tst_lengths, sid, args.noise_scale, args.length_scale, args.noise_scale_w, emotion_embedding=emotion_embedding,
+                             max_y_length_estimate_scale=Y_LENGTH_ESTIMATE_SCALARS[args.model_to_use] if args.estimate_max_y_length else None)[0, 0].realize()
+  logging.info(f"Inference took {(time.time() - start_time):.2f}s")
+
+  # Save the audio output.
+  audio_data = (np.clip(audio_tensor.numpy(), -1.0, 1.0) * 32767).astype(np.int16)
+  out_path = Path(args.out_path or Path(args.out_dir)/f"{args.model_to_use}{f'_sid_{args.speaker_id}' if model_has_multiple_speakers else ''}_{args.base_name}.wav")
+  out_path.parent.mkdir(parents=True, exist_ok=True)
+  with wave.open(str(out_path), 'wb') as wav_file:
+    wav_file.setnchannels(args.num_channels)
+    wav_file.setsampwidth(args.sample_width)
+    wav_file.setframerate(hps.data.sampling_rate)
+    wav_file.setnframes(len(audio_data))
+    wav_file.writeframes(audio_data.tobytes())
+  logging.info(f"Saved audio output to {out_path}")
@@ -3,7 +3,7 @@
 import sys, base64, multiprocessing, itertools, collections
 from typing import Optional, Union, Literal, List

-from tinygrad import Tensor, TinyJit, Variable, nn, dtypes
+from tinygrad import Tensor, TinyJit, Variable, nn
 from tinygrad.nn.state import torch_load, load_state_dict
 from tinygrad.helpers import getenv, fetch

@@ -244,16 +244,15 @@ def transcribe_waveform(model: Whisper, enc, waveforms, truncate=False):

  log_spec = prep_audio(waveforms, model.batch_size, truncate)
  nsample = model.decoder.max_tokens_to_sample
-  nctx = model.decoder.max_self_attn_cache_len

  def inferloop(ctx: Union[np.ndarray, List[np.ndarray]], encoded_audio):
    pos, next_tokens = 0, ctx
-    for i in range(nsample):
-      next_tokens = model.decoder(Tensor(next_tokens, dtype=dtypes.int32), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
+    for i in range((nsample-len(start_tokens))*2):
+      next_tokens = model.decoder(Tensor(next_tokens), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
      next_tokens[ctx[:, -1] == eot] = eot
      ctx = np.concatenate((ctx, next_tokens), axis=1)
      pos = ctx.shape[-1] - 1
-      if (next_tokens == eot).all() or pos == nctx: break
+      if (next_tokens == eot).all(): break
    return ctx

  def gettexttoks(line): return [tok for tok in line if tok < eot or tok > enc._special_tokens["<|notimestamps|>"]][-nsample+len(start_tokens):]
@@ -26,13 +26,11 @@ def color_temp(temp):
 def color_voltage(voltage): return colored(f"{voltage/1000:>5.3f}V", "cyan")

 def draw_bar(percentage, width=40, fill='|', empty=' ', opt_text='', color='cyan'):
-  percentage = 0.0 if percentage != percentage else percentage  # NaN guard
-  percentage = max(0.0, min(1.0, float(percentage)))
  filled_width = int(width * percentage)
  if not opt_text: opt_text = f'{percentage*100:.1f}%'

  bar = fill * filled_width + empty * (width - filled_width)
-  if opt_text and len(opt_text) <= len(bar): bar = (bar[:-len(opt_text)] + opt_text)
+  bar = (bar[:-len(opt_text)] + opt_text) if opt_text else bar
  bar = colored(bar[:filled_width], color) + bar[filled_width:]
  return f'[{bar}]'

@@ -90,7 +88,6 @@ class SMICtx:
    self.opened_pci_resources = {}
    self.prev_lines_cnt = 0
    self.prev_terminal_width = 0
-    self.prev_terminal_height = 0

    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
    lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
@@ -98,20 +95,6 @@ class SMICtx:
    for k,v in self.lspci.items():
      for part in remove_parts: self.lspci[k] = self.lspci[k].replace(part, "").strip().rstrip()

-  def _smuq10_round(self, v:int) -> int:
-    v = int(v)
-    return (v + 512) >> 10  # SMUQ10_ROUND
-
-  def _fmt_kb(self, kb:int) -> str:
-    kb = int(kb)
-    if kb < 1024: return f"{kb}KB"
-    mb = kb / 1024.0
-    if mb < 1024: return f"{mb:.1f}MB"
-    gb = mb / 1024.0
-    if gb < 1024: return f"{gb:.2f}GB"
-    tb = gb / 1024.0
-    return f"{tb:.2f}TB"
-
  def _open_am_device(self, pcibus):
    if pcibus not in self.opened_pci_resources:
      bar_fds = {bar: os.open(f"/sys/bus/pci/devices/{pcibus}/resource{bar}", os.O_RDWR | os.O_SYNC) for bar in [0, 2, 5]}
@@ -133,7 +116,6 @@ class SMICtx:
  def rescan_devs(self):
    pattern = os.path.join('/tmp', 'am_*.lock')
    for d in [f[8:-5] for f in glob.glob(pattern)]:
-      if d.startswith("usb"): continue
      if d not in self.opened_pcidevs:
        self._open_am_device(d)

@@ -149,53 +131,21 @@ class SMICtx:
        os.system('clear')
        if DEBUG >= 2: print(f"Removed AM device {d.pcibus}")

-  def collect(self):
-    tables = {}
-    for dev in self.devs:
-      match dev.ip_ver[am.MP1_HWIP]:
-        case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableX_t
-        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTableV2_t
-        case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
-      tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
-    return tables
+  def collect(self): return {d: d.smu.read_metrics() if d.pci_state == "D0" else None for d in self.devs}

-  def _pick_nonzero_avg(self, vals) -> int:
-    xs = [x for x in vals if x > 0]
-    return int(sum(xs) / len(xs)) if xs else 0
-
-  def get_gfx_activity(self, dev, metrics):
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return max(0, min(100, self._smuq10_round(metrics.SocketGfxBusy)))
-      case _: return metrics.SmuMetrics.AverageGfxActivity
-
-  def get_mem_activity(self, dev, metrics):
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
-      case _: return metrics.SmuMetrics.AverageUclkActivity
+  def get_gfx_activity(self, dev, metrics): return metrics.SmuMetrics.AverageGfxActivity
+  def get_mem_activity(self, dev, metrics): return metrics.SmuMetrics.AverageUclkActivity

  def get_temps(self, dev, metrics, compact=False):
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6):
-        temps = {
-          "Hotspot": self._smuq10_round(metrics.MaxSocketTemperature),
-          "HBM": self._smuq10_round(metrics.MaxHbmTemperature),
-          "VR": self._smuq10_round(metrics.MaxVrTemperature),
-        }
-        if compact: return {k: temps[k] for k in ("Hotspot", "HBM") if temps.get(k, 0) != 0}
-        return {k: v for k, v in temps.items() if v != 0}
-      case _:
-        temps_keys = [(k, name) for k, name in dev.smu.smu_mod.TEMP_e.items()
-                      if k < dev.smu.smu_mod.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
-        if compact: temps_keys = [(k, name) for k, name in temps_keys if k in (dev.smu.smu_mod.TEMP_HOTSPOT, dev.smu.smu_mod.TEMP_MEM)]
-        return {name: metrics.SmuMetrics.AvgTemperature[k] for k, name in temps_keys}
+    temps_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_TEMP_e__enumvalues.items()
+                  if k < dev.smu.smu_mod.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
+    if compact: temps_keys = [(k, name) for k, name in temps_keys if k in (dev.smu.smu_mod.TEMP_HOTSPOT, dev.smu.smu_mod.TEMP_MEM)]
+    return {name: metrics.SmuMetrics.AvgTemperature[k] for k, name in temps_keys}

  def get_voltage(self, dev, metrics, compact=False):
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return {}
-      case _:
-        voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.SVI_PLANE_e.items()
+    voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_SVI_PLANE_e__enumvalues.items()
                        if k < dev.smu.smu_mod.SVI_PLANE_COUNT and metrics.SmuMetrics.AvgVoltage[k] != 0]
-        return {name: metrics.SmuMetrics.AvgVoltage[k] for k, name in voltage_keys}
+    return {name: metrics.SmuMetrics.AvgVoltage[k] for k, name in voltage_keys}

  def get_busy_threshold(self, dev):
    match dev.ip_ver[am.MP1_HWIP]:
@@ -203,40 +153,22 @@ class SMICtx:
      case _: return 15

  def get_gfx_freq(self, dev, metrics):
-    if metrics is None: return 0
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.GfxclkFrequency[0])
-      case _:
-        return metrics.SmuMetrics.AverageGfxclkFrequencyPostDs if self.get_gfx_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
-               metrics.SmuMetrics.AverageGfxclkFrequencyPreDs
+    return metrics.SmuMetrics.AverageGfxclkFrequencyPostDs if self.get_gfx_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+          metrics.SmuMetrics.AverageGfxclkFrequencyPreDs

  def get_mem_freq(self, dev, metrics):
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.UclkFrequency)
-      case _:
-        return metrics.SmuMetrics.AverageMemclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
-               metrics.SmuMetrics.AverageMemclkFrequencyPreDs
+    return metrics.SmuMetrics.AverageMemclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+           metrics.SmuMetrics.AverageMemclkFrequencyPreDs

  def get_fckl_freq(self, dev, metrics):
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.FclkFrequency)
-      case _:
-        return metrics.SmuMetrics.AverageFclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
-               metrics.SmuMetrics.AverageFclkFrequencyPreDs
+    return metrics.SmuMetrics.AverageFclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+           metrics.SmuMetrics.AverageFclkFrequencyPreDs

-  def get_fan_rpm_pwm(self, dev, metrics):
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return None, None
-      case _: return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm
+  def get_fan_rpm_pwm(self, dev, metrics): return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm

-  def get_power(self, dev, metrics):
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
-      case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
+  def get_power(self, dev, metrics): return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX

  def get_mem_usage(self, dev):
-    return 0
-
    usage = 0
    pt_stack = [dev.mm.root_page_table]
    while len(pt_stack) > 0:
@@ -245,7 +177,7 @@ class SMICtx:
        entry = pt.entries[i]

        if (entry & am.AMDGPU_PTE_VALID) == 0: continue
-        if pt.lv!=am.AMDGPU_VM_PTB and not dev.gmc.is_pte_huge_page(pt.lv, entry):
+        if pt.lv!=am.AMDGPU_VM_PTB and not dev.gmc.is_pte_huge_page(entry):
          pt_stack.append(AMPageTableEntry(dev, entry & 0x0000FFFFFFFFF000, lv=pt.lv+1))
          continue
        if (entry & am.AMDGPU_PTE_SYSTEM) != 0: continue
@@ -287,28 +219,23 @@ class SMICtx:
      temps_table_compact = ["Temps (°C):" + '/'.join([f"{color_temp(val)} {name}" for name, val in temps_data_compact.items()])]

      fan_rpm, fan_pwm = self.get_fan_rpm_pwm(dev, metrics)
-      power_table = ["=== Power ==="]
-      power_table += ["Fan: N/A"] if fan_rpm is None or fan_pwm is None else [f"Fan Speed: {fan_rpm} RPM", f"Fan Power: {fan_pwm}%"]
+      power_table = ["=== Power ==="] + [f"Fan Speed: {fan_rpm} RPM"] + [f"Fan Power: {fan_pwm}%"]

      total_power, max_power = self.get_power(dev, metrics)
-      if max_power > 0:
-        power_line = [f"Power: " + draw_bar(total_power / max_power, 16, opt_text=f"{total_power}/{max_power}W")]
-        power_line_compact = [f"Power:       " + draw_bar(total_power / max_power, activity_line_width, opt_text=f"{total_power}/{max_power}W")]
-      else:
-        power_line = ["Power: N/A"]
-        power_line_compact = ["Power: N/A"]
+      power_line = [f"Power: " + draw_bar(total_power / max_power, 16, opt_text=f"{total_power}/{max_power}W")]
+      power_line_compact = [f"Power:       " + draw_bar(total_power / max_power, activity_line_width, opt_text=f"{total_power}/{max_power}W")]

      voltage_data = self.get_voltage(dev, metrics)
-      voltage_table = None if not voltage_data else (["=== Voltages ==="] + [f"{name:<20}: {color_voltage(voltage)}" for name, voltage in voltage_data.items()])
+      voltage_table = ["=== Voltages ==="] + [f"{name:<20}: {color_voltage(voltage)}" for name, voltage in voltage_data.items()]

      gfx_freq = self.get_gfx_freq(dev, metrics)
      mclk_freq = self.get_mem_freq(dev, metrics)
      fclk_freq = self.get_fckl_freq(dev, metrics)
+
      frequency_table = ["=== Frequencies ===", f"GFXCLK: {gfx_freq:>4} MHz", f"FCLK  : {fclk_freq:>4} MHz", f"MCLK  : {mclk_freq:>4} MHz"]

      if self.prev_terminal_width >= 231:
-        power_table += power_line
-        if voltage_table is not None: power_table += [""] + voltage_table
+        power_table += power_line + [""] + voltage_table
        activity_line += [""]
      elif self.prev_terminal_width >= 171:
        power_table += power_line + [""] + frequency_table
@@ -380,5 +307,4 @@ if __name__ == "__main__":
      smi_ctx.draw(args.list)
      if args.list: break
      time.sleep(1)
-  except KeyboardInterrupt:
-    print("Exiting...")
+  except KeyboardInterrupt: print("Exiting...")
@@ -1,14 +0,0 @@
-#!/usr/bin/env python3
-
-from tinygrad.helpers import Context
-from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
-from tinygrad.runtime.support.am.amdev import AMDev
-
-if __name__ == "__main__":
-  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
-  pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
-  amdevs = []
-  with Context(DEBUG=2):
-    for pcidev in pcidevs:
-      amdevs.append(AMDev(pcidev, reset_mode=True))
-    for amdev in amdevs: amdev.smu.mode1_reset()
@@ -1,65 +1,48 @@
 import re, ctypes, sys, importlib
-from tinygrad.helpers import getenv

 from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
-
-class GFXFake:
-  def __init__(self): self.xccs = 8
-
 class AMDFake(AMDev):
-  def __init__(self, pci_dev, dma_regions=None):
-    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
-    self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
+  def __init__(self, devfmt, vram, doorbell, mmio, dma_regions=None):
+    self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
    self._run_discovery()
    self._build_regs()

-    self.gfx = GFXFake()
-
 amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
 amdev.AMDev = AMDFake
+
 from tinygrad.runtime.ops_amd import PCIIface

-def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
-  register_map = register_names or {}
+def parse_amdgpu_logs(log_content, register_names=None):
+  register_map = register_names

+  final = ""
  def replace_register(match):
-    reg = match.group(1)
-    return f"Reading register {register_map.get(int(reg, 16), reg)}"
+    register = match.group(1)
+    return f"Reading register {register_map.get(int(register, base=16), register)}"

-  processed_log = re.sub(r'Reading register (0x[0-9a-fA-F]+)', replace_register, log_content)
+  pattern = r'Reading register (0x[0-9a-fA-F]+)'
+
+  processed_log = re.sub(pattern, replace_register, log_content)

  def replace_register_2(match):
-    reg = match.group(1)
-    return f"Writing register {register_map.get(int(reg, 16), reg)}"
-
-  processed_log = re.sub(r'Writing register (0x[0-9a-fA-F]+)', replace_register_2, processed_log)
-
-  # remove timing prefix
-  processed_log = re.sub(r'^\[\s*\d+(?:\.\d+)?\]\s*', '', processed_log, flags=re.MULTILINE)
-
-  # keep only xcc=0 lines (but keep lines with no xcc at all)
-  if only_xcc0:
-    kept = []
-    for line in processed_log.splitlines(True):
-      if "xcc=" not in line or re.search(r'\bxcc=0\b', line): kept.append(line)
-    processed_log = "".join(kept)
+    register = match.group(1)
+    return f"Writing register {register_map.get(int(register, base=16), register)}"

+  pattern = r'Writing register (0x[0-9a-fA-F]+)'
+  processed_log = re.sub(pattern, replace_register_2, processed_log)
  return processed_log

 def main():
-  only_xcc0 = bool(getenv("ONLY_XCC0", 0))
-
  reg_names = {}
  dev = PCIIface(None, 0)
  for x, y in dev.dev_impl.__dict__.items():
    if isinstance(y, AMRegister):
-      for xcc, addr in y.addr.items():
-        reg_names[addr] = f"{x}, xcc={xcc}"
+      for inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"

  with open(sys.argv[1], 'r') as f:
-    log_content = f.read()
+    log_content = log_content_them = f.read()

-  processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
+  processed_log = parse_amdgpu_logs(log_content, reg_names)

  with open(sys.argv[2], 'w') as f:
    f.write(processed_log)
@@ -68,4 +51,5 @@ if __name__ == '__main__':
  if len(sys.argv) != 3:
    print("Usage: <input_file_path> <output_file_path>")
    sys.exit(1)
-  main()
+
+  main()
@@ -1,31 +0,0 @@
-An integrated environment for AMD GPU assembly and emulation
-
-Test with `PYTHONPATH="." pytest -n12 extra/assembly/amd/`
-`AMD_LLVM=1 PYTHONPATH="." pytest -n12 extra/assembly/amd/`
-
-* pdf.py -- extract assembly format + instruction psuedocode from AMD PDF
-* dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
-* pcode.py -- psuedocode execution environment. psuedocode should be transformed as little as possible.
-* asm.py -- an asm/disasm function to transform to and from AMD assembly syntax
-* emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
-
-The code should be as readable and deduplicated as possible. asm and emu shouldn't be required for dsl.
-
-test_emu.py has a good set of instruction tests for the emulation, with USE_HW=1 it will compare to real hardware.
-Whenever an instruction is fixed, regression tests should be added here and confirmed with real hardware.
-
-test_llvm.py tests asm/disasm on the LLVM tests, confirming it behaves the same as LLVM.
-
-tinygrad's dtype tests should pass with and without LLVM. they run in about 12 seconds.
-
-`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
-`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
-
-The ops tests also pass, but they are very slow, so you should run them one at a time.
-
-`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_ops.py`
-`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_ops.py`
-
-When something is caught by main tinygrad tests, a local regression test should be added to `extra/assembly/amd/test`. While working with tinygrad, you can dump the assembly with `DEBUG=7`. These tests all pass on real hardware, so if a test is failing with `AMD=1 PYTHON_REMU=1 MOCKGPU=1` it's likely because an instruction is emulated incorrectly. You can test without `MOCKGPU=1` to test on real hardware, if it works on real hardware there's a bug in the emulator.
-
-Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~2000 lines. Count lines with `cloc --by-file extra/assembly/amd/*.py`
@@ -1,581 +0,0 @@
-# RDNA3 assembler and disassembler
-from __future__ import annotations
-import re
-from extra.assembly.amd.dsl import Inst, RawImm, Reg, SrcMod, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory
-from extra.assembly.amd.dsl import VCC_LO, VCC_HI, VCC, EXEC_LO, EXEC_HI, EXEC, SCC, M0, NULL, OFF
-from extra.assembly.amd.dsl import SPECIAL_GPRS, SPECIAL_PAIRS, FLOAT_DEC, FLOAT_ENC, decode_src
-from extra.assembly.amd.autogen.rdna3 import ins
-from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, VOPD, VINTERP, SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT, MUBUF, MTBUF, MIMG, EXP,
-  VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOPDOp, SOP1Op, SOPKOp, SOPPOp, SMEMOp, DSOp, MUBUFOp)
-
-def _matches_encoding(word: int, cls: type[Inst]) -> bool:
-  """Check if word matches the encoding pattern of an instruction class."""
-  if cls._encoding is None: return False
-  bf, val = cls._encoding
-  return ((word >> bf.lo) & bf.mask()) == val
-
-# Order matters: more specific encodings first, VOP2 last (it's a catch-all for bit31=0)
-_FORMATS_64 = [VOPD, VOP3P, VINTERP, VOP3, DS, FLAT, MUBUF, MTBUF, MIMG, SMEM, EXP]
-_FORMATS_32 = [SOP1, SOPC, SOPP, SOPK, VOPC, VOP1, SOP2, VOP2]  # SOP2/VOP2 are catch-alls
-
-def detect_format(data: bytes) -> type[Inst]:
-  """Detect instruction format from machine code bytes."""
-  assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
-  word = int.from_bytes(data[:4], 'little')
-  # Check 64-bit formats first (bits[31:30] == 0b11)
-  if (word >> 30) == 0b11:
-    for cls in _FORMATS_64:
-      if _matches_encoding(word, cls):
-        return VOP3SD if cls is VOP3 and ((word >> 16) & 0x3ff) in Inst._VOP3SD_OPS else cls
-    raise ValueError(f"unknown 64-bit format word={word:#010x}")
-  # 32-bit formats
-  for cls in _FORMATS_32:
-    if _matches_encoding(word, cls): return cls
-  raise ValueError(f"unknown 32-bit format word={word:#010x}")
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# CONSTANTS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-HWREG = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
-         6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
-         19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI', 22: 'HW_REG_XNACK_MASK',
-         23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
-HWREG_IDS = {v.lower(): k for k, v in HWREG.items()}
-MSG = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
-       131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA'}
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# HELPERS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{b}" if n == 1 else f"{p}[{b}:{b+n-1}]"
-def _sreg(b: int, n: int = 1) -> str: return _reg("s", b, n)
-def _vreg(b: int, n: int = 1) -> str: return _reg("v", b, n)
-def _ttmp(b: int, n: int = 1) -> str: return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
-def _sreg_or_ttmp(b: int, n: int = 1) -> str: return _ttmp(b, n) or _sreg(b, n)
-
-def _fmt_sdst(v: int, n: int = 1) -> str:
-  if v == 124: return "null"
-  if t := _ttmp(v, n): return t
-  if n > 1: return SPECIAL_PAIRS.get(v) or _sreg(v, n)
-  return SPECIAL_GPRS.get(v, f"s{v}")
-
-def _fmt_src(v: int, n: int = 1) -> str:
-  if n == 1: return decode_src(v)
-  if v >= 256: return _vreg(v - 256, n)
-  if v <= 105: return _sreg(v, n)
-  if n == 2 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
-  if t := _ttmp(v, n): return t
-  return decode_src(v)
-
-def _fmt_v16(v: int, base: int = 256, hi_thresh: int = 384) -> str:
-  return f"v{(v - base) & 0x7f}.{'h' if v >= hi_thresh else 'l'}"
-
-def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
-  return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
-
-def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
-def _omod(v: int) -> str: return {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(v, "")
-def _src16(inst, v: int) -> str: return _fmt_v16(v) if v >= 256 else inst.lit(v)  # format 16-bit src: vgpr.h/l or literal
-def _mods(*pairs) -> str: return " ".join(m for c, m in pairs if c)
-def _fmt_bits(label: str, val: int, count: int) -> str: return f"{label}:[{','.join(str((val >> i) & 1) for i in range(count))}]"
-
-def _vop3_src(inst, v: int, neg: int, abs_: int, hi: int, n: int, f16: bool, any_hi: bool) -> str:
-  """Format VOP3 source operand with modifiers."""
-  if n > 1: s = _fmt_src(v, n)
-  elif f16 and v >= 256: s = f"v{v - 256}.h" if hi else (f"v{v - 256}.l" if any_hi else inst.lit(v))
-  else: s = inst.lit(v)
-  if abs_: s = f"|{s}|"
-  return f"-{s}" if neg else s
-
-def _opsel_str(opsel: int, n: int, need: bool, is16_d: bool) -> str:
-  """Format op_sel modifier string."""
-  if not need: return ""
-  if is16_d and (opsel & 8): return f" op_sel:[1,1,1{',1' if n == 3 else ''}]"
-  if n == 3: return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{(opsel >> 3) & 1}]"
-  return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1}]"
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# DISASSEMBLER
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def _disasm_vop1(inst: VOP1) -> str:
-  name = inst.op_name.lower()
-  if inst.op in (VOP1Op.V_NOP, VOP1Op.V_PIPEFLUSH): return name
-  if inst.op == VOP1Op.V_READFIRSTLANE_B32: return f"v_readfirstlane_b32 {decode_src(inst.vdst)}, v{inst.src0 - 256 if inst.src0 >= 256 else inst.src0}"
-  # 16-bit dst: uses .h/.l suffix (determined by name pattern, not dtype - e.g. sat_pk_u8_i16 outputs 8-bit but uses 16-bit encoding)
-  parts = name.split('_')
-  is_16d = any(p in ('f16','i16','u16','b16') for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in ('f16','i16','u16','b16') and 'cvt' not in name)
-  dst = _vreg(inst.vdst, inst.dst_regs()) if inst.dst_regs() > 1 else _fmt_v16(inst.vdst, 0, 128) if is_16d else f"v{inst.vdst}"
-  src = _fmt_src(inst.src0, inst.src_regs(0)) if inst.src_regs(0) > 1 else _src16(inst, inst.src0) if inst.is_src_16(0) and 'sat_pk' not in name else inst.lit(inst.src0)
-  return f"{name}_e32 {dst}, {src}"
-
-def _disasm_vop2(inst: VOP2) -> str:
-  name = inst.op_name.lower()
-  suf = "" if inst.op == VOP2Op.V_DOT2ACC_F32_F16 else "_e32"
-  # fmaak: dst = src0 * vsrc1 + K, fmamk: dst = src0 * K + vsrc1
-  if inst.op in (VOP2Op.V_FMAAK_F32, VOP2Op.V_FMAAK_F16): return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}, 0x{inst._literal:x}"
-  if inst.op in (VOP2Op.V_FMAMK_F32, VOP2Op.V_FMAMK_F16): return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, 0x{inst._literal:x}, v{inst.vsrc1}"
-  if inst.is_16bit(): return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1, 0, 128)}"
-  return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}" + (", vcc_lo" if inst.op == VOP2Op.V_CNDMASK_B32 else "")
-
-def _disasm_vopc(inst: VOPC) -> str:
-  name = inst.op_name.lower()
-  s0 = _fmt_src(inst.src0, inst.src_regs(0)) if inst.src_regs(0) > 1 else _src16(inst, inst.src0) if inst.is_16bit() else inst.lit(inst.src0)
-  s1 = _vreg(inst.vsrc1, inst.src_regs(1)) if inst.src_regs(1) > 1 else _fmt_v16(inst.vsrc1, 0, 128) if inst.is_16bit() else f"v{inst.vsrc1}"
-  return f"{name}_e32 {s0}, {s1}" if inst.op.value >= 128 else f"{name}_e32 vcc_lo, {s0}, {s1}"
-
-NO_ARG_SOPP = {SOPPOp.S_ENDPGM, SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
-               SOPPOp.S_WAIT_IDLE, SOPPOp.S_ENDPGM_SAVED, SOPPOp.S_CODE_END, SOPPOp.S_ENDPGM_ORDERED_PS_DONE}
-
-def _disasm_sopp(inst: SOPP) -> str:
-  name = inst.op_name.lower()
-  if inst.op in NO_ARG_SOPP: return name
-  if inst.op == SOPPOp.S_WAITCNT:
-    vm, exp, lgkm = (inst.simm16 >> 10) & 0x3f, inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x3f
-    p = [f"vmcnt({vm})" if vm != 0x3f else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
-    return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
-  if inst.op == SOPPOp.S_DELAY_ALU:
-    deps, skips = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3'], ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
-    id0, skip, id1 = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 7) & 0xf
-    dep = lambda v: deps[v-1] if 0 < v <= len(deps) else str(v)
-    p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
-    return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
-  return f"{name} {inst.simm16}" if name.startswith(('s_cbranch', 's_branch')) else f"{name} 0x{inst.simm16:x}"
-
-def _disasm_smem(inst: SMEM) -> str:
-  name = inst.op_name.lower()
-  if inst.op in (SMEMOp.S_GL1_INV, SMEMOp.S_DCACHE_INV): return name
-  off_s = f"{decode_src(inst.soffset)} offset:0x{inst.offset:x}" if inst.offset and inst.soffset != 124 else f"0x{inst.offset:x}" if inst.offset else decode_src(inst.soffset)
-  sbase_idx, sbase_count = inst.sbase * 2, 4 if (8 <= inst.op.value <= 12 or name == 's_atc_probe_buffer') else 2
-  sbase_str = _fmt_src(sbase_idx, sbase_count) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
-  if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {inst.sdata}, {sbase_str}, {off_s}"
-  return f"{name} {_fmt_sdst(inst.sdata, inst.dst_regs())}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (inst.dlc, " dlc"))
-
-def _disasm_flat(inst: FLAT) -> str:
-  name = inst.op_name.lower()
-  seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
-  instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
-  off_val = inst.offset if seg == 'flat' else (inst.offset if inst.offset < 4096 else inst.offset - 8192)
-  w = inst.dst_regs() * (2 if 'cmpswap' in name else 1)
-  mods = f"{f' offset:{off_val}' if off_val else ''}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
-  # saddr
-  if seg == 'flat' or inst.saddr == 0x7F: saddr_s = ""
-  elif inst.saddr == 124: saddr_s = ", off"
-  elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr)}"
-  elif inst.saddr in SPECIAL_PAIRS: saddr_s = f", {SPECIAL_PAIRS[inst.saddr]}"
-  elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
-  else: saddr_s = f", {_sreg(inst.saddr, 2) if inst.saddr < 106 else decode_src(inst.saddr)}"
-  # addtid: no addr
-  if 'addtid' in name: return f"{instr} v{inst.data if 'store' in name else inst.vdst}{saddr_s}{mods}"
-  # addr width
-  addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, 1 if seg == 'scratch' or (inst.saddr not in (0x7F, 124)) else 2)
-  data_s, vdst_s = _vreg(inst.data, w), _vreg(inst.vdst, w // 2 if 'cmpswap' in name else w)
-  if 'atomic' in name:
-    return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if inst.glc else f"{instr} {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}"
-  if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
-  return f"{instr} {_vreg(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
-
-def _disasm_ds(inst: DS) -> str:
-  op, name = inst.op, inst.op_name.lower()
-  gds = " gds" if inst.gds else ""
-  off = f" offset:{inst.offset0 | (inst.offset1 << 8)}" if inst.offset0 or inst.offset1 else ""
-  off2 = f" offset0:{inst.offset0} offset1:{inst.offset1}" if inst.offset0 or inst.offset1 else ""
-  w = inst.dst_regs()
-  d0, d1, dst, addr = _vreg(inst.data0, w), _vreg(inst.data1, w), _vreg(inst.vdst, w), f"v{inst.addr}"
-
-  if op == DSOp.DS_NOP: return name
-  if op == DSOp.DS_BVH_STACK_RTN_B32: return f"{name} v{inst.vdst}, {addr}, v{inst.data0}, {_vreg(inst.data1, 4)}{off}{gds}"
-  if 'gws_sema' in name and op != DSOp.DS_GWS_SEMA_BR: return f"{name}{off}{gds}"
-  if 'gws_' in name: return f"{name} {addr}{off}{gds}"
-  if op in (DSOp.DS_CONSUME, DSOp.DS_APPEND): return f"{name} v{inst.vdst}{off}{gds}"
-  if 'gs_reg' in name: return f"{name} {_vreg(inst.vdst, 2)}, v{inst.data0}{off}{gds}"
-  if '2addr' in name:
-    if 'load' in name: return f"{name} {_vreg(inst.vdst, w*2)}, {addr}{off2}{gds}"
-    if 'store' in name and 'xchg' not in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
-    return f"{name} {_vreg(inst.vdst, w*2)}, {addr}, {d0}, {d1}{off2}{gds}"
-  if 'load' in name: return f"{name} v{inst.vdst}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
-  if 'store' in name and not _has(name, 'cmp', 'xchg'):
-    return f"{name} v{inst.data0}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
-  if 'swizzle' in name or op == DSOp.DS_ORDERED_COUNT: return f"{name} v{inst.vdst}, {addr}{off}{gds}"
-  if 'permute' in name: return f"{name} v{inst.vdst}, {addr}, v{inst.data0}{off}{gds}"
-  if 'condxchg' in name: return f"{name} {_vreg(inst.vdst, 2)}, {addr}, {_vreg(inst.data0, 2)}{off}{gds}"
-  if _has(name, 'cmpstore', 'mskor', 'wrap'):
-    return f"{name} {dst}, {addr}, {d0}, {d1}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}, {d1}{off}{gds}"
-  return f"{name} {dst}, {addr}, {d0}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}{off}{gds}"
-
-def _disasm_vop3(inst: VOP3) -> str:
-  op, name = inst.op, inst.op_name.lower()
-
-  # VOP3SD (shared encoding)
-  if isinstance(op, VOP3SDOp):
-    sdst = (inst.clmp << 7) | (inst.opsel << 3) | inst.abs
-    def src(v, neg, n): s = _fmt_src(v, n) if n > 1 else inst.lit(v); return f"-{s}" if neg else s
-    s0, s1, s2 = src(inst.src0, inst.neg & 1, inst.src_regs(0)), src(inst.src1, inst.neg & 2, inst.src_regs(1)), src(inst.src2, inst.neg & 4, inst.src_regs(2))
-    dst = _vreg(inst.vdst, inst.dst_regs()) if inst.dst_regs() > 1 else f"v{inst.vdst}"
-    srcs = f"{s0}, {s1}, {s2}" if inst.num_srcs() == 3 else f"{s0}, {s1}"
-    return f"{name} {dst}, {_fmt_sdst(sdst, 1)}, {srcs}" + _omod(inst.omod)
-
-  # Detect 16-bit operand sizes (for .h/.l suffix handling)
-  is16_d = is16_s = is16_s2 = False
-  if 'cvt_pk' in name: is16_s = name.endswith('16')
-  elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', name):
-    is16_d, is16_s = _has(m.group(1), 'f16','i16','u16','b16'), _has(m.group(2), 'f16','i16','u16','b16')
-    is16_s2 = is16_s
-  elif re.match(r'v_mad_[iu]32_[iu]16', name): is16_s = True
-  elif 'pack_b32' in name: is16_s = is16_s2 = True
-  else: is16_d = is16_s = is16_s2 = inst.is_16bit()
-
-  any_hi = inst.opsel != 0
-  s0 = _vop3_src(inst, inst.src0, inst.neg&1, inst.abs&1, inst.opsel&1, inst.src_regs(0), is16_s, any_hi)
-  s1 = _vop3_src(inst, inst.src1, inst.neg&2, inst.abs&2, inst.opsel&2, inst.src_regs(1), is16_s, any_hi)
-  s2 = _vop3_src(inst, inst.src2, inst.neg&4, inst.abs&4, inst.opsel&4, inst.src_regs(2), is16_s2, any_hi)
-
-  # Destination
-  dn = inst.dst_regs()
-  if op == VOP3Op.V_READLANE_B32: dst = _fmt_sdst(inst.vdst, 1)
-  elif dn > 1: dst = _vreg(inst.vdst, dn)
-  elif is16_d: dst = f"v{inst.vdst}.h" if (inst.opsel & 8) else f"v{inst.vdst}.l" if any_hi else f"v{inst.vdst}"
-  else: dst = f"v{inst.vdst}"
-
-  cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
-  nonvgpr_opsel = (inst.src0 < 256 and (inst.opsel & 1)) or (inst.src1 < 256 and (inst.opsel & 2)) or (inst.src2 < 256 and (inst.opsel & 4))
-  need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)
-
-  if inst.op < 256:  # VOPC
-    return f"{name}_e64 {s0}, {s1}" if name.startswith('v_cmpx') else f"{name}_e64 {_fmt_sdst(inst.vdst, 1)}, {s0}, {s1}"
-  if inst.op < 384:  # VOP2
-    n = inst.num_srcs()
-    os = _opsel_str(inst.opsel, n, need_opsel, is16_d)
-    return f"{name}_e64 {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name}_e64 {dst}, {s0}, {s1}{os}{cl}{om}"
-  if inst.op < 512:  # VOP1
-    return f"{name}_e64" if op in (VOP3Op.V_NOP, VOP3Op.V_PIPEFLUSH) else f"{name}_e64 {dst}, {s0}{_opsel_str(inst.opsel, 1, need_opsel, is16_d)}{cl}{om}"
-  # Native VOP3
-  n = inst.num_srcs()
-  os = _opsel_str(inst.opsel, n, need_opsel, is16_d)
-  return f"{name} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{os}{cl}{om}"
-
-def _disasm_vop3sd(inst: VOP3SD) -> str:
-  name = inst.op_name.lower()
-  def src(v, neg, n): s = _fmt_src(v, n) if n > 1 else inst.lit(v); return f"-{s}" if neg else s
-  s0, s1, s2 = src(inst.src0, inst.neg & 1, inst.src_regs(0)), src(inst.src1, inst.neg & 2, inst.src_regs(1)), src(inst.src2, inst.neg & 4, inst.src_regs(2))
-  dst = _vreg(inst.vdst, inst.dst_regs()) if inst.dst_regs() > 1 else f"v{inst.vdst}"
-  srcs = f"{s0}, {s1}, {s2}" if inst.num_srcs() == 3 else f"{s0}, {s1}"
-  suffix = "_e64" if name.startswith('v_') and 'co_' in name else ""
-  return f"{name}{suffix} {dst}, {_fmt_sdst(inst.sdst, 1)}, {srcs}{' clamp' if inst.clmp else ''}{_omod(inst.omod)}"
-
-def _disasm_vopd(inst: VOPD) -> str:
-  lit = inst._literal or inst.literal
-  vdst_y, nx, ny = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), VOPDOp(inst.opx).name.lower(), VOPDOp(inst.opy).name.lower()
-  def half(n, vd, s0, vs1): return f"{n} v{vd}, {inst.lit(s0)}{f', 0x{lit:x}' if lit and _has(n, 'fmaak', 'fmamk') else ''}" if 'mov' in n else f"{n} v{vd}, {inst.lit(s0)}, v{vs1}{f', 0x{lit:x}' if lit and _has(n, 'fmaak', 'fmamk') else ''}"
-  return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
-
-def _disasm_vop3p(inst: VOP3P) -> str:
-  name = inst.op_name.lower()
-  is_wmma, n, is_fma_mix = 'wmma' in name, inst.num_srcs(), 'fma_mix' in name
-  if is_wmma:
-    sc = 2 if 'iu4' in name else 4 if 'iu8' in name else 8
-    src0, src1, src2, dst = _fmt_src(inst.src0, sc), _fmt_src(inst.src1, sc), _fmt_src(inst.src2, 8), _vreg(inst.vdst, 8)
-  else: src0, src1, src2, dst = _fmt_src(inst.src0, 1), _fmt_src(inst.src1, 1), _fmt_src(inst.src2, 1), f"v{inst.vdst}"
-  opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
-  if is_fma_mix:
-    def m(s, neg, abs_): return f"-{f'|{s}|' if abs_ else s}" if neg else (f"|{s}|" if abs_ else s)
-    src0, src1, src2 = m(src0, inst.neg & 1, inst.neg_hi & 1), m(src1, inst.neg & 2, inst.neg_hi & 2), m(src2, inst.neg & 4, inst.neg_hi & 4)
-    mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if inst.clmp else [])
-  else:
-    mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != (7 if n == 3 else 3) else []) + \
-           ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
-  return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
-
-def _disasm_buf(inst: MUBUF | MTBUF) -> str:
-  name = inst.op_name.lower()
-  if inst.op in (MUBUFOp.BUFFER_GL0_INV, MUBUFOp.BUFFER_GL1_INV): return name
-  w = (2 if _has(name, 'xyz', 'xyzw') else 1) if 'd16' in name else \
-      ((2 if _has(name, 'b64', 'u64', 'i64') else 1) * (2 if 'cmpswap' in name else 1)) if 'atomic' in name else \
-      {'b32':1,'b64':2,'b96':3,'b128':4,'b16':1,'x':1,'xy':2,'xyz':3,'xyzw':4}.get(name.split('_')[-1], 1)
-  if inst.tfe: w += 1
-  vaddr = _vreg(inst.vaddr, 2) if inst.offen and inst.idxen else f"v{inst.vaddr}" if inst.offen or inst.idxen else "off"
-  srsrc = _sreg_or_ttmp(inst.srsrc*4, 4)
-  mods = ([f"format:{inst.format}"] if isinstance(inst, MTBUF) else []) + [m for c, m in [(inst.idxen,"idxen"),(inst.offen,"offen"),(inst.offset,f"offset:{inst.offset}"),(inst.glc,"glc"),(inst.dlc,"dlc"),(inst.slc,"slc"),(inst.tfe,"tfe")] if c]
-  return f"{name} {_vreg(inst.vdata, w)}, {vaddr}, {srsrc}, {decode_src(inst.soffset)}{' ' + ' '.join(mods) if mods else ''}"
-
-def _mimg_vaddr_width(name: str, dim: int, a16: bool) -> int:
-  """Calculate vaddr register count for MIMG sample/gather operations."""
-  #                    1d,2d,3d,cube,1d_arr,2d_arr,2d_msaa,2d_msaa_arr
-  base =              [1, 2, 3, 3,   2,     3,     3,      4][dim]  # address coords
-  grad =              [1, 2, 3, 2,   1,     2,     2,      2][dim]  # gradient coords (for derivatives)
-  if 'get_resinfo' in name: return 1  # only mip level
-  packed, unpacked = 0, 0
-  if '_mip' in name: packed += 1
-  elif 'sample' in name or 'gather' in name:
-    if '_o' in name: unpacked += 1                                              # offset
-    if re.search(r'_c(_|$)', name): unpacked += 1                               # compare (not _cl)
-    if '_d' in name: unpacked += (grad + 1) & ~1 if '_g16' in name else grad*2  # derivatives
-    if '_b' in name: unpacked += 1                                              # bias
-    if '_l' in name and '_cl' not in name and '_lz' not in name: packed += 1    # LOD
-    if '_cl' in name: packed += 1                                               # clamp
-  return (base + packed + 1) // 2 + unpacked if a16 else base + packed + unpacked
-
-def _disasm_mimg(inst: MIMG) -> str:
-  name = inst.op_name.lower()
-  srsrc_base = inst.srsrc * 4
-  srsrc_str = _sreg_or_ttmp(srsrc_base, 8)
-  # BVH intersect ray: special case with 4 SGPR srsrc
-  if 'bvh' in name:
-    vaddr = (9 if '64' in name else 8) if inst.a16 else (12 if '64' in name else 11)
-    return f"{name} {_vreg(inst.vdata, 4)}, {_vreg(inst.vaddr, vaddr)}, {_sreg_or_ttmp(srsrc_base, 4)}{' a16' if inst.a16 else ''}"
-  # vdata width from dmask (gather4/msaa_load always 4), d16 packs, tfe adds 1
-  vdata = 4 if 'gather4' in name or 'msaa_load' in name else (bin(inst.dmask).count('1') or 1)
-  if inst.d16: vdata = (vdata + 1) // 2
-  if inst.tfe: vdata += 1
-  # vaddr width
-  dim_names = ['1d', '2d', '3d', 'cube', '1d_array', '2d_array', '2d_msaa', '2d_msaa_array']
-  dim = dim_names[inst.dim] if inst.dim < len(dim_names) else f"dim_{inst.dim}"
-  vaddr = _mimg_vaddr_width(name, inst.dim, inst.a16)
-  vaddr_str = f"v{inst.vaddr}" if vaddr == 1 else _vreg(inst.vaddr, vaddr)
-  # modifiers
-  mods = [f"dmask:0x{inst.dmask:x}"] if inst.dmask and (inst.dmask != 15 or 'atomic' in name) else []
-  mods.append(f"dim:SQ_RSRC_IMG_{dim.upper()}")
-  for flag, mod in [(inst.unrm,"unorm"),(inst.glc,"glc"),(inst.slc,"slc"),(inst.dlc,"dlc"),(inst.r128,"r128"),
-                    (inst.a16,"a16"),(inst.tfe,"tfe"),(inst.lwe,"lwe"),(inst.d16,"d16")]:
-    if flag: mods.append(mod)
-  # ssamp for sample/gather/get_lod
-  ssamp_str = ""
-  if 'sample' in name or 'gather' in name or 'get_lod' in name:
-    ssamp_str = ", " + _sreg_or_ttmp(inst.ssamp * 4, 4)
-  return f"{name} {_vreg(inst.vdata, vdata)}, {vaddr_str}, {srsrc_str}{ssamp_str} {' '.join(mods)}"
-
-def _disasm_sop1(inst: SOP1) -> str:
-  op, name = inst.op, inst.op_name.lower()
-  if op == SOP1Op.S_GETPC_B64: return f"{name} {_fmt_sdst(inst.sdst, 2)}"
-  if op in (SOP1Op.S_SETPC_B64, SOP1Op.S_RFE_B64): return f"{name} {_fmt_src(inst.ssrc0, 2)}"
-  if op == SOP1Op.S_SWAPPC_B64: return f"{name} {_fmt_sdst(inst.sdst, 2)}, {_fmt_src(inst.ssrc0, 2)}"
-  if op in (SOP1Op.S_SENDMSG_RTN_B32, SOP1Op.S_SENDMSG_RTN_B64): return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs())}, sendmsg({MSG.get(inst.ssrc0, str(inst.ssrc0))})"
-  return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs())}, {inst.lit(inst.ssrc0) if inst.src_regs(0) == 1 else _fmt_src(inst.ssrc0, inst.src_regs(0))}"
-
-def _disasm_sop2(inst: SOP2) -> str:
-  return f"{inst.op_name.lower()} {_fmt_sdst(inst.sdst, inst.dst_regs())}, {inst.lit(inst.ssrc0) if inst.ssrc0 == 255 else _fmt_src(inst.ssrc0, inst.src_regs(0))}, {inst.lit(inst.ssrc1) if inst.ssrc1 == 255 else _fmt_src(inst.ssrc1, inst.src_regs(1))}"
-
-def _disasm_sopc(inst: SOPC) -> str:
-  return f"{inst.op_name.lower()} {_fmt_src(inst.ssrc0, inst.src_regs(0))}, {_fmt_src(inst.ssrc1, inst.src_regs(1))}"
-
-def _disasm_sopk(inst: SOPK) -> str:
-  op, name = inst.op, inst.op_name.lower()
-  if op == SOPKOp.S_VERSION: return f"{name} 0x{inst.simm16:x}"
-  if op in (SOPKOp.S_SETREG_B32, SOPKOp.S_GETREG_B32):
-    hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
-    hs = f"0x{inst.simm16:x}" if hid in (16, 17) else f"hwreg({HWREG.get(hid, str(hid))}, {hoff}, {hsz})"
-    return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1)}" if op == SOPKOp.S_SETREG_B32 else f"{name} {_fmt_sdst(inst.sdst, 1)}, {hs}"
-  return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs())}, 0x{inst.simm16:x}"
-
-def _disasm_vinterp(inst: VINTERP) -> str:
-  mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
-  return f"{inst.op_name.lower()} v{inst.vdst}, {inst.lit(inst.src0, inst.neg & 1)}, {inst.lit(inst.src1, inst.neg & 2)}, {inst.lit(inst.src2, inst.neg & 4)}" + (" " + mods if mods else "")
-
-DISASM_HANDLERS = {VOP1: _disasm_vop1, VOP2: _disasm_vop2, VOPC: _disasm_vopc, VOP3: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOPD: _disasm_vopd, VOP3P: _disasm_vop3p,
-                   VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, MUBUF: _disasm_buf, MTBUF: _disasm_buf,
-                   MIMG: _disasm_mimg, SOP1: _disasm_sop1, SOP2: _disasm_sop2, SOPC: _disasm_sopc, SOPK: _disasm_sopk}
-
-def disasm(inst: Inst) -> str: return DISASM_HANDLERS[type(inst)](inst)
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# ASSEMBLER
-# ═══════════════════════════════════════════════════════════════════════════════
-
-SPEC_REGS = {'vcc_lo': RawImm(106), 'vcc_hi': RawImm(107), 'vcc': RawImm(106), 'null': RawImm(124), 'off': RawImm(124), 'm0': RawImm(125),
-             'exec_lo': RawImm(126), 'exec_hi': RawImm(127), 'exec': RawImm(126), 'scc': RawImm(253), 'src_scc': RawImm(253)}
-FLOATS = {str(k): k for k in FLOAT_ENC}  # Valid float literal strings: '0.5', '-0.5', '1.0', etc.
-REG_MAP: dict[str, _RegFactory] = {'s': s, 'v': v, 't': ttmp, 'ttmp': ttmp}
-SMEM_OPS = {'s_load_b32', 's_load_b64', 's_load_b128', 's_load_b256', 's_load_b512',
-            's_buffer_load_b32', 's_buffer_load_b64', 's_buffer_load_b128', 's_buffer_load_b256', 's_buffer_load_b512'}
-SPEC_DSL = {'vcc_lo': 'VCC_LO', 'vcc_hi': 'VCC_HI', 'vcc': 'VCC_LO', 'null': 'NULL', 'off': 'OFF', 'm0': 'M0',
-            'exec_lo': 'EXEC_LO', 'exec_hi': 'EXEC_HI', 'exec': 'EXEC_LO', 'scc': 'SCC', 'src_scc': 'SCC'}
-
-def _op2dsl(op: str) -> str:
-  op = op.strip()
-  neg = op.startswith('-') and not (op[1:2].isdigit() or (len(op) > 2 and op[1] == '0' and op[2] in 'xX'))
-  if neg: op = op[1:]
-  abs_ = (op.startswith('|') and op.endswith('|')) or (op.startswith('abs(') and op.endswith(')'))
-  if abs_: op = op[1:-1] if op.startswith('|') else op[4:-1]
-  hi = ".h" if op.endswith('.h') else ".l" if op.endswith('.l') else ""
-  if hi: op = op[:-2]
-  lo = op.lower()
-  def wrap(b): return f"{'-' if neg else ''}abs({b}){hi}" if abs_ else f"-{b}{hi}" if neg else f"{b}{hi}"
-  if lo in SPEC_DSL: return wrap(SPEC_DSL[lo])
-  if op in FLOATS: return wrap(op)
-  rp = {'s': 's', 'v': 'v', 't': 'ttmp', 'ttmp': 'ttmp'}
-  if m := re.match(r'^([svt](?:tmp)?)\[(\d+):(\d+)\]$', lo): return wrap(f"{rp[m.group(1)]}[{m.group(2)}:{m.group(3)}]")
-  if m := re.match(r'^([svt](?:tmp)?)(\d+)$', lo): return wrap(f"{rp[m.group(1)]}[{m.group(2)}]")
-  if re.match(r'^-?\d+$|^-?0x[0-9a-fA-F]+$', op): return f"SrcMod({op}, neg={neg}, abs_={abs_})" if neg or abs_ else op
-  return wrap(op)
-
-def _parse_ops(s: str) -> list[str]:
-  ops, cur, depth, pipe = [], "", 0, False
-  for c in s:
-    if c in '[(': depth += 1
-    elif c in '])': depth -= 1
-    elif c == '|': pipe = not pipe
-    if c == ',' and depth == 0 and not pipe: ops.append(cur.strip()); cur = ""
-    else: cur += c
-  if cur.strip(): ops.append(cur.strip())
-  return ops
-
-def _extract(text: str, pat: str, flags=re.I):
-  if m := re.search(pat, text, flags): return m, text[:m.start()] + text[m.end():]
-  return None, text
-
-def get_dsl(text: str) -> str:
-  text, kw = text.strip(), []
-  # Extract modifiers
-  for pat, val in [(r'\s+mul:2(?:\s|$)', 1), (r'\s+mul:4(?:\s|$)', 2), (r'\s+div:2(?:\s|$)', 3)]:
-    if (m := _extract(text, pat))[0]: kw.append(f'omod={val}'); text = m[1]; break
-  if (m := _extract(text, r'\s+clamp(?:\s|$)'))[0]: kw.append('clmp=1'); text = m[1]
-  opsel, m, text = None, *_extract(text, r'\s+op_sel:\[([^\]]+)\]')
-  if m:
-    bits, mn = [int(x.strip()) for x in m.group(1).split(',')], text.split()[0].lower()
-    is3p = mn.startswith(('v_pk_', 'v_wmma_', 'v_dot'))
-    opsel = (bits[0] | (bits[1] << 1) | (bits[2] << 2)) if len(bits) == 3 and is3p else \
-            (bits[0] | (bits[1] << 1) | (bits[2] << 3)) if len(bits) == 3 else sum(b << i for i, b in enumerate(bits))
-  m, text = _extract(text, r'\s+wait_exp:(\d+)'); waitexp = m.group(1) if m else None
-  m, text = _extract(text, r'\s+offset:(0x[0-9a-fA-F]+|-?\d+)'); off_val = m.group(1) if m else None
-  m, text = _extract(text, r'\s+dlc(?:\s|$)'); dlc = 1 if m else None
-  m, text = _extract(text, r'\s+glc(?:\s|$)'); glc = 1 if m else None
-  m, text = _extract(text, r'\s+slc(?:\s|$)'); slc = 1 if m else None
-  m, text = _extract(text, r'\s+neg_lo:\[([^\]]+)\]'); neg_lo = sum(int(x.strip()) << i for i, x in enumerate(m.group(1).split(','))) if m else None
-  m, text = _extract(text, r'\s+neg_hi:\[([^\]]+)\]'); neg_hi = sum(int(x.strip()) << i for i, x in enumerate(m.group(1).split(','))) if m else None
-  if waitexp: kw.append(f'waitexp={waitexp}')
-
-  parts = text.replace(',', ' ').split()
-  if not parts: raise ValueError("empty instruction")
-  mn, op_str = parts[0].lower(), text[len(parts[0]):].strip()
-  ops, args = _parse_ops(op_str), [_op2dsl(o) for o in _parse_ops(op_str)]
-
-  # s_waitcnt
-  if mn == 's_waitcnt':
-    vm, exp, lgkm = 0x3f, 0x7, 0x3f
-    for p in op_str.replace(',', ' ').split():
-      if m := re.match(r'vmcnt\((\d+)\)', p): vm = int(m.group(1))
-      elif m := re.match(r'expcnt\((\d+)\)', p): exp = int(m.group(1))
-      elif m := re.match(r'lgkmcnt\((\d+)\)', p): lgkm = int(m.group(1))
-      elif re.match(r'^0x[0-9a-f]+$|^\d+$', p): return f"s_waitcnt(simm16={int(p, 0)})"
-    return f"s_waitcnt(simm16={waitcnt(vm, exp, lgkm)})"
-
-  # VOPD
-  if '::' in text:
-    xp, yp = text.split('::')
-    xps, yps = xp.strip().replace(',', ' ').split(), yp.strip().replace(',', ' ').split()
-    xo, yo = [_op2dsl(p) for p in xps[1:]], [_op2dsl(p) for p in yps[1:]]
-    vdx, sx0, vsx1 = xo[0], xo[1] if len(xo) > 1 else '0', xo[2] if len(xo) > 2 else 'v[0]'
-    vdy, sy0, vsy1 = yo[0], yo[1] if len(yo) > 1 else '0', yo[2] if len(yo) > 2 else 'v[0]'
-    lit = xo[3] if 'fmaak' in xps[0].lower() and len(xo) > 3 else yo[3] if 'fmaak' in yps[0].lower() and len(yo) > 3 else None
-    if 'fmamk' in xps[0].lower() and len(xo) > 3: lit, vsx1 = xo[2], xo[3]
-    elif 'fmamk' in yps[0].lower() and len(yo) > 3: lit, vsy1 = yo[2], yo[3]
-    return f"VOPD(VOPDOp.{xps[0].upper()}, VOPDOp.{yps[0].upper()}, vdstx={vdx}, vdsty={vdy}, srcx0={sx0}, vsrcx1={vsx1}, srcy0={sy0}, vsrcy1={vsy1}{f', literal={lit}' if lit else ''})"
-
-  # Special instructions
-  if mn == 's_setreg_imm32_b32': raise ValueError(f"unsupported: {mn}")
-  if mn in ('s_setpc_b64', 's_rfe_b64'): return f"{mn}(ssrc0={args[0]})"
-  if mn in ('s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'): return f"{mn}(sdst={args[0]}, ssrc0=RawImm({args[1].strip()}))"
-  if mn == 's_version': return f"{mn}(simm16={args[0]})"
-  if mn == 's_setreg_b32': return f"{mn}(simm16={args[0]}, sdst={args[1]})"
-
-  # SMEM
-  if mn in SMEM_OPS:
-    gs, ds = ", glc=1" if glc else "", ", dlc=1" if dlc else ""
-    if len(ops) >= 3 and re.match(r'^-?[0-9]|^-?0x', ops[2].strip().lower()):
-      return f"{mn}(sdata={args[0]}, sbase={args[1]}, offset={args[2]}, soffset=RawImm(124){gs}{ds})"
-    if off_val and len(ops) >= 3: return f"{mn}(sdata={args[0]}, sbase={args[1]}, offset={off_val}, soffset={args[2]}{gs}{ds})"
-    if len(ops) >= 3: return f"{mn}(sdata={args[0]}, sbase={args[1]}, soffset={args[2]}{gs}{ds})"
-
-  # Buffer
-  if mn.startswith('buffer_') and len(ops) >= 2 and ops[1].strip().lower() == 'off':
-    return f"{mn}(vdata={args[0]}, vaddr=0, srsrc={args[2]}, soffset={f'RawImm({args[3].strip()})' if len(args) > 3 else 'RawImm(0)'})"
-
-  # FLAT/GLOBAL/SCRATCH load/store/atomic - saddr needs RawImm(124) for off/null
-  def _saddr(a): return 'RawImm(124)' if a in ('OFF', 'NULL') else a
-  flat_mods = f"{f', offset={off_val}' if off_val else ''}{', glc=1' if glc else ''}{', slc=1' if slc else ''}{', dlc=1' if dlc else ''}"
-  for pre, flds in [('flat_load','vdst,addr,saddr'), ('global_load','vdst,addr,saddr'), ('scratch_load','vdst,addr,saddr'),
-                    ('flat_store','addr,data,saddr'), ('global_store','addr,data,saddr'), ('scratch_store','addr,data,saddr')]:
-    if mn.startswith(pre) and len(args) >= 2:
-      f0, f1, f2 = flds.split(',')
-      return f"{mn}({f0}={args[0]}, {f1}={args[1]}{f', {f2}={_saddr(args[2])}' if len(args) >= 3 else ', saddr=RawImm(124)'}{flat_mods})"
-  for pre in ('flat_atomic', 'global_atomic', 'scratch_atomic'):
-    if mn.startswith(pre):
-      if glc and len(args) >= 3: return f"{mn}(vdst={args[0]}, addr={args[1]}, data={args[2]}{f', saddr={_saddr(args[3])}' if len(args) >= 4 else ', saddr=RawImm(124)'}{flat_mods})"
-      if len(args) >= 2: return f"{mn}(addr={args[0]}, data={args[1]}{f', saddr={_saddr(args[2])}' if len(args) >= 3 else ', saddr=RawImm(124)'}{flat_mods})"
-
-  # DS instructions
-  if mn.startswith('ds_'):
-    off0, off1 = (str(int(off_val, 0) & 0xff), str((int(off_val, 0) >> 8) & 0xff)) if off_val else ("0", "0")
-    gds_s = ", gds=1" if 'gds' in text.lower().split()[-1:] else ""
-    off_kw = f", offset0={off0}, offset1={off1}{gds_s}"
-    if mn == 'ds_nop' or mn in ('ds_gws_sema_v', 'ds_gws_sema_p', 'ds_gws_sema_release_all'): return f"{mn}({off_kw.lstrip(', ')})"
-    if 'gws_' in mn: return f"{mn}(addr={args[0]}{off_kw})"
-    if 'consume' in mn or 'append' in mn: return f"{mn}(vdst={args[0]}{off_kw})"
-    if 'gs_reg' in mn: return f"{mn}(vdst={args[0]}, data0={args[1]}{off_kw})"
-    if '2addr' in mn:
-      if 'load' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}{off_kw})"
-      if 'store' in mn and 'xchg' not in mn: return f"{mn}(addr={args[0]}, data0={args[1]}, data1={args[2]}{off_kw})"
-      return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}, data1={args[3]}{off_kw})"
-    if 'load' in mn: return f"{mn}(vdst={args[0]}{off_kw})" if 'addtid' in mn else f"{mn}(vdst={args[0]}, addr={args[1]}{off_kw})"
-    if 'store' in mn and not _has(mn, 'cmp', 'xchg'):
-      return f"{mn}(data0={args[0]}{off_kw})" if 'addtid' in mn else f"{mn}(addr={args[0]}, data0={args[1]}{off_kw})"
-    if 'swizzle' in mn or 'ordered_count' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}{off_kw})"
-    if 'permute' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}{off_kw})"
-    if 'bvh' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}, data1={args[3]}{off_kw})"
-    if 'condxchg' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}{off_kw})"
-    if _has(mn, 'cmpstore', 'mskor', 'wrap'):
-      return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}, data1={args[3]}{off_kw})" if '_rtn' in mn else f"{mn}(addr={args[0]}, data0={args[1]}, data1={args[2]}{off_kw})"
-    return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}{off_kw})" if '_rtn' in mn else f"{mn}(addr={args[0]}, data0={args[1]}{off_kw})"
-
-  # v_fmaak/v_fmamk literal extraction
-  lit_s = ""
-  if mn in ('v_fmaak_f32', 'v_fmaak_f16') and len(args) == 4: lit_s, args = f", literal={args[3].strip()}", args[:3]
-  elif mn in ('v_fmamk_f32', 'v_fmamk_f16') and len(args) == 4: lit_s, args = f", literal={args[2].strip()}", [args[0], args[1], args[3]]
-
-  # VCC ops cleanup
-  vcc_ops = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'}
-  if mn.replace('_e32', '') in vcc_ops and len(args) >= 5: mn, args = mn.replace('_e32', '') + '_e32', [args[0], args[2], args[3]]
-  if mn.replace('_e64', '') in vcc_ops and mn.endswith('_e64'): mn = mn.replace('_e64', '')
-  if mn.startswith('v_cmp') and not mn.endswith('_e64') and len(args) >= 3 and ops[0].strip().lower() in ('vcc_lo', 'vcc_hi', 'vcc'): args = args[1:]
-  if 'cmpx' in mn and mn.endswith('_e64') and len(args) == 2: args = ['RawImm(126)'] + args
-
-  fn = mn.replace('.', '_')
-  if opsel is not None: args = [re.sub(r'\.[hl]$', '', a) for a in args]
-
-  # v_fma_mix*: extract inline neg/abs modifiers
-  if 'fma_mix' in mn and neg_lo is None and neg_hi is None:
-    inline_neg, inline_abs, clean_args = 0, 0, [args[0]]
-    for i, op in enumerate(ops[1:4]):
-      op = op.strip()
-      neg = op.startswith('-') and not (op[1:2].isdigit() or (len(op) > 2 and op[1] == '0' and op[2] in 'xX'))
-      if neg: op = op[1:]
-      abs_ = op.startswith('|') and op.endswith('|')
-      if abs_: op = op[1:-1]
-      if neg: inline_neg |= (1 << i)
-      if abs_: inline_abs |= (1 << i)
-      clean_args.append(_op2dsl(op))
-    args = clean_args + args[4:]
-    if inline_neg: neg_lo = inline_neg
-    if inline_abs: neg_hi = inline_abs
-
-  all_kw = list(kw)
-  if lit_s: all_kw.append(lit_s.lstrip(', '))
-  if opsel is not None: all_kw.append(f'opsel={opsel}')
-  if neg_lo is not None: all_kw.append(f'neg={neg_lo}')
-  if neg_hi is not None: all_kw.append(f'neg_hi={neg_hi}')
-  if 'bvh' in mn and 'intersect_ray' in mn: all_kw.extend(['dmask=15', 'unrm=1', 'r128=1'])
-
-  a_str, kw_str = ', '.join(args), ', '.join(all_kw)
-  return f"{fn}({a_str}, {kw_str})" if kw_str and a_str else f"{fn}({kw_str})" if kw_str else f"{fn}({a_str})"
-
-def asm(text: str) -> Inst:
-  dsl = get_dsl(text)
-  ns = {n: getattr(ins, n) for n in dir(ins) if not n.startswith('_')}
-  ns.update({'s': s, 'v': v, 'ttmp': ttmp, 'abs': abs, 'RawImm': RawImm, 'SrcMod': SrcMod, 'VGPR': VGPR, 'SGPR': SGPR, 'TTMP': TTMP,
-             'VCC_LO': VCC_LO, 'VCC_HI': VCC_HI, 'VCC': VCC, 'EXEC_LO': EXEC_LO, 'EXEC_HI': EXEC_HI, 'EXEC': EXEC, 'SCC': SCC, 'M0': M0, 'NULL': NULL, 'OFF': OFF})
-  try: return eval(dsl, ns)
-  except NameError:
-    if m := re.match(r'^(v_\w+)(\(.*\))$', dsl): return eval(f"{m.group(1)}_e32{m.group(2)}", ns)
-    raise
@@ -1,533 +0,0 @@
-# library for RDNA3 assembly DSL
-# mypy: ignore-errors
-from __future__ import annotations
-import struct, math, re
-from enum import IntEnum
-from functools import cache, cached_property
-from typing import overload, Annotated, TypeVar, Generic
-from extra.assembly.amd.autogen.rdna3.enum import (VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, VOPDOp, SOP1Op, SOP2Op,
-  SOPCOp, SOPKOp, SOPPOp, SMEMOp, DSOp, FLATOp, MUBUFOp, MTBUFOp, MIMGOp, VINTERPOp)
-
-# Common masks and bit conversion functions
-MASK32, MASK64 = 0xffffffff, 0xffffffffffffffff
-_struct_f, _struct_I = struct.Struct("<f"), struct.Struct("<I")
-_struct_e, _struct_H = struct.Struct("<e"), struct.Struct("<H")
-_struct_d, _struct_Q = struct.Struct("<d"), struct.Struct("<Q")
-def _f32(i): return _struct_f.unpack(_struct_I.pack(i & MASK32))[0]
-def _i32(f):
-  if isinstance(f, int): f = float(f)
-  if math.isnan(f): return 0xffc00000 if math.copysign(1.0, f) < 0 else 0x7fc00000
-  if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
-  try: return _struct_I.unpack(_struct_f.pack(f))[0]
-  except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
-def _sext(v, b): return v - (1 << b) if v & (1 << (b - 1)) else v
-def _f16(i): return _struct_e.unpack(_struct_H.pack(i & 0xffff))[0]
-def _i16(f):
-  if math.isnan(f): return 0x7e00
-  if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
-  try: return _struct_H.unpack(_struct_e.pack(f))[0]
-  except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
-def _f64(i): return _struct_d.unpack(_struct_Q.pack(i & MASK64))[0]
-def _i64(f):
-  if math.isnan(f): return 0x7ff8000000000000
-  if math.isinf(f): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
-  try: return _struct_Q.unpack(_struct_d.pack(f))[0]
-  except (OverflowError, struct.error): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
-
-# Instruction spec - register counts and dtypes derived from instruction names
-_REGS = {'B32': 1, 'B64': 2, 'B96': 3, 'B128': 4, 'B256': 8, 'B512': 16,
-         'F32': 1, 'I32': 1, 'U32': 1, 'F64': 2, 'I64': 2, 'U64': 2,
-         'F16': 1, 'I16': 1, 'U16': 1, 'B16': 1, 'I8': 1, 'U8': 1, 'B8': 1}
-_CVT_RE = re.compile(r'CVT_([FIUB]\d+)_([FIUB]\d+)$')
-_MAD_MUL_RE = re.compile(r'(?:MAD|MUL)_([IU]\d+)_([IU]\d+)$')
-_PACK_RE = re.compile(r'PACK_([FIUB]\d+)_([FIUB]\d+)$')
-_DST_SRC_RE = re.compile(r'_([FIUB]\d+)_([FIUB]\d+)$')
-_SINGLE_RE = re.compile(r'_([FIUB](?:32|64|16|8|96|128|256|512))$')
-@cache
-def _suffix(name: str) -> tuple[str | None, str | None]:
-  name = name.upper()
-  if m := _CVT_RE.search(name): return m.group(1), m.group(2)
-  if m := _MAD_MUL_RE.search(name): return m.group(1), m.group(2)
-  if m := _PACK_RE.search(name): return m.group(1), m.group(2)
-  if m := _DST_SRC_RE.search(name): return m.group(1), m.group(2)
-  if m := _SINGLE_RE.search(name): return m.group(1), m.group(1)
-  return None, None
-_SPECIAL_REGS = {
-  'V_LSHLREV_B64': (2, 1, 2, 1), 'V_LSHRREV_B64': (2, 1, 2, 1), 'V_ASHRREV_I64': (2, 1, 2, 1),
-  'S_LSHL_B64': (2, 2, 1, 1), 'S_LSHR_B64': (2, 2, 1, 1), 'S_ASHR_I64': (2, 2, 1, 1),
-  'S_BFE_U64': (2, 2, 1, 1), 'S_BFE_I64': (2, 2, 1, 1), 'S_BFM_B64': (2, 1, 1, 1),
-  'S_BITSET0_B64': (2, 1, 1, 1), 'S_BITSET1_B64': (2, 1, 1, 1),
-  'S_BITCMP0_B64': (1, 2, 1, 1), 'S_BITCMP1_B64': (1, 2, 1, 1),
-  'V_LDEXP_F64': (2, 2, 1, 1), 'V_TRIG_PREOP_F64': (2, 2, 1, 1),
-  'V_CMP_CLASS_F64': (1, 2, 1, 1), 'V_CMPX_CLASS_F64': (1, 2, 1, 1),
-  'V_CMP_CLASS_F32': (1, 1, 1, 1), 'V_CMPX_CLASS_F32': (1, 1, 1, 1),
-  'V_CMP_CLASS_F16': (1, 1, 1, 1), 'V_CMPX_CLASS_F16': (1, 1, 1, 1),
-  'V_MAD_U64_U32': (2, 1, 1, 2), 'V_MAD_I64_I32': (2, 1, 1, 2),
-  'V_QSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_U32_U8': (4, 2, 1, 4),
-}
-_SPECIAL_DTYPE = {
-  'V_LSHLREV_B64': ('B64', 'U32', 'B64', None), 'V_LSHRREV_B64': ('B64', 'U32', 'B64', None), 'V_ASHRREV_I64': ('I64', 'U32', 'I64', None),
-  'S_LSHL_B64': ('B64', 'B64', 'U32', None), 'S_LSHR_B64': ('B64', 'B64', 'U32', None), 'S_ASHR_I64': ('I64', 'I64', 'U32', None),
-  'S_BFE_U64': ('U64', 'U64', 'U32', None), 'S_BFE_I64': ('I64', 'I64', 'U32', None),
-  'S_BFM_B64': ('B64', 'U32', 'U32', None), 'S_BITSET0_B64': ('B64', 'U32', None, None), 'S_BITSET1_B64': ('B64', 'U32', None, None),
-  'S_BITCMP0_B64': ('SCC', 'B64', 'U32', None), 'S_BITCMP1_B64': ('SCC', 'B64', 'U32', None),
-  'V_LDEXP_F64': ('F64', 'F64', 'I32', None), 'V_TRIG_PREOP_F64': ('F64', 'F64', 'U32', None),
-  'V_CMP_CLASS_F64': ('VCC', 'F64', 'U32', None), 'V_CMPX_CLASS_F64': ('EXEC', 'F64', 'U32', None),
-  'V_CMP_CLASS_F32': ('VCC', 'F32', 'U32', None), 'V_CMPX_CLASS_F32': ('EXEC', 'F32', 'U32', None),
-  'V_CMP_CLASS_F16': ('VCC', 'F16', 'U32', None), 'V_CMPX_CLASS_F16': ('EXEC', 'F16', 'U32', None),
-  'V_MAD_U64_U32': ('U64', 'U32', 'U32', 'U64'), 'V_MAD_I64_I32': ('I64', 'I32', 'I32', 'I64'),
-  'V_QSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'), 'V_MQSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'),
-  'V_MQSAD_U32_U8': ('B128', 'B64', 'B64', 'B128'),
-}
-@cache
-def spec_regs(name: str) -> tuple[int, int, int, int]:
-  uname = name.upper()
-  if uname in _SPECIAL_REGS: return _SPECIAL_REGS[uname]
-  if 'SAD' in uname and 'U8' in uname and 'QSAD' not in uname and 'MQSAD' not in uname: return 1, 1, 1, 1
-  dst_suf, src_suf = _suffix(name)
-  return _REGS.get(dst_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1)
-@cache
-def spec_dtype(name: str) -> tuple[str | None, str | None, str | None, str | None]:
-  uname = name.upper()
-  if uname in _SPECIAL_DTYPE: return _SPECIAL_DTYPE[uname]
-  if 'SAD' in uname and ('U8' in uname or 'U16' in uname) and 'QSAD' not in uname and 'MQSAD' not in uname: return 'U32', 'U32', 'U32', 'U32'
-  if '_CMP_' in uname or '_CMPX_' in uname:
-    dst_suf, src_suf = _suffix(name)
-    return 'EXEC' if '_CMPX_' in uname else 'VCC', src_suf, src_suf, None
-  dst_suf, src_suf = _suffix(name)
-  return dst_suf, src_suf, src_suf, src_suf
-_F16_RE = re.compile(r'_[FIUB]16(?:_|$)')
-_F64_RE = re.compile(r'_[FIUB]64(?:_|$)')
-@cache
-def spec_is_16bit(name: str) -> bool:
-  uname = name.upper()
-  if 'SAD' in uname or 'PACK' in uname or '_PK_' in uname or 'SAT_PK' in uname or 'DOT2' in uname: return False
-  if '_F32' in uname or '_I32' in uname or '_U32' in uname or '_B32' in uname: return False
-  return bool(_F16_RE.search(uname))
-@cache
-def spec_is_64bit(name: str) -> bool: return bool(_F64_RE.search(name.upper()))
-_3SRC = {'FMA', 'MAD', 'MIN3', 'MAX3', 'MED3', 'DIV_FIX', 'DIV_FMAS', 'DIV_SCALE', 'SAD', 'LERP', 'ALIGN', 'CUBE', 'BFE', 'BFI',
-         'PERM_B32', 'PERMLANE', 'CNDMASK', 'XOR3', 'OR3', 'ADD3', 'LSHL_OR', 'AND_OR', 'LSHL_ADD', 'ADD_LSHL', 'XAD', 'MAXMIN',
-         'MINMAX', 'DOT2', 'DOT4', 'DOT8', 'WMMA', 'CVT_PK_U8', 'MULLIT', 'CO_CI'}
-_2SRC = {'FMAC'}  # FMAC uses dst as implicit accumulator, so only 2 explicit sources
-def spec_num_srcs(name: str) -> int:
-  name = name.upper()
-  if any(k in name for k in _2SRC): return 2
-  return 3 if any(k in name for k in _3SRC) else 2
-def is_dtype_16(dt: str | None) -> bool: return dt is not None and '16' in dt
-def is_dtype_64(dt: str | None) -> bool: return dt is not None and '64' in dt
-
-# Bit field DSL
-class BitField:
-  def __init__(self, hi: int, lo: int, name: str | None = None): self.hi, self.lo, self.name, self._marker = hi, lo, name, None
-  def __set_name__(self, owner, name):
-    import typing
-    self.name, self._owner = name, owner
-    # Cache marker at class definition time
-    hints = typing.get_type_hints(owner, include_extras=True)
-    if name in hints:
-      hint = hints[name]
-      if typing.get_origin(hint) is Annotated:
-        args = typing.get_args(hint)
-        self._marker = args[1] if len(args) > 1 else None
-  def __eq__(self, val: int) -> tuple[BitField, int]: return (self, val)  # type: ignore
-  def mask(self) -> int: return (1 << (self.hi - self.lo + 1)) - 1
-  @property
-  def marker(self) -> type | None: return self._marker
-  @overload
-  def __get__(self, obj: None, objtype: type) -> BitField: ...
-  @overload
-  def __get__(self, obj: object, objtype: type | None = None) -> int: ...
-  def __get__(self, obj, objtype=None):
-    if obj is None: return self
-    val = unwrap(obj._values.get(self.name, 0))
-    # Convert to IntEnum if marker is an IntEnum subclass
-    if self.marker and isinstance(self.marker, type) and issubclass(self.marker, IntEnum):
-      # VOP3 with VOPC opcodes (0-255) -> VOPCOp, VOP3SD opcodes -> VOP3SDOp
-      if self.marker is VOP3Op:
-        if val < 256: return VOPCOp(val)
-        if val in Inst._VOP3SD_OPS: return VOP3SDOp(val)
-      try: return self.marker(val)
-      except ValueError: pass
-    return val
-
-class _Bits:
-  def __getitem__(self, key) -> BitField: return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
-bits = _Bits()
-
-# Source operand with modifiers - base class for anything that can be a src with neg/abs
-class SrcMod:
-  __slots__ = ('val', 'neg', 'abs_')
-  def __init__(self, val: int, neg: bool = False, abs_: bool = False): self.val, self.neg, self.abs_ = val, neg, abs_
-  def __repr__(self): return f"{'-' if self.neg else ''}{'|' if self.abs_ else ''}{self.val}{'|' if self.abs_ else ''}"
-  def __neg__(self): return SrcMod(self.val, not self.neg, self.abs_)
-  def __abs__(self): return SrcMod(self.val, self.neg, True)
-
-# Register types
-class Reg(SrcMod):
-  __slots__ = ('idx', 'count', 'hi')
-  def __init__(self, idx: int, count: int = 1, hi: bool = False, neg: bool = False, abs_: bool = False):
-    self.idx, self.count, self.hi = idx, count, hi
-    super().__init__(idx, neg, abs_)
-  def __repr__(self): return f"{self.__class__.__name__.lower()[0]}[{self.idx}]" if self.count == 1 else f"{self.__class__.__name__.lower()[0]}[{self.idx}:{self.idx + self.count}]"
-  def __neg__(self): return self.__class__(self.idx, self.count, self.hi, not self.neg, self.abs_)
-  def __abs__(self): return self.__class__(self.idx, self.count, self.hi, self.neg, True)
-  @property
-  def l(self): return self.__class__(self.idx, self.count, False, self.neg, self.abs_)
-  @property
-  def h(self): return self.__class__(self.idx, self.count, True, self.neg, self.abs_)
-
-T = TypeVar('T', bound=Reg)
-class _RegFactory(Generic[T]):
-  def __init__(self, cls: type[T], name: str): self._cls, self._name = cls, name
-  @overload
-  def __getitem__(self, key: int) -> Reg: ...
-  @overload
-  def __getitem__(self, key: slice) -> Reg: ...
-  def __getitem__(self, key: int | slice) -> Reg:
-    return self._cls(key.start, key.stop - key.start + 1) if isinstance(key, slice) else self._cls(key)
-  def __repr__(self): return f"<{self._name} factory>"
-
-class SGPR(Reg): pass
-class VGPR(Reg): pass
-class TTMP(Reg): pass
-s: _RegFactory[SGPR] = _RegFactory(SGPR, "SGPR")
-v: _RegFactory[VGPR] = _RegFactory(VGPR, "VGPR")
-ttmp: _RegFactory[TTMP] = _RegFactory(TTMP, "TTMP")
-
-# Special registers as SrcMod objects (support -VCC_LO, abs(EXEC_LO), etc.)
-VCC_LO, VCC_HI, VCC = SrcMod(106), SrcMod(107), SrcMod(106)
-EXEC_LO, EXEC_HI, EXEC = SrcMod(126), SrcMod(127), SrcMod(126)
-SCC, M0, NULL, OFF = SrcMod(253), SrcMod(125), SrcMod(124), SrcMod(124)
-
-# Field type markers (runtime classes for validation)
-class _SSrc: pass
-class _Src: pass
-class _Imm: pass
-class _SImm: pass
-class _VDSTYEnc: pass  # VOPD vdsty: encoded = actual >> 1, actual = (encoded << 1) | ((vdstx & 1) ^ 1)
-class _SGPRField: pass
-class _VGPRField: pass
-
-# Type aliases for annotations - tells mypy it's a BitField while preserving marker info
-SSrc = Annotated[BitField, _SSrc]
-Src = Annotated[BitField, _Src]
-Imm = Annotated[BitField, _Imm]
-SImm = Annotated[BitField, _SImm]
-VDSTYEnc = Annotated[BitField, _VDSTYEnc]
-SGPRField = Annotated[BitField, _SGPRField]
-VGPRField = Annotated[BitField, _VGPRField]
-class RawImm:
-  def __init__(self, val: int): self.val = val
-  def __repr__(self): return f"RawImm({self.val})"
-  def __eq__(self, other): return isinstance(other, RawImm) and self.val == other.val
-
-def unwrap(val) -> int:
-  if isinstance(val, RawImm): return val.val
-  if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val  # Special registers like VCC_LO, NULL
-  if hasattr(val, 'value'): return val.value  # IntEnum
-  if hasattr(val, 'idx'): return val.idx  # Reg
-  return val
-
-# Encoding/decoding constants
-FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
-FLOAT_DEC = {v: str(k) for k, v in FLOAT_ENC.items()}
-SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
-SPECIAL_PAIRS = {106: "vcc", 126: "exec"}
-SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
-RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata', 'vsrc1'}
-
-def _encode_reg(val: Reg) -> int: return (108 if isinstance(val, TTMP) else 0) + val.idx
-
-def _is_inline_const(v: int) -> bool: return 0 <= v <= 127 or 128 <= v <= 208 or 240 <= v <= 255
-
-def encode_src(val) -> int:
-  if isinstance(val, VGPR): return 256 + _encode_reg(val)
-  if isinstance(val, Reg): return _encode_reg(val)
-  if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val if _is_inline_const(val.val) else 255
-  if hasattr(val, 'value'): return val.value  # IntEnum
-  if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
-  if isinstance(val, int): return 128 + val if 0 <= val <= 64 else 192 - val if -16 <= val <= -1 else 255
-  return 255
-
-def decode_src(val: int) -> str:
-  if val <= 105: return f"s{val}"
-  if val in SPECIAL_GPRS: return SPECIAL_GPRS[val]
-  if val in FLOAT_DEC: return FLOAT_DEC[val]
-  if 108 <= val <= 123: return f"ttmp{val - 108}"
-  if 128 <= val <= 192: return str(val - 128)
-  if 193 <= val <= 208: return str(-(val - 192))
-  if 256 <= val <= 511: return f"v{val - 256}"
-  return "lit" if val == 255 else f"?{val}"
-
-# Instruction base class
-class Inst:
-  _fields: dict[str, BitField]
-  _encoding: tuple[BitField, int] | None = None
-  _defaults: dict[str, int] = {}
-  _values: dict[str, int | RawImm]
-  _words: int  # size in 32-bit words, set by decode_program
-  _literal: int | None
-
-  def __init_subclass__(cls, **kwargs):
-    super().__init_subclass__(**kwargs)
-    cls._fields = {n: v[0] if isinstance(v, tuple) else v for n, v in cls.__dict__.items() if isinstance(v, BitField) or (isinstance(v, tuple) and len(v) == 2 and isinstance(v[0], BitField))}
-    if 'encoding' in cls._fields and isinstance(cls.__dict__.get('encoding'), tuple): cls._encoding = cls.__dict__['encoding']
-
-  def _or_field(self, name: str, bit: int):
-    cur = self._values.get(name, 0)
-    self._values[name] = (cur.val if isinstance(cur, RawImm) else cur) | bit
-
-  def _encode_src(self, name: str, val):
-    """Encode a source field, handling modifiers and literals."""
-    encoded = encode_src(val)
-    has_opsel = 'opsel' in self._fields
-    if isinstance(val, Reg) and val.hi and not has_opsel: encoded |= 0x80  # hi bit in src for VOP1/2/C
-    self._values[name] = RawImm(encoded)
-    # Handle neg/abs/opsel modifiers
-    if isinstance(val, SrcMod):
-      mod_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
-      if val.neg and 'neg' in self._fields: self._or_field('neg', mod_bit)
-      if val.abs_ and 'abs' in self._fields: self._or_field('abs', mod_bit)
-    if isinstance(val, Reg) and val.hi and has_opsel:
-      self._or_field('opsel', {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0))
-    # Track literal value if needed
-    if encoded == 255 and self._literal is None:
-      import struct
-      # Check if THIS source uses 64-bit encoding (not just src0)
-      src_idx = {'src0': 0, 'src1': 1, 'src2': 2, 'ssrc0': 0, 'ssrc1': 1}.get(name, 0)
-      src_regs = self.src_regs(src_idx)
-      is_64 = src_regs == 2
-      if isinstance(val, SrcMod) and not isinstance(val, Reg): lit32 = val.val & MASK32
-      elif isinstance(val, int) and not isinstance(val, IntEnum): lit32 = val & MASK32
-      elif isinstance(val, float): lit32 = (_i64(val) >> 32) if is_64 else _i32(val)  # f64: high 32 bits of f64 repr
-      else: return
-      self._literal = (lit32 << 32) if is_64 else lit32
-
-  def _encode_raw(self, name: str, val):
-    """Encode a raw register field (vdst, vdata, etc.)."""
-    if isinstance(val, Reg):
-      encoded = _encode_reg(val)
-      if val.hi and 'opsel' not in self._fields: encoded |= 0x80
-      self._values[name] = encoded
-      if name == 'vdst' and val.hi and 'opsel' in self._fields: self._or_field('opsel', 8)
-    elif hasattr(val, 'value'): self._values[name] = val.value
-
-  def _validate(self, orig_args: dict):
-    """Format-specific validation. Override in subclass or check by class name."""
-    cls_name, op = self.__class__.__name__, orig_args.get('op')
-    if hasattr(op, 'value'): op = op.value
-    # SMEM: register count must match opcode
-    if cls_name == 'SMEM' and op is not None:
-      expected = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op)
-      sdata = orig_args.get('sdata')
-      if expected and isinstance(sdata, Reg) and sdata.count != expected:
-        raise ValueError(f"SMEM op {op} expects {expected} registers, got {sdata.count}")
-    # SOP1: b32=1 reg, b64=2 regs
-    if cls_name == 'SOP1' and hasattr(orig_args.get('op'), 'name'):
-      expected = 2 if orig_args['op'].name.endswith('_B64') else 1
-      for fld in ('sdst', 'ssrc0'):
-        if isinstance(orig_args.get(fld), Reg) and orig_args[fld].count != expected:
-          raise ValueError(f"SOP1 {orig_args['op'].name} expects {expected} register(s) for {fld}, got {orig_args[fld].count}")
-
-  def __init__(self, *args, literal: int | None = None, **kwargs):
-    self._values, self._literal = dict(self._defaults), None
-    field_names = [n for n in self._fields if n != 'encoding']
-    orig_args = dict(zip(field_names, args)) | kwargs
-    self._values.update(orig_args)
-    self._validate(orig_args)
-    # Pre-shift literal for 64-bit sources (literal param is always raw 32-bit value from user)
-    if literal is not None:
-      # Find which source uses the literal (255) and check its register count
-      for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
-        v = orig_args.get(n)
-        if (isinstance(v, RawImm) and v.val == 255) or (isinstance(v, int) and v == 255):
-          self._literal = (literal << 32) if self.src_regs(idx) == 2 else literal
-          break
-      else:
-        self._literal = literal  # fallback if no literal source found
-    cls_name = self.__class__.__name__
-
-    # Format-specific setup
-    if cls_name == 'FLAT' and 'sve' in self._fields:
-      seg = self._values.get('seg', 0)
-      if (seg.val if isinstance(seg, RawImm) else seg) == 1 and isinstance(orig_args.get('addr'), VGPR): self._values['sve'] = 1
-    if cls_name == 'VOP3P':
-      op = orig_args.get('op')
-      if hasattr(op, 'value'): op = op.value
-      if op in (32, 33, 34) and 'opsel_hi' not in orig_args: self._values['opsel_hi'] = self._values['opsel_hi2'] = 0
-
-    # Encode all fields
-    for name, val in list(self._values.items()):
-      if name == 'encoding': continue
-      if isinstance(val, RawImm):
-        if name in RAW_FIELDS: self._values[name] = val.val
-        continue
-      field = self._fields.get(name)
-      marker = field.marker if field else None
-      # Type validation
-      if marker is _SGPRField and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
-      if marker is _VGPRField and not isinstance(val, VGPR): raise TypeError(f"field '{name}' requires VGPR, got {type(val).__name__}")
-      if marker is _SSrc and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires scalar source, got VGPR")
-      # Encode by field type
-      if name in SRC_FIELDS: self._encode_src(name, val)
-      elif name in RAW_FIELDS: self._encode_raw(name, val)
-      elif name == 'sbase': self._values[name] = (val.idx if isinstance(val, Reg) else val.val if isinstance(val, SrcMod) else val * 2) // 2
-      elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg): self._values[name] = val.idx // 4
-      elif marker is _VDSTYEnc and isinstance(val, VGPR): self._values[name] = val.idx >> 1
-
-  def _encode_field(self, name: str, val) -> int:
-    if isinstance(val, RawImm): return val.val
-    if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val  # Special regs like VCC_LO
-    if name in {'srsrc', 'ssamp'}: return val.idx // 4 if isinstance(val, Reg) else val
-    if name == 'sbase': return val.idx // 2 if isinstance(val, Reg) else val.val // 2 if isinstance(val, SrcMod) else val
-    if name in RAW_FIELDS: return _encode_reg(val) if isinstance(val, Reg) else val
-    if isinstance(val, Reg) or name in SRC_FIELDS: return encode_src(val)
-    return val.value if hasattr(val, 'value') else val
-
-  def to_int(self) -> int:
-    word = (self._encoding[1] & self._encoding[0].mask()) << self._encoding[0].lo if self._encoding else 0
-    for n, bf in self._fields.items():
-      if n != 'encoding' and n in self._values: word |= (self._encode_field(n, self._values[n]) & bf.mask()) << bf.lo
-    return word
-
-  def _get_literal(self) -> int | None:
-    for n in SRC_FIELDS:
-      if n in self._values and not isinstance(v := self._values[n], RawImm) and isinstance(v, int) and not isinstance(v, IntEnum) and not (0 <= v <= 64 or -16 <= v <= -1): return v
-    return None
-
-  def _is_64bit_op(self) -> bool:
-    """Check if this instruction uses 64-bit operands (and thus 64-bit literals)."""
-    op = self._values.get('op')
-    if op is None: return False
-    op_name = op.name if hasattr(op, 'name') else None
-    # Look up op name from int if needed (happens in from_bytes path)
-    if op_name is None and self.__class__.__name__ == 'VOP3':
-      try: op_name = VOP3Op(op).name
-      except ValueError: pass
-    if op_name is None and self.__class__.__name__ == 'VOPC':
-      try: op_name = VOPCOp(op).name
-      except ValueError: pass
-    if op_name is None: return False
-    # V_LDEXP_F64 has 32-bit integer src1, so literal is 32-bit
-    return op_name != 'V_LDEXP_F64' and op_name.endswith(('_F64', '_B64', '_I64', '_U64'))
-
-  def to_bytes(self) -> bytes:
-    result = self.to_int().to_bytes(self._size(), 'little')
-    lit = self._get_literal() or getattr(self, '_literal', None)
-    if lit is None: return result
-    # For 64-bit sources, literal is stored in high 32 bits internally, but encoded as 4 bytes
-    # Find which source uses the literal (255) and check its register count
-    lit_src_is_64 = False
-    for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
-      if n not in self._values: continue
-      v = self._values[n]
-      if (isinstance(v, RawImm) and v.val == 255) or (isinstance(v, int) and v == 255):
-        lit_src_is_64 = self.is_src_64(idx)
-        break
-    lit32 = (lit >> 32) if lit_src_is_64 else lit
-    return result + (lit32 & MASK32).to_bytes(4, 'little')
-
-  @classmethod
-  def _size(cls) -> int: return 4 if issubclass(cls, Inst32) else 8
-  def size(self) -> int:
-    # Literal is always 4 bytes in the binary (for 64-bit ops, it's in high 32 bits)
-    return self._size() + (4 if self._literal is not None else 0)
-
-  @classmethod
-  def from_int(cls, word: int):
-    inst = object.__new__(cls)
-    inst._values = {n: RawImm(v) if n in SRC_FIELDS else v for n, bf in cls._fields.items() if n != 'encoding' for v in [(word >> bf.lo) & bf.mask()]}
-    inst._literal = None
-    return inst
-
-  @classmethod
-  def from_bytes(cls, data: bytes):
-    inst = cls.from_int(int.from_bytes(data[:cls._size()], 'little'))
-    op_val = inst._values.get('op', 0)
-    has_literal = cls.__name__ == 'VOP2' and op_val in (44, 45, 55, 56)
-    has_literal = has_literal or (cls.__name__ == 'SOP2' and op_val in (69, 70))
-    # VOPD fmaak/fmamk always have a literal (opx/opy value 1 or 2)
-    opx, opy = inst._values.get('opx', 0), inst._values.get('opy', 0)
-    has_literal = has_literal or (cls.__name__ == 'VOPD' and (opx in (1, 2) or opy in (1, 2)))
-    for n in SRC_FIELDS:
-      if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255: has_literal = True
-    if has_literal:
-      # For 64-bit ops, the literal is 32 bits placed in the HIGH 32 bits of the 64-bit value
-      # (low 32 bits are zero). This is how AMD hardware interprets 32-bit literals for 64-bit ops.
-      # Check which source uses the literal and whether THAT source is 64-bit
-      if len(data) >= cls._size() + 4:
-        lit32 = int.from_bytes(data[cls._size():cls._size()+4], 'little')
-        # Find which source has literal (255) and check its register count
-        lit_src_is_64 = False
-        for n, idx in [('src0', 0), ('src1', 1), ('src2', 2)]:
-          if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255:
-            lit_src_is_64 = inst.src_regs(idx) == 2
-            break
-        inst._literal = (lit32 << 32) if lit_src_is_64 else lit32
-    return inst
-
-  def __repr__(self):
-    # Use _fields order and exclude fields that are 0/default (for consistent repr after roundtrip)
-    def is_zero(v): return (isinstance(v, int) and v == 0) or (isinstance(v, VGPR) and v.idx == 0 and v.count == 1)
-    items = [(k, self._values[k]) for k in self._fields if k in self._values and k != 'encoding'
-             and not (is_zero(self._values[k]) and k not in {'op'})]
-    lit = f", literal={hex(self._literal)}" if self._literal is not None else ""
-    return f"{self.__class__.__name__}({', '.join(f'{k}={v}' for k, v in items)}{lit})"
-
-  def __getattr__(self, name: str):
-    if name.startswith('_'): raise AttributeError(name)
-    return unwrap(self._values.get(name, 0))
-
-  def lit(self, v: int, neg: bool = False) -> str:
-    s = f"0x{self._literal:x}" if v == 255 and self._literal else decode_src(v)
-    return f"-{s}" if neg else s
-
-  def __eq__(self, other):
-    if not isinstance(other, Inst): return NotImplemented
-    return self.__class__ == other.__class__ and self._values == other._values and self._literal == other._literal
-
-  def __hash__(self): return hash((self.__class__.__name__, tuple(sorted((k, repr(v)) for k, v in self._values.items())), self._literal))
-
-  def disasm(self) -> str:
-    from extra.assembly.amd.asm import disasm
-    return disasm(self)
-
-  _enum_map = {'VOP1': VOP1Op, 'VOP2': VOP2Op, 'VOP3': VOP3Op, 'VOP3SD': VOP3SDOp, 'VOP3P': VOP3POp, 'VOPC': VOPCOp,
-               'SOP1': SOP1Op, 'SOP2': SOP2Op, 'SOPC': SOPCOp, 'SOPK': SOPKOp, 'SOPP': SOPPOp,
-               'SMEM': SMEMOp, 'DS': DSOp, 'FLAT': FLATOp, 'MUBUF': MUBUFOp, 'MTBUF': MTBUFOp, 'MIMG': MIMGOp,
-               'VOPD': VOPDOp, 'VINTERP': VINTERPOp}
-  _VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
-
-  @property
-  def op(self):
-    """Return the op as an enum (e.g., VOP1Op.V_MOV_B32). VOP3 returns VOPCOp/VOP3SDOp for those op ranges."""
-    val = self._values.get('op')
-    if val is None: return None
-    if hasattr(val, 'name'): return val  # already an enum
-    cls_name = self.__class__.__name__
-    assert cls_name in self._enum_map, f"no enum map for {cls_name}"
-    return self._enum_map[cls_name](val)
-
-  @cached_property
-  def op_name(self) -> str:
-    op = self.op
-    return op.name if hasattr(op, 'name') else ''
-
-  @cached_property
-  def _spec_regs(self) -> tuple[int, int, int, int]: return spec_regs(self.op_name)
-  @cached_property
-  def _spec_dtype(self) -> tuple[str | None, str | None, str | None, str | None]: return spec_dtype(self.op_name)
-  def dst_regs(self) -> int: return self._spec_regs[0]
-  def src_regs(self, n: int) -> int: return self._spec_regs[n + 1]
-  def num_srcs(self) -> int: return spec_num_srcs(self.op_name)
-  def dst_dtype(self) -> str | None: return self._spec_dtype[0]
-  def src_dtype(self, n: int) -> str | None: return self._spec_dtype[n + 1]
-  def is_src_16(self, n: int) -> bool: return self._spec_regs[n + 1] == 1 and is_dtype_16(self._spec_dtype[n + 1])
-  def is_src_64(self, n: int) -> bool: return self._spec_regs[n + 1] == 2
-  def is_16bit(self) -> bool: return spec_is_16bit(self.op_name)
-  def is_64bit(self) -> bool: return spec_is_64bit(self.op_name)
-  def is_dst_16(self) -> bool: return self._spec_regs[0] == 1 and is_dtype_16(self._spec_dtype[0])
-
-class Inst32(Inst): pass
-class Inst64(Inst): pass
@@ -1,462 +0,0 @@
-# RDNA3 emulator - executes compiled pseudocode from AMD ISA PDF
-# mypy: ignore-errors
-from __future__ import annotations
-import ctypes
-from extra.assembly.amd.dsl import Inst, unwrap, FLOAT_ENC, MASK32, MASK64, _f32, _i32, _sext, _f16, _i16, _f64, _i64
-from extra.assembly.amd.pcode import Reg
-from extra.assembly.amd.asm import detect_format
-from extra.assembly.amd.autogen.rdna3.gen_pcode import get_compiled_functions
-from extra.assembly.amd.autogen.rdna3.ins import (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD,
-  SrcEnum, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, GLOBALOp, FLATOp, DSOp, VOPDOp)
-
-Program = dict[int, Inst]
-WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
-VCC_LO, VCC_HI, NULL, EXEC_LO, EXEC_HI, SCC = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.SCC
-
-# Inline constants for src operands 128-254. Build tables for f32, f16, and f64 formats.
-_FLOAT_CONSTS = {v: k for k, v in FLOAT_ENC.items()} | {248: 0.15915494309189535}  # INV_2PI
-def _build_inline_consts(mask, to_bits):
-  tbl = list(range(65)) + [((-i) & mask) for i in range(1, 17)] + [0] * (127 - 81)
-  for k, v in _FLOAT_CONSTS.items(): tbl[k - 128] = to_bits(v)
-  return tbl
-_INLINE_CONSTS = _build_inline_consts(MASK32, _i32)
-_INLINE_CONSTS_F16 = _build_inline_consts(0xffff, _i16)
-_INLINE_CONSTS_F64 = _build_inline_consts(MASK64, _i64)
-
-# Helper: extract/write 16-bit half from/to 32-bit value
-def _src16(raw: int, is_hi: bool) -> int: return ((raw >> 16) & 0xffff) if is_hi else (raw & 0xffff)
-def _dst16(cur: int, val: int, is_hi: bool) -> int: return (cur & 0x0000ffff) | ((val & 0xffff) << 16) if is_hi else (cur & 0xffff0000) | (val & 0xffff)
-def _vgpr_hi(src: int) -> bool: return src >= 256 and ((src - 256) & 0x80) != 0
-def _vgpr_masked(src: int) -> int: return ((src - 256) & 0x7f) + 256 if src >= 256 else src
-
-# Memory access
-_valid_mem_ranges: list[tuple[int, int]] = []
-def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: _valid_mem_ranges.clear(); _valid_mem_ranges.extend(ranges)
-def _mem_valid(addr: int, size: int) -> bool:
-  return not _valid_mem_ranges or any(s <= addr and addr + size <= s + z for s, z in _valid_mem_ranges)
-def _ctypes_at(addr: int, size: int): return (ctypes.c_uint8 if size == 1 else ctypes.c_uint16 if size == 2 else ctypes.c_uint32).from_address(addr)
-def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
-def mem_write(addr: int, size: int, val: int) -> None:
-  if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
-
-# Memory op tables (not pseudocode - these are format descriptions)
-def _mem_ops(ops, suffix_map):
-  return {getattr(e, f"{p}_{s}"): v for e in ops for s, v in suffix_map.items() for p in [e.__name__.replace("Op", "")]}
-_LOAD_MAP = {'LOAD_B32': (1,4,0), 'LOAD_B64': (2,4,0), 'LOAD_B96': (3,4,0), 'LOAD_B128': (4,4,0), 'LOAD_U8': (1,1,0), 'LOAD_I8': (1,1,1), 'LOAD_U16': (1,2,0), 'LOAD_I16': (1,2,1)}
-_STORE_MAP = {'STORE_B32': (1,4), 'STORE_B64': (2,4), 'STORE_B96': (3,4), 'STORE_B128': (4,4), 'STORE_B8': (1,1), 'STORE_B16': (1,2)}
-FLAT_LOAD, FLAT_STORE = _mem_ops([GLOBALOp, FLATOp], _LOAD_MAP), _mem_ops([GLOBALOp, FLATOp], _STORE_MAP)
-# D16 ops: load/store 16-bit to lower or upper half of VGPR. Format: (size, sign, hi) where hi=1 means upper 16 bits
-_D16_LOAD_MAP = {'LOAD_D16_U8': (1,0,0), 'LOAD_D16_I8': (1,1,0), 'LOAD_D16_B16': (2,0,0),
-                 'LOAD_D16_HI_U8': (1,0,1), 'LOAD_D16_HI_I8': (1,1,1), 'LOAD_D16_HI_B16': (2,0,1)}
-_D16_STORE_MAP = {'STORE_D16_HI_B8': (1,1), 'STORE_D16_HI_B16': (2,1)}  # (size, hi)
-FLAT_D16_LOAD = _mem_ops([GLOBALOp, FLATOp], _D16_LOAD_MAP)
-FLAT_D16_STORE = _mem_ops([GLOBALOp, FLATOp], _D16_STORE_MAP)
-SMEM_LOAD = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}
-
-# VOPD op -> VOP3 op mapping (VOPD is dual-issue of VOP1/VOP2 ops, use VOP3 enums for pseudocode lookup)
-_VOPD_TO_VOP = {
-  VOPDOp.V_DUAL_FMAC_F32: VOP3Op.V_FMAC_F32, VOPDOp.V_DUAL_FMAAK_F32: VOP2Op.V_FMAAK_F32, VOPDOp.V_DUAL_FMAMK_F32: VOP2Op.V_FMAMK_F32,
-  VOPDOp.V_DUAL_MUL_F32: VOP3Op.V_MUL_F32, VOPDOp.V_DUAL_ADD_F32: VOP3Op.V_ADD_F32, VOPDOp.V_DUAL_SUB_F32: VOP3Op.V_SUB_F32,
-  VOPDOp.V_DUAL_SUBREV_F32: VOP3Op.V_SUBREV_F32, VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: VOP3Op.V_MUL_DX9_ZERO_F32,
-  VOPDOp.V_DUAL_MOV_B32: VOP3Op.V_MOV_B32, VOPDOp.V_DUAL_CNDMASK_B32: VOP3Op.V_CNDMASK_B32,
-  VOPDOp.V_DUAL_MAX_F32: VOP3Op.V_MAX_F32, VOPDOp.V_DUAL_MIN_F32: VOP3Op.V_MIN_F32,
-  VOPDOp.V_DUAL_ADD_NC_U32: VOP3Op.V_ADD_NC_U32, VOPDOp.V_DUAL_LSHLREV_B32: VOP3Op.V_LSHLREV_B32, VOPDOp.V_DUAL_AND_B32: VOP3Op.V_AND_B32,
-}
-
-# Compiled pseudocode functions (lazy loaded)
-_COMPILED: dict | None = None
-
-def _get_compiled() -> dict:
-  global _COMPILED
-  if _COMPILED is None: _COMPILED = get_compiled_functions()
-  return _COMPILED
-
-class WaveState:
-  __slots__ = ('sgpr', 'vgpr', 'scc', 'pc', 'literal', '_pend_sgpr')
-  def __init__(self):
-    self.sgpr, self.vgpr = [0] * SGPR_COUNT, [[0] * VGPR_COUNT for _ in range(WAVE_SIZE)]
-    self.sgpr[EXEC_LO], self.scc, self.pc, self.literal, self._pend_sgpr = 0xffffffff, 0, 0, 0, {}
-
-  @property
-  def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
-  @vcc.setter
-  def vcc(self, v: int): self.sgpr[VCC_LO], self.sgpr[VCC_HI] = v & MASK32, (v >> 32) & MASK32
-  @property
-  def exec_mask(self) -> int: return self.sgpr[EXEC_LO] | (self.sgpr[EXEC_HI] << 32)
-  @exec_mask.setter
-  def exec_mask(self, v: int): self.sgpr[EXEC_LO], self.sgpr[EXEC_HI] = v & MASK32, (v >> 32) & MASK32
-
-  def rsgpr(self, i: int) -> int: return 0 if i == NULL else self.scc if i == SCC else self.sgpr[i] if i < SGPR_COUNT else 0
-  def wsgpr(self, i: int, v: int):
-    if i < SGPR_COUNT and i != NULL: self.sgpr[i] = v & MASK32
-  def rsgpr64(self, i: int) -> int: return self.rsgpr(i) | (self.rsgpr(i+1) << 32)
-  def wsgpr64(self, i: int, v: int): self.wsgpr(i, v & MASK32); self.wsgpr(i+1, (v >> 32) & MASK32)
-
-  def _rsrc_base(self, v: int, lane: int, consts):
-    if v < SGPR_COUNT: return self.sgpr[v]
-    if v == SCC: return self.scc
-    if v < 255: return consts[v - 128]
-    if v == 255: return self.literal
-    return self.vgpr[lane][v - 256] if v <= 511 else 0
-  def rsrc(self, v: int, lane: int) -> int: return self._rsrc_base(v, lane, _INLINE_CONSTS)
-  def rsrc_f16(self, v: int, lane: int) -> int: return self._rsrc_base(v, lane, _INLINE_CONSTS_F16)
-  def rsrc64(self, v: int, lane: int) -> int:
-    if 128 <= v < 255: return _INLINE_CONSTS_F64[v - 128]
-    if v == 255: return self.literal  # literal is already shifted in from_bytes for 64-bit ops
-    return self.rsrc(v, lane) | ((self.rsrc(v+1, lane) if v < VCC_LO or 256 <= v <= 511 else 0) << 32)
-
-  def pend_sgpr_lane(self, reg: int, lane: int, val: int):
-    if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
-    if val: self._pend_sgpr[reg] |= (1 << lane)
-  def commit_pends(self):
-    for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
-    self._pend_sgpr.clear()
-
-
-def decode_program(data: bytes) -> Program:
-  result: Program = {}
-  i = 0
-  while i < len(data):
-    try: inst_class = detect_format(data[i:])
-    except ValueError: break  # stop at invalid instruction (padding/metadata after code)
-    if inst_class is None: i += 4; continue
-    base_size = inst_class._size()
-    # Pass enough data for potential 64-bit literal (base + 8 bytes max)
-    inst = inst_class.from_bytes(data[i:i+base_size+8])
-    for name, val in inst._values.items():
-      if name != 'op': setattr(inst, name, unwrap(val))  # skip op to preserve property access
-    inst._words = inst.size() // 4
-    result[i // 4] = inst
-    i += inst._words * 4
-  return result
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# EXECUTION - All ALU ops use pseudocode from PDF
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def exec_scalar(st: WaveState, inst: Inst) -> int:
-  """Execute scalar instruction. Returns PC delta or negative for special cases."""
-  compiled = _get_compiled()
-
-  # SOPP: special cases for control flow that has no pseudocode
-  if isinstance(inst, SOPP):
-    if inst.op == SOPPOp.S_ENDPGM: return -1
-    if inst.op == SOPPOp.S_BARRIER: return -2
-
-  # SMEM: memory loads (not ALU)
-  if isinstance(inst, SMEM):
-    addr = st.rsgpr64(inst.sbase * 2) + _sext(inst.offset, 21)
-    if inst.soffset not in (NULL, 0x7f): addr += st.rsrc(inst.soffset, 0)
-    if (cnt := SMEM_LOAD.get(inst.op)) is None: raise NotImplementedError(f"SMEM op {inst.op}")
-    for i in range(cnt): st.wsgpr(inst.sdata + i, mem_read((addr + i * 4) & MASK64, 4))
-    return 0
-
-  # Get op enum and lookup compiled function
-  if isinstance(inst, SOP1): ssrc0, sdst = inst.ssrc0, inst.sdst
-  elif isinstance(inst, SOP2): ssrc0, sdst = inst.ssrc0, inst.sdst
-  elif isinstance(inst, SOPC): ssrc0, sdst = inst.ssrc0, None
-  elif isinstance(inst, SOPK): ssrc0, sdst = inst.sdst, inst.sdst  # sdst is both src and dst
-  elif isinstance(inst, SOPP): ssrc0, sdst = None, None
-  else: raise NotImplementedError(f"Unknown scalar type {type(inst)}")
-
-  # SOPP has gaps in the opcode enum - treat unknown opcodes as no-ops
-  try: op = inst.op
-  except ValueError:
-    if isinstance(inst, SOPP): return 0
-    raise
-  fn = compiled.get(type(op), {}).get(op)
-  if fn is None:
-    # SOPP instructions without pseudocode (waits, hints, nops) are no-ops
-    if isinstance(inst, SOPP): return 0
-    raise NotImplementedError(f"{op.name} not in pseudocode")
-
-  # Build context - use inst methods to determine operand sizes
-  s0 = st.rsrc64(ssrc0, 0) if inst.is_src_64(0) else (st.rsrc(ssrc0, 0) if not isinstance(inst, (SOPK, SOPP)) else (st.rsgpr(inst.sdst) if isinstance(inst, SOPK) else 0))
-  s1 = st.rsrc64(inst.ssrc1, 0) if inst.is_src_64(1) else (st.rsrc(inst.ssrc1, 0) if isinstance(inst, (SOP2, SOPC)) else inst.simm16 if isinstance(inst, SOPK) else 0)
-  d0 = st.rsgpr64(sdst) if inst.dst_regs() == 2 and sdst is not None else (st.rsgpr(sdst) if sdst is not None else 0)
-  literal = inst.simm16 if isinstance(inst, (SOPK, SOPP)) else st.literal
-
-  # Create Reg objects for compiled function - mask VCC/EXEC to 32 bits for wave32
-  result = fn(Reg(s0), Reg(s1), None, Reg(d0), Reg(st.scc), Reg(st.vcc & MASK32), 0, Reg(st.exec_mask & MASK32), literal, None, PC=Reg(st.pc * 4))
-
-  # Apply results - extract values from returned Reg objects
-  if sdst is not None and 'D0' in result:
-    (st.wsgpr64 if inst.dst_regs() == 2 else st.wsgpr)(sdst, result['D0']._val)
-  if 'SCC' in result: st.scc = result['SCC']._val & 1
-  if 'EXEC' in result: st.exec_mask = result['EXEC']._val
-  if 'PC' in result:
-    # Convert absolute byte address to word delta
-    pc_val = result['PC']._val
-    new_pc = pc_val if pc_val < 0x8000000000000000 else pc_val - 0x10000000000000000
-    new_pc_words = new_pc // 4
-    return new_pc_words - st.pc - 1  # -1 because emulator adds inst_words (1 for scalar)
-  return 0
-
-def exec_vector(st: WaveState, inst: Inst, lane: int, lds: bytearray | None = None) -> None:
-  """Execute vector instruction for one lane."""
-  compiled = _get_compiled()
-  V = st.vgpr[lane]
-
-  # Memory ops (not ALU pseudocode)
-  if isinstance(inst, FLAT):
-    op, addr_reg, data_reg, vdst, offset, saddr = inst.op, inst.addr, inst.data, inst.vdst, _sext(inst.offset, 13), inst.saddr
-    addr = V[addr_reg] | (V[addr_reg+1] << 32)
-    addr = (st.rsgpr64(saddr) + V[addr_reg] + offset) & MASK64 if saddr not in (NULL, 0x7f) else (addr + offset) & MASK64
-    if op in FLAT_LOAD:
-      cnt, sz, sign = FLAT_LOAD[op]
-      for i in range(cnt): val = mem_read(addr + i * sz, sz); V[vdst + i] = _sext(val, sz * 8) & MASK32 if sign else val
-    elif op in FLAT_STORE:
-      cnt, sz = FLAT_STORE[op]
-      for i in range(cnt): mem_write(addr + i * sz, sz, V[data_reg + i] & ((1 << (sz * 8)) - 1))
-    elif op in FLAT_D16_LOAD:
-      sz, sign, hi = FLAT_D16_LOAD[op]
-      val = mem_read(addr, sz)
-      if sign: val = _sext(val, sz * 8) & 0xffff
-      V[vdst] = _dst16(V[vdst], val, hi)
-    elif op in FLAT_D16_STORE:
-      sz, hi = FLAT_D16_STORE[op]
-      mem_write(addr, sz, _src16(V[data_reg], hi) & ((1 << (sz * 8)) - 1))
-    else: raise NotImplementedError(f"FLAT op {op}")
-    return
-
-  if isinstance(inst, DS):
-    fn = compiled.get(DSOp, {}).get(inst.op)
-    if fn is None: raise NotImplementedError(f"DS op {inst.op.name} not in pseudocode")
-    # Prepare data registers as lists of dwords
-    data0 = [V[inst.data0 + i] for i in range(4)]  # up to 4 dwords
-    data1 = [V[inst.data1 + i] for i in range(4)] if inst.data1 else [0, 0, 0, 0]
-    result = fn(lds, V[inst.addr], data0, data1, inst.vdst, inst.offset0, inst.offset1)
-    # Write results for loads
-    if 'vdst' in result:
-      for i, val in enumerate(result['vdst']): V[inst.vdst + i] = val & MASK32
-    return
-
-  # VOPD: dual-issue, execute two ops simultaneously (read all inputs before writes)
-  if isinstance(inst, VOPD):
-    vdsty = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1)
-    inputs = [(inst.opx, st.rsrc(inst.srcx0, lane), V[inst.vsrcx1], V[inst.vdstx], inst.vdstx),
-              (inst.opy, st.rsrc(inst.srcy0, lane), V[inst.vsrcy1], V[vdsty], vdsty)]
-    def exec_vopd(vopd_op, s0, s1, d0):
-      op = _VOPD_TO_VOP[vopd_op]
-      return compiled[type(op)][op](Reg(s0), Reg(s1), None, Reg(d0), Reg(st.scc), Reg(st.vcc), lane, Reg(st.exec_mask), st.literal, None)['D0']._val
-    for vopd_op, s0, s1, d0, dst in inputs: V[dst] = exec_vopd(vopd_op, s0, s1, d0)
-    return
-
-  # VOP3SD: has extra scalar dest for carry output
-  if isinstance(inst, VOP3SD):
-    fn = compiled[VOP3SDOp][inst.op]
-    # Read sources based on register counts from inst properties
-    def rsrc_n(src, regs): return st.rsrc64(src, lane) if regs == 2 else st.rsrc(src, lane)
-    s0, s1, s2 = rsrc_n(inst.src0, inst.src_regs(0)), rsrc_n(inst.src1, inst.src_regs(1)), rsrc_n(inst.src2, inst.src_regs(2))
-    # Carry-in ops use src2 as carry bitmask instead of VCC
-    vcc = st.rsgpr64(inst.src2) if 'CO_CI' in inst.op_name else st.vcc
-    result = fn(Reg(s0), Reg(s1), Reg(s2), Reg(V[inst.vdst]), Reg(st.scc), Reg(vcc), lane, Reg(st.exec_mask), st.literal, None)
-    d0_val = result['D0']._val
-    V[inst.vdst] = d0_val & MASK32
-    if inst.dst_regs() == 2: V[inst.vdst + 1] = (d0_val >> 32) & MASK32
-    if 'VCC' in result: st.pend_sgpr_lane(inst.sdst, lane, (result['VCC']._val >> lane) & 1)
-    return
-
-  # Get op enum and sources (None means "no source" for that operand)
-  # dst_hi: for VOP1/VOP2 16-bit dst ops, bit 7 of vdst indicates .h (high 16-bit) destination
-  dst_hi = False
-  if isinstance(inst, VOP1):
-    if inst.op == VOP1Op.V_NOP: return
-    src0, src1, src2 = inst.src0, None, None
-    dst_hi = (inst.vdst & 0x80) != 0 and inst.is_dst_16()
-    vdst = inst.vdst & 0x7f if inst.is_dst_16() else inst.vdst
-  elif isinstance(inst, VOP2):
-    src0, src1, src2 = inst.src0, inst.vsrc1 + 256, None
-    dst_hi = (inst.vdst & 0x80) != 0 and inst.is_dst_16()
-    vdst = inst.vdst & 0x7f if inst.is_dst_16() else inst.vdst
-  elif isinstance(inst, VOP3):
-    # VOP3 ops 0-255 are VOPC comparisons encoded as VOP3 - inst.op returns VOPCOp for these
-    src0, src1, src2, vdst = inst.src0, inst.src1, (None if inst.op.value < 256 else inst.src2), inst.vdst
-  elif isinstance(inst, VOPC):
-    # For 16-bit VOPC, vsrc1 uses same encoding as VOP2 16-bit: bit 7 selects hi(1) or lo(0) half
-    # vsrc1 field is 8 bits: [6:0] = VGPR index, [7] = hi flag
-    src0, src1, src2, vdst = inst.src0, inst.vsrc1 + 256, None, VCC_LO
-  elif isinstance(inst, VOP3P):
-    # VOP3P: Packed 16-bit operations using compiled functions
-    # WMMA: wave-level matrix multiply-accumulate (special handling - needs cross-lane access)
-    if 'WMMA' in inst.op_name:
-      if lane == 0:  # Only execute once per wave, write results for all lanes
-        exec_wmma(st, inst, inst.op)
-      return
-    # V_FMA_MIX: Mixed precision FMA - opsel_hi controls f32(0) vs f16(1), opsel selects which f16 half
-    # Handle inline because abs/neg must be applied AFTER type conversion
-    if inst.op in (VOP3POp.V_FMA_MIX_F32, VOP3POp.V_FMA_MIXLO_F16, VOP3POp.V_FMA_MIXHI_F16):
-      opsel, opsel_hi, opsel_hi2 = getattr(inst, 'opsel', 0), getattr(inst, 'opsel_hi', 0), getattr(inst, 'opsel_hi2', 0)
-      neg, abs_ = getattr(inst, 'neg', 0), getattr(inst, 'neg_hi', 0)  # neg_hi reused as abs for FMA_MIX
-      raws = [st.rsrc(inst.src0, lane), st.rsrc(inst.src1, lane), st.rsrc(inst.src2, lane) if inst.src2 is not None else 0]
-      is_f16 = [opsel_hi & 1, opsel_hi & 2, opsel_hi2]
-      srcs = [_f16(_src16(raws[i], bool(opsel & (1<<i)))) if is_f16[i] else _f32(raws[i]) for i in range(3)]
-      for i in range(3):
-        if abs_ & (1<<i): srcs[i] = abs(srcs[i])
-        if neg & (1<<i): srcs[i] = -srcs[i]
-      result_f = srcs[0] * srcs[1] + srcs[2]
-      V = st.vgpr[lane]
-      V[inst.vdst] = _i32(result_f) if inst.op == VOP3POp.V_FMA_MIX_F32 else _dst16(V[inst.vdst], _i16(result_f), inst.op == VOP3POp.V_FMA_MIXHI_F16)
-      return
-    # VOP3P packed ops: opsel selects halves for lo, opsel_hi for hi; neg toggles f16 sign
-    raws = [st.rsrc_f16(inst.src0, lane), st.rsrc_f16(inst.src1, lane), st.rsrc_f16(inst.src2, lane) if inst.src2 is not None else 0]
-    opsel, opsel_hi, opsel_hi2 = getattr(inst, 'opsel', 0), getattr(inst, 'opsel_hi', 3), getattr(inst, 'opsel_hi2', 1)
-    neg, neg_hi = getattr(inst, 'neg', 0), getattr(inst, 'neg_hi', 0)
-    hi_sels = [opsel_hi & 1, opsel_hi & 2, opsel_hi2]
-    srcs = [((_src16(raws[i], hi_sels[i]) ^ (0x8000 if neg_hi & (1<<i) else 0)) << 16) |
-            (_src16(raws[i], opsel & (1<<i)) ^ (0x8000 if neg & (1<<i) else 0)) for i in range(3)]
-    result = compiled[VOP3POp][inst.op](Reg(srcs[0]), Reg(srcs[1]), Reg(srcs[2]), Reg(0), Reg(st.scc), Reg(st.vcc), lane, Reg(st.exec_mask), st.literal, None)
-    st.vgpr[lane][inst.vdst] = result['D0']._val & MASK32
-    return
-  else: raise NotImplementedError(f"Unknown vector type {type(inst)}")
-
-  op_cls = type(inst.op)
-  if (fn := compiled.get(op_cls, {}).get(inst.op)) is None: raise NotImplementedError(f"{inst.op_name} not in pseudocode")
-
-  # Read sources (with VOP3 modifiers if applicable)
-  neg, abs_ = (getattr(inst, 'neg', 0), getattr(inst, 'abs', 0)) if isinstance(inst, VOP3) else (0, 0)
-  opsel = getattr(inst, 'opsel', 0) if isinstance(inst, VOP3) else 0
-  def mod_src(val: int, idx: int, is64=False) -> int:
-    to_f, to_i = (_f64, _i64) if is64 else (_f32, _i32)
-    if (abs_ >> idx) & 1: val = to_i(abs(to_f(val)))
-    if (neg >> idx) & 1: val = to_i(-to_f(val))
-    return val
-
-  # Use inst methods to determine operand sizes (inst.is_src_16, inst.is_src_64, etc.)
-  is_vop2_16bit = isinstance(inst, VOP2) and inst.is_16bit()
-
-  # Read sources based on register counts and dtypes from inst properties
-  def read_src(src, idx, regs, is_src_16):
-    if src is None: return 0
-    if regs == 2: return mod_src(st.rsrc64(src, lane), idx, is64=True)
-    if is_src_16 and isinstance(inst, VOP3):
-      raw = st.rsrc_f16(src, lane) if 128 <= src < 255 else st.rsrc(src, lane)
-      val = _src16(raw, bool(opsel & (1 << idx)))
-      if abs_ & (1 << idx): val &= 0x7fff
-      if neg & (1 << idx): val ^= 0x8000
-      return val
-    if is_src_16 and isinstance(inst, (VOP1, VOP2, VOPC)):
-      if src >= 256: return _src16(mod_src(st.rsrc(_vgpr_masked(src), lane), idx), _vgpr_hi(src))
-      return mod_src(st.rsrc_f16(src, lane), idx) & 0xffff
-    return mod_src(st.rsrc(src, lane), idx)
-
-  s0 = read_src(src0, 0, inst.src_regs(0), inst.is_src_16(0))
-  s1 = read_src(src1, 1, inst.src_regs(1), inst.is_src_16(1)) if src1 is not None else 0
-  s2 = read_src(src2, 2, inst.src_regs(2), inst.is_src_16(2)) if src2 is not None else 0
-  # Read destination (accumulator for VOP2 f16, 64-bit for 64-bit ops)
-  d0 = _src16(V[vdst], dst_hi) if is_vop2_16bit else (V[vdst] | (V[vdst + 1] << 32)) if inst.dst_regs() == 2 else V[vdst]
-
-  # V_CNDMASK_B32/B16: VOP3 encoding uses src2 as mask (not VCC); VOP2 uses VCC implicitly
-  # Pass the correct mask as vcc to the function so pseudocode VCC.u64[laneId] works correctly
-  vcc_for_fn = st.rsgpr64(src2) if inst.op in (VOP3Op.V_CNDMASK_B32, VOP3Op.V_CNDMASK_B16) and isinstance(inst, VOP3) and src2 is not None and src2 < 256 else st.vcc
-
-  # Execute compiled function - pass src0_idx and vdst_idx for lane instructions
-  # For VGPR access: src0 index is the VGPR number (src0 - 256 if VGPR, else src0 for SGPR)
-  src0_idx = (src0 - 256) if src0 is not None and src0 >= 256 else (src0 if src0 is not None else 0)
-  result = fn(Reg(s0), Reg(s1), Reg(s2), Reg(d0), Reg(st.scc), Reg(vcc_for_fn), lane, Reg(st.exec_mask), st.literal, st.vgpr, src0_idx, vdst)
-
-  # Apply results - extract values from returned Reg objects
-  if 'vgpr_write' in result:
-    # Lane instruction wrote to VGPR: (lane, vgpr_idx, value)
-    wr_lane, wr_idx, wr_val = result['vgpr_write']
-    st.vgpr[wr_lane][wr_idx] = wr_val
-  if 'VCC' in result:
-    # VOP2 carry ops write to VCC implicitly; VOPC/VOP3 write to vdst
-    st.pend_sgpr_lane(VCC_LO if isinstance(inst, VOP2) and 'CO_CI' in inst.op_name else vdst, lane, (result['VCC']._val >> lane) & 1)
-  if 'EXEC' in result:
-    # V_CMPX instructions write to EXEC per-lane (not to vdst)
-    st.pend_sgpr_lane(EXEC_LO, lane, (result['EXEC']._val >> lane) & 1)
-  elif op_cls is VOPCOp:
-    # VOPC comparison result stored in D0 bitmask, extract lane bit (non-CMPX only)
-    st.pend_sgpr_lane(vdst, lane, (result['D0']._val >> lane) & 1)
-  if op_cls is not VOPCOp and 'vgpr_write' not in result:
-    writes_to_sgpr = 'READFIRSTLANE' in inst.op_name or 'READLANE' in inst.op_name
-    d0_val = result['D0']._val
-    if writes_to_sgpr: st.wsgpr(vdst, d0_val & MASK32)
-    elif inst.dst_regs() == 2: V[vdst], V[vdst + 1] = d0_val & MASK32, (d0_val >> 32) & MASK32
-    elif inst.is_dst_16(): V[vdst] = _dst16(V[vdst], d0_val, bool(opsel & 8) if isinstance(inst, VOP3) else dst_hi)
-    else: V[vdst] = d0_val & MASK32
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# WMMA (Wave Matrix Multiply-Accumulate)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def exec_wmma(st: WaveState, inst, op: VOP3POp) -> None:
-  """Execute WMMA instruction - 16x16x16 matrix multiply across the wave."""
-  src0, src1, src2, vdst = inst.src0, inst.src1, inst.src2, inst.vdst
-  # Read 16x16 f16 matrix from 16 lanes × 8 VGPRs (2 f16 per VGPR)
-  def read_f16_mat(src):
-    return [f for l in range(16) for r in range(8) for v in [st.vgpr[l][src-256+r] if src >= 256 else st.rsgpr(src+r)] for f in [_f16(v&0xffff), _f16((v>>16)&0xffff)]]
-  mat_a, mat_b = read_f16_mat(src0), read_f16_mat(src1)
-  # Read matrix C (16x16 f32) from lanes 0-31, VGPRs src2 to src2+7
-  mat_c = [_f32(st.vgpr[i % 32][src2 - 256 + i // 32] if src2 >= 256 else st.rsgpr(src2 + i // 32)) for i in range(256)]
-  # Compute D = A × B + C (16x16 matrix multiply)
-  mat_d = [sum(mat_a[row*16+k] * mat_b[col*16+k] for k in range(16)) + mat_c[row*16+col] for row in range(16) for col in range(16)]
-  # Write result - f16 packed or f32
-  if op == VOP3POp.V_WMMA_F16_16X16X16_F16:
-    for i in range(0, 256, 2):
-      st.vgpr[(i//2) % 32][vdst + (i//2)//32] = ((_i16(mat_d[i+1]) & 0xffff) << 16) | (_i16(mat_d[i]) & 0xffff)
-  else:
-    for i in range(256): st.vgpr[i % 32][vdst + i//32] = _i32(mat_d[i])
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# MAIN EXECUTION LOOP
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def step_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> int:
-  inst = program.get(st.pc)
-  if inst is None: return 1
-  inst_words, st.literal = inst._words, getattr(inst, '_literal', None) or 0
-
-  if isinstance(inst, (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM)):
-    delta = exec_scalar(st, inst)
-    if delta == -1: return -1  # endpgm
-    if delta == -2: st.pc += inst_words; return -2  # barrier
-    st.pc += inst_words + delta
-  else:
-    # V_READFIRSTLANE/V_READLANE write to SGPR, execute once; others execute per-lane with exec_mask
-    is_readlane = isinstance(inst, (VOP1, VOP3)) and ('READFIRSTLANE' in inst.op_name or 'READLANE' in inst.op_name)
-    exec_mask = 1 if is_readlane else st.exec_mask
-    for lane in range(1 if is_readlane else n_lanes):
-      if exec_mask & (1 << lane): exec_vector(st, inst, lane, lds)
-    st.commit_pends()
-    st.pc += inst_words
-  return 0
-
-def exec_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> int:
-  while st.pc in program:
-    result = step_wave(program, st, lds, n_lanes)
-    if result == -1: return 0
-    if result == -2: return -2
-  return 0
-
-def exec_workgroup(program: Program, workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int,
-                   wg_id_sgpr_base: int, wg_id_enables: tuple[bool, bool, bool]) -> None:
-  lx, ly, lz = local_size
-  total_threads, lds = lx * ly * lz, bytearray(65536)
-  waves: list[tuple[WaveState, int, int]] = []
-  for wave_start in range(0, total_threads, WAVE_SIZE):
-    n_lanes, st = min(WAVE_SIZE, total_threads - wave_start), WaveState()
-    st.exec_mask = (1 << n_lanes) - 1
-    st.wsgpr64(0, args_ptr)
-    # Set workgroup IDs in SGPRs based on USER_SGPR_COUNT and enable flags from COMPUTE_PGM_RSRC2
-    sgpr_idx = wg_id_sgpr_base
-    for wg_id, enabled in zip(workgroup_id, wg_id_enables):
-      if enabled: st.sgpr[sgpr_idx] = wg_id; sgpr_idx += 1
-    # Set workitem IDs in VGPR0 using packed method: v0 = (Z << 20) | (Y << 10) | X
-    for i in range(n_lanes):
-      tid = wave_start + i
-      st.vgpr[i][0] = ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
-    waves.append((st, n_lanes, wave_start))
-  has_barrier = any(isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER for inst in program.values())
-  for _ in range(2 if has_barrier else 1):
-    for st, n_lanes, _ in waves: exec_wave(program, st, lds, n_lanes)
-
-def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int, rsrc2: int = 0x19c) -> int:
-  program = decode_program((ctypes.c_char * lib_sz).from_address(lib).raw)
-  if not program: return -1
-  wg_id_enables = tuple(bool((rsrc2 >> (7+i)) & 1) for i in range(3))
-  for gidz in range(gz):
-    for gidy in range(gy):
-      for gidx in range(gx): exec_workgroup(program, (gidx, gidy, gidz), (lx, ly, lz), args_ptr, (rsrc2 >> 1) & 0x1f, wg_id_enables)
-  return 0
@@ -1,542 +0,0 @@
-# DSL for RDNA3 pseudocode - makes pseudocode expressions work directly as Python
-import struct, math
-from extra.assembly.amd.dsl import MASK32, MASK64, _f32, _i32, _sext, _f16, _i16, _f64, _i64
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# HELPER FUNCTIONS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def _div(a, b):
-  try: return a / b
-  except ZeroDivisionError:
-    if a == 0.0 or math.isnan(a): return float("nan")
-    return math.copysign(float("inf"), a * b) if b == 0.0 else float("inf") if a > 0 else float("-inf")
-def _to_f16_bits(v): return v if isinstance(v, int) else _i16(v)
-def _isnan(x):
-  try: return math.isnan(float(x))
-  except (TypeError, ValueError): return False
-def _check_nan_type(x, quiet_bit_expected, default):
-  """Check NaN type by examining quiet bit. Returns default if can't determine."""
-  try:
-    if not math.isnan(float(x)): return False
-    if hasattr(x, '_reg') and hasattr(x, '_bits'):
-      bits = x._reg._val & ((1 << x._bits) - 1)
-      # NaN format: exponent all 1s, quiet bit, mantissa != 0
-      # f16: exp[14:10]=31, quiet=bit9, mant[8:0]  |  f32: exp[30:23]=255, quiet=bit22, mant[22:0]  |  f64: exp[62:52]=2047, quiet=bit51, mant[51:0]
-      exp_bits, quiet_pos, mant_mask = {16: (0x1f, 9, 0x3ff), 32: (0xff, 22, 0x7fffff), 64: (0x7ff, 51, 0xfffffffffffff)}.get(x._bits, (0,0,0))
-      exp_shift = {16: 10, 32: 23, 64: 52}.get(x._bits, 0)
-      if exp_bits and ((bits >> exp_shift) & exp_bits) == exp_bits and (bits & mant_mask) != 0:
-        return ((bits >> quiet_pos) & 1) == quiet_bit_expected
-    return default
-  except (TypeError, ValueError): return False
-def _isquietnan(x): return _check_nan_type(x, 1, True)  # quiet NaN has quiet bit = 1
-def _issignalnan(x): return _check_nan_type(x, 0, False)  # signaling NaN has quiet bit = 0
-def _gt_neg_zero(a, b): return (a > b) or (a == 0 and b == 0 and not math.copysign(1, a) < 0 and math.copysign(1, b) < 0)
-def _lt_neg_zero(a, b): return (a < b) or (a == 0 and b == 0 and math.copysign(1, a) < 0 and not math.copysign(1, b) < 0)
-def _fma(a, b, c): return a * b + c
-def _signext(v): return v
-def _fpop(fn): return lambda x: (x := float(x), x if math.isnan(x) or math.isinf(x) else float(fn(x)))[1]
-trunc, floor, ceil = _fpop(math.trunc), _fpop(math.floor), _fpop(math.ceil)
-class _SafeFloat(float):
-  """Float subclass that uses _div for division to handle 0/inf correctly."""
-  def __truediv__(self, o): return _div(float(self), float(o))
-  def __rtruediv__(self, o): return _div(float(o), float(self))
-def sqrt(x): return _SafeFloat(math.sqrt(x)) if x >= 0 else _SafeFloat(float("nan"))
-def log2(x): return math.log2(x) if x > 0 else (float("-inf") if x == 0 else float("nan"))
-i32_to_f32 = u32_to_f32 = i32_to_f64 = u32_to_f64 = f32_to_f64 = f64_to_f32 = float
-def _f_to_int(f, lo, hi): f = float(f); return 0 if math.isnan(f) else (hi if f >= hi else lo if f <= lo else int(f))
-def f32_to_i32(f): return _f_to_int(f, -2147483648, 2147483647)
-def f32_to_u32(f): return _f_to_int(f, 0, 4294967295)
-f64_to_i32, f64_to_u32 = f32_to_i32, f32_to_u32
-def f32_to_f16(f):
-  f = float(f)
-  if math.isnan(f): return 0x7e00  # f16 NaN
-  if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00  # f16 ±infinity
-  try: return struct.unpack("<H", struct.pack("<e", f))[0]
-  except OverflowError: return 0x7c00 if f > 0 else 0xfc00  # overflow -> ±infinity
-def _f16_to_f32_bits(bits): return struct.unpack("<e", struct.pack("<H", int(bits) & 0xffff))[0]
-def f16_to_f32(v): return v if isinstance(v, float) else _f16_to_f32_bits(v)
-def i16_to_f16(v): return f32_to_f16(float(_sext(int(v) & 0xffff, 16)))
-def u16_to_f16(v): return f32_to_f16(float(int(v) & 0xffff))
-def f16_to_i16(bits): f = _f16_to_f32_bits(bits); return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
-def f16_to_u16(bits): f = _f16_to_f32_bits(bits); return max(0, min(65535, int(f))) if not math.isnan(f) else 0
-def u8_to_u32(v): return int(v) & 0xff
-def u4_to_u32(v): return int(v) & 0xf
-def _sign(f): return 1 if math.copysign(1.0, f) < 0 else 0
-def _mantissa_f32(f): return struct.unpack("<I", struct.pack("<f", f))[0] & 0x7fffff if not (math.isinf(f) or math.isnan(f)) else 0
-def _ldexp(m, e): return math.ldexp(m, e)
-def isEven(x):
-  x = float(x)
-  if math.isinf(x) or math.isnan(x): return False
-  return int(x) % 2 == 0
-def fract(x): return x - math.floor(x)
-PI = math.pi
-def _trig(fn, x):
-  # V_SIN/COS_F32: hardware does frac on input cycles before computing
-  if math.isinf(x) or math.isnan(x): return float("nan")
-  frac_cycles = fract(x / (2 * math.pi))
-  return fn(frac_cycles * 2 * math.pi)
-def sin(x): return _trig(math.sin, x)
-def cos(x): return _trig(math.cos, x)
-def pow(a, b):
-  try: return a ** b
-  except OverflowError: return float("inf") if b > 0 else 0.0
-def _brev(v, bits): return int(bin(v & ((1 << bits) - 1))[2:].zfill(bits)[::-1], 2)
-def _brev32(v): return _brev(v, 32)
-def _brev64(v): return _brev(v, 64)
-def _ctz(v, bits):
-  v, n = int(v) & ((1 << bits) - 1), 0
-  if v == 0: return bits
-  while (v & 1) == 0: v >>= 1; n += 1
-  return n
-def _ctz32(v): return _ctz(v, 32)
-def _ctz64(v): return _ctz(v, 64)
-def _exponent(f):
-  # Handle TypedView (f16/f32/f64) to get correct exponent for that type
-  if hasattr(f, '_bits') and hasattr(f, '_float') and f._float:
-    raw = f._val
-    if f._bits == 16: return (raw >> 10) & 0x1f  # f16: 5-bit exponent
-    if f._bits == 32: return (raw >> 23) & 0xff  # f32: 8-bit exponent
-    if f._bits == 64: return (raw >> 52) & 0x7ff  # f64: 11-bit exponent
-  # Fallback: convert to f32 and get exponent
-  f = float(f)
-  if math.isinf(f) or math.isnan(f): return 255
-  if f == 0.0: return 0
-  try: bits = struct.unpack("<I", struct.pack("<f", f))[0]; return (bits >> 23) & 0xff
-  except: return 0
-def _is_denorm_f32(f):
-  if not isinstance(f, float): f = _f32(int(f) & 0xffffffff)
-  if math.isinf(f) or math.isnan(f) or f == 0.0: return False
-  bits = struct.unpack("<I", struct.pack("<f", float(f)))[0]
-  return (bits >> 23) & 0xff == 0
-def _is_denorm_f64(f):
-  if not isinstance(f, float): f = _f64(int(f) & 0xffffffffffffffff)
-  if math.isinf(f) or math.isnan(f) or f == 0.0: return False
-  bits = struct.unpack("<Q", struct.pack("<d", float(f)))[0]
-  return (bits >> 52) & 0x7ff == 0
-def v_min_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _lt_neg_zero(a, b) else b)
-def v_max_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _gt_neg_zero(a, b) else b)
-v_min_f16, v_max_f16 = v_min_f32, v_max_f32
-v_min_i32, v_max_i32 = min, max
-v_min_i16, v_max_i16 = min, max
-def v_min_u32(a, b): return min(a & MASK32, b & MASK32)
-def v_max_u32(a, b): return max(a & MASK32, b & MASK32)
-def v_min_u16(a, b): return min(a & 0xffff, b & 0xffff)
-def v_max_u16(a, b): return max(a & 0xffff, b & 0xffff)
-def v_min3_f32(a, b, c): return v_min_f32(v_min_f32(a, b), c)
-def v_max3_f32(a, b, c): return v_max_f32(v_max_f32(a, b), c)
-v_min3_f16, v_max3_f16 = v_min3_f32, v_max3_f32
-v_min3_i32, v_max3_i32, v_min3_i16, v_max3_i16 = min, max, min, max
-def v_min3_u32(a, b, c): return min(a & MASK32, b & MASK32, c & MASK32)
-def v_max3_u32(a, b, c): return max(a & MASK32, b & MASK32, c & MASK32)
-def v_min3_u16(a, b, c): return min(a & 0xffff, b & 0xffff, c & 0xffff)
-def v_max3_u16(a, b, c): return max(a & 0xffff, b & 0xffff, c & 0xffff)
-def ABSDIFF(a, b): return abs(int(a) - int(b))
-
-# BF16 (bfloat16) conversion functions
-def _bf16(i):
-  """Convert bf16 bits to float. BF16 is just the top 16 bits of f32."""
-  return struct.unpack("<f", struct.pack("<I", (i & 0xffff) << 16))[0]
-def _ibf16(f):
-  """Convert float to bf16 bits (truncate to top 16 bits of f32)."""
-  if math.isnan(f): return 0x7fc0  # bf16 quiet NaN
-  if math.isinf(f): return 0x7f80 if f > 0 else 0xff80  # bf16 ±infinity
-  try: return (struct.unpack("<I", struct.pack("<f", float(f)))[0] >> 16) & 0xffff
-  except (OverflowError, struct.error): return 0x7f80 if f > 0 else 0xff80
-def bf16_to_f32(v): return _bf16(v) if isinstance(v, int) else float(v)
-def f32_to_bf16(f): return _ibf16(f)
-
-# BYTE_PERMUTE for V_PERM_B32 - select bytes from 64-bit data based on selector
-def BYTE_PERMUTE(data, sel):
-  """Select a byte from 64-bit data based on selector value.
-  sel 0-7: select byte from data (S1 is bytes 0-3, S0 is bytes 4-7 in {S0,S1})
-  sel 8-11: sign-extend from specific bytes (8->byte1, 9->byte3, 10->byte5, 11->byte7)
-  sel 12: constant 0x00
-  sel >= 13: constant 0xFF"""
-  sel = int(sel) & 0xff
-  if sel <= 7: return (int(data) >> (sel * 8)) & 0xff
-  if sel == 8: return 0xff if ((int(data) >> 15) & 1) else 0x00  # sign of byte 1
-  if sel == 9: return 0xff if ((int(data) >> 31) & 1) else 0x00  # sign of byte 3
-  if sel == 10: return 0xff if ((int(data) >> 47) & 1) else 0x00  # sign of byte 5
-  if sel == 11: return 0xff if ((int(data) >> 63) & 1) else 0x00  # sign of byte 7
-  if sel == 12: return 0x00
-  return 0xff  # sel >= 13
-
-# v_sad_u8 helper for V_SAD instructions (sum of absolute differences of 4 bytes)
-def v_sad_u8(s0, s1, s2):
-  """V_SAD_U8: Sum of absolute differences of 4 byte pairs plus accumulator."""
-  s0, s1, s2 = int(s0), int(s1), int(s2)
-  result = s2
-  for i in range(4):
-    a = (s0 >> (i * 8)) & 0xff
-    b = (s1 >> (i * 8)) & 0xff
-    result += abs(a - b)
-  return result & 0xffffffff
-
-# v_msad_u8 helper (masked SAD - skip when reference byte is 0)
-def v_msad_u8(s0, s1, s2):
-  """V_MSAD_U8: Masked sum of absolute differences (skip if reference byte is 0)."""
-  s0, s1, s2 = int(s0), int(s1), int(s2)
-  result = s2
-  for i in range(4):
-    a = (s0 >> (i * 8)) & 0xff
-    b = (s1 >> (i * 8)) & 0xff
-    if b != 0:  # Only add diff if reference (s1) byte is non-zero
-      result += abs(a - b)
-  return result & 0xffffffff
-def f16_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
-def f16_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
-def f32_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
-def f32_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
-def v_cvt_i16_f32(f): return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
-def v_cvt_u16_f32(f): return max(0, min(65535, int(f))) if not math.isnan(f) else 0
-def u32_to_u16(u): return int(u) & 0xffff
-def i32_to_i16(i): return ((int(i) + 32768) & 0xffff) - 32768
-def SAT8(v): return max(0, min(255, int(v)))
-def f32_to_u8(f): return max(0, min(255, int(f))) if not math.isnan(f) else 0
-def mantissa(f):
-  if f == 0.0 or math.isinf(f) or math.isnan(f): return f
-  m, _ = math.frexp(f)
-  return m  # AMD V_FREXP_MANT returns mantissa in [0.5, 1.0) range
-def signext_from_bit(val, bit):
-  bit = int(bit)
-  if bit == 0: return 0
-  mask = (1 << bit) - 1
-  val = int(val) & mask
-  if val & (1 << (bit - 1)): return val - (1 << bit)
-  return val
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# DSL EXPORTS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-__all__ = [
-  # Classes
-  'Reg', 'SliceProxy', 'TypedView',
-  # Pack functions
-  '_pack', '_pack32', 'pack', 'pack32',
-  # Constants
-  'WAVE32', 'WAVE64', 'MASK32', 'MASK64', 'WAVE_MODE', 'DENORM', 'OVERFLOW_F32', 'UNDERFLOW_F32',
-  'OVERFLOW_F64', 'UNDERFLOW_F64', 'MAX_FLOAT_F32', 'ROUND_MODE', 'cvtToQuietNAN', 'DST', 'INF', 'PI',
-  'TWO_OVER_PI_1201',
-  # Aliases for pseudocode
-  's_ff1_i32_b32', 's_ff1_i32_b64', 'GT_NEG_ZERO', 'LT_NEG_ZERO',
-  'isNAN', 'isQuietNAN', 'isSignalNAN', 'fma', 'ldexp', 'sign', 'exponent', 'F', 'signext',
-  # Conversion functions
-  '_f32', '_i32', '_f16', '_i16', '_f64', '_i64', '_sext', '_to_f16_bits', '_f16_to_f32_bits',
-  'i32_to_f32', 'u32_to_f32', 'i32_to_f64', 'u32_to_f64', 'f32_to_f64', 'f64_to_f32',
-  'f32_to_i32', 'f32_to_u32', 'f64_to_i32', 'f64_to_u32', 'f32_to_f16', 'f16_to_f32',
-  'i16_to_f16', 'u16_to_f16', 'f16_to_i16', 'f16_to_u16', 'u32_to_u16', 'i32_to_i16',
-  'f16_to_snorm', 'f16_to_unorm', 'f32_to_snorm', 'f32_to_unorm', 'v_cvt_i16_f32', 'v_cvt_u16_f32',
-  'SAT8', 'f32_to_u8', 'u8_to_u32', 'u4_to_u32',
-  # BF16 conversion functions
-  '_bf16', '_ibf16', 'bf16_to_f32', 'f32_to_bf16',
-  # Math functions
-  'trunc', 'floor', 'ceil', 'sqrt', 'log2', 'sin', 'cos', 'pow', 'fract', 'isEven', 'mantissa',
-  # Min/max functions
-  'v_min_f32', 'v_max_f32', 'v_min_i32', 'v_max_i32', 'v_min_u32', 'v_max_u32',
-  'v_min_f16', 'v_max_f16', 'v_min_i16', 'v_max_i16', 'v_min_u16', 'v_max_u16',
-  'v_min3_f32', 'v_max3_f32', 'v_min3_i32', 'v_max3_i32', 'v_min3_u32', 'v_max3_u32',
-  'v_min3_f16', 'v_max3_f16', 'v_min3_i16', 'v_max3_i16', 'v_min3_u16', 'v_max3_u16',
-  'ABSDIFF',
-  # Byte/SAD helper functions
-  'BYTE_PERMUTE', 'v_sad_u8', 'v_msad_u8',
-  # Bit manipulation
-  '_brev32', '_brev64', '_ctz32', '_ctz64', '_exponent', '_is_denorm_f32', '_is_denorm_f64',
-  '_sign', '_mantissa_f32', '_div', '_isnan', '_isquietnan', '_issignalnan', '_gt_neg_zero', '_lt_neg_zero', '_fma', '_ldexp', '_signext',
-  'signext_from_bit',
-]
-
-# Aliases used in pseudocode
-s_ff1_i32_b32, s_ff1_i32_b64 = _ctz32, _ctz64
-GT_NEG_ZERO, LT_NEG_ZERO = _gt_neg_zero, _lt_neg_zero
-isNAN = _isnan
-isQuietNAN = _isquietnan
-isSignalNAN = _issignalnan
-fma, ldexp, sign, exponent = _fma, _ldexp, _sign, _exponent
-def F(x):
-  """32'F(x) or 64'F(x) - interpret x as float. If x is int, treat as bit pattern."""
-  if isinstance(x, int): return _f32(x)  # int -> interpret as f32 bits
-  if isinstance(x, TypedView): return x  # preserve TypedView for bit-pattern checks
-  return float(x)  # already a float or float-like
-signext = lambda x: int(x)  # sign-extend to full width - already handled by Python's arbitrary precision ints
-pack = lambda hi, lo: ((int(hi) & 0xffff) << 16) | (int(lo) & 0xffff)
-pack32 = lambda hi, lo: ((int(hi) & 0xffffffff) << 32) | (int(lo) & 0xffffffff)
-_pack, _pack32 = pack, pack32  # Aliases for internal use
-WAVE32, WAVE64 = True, False
-
-# Float overflow/underflow constants
-OVERFLOW_F32 = float('inf')
-UNDERFLOW_F32 = 0.0
-OVERFLOW_F64 = float('inf')
-UNDERFLOW_F64 = 0.0
-MAX_FLOAT_F32 = 3.4028235e+38  # Largest finite float32
-
-# INF object that supports .f16/.f32/.f64 access and comparison with floats
-class _Inf:
-  f16 = f32 = f64 = float('inf')
-  def __neg__(self): return _NegInf()
-  def __pos__(self): return self
-  def __float__(self): return float('inf')
-  def __eq__(self, other): return float(other) == float('inf') if not isinstance(other, _NegInf) else False
-  def __req__(self, other): return self.__eq__(other)
-class _NegInf:
-  f16 = f32 = f64 = float('-inf')
-  def __neg__(self): return _Inf()
-  def __pos__(self): return self
-  def __float__(self): return float('-inf')
-  def __eq__(self, other): return float(other) == float('-inf') if not isinstance(other, _Inf) else False
-  def __req__(self, other): return self.__eq__(other)
-INF = _Inf()
-
-# Rounding mode placeholder
-class _RoundMode:
-  NEAREST_EVEN = 0
-ROUND_MODE = _RoundMode()
-
-# Helper functions for pseudocode
-def cvtToQuietNAN(x): return float('nan')
-DST = None  # Placeholder, will be set in context
-
-# 2/PI with 1201 bits of precision for V_TRIG_PREOP_F64
-# Computed as: int((2/pi) * 2^1201) - this is the fractional part of 2/pi scaled to integer
-# The MSB (bit 1200) corresponds to 2^0 position in the fraction 0.b1200 b1199 ... b1 b0
-_TWO_OVER_PI_1201_RAW = 0x0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6
-
-class _BigInt:
-  """Wrapper for large integers that supports bit slicing [high:low]."""
-  __slots__ = ('_val',)
-  def __init__(self, val): self._val = val
-  def __getitem__(self, key):
-    if isinstance(key, slice):
-      high, low = key.start, key.stop
-      if high < low: high, low = low, high  # Handle reversed slice
-      mask = (1 << (high - low + 1)) - 1
-      return (self._val >> low) & mask
-    return (self._val >> key) & 1
-  def __int__(self): return self._val
-  def __index__(self): return self._val
-  def __lshift__(self, n): return self._val << int(n)
-  def __rshift__(self, n): return self._val >> int(n)
-  def __and__(self, n): return self._val & int(n)
-  def __or__(self, n): return self._val | int(n)
-
-TWO_OVER_PI_1201 = _BigInt(_TWO_OVER_PI_1201_RAW)
-
-class _WaveMode:
-  IEEE = False
-WAVE_MODE = _WaveMode()
-
-class _DenormChecker:
-  """Comparator for denormalized floats. x == DENORM.f32 checks if x is denormalized."""
-  def __init__(self, bits): self._bits = bits
-  def _check(self, other):
-    return _is_denorm_f64(float(other)) if self._bits == 64 else _is_denorm_f32(float(other))
-  def __eq__(self, other): return self._check(other)
-  def __req__(self, other): return self._check(other)
-  def __ne__(self, other): return not self._check(other)
-
-class _Denorm:
-  f32 = _DenormChecker(32)
-  f64 = _DenormChecker(64)
-DENORM = _Denorm()
-
-def _brev(v, bits):
-  """Bit-reverse a value."""
-  result = 0
-  for i in range(bits): result |= ((v >> i) & 1) << (bits - 1 - i)
-  return result
-
-class SliceProxy:
-  """Proxy for D0[31:16] that supports .f16/.u16 etc getters and setters."""
-  __slots__ = ('_reg', '_high', '_low', '_reversed')
-  def __init__(self, reg, high, low):
-    self._reg = reg
-    # Handle reversed slices like [0:31] which means bit-reverse
-    if high < low: self._high, self._low, self._reversed = low, high, True
-    else: self._high, self._low, self._reversed = high, low, False
-  def _nbits(self): return self._high - self._low + 1
-  def _mask(self): return (1 << self._nbits()) - 1
-  def _get(self):
-    v = (self._reg._val >> self._low) & self._mask()
-    return _brev(v, self._nbits()) if self._reversed else v
-  def _set(self, v):
-    v = int(v)
-    if self._reversed: v = _brev(v, self._nbits())
-    self._reg._val = (self._reg._val & ~(self._mask() << self._low)) | ((v & self._mask()) << self._low)
-
-  u8 = property(lambda s: s._get() & 0xff)
-  u16 = property(lambda s: s._get() & 0xffff, lambda s, v: s._set(v))
-  u32 = property(lambda s: s._get() & MASK32, lambda s, v: s._set(v))
-  i16 = property(lambda s: _sext(s._get() & 0xffff, 16), lambda s, v: s._set(v))
-  i32 = property(lambda s: _sext(s._get() & MASK32, 32), lambda s, v: s._set(v))
-  f16 = property(lambda s: _f16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _i16(float(v))))
-  f32 = property(lambda s: _f32(s._get()), lambda s, v: s._set(_i32(float(v))))
-  bf16 = property(lambda s: _bf16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _ibf16(float(v))))
-  b16, b32 = u16, u32
-
-  def __int__(self): return self._get()
-  def __index__(self): return self._get()
-
-  # Comparison operators (compare as integers)
-  def __eq__(s, o): return s._get() == int(o)
-  def __ne__(s, o): return s._get() != int(o)
-  def __lt__(s, o): return s._get() < int(o)
-  def __le__(s, o): return s._get() <= int(o)
-  def __gt__(s, o): return s._get() > int(o)
-  def __ge__(s, o): return s._get() >= int(o)
-
-class TypedView:
-  """View for S0.u32 that supports [4:0] slicing and [bit] access."""
-  __slots__ = ('_reg', '_bits', '_signed', '_float', '_bf16')
-  def __init__(self, reg, bits, signed=False, is_float=False, is_bf16=False):
-    self._reg, self._bits, self._signed, self._float, self._bf16 = reg, bits, signed, is_float, is_bf16
-
-  @property
-  def _val(self):
-    mask = MASK64 if self._bits == 64 else MASK32 if self._bits == 32 else (1 << self._bits) - 1
-    return self._reg._val & mask
-
-  def __getitem__(self, key):
-    if isinstance(key, slice):
-      high, low = int(key.start), int(key.stop)
-      return SliceProxy(self._reg, high, low)
-    return (self._val >> int(key)) & 1
-
-  def __setitem__(self, key, value):
-    if isinstance(key, slice):
-      high, low = int(key.start), int(key.stop)
-      if high < low: high, low, value = low, high, _brev(int(value), low - high + 1)
-      mask = (1 << (high - low + 1)) - 1
-      self._reg._val = (self._reg._val & ~(mask << low)) | ((int(value) & mask) << low)
-    elif value: self._reg._val |= (1 << int(key))
-    else: self._reg._val &= ~(1 << int(key))
-
-  def __int__(self): return _sext(self._val, self._bits) if self._signed else self._val
-  def __index__(self): return int(self)
-  def __trunc__(self): return int(float(self)) if self._float else int(self)
-  def __float__(self):
-    if self._float:
-      if self._bf16: return _bf16(self._val)  # bf16 uses different conversion
-      return _f16(self._val) if self._bits == 16 else _f32(self._val) if self._bits == 32 else _f64(self._val)
-    return float(int(self))
-
-  # Arithmetic - floats use float(), ints use int()
-  def __add__(s, o): return float(s) + float(o) if s._float else int(s) + int(o)
-  def __radd__(s, o): return float(o) + float(s) if s._float else int(o) + int(s)
-  def __sub__(s, o): return float(s) - float(o) if s._float else int(s) - int(o)
-  def __rsub__(s, o): return float(o) - float(s) if s._float else int(o) - int(s)
-  def __mul__(s, o): return float(s) * float(o) if s._float else int(s) * int(o)
-  def __rmul__(s, o): return float(o) * float(s) if s._float else int(o) * int(s)
-  def __truediv__(s, o): return _div(float(s), float(o)) if s._float else _div(int(s), int(o))
-  def __rtruediv__(s, o): return _div(float(o), float(s)) if s._float else _div(int(o), int(s))
-  def __pow__(s, o): return float(s) ** float(o) if s._float else int(s) ** int(o)
-  def __rpow__(s, o): return float(o) ** float(s) if s._float else int(o) ** int(s)
-  def __neg__(s): return -float(s) if s._float else -int(s)
-  def __abs__(s): return abs(float(s)) if s._float else abs(int(s))
-
-  # Bitwise - GPU shifts mask the shift amount to valid range
-  def __and__(s, o): return int(s) & int(o)
-  def __or__(s, o): return int(s) | int(o)
-  def __xor__(s, o): return int(s) ^ int(o)
-  def __invert__(s): return ~int(s)
-  def __lshift__(s, o): n = int(o); return int(s) << n if 0 <= n < 64 else 0
-  def __rshift__(s, o): n = int(o); return int(s) >> n if 0 <= n < 64 else 0
-  def __rand__(s, o): return int(o) & int(s)
-  def __ror__(s, o): return int(o) | int(s)
-  def __rxor__(s, o): return int(o) ^ int(s)
-  def __rlshift__(s, o): n = int(s); return int(o) << n if 0 <= n < 64 else 0
-  def __rrshift__(s, o): n = int(s); return int(o) >> n if 0 <= n < 64 else 0
-
-  # Comparison - handle _DenormChecker specially
-  def __eq__(s, o):
-    if isinstance(o, _DenormChecker): return o._check(s)
-    return float(s) == float(o) if s._float else int(s) == int(o)
-  def __ne__(s, o):
-    if isinstance(o, _DenormChecker): return not o._check(s)
-    return float(s) != float(o) if s._float else int(s) != int(o)
-  def __lt__(s, o): return float(s) < float(o) if s._float else int(s) < int(o)
-  def __le__(s, o): return float(s) <= float(o) if s._float else int(s) <= int(o)
-  def __gt__(s, o): return float(s) > float(o) if s._float else int(s) > int(o)
-  def __ge__(s, o): return float(s) >= float(o) if s._float else int(s) >= int(o)
-
-  def __bool__(s): return bool(int(s))
-
-  # Allow chained type access like jump_addr.i64 when jump_addr is already a TypedView
-  # These just return self or convert appropriately
-  @property
-  def i64(s): return s if s._bits == 64 and s._signed else int(s)
-  @property
-  def u64(s): return s if s._bits == 64 and not s._signed else int(s) & MASK64
-  @property
-  def i32(s): return s if s._bits == 32 and s._signed else _sext(int(s) & MASK32, 32)
-  @property
-  def u32(s): return s if s._bits == 32 and not s._signed else int(s) & MASK32
-
-class Reg:
-  """GPU register: D0.f32 = S0.f32 + S1.f32 just works."""
-  __slots__ = ('_val',)
-  def __init__(self, val=0): self._val = int(val) & MASK64
-
-  # Typed views
-  u64 = property(lambda s: TypedView(s, 64), lambda s, v: setattr(s, '_val', int(v) & MASK64))
-  i64 = property(lambda s: TypedView(s, 64, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK64))
-  b64 = property(lambda s: TypedView(s, 64), lambda s, v: setattr(s, '_val', int(v) & MASK64))
-  f64 = property(lambda s: TypedView(s, 64, is_float=True), lambda s, v: setattr(s, '_val', v if isinstance(v, int) else _i64(float(v))))
-  u32 = property(lambda s: TypedView(s, 32), lambda s, v: setattr(s, '_val', int(v) & MASK32))
-  i32 = property(lambda s: TypedView(s, 32, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK32))
-  b32 = property(lambda s: TypedView(s, 32), lambda s, v: setattr(s, '_val', int(v) & MASK32))
-  f32 = property(lambda s: TypedView(s, 32, is_float=True), lambda s, v: setattr(s, '_val', _i32(float(v))))
-  u24 = property(lambda s: TypedView(s, 24))
-  i24 = property(lambda s: TypedView(s, 24, signed=True))
-  u16 = property(lambda s: TypedView(s, 16), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
-  i16 = property(lambda s: TypedView(s, 16, signed=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
-  b16 = property(lambda s: TypedView(s, 16), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
-  f16 = property(lambda s: TypedView(s, 16, is_float=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _i16(float(v))) & 0xffff)))
-  bf16 = property(lambda s: TypedView(s, 16, is_float=True, is_bf16=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _ibf16(float(v))) & 0xffff)))
-  u8 = property(lambda s: TypedView(s, 8))
-  i8 = property(lambda s: TypedView(s, 8, signed=True))
-  u1 = property(lambda s: TypedView(s, 1))  # single bit
-
-  def __getitem__(s, key):
-    if isinstance(key, slice): return SliceProxy(s, int(key.start), int(key.stop))
-    return (s._val >> int(key)) & 1
-
-  def __setitem__(s, key, value):
-    if isinstance(key, slice):
-      high, low = int(key.start), int(key.stop)
-      mask = (1 << (high - low + 1)) - 1
-      s._val = (s._val & ~(mask << low)) | ((int(value) & mask) << low)
-    elif value: s._val |= (1 << int(key))
-    else: s._val &= ~(1 << int(key))
-
-  def __int__(s): return s._val
-  def __index__(s): return s._val
-  def __bool__(s): return bool(s._val)
-
-  # Arithmetic (for tmp = tmp + 1 patterns). Float operands trigger f32 interpretation.
-  def __add__(s, o): return (_f32(s._val) + float(o)) if isinstance(o, float) else s._val + int(o)
-  def __radd__(s, o): return (float(o) + _f32(s._val)) if isinstance(o, float) else int(o) + s._val
-  def __sub__(s, o): return (_f32(s._val) - float(o)) if isinstance(o, float) else s._val - int(o)
-  def __rsub__(s, o): return (float(o) - _f32(s._val)) if isinstance(o, float) else int(o) - s._val
-  def __mul__(s, o): return (_f32(s._val) * float(o)) if isinstance(o, float) else s._val * int(o)
-  def __rmul__(s, o): return (float(o) * _f32(s._val)) if isinstance(o, float) else int(o) * s._val
-  def __and__(s, o): return s._val & int(o)
-  def __rand__(s, o): return int(o) & s._val
-  def __or__(s, o): return s._val | int(o)
-  def __ror__(s, o): return int(o) | s._val
-  def __xor__(s, o): return s._val ^ int(o)
-  def __rxor__(s, o): return int(o) ^ s._val
-  def __lshift__(s, o): n = int(o); return s._val << n if 0 <= n < 64 else 0
-  def __rshift__(s, o): n = int(o); return s._val >> n if 0 <= n < 64 else 0
-  def __invert__(s): return ~s._val
-
-  # Comparison (for tmp >= 0x100000000 patterns)
-  def __lt__(s, o): return s._val < int(o)
-  def __le__(s, o): return s._val <= int(o)
-  def __gt__(s, o): return s._val > int(o)
-  def __ge__(s, o): return s._val >= int(o)
-  def __eq__(s, o): return s._val == int(o)
-  def __ne__(s, o): return s._val != int(o)
-
-
@@ -1,670 +0,0 @@
-# Generate AMD ISA autogen files from PDF documentation
-# Combines format/enum generation (previously in dsl.py) and pseudocode compilation (previously in pcode.py)
-# Usage: python -m extra.assembly.amd.pdf [--arch rdna3|rdna4|cdna|all]
-import re, functools
-from pathlib import Path
-from concurrent.futures import ProcessPoolExecutor
-
-PDF_URLS = {
-  "rdna3": "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content",
-  "rdna4": "https://docs.amd.com/api/khub/documents/uQpkEvk3pv~kfAb2x~j4uw/content",
-  "cdna": ["https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-mi300-cdna3-instruction-set-architecture.pdf",
-           "https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-cdna4-instruction-set-architecture.pdf"],
-}
-
-# Field type mappings and ordering
-FIELD_TYPES = {'SSRC0': 'SSrc', 'SSRC1': 'SSrc', 'SOFFSET': 'SSrc', 'SADDR': 'SSrc', 'SRC0': 'Src', 'SRC1': 'Src', 'SRC2': 'Src',
-  'SDST': 'SGPRField', 'SBASE': 'SGPRField', 'SDATA': 'SGPRField', 'SRSRC': 'SGPRField', 'VDST': 'VGPRField', 'VSRC1': 'VGPRField',
-  'VDATA': 'VGPRField', 'VADDR': 'VGPRField', 'ADDR': 'VGPRField', 'DATA': 'VGPRField', 'DATA0': 'VGPRField', 'DATA1': 'VGPRField',
-  'SIMM16': 'SImm', 'OFFSET': 'Imm', 'OPX': 'VOPDOp', 'OPY': 'VOPDOp', 'SRCX0': 'Src', 'SRCY0': 'Src',
-  'VSRCX1': 'VGPRField', 'VSRCY1': 'VGPRField', 'VDSTX': 'VGPRField', 'VDSTY': 'VDSTYEnc'}
-FIELD_ORDER = {
-  'SOP2': ['op', 'sdst', 'ssrc0', 'ssrc1'], 'SOP1': ['op', 'sdst', 'ssrc0'], 'SOPC': ['op', 'ssrc0', 'ssrc1'],
-  'SOPK': ['op', 'sdst', 'simm16'], 'SOPP': ['op', 'simm16'], 'VOP1': ['op', 'vdst', 'src0'], 'VOPC': ['op', 'src0', 'vsrc1'],
-  'VOP2': ['op', 'vdst', 'src0', 'vsrc1'], 'VOP3SD': ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2', 'clmp'],
-  'SMEM': ['op', 'sdata', 'sbase', 'soffset', 'offset', 'glc', 'dlc'], 'DS': ['op', 'vdst', 'addr', 'data0', 'data1'],
-  'VOP3': ['op', 'vdst', 'src0', 'src1', 'src2', 'omod', 'neg', 'abs', 'clmp', 'opsel'],
-  'VOP3P': ['op', 'vdst', 'src0', 'src1', 'src2', 'neg', 'neg_hi', 'opsel', 'opsel_hi', 'clmp'],
-  'FLAT': ['op', 'vdst', 'addr', 'data', 'saddr', 'offset', 'seg', 'dlc', 'glc', 'slc'],
-  'MUBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
-  'MTBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
-  'MIMG': ['op', 'vdata', 'vaddr', 'srsrc', 'ssamp', 'dmask', 'dim', 'unrm', 'dlc', 'glc', 'slc'],
-  'EXP': ['en', 'target', 'vsrc0', 'vsrc1', 'vsrc2', 'vsrc3', 'done', 'row'],
-  'VINTERP': ['op', 'vdst', 'src0', 'src1', 'src2', 'waitexp', 'clmp', 'opsel', 'neg'],
-  'VOPD': ['opx', 'opy', 'vdstx', 'vdsty', 'srcx0', 'vsrcx1', 'srcy0', 'vsrcy1'],
-  'LDSDIR': ['op', 'vdst', 'attr', 'attr_chan', 'wait_va']}
-SRC_EXTRAS = {233: 'DPP8', 234: 'DPP8FI', 250: 'DPP16', 251: 'VCCZ', 252: 'EXECZ', 254: 'LDS_DIRECT'}
-FLOAT_MAP = {'0.5': 'POS_HALF', '-0.5': 'NEG_HALF', '1.0': 'POS_ONE', '-1.0': 'NEG_ONE', '2.0': 'POS_TWO', '-2.0': 'NEG_TWO',
-  '4.0': 'POS_FOUR', '-4.0': 'NEG_FOUR', '1/(2*PI)': 'INV_2PI', '0': 'ZERO'}
-INST_PATTERN = re.compile(r'^([SVD]S?_[A-Z0-9_]+)\s+(\d+)\s*$', re.M)
-
-# Patterns that can't be handled by the DSL (require special handling in emu.py)
-UNSUPPORTED = ['SGPR[', 'V_SWAP', 'eval ', 'FATAL_HALT', 'HW_REGISTERS',
-               'vscnt', 'vmcnt', 'expcnt', 'lgkmcnt',
-               'CVT_OFF_TABLE', 'ThreadMask',
-               'S1[i', 'C.i32',
-               'if n.', 'DST.u32', 'addrd = DST', 'addr = DST',
-               'BARRIER_STATE', 'ReallocVgprs',
-               'GPR_IDX', 'VSKIP', 'specified in', 'TTBL',
-               'fp6', 'bf6']  # Malformed pseudocode from PDF
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# COMPILER: pseudocode -> Python (minimal transforms)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def compile_pseudocode(pseudocode: str) -> str:
-  """Compile pseudocode to Python. Transforms are minimal - most syntax just works."""
-  pseudocode = re.sub(r'\bpass\b', 'pass_', pseudocode)  # 'pass' is Python keyword
-  raw_lines = pseudocode.strip().split('\n')
-  joined_lines: list[str] = []
-  for line in raw_lines:
-    line = line.strip()
-    if joined_lines and (joined_lines[-1].rstrip().endswith(('||', '&&', '(', ',')) or
-                         (joined_lines[-1].count('(') > joined_lines[-1].count(')'))):
-      joined_lines[-1] = joined_lines[-1].rstrip() + ' ' + line
-    else:
-      joined_lines.append(line)
-
-  lines = []
-  indent, need_pass, in_first_match_loop = 0, False, False
-  declared_arrays: dict[str, int] = {}  # Track declared arrays: name -> size
-  for line in joined_lines:
-    line = line.strip()
-    if not line or line.startswith('//'): continue
-    if line.startswith('if '):
-      lines.append('  ' * indent + f"if {_expr(line[3:].rstrip(' then'), declared_arrays)}:")
-      indent += 1
-      need_pass = True
-    elif line.startswith('elsif '):
-      if need_pass: lines.append('  ' * indent + "pass")
-      indent -= 1
-      lines.append('  ' * indent + f"elif {_expr(line[6:].rstrip(' then'), declared_arrays)}:")
-      indent += 1
-      need_pass = True
-    elif line == 'else':
-      if need_pass: lines.append('  ' * indent + "pass")
-      indent -= 1
-      lines.append('  ' * indent + "else:")
-      indent += 1
-      need_pass = True
-    elif line.startswith('endif'):
-      if need_pass: lines.append('  ' * indent + "pass")
-      indent -= 1
-      need_pass = False
-    elif line.startswith('endfor'):
-      if need_pass: lines.append('  ' * indent + "pass")
-      indent -= 1
-      need_pass, in_first_match_loop = False, False
-    elif m := re.match(r'declare\s+(\w+)\s*:\s*\d+\'[FBU]\[(\d+)\]', line):
-      # Handle array declarations: declare in : 32'F[3] or declare S : 32'B[3]
-      arr_name, arr_size = m[1], int(m[2])
-      declared_arrays[arr_name] = arr_size
-      py_name = f"{arr_name}_" if arr_name == 'in' else arr_name  # 'in' is Python keyword
-      if arr_name == 'S':
-        lines.append('  ' * indent + f"{py_name} = [S0, S1, S2]")  # Map to source registers
-      else:
-        lines.append('  ' * indent + f"{py_name} = [Reg(0) for _ in range({arr_size})]")
-    elif line.startswith('declare '):
-      pass  # Ignore other declare statements
-    elif m := re.match(r'for (\w+) in (.+?)\s*:\s*(.+?) do', line):
-      start, end = _expr(m[2].strip(), declared_arrays), _expr(m[3].strip(), declared_arrays)
-      lines.append('  ' * indent + f"for {m[1]} in range({start}, int({end})+1):")
-      indent += 1
-      need_pass, in_first_match_loop = True, True
-    elif '=' in line and not line.startswith('=='):
-      need_pass = False
-      line = line.rstrip(';')
-      if m := re.match(r'\{\s*D1\.[ui]1\s*,\s*D0\.[ui]64\s*\}\s*=\s*(.+)', line):
-        rhs = _expr(m[1], declared_arrays)
-        lines.append('  ' * indent + f"_full = {rhs}")
-        lines.append('  ' * indent + f"D0.u64 = int(_full) & 0xffffffffffffffff")
-        lines.append('  ' * indent + f"D1 = Reg((int(_full) >> 64) & 1)")
-      elif any(op in line for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^=')):
-        for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^='):
-          if op in line:
-            lhs, rhs = line.split(op, 1)
-            lhs_s = _expr(lhs.strip(), declared_arrays)  # Transform LHS too for array access
-            lines.append('  ' * indent + f"{lhs_s} {op} {_expr(rhs.strip(), declared_arrays)}")
-            break
-      else:
-        lhs, rhs = line.split('=', 1)
-        lhs_s, rhs_s = lhs.strip(), rhs.strip()
-        lhs_t = _expr(lhs_s, declared_arrays)  # Transform LHS for array access
-        stmt = _assign(lhs_t, _expr(rhs_s, declared_arrays), declared_arrays)
-        if in_first_match_loop and rhs_s == 'i' and (lhs_s == 'tmp' or lhs_s == 'D0.i32'):
-          stmt += "; break"
-        lines.append('  ' * indent + stmt)
-  if need_pass: lines.append('  ' * indent + "pass")
-  return '\n'.join(lines)
-
-def _assign(lhs: str, rhs: str, declared_arrays: dict[str, int] | None = None) -> str:
-  # Check for array element assignment: in_[i] should not wrap in Reg()
-  if declared_arrays and re.match(r'\w+_?\[\w+\]', lhs):
-    return f"{lhs} = {rhs}"
-  if lhs in ('tmp', 'SCC', 'VCC', 'EXEC', 'D0', 'D1', 'saveexec', 'PC'):
-    return f"{lhs} = Reg({rhs})"
-  return f"{lhs} = {rhs}"
-
-def _expr(e: str, declared_arrays: dict[str, int] | None = None) -> str:
-  e = e.strip()
-  # Handle OPSEL_HI.u3[i] and OPSEL.u3[i] - bit extraction from opsel fields
-  e = re.sub(r'(OPSEL(?:_HI)?)\.u\d+\[(\w+)\]', r'((\1 >> \2) & 1)', e)
-  # Rename 'in' to 'in_' to avoid Python keyword conflict
-  e = re.sub(r'\bin\[', 'in_[', e)
-  e = e.replace('&&', ' and ').replace('||', ' or ').replace('<>', ' != ')
-  e = re.sub(r'!([^=])', r' not \1', e)
-  e = re.sub(r'\{\s*(\w+\.u32)\s*,\s*(\w+\.u32)\s*\}', r'_pack32(\1, \2)', e)
-  def pack(m):
-    hi, lo = _expr(m[1].strip(), declared_arrays), _expr(m[2].strip(), declared_arrays)
-    return f'_pack({hi}, {lo})'
-  e = re.sub(r'\{\s*([^,{}]+)\s*,\s*([^,{}]+)\s*\}', pack, e)
-  e = re.sub(r"1201'B\(2\.0\s*/\s*PI\)", "TWO_OVER_PI_1201", e)
-  e = re.sub(r"\d+'([0-9a-fA-Fx]+)[UuFf]*", r'\1', e)
-  e = re.sub(r"\d+'[FIBU]\(", "(", e)
-  e = re.sub(r'\bB\(', '(', e)
-  e = re.sub(r'([0-9a-fA-Fx])ULL\b', r'\1', e)
-  e = re.sub(r'([0-9a-fA-Fx])LL\b', r'\1', e)
-  e = re.sub(r'([0-9a-fA-Fx])U\b', r'\1', e)
-  e = re.sub(r'(\d\.?\d*)F\b', r'\1', e)
-  e = re.sub(r'(\[laneId\])\.[uib]\d+', r'\1', e)
-  e = e.replace('+INF', 'INF').replace('-INF', '(-INF)')
-  e = re.sub(r'NAN\.f\d+', 'float("nan")', e)
-  def convert_verilog_slice(m):
-    start, width = m.group(1).strip(), m.group(2).strip()
-    return f'[({start}) + ({width}) - 1 : ({start})]'
-  e = re.sub(r'\[([^:\[\]]+)\s*\+:\s*([^:\[\]]+)\]', convert_verilog_slice, e)
-  def process_brackets(s):
-    result, i = [], 0
-    while i < len(s):
-      if s[i] == '[':
-        depth, start = 1, i + 1
-        j = start
-        while j < len(s) and depth > 0:
-          if s[j] == '[': depth += 1
-          elif s[j] == ']': depth -= 1
-          j += 1
-        inner = _expr(s[start:j-1], declared_arrays)
-        result.append('[' + inner + ']')
-        i = j
-      else:
-        result.append(s[i])
-        i += 1
-    return ''.join(result)
-  e = process_brackets(e)
-  while '?' in e:
-    depth, bracket, q = 0, 0, -1
-    for i, c in enumerate(e):
-      if c == '(': depth += 1
-      elif c == ')': depth -= 1
-      elif c == '[': bracket += 1
-      elif c == ']': bracket -= 1
-      elif c == '?' and depth == 0 and bracket == 0: q = i; break
-    if q < 0: break
-    depth, bracket, col = 0, 0, -1
-    for i in range(q + 1, len(e)):
-      if e[i] == '(': depth += 1
-      elif e[i] == ')': depth -= 1
-      elif e[i] == '[': bracket += 1
-      elif e[i] == ']': bracket -= 1
-      elif e[i] == ':' and depth == 0 and bracket == 0: col = i; break
-    if col < 0: break
-    cond, t, f = e[:q].strip(), e[q+1:col].strip(), e[col+1:].strip()
-    e = f'(({t}) if ({cond}) else ({f}))'
-  return e
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# PDF PARSING WITH PAGE CACHING
-# ═══════════════════════════════════════════════════════════════════════════════
-
-class CachedPDF:
-  """PDF wrapper with page text/table caching for faster repeated access."""
-  def __init__(self, pdf):
-    self._pdf, self._text_cache, self._table_cache = pdf, {}, {}
-  def __len__(self): return len(self._pdf.pages)
-  def text(self, i):
-    if i not in self._text_cache: self._text_cache[i] = self._pdf.pages[i].extract_text() or ''
-    return self._text_cache[i]
-  def tables(self, i):
-    if i not in self._table_cache: self._table_cache[i] = [t.extract() for t in self._pdf.pages[i].find_tables()]
-    return self._table_cache[i]
-
-def _parse_bits(s: str) -> tuple[int, int] | None:
-  return (int(m.group(1)), int(m.group(2) or m.group(1))) if (m := re.match(r'\[(\d+)(?::(\d+))?\]', s)) else None
-
-def _parse_fields_table(table: list, fmt: str, enums: set[str]) -> list[tuple]:
-  fields = []
-  for row in table[1:]:
-    if not row or not row[0]: continue
-    name, bits_str = row[0].split('\n')[0].strip(), (row[1] or '').split('\n')[0].strip()
-    if not (bits := _parse_bits(bits_str)): continue
-    enc_val, hi, lo = None, bits[0], bits[1]
-    if name == 'ENCODING' and row[2]:
-      if m := re.search(r"(?:'b|Must be:\s*)([01_]+)", row[2]):
-        enc_bits = m.group(1).replace('_', '')
-        enc_val, declared_width, actual_width = int(enc_bits, 2), hi - lo + 1, len(enc_bits)
-        if actual_width > declared_width: lo = hi - actual_width + 1
-    ftype = f"{fmt}Op" if name == 'OP' and f"{fmt}Op" in enums else FIELD_TYPES.get(name.upper())
-    fields.append((name, hi, lo, enc_val, ftype))
-  return fields
-
-def _parse_single_pdf(url: str):
-  """Parse a single PDF and return (formats, enums, src_enum, doc_name, instructions)."""
-  import pdfplumber
-  from tinygrad.helpers import fetch
-
-  pdf = CachedPDF(pdfplumber.open(fetch(url)))
-  total_pages = len(pdf)
-
-  # Auto-detect document type
-  first_page = pdf.text(0)
-  is_cdna4, is_cdna3 = 'CDNA4' in first_page or 'CDNA 4' in first_page, 'CDNA3' in first_page or 'MI300' in first_page
-  is_cdna, is_rdna4 = is_cdna3 or is_cdna4, 'RDNA4' in first_page or 'RDNA 4' in first_page
-  is_rdna35, is_rdna3 = 'RDNA3.5' in first_page or 'RDNA 3.5' in first_page, 'RDNA3' in first_page and 'RDNA3.5' not in first_page
-  doc_name = "CDNA4" if is_cdna4 else "CDNA3" if is_cdna3 else "RDNA4" if is_rdna4 else "RDNA3.5" if is_rdna35 else "RDNA3" if is_rdna3 else "Unknown"
-
-  # Find Microcode Formats section (for formats/enums)
-  microcode_start = next((i for i in range(int(total_pages * 0.2), total_pages)
-                          if re.search(r'\d+\.\d+\.\d+\.\s+SOP2\b|Chapter \d+\.\s+Microcode Formats', pdf.text(i))), int(total_pages * 0.9))
-  # Find Instructions section (for pseudocode)
-  instr_start = next((i for i in range(int(total_pages * 0.1), int(total_pages * 0.5))
-                      if re.search(r'Chapter \d+\.\s+Instructions\b', pdf.text(i))), total_pages // 3)
-  instr_end = next((i for start in [int(total_pages * 0.6), int(total_pages * 0.5), instr_start]
-                    for i in range(start, min(start + 100, total_pages))
-                    if re.search(r'Chapter \d+\.\s+Microcode Formats', pdf.text(i))), total_pages)
-
-  # Parse src enum from SSRC encoding table
-  src_enum = dict(SRC_EXTRAS)
-  for i in range(microcode_start, min(microcode_start + 10, total_pages)):
-    text = pdf.text(i)
-    if 'SSRC0' in text and 'VCC_LO' in text:
-      for m in re.finditer(r'^(\d+)\s+(\S+)', text, re.M):
-        val, name = int(m.group(1)), m.group(2).rstrip('.:')
-        if name in FLOAT_MAP: src_enum[val] = FLOAT_MAP[name]
-        elif re.match(r'^[A-Z][A-Z0-9_]*$', name): src_enum[val] = name
-      break
-
-  # Parse opcode tables
-  full_text = '\n'.join(pdf.text(i) for i in range(microcode_start, min(microcode_start + 50, total_pages)))
-  enums: dict[str, dict[int, str]] = {}
-  for m in re.finditer(r'Table \d+\. (\w+) Opcodes(.*?)(?=Table \d+\.|\n\d+\.\d+\.\d+\.\s+\w+\s*\nDescription|$)', full_text, re.S):
-    if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+([A-Z][A-Z0-9_]+)', m.group(2))}:
-      enums[m.group(1) + "Op"] = ops
-  if vopd_m := re.search(r'Table \d+\. VOPD Y-Opcodes\n(.*?)(?=Table \d+\.|15\.\d)', full_text, re.S):
-    if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+(V_DUAL_\w+)', vopd_m.group(1))}:
-      enums["VOPDOp"] = ops
-  enum_names = set(enums.keys())
-
-  # Parse instruction formats
-  def is_fields_table(t): return t and len(t) > 1 and t[0] and 'Field' in str(t[0][0] or '')
-  def has_encoding(fields): return any(f[0] == 'ENCODING' for f in fields)
-  def has_header_before_fields(text): return (pos := text.find('Field Name')) != -1 and bool(re.search(r'\d+\.\d+\.\d+\.\s+\w+\s*\n', text[:pos]))
-
-  format_headers = []
-  for i in range(50):
-    if microcode_start + i >= total_pages: break
-    text = pdf.text(microcode_start + i)
-    for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n?Description', text): format_headers.append((m.group(1), i, m.start()))
-    for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n', text):
-      fmt_name = m.group(1)
-      if is_cdna and fmt_name.isupper() and len(fmt_name) >= 2: format_headers.append((fmt_name, i, m.start()))
-      elif m.start() > len(text) - 200 and 'Description' not in text[m.end():] and i + 1 < 50:
-        next_text = pdf.text(microcode_start + i + 1).lstrip()
-        if next_text.startswith('Description') or (next_text.startswith('"RDNA') and 'Description' in next_text[:200]):
-          format_headers.append((fmt_name, i, m.start()))
-
-  formats: dict[str, list] = {}
-  for fmt_name, rel_idx, header_pos in format_headers:
-    if fmt_name in formats: continue
-    page_idx = microcode_start + rel_idx
-    text = pdf.text(page_idx)
-    field_pos = text.find('Field Name', header_pos)
-    fields = None
-    for offset in range(3):
-      if page_idx + offset >= total_pages: break
-      if offset > 0 and has_header_before_fields(pdf.text(page_idx + offset)): break
-      for t in pdf.tables(page_idx + offset) if offset > 0 or field_pos > header_pos else []:
-        if is_fields_table(t) and (f := _parse_fields_table(t, fmt_name, enum_names)) and has_encoding(f): fields = f; break
-      if fields: break
-    if not fields and field_pos > header_pos:
-      for t in pdf.tables(page_idx):
-        if is_fields_table(t) and (f := _parse_fields_table(t, fmt_name, enum_names)): fields = f; break
-    if not fields: continue
-    field_names = {f[0] for f in fields}
-    for pg_offset in range(1, 3):
-      if page_idx + pg_offset >= total_pages or has_header_before_fields(pdf.text(page_idx + pg_offset)): break
-      for t in pdf.tables(page_idx + pg_offset):
-        if is_fields_table(t) and (extra := _parse_fields_table(t, fmt_name, enum_names)) and not has_encoding(extra):
-          for ef in extra:
-            if ef[0] not in field_names: fields.append(ef); field_names.add(ef[0])
-          break
-    formats[fmt_name] = fields
-
-  # Fix known PDF errors
-  if 'SMEM' in formats:
-    formats['SMEM'] = [(n, 13 if n == 'DLC' else 14 if n == 'GLC' else h, 13 if n == 'DLC' else 14 if n == 'GLC' else l, e, t)
-                       for n, h, l, e, t in formats['SMEM']]
-  if doc_name in ('RDNA3', 'RDNA3.5'):
-    if 'SOPPOp' in enums: assert 8 not in enums['SOPPOp']; enums['SOPPOp'][8] = 'S_WAITCNT_DEPCTR'
-    if 'DSOp' in enums:
-      for k, v in {24: 'DS_GWS_SEMA_RELEASE_ALL', 25: 'DS_GWS_INIT', 26: 'DS_GWS_SEMA_V', 27: 'DS_GWS_SEMA_BR', 28: 'DS_GWS_SEMA_P', 29: 'DS_GWS_BARRIER'}.items():
-        assert k not in enums['DSOp']; enums['DSOp'][k] = v
-    if 'FLATOp' in enums:
-      for k, v in {40: 'GLOBAL_LOAD_ADDTID_B32', 41: 'GLOBAL_STORE_ADDTID_B32', 55: 'FLAT_ATOMIC_CSUB_U32'}.items():
-        assert k not in enums['FLATOp']; enums['FLATOp'][k] = v
-
-  # Extract pseudocode for instructions
-  all_text = '\n'.join(pdf.text(i) for i in range(instr_start, instr_end))
-  matches = list(INST_PATTERN.finditer(all_text))
-  raw_pseudocode: dict[tuple[str, int], str] = {}
-  for i, match in enumerate(matches):
-    name, opcode = match.group(1), int(match.group(2))
-    start, end = match.end(), matches[i + 1].start() if i + 1 < len(matches) else match.end() + 2000
-    snippet = all_text[start:end].strip()
-    if pseudocode := _extract_pseudocode(snippet): raw_pseudocode[(name, opcode)] = pseudocode
-
-  return {"formats": formats, "enums": enums, "src_enum": src_enum, "doc_name": doc_name, "pseudocode": raw_pseudocode, "is_cdna": is_cdna}
-
-def _extract_pseudocode(text: str) -> str | None:
-  """Extract pseudocode from an instruction description snippet."""
-  lines, result, depth, in_lambda = text.split('\n'), [], 0, 0
-  for line in lines:
-    s = line.strip()
-    if not s or re.match(r'^\d+ of \d+$', s) or re.match(r'^\d+\.\d+\..*Instructions', s): continue
-    if s.startswith(('Notes', 'Functional examples')): break
-    if s.startswith(('"RDNA', 'AMD ', 'CDNA')): continue
-    if '= lambda(' in s: in_lambda += 1; continue
-    if in_lambda > 0:
-      if s.endswith(');'): in_lambda -= 1
-      continue
-    if s.startswith('if '): depth += 1
-    elif s.startswith('endif'): depth = max(0, depth - 1)
-    if s.endswith('.') and not any(p in s for p in ['D0', 'D1', 'S0', 'S1', 'S2', 'SCC', 'VCC', 'tmp', '=']): continue
-    if re.match(r'^[a-z].*\.$', s) and '=' not in s: continue
-    is_code = (any(p in s for p in ['D0.', 'D1.', 'S0.', 'S1.', 'S2.', 'SCC =', 'SCC ?', 'VCC', 'EXEC', 'tmp =', 'tmp[', 'lane =', 'PC =',
-                                    'D0[', 'D1[', 'S0[', 'S1[', 'S2[', 'MEM[', 'RETURN_DATA', 'DATA.', 'DATA0', 'DATA1', 'ADDR']) or
-               s.startswith(('if ', 'else', 'elsif', 'endif', 'declare ', 'for ', 'endfor', '//')) or
-               re.match(r'^[a-z_]+\s*=', s) or re.match(r'^[a-z_]+\[', s) or (depth > 0 and '=' in s))
-    if is_code: result.append(s)
-  return '\n'.join(result) if result else None
-
-def _merge_results(results: list[dict]) -> dict:
-  """Merge multiple PDF parse results into a superset."""
-  merged = {"formats": {}, "enums": {}, "src_enum": dict(SRC_EXTRAS), "doc_names": [], "pseudocode": {}, "is_cdna": False}
-  for r in results:
-    merged["doc_names"].append(r["doc_name"])
-    merged["is_cdna"] = merged["is_cdna"] or r["is_cdna"]
-    for val, name in r["src_enum"].items():
-      if val in merged["src_enum"]: assert merged["src_enum"][val] == name
-      else: merged["src_enum"][val] = name
-    for enum_name, ops in r["enums"].items():
-      if enum_name not in merged["enums"]: merged["enums"][enum_name] = {}
-      for val, name in ops.items():
-        if val in merged["enums"][enum_name]: assert merged["enums"][enum_name][val] == name
-        else: merged["enums"][enum_name][val] = name
-    for fmt_name, fields in r["formats"].items():
-      if fmt_name not in merged["formats"]: merged["formats"][fmt_name] = list(fields)
-      else:
-        existing = {f[0]: (f[1], f[2]) for f in merged["formats"][fmt_name]}
-        for f in fields:
-          if f[0] in existing: assert existing[f[0]] == (f[1], f[2])
-          else: merged["formats"][fmt_name].append(f)
-    for key, pc in r["pseudocode"].items():
-      if key not in merged["pseudocode"]: merged["pseudocode"][key] = pc
-  return merged
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# CODE GENERATION
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def _generate_enum_py(enums, src_enum, doc_name) -> str:
-  """Generate enum.py content (just enums, no dsl.py dependency)."""
-  def enum_lines(name, items): return [f"class {name}(IntEnum):"] + [f"  {n} = {v}" for v, n in sorted(items.items())] + [""]
-  lines = [f"# autogenerated from AMD {doc_name} ISA PDF by pdf.py - do not edit", "from enum import IntEnum", ""]
-  lines += enum_lines("SrcEnum", src_enum) + sum([enum_lines(n, ops) for n, ops in sorted(enums.items())], [])
-  return '\n'.join(lines)
-
-def _generate_ins_py(formats, enums, src_enum, doc_name) -> str:
-  """Generate ins.py content (instruction formats and helpers, imports dsl.py and enum.py)."""
-  def field_key(f, order): return order.index(f[0].lower()) if f[0].lower() in order else 1000
-  lines = [f"# autogenerated from AMD {doc_name} ISA PDF by pdf.py - do not edit",
-           "# ruff: noqa: F401,F403", "from typing import Annotated",
-           "from extra.assembly.amd.dsl import bits, BitField, Inst32, Inst64, SGPR, VGPR, TTMP as TTMP, s as s, v as v, ttmp as ttmp, SSrc, Src, SImm, Imm, VDSTYEnc, SGPRField, VGPRField",
-           "from extra.assembly.amd.autogen.{arch}.enum import *",
-           "import functools", ""]
-  format_defaults = {'VOP3P': {'opsel_hi': 3, 'opsel_hi2': 1}}
-  lines.append("# instruction formats")
-  for fmt_name, fields in sorted(formats.items()):
-    base = "Inst64" if max(f[1] for f in fields) > 31 or fmt_name == 'VOP3SD' else "Inst32"
-    order = FIELD_ORDER.get(fmt_name, [])
-    lines.append(f"class {fmt_name}({base}):")
-    if enc := next((f for f in fields if f[0] == 'ENCODING'), None):
-      lines.append(f"  encoding = bits[{enc[1]}:{enc[2]}] == 0b{enc[3]:b}" if enc[1] != enc[2] else f"  encoding = bits[{enc[1]}] == {enc[3]}")
-    if defaults := format_defaults.get(fmt_name): lines.append(f"  _defaults = {defaults}")
-    for name, hi, lo, _, ftype in sorted([f for f in fields if f[0] != 'ENCODING'], key=lambda f: field_key(f, order)):
-      ann = f":Annotated[BitField, {ftype}]" if ftype and ftype.endswith('Op') else f":{ftype}" if ftype else ""
-      lines.append(f"  {name.lower()}{ann} = bits[{hi}]" if hi == lo else f"  {name.lower()}{ann} = bits[{hi}:{lo}]")
-    lines.append("")
-  lines.append("# instruction helpers")
-  for cls_name, ops in sorted(enums.items()):
-    fmt = cls_name[:-2]
-    for op_val, name in sorted(ops.items()):
-      seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=1"}.get(fmt, "")
-      tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt, f"{fmt}, {cls_name}")
-      if fmt in formats or fmt in ("GLOBAL", "SCRATCH"):
-        suffix = "_e32" if fmt in ("VOP1", "VOP2", "VOPC") else "_e64" if fmt == "VOP3" and op_val < 512 else ""
-        if name in ('V_FMAMK_F32', 'V_FMAMK_F16'):
-          lines.append(f"def {name.lower()}{suffix}(vdst, src0, K, vsrc1): return {fmt}({cls_name}.{name}, vdst, src0, vsrc1, literal=K)")
-        elif name in ('V_FMAAK_F32', 'V_FMAAK_F16'):
-          lines.append(f"def {name.lower()}{suffix}(vdst, src0, vsrc1, K): return {fmt}({cls_name}.{name}, vdst, src0, vsrc1, literal=K)")
-        else: lines.append(f"{name.lower()}{suffix} = functools.partial({tgt}.{name}{seg})")
-  src_names = {name for _, name in src_enum.items()}
-  lines += [""] + [f"{name} = SrcEnum.{name}" for _, name in sorted(src_enum.items()) if name not in {'DPP8', 'DPP16'}]
-  if "NULL" in src_names: lines.append("OFF = NULL\n")
-  return '\n'.join(lines)
-
-def _generate_gen_pcode_py(enums, pseudocode, arch) -> str:
-  """Generate gen_pcode.py content (compiled pseudocode functions)."""
-  # Get op enums for this arch (import from .ins which re-exports from .enum)
-  import importlib
-  autogen = importlib.import_module(f"extra.assembly.amd.autogen.{arch}.ins")
-  OP_ENUMS = [getattr(autogen, name) for name in ['SOP1Op', 'SOP2Op', 'SOPCOp', 'SOPKOp', 'SOPPOp', 'VOP1Op', 'VOP2Op', 'VOP3Op', 'VOP3SDOp', 'VOP3POp', 'VOPCOp', 'VOP3AOp', 'VOP3BOp', 'DSOp'] if hasattr(autogen, name)]
-
-  # Build defined ops mapping
-  defined_ops: dict[tuple, list] = {}
-  for enum_cls in OP_ENUMS:
-    for op in enum_cls:
-      if op.name.startswith(('S_', 'V_', 'DS_')): defined_ops.setdefault((op.name, op.value), []).append((enum_cls, op))
-
-  enum_names = [e.__name__ for e in OP_ENUMS]
-  lines = [f'''# autogenerated by pdf.py - do not edit
-# to regenerate: python -m extra.assembly.amd.pdf --arch {arch}
-# ruff: noqa: E501,F405,F403
-# mypy: ignore-errors
-from extra.assembly.amd.autogen.{arch}.enum import {", ".join(enum_names)}
-from extra.assembly.amd.pcode import *
-''']
-
-  instructions: dict = {cls: {} for cls in OP_ENUMS}
-  for key, pc in pseudocode.items():
-    if key in defined_ops:
-      for enum_cls, enum_val in defined_ops[key]: instructions[enum_cls][enum_val] = pc
-
-  for enum_cls in OP_ENUMS:
-    cls_name = enum_cls.__name__
-    if not instructions.get(enum_cls): continue
-    fn_entries = []
-    for op, pc in instructions[enum_cls].items():
-      if any(p in pc for p in UNSUPPORTED): continue
-      try:
-        code = compile_pseudocode(pc)
-        code = _apply_pseudocode_fixes(op, code)
-        fn_name, fn_code = _generate_function(cls_name, op, pc, code)
-        lines.append(fn_code)
-        fn_entries.append((op, fn_name))
-      except Exception as e: print(f"  Warning: Failed to compile {op.name}: {e}")
-    if fn_entries:
-      lines.append(f'{cls_name}_FUNCTIONS = {{')
-      for op, fn_name in fn_entries: lines.append(f"  {cls_name}.{op.name}: {fn_name},")
-      lines.append('}\n')
-
-  # Add V_WRITELANE_B32 if VOP3Op exists
-  if 'VOP3Op' in enum_names:
-    lines.append('''
-# V_WRITELANE_B32: Write scalar to specific lane's VGPR (not in PDF pseudocode)
-def _VOP3Op_V_WRITELANE_B32(s0, s1, s2, d0, scc, vcc, lane, exec_mask, literal, VGPR, _vars, src0_idx=0, vdst_idx=0):
-  wr_lane = s1 & 0x1f
-  return {'d0': d0, 'scc': scc, 'vgpr_write': (wr_lane, vdst_idx, s0 & 0xffffffff)}
-VOP3Op_FUNCTIONS[VOP3Op.V_WRITELANE_B32] = _VOP3Op_V_WRITELANE_B32
-''')
-
-  lines.append('COMPILED_FUNCTIONS = {')
-  for enum_cls in OP_ENUMS:
-    if instructions.get(enum_cls): lines.append(f'  {enum_cls.__name__}: {enum_cls.__name__}_FUNCTIONS,')
-  lines.append('}\n\ndef get_compiled_functions(): return COMPILED_FUNCTIONS')
-  return '\n'.join(lines)
-
-def _apply_pseudocode_fixes(op, code: str) -> str:
-  """Apply known fixes for PDF pseudocode bugs."""
-  if op.name == 'V_DIV_FMAS_F32':
-    code = code.replace('D0.f32 = 2.0 ** 32 * fma(S0.f32, S1.f32, S2.f32)',
-                        'D0.f32 = (2.0 ** 64 if exponent(S2.f32) > 127 else 2.0 ** -64) * fma(S0.f32, S1.f32, S2.f32)')
-  if op.name == 'V_DIV_FMAS_F64':
-    code = code.replace('D0.f64 = 2.0 ** 64 * fma(S0.f64, S1.f64, S2.f64)',
-                        'D0.f64 = (2.0 ** 128 if exponent(S2.f64) > 1023 else 2.0 ** -128) * fma(S0.f64, S1.f64, S2.f64)')
-  if op.name == 'V_DIV_SCALE_F32':
-    code = code.replace('D0.f32 = float("nan")', 'VCC = Reg(0x1); D0.f32 = float("nan")')
-    code = code.replace('elif S1.f32 == DENORM.f32:\n  D0.f32 = ldexp(S0.f32, 64)', 'elif False:\n  pass')
-    code += '\nif S1.f32 == DENORM.f32:\n  D0.f32 = float("nan")'
-    code = code.replace('elif exponent(S2.f32) <= 23:\n  D0.f32 = ldexp(S0.f32, 64)', 'elif exponent(S2.f32) <= 23:\n  VCC = Reg(0x1); D0.f32 = ldexp(S0.f32, 64)')
-    code = code.replace('elif S2.f32 / S1.f32 == DENORM.f32:\n  VCC = Reg(0x1)\n  if S0.f32 == S2.f32:\n    D0.f32 = ldexp(S0.f32, 64)', 'elif S2.f32 / S1.f32 == DENORM.f32:\n  VCC = Reg(0x1)')
-  if op.name == 'V_DIV_SCALE_F64':
-    code = code.replace('D0.f64 = float("nan")', 'VCC = Reg(0x1); D0.f64 = float("nan")')
-    code = code.replace('elif S1.f64 == DENORM.f64:\n  D0.f64 = ldexp(S0.f64, 128)', 'elif False:\n  pass')
-    code += '\nif S1.f64 == DENORM.f64:\n  D0.f64 = float("nan")'
-    code = code.replace('elif exponent(S2.f64) <= 52:\n  D0.f64 = ldexp(S0.f64, 128)', 'elif exponent(S2.f64) <= 52:\n  VCC = Reg(0x1); D0.f64 = ldexp(S0.f64, 128)')
-    code = code.replace('elif S2.f64 / S1.f64 == DENORM.f64:\n  VCC = Reg(0x1)\n  if S0.f64 == S2.f64:\n    D0.f64 = ldexp(S0.f64, 128)', 'elif S2.f64 / S1.f64 == DENORM.f64:\n  VCC = Reg(0x1)')
-  if op.name == 'V_DIV_FIXUP_F32':
-    code = code.replace('D0.f32 = ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))',
-                        'D0.f32 = ((-OVERFLOW_F32) if (sign_out) else (OVERFLOW_F32)) if isNAN(S0.f32) else ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))')
-  if op.name == 'V_DIV_FIXUP_F64':
-    code = code.replace('D0.f64 = ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))',
-                        'D0.f64 = ((-OVERFLOW_F64) if (sign_out) else (OVERFLOW_F64)) if isNAN(S0.f64) else ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))')
-  if op.name == 'V_TRIG_PREOP_F64':
-    code = code.replace('result = F((TWO_OVER_PI_1201[1200 : 0] << shift.u32) & 0x1fffffffffffff)',
-                        'result = float(((TWO_OVER_PI_1201[1200 : 0] << int(shift)) >> (1201 - 53)) & 0x1fffffffffffff)')
-  return code
-
-def _generate_function(cls_name: str, op, pc: str, code: str) -> tuple[str, str]:
-  """Generate a single compiled pseudocode function."""
-  has_d1 = '{ D1' in pc
-  is_cmpx = (cls_name in ('VOPCOp', 'VOP3Op')) and 'EXEC.u64[laneId]' in pc
-  is_div_scale = 'DIV_SCALE' in op.name
-  has_sdst = cls_name == 'VOP3SDOp' and ('VCC.u64[laneId]' in pc or is_div_scale)
-  has_opsel = 'OPSEL' in pc  # FMA_MIX and similar instructions need OPSEL/OPSEL_HI
-  combined = code + pc
-
-  fn_name = f"_{cls_name}_{op.name}"
-  # Function accepts Reg objects directly (uppercase names), laneId is passed directly as int
-  params = "S0, S1, S2, D0, SCC, VCC, laneId, EXEC, literal, VGPR, src0_idx=0, vdst_idx=0, PC=None"
-  if has_opsel: params += ", OPSEL=0, OPSEL_HI=0"
-  lines = [f"def {fn_name}({params}):"]
-
-  # Registers that need special handling (not passed directly)
-  # Only init if used but not first assigned as `name = Reg(...)` in the compiled code
-  def needs_init(name): return name in combined and not re.search(rf'^\s*{name}\s*=\s*Reg\(', code, re.MULTILINE)
-  special_regs = [('D1', 'Reg(0)'), ('SIMM16', 'Reg(literal)'), ('SIMM32', 'Reg(literal)'),
-                  ('SRC0', 'Reg(src0_idx)'), ('VDST', 'Reg(vdst_idx)')]
-  if needs_init('tmp'): special_regs.insert(0, ('tmp', 'Reg(0)'))
-  if needs_init('saveexec'): special_regs.insert(0, ('saveexec', 'Reg(EXEC._val)'))
-  used = {name for name, _ in special_regs if name in combined}
-
-  # Detect which registers are modified (not just read) - look for assignments
-  modifies_d0 = is_div_scale or bool(re.search(r'\bD0\b[.\[]', combined))
-  modifies_exec = is_cmpx or bool(re.search(r'EXEC\.(u32|u64|b32|b64)\s*=', combined))
-  modifies_vcc = has_sdst or bool(re.search(r'VCC\.(u32|u64|b32|b64)\s*=|VCC\.u64\[laneId\]\s*=', combined))
-  modifies_scc = bool(re.search(r'\bSCC\s*=', combined))
-  modifies_pc = bool(re.search(r'\bPC\s*=', combined))
-
-  # Build init code for special registers
-  init_lines = []
-  if is_div_scale: init_lines.append("  D0 = Reg(S0._val)")
-  for name, init in special_regs:
-    if name in used: init_lines.append(f"  {name} = {init}")
-  if 'EXEC_LO' in code: init_lines.append("  EXEC_LO = SliceProxy(EXEC, 31, 0)")
-  if 'EXEC_HI' in code: init_lines.append("  EXEC_HI = SliceProxy(EXEC, 63, 32)")
-  if 'VCCZ' in code and not re.search(r'^\s*VCCZ\s*=', code, re.MULTILINE): init_lines.append("  VCCZ = Reg(1 if VCC._val == 0 else 0)")
-  if 'EXECZ' in code and not re.search(r'^\s*EXECZ\s*=', code, re.MULTILINE): init_lines.append("  EXECZ = Reg(1 if EXEC._val == 0 else 0)")
-  code_lines = [line for line in code.split('\n') if line.strip()]
-  if init_lines:
-    lines.extend(init_lines)
-    if code_lines: lines.append("  # --- compiled pseudocode ---")
-  for line in code_lines:
-    lines.append(f"  {line}")
-
-  # Build result dict - only include registers that are modified
-  result_items = []
-  if modifies_d0: result_items.append("'D0': D0")
-  if modifies_scc: result_items.append("'SCC': SCC")
-  if modifies_vcc: result_items.append("'VCC': VCC")
-  if modifies_exec: result_items.append("'EXEC': EXEC")
-  if has_d1: result_items.append("'D1': D1")
-  if modifies_pc: result_items.append("'PC': PC")
-  lines.append(f"  return {{{', '.join(result_items)}}}\n")
-  return fn_name, '\n'.join(lines)
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# MAIN GENERATION
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def generate_arch(arch: str) -> dict:
-  """Generate enum.py, ins.py and gen_pcode.py for a single architecture."""
-  urls = PDF_URLS[arch]
-  if isinstance(urls, str): urls = [urls]
-
-  print(f"\n{'='*60}\nGenerating {arch}...")
-  print(f"Parsing {len(urls)} PDF(s)...")
-  results = [_parse_single_pdf(url) for url in urls]
-  merged = _merge_results(results) if len(results) > 1 else results[0]
-  doc_name = "+".join(merged["doc_names"]) if len(results) > 1 else merged["doc_name"]
-
-  base_path = Path(f"extra/assembly/amd/autogen/{arch}")
-  base_path.mkdir(parents=True, exist_ok=True)
-  (base_path / "__init__.py").touch()
-
-  # Write enum.py (enums only, no dsl.py dependency)
-  enum_path = base_path / "enum.py"
-  enum_content = _generate_enum_py(merged["enums"], merged["src_enum"], doc_name)
-  enum_path.write_text(enum_content)
-  print(f"Generated {enum_path}: SrcEnum ({len(merged['src_enum'])}) + {len(merged['enums'])} enums")
-
-  # Write ins.py (instruction formats and helpers, imports dsl.py and enum.py)
-  ins_path = base_path / "ins.py"
-  ins_content = _generate_ins_py(merged["formats"], merged["enums"], merged["src_enum"], doc_name).replace("{arch}", arch)
-  ins_path.write_text(ins_content)
-  print(f"Generated {ins_path}: {len(merged['formats'])} formats")
-
-  # Write gen_pcode.py (needs enum.py to exist first for imports)
-  pcode_path = base_path / "gen_pcode.py"
-  pcode_content = _generate_gen_pcode_py(merged["enums"], merged["pseudocode"], arch)
-  pcode_path.write_text(pcode_content)
-  print(f"Generated {pcode_path}: {len(merged['pseudocode'])} instructions")
-
-  return merged
-
-def _generate_arch_wrapper(arch: str):
-  """Wrapper for multiprocessing - returns arch name for ordering."""
-  generate_arch(arch)
-  return arch
-
-def generate_all():
-  """Generate all architectures in parallel."""
-  with ProcessPoolExecutor() as executor:
-    list(executor.map(_generate_arch_wrapper, PDF_URLS.keys()))
-
-if __name__ == "__main__":
-  import argparse
-  parser = argparse.ArgumentParser(description="Generate AMD ISA autogen files from PDF documentation")
-  parser.add_argument("--arch", choices=list(PDF_URLS.keys()) + ["all"], default="rdna3")
-  args = parser.parse_args()
-  if args.arch == "all": generate_all()
-  else: generate_arch(args.arch)
@@ -1,294 +0,0 @@
-#!/usr/bin/env python3
-"""Benchmark comparing Python vs Rust RDNA3 emulators on synthetic and real tinygrad kernels."""
-import ctypes, time, os, struct, cProfile, pstats, io
-from pathlib import Path
-from typing import Callable
-
-# Set AMD=1 before importing tinygrad
-os.environ["AMD"] = "1"
-
-from extra.assembly.amd.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program, step_wave, WaveState, WAVE_SIZE
-
-REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
-if not REMU_PATH.exists():
-  REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.dylib"
-
-def get_rust_remu():
-  """Load the Rust libremu shared library."""
-  if not REMU_PATH.exists(): return None
-  remu = ctypes.CDLL(str(REMU_PATH))
-  remu.run_asm.restype = ctypes.c_int32
-  remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
-                           ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
-  return remu
-
-def count_instructions(kernel: bytes) -> int:
-  """Count instructions in a kernel."""
-  return len(decode_program(kernel))
-
-def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = None):
-  """Allocate buffers and return args pointer + valid ranges."""
-  if init_data is None: init_data = {}
-  buffers = []
-  for i, size in enumerate(buf_sizes):
-    padded = ((size + 15) // 16) * 16 + 16
-    data = init_data.get(i, b'\x00' * padded)
-    data_list = list(data) + [0] * (padded - len(data))
-    buf = (ctypes.c_uint8 * padded)(*data_list[:padded])
-    buffers.append(buf)
-  args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
-  args_ptr = ctypes.addressof(args)
-  ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
-  ranges.add((args_ptr, ctypes.sizeof(args)))
-  return buffers, args, args_ptr, ranges
-
-def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, iterations: int = 5):
-  """Benchmark an emulator and return average time."""
-  gx, gy, gz = global_size
-  lx, ly, lz = local_size
-  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-  lib_ptr = ctypes.addressof(kernel_buf)
-
-  # Warmup
-  run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
-
-  # Timed runs
-  times = []
-  for _ in range(iterations):
-    start = time.perf_counter()
-    result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
-    end = time.perf_counter()
-    if result != 0:
-      print(f"  {name} returned error: {result}")
-      return None
-    times.append(end - start)
-
-  return sum(times) / len(times)
-
-def create_synthetic_kernel(n_ops: int) -> bytes:
-  """Create a synthetic kernel with n_ops vector operations."""
-  instructions = []
-  # VOP2 instructions: v_add_f32, v_mul_f32, v_max_f32, v_min_f32
-  ops = [
-    (0b0000011 << 25) | (1 << 17) | (0 << 9) | 256,  # v_add_f32 v0, v0, v1
-    (0b0001000 << 25) | (1 << 17) | (0 << 9) | 256,  # v_mul_f32 v0, v0, v1
-    (0b0010000 << 25) | (1 << 17) | (0 << 9) | 256,  # v_max_f32 v0, v0, v1
-    (0b0001111 << 25) | (1 << 17) | (0 << 9) | 256,  # v_min_f32 v0, v0, v1
-  ]
-  for i in range(n_ops):
-    instructions.append(ops[i % len(ops)])
-  # S_ENDPGM
-  instructions.append((0b101111111 << 23) | (48 << 16) | 0)
-  return b''.join(struct.pack('<I', inst) for inst in instructions)
-
-def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes]] | None:
-  """Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data)."""
-  try:
-    from tinygrad import Tensor
-    from tinygrad.runtime.support.elf import elf_loader
-    import numpy as np
-    np.random.seed(42)
-
-    ops = {
-      "add": lambda: Tensor.empty(1024) + Tensor.empty(1024),
-      "mul": lambda: Tensor.empty(1024) * Tensor.empty(1024),
-      "matmul_small": lambda: Tensor.empty(16, 16) @ Tensor.empty(16, 16),
-      "matmul_medium": lambda: Tensor.empty(64, 64) @ Tensor.empty(64, 64),
-      "reduce_sum": lambda: Tensor.empty(4096).sum(),
-      "reduce_max": lambda: Tensor.empty(4096).max(),
-      "softmax": lambda: Tensor.empty(256).softmax(),
-      "layernorm": lambda: Tensor.empty(32, 64).layernorm(),
-      "conv2d": lambda: Tensor.empty(1, 4, 16, 16).conv2d(Tensor.empty(4, 4, 3, 3)),
-      "gelu": lambda: Tensor.empty(1024).gelu(),
-      "exp": lambda: Tensor.empty(1024).exp(),
-      "sin": lambda: Tensor.empty(1024).sin(),
-    }
-
-    if op_name not in ops: return None
-    out = ops[op_name]()
-    sched = out.schedule()
-
-    for ei in sched:
-      lowered = ei.lower()
-      if ei.ast.op.name == 'SINK' and lowered.prg and lowered.prg.p.lib:
-        lib = bytes(lowered.prg.p.lib)
-        _, sections, _ = elf_loader(lib)
-        for sec in sections:
-          if sec.name == '.text':
-            buf_sizes = [b.nbytes for b in lowered.bufs]
-            # Get initial data from numpy arrays if available
-            buf_data = {}
-            for i, buf in enumerate(lowered.bufs):
-              if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
-                try: buf_data[i] = bytes(buf.base._buf)
-                except: pass
-            return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data)
-    return None
-  except Exception as e:
-    print(f"  Error getting kernel: {e}")
-    return None
-
-def profile_python_emu(kernel: bytes, global_size, local_size, args_ptr, n_runs: int = 1):
-  """Profile the Python emulator to find bottlenecks."""
-  gx, gy, gz = global_size
-  lx, ly, lz = local_size
-  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-  lib_ptr = ctypes.addressof(kernel_buf)
-
-  pr = cProfile.Profile()
-  pr.enable()
-  for _ in range(n_runs):
-    python_run_asm(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
-  pr.disable()
-
-  s = io.StringIO()
-  ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
-  ps.print_stats(20)
-  return s.getvalue()
-
-def measure_step_rate(kernel: bytes, n_steps: int = 10000) -> float:
-  """Measure raw step_wave() performance (steps per second)."""
-  program = decode_program(kernel)
-  if not program: return 0.0
-
-  st = WaveState()
-  st.exec_mask = 0xffffffff
-  lds = bytearray(65536)
-  n_lanes = 32
-
-  # Reset PC for each measurement
-  start = time.perf_counter()
-  for _ in range(n_steps):
-    st.pc = 0
-    while st.pc in program:
-      result = step_wave(program, st, lds, n_lanes)
-      if result == -1: break
-  elapsed = time.perf_counter() - start
-  return n_steps / elapsed if elapsed > 0 else 0
-
-# Test configurations
-SYNTHETIC_TESTS = [
-  ("synthetic_10ops", 10, (1, 1, 1), (32, 1, 1)),
-  ("synthetic_100ops", 100, (1, 1, 1), (32, 1, 1)),
-  ("synthetic_500ops", 500, (1, 1, 1), (32, 1, 1)),
-  ("synthetic_100ops_4wg", 100, (4, 1, 1), (32, 1, 1)),
-  ("synthetic_100ops_16wg", 100, (16, 1, 1), (32, 1, 1)),
-]
-
-TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
-
-def main():
-  import argparse
-  parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
-  parser.add_argument("--profile", action="store_true", help="Profile Python emulator")
-  parser.add_argument("--synthetic-only", action="store_true", help="Only run synthetic tests")
-  parser.add_argument("--tinygrad-only", action="store_true", help="Only run tinygrad tests")
-  parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
-  args = parser.parse_args()
-
-  rust_remu = get_rust_remu()
-  if rust_remu is None:
-    print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
-    print("Running Python-only benchmarks...\n")
-
-  print("=" * 90)
-  print("RDNA3 Emulator Benchmark: Python vs Rust")
-  print("=" * 90)
-
-  results = []
-
-  # Synthetic workloads
-  if not args.tinygrad_only:
-    print("\n[SYNTHETIC WORKLOADS]")
-    print("-" * 90)
-
-    for name, n_ops, global_size, local_size in SYNTHETIC_TESTS:
-      kernel = create_synthetic_kernel(n_ops)
-      n_insts = count_instructions(kernel)
-      n_workgroups = global_size[0] * global_size[1] * global_size[2]
-      n_threads = local_size[0] * local_size[1] * local_size[2]
-      total_work = n_insts * n_workgroups * n_threads
-
-      print(f"\n{name}: {n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
-
-      buf_sizes = [4096]
-      buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes)
-      set_valid_mem_ranges(ranges)
-
-      # Benchmark
-      py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
-      rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
-
-      if py_time:
-        py_rate = total_work / py_time / 1e6
-        print(f"  Python: {py_time*1000:8.3f} ms  ({py_rate:7.2f} M ops/s)")
-      if rust_time:
-        rust_rate = total_work / rust_time / 1e6
-        speedup = py_time / rust_time if py_time else 0
-        print(f"  Rust:   {rust_time*1000:8.3f} ms  ({rust_rate:7.2f} M ops/s)  [{speedup:.1f}x faster]")
-
-      results.append(("synthetic", name, n_insts, n_workgroups, py_time, rust_time))
-
-  # Tinygrad kernels
-  if not args.synthetic_only:
-    print("\n[TINYGRAD KERNELS]")
-    print("-" * 90)
-
-    for op_name in TINYGRAD_TESTS:
-      print(f"\n{op_name}:", end=" ", flush=True)
-      kernel_info = get_tinygrad_kernel(op_name)
-      if kernel_info is None:
-        print("failed to compile")
-        continue
-
-      kernel, global_size, local_size, buf_sizes, buf_data = kernel_info
-      n_insts = count_instructions(kernel)
-      n_workgroups = global_size[0] * global_size[1] * global_size[2]
-      n_threads = local_size[0] * local_size[1] * local_size[2]
-      total_work = n_insts * n_workgroups * n_threads
-
-      print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
-
-      buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
-      set_valid_mem_ranges(ranges)
-
-      py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
-      rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
-
-      if py_time:
-        py_rate = total_work / py_time / 1e6
-        print(f"  Python: {py_time*1000:8.3f} ms  ({py_rate:7.2f} M ops/s)")
-      if rust_time:
-        rust_rate = total_work / rust_time / 1e6
-        speedup = py_time / rust_time if py_time else 0
-        print(f"  Rust:   {rust_time*1000:8.3f} ms  ({rust_rate:7.2f} M ops/s)  [{speedup:.1f}x faster]")
-
-      results.append(("tinygrad", op_name, n_insts, n_workgroups, py_time, rust_time))
-
-      # Optional profiling
-      if args.profile and py_time:
-        print("\n  [PROFILE - Top 10 functions]")
-        profile_output = profile_python_emu(kernel, global_size, local_size, args_ptr)
-        for line in profile_output.split('\n')[5:15]:
-          if line.strip(): print(f"    {line}")
-
-  # Summary table
-  print("\n" + "=" * 90)
-  print("SUMMARY")
-  print("=" * 90)
-  print(f"{'Type':<10} {'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
-  print("-" * 90)
-
-  for test_type, name, n_insts, n_wgs, py_time, rust_time in results:
-    py_ms = f"{py_time*1000:.3f}" if py_time else "error"
-    if rust_time:
-      rust_ms = f"{rust_time*1000:.3f}"
-      speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
-    else:
-      rust_ms, speedup = "N/A", "N/A"
-    print(f"{test_type:<10} {name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
-
-
-
-if __name__ == "__main__":
-  main()
@@ -1,196 +0,0 @@
-# Usability tests for the RDNA3 ASM DSL
-# These tests demonstrate how the DSL *should* work for a good user experience
-# Currently many of these tests fail - they document desired behavior
-
-import unittest
-from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.dsl import Inst, RawImm, SGPR, VGPR
-
-class TestRegisterSliceSyntax(unittest.TestCase):
-  """
-  Issue: Register slice syntax should use AMD assembly convention (inclusive end).
-
-  In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
-  The DSL should match this convention so that:
-  - s[4:7] gives 4 registers
-  - Disassembler output can be copied directly back into DSL code
-
-  Fix: Change _RegFactory.__getitem__ to use inclusive end:
-    key.stop - key.start + 1  (instead of key.stop - key.start)
-  """
-  def test_register_slice_count(self):
-    # s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
-    reg = s[4:7]
-    self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
-
-  def test_register_slice_roundtrip(self):
-    # Round-trip: DSL -> disasm -> DSL should preserve register count
-    reg = s[4:7]  # 4 registers in AMD convention
-    inst = s_load_b128(reg, s[0:1], NULL, 0)
-    disasm = inst.disasm()
-    # Disasm shows s[4:7] - user should be able to copy this back
-    self.assertIn("s[4:7]", disasm)
-    # And s[4:7] in DSL should give the same 4 registers
-    reg_from_disasm = s[4:7]
-    self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
-
-
-class TestReprReadability(unittest.TestCase):
-  """
-  Issue: repr() leaks internal RawImm type and omits zero-valued fields.
-
-  When you create v_mov_b32_e32(v[0], v[1]), the repr shows:
-    VOP1(op=1, src0=RawImm(257))
-
-  Problems:
-  1. vdst=v[0] is omitted because 0 is treated as "default"
-  2. src0 shows RawImm(257) instead of v[1]
-  3. User sees encoded values (257 = 256 + 1) instead of register names
-
-  Expected repr: VOP1(op=1, vdst=v[0], src0=v[1])
-  """
-  def test_repr_shows_registers_not_raw_imm(self):
-    inst = v_mov_b32_e32(v[0], v[1])
-    # Should show v[1], not RawImm(257)
-    self.assertNotIn("RawImm", repr(inst), "repr should not expose RawImm internal type")
-    self.assertIn("v[1]", repr(inst), "repr should show register name")
-
-  def test_repr_includes_zero_dst(self):
-    inst = v_mov_b32_e32(v[0], v[1])
-    # v[0] is a valid destination register, should be shown
-    self.assertIn("vdst", repr(inst), "repr should include vdst even when 0")
-
-  def test_repr_roundtrip(self):
-    # repr should produce something that can be eval'd back
-    inst = v_mov_b32_e32(v[0], v[1])
-    # This would require repr to output valid Python, e.g.:
-    # "VOP1(op=VOP1Op.V_MOV_B32, vdst=v[0], src0=v[1])"
-    r = repr(inst)
-    # At minimum, it should be human-readable
-    self.assertIn("v[", r, "repr should show register syntax")
-
-
-class TestInstructionEquality(unittest.TestCase):
-  """
-  Issue: No __eq__ method - instruction comparison requires repr() workaround.
-
-  Two identical instructions should compare equal with ==, but currently:
-    inst1 == inst2 returns False
-
-  The test_handwritten.py works around this with:
-    self.assertEqual(repr(self.inst), repr(reasm))
-  """
-  def test_identical_instructions_equal(self):
-    inst1 = v_mov_b32_e32(v[0], v[1])
-    inst2 = v_mov_b32_e32(v[0], v[1])
-    self.assertEqual(inst1, inst2, "identical instructions should be equal")
-
-  def test_different_instructions_not_equal(self):
-    inst1 = v_mov_b32_e32(v[0], v[1])
-    inst2 = v_mov_b32_e32(v[0], v[2])
-    self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
-
-
-class TestVOPDHelperSignature(unittest.TestCase):
-  """
-  Issue: VOPD helper functions have confusing semantics.
-
-  v_dual_mul_f32 is defined as:
-    v_dual_mul_f32 = functools.partial(VOPD, VOPDOp.V_DUAL_MUL_F32)
-
-  This binds VOPDOp.V_DUAL_MUL_F32 to the FIRST positional arg of VOPD.__init__,
-  which is 'opx'. So v_dual_mul_f32 sets the X operation.
-
-  But then test_dual_mul in test_handwritten.py does:
-    v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], ...)
-
-  This passes V_DUAL_MUL_F32 as the SECOND positional arg (opy), making both
-  X and Y operations the same. This is confusing because:
-  1. The function name suggests it handles the X operation
-  2. But you still pass an opcode as the first arg (which becomes opy)
-
-  Expected: Either make the helper fully specify both ops, or make the
-  signature clearer about what the positional arg means.
-  """
-  def test_vopd_helper_opy_should_be_required(self):
-    # Using only keyword args "works" but opy silently defaults to 0
-    inst = v_dual_mul_f32(vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
-    self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32)
-    # Bug: opy defaults to 0 (V_DUAL_FMAC_F32) silently - should require explicit opy
-    # This test documents the bug - it should fail once fixed
-    self.assertNotEqual(inst.opy, VOPDOp.V_DUAL_FMAC_F32, "opy should not silently default to FMAC")
-
-  def test_vopd_helper_positional_arg_is_opy(self):
-    # The first positional arg after the partial becomes opy, not a second opx
-    inst = v_dual_mul_f32(VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
-    self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32)  # From partial
-    self.assertEqual(inst.opy, VOPDOp.V_DUAL_MOV_B32)  # From first positional arg
-
-
-class TestFieldAccessPreservesType(unittest.TestCase):
-  """
-  Issue: Field access loses type information.
-
-  After creating an instruction, accessing fields returns encoded int values:
-    inst = v_mov_b32_e32(v[0], v[1])
-    inst.vdst  # returns 0, not VGPR(0)
-
-  This makes it impossible to round-trip register types through field access.
-  """
-  def test_vdst_returns_register(self):
-    inst = v_mov_b32_e32(v[5], v[1])
-    vdst = inst.vdst
-    # Should return a VGPR, not an int
-    self.assertIsInstance(vdst, (VGPR, int), "vdst should return VGPR or at least be usable")
-    # Ideally: self.assertIsInstance(vdst, VGPR)
-
-  def test_src_returns_register_for_vgpr_source(self):
-    inst = v_mov_b32_e32(v[0], v[1])
-    # src0 is encoded as 257 (256 + 1 for v1)
-    # Ideally it should decode back to v[1]
-    src0_raw = inst._values.get('src0')
-    # Currently returns RawImm(257), should return VGPR(1) or similar
-    self.assertNotIsInstance(src0_raw, RawImm, "source should not be RawImm internally")
-
-
-class TestArgumentDiscoverability(unittest.TestCase):
-  """
-  Issue: No clear signature for positional arguments.
-
-  inspect.signature(s_load_b128) shows: (*args, literal=None, **kwargs)
-
-  Users have no way to know the argument order without reading source code.
-  The order is implicitly defined by the class field definition order.
-
-  Possible fixes:
-  1. Add explicit parameter names to functools.partial
-  2. Generate type stubs with proper signatures
-  3. Add docstrings listing the expected arguments
-  """
-  def test_signature_has_named_params(self):
-    import inspect
-    sig = inspect.signature(s_load_b128)
-    params = list(sig.parameters.keys())
-    # Currently: ['args', 'literal', 'kwargs'] (from *args, literal=None, **kwargs)
-    # Expected: something like ['sdata', 'sbase', 'soffset', 'offset', 'literal']
-    self.assertIn('sdata', params, "signature should show field names")
-
-
-class TestSpecialConstants(unittest.TestCase):
-  """
-  Issue: NULL and other constants are IntEnum values that might be confusing.
-
-  NULL = SrcEnum.NULL = 124, but users might expect NULL to be a special object
-  that clearly represents "no register" rather than a magic number.
-  """
-  def test_null_has_clear_repr(self):
-    # NULL should have a clear string representation
-    self.assertIn("NULL", str(NULL) or repr(NULL), "NULL should be clearly identifiable")
-
-  def test_null_is_distinguishable_from_int(self):
-    # NULL should be distinguishable from the raw integer 124
-    self.assertNotEqual(type(NULL), int, "NULL should not be plain int")
-
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,66 +0,0 @@
-"""Shared test helpers for RDNA3 tests."""
-import shutil
-from dataclasses import dataclass
-
-@dataclass
-class KernelInfo:
-  code: bytes
-  global_size: tuple[int, int, int]
-  local_size: tuple[int, int, int]
-  buf_idxs: list[int]  # indices into shared buffer pool
-  buf_sizes: list[int]  # sizes for each buffer index
-
-# LLVM tool detection (shared across test files)
-def get_llvm_mc():
-  """Find llvm-mc executable, preferring newer versions."""
-  for p in ['llvm-mc', 'llvm-mc-21', 'llvm-mc-20']:
-    if shutil.which(p): return p
-  raise FileNotFoundError("llvm-mc not found")
-
-def get_llvm_objdump():
-  """Find llvm-objdump executable, preferring newer versions."""
-  for p in ['llvm-objdump', 'llvm-objdump-21', 'llvm-objdump-20']:
-    if shutil.which(p): return p
-  raise FileNotFoundError("llvm-objdump not found")
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# EXECUTION CONTEXT (for testing compiled pseudocode)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-class ExecContext:
-  """Context for running compiled pseudocode in tests."""
-  def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=0xffffffff, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
-    from extra.assembly.amd.pcode import Reg, MASK32, MASK64, SliceProxy
-    self._Reg, self._MASK64, self._SliceProxy = Reg, MASK64, SliceProxy
-    self.S0, self.S1, self.S2 = Reg(s0), Reg(s1), Reg(s2)
-    self.D0, self.D1 = Reg(d0), Reg(0)
-    self.SCC, self.VCC, self.EXEC = Reg(scc), Reg(vcc), Reg(exec_mask)
-    self.tmp, self.saveexec = Reg(0), Reg(exec_mask)
-    self.lane, self.laneId, self.literal = lane, lane, literal
-    self.SIMM16, self.SIMM32 = Reg(literal), Reg(literal)
-    self.VGPR = vgprs if vgprs is not None else {}
-    self.SRC0, self.VDST = Reg(src0_idx), Reg(vdst_idx)
-
-  def run(self, code: str):
-    """Execute compiled code."""
-    import extra.assembly.amd.pcode as pcode
-    ns = {k: getattr(pcode, k) for k in dir(pcode) if not k.startswith('_')}
-    # Also include underscore-prefixed helpers that compiled pseudocode uses
-    for k in ['_pack', '_pack32']:
-      if hasattr(pcode, k): ns[k] = getattr(pcode, k)
-    ns.update({
-      'S0': self.S0, 'S1': self.S1, 'S2': self.S2, 'D0': self.D0, 'D1': self.D1,
-      'SCC': self.SCC, 'VCC': self.VCC, 'EXEC': self.EXEC,
-      'EXEC_LO': self._SliceProxy(self.EXEC, 31, 0), 'EXEC_HI': self._SliceProxy(self.EXEC, 63, 32),
-      'tmp': self.tmp, 'saveexec': self.saveexec,
-      'lane': self.lane, 'laneId': self.laneId, 'literal': self.literal,
-      'SIMM16': self.SIMM16, 'SIMM32': self.SIMM32, 'VGPR': self.VGPR, 'SRC0': self.SRC0, 'VDST': self.VDST,
-    })
-    exec(code, ns)
-    def _sync(ctx_reg, ns_val):
-      if isinstance(ns_val, self._Reg): ctx_reg._val = ns_val._val
-      else: ctx_reg._val = int(ns_val) & self._MASK64
-    for name in ('SCC', 'VCC', 'EXEC', 'D0', 'D1', 'tmp', 'saveexec'):
-      if ns.get(name) is not getattr(self, name): _sync(getattr(self, name), ns[name])
-
-  def result(self) -> dict: return {"d0": self.D0._val, "scc": self.SCC._val & 1}
@@ -1,402 +0,0 @@
-# Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
-import unittest, ctypes, os
-from dataclasses import dataclass
-from pathlib import Path
-
-# Set environment before any tinygrad imports to use MOCKGPU
-# This allows generating AMD GPU kernels without requiring real hardware
-os.environ["AMD"] = "1"
-os.environ["MOCKGPU"] = "1"
-os.environ["PYTHON_REMU"] = "1"
-
-from extra.assembly.amd.emu import WaveState, decode_program, step_wave, WAVE_SIZE, set_valid_mem_ranges
-from extra.assembly.amd.test.helpers import KernelInfo
-
-REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
-
-def _is_f32_nan(bits: int) -> bool:
-  """Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
-  return (bits & 0x7f800000) == 0x7f800000 and (bits & 0x007fffff) != 0
-
-def _vals_equal(a: int, b: int) -> bool:
-  """Compare two 32-bit values, treating all NaN bit patterns as equal."""
-  if a == b: return True
-  return _is_f32_nan(a) and _is_f32_nan(b)
-
-@dataclass
-class StateSnapshot:
-  pc: int
-  scc: int
-  vcc: int
-  exec_mask: int
-  sgpr: list[int]
-  vgpr: list[list[int]]
-
-  def diff(self, other: 'StateSnapshot', n_lanes: int, arrow: str = " vs ") -> list[str]:
-    """Return list of differences between two states."""
-    diffs = []
-    if self.pc != other.pc: diffs.append(f"pc: {self.pc}{arrow}{other.pc}")
-    if self.scc != other.scc: diffs.append(f"scc: {self.scc}{arrow}{other.scc}")
-    if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x}{arrow}0x{other.vcc:08x}")
-    if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x}{arrow}0x{other.exec_mask:08x}")
-    for i, (a, b) in enumerate(zip(self.sgpr, other.sgpr)):
-      # Skip VCC_LO/HI (106/107) and EXEC_LO/HI (126/127) as they alias vcc/exec_mask which are compared separately
-      if i in (106, 107, 126, 127): continue
-      if not _vals_equal(a, b): diffs.append(f"sgpr[{i}]: 0x{a:08x}{arrow}0x{b:08x}")
-    for lane in range(n_lanes):
-      for i, (a, b) in enumerate(zip(self.vgpr[lane], other.vgpr[lane])):
-        if not _vals_equal(a, b): diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x}{arrow}0x{b:08x}")
-    return diffs
-
-class CStateSnapshot(ctypes.Structure):
-  _fields_ = [("pc", ctypes.c_uint32), ("scc", ctypes.c_uint32), ("vcc", ctypes.c_uint32), ("exec_mask", ctypes.c_uint32),
-              ("sgpr", ctypes.c_uint32 * 128), ("vgpr", (ctypes.c_uint32 * 256) * 32)]
-
-  def to_snapshot(self) -> StateSnapshot:
-    return StateSnapshot(pc=self.pc, scc=self.scc, vcc=self.vcc, exec_mask=self.exec_mask,
-                         sgpr=list(self.sgpr), vgpr=[list(self.vgpr[i]) for i in range(32)])
-
-class RustEmulator:
-  def __init__(self):
-    self.lib = ctypes.CDLL(str(REMU_PATH))
-    self.lib.wave_create.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_create.restype = ctypes.c_void_p
-    self.lib.wave_step.argtypes = [ctypes.c_void_p]
-    self.lib.wave_step.restype = ctypes.c_int32
-    self.lib.wave_get_snapshot.argtypes = [ctypes.c_void_p, ctypes.POINTER(CStateSnapshot)]
-    self.lib.wave_set_sgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_set_vgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_init_lds.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
-    self.lib.wave_free.argtypes = [ctypes.c_void_p]
-    self.ctx = None
-
-  def create(self, kernel: bytes, n_lanes: int):
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    self.ctx = self.lib.wave_create(ctypes.addressof(kernel_buf), len(kernel), n_lanes)
-    self._kernel_buf = kernel_buf
-
-  def step(self) -> int: return self.lib.wave_step(self.ctx)
-  def set_sgpr(self, idx: int, val: int): self.lib.wave_set_sgpr(self.ctx, idx, val)
-  def set_vgpr(self, lane: int, idx: int, val: int): self.lib.wave_set_vgpr(self.ctx, lane, idx, val)
-  def init_lds(self, size: int): self.lib.wave_init_lds(self.ctx, size)
-
-  def get_snapshot(self) -> StateSnapshot:
-    snap = CStateSnapshot()
-    self.lib.wave_get_snapshot(self.ctx, ctypes.byref(snap))
-    return snap.to_snapshot()
-
-  def free(self):
-    if self.ctx: self.lib.wave_free(self.ctx); self.ctx = None
-
-class PythonEmulator:
-  def __init__(self):
-    self.state: WaveState | None = None
-    self.program: dict | None = None
-    self.lds: bytearray | None = None
-    self.n_lanes = 0
-
-  def create(self, kernel: bytes, n_lanes: int):
-    self.program = decode_program(kernel)
-    self.state = WaveState()
-    self.state.exec_mask = (1 << n_lanes) - 1
-    self.lds = bytearray(65536)
-    self.n_lanes = n_lanes
-
-  def step(self) -> int:
-    assert self.program is not None and self.state is not None and self.lds is not None
-    return step_wave(self.program, self.state, self.lds, self.n_lanes)
-  def set_sgpr(self, idx: int, val: int):
-    assert self.state is not None
-    self.state.sgpr[idx] = val & 0xffffffff
-  def set_vgpr(self, lane: int, idx: int, val: int):
-    assert self.state is not None
-    self.state.vgpr[lane][idx] = val & 0xffffffff
-
-  def get_snapshot(self) -> StateSnapshot:
-    assert self.state is not None
-    return StateSnapshot(pc=self.state.pc, scc=self.state.scc, vcc=self.state.vcc & 0xffffffff,
-                         exec_mask=self.state.exec_mask & 0xffffffff, sgpr=list(self.state.sgpr),
-                         vgpr=[list(self.state.vgpr[i]) for i in range(WAVE_SIZE)])
-
-def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
-                      program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
-                      max_workgroups: int = 8) -> tuple[bool, str, int]:
-  """Run a single kernel through both emulators. Returns (success, message, total_steps)."""
-  gx, gy, gz = global_size
-  total_steps = 0
-  wg_count = 0
-
-  for gidz in range(gz):
-    for gidy in range(gy):
-      for gidx in range(gx):
-        if wg_count >= max_workgroups: return True, f"Completed {wg_count} workgroups (limit reached)", total_steps
-        wg_count += 1
-        rust = RustEmulator()
-        python = PythonEmulator()
-        rust.create(kernel, n_lanes)
-        python.create(kernel, n_lanes)
-
-        # Initialize LDS (64KB, standard size for AMD GPUs)
-        rust.init_lds(65536)
-
-        for emu in (rust, python):
-          emu.set_sgpr(0, args_ptr & 0xffffffff)
-          emu.set_sgpr(1, (args_ptr >> 32) & 0xffffffff)
-          emu.set_sgpr(13, gidx)
-          emu.set_sgpr(14, gidy)
-          emu.set_sgpr(15, gidz)
-
-        step = 0
-        trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
-        try:
-          while step < max_steps:
-            rust_before = rust.get_snapshot()
-            python_before = python.get_snapshot()
-
-            inst = program.get(python_before.pc)
-            inst_str = inst.disasm() if inst else f"unknown at PC={python_before.pc}"
-            trace.append((step, python_before.pc, inst_str, rust_before, python_before))
-            if len(trace) > trace_len: trace.pop(0)
-
-            if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
-
-            # Instructions with known Rust emulator bugs - sync Python to Rust after execution
-            # v_div_scale/v_div_fixup: Rust has different VCC handling
-            # v_cvt_f16_f32: Rust clears high 16 bits, but hardware (and Python) preserves them
-            sync_after = any(x in inst_str for x in ('v_div_scale_f32', 'v_div_scale_f64', 'v_div_fixup_f32', 'v_div_fixup_f64',
-                                                      'v_cvt_f16_f32'))
-            diffs = rust_before.diff(python_before, n_lanes)
-            if diffs:
-              trace_lines = []
-              for idx, (s, pc, d, rb, pb) in enumerate(trace):
-                trace_lines.append(f"    step {s}: PC={pc:3d} {d}")
-                if idx < len(trace) - 1:
-                  next_rb, next_pb = trace[idx + 1][3:5]
-                  rust_diffs = rb.diff(next_rb, n_lanes, "->")
-                  python_diffs = pb.diff(next_pb, n_lanes, "->")
-                  if rust_diffs: trace_lines.append(f"             rust:   {', '.join(rust_diffs[:5])}")
-                  if python_diffs: trace_lines.append(f"             python: {', '.join(python_diffs[:5])}")
-                  elif rust_diffs: trace_lines.append(f"             python: (no changes)")
-                else:
-                  # Last traced instruction - compare with current state
-                  rust_diffs = rb.diff(rust_before, n_lanes, "->")
-                  python_diffs = pb.diff(python_before, n_lanes, "->")
-                  if rust_diffs: trace_lines.append(f"             rust:   {', '.join(rust_diffs[:5])}")
-                  if python_diffs: trace_lines.append(f"             python: {', '.join(python_diffs[:5])}")
-                  elif rust_diffs: trace_lines.append(f"             python: (no changes)")
-              trace_str = "\n".join(trace_lines)
-              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ (rust vs python):\n  " + "\n  ".join(diffs[:10]) + f"\n  Recent instructions:\n{trace_str}", total_steps
-
-            rust_result = rust.step()
-            python_result = python.step()
-
-            if rust_result != python_result:
-              # Rust returns 1 for unsupported instructions - skip test
-              if rust_result == 1 and python_result == 0:
-                raise unittest.SkipTest(f"Rust emulator doesn't support instruction: {inst_str}")
-              trace_str = "\n".join(f"    step {s}: PC={pc:3d} {d}" for s, pc, d, _, _ in trace)
-              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: rust={rust_result}, python={python_result}, inst={inst_str}\n  Recent instructions:\n{trace_str}", total_steps
-
-            # Sync Python state to Rust after instructions with known Rust emulator differences
-            if sync_after:
-              rust_after = rust.get_snapshot()
-              for i in range(128): python.set_sgpr(i, rust_after.sgpr[i])
-              for lane in range(n_lanes):
-                for i in range(256): python.set_vgpr(lane, i, rust_after.vgpr[lane][i])
-              assert python.state is not None
-              python.state.pc, python.state.scc, python.state.vcc, python.state.exec_mask = rust_after.pc, rust_after.scc, rust_after.vcc, rust_after.exec_mask
-
-            if rust_result == -1:
-              total_steps += step + 1
-              break
-            if rust_result == 1:
-              total_steps += step + 1
-              break
-            if rust_result < 0 and rust_result != -2:
-              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: error code {rust_result}", total_steps
-
-            step += 1
-          else:
-            return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Max steps ({max_steps}) reached", total_steps
-        finally:
-          rust.free()
-
-  return True, f"Completed {gx*gy*gz} workgroups", total_steps
-
-def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int, int], max_steps: int = 1000,
-                                    debug: bool = False, trace_len: int = 10, buf_data: dict[int, bytes] | None = None) -> tuple[bool, str]:
-  """Run all kernels through both emulators with shared buffer pool."""
-  if buf_data is None: buf_data = {}
-
-  # Allocate shared buffer pool with padding for over-reads (GPU loads up to 16 bytes at once)
-  buf_id_to_ptr: dict[int, int] = {}
-  buffers = []
-  for buf_id, size in buf_pool.items():
-    padded_size = ((size + 15) // 16) * 16 + 16  # round up to 16 bytes + extra padding
-    # Initialize with data from COPY if available
-    init_data = buf_data.get(buf_id, b'\x00' * padded_size)
-    init_list = list(init_data) + [0] * (padded_size - len(init_data))
-    buf = (ctypes.c_uint8 * padded_size)(*init_list[:padded_size])
-    buffers.append((buf, padded_size))
-    buf_id_to_ptr[buf_id] = ctypes.addressof(buf)
-
-  # Set up valid memory ranges
-  ranges = {(ctypes.addressof(b), size) for b, size in buffers}
-
-  total_steps = 0
-  for ki, kernel in enumerate(kernels):
-    # Create args array for this kernel's buffers
-    args = (ctypes.c_uint64 * len(kernel.buf_idxs))(*[buf_id_to_ptr[bid] for bid in kernel.buf_idxs])
-    args_ptr = ctypes.addressof(args)
-
-    # Update valid ranges to include this args array
-    kernel_ranges = ranges | {(args_ptr, ctypes.sizeof(args))}
-    set_valid_mem_ranges(kernel_ranges)
-
-    program = decode_program(kernel.code)
-    n_lanes = kernel.local_size[0] * kernel.local_size[1] * kernel.local_size[2]
-
-    ok, msg, steps = run_single_kernel(
-      kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
-      program, max_steps, debug, trace_len, ki
-    )
-    total_steps += steps
-    if not ok:
-      return False, msg
-
-  return True, f"Completed {len(kernels)} kernels, {total_steps} total steps"
-
-def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list, max_steps: int = 1000, debug: bool = False,
-                                   global_size: tuple[int, int, int] = (1, 1, 1), trace_len: int = 10) -> tuple[bool, str]:
-  """Run both emulators with memory set up for tinygrad kernels, executing all workgroups. Legacy wrapper."""
-  # Allocate buffers
-  buffers = []
-  for size in buf_sizes:
-    buf = (ctypes.c_uint8 * size)(*[0] * size)
-    buffers.append(buf)
-
-  # Create args array with buffer pointers
-  args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
-  args_ptr = ctypes.addressof(args)
-
-  # Set up valid memory ranges for Python emulator
-  ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
-  ranges.add((args_ptr, ctypes.sizeof(args)))
-  set_valid_mem_ranges(ranges)
-
-  program = decode_program(kernel)
-  ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, program, max_steps, debug, trace_len)
-  return ok, msg
-
-def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int], dict[int, bytes]]:
-  """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
-  from tinygrad import Tensor
-  from tinygrad.runtime.support.elf import elf_loader
-
-  out = op_fn(Tensor)
-  sched = out.schedule()
-  kernels = []
-  buf_pool: dict[int, int] = {}  # buffer id -> size
-  buf_data: dict[int, bytes] = {}  # buffer id -> initial data from COPY
-
-  for ei in sched:
-    lowered = ei.lower()
-    if ei.ast.op.name == 'COPY':
-      # Handle COPY: extract source data to initialize destination buffer
-      if len(lowered.bufs) >= 2:
-        dst_buf, src_buf = lowered.bufs[0], lowered.bufs[1]
-        dst_id = id(dst_buf)
-        if dst_id not in buf_pool:
-          buf_pool[dst_id] = dst_buf.nbytes
-        # Get source data if it's from numpy/CPU
-        if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'):
-          src_data = bytes(src_buf.base._buf)
-          buf_data[dst_id] = src_data
-    elif ei.ast.op.name == 'SINK':
-      if lowered.prg and lowered.prg.p.lib:
-        lib = bytes(lowered.prg.p.lib)
-        _, sections, _ = elf_loader(lib)
-        for sec in sections:
-          if sec.name == '.text':
-            buf_idxs = []
-            buf_sizes = []
-            for b in lowered.bufs:
-              buf_id = id(b)
-              if buf_id not in buf_pool:
-                buf_pool[buf_id] = b.nbytes
-              buf_idxs.append(buf_id)
-              buf_sizes.append(b.nbytes)
-            kernels.append(KernelInfo(
-              code=bytes(sec.content),
-              global_size=tuple(lowered.prg.p.global_size),
-              local_size=tuple(lowered.prg.p.local_size),
-              buf_idxs=buf_idxs,
-              buf_sizes=buf_sizes
-            ))
-  if not kernels: raise RuntimeError("No kernel found")
-  return kernels, buf_pool, buf_data
-
-def get_kernel_from_tinygrad(op_fn) -> tuple[bytes, tuple[int, int, int], tuple[int, int, int], list]:
-  """Compile a tinygrad operation and extract the last (main) kernel binary. Legacy wrapper."""
-  kernels, _, _ = get_kernels_from_tinygrad(op_fn)
-  k = kernels[-1]
-  return k.code, k.global_size, k.local_size, k.buf_sizes
-
-class TestTinygradKernels(unittest.TestCase):
-  """Compare emulators on real tinygrad-compiled kernels."""
-
-  def _test_kernel(self, op_fn, max_steps=10000):
-    kernels, buf_pool, buf_data = get_kernels_from_tinygrad(op_fn)
-    ok, msg = compare_emulators_multi_kernel(kernels, buf_pool, max_steps=max_steps, buf_data=buf_data)
-    self.assertTrue(ok, msg)
-
-  # Basic ops - consolidated tests covering key instruction patterns
-  def test_unary_ops(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu().exp().log().sqrt().reciprocal())
-  def test_binary_ops(self): self._test_kernel(lambda T: (T([1.0, 2.0]) + T([3.0, 4.0])) * T([0.5, 0.5]) - T([1.0, 1.0]))
-  def test_trig(self): self._test_kernel(lambda T: T([0.1, 1.0, 3.14, -1.0]*8).sin() + T([0.1, 1.0, 3.14, -1.0]*8).cos())
-  def test_compare(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
-  def test_bitwise(self): self._test_kernel(lambda T: (T([0xF0, 0x0F, 0xFF]*11).int() & T([0x0F, 0x0F, 0x00]*11).int()) | T([1]*33).int())
-  def test_int_ops(self): self._test_kernel(lambda T: ((T.empty(64).int() + T.empty(64).int()) * T.empty(64).int()).float())
-
-  # Reductions
-  def test_reduce(self): self._test_kernel(lambda T: T.empty(64).sum() + T.empty(64).max())
-  def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
-
-  # Matmul
-  def test_gemm(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=100000)
-  @unittest.skip("Rust emulator crashes on this kernel (assertion failure in thread.rs)")
-  def test_gemm_fp16(self): self._test_kernel(lambda T: T.empty(16, 16).half() @ T.empty(16, 16).half(), max_steps=100000)
-
-  # Complex ops
-  def test_softmax(self): self._test_kernel(lambda T: T.empty(16).softmax())
-  def test_layernorm(self): self._test_kernel(lambda T: T.empty(8, 8).layernorm())
-
-  # Memory patterns
-  def test_memory(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous() + T.empty(4, 1).expand(4, 4))
-
-  # Cast ops
-  def test_cast(self): self._test_kernel(lambda T: T.empty(32).half().float() + T.empty(32).int().float())
-
-  # Pooling - regression for VCC wave32 mode
-  def test_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4)) + T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4)))
-
-  # Convolution
-  def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 2, 8, 8).conv2d(T.empty(2, 2, 3, 3)), max_steps=50000)
-
-  # Regression tests
-  def test_topk(self): self._test_kernel(lambda T: T.empty(64).topk(3)[0])
-  def test_interpolate(self): self._test_kernel(lambda T: T.empty(1,2,16,16).relu().cast('uint8').interpolate((8,8), mode="linear"))
-  def test_index_int64(self):
-    from tinygrad import dtypes
-    self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
-  def test_gelu(self): self._test_kernel(lambda T: T.empty(32, 32).gelu())
-  def test_cross_entropy(self):
-    import numpy as np
-    np.random.seed(0)
-    classes = np.random.randint(0, 10, (16,), dtype=np.int32).tolist()
-    x_np = np.random.randn(16, 10).astype(np.float32)
-    self._test_kernel(lambda T: (T(x_np.tolist()).reshape(16,10) + 0).cross_entropy((T(classes).int().reshape(16) + 0)))
-  def test_isinf(self): self._test_kernel(lambda T: T([float('-inf'), 0., float('inf'), 1.1]*8).isinf())
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,407 +0,0 @@
-#!/usr/bin/env python3
-"""Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
-import unittest
-from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.dsl import encode_src, RawImm
-from extra.assembly.amd.asm import detect_format
-
-class TestMUBUF(unittest.TestCase):
-  """Test MUBUF (buffer) instructions."""
-
-  def test_buffer_load_b32_basic(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_idxen(self):
-    # buffer_load_b32 v5, v0, s[8:11], s3 idxen offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, idxen=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]))
-
-  def test_buffer_load_b32_offen(self):
-    # buffer_load_b32 v5, v0, s[8:11], s3 offen offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, offen=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]))
-
-  def test_buffer_load_b32_glc(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc
-    # GFX11: encoding: [0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_slc(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 slc
-    # GFX11: encoding: [0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, slc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_dlc(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 dlc
-    # GFX11: encoding: [0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_all_flags(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc slc dlc
-    # GFX11: encoding: [0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1, slc=1, dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_store_b32(self):
-    # buffer_store_b32 v1, off, s[12:15], s4 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]
-    inst = buffer_store_b32(vdata=v[1], vaddr=v[0], srsrc=s[12:16], soffset=s[4], offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]))
-
-  def test_buffer_load_b64(self):
-    # buffer_load_b64 v[5:6], off, s[8:11], s3 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b64(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_soffset_m0(self):
-    # buffer_load_b32 v5, off, s[8:11], m0 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=M0, offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]))
-
-  def test_buffer_load_soffset_inline_const(self):
-    # buffer_load_b32 v5, off, s[8:11], 0 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=0, offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]))
-
-  def test_buffer_disasm_roundtrip(self):
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
-    decoded = MUBUF.from_bytes(inst.to_bytes())
-    self.assertEqual(decoded.to_bytes(), inst.to_bytes())
-
-
-class TestMTBUF(unittest.TestCase):
-  """Test MTBUF (typed buffer) instructions."""
-
-  def test_tbuffer_load_format_x(self):
-    # tbuffer_load_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
-    # BUF_FMT_32_FLOAT = 22
-    # GFX11: encoding: [0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]
-    inst = tbuffer_load_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]))
-
-  def test_tbuffer_store_format_x(self):
-    # tbuffer_store_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
-    # BUF_FMT_32_FLOAT = 22
-    # GFX11: encoding: [0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]
-    inst = tbuffer_store_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]))
-
-  def test_tbuffer_load_format_xy(self):
-    # tbuffer_load_format_xy v[5:6], off, s[8:11], s3 format:[BUF_FMT_32_32_FLOAT] offset:4095
-    # BUF_FMT_32_32_FLOAT = 50
-    # GFX11: encoding: [0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]
-    inst = tbuffer_load_format_xy(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=50)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]))
-
-
-class TestMIMG(unittest.TestCase):
-  """Test MIMG (image) instructions."""
-
-  def test_image_load_2d(self):
-    # image_load v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
-    # GFX11: encoding: [0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_load(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1)  # dim=1 is SQ_RSRC_IMG_2D
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
-
-  def test_image_store_2d(self):
-    # image_store v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
-    # GFX11: encoding: [0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_store(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]))
-
-  def test_image_load_1d(self):
-    # image_load v[0:3], v4, s[0:7] dmask:0xf dim:SQ_RSRC_IMG_1D
-    # GFX11: encoding: [0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_load(vdata=v[0:4], vaddr=v[4], srsrc=s[0:8], dmask=0xf, dim=0)  # dim=0 is SQ_RSRC_IMG_1D
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
-
-  def test_image_sample(self):
-    # image_sample v[0:3], v[4:5], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
-    # GFX11: encoding: [0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]
-    inst = image_sample(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], ssamp=s[8:12], dmask=0xf, dim=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]))
-
-  def test_image_load_d16(self):
-    # image_load v[0:1], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D d16
-    # GFX11: encoding: [0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_load(vdata=v[0:2], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1, d16=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]))
-
-
-class TestEXP(unittest.TestCase):
-  """Test EXP (export) instructions."""
-
-  def test_exp_mrt0(self):
-    # exp mrt0 v0, v1, v2, v3
-    # GFX11: encoding: [0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]
-    inst = EXP(en=0xf, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[2], vsrc3=v[3])
-    self.assertEqual(inst.to_bytes(), bytes([0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]))
-
-  def test_exp_mrtz(self):
-    # exp mrtz v4, v3, v2, v1
-    # GFX11: encoding: [0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]
-    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]))
-
-  def test_exp_mrtz_done(self):
-    # exp mrtz v4, v3, v2, v1 done
-    # GFX11: encoding: [0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x01]
-    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[3], done=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x03]))
-
-  def test_exp_partial_mask(self):
-    # exp mrt0 v0, v1, off, off (en=0x3, only first two components)
-    # GFX11: encoding: [0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]
-    inst = EXP(en=0x3, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[0], vsrc3=v[0])
-    self.assertEqual(inst.to_bytes(), bytes([0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]))
-
-  def test_exp_row_en(self):
-    # exp mrtz v4, v3, v2, v1 row_en
-    # GFX11: encoding: [0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]
-    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1], row=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]))
-
-
-class TestDS(unittest.TestCase):
-  """Test DS (data share / LDS) instructions."""
-
-  def test_ds_store_b32(self):
-    # ds_store_b32 v0, v1
-    # GFX11: encoding: [0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_store_b32(addr=v[0], data0=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
-
-  def test_ds_load_b32(self):
-    # ds_load_b32 v0, v1
-    # GFX11: encoding: [0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]
-    inst = ds_load_b32(vdst=v[0], addr=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]))
-
-  def test_ds_store_b32_offset(self):
-    # ds_store_b32 v0, v1 offset:64
-    # GFX11: encoding: [0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_store_b32(addr=v[0], data0=v[1], offset0=64)
-    self.assertEqual(inst.to_bytes(), bytes([0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
-
-  def test_ds_load_b64(self):
-    # ds_load_b64 v[0:1], v2
-    # GFX11: encoding: [0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]
-    inst = ds_load_b64(vdst=v[0:2], addr=v[2])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]))
-
-  def test_ds_add_u32(self):
-    # ds_add_u32 v0, v1
-    # GFX11: encoding: [0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_add_u32(addr=v[0], data0=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]))
-
-  def test_ds_store_b32_gds(self):
-    # ds_store_b32 v0, v1 gds
-    # GFX11: encoding: [0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_store_b32(addr=v[0], data0=v[1], gds=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]))
-
-
-class TestVOP3(unittest.TestCase):
-  """Test VOP3 (3-operand vector) instructions."""
-
-  def test_v_fma_f32(self):
-    # v_fma_f32 v0, v1, v2, v3
-    # GFX11: encoding: [0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]
-    inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]))
-
-  def test_v_mad_f32(self):
-    # v_fmac_f32_e64 v0, v1, v2 (fmac is fma with implicit dst as src2)
-    # Use v_fma_f32 with vdst == src2
-    inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[0])
-    self.assertEqual(inst.to_bytes()[:4], bytes([0x00,0x00,0x13,0xd6]))
-
-  def test_v_add3_u32(self):
-    # v_add3_u32 v0, v1, v2, v3
-    # GFX11: encoding: [0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]
-    inst = v_add3_u32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]))
-
-
-class TestFLAT(unittest.TestCase):
-  """Test FLAT/GLOBAL/SCRATCH memory instructions."""
-
-  def test_global_load_b32(self):
-    # global_load_b32 v0, v[1:2], off (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]
-    inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]))
-
-  def test_global_store_b32(self):
-    # global_store_b32 v[0:1], v2, off (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]
-    inst = global_store_b32(addr=v[0:2], data=v[2], saddr=OFF)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]))
-
-  def test_global_load_b32_saddr(self):
-    # global_load_b32 v0, v1, s[0:1] (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]
-    inst = global_load_b32(vdst=v[0], addr=v[1], saddr=s[0:2])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]))
-
-  def test_global_load_b32_offset(self):
-    # global_load_b32 v0, v[1:2], off offset:256 (seg=2 for global)
-    # GFX11: encoding: [0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]
-    inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF, offset=256)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]))
-
-  def test_global_load_b64(self):
-    # global_load_b64 v[0:1], v[2:3], off (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]
-    inst = global_load_b64(vdst=v[0:2], addr=v[2:4], saddr=OFF)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]))
-
-
-class TestSMEM(unittest.TestCase):
-  """Test SMEM (scalar memory) instructions - regression tests for glc/dlc bit positions."""
-
-  def test_smem_dlc_bit_position(self):
-    # s_load_b32 s5, s[2:3], s0 dlc - tests that DLC is at bit 13 (not bit 14)
-    # GFX11: encoding: [0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]
-    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]))
-
-  def test_smem_glc_bit_position(self):
-    # s_load_b32 s5, s[2:3], s0 glc - tests that GLC is at bit 14 (not bit 16)
-    # GFX11: encoding: [0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]
-    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]))
-
-  def test_smem_glc_dlc_combined(self):
-    # s_load_b32 s5, s[2:3], s0 glc dlc - tests both flags together
-    # GFX11: encoding: [0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]
-    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1, dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]))
-
-  def test_smem_disasm_roundtrip_dlc(self):
-    # Test that disassembly/reassembly preserves DLC bit correctly
-    data = bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00])
-    decoded = SMEM.from_bytes(data)
-    self.assertEqual(decoded.to_bytes(), data)
-
-  def test_smem_disasm_roundtrip_glc_dlc(self):
-    # Test that disassembly/reassembly preserves GLC+DLC bits correctly
-    data = bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00])
-    decoded = SMEM.from_bytes(data)
-    self.assertEqual(decoded.to_bytes(), data)
-
-
-class TestVOP3Literal(unittest.TestCase):
-  """Test VOP3 literal handling - regression tests for Inst64 literal encoding."""
-
-  def test_vop3_with_literal(self):
-    # v_add3_u32 v5, vcc_hi, 0xaf123456, v255
-    # GFX11: encoding: [0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf]
-    from extra.assembly.amd.dsl import RawImm
-    inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=RawImm(107), src1=0xaf123456, src2=v[255])
-    expected = bytes([0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf])
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_vop3_literal_null_operand(self):
-    # v_add3_u32 v5, null, exec_lo, 0xaf123456
-    # GFX11: encoding: [0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf]
-    from extra.assembly.amd.dsl import RawImm
-    inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=NULL, src1=RawImm(126), src2=0xaf123456)
-    expected = bytes([0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf])
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_vop3p_with_literal(self):
-    # Test VOP3P literal encoding (also uses Inst64)
-    from extra.assembly.amd.dsl import RawImm
-    inst = VOP3P(VOP3POp.V_PK_ADD_F16, vdst=v[5], src0=RawImm(240), src1=0x12345678, src2=v[0])
-    self.assertEqual(len(inst.to_bytes()), 12)  # 8 bytes + 4 byte literal
-
-
-class TestDetectFormat(unittest.TestCase):
-  """Test detect_format uses encoding from autogen classes."""
-
-  def test_detect_sopp(self):
-    self.assertEqual(detect_format(s_endpgm().to_bytes()), SOPP)
-    self.assertEqual(detect_format(s_nop(0).to_bytes()), SOPP)
-    self.assertEqual(detect_format(s_barrier().to_bytes()), SOPP)
-
-  def test_detect_sop1(self):
-    self.assertEqual(detect_format(s_mov_b32(s[0], 0).to_bytes()), SOP1)
-    self.assertEqual(detect_format(s_mov_b64(s[0:1], 0).to_bytes()), SOP1)
-
-  def test_detect_sop2(self):
-    self.assertEqual(detect_format(s_add_u32(s[0], s[1], s[2]).to_bytes()), SOP2)
-    self.assertEqual(detect_format(s_mul_i32(s[0], s[1], s[2]).to_bytes()), SOP2)
-
-  def test_detect_sopc(self):
-    self.assertEqual(detect_format(s_cmp_eq_i32(s[0], s[1]).to_bytes()), SOPC)
-
-  def test_detect_sopk(self):
-    self.assertEqual(detect_format(s_movk_i32(s[0], 0x1234).to_bytes()), SOPK)
-
-  def test_detect_vop1(self):
-    self.assertEqual(detect_format(v_mov_b32_e32(v[0], 0).to_bytes()), VOP1)
-    self.assertEqual(detect_format(v_rcp_f32_e32(v[0], v[1]).to_bytes()), VOP1)
-
-  def test_detect_vop2(self):
-    self.assertEqual(detect_format(v_add_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
-    self.assertEqual(detect_format(v_mul_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
-
-  def test_detect_vopc(self):
-    self.assertEqual(detect_format(v_cmp_eq_f32_e32(v[0], v[1]).to_bytes()), VOPC)
-    self.assertEqual(detect_format(v_cmp_lt_i32_e32(v[0], v[1]).to_bytes()), VOPC)
-
-  def test_detect_vop3(self):
-    self.assertEqual(detect_format(v_add_f32_e64(v[0], v[1], v[2]).to_bytes()), VOP3)
-    self.assertEqual(detect_format(v_fma_f32(v[0], v[1], v[2], v[3]).to_bytes()), VOP3)
-
-  def test_detect_vop3p(self):
-    self.assertEqual(detect_format(VOP3P(VOP3POp.V_PK_ADD_F16, v[0], v[1], v[2], v[3]).to_bytes()), VOP3P)
-
-  def test_detect_smem(self):
-    self.assertEqual(detect_format(s_load_b32(s[0], s[2:3], 0).to_bytes()), SMEM)
-    self.assertEqual(detect_format(s_load_b64(s[0:1], s[2:3], s[5]).to_bytes()), SMEM)
-
-  def test_detect_ds(self):
-    self.assertEqual(detect_format(ds_load_b32(v[0], v[1]).to_bytes()), DS)
-    self.assertEqual(detect_format(ds_store_b32(v[0], v[1]).to_bytes()), DS)
-
-  def test_detect_flat(self):
-    self.assertEqual(detect_format(global_load_b32(v[0], v[1:3], RawImm(124)).to_bytes()), FLAT)
-    self.assertEqual(detect_format(global_store_b32(v[0:2], v[2], RawImm(124)).to_bytes()), FLAT)
-
-  def test_detect_mubuf(self):
-    self.assertEqual(detect_format(buffer_load_b32(v[0], v[1], s[0:4], s[5]).to_bytes()), MUBUF)
-
-  def test_detect_mtbuf(self):
-    self.assertEqual(detect_format(tbuffer_load_format_x(v[0], v[1], s[0:4], s[5], format=22).to_bytes()), MTBUF)
-
-  def test_detect_mimg(self):
-    self.assertEqual(detect_format(image_load(v[0:4], v[4:6], s[0:8], dmask=0xf, dim=1).to_bytes()), MIMG)
-
-  def test_detect_exp(self):
-    self.assertEqual(detect_format(EXP(en=0xf, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[2], vsrc3=v[3]).to_bytes()), EXP)
-
-  def test_detect_vopd(self):
-    inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=0, srcy0=0)
-    self.assertEqual(detect_format(inst.to_bytes()), VOPD)
-
-  def test_detect_vinterp(self):
-    inst = VINTERP(VINTERPOp.V_INTERP_P10_F32, vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
-    self.assertEqual(detect_format(inst.to_bytes()), VINTERP)
-
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,178 +0,0 @@
-# do not change these tests. we need to fix bugs to make them pass
-# the Inst constructor should be looking at the types of the fields to correctly set the value
-
-import unittest, struct
-from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.dsl import Inst
-from extra.assembly.amd.asm import asm
-from extra.assembly.amd.test.test_roundtrip import compile_asm
-
-class TestIntegration(unittest.TestCase):
-  inst: Inst
-  def tearDown(self):
-    if not hasattr(self, 'inst'): return
-    b = self.inst.to_bytes()
-    st = self.inst.disasm()
-    reasm = asm(st)
-    desc = f"{st:25s} {self.inst} {b!r} {reasm}"
-    self.assertEqual(b, compile_asm(st), desc)
-    # TODO: this compare should work for valid things
-    #self.assertEqual(self.inst, reasm)
-    self.assertEqual(repr(self.inst), repr(reasm))
-    print(desc)
-
-  def test_load_b128(self):
-    self.inst = s_load_b128(s[4:7], s[0:1], NULL, 0)
-
-  def test_load_b128_wrong_size(self):
-    # this should have to be 4 regs on the loaded to
-    with self.assertRaises(Exception):
-      self.inst = s_load_b128(s[4:6], s[0:1], NULL, 0)
-
-  def test_mov_b32(self):
-    self.inst = s_mov_b32(s[80], s[0])
-
-  def test_mov_b64(self):
-    self.inst = s_mov_b64(s[80:81], s[0:1])
-
-  def test_mov_b32_wrong(self):
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b32(s[80:81], s[0:1])
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b32(s[80:81], s[0])
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b32(s[80], s[0:1])
-
-  def test_mov_b64_wrong(self):
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b64(s[80], s[0])
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b64(s[80], s[0:1])
-    with self.assertRaises(Exception):
-      self.inst = s_mov_b64(s[80:81], s[0])
-
-  def test_load_b128_no_0(self):
-    self.inst = s_load_b128(s[4:7], s[0:1], NULL)
-
-  def test_load_b128_s(self):
-    self.inst = s_load_b128(s[4:7], s[0:1], s[8], 0)
-
-  def test_load_b128_v(self):
-    with self.assertRaises(TypeError):
-      self.inst = s_load_b128(s[4:7], s[0:1], v[8], 0)
-
-  def test_load_b128_off(self):
-    self.inst = s_load_b128(s[4:7], s[0:1], NULL, 3)
-
-  def test_simple_stos(self):
-    self.inst = s_mov_b32(s[0], s[1])
-
-  def test_simple_wrong(self):
-    with self.assertRaises(TypeError):
-      self.inst = s_mov_b32(v[0], s[1])
-
-  def test_simple_vtov(self):
-    self.inst = v_mov_b32_e32(v[0], v[1])
-
-  def test_simple_stov(self):
-    self.inst = v_mov_b32_e32(v[0], s[2])
-
-  def test_simple_float_to_v(self):
-    self.inst = v_mov_b32_e32(v[0], 1.0)
-
-  def test_simple_v_to_float(self):
-    with self.assertRaises(TypeError):
-      self.inst = v_mov_b32_e32(1, v[0])
-
-  def test_simple_int_to_v(self):
-    self.inst = v_mov_b32_e32(v[0], 1)
-
-  def test_three_add(self):
-    self.inst = v_add_co_ci_u32_e32(v[3], s[7], v[3])
-
-  def test_three_add_v(self):
-    self.inst = v_add_co_ci_u32_e32(v[3], v[7], v[3])
-
-  def test_three_add_const(self):
-    self.inst = v_add_co_ci_u32_e32(v[3], 2.0, v[3])
-
-  def test_swaitcnt_lgkm(self): self.inst = s_waitcnt(0xfc07)
-  def test_swaitcnt_vm(self): self.inst = s_waitcnt(0x03f7)
-
-  def test_vmad(self):
-    self.inst = v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2])
-
-  def test_large_imm(self):
-    self.inst = v_mov_b32_e32(v[0], 0x1234)
-
-  def test_dual_mov(self):
-    self.inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], srcy0=v[4])
-
-  def test_dual_mul(self):
-    self.inst = v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
-
-  def test_simple_int_to_s(self):
-    self.inst = s_mov_b32(s[0], 3)
-
-  def test_complex_int_to_s(self):
-    self.inst = s_mov_b32(s[0], 0x235646)
-
-  def test_simple_float_to_s(self):
-    self.inst = s_mov_b32(s[0], 1.0)
-
-  def test_complex_float_to_s(self):
-    self.inst = s_mov_b32(s[0], 1337.0)
-    int_inst = s_mov_b32(s[0], struct.unpack("I", struct.pack("f", 1337.0))[0])
-    self.assertEqual(self.inst, int_inst)
-
-class TestRegisterSliceSyntax(unittest.TestCase):
-  """
-  Issue: Register slice syntax should use AMD assembly convention (inclusive end).
-
-  In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
-  The DSL should match this convention so that:
-  - s[4:7] gives 4 registers
-  - Disassembler output can be copied directly back into DSL code
-
-  Fix: Change _RegFactory.__getitem__ to use inclusive end:
-    key.stop - key.start + 1  (instead of key.stop - key.start)
-  """
-  def test_register_slice_count(self):
-    # s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
-    reg = s[4:7]
-    self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
-
-  def test_register_slice_roundtrip(self):
-    # Round-trip: DSL -> disasm -> DSL should preserve register count
-    reg = s[4:7]  # 4 registers in AMD convention
-    inst = s_load_b128(reg, s[0:1], NULL, 0)
-    disasm = inst.disasm()
-    # Disasm shows s[4:7] - user should be able to copy this back
-    self.assertIn("s[4:7]", disasm)
-    # And s[4:7] in DSL should give the same 4 registers
-    reg_from_disasm = s[4:7]
-    self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
-
-class TestInstructionEquality(unittest.TestCase):
-  """
-  Issue: No __eq__ method - instruction comparison requires repr() workaround.
-
-  Two identical instructions should compare equal with ==, but currently:
-    inst1 == inst2 returns False
-
-  The test_handwritten.py works around this with:
-    self.assertEqual(repr(self.inst), repr(reasm))
-  """
-  def test_identical_instructions_equal(self):
-    inst1 = v_mov_b32_e32(v[0], v[1])
-    inst2 = v_mov_b32_e32(v[0], v[1])
-    self.assertEqual(inst1, inst2, "identical instructions should be equal")
-
-  def test_different_instructions_not_equal(self):
-    inst1 = v_mov_b32_e32(v[0], v[1])
-    inst2 = v_mov_b32_e32(v[0], v[2])
-    self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
-
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,330 +0,0 @@
-#!/usr/bin/env python3
-"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
-import unittest, re, io, sys, subprocess
-from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.asm import waitcnt, asm
-from extra.assembly.amd.test.helpers import get_llvm_mc
-
-def disassemble(lib: bytes, arch: str = "gfx1100") -> str:
-  """Disassemble ELF binary using tinygrad's compiler, return raw output."""
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-  old_stdout = sys.stdout
-  sys.stdout = io.StringIO()
-  HIPCompiler(arch).disassemble(lib)
-  output = sys.stdout.getvalue()
-  sys.stdout = old_stdout
-  return output
-
-def parse_disassembly(raw: str) -> list[str]:
-  """Parse disassembly output to list of instruction mnemonics."""
-  lines = []
-  for line in raw.splitlines():
-    if line.startswith('\t'):
-      instr = line.split('//')[0].strip()
-      if instr: lines.append(instr)
-  return lines
-
-def assemble_and_disassemble(instructions: list, arch: str = "gfx1100") -> list[str]:
-  """Assemble instructions with our DSL, then disassemble with AMD toolchain."""
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-  # Generate bytes from our DSL
-  code_bytes = b''.join(inst.to_bytes() for inst in instructions)
-
-  # Wrap in minimal ELF-compatible assembly with .byte directives
-  byte_str = ', '.join(f'0x{b:02x}' for b in code_bytes)
-  asm_src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n.byte {byte_str}\n"
-
-  # Assemble with AMD COMGR and disassemble
-  lib = HIPCompiler(arch).compile(asm_src)
-  return parse_disassembly(disassemble(lib, arch))
-
-class TestIntegration(unittest.TestCase):
-  """Test our assembler output matches LLVM disassembly."""
-
-  def test_simple_sop1(self):
-    """Test SOP1 instructions round-trip."""
-    instructions = [
-      s_mov_b32(s[0], s[1]),
-      s_mov_b32(s[2], 0),
-      s_not_b32(s[3], s[4]),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_mov_b32', disasm[0])
-    self.assertIn('s_mov_b32', disasm[1])
-    self.assertIn('s_not_b32', disasm[2])
-
-  def test_simple_sop2(self):
-    """Test SOP2 instructions round-trip."""
-    instructions = [
-      s_add_u32(s[0], s[1], s[2]),
-      s_sub_u32(s[3], s[4], 10),
-      s_and_b32(s[5], s[6], s[7]),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_add_u32', disasm[0])
-    self.assertIn('s_sub_u32', disasm[1])
-    self.assertIn('s_and_b32', disasm[2])
-
-  def test_simple_vop2(self):
-    """Test VOP2 instructions round-trip."""
-    instructions = [
-      v_add_f32_e32(v[0], v[1], v[2]),
-      v_mul_f32_e32(v[3], 1.0, v[4]),  # 1.0 is inline constant
-      v_and_b32_e32(v[5], 10, v[6]),  # small inline constant
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('v_add_f32', disasm[0])
-    self.assertIn('v_mul_f32', disasm[1])
-
-  def test_control_flow(self):
-    """Test control flow instructions."""
-    instructions = [
-      s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
-      s_endpgm(),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_waitcnt', disasm[0])
-    self.assertIn('s_endpgm', disasm[1])
-
-  def test_memory_ops(self):
-    """Test memory instructions."""
-    instructions = [
-      s_load_b32(s[0], s[0:2], NULL),
-      s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
-      global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
-      s_endpgm(),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_load_b32', disasm[0])
-    self.assertIn('s_waitcnt', disasm[1])
-    self.assertIn('global_store_b32', disasm[2])
-
-  def test_full_kernel(self):
-    """Test a complete kernel similar to tinygrad output."""
-    # Simple kernel: load value, add 1, store back
-    instructions = [
-      # Get thread ID
-      v_mov_b32_e32(v[0], s[0]),  # base addr low
-      v_mov_b32_e32(v[1], s[1]),  # base addr high
-      # Load value
-      global_load_b32(vdst=v[2], addr=v[0:2], saddr=OFF),
-      s_waitcnt(simm16=waitcnt(vmcnt=0)),
-      # Add 1.0
-      v_add_f32_e32(v[2], 1.0, v[2]),
-      # Store result
-      global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
-      s_endpgm(),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    # Verify key instructions are present
-    self.assertTrue(any('global_load' in d for d in disasm))
-    self.assertTrue(any('v_add_f32' in d for d in disasm))
-    self.assertTrue(any('global_store' in d for d in disasm))
-    self.assertTrue(any('s_endpgm' in d for d in disasm))
-
-  def test_bytes_roundtrip(self):
-    """Test that our bytes match what AMD assembler produces."""
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-    # Simple instruction
-    inst = s_mov_b32(s[0], s[1])
-    our_bytes = inst.to_bytes()
-
-    # Assemble same instruction with AMD toolchain
-    asm_src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\ns_mov_b32 s0, s1\n"
-    compiler = HIPCompiler("gfx1100")
-    lib = compiler.compile(asm_src)
-    raw = disassemble(lib)
-
-    for line in raw.splitlines():
-      if 's_mov_b32' in line and '//' in line:
-        # Extract hex bytes from comment: "// 000000001300: BE800001"
-        comment = line.split('//')[1].strip()
-        hex_str = comment.split(':')[1].strip()
-        # Convert big-endian hex string to little-endian bytes
-        amd_bytes = bytes.fromhex(hex_str)[::-1]  # reverse for little-endian
-        self.assertEqual(our_bytes, amd_bytes, f"Bytes mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
-        return
-    self.fail("Could not find s_mov_b32 in disassembly")
-
-class TestAsm(unittest.TestCase):
-  """Test asm() string parsing."""
-
-  def test_asm_basic(self):
-    """Test basic instruction parsing."""
-    inst = asm('s_mov_b32 s0, s1')
-    self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], s[1]).to_bytes())
-
-  def test_asm_with_immediates(self):
-    """Test parsing with immediate values."""
-    inst = asm('s_add_u32 s0, s1, 10')
-    self.assertEqual(inst.to_bytes(), s_add_u32(s[0], s[1], 10).to_bytes())
-
-  def test_asm_float_const(self):
-    """Test parsing float constants."""
-    inst = asm('v_mul_f32_e32 v0, 1.0, v1')
-    self.assertEqual(inst.to_bytes(), v_mul_f32_e32(v[0], 1.0, v[1]).to_bytes())
-
-  def test_asm_hex_immediate(self):
-    """Test parsing hex immediates."""
-    inst = asm('s_waitcnt 0xfc07')
-    self.assertEqual(inst.to_bytes(), s_waitcnt(simm16=0xfc07).to_bytes())
-
-  def test_asm_special_regs(self):
-    """Test parsing special registers."""
-    inst = asm('s_mov_b32 s0, vcc_lo')
-    self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], VCC_LO).to_bytes())
-
-  def test_asm_register_range(self):
-    """Test parsing register ranges."""
-    inst = asm('s_load_b128 s[4:7], s[0:1], null')
-    self.assertEqual(inst.to_bytes(), s_load_b128(s[4:7], s[0:1], NULL).to_bytes())
-
-  def test_asm_matches_llvm(self):
-    """Test asm() output matches LLVM assembler."""
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    compiler = HIPCompiler('gfx1100')
-
-    def get_llvm_bytes(instr: str) -> bytes:
-      src = f'.text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n{instr}\n'
-      lib = compiler.compile(src)
-      raw = disassemble(lib)
-      for line in raw.splitlines():
-        if instr.split()[0] in line and '//' in line:
-          hex_str = line.split('//')[1].strip().split(':')[1].strip()
-          return bytes.fromhex(hex_str)[::-1]
-      return b''
-
-    tests = ['s_mov_b32 s0, s1', 's_endpgm', 'v_add_f32_e32 v0, v1, v2']
-    for t in tests:
-      self.assertEqual(asm(t).to_bytes(), get_llvm_bytes(t), f"mismatch for: {t}")
-
-  def test_asm_vop3_modifiers(self):
-    """Test asm() with VOP3 modifiers (neg, abs, clamp)."""
-    def get_llvm_encoding(instr: str) -> str:
-      result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', '-mcpu=gfx1100', '-show-encoding'],
-                              input=instr, capture_output=True, text=True)
-      if m := re.search(r'encoding:\s*\[(.*?)\]', result.stdout):
-        return m.group(1).replace('0x','').replace(',','').replace(' ','')
-      return ''
-
-    tests = [
-      'v_fma_f32 v0, -v1, v2, v3',       # neg on src0
-      'v_fma_f32 v0, v1, |v2|, v3',      # abs on src1
-      'v_fma_f32 v0, v1, v2, v3 clamp',  # clamp
-      'v_fma_f32 v0, -v1, |v2|, v3 clamp',  # all modifiers
-      'v_fma_f32 v0, -|v1|, v2, v3',     # neg+abs on same operand
-    ]
-    for t in tests:
-      our_hex = asm(t).to_bytes().hex()
-      llvm_hex = get_llvm_encoding(t)
-      self.assertEqual(our_hex, llvm_hex, f"mismatch for: {t}")
-
-class TestTinygradIntegration(unittest.TestCase):
-  """Test that we can parse disassembled tinygrad kernels."""
-
-  def test_simple_add_kernel(self):
-    """Generate a simple add kernel from tinygrad and verify disassembly."""
-    from tinygrad import Tensor
-    from tinygrad.codegen import get_program
-    from tinygrad.renderer.cstyle import AMDHIPRenderer
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    from tinygrad.uop.ops import Ops
-
-    # Create a computation that generates a real kernel
-    a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
-    b = Tensor([5.0, 6.0, 7.0, 8.0]).realize()
-    c = a + b
-
-    # Get schedule and find SINK
-    schedule = c.schedule()
-    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
-    self.assertTrue(len(sink_items) > 0, "No SINK in schedule")
-
-    # Generate program
-    renderer = AMDHIPRenderer('gfx1100')
-    prg = get_program(sink_items[0].ast, renderer)
-    self.assertIsNotNone(prg.src)
-
-    # Compile and disassemble
-    compiler = HIPCompiler('gfx1100')
-    lib = compiler.compile(prg.src)
-    raw_disasm = disassemble(lib)
-    instrs = parse_disassembly(raw_disasm)
-
-    # Verify we got some instructions
-    self.assertTrue(len(instrs) > 0, "No instructions in disassembly")
-    # Should have an endpgm
-    self.assertTrue(any('s_endpgm' in i for i in instrs), "Missing s_endpgm")
-
-  def test_matmul_kernel(self):
-    """Generate a matmul kernel and verify disassembly has expected patterns."""
-    from tinygrad import Tensor
-    from tinygrad.codegen import get_program
-    from tinygrad.renderer.cstyle import AMDHIPRenderer
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    from tinygrad.uop.ops import Ops
-
-    # Create a small matmul
-    a = Tensor.rand(4, 4).realize()
-    b = Tensor.rand(4, 4).realize()
-    c = a @ b
-
-    # Get schedule
-    schedule = c.schedule()
-    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
-    self.assertTrue(len(sink_items) > 0)
-
-    # Generate and compile
-    renderer = AMDHIPRenderer('gfx1100')
-    prg = get_program(sink_items[0].ast, renderer)
-    compiler = HIPCompiler('gfx1100')
-    lib = compiler.compile(prg.src)
-    raw_disasm = disassemble(lib)
-    instrs = parse_disassembly(raw_disasm)
-
-    # Matmul should have multiply and add instructions
-    has_mul = any('mul' in i.lower() for i in instrs)
-    has_add = any('add' in i.lower() for i in instrs)
-    self.assertTrue(has_mul or has_add, "Matmul should have mul/add ops")
-
-  def test_disasm_to_bytes_roundtrip(self):
-    """Parse disassembled instructions and verify we can re-encode some of them."""
-    from tinygrad import Tensor
-    from tinygrad.codegen import get_program
-    from tinygrad.renderer.cstyle import AMDHIPRenderer
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    from tinygrad.uop.ops import Ops
-
-    # Simple kernel
-    a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
-    b = (a * 2.0)
-
-    schedule = b.schedule()
-    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
-    if not sink_items: return  # skip if no kernel
-
-    renderer = AMDHIPRenderer('gfx1100')
-    prg = get_program(sink_items[0].ast, renderer)
-    compiler = HIPCompiler('gfx1100')
-    lib = compiler.compile(prg.src)
-    raw_disasm = disassemble(lib)
-
-    # Find s_endpgm and verify we can encode it
-    for line in raw_disasm.splitlines():
-      if 's_endpgm' in line and '//' in line:
-        # Extract bytes from comment
-        comment = line.split('//')[1].strip()
-        hex_str = comment.split(':')[1].strip()
-        amd_bytes = bytes.fromhex(hex_str)[::-1]
-
-        # Our encoding
-        our_inst = s_endpgm()
-        our_bytes = our_inst.to_bytes()
-
-        self.assertEqual(our_bytes, amd_bytes, f"s_endpgm mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
-        return
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,201 +0,0 @@
-#!/usr/bin/env python3
-"""Test RDNA3 assembler/disassembler against LLVM test vectors."""
-import unittest, re, subprocess
-from tinygrad.helpers import fetch
-from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.asm import asm
-from extra.assembly.amd.test.helpers import get_llvm_mc
-
-LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/main/llvm/test/MC/AMDGPU"
-
-# Format info: (filename, format_class, op_enum)
-LLVM_TEST_FILES = {
-  # Scalar ALU
-  'sop1': ('gfx11_asm_sop1.s', SOP1, SOP1Op),
-  'sop2': ('gfx11_asm_sop2.s', SOP2, SOP2Op),
-  'sopp': ('gfx11_asm_sopp.s', SOPP, SOPPOp),
-  'sopk': ('gfx11_asm_sopk.s', SOPK, SOPKOp),
-  'sopc': ('gfx11_asm_sopc.s', SOPC, SOPCOp),
-  # Vector ALU
-  'vop1': ('gfx11_asm_vop1.s', VOP1, VOP1Op),
-  'vop2': ('gfx11_asm_vop2.s', VOP2, VOP2Op),
-  'vopc': ('gfx11_asm_vopc.s', VOPC, VOPCOp),
-  'vop3': ('gfx11_asm_vop3.s', VOP3, VOP3Op),
-  'vop3p': ('gfx11_asm_vop3p.s', VOP3P, VOP3POp),
-  'vop3sd': ('gfx11_asm_vop3.s', VOP3SD, VOP3SDOp),  # VOP3SD shares file with VOP3
-  'vinterp': ('gfx11_asm_vinterp.s', VINTERP, VINTERPOp),
-  'vopd': ('gfx11_asm_vopd.s', VOPD, VOPDOp),
-  'vopcx': ('gfx11_asm_vopcx.s', VOPC, VOPCOp),  # VOPCX uses VOPC format
-  # VOP3 promotions (VOP1/VOP2/VOPC promoted to VOP3 encoding)
-  'vop3_from_vop1': ('gfx11_asm_vop3_from_vop1.s', VOP3, VOP3Op),
-  'vop3_from_vop2': ('gfx11_asm_vop3_from_vop2.s', VOP3, VOP3Op),
-  'vop3_from_vopc': ('gfx11_asm_vop3_from_vopc.s', VOP3, VOP3Op),
-  'vop3_from_vopcx': ('gfx11_asm_vop3_from_vopcx.s', VOP3, VOP3Op),
-  # Memory
-  'ds': ('gfx11_asm_ds.s', DS, DSOp),
-  'smem': ('gfx11_asm_smem.s', SMEM, SMEMOp),
-  'flat': ('gfx11_asm_flat.s', FLAT, FLATOp),
-  'mubuf': ('gfx11_asm_mubuf.s', MUBUF, MUBUFOp),
-  'mtbuf': ('gfx11_asm_mtbuf.s', MTBUF, MTBUFOp),
-  'mimg': ('gfx11_asm_mimg.s', MIMG, MIMGOp),
-  # WMMA (matrix multiply)
-  'wmma': ('gfx11_asm_wmma.s', VOP3P, VOP3POp),
-  # Additional features
-  'vop3_features': ('gfx11_asm_vop3_features.s', VOP3, VOP3Op),
-  'vop3p_features': ('gfx11_asm_vop3p_features.s', VOP3P, VOP3POp),
-  'vopd_features': ('gfx11_asm_vopd_features.s', VOPD, VOPDOp),
-  # Alias files (alternative mnemonics)
-  'vop3_alias': ('gfx11_asm_vop3_alias.s', VOP3, VOP3Op),
-  'vop3p_alias': ('gfx11_asm_vop3p_alias.s', VOP3P, VOP3POp),
-  'vopc_alias': ('gfx11_asm_vopc_alias.s', VOPC, VOPCOp),
-  'vopcx_alias': ('gfx11_asm_vopcx_alias.s', VOPC, VOPCOp),
-  'vinterp_alias': ('gfx11_asm_vinterp_alias.s', VINTERP, VINTERPOp),
-  'smem_alias': ('gfx11_asm_smem_alias.s', SMEM, SMEMOp),
-  'mubuf_alias': ('gfx11_asm_mubuf_alias.s', MUBUF, MUBUFOp),
-  'mtbuf_alias': ('gfx11_asm_mtbuf_alias.s', MTBUF, MTBUFOp),
-}
-
-def parse_llvm_tests(text: str) -> list[tuple[str, bytes]]:
-  """Parse LLVM test format into (asm, expected_bytes) pairs."""
-  tests, lines = [], text.split('\n')
-  for i, line in enumerate(lines):
-    line = line.strip()
-    if not line or line.startswith(('//', '.', ';')): continue
-    asm_text = line.split('//')[0].strip()
-    if not asm_text: continue
-    for j in range(i, min(i + 3, len(lines))):
-      # Match GFX11, W32, or W64 encodings (all valid for gfx11)
-      # Format 1: "// GFX11: v_foo ... ; encoding: [0x01,0x02,...]"
-      # Format 2: "// GFX11: [0x01,0x02,...]" (used by DS, older files)
-      if m := re.search(r'(?:GFX11|W32|W64)[^:]*:.*?encoding:\s*\[(.*?)\]', lines[j]):
-        hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
-      elif m := re.search(r'(?:GFX11|W32|W64)[^:]*:\s*\[(0x[0-9a-fA-F,x\s]+)\]', lines[j]):
-        hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
-      else:
-        continue
-      if hex_bytes:
-        try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
-        except ValueError: pass
-      break
-  return tests
-
-def try_assemble(text: str):
-  """Try to assemble instruction text, return bytes or None on failure."""
-  try: return asm(text).to_bytes()
-  except: return None
-
-def compile_asm_batch(instrs: list[str]) -> list[bytes]:
-  """Compile multiple instructions with a single llvm-mc call."""
-  if not instrs: return []
-  asm_text = ".text\n" + "\n".join(instrs) + "\n"
-  result = subprocess.run(
-    [get_llvm_mc(), '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
-    input=asm_text, capture_output=True, text=True, timeout=30)
-  if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
-  # Parse all encodings from output
-  results = []
-  for line in result.stdout.split('\n'):
-    if 'encoding:' not in line: continue
-    enc = line.split('encoding:')[1].strip()
-    if enc.startswith('[') and enc.endswith(']'):
-      results.append(bytes.fromhex(enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')))
-  if len(results) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(results)}")
-  return results
-
-class TestLLVM(unittest.TestCase):
-  """Test assembler and disassembler against all LLVM test vectors."""
-  tests: dict[str, list[tuple[str, bytes]]] = {}
-
-  @classmethod
-  def setUpClass(cls):
-    for name, (filename, _, _) in LLVM_TEST_FILES.items():
-      try:
-        data = fetch(f"{LLVM_BASE}/{filename}").read_bytes()
-        cls.tests[name] = parse_llvm_tests(data.decode('utf-8', errors='ignore'))
-      except Exception as e:
-        print(f"Warning: couldn't fetch {filename}: {e}")
-        cls.tests[name] = []
-
-# Generate test methods dynamically for each format
-def _make_asm_test(name):
-  def test(self):
-    passed, failed, skipped = 0, 0, 0
-    for asm_text, expected in self.tests.get(name, []):
-      result = try_assemble(asm_text)
-      if result is None: skipped += 1
-      elif result == expected: passed += 1
-      else: failed += 1
-    print(f"{name.upper()} asm: {passed} passed, {failed} failed, {skipped} skipped")
-    self.assertEqual(failed, 0)
-  return test
-
-def _make_disasm_test(name):
-  def test(self):
-    _, fmt_cls, op_enum = LLVM_TEST_FILES[name]
-    # VOP3SD opcodes that share encoding with VOP3 (only for vop3sd test, not vopc promotions)
-    vop3sd_opcodes = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
-    is_vopc_promotion = name in ('vop3_from_vopc', 'vop3_from_vopcx')
-    undocumented = {'smem': {34, 35}, 'sopk': {22, 23}, 'sopp': {8, 58, 59}}
-
-    # First pass: decode all instructions and collect disasm strings
-    to_test: list[tuple[str, bytes, str | None, str | None]] = []  # (asm_text, data, disasm_str, error)
-    skipped = 0
-    for asm_text, data in self.tests.get(name, []):
-      if len(data) > fmt_cls._size(): continue
-      temp_inst = fmt_cls.from_bytes(data)
-      temp_op = temp_inst._values.get('op', 0)
-      temp_op = temp_op.val if hasattr(temp_op, 'val') else temp_op
-      if temp_op in undocumented.get(name, set()): skipped += 1; continue
-      if name == 'sopp':
-        simm16 = temp_inst._values.get('simm16', 0)
-        simm16 = simm16.val if hasattr(simm16, 'val') else simm16
-        sopp_no_imm = {48, 54, 53, 55, 60, 61, 62}
-        if temp_op in sopp_no_imm and simm16 != 0: skipped += 1; continue
-      try:
-        if fmt_cls.__name__ in ('VOP3', 'VOP3SD'):
-          temp = VOP3.from_bytes(data)
-          op_val = temp._values.get('op', 0)
-          op_val = op_val.val if hasattr(op_val, 'val') else op_val
-          is_vop3sd = (op_val in vop3sd_opcodes) and not is_vopc_promotion
-          decoded = VOP3SD.from_bytes(data) if is_vop3sd else VOP3.from_bytes(data)
-          if is_vop3sd: VOP3SDOp(op_val)
-          else: VOP3Op(op_val)
-        else:
-          decoded = fmt_cls.from_bytes(data)
-          op_val = decoded._values.get('op', 0)
-          op_val = op_val.val if hasattr(op_val, 'val') else op_val
-          op_enum(op_val)
-        if decoded.to_bytes()[:len(data)] != data:
-          to_test.append((asm_text, data, None, "decode roundtrip failed"))
-          continue
-        to_test.append((asm_text, data, decoded.disasm(), None))
-      except Exception as e:
-        to_test.append((asm_text, data, None, f"exception: {e}"))
-
-    # Batch compile all disasm strings with single llvm-mc call
-    disasm_strs = [(i, t[2]) for i, t in enumerate(to_test) if t[2] is not None]
-    llvm_results = compile_asm_batch([s for _, s in disasm_strs]) if disasm_strs else []
-    llvm_map = {i: llvm_results[j] for j, (i, _) in enumerate(disasm_strs)}
-
-    # Match results back
-    passed, failed = 0, 0
-    failures: list[str] = []
-    for idx, (asm_text, data, disasm_str, error) in enumerate(to_test):
-      if error:
-        failed += 1; failures.append(f"{error} for {data.hex()}")
-      elif disasm_str is not None and idx in llvm_map:
-        llvm_bytes = llvm_map[idx]
-        if llvm_bytes is not None and llvm_bytes == data: passed += 1
-        elif llvm_bytes is not None: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
-
-    print(f"{name.upper()} disasm: {passed} passed, {failed} failed" + (f", {skipped} skipped" if skipped else ""))
-    if failures[:10]: print("  " + "\n  ".join(failures[:10]))
-    self.assertEqual(failed, 0)
-  return test
-
-for name in LLVM_TEST_FILES:
-  setattr(TestLLVM, f'test_{name}_asm', _make_asm_test(name))
-  setattr(TestLLVM, f'test_{name}_disasm', _make_disasm_test(name))
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,55 +0,0 @@
-#!/usr/bin/env python3
-"""Test that invalid instructions raise exceptions through the mock GPU stack."""
-import unittest, subprocess, os, time
-
-class TestMockGPUInvalidInstruction(unittest.TestCase):
-  def test_unsupported_instruction_raises(self):
-    """Test that unsupported instructions raise immediately through the full MOCKGPU stack."""
-    test_code = '''
-import struct
-from tinygrad import Device, Tensor
-from tinygrad.engine.realize import get_runner
-from tinygrad.runtime.ops_amd import AMDProgram
-
-dev = Device["AMD"]
-a = Tensor([1.0]).realize()
-b = a + 1
-si = b.schedule()[-1]
-runner = get_runner(dev.device, si.ast)
-
-prg = runner._prg
-lib = bytearray(prg.lib)
-
-# Find s_endpgm (0xBFB00000) and replace with invalid SOPP op=127 (0xBFFF0000)
-found = False
-for i in range(0, len(lib) - 4, 4):
-  if struct.unpack("<I", lib[i:i+4])[0] == 0xBFB00000:
-    lib[i:i+4] = struct.pack("<I", 0xBFFF0000)
-    found = True
-    break
-assert found, "s_endpgm not found"
-
-patched_prg = AMDProgram(dev, "patched", bytes(lib))
-b.uop.buffer.allocate()
-patched_prg(b.uop.buffer._buf, a.uop.buffer._buf, global_size=(1,1,1), local_size=(1,1,1))
-dev.synchronize()
-'''
-
-    env = os.environ.copy()
-    env["AMD"] = "1"
-    env["MOCKGPU"] = "1"
-    env["PYTHON_REMU"] = "1"
-    env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
-
-    st = time.perf_counter()
-    result = subprocess.run(["python", "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
-    elapsed = time.perf_counter() - st
-
-    self.assertNotEqual(result.returncode, 0, "should have raised")
-    self.assertTrue("NotImplementedError" in result.stderr or "ValueError" in result.stderr,
-                    f"expected NotImplementedError or ValueError in stderr")
-    # Should exit immediately, not wait for the full timeout
-    self.assertLess(elapsed, 9.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,404 +0,0 @@
-#!/usr/bin/env python3
-"""Tests for the RDNA3 pseudocode DSL."""
-import unittest
-from extra.assembly.amd.pcode import (Reg, TypedView, SliceProxy, MASK32, MASK64,
-                                       _f32, _i32, _f16, _i16, f32_to_f16, _isnan, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
-                                       BYTE_PERMUTE, v_sad_u8, v_msad_u8)
-from extra.assembly.amd.pdf import compile_pseudocode, _expr
-from extra.assembly.amd.test.helpers import ExecContext
-from extra.assembly.amd.autogen.rdna3.gen_pcode import _VOP3SDOp_V_DIV_SCALE_F32, _VOPCOp_V_CMP_CLASS_F32
-
-class TestReg(unittest.TestCase):
-  def test_u32_read(self):
-    r = Reg(0xDEADBEEF)
-    self.assertEqual(int(r.u32), 0xDEADBEEF)
-
-  def test_u32_write(self):
-    r = Reg(0)
-    r.u32 = 0x12345678
-    self.assertEqual(r._val, 0x12345678)
-
-  def test_f32_read(self):
-    r = Reg(0x40400000)  # 3.0f
-    self.assertAlmostEqual(float(r.f32), 3.0)
-
-  def test_f32_write(self):
-    r = Reg(0)
-    r.f32 = 3.0
-    self.assertEqual(r._val, 0x40400000)
-
-  def test_i32_signed(self):
-    r = Reg(0xFFFFFFFF)  # -1 as signed
-    self.assertEqual(int(r.i32), -1)
-
-  def test_u64(self):
-    r = Reg(0xDEADBEEFCAFEBABE)
-    self.assertEqual(int(r.u64), 0xDEADBEEFCAFEBABE)
-
-  def test_f64(self):
-    r = Reg(0x4008000000000000)  # 3.0 as f64
-    self.assertAlmostEqual(float(r.f64), 3.0)
-
-class TestTypedView(unittest.TestCase):
-  def test_bit_slice(self):
-    r = Reg(0xDEADBEEF)
-    # Slices return SliceProxy which supports .u32, .u16 etc (matching pseudocode like S1.u32[1:0].u32)
-    self.assertEqual(r.u32[7:0].u32, 0xEF)
-    self.assertEqual(r.u32[15:8].u32, 0xBE)
-    self.assertEqual(r.u32[23:16].u32, 0xAD)
-    self.assertEqual(r.u32[31:24].u32, 0xDE)
-    # Also works with int() for arithmetic
-    self.assertEqual(int(r.u32[7:0]), 0xEF)
-
-  def test_single_bit_read(self):
-    r = Reg(0b11010101)
-    self.assertEqual(r.u32[0], 1)
-    self.assertEqual(r.u32[1], 0)
-    self.assertEqual(r.u32[2], 1)
-    self.assertEqual(r.u32[3], 0)
-
-  def test_single_bit_write(self):
-    r = Reg(0)
-    r.u32[5] = 1
-    r.u32[3] = 1
-    self.assertEqual(r._val, 0b00101000)
-
-  def test_nested_bit_access(self):
-    # S0.u32[S1.u32[4:0]] - access bit at position from another register
-    s0 = Reg(0b11010101)
-    s1 = Reg(3)
-    bit_pos = s1.u32[4:0]  # SliceProxy, int value = 3
-    bit_val = s0.u32[int(bit_pos)]  # bit 3 of s0 = 0
-    self.assertEqual(int(bit_pos), 3)
-    self.assertEqual(bit_val, 0)
-
-  def test_arithmetic(self):
-    r1 = Reg(0x40400000)  # 3.0f
-    r2 = Reg(0x40800000)  # 4.0f
-    result = r1.f32 + r2.f32
-    self.assertAlmostEqual(result, 7.0)
-
-  def test_comparison(self):
-    r1 = Reg(5)
-    r2 = Reg(3)
-    self.assertTrue(r1.u32 > r2.u32)
-    self.assertFalse(r1.u32 < r2.u32)
-    self.assertTrue(r1.u32 != r2.u32)
-
-class TestSliceProxy(unittest.TestCase):
-  def test_slice_read(self):
-    r = Reg(0x56781234)
-    self.assertEqual(r[15:0].u16, 0x1234)
-    self.assertEqual(r[31:16].u16, 0x5678)
-
-  def test_slice_write(self):
-    r = Reg(0)
-    r[15:0].u16 = 0x1234
-    r[31:16].u16 = 0x5678
-    self.assertEqual(r._val, 0x56781234)
-
-  def test_slice_f16(self):
-    r = Reg(0)
-    r[15:0].f16 = 3.0
-    self.assertAlmostEqual(_f16(r._val & 0xffff), 3.0, places=2)
-
-class TestCompiler(unittest.TestCase):
-  def test_ternary(self):
-    result = _expr("a > b ? 1 : 0")
-    self.assertIn("if", result)
-    self.assertIn("else", result)
-
-  def test_type_prefix_strip(self):
-    self.assertEqual(_expr("1'0U"), "0")
-    self.assertEqual(_expr("32'1"), "1")
-    self.assertEqual(_expr("16'0xFFFF"), "0xFFFF")
-
-  def test_suffix_strip(self):
-    self.assertEqual(_expr("0ULL"), "0")
-    self.assertEqual(_expr("1LL"), "1")
-    self.assertEqual(_expr("5U"), "5")
-    self.assertEqual(_expr("3.14F"), "3.14")
-
-  def test_boolean_ops(self):
-    self.assertIn("and", _expr("a && b"))
-    self.assertIn("or", _expr("a || b"))
-    self.assertIn("!=", _expr("a <> b"))
-
-  def test_pack16(self):
-    result = _expr("{ a, b }")
-    self.assertIn("_pack", result)
-
-  def test_type_cast_strip(self):
-    self.assertEqual(_expr("64'U(x)"), "(x)")
-    self.assertEqual(_expr("32'I(y)"), "(y)")
-
-class TestExecContext(unittest.TestCase):
-  def test_float_add(self):
-    ctx = ExecContext(s0=0x40400000, s1=0x40800000)  # 3.0f, 4.0f
-    ctx.D0.f32 = ctx.S0.f32 + ctx.S1.f32
-    self.assertAlmostEqual(_f32(ctx.D0._val), 7.0)
-
-  def test_float_mul(self):
-    ctx = ExecContext(s0=0x40400000, s1=0x40800000)  # 3.0f, 4.0f
-    ctx.run("D0.f32 = S0.f32 * S1.f32")
-    self.assertAlmostEqual(_f32(ctx.D0._val), 12.0)
-
-  def test_scc_comparison(self):
-    ctx = ExecContext(s0=42, s1=42)
-    ctx.run("SCC = S0.u32 == S1.u32")
-    self.assertEqual(ctx.SCC._val, 1)
-
-  def test_scc_comparison_false(self):
-    ctx = ExecContext(s0=42, s1=43)
-    ctx.run("SCC = S0.u32 == S1.u32")
-    self.assertEqual(ctx.SCC._val, 0)
-
-  def test_ternary(self):
-    code = compile_pseudocode("D0.u32 = S0.u32 > S1.u32 ? 1'1U : 1'0U")
-    ctx = ExecContext(s0=5, s1=3)
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val, 1)
-
-  def test_pack(self):
-    code = compile_pseudocode("D0 = { S1[15:0].u16, S0[15:0].u16 }")
-    ctx = ExecContext(s0=0x1234, s1=0x5678)
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val, 0x56781234)
-
-  def test_tmp_with_typed_access(self):
-    code = compile_pseudocode("""tmp = S0.u32 + S1.u32
-D0.u32 = tmp.u32""")
-    ctx = ExecContext(s0=100, s1=200)
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val, 300)
-
-  def test_s_add_u32_pattern(self):
-    # Real pseudocode pattern from S_ADD_U32
-    code = compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
-SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
-D0.u32 = tmp.u32""")
-    # Test overflow case
-    ctx = ExecContext(s0=0xFFFFFFFF, s1=0x00000001)
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val, 0)  # Wraps to 0
-    self.assertEqual(ctx.SCC._val, 1)  # Carry set
-
-  def test_s_add_u32_no_overflow(self):
-    code = compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
-SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
-D0.u32 = tmp.u32""")
-    ctx = ExecContext(s0=100, s1=200)
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val, 300)
-    self.assertEqual(ctx.SCC._val, 0)  # No carry
-
-  def test_vcc_lane_read(self):
-    ctx = ExecContext(vcc=0b1010, lane=1)
-    # Lane 1 is set
-    self.assertEqual(ctx.VCC.u64[1], 1)
-    self.assertEqual(ctx.VCC.u64[2], 0)
-
-  def test_vcc_lane_write(self):
-    ctx = ExecContext(vcc=0, lane=0)
-    ctx.VCC.u64[3] = 1
-    ctx.VCC.u64[1] = 1
-    self.assertEqual(ctx.VCC._val, 0b1010)
-
-  def test_for_loop(self):
-    # CTZ pattern - find first set bit
-    code = compile_pseudocode("""tmp = -1
-for i in 0 : 31 do
-  if S0.u32[i] == 1 then
-    tmp = i
-  endif
-endfor
-D0.i32 = tmp""")
-    ctx = ExecContext(s0=0b1000)  # Bit 3 is set
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val & MASK32, 3)
-
-  def test_result_dict(self):
-    ctx = ExecContext(s0=5, s1=3)
-    ctx.D0.u32 = 42
-    ctx.SCC._val = 1
-    result = ctx.result()
-    self.assertEqual(result['d0'], 42)
-    self.assertEqual(result['scc'], 1)
-
-class TestPseudocodeRegressions(unittest.TestCase):
-  """Regression tests for pseudocode instruction emulation bugs."""
-
-  def test_v_div_scale_f32_vcc_always_returned(self):
-    """V_DIV_SCALE_F32 must always return VCC, even when VCC=0 (no scaling needed).
-    Bug: when VCC._val == vcc (both 0), VCC wasn't returned, so VCC bits weren't written.
-    This caused division to produce wrong results for multiple lanes."""
-    # Normal case: 1.0 / 3.0, no scaling needed, VCC should be 0
-    S0 = Reg(0x3f800000)  # 1.0
-    S1 = Reg(0x40400000)  # 3.0
-    S2 = Reg(0x3f800000)  # 1.0 (numerator)
-    D0, SCC, VCC, EXEC = Reg(0), Reg(0), Reg(0), Reg(0xffffffff)
-    result = _VOP3SDOp_V_DIV_SCALE_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
-    # Must always have VCC in result
-    self.assertIn('VCC', result, "V_DIV_SCALE_F32 must always return VCC")
-    self.assertEqual(result['VCC']._val & 1, 0, "VCC lane 0 should be 0 when no scaling needed")
-
-  def test_v_cmp_class_f32_detects_quiet_nan(self):
-    """V_CMP_CLASS_F32 must correctly identify quiet NaN vs signaling NaN.
-    Bug: isQuietNAN and isSignalNAN both used math.isnan which can't distinguish them."""
-    quiet_nan = 0x7fc00000   # quiet NaN: exponent=255, bit22=1
-    signal_nan = 0x7f800001  # signaling NaN: exponent=255, bit22=0
-    # Test quiet NaN detection (bit 1 in mask)
-    s1_quiet = 0b0000000010  # bit 1 = quiet NaN
-    S0, S1, S2, D0, SCC, VCC, EXEC = Reg(quiet_nan), Reg(s1_quiet), Reg(0), Reg(0), Reg(0), Reg(0), Reg(0xffffffff)
-    result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
-    self.assertEqual(result['D0']._val & 1, 1, "Should detect quiet NaN with quiet NaN mask")
-    # Test signaling NaN detection (bit 0 in mask)
-    s1_signal = 0b0000000001  # bit 0 = signaling NaN
-    S0, S1 = Reg(signal_nan), Reg(s1_signal)
-    result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
-    self.assertEqual(result['D0']._val & 1, 1, "Should detect signaling NaN with signaling NaN mask")
-    # Test that quiet NaN doesn't match signaling NaN mask
-    S0, S1 = Reg(quiet_nan), Reg(s1_signal)
-    result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
-    self.assertEqual(result['D0']._val & 1, 0, "Quiet NaN should not match signaling NaN mask")
-    # Test that signaling NaN doesn't match quiet NaN mask
-    S0, S1 = Reg(signal_nan), Reg(s1_quiet)
-    result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
-    self.assertEqual(result['D0']._val & 1, 0, "Signaling NaN should not match quiet NaN mask")
-
-  def test_isnan_with_typed_view(self):
-    """_isnan must work with TypedView objects, not just Python floats.
-    Bug: _isnan checked isinstance(x, float) which returned False for TypedView."""
-    nan_reg = Reg(0x7fc00000)  # quiet NaN
-    normal_reg = Reg(0x3f800000)  # 1.0
-    inf_reg = Reg(0x7f800000)  # +inf
-    self.assertTrue(_isnan(nan_reg.f32), "_isnan should return True for NaN TypedView")
-    self.assertFalse(_isnan(normal_reg.f32), "_isnan should return False for normal TypedView")
-    self.assertFalse(_isnan(inf_reg.f32), "_isnan should return False for inf TypedView")
-
-class TestBF16(unittest.TestCase):
-  """Tests for BF16 (bfloat16) support."""
-
-  def test_bf16_conversion(self):
-    """Test bf16 <-> f32 conversion."""
-    # bf16 is just the top 16 bits of f32
-    # 1.0f = 0x3f800000, bf16 = 0x3f80
-    self.assertAlmostEqual(_bf16(0x3f80), 1.0, places=2)
-    self.assertEqual(_ibf16(1.0), 0x3f80)
-    # 2.0f = 0x40000000, bf16 = 0x4000
-    self.assertAlmostEqual(_bf16(0x4000), 2.0, places=2)
-    self.assertEqual(_ibf16(2.0), 0x4000)
-    # -1.0f = 0xbf800000, bf16 = 0xbf80
-    self.assertAlmostEqual(_bf16(0xbf80), -1.0, places=2)
-    self.assertEqual(_ibf16(-1.0), 0xbf80)
-
-  def test_bf16_special_values(self):
-    """Test bf16 special values (inf, nan)."""
-    import math
-    # +inf: f32 = 0x7f800000, bf16 = 0x7f80
-    self.assertTrue(math.isinf(_bf16(0x7f80)))
-    self.assertEqual(_ibf16(float('inf')), 0x7f80)
-    # -inf: f32 = 0xff800000, bf16 = 0xff80
-    self.assertTrue(math.isinf(_bf16(0xff80)))
-    self.assertEqual(_ibf16(float('-inf')), 0xff80)
-    # NaN: quiet NaN bf16 = 0x7fc0
-    self.assertTrue(math.isnan(_bf16(0x7fc0)))
-    self.assertEqual(_ibf16(float('nan')), 0x7fc0)
-
-  def test_bf16_register_property(self):
-    """Test Reg.bf16 property."""
-    r = Reg(0)
-    r.bf16 = 3.0  # 3.0f = 0x40400000, bf16 = 0x4040
-    self.assertEqual(r._val & 0xffff, 0x4040)
-    self.assertAlmostEqual(float(r.bf16), 3.0, places=1)
-
-  def test_bf16_slice_property(self):
-    """Test SliceProxy.bf16 property."""
-    r = Reg(0x40404040)  # Two bf16 3.0 values
-    self.assertAlmostEqual(r[15:0].bf16, 3.0, places=1)
-    self.assertAlmostEqual(r[31:16].bf16, 3.0, places=1)
-
-class TestBytePermute(unittest.TestCase):
-  """Tests for BYTE_PERMUTE helper function (V_PERM_B32)."""
-
-  def test_byte_select_0_to_7(self):
-    """Test selecting bytes 0-7 from 64-bit data."""
-    # data = {s0, s1} where s0 is bytes 0-3, s1 is bytes 4-7
-    # Combined: 0x0706050403020100 (byte 0 = 0x00, byte 7 = 0x07)
-    data = 0x0706050403020100
-    for i in range(8):
-      self.assertEqual(BYTE_PERMUTE(data, i), i, f"byte {i} should be {i}")
-
-  def test_sign_extend_bytes(self):
-    """Test sign extension selectors 8-11."""
-    # sel 8: sign of byte 1 (bits 15:8)
-    # sel 9: sign of byte 3 (bits 31:24)
-    # sel 10: sign of byte 5 (bits 47:40)
-    # sel 11: sign of byte 7 (bits 63:56)
-    data = 0x8000800080008000  # All relevant bytes have sign bit set
-    self.assertEqual(BYTE_PERMUTE(data, 8), 0xff)
-    self.assertEqual(BYTE_PERMUTE(data, 9), 0xff)
-    self.assertEqual(BYTE_PERMUTE(data, 10), 0xff)
-    self.assertEqual(BYTE_PERMUTE(data, 11), 0xff)
-    data = 0x7f007f007f007f00  # No sign bits set
-    self.assertEqual(BYTE_PERMUTE(data, 8), 0x00)
-    self.assertEqual(BYTE_PERMUTE(data, 9), 0x00)
-    self.assertEqual(BYTE_PERMUTE(data, 10), 0x00)
-    self.assertEqual(BYTE_PERMUTE(data, 11), 0x00)
-
-  def test_constant_zero(self):
-    """Test selector 12 returns 0x00."""
-    self.assertEqual(BYTE_PERMUTE(0xffffffffffffffff, 12), 0x00)
-
-  def test_constant_ff(self):
-    """Test selectors >= 13 return 0xFF."""
-    for sel in [13, 14, 15, 255]:
-      self.assertEqual(BYTE_PERMUTE(0, sel), 0xff, f"sel {sel} should be 0xff")
-
-class TestSADHelpers(unittest.TestCase):
-  """Tests for V_SAD_U8 and V_MSAD_U8 helper functions."""
-
-  def test_v_sad_u8_basic(self):
-    """Test v_sad_u8 with simple values."""
-    # s0 = 0x04030201, s1 = 0x04030201 -> diff = 0 for all bytes
-    result = v_sad_u8(0x04030201, 0x04030201, 0)
-    self.assertEqual(result, 0)
-    # s0 = 0x05040302, s1 = 0x04030201 -> diff = 1+1+1+1 = 4
-    result = v_sad_u8(0x05040302, 0x04030201, 0)
-    self.assertEqual(result, 4)
-
-  def test_v_sad_u8_with_accumulator(self):
-    """Test v_sad_u8 with non-zero accumulator."""
-    # s0 = 0x05040302, s1 = 0x04030201, s2 = 100 -> 4 + 100 = 104
-    result = v_sad_u8(0x05040302, 0x04030201, 100)
-    self.assertEqual(result, 104)
-
-  def test_v_sad_u8_large_diff(self):
-    """Test v_sad_u8 with maximum byte differences."""
-    # s0 = 0xffffffff, s1 = 0x00000000 -> diff = 255*4 = 1020
-    result = v_sad_u8(0xffffffff, 0x00000000, 0)
-    self.assertEqual(result, 1020)
-
-  def test_v_msad_u8_basic(self):
-    """Test v_msad_u8 masks when reference byte is 0."""
-    # s0 = 0x10101010, s1 = 0x00000000 -> all masked, result = 0
-    result = v_msad_u8(0x10101010, 0x00000000, 0)
-    self.assertEqual(result, 0)
-    # s0 = 0x10101010, s1 = 0x01010101 -> diff = |0x10-0x01|*4 = 15*4 = 60
-    result = v_msad_u8(0x10101010, 0x01010101, 0)
-    self.assertEqual(result, 60)
-
-  def test_v_msad_u8_partial_mask(self):
-    """Test v_msad_u8 with partial masking."""
-    # s0 = 0x10101010, s1 = 0x00010001 -> bytes 1 and 3 masked
-    # diff = |0x10-0x01| + |0x10-0x01| = 15 + 15 = 30
-    result = v_msad_u8(0x10101010, 0x00010001, 0)
-    self.assertEqual(result, 30)
-
-  def test_v_msad_u8_with_accumulator(self):
-    """Test v_msad_u8 with non-zero accumulator."""
-    result = v_msad_u8(0x10101010, 0x01010101, 50)
-    self.assertEqual(result, 110)  # 60 + 50
-
-if __name__ == '__main__':
-  unittest.main()
@@ -1,150 +0,0 @@
-#!/usr/bin/env python3
-"""Test that PDF parser correctly extracts format fields."""
-import unittest, os
-from extra.assembly.amd.autogen.rdna3.ins import SOP1, SOP2, SOPK, SOPP, VOP1, VOP2, VOP3SD, VOPC, FLAT, VOPD, SOP1Op, SOP2Op, VOP1Op, VOP3Op
-
-# expected formats with key fields and whether they have ENCODING
-EXPECTED_FORMATS = {
-  'DPP16': (['SRC0', 'DPP_CTRL', 'BANK_MASK', 'ROW_MASK'], False),
-  'DPP8': (['SRC0', 'LANE_SEL0', 'LANE_SEL7'], False),
-  'DS': (['OP', 'ADDR', 'DATA0', 'DATA1', 'VDST'], True),
-  'EXP': (['EN', 'TARGET', 'VSRC0', 'VSRC1', 'VSRC2', 'VSRC3'], True),
-  'FLAT': (['OP', 'ADDR', 'DATA', 'SADDR', 'VDST', 'OFFSET'], True),
-  'LDSDIR': (['VDST', 'OP'], True),
-  'MIMG': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'DMASK'], True),
-  'MTBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'FORMAT', 'SOFFSET'], True),
-  'MUBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'SOFFSET'], True),
-  'SMEM': (['OP', 'SBASE', 'SDATA', 'OFFSET', 'SOFFSET'], True),
-  'SOP1': (['OP', 'SDST', 'SSRC0'], True),
-  'SOP2': (['OP', 'SDST', 'SSRC0', 'SSRC1'], True),
-  'SOPC': (['OP', 'SSRC0', 'SSRC1'], True),
-  'SOPK': (['OP', 'SDST', 'SIMM16'], True),
-  'SOPP': (['OP', 'SIMM16'], True),
-  'VINTERP': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOP1': (['OP', 'VDST', 'SRC0'], True),
-  'VOP2': (['OP', 'VDST', 'SRC0', 'VSRC1'], True),
-  'VOP3': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOP3P': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOP3SD': (['OP', 'VDST', 'SDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOPC': (['OP', 'SRC0', 'VSRC1'], True),
-  'VOPD': (['OPX', 'OPY', 'SRCX0', 'SRCY0', 'VDSTX', 'VDSTY'], True),
-}
-
-# Skip PDF parsing tests by default - only run with TEST_PDF_PARSER=1
-# These are slow (~5s) and only needed when regenerating autogen/
-@unittest.skipUnless(os.environ.get("TEST_PDF_PARSER"), "set TEST_PDF_PARSER=1 to run PDF parser tests")
-class TestPDFParserGenerate(unittest.TestCase):
-  """Test the PDF parser by running generate() and checking results."""
-
-  def test_pdf_parser(self):
-    """Single test that validates all PDF parser outputs."""
-    from extra.assembly.amd.dsl import generate
-    result = generate()
-
-    # test_all_formats_present
-    for fmt_name in EXPECTED_FORMATS:
-      self.assertIn(fmt_name, result["formats"], f"missing format {fmt_name}")
-
-    # test_format_count
-    self.assertEqual(len(result["formats"]), 23)
-
-    # test_no_duplicate_fields
-    for fmt_name, fields in result["formats"].items():
-      field_names = [f[0] for f in fields]
-      self.assertEqual(len(field_names), len(set(field_names)), f"{fmt_name} has duplicate fields: {field_names}")
-
-    # test_expected_fields
-    for fmt_name, (expected_fields, has_encoding) in EXPECTED_FORMATS.items():
-      fields = {f[0] for f in result["formats"].get(fmt_name, [])}
-      for field in expected_fields:
-        self.assertIn(field, fields, f"{fmt_name} missing {field}")
-      if has_encoding:
-        self.assertIn("ENCODING", fields, f"{fmt_name} should have ENCODING")
-      else:
-        self.assertNotIn("ENCODING", fields, f"{fmt_name} should not have ENCODING")
-
-    # test_vopd_no_dpp16_fields
-    vopd_fields = {f[0] for f in result["formats"].get("VOPD", [])}
-    for field in ['DPP_CTRL', 'BANK_MASK', 'ROW_MASK']:
-      self.assertNotIn(field, vopd_fields, f"VOPD should not have {field}")
-
-    # test_dpp16_no_vinterp_fields
-    dpp16_fields = {f[0] for f in result["formats"].get("DPP16", [])}
-    for field in ['VDST', 'WAITEXP']:
-      self.assertNotIn(field, dpp16_fields, f"DPP16 should not have {field}")
-
-    # test_sopp_no_smem_fields
-    sopp_fields = {f[0] for f in result["formats"].get("SOPP", [])}
-    for field in ['SBASE', 'SDATA']:
-      self.assertNotIn(field, sopp_fields, f"SOPP should not have {field}")
-
-class TestPDFParser(unittest.TestCase):
-  """Verify format classes have correct fields from PDF parsing."""
-
-  def test_sop2_fields(self):
-    """SOP2 should have op, sdst, ssrc0, ssrc1."""
-    for field in ['op', 'sdst', 'ssrc0', 'ssrc1']:
-      self.assertIn(field, SOP2._fields)
-    self.assertEqual(SOP2._fields['op'].hi, 29)
-    self.assertEqual(SOP2._fields['op'].lo, 23)
-
-  def test_sop1_fields(self):
-    """SOP1 should have op, sdst, ssrc0 with correct bit positions."""
-    for field in ['op', 'sdst', 'ssrc0']:
-      self.assertIn(field, SOP1._fields)
-    self.assertNotIn('simm16', SOP1._fields)
-    self.assertEqual(SOP1._fields['ssrc0'].hi, 7)
-    self.assertEqual(SOP1._fields['ssrc0'].lo, 0)
-    assert SOP1._encoding is not None
-    self.assertEqual(SOP1._encoding[0].hi, 31)
-    self.assertEqual(SOP1._encoding[1], 0b101111101)
-
-  def test_vop3sd_fields(self):
-    """VOP3SD should have all fields including src0/src1/src2 from page continuation."""
-    for field in ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2']:
-      self.assertIn(field, VOP3SD._fields)
-    self.assertEqual(VOP3SD._fields['src0'].hi, 40)
-    self.assertEqual(VOP3SD._fields['src0'].lo, 32)
-    self.assertEqual(VOP3SD._size(), 8)
-
-  def test_flat_has_vdst(self):
-    """FLAT should have vdst field."""
-    self.assertIn('vdst', FLAT._fields)
-    self.assertEqual(FLAT._fields['vdst'].hi, 63)
-    self.assertEqual(FLAT._fields['vdst'].lo, 56)
-
-  def test_encoding_bits(self):
-    """Verify encoding bits are correct for major formats."""
-    tests = [
-      (SOP2, 31, 30, 0b10),
-      (SOPK, 31, 28, 0b1011),
-      (SOPP, 31, 23, 0b101111111),
-      (VOP1, 31, 25, 0b0111111),
-      (VOP2, 31, 31, 0b0),
-      (VOPC, 31, 25, 0b0111110),
-      (FLAT, 31, 26, 0b110111),
-    ]
-    for cls, hi, lo, val in tests:
-      assert cls._encoding is not None
-      self.assertEqual(cls._encoding[0].hi, hi, f"{cls.__name__} encoding hi")
-      self.assertEqual(cls._encoding[0].lo, lo, f"{cls.__name__} encoding lo")
-      self.assertEqual(cls._encoding[1], val, f"{cls.__name__} encoding val")
-
-  def test_opcode_enums_exist(self):
-    """Verify opcode enums are generated with expected counts."""
-    self.assertGreater(len(SOP1Op), 50)
-    self.assertGreater(len(SOP2Op), 50)
-    self.assertGreater(len(VOP1Op), 50)
-    self.assertGreater(len(VOP3Op), 200)
-
-  def test_vopd_no_duplicate_fields(self):
-    """VOPD should not have duplicate fields and should not include DPP16 fields."""
-    field_names = list(VOPD._fields.keys())
-    self.assertEqual(len(field_names), len(set(field_names)))
-    for field in ['srcx0', 'srcy0', 'opx', 'opy']:
-      self.assertIn(field, VOPD._fields)
-    for field in ['dpp_ctrl', 'bank_mask', 'row_mask']:
-      self.assertNotIn(field, VOPD._fields)
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,95 +0,0 @@
-#!/usr/bin/env python3
-import unittest, subprocess
-from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.test.helpers import get_llvm_mc
-
-def llvm_assemble(asm: str) -> bytes:
-  """Assemble using llvm-mc and return bytes."""
-  result = subprocess.run(
-    [get_llvm_mc(), "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
-    input=asm, capture_output=True, text=True
-  )
-  out = b''
-  for line in result.stdout.split('\n'):
-    if 'encoding:' in line:
-      enc = line.split('encoding:')[1].strip()
-      enc = enc.strip('[]').replace('0x', '').replace(',', '')
-      out += bytes.fromhex(enc)
-  if not out: raise ValueError(f"no encoding found: {result.stdout} {result.stderr}")
-  return out
-
-class TestRDNA3Asm(unittest.TestCase):
-  def test_full_program(self):
-    """Test the full program from rdna3fun.py matches llvm-mc output."""
-    program = [
-      v_bfe_u32(v[1], v[0], 10, 10),
-      s_load_b128(s[4:7], s[0:1], NULL),
-      v_and_b32_e32(v[0], 0x3FF, v[0]),
-      s_mulk_i32(s[3], 0x87),
-      v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2]),
-      v_mul_u32_u24_e32(v[0], 45, v[0]),
-      v_ashrrev_i32_e32(v[2], 31, v[1]),
-      v_add3_u32(v[0], v[0], s[3], v[1]),
-      v_lshlrev_b64(v[2:3], 2, v[1:2]),
-      v_ashrrev_i32_e32(v[1], 31, v[0]),
-      v_lshlrev_b64(v[0:1], 2, v[0:1]),
-      s_waitcnt(0xfc07),  # lgkmcnt(0)
-      v_add_co_u32(v[2], VCC_LO, s[6], v[2]),
-      v_add_co_ci_u32_e32(v[3], s[7], v[3]),
-      v_add_co_u32(v[0], VCC_LO, s[4], v[0]),
-      global_load_b32(vdst=v[2], addr=v[2], saddr=OFF),
-      v_add_co_ci_u32_e32(v[1], s[5], v[1]),
-      s_waitcnt(0x03f7),  # vmcnt(0)
-      global_store_b32(addr=v[0], data=v[2], saddr=OFF),
-      s_endpgm(),
-    ]
-
-    asm = """
-v_bfe_u32 v1, v0, 10, 10
-s_load_b128 s[4:7], s[0:1], null
-v_and_b32_e32 v0, 0x3FF, v0
-s_mulk_i32 s3, 0x87
-v_mad_u64_u32 v[1:2], null, s2, 3, v[1:2]
-v_mul_u32_u24_e32 v0, 45, v0
-v_ashrrev_i32_e32 v2, 31, v1
-v_add3_u32 v0, v0, s3, v1
-v_lshlrev_b64 v[2:3], 2, v[1:2]
-v_ashrrev_i32_e32 v1, 31, v0
-v_lshlrev_b64 v[0:1], 2, v[0:1]
-s_waitcnt lgkmcnt(0)
-v_add_co_u32 v2, vcc_lo, s6, v2
-v_add_co_ci_u32_e32 v3, vcc_lo, s7, v3, vcc_lo
-v_add_co_u32 v0, vcc_lo, s4, v0
-global_load_b32 v2, v[2:3], off
-v_add_co_ci_u32_e32 v1, vcc_lo, s5, v1, vcc_lo
-s_waitcnt vmcnt(0)
-global_store_b32 v[0:1], v2, off
-s_endpgm
-"""
-    expected = llvm_assemble(asm)
-    for inst,rt in zip(program, asm.strip().split("\n")): print(f"{inst.disasm():50s} {rt}")
-    actual = b''.join(inst.to_bytes() for inst in program)
-    self.assertEqual(actual, expected)
-
-  def test_sop2_s_add_u32(self):
-    inst = SOP2(SOP2Op.S_ADD_U32, s[3], s[0], s[1])
-    expected = llvm_assemble("s_add_u32 s3, s0, s1")
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_vop2_v_and_b32_inline_const(self):
-    inst = v_and_b32_e32(v[0], 10, v[0])
-    expected = llvm_assemble("v_and_b32_e32 v0, 10, v0")
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_sopp_s_endpgm(self):
-    inst = s_endpgm()
-    expected = llvm_assemble("s_endpgm")
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_sop1_s_mov_b32(self):
-    inst = s_mov_b32(s[0], s[1])
-    expected = llvm_assemble("s_mov_b32 s0, s1")
-    self.assertEqual(inst.to_bytes(), expected)
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,263 +0,0 @@
-#!/usr/bin/env python3
-"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
-import unittest, io, sys, re, subprocess, os
-from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.dsl import Inst
-from extra.assembly.amd.asm import asm
-from extra.assembly.amd.asm import detect_format
-from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump
-
-def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
-  """Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
-  old_stdout = sys.stdout
-  sys.stdout = io.StringIO()
-  compiler.disassemble(lib)
-  output = sys.stdout.getvalue()
-  sys.stdout = old_stdout
-
-  results = []
-  for line in output.splitlines():
-    if '//' not in line: continue
-    instr = line.split('//')[0].strip()
-    if not instr: continue
-    comment = line.split('//')[1].strip()
-    if ':' not in comment: continue
-    hex_str = comment.split(':')[1].strip().split()[0]
-    try:
-      machine_bytes = bytes.fromhex(hex_str)[::-1]  # big-endian to little-endian
-      results.append((instr, machine_bytes))
-    except ValueError:
-      continue
-  return results
-
-def compile_asm(instr: str, compiler=None) -> bytes:
-  """Compile a single instruction with llvm-mc and return the machine code bytes."""
-  llvm_mc = get_llvm_mc()
-  result = subprocess.run(
-    [llvm_mc, '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
-    input=f".text\n{instr}\n", capture_output=True, text=True)
-  if result.returncode != 0: raise RuntimeError(f"llvm-mc failed for '{instr}': {result.stderr.strip()}")
-  # Parse encoding: [0x01,0x39,0x0a,0x7e]
-  for line in result.stdout.split('\n'):
-    if 'encoding:' in line:
-      enc = line.split('encoding:')[1].strip()
-      if enc.startswith('[') and enc.endswith(']'):
-        hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
-        return bytes.fromhex(hex_vals)
-  raise RuntimeError(f"no encoding found in llvm-mc output for: {instr}")
-
-def compile_asm_batch(instrs: list[str]) -> list[bytes]:
-  """Compile multiple instructions with a single llvm-mc call."""
-  if not instrs: return []
-  llvm_mc = get_llvm_mc()
-  src = ".text\n" + "\n".join(instrs) + "\n"
-  result = subprocess.run(
-    [llvm_mc, '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
-    input=src, capture_output=True, text=True)
-  if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
-  # Parse all encodings in order
-  encodings = []
-  for line in result.stdout.split('\n'):
-    if 'encoding:' in line:
-      enc = line.split('encoding:')[1].strip()
-      if enc.startswith('[') and enc.endswith(']'):
-        hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
-        encodings.append(bytes.fromhex(hex_vals))
-  if len(encodings) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(encodings)}")
-  return encodings
-
-def compile_and_disasm_batch(instrs: list[str], compiler) -> list[str]:
-  """Compile instructions with LLVM and get LLVM's disassembly."""
-  import tempfile, os
-  if not instrs: return []
-  # Build assembly source with all instructions
-  src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n"
-  src += "\n".join(f"  {instr}" for instr in instrs) + "\n"
-  # Use llvm-mc to assemble to object file
-  with tempfile.NamedTemporaryFile(suffix='.o', delete=False) as f:
-    obj_path = f.name
-  try:
-    result = subprocess.run(
-      [get_llvm_mc(), '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-filetype=obj', '-o', obj_path],
-      input=src, capture_output=True, text=True)
-    if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
-    # Disassemble with llvm-objdump
-    result = subprocess.run([get_llvm_objdump(), '-d', '--mcpu=gfx1100', obj_path], capture_output=True, text=True)
-    if result.returncode != 0: raise RuntimeError(f"llvm-objdump failed: {result.stderr.strip()}")
-    # Parse disassembly output
-    results: list[str] = []
-    for line in result.stdout.splitlines():
-      if '//' not in line: continue
-      instr = line.split('//')[0].strip()
-      if instr: results.append(instr)
-    return results[:len(instrs)]
-  finally:
-    os.unlink(obj_path)
-
-class TestTinygradKernelRoundtrip(unittest.TestCase):
-  """Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
-
-  def _test_kernel_roundtrip(self, op_fn):
-    """Generate kernel from op_fn, test:
-    1. decode -> reencode matches original bytes
-    2. asm(disasm()) matches LLVM output
-    3. our disasm() matches LLVM's disassembly string exactly
-    """
-    from extra.assembly.amd.test.test_compare_emulators import get_kernels_from_tinygrad
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-    kernels, _, _ = get_kernels_from_tinygrad(op_fn)
-    compiler = HIPCompiler('gfx1100')
-
-    # First pass: decode all instructions and collect info
-    decoded_instrs: list[tuple] = []  # list of (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err)
-    for ki, kernel in enumerate(kernels):
-      offset = 0
-      while offset < len(kernel.code):
-        remaining = kernel.code[offset:]
-        fmt = detect_format(remaining)
-        if fmt is None:
-          decoded_instrs.append((ki, offset, None, None, None, False, "no format"))
-          offset += 4
-          continue
-
-        base_size = fmt._size()
-        if len(remaining) < base_size:
-          break
-
-        try:
-          decoded = fmt.from_bytes(remaining)  # pass all remaining bytes so from_bytes can read literal
-          size = decoded.size()  # actual size including literal
-          orig_bytes = remaining[:size]
-          reencoded = decoded.to_bytes()
-          our_disasm = decoded.disasm()
-          decode_ok = reencoded == orig_bytes
-          decode_err: str | None = None if decode_ok else f"orig={orig_bytes.hex()} reenc={reencoded.hex()}"
-          decoded_instrs.append((ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err))
-        except Exception as e:
-          decoded_instrs.append((ki, offset, remaining[:base_size], None, None, False, str(e)))
-          size = base_size
-
-        offset += size
-
-    # Collect disasm strings for batched LLVM calls - skip unknown opcodes (op_X) that LLVM can't compile
-    asm_test_instrs: list[tuple[int, str]] = []  # (idx, our_disasm) for asm test
-    disasm_test_instrs: list[tuple[int, str]] = []  # (idx, our_disasm) for disasm comparison test
-
-    for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
-      if our_disasm is None: continue
-      # Skip unknown opcodes and malformed instructions for both tests
-      if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm): continue
-      asm_test_instrs.append((idx, our_disasm))
-      disasm_test_instrs.append((idx, our_disasm))
-
-    # Batch compile for asm test
-    asm_llvm_results = compile_asm_batch([d for _, d in asm_test_instrs])
-    asm_llvm_map = {idx: result for (idx, _), result in zip(asm_test_instrs, asm_llvm_results)}
-
-    # Batch compile+disasm for disasm comparison test
-    disasm_llvm_results = compile_and_disasm_batch([d for _, d in disasm_test_instrs], compiler)
-    disasm_llvm_map = {idx: result for (idx, _), result in zip(disasm_test_instrs, disasm_llvm_results)}
-
-    # Now evaluate results
-    decode_passed, decode_failed, decode_skipped = 0, 0, 0
-    asm_passed, asm_failed, asm_skipped = 0, 0, 0
-    disasm_passed, disasm_failed, disasm_skipped = 0, 0, 0
-    decode_failures: list[str] = []
-    asm_failures: list[str] = []
-    disasm_failures: list[str] = []
-
-    for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
-      # Decode test
-      if decode_ok:
-        decode_passed += 1
-      elif decode_err == "no format":
-        decode_skipped += 1
-      else:
-        decode_failed += 1
-        decode_failures.append(f"K{ki}@{offset}: {our_disasm}: {decode_err}")
-
-      # Asm test
-      if our_disasm is None:
-        asm_skipped += 1
-      elif idx in asm_llvm_map:
-        llvm_bytes = asm_llvm_map[idx]
-        try:
-          our_bytes = asm(our_disasm).to_bytes()
-          if our_bytes[:len(llvm_bytes)] == llvm_bytes:
-            asm_passed += 1
-          else:
-            asm_failed += 1
-            asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': ours={our_bytes[:len(llvm_bytes)].hex()} llvm={llvm_bytes.hex()}")
-        except Exception:
-          asm_skipped += 1
-      else:
-        asm_skipped += 1
-
-      # Disasm comparison test
-      if our_disasm is None:
-        disasm_skipped += 1
-      elif idx in disasm_llvm_map:
-        llvm_disasm = disasm_llvm_map[idx]
-        if our_disasm == llvm_disasm:
-          disasm_passed += 1
-        else:
-          disasm_failed += 1
-          disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
-      else:
-        disasm_skipped += 1
-
-    print(f"decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
-    print(f"asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
-    print(f"disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
-    self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
-    self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
-    # Note: disasm string comparison is informational only - formatting differences between LLVM versions are expected
-
-  # Basic unary ops
-  def test_neg(self): self._test_kernel_roundtrip(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
-  def test_relu(self): self._test_kernel_roundtrip(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
-  def test_exp(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).exp())
-  def test_log(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 3.0]).log())
-  def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
-  def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
-  def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
-
-  # Binary ops
-  def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
-  def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
-  def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
-  def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
-  def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
-
-  # Reductions
-  def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
-  def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
-  def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
-
-  # Matmul
-  def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
-  def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
-
-  # Complex ops
-  def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
-  def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
-
-  # Memory patterns
-  def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
-  def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
-  def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
-
-  # Cast ops
-  def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
-  def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
-
-  # Comparison ops
-  def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
-  def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
-
-  # Fused ops
-  def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
-
-if __name__ == "__main__":
-  unittest.main()
@@ -0,0 +1,189 @@
+from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
+from tinygrad.codegen.opt.kernel import Ops, MemOp, UOp
+from tinygrad.uop.ops import BinaryOps, UnaryOps
+from tinygrad.dtype import DType, dtypes
+from tinygrad.helpers import DEBUG
+from tinygrad.uop.ops import Variable, NumNode, MulNode, DivNode, ModNode, LtNode, SumNode, AndNode
+import functools
+import math
+from collections import defaultdict
+
+_type_to_letter = {dtypes.float32: 'f', dtypes.bool: 'p', dtypes.int32: 'i', dtypes.int64: 'a', dtypes.uint32: 'u', dtypes.uint64: 'b', dtypes.float.vec(4): 'x', dtypes.uint8: 'uc', dtypes.float16: 'h',
+                   dtypes.int8: 'c', dtypes.uint16: 'us', dtypes.float64: 'd'}
+
+class Register(NamedTuple):
+  nm:str
+  dtype:DType
+  scalar:bool
+  off:Optional[int] = None
+  def __repr__(self): return self.nm if self.off is None else f"{self.nm}:{self.off}"
+  def subregs(self):
+    if self.dtype == dtypes.float.vec(4):
+      return [Register(self.nm, dtypes.float, False, off=off) for off in range(4)]
+    return []
+
+class AssemblyInstruction(NamedTuple):
+  op: Ops
+  out: Optional[Register]
+  vin: List[Union[Register, int, float]]
+  arg: Any = None
+
+# warp size of 32, s registers are shared across the warp, v are 32-wide vectors
+class AssemblyLanguage:
+  supports_load3: bool = False
+  sin_is_sin2pi: bool = False
+  no_div: bool = False
+  #TODO: these should be global vars
+  cnts:DefaultDict[Tuple[DType, bool], int] = defaultdict(int)
+  tor: Dict[Any, Register] = {}
+  ins: List[AssemblyInstruction] = []
+
+  def type_to_letter(self,x): return _type_to_letter[x[0]].upper() if x[1] else _type_to_letter[x[0]]
+  def newreg(self, tok, dtype=dtypes.float32, scalar=False) -> Register:
+    self.tor[tok] = ret = Register(f"%{self.type_to_letter((dtype, scalar))}{self.cnts[(dtype, scalar)]}", dtype, scalar)
+    if dtype == dtypes.float.vec(4):
+      for off in range(4):
+        self.tor[tok] = Register(ret.nm, dtypes.float, ret.scalar, off)
+    self.cnts[(dtype, scalar)] += 1
+    return ret
+
+  def render_numnode(self, b) -> Register:
+    key = ("num", b)
+    if key not in self.tor: self.ins.append(AssemblyInstruction(Ops.LOAD, self.newreg(key, scalar=True, dtype=dtypes.int32), [], b))
+    return self.tor[key]
+
+  def render_alu(self, op, a:Register, b:Union[Register, int, float], dtype=dtypes.int32) -> Register:
+    key = (op, a, b)
+    if key not in self.tor:
+      #if not isinstance(b, Register): b = render_numnode(b)
+      self.ins.append(AssemblyInstruction(Ops.ALU, self.newreg(key, dtype=dtype, scalar=a.scalar and (not isinstance(b, Register) or b.scalar)), [a, b], op))
+    return self.tor[key]
+
+  def render_cast(self, a:Register, new_dtype:DType) -> Register:
+    if a.dtype == new_dtype: return a
+    key = (a, new_dtype)
+    if key not in self.tor:
+      self.ins.append(AssemblyInstruction(Ops.CAST, self.newreg(key, dtype=new_dtype), [a]))
+    return self.tor[key]
+
+  render_ops: Any = { Variable: lambda self, ops, ctx: ctx.tor[self], NumNode: lambda self, ops, ctx: ctx.render_numnode(self.b),
+                 MulNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.MUL, self.a.render(ops, ctx), self.b),
+                 DivNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.DIV, self.a.render(ops, ctx), self.b),
+                 ModNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.MOD, self.a.render(ops, ctx), self.b),
+                 LtNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.CMPLT, self.a.render(ops, ctx), self.b, dtype=dtypes.bool),
+    SumNode: lambda self,ops,ctx: functools.reduce(lambda a,b: ctx.render_alu(BinaryOps.ADD, a, b.render(ops,ctx)), self.nodes[1:], self.nodes[0].render(ops,ctx)),
+    AndNode: lambda self,ops,ctx: functools.reduce(lambda a,b: ctx.render_alu(BinaryOps.MUL, a, b.render(ops,ctx), dtype=dtypes.bool), self.nodes[1:], self.nodes[0].render(ops,ctx)) }
+
+  def addr_w_offset(self, args):
+    assert isinstance(args, MemOp)
+    idx = args.idx*args.memory_dtype.itemsize
+    off = 0  # TODO: should this be None?
+    if isinstance(idx, SumNode):
+      nums = [n.b for n in idx.nodes if isinstance(n, NumNode)]
+      if nums and nums[0] < 4096 and (idx-nums[0]).min >= 0:  # TODO: different for each GPU?
+        idx -= nums[0]
+        off = cast(int, nums[0])
+    reg = idx.render(self.render_ops, self)
+    if self.supports_load3:
+      if reg.scalar:
+        new_reg = self.newreg((reg.nm, 'vec'), dtype=reg.dtype)
+        self.ins.append(AssemblyInstruction(Ops.ALU, new_reg, [reg], UnaryOps.NOOP))
+        reg = new_reg
+      return self.tor[args.name], reg, off
+    reg = self.render_alu(BinaryOps.ADD, self.render_cast(reg, dtypes.uint64), self.tor[args.name], dtype=dtypes.uint64)
+    return reg, None, off
+
+def uops_to_asmstyle(lang, function_name:str, uops:List[UOp]):
+  #TODO: Do not use clear()
+  lang.ins.clear()
+  lang.tor.clear()
+  lang.cnts.clear()
+  buf_to_dtype = {args:dtype for uop,dtype,_,args,_ in uops if uop == Ops.DEFINE_GLOBAL}
+  global_size, local_size = [], []
+  skipload_branch = 0
+  lang.ins += [AssemblyInstruction(Ops.SPECIAL, lang.newreg(buf, dtype=dtypes.uint64, scalar=True), [], buf) for buf in buf_to_dtype]
+  for u in uops:
+    uop,dtype,vin,args,_ = u
+    if uop == Ops.DEFINE_LOCAL:
+      lang.ins.append(AssemblyInstruction(Ops.DEFINE_LOCAL, None, [], args))
+      lang.ins.append(AssemblyInstruction(Ops.ALU, lang.newreg(args[0], dtype=dtypes.uint64), [args[0]], UnaryOps.NOOP))
+    elif uop == Ops.LOOP:
+      if args[1] == "global":
+        for i,var in enumerate(args[0]):
+          global_size.append(var.max+1)
+          lang.ins.append(AssemblyInstruction(Ops.SPECIAL, lang.newreg(var, dtype=dtypes.int32), [], f"gid{len(args[0])-1-i}"))
+      elif args[1] == "local":
+        for i,var in enumerate(args[0]):
+          local_size.append(var.max+1)
+          lang.ins.append(AssemblyInstruction(Ops.SPECIAL, lang.newreg(var, dtype=dtypes.int32), [], f"lid{len(args[0])-1-i}"))
+      else:
+        for var in args[0]:
+          if not isinstance(var, NumNode):  # TODO: why is this coming through?
+            lang.ins.append(AssemblyInstruction(Ops.LOAD, lang.newreg(var, dtype=dtypes.int32, scalar=True), [], 0))
+            lang.ins.append(AssemblyInstruction(Ops.LABEL, None, [], "$loop_"+var.expr))
+    elif uop == Ops.ENDLOOP:
+      if args[1] not in ["global", "local", "global+local"]:
+        for var in reversed(args[0]):
+          if not isinstance(var, NumNode):  # TODO: why is this coming through?
+            lang.ins.append(AssemblyInstruction(Ops.ALU, lang.tor[var], [lang.tor[var], 1], BinaryOps.ADD))
+            pred = lang.render_alu(BinaryOps.CMPLT, lang.tor[var], var.max+1, dtypes.bool)
+            lang.ins.append(AssemblyInstruction(Ops.COND_BRANCH, None, [pred], ("$loop_"+var.expr, True)))
+      elif args[1] == "global+local":
+        for i, var in enumerate(reversed(args[0])):
+          lang.ins.append(AssemblyInstruction(Ops.ENDLOOP, None, [lang.tor[var]], (var.max+1, f"gid{i}")))
+      elif args[1] == 'local':
+        for i, var in enumerate(reversed(args[0])):
+          lang.ins.append(AssemblyInstruction(Ops.ENDLOOP, None, [lang.tor[var]], (var.max+1, f"lid{i}")))
+    elif uop == Ops.CAST:
+      # TODO: we should reconsider outputting CAST in the linearizer. these are needless copies
+      out = lang.newreg(u, dtype)
+      for i,sr in enumerate(out.subregs()):
+        lang.ins.append(AssemblyInstruction(Ops.ALU, sr, [lang.tor[vin[i]]], UnaryOps.NOOP))
+    elif uop == Ops.ALU:
+      out = lang.newreg(u, dtype) if u not in lang.tor else lang.tor[u]
+      # this is the only thing that can violate SSA
+      if args in [BinaryOps.CMPLT]:
+        pred_reg = lang.newreg((u, 'pred'), dtype=dtypes.bool)
+        lang.ins.append(AssemblyInstruction(Ops.ALU, pred_reg, [lang.tor[x] for x in vin], args))
+        lang.ins.append(AssemblyInstruction(Ops.CAST, out, [pred_reg], args))
+      elif args == BinaryOps.DIV and lang.no_div:
+        tmp = lang.newreg((u, "rcp"))
+        lang.ins.append(AssemblyInstruction(Ops.ALU, tmp, [lang.tor[vin[1]]], UnaryOps.RECIP))
+        lang.ins.append(AssemblyInstruction(Ops.ALU, out, [lang.tor[vin[0]], tmp], BinaryOps.MUL))
+      elif args == UnaryOps.SIN and lang.sin_is_sin2pi:
+        tmp = lang.newreg((u, "2pi"))
+        lang.ins.append(AssemblyInstruction(Ops.ALU, tmp, [lang.tor[vin[0]], 1/(math.pi*2)], BinaryOps.MUL))
+        lang.ins.append(AssemblyInstruction(Ops.ALU, out, [tmp], args))
+      else:
+        lang.ins.append(AssemblyInstruction(Ops.ALU, out, [lang.tor[x] for x in vin], args))
+    elif uop == Ops.DEFINE_REG:
+      reg = lang.newreg(u, dtype=dtype)
+      lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [], args))
+    elif uop == Ops.SPECIAL:
+      lang.tor[u] = lang.tor[args]
+    elif uop == Ops.CONST:
+      lang.ins.append(AssemblyInstruction(Ops.LOAD, lang.newreg(u, dtype=dtype), [], args))
+    elif uop == Ops.LOAD:
+      idx, treg, off = lang.addr_w_offset(args)
+      reg = lang.newreg(u, dtype=dtype, scalar=(idx.scalar and (not isinstance(treg, Register) or treg.scalar)))
+      if args.valid.min == 0:
+        lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [], 0))
+        if args.valid.max == 1:
+          pred = args.valid.render(lang.render_ops, lang)
+          lang.ins.append(AssemblyInstruction(Ops.COND_BRANCH, None, [pred], (f"$skipload_{skipload_branch}", False)))
+      if args.valid.max == 1:
+          # NOTE: you can't compute the index in here, because it assumes it's all available later
+        lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [idx] + ([treg] if treg is not None else []), (off, 'global' if not args.local else 'shared', args.memory_dtype if args.memory_dtype != dtypes.float else None)))
+      if args.valid.min == 0 and args.valid.max == 1:
+        lang.ins.append(AssemblyInstruction(Ops.LABEL, None, [], f"$skipload_{skipload_branch}"))
+        skipload_branch += 1
+    elif uop == Ops.STORE:
+      if args is None:
+        lang.ins.append(AssemblyInstruction(Ops.ALU, lang.tor[vin[0]], [lang.tor[vin[1]]], UnaryOps.NOOP))
+      else:
+        idx, treg, off = lang.addr_w_offset(args)
+        lang.ins.append(AssemblyInstruction(Ops.STORE, None, [idx, lang.tor[vin[0]]] + ([treg] if treg is not None else []), (off, 'global' if not args.local else 'shared', args.memory_dtype if args.memory_dtype != dtypes.float else None)))
+
+  if DEBUG >= 4:
+    for tins in lang.ins: print(tins)
+  return global_size, local_size
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	30ff87eab4	realize sched	2025-10-14 14:44:56 +08:00
George HotzandGitHub	fe683bafa6	Merge branch 'master' into outerworld_work	2025-10-14 14:26:52 +08:00
geohot	ab9064c411	train loop	2025-10-10 20:20:19 +08:00
George HotzandGitHub	8832f08af3	Merge branch 'master' into outerworld_work	2025-10-10 20:07:44 +08:00
geohot	402e1cf48f	work	2025-10-10 19:49:24 +08:00
geohot	b2490b6e31	test assign/reduce	2025-10-10 18:50:16 +08:00
George HotzandGitHub	33e8babdd8	Merge branch 'master' into outerworld_work	2025-10-10 18:25:08 +08:00
geohot	67a409343d	work	2025-10-10 18:10:42 +08:00
geohot	5b24999a36	work on outerworld	2025-10-10 14:48:05 +08:00