no long

viz slowness
remove the device when we render
2026-01-21 21:53:23 +09:00 · 2026-01-21 21:50:09 +09:00 · 2026-01-21 21:41:29 +09:00 · 2026-01-21 21:30:04 +09:00 · 2026-01-21 19:11:49 +09:00 · 2026-01-20 20:34:47 -08:00
577 changed files with 197648 additions and 240330 deletions
@@ -0,0 +1,3 @@
+[run]
+source = tinygrad
+branch = True
@@ -61,7 +61,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 }}-${{ hashFiles('**/setup.py') }}-${{ env.PYTHON_CACHE_VERSION }}
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}

    # **** Caching downloads ****

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

    # **** Python deps ****

@@ -187,7 +187,7 @@ runs:
      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.APT_CACHE_VERSION }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.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 +221,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 -L -o /usr/local/lib/libamd_comgr.dylib
+          sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
        cargo build --release --manifest-path ./extra/remu/Cargo.toml

    # **** gpuocelot ****
@@ -247,7 +247,7 @@ runs:
        cache-name: cache-gpuocelot-build-1
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
    - name: Clone/compile gpuocelot
      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
      shell: bash
@@ -278,7 +278,7 @@ runs:
      if: inputs.webgpu == 'true' && runner.os == 'Linux'
      shell: bash
      run: |
-        sudo curl -L https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/local/lib/libwebgpu_dawn.so
+        sudo curl -fL 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'
@@ -298,7 +298,7 @@ runs:
    - name: Install mesa (linux)
      if: inputs.mesa == 'true' && runner.os == 'Linux'
      shell: bash
-      run: sudo curl -L https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
+      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
@@ -1,10 +1,7 @@
 name: Autogen
 env:
  # increment this when downloads substantially change to avoid the internet
-  DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '4'
-  APT_CACHE_VERSION: '1'
-  BUILD_CACHE_VERSION: '1'
+  CACHE_VERSION: '13'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
@@ -14,15 +11,19 @@ on:
    branches:
      - master
  pull_request:
-    paths: 
+    paths:
    - 'tinygrad/runtime/autogen/**/*'
+    - 'tinygrad/runtime/support/autogen.py'
+    - '.github/workflows/autogen.yml'
  workflow_dispatch:
-    paths: 
+    paths:
    - 'tinygrad/runtime/autogen/**/*'
+    - 'tinygrad/runtime/support/autogen.py'
+    - '.github/workflows/autogen.yml'

 jobs:
  autogen:
-    name: Autogen
+    name: In-tree Autogen
    runs-on: ubuntu-24.04
    timeout-minutes: 15
    steps:
@@ -34,64 +35,102 @@ jobs:
        opencl: 'true'
        amd: 'true'
        cuda: 'true'
-        webgpu: 'true'
        llvm: 'true'
+        webgpu: 'true'
+        mesa: 'true'
        pydeps: 'pyyaml mako'
    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev llvm-20-dev
-    - name: Verify OpenCL autogen
+      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
+    - name: Regenerate autogen files
      run: |
-        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
+        find tinygrad/runtime/autogen -type f -name "*.py" -not -name "__init__.py" -not -name "comgr_3.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
+        python3 -c "from tinygrad.runtime.autogen import opencl"
+        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
+        python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd"
+        python3 -c "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_v13_0_6, smu_v14_0_2"
+        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
+        python3 -c "from tinygrad.runtime.autogen import llvm"
+        python3 -c "from tinygrad.runtime.autogen import webgpu"
+        python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
+        python3 -c "from tinygrad.runtime.autogen import libusb"
+        python3 -c "from tinygrad.runtime.autogen import mesa"
+        python3 -c "from tinygrad.runtime.autogen import avcodec"
+        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
+    - name: Check for differences
      run: |
-        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/nv_gpu.py.bak tinygrad/runtime/autogen/nv_gpu.py
-    - name: Verify AMD autogen
+        if ! git diff --quiet; then
+          git diff > autogen-ubuntu.patch
+          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          exit 1
+        fi
+    - name: Upload patch artifact
+      if: failure()
+      uses: actions/upload-artifact@v4
+      with:
+        name: autogen-ubuntu-patch
+        path: autogen-ubuntu.patch
+
+  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: Regenerate autogen files
      run: |
-        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/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
-    - name: Verify Linux autogen
+        rm tinygrad/runtime/autogen/metal.py tinygrad/runtime/autogen/iokit.py tinygrad/runtime/autogen/corefoundation.py
+        LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
+    - name: Check for differences
      run: |
-        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/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
-        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
-    - name: Verify WebGPU autogen
+        if ! git diff --quiet; then
+          git diff > autogen-macos.patch
+          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          exit 1
+        fi
+    - name: Upload patch artifact
+      if: failure()
+      uses: actions/upload-artifact@v4
+      with:
+        name: autogen-macos-patch
+        path: autogen-macos.patch
+
+  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: |
-        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 LLVM autogen
+        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: Regenerate autogen files
      run: |
-        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
-        ./autogen_stubs.sh llvm
-        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
-    - name: Verify mesa autogen
+        rm tinygrad/runtime/autogen/comgr_3.py
+        python3 -c "from tinygrad.runtime.autogen import comgr_3"
+    - name: Check for differences
      run: |
-        cp tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
-        ./autogen_stubs.sh mesa
-        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
+        if ! git diff --quiet; then
+          git diff > autogen-comgr3.patch
+          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          exit 1
+        fi
+    - name: Upload patch artifact
+      if: failure()
+      uses: actions/upload-artifact@v4
+      with:
+        name: autogen-comgr3-patch
+        path: autogen-comgr3.patch
@@ -14,12 +14,6 @@ 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:
@@ -39,6 +33,7 @@ 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
@@ -54,19 +49,19 @@ jobs:
    - 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=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
    - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 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=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing
    - name: Run Stable Diffusion v2
      # TODO: very slow step time
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing
    # process replay can't capture this, the graph is too large
    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing
    - name: Run model inference benchmark
-      run: METAL=1 python3.11 test/external/external_model_benchmark.py
+      run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
-      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
+      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py
    - name: Test tensor cores
      run: METAL=1 python3.11 test/opt/test_tensor_cores.py
    - name: Test AMX tensor cores
@@ -76,54 +71,80 @@ jobs:
        DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
        DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
    - name: Run Tensor Core GEMM (float)
-      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (half)
-      run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (bfloat16)
-      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py
    - name: Fuzz Padded Tensor Core GEMM
      run: METAL=1 M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3.11 ./extra/gemm/fuzz_matmul.py
    - name: Run LLaMA
      run: |
-        BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
-        BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run LLaMA with BEAM
-      run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run quantized LLaMA
      run: |
-        BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8 | tee llama_int8.txt
-        BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4 | tee llama_nf4.txt
+        BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8
+        BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4
    - name: Run quantized LLaMA3
      run: |
-        BENCHMARK_LOG=llama3_int8 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8 | tee llama3_int8.txt
-        BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4 | tee llama3_nf4.txt
+        BENCHMARK_LOG=llama3_int8 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8
+        BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4
    #- name: Run LLaMA 7B on 4 (virtual) GPUs
-    #  run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    #  run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run GPT2
      run: |
-        BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF/BEAM
-      run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+      run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
    - name: Run OLMoE
      run: BENCHMARK_LOG=olmoe python3.11 examples/olmoe.py
    - name: Train MNIST
-      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
+      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py

    # NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
    # the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
    #- name: Run 10 CIFAR training steps
-    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
+    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py
    #- name: Run 10 CIFAR training steps w HALF
-    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
+    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py

    #- name: Run 10 CIFAR training steps w BF16
-    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py
+    - uses: actions/upload-artifact@v4
+      with:
+        name: Speed (Mac)
+        path: |
+          onnx_inference_speed.csv
+    - 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
@@ -132,38 +153,10 @@ jobs:
      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
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (Mac)
-        path: |
-          onnx_inference_speed.csv
-          torch_speed.txt
-          llama_unjitted.txt
-          llama_jitted.txt
-          llama_beam.txt
-          llama_int8.txt
-          llama_nf4.txt
-          llama3_int8.txt
-          llama3_nf4.txt
-          llama_four_gpu.txt
-          gpt2_unjitted.txt
-          gpt2_jitted.txt
-          gpt2_half.txt
-          gpt2_half_beam.txt
-          matmul.txt
-          matmul_half.txt
-          matmul_bfloat16.txt
-          sd.txt
-          sd_no_fp16.txt
-          sdv2.txt
-          sdxl.txt
-          beautiful_mnist.txt
-          train_cifar.txt
-          train_cifar_half.txt
-          train_cifar_bf16.txt
-          train_cifar_wino.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.11 process_replay.py
+    - 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
@@ -197,9 +190,9 @@ jobs:
    - name: Run model inference benchmark
      run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
    - 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
+      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
    - name: Test speed vs theoretical
-      run: NV=1 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      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
    - name: Test benchmark allreduce
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
@@ -208,79 +201,58 @@ jobs:
        NV=1 NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
    - name: Run Tensor Core GEMM (CUDA)
      run: |
-        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
+        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - 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
+      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (NV)
-      run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
+      run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Test NV=1
      run: DEBUG=2 NV=1 python -m pytest -rA test/test_tiny.py
    - name: Test CUDA=1
      run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
    # TODO: too slow
    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
    - name: Run LLaMA
      run: |
-        BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
-        BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run LLaMA with BEAM
-      run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    # - name: Run LLaMA 7B on 4 GPUs
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
    # - name: Run LLaMA 7B on 6 GPUs
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run LLaMA-3 8B BEAM
-      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
+      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
    - 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
+      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
    - 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
+      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8
    # - 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
+    #   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
    # - name: Run LLaMA-2 70B
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run Mixtral 8x7B
-      run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
+      run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing
    - name: Run GPT2
      run: |
-        BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF/BEAM
-      run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+      run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA)
        path: |
          onnx_inference_speed.csv
-          torch_speed.txt
-          matmul.txt
-          matmul_bfloat16.txt
-          matmul_tf32.txt
-          matmul_ptx.txt
-          matmul_nv.txt
-          sd.txt
-          sdxl.txt
-          llama_unjitted.txt
-          llama_jitted.txt
-          llama_beam.txt
-          llama3_beam.txt
-          llama3_four_gpu.txt
-          llama3_six_gpu.txt
-          llama3_fp8.txt
-          llama_2_70B.txt
-          mixtral.txt
-          gpt2_unjitted.txt
-          gpt2_jitted.txt
-          gpt2_half.txt
-          gpt2_half_beam.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

@@ -318,44 +290,31 @@ 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 JITBEAM=1 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
+      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 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=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    # TODO: too slow
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 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
+      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
    - 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
+    - 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
    - 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=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)
-        path: |
-          beautiful_mnist.txt
-          train_cifar.txt
-          train_cifar_half.txt
-          train_cifar_bf16.txt
-          train_cifar_wino.txt
-          train_cifar_one_gpu.txt
-          train_cifar_six_gpu.txt
-          train_resnet.txt
-          train_resnet_one_gpu.txt
-          train_bert.txt
+      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
    - 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

@@ -370,10 +329,12 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
-    - name: Remove amdgpu
-      run: sudo rmmod amdgpu || true
-    - name: Cleanup running AM processes
-      run: python extra/amdpci/am_smi.py --pids --kill
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
+    - name: Remove amd modules
+      run: ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: ./extra/hcq/hcq_smi.py amd kill_pids
    #- name: Insert amdgpu
    #  run: sudo modprobe amdgpu
    - name: Symlink models and datasets
@@ -407,16 +368,18 @@ jobs:
    #- name: Test speed vs torch
    #  run: |
    #    python3 -c "import torch; print(torch.__version__)"
-    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
+    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
    - name: Test speed vs theoretical
-      run: AMD=1 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_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
+      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 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
+      run: |
+        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
    - name: Test AMD=1
      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
    #- name: Test HIP=1
@@ -431,62 +394,39 @@ jobs:
    - name: Test AM warm start time
      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
-    # 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
+      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
+    - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing
    - 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
-        BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run LLaMA 7B with BEAM
-      run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    # - name: Run LLaMA 7B on 4 GPUs
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
    # - name: Run LLaMA 7B on 6 GPUs
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run LLaMA-3 8B BEAM
-      run: BENCHMARK_LOG=llama3_beam AMD=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
+      run: BENCHMARK_LOG=llama3_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu AMD=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
+      run: BENCHMARK_LOG=llama3_beam_4gpu AMD=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
    # - name: Run LLaMA-3 8B on 6 GPUs
-    #   run: AMD=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
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    #- name: Restore amdgpu
    #  run: sudo modprobe amdgpu
    # - name: Run LLaMA-2 70B
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run Mixtral 8x7B
-      run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
+      run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing
    - name: Run GPT2
      run: |
-        BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF/BEAM
-      run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (AMD)
-        path: |
-          onnx_inference_speed.csv
-          torch_speed.txt
-          llama_unjitted.txt
-          llama_jitted.txt
-          llama_beam.txt
-          llama3_beam.txt
-          llama3_four_gpu.txt
-          llama3_six_gpu.txt
-          llama_2_70B.txt
-          gpt2_unjitted.txt
-          gpt2_jitted.txt
-          gpt2_half.txt
-          gpt2_half_beam.txt
-          matmul.txt
-          matmul_amd.txt
-          sd.txt
-          sdxl.txt
-          mixtral.txt
+      run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
    - 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

@@ -501,10 +441,12 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
-    - name: Remove amdgpu
-      run: sudo rmmod amdgpu || true
-    - name: Cleanup running AM processes
-      run: python extra/amdpci/am_smi.py --pids --kill
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
+    - name: Remove amd modules
+      run: ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -523,34 +465,20 @@ jobs:
    - name: reset process replay
      run: test/external/process_replay/reset.py
    - name: Train MNIST
-      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
+      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=350 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=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    # - name: Run 10 CIFAR training steps w BF16
-    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=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=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    #   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
    - 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 (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)
-        path: |
-          beautiful_mnist.txt
-          train_cifar.txt
-          train_cifar_half.txt
-          train_cifar_bf16.txt
-          train_cifar_wino.txt
-          train_cifar_one_gpu.txt
-          train_cifar_six_gpu.txt
-          train_cifar_six_gpu_remote.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run full CIFAR training steps w 6 GPUS
+      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - 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

@@ -565,10 +493,12 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
-    - name: Remove amdgpu
-      run: sudo rmmod amdgpu || true
-    - name: Cleanup running AM processes
-      run: python extra/amdpci/am_smi.py --pids --kill
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
+    - name: Remove amd modules
+      run: ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -588,20 +518,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
+    - 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
    - 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=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)
-        path: |
-          train_resnet.txt
-          train_resnet_one_gpu.txt
-          train_bert.txt
+      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
    - 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

@@ -623,24 +546,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.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 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: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 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: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 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 0.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 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
+      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 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=13 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
+      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: IMAGE=1 openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=image_1_openpilot_0_10_1_vision PYTHONPATH="." 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
-      # TODO: ASSERT_MIN_STEP_TIME=17
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=25 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
+      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=5 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
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 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
-      # TODO: ASSERT_MIN_STEP_TIME=10
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 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
+      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 MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -663,6 +584,8 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove amd modules
      run: ./extra/hcq/hcq_smi.py amd rmmod
    - name: Kill stale pids
@@ -695,7 +618,7 @@ jobs:
    #     AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
    #     AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - 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 am_matmul_amd.txt
+      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
    - name: Test AMD=1
      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
    - name: Test DISK copy time
@@ -705,21 +628,12 @@ jobs:
        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
        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
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #   run: BENCHMARK_LOG=resnet_10steps 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
+    #   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
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps AMD=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 am_train_bert_one_gpu.txt
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (AM Driver)
-        path: |
-          am_matmul_amd.txt
-          am_train_cifar_one_gpu.txt
-          am_train_resnet_one_gpu.txt
-          am_train_bert_one_gpu.txt
+      run: BENCHMARK_LOG=bert_10steps AMD=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
    - 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

@@ -734,6 +648,8 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove nv modules
      run: ./extra/hcq/hcq_smi.py nv rmmod
    - name: Kill stale pids
@@ -766,21 +682,13 @@ jobs:
        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Test LLAMA-3
-      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
+      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar 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
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps 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
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (NV Driver)
-        path: |
-          nv_llama3_beam.txt
-          nv_train_cifar_one_gpu.txt
-          nv_train_resnet_one_gpu.txt
-          nv_train_bert_one_gpu.txt
+      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
    - 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
@@ -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 DISABLE_COMPILER_CACHE=1 python examples/sdxl.py --noshow --timing --seed 0
+        BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CCACHE=0 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 DISABLE_COMPILER_CACHE=1 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 CCACHE=0 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
@@ -20,11 +20,11 @@ jobs:
    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
-        pip install setuptools wheel twine
+        pip install setuptools wheel build twine
    - name: Build and publish
      env:
        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
      run: |
-        python setup.py sdist bdist_wheel
+        python -m build
        twine upload dist/*
@@ -56,15 +56,15 @@ jobs:
        uses: actions/checkout@v4
        with:
          path: base
-      - name: Set up Python 3.10
+      - name: Set up Python 3.12
        uses: actions/setup-python@v5
        with:
-          python-version: '3.10'
+          python-version: '3.12'
      - 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,13 +1,11 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '4'
-  APT_CACHE_VERSION: '1'
-  BUILD_CACHE_VERSION: '1'
+  CACHE_VERSION: '15'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
+  IGNORE_OOB: 0

 on:
  push:
@@ -39,6 +37,8 @@ jobs:
    name: Docs
    runs-on: ubuntu-22.04
    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 1
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -74,9 +74,7 @@ jobs:
    - name: Test Docs Build
      run: python -m mkdocs build --strict
    - name: Test Docs
-      run: |
-        python docs/abstractions2.py
-        python docs/abstractions3.py
+      run: 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
@@ -89,65 +87,63 @@ jobs:
        clang -O2 recognize.c -lm -o recognize
        cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock

-  # TODO: fix the torch backend and reenable
-  # torchbackend:
-  #   name: Torch Backend Tests
-  #   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: torch-backend-pillow-torchvision-et-pt
-  #       deps: testing_minimal
-  #       pydeps: "pillow torchvision expecttest"
-  #       llvm: 'true'
-  #   - name: Install ninja
-  #     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: 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
-  #   - name: Test Ops with TINY_BACKEND
-  #     run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
-  #   - name: Test in-place operations on views
-  #     run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
-  #   - name: Test multi-gpu
-  #     run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+  torchbackend:
+    name: Torch Backend Tests
+    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: torch-backend-pillow-torchvision-et-pt
+        deps: testing_minimal
+        pydeps: "pillow torchvision expecttest"
+        llvm: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - name: Test one op
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+    - name: Test ResNet-18
+      run: DEBUG=2 python3 extra/torch_backend/example.py
+    - name: custom tests
+      run: python3 -m pytest -n auto extra/torch_backend/test.py --durations=20
+    - name: Test one op in torch tests
+      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+    - name: Test Ops with TINY_BACKEND
+      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+    - name: Test in-place operations on views
+      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+    - name: Test multi-gpu
+      run: 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
-  #   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: torch-backend-pillow-torchvision-et-pt
-  #       deps: testing_minimal
-  #       llvm: 'true'
-  #   - name: Install ninja
-  #     run: |
-  #       sudo apt update || true
-  #       sudo apt install -y --no-install-recommends ninja-build
-  #   - name: Test beautiful_mnist in torch with TINY_BACKEND
-  #     run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.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
+
+  torchbackendmore:
+    name: Torch Backend Tests More
+    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: torch-backend-pillow-torchvision-et-pt
+        deps: testing_minimal
+        llvm: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - name: Test beautiful_mnist in torch with TINY_BACKEND
+      run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.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

  bepython:
    name: Python Backend
@@ -222,7 +218,6 @@ jobs:
    runs-on: ubuntu-latest
    timeout-minutes: 10

-    # TODO: run the pre-commit hook to replace a lot of this
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -234,18 +229,19 @@ jobs:
        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
+    - name: Run pre-commit linting hooks
+      run: SKIP=tiny,tests,example pre-commit run --all-files
+    - name: Lint additional files with ruff
      run: |
-        pip3 install --upgrade --force-reinstall ruff==0.11.0
-        python3 -m ruff check .
        python3 -m ruff check examples/mlperf/ --ignore E501
-    - name: Run mypy
+        python3 -m ruff check extra/thunder/tiny/ --ignore E501 --ignore F841 --ignore E722
+        python3 -m ruff check extra/torch_backend/backend.py
+    - name: Run mypy with lineprecision report
      run: |
-        python -m mypy --strict-equality --lineprecision-report .
+        python -m mypy --lineprecision-report .
        cat lineprecision.txt
-    # broken because of UPatAny
-    #- name: Run TYPED=1
-    #  run: TYPED=1 python -c "import tinygrad"
+    - name: Run TYPED=1
+      run: TYPED=1 python -c "import tinygrad"

  unittest:
    name: Unit Tests
@@ -258,13 +254,18 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: unittest-12
-        pydeps: "pillow numpy ftfy regex"
+        key: unittest-13
+        pydeps: "pillow numpy ftfy regex pre-commit"
        deps: testing_unit
+        llvm: 'true'
+    - name: Run pre-commit test hooks
+      run: SKIP=ruff,mypy pre-commit run --all-files
    - 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 -m pytest -n=auto test/unit/ --durations=20
+      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
    - 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
@@ -290,8 +291,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 < 18000 lines
-      run: MAX_LINE_COUNT=18000 python sz.py
+    - name: Repo line count < 20000 lines
+      run: MAX_LINE_COUNT=20000 python sz.py

  spec:
    strategy:
@@ -309,8 +310,9 @@ jobs:
      with:
        key: spec-unit
        deps: testing_unit
+        python-version: '3.14'
    - name: Test SPEC=2
-      run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+      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 --ignore test/unit/test_autogen.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}

  fuzzing:
    name: Fuzzing
@@ -326,6 +328,8 @@ 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 shape ops
@@ -344,10 +348,11 @@ jobs:
          key: gpu-image
          deps: testing_minimal
          opencl: 'true'
-      - name: Test CL IMAGE=2 ops + training
+      - name: Test CL IMAGE=2 ops
        run: |
          CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
-          CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
+        # TODO: training is broken
+        # CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -392,7 +397,7 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1452 ALLOWED_GATED_READ_IMAGE=122 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          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)
@@ -444,7 +449,7 @@ jobs:
        with:
          key: onnxoptl
          deps: testing
-          pydeps: "tensorflow==2.15.1 tensorflow_addons"
+          pydeps: "tensorflow==2.19"
          python-version: '3.11'
          opencl: 'true'
      - name: Test ONNX (CL)
@@ -462,7 +467,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=8 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=1 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

@@ -470,6 +475,8 @@ jobs:
    name: Test LLM
    runs-on: ubuntu-24.04
    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 1
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -600,9 +607,7 @@ jobs:
        WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
    - name: Run selected webgpu tests
      run: |
-          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit \
-          --ignore=test/test_copy_speed.py --ignore=test/test_rearrange_einops.py \
-          --ignore=test/test_fuzz_shape_ops.py --durations=20
+          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -638,7 +643,7 @@ 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 --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 test/testextra/test_cfg_viz.py --durations=20
      - name: Run pytest (amd)
        run: python -m pytest test/external/external_test_am.py --durations=20
      - name: Run TRANSCENDENTAL math
@@ -651,6 +656,46 @@ jobs:
      - 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'
+          python-version: '3.13'
+      - name: Verify AMD autogen is up to date
+        run: |
+          python -m extra.assembly.amd.amdxml
+          git diff --exit-code extra/assembly/amd/autogen/
+      - 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: Install rocprof-trace-decoder
+        run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.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
+      # TODO: run all once emulator is faster
+      - name: Run RDNA3 ops tests
+        run: SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril"
+
  testnvidia:
    strategy:
      fail-fast: false
@@ -718,71 +763,6 @@ 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:
@@ -803,6 +783,8 @@ jobs:
        ocelot: 'true'
        llvm: 'true'
    - name: Run unit tests
+      env:
+        LIBCLANG_PATH: '/opt/homebrew/opt/llvm@20/lib/libclang.dylib'
      run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
    - name: Run ONNX
      run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -880,30 +862,6 @@ 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
@@ -969,3 +927,33 @@ 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
@@ -58,8 +58,11 @@ weights
 *.lprof
 comgr_*
 *.pkl
+!extra/sqtt/examples/**/*.pkl
 site/
 profile_stats
 *.log
 target
 .mypy_cache
+mutants
+.mutmut-cache
@@ -16,7 +16,7 @@ repos:
        pass_filenames: false
      - id: mypy
        name: mypy
-        entry: python3 -m mypy tinygrad/ --strict-equality
+        entry: python3 -m mypy
        language: system
        always_run: true
        pass_filenames: false
@@ -27,8 +27,8 @@ repos:
        always_run: true
        pass_filenames: false
      - id: tests
-        name: subset of tests
-        entry: env OMP_NUM_THREADS=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
+        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/unit/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
        language: system
        always_run: true
        pass_filenames: false
@@ -0,0 +1,227 @@
+# 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 Analysis
+
+**Use the right tool:**
+
+- `TRACK_MATCH_STATS=2` - **Profiling**: identify expensive patterns
+- `VIZ=-1` - **Inspection**: see all transformations, what every match pattern does, the before/after diffs
+
+```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.
+
+```bash
+# Save the trace
+VIZ=-1 python test/test_tiny.py TestTiny.test_gemm
+
+# Explore it
+./extra/viz/cli.py --help
+```
+
+## 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,17 +21,38 @@ tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) an

 ---

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

-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.
+- **Tensor library** with autograd
+- **IR and compiler** that fuse and lower kernels
+- **JIT + graph execution**
+- **nn / optim / datasets** for real training

-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.
+It’s inspired by PyTorch (ergonomics), JAX (functional transforms and IR-based AD), and TVM (scheduling and codegen), but stays intentionally tiny and hackable.

-## Features
+---

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

-tinygrad can run [LLaMA](/docs/showcase.md#llama) and [Stable Diffusion](/docs/showcase.md#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.
+
+---

 ### Laziness

@@ -1,568 +0,0 @@
-#!/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:
-    try: return ctypes.CDLL(library)
-    except OSError: pass
-  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 $BASE/rocprof.py
-  fixup $BASE/rocprof.py
-  sed -i '1s/^/# pylint: skip-file\n/' $BASE/rocprof.py
-  sed -i "s/import ctypes/import ctypes, ctypes.util/g" $BASE/rocprof.py
-  patch_dlopen $BASE/rocprof.py rocprof-trace-decoder "'/usr/local/lib/librocprof-trace-decoder.so'" "'/usr/local/lib/librocprof-trace-decoder.dylib'"
-  sed -i "s/def _try_dlopen_rocprof-trace-decoder():/def _try_dlopen_rocprof_trace_decoder():/g" $BASE/rocprof.py
-  sed -i "s|FunctionFactoryStub()|_try_dlopen_rocprof_trace_decoder()|g" $BASE/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"
-}
-
-generate_mesa() {
-  MESA_TAG="mesa-25.2.4"
-  MESA_SRC=/tmp/mesa-$MESA_TAG
-  TINYMESA_TAG=tinymesa-32dc66c
-  TINYMESA_DIR=/tmp/tinymesa-$MESA_TAG-$TINYMESA_TAG/
-  TINYMESA_SO=$TINYMESA_DIR/libtinymesa_cpu.so
-  if [ ! -d "$MESA_SRC" ]; then
-    git clone --depth 1 --branch $MESA_TAG https://gitlab.freedesktop.org/mesa/mesa.git $MESA_SRC
-    pushd .
-    cd $MESA_SRC
-    git reset --hard $MESA_COMMIT_HASH
-    # clang 14 doesn't support packed enums
-    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/headers/nv_device_info.h
-    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/compiler/nak.h
-    sed -i "s/nir_instr_type \(\w\+\);/uint8_t \1;/" $MESA_SRC/src/compiler/nir/nir.h
-    mkdir -p gen/util/format
-    python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
-    python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
-    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
-    python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
-    python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
-    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
-    python3 src/compiler/builtin_types_h.py gen/builtin_types.h
-    popd
-  fi
-
-  if [ ! -d "$TINYMESA_DIR" ]; then
-    mkdir $TINYMESA_DIR
-    curl -L https://github.com/sirhcm/tinymesa/releases/download/$TINYMESA_TAG/libtinymesa_cpu-$MESA_TAG-linux-amd64.so -o $TINYMESA_SO
-  fi
-
-  clang2py -k cdefstu \
-    $MESA_SRC/src/compiler/nir/nir.h \
-    $MESA_SRC/src/compiler/nir/nir_builder.h \
-    $MESA_SRC/src/compiler/nir/nir_shader_compiler_options.h \
-    $MESA_SRC/src/compiler/nir/nir_serialize.h \
-    $MESA_SRC/gen/nir_intrinsics.h \
-    $MESA_SRC/src/nouveau/headers/nv_device_info.h \
-    $MESA_SRC/src/nouveau/compiler/nak.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_passmgr.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_misc.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_type.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_init.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_nir.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_struct.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_jit_types.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_flow.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_const.h \
-    $MESA_SRC/src/compiler/glsl_types.h \
-    $MESA_SRC/src/util/blob.h \
-    $MESA_SRC/src/util/ralloc.h \
-    --clang-args="-DHAVE_ENDIAN_H -DHAVE_STRUCT_TIMESPEC -DHAVE_PTHREAD -I$MESA_SRC/src -I$MESA_SRC/include -I$MESA_SRC/gen -I$MESA_SRC/src/compiler/nir -I$MESA_SRC/src/gallium/auxiliary -I$MESA_SRC/src/gallium/include -I$(llvm-config-20 --includedir)" \
-    -l $TINYMESA_SO \
-    -o $BASE/mesa.py
-
-  LVP_NIR_OPTIONS=$(./extra/mesa/lvp_nir_options.sh $MESA_SRC)
-
-  fixup $BASE/mesa.py
-  patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "'/opt/homebrew/lib/libtinymesa_cpu.dylib'" "'/opt/homebrew/lib/libtinymesa.dylib'"
-  echo "lvp_nir_options = gzip.decompress(base64.b64decode('$LVP_NIR_OPTIONS'))" >> $BASE/mesa.py
-  sed -i "/in_dll/s/.*/try: &\nexcept (AttributeError, ValueError): pass/" $BASE/mesa.py
-  sed -i "s/import ctypes/import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers/" $BASE/mesa.py
-  sed -i "s/ctypes.CDLL('.\+')/(dll := _try_dlopen_tinymesa_cpu())/" $BASE/mesa.py
-  echo "def __getattr__(nm): raise AttributeError('LLVMpipe requires tinymesa_cpu' if 'tinymesa_cpu' not in dll._name else f'attribute {nm} not found') if dll else FileNotFoundError(f'libtinymesa not found (MESA_PATH={BASE}). See https://github.com/sirhcm/tinymesa ($TINYMESA_TAG, $MESA_TAG)')" >> $BASE/mesa.py
-  sed -i "s/ctypes.glsl_base_type/glsl_base_type/" $BASE/mesa.py
-  # bitfield bug in clang2py
-  sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
-  sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
-  sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
-  sed -i '/struct_nir_builder._pack_ = 1 # source:False/d' "$BASE/mesa.py"
-  python3 -c "import tinygrad.runtime.autogen.mesa"
-}
-
-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" == "mesa" ]; then generate_mesa
-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; generate_mesa
-else echo "usage: $0 <type>"
-fi
@@ -1,135 +0,0 @@
-# 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)
-alu = buf_1.index(idx) + buf_2.index(idx)
-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,25 +38,19 @@ 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 ScheduleItem
-schedule: List[ScheduleItem] = Tensor.schedule(l1, l2)
+from tinygrad.engine.schedule import ExecItem
+schedule: List[ExecItem] = Tensor.schedule(l1, l2)

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

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

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

 # *****
-# 5. Run the schedule
-
-for ei in tqdm(lowered): ei.run()
-
-# *****
-# 6. Print the weight change
+# 5. 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 (specified by a ShapeTracker). 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. 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 `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.
+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.

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

 ## Lowering

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

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

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

@@ -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 list of UOps.
+Transform the optimized ast into a linearized and rendered program.

-::: tinygrad.codegen.full_rewrite
+::: tinygrad.codegen.get_program
    options:
        members: false
        show_labels: false
@@ -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 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.
+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.

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

@@ -1,293 +0,0 @@
-#!/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
@@ -70,7 +70,7 @@ AMD backend supports several interfaces for communicating with devices:

 * `KFD`: uses the amdgpu driver
 * `PCI`: uses the [AM driver](developer/am.md)
-* `USB`: USB3 interafce for asm24xx chips.
+* `USB`: USB3 interface for asm24xx chips.

 You can force an interface by setting `AMD_IFACE` to one of these values. In the case of `AMD_IFACE=PCI`, this may unbind your GPU from the amdgpu driver.

@@ -6,6 +6,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.neg
 ::: tinygrad.Tensor.log
 ::: tinygrad.Tensor.log2
+::: tinygrad.Tensor.log10
 ::: tinygrad.Tensor.exp
 ::: tinygrad.Tensor.exp2
 ::: tinygrad.Tensor.sqrt
@@ -87,4 +88,8 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.float
 ::: tinygrad.Tensor.half
 ::: tinygrad.Tensor.int
-::: tinygrad.Tensor.bool
+::: tinygrad.Tensor.bool
+::: tinygrad.Tensor.bfloat16
+::: tinygrad.Tensor.double
+::: tinygrad.Tensor.long
+::: tinygrad.Tensor.short
@@ -27,5 +27,6 @@
 ::: tinygrad.Tensor.flatten
 ::: tinygrad.Tensor.unflatten
 ::: tinygrad.Tensor.diag
+::: tinygrad.Tensor.diagonal
 ::: tinygrad.Tensor.roll
 ::: tinygrad.Tensor.rearrange
@@ -7,6 +7,7 @@
 ::: tinygrad.Tensor.any
 ::: tinygrad.Tensor.all
 ::: tinygrad.Tensor.isclose
+::: tinygrad.Tensor.allclose
 ::: tinygrad.Tensor.mean
 ::: tinygrad.Tensor.var
 ::: tinygrad.Tensor.var_mean
@@ -30,7 +31,9 @@
 ::: tinygrad.Tensor.matmul
 ::: tinygrad.Tensor.einsum
 ::: tinygrad.Tensor.cumsum
+::: tinygrad.Tensor.cumprod
 ::: tinygrad.Tensor.cummax
+::: tinygrad.Tensor.cummin
 ::: tinygrad.Tensor.triu
 ::: tinygrad.Tensor.tril
 ::: tinygrad.Tensor.interpolate
@@ -38,7 +41,9 @@
 ::: tinygrad.Tensor.scatter_reduce
 ::: tinygrad.Tensor.masked_select
 ::: tinygrad.Tensor.masked_fill
+::: tinygrad.Tensor.nonzero
 ::: tinygrad.Tensor.sort
+::: tinygrad.Tensor.argsort
 ::: tinygrad.Tensor.topk
 ::: tinygrad.Tensor.multinomial

@@ -56,3 +61,8 @@
 ::: tinygrad.Tensor.sparse_categorical_crossentropy
 ::: tinygrad.Tensor.cross_entropy
 ::: tinygrad.Tensor.nll_loss
+
+## Linear Algebra
+
+::: tinygrad.Tensor.qr
+::: tinygrad.Tensor.svd
@@ -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 `displayservice.service` service.
+Reboot after making these changes or restart the `tinybox-display.service` service.

 ## What do I use it for?

@@ -1,9 +0,0 @@
-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,
-];
@@ -0,0 +1,79 @@
+from typing import Optional
+from tinygrad import Tensor
+from tinygrad.dtype import DTypeLike, dtypes
+import math
+
+# rewritten from numpy
+def rfftfreq(n: int, d: float = 1.0, device=None) -> Tensor:
+  val = 1.0 / (n * d)
+  N = n // 2 + 1
+  results = Tensor.arange(N, device=device)
+  return results * val
+
+# just like in librosa
+def fft_frequencies(sr: float, n_fft: int) -> Tensor:
+  return rfftfreq(n=n_fft, d=1.0 / sr)
+
+def hz_to_mel(freq: Tensor) -> Tensor:
+  # linear part
+  f_min = 0.0
+  f_sp = 200.0 / 3
+  mels = (freq - f_min) / f_sp
+
+  # log-scale part
+  min_log_hz = 1000.0  # beginning of log region (Hz)
+  mask = freq >= min_log_hz
+  return mask.where(((min_log_hz - f_min) / f_sp) + (freq / min_log_hz).log() / (math.log(6.4) / 27.0), mels)
+
+def mel_to_hz(mels: Tensor) -> Tensor:
+  # linear scale
+  f_min = 0.0
+  f_sp = 200.0 / 3
+  freqs = f_min + f_sp * mels
+
+  # nonlinear scale
+  min_log_hz = 1000.0  # beginning of log region (Hz)
+  min_log_mel = (min_log_hz - f_min) / f_sp  # same (Mels)
+  logstep = math.log(6.4) / 27.0  # step size for log region
+
+  log_t = mels >= min_log_mel
+  freqs = log_t.where(min_log_hz * ((logstep * (mels - min_log_mel)).exp()), freqs)
+  return freqs
+
+def mel_frequencies(n_mels: int = 128, *, fmin: float = 0.0, fmax: float = 11025.0) -> Tensor:
+  # center freqs of mel bands - uniformly spaced between limits
+  min_max_mel = hz_to_mel(Tensor([fmin, fmax]))
+
+  mels = Tensor.linspace(min_max_mel[0], min_max_mel[1], n_mels)
+  hz = mel_to_hz(mels)
+  return hz
+
+def mel(
+  *,
+  sr: float,
+  n_fft: int,
+  n_mels: int = 128,
+  fmin: float = 0.0,
+  fmax: Optional[float] = None,
+  dtype: DTypeLike = dtypes.default_float,
+) -> Tensor:
+  if fmax is None:
+    fmax = float(sr) / 2
+
+  n_mels = int(n_mels)
+
+  fftfreqs = fft_frequencies(sr=sr, n_fft=n_fft)  # center freqs of each FFT bin
+  mel_f = mel_frequencies(n_mels + 2, fmin=fmin, fmax=fmax)  # center freqs of mel bands
+
+  fdiff = mel_f[1:] - mel_f[:-1]
+  ramps = mel_f[None].T.expand(-1, fftfreqs.shape[-1]) - fftfreqs
+
+  lower = -ramps[:n_mels] / fdiff[:n_mels][None].T
+  upper = ramps[2 : n_mels + 2] / fdiff[1 : n_mels + 1][None].T
+  weights = lower.minimum(upper).maximum(0)
+
+  # Slaney-style mel is scaled to be approx constant energy per channel
+  enorm = 2.0 / (mel_f[2 : n_mels + 2] - mel_f[:n_mels])
+  weights *= enorm[:, None]
+
+  return weights
@@ -21,7 +21,7 @@ if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))

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

  @TinyJit
  @Tensor.train()
@@ -1,93 +0,0 @@
-#!/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("")
@@ -1,341 +0,0 @@
-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}")
@@ -1,89 +0,0 @@
-# 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)])
@@ -1,498 +0,0 @@
-# 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"))
@@ -0,0 +1,108 @@
+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,8 +1,6 @@
 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
