symbolic llm with prefill (ai slop)

* add better printing to tinygrad.apps.llm

* add gc.collect

* comment

.github/actions/process-replay/action.yml

@@ -11,5 +11,5 @@ runs:
         git fetch origin $CURRENT_SHA
         export COMMIT_MESSAGE=$(git show -s --format=%B "$CURRENT_SHA")
         export CURRENT_HEAD=$(git rev-parse HEAD)
-        cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && IGNORE_OOB=1 PYTHONPATH=. python3 process_replay.py
+        cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && CHECK_OOB=0 PYTHONPATH=. python3 process_replay.py
         git checkout $CURRENT_HEAD  # restore to branch

.github/actions/setup-tinygrad/action.yml

@@ -56,7 +56,15 @@ runs:
 
     # **** Caching packages ****
 
+    - name: Cache Python packages (PR)
+      if: github.event_name == 'pull_request'
+      id: restore-venv-pr
+      uses: actions/cache/restore@v4
+      with:
+        path: ${{ github.workspace }}/.venv
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
     - name: Cache Python packages
+      if: github.event_name != 'pull_request'
       id: restore-venv
       uses: actions/cache@v4
       with:
@@ -65,23 +73,23 @@ runs:
 
     # **** Caching downloads ****
 
-    - name: Cache downloads (Linux)
-      if: inputs.key != '' && runner.os == 'Linux'
-      uses: actions/cache@v4
+    - name: Cache downloads (PR)
+      if: inputs.key != '' && github.event_name == 'pull_request'
+      uses: actions/cache/restore@v4
       with:
-        path: ~/.cache/tinygrad/downloads/
+        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
         key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
-    - name: Cache downloads (macOS)
-      if: inputs.key != '' && runner.os == 'macOS'
+    - name: Cache downloads
+      if: inputs.key != '' && github.event_name != 'pull_request'
       uses: actions/cache@v4
       with:
-        path: ~/Library/Caches/tinygrad/downloads/
+        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
         key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
 
     # **** Python deps ****
 
     - name: Install dependencies in venv (with extra)
-      if: inputs.deps != '' && steps.restore-venv.outputs.cache-hit != 'true'
+      if: inputs.deps != '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
       shell: bash
       run: |
         python -m venv .venv
@@ -92,7 +100,7 @@ runs:
         fi
         python -m pip install -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
     - name: Install dependencies in venv (without extra)
-      if: inputs.deps == '' && steps.restore-venv.outputs.cache-hit != 'true'
+      if: inputs.deps == '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
       shell: bash
       run: |
         python -m venv .venv
@@ -137,7 +145,7 @@ runs:
       run: |
         wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
         sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
-        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.2 $(lsb_release -cs) main
+        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/7.1 $(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
 
@@ -182,8 +190,14 @@ runs:
         echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
         echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
 
+    - name: Cache apt (PR)
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name == 'pull_request'
+      uses: actions/cache/restore@v4
+      with:
+        path: /var/cache/apt/archives/
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
     - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name != 'pull_request'
       uses: actions/cache@v4
       with:
         path: /var/cache/apt/archives/
@@ -219,7 +233,7 @@ runs:
       shell: bash
       run: |
         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 | \
+        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
           jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
           sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
         cargo build --release --manifest-path ./extra/remu/Cargo.toml
@@ -239,8 +253,17 @@ runs:
         ln -s /opt/homebrew/opt/[email protected] /opt/homebrew/opt/boost || true
         ln -s /opt/homebrew/opt/boost/lib/libboost_atomic-mt.dylib /opt/homebrew/opt/boost/lib/libboost_atomic.dylib || true
         ln -s /opt/homebrew/opt/boost/lib/libboost_thread-mt.dylib /opt/homebrew/opt/boost/lib/libboost_thread.dylib || true
+    - name: Cache gpuocelot (PR)
+      if: inputs.ocelot == 'true' && github.event_name == 'pull_request'
+      id: cache-build-pr
+      uses: actions/cache/restore@v4
+      env:
+        cache-name: cache-gpuocelot-build-1
+      with:
+        path: ${{ github.workspace }}/gpuocelot/ocelot
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
     - name: Cache gpuocelot
-      if: inputs.ocelot == 'true'
+      if: inputs.ocelot == 'true' && github.event_name != 'pull_request'
       id: cache-build
       uses: actions/cache@v4
       env:
@@ -249,7 +272,7 @@ runs:
         path: ${{ github.workspace }}/gpuocelot/ocelot
         key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
     - name: Clone/compile gpuocelot
-      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
+      if: inputs.ocelot == 'true' && steps.cache-build-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true'
       shell: bash
       run: |
         git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot

.github/workflows/autogen.yml

@@ -14,10 +14,12 @@ on:
     paths:
     - 'tinygrad/runtime/autogen/**/*'
     - 'tinygrad/runtime/support/autogen.py'
+    - '.github/workflows/autogen.yml'
   workflow_dispatch:
     paths:
     - 'tinygrad/runtime/autogen/**/*'
     - 'tinygrad/runtime/support/autogen.py'
+    - '.github/workflows/autogen.yml'
 
 jobs:
   autogen:
@@ -30,6 +32,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
+        key: 'autogen'
         opencl: 'true'
         amd: 'true'
         cuda: 'true'
@@ -38,103 +41,37 @@ jobs:
         mesa: 'true'
         pydeps: 'pyyaml mako'
     - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
-    - 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 libdrm-dev
+    - name: Regenerate autogen files
       run: |
-        mv tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
+        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "comgr.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"
-        diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
-    - name: Verify CUDA autogen
-      run: |
-        mv tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
-        mv tinygrad/runtime/autogen/nvrtc.py /tmp/nvrtc.py.bak
-        mv tinygrad/runtime/autogen/nvjitlink.py /tmp/nvjitlink.py.bak
-        mv tinygrad/runtime/autogen/nv_570.py /tmp/nv_570.py.bak
-        mv tinygrad/runtime/autogen/nv.py /tmp/nv.py.bak
-        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv"
-        diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
-        diff /tmp/nvrtc.py.bak tinygrad/runtime/autogen/nvrtc.py
-        diff /tmp/nvjitlink.py.bak tinygrad/runtime/autogen/nvjitlink.py
-        diff /tmp/nv_570.py.bak tinygrad/runtime/autogen/nv_570.py
-        diff /tmp/nv.py.bak tinygrad/runtime/autogen/nv.py
-    - name: Verify AMD autogen
-      run: |
-        mv tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
-        mv tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
-        mv tinygrad/runtime/autogen/hip.py /tmp/hip.py.bak
-        mv tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
-        mv tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
-        mv tinygrad/runtime/autogen/rocprof.py /tmp/rocprof.py.bak
-        mv tinygrad/runtime/autogen/am/am.py /tmp/am_am.py.bak
-        mv tinygrad/runtime/autogen/am/pm4_soc15.py /tmp/am_pm4_soc15.py.bak
-        mv tinygrad/runtime/autogen/am/pm4_nv.py /tmp/am_pm4_nv.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_4_0_0.py /tmp/am_sdma_4_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_5_0_0.py /tmp/am_sdma_5_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_6_0_0.py /tmp/am_sdma_6_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/smu_v13_0_0.py /tmp/am_smu_v13_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/smu_v14_0_2.py /tmp/am_smu_v14_0_2.py.bak
-        python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof; from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v14_0_2"
-        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
-        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
-        diff /tmp/hip.py.bak tinygrad/runtime/autogen/hip.py
-        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
-        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
-        diff /tmp/rocprof.py.bak tinygrad/runtime/autogen/rocprof.py
-        diff /tmp/am_am.py.bak tinygrad/runtime/autogen/am/am.py
-        diff /tmp/am_pm4_soc15.py.bak tinygrad/runtime/autogen/am/pm4_soc15.py
-        diff /tmp/am_pm4_nv.py.bak tinygrad/runtime/autogen/am/pm4_nv.py
-        diff /tmp/am_sdma_4_0_0.py.bak tinygrad/runtime/autogen/am/sdma_4_0_0.py
-        diff /tmp/am_sdma_5_0_0.py.bak tinygrad/runtime/autogen/am/sdma_5_0_0.py
-        diff /tmp/am_sdma_6_0_0.py.bak tinygrad/runtime/autogen/am/sdma_6_0_0.py
-        diff /tmp/am_smu_v13_0_0.py.bak tinygrad/runtime/autogen/am/smu_v13_0_0.py
-        diff /tmp/am_smu_v14_0_2.py.bak tinygrad/runtime/autogen/am/smu_v14_0_2.py
-    - name: Verify Linux autogen
-      run: |
-        mv tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
-        mv tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
-        mv tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
-        mv tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
-        mv tinygrad/runtime/autogen/pci.py /tmp/pci.py.bak
-        mv tinygrad/runtime/autogen/vfio.py /tmp/vfio.py.bak
+        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
+        python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
+        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_v13_0_12, smu_v14_0_2"
         python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
-        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
-        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
-        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
-        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
-        diff /tmp/pci.py.bak tinygrad/runtime/autogen/pci.py
-        diff /tmp/vfio.py.bak tinygrad/runtime/autogen/vfio.py
-    - name: Verify LLVM autogen
-      run: |
-        mv tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
         python3 -c "from tinygrad.runtime.autogen import llvm"
-        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
-    - name: Verify WebGPU autogen
-      run: |
-        mv tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
         python3 -c "from tinygrad.runtime.autogen import webgpu"
-        diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
-    - name: Verify Qualcomm autogen
-      run: |
-        mv tinygrad/runtime/autogen/kgsl.py /tmp/kgsl.py.bak
-        mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
         python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
-        diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
-        diff /tmp/qcom_dsp.py.bak tinygrad/runtime/autogen/qcom_dsp.py
-    - name: Verify libusb autogen
-      run: |
-        mv tinygrad/runtime/autogen/libusb.py /tmp/libusb.py.bak
         python3 -c "from tinygrad.runtime.autogen import libusb"
-        diff /tmp/libusb.py.bak tinygrad/runtime/autogen/libusb.py
-    - name: Verify mesa autogen
-      run: |
-        mv tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
         python3 -c "from tinygrad.runtime.autogen import mesa"
-        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
-    - name: Verify libclang autogen
-      run: |
-        cp tinygrad/runtime/autogen/libclang.py /tmp/libclang.py.bak
+        python3 -c "from tinygrad.runtime.autogen import avcodec"
         REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
-        diff /tmp/libclang.py.bak tinygrad/runtime/autogen/libclang.py
+    - name: Check for differences
+      run: |
+        if ! git diff --quiet; then
+          git diff
+          git diff > autogen-ubuntu.patch
+          echo "Autogen mismatch detected. Patch available at: ${{ 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
@@ -145,14 +82,29 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
+        key: 'autogen-mac'
         llvm: 'true'
-    - name: Verify macos autogen
+    - name: Regenerate autogen files
       run: |
-        mv tinygrad/runtime/autogen/metal.py /tmp/metal.py.bak
-        LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
-        diff /tmp/metal.py.bak tinygrad/runtime/autogen/metal.py
-  autogen-comgr-3:
-    name: In-tree Autogen (comgr 3)
+        rm tinygrad/runtime/autogen/metal.py tinygrad/runtime/autogen/iokit.py tinygrad/runtime/autogen/corefoundation.py
+        python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
+    - name: Check for differences
+      run: |
+        if ! git diff --quiet; then
+          git diff
+          git diff > autogen-macos.patch
+          echo "Autogen mismatch detected. Patch available at: ${{ 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-2:
+    name: In-tree Autogen (comgr 2)
     runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
@@ -160,17 +112,32 @@ jobs:
       uses: actions/checkout@v4
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
+      with:
+        key: 'autogen-comgr'
     - name: Install autogen support packages
       run: |
         wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
         sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
-        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.4 $(lsb_release -cs) main
+        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.2 $(lsb_release -cs) main
         EOF
         echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
         sudo apt -qq update || true
         sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
-    - name: Verify comgr (3) autogen
+    - name: Regenerate autogen files
       run: |
-        mv tinygrad/runtime/autogen/comgr_3.py /tmp/comgr_3.py.bak
-        python3 -c "from tinygrad.runtime.autogen import comgr_3"
-        diff /tmp/comgr_3.py.bak tinygrad/runtime/autogen/comgr_3.py
+        rm tinygrad/runtime/autogen/comgr.py
+        python3 -c "from tinygrad.runtime.autogen import comgr"
+    - name: Check for differences
+      run: |
+        if ! git diff --quiet; then
+          git diff
+          git diff > autogen-comgr2.patch
+          echo "Autogen mismatch detected. Patch available at: ${{ 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-comgr2-patch
+        path: autogen-comgr2.patch

.github/workflows/benchmark.yml

@@ -16,6 +16,43 @@ on:
   workflow_dispatch:
 
 jobs:
+  # the goal of this test is to replicate a normal person on a laptop running the test
+  # no process replay, no benchmarks, no CI, just a normal laptop person
+  # the 3 minute timeout should not be raised
+  testmacpytest:
+    name: Mac pytest
+    env:
+      CI: ""
+      CAPTURE_PROCESS_REPLAY: "0"
+    runs-on: [self-hosted, macOS]
+    timeout-minutes: 3
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    # brew install uv
+    - name: setup python environment
+      run: |
+        rm -rf /tmp/tinygrad_pytest_ci
+        uv venv /tmp/tinygrad_pytest_ci
+        source /tmp/tinygrad_pytest_ci/bin/activate
+        uv pip install .[testing]
+    - name: setup staging db
+      run: |
+        echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
+        rm -f /tmp/pytest-db-ci*
+    - name: Run pytest -nauto
+      run: |
+        source /tmp/tinygrad_pytest_ci/bin/activate
+        pytest -nauto --durations=20
+    - name: openpilot compile3 0.10.1 driving_vision
+      run: FLOAT16=1 CL=1 IMAGE=2 python3.11 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: FLOAT16=1 CL=1 IMAGE=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+
   testmacbenchmark:
     name: Mac Benchmark
     env:
@@ -49,19 +86,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=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=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 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
@@ -71,84 +108,59 @@ 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
-          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
 
@@ -170,6 +182,10 @@ jobs:
       run: |
         echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: Kill stale pids
+      run: |
+        PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+        PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
     - 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
@@ -215,7 +231,7 @@ 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 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Test benchmark allreduce
@@ -226,79 +242,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
 
@@ -337,44 +332,30 @@ jobs:
     # - name: Fuzz Padded Tensor Core GEMM (PTX)
     #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
     - name: HEVC Decode Benchmark
-      run: VALIDATE=1 MAX_FRAMES=100 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
+      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=120 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=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=120 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=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=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
+      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 | tee train_resnet.txt
+      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=72 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
 
@@ -389,10 +370,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: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
     #- name: Insert amdgpu
     #  run: sudo modprobe amdgpu
     - name: Symlink models and datasets
@@ -426,7 +409,7 @@ 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 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
@@ -437,7 +420,7 @@ jobs:
     - name: Run Tensor Core GEMM (AMD)
       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 | tee matmul_amd.txt
+        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
@@ -452,61 +435,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
+      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 | tee sdxl.txt
+      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
 
@@ -521,10 +482,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: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -543,31 +506,23 @@ 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=200 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=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=230 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
+      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 | tee train_cifar_six_gpu.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
+      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
+    # TODO: broken on some of the machines
+    #- name: Test full tinyfs load
+    #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
     - name: Run process replay tests
       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
 
@@ -582,10 +537,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: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -606,19 +563,12 @@ jobs:
     - 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
+      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 | tee train_resnet.txt
+      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=72 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
 
@@ -641,19 +591,21 @@ jobs:
     - name: reset process replay
       run: test/external/process_replay/reset.py
     - name: openpilot compile3 0.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
+      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=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
       run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      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
-      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
+      run: BENCHMARK_LOG=openpilot_0_10_1_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/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: benchmark MobileNetV2 on DSP
       run: |
         # generate quantized weights
@@ -676,10 +628,12 @@ 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
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
     - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py amd kill_pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -708,7 +662,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
@@ -718,20 +672,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
+      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
 
@@ -746,10 +692,12 @@ 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
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv rmmod
     - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py nv kill_pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -778,21 +726,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
+      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 | tee nv_train_resnet_one_gpu.txt
+      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

.github/workflows/test.yml

@@ -1,10 +1,11 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '15'
+  CACHE_VERSION: '17'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
+  CHECK_OOB: 1
 
 on:
   push:
@@ -25,17 +26,19 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: llvm-speed
-        deps: testing_minimal
+        deps: testing_unit
         llvm: 'true'
     - name: Speed Test
-      run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
     - name: Speed Test (BEAM=2)
-      run: BEAM=2 CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
 
   docs:
     name: Docs
     runs-on: ubuntu-22.04
     timeout-minutes: 10
+    env:
+      CHECK_OOB: 0
     steps:
     - name: Checkout Code
       uses: actions/checkout@v4
@@ -95,27 +98,23 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
+        deps: testing_unit
         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
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_tiny.py TestTiny.test_plus
     - name: Test ResNet-18
       run: DEBUG=2 python3 extra/torch_backend/example.py
-    - name: My (custom) tests
-      run: python3 extra/torch_backend/test.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
+      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/backend/test_ops.py --durations=20
     - name: Test in-place operations on views
       run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
     - name: Test multi-gpu
@@ -135,7 +134,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
+        deps: testing_unit
         llvm: 'true'
     - name: Install ninja
       run: |
@@ -157,27 +156,27 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: be-minimal
-        deps: testing_minimal
+        deps: testing_unit
     - name: Test dtype with Python emulator
-      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
+      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/backend/test_dtype.py test/backend/test_dtype_alu.py
     - name: Test ops with Python emulator
-      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20
+      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/backend/test_ops.py --durations=20
     - name: Test uops with Python emulator
-      run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
+      run: PYTHON=1 python3 -m pytest test/backend/test_uops.py --durations=20
     - name: Test symbolic with Python emulator
-      run: PYTHON=1 python3 test/test_symbolic_ops.py
+      run: PYTHON=1 python3 test/backend/test_symbolic_ops.py
     - name: test_renderer_failures with Python emulator
-      run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
+      run: PYTHON=1 python3 -m pytest -rA test/backend/test_renderer_failures.py::TestRendererFailures
     - name: Test IMAGE=2 support
       run: |
-        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_simple_conv2d
+        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
+        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
     - name: Test emulated METAL tensor cores
       run: |
-        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
+        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_big_gemm
         DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test emulated AMX tensor cores
-      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
     - name: Test emulated AMD tensor cores
       run: |
         DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
@@ -198,9 +197,9 @@ jobs:
         DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test emulated CUDA tensor cores
       run: |
-        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
+        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
         DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test emulated INTEL OpenCL tensor cores
       run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
@@ -208,18 +207,17 @@ jobs:
       run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - name: Test device flop counts
       run: |
-        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul
+        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStats.test_simple_matmul
 
   linter:
     name: Linters
     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
@@ -231,18 +229,20 @@ 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
         python3 -m ruff check extra/thunder/tiny/ --ignore E501 --ignore F841 --ignore E722
-    - name: Run mypy
+        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 .
+        grep -v autogen lineprecision.txt | awk 'NR>2 {lines+=$2; precise+=$3; imprecise+=$4; any+=$5; empty+=$6} END {t=lines-empty; printf "TOTAL: %d lines, %d precise (%.1f%%), %d imprecise (%.1f%%), %d any (%.1f%%)\n", t, precise, 100*precise/t, imprecise, 100*imprecise/t, any, 100*any/t}'
         cat lineprecision.txt
     - name: Run TYPED=1
-      run: TYPED=1 python -c "import tinygrad"
+      run: CHECK_OOB=0 DEV=CPU TYPED=1 python test/test_tiny.py
 
   unittest:
     name: Unit Tests
@@ -255,24 +255,30 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
-        key: unittest-12
-        pydeps: "pillow numpy ftfy regex"
+        key: unittest-13
+        pydeps: "pillow ftfy regex pre-commit"
         deps: testing_unit
+        llvm: 'true'
+        amd: '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 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
+        CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
+        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
     - 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
+      run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
     # TODO: too slow
     # - name: Run SDXL on NULL backend
     #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
     - name: Run Clip tests for SD MLPerf on NULL backend
       run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
     - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     # TODO: support fake weights
     #- name: Run LLaMA 7B on 4 fake devices
     #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -289,8 +295,8 @@ jobs:
         python extra/optimization/extract_dataset.py
         gzip -c /tmp/sops > extra/datasets/sops.gz
         #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 20000 lines
-      run: MAX_LINE_COUNT=20000 python sz.py
+    - name: Repo line count < 24000 lines
+      run: MAX_LINE_COUNT=24000 python sz.py
 
   spec:
     strategy:
@@ -310,7 +316,7 @@ jobs:
         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/test_custom_kernel.py --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 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
 
   fuzzing:
     name: Fuzzing
@@ -344,11 +350,11 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: gpu-image
-          deps: testing_minimal
+          deps: testing_unit
           opencl: 'true'
       - 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 -m pytest -n=auto test/backend/test_ops.py --durations=20
         # TODO: training is broken
         # CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
       - name: Run process replay tests
@@ -365,14 +371,14 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: gen-dataset
-          deps: testing_minimal
+          deps: testing
           opencl: 'true'
       - name: Generate Dataset
         run: CL=1 extra/optimization/generate_dataset.sh
       - name: Run Kernel Count Test
         run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
       - name: Run fused optimizer tests
-        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/test_optim.py -k "not muon"
+        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/backend/test_optim.py -k "not muon"
       - name: Upload artifact
         uses: actions/upload-artifact@v4
         with:
@@ -420,7 +426,7 @@ jobs:
         with:
           key: onnxoptc
           deps: testing
-          python-version: '3.11'
+          python-version: '3.12'
           llvm: 'true'
       - name: Test ONNX (CPU)
         run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -431,7 +437,7 @@ jobs:
       - name: Test Additional ONNX Ops (CPU)
         run: CPU=1 CPU_LLVM=0 python3 test/external/external_test_onnx_ops.py
       - name: Test Quantize ONNX
-        run: CPU=1 CPU_LLVM=0 python3 test/test_quantize_onnx.py
+        run: CPU=1 CPU_LLVM=0 python3 test/backend/test_quantize_onnx.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -448,7 +454,7 @@ jobs:
           key: onnxoptl
           deps: testing
           pydeps: "tensorflow==2.19"
-          python-version: '3.11'
+          python-version: '3.12'
           opencl: 'true'
       - name: Test ONNX (CL)
         run: CL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -461,11 +467,11 @@ jobs:
       - name: Test MLPerf stuff
         run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
       - name: NULL=1 beautiful_mnist_multigpu
-        run: NULL=1 python examples/beautiful_mnist_multigpu.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py
       - 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
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
       - name: Test llama 3 training
-        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=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
 
@@ -473,6 +479,8 @@ jobs:
     name: Test LLM
     runs-on: ubuntu-24.04
     timeout-minutes: 15
+    env:
+      CHECK_OOB: 0
     steps:
       - name: Checkout Code
         uses: actions/checkout@v4
@@ -520,7 +528,7 @@ jobs:
       with:
         key: metal
         deps: testing
-        python-version: '3.11'
+        python-version: '3.12'
     - name: Test models (Metal)
       run: METAL=1 python -m pytest -n=auto test/models --durations=20
     - name: Test LLaMA compile speed
@@ -539,15 +547,15 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: devectorize-minimal
-        deps: testing_minimal
+        deps: testing_unit
         pydeps: "pillow"
         llvm: "true"
     - name: Test LLVM=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py
     - name: Test LLVM=1 DEVECTORIZE=0 for model
       run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
     - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
+      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py
 
   testdsp:
     name: Linux (DSP)
@@ -560,8 +568,8 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: dsp-minimal
-        deps: testing_minimal
-        pydeps: "onnx==1.18.0 onnxruntime pillow"
+        deps: testing_unit
+        pydeps: "onnx==1.18.0 onnxruntime ml_dtypes"
         llvm: "true"
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v3
@@ -573,15 +581,15 @@ jobs:
         load: true
         tags: qemu-hexagon:latest
         cache-from: type=gha
-        cache-to: type=gha,mode=min
+        cache-to: ${{ github.event_name != 'pull_request' && 'type=gha,mode=min' || '' }}
     - name: Set MOCKDSP env
       run: printf "MOCKDSP=1" >> $GITHUB_ENV
     - name: Run test_tiny on DSP
       run: DEBUG=2 DSP=1 python test/test_tiny.py
     - name: Test transcendentals
-      run: CC=clang-20 DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
+      run: CC=clang-20 DEBUG=2 DSP=1 python test/backend/test_transcendental.py TestTranscendentalVectorized
     - name: Test quantize onnx
-      run: DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
+      run: DEBUG=2 DSP=1 python3 test/backend/test_quantize_onnx.py
 
   testwebgpu:
     name: Linux (WebGPU)
@@ -594,34 +602,106 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: webgpu-minimal
-        deps: testing_minimal
-        python-version: '3.11'
+        deps: testing_unit
+        python-version: '3.12'
         webgpu: 'true'
     - name: Check Device.DEFAULT (WEBGPU) and print some source
       run: |
         WEBGPU=1 python -c "from tinygrad import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT"
-        WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+        WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
     - 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/backend --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
+  testamdasm:
+    name: AMD ASM IDE
+    runs-on: ubuntu-24.04
+    timeout-minutes: 20
+    env:
+      AMD: 1
+      PYTHON_REMU: 1
+      MOCKGPU: 1
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: rdna3-emu
+          deps: testing_unit
+          amd: 'true'
+          python-version: '3.14'
+      - name: Verify AMD autogen is up to date
+        run: |
+          python -m tinygrad.renderer.amd.generate
+          git diff --exit-code tinygrad/runtime/autogen/amd/
+      - 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: Install rocprof-trace-decoder
+        run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
+      - name: Run AMD renderer tests
+        run: AMD_LLVM=0 python -m pytest -n=auto test/amd/ --durations 20
+      - name: Run AMD renderer tests (AMD_LLVM=1)
+        run: AMD_LLVM=1 python -m pytest -n=auto test/amd/ --durations 20
+      - name: Run SQTT profiling tests
+        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+      - name: Run AMD emulated tests on NULL backend
+        env:
+          AMD: 0
+        run: |
+          PYTHONPATH=. NULL=1 EMULATE=AMD python extra/mmapeak/mmapeak.py
+          PYTHONPATH=. NULL=1 EMULATE=AMD_CDNA4 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
+      - name: Run ASM matmul on MOCKGPU
+        run: PYTHONPATH="." AMD=1 MOCKGPU=1 N=256 python3 extra/gemm/amd_asm_matmul.py
+      - name: Run LLVM test
+        run: AMD_LLVM=1 python test/device/test_amd_llvm.py
+
+  testmockam:
+    name: Linux (am)
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    env:
+      AMD: 1
+      MOCKGPU: 1
+      AMD_IFACE: PCI
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: mockam
+          deps: testing_unit
+          amd: 'true'
+      - name: Run test_tiny on MOCKAM
+        run: python test/test_tiny.py
+      - name: Run test_tiny on MOCKAM USB
+        run: AMD_IFACE=USB python test/test_tiny.py
+      - name: Run test_hcq on MOCKAM
+        run: python -m pytest test/device/test_hcq.py
+
   testamd:
     strategy:
       fail-fast: false
       matrix:
         backend: [amd, amdllvm]
+        arch: [rdna3, rdna4]
+        #arch: [rdna3, rdna4, cdna4]
 
-    name: Linux (${{ matrix.backend }})
+    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
     runs-on: ubuntu-22.04
-    timeout-minutes: 20
+    timeout-minutes: 15
     env:
       AMD: 1
       MOCKGPU: 1
-      FORWARD_ONLY: 1
+      MOCKGPU_ARCH: ${{ matrix.arch }}
+      SKIP_SLOW_TEST: 1
       AMD_LLVM: ${{ matrix.backend == 'amdllvm' && '1' || matrix.backend != 'amdllvm' && '0' }}
     steps:
       - name: Checkout Code
@@ -630,27 +710,17 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: ${{ matrix.backend }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
           amd: 'true'
           llvm: ${{ matrix.backend == 'amdllvm' && 'true' }}
       - name: Check Device.DEFAULT and print some source
         run: |
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
-          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
-      - name: Run LLVM test
-        if: matrix.backend=='amdllvm'
-        run: python test/device/test_amd_llvm.py
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py --durations=20
-      - name: Run pytest (amd)
-        run: python -m pytest test/external/external_test_am.py --durations=20
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py test/external/external_test_am.py --durations=20
       - name: Run TRANSCENDENTAL math
-        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
-      - name: Run TestOps.test_add with SQTT
-        run: |
-          VIZ=1 PMC=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
-          VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
-          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
+        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -673,7 +743,7 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: ${{ matrix.backend }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
           cuda: 'true'
           ocelot: 'true'
       - name: Set env
@@ -681,10 +751,12 @@ jobs:
       - name: Check Device.DEFAULT and print some source
         run: |
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
-          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run pytest (cuda)
         # skip multitensor because it's slow
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
+        run: python -m pytest -n=auto test/backend --ignore test/backend/test_multitensor.py --durations=20
+      - name: Run TestOps.test_add with PMA
+        run: VIZ=-1 PMA=1 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -704,7 +776,7 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: ${{ matrix.backend }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
           opencl: ${{ matrix.backend == 'opencl' && 'true' }}
           llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
           mesa: ${{ matrix.backend == 'lvp' && 'true' }}
@@ -713,11 +785,11 @@ jobs:
       - name: Check Device.DEFAULT and print some source
         run: |
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
-          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run pytest (${{ matrix.backend }})
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
+        run: python -m pytest -n=auto test/backend --durations=20
       - name: Run TRANSCENDENTAL math
-        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
+        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -735,25 +807,29 @@ jobs:
       with:
         key: metal
         deps: testing
-        python-version: '3.11'
+        python-version: '3.12'
         amd: 'true'
         cuda: 'true'
         ocelot: 'true'
         llvm: 'true'
     - name: Run unit tests
       run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
     - name: Run ONNX
       run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
     - name: Test tensor core ops (fake)
-      run: METAL=1 DEBUG=3 TC=2 python test/test_ops.py TestOps.test_gemm
+      run: METAL=1 DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
     - name: Test tensor core ops (real)
-      run: METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_big_gemm
+      run: METAL=1 DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
     - name: Test Beam Search
       run: METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
+    - name: Test Device Specific
+      run: METAL=1 python3 -m pytest test/device/test_metal.py
     #- name: Fuzz Test linearizer
     #  run: METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
     - name: Run TRANSCENDENTAL math
-      run: METAL=1 TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
+      run: METAL=1 TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
     - name: Run pytest (amd)
       env:
         MOCKGPU: 1
@@ -776,6 +852,8 @@ jobs:
         NV_PTX: 1
         NV: 1
         FORWARD_ONLY: 1
+        # TODO: failing due to library loading error
+        CAPTURE_PROCESS_REPLAY: 0
       run: |
         python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
     - name: Run process replay tests
@@ -794,14 +872,14 @@ jobs:
         key: osx-webgpu
         deps: testing
         webgpu: 'true'
-    - name: Test infinity math in WGSL
-      run: WEBGPU=1 python -m pytest -n=auto test/test_renderer_failures.py::TestWGSLFailures::test_multiply_infinity --durations=20
     - name: Build WEBGPU Efficientnet
       run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m examples.compile_efficientnet
-    - name: Clean npm cache
-      run: npm cache clean --force
-    - name: Install Puppeteer
-      run: npm install puppeteer
+    - name: Run selected webgpu tests
+      run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m pytest -n=auto test/backend --durations=20
+    #- name: Clean npm cache
+    #  run: npm cache clean --force
+    #- name: Install Puppeteer
+    #  run: npm install puppeteer
     # this is also flaky
     #- name: Run WEBGPU Efficientnet
     #  run: node test/web/test_webgpu.js
@@ -833,8 +911,7 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: macos-${{ matrix.backend }}-minimal
-          deps: testing_minimal
-          pydeps: "capstone"
+          deps: testing_unit
           llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
           mesa: ${{ matrix.backend == 'lvp' && 'true' }}
       - name: Set env
@@ -844,7 +921,7 @@ jobs:
           python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
           DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run pytest (${{ matrix.backend }})
-        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
+        run: python3 -m pytest -n=auto test/backend --durations=20
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
       - name: Run macOS-specific unit test
@@ -877,12 +954,16 @@ jobs:
       - name: Run unit tests
         if: matrix.backend=='llvm'
         # test_newton_schulz hits RecursionError
-        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
+        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
+      - name: Run NULL backend tests
+        if: matrix.backend=='llvm'
+        shell: bash
+        run: CPU=0 CPU_LLVM=0 NULL=1 python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
       - name: Run pytest (${{ matrix.backend }})
         shell: bash
         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
+          python -m pytest -n=auto test/test_tiny.py test/backend/test_ops.py --durations=20
 
 # ****** Compile-only Tests ******
 
@@ -901,15 +982,15 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: compile-${{ matrix.backend }}
-          deps: testing_minimal
+          deps: testing_unit
           mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
-          python-version: '3.14'
+          python-version: '3.12'
       - name: Set env
         shell: bash
-        run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
+        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=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
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20

.gitignore

@@ -58,10 +58,13 @@ weights
 *.lprof
 comgr_*
 *.pkl
+!extra/sqtt/examples/**/*.pkl
 site/
 profile_stats
 *.log
 target
 .mypy_cache
 mutants
-.mutmut-cache
+.mutmut-cache
+dagre/
+graphlib/

.pre-commit-config.yaml

@@ -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
@@ -28,7 +28,7 @@ repos:
         pass_filenames: false
       - id: tests
         name: comprehensive test suite
-        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
+        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/unit/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
         language: system
         always_run: true
         pass_filenames: false

AGENTS.md

@@ -1,17 +0,0 @@
-# tinygrad agents
-
-Hello agent. You are one of the most talented programmers of your generation.
-
-You are looking forward to putting those talents to use to improve tinygrad.
-
-## philosophy
-
-tinygrad is a **tensor** library focused on beauty and minimalism, while still matching the functionality of PyTorch and JAX.
-
-Every line must earn its keep. Prefer readability over cleverness. We believe that if carefully designed, 10 lines can have the impact of 1000.
-
-Never mix functionality changes with whitespace changes. All functionality changes must be tested.
-
-## style
-
-Use **2-space indentation**, and keep lines to a maximum of **150 characters**. Match the existing style.

CLAUDE.md

@@ -1,227 +0,0 @@
-# Claude Code Guide for tinygrad
-
-## Architecture Overview
-
-tinygrad compiles tensor operations into optimized kernels. The pipeline:
-
-1. **Tensor** (`tensor.py`) - User-facing API, creates UOp graph
-2. **UOp** (`uop/ops.py`) - Unified IR for all operations (both tensor and kernel level)
-3. **Schedule** (`engine/schedule.py`, `schedule/`) - Converts tensor UOps to kernel UOps
-4. **Codegen** (`codegen/`) - Converts kernel UOps to device code
-5. **Runtime** (`runtime/`) - Device-specific execution
-
-## Key Concepts
-
-### UOp (Universal Operation)
-Everything is a UOp - tensors, operations, buffers, kernels. Key properties:
-- `op`: The operation type (Ops enum)
-- `dtype`: Data type
-- `src`: Tuple of source UOps
-- `arg`: Operation-specific argument
-- `tag`: Optional tag for graph transformations
-
-UOps are **immutable and cached** - creating the same UOp twice returns the same object (ucache).
-
-### PatternMatcher
-Used extensively for graph transformations:
-```python
-pm = PatternMatcher([
-  (UPat(Ops.ADD, src=(UPat.cvar("x"), UPat.cvar("x"))), lambda x: x * 2),
-])
-result = graph_rewrite(uop, pm)
-```
-
-### Schedule Cache
-Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.
-
-## Directory Structure
-
-```
-tinygrad/
-├── tensor.py          # Tensor class, user API
-├── device.py          # Buffer, device management
-├── dtype.py           # Data types
-├── helpers.py         # Utilities, environment vars
-├── uop/
-│   ├── ops.py         # UOp class, Ops enum, PatternMatcher
-│   ├── spec.py        # UOp type verification
-│   └── symbolic.py    # Symbolic math simplification
-├── engine/
-│   ├── schedule.py    # Schedule creation, caching
-│   ├── realize.py     # Tensor realization
-│   ├── jit.py         # JIT compilation
-│   └── memory.py      # Memory planning
-├── schedule/
-│   ├── rangeify.py    # Convert movements to ranges
-│   └── indexing.py    # Index calculations
-├── codegen/
-│   ├── kernel.py      # Kernel optimization
-│   └── uopgraph.py    # UOp graph transformations
-├── renderer/          # Code generation (CUDA, Metal, etc.)
-└── runtime/           # Device backends
-```
-
-## 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-4` - Increasing verbosity
-- `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
-
-## Style Notes
-
-- 2-space indentation, 150 char line limit
-- PatternMatchers should be defined at module level (slow to construct)
-- Prefer `graph_rewrite` over manual graph traversal
-- UOp methods like `.replace()` preserve tags unless explicitly changed
-- Use `.rtag(value)` to add tags to UOps
-
-## Lessons Learned
-
-### UOp ucache Behavior
-UOps are cached by their contents - creating a UOp with identical (op, dtype, src, arg) returns the **same object**. This means:
-- `uop.replace(tag=None)` on a tagged UOp returns the original untagged UOp if it exists in cache
-- Two UOps with same structure are identical (`is` comparison works)
-
-### Spec Validation
-When adding new UOp patterns, update `tinygrad/uop/spec.py`. Test with:
-```bash
-SPEC=2 python3 test/unit/test_something.py
-```
-Spec issues appear as `RuntimeError: SPEC ISSUE None: UOp(...)`.
-
-### Schedule Cache Key Normalization
-The schedule cache strips values from BIND nodes so different bound values (e.g., KV cache positions) hit the same cache entry:
-- `pm_pre_sched_cache`: BIND(DEFINE_VAR, CONST) → BIND(DEFINE_VAR) for cache key
-- `pm_post_sched_cache`: restores original BIND from context
-- When accessing `bind.src[1]`, check `len(bind.src) > 1` first (might be stripped)
-- Extract var_vals from `input_buffers` dict after graph_rewrite (avoids extra toposort)
-
-### Avoiding Extra Work
-- Use ctx dict from graph_rewrite to collect info during traversal instead of separate toposort
-- Only extract var_vals when schedule is non-empty (no kernels = no vars needed)
-- PatternMatchers are slow to construct - define at module level, not in functions
-
-### Readability Over Speed
-Don't add complexity for marginal performance gains. Simpler code that's slightly slower is often better:
-```python
-# BAD: "optimized" with extra complexity
-if has_afters:  # skip toposort if no AFTERs
-  after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
-
-# GOOD: simple, always works
-after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
-```
-The conditional check adds complexity, potential bugs, and often negligible speedup. Only optimize when profiling shows a real bottleneck.
-
-### Testing LLM Changes
-```bash
-# Quick smoke test
-echo "Hello" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b"
-
-# Check cache hits (should see "cache hit" after warmup)
-echo "Hello world" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b" 2>&1 | grep cache
-
-# Test with beam search
-echo "Hello" | BEAM=2 python tinygrad/apps/llm.py --model "llama3.2:1b"
-```
-
-## Common Patterns
-
-### Graph Transformation
-```python
-def my_transform(ctx, x):
-  # Return new UOp or None to skip
-  return x.replace(arg=new_arg)
-
-pm = PatternMatcher([
-  (UPat(Ops.SOMETHING, name="x"), my_transform),
-])
-result = graph_rewrite(input_uop, pm, ctx={})
-```
-
-### Finding Variables
-```python
-# Get all variables in a UOp graph
-variables = uop.variables()
-
-# Get bound variable values
-var, val = bind_uop.unbind()
-```
-
-### Shape Handling
-```python
-# Shapes can be symbolic (contain UOps)
-shape = tensor.shape  # tuple[sint, ...] where sint = int | UOp
-```
-
-## Performance Optimization
-
-When optimizing tinygrad internals:
-
-1. **Measure wall time, not just call counts** - Reducing `graph_rewrite` calls doesn't always improve wall time. The overhead of conditional checks can exceed the cost of the operation being skipped.
-
-2. **Profile each optimization individually** - Run benchmarks with and without each change to measure actual impact. Use `test/external/external_benchmark_schedule.py` for schedule/rewrite timing.
-
-3. **Early exits in hot paths are effective** - Simple checks like `if self.op is Ops.CONST: return self` in `simplify()` can eliminate many unnecessary `graph_rewrite` calls.
-
-4. **`graph_rewrite` is expensive** - Each call has overhead even for small graphs. Avoid calling it when the result is trivially known (e.g., simplifying a CONST returns itself).
-
-5. **Beware iterator overhead** - Checks like `all(x.op is Ops.CONST for x in self.src)` can be slower than just running the operation, especially for small sequences.
-
-6. **Verify cache hit rates before adding/keeping caches** - Measure actual hit rates with real workloads. A cache with 0% hit rate is pure overhead (e.g., `pm_cache` was removed because the algorithm guarantees each UOp is only passed to `pm_rewrite` once).
-
-7. **Use `TRACK_MATCH_STATS=2` to profile pattern matching** - This shows match rates and time per pattern. Look for patterns with 0% match rate that still cost significant time - these are pure overhead for that workload.
-
-8. **Cached properties beat manual traversal** - `backward_slice` uses `@functools.cached_property`. A DFS with early-exit sounds faster but is actually slower because it doesn't benefit from caching. The cache hit benefit often outweighs algorithmic improvements.
-
-9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.
-
-## Pattern Matching Profiling
-
-Use `TRACK_MATCH_STATS=2` to identify expensive patterns:
-
-```bash
-TRACK_MATCH_STATS=2 PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
-```
-
-Output format: `matches / attempts -- match_time / total_time ms -- location`
-
-Key patterns to watch (from ResNet50 benchmark):
-- `split_load_store`: ~146ms, 31% match rate - does real work
-- `simplify_valid`: ~75ms, 0% match rate in this workload - checks AND ops for INDEX in backward slice
-- `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches
-
-Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.

README.md

@@ -192,7 +192,7 @@ For more examples on how to run the full test suite please refer to the [CI work
 Some examples of running tests locally:
 ```sh
 python3 -m pip install -e '.[testing]'  # install extra deps for testing
-python3 test/test_ops.py                # just the ops tests
+python3 test/backend/test_ops.py        # just the ops tests
 python3 -m pytest test/                 # whole test suite
 ```
 

docs/developer/layout.md

@@ -10,7 +10,7 @@ Directories are listed in order of how they are processed.
 
 Group UOps into kernels.
 
-::: tinygrad.schedule.rangeify.get_rangeify_map
+::: tinygrad.schedule.rangeify.get_kernel_graph
     options:
         members: false
         show_labels: false

docs/runtime.md

@@ -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.
 

docs/tensor/elementwise.md

@@ -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

docs/tensor/movement.md

@@ -27,5 +27,6 @@
 ::: tinygrad.Tensor.flatten
 ::: tinygrad.Tensor.unflatten
 ::: tinygrad.Tensor.diag
+::: tinygrad.Tensor.diagonal
 ::: tinygrad.Tensor.roll
 ::: tinygrad.Tensor.rearrange

docs/tensor/ops.md

@@ -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

examples/anthropic_challenge.py

@@ -0,0 +1,196 @@
+from tinygrad import Tensor, dtypes, Context, getenv, UOp, fetch
+from tinygrad.uop.ops import Ops, PatternMatcher, UPat
+from tinygrad.uop.symbolic import symbolic
+from tinygrad.codegen import Renderer
+from tinygrad.codegen.opt import Opt, OptOps
+
+# ************************* implementation of the problem ************************
+
+def myhash(a: Tensor) -> Tensor:
+  a = (a + 0x7ED55D16) + (a << 12)
+  a = (a ^ 0xC761C23C) ^ (a >> 19)
+  a = (a + 0x165667B1) + (a << 5)
+  a = (a + 0xD3A2646C) ^ (a << 9)
+  a = (a + 0xFD7046C5) + (a << 3)
+  a = (a ^ 0xB55A4F09) ^ (a >> 16)
+  return a
+
+def select_with_where_tree(values: Tensor, relative_idx: Tensor) -> Tensor:
+  n = values.shape[0]
+  if n == 1: return values[0].expand(relative_idx.shape)
+
+  mid = n // 2
+  left = select_with_where_tree(values[:mid], relative_idx)
+  right = select_with_where_tree(values[mid:], relative_idx - mid)
+
+  go_left = relative_idx < mid
+  return go_left.where(left, right)
+
+def tree_traversal(forest: Tensor, val: Tensor, height: int, rounds: int, where_tree_threshold=3) -> Tensor:
+  # All walkers start at idx=0
+  idx = Tensor.zeros(val.shape, device=val.device, dtype=dtypes.uint32)
+
+  for r in range(rounds):
+    level = r % (height + 1)
+    level_start = (1 << level) - 1
+    level_size = 1 << level
+
+    if level == 0:
+      # At root (level 0), all walkers are at idx=0
+      # No gather needed, just broadcast the root value
+      node_val = forest[0].expand(val.shape)
+      idx = idx * 0  # Reset to 0
+    elif level <= where_tree_threshold:
+      # Small level: use where-tree
+      level_values = forest[level_start : level_start + level_size]
+      relative_idx = (idx - level_start)
+      node_val = select_with_where_tree(level_values, relative_idx)
+    else:
+      # Large level: use gather
+      node_val = forest.gather(0, idx)
+
+    val = myhash(val ^ node_val)
+    idx = (idx << 1) + (1 + (val & 1))
+
+    # No wrap check needed! At round 10 (level becomes 0), we reset idx above.
+
+  return val.contiguous(arg=(Opt(OptOps.UPCAST, 0, 8),))
+
+# ************************* renderer for VLIW machine *************************
+
+def loop_unrolling(sink:UOp):
+  rng = [x for x in sink.toposort() if x.op is Ops.RANGE]
+  if len(rng) == 0: return None
+  print(f"unrolling loop with size {rng[0].vmax+1}")
+  unrolled_sinks = [sink.substitute({rng[0]:rng[0].const_like(i)}).src[0] for i in range(rng[0].vmax+1)]
+  return UOp.sink(*unrolled_sinks, arg=sink.arg)
+
+global_addrs = []
+vliw_prepare = PatternMatcher([
+  # loop unrolling (should be a part of tinygrad)
+  (UPat(Ops.SINK, name="sink"), loop_unrolling),
+  # cast is fake
+  (UPat(Ops.CAST, name="c"), lambda c: c.src[0]),
+  # rewrites to hardcode the addresses in memory
+  (UPat(Ops.PARAM, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
+  # INDEX is just plus
+  (UPat(Ops.INDEX, name="i"), lambda i: i.src[0]+i.src[1]),
+])+symbolic
+
+class VLIWRenderer(Renderer):
+  has_local = False  # TODO: this should be the default / cleaned up
+  # this says this backend supports MULACC + more. decompositions uses this
+  code_for_op: dict = {Ops.MULACC: None, Ops.ADD: "+", Ops.MUL: "*",
+                       Ops.XOR: "^", Ops.AND: "&", Ops.OR: "|",
+                       Ops.SHL: "<<", Ops.SHR: ">>", Ops.CMPLT: "<"}
+  # this matcher runs while still in graph form
+  pre_matcher = vliw_prepare
+
+  def render(self, uops:list[UOp]):
+
+    # TODO: this is a minimal renderer. for low cycle count, make it good
+    # to get speed, you need to add VLIW packing
+    # to get under 1536 regs, you need to add a register allocator
+    # we left the fun parts to you
+
+    print(f"rendering with {len(uops)} uops")
+    reg, inst = 0, []
+    r: dict[UOp, int] = {}
+    for u in uops:
+      assert u.dtype.count in (1,8), "dtype count must be 1 or 8"
+
+      # dumb register allocator
+      if u.op not in {Ops.STORE, Ops.SINK, Ops.GEP}:
+        r[u] = reg
+        reg += u.dtype.count
+
+      # render UOps to instructions
+      match u.op:
+        case Ops.SINK:
+          inst.append({"flow": [("halt",)]})
+        case Ops.CONST:
+          inst.append({"load": [("const", r[u], u.arg)]})
+        case Ops.GEP:
+          # a GEP is just an alias to a special register in the vector
+          r[u] = r[u.src[0]] + u.arg[0]
+        case Ops.VECTORIZE:
+          if all(s == u.src[0] for s in u.src):
+            # if all sources are the same, we can broadcast
+            inst.append({"valu": [("vbroadcast", r[u], r[u.src[0]])]})
+          else:
+            # this is a copy into a contiguous chunk of registers
+            inst.extend({"flow": [("add_imm", r[u]+i, r[s], 0)]} for i,s in enumerate(u.src) if r[s] != r[u]+i)
+        case Ops.LOAD:
+          op = "vload" if u.dtype.count > 1 else "load"
+          inst.append({"load": [(op, r[u], r[u.src[0]])]})
+        case Ops.STORE:
+          op = "vstore" if u.src[1].dtype.count > 1 else "store"
+          inst.append({"store": [(op, r[u.src[0]], r[u.src[1]])]})
+        case Ops.MULACC:
+          assert u.dtype.count == 8
+          inst.append({"valu": [("multiply_add", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
+        case Ops.WHERE:
+          assert u.dtype.count == 8
+          inst.append({"flow": [("vselect", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
+        case _ if u.op in self.code_for_op:
+          cat = "valu" if u.dtype.count > 1 else "alu"
+          inst.append({cat: [(self.code_for_op[u.op], r[u], r[u.src[0]], r[u.src[1]])]})
+        case _:
+          raise NotImplementedError(f"unhandled op {u.op}")
+    return repr(inst)
+
+# ************************* test and render *************************
+
+import sys, types
+PROBLEM_URL = "https://raw.githubusercontent.com/anthropics/original_performance_takehome/refs/heads/main/tests/frozen_problem.py"
+sys.modules["problem"] = problem = types.ModuleType("problem")
+exec(fetch(PROBLEM_URL).read_text(), problem.__dict__)
+
+if __name__ == "__main__":
+  batch_size = getenv("BS", 256)
+  height = 10
+  rounds = getenv("ROUNDS", 16)
+
+  # build problem
+  tree = problem.Tree.generate(height)
+  inp = problem.Input.generate(tree, batch_size, rounds)
+  mem = problem.build_mem_image(tree, inp)
+  global_addrs.extend([mem[6], mem[6], mem[4]])  # output, input, forest
+
+  # *** verify the kernel in tinygrad compared to reference ***
+
+  forest_t = Tensor(tree.values, dtype=dtypes.uint32)
+  val_t = Tensor(inp.values, dtype=dtypes.uint32)
+
+  if getenv("VERIFY", 1):
+    # verify on normal tinygrad device
+    with Context(PCONTIG=2):
+      out = tree_traversal(forest_t, val_t, height, rounds)
+      val_out = out.tolist()
+    problem.reference_kernel(tree, inp)
+    assert val_out == inp.values
+    print("verification passed")
+
+  # *** render to device ***
+
+  from tinygrad.codegen import get_program
+  with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0):
+    out = tree_traversal(forest_t, val_t, height, rounds)
+    sink = out.schedule()[-1].ast
+    prg = get_program(sink, VLIWRenderer())
+
+  # *** run on Machine and compare ***
+
+  # NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
+  src = eval(prg.src)
+  max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
+  print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else "       <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
+  machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)
+  machine.run()
+  print(f"ran for {machine.cycle:5d} cycles" + ("" if machine.cycle <= 1363 else "  <-- EVEN CLAUDE GOT 1363"))
+
+  # compare to reference
+  ref_mem = mem.copy()
+  for _ in problem.reference_kernel2(ref_mem, {}): pass
+  assert machine.mem[mem[6]:mem[6]+mem[2]] == ref_mem[mem[6]:mem[6]+mem[2]]
+  print("compare passed!")

examples/audio_helpers.py

@@ -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

examples/hlb_cifar10.py

@@ -19,8 +19,8 @@ cifar_std = [0.24703225141799082, 0.24348516474564, 0.26158783926049628]
 BS, STEPS = getenv("BS", 512), getenv("STEPS", 1000)
 EVAL_BS = getenv("EVAL_BS", BS)
 GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
-assert BS % len(GPUS) == 0, f"{BS=} is not a multiple of {len(GPUS)=}, uneven multi GPU is slow"
-assert EVAL_BS % len(GPUS) == 0, f"{EVAL_BS=} is not a multiple of {len(GPUS)=}, uneven multi GPU is slow"
+assert BS % len(GPUS) == 0, f"{BS=} is not a multiple of {len(GPUS)=}"
+assert EVAL_BS % len(GPUS) == 0, f"{EVAL_BS=} is not a multiple of {len(GPUS)=}"
 
 class UnsyncedBatchNorm:
   def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1, num_devices=len(GPUS)):

examples/mlperf/dataloader.py

@@ -65,17 +65,7 @@ def loader_process(q_in, q_out, X:Tensor, seed):
       else:
         # pad data with training mean
         img = np.tile(np.array([[[123.68, 116.78, 103.94]]], dtype=np.uint8), (224, 224, 1))
-
-      # broken out
-      #img_tensor = Tensor(img.tobytes(), device='CPU')
-      #storage_tensor = X[idx].contiguous().realize().lazydata.base.realized
-      #storage_tensor._copyin(img_tensor.numpy())
-
-      # faster
-      X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
-
-      # ideal
-      #X[idx].assign(img.tobytes())   # NOTE: this is slow!
+      X[idx].flatten().assign(img.tobytes())
       q_out.put(idx)
     q_out.put(None)
 
@@ -213,12 +203,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))
@@ -263,8 +254,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
       x = random_brightness_augmentation(x)
       x = gaussian_noise(x)
 
-    X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
-    Y[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = y.tobytes()
+    X[idx].flatten().assign(x.tobytes())
+    Y[idx].flatten().assign(y.tobytes())
 
     queue_out.put(idx)
   queue_out.put(None)
@@ -378,12 +369,12 @@ def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue
       clipped_match_idxs = np.clip(match_idxs, 0, None)
       clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]
 
-      boxes[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
-      labels[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
-      matches[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
-      anchors[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = anchor.tobytes()
+      boxes[idx].flatten().assign(clipped_boxes.tobytes())
+      labels[idx].flatten().assign(clipped_labels.tobytes())
+      matches[idx].flatten().assign(match_idxs.tobytes())
+      anchors[idx].flatten().assign(anchor.tobytes())
 
-    imgs[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
+    imgs[idx].flatten().assign(img.tobytes())
 
     queue_out.put(idx)
   queue_out.put(None)
@@ -405,6 +396,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
       queue_in.put((idx, img, tgt))
 
   def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
+    shm_name = f"{shm_name}_{os.getpid()}"
     if os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
     shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(size))
     shm_tensor = Tensor.empty(*size, dtype=dtype, device=f"disk:/dev/shm/{shm_name}")
@@ -551,7 +543,7 @@ class BinIdxDataset:
     version, = struct.unpack("<Q", self.idx.read(8))
     assert version == 1, "unsupported index version"
     dtype_code, = struct.unpack("<B", self.idx.read(1))
-    self.dtype = {1:dtypes.uint8, 2:dtypes.int8, 3:dtypes.int16, 4:dtypes.int32, 5:dtypes.int64, 6:dtypes.float64, 7:dtypes.double, 8:dtypes.uint16}[dtype_code]
+    self.dtype = {1:np.dtype(np.uint8), 2:np.dtype(np.int8), 3:np.dtype(np.int16), 4:np.dtype(np.int32), 5:np.dtype(np.int64), 6:np.dtype(np.float64), 7:np.dtype(np.double), 8:np.dtype(np.uint16)}[dtype_code]
     self.count, = struct.unpack("<Q", self.idx.read(8))
     doc_count, = struct.unpack("<Q", self.idx.read(8))
 
@@ -568,7 +560,7 @@ class BinIdxDataset:
     self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
 
     # bin file
-    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
+    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin")).numpy()
 
   def _index(self, idx) -> tuple[int, int]:
     return int(self.pointers[idx]), int(self.sizes[idx])
@@ -577,7 +569,7 @@ class BinIdxDataset:
     ptr, size = self._index(idx)
     if length is None: length = size - offset
     ptr += offset * self.dtype.itemsize
-    return self.bin_t[ptr:ptr+length*self.dtype.itemsize].bitcast(self.dtype).to(None)
+    return self.bin_t[ptr:ptr+length*self.dtype.itemsize].view(self.dtype)
 
 # https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
 class GPTDataset:
@@ -636,7 +628,7 @@ class GPTDataset:
         sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
 
     # concat all parts
-    text = Tensor.cat(*sample_parts)
+    text = np.concatenate(sample_parts, axis=0)
 
     return text
 
@@ -779,7 +771,8 @@ def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:b
 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)
+    stacked = np.stack(batch, axis=0)
+    yield Tensor(stacked, device="NPY")
 
 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)

examples/mlperf/helpers.py

@@ -219,7 +219,18 @@ 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 {

examples/mlperf/model_train.py

@@ -3,7 +3,7 @@ from pathlib import Path
 import multiprocessing
 
 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
-from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker
 from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
 
@@ -1008,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
 
@@ -1085,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
@@ -1147,7 +1148,7 @@ def train_bert():
 
         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()
@@ -1160,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
@@ -1284,18 +1285,26 @@ def train_llama3():
   from extra.models.llama import Transformer
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
+  from examples.mlperf.optim import GradAccClipAdamW
+
+  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)
+  DATA_SEED          = config["DATA_SEED"]              = getenv("DATA_SEED", SEED)
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
   TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
   SMALL              = config["SMALL"]                  = getenv("SMALL", 0)
   SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
+  EVAL_SAMPLES       = config["EVAL_SAMPLES"]           = getenv("EVAL_SAMPLES", 5760 if not SMALL else 1024)
+  MAX_STEPS          = config["MAX_STEPS"]              = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS))
+  WARMUP_STEPS       = config["WARMUP_STEPS"]           = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
+  LR                 = config["LR"]                     = getenv("LR", 8e-5 * GBS / 1152)
+  END_LR             = config["END_LR"]                 = getenv("END_LR", 8e-7)
   EVAL_FREQ          = config["EVAL_FREQ"]              = getenv("EVAL_FREQ", 46080)
   EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
   EVAL_TARGET        = config["EVAL_TARGET"]            = getenv("EVAL_TARGET", 5.6)
@@ -1309,10 +1318,12 @@ def train_llama3():
   opt_adamw_weight_decay = 0.1
 
   opt_gradient_clip_norm = 1.0
-  opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
-  opt_learning_rate_decay_steps = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS)) - opt_learning_rate_warmup_steps
-  opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
-  opt_end_learning_rate = getenv("END_LR", 8e-7)
+  opt_learning_rate_warmup_steps = WARMUP_STEPS
+  opt_learning_rate_decay_steps = MAX_STEPS - opt_learning_rate_warmup_steps
+  opt_base_learning_rate = LR
+  opt_end_learning_rate = END_LR
+
+  Tensor.manual_seed(SEED)  # seed for weight initialization
 
   # ** init wandb **
   WANDB = getenv("WANDB")
@@ -1324,7 +1335,12 @@ def train_llama3():
   model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
   # vocab_size from the mixtral tokenizer
   if not SMALL: model_params |= {"vocab_size": 32000}
+  real_vocab_size = model_params['vocab_size']
+  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
+  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
   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
@@ -1339,6 +1355,8 @@ def train_llama3():
     for v in get_parameters(model):
       v.shard_(device, axis=None)
 
+    vocab_mask.shard_(device, axis=None)
+
   if (MP := getenv("MP", 1)) > 1:
     device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
     for k,v in get_state_dict(model).items():
@@ -1346,6 +1364,7 @@ def train_llama3():
       elif '.attention.wq' in k: v.shard_(device, axis=0)
       elif '.attention.wk' in k: v.shard_(device, axis=0)
       elif '.attention.wv' in k: v.shard_(device, axis=0)
+      elif '.attention.wqkv' in k: v.shard_(device, axis=0)
       elif '.attention.wo' in k: v.shard_(device, axis=1)
       elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
       elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
@@ -1358,8 +1377,17 @@ def train_llama3():
       # prevents memory spike on device 0
       v.realize()
 
-  optim = AdamW(get_parameters(model), lr=0.0,
-                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
+    vocab_mask.shard_(device, axis=2).realize()
+
+  optim_device = "CPU" if getenv("OFFLOAD_OPTIM") else None
+  optim = GradAccClipAdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
+                           eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
+
+  # init grads
+  for p in optim.params:
+    p.grad = p.zeros_like().contiguous().realize()
+  grads: list[Tensor] = [p.grad for p in optim.params]
+
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
 
   if resume_ckpt := getenv("RESUME_CKPT"):
@@ -1372,9 +1400,8 @@ def train_llama3():
     load_state_dict(scheduler, safe_load(fn), realize=False)
 
   @TinyJit
-  @Tensor.train()
-  def train_step(model, tokens:Tensor):
-    optim.zero_grad()
+  def minibatch(tokens:Tensor):
+    tokens = tokens.to(None)
     if (DP := getenv("DP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
       tokens = tokens.shard(device, 0)
@@ -1382,29 +1409,29 @@ def train_llama3():
       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 = vocab_mask.where(-float("inf"), 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
-    if not getenv("DISABLE_GRAD_CLIP_NORM"):
-      total_norm = Tensor(0.0, dtype=dtypes.float32, device=optim.params[0].device)
-      for p in optim.params:
-        total_norm += p.grad.float().square().sum()
-      total_norm = total_norm.sqrt().contiguous()
-      for p in optim.params:
-        p.grad = p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
+    assert all(p.grad is g for p,g in zip(optim.params, grads))
+    Tensor.realize(loss, *grads)
+    return loss.flatten().float().to("CPU")
 
+  @TinyJit
+  def optim_step():
     optim.step()
     scheduler.step()
 
+    for g in grads:
+      g.assign(g.zeros_like())
+
     lr = optim.lr
-    loss.realize(lr)
-    return loss, lr
+    Tensor.realize(lr, *grads)
+
+    return lr.float().to("CPU")
 
   @TinyJit
   @Tensor.train(False)
-  def eval_step(model, tokens:Tensor):
+  def eval_step(tokens:Tensor):
+    tokens = tokens.to(None)
     if (DP := getenv("DP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
       tokens = tokens.shard(device, 0)
@@ -1412,8 +1439,8 @@ def train_llama3():
       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:])
-    return loss.flatten().float()
+    loss = vocab_mask.where(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    return loss.flatten().float().to("CPU")
 
   # ** data iters **
   def fake_data(bs, samples):
@@ -1425,45 +1452,78 @@ def train_llama3():
       return fake_data(BS, SAMPLES)
     else:
       from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
+      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=DATA_SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
 
   if getenv("FAKEDATA", 0):
     eval_dataset = None
   else:
     from examples.mlperf.dataloader import get_llama3_dataset
-    eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
+    eval_dataset = get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
 
   def get_eval_iter():
     if eval_dataset is None:
-      return fake_data(EVAL_BS, 5760)
+      return fake_data(EVAL_BS, EVAL_SAMPLES)
     from examples.mlperf.dataloader import iterate_llama3_dataset
     return iterate_llama3_dataset(eval_dataset, EVAL_BS)
 
-  iter = get_train_iter()
+  num_params = sum(p.numel() for p in params) - model_params["vocab_size"]*model_params["dim"]
+  train_iter = get_train_iter()
   i, sequences_seen = resume_ckpt, 0
-  for tokens in tqdm(iter, total=SAMPLES//GBS):
+  step_times = []
+  while i < MAX_STEPS:
     GlobalCounters.reset()
     if getenv("TRAIN", 1):
-      t = time.perf_counter()
-      loss, lr = train_step(model, tokens)
-      loss = loss.float().item()
-      lr = lr.item()
+      profile_marker(f"train @ {i}")
+      st = time.perf_counter()
+
+      stopped = False
+      losses, data_time, dev_time = [], 0, 0
+      for _ in range(grad_acc):
+        ist = time.perf_counter()
+        try: tokens = next(train_iter)
+        except StopIteration:
+          stopped = True
+          break
+        mst = time.perf_counter()
+        data_time += mst - ist
+        losses.append(minibatch(tokens).item())
+        dev_time += time.perf_counter() - mst
+      if stopped: break
+
+      gt = time.perf_counter()
+      lr = optim_step().item()
+      et = time.perf_counter()
+
+      loss = sum(losses) / len(losses)
+      optim_time = et - gt
+      dev_time += optim_time
+      step_time = et - st
+      gbs_time = gt - st
+      if BENCHMARK: step_times.append(step_time)
 
       i += 1
-      sequences_seen += tokens.shape[0]
+      sequences_seen += GBS
 
-      sec = time.perf_counter()-t
       mem_gb = GlobalCounters.mem_used / 1e9
-      gflops = GlobalCounters.global_ops / 1e9 / sec
+      gflops = GlobalCounters.global_ops / 1e9 / dev_time
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
       tqdm.write(
-        f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS")
-
-      if (fname:=getenv("LOSS_FILE", "")):
-        with open(fname, "a") as f:
-          f.write(f"{i} {loss:.4f} {lr:.12f} {mem_gb:.2f}\n")
+          f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
+          f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
 
       if WANDB:
-        wandb.log({"lr": lr, "train/loss": loss, "train/step_time": sec, "train/GFLOPS": gflops, "train/sequences_seen": sequences_seen})
+        wandb.log({
+          "lr": lr, "train/loss": loss,
+          "train/step_time": step_time,
+          "train/gbs_time": gbs_time,
+          "train/optim_time": optim_time,
+          "train/dev_time": dev_time,
+          "train/data_time": data_time,
+          "train/mem": mem_gb,
+          "train/GFLOPS": gflops,
+          "train/MFU": mfu,
+          "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")
@@ -1475,16 +1535,29 @@ def train_llama3():
         fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
         safe_save(get_state_dict(scheduler), fn)
 
-    if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
+      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):
+      if EVAL_BS == 0: return
       tqdm.write(f"evaluating after {sequences_seen} sequences")
+      profile_marker(f"eval @ {i}")
 
       # run eval
       eval_losses = []
       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):
-        eval_losses += eval_step(model, tokens).tolist()
+      for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
+        eval_losses += eval_step(tokens).tolist()
+
+        if BENCHMARK and (j+1) == min(BENCHMARK, EVAL_SAMPLES//EVAL_BS):
+          return
+
       log_perplexity = Tensor(eval_losses).mean().float().item()
 
       tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
@@ -1570,7 +1643,7 @@ def train_stable_diffusion():
     loss, out_lr = loss.detach().to("CPU"), optimizer.lr.to("CPU")
     Tensor.realize(loss, out_lr)
     return loss, out_lr
-    
+
   # checkpointing takes ~9 minutes without this, and ~1 minute with this
   @TinyJit
   def ckpt_to_cpu():
@@ -1609,7 +1682,7 @@ def train_stable_diffusion():
     if i == 3:
       for _ in range(3): ckpt_to_cpu() # do this at the beginning of run to prevent OOM surprises when checkpointing
       print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
-      
+
     total_train_time = time.perf_counter() - train_start_time
     if WANDB:
       wandb.log({"train/loss": loss_item, "train/lr": lr_item, "train/loop_time_prev": loop_time, "train/dl_time": dl_time, "train/step": i,

examples/mlperf/optim.py

@@ -0,0 +1,47 @@
+from tinygrad.tensor import Tensor
+from tinygrad.dtype import dtypes
+from tinygrad.nn.optim import Optimizer
+from tinygrad.helpers import FUSE_OPTIM
+
+class GradAccClipAdamW(Optimizer):
+  def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
+    super().__init__(params, lr, device, fused)
+    self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay
+    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False).contiguous() for _ in [b1, b2])
+    self.m = self._new_optim_param()
+    self.v = self._new_optim_param()
+    self.grad_acc, self.clip_norm = grad_acc, clip_norm
+
+  def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
+    for i in range(len(grads)):
+      if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
+
+    if self.fused:
+      grads[0] = grads[0] / self.grad_acc
+      total_norm = grads[0].float().square().sum().sqrt()
+      grads[0] = (grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype)
+    else:
+      for i in range(len(grads)):
+        grads[i] = grads[i] / self.grad_acc
+      total_norm = Tensor.zeros((), dtype=dtypes.float32, device=self.device)
+      for g in grads:
+        total_norm += g.float().square().sum()
+      total_norm = total_norm.sqrt()
+      for i in range(len(grads)):
+        grads[i] = (grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)
+
+    ret = []
+    self.b1_t *= self.b1
+    self.b2_t *= self.b2
+    for i, (t, g) in enumerate(zip(params, grads)):
+      self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
+      self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
+      m_hat = self.m[i] / (1.0 - self.b1_t)
+      v_hat = self.v[i] / (1.0 - self.b2_t)
+      up = m_hat / (v_hat.sqrt() + self.eps)
+      ret.append((self.lr * up).cast(t.dtype))
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v
+
+  def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
+    up = up.shard_like(t) + self.lr.to(t.device) * self.wd * t.detach()
+    return t.detach() - up.cast(t.dtype)

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_1xMI300X/dev_beam.sh

@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
 
-export IGNORE_OOB=1
+export CHECK_OOB=0
 
 export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_beam.sh

@@ -5,7 +5,7 @@ export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
 
-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000
 
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_run.sh

@@ -8,7 +8,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
 export TRAIN_STEPS=3900
 
-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000
 
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh

@@ -11,7 +11,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
 export TRAIN_STEPS=3900
 
-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000
 
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI350X/dev_run_fa.sh

@@ -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 CHECK_OOB=0
+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

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI350X/dev_run_fp8.sh

@@ -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 CHECK_OOB=0
+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

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI350X/run_and_time.sh

@@ -11,7 +11,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
 export TRAIN_STEPS=3900
 
-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000
 
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_beam.sh

@@ -4,7 +4,7 @@ export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
 
-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000
 
 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_run.sh

@@ -4,7 +4,7 @@ export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
 
-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000
 
 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh

@@ -7,7 +7,7 @@ export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
 
-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000
 
 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_beam.sh

@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 
-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000
 
 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_run.sh

@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 
-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000
 
 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh

@@ -7,7 +7,7 @@ export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
 
-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000
 
 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+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="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export FAKEDATA=1 BENCHMARK=10
+if [ -z "$FULL_LAYERS" ]; then
+  export LLAMA_LAYERS=2
+fi
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -0,0 +1,37 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+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="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+export VIZ=${VIZ:--1}
+examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/startup_walltime.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+export FAKEDATA=1
+export NULL_ALLOW_COPYOUT=1
+export HIP_VISIBLE_DEVICES=""
+export DEV=NULL
+export JITBEAM=0
+export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

examples/stable_diffusion.py

@@ -6,7 +6,6 @@ import argparse, time
 from collections import namedtuple
 from typing import Dict, Any
 
-from PIL import Image
 import numpy as np
 from tinygrad import Device, GlobalCounters, dtypes, Tensor, TinyJit
 from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten, profile_marker
@@ -336,6 +335,7 @@ if __name__ == "__main__":
   print(x.shape)
 
   profile_marker("save image")
+  from PIL import Image
   im = Image.fromarray(x.numpy())
   print(f"saving {args.out}")
   im.save(args.out)

examples/test_pkl_imagenet.py

@@ -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))):

examples/tinychat/tinychat-browser/compile.py

@@ -48,7 +48,7 @@ def prepare_browser_chunks(model):
         weight_metadata = metadata.get(name, default)
         weight_metadata["parts"][part_num] = {"file": i, "file_start_pos": cursor, "size": size}
         metadata[name] = weight_metadata
-        data = bytes(state_dict[name].uop.base.realized.as_buffer())
+        data = bytes(state_dict[name].uop.base.realized.as_memoryview())
         data = data if not offsets else data[offsets[0]:offsets[1]]
         writer.write(data)
         cursor += size

examples/webgpu/stable_diffusion/compile.py

@@ -93,7 +93,7 @@ if __name__ == "__main__":
     forward: Any = None
 
   sub_steps = [
-    Step(name = "textModel", input = [Tensor.randn(1, 77)], forward = model.cond_stage_model.transformer.text_model),
+    Step(name = "textModel", input = [Tensor.randint(1, 77, low=0, high=49408, dtype=dtypes.int32)], forward = model.cond_stage_model.transformer.text_model),
     Step(name = "diffusor", input = [Tensor.randn(1, 77, 768), Tensor.randn(1, 77, 768), Tensor.randn(1,4,64,64), Tensor.rand(1), Tensor.randn(1), Tensor.randn(1), Tensor.randn(1)], forward = model),
     Step(name = "decoder", input = [Tensor.randn(1,4,64,64)], forward = model.decode),
     Step(name = "f16tof32", input = [Tensor.randn(2097120, dtype=dtypes.uint32)], forward = u32_to_f16)

examples/whisper.py

@@ -7,6 +7,7 @@ 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)

extra/amdpci/am_smi.py

@@ -92,7 +92,7 @@ class SMICtx:
     self.prev_terminal_width = 0
     self.prev_terminal_height = 0
 
-    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
+    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
     lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
     self.lspci = {l.split()[0]: l.split(" ", 1)[1] for l in lspci}
     for k,v in self.lspci.items():
@@ -153,8 +153,8 @@ class SMICtx:
     tables = {}
     for dev in self.devs:
       match dev.ip_ver[am.MP1_HWIP]:
-        case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableX_t
-        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTableV2_t
+        case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableV0_t
+        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTable_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
@@ -165,17 +165,17 @@ class SMICtx:
 
   def get_gfx_activity(self, dev, metrics):
     match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return max(0, min(100, self._smuq10_round(metrics.SocketGfxBusy)))
+      case (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): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
+      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):
     match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6):
+      case (13,0,6)|(13,0,12):
         temps = {
           "Hotspot": self._smuq10_round(metrics.MaxSocketTemperature),
           "HBM": self._smuq10_round(metrics.MaxHbmTemperature),
@@ -191,7 +191,7 @@ class SMICtx:
 
   def get_voltage(self, dev, metrics, compact=False):
     match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return {}
+      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]
@@ -205,38 +205,37 @@ class SMICtx:
   def get_gfx_freq(self, dev, metrics):
     if metrics is None: return 0
     match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.GfxclkFrequency[0])
+      case (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):
     match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.UclkFrequency)
+      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):
     match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.FclkFrequency)
+      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):
     match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return None, None
+      case (13,0,6)|(13,0,12): return None, None
       case _: return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm
 
   def get_power(self, dev, metrics):
     match dev.ip_ver[am.MP1_HWIP]:
       case (13,0,6): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
+      case (13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.SocketPowerLimit)
       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:
@@ -245,8 +244,8 @@ class SMICtx:
         entry = pt.entries[i]
 
         if (entry & am.AMDGPU_PTE_VALID) == 0: continue
-        if pt.lv!=am.AMDGPU_VM_PTB and not dev.gmc.is_pte_huge_page(pt.lv, entry):
-          pt_stack.append(AMPageTableEntry(dev, entry & 0x0000FFFFFFFFF000, lv=pt.lv+1))
+        if pt.lv < am.AMDGPU_VM_PDB0 and not dev.gmc.is_pte_huge_page(pt.lv, entry):
+          pt_stack.append(AMPageTableEntry(dev, dev.xgmi2paddr(entry & 0x0000FFFFFFFFF000), lv=pt.lv+1))
           continue
         if (entry & am.AMDGPU_PTE_SYSTEM) != 0: continue
         usage += (1 << ((9 * (3-pt.lv)) + 12))
@@ -280,7 +279,7 @@ class SMICtx:
       device_line = [f"{bold(dev.pcibus)} {trim(self.lspci[dev.pcibus[5:]], col_size - 20)}"] + [pad("", col_size)]
       activity_line = [f"GFX Activity {draw_bar(self.get_gfx_activity(dev, metrics) / 100, activity_line_width)}"] \
                     + [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
-                    + [f"MEM Usage    {draw_bar((mem_used / mem_total) / 100, activity_line_width, opt_text=mem_fmt)}"] \
+                    + [f"MEM Usage    {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \
 
       temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
       temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]

extra/amdpci/hive_reset.py

@@ -1,12 +1,18 @@
 #!/usr/bin/env python3
 
+import os
 from tinygrad.helpers import Context
 from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
+from tinygrad.runtime.support.hcq import FileIOInterface
 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]
+  for gpu in gpus:
+    drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
+    if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
+      raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
+  pcidevs = [PCIDevice("AM", gpu, bars=[0, 2, 5]) for gpu in gpus]
   amdevs = []
   with Context(DEBUG=2):
     for pcidev in pcidevs:

extra/amdpci/proclogs.py

@@ -19,8 +19,9 @@ amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
 amdev.AMDev = AMDFake
 from tinygrad.runtime.ops_amd import PCIIface
 
-def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
+def parse_amdgpu_logs(log_content, register_names=None, register_objects=None, *, only_xcc0: bool = False):
   register_map = register_names or {}
+  register_objs = register_objects or {}
 
   def replace_register(match):
     reg = match.group(1)
@@ -37,6 +38,28 @@ def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = Fal
   # remove timing prefix
   processed_log = re.sub(r'^\[\s*\d+(?:\.\d+)?\]\s*', '', processed_log, flags=re.MULTILINE)
 
+  # decode register values into field dicts
+  def decode_value(match):
+    reg_name = match.group(1)
+    xcc_part = match.group(2)  # "xcc=0 " or ""
+    val_str = match.group(3)
+    val = int(val_str, 16)
+
+    reg_obj = register_objs.get(reg_name)
+    if reg_obj is not None and reg_obj.fields:
+      fields = reg_obj.decode(val)
+      # show raw for unaccounted bits
+      accounted = 0
+      for name, (start, end) in reg_obj.fields.items():
+        accounted |= (((1 << (end - start + 1)) - 1) << start)
+      unaccounted = val & ~accounted
+      parts = {k: v for k, v in fields.items() if v != 0}
+      if unaccounted: parts['_raw_unaccounted'] = hex(unaccounted)
+      return f"register {reg_name}, {xcc_part}with value {val_str} {parts}"
+    return match.group(0)
+
+  processed_log = re.sub(r'register (reg\w+), ((?:xcc=\d+ )?)with value (0x[0-9a-fA-F]+)', decode_value, processed_log)
+
   # keep only xcc=0 lines (but keep lines with no xcc at all)
   if only_xcc0:
     kept = []
@@ -50,16 +73,18 @@ def main():
   only_xcc0 = bool(getenv("ONLY_XCC0", 0))
 
   reg_names = {}
+  reg_objs = {}
   dev = PCIIface(None, 0)
   for x, y in dev.dev_impl.__dict__.items():
     if isinstance(y, AMRegister):
       for xcc, addr in y.addr.items():
         reg_names[addr] = f"{x}, xcc={xcc}"
+      reg_objs[x] = y
 
   with open(sys.argv[1], 'r') as f:
     log_content = f.read()
 
-  processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
+  processed_log = parse_amdgpu_logs(log_content, reg_names, reg_objs, only_xcc0=only_xcc0)
 
   with open(sys.argv[2], 'w') as f:
     f.write(processed_log)

extra/assembly/rdna3/alu.py

@@ -1,254 +0,0 @@
-# Pure combinational ALU functions for RDNA3 emulation
-from __future__ import annotations
-import struct, math
-from typing import Callable
-from extra.assembly.rdna3.autogen import SOP1Op, SOP2Op, SOPCOp, SOPKOp, VOP1Op, VOP2Op, VOP3Op
-
-# Format base offsets for unified opcode space
-SOP2_BASE, SOP1_BASE, SOPC_BASE, SOPK_BASE = 0x000, 0x100, 0x200, 0x300
-VOP2_BASE, VOP1_BASE = 0x100, 0x180
-
-# Float conversion helpers
-_I, _f, _H, _e = struct.Struct('<I'), struct.Struct('<f'), struct.Struct('<H'), struct.Struct('<e')
-def f32(i: int) -> float: return _f.unpack(_I.pack(i & 0xffffffff))[0]
-def i32(f: float) -> int:
-  if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
-  try: return _I.unpack(_f.pack(f))[0]
-  except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
-def f16(i: int) -> float: return _e.unpack(_H.pack(i & 0xffff))[0]
-def i16(f: float) -> int:
-  if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
-  try: return _H.unpack(_e.pack(f))[0]
-  except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
-def sext(v: int, b: int) -> int: return v - (1 << b) if v & (1 << (b-1)) else v
-def clz(x: int) -> int: return 32 - x.bit_length() if x else 32
-def cls(x: int) -> int: x &= 0xffffffff; return 31 if x in (0, 0xffffffff) else clz(~x & 0xffffffff if x >> 31 else x) - 1
-def _cvt_i32_f32(v): return (0x7fffffff if v > 0 else 0x80000000) if math.isinf(v) else (0 if math.isnan(v) else max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff)
-def _cvt_u32_f32(v): return (0xffffffff if v > 0 else 0) if math.isinf(v) else (0 if math.isnan(v) or v < 0 else min(0xffffffff, int(v)))
-
-# SALU: op -> fn(s0, s1, scc_in) -> (result, scc_out)
-SALU: dict[int, Callable] = {
-  # SOP2
-  SOP2_BASE + SOP2Op.S_ADD_U32: lambda a, b, scc: ((a + b) & 0xffffffff, int((a + b) >= 0x100000000)),
-  SOP2_BASE + SOP2Op.S_SUB_U32: lambda a, b, scc: ((a - b) & 0xffffffff, int(b > a)),
-  SOP2_BASE + SOP2Op.S_ADDC_U32: lambda a, b, scc: ((r := a + b + scc) & 0xffffffff, int(r >= 0x100000000)),
-  SOP2_BASE + SOP2Op.S_SUBB_U32: lambda a, b, scc: ((a - b - scc) & 0xffffffff, int((b + scc) > a)),
-  SOP2_BASE + SOP2Op.S_ADD_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 32)) & 0xffffffff, int(((a >> 31) == (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
-  SOP2_BASE + SOP2Op.S_SUB_I32: lambda a, b, scc: ((r := sext(a, 32) - sext(b, 32)) & 0xffffffff, int(((a >> 31) != (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
-  SOP2_BASE + SOP2Op.S_AND_B32: lambda a, b, scc: ((r := a & b), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_OR_B32: lambda a, b, scc: ((r := a | b), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_XOR_B32: lambda a, b, scc: ((r := a ^ b), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_AND_NOT1_B32: lambda a, b, scc: ((r := a & (~b & 0xffffffff)), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_OR_NOT1_B32: lambda a, b, scc: ((r := a | (~b & 0xffffffff)), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_LSHL_B32: lambda a, b, scc: ((r := (a << (b & 0x1f)) & 0xffffffff), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_LSHR_B32: lambda a, b, scc: ((r := a >> (b & 0x1f)), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_ASHR_I32: lambda a, b, scc: ((r := sext(a, 32) >> (b & 0x1f)) & 0xffffffff, int(r != 0)),
-  SOP2_BASE + SOP2Op.S_MUL_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 32)) & 0xffffffff, scc),
-  SOP2_BASE + SOP2Op.S_MUL_HI_U32: lambda a, b, scc: (((a * b) >> 32) & 0xffffffff, scc),
-  SOP2_BASE + SOP2Op.S_MUL_HI_I32: lambda a, b, scc: (((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff, scc),
-  SOP2_BASE + SOP2Op.S_MIN_I32: lambda a, b, scc: (a, 1) if sext(a, 32) < sext(b, 32) else (b, 0),
-  SOP2_BASE + SOP2Op.S_MIN_U32: lambda a, b, scc: (a, 1) if a < b else (b, 0),
-  SOP2_BASE + SOP2Op.S_MAX_I32: lambda a, b, scc: (a, 1) if sext(a, 32) > sext(b, 32) else (b, 0),
-  SOP2_BASE + SOP2Op.S_MAX_U32: lambda a, b, scc: (a, 1) if a > b else (b, 0),
-  SOP2_BASE + SOP2Op.S_CSELECT_B32: lambda a, b, scc: (a if scc else b, scc),
-  SOP2_BASE + SOP2Op.S_BFE_U32: lambda a, b, scc: ((r := ((a >> (b & 0x1f)) & ((1 << ((b >> 16) & 0x7f)) - 1)) if (b >> 16) & 0x7f else 0), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_BFE_I32: lambda a, b, scc: ((r := sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := (b >> 16) & 0x7f) else 0), int(r != 0)),
-  SOP2_BASE + SOP2Op.S_PACK_LL_B32_B16: lambda a, b, scc: ((a & 0xffff) | ((b & 0xffff) << 16), scc),
-  SOP2_BASE + SOP2Op.S_PACK_LH_B32_B16: lambda a, b, scc: ((a & 0xffff) | (b & 0xffff0000), scc),
-  SOP2_BASE + SOP2Op.S_PACK_HH_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | (b & 0xffff0000), scc),
-  SOP2_BASE + SOP2Op.S_PACK_HL_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | ((b & 0xffff) << 16), scc),
-  SOP2_BASE + SOP2Op.S_ADD_F32: lambda a, b, scc: (i32(f32(a) + f32(b)), scc),
-  SOP2_BASE + SOP2Op.S_SUB_F32: lambda a, b, scc: (i32(f32(a) - f32(b)), scc),
-  SOP2_BASE + SOP2Op.S_MUL_F32: lambda a, b, scc: (i32(f32(a) * f32(b)), scc),
-  # SOP1
-  SOP1_BASE + SOP1Op.S_MOV_B32: lambda a, b, scc: (a, scc),
-  SOP1_BASE + SOP1Op.S_NOT_B32: lambda a, b, scc: ((r := (~a) & 0xffffffff), int(r != 0)),
-  SOP1_BASE + SOP1Op.S_BREV_B32: lambda a, b, scc: (int(f'{a & 0xffffffff:032b}'[::-1], 2), scc),
-  SOP1_BASE + SOP1Op.S_CLZ_I32_U32: lambda a, b, scc: (clz(a), scc),
-  SOP1_BASE + SOP1Op.S_CLS_I32: lambda a, b, scc: (cls(a), scc),
-  SOP1_BASE + SOP1Op.S_SEXT_I32_I8: lambda a, b, scc: (sext(a & 0xff, 8) & 0xffffffff, scc),
-  SOP1_BASE + SOP1Op.S_SEXT_I32_I16: lambda a, b, scc: (sext(a & 0xffff, 16) & 0xffffffff, scc),
-  SOP1_BASE + SOP1Op.S_ABS_I32: lambda a, b, scc: ((r := abs(sext(a, 32)) & 0xffffffff), int(r != 0)),
-  SOP1_BASE + SOP1Op.S_CVT_F32_I32: lambda a, b, scc: (i32(float(sext(a, 32))), scc),
-  SOP1_BASE + SOP1Op.S_CVT_F32_U32: lambda a, b, scc: (i32(float(a)), scc),
-  SOP1_BASE + SOP1Op.S_CVT_I32_F32: lambda a, b, scc: (_cvt_i32_f32(f32(a)), scc),
-  SOP1_BASE + SOP1Op.S_CVT_U32_F32: lambda a, b, scc: (_cvt_u32_f32(f32(a)), scc),
-  SOP1_BASE + SOP1Op.S_CEIL_F32: lambda a, b, scc: (i32(math.ceil(f32(a))), scc),
-  SOP1_BASE + SOP1Op.S_FLOOR_F32: lambda a, b, scc: (i32(math.floor(f32(a))), scc),
-  SOP1_BASE + SOP1Op.S_TRUNC_F32: lambda a, b, scc: (i32(math.trunc(f32(a))), scc),
-  SOP1_BASE + SOP1Op.S_RNDNE_F32: lambda a, b, scc: (i32(round(f32(a))), scc),
-  SOP1_BASE + SOP1Op.S_CVT_F16_F32: lambda a, b, scc: (i16(f32(a)), scc),
-  SOP1_BASE + SOP1Op.S_CVT_F32_F16: lambda a, b, scc: (i32(f16(a)), scc),
-  # SOPC
-  SOPC_BASE + SOPCOp.S_CMP_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 32))),
-  SOPC_BASE + SOPCOp.S_CMP_EQ_U32: lambda a, b, scc: (0, int(a == b)),
-  SOPC_BASE + SOPCOp.S_CMP_LG_U32: lambda a, b, scc: (0, int(a != b)),
-  SOPC_BASE + SOPCOp.S_CMP_GT_U32: lambda a, b, scc: (0, int(a > b)),
-  SOPC_BASE + SOPCOp.S_CMP_GE_U32: lambda a, b, scc: (0, int(a >= b)),
-  SOPC_BASE + SOPCOp.S_CMP_LT_U32: lambda a, b, scc: (0, int(a < b)),
-  SOPC_BASE + SOPCOp.S_CMP_LE_U32: lambda a, b, scc: (0, int(a <= b)),
-  SOPC_BASE + SOPCOp.S_BITCMP0_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) == 0)),
-  SOPC_BASE + SOPCOp.S_BITCMP1_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) != 0)),
-  # SOPK
-  SOPK_BASE + SOPKOp.S_MOVK_I32: lambda a, b, scc: (sext(b, 16) & 0xffffffff, scc),
-  SOPK_BASE + SOPKOp.S_CMOVK_I32: lambda a, b, scc: ((sext(b, 16) & 0xffffffff) if scc else a, scc),
-  SOPK_BASE + SOPKOp.S_ADDK_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 16)) & 0xffffffff, int(((a >> 31) == ((b >> 15) & 1)) and ((a >> 31) != ((r >> 31) & 1)))),
-  SOPK_BASE + SOPKOp.S_MULK_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 16)) & 0xffffffff, scc),
-  SOPK_BASE + SOPKOp.S_CMPK_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 16))),
-  SOPK_BASE + SOPKOp.S_CMPK_EQ_U32: lambda a, b, scc: (0, int(a == (b & 0xffff))),
-  SOPK_BASE + SOPKOp.S_CMPK_LG_U32: lambda a, b, scc: (0, int(a != (b & 0xffff))),
-  SOPK_BASE + SOPKOp.S_CMPK_GT_U32: lambda a, b, scc: (0, int(a > (b & 0xffff))),
-  SOPK_BASE + SOPKOp.S_CMPK_GE_U32: lambda a, b, scc: (0, int(a >= (b & 0xffff))),
-  SOPK_BASE + SOPKOp.S_CMPK_LT_U32: lambda a, b, scc: (0, int(a < (b & 0xffff))),
-  SOPK_BASE + SOPKOp.S_CMPK_LE_U32: lambda a, b, scc: (0, int(a <= (b & 0xffff))),
-}
-
-# VALU: op -> fn(s0, s1, s2) -> result
-VALU: dict[int, Callable] = {
-  # VOP2
-  VOP2_BASE + VOP2Op.V_ADD_F32: lambda a, b, c: i32(f32(a) + f32(b)),
-  VOP2_BASE + VOP2Op.V_SUB_F32: lambda a, b, c: i32(f32(a) - f32(b)),
-  VOP2_BASE + VOP2Op.V_SUBREV_F32: lambda a, b, c: i32(f32(b) - f32(a)),
-  VOP2_BASE + VOP2Op.V_MUL_F32: lambda a, b, c: i32(f32(a) * f32(b)),
-  VOP2_BASE + VOP2Op.V_MIN_F32: lambda a, b, c: i32(min(f32(a), f32(b))),
-  VOP2_BASE + VOP2Op.V_MAX_F32: lambda a, b, c: i32(max(f32(a), f32(b))),
-  VOP2_BASE + VOP2Op.V_ADD_NC_U32: lambda a, b, c: (a + b) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_SUB_NC_U32: lambda a, b, c: (a - b) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_SUBREV_NC_U32: lambda a, b, c: (b - a) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_AND_B32: lambda a, b, c: a & b,
-  VOP2_BASE + VOP2Op.V_OR_B32: lambda a, b, c: a | b,
-  VOP2_BASE + VOP2Op.V_XOR_B32: lambda a, b, c: a ^ b,
-  VOP2_BASE + VOP2Op.V_XNOR_B32: lambda a, b, c: (~(a ^ b)) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_LSHLREV_B32: lambda a, b, c: (b << (a & 0x1f)) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_LSHRREV_B32: lambda a, b, c: b >> (a & 0x1f),
-  VOP2_BASE + VOP2Op.V_ASHRREV_I32: lambda a, b, c: (sext(b, 32) >> (a & 0x1f)) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_MIN_I32: lambda a, b, c: a if sext(a, 32) < sext(b, 32) else b,
-  VOP2_BASE + VOP2Op.V_MAX_I32: lambda a, b, c: a if sext(a, 32) > sext(b, 32) else b,
-  VOP2_BASE + VOP2Op.V_MIN_U32: lambda a, b, c: min(a, b),
-  VOP2_BASE + VOP2Op.V_MAX_U32: lambda a, b, c: max(a, b),
-  VOP2_BASE + VOP2Op.V_MUL_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_MUL_HI_I32_I24: lambda a, b, c: ((sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) >> 32) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_MUL_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff)) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_MUL_HI_U32_U24: lambda a, b, c: (((a & 0xffffff) * (b & 0xffffff)) >> 32) & 0xffffffff,
-  VOP2_BASE + VOP2Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
-  VOP2_BASE + VOP2Op.V_LDEXP_F16: lambda a, b, c: i16(math.ldexp(f16(a), sext(b, 32))),
-  VOP2_BASE + VOP2Op.V_ADD_F16: lambda a, b, c: i16(f16(a) + f16(b)),
-  VOP2_BASE + VOP2Op.V_SUB_F16: lambda a, b, c: i16(f16(a) - f16(b)),
-  VOP2_BASE + VOP2Op.V_MUL_F16: lambda a, b, c: i16(f16(a) * f16(b)),
-  VOP2_BASE + VOP2Op.V_MIN_F16: lambda a, b, c: i16(min(f16(a), f16(b))),
-  VOP2_BASE + VOP2Op.V_MAX_F16: lambda a, b, c: i16(max(f16(a), f16(b))),
-  # VOP1
-  VOP1_BASE + VOP1Op.V_MOV_B32: lambda a, b, c: a,
-  VOP1_BASE + VOP1Op.V_NOT_B32: lambda a, b, c: (~a) & 0xffffffff,
-  VOP1_BASE + VOP1Op.V_BFREV_B32: lambda a, b, c: int(f'{a & 0xffffffff:032b}'[::-1], 2),
-  VOP1_BASE + VOP1Op.V_CLZ_I32_U32: lambda a, b, c: clz(a),
-  VOP1_BASE + VOP1Op.V_CLS_I32: lambda a, b, c: cls(a),
-  VOP1_BASE + VOP1Op.V_CVT_F32_I32: lambda a, b, c: i32(float(sext(a, 32))),
-  VOP1_BASE + VOP1Op.V_CVT_F32_U32: lambda a, b, c: i32(float(a)),
-  VOP1_BASE + VOP1Op.V_CVT_I32_F32: lambda a, b, c: _cvt_i32_f32(f32(a)),
-  VOP1_BASE + VOP1Op.V_CVT_U32_F32: lambda a, b, c: _cvt_u32_f32(f32(a)),
-  VOP1_BASE + VOP1Op.V_CVT_F16_F32: lambda a, b, c: i16(f32(a)),
-  VOP1_BASE + VOP1Op.V_CVT_F32_F16: lambda a, b, c: i32(f16(a)),
-  VOP1_BASE + VOP1Op.V_RCP_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
-  VOP1_BASE + VOP1Op.V_RCP_IFLAG_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
-  VOP1_BASE + VOP1Op.V_RSQ_F32: lambda a, b, c: i32(1.0 / math.sqrt(f32(a)) if f32(a) > 0 else (float('nan') if f32(a) < 0 else float('inf'))),
-  VOP1_BASE + VOP1Op.V_SQRT_F32: lambda a, b, c: i32(math.sqrt(f32(a)) if f32(a) >= 0 else float('nan')),
-  VOP1_BASE + VOP1Op.V_LOG_F32: lambda a, b, c: i32(math.log2(f32(a)) if f32(a) > 0 else (float('-inf') if f32(a) == 0 else float('nan'))),
-  VOP1_BASE + VOP1Op.V_EXP_F32: lambda a, b, c: i32(float('inf') if f32(a) > 128 else (0.0 if f32(a) < -150 else math.pow(2.0, f32(a)))),
-  VOP1_BASE + VOP1Op.V_SIN_F32: lambda a, b, c: i32(math.sin(f32(a) * 2 * math.pi)),
-  VOP1_BASE + VOP1Op.V_COS_F32: lambda a, b, c: i32(math.cos(f32(a) * 2 * math.pi)),
-  VOP1_BASE + VOP1Op.V_FLOOR_F32: lambda a, b, c: i32(math.floor(f32(a))),
-  VOP1_BASE + VOP1Op.V_CEIL_F32: lambda a, b, c: i32(math.ceil(f32(a))),
-  VOP1_BASE + VOP1Op.V_TRUNC_F32: lambda a, b, c: i32(math.trunc(f32(a))),
-  VOP1_BASE + VOP1Op.V_RNDNE_F32: lambda a, b, c: i32(round(f32(a))),
-  VOP1_BASE + VOP1Op.V_FRACT_F32: lambda a, b, c: i32((v := f32(a)) - math.floor(v)),
-  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE0: lambda a, b, c: i32(float(a & 0xff)),
-  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE1: lambda a, b, c: i32(float((a >> 8) & 0xff)),
-  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE2: lambda a, b, c: i32(float((a >> 16) & 0xff)),
-  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE3: lambda a, b, c: i32(float((a >> 24) & 0xff)),
-  VOP1_BASE + VOP1Op.V_FREXP_MANT_F32: lambda a, b, c: i32(math.frexp(v)[0] if (v := f32(a)) != 0 else 0.0),
-  VOP1_BASE + VOP1Op.V_FREXP_EXP_I32_F32: lambda a, b, c: (math.frexp(v)[1] if (v := f32(a)) != 0 else 0) & 0xffffffff,
-  # VOP3
-  VOP3Op.V_FMA_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
-  VOP3Op.V_DIV_FMAS_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
-  VOP3Op.V_ADD3_U32: lambda a, b, c: (a + b + c) & 0xffffffff,
-  VOP3Op.V_LSHL_ADD_U32: lambda a, b, c: ((a << (b & 0x1f)) + c) & 0xffffffff,
-  VOP3Op.V_ADD_LSHL_U32: lambda a, b, c: ((a + b) << (c & 0x1f)) & 0xffffffff,
-  VOP3Op.V_XOR3_B32: lambda a, b, c: a ^ b ^ c,
-  VOP3Op.V_OR3_B32: lambda a, b, c: a | b | c,
-  VOP3Op.V_AND_OR_B32: lambda a, b, c: (a & b) | c,
-  VOP3Op.V_LSHL_OR_B32: lambda a, b, c: ((a << (b & 0x1f)) | c) & 0xffffffff,
-  VOP3Op.V_XAD_U32: lambda a, b, c: ((a ^ b) + c) & 0xffffffff,
-  VOP3Op.V_MAD_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff) + c) & 0xffffffff,
-  VOP3Op.V_MAD_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24) + sext(c, 32)) & 0xffffffff,
-  VOP3Op.V_BFE_U32: lambda a, b, c: (a >> (b & 0x1f)) & ((1 << (c & 0x1f)) - 1) if c & 0x1f else 0,
-  VOP3Op.V_BFE_I32: lambda a, b, c: sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := c & 0x1f) else 0,
-  VOP3Op.V_ALIGNBIT_B32: lambda a, b, c: (((a << 32) | b) >> (c & 0x1f)) & 0xffffffff,
-  VOP3Op.V_MUL_LO_U32: lambda a, b, c: (a * b) & 0xffffffff,
-  VOP3Op.V_MUL_HI_U32: lambda a, b, c: ((a * b) >> 32) & 0xffffffff,
-  VOP3Op.V_MUL_HI_I32: lambda a, b, c: ((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff,
-  VOP3Op.V_LDEXP_F32: lambda a, b, c: i32(math.ldexp(f32(a), sext(b, 32))),
-  VOP3Op.V_DIV_FIXUP_F32: lambda a, b, c: i32(math.copysign(float('inf'), f32(c)) if f32(b) == 0.0 else f32(c) / f32(b)),
-  VOP3Op.V_PACK_B32_F16: lambda a, b, c: (a & 0xffff) | ((b & 0xffff) << 16),
-  VOP3Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
-  VOP3Op.V_LSHLREV_B16: lambda a, b, c: ((b & 0xffff) << (a & 0xf)) & 0xffff,
-  VOP3Op.V_LSHRREV_B16: lambda a, b, c: (b & 0xffff) >> (a & 0xf),
-  VOP3Op.V_ASHRREV_I16: lambda a, b, c: (sext(b & 0xffff, 16) >> (a & 0xf)) & 0xffff,
-  VOP3Op.V_ADD_NC_U16: lambda a, b, c: ((a & 0xffff) + (b & 0xffff)) & 0xffff,
-  VOP3Op.V_SUB_NC_U16: lambda a, b, c: ((a & 0xffff) - (b & 0xffff)) & 0xffff,
-  VOP3Op.V_MUL_LO_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff)) & 0xffff,
-  VOP3Op.V_MIN_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff),
-  VOP3Op.V_MAX_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff),
-  VOP3Op.V_MIN_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) < sext(b & 0xffff, 16) else (b & 0xffff),
-  VOP3Op.V_MAX_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) > sext(b & 0xffff, 16) else (b & 0xffff),
-  VOP3Op.V_MAD_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff) + (c & 0xffff)) & 0xffff,
-  VOP3Op.V_MAD_I16: lambda a, b, c: (sext(a & 0xffff, 16) * sext(b & 0xffff, 16) + sext(c & 0xffff, 16)) & 0xffff,
-  VOP3Op.V_FMA_F16: lambda a, b, c: i16(f16(a) * f16(b) + f16(c)),
-  VOP3Op.V_MIN3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[0] & 0xffffffff,
-  VOP3Op.V_MAX3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[2] & 0xffffffff,
-  VOP3Op.V_MED3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[1] & 0xffffffff,
-  VOP3Op.V_MIN3_F16: lambda a, b, c: i16(min(f16(a), f16(b), f16(c))),
-  VOP3Op.V_MAX3_F16: lambda a, b, c: i16(max(f16(a), f16(b), f16(c))),
-  VOP3Op.V_MED3_F16: lambda a, b, c: i16(sorted([f16(a), f16(b), f16(c)])[1]),
-  VOP3Op.V_MIN3_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff, c & 0xffff),
-  VOP3Op.V_MAX3_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff, c & 0xffff),
-  VOP3Op.V_MED3_U16: lambda a, b, c: sorted([a & 0xffff, b & 0xffff, c & 0xffff])[1],
-  VOP3Op.V_MIN3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[0] & 0xffff,
-  VOP3Op.V_MAX3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[2] & 0xffff,
-  VOP3Op.V_MED3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[1] & 0xffff,
-}
-
-def _cmp8(a, b): return [False, a < b, a == b, a <= b, a > b, a != b, a >= b, True]
-def _cmp6(a, b): return [a < b, a == b, a <= b, a > b, a != b, a >= b]
-
-def vopc(op: int, s0: int, s1: int, s0_hi: int = 0, s1_hi: int = 0) -> int:
-  base = op & 0x7f
-  if 16 <= base <= 31:  # F32
-    f0, f1, cmp, nan = f32(s0), f32(s1), base - 16, math.isnan(f32(s0)) or math.isnan(f32(s1))
-    return int([False, f0<f1, f0==f1, f0<=f1, f0>f1, f0!=f1, f0>=f1, not nan, nan, f0<f1 or nan, f0==f1 or nan, f0<=f1 or nan, f0>f1 or nan, f0!=f1 or nan, f0>=f1 or nan, True][cmp])
-  if 49 <= base <= 54: return int(_cmp6(sext(s0 & 0xffff, 16), sext(s1 & 0xffff, 16))[base - 49])  # I16
-  if 57 <= base <= 62: return int(_cmp6(s0 & 0xffff, s1 & 0xffff)[base - 57])  # U16
-  if 64 <= base <= 79:  # I32/U32
-    cmp = (base - 64) % 8
-    return int(_cmp8(sext(s0, 32), sext(s1, 32))[cmp] if base < 72 else _cmp8(s0, s1)[cmp])
-  if 80 <= base <= 95:  # I64/U64
-    s0_64, s1_64 = s0 | (s0_hi << 32), s1 | (s1_hi << 32)
-    return int(_cmp8(sext(s0_64, 64), sext(s1_64, 64))[(base - 80) % 8] if base < 88 else _cmp8(s0_64, s1_64)[(base - 80) % 8])
-  if base == 126:  # CLASS_F32
-    f, mask = f32(s0), s1
-    if math.isnan(f): return int(bool(mask & 0x3))
-    if math.isinf(f): return int(bool(mask & (0x4 if f < 0 else 0x200)))
-    if f == 0.0: return int(bool(mask & (0x20 if (s0 >> 31) & 1 else 0x40)))
-    exp, sign = (s0 >> 23) & 0xff, (s0 >> 31) & 1
-    return int(bool(mask & ((0x10 if sign else 0x80) if exp == 0 else (0x8 if sign else 0x100))))
-  raise NotImplementedError(f"VOPC op {op} (base {base})")

extra/assembly/rdna3/asm.py

@@ -1,600 +0,0 @@
-# RDNA3 assembler and disassembler
-from __future__ import annotations
-import re
-from extra.assembly.rdna3.lib import Inst, RawImm, Reg, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory, FLOAT_ENC, SRC_FIELDS, unwrap
-
-# Decoding helpers
-SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
-SPECIAL_DEC = {**SPECIAL_GPRS, **{v: str(k) for k, v in FLOAT_ENC.items()}}
-SPECIAL_PAIRS = {106: "vcc", 126: "exec"}  # Special register pairs (for 64-bit ops)
-# GFX11 hwreg names (IDs 16-17 are TBA - not supported, IDs 18-19 are PERF_SNAPSHOT)
-HWREG_NAMES = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
-               6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
-               19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI',
-               22: 'HW_REG_XNACK_MASK', 23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
-HWREG_IDS = {v.lower(): k for k, v in HWREG_NAMES.items()}  # Reverse map for assembler
-MSG_NAMES = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
-             131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA'}
-_16BIT_TYPES = ('f16', 'i16', 'u16', 'b16')
-def _is_16bit(s: str) -> bool: return any(s.endswith(x) for x in _16BIT_TYPES)
-
-def decode_src(val: int) -> str:
-  if val <= 105: return f"s{val}"
-  if val in SPECIAL_DEC: return SPECIAL_DEC[val]
-  if 108 <= val <= 123: return f"ttmp{val - 108}"
-  if 128 <= val <= 192: return str(val - 128)
-  if 193 <= val <= 208: return str(-(val - 192))
-  if 256 <= val <= 511: return f"v{val - 256}"
-  return "lit" if val == 255 else f"?{val}"
-
-def _reg(prefix: str, base: int, cnt: int = 1) -> str: return f"{prefix}{base}" if cnt == 1 else f"{prefix}[{base}:{base+cnt-1}]"
-def _sreg(base: int, cnt: int = 1) -> str: return _reg("s", base, cnt)
-def _vreg(base: int, cnt: int = 1) -> str: return _reg("v", base, cnt)
-
-def _fmt_sdst(v: int, cnt: int = 1) -> str:
-  """Format SGPR destination with special register names."""
-  if v == 124: return "null"
-  if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
-  if cnt > 1 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
-  if cnt > 1: return _sreg(v, cnt)
-  return {126: "exec_lo", 127: "exec_hi", 106: "vcc_lo", 107: "vcc_hi", 125: "m0"}.get(v, f"s{v}")
-
-def _fmt_ssrc(v: int, cnt: int = 1) -> str:
-  """Format SGPR source with special register names and pairs."""
-  if cnt == 2:
-    if v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
-    if v <= 105: return _sreg(v, 2)
-    if 108 <= v <= 123: return _reg("ttmp", v - 108, 2)
-  return decode_src(v)
-
-def _fmt_src_n(v: int, cnt: int) -> str:
-  """Format source with given register count (1, 2, or 4)."""
-  if cnt == 1: return decode_src(v)
-  if v >= 256: return _vreg(v - 256, cnt)
-  if v <= 105: return _sreg(v, cnt)
-  if cnt == 2 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
-  if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
-  return decode_src(v)
-
-def _fmt_src64(v: int) -> str:
-  """Format 64-bit source (VGPR pair, SGPR pair, or special pair)."""
-  return _fmt_src_n(v, 2)
-
-def _parse_sop_sizes(op_name: str) -> tuple[int, ...]:
-  """Parse dst and src sizes from SOP instruction name. Returns (dst_cnt, src0_cnt) or (dst_cnt, src0_cnt, src1_cnt)."""
-  if op_name in ('s_bitset0_b64', 's_bitset1_b64'): return (2, 1)
-  if op_name in ('s_lshl_b64', 's_lshr_b64', 's_ashr_i64', 's_bfe_u64', 's_bfe_i64'): return (2, 2, 1)
-  if op_name in ('s_bfm_b64',): return (2, 1, 1)
-  # SOPC: s_bitcmp0_b64, s_bitcmp1_b64 - 64-bit src0, 32-bit src1 (bit index)
-  if op_name in ('s_bitcmp0_b64', 's_bitcmp1_b64'): return (1, 2, 1)
-  if m := re.search(r'_(b|i|u)(32|64)_(b|i|u)(32|64)$', op_name):
-    return (2 if m.group(2) == '64' else 1, 2 if m.group(4) == '64' else 1)
-  if m := re.search(r'_(b|i|u)(32|64)$', op_name):
-    sz = 2 if m.group(2) == '64' else 1
-    return (sz, sz)
-  return (1, 1)
-
-# Waitcnt helpers (RDNA3 format: bits 15:10=vmcnt, bits 9:4=lgkmcnt, bits 3:0=expcnt)
-def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
-  return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
-def decode_waitcnt(val: int) -> tuple[int, int, int]:
-  return (val >> 10) & 0x3f, val & 0xf, (val >> 4) & 0x3f  # vmcnt, expcnt, lgkmcnt
-
-# VOP3SD opcodes (shared encoding with VOP3 but different field layout)
-# Note: opcodes 0-255 are VOPC promoted to VOP3 - never treat as VOP3SD
-VOP3SD_OPCODES = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
-
-# Disassembler
-def disasm(inst: Inst) -> str:
-  op_val = unwrap(inst._values.get('op', 0))
-  cls_name = inst.__class__.__name__
-  # VOP3 and VOP3SD share encoding - check opcode to determine which
-  is_vop3sd = cls_name == 'VOP3' and op_val in VOP3SD_OPCODES
-  try:
-    from extra.assembly.rdna3 import autogen
-    if is_vop3sd:
-      op_name = autogen.VOP3SDOp(op_val).name.lower()
-    else:
-      op_name = getattr(autogen, f"{cls_name}Op")(op_val).name.lower() if hasattr(autogen, f"{cls_name}Op") else f"op_{op_val}"
-  except (ValueError, KeyError): op_name = f"op_{op_val}"
-  def fmt_src(v): return f"0x{inst._literal:x}" if v == 255 and getattr(inst, '_literal', None) else decode_src(v)
-
-  # VOP1
-  if cls_name == 'VOP1':
-    vdst, src0 = unwrap(inst._values['vdst']), unwrap(inst._values['src0'])
-    if op_name == 'v_nop': return 'v_nop'
-    if op_name == 'v_pipeflush': return 'v_pipeflush'
-    parts = op_name.split('_')
-    is_16bit_dst = any(p in _16BIT_TYPES for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in _16BIT_TYPES and 'cvt' not in op_name)
-    is_16bit_src = parts[-1] in _16BIT_TYPES and 'sat_pk' not in op_name
-    _F64_OPS = ('v_ceil_f64', 'v_floor_f64', 'v_fract_f64', 'v_frexp_mant_f64', 'v_rcp_f64', 'v_rndne_f64', 'v_rsq_f64', 'v_sqrt_f64', 'v_trunc_f64')
-    is_f64_dst = op_name in _F64_OPS or op_name in ('v_cvt_f64_f32', 'v_cvt_f64_i32', 'v_cvt_f64_u32')
-    is_f64_src = op_name in _F64_OPS or op_name in ('v_cvt_f32_f64', 'v_cvt_i32_f64', 'v_cvt_u32_f64', 'v_frexp_exp_i32_f64')
-    if op_name == 'v_readfirstlane_b32':
-      return f"v_readfirstlane_b32 {decode_src(vdst)}, v{src0 - 256 if src0 >= 256 else src0}"
-    dst_str = _vreg(vdst, 2) if is_f64_dst else f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}" if is_16bit_dst else f"v{vdst}"
-    src_str = _fmt_src64(src0) if is_f64_src else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit_src and src0 >= 256 else fmt_src(src0)
-    return f"{op_name}_e32 {dst_str}, {src_str}"
-
-  # VOP2
-  if cls_name == 'VOP2':
-    vdst, src0_raw, vsrc1 = unwrap(inst._values['vdst']), unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
-    suffix = "" if op_name == "v_dot2acc_f32_f16" else "_e32"
-    is_16bit_op = ('_f16' in op_name or '_i16' in op_name or '_u16' in op_name) and '_f32' not in op_name and '_i32' not in op_name and 'pk_' not in op_name
-    if is_16bit_op:
-      dst_str = f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}"
-      src0_str = f"v{(src0_raw - 256) & 0x7f}.{'h' if src0_raw >= 384 else 'l'}" if src0_raw >= 256 else fmt_src(src0_raw)
-      vsrc1_str = f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}"
-    else:
-      dst_str, src0_str, vsrc1_str = f"v{vdst}", fmt_src(src0_raw), f"v{vsrc1}"
-    return f"{op_name}{suffix} {dst_str}, {src0_str}, {vsrc1_str}" + (", vcc_lo" if op_name == "v_cndmask_b32" else "")
-
-  # VOPC
-  if cls_name == 'VOPC':
-    src0, vsrc1 = unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
-    is_64bit = any(x in op_name for x in ('f64', 'i64', 'u64'))
-    is_64bit_vsrc1 = is_64bit and 'class' not in op_name
-    is_16bit = any(x in op_name for x in ('_f16', '_i16', '_u16')) and 'f32' not in op_name
-    is_cmpx = op_name.startswith('v_cmpx')  # VOPCX writes to exec, no vcc destination
-    src0_str = _fmt_src64(src0) if is_64bit else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit and src0 >= 256 else fmt_src(src0)
-    vsrc1_str = _vreg(vsrc1, 2) if is_64bit_vsrc1 else f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}" if is_16bit else f"v{vsrc1}"
-    return f"{op_name}_e32 {src0_str}, {vsrc1_str}" if is_cmpx else f"{op_name}_e32 vcc_lo, {src0_str}, {vsrc1_str}"
-
-  # SOPP
-  if cls_name == 'SOPP':
-    simm16 = unwrap(inst._values.get('simm16', 0))
-    # No-operand instructions (simm16 is ignored)
-    no_imm_ops = ('s_endpgm', 's_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_ttracedata_imm',
-                  's_wait_idle', 's_endpgm_saved', 's_code_end', 's_endpgm_ordered_ps_done')
-    if op_name in no_imm_ops: return op_name
-    if op_name == 's_waitcnt':
-      vmcnt, expcnt, lgkmcnt = decode_waitcnt(simm16)
-      parts = []
-      if vmcnt != 0x3f: parts.append(f"vmcnt({vmcnt})")
-      if expcnt != 0x7: parts.append(f"expcnt({expcnt})")
-      if lgkmcnt != 0x3f: parts.append(f"lgkmcnt({lgkmcnt})")
-      return f"s_waitcnt {' '.join(parts)}" if parts else "s_waitcnt 0"
-    if op_name == 's_delay_alu':
-      dep_names = ['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']
-      skip_names = ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
-      id0, skip, id1 = simm16 & 0xf, (simm16 >> 4) & 0x7, (simm16 >> 7) & 0xf
-      def dep_name(v): return dep_names[v-1] if 0 < v <= len(dep_names) else str(v)
-      parts = [f"instid0({dep_name(id0)})"] if id0 else []
-      if skip: parts.append(f"instskip({skip_names[skip]})")
-      if id1: parts.append(f"instid1({dep_name(id1)})")
-      return f"s_delay_alu {' | '.join(p for p in parts if p)}" if parts else "s_delay_alu 0"
-    if op_name.startswith('s_cbranch') or op_name.startswith('s_branch'):
-      return f"{op_name} {simm16}"
-    # Most SOPP ops require immediate (s_nop, s_setkill, s_sethalt, s_sleep, s_setprio, s_sendmsg*, etc.)
-    return f"{op_name} 0x{simm16:x}"
-
-  # SMEM
-  if cls_name == 'SMEM':
-    if op_name in ('s_gl1_inv', 's_dcache_inv'): return op_name
-    sdata, sbase, soffset, offset = unwrap(inst._values['sdata']), unwrap(inst._values['sbase']), unwrap(inst._values['soffset']), unwrap(inst._values.get('offset', 0))
-    glc, dlc = unwrap(inst._values.get('glc', 0)), unwrap(inst._values.get('dlc', 0))
-    # Format offset: "soffset offset:X" if both, "0x{offset:x}" if only imm, or decode_src(soffset)
-    off_str = f"{decode_src(soffset)} offset:0x{offset:x}" if offset and soffset != 124 else f"0x{offset:x}" if offset else decode_src(soffset)
-    sbase_idx, sbase_cnt = sbase * 2, 4 if (8 <= op_val <= 12 or op_name == 's_atc_probe_buffer') else 2
-    sbase_str = _fmt_ssrc(sbase_idx, sbase_cnt) if sbase_cnt == 2 else _sreg(sbase_idx, sbase_cnt) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_cnt)
-    if op_name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{op_name} {sdata}, {sbase_str}, {off_str}"
-    width = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val, 1)
-    mods = [m for m in ["glc" if glc else "", "dlc" if dlc else ""] if m]
-    return f"{op_name} {_fmt_sdst(sdata, width)}, {sbase_str}, {off_str}" + (" " + " ".join(mods) if mods else "")
-
-  # FLAT
-  if cls_name == 'FLAT':
-    vdst, addr, data, saddr, offset, seg = [unwrap(inst._values.get(f, 0)) for f in ['vdst', 'addr', 'data', 'saddr', 'offset', 'seg']]
-    instr = f"{['flat', 'scratch', 'global'][seg] if seg < 3 else 'flat'}_{op_name.split('_', 1)[1] if '_' in op_name else op_name}"
-    width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'u8':1, 'i8':1, 'u16':1, 'i16':1}.get(op_name.split('_')[-1], 1)
-    addr_str = _vreg(addr, 2) if saddr == 0x7F else _vreg(addr)
-    saddr_str = "" if saddr == 0x7F else f", {_sreg(saddr, 2)}" if saddr < 106 else ", off" if saddr == 124 else f", {decode_src(saddr)}"
-    off_str = f" offset:{offset}" if offset else ""
-    vdata_str = _vreg(data if 'store' in op_name else vdst, width)
-    return f"{instr} {addr_str}, {vdata_str}{saddr_str}{off_str}" if 'store' in op_name else f"{instr} {vdata_str}, {addr_str}{saddr_str}{off_str}"
-
-  # VOP3: vector ops with modifiers (can be 1, 2, or 3 sources depending on opcode range)
-  if cls_name == 'VOP3':
-    # Handle VOP3SD opcodes (same encoding, different field layout)
-    if is_vop3sd:
-      vdst = unwrap(inst._values.get('vdst', 0))
-      # VOP3SD: sdst is at bits [14:8], but VOP3 decodes opsel at [14:11], abs at [10:8], clmp at [15]
-      # We need to reconstruct sdst from these fields
-      opsel_raw = unwrap(inst._values.get('opsel', 0))
-      abs_raw = unwrap(inst._values.get('abs', 0))
-      clmp_raw = unwrap(inst._values.get('clmp', 0))
-      sdst = (clmp_raw << 7) | (opsel_raw << 3) | abs_raw
-      src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
-      neg = unwrap(inst._values.get('neg', 0))
-      omod = unwrap(inst._values.get('omod', 0))
-      omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
-      is_f64 = 'f64' in op_name
-      # v_mad_i64_i32/v_mad_u64_u32: 64-bit dst and src2, 32-bit src0/src1
-      is_mad64 = 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
-      def fmt_sd_src(v, neg_bit, is_64bit=False):
-        s = _fmt_src64(v) if (is_64bit or is_f64) else fmt_src(v)
-        return f"-{s}" if neg_bit else s
-      src0_str, src1_str = fmt_sd_src(src0, neg & 1), fmt_sd_src(src1, neg & 2)
-      src2_str = fmt_sd_src(src2, neg & 4, is_mad64)
-      dst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}"
-      sdst_str = _fmt_sdst(sdst, 1)
-      # v_add_co_u32, v_sub_co_u32, v_subrev_co_u32, v_add_co_ci_u32, etc. only use 2 sources
-      if op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32', 'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'):
-        return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}"
-      # v_div_scale uses 3 sources
-      return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}, {src2_str}" + omod_str
-
-    vdst = unwrap(inst._values.get('vdst', 0))
-    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
-    neg, abs_, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('abs', 0)), unwrap(inst._values.get('clmp', 0))
-    opsel = unwrap(inst._values.get('opsel', 0))
-    # Check if 64-bit op (needs register pairs)
-    is_f64 = 'f64' in op_name or 'i64' in op_name or 'u64' in op_name or 'b64' in op_name
-    # v_cmp_class_* has 64-bit src0 but 32-bit src1 (class mask)
-    is_class = 'class' in op_name
-    # Shift ops: v_*rev_*64 have 32-bit shift amount (src0), 64-bit value (src1)
-    is_shift64 = 'rev' in op_name and '64' in op_name and op_name.startswith('v_')
-    # v_ldexp_f64: 64-bit src0 (mantissa), 32-bit src1 (exponent)
-    is_ldexp64 = op_name == 'v_ldexp_f64'
-    # v_trig_preop_f64: 64-bit dst/src0, 32-bit src1 (exponent/scale)
-    is_trig_preop = op_name == 'v_trig_preop_f64'
-    # v_readlane_b32: destination is SGPR (despite vdst field)
-    is_readlane = op_name == 'v_readlane_b32'
-    # SAD/QSAD/MQSAD instructions have mixed sizes
-    # v_qsad_pk_u16_u8, v_mqsad_pk_u16_u8: 64-bit dst/src0/src2, 32-bit src1
-    # v_mqsad_u32_u8: 128-bit (4 reg) dst/src2, 64-bit src0, 32-bit src1
-    is_sad64 = any(x in op_name for x in ('qsad_pk', 'mqsad_pk'))
-    is_mqsad_u32 = 'mqsad_u32' in op_name
-    # Detect 16-bit and 64-bit operand sizes for various instruction patterns
-    if 'cvt_pk' in op_name:
-      is_f16_dst, is_f16_src, is_f16_src2 = False, op_name.endswith('16'), False
-    elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', op_name):
-      dst_type, src_type = m.group(1), m.group(2)
-      is_f16_dst, is_f16_src, is_f16_src2 = _is_16bit(dst_type), _is_16bit(src_type), _is_16bit(src_type)
-      is_f64_dst, is_f64_src, is_f64 = '64' in dst_type, '64' in src_type, False
-    elif re.match(r'v_mad_[iu]32_[iu]16', op_name):
-      is_f16_dst, is_f16_src, is_f16_src2 = False, True, False  # 32-bit dst, 16-bit src0/src1, 32-bit src2
-    elif 'pack_b32' in op_name:
-      is_f16_dst, is_f16_src, is_f16_src2 = False, True, True  # 32-bit dst, 16-bit sources
-    else:
-      is_16bit_op = any(x in op_name for x in _16BIT_TYPES) and not any(x in op_name for x in ('dot2', 'pk_', 'sad', 'msad', 'qsad', 'mqsad'))
-      is_f16_dst = is_f16_src = is_f16_src2 = is_16bit_op
-    def fmt_vop3_src(v, neg_bit, abs_bit, hi_bit=False, reg_cnt=1, is_16=False):
-      s = _fmt_src_n(v, reg_cnt) if reg_cnt > 1 else f"v{v - 256}.h" if is_16 and v >= 256 and hi_bit else f"v{v - 256}.l" if is_16 and v >= 256 else fmt_src(v)
-      if abs_bit: s = f"|{s}|"
-      return f"-{s}" if neg_bit else s
-    # Determine register count for each source (check for cvt-specific 64-bit flags first)
-    is_src0_64 = locals().get('is_f64_src', is_f64 and not is_shift64) or is_sad64 or is_mqsad_u32
-    is_src1_64 = is_f64 and not is_class and not is_ldexp64 and not is_trig_preop
-    src0_cnt = 2 if is_src0_64 else 1
-    src1_cnt = 2 if is_src1_64 else 1
-    src2_cnt = 4 if is_mqsad_u32 else 2 if (is_f64 or is_sad64) else 1
-    src0_str = fmt_vop3_src(src0, neg & 1, abs_ & 1, opsel & 1, src0_cnt, is_f16_src)
-    src1_str = fmt_vop3_src(src1, neg & 2, abs_ & 2, opsel & 2, src1_cnt, is_f16_src)
-    src2_str = fmt_vop3_src(src2, neg & 4, abs_ & 4, opsel & 4, src2_cnt, is_f16_src2)
-    # Format destination - for 16-bit ops, use .h/.l suffix; readlane uses SGPR dest
-    is_dst_64 = locals().get('is_f64_dst', is_f64) or is_sad64
-    dst_cnt = 4 if is_mqsad_u32 else 2 if is_dst_64 else 1
-    if is_readlane:
-      dst_str = _fmt_sdst(vdst, 1)
-    elif dst_cnt > 1:
-      dst_str = _vreg(vdst, dst_cnt)
-    elif is_f16_dst:
-      dst_str = f"v{vdst}.h" if (opsel & 8) else f"v{vdst}.l"
-    else:
-      dst_str = f"v{vdst}"
-    clamp_str = " clamp" if clmp else ""
-    omod = unwrap(inst._values.get('omod', 0))
-    omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
-    # op_sel for non-VGPR sources (when opsel bits are set but source is not a VGPR)
-    # For 16-bit ops with VGPR sources, opsel is encoded in .h/.l suffix
-    # For non-VGPR sources or non-16-bit ops, we need explicit op_sel
-    has_nonvgpr_opsel = (src0 < 256 and (opsel & 1)) or (src1 < 256 and (opsel & 2)) or (src2 < 256 and (opsel & 4))
-    need_opsel = has_nonvgpr_opsel or (opsel and not is_f16_src)
-    # Helper to format opsel string based on source count
-    def fmt_opsel(num_src):
-      if not need_opsel: return ""
-      # When dst is .h (for 16-bit ops) and non-VGPR sources have opsel, use all 1s
-      if is_f16_dst and (opsel & 8):  # dst is .h
-        return f" op_sel:[1,1,1{',1' if num_src == 3 else ''}]"
-      # Otherwise output actual opsel values
-      if num_src == 3:
-        return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{(opsel >> 3) & 1}]"
-      return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1}]"
-    # Determine number of sources based on opcode range:
-    # 0-255: VOPC promoted (comparison, 2 src, sdst)
-    # 256-383: VOP2 promoted (2 src)
-    # 384-511: VOP1 promoted (1 src)
-    # 512+: Native VOP3 (2 or 3 src depending on instruction)
-    if op_val < 256:  # VOPC promoted
-      # VOPCX (v_cmpx_*) writes to exec, no explicit destination
-      if op_name.startswith('v_cmpx'):
-        return f"{op_name}_e64 {src0_str}, {src1_str}"
-      return f"{op_name}_e64 {_fmt_sdst(vdst, 1)}, {src0_str}, {src1_str}"
-    elif op_val < 384:  # VOP2 promoted
-      # v_cndmask_b32 in VOP3 format has 3 sources (src2 is mask selector)
-      if 'cndmask' in op_name:
-        return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
-      return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
-    elif op_val < 512:  # VOP1 promoted
-      if op_name in ('v_nop', 'v_pipeflush'): return f"{op_name}_e64"
-      return f"{op_name}_e64 {dst_str}, {src0_str}" + fmt_opsel(1) + clamp_str + omod_str
-    else:  # Native VOP3 - determine 2 vs 3 sources based on instruction name
-      # 3-source ops: fma, mad, min3, max3, med3, div_fixup, div_fmas, sad, msad, qsad, mqsad, lerp, alignbit/byte, cubeid/sc/tc/ma, bfe, bfi, perm_b32, permlane, cndmask
-      # Note: v_writelane_b32 is 2-src (src0, src1 with vdst as 3rd operand - read-modify-write)
-      is_3src = any(x in op_name for x in ('fma', 'mad', 'min3', 'max3', 'med3', 'div_fix', 'div_fmas', 'sad', 'lerp', 'align', 'cube',
-                                            'bfe', 'bfi', 'perm_b32', 'permlane', 'cndmask', 'xor3', 'or3', 'add3', 'lshl_or', 'and_or', 'lshl_add',
-                                            'add_lshl', 'xad', 'maxmin', 'minmax', 'dot2', 'cvt_pk_u8', 'mullit'))
-      if is_3src:
-        return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
-      return f"{op_name} {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
-
-  # VOP3SD: 3-source with scalar destination (v_div_scale_*, v_add_co_u32, v_mad_*64_*32, etc.)
-  if cls_name == 'VOP3SD':
-    vdst, sdst = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('sdst', 0))
-    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
-    neg, omod, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('omod', 0)), unwrap(inst._values.get('clmp', 0))
-    is_f64, is_mad64 = 'f64' in op_name, 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
-    def fmt_neg(v, neg_bit, is_64=False): return f"-{_fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)}" if neg_bit else _fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)
-    srcs = [fmt_neg(src0, neg & 1), fmt_neg(src1, neg & 2), fmt_neg(src2, neg & 4, is_mad64)]
-    dst_str, sdst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}", _fmt_sdst(sdst, 1)
-    clamp_str, omod_str = " clamp" if clmp else "", {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
-    is_2src = op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32')
-    suffix = "_e64" if op_name.startswith('v_') and 'co_' in op_name else ""
-    return f"{op_name}{suffix} {dst_str}, {sdst_str}, {', '.join(srcs[:2] if is_2src else srcs)}" + clamp_str + omod_str
-
-  # VOPD: dual-issue instructions
-  if cls_name == 'VOPD':
-    from extra.assembly.rdna3 import autogen
-    opx, opy, vdstx, vdsty_enc = [unwrap(inst._values.get(f, 0)) for f in ('opx', 'opy', 'vdstx', 'vdsty')]
-    srcx0, vsrcx1, srcy0, vsrcy1 = [unwrap(inst._values.get(f, 0)) for f in ('srcx0', 'vsrcx1', 'srcy0', 'vsrcy1')]
-    vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1)  # Decode vdsty
-    def fmt_vopd(op, vdst, src0, vsrc1):
-      try: name = autogen.VOPDOp(op).name.lower()
-      except (ValueError, KeyError): name = f"op_{op}"
-      return f"{name} v{vdst}, {fmt_src(src0)}" if 'mov' in name else f"{name} v{vdst}, {fmt_src(src0)}, v{vsrc1}"
-    return f"{fmt_vopd(opx, vdstx, srcx0, vsrcx1)} :: {fmt_vopd(opy, vdsty, srcy0, vsrcy1)}"
-
-  # VOP3P: packed vector ops
-  if cls_name == 'VOP3P':
-    vdst, clmp = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('clmp', 0))
-    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
-    neg, neg_hi = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('neg_hi', 0))
-    opsel, opsel_hi, opsel_hi2 = unwrap(inst._values.get('opsel', 0)), unwrap(inst._values.get('opsel_hi', 0)), unwrap(inst._values.get('opsel_hi2', 0))
-    is_wmma, is_3src = 'wmma' in op_name, any(x in op_name for x in ('fma', 'mad', 'dot', 'wmma'))
-    def fmt_bits(name, val, n): return f"{name}:[{','.join(str((val >> i) & 1) for i in range(n))}]"
-    # WMMA: f16/bf16 use 8-reg sources, iu8 uses 4-reg, iu4 uses 2-reg; all have 8-reg dst
-    if is_wmma:
-      src_cnt = 2 if 'iu4' in op_name else 4 if 'iu8' in op_name else 8
-      src0_str, src1_str, src2_str = _fmt_src_n(src0, src_cnt), _fmt_src_n(src1, src_cnt), _fmt_src_n(src2, 8)
-      dst_str = _vreg(vdst, 8)
-    else:
-      src0_str, src1_str, src2_str = _fmt_src_n(src0, 1), _fmt_src_n(src1, 1), _fmt_src_n(src2, 1)
-      dst_str = f"v{vdst}"
-    n = 3 if is_3src else 2
-    full_opsel_hi = opsel_hi | (opsel_hi2 << 2)
-    mods = [fmt_bits("op_sel", opsel, n)] if opsel else []
-    if full_opsel_hi != (0b111 if is_3src else 0b11): mods.append(fmt_bits("op_sel_hi", full_opsel_hi, n))
-    if neg: mods.append(fmt_bits("neg_lo", neg, n))
-    if neg_hi: mods.append(fmt_bits("neg_hi", neg_hi, n))
-    if clmp: mods.append("clamp")
-    mod_str = " " + " ".join(mods) if mods else ""
-    return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}{mod_str}" if is_3src else f"{op_name} {dst_str}, {src0_str}, {src1_str}{mod_str}"
-
-  # VINTERP: interpolation instructions
-  if cls_name == 'VINTERP':
-    vdst = unwrap(inst._values.get('vdst', 0))
-    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
-    neg, waitexp, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('waitexp', 0)), unwrap(inst._values.get('clmp', 0))
-    def fmt_neg_vi(v, neg_bit): return f"-{v}" if neg_bit else v
-    srcs = [fmt_neg_vi(f"v{s - 256}" if s >= 256 else fmt_src(s), neg & (1 << i)) for i, s in enumerate([src0, src1, src2])]
-    mods = [m for m in [f"wait_exp:{waitexp}" if waitexp else "", "clamp" if clmp else ""] if m]
-    return f"{op_name} v{vdst}, {', '.join(srcs)}" + (" " + " ".join(mods) if mods else "")
-
-  # MUBUF/MTBUF helpers
-  def _buf_vaddr(vaddr, offen, idxen): return _vreg(vaddr, 2) if offen and idxen else f"v{vaddr}" if offen or idxen else "off"
-  def _buf_srsrc(srsrc): srsrc_base = srsrc * 4; return _reg("ttmp", srsrc_base - 108, 4) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 4)
-
-  # MUBUF: buffer load/store
-  if cls_name == 'MUBUF':
-    vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
-    offset, offen, idxen = unwrap(inst._values.get('offset', 0)), unwrap(inst._values.get('offen', 0)), unwrap(inst._values.get('idxen', 0))
-    glc, dlc, slc, tfe = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc', 'tfe')]
-    if op_name in ('buffer_gl0_inv', 'buffer_gl1_inv'): return op_name
-    # Determine data width from op name
-    if 'd16' in op_name: width = 2 if any(x in op_name for x in ('xyz', 'xyzw')) else 1
-    elif 'atomic' in op_name:
-      base_width = 2 if any(x in op_name for x in ('b64', 'u64', 'i64')) else 1
-      width = base_width * 2 if 'cmpswap' in op_name else base_width
-    else: width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'b16':1, 'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
-    if tfe: width += 1
-    mods = [m for m in ["offen" if offen else "", "idxen" if idxen else "", f"offset:{offset}" if offset else "",
-                        "glc" if glc else "", "dlc" if dlc else "", "slc" if slc else "", "tfe" if tfe else ""] if m]
-    return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)}" + (" " + " ".join(mods) if mods else "")
-
-  # MTBUF: typed buffer load/store
-  if cls_name == 'MTBUF':
-    vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
-    offset, tbuf_fmt, offen, idxen = [unwrap(inst._values.get(f, 0)) for f in ('offset', 'format', 'offen', 'idxen')]
-    glc, dlc, slc = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc')]
-    mods = [f"format:{tbuf_fmt}"] + [m for m in ["idxen" if idxen else "", "offen" if offen else "", f"offset:{offset}" if offset else "",
-                                                  "glc" if glc else "", "dlc" if dlc else "", "slc" if slc else ""] if m]
-    width = 2 if 'd16' in op_name and any(x in op_name for x in ('xyz', 'xyzw')) else 1 if 'd16' in op_name else {'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
-    return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)} {' '.join(mods)}"
-
-  # SOP1/SOP2/SOPC/SOPK
-  if cls_name in ('SOP1', 'SOP2', 'SOPC', 'SOPK'):
-    sizes = _parse_sop_sizes(op_name)
-    dst_cnt, src0_cnt = sizes[0], sizes[1]
-    src1_cnt = sizes[2] if len(sizes) > 2 else src0_cnt
-    if cls_name == 'SOP1':
-      sdst, ssrc0 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('ssrc0', 0))
-      if op_name == 's_getpc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}"
-      if op_name in ('s_setpc_b64', 's_rfe_b64'): return f"{op_name} {_fmt_ssrc(ssrc0, 2)}"
-      if op_name == 's_swappc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}, {_fmt_ssrc(ssrc0, 2)}"
-      if op_name in ('s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'):
-        return f"{op_name} {_fmt_sdst(sdst, 2 if 'b64' in op_name else 1)}, sendmsg({MSG_NAMES.get(ssrc0, str(ssrc0))})"
-      ssrc0_str = fmt_src(ssrc0) if src0_cnt == 1 else _fmt_ssrc(ssrc0, src0_cnt)
-      return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {ssrc0_str}"
-    if cls_name == 'SOP2':
-      sdst, ssrc0, ssrc1 = [unwrap(inst._values.get(f, 0)) for f in ('sdst', 'ssrc0', 'ssrc1')]
-      return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {_fmt_ssrc(ssrc0, src0_cnt)}, {_fmt_ssrc(ssrc1, src1_cnt)}"
-    if cls_name == 'SOPC':
-      return f"{op_name} {_fmt_ssrc(unwrap(inst._values.get('ssrc0', 0)), src0_cnt)}, {_fmt_ssrc(unwrap(inst._values.get('ssrc1', 0)), src1_cnt)}"
-    if cls_name == 'SOPK':
-      sdst, simm16 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('simm16', 0))
-      if op_name == 's_version': return f"{op_name} 0x{simm16:x}"
-      if op_name in ('s_setreg_b32', 's_getreg_b32'):
-        hwreg_id, hwreg_offset, hwreg_size = simm16 & 0x3f, (simm16 >> 6) & 0x1f, ((simm16 >> 11) & 0x1f) + 1
-        hwreg_str = f"0x{simm16:x}" if hwreg_id in (16, 17) else f"hwreg({HWREG_NAMES.get(hwreg_id, str(hwreg_id))}, {hwreg_offset}, {hwreg_size})"
-        return f"{op_name} {hwreg_str}, {_fmt_sdst(sdst, 1)}" if op_name == 's_setreg_b32' else f"{op_name} {_fmt_sdst(sdst, 1)}, {hwreg_str}"
-      return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, 0x{simm16:x}"
-
-  # Generic fallback
-  def fmt_field(n, v):
-    v = unwrap(v)
-    if n in SRC_FIELDS: return fmt_src(v) if v != 255 else "0xff"
-    if n in ('sdst', 'vdst'): return f"{'s' if n == 'sdst' else 'v'}{v}"
-    return f"v{v}" if n == 'vsrc1' else f"0x{v:x}" if n == 'simm16' else str(v)
-  ops = [fmt_field(n, inst._values.get(n, 0)) for n in inst._fields if n not in ('encoding', 'op')]
-  return f"{op_name} {', '.join(ops)}" if ops else op_name
-
-# Assembler
-SPECIAL_REGS = {'vcc_lo': RawImm(106), 'vcc_hi': RawImm(107), 'null': RawImm(124), 'off': RawImm(124), 'm0': RawImm(125), 'exec_lo': RawImm(126), 'exec_hi': RawImm(127), 'scc': RawImm(253)}
-FLOAT_CONSTS = {'0.5': 0.5, '-0.5': -0.5, '1.0': 1.0, '-1.0': -1.0, '2.0': 2.0, '-2.0': -2.0, '4.0': 4.0, '-4.0': -4.0}
-REG_MAP: dict[str, _RegFactory] = {'s': s, 'v': v, 't': ttmp, 'ttmp': ttmp}
-
-def parse_operand(op: str) -> tuple:
-  op = op.strip().lower()
-  neg = op.startswith('-') and not op[1:2].isdigit(); op = op[1:] if neg else op
-  abs_ = op.startswith('|') and op.endswith('|') or op.startswith('abs(') and op.endswith(')')
-  op = op[1:-1] if op.startswith('|') else op[4:-1] if op.startswith('abs(') else op
-  hi_half = op.endswith('.h')
-  op = re.sub(r'\.[lh]$', '', op)
-  if op in FLOAT_CONSTS: return (FLOAT_CONSTS[op], neg, abs_, hi_half)
-  if re.match(r'^-?\d+$', op): return (int(op), neg, abs_, hi_half)
-  if m := re.match(r'^-?0x([0-9a-f]+)$', op):
-    v = -int(m.group(1), 16) if op.startswith('-') else int(m.group(1), 16)
-    return (v, neg, abs_, hi_half)
-  if op in SPECIAL_REGS: return (SPECIAL_REGS[op], neg, abs_, hi_half)
-  if op == 'lit': return (RawImm(255), neg, abs_, hi_half)  # literal marker (actual value comes from literal word)
-  if m := re.match(r'^([svt](?:tmp)?)\[(\d+):(\d+)\]$', op): return (REG_MAP[m.group(1)][int(m.group(2)):int(m.group(3))], neg, abs_, hi_half)
-  if m := re.match(r'^([svt](?:tmp)?)(\d+)$', op):
-    reg = REG_MAP[m.group(1)][int(m.group(2))]
-    reg.hi = hi_half
-    return (reg, neg, abs_, hi_half)
-  # hwreg(name, offset, size) or hwreg(name) -> simm16 encoding
-  if m := re.match(r'^hwreg\((\w+)(?:,\s*(\d+),\s*(\d+))?\)$', op):
-    name_str = m.group(1).lower()
-    hwreg_id = HWREG_IDS.get(name_str, int(name_str) if name_str.isdigit() else None)
-    if hwreg_id is None: raise ValueError(f"unknown hwreg name: {name_str}")
-    offset, size = int(m.group(2)) if m.group(2) else 0, int(m.group(3)) if m.group(3) else 32
-    return (((size - 1) << 11) | (offset << 6) | hwreg_id, neg, abs_, hi_half)
-  raise ValueError(f"cannot parse operand: {op}")
-
-SMEM_OPS = {'s_load_b32', 's_load_b64', 's_load_b128', 's_load_b256', 's_load_b512',
-            's_buffer_load_b32', 's_buffer_load_b64', 's_buffer_load_b128', 's_buffer_load_b256', 's_buffer_load_b512'}
-SOP1_SRC_ONLY = {'s_setpc_b64', 's_rfe_b64'}
-SOP1_MSG_IMM = {'s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'}
-SOPK_IMM_ONLY = {'s_version'}
-SOPK_IMM_FIRST = {'s_setreg_b32'}
-SOPK_UNSUPPORTED = {'s_setreg_imm32_b32'}
-
-def asm(text: str) -> Inst:
-  from extra.assembly.rdna3 import autogen
-  text = text.strip()
-  clamp = 'clamp' in text.lower()
-  if clamp: text = re.sub(r'\s+clamp\s*$', '', text, flags=re.I)
-  modifiers = {}
-  if m := re.search(r'\s+wait_exp:(\d+)', text, re.I): modifiers['waitexp'] = int(m.group(1)); text = text[:m.start()] + text[m.end():]
-  parts = text.replace(',', ' ').split()
-  if not parts: raise ValueError("empty instruction")
-  mnemonic, op_str = parts[0].lower(), text[len(parts[0]):].strip()
-  # Handle s_waitcnt specially before operand parsing
-  if mnemonic == 's_waitcnt':
-    vmcnt, expcnt, lgkmcnt = 0x3f, 0x7, 0x3f
-    for part in op_str.replace(',', ' ').split():
-      if m := re.match(r'vmcnt\((\d+)\)', part): vmcnt = int(m.group(1))
-      elif m := re.match(r'expcnt\((\d+)\)', part): expcnt = int(m.group(1))
-      elif m := re.match(r'lgkmcnt\((\d+)\)', part): lgkmcnt = int(m.group(1))
-      elif re.match(r'^0x[0-9a-f]+$|^\d+$', part): return autogen.s_waitcnt(simm16=int(part, 0))
-    return autogen.s_waitcnt(simm16=waitcnt(vmcnt, expcnt, lgkmcnt))
-  # Handle VOPD dual-issue instructions: opx dst, src :: opy dst, src
-  if '::' in text:
-    x_part, y_part = text.split('::')
-    x_parts, y_parts = x_part.strip().replace(',', ' ').split(), y_part.strip().replace(',', ' ').split()
-    opx_name, opy_name = x_parts[0].upper(), y_parts[0].upper()
-    opx, opy = autogen.VOPDOp[opx_name], autogen.VOPDOp[opy_name]
-    x_ops, y_ops = [parse_operand(p)[0] for p in x_parts[1:]], [parse_operand(p)[0] for p in y_parts[1:]]
-    vdstx, srcx0 = x_ops[0], x_ops[1] if len(x_ops) > 1 else 0
-    vsrcx1 = x_ops[2] if len(x_ops) > 2 else VGPR(0)
-    vdsty, srcy0 = y_ops[0], y_ops[1] if len(y_ops) > 1 else 0
-    vsrcy1 = y_ops[2] if len(y_ops) > 2 else VGPR(0)
-    # Handle fmaak/fmamk literals (4th operand on x or y side)
-    lit = None
-    if 'fmaak' in opx_name.lower() and len(x_ops) > 3: lit = unwrap(x_ops[3])
-    elif 'fmamk' in opx_name.lower() and len(x_ops) > 3: lit, vsrcx1 = unwrap(x_ops[2]), x_ops[3]
-    elif 'fmaak' in opy_name.lower() and len(y_ops) > 3: lit = unwrap(y_ops[3])
-    elif 'fmamk' in opy_name.lower() and len(y_ops) > 3: lit, vsrcy1 = unwrap(y_ops[2]), y_ops[3]
-    return autogen.VOPD(opx, opy, vdstx=vdstx, vdsty=vdsty, srcx0=srcx0, vsrcx1=vsrcx1, srcy0=srcy0, vsrcy1=vsrcy1, literal=lit)
-  operands, current, depth, in_pipe = [], "", 0, False
-  for ch in op_str:
-    if ch in '[(': depth += 1
-    elif ch in '])': depth -= 1
-    elif ch == '|': in_pipe = not in_pipe
-    if ch == ',' and depth == 0 and not in_pipe: operands.append(current.strip()); current = ""
-    else: current += ch
-  if current.strip(): operands.append(current.strip())
-  parsed = [parse_operand(op) for op in operands]
-  values = [p[0] for p in parsed]
-  neg_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[1])
-  abs_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[2])
-  opsel_bits = (8 if len(parsed) > 0 and parsed[0][3] else 0) | sum((1 << i) for i, p in enumerate(parsed[1:4]) if p[3])
-  lit = None
-  if mnemonic in ('v_fmaak_f32', 'v_fmaak_f16') and len(values) == 4: lit, values = unwrap(values[3]), values[:3]
-  elif mnemonic in ('v_fmamk_f32', 'v_fmamk_f16') and len(values) == 4: lit, values = unwrap(values[2]), [values[0], values[1], values[3]]
-  vcc_ops = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32', 'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'}
-  if mnemonic.replace('_e32', '') in vcc_ops and len(values) >= 5: values = [values[0], values[2], values[3]]
-  if mnemonic.startswith('v_cmp') and len(values) >= 3 and operands[0].strip().lower() in ('vcc_lo', 'vcc_hi', 'vcc'):
-    values = values[1:]
-  # CMPX instructions with _e64 suffix: prepend implicit EXEC_LO destination (vdst=126)
-  if 'cmpx' in mnemonic and mnemonic.endswith('_e64') and len(values) == 2:
-    values = [VGPR(126, 1)] + values
-    # Recalculate modifiers: parsed[0]=src0, parsed[1]=src1 (no vdst in user input)
-    neg_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[1])
-    abs_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[2])
-    opsel_bits = sum((1 << i) for i, p in enumerate(parsed[:2]) if p[3])
-  vop3sd_ops = {'v_div_scale_f32', 'v_div_scale_f64'}
-  if mnemonic in vop3sd_ops and len(parsed) >= 5:
-    neg_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[1])
-    abs_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[2])
-  if mnemonic in SOPK_UNSUPPORTED: raise ValueError(f"unsupported instruction: {mnemonic}")
-  elif mnemonic in SOP1_SRC_ONLY:
-    return getattr(autogen, mnemonic)(ssrc0=values[0])
-  elif mnemonic in SOP1_MSG_IMM:
-    return getattr(autogen, mnemonic)(sdst=values[0], ssrc0=RawImm(unwrap(values[1])))
-  elif mnemonic in SOPK_IMM_ONLY:
-    return getattr(autogen, mnemonic)(simm16=values[0])
-  elif mnemonic in SOPK_IMM_FIRST:
-    return getattr(autogen, mnemonic)(simm16=values[0], sdst=values[1])
-  elif mnemonic in SMEM_OPS and len(operands) >= 3 and re.match(r'^-?[0-9]|^-?0x', operands[2].strip().lower()):
-    return getattr(autogen, mnemonic)(sdata=values[0], sbase=values[1], offset=values[2], soffset=RawImm(124))
-  elif mnemonic.startswith('buffer_') and len(operands) >= 2 and operands[1].strip().lower() == 'off':
-    return getattr(autogen, mnemonic)(vdata=values[0], vaddr=0, srsrc=values[2], soffset=RawImm(unwrap(values[3])) if len(values) > 3 else RawImm(0))
-  elif (mnemonic.startswith('flat_load') or mnemonic.startswith('global_load') or mnemonic.startswith('scratch_load')) and len(values) >= 3:
-    offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
-    return getattr(autogen, mnemonic)(vdst=values[0], addr=values[1], saddr=values[2], offset=offset)
-  elif (mnemonic.startswith('flat_store') or mnemonic.startswith('global_store') or mnemonic.startswith('scratch_store')) and len(values) >= 3:
-    offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
-    return getattr(autogen, mnemonic)(addr=values[0], data=values[1], saddr=values[2], offset=offset)
-  for suffix in (['_e32', ''] if not (neg_bits or abs_bits or clamp) else ['', '_e32']):
-    if hasattr(autogen, name := mnemonic.replace('.', '_') + suffix):
-      use_opsel = 'opsel' in getattr(autogen, name).func._fields
-      vals = [type(v)(v.idx, v.count, False) if isinstance(v, Reg) and v.hi and use_opsel else v for v in values]
-      inst = getattr(autogen, name)(*vals, literal=lit, **modifiers)
-      if neg_bits and 'neg' in inst._fields: inst._values['neg'] = neg_bits
-      if opsel_bits and use_opsel: inst._values['opsel'] = opsel_bits
-      if abs_bits and 'abs' in inst._fields: inst._values['abs'] = abs_bits
-      if clamp and 'clmp' in inst._fields: inst._values['clmp'] = 1
-      return inst
-  raise ValueError(f"unknown instruction: {mnemonic}")

extra/assembly/rdna3/emu.py

@@ -1,505 +0,0 @@
-# RDNA3 emulator - pure Python implementation for testing
-from __future__ import annotations
-import ctypes, struct, math
-from typing import Callable
-from extra.assembly.rdna3.lib import Inst, Inst32, Inst64, RawImm
-
-Program = dict[int, Inst]  # pc (word offset) -> instruction
-from extra.assembly.rdna3.autogen import (
-  SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD, SrcEnum,
-  SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, VOPDOp
-)
-from extra.assembly.rdna3.alu import (
-  f32, i32, f16, i16, sext, vopc, SALU, VALU,
-  SOP1_BASE, SOP2_BASE, SOPC_BASE, SOPK_BASE, VOP1_BASE, VOP2_BASE
-)
-
-WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
-VCC_LO, VCC_HI, NULL, M0, EXEC_LO, EXEC_HI, SCC = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.M0, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.SCC
-# Pre-computed inline constant table for src operands 128-254 (index = src - 128)
-_INLINE_CONSTS = [0] * 127
-for _i in range(65): _INLINE_CONSTS[_i] = _i  # 128-192 -> 0-64
-for _i in range(1, 17): _INLINE_CONSTS[64 + _i] = ((-_i) & 0xffffffff)  # 193-208 -> -1 to -16
-for _k, _v in {SrcEnum.POS_HALF: 0x3f000000, SrcEnum.NEG_HALF: 0xbf000000, SrcEnum.POS_ONE: 0x3f800000, SrcEnum.NEG_ONE: 0xbf800000,
-               SrcEnum.POS_TWO: 0x40000000, SrcEnum.NEG_TWO: 0xc0000000, SrcEnum.POS_FOUR: 0x40800000, SrcEnum.NEG_FOUR: 0xc0800000,
-               SrcEnum.INV_2PI: 0x3e22f983}.items(): _INLINE_CONSTS[_k - 128] = _v
-
-_valid_mem_ranges: list[tuple[int, int]] = []
-def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: global _valid_mem_ranges; _valid_mem_ranges = list(ranges)
-def _mem_valid(addr: int, size: int) -> bool:
-  for s, z in _valid_mem_ranges:
-    if s <= addr and addr + size <= s + z: return True
-  return not _valid_mem_ranges
-def _ctypes_at(addr: int, size: int): return (ctypes.c_uint8 if size == 1 else ctypes.c_uint16 if size == 2 else ctypes.c_uint32).from_address(addr)
-def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
-def mem_write(addr: int, size: int, val: int) -> None:
-  if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
-
-# Memory op tables - (cnt, sz, sign) for loads, (cnt, sz) for stores
-def _mem_ops(ops, suffix_map):
-  return {getattr(e, f"{p}_{s}"): v for e in ops for s, v in suffix_map.items() for p in [e.__name__.replace("Op", "")]}
-_LOAD_MAP = {'LOAD_B32': (1,4,0), 'LOAD_B64': (2,4,0), 'LOAD_B96': (3,4,0), 'LOAD_B128': (4,4,0), 'LOAD_U8': (1,1,0), 'LOAD_I8': (1,1,1), 'LOAD_U16': (1,2,0), 'LOAD_I16': (1,2,1)}
-_STORE_MAP = {'STORE_B32': (1,4), 'STORE_B64': (2,4), 'STORE_B96': (3,4), 'STORE_B128': (4,4), 'STORE_B8': (1,1), 'STORE_B16': (1,2)}
-FLAT_LOAD = _mem_ops([GLOBALOp, FLATOp], _LOAD_MAP)
-FLAT_STORE = _mem_ops([GLOBALOp, FLATOp], _STORE_MAP)
-DS_LOAD: dict[int, tuple[int,int,int]] = {DSOp.DS_LOAD_B32: (1,4,0), DSOp.DS_LOAD_B64: (2,4,0), DSOp.DS_LOAD_B128: (4,4,0), DSOp.DS_LOAD_U8: (1,1,0), DSOp.DS_LOAD_I8: (1,1,1), DSOp.DS_LOAD_U16: (1,2,0), DSOp.DS_LOAD_I16: (1,2,1)}
-DS_STORE: dict[int, tuple[int,int]] = {DSOp.DS_STORE_B32: (1,4), DSOp.DS_STORE_B64: (2,4), DSOp.DS_STORE_B128: (4,4), DSOp.DS_STORE_B8: (1,1), DSOp.DS_STORE_B16: (1,2)}
-FLAT_D16_LO = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_U8', (1, 0)), ('LOAD_D16_I8', (1, 1)), ('LOAD_D16_B16', (2, 0))]}
-FLAT_D16_HI = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_HI_U8', (1, 0)), ('LOAD_D16_HI_I8', (1, 1)), ('LOAD_D16_HI_B16', (2, 0))]}
-FLAT_D16_STORE = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('STORE_D16_HI_B8', 1), ('STORE_D16_HI_B16', 2)]}
-SMEM_LOAD: dict[int, int] = {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}
-SOPK_WAIT = {SOPKOp.S_WAITCNT_VSCNT, SOPKOp.S_WAITCNT_VMCNT, SOPKOp.S_WAITCNT_EXPCNT, SOPKOp.S_WAITCNT_LGKMCNT}
-
-class WaveState:
-  __slots__ = ('sgpr', 'vgpr', 'scc', 'pc', 'literal', '_pend_sgpr')
-  def __init__(self):
-    self.sgpr, self.vgpr = [0] * SGPR_COUNT, [[0] * VGPR_COUNT for _ in range(WAVE_SIZE)]
-    self.sgpr[EXEC_LO] = 0xffffffff  # wave32: all lanes active
-    self.scc = self.pc = self.literal = 0
-    self._pend_sgpr = {}
-
-  @property
-  def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
-  @vcc.setter
-  def vcc(self, v: int) -> None: self.sgpr[VCC_LO] = v & 0xffffffff; self.sgpr[VCC_HI] = (v >> 32) & 0xffffffff
-  @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) -> None: self.sgpr[EXEC_LO] = v & 0xffffffff; self.sgpr[EXEC_HI] = (v >> 32) & 0xffffffff
-
-  def rsgpr(self, i: int) -> int:
-    if i == NULL: return 0
-    if i == SCC: return self.scc
-    return self.sgpr[i] if i < SGPR_COUNT else 0
-  def wsgpr(self, i: int, v: int) -> None:
-    if i < SGPR_COUNT and i != NULL: self.sgpr[i] = v & 0xffffffff
-  def rsgpr64(self, i: int) -> int: return self.rsgpr(i) | (self.rsgpr(i+1) << 32)
-  def wsgpr64(self, i: int, v: int) -> None: self.wsgpr(i, v & 0xffffffff); self.wsgpr(i+1, (v >> 32) & 0xffffffff)
-
-  def rsrc(self, v: int, lane: int) -> int:
-    if v < SGPR_COUNT: return self.sgpr[v]
-    if v == SCC: return self.scc
-    if v < 255: return _INLINE_CONSTS[v - 128]
-    if v == 255: return self.literal
-    return self.vgpr[lane][v - 256] if v <= 511 else 0
-
-  def rsrc64(self, v: int, lane: int) -> int:
-    return self.rsrc(v, lane) | ((self.rsrc(v+1, lane) if v < VCC_LO or 256 <= v <= 511 else 0) << 32)
-
-  def pend_sgpr_lane(self, reg: int, lane: int, val: int) -> None:
-    if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
-    if val: self._pend_sgpr[reg] |= (1 << lane)
-
-  def commit_pends(self) -> None:
-    for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
-    self._pend_sgpr.clear()
-
-def decode_format(word: int) -> tuple[type[Inst] | None, bool]:
-  hi2 = (word >> 30) & 0x3
-  if hi2 == 0b11:
-    enc = (word >> 26) & 0xf
-    if enc == 0b1101: return SMEM, True
-    if enc == 0b0101:
-      op = (word >> 16) & 0x3ff
-      return (VOP3SD, True) if op in (288, 289, 290, 764, 765, 766, 767, 768, 769, 770) else (VOP3, True)
-    return {0b0011: (VOP3P, True), 0b0110: (DS, True), 0b0111: (FLAT, True), 0b0010: (VOPD, True)}.get(enc, (None, True))
-  if hi2 == 0b10:
-    enc = (word >> 23) & 0x7f
-    return {0b1111101: (SOP1, False), 0b1111110: (SOPC, False), 0b1111111: (SOPP, False)}.get(enc, (SOPK, False) if ((word >> 28) & 0xf) == 0b1011 else (SOP2, False))
-  enc = (word >> 25) & 0x7f
-  return (VOPC, False) if enc == 0b0111110 else (VOP1, False) if enc == 0b0111111 else (VOP2, False)
-
-def _unwrap(v) -> int: return v.val if isinstance(v, RawImm) else v.value if hasattr(v, 'value') else v
-
-def decode_program(data: bytes) -> Program:
-  result: Program = {}
-  i = 0
-  while i < len(data):
-    word = int.from_bytes(data[i:i+4], 'little')
-    inst_class, is_64 = decode_format(word)
-    if inst_class is None: i += 4; continue
-    base_size = 8 if is_64 else 4
-    inst = inst_class.from_bytes(data[i:i+base_size])
-    for name, val in inst._values.items(): setattr(inst, name, _unwrap(val))
-    has_literal = any(getattr(inst, fld, None) == 255 for fld in ('src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'srcx0', 'srcy0'))
-    if inst_class == VOP2 and inst.op in (44, 45, 55, 56): has_literal = True
-    if inst_class == VOPD and (inst.opx in (1, 2) or inst.opy in (1, 2)): has_literal = True
-    if inst_class == SOP2 and inst.op in (69, 70): has_literal = True
-    if has_literal: inst._literal = int.from_bytes(data[i+base_size:i+base_size+4], 'little')
-    inst._words = inst.size() // 4  # cache size for step_wave
-    result[i // 4] = inst
-    i += inst._words * 4
-  return result
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# SCALAR EXECUTION
-# ═══════════════════════════════════════════════════════════════════════════════
-def exec_sop1(st: WaveState, inst: SOP1) -> int:
-  s0, op = st.rsrc(inst.ssrc0, 0), inst.op
-  # 64-bit and special ops handled inline
-  if op == SOP1Op.S_MOV_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0)); return 0
-  if op == SOP1Op.S_NOT_B64: r = (~st.rsrc64(inst.ssrc0, 0)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
-  if op == SOP1Op.S_BITSET0_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) & ~(1 << (s0 & 0x1f))); return 0
-  if op == SOP1Op.S_BITSET1_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) | (1 << (s0 & 0x1f))); return 0
-  if op == SOP1Op.S_AND_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
-  if op == SOP1Op.S_OR_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 | old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
-  if op == SOP1Op.S_AND_NOT1_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & (~old & 0xffffffff); st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
-  if op == SOP1Op.S_GETPC_B64: return -3
-  if op == SOP1Op.S_SETPC_B64: return -4
-  if op == SOP1Op.S_SWAPPC_B64: return -5
-  if (fn := SALU.get(SOP1_BASE + op)) is None: raise NotImplementedError(f"SOP1 op {op}")
-  r, scc = fn(s0, 0, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
-
-_SOP2_64: dict[int, Callable[[int, int], int]] = {SOP2Op.S_AND_B64: lambda a, b: a & b, SOP2Op.S_OR_B64: lambda a, b: a | b, SOP2Op.S_XOR_B64: lambda a, b: a ^ b}
-def exec_sop2(st: WaveState, inst: SOP2) -> int:
-  s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
-  # 64-bit ops handled inline
-  if op == SOP2Op.S_LSHL_B64: r = (st.rsrc64(inst.ssrc0, 0) << (s1 & 0x3f)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
-  if op == SOP2Op.S_LSHR_B64: r = st.rsrc64(inst.ssrc0, 0) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
-  if op == SOP2Op.S_ASHR_I64: r = sext(st.rsrc64(inst.ssrc0, 0), 64) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r & 0xffffffffffffffff); st.scc = int(r != 0); return 0
-  if (fn := _SOP2_64.get(op)): r = fn(st.rsrc64(inst.ssrc0, 0), st.rsrc64(inst.ssrc1, 0)); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
-  if op == SOP2Op.S_CSELECT_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0) if st.scc else st.rsrc64(inst.ssrc1, 0)); return 0
-  if op == SOP2Op.S_FMAC_F32: st.wsgpr(inst.sdst, i32(f32(st.rsgpr(inst.sdst)) + f32(s0) * f32(s1))); return 0
-  if op == SOP2Op.S_FMAAK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(s1) + f32(inst._literal or 0))); return 0
-  if op == SOP2Op.S_FMAMK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(inst._literal or 0) + f32(s1))); return 0
-  if (fn := SALU.get(SOP2_BASE + op)) is None: raise NotImplementedError(f"SOP2 op {op}")
-  r, scc = fn(s0, s1, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
-
-def exec_sopc(st: WaveState, inst: SOPC) -> int:
-  s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
-  if op == SOPCOp.S_CMP_EQ_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) == st.rsrc64(inst.ssrc1, 0)); return 0
-  if op == SOPCOp.S_CMP_LG_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) != st.rsrc64(inst.ssrc1, 0)); return 0
-  if (fn := SALU.get(SOPC_BASE + op)) is None: raise NotImplementedError(f"SOPC op {op}")
-  st.scc = fn(s0, s1, st.scc)[1]; return 0
-
-_SOPK_CMP = frozenset((SOPKOp.S_CMPK_EQ_I32, SOPKOp.S_CMPK_LG_I32, SOPKOp.S_CMPK_GT_I32, SOPKOp.S_CMPK_GE_I32,
-  SOPKOp.S_CMPK_LT_I32, SOPKOp.S_CMPK_LE_I32, SOPKOp.S_CMPK_EQ_U32, SOPKOp.S_CMPK_LG_U32,
-  SOPKOp.S_CMPK_GT_U32, SOPKOp.S_CMPK_GE_U32, SOPKOp.S_CMPK_LT_U32, SOPKOp.S_CMPK_LE_U32))
-def exec_sopk(st: WaveState, inst: SOPK) -> int:
-  simm, s0, op = inst.simm16, st.rsgpr(inst.sdst), inst.op
-  if op in SOPK_WAIT: return 0
-  if (fn := SALU.get(SOPK_BASE + op)) is None: raise NotImplementedError(f"SOPK op {op}")
-  r, scc = fn(s0, simm, st.scc)
-  if op not in _SOPK_CMP: st.wsgpr(inst.sdst, r)
-  st.scc = scc; return 0
-
-def exec_sopp(st: WaveState, inst: SOPP) -> int:
-  if inst.op == SOPPOp.S_ENDPGM: return -1
-  if inst.op == SOPPOp.S_BARRIER: return -2
-  if inst.op == SOPPOp.S_BRANCH: return sext(inst.simm16, 16)
-  if inst.op == SOPPOp.S_CBRANCH_SCC0: return sext(inst.simm16, 16) if st.scc == 0 else 0
-  if inst.op == SOPPOp.S_CBRANCH_SCC1: return sext(inst.simm16, 16) if st.scc == 1 else 0
-  # In wave32 mode, only VCC_LO is used for lane masks; VCC_HI is a free SGPR
-  if inst.op == SOPPOp.S_CBRANCH_VCCZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) == 0 else 0
-  if inst.op == SOPPOp.S_CBRANCH_VCCNZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) != 0 else 0
-  if inst.op == SOPPOp.S_CBRANCH_EXECZ: return sext(inst.simm16, 16) if st.exec_mask == 0 else 0
-  if inst.op == SOPPOp.S_CBRANCH_EXECNZ: return sext(inst.simm16, 16) if st.exec_mask != 0 else 0
-  # Scheduling hints and wait instructions are no-ops in emulation
-  if inst.op <= 31: return 0  # S_NOP, S_CLAUSE, S_DELAY_ALU, S_WAITCNT, etc.
-  # S_WAKEUP(52), S_SETPRIO(53), S_SENDMSG(54), S_SENDMSGHALT(55), perf counters, S_ICACHE_INV(60) are no-ops
-  if inst.op in (52, 53, 54, 55, 56, 57, 60): return 0
-  raise NotImplementedError(f"SOPP op {inst.op}")
-
-def exec_smem(st: WaveState, inst: SMEM) -> int:
-  addr = st.rsgpr64(inst.sbase * 2) + sext(inst.offset, 21)
-  if inst.soffset not in (NULL, 0x7f): addr += st.rsrc(inst.soffset, 0)
-  if (cnt := SMEM_LOAD.get(inst.op)) is None: raise NotImplementedError(f"SMEM op {inst.op}")
-  for i in range(cnt): st.wsgpr(inst.sdata + i, mem_read((addr + i * 4) & 0xffffffffffffffff, 4))
-  return 0
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# VECTOR EXECUTION
-# ═══════════════════════════════════════════════════════════════════════════════
-def f64(hi: int, lo: int) -> float: return struct.unpack('<d', struct.pack('<Q', (hi << 32) | lo))[0]
-def i64_parts(f: float) -> tuple[int, int]:
-  if math.isnan(f): val = 0x7ff8000000000000
-  elif math.isinf(f): val = 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
-  else: val = struct.unpack('<Q', struct.pack('<d', f))[0]
-  return val & 0xffffffff, (val >> 32) & 0xffffffff
-
-def exec_vop1(st: WaveState, inst: VOP1, lane: int) -> None:
-  if inst.op == VOP1Op.V_NOP: return
-  V, s0 = st.vgpr[lane], st.rsrc(inst.src0, lane)
-  if inst.op == VOP1Op.V_READFIRSTLANE_B32:
-    first = (st.exec_mask & -st.exec_mask).bit_length() - 1 if st.exec_mask else 0
-    st.wsgpr(inst.vdst, st.rsrc(inst.src0, first) if inst.src0 >= 256 else s0); return
-  # F64 ops handled inline
-  if inst.op == VOP1Op.V_CVT_F64_F32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(f32(s0))); return
-  if inst.op == VOP1Op.V_CVT_F64_I32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(sext(s0, 32))); return
-  if inst.op == VOP1Op.V_CVT_F64_U32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(s0)); return
-  if inst.op in (VOP1Op.V_CVT_F32_F64, VOP1Op.V_CVT_I32_F64, VOP1Op.V_CVT_U32_F64):
-    src = inst.src0 - 256 if inst.src0 >= 256 else inst.src0
-    lo, hi = (V[src], V[src+1]) if inst.src0 >= 256 else (st.sgpr[src], st.sgpr[src+1])
-    v = f64(hi, lo)
-    if inst.op == VOP1Op.V_CVT_F32_F64: V[inst.vdst] = i32(v)
-    elif inst.op == VOP1Op.V_CVT_I32_F64: V[inst.vdst] = (max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff) if math.isfinite(v) else 0
-    else: V[inst.vdst] = max(0, min(0xffffffff, int(v))) if math.isfinite(v) and v == v else 0
-    return
-  if (fn := VALU.get(VOP1_BASE + inst.op)): V[inst.vdst] = fn(s0, 0, 0); return
-  raise NotImplementedError(f"VOP1 op {inst.op}")
-
-def exec_vop2(st: WaveState, inst: VOP2, lane: int) -> None:
-  V, s0, s1, op = st.vgpr[lane], st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1], inst.op
-  if op == VOP2Op.V_CNDMASK_B32: V[inst.vdst] = s1 if (st.vcc >> lane) & 1 else s0; return
-  if op == VOP2Op.V_FMAC_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(V[inst.vdst])); return
-  if op == VOP2Op.V_FMAMK_F32: V[inst.vdst] = i32(f32(s0)*f32(st.literal)+f32(s1)); return
-  if op == VOP2Op.V_FMAAK_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(st.literal)); return
-  if op == VOP2Op.V_FMAC_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(V[inst.vdst])); return
-  if op == VOP2Op.V_FMAMK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(st.literal)+f16(s1)); return
-  if op == VOP2Op.V_FMAAK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(st.literal)); return
-  if op == VOP2Op.V_PK_FMAC_F16:
-    lo = i16(f16(s0 & 0xffff) * f16(s1 & 0xffff) + f16(V[inst.vdst] & 0xffff))
-    hi = i16(f16((s0 >> 16) & 0xffff) * f16((s1 >> 16) & 0xffff) + f16((V[inst.vdst] >> 16) & 0xffff))
-    V[inst.vdst] = lo | (hi << 16); return
-  if op == VOP2Op.V_ADD_CO_CI_U32: r = s0+s1+((st.vcc>>lane)&1); st.pend_sgpr_lane(VCC_LO, lane, r >= 0x100000000); V[inst.vdst] = r & 0xffffffff; return
-  if op == VOP2Op.V_SUB_CO_CI_U32: b = (st.vcc>>lane)&1; st.pend_sgpr_lane(VCC_LO, lane, s1+b > s0); V[inst.vdst] = (s0-s1-b) & 0xffffffff; return
-  if (fn := VALU.get(VOP2_BASE + op)): V[inst.vdst] = fn(s0, s1, 0); return
-  raise NotImplementedError(f"VOP2 op {op}")
-
-def vop3_mod(val: int, neg: int, abs_: int, idx: int) -> int:
-  if (abs_ >> idx) & 1: val = i32(abs(f32(val)))
-  if (neg >> idx) & 1: val = i32(-f32(val))
-  return val
-
-def exec_vop3(st: WaveState, inst: VOP3, lane: int) -> None:
-  op, src0, src1, src2, vdst, neg, abs_ = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.neg, getattr(inst, 'abs', 0)
-  V = st.vgpr[lane]
-  # VOPC encoded in VOP3 (0-255)
-  if 0 <= op <= 255:
-    base = op & 0x7f
-    # For 64-bit comparisons (I64: 80-87, U64: 88-95), read raw 64-bit values (no float modifiers)
-    if 80 <= base <= 95:
-      s0_64, s1_64 = st.rsrc64(src0, lane), st.rsrc64(src1, lane)
-      result = vopc(op, s0_64 & 0xffffffff, s1_64 & 0xffffffff, (s0_64 >> 32) & 0xffffffff, (s1_64 >> 32) & 0xffffffff)
-    else:
-      s0, s1 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1)
-      result = vopc(op, s0, s1)
-    is_cmpx = op >= 128
-    st.pend_sgpr_lane(vdst, lane, result)
-    if is_cmpx: st.pend_sgpr_lane(EXEC_LO, lane, result)
-    return
-  s0, s1, s2 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1), vop3_mod(st.rsrc(src2, lane), neg, abs_, 2)
-  # Special ops
-  if op == VOP3Op.V_FMAC_F32: V[vdst] = i32(f32(s0)*f32(s1)+f32(V[vdst])); return
-  if op == VOP3Op.V_READLANE_B32: st.wsgpr(vdst, st.vgpr[s1 & 0x1f][src0 - 256] if src0 >= 256 else s0); return
-  if op == VOP3Op.V_WRITELANE_B32: st.vgpr[s1 & 0x1f][vdst] = s0; return
-  if op == VOP3Op.V_CNDMASK_B32:
-    mask = st.rsgpr(src2) if src2 < 256 else st.vcc
-    V[vdst] = s1 if (mask >> lane) & 1 else s0; return
-  if op in (VOP3Op.V_LSHLREV_B64, VOP3Op.V_LSHRREV_B64, VOP3Op.V_ASHRREV_I64):
-    v64 = st.rsrc64(src1, lane)
-    r = ((v64 << (s0 & 0x3f)) & 0xffffffffffffffff if op == VOP3Op.V_LSHLREV_B64 else
-         v64 >> (s0 & 0x3f) if op == VOP3Op.V_LSHRREV_B64 else sext(v64, 64) >> (s0 & 0x3f))
-    V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff; return
-  if op in (VOP3Op.V_ADD_F64, VOP3Op.V_MUL_F64, VOP3Op.V_FMA_F64, VOP3Op.V_MAX_F64, VOP3Op.V_MIN_F64):
-    a, b = f64(st.rsrc(src0+1, lane), s0), f64(st.rsrc(src1+1, lane), s1)
-    c = f64(st.rsrc(src2+1, lane), s2) if op == VOP3Op.V_FMA_F64 else 0.0
-    rf = a + b if op == VOP3Op.V_ADD_F64 else a * b if op == VOP3Op.V_MUL_F64 else a * b + c if op == VOP3Op.V_FMA_F64 else max(a, b) if op == VOP3Op.V_MAX_F64 else min(a, b)
-    V[vdst], V[vdst+1] = i64_parts(rf); return
-  if (fn := VALU.get(op)): V[vdst] = fn(s0, s1, s2); return
-  raise NotImplementedError(f"VOP3 op {op}")
-
-def exec_vopc(st: WaveState, inst: VOPC, lane: int) -> None:
-  result, is_cmpx = vopc(inst.op, st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1]), inst.op >= 128
-  st.pend_sgpr_lane(EXEC_LO if is_cmpx else VCC_LO, lane, result)
-
-def exec_vop3sd(st: WaveState, inst: VOP3SD, lane: int) -> None:
-  op, src0, src1, src2, vdst, sdst, neg = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.sdst, inst.neg
-  s0, s1, s2 = st.rsrc(src0, lane), st.rsrc(src1, lane), st.rsrc(src2, lane)
-  if (neg >> 0) & 1: s0 = i32(-f32(s0))
-  if (neg >> 1) & 1: s1 = i32(-f32(s1))
-  if (neg >> 2) & 1: s2 = i32(-f32(s2))
-  V = st.vgpr[lane]
-  if op == VOP3SDOp.V_ADD_CO_U32: r = s0 + s1; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
-  elif op == VOP3SDOp.V_SUB_CO_U32: V[vdst] = (s0 - s1) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 > s0)
-  elif op == VOP3SDOp.V_SUBREV_CO_U32: V[vdst] = (s1 - s0) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s0 > s1)
-  elif op == VOP3SDOp.V_ADD_CO_CI_U32:
-    cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
-    r = s0 + s1 + cin; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
-  elif op == VOP3SDOp.V_SUB_CO_CI_U32:
-    cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
-    V[vdst] = (s0 - s1 - cin) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 + cin > s0)
-  elif op == VOP3SDOp.V_MAD_U64_U32:
-    s2_64 = s2 | (st.rsrc(src2+1, lane) << 32); r = s0 * s1 + s2_64
-    V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
-  elif op == VOP3SDOp.V_MAD_I64_I32:
-    s2_64 = sext(s2 | (st.rsrc(src2+1, lane) << 32), 64)
-    r = (sext(s0, 32) * sext(s1, 32) + s2_64) & 0xffffffffffffffff
-    V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
-  elif op == VOP3SDOp.V_DIV_SCALE_F32: V[vdst] = 0; st.pend_sgpr_lane(sdst, lane, False)
-  elif op == VOP3SDOp.V_DIV_SCALE_F64: V[vdst], V[vdst+1] = s0, st.rsrc(src0+1, lane); st.pend_sgpr_lane(VCC_LO, lane, s0 == s2)
-  else: raise NotImplementedError(f"VOP3SD op {op}")
-
-def exec_flat(st: WaveState, inst: FLAT, lane: int) -> None:
-  op, addr_reg, data_reg, vdst, offset, saddr, V = inst.op, inst.addr, inst.data, inst.vdst, sext(inst.offset, 13), inst.saddr, st.vgpr[lane]
-  addr = V[addr_reg] | (V[addr_reg+1] << 32)
-  addr = (st.rsgpr64(saddr) + V[addr_reg] + offset) & 0xffffffffffffffff if saddr not in (NULL, 0x7f) else (addr + offset) & 0xffffffffffffffff
-  if op in FLAT_LOAD:
-    cnt, sz, sign = FLAT_LOAD[op]
-    for i in range(cnt): val = mem_read(addr + i * sz, sz); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
-  elif op in FLAT_STORE:
-    cnt, sz = FLAT_STORE[op]
-    for i in range(cnt): mem_write(addr + i * sz, sz, V[data_reg + i] & ((1 << (sz * 8)) - 1))
-  elif op in FLAT_D16_LO: sz, sign = FLAT_D16_LO[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0xffff0000) | ((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff))
-  elif op in FLAT_D16_HI: sz, sign = FLAT_D16_HI[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0x0000ffff) | (((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff)) << 16)
-  elif op in FLAT_D16_STORE: mem_write(addr, FLAT_D16_STORE[op], (V[data_reg] >> 16) & ((1 << (FLAT_D16_STORE[op] * 8)) - 1))
-  else: raise NotImplementedError(f"FLAT op {op}")
-
-def exec_ds(st: WaveState, inst: DS, lane: int, lds: bytearray) -> None:
-  op, addr, vdst, V = inst.op, (st.vgpr[lane][inst.addr] + inst.offset0) & 0xffff, inst.vdst, st.vgpr[lane]
-  if op in DS_LOAD:
-    cnt, sz, sign = DS_LOAD[op]
-    for i in range(cnt): val = int.from_bytes(lds[addr+i*sz:addr+i*sz+sz], 'little'); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
-  elif op in DS_STORE:
-    cnt, sz = DS_STORE[op]
-    for i in range(cnt): lds[addr+i*sz:addr+i*sz+sz] = (V[inst.data0 + i] & ((1 << (sz * 8)) - 1)).to_bytes(sz, 'little')
-  else: raise NotImplementedError(f"DS op {op}")
-
-VOPD_OPS: dict[int, Callable[[int, int, int, int, int], int]] = {
-  VOPDOp.V_DUAL_MUL_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)), VOPDOp.V_DUAL_ADD_F32: lambda a, b, d, l, lit: i32(f32(a)+f32(b)),
-  VOPDOp.V_DUAL_SUB_F32: lambda a, b, d, l, lit: i32(f32(a)-f32(b)), VOPDOp.V_DUAL_SUBREV_F32: lambda a, b, d, l, lit: i32(f32(b)-f32(a)),
-  VOPDOp.V_DUAL_MAX_F32: lambda a, b, d, l, lit: i32(max(f32(a), f32(b))), VOPDOp.V_DUAL_MIN_F32: lambda a, b, d, l, lit: i32(min(f32(a), f32(b))),
-  VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: lambda a, b, d, l, lit: i32(0.0 if f32(a) == 0.0 or f32(b) == 0.0 else f32(a)*f32(b)),
-  VOPDOp.V_DUAL_MOV_B32: lambda a, b, d, l, lit: a, VOPDOp.V_DUAL_ADD_NC_U32: lambda a, b, d, l, lit: (a + b) & 0xffffffff,
-  VOPDOp.V_DUAL_LSHLREV_B32: lambda a, b, d, l, lit: (b << (a & 0x1f)) & 0xffffffff, VOPDOp.V_DUAL_AND_B32: lambda a, b, d, l, lit: a & b,
-  VOPDOp.V_DUAL_FMAC_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(d)), VOPDOp.V_DUAL_FMAAK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(lit)),
-  VOPDOp.V_DUAL_FMAMK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(lit)+f32(b)), VOPDOp.V_DUAL_CNDMASK_B32: lambda a, b, d, l, lit: b if l else a,
-}
-def exec_vopd(st: WaveState, inst: VOPD, lane: int) -> None:
-  V, vdsty, vcc_lane = st.vgpr[lane], (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), (st.vcc >> lane) & 1
-  sx0, sx1, sy0, sy1, dstx = st.rsrc(inst.srcx0, lane), V[inst.vsrcx1], st.rsrc(inst.srcy0, lane), V[inst.vsrcy1], inst.vdstx
-  if (fn := VOPD_OPS.get(inst.opx)): V[dstx] = fn(sx0, sx1, V[dstx], vcc_lane, st.literal)
-  else: raise NotImplementedError(f"VOPD opx {inst.opx}")
-  if (fn := VOPD_OPS.get(inst.opy)): V[vdsty] = fn(sy0, sy1, V[vdsty], vcc_lane, st.literal)
-  else: raise NotImplementedError(f"VOPD opy {inst.opy}")
-
-def exec_vop3p(st: WaveState, inst: VOP3P, lane: int) -> None:
-  op, vdst, V = inst.op, inst.vdst, st.vgpr[lane]
-  s0, s1, s2 = st.rsrc(inst.src0, lane), st.rsrc(inst.src1, lane), st.rsrc(inst.src2, lane)
-  opsel, opsel_hi = [(inst.opsel >> i) & 1 for i in range(3)], [(inst.opsel_hi >> i) & 1 for i in range(2)] + [inst.opsel_hi2]
-  neg, neg_hi = inst.neg, inst.neg_hi
-  def get_src(src: int, idx: int, for_mix: bool = False) -> float:
-    if for_mix:
-      if not opsel_hi[idx]: return abs(f32(src)) if (neg_hi >> idx) & 1 else f32(src)
-      return float(f16((src >> 16) & 0xffff) if opsel[idx] else f16(src & 0xffff))
-    use_hi = opsel[idx]
-    val = ((src >> 16) & 0xffff) if use_hi else (src & 0xffff)
-    f = f16(val)
-    if use_hi and (neg >> idx) & 1: f = -f
-    elif not use_hi and (neg_hi >> idx) & 1: f = -f
-    return f
-  if op == VOP3POp.V_FMA_MIX_F32: V[vdst] = i32(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
-  elif op == VOP3POp.V_FMA_MIXLO_F16: V[vdst] = (V[vdst] & 0xffff0000) | i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
-  elif op == VOP3POp.V_FMA_MIXHI_F16: V[vdst] = (V[vdst] & 0x0000ffff) | (i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True)) << 16)
-  else: raise NotImplementedError(f"VOP3P op {op}")
-
-def exec_wmma_f32_16x16x16_f16(st: WaveState, inst: VOP3P, n_lanes: int) -> None:
-  src0_base, src1_base, src2_base = (inst.src0 - 256) if inst.src0 >= 256 else inst.src0, (inst.src1 - 256) if inst.src1 >= 256 else inst.src1, (inst.src2 - 256) if inst.src2 >= 256 else inst.src2
-  src0_is_vgpr, src1_is_vgpr, src2_is_vgpr, vdst = inst.src0 >= 256, inst.src1 >= 256, inst.src2 >= 256, inst.vdst
-  A, B, C = [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)]
-  for lane in range(min(n_lanes, 16)):
-    V = st.vgpr[lane]
-    for reg in range(8):
-      val = V[src0_base + reg] if src0_is_vgpr else st.sgpr[src0_base + reg]
-      A[lane][reg * 2], A[lane][reg * 2 + 1] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
-      val = V[src1_base + reg] if src1_is_vgpr else st.sgpr[src1_base + reg]
-      B[reg * 2][lane], B[reg * 2 + 1][lane] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
-  for row in range(16):
-    for col in range(16):
-      idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
-      if lane_idx < n_lanes:
-        val = st.vgpr[lane_idx][src2_base + reg] if src2_is_vgpr else st.sgpr[src2_base + reg]
-        C[row][col] = f32(val)
-  for row in range(16):
-    for col in range(16):
-      for k in range(16): C[row][col] += A[row][k] * B[k][col]
-  for row in range(16):
-    for col in range(16):
-      idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
-      if lane_idx < n_lanes and (st.exec_mask & (1 << lane_idx)): st.vgpr[lane_idx][vdst + reg] = i32(C[row][col])
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# MAIN EXECUTION LOOP
-# ═══════════════════════════════════════════════════════════════════════════════
-SCALAR: dict[type, Callable[..., int]] = {SOP1: exec_sop1, SOP2: exec_sop2, SOPC: exec_sopc, SOPK: exec_sopk, SOPP: exec_sopp, SMEM: exec_smem}
-VECTOR: dict[type, Callable[..., None]] = {VOP1: exec_vop1, VOP2: exec_vop2, VOP3: exec_vop3, VOP3SD: exec_vop3sd, VOPC: exec_vopc, FLAT: exec_flat, DS: exec_ds, VOPD: exec_vopd, VOP3P: exec_vop3p}
-
-_WMMA_OPS = frozenset((VOP3POp.V_WMMA_F32_16X16X16_F16, VOP3POp.V_WMMA_F32_16X16X16_BF16, VOP3POp.V_WMMA_F16_16X16X16_F16,
-                       VOP3POp.V_WMMA_BF16_16X16X16_BF16, VOP3POp.V_WMMA_I32_16X16X16_IU8, VOP3POp.V_WMMA_I32_16X16X16_IU4))
-
-def step_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> int:
-  inst = program.get(st.pc)
-  if inst is None: return 1
-  inst_words, st.literal, inst_type = inst._words, inst._literal or 0, type(inst)
-  if (handler := SCALAR.get(inst_type)) is not None:
-    delta = handler(st, inst)
-    if delta == -1: return -1
-    if delta == -2: st.pc += inst_words; return -2
-    if delta == -3:  # S_GETPC_B64
-      sop1 = inst if isinstance(inst, SOP1) else None
-      assert sop1 is not None
-      next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc += inst_words; return 0
-    if delta == -4:  # S_SETPC_B64
-      sop1 = inst if isinstance(inst, SOP1) else None
-      assert sop1 is not None
-      st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
-    if delta == -5:  # S_SWAPPC_B64
-      sop1 = inst if isinstance(inst, SOP1) else None
-      assert sop1 is not None
-      next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
-    st.pc += inst_words + delta
-  else:
-    vec_handler, exec_mask = VECTOR[inst_type], st.exec_mask
-    if inst_type is DS:
-      for lane in range(n_lanes):
-        if exec_mask & (1 << lane): vec_handler(st, inst, lane, lds)
-    elif inst_type is VOP3P:
-      vop3p = inst if isinstance(inst, VOP3P) else None
-      assert vop3p is not None
-      if vop3p.op in _WMMA_OPS:
-        exec_wmma_f32_16x16x16_f16(st, vop3p, n_lanes)
-      else:
-        for lane in range(n_lanes):
-          if exec_mask & (1 << lane): vec_handler(st, vop3p, lane)
-    else:
-      for lane in range(n_lanes):
-        if exec_mask & (1 << lane): vec_handler(st, inst, lane)
-    st.commit_pends(); st.pc += inst_words
-  return 0
-
-def exec_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int, wg_id: tuple[int,int,int]=(0,0,0), local_size: tuple[int,int,int]=(1,1,1), wave_start: int=0) -> int:
-  while st.pc in program:
-    result = step_wave(program, st, lds, n_lanes)
-    if result == -1: return 0
-    if result == -2: return -2
-  return 0
-
-def exec_workgroup(program: Program, workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int, dispatch_dim: int) -> None:
-  lx, ly, lz = local_size
-  total_threads, lds = lx * ly * lz, bytearray(65536)
-  waves: list[tuple[WaveState, int, int]] = []
-  for wave_start in range(0, total_threads, WAVE_SIZE):
-    n_lanes, st = min(WAVE_SIZE, total_threads - wave_start), WaveState()
-    st.exec_mask = (1 << n_lanes) - 1
-    st.wsgpr64(0, args_ptr)
-    gx, gy, gz = workgroup_id
-    if dispatch_dim >= 3: st.sgpr[13], st.sgpr[14], st.sgpr[15] = gx, gy, gz
-    elif dispatch_dim == 2: st.sgpr[14], st.sgpr[15] = gx, gy
-    else: st.sgpr[15] = gx
-    for i in range(n_lanes):
-      tid = wave_start + i
-      st.vgpr[i][0] = tid if local_size == (lx, 1, 1) else ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
-    waves.append((st, n_lanes, wave_start))
-  has_barrier = any(isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER for inst in program.values())
-  for _ in range(2 if has_barrier else 1):
-    for st, n_lanes, wave_start in waves: exec_wave(program, st, lds, n_lanes, workgroup_id, local_size, wave_start)
-
-def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
-  data = (ctypes.c_char * lib_sz).from_address(lib).raw
-  program = decode_program(data)
-  if not program: return -1
-  dispatch_dim = 3 if gz > 1 else (2 if gy > 1 else 1)
-  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, dispatch_dim)
-  return 0

extra/assembly/rdna3/gen.py

@@ -1,191 +0,0 @@
-#!/usr/bin/env python3
-# generates autogen/__init__.py by parsing the AMD RDNA3.5 ISA PDF
-import re, pdfplumber, pathlib
-from tinygrad.helpers import fetch
-
-PDF_URL = "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content"
-FIELD_TYPES = {'SSRC0': 'SSrc', 'SSRC1': 'SSrc', 'SOFFSET': 'SSrc', 'SADDR': 'SSrc', 'SRC0': 'Src', 'SRC1': 'Src', 'SRC2': 'Src',
-  'SDST': 'SGPRField', 'SBASE': 'SGPRField', 'SDATA': 'SGPRField', 'SRSRC': 'SGPRField', 'VDST': 'VGPRField', 'VSRC1': 'VGPRField', 'VDATA': 'VGPRField',
-  'VADDR': 'VGPRField', 'ADDR': 'VGPRField', 'DATA': 'VGPRField', 'DATA0': 'VGPRField', 'DATA1': 'VGPRField', 'SIMM16': 'SImm', 'OFFSET': 'Imm',
-  'OPX': 'VOPDOp', 'OPY': 'VOPDOp', 'SRCX0': 'Src', 'SRCY0': 'Src', 'VSRCX1': 'VGPRField', 'VSRCY1': 'VGPRField', 'VDSTX': 'VGPRField', 'VDSTY': 'VDSTYEnc'}
-FIELD_ORDER = {
-  'SOP2': ['op', 'sdst', 'ssrc0', 'ssrc1'], 'SOP1': ['op', 'sdst', 'ssrc0'], 'SOPC': ['op', 'ssrc0', 'ssrc1'],
-  'SOPK': ['op', 'sdst', 'simm16'], 'SOPP': ['op', 'simm16'], 'VOP1': ['op', 'vdst', 'src0'], 'VOPC': ['op', 'src0', 'vsrc1'],
-  'VOP2': ['op', 'vdst', 'src0', 'vsrc1'], 'VOP3SD': ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2', 'clmp'],
-  'SMEM': ['op', 'sdata', 'sbase', 'soffset', 'offset', 'glc', 'dlc'], 'DS': ['op', 'vdst', 'addr', 'data0', 'data1'],
-  'VOP3': ['op', 'vdst', 'src0', 'src1', 'src2', 'omod', 'neg', 'abs', 'clmp', 'opsel'],
-  'VOP3P': ['op', 'vdst', 'src0', 'src1', 'src2', 'neg', 'neg_hi', 'opsel', 'opsel_hi', 'clmp'],
-  'FLAT': ['op', 'vdst', 'addr', 'data', 'saddr', 'offset', 'seg', 'dlc', 'glc', 'slc'],
-  'MUBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
-  'MTBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
-  'MIMG': ['op', 'vdata', 'vaddr', 'srsrc', 'ssamp', 'dmask', 'dim', 'unrm', 'dlc', 'glc', 'slc'],
-  'EXP': ['en', 'target', 'vsrc0', 'vsrc1', 'vsrc2', 'vsrc3', 'done', 'row'],
-  'VINTERP': ['op', 'vdst', 'src0', 'src1', 'src2', 'waitexp', 'clmp', 'opsel', 'neg'],
-  'VOPD': ['opx', 'opy', 'vdstx', 'vdsty', 'srcx0', 'vsrcx1', 'srcy0', 'vsrcy1'],
-  'LDSDIR': ['op', 'vdst', 'attr', 'attr_chan', 'wait_va']}
-SRC_EXTRAS = {233: 'DPP8', 234: 'DPP8FI', 250: 'DPP16', 251: 'VCCZ', 252: 'EXECZ', 254: 'LDS_DIRECT'}
-FLOAT_MAP = {'0.5': 'POS_HALF', '-0.5': 'NEG_HALF', '1.0': 'POS_ONE', '-1.0': 'NEG_ONE', '2.0': 'POS_TWO', '-2.0': 'NEG_TWO',
-  '4.0': 'POS_FOUR', '-4.0': 'NEG_FOUR', '1/(2*PI)': 'INV_2PI', '0': 'ZERO'}
-
-def parse_bits(s: str) -> tuple[int, int] | None:
-  return (int(m.group(1)), int(m.group(2) or m.group(1))) if (m := re.match(r'\[(\d+)(?::(\d+))?\]', s)) else None
-
-def parse_fields_table(table: list, fmt: str, enums: set[str]) -> list[tuple]:
-  fields = []
-  for row in table[1:]:
-    if not row or not row[0]: continue
-    name, bits_str = row[0].split('\n')[0].strip(), (row[1] or '').split('\n')[0].strip()
-    if not (bits := parse_bits(bits_str)): continue
-    enc_val, hi, lo = None, bits[0], bits[1]
-    if name == 'ENCODING' and row[2] and (m := re.search(r"'b([01_]+)", row[2])):
-      enc_bits = m.group(1).replace('_', '')
-      enc_val = int(enc_bits, 2)
-      declared_width, actual_width = hi - lo + 1, len(enc_bits)
-      if actual_width > declared_width: lo = hi - actual_width + 1
-    ftype = f"{fmt}Op" if name == 'OP' and f"{fmt}Op" in enums else FIELD_TYPES.get(name.upper())
-    fields.append((name, hi, lo, enc_val, ftype))
-  return fields
-
-def generate(output_path: pathlib.Path|str|None = None) -> dict:
-  """Generate RDNA3.5 instruction definitions from the AMD ISA PDF. Returns dict with formats for testing."""
-  pdf = pdfplumber.open(fetch(PDF_URL))
-  pages = pdf.pages[150:200]
-  page_texts = [p.extract_text() or '' for p in pages]
-  page_tables = [[t.extract() for t in p.find_tables()] for p in pages]
-  full_text = '\n'.join(page_texts)
-
-  # parse SSRC encoding from first page with VCC_LO
-  src_enum = dict(SRC_EXTRAS)
-  for text in page_texts[:10]:
-    if 'SSRC0' in text and 'VCC_LO' in text:
-      for m in re.finditer(r'^(\d+)\s+(\S+)', text, re.M):
-        val, name = int(m.group(1)), m.group(2).rstrip('.:')
-        if name in FLOAT_MAP: src_enum[val] = FLOAT_MAP[name]
-        elif re.match(r'^[A-Z][A-Z0-9_]*$', name): src_enum[val] = name
-      break
-
-  # parse opcode tables
-  enums: dict[str, dict[int, str]] = {}
-  for m in re.finditer(r'Table \d+\. (\w+) Opcodes(.*?)(?=Table \d+\.|\n\d+\.\d+\.\d+\.\s+\w+\s*\nDescription|$)', full_text, re.S):
-    if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+([A-Z][A-Z0-9_]+)', m.group(2))}:
-      enums[m.group(1) + "Op"] = ops
-  if vopd_m := re.search(r'Table \d+\. VOPD Y-Opcodes\n(.*?)(?=Table \d+\.|15\.\d)', full_text, re.S):
-    if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+(V_DUAL_\w+)', vopd_m.group(1))}:
-      enums["VOPDOp"] = ops
-  enum_names = set(enums.keys())
-
-  def is_fields_table(t) -> bool: return t and len(t) > 1 and t[0] and 'Field' in str(t[0][0] or '')
-  def has_encoding(fields) -> bool: return any(f[0] == 'ENCODING' for f in fields)
-  def has_header_before_fields(text) -> bool:
-    return (pos := text.find('Field Name')) != -1 and bool(re.search(r'\d+\.\d+\.\d+\.\s+\w+\s*\n', text[:pos]))
-
-  # find format headers with their page indices
-  format_headers = []  # (fmt_name, page_idx)
-  for i, text in enumerate(page_texts):
-    for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n?Description', text): format_headers.append((m.group(1), i, m.start()))
-    for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n', text):
-      if m.start() > len(text) - 200 and 'Description' not in text[m.end():] and i + 1 < len(page_texts):
-        next_text = page_texts[i + 1].lstrip()
-        if next_text.startswith('Description') or (next_text.startswith('"RDNA') and 'Description' in next_text[:200]):
-          format_headers.append((m.group(1), i, m.start()))
-
-  # parse instruction formats
-  formats: dict[str, list] = {}
-  for fmt_name, page_idx, header_pos in format_headers:
-    if fmt_name in formats: continue
-    text, tables = page_texts[page_idx], page_tables[page_idx]
-    field_pos = text.find('Field Name', header_pos)
-
-    # find fields table with ENCODING (same page or up to 2 pages ahead)
-    fields = None
-    for offset in range(3):
-      if page_idx + offset >= len(pages): break
-      if offset > 0 and has_header_before_fields(page_texts[page_idx + offset]): break
-      for t in page_tables[page_idx + offset] if offset > 0 or field_pos > header_pos else []:
-        if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)) and has_encoding(f):
-          fields = f
-          break
-      if fields: break
-
-    # for modifier formats (no ENCODING), accept first fields table on same page
-    if not fields and field_pos > header_pos:
-      for t in tables:
-        if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)):
-          fields = f
-          break
-
-    if not fields: continue
-    field_names = {f[0] for f in fields}
-
-    # check next pages for continuation fields (tables without ENCODING)
-    for pg_offset in range(1, 3):
-      if page_idx + pg_offset >= len(pages) or has_header_before_fields(page_texts[page_idx + pg_offset]): break
-      for t in page_tables[page_idx + pg_offset]:
-        if is_fields_table(t) and (extra := parse_fields_table(t, fmt_name, enum_names)) and not has_encoding(extra):
-          for ef in extra:
-            if ef[0] not in field_names:
-              fields.append(ef)
-              field_names.add(ef[0])
-          break
-    formats[fmt_name] = fields
-
-  # fix known PDF errors (verified against LLVM test vectors)
-  # SMEM: PDF says DLC=bit14, GLC=bit16 but actual encoding is DLC=bit13, GLC=bit14
-  if 'SMEM' in formats:
-    formats['SMEM'] = [(n, 13 if n == 'DLC' else 14 if n == 'GLC' else h, 13 if n == 'DLC' else 14 if n == 'GLC' else l, e, t)
-                       for n, h, l, e, t in formats['SMEM']]
-
-  # generate output
-  def enum_lines(name, items):
-    return [f"class {name}(IntEnum):"] + [f"  {n} = {v}" for v, n in sorted(items.items())] + [""]
-  def field_key(f): return order.index(f[0].lower()) if f[0].lower() in order else 1000
-  lines = ["# autogenerated from AMD RDNA3.5 ISA PDF by gen.py - do not edit", "from enum import IntEnum",
-           "from typing import Annotated",
-           "from extra.assembly.rdna3.lib import bits, BitField, Inst32, Inst64, SGPR, VGPR, TTMP as TTMP, s as s, v as v, ttmp as ttmp, SSrc, Src, SImm, Imm, VDSTYEnc, SGPRField, VGPRField",
-           "import functools", ""]
-  lines += enum_lines("SrcEnum", src_enum) + sum([enum_lines(n, ops) for n, ops in sorted(enums.items())], [])
-  # Format-specific field defaults (verified against LLVM test vectors)
-  format_defaults = {'VOP3P': {'opsel_hi': 3, 'opsel_hi2': 1}}
-  lines.append("# instruction formats")
-  for fmt_name, fields in sorted(formats.items()):
-    base = "Inst64" if max(f[1] for f in fields) > 31 or fmt_name == 'VOP3SD' else "Inst32"
-    order = FIELD_ORDER.get(fmt_name, [])
-    lines.append(f"class {fmt_name}({base}):")
-    if enc := next((f for f in fields if f[0] == 'ENCODING'), None):
-      enc_str = f"bits[{enc[1]}:{enc[2]}] == 0b{enc[3]:b}" if enc[1] != enc[2] else f"bits[{enc[1]}] == {enc[3]}"
-      lines.append(f"  encoding = {enc_str}")
-    if defaults := format_defaults.get(fmt_name):
-      lines.append(f"  _defaults = {defaults}")
-    for name, hi, lo, _, ftype in sorted([f for f in fields if f[0] != 'ENCODING'], key=field_key):
-      # Wrap IntEnum types (ending in Op) with Annotated[BitField, ...] for correct typing
-      if ftype and ftype.endswith('Op'):
-        ann = f":Annotated[BitField, {ftype}]"
-      else:
-        ann = f":{ftype}" if ftype else ""
-      lines.append(f"  {name.lower()}{ann} = bits[{hi}]" if hi == lo else f"  {name.lower()}{ann} = bits[{hi}:{lo}]")
-    lines.append("")
-  lines.append("# instruction helpers")
-  for cls_name, ops in sorted(enums.items()):
-    fmt = cls_name[:-2]
-    for op_val, name in sorted(ops.items()):
-      seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=2"}.get(fmt, "")
-      tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt, f"{fmt}, {cls_name}")
-      if fmt in formats or fmt in ("GLOBAL", "SCRATCH"):
-        # VOP1/VOP2/VOPC get _e32 suffix, VOP3 promoted ops (< 512) get _e64 suffix
-        if fmt in ("VOP1", "VOP2", "VOPC"):
-          suffix = "_e32"
-        elif fmt == "VOP3" and op_val < 512:
-          suffix = "_e64"
-        else:
-          suffix = ""
-        lines.append(f"{name.lower()}{suffix} = functools.partial({tgt}.{name}{seg})")
-  # export SrcEnum values, but skip DPP8/DPP16 which conflict with class names
-  skip_exports = {'DPP8', 'DPP16'}
-  lines += [""] + [f"{name} = SrcEnum.{name}" for _, name in sorted(src_enum.items()) if name not in skip_exports] + ["OFF = NULL\n"]
-
-  if output_path is not None: pathlib.Path(output_path).write_text('\n'.join(lines))
-  return {"formats": formats, "enums": enums, "src_enum": src_enum}
-
-if __name__ == "__main__":
-  result = generate("extra/assembly/rdna3/autogen/__init__.py")
-  print(f"generated SrcEnum ({len(result['src_enum'])}) + {len(result['enums'])} opcode enums + {len(result['formats'])} format classes")

extra/assembly/rdna3/lib.py

@@ -1,248 +0,0 @@
-# library for RDNA3 assembly DSL
-from __future__ import annotations
-from enum import IntEnum
-from typing import overload, Annotated, TypeVar, Generic
-
-# Bit field DSL
-class BitField:
-  def __init__(self, hi: int, lo: int, name: str | None = None): self.hi, self.lo, self.name = hi, lo, name
-  def __set_name__(self, owner, name): self.name, self._owner = name, owner
-  def __eq__(self, val: int) -> tuple[BitField, int]: return (self, val)  # type: ignore
-  def mask(self) -> int: return (1 << (self.hi - self.lo + 1)) - 1
-  @property
-  def marker(self) -> type | None:
-    # Get marker from Annotated type hint if present
-    import typing
-    if hasattr(self, '_owner') and self.name:
-      hints = typing.get_type_hints(self._owner, include_extras=True)
-      if self.name in hints:
-        hint = hints[self.name]
-        if typing.get_origin(hint) is Annotated:
-          args = typing.get_args(hint)
-          return args[1] if len(args) > 1 else None
-    return None
-  @overload
-  def __get__(self, obj: None, objtype: type) -> BitField: ...
-  @overload
-  def __get__(self, obj: object, objtype: type | None = None) -> int: ...
-  def __get__(self, obj, objtype=None):
-    if obj is None: return self
-    val = unwrap(obj._values.get(self.name, 0))
-    # Convert to IntEnum if marker is an IntEnum subclass
-    if self.marker and isinstance(self.marker, type) and issubclass(self.marker, IntEnum):
-      try: return self.marker(val)
-      except ValueError: pass
-    return val
-
-class _Bits:
-  def __getitem__(self, key) -> BitField: return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
-bits = _Bits()
-
-# Register types
-class Reg:
-  def __init__(self, idx: int, count: int = 1, hi: bool = False): self.idx, self.count, self.hi = idx, count, hi
-  def __repr__(self): return f"{self.__class__.__name__.lower()[0]}[{self.idx}]" if self.count == 1 else f"{self.__class__.__name__.lower()[0]}[{self.idx}:{self.idx + self.count}]"
-
-T = TypeVar('T', bound=Reg)
-class _RegFactory(Generic[T]):
-  def __init__(self, cls: type[T], name: str): self._cls, self._name = cls, name
-  @overload
-  def __getitem__(self, key: int) -> Reg: ...
-  @overload
-  def __getitem__(self, key: slice) -> Reg: ...
-  def __getitem__(self, key: int | slice) -> Reg:
-    return self._cls(key.start, key.stop - key.start + 1) if isinstance(key, slice) else self._cls(key)
-  def __repr__(self): return f"<{self._name} factory>"
-
-class SGPR(Reg): pass
-class VGPR(Reg): pass
-class TTMP(Reg): pass
-s: _RegFactory[SGPR] = _RegFactory(SGPR, "SGPR")
-v: _RegFactory[VGPR] = _RegFactory(VGPR, "VGPR")
-ttmp: _RegFactory[TTMP] = _RegFactory(TTMP, "TTMP")
-
-# Field type markers (runtime classes for validation)
-class _SSrc: pass
-class _Src: pass
-class _Imm: pass
-class _SImm: pass
-class _VDSTYEnc: pass  # VOPD vdsty: encoded = actual >> 1, actual = (encoded << 1) | ((vdstx & 1) ^ 1)
-class _SGPRField: pass
-class _VGPRField: pass
-
-# Type aliases for annotations - tells mypy it's a BitField while preserving marker info
-SSrc = Annotated[BitField, _SSrc]
-Src = Annotated[BitField, _Src]
-Imm = Annotated[BitField, _Imm]
-SImm = Annotated[BitField, _SImm]
-VDSTYEnc = Annotated[BitField, _VDSTYEnc]
-SGPRField = Annotated[BitField, _SGPRField]
-VGPRField = Annotated[BitField, _VGPRField]
-class RawImm:
-  def __init__(self, val: int): self.val = val
-  def __repr__(self): return f"RawImm({self.val})"
-  def __eq__(self, other): return isinstance(other, RawImm) and self.val == other.val
-
-def unwrap(val) -> int:
-  return val.val if isinstance(val, RawImm) else val.value if hasattr(val, 'value') else val.idx if hasattr(val, 'idx') else val
-
-# Encoding helpers
-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}
-SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
-RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata'}
-
-def _encode_reg(val) -> int:
-  if isinstance(val, TTMP): return 108 + val.idx
-  return val.idx | (0x80 if val.hi else 0)
-
-def encode_src(val) -> int:
-  if isinstance(val, VGPR): return 256 + _encode_reg(val)
-  if isinstance(val, Reg): return _encode_reg(val)
-  if hasattr(val, 'value'): return val.value
-  if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
-  return 128 + val if isinstance(val, int) and 0 <= val <= 64 else 192 + (-val) if isinstance(val, int) and -16 <= val <= -1 else 255
-
-# Instruction base class
-class Inst:
-  _fields: dict[str, BitField]
-  _encoding: tuple[BitField, int] | None = None
-  _defaults: dict[str, int] = {}
-  _values: dict[str, int | RawImm]
-  _words: int  # size in 32-bit words, set by decode_program
-  _literal: int | None
-
-  def __init_subclass__(cls, **kwargs):
-    super().__init_subclass__(**kwargs)
-    cls._fields = {n: v[0] if isinstance(v, tuple) else v for n, v in cls.__dict__.items() if isinstance(v, BitField) or (isinstance(v, tuple) and len(v) == 2 and isinstance(v[0], BitField))}
-    if 'encoding' in cls._fields and isinstance(cls.__dict__.get('encoding'), tuple): cls._encoding = cls.__dict__['encoding']
-
-  def __init__(self, *args, literal: int | None = None, **kwargs):
-    self._values, self._literal = dict(self._defaults), literal
-    # Map positional args to field names
-    field_names = [n for n in self._fields if n != 'encoding']
-    orig_args = dict(zip(field_names, args))
-    orig_args.update(kwargs)
-    self._values.update(orig_args)
-    # Validate register counts for SMEM instructions (before encoding)
-    if self.__class__.__name__ == 'SMEM':
-      op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
-      if op_val is not None:
-        if hasattr(op_val, 'value'): op_val = op_val.value
-        expected_cnt = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val)
-        sdata_val = orig_args.get('sdata')
-        if expected_cnt is not None and isinstance(sdata_val, Reg) and sdata_val.count != expected_cnt:
-          raise ValueError(f"SMEM op {op_val} expects {expected_cnt} registers, got {sdata_val.count}")
-    # Validate register counts for SOP1 instructions (b32 = 1 reg, b64 = 2 regs)
-    if self.__class__.__name__ == 'SOP1':
-      op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
-      if op_val is not None and hasattr(op_val, 'name'):
-        expected = 2 if op_val.name.endswith('_B64') else 1
-        sdst_val, ssrc0_val = orig_args.get('sdst'), orig_args.get('ssrc0')
-        if isinstance(sdst_val, Reg) and sdst_val.count != expected:
-          raise ValueError(f"SOP1 {op_val.name} expects {expected} destination register(s), got {sdst_val.count}")
-        if isinstance(ssrc0_val, Reg) and ssrc0_val.count != expected:
-          raise ValueError(f"SOP1 {op_val.name} expects {expected} source register(s), got {ssrc0_val.count}")
-    # Type check and encode values
-    for name, val in list(self._values.items()):
-      if name == 'encoding': continue
-      # For RawImm, only process RAW_FIELDS to unwrap to int
-      if isinstance(val, RawImm):
-        if name in RAW_FIELDS: self._values[name] = val.val
-        continue
-      field = self._fields.get(name)
-      marker = field.marker if field else None
-      # Type validation
-      if marker is _SGPRField:
-        if isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
-        if not isinstance(val, (SGPR, TTMP, int, RawImm)): raise TypeError(f"field '{name}' requires SGPR, got {type(val).__name__}")
-      if marker is _VGPRField:
-        if not isinstance(val, VGPR): raise TypeError(f"field '{name}' requires VGPR, got {type(val).__name__}")
-      if marker is _SSrc and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires scalar source, got VGPR")
-      # Encode source fields as RawImm for consistent disassembly
-      if name in SRC_FIELDS:
-        encoded = encode_src(val)
-        self._values[name] = RawImm(encoded)
-        # Track literal value if needed (encoded as 255)
-        if encoded == 255 and self._literal is None and isinstance(val, int) and not isinstance(val, IntEnum):
-          self._literal = val
-        elif encoded == 255 and self._literal is None and isinstance(val, float):
-          import struct
-          self._literal = struct.unpack('<I', struct.pack('<f', val))[0]
-      # Encode raw register fields for consistent repr
-      elif name in RAW_FIELDS:
-        if isinstance(val, Reg): self._values[name] = _encode_reg(val)
-        elif hasattr(val, 'value'): self._values[name] = val.value  # IntEnum like SrcEnum.NULL
-      # Encode sbase (divided by 2) and srsrc/ssamp (divided by 4)
-      elif name == 'sbase' and isinstance(val, Reg):
-        self._values[name] = val.idx // 2
-      elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg):
-        self._values[name] = val.idx // 4
-      # VOPD vdsty: encode as actual >> 1 (constraint: vdsty parity must be opposite of vdstx)
-      elif marker is _VDSTYEnc and isinstance(val, VGPR):
-        self._values[name] = val.idx >> 1
-
-  def _encode_field(self, name: str, val) -> int:
-    if isinstance(val, RawImm): return val.val
-    if name in {'srsrc', 'ssamp'}: return val.idx // 4 if isinstance(val, Reg) else val
-    if name == 'sbase': return val.idx // 2 if isinstance(val, Reg) else val
-    if name in RAW_FIELDS: return _encode_reg(val) if isinstance(val, Reg) else val
-    if isinstance(val, Reg) or name in SRC_FIELDS: return encode_src(val)
-    return val.value if hasattr(val, 'value') else val
-
-  def to_int(self) -> int:
-    word = (self._encoding[1] & self._encoding[0].mask()) << self._encoding[0].lo if self._encoding else 0
-    for n, bf in self._fields.items():
-      if n != 'encoding' and n in self._values: word |= (self._encode_field(n, self._values[n]) & bf.mask()) << bf.lo
-    return word
-
-  def _get_literal(self) -> int | None:
-    for n in SRC_FIELDS:
-      if n in self._values and not isinstance(v := self._values[n], RawImm) and isinstance(v, int) and not isinstance(v, IntEnum) and not (0 <= v <= 64 or -16 <= v <= -1): return v
-    return None
-
-  def to_bytes(self) -> bytes:
-    result = self.to_int().to_bytes(self._size(), 'little')
-    return result + (lit & 0xffffffff).to_bytes(4, 'little') if (lit := self._get_literal() or getattr(self, '_literal', None)) else result
-
-  @classmethod
-  def _size(cls) -> int: return 4 if issubclass(cls, Inst32) else 8
-  def size(self) -> int: return self._size() + (4 if self._literal is not None else 0)
-
-  @classmethod
-  def from_int(cls, word: int):
-    inst = object.__new__(cls)
-    inst._values = {n: RawImm(v) if n in SRC_FIELDS else v for n, bf in cls._fields.items() if n != 'encoding' for v in [(word >> bf.lo) & bf.mask()]}
-    inst._literal = None
-    return inst
-
-  @classmethod
-  def from_bytes(cls, data: bytes):
-    inst = cls.from_int(int.from_bytes(data[:cls._size()], 'little'))
-    op_val = inst._values.get('op', 0)
-    has_literal = cls.__name__ == 'VOP2' and op_val in (44, 45, 55, 56)
-    has_literal = has_literal or (cls.__name__ == 'SOP2' and op_val in (69, 70))
-    for n in SRC_FIELDS:
-      if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255: has_literal = True
-    if has_literal and len(data) >= cls._size() + 4: inst._literal = int.from_bytes(data[cls._size():cls._size()+4], 'little')
-    return inst
-
-  def __repr__(self):
-    # Use _fields order and exclude fields that are 0/default (for consistent repr after roundtrip)
-    def is_zero(v): return (isinstance(v, int) and v == 0) or (isinstance(v, VGPR) and v.idx == 0 and v.count == 1)
-    items = [(k, self._values[k]) for k in self._fields if k in self._values and k != 'encoding'
-             and not (is_zero(self._values[k]) and k not in {'op'})]
-    lit = f", literal={hex(self._literal)}" if self._literal is not None else ""
-    return f"{self.__class__.__name__}({', '.join(f'{k}={v}' for k, v in items)}{lit})"
-
-  def __eq__(self, other):
-    if not isinstance(other, Inst): return NotImplemented
-    return self.__class__ == other.__class__ and self._values == other._values and self._literal == other._literal
-
-  def __hash__(self): return hash((self.__class__.__name__, tuple(sorted((k, repr(v)) for k, v in self._values.items())), self._literal))
-
-  def disasm(self) -> str:
-    from extra.assembly.rdna3.asm import disasm
-    return disasm(self)
-
-class Inst32(Inst): pass
-class Inst64(Inst): pass

extra/assembly/rdna3/test/bench_emu.py

@@ -1,294 +0,0 @@
-#!/usr/bin/env python3
-"""Benchmark comparing Python vs Rust RDNA3 emulators on synthetic and real tinygrad kernels."""
-import ctypes, time, os, struct, cProfile, pstats, io
-from pathlib import Path
-from typing import Callable
-
-# Set AMD=1 before importing tinygrad
-os.environ["AMD"] = "1"
-
-from extra.assembly.rdna3.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program, step_wave, WaveState, WAVE_SIZE
-
-REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
-if not REMU_PATH.exists():
-  REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.dylib"
-
-def get_rust_remu():
-  """Load the Rust libremu shared library."""
-  if not REMU_PATH.exists(): return None
-  remu = ctypes.CDLL(str(REMU_PATH))
-  remu.run_asm.restype = ctypes.c_int32
-  remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
-                           ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
-  return remu
-
-def count_instructions(kernel: bytes) -> int:
-  """Count instructions in a kernel."""
-  return len(decode_program(kernel))
-
-def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = None):
-  """Allocate buffers and return args pointer + valid ranges."""
-  if init_data is None: init_data = {}
-  buffers = []
-  for i, size in enumerate(buf_sizes):
-    padded = ((size + 15) // 16) * 16 + 16
-    data = init_data.get(i, b'\x00' * padded)
-    data_list = list(data) + [0] * (padded - len(data))
-    buf = (ctypes.c_uint8 * padded)(*data_list[:padded])
-    buffers.append(buf)
-  args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
-  args_ptr = ctypes.addressof(args)
-  ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
-  ranges.add((args_ptr, ctypes.sizeof(args)))
-  return buffers, args, args_ptr, ranges
-
-def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, iterations: int = 5):
-  """Benchmark an emulator and return average time."""
-  gx, gy, gz = global_size
-  lx, ly, lz = local_size
-  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-  lib_ptr = ctypes.addressof(kernel_buf)
-
-  # Warmup
-  run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
-
-  # Timed runs
-  times = []
-  for _ in range(iterations):
-    start = time.perf_counter()
-    result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
-    end = time.perf_counter()
-    if result != 0:
-      print(f"  {name} returned error: {result}")
-      return None
-    times.append(end - start)
-
-  return sum(times) / len(times)
-
-def create_synthetic_kernel(n_ops: int) -> bytes:
-  """Create a synthetic kernel with n_ops vector operations."""
-  instructions = []
-  # VOP2 instructions: v_add_f32, v_mul_f32, v_max_f32, v_min_f32
-  ops = [
-    (0b0000011 << 25) | (1 << 17) | (0 << 9) | 256,  # v_add_f32 v0, v0, v1
-    (0b0001000 << 25) | (1 << 17) | (0 << 9) | 256,  # v_mul_f32 v0, v0, v1
-    (0b0010000 << 25) | (1 << 17) | (0 << 9) | 256,  # v_max_f32 v0, v0, v1
-    (0b0001111 << 25) | (1 << 17) | (0 << 9) | 256,  # v_min_f32 v0, v0, v1
-  ]
-  for i in range(n_ops):
-    instructions.append(ops[i % len(ops)])
-  # S_ENDPGM
-  instructions.append((0b101111111 << 23) | (48 << 16) | 0)
-  return b''.join(struct.pack('<I', inst) for inst in instructions)
-
-def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes]] | None:
-  """Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data)."""
-  try:
-    from tinygrad import Tensor
-    from tinygrad.runtime.support.elf import elf_loader
-    import numpy as np
-    np.random.seed(42)
-
-    ops = {
-      "add": lambda: Tensor.empty(1024) + Tensor.empty(1024),
-      "mul": lambda: Tensor.empty(1024) * Tensor.empty(1024),
-      "matmul_small": lambda: Tensor.empty(16, 16) @ Tensor.empty(16, 16),
-      "matmul_medium": lambda: Tensor.empty(64, 64) @ Tensor.empty(64, 64),
-      "reduce_sum": lambda: Tensor.empty(4096).sum(),
-      "reduce_max": lambda: Tensor.empty(4096).max(),
-      "softmax": lambda: Tensor.empty(256).softmax(),
-      "layernorm": lambda: Tensor.empty(32, 64).layernorm(),
-      "conv2d": lambda: Tensor.empty(1, 4, 16, 16).conv2d(Tensor.empty(4, 4, 3, 3)),
-      "gelu": lambda: Tensor.empty(1024).gelu(),
-      "exp": lambda: Tensor.empty(1024).exp(),
-      "sin": lambda: Tensor.empty(1024).sin(),
-    }
-
-    if op_name not in ops: return None
-    out = ops[op_name]()
-    sched = out.schedule()
-
-    for ei in sched:
-      lowered = ei.lower()
-      if ei.ast.op.name == 'SINK' and lowered.prg and lowered.prg.p.lib:
-        lib = bytes(lowered.prg.p.lib)
-        _, sections, _ = elf_loader(lib)
-        for sec in sections:
-          if sec.name == '.text':
-            buf_sizes = [b.nbytes for b in lowered.bufs]
-            # Get initial data from numpy arrays if available
-            buf_data = {}
-            for i, buf in enumerate(lowered.bufs):
-              if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
-                try: buf_data[i] = bytes(buf.base._buf)
-                except: pass
-            return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data)
-    return None
-  except Exception as e:
-    print(f"  Error getting kernel: {e}")
-    return None
-
-def profile_python_emu(kernel: bytes, global_size, local_size, args_ptr, n_runs: int = 1):
-  """Profile the Python emulator to find bottlenecks."""
-  gx, gy, gz = global_size
-  lx, ly, lz = local_size
-  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-  lib_ptr = ctypes.addressof(kernel_buf)
-
-  pr = cProfile.Profile()
-  pr.enable()
-  for _ in range(n_runs):
-    python_run_asm(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
-  pr.disable()
-
-  s = io.StringIO()
-  ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
-  ps.print_stats(20)
-  return s.getvalue()
-
-def measure_step_rate(kernel: bytes, n_steps: int = 10000) -> float:
-  """Measure raw step_wave() performance (steps per second)."""
-  program = decode_program(kernel)
-  if not program: return 0.0
-
-  st = WaveState()
-  st.exec_mask = 0xffffffff
-  lds = bytearray(65536)
-  n_lanes = 32
-
-  # Reset PC for each measurement
-  start = time.perf_counter()
-  for _ in range(n_steps):
-    st.pc = 0
-    while st.pc in program:
-      result = step_wave(program, st, lds, n_lanes)
-      if result == -1: break
-  elapsed = time.perf_counter() - start
-  return n_steps / elapsed if elapsed > 0 else 0
-
-# Test configurations
-SYNTHETIC_TESTS = [
-  ("synthetic_10ops", 10, (1, 1, 1), (32, 1, 1)),
-  ("synthetic_100ops", 100, (1, 1, 1), (32, 1, 1)),
-  ("synthetic_500ops", 500, (1, 1, 1), (32, 1, 1)),
-  ("synthetic_100ops_4wg", 100, (4, 1, 1), (32, 1, 1)),
-  ("synthetic_100ops_16wg", 100, (16, 1, 1), (32, 1, 1)),
-]
-
-TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
-
-def main():
-  import argparse
-  parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
-  parser.add_argument("--profile", action="store_true", help="Profile Python emulator")
-  parser.add_argument("--synthetic-only", action="store_true", help="Only run synthetic tests")
-  parser.add_argument("--tinygrad-only", action="store_true", help="Only run tinygrad tests")
-  parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
-  args = parser.parse_args()
-
-  rust_remu = get_rust_remu()
-  if rust_remu is None:
-    print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
-    print("Running Python-only benchmarks...\n")
-
-  print("=" * 90)
-  print("RDNA3 Emulator Benchmark: Python vs Rust")
-  print("=" * 90)
-
-  results = []
-
-  # Synthetic workloads
-  if not args.tinygrad_only:
-    print("\n[SYNTHETIC WORKLOADS]")
-    print("-" * 90)
-
-    for name, n_ops, global_size, local_size in SYNTHETIC_TESTS:
-      kernel = create_synthetic_kernel(n_ops)
-      n_insts = count_instructions(kernel)
-      n_workgroups = global_size[0] * global_size[1] * global_size[2]
-      n_threads = local_size[0] * local_size[1] * local_size[2]
-      total_work = n_insts * n_workgroups * n_threads
-
-      print(f"\n{name}: {n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
-
-      buf_sizes = [4096]
-      buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes)
-      set_valid_mem_ranges(ranges)
-
-      # Benchmark
-      py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
-      rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
-
-      if py_time:
-        py_rate = total_work / py_time / 1e6
-        print(f"  Python: {py_time*1000:8.3f} ms  ({py_rate:7.2f} M ops/s)")
-      if rust_time:
-        rust_rate = total_work / rust_time / 1e6
-        speedup = py_time / rust_time if py_time else 0
-        print(f"  Rust:   {rust_time*1000:8.3f} ms  ({rust_rate:7.2f} M ops/s)  [{speedup:.1f}x faster]")
-
-      results.append(("synthetic", name, n_insts, n_workgroups, py_time, rust_time))
-
-  # Tinygrad kernels
-  if not args.synthetic_only:
-    print("\n[TINYGRAD KERNELS]")
-    print("-" * 90)
-
-    for op_name in TINYGRAD_TESTS:
-      print(f"\n{op_name}:", end=" ", flush=True)
-      kernel_info = get_tinygrad_kernel(op_name)
-      if kernel_info is None:
-        print("failed to compile")
-        continue
-
-      kernel, global_size, local_size, buf_sizes, buf_data = kernel_info
-      n_insts = count_instructions(kernel)
-      n_workgroups = global_size[0] * global_size[1] * global_size[2]
-      n_threads = local_size[0] * local_size[1] * local_size[2]
-      total_work = n_insts * n_workgroups * n_threads
-
-      print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
-
-      buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
-      set_valid_mem_ranges(ranges)
-
-      py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
-      rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
-
-      if py_time:
-        py_rate = total_work / py_time / 1e6
-        print(f"  Python: {py_time*1000:8.3f} ms  ({py_rate:7.2f} M ops/s)")
-      if rust_time:
-        rust_rate = total_work / rust_time / 1e6
-        speedup = py_time / rust_time if py_time else 0
-        print(f"  Rust:   {rust_time*1000:8.3f} ms  ({rust_rate:7.2f} M ops/s)  [{speedup:.1f}x faster]")
-
-      results.append(("tinygrad", op_name, n_insts, n_workgroups, py_time, rust_time))
-
-      # Optional profiling
-      if args.profile and py_time:
-        print("\n  [PROFILE - Top 10 functions]")
-        profile_output = profile_python_emu(kernel, global_size, local_size, args_ptr)
-        for line in profile_output.split('\n')[5:15]:
-          if line.strip(): print(f"    {line}")
-
-  # Summary table
-  print("\n" + "=" * 90)
-  print("SUMMARY")
-  print("=" * 90)
-  print(f"{'Type':<10} {'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
-  print("-" * 90)
-
-  for test_type, name, n_insts, n_wgs, py_time, rust_time in results:
-    py_ms = f"{py_time*1000:.3f}" if py_time else "error"
-    if rust_time:
-      rust_ms = f"{rust_time*1000:.3f}"
-      speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
-    else:
-      rust_ms, speedup = "N/A", "N/A"
-    print(f"{test_type:<10} {name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
-
-
-
-if __name__ == "__main__":
-  main()

extra/assembly/rdna3/test/test_compare_emulators.py

@@ -1,466 +0,0 @@
-# Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
-import unittest, ctypes, os
-from dataclasses import dataclass
-from pathlib import Path
-
-# Set environment before any tinygrad imports to use MOCKGPU
-# This allows generating AMD GPU kernels without requiring real hardware
-os.environ["AMD"] = "1"
-os.environ["MOCKGPU"] = "1"
-
-from extra.assembly.rdna3.emu import WaveState, decode_program, step_wave, WAVE_SIZE
-
-REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
-
-@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
-
-@dataclass
-class StateSnapshot:
-  pc: int
-  scc: int
-  vcc: int
-  exec_mask: int
-  sgpr: list[int]
-  vgpr: list[list[int]]
-
-  def diff(self, other: 'StateSnapshot', n_lanes: int) -> list[str]:
-    """Return list of differences between two states."""
-    diffs = []
-    if self.pc != other.pc: diffs.append(f"pc: {self.pc} vs {other.pc}")
-    if self.scc != other.scc: diffs.append(f"scc: {self.scc} vs {other.scc}")
-    if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x} vs 0x{other.vcc:08x}")
-    if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x} vs 0x{other.exec_mask:08x}")
-    for i, (a, b) in enumerate(zip(self.sgpr, other.sgpr)):
-      # Skip VCC_LO/HI (106/107) and EXEC_LO/HI (126/127) as they alias vcc/exec_mask which are compared separately
-      if i in (106, 107, 126, 127): continue
-      if a != b: diffs.append(f"sgpr[{i}]: 0x{a:08x} vs 0x{b:08x}")
-    for lane in range(n_lanes):
-      for i, (a, b) in enumerate(zip(self.vgpr[lane], other.vgpr[lane])):
-        if a != b: diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x} vs 0x{b:08x}")
-    return diffs
-
-class CStateSnapshot(ctypes.Structure):
-  _fields_ = [("pc", ctypes.c_uint32), ("scc", ctypes.c_uint32), ("vcc", ctypes.c_uint32), ("exec_mask", ctypes.c_uint32),
-              ("sgpr", ctypes.c_uint32 * 128), ("vgpr", (ctypes.c_uint32 * 256) * 32)]
-
-  def to_snapshot(self) -> StateSnapshot:
-    return StateSnapshot(pc=self.pc, scc=self.scc, vcc=self.vcc, exec_mask=self.exec_mask,
-                         sgpr=list(self.sgpr), vgpr=[list(self.vgpr[i]) for i in range(32)])
-
-class RustEmulator:
-  def __init__(self):
-    self.lib = ctypes.CDLL(str(REMU_PATH))
-    self.lib.wave_create.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_create.restype = ctypes.c_void_p
-    self.lib.wave_step.argtypes = [ctypes.c_void_p]
-    self.lib.wave_step.restype = ctypes.c_int32
-    self.lib.wave_get_snapshot.argtypes = [ctypes.c_void_p, ctypes.POINTER(CStateSnapshot)]
-    self.lib.wave_set_sgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_set_vgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_init_lds.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
-    self.lib.wave_free.argtypes = [ctypes.c_void_p]
-    self.ctx = None
-
-  def create(self, kernel: bytes, n_lanes: int):
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    self.ctx = self.lib.wave_create(ctypes.addressof(kernel_buf), len(kernel), n_lanes)
-    self._kernel_buf = kernel_buf
-
-  def step(self) -> int: return self.lib.wave_step(self.ctx)
-  def set_sgpr(self, idx: int, val: int): self.lib.wave_set_sgpr(self.ctx, idx, val)
-  def set_vgpr(self, lane: int, idx: int, val: int): self.lib.wave_set_vgpr(self.ctx, lane, idx, val)
-  def init_lds(self, size: int): self.lib.wave_init_lds(self.ctx, size)
-
-  def get_snapshot(self) -> StateSnapshot:
-    snap = CStateSnapshot()
-    self.lib.wave_get_snapshot(self.ctx, ctypes.byref(snap))
-    return snap.to_snapshot()
-
-  def free(self):
-    if self.ctx: self.lib.wave_free(self.ctx); self.ctx = None
-
-class PythonEmulator:
-  def __init__(self):
-    self.state: WaveState | None = None
-    self.program: dict | None = None
-    self.lds: bytearray | None = None
-    self.n_lanes = 0
-
-  def create(self, kernel: bytes, n_lanes: int):
-    self.program = decode_program(kernel)
-    self.state = WaveState()
-    self.state.exec_mask = (1 << n_lanes) - 1
-    self.lds = bytearray(65536)
-    self.n_lanes = n_lanes
-
-  def step(self) -> int:
-    assert self.program is not None and self.state is not None and self.lds is not None
-    return step_wave(self.program, self.state, self.lds, self.n_lanes)
-  def set_sgpr(self, idx: int, val: int):
-    assert self.state is not None
-    self.state.sgpr[idx] = val & 0xffffffff
-  def set_vgpr(self, lane: int, idx: int, val: int):
-    assert self.state is not None
-    self.state.vgpr[lane][idx] = val & 0xffffffff
-
-  def get_snapshot(self) -> StateSnapshot:
-    assert self.state is not None
-    return StateSnapshot(pc=self.state.pc, scc=self.state.scc, vcc=self.state.vcc & 0xffffffff,
-                         exec_mask=self.state.exec_mask & 0xffffffff, sgpr=list(self.state.sgpr),
-                         vgpr=[list(self.state.vgpr[i]) for i in range(WAVE_SIZE)])
-
-def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
-                      program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
-                      max_workgroups: int = 64) -> tuple[bool, str, int]:
-  """Run a single kernel through both emulators. Returns (success, message, total_steps)."""
-  gx, gy, gz = global_size
-  total_steps = 0
-  wg_count = 0
-
-  for gidz in range(gz):
-    for gidy in range(gy):
-      for gidx in range(gx):
-        if wg_count >= max_workgroups: return True, f"Completed {wg_count} workgroups (limit reached)", total_steps
-        wg_count += 1
-        rust = RustEmulator()
-        python = PythonEmulator()
-        rust.create(kernel, n_lanes)
-        python.create(kernel, n_lanes)
-
-        # Initialize LDS (64KB, standard size for AMD GPUs)
-        rust.init_lds(65536)
-
-        for emu in (rust, python):
-          emu.set_sgpr(0, args_ptr & 0xffffffff)
-          emu.set_sgpr(1, (args_ptr >> 32) & 0xffffffff)
-          emu.set_sgpr(13, gidx)
-          emu.set_sgpr(14, gidy)
-          emu.set_sgpr(15, gidz)
-
-        step = 0
-        trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
-        try:
-          while step < max_steps:
-            rust_before = rust.get_snapshot()
-            python_before = python.get_snapshot()
-
-            inst = program.get(python_before.pc)
-            inst_str = inst.disasm() if inst else f"unknown at PC={python_before.pc}"
-            trace.append((step, python_before.pc, inst_str, rust_before, python_before))
-            if len(trace) > trace_len: trace.pop(0)
-
-            if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
-
-            diffs = rust_before.diff(python_before, n_lanes)
-            if diffs:
-              trace_lines = []
-              for s, pc, d, rb, pb in trace[:-1]:
-                trace_lines.append(f"    step {s}: PC={pc:3d} {d}")
-                if trace.index((s, pc, d, rb, pb)) < len(trace) - 2:
-                  next_rb, next_pb = trace[trace.index((s, pc, d, rb, pb)) + 1][3:5]
-                  inst_diffs = rb.diff(next_rb, n_lanes)
-                  if inst_diffs: trace_lines.append(f"             rust changes: {', '.join(inst_diffs[:3])}")
-              trace_str = "\n".join(trace_lines)
-              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ:\n  " + "\n  ".join(diffs[:10]) + f"\n  Recent instructions:\n{trace_str}", total_steps
-
-            rust_result = rust.step()
-            python_result = python.step()
-
-            if rust_result != python_result:
-              trace_str = "\n".join(f"    step {s}: PC={pc:3d} {d}" for s, pc, d, _, _ in trace)
-              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: rust={rust_result}, python={python_result}, inst={inst_str}\n  Recent instructions:\n{trace_str}", total_steps
-
-            if rust_result == -1:
-              total_steps += step + 1
-              break
-            if rust_result == 1:
-              total_steps += step + 1
-              break
-            if rust_result < 0 and rust_result != -2:
-              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: error code {rust_result}", total_steps
-
-            step += 1
-          else:
-            return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Max steps ({max_steps}) reached", total_steps
-        finally:
-          rust.free()
-
-  return True, f"Completed {gx*gy*gz} workgroups", total_steps
-
-def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int, int], max_steps: int = 1000,
-                                    debug: bool = False, trace_len: int = 10, buf_data: dict[int, bytes] | None = None) -> tuple[bool, str]:
-  """Run all kernels through both emulators with shared buffer pool."""
-  from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
-  if buf_data is None: buf_data = {}
-
-  # Allocate shared buffer pool with padding for over-reads (GPU loads up to 16 bytes at once)
-  buf_id_to_ptr: dict[int, int] = {}
-  buffers = []
-  for buf_id, size in buf_pool.items():
-    padded_size = ((size + 15) // 16) * 16 + 16  # round up to 16 bytes + extra padding
-    # Initialize with data from COPY if available
-    init_data = buf_data.get(buf_id, b'\x00' * padded_size)
-    init_list = list(init_data) + [0] * (padded_size - len(init_data))
-    buf = (ctypes.c_uint8 * padded_size)(*init_list[:padded_size])
-    buffers.append((buf, padded_size))
-    buf_id_to_ptr[buf_id] = ctypes.addressof(buf)
-
-  # Set up valid memory ranges
-  ranges = {(ctypes.addressof(b), size) for b, size in buffers}
-
-  total_steps = 0
-  for ki, kernel in enumerate(kernels):
-    # Create args array for this kernel's buffers
-    args = (ctypes.c_uint64 * len(kernel.buf_idxs))(*[buf_id_to_ptr[bid] for bid in kernel.buf_idxs])
-    args_ptr = ctypes.addressof(args)
-
-    # Update valid ranges to include this args array
-    kernel_ranges = ranges | {(args_ptr, ctypes.sizeof(args))}
-    set_valid_mem_ranges(kernel_ranges)
-
-    program = decode_program(kernel.code)
-    n_lanes = kernel.local_size[0] * kernel.local_size[1] * kernel.local_size[2]
-
-    ok, msg, steps = run_single_kernel(
-      kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
-      program, max_steps, debug, trace_len, ki
-    )
-    total_steps += steps
-    if not ok:
-      return False, msg
-
-  return True, f"Completed {len(kernels)} kernels, {total_steps} total steps"
-
-def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list, max_steps: int = 1000, debug: bool = False,
-                                   global_size: tuple[int, int, int] = (1, 1, 1), trace_len: int = 10) -> tuple[bool, str]:
-  """Run both emulators with memory set up for tinygrad kernels, executing all workgroups. Legacy wrapper."""
-  from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
-
-  # Allocate buffers
-  buffers = []
-  for size in buf_sizes:
-    buf = (ctypes.c_uint8 * size)(*[0] * size)
-    buffers.append(buf)
-
-  # Create args array with buffer pointers
-  args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
-  args_ptr = ctypes.addressof(args)
-
-  # Set up valid memory ranges for Python emulator
-  ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
-  ranges.add((args_ptr, ctypes.sizeof(args)))
-  set_valid_mem_ranges(ranges)
-
-  program = decode_program(kernel)
-  ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, program, max_steps, debug, trace_len)
-  return ok, msg
-
-def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int], dict[int, bytes]]:
-  """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
-  from tinygrad import Tensor
-  from tinygrad.runtime.support.elf import elf_loader
-
-  out = op_fn(Tensor)
-  sched = out.schedule()
-  kernels = []
-  buf_pool: dict[int, int] = {}  # buffer id -> size
-  buf_data: dict[int, bytes] = {}  # buffer id -> initial data from COPY
-
-  for ei in sched:
-    lowered = ei.lower()
-    if ei.ast.op.name == 'COPY':
-      # Handle COPY: extract source data to initialize destination buffer
-      if len(lowered.bufs) >= 2:
-        dst_buf, src_buf = lowered.bufs[0], lowered.bufs[1]
-        dst_id = id(dst_buf)
-        if dst_id not in buf_pool:
-          buf_pool[dst_id] = dst_buf.nbytes
-        # Get source data if it's from numpy/CPU
-        if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'):
-          src_data = bytes(src_buf.base._buf)
-          buf_data[dst_id] = src_data
-    elif ei.ast.op.name == 'SINK':
-      if lowered.prg and lowered.prg.p.lib:
-        lib = bytes(lowered.prg.p.lib)
-        _, sections, _ = elf_loader(lib)
-        for sec in sections:
-          if sec.name == '.text':
-            buf_idxs = []
-            buf_sizes = []
-            for b in lowered.bufs:
-              buf_id = id(b)
-              if buf_id not in buf_pool:
-                buf_pool[buf_id] = b.nbytes
-              buf_idxs.append(buf_id)
-              buf_sizes.append(b.nbytes)
-            kernels.append(KernelInfo(
-              code=bytes(sec.content),
-              global_size=tuple(lowered.prg.p.global_size),
-              local_size=tuple(lowered.prg.p.local_size),
-              buf_idxs=buf_idxs,
-              buf_sizes=buf_sizes
-            ))
-  if not kernels: raise RuntimeError("No kernel found")
-  return kernels, buf_pool, buf_data
-
-def get_kernel_from_tinygrad(op_fn) -> tuple[bytes, tuple[int, int, int], tuple[int, int, int], list]:
-  """Compile a tinygrad operation and extract the last (main) kernel binary. Legacy wrapper."""
-  kernels, _, _ = get_kernels_from_tinygrad(op_fn)
-  k = kernels[-1]
-  return k.code, k.global_size, k.local_size, k.buf_sizes
-
-@unittest.skipUnless(REMU_PATH.exists(), "libremu.so not found")
-class TestTinygradKernels(unittest.TestCase):
-  """Compare emulators on real tinygrad-compiled kernels."""
-
-  def _test_kernel(self, op_fn, max_steps=10000):
-    kernels, buf_pool, buf_data = get_kernels_from_tinygrad(op_fn)
-    ok, msg = compare_emulators_multi_kernel(kernels, buf_pool, max_steps=max_steps, buf_data=buf_data)
-    self.assertTrue(ok, msg)
-
-  # Basic unary ops
-  def test_neg(self): self._test_kernel(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
-  def test_relu(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
-  def test_exp(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).exp())
-  def test_log(self): self._test_kernel(lambda T: T([1.0, 2.0, 3.0]).log())
-  def test_sin(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).sin())
-  def test_sqrt(self): self._test_kernel(lambda T: T([1.0, 4.0, 9.0]).sqrt())
-  def test_recip(self): self._test_kernel(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
-
-  # Binary ops
-  def test_add(self): self._test_kernel(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
-  def test_sub(self): self._test_kernel(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
-  def test_mul(self): self._test_kernel(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
-  def test_div(self): self._test_kernel(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
-  def test_max_binary(self): self._test_kernel(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
-
-  # Reductions
-  def test_sum_reduce(self): self._test_kernel(lambda T: T.empty(64).sum())
-  def test_max_reduce(self): self._test_kernel(lambda T: T.empty(64).max())
-  def test_mean_reduce(self): self._test_kernel(lambda T: T.empty(32).mean())
-
-  # Matmul - various sizes
-  def test_gemm_4x4(self): self._test_kernel(lambda T: T.empty(4, 4) @ T.empty(4, 4), max_steps=100000)
-  def test_gemm_8x8(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=200000)
-  @unittest.skip("too slow")
-  def test_gemm_16x16(self): self._test_kernel(lambda T: T.empty(16, 16) @ T.empty(16, 16), max_steps=500000)
-  def test_gemv(self): self._test_kernel(lambda T: T.empty(1, 16) @ T.empty(16, 16), max_steps=100000)
-
-  # Complex ops
-  def test_softmax(self): self._test_kernel(lambda T: T.empty(16).softmax())
-  def test_layernorm(self): self._test_kernel(lambda T: T.empty(8, 8).layernorm())
-
-  # Memory patterns
-  def test_contiguous(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
-  def test_reshape(self): self._test_kernel(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
-  def test_expand(self): self._test_kernel(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
-
-  # Cast ops
-  def test_cast_int(self): self._test_kernel(lambda T: T.empty(16).int().float())
-  def test_cast_half(self): self._test_kernel(lambda T: T.empty(16).half().float())
-
-  # Min/max (uses comparison internally)
-  def test_min_binary(self): self._test_kernel(lambda T: T([1.0, 5.0, 3.0]).minimum(T([3.0, 2.0, 4.0])))
-
-  # Comparison ops (test VOPC instructions) - use 32+ elements to force vector instructions
-  def test_cmp_lt(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
-  def test_cmp_eq(self): self._test_kernel(lambda T: (T.empty(64) == T.empty(64)).where(T.empty(64), T.empty(64)))
-  def test_where(self): self._test_kernel(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
-
-  # Bitwise ops
-  def test_bitwise_and(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0xFF]).int() & T([0x0F, 0x0F, 0x00]).int())
-  def test_bitwise_or(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0x00]).int() | T([0x0F, 0x0F, 0xFF]).int())
-  def test_bitwise_xor(self): self._test_kernel(lambda T: T([0xFF, 0x0F, 0xF0]).int() ^ T([0x0F, 0xF0, 0xF0]).int())
-
-  # Integer ops - use 32+ elements to force vector instructions
-  def test_int_add(self): self._test_kernel(lambda T: (T.empty(64).int() + T.empty(64).int()).float())
-  def test_int_mul(self): self._test_kernel(lambda T: (T.empty(64).int() * T.empty(64).int()).float())
-  def test_int_mod(self): self._test_kernel(lambda T: (T.empty(64).int().abs() % (T.empty(64).int().abs() + 1)).float())
-
-  # More math ops - use 32+ elements to force vector instructions
-  def test_abs(self): self._test_kernel(lambda T: T.empty(64).abs())
-  def test_floor(self): self._test_kernel(lambda T: T.empty(64).floor())
-  def test_ceil(self): self._test_kernel(lambda T: T.empty(64).ceil())
-  def test_trunc(self): self._test_kernel(lambda T: T.empty(64).trunc())
-
-  # Fused ops
-  def test_fma(self): self._test_kernel(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
-
-  # Argmax/argmin (tests different reduction pattern) - use 32+ elements to force vector instructions
-  def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
-  def test_argmin(self): self._test_kernel(lambda T: T.empty(64).argmin())
-
-  # Exact value tests - use 32+ elements to force vector instructions (small tensors use scalar ops which Rust emu doesn't fully support)
-  def test_abs_exact(self): self._test_kernel(lambda T: T([-1., 0., 1.]*11).abs())  # 33 elements
-  def test_neg_exact(self): self._test_kernel(lambda T: -T([-1., 0., 1.]*11))
-  def test_log_special(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).log())
-  def test_exp_exact(self): self._test_kernel(lambda T: T([0., 1., -1.]*11).exp())
-  def test_reciprocal_exact(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).reciprocal())
-
-  # Integer division and mod - use 32+ elements
-  def test_int_div(self): self._test_kernel(lambda T: (T([10, 20, 30]*11).int() // T([3, 4, 5]*11).int()).float())
-  def test_int_neg(self): self._test_kernel(lambda T: (-T([1, -2, 3]*11).int()).float())
-
-  # Mixed precision - use 32+ elements
-  def test_half_add(self): self._test_kernel(lambda T: (T([1., 2.]*16).half() + T([3., 4.]*16).half()).float())
-  def test_half_mul(self): self._test_kernel(lambda T: (T([2., 3.]*16).half() * T([4., 5.]*16).half()).float())
-
-  # Matrix ops - patterns from test_ops.py failures
-  def test_cat(self): self._test_kernel(lambda T: T.empty(32, 64).cat(T.empty(32, 64), dim=1))
-  def test_gather(self): self._test_kernel(lambda T: T.empty(64).gather(0, T.arange(32).int()))
-
-  # Tests from test_ops.py that are failing
-  def test_permute(self): self._test_kernel(lambda T: T.empty(3, 4, 5, 6).permute((3, 2, 1, 0)).contiguous())
-  def test_cat_large(self): self._test_kernel(lambda T: T.empty(45, 65, 9).cat(T.empty(45, 65, 9), T.empty(45, 65, 9), dim=1))
-  def test_gather_small(self): self._test_kernel(lambda T: T.empty(10).gather(0, T.arange(5).int()))
-  @unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
-  def test_cross_entropy(self): self._test_kernel(lambda T: T.randn(32, 10).softmax().log().sum())
-  def test_cross_entropy_class(self):
-    import numpy as np
-    np.random.seed(0)
-    classes = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
-    x_np = np.random.randn(32, 10).astype(np.float32)
-    self._test_kernel(lambda T: (T(x_np.tolist()).reshape(32,10) + 0).cross_entropy((T(classes).int().reshape(32) + 0)))
-
-  # Regression tests for BFE operations with width=0 (walrus operator bug)
-  def test_topk(self): self._test_kernel(lambda T: T.empty(64).topk(3)[0])
-  def test_interpolate_uint8(self): self._test_kernel(lambda T: T.empty(2,3,64,64).relu().cast('uint8').interpolate((10,10), mode="linear"))
-
-  # Regression test for 64-bit comparison (V_CMP_GT_I64, V_CMP_LT_U64, etc.) with rsrc64
-  def test_index_int64(self):
-    from tinygrad import dtypes
-    self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
-
-  @unittest.skip("only works with mock GPU")
-  def test_index_int64_2d(self):
-    from tinygrad import dtypes
-    # Tests 64-bit compare with inline constants (comparing against 0)
-    self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), T.arange(4).cast(dtypes.int64)])
-
-  # Pooling operations - regression test for VCC wave32 mode (S_CBRANCH_VCCZ should only check VCC_LO)
-  def test_avg_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4), stride=2))
-  @unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
-  def test_avg_pool3d(self):
-    import numpy as np
-    np.random.seed(0)
-    self._test_kernel(lambda T: T(np.random.randn(1, 1, 16, 16, 16).astype(np.float32).tolist()).avg_pool2d(kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False))
-  def test_max_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4), stride=2))
-
-  # Convolution operations - multi-kernel tests
-  def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv2d(T.empty(4, 4, 3, 3)), max_steps=100000)
-  def test_conv_transpose2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv_transpose2d(T.empty(4, 4, 3, 3)), max_steps=200000)
-  @unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
-  def test_conv_transpose3d(self):
-    import numpy as np
-    np.random.seed(0)
-    self._test_kernel(lambda T: T(np.random.randn(2, 4, 9, 9, 9).astype(np.float32).tolist()).conv_transpose2d(
-      T(np.random.randn(4, 4, 3, 3, 3).astype(np.float32).tolist())), max_steps=500000)
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_emu.py

@@ -1,845 +0,0 @@
-# Unit tests for RDNA3 Python emulator
-import unittest
-import ctypes
-import struct
-import math
-from extra.assembly.rdna3.emu import (
-  WaveState, decode_program, exec_wave, exec_workgroup, run_asm,
-  i32, f32, sext, WAVE_SIZE, set_valid_mem_ranges
-)
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.lib import RawImm
-
-def run_kernel(kernel: bytes, n_threads: int = 1, n_outputs: int = 1) -> list[int]:
-  """Helper to run a kernel and return output values."""
-  output = (ctypes.c_uint32 * (n_threads * n_outputs))(*[0xdead] * (n_threads * n_outputs))
-  output_ptr = ctypes.addressof(output)
-  args = (ctypes.c_uint64 * 1)(output_ptr)
-  args_ptr = ctypes.addressof(args)
-  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-  kernel_ptr = ctypes.addressof(kernel_buf)
-  # Register valid memory ranges for bounds checking
-  set_valid_mem_ranges({
-    (output_ptr, ctypes.sizeof(output)),
-    (args_ptr, ctypes.sizeof(args)),
-    (kernel_ptr, len(kernel)),
-  })
-  result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
-  assert result == 0, f"run_asm failed with {result}"
-  return [output[i] for i in range(n_threads * n_outputs)]
-
-def make_store_kernel(setup_instrs: list, store_vreg: int = 1) -> bytes:
-  """Create a kernel that runs setup instructions then stores v[store_vreg] to output[tid]."""
-  kernel = b''
-  # Load output pointer
-  kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
-  kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-  # Run setup instructions
-  for instr in setup_instrs:
-    kernel += instr.to_bytes()
-  # Compute offset: v3 = tid * 4
-  kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
-  # Store result
-  kernel += global_store_b32(addr=v[3], data=v[store_vreg], saddr=s[2]).to_bytes()
-  kernel += s_endpgm().to_bytes()
-  return kernel
-
-class TestScalarOps(unittest.TestCase):
-  def test_s_mov_b32(self):
-    state = WaveState()
-    kernel = s_mov_b32(s[5], 42).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[5], 42)
-
-  def test_s_add_u32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 100, 50
-    kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 150)
-    self.assertEqual(state.scc, 0)  # no carry
-
-  def test_s_add_u32_carry(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 0xffffffff, 1
-    kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 0)
-    self.assertEqual(state.scc, 1)  # carry
-
-  def test_s_sub_u32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 100, 30
-    kernel = s_sub_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 70)
-    self.assertEqual(state.scc, 0)  # no borrow
-
-  def test_s_and_b32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 0xff00, 0x0ff0
-    kernel = s_and_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 0x0f00)
-
-  def test_s_or_b32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 0xff00, 0x00ff
-    kernel = s_or_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 0xffff)
-
-  def test_s_lshl_b32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 1, 4
-    kernel = s_lshl_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 16)
-
-  def test_s_lshr_b32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 256, 4
-    kernel = s_lshr_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 16)
-
-  def test_s_mul_i32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 7, 6
-    kernel = s_mul_i32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[2], 42)
-
-  def test_s_cmp_eq_u32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 42, 42
-    kernel = s_cmp_eq_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.scc, 1)
-
-  def test_s_cmp_lg_u32(self):
-    state = WaveState()
-    state.sgpr[0], state.sgpr[1] = 42, 43
-    kernel = s_cmp_lg_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.scc, 1)
-
-class TestVectorOps(unittest.TestCase):
-  def test_v_mov_b32(self):
-    kernel = make_store_kernel([v_mov_b32_e32(v[1], 42)])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [42])
-
-  def test_v_add_nc_u32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 10),
-      v_mov_b32_e32(v[2], 32),
-      v_add_nc_u32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [42])
-
-  def test_v_sub_nc_u32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 50),
-      v_mov_b32_e32(v[2], 8),
-      v_sub_nc_u32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [42])
-
-  def test_v_mul_lo_u32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 6),
-      v_mov_b32_e32(v[2], 7),
-      v_mul_lo_u32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [42])
-
-  def test_v_and_b32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 0xff0f),
-      v_mov_b32_e32(v[2], 0x0fff),
-      v_and_b32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [0x0f0f])
-
-  def test_v_or_b32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 0xf000),
-      v_mov_b32_e32(v[2], 0x000f),
-      v_or_b32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [0xf00f])
-
-  def test_v_lshlrev_b32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 1),
-      v_lshlrev_b32_e32(v[1], 5, v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [32])
-
-  def test_v_lshrrev_b32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 128),
-      v_lshrrev_b32_e32(v[1], 3, v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out, [16])
-
-  def test_v_add_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(1.5)),
-      v_mov_b32_e32(v[2], i32(2.5)),
-      v_add_f32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 4.0)
-
-  def test_v_mul_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(3.0)),
-      v_mov_b32_e32(v[2], i32(4.0)),
-      v_mul_f32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 12.0)
-
-  def test_v_max_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(3.0)),
-      v_mov_b32_e32(v[2], i32(5.0)),
-      v_max_f32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 5.0)
-
-  def test_v_min_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(3.0)),
-      v_mov_b32_e32(v[2], i32(5.0)),
-      v_min_f32_e32(v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 3.0)
-
-class TestThreading(unittest.TestCase):
-  def test_thread_id(self):
-    """Each thread should get its own thread ID in v0."""
-    kernel = make_store_kernel([v_mov_b32_e32(v[1], v[0])], store_vreg=1)
-    out = run_kernel(kernel, n_threads=4)
-    self.assertEqual(out, [0, 1, 2, 3])
-
-  def test_thread_local_ops(self):
-    """Each thread computes tid * 10."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[2], 10),
-      v_mul_lo_u32(v[1], v[0], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=4)
-    self.assertEqual(out, [0, 10, 20, 30])
-
-  def test_exec_mask(self):
-    """Test that exec mask controls which lanes execute."""
-    kernel = b''
-    kernel += s_load_b64(s[2:3], s[0:1], 0, soffset=NULL).to_bytes()
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    kernel += v_mov_b32_e32(v[1], 100).to_bytes()  # default value
-    kernel += s_mov_b32(EXEC_LO, 0b0101).to_bytes()  # only lanes 0 and 2
-    kernel += v_mov_b32_e32(v[1], 42).to_bytes()   # only for active lanes
-    kernel += s_mov_b32(EXEC_LO, 0xf).to_bytes()   # restore all lanes
-    kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
-    kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[2]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    out = run_kernel(kernel, n_threads=4)
-    self.assertEqual(out, [42, 100, 42, 100])
-
-class TestBranching(unittest.TestCase):
-  def test_s_branch(self):
-    """Test unconditional branch."""
-    state = WaveState()
-    kernel = b''
-    kernel += s_mov_b32(s[0], 1).to_bytes()
-    kernel += s_branch(1).to_bytes()  # skip next instruction
-    kernel += s_mov_b32(s[0], 2).to_bytes()  # should be skipped
-    kernel += s_mov_b32(s[1], 3).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[0], 1)  # not overwritten
-    self.assertEqual(state.sgpr[1], 3)
-
-  def test_s_cbranch_scc0(self):
-    """Test conditional branch on SCC=0."""
-    state = WaveState()
-    state.scc = 0
-    kernel = b''
-    kernel += s_mov_b32(s[0], 1).to_bytes()
-    kernel += s_cbranch_scc0(1).to_bytes()  # branch if scc=0
-    kernel += s_mov_b32(s[0], 2).to_bytes()  # should be skipped
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[0], 1)
-
-  def test_s_cbranch_scc1(self):
-    """Test conditional branch on SCC=1."""
-    state = WaveState()
-    state.scc = 1
-    kernel = b''
-    kernel += s_mov_b32(s[0], 1).to_bytes()
-    kernel += s_cbranch_scc1(1).to_bytes()  # branch if scc=1
-    kernel += s_mov_b32(s[0], 2).to_bytes()  # should be skipped
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.sgpr[0], 1)
-
-  def test_unknown_sopp_opcode(self):
-    """Regression test: unknown SOPP opcodes should be ignored, not crash."""
-    state = WaveState()
-    # Create a raw SOPP instruction with opcode 8 (undefined in our enum)
-    # SOPP format: bits[31:23] = 0b101111111, bits[22:16] = op, bits[15:0] = simm16
-    unknown_sopp = (0b101111111 << 23) | (8 << 16) | 0  # op=8, simm16=0
-    kernel = unknown_sopp.to_bytes(4, 'little') + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    # Should not raise an exception
-    exec_wave(prog, state, bytearray(65536), 1)
-
-class TestMemory(unittest.TestCase):
-  def test_global_load_store(self):
-    """Test global load followed by store."""
-    # Create input buffer
-    input_buf = (ctypes.c_uint32 * 4)(10, 20, 30, 40)
-    input_ptr = ctypes.addressof(input_buf)
-    output_buf = (ctypes.c_uint32 * 4)(*[0]*4)
-    output_ptr = ctypes.addressof(output_buf)
-    args = (ctypes.c_uint64 * 2)(output_ptr, input_ptr)
-    args_ptr = ctypes.addressof(args)
-
-    # Kernel: load from input[tid], add 1, store to output[tid]
-    kernel = b''
-    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()  # output ptr
-    kernel += s_load_b64(s[4:5], s[0:1], soffset=NULL, offset=8).to_bytes()  # input ptr
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    kernel += v_lshlrev_b32_e32(v[2], 2, v[0]).to_bytes()  # offset = tid * 4
-    kernel += global_load_b32(vdst=v[1], addr=v[2], saddr=s[4]).to_bytes()
-    kernel += s_waitcnt(vmcnt=0).to_bytes()
-    kernel += v_add_nc_u32_e32(v[1], 1, v[1]).to_bytes()  # add 1
-    kernel += global_store_b32(addr=v[2], data=v[1], saddr=s[2]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    kernel_ptr = ctypes.addressof(kernel_buf)
-    set_valid_mem_ranges({
-      (input_ptr, ctypes.sizeof(input_buf)),
-      (output_ptr, ctypes.sizeof(output_buf)),
-      (args_ptr, ctypes.sizeof(args)),
-      (kernel_ptr, len(kernel)),
-    })
-    result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, 4, 1, 1, args_ptr)
-    self.assertEqual(result, 0)
-    self.assertEqual([output_buf[i] for i in range(4)], [11, 21, 31, 41])
-
-class TestFloatOps(unittest.TestCase):
-  def test_v_rcp_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(4.0)),
-      v_rcp_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertAlmostEqual(f32(out[0]), 0.25, places=5)
-
-  def test_v_sqrt_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(16.0)),
-      v_sqrt_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertAlmostEqual(f32(out[0]), 4.0, places=5)
-
-  def test_v_floor_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(3.7)),
-      v_floor_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 3.0)
-
-  def test_v_ceil_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(3.2)),
-      v_ceil_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 4.0)
-
-  def test_v_cvt_f32_i32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 42),
-      v_cvt_f32_i32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 42.0)
-
-  def test_v_cvt_i32_f32(self):
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(42.9)),
-      v_cvt_i32_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0], 42)
-
-class TestVOP3(unittest.TestCase):
-  def test_v_fma_f32(self):
-    """Test fused multiply-add: a*b + c"""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(2.0)),
-      v_mov_b32_e32(v[2], i32(3.0)),
-      v_mov_b32_e32(v[4], i32(4.0)),
-      v_fma_f32(v[1], v[1], v[2], v[4]),  # 2*3+4 = 10
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 10.0)
-
-  def test_v_add3_u32(self):
-    """Test 3-operand add."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 10),
-      v_mov_b32_e32(v[2], 20),
-      v_mov_b32_e32(v[4], 12),
-      v_add3_u32(v[1], v[1], v[2], v[4]),  # 10+20+12 = 42
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0], 42)
-
-  def test_v_neg_modifier(self):
-    """Test VOP3 negation modifier."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(5.0)),
-      v_mov_b32_e32(v[2], i32(3.0)),
-      # v_add_f32 with neg on src1: 5 + (-3) = 2
-      v_add_f32_e64(v[1], v[1], v[2], neg=0b010),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 2.0)
-
-  def test_v_ldexp_f32(self):
-    """Regression test: V_LDEXP_F32 used by exp()."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(1.5)),
-      v_mov_b32_e32(v[2], 3),  # exponent
-      v_ldexp_f32(v[1], v[1], v[2]),  # 1.5 * 2^3 = 12.0
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(f32(out[0]), 12.0)
-
-  def test_v_xad_u32(self):
-    """Regression test: V_XAD_U32 (xor-add) used by random number generation."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 3),
-      v_mov_b32_e32(v[2], 4),
-      v_mov_b32_e32(v[4], 5),
-      v_xad_u32(v[1], v[1], v[2], v[4]),  # (3^4)+5 = 7+5 = 12
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0], 12)
-
-  def test_v_lshl_or_b32(self):
-    """Regression test: V_LSHL_OR_B32 operand order is (s0 << s1) | s2, not (s0 << s2) | s1."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 5),   # s0 = value to shift
-      v_mov_b32_e32(v[2], 2),   # s1 = shift amount
-      v_mov_b32_e32(v[4], 3),   # s2 = value to OR
-      v_lshl_or_b32(v[1], v[1], v[2], v[4]),  # (5 << 2) | 3 = 20 | 3 = 23
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0], 23)
-
-  def test_v_sqrt_f32_negative(self):
-    """Regression test: V_SQRT_F32 should return NaN for negative inputs, not 0."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(-1.0)),
-      v_sqrt_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertTrue(math.isnan(f32(out[0])))
-
-  def test_v_rsq_f32_negative(self):
-    """Regression test: V_RSQ_F32 should return NaN for negative inputs, not inf."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i32(-1.0)),
-      v_rsq_f32_e32(v[1], v[1]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertTrue(math.isnan(f32(out[0])))
-
-class TestVOPD(unittest.TestCase):
-  def test_vopd_add_nc_u32(self):
-    """Test VOPD V_DUAL_ADD_NC_U32."""
-    state = WaveState()
-    state.vgpr[0][1] = 100
-    state.vgpr[0][2] = 50
-    # vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
-    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
-                  opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.vgpr[0][3], 100)  # MOV result
-    self.assertEqual(state.vgpr[0][4], 150)  # 100 + 50
-
-  def test_vopd_lshlrev(self):
-    """Test VOPD V_DUAL_LSHLREV_B32."""
-    state = WaveState()
-    state.vgpr[0][1] = 0x10
-    state.vgpr[0][2] = 0
-    # vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
-    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
-                  opy=VOPDOp.V_DUAL_LSHLREV_B32, srcy0=4, vsrcy1=VGPR(1), vdsty=VGPR(4)).to_bytes()  # V4 = V1 << 4
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.vgpr[0][3], 0x10)  # MOV result
-    self.assertEqual(state.vgpr[0][4], 0x100)  # 0x10 << 4 = 0x100
-
-  def test_vopd_and(self):
-    """Test VOPD V_DUAL_AND_B32."""
-    state = WaveState()
-    state.vgpr[0][1] = 0xff
-    state.vgpr[0][2] = 0x0f
-    # vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
-    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
-                  opy=VOPDOp.V_DUAL_AND_B32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.vgpr[0][3], 0xff)
-    self.assertEqual(state.vgpr[0][4], 0x0f)  # 0xff & 0x0f = 0x0f
-
-  def test_vopd_parallel_read(self):
-    """Regression: VOPD must read all inputs before writing - Y op reads register that X op writes."""
-    state = WaveState()
-    state.vgpr[0][4] = 0
-    state.vgpr[0][7] = 5  # Y op reads v7 as vsrcy1, X op writes to v7
-    # X: MOV v7, v0 (v0=0, so v7 becomes 0)
-    # Y: ADD v6, v4, v7 (should use original v7=5, not the overwritten 0)
-    # vdsty_enc=3 with vdstx=7 (odd) -> vdsty = (3 << 1) | (7&1)^1 = 6 | 0 = 6
-    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[0], vsrcx1=VGPR(0), vdstx=VGPR(7),
-                  opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[4], vsrcy1=VGPR(7), vdsty=VGPR(6)).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.vgpr[0][7], 0)  # X op: v7 = v0 = 0
-    self.assertEqual(state.vgpr[0][6], 5)  # Y op: v6 = v4 + v7 = 0 + 5 (original v7)
-
-class TestDecoder(unittest.TestCase):
-  def test_vopd_literal_handling(self):
-    """Regression test: VOPD srcx0/srcy0 with literal (255) wasn't consuming the literal dword."""
-    state = WaveState()
-    # Create VOPD with srcx0=255 (literal), followed by literal value 0x12345678
-    vopd_bytes = VOPD(opx=8, srcx0=RawImm(255), vsrcx1=VGPR(0), vdstx=VGPR(1),  # MOV: V1 = literal
-                      opy=8, srcy0=RawImm(128), vsrcy1=VGPR(0), vdsty=VGPR(2)).to_bytes()  # MOV: V2 = 0
-    literal_bytes = (0x12345678).to_bytes(4, 'little')
-    kernel = vopd_bytes + literal_bytes + s_endpgm().to_bytes()
-    prog = decode_program(kernel)
-    # Should decode as 3 instructions: VOPD (with literal), then S_ENDPGM
-    # The literal should NOT be decoded as a separate instruction
-    self.assertEqual(len(prog), 2)  # VOPD + S_ENDPGM
-    exec_wave(prog, state, bytearray(65536), 1)
-    self.assertEqual(state.vgpr[0][1], 0x12345678)
-
-  def test_s_endpgm_stops_decode(self):
-    """Regression test: decoder should stop at S_ENDPGM, not read past into metadata."""
-    # Create a kernel followed by garbage that looks like an invalid instruction
-    kernel = s_mov_b32(s[0], 42).to_bytes() + s_endpgm().to_bytes()
-    garbage = bytes([0xff] * 16)  # garbage after kernel
-    prog = decode_program(kernel + garbage)
-    # Should only have 2 instructions (s_mov_b32 and s_endpgm)
-    self.assertEqual(len(prog), 2)
-
-class TestFloatConversion(unittest.TestCase):
-  """Unit tests for i32/i16/f32/f16 float conversion functions."""
-
-  def test_i32_preserves_nan_sign(self):
-    """NaN sign bit should be preserved when converting float to int bits."""
-    from extra.assembly.rdna3.emu import i32, f32
-    # 0 * -inf produces a negative NaN
-    neg_nan = 0.0 * float('-inf')
-    bits = i32(neg_nan)
-    # Should have sign bit set (0xffc00000), not canonical positive NaN (0x7fc00000)
-    self.assertEqual(bits & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{bits:08x}")
-    self.assertTrue(math.isnan(f32(bits)))
-
-  def test_i32_preserves_positive_nan(self):
-    """Positive NaN should remain positive."""
-    from extra.assembly.rdna3.emu import i32, f32
-    pos_nan = float('nan')
-    bits = i32(pos_nan)
-    # Standard Python NaN is positive (0x7fc00000)
-    self.assertEqual(bits & 0x80000000, 0, f"Expected positive NaN, got 0x{bits:08x}")
-    self.assertTrue(math.isnan(f32(bits)))
-
-  def test_i32_overflow_to_inf(self):
-    """Values too large for f32 should become inf."""
-    from extra.assembly.rdna3.emu import i32, f32
-    big = 2.0 ** 200
-    self.assertEqual(i32(big), 0x7f800000)  # +inf
-    self.assertEqual(i32(-big), 0xff800000)  # -inf
-
-  def test_i32_inf(self):
-    """Infinity should be preserved."""
-    from extra.assembly.rdna3.emu import i32
-    self.assertEqual(i32(float('inf')), 0x7f800000)
-    self.assertEqual(i32(float('-inf')), 0xff800000)
-
-  def test_i32_normal_values(self):
-    """Normal float values should round-trip correctly (within f32 precision)."""
-    from extra.assembly.rdna3.emu import i32, f32
-    # Use values exactly representable in float32
-    for val in [0.0, 1.0, -1.0, 0.5, -0.5, 100.0, -100.0, 1e10]:
-      bits = i32(val)
-      self.assertAlmostEqual(f32(bits), val, places=5)
-
-  def test_i16_overflow_to_inf(self):
-    """Values too large for f16 should become inf."""
-    from extra.assembly.rdna3.emu import i16
-    big = 100000.0  # way larger than f16 max (65504)
-    self.assertEqual(i16(big), 0x7c00)  # +inf
-    self.assertEqual(i16(-big), 0xfc00)  # -inf
-
-  def test_i16_inf(self):
-    """Infinity should be preserved."""
-    from extra.assembly.rdna3.emu import i16
-    self.assertEqual(i16(float('inf')), 0x7c00)
-    self.assertEqual(i16(float('-inf')), 0xfc00)
-
-  def test_fma_nan_sign_preserved(self):
-    """FMA producing NaN should preserve the correct sign bit."""
-    from extra.assembly.rdna3.emu import i32, f32
-    # 0 * (-inf) + 1.0 = NaN (from 0 * -inf)
-    a, b, c = 0.0, float('-inf'), 1.0
-    result = i32(a * b + c)
-    # The NaN should be negative since 0 * -inf produces negative NaN
-    self.assertEqual(result & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{result:08x}")
-
-class TestMultiWave(unittest.TestCase):
-  def test_all_waves_execute(self):
-    """Regression test: all waves in a workgroup must execute, not just the first."""
-    n_threads = 64  # 2 waves of 32 threads each
-    output = (ctypes.c_uint32 * n_threads)(*[0xdead] * n_threads)
-    output_ptr = ctypes.addressof(output)
-    args = (ctypes.c_uint64 * 1)(output_ptr)
-    args_ptr = ctypes.addressof(args)
-
-    # Simple kernel: store tid to output[tid]
-    kernel = b''
-    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    kernel += v_lshlrev_b32_e32(v[1], 2, v[0]).to_bytes()  # offset = tid * 4
-    kernel += global_store_b32(addr=v[1], data=v[0], saddr=s[2]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    kernel_ptr = ctypes.addressof(kernel_buf)
-    set_valid_mem_ranges({
-      (output_ptr, ctypes.sizeof(output)),
-      (args_ptr, ctypes.sizeof(args)),
-      (kernel_ptr, len(kernel)),
-    })
-    result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
-    self.assertEqual(result, 0)
-    # All threads should have written their tid
-    for i in range(n_threads):
-      self.assertEqual(output[i], i, f"Thread {i} didn't execute")
-
-class TestRegressions(unittest.TestCase):
-  """Regression tests for bugs fixed in the emulator."""
-
-  def test_v_fmac_f16(self):
-    """V_FMAC_F16: fused multiply-add for FP16. Regression for VOP2 op 54."""
-    from extra.assembly.rdna3.emu import i16, f16
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], i16(2.0)),  # v1.lo = 2.0 (fp16)
-      v_mov_b32_e32(v[2], i16(3.0)),  # v2.lo = 3.0 (fp16)
-      # v1 = v1 * v2 + v1 = 2.0 * 3.0 + 2.0 = 8.0
-      VOP2(VOP2Op.V_FMAC_F16, v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertAlmostEqual(f16(out[0] & 0xffff), 8.0, places=2)
-
-  def test_v_cvt_f64_f32(self):
-    """V_CVT_F64_F32: convert float32 to float64. Regression for VOP1 op 16."""
-    kernel = b''
-    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    kernel += v_mov_b32_e32(v[1], i32(3.14159)).to_bytes()
-    kernel += VOP1(VOP1Op.V_CVT_F64_F32, v[4], v[1]).to_bytes()  # v4:v5 = f64(v1)
-    kernel += v_lshlrev_b32_e32(v[3], 3, v[0]).to_bytes()  # offset = tid * 8
-    kernel += global_store_b64(addr=v[3], data=v[4], saddr=s[2]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    output = (ctypes.c_double * 1)(0.0)
-    output_ptr = ctypes.addressof(output)
-    args = (ctypes.c_uint64 * 1)(output_ptr)
-    args_ptr = ctypes.addressof(args)
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    kernel_ptr = ctypes.addressof(kernel_buf)
-    set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
-    run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
-    self.assertAlmostEqual(output[0], 3.14159, places=4)
-
-  def test_v_add_f64(self):
-    """V_ADD_F64: add two float64 values. Regression for VOP3 op 807."""
-    from extra.assembly.rdna3.emu import i64_parts
-    kernel = b''
-    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    # Load 1.5 into v1:v2
-    lo, hi = i64_parts(1.5)
-    kernel += v_mov_b32_e32(v[1], lo).to_bytes()
-    kernel += v_mov_b32_e32(v[2], hi).to_bytes()
-    # Load 2.5 into v3:v4
-    lo, hi = i64_parts(2.5)
-    kernel += v_mov_b32_e32(v[3], lo).to_bytes()
-    kernel += v_mov_b32_e32(v[4], hi).to_bytes()
-    # v5:v6 = v1:v2 + v3:v4 = 1.5 + 2.5 = 4.0
-    kernel += VOP3(VOP3Op.V_ADD_F64, v[5], v[1], v[3]).to_bytes()
-    kernel += v_lshlrev_b32_e32(v[7], 3, v[0]).to_bytes()
-    kernel += global_store_b64(addr=v[7], data=v[5], saddr=s[2]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-    output = (ctypes.c_double * 1)(0.0)
-    output_ptr = ctypes.addressof(output)
-    args = (ctypes.c_uint64 * 1)(output_ptr)
-    args_ptr = ctypes.addressof(args)
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    kernel_ptr = ctypes.addressof(kernel_buf)
-    set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
-    run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
-    self.assertAlmostEqual(output[0], 4.0, places=10)
-
-  def test_flat_load_d16_hi_b16(self):
-    """FLAT_LOAD_D16_HI_B16: load 16-bit to high half. Regression for FLAT op 35."""
-    from extra.assembly.rdna3.emu import i16
-    # Create a buffer with test data
-    src_data = (ctypes.c_uint16 * 1)(0x1234)
-    src_ptr = ctypes.addressof(src_data)
-    output = (ctypes.c_uint32 * 1)(0xABCD0000)  # preset low bits
-    output_ptr = ctypes.addressof(output)
-    args = (ctypes.c_uint64 * 2)(output_ptr, src_ptr)
-    args_ptr = ctypes.addressof(args)
-
-    kernel = b''
-    kernel += s_load_b128(s[0:3], s[0:1], soffset=NULL, offset=0).to_bytes()
-    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
-    kernel += v_mov_b32_e32(v[1], 0xDEAD).to_bytes()  # initial value with low bits set
-    kernel += v_mov_b32_e32(v[2], 0).to_bytes()  # offset = 0
-    kernel += FLAT(FLATOp.FLAT_LOAD_D16_HI_B16, v[1], v[2], saddr=s[2], offset=0).to_bytes()
-    kernel += s_waitcnt(vmcnt=0).to_bytes()
-    kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
-    kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[0]).to_bytes()
-    kernel += s_endpgm().to_bytes()
-
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    kernel_ptr = ctypes.addressof(kernel_buf)
-    set_valid_mem_ranges({(output_ptr, 4), (src_ptr, 2), (args_ptr, 16), (kernel_ptr, len(kernel))})
-    run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
-    # High 16 bits should be 0x1234, low 16 bits preserved as 0xDEAD
-    self.assertEqual(output[0], 0x1234DEAD)
-
-  def test_v_mad_u16(self):
-    """V_MAD_U16: multiply-add unsigned 16-bit. Regression for VOP3 op 577."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 10),  # a = 10
-      v_mov_b32_e32(v[2], 20),  # b = 20
-      v_mov_b32_e32(v[4], 5),   # c = 5
-      VOP3(VOP3Op.V_MAD_U16, v[1], v[1], v[2], v[4]),  # v1 = 10*20+5 = 205
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0] & 0xffff, 205)
-
-  def test_v_lshrrev_b16(self):
-    """V_LSHRREV_B16: logical shift right 16-bit. Regression for VOP3 op 825."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 0x8000),  # value to shift
-      v_mov_b32_e32(v[2], 4),       # shift amount
-      VOP3(VOP3Op.V_LSHRREV_B16, v[1], v[2], v[1]),  # v1 = 0x8000 >> 4 = 0x0800
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0] & 0xffff, 0x0800)
-
-  def test_v_min_u16(self):
-    """V_MIN_U16: minimum of two unsigned 16-bit values. Regression for VOP3 op 779."""
-    kernel = make_store_kernel([
-      v_mov_b32_e32(v[1], 100),
-      v_mov_b32_e32(v[2], 50),
-      VOP3(VOP3Op.V_MIN_U16, v[1], v[1], v[2]),
-    ])
-    out = run_kernel(kernel, n_threads=1)
-    self.assertEqual(out[0] & 0xffff, 50)
-
-class TestWMMA(unittest.TestCase):
-  """Tests for WMMA (Wave Matrix Multiply Accumulate) instructions."""
-
-  def test_wmma_f32_16x16x16_f16_identity(self):
-    """V_WMMA_F32_16X16X16_F16 with identity matrix. Regression for VOP3P op 64."""
-    from extra.assembly.rdna3.emu import i16, f16, exec_wmma_f32_16x16x16_f16, WaveState
-    # Test using direct emulator call rather than full kernel to simplify
-    st = WaveState()
-    st.exec_mask = 0xffffffff  # all 32 lanes active
-
-    # Set up A as identity matrix: A[i][i] = 1.0, rest = 0.0
-    # Lane i holds row i of A in 8 regs (2 fp16 per reg)
-    for lane in range(16):
-      for reg in range(8):
-        col0, col1 = reg * 2, reg * 2 + 1
-        val0 = i16(1.0) if col0 == lane else 0
-        val1 = i16(1.0) if col1 == lane else 0
-        st.vgpr[lane][0 + reg] = val0 | (val1 << 16)  # src0 = v0:v7
-
-    # Set up B as identity matrix: lane i holds column i of B
-    for lane in range(16):
-      for reg in range(8):
-        row0, row1 = reg * 2, reg * 2 + 1
-        val0 = i16(1.0) if row0 == lane else 0
-        val1 = i16(1.0) if row1 == lane else 0
-        st.vgpr[lane][8 + reg] = val0 | (val1 << 16)  # src1 = v8:v15
-
-    # Set up C as zeros
-    for lane in range(32):
-      for reg in range(8):
-        st.vgpr[lane][16 + reg] = 0  # src2 = v16:v23
-
-    # Create a fake VOP3P instruction
-    inst = VOP3P(VOP3POp.V_WMMA_F32_16X16X16_F16, v[24], src0=VGPR(0), src1=VGPR(8), src2=VGPR(16))
-
-    # Execute WMMA
-    exec_wmma_f32_16x16x16_f16(st, inst, 32)
-
-    # Check result: C should be identity (since A @ B where both are identity)
-    # Output i = row*16+col goes to lane (i%32), reg (i//32)
-    for row in range(16):
-      for col in range(16):
-        idx = row * 16 + col
-        lane, reg = idx % 32, idx // 32
-        result = st.vgpr[lane][24 + reg]
-        expected = 1.0 if row == col else 0.0
-        self.assertAlmostEqual(f32(result), expected, places=3,
-          msg=f"C[{row},{col}] = {f32(result)}, expected {expected}")
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_formats.py

@@ -1,332 +0,0 @@
-#!/usr/bin/env python3
-"""Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
-import unittest
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.lib import encode_src
-
-class TestMUBUF(unittest.TestCase):
-  """Test MUBUF (buffer) instructions."""
-
-  def test_buffer_load_b32_basic(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_idxen(self):
-    # buffer_load_b32 v5, v0, s[8:11], s3 idxen offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, idxen=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]))
-
-  def test_buffer_load_b32_offen(self):
-    # buffer_load_b32 v5, v0, s[8:11], s3 offen offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, offen=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]))
-
-  def test_buffer_load_b32_glc(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc
-    # GFX11: encoding: [0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_slc(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 slc
-    # GFX11: encoding: [0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, slc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_dlc(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 dlc
-    # GFX11: encoding: [0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_b32_all_flags(self):
-    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc slc dlc
-    # GFX11: encoding: [0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1, slc=1, dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_store_b32(self):
-    # buffer_store_b32 v1, off, s[12:15], s4 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]
-    inst = buffer_store_b32(vdata=v[1], vaddr=v[0], srsrc=s[12:16], soffset=s[4], offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]))
-
-  def test_buffer_load_b64(self):
-    # buffer_load_b64 v[5:6], off, s[8:11], s3 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]
-    inst = buffer_load_b64(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]))
-
-  def test_buffer_load_soffset_m0(self):
-    # buffer_load_b32 v5, off, s[8:11], m0 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=M0, offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]))
-
-  def test_buffer_load_soffset_inline_const(self):
-    # buffer_load_b32 v5, off, s[8:11], 0 offset:4095
-    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=0, offset=4095)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]))
-
-  def test_buffer_disasm_roundtrip(self):
-    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
-    decoded = MUBUF.from_bytes(inst.to_bytes())
-    self.assertEqual(decoded.to_bytes(), inst.to_bytes())
-
-
-class TestMTBUF(unittest.TestCase):
-  """Test MTBUF (typed buffer) instructions."""
-
-  def test_tbuffer_load_format_x(self):
-    # tbuffer_load_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
-    # BUF_FMT_32_FLOAT = 22
-    # GFX11: encoding: [0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]
-    inst = tbuffer_load_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]))
-
-  def test_tbuffer_store_format_x(self):
-    # tbuffer_store_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
-    # BUF_FMT_32_FLOAT = 22
-    # GFX11: encoding: [0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]
-    inst = tbuffer_store_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]))
-
-  def test_tbuffer_load_format_xy(self):
-    # tbuffer_load_format_xy v[5:6], off, s[8:11], s3 format:[BUF_FMT_32_32_FLOAT] offset:4095
-    # BUF_FMT_32_32_FLOAT = 50
-    # GFX11: encoding: [0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]
-    inst = tbuffer_load_format_xy(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=50)
-    self.assertEqual(inst.to_bytes(), bytes([0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]))
-
-
-class TestMIMG(unittest.TestCase):
-  """Test MIMG (image) instructions."""
-
-  def test_image_load_2d(self):
-    # image_load v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
-    # GFX11: encoding: [0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_load(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1)  # dim=1 is SQ_RSRC_IMG_2D
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
-
-  def test_image_store_2d(self):
-    # image_store v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
-    # GFX11: encoding: [0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_store(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]))
-
-  def test_image_load_1d(self):
-    # image_load v[0:3], v4, s[0:7] dmask:0xf dim:SQ_RSRC_IMG_1D
-    # GFX11: encoding: [0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_load(vdata=v[0:4], vaddr=v[4], srsrc=s[0:8], dmask=0xf, dim=0)  # dim=0 is SQ_RSRC_IMG_1D
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
-
-  def test_image_sample(self):
-    # image_sample v[0:3], v[4:5], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
-    # GFX11: encoding: [0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]
-    inst = image_sample(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], ssamp=s[8:12], dmask=0xf, dim=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]))
-
-  def test_image_load_d16(self):
-    # image_load v[0:1], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D d16
-    # GFX11: encoding: [0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]
-    inst = image_load(vdata=v[0:2], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1, d16=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]))
-
-
-class TestEXP(unittest.TestCase):
-  """Test EXP (export) instructions."""
-
-  def test_exp_mrt0(self):
-    # exp mrt0 v0, v1, v2, v3
-    # GFX11: encoding: [0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]
-    inst = EXP(en=0xf, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[2], vsrc3=v[3])
-    self.assertEqual(inst.to_bytes(), bytes([0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]))
-
-  def test_exp_mrtz(self):
-    # exp mrtz v4, v3, v2, v1
-    # GFX11: encoding: [0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]
-    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]))
-
-  def test_exp_mrtz_done(self):
-    # exp mrtz v4, v3, v2, v1 done
-    # GFX11: encoding: [0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x01]
-    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[3], done=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x03]))
-
-  def test_exp_partial_mask(self):
-    # exp mrt0 v0, v1, off, off (en=0x3, only first two components)
-    # GFX11: encoding: [0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]
-    inst = EXP(en=0x3, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[0], vsrc3=v[0])
-    self.assertEqual(inst.to_bytes(), bytes([0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]))
-
-  def test_exp_row_en(self):
-    # exp mrtz v4, v3, v2, v1 row_en
-    # GFX11: encoding: [0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]
-    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1], row=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]))
-
-
-class TestDS(unittest.TestCase):
-  """Test DS (data share / LDS) instructions."""
-
-  def test_ds_store_b32(self):
-    # ds_store_b32 v0, v1
-    # GFX11: encoding: [0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_store_b32(addr=v[0], data0=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
-
-  def test_ds_load_b32(self):
-    # ds_load_b32 v0, v1
-    # GFX11: encoding: [0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]
-    inst = ds_load_b32(vdst=v[0], addr=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]))
-
-  def test_ds_store_b32_offset(self):
-    # ds_store_b32 v0, v1 offset:64
-    # GFX11: encoding: [0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_store_b32(addr=v[0], data0=v[1], offset0=64)
-    self.assertEqual(inst.to_bytes(), bytes([0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
-
-  def test_ds_load_b64(self):
-    # ds_load_b64 v[0:1], v2
-    # GFX11: encoding: [0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]
-    inst = ds_load_b64(vdst=v[0:2], addr=v[2])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]))
-
-  def test_ds_add_u32(self):
-    # ds_add_u32 v0, v1
-    # GFX11: encoding: [0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_add_u32(addr=v[0], data0=v[1])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]))
-
-  def test_ds_store_b32_gds(self):
-    # ds_store_b32 v0, v1 gds
-    # GFX11: encoding: [0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]
-    inst = ds_store_b32(addr=v[0], data0=v[1], gds=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]))
-
-
-class TestVOP3(unittest.TestCase):
-  """Test VOP3 (3-operand vector) instructions."""
-
-  def test_v_fma_f32(self):
-    # v_fma_f32 v0, v1, v2, v3
-    # GFX11: encoding: [0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]
-    inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]))
-
-  def test_v_mad_f32(self):
-    # v_fmac_f32_e64 v0, v1, v2 (fmac is fma with implicit dst as src2)
-    # Use v_fma_f32 with vdst == src2
-    inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[0])
-    self.assertEqual(inst.to_bytes()[:4], bytes([0x00,0x00,0x13,0xd6]))
-
-  def test_v_add3_u32(self):
-    # v_add3_u32 v0, v1, v2, v3
-    # GFX11: encoding: [0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]
-    inst = v_add3_u32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]))
-
-
-class TestFLAT(unittest.TestCase):
-  """Test FLAT/GLOBAL/SCRATCH memory instructions."""
-
-  def test_global_load_b32(self):
-    # global_load_b32 v0, v[1:2], off (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]
-    inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]))
-
-  def test_global_store_b32(self):
-    # global_store_b32 v[0:1], v2, off (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]
-    inst = global_store_b32(addr=v[0:2], data=v[2], saddr=OFF)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]))
-
-  def test_global_load_b32_saddr(self):
-    # global_load_b32 v0, v1, s[0:1] (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]
-    inst = global_load_b32(vdst=v[0], addr=v[1], saddr=s[0:2])
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]))
-
-  def test_global_load_b32_offset(self):
-    # global_load_b32 v0, v[1:2], off offset:256 (seg=2 for global)
-    # GFX11: encoding: [0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]
-    inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF, offset=256)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]))
-
-  def test_global_load_b64(self):
-    # global_load_b64 v[0:1], v[2:3], off (seg=2 for global)
-    # GFX11: encoding: [0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]
-    inst = global_load_b64(vdst=v[0:2], addr=v[2:4], saddr=OFF)
-    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]))
-
-
-class TestSMEM(unittest.TestCase):
-  """Test SMEM (scalar memory) instructions - regression tests for glc/dlc bit positions."""
-
-  def test_smem_dlc_bit_position(self):
-    # s_load_b32 s5, s[2:3], s0 dlc - tests that DLC is at bit 13 (not bit 14)
-    # GFX11: encoding: [0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]
-    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]))
-
-  def test_smem_glc_bit_position(self):
-    # s_load_b32 s5, s[2:3], s0 glc - tests that GLC is at bit 14 (not bit 16)
-    # GFX11: encoding: [0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]
-    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]))
-
-  def test_smem_glc_dlc_combined(self):
-    # s_load_b32 s5, s[2:3], s0 glc dlc - tests both flags together
-    # GFX11: encoding: [0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]
-    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1, dlc=1)
-    self.assertEqual(inst.to_bytes(), bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]))
-
-  def test_smem_disasm_roundtrip_dlc(self):
-    # Test that disassembly/reassembly preserves DLC bit correctly
-    data = bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00])
-    decoded = SMEM.from_bytes(data)
-    self.assertEqual(decoded.to_bytes(), data)
-
-  def test_smem_disasm_roundtrip_glc_dlc(self):
-    # Test that disassembly/reassembly preserves GLC+DLC bits correctly
-    data = bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00])
-    decoded = SMEM.from_bytes(data)
-    self.assertEqual(decoded.to_bytes(), data)
-
-
-class TestVOP3Literal(unittest.TestCase):
-  """Test VOP3 literal handling - regression tests for Inst64 literal encoding."""
-
-  def test_vop3_with_literal(self):
-    # v_add3_u32 v5, vcc_hi, 0xaf123456, v255
-    # GFX11: encoding: [0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf]
-    from extra.assembly.rdna3.lib import RawImm
-    inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=RawImm(107), src1=0xaf123456, src2=v[255])
-    expected = bytes([0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf])
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_vop3_literal_null_operand(self):
-    # v_add3_u32 v5, null, exec_lo, 0xaf123456
-    # GFX11: encoding: [0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf]
-    from extra.assembly.rdna3.lib import RawImm
-    inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=NULL, src1=RawImm(126), src2=0xaf123456)
-    expected = bytes([0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf])
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_vop3p_with_literal(self):
-    # Test VOP3P literal encoding (also uses Inst64)
-    from extra.assembly.rdna3.lib import RawImm
-    inst = VOP3P(VOP3POp.V_PK_ADD_F16, vdst=v[5], src0=RawImm(240), src1=0x12345678, src2=v[0])
-    self.assertEqual(len(inst.to_bytes()), 12)  # 8 bytes + 4 byte literal
-
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_integration.py

@@ -1,343 +0,0 @@
-#!/usr/bin/env python3
-"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
-import unittest, re, io, sys
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.asm import waitcnt, asm
-
-def get_amd_toolchain():
-  """Check if AMD toolchain is available."""
-  try:
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    HIPCompiler("gfx1100").compile(".text\ns_endpgm")
-    return True
-  except Exception:
-    return False
-
-def disassemble(lib: bytes, arch: str = "gfx1100") -> str:
-  """Disassemble ELF binary using tinygrad's compiler, return raw output."""
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-  old_stdout = sys.stdout
-  sys.stdout = io.StringIO()
-  HIPCompiler(arch).disassemble(lib)
-  output = sys.stdout.getvalue()
-  sys.stdout = old_stdout
-  return output
-
-def parse_disassembly(raw: str) -> list[str]:
-  """Parse disassembly output to list of instruction mnemonics."""
-  lines = []
-  for line in raw.splitlines():
-    if line.startswith('\t'):
-      instr = line.split('//')[0].strip()
-      if instr: lines.append(instr)
-  return lines
-
-def assemble_and_disassemble(instructions: list, arch: str = "gfx1100") -> list[str]:
-  """Assemble instructions with our DSL, then disassemble with AMD toolchain."""
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-  # Generate bytes from our DSL
-  code_bytes = b''.join(inst.to_bytes() for inst in instructions)
-
-  # Wrap in minimal ELF-compatible assembly with .byte directives
-  byte_str = ', '.join(f'0x{b:02x}' for b in code_bytes)
-  asm_src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n.byte {byte_str}\n"
-
-  # Assemble with AMD COMGR and disassemble
-  lib = HIPCompiler(arch).compile(asm_src)
-  return parse_disassembly(disassemble(lib, arch))
-
-@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
-class TestIntegration(unittest.TestCase):
-  """Test our assembler output matches LLVM disassembly."""
-
-  def test_simple_sop1(self):
-    """Test SOP1 instructions round-trip."""
-    instructions = [
-      s_mov_b32(s[0], s[1]),
-      s_mov_b32(s[2], 0),
-      s_not_b32(s[3], s[4]),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_mov_b32', disasm[0])
-    self.assertIn('s_mov_b32', disasm[1])
-    self.assertIn('s_not_b32', disasm[2])
-
-  def test_simple_sop2(self):
-    """Test SOP2 instructions round-trip."""
-    instructions = [
-      s_add_u32(s[0], s[1], s[2]),
-      s_sub_u32(s[3], s[4], 10),
-      s_and_b32(s[5], s[6], s[7]),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_add_u32', disasm[0])
-    self.assertIn('s_sub_u32', disasm[1])
-    self.assertIn('s_and_b32', disasm[2])
-
-  def test_simple_vop2(self):
-    """Test VOP2 instructions round-trip."""
-    instructions = [
-      v_add_f32_e32(v[0], v[1], v[2]),
-      v_mul_f32_e32(v[3], 1.0, v[4]),  # 1.0 is inline constant
-      v_and_b32_e32(v[5], 10, v[6]),  # small inline constant
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('v_add_f32', disasm[0])
-    self.assertIn('v_mul_f32', disasm[1])
-
-  def test_control_flow(self):
-    """Test control flow instructions."""
-    instructions = [
-      s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
-      s_endpgm(),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_waitcnt', disasm[0])
-    self.assertIn('s_endpgm', disasm[1])
-
-  def test_memory_ops(self):
-    """Test memory instructions."""
-    instructions = [
-      s_load_b32(s[0], s[0:2], NULL),
-      s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
-      global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
-      s_endpgm(),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    self.assertIn('s_load_b32', disasm[0])
-    self.assertIn('s_waitcnt', disasm[1])
-    self.assertIn('global_store_b32', disasm[2])
-
-  def test_full_kernel(self):
-    """Test a complete kernel similar to tinygrad output."""
-    # Simple kernel: load value, add 1, store back
-    instructions = [
-      # Get thread ID
-      v_mov_b32_e32(v[0], s[0]),  # base addr low
-      v_mov_b32_e32(v[1], s[1]),  # base addr high
-      # Load value
-      global_load_b32(vdst=v[2], addr=v[0:2], saddr=OFF),
-      s_waitcnt(simm16=waitcnt(vmcnt=0)),
-      # Add 1.0
-      v_add_f32_e32(v[2], 1.0, v[2]),
-      # Store result
-      global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
-      s_endpgm(),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    # Verify key instructions are present
-    self.assertTrue(any('global_load' in d for d in disasm))
-    self.assertTrue(any('v_add_f32' in d for d in disasm))
-    self.assertTrue(any('global_store' in d for d in disasm))
-    self.assertTrue(any('s_endpgm' in d for d in disasm))
-
-  def test_bytes_roundtrip(self):
-    """Test that our bytes match what AMD assembler produces."""
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-    # Simple instruction
-    inst = s_mov_b32(s[0], s[1])
-    our_bytes = inst.to_bytes()
-
-    # Assemble same instruction with AMD toolchain
-    asm_src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\ns_mov_b32 s0, s1\n"
-    compiler = HIPCompiler("gfx1100")
-    lib = compiler.compile(asm_src)
-    raw = disassemble(lib)
-
-    for line in raw.splitlines():
-      if 's_mov_b32' in line and '//' in line:
-        # Extract hex bytes from comment: "// 000000001300: BE800001"
-        comment = line.split('//')[1].strip()
-        hex_str = comment.split(':')[1].strip()
-        # Convert big-endian hex string to little-endian bytes
-        amd_bytes = bytes.fromhex(hex_str)[::-1]  # reverse for little-endian
-        self.assertEqual(our_bytes, amd_bytes, f"Bytes mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
-        return
-    self.fail("Could not find s_mov_b32 in disassembly")
-
-@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
-class TestAsm(unittest.TestCase):
-  """Test asm() string parsing."""
-
-  def test_asm_basic(self):
-    """Test basic instruction parsing."""
-    inst = asm('s_mov_b32 s0, s1')
-    self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], s[1]).to_bytes())
-
-  def test_asm_with_immediates(self):
-    """Test parsing with immediate values."""
-    inst = asm('s_add_u32 s0, s1, 10')
-    self.assertEqual(inst.to_bytes(), s_add_u32(s[0], s[1], 10).to_bytes())
-
-  def test_asm_float_const(self):
-    """Test parsing float constants."""
-    inst = asm('v_mul_f32_e32 v0, 1.0, v1')
-    self.assertEqual(inst.to_bytes(), v_mul_f32_e32(v[0], 1.0, v[1]).to_bytes())
-
-  def test_asm_hex_immediate(self):
-    """Test parsing hex immediates."""
-    inst = asm('s_waitcnt 0xfc07')
-    self.assertEqual(inst.to_bytes(), s_waitcnt(simm16=0xfc07).to_bytes())
-
-  def test_asm_special_regs(self):
-    """Test parsing special registers."""
-    inst = asm('s_mov_b32 s0, vcc_lo')
-    self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], VCC_LO).to_bytes())
-
-  def test_asm_register_range(self):
-    """Test parsing register ranges."""
-    inst = asm('s_load_b128 s[4:7], s[0:1], null')
-    self.assertEqual(inst.to_bytes(), s_load_b128(s[4:7], s[0:1], NULL).to_bytes())
-
-  def test_asm_matches_llvm(self):
-    """Test asm() output matches LLVM assembler."""
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    compiler = HIPCompiler('gfx1100')
-
-    def get_llvm_bytes(instr: str) -> bytes:
-      src = f'.text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n{instr}\n'
-      lib = compiler.compile(src)
-      raw = disassemble(lib)
-      for line in raw.splitlines():
-        if instr.split()[0] in line and '//' in line:
-          hex_str = line.split('//')[1].strip().split(':')[1].strip()
-          return bytes.fromhex(hex_str)[::-1]
-      return b''
-
-    tests = ['s_mov_b32 s0, s1', 's_endpgm', 'v_add_f32_e32 v0, v1, v2']
-    for t in tests:
-      self.assertEqual(asm(t).to_bytes(), get_llvm_bytes(t), f"mismatch for: {t}")
-
-  def test_asm_vop3_modifiers(self):
-    """Test asm() with VOP3 modifiers (neg, abs, clamp)."""
-    import subprocess, re
-
-    def get_llvm_encoding(instr: str) -> str:
-      result = subprocess.run(['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-show-encoding'],
-                              input=instr, capture_output=True, text=True)
-      if m := re.search(r'encoding:\s*\[(.*?)\]', result.stdout):
-        return m.group(1).replace('0x','').replace(',','').replace(' ','')
-      return ''
-
-    tests = [
-      'v_fma_f32 v0, -v1, v2, v3',       # neg on src0
-      'v_fma_f32 v0, v1, |v2|, v3',      # abs on src1
-      'v_fma_f32 v0, v1, v2, v3 clamp',  # clamp
-      'v_fma_f32 v0, -v1, |v2|, v3 clamp',  # all modifiers
-      'v_fma_f32 v0, -|v1|, v2, v3',     # neg+abs on same operand
-    ]
-    for t in tests:
-      our_hex = asm(t).to_bytes().hex()
-      llvm_hex = get_llvm_encoding(t)
-      self.assertEqual(our_hex, llvm_hex, f"mismatch for: {t}")
-
-@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
-class TestTinygradIntegration(unittest.TestCase):
-  """Test that we can parse disassembled tinygrad kernels."""
-
-  def test_simple_add_kernel(self):
-    """Generate a simple add kernel from tinygrad and verify disassembly."""
-    from tinygrad import Tensor
-    from tinygrad.codegen import get_program
-    from tinygrad.renderer.cstyle import AMDHIPRenderer
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    from tinygrad.uop.ops import Ops
-
-    # Create a computation that generates a real kernel
-    a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
-    b = Tensor([5.0, 6.0, 7.0, 8.0]).realize()
-    c = a + b
-
-    # Get schedule and find SINK
-    schedule = c.schedule()
-    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
-    self.assertTrue(len(sink_items) > 0, "No SINK in schedule")
-
-    # Generate program
-    renderer = AMDHIPRenderer('gfx1100')
-    prg = get_program(sink_items[0].ast, renderer)
-    self.assertIsNotNone(prg.src)
-
-    # Compile and disassemble
-    compiler = HIPCompiler('gfx1100')
-    lib = compiler.compile(prg.src)
-    raw_disasm = disassemble(lib)
-    instrs = parse_disassembly(raw_disasm)
-
-    # Verify we got some instructions
-    self.assertTrue(len(instrs) > 0, "No instructions in disassembly")
-    # Should have an endpgm
-    self.assertTrue(any('s_endpgm' in i for i in instrs), "Missing s_endpgm")
-
-  def test_matmul_kernel(self):
-    """Generate a matmul kernel and verify disassembly has expected patterns."""
-    from tinygrad import Tensor
-    from tinygrad.codegen import get_program
-    from tinygrad.renderer.cstyle import AMDHIPRenderer
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    from tinygrad.uop.ops import Ops
-
-    # Create a small matmul
-    a = Tensor.rand(4, 4).realize()
-    b = Tensor.rand(4, 4).realize()
-    c = a @ b
-
-    # Get schedule
-    schedule = c.schedule()
-    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
-    self.assertTrue(len(sink_items) > 0)
-
-    # Generate and compile
-    renderer = AMDHIPRenderer('gfx1100')
-    prg = get_program(sink_items[0].ast, renderer)
-    compiler = HIPCompiler('gfx1100')
-    lib = compiler.compile(prg.src)
-    raw_disasm = disassemble(lib)
-    instrs = parse_disassembly(raw_disasm)
-
-    # Matmul should have multiply and add instructions
-    has_mul = any('mul' in i.lower() for i in instrs)
-    has_add = any('add' in i.lower() for i in instrs)
-    self.assertTrue(has_mul or has_add, "Matmul should have mul/add ops")
-
-  def test_disasm_to_bytes_roundtrip(self):
-    """Parse disassembled instructions and verify we can re-encode some of them."""
-    from tinygrad import Tensor
-    from tinygrad.codegen import get_program
-    from tinygrad.renderer.cstyle import AMDHIPRenderer
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    from tinygrad.uop.ops import Ops
-
-    # Simple kernel
-    a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
-    b = (a * 2.0)
-
-    schedule = b.schedule()
-    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
-    if not sink_items: return  # skip if no kernel
-
-    renderer = AMDHIPRenderer('gfx1100')
-    prg = get_program(sink_items[0].ast, renderer)
-    compiler = HIPCompiler('gfx1100')
-    lib = compiler.compile(prg.src)
-    raw_disasm = disassemble(lib)
-
-    # Find s_endpgm and verify we can encode it
-    for line in raw_disasm.splitlines():
-      if 's_endpgm' in line and '//' in line:
-        # Extract bytes from comment
-        comment = line.split('//')[1].strip()
-        hex_str = comment.split(':')[1].strip()
-        amd_bytes = bytes.fromhex(hex_str)[::-1]
-
-        # Our encoding
-        our_inst = s_endpgm()
-        our_bytes = our_inst.to_bytes()
-
-        self.assertEqual(our_bytes, amd_bytes, f"s_endpgm mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
-        return
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_llvm.py

@@ -1,173 +0,0 @@
-#!/usr/bin/env python3
-"""Test RDNA3 assembler/disassembler against LLVM test vectors."""
-import unittest, re
-from tinygrad.helpers import fetch
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.asm import asm
-from extra.assembly.rdna3.test.test_roundtrip import compile_asm, disassemble_lib
-
-LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/main/llvm/test/MC/AMDGPU"
-
-# Format info: (filename, format_class, op_enum)
-LLVM_TEST_FILES = {
-  # Scalar ALU
-  'sop1': ('gfx11_asm_sop1.s', SOP1, SOP1Op),
-  'sop2': ('gfx11_asm_sop2.s', SOP2, SOP2Op),
-  'sopp': ('gfx11_asm_sopp.s', SOPP, SOPPOp),
-  'sopk': ('gfx11_asm_sopk.s', SOPK, SOPKOp),
-  'sopc': ('gfx11_asm_sopc.s', SOPC, SOPCOp),
-  # Vector ALU
-  'vop1': ('gfx11_asm_vop1.s', VOP1, VOP1Op),
-  'vop2': ('gfx11_asm_vop2.s', VOP2, VOP2Op),
-  'vopc': ('gfx11_asm_vopc.s', VOPC, VOPCOp),
-  'vop3': ('gfx11_asm_vop3.s', VOP3, VOP3Op),
-  'vop3p': ('gfx11_asm_vop3p.s', VOP3P, VOP3POp),
-  'vop3sd': ('gfx11_asm_vop3.s', VOP3SD, VOP3SDOp),  # VOP3SD shares file with VOP3
-  'vinterp': ('gfx11_asm_vinterp.s', VINTERP, VINTERPOp),
-  'vopd': ('gfx11_asm_vopd.s', VOPD, VOPDOp),
-  'vopcx': ('gfx11_asm_vopcx.s', VOPC, VOPCOp),  # VOPCX uses VOPC format
-  # VOP3 promotions (VOP1/VOP2/VOPC promoted to VOP3 encoding)
-  'vop3_from_vop1': ('gfx11_asm_vop3_from_vop1.s', VOP3, VOP3Op),
-  'vop3_from_vop2': ('gfx11_asm_vop3_from_vop2.s', VOP3, VOP3Op),
-  'vop3_from_vopc': ('gfx11_asm_vop3_from_vopc.s', VOP3, VOP3Op),
-  'vop3_from_vopcx': ('gfx11_asm_vop3_from_vopcx.s', VOP3, VOP3Op),
-  # Memory
-  'ds': ('gfx11_asm_ds.s', DS, DSOp),
-  'smem': ('gfx11_asm_smem.s', SMEM, SMEMOp),
-  'flat': ('gfx11_asm_flat.s', FLAT, FLATOp),
-  'mubuf': ('gfx11_asm_mubuf.s', MUBUF, MUBUFOp),
-  'mtbuf': ('gfx11_asm_mtbuf.s', MTBUF, MTBUFOp),
-  'mimg': ('gfx11_asm_mimg.s', MIMG, MIMGOp),
-  # WMMA (matrix multiply)
-  'wmma': ('gfx11_asm_wmma.s', VOP3P, VOP3POp),
-  # Additional features
-  'vop3_features': ('gfx11_asm_vop3_features.s', VOP3, VOP3Op),
-  'vop3p_features': ('gfx11_asm_vop3p_features.s', VOP3P, VOP3POp),
-  'vopd_features': ('gfx11_asm_vopd_features.s', VOPD, VOPDOp),
-  # Alias files (alternative mnemonics)
-  'vop3_alias': ('gfx11_asm_vop3_alias.s', VOP3, VOP3Op),
-  'vop3p_alias': ('gfx11_asm_vop3p_alias.s', VOP3P, VOP3POp),
-  'vopc_alias': ('gfx11_asm_vopc_alias.s', VOPC, VOPCOp),
-  'vopcx_alias': ('gfx11_asm_vopcx_alias.s', VOPC, VOPCOp),
-  'vinterp_alias': ('gfx11_asm_vinterp_alias.s', VINTERP, VINTERPOp),
-  'smem_alias': ('gfx11_asm_smem_alias.s', SMEM, SMEMOp),
-  'mubuf_alias': ('gfx11_asm_mubuf_alias.s', MUBUF, MUBUFOp),
-  'mtbuf_alias': ('gfx11_asm_mtbuf_alias.s', MTBUF, MTBUFOp),
-}
-
-def parse_llvm_tests(text: str) -> list[tuple[str, bytes]]:
-  """Parse LLVM test format into (asm, expected_bytes) pairs."""
-  tests, lines = [], text.split('\n')
-  for i, line in enumerate(lines):
-    line = line.strip()
-    if not line or line.startswith(('//', '.', ';')): continue
-    asm_text = line.split('//')[0].strip()
-    if not asm_text: continue
-    for j in range(i, min(i + 3, len(lines))):
-      # Match GFX11, W32, or W64 encodings (all valid for gfx11)
-      if m := re.search(r'(?:GFX11|W32|W64)[^:]*:.*?encoding:\s*\[(.*?)\]', lines[j]):
-        hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
-        if hex_bytes:
-          try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
-          except ValueError: pass
-        break
-  return tests
-
-def try_assemble(text: str):
-  """Try to assemble instruction text, return bytes or None on failure."""
-  try: return asm(text).to_bytes()
-  except: return None
-
-class TestLLVM(unittest.TestCase):
-  """Test assembler and disassembler against all LLVM test vectors."""
-  tests: dict[str, list[tuple[str, bytes]]] = {}
-
-  @classmethod
-  def setUpClass(cls):
-    for name, (filename, _, _) in LLVM_TEST_FILES.items():
-      try:
-        data = fetch(f"{LLVM_BASE}/{filename}").read_bytes()
-        cls.tests[name] = parse_llvm_tests(data.decode('utf-8', errors='ignore'))
-      except Exception as e:
-        print(f"Warning: couldn't fetch {filename}: {e}")
-        cls.tests[name] = []
-
-# Generate test methods dynamically for each format
-def _make_asm_test(name):
-  def test(self):
-    passed, failed, skipped = 0, 0, 0
-    for asm_text, expected in self.tests.get(name, []):
-      result = try_assemble(asm_text)
-      if result is None: skipped += 1
-      elif result == expected: passed += 1
-      else: failed += 1
-    print(f"{name.upper()} asm: {passed} passed, {failed} failed, {skipped} skipped")
-    self.assertEqual(failed, 0)
-  return test
-
-def _make_disasm_test(name):
-  def test(self):
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-    compiler = HIPCompiler('gfx1100')
-    _, fmt_cls, op_enum = LLVM_TEST_FILES[name]
-    passed, failed, skipped, failures = 0, 0, 0, []
-    # VOP3SD opcodes that share encoding with VOP3 (only for vop3sd test, not vopc promotions)
-    # Note: opcodes 0-255 are VOPC promoted to VOP3, never VOP3SD
-    vop3sd_opcodes = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
-    # vop3_from_vopc/vopcx tests have VOPC opcodes 0-255, not VOP3SD - don't detect as VOP3SD
-    is_vopc_promotion = name in ('vop3_from_vopc', 'vop3_from_vopcx')
-    # Undocumented opcodes not in AMD ISA PDF - skip these
-    undocumented = {'smem': {34, 35}, 'sopk': {22, 23}, 'sopp': {8, 58, 59}}  # s_atc_probe*, s_subvector_loop*, s_waitcnt_depctr, unknown
-    for asm_text, data in self.tests.get(name, []):
-      if len(data) > fmt_cls._size(): continue  # skip literals (need different handling)
-      # Skip undocumented opcodes
-      temp_inst = fmt_cls.from_bytes(data)
-      temp_op = temp_inst._values.get('op', 0)
-      temp_op = temp_op.val if hasattr(temp_op, 'val') else temp_op
-      if temp_op in undocumented.get(name, set()): skipped += 1; continue
-      # Skip SOPP no-imm instructions with non-zero simm16 (can't roundtrip through LLVM)
-      if name == 'sopp':
-        simm16 = temp_inst._values.get('simm16', 0)
-        simm16 = simm16.val if hasattr(simm16, 'val') else simm16
-        sopp_no_imm = {48, 54, 53, 55, 60, 61, 62}  # s_endpgm, s_barrier, s_wakeup, s_icache_inv, s_wait_idle, s_endpgm_saved, s_code_end
-        if temp_op in sopp_no_imm and simm16 != 0: skipped += 1; continue
-      try:
-        # VOP3 and VOP3SD share encoding - peek at opcode to determine which class to use
-        if fmt_cls.__name__ in ('VOP3', 'VOP3SD'):
-          temp = VOP3.from_bytes(data)
-          op_val = temp._values.get('op', 0)
-          op_val = op_val.val if hasattr(op_val, 'val') else op_val
-          is_vop3sd = (op_val in vop3sd_opcodes) and not is_vopc_promotion
-          decoded = VOP3SD.from_bytes(data) if is_vop3sd else VOP3.from_bytes(data)
-          # Validate opcode with appropriate enum
-          if is_vop3sd:
-            VOP3SDOp(op_val)
-          else:
-            VOP3Op(op_val)
-        else:
-          decoded = fmt_cls.from_bytes(data)
-          op_val = decoded._values.get('op', 0)
-          op_val = op_val.val if hasattr(op_val, 'val') else op_val
-          op_enum(op_val)  # validate opcode
-        if decoded.to_bytes()[:len(data)] != data:
-          failed += 1; failures.append(f"decode roundtrip failed for {data.hex()}"); continue
-        disasm_str = decoded.disasm()
-        # Test: LLVM should assemble our disasm output to the same bytes
-        llvm_bytes = compile_asm(disasm_str, compiler)
-        if llvm_bytes is None:
-          failed += 1; failures.append(f"LLVM failed to assemble: '{disasm_str}' (from '{asm_text}')")
-        elif llvm_bytes == data: passed += 1
-        else: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
-      except Exception as e:
-        failed += 1; failures.append(f"exception for {data.hex()}: {e}")
-    print(f"{name.upper()} disasm: {passed} passed, {failed} failed" + (f", {skipped} skipped" if skipped else ""))
-    if failures[:10]: print("  " + "\n  ".join(failures[:10]))
-    self.assertEqual(failed, 0)
-  return test
-
-for name in LLVM_TEST_FILES:
-  setattr(TestLLVM, f'test_{name}_asm', _make_asm_test(name))
-  setattr(TestLLVM, f'test_{name}_disasm', _make_disasm_test(name))
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_pdf_parser.py

@@ -1,158 +0,0 @@
-#!/usr/bin/env python3
-"""Test that PDF parser correctly extracts format fields."""
-import unittest
-from extra.assembly.rdna3.autogen import (
-  SOP1, SOP2, SOPK, SOPP, VOP1, VOP2, VOP3SD, VOPC, FLAT, VOPD,
-  SOP1Op, SOP2Op, VOP1Op, VOP3Op
-)
-
-# expected formats with key fields and whether they have ENCODING
-EXPECTED_FORMATS = {
-  'DPP16': (['SRC0', 'DPP_CTRL', 'BANK_MASK', 'ROW_MASK'], False),
-  'DPP8': (['SRC0', 'LANE_SEL0', 'LANE_SEL7'], False),
-  'DS': (['OP', 'ADDR', 'DATA0', 'DATA1', 'VDST'], True),
-  'EXP': (['EN', 'TARGET', 'VSRC0', 'VSRC1', 'VSRC2', 'VSRC3'], True),
-  'FLAT': (['OP', 'ADDR', 'DATA', 'SADDR', 'VDST', 'OFFSET'], True),
-  'LDSDIR': (['VDST', 'OP'], True),
-  'MIMG': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'DMASK'], True),
-  'MTBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'FORMAT', 'SOFFSET'], True),
-  'MUBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'SOFFSET'], True),
-  'SMEM': (['OP', 'SBASE', 'SDATA', 'OFFSET', 'SOFFSET'], True),
-  'SOP1': (['OP', 'SDST', 'SSRC0'], True),
-  'SOP2': (['OP', 'SDST', 'SSRC0', 'SSRC1'], True),
-  'SOPC': (['OP', 'SSRC0', 'SSRC1'], True),
-  'SOPK': (['OP', 'SDST', 'SIMM16'], True),
-  'SOPP': (['OP', 'SIMM16'], True),
-  'VINTERP': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOP1': (['OP', 'VDST', 'SRC0'], True),
-  'VOP2': (['OP', 'VDST', 'SRC0', 'VSRC1'], True),
-  'VOP3': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOP3P': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOP3SD': (['OP', 'VDST', 'SDST', 'SRC0', 'SRC1', 'SRC2'], True),
-  'VOPC': (['OP', 'SRC0', 'VSRC1'], True),
-  'VOPD': (['OPX', 'OPY', 'SRCX0', 'SRCY0', 'VDSTX', 'VDSTY'], True),
-}
-
-class TestPDFParserGenerate(unittest.TestCase):
-  """Test the PDF parser by running generate() and checking results."""
-  result: dict
-
-  @classmethod
-  def setUpClass(cls):
-    from extra.assembly.rdna3.gen import generate
-    cls.result = generate()
-
-  def test_all_formats_present(self):
-    """All expected formats should be parsed."""
-    for fmt_name in EXPECTED_FORMATS:
-      self.assertIn(fmt_name, self.result["formats"], f"missing format {fmt_name}")
-
-  def test_format_count(self):
-    """Should have exactly 23 formats."""
-    self.assertEqual(len(self.result["formats"]), 23)
-
-  def test_no_duplicate_fields(self):
-    """No format should have duplicate field names."""
-    for fmt_name, fields in self.result["formats"].items():
-      field_names = [f[0] for f in fields]
-      self.assertEqual(len(field_names), len(set(field_names)), f"{fmt_name} has duplicate fields: {field_names}")
-
-  def test_expected_fields(self):
-    """Each format should have its expected key fields."""
-    for fmt_name, (expected_fields, has_encoding) in EXPECTED_FORMATS.items():
-      fields = {f[0] for f in self.result["formats"].get(fmt_name, [])}
-      for field in expected_fields:
-        self.assertIn(field, fields, f"{fmt_name} missing {field}")
-      if has_encoding:
-        self.assertIn("ENCODING", fields, f"{fmt_name} should have ENCODING")
-      else:
-        self.assertNotIn("ENCODING", fields, f"{fmt_name} should not have ENCODING")
-
-  def test_vopd_no_dpp16_fields(self):
-    """VOPD should not have DPP16-specific fields (parser boundary bug)."""
-    vopd_fields = {f[0] for f in self.result["formats"].get("VOPD", [])}
-    for field in ['DPP_CTRL', 'BANK_MASK', 'ROW_MASK']:
-      self.assertNotIn(field, vopd_fields, f"VOPD should not have {field}")
-
-  def test_dpp16_no_vinterp_fields(self):
-    """DPP16 should not have VINTERP-specific fields."""
-    dpp16_fields = {f[0] for f in self.result["formats"].get("DPP16", [])}
-    for field in ['VDST', 'WAITEXP']:
-      self.assertNotIn(field, dpp16_fields, f"DPP16 should not have {field}")
-
-  def test_sopp_no_smem_fields(self):
-    """SOPP should not have SMEM fields (page break bug)."""
-    sopp_fields = {f[0] for f in self.result["formats"].get("SOPP", [])}
-    for field in ['SBASE', 'SDATA']:
-      self.assertNotIn(field, sopp_fields, f"SOPP should not have {field}")
-
-class TestPDFParser(unittest.TestCase):
-  """Verify format classes have correct fields from PDF parsing."""
-
-  def test_sop2_fields(self):
-    """SOP2 should have op, sdst, ssrc0, ssrc1."""
-    for field in ['op', 'sdst', 'ssrc0', 'ssrc1']:
-      self.assertIn(field, SOP2._fields)
-    self.assertEqual(SOP2._fields['op'].hi, 29)
-    self.assertEqual(SOP2._fields['op'].lo, 23)
-
-  def test_sop1_fields(self):
-    """SOP1 should have op, sdst, ssrc0 with correct bit positions."""
-    for field in ['op', 'sdst', 'ssrc0']:
-      self.assertIn(field, SOP1._fields)
-    self.assertNotIn('simm16', SOP1._fields)
-    self.assertEqual(SOP1._fields['ssrc0'].hi, 7)
-    self.assertEqual(SOP1._fields['ssrc0'].lo, 0)
-    assert SOP1._encoding is not None
-    self.assertEqual(SOP1._encoding[0].hi, 31)
-    self.assertEqual(SOP1._encoding[1], 0b101111101)
-
-  def test_vop3sd_fields(self):
-    """VOP3SD should have all fields including src0/src1/src2 from page continuation."""
-    for field in ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2']:
-      self.assertIn(field, VOP3SD._fields)
-    self.assertEqual(VOP3SD._fields['src0'].hi, 40)
-    self.assertEqual(VOP3SD._fields['src0'].lo, 32)
-    self.assertEqual(VOP3SD._size(), 8)
-
-  def test_flat_has_vdst(self):
-    """FLAT should have vdst field."""
-    self.assertIn('vdst', FLAT._fields)
-    self.assertEqual(FLAT._fields['vdst'].hi, 63)
-    self.assertEqual(FLAT._fields['vdst'].lo, 56)
-
-  def test_encoding_bits(self):
-    """Verify encoding bits are correct for major formats."""
-    tests = [
-      (SOP2, 31, 30, 0b10),
-      (SOPK, 31, 28, 0b1011),
-      (SOPP, 31, 23, 0b101111111),
-      (VOP1, 31, 25, 0b0111111),
-      (VOP2, 31, 31, 0b0),
-      (VOPC, 31, 25, 0b0111110),
-      (FLAT, 31, 26, 0b110111),
-    ]
-    for cls, hi, lo, val in tests:
-      assert cls._encoding is not None
-      self.assertEqual(cls._encoding[0].hi, hi, f"{cls.__name__} encoding hi")
-      self.assertEqual(cls._encoding[0].lo, lo, f"{cls.__name__} encoding lo")
-      self.assertEqual(cls._encoding[1], val, f"{cls.__name__} encoding val")
-
-  def test_opcode_enums_exist(self):
-    """Verify opcode enums are generated with expected counts."""
-    self.assertGreater(len(SOP1Op), 50)
-    self.assertGreater(len(SOP2Op), 50)
-    self.assertGreater(len(VOP1Op), 50)
-    self.assertGreater(len(VOP3Op), 200)
-
-  def test_vopd_no_duplicate_fields(self):
-    """VOPD should not have duplicate fields and should not include DPP16 fields."""
-    field_names = list(VOPD._fields.keys())
-    self.assertEqual(len(field_names), len(set(field_names)))
-    for field in ['srcx0', 'srcy0', 'opx', 'opy']:
-      self.assertIn(field, VOPD._fields)
-    for field in ['dpp_ctrl', 'bank_mask', 'row_mask']:
-      self.assertNotIn(field, VOPD._fields)
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_rdna3_asm.py

@@ -1,94 +0,0 @@
-#!/usr/bin/env python3
-import unittest, subprocess
-from extra.assembly.rdna3.autogen import *
-
-def llvm_assemble(asm: str) -> bytes:
-  """Assemble using llvm-mc and return bytes."""
-  result = subprocess.run(
-    ["llvm-mc", "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
-    input=asm, capture_output=True, text=True
-  )
-  out = b''
-  for line in result.stdout.split('\n'):
-    if 'encoding:' in line:
-      enc = line.split('encoding:')[1].strip()
-      enc = enc.strip('[]').replace('0x', '').replace(',', '')
-      out += bytes.fromhex(enc)
-  if not out: raise ValueError(f"no encoding found: {result.stdout} {result.stderr}")
-  return out
-
-class TestRDNA3Asm(unittest.TestCase):
-  def test_full_program(self):
-    """Test the full program from rdna3fun.py matches llvm-mc output."""
-    program = [
-      v_bfe_u32(v[1], v[0], 10, 10),
-      s_load_b128(s[4:7], s[0:1], NULL),
-      v_and_b32_e32(v[0], 0x3FF, v[0]),
-      s_mulk_i32(s[3], 0x87),
-      v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2]),
-      v_mul_u32_u24_e32(v[0], 45, v[0]),
-      v_ashrrev_i32_e32(v[2], 31, v[1]),
-      v_add3_u32(v[0], v[0], s[3], v[1]),
-      v_lshlrev_b64(v[2:3], 2, v[1:2]),
-      v_ashrrev_i32_e32(v[1], 31, v[0]),
-      v_lshlrev_b64(v[0:1], 2, v[0:1]),
-      s_waitcnt(0xfc07),  # lgkmcnt(0)
-      v_add_co_u32(v[2], VCC_LO, s[6], v[2]),
-      v_add_co_ci_u32_e32(v[3], s[7], v[3]),
-      v_add_co_u32(v[0], VCC_LO, s[4], v[0]),
-      global_load_b32(vdst=v[2], addr=v[2], saddr=OFF),
-      v_add_co_ci_u32_e32(v[1], s[5], v[1]),
-      s_waitcnt(0x03f7),  # vmcnt(0)
-      global_store_b32(addr=v[0], data=v[2], saddr=OFF),
-      s_endpgm(),
-    ]
-
-    asm = """
-v_bfe_u32 v1, v0, 10, 10
-s_load_b128 s[4:7], s[0:1], null
-v_and_b32_e32 v0, 0x3FF, v0
-s_mulk_i32 s3, 0x87
-v_mad_u64_u32 v[1:2], null, s2, 3, v[1:2]
-v_mul_u32_u24_e32 v0, 45, v0
-v_ashrrev_i32_e32 v2, 31, v1
-v_add3_u32 v0, v0, s3, v1
-v_lshlrev_b64 v[2:3], 2, v[1:2]
-v_ashrrev_i32_e32 v1, 31, v0
-v_lshlrev_b64 v[0:1], 2, v[0:1]
-s_waitcnt lgkmcnt(0)
-v_add_co_u32 v2, vcc_lo, s6, v2
-v_add_co_ci_u32_e32 v3, vcc_lo, s7, v3, vcc_lo
-v_add_co_u32 v0, vcc_lo, s4, v0
-global_load_b32 v2, v[2:3], off
-v_add_co_ci_u32_e32 v1, vcc_lo, s5, v1, vcc_lo
-s_waitcnt vmcnt(0)
-global_store_b32 v[0:1], v2, off
-s_endpgm
-"""
-    expected = llvm_assemble(asm)
-    for inst,rt in zip(program, asm.strip().split("\n")): print(f"{inst.disasm():50s} {rt}")
-    actual = b''.join(inst.to_bytes() for inst in program)
-    self.assertEqual(actual, expected)
-
-  def test_sop2_s_add_u32(self):
-    inst = SOP2(SOP2Op.S_ADD_U32, s[3], s[0], s[1])
-    expected = llvm_assemble("s_add_u32 s3, s0, s1")
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_vop2_v_and_b32_inline_const(self):
-    inst = v_and_b32_e32(v[0], 10, v[0])
-    expected = llvm_assemble("v_and_b32_e32 v0, 10, v0")
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_sopp_s_endpgm(self):
-    inst = s_endpgm()
-    expected = llvm_assemble("s_endpgm")
-    self.assertEqual(inst.to_bytes(), expected)
-
-  def test_sop1_s_mov_b32(self):
-    inst = s_mov_b32(s[0], s[1])
-    expected = llvm_assemble("s_mov_b32 s0, s1")
-    self.assertEqual(inst.to_bytes(), expected)
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rdna3/test/test_roundtrip.py

@@ -1,234 +0,0 @@
-#!/usr/bin/env python3
-"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
-import unittest, io, sys, re
-from extra.assembly.rdna3.autogen import *
-from extra.assembly.rdna3.lib import Inst
-from extra.assembly.rdna3.asm import asm
-
-# Instruction format detection based on encoding bits
-def detect_format(data: bytes) -> type[Inst] | None:
-  """Detect instruction format from machine code bytes."""
-  if len(data) < 4: return None
-  word = int.from_bytes(data[:4], 'little')
-  enc_9bit = (word >> 23) & 0x1FF  # 9-bit encoding for SOP1/SOPC/SOPP
-  enc_8bit = (word >> 24) & 0xFF
-
-  # Check 9-bit encodings first (most specific)
-  if enc_9bit == 0x17D: return SOP1  # bits 31:23 = 101111101
-  if enc_9bit == 0x17E: return SOPC  # bits 31:23 = 101111110
-  if enc_9bit == 0x17F: return SOPP  # bits 31:23 = 101111111
-  # SOPK: bits 31:28 = 1011, bits 27:23 = opcode (check after SOP1/SOPC/SOPP)
-  if enc_8bit in range(0xB0, 0xC0): return SOPK
-  # SOP2: bits 31:23 in range 0x100-0x17C (0x80-0xBE in bits 31:24, but not SOPK)
-  if 0x80 <= enc_8bit <= 0x9F: return SOP2
-  # VOP1: bits 31:25 = 0111111 (0x3F)
-  if (word >> 25) == 0x3F: return VOP1
-  # VOPC: bits 31:25 = 0111110 (0x3E)
-  if (word >> 25) == 0x3E: return VOPC
-  # VOP2: bits 31:30 = 00
-  if (word >> 30) == 0: return VOP2
-
-  # Check 64-bit formats
-  if len(data) >= 8:
-    if enc_8bit in (0xD4, 0xD5, 0xD7): return VOP3
-    if enc_8bit == 0xD6: return VOP3SD
-    if enc_8bit == 0xCC: return VOP3P
-    if enc_8bit == 0xCD: return VINTERP
-    if enc_8bit in (0xC8, 0xC9): return VOPD
-    if enc_8bit == 0xF4: return SMEM
-    if enc_8bit == 0xD8: return DS
-    if enc_8bit in (0xDC, 0xDD, 0xDE, 0xDF): return FLAT
-    if enc_8bit in (0xE0, 0xE1, 0xE2, 0xE3): return MUBUF
-    if enc_8bit in (0xE8, 0xE9, 0xEA, 0xEB): return MTBUF
-
-  return None
-
-def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
-  """Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
-  old_stdout = sys.stdout
-  sys.stdout = io.StringIO()
-  compiler.disassemble(lib)
-  output = sys.stdout.getvalue()
-  sys.stdout = old_stdout
-
-  results = []
-  for line in output.splitlines():
-    if '//' not in line: continue
-    instr = line.split('//')[0].strip()
-    if not instr: continue
-    comment = line.split('//')[1].strip()
-    if ':' not in comment: continue
-    hex_str = comment.split(':')[1].strip().split()[0]
-    try:
-      machine_bytes = bytes.fromhex(hex_str)[::-1]  # big-endian to little-endian
-      results.append((instr, machine_bytes))
-    except ValueError:
-      continue
-  return results
-
-def compile_asm(instr: str, compiler=None) -> bytes | None:
-  """Compile a single instruction with llvm-mc and return the machine code bytes."""
-  import subprocess
-  try:
-    result = subprocess.run(
-      ['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
-      input=f".text\n{instr}\n", capture_output=True, text=True)
-    if result.returncode != 0: return None
-    # Parse encoding: [0x01,0x39,0x0a,0x7e]
-    for line in result.stdout.split('\n'):
-      if 'encoding:' in line:
-        enc = line.split('encoding:')[1].strip()
-        if enc.startswith('[') and enc.endswith(']'):
-          hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
-          return bytes.fromhex(hex_vals)
-  except Exception:
-    pass
-  return None
-
-class TestTinygradKernelRoundtrip(unittest.TestCase):
-  """Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
-
-  def _test_kernel_roundtrip(self, op_fn):
-    """Generate kernel from op_fn, test:
-    1. decode -> reencode matches original bytes
-    2. asm(disasm()) matches LLVM output
-    3. our disasm() matches LLVM's disassembly string exactly
-    """
-    from extra.assembly.rdna3.test.test_compare_emulators import get_kernels_from_tinygrad
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-    kernels, _, _ = get_kernels_from_tinygrad(op_fn)
-    compiler = HIPCompiler('gfx1100')
-
-    decode_passed, decode_failed, decode_skipped = 0, 0, 0
-    asm_passed, asm_failed, asm_skipped = 0, 0, 0
-    disasm_passed, disasm_failed, disasm_skipped = 0, 0, 0
-    decode_failures, asm_failures, disasm_failures = [], [], []
-
-    for ki, kernel in enumerate(kernels):
-      offset = 0
-      while offset < len(kernel.code):
-        remaining = kernel.code[offset:]
-        fmt = detect_format(remaining)
-        if fmt is None:
-          decode_skipped += 1
-          asm_skipped += 1
-          disasm_skipped += 1
-          offset += 4
-          continue
-
-        size = fmt._size()
-        if len(remaining) < size:
-          break
-
-        orig_bytes = remaining[:size]
-
-        # Test 1: decode -> reencode roundtrip
-        try:
-          decoded = fmt.from_bytes(orig_bytes)
-          reencoded = decoded.to_bytes()
-          if reencoded[:size] == orig_bytes:
-            decode_passed += 1
-          else:
-            decode_failed += 1
-            decode_failures.append(f"K{ki}@{offset}: {decoded.disasm()}: orig={orig_bytes.hex()} reenc={reencoded[:size].hex()}")
-
-          our_disasm = decoded.disasm()
-
-          # Test 2: asm(disasm()) matches LLVM output
-          try:
-            our_bytes = asm(our_disasm).to_bytes()
-            llvm_bytes = compile_asm(our_disasm, compiler)
-            if llvm_bytes is None:
-              asm_skipped += 1
-            elif our_bytes[:len(llvm_bytes)] == llvm_bytes:
-              asm_passed += 1
-            else:
-              asm_failed += 1
-              asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': ours={our_bytes[:len(llvm_bytes)].hex()} llvm={llvm_bytes.hex()}")
-          except Exception:
-            asm_skipped += 1
-
-          # Test 3: our disasm() matches LLVM's disassembly string exactly
-          # Skip if instruction uses op_XX (unknown opcode) or looks malformed (many raw field values)
-          if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm):
-            disasm_skipped += 1
-          else:
-            try:
-              # Get LLVM's disassembly of our instruction
-              src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n  {our_disasm}\n"
-              lib = compiler.compile(src)
-              llvm_instrs = disassemble_lib(lib, compiler)
-              if llvm_instrs:
-                llvm_disasm = llvm_instrs[0][0]
-                if our_disasm == llvm_disasm:
-                  disasm_passed += 1
-                else:
-                  disasm_failed += 1
-                  disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
-              else:
-                disasm_skipped += 1
-            except Exception:
-              disasm_skipped += 1
-
-        except Exception:
-          decode_skipped += 1
-          asm_skipped += 1
-          disasm_skipped += 1
-
-        offset += size
-
-    print(f"decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
-    print(f"asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
-    print(f"disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
-    self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
-    self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
-    self.assertEqual(disasm_failed, 0, f"Disasm failures:\n" + "\n".join(disasm_failures[:20]))
-
-  # Basic unary ops
-  def test_neg(self): self._test_kernel_roundtrip(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
-  def test_relu(self): self._test_kernel_roundtrip(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
-  def test_exp(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).exp())
-  def test_log(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 3.0]).log())
-  def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
-  def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
-  def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
-
-  # Binary ops
-  def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
-  def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
-  def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
-  def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
-  def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
-
-  # Reductions
-  def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
-  def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
-  def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
-
-  # Matmul
-  def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
-  def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
-
-  # Complex ops
-  def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
-  def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
-
-  # Memory patterns
-  def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
-  def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
-  def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
-
-  # Cast ops
-  def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
-  def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
-
-  # Comparison ops
-  def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
-  def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
-
-  # Fused ops
-  def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/rocm/.gitignore

@@ -1,4 +0,0 @@
-*.deb
-build
-src
-sniffer/sniff.so

extra/assembly/rocm/NOTES

@@ -1,20 +0,0 @@
-Built ROCT-Thunk-Interface (hsakmt)
-  hsakmt-roct-dev_5.4.4.99999-local_amd64.deb
-  note: installs to /opt/rocm
-Built ROCm-Device-Libs
-  Works with ROCM_PATH=/home/tiny/build/ROCm-Device-Libs/build/dist
-  rocm-device-libs_1.0.0.99999-local_amd64.deb
-Built ROCm-CompilerSupport (amd_comgr)
-  no deb, sudo make install to /usr/local
-Built ROCR-Runtime
-  hsa-rocr_1.8.0-local_amd64.deb
-  hsa-rocr-dev_1.8.0-local_amd64.deb
-Built ROCm-OpenCL-Runtime
-  rocm-ocl-icd_2.0.0-local_amd64.deb
-  ISSUE: these depend on "comgr"
-  rocm-opencl_2.0.0-local_amd64.deb
-  rocm-opencl-dev_2.0.0-local_amd64.deb
-  Did sudo make install
-
-
-

extra/assembly/rocm/kernel_crashes/dump

@@ -1,41 +0,0 @@
-# run two "rocm-bandwidth-test" in a loop
-# amdgpu-6.0.5-1581431.20.04
-# fixed in kernel 6.2.14
-
-[   72.153646] RIP: 0010:pm_send_runlist+0x4a/0x630 [amdgpu]
-[   72.153815] Code: 30 65 48 8b 04 25 28 00 00 00 48 89 45 d0 31 c0 80 fb 01 0f 87 aa 9d 49 00 83 e3 01 0f 85 1c 05 00 00 49 8b 3f b8 01 00 00 00 <48> 8b 97 30 01 00 00 44 8b b7 6c 01 00 00 8b 9f 70 01 00 00 8b 8a
-[   72.153900] RSP: 0018:ffffb48445c03c30 EFLAGS: 00010246
-[   72.153928] RAX: 0000000000000001 RBX: 0000000000000000 RCX: 0000000000000000
-[   72.153962] RDX: 000000000000007b RSI: ffff9395e1562558 RDI: 0000000000000000
-[   72.153996] RBP: ffffb48445c03cb8 R08: 0000000000000000 R09: 0000000000000001
-[   72.154030] R10: ffff9395c900d840 R11: 0000000000000000 R12: 0000000000000000
-[   72.154065] R13: ffff9395c9e00400 R14: 0000000000000001 R15: ffff9395e15624e0
-[   72.154099] FS:  00007f345c6463c0(0000) GS:ffff93a4aee80000(0000) knlGS:0000000000000000
-[   72.154137] CS:  0010 DS: 0000 ES: 0000 CR0: 0000000080050033
-[   72.154165] CR2: 0000000000000130 CR3: 0000000112840000 CR4: 0000000000750ee0
-[   72.154201] PKRU: 55555554
-[   72.154215] Call Trace:
-[   72.154230]  <TASK>
-[   72.154244]  map_queues_cpsch+0x75/0xc0 [amdgpu]
-[   72.154365]  debug_map_and_unlock+0x51/0x90 [amdgpu]
-[   72.154480]  debug_refresh_runlist+0x1f/0x30 [amdgpu]
-[   72.154591]  kfd_dbg_runtime_disable+0x13c/0x240 [amdgpu]
-[   72.154705]  kfd_ioctl_dbg_set_debug_trap+0x69d/0x8b0 [amdgpu]
-[   72.154820]  kfd_ioctl+0x24a/0x5b0 [amdgpu]
-[   72.154925]  ? kfd_ioctl_create_queue+0x770/0x770 [amdgpu]
-[   72.155035]  ? syscall_exit_to_user_mode+0x27/0x50
-[   72.155061]  ? exit_to_user_mode_prepare+0x3d/0x1c0
-[   72.155088]  __x64_sys_ioctl+0x95/0xd0
-[   72.155109]  do_syscall_64+0x5c/0xc0
-[   72.155128]  ? syscall_exit_to_user_mode+0x27/0x50
-[   72.155151]  ? do_syscall_64+0x69/0xc0
-[   72.155172]  entry_SYSCALL_64_after_hwframe+0x61/0xcb
-[   72.155198] RIP: 0033:0x7f345c7f63ab
-[   72.155218] Code: 0f 1e fa 48 8b 05 e5 7a 0d 00 64 c7 00 26 00 00 00 48 c7 c0 ff ff ff ff c3 66 0f 1f 44 00 00 f3 0f 1e fa b8 10 00 00 00 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 8b 0d b5 7a 0d 00 f7 d8 64 89 01 48
-[   72.155301] RSP: 002b:00007ffc97cc89f8 EFLAGS: 00000246 ORIG_RAX: 0000000000000010
-[   72.155339] RAX: ffffffffffffffda RBX: 00007ffc97cc8a30 RCX: 00007f345c7f63ab
-[   72.155375] RDX: 00007ffc97cc8a30 RSI: 00000000c0284b82 RDI: 0000000000000003
-[   72.155411] RBP: 00000000c0284b82 R08: 0000000000000000 R09: 0000000000000000
-[   72.155447] R10: 00007f345cd4ddb0 R11: 0000000000000246 R12: 00007ffc97cc8a30
-[   72.155481] R13: 0000000000000003 R14: 00007ffc97cc8d20 R15: 0000000000000000
-[   72.155517]  </TASK>

extra/assembly/rocm/kernel_crashes/dump2

@@ -1,41 +0,0 @@
-# run two tinygrad matrix example in a loop
-# amdgpu-6.0.5-1581431.20.04
-# NOT fixed in kernel 6.2.14
-
-[  553.016624] gmc_v11_0_process_interrupt: 30 callbacks suppressed
-[  553.016631] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:9 pasid:32770, for process python3 pid 10001 thread python3 pid 10001)
-[  553.016790] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f0000000000 from client 10
-[  553.016892] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00901A30
-[  553.016974] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
-[  553.017051] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x0
-[  553.017111] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
-[  553.017173] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
-[  553.017238] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
-[  553.017300] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
-[  553.123921] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
-[  553.124153] amdgpu: failed to add hardware queue to MES, doorbell=0x1a16
-[  553.124195] amdgpu: MES might be in unrecoverable state, issue a GPU reset
-[  553.124237] amdgpu: Failed to restore queue 2
-[  553.124266] amdgpu: Failed to restore process queues
-[  553.124270] amdgpu: Failed to evict queue 3
-[  553.124297] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
-
-# alternative crash in kernel 6.2.14
-
-[  151.097948] gmc_v11_0_process_interrupt: 30 callbacks suppressed
-[  151.097953] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:8 pasid:32771, for process python3 pid 7525 thread python3 pid 7525)
-[  151.097993] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f0000000000 from client 10
-[  151.098008] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00801A30
-[  151.098020] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
-[  151.098032] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x0
-[  151.098042] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
-[  151.098052] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
-[  151.098062] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
-[  151.098071] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
-[  151.209517] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
-[  151.209724] amdgpu: failed to add hardware queue to MES, doorbell=0x1002
-[  151.209734] amdgpu: MES might be in unrecoverable state, issue a GPU reset
-[  151.209743] amdgpu: Failed to restore queue 1
-[  151.209751] amdgpu: Failed to restore process queues
-[  151.209759] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
-[  151.209858] amdgpu 0000:0b:00.0: amdgpu: GPU reset begin!

extra/assembly/rocm/kernel_crashes/dump3

@@ -1,20 +0,0 @@
-# two tinygrad + two bandwidth test
-# RDNA2, driver 6.0.5
-# recovered from this!
-
-[  136.971209] gmc_v10_0_process_interrupt: 39 callbacks suppressed
-[  136.971218] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:11 pasid:32773, for process rocm-bandwidth- pid 20281 thread rocm-bandwidth- pid 20281)
-[  136.971228] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f5c2b800000 from client 0x1b (UTCL2)
-[  136.971232] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00B01A31
-[  136.971233] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
-[  136.971235] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x1
-[  136.971236] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
-[  136.971236] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
-[  136.971237] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
-[  136.971238] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
-...
-[  136.993979] amdgpu 0000:0b:00.0: amdgpu: IH ring buffer overflow (0x000BE5A0, 0x0003C480, 0x0003E5C0)
-[  138.209072] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004000 flags=0x0000]
-[  138.209078] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004d80 flags=0x0000]
-[  138.209081] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005000 flags=0x0000]
-[  138.209084] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005d80 flags=0x0000]

extra/assembly/rocm/kernel_crashes/dump4

@@ -1,33 +0,0 @@
-# ROCK-Kernel-Driver 0b579de9622f5c93021dcb7927d13926313740a2
-# non fatal "crash"
-
-[  127.418045] ------------[ cut here ]------------
-[  127.418046] User pages unexpectedly invalid
-[  127.418056] WARNING: CPU: 16 PID: 260 at drivers/gpu/drm/amd/amdgpu/amdgpu_amdkfd_gpuvm.c:3000 amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
-[  127.418235] Modules linked in: rfcomm cmac algif_hash algif_skcipher af_alg bnep nls_iso8859_1 iwlmvm mac80211 intel_rapl_msr intel_rapl_common edac_mce_amd snd_hda_codec_realtek snd_hda_codec_generic snd_hda_codec_hdmi kvm_amd binfmt_misc snd_hda_intel snd_intel_dspcfg kvm libarc4 snd_intel_sdw_acpi snd_hda_codec btusb iwlwifi btrtl snd_hda_core btbcm btintel irqbypass btmtk snd_hwdep crct10dif_pclmul snd_pcm polyval_clmulni bluetooth snd_seq_midi snd_seq_midi_event snd_rawmidi snd_seq polyval_generic cfg80211 ghash_clmulni_intel eeepc_wmi snd_seq_device snd_timer aesni_intel asus_wmi ecdh_generic snd platform_profile crypto_simd ledtrig_audio cryptd ecc ccp soundcore sparse_keymap rapl k10temp wmi_bmof mac_hid sch_fq_codel msr parport_pc ppdev lp parport ramoops pstore_blk efi_pstore reed_solomon pstore_zone ip_tables x_tables autofs4 amdgpu hid_generic usbhid hid i2c_algo_bit drm_ttm_helper ttm video iommu_v2 drm_buddy gpu_sched drm_display_helper drm_kms_helper syscopyarea
-[  127.418276]  sysfillrect sysimgblt fb_sys_fops drm nvme nvme_core cec r8169 ahci crc32_pclmul rc_core i2c_piix4 xhci_pci libahci nvme_common xhci_pci_renesas realtek wmi
-[  127.418284] CPU: 16 PID: 260 Comm: kworker/16:1 Tainted: G        W          6.0.0 #4
-[  127.418286] Hardware name: System manufacturer System Product Name/TUF GAMING X570-PLUS (WI-FI), BIOS 3603 03/20/2021
-[  127.418287] Workqueue: events amdgpu_amdkfd_restore_userptr_worker [amdgpu]
-[  127.418455] RIP: 0010:amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
-[  127.418601] Code: ff e8 2b 8a 96 d1 e9 66 fe ff ff 48 c7 c7 40 4f f5 c0 e8 56 7b 8a d1 0f 0b e9 2e ff ff ff 48 c7 c7 d8 d0 ed c0 e8 43 7b 8a d1 <0f> 0b e9 0a fe ff ff 4c 89 ef e8 f8 89 96 d1 e9 cb fd ff ff e8 ce
-[  127.418603] RSP: 0018:ffffb36740a83dc8 EFLAGS: 00010282
-[  127.418604] RAX: 0000000000000000 RBX: ffff9d159ee9df30 RCX: 0000000000000027
-[  127.418605] RDX: 0000000000000027 RSI: ffffb36740a83c88 RDI: ffff9d242a220568
-[  127.418606] RBP: ffffb36740a83e58 R08: ffff9d242a220560 R09: 0000000000000001
-[  127.418607] R10: 0000000000000001 R11: 0000000000000020 R12: ffff9d159ee9df98
-[  127.418607] R13: ffff9d159ee9df70 R14: ffff9d159ee9dee0 R15: ffff9d159ee9dee0
-[  127.418608] FS:  0000000000000000(0000) GS:ffff9d242a200000(0000) knlGS:0000000000000000
-[  127.418609] CS:  0010 DS: 0000 ES: 0000 CR0: 0000000080050033
-[  127.418610] CR2: 00007fd5d4715000 CR3: 0000000120ffe000 CR4: 0000000000750ee0
-[  127.418611] PKRU: 55555554
-[  127.418611] Call Trace:
-[  127.418612]  <TASK>
-[  127.418613]  process_one_work+0x21f/0x3f0
-[  127.418615]  worker_thread+0x4a/0x3c0
-[  127.418617]  ? process_one_work+0x3f0/0x3f0
-[  127.418618]  kthread+0xf0/0x120
-[  127.418619]  ? kthread_complete_and_exit+0x20/0x20
-[  127.418620]  ret_from_fork+0x22/0x30
-[  127.418622]  </TASK>
-[  127.418623] ---[ end trace 0000000000000000 ]---

extra/assembly/rocm/rdna3/asm.py

@@ -1,80 +0,0 @@
-import numpy as np
-import pathlib
-from hexdump import hexdump
-from tinygrad.helpers import colored
-from extra.helpers import enable_early_exec
-early_exec = enable_early_exec()
-
-from tinygrad.runtime.ops_cl import CLProgram, CLBuffer, ROCM_LLVM_PATH
-
-ENABLE_NON_ASM = False
-
-WMMA = True
-DUAL_ALU = True
-F32 = True
-
-if ENABLE_NON_ASM:
-  buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
-  prg_empty = CLProgram("code", "__kernel void code(__global float *a) { a[0] = 1; }")
-  asm_real = prg_empty.binary()
-  with open("/tmp/cc.elf", "wb") as f:
-    f.write(asm_real)
-  prg_empty([1], [1], buf, wait=True)
-  print(buf.toCPU())
-
-print(colored("creating CLBuffer", "green"))
-buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
-code = open(pathlib.Path(__file__).parent / "prog.s", "r").read()
-
-gen = []
-FLOPS = 0
-MAX_REG = 251
-for j in range(1):
-  if WMMA:
-    KY, KX = 4, 4
-    for y in range(KY):
-      for x in range(KX):
-        c = (y*KX+x)*8
-        a = (KY*KX*8) + y*8
-        b = (KY*KX*8) + (KY*8) + x*8
-        gen.append(f"v_wmma_f32_16x16x16_f16 v[{c}:{c+7}], v[{a}:{a+7}], v[{b}:{b+7}], v[{c}:{c+7}]")
-        FLOPS += 16*8*2
-  else:
-    for i in range(0, MAX_REG, 6):
-      if DUAL_ALU:
-        if F32:
-          gen.append(f"v_dual_fmac_f32 v{i+0}, v{i+1}, v{i+2} :: v_dual_fmac_f32 v{i+3}, v{i+4}, v{i+5}")
-          FLOPS += 4
-        else:
-          gen.append(f"v_dual_dot2acc_f32_f16 v{i+0}, v{i+1}, v{i+2} :: v_dual_dot2acc_f32_f16 v{i+3}, v{i+4}, v{i+5}")
-          FLOPS += 8
-      else:
-        assert F32
-        gen.append(f"v_fmac_f32 v{i+0}, v{i+1}, v{i+2}")
-        gen.append(f"v_fmac_f32 v{i+3}, v{i+4}, v{i+5}")
-code = code.replace("// FLOPS", '\n'.join(gen))
-print(code)
-
-
-# fix: COMGR failed to get code object ISA name. set triple to 'amdgcn-amd-amdhsa'
-
-object = early_exec(([ROCM_LLVM_PATH / "llvm-mc", '--arch=amdgcn', '--mcpu=gfx1100', '--triple=amdgcn-amd-amdhsa', '--filetype=obj', '-'], code.encode("utf-8")))
-asm = early_exec(([ROCM_LLVM_PATH / "ld.lld", "/dev/stdin", "-o", "/dev/stdout", "--pie"], object))
-
-with open("/tmp/cc2.o", "wb") as f:
-  f.write(object)
-with open("/tmp/cc2.elf", "wb") as f:
-  f.write(asm)
-
-print(colored("creating CLProgram", "green"))
-prg = CLProgram("code", asm)
-
-print(colored("running program", "green"))
-G = 512
-FLOPS *= 100000*G*G  # loop * global_size
-for i in range(3):
-  tm = prg(buf, global_size=[G//256, G, 1], local_size=[256, 1, 1], wait=True)
-  print(f"ran in {tm*1e3:.2f} ms, {FLOPS/(tm*1e9):.2f} GFLOPS")
-
-print(colored("transferring buffer", "green"))
-print(buf.toCPU())

extra/assembly/rocm/rdna3/prog.s

@@ -1,80 +0,0 @@
-.global _start
-_start:
-.rodata
-.align 0x10
-.global code.kd
-.type code.kd,STT_OBJECT
-# amd_kernel_code_t (must be at 0x440 for kernel_code_entry_byte_offset to be right)
-code.kd:
-# amd_kernel_..., amd_machine_...
-.long 0,0,0,0
-# kernel_code_entry_byte_offset, kernel_code_prefetch_byte_offset
-.long 0x00000bc0,0x00000000,0x00000000,0x00000000
-# kernel_code_prefetch_byte_size, max_scratch_backing_memory_byte_size
-.long 0,0,0,0
-# compute_pgm_rsrc1, compute_pgm_rsrc2, kernel_code_properties, workitem_private_segment_byte_size
-.long 0x60af0000,0x0000009e,0x00000408,0x00000000
-# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_32 | AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_16_64
-# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_DX10_CLAMP | AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_IEEE_MODE
-# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT = 0xF
-# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X
-# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_KERNARG_SEGMENT_PTR = 1
-# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_RESERVED1 = 1
-.text
-.global code
-.type code,STT_FUNC
-code:
-# https://llvm.org/docs/AMDGPUUsage.html#initial-kernel-execution-state
-# s[0:1] contains the kernarg_address
-# TODO: can we use s[2:3] if this was really a wave since we only alloced 2 SGPRs?
-s_load_b64 s[2:3], s[0:1], null
-
-s_mov_b32 s8, 0
-loop:
-s_addk_i32 s8, 1
-s_cmp_eq_u32 s8, 100000
-// FLOPS
-s_cbranch_scc0 loop
-
-# wait for the s_load_b64
-s_waitcnt lgkmcnt(0)
-
-v_dual_mov_b32 v0, 4 :: v_dual_mov_b32 v1, 2.0
-global_store_b32 v0, v1, s[2:3]
-
-# Deallocate all VGPRs for this wave. Use only when next instruction is S_ENDPGM.
-s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
-s_endpgm
-s_code_end
-
-.amdgpu_metadata
-amdhsa.kernels:
-  - .args:
-      - .address_space:  global
-        .name:           a
-        .offset:         0
-        .size:           8
-        .type_name:      'float*'
-        .value_kind:     global_buffer
-    .group_segment_fixed_size: 0
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 8
-    .language:       OpenCL C
-    .language_version:
-      - 1
-      - 2
-    .max_flat_workgroup_size: 256
-    .name:           code
-    .private_segment_fixed_size: 0
-    .sgpr_count:     2
-    .sgpr_spill_count: 0
-    .symbol:         code.kd
-    .uses_dynamic_stack: false
-    .vgpr_count:     256
-    .vgpr_spill_count: 0
-    .wavefront_size: 32
-amdhsa.target:   amdgcn-amd-amdhsa--gfx1100
-amdhsa.version:
-  - 1
-  - 2
-.end_amdgpu_metadata

extra/assembly/rocm/rocm_clone.sh

@@ -1,11 +0,0 @@
-#!/bin/bash
-mkdir -p src
-cd src
-git clone https://github.com/RadeonOpenCompute/ROCT-Thunk-Interface.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/ROCm-Device-Libs.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/llvm-project.git -b rocm-5.5.0 --depth 1
-git clone https://github.com/RadeonOpenCompute/ROCR-Runtime.git -b rocm-5.5.0
-git clone https://github.com/ROCm-Developer-Tools/ROCclr.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/ROCm-CompilerSupport.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/ROCm-OpenCL-Runtime.git -b rocm-5.5.0
-cd ../

extra/assembly/rocm/rocm_from_scratch.sh

@@ -1,69 +0,0 @@
-#!/bin/bash
-mkdir -p build/debs
-cd build
-
-# ROCT-Thunk-Interface (hsakmt)
-if [ ! -f debs/hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ]
-then
-  mkdir -p ROCT-Thunk-Interface
-  cd ROCT-Thunk-Interface
-  cmake ../../src/ROCT-Thunk-Interface
-  make -j32 package
-  cp hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ../debs
-  cd ../
-fi
-
-
-# build custom LLVM
-if [ ! -f llvm-project/bin/clang ]
-then
-  mkdir -p llvm-project
-  cd llvm-project
-  cmake -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_PROJECTS="llvm;clang;lld" -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" ../../src/llvm-project/llvm
-  make -j32
-  cd ..
-fi
-
-# use custom LLVM
-export PATH="$PWD/llvm-project/bin:$PATH"
-
-# ROCm-Device-Libs
-if [ ! -f debs/rocm-device-libs_1.0.0.99999-local_amd64.deb ]
-then
-  mkdir -p ROCm-Device-Libs
-  cd ROCm-Device-Libs
-  cmake ../../src/ROCm-Device-Libs
-  make -j32 package
-  cp rocm-device-libs_1.0.0.99999-local_amd64.deb ../debs
-  cd ../
-fi
-
-# ROCR-Runtime
-if [ ! -f debs/hsa-rocr_1.8.0-local_amd64.deb ]
-then
-  mkdir -p ROCR-Runtime
-  cd ROCR-Runtime
-  cmake ../../src/ROCR-Runtime/src
-  make -j32 package
-  cp hsa-rocr_1.8.0-local_amd64.deb ../debs
-  cp hsa-rocr-dev_1.8.0-local_amd64.deb ../debs
-  cd ../
-fi
-
-# ROCm-OpenCL-Runtime (needs ROCclr)
-if [ ! -f debs/rocm-opencl_2.0.0-local_amd64.deb ]
-then
-  mkdir -p ROCm-OpenCL-Runtime
-  cd ROCm-OpenCL-Runtime
-  cmake ../../src/ROCm-OpenCL-Runtime
-  make -j32 package
-  cp rocm-opencl_2.0.0-local_amd64.deb ../debs
-  cp rocm-opencl-dev_2.0.0-local_amd64.deb ../debs
-  cp rocm-ocl-icd_2.0.0-local_amd64.deb ../debs
-fi
-
-# ROCm-CompilerSupport (broken)
-#mkdir -p ROCm-CompilerSupport
-#cd ROCm-CompilerSupport
-#cmake ../../src/ROCm-CompilerSupport/lib/comgr
-#make -j32

extra/assembly/rocm/rocm_setup.sh

@@ -1,14 +0,0 @@
-#!/bin/bash
-rm amdgpu-install_5.5.50500-1_all.deb
-wget https://repo.radeon.com/amdgpu-install/5.5/ubuntu/$(lsb_release -cs)/amdgpu-install_5.5.50500-1_all.deb
-sudo dpkg -i amdgpu-install_5.5.50500-1_all.deb
-sudo apt-get update
-
-# kernel driver
-sudo apt-get install amdgpu-dkms
-
-# for opencl
-sudo apt-get install rocm-opencl-runtime
-
-# for HIP
-sudo apt-get install hip-runtime-amd rocm-device-libs hip-dev

extra/assembly/rocm/sniffer/build.sh

@@ -1,11 +0,0 @@
-#!/bin/bash -e
-clang sniff.cc -Werror -shared -fPIC -I../src/ -I../src/ROCT-Thunk-Interface/include -I../src/ROCm-Device-Libs/ockl/inc -o sniff.so -lstdc++
-#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so /home/tiny/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
-#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so $HOME/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
-#AMD_LOG_LEVEL=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
-DEBUG=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
-#AMD_LOG_LEVEL=5 HSAKMT_DEBUG_LEVEL=7 DEBUG=5 LD_PRELOAD=$PWD/sniff.so strace -F python3 ../rdna3/asm.py
-#LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
-#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so FORWARD_ONLY=1 DEBUG=2 python3 ../../../test/test_ops.py TestOps.test_add
-#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 0 -d 1 -m 1
-#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 1 -d 2 -m 1

extra/assembly/rocm/sniffer/sniff.cc

@@ -1,282 +0,0 @@
-// template copied from https://github.com/geohot/cuda_ioctl_sniffer/blob/master/sniff.cc
-
-#include <stdio.h>
-#include <stdlib.h>
-#include <string.h>
-#include <dlfcn.h>
-#include <signal.h>
-#include <ucontext.h>
-
-#include <sys/mman.h>
-
-// includes from the ROCm sources
-#include <linux/kfd_ioctl.h>
-#include <hsa.h>
-#include <amd_hsa_kernel_code.h>
-#include <ROCR-Runtime/src/core/inc/sdma_registers.h>
-using namespace rocr::AMD;
-
-#include <string>
-#include <map>
-std::map<int, std::string> files;
-std::map<uint64_t, uint64_t> ring_base_addresses;
-
-#define D(args...) fprintf(stderr, args)
-
-uint64_t doorbell_offset = -1;
-std::map<uint64_t, int> queue_types;
-
-void hexdump(void *d, int l) {
-  for (int i = 0; i < l; i++) {
-    if (i%0x10 == 0 && i != 0) printf("\n");
-    if (i%0x10 == 8) printf(" ");
-    if (i%0x10 == 0) printf("%8X: ", i);
-    printf("%2.2X ", ((uint8_t*)d)[i]);
-  }
-  printf("\n");
-}
-
-extern "C" {
-
-// https://defuse.ca/online-x86-assembler.htm#disassembly2
-static void handler(int sig, siginfo_t *si, void *unused) {
-  ucontext_t *u = (ucontext_t *)unused;
-  uint8_t *rip = (uint8_t*)u->uc_mcontext.gregs[REG_RIP];
-
-  int store_size = 0;
-  uint64_t value;
-  if (rip[0] == 0x48 && rip[1] == 0x89 && rip[2] == 0x30) {
-    // 0:  48 89 30                mov    QWORD PTR [rax],rsi
-    store_size = 8;
-    value = u->uc_mcontext.gregs[REG_RSI];
-    u->uc_mcontext.gregs[REG_RIP] += 3;
-  } else if (rip[0] == 0x4c && rip[1] == 0x89 && rip[2] == 0x28) {
-    // 0:  4c 89 28                mov    QWORD PTR [rax],r13
-    store_size = 8;
-    value = u->uc_mcontext.gregs[REG_R13];
-    u->uc_mcontext.gregs[REG_RIP] += 3;
-  } else {
-    D("segfault %02X %02X %02X %02X %02X %02X %02X %02X rip: %p addr: %p\n", rip[0], rip[1], rip[2], rip[3], rip[4], rip[5], rip[6], rip[7], rip, si->si_addr);
-    D("rax: %llx rcx: %llx rdx: %llx rsi: %llx rbx: %llx\n", u->uc_mcontext.gregs[REG_RAX], u->uc_mcontext.gregs[REG_RCX], u->uc_mcontext.gregs[REG_RDX], u->uc_mcontext.gregs[REG_RSI], u->uc_mcontext.gregs[REG_RBX]);
-    exit(-1);
-  }
-
-  uint64_t ring_base_address = ring_base_addresses[((uint64_t)si->si_addr)&0xFFF];
-  int queue_type = queue_types[((uint64_t)si->si_addr)&0xFFF];
-  D("%16p: \u001b[31mDING DONG\u001b[0m (queue_type %d) store(%d): 0x%8lx -> %p ring_base_address:0x%lx\n", rip, queue_type, store_size, value, si->si_addr, ring_base_address);
-
-  if (queue_type == KFD_IOC_QUEUE_TYPE_SDMA) {
-    uint8_t *sdma_ptr = (uint8_t*)(ring_base_address);
-    while (sdma_ptr < ((uint8_t*)(ring_base_address)+value)) {
-      D("0x%3lx: ", sdma_ptr-(uint8_t*)(ring_base_address));
-      if (sdma_ptr[0] == SDMA_OP_TIMESTAMP) {
-        D("SDMA_PKT_TIMESTAMP\n");
-        sdma_ptr += sizeof(SDMA_PKT_TIMESTAMP);
-      } else if (sdma_ptr[0] == SDMA_OP_GCR) {
-        D("SDMA_PKT_GCR\n");
-        sdma_ptr += sizeof(SDMA_PKT_GCR);
-      } else if (sdma_ptr[0] == SDMA_OP_ATOMIC) {
-        D("SDMA_PKT_ATOMIC\n");
-        sdma_ptr += sizeof(SDMA_PKT_ATOMIC);
-      } else if (sdma_ptr[0] == SDMA_OP_FENCE) {
-        D("SDMA_PKT_FENCE\n");
-        sdma_ptr += sizeof(SDMA_PKT_FENCE);
-      } else if (sdma_ptr[0] == SDMA_OP_TRAP) {
-        D("SDMA_PKT_TRAP\n");
-        sdma_ptr += sizeof(SDMA_PKT_TRAP);
-      } else if (sdma_ptr[0] == SDMA_OP_COPY && sdma_ptr[1] == SDMA_SUBOP_COPY_LINEAR) {
-        SDMA_PKT_COPY_LINEAR *pkt = (SDMA_PKT_COPY_LINEAR *)sdma_ptr;
-        D("SDMA_PKT_COPY_LINEAR: count:0x%x src:0x%lx dst:0x%lx\n", pkt->COUNT_UNION.count+1,
-          (uint64_t)pkt->SRC_ADDR_LO_UNION.src_addr_31_0 | ((uint64_t)pkt->SRC_ADDR_HI_UNION.src_addr_63_32 << 32),
-          (uint64_t)pkt->DST_ADDR_LO_UNION.dst_addr_31_0 | ((uint64_t)pkt->DST_ADDR_HI_UNION.dst_addr_63_32 << 32)
-        );
-        sdma_ptr += sizeof(SDMA_PKT_COPY_LINEAR);
-      } else {
-        D("unhandled packet type %d %d, exiting\n", sdma_ptr[0], sdma_ptr[1]);
-        break;
-      }
-    }
-
-    //hexdump((void*)(ring_base_address), 0x100);
-  } else if (queue_type == KFD_IOC_QUEUE_TYPE_COMPUTE_AQL) {
-    hsa_kernel_dispatch_packet_t *pkt = (hsa_kernel_dispatch_packet_t *)(ring_base_address+value*0x40);
-    if ((pkt->header&0xFF) == HSA_PACKET_TYPE_KERNEL_DISPATCH) {
-      D("HSA_PACKET_TYPE_KERNEL_DISPATCH -- setup:%d workgroup[%d, %d, %d] grid[%d, %d, %d] kernel_object:0x%lx kernarg_address:%p\n", pkt->setup, pkt->workgroup_size_x, pkt->workgroup_size_y, pkt->workgroup_size_z, pkt->grid_size_x, pkt->grid_size_y, pkt->grid_size_z, pkt->kernel_object, pkt->kernarg_address);
-      amd_kernel_code_t *code = (amd_kernel_code_t *)pkt->kernel_object;
-      D("kernel_code_entry_byte_offset:%lx\n", code->kernel_code_entry_byte_offset);
-      uint32_t *kernel_code = (uint32_t*)(pkt->kernel_object + code->kernel_code_entry_byte_offset);
-      int code_len = 0;
-      while (kernel_code[code_len] != 0xbf9f0000 && kernel_code[code_len] != 0) code_len++;
-      hexdump(kernel_code, code_len*4);
-      /*FILE *f = fopen("/tmp/kernel_code", "wb");
-      fwrite(kernel_code, 4, code_len, f);
-      fclose(f);
-      system("python -c 'print(\" \".join([(\"0x%02X\"%x) for x in open(\"/tmp/kernel_code\", \"rb\").read()]))' | ../build/llvm-project/bin/llvm-mc --disassemble --arch=amdgcn --mcpu=gfx1100 --show-encoding");*/
-      D("kernargs (kernarg_segment_byte_size:0x%lx)\n", code->kernarg_segment_byte_size);
-      // get length
-      int i;
-      for (i = 0; i < 0x400; i+=0x10) {
-        if (memcmp((void*)((uint64_t)pkt->kernarg_address+i), "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", 0x10) == 0) break;
-      }
-      hexdump((void*)pkt->kernarg_address, i+0x10);
-    } else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_BARRIER_AND) {
-      hsa_barrier_and_packet_t *pkt_and = (hsa_barrier_and_packet_t *)(ring_base_address+value*0x40);
-      D("HSA_PACKET_TYPE_BARRIER_AND completion_signal:0x%lx\n", pkt_and->completion_signal.handle);
-      //hexdump((void*)(ring_base_address+value*0x40), 0x40);
-    } else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_VENDOR_SPECIFIC) {
-      D("HSA_PACKET_TYPE_VENDOR_SPECIFIC\n");
-      hexdump((void*)(ring_base_address+value*0x40), 0x40);
-    } else {
-      hexdump((void*)(ring_base_address+value*0x40), 0x40);
-    }
-  }
-
-  mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_READ | PROT_WRITE);
-  if (store_size == 8) {
-    *(volatile uint64_t*)(si->si_addr) = value;
-  } else if (store_size == 4) {
-    *(volatile uint32_t*)(si->si_addr) = value;
-  } else if (store_size == 2) {
-    *(volatile uint16_t*)(si->si_addr) = value;
-  } else {
-    D("store size not supported\n");
-    exit(-1);
-  }
-  mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_NONE);
-}
-
-void register_sigsegv_handler() {
-  struct sigaction sa = {0};
-  sa.sa_flags = SA_SIGINFO;
-  sigemptyset(&sa.sa_mask);
-  sa.sa_sigaction = handler;
-  if (sigaction(SIGSEGV, &sa, NULL) == -1) {
-    D("ERROR: failed to register sigsegv handler");
-    exit(-1);
-  }
-  // NOTE: python (or ocl runtime?) blocks the SIGSEGV signal
-  sigset_t x;
-  sigemptyset(&x);
-  sigaddset(&x, SIGSEGV);
-  sigprocmask(SIG_UNBLOCK, &x, NULL);
-}
-
-int (*my_open)(const char *pathname, int flags, mode_t mode);
-#undef open
-int open(const char *pathname, int flags, mode_t mode) {
-  if (my_open == NULL) my_open = reinterpret_cast<decltype(my_open)>(dlsym(RTLD_NEXT, "open"));
-  int ret = my_open(pathname, flags, mode);
-  //D("open %s (0o%o) = %d\n", pathname, flags, ret);
-  files[ret] = pathname;
-  return ret;
-}
-
-
-int (*my_open64)(const char *pathname, int flags, mode_t mode);
-#undef open
-int open64(const char *pathname, int flags, mode_t mode) {
-  if (my_open64 == NULL) my_open64 = reinterpret_cast<decltype(my_open64)>(dlsym(RTLD_NEXT, "open64"));
-  int ret = my_open64(pathname, flags, mode);
-  //D("open %s (0o%o) = %d\n", pathname, flags, ret);
-  files[ret] = pathname;
-  return ret;
-}
-
-void *(*my_mmap)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
-#undef mmap
-void *mmap(void *addr, size_t length, int prot, int flags, int fd, off_t offset) {
-  if (my_mmap == NULL) my_mmap = reinterpret_cast<decltype(my_mmap)>(dlsym(RTLD_NEXT, "mmap"));
-  void *ret = my_mmap(addr, length, prot, flags, fd, offset);
-
-  if (doorbell_offset != -1 && offset == doorbell_offset) {
-    D("HIDDEN DOORBELL %p, handled by %p\n", addr, handler);
-    register_sigsegv_handler();
-    mprotect(addr, length, PROT_NONE);
-  }
-
-  if (fd != -1) D("mmapped %p (target %p) with flags 0x%x length 0x%zx fd %d %s offset 0x%lx\n", ret, addr, flags, length, fd, files[fd].c_str(), offset);
-  return ret;
-}
-
-void *(*my_mmap64)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
-#undef mmap64
-void *mmap64(void *addr, size_t length, int prot, int flags, int fd, off_t offset) { return mmap(addr, length, prot, flags, fd, offset); }
-
-int ioctl_num = 1;
-int (*my_ioctl)(int filedes, unsigned long request, void *argp) = NULL;
-#undef ioctl
-int ioctl(int filedes, unsigned long request, void *argp) {
-  if (my_ioctl == NULL) my_ioctl = reinterpret_cast<decltype(my_ioctl)>(dlsym(RTLD_NEXT, "ioctl"));
-  int ret = 0;
-  ret = my_ioctl(filedes, request, argp);
-  if (!files.count(filedes)) return ret;
-
-  uint8_t type = (request >> 8) & 0xFF;
-  uint8_t nr = (request >> 0) & 0xFF;
-  uint16_t size = (request >> 16) & 0xFFF;
-
-  D("%3d: %d = %3d(%20s) 0x%3x ", ioctl_num, ret, filedes, files[filedes].c_str(), size);
-
-  if (request == AMDKFD_IOC_SET_EVENT) {
-    kfd_ioctl_set_event_args *args = (kfd_ioctl_set_event_args *)argp;
-    D("AMDKFD_IOC_SET_EVENT event_id:%d", args->event_id);
-  } else if (request == AMDKFD_IOC_ALLOC_MEMORY_OF_GPU) {
-    kfd_ioctl_alloc_memory_of_gpu_args *args = (kfd_ioctl_alloc_memory_of_gpu_args *)argp;
-    D("AMDKFD_IOC_ALLOC_MEMORY_OF_GPU va_addr:0x%llx size:0x%llx handle:%llX gpu_id:0x%x", args->va_addr, args->size, args->handle, args->gpu_id);
-  } else if (request == AMDKFD_IOC_MAP_MEMORY_TO_GPU) {
-    kfd_ioctl_map_memory_to_gpu_args *args = (kfd_ioctl_map_memory_to_gpu_args *)argp;
-    D("AMDKFD_IOC_MAP_MEMORY_TO_GPU handle:%llX", args->handle);
-  } else if (request == AMDKFD_IOC_CREATE_EVENT) {
-    kfd_ioctl_create_event_args *args = (kfd_ioctl_create_event_args *)argp;
-    D("AMDKFD_IOC_CREATE_EVENT event_page_offset:0x%llx event_type:%d event_id:%d", args->event_page_offset, args->event_type, args->event_id);
-  } else if (request == AMDKFD_IOC_WAIT_EVENTS) {
-    D("AMDKFD_IOC_WAIT_EVENTS");
-  } else if (request == AMDKFD_IOC_SET_XNACK_MODE) {
-    D("AMDKFD_IOC_SET_XNACK_MODE");
-  } else if (request == AMDKFD_IOC_SVM || (type == 0x4b && nr == 0x20)) {
-    // NOTE: this one is variable length
-    kfd_ioctl_svm_args *args = (kfd_ioctl_svm_args *)argp;
-    D("AMDKFD_IOC_SVM start_addr:0x%llx size:0x%llx op:%d", args->start_addr, args->size, args->op);
-  } else if (request == AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU) {
-    kfd_ioctl_unmap_memory_from_gpu_args *args = (kfd_ioctl_unmap_memory_from_gpu_args *)argp;
-    D("AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU handle:%llX", args->handle);
-  } else if (request == AMDKFD_IOC_FREE_MEMORY_OF_GPU) {
-    D("AMDKFD_IOC_FREE_MEMORY_OF_GPU");
-  } else if (request == AMDKFD_IOC_SET_SCRATCH_BACKING_VA) {
-    D("AMDKFD_IOC_SET_SCRATCH_BACKING_VA");
-  } else if (request == AMDKFD_IOC_GET_TILE_CONFIG) {
-    D("AMDKFD_IOC_GET_TILE_CONFIG");
-  } else if (request == AMDKFD_IOC_SET_TRAP_HANDLER) {
-    D("AMDKFD_IOC_SET_TRAP_HANDLER");
-  } else if (request == AMDKFD_IOC_GET_VERSION) {
-    kfd_ioctl_get_version_args *args = (kfd_ioctl_get_version_args *)argp;
-    D("AMDKFD_IOC_GET_VERSION major_version:%d minor_version:%d", args->major_version, args->minor_version);
-  } else if (request == AMDKFD_IOC_GET_PROCESS_APERTURES_NEW) {
-    D("AMDKFD_IOC_GET_PROCESS_APERTURES_NEW");
-  } else if (request == AMDKFD_IOC_ACQUIRE_VM) {
-    D("AMDKFD_IOC_ACQUIRE_VM");
-  } else if (request == AMDKFD_IOC_SET_MEMORY_POLICY) {
-    D("AMDKFD_IOC_SET_MEMORY_POLICY");
-  } else if (request == AMDKFD_IOC_GET_CLOCK_COUNTERS) {
-    D("AMDKFD_IOC_GET_CLOCK_COUNTERS");
-  } else if (request == AMDKFD_IOC_CREATE_QUEUE) {
-    kfd_ioctl_create_queue_args *args = (kfd_ioctl_create_queue_args *)argp;
-    D("AMDKFD_IOC_CREATE_QUEUE\n");
-    D("queue_type:%d ring_base_address:0x%llx\n", args->queue_type, args->ring_base_address);
-    D("eop_buffer_address:0x%llx ctx_save_restore_address:0x%llx\n", args->eop_buffer_address, args->ctx_save_restore_address);
-    D("ring_size:0x%x queue_priority:%d\n", args->ring_size, args->queue_priority);
-    D("RETURNS write_pointer_address:0x%llx read_pointer_address:0x%llx doorbell_offset:0x%llx queue_id:%d\n", args->write_pointer_address, args->read_pointer_address, args->doorbell_offset, args->queue_id);
-    //D("RETURNS *write_pointer_address:0x%llx *read_pointer_address:0x%llx\n", *(uint64_t*)args->write_pointer_address, *(uint64_t*)args->read_pointer_address);
-    ring_base_addresses[args->doorbell_offset&0xFFF] = args->ring_base_address;
-    queue_types[args->doorbell_offset&0xFFF] = args->queue_type;
-    doorbell_offset = args->doorbell_offset&~0xFFF;
-  } else {
-    D("type:0x%x nr:0x%x size:0x%x", type, nr, size);
-  }
-
-  D("\n");
-  ioctl_num++;
-  return ret;
-}
-
-}

extra/export_model.py

@@ -67,12 +67,11 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
   forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs))
 
   if not wasm:
-    thread_id = 0 # NOTE: export does not support threading, thread_id is always 0
     for name,cl in bufs_to_save.items():
       weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
       cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
     cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
-    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)}, {thread_id});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
     return '\n'.join(headers + cprog)
   else:
     if bufs_to_save:

extra/fp8/fp8_linear.py

@@ -0,0 +1,101 @@
+from typing import Callable, Any
+from tinygrad import Tensor, dtypes, nn, UOp
+from tinygrad.uop.ops import KernelInfo, AxisType, Ops
+
+def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
+  fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
+  fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
+  x_abs_max = x.abs().max().detach()
+  scale = fp8_max / (x_abs_max + 1e-8)
+  x_scaled = x * scale
+  x_det = x_scaled.detach()
+  x_clamped = x_det.clamp(fp8_min, fp8_max)
+  x_clamped_ste = x_scaled + (x_clamped - x_det)
+  res = x_clamped_ste.cast(dtype)
+  return res, scale.float().reciprocal()
+
+def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
+  SEQ = inp.shape[1]
+  OUT = weight.shape[0]
+  IN = weight.shape[-1]
+  seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
+  out_idx = UOp.range(OUT, 3, AxisType.LOOP)
+  batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
+  reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
+  product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
+  reduced = product.reduce(reduce_idx, arg=Ops.ADD)
+  store_op = output.index((seq_idx*OUT+out_idx+batch_idx*OUT*SEQ), ptr=True).store(reduced).end(batch_idx, seq_idx, out_idx)
+  return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))
+
+def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
+  _, input_uop, weight_uop = kernel.src[1:]
+  input_tensor = Tensor(input_uop, device=input_uop.device)
+  grad_tensor = Tensor(gradient, device=gradient.device)
+  weight_tensor = Tensor(weight_uop, device=weight_uop.device)
+  grad_quantized, scale = quantize_to_fp8(grad_tensor)
+  scale_scalar = scale.reshape(())
+  grad_weight = Tensor.einsum("bso,bsi->oi", grad_quantized, input_tensor, dtype=dtypes.float)
+  grad_weight = grad_weight * scale_scalar
+  grad_2d = grad_quantized.reshape(grad_tensor.shape[0] * grad_tensor.shape[1], grad_tensor.shape[-1])
+  grad_input = (grad_2d.dot(weight_tensor, dtype=dtypes.float)).contiguous().reshape(input_tensor.shape) * scale
+  return (None, grad_input.uop, grad_weight.uop)
+
+class FP8Linear:
+  def __init__(self, in_features:int, out_features:int, bias:bool=True):
+    self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.float32)
+    self.bias = Tensor.empty(out_features, dtype=dtypes.float32) if bias else None
+
+  def __call__(self, x: Tensor) -> Tensor:
+    original_ndim = len(x.shape)
+    if original_ndim == 2: x = x.reshape(x.shape[0], 1, x.shape[1])
+    batch, seq, _ = x.shape
+    w_fp8, w_scale = quantize_to_fp8(self.weight)
+    x_fp8, x_scale = quantize_to_fp8(x)
+    GPUS = self.weight.device
+    if isinstance(GPUS, tuple) and len(GPUS) > 1:
+      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.multi(0), device=GPUS)
+    else:
+      y = Tensor.empty((batch, seq, self.weight.shape[0]), dtype=dtypes.float)
+    y = Tensor.custom_kernel(y, x_fp8, w_fp8, fxn=custom_matmul, grad_fxn=custom_matmul_backward)[0]
+    y = y * w_scale * x_scale
+    if self.bias is not None: y = y + self.bias
+    if original_ndim == 2: y = y.reshape(batch, self.weight.shape[0])
+    return y.cast(x.dtype)
+
+def _replace_linear(layer: nn.Linear):
+  fp8_linear = FP8Linear(layer.weight.shape[1], layer.weight.shape[0], layer.bias is not None)
+  fp8_linear.weight = layer.weight
+  if layer.bias is not None: fp8_linear.bias = layer.bias
+  return fp8_linear
+
+def _swap_linear_with_fp8(model, module_filter_fn:Callable[[Any, str],bool]|None=None, fqn:str="", parent:Any|None=None,
+                          attr_name:str="", visited:set|None=None):
+  if visited is None: visited = set()
+  if id(model) in visited: return
+  visited.add(id(model))
+  if isinstance(model, (str, int, float, bool, type(None), Tensor, UOp)): return
+  elif isinstance(model, nn.Linear):
+    if module_filter_fn is not None and not module_filter_fn(model, fqn): return
+    fp8_linear = _replace_linear(model)
+    if parent is not None and attr_name:
+      setattr(parent, attr_name, fp8_linear)
+  elif isinstance(model, list):
+    for i, item in enumerate(model):
+      child_fqn = f"{fqn}.{i}" if fqn else str(i)
+      if isinstance(item, nn.Linear) and (module_filter_fn is None or module_filter_fn(item, child_fqn)): model[i] = _replace_linear(item)
+      else: _swap_linear_with_fp8(item, module_filter_fn, child_fqn, None, "", visited)
+  elif isinstance(model, dict):
+    for key, item in list(model.items()):
+      child_fqn = f"{fqn}.{key}" if fqn else str(key)
+      if isinstance(item, nn.Linear) and (module_filter_fn is None or module_filter_fn(item, child_fqn)): model[key] = _replace_linear(item)
+      else: _swap_linear_with_fp8(item, module_filter_fn, child_fqn, None, "", visited)
+  elif hasattr(model, "__dict__"):
+    for attr_key in list(vars(model).keys()):
+      try: attr = getattr(model, attr_key)
+      except Exception: continue
+      child_fqn = f"{fqn}.{attr_key}" if fqn else attr_key
+      _swap_linear_with_fp8(attr, module_filter_fn, child_fqn, model, attr_key, visited)
+
+def convert_to_float8_training(model, module_filter_fn:Callable[[Any,str],bool]|None=None):
+  _swap_linear_with_fp8(model, module_filter_fn, "", None, "")
+  return model

extra/gemm/amd_asm_matmul.py

@@ -0,0 +1,496 @@
+# RDNA3 128x128 tiled GEMM kernel - DSL version
+# Computes C = A @ B for NxN float32 matrices using 128x128 tiles
+#
+# Architecture: RDNA3 (gfx1100)
+# Tile size: 128x128 (each workgroup computes one tile of C)
+# Workgroup: 128 threads (arranged as 32x4 for coalesced memory access)
+# Inner loop: 8 iterations per K-block, processing 8 columns of A and 8 rows of B
+#
+# Accumulators: 128 vgprs (v[2-129])
+
+import numpy as np
+from tinygrad import Tensor, Device, Context, GlobalCounters
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.helpers import getenv, colored
+from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.engine.realize import Estimates
+from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL
+from tinygrad.runtime.autogen.amd.rdna3.ins import *
+
+# =============================================================================
+# Kernel constants
+# =============================================================================
+LDS_SIZE = 8320       # Local data share size in bytes
+LDS_A_STRIDE = 0x210  # LDS stride for A tile (528 bytes)
+LDS_B_STRIDE = 0x200  # LDS stride for B tile (512 bytes)
+LDS_BASE_OFFSET = 0x1080  # Base LDS offset for tiles
+ADDR_MASK = 0x3fffff80    # Address alignment mask
+
+# =============================================================================
+# Named register assignments (VGPRs)
+# =============================================================================
+V_LANE_ID = 0             # lane_id set on startup
+# Use tile gaps (v146-159) for named regs to minimize max VGPR
+V_LANE_ID_MOD8 = 146      # lane_id & 7
+V_LANE_MOD8_X4 = 147      # (lane_id & 7) << 2
+V_LANE_DIV8_X4 = 150      # ((lane_id >> 3) & 3) << 2
+V_LDS_B_BASE = 151        # LDS B-tile base address for inner loop
+V_LDS_A_BASE = 154        # LDS A-tile base address for inner loop
+V_GLOBAL_A_ADDR = 155     # global memory A prefetch address
+V_GLOBAL_B_ADDR = 158     # global memory B prefetch address
+V_LDS_A_ADDR = 159        # single base register for A stores
+V_LDS_B_ADDR = 162        # single base register for B stores
+
+# LDS tile register destinations - SEPARATE from DATA to avoid overlap
+# A on banks 2-3, B on banks 0-1 to avoid bank conflicts in VOPD
+V_A_TILE_REGS = [130, 134, 138, 142]  # A tile: banks 2,2,2,2 (130%4=2, etc.)
+V_B_TILE_REGS = [132, 136, 140, 144, 148, 152, 156, 160]  # B tile: banks 0,0,0,0,0,0,0,0
+
+# =============================================================================
+# Named register assignments (SGPRs)
+# =============================================================================
+S_OUT_PTR = (0, 1)        # output C matrix base pointer
+S_WORKGROUP_X = 2         # workgroup_id_x (system SGPR, follows user SGPRs)
+S_WORKGROUP_Y = 3         # workgroup_id_y (system SGPR)
+S_DIM_N = 4               # matrix dimension N
+S_LOOP_BOUND = 7          # K-8 (loop termination bound)
+S_LOOP_CTR = 12           # loop counter (increments by 8)
+S_PREFETCH_FLAG = 13      # prefetch condition flag / row stride in epilogue
+S_TILE_X = 14             # workgroup_x << 7
+S_TILE_Y = 15             # workgroup_y << 7
+# Kernarg load destinations
+S_KERNARG_A = (20, 21)    # A pointer from kernarg
+S_KERNARG_B = (22, 23)    # B pointer from kernarg
+# Prefetch base pointers (8 pairs each, B: N*4 bytes apart, A: N*64 bytes apart)
+S_PREFETCH_B = 24         # s[24:39] - 8 B tile pointers
+S_PREFETCH_A = 40         # s[40:55] - 8 A tile pointers
+
+# =============================================================================
+# Data tables
+# =============================================================================
+
+# Accumulator grid: ACC_GRID[a_idx][b_idx] = vgpr for C[a,b]
+# a_idx: which A value (0-7), b_idx: which B value (0-15)
+# Scattered due to VOPD bank constraints (vdst_x % 4 != vdst_y % 4)
+# Range is from v2 - v129
+ACC_GRID = [
+  [  5,  3,  9,  8,   37, 35, 41, 40,   69, 67, 73, 72,  101, 99,105,104],  # a0
+  [  4,  2,  7,  6,   36, 34, 39, 38,   68, 66, 71, 70,  100, 98,103,102],  # a1
+  [ 17, 16, 13, 11,   49, 48, 45, 43,   81, 80, 77, 75,  113,112,109,107],  # a2
+  [ 15, 14, 12, 10,   47, 46, 44, 42,   79, 78, 76, 74,  111,110,108,106],  # a3
+  [ 21, 19, 25, 24,   53, 51, 57, 56,   85, 83, 89, 88,  117,115,121,120],  # a4
+  [ 20, 18, 23, 22,   52, 50, 55, 54,   84, 82, 87, 86,  116,114,123,122],  # a5
+  [125,128, 29, 27,   33, 32, 61, 59,   65, 64, 93, 91,   97, 96,129,127],  # a6
+  [119,118, 28, 26,   31, 30, 60, 58,   63, 62, 92, 90,   95, 94,124,126],  # a7
+]
+
+# Optimized (a_pair, b_pair) iteration order for better GPU scheduling
+# Interleaves A and B pairs to maximize instruction-level parallelism
+FMAC_PAIR_ORDER = [
+  (0,0),(0,1),(1,1),(1,0), (2,0),(2,1),(3,1),(3,2), (0,2),(0,3),(1,3),(1,2), (2,2),(2,3),(3,3),(3,4),
+  (0,4),(0,5),(1,5),(1,4), (2,4),(2,5),(3,5),(3,6), (0,6),(0,7),(1,7),(1,6), (2,6),(2,7),(3,7),(3,0),
+]
+
+def derive_fmac_pattern(acc_grid, a_tile_regs=None, b_tile_regs=None):
+  """Generate 64 dual FMAC ops from accumulator grid with optimized iteration order."""
+  pattern = []
+  for idx, (a_pair, b_pair) in enumerate(FMAC_PAIR_ORDER):
+    a_even, a_odd = a_pair * 2, a_pair * 2 + 1
+    b_even, b_odd = b_pair * 2, b_pair * 2 + 1
+    a_base, b_base = a_tile_regs[a_pair], b_tile_regs[b_pair]
+    # Op 1: normal order -> C[a_even, b_even] + C[a_odd, b_odd]
+    pattern.append((acc_grid[a_even][b_even], acc_grid[a_odd][b_odd],
+                   a_base, b_base, a_base+1, b_base+1))
+    # Op 2: alternate swapping A vs B to vary register banks
+    if idx % 2 == 0:  # swap B
+      pattern.append((acc_grid[a_even][b_odd], acc_grid[a_odd][b_even],
+                     a_base, b_base+1, a_base+1, b_base))
+    else:  # swap A
+      pattern.append((acc_grid[a_odd][b_even], acc_grid[a_even][b_odd],
+                     a_base+1, b_base, a_base, b_base+1))
+  return pattern
+
+# Derived: 64 dual FMAC operations
+FMAC_PATTERN = derive_fmac_pattern(ACC_GRID, V_A_TILE_REGS, V_B_TILE_REGS)
+
+def derive_permute_swaps(acc_grid, out_regs):
+  """Derive swap sequence to permute accumulators from FMAC layout to output order.
+
+  After FMAC loop: acc_grid[a][b] holds C[a,b]
+  Output order: for row_half in 0,1; col_group in 0-3; row_in_group in 0-3; b_off in 0-3
+    -> need C[row_half*4 + row_in_group, col_group*4 + b_off] in specified reg order
+  """
+  def target_ab(i):
+    row_half, col_group = i // 64, (i // 16) % 4
+    row_in_group, b_off = (i // 4) % 4, i % 4
+    return (row_half * 4 + row_in_group, col_group * 4 + b_off)
+
+  reg_contents = {acc_grid[a][b]: (a, b) for a in range(8) for b in range(16)}
+  ab_location = {ab: r for r, ab in reg_contents.items()}
+
+  swaps = []
+  for i in range(128):
+    target_reg, needed_ab = out_regs[i], target_ab(i)
+    current_reg = ab_location[needed_ab]
+    if current_reg != target_reg:
+      swaps.append((current_reg, target_reg))
+      ab_at_target = reg_contents.get(target_reg)
+      reg_contents[target_reg], ab_location[needed_ab] = needed_ab, target_reg
+      if ab_at_target is not None:
+        reg_contents[current_reg], ab_location[ab_at_target] = ab_at_target, current_reg
+  return swaps
+
+# Derived: swap sequence to arrange accumulators for output
+# Each group of 4 registers is ascending for direct global_store_b128
+OUT_REGS = [r for i in range(32) for r in range(126 - i*4, 130 - i*4)]
+PERMUTE_SWAPS = derive_permute_swaps(ACC_GRID, OUT_REGS)
+
+# =============================================================================
+# LDS tile staging registers
+# =============================================================================
+# DATA regs receive contiguous global prefetch, then write to LDS
+# TILE regs receive scattered LDS loads (ds_load_b64 pairs), then feed FMACs
+# Contiguous layout with mod4=[3,0,1,2,3,0,1,2] for bank conflict avoidance
+V_LDS_A_DATA = [163, 164, 165, 166, 167, 168, 169, 170]
+V_LDS_B_DATA = [171, 172, 173, 174, 175, 176, 177, 178]
+
+# Initial tile prefetch: (vdst, saddr_lo) - load into A data regs using B prefetch pointers (s[24:31])
+INIT_PREFETCH = [(V_LDS_A_DATA[i], S_PREFETCH_B+2*i) for i in range(4)]
+
+# Global memory prefetch schedule: (vdst1, vdst2, addr_vreg, saddr_lo1, saddr_lo2)
+# First 2 pairs from B prefetch pointers (s[32:39]), next 4 pairs from A prefetch pointers (s[40:55])
+PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR, S_PREFETCH_B+8+4*i, S_PREFETCH_B+10+4*i) for i in range(2)] + \
+                 [(V_LDS_B_DATA[2*(i-2)], V_LDS_B_DATA[2*(i-2)+1], V_GLOBAL_A_ADDR, S_PREFETCH_A+4*(i-2), S_PREFETCH_A+2+4*(i-2)) for i in range(2, 6)]
+
+# =============================================================================
+# Kernel class
+# =============================================================================
+
+class Kernel:
+  def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
+  def label(self, name): self.labels[name] = self.pos
+
+  def emit(self, inst, target=None):
+    self.instructions.append(inst)
+    inst._target, inst._pos = target, self.pos
+    self.pos += inst.size()
+    return inst
+
+  def waitcnt(self, lgkm=None, vm=None):
+    """Wait for memory operations. lgkm=N waits until N lgkm ops remain, vm=N waits until N vmem ops remain."""
+    vmcnt, lgkmcnt, expcnt = vm if vm is not None else 63, lgkm if lgkm is not None else 63, 7
+    waitcnt = (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
+    self.emit(s_waitcnt(simm16=waitcnt))
+
+  def finalize(self):
+    """Patch branch offsets and return the finalized instruction list."""
+    for inst in self.instructions:
+      if inst._target is None: continue
+      offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
+      if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
+      inst.simm16 = offset_dwords
+    return self.instructions
+
+
+# =============================================================================
+# Kernel builder
+# =============================================================================
+
+def build_kernel(N, arch='gfx1100'):
+  assert N % 128 == 0, f"N must be a multiple of 128 (tile size), got {N}"
+  assert N >= 256, f"N must be >= 256 (prefetch pipeline requires at least 2 K-blocks), got {N}"
+  k = Kernel(arch)
+
+  # ===========================================================================
+  # PROLOGUE: Load kernel arguments, compute tile coordinates and addresses
+  # ===========================================================================
+  k.emit(s_load_b128(sdata=s[S_KERNARG_A[0]:S_KERNARG_B[1]], sbase=s[0:1], offset=0x0, soffset=NULL))
+  k.emit(s_load_b64(sdata=s[S_OUT_PTR[0]:S_OUT_PTR[1]], sbase=s[0:1], offset=0x10, soffset=NULL))
+  k.emit(s_mov_b32(s[S_DIM_N], N))
+  k.emit(s_mov_b32(s[S_LOOP_CTR], 0))  # used by LDS swizzle, always 0 for valid workgroups
+  k.emit(s_lshl_b32(s[S_TILE_X], s[S_WORKGROUP_X], 7))
+  k.emit(s_lshl_b32(s[S_TILE_Y], s[S_WORKGROUP_Y], 7))
+
+  # Lane-derived values
+  k.emit(v_and_b32_e32(v[V_LANE_ID_MOD8], 7, v[V_LANE_ID]))
+  k.emit(v_lshrrev_b32_e32(v[4], 3, v[V_LANE_ID]))
+  k.emit(v_or_b32_e32(v[1], s[S_TILE_X], v[V_LANE_ID]))
+  k.emit(v_or_b32_e32(v[22], s[S_TILE_Y], v[4]))
+  k.emit(v_lshlrev_b32_e32(v[V_LANE_MOD8_X4], 2, v[V_LANE_ID_MOD8]))
+  k.waitcnt(lgkm=0)
+
+  # Compute 8 A and B matrix tile base pointers for prefetch
+  k.emit(s_mov_b64(s[S_PREFETCH_B:S_PREFETCH_B+1], s[S_KERNARG_B[0]:S_KERNARG_B[1]]))  # B[0]: no offset
+  for i in range(1, 8):  # B: each pointer 1 row of B apart (N*4 bytes)
+    k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_KERNARG_B[0]], i * N * 4))
+    k.emit(s_addc_u32(s[S_PREFETCH_B+i*2+1], s[S_KERNARG_B[1]], 0))
+  k.emit(s_mov_b64(s[S_PREFETCH_A:S_PREFETCH_A+1], s[S_KERNARG_A[0]:S_KERNARG_A[1]]))  # A[0]: no offset
+  for i in range(1, 8):  # A: each pointer 16 rows of A apart (16*N*4 bytes)
+    k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_KERNARG_A[0]], i * N * 64))
+    k.emit(s_addc_u32(s[S_PREFETCH_A+i*2+1], s[S_KERNARG_A[1]], 0))
+
+  # Global prefetch addresses: B = (tile_x + lane_id) * 4, A = (tile_y*N + (lane_id/8)*N + lane_id%8) * 4
+  k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], s[S_TILE_X], v[V_LANE_ID]))
+  k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_B_ADDR], 2, v[V_GLOBAL_B_ADDR]))
+  k.emit(s_mul_i32(s[19], s[S_TILE_Y], N))
+  k.emit(v_mul_lo_u32(v[V_GLOBAL_A_ADDR], v[4], N))                          # (lane_id/8)*N
+  k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], v[V_LANE_ID_MOD8], v[V_GLOBAL_A_ADDR]))  # + lane_id%8
+  k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], s[19], v[V_GLOBAL_A_ADDR]))
+  k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_A_ADDR], 2, v[V_GLOBAL_A_ADDR]))
+
+  # Do initial loads
+  for vdst, saddr_lo in INIT_PREFETCH:
+    k.emit(global_load_b32(vdst=v[vdst], addr=v[V_GLOBAL_B_ADDR], saddr=s[saddr_lo:saddr_lo+1]))
+  for iter in range(6):
+    vdst1, vdst2, addr, slo1, slo2 = PREFETCH_LOADS[iter]
+    k.emit(global_load_b32(vdst=v[vdst1], addr=v[addr], saddr=s[slo1:slo1+1]))
+    k.emit(global_load_b32(vdst=v[vdst2], addr=v[addr], saddr=s[slo2:slo2+1]))
+
+  # ===========================================================================
+  # LDS store address computation (bank-conflict-avoiding swizzle)
+  # ===========================================================================
+  # This section computes LDS store addresses with a swizzle pattern to avoid bank conflicts.
+  # The swizzle ensures that threads in the same wavefront write to different LDS banks.
+  # Formula: swizzled_addr = base + (lane_id & 7) * LDS_A_STRIDE + swizzle_offset
+  # where swizzle_offset depends on (lane_id >> 3) to distribute across banks.
+  k.emit(v_add_nc_u32_e32(v[9], s[S_LOOP_CTR], v[22]))  # row 0 base
+  k.emit(v_and_b32_e32(v[9], ADDR_MASK, v[9]))
+  k.emit(v_sub_nc_u32_e32(v[9], v[22], v[9]))  # row 0 swizzle offset
+  k.emit(v_lshlrev_b32_e32(v[9], 2, v[9]))  # * 4
+  k.emit(v_mad_u32_u24(v[V_LDS_B_ADDR], LDS_A_STRIDE, v[V_LANE_ID_MOD8], v[9]))
+
+  # For V_LDS_A_BASE and epilogue
+  k.emit(v_bfe_u32(v[2], v[V_LANE_ID], 3, 2))  # v[2] = (lane_id >> 3) & 3
+  k.emit(v_lshlrev_b32_e32(v[V_LANE_DIV8_X4], 2, v[2]))
+
+  # Compute LDS load/store base addresses for inner loop
+  k.emit(v_lshlrev_b32_e32(v[2], 4, v[2]))
+  k.emit(v_and_b32_e32(v[3], 0x7F, v[1]))  # simplified from 3 lines
+  k.emit(v_lshl_or_b32(v[V_LDS_B_BASE], v[V_LANE_ID_MOD8], 4, LDS_BASE_OFFSET))
+  k.emit(v_lshl_add_u32(v[V_LDS_A_ADDR], v[3], 2, LDS_BASE_OFFSET))
+  k.emit(v_lshlrev_b32_e32(v[3], 2, v[V_LANE_ID]))
+  k.emit(v_and_or_b32(v[V_LDS_A_BASE], 0x180, v[3], v[2]))
+
+  # Do initial stores
+  k.waitcnt(vm=0)
+  for i in range(4):  # A tile: 8 values via 4 stride64 stores
+    k.emit(ds_store_2addr_stride64_b32(addr=v[V_LDS_A_ADDR], data0=v[V_LDS_A_DATA[i*2]], data1=v[V_LDS_A_DATA[i*2+1]], offset0=i*4, offset1=i*4+2))
+  for i in range(8):  # B tile: 8 values via 8 scalar stores with 64-byte spacing
+    offset = i * 64
+    k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[V_LDS_B_DATA[i]], offset0=offset & 0xFF, offset1=offset >> 8))
+
+  # Zero all 128 accumulators using VOPD dual moves (64 instructions instead of 128)
+  for i in range(0, len(OUT_REGS), 2):
+    k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[OUT_REGS[i]], vdsty=v[OUT_REGS[i+1]], srcx0=0, srcy0=0))
+  k.emit(s_add_i32(s[S_LOOP_BOUND], s[S_DIM_N], -8))
+
+  # S_LOOP_CTR is already 0 from prologue initialization
+  k.emit(s_branch(), target='LOOP_ENTRY')
+
+  # ===========================================================================
+  # MAIN GEMM LOOP
+  # ===========================================================================
+
+  NO_DS, NO_GLOBAL = getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
+
+  k.label('LOOP_INC')
+  k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
+  k.emit(s_cmp_ge_i32(s[S_LOOP_CTR], s[S_DIM_N]))
+  k.emit(s_cbranch_scc1(), target='EPILOGUE')
+
+  k.label('LOOP_ENTRY')
+  k.emit(s_cmp_lt_i32(s[S_LOOP_CTR], s[S_LOOP_BOUND]))
+  k.emit(s_cselect_b32(s[S_PREFETCH_FLAG], -1, 0))  # s_cselect doesn't modify SCC
+  k.emit(s_cbranch_scc0(), target='SKIP_PREFETCH')  # branch if loop_ctr >= loop_bound
+
+  if not NO_GLOBAL:
+    # Advance prefetch pointers (VGPR)
+    #k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], N * 32, v[V_GLOBAL_B_ADDR]))
+    #k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], 0x20, v[V_GLOBAL_A_ADDR]))
+
+    # Advance prefetch pointers (64-bit adds): B advances 8 rows (8*N*4 bytes), A advances 8 cols (8*4 bytes)
+    k.emit(s_clause(simm16=31))
+    for i in range(8):
+      k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_PREFETCH_B+i*2], N * 32))
+      k.emit(s_addc_u32(s[S_PREFETCH_B+i*2+1], s[S_PREFETCH_B+i*2+1], 0))
+    for i in range(8):
+      k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_PREFETCH_A+i*2], 0x20))
+      k.emit(s_addc_u32(s[S_PREFETCH_A+i*2+1], s[S_PREFETCH_A+i*2+1], 0))
+
+    # do the fetch
+    for vdst, saddr_lo in INIT_PREFETCH:
+      k.emit(global_load_b32(vdst=v[vdst], addr=v[V_GLOBAL_B_ADDR], saddr=s[saddr_lo:saddr_lo+1]))
+
+  k.label('SKIP_PREFETCH')
+
+  # wait for local stores to finish (either initial or loop)
+  # then sync the warp so it's safe to load local
+  k.waitcnt(lgkm=0)
+  k.emit(s_barrier())
+
+  # 8 inner loop iterations
+  for iter in range(8):
+    # Load A tile (4 pairs) and B tile (8 pairs) from LDS
+    if not NO_DS:
+      k.emit(s_clause(simm16=len(V_A_TILE_REGS) + len(V_B_TILE_REGS) - 1))  # 12 loads total: 4 A + 8 B
+      # A tile: 4 ds_load_b64
+      for i, vdst in enumerate(V_A_TILE_REGS):
+        a_off = (i & 1) * 8 + (i >> 1) * 64 + iter * LDS_A_STRIDE
+        k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_A_BASE], offset0=a_off & 0xFF, offset1=a_off >> 8))
+      # B tile: 8 ds_load_b64
+      for i, vdst in enumerate(V_B_TILE_REGS):
+        b_off = (i & 1) * 8 + (i & 2) * 64 + (i >> 2) * 256 + iter * LDS_B_STRIDE
+        k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_B_BASE], offset0=b_off & 0xFF, offset1=b_off >> 8))
+
+    # Issue global prefetch (first 6 iterations only)
+    if iter < 6 and not NO_GLOBAL:
+      vdst1, vdst2, addr, slo1, slo2 = PREFETCH_LOADS[iter]
+      k.emit(global_load_b32(vdst=v[vdst1], addr=v[addr], saddr=s[slo1:slo1+1]))
+      k.emit(global_load_b32(vdst=v[vdst2], addr=v[addr], saddr=s[slo2:slo2+1]))
+
+    # 64 dual FMACs
+    k.waitcnt(lgkm=0)
+    k.emit(s_clause(simm16=len(FMAC_PATTERN)-1))
+    for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
+      k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
+                  vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))
+
+  # wait for all global loads to finish
+  # then sync the warp so it's safe to store local
+  k.waitcnt(vm=0)
+  k.emit(s_barrier())
+
+  # Store prefetched data to LDS
+  # NOTE: Register naming reflects LDS tile organization, not source matrix:
+  #   V_LDS_A_DATA (v155-162) holds data that goes to LDS A-tile region
+  #   V_LDS_B_DATA (v163-170) holds data that goes to LDS B-tile region
+  # The data sources are swapped: A-tile receives B matrix rows, B-tile receives A matrix columns
+  if not NO_DS:
+    for i in range(4):  # A tile: 8 values via 4 stride64 stores
+      k.emit(ds_store_2addr_stride64_b32(addr=v[V_LDS_A_ADDR], data0=v[V_LDS_A_DATA[i*2]], data1=v[V_LDS_A_DATA[i*2+1]], offset0=i*4, offset1=i*4+2))
+    for i in range(8):  # B tile: 8 values via 8 scalar stores with 64-byte spacing
+      offset = i * 64
+      k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[V_LDS_B_DATA[i]], offset0=offset & 0xFF, offset1=offset >> 8))
+
+  k.emit(s_branch(), target='LOOP_INC')
+
+  # ===========================================================================
+  # EPILOGUE: Permute and store results
+  # ===========================================================================
+  k.label('EPILOGUE')
+
+  # Rearrange accumulators from FMAC layout to contiguous output order
+  for a, b in PERMUTE_SWAPS:
+    k.emit(v_swap_b32_e32(v[a], v[b]))
+
+  # Compute output base coordinates
+  # v[130] = col_base = tile_x + (lane_id & 7) * 4
+  # v[131] = row_base = tile_y + (lane_id & 0x60) + ((lane_id >> 3) & 3) * 4
+  # v[132] = 0 (for 64-bit address high part)
+  k.emit(v_add_nc_u32_e32(v[130], s[S_TILE_X], v[V_LANE_MOD8_X4]))
+  k.emit(v_and_b32_e32(v[131], 0x60, v[V_LANE_ID]))
+  k.emit(v_add_nc_u32_e32(v[131], s[S_TILE_Y], v[131]))
+  k.emit(v_add_nc_u32_e32(v[131], v[V_LANE_DIV8_X4], v[131]))
+  k.emit(v_mov_b32_e32(v[132], 0))
+
+  # Precompute row offsets: v[133-136] for rows 0-3, v[137-140] for rows 16-19
+  for base, row_off in [(133, 0), (137, 16)]:
+    if row_off: k.emit(v_add_nc_u32_e32(v[141], row_off, v[131]))
+    k.emit(v_mul_lo_u32(v[base], v[141] if row_off else v[131], s[S_DIM_N]))
+    for j in range(3): k.emit(v_add_nc_u32_e32(v[base + 1 + j], s[S_DIM_N], v[base + j]))
+
+  # s[S_PREFETCH_FLAG] = row stride in bytes (N * 4)
+  k.emit(s_lshl_b32(s[S_PREFETCH_FLAG], s[S_DIM_N], 2))
+
+  # Store 128 output values as 32 groups of 4 (128-bit stores)
+  # Layout: 2 row halves (0-3, 16-19) x 4 col groups x 4 rows = 32 stores of 4 floats
+  for i, (row_half, col_off, row_in_group) in enumerate([(rh, co, ri)
+      for rh in range(2) for co in [0, 32, 64, 96] for ri in range(4)]):
+    row = row_half * 16 + row_in_group
+    src = OUT_REGS[i*4]  # first reg of ascending group of 4
+
+    if row_in_group == 0:
+      # First row of group: compute full address
+      if col_off == 0: k.emit(v_mov_b32_e32(v[141], v[130]))
+      else: k.emit(v_add_nc_u32_e32(v[141], col_off, v[130]))
+      row_base = 133 + row if row < 4 else 137 + row - 16
+      k.emit(v_add_nc_u32_e32(v[141], v[row_base], v[141]))
+      k.emit(v_lshlrev_b32_e32(v[141], 2, v[141]))
+      k.emit(v_add_co_u32(v[141], VCC_LO, s[S_OUT_PTR[0]], v[141]))
+      k.emit(v_add_co_ci_u32_e32(v[142], s[S_OUT_PTR[1]], v[132]))
+    else:
+      # Subsequent rows: add stride
+      k.emit(v_add_co_u32(v[141], VCC_LO, s[S_PREFETCH_FLAG], v[141]))
+      k.emit(v_add_co_ci_u32_e32(v[142], v[142], v[132]))
+
+    k.emit(global_store_b128(addr=v[141:142], data=v[src:src+3], saddr=NULL))
+
+  k.emit(s_sendmsg(simm16=3))  # DEALLOC_VGPRS
+  k.emit(s_endpgm())
+
+  return k.finalize()
+
+# =============================================================================
+# Test harness
+# =============================================================================
+
+N = getenv("N", 4096)
+BLOCK_M, BLOCK_N = 128, 128
+THREADS = 128
+
+def test_matmul():
+  dev = Device[Device.DEFAULT]
+  print(f"Device arch: {dev.renderer.arch}")
+
+  insts = build_kernel(N, dev.renderer.arch)
+
+  rng = np.random.default_rng(42)
+  a = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
+  b = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
+  c = Tensor.empty(N, N)
+  Tensor.realize(a, b, c)
+
+  grid, local = (N // BLOCK_N, N // BLOCK_M, 1), (THREADS, 1, 1)
+  print(f"Grid: {grid}, Local: {local}")
+
+  dname:str = Device.DEFAULT
+  def asm_kernel(A:UOp, B:UOp, C:UOp) -> UOp:
+    gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
+    lidxs = [UOp.special(n, f"lidx{i}") for i,n in enumerate(local)]
+    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536)), addrspace=AddrSpace.LOCAL), (), 'lds')
+    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
+                                                                                  estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
+  ei = c.schedule()[0].lower()
+
+  ets = []
+  with Context(DEBUG=2):
+    for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True))
+  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
+
+  if getenv("VERIFY", 1):
+    GlobalCounters.reset()
+    with Context(DEBUG=2): tc = (a @ b).realize()
+    with Context(DEBUG=0): err = (c - tc).square().mean().item()
+    print(f"mean squared error {err}")
+    if err != err or err > 1e-06:
+      c_np, tc_np = c.numpy(), tc.numpy()
+      for bi in range(N // 128):
+        for bj in range(N // 128):
+          blk_c = c_np[bi*128:(bi+1)*128, bj*128:(bj+1)*128]
+          blk_ref = tc_np[bi*128:(bi+1)*128, bj*128:(bj+1)*128]
+          blk_diff = blk_c - blk_ref
+          zero_rows = [i for i in range(128) if np.all(np.abs(blk_c[i,:]) < 1e-10)]
+          nz_rows = [i for i in range(128) if i not in zero_rows]
+          nz_mse = float(np.mean(blk_diff[nz_rows,:]**2)) if nz_rows else 0
+          print(f"Block ({bi},{bj}): zero_rows={zero_rows}, nz_rows_mse={nz_mse:.2e}")
+          # show first few non-zero row comparisons
+          if nz_rows and nz_mse > 1e-6:
+            for r in nz_rows[:3]:
+              print(f"  row {r} asm[0:8]:  {blk_c[r,:8]}")
+              print(f"  row {r} ref[0:8]:  {blk_ref[r,:8]}")
+      raise RuntimeError("matmul is wrong!")
+
+if __name__ == "__main__":
+  test_matmul()

extra/gemm/amd_uop_matmul.py

@@ -140,11 +140,11 @@ def hand_spec_kernel3():
 
   return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
 
-def test_matmul(sink:UOp, N=N):
+def test_matmul(sink:UOp, dtype=dtypes.float32, N=N):
   rng = np.random.default_rng()
-  a = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
-  b = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
-  hc = Tensor.empty(N, N)
+  a = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
+  b = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
+  hc = Tensor.empty(N, N, dtype=dtype)
   Tensor.realize(a, b, hc)
 
   ei = ExecItem(sink, [t.uop.buffer for t in [hc, a, b]], prg=get_runner(Device.DEFAULT, sink))

extra/gemm/asm/cdna/gemm.py

@@ -0,0 +1,101 @@
+import atexit, functools
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.renderer import Estimates
+from tinygrad.helpers import getenv, all_same, dedup
+from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG
+
+# ** CDNA4 assembly gemm
+
+WORKGROUP_SIZE = 256
+
+@functools.cache
+def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
+  batch, M, K = A.shape
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
+  gidx = UOp.special(NUM_WG, "gidx0")
+  insts = build_kernel(batch, M, N, K, A.dtype.base)
+  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
+  sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
+                  arg=KernelInfo(name=f"gemm_{batch}_{M}_{N}_{K}", estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname),
+                                UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+
+counters = {"used":0, "todos":[]}
+def todo(msg:str) -> bool: counters["todos"].append(msg); return False
+atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))
+
+def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
+  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
+  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
+  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
+  N = b.shape[1]
+  # only sharding on the batch or K is tested, others might work too
+  if isinstance(a.device, tuple):
+    if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
+    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
+    dname = a.device[0]
+  else: dname = a.device
+  arch = getattr(Device[dname].renderer, "arch", "")
+  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
+  if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
+    return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
+  return True
+
+# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
+# note: this can be removed after we have GEMM on mixins
+
+def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
+  M, K = A.shape[0]*A.shape[1], A.shape[2]
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  m = UOp.range(M, 1, AxisType.LOOP)
+  n = UOp.range(N, 2, AxisType.LOOP)
+  k = UOp.range(K, 0, AxisType.REDUCE)
+  mul = (A.index((m*UOp.const(dtypes.index, K)+k))*B.index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
+  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
+  store = C.index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
+  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
+
+# ** backward gemm, might use the asm gemm
+
+def custom_gemm_bw(gradient:UOp, kernel:UOp):
+  out, a, b = kernel.src[1:]
+  assert all_same([gradient.device, a.device, b.device, out.device])
+  a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
+  # TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
+  g_t = g_t[:a.shape[0]]
+  grad_a = (g_t @ b_t.T).uop
+  grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
+  return (None, grad_a, grad_b)
+
+# ** main gemm function
+
+def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
+  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
+  counters["used"] += 1
+  squeeze = a.ndim == 2
+  if squeeze: a = a.unsqueeze(0)
+
+  batch, M, K = a.shape
+  N = b.shape[1]
+  is_multi = isinstance(a.device, tuple)
+  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
+
+  if is_multi:
+    out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+  else:
+    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
+
+  renderer = Device[a.device[0] if is_multi else a.device].renderer
+  dname, arch = renderer.device, getattr(renderer, "arch", "")
+  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
+    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
+  else:
+    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
+  if k_sharded: out = out.sum(0)
+  return out.squeeze(0) if squeeze else out

extra/gemm/asm/rdna3/template.s

@@ -0,0 +1,76 @@
+.text
+.section	.text.
+.global	gemm
+.p2align	8
+.type	gemm,@function
+
+gemm:
+INSTRUCTIONS
+
+.section .rodata,"a",@progbits
+.p2align 6, 0x0
+.amdhsa_kernel gemm
+  # basic memory requirements
+  .amdhsa_group_segment_fixed_size 30336
+  .amdhsa_private_segment_fixed_size 0
+  .amdhsa_kernarg_size 32
+  # register usage (RSRC1)
+  .amdhsa_next_free_vgpr 256
+  .amdhsa_next_free_sgpr 100
+  # workgroup / workitem IDs (RSRC2)
+  .amdhsa_system_sgpr_workgroup_id_x 1
+  .amdhsa_system_sgpr_workgroup_id_y 1
+  .amdhsa_system_sgpr_workgroup_id_z 1
+  # user SGPRs: kernarg ptr in s[0:1]
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_user_sgpr_count 2
+  # gfx10+ / gfx11 specifics (RSRC1[29..31])
+  .amdhsa_wavefront_size32 1
+  .amdhsa_workgroup_processor_mode 1
+  .amdhsa_memory_ordered 1
+  .amdhsa_forward_progress 1
+  # misc for gfx11
+  .amdhsa_dx10_clamp 1
+  .amdhsa_ieee_mode 1
+  .amdhsa_uses_dynamic_stack 0
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.kernels:
+  - .args:
+      - .address_space:  generic
+        .name:           C
+        .offset:         0
+        .size:           8
+        .value_kind:     global_buffer
+        .value_type:     f16
+      - .address_space:  generic
+        .name:           A
+        .offset:         8
+        .size:           8
+        .value_kind:     global_buffer
+        .value_type:     f16
+      - .address_space:  generic
+        .name:           B
+        .offset:         16
+        .size:           8
+        .value_kind:     global_buffer
+        .value_type:     f16
+    .group_segment_fixed_size: 30336
+    .kernarg_segment_align: 8
+    .kernarg_segment_size: 32
+    .max_flat_workgroup_size: 128
+    .name:           gemm
+    .private_segment_fixed_size: 0
+    .sgpr_count:     70
+    .sgpr_spill_count: 0
+    .symbol:         gemm.kd
+    .vgpr_count:     256
+    .vgpr_spill_count: 0
+    .wavefront_size: 32
+amdhsa.version:
+  - 1
+  - 1
+...
+.end_amdgpu_metadata

extra/gemm/asm/rdna3/test.py

@@ -0,0 +1,30 @@
+import math, pathlib
+
+from tinygrad import Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+
+from extra.gemm.amd_uop_matmul import test_matmul
+
+N = 4096
+TN = 96
+THREADS_PER_WG = 128
+NUM_WG = math.ceil(N / TN) * math.ceil(N / TN)
+
+dname:str = Device.DEFAULT
+template:str = (pathlib.Path(__file__).parent/"template.s").read_text()
+
+def asm_kernel() -> UOp:
+  lidx = UOp.special(THREADS_PER_WG, "lidx0")
+  gidx = UOp.special(NUM_WG, "gidx0")
+
+  a = UOp.placeholder((N*N,), dtypes.half, slot=1)
+  b = UOp.placeholder((N*N,), dtypes.half, slot=2)
+  c = UOp.placeholder((N*N,), dtypes.half, slot=0)
+
+  src = template.replace("INSTRUCTIONS", (pathlib.Path(__file__).parent/"gemm.s").read_text())
+
+  sink = UOp.sink(a, b, c, lidx, gidx, arg=KernelInfo(name="gemm"))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
+
+if __name__ == "__main__":
+  test_matmul(asm_kernel(), dtype=dtypes.half, N=N)

extra/gemm/asm/test.py

@@ -1,65 +0,0 @@
-# Run assembly on the AMD runtime and check correctness
-# VIZ=2 to profile
-import pathlib
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.engine.realize import ExecItem, CompiledRunner
-from tinygrad.renderer import ProgramSpec
-from tinygrad.uop.ops import track_rewrites, UOp
-from tinygrad.helpers import TracingKey
-
-fp = pathlib.Path(__file__).parent/"gemm.s"
-
-# ** generate inputs on CPU
-
-N = 8192
-scale = 10.0
-
-import torch
-torch.manual_seed(0)
-A = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
-B = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
-Bt = B.t().contiguous() # transpose B for the baseline gemm
-C_torch = A@Bt
-
-# ** copy buffers to AMD
-
-# input creation and validation run on the copy engine for simpler tracing
-
-def from_torch(t:torch.Tensor) -> Tensor:
-  return Tensor.from_blob(t.data_ptr(), t.shape, dtype=dtypes.bfloat16, device="cpu").to(Device.DEFAULT).realize()
-
-C_tiny = Tensor.matmul(from_torch(A), from_torch(Bt), dtype=dtypes.float32).cast(dtypes.bfloat16)
-C_asm = Tensor.empty_like(C_tiny)
-C_asm.uop.buffer.allocate()
-
-# ** run gemms
-
-@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name,), ret=ret))
-def get_asm_gemm(ast:UOp, fp:pathlib.Path) -> ProgramSpec:
-  src = fp.read_text()
-  lib = Device[Device.DEFAULT].compiler.compile(src)
-  return ProgramSpec("gemm", src, Device.DEFAULT, ast, lib=lib, global_size=[1024, 1, 1], local_size=[256, 1, 1], globals=[0, 1, 2])
-
-sched = C_tiny.schedule()
-assert len(sched) == 1
-eis:list[ExecItem] = [sched[-1].lower()]
-ast = eis[0].ast
-prg = get_asm_gemm(ast, fp)
-eis.append(ExecItem(ast, [C_asm.uop.buffer, from_torch(B).uop.buffer, from_torch(A).uop.buffer], prg=CompiledRunner(prg)))
-
-
-for ei in eis:
-  et = ei.run(wait=True)
-  print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
-
-# ** correctness
-
-import ctypes
-
-def torch_bf16(t:Tensor) -> torch.tensor:
-  asm_out = t.to("cpu").realize().uop.buffer._buf
-  buf = (ctypes.c_uint16*C_asm.uop.size).from_address(asm_out.va_addr)
-  return torch.frombuffer(buf, dtype=torch.bfloat16, count=C_asm.uop.size).reshape(C_asm.shape)
-
-assert torch.allclose(torch_bf16(C_asm), C_torch, rtol=1e-2, atol=1e-3)
-assert torch.allclose(torch_bf16(C_tiny), C_torch, rtol=1e-2, atol=1e-3)

extra/gemm/asm/unpack_kd.py

@@ -1,12 +1,12 @@
-# unpack the complete kernel descriptor of an amdgpu ELF of for gfx950
+# unpack the complete kernel descriptor of an amdgpu ELF
 # https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
-import struct, pathlib
+import struct, pathlib, sys
 from tinygrad.runtime.support.elf import elf_loader
 
 def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
 def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
 
-with open(fp:=pathlib.Path(__file__).parent/"lib", "rb") as f:
+with open(sys.argv[1], "rb") as f:
   lib = f.read()
 
 image, sections, relocs = elf_loader(lib)
@@ -49,7 +49,7 @@ print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
 print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
 print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
 
-# rsrc 3
+# rsrc 3 (gfx950)
 
 accum_offset_raw = bits(pgm_rsrc3, 0, 5)
 assert_zero(pgm_rsrc3, 6, 15)
@@ -169,10 +169,10 @@ assert_zero(desc, 458, 459)
 uses_dynamic_stack = bits(desc, 459, 460)
 print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
 
+# gfx950 only
 assert_zero(desc, 460, 463)
 kernarg_preload_spec_length = bits(desc, 464, 470)
 print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
-
 kernarg_preload_spec_offset = bits(desc, 471, 479)
 print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
 

extra/gemm/intel_xmx.py

@@ -37,7 +37,7 @@ b.copyin(row.data)
 c.copyin(mat.data)
 ret = prog(a._buf, b._buf, c._buf, global_size=[1,1,1], local_size=[8,1,1], wait=True)
 print(ret)
-out = np.frombuffer(a.as_buffer(), np.float32)
+out = np.frombuffer(a.as_memoryview(), np.float32)
 real = row.astype(np.float32)@mat.T.astype(np.float32)
 print("out:", out)
 print("real", real)

extra/gemm/triton_nv_matmul.py

@@ -98,10 +98,10 @@ if __name__ == "__main__":
   # check correctness
   if getenv("VERIFY"):
     from tinygrad.engine.realize import run_schedule
-    triton_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
+    triton_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
     print(triton_buf)
     run_schedule(sched)
-    tinygrad_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
+    tinygrad_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
     print(tinygrad_buf)
     np.testing.assert_allclose(triton_buf, tinygrad_buf)
     print("correct!")

extra/hcq/hcq_smi.py

@@ -1,14 +1,15 @@
 #!/usr/bin/env python3
 
 import argparse, glob, os, time, subprocess, sys
+from tinygrad.helpers import temp
 
 def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
   target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
 
   devs = []
-  for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
-    dev_id = dev[8:-5]
-    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
+  for dev in glob.glob(temp(f'{prefix}_*.lock')):
+    dev_id = dev.split('/')[-1][len(prefix)+1:-5]
+    if dev_id.startswith(target_dev): devs.append(dev_id)
   return devs
 
 def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{pci_bus}/reset'")
@@ -53,16 +54,7 @@ def cmd_show_pids(args):
 
   for dev in devs:
     try:
-      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
-      print(f"{dev}: {pid}")
-    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
-
-def cmd_kill_pids(args):
-  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
-
-  for dev in devs:
-    try:
-      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+      pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
       print(f"{dev}: {pid}")
     except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
 
@@ -74,7 +66,7 @@ def cmd_kill_pids(args):
       if i > 0: time.sleep(0.2)
 
       try:
-        try: pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+        try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
         except subprocess.CalledProcessError: break
 
         print(f"Killing process {pid} (which uses {dev})")

extra/hevc/decode.py

@@ -1,8 +1,42 @@
-import argparse, os, hashlib
-from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch
+import argparse, os, hashlib, functools
+from typing import Iterator, Callable
+from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch, ceildiv
 from extra.hevc.hevc import parse_hevc_file_headers, untile_nv12, to_bgr, nv_gpu
 from tinygrad import Tensor, dtypes, Device, Variable, TinyJit
 
+# rounds up hevc input data to 32 bytes, so more optimal kernels can be generated
+HEVC_ROUNDUP = getenv("DATA_ROUNDUP", 32)
+
+@functools.cache
+def _hevc_jitted_decoder(out_image_size:tuple[int, int], max_hist:int, inplace:bool):
+  def hevc_decode_frame(pos:Variable, hevc_tensor:Tensor, offset:Variable, sz:Variable, opaque:Tensor, i:Variable, *hist:Tensor, outbuf:Tensor|None=None):
+    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist)
+    if outbuf is not None: outbuf.assign(x).realize()
+    return x.realize()
+  return TinyJit(hevc_decode_frame)
+
+def hevc_decode(hevc_tensor:Tensor, opaque:Tensor, frame_info:list, luma_h:int, luma_w:int,
+                history:list[Tensor]|None=None, preallocated_outputs:list[Tensor]|None=None, warmup=False) -> Iterator[Tensor]:
+  out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
+  max_hist = max((hs for _, _, _, hs, _ in frame_info), default=0)
+
+  v_pos = Variable("pos", 0, max_hist + 1)
+  v_offset = Variable("offset", 0, hevc_tensor.numel()-1)
+  v_sz = Variable("sz", 1, ceildiv(hevc_tensor.numel(), HEVC_ROUNDUP))
+  v_i = Variable("i", 0, len(frame_info)-1)
+
+  decode_jit = _hevc_jitted_decoder(out_image_size, max_hist, preallocated_outputs is not None)
+  history = history or [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(max_hist)]
+  assert len(history) == max_hist, f"history length {len(history)} does not match max_hist {max_hist}"
+
+  for i, (offset, sz, frame_pos, _, is_hist) in enumerate(frame_info):
+    history = history[-max_hist:] if max_hist > 0 else []
+    img = decode_jit(v_pos.bind(frame_pos), hevc_tensor, v_offset.bind(offset), v_sz.bind(ceildiv(sz, HEVC_ROUNDUP)),
+                     opaque, v_i.bind(i), *history, outbuf=preallocated_outputs[i] if preallocated_outputs else None)
+    res = preallocated_outputs[i] if preallocated_outputs else img.clone().realize()
+    if is_hist: history.append(res)
+    yield res
+
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--input_file", type=str, default="")
@@ -19,49 +53,32 @@ if __name__ == "__main__":
   dat_hash = hashlib.md5(dat).hexdigest()
 
   with Timing("prep infos: "):
-    dat_nv = hevc_tensor.to("NV")
     opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat)
 
   frame_info = frame_info[:getenv("MAX_FRAMES", len(frame_info))]
 
   # move all needed data to gpu
-  #all_slices = []
   with Timing("copy to gpu: "):
     opaque_nv = opaque.to("NV").contiguous().realize()
     hevc_tensor = hevc_tensor.to("NV")
 
   out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
-  max_hist = max(history_sz for _, _, _, history_sz, _ in frame_info)
 
-  # define variables
-  v_pos = Variable("pos", 0, max_hist + 1)
-  v_offset = Variable("offset", 0, hevc_tensor.numel()-1)
-  v_sz = Variable("sz", 0, hevc_tensor.numel())
-  v_i = Variable("i", 0, len(frame_info)-1)
+  # preallocate output/hist buffers
+  max_hist = max((hs for _, _, _, hs, _ in frame_info), default=0)
+  hist = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(max_hist)]
+  out_images = [Tensor.zeros(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(len(frame_info))]
 
-  @TinyJit
-  def decode_jit(pos:Variable, src:Tensor, data:Tensor, *hist:Tensor):
-    return src.decode_hevc_frame(pos, out_image_size, data, hist).realize()
+  # warmup decode
+  _ = list(hevc_decode(hevc_tensor, opaque_nv, frame_info[:3], luma_h, luma_w, history=hist, preallocated_outputs=out_images))
+  Device.default.synchronize()
 
-  # warm up
-  history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)]
-  for i in range(3):
-    hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(frame_info[0][0])), bound_offset+v_sz.bind(frame_info[0][1])),))
-    decode_jit(v_pos.bind(0), hevc_frame, opaque_nv[v_i.bind(0)], *history)
-
-  out_images = []
+  # decode all frames using the iterator
   with Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps")):
-    for i, (offset, sz, frame_pos, history_sz, is_hist) in enumerate(frame_info):
-      history = history[-max_hist:] if max_hist > 0 else []
-      # TODO: this shrink should work as a slice
-      hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(offset)), bound_offset+v_sz.bind(sz)),))
-
-      outimg = decode_jit(v_pos.bind(frame_pos), hevc_frame, opaque_nv[v_i.bind(i)], *history).clone()
-      out_images.append(outimg)
-      if is_hist: history.append(outimg)
-
+    images = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w, history=hist, preallocated_outputs=out_images))
     Device.default.synchronize()
 
+  # validation
   if getenv("VALIDATE", 0):
     import pickle
     if dat_hash == "b813bfdbec194fd17fdf0e3ceb8cea1c":
@@ -70,7 +87,7 @@ if __name__ == "__main__":
     else: decoded_frames = pickle.load(open(f"extra/hevc/decoded_frames_{dat_hash}.pkl", "rb"))
   else: import cv2
 
-  for i, img in tqdm(enumerate(out_images)):
+  for i, img in tqdm(enumerate(images)):
     if getenv("VALIDATE", 0):
       if i < len(decoded_frames) and len(decoded_frames[i]) > 0:
         img = untile_nv12(img, h, w, luma_w, chroma_off).realize()

extra/hip_gpu_driver/hip_ioctl.py

@@ -13,7 +13,7 @@ def get_struct(argp, stype):
 
 def format_struct(s):
   sdats = []
-  for field_name, field_type in s._fields_:
+  for field_name, *_ in s._real_fields_:
     dat = getattr(s, field_name)
     if isinstance(dat, int): sdats.append(f"{field_name}:0x{dat:X}")
     else: sdats.append(f"{field_name}:{dat}")
@@ -46,6 +46,22 @@ def install_hook(c_function, python_function):
 # *** ioctl lib end ***
 
 import tinygrad.runtime.autogen.kfd as kfd_ioctl
+import tinygrad.runtime.autogen.hsa as hsa
+
+def print_aql_queue(read_pointer_address):
+  rptr_offset = getattr(hsa.amd_queue_v2_t, 'read_dispatch_id').offset
+  queue_base = read_pointer_address - rptr_offset
+  queue = hsa.amd_queue_v2_t.from_address(queue_base)
+  print(f"  AQL Queue @ 0x{queue_base:X}:")
+  for field_name, *_ in hsa.amd_queue_v2_t._real_fields_:
+    val = getattr(queue, field_name)
+    if isinstance(val, int): print(f"    {field_name}: 0x{val:X}")
+    elif hasattr(val, '_length_'):
+      arr_vals = [f"{format_struct(v)}" if hasattr(v, '_real_fields_') else f"{v:#X}" for v in val]
+      print(f"    {field_name}: [{', '.join(arr_vals)}]")
+    elif hasattr(val, '_real_fields_'): print(f"    {field_name}: {format_struct(val)}")
+    else: print(f"    {field_name}: {val}")
+
 def ioctls_from_header():
   hdr = (pathlib.Path(__file__).parent / "kfd_ioctl.h").read_text().replace("\\\n", "")
   pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_IOW?R?\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
@@ -67,6 +83,7 @@ def ioctl(fd, request, argp):
     if name == "AMDKFD_IOC_SVM":
       out = ctypes.cast(s.attrs, ctypes.POINTER(kfd_ioctl.struct_kfd_ioctl_svm_attribute))
       for i in range(s.nattr): print(f"{i}: {kfd_ioctl.enum_kfd_ioctl_svm_attr_type.get(out[i].type):40s}: {out[i].value:#x}")
+    if name == "AMDKFD_IOC_CREATE_QUEUE" and s.queue_type == kfd_ioctl.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL: print_aql_queue(s.read_pointer_address)
   else:
     print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : ioctl",
           f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", os.readlink(f"/proc/self/fd/{fd}") if fd >= 0 else "")

extra/hip_large_kernel.py

@@ -18,7 +18,7 @@ prg = dev.runtime("write_ones", mbin)
 prg(buf0._buf, global_size=(1,65537,1), local_size=(1,1,1), wait=True)
 
 import numpy as np
-def to_np(buf): return np.frombuffer(buf.as_buffer().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))
+def to_np(buf): return np.frombuffer(buf.as_memoryview().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))
 
 big = to_np(buf0)
 print(big)

extra/huggingface_onnx/huggingface_manager.py

@@ -8,14 +8,14 @@ from tinygrad.helpers import _ensure_downloads_dir
 DOWNLOADS_DIR = _ensure_downloads_dir() / "models"
 from tinygrad.helpers import tqdm
 
-def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, cache_dir: str|Path|None=None,
+def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, local_dir: str|Path|None=None,
                                  tries: int=2, **kwargs) -> Path:
   for attempt in range(tries):
     try:
       return Path(snapshot_download(
         repo_id=repo_id,
         allow_patterns=allow_patterns,
-        cache_dir=str(cache_dir) if cache_dir is not None else None,
+        local_dir=str(local_dir) if local_dir is not None else None,
         **kwargs
       ))
     except Exception as e:
@@ -144,14 +144,14 @@ class HuggingFaceONNXManager:
         root_path = snapshot_download_with_retry(
           repo_id=model_id,
           allow_patterns=allow_patterns,
-          cache_dir=str(self.models_dir)
+          local_dir=str(self.models_dir / model_id)
         )
 
         # Download config files (usually small)
         snapshot_download_with_retry(
           repo_id=model_id,
           allow_patterns=["*config.json"],
-          cache_dir=str(self.models_dir)
+          local_dir=str(self.models_dir / model_id)
         )
 
         model_data["download_path"] = str(root_path)

extra/huggingface_onnx/run_models.py

@@ -88,8 +88,8 @@ if __name__ == "__main__":
       # repo id
       # validates all onnx models inside repo
       repo_id = "/".join(path)
-      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=DOWNLOADS_DIR)
-      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], local_dir=DOWNLOADS_DIR / repo_id)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], local_dir=DOWNLOADS_DIR / repo_id)
       config = get_config(root_path)
       for onnx_model in root_path.rglob("*.onnx"):
         rtol, atol = get_tolerances(onnx_model.name)
@@ -101,8 +101,8 @@ if __name__ == "__main__":
       onnx_model = path[-1]
       assert path[-1].endswith(".onnx")
       repo_id, relative_path = "/".join(path[:2]), "/".join(path[2:])
-      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=DOWNLOADS_DIR)
-      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], local_dir=DOWNLOADS_DIR / repo_id)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], local_dir=DOWNLOADS_DIR / repo_id)
       config = get_config(root_path)
       rtol, atol = get_tolerances(onnx_model)
       print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")

extra/mmapeak/mmapeak.py

@@ -1,99 +1,105 @@
-import os, pathlib
+import os
 
 # TODO: there is a timing bug without this
 os.environ["AMD_AQL"] = "1"
 
-from tinygrad.device import Device
-from tinygrad.runtime.ops_amd import AMDProgram, HIPCompiler
+from tinygrad import Tensor, Device
+from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from tinygrad.renderer.amd.dsl import Reg, Inst, s, v
 
 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
-NUM_WAVES = 2
+NUM_WAVES = 4
 FLOPS_PER_MATMUL = 16*16*16*2
-INTERNAL_LOOP = 1_000_00
+INTERNAL_LOOP = getenv("LOOP", 10_000)
 INSTRUCTIONS_PER_LOOP = 200
-DIRECTIVE = ".amdhsa_wavefront_size32 1"
 
-assemblyTemplate = (pathlib.Path(__file__).parent / "template.s").read_text()
+def repeat(insts:list[Inst], n:int, counter_sreg:Reg) -> list[Inst]:
+  insts_bytes = b"".join([inst.to_bytes() for inst in insts])
+  sub_inst, cmp_inst = s_sub_u32(counter_sreg, counter_sreg, 1), s_cmp_lg_i32(counter_sreg, 0)
+  loop_sz = len(insts_bytes) + sub_inst.size() + cmp_inst.size()
+  branch_inst = s_cbranch_scc1(simm16=-((loop_sz // 4) + 1) & 0xFFFF)
+  return [s_mov_b32(counter_sreg, n)] + insts + [sub_inst, cmp_inst, branch_inst, s_endpgm()]
 
-def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, extra=""):
+def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs):
   if accum:
-    instructions = "{} a[0:{}], v[{}:{}], v[{}:{}], 1{}\n".format(instruction, vgprIndices[0],
-                                                                vgprIndices[1], vgprIndices[2],
-                                                                vgprIndices[1], vgprIndices[2], extra)
+    inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[1]:vgprIndices[2]], 1, acc_cd=1, **kwargs)
   elif dense:
-    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], 1\n".format(instruction, vgprIndices[0],
-                                                                vgprIndices[1], vgprIndices[2],
-                                                                vgprIndices[1], vgprIndices[2])
+    inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[1]:vgprIndices[2]], 1)
   else:
-    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], v{}\n".format(instruction, vgprIndices[0],
-                                                                  vgprIndices[1], vgprIndices[2],
-                                                                  vgprIndices[3], vgprIndices[4],
-                                                                  vgprIndices[5])
-  src = assemblyTemplate.replace("INTERNAL_LOOP", str(INTERNAL_LOOP)).replace("INSTRUCTION", instructions*INSTRUCTIONS_PER_LOOP)
-  src = src.replace("DIRECTIVE", DIRECTIVE)
-  lib = COMPILER.compile(src)
-  fxn = AMDProgram(DEV, "matmul", lib)
-  elapsed = min([fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) for _ in range(2)])
+    inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[3]:vgprIndices[4]], v[vgprIndices[5]])
+  insts = repeat([inst for _ in range(INSTRUCTIONS_PER_LOOP)], n=INTERNAL_LOOP, counter_sreg=s[1])
+  def fxn(A:UOp) -> UOp:
+    threads = UOp.special(WAVE_SIZE * NUM_WAVES, "lidx0")
+    gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
+    FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
+    sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  dummy = Tensor.zeros(1).contiguous().realize()
+  out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
+  ei = out.schedule()[-1].lower()
+  elapsed = min([ei.run(wait=True) for _ in range(2)])
   FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
-  print(f"{instruction:<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")
+  print(f"{inst.op_name.lower():<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")
 
 if __name__=="__main__":
-  DEVICENUM = os.getenv("DEVICENUM", "0")
-  try:
-    DEV = Device['AMD:' + DEVICENUM]
-  except:
-    raise RuntimeError("Error while initiating AMD device")
+  DEV = Device[Device.DEFAULT]
+  arch = DEV.renderer.arch
 
-  COMPILER = HIPCompiler(DEV.arch)
-  if DEV.arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
-    if DEV.arch == 'gfx1103': NUM_WORKGROUPS = 8
-    if DEV.arch == 'gfx1151': NUM_WORKGROUPS = 32
-    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
-    launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
-    launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))
-    launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,15))
-    launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,9))
-    launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,11))
-  elif DEV.arch == 'gfx1201':
+  if arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
+    from tinygrad.runtime.autogen.amd.rdna3.ins import *
+    if arch == 'gfx1103': NUM_WORKGROUPS = 8
+    if arch == 'gfx1151': NUM_WORKGROUPS = 32
+    launchBenchmark(v_wmma_bf16_16x16x16_bf16, (7,8,15))
+    launchBenchmark(v_wmma_f16_16x16x16_f16, (7,8,15))
+    launchBenchmark(v_wmma_f32_16x16x16_bf16, (7,8,15))
+    launchBenchmark(v_wmma_f32_16x16x16_f16, (7,8,15))
+    launchBenchmark(v_wmma_i32_16x16x16_iu4, (7,8,9))
+    launchBenchmark(v_wmma_i32_16x16x16_iu8, (7,8,11))
+  elif arch in {'gfx1200', 'gfx1201'}:
+    from tinygrad.runtime.autogen.amd.rdna4.ins import *
+    # this instruction does not exist in the rdna4 isa, use the co version
+    s_sub_u32 = s_sub_co_u32
     NUM_WORKGROUPS = 64
-    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (3,4,7))
-    launchBenchmark("v_wmma_f16_16x16x16_f16", (3,4,7))
-    launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,11))
-    launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,11))
-    launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,8))
-    launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,9))
-    launchBenchmark("v_wmma_f32_16x16x16_fp8_fp8", (7,8,9))
-    launchBenchmark("v_wmma_f32_16x16x16_fp8_bf8", (7,8,9))
-    launchBenchmark("v_wmma_f32_16x16x16_bf8_fp8", (7,8,9))
-    launchBenchmark("v_wmma_f32_16x16x16_bf8_bf8", (7,8,9))
+    launchBenchmark(v_wmma_bf16_16x16x16_bf16, (3,4,7))
+    launchBenchmark(v_wmma_f16_16x16x16_f16, (3,4,7))
+    launchBenchmark(v_wmma_f32_16x16x16_bf16, (7,8,11))
+    launchBenchmark(v_wmma_f32_16x16x16_f16, (7,8,11))
+    launchBenchmark(v_wmma_i32_16x16x16_iu4, (7,8,8))
+    launchBenchmark(v_wmma_i32_16x16x16_iu8, (7,8,9))
+    launchBenchmark(v_wmma_f32_16x16x16_fp8_fp8, (7,8,9))
+    launchBenchmark(v_wmma_f32_16x16x16_fp8_bf8, (7,8,9))
+    launchBenchmark(v_wmma_f32_16x16x16_bf8_fp8, (7,8,9))
+    launchBenchmark(v_wmma_f32_16x16x16_bf8_bf8, (7,8,9))
     FLOPS_PER_MATMUL = 16*16*32*2
-    launchBenchmark("v_wmma_i32_16X16X32_iu4", (7,8,9))
-    launchBenchmark("v_swmmac_f32_16x16x32_f16", (7,8,11,12,19,20), False)
-    launchBenchmark("v_swmmac_f32_16x16x32_bf16", (7,8,11,12,19,20), False)
-    launchBenchmark("v_swmmac_f16_16x16x32_f16", (3,4,7,8,15,16), False)
-    launchBenchmark("v_swmmac_bf16_16x16x32_bf16", (3,4,7,8,15,16), False)
-    launchBenchmark("v_swmmac_i32_16x16x32_iu8", (7,8,9,10,13,14), False)
-    launchBenchmark("v_swmmac_i32_16x16x32_iu4", (7,8,8,9,10,11), False)
-    launchBenchmark("v_swmmac_f32_16x16x32_fp8_fp8", (7,8,9,10,13,14), False)
-    launchBenchmark("v_swmmac_f32_16x16x32_fp8_bf8", (7,8,9,10,13,14), False)
-    launchBenchmark("v_swmmac_f32_16x16x32_bf8_fp8", (7,8,9,10,13,14), False)
-    launchBenchmark("v_swmmac_f32_16x16x32_bf8_bf8", (7,8,9,10,13,14), False)
+    launchBenchmark(v_wmma_i32_16x16x32_iu4, (7,8,9))
+    launchBenchmark(v_swmmac_f32_16x16x32_f16, (7,8,11,12,19,20), False)
+    launchBenchmark(v_swmmac_f32_16x16x32_bf16, (7,8,11,12,19,20), False)
+    launchBenchmark(v_swmmac_f16_16x16x32_f16, (3,4,7,8,15,16), False)
+    launchBenchmark(v_swmmac_bf16_16x16x32_bf16, (3,4,7,8,15,16), False)
+    launchBenchmark(v_swmmac_i32_16x16x32_iu8, (7,8,9,10,13,14), False)
+    launchBenchmark(v_swmmac_i32_16x16x32_iu4, (7,8,8,9,10,11), False)
+    launchBenchmark(v_swmmac_f32_16x16x32_fp8_fp8, (7,8,9,10,13,14), False)
+    launchBenchmark(v_swmmac_f32_16x16x32_fp8_bf8, (7,8,9,10,13,14), False)
+    launchBenchmark(v_swmmac_f32_16x16x32_bf8_fp8, (7,8,9,10,13,14), False)
+    launchBenchmark(v_swmmac_f32_16x16x32_bf8_bf8, (7,8,9,10,13,14), False)
     FLOPS_PER_MATMUL = 16*16*64*2
-    launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
-  elif DEV.arch == 'gfx950':
-    DIRECTIVE = ".amdhsa_accum_offset 4"
+    launchBenchmark(v_swmmac_i32_16x16x64_iu4, (7,8,9,10,13,14), False)
+  elif arch == 'gfx950':
+    from tinygrad.runtime.autogen.amd.cdna.ins import *
     NUM_WORKGROUPS = 256
     WAVE_SIZE = 64
     NUM_WAVES = 4
-    launchBenchmark("v_mfma_f32_16x16x16_f16", (3,0,1), accum=True)
-    launchBenchmark("v_mfma_f32_16x16x16_bf16", (3,0,1), accum=True)
+    launchBenchmark(v_mfma_f32_16x16x16_f16, (3,0,1), accum=True)
+    launchBenchmark(v_mfma_f32_16x16x16_bf16, (3,0,1), accum=True)
     FLOPS_PER_MATMUL = 16*16*32*2
-    launchBenchmark("v_mfma_f32_16x16x32_f16", (3,0,3), accum=True)
-    launchBenchmark("v_mfma_f32_16x16x32_bf16", (3,0,3), accum=True)
+    launchBenchmark(v_mfma_f32_16x16x32_f16, (3,0,3), accum=True)
+    launchBenchmark(v_mfma_f32_16x16x32_bf16, (3,0,3), accum=True)
     FLOPS_PER_MATMUL = 16*16*128*2
-    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,7), accum=True) # fp8
-    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,5), accum=True, extra=", cbsz:2 blgp:2") # fp6
-    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,3), accum=True, extra=", cbsz:4 blgp:4") # fp4
+    launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,7), accum=True) # fp8
+    launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,5), accum=True, cbsz=2, blgp=2) # fp6
+    launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,3), accum=True, cbsz=4, blgp=4) # fp4
   else:
-    raise RuntimeError(f"arch {DEV.arch} not supported.")
+    raise RuntimeError(f"arch {arch} not supported.")

extra/mmapeak/template.s

@@ -1,40 +0,0 @@
-    .text
-    .globl matmul
-    .p2align 8
-    .type matmul,@function
-matmul:
-  s_mov_b32 s1, INTERNAL_LOOP
-  s_mov_b32 s2, 0
-  inner_loop:
-    INSTRUCTION
-    s_sub_u32 s1, s1, 1
-    s_cmp_lg_i32 s1, s2
-    s_cbranch_scc1 inner_loop
-  s_endpgm
-
-.rodata
-.p2align 6
-.amdhsa_kernel matmul
-  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
-  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  DIRECTIVE
-.end_amdhsa_kernel
-
-.amdgpu_metadata
----
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: matmul
-    .symbol: matmul.kd
-    .kernarg_segment_size: 0
-    .group_segment_fixed_size: 0
-    .private_segment_fixed_size: 0
-    .kernarg_segment_align: 4
-    .wavefront_size: 32
-    .sgpr_count: 8
-    .vgpr_count: 32
-    .max_flat_workgroup_size: 1024
-...
-.end_amdgpu_metadata

extra/models/clip.py

@@ -5,7 +5,6 @@ from tinygrad.nn import Linear, LayerNorm, Embedding, Conv2d
 from typing import List, Optional, Union, Tuple, Dict
 from abc import ABC, abstractmethod
 from functools import lru_cache
-from PIL import Image
 import numpy as np
 import re, gzip
 
@@ -444,7 +443,8 @@ class OpenClipEncoder:
   # TODO:
   # Should be doable in pure tinygrad, would just require some work and verification.
   # This is very desirable since it would allow for full generation->evaluation in a single JIT call.
-  def prepare_image(self, image:Image.Image) -> Tensor:
+  def prepare_image(self, image) -> Tensor:
+    from PIL import Image
     SIZE = 224
     w, h = image.size
     scale = min(SIZE / h, SIZE / w)

extra/models/llama.py

@@ -41,9 +41,13 @@ class Attention:
     self.n_rep = self.n_heads // self.n_kv_heads
     self.max_context = max_context
 
-    self.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
-    self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-    self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+    if getenv("WQKV"):
+      self.wqkv = linear(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2, bias=False)
+    else:
+      self.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
+      self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+      self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+
     self.wo = linear(self.n_heads * self.head_dim, dim, bias=False)
 
     self.q_norm = nn.RMSNorm(dim, qk_norm) if qk_norm is not None else None
@@ -51,16 +55,21 @@ class Attention:
 
   def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]=None) -> Tensor:
     if getenv("WQKV"):
-      if not hasattr(self, 'wqkv'): self.wqkv = Tensor.cat(self.wq.weight, self.wk.weight, self.wv.weight)
-      xqkv = x @ self.wqkv.T
-      xq, xk, xv = xqkv.split([self.wq.weight.shape[0], self.wk.weight.shape[0], self.wv.weight.shape[0]], dim=2)
+      xqkv = self.wqkv(x)
+      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
+      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
     else:
-      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+      xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)
 
     if self.q_norm is not None and self.k_norm is not None:
       xq = self.q_norm(xq)
       xk = self.k_norm(xk)
 
+    # cast_float_to_bf16 is expensive in reduction loops, break it out
+    if x.dtype == dtypes.bfloat16: xq, xk = xq.contiguous_backward(), xk.contiguous_backward()
+
     xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
     xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
     xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
@@ -86,20 +95,23 @@ class Attention:
       assert start_pos == 0
       keys, values = xk, xv
 
-    keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
-    xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
-    attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
+    if self.max_context:
+      keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
+      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
+    else:
+      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(keys, values, is_causal=True, enable_gqa=True).transpose(1, 2)
     if getenv("STUB_ATTENTION"):
-      # TODO: do we need mask?
       from tinygrad.uop.ops import UOp, KernelInfo
       def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
         return UOp.sink(arg=KernelInfo(name="fa_custom_forward"))
       def fa_custom_backward(out_q:UOp, out_k:UOp, out_v:UOp, grad:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
         return UOp.sink(arg=KernelInfo(name="fa_custom_backward"))
       def fa_backward(grad:UOp, kernel:UOp) -> tuple[None, UOp, UOp, UOp]:
-        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[1]))
-        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[2]))
-        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[3]))
+        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[2]))
+        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[3]))
+        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[4]))
         ck = Tensor.custom_kernel(grad_q, grad_k, grad_v, Tensor(grad), q, k, v, fxn=fa_custom_backward)[:3]
         return (None, ck[0].uop, ck[1].uop, ck[2].uop)
       attn = Tensor.empty_like(attn).custom_kernel(xq, keys, values, fxn=fa_custom_forward, grad_fxn=fa_backward)[0]
@@ -194,12 +206,14 @@ class Transformer:
 
   def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):
     _bsz, seqlen = tokens.shape
-    h = self.tok_embeddings(tokens)
+    h = self.tok_embeddings(tokens).contiguous()
     freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :]
 
-    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
+    if self.max_context != 0 and seqlen > 1:
+      mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1)
+    else: mask = None
     for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
-    logits = self.output(self.norm(h))
+    logits = self.output(self.norm(h).contiguous().contiguous_backward()).contiguous_backward()
     if math.isnan(temperature): return logits
 
     return sample(logits[:, -1, :].flatten(), temperature, top_k, top_p, alpha_f, alpha_p)