@@ -12,6 +10,8 @@ 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 = [
@@ -1,299 +0,0 @@
-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()
@@ -213,12 +213,13 @@ class InterleavedDataset:
    self.queues[queue_index].queue.extend(load_file(file))

 # Reference: https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/run_pretraining.py, Line 394
-def batch_load_train_bert(BS:int):
+def batch_load_train_bert(BS:int, seed:int|None=None):
  from extra.datasets.wikipedia import get_wiki_train_files
+  rng = random.Random(seed)
  fs = sorted(get_wiki_train_files())
  train_files = []
  while fs: # TF shuffle
-    random.shuffle(fs)
+    rng.shuffle(fs)
    train_files.append(fs.pop(0))

  cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
@@ -763,48 +764,26 @@ class BlendedGPTDataset:

    return dataset_idx, dataset_sample_idx

-def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
+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)
  if val:
-    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)
+    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)

-  for b in range(math.ceil(samples / bs)):
-    batch = []
-    for i in range(bs):
-      tokens = dataset.get(b * bs + i)
-      batch.append(tokens)
+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)]
    yield Tensor.stack(batch, dim=0)

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

 if __name__ == "__main__":
  def load_unet3d(val):
@@ -219,17 +219,28 @@ def get_mlperf_bert_model():
  config = get_mlperf_bert_config()
  if getenv("DISABLE_DROPOUT", 0):
    config["hidden_dropout_prob"] = config["attention_probs_dropout_prob"] = 0.0
-  return BertForPretraining(**config)
+  model = BertForPretraining(**config)
+  if getenv("FP8_TRAIN"):
+    from extra.fp8.fp8_linear import convert_to_float8_training
+    def module_filter_fn(mod, fqn):
+      if isinstance(mod, LinearBert):
+        skip_layers = [] if (ln:=config["num_hidden_layers"]) <= 2 else ["bert.encoder.layer.0.", f"bert.encoder.layer.{ln-1}"]
+        if mod.weight.shape[-1] >= 1024 and "encoder" in fqn and not any(name in fqn for name in skip_layers):
+          print(f"replacing linear with fp8: {fqn} {mod.weight.shape}")
+          return True
+      return False
+    convert_to_float8_training(model, module_filter_fn)
+  return model

 def get_fake_data_bert(BS:int):
  return {
-    "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"),
+    "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(),
  }

 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,9 +59,7 @@ 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)
-    # 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()
+    return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype)

 class LayerNormBert:
  def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
@@ -204,43 +204,6 @@ 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
@@ -271,12 +234,9 @@ def eval_llama3():
    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
    return loss.flatten().float()

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

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

-  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(",")
+  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
  for m in models:
    nm = f"eval_{m}"
    if nm in globals():
@@ -918,40 +918,6 @@ 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):
@@ -1014,7 +980,8 @@ 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: mlperf logging
+  # TODO: implement grad accumulation + mlperf logging
+  assert grad_acc == 1
  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))
@@ -1041,6 +1008,7 @@ def train_bert():
  config["DISABLE_DROPOUT"] = getenv("DISABLE_DROPOUT", 0)
  config["TRAIN_BEAM"]    = TRAIN_BEAM = getenv("TRAIN_BEAM", BEAM.value)
  config["EVAL_BEAM"]     = EVAL_BEAM  = getenv("EVAL_BEAM", BEAM.value)
+  config["FP8_TRAIN"]     = getenv("FP8_TRAIN", 0)

  Tensor.manual_seed(seed)  # seed for weight initialization

@@ -1073,8 +1041,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 = [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)
+  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)
  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)

@@ -1118,7 +1086,7 @@ def train_bert():
  if RUNMLPERF:
    # only load real data with RUNMLPERF
    eval_it = iter(batch_load_val_bert(EVAL_BS))
-    train_it = iter(tqdm(batch_load_train_bert(BS), total=train_steps, disable=BENCHMARK))
+    train_it = iter(tqdm(batch_load_train_bert(BS, seed=seed), total=train_steps, disable=BENCHMARK))
    for _ in range(start_step): next(train_it) # Fast forward
  else:
    # repeat fake data
@@ -1131,12 +1099,38 @@ def train_bert():
  # ** train loop **
  wc_start = time.perf_counter()

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

  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
@@ -1144,21 +1138,17 @@ def train_bert():
      st = time.perf_counter()
      GlobalCounters.reset()
      with WallTimeEvent(BenchEvent.STEP):
-        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)
+        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"])

        pt = time.perf_counter()
-
-        try:
-          next_data = [next(train_it) for _ in range(grad_acc)]
-        except StopIteration:
-          next_data = None
-
+        next_data = next(train_it)
        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)}"
        loss = loss.item()
-        assert not math.isnan(loss)
+        if not getenv("FP8_TRAIN"): assert not math.isnan(loss)
        lr = lr.item()

      cl = time.perf_counter()
@@ -1171,7 +1161,7 @@ def train_bert():
      if WANDB:
        wandb.log({"lr": lr, "train/loss": loss, "train/global_norm": global_norm.item(), "train/step_time": cl - st,
                    "train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
-                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
+                    "train/mem":GlobalCounters.mem_used / 1e9, "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})

      train_data, next_data = next_data, None
      i += 1
@@ -1188,8 +1178,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", 0) and train_step_bert.captured is not None:
-        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
+      elif getenv("FREE_INTERMEDIATE") and train_step_bert.captured is not None:
+        # TODO: this hangs on tiny green after 90 minutes of training
        train_step_bert.captured.free_intermediates()
      eval_lm_losses = []
      eval_clsf_losses = []
@@ -1224,7 +1214,7 @@ def train_bert():
          return

      if getenv("RESET_STEP"): eval_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE") 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)
@@ -1296,10 +1286,13 @@ def train_llama3():
  from examples.llama3 import MODEL_PARAMS
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup

+  BENCHMARK = getenv("BENCHMARK")
+
  config = {}
  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)
@@ -1324,12 +1317,23 @@ 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)

-  # TODO: confirm weights are in bf16
+  # ** 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"]
  # vocab_size from the mixtral tokenizer
-  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 not SMALL: model_params |= {"vocab_size": 32000}
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
+  print(f"model parameters: {model_params}")
+
+  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)

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

  @TinyJit
  @Tensor.train()
-  def train_step(model, tokens:Tensor, grad_acc:int):
+  def train_step(model, tokens:Tensor):
    optim.zero_grad()
-    # 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])
+    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()
    # 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
@@ -1422,57 +1423,74 @@ def train_llama3():
  # ** data iters **
  def fake_data(bs, samples):
    for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)

  def get_train_iter():
    if getenv("FAKEDATA", 0):
-      return fake_data(GBS, SAMPLES)
+      return fake_data(BS, SAMPLES)
    else:
-      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))
+      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(1024 if SMALL else 5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))

  def get_eval_iter():
-    if getenv("FAKEDATA", 0):
+    if eval_dataset is None:
      return fake_data(EVAL_BS, 5760)
-    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)
+    from examples.mlperf.dataloader import iterate_llama3_dataset
+    return iterate_llama3_dataset(eval_dataset, EVAL_BS)

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

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

-    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")
+      sec = time.perf_counter()-t
+      if BENCHMARK: step_times.append(sec)

-    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)
+      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("saving optim checkpoint")
-      fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
-      safe_save(get_state_dict(scheduler), fn)
+      if (fname:=getenv("LOSS_FILE", "")):
+        with open(fname, "a") as f:
+          f.write(f"{i} {loss:.4f} {lr:.12f} {mem_gb:.2f}\n")

-    if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
+      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)
+
+      if i == BENCHMARK:
+        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]
+        estimated_total_minutes = int(median_step_time * (SAMPLES // GBS) / 60)
+        print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
+        print(f"epoch global_ops: {GlobalCounters.global_ops:_}, "
+              f"epoch global_mem: {GlobalCounters.global_mem:_}")
+
+    if (sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
      tqdm.write(f"evaluating after {sequences_seen} sequences")

      # run eval
@@ -1480,12 +1498,19 @@ def train_llama3():
      eval_iter = get_eval_iter()
      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")

-      for tokens in tqdm(eval_iter, total=5760//EVAL_BS):
+      for j,tokens in tqdm(enumerate(eval_iter), total=5760//EVAL_BS):
        eval_losses += eval_step(model, tokens).tolist()
+        
+        if BENCHMARK and (j+1) == min(BENCHMARK, 5760//EVAL_BS):
+          return
+
      log_perplexity = Tensor(eval_losses).mean().float().item()

      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"):
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1 DEBUG=0 JIT=1 FLASH_ATTENTION=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=5000000
+
+export BEAM=0 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,24 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=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"
+export BEAM_TIMEOUT_SEC=15
+export FP8_TRAIN=1
+# search
+IGNORE_BEAM_CACHE=1 BENCHMARK=10 BERT_LAYERS=2 RUNMLPERF=0 python3 examples/mlperf/model_train.py
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,31 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_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=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72

 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=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72

 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=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72

 export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=500000
@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="." AMD=1
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+
+export DEBUG=${DEBUG:-0}
+export FLASH_ATTENTION=1
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=8 BS=8 EVAL_BS=8
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8b"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_STEPS=1024 MAX_STEPS=1200000
+export SAMPLES=$((MAX_STEPS * BS))
+
+export SEED=5760
+
+export JITBEAM=3
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+
+python3 examples/mlperf/model_train.py
@@ -1,118 +0,0 @@
-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,11 +1,10 @@
 import os, sys, pickle, time, re
 import numpy as np
+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"
@@ -40,7 +39,7 @@ def compile(onnx_file):
  np.testing.assert_equal(test_val, ret, "JIT run failed")
  print("jit run validated")

-  # checks from compile2
+  # check gated read_image usage
  kernel_count = 0
  read_image_count = 0
  gated_read_image_count = 0
@@ -96,6 +95,7 @@ def test_vs_compile(run, inputs, test_val=None):
  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()}
@@ -1,55 +0,0 @@
-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}%")
@@ -1,45 +0,0 @@
-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,
-)
@@ -115,7 +115,7 @@ if __name__ == "__main__":

  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
    if not args.fakeweights:
-      default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
+      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
@@ -1,48 +0,0 @@
-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}")
@@ -1,136 +0,0 @@
-#!/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}')
@@ -1,17 +0,0 @@
-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()
@@ -1,669 +0,0 @@
-# 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}")
@@ -1,204 +0,0 @@
-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
+from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten, profile_marker
 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
@@ -266,13 +266,16 @@ if __name__ == "__main__":
  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()

-  # load in weights
+  profile_marker("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')
-      load_state_dict(model, torch_load(model_bin)['state_dict'], verbose=False, strict=False, realize=False)
+      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)

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

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

-  # run through CLIP to get context
+  profile_marker("run clip (conditional)")
  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)
@@ -310,6 +314,7 @@ 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))
@@ -319,24 +324,26 @@ 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"
-  # upsample latent space to image with autoencoder
-  x = model.decode(latent)
+  profile_marker("run decoder") # upsample latent space to image with autoencoder
+  x = model.decode(latent).realize()
  print(x.shape)

-  # save image
+  profile_marker("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
@@ -7,7 +7,7 @@ if __name__ == "__main__":
  with open(fetch(sys.argv[1]), "rb") as f:
    run_onnx_jit = pickle.load(f)
  input_name = run_onnx_jit.captured.expected_names[0]
-  device = run_onnx_jit.captured.expected_st_vars_dtype_device[0][-1]
+  device = run_onnx_jit.captured.expected_input_info[0][-1]
  print(f"input goes into {input_name=} on {device=}")
  hit = 0
  for i,(img,y) in enumerate(imagenet_dataloader(cnt=getenv("CNT", 100))):
@@ -0,0 +1,34 @@
+#!/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]))
@@ -0,0 +1,16 @@
+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)
@@ -1,104 +0,0 @@
-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
@@ -1,46 +0,0 @@
-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])
@@ -1,740 +0,0 @@
-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,10 +3,11 @@
 import sys, base64, multiprocessing, itertools, collections
 from typing import Optional, Union, Literal, List

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

+from examples.audio_helpers import mel
 import numpy as np
 import librosa

@@ -159,7 +160,7 @@ def prep_audio(waveforms: List[np.ndarray], batch_size: int, truncate=False) ->

  stft = librosa.stft(waveforms, n_fft=N_FFT, hop_length=HOP_LENGTH, window='hann', dtype=np.csingle)
  magnitudes = np.absolute(stft[..., :-1]) ** 2
-  mel_spec = librosa.filters.mel(sr=RATE, n_fft=N_FFT, n_mels=N_MELS) @ magnitudes
+  mel_spec = mel(sr=RATE, n_fft=N_FFT, n_mels=N_MELS).numpy() @ magnitudes

  log_spec = np.log10(np.clip(mel_spec, 1e-10, None))
  log_spec = np.maximum(log_spec, log_spec.max((1,2), keepdims=True) - 8.0)
@@ -244,15 +245,16 @@ 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-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)
+    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)
      next_tokens[ctx[:, -1] == eot] = eot
      ctx = np.concatenate((ctx, next_tokens), axis=1)
      pos = ctx.shape[-1] - 1
-      if (next_tokens == eot).all(): break
+      if (next_tokens == eot).all() or pos == nctx: 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,11 +26,13 @@ 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)
-  bar = (bar[:-len(opt_text)] + opt_text) if opt_text else bar
+  if opt_text and len(opt_text) <= len(bar): bar = (bar[:-len(opt_text)] + opt_text)
  bar = colored(bar[:filled_width], color) + bar[filled_width:]
  return f'[{bar}]'

@@ -88,6 +90,7 @@ 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()
@@ -95,6 +98,20 @@ 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]}
@@ -116,6 +133,7 @@ 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)

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

-  def collect(self): return {d: d.smu.read_metrics() if d.pci_state == "D0" else None for d in self.devs}
+  def collect(self):
+    tables = {}
+    for dev in self.devs:
+      match dev.ip_ver[am.MP1_HWIP]:
+        case (13,0,6)|(13,0,12): table_t = dev.smu.smu_mod.MetricsTableX_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 get_gfx_activity(self, dev, metrics): return metrics.SmuMetrics.AverageGfxActivity
-  def get_mem_activity(self, dev, metrics): return metrics.SmuMetrics.AverageUclkActivity
+  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)|(13,0,12): 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)|(13,0,12): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
+      case _: return metrics.SmuMetrics.AverageUclkActivity

  def get_temps(self, dev, metrics, compact=False):
-    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}
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12):
+        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}

  def get_voltage(self, dev, metrics, compact=False):
-    voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_SVI_PLANE_e__enumvalues.items()
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): return {}
+      case _:
+        voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.SVI_PLANE_e.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]:
@@ -153,22 +202,40 @@ class SMICtx:
      case _: return 15

  def get_gfx_freq(self, dev, metrics):
-    return metrics.SmuMetrics.AverageGfxclkFrequencyPostDs if self.get_gfx_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
-          metrics.SmuMetrics.AverageGfxclkFrequencyPreDs
+    if metrics is None: return 0
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): 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

  def get_mem_freq(self, dev, metrics):
-    return metrics.SmuMetrics.AverageMemclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
-           metrics.SmuMetrics.AverageMemclkFrequencyPreDs
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): 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

  def get_fckl_freq(self, dev, metrics):
-    return metrics.SmuMetrics.AverageFclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
-           metrics.SmuMetrics.AverageFclkFrequencyPreDs
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): 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

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

-  def get_power(self, dev, metrics): return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
+  def get_power(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
+      case _: 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:
@@ -177,7 +244,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(entry):
+        if pt.lv!=am.AMDGPU_VM_PTB and not dev.gmc.is_pte_huge_page(pt.lv, entry):
          pt_stack.append(AMPageTableEntry(dev, entry & 0x0000FFFFFFFFF000, lv=pt.lv+1))
          continue
        if (entry & am.AMDGPU_PTE_SYSTEM) != 0: continue
@@ -219,23 +286,28 @@ 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 ==="] + [f"Fan Speed: {fan_rpm} RPM"] + [f"Fan Power: {fan_pwm}%"]
+      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}%"]

      total_power, max_power = self.get_power(dev, metrics)
-      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")]
+      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"]

      voltage_data = self.get_voltage(dev, metrics)
-      voltage_table = ["=== Voltages ==="] + [f"{name:<20}: {color_voltage(voltage)}" for name, voltage in voltage_data.items()]
+      voltage_table = None if not voltage_data else (["=== 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 + [""] + voltage_table
+        power_table += power_line
+        if voltage_table is not None: power_table += [""] + voltage_table
        activity_line += [""]
      elif self.prev_terminal_width >= 171:
        power_table += power_line + [""] + frequency_table
@@ -307,4 +379,5 @@ if __name__ == "__main__":
      smi_ctx.draw(args.list)
      if args.list: break
      time.sleep(1)
-  except KeyboardInterrupt: print("Exiting...")
+  except KeyboardInterrupt:
+    print("Exiting...")
@@ -0,0 +1,14 @@
+#!/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,48 +1,65 @@
 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, devfmt, vram, doorbell, mmio, dma_regions=None):
-    self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
+  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')
    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):
-  register_map = register_names
+def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
+  register_map = register_names or {}

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

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

  def replace_register_2(match):
-    register = match.group(1)
-    return f"Writing register {register_map.get(int(register, base=16), register)}"
+    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)

-  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 inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"
+      for xcc, addr in y.addr.items():
+        reg_names[addr] = f"{x}, xcc={xcc}"

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

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

  with open(sys.argv[2], 'w') as f:
    f.write(processed_log)
@@ -51,5 +68,4 @@ if __name__ == '__main__':
  if len(sys.argv) != 3:
    print("Usage: <input_file_path> <output_file_path>")
    sys.exit(1)
-
-  main()
+  main()
@@ -0,0 +1,39 @@
+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 pseudocode from AMD PDF
+* dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
+* pcode.py -- pseudocode execution environment. pseudocode 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.
+
+The autogen folder is autogenerated from the AMD PDFs with `python3 -m extra.assembly.amd.pdf --arch all`
+
+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
+If a test is failing with `AMD=1 PYTHON_REMU=1 MOCKGPU=1` it's 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.
+IMPORTANT: if a test is failing in the emulator, it's an instruction bug. Use DEBUG=7, get the instructions, and debug.
+
+Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~2000 lines.
+Get line count with `cloc --by-file extra/assembly/amd/*.py`
+
@@ -0,0 +1,463 @@
+# AMD ISA code generator - generates enum.py, ins.py, operands.py, str_pcode.py
+# Sources: XML from https://gpuopen.com/download/machine-readable-isa/latest/
+#          PDF manuals from AMD documentation
+import re, zlib, xml.etree.ElementTree as ET, zipfile
+from tinygrad.helpers import fetch
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Configuration
+# ═══════════════════════════════════════════════════════════════════════════════
+
+ARCHS = {
+  "rdna3": {"xml": "amdgpu_isa_rdna3_5.xml", "pdf": "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content"},
+  "rdna4": {"xml": "amdgpu_isa_rdna4.xml", "pdf": "https://docs.amd.com/api/khub/documents/uQpkEvk3pv~kfAb2x~j4uw/content"},
+  "cdna": {"xml": "amdgpu_isa_cdna4.xml", "pdf": "https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-cdna4-instruction-set-architecture.pdf"},
+}
+XML_URL = "https://gpuopen.com/download/machine-readable-isa/latest/"
+# Map XML encoding names to codebase names
+NAME_MAP = {"VOP3_SDST_ENC": "VOP3SD", "VOP3_SDST_ENC_LIT": "VOP3SD_LIT", "VOP3_SDST_ENC_DPP16": "VOP3SD_DPP16",
+            "VOP3_SDST_ENC_DPP8": "VOP3SD_DPP8", "VOPDXY": "VOPD", "VOPDXY_LIT": "VOPD_LIT", "VDS": "DS"}
+# Instructions missing from XML but present in PDF
+FIXES = {"rdna3": {"SOPK": {22: "S_SUBVECTOR_LOOP_BEGIN", 23: "S_SUBVECTOR_LOOP_END"}, "FLAT": {55: "FLAT_ATOMIC_CSUB_U32"}},
+         "rdna4": {"SOP1": {80: "S_GET_BARRIER_STATE", 81: "S_BARRIER_INIT", 82: "S_BARRIER_JOIN"}, "SOPP": {9: "S_WAITCNT", 21: "S_BARRIER_LEAVE"}},
+         "cdna": {"DS": {152: "DS_GWS_SEMA_RELEASE_ALL", 154: "DS_GWS_SEMA_V", 156: "DS_GWS_SEMA_P"},
+                  "VOP3P": {44: "V_MFMA_LD_SCALE_B32", 62: "V_MFMA_F32_16X16X8_XF32", 63: "V_MFMA_F32_32X32X4_XF32"}}}
+# Encoding suffixes to strip (variants we don't generate separate classes for)
+_ENC_SUFFIXES = ("_NSA1",)
+# Encoding suffix to class suffix mapping (for variants we DO generate)
+_ENC_SUFFIX_MAP = {"_INST_LITERAL": "_LIT", "_VOP_DPP16": "_DPP16", "_VOP_DPP": "_DPP16", "_VOP_DPP8": "_DPP8",
+                   "_VOP_SDWA": "_SDWA", "_VOP_SDWA_SDST_ENC": "_SDWA_SDST", "_MFMA": "_MFMA"}
+# Field name normalization
+_FIELD_RENAMES = {"opsel_hi_2": "opsel_hi2", "op_sel_hi_2": "opsel_hi2", "op_sel": "opsel", "bound_ctrl": "bc",
+                  "tgt": "target", "row_en": "row", "unorm": "unrm", "clamp": "clmp", "wait_exp": "waitexp",
+                  "simm32": "literal", "dpp_ctrl": "dpp", "acc_cd": "acc_cd", "acc": "acc",
+                  "dst_sel": "dst_sel", "dst_unused": "dst_unused", "src0_sel": "src0_sel", "src1_sel": "src1_sel"}
+# Encoding variants to skip entirely (NSA is for MIMG graphics instructions)
+_SKIP_ENCODINGS = ("NSA",)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# XML parsing helpers
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def _strip_enc(name: str) -> str:
+  """Strip ENC_ prefix and normalize encoding suffixes."""
+  name = name.removeprefix("ENC_")
+  for sfx in _ENC_SUFFIXES: name = name.replace(sfx, "")
+  # Process longer suffixes first to avoid partial matches (e.g., _VOP_DPP8 before _VOP_DPP)
+  for old, new in sorted(_ENC_SUFFIX_MAP.items(), key=lambda x: -len(x[0])): name = name.replace(old, new)
+  return name
+
+def _norm_field(name: str) -> str:
+  """Normalize field name to match expected names."""
+  for old, new in _FIELD_RENAMES.items(): name = name.replace(old, new)
+  return name
+
+def _map_flat(enc_name: str, instr_name: str) -> str:
+  """Map FLAT/GLOBAL/SCRATCH encoding to correct enum based on instruction prefix."""
+  if enc_name in ("FLAT_GLBL", "FLAT_GLOBAL"): return "GLOBAL"
+  if enc_name == "FLAT_SCRATCH": return "SCRATCH"
+  if enc_name in ("FLAT", "VFLAT", "VGLOBAL", "VSCRATCH"):
+    v = "V" if enc_name.startswith("V") else ""
+    if instr_name.startswith("GLOBAL_"): return f"{v}GLOBAL"
+    if instr_name.startswith("SCRATCH_"): return f"{v}SCRATCH"
+    return f"{v}FLAT"
+  return enc_name
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# XML parsing
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def parse_xml(filename: str):
+  root = ET.fromstring(zipfile.ZipFile(fetch(XML_URL)).read(filename))
+  encodings, enums, types, fmts, op_types_set = {}, {}, {}, {}, set()
+  # Extract HWREG and MSG enums from OperandTypes
+  op_enum_map = {("OPR_HWREG", "ID"): "HWREG", ("OPR_SENDMSG_RTN", "MSG"): "MSG"}
+  for ot in root.findall(".//OperandTypes/OperandType"):
+    ot_name = ot.findtext("OperandTypeName")
+    for field in ot.findall(".//Field"):
+      if (enum_name := op_enum_map.get((ot_name, field.findtext("FieldName")))):
+        enums[enum_name] = {int(pv.findtext("Value")): pv.findtext("Name").upper() for pv in field.findall(".//PredefinedValue")}
+  # Extract DataFormats with BitCount
+  for df in root.findall("ISA/DataFormats/DataFormat"):
+    name, bits = df.findtext("DataFormatName"), df.findtext("BitCount")
+    if name and bits: fmts[name] = int(bits)
+  # Extract encoding definitions
+  for enc in root.findall("ISA/Encodings/Encoding"):
+    name = enc.findtext("EncodingName")
+    is_base = name.startswith("ENC_") or name in ("VOP3_SDST_ENC", "VOPDXY")
+    is_variant = any(sfx in name for sfx in _ENC_SUFFIX_MAP)
+    if not is_base and not is_variant: continue
+    if any(s in name for s in _SKIP_ENCODINGS): continue
+    fields = [(_norm_field(f.findtext("FieldName").lower()), int(f.find("BitLayout/Range").findtext("BitOffset") or 0) + int(f.find("BitLayout/Range").findtext("BitCount") or 0) - 1,
+               int(f.find("BitLayout/Range").findtext("BitOffset") or 0))
+              for f in enc.findall(".//MicrocodeFormat/BitMap/Field") if f.find("BitLayout/Range") is not None]
+    ident = (enc.findall("EncodingIdentifiers/EncodingIdentifier") or [None])[0]
+    enc_field = next((f for f in fields if f[0] == "encoding"), None)
+    # For multi-dword formats, encoding field may be in higher dword but identifier pattern is always in dword0; use % 32
+    enc_bits = "".join(ident.text[len(ident.text)-1-b] for b in range(enc_field[1] % 32, (enc_field[2] % 32)-1, -1)) if ident is not None and enc_field else None
+    base_name = _strip_enc(name)
+    encodings[NAME_MAP.get(base_name, base_name)] = (fields, enc_bits)
+  # Extract instruction opcodes and operand info
+  # Track which encodings each opcode appears in (for detecting LIT-only ops)
+  opcode_encs: dict[str, dict[int, set[str]]] = {}  # {base_fmt: {opcode: {enc_names}}}
+  for instr in root.findall("ISA/Instructions/Instruction"):
+    name = instr.findtext("InstructionName")
+    for enc in instr.findall("InstructionEncodings/InstructionEncoding"):
+      if enc.findtext("EncodingCondition") != "default": continue
+      base, opcode = _map_flat(_strip_enc(enc.findtext("EncodingName")), name), int(enc.findtext("Opcode") or 0)
+      enc_name = NAME_MAP.get(base, base)
+      # Encoding variants use the same Op enum as the base format
+      base_enum = enc_name
+      for sfx in ("_SDWA_SDST", "_DPP16", "_DPP8", "_SDWA", "_LIT", "_MFMA"):
+        base_enum = base_enum.replace(sfx, "")
+      # Track which encodings this opcode appears in
+      opcode_encs.setdefault(base_enum, {}).setdefault(opcode, set()).add(enc_name)
+      # ADDTID instructions go in both FLAT and GLOBAL enums (pcode uses FLATOp for these)
+      if "ADDTID" in name:
+        if base == "GLOBAL": enums.setdefault("FLAT", {})[opcode] = name
+        elif base == "VGLOBAL": enums.setdefault("VFLAT", {})[opcode] = name
+      enums.setdefault(base_enum, {})[opcode] = name
+      # Extract operand info
+      op_info = {op.findtext("FieldName").lower(): (op.findtext("DataFormatName"), int(op.findtext("OperandSize") or 0), op.findtext("OperandType"))
+                 for op in enc.findall("Operands/Operand") if op.findtext("FieldName")}
+      for fmt, _, otype in op_info.values():
+        if fmt and fmt not in fmts: fmts[fmt] = 0
+        if otype: op_types_set.add(otype)
+      if op_info: types[(name, base_enum)] = op_info
+  # Find opcodes that only exist in _LIT encoding (no base format version)
+  lit_only_ops: dict[str, set[int]] = {}
+  for base_fmt, opcodes in opcode_encs.items():
+    for opcode, encs in opcodes.items():
+      if all("_LIT" in e for e in encs):
+        lit_only_ops.setdefault(base_fmt, set()).add(opcode)
+  return encodings, enums, types, fmts, op_types_set, lit_only_ops
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PDF parsing
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def extract_pdf_text(url: str) -> list[list[tuple[float, float, str, str]]]:
+  """Extract positioned text from PDF. Returns list of text elements (x, y, text, font) per page."""
+  data = fetch(url).read_bytes()
+  # Parse xref table to locate objects
+  xref: dict[int, int] = {}
+  pos = int(re.search(rb'startxref\s+(\d+)', data).group(1)) + 4
+  while data[pos:pos+7] != b'trailer':
+    while data[pos:pos+1] in b' \r\n': pos += 1
+    line_end = data.find(b'\n', pos)
+    start_obj, count = map(int, data[pos:line_end].split()[:2])
+    pos = line_end + 1
+    for i in range(count):
+      if data[pos+17:pos+18] == b'n' and (off := int(data[pos:pos+10])) > 0: xref[start_obj + i] = off
+      pos += 20
+
+  def get_stream(n: int) -> bytes:
+    obj = data[xref[n]:data.find(b'endobj', xref[n])]
+    raw = obj[obj.find(b'stream\n') + 7:obj.find(b'\nendstream')]
+    return zlib.decompress(raw) if b'/FlateDecode' in obj else raw
+
+  pages = []
+  for n in sorted(xref):
+    if b'/Type /Page' not in data[xref[n]:xref[n]+500]: continue
+    if not (m := re.search(rb'/Contents (\d+) 0 R', data[xref[n]:xref[n]+500])): continue
+    stream = get_stream(int(m.group(1))).decode('latin-1')
+    elements, font = [], ''
+    for bt in re.finditer(r'BT(.*?)ET', stream, re.S):
+      x, y = 0.0, 0.0
+      for m in re.finditer(r'(/F[\d.]+) [\d.]+ Tf|([\d.+-]+) ([\d.+-]+) Td|[\d.+-]+ [\d.+-]+ [\d.+-]+ [\d.+-]+ ([\d.+-]+) ([\d.+-]+) Tm|<([0-9A-Fa-f]+)>.*?Tj|\[([^\]]+)\] TJ', bt.group(1)):
+        if m.group(1): font = m.group(1)
+        elif m.group(2): x, y = x + float(m.group(2)), y + float(m.group(3))
+        elif m.group(4): x, y = float(m.group(4)), float(m.group(5))
+        elif m.group(6) and (t := bytes.fromhex(m.group(6)).decode('latin-1')).strip(): elements.append((x, y, t, font))
+        elif m.group(7) and (t := ''.join(bytes.fromhex(h).decode('latin-1') for h in re.findall(r'<([0-9A-Fa-f]+)>', m.group(7)))).strip(): elements.append((x, y, t, font))
+    pages.append(sorted(elements, key=lambda e: (-e[1], e[0])))
+  return pages
+
+def extract_pcode(pages: list[list[tuple[float, float, str, str]]], name_to_op: dict[str, int]) -> dict[tuple[str, int], str]:
+  """Extract pseudocode for instructions. Returns {(name, opcode): pseudocode}."""
+  # First pass: find all instruction headers across all pages
+  all_instructions: list[tuple[int, float, str, int]] = []  # (page_idx, y, name, opcode)
+  for page_idx, page in enumerate(pages):
+    by_y: dict[int, list[tuple[float, str]]] = {}
+    for x, y, t, _ in page:
+      by_y.setdefault(round(y), []).append((x, t))
+    for y, items in sorted(by_y.items(), reverse=True):
+      left = [(x, t) for x, t in items if 55 < x < 65]
+      right = [(x, t) for x, t in items if 535 < x < 550]
+      if left and right and left[0][1] in name_to_op and right[0][1].isdigit():
+        all_instructions.append((page_idx, y, left[0][1], int(right[0][1])))
+
+  # Second pass: extract pseudocode between consecutive instructions
+  pcode: dict[tuple[str, int], str] = {}
+  for i, (page_idx, y, name, opcode) in enumerate(all_instructions):
+    if i + 1 < len(all_instructions):
+      next_page, next_y = all_instructions[i + 1][0], all_instructions[i + 1][1]
+    else:
+      next_page, next_y = page_idx, 0
+    # Collect F6 text from current position to next instruction (pseudocode is at x ≈ 69)
+    lines = []
+    for p in range(page_idx, next_page + 1):
+      start_y = y if p == page_idx else 800
+      end_y = next_y if p == next_page else 0
+      lines.extend((p, y2, t) for x, y2, t, f in pages[p] if f in ('/F6.0', '/F7.0') and end_y < y2 < start_y and 60 < x < 80)
+    if lines:
+      sorted_lines = sorted(lines, key=lambda x: (x[0], -x[1]))
+      # Stop at large Y gaps (>30) - indicates section break
+      filtered = [sorted_lines[0]]
+      for j in range(1, len(sorted_lines)):
+        prev_page, prev_y, _ = sorted_lines[j-1]
+        curr_page, curr_y, _ = sorted_lines[j]
+        if curr_page == prev_page and prev_y - curr_y > 30: break
+        if curr_page != prev_page and prev_y > 60 and curr_y < 730: break
+        filtered.append(sorted_lines[j])
+      pcode_lines = [t.replace('Ê', '').strip() for _, _, t in filtered]
+      if pcode_lines: pcode[(name, opcode)] = '\n'.join(pcode_lines)
+  return pcode
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Code generation
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def write_common(all_fmts, all_op_types, path):
+  lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
+  lines.append("class Fmt(Enum):")
+  for fmt in sorted(all_fmts.keys()): lines.append(f"  {fmt} = auto()")
+  lines.append("")
+  lines.append("FMT_BITS = {")
+  for fmt, bits in sorted(all_fmts.items()): lines.append(f"  Fmt.{fmt}: {bits},")
+  lines.append("}")
+  lines.append("")
+  lines.append("class OpType(Enum):")
+  for ot in sorted(all_op_types): lines.append(f"  {ot} = auto()")
+  with open(path, "w") as f: f.write("\n".join(lines))
+
+def write_enum(enums, path):
+  lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum", "from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType  # noqa: F401", ""]
+  for name, ops in sorted(enums.items()):
+    if not ops: continue
+    suffix = "_E32" if name in ("VOP1", "VOP2", "VOPC") else "_E64" if name == "VOP3" else ""
+    lines.append(f"class {name}(Enum):" if name in ("HWREG", "MSG") else f"class {name}Op(Enum):")
+    aliases = []
+    for op, mem in sorted(ops.items()):
+      msuf = suffix if name != "VOP3" or op < 512 else ""
+      lines.append(f"  {mem}{msuf} = {op}")
+      if msuf: aliases.append((mem, msuf))
+    for mem, msuf in aliases: lines.append(f"  {mem} = {mem}{msuf}")
+    lines.append("")
+  with open(path, "w") as f: f.write("\n".join(lines))
+
+def write_ins(encodings, enums, lit_only_ops, types, arch, path):
+  _VGPR_FIELDS = {"vdst", "vdstx", "vsrc0", "vsrc1", "vsrc2", "vsrc3", "vsrcx1", "vsrcy1", "vaddr", "vdata", "data", "data0", "data1", "addr"}
+  _VARIANT_SUFFIXES = ("_LIT", "_DPP16", "_DPP8", "_SDWA_SDST", "_SDWA", "_MFMA")
+  def get_base_fmt(fmt):
+    for sfx in _VARIANT_SUFFIXES: fmt = fmt.replace(sfx, "")
+    return fmt
+  def field_def(name, hi, lo, fmt, enc_bits=None):
+    bits = hi - lo + 1
+    base_fmt = get_base_fmt(fmt)
+    if name == "encoding" and enc_bits: return f"FixedBitField({hi}, {lo}, 0b{enc_bits})"
+    if name == "op" and fmt not in ("DPP", "SDWA"): return f"EnumBitField({hi}, {lo}, {base_fmt}Op)"
+    if name in ("opx", "opy"): return f"EnumBitField({hi}, {lo}, VOPDOp)"
+    if name == "vdsty": return f"VDSTYField({hi}, {lo})"
+    if name in _VGPR_FIELDS and bits == 8: return f"VGPRField({hi}, {lo})"
+    if name == "sbase" and bits == 6: return f"SBaseField({hi}, {lo})"
+    if name in ("srsrc", "ssamp") and bits == 5: return f"SRsrcField({hi}, {lo})"
+    if name in ("sdst", "sdata") and bits == 7: return f"SGPRField({hi}, {lo})"
+    if name in ("soffset", "saddr") and bits == 7: return f"SGPRField({hi}, {lo}, default=NULL)"
+    if name.startswith("ssrc") and bits == 8: return f"SSrcField({hi}, {lo})"
+    if name in ("saddr", "soffset") and bits == 8: return f"SSrcField({hi}, {lo}, default=NULL)"
+    if name.startswith("src") and bits == 9: return f"SrcField({hi}, {lo})"
+    if base_fmt == "VOP3P" and name == "opsel_hi": return f"BitField({hi}, {lo}, default=3)"
+    if base_fmt == "VOP3P" and name == "opsel_hi2": return f"BitField({hi}, {lo}, default=1)"
+    return f"BitField({hi}, {lo})"
+  ORDER = ['encoding', 'op', 'opx', 'opy', 'vdst', 'vdstx', 'vdsty', 'sdst', 'vdata', 'sdata', 'addr', 'vaddr', 'data', 'data0', 'data1',
+           'src0', 'srcx0', 'srcy0', 'vsrc0', 'ssrc0', 'src1', 'vsrc1', 'vsrcx1', 'vsrcy1', 'ssrc1', 'src2', 'vsrc2', 'src3', 'vsrc3',
+           'saddr', 'sbase', 'srsrc', 'ssamp', 'soffset', 'offset', 'simm16', 'literal', 'en', 'target', 'attr', 'attr_chan',
+           'omod', 'neg', 'neg_hi', 'abs', 'clmp', 'opsel', 'opsel_hi', 'waitexp', 'wait_va',
+           'dmask', 'dim', 'seg', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe', 'unrm', 'done', 'row',
+           'dpp', 'fi', 'bc', 'row_mask', 'bank_mask', 'src0_neg', 'src0_abs', 'src1_neg', 'src1_abs',
+           'cbsz', 'abid', 'acc_cd', 'acc', 'blgp', 'lane_sel_0', 'lane_sel_1', 'lane_sel_2', 'lane_sel_3',
+           'lane_sel_4', 'lane_sel_5', 'lane_sel_6', 'lane_sel_7', 'dst_sel', 'dst_unused', 'src0_sel', 'src1_sel']
+  sort_fields = lambda fields: sorted(fields, key=lambda f: (ORDER.index(f[0]) if f[0] in ORDER else 999, f[2]))
+
+  # Separate base encodings from variants
+  base_encodings, variant_encodings = {}, {}
+  for enc_name, data in encodings.items():
+    base = get_base_fmt(enc_name)
+    if base == enc_name: base_encodings[enc_name] = data
+    else: variant_encodings[enc_name] = data
+
+  # Build sets of ops by their vdst type from operand metadata
+  sdst_opcodes = {}  # ops where vdst is OPR_SREG (writes to SGPR)
+  for fmt, ops in enums.items():
+    for op, name in ops.items():
+      op_types = types.get((name, fmt), {})
+      vdst_type = op_types.get("vdst", (None, None, None))[2]
+      if vdst_type == "OPR_SREG": sdst_opcodes.setdefault(fmt, set()).add(op)
+
+  lines = ["# autogenerated from AMD ISA XML - do not edit", "# ruff: noqa: F401,F403",
+           "from extra.assembly.amd.dsl import *", f"from extra.assembly.amd.autogen.{arch}.enum import *", "import functools", ""]
+
+  def fmt_allowed(op_enum: str, ops: set[int]) -> str:
+    """Format allowed ops as {EnumName.MEMBER, ...}."""
+    names = [f"{op_enum}.{enums[op_enum.removesuffix('Op')][op]}" for op in sorted(ops)]
+    return "{" + ", ".join(names) + "}"
+
+  # Generate base classes first
+  for enc_name, (fields, enc_bits) in sorted(base_encodings.items()):
+    # Get lit-only ops for this format (these can't be used in base class)
+    base_lit_ops = lit_only_ops.get(enc_name, set())
+    all_ops = set(enums.get(enc_name, {}).keys())
+    # Exclude SDST ops from base class (they need VOP1_SDST/VOP3_SDST/VOP3B)
+    base_allowed = all_ops - base_lit_ops - sdst_opcodes.get(enc_name, set())
+    if enc_name in ("FLAT", "VFLAT"):
+      prefix = "V" if enc_name == "VFLAT" else ""
+      for cls, seg, op_enum in [(f"{prefix}FLAT", 0, f"{prefix}FLATOp"), (f"{prefix}GLOBAL", 2, f"{prefix}GLOBALOp"), (f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]:
+        cls_ops = set(enums.get(cls, {}).keys())
+        lines.append(f"class {cls}(Inst):")
+        for fn, hi, lo in sort_fields(fields):
+          if fn == "seg": lines.append(f"  seg = FixedBitField({hi}, {lo}, {seg})")
+          elif fn == "op": lines.append(f"  op = EnumBitField({hi}, {lo}, {op_enum}, {fmt_allowed(op_enum, cls_ops)})")
+          else: lines.append(f"  {fn} = {field_def(fn, hi, lo, cls, enc_bits)}")
+        lines.append("")
+    elif enc_name not in ("FLAT_GLOBAL", "FLAT_SCRATCH", "FLAT_GLBL", "VGLOBAL", "VSCRATCH", "DPP", "SDWA"):
+      lines.append(f"class {enc_name}(Inst):")
+      for fn, hi, lo in sort_fields(fields):
+        if fn == "op":
+          base_fmt = get_base_fmt(enc_name)
+          lines.append(f"  op = EnumBitField({hi}, {lo}, {base_fmt}Op, {fmt_allowed(f'{base_fmt}Op', base_allowed)})")
+        else:
+          lines.append(f"  {fn} = {field_def(fn, hi, lo, enc_name, enc_bits if fn == 'encoding' else None)}")
+      lines.append("")
+
+  # Generate variant classes that inherit from base (only add extra fields)
+  for enc_name, (fields, enc_bits) in sorted(variant_encodings.items()):
+    base = get_base_fmt(enc_name)
+    if base not in base_encodings: continue  # skip if no base class
+    base_fields = {f[0] for f in base_encodings[base][0]}
+    extra_fields = [(fn, hi, lo) for fn, hi, lo in fields if fn not in base_fields]
+    is_lit = enc_name.endswith("_LIT")
+    all_ops = set(enums.get(base, {}).keys())
+    if extra_fields or is_lit:
+      lines.append(f"class {enc_name}({base}):")
+      op_field = next((f for f in base_encodings[base][0] if f[0] == "op"), None)
+      # _LIT classes: override op to allow all opcodes (base excludes lit-only ops)
+      if op_field and is_lit:
+        _, hi, lo = op_field
+        lines.append(f"  op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', all_ops)})")
+      for fn, hi, lo in sort_fields(extra_fields):
+        lines.append(f"  {fn} = {field_def(fn, hi, lo, enc_name)}")
+      lines.append("")
+
+  # SDST variants (special case - redefine vdst field type, restrict to SDST ops)
+  for base, field_hi, field_lo in [("VOP1", 24, 17), ("VOP3", 7, 0)]:
+    if base not in base_encodings: continue
+    sdst_ops = sdst_opcodes.get(base, set())
+    if not sdst_ops: continue
+    # For VOP3, all ops < 256 (compare/cmpx ops) use SDST encoding
+    all_base_ops = set(enums.get(base, {}).keys())
+    if base == "VOP3": sdst_ops = sdst_ops | {op for op in all_base_ops if op < 256}
+    op_field = next((f for f in base_encodings[base][0] if f[0] == "op"), None)
+    lines.append(f"class {base}_SDST({base}):")
+    if op_field:
+      _, hi, lo = op_field
+      lines.append(f"  op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', sdst_ops)})")
+    lines.append(f"  vdst = SSrcField({field_hi}, {field_lo})")
+    lines.append("")
+    # SDST_LIT class (for literals with SDST destination) - same ops, just adds literal field
+    lit_enc = variant_encodings.get(f"{base}_LIT")
+    if lit_enc:
+      lit_field = next((f for f in lit_enc[0] if f[0] == "literal"), None)
+      if lit_field:
+        lines.append(f"class {base}_SDST_LIT({base}_SDST):")
+        lines.append(f"  literal = BitField({lit_field[1]}, {lit_field[2]})")
+        lines.append("")
+
+  # Instruction helpers
+  lines.append("# instruction helpers")
+  for fmt, ops in sorted(enums.items()):
+    if fmt not in base_encodings and fmt not in ("GLOBAL", "SCRATCH", "VGLOBAL", "VSCRATCH"): continue
+    suffix = "_E32" if fmt in ("VOP1", "VOP2", "VOPC") else "_E64" if fmt == "VOP3" else ""
+    lit_ops = lit_only_ops.get(fmt, set())
+    fmt_sdst_ops = sdst_opcodes.get(fmt, set())
+    for op, name in sorted(ops.items()):
+      msuf = suffix if fmt != "VOP3" or op < 512 else ""
+      # Determine class: SDST variants, LIT-only instructions, or base
+      if fmt == "VOP1" and op in fmt_sdst_ops: cls = "VOP1_SDST"
+      elif fmt == "VOP3" and (op in fmt_sdst_ops or op < 256): cls = "VOP3_SDST"
+      elif op in lit_ops: cls = f"{fmt}_LIT"
+      else: cls = fmt
+      lines.append(f"{name.lower()}{msuf.lower()} = functools.partial({cls}, {fmt}Op.{name}{msuf})")
+  with open(path, "w") as f: f.write("\n".join(lines))
+
+def write_operands(types, enums, arch, path):
+  valid = {(name, fmt) for fmt, ops in enums.items() for name in ops.values()}
+  lines = ["# autogenerated from AMD ISA XML - do not edit",
+           "from extra.assembly.amd.autogen.common import Fmt, OpType",
+           f"from extra.assembly.amd.autogen.{arch}.enum import *", ""]
+  lines.append("# instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}")
+  lines.append("OPERANDS = {")
+  def fmt_val(v):
+    fmt, size, otype = v
+    return f"({f'Fmt.{fmt}' if fmt else 'None'}, {size}, {f'OpType.{otype}' if otype else 'None'})"
+  for (name, enc_base), fields in sorted(types.items()):
+    if (name, enc_base) not in valid: continue
+    fstr = ", ".join(f'"{k}": {fmt_val(v)}' for k, v in sorted(fields.items()))
+    lines.append(f'  {enc_base}Op.{name}: {{{fstr}}},')
+  lines.append("}")
+  with open(path, "w") as f: f.write("\n".join(lines))
+
+def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, str]], arch: str, path: str):
+  """Write str_pcode.py file from extracted pseudocode."""
+  entries: list[tuple[str, str, int, str]] = []
+  for fmt_name, ops in enums.items():
+    member_suffix = "_E32" if fmt_name in ("VOP1", "VOP2", "VOPC") else "_E64" if fmt_name == "VOP3" else ""
+    for opcode, name in ops.items():
+      if (name, opcode) in pcode:
+        msuf = member_suffix if fmt_name != "VOP3" or opcode < 512 else ""
+        entries.append((f"{fmt_name}Op", f"{name}{msuf}", opcode, pcode[(name, opcode)]))
+  enum_names = sorted(set(e[0] for e in entries))
+  lines = ["# autogenerated from AMD ISA PDF - do not edit", "# ruff: noqa: E501",
+           f"from extra.assembly.amd.autogen.{arch}.enum import {', '.join(enum_names)}", "", "PCODE = {"]
+  for enum_name, name, opcode, code in sorted(entries, key=lambda x: (x[0], x[2])):
+    lines.append(f"  {enum_name}.{name}: {code!r},")
+  lines.append("}")
+  with open(path, "w") as f: f.write("\n".join(lines))
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Main
+# ═══════════════════════════════════════════════════════════════════════════════
+
+if __name__ == "__main__":
+  import pathlib
+  all_fmts, all_op_types, arch_data = {}, set(), {}
+  # First pass: parse XML for all architectures
+  for arch, cfg in ARCHS.items():
+    print(f"Parsing XML: {cfg['xml']} -> {arch}")
+    encodings, enums, types, fmts, op_types_set, lit_only_ops = parse_xml(cfg["xml"])
+    for fmt, ops in FIXES.get(arch, {}).items(): enums.setdefault(fmt, {}).update(ops)
+    arch_data[arch] = {"encodings": encodings, "enums": enums, "types": types, "lit_only_ops": lit_only_ops}
+    for fmt, bits in fmts.items():
+      assert fmt not in all_fmts or all_fmts[fmt] == bits, f"FMT_BITS mismatch for {fmt}: {all_fmts[fmt]} vs {bits}"
+      all_fmts[fmt] = bits
+    all_op_types.update(op_types_set)
+  # Write common.py
+  common_path = pathlib.Path(__file__).parent / "autogen" / "common.py"
+  write_common(all_fmts, all_op_types, common_path)
+  print(f"Wrote common.py: {len(all_fmts)} formats, {len(all_op_types)} op types")
+  # Write per-arch files from XML
+  for arch, data in arch_data.items():
+    base = pathlib.Path(__file__).parent / "autogen" / arch
+    write_enum(data["enums"], base / "enum.py")
+    write_ins(data["encodings"], data["enums"], data["lit_only_ops"], data["types"], arch, base / "ins.py")
+    write_operands(data["types"], data["enums"], arch, base / "operands.py")
+    print(f"  {arch}: {len(data['encodings'])} encodings, {sum(len(v) for v in data['enums'].values())} instructions")
+  # Second pass: parse PDFs and write pcode
+  for arch, cfg in ARCHS.items():
+    print(f"Parsing PDF: {arch}...")
+    pages = extract_pdf_text(cfg["pdf"])
+    name_to_op = {name: op for ops in arch_data[arch]["enums"].values() for op, name in ops.items()}
+    pcode = extract_pcode(pages, name_to_op)
+    base = pathlib.Path(__file__).parent / "autogen" / arch
+    write_pcode(pcode, arch_data[arch]["enums"], arch, base / "str_pcode.py")
+    print(f"  {arch}: {len(pcode)} pcode entries")
@@ -0,0 +1,320 @@
+# autogenerated from AMD ISA XML - do not edit
+from enum import Enum, auto
+
+class Fmt(Enum):
+  FMT_ANY = auto()
+  FMT_BUF = auto()
+  FMT_IMG = auto()
+  FMT_IMG_BVH = auto()
+  FMT_NUM_B1 = auto()
+  FMT_NUM_B1024 = auto()
+  FMT_NUM_B128 = auto()
+  FMT_NUM_B16 = auto()
+  FMT_NUM_B256 = auto()
+  FMT_NUM_B32 = auto()
+  FMT_NUM_B512 = auto()
+  FMT_NUM_B64 = auto()
+  FMT_NUM_B8 = auto()
+  FMT_NUM_B96 = auto()
+  FMT_NUM_BF16 = auto()
+  FMT_NUM_BF6 = auto()
+  FMT_NUM_BF8 = auto()
+  FMT_NUM_F16 = auto()
+  FMT_NUM_F32 = auto()
+  FMT_NUM_F64 = auto()
+  FMT_NUM_FP4 = auto()
+  FMT_NUM_FP6 = auto()
+  FMT_NUM_FP8 = auto()
+  FMT_NUM_I16 = auto()
+  FMT_NUM_I24 = auto()
+  FMT_NUM_I32 = auto()
+  FMT_NUM_I4 = auto()
+  FMT_NUM_I64 = auto()
+  FMT_NUM_I8 = auto()
+  FMT_NUM_IU4 = auto()
+  FMT_NUM_IU8 = auto()
+  FMT_NUM_M64 = auto()
+  FMT_NUM_PK16_BF16 = auto()
+  FMT_NUM_PK16_BF8 = auto()
+  FMT_NUM_PK16_F16 = auto()
+  FMT_NUM_PK16_F32 = auto()
+  FMT_NUM_PK16_FP8 = auto()
+  FMT_NUM_PK16_I32 = auto()
+  FMT_NUM_PK16_I8 = auto()
+  FMT_NUM_PK2_B16 = auto()
+  FMT_NUM_PK2_B32 = auto()
+  FMT_NUM_PK2_B64 = auto()
+  FMT_NUM_PK2_BF16 = auto()
+  FMT_NUM_PK2_BF8 = auto()
+  FMT_NUM_PK2_F16 = auto()
+  FMT_NUM_PK2_F32 = auto()
+  FMT_NUM_PK2_FP4 = auto()
+  FMT_NUM_PK2_FP8 = auto()
+  FMT_NUM_PK2_I16 = auto()
+  FMT_NUM_PK2_I8 = auto()
+  FMT_NUM_PK2_U16 = auto()
+  FMT_NUM_PK2_U8 = auto()
+  FMT_NUM_PK32_BF16 = auto()
+  FMT_NUM_PK32_BF6 = auto()
+  FMT_NUM_PK32_BF8 = auto()
+  FMT_NUM_PK32_F16 = auto()
+  FMT_NUM_PK32_F32 = auto()
+  FMT_NUM_PK32_FP6 = auto()
+  FMT_NUM_PK32_FP8 = auto()
+  FMT_NUM_PK32_I32 = auto()
+  FMT_NUM_PK32_I8 = auto()
+  FMT_NUM_PK4_B8 = auto()
+  FMT_NUM_PK4_BF16 = auto()
+  FMT_NUM_PK4_BF8 = auto()
+  FMT_NUM_PK4_F16 = auto()
+  FMT_NUM_PK4_F32 = auto()
+  FMT_NUM_PK4_F64 = auto()
+  FMT_NUM_PK4_FP8 = auto()
+  FMT_NUM_PK4_I32 = auto()
+  FMT_NUM_PK4_I8 = auto()
+  FMT_NUM_PK4_IU8 = auto()
+  FMT_NUM_PK4_U8 = auto()
+  FMT_NUM_PK8_B32 = auto()
+  FMT_NUM_PK8_BF16 = auto()
+  FMT_NUM_PK8_BF8 = auto()
+  FMT_NUM_PK8_F16 = auto()
+  FMT_NUM_PK8_FP8 = auto()
+  FMT_NUM_PK8_I4 = auto()
+  FMT_NUM_PK8_I8 = auto()
+  FMT_NUM_PK8_IU4 = auto()
+  FMT_NUM_PK8_U4 = auto()
+  FMT_NUM_PK8_U8 = auto()
+  FMT_NUM_PK_F16 = auto()
+  FMT_NUM_PK_I16 = auto()
+  FMT_NUM_PK_I8 = auto()
+  FMT_NUM_PK_U16 = auto()
+  FMT_NUM_PK_U8 = auto()
+  FMT_NUM_U16 = auto()
+  FMT_NUM_U24 = auto()
+  FMT_NUM_U32 = auto()
+  FMT_NUM_U4 = auto()
+  FMT_NUM_U64 = auto()
+  FMT_NUM_U8 = auto()
+  FMT_RSRC = auto()
+  FMT_RSRC_SCALAR = auto()
+  FMT_RSRC_SCRATCH = auto()
+  FMT_RSRC_SCRATCH_BYTE = auto()
+  FMT_RSRC_SCRATCH_STRIDE = auto()
+  FMT_RSRC_TYPED = auto()
+  FMT_RSRC_TYPED_BYTE = auto()
+  FMT_RSRC_TYPED_SCRATCH = auto()
+  FMT_RSRC_TYPED_STRIDE = auto()
+  FMT_RSRC_VECTOR = auto()
+  FMT_RSRC_VECTOR_BYTE = auto()
+  FMT_RSRC_VECTOR_STRIDE = auto()
+  FMT_SAMP = auto()
+  FMT_WMMA_AB_16X16_BF16 = auto()
+  FMT_WMMA_AB_16X16_BF8 = auto()
+  FMT_WMMA_AB_16X16_F16 = auto()
+  FMT_WMMA_AB_16X16_FP8 = auto()
+  FMT_WMMA_AB_16X16_IU4 = auto()
+  FMT_WMMA_AB_16X16_IU8 = auto()
+  FMT_WMMA_AB_16X32_BF16 = auto()
+  FMT_WMMA_AB_16X32_BF8 = auto()
+  FMT_WMMA_AB_16X32_F16 = auto()
+  FMT_WMMA_AB_16X32_FP8 = auto()
+  FMT_WMMA_AB_16X32_IU4 = auto()
+  FMT_WMMA_AB_16X32_IU8 = auto()
+  FMT_WMMA_AB_16X64_IU4 = auto()
+  FMT_WMMA_AB_BF16 = auto()
+  FMT_WMMA_AB_F16 = auto()
+  FMT_WMMA_AB_IU4 = auto()
+  FMT_WMMA_AB_IU8 = auto()
+  FMT_WMMA_DC_16X16_BF16 = auto()
+  FMT_WMMA_DC_16X16_F16 = auto()
+  FMT_WMMA_DC_16X16_F32 = auto()
+  FMT_WMMA_DC_16X16_I32 = auto()
+  FMT_WMMA_DC_BF16 = auto()
+  FMT_WMMA_DC_F16 = auto()
+  FMT_WMMA_DC_F32 = auto()
+  FMT_WMMA_DC_I32 = auto()
+  FMT_WMMA_INDEX_SET = auto()
+
+FMT_BITS = {
+  Fmt.FMT_ANY: 1,
+  Fmt.FMT_BUF: 64,
+  Fmt.FMT_IMG: 256,
+  Fmt.FMT_IMG_BVH: 128,
+  Fmt.FMT_NUM_B1: 1,
+  Fmt.FMT_NUM_B1024: 1024,
+  Fmt.FMT_NUM_B128: 128,
+  Fmt.FMT_NUM_B16: 16,
+  Fmt.FMT_NUM_B256: 256,
+  Fmt.FMT_NUM_B32: 32,
+  Fmt.FMT_NUM_B512: 512,
+  Fmt.FMT_NUM_B64: 64,
+  Fmt.FMT_NUM_B8: 8,
+  Fmt.FMT_NUM_B96: 96,
+  Fmt.FMT_NUM_BF16: 16,
+  Fmt.FMT_NUM_BF6: 6,
+  Fmt.FMT_NUM_BF8: 8,
+  Fmt.FMT_NUM_F16: 16,
+  Fmt.FMT_NUM_F32: 32,
+  Fmt.FMT_NUM_F64: 64,
+  Fmt.FMT_NUM_FP4: 4,
+  Fmt.FMT_NUM_FP6: 6,
+  Fmt.FMT_NUM_FP8: 8,
+  Fmt.FMT_NUM_I16: 16,
+  Fmt.FMT_NUM_I24: 24,
+  Fmt.FMT_NUM_I32: 32,
+  Fmt.FMT_NUM_I4: 4,
+  Fmt.FMT_NUM_I64: 64,
+  Fmt.FMT_NUM_I8: 8,
+  Fmt.FMT_NUM_IU4: 4,
+  Fmt.FMT_NUM_IU8: 8,
+  Fmt.FMT_NUM_M64: 64,
+  Fmt.FMT_NUM_PK16_BF16: 256,
+  Fmt.FMT_NUM_PK16_BF8: 128,
+  Fmt.FMT_NUM_PK16_F16: 256,
+  Fmt.FMT_NUM_PK16_F32: 512,
+  Fmt.FMT_NUM_PK16_FP8: 128,
+  Fmt.FMT_NUM_PK16_I32: 512,
+  Fmt.FMT_NUM_PK16_I8: 128,
+  Fmt.FMT_NUM_PK2_B16: 32,
+  Fmt.FMT_NUM_PK2_B32: 64,
+  Fmt.FMT_NUM_PK2_B64: 128,
+  Fmt.FMT_NUM_PK2_BF16: 32,
+  Fmt.FMT_NUM_PK2_BF8: 16,
+  Fmt.FMT_NUM_PK2_F16: 32,
+  Fmt.FMT_NUM_PK2_F32: 64,
+  Fmt.FMT_NUM_PK2_FP4: 8,
+  Fmt.FMT_NUM_PK2_FP8: 16,
+  Fmt.FMT_NUM_PK2_I16: 32,
+  Fmt.FMT_NUM_PK2_I8: 16,
+  Fmt.FMT_NUM_PK2_U16: 32,
+  Fmt.FMT_NUM_PK2_U8: 16,
+  Fmt.FMT_NUM_PK32_BF16: 512,
+  Fmt.FMT_NUM_PK32_BF6: 192,
+  Fmt.FMT_NUM_PK32_BF8: 256,
+  Fmt.FMT_NUM_PK32_F16: 512,
+  Fmt.FMT_NUM_PK32_F32: 1024,
+  Fmt.FMT_NUM_PK32_FP6: 192,
+  Fmt.FMT_NUM_PK32_FP8: 256,
+  Fmt.FMT_NUM_PK32_I32: 1024,
+  Fmt.FMT_NUM_PK32_I8: 256,
+  Fmt.FMT_NUM_PK4_B8: 32,
+  Fmt.FMT_NUM_PK4_BF16: 64,
+  Fmt.FMT_NUM_PK4_BF8: 32,
+  Fmt.FMT_NUM_PK4_F16: 64,
+  Fmt.FMT_NUM_PK4_F32: 128,
+  Fmt.FMT_NUM_PK4_F64: 256,
+  Fmt.FMT_NUM_PK4_FP8: 32,
+  Fmt.FMT_NUM_PK4_I32: 128,
+  Fmt.FMT_NUM_PK4_I8: 32,
+  Fmt.FMT_NUM_PK4_IU8: 32,
+  Fmt.FMT_NUM_PK4_U8: 32,
+  Fmt.FMT_NUM_PK8_B32: 256,
+  Fmt.FMT_NUM_PK8_BF16: 128,
+  Fmt.FMT_NUM_PK8_BF8: 64,
+  Fmt.FMT_NUM_PK8_F16: 128,
+  Fmt.FMT_NUM_PK8_FP8: 64,
+  Fmt.FMT_NUM_PK8_I4: 32,
+  Fmt.FMT_NUM_PK8_I8: 64,
+  Fmt.FMT_NUM_PK8_IU4: 32,
+  Fmt.FMT_NUM_PK8_U4: 32,
+  Fmt.FMT_NUM_PK8_U8: 64,
+  Fmt.FMT_NUM_PK_F16: 32,
+  Fmt.FMT_NUM_PK_I16: 32,
+  Fmt.FMT_NUM_PK_I8: 32,
+  Fmt.FMT_NUM_PK_U16: 32,
+  Fmt.FMT_NUM_PK_U8: 32,
+  Fmt.FMT_NUM_U16: 16,
+  Fmt.FMT_NUM_U24: 24,
+  Fmt.FMT_NUM_U32: 32,
+  Fmt.FMT_NUM_U4: 4,
+  Fmt.FMT_NUM_U64: 64,
+  Fmt.FMT_NUM_U8: 8,
+  Fmt.FMT_RSRC: 128,
+  Fmt.FMT_RSRC_SCALAR: 128,
+  Fmt.FMT_RSRC_SCRATCH: 128,
+  Fmt.FMT_RSRC_SCRATCH_BYTE: 128,
+  Fmt.FMT_RSRC_SCRATCH_STRIDE: 128,
+  Fmt.FMT_RSRC_TYPED: 128,
+  Fmt.FMT_RSRC_TYPED_BYTE: 128,
+  Fmt.FMT_RSRC_TYPED_SCRATCH: 128,
+  Fmt.FMT_RSRC_TYPED_STRIDE: 128,
+  Fmt.FMT_RSRC_VECTOR: 128,
+  Fmt.FMT_RSRC_VECTOR_BYTE: 128,
+  Fmt.FMT_RSRC_VECTOR_STRIDE: 128,
+  Fmt.FMT_SAMP: 128,
+  Fmt.FMT_WMMA_AB_16X16_BF16: 128,
+  Fmt.FMT_WMMA_AB_16X16_BF8: 64,
+  Fmt.FMT_WMMA_AB_16X16_F16: 128,
+  Fmt.FMT_WMMA_AB_16X16_FP8: 64,
+  Fmt.FMT_WMMA_AB_16X16_IU4: 32,
+  Fmt.FMT_WMMA_AB_16X16_IU8: 64,
+  Fmt.FMT_WMMA_AB_16X32_BF16: 256,
+  Fmt.FMT_WMMA_AB_16X32_BF8: 128,
+  Fmt.FMT_WMMA_AB_16X32_F16: 256,
+  Fmt.FMT_WMMA_AB_16X32_FP8: 128,
+  Fmt.FMT_WMMA_AB_16X32_IU4: 64,
+  Fmt.FMT_WMMA_AB_16X32_IU8: 128,
+  Fmt.FMT_WMMA_AB_16X64_IU4: 128,
+  Fmt.FMT_WMMA_AB_BF16: 256,
+  Fmt.FMT_WMMA_AB_F16: 256,
+  Fmt.FMT_WMMA_AB_IU4: 64,
+  Fmt.FMT_WMMA_AB_IU8: 128,
+  Fmt.FMT_WMMA_DC_16X16_BF16: 128,
+  Fmt.FMT_WMMA_DC_16X16_F16: 128,
+  Fmt.FMT_WMMA_DC_16X16_F32: 256,
+  Fmt.FMT_WMMA_DC_16X16_I32: 256,
+  Fmt.FMT_WMMA_DC_BF16: 256,
+  Fmt.FMT_WMMA_DC_F16: 256,
+  Fmt.FMT_WMMA_DC_F32: 256,
+  Fmt.FMT_WMMA_DC_I32: 256,
+  Fmt.FMT_WMMA_INDEX_SET: 32,
+}
+
+class OpType(Enum):
+  OPR_ACCVGPR = auto()
+  OPR_ATTR = auto()
+  OPR_CLAUSE = auto()
+  OPR_DELAY = auto()
+  OPR_EXEC = auto()
+  OPR_HWREG = auto()
+  OPR_LABEL = auto()
+  OPR_SDST = auto()
+  OPR_SDST_NULL = auto()
+  OPR_SENDMSG = auto()
+  OPR_SENDMSG_RTN = auto()
+  OPR_SIMM16 = auto()
+  OPR_SIMM24 = auto()
+  OPR_SIMM4 = auto()
+  OPR_SIMM5 = auto()
+  OPR_SIMM8 = auto()
+  OPR_SLEEP = auto()
+  OPR_SMEM_OFFSET = auto()
+  OPR_SMEM_OFFSET_NOK = auto()
+  OPR_SRC = auto()
+  OPR_SRC_ACCVGPR = auto()
+  OPR_SRC_NOLDS = auto()
+  OPR_SRC_NOLIT = auto()
+  OPR_SRC_SIMPLE = auto()
+  OPR_SRC_VGPR = auto()
+  OPR_SRC_VGPR_OR_ACCVGPR = auto()
+  OPR_SRC_VGPR_OR_ACCVGPR_OR_CONST = auto()
+  OPR_SRC_VGPR_OR_INLINE = auto()
+  OPR_SREG = auto()
+  OPR_SREG_LITERAL = auto()
+  OPR_SREG_M0 = auto()
+  OPR_SREG_M0_INL = auto()
+  OPR_SREG_NOVCC = auto()
+  OPR_SSRC = auto()
+  OPR_SSRC_BARRIER_ID = auto()
+  OPR_SSRC_LANESEL = auto()
+  OPR_SSRC_NOLIT = auto()
+  OPR_TGT = auto()
+  OPR_VERSION = auto()
+  OPR_VGPR = auto()
+  OPR_VGPR_OR_ACCVGPR = auto()
+  OPR_VGPR_OR_LDS = auto()
+  OPR_WAITCNT = auto()
+  OPR_WAITCNT_DEPCTR = auto()
+  OPR_WAIT_ALU = auto()
+  OPR_WAIT_EVENT = auto()
+  OPR_WAIT_MEM_DS = auto()
@@ -0,0 +1,65 @@
+# Instruction format detection and decoding
+from __future__ import annotations
+from extra.assembly.amd.dsl import Inst, FixedBitField, EnumBitField
+
+# SDWA/DPP variant detection: src0 field (bits 0-8) encodes the variant
+# 0xf9 (249) = SDWA, 0xfa (250) = DPP16 for CDNA (GFX9)
+_VARIANT_SRC0 = {"_SDWA_SDST": 0xf9, "_SDWA": 0xf9, "_DPP16": 0xfa}
+
+def _matches(data: bytes, cls: type[Inst]) -> bool:
+  """Check if data matches all FixedBitFields and op is in allowed."""
+  for _, field in cls._fields:
+    dword_idx = field.lo // 32
+    if len(data) < (dword_idx + 1) * 4: return False
+    word = int.from_bytes(data[dword_idx*4:(dword_idx+1)*4], 'little')
+    field_lo = field.lo % 32
+    if isinstance(field, FixedBitField):
+      if ((word >> field_lo) & field.mask) != field.default: return False
+    if isinstance(field, EnumBitField) and field.allowed is not None:
+      try: opcode = field.decode((word >> field_lo) & field.mask)
+      except ValueError: return False  # opcode not in enum
+      if opcode not in field.allowed: return False
+  # Check SDWA/DPP variant based on src0 field (bits 0-8) - only for variant classes
+  name = cls.__name__
+  word = int.from_bytes(data[:4], 'little')
+  for suffix, expected_src0 in _VARIANT_SRC0.items():
+    if name.endswith(suffix): return (word & 0x1ff) == expected_src0
+  return True
+
+# Import instruction classes for each architecture
+from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3SD, VOP3P, VOPC, VOPD, VINTERP,
+  SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH)
+from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT,
+  VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
+  VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
+  SOPC as R4_SOPC, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP,
+  SMEM as R4_SMEM, DS as R4_DS, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
+from extra.assembly.amd.autogen.cdna.ins import (VOP1 as C_VOP1, VOP1_SDWA as C_VOP1_SDWA, VOP1_DPP16 as C_VOP1_DPP16,
+  VOP2 as C_VOP2, VOP2_LIT as C_VOP2_LIT, VOP2_SDWA as C_VOP2_SDWA, VOP2_DPP16 as C_VOP2_DPP16,
+  VOPC as C_VOPC, VOPC_SDWA_SDST as C_VOPC_SDWA_SDST,
+  VOP3 as C_VOP3, VOP3_SDST as C_VOP3_SDST, VOP3SD as C_VOP3SD, VOP3P as C_VOP3P, VOP3PX2 as C_VOP3PX2,
+  SOP1 as C_SOP1, SOP2 as C_SOP2, SOPC as C_SOPC, SOPK as C_SOPK, SOPK_LIT as C_SOPK_LIT, SOPP as C_SOPP, SMEM as C_SMEM, DS as C_DS,
+  FLAT as C_FLAT, GLOBAL as C_GLOBAL, SCRATCH as C_SCRATCH, MUBUF as C_MUBUF)
+
+# Order matters: more specific encodings first, catch-alls (SOP2, VOP2) last
+# Order: base before _LIT (base matches regular ops, _LIT catches lit-only ops excluded from base)
+_FORMATS = {
+  "rdna3": [VOPD, VOP3P, VINTERP, VOP3SD, VOP3_SDST, VOP3, DS, GLOBAL, SCRATCH, FLAT, SMEM,
+            SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, VOPC, VOP1_SDST, VOP1, VOP1_LIT, VOP2, VOP2_LIT],
+  "rdna4": [R4_VOPD, R4_VOP3P, R4_VINTERP, R4_VOP3SD, R4_VOP3_SDST, R4_VOP3, R4_DS, R4_GLOBAL, R4_SCRATCH, R4_FLAT, R4_SMEM,
+            R4_SOP1, R4_SOPC, R4_SOPP, R4_SOPK, R4_SOPK_LIT, R4_VOPC, R4_VOP1_SDST, R4_VOP1, R4_SOP2, R4_SOP2_LIT, R4_VOP2, R4_VOP2_LIT],
+  "cdna": [C_VOP3PX2, C_VOP3P, C_VOP3SD, C_VOP3_SDST, C_VOP3, C_DS, C_GLOBAL, C_SCRATCH, C_FLAT, C_MUBUF, C_SMEM,
+            C_SOP1, C_SOPC, C_SOPP, C_SOPK, C_SOPK_LIT, C_VOPC_SDWA_SDST, C_VOPC,
+            C_VOP1_DPP16, C_VOP1_SDWA, C_VOP1, C_VOP2_DPP16, C_VOP2_SDWA, C_SOP2, C_VOP2, C_VOP2_LIT],
+}
+
+def detect_format(data: bytes, arch: str = "rdna3") -> type[Inst]:
+  """Detect instruction format from machine code bytes."""
+  assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
+  for cls in _FORMATS[arch]:
+    if _matches(data, cls): return cls
+  raise ValueError(f"unknown {arch} format word={int.from_bytes(data[:4], 'little'):#010x}")
+
+def decode_inst(data: bytes, arch: str = "rdna3") -> Inst:
+  """Decode machine code bytes into an instruction."""
+  return detect_format(data, arch).from_bytes(data)
@@ -0,0 +1,899 @@
+# RDNA3/RDNA4/CDNA disassembler
+from __future__ import annotations
+import re, struct
+from extra.assembly.amd.dsl import Inst, Reg
+
+# Special register mappings for disassembly
+SPECIAL_GPRS = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'null', 125: 'm0', 126: 'exec_lo', 127: 'exec_hi',
+                128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc'}
+SPECIAL_GPRS_CDNA = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'm0', 126: 'exec_lo', 127: 'exec_hi',
+                     128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc',
+                     102: 'flat_scratch_lo', 103: 'flat_scratch_hi', 104: 'xnack_mask_lo', 105: 'xnack_mask_hi',
+                     251: 'src_vccz', 252: 'src_execz'}
+SPECIAL_PAIRS = {106: 'vcc', 126: 'exec'}
+SPECIAL_PAIRS_CDNA = {106: 'vcc', 126: 'exec', 102: 'flat_scratch', 104: 'xnack_mask'}
+
+def decode_src(v, cdna: bool = False) -> str:
+  """Decode a source operand encoding to its string representation."""
+  v = _unwrap(v)
+  gprs = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
+  if v in gprs: return gprs[v]
+  if v < 106: return f's{v}'
+  if 108 <= v < 124: return f'ttmp{v - 108}'
+  if 129 <= v <= 192: return str(v - 128)  # positive integers 1-64
+  if 193 <= v <= 208: return str(-(v - 192))  # negative integers -1 to -16
+  if v >= 256: return f'v{v - 256}'
+  return f's{v}'
+
+def _unwrap(v) -> int:
+  """Unwrap Reg to int offset, or return int as-is."""
+  return v.offset if isinstance(v, Reg) else v
+
+def _vi(v) -> int:
+  """Get VGPR index from Reg or int (for v[N] fields that encode as 256+N)."""
+  off = _unwrap(v)
+  return off - 256 if off >= 256 else off
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# LITERAL FORMATTING
+# ═══════════════════════════════════════════════════════════════════════════════
+
+_FLOAT_DEC = {240: 0.5, 241: -0.5, 242: 1.0, 243: -1.0, 244: 2.0, 245: -2.0, 246: 4.0, 247: -4.0}
+
+def _lit(inst, v, neg=0, cdna=None) -> str:
+  """Format literal/inline constant value."""
+  if cdna is None: cdna = _is_cdna(inst)
+  v = _unwrap(v)
+  if v == 255:
+    lit = inst._literal
+    if lit is None: return "0"
+    s = f"0x{lit:x}"
+  elif v in _FLOAT_DEC: s = str(_FLOAT_DEC[v])
+  elif 128 <= v <= 192: s = str(v - 128)
+  elif 193 <= v <= 208: s = str(-(v - 192))
+  elif v < 128: s = decode_src(v, cdna)
+  elif v >= 256: s = f"v{v - 256}"
+  else: s = decode_src(v, cdna)
+  return f"-{s}" if neg else s
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# INSTRUCTION METADATA - fallback functions when inst.num_srcs()/inst.operands unavailable
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def _num_srcs(inst) -> int:
+  """Fallback: get number of source operands from instruction name."""
+  name = getattr(inst, 'op_name', '') or ''
+  n = name.upper()
+  # FMAC/MAC ops are 2-source (dst is implicit accumulator), but FMA/MAD ops are 3-source
+  if 'FMAC' in n or 'V_MAC_' in n: return 2
+  if any(x in n for x in ('FMA', 'MAD', 'CNDMASK', 'BFE', 'BFI', 'LERP', 'MED3', 'SAD', 'DIV_FMAS', 'DIV_FIXUP', 'DIV_SCALE', 'CUBE')): return 3
+  # PERMLANE_VAR ops are 2-source, but PERMLANE (non-VAR) are 3-source
+  if 'PERMLANE' in n and '_VAR' not in n: return 3
+  if any(x in n for x in ('_ADD3', '_LSHL_ADD', '_ADD_LSHL', '_LSHL_OR', '_AND_OR', 'OR3_B32', 'AND_OR_B32', 'ALIGNBIT', 'ALIGNBYTE', 'V_PERM_', 'XOR3', 'XAD', 'MULLIT', 'MINMAX', 'MAXMIN', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM', 'MINIMUM3', 'MAXIMUM3', 'MIN3', 'MAX3', 'DOT2', 'CVT_PK_U8_F32', 'DOT4', 'DOT8', 'WMMA', 'SWMMAC')): return 3
+  return 2
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# IMPORTS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP1_SDST_LIT, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3_SDST_LIT,
+  VOP3_LIT, VOP3SD, VOP3SD_LIT, VOP3P, VOP3P_LIT, VOPC, VOPC_LIT, VOPD, VOPD_LIT, VINTERP, SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPC_LIT,
+  SOPK, SOPK_LIT, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH, VOP2Op, VOPDOp, SOPPOp, HWREG, MSG)
+from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP1_SDST_LIT as R4_VOP1_SDST_LIT, VOP1_LIT as R4_VOP1_LIT,
+  VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT, VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3_SDST_LIT as R4_VOP3_SDST_LIT, VOP3_LIT as R4_VOP3_LIT,
+  VOP3SD as R4_VOP3SD, VOP3SD_LIT as R4_VOP3SD_LIT, VOP3P as R4_VOP3P, VOP3P_LIT as R4_VOP3P_LIT, VOPC as R4_VOPC, VOPC_LIT as R4_VOPC_LIT,
+  VOPD as R4_VOPD, VOPD_LIT as R4_VOPD_LIT, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP1_LIT as R4_SOP1_LIT, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
+  SOPC as R4_SOPC, SOPC_LIT as R4_SOPC_LIT, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP, SMEM as R4_SMEM, DS as R4_DS,
+  VOPDOp as R4_VOPDOp, HWREG as HWREG_RDNA4)
+from extra.assembly.amd.autogen.cdna.ins import FLAT as C_FLAT, HWREG as HWREG_CDNA
+
+def _is_cdna(inst: Inst) -> bool: return 'cdna' in inst.__class__.__module__
+
+# CDNA opcode name aliases for disasm (new name -> old name expected by tests)
+_CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32', 'v_fmamk_f32': 'v_madmk_f32', 'v_fmaak_f32': 'v_madak_f32'}
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# HELPERS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{_unwrap(b)}" if n == 1 else f"{p}[{_unwrap(b)}:{_unwrap(b)+n-1}]"
+def _sreg(b: int, n: int = 1) -> str: return _reg("s", _unwrap(b), n)
+def _vreg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("v", b - 256 if b >= 256 else b, n)
+def _areg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("a", b - 256 if b >= 256 else b, n)  # accumulator registers for GFX90a
+def _ttmp(b, n: int = 1) -> str: b = _unwrap(b); return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
+
+def _fmt_sdst(v, n: int = 1, cdna: bool = False) -> str:
+  v = _unwrap(v)
+  if t := _ttmp(v, n): return t
+  pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
+  gprs = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
+  if n > 1: return pairs.get(v) or gprs.get(v) or _sreg(v, n)  # also check gprs for null/m0
+  return gprs.get(v, f"s{v}")
+
+def _fmt_src(v, n: int = 1, cdna: bool = False) -> str:
+  v = _unwrap(v)
+  if v == 253: return "src_scc"  # SCC as source operand
+  if n == 1: return decode_src(v, cdna)
+  if v >= 256: return _vreg(v, n)
+  if v <= 101: return _sreg(v, n)  # s0-s101 can be pairs, but 102+ are special on CDNA
+  pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
+  if n == 2 and v in pairs: return pairs[v]
+  if v <= 105: return _sreg(v, n)  # s102-s105 regular pairs for RDNA
+  if t := _ttmp(v, n): return t
+  return decode_src(v, cdna)
+
+def _fmt_v16(v, base: int = 256, hi_thresh: int = 384) -> str:
+  v = _unwrap(v)
+  return f"v{(v - base) & 0x7f}.{'h' if v >= hi_thresh else 'l'}"
+
+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: v = _unwrap(v); return _fmt_v16(v) if v >= 256 else _lit(inst, 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) -> str:
+  """Format VOP3 source operand with modifiers."""
+  v = _unwrap(v)
+  if v == 255: s = _lit(inst, v)  # literal constant takes priority
+  elif 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"
+  elif v == 253: s = "src_scc"  # VOP3 sources use src_scc not scc
+  else: s = _lit(inst, 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 ""
+  dst_hi = (opsel >> 3) & 1
+  if n == 1: return f" op_sel:[{opsel & 1},{dst_hi}]"
+  # Use 4-element format if bit 2 is set (src2 selection used) or if 3+ sources
+  if n == 2 and not ((opsel >> 2) & 1): return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{dst_hi}]"
+  return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{dst_hi}]"
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# DISASSEMBLER
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def _disasm_vop1(inst: VOP1) -> str:
+  name, cdna = inst.op_name.lower() or f'vop1_op_{inst.op}', _is_cdna(inst)
+  name = name.replace('_e32', '')  # Strip _e32 suffix
+  if any(x in name for x in ('v_nop', 'v_pipeflush', 'v_clrexcp')): return name  # no operands
+  if 'readfirstlane' in name:
+    src = inst.src0.fmt() if inst.src0.offset >= 256 else decode_src(inst.src0.offset, cdna)
+    vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
+    return f"{name} {_fmt_sdst(vdst_off, 1, cdna)}, {src}"
+  bits = inst.canonical_op_bits
+  is16_dst, is16_src = not cdna and bits['d'] == 16, not cdna and bits['s0'] == 16
+  # Format dst
+  if is16_dst: dst = _fmt_v16(inst.vdst)
+  else: dst = inst.vdst.fmt()
+  # Format src
+  if inst.src0.offset == 255: src = _lit(inst, inst.src0)
+  elif is16_src and inst.src0.offset >= 256: src = _fmt_v16(inst.src0)
+  elif inst.src0.sz > 1: src = _fmt_src(inst.src0, inst.src0.sz, cdna)
+  else: src = _lit(inst, inst.src0)
+  return f"{name} {dst}, {src}"
+
+_VOP2_CARRY_OUT = {'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'}  # carry out only
+_VOP2_CARRY_INOUT = {'v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'}  # carry in and out (CDNA)
+_VOP2_CARRY_INOUT_RDNA = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'}  # carry in and out (RDNA)
+def _disasm_vop2(inst: VOP2) -> str:
+  name, cdna = inst.op_name.lower(), _is_cdna(inst)
+  if cdna: name = _CDNA_DISASM_ALIASES.get(name, name)  # apply CDNA aliases
+  suf = "" if cdna or name.endswith('_e32') or (not cdna and inst.op == VOP2Op.V_DOT2ACC_F32_F16_E32) else "_e32"
+  lit = inst._literal
+  is16 = not cdna and inst.canonical_op_bits['d'] == 16
+  # fmaak/madak: dst = src0 * vsrc1 + K, fmamk/madmk: dst = src0 * K + vsrc1
+  if 'fmaak' in name or 'madak' in name or (not cdna and inst.op in (VOP2Op.V_FMAAK_F32_E32, VOP2Op.V_FMAAK_F16_E32)):
+    if lit is None: return f"op_{inst.op.value if hasattr(inst.op, 'value') else inst.op}"
+    if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1)}, 0x{lit:x}"
+    return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, 0x{lit:x}"
+  if 'fmamk' in name or 'madmk' in name or (not cdna and inst.op in (VOP2Op.V_FMAMK_F32_E32, VOP2Op.V_FMAMK_F16_E32)):
+    if lit is None: return f"op_{inst.op.value if hasattr(inst.op, 'value') else inst.op}"
+    if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, 0x{lit:x}, {_fmt_v16(inst.vsrc1)}"
+    return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, 0x{lit:x}, {inst.vsrc1.fmt()}"
+  if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1)}"
+  vcc = "vcc" if cdna else "vcc_lo"
+  basename = name.replace('_e32', '')
+  if cdna and basename in _VOP2_CARRY_OUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}"
+  if cdna and basename in _VOP2_CARRY_INOUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
+  if not cdna and basename in _VOP2_CARRY_INOUT_RDNA: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
+  sn0 = inst.canonical_op_regs.get('s0', 1)
+  if inst.vdst.sz > 1 or sn0 > 1 or inst.vsrc1.sz > 1:
+    src0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else _fmt_src(inst.src0, sn0, cdna)
+    return f"{name.replace('_e32', '')} {inst.vdst.fmt()}, {src0}, {inst.vsrc1.fmt()}"
+  return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}" + (f", {vcc}" if name == 'v_cndmask_b32' else "")
+
+def _disasm_vopc(inst: VOPC) -> str:
+  name, cdna = inst.op_name.lower(), _is_cdna(inst)
+  bits = inst.canonical_op_bits
+  is16 = bits['s0'] == 16
+  if cdna:
+    s0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else _fmt_src(inst.src0, inst.src0.sz, cdna)
+    return f"{name} vcc, {s0}, {inst.vsrc1.fmt()}"  # CDNA VOPC always outputs vcc
+  # RDNA: v_cmpx_* writes to exec (no vcc), v_cmp_* writes to vcc_lo
+  has_vcc = 'cmpx' not in name
+  s0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else inst.src0.fmt() if inst.src0.sz > 1 else _src16(inst, inst.src0.offset) if is16 else _lit(inst, inst.src0)
+  s1 = inst.vsrc1.fmt() if inst.vsrc1.sz > 1 else _fmt_v16(inst.vsrc1) if is16 else inst.vsrc1.fmt()
+  suf = "" if name.endswith('_e32') else "_e32"
+  return f"{name}{suf} vcc_lo, {s0}, {s1}" if has_vcc else f"{name}{suf} {s0}, {s1}"
+
+NO_ARG_SOPP = {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, SOPPOp.S_TTRACEDATA}
+
+def _disasm_sopp(inst: SOPP) -> str:
+  name, cdna = inst.op_name.lower(), _is_cdna(inst)
+  is_rdna4 = 'rdna4' in inst.__class__.__module__
+  # Ops that have no argument when simm16 == 0
+  no_arg_zero = {'s_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_wait_idle', 's_endpgm_saved',
+                 's_endpgm_ordered_ps_done', 's_code_end'}
+  if name in no_arg_zero: return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
+  if name == 's_endpgm': return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
+  if cdna:
+    if name == 's_waitcnt':
+      # GFX9 format: vmcnt[3:0]=bits[3:0], vmcnt[5:4]=bits[15:14], expcnt=bits[6:4], lgkmcnt=bits[11:8] (4 bits, max 15)
+      vm_lo, exp, lgkm, vm_hi = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 8) & 0xf, (inst.simm16 >> 14) & 0x3
+      vm = vm_lo | (vm_hi << 4)
+      p = [f"vmcnt({vm})" if vm != 0x3f else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0xf else ""]
+      return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
+    if name.startswith(('s_cbranch', 's_branch')): return f"{name} {inst.simm16}"
+    if name == 's_set_gpr_idx_mode':
+      flags = [n for i, n in enumerate(['SRC0', 'SRC1', 'SRC2', 'DST']) if inst.simm16 & (1 << i)]
+      return f"{name} gpr_idx({','.join(flags)})"
+    return f"{name} 0x{inst.simm16:x}" if inst.simm16 else name
+  # RDNA (use name-based checks instead of enum-based for cross-arch compatibility)
+  if name == 's_waitcnt':
+    if is_rdna4:
+      return f"{name} {inst.simm16}" if inst.simm16 else f"{name} 0"
+    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 name == 's_delay_alu':
+    deps = ['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']
+    skips = ['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'}"
+  if name.startswith(('s_cbranch', 's_branch')): return f"{name} 0x{inst.simm16:x}"
+  return f"{name} 0x{inst.simm16:x}"
+
+def _disasm_smem(inst: SMEM) -> str:
+  name, cdna = inst.op_name.lower(), _is_cdna(inst)
+  if name in ('s_gl1_inv', 's_dcache_inv', 's_dcache_inv_vol', 's_dcache_wb', 's_dcache_wb_vol', 's_icache_inv'): return name
+  soe, imm = getattr(inst, 'soe', 0) or getattr(inst, 'soffset_en', 0), getattr(inst, 'imm', 1)
+  is_rdna4 = 'rdna4' in inst.__class__.__module__
+  offset = inst.ioffset if is_rdna4 else getattr(inst, 'offset', 0)
+  if cdna:
+    if soe and imm: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
+    elif imm: off_s = f"0x{offset:x}"
+    elif offset < 256: off_s = decode_src(offset, cdna)
+    else: off_s = decode_src(inst.soffset, cdna)
+  elif offset and inst.soffset != 124: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
+  elif offset: off_s = f"0x{offset:x}"
+  else: off_s = decode_src(inst.soffset, cdna)
+  is_buffer = 'buffer' in name or 's_atc_probe_buffer' == name
+  sbase_idx, sbase_count = _unwrap(inst.sbase), 4 if is_buffer else 2
+  sbase_str = _fmt_src(sbase_idx, sbase_count, cdna) 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} {_unwrap(inst.sdata)}, {sbase_str}, {off_s}"
+  if 'prefetch' in name:
+    off = getattr(inst, 'ioffset', getattr(inst, 'offset', 0))
+    if off >= 0x800000: off = off - 0x1000000
+    off_s = f"0x{off:x}" if off > 255 else str(off)
+    soff_s = decode_src(inst.soffset, cdna) if inst.soffset != 124 else ("m0" if cdna else "null")
+    if 'pc_rel' in name: return f"{name} {off_s}, {soff_s}, {_unwrap(inst.sdata)}"
+    return f"{name} {sbase_str}, {off_s}, {soff_s}, {_unwrap(inst.sdata)}"
+  # Use get_field_bits for register count
+  dst_n = inst.canonical_op_regs.get('d', 1)
+  th, scope = getattr(inst, 'th', 0), getattr(inst, 'scope', 0)
+  if is_rdna4:  # RDNA4 uses th/scope instead of glc/dlc
+    th_names = ['TH_LOAD_RT', 'TH_LOAD_NT', 'TH_LOAD_HT', 'TH_LOAD_LU']
+    scope_names = ['SCOPE_CU', 'SCOPE_SE', 'SCOPE_DEV', 'SCOPE_SYS']
+    mods = (f" th:{th_names[th]}" if th else "") + (f" scope:{scope_names[scope]}" if scope else "")
+    return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}{mods}"
+  if th or scope:
+    th_names = ['TH_LOAD_RT', 'TH_LOAD_NT', 'TH_LOAD_HT', 'TH_LOAD_LU']
+    scope_names = ['SCOPE_CU', 'SCOPE_SE', 'SCOPE_DEV', 'SCOPE_SYS']
+    mods = (f" th:{th_names[th]}" if th else "") + (f" scope:{scope_names[scope]}" if scope else "")
+    return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}{mods}"
+  if 'discard' in name: return f"{name} {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
+  if name in ('s_memrealtime', 's_memtime'): return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}"
+  return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
+
+def _disasm_flat(inst: FLAT) -> str:
+  name, cdna = inst.op_name.lower(), _is_cdna(inst)
+  acc = getattr(inst, 'acc', 0)
+  reg_fn = _areg if acc else _vreg
+  seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
+  instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
+  # Global/scratch uses 13-bit signed offset
+  if seg != 'flat':
+    if cdna:
+      # CDNA: bit 12 is sign bit but not in offset field
+      raw = int.from_bytes(inst.to_bytes(), 'little')
+      off_val = inst.offset | ((raw >> 12) & 1) << 12  # get bit 12
+    else:
+      off_val = inst.offset
+    off_val = off_val if off_val < 4096 else off_val - 8192  # sign extend 13-bit
+  else:
+    off_val = inst.offset
+  # Use get_field_bits: data for stores/atomics, d for loads
+  regs = inst.canonical_op_regs
+  w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
+  off_s = f" offset:{off_val}" if off_val else ""
+  if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"
+  else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
+  if seg == 'flat': saddr_s = ""
+  elif _unwrap(inst.saddr) in (0x7F, 124): saddr_s = ", off"
+  elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr, cdna)}"
+  elif _unwrap(inst.saddr) in (SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS): saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
+  elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
+  else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
+  if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
+  # load_lds_* instructions: vaddr, saddr (no vdst, data goes to LDS)
+  if 'load_lds' in name:
+    addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
+    addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
+    return f"{instr} {addr_s}{saddr_s}{mods}"
+  if seg == 'flat': addr_w = 2  # flat always uses 64-bit vaddr
+  elif cdna: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
+  else: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
+  addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
+  data_s, vdst_s = reg_fn(inst.data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
+  glc_or_sc0 = inst.sc0 if cdna else inst.glc
+  if 'atomic' in name:
+    return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if glc_or_sc0 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} {reg_fn(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
+
+def _disasm_ds(inst: DS) -> str:
+  op, name = inst.op, inst.op_name.lower()
+  acc = getattr(inst, 'acc', 0)
+  reg_fn = _areg if acc else _vreg
+  gds = " gds" if getattr(inst, 'gds', 0) else ""
+  off = f" offset:{inst.offset0 | (inst.offset1 << 8)}" if inst.offset0 or inst.offset1 else ""
+  off2 = (" offset0:" + str(inst.offset0) if inst.offset0 else "") + (" offset1:" + str(inst.offset1) if inst.offset1 else "")
+  # Use get_field_bits: data for stores/writes/atomics, d for loads
+  regs = inst.canonical_op_regs
+  w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'write' in name or ('load' not in name and 'read' not in name) else regs.get('d', 1)
+  d0, d1, dst, addr = reg_fn(inst.data0, w), reg_fn(inst.data1, w), reg_fn(inst.vdst, w), _vreg(inst.addr)
+
+  if name == 'ds_nop': return name
+  if name == 'ds_bvh_stack_rtn_b32': return f"{name} {_vreg(inst.vdst)}, {addr}, {_vreg(inst.data0)}, {_vreg(inst.data1, 4)}{off}{gds}"
+  if 'bvh_stack_push' in name:
+    d1_regs = 8 if 'push8' in name else 4
+    vdst_regs = 2 if 'pop2' in name else 1
+    vdst_s = _vreg(inst.vdst, vdst_regs) if vdst_regs > 1 else _vreg(inst.vdst)
+    return f"{name} {vdst_s}, {addr}, {_vreg(inst.data0)}, {_vreg(inst.data1, d1_regs)}{off}{gds}"
+  if 'gws_sema' in name and 'sema_br' not in name: return f"{name}{off}{gds}"
+  if 'gws_' in name: return f"{name} {addr}{off}{gds}"
+  if name in ('ds_consume', 'ds_append'): return f"{name} {reg_fn(inst.vdst)}{off}{gds}"
+  if 'gs_reg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {reg_fn(inst.data0)}{off}{gds}"
+  if '2addr' in name:
+    if 'load' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}{off2}{gds}"
+    if 'store' in name and 'xchg' not in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
+    return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}, {d0}, {d1}{off2}{gds}"
+  if 'write2' in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
+  if 'read2' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}{off2}{gds}"
+  if 'xchg2' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}, {d0}, {d1}{off2}{gds}"
+  if 'load' in name or ('read' in name and 'read2' not in name): return f"{name} {reg_fn(inst.vdst)}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
+  if ('store' in name or 'write' in name) and not _has(name, 'cmp', 'xchg', 'write2'):
+    return f"{name} {reg_fn(inst.data0)}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
+  if 'swizzle' in name or name == 'ds_ordered_count': return f"{name} {reg_fn(inst.vdst)}, {addr}{off}{gds}"
+  if 'permute' in name: return f"{name} {reg_fn(inst.vdst)}, {addr}, {reg_fn(inst.data0)}{off}{gds}"
+  if 'condxchg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {addr}, {reg_fn(inst.data0, 2)}{off}{gds}"
+  if _has(name, 'cmpst', '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()
+  n_up = name.upper()
+  bits = inst.canonical_op_bits
+
+  # RDNA4 v_s_* scalar VOP3 instructions - vdst is SGPR (VGPRField adds 256)
+  if name.startswith('v_s_'):
+    src = _lit(inst, inst.src0) if _unwrap(inst.src0) == 255 else ("src_scc" if _unwrap(inst.src0) == 253 else _fmt_src(inst.src0, max(1, bits['s0'] // 32)))
+    if inst.neg & 1: src = f"-{src}"
+    if inst.abs & 1: src = f"|{src}|"
+    clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
+    vdst_raw = _unwrap(inst.vdst)
+    return f"{name} s{vdst_raw - 256 if vdst_raw >= 256 else vdst_raw}, {src}" + (" clamp" if clamp else "") + _omod(inst.omod)
+
+  # Use get_field_bits for register sizes and 16-bit detection
+  r0, r1, r2 = max(1, bits['s0'] // 32), max(1, bits['s1'] // 32), max(1, bits['s2'] // 32)
+  dn = max(1, bits['d'] // 32)
+  is16_d, is16_s, is16_s2 = bits['d'] == 16, bits['s0'] == 16, bits['s2'] == 16
+
+  s0 = _vop3_src(inst, inst.src0, inst.neg&1, inst.abs&1, inst.opsel&1, r0, is16_s)
+  s1 = _vop3_src(inst, inst.src1, inst.neg&2, inst.abs&2, inst.opsel&2, r1, is16_s)
+  s2 = _vop3_src(inst, inst.src2, inst.neg&4, inst.abs&4, inst.opsel&4, r2, is16_s2)
+
+  # Format destination
+  if 'readlane' in name:
+    vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
+    dst = _fmt_sdst(vdst_off, 1)
+  elif is16_d: dst = f"{inst.vdst.fmt()}.h" if (inst.opsel & 8) else f"{inst.vdst.fmt()}.l"
+  else: dst = inst.vdst.fmt()
+
+  clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
+  cl, om = " clamp" if clamp else "", _omod(inst.omod)
+  nonvgpr_opsel = (inst.src0.offset < 256 and (inst.opsel & 1)) or (inst.src1.offset < 256 and (inst.opsel & 2)) or (inst.src2.offset < 256 and (inst.opsel & 4))
+  need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)
+
+  op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
+  e64 = "" if name.endswith('_e64') else "_e64"
+  if op_val < 256:  # VOPC
+    vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
+    return f"{name}{e64} {s0}, {s1}{cl}" if name.startswith('v_cmpx') else f"{name}{e64} {_fmt_sdst(vdst_off, 1)}, {s0}, {s1}{cl}"
+  if op_val < 384:  # VOP2
+    n = inst.num_srcs() or 2
+    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 op_val < 512:  # VOP1
+    if re.match(r'v_cvt_f32_(bf|fp)8', name) and inst.opsel:
+      os = f" byte_sel:{((inst.opsel & 1) << 1) | ((inst.opsel >> 1) & 1)}"
+    else:
+      os = _opsel_str(inst.opsel, 1, need_opsel, is16_d)
+    if 'v_nop' in name or 'v_pipeflush' in name: return f"{name}{e64}"
+    return f"{name}{e64} {dst}, {s0}{os}{cl}{om}"
+  # Native VOP3
+  n = inst.num_srcs() or 2
+  os = f" byte_sel:{inst.opsel >> 2}" if 'cvt_sr' in name and inst.opsel else _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(reg, neg):
+    s = _lit(inst, reg.offset) if reg.offset == 255 else ("src_scc" if reg.offset == 253 else (reg.fmt() if reg.sz > 1 else _lit(inst, reg.offset)))
+    return f"neg({s})" if neg and reg.offset == 255 else (f"-{s}" if neg else s)
+  s0, s1, s2 = src(inst.src0, inst.neg & 1), src(inst.src1, inst.neg & 2), src(inst.src2, inst.neg & 4)
+  # VOP3SD: _co_ ops (add/sub) without _ci_ have only 2 sources, all others (mad, div_scale, _co_ci_) have 3 sources
+  has_only_two_srcs = '_co_' in name and '_ci_' not in name and 'mad' not in name
+  srcs = f"{s0}, {s1}" if has_only_two_srcs else f"{s0}, {s1}, {s2}"
+  clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
+  return f"{name} {inst.vdst.fmt()}, {_fmt_sdst(inst.sdst, 1)}, {srcs}{' clamp' if clamp else ''}{_omod(inst.omod)}"
+
+def _disasm_vopd(inst: VOPD) -> str:
+  lit = inst._literal
+  is_rdna4 = 'rdna4' in inst.__class__.__module__
+  op_enum = R4_VOPDOp if is_rdna4 else VOPDOp
+  vdst_y, nx, ny = (_unwrap(inst.vdsty) << 1) | ((_unwrap(inst.vdstx) & 1) ^ 1), op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
+  def half(n, vd, s0, vs1):
+    vd, vs1 = _vi(vd), _vi(vs1)
+    if 'mov' in n: return f"{n} v{vd}, {_lit(inst, s0)}"
+    if 'fmamk' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, 0x{lit:x}, v{vs1}"
+    if 'fmaak' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}, 0x{lit:x}"
+    return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}"
+  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, is_swmmac, n, is_fma_mix = 'wmma' in name, 'swmmac' in name, inst.num_srcs() or 2, 'fma_mix' in name
+  def get_src(reg):
+    return _lit(inst, reg.offset) if reg.offset == 255 else reg.fmt()
+  src0, src1, src2, dst = get_src(inst.src0), get_src(inst.src1), get_src(inst.src2), inst.vdst.fmt()
+  opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
+  clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
+  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 clamp else [])
+  elif is_swmmac:
+    mods = ([f"index_key:{inst.opsel}"] if inst.opsel 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 clamp else [])
+  else:
+    opsel_hi_default = 7 if n == 3 else 3
+    mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default 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 clamp 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_sop1(inst: SOP1) -> str:
+  op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
+  # Use get_field_bits for register sizes
+  regs = inst.canonical_op_regs
+  dst_regs, src_regs = regs.get('d', 1), regs.get('s0', 1)
+  src = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, src_regs, cdna)
+  if not cdna:
+    if 'getpc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2)}"
+    if 'setpc_b64' in name or 'rfe_b64' in name: return f"{name} {src}"
+    if 'swappc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2)}, {src}"
+    if 'sendmsg_rtn' in name:
+      v = _unwrap(inst.ssrc0)
+      try: msg_str = MSG(v).name if v != 255 else None  # MSG_RTN_ILLEGAL_MSG (255) not supported by LLVM
+      except ValueError: msg_str = None
+      return f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, sendmsg({msg_str})" if msg_str else f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, 0x{v:x}"
+  sop1_src_only = ('S_ALLOC_VGPR', 'S_SLEEP_VAR', 'S_BARRIER_SIGNAL', 'S_BARRIER_SIGNAL_ISFIRST', 'S_BARRIER_INIT', 'S_BARRIER_JOIN', 'S_SET_GPR_IDX_IDX',
+                   'S_CBRANCH_JOIN')
+  if inst.op_name in sop1_src_only: return f"{name} {src}"
+  if cdna:
+    if 'getpc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2, cdna)}"
+    if 'setpc_b64' in name or 'rfe_b64' in name: return f"{name} {src}"
+    if 'swappc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2, cdna)}, {src}"
+  return f"{name} {_fmt_sdst(inst.sdst, dst_regs, cdna)}, {src}"
+
+def _disasm_sop2(inst: SOP2) -> str:
+  cdna, name = _is_cdna(inst), inst.op_name.lower()
+  lit = inst._literal
+  # Use get_field_bits for register sizes
+  regs = inst.canonical_op_regs
+  dn, s0n, s1n = regs['d'], regs['s0'], regs['s1']
+  s0 = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, s0n, cdna)
+  s1 = _lit(inst, inst.ssrc1) if _unwrap(inst.ssrc1) == 255 else _fmt_src(inst.ssrc1, s1n, cdna)
+  dst = _fmt_sdst(inst.sdst, dn, cdna)
+  if 'fmamk' in name and lit is not None: return f"{name} {dst}, {s0}, 0x{lit:x}, {s1}"
+  if 'fmaak' in name and lit is not None: return f"{name} {dst}, {s0}, {s1}, 0x{lit:x}"
+  if name in ('s_cbranch_g_fork', 's_rfe_restore_b64'): return f"{name} {s0}, {s1}"  # no destination
+  return f"{name} {dst}, {s0}, {s1}"
+
+def _disasm_sopc(inst: SOPC) -> str:
+  cdna, regs, name = _is_cdna(inst), inst.canonical_op_regs, inst.op_name.lower()
+  s0 = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, regs['s0'], cdna)
+  if name == 's_set_gpr_idx_on':
+    imm = _unwrap(inst.ssrc1) & 0xf
+    flags = [n for i, n in enumerate(['SRC0', 'SRC1', 'SRC2', 'DST']) if imm & (1 << i)]
+    return f"{name} {s0}, gpr_idx({','.join(flags)})"
+  s1 = _lit(inst, inst.ssrc1) if _unwrap(inst.ssrc1) == 255 else _fmt_src(inst.ssrc1, regs['s1'], cdna)
+  return f"{name} {s0}, {s1}"
+
+_HWREG_BLACKLIST = {'HW_REG_PC_LO', 'HW_REG_PC_HI', 'HW_REG_IB_DBG1', 'HW_REG_FLUSH_IB', 'HW_REG_SHADER_TBA_LO', 'HW_REG_SHADER_TBA_HI',
+                    'HW_REG_SHADER_FLAT_SCRATCH_LO', 'HW_REG_SHADER_FLAT_SCRATCH_HI', 'HW_REG_SHADER_CYCLES'}
+_HWREG_BLACKLIST_CDNA = {'HW_REG_PC_LO', 'HW_REG_PC_HI', 'HW_REG_IB_DBG1', 'HW_REG_FLUSH_IB', 'HW_REG_SQ_SHADER_TBA_LO', 'HW_REG_SQ_SHADER_TBA_HI',
+                         'HW_REG_SQ_SHADER_TMA_LO', 'HW_REG_SQ_SHADER_TMA_HI', 'HW_REG_SQ_PERF_SNAPSHOT_DATA', 'HW_REG_SQ_PERF_SNAPSHOT_DATA1',
+                         'HW_REG_SQ_PERF_SNAPSHOT_PC_LO', 'HW_REG_SQ_PERF_SNAPSHOT_PC_HI', 'HW_REG_XCC_ID'}
+def _disasm_sopk(inst: SOPK) -> str:
+  op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
+  is_rdna4 = 'rdna4' in inst.__class__.__module__
+  hw = HWREG_CDNA if cdna else (HWREG_RDNA4 if is_rdna4 else HWREG)
+  blacklist = _HWREG_BLACKLIST_CDNA if cdna else _HWREG_BLACKLIST
+  def fmt_hwreg(hid, hoff, hsz):
+    try: hr_name = hw(hid).name.replace("HW_REG_WAVE_", "HW_REG_")
+    except ValueError: return f"0x{inst.simm16:x}"
+    if hr_name in blacklist: return f"0x{inst.simm16:x}"
+    return f"hwreg({hr_name})" if hoff == 0 and hsz == 32 else f"hwreg({hr_name}, {hoff}, {hsz})"
+  if name == 's_setreg_imm32_b32':
+    hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
+    return f"{name} {fmt_hwreg(hid, hoff, hsz)}, 0x{inst._literal:x}"
+  if name == 's_version': return f"{name} 0x{inst.simm16:x}"
+  if name in ('s_setreg_b32', 's_getreg_b32'):
+    hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
+    hs = fmt_hwreg(hid, hoff, hsz)
+    return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1, cdna)}" if 'setreg' in name else f"{name} {_fmt_sdst(inst.sdst, 1, cdna)}, {hs}"
+  if name in ('s_subvector_loop_begin', 's_subvector_loop_end'):
+    return f"{name} {_fmt_sdst(inst.sdst, 1)}, 0x{inst.simm16:x}"
+  return f"{name} {_fmt_sdst(inst.sdst, inst.canonical_op_regs['d'], cdna)}, 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()} {inst.vdst.fmt()}, {_lit(inst, inst.src0, inst.neg & 1)}, {_lit(inst, inst.src1, inst.neg & 2)}, {_lit(inst, inst.src2, inst.neg & 4)}" + (" " + mods if mods else "")
+
+DISASM_HANDLERS: dict[type, callable] = {
+  VOP1: _disasm_vop1, VOP1_SDST: _disasm_vop1, VOP1_SDST_LIT: _disasm_vop1, VOP1_LIT: _disasm_vop1,
+  VOP2: _disasm_vop2, VOP2_LIT: _disasm_vop2, VOPC: _disasm_vopc, VOPC_LIT: _disasm_vopc,
+  VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3_SDST_LIT: _disasm_vop3, VOP3_LIT: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOP3SD_LIT: _disasm_vop3sd,
+  VOPD: _disasm_vopd, VOPD_LIT: _disasm_vopd, VOP3P: _disasm_vop3p, VOP3P_LIT: _disasm_vop3p,
+  VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, GLOBAL: _disasm_flat, SCRATCH: _disasm_flat,
+  SOP1: _disasm_sop1, SOP1_LIT: _disasm_sop1, SOP2: _disasm_sop2, SOP2_LIT: _disasm_sop2,
+  SOPC: _disasm_sopc, SOPC_LIT: _disasm_sopc, SOPK: _disasm_sopk, SOPK_LIT: _disasm_sopk,
+  # RDNA4
+  R4_VOP1: _disasm_vop1, R4_VOP1_SDST: _disasm_vop1, R4_VOP1_SDST_LIT: _disasm_vop1, R4_VOP1_LIT: _disasm_vop1,
+  R4_VOP2: _disasm_vop2, R4_VOP2_LIT: _disasm_vop2, R4_VOPC: _disasm_vopc, R4_VOPC_LIT: _disasm_vopc,
+  R4_VOP3: _disasm_vop3, R4_VOP3_SDST: _disasm_vop3, R4_VOP3_SDST_LIT: _disasm_vop3, R4_VOP3_LIT: _disasm_vop3,
+  R4_VOP3SD: _disasm_vop3sd, R4_VOP3SD_LIT: _disasm_vop3sd, R4_VOP3P: _disasm_vop3p, R4_VOP3P_LIT: _disasm_vop3p,
+  R4_VOPD: _disasm_vopd, R4_VOPD_LIT: _disasm_vopd, R4_VINTERP: _disasm_vinterp, R4_SOPP: _disasm_sopp, R4_SMEM: _disasm_smem, R4_DS: _disasm_ds,
+  R4_SOP1: _disasm_sop1, R4_SOP1_LIT: _disasm_sop1, R4_SOP2: _disasm_sop2, R4_SOP2_LIT: _disasm_sop2,
+  R4_SOPC: _disasm_sopc, R4_SOPC_LIT: _disasm_sopc, R4_SOPK: _disasm_sopk, R4_SOPK_LIT: _disasm_sopk}
+
+def disasm(inst: Inst) -> str: return DISASM_HANDLERS[type(inst)](inst)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# CDNA DISASSEMBLER SUPPORT
+# ═══════════════════════════════════════════════════════════════════════════════
+
+from extra.assembly.amd.autogen.cdna.ins import (VOP1 as CDNA_VOP1, VOP1_LIT as CDNA_VOP1_LIT,
+  VOP1_SDWA as CDNA_VOP1_SDWA, VOP1_DPP16 as CDNA_VOP1_DPP16,
+  VOP2 as CDNA_VOP2, VOP2_LIT as CDNA_VOP2_LIT, VOP2_SDWA as CDNA_VOP2_SDWA, VOP2_DPP16 as CDNA_VOP2_DPP16,
+  VOPC as CDNA_VOPC, VOPC_LIT as CDNA_VOPC_LIT, VOPC_SDWA_SDST as CDNA_VOPC_SDWA_SDST,
+  VOP3 as CDNA_VOP3, VOP3_SDST as CDNA_VOP3_SDST, VOP3SD as CDNA_VOP3SD, VOP3P as CDNA_VOP3P, VOP3PX2 as CDNA_VOP3PX2,
+  SOP1 as CDNA_SOP1, SOP1_LIT as CDNA_SOP1_LIT, SOP2 as CDNA_SOP2, SOP2_LIT as CDNA_SOP2_LIT,
+  SOPC as CDNA_SOPC, SOPC_LIT as CDNA_SOPC_LIT, SOPK as CDNA_SOPK, SOPK_LIT as CDNA_SOPK_LIT,
+  SOPP as CDNA_SOPP, SMEM as CDNA_SMEM, DS as CDNA_DS,
+  FLAT as CDNA_FLAT, GLOBAL as CDNA_GLOBAL, SCRATCH as CDNA_SCRATCH, MUBUF as CDNA_MUBUF)
+
+def _cdna_src(inst, v, neg, abs_=0, n=1):
+  s = _lit(inst, v) if v == 255 else _fmt_src(v, n, cdna=True)
+  if abs_: s = f"|{s}|"
+  return f"neg({s})" if neg and v == 255 else (f"-{s}" if neg else s)
+
+_CDNA_VOP3_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32'}
+
+def _disasm_vop3a(inst) -> str:
+  op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
+  name = inst.op_name.lower() or f'vop3a_op_{op_val}'
+  n = inst.num_srcs() or _num_srcs(inst)
+  cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
+  # _sr_ instructions use 4-element op_sel (src2 for byte selection)
+  opsel_n = 3 if '_sr_' in name and n == 2 else n
+  opsel = _opsel_str(inst.opsel, opsel_n, inst.opsel != 0, False)
+  orig_name = name
+  name = _CDNA_VOP3_ALIASES.get(name, name)
+  if name != orig_name:
+    s0, s1 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, 1), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, 1)
+    s2 = ""
+    dst = _vreg(inst.vdst)
+  else:
+    regs = inst.canonical_op_regs
+    dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
+    s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, r0), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, r1), _cdna_src(inst, inst.src2, inst.neg&4, inst.abs&4, r2)
+    dst = _vreg(inst.vdst, dregs) if dregs > 1 else _vreg(inst.vdst)
+  if op_val >= 512:
+    return f"{name} {dst}, {s0}, {s1}, {s2}{opsel}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{opsel}{cl}{om}"
+  if op_val < 256:
+    # VOPC: vdst is actually sdst (SGPR pair), but VGPRField adds 256 to the offset
+    sdst_val = _unwrap(inst.vdst)
+    if sdst_val >= 256: sdst_val -= 256
+    sdst = _fmt_sdst(sdst_val, 2, cdna=True)
+    return f"{name} {sdst}, {s0}, {s1}{cl}"
+  if 320 <= op_val < 512:
+    if name in ('v_nop', 'v_clrexcp', 'v_nop_e64', 'v_clrexcp_e64'): return name.replace('_e64', '')
+    return f"{name} {dst}, {s0}{cl}{om}"
+  if name == 'v_cndmask_b32':
+    s2 = _fmt_src(inst.src2, 2, cdna=True)
+    return f"{name} {dst}, {s0}, {s1}, {s2}{cl}{om}"
+  return f"{name} {dst}, {s0}, {s1}, {s2}{opsel}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{opsel}{cl}{om}"
+
+def _disasm_vop3b(inst) -> str:
+  op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
+  name, cdna = inst.op_name.lower() or f'vop3b_op_{op_val}', _is_cdna(inst)
+  n = inst.num_srcs() or _num_srcs(inst)
+  regs = inst.canonical_op_regs
+  dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
+  s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, n=r0), _cdna_src(inst, inst.src1, inst.neg&2, n=r1), _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
+  # CDNA VOP3_SDST uses vdst field for sdst (but vdst adds 256), RDNA uses separate sdst field
+  sdst_val = getattr(inst, 'sdst', None)
+  if sdst_val is None and hasattr(inst, 'vdst'):
+    sdst_val = _unwrap(inst.vdst)
+    if sdst_val >= 256: sdst_val -= 256  # VGPRField adds 256, remove it for SGPR
+  # For CDNA VOP3_SDST (VOPC->VOP3), vdst is the scalar dest (sdst), there's no vdst output
+  if cdna and 'v_cmp' in name:
+    sdst = _fmt_sdst(sdst_val, 2, cdna=True)
+    cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
+    return f"{name} {sdst}, {s0}, {s1}{cl}{om}"
+  dst = _vreg(inst.vdst, dregs) if dregs > 1 else _vreg(inst.vdst)
+  sdst = _fmt_sdst(sdst_val, 2, cdna=cdna)
+  cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
+  if name in ('v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'):
+    s2 = _fmt_src(inst.src2, 2, cdna=cdna)
+  return f"{name} {dst}, {sdst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name} {dst}, {sdst}, {s0}, {s1}{cl}{om}"
+
+def _disasm_cdna_vop3p(inst) -> str:
+  name, n = inst.op_name.lower(), inst.num_srcs() or 2
+  is_mfma = 'mfma' in name or 'smfmac' in name
+  is_accvgpr = 'accvgpr' in name
+  get_src = lambda v, sc: _lit(inst, v) if v == 255 else _fmt_src(v, sc, cdna=True)
+
+  # Handle accvgpr read/write (accumulator register operations)
+  if is_accvgpr:
+    src0_off = _unwrap(inst.src0)
+    vdst_off = _vi(inst.vdst)
+    if 'read' in name:
+      # v_accvgpr_read_b32 vN, aM - reads from accumulator to VGPR
+      return f"{name}_b32 v{vdst_off}, a{src0_off - 256 if src0_off >= 256 else src0_off}"
+    if 'write' in name:
+      # v_accvgpr_write_b32 aM, src - writes to accumulator from source
+      src = _lit(inst, inst.src0) if src0_off == 255 else (f"v{src0_off - 256}" if src0_off >= 256 else decode_src(src0_off, cdna=True))
+      return f"{name}_b32 a{vdst_off}, {src}"
+
+  # Handle v_mfma_ld_scale_b32 - special 2-operand format: v_mfma_ld_scale_b32 src0, src1
+  if 'ld_scale' in name:
+    src0, src1 = get_src(inst.src0, 1), get_src(inst.src1, 1)
+    mods = ([_fmt_bits("op_sel", inst.opsel, 2)] if inst.opsel else []) + \
+           ([_fmt_bits("op_sel_hi", inst.opsel_hi, 2)] if inst.opsel_hi != 3 else [])
+    return f"{name} {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
+
+  # Handle MFMA instructions with accumulator destinations
+  if is_mfma:
+    regs = inst.canonical_op_regs
+    dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
+    # Infer register counts from instruction name if not in operands table (e.g., v_mfma_f32_32x32x4_xf32)
+    if dregs == 1:
+      if '32x32' in name: dregs, r0, r1, r2 = 16, 2, 2, 16
+      elif '16x16' in name: dregs, r0, r1, r2 = 4, 2, 2, 4
+    # MFMA reuses VOP3P fields differently: clmp -> acc_cd (dest is acc), opsel_hi -> acc (src1/src2 are acc)
+    # acc field (bits 60-59): bit 0 = src2 is acc (always for MFMA), bit 1 = src1 is acc
+    acc = inst.opsel_hi  # opsel_hi field maps to acc for MFMA
+    acc_cd = inst.clmp   # clmp field maps to acc_cd for MFMA (dest is accumulator)
+    is_smfmac = 'smfmac' in name  # SMFMAC has different operand semantics
+    # Format sources: src0 is always VGPR, src1/src2 depend on acc bits
+    def mfma_src(v, sc, is_acc):
+      v = _unwrap(v)
+      if v == 255: return _lit(inst, v)
+      if 128 <= v <= 208 or 240 <= v <= 248: return _lit(inst, v)
+      base = v - 256 if v >= 256 else v
+      if is_acc: return _areg(base, sc)
+      return _vreg(base, sc)
+    src0 = get_src(inst.src0, r0)  # src0 is always VGPR
+    src1 = mfma_src(inst.src1, r1, acc & 2)  # bit 1 = src1 is acc
+    # For SMFMAC, src2 is always a VGPR index (1 register), not accumulator
+    src2 = _vreg(inst.src2) if is_smfmac else mfma_src(inst.src2, r2, acc_cd)
+    dst = _areg(inst.vdst, dregs) if acc_cd else _vreg(inst.vdst, dregs)
+    # MFMA uses neg:[...] not neg_lo:[...], and doesn't support op_sel_hi or clamp
+    # Only f64 MFMA instructions support neg modifier
+    # f8f6f4 MFMA instructions support cbsz/blgp modifiers
+    mods = []
+    if 'f8f6f4' in name:
+      if inst.neg_hi: mods.append(f"cbsz:{inst.neg_hi}")
+      if inst.neg: mods.append(f"blgp:{inst.neg}")
+    elif inst.neg and 'f64' in name:
+      mods.append(_fmt_bits("neg", inst.neg, n))
+    return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}"
+
+  # Standard VOP3P instructions
+  src0, src1, src2, dst = get_src(inst.src0, 1), get_src(inst.src1, 1), get_src(inst.src2, 1), _vreg(inst.vdst)
+  opsel_hi = inst.opsel_hi  # CDNA VOP3P only has 2 bits for opsel_hi (no opsel_hi2)
+  opsel_hi_default = 3  # CDNA default is 0b11 (2 bits), not 0b111 like RDNA
+  mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default 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_mubuf(inst) -> str:
+  name = inst.op_name.lower()
+  # Determine vdata register count from instruction name
+  nregs = 4 if 'xyzw' in name else 3 if 'xyz' in name else 2 if 'xy' in name or 'x2' in name or 'f64' in name or 'dwordx2' in name else 1
+  vdata = _vreg(inst.vdata, nregs)
+  vaddr = _vreg(inst.vaddr) if inst.offen or inst.idxen else None
+  srsrc = str(inst.srsrc)
+  soffset_val = _unwrap(inst.soffset)
+  soffset = f"s{soffset_val}" if soffset_val < 128 else "off"
+  offset = f" offset:{inst.offset}" if inst.offset else ""
+  offen = " offen" if inst.offen else ""
+  idxen = " idxen" if inst.idxen else ""
+  lds = " lds" if inst.lds else ""
+  sc0 = " sc0" if inst.sc0 else ""
+  sc1 = " sc1" if inst.sc1 else ""
+  nt = " nt" if inst.nt else ""
+  # Handle special cases
+  if name in ('buffer_wbl2', 'buffer_inv'):
+    return f"{name}{sc0}{sc1}"
+  if vaddr:
+    return f"{name} {vdata}, {vaddr}, {srsrc}, {soffset}{offen}{idxen}{offset}{sc0}{nt}{sc1}{lds}"
+  return f"{name} {vdata}, off, {srsrc}, {soffset}{offset}{sc0}{nt}{sc1}{lds}"
+
+_SDWA_SEL = {0: 'BYTE_0', 1: 'BYTE_1', 2: 'BYTE_2', 3: 'BYTE_3', 4: 'WORD_0', 5: 'WORD_1', 6: 'DWORD'}
+
+def _disasm_vop1_sdwa(inst) -> str:
+  name = inst.op_name.lower().replace('_e32', '')
+  regs = inst.canonical_op_regs
+  dst = _vreg(inst.vdst, regs['d'])
+  # When s0=1, vsrc0 is SGPR/constant (VGPRField adds 256, so subtract it back)
+  if inst.s0 == 0: src0 = _vreg(inst.vsrc0, regs['s0'])
+  else:
+    raw = _unwrap(inst.vsrc0) - 256  # VGPRField adds 256
+    src0 = decode_src(raw, cdna=True)  # handles SGPRs, constants, specials
+  src0_sel = _SDWA_SEL.get(inst.src0_sel, f'SEL{inst.src0_sel}')
+  mods = []
+  if inst.clmp: mods.append("clamp")
+  if inst.omod == 1: mods.append("mul:2")
+  elif inst.omod == 2: mods.append("mul:4")
+  elif inst.omod == 3: mods.append("div:2")
+  mods.append(f"src0_sel:{src0_sel}")
+  return f"{name}_sdwa {dst}, {src0} {' '.join(mods)}"
+
+def _decode_dpp(dpp: int) -> str:
+  """Decode DPP control value to string."""
+  if dpp < 0x100: return f"quad_perm:[{dpp&3},{(dpp>>2)&3},{(dpp>>4)&3},{(dpp>>6)&3}]"
+  if 0x100 <= dpp <= 0x10f: return f"row_shl:{dpp & 0xf}"
+  if 0x110 <= dpp <= 0x11f: return f"row_shr:{dpp & 0xf}"
+  if 0x120 <= dpp <= 0x12f: return f"row_ror:{dpp & 0xf}"
+  if dpp == 0x130: return "wave_shl:1"
+  if dpp == 0x134: return "wave_rol:1"
+  if dpp == 0x138: return "wave_shr:1"
+  if dpp == 0x13c: return "wave_ror:1"
+  if dpp == 0x140: return "row_mirror"
+  if dpp == 0x141: return "row_half_mirror"
+  if dpp == 0x142: return "row_bcast:15"
+  if dpp == 0x143: return "row_bcast:31"
+  if 0x150 <= dpp <= 0x15f: return f"row_newbcast:{dpp & 0xf}"
+  if 0x160 <= dpp <= 0x16f: return f"row_share:{dpp & 0xf}"
+  if 0x170 <= dpp <= 0x17f: return f"row_xmask:{dpp & 0xf}"
+  return f"dpp:{dpp:#x}"
+
+def _disasm_vop1_dpp(inst) -> str:
+  name = inst.op_name.lower().replace('_e32', '')
+  regs = inst.canonical_op_regs
+  dst, src0 = _vreg(inst.vdst, regs['d']), _vreg(inst.vsrc0, regs['s0'])
+  dpp_str = _decode_dpp(inst.dpp)
+  mods = [dpp_str]
+  if inst.row_mask != 0xf: mods.append(f"row_mask:{inst.row_mask:#x}")
+  if inst.bank_mask != 0xf: mods.append(f"bank_mask:{inst.bank_mask:#x}")
+  if inst.bc: mods.append("bound_ctrl:1")
+  return f"{name}_dpp {dst}, {src0} {' '.join(mods)}"
+
+def _disasm_vop2_sdwa(inst) -> str:
+  name, cdna = inst.op_name.lower().replace('_e32', ''), _is_cdna(inst)
+  regs = inst.canonical_op_regs
+  dst = _vreg(inst.vdst, regs['d'])
+  # When s0/s1=1, vsrc is SGPR/constant (VGPRField adds 256, so subtract it back)
+  src0 = _vreg(inst.vsrc0, regs['s0']) if inst.s0 == 0 else decode_src(_unwrap(inst.vsrc0) - 256, cdna)
+  src1 = _vreg(inst.vsrc1, regs['s1']) if inst.s1 == 0 else decode_src(_unwrap(inst.vsrc1) - 256, cdna)
+  src0_sel = _SDWA_SEL.get(inst.src0_sel, f'SEL{inst.src0_sel}')
+  src1_sel = _SDWA_SEL.get(inst.src1_sel, f'SEL{inst.src1_sel}')
+  mods = []
+  if inst.clmp: mods.append("clamp")
+  if inst.omod == 1: mods.append("mul:2")
+  elif inst.omod == 2: mods.append("mul:4")
+  elif inst.omod == 3: mods.append("div:2")
+  if inst.src0_sel != 6: mods.append(f"src0_sel:{src0_sel}")
+  if inst.src1_sel != 6: mods.append(f"src1_sel:{src1_sel}")
+  mods_str = ' '.join(mods) if mods else ""
+  # CDNA carry instructions and cndmask need vcc operands
+  if cdna and name in _VOP2_CARRY_OUT: return f"{name}_sdwa {dst}, vcc, {src0}, {src1} {mods_str}".strip()
+  if cdna and name in _VOP2_CARRY_INOUT: return f"{name}_sdwa {dst}, vcc, {src0}, {src1}, vcc {mods_str}".strip()
+  if cdna and name == 'v_cndmask_b32': return f"{name}_sdwa {dst}, {src0}, {src1}, vcc {mods_str}".strip()
+  return f"{name}_sdwa {dst}, {src0}, {src1} {mods_str}".strip()
+
+def _disasm_vop2_dpp(inst) -> str:
+  name, cdna = inst.op_name.lower().replace('_e32', ''), _is_cdna(inst)
+  regs = inst.canonical_op_regs
+  dst, src0, src1 = _vreg(inst.vdst, regs['d']), _vreg(inst.vsrc0, regs['s0']), _vreg(inst.vsrc1, regs['s1'])
+  dpp_str = _decode_dpp(inst.dpp)
+  mods = [dpp_str]
+  if inst.row_mask != 0xf: mods.append(f"row_mask:{inst.row_mask:#x}")
+  if inst.bank_mask != 0xf: mods.append(f"bank_mask:{inst.bank_mask:#x}")
+  if inst.bc: mods.append("bound_ctrl:1")
+  # CDNA carry instructions and cndmask need vcc operands
+  if cdna and name in _VOP2_CARRY_OUT: return f"{name}_dpp {dst}, vcc, {src0}, {src1} {' '.join(mods)}"
+  if cdna and name in _VOP2_CARRY_INOUT: return f"{name}_dpp {dst}, vcc, {src0}, {src1}, vcc {' '.join(mods)}"
+  if cdna and name == 'v_cndmask_b32': return f"{name}_dpp {dst}, {src0}, {src1}, vcc {' '.join(mods)}"
+  return f"{name}_dpp {dst}, {src0}, {src1} {' '.join(mods)}"
+
+def _disasm_vopc_sdwa(inst) -> str:
+  name = inst.op_name.lower().replace('_e32', '')
+  regs = inst.canonical_op_regs
+  sdst = _fmt_sdst(inst.sdst, 2, cdna=True)
+  src0 = _vreg(inst.vsrc0, regs['s0']) if getattr(inst, 's0', 0) == 0 else decode_src(_unwrap(inst.vsrc0) - 256, cdna=True)
+  src1 = _vreg(inst.vsrc1, regs['s1']) if getattr(inst, 's1', 0) == 0 else decode_src(_unwrap(inst.vsrc1) - 256, cdna=True)
+  src0_sel = _SDWA_SEL.get(inst.src0_sel, f'SEL{inst.src0_sel}')
+  src1_sel = _SDWA_SEL.get(inst.src1_sel, f'SEL{inst.src1_sel}')
+  mods = []
+  if inst.src0_sel != 6: mods.append(f"src0_sel:{src0_sel}")
+  if inst.src1_sel != 6: mods.append(f"src1_sel:{src1_sel}")
+  return f"{name}_sdwa {sdst}, {src0}, {src1} {' '.join(mods)}".strip()
+
+def _disasm_vop3px2(inst) -> str:
+  """VOP3PX2 disassembler for scaled MFMA instructions."""
+  name = inst.op_name.lower()
+  regs = inst.canonical_op_regs
+  dregs, r2 = regs['d'], regs['s2']
+  # F8F6F4 MFMA: CBSZ selects matrix A format, BLGP selects matrix B format
+  # VGPRs: FP8/BF8(0,1)=8, FP6/BF6(2,3)=6, FP4(4)=4
+  vgprs = {0: 8, 1: 8, 2: 6, 3: 6, 4: 4}
+  r0, r1 = vgprs.get(inst.cbsz, 8), vgprs.get(inst.blgp, 8)
+  def mfma_src(v, sc, is_acc):
+    v = _unwrap(v)
+    if v == 255: return _lit(inst, v)
+    base = v - 256 if v >= 256 else v
+    return _areg(base, sc) if is_acc else _vreg(base, sc)
+  src0, src1, src2 = mfma_src(inst.src0, r0, False), mfma_src(inst.src1, r1, inst.acc & 2), mfma_src(inst.src2, r2, inst.acc_cd)
+  dst = _areg(inst.vdst, dregs) if inst.acc_cd else _vreg(inst.vdst, dregs)
+  scale_src0, scale_src1 = _vreg(inst.scale_src0), _vreg(inst.scale_src1)
+  mods = []
+  if inst.opsel: mods.append(_fmt_bits("op_sel", inst.opsel, 3))
+  if inst.opsel_hi != 0: mods.append(_fmt_bits("op_sel_hi", inst.opsel_hi, 3))
+  if inst.neg: mods.append(_fmt_bits("neg", inst.neg, 3))
+  if inst.cbsz: mods.append(f"cbsz:{inst.cbsz}")
+  if inst.blgp: mods.append(f"blgp:{inst.blgp}")
+  return f"{name} {dst}, {src0}, {src1}, {src2}, {scale_src0}, {scale_src1}{' ' + ' '.join(mods) if mods else ''}"
+
+DISASM_HANDLERS.update({CDNA_VOP1: _disasm_vop1, CDNA_VOP1_LIT: _disasm_vop1,
+  CDNA_VOP1_SDWA: _disasm_vop1_sdwa, CDNA_VOP1_DPP16: _disasm_vop1_dpp,
+  CDNA_VOP2: _disasm_vop2, CDNA_VOP2_LIT: _disasm_vop2,
+  CDNA_VOP2_SDWA: _disasm_vop2_sdwa, CDNA_VOP2_DPP16: _disasm_vop2_dpp,
+  CDNA_VOPC: _disasm_vopc, CDNA_VOPC_LIT: _disasm_vopc, CDNA_VOPC_SDWA_SDST: _disasm_vopc_sdwa,
+  CDNA_SOP1: _disasm_sop1, CDNA_SOP1_LIT: _disasm_sop1, CDNA_SOP2: _disasm_sop2, CDNA_SOP2_LIT: _disasm_sop2,
+  CDNA_SOPC: _disasm_sopc, CDNA_SOPC_LIT: _disasm_sopc, CDNA_SOPK: _disasm_sopk, CDNA_SOPK_LIT: _disasm_sopk, CDNA_SOPP: _disasm_sopp,
+  CDNA_SMEM: _disasm_smem, CDNA_DS: _disasm_ds, CDNA_FLAT: _disasm_flat, CDNA_GLOBAL: _disasm_flat, CDNA_SCRATCH: _disasm_flat,
+  CDNA_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p,
+  CDNA_MUBUF: _disasm_mubuf, CDNA_VOP3PX2: _disasm_vop3px2})
@@ -0,0 +1,424 @@
+# dsl.py - clean DSL for AMD assembly
+
+# ══════════════════════════════════════════════════════════════
+# Registers - unified src encoding space (0-511)
+# ══════════════════════════════════════════════════════════════
+
+def _reg_size(t: str | None) -> int: return {'b64': 2, 'f64': 2, 'u64': 2, 'i64': 2, 'b128': 4}.get(t, 1)
+
+class Reg:
+  _NAMES = {106: "VCC_LO", 107: "VCC_HI", 124: "NULL", 125: "M0", 126: "EXEC_LO", 127: "EXEC_HI",
+            240: "0.5", 241: "-0.5", 242: "1.0", 243: "-1.0", 244: "2.0", 245: "-2.0", 246: "4.0", 247: "-4.0",
+            248: "INV_2PI", 250: "DPP16", 253: "SCC", 255: "LIT"}
+  _PAIRS = {106: "VCC", 126: "EXEC"}
+
+  def __init__(self, offset: int = 0, sz: int = 512, *, neg: bool = False, abs_: bool = False, hi: bool = False):
+    self.offset, self.sz = offset, sz
+    self.neg, self.abs_, self.hi = neg, abs_, hi
+
+  # TODO: remove these legacy aliases
+  @property
+  def count(self): return self.sz
+  @property
+  def idx(self): return self.offset
+
+  def __hash__(self): return hash((self.offset, self.sz, self.neg, self.abs_, self.hi))
+  def __getitem__(self, key):
+    if isinstance(key, slice):
+      start, stop = key.start or 0, key.stop or (self.sz - 1)
+      if start < 0 or stop >= self.sz: raise RuntimeError(f"slice [{start}:{stop}] out of bounds for size {self.sz}")
+      return Reg(self.offset + start, stop - start + 1)
+    if key < 0 or key >= self.sz: raise RuntimeError(f"index {key} out of bounds for size {self.sz}")
+    return Reg(self.offset + key, 1)
+  def __eq__(self, other):
+    if isinstance(other, Reg):
+      return (self.offset == other.offset and self.sz == other.sz and
+              self.neg == other.neg and self.abs_ == other.abs_ and self.hi == other.hi)
+    return NotImplemented
+  def __add__(self, other):
+    if isinstance(other, int): return Reg(self.offset + other, self.sz)
+    return NotImplemented
+  def __neg__(self) -> 'Reg': return Reg(self.offset, self.sz, neg=not self.neg, abs_=self.abs_, hi=self.hi)
+  def __abs__(self) -> 'Reg': return Reg(self.offset, self.sz, neg=self.neg, abs_=True, hi=self.hi)
+  @property
+  def h(self) -> 'Reg': return Reg(self.offset, self.sz, neg=self.neg, abs_=self.abs_, hi=True)
+  @property
+  def l(self) -> 'Reg': return Reg(self.offset, self.sz, neg=self.neg, abs_=self.abs_, hi=False)
+  def fmt(self, sz=None, parens=False, upper=False) -> str:
+    o, sz = self.offset, sz or self.sz
+    l, r = ("[", "]") if parens or sz > 1 else ("", "")  # brackets for multi-reg or when parens=True
+    if 256 <= o < 512: idx = o - 256; base = f"v{l}{idx}{r}" if sz == 1 else f"v[{idx}:{idx + sz - 1}]"
+    elif o < 106: base = f"s{l}{o}{r}" if sz == 1 else f"s[{o}:{o + sz - 1}]"
+    elif sz == 2 and o in self._PAIRS: base = self._PAIRS[o] if upper else self._PAIRS[o].lower()
+    elif o in self._NAMES: base = self._NAMES[o] if upper else self._NAMES[o].lower()  # special regs (any sz)
+    elif 108 <= o < 124: idx = o - 108; base = f"ttmp{l}{idx}{r}" if sz == 1 else f"ttmp[{idx}:{idx + sz - 1}]"
+    elif 128 <= o <= 192: base = str(o - 128)  # inline int constants (0-64)
+    elif 193 <= o <= 208: base = str(-(o - 192))  # inline negative int constants (-1 to -16)
+    else: raise RuntimeError(f"unknown register: offset={o}, sz={sz}")
+    if self.hi: base += ".h"
+    if self.abs_: base = f"abs({base})" if upper else f"|{base}|"
+    if self.neg: base = f"-{base}"
+    return base
+  def __repr__(self): return self.fmt(parens=True, upper=True)
+
+# Full src encoding space
+src = Reg(0, 512)
+
+# Slices for each region (inclusive end)
+s = src[0:105]           # SGPR0-105
+VCC_LO = src[106]
+VCC_HI = src[107]
+VCC = src[106:107]
+ttmp = src[108:123]      # TTMP0-15
+NULL = OFF = src[124]
+M0 = src[125]
+EXEC_LO = src[126]
+EXEC_HI = src[127]
+EXEC = src[126:127]
+# 128: 0, 129-192: integers 1-64, 193-208: integers -1 to -16
+# 240-248: float constants (0.5, -0.5, 1.0, -1.0, 2.0, -2.0, 4.0, -4.0, 1/(2*PI))
+INV_2PI = src[248]
+DPP16 = src[250]
+SCC = src[253]
+# 255: literal constant
+v = src[256:511]         # VGPR0-255
+
+# ══════════════════════════════════════════════════════════════
+# BitField
+# ══════════════════════════════════════════════════════════════
+
+class _Bits:
+  """Helper for defining bit fields with slice syntax: bits[hi:lo] or bits[n]."""
+  def __getitem__(self, key) -> 'BitField': return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
+bits = _Bits()
+
+class BitField:
+  def __init__(self, hi: int, lo: int, default: int = 0):
+    self.hi, self.lo, self.default, self.name, self.mask = hi, lo, default, None, (1 << (hi - lo + 1)) - 1
+  def __set_name__(self, owner, name): self.name = name
+  def __eq__(self, other) -> 'FixedBitField':
+    if isinstance(other, int): return FixedBitField(self.hi, self.lo, other)
+    return NotImplemented
+  def enum(self, enum_cls) -> 'EnumBitField': return EnumBitField(self.hi, self.lo, enum_cls)
+  def encode(self, val) -> int:
+    assert isinstance(val, int), f"BitField.encode expects int, got {type(val).__name__}"
+    return val
+  def decode(self, val): return val
+  def set(self, raw: int, val) -> int:
+    if val is None: val = self.default
+    encoded = self.encode(val)
+    if encoded < 0 or encoded > self.mask: raise RuntimeError(f"field '{self.name}': value {encoded} doesn't fit in {self.hi - self.lo + 1} bits")
+    return (raw & ~(self.mask << self.lo)) | (encoded << self.lo)
+  def __get__(self, obj, objtype=None):
+    if obj is None: return self
+    return self.decode((obj._raw >> self.lo) & self.mask)
+
+class FixedBitField(BitField):
+  def set(self, raw: int, val=None) -> int:
+    assert val is None, f"FixedBitField does not accept values, got {val}"
+    return super().set(raw, self.default)
+
+class EnumBitField(BitField):
+  def __init__(self, hi: int, lo: int, enum_cls, allowed: set | None = None):
+    super().__init__(hi, lo)
+    self._enum = enum_cls
+    self.allowed = allowed  # if set, only these enum values are valid for this encoding
+  def encode(self, val) -> int:
+    if not isinstance(val, self._enum): raise RuntimeError(f"expected {self._enum.__name__}, got {type(val).__name__}")
+    if self.allowed is not None and val not in self.allowed:
+      raise RuntimeError(f"opcode {val.name} not allowed in this encoding")
+    return val.value
+  def decode(self, raw): return self._enum(raw)
+
+# ══════════════════════════════════════════════════════════════
+# Typed fields
+# ══════════════════════════════════════════════════════════════
+
+import struct
+def _f32(f: float) -> int: return struct.unpack('I', struct.pack('f', f))[0]
+
+class SrcField(BitField):
+  _valid_range = (0, 511)  # inclusive
+  _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}
+
+  def __init__(self, hi: int, lo: int, default=s[0]):
+    super().__init__(hi, lo, default)
+    expected_size = self._valid_range[1] - self._valid_range[0] + 1
+    actual_size = 1 << (hi - lo + 1)
+    if actual_size != expected_size:
+      raise RuntimeError(f"{self.__class__.__name__}: field size {hi - lo + 1} bits ({actual_size}) doesn't match range {self._valid_range} ({expected_size})")
+
+  def encode(self, val) -> int:
+    """Encode value. Returns 255 (literal marker) for out-of-range values."""
+    if isinstance(val, Reg): offset = val.offset
+    elif isinstance(val, float): offset = self._FLOAT_ENC.get(val, 255)
+    elif isinstance(val, int) and 0 <= val <= 64: offset = 128 + val
+    elif isinstance(val, int) and -16 <= val < 0: offset = 192 - val
+    elif isinstance(val, int): offset = 255  # literal
+    else: raise TypeError(f"invalid src value {val}")
+    if not (self._valid_range[0] <= offset <= self._valid_range[1]):
+      raise TypeError(f"{self.__class__.__name__}: {val} (offset {offset}) out of range {self._valid_range}")
+    return offset - self._valid_range[0]
+
+  def decode(self, raw): return src[raw + self._valid_range[0]]
+
+  def __get__(self, obj, objtype=None):
+    if obj is None: return self
+    reg = self.decode((obj._raw >> self.lo) & self.mask)
+    # Resize register based on operand info (skip non-resizable special registers)
+    # VCC/EXEC pairs (106, 126), NULL (124), M0 (125), float constants (240-255)
+    if reg.offset not in (124, 125) and not 240 <= reg.offset <= 255:
+      if sz := obj.op_regs.get(self.name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
+    return reg
+
+class VGPRField(SrcField):
+  _valid_range = (256, 511)
+  def __init__(self, hi: int, lo: int, default=v[0]): super().__init__(hi, lo, default)
+  def encode(self, val) -> int:
+    if not isinstance(val, Reg): raise TypeError(f"VGPRField requires Reg, got {type(val).__name__}")
+    # For 8-bit vdst fields in VOP1/VOP2 16-bit ops, bit 7 is opsel for dest half
+    encoded = super().encode(val)
+    if val.hi and (self.hi - self.lo + 1) == 8:
+      if encoded >= 128:
+        raise ValueError(f"VGPRField: v[{encoded}].h not encodable in 8-bit field (v[0:127] only for .h)")
+      encoded |= 0x80
+    return encoded
+class SGPRField(SrcField): _valid_range = (0, 127)
+class SSrcField(SrcField): _valid_range = (0, 255)
+
+class AlignedSGPRField(BitField):
+  """SGPR field with alignment requirement. Encoded as sgpr_index // alignment."""
+  _align: int = 2
+  def encode(self, val):
+    if isinstance(val, int) and val == 0: return 0  # default: encode as s[0]
+    if not isinstance(val, Reg): raise TypeError(f"{self.__class__.__name__} requires Reg, got {type(val).__name__}")
+    if not (0 <= val.offset < 128): raise ValueError(f"{self.__class__.__name__} requires SGPR, got offset {val.offset}")
+    if val.offset & (self._align - 1): raise ValueError(f"{self.__class__.__name__} requires {self._align}-aligned SGPR, got s[{val.offset}]")
+    return val.offset >> (self._align.bit_length() - 1)
+  def decode(self, raw): return src[raw << (self._align.bit_length() - 1)]
+  def __get__(self, obj, objtype=None):
+    if obj is None: return self
+    reg = self.decode((obj._raw >> self.lo) & self.mask)
+    if sz := obj.op_regs.get(self.name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
+    return reg
+
+class SBaseField(AlignedSGPRField): _align = 2
+class SRsrcField(AlignedSGPRField): _align = 4
+
+class VDSTYField(BitField):
+  """VOPD vdsty: encoded = vgpr_idx >> 1. Actual vgpr = (encoded << 1) | ((vdstx & 1) ^ 1)."""
+  def encode(self, val):
+    if not isinstance(val, Reg): raise TypeError(f"VDSTYField requires Reg, got {type(val).__name__}")
+    if not (256 <= val.offset < 512): raise ValueError(f"VDSTYField requires VGPR, got offset {val.offset}")
+    return (val.offset - 256) >> 1
+  def decode(self, raw): return raw  # raw value, actual vdsty = (raw << 1) | ((vdstx & 1) ^ 1)
+
+# ══════════════════════════════════════════════════════════════
+# Operand info from XML
+# ══════════════════════════════════════════════════════════════
+
+import functools
+from extra.assembly.amd.autogen.rdna3.operands import OPERANDS as OPERANDS_RDNA3
+from extra.assembly.amd.autogen.rdna4.operands import OPERANDS as OPERANDS_RDNA4
+from extra.assembly.amd.autogen.cdna.operands import OPERANDS as OPERANDS_CDNA
+OPERANDS = {**OPERANDS_CDNA, **OPERANDS_RDNA3, **OPERANDS_RDNA4}
+
+# ══════════════════════════════════════════════════════════════
+# Inst base class
+# ══════════════════════════════════════════════════════════════
+
+def _needs_literal(val) -> bool:
+  """Check if a value needs a literal constant (can't be encoded inline)."""
+  if val is None or isinstance(val, Reg): return False
+  if isinstance(val, float): return val not in SrcField._FLOAT_ENC
+  if isinstance(val, int): return not (0 <= val <= 64 or -16 <= val < 0)
+  return False
+
+def _get_variant(cls, suffix: str):
+  """Get a variant class by suffix (e.g., '_LIT') via module lookup."""
+  import sys
+  module = sys.modules.get(cls.__module__)
+  return getattr(module, f"{cls.__name__}{suffix}", None) if module else None
+
+class Inst:
+  _fields: list[tuple[str, BitField]]
+  _base_size: int
+
+  def __init_subclass__(cls):
+    # Collect fields from all parent classes, then override with this class's fields
+    inherited = {}
+    for base in reversed(cls.__mro__[1:]):
+      if hasattr(base, '_fields'):
+        inherited.update({name: field for name, field in base._fields})
+    inherited.update({name: val for name, val in cls.__dict__.items() if isinstance(val, BitField)})
+    cls._fields = list(inherited.items())
+    cls._base_size = (max(f.hi for _, f in cls._fields) + 8) // 8
+
+  def __new__(cls, *args, **kwargs):
+    # Auto-upgrade to _LIT variant if needed (only for base classes, not variants)
+    if not any(cls.__name__.endswith(sfx) for sfx in ('_LIT', '_DPP16', '_DPP8', '_SDWA', '_SDWA_SDST', '_MFMA')):
+      lit_cls = _get_variant(cls, '_LIT')
+      if lit_cls is not None:
+        # Check if any src field needs a literal
+        # Map positional args to field names to find src values
+        args_iter = iter(args)
+        for name, field in cls._fields:
+          if isinstance(field, FixedBitField): continue
+          val = kwargs.get(name) if name in kwargs else next(args_iter, None)
+          if isinstance(field, SrcField) and _needs_literal(val):
+            return lit_cls(*args, **kwargs)
+    return object.__new__(cls)
+
+  def __init__(self, *args, **kwargs):
+    self._raw = 0
+    # Map positional args to field names (skip FixedBitFields)
+    args_iter = iter(args)
+    vals = {}
+    for name, field in self._fields:
+      if isinstance(field, FixedBitField): vals[name] = None
+      elif name in kwargs: vals[name] = kwargs[name]
+      else: vals[name] = next(args_iter, None)
+    remaining = list(args_iter)
+    assert not remaining, f"too many positional args: {remaining}"
+    # Extract modifiers from Reg objects and merge into neg/abs/opsel
+    neg_bits, abs_bits, opsel_bits = 0, 0, 0
+    for name, bit in [('src0', 0), ('src1', 1), ('src2', 2)]:
+      if name in vals and isinstance(vals[name], Reg):
+        reg = vals[name]
+        if reg.neg: neg_bits |= (1 << bit)
+        if reg.abs_: abs_bits |= (1 << bit)
+        if reg.hi: opsel_bits |= (1 << bit)
+    if 'vdst' in vals and isinstance(vals['vdst'], Reg) and vals['vdst'].hi:
+      opsel_bits |= (1 << 3)
+    if neg_bits: vals['neg'] = (vals.get('neg') or 0) | neg_bits
+    if abs_bits: vals['abs'] = (vals.get('abs') or 0) | abs_bits
+    if opsel_bits: vals['opsel'] = (vals.get('opsel') or 0) | opsel_bits
+    # For _LIT classes, capture literal value from SrcFields that encode to 255
+    literal_val = None
+    for name, field in self._fields:
+      val = vals[name]
+      if isinstance(field, SrcField) and val is not None and _needs_literal(val):
+        literal_val = _f32(val) if isinstance(val, float) else val & 0xFFFFFFFF
+    if literal_val is not None and 'literal' in vals:
+      vals['literal'] = literal_val
+    # Set all field values
+    for name, field in self._fields:
+      self._raw = field.set(self._raw, vals[name])
+    # Validate register sizes against operand info (skip special registers like NULL, VCC, EXEC)
+    for name, expected in self.op_regs.items():
+      if (val := vals.get(name)) is None: continue
+      if isinstance(val, Reg) and val.sz != expected and not (106 <= val.offset <= 127 or val.offset == 253):
+        raise TypeError(f"{name} expects {expected} register(s), got {val.sz}")
+
+  @property
+  def op_name(self) -> str: return self.op.name
+  @property
+  def operands(self) -> dict: return OPERANDS.get(self.op, {}) if hasattr(self, 'op') else {}
+  def _is_cdna(self) -> bool: return 'cdna' in type(self).__module__
+
+  @functools.cached_property
+  def op_bits(self) -> dict[str, int]:
+    """Get bit widths for each operand field, with WAVE32 and addr/saddr adjustments."""
+    if not hasattr(self, 'op'): return {k: v[1] for k, v in self.operands.items()}
+    bits = {k: v[1] for k, v in self.operands.items()}
+    # RDNA (WAVE32): condition masks, carry flags, and compare results are 32-bit
+    if not self._is_cdna():
+      name = self.op_name.lower()
+      if 'cndmask' in name and 'src2' in bits: bits['src2'] = 32
+      if '_co_ci_' in name:
+        if 'src2' in bits: bits['src2'] = 32
+        if 'sdst' in bits: bits['sdst'] = 32
+      if 'cmp' in name and 'vdst' in bits: bits['vdst'] = 32
+    # GLOBAL/FLAT: addr is 32-bit if saddr is valid SGPR, 64-bit if saddr is NULL
+    # SCRATCH: addr is always 32-bit (offset from scratch base, not absolute address)
+    if 'addr' in bits and (saddr_field := getattr(type(self), 'saddr', None)) and type(self).__name__ not in ('SCRATCH', 'VSCRATCH'):
+      saddr_val = (self._raw >> saddr_field.lo) & saddr_field.mask  # access _raw directly to avoid recursion
+      bits['addr'] = 64 if saddr_val in (124, 125) else 32  # 124=NULL, 125=M0
+    # MUBUF/MTBUF: vaddr size depends on offen/idxen (1 or 2 regs)
+    if 'vaddr' in bits and hasattr(self, 'offen') and hasattr(self, 'idxen'):
+      bits['vaddr'] = max(1, self.offen + self.idxen) * 32
+    # F8F6F4 MFMA: CBSZ selects matrix A format, BLGP selects matrix B format
+    # VGPRs: FP8/BF8(0,1)=8, FP6/BF6(2,3)=6, FP4(4)=4
+    if 'f8f6f4' in getattr(self, 'op_name', '').lower():
+      # Use explicit fields if available (VOP3PX2), else extract from VOP3P-MAI bit positions
+      cbsz = self.cbsz if hasattr(type(self), 'cbsz') else (self._raw >> 8) & 0x7
+      blgp = self.blgp if hasattr(type(self), 'blgp') else (self._raw >> 61) & 0x7
+      vgprs = {0: 8, 1: 8, 2: 6, 3: 6, 4: 4}
+      bits['src0'], bits['src1'] = vgprs.get(cbsz, 8) * 32, vgprs.get(blgp, 8) * 32
+    return bits
+  @property
+  def op_regs(self) -> dict[str, int]:
+    """Get register counts for each operand field."""
+    return {k: max(1, v // 32) for k, v in self.op_bits.items()}
+
+  @functools.cached_property
+  def canonical_op_bits(self) -> dict[str, int]:
+    """Get bit widths with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
+    bits = {'d': 32, 's0': 32, 's1': 32, 's2': 32, 'data': 32}
+    for name, val in self.op_bits.items():
+      if name in ('src0', 'vsrc0', 'ssrc0'): bits['s0'] = val
+      elif name in ('src1', 'vsrc1', 'ssrc1'): bits['s1'] = val
+      elif name == 'src2': bits['s2'] = val
+      elif name in ('vdst', 'sdst', 'sdata'): bits['d'] = val
+      elif name in ('data', 'vdata', 'data0'): bits['data'] = val
+    return bits
+  @property
+  def canonical_op_regs(self) -> dict[str, int]:
+    """Get register counts with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
+    return {k: max(1, v // 32) for k, v in self.canonical_op_bits.items()}
+
+  def num_srcs(self) -> int:
+    """Get number of source operands from operand info."""
+    ops = self.operands
+    if 'src2' in ops: return 3
+    if 'src1' in ops or 'vsrc1' in ops or 'ssrc1' in ops: return 2
+    if 'src0' in ops or 'vsrc0' in ops or 'ssrc0' in ops: return 1
+    return 0
+  @classmethod
+  def _size(cls) -> int: return cls._base_size
+  def size(self) -> int: return self._base_size
+  def disasm(self) -> str:
+    from extra.assembly.amd.disasm import disasm
+    return disasm(self)
+
+  def to_bytes(self) -> bytes: return self._raw.to_bytes(self._base_size, 'little')
+
+  @property
+  def _literal(self) -> int | None:
+    """Get the literal value if this instruction has one."""
+    return getattr(self, 'literal', None)
+
+  def _variant_suffix(self) -> str | None:
+    """Check if instruction needs a variant class (_LIT, _DPP8, _DPP16, _SDWA). Returns suffix or None."""
+    cls_name = type(self).__name__
+    # Don't check for variants if we're already a variant class
+    if any(s in cls_name for s in ('_LIT', '_DPP8', '_DPP16', '_SDWA')): return None
+    # VOPD: FMAMK/FMAAK opcodes always require literal (check by name since enum may differ across archs)
+    for name in ('opx', 'opy'):
+      if hasattr(self, name) and any(x in getattr(self, name).name for x in ('FMAMK', 'FMAAK')): return '_LIT'
+    for name, field in self._fields:
+      if isinstance(field, SrcField):
+        off = getattr(self, name).offset
+        if off == 255: return '_LIT'
+        if off == 249: return '_SDWA' if self._is_cdna() else '_DPP8'
+        if off == 250: return '_DPP16'
+    return None
+
+  @classmethod
+  def from_bytes(cls, data: bytes):
+    inst = object.__new__(cls)
+    inst._raw = int.from_bytes(data[:cls._base_size], 'little')
+    # Upgrade to variant class if needed (_LIT, _DPP8, _DPP16, _SDWA)
+    if (suffix := inst._variant_suffix()) and (var_cls := _get_variant(cls, suffix)) is not None:
+      return var_cls.from_bytes(data)
+    return inst
+
+  def __eq__(self, other): return type(self) is type(other) and self._raw == other._raw
+  def __hash__(self): return hash((type(self), self._raw))
+
+  def __repr__(self):
+    # collect (repr, is_default) pairs, strip trailing defaults so repr roundtrips with eval
+    name = self.op.name.lower() if hasattr(self, 'op') else type(self).__name__
+    parts = [(repr(v := getattr(self, n)), v == f.default) for n, f in self._fields if n != 'op' and not isinstance(f, FixedBitField)]
+    while parts and parts[-1][1]: parts.pop()
+    return f"{name}({', '.join(p[0] for p in parts)})"
@@ -0,0 +1,479 @@
+# RDNA3 emulator - executes compiled pseudocode from AMD ISA PDF
+# mypy: ignore-errors
+from __future__ import annotations
+import ctypes, functools
+from enum import IntEnum
+from tinygrad.runtime.autogen import hsa
+from extra.assembly.amd.dsl import Inst, NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, v, s
+from extra.assembly.amd.pcode import _f32, _i32, _sext, _f16, _i16, _f64, _i64
+from extra.assembly.amd.decode import decode_inst
+from extra.assembly.amd.pcode import compile_pseudocode
+from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
+from extra.assembly.amd.autogen.rdna3.ins import (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, GLOBAL, SCRATCH, VOPD,
+  SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, SCRATCHOp, VOPDOp)
+
+# Constants and helpers defined locally (not imported from dsl.py)
+MASK32, MASK64 = 0xFFFFFFFF, 0xFFFFFFFFFFFFFFFF
+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}
+
+class SGPRArray:
+  """SGPR array indexed by Reg or int."""
+  __slots__ = ('_data',)
+  def __init__(self, size: int): self._data = [0] * size
+  def __getitem__(self, key): return self._data[getattr(key, 'offset', key)]
+  def __setitem__(self, key, val): self._data[getattr(key, 'offset', key)] = val
+  def __len__(self): return len(self._data)
+  def __iter__(self): return iter(self._data)
+
+class VGPRLane:
+  """Single lane of VGPRs indexed by Reg (offset 256-511) or int (0-255)."""
+  __slots__ = ('_data',)
+  def __init__(self, size: int): self._data = [0] * size
+  def __getitem__(self, key):
+    i = getattr(key, 'offset', key)
+    return self._data[i - 256 if i >= 256 else i]
+  def __setitem__(self, key, val):
+    i = getattr(key, 'offset', key)
+    self._data[i - 256 if i >= 256 else i] = val
+  def __len__(self): return len(self._data)
+  def __iter__(self): return iter(self._data)
+
+WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
+
+# 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) -> bool: return src.offset >= 256 and ((src.offset - 256) & 0x80) != 0
+def _vgpr_masked(src): return v[(src.offset - 256) & 0x7f] if src.offset >= 256 else src
+
+# VOP3 source modifier: apply abs/neg to value
+def _mod_src(val: int, idx: int, neg: int, abs_: int, is64: bool = 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
+
+# Read source operand with VOP3 modifiers
+def _read_src(st, inst, src, idx: int, lane: int, neg: int, abs_: int, opsel: int) -> int:
+  if src is None: return 0
+  src_off = src.offset
+  src_bits = inst.canonical_op_bits[f's{idx}']
+  literal, is_src_64, is_src_16 = inst._literal, src_bits == 64, src_bits == 16
+  if is_src_64: return _mod_src(st.rsrc64(src, lane, literal), idx, neg, abs_, is64=True)
+  if isinstance(inst, VOP3P):
+    opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
+    if 'FMA_MIX' in inst.op_name:
+      raw = st.rsrc(src, lane, literal)
+      sign_bit = (15 if not (opsel & (1 << idx)) else 31) if (opsel_hi >> idx) & 1 else 31
+      if inst.neg_hi & (1 << idx): raw &= ~(1 << sign_bit)
+      if neg & (1 << idx): raw ^= (1 << sign_bit)
+      return raw
+    raw = st.rsrc_f16(src, lane, literal)
+    hi = _src16(raw, opsel_hi & (1 << idx)) ^ (0x8000 if inst.neg_hi & (1 << idx) else 0)
+    lo = _src16(raw, opsel & (1 << idx)) ^ (0x8000 if neg & (1 << idx) else 0)
+    return (hi << 16) | lo
+  if is_src_16 and isinstance(inst, VOP3):
+    raw = st.rsrc_f16(src, lane, literal) if 128 <= src_off < 255 else st.rsrc(src, lane, literal)
+    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_off >= 256: return _src16(_mod_src(st.rsrc(_vgpr_masked(src), lane, literal), idx, neg, abs_), _vgpr_hi(src))
+    return _mod_src(st.rsrc_f16(src, lane, literal), idx, neg, abs_) & 0xffff
+  return _mod_src(st.rsrc(src, lane, literal), idx, neg, abs_)
+
+# Helper: get number of dwords from memory op name
+def _op_ndwords(name: str) -> int:
+  if '_B128' in name: return 4
+  if '_B96' in name: return 3
+  if any(s in name for s in ('_B64', '_U64', '_I64', '_F64')): return 2
+  return 1
+
+# Helper: build multi-dword int from consecutive VGPRs
+def _vgpr_read(V: VGPRLane, reg, ndwords: int) -> int:
+  return sum(V[reg + i] << (32 * i) for i in range(ndwords))
+
+# Helper: write multi-dword value to consecutive VGPRs
+def _vgpr_write(V: VGPRLane, reg, val: int, ndwords: int):
+  for i in range(ndwords): V[reg + i] = (val >> (32 * i)) & MASK32
+
+# 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_uint64 if size == 8 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
+
+def _make_mem_accessor(read_fn, write_fn):
+  """Create a memory accessor class with the given read/write functions."""
+  class _MemAccessor:
+    __slots__ = ('_addr',)
+    def __init__(self, addr: int): self._addr = int(addr)
+    u8 = property(lambda s: read_fn(s._addr, 1), lambda s, v: write_fn(s._addr, 1, int(v)))
+    u16 = property(lambda s: read_fn(s._addr, 2), lambda s, v: write_fn(s._addr, 2, int(v)))
+    u32 = property(lambda s: read_fn(s._addr, 4), lambda s, v: write_fn(s._addr, 4, int(v)))
+    u64 = property(lambda s: read_fn(s._addr, 8), lambda s, v: write_fn(s._addr, 8, int(v)))
+    i8 = property(lambda s: _sext(read_fn(s._addr, 1), 8), lambda s, v: write_fn(s._addr, 1, int(v)))
+    i16 = property(lambda s: _sext(read_fn(s._addr, 2), 16), lambda s, v: write_fn(s._addr, 2, int(v)))
+    i32 = property(lambda s: _sext(read_fn(s._addr, 4), 32), lambda s, v: write_fn(s._addr, 4, int(v)))
+    i64 = property(lambda s: _sext(read_fn(s._addr, 8), 64), lambda s, v: write_fn(s._addr, 8, int(v)))
+    b8, b16, b32, b64 = u8, u16, u32, u64
+  return _MemAccessor
+
+_GlobalMemAccessor = _make_mem_accessor(mem_read, mem_write)
+
+class _GlobalMem:
+  """Global memory wrapper that supports MEM[addr].u32 style access."""
+  def __getitem__(self, addr) -> _GlobalMemAccessor: return _GlobalMemAccessor(addr)
+GlobalMem = _GlobalMem()
+
+class LDSMem:
+  """LDS memory wrapper that supports MEM[addr].u32 style access."""
+  __slots__ = ('_lds',)
+  def __init__(self, lds: bytearray): self._lds = lds
+  def _read(self, addr: int, size: int) -> int:
+    addr = addr & 0xffff
+    return int.from_bytes(self._lds[addr:addr+size], 'little') if addr + size <= len(self._lds) else 0
+  def _write(self, addr: int, size: int, val: int):
+    addr = addr & 0xffff
+    if addr + size <= len(self._lds): self._lds[addr:addr+size] = (int(val) & ((1 << (size*8)) - 1)).to_bytes(size, 'little')
+  def __getitem__(self, addr): return _make_mem_accessor(self._read, self._write)(addr)
+
+# SMEM dst register count (for writing result back to SGPRs)
+SMEM_DST_COUNT = {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_E64, VOPDOp.V_DUAL_FMAAK_F32: VOP2Op.V_FMAAK_F32_E32, VOPDOp.V_DUAL_FMAMK_F32: VOP2Op.V_FMAMK_F32_E32,
+  VOPDOp.V_DUAL_MUL_F32: VOP3Op.V_MUL_F32_E64, VOPDOp.V_DUAL_ADD_F32: VOP3Op.V_ADD_F32_E64, VOPDOp.V_DUAL_SUB_F32: VOP3Op.V_SUB_F32_E64,
+  VOPDOp.V_DUAL_SUBREV_F32: VOP3Op.V_SUBREV_F32_E64, VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: VOP3Op.V_MUL_DX9_ZERO_F32_E64,
+  VOPDOp.V_DUAL_MOV_B32: VOP3Op.V_MOV_B32_E64, VOPDOp.V_DUAL_CNDMASK_B32: VOP3Op.V_CNDMASK_B32_E64,
+  VOPDOp.V_DUAL_MAX_F32: VOP3Op.V_MAX_F32_E64, VOPDOp.V_DUAL_MIN_F32: VOP3Op.V_MIN_F32_E64,
+  VOPDOp.V_DUAL_ADD_NC_U32: VOP3Op.V_ADD_NC_U32_E64, VOPDOp.V_DUAL_LSHLREV_B32: VOP3Op.V_LSHLREV_B32_E64, VOPDOp.V_DUAL_AND_B32: VOP3Op.V_AND_B32_E64,
+}
+
+
+class WaveState:
+  __slots__ = ('sgpr', 'vgpr', 'scc', 'pc', '_pend_sgpr', 'lds', 'n_lanes')
+  def __init__(self, lds: LDSMem | None = None, n_lanes: int = WAVE_SIZE):
+    self.sgpr, self.vgpr = SGPRArray(SGPR_COUNT), [VGPRLane(VGPR_COUNT) for _ in range(WAVE_SIZE)]
+    self.sgpr[EXEC_LO], self.scc, self.pc, self._pend_sgpr, self.lds, self.n_lanes = 0xffffffff, 0, 0, {}, lds, n_lanes
+
+  @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, reg) -> int:
+    if reg == NULL: return 0
+    if reg == SCC: return self.scc
+    return self.sgpr[reg]
+  def wsgpr(self, reg, v: int):
+    if reg != NULL: self.sgpr[reg] = v & MASK32
+  def rsgpr64(self, reg) -> int:
+    off = reg.offset
+    return self.sgpr._data[off] | (self.sgpr._data[off + 1] << 32)
+  def wsgpr64(self, reg, v: int):
+    off = reg.offset
+    self.sgpr._data[off] = v & MASK32; self.sgpr._data[off + 1] = (v >> 32) & MASK32
+
+  def _rsrc_base(self, reg, lane: int, consts, literal: int):
+    off = reg.offset
+    if off < SGPR_COUNT: return self.sgpr._data[off]
+    if off == SCC.offset: return self.scc
+    if off < 255: return consts[off - 128]
+    if off == 255: return literal
+    return self.vgpr[lane]._data[off - 256] if off <= 511 else 0
+  def rsrc(self, reg, lane: int, literal: int = 0) -> int: return self._rsrc_base(reg, lane, _INLINE_CONSTS, literal)
+  def rsrc_f16(self, reg, lane: int, literal: int = 0) -> int: return self._rsrc_base(reg, lane, _INLINE_CONSTS_F16, literal)
+  def rsrc64(self, reg, lane: int, literal: int = 0) -> int:
+    off = reg.offset
+    if 128 <= off < 255: return _INLINE_CONSTS_F64[off - 128]
+    if off == 255: return literal << 32  # 32-bit literal forms upper 32 bits of 64-bit value
+    return self.rsrc(reg, lane, literal) | ((self.rsrc(reg + 1, lane, literal) if off < VCC_LO.offset or 256 <= off <= 511 else 0) << 32)
+
+  def pend_sgpr_lane(self, reg, 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()
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# EXECUTION - All ops use pseudocode from PDF
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def exec_scalar(st: WaveState, inst: Inst):
+  """Execute scalar instruction. Returns 0 to continue execution."""
+  # Get op enum and lookup compiled function
+  if isinstance(inst, SMEM): ssrc0, sdst = None, None
+  elif 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)}")
+
+  # SMEM: memory loads
+  if isinstance(inst, SMEM):
+    addr = st.rsgpr64(inst.sbase) + _sext(inst.offset, 21)
+    if inst.soffset != NULL: addr += st.rsrc(inst.soffset, 0, inst._literal)
+    result = inst._fn(GlobalMem, addr & MASK64)
+    if 'SDATA' in result:
+      sdata = result['SDATA']
+      for i in range(SMEM_DST_COUNT.get(inst.op, 1)): st.wsgpr(inst.sdata + i, (sdata >> (i * 32)) & MASK32)
+    st.pc += inst._words
+    return 0
+
+  # Build context - use canonical_op_bits to determine operand sizes
+  literal = inst._literal
+  s0 = st.rsrc64(ssrc0, 0, literal) if inst.canonical_op_bits['s0'] == 64 else (st.rsrc(ssrc0, 0, literal) if not isinstance(inst, (SOPK, SOPP)) else (st.rsgpr(inst.sdst) if isinstance(inst, SOPK) else 0))
+  s1 = st.rsrc64(inst.ssrc1, 0, literal) if inst.canonical_op_bits['s1'] == 64 else (st.rsrc(inst.ssrc1, 0, literal) if isinstance(inst, (SOP2, SOPC)) else inst.simm16 if isinstance(inst, SOPK) else 0)
+  d0 = st.rsgpr64(sdst) if inst.canonical_op_bits['d'] == 64 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 inst._literal
+
+  # Call compiled function with int parameters
+  result = inst._fn(s0, s1, 0, d0, st.scc, st.vcc & MASK32, 0, st.exec_mask & MASK32, literal, None, pc=st.pc * 4)
+
+  # Apply results (already int values)
+  if sdst is not None and 'D0' in result:
+    (st.wsgpr64 if inst.canonical_op_bits['d'] == 64 else st.wsgpr)(sdst, result['D0'])
+  if 'SCC' in result: st.scc = result['SCC'] & 1
+  if 'EXEC' in result: st.exec_mask = result['EXEC']
+  if 'PC' in result:
+    # Convert absolute byte address to word offset
+    pc_val = result['PC']
+    new_pc = pc_val if pc_val < 0x8000000000000000 else pc_val - 0x10000000000000000
+    st.pc = new_pc // 4
+  else:
+    st.pc += inst._words
+  return 0
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# VECTOR INSTRUCTIONS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def exec_vopd(st: WaveState, inst, V: VGPRLane, lane: int) -> None:
+  """VOPD: dual-issue, execute two ops simultaneously (read all inputs before writes)."""
+  literal, vdstx = inst._literal, inst.vdstx
+  vdsty = v[(inst.vdsty << 1) | ((inst.vdstx.offset & 1) ^ 1)]  # vdsty is raw int from VDSTYField.decode
+  sx0, sx1, dx, sy0, sy1, dy = st.rsrc(inst.srcx0, lane, literal), V[inst.vsrcx1], V[vdstx], st.rsrc(inst.srcy0, lane, literal), V[inst.vsrcy1], V[vdsty]
+  V[vdstx] = inst._fnx(sx0, sx1, 0, dx, st.scc, st.vcc, lane, st.exec_mask, literal, None)['D0']
+  V[vdsty] = inst._fny(sy0, sy1, 0, dy, st.scc, st.vcc, lane, st.exec_mask, literal, None)['D0']
+
+def exec_flat(st: WaveState, inst, V: VGPRLane, lane: int) -> None:
+  """FLAT/GLOBAL/SCRATCH memory ops."""
+  ndwords = _op_ndwords(inst.op_name)
+  addr = V[inst.addr] | (V[inst.addr + 1] << 32)
+  ADDR = (st.rsgpr64(inst.saddr) + V[inst.addr] + _sext(inst.offset, 13)) & MASK64 if inst.saddr != NULL else (addr + _sext(inst.offset, 13)) & MASK64
+  vdata_src = inst.vdst if 'LOAD' in inst.op_name else inst.data
+  result = inst._fn(GlobalMem, ADDR, _vgpr_read(V, vdata_src, ndwords), V[inst.vdst])
+  if 'VDATA' in result: _vgpr_write(V, inst.vdst, result['VDATA'], ndwords)
+  if 'RETURN_DATA' in result: _vgpr_write(V, inst.vdst, result['RETURN_DATA'], ndwords)
+
+def exec_ds(st: WaveState, inst, V: VGPRLane, lane: int) -> None:
+  """DS (LDS) memory ops."""
+  ndwords = _op_ndwords(inst.op_name)
+  data0, data1 = _vgpr_read(V, inst.data0, ndwords), _vgpr_read(V, inst.data1, ndwords) if inst.data1 is not None else 0
+  result = inst._fn(st.lds, V[inst.addr], data0, data1, inst.offset0, inst.offset1)
+  if 'RETURN_DATA' in result and ('_RTN' in inst.op_name or '_LOAD' in inst.op_name):
+    _vgpr_write(V, inst.vdst, result['RETURN_DATA'], ndwords * 2 if '_2ADDR_' in inst.op_name else ndwords)
+
+def exec_vop(st: WaveState, inst: Inst, V: VGPRLane, lane: int) -> None:
+  """VOP1/VOP2/VOP3/VOP3SD/VOP3P/VOPC: standard ALU ops."""
+  is_dst_16 = inst.canonical_op_bits['d'] == 16
+  if isinstance(inst, VOP3P):
+    src0, src1, src2, vdst, dst_hi = inst.src0, inst.src1, inst.src2, inst.vdst, False
+    neg, abs_, opsel = inst.neg, 0, inst.opsel
+  elif isinstance(inst, VOP1):
+    src0, src1, src2, vdst = inst.src0, None, None, inst.vdst
+    neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst.offset & 0x80) != 0 and is_dst_16
+    if is_dst_16: vdst = v[inst.vdst.offset & 0x7f]
+  elif isinstance(inst, VOP2):
+    src0, src1, src2, vdst = inst.src0, inst.vsrc1, None, inst.vdst
+    neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst.offset & 0x80) != 0 and is_dst_16
+    if is_dst_16: vdst = v[inst.vdst.offset & 0x7f]
+  elif isinstance(inst, (VOP3, VOP3SD)):
+    src0, src1, src2, vdst = inst.src0, inst.src1, (None if isinstance(inst, VOP3) and inst.op.value < 256 else inst.src2), inst.vdst
+    neg, abs_, opsel, dst_hi = (inst.neg, inst.abs, inst.opsel, False) if isinstance(inst, VOP3) else (0, 0, 0, False)
+  elif isinstance(inst, VOPC):
+    src0, src1, src2, vdst, neg, abs_, opsel, dst_hi = inst.src0, inst.vsrc1, None, VCC_LO, 0, 0, 0, False
+  else:
+    raise NotImplementedError(f"exec_vop: unhandled instruction type {type(inst).__name__}")
+
+  s0 = _read_src(st, inst, src0, 0, lane, neg, abs_, opsel)
+  s1 = _read_src(st, inst, src1, 1, lane, neg, abs_, opsel)
+  s2 = _read_src(st, inst, src2, 2, lane, neg, abs_, opsel)
+  if isinstance(inst, VOP2) and is_dst_16: d0 = _src16(V[vdst], dst_hi)
+  elif inst.canonical_op_bits['d'] == 64: d0 = V[vdst] | (V[vdst + 1] << 32)
+  else: d0 = V[vdst]
+
+  if isinstance(inst, VOP3SD) and 'CO_CI' in inst.op_name: vcc_for_fn = st.rsgpr64(inst.src2)
+  elif isinstance(inst, VOP3) and inst.op in (VOP3Op.V_CNDMASK_B32_E64, VOP3Op.V_CNDMASK_B16) and src2 is not None and src2.offset < 256: vcc_for_fn = st.rsgpr64(src2)
+  else: vcc_for_fn = st.vcc
+  src0_off = src0.offset if src0 is not None else 0
+  src0_idx = (src0_off - 256) if src0_off >= 256 else src0_off
+  vdst_off = vdst.offset
+  extra_kwargs = {'opsel': opsel, 'opsel_hi': inst.opsel_hi | (inst.opsel_hi2 << 2)} if isinstance(inst, VOP3P) and 'FMA_MIX' in inst.op_name else {}
+  result = inst._fn(s0, s1, s2, d0, st.scc, vcc_for_fn, lane, st.exec_mask, inst._literal, st.vgpr, src0_idx, vdst_off, **extra_kwargs)
+
+  # Check if this is a VOPC instruction (either standalone VOPC or VOP3 with VOPC opcode)
+  is_vopc = isinstance(inst.op, VOPCOp) or (isinstance(inst, VOP3) and inst.op.value < 256)
+  if 'VCC' in result:
+    if isinstance(inst, VOP3SD): st.pend_sgpr_lane(inst.sdst, lane, (result['VCC'] >> lane) & 1)
+    elif isinstance(inst, VOP2) and 'CO_CI' in inst.op_name: st.pend_sgpr_lane(VCC_LO, lane, (result['VCC'] >> lane) & 1)
+    elif is_vopc: st.pend_sgpr_lane(vdst, lane, (result['VCC'] >> lane) & 1)  # vdst is VCC_LO for VOPC
+    else: st.pend_sgpr_lane(VCC_LO, lane, (result['VCC'] >> lane) & 1)
+  if 'EXEC' in result:
+    st.pend_sgpr_lane(EXEC_LO, lane, (result['EXEC'] >> lane) & 1)
+  elif is_vopc:
+    st.pend_sgpr_lane(vdst, lane, (result['D0'] >> lane) & 1)
+  if not is_vopc:
+    d0_val = result['D0']
+    if inst.canonical_op_bits['d'] == 64: V[vdst], V[vdst + 1] = d0_val & MASK32, (d0_val >> 32) & MASK32
+    elif not isinstance(inst, VOP3P) and 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.offset, inst.src1.offset, inst.src2.offset, inst.vdst.offset
+  # 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 - 256 + (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 - 256 + i//32] = _i32(mat_d[i])
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PROGRAM DECODE
+# ═══════════════════════════════════════════════════════════════════════════════
+
+# Wave-level dispatch functions: (st, inst) -> return_code (0 = continue, -1 = end, -2 = barrier)
+def dispatch_endpgm(st, inst): return -1
+def dispatch_barrier(st, inst): st.pc += inst._words; return -2
+def dispatch_nop(st, inst): st.pc += inst._words; return 0
+def dispatch_wmma(st, inst): exec_wmma(st, inst, inst.op); st.pc += inst._words; return 0
+def dispatch_writelane(st, inst): st.vgpr[st.rsrc(inst.src1, 0, inst._literal) & 0x1f][inst.vdst.offset - 256] = st.rsrc(inst.src0, 0, inst._literal) & MASK32; st.pc += inst._words; return 0
+def dispatch_readlane(st, inst):
+  src0_off = inst.src0.offset
+  src0_idx = (src0_off - 256) if src0_off >= 256 else src0_off
+  s1 = st.rsrc(inst.src1, 0, inst._literal) if getattr(inst, 'src1', None) is not None else 0
+  result = inst._fn(0, s1, 0, 0, st.scc, st.vcc, 0, st.exec_mask, inst._literal, st.vgpr, src0_idx, inst.vdst.offset)
+  st.wsgpr(inst.vdst.offset, result['D0'])
+  st.pc += inst._words; return 0
+
+# Per-lane dispatch wrapper: wraps per-lane exec functions into wave-level dispatch
+@functools.cache
+def dispatch_lane(exec_fn):
+  def dispatch(st, inst):
+    exec_mask, vgpr, n_lanes = st.exec_mask, st.vgpr, st.n_lanes
+    for lane in range(n_lanes):
+      if exec_mask >> lane & 1: exec_fn(st, inst, vgpr[lane], lane)
+    st.commit_pends()
+    st.pc += inst._words
+    return 0
+  return dispatch
+
+def decode_program(data: bytes) -> dict[int, Inst]:
+  result: dict[int, Inst] = {}
+  i = 0
+  while i < len(data):
+    inst = decode_inst(data[i:])
+    inst._words = inst.size() // 4
+
+    # Determine dispatch function and pcode function
+    if isinstance(inst, SOPP) and inst.op == SOPPOp.S_CODE_END: break
+    elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_ENDPGM: inst._dispatch = dispatch_endpgm
+    elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER: inst._dispatch = dispatch_barrier
+    elif isinstance(inst, SOPP) and inst.op in (SOPPOp.S_CLAUSE, SOPPOp.S_WAITCNT, SOPPOp.S_WAITCNT_DEPCTR, SOPPOp.S_SENDMSG, SOPPOp.S_SET_INST_PREFETCH_DISTANCE, SOPPOp.S_DELAY_ALU): inst._dispatch = dispatch_nop
+    elif isinstance(inst, (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM)): inst._dispatch = exec_scalar
+    elif isinstance(inst, VOP1) and inst.op == VOP1Op.V_NOP_E32: inst._dispatch = dispatch_nop
+    elif isinstance(inst, VOP3P) and 'WMMA' in inst.op_name: inst._dispatch = dispatch_wmma
+    elif isinstance(inst, VOP3) and inst.op == VOP3Op.V_WRITELANE_B32: inst._dispatch = dispatch_writelane
+    elif isinstance(inst, (VOP1, VOP3)) and inst.op in (VOP1Op.V_READFIRSTLANE_B32_E32, VOP3Op.V_READFIRSTLANE_B32, VOP3Op.V_READLANE_B32): inst._dispatch = dispatch_readlane
+    elif isinstance(inst, VOPD): inst._dispatch = dispatch_lane(exec_vopd)
+    elif isinstance(inst, (FLAT, GLOBAL, SCRATCH)): inst._dispatch = dispatch_lane(exec_flat)
+    elif isinstance(inst, DS): inst._dispatch = dispatch_lane(exec_ds)
+    else: inst._dispatch = dispatch_lane(exec_vop)
+
+    # Compile pcode for instructions that use it (not VOPD which has _fnx/_fny, not special dispatches)
+    # VOPD needs separate functions for X and Y ops
+    if isinstance(inst, VOPD):
+      def _compile_vopd_op(op): return compile_pseudocode(type(op).__name__, op.name, PCODE[op])
+      inst._fnx, inst._fny = _compile_vopd_op(_VOPD_TO_VOP[inst.opx]), _compile_vopd_op(_VOPD_TO_VOP[inst.opy])
+    elif inst._dispatch not in (dispatch_endpgm, dispatch_barrier, dispatch_nop, dispatch_wmma, dispatch_writelane):
+      assert type(inst.op) != int, f"inst op of {inst} is int"
+      inst._fn = compile_pseudocode(type(inst.op).__name__, inst.op.name, PCODE[inst.op])
+    result[i // 4] = inst
+    i += inst._words * 4
+  return result
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# MAIN EXECUTION LOOP
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def exec_wave(program: dict[int, Inst], st: WaveState) -> int:
+  while (inst := program.get(st.pc)) and (result := inst._dispatch(st, inst)) == 0: pass
+  return result
+
+def exec_workgroup(program: dict[int, Inst], workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int, rsrc2: int) -> None:
+  lx, ly, lz = local_size
+  total_threads = lx * ly * lz
+  # GRANULATED_LDS_SIZE is in 512-byte units (see ops_amd.py: lds_size = ((group_segment_size + 511) // 512))
+  lds_size = ((rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE_SHIFT) * 512
+  lds = LDSMem(bytearray(lds_size)) if lds_size else None
+  waves: list[WaveState] = []
+  for wave_start in range(0, total_threads, WAVE_SIZE):
+    n_lanes = min(WAVE_SIZE, total_threads - wave_start)
+    st = WaveState(lds, n_lanes)
+    st.exec_mask = (1 << n_lanes) - 1
+    st.wsgpr64(s[0:1], args_ptr)  # s[0:1] = kernel arguments pointer
+    # COMPUTE_PGM_RSRC2: USER_SGPR_COUNT is where workgroup IDs start, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z control which are passed
+    sgpr_idx = (rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT_SHIFT
+    if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X: st.sgpr[sgpr_idx] = workgroup_id[0]; sgpr_idx += 1
+    if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Y: st.sgpr[sgpr_idx] = workgroup_id[1]; sgpr_idx += 1
+    if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Z: st.sgpr[sgpr_idx] = workgroup_id[2]
+    # VGPR0 = packed workitem IDs: (Z << 20) | (Y << 10) | X
+    for tid in range(wave_start, wave_start + n_lanes):
+      st.vgpr[tid - wave_start][0] = ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
+    waves.append(st)
+  while waves:
+    waves = [st for st in waves if exec_wave(program, st) != -1]
+
+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)
+  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)
+  return 0
@@ -0,0 +1,822 @@
+# DSL for RDNA3 pseudocode - makes pseudocode expressions work directly as Python
+import struct, math, re, functools
+
+MASK32, MASK64 = 0xFFFFFFFF, 0xFFFFFFFFFFFFFFFF
+
+# Float/int bit conversion functions
+_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):
+  i = i & MASK32
+  # RDNA3 default mode: flush f32 denormals to zero (FTZ)
+  # Denormal: exponent=0 (bits 23-30) and mantissa!=0 (bits 0-22)
+  if (i & 0x7f800000) == 0 and (i & 0x007fffff) != 0: return 0.0
+  return _struct_f.unpack(_struct_I.pack(i))[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:
+    bits = _struct_I.unpack(_struct_f.pack(f))[0]
+    # RDNA3 default mode: flush f32 denormals to zero (FTZ)
+    if (bits & 0x7f800000) == 0 and (bits & 0x007fffff) != 0: return 0x80000000 if bits & 0x80000000 else 0
+    return bits
+  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
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# INTERNAL HELPERS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+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 _check_nan_type(x, quiet_bit_expected, default):
+  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)
+      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 _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 _fpop(fn):
+  def wrapper(x):
+    x = float(x)
+    if math.isnan(x) or math.isinf(x): return x
+    result = float(fn(x))
+    return math.copysign(0.0, x) if result == 0.0 else result
+  return wrapper
+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 _f16_to_f32_bits(bits): return struct.unpack("<e", struct.pack("<H", int(bits) & 0xffff))[0]
+def _brev(v, bits): return int(bin(v & ((1 << bits) - 1))[2:].zfill(bits)[::-1], 2)
+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 _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 _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))
+  result = fn(frac_cycles * 2 * math.pi)
+  # Hardware returns exactly 0 for cos(π/2), sin(π), etc. due to lookup table
+  # Round very small results (below f32 precision) to exactly 0
+  if abs(result) < 1e-7: return 0.0
+  return result
+
+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))
+
+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)
+
+class _RoundMode:
+  NEAREST_EVEN = 0
+
+class _WaveMode:
+  IEEE = False
+
+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):
+    f = float(other)
+    if math.isinf(f) or math.isnan(f) or f == 0.0: return False
+    if self._bits == 64:
+      bits = struct.unpack("<Q", struct.pack("<d", f))[0]
+      return (bits >> 52) & 0x7ff == 0
+    bits = struct.unpack("<I", struct.pack("<f", f))[0]
+    return (bits >> 23) & 0xff == 0
+  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)
+
+_pack = lambda hi, lo: ((int(hi) & 0xffff) << 16) | (int(lo) & 0xffff)
+_pack32 = lambda hi, lo: ((int(hi) & 0xffffffff) << 32) | (int(lo) & 0xffffffff)
+
+class TypedView:
+  """View into a Reg with typed access. Used for both full-width (Reg.u32) and slices (Reg[31:16])."""
+  __slots__ = ('_reg', '_high', '_low', '_signed', '_float', '_bf16', '_reversed')
+  def __init__(self, reg, high, low=0, signed=False, is_float=False, is_bf16=False):
+    # Handle reversed slices like [0:31] which means bit-reverse
+    if high < low: high, low, reversed = low, high, True
+    else: reversed = False
+    self._reg, self._high, self._low, self._reversed = reg, high, low, reversed
+    self._signed, self._float, self._bf16 = signed, is_float, is_bf16
+
+  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)
+
+  @property
+  def _val(self): return self._get()
+  @property
+  def _bits(self): return self._nbits()
+
+  # Type accessors for slices (e.g., D0[31:16].f16)
+  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
+
+  # Chained type access (e.g., jump_addr.i64 when jump_addr is already TypedView)
+  @property
+  def i64(s): return s if s._nbits() == 64 and s._signed else int(s)
+  @property
+  def u64(s): return s if s._nbits() == 64 and not s._signed else int(s) & MASK64
+
+  def __getitem__(self, key):
+    if isinstance(key, slice):
+      high, low = int(key.start), int(key.stop)
+      return TypedView(self._reg, high, low)
+    return (self._get() >> 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._get(), self._nbits()) if self._signed else self._get()
+  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._get())
+      bits = self._nbits()
+      return _f16(self._get()) if bits == 16 else _f32(self._get()) if bits == 32 else _f64(self._get())
+    return float(int(self))
+  def __bool__(s): return bool(int(s))
+
+  # 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 or s._nbits() > 64 else 0
+  def __rshift__(s, o): n = int(o); return int(s) >> n if 0 <= n < 64 or s._nbits() > 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)
+
+class Reg:
+  """GPU register: D0.f32 = S0.f32 + S1.f32 just works. Supports up to 128 bits for DS_LOAD_B128."""
+  __slots__ = ('_val',)
+  def __init__(self, val=0): self._val = int(val)
+
+  # Typed views - TypedView(reg, high, signed, is_float, is_bf16)
+  u64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
+  i64 = property(lambda s: TypedView(s, 63, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK64))
+  b64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
+  f64 = property(lambda s: TypedView(s, 63, is_float=True), lambda s, v: setattr(s, '_val', v if isinstance(v, int) else _i64(float(v))))
+  u32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
+  i32 = property(lambda s: TypedView(s, 31, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK32))
+  b32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
+  f32 = property(lambda s: TypedView(s, 31, is_float=True), lambda s, v: setattr(s, '_val', _i32(float(v))))
+  u24 = property(lambda s: TypedView(s, 23))
+  i24 = property(lambda s: TypedView(s, 23, signed=True))
+  u16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
+  i16 = property(lambda s: TypedView(s, 15, signed=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
+  b16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
+  f16 = property(lambda s: TypedView(s, 15, 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, 15, 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, 7))
+  i8 = property(lambda s: TypedView(s, 7, signed=True))
+  u3 = property(lambda s: TypedView(s, 2))  # 3-bit for opsel fields
+  u1 = property(lambda s: TypedView(s, 0))  # single bit
+
+  def __getitem__(s, key):
+    if isinstance(key, slice): return TypedView(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)
+      if high < low: high, low = low, high
+      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)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PSEUDOCODE API - Functions and constants from AMD ISA pseudocode
+# ═══════════════════════════════════════════════════════════════════════════════
+
+# Rounding and float operations
+trunc, floor, ceil = _fpop(math.trunc), _fpop(math.floor), _fpop(math.ceil)
+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"))
+def fract(x): return x - math.floor(x)
+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 isEven(x):
+  x = float(x)
+  if math.isinf(x) or math.isnan(x): return False
+  return int(x) % 2 == 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
+
+# Type conversions
+i32_to_f32 = u32_to_f32 = i32_to_f64 = u32_to_f64 = f32_to_f64 = f64_to_f32 = float
+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(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 bf16_to_f32(v): return _bf16(v) if isinstance(v, int) else float(v)
+def f32_to_bf16(f): return _ibf16(f)
+def u8_to_u32(v): return int(v) & 0xff
+def u4_to_u32(v): return int(v) & 0xf
+def u32_to_u16(u): return int(u) & 0xffff
+def i32_to_i16(i): return ((int(i) + 32768) & 0xffff) - 32768
+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 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
+
+# Min/max operations
+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)
+
+# SAD/MSAD operations
+def ABSDIFF(a, b): return abs(int(a) - int(b))
+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
+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 BYTE_PERMUTE(data, sel):
+  """Select a byte from 64-bit data based on selector value."""
+  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
+  if sel == 9: return 0xff if ((int(data) >> 31) & 1) else 0x00
+  if sel == 10: return 0xff if ((int(data) >> 47) & 1) else 0x00
+  if sel == 11: return 0xff if ((int(data) >> 63) & 1) else 0x00
+  if sel == 12: return 0x00
+  return 0xff
+
+# Pseudocode functions
+def s_ff1_i32_b32(v): return _ctz(v, 32)
+def s_ff1_i32_b64(v): return _ctz(v, 64)
+GT_NEG_ZERO, LT_NEG_ZERO = _gt_neg_zero, _lt_neg_zero
+def isNAN(x):
+  try: return math.isnan(float(x))
+  except (TypeError, ValueError): return False
+def isQuietNAN(x): return _check_nan_type(x, 1, True)
+def isSignalNAN(x): return _check_nan_type(x, 0, False)
+def fma(a, b, c):
+  try: return math.fma(a, b, c)
+  except ValueError: return float('nan')
+def ldexp(m, e): return math.ldexp(m, e)
+def sign(f): return 1 if math.copysign(1.0, f) < 0 else 0
+def exponent(f):
+  if hasattr(f, '_bits') and hasattr(f, '_float') and f._float:
+    raw = f._val
+    if f._bits == 16: return (raw >> 10) & 0x1f
+    if f._bits == 32: return (raw >> 23) & 0xff
+    if f._bits == 64: return (raw >> 52) & 0x7ff
+  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 signext(x): return int(x)
+def cvtToQuietNAN(x): return float('nan')
+
+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)
+  if isinstance(x, TypedView): return x
+  return float(x)
+
+# Constants
+PI = math.pi
+WAVE32, WAVE64 = True, False
+OVERFLOW_F32, UNDERFLOW_F32 = float('inf'), 0.0
+OVERFLOW_F64, UNDERFLOW_F64 = float('inf'), 0.0
+MAX_FLOAT_F32 = 3.4028235e+38
+INF = _Inf()
+ROUND_MODE = _RoundMode()
+WAVE_MODE = _WaveMode()
+DENORM = _Denorm()
+
+# 2/PI with 1201 bits of precision for V_TRIG_PREOP_F64
+TWO_OVER_PI_1201 = Reg(0x0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# COMPILER: pseudocode -> Python (minimal transforms)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def _filter_pseudocode(pseudocode: str) -> str:
+  """Filter raw PDF pseudocode to only include actual code lines."""
+  pcode_lines, in_lambda, depth = [], 0, 0
+  for line in pseudocode.split('\n'):
+    s = line.strip()
+    if not s: continue
+    if '=>' in s or re.match(r'^[A-Z_]+\(', s): continue  # Skip example lines
+    if '= lambda(' in s: in_lambda += 1; continue  # Skip lambda definitions
+    if in_lambda > 0:
+      if s.endswith(');'): in_lambda -= 1
+      continue
+    # Only include lines that look like pseudocode
+    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', 'VADDR', 'VDATA', 'VDST', 'SADDR', 'OFFSET']) 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 s.startswith('if '): depth += 1
+    elif s.startswith('endif'): depth = max(0, depth - 1)
+    if is_code: pcode_lines.append(s)
+  return '\n'.join(pcode_lines)
+
+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
+  for line in joined_lines:
+    line = line.split('//')[0].strip()  # Strip C-style comments
+    if not line: continue
+    if line.startswith('if '):
+      lines.append('  ' * indent + f"if {_expr(line[3:].rstrip(' then'))}:")
+      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'))}:")
+      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 line.startswith('declare '):
+      pass
+    elif m := re.match(r'for (\w+) in (.+?)\s*:\s*(.+?) do', line):
+      start, end = _expr(m[2].strip()), _expr(m[3].strip())
+      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])
+        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)
+            lines.append('  ' * indent + f"{lhs.strip()} {op} {_expr(rhs.strip())}")
+            break
+      else:
+        lhs, rhs = line.split('=', 1)
+        lhs_s, rhs_s = _expr(lhs.strip()), rhs.strip()
+        stmt = _assign(lhs_s, _expr(rhs_s))
+        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) -> str:
+  if lhs in ('tmp', 'SCC', 'VCC', 'EXEC', 'D0', 'D1', 'saveexec', 'PC'):
+    return f"{lhs} = Reg({rhs})"
+  return f"{lhs} = {rhs}"
+
+def _expr(e: str) -> str:
+  e = e.strip()
+  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()), _expr(m[2].strip())
+    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])
+        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
+
+def _apply_pseudocode_fixes(op_name: str, 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(1 << laneId); 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(1 << laneId); 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(1 << laneId)')
+  if op_name == 'V_DIV_SCALE_F64':
+    code = code.replace('D0.f64 = float("nan")', 'VCC = Reg(1 << laneId); 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(1 << laneId); 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(1 << laneId)')
+  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_name: str, pc: str, code: str) -> str:
+  """Generate a single compiled pseudocode function.
+  Functions take int parameters and return dict of int values.
+  Reg wrapping happens inside the function, only for registers actually used."""
+  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)
+  is_ds = cls_name == 'DSOp'
+  is_flat = cls_name in ('FLATOp', 'GLOBALOp', 'SCRATCHOp')
+  is_smem = cls_name == 'SMEMOp'
+  has_s_array = 'S[i]' in pc  # FMA_MIX style: S[0], S[1], S[2] array access
+  combined = code + pc
+
+  fn_name = f"_{cls_name}_{op_name}"
+
+  # Detect which registers are used/modified
+  def needs_init(name): return name in combined and not re.search(rf'^\s*{name}\s*=\s*Reg\(', code, re.MULTILINE)
+  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 function signature and Reg init lines
+  if is_smem:
+    lines = [f"def {fn_name}(MEM, addr):"]
+    reg_inits = ["ADDR=Reg(addr)", "SDATA=Reg(0)"]
+    special_regs = []
+  elif is_ds:
+    lines = [f"def {fn_name}(MEM, addr, data0, data1, offset0, offset1):"]
+    reg_inits = ["ADDR=Reg(addr)", "DATA0=Reg(data0)", "DATA1=Reg(data1)", "OFFSET0=Reg(offset0)", "OFFSET1=Reg(offset1)", "RETURN_DATA=Reg(0)"]
+    special_regs = [('DATA', 'DATA0'), ('DATA2', 'DATA1'), ('OFFSET', 'OFFSET0'), ('ADDR_BASE', 'ADDR')]
+  elif is_flat:
+    lines = [f"def {fn_name}(MEM, addr, vdata, vdst):"]
+    reg_inits = ["ADDR=addr", "VDATA=Reg(vdata)", "VDST=Reg(vdst)", "RETURN_DATA=Reg(0)"]
+    special_regs = [('DATA', 'VDATA')]
+  elif has_s_array:
+    # FMA_MIX style: needs S[i] array, opsel, opsel_hi for source selection (neg/neg_hi applied in emu.py before call)
+    lines = [f"def {fn_name}(s0, s1, s2, d0, scc, vcc, laneId, exec_mask, literal, VGPR, src0_idx=0, vdst_idx=0, pc=None, opsel=0, opsel_hi=0):"]
+    reg_inits = ["S0=Reg(s0)", "S1=Reg(s1)", "S2=Reg(s2)", "S=[S0,S1,S2]", "D0=Reg(d0)", "OPSEL=Reg(opsel)", "OPSEL_HI=Reg(opsel_hi)"]
+    special_regs = []
+    # Detect array declarations like "declare in : 32'F[3]" and create them (rename 'in' to 'ins' since 'in' is a keyword)
+    if "in[" in combined:
+      reg_inits.append("ins=[Reg(0),Reg(0),Reg(0)]")
+      code = code.replace("in[", "ins[")
+  else:
+    lines = [f"def {fn_name}(s0, s1, s2, d0, scc, vcc, laneId, exec_mask, literal, VGPR, src0_idx=0, vdst_idx=0, pc=None):"]
+    # Only create Regs for registers actually used in the pseudocode
+    reg_inits = []
+    if 'S0' in combined: reg_inits.append("S0=Reg(s0)")
+    if 'S1' in combined: reg_inits.append("S1=Reg(s1)")
+    if 'S2' in combined: reg_inits.append("S2=Reg(s2)")
+    if modifies_d0 or 'D0' in combined: reg_inits.append("D0=Reg(s0)" if is_div_scale else "D0=Reg(d0)")
+    if modifies_scc or 'SCC' in combined: reg_inits.append("SCC=Reg(scc)")
+    if modifies_vcc or 'VCC' in combined: reg_inits.append("VCC=Reg(vcc)")
+    if modifies_exec or 'EXEC' in combined: reg_inits.append("EXEC=Reg(exec_mask)")
+    if modifies_pc or 'PC' in combined: reg_inits.append("PC=Reg(pc) if pc is not None else None")
+    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)'))
+
+  # Build init code
+  init_parts = reg_inits.copy()
+  for name, init in special_regs:
+    if name in combined: init_parts.append(f"{name}={init}")
+  if 'EXEC_LO' in code: init_parts.append("EXEC_LO=TypedView(EXEC, 31, 0)")
+  if 'EXEC_HI' in code: init_parts.append("EXEC_HI=TypedView(EXEC, 63, 32)")
+  if 'VCCZ' in code and not re.search(r'^\s*VCCZ\s*=', code, re.MULTILINE): init_parts.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_parts.append("EXECZ=Reg(1 if EXEC._val == 0 else 0)")
+
+  # Add init line and separator
+  if init_parts: lines.append(f"  {'; '.join(init_parts)}")
+
+  # Add compiled pseudocode
+  for line in code.split('\n'):
+    if line.strip(): lines.append(f"  {line}")
+
+  # Build result dict
+  result_items = []
+  if modifies_d0: result_items.append("'D0': D0._val")
+  if modifies_scc: result_items.append("'SCC': SCC._val")
+  if modifies_vcc: result_items.append("'VCC': VCC._val")
+  if modifies_exec: result_items.append("'EXEC': EXEC._val")
+  if has_d1: result_items.append("'D1': D1._val")
+  if modifies_pc: result_items.append("'PC': PC._val")
+  if is_smem and 'SDATA' in combined and re.search(r'^\s*SDATA[\.\[].*=', code, re.MULTILINE):
+    result_items.append("'SDATA': SDATA._val")
+  if is_ds and 'RETURN_DATA' in combined and re.search(r'^\s*RETURN_DATA[\.\[].*=', code, re.MULTILINE):
+    result_items.append("'RETURN_DATA': RETURN_DATA._val")
+  if is_flat:
+    if 'RETURN_DATA' in combined and re.search(r'^\s*RETURN_DATA[\.\[].*=', code, re.MULTILINE):
+      result_items.append("'RETURN_DATA': RETURN_DATA._val")
+    if re.search(r'^\s*VDATA[\.\[].*=', code, re.MULTILINE):
+      result_items.append("'VDATA': VDATA._val")
+  lines.append(f"  return {{{', '.join(result_items)}}}")
+  return '\n'.join(lines)
+
+# Build the globals dict for exec() - includes all pcode symbols
+_PCODE_GLOBALS = {
+  'Reg': Reg, 'TypedView': TypedView, '_pack': _pack, '_pack32': _pack32,
+  'ABSDIFF': ABSDIFF, 'BYTE_PERMUTE': BYTE_PERMUTE, 'DENORM': DENORM, 'F': F,
+  'GT_NEG_ZERO': GT_NEG_ZERO, 'LT_NEG_ZERO': LT_NEG_ZERO, 'INF': INF,
+  'MAX_FLOAT_F32': MAX_FLOAT_F32, 'OVERFLOW_F32': OVERFLOW_F32, 'OVERFLOW_F64': OVERFLOW_F64,
+  'UNDERFLOW_F32': UNDERFLOW_F32, 'UNDERFLOW_F64': UNDERFLOW_F64,
+  'PI': PI, 'ROUND_MODE': ROUND_MODE, 'WAVE_MODE': WAVE_MODE,
+  'WAVE32': WAVE32, 'WAVE64': WAVE64, 'TWO_OVER_PI_1201': TWO_OVER_PI_1201,
+  'SAT8': SAT8, 'trunc': trunc, 'floor': floor, 'ceil': ceil, 'sqrt': sqrt,
+  'log2': log2, 'fract': fract, 'sin': sin, 'cos': cos, 'pow': pow,
+  'isEven': isEven, 'mantissa': mantissa, 'signext_from_bit': signext_from_bit,
+  'i32_to_f32': i32_to_f32, 'u32_to_f32': u32_to_f32, 'i32_to_f64': i32_to_f64,
+  'u32_to_f64': u32_to_f64, 'f32_to_f64': f32_to_f64, 'f64_to_f32': f64_to_f32,
+  'f32_to_i32': f32_to_i32, 'f32_to_u32': f32_to_u32, 'f64_to_i32': f64_to_i32,
+  'f64_to_u32': f64_to_u32, 'f32_to_f16': f32_to_f16, 'f16_to_f32': f16_to_f32,
+  'i16_to_f16': i16_to_f16, 'u16_to_f16': u16_to_f16, 'f16_to_i16': f16_to_i16,
+  'f16_to_u16': f16_to_u16, 'bf16_to_f32': bf16_to_f32, 'f32_to_bf16': f32_to_bf16,
+  'u8_to_u32': u8_to_u32, 'u4_to_u32': u4_to_u32, 'u32_to_u16': u32_to_u16,
+  'i32_to_i16': i32_to_i16, 'f16_to_snorm': f16_to_snorm, 'f16_to_unorm': f16_to_unorm,
+  'f32_to_snorm': f32_to_snorm, 'f32_to_unorm': f32_to_unorm,
+  'v_cvt_i16_f32': v_cvt_i16_f32, 'v_cvt_u16_f32': v_cvt_u16_f32, 'f32_to_u8': f32_to_u8,
+  'v_min_f32': v_min_f32, 'v_max_f32': v_max_f32, 'v_min_f16': v_min_f16, 'v_max_f16': v_max_f16,
+  'v_min_i32': v_min_i32, 'v_max_i32': v_max_i32, 'v_min_i16': v_min_i16, 'v_max_i16': v_max_i16,
+  'v_min_u32': v_min_u32, 'v_max_u32': v_max_u32, 'v_min_u16': v_min_u16, 'v_max_u16': v_max_u16,
+  'v_min3_f32': v_min3_f32, 'v_max3_f32': v_max3_f32, 'v_min3_f16': v_min3_f16, 'v_max3_f16': v_max3_f16,
+  'v_min3_i32': v_min3_i32, 'v_max3_i32': v_max3_i32, 'v_min3_i16': v_min3_i16, 'v_max3_i16': v_max3_i16,
+  'v_min3_u32': v_min3_u32, 'v_max3_u32': v_max3_u32, 'v_min3_u16': v_min3_u16, 'v_max3_u16': v_max3_u16,
+  'v_sad_u8': v_sad_u8, 'v_msad_u8': v_msad_u8,
+  's_ff1_i32_b32': s_ff1_i32_b32, 's_ff1_i32_b64': s_ff1_i32_b64,
+  'isNAN': isNAN, 'isQuietNAN': isQuietNAN, 'isSignalNAN': isSignalNAN,
+  'fma': fma, 'ldexp': ldexp, 'sign': sign, 'exponent': exponent,
+  'signext': signext, 'cvtToQuietNAN': cvtToQuietNAN,
+}
+
+@functools.cache
+def compile_pseudocode(cls_name: str, op_name: str, pseudocode: str):
+  """Compile pseudocode string to executable function. Cached for performance."""
+  filtered = _filter_pseudocode(pseudocode)
+  code = _compile_pseudocode(filtered)
+  code = _apply_pseudocode_fixes(op_name, code)
+  fn_code = _generate_function(cls_name, op_name, filtered, code)
+  fn_name = f"_{cls_name}_{op_name}"
+  local_ns = {}
+  exec(fn_code, _PCODE_GLOBALS, local_ns)
+  return local_ns[fn_name]
@@ -0,0 +1,386 @@
+"""SQTT (SQ Thread Trace) packet encoder and decoder for AMD GPUs.
+
+This module provides encoding and decoding of raw SQTT byte streams.
+The format is nibble-based with variable-width packets determined by a state machine.
+Uses BitField infrastructure from dsl.py, similar to GPU instruction encoding.
+"""
+from __future__ import annotations
+from typing import Iterator
+from enum import Enum
+from extra.assembly.amd.dsl import BitField, FixedBitField, bits
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# FIELD ENUMS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class MemSrc(Enum):
+  LDS = 0
+  LDS_ALT = 1
+  VMEM = 2
+  VMEM_ALT = 3
+
+class AluSrc(Enum):
+  NONE = 0
+  SALU = 1
+  VALU = 2
+  VALU_SALU = 3
+
+class InstOp(Enum):
+  """SQTT instruction operation types.
+
+  Memory ops appear in two ranges depending on which SIMD executes them:
+  - 0x1x-0x2x range: ops on traced SIMD
+  - 0x5x range: ops on other SIMD (OTHER_ prefix)
+
+  GLOBAL memory ops encoding depends on addressing mode AND size:
+  - Loads: 0x21 (saddr=SGPR) or 0x22 (saddr=NULL), all sizes same
+  - Stores: base + size_offset, where VADDR is shifted +1 from SADDR
+    SADDR: 0x24(32) 0x25(64) 0x26(96) 0x27(128)
+    VADDR: 0x25(32) 0x26(64) 0x27(96) 0x28(128)
+
+  OTHER_ range follows same pattern but values overlap differently.
+  """
+  SALU = 0x0
+  SMEM = 0x1
+  JUMP = 0x3              # branch taken
+  JUMP_NO = 0x4           # branch not taken
+  MESSAGE = 0x9
+  VALU_TRANS = 0xb        # transcendental: exp, log, rcp, sqrt, sin, cos
+  VALU_64_SHIFT = 0xd     # 64-bit shifts: lshl, lshr, ashr
+  VALU_MAD64 = 0xe        # 64-bit multiply-add
+  VALU_64 = 0xf           # 64-bit: add, mul, fma, rcp, sqrt, rounding, frexp, div helpers
+  VINTERP = 0x12          # interpolation: v_interp_p10_f32, v_interp_p2_f32
+  BARRIER = 0x13
+
+  # FLAT memory ops on traced SIMD (0x1x range)
+  FLAT_LOAD = 0x1c
+  FLAT_STORE = 0x1d
+  FLAT_STORE_64 = 0x1e
+  FLAT_STORE_96 = 0x1f
+  FLAT_STORE_128 = 0x20
+
+  # GLOBAL memory ops on traced SIMD (0x2x range)
+  GLOBAL_LOAD = 0x21             # saddr=SGPR, all sizes
+  GLOBAL_LOAD_VADDR = 0x22       # saddr=NULL, all sizes
+  GLOBAL_STORE = 0x24            # saddr=SGPR, 32-bit
+  GLOBAL_STORE_64 = 0x25         # saddr=SGPR 64 or saddr=NULL 32
+  GLOBAL_STORE_96 = 0x26         # saddr=SGPR 96 or saddr=NULL 64
+  GLOBAL_STORE_128 = 0x27        # saddr=SGPR 128 or saddr=NULL 96
+  GLOBAL_STORE_VADDR_128 = 0x28  # saddr=NULL, 128-bit
+
+  # LDS ops on traced SIMD
+  LDS_LOAD = 0x29
+  LDS_STORE = 0x2b
+  LDS_STORE_64 = 0x2c
+  LDS_STORE_128 = 0x2e
+
+  # Memory ops on other SIMD (0x5x range)
+  OTHER_LDS_LOAD = 0x50
+  OTHER_LDS_STORE = 0x51
+  OTHER_LDS_STORE_64 = 0x52
+  OTHER_LDS_STORE_128 = 0x54
+  OTHER_FLAT_LOAD = 0x55
+  OTHER_FLAT_STORE = 0x56
+  OTHER_FLAT_STORE_64 = 0x57
+  OTHER_FLAT_STORE_96 = 0x58
+  OTHER_FLAT_STORE_128 = 0x59
+  OTHER_GLOBAL_LOAD = 0x5a             # saddr=SGPR, all sizes
+  OTHER_GLOBAL_LOAD_VADDR = 0x5b       # saddr=NULL or saddr=SGPR store 32
+  OTHER_GLOBAL_STORE_64 = 0x5c         # saddr=SGPR 64 or saddr=NULL 32
+  OTHER_GLOBAL_STORE_96 = 0x5d         # saddr=SGPR 96 or saddr=NULL 64
+  OTHER_GLOBAL_STORE_128 = 0x5e        # saddr=SGPR 128 or saddr=NULL 96
+  OTHER_GLOBAL_STORE_VADDR_128 = 0x5f  # saddr=NULL, 128-bit
+
+  # EXEC-modifying ops (0x7x range)
+  SALU_SAVEEXEC = 0x72    # s_*_saveexec_b32/b64
+  VALU_CMPX = 0x73        # v_cmpx_*
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PACKET TYPE BASE CLASS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class PacketType:
+  """Base class for SQTT packet types."""
+  encoding: FixedBitField
+  _raw: int
+  _time: int
+
+  def __init_subclass__(cls, **kwargs):
+    super().__init_subclass__(**kwargs)
+    cls._fields = {k: v for k, v in cls.__dict__.items() if isinstance(v, BitField)}
+    cls._size_nibbles = ((max((f.hi for f in cls._fields.values()), default=0) + 4) // 4)
+
+  @classmethod
+  def from_raw(cls, raw: int, time: int = 0):
+    inst = object.__new__(cls)
+    inst._raw, inst._time = raw, time
+    return inst
+
+  def __repr__(self) -> str:
+    fields_str = ", ".join(f"{k}={getattr(self, k)}" for k in self._fields if not k.startswith('_'))
+    return f"{self.__class__.__name__}({fields_str})"
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PACKET TYPE DEFINITIONS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class VALUINST(PacketType):  # exclude: 1 << 2
+  encoding = bits[2:0] == 0b011
+  delta = bits[5:3]
+  flag = bits[6:6]
+  wave = bits[11:7]
+
+class VMEMEXEC(PacketType):  # exclude: 1 << 0
+  encoding = bits[3:0] == 0b1111
+  delta = bits[5:4]
+  src = bits[7:6].enum(MemSrc)
+
+class ALUEXEC(PacketType):  # exclude: 1 << 1
+  encoding = bits[3:0] == 0b1110
+  delta = bits[5:4]
+  src = bits[7:6].enum(AluSrc)
+
+class IMMEDIATE(PacketType):  # exclude: 1 << 5
+  encoding = bits[3:0] == 0b1101
+  delta = bits[6:4]
+  wave = bits[11:7]
+
+class IMMEDIATE_MASK(PacketType):  # exclude: 1 << 5
+  encoding = bits[4:0] == 0b00100
+  delta = bits[7:5]
+  mask = bits[23:8]
+
+class WAVERDY(PacketType):  # exclude: 1 << 3
+  encoding = bits[4:0] == 0b10100
+  delta = bits[7:5]
+  mask = bits[23:8]
+
+class TS_DELTA_S8_W3(PacketType):
+  encoding = bits[6:0] == 0b0100001
+  delta = bits[10:8]
+  _padding = bits[63:11]
+
+class WAVEEND(PacketType):  # exclude: 1 << 4
+  encoding = bits[4:0] == 0b10101
+  delta = bits[7:5]
+  flag7 = bits[8:8]
+  simd = bits[10:9]
+  cu_lo = bits[13:11]
+  wave = bits[19:15]
+  @property
+  def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
+
+class WAVESTART(PacketType):  # exclude: 1 << 4
+  encoding = bits[4:0] == 0b01100
+  delta = bits[6:5]
+  flag7 = bits[7:7]
+  simd = bits[9:8]
+  cu_lo = bits[12:10]
+  wave = bits[17:13]
+  id7 = bits[31:18]
+  @property
+  def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
+
+class TS_DELTA_S5_W2(PacketType):
+  encoding = bits[4:0] == 0b11100
+  delta = bits[6:5]
+  _padding = bits[47:7]
+
+class WAVEALLOC(PacketType):  # exclude: 1 << 10
+  encoding = bits[4:0] == 0b00101
+  delta = bits[7:5]
+  _padding = bits[19:8]
+
+class TS_DELTA_S5_W3(PacketType):
+  encoding = bits[4:0] == 0b00110
+  delta = bits[7:5]
+  _padding = bits[51:8]
+
+class PERF(PacketType):  # exclude: 1 << 11
+  encoding = bits[4:0] == 0b10110
+  delta = bits[7:5]
+  arg = bits[27:8]
+
+class TS_DELTA_SHORT(PacketType):
+  encoding = bits[3:0] == 0b1000
+  delta = bits[7:4]
+
+class NOP(PacketType):
+  encoding = bits[3:0] == 0b0000
+  delta = None  # type: ignore
+  _padding = bits[3:0]
+
+class TS_WAVE_STATE(PacketType):
+  encoding = bits[6:0] == 0b1010001
+  delta = bits[15:7]
+  coarse = bits[23:16]
+  @property
+  def wave_interest(self) -> bool: return bool(self.coarse & 1)
+  @property
+  def terminate_all(self) -> bool: return bool(self.coarse & 8)
+
+class EVENT(PacketType):  # exclude: 1 << 7
+  encoding = bits[7:0] == 0b01100001
+  delta = bits[10:8]
+  event = bits[23:11]
+
+class EVENT_BIG(PacketType):
+  encoding = bits[7:0] == 0b11100001
+  delta = bits[10:8]
+  event = bits[31:11]
+
+class REG(PacketType):
+  encoding = bits[3:0] == 0b1001
+  delta = bits[6:4]
+  slot = bits[9:7]
+  hi_byte = bits[15:8]
+  subop = bits[31:16]
+  val32 = bits[63:32]
+  @property
+  def is_config(self) -> bool: return bool(self.hi_byte & 0x80)
+
+class SNAPSHOT(PacketType):
+  encoding = bits[6:0] == 0b1110001
+  delta = bits[9:7]
+  snap = bits[63:10]
+
+class TS_DELTA_OR_MARK(PacketType):
+  encoding = bits[6:0] == 0b0000001
+  delta = bits[47:12]
+  bit8 = bits[8:8]
+  bit9 = bits[9:9]
+  @property
+  def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
+
+class LAYOUT_HEADER(PacketType):
+  encoding = bits[6:0] == 0b0010001
+  delta = None  # type: ignore
+  layout = bits[12:7]
+  simd = bits[14:13]
+  group = bits[17:15]
+  sel_a = bits[31:28]
+  sel_b = bits[36:33]
+  flag4 = bits[59:59]
+  _padding = bits[63:60]
+
+class INST(PacketType):
+  encoding = bits[2:0] == 0b010
+  delta = bits[6:4]
+  flag1 = bits[3:3]
+  flag2 = bits[7:7]
+  wave = bits[12:8]
+  op = bits[19:13].enum(InstOp)
+
+class UTILCTR(PacketType):
+  encoding = bits[6:0] == 0b0110001
+  delta = bits[8:7]
+  ctr = bits[47:9]
+
+# All packet types in encoding priority order (more specific masks first, NOP last as fallback)
+PACKET_TYPES: list[type[PacketType]] = [
+  EVENT, EVENT_BIG,
+  TS_DELTA_S8_W3, TS_WAVE_STATE, SNAPSHOT, TS_DELTA_OR_MARK, LAYOUT_HEADER, UTILCTR,
+  IMMEDIATE_MASK, WAVERDY, WAVEEND, WAVESTART, TS_DELTA_S5_W2, WAVEALLOC, TS_DELTA_S5_W3, PERF,
+  VMEMEXEC, ALUEXEC, IMMEDIATE, TS_DELTA_SHORT, REG,
+  VALUINST, INST,
+  NOP,
+]
+
+def _build_state_table() -> tuple[bytes, dict[int, type[PacketType]]]:
+  table = [len(PACKET_TYPES) - 1] * 256  # default to NOP
+  opcode_to_class: dict[int, type[PacketType]] = {i: cls for i, cls in enumerate(PACKET_TYPES)}
+
+  for byte_val in range(256):
+    for opcode, pkt_cls in enumerate(PACKET_TYPES):
+      if (byte_val & pkt_cls.encoding.mask) == pkt_cls.encoding.default:
+        table[byte_val] = opcode
+        break
+
+  return bytes(table), opcode_to_class
+
+STATE_TO_OPCODE, OPCODE_TO_CLASS = _build_state_table()
+
+# Precompute special case opcodes
+_TS_DELTA_OR_MARK_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_OR_MARK)
+_TS_DELTA_SHORT_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_SHORT)
+
+# Combined lookup: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
+# special_case: 0=none, 1=TS_DELTA_OR_MARK, 2=TS_DELTA_SHORT
+_DECODE_INFO: dict[int, tuple] = {}
+for _opcode, _pkt_cls in OPCODE_TO_CLASS.items():
+  _delta_field = getattr(_pkt_cls, 'delta', None)
+  _delta_lo = _delta_field.lo if _delta_field else 0
+  _delta_mask = _delta_field.mask if _delta_field else 0
+  _special = 1 if _opcode == _TS_DELTA_OR_MARK_OPCODE else (2 if _opcode == _TS_DELTA_SHORT_OPCODE else 0)
+  _DECODE_INFO[_opcode] = (_pkt_cls, _pkt_cls._size_nibbles, _delta_lo, _delta_mask, _special)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# DECODER
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def decode(data: bytes) -> Iterator[PacketType]:
+  """Decode raw SQTT blob, yielding packet instances."""
+  n, reg, pos, nib_off, nib_count, time = len(data), 0, 0, 0, 16, 0
+
+  while pos + ((nib_count + nib_off + 1) >> 1) <= n:
+    need = nib_count - nib_off
+    # 1. if unaligned, read high nibble to align
+    if nib_off: reg, pos = (reg >> 4) | ((data[pos] >> 4) << 60), pos + 1
+    # 2. read all full bytes at once
+    if (byte_count := need >> 1):
+      chunk = int.from_bytes(data[pos:pos + byte_count], 'little')
+      reg, pos = (reg >> (byte_count * 8)) | (chunk << (64 - byte_count * 8)), pos + byte_count
+    # 3. if odd, read low nibble
+    if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)
+
+    opcode = STATE_TO_OPCODE[reg & 0xFF]
+    pkt_cls, nib_count, delta_lo, delta_mask, special = _DECODE_INFO[opcode]
+    delta = (reg >> delta_lo) & delta_mask
+    if special == 1 and (reg >> 9) & 1 and not (reg >> 8) & 1: delta = 0  # TS_DELTA_OR_MARK marker
+    elif special == 2: delta += 8  # TS_DELTA_SHORT
+    time += delta
+    yield pkt_cls.from_raw(reg, time)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PRINTER
+# ═══════════════════════════════════════════════════════════════════════════════
+
+PACKET_COLORS = {
+  "INST": "WHITE", "VALUINST": "BLACK", "VMEMEXEC": "yellow", "ALUEXEC": "yellow",
+  "IMMEDIATE": "YELLOW", "IMMEDIATE_MASK": "YELLOW", "WAVERDY": "cyan", "WAVEALLOC": "cyan",
+  "WAVEEND": "blue", "WAVESTART": "blue", "PERF": "magenta", "EVENT": "red", "EVENT_BIG": "red",
+  "REG": "green", "LAYOUT_HEADER": "white", "SNAPSHOT": "white", "UTILCTR": "green",
+}
+
+def format_packet(p) -> str:
+  from tinygrad.helpers import colored
+  name = type(p).__name__
+  if isinstance(p, INST):
+    op_name = p.op.name if isinstance(p.op, InstOp) else f"0x{p.op:02x}"
+    fields = f"wave={p.wave} op={op_name}" + (" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "")
+  elif isinstance(p, VALUINST): fields = f"wave={p.wave}" + (" flag" if p.flag else "")
+  elif isinstance(p, ALUEXEC): fields = f"src={p.src.name if isinstance(p.src, AluSrc) else p.src}"
+  elif isinstance(p, VMEMEXEC): fields = f"src={p.src.name if isinstance(p.src, MemSrc) else p.src}"
+  elif isinstance(p, (WAVESTART, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
+  elif hasattr(p, '_fields'):
+    fields = " ".join(f"{k}=0x{getattr(p, k):x}" if k in {'snap', 'val32'} else f"{k}={getattr(p, k)}"
+                      for k in p._fields if not k.startswith('_') and k not in {'delta', 'encoding'})
+  else: fields = ""
+  return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name, 'white'))} {fields}"
+
+def print_packets(packets) -> None:
+  skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"}
+  for p in packets:
+    if type(p).__name__ not in skip: print(format_packet(p))
+
+if __name__ == "__main__":
+  import sys, pickle
+  if len(sys.argv) < 2:
+    print("Usage: python sqtt.py <pkl_file>")
+    sys.exit(1)
+  with open(sys.argv[1], "rb") as f:
+    data = pickle.load(f)
+  sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
+  for i, event in enumerate(sqtt_events):
+    print(f"\n=== event {i} ===")
+    print_packets(decode(event.blob))
@@ -0,0 +1,132 @@
+# maps SQTT trace packets to instructions.
+from dataclasses import dataclass
+from typing import Iterator
+
+from tinygrad.runtime.support.elf import elf_loader
+
+from extra.assembly.amd.sqtt import decode, print_packets, INST, VALUINST, IMMEDIATE, WAVESTART, WAVEEND, InstOp, PacketType, IMMEDIATE_MASK
+from extra.assembly.amd.dsl import Inst
+from extra.assembly.amd.decode import decode_inst
+from extra.assembly.amd.autogen.rdna3.ins import SOPP, s_endpgm
+from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
+
+@dataclass(frozen=True)
+class InstructionInfo:
+  pc: int
+  wave: int
+  inst: Inst
+
+def map_insts(data:bytes, lib:bytes) -> Iterator[tuple[PacketType, InstructionInfo|None]]:
+  """maps SQTT packets to instructions, yields (packet, instruction_info or None)"""
+  # map pcs to insts
+  pc_map:dict[int, Inst] = {}
+  image, sections, _ = elf_loader(lib)
+  text = next((sh for sh in sections if sh.name == ".text"), None)
+  assert text is not None, "no .text section found"
+  text_off, text_size = text.header.sh_addr, text.header.sh_size
+  offset = text_off
+  while offset < text_off + text_size:
+    inst = decode_inst(image[offset:])
+    pc_map[offset-text_off] = inst
+    offset += inst.size()
+
+  wave_pc:dict[int, int] = {}
+  # only processing packets on one [CU, SIMD] unit
+  def simd_select(p) -> bool: return getattr(p, "cu", 0) == 0 and getattr(p, "simd", 0) == 0
+  for p in decode(data):
+    if not simd_select(p): continue
+    if isinstance(p, WAVESTART):
+      assert p.wave not in wave_pc, "only one inflight wave per unit"
+      wave_pc[p.wave] = 0
+      continue
+    if isinstance(p, WAVEEND):
+      pc = wave_pc.pop(p.wave)
+      yield (p, InstructionInfo(pc, p.wave, s_endpgm()))
+      continue
+    # skip OTHER_ instructions, they don't belong to this unit
+    if isinstance(p, INST) and p.op.name.startswith("OTHER_"): continue
+    if isinstance(p, IMMEDIATE_MASK):
+      # immediate mask may yield multiple times per packet
+      for wave in range(16):
+        if p.mask & (1 << wave):
+          inst = pc_map[pc:=wave_pc[wave]]
+          # can this assert be more strict?
+          assert isinstance(inst, SOPP), f"IMMEDIATE_MASK packet must map to SOPP, got {inst}"
+          wave_pc[wave] += inst.size()
+          yield (p, InstructionInfo(pc, wave, inst))
+      continue
+    if isinstance(p, (VALUINST, INST, IMMEDIATE)):
+      inst = pc_map[pc:=wave_pc[p.wave]]
+      # s_delay_alu doesn't get a packet?
+      if isinstance(inst, SOPP) and inst.op in {SOPPOp.S_DELAY_ALU}:
+        wave_pc[p.wave] += inst.size()
+        inst = pc_map[pc:=wave_pc[p.wave]]
+      # identify a branch instruction, only used for asserts
+      is_branch = isinstance(inst, SOPP) and "BRANCH" in inst.op_name
+      if is_branch: assert isinstance(p, INST) and p.op in {InstOp.JUMP_NO, InstOp.JUMP}, f"branch can only be folowed by jump packets, got {p}"
+      # JUMP handling
+      if isinstance(p, INST) and p.op is InstOp.JUMP:
+        assert is_branch, f"JUMP packet must map to a branch instruction, got {inst}"
+        x = inst.simm16 & 0xffff
+        wave_pc[p.wave] += inst.size() + (x - 0x10000 if x & 0x8000 else x)*4
+      else:
+        if is_branch: assert inst.op != SOPPOp.S_BRANCH, f"S_BRANCH must have a JUMP packet, got {p}"
+        wave_pc[p.wave] += inst.size()
+      yield (p, InstructionInfo(pc, p.wave, inst))
+      continue
+    # for all other packets (VMEMEXEC, ALUEXEC, etc.), yield with None
+    yield (p, None)
+
+# test to compare every packet with the rocprof decoder
+
+def test_rocprof_inst_traces_match(sqtt, prg, target):
+  from tinygrad.viz.serve import llvm_disasm
+  from extra.sqtt.roc import decode as roc_decode, InstExec
+  disasm = {addr+prg.base:inst_disasm for addr, inst_disasm in llvm_disasm(target, prg.lib).items()}
+  rctx = roc_decode([sqtt], {prg.name:disasm})
+  rwaves = rctx.inst_execs[(sqtt.kern, sqtt.exec_tag)]
+  rwaves_iter:dict[int, list[Iterator[InstExec]]] = {} # wave unit (0-15) -> list of inst trace iterators for all executions on that unit
+  for w in rwaves: rwaves_iter.setdefault(w.wave_id, []).append(w.unpack_insts())
+  rwaves_base = next(iter(disasm)) # base program counter
+
+  passed_insts = 0
+  for pkt, info in map_insts(sqtt.blob, prg.lib):
+    if DEBUG >= 2: print_packets([pkt])
+    if info is None: continue
+    if DEBUG >= 2: print(f"{' '*29}{info.inst.disasm()}")
+    rocprof_inst = next(rwaves_iter[info.wave][0])
+    ref_pc = rocprof_inst.pc-rwaves_base
+    # always check pc matches
+    assert ref_pc == info.pc, f"pc mismatch {ref_pc}:{disasm[rocprof_inst.pc][0]} != {info.pc}:{info.inst.disasm()}"
+    # special handling for s_endpgm, it marks the wave completion.
+    if info.inst == s_endpgm():
+      completed_wave = list(rwaves_iter[info.wave].pop(0))
+      assert len(completed_wave) == 0, f"incomplete instructions in wave {info.wave}"
+    # otherwise the packet timestamp is time + "stall"
+    else:
+      assert pkt._time == rocprof_inst.time+rocprof_inst.stall
+    passed_insts += 1
+
+  for k,v in rwaves_iter.items():
+    assert len(v) == 0, f"incomplete wave {k}"
+
+  print(f"passed for {passed_insts} instructions across {len(rwaves)} waves scheduled on {len(rwaves_iter)} wave units")
+
+if __name__ == "__main__":
+  import argparse, pickle, pathlib
+  from tinygrad.helpers import temp, DEBUG
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--profile', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
+                      default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  parser.add_argument('--kernel', type=str, default=None, metavar="NAME", help='Kernel to focus on (optional name, default: all kernels)')
+  args = parser.parse_args()
+  with open(args.profile, "rb") as f:
+    data = pickle.load(f)
+  sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
+  kern_events = {e.name:e for e in data if type(e).__name__ == "ProfileProgramEvent"}
+  target = next((e for e in data if type(e).__name__ == "ProfileDeviceEvent" and e.device.startswith("AMD"))).props["gfx_target_version"]
+  for e in sqtt_events:
+    if args.kernel is not None and args.kernel != e.kern: continue
+    if not e.itrace: continue
+    print(f"==== {e.kern}")
+    test_rocprof_inst_traces_match(e, kern_events[e.kern], target)
@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+"""Benchmark comparing Python vs Rust RDNA3 emulators on real tinygrad kernels."""
+import ctypes, time, os
+from pathlib import Path
+
+# 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
+
+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, rsrc2: int, 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, rsrc2)
+
+  # 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, rsrc2)
+    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 get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes], int] | None:
+  """Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data, rsrc2)."""
+  try:
+    from tinygrad import Tensor
+    from tinygrad.runtime.support.elf import elf_loader
+    from tinygrad.runtime.autogen import hsa
+    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)
+        image = memoryview(bytearray(lib))
+        _, sections, _ = elf_loader(lib)
+        rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
+        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
+            # Extract rsrc2 from ELF (same as ops_amd.py)
+            group_segment_size = image[rodata_entry:rodata_entry+4].cast("I")[0]
+            lds_size = ((group_segment_size + 511) // 512) & 0x1FF
+            code = hsa.amd_kernel_code_t.from_buffer_copy(bytes(image[rodata_entry:rodata_entry+256]) + b'\x00'*256)
+            rsrc2 = code.compute_pgm_rsrc2 | (lds_size << 15)
+            return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data, rsrc2)
+    return None
+  except Exception as e:
+    print(f"  Error getting kernel: {e}")
+    return None
+
+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("--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 = []
+
+  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, rsrc2 = 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, rsrc2, args.iterations)
+    rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, rsrc2, 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((op_name, n_insts, n_workgroups, py_time, rust_time))
+
+  # Summary table
+  print("\n" + "=" * 90)
+  print("SUMMARY")
+  print("=" * 90)
+  print(f"{'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
+  print("-" * 90)
+
+  for 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"{name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
+
+if __name__ == "__main__":
+  main()
@@ -0,0 +1,196 @@
+# 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()
@@ -0,0 +1,79 @@
+"""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")
+
+ARCH_TO_TARGET:dict[str, list[str]] = {
+  "rdna3":["gfx1100"],
+  "rdna4":["gfx1200"],
+  "cdna":["gfx950", "gfx942"],
+}
+
+TARGET_TO_ARCH:dict[str, str] = {t:arch for arch,targets in ARCH_TO_TARGET.items() for t in targets}
+
+def get_target(arch:str) -> str: return ARCH_TO_TARGET[arch][0]
+
+def get_mattr(arch:str) -> str:
+  return {"rdna3":"+real-true16,+wavefrontsize32", "rdna4":"+real-true16,+wavefrontsize32", "cdna":"+wavefrontsize64"}[arch]
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# 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, TypedView
+    self._Reg, self._MASK64, self._TypedView = Reg, MASK64, TypedView
+    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._TypedView(self.EXEC, 31, 0), 'EXEC_HI': self._TypedView(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}
@@ -0,0 +1 @@
+"""Hardware-validated emulator tests for RDNA3 instructions."""
@@ -0,0 +1,251 @@
+"""Test infrastructure for hardware-validated RDNA3 emulator tests.
+
+Uses run_asm() with memory output, so tests can run on both emulator and real hardware.
+Set USE_HW=1 to run on both emulator and hardware, comparing results.
+"""
+import ctypes, math, os, struct
+from extra.assembly.amd.autogen.rdna3.ins import *
+
+from extra.assembly.amd.emu import run_asm
+from extra.assembly.amd.dsl import NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, M0
+
+def _i32(f: float) -> int: return struct.unpack('<I', struct.pack('<f', f))[0]
+def _f32(i: int) -> float: return struct.unpack('<f', struct.pack('<I', i & 0xFFFFFFFF))[0]
+
+# f16 conversion helpers
+def _f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
+def f32_to_f16(f: float) -> int:
+  f = float(f)
+  if math.isnan(f): return 0x7e00
+  if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
+  try: return struct.unpack('<H', struct.pack('<e', f))[0]
+  except OverflowError: return 0x7c00 if f > 0 else 0xfc00
+
+# For backwards compatibility with tests using SrcEnum.NULL etc.
+class SrcEnum:
+  NULL = NULL
+  VCC_LO = VCC_LO
+  VCC_HI = VCC_HI
+  EXEC_LO = EXEC_LO
+  EXEC_HI = EXEC_HI
+  SCC = SCC
+  M0 = M0
+  POS_HALF = 0.5
+  NEG_HALF = -0.5
+  POS_ONE = 1.0
+  NEG_ONE = -1.0
+  POS_TWO = 2.0
+  NEG_TWO = -2.0
+  POS_FOUR = 4.0
+  NEG_FOUR = -4.0
+
+VCC = VCC_LO  # For VOP3SD sdst field (VCC_LO is exported from dsl)
+USE_HW = os.environ.get("USE_HW", "0") == "1"
+FLOAT_TOLERANCE = 1e-5
+
+# Output buffer layout: vgpr[16][32], sgpr[16], vcc, scc, exec
+N_VGPRS, N_SGPRS, WAVE_SIZE = 16, 16, 32
+VGPR_BYTES = N_VGPRS * WAVE_SIZE * 4  # 16 regs * 32 lanes * 4 bytes = 2048
+SGPR_BYTES = N_SGPRS * 4  # 16 regs * 4 bytes = 64
+OUT_BYTES = VGPR_BYTES + SGPR_BYTES + 12  # + vcc + scc + exec
+
+# Float conversion helpers
+def f2i(f: float) -> int: return _i32(f)
+def i2f(i: int) -> float: return _f32(i)
+def f2i64(f: float) -> int: return struct.unpack('<Q', struct.pack('<d', f))[0]
+def i642f(i: int) -> float: return struct.unpack('<d', struct.pack('<Q', i))[0]
+
+def assemble(instructions: list) -> bytes:
+  return b''.join(inst.to_bytes() for inst in instructions)
+
+# Simple WaveState class for test output parsing (mirrors emu.py interface for tests)
+class WaveState:
+  def __init__(self):
+    self.vgpr = [[0] * 256 for _ in range(32)]  # vgpr[lane][reg]
+    self.sgpr = [0] * 128
+    self.vcc = 0
+    self.scc = 0
+
+def get_prologue_epilogue(n_lanes: int) -> tuple[list, list]:
+  """Generate prologue and epilogue instructions for state capture."""
+  prologue = [
+    s_mov_b32(s[80], s[0]),
+    s_mov_b32(s[81], s[1]),
+    v_mov_b32_e32(v[255], v[0]),
+  ]
+  for i in range(N_VGPRS):
+    prologue.append(v_mov_b32_e32(v[i], 0))
+  for i in range(N_SGPRS):
+    prologue.append(s_mov_b32(s[i], 0))
+  prologue.append(s_mov_b32(VCC_LO, 0))
+
+  epilogue = [
+    s_mov_b32(s[90], VCC_LO),
+    s_cselect_b32(s[91], 1, 0),
+    # Save EXEC early (before we modify it for VGPR stores)
+    s_mov_b32(s[95], EXEC_LO),
+    # Restore EXEC to all active lanes for VGPR stores (test may have modified EXEC)
+    s_mov_b32(EXEC_LO, (1 << n_lanes) - 1),
+    s_load_b64(s[92:93], s[80:81], 0, soffset=NULL),
+    s_waitcnt(0),  # simm16=0 waits for all
+    v_lshlrev_b32_e32(v[240], 2, v[255]),
+  ]
+  for i in range(N_VGPRS):
+    epilogue.append(global_store_b32(addr=v[240], data=v[i], saddr=s[92:93], offset=i * WAVE_SIZE * 4))
+  epilogue.append(v_mov_b32_e32(v[241], 0))
+  epilogue.append(v_cmp_eq_u32_e32(v[255], v[241]))
+  epilogue.append(s_and_saveexec_b32(s[94], VCC_LO))
+  epilogue.append(v_mov_b32_e32(v[240], 0))
+  for i in range(N_SGPRS):
+    epilogue.append(v_mov_b32_e32(v[243], s[i]))
+    epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + i * 4))
+  epilogue.append(v_mov_b32_e32(v[243], s[90]))
+  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + SGPR_BYTES))
+  epilogue.append(v_mov_b32_e32(v[243], s[91]))
+  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + SGPR_BYTES + 4))
+  # Store EXEC (saved earlier in s[95])
+  epilogue.append(v_mov_b32_e32(v[243], s[95]))
+  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + SGPR_BYTES + 8))
+  epilogue.append(s_mov_b32(EXEC_LO, s[94]))
+  epilogue.append(s_endpgm())
+  return prologue, epilogue
+
+def parse_output(out_buf: bytes, n_lanes: int) -> WaveState:
+  """Parse output buffer into WaveState."""
+  st = WaveState()
+  for i in range(N_VGPRS):
+    for lane in range(n_lanes):
+      off = i * WAVE_SIZE * 4 + lane * 4
+      st.vgpr[lane][i] = struct.unpack_from('<I', out_buf, off)[0]
+  for i in range(N_SGPRS):
+    st.sgpr[i] = struct.unpack_from('<I', out_buf, VGPR_BYTES + i * 4)[0]
+  st.vcc = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES)[0]
+  st.scc = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES + 4)[0]
+  # Store EXEC in its proper location (index 126)
+  st.sgpr[EXEC_LO.offset] = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES + 8)[0]
+  return st
+
+def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
+  """Run instructions via emulator run_asm, dump state to memory, return WaveState."""
+  out_buf = (ctypes.c_uint8 * OUT_BYTES)(*([0] * OUT_BYTES))
+  out_addr = ctypes.addressof(out_buf)
+
+  prologue, epilogue = get_prologue_epilogue(n_lanes)
+  code = assemble(prologue + instructions + epilogue)
+
+  args = (ctypes.c_uint64 * 1)(out_addr)
+  args_ptr = ctypes.addressof(args)
+  kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
+  lib_ptr = ctypes.addressof(kernel_buf)
+
+  # rsrc2: USER_SGPR_COUNT=2, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z=1, LDS_SIZE=128 (64KB)
+  rsrc2 = 0x19c | (128 << 15)
+  scratch_size = 0x10000  # 64KB per lane, matches .amdhsa_private_segment_fixed_size in run_program_hw
+  result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2)
+  assert result == 0, f"run_asm failed with {result}"
+
+  return parse_output(bytes(out_buf), n_lanes)
+
+def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
+  """Run instructions on real AMD hardware via HIPCompiler and AMDProgram."""
+  from tinygrad.device import Device
+  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.runtime.support.compiler_amd import HIPCompiler
+  from tinygrad.helpers import flat_mv
+
+  dev = Device["AMD"]
+  compiler = HIPCompiler(dev.arch)
+
+  prologue, epilogue = get_prologue_epilogue(n_lanes)
+  code = assemble(prologue + instructions + epilogue)
+
+  byte_str = ', '.join(f'0x{b:02x}' for b in code)
+  asm_src = f""".text
+.globl test
+.p2align 8
+.type test,@function
+test:
+.byte {byte_str}
+
+.rodata
+.p2align 6
+.amdhsa_kernel test
+  .amdhsa_next_free_vgpr 256
+  .amdhsa_next_free_sgpr 96
+  .amdhsa_wavefront_size32 1
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_kernarg_size 8
+  .amdhsa_group_segment_fixed_size 65536
+  .amdhsa_private_segment_fixed_size 65536
+  .amdhsa_enable_private_segment 1
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.version:
+  - 1
+  - 0
+amdhsa.kernels:
+  - .name: test
+    .symbol: test.kd
+    .kernarg_segment_size: 8
+    .group_segment_fixed_size: 65536
+    .private_segment_fixed_size: 65536
+    .kernarg_segment_align: 8
+    .wavefront_size: 32
+    .sgpr_count: 96
+    .vgpr_count: 256
+    .max_flat_workgroup_size: 1024
+...
+.end_amdgpu_metadata
+"""
+
+  lib = compiler.compile(asm_src)
+  prg = AMDProgram(dev, "test", lib)
+
+  out_gpu = dev.allocator.alloc(OUT_BYTES)
+  prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
+
+  out_buf = bytearray(OUT_BYTES)
+  dev.allocator._copyout(flat_mv(memoryview(out_buf)), out_gpu)
+
+  return parse_output(bytes(out_buf), n_lanes)
+
+def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS) -> list[str]:
+  """Compare two WaveStates and return list of differences."""
+  import math
+  diffs = []
+  for i in range(n_vgprs):
+    for lane in range(n_lanes):
+      emu_val = emu_st.vgpr[lane][i]
+      hw_val = hw_st.vgpr[lane][i]
+      if emu_val != hw_val:
+        emu_f, hw_f = _f32(emu_val), _f32(hw_val)
+        if math.isnan(emu_f) and math.isnan(hw_f):
+          continue
+        diffs.append(f"v[{i}] lane {lane}: emu=0x{emu_val:08x} ({emu_f:.6g}) hw=0x{hw_val:08x} ({hw_f:.6g})")
+  for i in range(N_SGPRS):
+    emu_val = emu_st.sgpr[i]
+    hw_val = hw_st.sgpr[i]
+    if emu_val != hw_val:
+      diffs.append(f"s[{i}]: emu=0x{emu_val:08x} hw=0x{hw_val:08x}")
+  if emu_st.vcc != hw_st.vcc:
+    diffs.append(f"vcc: emu=0x{emu_st.vcc:08x} hw=0x{hw_st.vcc:08x}")
+  if emu_st.scc != hw_st.scc:
+    diffs.append(f"scc: emu={emu_st.scc} hw={hw_st.scc}")
+  return diffs
+
+def run_program(instructions: list, n_lanes: int = 1) -> WaveState:
+  """Run instructions and return WaveState.
+
+  If USE_HW=1, runs on both emulator and hardware, compares results, and raises if they differ.
+  Otherwise, runs only on emulator.
+  """
+  emu_st = run_program_emu(instructions, n_lanes)
+  if USE_HW:
+    hw_st = run_program_hw(instructions, n_lanes)
+    diffs = compare_wave_states(emu_st, hw_st, n_lanes)
+    if diffs:
+      raise AssertionError(f"Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
+    return hw_st
+  return emu_st
@@ -0,0 +1,679 @@
+"""Tests for DS instructions - data share (LDS) operations.
+
+Includes: ds_store_b32, ds_load_b32, ds_store_2addr_*, ds_load_2addr_*,
+          ds_add_*, ds_max_*, ds_min_*, ds_and_*, ds_or_*, ds_xor_*,
+          ds_inc_*, ds_dec_*, ds_cmpstore_*, ds_storexchg_*
+"""
+import unittest
+from extra.assembly.amd.test.hw.helpers import *
+
+class TestDS2Addr(unittest.TestCase):
+  """Tests for DS_*_2ADDR instructions."""
+
+  def test_ds_store_load_2addr_b32(self):
+    """DS_STORE_2ADDR_B32 and DS_LOAD_2ADDR_B32 with offset * 4."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0xAAAAAAAA),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[0], 0xBBBBBBBB),
+      v_mov_b32_e32(v[1], s[0]),
+      DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2:3], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
+    self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
+
+  def test_ds_store_load_2addr_b64(self):
+    """DS_STORE_2ADDR_B64 and DS_LOAD_2ADDR_B64."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0xDEADBEEF),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[0], 0xCAFEBABE),
+      v_mov_b32_e32(v[1], s[0]),
+      s_mov_b32(s[0], 0x12345678),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0x9ABCDEF0),
+      v_mov_b32_e32(v[3], s[0]),
+      DS(DSOp.DS_STORE_2ADDR_B64, addr=v[10], data0=v[0:1], data1=v[2:3], vdst=v[0], offset0=0, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4:7], offset0=0, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
+    self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
+    self.assertEqual(st.vgpr[0][6], 0x12345678)
+    self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
+
+
+class TestDS2AddrMore(unittest.TestCase):
+  """Additional DS_*_2ADDR tests."""
+
+  def test_ds_store_load_2addr_b32_nonzero_offsets(self):
+    """DS_STORE_2ADDR_B32 with non-zero offsets (offset*4 scaling)."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0x11111111),
+      v_mov_b32_e32(v[0], s[2]),
+      s_mov_b32(s[2], 0x22222222),
+      v_mov_b32_e32(v[1], s[2]),
+      DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=2, offset1=5),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2:3], offset0=2, offset1=5),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have value from offset 8 (2*4)")
+    self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have value from offset 20 (5*4)")
+
+  def test_ds_2addr_b64_no_overlap(self):
+    """DS_LOAD_2ADDR_B64 with adjacent offsets should not overlap."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0x11111111),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_mov_b32(s[2], 0x22222222),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=4),
+      s_mov_b32(s[2], 0x33333333),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=8),
+      s_mov_b32(s[2], 0x44444444),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=12),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4:7], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should be 0x11111111")
+    self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should be 0x22222222")
+    self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should be 0x33333333")
+    self.assertEqual(st.vgpr[0][7], 0x44444444, "v7 should be 0x44444444")
+
+  def test_ds_load_2addr_b32_no_overwrite(self):
+    """DS_LOAD_2ADDR_B32 should only write 2 VGPRs."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0xAAAAAAAA),
+      v_mov_b32_e32(v[0], s[2]),
+      s_mov_b32(s[2], 0xBBBBBBBB),
+      v_mov_b32_e32(v[1], s[2]),
+      DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 0xDEADBEEF),
+      v_mov_b32_e32(v[4], s[2]),  # Sentinel
+      DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2:3], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
+    self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
+    self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should be untouched")
+
+  def test_ds_load_b64_no_overwrite(self):
+    """DS_LOAD_B64 should only write 2 VGPRs."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0xDEADBEEF),
+      v_mov_b32_e32(v[0], s[2]),
+      s_mov_b32(s[2], 0xCAFEBABE),
+      v_mov_b32_e32(v[1], s[2]),
+      ds_store_b64(addr=v[10], data0=v[0:1], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 0x12345678),
+      v_mov_b32_e32(v[4], s[2]),  # Sentinel
+      ds_load_b64(addr=v[10], vdst=v[2:3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
+    self.assertEqual(st.vgpr[0][3], 0xCAFEBABE)
+    self.assertEqual(st.vgpr[0][4], 0x12345678, "v4 should be untouched")
+
+
+class TestDSB128(unittest.TestCase):
+  """Tests for DS_STORE_B128 and DS_LOAD_B128 (128-bit / 4 dwords)."""
+
+  def test_ds_store_load_b128(self):
+    """DS_STORE_B128 stores 4 VGPRs, DS_LOAD_B128 loads them back."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0x11111111),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[0], 0x22222222),
+      v_mov_b32_e32(v[1], s[0]),
+      s_mov_b32(s[0], 0x33333333),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0x44444444),
+      v_mov_b32_e32(v[3], s[0]),
+      ds_store_b128(addr=v[10], data0=v[0:3]),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b128(addr=v[10], vdst=v[4:7]),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have first dword")
+    self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have second dword")
+    self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should have third dword")
+    self.assertEqual(st.vgpr[0][7], 0x44444444, "v7 should have fourth dword")
+
+  def test_ds_store_b128_with_offset(self):
+    """DS_STORE_B128 with non-zero offset."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0xAAAAAAAA),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[0], 0xBBBBBBBB),
+      v_mov_b32_e32(v[1], s[0]),
+      s_mov_b32(s[0], 0xCCCCCCCC),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0xDDDDDDDD),
+      v_mov_b32_e32(v[3], s[0]),
+      DS(DSOp.DS_STORE_B128, addr=v[10], data0=v[0:3], offset0=16),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_B128, addr=v[10], vdst=v[4:7], offset0=16),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
+    self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB)
+    self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC)
+    self.assertEqual(st.vgpr[0][7], 0xDDDDDDDD)
+
+
+class TestDSAtomic(unittest.TestCase):
+  """Tests for DS atomic operations."""
+
+  def test_ds_max_rtn_u32(self):
+    """DS_MAX_RTN_U32: atomically store max and return old value."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 100),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 200),
+      v_mov_b32_e32(v[1], s[2]),
+      ds_max_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
+    self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
+
+  def test_ds_min_rtn_u32(self):
+    """DS_MIN_RTN_U32: atomically store min and return old value."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 200),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 100),
+      v_mov_b32_e32(v[1], s[2]),
+      ds_min_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 200)
+    self.assertEqual(st.vgpr[0][3], 100)
+
+  def test_ds_and_rtn_b32(self):
+    """DS_AND_RTN_B32: atomically AND and return old value."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0xFF00FF00),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 0xFFFF0000),
+      v_mov_b32_e32(v[1], s[2]),
+      ds_and_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xFF00FF00)
+    self.assertEqual(st.vgpr[0][3], 0xFF000000)
+
+  def test_ds_or_rtn_b32(self):
+    """DS_OR_RTN_B32: atomically OR and return old value."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0x00FF0000),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 0x000000FF),
+      v_mov_b32_e32(v[1], s[2]),
+      ds_or_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x00FF0000)
+    self.assertEqual(st.vgpr[0][3], 0x00FF00FF)
+
+  def test_ds_xor_rtn_b32(self):
+    """DS_XOR_RTN_B32: atomically XOR and return old value."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0xAAAAAAAA),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 0xFFFFFFFF),
+      v_mov_b32_e32(v[1], s[2]),
+      ds_xor_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
+    self.assertEqual(st.vgpr[0][3], 0x55555555)
+
+  def test_ds_inc_rtn_u32(self):
+    """DS_INC_RTN_U32: increment with wrap."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 5),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 10),  # limit
+      v_mov_b32_e32(v[1], s[2]),
+      ds_inc_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 5)
+    self.assertEqual(st.vgpr[0][3], 6)
+
+  def test_ds_dec_rtn_u32(self):
+    """DS_DEC_RTN_U32: decrement with wrap."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 5),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 10),  # limit
+      v_mov_b32_e32(v[1], s[2]),
+      ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 5)
+    self.assertEqual(st.vgpr[0][3], 4)
+
+  def test_ds_cmpstore_b32_match(self):
+    """DS_CMPSTORE_B32: conditional store when compare matches."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 100),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 200),
+      v_mov_b32_e32(v[1], s[2]),  # new value
+      s_mov_b32(s[2], 100),
+      v_mov_b32_e32(v[2], s[2]),  # compare = 100 (matches)
+      ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 200)
+
+  def test_ds_cmpstore_b32_no_match(self):
+    """DS_CMPSTORE_B32: no store when compare doesn't match."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 100),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 200),
+      v_mov_b32_e32(v[1], s[2]),  # new value
+      s_mov_b32(s[2], 50),
+      v_mov_b32_e32(v[2], s[2]),  # compare = 50 (doesn't match)
+      ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 100)
+
+  def test_ds_max_u32_no_rtn(self):
+    """DS_MAX_U32 (no RTN): atomically store max, no return value."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 100),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 200),
+      v_mov_b32_e32(v[1], s[2]),
+      ds_max_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
+
+  def test_ds_add_u32_no_rtn_preserves_vdst(self):
+    """DS_ADD_U32 (no RTN) should NOT write to vdst."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0xDEADBEEF),
+      v_mov_b32_e32(v[2], s[2]),  # sentinel
+      s_mov_b32(s[2], 100),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 50),
+      v_mov_b32_e32(v[1], s[2]),
+      ds_add_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xDEADBEEF, "v2 should preserve sentinel")
+    self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
+
+  def test_ds_add_rtn_u32_writes_vdst(self):
+    """DS_ADD_RTN_U32 should write old value to vdst."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0xDEADBEEF),
+      v_mov_b32_e32(v[2], s[2]),  # sentinel
+      s_mov_b32(s[2], 100),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 50),
+      v_mov_b32_e32(v[1], s[2]),
+      ds_add_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
+    self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
+
+  def test_ds_dec_rtn_u32_wrap(self):
+    """DS_DEC_RTN_U32: decrement wraps when value is 0 or > limit."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0),  # Start at 0
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[2], 10),  # limit
+      v_mov_b32_e32(v[1], s[2]),
+      ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0, "v2 should have old value (0)")
+    # When mem == 0 or mem > limit, result = limit
+    self.assertEqual(st.vgpr[0][3], 10, "v3 should wrap to limit (10)")
+
+
+class TestDSStorexchg(unittest.TestCase):
+  """Tests for DS_STOREXCHG instructions."""
+
+  def test_ds_storexchg_rtn_b32(self):
+    """DS_STOREXCHG_RTN_B32: exchange value and return old."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0xAAAAAAAA),
+      v_mov_b32_e32(v[0], s[0]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[0], 0xBBBBBBBB),
+      v_mov_b32_e32(v[1], s[0]),
+      DS(DSOp.DS_STOREXCHG_RTN_B32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
+    self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
+
+
+class TestDSRegisterWidth(unittest.TestCase):
+  """Regression tests: DS loads should only write correct number of VGPRs."""
+
+  def test_ds_load_b32_no_overwrite(self):
+    """DS_LOAD_B32 should only write 1 VGPR."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      s_mov_b32(s[0], 0xDEADBEEF),
+      v_mov_b32_e32(v[1], s[0]),
+      s_mov_b32(s[0], 0x11111111),
+      v_mov_b32_e32(v[2], s[0]),  # sentinel
+      ds_store_b32(addr=v[0], data0=v[1], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(addr=v[0], vdst=v[1], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
+    self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should be untouched")
+
+
+class TestDS2AddrStride64(unittest.TestCase):
+  """Tests for DS_*_2ADDR_STRIDE64 (offset * 256 for B32, offset * 512 for B64)."""
+
+  def test_ds_store_load_2addr_stride64_b32(self):
+    """DS_STORE_2ADDR_STRIDE64_B32: stores at ADDR + offset*256."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0xAAAAAAAA),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[0], 0xBBBBBBBB),
+      v_mov_b32_e32(v[1], s[0]),
+      DS(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[2:3], offset0=1, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 from addr 256")
+    self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 from addr 512")
+
+  def test_ds_store_load_2addr_stride64_b64(self):
+    """DS_STORE_2ADDR_STRIDE64_B64: stores at ADDR + offset*512."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0xDEADBEEF),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[0], 0xCAFEBABE),
+      v_mov_b32_e32(v[1], s[0]),
+      s_mov_b32(s[0], 0x12345678),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0x9ABCDEF0),
+      v_mov_b32_e32(v[3], s[0]),
+      DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0:1], data1=v[2:3], vdst=v[0], offset0=1, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[4:7], offset0=1, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
+    self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
+    self.assertEqual(st.vgpr[0][6], 0x12345678)
+    self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
+
+  def test_ds_storexchg_2addr_rtn_b32(self):
+    """DS_STOREXCHG_2ADDR_RTN_B32: exchange at two addresses."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0x11111111),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[0], 0x22222222),
+      v_mov_b32_e32(v[1], s[0]),
+      DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[0], 0xAAAAAAAA),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0xBBBBBBBB),
+      v_mov_b32_e32(v[3], s[0]),
+      DS(DSOp.DS_STOREXCHG_2ADDR_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4:5], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[6:7], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0x11111111, "old val 0")
+    self.assertEqual(st.vgpr[0][5], 0x22222222, "old val 1")
+    self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "new val 0")
+    self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "new val 1")
+
+
+  def test_ds_storexchg_rtn_b64(self):
+    """DS_STOREXCHG_RTN_B64: exchange 64-bit value and return old."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0xDEADBEEF),
+      v_mov_b32_e32(v[0], s[0]),   # initial low
+      s_mov_b32(s[0], 0xCAFEBABE),
+      v_mov_b32_e32(v[1], s[0]),   # initial high
+      DS(DSOp.DS_STORE_B64, addr=v[10], data0=v[0:1], vdst=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[0], 0x12345678),
+      v_mov_b32_e32(v[2], s[0]),   # new low
+      s_mov_b32(s[0], 0x9ABCDEF0),
+      v_mov_b32_e32(v[3], s[0]),   # new high
+      DS(DSOp.DS_STOREXCHG_RTN_B64, addr=v[10], data0=v[2:3], vdst=v[4:5], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_B64, addr=v[10], vdst=v[6:7], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should have old low dword")
+    self.assertEqual(st.vgpr[0][5], 0xCAFEBABE, "v5 should have old high dword")
+    self.assertEqual(st.vgpr[0][6], 0x12345678, "v6 should have new low dword")
+    self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0, "v7 should have new high dword")
+
+  def test_ds_store_load_2addr_stride64_b64_roundtrip(self):
+    """DS_STORE_2ADDR_STRIDE64_B64 followed by DS_LOAD_2ADDR_STRIDE64_B64 works correctly."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0x11111111),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[0], 0x22222222),
+      v_mov_b32_e32(v[1], s[0]),
+      DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0:1], data1=v[0:1], vdst=v[0], offset0=1, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[2:5], offset0=1, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have val1 low")
+    self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have val1 high")
+    self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have val2 low")
+    self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have val2 high")
+
+  def test_ds_storexchg_2addr_stride64_rtn_b32(self):
+    """DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32: exchange at two addresses (offset*256)."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0x11111111),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[0], 0x22222222),
+      v_mov_b32_e32(v[1], s[0]),
+      DS(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[0], 0xAAAAAAAA),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0xBBBBBBBB),
+      v_mov_b32_e32(v[3], s[0]),
+      DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4:5], offset0=1, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[6:7], offset0=1, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have old value")
+    self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have old value")
+    self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "v6 should have new value")
+    self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "v7 should have new value")
+
+  def test_ds_storexchg_2addr_stride64_rtn_b64_returns_old(self):
+    """DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64: returns old values correctly."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[0], 0x11111111),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[0], 0x22222222),
+      v_mov_b32_e32(v[1], s[0]),
+      DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0:1], data1=v[0:1], vdst=v[0], offset0=1, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+      s_mov_b32(s[0], 0xAAAAAAAA),
+      v_mov_b32_e32(v[6], s[0]),
+      s_mov_b32(s[0], 0xBBBBBBBB),
+      v_mov_b32_e32(v[7], s[0]),
+      DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64, addr=v[10], data0=v[6:7], data1=v[6:7], vdst=v[8:11], offset0=1, offset1=2),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][8], 0x11111111, "v8 should have old val1 low")
+    self.assertEqual(st.vgpr[0][9], 0x22222222, "v9 should have old val1 high")
+    self.assertEqual(st.vgpr[0][10], 0x11111111, "v10 should have old val2 low")
+    self.assertEqual(st.vgpr[0][11], 0x22222222, "v11 should have old val2 high")
+
+
+class TestAtomicOrdering(unittest.TestCase):
+  """Tests for atomic operation return values and ordering."""
+
+  def test_ds_add_rtn_sequence(self):
+    """DS_ADD_RTN returns correct old values in sequence."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      v_mov_b32_e32(v[0], 100),
+      DS(DSOp.DS_STORE_B32, addr=v[10], data0=v[0], vdst=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      v_mov_b32_e32(v[1], 25),
+      DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[3], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+      DS(DSOp.DS_LOAD_B32, addr=v[10], vdst=v[4], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 100, "First add should return 100")
+    self.assertEqual(st.vgpr[0][3], 125, "Second add should return 125")
+    self.assertEqual(st.vgpr[0][4], 150, "Final value should be 150")
+
+
+if __name__ == '__main__':
+  unittest.main()
@@ -0,0 +1,363 @@
+"""Tests for FLAT instructions - flat memory operations.
+
+Includes: flat_load_*, flat_store_*, flat_atomic_*
+"""
+import unittest
+from extra.assembly.amd.test.hw.helpers import *
+
+class TestFlatAtomic(unittest.TestCase):
+  """Tests for FLAT atomic instructions."""
+
+  def _make_test(self, setup_instrs, atomic_instr, check_fn, test_offset=2000):
+    """Helper to create atomic test instructions."""
+    instructions = [
+      s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
+      s_waitcnt(lgkmcnt=0),
+      v_mov_b32_e32(v[0], s[2]),
+      v_mov_b32_e32(v[1], s[3]),
+    ] + setup_instrs + [atomic_instr, s_waitcnt(vmcnt=0),
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 0),
+      s_mov_b32(s[2], 0),
+      s_mov_b32(s[3], 0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    check_fn(st)
+
+  def test_flat_atomic_add_u32(self):
+    """FLAT_ATOMIC_ADD_U32 adds to memory and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 100),
+      v_mov_b32_e32(v[2], s[0]),
+      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 50),
+      v_mov_b32_e32(v[3], s[0]),
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_ADD_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][4], 100)
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_swap_b32(self):
+    """FLAT_ATOMIC_SWAP_B32 swaps memory value and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 0xAAAAAAAA),
+      v_mov_b32_e32(v[2], s[0]),
+      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 0xBBBBBBBB),
+      v_mov_b32_e32(v[3], s[0]),
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_SWAP_B32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_and_b32(self):
+    """FLAT_ATOMIC_AND_B32 ANDs with memory and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 0xFF00FF00),
+      v_mov_b32_e32(v[2], s[0]),
+      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 0xFFFF0000),
+      v_mov_b32_e32(v[3], s[0]),
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_AND_B32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][4], 0xFF00FF00)
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_or_b32(self):
+    """FLAT_ATOMIC_OR_B32 ORs with memory and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 0x00FF0000),
+      v_mov_b32_e32(v[2], s[0]),
+      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 0x0000FF00),
+      v_mov_b32_e32(v[3], s[0]),
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_OR_B32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][4], 0x00FF0000)
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_inc_u32(self):
+    """FLAT_ATOMIC_INC_U32 increments and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 10),
+      v_mov_b32_e32(v[2], s[0]),
+      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 100),  # threshold
+      v_mov_b32_e32(v[3], s[0]),
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_INC_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][4], 10)
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_dec_u32(self):
+    """FLAT_ATOMIC_DEC_U32 decrements and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 10),
+      v_mov_b32_e32(v[2], s[0]),
+      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 100),
+      v_mov_b32_e32(v[3], s[0]),
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_DEC_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][4], 10)
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_sub_u32(self):
+    """FLAT_ATOMIC_SUB_U32 subtracts from memory and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 100),
+      v_mov_b32_e32(v[2], s[0]),
+      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 30),
+      v_mov_b32_e32(v[3], s[0]),  # sub 30
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_SUB_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][4], 100, "v4 should have old value (100)")
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_xor_b32(self):
+    """FLAT_ATOMIC_XOR_B32 XORs with memory and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 0xAAAAAAAA),
+      v_mov_b32_e32(v[2], s[0]),
+      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 0xFFFFFFFF),
+      v_mov_b32_e32(v[3], s[0]),  # XOR mask
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_XOR_B32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA, "v4 should have old value")
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_min_u32(self):
+    """FLAT_ATOMIC_MIN_U32 stores min and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 100),
+      v_mov_b32_e32(v[2], s[0]),
+      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 50),
+      v_mov_b32_e32(v[3], s[0]),  # compare value (smaller)
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_MIN_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][4], 100, "v4 should have old value (100)")
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_max_u32(self):
+    """FLAT_ATOMIC_MAX_U32 stores max and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 50),
+      v_mov_b32_e32(v[2], s[0]),
+      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 100),
+      v_mov_b32_e32(v[3], s[0]),  # compare value (larger)
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_MAX_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][4], 50, "v4 should have old value (50)")
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_inc_u64_returns_old_value(self):
+    """FLAT_ATOMIC_INC_U64 should return full 64-bit old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      # Store initial 64-bit value: 0xCAFEBABE_DEADBEEF
+      s_mov_b32(s[0], 0xDEADBEEF),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0xCAFEBABE),
+      v_mov_b32_e32(v[3], s[0]),
+      global_store_b64(addr=v[0:1], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      # Threshold: 0xFFFFFFFF_FFFFFFFF
+      s_mov_b32(s[0], 0xFFFFFFFF),
+      v_mov_b32_e32(v[4], s[0]),
+      v_mov_b32_e32(v[5], s[0]),
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_INC_U64, addr=v[0:1], data=v[4:5], vdst=v[6:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][6], 0xDEADBEEF, "v6 should have old value low dword")
+      self.assertEqual(st.vgpr[0][7], 0xCAFEBABE, "v7 should have old value high dword")
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_add_u64(self):
+    """FLAT_ATOMIC_ADD_U64 adds 64-bit value and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 0x11111111),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0x22222222),
+      v_mov_b32_e32(v[3], s[0]),
+      global_store_b64(addr=v[0:1], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 0x00000001),  # add 1
+      v_mov_b32_e32(v[4], s[0]),
+      s_mov_b32(s[0], 0x00000000),
+      v_mov_b32_e32(v[5], s[0]),
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_ADD_U64, addr=v[0:1], data=v[4:5], vdst=v[6:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][6], 0x11111111, "v6 should have old value low")
+      self.assertEqual(st.vgpr[0][7], 0x22222222, "v7 should have old value high")
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+  def test_flat_atomic_swap_b64(self):
+    """FLAT_ATOMIC_SWAP_B64 swaps 64-bit value and returns old value."""
+    TEST_OFFSET = 2000
+    setup = [
+      s_mov_b32(s[0], 0xAAAAAAAA),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0xBBBBBBBB),
+      v_mov_b32_e32(v[3], s[0]),
+      global_store_b64(addr=v[0:1], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      s_mov_b32(s[0], 0xCCCCCCCC),
+      v_mov_b32_e32(v[4], s[0]),
+      s_mov_b32(s[0], 0xDDDDDDDD),
+      v_mov_b32_e32(v[5], s[0]),
+    ]
+    atomic = FLAT(FLATOp.FLAT_ATOMIC_SWAP_B64, addr=v[0:1], data=v[4:5], vdst=v[6:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
+    def check(st):
+      self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "v6 should have old value low")
+      self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "v7 should have old value high")
+    self._make_test(setup, atomic, check, TEST_OFFSET)
+
+
+class TestFlatLoad(unittest.TestCase):
+  """Tests for FLAT load instructions."""
+
+  def test_flat_load_b32(self):
+    """FLAT_LOAD_B32 loads 32-bit value correctly."""
+    TEST_OFFSET = 2000
+    instructions = [
+      s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
+      s_waitcnt(lgkmcnt=0),
+      v_mov_b32_e32(v[0], s[2]),
+      v_mov_b32_e32(v[1], s[3]),
+      s_mov_b32(s[0], 0xDEADBEEF),
+      v_mov_b32_e32(v[2], s[0]),
+      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      FLAT(FLATOp.FLAT_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 0),
+      s_mov_b32(s[2], 0),
+      s_mov_b32(s[3], 0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
+
+  def test_flat_load_b64(self):
+    """FLAT_LOAD_B64 loads 64-bit value correctly."""
+    TEST_OFFSET = 2000
+    instructions = [
+      s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
+      s_waitcnt(lgkmcnt=0),
+      v_mov_b32_e32(v[0], s[2]),
+      v_mov_b32_e32(v[1], s[3]),
+      s_mov_b32(s[0], 0xDEADBEEF),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0xCAFEBABE),
+      v_mov_b32_e32(v[3], s[0]),
+      global_store_b64(addr=v[0:1], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      FLAT(FLATOp.FLAT_LOAD_B64, addr=v[0:1], vdst=v[4:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 0),
+      s_mov_b32(s[2], 0),
+      s_mov_b32(s[3], 0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
+    self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
+
+  def test_flat_load_b96(self):
+    """FLAT_LOAD_B96 loads 96-bit (3 dword) value correctly."""
+    TEST_OFFSET = 2000
+    instructions = [
+      s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
+      s_waitcnt(lgkmcnt=0),
+      v_mov_b32_e32(v[0], s[2]),
+      v_mov_b32_e32(v[1], s[3]),
+      s_mov_b32(s[0], 0x11111111),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0x22222222),
+      v_mov_b32_e32(v[3], s[0]),
+      s_mov_b32(s[0], 0x33333333),
+      v_mov_b32_e32(v[4], s[0]),
+      global_store_b96(addr=v[0:1], data=v[2:4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      FLAT(FLATOp.FLAT_LOAD_B96, addr=v[0:1], vdst=v[5:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 0),
+      s_mov_b32(s[2], 0),
+      s_mov_b32(s[3], 0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][5], 0x11111111)
+    self.assertEqual(st.vgpr[0][6], 0x22222222)
+    self.assertEqual(st.vgpr[0][7], 0x33333333)
+
+  def test_flat_load_b128(self):
+    """FLAT_LOAD_B128 loads 128-bit value correctly."""
+    TEST_OFFSET = 2000
+    instructions = [
+      s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
+      s_waitcnt(lgkmcnt=0),
+      v_mov_b32_e32(v[0], s[2]),
+      v_mov_b32_e32(v[1], s[3]),
+      s_mov_b32(s[0], 0x11111111),
+      v_mov_b32_e32(v[2], s[0]),
+      s_mov_b32(s[0], 0x22222222),
+      v_mov_b32_e32(v[3], s[0]),
+      s_mov_b32(s[0], 0x33333333),
+      v_mov_b32_e32(v[4], s[0]),
+      s_mov_b32(s[0], 0x44444444),
+      v_mov_b32_e32(v[5], s[0]),
+      global_store_b128(addr=v[0:1], data=v[2:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      FLAT(FLATOp.FLAT_LOAD_B128, addr=v[0:1], vdst=v[6:9], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
+      s_waitcnt(vmcnt=0),
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 0),
+      s_mov_b32(s[2], 0),
+      s_mov_b32(s[3], 0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][6], 0x11111111)
+    self.assertEqual(st.vgpr[0][7], 0x22222222)
+    self.assertEqual(st.vgpr[0][8], 0x33333333)
+    self.assertEqual(st.vgpr[0][9], 0x44444444)
+
+
+if __name__ == '__main__':
+  unittest.main()
--- a/Show More
+++ b/Show More
				`@@ -0,0 +1 @@`
				`"""Hardware-validated emulator tests for RDNA3 instructions."""`