more test

* viz: variable length rdna barriers

* work

* tiny changes

* simple wave simd test

* small wave sync test

* good multi barrier bug find

* simple fix

* wave_sync asserts

* rdna4 work

* more rdna4

* find more bugs in my model

* it's so much simpler

* wave_sync tests duration

* r4

* should just call this rdna4

.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

@@ -45,6 +45,10 @@ inputs:
     description: "Install mesa"
     required: false
     default: 'false'
+  tinydreno:
+    description: "Install tinydreno"
+    required: false
+    default: 'false'
 runs:
   using: "composite"
   steps:
@@ -56,32 +60,40 @@ 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 }}-${{ runner.arch }}-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:
         path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
 
     # **** 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 +104,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 +149,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,12 +194,18 @@ 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 }}-${{ runner.arch }}-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/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
 
     - name: Run apt Update + Install
       if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
@@ -219,7 +237,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 +257,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 +276,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
@@ -303,3 +330,9 @@ runs:
       if: inputs.mesa == 'true' && runner.os == 'macOS'
       shell: bash
       run: brew install sirhcm/tinymesa/tinymesa_cpu
+
+    # *** tinydreno ***
+    - name: Install tinydreno (linux)
+      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so

.github/workflows/autogen.yml

@@ -32,6 +32,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
+        key: 'autogen'
         opencl: 'true'
         amd: 'true'
         cuda: 'true'
@@ -40,14 +41,14 @@ 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
+      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: |
-        find tinygrad/runtime/autogen -type f -name "*.py" -not -name "__init__.py" -not -name "comgr_3.py" -not -name "metal.py" -not -name "libclang.py" -delete
+        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"
         python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
-        python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd"
-        python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v14_0_2"
+        python3 -c "from tinygrad.runtime.autogen import 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"
         python3 -c "from tinygrad.runtime.autogen import llvm"
         python3 -c "from tinygrad.runtime.autogen import webgpu"
@@ -59,8 +60,9 @@ jobs:
     - name: Check for differences
       run: |
         if ! git diff --quiet; then
+          git diff
           git diff > autogen-ubuntu.patch
-          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          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
@@ -80,16 +82,18 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
+        key: 'autogen-mac'
         llvm: 'true'
     - name: Regenerate autogen files
       run: |
-        rm tinygrad/runtime/autogen/metal.py
-        LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
+        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 files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          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
@@ -99,8 +103,8 @@ jobs:
         name: autogen-macos-patch
         path: autogen-macos.patch
 
-  autogen-comgr-3:
-    name: In-tree Autogen (comgr 3)
+  autogen-comgr-2:
+    name: In-tree Autogen (comgr 2)
     runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
@@ -108,29 +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: Regenerate autogen files
       run: |
-        rm tinygrad/runtime/autogen/comgr_3.py
-        python3 -c "from tinygrad.runtime.autogen import comgr_3"
+        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 > autogen-comgr3.patch
-          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          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-comgr3-patch
-        path: autogen-comgr3.patch
+        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:
@@ -145,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
@@ -291,13 +332,13 @@ 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 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
+      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 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
     - 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
     - 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
+      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
     # - name: Run 10 CIFAR training steps w winograd
@@ -332,9 +373,9 @@ jobs:
     - 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: Insert amdgpu
     #  run: sudo modprobe amdgpu
     - name: Symlink models and datasets
@@ -444,9 +485,9 @@ jobs:
     - 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
@@ -469,7 +510,7 @@ jobs:
     - 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
     - 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
+      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
     # TODO: too slow
@@ -479,6 +520,9 @@ jobs:
       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
+    # 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
 
@@ -496,9 +540,9 @@ jobs:
     - 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
@@ -561,7 +605,7 @@ jobs:
     - name: openpilot compile3 0.10.1 driving_policy
       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
@@ -573,6 +617,27 @@ jobs:
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
+  testcommausbgpubenchmark:
+    name: UsbGPU Benchmark (comma)
+    runs-on: [self-hosted, Linux, comma4]
+    timeout-minutes: 20
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot load_pickle 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
+
   testreddriverbenchmark:
     name: AM Benchmark
     runs-on: [self-hosted, Linux, tinyboxrandom]
@@ -587,9 +652,9 @@ jobs:
     - 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
@@ -651,9 +716,9 @@ jobs:
     - 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

.github/workflows/test.yml

@@ -1,11 +1,11 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '15'
+  CACHE_VERSION: '18'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
-  IGNORE_OOB: 0
+  CHECK_OOB: 1
 
 on:
   push:
@@ -26,19 +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:
-      IGNORE_OOB: 1
+      CHECK_OOB: 0
     steps:
     - name: Checkout Code
       uses: actions/checkout@v4
@@ -98,7 +98,7 @@ 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
@@ -106,7 +106,7 @@ jobs:
         sudo apt update || true
         sudo apt install -y --no-install-recommends ninja-build
     - name: Test one op
-      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_tiny.py TestTiny.test_plus
     - name: Test ResNet-18
       run: DEBUG=2 python3 extra/torch_backend/example.py
     - name: custom tests
@@ -114,7 +114,7 @@ jobs:
     - 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
@@ -134,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: |
@@ -156,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
@@ -197,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
@@ -207,11 +207,11 @@ 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
@@ -229,19 +229,51 @@ 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.14.10
-        pre-commit run ruff --all-files
         python3 -m ruff check examples/mlperf/ --ignore E501
         python3 -m ruff check extra/thunder/tiny/ --ignore E501 --ignore F841 --ignore E722
         python3 -m ruff check extra/torch_backend/backend.py
-    - name: Run mypy
+    - name: Run mypy with lineprecision report
       run: |
         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
+
+  nulltest:
+    name: Null Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: unittest-13
+        pydeps: "pillow ftfy regex pre-commit"
+        deps: testing_unit
+        llvm: 'true'
+        amd: 'true'
+    - 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.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 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
 
   unittest:
     name: Unit Tests
@@ -254,28 +286,19 @@ 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
-    - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
-    # TODO: too slow
-    # - 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
-    # 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
+        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 GC tests
       run: python test/external/external_uop_gc.py
     - name: External Benchmark Schedule
@@ -289,8 +312,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 +333,7 @@ jobs:
         deps: testing_unit
         python-version: '3.14'
     - name: Test SPEC=2
-      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --ignore test/unit/test_autogen.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 +367,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 +388,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 +443,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 +454,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 +471,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 +484,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
 
@@ -474,7 +497,7 @@ jobs:
     runs-on: ubuntu-24.04
     timeout-minutes: 15
     env:
-      IGNORE_OOB: 1
+      CHECK_OOB: 0
     steps:
       - name: Checkout Code
         uses: actions/checkout@v4
@@ -522,7 +545,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
@@ -541,15 +564,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)
@@ -562,8 +585,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
@@ -575,15 +598,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)
@@ -596,32 +619,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 --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_rocprof_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,70 +727,20 @@ 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
 
-  testamdasm:
-    name: AMD ASM IDE
-    runs-on: ubuntu-24.04
-    timeout-minutes: 10
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: rdna3-emu
-          deps: testing_minimal
-          amd: 'true'
-          python-version: '3.13'
-      - name: Verify AMD autogen is up to date
-        run: |
-          python -m extra.assembly.amd.amdxml
-          git diff --exit-code extra/assembly/amd/autogen/
-      - name: Install LLVM 21
-        run: |
-          wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
-          echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
-          sudo apt-get update
-          sudo apt-get install llvm-21 llvm-21-tools cloc
-      - name: RDNA3 Line Count
-        run: cloc --by-file extra/assembly/amd/*.py
-      - name: Install rocprof-trace-decoder
-        run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
-      - name: Run RDNA3 emulator tests
-        run: python -m pytest -n=auto extra/assembly/amd/ --durations 20
-      - name: Run RDNA3 emulator tests (AMD_LLVM=1)
-        run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
-      - name: Run RDNA3 dtype tests
-        run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
-      - name: Run RDNA3 dtype tests (AMD_LLVM=1)
-        run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
-      # TODO: run all once emulator is faster
-      - name: Run RDNA3 ops tests
-        run: SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril"
-
   testnvidia:
     strategy:
       fail-fast: false
@@ -713,7 +760,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
@@ -721,10 +768,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
 
@@ -744,7 +793,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' }}
@@ -753,11 +802,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
 
@@ -775,27 +824,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
-      env:
-        LIBCLANG_PATH: '/opt/homebrew/opt/llvm@20/lib/libclang.dylib'
       run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run 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
@@ -818,6 +869,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
@@ -836,14 +889,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
@@ -875,8 +928,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
@@ -886,7 +938,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
@@ -919,12 +971,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 ******
 
@@ -943,15 +999,38 @@ 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
+  qcomclcompiletests:
+    name: Compile-only (QCOM CL)
+    runs-on: ubuntu-24.04-arm
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: compile-qcomcl
+          deps: testing_unit
+          tinydreno: 'true'
+          python-version: '3.12'
+      - name: Set env
+        shell: bash
+        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\nNULL_QCOMCL=1" >> $GITHUB_ENV
+      - name: Run test_ops
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
+          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

@@ -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.
-
-## Testing
-
-```bash
-# Run specific test
-python -m pytest test/unit/test_schedule_cache.py -xvs
-
-# Run with timeout
-python -m pytest test/test_symbolic_ops.py -x --timeout=60
-
-# Debug with print
-DEBUG=2 python -m pytest test/test_schedule.py::test_name -xvs
-
-# Visualize UOp graphs
-VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
-```
-
-## Common Environment Variables
-
-- `DEBUG=1-7` - Increasing verbosity (7 shows assembly output)
-- `VIZ=1` - Enable graph visualization
-- `SPEC=1` - Enable UOp spec verification
-- `NOOPT=1` - Disable optimizations
-- `DEVICE=CPU/CUDA/AMD/METAL` - Set default device
-
-## Debugging Tips
-
-1. **Print UOp graphs**: `print(tensor.uop)` or `print(tensor.uop.sink())`
-2. **Check schedule**: `tensor.schedule()` returns list of ExecItems
-3. **Trace graph rewrites**: Use `VIZ=1` or add print in PatternMatcher callbacks
-4. **Find UOps by type**: `[u for u in uop.toposort() if u.op is Ops.SOMETHING]`
-
-## Workflow Rules
-
-- **NEVER commit without explicit user approval** - always show the diff and wait for approval
-- **NEVER amend commits** - always create a new commit instead
-- Run `pre-commit run --all-files` before committing to catch linting/type errors
-- Run tests before proposing commits
-- Test with `SPEC=2` when modifying UOp-related code
-
-## Auto-generated Files (DO NOT EDIT)
-
-The following files are auto-generated and should never be edited manually:
-- `extra/assembly/amd/autogen/{arch}/__init__.py` - Generated by `python -m extra.assembly.amd.dsl --arch {arch}`
-- `extra/assembly/amd/autogen/{arch}/gen_pcode.py` - Generated by `python -m extra.assembly.amd.pcode --arch {arch}`
-
-Where `{arch}` is one of: `rdna3`, `rdna4`, `cdna`
-
-To add missing instruction implementations, add them to `extra/assembly/amd/emu.py` instead.
-
-## Style Notes
-
-- 2-space indentation, 150 char line limit
-- PatternMatchers should be defined at module level (slow to construct)
-- Prefer `graph_rewrite` over manual graph traversal
-- UOp methods like `.replace()` preserve tags unless explicitly changed
-- Use `.rtag(value)` to add tags to UOps
-
-## Lessons Learned
-
-### UOp ucache Behavior
-UOps are cached by their contents - creating a UOp with identical (op, dtype, src, arg) returns the **same object**. This means:
-- `uop.replace(tag=None)` on a tagged UOp returns the original untagged UOp if it exists in cache
-- Two UOps with same structure are identical (`is` comparison works)
-
-### Spec Validation
-When adding new UOp patterns, update `tinygrad/uop/spec.py`. Test with:
-```bash
-SPEC=2 python3 test/unit/test_something.py
-```
-Spec issues appear as `RuntimeError: SPEC ISSUE None: UOp(...)`.
-
-### Schedule Cache Key Normalization
-The schedule cache strips values from BIND nodes so different bound values (e.g., KV cache positions) hit the same cache entry:
-- `pm_pre_sched_cache`: BIND(DEFINE_VAR, CONST) → BIND(DEFINE_VAR) for cache key
-- `pm_post_sched_cache`: restores original BIND from context
-- When accessing `bind.src[1]`, check `len(bind.src) > 1` first (might be stripped)
-- Extract var_vals from `input_buffers` dict after graph_rewrite (avoids extra toposort)
-
-### Avoiding Extra Work
-- Use ctx dict from graph_rewrite to collect info during traversal instead of separate toposort
-- Only extract var_vals when schedule is non-empty (no kernels = no vars needed)
-- PatternMatchers are slow to construct - define at module level, not in functions
-
-### Readability Over Speed
-Don't add complexity for marginal performance gains. Simpler code that's slightly slower is often better:
-```python
-# BAD: "optimized" with extra complexity
-if has_afters:  # skip toposort if no AFTERs
-  after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
-
-# GOOD: simple, always works
-after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
-```
-The conditional check adds complexity, potential bugs, and often negligible speedup. Only optimize when profiling shows a real bottleneck.
-
-### Testing LLM Changes
-```bash
-# Quick smoke test
-echo "Hello" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b"
-
-# Check cache hits (should see "cache hit" after warmup)
-echo "Hello world" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b" 2>&1 | grep cache
-
-# Test with beam search
-echo "Hello" | BEAM=2 python tinygrad/apps/llm.py --model "llama3.2:1b"
-```
-
-## Common Patterns
-
-### Graph Transformation
-```python
-def my_transform(ctx, x):
-  # Return new UOp or None to skip
-  return x.replace(arg=new_arg)
-
-pm = PatternMatcher([
-  (UPat(Ops.SOMETHING, name="x"), my_transform),
-])
-result = graph_rewrite(input_uop, pm, ctx={})
-```
-
-### Finding Variables
-```python
-# Get all variables in a UOp graph
-variables = uop.variables()
-
-# Get bound variable values
-var, val = bind_uop.unbind()
-```
-
-### Shape Handling
-```python
-# Shapes can be symbolic (contain UOps)
-shape = tensor.shape  # tuple[sint, ...] where sint = int | UOp
-```
-
-## Performance Optimization
-
-When optimizing tinygrad internals:
-
-1. **Measure wall time, not just call counts** - Reducing `graph_rewrite` calls doesn't always improve wall time. The overhead of conditional checks can exceed the cost of the operation being skipped.
-
-2. **Profile each optimization individually** - Run benchmarks with and without each change to measure actual impact. Use `test/external/external_benchmark_schedule.py` for schedule/rewrite timing.
-
-3. **Early exits in hot paths are effective** - Simple checks like `if self.op is Ops.CONST: return self` in `simplify()` can eliminate many unnecessary `graph_rewrite` calls.
-
-4. **`graph_rewrite` is expensive** - Each call has overhead even for small graphs. Avoid calling it when the result is trivially known (e.g., simplifying a CONST returns itself).
-
-5. **Beware iterator overhead** - Checks like `all(x.op is Ops.CONST for x in self.src)` can be slower than just running the operation, especially for small sequences.
-
-6. **Verify cache hit rates before adding/keeping caches** - Measure actual hit rates with real workloads. A cache with 0% hit rate is pure overhead (e.g., `pm_cache` was removed because the algorithm guarantees each UOp is only passed to `pm_rewrite` once).
-
-7. **Use `TRACK_MATCH_STATS=2` to profile pattern matching** - This shows match rates and time per pattern. Look for patterns with 0% match rate that still cost significant time - these are pure overhead for that workload.
-
-8. **Cached properties beat manual traversal** - `backward_slice` uses `@functools.cached_property`. A DFS with early-exit sounds faster but is actually slower because it doesn't benefit from caching. The cache hit benefit often outweighs algorithmic improvements.
-
-9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.
-
-## Pattern Matching Analysis
-
-**Use the right tool:**
-
-- `TRACK_MATCH_STATS=2` - **Profiling**: identify expensive patterns
-- `VIZ=-1` - **Inspection**: see all transformations, what every match pattern does, the before/after diffs
-
-```bash
-TRACK_MATCH_STATS=2 PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
-```
-
-Output format: `matches / attempts -- match_time / total_time ms -- location`
-
-Key patterns to watch (from ResNet50 benchmark):
-- `split_load_store`: ~146ms, 31% match rate - does real work
-- `simplify_valid`: ~75ms, 0% match rate in this workload - checks AND ops for INDEX in backward slice
-- `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches
-
-Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.
-
-```bash
-# Save the trace
-VIZ=-1 python test/test_tiny.py TestTiny.test_gemm
-
-# Explore it
-./extra/viz/cli.py --help
-```
-
-## AMD Performance Counter Profiling
-
-Set VIZ to `-2` to save performance counters traces for the AMD backend.
-
-Use the CLI in `./extra/sqtt/roc.py` to explore the trace.

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

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)
 
@@ -264,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)
@@ -379,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)
@@ -406,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}")
@@ -552,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))
 
@@ -569,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])
@@ -578,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:
@@ -637,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
 
@@ -780,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/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, DEBUG
 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
 
@@ -1282,21 +1282,29 @@ def train_bert():
         previous_step = i
 
 def train_llama3():
-  from extra.models.llama import Transformer
+  from examples.mlperf.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)
@@ -1310,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")
@@ -1325,8 +1335,16 @@ 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 (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
-  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+  print(f"model parameters: {model_params}")
+
+  # pad vocab
+  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
+
+  model = Transformer(**model_params, max_context=SEQLEN)
+
   params = get_parameters(model)
   # weights are all bfloat16 for now
   assert params and all(p.dtype == dtypes.bfloat16 for p in params)
@@ -1340,6 +1358,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():
@@ -1347,6 +1367,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)
@@ -1359,8 +1380,23 @@ 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()
+
+  is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
+  is_fake_offload = Device.DEFAULT == "NULL"
+  optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_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
+  if is_offload_optim:
+    for p in optim.params:
+      p.grad = Tensor.zeros(p.shape, dtype=p.dtype, device=optim_device, requires_grad=False).contiguous().realize()
+  else:
+    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"):
@@ -1373,98 +1409,139 @@ 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):
     if (DP := getenv("DP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
+      tokens = tokens.to(None).shard(device, 0)
     if (MP := getenv("MP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
       tokens = tokens.shard(device)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    if DP == 1 and MP == 1: tokens = tokens.to(None)
+    logits:Tensor = model(tokens[:, :-1])
+    loss = vocab_mask.where(-1e9, 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))
+    loss_cpu = loss.flatten().float().to("CPU")
+    Tensor.realize(loss_cpu, *grads)
+    return loss_cpu
 
-    optim.step()
+  @TinyJit
+  def optim_step():
+    grad_norm = optim.fstep(grads)
     scheduler.step()
 
-    lr = optim.lr
-    loss.realize(lr)
-    return loss, lr
+    for g in grads:
+      g.assign(g.zeros_like())
+
+    lr_cpu = optim.lr.float().to("CPU")
+    grad_norm_cpu = grad_norm.float().to("CPU")
+    Tensor.realize(lr_cpu, grad_norm_cpu, *grads)
+
+    return lr_cpu, grad_norm_cpu
 
   @TinyJit
   @Tensor.train(False)
-  def eval_step(model, tokens:Tensor):
+  def eval_step(tokens:Tensor):
     if (DP := getenv("DP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
+      tokens = tokens.to(None).shard(device, 0)
     if (MP := getenv("MP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
       tokens = tokens.shard(device)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten().float()
+    if DP == 1 and MP == 1: tokens = tokens.to(None)
+    logits:Tensor = model(tokens[:, :-1])
+    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    return loss.flatten().float().to("CPU")
 
   # ** data iters **
   def fake_data(bs, samples):
+    import numpy as np
     for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
+      fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
+      yield Tensor(fake_data_np, device="NPY")
 
   def get_train_iter():
     if getenv("FAKEDATA", 0):
       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()
+    actual_gbs = GBS if i >= 2 else BS
     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 if i >= 2 else 1):
+        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()
+      ret = optim_step()
+      lr, grad_norm = ret[0].item(), ret[1].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 += actual_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, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
+      if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
 
       if WANDB:
-        wandb.log({"lr": lr, "train/loss": loss, "train/step_time": sec, "train/GFLOPS": gflops, "train/sequences_seen": sequences_seen})
+        wandb.log({
+          "train/loss": loss,
+          "train/lr": lr,
+          "train/grad_norm": grad_norm,
+          "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")
@@ -1476,16 +1553,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 != (sequences_seen - actual_gbs) // EVAL_FREQ 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")
+      tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences")
+
+      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
 
-      for tokens in tqdm(eval_iter, total=5760//EVAL_BS):
-        eval_losses += eval_step(model, tokens).tolist()
       log_perplexity = Tensor(eval_losses).mean().float().item()
 
       tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
@@ -1571,7 +1661,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():
@@ -1610,7 +1700,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/models/llama.py

@@ -0,0 +1,80 @@
+from tinygrad import Tensor, nn
+from tinygrad.helpers import getenv
+from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+
+class Attention:
+  def __init__(self, dim:int, n_heads:int, n_kv_heads:int|None=None, linear=nn.Linear):
+    self.n_heads = n_heads
+    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
+    self.head_dim = dim // n_heads
+    self.n_rep = self.n_heads // self.n_kv_heads
+
+    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)
+
+  def __call__(self, x:Tensor, freqs_cis:Tensor) -> Tensor:
+    if getenv("WQKV"):
+      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 = 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)
+
+    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+    bsz, seqlen, _, _ = xq.shape
+
+    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
+    attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
+
+    attn = attn.reshape(bsz, seqlen, -1)
+    return self.wo(attn)
+
+class FeedForward:
+  def __init__(self, dim:int, hidden_dim:int, linear=nn.Linear):
+    self.w1 = linear(dim, hidden_dim, bias=False)
+    self.w2 = linear(hidden_dim, dim, bias=False)
+    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit
+
+  def __call__(self, x:Tensor) -> Tensor:
+    w1 = self.w1(x).silu()
+    w3 = self.w3(x)
+    return self.w2(w1 * w3)
+
+class TransformerBlock:
+  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int|None, norm_eps:float, linear=nn.Linear):
+    self.attention = Attention(dim, n_heads, n_kv_heads, linear)
+    self.feed_forward = FeedForward(dim, hidden_dim, linear)
+    self.attention_norm = nn.RMSNorm(dim, norm_eps)
+    self.ffn_norm = nn.RMSNorm(dim, norm_eps)
+
+  def __call__(self, x:Tensor, freqs_cis:Tensor):
+    h = x + self.attention(self.attention_norm(x), freqs_cis)
+    return h + self.feed_forward(self.ffn_norm(h))
+
+class Transformer:
+  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
+               rope_theta:int=10000, max_context:int=1024, linear=nn.Linear, embedding=nn.Embedding):
+    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, linear) for _ in range(n_layers)]
+    self.norm = nn.RMSNorm(dim, norm_eps)
+    self.tok_embeddings = embedding(vocab_size, dim)
+    self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
+
+  def __call__(self, tokens:Tensor):
+    h = self.tok_embeddings(tokens)
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
+    for layer in self.layers: h = layer(h, freqs_cis)
+    logits = self.output(self.norm(h))
+    return logits

examples/mlperf/optim.py

@@ -0,0 +1,57 @@
+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) 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 fstep(self, grads:list[Tensor]):
+    if self.fused:
+      out, extra = self._step([], grads)
+      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
+    else:
+      updates, extra = self._step([], grads)
+    for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i]))
+    to_realize = extra+self.params+self.buffers
+
+    Tensor.realize(*to_realize)
+    return extra[-1]
+
+  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].assign(grads[i].to(self.m[i].device))
+
+    if self.fused:
+      grads[0].assign(grads[0] / self.grad_acc)
+      total_norm = grads[0].float().square().sum().sqrt()
+      grads[0].assign((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].assign(grads[i] / self.grad_acc)
+      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+      for i in range(len(grads)):
+        grads[i].assign((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, g in enumerate(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)).cast(self.m[i].dtype)
+      v_hat = (self.v[i] / (1.0 - self.b2_t)).cast(self.v[i].dtype)
+      up = m_hat / (v_hat.sqrt() + self.eps)
+      ret.append((self.lr * up).cast(g.dtype))
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
+
+  def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
+    wd = self.wd if t.ndim >= 2 else 0.0
+    up = up.shard_like(t) + self.lr.to(t.device) * 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

@@ -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=5000000
 
 export BEAM=0 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_fp8.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=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_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/llama405b/implementations/tinybox_8xMI350X/dev_beam.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 DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8}
+export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4/"
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
+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/llama405b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -0,0 +1,32 @@
+#!/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:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8}
+export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1152}
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4/"
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
+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/dev_beam.sh

@@ -0,0 +1,42 @@
+#!/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:-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} MP=${MP:-1} 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="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 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_beam_mp.sh

@@ -0,0 +1,43 @@
+#!/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:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} 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="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 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} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
+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="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 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/dev_run_mp.sh

@@ -0,0 +1,38 @@
+#!/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:-0}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-32}
+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="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 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
+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/openpilot/compile3.py

@@ -31,7 +31,7 @@ def compile(onnx_file):
   for i in range(3):
     GlobalCounters.reset()
     print(f"run {i}")
-    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
+    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1), OPENPILOT_HACKS=1):
       ret = run_onnx_jit(**inputs).numpy()
     # copy i == 1 so use of JITBEAM is okay
     if i == 1: test_val = np.copy(ret)

examples/openpilot/load_pickle.py

@@ -0,0 +1,16 @@
+import sys, pickle
+from extra.bench_log import WallTimeEvent, BenchEvent
+from tinygrad.helpers import getenv
+
+PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
+
+load_times = []
+
+for _ in range(10):
+  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
+  load_times.append(wte.time)
+  print(f"pickle load: {wte.time:6.2f} s")
+
+if (assert_time:=getenv("ASSERT_MIN_LOAD_TIME")):
+  min_time = min(load_times)
+  assert min_time < assert_time, f"Speed regression, expected min load time of < {assert_time} s but took: {min_time} s"

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

@@ -65,7 +65,7 @@ def get_bar0_size(pcibus):
 class AMSMI(AMDev):
   def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
     self.pcibus = pcibus
-    self.vram, self.doorbell64, self.mmio, self.dma_regions = vram_bar, doorbell_bar, mmio_bar, None
+    self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
     self.pci_state = self.read_pci_state()
     if self.pci_state == "D0": self._init_from_d0()
 
@@ -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,7 +153,8 @@ class SMICtx:
     tables = {}
     for dev in self.devs:
       match dev.ip_ver[am.MP1_HWIP]:
-        case (13,0,6)|(13,0,12): table_t = dev.smu.smu_mod.MetricsTableX_t
+        case (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
@@ -230,12 +231,11 @@ class SMICtx:
 
   def get_power(self, dev, metrics):
     match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
+      case (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:
@@ -244,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))
@@ -279,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

@@ -7,8 +7,8 @@ class GFXFake:
   def __init__(self): self.xccs = 8
 
 class AMDFake(AMDev):
-  def __init__(self, pci_dev, dma_regions=None):
-    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
+  def __init__(self, pci_dev):
+    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
     self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
     self._run_discovery()
     self._build_regs()
@@ -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/amd/decode.py

@@ -1,83 +0,0 @@
-# Instruction format detection and decoding
-from __future__ import annotations
-from extra.assembly.amd.dsl import Inst, FixedBitField
-from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP2, VOP3, VOP3_SDST, VOP3SD, VOP3P, VOPC, VOPD, VINTERP,
-  SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT)
-from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP2 as R4_VOP2,
-  VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
-  VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOPC as R4_SOPC, SOPK as R4_SOPK, SOPP as R4_SOPP,
-  SMEM as R4_SMEM, DS as R4_DS)
-from extra.assembly.amd.autogen.cdna.ins import (VOP1 as C_VOP1, VOP2 as C_VOP2, VOPC as C_VOPC, VOP3A, VOP3B, VOP3P as C_VOP3P,
-  SOP1 as C_SOP1, SOP2 as C_SOP2, SOPC as C_SOPC, SOPK as C_SOPK, SOPP as C_SOPP, SMEM as C_SMEM, DS as C_DS, FLAT as C_FLAT)
-
-def _matches_encoding(word: int, cls: type[Inst]) -> bool:
-  """Check if word matches the encoding pattern of an instruction class."""
-  enc = next(((n, f) for n, f in cls._fields if isinstance(f, FixedBitField) and n == 'encoding'), None)
-  if enc is None: return False
-  bf = enc[1]
-  return ((word >> bf.lo) & bf.mask) == bf.default
-
-# Order matters: more specific encodings first, VOP2 last (it's a catch-all for bit31=0)
-_RDNA_FORMATS_64 = [VOPD, VOP3P, VINTERP, VOP3, DS, FLAT, SMEM]
-_RDNA_FORMATS_32 = [SOP1, SOPC, SOPP, SOPK, VOPC, VOP1, SOP2, VOP2]  # SOP2/VOP2 are catch-alls
-_CDNA_FORMATS_64 = [C_VOP3P, VOP3A, C_DS, C_FLAT, C_SMEM]
-_CDNA_FORMATS_32 = [C_SOP1, C_SOPC, C_SOPP, C_SOPK, C_VOPC, C_VOP1, C_SOP2, C_VOP2]
-_CDNA_VOP3B_OPS = {281, 282, 283, 284, 285, 286, 480, 481, 488, 489}  # VOP3B opcodes
-_RDNA4_FORMATS_64 = [R4_VOPD, R4_VOP3P, R4_VINTERP, R4_VOP3, R4_DS, R4_SMEM]
-_RDNA4_FORMATS_32 = [R4_SOP1, R4_SOPC, R4_SOPP, R4_SOPK, R4_VOPC, R4_VOP1, R4_SOP2, R4_VOP2]
-_RDNA4_VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
-_RDNA3_VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
-# Instructions with SGPR destination (READLANE, READFIRSTLANE, and VOP3-encoded VOPC)
-_VOP1_SDST_OPS = {2}  # V_READFIRSTLANE_B32_E32
-_VOP3_SDST_OPS = {386, 864}  # V_READFIRSTLANE_B32_E64, V_READLANE_B32 (V_WRITELANE_B32=865 writes to VGPR)
-# VOP3-encoded VOPC instructions (opcodes < 256) also have SGPR destination
-
-def detect_format(data: bytes, arch: str = "rdna3") -> type[Inst]:
-  """Detect instruction format from machine code bytes."""
-  assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
-  word = int.from_bytes(data[:4], 'little')
-  if arch == "cdna":
-    # SDWA (0xf9) and DPP (0xfa) use special src0 encodings - not supported
-    if (word & 0x1ff) in (0xf9, 0xfa): raise ValueError(f"CDNA SDWA/DPP not supported word={word:#010x}")
-    if (word >> 30) == 0b11:
-      for cls in _CDNA_FORMATS_64:
-        if _matches_encoding(word, cls):
-          return VOP3B if cls is VOP3A and ((word >> 16) & 0x3ff) in _CDNA_VOP3B_OPS else cls
-      raise ValueError(f"unknown CDNA 64-bit format word={word:#010x}")
-    for cls in _CDNA_FORMATS_32:
-      if _matches_encoding(word, cls): return cls
-    raise ValueError(f"unknown CDNA 32-bit format word={word:#010x}")
-  if arch == "rdna4":
-    if (word >> 30) == 0b11:
-      for cls in _RDNA4_FORMATS_64:
-        if _matches_encoding(word, cls):
-          if cls is R4_VOP3:
-            opcode = (word >> 16) & 0x3ff
-            if opcode in _RDNA4_VOP3SD_OPS: return R4_VOP3SD
-            if opcode in _VOP3_SDST_OPS or opcode < 256: return R4_VOP3_SDST  # VOP3-encoded VOPC (op < 256) writes to SGPR
-          return cls
-      raise ValueError(f"unknown RDNA4 64-bit format word={word:#010x}")
-    for cls in _RDNA4_FORMATS_32:
-      if _matches_encoding(word, cls):
-        if cls is R4_VOP1 and ((word >> 9) & 0xff) in _VOP1_SDST_OPS: return R4_VOP1_SDST
-        return cls
-    raise ValueError(f"unknown RDNA4 32-bit format word={word:#010x}")
-  # RDNA3 (default)
-  if (word >> 30) == 0b11:
-    for cls in _RDNA_FORMATS_64:
-      if _matches_encoding(word, cls):
-        if cls is VOP3:
-          opcode = (word >> 16) & 0x3ff
-          if opcode in _RDNA3_VOP3SD_OPS: return VOP3SD
-          if opcode in _VOP3_SDST_OPS or opcode < 256: return VOP3_SDST  # VOP3-encoded VOPC (op < 256) writes to SGPR
-        return cls
-    raise ValueError(f"unknown 64-bit format word={word:#010x}")
-  for cls in _RDNA_FORMATS_32:
-    if _matches_encoding(word, cls):
-      if cls is VOP1 and ((word >> 9) & 0xff) in _VOP1_SDST_OPS: return VOP1_SDST
-      return cls
-  raise ValueError(f"unknown 32-bit format word={word:#010x}")
-
-def decode_inst(data: bytes, arch: str = "rdna3") -> Inst:
-  """Decode machine code bytes into an instruction."""
-  return detect_format(data, arch).from_bytes(data)

extra/assembly/amd/disasm.py

@@ -1,671 +0,0 @@
-# RDNA3/RDNA4/CDNA disassembler
-from __future__ import annotations
-import re, struct
-from extra.assembly.amd.dsl import Inst, Reg
-
-# Special register mappings for disassembly
-SPECIAL_GPRS = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'null', 125: 'm0', 126: 'exec_lo', 127: 'exec_hi',
-                128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc'}
-SPECIAL_GPRS_CDNA = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'null', 125: 'm0', 126: 'exec_lo', 127: 'exec_hi',
-                     128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc',
-                     102: 'flat_scratch_lo', 103: 'flat_scratch_hi', 104: 'xnack_mask_lo', 105: 'xnack_mask_hi'}
-SPECIAL_PAIRS = {106: 'vcc', 126: 'exec'}
-SPECIAL_PAIRS_CDNA = {106: 'vcc', 126: 'exec', 102: 'flat_scratch', 104: 'xnack_mask'}
-
-def decode_src(v, cdna: bool = False) -> str:
-  """Decode a source operand encoding to its string representation."""
-  v = _unwrap(v)
-  gprs = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
-  if v in gprs: return gprs[v]
-  if v < 106: return f's{v}'
-  if 108 <= v < 124: return f'ttmp{v - 108}'
-  if 129 <= v <= 192: return str(v - 128)  # positive integers 1-64
-  if 193 <= v <= 208: return str(-(v - 192))  # negative integers -1 to -16
-  if v >= 256: return f'v{v - 256}'
-  return f's{v}'
-
-def _unwrap(v) -> int:
-  """Unwrap Reg to int offset, or return int as-is."""
-  return v.offset if isinstance(v, Reg) else v
-
-def _vi(v) -> int:
-  """Get VGPR index from Reg or int (for v[N] fields that encode as 256+N)."""
-  off = _unwrap(v)
-  return off - 256 if off >= 256 else off
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# LITERAL FORMATTING
-# ═══════════════════════════════════════════════════════════════════════════════
-
-_FLOAT_DEC = {240: 0.5, 241: -0.5, 242: 1.0, 243: -1.0, 244: 2.0, 245: -2.0, 246: 4.0, 247: -4.0}
-
-def _lit(inst, v, neg=0) -> str:
-  """Format literal/inline constant value."""
-  v = _unwrap(v)
-  if v == 255:
-    lit = inst._literal
-    if lit is None: return "0"
-    s = f"0x{lit:x}"
-  elif v in _FLOAT_DEC: s = str(_FLOAT_DEC[v])
-  elif 128 <= v <= 192: s = str(v - 128)
-  elif 193 <= v <= 208: s = str(-(v - 192))
-  elif v < 128: s = decode_src(v)
-  elif v >= 256: s = f"v{v - 256}"
-  else: s = decode_src(v)
-  return f"-{s}" if neg else s
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# INSTRUCTION METADATA - fallback functions when inst.num_srcs()/inst.operands unavailable
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def _num_srcs(inst) -> int:
-  """Fallback: get number of source operands from instruction name."""
-  name = getattr(inst, 'op_name', '') or ''
-  n = name.upper()
-  # FMAC/MAC ops are 2-source (dst is implicit accumulator), but FMA/MAD ops are 3-source
-  if 'FMAC' in n or 'V_MAC_' in n: return 2
-  if any(x in n for x in ('FMA', 'MAD', 'CNDMASK', 'BFE', 'BFI', 'LERP', 'MED3', 'SAD', 'DIV_FMAS', 'DIV_FIXUP', 'DIV_SCALE', 'CUBE')): return 3
-  # PERMLANE_VAR ops are 2-source, but PERMLANE (non-VAR) are 3-source
-  if 'PERMLANE' in n and '_VAR' not in n: return 3
-  if any(x in n for x in ('_ADD3', '_LSHL_ADD', '_ADD_LSHL', '_LSHL_OR', '_AND_OR', 'OR3_B32', 'AND_OR_B32', 'ALIGNBIT', 'ALIGNBYTE', 'V_PERM_', 'XOR3', 'XAD', 'MULLIT', 'MINMAX', 'MAXMIN', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM', 'MINIMUM3', 'MAXIMUM3', 'MIN3', 'MAX3', 'DOT2', 'CVT_PK_U8_F32', 'DOT4', 'DOT8', 'WMMA', 'SWMMAC')): return 3
-  return 2
-
-# SWMMAC register counts: (dst, src0, src1, src2)
-def _swmmac_regs(name: str) -> tuple[int, int, int, int]:
-  """Return (dst, src0, src1, src2) register counts for SWMMAC instructions."""
-  if 'f16_16x16x32' in name or 'bf16_16x16x32' in name: return (4, 4, 8, 1)
-  if 'f32_16x16x32_f16' in name or 'f32_16x16x32_bf16' in name: return (8, 4, 8, 1)
-  if 'i32_16x16x32_iu4' in name: return (8, 1, 2, 1)
-  if 'i32_16x16x64_iu4' in name: return (8, 2, 4, 1)
-  if 'i32_16x16x32_iu8' in name or 'f32_16x16x32_fp8' in name or 'f32_16x16x32_bf8' in name: return (8, 2, 4, 1)
-  return (8, 8, 8, 8)
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# IMPORTS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP2, VOP3, VOP3_SDST, VOP3SD, VOP3P, VOPC, VOPD, VINTERP, SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH,
-  VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOPDOp, SOP1Op, SOPKOp, SOPPOp, SMEMOp, DSOp)
-from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP2 as R4_VOP2, VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
-  VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOPC as R4_SOPC, SOPK as R4_SOPK, SOPP as R4_SOPP,
-  SMEM as R4_SMEM, DS as R4_DS, VOPDOp as R4_VOPDOp)
-from extra.assembly.amd.autogen.cdna.ins import FLAT as C_FLAT
-
-def _is_cdna(inst: Inst) -> bool: return 'cdna' in inst.__class__.__module__
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# CONSTANTS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-HWREG = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
-         6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
-         19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI', 22: 'HW_REG_XNACK_MASK',
-         23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
-HWREG_RDNA4 = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 4: 'HW_REG_STATE_PRIV', 5: 'HW_REG_GPR_ALLOC',
-               6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 10: 'HW_REG_PERF_SNAPSHOT_DATA', 11: 'HW_REG_PERF_SNAPSHOT_PC_LO',
-               12: 'HW_REG_PERF_SNAPSHOT_PC_HI', 15: 'HW_REG_PERF_SNAPSHOT_DATA1', 16: 'HW_REG_PERF_SNAPSHOT_DATA2',
-               17: 'HW_REG_EXCP_FLAG_PRIV', 18: 'HW_REG_EXCP_FLAG_USER', 19: 'HW_REG_TRAP_CTRL',
-               20: 'HW_REG_SCRATCH_BASE_LO', 21: 'HW_REG_SCRATCH_BASE_HI', 23: 'HW_REG_HW_ID1',
-               24: 'HW_REG_HW_ID2', 26: 'HW_REG_SCHED_MODE', 29: 'HW_REG_SHADER_CYCLES_LO',
-               30: 'HW_REG_SHADER_CYCLES_HI', 31: 'HW_REG_DVGPR_ALLOC_LO', 32: 'HW_REG_DVGPR_ALLOC_HI'}
-MSG = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
-       131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA',
-       134: 'MSG_RTN_GET_TBA_TO_PC', 135: 'MSG_RTN_GET_SE_AID_ID'}
-# CDNA opcode name aliases for disasm (new name -> old name expected by tests)
-_CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32', 'v_fmamk_f32': 'v_madmk_f32', 'v_fmaak_f32': 'v_madak_f32'}
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# HELPERS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{_unwrap(b)}" if n == 1 else f"{p}[{_unwrap(b)}:{_unwrap(b)+n-1}]"
-def _sreg(b: int, n: int = 1) -> str: return _reg("s", _unwrap(b), n)
-def _vreg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("v", b - 256 if b >= 256 else b, n)
-def _areg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("a", b - 256 if b >= 256 else b, n)  # accumulator registers for GFX90a
-def _ttmp(b, n: int = 1) -> str: b = _unwrap(b); return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
-def _sreg_or_ttmp(b, n: int = 1) -> str: return _ttmp(b, n) or _sreg(b, n)
-
-def _fmt_sdst(v, n: int = 1, cdna: bool = False) -> str:
-  v = _unwrap(v)
-  if t := _ttmp(v, n): return t
-  pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
-  gprs = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
-  if n > 1: return pairs.get(v) or gprs.get(v) or _sreg(v, n)  # also check gprs for null/m0
-  return gprs.get(v, f"s{v}")
-
-def _fmt_src(v, n: int = 1, cdna: bool = False) -> str:
-  v = _unwrap(v)
-  if v == 253: return "src_scc"  # SCC as source operand
-  if n == 1: return decode_src(v, cdna)
-  if v >= 256: return _vreg(v, n)
-  if v <= 101: return _sreg(v, n)  # s0-s101 can be pairs, but 102+ are special on CDNA
-  pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
-  if n == 2 and v in pairs: return pairs[v]
-  if v <= 105: return _sreg(v, n)  # s102-s105 regular pairs for RDNA
-  if t := _ttmp(v, n): return t
-  return decode_src(v, cdna)
-
-def _fmt_v16(v, base: int = 256, hi_thresh: int = 384) -> str:
-  v = _unwrap(v)
-  return f"v{(v - base) & 0x7f}.{'h' if v >= hi_thresh else 'l'}"
-
-def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
-def _omod(v: int) -> str: return {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(v, "")
-def _src16(inst, v: int) -> str: v = _unwrap(v); return _fmt_v16(v) if v >= 256 else _lit(inst, v)  # format 16-bit src: vgpr.h/l or literal
-def _mods(*pairs) -> str: return " ".join(m for c, m in pairs if c)
-def _fmt_bits(label: str, val: int, count: int) -> str: return f"{label}:[{','.join(str((val >> i) & 1) for i in range(count))}]"
-
-def _vop3_src(inst, v: int, neg: int, abs_: int, hi: int, n: int, f16: bool) -> str:
-  """Format VOP3 source operand with modifiers."""
-  v = _unwrap(v)
-  if v == 255: s = _lit(inst, v)  # literal constant takes priority
-  elif n > 1: s = _fmt_src(v, n)
-  elif f16 and v >= 256: s = f"v{v - 256}.h" if hi else f"v{v - 256}.l"
-  elif v == 253: s = "src_scc"  # VOP3 sources use src_scc not scc
-  else: s = _lit(inst, v)
-  if abs_: s = f"|{s}|"
-  return f"-{s}" if neg else s
-
-def _opsel_str(opsel: int, n: int, need: bool, is16_d: bool) -> str:
-  """Format op_sel modifier string."""
-  if not need: return ""
-  dst_hi = (opsel >> 3) & 1
-  if n == 1: return f" op_sel:[{opsel & 1},{dst_hi}]"
-  if n == 2: return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{dst_hi}]"
-  return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{dst_hi}]"
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# DISASSEMBLER
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def _disasm_vop1(inst: VOP1) -> str:
-  name, cdna = inst.op_name.lower() or f'vop1_op_{inst.op}', _is_cdna(inst)
-  name = name.replace('_e32', '')  # Strip _e32 suffix
-  if any(x in name for x in ('v_nop', 'v_pipeflush', 'v_clrexcp')): return name  # no operands
-  if 'readfirstlane' in name:
-    src = inst.src0.fmt() if inst.src0.offset >= 256 else decode_src(inst.src0.offset, cdna)
-    vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
-    return f"{name} {_fmt_sdst(vdst_off, 1, cdna)}, {src}"
-  bits = inst.canonical_op_bits
-  is16_dst, is16_src = not cdna and bits['d'] == 16, not cdna and bits['s0'] == 16
-  # v_cvt_pk_f32_fp8/bf8: pcode has None dst type but outputs 2 VGPRs
-  if 'cvt_pk_f32_fp8' in name or 'cvt_pk_f32_bf8' in name: is16_src = True
-  # Format dst
-  if is16_dst: dst = _fmt_v16(inst.vdst)
-  else: dst = inst.vdst.fmt()
-  # Format src
-  if inst.src0.offset == 255: src = _lit(inst, inst.src0)
-  elif is16_src and inst.src0.offset >= 256: src = _fmt_v16(inst.src0)
-  elif inst.src0.sz > 1: src = _fmt_src(inst.src0, inst.src0.sz, cdna)
-  else: src = _lit(inst, inst.src0)
-  return f"{name} {dst}, {src}"
-
-_VOP2_CARRY_OUT = {'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'}  # carry out only
-_VOP2_CARRY_INOUT = {'v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'}  # carry in and out (CDNA)
-_VOP2_CARRY_INOUT_RDNA = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'}  # carry in and out (RDNA)
-def _disasm_vop2(inst: VOP2) -> str:
-  name, cdna = inst.op_name.lower(), _is_cdna(inst)
-  if cdna: name = _CDNA_DISASM_ALIASES.get(name, name)  # apply CDNA aliases
-  suf = "" if cdna or name.endswith('_e32') or (not cdna and inst.op == VOP2Op.V_DOT2ACC_F32_F16_E32) else "_e32"
-  lit = getattr(inst, '_literal', None)
-  is16 = not cdna and inst.canonical_op_bits['d'] == 16
-  # fmaak/madak: dst = src0 * vsrc1 + K, fmamk/madmk: dst = src0 * K + vsrc1
-  if 'fmaak' in name or 'madak' in name or (not cdna and inst.op in (VOP2Op.V_FMAAK_F32_E32, VOP2Op.V_FMAAK_F16_E32)):
-    if lit is None: return f"op_{inst.op.value if hasattr(inst.op, 'value') else inst.op}"
-    if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1)}, 0x{lit:x}"
-    return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, 0x{lit:x}"
-  if 'fmamk' in name or 'madmk' in name or (not cdna and inst.op in (VOP2Op.V_FMAMK_F32_E32, VOP2Op.V_FMAMK_F16_E32)):
-    if lit is None: return f"op_{inst.op.value if hasattr(inst.op, 'value') else inst.op}"
-    if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, 0x{lit:x}, {_fmt_v16(inst.vsrc1)}"
-    return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, 0x{lit:x}, {inst.vsrc1.fmt()}"
-  if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1)}"
-  vcc = "vcc" if cdna else "vcc_lo"
-  if cdna and name in _VOP2_CARRY_OUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}"
-  if cdna and name in _VOP2_CARRY_INOUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
-  if not cdna and name in _VOP2_CARRY_INOUT_RDNA: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
-  sn0 = inst.canonical_op_regs.get('s0', 1)
-  if inst.vdst.sz > 1 or sn0 > 1 or inst.vsrc1.sz > 1:
-    src0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else _fmt_src(inst.src0, sn0, cdna)
-    return f"{name.replace('_e32', '')} {inst.vdst.fmt()}, {src0}, {inst.vsrc1.fmt()}"
-  return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}" + (f", {vcc}" if name == 'v_cndmask_b32' else "")
-
-def _disasm_vopc(inst: VOPC) -> str:
-  name, cdna = inst.op_name.lower(), _is_cdna(inst)
-  bits = inst.canonical_op_bits
-  is16 = bits['s0'] == 16
-  if cdna:
-    s0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else _fmt_src(inst.src0, inst.src0.sz, cdna)
-    return f"{name} vcc, {s0}, {inst.vsrc1.fmt()}"  # CDNA VOPC always outputs vcc
-  # RDNA: v_cmpx_* writes to exec (no vcc), v_cmp_* writes to vcc_lo
-  has_vcc = 'cmpx' not in name
-  s0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else inst.src0.fmt() if inst.src0.sz > 1 else _src16(inst, inst.src0.offset) if is16 else _lit(inst, inst.src0)
-  s1 = inst.vsrc1.fmt() if inst.vsrc1.sz > 1 else _fmt_v16(inst.vsrc1) if is16 else inst.vsrc1.fmt()
-  suf = "" if name.endswith('_e32') else "_e32"
-  return f"{name}{suf} vcc_lo, {s0}, {s1}" if has_vcc else f"{name}{suf} {s0}, {s1}"
-
-NO_ARG_SOPP = {SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
-               SOPPOp.S_WAIT_IDLE, SOPPOp.S_ENDPGM_SAVED, SOPPOp.S_CODE_END, SOPPOp.S_ENDPGM_ORDERED_PS_DONE, SOPPOp.S_TTRACEDATA}
-_CDNA_NO_ARG_SOPP = {'s_endpgm', 's_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_nop', 's_sethalt', 's_sleep',
-                     's_setprio', 's_trap', 's_incperflevel', 's_decperflevel', 's_sendmsg', 's_sendmsghalt'}
-
-def _disasm_sopp(inst: SOPP) -> str:
-  name, cdna = inst.op_name.lower(), _is_cdna(inst)
-  is_rdna4 = 'rdna4' in inst.__class__.__module__
-  # Ops that have no argument when simm16 == 0
-  no_arg_zero = {'s_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_wait_idle', 's_endpgm_saved',
-                 's_endpgm_ordered_ps_done', 's_code_end'}
-  if name in no_arg_zero: return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
-  if name == 's_endpgm': return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
-  if cdna:
-    if name == 's_waitcnt':
-      vm, lgkm, exp = inst.simm16 & 0xf, (inst.simm16 >> 8) & 0x3f, (inst.simm16 >> 4) & 0x7
-      p = [f"vmcnt({vm})" if vm != 0xf else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
-      return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
-    if name.startswith(('s_cbranch', 's_branch')): return f"{name} {inst.simm16}"
-    return f"{name} 0x{inst.simm16:x}" if inst.simm16 else name
-  # RDNA (use name-based checks instead of enum-based for cross-arch compatibility)
-  if name == 's_waitcnt':
-    if is_rdna4:
-      return f"{name} {inst.simm16}" if inst.simm16 else f"{name} 0"
-    vm, exp, lgkm = (inst.simm16 >> 10) & 0x3f, inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x3f
-    p = [f"vmcnt({vm})" if vm != 0x3f else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
-    return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
-  if name == 's_delay_alu':
-    deps = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
-    skips = ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
-    id0, skip, id1 = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 7) & 0xf
-    dep = lambda v: deps[v-1] if 0 < v <= len(deps) else str(v)
-    p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
-    return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
-  if name.startswith(('s_cbranch', 's_branch')): return f"{name} 0x{inst.simm16:x}"
-  return f"{name} 0x{inst.simm16:x}"
-
-def _disasm_smem(inst: SMEM) -> str:
-  name, cdna = inst.op_name.lower(), _is_cdna(inst)
-  if name in ('s_gl1_inv', 's_dcache_inv'): return name
-  soe, imm = getattr(inst, 'soe', 0), getattr(inst, 'imm', 1)
-  is_rdna4 = 'rdna4' in inst.__class__.__module__
-  offset = inst.ioffset if is_rdna4 else getattr(inst, 'offset', 0)
-  if cdna:
-    if soe and imm: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
-    elif imm: off_s = f"0x{offset:x}"
-    elif offset < 256: off_s = decode_src(offset, cdna)
-    else: off_s = decode_src(inst.soffset, cdna)
-  elif offset and inst.soffset != 124: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
-  elif offset: off_s = f"0x{offset:x}"
-  else: off_s = decode_src(inst.soffset, cdna)
-  is_buffer = 'buffer' in name or 's_atc_probe_buffer' == name
-  sbase_idx, sbase_count = _unwrap(inst.sbase), 4 if is_buffer else 2
-  sbase_str = _fmt_src(sbase_idx, sbase_count, cdna) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
-  if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {_unwrap(inst.sdata)}, {sbase_str}, {off_s}"
-  if 'prefetch' in name:
-    off = getattr(inst, 'ioffset', getattr(inst, 'offset', 0))
-    if off >= 0x800000: off = off - 0x1000000
-    off_s = f"0x{off:x}" if off > 255 else str(off)
-    soff_s = decode_src(inst.soffset, cdna) if inst.soffset != 124 else "null"
-    if 'pc_rel' in name: return f"{name} {off_s}, {soff_s}, {_unwrap(inst.sdata)}"
-    return f"{name} {sbase_str}, {off_s}, {soff_s}, {_unwrap(inst.sdata)}"
-  # Use get_field_bits for register count
-  dst_n = inst.canonical_op_regs.get('d', 1)
-  th, scope = getattr(inst, 'th', 0), getattr(inst, 'scope', 0)
-  if is_rdna4:  # RDNA4 uses th/scope instead of glc/dlc
-    th_names = ['TH_LOAD_RT', 'TH_LOAD_NT', 'TH_LOAD_HT', 'TH_LOAD_LU']
-    scope_names = ['SCOPE_CU', 'SCOPE_SE', 'SCOPE_DEV', 'SCOPE_SYS']
-    mods = (f" th:{th_names[th]}" if th else "") + (f" scope:{scope_names[scope]}" if scope else "")
-    return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}{mods}"
-  if th or scope:
-    th_names = ['TH_LOAD_RT', 'TH_LOAD_NT', 'TH_LOAD_HT', 'TH_LOAD_LU']
-    scope_names = ['SCOPE_CU', 'SCOPE_SE', 'SCOPE_DEV', 'SCOPE_SYS']
-    mods = (f" th:{th_names[th]}" if th else "") + (f" scope:{scope_names[scope]}" if scope else "")
-    return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}{mods}"
-  return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
-
-def _disasm_flat(inst: FLAT) -> str:
-  name, cdna = inst.op_name.lower(), _is_cdna(inst)
-  acc = getattr(inst, 'acc', 0)
-  reg_fn = _areg if acc else _vreg
-  seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
-  instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
-  off_val = inst.offset if seg == 'flat' else (inst.offset if inst.offset < 4096 else inst.offset - 8192)
-  # Use get_field_bits: data for stores/atomics, d for loads
-  regs = inst.canonical_op_regs
-  w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
-  off_s = f" offset:{off_val}" if off_val else ""
-  if cdna: mods = f"{off_s}{' glc' if inst.sc0 else ''}{' slc' if inst.nt else ''}"
-  else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
-  if seg == 'flat' or _unwrap(inst.saddr) == 0x7F: saddr_s = ""
-  elif _unwrap(inst.saddr) == 124: saddr_s = ", off"
-  elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr, cdna)}"
-  elif _unwrap(inst.saddr) in (SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS): saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
-  elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
-  else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
-  if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
-  if cdna: addr_w = 1 if seg == 'scratch' else 2
-  else: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
-  addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
-  data_s, vdst_s = reg_fn(inst.data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
-  glc_or_sc0 = inst.sc0 if cdna else inst.glc
-  if 'atomic' in name:
-    return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}"
-  if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
-  return f"{instr} {reg_fn(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
-
-def _disasm_ds(inst: DS) -> str:
-  op, name = inst.op, inst.op_name.lower()
-  acc = getattr(inst, 'acc', 0)
-  reg_fn = _areg if acc else _vreg
-  gds = " gds" if getattr(inst, 'gds', 0) else ""
-  off = f" offset:{inst.offset0 | (inst.offset1 << 8)}" if inst.offset0 or inst.offset1 else ""
-  off2 = (" offset0:" + str(inst.offset0) if inst.offset0 else "") + (" offset1:" + str(inst.offset1) if inst.offset1 else "")
-  # Use get_field_bits: data for stores/writes/atomics, d for loads
-  regs = inst.canonical_op_regs
-  w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'write' in name or ('load' not in name and 'read' not in name) else regs.get('d', 1)
-  d0, d1, dst, addr = reg_fn(inst.data0, w), reg_fn(inst.data1, w), reg_fn(inst.vdst, w), _vreg(inst.addr)
-
-  if name == 'ds_nop': return name
-  if name == 'ds_bvh_stack_rtn_b32': return f"{name} {_vreg(inst.vdst)}, {addr}, {_vreg(inst.data0)}, {_vreg(inst.data1, 4)}{off}{gds}"
-  if 'bvh_stack_push' in name:
-    d1_regs = 8 if 'push8' in name else 4
-    vdst_regs = 2 if 'pop2' in name else 1
-    vdst_s = _vreg(inst.vdst, vdst_regs) if vdst_regs > 1 else _vreg(inst.vdst)
-    return f"{name} {vdst_s}, {addr}, {_vreg(inst.data0)}, {_vreg(inst.data1, d1_regs)}{off}{gds}"
-  if 'gws_sema' in name and 'sema_br' not in name: return f"{name}{off}{gds}"
-  if 'gws_' in name: return f"{name} {addr}{off}{gds}"
-  if name in ('ds_consume', 'ds_append'): return f"{name} {reg_fn(inst.vdst)}{off}{gds}"
-  if 'gs_reg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {reg_fn(inst.data0)}{off}{gds}"
-  if '2addr' in name:
-    if 'load' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}{off2}{gds}"
-    if 'store' in name and 'xchg' not in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
-    return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}, {d0}, {d1}{off2}{gds}"
-  if 'write2' in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
-  if 'read2' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}{off2}{gds}"
-  if 'load' in name: return f"{name} {reg_fn(inst.vdst)}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
-  if 'store' in name and not _has(name, 'cmp', 'xchg'):
-    return f"{name} {reg_fn(inst.data0)}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
-  if 'swizzle' in name or name == 'ds_ordered_count': return f"{name} {reg_fn(inst.vdst)}, {addr}{off}{gds}"
-  if 'permute' in name: return f"{name} {reg_fn(inst.vdst)}, {addr}, {reg_fn(inst.data0)}{off}{gds}"
-  if 'condxchg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {addr}, {reg_fn(inst.data0, 2)}{off}{gds}"
-  if _has(name, 'cmpstore', 'mskor', 'wrap'):
-    return f"{name} {dst}, {addr}, {d0}, {d1}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}, {d1}{off}{gds}"
-  return f"{name} {dst}, {addr}, {d0}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}{off}{gds}"
-
-def _disasm_vop3(inst: VOP3) -> str:
-  op, name = inst.op, inst.op_name.lower()
-  n_up = name.upper()
-  bits = inst.canonical_op_bits
-
-  # RDNA4 v_s_* scalar VOP3 instructions - vdst is SGPR (VGPRField adds 256)
-  if name.startswith('v_s_'):
-    src = _lit(inst, inst.src0) if _unwrap(inst.src0) == 255 else ("src_scc" if _unwrap(inst.src0) == 253 else _fmt_src(inst.src0, max(1, bits['s0'] // 32)))
-    if inst.neg & 1: src = f"-{src}"
-    if inst.abs & 1: src = f"|{src}|"
-    clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
-    vdst_raw = _unwrap(inst.vdst)
-    return f"{name} s{vdst_raw - 256 if vdst_raw >= 256 else vdst_raw}, {src}" + (" clamp" if clamp else "") + _omod(inst.omod)
-
-  # Use get_field_bits for register sizes and 16-bit detection
-  r0, r1, r2 = max(1, bits['s0'] // 32), max(1, bits['s1'] // 32), max(1, bits['s2'] // 32)
-  dn = max(1, bits['d'] // 32)
-  is16_d, is16_s, is16_s2 = bits['d'] == 16, bits['s0'] == 16, bits['s2'] == 16
-
-  s0 = _vop3_src(inst, inst.src0, inst.neg&1, inst.abs&1, inst.opsel&1, r0, is16_s)
-  s1 = _vop3_src(inst, inst.src1, inst.neg&2, inst.abs&2, inst.opsel&2, r1, is16_s)
-  s2 = _vop3_src(inst, inst.src2, inst.neg&4, inst.abs&4, inst.opsel&4, r2, is16_s2)
-
-  # Format destination
-  if 'readlane' in name:
-    vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
-    dst = _fmt_sdst(vdst_off, 1)
-  elif is16_d: dst = f"{inst.vdst.fmt()}.h" if (inst.opsel & 8) else f"{inst.vdst.fmt()}.l"
-  else: dst = inst.vdst.fmt()
-
-  clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
-  cl, om = " clamp" if clamp else "", _omod(inst.omod)
-  nonvgpr_opsel = (inst.src0.offset < 256 and (inst.opsel & 1)) or (inst.src1.offset < 256 and (inst.opsel & 2)) or (inst.src2.offset < 256 and (inst.opsel & 4))
-  need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)
-
-  op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
-  e64 = "" if name.endswith('_e64') else "_e64"
-  if op_val < 256:  # VOPC
-    vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
-    return f"{name}{e64} {s0}, {s1}{cl}" if name.startswith('v_cmpx') else f"{name}{e64} {_fmt_sdst(vdst_off, 1)}, {s0}, {s1}{cl}"
-  if op_val < 384:  # VOP2
-    n = inst.num_srcs() or 2
-    os = _opsel_str(inst.opsel, n, need_opsel, is16_d)
-    return f"{name}{e64} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name}{e64} {dst}, {s0}, {s1}{os}{cl}{om}"
-  if op_val < 512:  # VOP1
-    if re.match(r'v_cvt_f32_(bf|fp)8', name) and inst.opsel:
-      os = f" byte_sel:{((inst.opsel & 1) << 1) | ((inst.opsel >> 1) & 1)}"
-    else:
-      os = _opsel_str(inst.opsel, 1, need_opsel, is16_d)
-    if 'v_nop' in name or 'v_pipeflush' in name: return f"{name}{e64}"
-    return f"{name}{e64} {dst}, {s0}{os}{cl}{om}"
-  # Native VOP3
-  n = inst.num_srcs() or 2
-  os = f" byte_sel:{inst.opsel >> 2}" if 'cvt_sr' in name and inst.opsel else _opsel_str(inst.opsel, n, need_opsel, is16_d)
-  return f"{name} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{os}{cl}{om}"
-
-def _disasm_vop3sd(inst: VOP3SD) -> str:
-  name = inst.op_name.lower()
-  def src(reg, neg):
-    s = _lit(inst, reg.offset) if reg.offset == 255 else ("src_scc" if reg.offset == 253 else (reg.fmt() if reg.sz > 1 else _lit(inst, reg.offset)))
-    return f"neg({s})" if neg and reg.offset == 255 else (f"-{s}" if neg else s)
-  s0, s1, s2 = src(inst.src0, inst.neg & 1), src(inst.src1, inst.neg & 2), src(inst.src2, inst.neg & 4)
-  # VOP3SD: _co_ ops (add/sub) without _ci_ have only 2 sources, all others (mad, div_scale, _co_ci_) have 3 sources
-  has_only_two_srcs = '_co_' in name and '_ci_' not in name and 'mad' not in name
-  srcs = f"{s0}, {s1}" if has_only_two_srcs else f"{s0}, {s1}, {s2}"
-  clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
-  return f"{name} {inst.vdst.fmt()}, {_fmt_sdst(inst.sdst, 1)}, {srcs}{' clamp' if clamp else ''}{_omod(inst.omod)}"
-
-def _disasm_vopd(inst: VOPD) -> str:
-  lit = inst._literal or getattr(inst, 'literal', None)
-  is_rdna4 = 'rdna4' in inst.__class__.__module__
-  op_enum = R4_VOPDOp if is_rdna4 else VOPDOp
-  vdst_y, nx, ny = (_unwrap(inst.vdsty) << 1) | ((_unwrap(inst.vdstx) & 1) ^ 1), op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
-  def half(n, vd, s0, vs1):
-    vd, vs1 = _vi(vd), _vi(vs1)
-    if 'mov' in n: return f"{n} v{vd}, {_lit(inst, s0)}"
-    if 'fmamk' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, 0x{lit:x}, v{vs1}"
-    if 'fmaak' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}, 0x{lit:x}"
-    return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}"
-  return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
-
-def _disasm_vop3p(inst: VOP3P) -> str:
-  name = inst.op_name.lower()
-  is_wmma, is_swmmac, n, is_fma_mix = 'wmma' in name, 'swmmac' in name, inst.num_srcs() or 2, 'fma_mix' in name
-  def get_src(reg):
-    return _lit(inst, reg.offset) if reg.offset == 255 else reg.fmt()
-  src0, src1, src2, dst = get_src(inst.src0), get_src(inst.src1), get_src(inst.src2), inst.vdst.fmt()
-  opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
-  clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
-  if is_fma_mix:
-    def m(s, neg, abs_): return f"-{f'|{s}|' if abs_ else s}" if neg else (f"|{s}|" if abs_ else s)
-    src0, src1, src2 = m(src0, inst.neg & 1, inst.neg_hi & 1), m(src1, inst.neg & 2, inst.neg_hi & 2), m(src2, inst.neg & 4, inst.neg_hi & 4)
-    mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if clamp else [])
-  elif is_swmmac:
-    mods = ([f"index_key:{inst.opsel}"] if inst.opsel else []) + ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + \
-           ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
-  else:
-    opsel_hi_default = 7 if n == 3 else 3
-    mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else []) + \
-           ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
-  return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
-
-def _disasm_sop1(inst: SOP1) -> str:
-  op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
-  # Use get_field_bits for register sizes
-  regs = inst.canonical_op_regs
-  dst_regs, src_regs = regs.get('d', 1), regs.get('s0', 1)
-  src = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, src_regs, cdna)
-  if not cdna:
-    if 'getpc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2)}"
-    if 'setpc_b64' in name or 'rfe_b64' in name: return f"{name} {src}"
-    if 'swappc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2)}, {src}"
-    if 'sendmsg_rtn' in name:
-      v = _unwrap(inst.ssrc0)
-      msg_str = MSG.get(v)
-      return f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, sendmsg({msg_str})" if msg_str else f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, 0x{v:x}"
-  sop1_src_only = ('S_ALLOC_VGPR', 'S_SLEEP_VAR', 'S_BARRIER_SIGNAL', 'S_BARRIER_SIGNAL_ISFIRST', 'S_BARRIER_INIT', 'S_BARRIER_JOIN')
-  if inst.op_name in sop1_src_only: return f"{name} {src}"
-  return f"{name} {_fmt_sdst(inst.sdst, dst_regs, cdna)}, {src}"
-
-def _disasm_sop2(inst: SOP2) -> str:
-  cdna, name = _is_cdna(inst), inst.op_name.lower()
-  lit = getattr(inst, '_literal', None)
-  # Use get_field_bits for register sizes
-  regs = inst.canonical_op_regs
-  dn, s0n, s1n = regs['d'], regs['s0'], regs['s1']
-  s0 = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, s0n, cdna)
-  s1 = _lit(inst, inst.ssrc1) if _unwrap(inst.ssrc1) == 255 else _fmt_src(inst.ssrc1, s1n, cdna)
-  dst = _fmt_sdst(inst.sdst, dn, cdna)
-  if 'fmamk' in name and lit is not None: return f"{name} {dst}, {s0}, 0x{lit:x}, {s1}"
-  if 'fmaak' in name and lit is not None: return f"{name} {dst}, {s0}, {s1}, 0x{lit:x}"
-  return f"{name} {dst}, {s0}, {s1}"
-
-def _disasm_sopc(inst: SOPC) -> str:
-  cdna, regs = _is_cdna(inst), inst.canonical_op_regs
-  s0 = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, regs['s0'], cdna)
-  s1 = _lit(inst, inst.ssrc1) if _unwrap(inst.ssrc1) == 255 else _fmt_src(inst.ssrc1, regs['s1'], cdna)
-  return f"{inst.op_name.lower()} {s0}, {s1}"
-
-def _disasm_sopk(inst: SOPK) -> str:
-  op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
-  is_rdna4 = 'rdna4' in inst.__class__.__module__
-  hw = HWREG_RDNA4 if is_rdna4 else HWREG
-  def fmt_hwreg(hid, hoff, hsz):
-    if hid not in hw: return f"0x{inst.simm16:x}"
-    hr_name = hw[hid]
-    return f"hwreg({hr_name})" if hoff == 0 and hsz == 32 else f"hwreg({hr_name}, {hoff}, {hsz})"
-  if name == 's_setreg_imm32_b32':
-    hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
-    return f"{name} {fmt_hwreg(hid, hoff, hsz)}, 0x{inst._literal:x}"
-  if name == 's_version': return f"{name} 0x{inst.simm16:x}"
-  if name in ('s_setreg_b32', 's_getreg_b32'):
-    hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
-    hs = fmt_hwreg(hid, hoff, hsz)
-    return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1, cdna)}" if 'setreg' in name else f"{name} {_fmt_sdst(inst.sdst, 1, cdna)}, {hs}"
-  if name in ('s_subvector_loop_begin', 's_subvector_loop_end'):
-    return f"{name} {_fmt_sdst(inst.sdst, 1)}, 0x{inst.simm16:x}"
-  return f"{name} {_fmt_sdst(inst.sdst, inst.canonical_op_regs['d'], cdna)}, 0x{inst.simm16:x}"
-
-def _disasm_vinterp(inst: VINTERP) -> str:
-  mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
-  return f"{inst.op_name.lower()} {inst.vdst.fmt()}, {_lit(inst, inst.src0, inst.neg & 1)}, {_lit(inst, inst.src1, inst.neg & 2)}, {_lit(inst, inst.src2, inst.neg & 4)}" + (" " + mods if mods else "")
-
-DISASM_HANDLERS: dict[type, callable] = {
-  VOP1: _disasm_vop1, VOP1_SDST: _disasm_vop1, VOP2: _disasm_vop2, VOPC: _disasm_vopc, VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOPD: _disasm_vopd, VOP3P: _disasm_vop3p,
-  VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, GLOBAL: _disasm_flat, SCRATCH: _disasm_flat,
-  SOP1: _disasm_sop1, SOP2: _disasm_sop2, SOPC: _disasm_sopc, SOPK: _disasm_sopk,
-  # RDNA4
-  R4_VOP1: _disasm_vop1, R4_VOP1_SDST: _disasm_vop1, R4_VOP2: _disasm_vop2, R4_VOPC: _disasm_vopc, R4_VOP3: _disasm_vop3, R4_VOP3_SDST: _disasm_vop3, R4_VOP3SD: _disasm_vop3sd,
-  R4_VOPD: _disasm_vopd, R4_VOP3P: _disasm_vop3p, R4_VINTERP: _disasm_vinterp, R4_SOPP: _disasm_sopp, R4_SMEM: _disasm_smem,
-  R4_DS: _disasm_ds, R4_SOP1: _disasm_sop1, R4_SOP2: _disasm_sop2, R4_SOPC: _disasm_sopc, R4_SOPK: _disasm_sopk}
-
-def disasm(inst: Inst) -> str: return DISASM_HANDLERS[type(inst)](inst)
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# CDNA DISASSEMBLER SUPPORT
-# ═══════════════════════════════════════════════════════════════════════════════
-
-try:
-  from extra.assembly.amd.autogen.cdna.ins import (VOP1 as CDNA_VOP1, VOP2 as CDNA_VOP2, VOPC as CDNA_VOPC, VOP3A, VOP3B, VOP3P as CDNA_VOP3P,
-    SOP1 as CDNA_SOP1, SOP2 as CDNA_SOP2, SOPC as CDNA_SOPC, SOPK as CDNA_SOPK, SOPP as CDNA_SOPP, SMEM as CDNA_SMEM, DS as CDNA_DS,
-    FLAT as CDNA_FLAT, VOP1Op as CDNA_VOP1Op, VOP2Op as CDNA_VOP2Op, VOPCOp as CDNA_VOPCOp)
-
-  def _cdna_src(inst, v, neg, abs_=0, n=1):
-    s = _lit(inst, v) if v == 255 else _fmt_src(v, n, cdna=True)
-    if abs_: s = f"|{s}|"
-    return f"neg({s})" if neg and v == 255 else (f"-{s}" if neg else s)
-
-  _CDNA_VOP3_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32'}
-
-  def _disasm_vop3a(inst) -> str:
-    op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
-    name = inst.op_name.lower() or f'vop3a_op_{op_val}'
-    n = inst.num_srcs() or _num_srcs(inst)
-    cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
-    orig_name = name
-    name = _CDNA_VOP3_ALIASES.get(name, name)
-    if name != orig_name:
-      s0, s1 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, 1), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, 1)
-      s2 = ""
-      dst = _vreg(inst.vdst)
-    else:
-      regs = inst.canonical_op_regs
-      dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
-      s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, r0), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, r1), _cdna_src(inst, inst.src2, inst.neg&4, inst.abs&4, r2)
-      dst = _vreg(inst.vdst, dregs) if dregs > 1 else _vreg(inst.vdst)
-    if op_val >= 512:
-      return f"{name} {dst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{cl}{om}"
-    if op_val < 256:
-      sdst = _fmt_sdst(inst.vdst, 2, cdna=True)
-      return f"{name}_e64 {sdst}, {s0}, {s1}{cl}"
-    if 320 <= op_val < 512:
-      if name in ('v_nop', 'v_clrexcp'): return f"{name}_e64"
-      return f"{name}_e64 {dst}, {s0}{cl}{om}"
-    if name == 'v_cndmask_b32':
-      s2 = _fmt_src(inst.src2, 2, cdna=True)
-      return f"{name}_e64 {dst}, {s0}, {s1}, {s2}{cl}{om}"
-    if name in ('v_mul_legacy_f32', 'v_mac_f32'):
-      return f"{name}_e64 {dst}, {s0}, {s1}{cl}{om}"
-    suf = "_e64" if op_val < 512 else ""
-    return f"{name}{suf} {dst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name}{suf} {dst}, {s0}, {s1}{cl}{om}"
-
-  def _disasm_vop3b(inst) -> str:
-    op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
-    name = inst.op_name.lower() or f'vop3b_op_{op_val}'
-    n = inst.num_srcs() or _num_srcs(inst)
-    regs = inst.canonical_op_regs
-    dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
-    s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, n=r0), _cdna_src(inst, inst.src1, inst.neg&2, n=r1), _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
-    dst = _vreg(inst.vdst, dregs) if dregs > 1 else _vreg(inst.vdst)
-    sdst = _fmt_sdst(inst.sdst, 2, cdna=True)
-    cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
-    if name in ('v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'):
-      s2 = _fmt_src(inst.src2, 2, cdna=True)
-      return f"{name}_e64 {dst}, {sdst}, {s0}, {s1}, {s2}{cl}{om}"
-    suf = "_e64" if 'co_' in name else ""
-    return f"{name}{suf} {dst}, {sdst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name}{suf} {dst}, {sdst}, {s0}, {s1}{cl}{om}"
-
-  def _disasm_cdna_vop3p(inst) -> str:
-    name, n = inst.op_name.lower(), inst.num_srcs() or 2
-    is_mfma = 'mfma' in name or 'smfmac' in name
-    is_accvgpr = 'accvgpr' in name
-    get_src = lambda v, sc: _lit(inst, v) if v == 255 else _fmt_src(v, sc, cdna=True)
-
-    # Handle accvgpr read/write (accumulator register operations)
-    if is_accvgpr:
-      src0_off = _unwrap(inst.src0)
-      vdst_off = _vi(inst.vdst)
-      if 'read' in name:
-        # v_accvgpr_read_b32 vN, aM - reads from accumulator to VGPR
-        return f"{name}_b32 v{vdst_off}, a{src0_off - 256 if src0_off >= 256 else src0_off}"
-      if 'write' in name:
-        # v_accvgpr_write_b32 aM, src - writes to accumulator from source
-        src = _lit(inst, inst.src0) if src0_off == 255 else (f"v{src0_off - 256}" if src0_off >= 256 else decode_src(src0_off, cdna=True))
-        return f"{name}_b32 a{vdst_off}, {src}"
-
-    # Handle MFMA instructions with accumulator destinations
-    if is_mfma:
-      sc = 2 if 'iu4' in name else 4 if 'iu8' in name or 'i4' in name else 8 if 'f16' in name or 'bf16' in name else 4
-      src0, src1, src2 = get_src(inst.src0, sc), get_src(inst.src1, sc), get_src(inst.src2, 16)
-      dst = _areg(inst.vdst, 16)  # MFMA uses accumulator registers
-      opsel_hi = inst.opsel_hi
-      mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != 3 else []) + \
-             ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
-      return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}"
-
-    # Standard VOP3P instructions
-    src0, src1, src2, dst = get_src(inst.src0, 1), get_src(inst.src1, 1), get_src(inst.src2, 1), _vreg(inst.vdst)
-    opsel_hi = inst.opsel_hi  # CDNA VOP3P only has 2 bits for opsel_hi (no opsel_hi2)
-    opsel_hi_default = 3  # CDNA default is 0b11 (2 bits), not 0b111 like RDNA
-    mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else []) + \
-           ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
-    return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
-
-  DISASM_HANDLERS.update({CDNA_VOP1: _disasm_vop1, CDNA_VOP2: _disasm_vop2, CDNA_VOPC: _disasm_vopc,
-    CDNA_SOP1: _disasm_sop1, CDNA_SOP2: _disasm_sop2, CDNA_SOPC: _disasm_sopc, CDNA_SOPK: _disasm_sopk, CDNA_SOPP: _disasm_sopp,
-    CDNA_SMEM: _disasm_smem, CDNA_DS: _disasm_ds, CDNA_FLAT: _disasm_flat,
-    VOP3A: _disasm_vop3a, VOP3B: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p})
-except ImportError:
-  pass

extra/assembly/amd/emu.py

@@ -1,479 +0,0 @@
-# RDNA3 emulator - executes compiled pseudocode from AMD ISA PDF
-# mypy: ignore-errors
-from __future__ import annotations
-import ctypes, functools
-from enum import IntEnum
-from tinygrad.runtime.autogen import hsa
-from extra.assembly.amd.dsl import Inst, NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, v, s
-from extra.assembly.amd.pcode import _f32, _i32, _sext, _f16, _i16, _f64, _i64
-from extra.assembly.amd.decode import decode_inst
-from extra.assembly.amd.pcode import compile_pseudocode
-from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
-from extra.assembly.amd.autogen.rdna3.ins import (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD,
-  SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, SCRATCHOp, VOPDOp)
-
-# Constants and helpers defined locally (not imported from dsl.py)
-MASK32, MASK64 = 0xFFFFFFFF, 0xFFFFFFFFFFFFFFFF
-FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
-
-class SGPRArray:
-  """SGPR array indexed by Reg or int."""
-  __slots__ = ('_data',)
-  def __init__(self, size: int): self._data = [0] * size
-  def __getitem__(self, key): return self._data[getattr(key, 'offset', key)]
-  def __setitem__(self, key, val): self._data[getattr(key, 'offset', key)] = val
-  def __len__(self): return len(self._data)
-  def __iter__(self): return iter(self._data)
-
-class VGPRLane:
-  """Single lane of VGPRs indexed by Reg (offset 256-511) or int (0-255)."""
-  __slots__ = ('_data',)
-  def __init__(self, size: int): self._data = [0] * size
-  def __getitem__(self, key):
-    i = getattr(key, 'offset', key)
-    return self._data[i - 256 if i >= 256 else i]
-  def __setitem__(self, key, val):
-    i = getattr(key, 'offset', key)
-    self._data[i - 256 if i >= 256 else i] = val
-  def __len__(self): return len(self._data)
-  def __iter__(self): return iter(self._data)
-
-WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
-
-# Inline constants for src operands 128-254. Build tables for f32, f16, and f64 formats.
-_FLOAT_CONSTS = {v: k for k, v in FLOAT_ENC.items()} | {248: 0.15915494309189535}  # INV_2PI
-def _build_inline_consts(mask, to_bits):
-  tbl = list(range(65)) + [((-i) & mask) for i in range(1, 17)] + [0] * (127 - 81)
-  for k, v in _FLOAT_CONSTS.items(): tbl[k - 128] = to_bits(v)
-  return tbl
-_INLINE_CONSTS = _build_inline_consts(MASK32, _i32)
-_INLINE_CONSTS_F16 = _build_inline_consts(0xffff, _i16)
-_INLINE_CONSTS_F64 = _build_inline_consts(MASK64, _i64)
-
-# Helper: extract/write 16-bit half from/to 32-bit value
-def _src16(raw: int, is_hi: bool) -> int: return ((raw >> 16) & 0xffff) if is_hi else (raw & 0xffff)
-def _dst16(cur: int, val: int, is_hi: bool) -> int: return (cur & 0x0000ffff) | ((val & 0xffff) << 16) if is_hi else (cur & 0xffff0000) | (val & 0xffff)
-def _vgpr_hi(src) -> bool: return src.offset >= 256 and ((src.offset - 256) & 0x80) != 0
-def _vgpr_masked(src): return v[(src.offset - 256) & 0x7f] if src.offset >= 256 else src
-
-# VOP3 source modifier: apply abs/neg to value
-def _mod_src(val: int, idx: int, neg: int, abs_: int, is64: bool = False) -> int:
-  to_f, to_i = (_f64, _i64) if is64 else (_f32, _i32)
-  if (abs_ >> idx) & 1: val = to_i(abs(to_f(val)))
-  if (neg >> idx) & 1: val = to_i(-to_f(val))
-  return val
-
-# Read source operand with VOP3 modifiers
-def _read_src(st, inst, src, idx: int, lane: int, neg: int, abs_: int, opsel: int) -> int:
-  if src is None: return 0
-  src_off = src.offset
-  src_bits = inst.canonical_op_bits[f's{idx}']
-  literal, is_src_64, is_src_16 = inst._literal, src_bits == 64, src_bits == 16
-  if is_src_64: return _mod_src(st.rsrc64(src, lane, literal), idx, neg, abs_, is64=True)
-  if isinstance(inst, VOP3P):
-    opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
-    if 'FMA_MIX' in inst.op_name:
-      raw = st.rsrc(src, lane, literal)
-      sign_bit = (15 if not (opsel & (1 << idx)) else 31) if (opsel_hi >> idx) & 1 else 31
-      if inst.neg_hi & (1 << idx): raw &= ~(1 << sign_bit)
-      if neg & (1 << idx): raw ^= (1 << sign_bit)
-      return raw
-    raw = st.rsrc_f16(src, lane, literal)
-    hi = _src16(raw, opsel_hi & (1 << idx)) ^ (0x8000 if inst.neg_hi & (1 << idx) else 0)
-    lo = _src16(raw, opsel & (1 << idx)) ^ (0x8000 if neg & (1 << idx) else 0)
-    return (hi << 16) | lo
-  if is_src_16 and isinstance(inst, VOP3):
-    raw = st.rsrc_f16(src, lane, literal) if 128 <= src_off < 255 else st.rsrc(src, lane, literal)
-    val = _src16(raw, bool(opsel & (1 << idx)))
-    if abs_ & (1 << idx): val &= 0x7fff
-    if neg & (1 << idx): val ^= 0x8000
-    return val
-  if is_src_16 and isinstance(inst, (VOP1, VOP2, VOPC)):
-    if src_off >= 256: return _src16(_mod_src(st.rsrc(_vgpr_masked(src), lane, literal), idx, neg, abs_), _vgpr_hi(src))
-    return _mod_src(st.rsrc_f16(src, lane, literal), idx, neg, abs_) & 0xffff
-  return _mod_src(st.rsrc(src, lane, literal), idx, neg, abs_)
-
-# Helper: get number of dwords from memory op name
-def _op_ndwords(name: str) -> int:
-  if '_B128' in name: return 4
-  if '_B96' in name: return 3
-  if any(s in name for s in ('_B64', '_U64', '_I64', '_F64')): return 2
-  return 1
-
-# Helper: build multi-dword int from consecutive VGPRs
-def _vgpr_read(V: VGPRLane, reg, ndwords: int) -> int:
-  return sum(V[reg + i] << (32 * i) for i in range(ndwords))
-
-# Helper: write multi-dword value to consecutive VGPRs
-def _vgpr_write(V: VGPRLane, reg, val: int, ndwords: int):
-  for i in range(ndwords): V[reg + i] = (val >> (32 * i)) & MASK32
-
-# Memory access
-_valid_mem_ranges: list[tuple[int, int]] = []
-def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: _valid_mem_ranges.clear(); _valid_mem_ranges.extend(ranges)
-def _mem_valid(addr: int, size: int) -> bool:
-  return not _valid_mem_ranges or any(s <= addr and addr + size <= s + z for s, z in _valid_mem_ranges)
-def _ctypes_at(addr: int, size: int): return (ctypes.c_uint8 if size == 1 else ctypes.c_uint16 if size == 2 else ctypes.c_uint64 if size == 8 else ctypes.c_uint32).from_address(addr)
-def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
-def mem_write(addr: int, size: int, val: int) -> None:
-  if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
-
-def _make_mem_accessor(read_fn, write_fn):
-  """Create a memory accessor class with the given read/write functions."""
-  class _MemAccessor:
-    __slots__ = ('_addr',)
-    def __init__(self, addr: int): self._addr = int(addr)
-    u8 = property(lambda s: read_fn(s._addr, 1), lambda s, v: write_fn(s._addr, 1, int(v)))
-    u16 = property(lambda s: read_fn(s._addr, 2), lambda s, v: write_fn(s._addr, 2, int(v)))
-    u32 = property(lambda s: read_fn(s._addr, 4), lambda s, v: write_fn(s._addr, 4, int(v)))
-    u64 = property(lambda s: read_fn(s._addr, 8), lambda s, v: write_fn(s._addr, 8, int(v)))
-    i8 = property(lambda s: _sext(read_fn(s._addr, 1), 8), lambda s, v: write_fn(s._addr, 1, int(v)))
-    i16 = property(lambda s: _sext(read_fn(s._addr, 2), 16), lambda s, v: write_fn(s._addr, 2, int(v)))
-    i32 = property(lambda s: _sext(read_fn(s._addr, 4), 32), lambda s, v: write_fn(s._addr, 4, int(v)))
-    i64 = property(lambda s: _sext(read_fn(s._addr, 8), 64), lambda s, v: write_fn(s._addr, 8, int(v)))
-    b8, b16, b32, b64 = u8, u16, u32, u64
-  return _MemAccessor
-
-_GlobalMemAccessor = _make_mem_accessor(mem_read, mem_write)
-
-class _GlobalMem:
-  """Global memory wrapper that supports MEM[addr].u32 style access."""
-  def __getitem__(self, addr) -> _GlobalMemAccessor: return _GlobalMemAccessor(addr)
-GlobalMem = _GlobalMem()
-
-class LDSMem:
-  """LDS memory wrapper that supports MEM[addr].u32 style access."""
-  __slots__ = ('_lds',)
-  def __init__(self, lds: bytearray): self._lds = lds
-  def _read(self, addr: int, size: int) -> int:
-    addr = addr & 0xffff
-    return int.from_bytes(self._lds[addr:addr+size], 'little') if addr + size <= len(self._lds) else 0
-  def _write(self, addr: int, size: int, val: int):
-    addr = addr & 0xffff
-    if addr + size <= len(self._lds): self._lds[addr:addr+size] = (int(val) & ((1 << (size*8)) - 1)).to_bytes(size, 'little')
-  def __getitem__(self, addr): return _make_mem_accessor(self._read, self._write)(addr)
-
-# SMEM dst register count (for writing result back to SGPRs)
-SMEM_DST_COUNT = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}
-
-# VOPD op -> VOP3 op mapping (VOPD is dual-issue of VOP1/VOP2 ops, use VOP3 enums for pseudocode lookup)
-_VOPD_TO_VOP = {
-  VOPDOp.V_DUAL_FMAC_F32: VOP3Op.V_FMAC_F32_E64, VOPDOp.V_DUAL_FMAAK_F32: VOP2Op.V_FMAAK_F32_E32, VOPDOp.V_DUAL_FMAMK_F32: VOP2Op.V_FMAMK_F32_E32,
-  VOPDOp.V_DUAL_MUL_F32: VOP3Op.V_MUL_F32_E64, VOPDOp.V_DUAL_ADD_F32: VOP3Op.V_ADD_F32_E64, VOPDOp.V_DUAL_SUB_F32: VOP3Op.V_SUB_F32_E64,
-  VOPDOp.V_DUAL_SUBREV_F32: VOP3Op.V_SUBREV_F32_E64, VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: VOP3Op.V_MUL_DX9_ZERO_F32_E64,
-  VOPDOp.V_DUAL_MOV_B32: VOP3Op.V_MOV_B32_E64, VOPDOp.V_DUAL_CNDMASK_B32: VOP3Op.V_CNDMASK_B32_E64,
-  VOPDOp.V_DUAL_MAX_F32: VOP3Op.V_MAX_F32_E64, VOPDOp.V_DUAL_MIN_F32: VOP3Op.V_MIN_F32_E64,
-  VOPDOp.V_DUAL_ADD_NC_U32: VOP3Op.V_ADD_NC_U32_E64, VOPDOp.V_DUAL_LSHLREV_B32: VOP3Op.V_LSHLREV_B32_E64, VOPDOp.V_DUAL_AND_B32: VOP3Op.V_AND_B32_E64,
-}
-
-
-class WaveState:
-  __slots__ = ('sgpr', 'vgpr', 'scc', 'pc', '_pend_sgpr', 'lds', 'n_lanes')
-  def __init__(self, lds: LDSMem | None = None, n_lanes: int = WAVE_SIZE):
-    self.sgpr, self.vgpr = SGPRArray(SGPR_COUNT), [VGPRLane(VGPR_COUNT) for _ in range(WAVE_SIZE)]
-    self.sgpr[EXEC_LO], self.scc, self.pc, self._pend_sgpr, self.lds, self.n_lanes = 0xffffffff, 0, 0, {}, lds, n_lanes
-
-  @property
-  def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
-  @vcc.setter
-  def vcc(self, v: int): self.sgpr[VCC_LO], self.sgpr[VCC_HI] = v & MASK32, (v >> 32) & MASK32
-  @property
-  def exec_mask(self) -> int: return self.sgpr[EXEC_LO] | (self.sgpr[EXEC_HI] << 32)
-  @exec_mask.setter
-  def exec_mask(self, v: int): self.sgpr[EXEC_LO], self.sgpr[EXEC_HI] = v & MASK32, (v >> 32) & MASK32
-
-  def rsgpr(self, reg) -> int:
-    if reg == NULL: return 0
-    if reg == SCC: return self.scc
-    return self.sgpr[reg]
-  def wsgpr(self, reg, v: int):
-    if reg != NULL: self.sgpr[reg] = v & MASK32
-  def rsgpr64(self, reg) -> int:
-    off = reg.offset
-    return self.sgpr._data[off] | (self.sgpr._data[off + 1] << 32)
-  def wsgpr64(self, reg, v: int):
-    off = reg.offset
-    self.sgpr._data[off] = v & MASK32; self.sgpr._data[off + 1] = (v >> 32) & MASK32
-
-  def _rsrc_base(self, reg, lane: int, consts, literal: int):
-    off = reg.offset
-    if off < SGPR_COUNT: return self.sgpr._data[off]
-    if off == SCC.offset: return self.scc
-    if off < 255: return consts[off - 128]
-    if off == 255: return literal
-    return self.vgpr[lane]._data[off - 256] if off <= 511 else 0
-  def rsrc(self, reg, lane: int, literal: int = 0) -> int: return self._rsrc_base(reg, lane, _INLINE_CONSTS, literal)
-  def rsrc_f16(self, reg, lane: int, literal: int = 0) -> int: return self._rsrc_base(reg, lane, _INLINE_CONSTS_F16, literal)
-  def rsrc64(self, reg, lane: int, literal: int = 0) -> int:
-    off = reg.offset
-    if 128 <= off < 255: return _INLINE_CONSTS_F64[off - 128]
-    if off == 255: return literal << 32  # 32-bit literal forms upper 32 bits of 64-bit value
-    return self.rsrc(reg, lane, literal) | ((self.rsrc(reg + 1, lane, literal) if off < VCC_LO.offset or 256 <= off <= 511 else 0) << 32)
-
-  def pend_sgpr_lane(self, reg, lane: int, val: int):
-    if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
-    if val: self._pend_sgpr[reg] |= (1 << lane)
-  def commit_pends(self):
-    for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
-    self._pend_sgpr.clear()
-
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# EXECUTION - All ops use pseudocode from PDF
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def exec_scalar(st: WaveState, inst: Inst):
-  """Execute scalar instruction. Returns 0 to continue execution."""
-  # Get op enum and lookup compiled function
-  if isinstance(inst, SMEM): ssrc0, sdst = None, None
-  elif isinstance(inst, SOP1): ssrc0, sdst = inst.ssrc0, inst.sdst
-  elif isinstance(inst, SOP2): ssrc0, sdst = inst.ssrc0, inst.sdst
-  elif isinstance(inst, SOPC): ssrc0, sdst = inst.ssrc0, None
-  elif isinstance(inst, SOPK): ssrc0, sdst = inst.sdst, inst.sdst  # sdst is both src and dst
-  elif isinstance(inst, SOPP): ssrc0, sdst = None, None
-  else: raise NotImplementedError(f"Unknown scalar type {type(inst)}")
-
-  # SMEM: memory loads
-  if isinstance(inst, SMEM):
-    addr = st.rsgpr64(inst.sbase) + _sext(inst.offset, 21)
-    if inst.soffset != NULL: addr += st.rsrc(inst.soffset, 0, inst._literal)
-    result = inst._fn(GlobalMem, addr & MASK64)
-    if 'SDATA' in result:
-      sdata = result['SDATA']
-      for i in range(SMEM_DST_COUNT.get(inst.op, 1)): st.wsgpr(inst.sdata + i, (sdata >> (i * 32)) & MASK32)
-    st.pc += inst._words
-    return 0
-
-  # Build context - use canonical_op_bits to determine operand sizes
-  literal = inst._literal
-  s0 = st.rsrc64(ssrc0, 0, literal) if inst.canonical_op_bits['s0'] == 64 else (st.rsrc(ssrc0, 0, literal) if not isinstance(inst, (SOPK, SOPP)) else (st.rsgpr(inst.sdst) if isinstance(inst, SOPK) else 0))
-  s1 = st.rsrc64(inst.ssrc1, 0, literal) if inst.canonical_op_bits['s1'] == 64 else (st.rsrc(inst.ssrc1, 0, literal) if isinstance(inst, (SOP2, SOPC)) else inst.simm16 if isinstance(inst, SOPK) else 0)
-  d0 = st.rsgpr64(sdst) if inst.canonical_op_bits['d'] == 64 and sdst is not None else (st.rsgpr(sdst) if sdst is not None else 0)
-  literal = inst.simm16 if isinstance(inst, (SOPK, SOPP)) else inst._literal
-
-  # Call compiled function with int parameters
-  result = inst._fn(s0, s1, 0, d0, st.scc, st.vcc & MASK32, 0, st.exec_mask & MASK32, literal, None, pc=st.pc * 4)
-
-  # Apply results (already int values)
-  if sdst is not None and 'D0' in result:
-    (st.wsgpr64 if inst.canonical_op_bits['d'] == 64 else st.wsgpr)(sdst, result['D0'])
-  if 'SCC' in result: st.scc = result['SCC'] & 1
-  if 'EXEC' in result: st.exec_mask = result['EXEC']
-  if 'PC' in result:
-    # Convert absolute byte address to word offset
-    pc_val = result['PC']
-    new_pc = pc_val if pc_val < 0x8000000000000000 else pc_val - 0x10000000000000000
-    st.pc = new_pc // 4
-  else:
-    st.pc += inst._words
-  return 0
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# VECTOR INSTRUCTIONS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def exec_vopd(st: WaveState, inst, V: VGPRLane, lane: int) -> None:
-  """VOPD: dual-issue, execute two ops simultaneously (read all inputs before writes)."""
-  literal, vdstx = inst._literal, inst.vdstx
-  vdsty = v[(inst.vdsty << 1) | ((inst.vdstx.offset & 1) ^ 1)]  # vdsty is raw int from VDSTYField.decode
-  sx0, sx1, dx, sy0, sy1, dy = st.rsrc(inst.srcx0, lane, literal), V[inst.vsrcx1], V[vdstx], st.rsrc(inst.srcy0, lane, literal), V[inst.vsrcy1], V[vdsty]
-  V[vdstx] = inst._fnx(sx0, sx1, 0, dx, st.scc, st.vcc, lane, st.exec_mask, literal, None)['D0']
-  V[vdsty] = inst._fny(sy0, sy1, 0, dy, st.scc, st.vcc, lane, st.exec_mask, literal, None)['D0']
-
-def exec_flat(st: WaveState, inst, V: VGPRLane, lane: int) -> None:
-  """FLAT/GLOBAL/SCRATCH memory ops."""
-  ndwords = _op_ndwords(inst.op_name)
-  addr = V[inst.addr] | (V[inst.addr + 1] << 32)
-  ADDR = (st.rsgpr64(inst.saddr) + V[inst.addr] + _sext(inst.offset, 13)) & MASK64 if inst.saddr != NULL else (addr + _sext(inst.offset, 13)) & MASK64
-  vdata_src = inst.vdst if 'LOAD' in inst.op_name else inst.data
-  result = inst._fn(GlobalMem, ADDR, _vgpr_read(V, vdata_src, ndwords), V[inst.vdst])
-  if 'VDATA' in result: _vgpr_write(V, inst.vdst, result['VDATA'], ndwords)
-  if 'RETURN_DATA' in result: _vgpr_write(V, inst.vdst, result['RETURN_DATA'], ndwords)
-
-def exec_ds(st: WaveState, inst, V: VGPRLane, lane: int) -> None:
-  """DS (LDS) memory ops."""
-  ndwords = _op_ndwords(inst.op_name)
-  data0, data1 = _vgpr_read(V, inst.data0, ndwords), _vgpr_read(V, inst.data1, ndwords) if inst.data1 is not None else 0
-  result = inst._fn(st.lds, V[inst.addr], data0, data1, inst.offset0, inst.offset1)
-  if 'RETURN_DATA' in result and ('_RTN' in inst.op_name or '_LOAD' in inst.op_name):
-    _vgpr_write(V, inst.vdst, result['RETURN_DATA'], ndwords * 2 if '_2ADDR_' in inst.op_name else ndwords)
-
-def exec_vop(st: WaveState, inst: Inst, V: VGPRLane, lane: int) -> None:
-  """VOP1/VOP2/VOP3/VOP3SD/VOP3P/VOPC: standard ALU ops."""
-  is_dst_16 = inst.canonical_op_bits['d'] == 16
-  if isinstance(inst, VOP3P):
-    src0, src1, src2, vdst, dst_hi = inst.src0, inst.src1, inst.src2, inst.vdst, False
-    neg, abs_, opsel = inst.neg, 0, inst.opsel
-  elif isinstance(inst, VOP1):
-    src0, src1, src2, vdst = inst.src0, None, None, inst.vdst
-    neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst.offset & 0x80) != 0 and is_dst_16
-    if is_dst_16: vdst = v[inst.vdst.offset & 0x7f]
-  elif isinstance(inst, VOP2):
-    src0, src1, src2, vdst = inst.src0, inst.vsrc1, None, inst.vdst
-    neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst.offset & 0x80) != 0 and is_dst_16
-    if is_dst_16: vdst = v[inst.vdst.offset & 0x7f]
-  elif isinstance(inst, (VOP3, VOP3SD)):
-    src0, src1, src2, vdst = inst.src0, inst.src1, (None if isinstance(inst, VOP3) and inst.op.value < 256 else inst.src2), inst.vdst
-    neg, abs_, opsel, dst_hi = (inst.neg, inst.abs, inst.opsel, False) if isinstance(inst, VOP3) else (0, 0, 0, False)
-  elif isinstance(inst, VOPC):
-    src0, src1, src2, vdst, neg, abs_, opsel, dst_hi = inst.src0, inst.vsrc1, None, VCC_LO, 0, 0, 0, False
-  else:
-    raise NotImplementedError(f"exec_vop: unhandled instruction type {type(inst).__name__}")
-
-  s0 = _read_src(st, inst, src0, 0, lane, neg, abs_, opsel)
-  s1 = _read_src(st, inst, src1, 1, lane, neg, abs_, opsel)
-  s2 = _read_src(st, inst, src2, 2, lane, neg, abs_, opsel)
-  if isinstance(inst, VOP2) and is_dst_16: d0 = _src16(V[vdst], dst_hi)
-  elif inst.canonical_op_bits['d'] == 64: d0 = V[vdst] | (V[vdst + 1] << 32)
-  else: d0 = V[vdst]
-
-  if isinstance(inst, VOP3SD) and 'CO_CI' in inst.op_name: vcc_for_fn = st.rsgpr64(inst.src2)
-  elif isinstance(inst, VOP3) and inst.op in (VOP3Op.V_CNDMASK_B32_E64, VOP3Op.V_CNDMASK_B16) and src2 is not None and src2.offset < 256: vcc_for_fn = st.rsgpr64(src2)
-  else: vcc_for_fn = st.vcc
-  src0_off = src0.offset if src0 is not None else 0
-  src0_idx = (src0_off - 256) if src0_off >= 256 else src0_off
-  vdst_off = vdst.offset
-  extra_kwargs = {'opsel': opsel, 'opsel_hi': inst.opsel_hi | (inst.opsel_hi2 << 2)} if isinstance(inst, VOP3P) and 'FMA_MIX' in inst.op_name else {}
-  result = inst._fn(s0, s1, s2, d0, st.scc, vcc_for_fn, lane, st.exec_mask, inst._literal, st.vgpr, src0_idx, vdst_off, **extra_kwargs)
-
-  # Check if this is a VOPC instruction (either standalone VOPC or VOP3 with VOPC opcode)
-  is_vopc = isinstance(inst.op, VOPCOp) or (isinstance(inst, VOP3) and inst.op.value < 256)
-  if 'VCC' in result:
-    if isinstance(inst, VOP3SD): st.pend_sgpr_lane(inst.sdst, lane, (result['VCC'] >> lane) & 1)
-    elif isinstance(inst, VOP2) and 'CO_CI' in inst.op_name: st.pend_sgpr_lane(VCC_LO, lane, (result['VCC'] >> lane) & 1)
-    elif is_vopc: st.pend_sgpr_lane(vdst, lane, (result['VCC'] >> lane) & 1)  # vdst is VCC_LO for VOPC
-    else: st.pend_sgpr_lane(VCC_LO, lane, (result['VCC'] >> lane) & 1)
-  if 'EXEC' in result:
-    st.pend_sgpr_lane(EXEC_LO, lane, (result['EXEC'] >> lane) & 1)
-  elif is_vopc:
-    st.pend_sgpr_lane(vdst, lane, (result['D0'] >> lane) & 1)
-  if not is_vopc:
-    d0_val = result['D0']
-    if inst.canonical_op_bits['d'] == 64: V[vdst], V[vdst + 1] = d0_val & MASK32, (d0_val >> 32) & MASK32
-    elif not isinstance(inst, VOP3P) and is_dst_16: V[vdst] = _dst16(V[vdst], d0_val, bool(opsel & 8) if isinstance(inst, VOP3) else dst_hi)
-    else: V[vdst] = d0_val & MASK32
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# WMMA (Wave Matrix Multiply-Accumulate)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def exec_wmma(st: WaveState, inst, op: VOP3POp) -> None:
-  """Execute WMMA instruction - 16x16x16 matrix multiply across the wave."""
-  src0, src1, src2, vdst = inst.src0.offset, inst.src1.offset, inst.src2.offset, inst.vdst.offset
-  # Read 16x16 f16 matrix from 16 lanes × 8 VGPRs (2 f16 per VGPR)
-  def read_f16_mat(src):
-    return [f for l in range(16) for r in range(8) for v in [st.vgpr[l][src-256+r] if src >= 256 else st.rsgpr(src+r)] for f in [_f16(v&0xffff), _f16((v>>16)&0xffff)]]
-  mat_a, mat_b = read_f16_mat(src0), read_f16_mat(src1)
-  # Read matrix C (16x16 f32) from lanes 0-31, VGPRs src2 to src2+7
-  mat_c = [_f32(st.vgpr[i % 32][src2 - 256 + i // 32] if src2 >= 256 else st.rsgpr(src2 + i // 32)) for i in range(256)]
-  # Compute D = A × B + C (16x16 matrix multiply)
-  mat_d = [sum(mat_a[row*16+k] * mat_b[col*16+k] for k in range(16)) + mat_c[row*16+col] for row in range(16) for col in range(16)]
-  # Write result - f16 packed or f32
-  if op == VOP3POp.V_WMMA_F16_16X16X16_F16:
-    for i in range(0, 256, 2):
-      st.vgpr[(i//2) % 32][vdst - 256 + (i//2)//32] = ((_i16(mat_d[i+1]) & 0xffff) << 16) | (_i16(mat_d[i]) & 0xffff)
-  else:
-    for i in range(256): st.vgpr[i % 32][vdst - 256 + i//32] = _i32(mat_d[i])
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# PROGRAM DECODE
-# ═══════════════════════════════════════════════════════════════════════════════
-
-# Wave-level dispatch functions: (st, inst) -> return_code (0 = continue, -1 = end, -2 = barrier)
-def dispatch_endpgm(st, inst): return -1
-def dispatch_barrier(st, inst): st.pc += inst._words; return -2
-def dispatch_nop(st, inst): st.pc += inst._words; return 0
-def dispatch_wmma(st, inst): exec_wmma(st, inst, inst.op); st.pc += inst._words; return 0
-def dispatch_writelane(st, inst): st.vgpr[st.rsrc(inst.src1, 0, inst._literal) & 0x1f][inst.vdst.offset - 256] = st.rsrc(inst.src0, 0, inst._literal) & MASK32; st.pc += inst._words; return 0
-def dispatch_readlane(st, inst):
-  src0_off = inst.src0.offset
-  src0_idx = (src0_off - 256) if src0_off >= 256 else src0_off
-  s1 = st.rsrc(inst.src1, 0, inst._literal) if getattr(inst, 'src1', None) is not None else 0
-  result = inst._fn(0, s1, 0, 0, st.scc, st.vcc, 0, st.exec_mask, inst._literal, st.vgpr, src0_idx, inst.vdst.offset)
-  st.wsgpr(inst.vdst.offset, result['D0'])
-  st.pc += inst._words; return 0
-
-# Per-lane dispatch wrapper: wraps per-lane exec functions into wave-level dispatch
-@functools.cache
-def dispatch_lane(exec_fn):
-  def dispatch(st, inst):
-    exec_mask, vgpr, n_lanes = st.exec_mask, st.vgpr, st.n_lanes
-    for lane in range(n_lanes):
-      if exec_mask >> lane & 1: exec_fn(st, inst, vgpr[lane], lane)
-    st.commit_pends()
-    st.pc += inst._words
-    return 0
-  return dispatch
-
-def decode_program(data: bytes) -> dict[int, Inst]:
-  result: dict[int, Inst] = {}
-  i = 0
-  while i < len(data):
-    inst = decode_inst(data[i:])
-    inst._words = inst.size() // 4
-
-    # Determine dispatch function and pcode function
-    if isinstance(inst, SOPP) and inst.op == SOPPOp.S_CODE_END: break
-    elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_ENDPGM: inst._dispatch = dispatch_endpgm
-    elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER: inst._dispatch = dispatch_barrier
-    elif isinstance(inst, SOPP) and inst.op in (SOPPOp.S_CLAUSE, SOPPOp.S_WAITCNT, SOPPOp.S_WAITCNT_DEPCTR, SOPPOp.S_SENDMSG, SOPPOp.S_SET_INST_PREFETCH_DISTANCE, SOPPOp.S_DELAY_ALU): inst._dispatch = dispatch_nop
-    elif isinstance(inst, (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM)): inst._dispatch = exec_scalar
-    elif isinstance(inst, VOP1) and inst.op == VOP1Op.V_NOP_E32: inst._dispatch = dispatch_nop
-    elif isinstance(inst, VOP3P) and 'WMMA' in inst.op_name: inst._dispatch = dispatch_wmma
-    elif isinstance(inst, VOP3) and inst.op == VOP3Op.V_WRITELANE_B32: inst._dispatch = dispatch_writelane
-    elif isinstance(inst, (VOP1, VOP3)) and inst.op in (VOP1Op.V_READFIRSTLANE_B32_E32, VOP3Op.V_READFIRSTLANE_B32, VOP3Op.V_READLANE_B32): inst._dispatch = dispatch_readlane
-    elif isinstance(inst, VOPD): inst._dispatch = dispatch_lane(exec_vopd)
-    elif isinstance(inst, FLAT): inst._dispatch = dispatch_lane(exec_flat)
-    elif isinstance(inst, DS): inst._dispatch = dispatch_lane(exec_ds)
-    else: inst._dispatch = dispatch_lane(exec_vop)
-
-    # Compile pcode for instructions that use it (not VOPD which has _fnx/_fny, not special dispatches)
-    # VOPD needs separate functions for X and Y ops
-    if isinstance(inst, VOPD):
-      def _compile_vopd_op(op): return compile_pseudocode(type(op).__name__, op.name, PCODE[op])
-      inst._fnx, inst._fny = _compile_vopd_op(_VOPD_TO_VOP[inst.opx]), _compile_vopd_op(_VOPD_TO_VOP[inst.opy])
-    elif inst._dispatch not in (dispatch_endpgm, dispatch_barrier, dispatch_nop, dispatch_wmma, dispatch_writelane):
-      assert type(inst.op) != int, f"inst op of {inst} is int"
-      inst._fn = compile_pseudocode(type(inst.op).__name__, inst.op.name, PCODE[inst.op])
-    result[i // 4] = inst
-    i += inst._words * 4
-  return result
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# MAIN EXECUTION LOOP
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def exec_wave(program: dict[int, Inst], st: WaveState) -> int:
-  while (inst := program.get(st.pc)) and (result := inst._dispatch(st, inst)) == 0: pass
-  return result
-
-def exec_workgroup(program: dict[int, Inst], workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int, rsrc2: int) -> None:
-  lx, ly, lz = local_size
-  total_threads = lx * ly * lz
-  # GRANULATED_LDS_SIZE is in 512-byte units (see ops_amd.py: lds_size = ((group_segment_size + 511) // 512))
-  lds_size = ((rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE_SHIFT) * 512
-  lds = LDSMem(bytearray(lds_size)) if lds_size else None
-  waves: list[WaveState] = []
-  for wave_start in range(0, total_threads, WAVE_SIZE):
-    n_lanes = min(WAVE_SIZE, total_threads - wave_start)
-    st = WaveState(lds, n_lanes)
-    st.exec_mask = (1 << n_lanes) - 1
-    st.wsgpr64(s[0:1], args_ptr)  # s[0:1] = kernel arguments pointer
-    # COMPUTE_PGM_RSRC2: USER_SGPR_COUNT is where workgroup IDs start, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z control which are passed
-    sgpr_idx = (rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT_SHIFT
-    if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X: st.sgpr[sgpr_idx] = workgroup_id[0]; sgpr_idx += 1
-    if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Y: st.sgpr[sgpr_idx] = workgroup_id[1]; sgpr_idx += 1
-    if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Z: st.sgpr[sgpr_idx] = workgroup_id[2]
-    # VGPR0 = packed workitem IDs: (Z << 20) | (Y << 10) | X
-    for tid in range(wave_start, wave_start + n_lanes):
-      st.vgpr[tid - wave_start][0] = ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
-    waves.append(st)
-  while waves:
-    waves = [st for st in waves if exec_wave(program, st) != -1]
-
-def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int, rsrc2: int = 0x19c) -> int:
-  program = decode_program((ctypes.c_char * lib_sz).from_address(lib).raw)
-  for gidz in range(gz):
-    for gidy in range(gy):
-      for gidx in range(gx): exec_workgroup(program, (gidx, gidy, gidz), (lx, ly, lz), args_ptr, rsrc2)
-  return 0

extra/assembly/amd/pcode.py

@@ -1,822 +0,0 @@
-# DSL for RDNA3 pseudocode - makes pseudocode expressions work directly as Python
-import struct, math, re, functools
-
-MASK32, MASK64 = 0xFFFFFFFF, 0xFFFFFFFFFFFFFFFF
-
-# Float/int bit conversion functions
-_struct_f, _struct_I = struct.Struct("<f"), struct.Struct("<I")
-_struct_e, _struct_H = struct.Struct("<e"), struct.Struct("<H")
-_struct_d, _struct_Q = struct.Struct("<d"), struct.Struct("<Q")
-def _f32(i):
-  i = i & MASK32
-  # RDNA3 default mode: flush f32 denormals to zero (FTZ)
-  # Denormal: exponent=0 (bits 23-30) and mantissa!=0 (bits 0-22)
-  if (i & 0x7f800000) == 0 and (i & 0x007fffff) != 0: return 0.0
-  return _struct_f.unpack(_struct_I.pack(i))[0]
-def _i32(f):
-  if isinstance(f, int): f = float(f)
-  if math.isnan(f): return 0xffc00000 if math.copysign(1.0, f) < 0 else 0x7fc00000
-  if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
-  try:
-    bits = _struct_I.unpack(_struct_f.pack(f))[0]
-    # RDNA3 default mode: flush f32 denormals to zero (FTZ)
-    if (bits & 0x7f800000) == 0 and (bits & 0x007fffff) != 0: return 0x80000000 if bits & 0x80000000 else 0
-    return bits
-  except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
-def _sext(v, b): return v - (1 << b) if v & (1 << (b - 1)) else v
-def _f16(i): return _struct_e.unpack(_struct_H.pack(i & 0xffff))[0]
-def _i16(f):
-  if math.isnan(f): return 0x7e00
-  if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
-  try: return _struct_H.unpack(_struct_e.pack(f))[0]
-  except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
-def _f64(i): return _struct_d.unpack(_struct_Q.pack(i & MASK64))[0]
-def _i64(f):
-  if math.isnan(f): return 0x7ff8000000000000
-  if math.isinf(f): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
-  try: return _struct_Q.unpack(_struct_d.pack(f))[0]
-  except (OverflowError, struct.error): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# INTERNAL HELPERS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def _div(a, b):
-  try: return a / b
-  except ZeroDivisionError:
-    if a == 0.0 or math.isnan(a): return float("nan")
-    return math.copysign(float("inf"), a * b) if b == 0.0 else float("inf") if a > 0 else float("-inf")
-def _check_nan_type(x, quiet_bit_expected, default):
-  try:
-    if not math.isnan(float(x)): return False
-    if hasattr(x, '_reg') and hasattr(x, '_bits'):
-      bits = x._reg._val & ((1 << x._bits) - 1)
-      exp_bits, quiet_pos, mant_mask = {16: (0x1f, 9, 0x3ff), 32: (0xff, 22, 0x7fffff), 64: (0x7ff, 51, 0xfffffffffffff)}.get(x._bits, (0,0,0))
-      exp_shift = {16: 10, 32: 23, 64: 52}.get(x._bits, 0)
-      if exp_bits and ((bits >> exp_shift) & exp_bits) == exp_bits and (bits & mant_mask) != 0:
-        return ((bits >> quiet_pos) & 1) == quiet_bit_expected
-    return default
-  except (TypeError, ValueError): return False
-def _gt_neg_zero(a, b): return (a > b) or (a == 0 and b == 0 and not math.copysign(1, a) < 0 and math.copysign(1, b) < 0)
-def _lt_neg_zero(a, b): return (a < b) or (a == 0 and b == 0 and math.copysign(1, a) < 0 and not math.copysign(1, b) < 0)
-def _fpop(fn):
-  def wrapper(x):
-    x = float(x)
-    if math.isnan(x) or math.isinf(x): return x
-    result = float(fn(x))
-    return math.copysign(0.0, x) if result == 0.0 else result
-  return wrapper
-def _f_to_int(f, lo, hi): f = float(f); return 0 if math.isnan(f) else (hi if f >= hi else lo if f <= lo else int(f))
-def _f16_to_f32_bits(bits): return struct.unpack("<e", struct.pack("<H", int(bits) & 0xffff))[0]
-def _brev(v, bits): return int(bin(v & ((1 << bits) - 1))[2:].zfill(bits)[::-1], 2)
-def _ctz(v, bits):
-  v, n = int(v) & ((1 << bits) - 1), 0
-  if v == 0: return bits
-  while (v & 1) == 0: v >>= 1; n += 1
-  return n
-
-def _bf16(i):
-  """Convert bf16 bits to float. BF16 is just the top 16 bits of f32."""
-  return struct.unpack("<f", struct.pack("<I", (i & 0xffff) << 16))[0]
-def _ibf16(f):
-  """Convert float to bf16 bits (truncate to top 16 bits of f32)."""
-  if math.isnan(f): return 0x7fc0  # bf16 quiet NaN
-  if math.isinf(f): return 0x7f80 if f > 0 else 0xff80  # bf16 ±infinity
-  try: return (struct.unpack("<I", struct.pack("<f", float(f)))[0] >> 16) & 0xffff
-  except (OverflowError, struct.error): return 0x7f80 if f > 0 else 0xff80
-def _trig(fn, x):
-  # V_SIN/COS_F32: hardware does frac on input cycles before computing
-  if math.isinf(x) or math.isnan(x): return float("nan")
-  frac_cycles = fract(x / (2 * math.pi))
-  result = fn(frac_cycles * 2 * math.pi)
-  # Hardware returns exactly 0 for cos(π/2), sin(π), etc. due to lookup table
-  # Round very small results (below f32 precision) to exactly 0
-  if abs(result) < 1e-7: return 0.0
-  return result
-
-class _SafeFloat(float):
-  """Float subclass that uses _div for division to handle 0/inf correctly."""
-  def __truediv__(self, o): return _div(float(self), float(o))
-  def __rtruediv__(self, o): return _div(float(o), float(self))
-
-class _Inf:
-  f16 = f32 = f64 = float('inf')
-  def __neg__(self): return _NegInf()
-  def __pos__(self): return self
-  def __float__(self): return float('inf')
-  def __eq__(self, other): return float(other) == float('inf') if not isinstance(other, _NegInf) else False
-  def __req__(self, other): return self.__eq__(other)
-class _NegInf:
-  f16 = f32 = f64 = float('-inf')
-  def __neg__(self): return _Inf()
-  def __pos__(self): return self
-  def __float__(self): return float('-inf')
-  def __eq__(self, other): return float(other) == float('-inf') if not isinstance(other, _Inf) else False
-  def __req__(self, other): return self.__eq__(other)
-
-class _RoundMode:
-  NEAREST_EVEN = 0
-
-class _WaveMode:
-  IEEE = False
-
-class _DenormChecker:
-  """Comparator for denormalized floats. x == DENORM.f32 checks if x is denormalized."""
-  def __init__(self, bits): self._bits = bits
-  def _check(self, other):
-    f = float(other)
-    if math.isinf(f) or math.isnan(f) or f == 0.0: return False
-    if self._bits == 64:
-      bits = struct.unpack("<Q", struct.pack("<d", f))[0]
-      return (bits >> 52) & 0x7ff == 0
-    bits = struct.unpack("<I", struct.pack("<f", f))[0]
-    return (bits >> 23) & 0xff == 0
-  def __eq__(self, other): return self._check(other)
-  def __req__(self, other): return self._check(other)
-  def __ne__(self, other): return not self._check(other)
-
-class _Denorm:
-  f32 = _DenormChecker(32)
-  f64 = _DenormChecker(64)
-
-_pack = lambda hi, lo: ((int(hi) & 0xffff) << 16) | (int(lo) & 0xffff)
-_pack32 = lambda hi, lo: ((int(hi) & 0xffffffff) << 32) | (int(lo) & 0xffffffff)
-
-class TypedView:
-  """View into a Reg with typed access. Used for both full-width (Reg.u32) and slices (Reg[31:16])."""
-  __slots__ = ('_reg', '_high', '_low', '_signed', '_float', '_bf16', '_reversed')
-  def __init__(self, reg, high, low=0, signed=False, is_float=False, is_bf16=False):
-    # Handle reversed slices like [0:31] which means bit-reverse
-    if high < low: high, low, reversed = low, high, True
-    else: reversed = False
-    self._reg, self._high, self._low, self._reversed = reg, high, low, reversed
-    self._signed, self._float, self._bf16 = signed, is_float, is_bf16
-
-  def _nbits(self): return self._high - self._low + 1
-  def _mask(self): return (1 << self._nbits()) - 1
-  def _get(self):
-    v = (self._reg._val >> self._low) & self._mask()
-    return _brev(v, self._nbits()) if self._reversed else v
-  def _set(self, v):
-    v = int(v)
-    if self._reversed: v = _brev(v, self._nbits())
-    self._reg._val = (self._reg._val & ~(self._mask() << self._low)) | ((v & self._mask()) << self._low)
-
-  @property
-  def _val(self): return self._get()
-  @property
-  def _bits(self): return self._nbits()
-
-  # Type accessors for slices (e.g., D0[31:16].f16)
-  u8 = property(lambda s: s._get() & 0xff)
-  u16 = property(lambda s: s._get() & 0xffff, lambda s, v: s._set(v))
-  u32 = property(lambda s: s._get() & MASK32, lambda s, v: s._set(v))
-  i16 = property(lambda s: _sext(s._get() & 0xffff, 16), lambda s, v: s._set(v))
-  i32 = property(lambda s: _sext(s._get() & MASK32, 32), lambda s, v: s._set(v))
-  f16 = property(lambda s: _f16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _i16(float(v))))
-  f32 = property(lambda s: _f32(s._get()), lambda s, v: s._set(_i32(float(v))))
-  bf16 = property(lambda s: _bf16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _ibf16(float(v))))
-  b16, b32 = u16, u32
-
-  # Chained type access (e.g., jump_addr.i64 when jump_addr is already TypedView)
-  @property
-  def i64(s): return s if s._nbits() == 64 and s._signed else int(s)
-  @property
-  def u64(s): return s if s._nbits() == 64 and not s._signed else int(s) & MASK64
-
-  def __getitem__(self, key):
-    if isinstance(key, slice):
-      high, low = int(key.start), int(key.stop)
-      return TypedView(self._reg, high, low)
-    return (self._get() >> int(key)) & 1
-
-  def __setitem__(self, key, value):
-    if isinstance(key, slice):
-      high, low = int(key.start), int(key.stop)
-      if high < low: high, low, value = low, high, _brev(int(value), low - high + 1)
-      mask = (1 << (high - low + 1)) - 1
-      self._reg._val = (self._reg._val & ~(mask << low)) | ((int(value) & mask) << low)
-    elif value: self._reg._val |= (1 << int(key))
-    else: self._reg._val &= ~(1 << int(key))
-
-  def __int__(self): return _sext(self._get(), self._nbits()) if self._signed else self._get()
-  def __index__(self): return int(self)
-  def __trunc__(self): return int(float(self)) if self._float else int(self)
-  def __float__(self):
-    if self._float:
-      if self._bf16: return _bf16(self._get())
-      bits = self._nbits()
-      return _f16(self._get()) if bits == 16 else _f32(self._get()) if bits == 32 else _f64(self._get())
-    return float(int(self))
-  def __bool__(s): return bool(int(s))
-
-  # Arithmetic - floats use float(), ints use int()
-  def __add__(s, o): return float(s) + float(o) if s._float else int(s) + int(o)
-  def __radd__(s, o): return float(o) + float(s) if s._float else int(o) + int(s)
-  def __sub__(s, o): return float(s) - float(o) if s._float else int(s) - int(o)
-  def __rsub__(s, o): return float(o) - float(s) if s._float else int(o) - int(s)
-  def __mul__(s, o): return float(s) * float(o) if s._float else int(s) * int(o)
-  def __rmul__(s, o): return float(o) * float(s) if s._float else int(o) * int(s)
-  def __truediv__(s, o): return _div(float(s), float(o)) if s._float else _div(int(s), int(o))
-  def __rtruediv__(s, o): return _div(float(o), float(s)) if s._float else _div(int(o), int(s))
-  def __pow__(s, o): return float(s) ** float(o) if s._float else int(s) ** int(o)
-  def __rpow__(s, o): return float(o) ** float(s) if s._float else int(o) ** int(s)
-  def __neg__(s): return -float(s) if s._float else -int(s)
-  def __abs__(s): return abs(float(s)) if s._float else abs(int(s))
-
-  # Bitwise - GPU shifts mask the shift amount to valid range
-  def __and__(s, o): return int(s) & int(o)
-  def __or__(s, o): return int(s) | int(o)
-  def __xor__(s, o): return int(s) ^ int(o)
-  def __invert__(s): return ~int(s)
-  def __lshift__(s, o): n = int(o); return int(s) << n if 0 <= n < 64 or s._nbits() > 64 else 0
-  def __rshift__(s, o): n = int(o); return int(s) >> n if 0 <= n < 64 or s._nbits() > 64 else 0
-  def __rand__(s, o): return int(o) & int(s)
-  def __ror__(s, o): return int(o) | int(s)
-  def __rxor__(s, o): return int(o) ^ int(s)
-  def __rlshift__(s, o): n = int(s); return int(o) << n if 0 <= n < 64 else 0
-  def __rrshift__(s, o): n = int(s); return int(o) >> n if 0 <= n < 64 else 0
-
-  # Comparison - handle _DenormChecker specially
-  def __eq__(s, o):
-    if isinstance(o, _DenormChecker): return o._check(s)
-    return float(s) == float(o) if s._float else int(s) == int(o)
-  def __ne__(s, o):
-    if isinstance(o, _DenormChecker): return not o._check(s)
-    return float(s) != float(o) if s._float else int(s) != int(o)
-  def __lt__(s, o): return float(s) < float(o) if s._float else int(s) < int(o)
-  def __le__(s, o): return float(s) <= float(o) if s._float else int(s) <= int(o)
-  def __gt__(s, o): return float(s) > float(o) if s._float else int(s) > int(o)
-  def __ge__(s, o): return float(s) >= float(o) if s._float else int(s) >= int(o)
-
-class Reg:
-  """GPU register: D0.f32 = S0.f32 + S1.f32 just works. Supports up to 128 bits for DS_LOAD_B128."""
-  __slots__ = ('_val',)
-  def __init__(self, val=0): self._val = int(val)
-
-  # Typed views - TypedView(reg, high, signed, is_float, is_bf16)
-  u64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
-  i64 = property(lambda s: TypedView(s, 63, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK64))
-  b64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
-  f64 = property(lambda s: TypedView(s, 63, is_float=True), lambda s, v: setattr(s, '_val', v if isinstance(v, int) else _i64(float(v))))
-  u32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
-  i32 = property(lambda s: TypedView(s, 31, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK32))
-  b32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
-  f32 = property(lambda s: TypedView(s, 31, is_float=True), lambda s, v: setattr(s, '_val', _i32(float(v))))
-  u24 = property(lambda s: TypedView(s, 23))
-  i24 = property(lambda s: TypedView(s, 23, signed=True))
-  u16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
-  i16 = property(lambda s: TypedView(s, 15, signed=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
-  b16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
-  f16 = property(lambda s: TypedView(s, 15, is_float=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _i16(float(v))) & 0xffff)))
-  bf16 = property(lambda s: TypedView(s, 15, is_float=True, is_bf16=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _ibf16(float(v))) & 0xffff)))
-  u8 = property(lambda s: TypedView(s, 7))
-  i8 = property(lambda s: TypedView(s, 7, signed=True))
-  u3 = property(lambda s: TypedView(s, 2))  # 3-bit for opsel fields
-  u1 = property(lambda s: TypedView(s, 0))  # single bit
-
-  def __getitem__(s, key):
-    if isinstance(key, slice): return TypedView(s, int(key.start), int(key.stop))
-    return (s._val >> int(key)) & 1
-
-  def __setitem__(s, key, value):
-    if isinstance(key, slice):
-      high, low = int(key.start), int(key.stop)
-      if high < low: high, low = low, high
-      mask = (1 << (high - low + 1)) - 1
-      s._val = (s._val & ~(mask << low)) | ((int(value) & mask) << low)
-    elif value: s._val |= (1 << int(key))
-    else: s._val &= ~(1 << int(key))
-
-  def __int__(s): return s._val
-  def __index__(s): return s._val
-  def __bool__(s): return bool(s._val)
-
-  # Arithmetic (for tmp = tmp + 1 patterns). Float operands trigger f32 interpretation.
-  def __add__(s, o): return (_f32(s._val) + float(o)) if isinstance(o, float) else s._val + int(o)
-  def __radd__(s, o): return (float(o) + _f32(s._val)) if isinstance(o, float) else int(o) + s._val
-  def __sub__(s, o): return (_f32(s._val) - float(o)) if isinstance(o, float) else s._val - int(o)
-  def __rsub__(s, o): return (float(o) - _f32(s._val)) if isinstance(o, float) else int(o) - s._val
-  def __mul__(s, o): return (_f32(s._val) * float(o)) if isinstance(o, float) else s._val * int(o)
-  def __rmul__(s, o): return (float(o) * _f32(s._val)) if isinstance(o, float) else int(o) * s._val
-  def __and__(s, o): return s._val & int(o)
-  def __rand__(s, o): return int(o) & s._val
-  def __or__(s, o): return s._val | int(o)
-  def __ror__(s, o): return int(o) | s._val
-  def __xor__(s, o): return s._val ^ int(o)
-  def __rxor__(s, o): return int(o) ^ s._val
-  def __lshift__(s, o): n = int(o); return s._val << n if 0 <= n < 64 else 0
-  def __rshift__(s, o): n = int(o); return s._val >> n if 0 <= n < 64 else 0
-  def __invert__(s): return ~s._val
-
-  # Comparison (for tmp >= 0x100000000 patterns)
-  def __lt__(s, o): return s._val < int(o)
-  def __le__(s, o): return s._val <= int(o)
-  def __gt__(s, o): return s._val > int(o)
-  def __ge__(s, o): return s._val >= int(o)
-  def __eq__(s, o): return s._val == int(o)
-  def __ne__(s, o): return s._val != int(o)
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# PSEUDOCODE API - Functions and constants from AMD ISA pseudocode
-# ═══════════════════════════════════════════════════════════════════════════════
-
-# Rounding and float operations
-trunc, floor, ceil = _fpop(math.trunc), _fpop(math.floor), _fpop(math.ceil)
-def sqrt(x): return _SafeFloat(math.sqrt(x)) if x >= 0 else _SafeFloat(float("nan"))
-def log2(x): return math.log2(x) if x > 0 else (float("-inf") if x == 0 else float("nan"))
-def fract(x): return x - math.floor(x)
-def sin(x): return _trig(math.sin, x)
-def cos(x): return _trig(math.cos, x)
-def pow(a, b):
-  try: return a ** b
-  except OverflowError: return float("inf") if b > 0 else 0.0
-def isEven(x):
-  x = float(x)
-  if math.isinf(x) or math.isnan(x): return False
-  return int(x) % 2 == 0
-def mantissa(f):
-  if f == 0.0 or math.isinf(f) or math.isnan(f): return f
-  m, _ = math.frexp(f)
-  return m  # AMD V_FREXP_MANT returns mantissa in [0.5, 1.0) range
-def signext_from_bit(val, bit):
-  bit = int(bit)
-  if bit == 0: return 0
-  mask = (1 << bit) - 1
-  val = int(val) & mask
-  if val & (1 << (bit - 1)): return val - (1 << bit)
-  return val
-
-# Type conversions
-i32_to_f32 = u32_to_f32 = i32_to_f64 = u32_to_f64 = f32_to_f64 = f64_to_f32 = float
-def f32_to_i32(f): return _f_to_int(f, -2147483648, 2147483647)
-def f32_to_u32(f): return _f_to_int(f, 0, 4294967295)
-f64_to_i32, f64_to_u32 = f32_to_i32, f32_to_u32
-def f32_to_f16(f):
-  f = float(f)
-  if math.isnan(f): return 0x7e00  # f16 NaN
-  if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00  # f16 ±infinity
-  try: return struct.unpack("<H", struct.pack("<e", f))[0]
-  except OverflowError: return 0x7c00 if f > 0 else 0xfc00  # overflow -> ±infinity
-def f16_to_f32(v): return v if isinstance(v, float) else _f16_to_f32_bits(v)
-def i16_to_f16(v): return f32_to_f16(float(_sext(int(v) & 0xffff, 16)))
-def u16_to_f16(v): return f32_to_f16(float(int(v) & 0xffff))
-def f16_to_i16(bits): f = _f16_to_f32_bits(bits); return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
-def f16_to_u16(bits): f = _f16_to_f32_bits(bits); return max(0, min(65535, int(f))) if not math.isnan(f) else 0
-def bf16_to_f32(v): return _bf16(v) if isinstance(v, int) else float(v)
-def f32_to_bf16(f): return _ibf16(f)
-def u8_to_u32(v): return int(v) & 0xff
-def u4_to_u32(v): return int(v) & 0xf
-def u32_to_u16(u): return int(u) & 0xffff
-def i32_to_i16(i): return ((int(i) + 32768) & 0xffff) - 32768
-def f16_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
-def f16_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
-def f32_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
-def f32_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
-def v_cvt_i16_f32(f): return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
-def v_cvt_u16_f32(f): return max(0, min(65535, int(f))) if not math.isnan(f) else 0
-def SAT8(v): return max(0, min(255, int(v)))
-def f32_to_u8(f): return max(0, min(255, int(f))) if not math.isnan(f) else 0
-
-# Min/max operations
-def v_min_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _lt_neg_zero(a, b) else b)
-def v_max_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _gt_neg_zero(a, b) else b)
-v_min_f16, v_max_f16 = v_min_f32, v_max_f32
-v_min_i32, v_max_i32 = min, max
-v_min_i16, v_max_i16 = min, max
-def v_min_u32(a, b): return min(a & MASK32, b & MASK32)
-def v_max_u32(a, b): return max(a & MASK32, b & MASK32)
-def v_min_u16(a, b): return min(a & 0xffff, b & 0xffff)
-def v_max_u16(a, b): return max(a & 0xffff, b & 0xffff)
-def v_min3_f32(a, b, c): return v_min_f32(v_min_f32(a, b), c)
-def v_max3_f32(a, b, c): return v_max_f32(v_max_f32(a, b), c)
-v_min3_f16, v_max3_f16 = v_min3_f32, v_max3_f32
-v_min3_i32, v_max3_i32, v_min3_i16, v_max3_i16 = min, max, min, max
-def v_min3_u32(a, b, c): return min(a & MASK32, b & MASK32, c & MASK32)
-def v_max3_u32(a, b, c): return max(a & MASK32, b & MASK32, c & MASK32)
-def v_min3_u16(a, b, c): return min(a & 0xffff, b & 0xffff, c & 0xffff)
-def v_max3_u16(a, b, c): return max(a & 0xffff, b & 0xffff, c & 0xffff)
-
-# SAD/MSAD operations
-def ABSDIFF(a, b): return abs(int(a) - int(b))
-def v_sad_u8(s0, s1, s2):
-  """V_SAD_U8: Sum of absolute differences of 4 byte pairs plus accumulator."""
-  s0, s1, s2 = int(s0), int(s1), int(s2)
-  result = s2
-  for i in range(4):
-    a = (s0 >> (i * 8)) & 0xff
-    b = (s1 >> (i * 8)) & 0xff
-    result += abs(a - b)
-  return result & 0xffffffff
-def v_msad_u8(s0, s1, s2):
-  """V_MSAD_U8: Masked sum of absolute differences (skip if reference byte is 0)."""
-  s0, s1, s2 = int(s0), int(s1), int(s2)
-  result = s2
-  for i in range(4):
-    a = (s0 >> (i * 8)) & 0xff
-    b = (s1 >> (i * 8)) & 0xff
-    if b != 0:  # Only add diff if reference (s1) byte is non-zero
-      result += abs(a - b)
-  return result & 0xffffffff
-
-def BYTE_PERMUTE(data, sel):
-  """Select a byte from 64-bit data based on selector value."""
-  sel = int(sel) & 0xff
-  if sel <= 7: return (int(data) >> (sel * 8)) & 0xff
-  if sel == 8: return 0xff if ((int(data) >> 15) & 1) else 0x00
-  if sel == 9: return 0xff if ((int(data) >> 31) & 1) else 0x00
-  if sel == 10: return 0xff if ((int(data) >> 47) & 1) else 0x00
-  if sel == 11: return 0xff if ((int(data) >> 63) & 1) else 0x00
-  if sel == 12: return 0x00
-  return 0xff
-
-# Pseudocode functions
-def s_ff1_i32_b32(v): return _ctz(v, 32)
-def s_ff1_i32_b64(v): return _ctz(v, 64)
-GT_NEG_ZERO, LT_NEG_ZERO = _gt_neg_zero, _lt_neg_zero
-def isNAN(x):
-  try: return math.isnan(float(x))
-  except (TypeError, ValueError): return False
-def isQuietNAN(x): return _check_nan_type(x, 1, True)
-def isSignalNAN(x): return _check_nan_type(x, 0, False)
-def fma(a, b, c):
-  try: return math.fma(a, b, c)
-  except ValueError: return float('nan')
-def ldexp(m, e): return math.ldexp(m, e)
-def sign(f): return 1 if math.copysign(1.0, f) < 0 else 0
-def exponent(f):
-  if hasattr(f, '_bits') and hasattr(f, '_float') and f._float:
-    raw = f._val
-    if f._bits == 16: return (raw >> 10) & 0x1f
-    if f._bits == 32: return (raw >> 23) & 0xff
-    if f._bits == 64: return (raw >> 52) & 0x7ff
-  f = float(f)
-  if math.isinf(f) or math.isnan(f): return 255
-  if f == 0.0: return 0
-  try: bits = struct.unpack("<I", struct.pack("<f", f))[0]; return (bits >> 23) & 0xff
-  except: return 0
-def signext(x): return int(x)
-def cvtToQuietNAN(x): return float('nan')
-
-def F(x):
-  """32'F(x) or 64'F(x) - interpret x as float. If x is int, treat as bit pattern."""
-  if isinstance(x, int): return _f32(x)
-  if isinstance(x, TypedView): return x
-  return float(x)
-
-# Constants
-PI = math.pi
-WAVE32, WAVE64 = True, False
-OVERFLOW_F32, UNDERFLOW_F32 = float('inf'), 0.0
-OVERFLOW_F64, UNDERFLOW_F64 = float('inf'), 0.0
-MAX_FLOAT_F32 = 3.4028235e+38
-INF = _Inf()
-ROUND_MODE = _RoundMode()
-WAVE_MODE = _WaveMode()
-DENORM = _Denorm()
-
-# 2/PI with 1201 bits of precision for V_TRIG_PREOP_F64
-TWO_OVER_PI_1201 = Reg(0x0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6)
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# COMPILER: pseudocode -> Python (minimal transforms)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def _filter_pseudocode(pseudocode: str) -> str:
-  """Filter raw PDF pseudocode to only include actual code lines."""
-  pcode_lines, in_lambda, depth = [], 0, 0
-  for line in pseudocode.split('\n'):
-    s = line.strip()
-    if not s: continue
-    if '=>' in s or re.match(r'^[A-Z_]+\(', s): continue  # Skip example lines
-    if '= lambda(' in s: in_lambda += 1; continue  # Skip lambda definitions
-    if in_lambda > 0:
-      if s.endswith(');'): in_lambda -= 1
-      continue
-    # Only include lines that look like pseudocode
-    is_code = (any(p in s for p in ['D0.', 'D1.', 'S0.', 'S1.', 'S2.', 'SCC =', 'SCC ?', 'VCC', 'EXEC', 'tmp =', 'tmp[', 'lane =', 'PC =',
-                                    'D0[', 'D1[', 'S0[', 'S1[', 'S2[', 'MEM[', 'RETURN_DATA', 'VADDR', 'VDATA', 'VDST', 'SADDR', 'OFFSET']) or
-               s.startswith(('if ', 'else', 'elsif', 'endif', 'declare ', 'for ', 'endfor', '//')) or
-               re.match(r'^[a-z_]+\s*=', s) or re.match(r'^[a-z_]+\[', s) or (depth > 0 and '=' in s))
-    if s.startswith('if '): depth += 1
-    elif s.startswith('endif'): depth = max(0, depth - 1)
-    if is_code: pcode_lines.append(s)
-  return '\n'.join(pcode_lines)
-
-def _compile_pseudocode(pseudocode: str) -> str:
-  """Compile pseudocode to Python. Transforms are minimal - most syntax just works."""
-  pseudocode = re.sub(r'\bpass\b', 'pass_', pseudocode)  # 'pass' is Python keyword
-  raw_lines = pseudocode.strip().split('\n')
-  joined_lines: list[str] = []
-  for line in raw_lines:
-    line = line.strip()
-    if joined_lines and (joined_lines[-1].rstrip().endswith(('||', '&&', '(', ',')) or
-                         (joined_lines[-1].count('(') > joined_lines[-1].count(')'))):
-      joined_lines[-1] = joined_lines[-1].rstrip() + ' ' + line
-    else:
-      joined_lines.append(line)
-
-  lines = []
-  indent, need_pass, in_first_match_loop = 0, False, False
-  for line in joined_lines:
-    line = line.split('//')[0].strip()  # Strip C-style comments
-    if not line: continue
-    if line.startswith('if '):
-      lines.append('  ' * indent + f"if {_expr(line[3:].rstrip(' then'))}:")
-      indent += 1
-      need_pass = True
-    elif line.startswith('elsif '):
-      if need_pass: lines.append('  ' * indent + "pass")
-      indent -= 1
-      lines.append('  ' * indent + f"elif {_expr(line[6:].rstrip(' then'))}:")
-      indent += 1
-      need_pass = True
-    elif line == 'else':
-      if need_pass: lines.append('  ' * indent + "pass")
-      indent -= 1
-      lines.append('  ' * indent + "else:")
-      indent += 1
-      need_pass = True
-    elif line.startswith('endif'):
-      if need_pass: lines.append('  ' * indent + "pass")
-      indent -= 1
-      need_pass = False
-    elif line.startswith('endfor'):
-      if need_pass: lines.append('  ' * indent + "pass")
-      indent -= 1
-      need_pass, in_first_match_loop = False, False
-    elif line.startswith('declare '):
-      pass
-    elif m := re.match(r'for (\w+) in (.+?)\s*:\s*(.+?) do', line):
-      start, end = _expr(m[2].strip()), _expr(m[3].strip())
-      lines.append('  ' * indent + f"for {m[1]} in range({start}, int({end})+1):")
-      indent += 1
-      need_pass, in_first_match_loop = True, True
-    elif '=' in line and not line.startswith('=='):
-      need_pass = False
-      line = line.rstrip(';')
-      if m := re.match(r'\{\s*D1\.[ui]1\s*,\s*D0\.[ui]64\s*\}\s*=\s*(.+)', line):
-        rhs = _expr(m[1])
-        lines.append('  ' * indent + f"_full = {rhs}")
-        lines.append('  ' * indent + f"D0.u64 = int(_full) & 0xffffffffffffffff")
-        lines.append('  ' * indent + f"D1 = Reg((int(_full) >> 64) & 1)")
-      elif any(op in line for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^=')):
-        for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^='):
-          if op in line:
-            lhs, rhs = line.split(op, 1)
-            lines.append('  ' * indent + f"{lhs.strip()} {op} {_expr(rhs.strip())}")
-            break
-      else:
-        lhs, rhs = line.split('=', 1)
-        lhs_s, rhs_s = _expr(lhs.strip()), rhs.strip()
-        stmt = _assign(lhs_s, _expr(rhs_s))
-        if in_first_match_loop and rhs_s == 'i' and (lhs_s == 'tmp' or lhs_s == 'D0.i32'):
-          stmt += "; break"
-        lines.append('  ' * indent + stmt)
-  if need_pass: lines.append('  ' * indent + "pass")
-  return '\n'.join(lines)
-
-def _assign(lhs: str, rhs: str) -> str:
-  if lhs in ('tmp', 'SCC', 'VCC', 'EXEC', 'D0', 'D1', 'saveexec', 'PC'):
-    return f"{lhs} = Reg({rhs})"
-  return f"{lhs} = {rhs}"
-
-def _expr(e: str) -> str:
-  e = e.strip()
-  e = e.replace('&&', ' and ').replace('||', ' or ').replace('<>', ' != ')
-  e = re.sub(r'!([^=])', r' not \1', e)
-  e = re.sub(r'\{\s*(\w+\.u32)\s*,\s*(\w+\.u32)\s*\}', r'_pack32(\1, \2)', e)
-  def pack(m):
-    hi, lo = _expr(m[1].strip()), _expr(m[2].strip())
-    return f'_pack({hi}, {lo})'
-  e = re.sub(r'\{\s*([^,{}]+)\s*,\s*([^,{}]+)\s*\}', pack, e)
-  e = re.sub(r"1201'B\(2\.0\s*/\s*PI\)", "TWO_OVER_PI_1201", e)
-  e = re.sub(r"\d+'([0-9a-fA-Fx]+)[UuFf]*", r'\1', e)
-  e = re.sub(r"\d+'[FIBU]\(", "(", e)
-  e = re.sub(r'\bB\(', '(', e)
-  e = re.sub(r'([0-9a-fA-Fx])ULL\b', r'\1', e)
-  e = re.sub(r'([0-9a-fA-Fx])LL\b', r'\1', e)
-  e = re.sub(r'([0-9a-fA-Fx])U\b', r'\1', e)
-  e = re.sub(r'(\d\.?\d*)F\b', r'\1', e)
-  e = re.sub(r'(\[laneId\])\.[uib]\d+', r'\1', e)
-  e = e.replace('+INF', 'INF').replace('-INF', '(-INF)')
-  e = re.sub(r'NAN\.f\d+', 'float("nan")', e)
-  def convert_verilog_slice(m):
-    start, width = m.group(1).strip(), m.group(2).strip()
-    return f'[({start}) + ({width}) - 1 : ({start})]'
-  e = re.sub(r'\[([^:\[\]]+)\s*\+:\s*([^:\[\]]+)\]', convert_verilog_slice, e)
-  def process_brackets(s):
-    result, i = [], 0
-    while i < len(s):
-      if s[i] == '[':
-        depth, start = 1, i + 1
-        j = start
-        while j < len(s) and depth > 0:
-          if s[j] == '[': depth += 1
-          elif s[j] == ']': depth -= 1
-          j += 1
-        inner = _expr(s[start:j-1])
-        result.append('[' + inner + ']')
-        i = j
-      else:
-        result.append(s[i])
-        i += 1
-    return ''.join(result)
-  e = process_brackets(e)
-  while '?' in e:
-    depth, bracket, q = 0, 0, -1
-    for i, c in enumerate(e):
-      if c == '(': depth += 1
-      elif c == ')': depth -= 1
-      elif c == '[': bracket += 1
-      elif c == ']': bracket -= 1
-      elif c == '?' and depth == 0 and bracket == 0: q = i; break
-    if q < 0: break
-    depth, bracket, col = 0, 0, -1
-    for i in range(q + 1, len(e)):
-      if e[i] == '(': depth += 1
-      elif e[i] == ')': depth -= 1
-      elif e[i] == '[': bracket += 1
-      elif e[i] == ']': bracket -= 1
-      elif e[i] == ':' and depth == 0 and bracket == 0: col = i; break
-    if col < 0: break
-    cond, t, f = e[:q].strip(), e[q+1:col].strip(), e[col+1:].strip()
-    e = f'(({t}) if ({cond}) else ({f}))'
-  return e
-
-def _apply_pseudocode_fixes(op_name: str, code: str) -> str:
-  """Apply known fixes for PDF pseudocode bugs."""
-  if op_name == 'V_DIV_FMAS_F32':
-    code = code.replace('D0.f32 = 2.0 ** 32 * fma(S0.f32, S1.f32, S2.f32)',
-                        'D0.f32 = (2.0 ** 64 if exponent(S2.f32) > 127 else 2.0 ** -64) * fma(S0.f32, S1.f32, S2.f32)')
-  if op_name == 'V_DIV_FMAS_F64':
-    code = code.replace('D0.f64 = 2.0 ** 64 * fma(S0.f64, S1.f64, S2.f64)',
-                        'D0.f64 = (2.0 ** 128 if exponent(S2.f64) > 1023 else 2.0 ** -128) * fma(S0.f64, S1.f64, S2.f64)')
-  if op_name == 'V_DIV_SCALE_F32':
-    code = code.replace('D0.f32 = float("nan")', 'VCC = Reg(0x1); D0.f32 = float("nan")')
-    code = code.replace('elif S1.f32 == DENORM.f32:\n  D0.f32 = ldexp(S0.f32, 64)', 'elif False:\n  pass')
-    code += '\nif S1.f32 == DENORM.f32:\n  D0.f32 = float("nan")'
-    code = code.replace('elif exponent(S2.f32) <= 23:\n  D0.f32 = ldexp(S0.f32, 64)', 'elif exponent(S2.f32) <= 23:\n  VCC = Reg(0x1); D0.f32 = ldexp(S0.f32, 64)')
-    code = code.replace('elif S2.f32 / S1.f32 == DENORM.f32:\n  VCC = Reg(0x1)\n  if S0.f32 == S2.f32:\n    D0.f32 = ldexp(S0.f32, 64)', 'elif S2.f32 / S1.f32 == DENORM.f32:\n  VCC = Reg(0x1)')
-  if op_name == 'V_DIV_SCALE_F64':
-    code = code.replace('D0.f64 = float("nan")', 'VCC = Reg(0x1); D0.f64 = float("nan")')
-    code = code.replace('elif S1.f64 == DENORM.f64:\n  D0.f64 = ldexp(S0.f64, 128)', 'elif False:\n  pass')
-    code += '\nif S1.f64 == DENORM.f64:\n  D0.f64 = float("nan")'
-    code = code.replace('elif exponent(S2.f64) <= 52:\n  D0.f64 = ldexp(S0.f64, 128)', 'elif exponent(S2.f64) <= 52:\n  VCC = Reg(0x1); D0.f64 = ldexp(S0.f64, 128)')
-    code = code.replace('elif S2.f64 / S1.f64 == DENORM.f64:\n  VCC = Reg(0x1)\n  if S0.f64 == S2.f64:\n    D0.f64 = ldexp(S0.f64, 128)', 'elif S2.f64 / S1.f64 == DENORM.f64:\n  VCC = Reg(0x1)')
-  if op_name == 'V_DIV_FIXUP_F32':
-    code = code.replace('D0.f32 = ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))',
-                        'D0.f32 = ((-OVERFLOW_F32) if (sign_out) else (OVERFLOW_F32)) if isNAN(S0.f32) else ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))')
-  if op_name == 'V_DIV_FIXUP_F64':
-    code = code.replace('D0.f64 = ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))',
-                        'D0.f64 = ((-OVERFLOW_F64) if (sign_out) else (OVERFLOW_F64)) if isNAN(S0.f64) else ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))')
-  if op_name == 'V_TRIG_PREOP_F64':
-    code = code.replace('result = F((TWO_OVER_PI_1201[1200 : 0] << shift.u32) & 0x1fffffffffffff)',
-                        'result = float(((TWO_OVER_PI_1201[1200 : 0] << int(shift)) >> (1201 - 53)) & 0x1fffffffffffff)')
-  return code
-
-def _generate_function(cls_name: str, op_name: str, pc: str, code: str) -> str:
-  """Generate a single compiled pseudocode function.
-  Functions take int parameters and return dict of int values.
-  Reg wrapping happens inside the function, only for registers actually used."""
-  has_d1 = '{ D1' in pc
-  is_cmpx = (cls_name in ('VOPCOp', 'VOP3Op')) and 'EXEC.u64[laneId]' in pc
-  is_div_scale = 'DIV_SCALE' in op_name
-  has_sdst = cls_name == 'VOP3SDOp' and ('VCC.u64[laneId]' in pc or is_div_scale)
-  is_ds = cls_name == 'DSOp'
-  is_flat = cls_name in ('FLATOp', 'GLOBALOp', 'SCRATCHOp')
-  is_smem = cls_name == 'SMEMOp'
-  has_s_array = 'S[i]' in pc  # FMA_MIX style: S[0], S[1], S[2] array access
-  combined = code + pc
-
-  fn_name = f"_{cls_name}_{op_name}"
-
-  # Detect which registers are used/modified
-  def needs_init(name): return name in combined and not re.search(rf'^\s*{name}\s*=\s*Reg\(', code, re.MULTILINE)
-  modifies_d0 = is_div_scale or bool(re.search(r'\bD0\b[.\[]', combined))
-  modifies_exec = is_cmpx or bool(re.search(r'EXEC\.(u32|u64|b32|b64)\s*=', combined))
-  modifies_vcc = has_sdst or bool(re.search(r'VCC\.(u32|u64|b32|b64)\s*=|VCC\.u64\[laneId\]\s*=', combined))
-  modifies_scc = bool(re.search(r'\bSCC\s*=', combined))
-  modifies_pc = bool(re.search(r'\bPC\s*=', combined))
-
-  # Build function signature and Reg init lines
-  if is_smem:
-    lines = [f"def {fn_name}(MEM, addr):"]
-    reg_inits = ["ADDR=Reg(addr)", "SDATA=Reg(0)"]
-    special_regs = []
-  elif is_ds:
-    lines = [f"def {fn_name}(MEM, addr, data0, data1, offset0, offset1):"]
-    reg_inits = ["ADDR=Reg(addr)", "DATA0=Reg(data0)", "DATA1=Reg(data1)", "OFFSET0=Reg(offset0)", "OFFSET1=Reg(offset1)", "RETURN_DATA=Reg(0)"]
-    special_regs = [('DATA', 'DATA0'), ('DATA2', 'DATA1'), ('OFFSET', 'OFFSET0'), ('ADDR_BASE', 'ADDR')]
-  elif is_flat:
-    lines = [f"def {fn_name}(MEM, addr, vdata, vdst):"]
-    reg_inits = ["ADDR=addr", "VDATA=Reg(vdata)", "VDST=Reg(vdst)", "RETURN_DATA=Reg(0)"]
-    special_regs = [('DATA', 'VDATA')]
-  elif has_s_array:
-    # FMA_MIX style: needs S[i] array, opsel, opsel_hi for source selection (neg/neg_hi applied in emu.py before call)
-    lines = [f"def {fn_name}(s0, s1, s2, d0, scc, vcc, laneId, exec_mask, literal, VGPR, src0_idx=0, vdst_idx=0, pc=None, opsel=0, opsel_hi=0):"]
-    reg_inits = ["S0=Reg(s0)", "S1=Reg(s1)", "S2=Reg(s2)", "S=[S0,S1,S2]", "D0=Reg(d0)", "OPSEL=Reg(opsel)", "OPSEL_HI=Reg(opsel_hi)"]
-    special_regs = []
-    # Detect array declarations like "declare in : 32'F[3]" and create them (rename 'in' to 'ins' since 'in' is a keyword)
-    if "in[" in combined:
-      reg_inits.append("ins=[Reg(0),Reg(0),Reg(0)]")
-      code = code.replace("in[", "ins[")
-  else:
-    lines = [f"def {fn_name}(s0, s1, s2, d0, scc, vcc, laneId, exec_mask, literal, VGPR, src0_idx=0, vdst_idx=0, pc=None):"]
-    # Only create Regs for registers actually used in the pseudocode
-    reg_inits = []
-    if 'S0' in combined: reg_inits.append("S0=Reg(s0)")
-    if 'S1' in combined: reg_inits.append("S1=Reg(s1)")
-    if 'S2' in combined: reg_inits.append("S2=Reg(s2)")
-    if modifies_d0 or 'D0' in combined: reg_inits.append("D0=Reg(s0)" if is_div_scale else "D0=Reg(d0)")
-    if modifies_scc or 'SCC' in combined: reg_inits.append("SCC=Reg(scc)")
-    if modifies_vcc or 'VCC' in combined: reg_inits.append("VCC=Reg(vcc)")
-    if modifies_exec or 'EXEC' in combined: reg_inits.append("EXEC=Reg(exec_mask)")
-    if modifies_pc or 'PC' in combined: reg_inits.append("PC=Reg(pc) if pc is not None else None")
-    special_regs = [('D1', 'Reg(0)'), ('SIMM16', 'Reg(literal)'), ('SIMM32', 'Reg(literal)'),
-                    ('SRC0', 'Reg(src0_idx)'), ('VDST', 'Reg(vdst_idx)')]
-    if needs_init('tmp'): special_regs.insert(0, ('tmp', 'Reg(0)'))
-    if needs_init('saveexec'): special_regs.insert(0, ('saveexec', 'Reg(EXEC._val)'))
-
-  # Build init code
-  init_parts = reg_inits.copy()
-  for name, init in special_regs:
-    if name in combined: init_parts.append(f"{name}={init}")
-  if 'EXEC_LO' in code: init_parts.append("EXEC_LO=TypedView(EXEC, 31, 0)")
-  if 'EXEC_HI' in code: init_parts.append("EXEC_HI=TypedView(EXEC, 63, 32)")
-  if 'VCCZ' in code and not re.search(r'^\s*VCCZ\s*=', code, re.MULTILINE): init_parts.append("VCCZ=Reg(1 if VCC._val == 0 else 0)")
-  if 'EXECZ' in code and not re.search(r'^\s*EXECZ\s*=', code, re.MULTILINE): init_parts.append("EXECZ=Reg(1 if EXEC._val == 0 else 0)")
-
-  # Add init line and separator
-  if init_parts: lines.append(f"  {'; '.join(init_parts)}")
-
-  # Add compiled pseudocode
-  for line in code.split('\n'):
-    if line.strip(): lines.append(f"  {line}")
-
-  # Build result dict
-  result_items = []
-  if modifies_d0: result_items.append("'D0': D0._val")
-  if modifies_scc: result_items.append("'SCC': SCC._val")
-  if modifies_vcc: result_items.append("'VCC': VCC._val")
-  if modifies_exec: result_items.append("'EXEC': EXEC._val")
-  if has_d1: result_items.append("'D1': D1._val")
-  if modifies_pc: result_items.append("'PC': PC._val")
-  if is_smem and 'SDATA' in combined and re.search(r'^\s*SDATA[\.\[].*=', code, re.MULTILINE):
-    result_items.append("'SDATA': SDATA._val")
-  if is_ds and 'RETURN_DATA' in combined and re.search(r'^\s*RETURN_DATA[\.\[].*=', code, re.MULTILINE):
-    result_items.append("'RETURN_DATA': RETURN_DATA._val")
-  if is_flat:
-    if 'RETURN_DATA' in combined and re.search(r'^\s*RETURN_DATA[\.\[].*=', code, re.MULTILINE):
-      result_items.append("'RETURN_DATA': RETURN_DATA._val")
-    if re.search(r'^\s*VDATA[\.\[].*=', code, re.MULTILINE):
-      result_items.append("'VDATA': VDATA._val")
-  lines.append(f"  return {{{', '.join(result_items)}}}")
-  return '\n'.join(lines)
-
-# Build the globals dict for exec() - includes all pcode symbols
-_PCODE_GLOBALS = {
-  'Reg': Reg, 'TypedView': TypedView, '_pack': _pack, '_pack32': _pack32,
-  'ABSDIFF': ABSDIFF, 'BYTE_PERMUTE': BYTE_PERMUTE, 'DENORM': DENORM, 'F': F,
-  'GT_NEG_ZERO': GT_NEG_ZERO, 'LT_NEG_ZERO': LT_NEG_ZERO, 'INF': INF,
-  'MAX_FLOAT_F32': MAX_FLOAT_F32, 'OVERFLOW_F32': OVERFLOW_F32, 'OVERFLOW_F64': OVERFLOW_F64,
-  'UNDERFLOW_F32': UNDERFLOW_F32, 'UNDERFLOW_F64': UNDERFLOW_F64,
-  'PI': PI, 'ROUND_MODE': ROUND_MODE, 'WAVE_MODE': WAVE_MODE,
-  'WAVE32': WAVE32, 'WAVE64': WAVE64, 'TWO_OVER_PI_1201': TWO_OVER_PI_1201,
-  'SAT8': SAT8, 'trunc': trunc, 'floor': floor, 'ceil': ceil, 'sqrt': sqrt,
-  'log2': log2, 'fract': fract, 'sin': sin, 'cos': cos, 'pow': pow,
-  'isEven': isEven, 'mantissa': mantissa, 'signext_from_bit': signext_from_bit,
-  'i32_to_f32': i32_to_f32, 'u32_to_f32': u32_to_f32, 'i32_to_f64': i32_to_f64,
-  'u32_to_f64': u32_to_f64, 'f32_to_f64': f32_to_f64, 'f64_to_f32': f64_to_f32,
-  'f32_to_i32': f32_to_i32, 'f32_to_u32': f32_to_u32, 'f64_to_i32': f64_to_i32,
-  'f64_to_u32': f64_to_u32, 'f32_to_f16': f32_to_f16, 'f16_to_f32': f16_to_f32,
-  'i16_to_f16': i16_to_f16, 'u16_to_f16': u16_to_f16, 'f16_to_i16': f16_to_i16,
-  'f16_to_u16': f16_to_u16, 'bf16_to_f32': bf16_to_f32, 'f32_to_bf16': f32_to_bf16,
-  'u8_to_u32': u8_to_u32, 'u4_to_u32': u4_to_u32, 'u32_to_u16': u32_to_u16,
-  'i32_to_i16': i32_to_i16, 'f16_to_snorm': f16_to_snorm, 'f16_to_unorm': f16_to_unorm,
-  'f32_to_snorm': f32_to_snorm, 'f32_to_unorm': f32_to_unorm,
-  'v_cvt_i16_f32': v_cvt_i16_f32, 'v_cvt_u16_f32': v_cvt_u16_f32, 'f32_to_u8': f32_to_u8,
-  'v_min_f32': v_min_f32, 'v_max_f32': v_max_f32, 'v_min_f16': v_min_f16, 'v_max_f16': v_max_f16,
-  'v_min_i32': v_min_i32, 'v_max_i32': v_max_i32, 'v_min_i16': v_min_i16, 'v_max_i16': v_max_i16,
-  'v_min_u32': v_min_u32, 'v_max_u32': v_max_u32, 'v_min_u16': v_min_u16, 'v_max_u16': v_max_u16,
-  'v_min3_f32': v_min3_f32, 'v_max3_f32': v_max3_f32, 'v_min3_f16': v_min3_f16, 'v_max3_f16': v_max3_f16,
-  'v_min3_i32': v_min3_i32, 'v_max3_i32': v_max3_i32, 'v_min3_i16': v_min3_i16, 'v_max3_i16': v_max3_i16,
-  'v_min3_u32': v_min3_u32, 'v_max3_u32': v_max3_u32, 'v_min3_u16': v_min3_u16, 'v_max3_u16': v_max3_u16,
-  'v_sad_u8': v_sad_u8, 'v_msad_u8': v_msad_u8,
-  's_ff1_i32_b32': s_ff1_i32_b32, 's_ff1_i32_b64': s_ff1_i32_b64,
-  'isNAN': isNAN, 'isQuietNAN': isQuietNAN, 'isSignalNAN': isSignalNAN,
-  'fma': fma, 'ldexp': ldexp, 'sign': sign, 'exponent': exponent,
-  'signext': signext, 'cvtToQuietNAN': cvtToQuietNAN,
-}
-
-@functools.cache
-def compile_pseudocode(cls_name: str, op_name: str, pseudocode: str):
-  """Compile pseudocode string to executable function. Cached for performance."""
-  filtered = _filter_pseudocode(pseudocode)
-  code = _compile_pseudocode(filtered)
-  code = _apply_pseudocode_fixes(op_name, code)
-  fn_code = _generate_function(cls_name, op_name, filtered, code)
-  fn_name = f"_{cls_name}_{op_name}"
-  local_ns = {}
-  exec(fn_code, _PCODE_GLOBALS, local_ns)
-  return local_ns[fn_name]

extra/assembly/amd/sqtt.py

@@ -1,386 +0,0 @@
-"""SQTT (SQ Thread Trace) packet encoder and decoder for AMD GPUs.
-
-This module provides encoding and decoding of raw SQTT byte streams.
-The format is nibble-based with variable-width packets determined by a state machine.
-Uses BitField infrastructure from dsl.py, similar to GPU instruction encoding.
-"""
-from __future__ import annotations
-from typing import Iterator
-from enum import Enum
-from extra.assembly.amd.dsl import BitField, FixedBitField, bits
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# FIELD ENUMS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-class MemSrc(Enum):
-  LDS = 0
-  LDS_ALT = 1
-  VMEM = 2
-  VMEM_ALT = 3
-
-class AluSrc(Enum):
-  NONE = 0
-  SALU = 1
-  VALU = 2
-  VALU_SALU = 3
-
-class InstOp(Enum):
-  """SQTT instruction operation types.
-
-  Memory ops appear in two ranges depending on which SIMD executes them:
-  - 0x1x-0x2x range: ops on traced SIMD
-  - 0x5x range: ops on other SIMD (OTHER_ prefix)
-
-  GLOBAL memory ops encoding depends on addressing mode AND size:
-  - Loads: 0x21 (saddr=SGPR) or 0x22 (saddr=NULL), all sizes same
-  - Stores: base + size_offset, where VADDR is shifted +1 from SADDR
-    SADDR: 0x24(32) 0x25(64) 0x26(96) 0x27(128)
-    VADDR: 0x25(32) 0x26(64) 0x27(96) 0x28(128)
-
-  OTHER_ range follows same pattern but values overlap differently.
-  """
-  SALU = 0x0
-  SMEM = 0x1
-  JUMP = 0x3              # branch taken
-  JUMP_NO = 0x4           # branch not taken
-  MESSAGE = 0x9
-  VALU_TRANS = 0xb        # transcendental: exp, log, rcp, sqrt, sin, cos
-  VALU_64_SHIFT = 0xd     # 64-bit shifts: lshl, lshr, ashr
-  VALU_MAD64 = 0xe        # 64-bit multiply-add
-  VALU_64 = 0xf           # 64-bit: add, mul, fma, rcp, sqrt, rounding, frexp, div helpers
-  VINTERP = 0x12          # interpolation: v_interp_p10_f32, v_interp_p2_f32
-  BARRIER = 0x13
-
-  # FLAT memory ops on traced SIMD (0x1x range)
-  FLAT_LOAD = 0x1c
-  FLAT_STORE = 0x1d
-  FLAT_STORE_64 = 0x1e
-  FLAT_STORE_96 = 0x1f
-  FLAT_STORE_128 = 0x20
-
-  # GLOBAL memory ops on traced SIMD (0x2x range)
-  GLOBAL_LOAD = 0x21             # saddr=SGPR, all sizes
-  GLOBAL_LOAD_VADDR = 0x22       # saddr=NULL, all sizes
-  GLOBAL_STORE = 0x24            # saddr=SGPR, 32-bit
-  GLOBAL_STORE_64 = 0x25         # saddr=SGPR 64 or saddr=NULL 32
-  GLOBAL_STORE_96 = 0x26         # saddr=SGPR 96 or saddr=NULL 64
-  GLOBAL_STORE_128 = 0x27        # saddr=SGPR 128 or saddr=NULL 96
-  GLOBAL_STORE_VADDR_128 = 0x28  # saddr=NULL, 128-bit
-
-  # LDS ops on traced SIMD
-  LDS_LOAD = 0x29
-  LDS_STORE = 0x2b
-  LDS_STORE_64 = 0x2c
-  LDS_STORE_128 = 0x2e
-
-  # Memory ops on other SIMD (0x5x range)
-  OTHER_LDS_LOAD = 0x50
-  OTHER_LDS_STORE = 0x51
-  OTHER_LDS_STORE_64 = 0x52
-  OTHER_LDS_STORE_128 = 0x54
-  OTHER_FLAT_LOAD = 0x55
-  OTHER_FLAT_STORE = 0x56
-  OTHER_FLAT_STORE_64 = 0x57
-  OTHER_FLAT_STORE_96 = 0x58
-  OTHER_FLAT_STORE_128 = 0x59
-  OTHER_GLOBAL_LOAD = 0x5a             # saddr=SGPR, all sizes
-  OTHER_GLOBAL_LOAD_VADDR = 0x5b       # saddr=NULL or saddr=SGPR store 32
-  OTHER_GLOBAL_STORE_64 = 0x5c         # saddr=SGPR 64 or saddr=NULL 32
-  OTHER_GLOBAL_STORE_96 = 0x5d         # saddr=SGPR 96 or saddr=NULL 64
-  OTHER_GLOBAL_STORE_128 = 0x5e        # saddr=SGPR 128 or saddr=NULL 96
-  OTHER_GLOBAL_STORE_VADDR_128 = 0x5f  # saddr=NULL, 128-bit
-
-  # EXEC-modifying ops (0x7x range)
-  SALU_SAVEEXEC = 0x72    # s_*_saveexec_b32/b64
-  VALU_CMPX = 0x73        # v_cmpx_*
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# PACKET TYPE BASE CLASS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-class PacketType:
-  """Base class for SQTT packet types."""
-  encoding: FixedBitField
-  _raw: int
-  _time: int
-
-  def __init_subclass__(cls, **kwargs):
-    super().__init_subclass__(**kwargs)
-    cls._fields = {k: v for k, v in cls.__dict__.items() if isinstance(v, BitField)}
-    cls._size_nibbles = ((max((f.hi for f in cls._fields.values()), default=0) + 4) // 4)
-
-  @classmethod
-  def from_raw(cls, raw: int, time: int = 0):
-    inst = object.__new__(cls)
-    inst._raw, inst._time = raw, time
-    return inst
-
-  def __repr__(self) -> str:
-    fields_str = ", ".join(f"{k}={getattr(self, k)}" for k in self._fields if not k.startswith('_'))
-    return f"{self.__class__.__name__}({fields_str})"
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# PACKET TYPE DEFINITIONS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-class VALUINST(PacketType):  # exclude: 1 << 2
-  encoding = bits[2:0] == 0b011
-  delta = bits[5:3]
-  flag = bits[6:6]
-  wave = bits[11:7]
-
-class VMEMEXEC(PacketType):  # exclude: 1 << 0
-  encoding = bits[3:0] == 0b1111
-  delta = bits[5:4]
-  src = bits[7:6].enum(MemSrc)
-
-class ALUEXEC(PacketType):  # exclude: 1 << 1
-  encoding = bits[3:0] == 0b1110
-  delta = bits[5:4]
-  src = bits[7:6].enum(AluSrc)
-
-class IMMEDIATE(PacketType):  # exclude: 1 << 5
-  encoding = bits[3:0] == 0b1101
-  delta = bits[6:4]
-  wave = bits[11:7]
-
-class IMMEDIATE_MASK(PacketType):  # exclude: 1 << 5
-  encoding = bits[4:0] == 0b00100
-  delta = bits[7:5]
-  mask = bits[23:8]
-
-class WAVERDY(PacketType):  # exclude: 1 << 3
-  encoding = bits[4:0] == 0b10100
-  delta = bits[7:5]
-  mask = bits[23:8]
-
-class TS_DELTA_S8_W3(PacketType):
-  encoding = bits[6:0] == 0b0100001
-  delta = bits[10:8]
-  _padding = bits[63:11]
-
-class WAVEEND(PacketType):  # exclude: 1 << 4
-  encoding = bits[4:0] == 0b10101
-  delta = bits[7:5]
-  flag7 = bits[8:8]
-  simd = bits[10:9]
-  cu_lo = bits[13:11]
-  wave = bits[19:15]
-  @property
-  def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
-
-class WAVESTART(PacketType):  # exclude: 1 << 4
-  encoding = bits[4:0] == 0b01100
-  delta = bits[6:5]
-  flag7 = bits[7:7]
-  simd = bits[9:8]
-  cu_lo = bits[12:10]
-  wave = bits[17:13]
-  id7 = bits[31:18]
-  @property
-  def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
-
-class TS_DELTA_S5_W2(PacketType):
-  encoding = bits[4:0] == 0b11100
-  delta = bits[6:5]
-  _padding = bits[47:7]
-
-class WAVEALLOC(PacketType):  # exclude: 1 << 10
-  encoding = bits[4:0] == 0b00101
-  delta = bits[7:5]
-  _padding = bits[19:8]
-
-class TS_DELTA_S5_W3(PacketType):
-  encoding = bits[4:0] == 0b00110
-  delta = bits[7:5]
-  _padding = bits[51:8]
-
-class PERF(PacketType):  # exclude: 1 << 11
-  encoding = bits[4:0] == 0b10110
-  delta = bits[7:5]
-  arg = bits[27:8]
-
-class TS_DELTA_SHORT(PacketType):
-  encoding = bits[3:0] == 0b1000
-  delta = bits[7:4]
-
-class NOP(PacketType):
-  encoding = bits[3:0] == 0b0000
-  delta = None  # type: ignore
-  _padding = bits[3:0]
-
-class TS_WAVE_STATE(PacketType):
-  encoding = bits[6:0] == 0b1010001
-  delta = bits[15:7]
-  coarse = bits[23:16]
-  @property
-  def wave_interest(self) -> bool: return bool(self.coarse & 1)
-  @property
-  def terminate_all(self) -> bool: return bool(self.coarse & 8)
-
-class EVENT(PacketType):  # exclude: 1 << 7
-  encoding = bits[7:0] == 0b01100001
-  delta = bits[10:8]
-  event = bits[23:11]
-
-class EVENT_BIG(PacketType):
-  encoding = bits[7:0] == 0b11100001
-  delta = bits[10:8]
-  event = bits[31:11]
-
-class REG(PacketType):
-  encoding = bits[3:0] == 0b1001
-  delta = bits[6:4]
-  slot = bits[9:7]
-  hi_byte = bits[15:8]
-  subop = bits[31:16]
-  val32 = bits[63:32]
-  @property
-  def is_config(self) -> bool: return bool(self.hi_byte & 0x80)
-
-class SNAPSHOT(PacketType):
-  encoding = bits[6:0] == 0b1110001
-  delta = bits[9:7]
-  snap = bits[63:10]
-
-class TS_DELTA_OR_MARK(PacketType):
-  encoding = bits[6:0] == 0b0000001
-  delta = bits[47:12]
-  bit8 = bits[8:8]
-  bit9 = bits[9:9]
-  @property
-  def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
-
-class LAYOUT_HEADER(PacketType):
-  encoding = bits[6:0] == 0b0010001
-  delta = None  # type: ignore
-  layout = bits[12:7]
-  simd = bits[14:13]
-  group = bits[17:15]
-  sel_a = bits[31:28]
-  sel_b = bits[36:33]
-  flag4 = bits[59:59]
-  _padding = bits[63:60]
-
-class INST(PacketType):
-  encoding = bits[2:0] == 0b010
-  delta = bits[6:4]
-  flag1 = bits[3:3]
-  flag2 = bits[7:7]
-  wave = bits[12:8]
-  op = bits[19:13].enum(InstOp)
-
-class UTILCTR(PacketType):
-  encoding = bits[6:0] == 0b0110001
-  delta = bits[8:7]
-  ctr = bits[47:9]
-
-# All packet types in encoding priority order (more specific masks first, NOP last as fallback)
-PACKET_TYPES: list[type[PacketType]] = [
-  EVENT, EVENT_BIG,
-  TS_DELTA_S8_W3, TS_WAVE_STATE, SNAPSHOT, TS_DELTA_OR_MARK, LAYOUT_HEADER, UTILCTR,
-  IMMEDIATE_MASK, WAVERDY, WAVEEND, WAVESTART, TS_DELTA_S5_W2, WAVEALLOC, TS_DELTA_S5_W3, PERF,
-  VMEMEXEC, ALUEXEC, IMMEDIATE, TS_DELTA_SHORT, REG,
-  VALUINST, INST,
-  NOP,
-]
-
-def _build_state_table() -> tuple[bytes, dict[int, type[PacketType]]]:
-  table = [len(PACKET_TYPES) - 1] * 256  # default to NOP
-  opcode_to_class: dict[int, type[PacketType]] = {i: cls for i, cls in enumerate(PACKET_TYPES)}
-
-  for byte_val in range(256):
-    for opcode, pkt_cls in enumerate(PACKET_TYPES):
-      if (byte_val & pkt_cls.encoding.mask) == pkt_cls.encoding.default:
-        table[byte_val] = opcode
-        break
-
-  return bytes(table), opcode_to_class
-
-STATE_TO_OPCODE, OPCODE_TO_CLASS = _build_state_table()
-
-# Precompute special case opcodes
-_TS_DELTA_OR_MARK_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_OR_MARK)
-_TS_DELTA_SHORT_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_SHORT)
-
-# Combined lookup: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
-# special_case: 0=none, 1=TS_DELTA_OR_MARK, 2=TS_DELTA_SHORT
-_DECODE_INFO: dict[int, tuple] = {}
-for _opcode, _pkt_cls in OPCODE_TO_CLASS.items():
-  _delta_field = getattr(_pkt_cls, 'delta', None)
-  _delta_lo = _delta_field.lo if _delta_field else 0
-  _delta_mask = _delta_field.mask if _delta_field else 0
-  _special = 1 if _opcode == _TS_DELTA_OR_MARK_OPCODE else (2 if _opcode == _TS_DELTA_SHORT_OPCODE else 0)
-  _DECODE_INFO[_opcode] = (_pkt_cls, _pkt_cls._size_nibbles, _delta_lo, _delta_mask, _special)
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# DECODER
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def decode(data: bytes) -> Iterator[PacketType]:
-  """Decode raw SQTT blob, yielding packet instances."""
-  n, reg, pos, nib_off, nib_count, time = len(data), 0, 0, 0, 16, 0
-
-  while pos + ((nib_count + nib_off + 1) >> 1) <= n:
-    need = nib_count - nib_off
-    # 1. if unaligned, read high nibble to align
-    if nib_off: reg, pos = (reg >> 4) | ((data[pos] >> 4) << 60), pos + 1
-    # 2. read all full bytes at once
-    if (byte_count := need >> 1):
-      chunk = int.from_bytes(data[pos:pos + byte_count], 'little')
-      reg, pos = (reg >> (byte_count * 8)) | (chunk << (64 - byte_count * 8)), pos + byte_count
-    # 3. if odd, read low nibble
-    if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)
-
-    opcode = STATE_TO_OPCODE[reg & 0xFF]
-    pkt_cls, nib_count, delta_lo, delta_mask, special = _DECODE_INFO[opcode]
-    delta = (reg >> delta_lo) & delta_mask
-    if special == 1 and (reg >> 9) & 1 and not (reg >> 8) & 1: delta = 0  # TS_DELTA_OR_MARK marker
-    elif special == 2: delta += 8  # TS_DELTA_SHORT
-    time += delta
-    yield pkt_cls.from_raw(reg, time)
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# PRINTER
-# ═══════════════════════════════════════════════════════════════════════════════
-
-PACKET_COLORS = {
-  "INST": "WHITE", "VALUINST": "BLACK", "VMEMEXEC": "yellow", "ALUEXEC": "yellow",
-  "IMMEDIATE": "YELLOW", "IMMEDIATE_MASK": "YELLOW", "WAVERDY": "cyan", "WAVEALLOC": "cyan",
-  "WAVEEND": "blue", "WAVESTART": "blue", "PERF": "magenta", "EVENT": "red", "EVENT_BIG": "red",
-  "REG": "green", "LAYOUT_HEADER": "white", "SNAPSHOT": "white", "UTILCTR": "green",
-}
-
-def format_packet(p) -> str:
-  from tinygrad.helpers import colored
-  name = type(p).__name__
-  if isinstance(p, INST):
-    op_name = p.op.name if isinstance(p.op, InstOp) else f"0x{p.op:02x}"
-    fields = f"wave={p.wave} op={op_name}" + (" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "")
-  elif isinstance(p, VALUINST): fields = f"wave={p.wave}" + (" flag" if p.flag else "")
-  elif isinstance(p, ALUEXEC): fields = f"src={p.src.name if isinstance(p.src, AluSrc) else p.src}"
-  elif isinstance(p, VMEMEXEC): fields = f"src={p.src.name if isinstance(p.src, MemSrc) else p.src}"
-  elif isinstance(p, (WAVESTART, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
-  elif hasattr(p, '_fields'):
-    fields = " ".join(f"{k}=0x{getattr(p, k):x}" if k in {'snap', 'val32'} else f"{k}={getattr(p, k)}"
-                      for k in p._fields if not k.startswith('_') and k not in {'delta', 'encoding'})
-  else: fields = ""
-  return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name, 'white'))} {fields}"
-
-def print_packets(packets) -> None:
-  skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"}
-  for p in packets:
-    if type(p).__name__ not in skip: print(format_packet(p))
-
-if __name__ == "__main__":
-  import sys, pickle
-  if len(sys.argv) < 2:
-    print("Usage: python sqtt.py <pkl_file>")
-    sys.exit(1)
-  with open(sys.argv[1], "rb") as f:
-    data = pickle.load(f)
-  sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
-  for i, event in enumerate(sqtt_events):
-    print(f"\n=== event {i} ===")
-    print_packets(decode(event.blob))

extra/assembly/amd/test/bench_emu.py

@@ -1,192 +0,0 @@
-#!/usr/bin/env python3
-"""Benchmark comparing Python vs Rust RDNA3 emulators on real tinygrad kernels."""
-import ctypes, time, os
-from pathlib import Path
-
-# Set AMD=1 before importing tinygrad
-os.environ["AMD"] = "1"
-
-from extra.assembly.amd.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program
-
-REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
-if not REMU_PATH.exists():
-  REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.dylib"
-
-def get_rust_remu():
-  """Load the Rust libremu shared library."""
-  if not REMU_PATH.exists(): return None
-  remu = ctypes.CDLL(str(REMU_PATH))
-  remu.run_asm.restype = ctypes.c_int32
-  remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
-                           ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
-  return remu
-
-def count_instructions(kernel: bytes) -> int:
-  """Count instructions in a kernel."""
-  return len(decode_program(kernel))
-
-def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = None):
-  """Allocate buffers and return args pointer + valid ranges."""
-  if init_data is None: init_data = {}
-  buffers = []
-  for i, size in enumerate(buf_sizes):
-    padded = ((size + 15) // 16) * 16 + 16
-    data = init_data.get(i, b'\x00' * padded)
-    data_list = list(data) + [0] * (padded - len(data))
-    buf = (ctypes.c_uint8 * padded)(*data_list[:padded])
-    buffers.append(buf)
-  args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
-  args_ptr = ctypes.addressof(args)
-  ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
-  ranges.add((args_ptr, ctypes.sizeof(args)))
-  return buffers, args, args_ptr, ranges
-
-def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, rsrc2: int, iterations: int = 5):
-  """Benchmark an emulator and return average time."""
-  gx, gy, gz = global_size
-  lx, ly, lz = local_size
-  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-  lib_ptr = ctypes.addressof(kernel_buf)
-
-  # Warmup
-  run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
-
-  # Timed runs
-  times = []
-  for _ in range(iterations):
-    start = time.perf_counter()
-    result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
-    end = time.perf_counter()
-    if result != 0:
-      print(f"  {name} returned error: {result}")
-      return None
-    times.append(end - start)
-
-  return sum(times) / len(times)
-
-def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes], int] | None:
-  """Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data, rsrc2)."""
-  try:
-    from tinygrad import Tensor
-    from tinygrad.runtime.support.elf import elf_loader
-    from tinygrad.runtime.autogen import hsa
-    import numpy as np
-    np.random.seed(42)
-
-    ops = {
-      "add": lambda: Tensor.empty(1024) + Tensor.empty(1024),
-      "mul": lambda: Tensor.empty(1024) * Tensor.empty(1024),
-      "matmul_small": lambda: Tensor.empty(16, 16) @ Tensor.empty(16, 16),
-      "matmul_medium": lambda: Tensor.empty(64, 64) @ Tensor.empty(64, 64),
-      "reduce_sum": lambda: Tensor.empty(4096).sum(),
-      "reduce_max": lambda: Tensor.empty(4096).max(),
-      "softmax": lambda: Tensor.empty(256).softmax(),
-      "layernorm": lambda: Tensor.empty(32, 64).layernorm(),
-      "conv2d": lambda: Tensor.empty(1, 4, 16, 16).conv2d(Tensor.empty(4, 4, 3, 3)),
-      "gelu": lambda: Tensor.empty(1024).gelu(),
-      "exp": lambda: Tensor.empty(1024).exp(),
-      "sin": lambda: Tensor.empty(1024).sin(),
-    }
-
-    if op_name not in ops: return None
-    out = ops[op_name]()
-    sched = out.schedule()
-
-    for ei in sched:
-      lowered = ei.lower()
-      if ei.ast.op.name == 'SINK' and lowered.prg and lowered.prg.p.lib:
-        lib = bytes(lowered.prg.p.lib)
-        image = memoryview(bytearray(lib))
-        _, sections, _ = elf_loader(lib)
-        rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
-        for sec in sections:
-          if sec.name == '.text':
-            buf_sizes = [b.nbytes for b in lowered.bufs]
-            # Get initial data from numpy arrays if available
-            buf_data = {}
-            for i, buf in enumerate(lowered.bufs):
-              if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
-                try: buf_data[i] = bytes(buf.base._buf)
-                except: pass
-            # Extract rsrc2 from ELF (same as ops_amd.py)
-            group_segment_size = image[rodata_entry:rodata_entry+4].cast("I")[0]
-            lds_size = ((group_segment_size + 511) // 512) & 0x1FF
-            code = hsa.amd_kernel_code_t.from_buffer_copy(bytes(image[rodata_entry:rodata_entry+256]) + b'\x00'*256)
-            rsrc2 = code.compute_pgm_rsrc2 | (lds_size << 15)
-            return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data, rsrc2)
-    return None
-  except Exception as e:
-    print(f"  Error getting kernel: {e}")
-    return None
-
-TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
-
-def main():
-  import argparse
-  parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
-  parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
-  args = parser.parse_args()
-
-  rust_remu = get_rust_remu()
-  if rust_remu is None:
-    print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
-    print("Running Python-only benchmarks...\n")
-
-  print("=" * 90)
-  print("RDNA3 Emulator Benchmark: Python vs Rust")
-  print("=" * 90)
-
-  results = []
-
-  print("\n[TINYGRAD KERNELS]")
-  print("-" * 90)
-
-  for op_name in TINYGRAD_TESTS:
-    print(f"\n{op_name}:", end=" ", flush=True)
-    kernel_info = get_tinygrad_kernel(op_name)
-    if kernel_info is None:
-      print("failed to compile")
-      continue
-
-    kernel, global_size, local_size, buf_sizes, buf_data, rsrc2 = kernel_info
-    n_insts = count_instructions(kernel)
-    n_workgroups = global_size[0] * global_size[1] * global_size[2]
-    n_threads = local_size[0] * local_size[1] * local_size[2]
-    total_work = n_insts * n_workgroups * n_threads
-
-    print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
-
-    buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
-    set_valid_mem_ranges(ranges)
-
-    py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
-    rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations) if rust_remu else None
-
-    if py_time:
-      py_rate = total_work / py_time / 1e6
-      print(f"  Python: {py_time*1000:8.3f} ms  ({py_rate:7.2f} M ops/s)")
-    if rust_time:
-      rust_rate = total_work / rust_time / 1e6
-      speedup = py_time / rust_time if py_time else 0
-      print(f"  Rust:   {rust_time*1000:8.3f} ms  ({rust_rate:7.2f} M ops/s)  [{speedup:.1f}x faster]")
-
-    results.append((op_name, n_insts, n_workgroups, py_time, rust_time))
-
-  # Summary table
-  print("\n" + "=" * 90)
-  print("SUMMARY")
-  print("=" * 90)
-  print(f"{'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
-  print("-" * 90)
-
-  for name, n_insts, n_wgs, py_time, rust_time in results:
-    py_ms = f"{py_time*1000:.3f}" if py_time else "error"
-    if rust_time:
-      rust_ms = f"{rust_time*1000:.3f}"
-      speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
-    else:
-      rust_ms, speedup = "N/A", "N/A"
-    print(f"{name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
-
-if __name__ == "__main__":
-  main()

extra/assembly/amd/test/external_test_usability.py

@@ -1,196 +0,0 @@
-# Usability tests for the RDNA3 ASM DSL
-# These tests demonstrate how the DSL *should* work for a good user experience
-# Currently many of these tests fail - they document desired behavior
-
-import unittest
-from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.dsl import Inst, RawImm, SGPR, VGPR
-
-class TestRegisterSliceSyntax(unittest.TestCase):
-  """
-  Issue: Register slice syntax should use AMD assembly convention (inclusive end).
-
-  In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
-  The DSL should match this convention so that:
-  - s[4:7] gives 4 registers
-  - Disassembler output can be copied directly back into DSL code
-
-  Fix: Change _RegFactory.__getitem__ to use inclusive end:
-    key.stop - key.start + 1  (instead of key.stop - key.start)
-  """
-  def test_register_slice_count(self):
-    # s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
-    reg = s[4:7]
-    self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
-
-  def test_register_slice_roundtrip(self):
-    # Round-trip: DSL -> disasm -> DSL should preserve register count
-    reg = s[4:7]  # 4 registers in AMD convention
-    inst = s_load_b128(reg, s[0:1], NULL, 0)
-    disasm = inst.disasm()
-    # Disasm shows s[4:7] - user should be able to copy this back
-    self.assertIn("s[4:7]", disasm)
-    # And s[4:7] in DSL should give the same 4 registers
-    reg_from_disasm = s[4:7]
-    self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
-
-
-class TestReprReadability(unittest.TestCase):
-  """
-  Issue: repr() leaks internal RawImm type and omits zero-valued fields.
-
-  When you create v_mov_b32_e32(v[0], v[1]), the repr shows:
-    VOP1(op=1, src0=RawImm(257))
-
-  Problems:
-  1. vdst=v[0] is omitted because 0 is treated as "default"
-  2. src0 shows RawImm(257) instead of v[1]
-  3. User sees encoded values (257 = 256 + 1) instead of register names
-
-  Expected repr: VOP1(op=1, vdst=v[0], src0=v[1])
-  """
-  def test_repr_shows_registers_not_raw_imm(self):
-    inst = v_mov_b32_e32(v[0], v[1])
-    # Should show v[1], not RawImm(257)
-    self.assertNotIn("RawImm", repr(inst), "repr should not expose RawImm internal type")
-    self.assertIn("v[1]", repr(inst), "repr should show register name")
-
-  def test_repr_includes_zero_dst(self):
-    inst = v_mov_b32_e32(v[0], v[1])
-    # v[0] is a valid destination register, should be shown
-    self.assertIn("vdst", repr(inst), "repr should include vdst even when 0")
-
-  def test_repr_roundtrip(self):
-    # repr should produce something that can be eval'd back
-    inst = v_mov_b32_e32(v[0], v[1])
-    # This would require repr to output valid Python, e.g.:
-    # "VOP1(op=VOP1Op.V_MOV_B32, vdst=v[0], src0=v[1])"
-    r = repr(inst)
-    # At minimum, it should be human-readable
-    self.assertIn("v[", r, "repr should show register syntax")
-
-
-class TestInstructionEquality(unittest.TestCase):
-  """
-  Issue: No __eq__ method - instruction comparison requires repr() workaround.
-
-  Two identical instructions should compare equal with ==, but currently:
-    inst1 == inst2 returns False
-
-  The test_handwritten.py works around this with:
-    self.assertEqual(repr(self.inst), repr(reasm))
-  """
-  def test_identical_instructions_equal(self):
-    inst1 = v_mov_b32_e32(v[0], v[1])
-    inst2 = v_mov_b32_e32(v[0], v[1])
-    self.assertEqual(inst1, inst2, "identical instructions should be equal")
-
-  def test_different_instructions_not_equal(self):
-    inst1 = v_mov_b32_e32(v[0], v[1])
-    inst2 = v_mov_b32_e32(v[0], v[2])
-    self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
-
-
-class TestVOPDHelperSignature(unittest.TestCase):
-  """
-  Issue: VOPD helper functions have confusing semantics.
-
-  v_dual_mul_f32 is defined as:
-    v_dual_mul_f32 = functools.partial(VOPD, VOPDOp.V_DUAL_MUL_F32)
-
-  This binds VOPDOp.V_DUAL_MUL_F32 to the FIRST positional arg of VOPD.__init__,
-  which is 'opx'. So v_dual_mul_f32 sets the X operation.
-
-  But then test_dual_mul in test_handwritten.py does:
-    v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], ...)
-
-  This passes V_DUAL_MUL_F32 as the SECOND positional arg (opy), making both
-  X and Y operations the same. This is confusing because:
-  1. The function name suggests it handles the X operation
-  2. But you still pass an opcode as the first arg (which becomes opy)
-
-  Expected: Either make the helper fully specify both ops, or make the
-  signature clearer about what the positional arg means.
-  """
-  def test_vopd_helper_opy_should_be_required(self):
-    # Using only keyword args "works" but opy silently defaults to 0
-    inst = v_dual_mul_f32(vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
-    self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32)
-    # Bug: opy defaults to 0 (V_DUAL_FMAC_F32) silently - should require explicit opy
-    # This test documents the bug - it should fail once fixed
-    self.assertNotEqual(inst.opy, VOPDOp.V_DUAL_FMAC_F32, "opy should not silently default to FMAC")
-
-  def test_vopd_helper_positional_arg_is_opy(self):
-    # The first positional arg after the partial becomes opy, not a second opx
-    inst = v_dual_mul_f32(VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
-    self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32)  # From partial
-    self.assertEqual(inst.opy, VOPDOp.V_DUAL_MOV_B32)  # From first positional arg
-
-
-class TestFieldAccessPreservesType(unittest.TestCase):
-  """
-  Issue: Field access loses type information.
-
-  After creating an instruction, accessing fields returns encoded int values:
-    inst = v_mov_b32_e32(v[0], v[1])
-    inst.vdst  # returns 0, not VGPR(0)
-
-  This makes it impossible to round-trip register types through field access.
-  """
-  def test_vdst_returns_register(self):
-    inst = v_mov_b32_e32(v[5], v[1])
-    vdst = inst.vdst
-    # Should return a VGPR, not an int
-    self.assertIsInstance(vdst, (VGPR, int), "vdst should return VGPR or at least be usable")
-    # Ideally: self.assertIsInstance(vdst, VGPR)
-
-  def test_src_returns_register_for_vgpr_source(self):
-    inst = v_mov_b32_e32(v[0], v[1])
-    # src0 is encoded as 257 (256 + 1 for v1)
-    # Ideally it should decode back to v[1]
-    src0_raw = inst._values.get('src0')
-    # Currently returns RawImm(257), should return VGPR(1) or similar
-    self.assertNotIsInstance(src0_raw, RawImm, "source should not be RawImm internally")
-
-
-class TestArgumentDiscoverability(unittest.TestCase):
-  """
-  Issue: No clear signature for positional arguments.
-
-  inspect.signature(s_load_b128) shows: (*args, literal=None, **kwargs)
-
-  Users have no way to know the argument order without reading source code.
-  The order is implicitly defined by the class field definition order.
-
-  Possible fixes:
-  1. Add explicit parameter names to functools.partial
-  2. Generate type stubs with proper signatures
-  3. Add docstrings listing the expected arguments
-  """
-  def test_signature_has_named_params(self):
-    import inspect
-    sig = inspect.signature(s_load_b128)
-    params = list(sig.parameters.keys())
-    # Currently: ['args', 'literal', 'kwargs'] (from *args, literal=None, **kwargs)
-    # Expected: something like ['sdata', 'sbase', 'soffset', 'offset', 'literal']
-    self.assertIn('sdata', params, "signature should show field names")
-
-
-class TestSpecialConstants(unittest.TestCase):
-  """
-  Issue: NULL and other constants are IntEnum values that might be confusing.
-
-  NULL = SrcEnum.NULL = 124, but users might expect NULL to be a special object
-  that clearly represents "no register" rather than a magic number.
-  """
-  def test_null_has_clear_repr(self):
-    # NULL should have a clear string representation
-    self.assertIn("NULL", str(NULL) or repr(NULL), "NULL should be clearly identifiable")
-
-  def test_null_is_distinguishable_from_int(self):
-    # NULL should be distinguishable from the raw integer 124
-    self.assertNotEqual(type(NULL), int, "NULL should not be plain int")
-
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/amd/test/helpers.py

@@ -1,66 +0,0 @@
-"""Shared test helpers for RDNA3 tests."""
-import shutil
-from dataclasses import dataclass
-
-@dataclass
-class KernelInfo:
-  code: bytes
-  global_size: tuple[int, int, int]
-  local_size: tuple[int, int, int]
-  buf_idxs: list[int]  # indices into shared buffer pool
-  buf_sizes: list[int]  # sizes for each buffer index
-
-# LLVM tool detection (shared across test files)
-def get_llvm_mc():
-  """Find llvm-mc executable, preferring newer versions."""
-  for p in ['llvm-mc', 'llvm-mc-21', 'llvm-mc-20']:
-    if shutil.which(p): return p
-  raise FileNotFoundError("llvm-mc not found")
-
-def get_llvm_objdump():
-  """Find llvm-objdump executable, preferring newer versions."""
-  for p in ['llvm-objdump', 'llvm-objdump-21', 'llvm-objdump-20']:
-    if shutil.which(p): return p
-  raise FileNotFoundError("llvm-objdump not found")
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# EXECUTION CONTEXT (for testing compiled pseudocode)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-class ExecContext:
-  """Context for running compiled pseudocode in tests."""
-  def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=0xffffffff, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
-    from extra.assembly.amd.pcode import Reg, MASK32, MASK64, TypedView
-    self._Reg, self._MASK64, self._TypedView = Reg, MASK64, TypedView
-    self.S0, self.S1, self.S2 = Reg(s0), Reg(s1), Reg(s2)
-    self.D0, self.D1 = Reg(d0), Reg(0)
-    self.SCC, self.VCC, self.EXEC = Reg(scc), Reg(vcc), Reg(exec_mask)
-    self.tmp, self.saveexec = Reg(0), Reg(exec_mask)
-    self.lane, self.laneId, self.literal = lane, lane, literal
-    self.SIMM16, self.SIMM32 = Reg(literal), Reg(literal)
-    self.VGPR = vgprs if vgprs is not None else {}
-    self.SRC0, self.VDST = Reg(src0_idx), Reg(vdst_idx)
-
-  def run(self, code: str):
-    """Execute compiled code."""
-    import extra.assembly.amd.pcode as pcode
-    ns = {k: getattr(pcode, k) for k in dir(pcode) if not k.startswith('_')}
-    # Also include underscore-prefixed helpers that compiled pseudocode uses
-    for k in ['_pack', '_pack32']:
-      if hasattr(pcode, k): ns[k] = getattr(pcode, k)
-    ns.update({
-      'S0': self.S0, 'S1': self.S1, 'S2': self.S2, 'D0': self.D0, 'D1': self.D1,
-      'SCC': self.SCC, 'VCC': self.VCC, 'EXEC': self.EXEC,
-      'EXEC_LO': self._TypedView(self.EXEC, 31, 0), 'EXEC_HI': self._TypedView(self.EXEC, 63, 32),
-      'tmp': self.tmp, 'saveexec': self.saveexec,
-      'lane': self.lane, 'laneId': self.laneId, 'literal': self.literal,
-      'SIMM16': self.SIMM16, 'SIMM32': self.SIMM32, 'VGPR': self.VGPR, 'SRC0': self.SRC0, 'VDST': self.VDST,
-    })
-    exec(code, ns)
-    def _sync(ctx_reg, ns_val):
-      if isinstance(ns_val, self._Reg): ctx_reg._val = ns_val._val
-      else: ctx_reg._val = int(ns_val) & self._MASK64
-    for name in ('SCC', 'VCC', 'EXEC', 'D0', 'D1', 'tmp', 'saveexec'):
-      if ns.get(name) is not getattr(self, name): _sync(getattr(self, name), ns[name])
-
-  def result(self) -> dict: return {"d0": self.D0._val, "scc": self.SCC._val & 1}

extra/assembly/amd/test/hw/test_sop.py

@@ -1,381 +0,0 @@
-"""Tests for SOP instructions - scalar operations.
-
-Includes: s_add_u32, s_mov_b32, s_and_b32, s_or_b32, s_quadmask_b32, s_wqm_b32,
-          s_cbranch_vccnz, s_cbranch_vccz
-"""
-import unittest
-from extra.assembly.amd.test.hw.helpers import *
-
-class TestBasicScalar(unittest.TestCase):
-  """Tests for basic scalar operations."""
-
-  def test_s_add_u32(self):
-    """S_ADD_U32 adds two scalar values."""
-    instructions = [
-      s_mov_b32(s[0], 100),
-      s_mov_b32(s[1], 200),
-      s_add_u32(s[2], s[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 300)
-
-  def test_s_add_u32_carry(self):
-    """S_ADD_U32 sets SCC on overflow."""
-    instructions = [
-      s_mov_b32(s[0], 64),
-      s_not_b32(s[0], s[0]),  # ~64 = 0xffffffbf
-      s_mov_b32(s[1], 64),
-      s_add_u32(s[2], s[0], s[1]),  # 0xffffffbf + 64 = 0xffffffff
-      s_mov_b32(s[3], 1),
-      s_add_u32(s[4], s[2], s[3]),  # 0xffffffff + 1 = overflow
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[4], 0)
-    self.assertEqual(st.scc, 1)
-
-  def test_s_brev_b32(self):
-    """S_BREV_B32 reverses bits of a 32-bit value."""
-    # 10 = 0b00000000000000000000000000001010
-    # reversed = 0b01010000000000000000000000000000 = 0x50000000
-    instructions = [
-      s_mov_b32(s[0], 10),
-      s_brev_b32(s[1], s[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[1], 0x50000000)
-
-  def test_s_brev_b32_all_ones(self):
-    """S_BREV_B32 with all ones stays all ones."""
-    instructions = [
-      s_mov_b32(s[0], 0xFFFFFFFF),
-      s_brev_b32(s[1], s[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
-
-  def test_s_brev_b32_single_bit(self):
-    """S_BREV_B32 with bit 0 set becomes bit 31."""
-    instructions = [
-      s_mov_b32(s[0], 1),
-      s_brev_b32(s[1], s[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[1], 0x80000000)
-
-
-class TestQuadmaskWqm(unittest.TestCase):
-  """Tests for S_QUADMASK_B32 and S_WQM_B32."""
-
-  def test_s_quadmask_b32_all_quads_active(self):
-    """S_QUADMASK_B32 with all quads active."""
-    instructions = [
-      s_mov_b32(s[0], 0xFFFFFFFF),  # All lanes active
-      s_quadmask_b32(s[1], s[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # Each quad (4 lanes) with any bit set -> 1 bit in result
-    # 32 lanes = 8 quads, all active -> 0xFF
-    self.assertEqual(st.sgpr[1], 0xFF)
-
-  def test_s_quadmask_b32_alternating_quads(self):
-    """S_QUADMASK_B32 with alternating quads active."""
-    instructions = [
-      s_mov_b32(s[0], 0x0F0F0F0F),  # Quads 0,2,4,6 active
-      s_quadmask_b32(s[1], s[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # Quads 0,2,4,6 have at least one bit -> 0b01010101 = 0x55
-    self.assertEqual(st.sgpr[1], 0x55)
-
-  def test_s_quadmask_b32_no_quads_active(self):
-    """S_QUADMASK_B32 with no quads active."""
-    instructions = [
-      s_mov_b32(s[0], 0),
-      s_quadmask_b32(s[1], s[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[1], 0)
-
-  def test_s_quadmask_b32_single_lane_per_quad(self):
-    """S_QUADMASK_B32 with single lane active in each quad."""
-    instructions = [
-      s_mov_b32(s[0], 0x11111111),  # Bit 0 of each nibble
-      s_quadmask_b32(s[1], s[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # All 8 quads have at least one lane -> 0xFF
-    self.assertEqual(st.sgpr[1], 0xFF)
-
-  def test_s_wqm_b32_all_active(self):
-    """S_WQM_B32 with all lanes active returns all 1s."""
-    instructions = [
-      s_mov_b32(s[0], 0xFFFFFFFF),
-      s_wqm_b32(s[1], s[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
-
-  def test_s_wqm_b32_alternating_quads(self):
-    """S_WQM_B32 with single lane per quad expands to full quads."""
-    instructions = [
-      s_mov_b32(s[0], 0x11111111),  # One lane per quad
-      s_wqm_b32(s[1], s[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # Each quad with any bit expands to all 4 bits
-    self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
-
-  def test_s_wqm_b32_zero(self):
-    """S_WQM_B32 with zero input returns zero."""
-    instructions = [
-      s_mov_b32(s[0], 0),
-      s_wqm_b32(s[1], s[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[1], 0)
-
-
-class TestBranch(unittest.TestCase):
-  """Tests for branch instructions."""
-
-  def test_cbranch_vccnz_ignores_vcc_hi(self):
-    """S_CBRANCH_VCCNZ should only check VCC_LO in wave32."""
-    instructions = [
-      # Set VCC_LO = 0, VCC_HI = 1
-      s_mov_b32(VCC_LO, 0),
-      s_mov_b32(VCC_HI, 1),
-      v_mov_b32_e32(v[0], 0),
-      # If VCC_HI is incorrectly used, branch will be taken
-      s_cbranch_vccnz(1),  # Skip next instruction if VCC != 0
-      v_mov_b32_e32(v[0], 42),  # This should execute
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 42, "Branch should NOT be taken (VCC_LO is 0)")
-
-  def test_cbranch_vccz_ignores_vcc_hi(self):
-    """S_CBRANCH_VCCZ should only check VCC_LO in wave32."""
-    instructions = [
-      # Set VCC_LO = 1, VCC_HI = 0
-      s_mov_b32(VCC_LO, 1),
-      s_mov_b32(VCC_HI, 0),
-      v_mov_b32_e32(v[0], 0),
-      # If VCC_HI is incorrectly used, branch will be taken
-      s_cbranch_vccz(1),  # Skip next instruction if VCC == 0
-      v_mov_b32_e32(v[0], 42),  # This should execute
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 42, "Branch should NOT be taken (VCC_LO is 1)")
-
-  def test_cbranch_vccnz_branches_on_vcc_lo(self):
-    """S_CBRANCH_VCCNZ branches when VCC_LO is non-zero."""
-    instructions = [
-      s_mov_b32(VCC_LO, 1),
-      v_mov_b32_e32(v[0], 0),
-      s_cbranch_vccnz(1),  # Skip next instruction if VCC != 0
-      v_mov_b32_e32(v[0], 42),  # This should be skipped
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0, "Branch should be taken (VCC_LO is 1)")
-
-
-class Test64BitLiterals(unittest.TestCase):
-  """Tests for 64-bit literal encoding in instructions."""
-
-  def test_64bit_literal_negative_encoding(self):
-    """64-bit literal -2^32 encodes correctly."""
-    lit = -4294967296.0  # -2^32
-    lit_bits = f2i64(lit)
-    instructions = [
-      s_mov_b32(s[0], lit_bits & 0xffffffff),
-      s_mov_b32(s[1], lit_bits >> 32),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = i642f(st.vgpr[0][0] | (st.vgpr[0][1] << 32))
-    self.assertAlmostEqual(result, -4294967296.0, places=5)
-
-class TestSCCBehavior(unittest.TestCase):
-  """Tests for SCC condition code behavior."""
-
-  def test_scc_from_s_cmp(self):
-    """SCC should be set by scalar compare."""
-    instructions = [
-      s_mov_b32(s[0], 10),
-      s_cmp_eq_u32(s[0], 10),
-      s_cselect_b32(s[1], 1, 0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[1], 1, "SCC should be true")
-    self.assertEqual(st.scc, 1)
-
-  def test_scc_clear(self):
-    """SCC should be cleared by failing compare."""
-    instructions = [
-      s_mov_b32(s[0], 10),
-      s_cmp_eq_u32(s[0], 20),
-      s_cselect_b32(s[1], 1, 0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[1], 0, "SCC should be false")
-    self.assertEqual(st.scc, 0)
-
-
-class TestSignedArithmetic(unittest.TestCase):
-  """Tests for S_ADD_I32, S_SUB_I32 and their SCC overflow behavior."""
-
-  def test_s_add_i32_no_overflow(self):
-    """S_ADD_I32: 1 + 1 = 2, no overflow, SCC=0."""
-    instructions = [
-      s_mov_b32(s[0], 1),
-      s_add_i32(s[1], s[0], 1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[1], 2)
-    self.assertEqual(st.scc, 0, "No overflow, SCC should be 0")
-
-  def test_s_add_i32_positive_overflow(self):
-    """S_ADD_I32: MAX_INT + 1 overflows, SCC=1."""
-    instructions = [
-      s_mov_b32(s[0], 0x7FFFFFFF),  # MAX_INT
-      s_add_i32(s[1], s[0], 1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[1], 0x80000000)  # Wraps to MIN_INT
-    self.assertEqual(st.scc, 1, "Overflow, SCC should be 1")
-
-  def test_s_add_i32_negative_no_overflow(self):
-    """S_ADD_I32: -10 + 20 = 10, no overflow."""
-    instructions = [
-      s_mov_b32(s[0], 0xFFFFFFF6),  # -10 in two's complement
-      s_mov_b32(s[1], 20),
-      s_add_i32(s[2], s[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 10)
-    self.assertEqual(st.scc, 0)
-
-  def test_s_add_i32_negative_overflow(self):
-    """S_ADD_I32: MIN_INT + (-1) underflows, SCC=1."""
-    instructions = [
-      s_mov_b32(s[0], 0x80000000),  # MIN_INT
-      s_mov_b32(s[1], 0xFFFFFFFF),  # -1
-      s_add_i32(s[2], s[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 0x7FFFFFFF)  # Wraps to MAX_INT
-    self.assertEqual(st.scc, 1, "Underflow, SCC should be 1")
-
-  def test_s_sub_i32_no_overflow(self):
-    """S_SUB_I32: 10 - 5 = 5, no overflow."""
-    instructions = [
-      s_mov_b32(s[0], 10),
-      s_mov_b32(s[1], 5),
-      s_sub_i32(s[2], s[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 5)
-    self.assertEqual(st.scc, 0)
-
-  def test_s_sub_i32_overflow(self):
-    """S_SUB_I32: MAX_INT - (-1) overflows, SCC=1."""
-    instructions = [
-      s_mov_b32(s[0], 0x7FFFFFFF),  # MAX_INT
-      s_mov_b32(s[1], 0xFFFFFFFF),  # -1
-      s_sub_i32(s[2], s[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 0x80000000)  # Wraps to MIN_INT
-    self.assertEqual(st.scc, 1, "Overflow, SCC should be 1")
-
-  def test_s_mul_hi_u32(self):
-    """S_MUL_HI_U32: high 32 bits of u32 * u32."""
-    instructions = [
-      s_mov_b32(s[0], 0x80000000),  # 2^31
-      s_mov_b32(s[1], 4),
-      s_mul_hi_u32(s[2], s[0], s[1]),  # (2^31 * 4) >> 32 = 2
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 2)
-
-  def test_s_mul_i32(self):
-    """S_MUL_I32: signed multiply low 32 bits."""
-    instructions = [
-      s_mov_b32(s[0], 0xFFFFFFFF),  # -1
-      s_mov_b32(s[1], 10),
-      s_mul_i32(s[2], s[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 0xFFFFFFF6)  # -10
-
-  def test_division_sequence_from_llvm(self):
-    """Test the division sequence pattern from LLVM-generated code."""
-    # This sequence is from the sin kernel and computes integer division
-    # s10 = dividend, s18 = divisor, result in s6/s14
-    dividend = 0x28BE60DB  # Some value from the sin kernel
-    divisor = 3  # Simplified divisor
-    instructions = [
-      s_mov_b32(s[10], dividend),
-      s_mov_b32(s[18], divisor),
-      # Compute reciprocal approximation: s6 = ~0 / divisor (approx)
-      s_mov_b32(s[11], 0),
-      s_sub_i32(s[11], s[11], s[18]),  # s11 = -divisor
-      # For testing, just verify basic arithmetic works
-      s_mul_i32(s[6], s[10], 2),
-      s_add_i32(s[7], s[6], 1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[6], (dividend * 2) & 0xFFFFFFFF)
-    self.assertEqual(st.sgpr[7], ((dividend * 2) + 1) & 0xFFFFFFFF)
-
-
-class Test64BitCompare(unittest.TestCase):
-  """Tests for 64-bit scalar compare instructions."""
-
-  def test_s_cmp_eq_u64_equal(self):
-    """S_CMP_EQ_U64: comparing equal 64-bit values sets SCC=1."""
-    val = 0x123456789ABCDEF0
-    instructions = [
-      s_mov_b32(s[0], val & 0xFFFFFFFF),
-      s_mov_b32(s[1], val >> 32),
-      s_mov_b32(s[2], val & 0xFFFFFFFF),
-      s_mov_b32(s[3], val >> 32),
-      s_cmp_eq_u64(s[0:1], s[2:3]),
-      s_cselect_b32(s[4], 1, 0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.scc, 1)
-    self.assertEqual(st.sgpr[4], 1)
-
-  def test_s_cmp_eq_u64_different_upper_bits(self):
-    """S_CMP_EQ_U64: values differing only in upper 32 bits are not equal."""
-    # This is the bug case - if only lower 32 bits are compared, these would be equal
-    instructions = [
-      s_mov_b32(s[0], 0),  # lower 32 bits of value 0
-      s_mov_b32(s[1], 0),  # upper 32 bits of value 0
-      s_mov_b32(s[2], 0),  # lower 32 bits of 0x100000000
-      s_mov_b32(s[3], 1),  # upper 32 bits of 0x100000000
-      s_cmp_eq_u64(s[0:1], s[2:3]),
-      s_cselect_b32(s[4], 1, 0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.scc, 0, "0 != 0x100000000, SCC should be 0")
-    self.assertEqual(st.sgpr[4], 0)
-
-  def test_s_cmp_lg_u64_different(self):
-    """S_CMP_LG_U64: different 64-bit values sets SCC=1."""
-    instructions = [
-      s_mov_b32(s[0], 0),
-      s_mov_b32(s[1], 0),  # s[0:1] = 0
-      s_mov_b32(s[2], 0),
-      s_mov_b32(s[3], 1),  # s[2:3] = 0x100000000
-      s_cmp_lg_u64(s[0:1], s[2:3]),
-      s_cselect_b32(s[4], 1, 0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.scc, 1, "0 != 0x100000000, SCC should be 1")
-    self.assertEqual(st.sgpr[4], 1)
-
-
-if __name__ == '__main__':
-  unittest.main()

extra/assembly/amd/test/hw/test_vop2.py

@@ -1,451 +0,0 @@
-"""Tests for VOP2 instructions - two operand vector operations.
-
-Includes: v_add_f32, v_mul_f32, v_and_b32, v_or_b32, v_xor_b32,
-          v_lshrrev_b32, v_lshlrev_b32, v_fmac_f32, v_fmaak_f32, v_fmamk_f32,
-          v_add_nc_u32, v_cndmask_b32, v_add_f16, v_mul_f16
-"""
-import unittest
-from extra.assembly.amd.test.hw.helpers import *
-
-class TestBasicArithmetic(unittest.TestCase):
-  """Tests for basic arithmetic VOP2 instructions."""
-
-  def test_v_add_f32(self):
-    """V_ADD_F32 adds two floats."""
-    instructions = [
-      v_mov_b32_e32(v[0], 1.0),
-      v_mov_b32_e32(v[1], 2.0),
-      v_add_f32_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][2]), 3.0, places=5)
-
-  def test_v_mul_f32(self):
-    """V_MUL_F32 multiplies two floats."""
-    instructions = [
-      v_mov_b32_e32(v[0], 2.0),
-      v_mov_b32_e32(v[1], 4.0),
-      v_mul_f32_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][2]), 8.0, places=5)
-
-  def test_v_fmac_f32(self):
-    """V_FMAC_F32: d = d + a*b using inline constants."""
-    instructions = [
-      v_mov_b32_e32(v[0], 2.0),
-      v_mov_b32_e32(v[1], 4.0),
-      v_mov_b32_e32(v[2], 1.0),
-      v_fmac_f32_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
-
-  def test_v_fmaak_f32(self):
-    """V_FMAAK_F32: d = a * b + K using inline constants."""
-    instructions = [
-      v_mov_b32_e32(v[0], 2.0),
-      v_mov_b32_e32(v[1], 4.0),
-      v_fmaak_f32_e32(v[2], v[0], v[1], literal=0x3f800000),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
-
-  def test_v_fmamk_f32_basic(self):
-    """V_FMAMK_F32: d = a * K + b."""
-    instructions = [
-      v_mov_b32_e32(v[0], 2.0),
-      v_mov_b32_e32(v[1], 1.0),
-      v_fmamk_f32_e32(v[2], v[0], v[1], literal=0x40800000),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
-
-  def test_v_fmamk_f32_small_constant(self):
-    """V_FMAMK_F32 with small constant."""
-    instructions = [
-      v_mov_b32_e32(v[0], 4.0),
-      v_mov_b32_e32(v[1], 1.0),
-      v_fmamk_f32_e32(v[2], v[0], v[1], literal=f2i(0.5)),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][2]), 3.0, places=5)
-
-
-class TestBitManipulation(unittest.TestCase):
-  """Tests for bit manipulation VOP2 instructions."""
-
-  def test_v_and_b32(self):
-    """V_AND_B32 bitwise and."""
-    instructions = [
-      s_mov_b32(s[0], 0xff),
-      s_mov_b32(s[1], 0x0f),
-      v_mov_b32_e32(v[0], s[0]),
-      v_and_b32_e32(v[1], s[1], v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][1], 0x0f)
-
-  def test_v_and_b32_quadrant(self):
-    """V_AND_B32 for quadrant extraction (n & 3)."""
-    instructions = [
-      s_mov_b32(s[0], 15915),
-      v_mov_b32_e32(v[0], s[0]),
-      v_and_b32_e32(v[1], 3, v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][1], 15915 & 3)
-
-  def test_v_lshrrev_b32(self):
-    """V_LSHRREV_B32 logical shift right."""
-    instructions = [
-      s_mov_b32(s[0], 0xff00),
-      v_mov_b32_e32(v[0], s[0]),
-      v_lshrrev_b32_e32(v[1], 8, v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][1], 0xff)
-
-  def test_v_lshlrev_b32(self):
-    """V_LSHLREV_B32 logical shift left."""
-    instructions = [
-      s_mov_b32(s[0], 0xff),
-      v_mov_b32_e32(v[0], s[0]),
-      v_lshlrev_b32_e32(v[1], 8, v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][1], 0xff00)
-
-  def test_v_xor_b32(self):
-    """V_XOR_B32 bitwise xor (used in sin for sign)."""
-    instructions = [
-      s_mov_b32(s[0], 0x80000000),
-      s_mov_b32(s[1], f2i(1.0)),
-      v_mov_b32_e32(v[0], s[1]),
-      v_xor_b32_e32(v[1], s[0], v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][1]), -1.0, places=5)
-
-  def test_v_xor_b32_sign_flip(self):
-    """V_XOR_B32 for sign flip pattern."""
-    instructions = [
-      s_mov_b32(s[0], 0x80000000),
-      v_mov_b32_e32(v[0], -2.0),
-      v_xor_b32_e32(v[1], s[0], v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][1]), 2.0, places=5)
-
-
-class TestSpecialValues(unittest.TestCase):
-  """Tests for special float values - inf, nan, zero handling."""
-
-  def test_v_mul_f32_zero_times_inf(self):
-    """V_MUL_F32: 0 * inf = NaN."""
-    import math
-    instructions = [
-      v_mov_b32_e32(v[0], 0),
-      s_mov_b32(s[0], 0x7f800000),
-      v_mov_b32_e32(v[1], s[0]),
-      v_mul_f32_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
-
-  def test_v_add_f32_inf_minus_inf(self):
-    """V_ADD_F32: inf + (-inf) = NaN."""
-    import math
-    instructions = [
-      s_mov_b32(s[0], 0x7f800000),
-      s_mov_b32(s[1], 0xff800000),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_add_f32_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
-
-
-class TestF16Ops(unittest.TestCase):
-  """Tests for 16-bit VOP2 operations."""
-
-  def test_v_add_f16_basic(self):
-    """V_ADD_F16 adds two f16 values."""
-    instructions = [
-      s_mov_b32(s[0], 0x3c00),  # f16 1.0
-      s_mov_b32(s[1], 0x4000),  # f16 2.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_add_f16_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    self.assertEqual(result, 0x4200, f"Expected 0x4200 (f16 3.0), got 0x{result:04x}")
-
-  def test_v_add_f16_negative(self):
-    """V_ADD_F16 with negative values."""
-    instructions = [
-      s_mov_b32(s[0], 0x3c00),  # f16 1.0
-      s_mov_b32(s[1], 0xc000),  # f16 -2.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_add_f16_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    self.assertEqual(result, 0xbc00, f"Expected 0xbc00 (f16 -1.0), got 0x{result:04x}")
-
-  def test_v_mul_f16_basic(self):
-    """V_MUL_F16 multiplies two f16 values."""
-    instructions = [
-      s_mov_b32(s[0], 0x4000),  # f16 2.0
-      s_mov_b32(s[1], 0x4200),  # f16 3.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mul_f16_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    self.assertEqual(result, 0x4600, f"Expected 0x4600 (f16 6.0), got 0x{result:04x}")
-
-  def test_v_mul_f16_by_zero(self):
-    """V_MUL_F16 by zero."""
-    instructions = [
-      s_mov_b32(s[0], 0x4000),  # f16 2.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0),
-      v_mul_f16_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    self.assertEqual(result, 0x0000, f"Expected 0x0000 (f16 0.0), got 0x{result:04x}")
-
-  def test_v_fmac_f16_basic(self):
-    """V_FMAC_F16: d = d + a*b."""
-    instructions = [
-      s_mov_b32(s[0], 0x4000),  # f16 2.0
-      s_mov_b32(s[1], 0x4200),  # f16 3.0
-      s_mov_b32(s[2], 0x3c00),  # f16 1.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], s[2]),
-      v_fmac_f16_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    # 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
-    self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
-
-  def test_v_fmaak_f16_basic(self):
-    """V_FMAAK_F16: d = a * b + K."""
-    instructions = [
-      s_mov_b32(s[0], 0x4000),  # f16 2.0
-      s_mov_b32(s[1], 0x4200),  # f16 3.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_fmaak_f16_e32(v[2], v[0], v[1], literal=0x3c00),  # + f16 1.0
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    # 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
-    self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
-
-
-class TestHiHalfOps(unittest.TestCase):
-  """Tests for VOP2 16-bit operations with hi-half operands."""
-
-  def test_v_add_f16_src0_hi_fold(self):
-    """V_ADD_F16 with src0 hi-half fold (same register, different halves)."""
-    instructions = [
-      s_mov_b32(s[0], 0x40003c00),  # lo=f16(1.0), hi=f16(2.0)
-      v_mov_b32_e32(v[0], s[0]),
-      VOP3(VOP3Op.V_ADD_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0001),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][1] & 0xffff
-    self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
-
-  def test_v_add_f16_src0_hi_different_reg(self):
-    """V_ADD_F16 with src0 hi-half from different register."""
-    instructions = [
-      s_mov_b32(s[0], 0x40000000),  # hi=f16(2.0), lo=0
-      s_mov_b32(s[1], 0x00003c00),  # hi=0, lo=f16(1.0)
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      VOP3(VOP3Op.V_ADD_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
-
-  def test_v_mul_f16_src0_hi(self):
-    """V_MUL_F16 with src0 from high half."""
-    instructions = [
-      s_mov_b32(s[0], 0x40000000),  # hi=f16(2.0), lo=0
-      s_mov_b32(s[1], 0x00004200),  # hi=0, lo=f16(3.0)
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      VOP3(VOP3Op.V_MUL_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    self.assertEqual(result, 0x4600, f"Expected f16(6.0)=0x4600, got 0x{result:04x}")
-
-  def test_v_mul_f16_hi_half(self):
-    """V_MUL_F16 reading from high half."""
-    instructions = [
-      s_mov_b32(s[0], 0x40003c00),  # lo=1.0, hi=2.0
-      v_mov_b32_e32(v[0], s[0]),
-      VOP3(VOP3Op.V_MUL_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0011),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][1] & 0xffff
-    self.assertEqual(result, 0x4400, f"Expected f16(4.0)=0x4400, got 0x{result:04x}")
-
-  def test_v_fma_f16_hi_dest(self):
-    """V_FMA_F16 writing to high half with opsel.
-
-    Uses V_FMA_F16 (not V_FMAC_F16) because it has explicit src2 operand
-    which makes opsel handling clearer.
-    """
-    instructions = [
-      s_mov_b32(s[0], 0x3c000000),  # hi=f16(1.0), lo=0
-      s_mov_b32(s[1], 0x4000),      # f16(2.0) in lo
-      s_mov_b32(s[2], 0x4200),      # f16(3.0) in lo
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], s[2]),
-      # V_FMA_F16: dst = src0 * src1 + src2
-      # opsel=0b1100: bit2=src2 hi, bit3=dst hi
-      # So: v[0].hi = v[1].lo * v[2].lo + v[0].hi = 2.0 * 3.0 + 1.0 = 7.0
-      VOP3(VOP3Op.V_FMA_F16, vdst=v[0], src0=v[1], src1=v[2], src2=v[0], opsel=0b1100),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    hi = (st.vgpr[0][0] >> 16) & 0xffff
-    # 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
-    self.assertEqual(hi, 0x4700, f"Expected f16(7.0)=0x4700 in hi, got 0x{hi:04x}")
-
-  def test_v_add_f16_multilane(self):
-    """V_ADD_F16 with multiple lanes."""
-    instructions = [
-      s_mov_b32(s[0], 0x3c00),  # f16 1.0
-      s_mov_b32(s[1], 0x4000),  # f16 2.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_add_f16_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=4)
-    for lane in range(4):
-      result = st.vgpr[lane][2] & 0xffff
-      self.assertEqual(result, 0x4200, f"Lane {lane}: expected 0x4200, got 0x{result:04x}")
-
-
-class TestCndmask(unittest.TestCase):
-  """Tests for V_CNDMASK_B32 and V_CNDMASK_B16."""
-
-  def test_v_cndmask_b16_select_src0(self):
-    """V_CNDMASK_B16 selects src0 when VCC bit is 0."""
-    instructions = [
-      s_mov_b32(VCC_LO, 0),  # VCC = 0
-      s_mov_b32(s[0], 0x3c00),  # f16 1.0
-      s_mov_b32(s[1], 0x4000),  # f16 2.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_cndmask_b16(v[2], v[0], v[1], VCC),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    self.assertEqual(result, 0x3c00, f"Expected src0=0x3c00, got 0x{result:04x}")
-
-  def test_v_cndmask_b16_select_src1(self):
-    """V_CNDMASK_B16 selects src1 when VCC bit is 1."""
-    instructions = [
-      s_mov_b32(VCC_LO, 1),  # VCC = 1
-      s_mov_b32(s[0], 0x3c00),  # f16 1.0
-      s_mov_b32(s[1], 0x4000),  # f16 2.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_cndmask_b16(v[2], v[0], v[1], VCC),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    self.assertEqual(result, 0x4000, f"Expected src1=0x4000, got 0x{result:04x}")
-
-  def test_v_cndmask_b16_write_hi(self):
-    """V_CNDMASK_B16 can write to high 16 bits with opsel."""
-    instructions = [
-      s_mov_b32(s[0], 0x3c003800),  # src0: hi=1.0, lo=0.5
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], 0x4000c000),  # src1: hi=2.0, lo=-2.0
-      v_mov_b32_e32(v[1], s[1]),
-      s_mov_b32(s[2], 0xDEAD0000),  # v2 initial: hi=0xDEAD, lo=0
-      v_mov_b32_e32(v[2], s[2]),
-      s_mov_b32(VCC_LO, 0),  # vcc = 0, select src0
-      # opsel=0b1011: bit0=src0 hi, bit1=src1 hi, bit3=dst hi
-      VOP3(VOP3Op.V_CNDMASK_B16, vdst=v[2], src0=v[0], src1=v[1], src2=SrcEnum.VCC_LO, opsel=0b1011),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    hi = (st.vgpr[0][2] >> 16) & 0xffff
-    lo = st.vgpr[0][2] & 0xffff
-    # vcc=0 selects src0.h = 1.0 = 0x3c00, writes to hi
-    self.assertEqual(hi, 0x3c00, f"Expected hi=0x3c00 (1.0), got 0x{hi:04x}")
-    self.assertEqual(lo, 0x0000, f"Expected lo preserved as 0, got 0x{lo:04x}")
-
-
-class TestSpecialFloatValues(unittest.TestCase):
-  """Tests for special float value handling in VOP2 instructions."""
-
-  def test_neg_zero_add(self):
-    """-0.0 + 0.0 = +0.0 (IEEE 754)."""
-    neg_zero = 0x80000000
-    instructions = [
-      s_mov_b32(s[0], neg_zero),
-      v_mov_b32_e32(v[0], s[0]),
-      v_add_f32_e32(v[1], 0.0, v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
-
-  def test_neg_zero_mul(self):
-    """-0.0 * -1.0 = +0.0."""
-    neg_zero = 0x80000000
-    instructions = [
-      s_mov_b32(s[0], neg_zero),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mul_f32_e32(v[1], -1.0, v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
-
-  def test_inf_minus_inf(self):
-    """+inf - inf = NaN."""
-    import math
-    pos_inf = 0x7f800000
-    neg_inf = 0xff800000
-    instructions = [
-      s_mov_b32(s[0], pos_inf),
-      s_mov_b32(s[1], neg_inf),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_sub_f32_e32(v[2], v[0], v[1]),  # inf - (-inf) = inf
-      v_add_f32_e32(v[3], v[0], v[1]),  # inf + (-inf) = NaN
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], pos_inf, "inf - (-inf) = inf")
-    self.assertTrue(math.isnan(i2f(st.vgpr[0][3])), "inf + (-inf) = NaN")
-
-  def test_denormal_f32_mul_ftz(self):
-    """Denormal * normal - RDNA3 flushes denormals to zero (FTZ mode)."""
-    smallest_denorm = 0x00000001  # Smallest positive denormal
-    instructions = [
-      s_mov_b32(s[0], smallest_denorm),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mul_f32_e32(v[1], 2.0, v[0]),  # Denormal input gets flushed to 0
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][1], 0x00000000)
-
-
-if __name__ == '__main__':
-  unittest.main()

extra/assembly/amd/test/hw/test_vop3p.py

@@ -1,532 +0,0 @@
-"""Tests for VOP3P instructions - packed 16-bit vector operations.
-
-Includes: v_pk_add_f16, v_pk_mul_f16, v_pk_fma_f16, v_pack_b32_f16, v_wmma_*, v_dot2_*
-"""
-import unittest
-from extra.assembly.amd.test.hw.helpers import *
-
-class TestPackInstructions(unittest.TestCase):
-  """Tests for pack instructions."""
-
-  def test_v_pack_b32_f16(self):
-    """V_PACK_B32_F16 packs two f16 values into one 32-bit register."""
-    instructions = [
-      s_mov_b32(s[0], 0x3c00),  # f16 1.0
-      s_mov_b32(s[1], 0x4000),  # f16 2.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pack_b32_f16(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2]
-    self.assertEqual(result, 0x40003c00, f"Expected 0x40003c00, got 0x{result:08x}")
-
-  def test_v_pack_b32_f16_opsel_hi_hi(self):
-    """V_PACK_B32_F16 with opsel to read high halves."""
-    instructions = [
-      s_mov_b32(s[0], 0x40003c00),  # hi=2.0, lo=1.0
-      s_mov_b32(s[1], 0x44004200),  # hi=4.0, lo=3.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pack_b32_f16(v[2], v[0], v[1], opsel=0b0011),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2]
-    self.assertEqual(result, 0x44004000, f"Expected 0x44004000, got 0x{result:08x}")
-
-
-class TestPackMore(unittest.TestCase):
-  """Additional pack instruction tests."""
-
-  def test_v_pack_b32_f16_basic(self):
-    """V_PACK_B32_F16 packs two f16 values."""
-    instructions = [
-      s_mov_b32(s[0], 0x3c00),  # f16 1.0
-      s_mov_b32(s[1], 0x4000),  # f16 2.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pack_b32_f16(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2]
-    self.assertEqual(result, 0x40003c00, f"Expected 0x40003c00, got 0x{result:08x}")
-
-  def test_v_pack_b32_f16_with_cvt(self):
-    """V_PACK_B32_F16 after V_CVT_F16_F32 conversions."""
-    instructions = [
-      s_mov_b32(s[0], 0x3f800000),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[0]),
-      v_cvt_f16_f32_e32(v[2], v[0]),
-      v_cvt_f16_f32_e32(v[3], v[1]),
-      v_pack_b32_f16(v[4], v[2], v[3]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][4]
-    self.assertEqual(result, 0x3c003c00, f"Expected 0x3c003c00, got 0x{result:08x}")
-
-  def test_v_pack_b32_f16_packed_sources(self):
-    """V_PACK_B32_F16 with packed f16 sources (reads lo halves)."""
-    instructions = [
-      s_mov_b32(s[0], 0x40003c00),  # hi=2.0, lo=1.0
-      s_mov_b32(s[1], 0x44004200),  # hi=4.0, lo=3.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pack_b32_f16(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2]
-    # Expected: hi=v1.lo=0x4200 (3.0), lo=v0.lo=0x3c00 (1.0) -> 0x42003c00
-    self.assertEqual(result, 0x42003c00, f"Expected 0x42003c00, got 0x{result:08x}")
-
-  def test_v_pack_b32_f16_opsel_lo_hi(self):
-    """V_PACK_B32_F16 with opsel=0b0010 to read lo from src0, hi from src1."""
-    instructions = [
-      s_mov_b32(s[0], 0x40003c00),
-      s_mov_b32(s[1], 0x44004200),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pack_b32_f16(v[2], v[0], v[1], opsel=0b0010),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2]
-    self.assertEqual(result, 0x44003c00, f"Expected 0x44003c00, got 0x{result:08x}")
-
-  def test_v_pack_b32_f16_opsel_hi_lo(self):
-    """V_PACK_B32_F16 with opsel=0b0001 to read hi from src0, lo from src1."""
-    instructions = [
-      s_mov_b32(s[0], 0x40003c00),
-      s_mov_b32(s[1], 0x44004200),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pack_b32_f16(v[2], v[0], v[1], opsel=0b0001),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2]
-    self.assertEqual(result, 0x42004000, f"Expected 0x42004000, got 0x{result:08x}")
-
-  def test_v_pack_b32_f16_zeros(self):
-    """V_PACK_B32_F16 with zero values."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0),
-      v_mov_b32_e32(v[1], 0),
-      v_pack_b32_f16(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0)
-
-  def test_v_pack_b32_f16_both_positive(self):
-    """V_PACK_B32_F16 with positive f16 values."""
-    instructions = [
-      s_mov_b32(s[0], 0x4200),  # f16 3.0
-      s_mov_b32(s[1], 0x4400),  # f16 4.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pack_b32_f16(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2]
-    self.assertEqual(result, 0x44004200, f"Expected 0x44004200, got 0x{result:08x}")
-
-
-class TestFmaMix(unittest.TestCase):
-  """Tests for V_FMA_MIX_F32 and V_FMA_MIXLO_F16."""
-
-  def test_v_fma_mix_f32_all_f32_sources(self):
-    """V_FMA_MIX_F32 with all f32 sources."""
-    instructions = [
-      s_mov_b32(s[0], f2i(2.0)),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], f2i(3.0)),
-      v_mov_b32_e32(v[1], s[1]),
-      s_mov_b32(s[2], f2i(1.0)),
-      v_mov_b32_e32(v[2], s[2]),
-      VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = i2f(st.vgpr[0][3])
-    self.assertAlmostEqual(result, 7.0, places=5)
-
-  def test_v_fma_mix_f32_src2_f16_lo(self):
-    """V_FMA_MIX_F32 with src2 as f16 from lo bits."""
-    from extra.assembly.amd.pcode import f32_to_f16
-    f16_2 = f32_to_f16(2.0)
-    instructions = [
-      s_mov_b32(s[0], f2i(1.0)),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], f2i(3.0)),
-      v_mov_b32_e32(v[1], s[1]),
-      s_mov_b32(s[2], f16_2),
-      v_mov_b32_e32(v[2], s[2]),
-      VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = i2f(st.vgpr[0][3])
-    self.assertAlmostEqual(result, 5.0, places=5)
-
-  def test_v_fma_mix_f32_src2_f16_hi(self):
-    """V_FMA_MIX_F32 with src2 as f16 from hi bits."""
-    from extra.assembly.amd.pcode import f32_to_f16
-    f16_2 = f32_to_f16(2.0)
-    val = (f16_2 << 16) | 0
-    instructions = [
-      s_mov_b32(s[0], f2i(1.0)),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], f2i(3.0)),
-      v_mov_b32_e32(v[1], s[1]),
-      s_mov_b32(s[2], val),
-      v_mov_b32_e32(v[2], s[2]),
-      VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=4, opsel_hi=0, opsel_hi2=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = i2f(st.vgpr[0][3])
-    self.assertAlmostEqual(result, 5.0, places=5)
-
-  def test_v_fma_mix_f32_with_abs(self):
-    """V_FMA_MIX_F32 with abs modifier on src2."""
-    instructions = [
-      s_mov_b32(s[0], f2i(2.0)),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], f2i(3.0)),
-      v_mov_b32_e32(v[1], s[1]),
-      s_mov_b32(s[2], f2i(-1.0)),
-      v_mov_b32_e32(v[2], s[2]),
-      VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0, neg_hi=4),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = i2f(st.vgpr[0][3])
-    self.assertAlmostEqual(result, 7.0, places=5)
-
-  def test_v_fma_mixlo_f16(self):
-    """V_FMA_MIXLO_F16 writes to low 16 bits of destination."""
-    from extra.assembly.amd.pcode import _f16
-    instructions = [
-      s_mov_b32(s[0], f2i(2.0)),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], f2i(3.0)),
-      v_mov_b32_e32(v[1], s[1]),
-      s_mov_b32(s[2], f2i(1.0)),
-      v_mov_b32_e32(v[2], s[2]),
-      s_mov_b32(s[3], 0xdead0000),
-      v_mov_b32_e32(v[3], s[3]),
-      VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    lo = _f16(st.vgpr[0][3] & 0xffff)
-    hi = (st.vgpr[0][3] >> 16) & 0xffff
-    self.assertAlmostEqual(lo, 7.0, places=1)
-    self.assertEqual(hi, 0xdead, f"hi should be preserved, got 0x{hi:04x}")
-
-  def test_v_fma_mixlo_f16_all_f32_sources(self):
-    """V_FMA_MIXLO_F16 with all f32 sources."""
-    from extra.assembly.amd.pcode import _f16
-    instructions = [
-      s_mov_b32(s[0], f2i(1.0)),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], f2i(2.0)),
-      v_mov_b32_e32(v[1], s[1]),
-      s_mov_b32(s[2], f2i(3.0)),
-      v_mov_b32_e32(v[2], s[2]),
-      v_mov_b32_e32(v[3], 0),
-      VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    lo = _f16(st.vgpr[0][3] & 0xffff)
-    # 1*2+3 = 5
-    self.assertAlmostEqual(lo, 5.0, places=1)
-
-  def test_v_fma_mixlo_f16_sin_case(self):
-    """V_FMA_MIXLO_F16 case from sin kernel."""
-    from extra.assembly.amd.pcode import _f16
-    instructions = [
-      s_mov_b32(s[0], 0x3f800000),  # f32 1.0
-      v_mov_b32_e32(v[3], s[0]),
-      s_mov_b32(s[1], 0xaf05a309),  # f32 tiny negative
-      s_mov_b32(s[6], s[1]),
-      s_mov_b32(s[2], 0xc0490fdb),  # f32 -π
-      v_mov_b32_e32(v[5], s[2]),
-      s_mov_b32(s[3], 0x3f800000),
-      v_mov_b32_e32(v[3], s[3]),
-      VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[3], src1=s[6], src2=v[5], opsel=0, opsel_hi=0, opsel_hi2=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    lo = _f16(st.vgpr[0][3] & 0xffff)
-    self.assertAlmostEqual(lo, -3.14159, delta=0.01)
-
-
-class TestVOP3P(unittest.TestCase):
-  """Tests for VOP3P packed 16-bit operations."""
-
-  def test_v_pk_add_f16_basic(self):
-    """V_PK_ADD_F16 adds two packed f16 values."""
-    from extra.assembly.amd.pcode import _f16
-    instructions = [
-      s_mov_b32(s[0], 0x40003c00),  # hi=2.0, lo=1.0
-      s_mov_b32(s[1], 0x44004200),  # hi=4.0, lo=3.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pk_add_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
-    self.assertAlmostEqual(lo, 4.0, places=2)
-    self.assertAlmostEqual(hi, 6.0, places=2)
-
-  def test_v_pk_mul_f16_basic(self):
-    """V_PK_MUL_F16 multiplies two packed f16 values."""
-    from extra.assembly.amd.pcode import _f16
-    instructions = [
-      s_mov_b32(s[0], 0x42004000),  # hi=3.0, lo=2.0
-      s_mov_b32(s[1], 0x45004400),  # hi=5.0, lo=4.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pk_mul_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
-    self.assertAlmostEqual(lo, 8.0, places=1)
-    self.assertAlmostEqual(hi, 15.0, places=1)
-
-  def test_v_pk_fma_f16_basic(self):
-    """V_PK_FMA_F16: D = A * B + C for packed f16."""
-    from extra.assembly.amd.pcode import _f16
-    instructions = [
-      s_mov_b32(s[0], 0x42004000),  # A: hi=3.0, lo=2.0
-      s_mov_b32(s[1], 0x45004400),  # B: hi=5.0, lo=4.0
-      s_mov_b32(s[2], 0x3c003c00),  # C: hi=1.0, lo=1.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], s[2]),
-      v_pk_fma_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][3]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
-    self.assertAlmostEqual(lo, 9.0, places=1)   # 2*4+1
-    self.assertAlmostEqual(hi, 16.0, places=0)  # 3*5+1
-
-  def test_v_pk_add_f16_with_inline_constant(self):
-    """V_PK_ADD_F16 with inline constant POS_ONE (1.0).
-    Inline constants for VOP3P are f16 values in the low 16 bits only.
-    hi half of inline constant is 0, so hi result = v0.hi + 0 = 1.0.
-    """
-    from extra.assembly.amd.pcode import _f16
-    instructions = [
-      s_mov_b32(s[0], 0x3c003c00),  # packed f16: hi=1.0, lo=1.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_pk_add_f16(v[1], v[0], SrcEnum.POS_ONE, opsel_hi=3, opsel_hi2=1),  # Add inline constant 1.0
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][1]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
-    # lo = 1.0 + 1.0 = 2.0, hi = 1.0 + 0.0 = 1.0 (inline const hi half is 0)
-    self.assertAlmostEqual(lo, 2.0, places=2)
-    self.assertAlmostEqual(hi, 1.0, places=2)
-
-  def test_v_pk_mul_f16_with_inline_constant(self):
-    """V_PK_MUL_F16 with inline constant POS_TWO (2.0).
-    Inline constant has value only in low 16 bits, hi is 0.
-    """
-    from extra.assembly.amd.pcode import _f16
-    # v0 = packed (3.0, 4.0), multiply by POS_TWO
-    # lo = 3.0 * 2.0 = 6.0, hi = 4.0 * 0.0 = 0.0 (inline const hi is 0)
-    instructions = [
-      s_mov_b32(s[0], 0x44004200),  # packed f16: hi=4.0, lo=3.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_pk_mul_f16(v[1], v[0], SrcEnum.POS_TWO, opsel_hi=3, opsel_hi2=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][1]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
-    self.assertAlmostEqual(lo, 6.0, places=1)
-    self.assertAlmostEqual(hi, 0.0, places=1)
-
-
-class TestWMMA(unittest.TestCase):
-  """Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions."""
-
-  def test_v_wmma_f32_16x16x16_f16_all_ones(self):
-    """V_WMMA_F32_16X16X16_F16 with all ones produces 16.0."""
-    instructions = []
-    instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
-    for i in range(16, 32):
-      instructions.append(v_mov_b32_e32(v[i], s[0]))
-    for i in range(8):
-      instructions.append(v_mov_b32_e32(v[i], 0))
-    instructions.append(v_wmma_f32_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
-    st = run_program(instructions, n_lanes=32)
-    expected = f2i(16.0)
-    for lane in range(32):
-      for reg in range(8):
-        result = st.vgpr[lane][reg]
-        self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
-
-  def test_v_wmma_f32_16x16x16_f16_with_accumulator(self):
-    """V_WMMA_F32_16X16X16_F16 with non-zero accumulator."""
-    instructions = []
-    instructions.append(s_mov_b32(s[0], 0x3c003c00))
-    instructions.append(s_mov_b32(s[1], f2i(5.0)))
-    for i in range(16, 32):
-      instructions.append(v_mov_b32_e32(v[i], s[0]))
-    for i in range(8):
-      instructions.append(v_mov_b32_e32(v[i], s[1]))
-    instructions.append(v_wmma_f32_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
-    st = run_program(instructions, n_lanes=32)
-    expected = f2i(21.0)  # 16 + 5
-    for lane in range(32):
-      for reg in range(8):
-        result = st.vgpr[lane][reg]
-        self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
-
-
-class TestSpecialOps(unittest.TestCase):
-  """Tests for special operations (SAD, PERM, DOT2)."""
-
-  def test_v_sad_u8_basic(self):
-    """V_SAD_U8 computes sum of absolute differences."""
-    instructions = [
-      s_mov_b32(s[0], 0x04030201),  # bytes: 1, 2, 3, 4
-      s_mov_b32(s[1], 0x05040302),  # bytes: 2, 3, 4, 5
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], 0),
-      v_sad_u8(v[3], v[0], v[1], v[2]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # |1-2| + |2-3| + |3-4| + |4-5| = 1 + 1 + 1 + 1 = 4
-    self.assertEqual(st.vgpr[0][3], 4)
-
-  def test_v_sad_u8_identical_bytes(self):
-    """V_SAD_U8 with identical inputs returns accumulator."""
-    instructions = [
-      s_mov_b32(s[0], 0x04030201),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], 10),
-      v_mov_b32_e32(v[2], s[1]),
-      v_sad_u8(v[3], v[0], v[0], v[2]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # Same inputs -> SAD = 0, result = accumulator = 10
-    self.assertEqual(st.vgpr[0][3], 10)
-
-  def test_v_sad_u16_basic(self):
-    """V_SAD_U16 computes sum of absolute differences of u16 pairs."""
-    instructions = [
-      s_mov_b32(s[0], 0x00030001),  # hi=3, lo=1
-      s_mov_b32(s[1], 0x00050002),  # hi=5, lo=2
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], 0),
-      v_sad_u16(v[3], v[0], v[1], v[2]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # |1-2| + |3-5| = 1 + 2 = 3
-    self.assertEqual(st.vgpr[0][3], 3)
-
-  def test_v_sad_u32_basic(self):
-    """V_SAD_U32 computes absolute difference of u32 values."""
-    instructions = [
-      s_mov_b32(s[0], 100),
-      s_mov_b32(s[1], 70),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], 0),
-      v_sad_u32(v[3], v[0], v[1], v[2]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # |100-70| = 30
-    self.assertEqual(st.vgpr[0][3], 30)
-
-  def test_v_msad_u8_masked(self):
-    """V_MSAD_U8 masked SAD operation."""
-    instructions = [
-      s_mov_b32(s[0], 0x04030201),
-      s_mov_b32(s[1], 0x05040302),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], 0),
-      v_msad_u8(v[3], v[0], v[1], v[2]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # V_MSAD_U8 skips bytes where src0 is 0
-    # Since no bytes are 0, result same as V_SAD_U8 = 4
-    self.assertEqual(st.vgpr[0][3], 4)
-
-  def test_v_perm_b32_select_bytes(self):
-    """V_PERM_B32 selects bytes from two sources.
-
-    V_PERM_B32 concatenates {S1, S0} as a 64-bit value with S1 in low 32 bits.
-    Selector byte values 0-3 select from S1, values 4-7 select from S0.
-    """
-    instructions = [
-      s_mov_b32(s[0], 0x44332211),  # src0: bytes 4-7 in 64-bit view
-      s_mov_b32(s[1], 0x88776655),  # src1: bytes 0-3 in 64-bit view
-      s_mov_b32(s[2], 0x07060504),  # select bytes 4,5,6,7 (from src0)
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_perm_b32(v[2], v[0], v[1], s[2]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x44332211)
-
-  def test_v_dot2_f32_bf16_basic(self):
-    """V_DOT2_F32_BF16 computes dot product of bf16 pairs."""
-    # bf16 1.0 = 0x3f80, bf16 2.0 = 0x4000
-    instructions = [
-      s_mov_b32(s[0], 0x3f803f80),  # packed bf16: lo=1.0, hi=1.0
-      s_mov_b32(s[1], 0x40003f80),  # packed bf16: lo=1.0, hi=2.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], 0),
-      v_dot2_f32_bf16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # 1.0*1.0 + 1.0*2.0 + 0 = 3.0
-    result = i2f(st.vgpr[0][3])
-    self.assertAlmostEqual(result, 3.0, places=4)
-
-
-class TestPackedMixedSigns(unittest.TestCase):
-  """Tests for packed operations with mixed sign values."""
-
-  def test_pk_add_f16_mixed_signs(self):
-    """V_PK_ADD_F16 with mixed positive/negative values."""
-    from extra.assembly.amd.pcode import _f16
-    instructions = [
-      s_mov_b32(s[0], 0xc0003c00),  # packed: hi=-2.0, lo=1.0
-      s_mov_b32(s[1], 0x3c003c00),  # packed: hi=1.0, lo=1.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pk_add_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2]
-    lo = _f16(result & 0xffff)
-    hi = _f16((result >> 16) & 0xffff)
-    self.assertAlmostEqual(lo, 2.0, places=2)   # 1.0 + 1.0
-    self.assertAlmostEqual(hi, -1.0, places=2)  # -2.0 + 1.0
-
-  def test_pk_mul_f16_zero(self):
-    """V_PK_MUL_F16 with zero."""
-    from extra.assembly.amd.pcode import _f16
-    instructions = [
-      s_mov_b32(s[0], 0x40004000),  # packed: 2.0, 2.0
-      s_mov_b32(s[1], 0x00000000),  # packed: 0.0, 0.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pk_mul_f16(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2]
-    self.assertEqual(result, 0x00000000, "2.0 * 0.0 should be 0.0")
-
-
-if __name__ == '__main__':
-  unittest.main()

extra/assembly/amd/test/hw/test_vopc.py

@@ -1,476 +0,0 @@
-"""Tests for VOPC instructions - vector compare operations.
-
-Includes: v_cmp_class_f32, v_cmp_class_f16, v_cmp_eq_*, v_cmp_lt_*, v_cmp_gt_*
-"""
-import unittest
-from extra.assembly.amd.test.hw.helpers import *
-
-VCC = 106  # SGPR index for VCC_LO
-
-class TestCmpClass(unittest.TestCase):
-  """Tests for V_CMP_CLASS_F32 float classification."""
-
-  def test_cmp_class_quiet_nan(self):
-    """V_CMP_CLASS_F32 detects quiet NaN."""
-    quiet_nan = 0x7fc00000
-    instructions = [
-      s_mov_b32(s[0], quiet_nan),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0b0000000010),  # bit 1 = quiet NaN
-      v_cmp_class_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
-
-  def test_cmp_class_signaling_nan(self):
-    """V_CMP_CLASS_F32 detects signaling NaN."""
-    signal_nan = 0x7f800001
-    instructions = [
-      s_mov_b32(s[0], signal_nan),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0b0000000001),  # bit 0 = signaling NaN
-      v_cmp_class_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
-
-  def test_cmp_class_positive_inf(self):
-    """V_CMP_CLASS_F32 detects +inf."""
-    pos_inf = 0x7f800000
-    instructions = [
-      s_mov_b32(s[0], pos_inf),
-      s_mov_b32(s[1], 0b1000000000),  # bit 9 = +inf
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_cmp_class_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect +inf")
-
-  def test_cmp_class_negative_inf(self):
-    """V_CMP_CLASS_F32 detects -inf."""
-    neg_inf = 0xff800000
-    instructions = [
-      s_mov_b32(s[0], neg_inf),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0b0000000100),  # bit 2 = -inf
-      v_cmp_class_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect -inf")
-
-  def test_cmp_class_normal_positive(self):
-    """V_CMP_CLASS_F32 detects positive normal."""
-    instructions = [
-      v_mov_b32_e32(v[0], 1.0),
-      s_mov_b32(s[1], 0b0100000000),  # bit 8 = positive normal
-      v_mov_b32_e32(v[1], s[1]),
-      v_cmp_class_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
-
-  def test_cmp_class_normal_negative(self):
-    """V_CMP_CLASS_F32 detects negative normal."""
-    instructions = [
-      v_mov_b32_e32(v[0], -1.0),
-      v_mov_b32_e32(v[1], 0b0000001000),  # bit 3 = negative normal
-      v_cmp_class_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
-
-  def test_cmp_class_quiet_nan_not_signaling(self):
-    """Quiet NaN does not match signaling NaN mask."""
-    quiet_nan = 0x7fc00000
-    instructions = [
-      s_mov_b32(s[0], quiet_nan),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0b0000000001),  # bit 0 = signaling NaN only
-      v_cmp_class_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 0, "Quiet NaN should not match signaling mask")
-
-  def test_cmp_class_signaling_nan_not_quiet(self):
-    """Signaling NaN does not match quiet NaN mask."""
-    signal_nan = 0x7f800001
-    instructions = [
-      s_mov_b32(s[0], signal_nan),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0b0000000010),  # bit 1 = quiet NaN only
-      v_cmp_class_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 0, "Signaling NaN should not match quiet mask")
-
-  def test_v_cmp_sets_vcc_bits(self):
-    """V_CMP_EQ sets VCC bits based on per-lane comparison."""
-    instructions = [
-      s_mov_b32(s[0], 5),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[0]),
-      v_cmp_eq_u32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=4)
-    self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
-
-
-class TestCmpClassF16(unittest.TestCase):
-  """Tests for V_CMP_CLASS_F16 float classification.
-
-  Class bit mapping:
-    bit 0 = signaling NaN
-    bit 1 = quiet NaN
-    bit 2 = -infinity
-    bit 3 = -normal
-    bit 4 = -denormal
-    bit 5 = -zero
-    bit 6 = +zero
-    bit 7 = +denormal
-    bit 8 = +normal
-    bit 9 = +infinity
-  """
-
-  def test_cmp_class_f16_positive_zero(self):
-    """V_CMP_CLASS_F16: +zero matches bit 6."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0x0000),  # f16 +0.0
-      v_mov_b32_e32(v[1], 0x40),     # bit 6 = +zero
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect positive zero")
-
-  def test_cmp_class_f16_negative_zero(self):
-    """V_CMP_CLASS_F16: -zero matches bit 5."""
-    instructions = [
-      s_mov_b32(s[0], 0x8000),       # f16 -0.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0x20),     # bit 5 = -zero
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect negative zero")
-
-  def test_cmp_class_f16_positive_normal(self):
-    """V_CMP_CLASS_F16: +1.0 (normal) matches bit 8."""
-    instructions = [
-      s_mov_b32(s[0], 0x3c00),       # f16 +1.0
-      s_mov_b32(s[1], 0x100),        # bit 8 = +normal
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
-
-  def test_cmp_class_f16_negative_normal(self):
-    """V_CMP_CLASS_F16: -1.0 (normal) matches bit 3."""
-    instructions = [
-      s_mov_b32(s[0], 0xbc00),       # f16 -1.0
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0x08),     # bit 3 = -normal
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
-
-  def test_cmp_class_f16_positive_infinity(self):
-    """V_CMP_CLASS_F16: +inf matches bit 9."""
-    instructions = [
-      s_mov_b32(s[0], 0x7c00),       # f16 +inf
-      s_mov_b32(s[1], 0x200),        # bit 9 = +inf
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect positive infinity")
-
-  def test_cmp_class_f16_negative_infinity(self):
-    """V_CMP_CLASS_F16: -inf matches bit 2."""
-    instructions = [
-      s_mov_b32(s[0], 0xfc00),       # f16 -inf
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0x04),     # bit 2 = -inf
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect negative infinity")
-
-  def test_cmp_class_f16_quiet_nan(self):
-    """V_CMP_CLASS_F16: quiet NaN matches bit 1."""
-    instructions = [
-      s_mov_b32(s[0], 0x7e00),       # f16 quiet NaN
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0x02),     # bit 1 = quiet NaN
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
-
-  def test_cmp_class_f16_signaling_nan(self):
-    """V_CMP_CLASS_F16: signaling NaN matches bit 0."""
-    instructions = [
-      s_mov_b32(s[0], 0x7c01),       # f16 signaling NaN
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0x01),     # bit 0 = signaling NaN
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
-
-  def test_cmp_class_f16_positive_denormal(self):
-    """V_CMP_CLASS_F16: positive denormal matches bit 7."""
-    instructions = [
-      v_mov_b32_e32(v[0], 1),        # f16 +denormal (0x0001)
-      v_mov_b32_e32(v[1], 0x80),     # bit 7 = +denormal
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect positive denormal")
-
-  def test_cmp_class_f16_negative_denormal(self):
-    """V_CMP_CLASS_F16: negative denormal matches bit 4."""
-    instructions = [
-      s_mov_b32(s[0], 0x8001),       # f16 -denormal
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], 0x10),     # bit 4 = -denormal
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Should detect negative denormal")
-
-  def test_cmp_class_f16_combined_mask_zeros(self):
-    """V_CMP_CLASS_F16: mask 0x60 covers both +zero and -zero."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0),         # f16 +0.0
-      v_mov_b32_e32(v[1], 0x60),      # bits 5 and 6 (+-zero)
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x60")
-
-  def test_cmp_class_f16_combined_mask_1f8(self):
-    """V_CMP_CLASS_F16: mask 0x1f8 covers -normal,-denorm,-zero,+zero,+denorm,+normal.
-
-    This is the exact mask used in the f16 sin kernel at PC=46.
-    """
-    instructions = [
-      v_mov_b32_e32(v[0], 0),         # f16 +0.0
-      s_mov_b32(s[0], 0x1f8),
-      v_mov_b32_e32(v[1], s[0]),      # mask 0x1f8
-      v_cmp_class_f16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x1f8")
-
-  def test_cmp_class_f16_vop3_encoding(self):
-    """V_CMP_CLASS_F16 in VOP3 encoding (v_cmp_class_f16_e64)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0),         # f16 +0.0
-      s_mov_b32(s[0], 0x1f8),         # class mask
-      v_cmp_class_f16_e64(VCC_LO, v[0], s[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with VOP3 encoding")
-
-  def test_cmp_class_f16_vop3_normal_positive(self):
-    """V_CMP_CLASS_F16 VOP3 encoding with +1.0 (normal)."""
-    instructions = [
-      s_mov_b32(s[0], 0x3c00),        # f16 +1.0
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], 0x1f8),         # class mask
-      v_cmp_class_f16_e64(VCC_LO, v[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +1.0 (normal) with mask 0x1f8")
-
-  def test_cmp_class_f16_vop3_nan_fails_mask(self):
-    """V_CMP_CLASS_F16 VOP3: NaN should NOT match mask 0x1f8 (no NaN bits set)."""
-    instructions = [
-      s_mov_b32(s[0], 0x7e00),        # f16 quiet NaN
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], 0x1f8),         # class mask
-      v_cmp_class_f16_e64(VCC_LO, v[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for NaN with mask 0x1f8 (no NaN bits)")
-
-  def test_cmp_class_f16_vop3_inf_fails_mask(self):
-    """V_CMP_CLASS_F16 VOP3: +inf should NOT match mask 0x1f8 (no inf bits set)."""
-    instructions = [
-      s_mov_b32(s[0], 0x7c00),        # f16 +inf
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], 0x1f8),         # class mask
-      v_cmp_class_f16_e64(VCC_LO, v[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for +inf with mask 0x1f8 (no inf bits)")
-
-
-class TestCmpInt(unittest.TestCase):
-  """Tests for integer comparison operations."""
-
-  def test_v_cmp_eq_u32(self):
-    """V_CMP_EQ_U32 sets VCC bits based on per-lane comparison."""
-    instructions = [
-      s_mov_b32(s[0], 5),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[0]),
-      v_cmp_eq_u32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=4)
-    self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
-
-  def test_cmp_eq_u16_opsel_lo_lo(self):
-    """V_CMP_EQ_U16 comparing lo halves."""
-    instructions = [
-      s_mov_b32(s[0], 0x12340005),  # lo=5, hi=0x1234
-      s_mov_b32(s[1], 0xABCD0005),  # lo=5, hi=0xABCD
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_cmp_eq_u16_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Lo halves should be equal")
-
-  def test_cmp_eq_u16_opsel_hi_hi(self):
-    """V_CMP_EQ_U16 comparing hi halves with VOP3 opsel."""
-    instructions = [
-      s_mov_b32(s[2], 0x00051234),  # hi=5, lo=0x1234
-      v_mov_b32_e32(v[0], s[2]),
-      s_mov_b32(s[2], 0x0005ABCD),  # hi=5, lo=0xABCD
-      v_mov_b32_e32(v[1], s[2]),
-      v_cmp_eq_u16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[0] & 1, 1, "Hi halves should be equal: 5==5")
-
-  def test_cmp_eq_u16_opsel_hi_hi_equal(self):
-    """V_CMP_EQ_U16 VOP3 with opsel=3 compares hi halves (equal case)."""
-    instructions = [
-      s_mov_b32(s[2], 0x12340005),  # lo=5, hi=0x1234
-      v_mov_b32_e32(v[0], s[2]),
-      s_mov_b32(s[2], 0x12340009),  # lo=9, hi=0x1234
-      v_mov_b32_e32(v[1], s[2]),
-      v_cmp_eq_u16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[0] & 1, 1, "hi==hi should be true: 0x1234==0x1234")
-
-  def test_cmp_gt_u16_opsel_hi(self):
-    """V_CMP_GT_U16 VOP3 with opsel=3 compares hi halves."""
-    instructions = [
-      s_mov_b32(s[2], 0x99990005),  # lo=5, hi=0x9999
-      v_mov_b32_e32(v[0], s[2]),
-      s_mov_b32(s[2], 0x12340005),  # lo=5, hi=0x1234
-      v_mov_b32_e32(v[1], s[2]),
-      v_cmp_gt_u16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[0] & 1, 1, "hi>hi should be true: 0x9999>0x1234")
-
-
-class TestCmpFloat(unittest.TestCase):
-  """Tests for float comparison operations."""
-
-  def test_v_cmp_lt_f16_vsrc1_hi(self):
-    """V_CMP_LT_F16 with both operands from high half using VOP3 opsel."""
-    instructions = [
-      s_mov_b32(s[2], 0x3c000000),  # hi=1.0 (f16), lo=0
-      v_mov_b32_e32(v[0], s[2]),
-      s_mov_b32(s[2], 0x40000000),  # hi=2.0 (f16), lo=0
-      v_mov_b32_e32(v[1], s[2]),
-      v_cmp_lt_f16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[0] & 1, 1, "1.0 < 2.0 should be true")
-
-  def test_v_cmp_gt_f16_vsrc1_hi(self):
-    """V_CMP_GT_F16 with both operands from high half using VOP3 opsel."""
-    instructions = [
-      s_mov_b32(s[2], 0x40000000),  # hi=2.0 (f16), lo=0
-      v_mov_b32_e32(v[0], s[2]),
-      s_mov_b32(s[2], 0x3c000000),  # hi=1.0 (f16), lo=0
-      v_mov_b32_e32(v[1], s[2]),
-      v_cmp_gt_f16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[0] & 1, 1, "2.0 > 1.0 should be true")
-
-  def test_v_cmp_eq_f16_vsrc1_hi_equal(self):
-    """v_cmp_eq_f16 with equal low and high halves."""
-    instructions = [
-      s_mov_b32(s[0], 0x42004200),  # hi=3.0 (0x4200), lo=3.0 (0x4200)
-      v_mov_b32_e32(v[0], s[0]),
-      v_cmp_eq_f16_e32(v[0], v[0].h),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (3.0 == 3.0)")
-
-  def test_v_cmp_neq_f16_vsrc1_hi(self):
-    """v_cmp_neq_f16 with different low and high halves."""
-    instructions = [
-      s_mov_b32(s[0], 0x40003c00),  # hi=2.0 (0x4000), lo=1.0 (0x3c00)
-      v_mov_b32_e32(v[0], s[0]),
-      v_cmp_lg_f16_e32(v[0], v[0].h),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (1.0 != 2.0)")
-
-  def test_v_cmp_nge_f16_inf_self(self):
-    """v_cmp_nge_f16 comparing -inf with itself (unordered less than).
-
-    Regression test: -inf < -inf should be false (IEEE 754).
-    """
-    instructions = [
-      s_mov_b32(s[0], 0xFC00FC00),  # both halves = -inf (0xFC00)
-      v_mov_b32_e32(v[0], s[0]),
-      v_cmp_nge_f16_e32(v[0], v[0].h),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 0, "Expected vcc=0 (-inf >= -inf)")
-
-  def test_v_cmp_f16_multilane(self):
-    """v_cmp_lt_f16 with vsrc1=v128 across multiple lanes."""
-    instructions = [
-      # Lane 0: v0 = 0x40003c00 (hi=2.0, lo=1.0) -> 1.0 < 2.0 = true
-      # Lane 1: v0 = 0x3c004000 (hi=1.0, lo=2.0) -> 2.0 < 1.0 = false
-      v_mov_b32_e32(v[0], 0x40003c00),  # default
-      v_cmp_eq_u32_e32(1, v[255]),  # vcc = (lane == 1)
-      v_cndmask_b32_e64(v[0], v[0], 0x3c004000, SrcEnum.VCC_LO),
-      v_cmp_lt_f16_e32(v[0], v[0].h),
-    ]
-    st = run_program(instructions, n_lanes=2)
-    self.assertEqual(st.vcc & 1, 1, "Lane 0: expected vcc=1 (1.0 < 2.0)")
-    self.assertEqual((st.vcc >> 1) & 1, 0, "Lane 1: expected vcc=0 (2.0 < 1.0)")
-
-
-class TestVCCBehavior(unittest.TestCase):
-  """Tests for VCC condition code behavior."""
-
-  def test_vcc_all_lanes_true(self):
-    """VCC should have all bits set when all lanes compare true."""
-    instructions = [
-      v_mov_b32_e32(v[0], 5),
-      v_mov_b32_e32(v[1], 5),
-      v_cmp_eq_u32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=32)
-    self.assertEqual(st.vcc, 0xFFFFFFFF, "All 32 lanes should be true")
-
-  def test_vcc_lane_dependent(self):
-    """VCC should differ per lane based on lane_id comparison."""
-    instructions = [
-      v_mov_b32_e32(v[0], 16),
-      v_cmp_lt_u32_e32(v[255], v[0]),  # lanes 0-15 are < 16
-    ]
-    st = run_program(instructions, n_lanes=32)
-    self.assertEqual(st.vcc & 0xFFFF, 0xFFFF, "Lanes 0-15 should be true")
-    self.assertEqual(st.vcc >> 16, 0x0000, "Lanes 16-31 should be false")
-
-
-if __name__ == '__main__':
-  unittest.main()

extra/assembly/amd/test/test_custom_kernel.py

@@ -1,48 +0,0 @@
-import unittest
-import functools
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.dsl import s, v, Inst
-
-def assemble_insts(insts:list[Inst], name:str, arch:str, kernarg_size:int=8) -> tuple[UOp, UOp]:
-  kd = {"kernarg_size":kernarg_size, "user_sgpr_kernarg_segment_ptr":1, "next_free_vgpr":8, "next_free_sgpr":8, "wavefront_size32":1}
-  disasm = "\n".join([inst.disasm() for inst in insts])
-  hsasrc = f".text\n.globl {name}\n.p2align 8\n.type fn_name,@function\n{name}:\n{disasm}\ns_code_end\n"
-  hsasrc += f".rodata\n.p2align 6\n.amdhsa_kernel {name}\n"+"\n".join([f".amdhsa_{k} {v}" for k,v in kd.items()])+"\n.end_amdhsa_kernel"
-  binary = HIPCompiler(arch).compile(hsasrc)
-  return UOp(Ops.SOURCE, arg=disasm), UOp(Ops.BINARY, arg=binary)
-
-def custom_add_one(A:UOp, arch:str) -> UOp:
-  A = A.flatten()
-  assert dtypes.is_float(A.dtype.base), f"buffer dtype must be float32, got {A.dtype}"
-  threads = UOp.special(A.size, "lidx0")
-  insts = [
-    s_load_b64(s[0:1], s[0:1], soffset=NULL),
-    s_waitcnt(lgkmcnt=0),
-    v_lshlrev_b32_e32(v[0], 2, v[0]), # element offset
-    global_load_b32(v[1], v[0], saddr=s[0:1]),
-    s_waitcnt(vmcnt=0),
-    v_mov_b32_e32(v[2], 1.0),
-    v_add_f32_e32(v[1], v[1], v[2]),
-    global_store_b32(addr=v[0], data=v[1], saddr=s[0:1]),
-    s_endpgm(),
-  ]
-  sink = UOp.sink(A.base, threads, arg=KernelInfo(name:=f"custom_add_one_{A.size}", estimates=Estimates(ops=A.size, mem=A.size*4*2)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=(*sink.src, sink)), *assemble_insts(insts, name, arch)))
-
-class TestCustomKernel(unittest.TestCase):
-  def test_simple(self):
-    a = Tensor.full((16, 16), 1.).contiguous().realize()
-    a = Tensor.custom_kernel(a, fxn=functools.partial(custom_add_one, arch=Device[Device.DEFAULT].arch))[0]
-    ei = a.schedule()[-1].lower()
-    self.assertEqual(ei.prg.estimates.ops, a.numel())
-    self.assertEqual(ei.prg.estimates.mem, a.nbytes()*2)
-    ei.run()
-    self.assertTrue((a.numpy() == 2.).all())
-
-if __name__ == "__main__":
-  unittest.main()

extra/assembly/amd/test/test_integration.py

@@ -1,258 +0,0 @@
-#!/usr/bin/env python3
-"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
-import unittest, io, sys
-from extra.assembly.amd.autogen.rdna3.ins import *
-
-def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
-  return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
-
-def disassemble(lib: bytes, arch: str = "gfx1100") -> str:
-  """Disassemble ELF binary using tinygrad's compiler, return raw output."""
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-  old_stdout = sys.stdout
-  sys.stdout = io.StringIO()
-  HIPCompiler(arch).disassemble(lib)
-  output = sys.stdout.getvalue()
-  sys.stdout = old_stdout
-  return output
-
-def parse_disassembly(raw: str) -> list[str]:
-  """Parse disassembly output to list of instruction mnemonics."""
-  lines = []
-  for line in raw.splitlines():
-    if line.startswith('\t'):
-      instr = line.split('//')[0].strip()
-      if instr: lines.append(instr)
-  return lines
-
-def assemble_and_disassemble(instructions: list, arch: str = "gfx1100") -> list[str]:
-  """Assemble instructions with our DSL, then disassemble with AMD toolchain."""
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-  # Generate bytes from our DSL
-  code_bytes = b''.join(inst.to_bytes() for inst in instructions)
-
-  # Wrap in minimal ELF-compatible assembly with .byte directives
-  byte_str = ', '.join(f'0x{b:02x}' for b in code_bytes)
-  asm_src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n.byte {byte_str}\n"
-
-  # Assemble with AMD COMGR and disassemble
-  lib = HIPCompiler(arch).compile(asm_src)
-  return parse_disassembly(disassemble(lib, arch))
-
-class TestIntegration(unittest.TestCase):
-  """Test our DSL 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:1], NULL),
-      s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
-      global_store_b32(addr=v[0:1], 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:1], 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:1], data=v[2], saddr=OFF),
-      s_endpgm(),
-    ]
-    disasm = assemble_and_disassemble(instructions)
-    # Verify key instructions are present
-    self.assertTrue(any('global_load' in d for d in disasm))
-    self.assertTrue(any('v_add_f32' in d for d in disasm))
-    self.assertTrue(any('global_store' in d for d in disasm))
-    self.assertTrue(any('s_endpgm' in d for d in disasm))
-
-  def test_bytes_roundtrip(self):
-    """Test that our bytes match what AMD assembler produces."""
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-    # Simple instruction
-    inst = s_mov_b32(s[0], s[1])
-    our_bytes = inst.to_bytes()
-
-    # Assemble same instruction with AMD toolchain
-    asm_src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\ns_mov_b32 s0, s1\n"
-    compiler = HIPCompiler("gfx1100")
-    lib = compiler.compile(asm_src)
-    raw = disassemble(lib)
-
-    for line in raw.splitlines():
-      if 's_mov_b32' in line and '//' in line:
-        # Extract hex bytes from comment: "// 000000001300: BE800001"
-        comment = line.split('//')[1].strip()
-        hex_str = comment.split(':')[1].strip()
-        # Convert big-endian hex string to little-endian bytes
-        amd_bytes = bytes.fromhex(hex_str)[::-1]  # reverse for little-endian
-        self.assertEqual(our_bytes, amd_bytes, f"Bytes mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
-        return
-    self.fail("Could not find s_mov_b32 in disassembly")
-
-class 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/amd/test/test_pcode.py

@@ -1,403 +0,0 @@
-#!/usr/bin/env python3
-"""Tests for the RDNA3 pseudocode DSL."""
-import unittest
-from extra.assembly.amd.pcode import (Reg, TypedView, TypedView, MASK32, MASK64,
-                                       _f32, _i32, _f16, _i16, f32_to_f16, isNAN, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
-                                       BYTE_PERMUTE, v_sad_u8, v_msad_u8, _compile_pseudocode, _expr, compile_pseudocode)
-from extra.assembly.amd.test.helpers import ExecContext
-from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
-from extra.assembly.amd.autogen.rdna3.enum import VOP3SDOp, VOPCOp
-
-# Compile pseudocode functions on demand for regression tests
-_VOP3SDOp_V_DIV_SCALE_F32 = compile_pseudocode('VOP3SDOp', 'V_DIV_SCALE_F32', PCODE[VOP3SDOp.V_DIV_SCALE_F32])
-_VOPCOp_V_CMP_CLASS_F32 = compile_pseudocode('VOPCOp', 'V_CMP_CLASS_F32', PCODE[VOPCOp.V_CMP_CLASS_F32_E32])
-
-class TestReg(unittest.TestCase):
-  def test_u32_read(self):
-    r = Reg(0xDEADBEEF)
-    self.assertEqual(int(r.u32), 0xDEADBEEF)
-
-  def test_u32_write(self):
-    r = Reg(0)
-    r.u32 = 0x12345678
-    self.assertEqual(r._val, 0x12345678)
-
-  def test_f32_read(self):
-    r = Reg(0x40400000)  # 3.0f
-    self.assertAlmostEqual(float(r.f32), 3.0)
-
-  def test_f32_write(self):
-    r = Reg(0)
-    r.f32 = 3.0
-    self.assertEqual(r._val, 0x40400000)
-
-  def test_i32_signed(self):
-    r = Reg(0xFFFFFFFF)  # -1 as signed
-    self.assertEqual(int(r.i32), -1)
-
-  def test_u64(self):
-    r = Reg(0xDEADBEEFCAFEBABE)
-    self.assertEqual(int(r.u64), 0xDEADBEEFCAFEBABE)
-
-  def test_f64(self):
-    r = Reg(0x4008000000000000)  # 3.0 as f64
-    self.assertAlmostEqual(float(r.f64), 3.0)
-
-class TestTypedView(unittest.TestCase):
-  def test_bit_slice(self):
-    r = Reg(0xDEADBEEF)
-    # Slices return TypedView which supports .u32, .u16 etc (matching pseudocode like S1.u32[1:0].u32)
-    self.assertEqual(r.u32[7:0].u32, 0xEF)
-    self.assertEqual(r.u32[15:8].u32, 0xBE)
-    self.assertEqual(r.u32[23:16].u32, 0xAD)
-    self.assertEqual(r.u32[31:24].u32, 0xDE)
-    # Also works with int() for arithmetic
-    self.assertEqual(int(r.u32[7:0]), 0xEF)
-
-  def test_single_bit_read(self):
-    r = Reg(0b11010101)
-    self.assertEqual(r.u32[0], 1)
-    self.assertEqual(r.u32[1], 0)
-    self.assertEqual(r.u32[2], 1)
-    self.assertEqual(r.u32[3], 0)
-
-  def test_single_bit_write(self):
-    r = Reg(0)
-    r.u32[5] = 1
-    r.u32[3] = 1
-    self.assertEqual(r._val, 0b00101000)
-
-  def test_nested_bit_access(self):
-    # S0.u32[S1.u32[4:0]] - access bit at position from another register
-    s0 = Reg(0b11010101)
-    s1 = Reg(3)
-    bit_pos = s1.u32[4:0]  # TypedView, int value = 3
-    bit_val = s0.u32[int(bit_pos)]  # bit 3 of s0 = 0
-    self.assertEqual(int(bit_pos), 3)
-    self.assertEqual(bit_val, 0)
-
-  def test_arithmetic(self):
-    r1 = Reg(0x40400000)  # 3.0f
-    r2 = Reg(0x40800000)  # 4.0f
-    result = r1.f32 + r2.f32
-    self.assertAlmostEqual(result, 7.0)
-
-  def test_comparison(self):
-    r1 = Reg(5)
-    r2 = Reg(3)
-    self.assertTrue(r1.u32 > r2.u32)
-    self.assertFalse(r1.u32 < r2.u32)
-    self.assertTrue(r1.u32 != r2.u32)
-
-class TestTypedView(unittest.TestCase):
-  def test_slice_read(self):
-    r = Reg(0x56781234)
-    self.assertEqual(r[15:0].u16, 0x1234)
-    self.assertEqual(r[31:16].u16, 0x5678)
-
-  def test_slice_write(self):
-    r = Reg(0)
-    r[15:0].u16 = 0x1234
-    r[31:16].u16 = 0x5678
-    self.assertEqual(r._val, 0x56781234)
-
-  def test_slice_f16(self):
-    r = Reg(0)
-    r[15:0].f16 = 3.0
-    self.assertAlmostEqual(_f16(r._val & 0xffff), 3.0, places=2)
-
-class TestCompiler(unittest.TestCase):
-  def test_ternary(self):
-    result = _expr("a > b ? 1 : 0")
-    self.assertIn("if", result)
-    self.assertIn("else", result)
-
-  def test_type_prefix_strip(self):
-    self.assertEqual(_expr("1'0U"), "0")
-    self.assertEqual(_expr("32'1"), "1")
-    self.assertEqual(_expr("16'0xFFFF"), "0xFFFF")
-
-  def test_suffix_strip(self):
-    self.assertEqual(_expr("0ULL"), "0")
-    self.assertEqual(_expr("1LL"), "1")
-    self.assertEqual(_expr("5U"), "5")
-    self.assertEqual(_expr("3.14F"), "3.14")
-
-  def test_boolean_ops(self):
-    self.assertIn("and", _expr("a && b"))
-    self.assertIn("or", _expr("a || b"))
-    self.assertIn("!=", _expr("a <> b"))
-
-  def test_pack16(self):
-    result = _expr("{ a, b }")
-    self.assertIn("_pack", result)
-
-  def test_type_cast_strip(self):
-    self.assertEqual(_expr("64'U(x)"), "(x)")
-    self.assertEqual(_expr("32'I(y)"), "(y)")
-
-class TestExecContext(unittest.TestCase):
-  def test_float_add(self):
-    ctx = ExecContext(s0=0x40400000, s1=0x40800000)  # 3.0f, 4.0f
-    ctx.D0.f32 = ctx.S0.f32 + ctx.S1.f32
-    self.assertAlmostEqual(_f32(ctx.D0._val), 7.0)
-
-  def test_float_mul(self):
-    ctx = ExecContext(s0=0x40400000, s1=0x40800000)  # 3.0f, 4.0f
-    ctx.run("D0.f32 = S0.f32 * S1.f32")
-    self.assertAlmostEqual(_f32(ctx.D0._val), 12.0)
-
-  def test_scc_comparison(self):
-    ctx = ExecContext(s0=42, s1=42)
-    ctx.run("SCC = S0.u32 == S1.u32")
-    self.assertEqual(ctx.SCC._val, 1)
-
-  def test_scc_comparison_false(self):
-    ctx = ExecContext(s0=42, s1=43)
-    ctx.run("SCC = S0.u32 == S1.u32")
-    self.assertEqual(ctx.SCC._val, 0)
-
-  def test_ternary(self):
-    code = _compile_pseudocode("D0.u32 = S0.u32 > S1.u32 ? 1'1U : 1'0U")
-    ctx = ExecContext(s0=5, s1=3)
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val, 1)
-
-  def test_pack(self):
-    code = _compile_pseudocode("D0 = { S1[15:0].u16, S0[15:0].u16 }")
-    ctx = ExecContext(s0=0x1234, s1=0x5678)
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val, 0x56781234)
-
-  def test_tmp_with_typed_access(self):
-    code = _compile_pseudocode("""tmp = S0.u32 + S1.u32
-D0.u32 = tmp.u32""")
-    ctx = ExecContext(s0=100, s1=200)
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val, 300)
-
-  def test_s_add_u32_pattern(self):
-    # Real pseudocode pattern from S_ADD_U32
-    code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
-SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
-D0.u32 = tmp.u32""")
-    # Test overflow case
-    ctx = ExecContext(s0=0xFFFFFFFF, s1=0x00000001)
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val, 0)  # Wraps to 0
-    self.assertEqual(ctx.SCC._val, 1)  # Carry set
-
-  def test_s_add_u32_no_overflow(self):
-    code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
-SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
-D0.u32 = tmp.u32""")
-    ctx = ExecContext(s0=100, s1=200)
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val, 300)
-    self.assertEqual(ctx.SCC._val, 0)  # No carry
-
-  def test_vcc_lane_read(self):
-    ctx = ExecContext(vcc=0b1010, lane=1)
-    # Lane 1 is set
-    self.assertEqual(ctx.VCC.u64[1], 1)
-    self.assertEqual(ctx.VCC.u64[2], 0)
-
-  def test_vcc_lane_write(self):
-    ctx = ExecContext(vcc=0, lane=0)
-    ctx.VCC.u64[3] = 1
-    ctx.VCC.u64[1] = 1
-    self.assertEqual(ctx.VCC._val, 0b1010)
-
-  def test_for_loop(self):
-    # CTZ pattern - find first set bit
-    code = _compile_pseudocode("""tmp = -1
-for i in 0 : 31 do
-  if S0.u32[i] == 1 then
-    tmp = i
-  endif
-endfor
-D0.i32 = tmp""")
-    ctx = ExecContext(s0=0b1000)  # Bit 3 is set
-    ctx.run(code)
-    self.assertEqual(ctx.D0._val & MASK32, 3)
-
-  def test_result_dict(self):
-    ctx = ExecContext(s0=5, s1=3)
-    ctx.D0.u32 = 42
-    ctx.SCC._val = 1
-    result = ctx.result()
-    self.assertEqual(result['d0'], 42)
-    self.assertEqual(result['scc'], 1)
-
-class TestPseudocodeRegressions(unittest.TestCase):
-  """Regression tests for pseudocode instruction emulation bugs."""
-
-  def test_v_div_scale_f32_vcc_always_returned(self):
-    """V_DIV_SCALE_F32 must always return VCC, even when VCC=0 (no scaling needed).
-    Bug: when VCC._val == vcc (both 0), VCC wasn't returned, so VCC bits weren't written.
-    This caused division to produce wrong results for multiple lanes."""
-    # Normal case: 1.0 / 3.0, no scaling needed, VCC should be 0
-    s0 = 0x3f800000  # 1.0
-    s1 = 0x40400000  # 3.0
-    s2 = 0x3f800000  # 1.0 (numerator)
-    result = _VOP3SDOp_V_DIV_SCALE_F32(s0, s1, s2, 0, 0, 0, 0, 0xffffffff, 0, None)
-    # Must always have VCC in result
-    self.assertIn('VCC', result, "V_DIV_SCALE_F32 must always return VCC")
-    self.assertEqual(result['VCC'] & 1, 0, "VCC lane 0 should be 0 when no scaling needed")
-
-  def test_v_cmp_class_f32_detects_quiet_nan(self):
-    """V_CMP_CLASS_F32 must correctly identify quiet NaN vs signaling NaN.
-    Bug: isQuietNAN and isSignalNAN both used math.isnan which can't distinguish them."""
-    quiet_nan = 0x7fc00000   # quiet NaN: exponent=255, bit22=1
-    signal_nan = 0x7f800001  # signaling NaN: exponent=255, bit22=0
-    # Test quiet NaN detection (bit 1 in mask)
-    s1_quiet = 0b0000000010  # bit 1 = quiet NaN
-    result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
-    self.assertEqual(result['D0'] & 1, 1, "Should detect quiet NaN with quiet NaN mask")
-    # Test signaling NaN detection (bit 0 in mask)
-    s1_signal = 0b0000000001  # bit 0 = signaling NaN
-    result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
-    self.assertEqual(result['D0'] & 1, 1, "Should detect signaling NaN with signaling NaN mask")
-    # Test that quiet NaN doesn't match signaling NaN mask
-    result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
-    self.assertEqual(result['D0'] & 1, 0, "Quiet NaN should not match signaling NaN mask")
-    # Test that signaling NaN doesn't match quiet NaN mask
-    result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
-    self.assertEqual(result['D0'] & 1, 0, "Signaling NaN should not match quiet NaN mask")
-
-  def testisNAN_with_typed_view(self):
-    """isNAN must work with TypedView objects, not just Python floats.
-    Bug: isNAN checked isinstance(x, float) which returned False for TypedView."""
-    nan_reg = Reg(0x7fc00000)  # quiet NaN
-    normal_reg = Reg(0x3f800000)  # 1.0
-    inf_reg = Reg(0x7f800000)  # +inf
-    self.assertTrue(isNAN(nan_reg.f32), "isNAN should return True for NaN TypedView")
-    self.assertFalse(isNAN(normal_reg.f32), "isNAN should return False for normal TypedView")
-    self.assertFalse(isNAN(inf_reg.f32), "isNAN should return False for inf TypedView")
-
-class TestBF16(unittest.TestCase):
-  """Tests for BF16 (bfloat16) support."""
-
-  def test_bf16_conversion(self):
-    """Test bf16 <-> f32 conversion."""
-    # bf16 is just the top 16 bits of f32
-    # 1.0f = 0x3f800000, bf16 = 0x3f80
-    self.assertAlmostEqual(_bf16(0x3f80), 1.0, places=2)
-    self.assertEqual(_ibf16(1.0), 0x3f80)
-    # 2.0f = 0x40000000, bf16 = 0x4000
-    self.assertAlmostEqual(_bf16(0x4000), 2.0, places=2)
-    self.assertEqual(_ibf16(2.0), 0x4000)
-    # -1.0f = 0xbf800000, bf16 = 0xbf80
-    self.assertAlmostEqual(_bf16(0xbf80), -1.0, places=2)
-    self.assertEqual(_ibf16(-1.0), 0xbf80)
-
-  def test_bf16_special_values(self):
-    """Test bf16 special values (inf, nan)."""
-    import math
-    # +inf: f32 = 0x7f800000, bf16 = 0x7f80
-    self.assertTrue(math.isinf(_bf16(0x7f80)))
-    self.assertEqual(_ibf16(float('inf')), 0x7f80)
-    # -inf: f32 = 0xff800000, bf16 = 0xff80
-    self.assertTrue(math.isinf(_bf16(0xff80)))
-    self.assertEqual(_ibf16(float('-inf')), 0xff80)
-    # NaN: quiet NaN bf16 = 0x7fc0
-    self.assertTrue(math.isnan(_bf16(0x7fc0)))
-    self.assertEqual(_ibf16(float('nan')), 0x7fc0)
-
-  def test_bf16_register_property(self):
-    """Test Reg.bf16 property."""
-    r = Reg(0)
-    r.bf16 = 3.0  # 3.0f = 0x40400000, bf16 = 0x4040
-    self.assertEqual(r._val & 0xffff, 0x4040)
-    self.assertAlmostEqual(float(r.bf16), 3.0, places=1)
-
-  def test_bf16_slice_property(self):
-    """Test TypedView.bf16 property."""
-    r = Reg(0x40404040)  # Two bf16 3.0 values
-    self.assertAlmostEqual(r[15:0].bf16, 3.0, places=1)
-    self.assertAlmostEqual(r[31:16].bf16, 3.0, places=1)
-
-class TestBytePermute(unittest.TestCase):
-  """Tests for BYTE_PERMUTE helper function (V_PERM_B32)."""
-
-  def test_byte_select_0_to_7(self):
-    """Test selecting bytes 0-7 from 64-bit data."""
-    # data = {s0, s1} where s0 is bytes 0-3, s1 is bytes 4-7
-    # Combined: 0x0706050403020100 (byte 0 = 0x00, byte 7 = 0x07)
-    data = 0x0706050403020100
-    for i in range(8):
-      self.assertEqual(BYTE_PERMUTE(data, i), i, f"byte {i} should be {i}")
-
-  def test_sign_extend_bytes(self):
-    """Test sign extension selectors 8-11."""
-    # sel 8: sign of byte 1 (bits 15:8)
-    # sel 9: sign of byte 3 (bits 31:24)
-    # sel 10: sign of byte 5 (bits 47:40)
-    # sel 11: sign of byte 7 (bits 63:56)
-    data = 0x8000800080008000  # All relevant bytes have sign bit set
-    self.assertEqual(BYTE_PERMUTE(data, 8), 0xff)
-    self.assertEqual(BYTE_PERMUTE(data, 9), 0xff)
-    self.assertEqual(BYTE_PERMUTE(data, 10), 0xff)
-    self.assertEqual(BYTE_PERMUTE(data, 11), 0xff)
-    data = 0x7f007f007f007f00  # No sign bits set
-    self.assertEqual(BYTE_PERMUTE(data, 8), 0x00)
-    self.assertEqual(BYTE_PERMUTE(data, 9), 0x00)
-    self.assertEqual(BYTE_PERMUTE(data, 10), 0x00)
-    self.assertEqual(BYTE_PERMUTE(data, 11), 0x00)
-
-  def test_constant_zero(self):
-    """Test selector 12 returns 0x00."""
-    self.assertEqual(BYTE_PERMUTE(0xffffffffffffffff, 12), 0x00)
-
-  def test_constant_ff(self):
-    """Test selectors >= 13 return 0xFF."""
-    for sel in [13, 14, 15, 255]:
-      self.assertEqual(BYTE_PERMUTE(0, sel), 0xff, f"sel {sel} should be 0xff")
-
-class TestSADHelpers(unittest.TestCase):
-  """Tests for V_SAD_U8 and V_MSAD_U8 helper functions."""
-
-  def test_v_sad_u8_basic(self):
-    """Test v_sad_u8 with simple values."""
-    # s0 = 0x04030201, s1 = 0x04030201 -> diff = 0 for all bytes
-    result = v_sad_u8(0x04030201, 0x04030201, 0)
-    self.assertEqual(result, 0)
-    # s0 = 0x05040302, s1 = 0x04030201 -> diff = 1+1+1+1 = 4
-    result = v_sad_u8(0x05040302, 0x04030201, 0)
-    self.assertEqual(result, 4)
-
-  def test_v_sad_u8_with_accumulator(self):
-    """Test v_sad_u8 with non-zero accumulator."""
-    # s0 = 0x05040302, s1 = 0x04030201, s2 = 100 -> 4 + 100 = 104
-    result = v_sad_u8(0x05040302, 0x04030201, 100)
-    self.assertEqual(result, 104)
-
-  def test_v_sad_u8_large_diff(self):
-    """Test v_sad_u8 with maximum byte differences."""
-    # s0 = 0xffffffff, s1 = 0x00000000 -> diff = 255*4 = 1020
-    result = v_sad_u8(0xffffffff, 0x00000000, 0)
-    self.assertEqual(result, 1020)
-
-  def test_v_msad_u8_basic(self):
-    """Test v_msad_u8 masks when reference byte is 0."""
-    # s0 = 0x10101010, s1 = 0x00000000 -> all masked, result = 0
-    result = v_msad_u8(0x10101010, 0x00000000, 0)
-    self.assertEqual(result, 0)
-    # s0 = 0x10101010, s1 = 0x01010101 -> diff = |0x10-0x01|*4 = 15*4 = 60
-    result = v_msad_u8(0x10101010, 0x01010101, 0)
-    self.assertEqual(result, 60)
-
-  def test_v_msad_u8_partial_mask(self):
-    """Test v_msad_u8 with partial masking."""
-    # s0 = 0x10101010, s1 = 0x00010001 -> bytes 1 and 3 masked
-    # diff = |0x10-0x01| + |0x10-0x01| = 15 + 15 = 30
-    result = v_msad_u8(0x10101010, 0x00010001, 0)
-    self.assertEqual(result, 30)
-
-  def test_v_msad_u8_with_accumulator(self):
-    """Test v_msad_u8 with non-zero accumulator."""
-    result = v_msad_u8(0x10101010, 0x01010101, 50)
-    self.assertEqual(result, 110)  # 60 + 50
-
-if __name__ == '__main__':
-  unittest.main()

extra/assembly/amd/test/test_rdna3_asm.py

@@ -1,95 +0,0 @@
-#!/usr/bin/env python3
-import unittest, subprocess
-from extra.assembly.amd.autogen.rdna3.ins import *
-from extra.assembly.amd.test.helpers import get_llvm_mc
-
-def llvm_assemble(asm: str) -> bytes:
-  """Assemble using llvm-mc and return bytes."""
-  result = subprocess.run(
-    [get_llvm_mc(), "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
-    input=asm, capture_output=True, text=True
-  )
-  out = b''
-  for line in result.stdout.split('\n'):
-    if 'encoding:' in line:
-      enc = line.split('encoding:')[1].strip()
-      enc = enc.strip('[]').replace('0x', '').replace(',', '')
-      out += bytes.fromhex(enc)
-  if not out: raise ValueError(f"no encoding found: {result.stdout} {result.stderr}")
-  return out
-
-class TestRDNA3Asm(unittest.TestCase):
-  def test_full_program(self):
-    """Test the full program from rdna3fun.py matches llvm-mc output."""
-    program = [
-      v_bfe_u32(v[1], v[0], 10, 10),
-      s_load_b128(s[4:7], s[0:1], NULL),
-      v_and_b32_e32(v[0], 0x3FF, v[0]),
-      s_mulk_i32(s[3], 0x87),
-      v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2]),
-      v_mul_u32_u24_e32(v[0], 45, v[0]),
-      v_ashrrev_i32_e32(v[2], 31, v[1]),
-      v_add3_u32(v[0], v[0], s[3], v[1]),
-      v_lshlrev_b64(v[2:3], 2, v[1:2]),
-      v_ashrrev_i32_e32(v[1], 31, v[0]),
-      v_lshlrev_b64(v[0:1], 2, v[0:1]),
-      s_waitcnt(0xfc07),  # lgkmcnt(0)
-      v_add_co_u32(v[2], VCC_LO, s[6], v[2]),
-      v_add_co_ci_u32_e32(v[3], s[7], v[3]),
-      v_add_co_u32(v[0], VCC_LO, s[4], v[0]),
-      global_load_b32(vdst=v[2], addr=v[2:3], 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:1], 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/bench_log.py

@@ -34,7 +34,8 @@ class WallTimeEvent:
     self.start = time.monotonic()
     return self
   def __exit__(self, *_):
-    _events[self.event]["wall"].append(time.monotonic() - self.start)
+    self.time = time.monotonic() - self.start
+    _events[self.event]["wall"].append(self.time)
     return False
 
 class KernelTimeEvent:

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

@@ -28,7 +28,7 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
   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
+  _, 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)

extra/gemm/amd_asm_matmul.py

@@ -1,49 +1,45 @@
 # RDNA3 128x128 tiled GEMM kernel - DSL version
-# Computes C = A @ B for 4096x4096 float32 matrices using 128x128 tiles
+# 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-117], v[120-124], v[126-129], v[131-133])
+# Accumulators: 128 vgprs (v[2-129])
 
 import numpy as np
-from pathlib import Path
 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 extra.assembly.amd.dsl import s, v, VCC_LO, NULL
-from extra.assembly.amd.autogen.rdna3.ins import *
+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
-MATRIX_DIM = 4096     # Matrix dimension N (assumes square NxN matrices)
 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
 
-# s_waitcnt encodings: wait for memory operations to complete
-WAIT_LGKM = 64519    # wait for LDS/GDS/KMEM (lgkm_cnt=0)
-WAIT_ALL = 0         # wait for everything
-WAIT_VMEM = 1015     # wait for VMEM only (vm_cnt=0, lgkm_cnt=63)
-
 # =============================================================================
-# Named register assignments (VGPRs) - COMPACT LAYOUT
+# Named register assignments (VGPRs)
 # =============================================================================
-V_LANE_ID_MOD8 = 182      # lane_id & 7 (column within 8-wide tile chunk)
-V_OUTPUT_ROW = 171        # output row coordinate
-V_LANE_MOD8_X4 = 174      # V_LANE_ID_MOD8 << 2 (byte offset)
-V_LANE_DIV8_X4 = 175      # (lane_id >> 3) << 2
-V_ADDR_HI_ZERO = 188      # always 0 (for 64-bit address high bits)
-V_LDS_A_BASE = 186        # LDS A-tile base address for inner loop
-V_LDS_B_BASE = 170        # LDS B-tile base address for inner loop
-V_GLOBAL_A_ADDR = 171     # global memory A prefetch address (reuses V_OUTPUT_ROW slot during main loop)
-V_GLOBAL_B_ADDR = 178     # global memory B prefetch address
+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
@@ -54,21 +50,18 @@ V_B_TILE_REGS = [132, 136, 140, 144, 148, 152, 156, 160]  # B tile: banks 0,0,0,
 # Named register assignments (SGPRs)
 # =============================================================================
 S_OUT_PTR = (0, 1)        # output C matrix base pointer
-S_TILE_X = 2              # workgroup_x << 7
-S_TILE_Y = 3              # workgroup_y << 7
+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_A_PTR = (8, 9)          # A matrix base pointer
-S_B_PTR = (10, 11)        # B matrix base pointer
 S_LOOP_CTR = 12           # loop counter (increments by 8)
 S_PREFETCH_FLAG = 13      # prefetch condition flag / row stride in epilogue
-S_WORKGROUP_X = 14        # workgroup_id_x
-S_WORKGROUP_Y = 15        # workgroup_id_y
-# Kernarg load destinations (before copy to working regs)
-S_KERNARG_OUT = (16, 17)  # output pointer from kernarg
+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, 16KB/256KB apart)
+# 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
 
@@ -100,8 +93,6 @@ FMAC_PAIR_ORDER = [
 
 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."""
-  if a_tile_regs is None: a_tile_regs = V_A_TILE_REGS
-  if b_tile_regs is None: b_tile_regs = V_B_TILE_REGS
   pattern = []
   for idx, (a_pair, b_pair) in enumerate(FMAC_PAIR_ORDER):
     a_even, a_odd = a_pair * 2, a_pair * 2 + 1
@@ -120,14 +111,14 @@ def derive_fmac_pattern(acc_grid, a_tile_regs=None, b_tile_regs=None):
   return pattern
 
 # Derived: 64 dual FMAC operations
-FMAC_PATTERN = derive_fmac_pattern(ACC_GRID)
+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 descending reg order
+    -> 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
@@ -150,45 +141,40 @@ def derive_permute_swaps(acc_grid, out_regs):
   return swaps
 
 # Derived: swap sequence to arrange accumulators for output
-OUT_REGS = list(range(129, 1, -1))
+# 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 - COMPACT LAYOUT
+# 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
-# These are SEPARATE - DATA lives during prefetch/store, TILE lives during inner loop
-V_LDS_A_ADDR = 189                            # single base register for A stores (use +512 offsets)
-V_LDS_A_DATA = [155, 172, 173, 154, 159, 176, 177, 158]  # 8 data registers for A prefetch (mod 4: 3,0,1,2,3,0,1,2)
-V_LDS_B_ADDR = 190                            # single base register for B stores (use 16-bit offsets)
-V_LDS_B_DATA = [163, 180, 181, 162, 167, 184, 185, 166]  # 8 data registers for B prefetch (mod 4: 3,0,1,2,3,0,1,2)
+# 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)]
 
-# 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)]
-
-# Initial tile loads: (vdst, addr_lo) pairs - use temp regs in accumulator gaps
-INIT_TILE_LOADS = [(23,5),(24,9),(25,7),(26,2),(27,11),(28,13),(29,6),(30,8),(31,10),(12,12),(13,14),(3,2),(4,4),(5,8),(6,6),(7,10)]
-
-# A matrix row offset registers (scattered to avoid accumulator conflicts)
-ROW_REGS = [165, 146, 147, 164, 169, 150, 151, 168]  # mod 4: 1,2,3,0,1,2,3,0
-
 # =============================================================================
 # Kernel class
 # =============================================================================
 
 class Kernel:
-  def __init__(self, arch='gfx1100'):
-    self.instructions, self.labels, self.branch_targets, self.arch = [], {}, {}, arch
+  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): self.instructions.append(inst); return inst
-  def label(self, name): self.labels[name] = len(self.instructions)
-  def branch_to(self, label): self.branch_targets[len(self.instructions) - 1] = label
+  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."""
@@ -196,231 +182,110 @@ class Kernel:
     waitcnt = (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
     self.emit(s_waitcnt(simm16=waitcnt))
 
-  def to_asm(self):
-    import re
-    # Instruction stream with labels
-    label_at = {pos: name for name, pos in self.labels.items()}
-    body = []
-    for i, inst in enumerate(self.instructions):
-      if i in label_at: body.append(f'.{label_at[i]}:')
-      asm = inst.disasm()
-      if i in self.branch_targets:
-        asm = re.sub(r'(s_cbranch_\w+|s_branch)\s+\S+', rf'\1 .{self.branch_targets[i]}', asm)
-      body.append('\t' + asm)
-
-    # limit wave occupancy by using more LDS
-    lds_size = max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536))
-
-    # HSA kernel descriptor attributes (zeros included for compatibility)
-    hsa = [
-      ('group_segment_fixed_size', lds_size), ('private_segment_fixed_size', 0), ('kernarg_size', 36),
-      ('user_sgpr_count', 14), ('user_sgpr_dispatch_ptr', 0), ('user_sgpr_queue_ptr', 0),
-      ('user_sgpr_kernarg_segment_ptr', 1), ('user_sgpr_dispatch_id', 0), ('user_sgpr_private_segment_size', 0),
-      ('wavefront_size32', 1), ('uses_dynamic_stack', 0), ('enable_private_segment', 0),
-      ('system_sgpr_workgroup_id_x', 1), ('system_sgpr_workgroup_id_y', 1), ('system_sgpr_workgroup_id_z', 0),
-      ('system_sgpr_workgroup_info', 0), ('system_vgpr_workitem_id', 0), ('next_free_vgpr', 192),
-      ('next_free_sgpr', 16), ('float_round_mode_32', 0), ('float_round_mode_16_64', 0),
-      ('float_denorm_mode_32', 3), ('float_denorm_mode_16_64', 3), ('dx10_clamp', 1), ('ieee_mode', 1),
-      ('fp16_overflow', 0), ('workgroup_processor_mode', 0), ('memory_ordered', 1), ('forward_progress', 0),
-      ('shared_vgpr_count', 0)]
-
-    return '\n'.join([
-      '\t.text', f'\t.amdgcn_target "amdgcn-amd-amdhsa--{self.arch}"',
-      '\t.protected\tkernel', '\t.globl\tkernel', '\t.p2align\t8', '\t.type\tkernel,@function', 'kernel:',
-      *body,
-      '\t.section\t.rodata,"a",@progbits', '\t.p2align\t6, 0x0', '\t.amdhsa_kernel kernel',
-      *[f'\t\t.amdhsa_{k} {v}' for k, v in hsa],
-      '\t.end_amdhsa_kernel', '\t.text', '.Lfunc_end0:', '\t.size\tkernel, .Lfunc_end0-kernel',
-      '\t.amdgpu_metadata', '---', 'amdhsa.kernels:', '  - .args:',
-      *[f'      - .address_space: global\n        .offset: {i*8}\n        .size: 8\n        .value_kind: global_buffer' for i in range(3)],
-      f'    .group_segment_fixed_size: {lds_size}', '    .kernarg_segment_align: 8',
-      '    .kernarg_segment_size: 24', '    .max_flat_workgroup_size: 128', '    .name: kernel',
-      '    .private_segment_fixed_size: 0', '    .sgpr_count: 60', '    .symbol: kernel.kd',
-      '    .vgpr_count: 192', '    .wavefront_size: 32', f'amdhsa.target: amdgcn-amd-amdhsa--{self.arch}',
-      'amdhsa.version:', '  - 1', '  - 2', '...', '\t.end_amdgpu_metadata'])
+  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(arch='gfx1100'):
+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_KERNARG_OUT[0]:S_KERNARG_OUT[1]], sbase=s[0:1], offset=0x10, soffset=NULL))
-  k.emit(s_mov_b32(s[S_DIM_N], MATRIX_DIM))
+  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[0]))
-  k.emit(v_lshrrev_b32_e32(v[4], 3, v[0]))
-  k.emit(v_or_b32_e32(v[1], s[S_TILE_X], v[0]))
+  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.emit(v_mov_b32_e32(v[2], 0))  # v[1] always positive, sign extension is 0
-  k.emit(v_lshlrev_b64(v[5:6], 2, v[1:2]))
   k.waitcnt(lgkm=0)
 
-  # Copy pointers to working registers
-  k.emit(s_mov_b64(s[S_OUT_PTR[0]:S_OUT_PTR[1]], s[S_KERNARG_OUT[0]:S_KERNARG_OUT[1]]))
-  k.emit(s_mov_b64(s[S_A_PTR[0]:S_A_PTR[1]], s[S_KERNARG_A[0]:S_KERNARG_A[1]]))
-  k.emit(s_mov_b64(s[S_B_PTR[0]:S_B_PTR[1]], s[S_KERNARG_B[0]:S_KERNARG_B[1]]))
-
   # 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: 16KB apart
-    k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_KERNARG_B[0]], i * 0x4000))
+  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: 256KB apart
-    k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_KERNARG_A[0]], i * 0x40000))
+  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 << 12) + (lane_id/8)*4K + lane_id%8) * 4
-  k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], s[S_TILE_X], v[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_lshl_b32(s[19], s[S_TILE_Y], 12))
-  k.emit(v_lshl_add_u32(v[V_GLOBAL_A_ADDR], v[4], 12, v[V_LANE_ID_MOD8]))  # (lane_id/8)*4K + lane_id%8
+  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]))
 
-  # ===========================================================================
-  # Tile address computation for initial A/B matrix loads
-  # ===========================================================================
-  k.emit(s_lshl_b32(s[S_LOOP_BOUND], s[S_DIM_N], 4))  # row stride = 16*N
-  k.emit(v_mul_lo_u32(v[ROW_REGS[0]], v[22], s[S_DIM_N]))  # A matrix row offsets
-  for i in range(1, 8): k.emit(v_add_nc_u32_e32(v[ROW_REGS[i]], s[S_LOOP_BOUND], v[ROW_REGS[i-1]]))
-
-  def addr64(dst, base_s):  # 64-bit address: v[dst:dst+1] = s[base_s:base_s+1] + v[dst]*4
-    k.emit(v_mov_b32_e32(v[dst+1], 0))  # offset always positive, sign ext = 0
-    k.emit(v_lshlrev_b64(v[dst:dst+1], 2, v[dst:dst+1]))
-    k.emit(v_add_co_u32(v[dst], VCC_LO, s[base_s], v[dst]))
-    k.emit(v_add_co_ci_u32_e32(v[dst+1], s[base_s+1], v[dst+1]))
-
-  def b_addr(dst, mult, tmp=None):  # B address for col + mult*N
-    tmp = tmp if tmp is not None else dst
-    k.emit(v_mad_u32_u24(v[tmp], s[S_DIM_N], mult, v[1]))
-    if tmp != dst:
-      k.emit(v_mov_b32_e32(v[tmp+1], 0))  # offset always positive
-      k.emit(v_lshlrev_b64(v[dst:dst+1], 2, v[tmp:tmp+1]))
-      k.emit(v_add_co_u32(v[dst], VCC_LO, s[S_B_PTR[0]], v[dst]))
-      k.emit(v_add_co_ci_u32_e32(v[dst+1], s[S_B_PTR[1]], v[dst+1]))
-    else: addr64(dst, S_B_PTR[0])
-
-  def a_addr(dst, row_reg, tmp):  # A address for row_reg + lane_id_mod8
-    k.emit(v_add_nc_u32_e32(v[tmp], v[row_reg], v[V_LANE_ID_MOD8]))
-    k.emit(v_mov_b32_e32(v[tmp+1], 0))  # offset always positive
-    k.emit(v_lshlrev_b64(v[dst:dst+1], 2, v[tmp:tmp+1]))
-    k.emit(v_add_co_u32(v[dst], VCC_LO, s[S_A_PTR[0]], v[dst]))
-    k.emit(v_add_co_ci_u32_e32(v[dst+1], s[S_A_PTR[1]], v[dst+1]))
-
-  # Batch 1: B addresses (cols 0-5) and loads
-  k.emit(v_add_co_u32(v[5], VCC_LO, s[S_B_PTR[0]], v[5]))
-  k.emit(v_add_co_ci_u32_e32(v[6], s[S_B_PTR[1]], v[6]))
-  for dst, mult in [(9,1), (7,2), (2,3), (11,4), (13,5)]: b_addr(dst, mult)
-  k.emit(s_clause(simm16=5))  # 6 consecutive global loads
-  for vdst, addr in INIT_TILE_LOADS[:6]: k.emit(global_load_b32(vdst=v[vdst], addr=v[addr:addr+1], saddr=NULL))
-
-  # Batch 2: A addresses (rows 0-4) and loads
-  for dst, ri in [(6,0), (8,1), (10,2), (12,3), (14,4)]:
-    k.emit(v_add_nc_u32_e32(v[dst], v[ROW_REGS[ri]], v[V_LANE_ID_MOD8]))
-    addr64(dst, S_A_PTR[0])
-  k.emit(s_clause(simm16=4))  # 5 consecutive global loads
-  for vdst, addr in INIT_TILE_LOADS[6:11]: k.emit(global_load_b32(vdst=v[vdst], addr=v[addr:addr+1], saddr=NULL))
-
-  # Batch 3: B cols 6-7, A rows 5-7, and loads
-  for dst, mult, tmp in [(2,6,15), (4,7,4)]: b_addr(dst, mult, tmp)
-  for dst, ri, tmp in [(8,5,16), (6,6,18), (10,7,20)]: a_addr(dst, ROW_REGS[ri], tmp)
-  k.emit(s_clause(simm16=4))  # 5 consecutive global loads
-  for vdst, addr in INIT_TILE_LOADS[11:]: k.emit(global_load_b32(vdst=v[vdst], addr=v[addr:addr+1], saddr=NULL))
+  # 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.
-  # Key outputs:
-  #   v[8]: A-tile initial store base (used only for initial stores with stride64)
-  #   V_LDS_B_ADDR (v145): B-tile store base (used for both initial and main loop)
-  #   V_LANE_DIV8_X4 (v135): (lane_id >> 3) << 2 for epilogue
-  #
   # 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.
-
-  # v[22] = tile_y | (lane_id >> 3) from prologue, used as base for row offsets
-  # Compute 7 row offsets for B-tile rows 1-7 (row 0 computed separately in v[9])
-  k.emit(v_add_nc_u32_e32(v[9], s[S_LOOP_CTR], v[22]))  # row 0 base (S_LOOP_CTR=0)
-  for i in range(7): k.emit(v_or_b32_e32(v[10 + i if i < 2 else 12 + i], 16 * (i + 1), v[22]))  # rows 1-7
-
-  # Extract sign bit of workgroup_x (always 0 for valid workgroups, used for masking)
-  k.emit(s_bfe_i32(s[S_LOOP_BOUND], s[S_WORKGROUP_X], 0x10018))
+  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(s_lshr_b32(s[S_LOOP_BOUND], s[S_LOOP_BOUND], 25))
-
-  # Compute masked row offsets for bank conflict avoidance pattern
-  # Pattern: v[row] = row_val - (row_val & ADDR_MASK) extracts lower bits
-  k.emit(v_add_nc_u32_e32(v[19], s[S_LOOP_CTR], v[10]))
-  k.emit(v_add_nc_u32_e32(v[8], s[S_LOOP_BOUND], v[1]))  # A-tile base computation
-  for d, r in zip([20, 21, 32, 33, 34, 35], [11, 14, 15, 16, 17, 18]):
-    k.emit(v_add_nc_u32_e32(v[d], s[S_LOOP_CTR], v[r]))
-  k.emit(v_and_b32_e32(v[8], ADDR_MASK, v[8]))
   k.emit(v_sub_nc_u32_e32(v[9], v[22], v[9]))  # row 0 swizzle offset
-  for d, s_ in zip([19, 20, 21, 22, 32, 33, 34], [20, 21, 22, 32, 33, 34, 35]):
-    k.emit(v_and_b32_e32(v[d], ADDR_MASK, v[s_]))
-  k.emit(v_sub_nc_u32_e32(v[8], v[1], v[8]))  # A-tile swizzle
-
-  # Apply swizzle offsets and scale to byte offsets
-  k.emit(v_lshlrev_b32_e32(v[9], 2, v[9]))  # row 0 offset * 4
-  for r, t in zip([10, 11, 14, 15, 16, 17, 18], [19, 20, 21, 22, 32, 33, 34]):
-    k.emit(v_sub_nc_u32_e32(v[r], v[r], v[t]))  # rows 1-7 swizzle
-  k.emit(v_bfe_u32(v[2], v[0], 3, 2))  # v[2] = (lane_id >> 3) & 3
-  k.emit(v_lshlrev_b32_e32(v[8], 2, v[8]))  # A-tile base * 4
-
-  # Compute B-tile base address: LDS_A_STRIDE * (lane_id % 8) + row0_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]))
-  # Scale row offsets 1-7 to byte offsets (row 0 already in v[9])
-  for d, r in zip([9, 10, 11, 14, 15, 16, 17], [10, 11, 14, 15, 16, 17, 18]):
-    k.emit(v_lshlrev_b32_e32(v[d], 2, v[r]))
+
+  # 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]))
-  k.emit(v_add_nc_u32_e32(v[8], 0x80, v[8]))  # A-tile initial store base + 128
 
-  # Store initial tile data to LDS
-  k.waitcnt(vm=0)
-  for i, (d0, d1) in enumerate([(0,1), (2,3), (4,5), (11,12)]):
-    k.emit(ds_store_2addr_stride64_b32(addr=v[8], data0=v[INIT_TILE_LOADS[d0][0]], data1=v[INIT_TILE_LOADS[d1][0]], offset0=16+i*4, offset1=18+i*4))
-  # B stores: single base with offsets 0,64,128,192,256,320,384,448
-  for i, idx in enumerate([6,7,8,9,10,13,14,15]):
-    offset = i * 64
-    k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[INIT_TILE_LOADS[idx][0]], offset0=offset & 0xFF, offset1=offset >> 8))
-
-  # ===========================================================================
-  # INIT: Compute LDS base addresses, then zero accumulators
-  # ===========================================================================
-  # v[3] = v[1] & 0x7F (lower 7 bits) since S_LOOP_BOUND=0 for valid workgroups
+  # Compute LDS load/store base addresses for inner loop
   k.emit(v_lshlrev_b32_e32(v[2], 4, v[2]))
-  k.emit(v_add_nc_u32_e32(v[3], s[S_LOOP_BOUND], v[1]))
-  k.emit(v_and_b32_e32(v[3], ADDR_MASK, v[3]))
-  k.emit(v_sub_nc_u32_e32(v[3], v[1], v[3]))
+  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[0]))
+  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))
-  k.emit(s_add_u32(s[S_A_PTR[0]], s[S_A_PTR[0]], 32))
-  k.emit(s_addc_u32(s[S_A_PTR[1]], s[S_A_PTR[1]], 0))
+
   # S_LOOP_CTR is already 0 from prologue initialization
-  k.emit(s_branch(simm16=0)); k.branch_to('LOOP_ENTRY')
+  k.emit(s_branch(), target='LOOP_ENTRY')
 
   # ===========================================================================
   # MAIN GEMM LOOP
@@ -431,22 +296,22 @@ def build_kernel(arch='gfx1100'):
   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(simm16=0)); k.branch_to('EPILOGUE')
+  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(simm16=0)); k.branch_to('SKIP_PREFETCH')  # branch if loop_ctr >= loop_bound
+  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], 0x20000, v[V_GLOBAL_B_ADDR]))
+    #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 (SGPRs, 64-bit adds)
+    # 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], 0x20000))
+      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))
@@ -490,7 +355,7 @@ def build_kernel(arch='gfx1100'):
       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 stores to finish
+  # 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())
@@ -507,7 +372,7 @@ def build_kernel(arch='gfx1100'):
       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(simm16=0)); k.branch_to('LOOP_INC')
+  k.emit(s_branch(), target='LOOP_INC')
 
   # ===========================================================================
   # EPILOGUE: Permute and store results
@@ -518,59 +383,52 @@ def build_kernel(arch='gfx1100'):
   for a, b in PERMUTE_SWAPS:
     k.emit(v_swap_b32_e32(v[a], v[b]))
 
-  # Compute output coordinates: v[V_LANE_ID_MOD8] = col, v[V_OUTPUT_ROW] = row
-  k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32,
-              vdstx=v[149], vdsty=v[150], srcx0=v[V_LANE_MOD8_X4], vsrcx1=v[0], srcy0=v[V_LANE_DIV8_X4], vsrcy1=v[0]))
-  k.emit(v_and_b32_e32(v[0], 0x60, v[0]))
-  k.emit(v_or_b32_e32(v[V_LANE_ID_MOD8], s[S_TILE_X], v[149]))
-  k.emit(v_add_nc_u32_e32(v[0], s[S_TILE_Y], v[0]))
-  k.emit(v_or_b32_e32(v[V_OUTPUT_ROW], v[0], v[150]))
+  # 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[144-147] for rows 0-3, v[148-151] for rows 16-19
-  for base, row_off in [(144, 0), (148, 16)]:
-    if row_off: k.emit(v_or_b32_e32(v[1], row_off, v[V_OUTPUT_ROW]))
-    k.emit(v_mul_lo_u32(v[base], v[1] if row_off else v[V_OUTPUT_ROW], s[S_DIM_N]))
-    for i in range(3): k.emit(v_add_nc_u32_e32(v[base + 1 + i], s[S_DIM_N], v[base + i]))
+  # 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]))
 
-  k.emit(v_mov_b32_e32(v[V_ADDR_HI_ZERO], 0))
-  k.emit(s_lshl_b32(s[S_PREFETCH_FLAG], s[S_DIM_N], 2))  # row stride in bytes
+  # 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
-  epilogue_reserved = {V_LANE_ID_MOD8, V_OUTPUT_ROW, V_LANE_MOD8_X4, V_LANE_DIV8_X4, V_ADDR_HI_ZERO}
-
   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
-    srcs = OUT_REGS[i*4:(i+1)*4]
+    src = OUT_REGS[i*4]  # first reg of ascending group of 4
 
-    # Find temp register for scaled values (must not conflict with reserved regs)
-    tmp = max(srcs) + 5
-    while any(r in epilogue_reserved for r in range(tmp, tmp + 4)): tmp += 1
+    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]))
 
-    # Copy values to temp regs for output (alpha=1.0 hardcoded, so just move)
-    for j, src in enumerate(srcs):
-      k.emit(v_mov_b32_e32(v[tmp + j], v[src]))
-
-    # Compute output address
-    if row_in_group == 0:  # first row: compute base address for this column group
-      if col_off == 0: k.emit(v_mov_b32_e32(v[0], v[V_LANE_ID_MOD8]))
-      else: k.emit(v_add_nc_u32_e32(v[0], col_off, v[V_LANE_ID_MOD8]))
-      row_base = 144 + row if row < 4 else 148 + row - 16
-      k.emit(v_add_nc_u32_e32(v[0], v[row_base], v[0]))
-      k.emit(v_lshlrev_b32_e32(v[0], 2, v[0]))
-      k.emit(v_add_co_u32(v[0], VCC_LO, s[S_OUT_PTR[0]], v[0]))
-      k.emit(v_add_co_ci_u32_e32(v[1], s[S_OUT_PTR[1]], v[V_ADDR_HI_ZERO]))
-    else:  # subsequent rows: just add stride
-      k.emit(v_add_co_u32(v[0], VCC_LO, s[S_PREFETCH_FLAG], v[0]))
-      k.emit(v_add_co_ci_u32_e32(v[1], v[1], v[V_ADDR_HI_ZERO]))
-
-    k.emit(global_store_b128(addr=v[0:1], data=v[tmp:tmp+3], saddr=NULL))
+    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.to_asm()
+  return k.finalize()
 
 # =============================================================================
 # Test harness
@@ -582,18 +440,9 @@ THREADS = 128
 
 def test_matmul():
   dev = Device[Device.DEFAULT]
-  print(f"Device arch: {dev.arch}")
+  print(f"Device arch: {dev.renderer.arch}")
 
-  if getenv("STOCK", 0):
-    # Load the stock kernel from amd_seb/kernel8_batched_gmem.s
-    stock_path = Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s"
-    asm = stock_path.read_text()
-    print(f"Loaded stock kernel from {stock_path}")
-  else:
-    asm = build_kernel(dev.arch)
-
-  binary = dev.compiler.compile(asm)
-  print(f"Compiled! Binary size: {len(binary)} bytes")
+  insts = build_kernel(N, dev.renderer.arch)
 
   rng = np.random.default_rng(42)
   a = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
@@ -608,10 +457,10 @@ def test_matmul():
   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)]
-    sink = UOp.sink(A.base, B.base, C.base, *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=(*sink.src, sink)), UOp(Ops.SOURCE, arg=asm),
-                                 UOp(Ops.BINARY, arg=binary)))
+    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()
 
@@ -625,33 +474,23 @@ def test_matmul():
     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: raise RuntimeError("matmul is wrong!")
-
-def run_sqtt():
-  """Run with SQTT profiling and write trace files."""
-  import subprocess, os
-
-  # Run test_matmul in a subprocess with SQTT enabled from the start (no verify)
-  env = {**os.environ, "AMD": "1", "SQTT": "1", "CNT": "1", "PROFILE": "1", "PYTHONPATH": ".", "VERIFY": "0"}
-  result = subprocess.run(
-    ["python", "-c", "from extra.gemm.amd_asm_matmul import test_matmul; test_matmul()"],
-    capture_output=True, text=True, env=env, timeout=120
-  )
-  print(result.stdout)
-
-  # Run roc.py to extract trace data
-  result = subprocess.run(
-    ["python", "extra/sqtt/roc.py", "--profile", "/tmp/profile.pkl.tiny", "--kernel", "kernel"],
-    capture_output=True, text=True, env={**os.environ, "DEBUG": "5"}, timeout=60
-  )
-  output = result.stdout + result.stderr
-
-  # Write full output to trace file
-  with open("/tmp/sqtt_trace.txt", "w") as f:
-    f.write(output)
-  print(f"Wrote {len(output)} bytes to /tmp/sqtt_trace.txt")
+    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__":
-  if getenv("ASM", 0): print(build_kernel(Device[Device.DEFAULT].arch))
-  elif getenv("SQTT", 0): run_sqtt()
-  else: test_matmul()
+  test_matmul()

extra/gemm/amd_copy_matmul.py

@@ -0,0 +1,84 @@
+from tinygrad import UOp, getenv
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+from tinygrad.dtype import AddrSpace, dtypes
+
+N = getenv("N", 4096)
+M = getenv("M", N)
+K = getenv("K", N)
+
+WARP_SIZE = 32
+BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 8
+TM, TN = 4, 4
+LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
+assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
+
+is_kernel5 = getenv("K5", 0)
+THREADS_PER_BLOCK = 128 if is_kernel5 else 256
+WAVES_PER_BLOCK_N = 1 if is_kernel5 else 2
+WAVES_PER_BLOCK_M = THREADS_PER_BLOCK // WARP_SIZE // WAVES_PER_BLOCK_N
+REG_TILES_PER_WAVE_N = BLOCK_N // (WAVES_PER_BLOCK_N * LANES_PER_WAVE_N * TN)
+REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
+
+assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
+assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
+
+consts = {"wpb_m":WAVES_PER_BLOCK_M, "lpw_m":LANES_PER_WAVE_M, "rt_m":REG_TILES_PER_WAVE_M, "t_m": TM,
+          "wpb_n":WAVES_PER_BLOCK_N, "lpw_n":LANES_PER_WAVE_N, "rt_n":REG_TILES_PER_WAVE_N, "t_n": TN}
+
+# 128x128 out, kx128, kx128 in
+def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
+  tid = UOp.range(THREADS_PER_BLOCK, 2, AxisType.LOCAL)
+
+  #tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
+  warp, lane = tid // WARP_SIZE, tid % WARP_SIZE
+  wave_n, wave_m = warp % WAVES_PER_BLOCK_N, warp // WAVES_PER_BLOCK_N
+  lane_n, lane_m = lane % LANES_PER_WAVE_N, lane // LANES_PER_WAVE_N
+
+  # define locals
+  A_local = UOp.placeholder((BLOCK_K, BLOCK_M), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
+  B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+
+  # open the main reduction range and copy in GLOBAL -> LOCAL
+  a = a.reshape(K // BLOCK_K, BLOCK_K, BLOCK_M)
+  b = b.reshape(K // BLOCK_K, BLOCK_K, BLOCK_N)
+  k_tile_range = UOp.range(K // BLOCK_K, 3, AxisType.REDUCE)
+  A_store = A_local.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile_range].reshape(-1, THREADS_PER_BLOCK)[:, tid])
+  B_store = B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile_range].reshape(-1, THREADS_PER_BLOCK)[:, tid])
+  barrier = UOp.barrier(A_store, B_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+
+  # define accumulator (128x128), but broadcast across tid
+  c_regs = UOp.placeholder((REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  c_regs = c_regs.after(c_regs.store(UOp.const(dtypes.float, 0).reshape((1,)*len(c_regs.shape)).expand(c_regs.shape)))
+
+  # define registers (NOTE: the thread count is the device count for this multi, it's sharded across the THREADS_PER_BLOCK)
+  A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+
+  # LOCAL -> REGS
+  k = UOp.range(BLOCK_K, 4, AxisType.REDUCE)
+  A_col = A_col.after(A_col.store(A_local[k].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[wave_m, :, lane_m, :]))
+  B_row = B_row.after(B_row.store(B_local[k].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[wave_n, :, lane_n, :]))
+
+  # do FMA
+  A_col = A_col.reshape(REG_TILES_PER_WAVE_M*TM, 1).expand(REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN)
+  B_row = B_row.reshape(1, REG_TILES_PER_WAVE_N*TN).expand(REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN)
+  c_regs = c_regs.after(c_regs.store(c_regs.after(k) + (A_col * B_row)).end(k).barrier().end(k_tile_range))
+
+  # store back to c
+  c_store = c.rearrange("(wpb_m rt_m lpw_m t_m) (wpb_n rt_n lpw_n t_n) -> (wpb_m wpb_n lpw_m lpw_n) (rt_m t_m) (rt_n t_n)", **consts)
+  return c_store[tid].store(c_regs).end(tid)
+
+def amd_copy_matmul(c:UOp, a:UOp, b:UOp) -> UOp:
+  block_id_n = UOp.range(N // BLOCK_N, 0, AxisType.GLOBAL)
+  block_id_m = UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL)
+
+  # index the output with the globals
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
+  a = a.T.reshape(K, M // BLOCK_M, BLOCK_M)[:, block_id_m, :]
+  b = b.reshape(K, N // BLOCK_N, BLOCK_N)[:, block_id_n, :]
+  return block_128x128_gemm(c, a, b).end(block_id_n, block_id_m).sink(arg=KernelInfo(opts_to_apply=()))
+
+if __name__ == "__main__":
+  from amd_uop_matmul import eval_custom_matmul
+  eval_custom_matmul(amd_copy_matmul)

extra/gemm/amd_uop_matmul.py

@@ -1,98 +1,74 @@
-import numpy as np
-from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
+from tinygrad import Tensor, Context, GlobalCounters, dtypes
 from tinygrad.uop.ops import UOp, KernelInfo, sint, AxisType
-from tinygrad.engine.realize import ExecItem, get_runner
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import getenv
+from tinygrad.helpers import DEBUG, getenv
 
 N = getenv("N", 4096)
-M = K = N
-run_count = getenv("CNT", 5)
+M = getenv("M", N)
+K = getenv("K", N)
+NUM_RUNS = getenv("CNT", 5)
 
 # ---------------------------
 # launch/config constants
 # ---------------------------
 
 WARP_SIZE = 32
-
-# Threadblock tile sizes (block-level tile of C that a block computes)
-BLOCK_N = 128   # columns of C (N-dim) per block
-BLOCK_M = 128   # rows of C (M-dim) per block
-BLOCK_K = 8     # K-slice per block iteration
-
-# Register tile sizes (per-thread accumulator tile of C)
-TN = 4     # columns per thread
-TM = 4     # rows per thread
+BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 8
+TM, TN = 4, 4
+LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
+assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
 
 is_kernel5 = getenv("K5", 0)
 THREADS_PER_BLOCK = 128 if is_kernel5 else 256
-assert THREADS_PER_BLOCK % BLOCK_N == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_N"
-assert THREADS_PER_BLOCK % BLOCK_K == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_K"
-assert (BLOCK_N * BLOCK_K) % THREADS_PER_BLOCK == 0
-assert (BLOCK_M * BLOCK_K) % THREADS_PER_BLOCK == 0
+WAVES_PER_BLOCK_N = 1 if is_kernel5 else 2
+WAVES_PER_BLOCK_M = THREADS_PER_BLOCK // WARP_SIZE // WAVES_PER_BLOCK_N
+REG_TILES_PER_WAVE_N = BLOCK_N // (WAVES_PER_BLOCK_N * LANES_PER_WAVE_N * TN)
+REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
 
-WARPS_PER_BLOCK = THREADS_PER_BLOCK // WARP_SIZE
-WAVE_TILE_N = 128 if is_kernel5 else 64
-WAVE_TILE_M = BLOCK_N * BLOCK_M // WARPS_PER_BLOCK // WAVE_TILE_N
-assert BLOCK_N % WAVE_TILE_N == 0, "BN must be a multiple of WN"
-assert BLOCK_M % WAVE_TILE_M == 0, "BM must be a multiple of WM"
-WAVES_IN_BLOCK_X = BLOCK_N // WAVE_TILE_N
-WAVES_IN_BLOCK_Y = BLOCK_M // WAVE_TILE_M
-assert WAVES_IN_BLOCK_X * WAVES_IN_BLOCK_Y == WARPS_PER_BLOCK, "wave grid must match warps/block"
-
-LANES_PER_WAVE_X = 8
-LANES_PER_WAVE_Y = 4
-ITERS_PER_WAVE_N = WAVE_TILE_N // (LANES_PER_WAVE_X * TN)
-ITERS_PER_WAVE_M = WAVE_TILE_M // (LANES_PER_WAVE_Y * TM)
-assert WAVE_TILE_N % (LANES_PER_WAVE_X * TN) == 0, "WAVE_TILE_N must be divisible by LANES_PER_WAVE_X*TN"
-assert WAVE_TILE_M % (LANES_PER_WAVE_Y * TM) == 0, "WAVE_TILE_M must be divisible by LANES_PER_WAVE_Y*TM"
+assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
+assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
 
 def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
-def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
+def copy(dest:UOp, src:UOp, rng:int, upcast=False):
   assert dest.shape == src.shape
   rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
-  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
-  return dest.after(copy) if set else copy
+  return dest[*rngs].store(src[*rngs]).end(*rngs)
 
-def hand_spec_kernel3():
+def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   # ---------------------------
-  # block indices & placeholders
+  # block indices
   # ---------------------------
-  blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
-  blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")
-
-  a = UOp.placeholder((N, N), dtypes.float, slot=1)
-  b = UOp.placeholder((N, N), dtypes.float, slot=2)
-  c = UOp.placeholder((N, N), dtypes.float, slot=0)
+  block_id_n = UOp.special(N // BLOCK_N, "gidx0")
+  block_id_m = UOp.special(M // BLOCK_M, "gidx1")
 
   # index the output with the globals
-  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
 
   # open the main reduction range
-  k_tile_range = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)
-  a = a.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_K, BLOCK_K)[blockIdx_y, :, k_tile_range, :]
-  b = b.reshape(N // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, blockIdx_x, :]
+  k_tile_range = UOp.range(K // BLOCK_K, 0, AxisType.REDUCE)
+  a = a.reshape(M // BLOCK_M, BLOCK_M, K // BLOCK_K, BLOCK_K)[block_id_m, :, k_tile_range, :]
+  b = b.reshape(K // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, block_id_n, :]
 
   # globals are no longer used, they are already in the indexes
-  del blockIdx_y, blockIdx_x
+  del block_id_m, block_id_n
 
   # ---------------------------
-  # GLOBAL -> LOCAL (As, Bs)
+  # GLOBAL -> LOCAL (A_local, B_local)
   # ---------------------------
   tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
 
   # A: read BM x BK tiles (permute on store into locals)
-  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
-  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
-  As_store = copy(As.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
+  BM_A_local_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
+  A_local = UOp.placeholder((BLOCK_K, BM_A_local_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
+  A_local_store = copy(A_local.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
 
   # B: read BK x BN tiles
-  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-  Bs_store = copy(Bs.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
+  B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
 
   # TODO: can we automate barrier?
-  barrier = UOp.barrier(As_store, Bs_store)
-  As, Bs = As.after(barrier), Bs.after(barrier)
+  barrier = UOp.barrier(A_local_store, B_local_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
 
   # open inner k range
   k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -100,31 +76,30 @@ def hand_spec_kernel3():
   # ---------------------------
   # LOCAL -> REG (per-wave tiles)
   # ---------------------------
-  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
-  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
-  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
+  warp, lane = tid // WARP_SIZE, tid % WARP_SIZE
+  waveIdx, waveIdy = warp % WAVES_PER_BLOCK_N, warp // WAVES_PER_BLOCK_N
+  laneIdx, laneIdy = lane % LANES_PER_WAVE_N, lane // LANES_PER_WAVE_N
+  assert waveIdy.vmax+1 == WAVES_PER_BLOCK_M and laneIdy.vmax+1 == LANES_PER_WAVE_M
 
-  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
-  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
-  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
+  A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  A_local_slice = A_local[k, :].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[waveIdy, :, laneIdy, :]
+  A_col = A_col.after(copy(A_col, A_local_slice, 300, upcast=True))
 
-  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  A_col = copy(A_col, As[k, :].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :], 300, set=True, upcast=True)
-
-  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  B_row = copy(B_row, Bs[k, :].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :], 400, set=True, upcast=True)
+  B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  B_local_slice = B_local[k, :].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[waveIdx, :, laneIdx, :]
+  B_row = B_row.after(copy(B_row, B_local_slice, 400, upcast=True))
 
   # ---------------------------
   # FMA: c_regs += A_col * B_row
   # ---------------------------
-  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  c_regs = UOp.placeholder((REG_TILES_PER_WAVE_M, TM, REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
   i = UOp.range(c_regs.size, 16)
   c_regs = c_regs.after(c_regs.flatten()[i].store(0.0).end(i))
 
   # TODO: why don't these work as upcast?
   # why if the ranges merge is it slow?!? (if you change the order on end, they will merge. big slowdown on METAL)
-  iterWaveM, yt, iterWaveN, xt = rngs = rngs_for_shape(c_regs.shape, 500)
-  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iterWaveM, yt] * B_row[iterWaveN, xt]).end(iterWaveM, iterWaveN, yt, xt)
+  iter_m, t_m, iter_n, t_n = rngs = rngs_for_shape(c_regs.shape, 500)
+  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)
 
   # Close k, sync, and close K tiles
   sink = sink.end(k).barrier().end(k_tile_range)
@@ -132,38 +107,37 @@ def hand_spec_kernel3():
   # ---------------------------
   # REG -> GLOBAL (epilogue)
   # ---------------------------
-  c = c.reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
-                WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
+  c = c.reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM,
+                WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)
   c = c[waveIdy, :, laneIdy, :,
         waveIdx, :, laneIdx, :]
   sink = copy(c, c_regs.after(sink), rng=600)
 
   return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
 
-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, 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))
+def eval_custom_matmul(fxn):
+  a = Tensor.randn(M, K, dtype=dtypes.float)
+  b = Tensor.randn(K, N, dtype=dtypes.float)
+  c = Tensor.empty(M, N, dtype=dtypes.float)
+  with Context(DEBUG=0): Tensor.realize(a, b)
 
   ets = []
-  with Context(DEBUG=2):
-    for _ in range(run_count):
-      ets.append(ei.run(wait=True))
-  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
+  with Context(DEBUG=max(2, DEBUG.value)):
+    for _ in range(NUM_RUNS):
+      GlobalCounters.reset()
+      tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize()
+      ets.append(GlobalCounters.time_sum_s)
+  print(f"REAL TFLOPS {M * N * K * 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 = (hc - tc).square().mean().item()
+      err = (tc - tst).square().mean().item()
     print(f"mean squared error {err}")
     if err > 1e-06:
       raise RuntimeError("matmul is wrong!")
 
 if __name__ == "__main__":
-  test_matmul(hand_spec_kernel3(), N=N)
+  eval_custom_matmul(hand_spec_kernel3)

extra/gemm/asm/cdna/template.s

@@ -1,78 +0,0 @@
-.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 133120
-  .amdhsa_private_segment_fixed_size 0
-  .amdhsa_kernarg_size 28
-  # register usage (RSRC1)
-  .amdhsa_next_free_vgpr 504
-  .amdhsa_next_free_sgpr 96
-  # 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, we only specify the kernel args ptr in s[0:1]
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_user_sgpr_count 2
-  .amdhsa_user_sgpr_kernarg_preload_length 0
-  .amdhsa_user_sgpr_kernarg_preload_offset 0
-  # gfx90a / gfx940 specifics (RSRC3)
-  .amdhsa_accum_offset 248
-  .amdhsa_uses_dynamic_stack 0
-  .amdhsa_tg_split 0
-.end_amdhsa_kernel
-
-.amdgpu_metadata
----
-amdhsa.kernels:
-  - .name: gemm
-    .symbol: gemm.kd
-    .args:
-      - .name: C
-        .address_space: global
-        .offset: 0
-        .size: 8
-        .value_kind: global_buffer
-        .value_type: bf16
-      - .name: B
-        .address_space: global
-        .offset: 8
-        .size: 8
-        .value_kind: global_buffer
-        .value_type: bf16
-      - .name: A
-        .address_space: global
-        .offset: 16
-        .size: 8
-        .value_kind: global_buffer
-        .value_type: bf16
-      - .name: sz
-        .offset: 24
-        .size: 4
-        .value_kind: by_value
-        .value_type: u32
-    .group_segment_fixed_size: 133120
-    .private_segment_fixed_size: 0
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 28
-    .max_flat_workgroup_size: 256
-    .sgpr_count: 88
-    .sgpr_spill_count: 0
-    .vgpr_count: 248
-    .vgpr_spill_count: 0
-    .wavefront_size: 64
-amdhsa.version:
-  - 1
-  - 0
-...
-.end_amdgpu_metadata

extra/gemm/asm/cdna/test.py

@@ -1,72 +0,0 @@
-# Run assembly on the AMD runtime and check correctness
-# VIZ=2 to profile
-import pathlib
-from tinygrad import Tensor, Device, dtypes, Context
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.helpers import getenv
-
-fp = pathlib.Path(__file__).parent/"gemm.s"
-
-N = getenv("N", 8192)
-THREADS_PER_WG = 256
-NUM_WG = N//THREADS_PER_WG * N//THREADS_PER_WG
-
-assert N % THREADS_PER_WG == 0, "N must be divisible by THREADS_PER_WG"
-
-# ** generate inputs on CPU
-
-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 asm gemm
-C_torch = A@B
-
-# ** 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 = from_torch(A) @ from_torch(B)
-C_asm = Tensor.empty_like(C_tiny)
-
-# ** assembly custom kernel
-
-def custom_asm_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  lidx = UOp.special(THREADS_PER_WG, "lidx0")
-  gidx = UOp.special(NUM_WG, "gidx0")
-
-  src = (pathlib.Path(__file__).parent/"template.s").read_text().replace("INSTRUCTIONS", fp.read_text())
-
-  sz = UOp.variable("SZ", 256, 8192)
-
-  sink = UOp.sink(C.base, A.base, B.base, sz, lidx, gidx, arg=KernelInfo(name="gemm"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
-
-C_asm = Tensor.custom_kernel(C_asm, from_torch(A), from_torch(Bt), fxn=custom_asm_gemm)[0]
-
-# ** run gemms
-
-sched = Tensor.schedule(C_tiny, C_asm)
-eis = [si.lower() for si in sched]
-
-with Context(DEBUG=2):
-  for ei in eis:
-    et = ei.run({"SZ":N}, 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/rdna3/template.s

@@ -1,76 +0,0 @@
-.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

@@ -1,30 +0,0 @@
-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/unpack_kd.py

@@ -1,179 +0,0 @@
-# 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, 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(sys.argv[1], "rb") as f:
-  lib = f.read()
-
-image, sections, relocs = elf_loader(lib)
-rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"))
-
-# rodata is exactly 64 bytes
-kd = image[rodata_entry:rodata_entry+64]
-desc = int.from_bytes(kd, byteorder="little")
-
-group_segment_fixed_size = bits(desc, 0, 31)
-private_segment_fixed_size = bits(desc, 32, 63)
-kernarg_size = bits(desc, 64, 95)
-reserved_127_96 = bits(desc, 96, 127)
-assert reserved_127_96 == 0
-
-print("GROUP_SEGMENT_FIXED_SIZE:", group_segment_fixed_size)
-print("PRIVATE_SEGMENT_FIXED_SIZE:", private_segment_fixed_size)
-print("KERNARG_SIZE:", kernarg_size)
-print("RESERVED 127:96:", reserved_127_96)
-
-entry_off = bits(desc, 128, 191)
-
-# sign-extend manually if needed
-if entry_off & (1 << 63):
-  entry_off -= 1 << 64
-
-print("KERNEL_CODE_ENTRY_BYTE_OFFSET:", entry_off)
-
-kd_addr = 0x1840
-entry_addr = kd_addr + entry_off
-
-print("Computed entry address: 0x%016x" % entry_addr)
-print("256B aligned:", entry_addr % 256 == 0)
-
-pgm_rsrc3 = bits(desc, 352, 383)
-pgm_rsrc1 = bits(desc, 384, 415)
-pgm_rsrc2 = bits(desc, 416, 447)
-
-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 (gfx950)
-
-accum_offset_raw = bits(pgm_rsrc3, 0, 5)
-assert_zero(pgm_rsrc3, 6, 15)
-tg_split = bits(pgm_rsrc3, 16, 16)
-accum_offset_vgprs = (accum_offset_raw + 1) * 4
-print("RSRC3.ACCUM_OFFSET (AccVGPR index):", accum_offset_vgprs)
-print("RSRC3.TG_SPLIT:", tg_split)
-
-# rsrc 1
-
-vgpr_gran = bits(pgm_rsrc1, 0, 5)
-sgpr_gran = bits(pgm_rsrc1, 6, 9)
-assert_zero(pgm_rsrc1, 27, 28)
-
-# NOTE: this is vgprs + agprs
-vgprs_used = (vgpr_gran + 1) * 8
-assert 0 <= vgprs_used <= 512
-
-k = sgpr_gran // 2
-sgprs_used = (k + 1) * 16
-
-print("RSRC1.VGPRS:", vgprs_used)
-print("RSRC1.SGPRS:", sgprs_used)
-
-assert_zero(pgm_rsrc1, 10, 11)
-
-float_round_mode_32 = bits(pgm_rsrc1, 12, 13)
-float_round_mode_16_64 = bits(pgm_rsrc1, 15, 14)
-float_denorm_mode_32 = bits(pgm_rsrc1, 16, 17)
-float_denorm_mode_16_64 = bits(pgm_rsrc1, 18, 19)
-
-priv = bits(pgm_rsrc1, 20, 20)
-assert priv == 0
-enable_dx10_clamp_wg_rr_en = bits(pgm_rsrc1, 21, 21)
-debug_mode = bits(pgm_rsrc1, 22, 22)
-enable_ieee_mode = bits(pgm_rsrc1, 23, 23)
-bulky = bits(pgm_rsrc1, 24, 24)
-assert bulky == 0
-cdbg_user = bits(pgm_rsrc1, 25, 25)
-assert cdbg_user == 0
-fp16_ovfl = bits(pgm_rsrc1, 26, 26)
-assert_zero(pgm_rsrc1, 27, 28)  # reserved
-assert_zero(pgm_rsrc1, 29, 29)  # WGP_MODE (reserved on gfx9)
-assert_zero(pgm_rsrc1, 30, 30)  # MEM_ORDERED (reserved on gfx9)
-assert_zero(pgm_rsrc1, 31, 31)  # FWD_PROGRESS (reserved on gfx9)
-
-# rsrc 2
-
-enable_private_segment = bits(pgm_rsrc2, 0, 0)  # SCRATCH_EN
-user_sgpr_count = bits(pgm_rsrc2, 1, 5)         # USER_SGPR
-enable_trap_handler = bits(pgm_rsrc2, 6, 6)     # TRAP_PRESENT (must be 0 here)
-assert enable_trap_handler == 0
-
-enable_sgpr_workgroup_id_x = bits(pgm_rsrc2, 7, 7)
-enable_sgpr_workgroup_id_y = bits(pgm_rsrc2, 8, 8)
-enable_sgpr_workgroup_id_z = bits(pgm_rsrc2, 9, 9)
-enable_sgpr_workgroup_info = bits(pgm_rsrc2, 10, 10)
-
-enable_vgpr_workitem_id = bits(pgm_rsrc2, 11, 12)  # TIDIG_CMP_CNT enum (0..3)
-
-enable_exception_address_watch = bits(pgm_rsrc2, 13, 13)
-assert enable_exception_address_watch == 0
-enable_exception_memory = bits(pgm_rsrc2, 14, 14)
-assert enable_exception_memory == 0
-
-granulated_lds_size = bits(pgm_rsrc2, 15, 23)
-assert granulated_lds_size == 0  # spec: must be 0; CP uses dispatch packet rounding
-
-enable_exception_fp_invalid = bits(pgm_rsrc2, 24, 24)
-enable_exception_fp_denorm_src = bits(pgm_rsrc2, 25, 25)
-enable_exception_fp_div0 = bits(pgm_rsrc2, 26, 26)
-enable_exception_fp_overflow = bits(pgm_rsrc2, 27, 27)
-enable_exception_fp_underflow = bits(pgm_rsrc2, 28, 28)
-enable_exception_fp_inexact = bits(pgm_rsrc2, 29, 29)
-enable_exception_int_div0 = bits(pgm_rsrc2, 30, 30)
-
-assert_zero(pgm_rsrc2, 31, 31)
-
-print("RSRC2.ENABLE_PRIVATE_SEGMENT:", enable_private_segment)
-print("RSRC2.USER_SGPR_COUNT:", user_sgpr_count)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_X:", enable_sgpr_workgroup_id_x)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Y:", enable_sgpr_workgroup_id_y)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Z:", enable_sgpr_workgroup_id_z)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_INFO:", enable_sgpr_workgroup_info)
-print("RSRC2.ENABLE_VGPR_WORKITEM_ID (enum):", enable_vgpr_workitem_id)
-
-print("RSRC2.EXC_FP_INVALID:", enable_exception_fp_invalid)
-print("RSRC2.EXC_FP_DENORM_SRC:", enable_exception_fp_denorm_src)
-print("RSRC2.EXC_FP_DIV0:", enable_exception_fp_div0)
-print("RSRC2.EXC_FP_OVERFLOW:", enable_exception_fp_overflow)
-print("RSRC2.EXC_FP_UNDERFLOW:", enable_exception_fp_underflow)
-print("RSRC2.EXC_FP_INEXACT:", enable_exception_fp_inexact)
-print("RSRC2.EXC_INT_DIV0:", enable_exception_int_div0)
-
-# user sgprs
-
-enable_sgpr_private_segment_buffer = bits(desc, 448, 448)
-enable_sgpr_dispatch_ptr = bits(desc, 449, 449)
-enable_sgpr_queue_ptr = bits(desc, 450, 450)
-enable_sgpr_kernarg_segment_ptr = bits(desc, 451, 451)
-enable_sgpr_dispatch_id = bits(desc, 452, 452)
-enable_sgpr_flat_scratch_init = bits(desc, 453, 453)
-enable_sgpr_private_segment_size = bits(desc, 454, 454)
-
-assert_zero(desc, 455, 457)
-
-print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER:", enable_sgpr_private_segment_buffer)
-print("DESC.ENABLE_SGPR_DISPATCH_PTR:", enable_sgpr_dispatch_ptr)
-print("DESC.ENABLE_SGPR_QUEUE_PTR:", enable_sgpr_queue_ptr)
-print("DESC.ENABLE_SGPR_KERNARG_SEGMENT_PTR:", enable_sgpr_kernarg_segment_ptr)
-print("DESC.ENABLE_SGPR_DISPATCH_ID:", enable_sgpr_dispatch_id)
-print("DESC.ENABLE_SGPR_FLAT_SCRATCH_INIT:", enable_sgpr_flat_scratch_init)
-print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_SIZE:", enable_sgpr_private_segment_size)
-
-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)
-
-assert_zero(desc, 480, 511)

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

@@ -10,9 +10,9 @@ 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)
+    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist).realize()
     if outbuf is not None: outbuf.assign(x).realize()
-    return x.realize()
+    return x
   return TinyJit(hevc_decode_frame)
 
 def hevc_decode(hevc_tensor:Tensor, opaque:Tensor, frame_info:list, luma_h:int, luma_w:int,
@@ -74,10 +74,14 @@ if __name__ == "__main__":
   Device.default.synchronize()
 
   # 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")):
+  tm = Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps"))
+  with tm:
     images = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w, history=hist, preallocated_outputs=out_images))
     Device.default.synchronize()
 
+  fps = len(frame_info)/(tm.et/1e9)
+  assert fps >= getenv("ASSERT_FPS", 0), f"HEVC decode too slow: {fps:.2f} fps"
+
   # validation
   if getenv("VALIDATE", 0):
     import pickle

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}")
@@ -53,13 +53,13 @@ def print_aql_queue(read_pointer_address):
   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, field_type in hsa.amd_queue_v2_t._fields_:
+  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, '_fields_') else f"{v:#X}" for v in val]
+      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, '_fields_'): print(f"    {field_name}: {format_struct(val)}")
+    elif hasattr(val, '_real_fields_'): print(f"    {field_name}: {format_struct(val)}")
     else: print(f"    {field_name}: {val}")
 
 def ioctls_from_header():

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,100 +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.support.compiler_amd import HIPCompiler
-from tinygrad.runtime.ops_amd import AMDProgram
+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)

extra/models/resnet.py

@@ -150,7 +150,7 @@ class ResNet:
         continue # Skip FC if transfer learning
 
       if 'bn' not in k and 'downsample' not in k: assert obj.shape == dat.shape, (k, obj.shape, dat.shape)
-      obj.assign(dat.to(obj.device).reshape(obj.shape))
+      obj.assign(dat.to(obj.device).cast(obj.dtype).reshape(obj.shape))
 
 ResNet18 = lambda num_classes=1000: ResNet(18, num_classes=num_classes)
 ResNet34 = lambda num_classes=1000: ResNet(34, num_classes=num_classes)

extra/nv_gpu_driver/nv_ioctl.py

@@ -1,6 +1,7 @@
 # type: ignore
 import ctypes, ctypes.util, struct, platform, pathlib, re, time, os, signal
-from tinygrad.helpers import from_mv, to_mv, getenv, init_c_struct_t
+from tinygrad.helpers import from_mv, to_mv, getenv
+from tinygrad.runtime.support.c import init_c_struct_t
 from hexdump import hexdump
 start = time.perf_counter()
 
@@ -10,18 +11,21 @@ processor = platform.processor()
 IOCTL_SYSCALL = {"aarch64": 0x1d, "x86_64":16}[processor]
 MMAP_SYSCALL = {"aarch64": 0xde, "x86_64":0x09}[processor]
 
+IOCTL_PRINT = getenv("IOCTL_PRINT", getenv("IOCTL", 0))
+GRAB_PMA = getenv("GRAB_PMA", 0)
+
 def get_struct(argp, stype):
   return ctypes.cast(ctypes.c_void_p(argp), ctypes.POINTER(stype)).contents
 
 def dump_struct(st):
-  if getenv("IOCTL", 0) == 0: return
+  if IOCTL_PRINT == 0: return
   print("\t", st.__class__.__name__, end=" { ")
-  for v in type(st)._fields_: print(f"{v[0]}={getattr(st, v[0])}", end=" ")
+  for v in type(st)._real_fields_: print(f"{v[0]}={getattr(st, v[0])}", end=" ")
   print("}")
 
 def format_struct(s):
   sdats = []
-  for field in s._fields_:
+  for field in s._real_fields_:
     dat = getattr(s, field[0])
     if isinstance(dat, int): sdats.append(f"{field[0]}:0x{dat:X}")
     else: sdats.append(f"{field[0]}:{dat}")
@@ -58,6 +62,29 @@ def install_hook(c_function, python_function):
   return orig_func
 
 # *** ioctl lib end ***
+
+# PMA buffer tracking for raw PC sampling data (only when GRAB_PMA is enabled)
+pma_mem_handle = 0  # hMemPmaBuffer from ALLOC_PMA_STREAM
+pma_buffer_size = 0
+pma_buffer_va = 0   # actual mapped VA (found via /proc/self/maps)
+pma_get_offset = 0  # current read offset in ring buffer
+pma_pending_map = False  # flag to check for new mapping on next ioctl
+pma_maps_before = set()  # mappings before MAP_MEMORY
+pma_raw_dumps: list[bytes] = []  # raw PMA buffer dumps
+
+def get_pma_raw_dumps() -> list[bytes]: return pma_raw_dumps
+def clear_pma_raw_dumps(): pma_raw_dumps.clear()
+
+def get_proc_maps():
+  """Read current process memory mappings as set of (start, end) tuples."""
+  result = set()
+  with open("/proc/self/maps", "r") as f:
+    for line in f:
+      addr_range = line.split()[0]
+      start, end = addr_range.split("-")
+      result.add((int(start, 16), int(end, 16)))
+  return result
+
 from tinygrad.runtime.autogen import nv_570 as nv_gpu
 nvescs = {getattr(nv_gpu, x):x for x in dir(nv_gpu) if x.startswith("NV_ESC")}
 nvcmds = {getattr(nv_gpu, x):(x, getattr(nv_gpu, "struct_"+x+"_PARAMS", getattr(nv_gpu, "struct_"+x.replace("_CMD_", "_")+"_PARAMS", None))) for x in dir(nv_gpu) if \
@@ -69,6 +96,7 @@ def get_classes():
                    "NV20_SUBDEVICE_0"}
   for nm,val in nv_gpu.__dict__.items():
     if not isinstance(val, int): continue
+    if nm.endswith("PARAMETERS_MESSAGE_ID"): continue
     if 0x3000 < val < 0xffff: res[val] = nm
     if nm in known_classes: res[val] = nm
   return res
@@ -80,37 +108,92 @@ global_ioctl_id = 0
 gpus_user_modes = []
 gpus_mmio = []
 gpus_fifo = []
+offset_load = 0
 
 @ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_void_p)
 def ioctl(fd, request, argp):
   global global_ioctl_id, gpus_user_modes, gpus_mmio
+  global pma_mem_handle, pma_buffer_size, pma_buffer_va, pma_get_offset, pma_pending_map, pma_maps_before
   global_ioctl_id += 1
+
+  # Check for new PMA buffer mapping from previous MAP_MEMORY call (only when GRAB_PMA is enabled)
+  if GRAB_PMA and pma_pending_map:
+    pma_pending_map = False
+    new_maps = get_proc_maps()
+    for start, end in new_maps - pma_maps_before:
+      if end - start == pma_buffer_size:
+        pma_buffer_va = start
+        if IOCTL_PRINT >= 1: print(f"\t  PMA buffer mapped at CPU VA=0x{pma_buffer_va:x}")
+        break
+
   st = time.perf_counter()
   ret = libc.syscall(IOCTL_SYSCALL, ctypes.c_int(fd), ctypes.c_ulong(request), ctypes.c_void_p(argp))
   et = time.perf_counter()-st
   fn = os.readlink(f"/proc/self/fd/{fd}")
-  #print(f"ioctl {request:8x} {fn:20s}")
+
   idir, size, itype, nr = (request>>30), (request>>16)&0x3FFF, (request>>8)&0xFF, request&0xFF
-  if getenv("IOCTL", 0) >= 1: print(f"#{global_ioctl_id}: ", end="")
+  if IOCTL_PRINT >= 1: print(f"#{global_ioctl_id}: ", end="")
   if itype == ord(nv_gpu.NV_IOCTL_MAGIC):
     if nr == nv_gpu.NV_ESC_RM_CONTROL:
       s = get_struct(argp, nv_gpu.NVOS54_PARAMETERS)
       if s.cmd in nvcmds:
         name, struc = nvcmds[s.cmd]
-        if getenv("IOCTL", 0) >= 1:
+        if IOCTL_PRINT >= 1:
           print(f"NV_ESC_RM_CONTROL    cmd={name:30s} hClient={s.hClient}, hObject={s.hObject}, flags={s.flags}, params={s.params}, paramsSize={s.paramsSize}, status={s.status}")
 
         if struc is not None: dump_struct(get_struct(s.params, struc))
         elif hasattr(nv_gpu, name+"_PARAMS"): dump_struct(get_struct(argp, getattr(nv_gpu, name+"_PARAMS")))
         elif name == "NVA06C_CTRL_CMD_GPFIFO_SCHEDULE": dump_struct(get_struct(argp, nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS))
         elif name == "NV83DE_CTRL_CMD_GET_MAPPINGS": dump_struct(get_struct(s.params, nv_gpu.NV83DE_CTRL_DEBUG_GET_MAPPINGS_PARAMETERS))
+        elif name == "NVB0CC_CTRL_CMD_SET_HS_CREDITS":
+          hs_params = get_struct(s.params, nv_gpu.NVB0CC_CTRL_SET_HS_CREDITS_PARAMS)
+          dump_struct(hs_params)
+          if IOCTL_PRINT >= 2:
+            for i in range(hs_params.numEntries):
+              print(f"\t\t", end="")
+              dump_struct(hs_params.creditInfo[i])
+
+        # PMA buffer tracking (only when GRAB_PMA is enabled)
+        if GRAB_PMA and name == "NVB0CC_CTRL_CMD_ALLOC_PMA_STREAM":
+          pma_params = get_struct(s.params, nv_gpu.struct_NVB0CC_CTRL_ALLOC_PMA_STREAM_PARAMS)
+          pma_mem_handle = pma_params.hMemPmaBuffer
+          pma_buffer_size = pma_params.pmaBufferSize
+          pma_get_offset = 0  # Reset read offset for new stream
+          if IOCTL_PRINT >= 1: print(f"\t  PMA buffer: hMem=0x{pma_mem_handle:x} size={pma_buffer_size}")
+        if GRAB_PMA and name == "NVB0CC_CTRL_CMD_PMA_STREAM_UPDATE_GET_PUT":
+          pma_update = get_struct(s.params, nv_gpu.struct_NVB0CC_CTRL_PMA_STREAM_UPDATE_GET_PUT_PARAMS)
+          if pma_update.bytesAvailable > 0 and pma_buffer_va and pma_buffer_size > 0:
+            avail = pma_update.bytesAvailable
+            read_offset = pma_get_offset
+            # Handle ring buffer wrap-around
+            if pma_get_offset + avail <= pma_buffer_size:
+              pma_data = bytes(to_mv(pma_buffer_va + pma_get_offset, avail))
+            else:
+              # Wrap around: read to end, then from start
+              first_part = pma_buffer_size - pma_get_offset
+              second_part = avail - first_part
+              pma_data = bytes(to_mv(pma_buffer_va + pma_get_offset, first_part)) + bytes(to_mv(pma_buffer_va, second_part))
+            pma_raw_dumps.append(pma_data)
+            pma_get_offset = (pma_get_offset + avail) % pma_buffer_size
+            if IOCTL_PRINT >= 2:
+              print(f"\t  PMA data: {avail} bytes from offset=0x{read_offset:x}, new offset=0x{pma_get_offset:x}")
+              hexdump(pma_data)
+
+        # Dump regOps for EXEC_REG_OPS when IOCTL >= 3
+        if name == "NVB0CC_CTRL_CMD_EXEC_REG_OPS" and struc is not None and IOCTL_PRINT >= 3:
+          reg_params = get_struct(s.params, struc)
+          for i in range(reg_params.regOpCount):
+            print(f"\t\t", end="")
+            dump_struct(reg_params.regOps[i])
+            # val = (op.regValueHi << 32) | op.regValueLo
+            # print(f"\t  regOps[{i:3d}]: op={op.regOp} type={op.regType} status={op.regStatus} offset=0x{op.regOffset:08x} value=0x{val:016x}")
       else:
-        if getenv("IOCTL", 0) >= 1: print("unhandled cmd", hex(s.cmd))
+        if IOCTL_PRINT >= 1: print("unhandled cmd", hex(s.cmd))
       # format_struct(s)
       # print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : {ret:2d} = {name:40s}", ' '.join(format_struct(s)))
     elif nr == nv_gpu.NV_ESC_RM_ALLOC:
       s = get_struct(argp, nv_gpu.NVOS21_PARAMETERS)
-      if getenv("IOCTL", 0) >= 1: print(f"NV_ESC_RM_ALLOC    hClass={nvclasses.get(s.hClass, f'unk=0x{s.hClass:X}'):30s}, hRoot={s.hRoot}, hObjectParent={s.hObjectParent}, pAllocParms={s.pAllocParms}, hObjectNew={s.hObjectNew} status={s.status}")
+      if IOCTL_PRINT >= 1: print(f"NV_ESC_RM_ALLOC    hClass={nvclasses.get(s.hClass, f'unk=0x{s.hClass:X}'):30s}, hRoot={s.hRoot}, hObjectParent={s.hObjectParent}, pAllocParms={s.pAllocParms}, hObjectNew={s.hObjectNew} status={s.status}")
       if s.pAllocParms is not None:
         if s.hClass == nv_gpu.NV01_DEVICE_0: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV0080_ALLOC_PARAMETERS))
         if s.hClass == nv_gpu.FERMI_VASPACE_A: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_VASPACE_ALLOCATION_PARAMETERS))
@@ -118,7 +201,8 @@ def ioctl(fd, request, argp):
         if s.hClass == nv_gpu.NV1_MEMORY_USER: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_MEMORY_ALLOCATION_PARAMS))
         if s.hClass == nv_gpu.NV1_MEMORY_SYSTEM: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_MEMORY_ALLOCATION_PARAMS))
         if s.hClass == nv_gpu.GT200_DEBUGGER: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV83DE_ALLOC_PARAMETERS))
-        if s.hClass == nv_gpu.AMPERE_CHANNEL_GPFIFO_A:
+        if s.hClass == nv_gpu.MAXWELL_PROFILER_DEVICE: dump_struct(get_struct(s.pAllocParms, nv_gpu.NVB2CC_ALLOC_PARAMETERS))
+        if s.hClass in {nv_gpu.AMPERE_CHANNEL_GPFIFO_A, nv_gpu.BLACKWELL_CHANNEL_GPFIFO_A}:
           sx = get_struct(s.pAllocParms, nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS)
           dump_struct(sx)
           gpus_fifo.append((sx.gpFifoOffset, sx.gpFifoEntries))
@@ -126,31 +210,35 @@ def ioctl(fd, request, argp):
       if s.hClass == nv_gpu.TURING_USERMODE_A: gpus_user_modes.append(s.hObjectNew)
     elif nr == nv_gpu.NV_ESC_RM_MAP_MEMORY:
       # nv_ioctl_nvos33_parameters_with_fd
-      if getenv("IOCTL", 0) >= 1:
-        s = get_struct(argp, nv_gpu.NVOS33_PARAMETERS)
+      s = get_struct(argp, nv_gpu.NVOS33_PARAMETERS)
+      if IOCTL_PRINT >= 1:
         print(f"NV_ESC_RM_MAP_MEMORY   hClient={s.hClient}, hDevice={s.hDevice}, hMemory={s.hMemory}, length={s.length} flags={s.flags} pLinearAddress={s.pLinearAddress}")
+      # Track PMA buffer mapping - save maps now, check for new mapping on next ioctl (after mmap happens)
+      if GRAB_PMA and pma_mem_handle and s.hMemory == pma_mem_handle:
+        pma_maps_before = get_proc_maps()
+        pma_pending_map = True
     elif nr == nv_gpu.NV_ESC_RM_UPDATE_DEVICE_MAPPING_INFO:
-      if getenv("IOCTL", 0) >= 1:
+      if IOCTL_PRINT >= 1:
         s = get_struct(argp, nv_gpu.NVOS56_PARAMETERS)
         print(f"NV_ESC_RM_UPDATE_DEVICE_MAPPING_INFO   hClient={s.hClient}, hDevice={s.hDevice}, hMemory={s.hMemory}, pOldCpuAddress={s.pOldCpuAddress} pNewCpuAddress={s.pNewCpuAddress} status={s.status}")
     elif nr == nv_gpu.NV_ESC_RM_ALLOC_MEMORY:
-      if getenv("IOCTL", 0) >= 1:
+      if IOCTL_PRINT >= 1:
         s = get_struct(argp, nv_gpu.nv_ioctl_nvos02_parameters_with_fd)
         print(f"NV_ESC_RM_ALLOC_MEMORY  fd={s.fd}, hRoot={s.params.hRoot}, hObjectParent={s.params.hObjectParent}, hObjectNew={s.params.hObjectNew}, hClass={s.params.hClass}, flags={s.params.flags}, pMemory={s.params.pMemory}, limit={s.params.limit}, status={s.params.status}")
     elif nr == nv_gpu.NV_ESC_ALLOC_OS_EVENT:
-      if getenv("IOCTL", 0) >= 1:
+      if IOCTL_PRINT >= 1:
         s = get_struct(argp, nv_gpu.nv_ioctl_alloc_os_event_t)
         print(f"NV_ESC_ALLOC_OS_EVENT  hClient={s.hClient} hDevice={s.hDevice} fd={s.fd} Status={s.Status}")
     elif nr == nv_gpu.NV_ESC_REGISTER_FD:
-      if getenv("IOCTL", 0) >= 1:
+      if IOCTL_PRINT >= 1:
         s = get_struct(argp, nv_gpu.nv_ioctl_register_fd_t)
         print(f"NV_ESC_REGISTER_FD  fd={s.ctl_fd}")
     elif nr in nvescs:
-      if getenv("IOCTL", 0) >= 1: print(nvescs[nr])
+      if IOCTL_PRINT >= 1: print(nvescs[nr])
     else:
-      if getenv("IOCTL", 0) >= 1: print("unhandled NR", nr)
+      if IOCTL_PRINT >= 1: print("unhandled NR", nr)
   elif fn.endswith("nvidia-uvm"):
-    if getenv("IOCTL", 0) >= 1:
+    if IOCTL_PRINT >= 1:
       print(f"{nvuvms.get(request, f'UVM UNKNOWN {request=}')}")
       if nvuvms.get(request) is not None: dump_struct(get_struct(argp, getattr(nv_gpu, nvuvms.get(request)+"_PARAMS")))
       if nvuvms.get(request) == "UVM_MAP_EXTERNAL_ALLOCATION":
@@ -159,7 +247,7 @@ def ioctl(fd, request, argp):
           print("perGpuAttributes[{i}] = ", end="")
           dump_struct(st.perGpuAttributes[i])
 
-  if getenv("IOCTL") >= 2: print("ioctl", f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", fn)
+  if IOCTL_PRINT >= 2: print("ioctl", f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", fn)
   return ret
 
 @ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long)
@@ -172,14 +260,14 @@ def _mmap(addr, length, prot, flags, fd, offset):
   return ret
 
 install_hook(libc.ioctl, ioctl)
-if getenv("IOCTL") >= 3: orig_mmap_mv = install_hook(libc.mmap, _mmap)
+if getenv("IOCTL") >= 4: orig_mmap_mv = install_hook(libc.mmap, _mmap)
 
 import collections
 old_gpputs = collections.defaultdict(int)
 def _dump_gpfifo(mark):
   launches = []
 
-  # print("_dump_gpfifo:", mark)
+  print("_dump_gpfifo:", mark)
   for start, size in gpus_fifo:
     gpfifo_controls = nv_gpu.AmpereAControlGPFifo.from_address(start+size*8)
     gpfifo = to_mv(start, size * 8).cast("Q")
@@ -205,7 +293,7 @@ def make_qmd_struct_type():
     fields.append((name.replace("NVC6C0_QMDV03_00_", "").lower(), ctypes.c_uint32, data[0]-data[1]+1))
     if len(fields) >= 2 and fields[-2][0].endswith('_lower') and fields[-1][0].endswith('_upper') and fields[-1][0][:-6] == fields[-2][0][:-6]:
       fields = fields[:-2] + [(fields[-1][0][:-6], ctypes.c_uint64, fields[-1][2] + fields[-2][2])]
-  return init_c_struct_t(tuple(fields))
+  return init_c_struct_t(0x40 * 4, tuple(fields))
 qmd_struct_t = make_qmd_struct_type()
 assert ctypes.sizeof(qmd_struct_t) == 0x40 * 4
 
@@ -222,7 +310,7 @@ def _dump_qmd(address, packets):
     subc = (dat>>13) & 7
     mthd = (dat<<2) & 0x7FFF
     method_name = nvqcmds.get(mthd, f"unknown method #{mthd}")
-    if getenv("IOCTL", 0) >= 1:
+    if IOCTL_PRINT >= 1:
       print(f"\t\t{method_name}, {typ=} {size=} {subc=} {mthd=}")
       for j in range(size): print(f"\t\t\t{j}: {gpfifo[i+j+1]} | 0x{gpfifo[i+j+1]:x}")
     if mthd == 792:

extra/nv_pma/.gitignore

@@ -0,0 +1 @@
+examples/

extra/nv_pma/collect.py

@@ -0,0 +1,135 @@
+import pickle, os, sys, functools, numpy as np
+from pathlib import Path
+
+os.environ["DEV"] = "CUDA"
+os.environ["PROFILE"] = os.environ.get("PROFILE", "2")
+from extra.nv_pma.cupti import cu_prof_ext
+cu_prof_ext.enable_auto()
+
+from tinygrad import Tensor, Device
+
+if not os.environ.get("IOCTL") or not os.environ.get("GRAB_PMA"):
+  print("Usage: GRAB_PMA=1 IOCTL=1 IOCTL_PRINT=0 python3 extra/nv_pma/collect.py")
+  sys.exit(1)
+
+assert Device.DEFAULT == "CUDA", "only works with CUDA"
+
+EXAMPLES_DIR = Path(__file__).parent / "examples"
+_collectors: list[tuple[str, callable]] = []
+
+def pcsampling_test(name: str):
+  def decorator(fn):
+    @functools.wraps(fn)
+    def wrapper():
+      cu_prof_ext.clear_pma_raw_dumps()
+      cu_prof_ext.clear_cupti_pc_samples()
+
+      fn()
+      Device["CUDA"].synchronize()
+
+      dumps = cu_prof_ext.get_pma_raw_dumps()
+      # from hexdump import hexdump
+      # hexdump(dumps[0][:0x40])
+
+      return {"test_name": name, "pma_raw_dumps": list(cu_prof_ext.get_pma_raw_dumps()), "cupti_pc_samples": list(cu_prof_ext.get_cupti_pc_samples())}
+    _collectors.append((name, wrapper))
+    return wrapper
+  return decorator
+
+# Refs
+
+@pcsampling_test("test_plus")
+def test_plus():
+  a = Tensor([1, 2, 3, 4])
+  b = Tensor([5, 6, 7, 8])
+  (a + b).realize()
+
+@pcsampling_test("test_matmul")
+def test_matmul():
+  a = Tensor(np.random.rand(12, 12).astype(np.float32))
+  b = Tensor(np.random.rand(12, 12).astype(np.float32))
+  (a @ b).realize()
+
+@pcsampling_test("test_reduce_sum")
+def test_reduce_sum():
+  a = Tensor(np.random.rand(1024).astype(np.float32))
+  a.sum().realize()
+
+@pcsampling_test("test_reduce_max")
+def test_reduce_max():
+  a = Tensor(np.random.rand(1024).astype(np.float32))
+  a.max().realize()
+
+@pcsampling_test("test_exp")
+def test_exp():
+  a = Tensor(np.random.rand(256).astype(np.float32))
+  a.exp().realize()
+
+@pcsampling_test("test_softmax")
+def test_softmax():
+  a = Tensor(np.random.rand(64, 64).astype(np.float32))
+  a.softmax().realize()
+
+@pcsampling_test("test_conv2d")
+def test_conv2d():
+  x = Tensor(np.random.rand(1, 3, 32, 32).astype(np.float32))
+  w = Tensor(np.random.rand(8, 3, 3, 3).astype(np.float32))
+  x.conv2d(w).realize()
+
+@pcsampling_test("test_large_matmul")
+def test_large_matmul():
+  a = Tensor(np.random.rand(128, 128).astype(np.float32))
+  b = Tensor(np.random.rand(128, 128).astype(np.float32))
+  (a @ b).realize()
+
+@pcsampling_test("test_elementwise_chain")
+def test_elementwise_chain():
+  a = Tensor(np.random.rand(512).astype(np.float32))
+  ((a + 1) * 2 - 0.5).relu().realize()
+
+@pcsampling_test("test_broadcast")
+def test_broadcast():
+  a = Tensor(np.random.rand(64, 1).astype(np.float32))
+  b = Tensor(np.random.rand(1, 64).astype(np.float32))
+  (a + b).realize()
+
+@pcsampling_test("test_plus_big")
+def test_plus_big():
+  a = Tensor(np.random.rand(64, 32).astype(np.float32))
+  b = Tensor(np.random.rand(64, 32).astype(np.float32))
+  (a + b).realize()
+
+def save_example(name: str, data: dict):
+  pma_bytes = sum(len(d) for d in data['pma_raw_dumps'])
+  cupti_samples = sum(r['samples'] for r in data['cupti_pc_samples'])
+  print(f"  PMA: {len(data['pma_raw_dumps'])} buffers, {pma_bytes} bytes")
+  print(f"  CUPTI: {len(data['cupti_pc_samples'])} records, {cupti_samples} samples")
+
+  outfile = EXAMPLES_DIR / f"{name}.pkl"
+  with open(outfile, "wb") as f:
+    pickle.dump(data, f)
+  print(f"  Saved to {outfile}")
+
+if __name__ == "__main__":
+  EXAMPLES_DIR.mkdir(exist_ok=True)
+
+  # Run specific tests if provided as arguments, otherwise run all
+  if len(sys.argv) > 1:
+    test_names = sys.argv[1:]
+    collectors = [(name, fn) for name, fn in _collectors if name in test_names]
+    if not collectors:
+      print(f"Unknown tests: {test_names}")
+      print(f"Available: {[name for name, _ in _collectors]}")
+      sys.exit(1)
+  else:
+    collectors = _collectors
+
+  for name, collect_fn in collectors:
+    print(f"\nCollecting {name}...")
+    try:
+      data = collect_fn()
+      save_example(name, data)
+    except Exception as e:
+      print(f"  ERROR: {e}")
+      import traceback
+      traceback.print_exc()

extra/nv_pma/cupti/__init__.py

@@ -0,0 +1,26 @@
+# CUPTI autogen loader for nv_pma
+# To regenerate: REGEN=1 python -c "import extra.nv_pma.cupti"
+import importlib, pathlib
+from tinygrad.helpers import getenv
+
+root = pathlib.Path(__file__).parents[3]
+here = pathlib.Path(__file__).parent
+
+def load(name, dll, files, **kwargs):
+  if not (f:=here/f"{name}.py").exists() or getenv('REGEN'):
+    kwargs['args'] = kwargs.get('args', [])
+    f.write_text(importlib.import_module("tinygrad.runtime.support.autogen").gen(name, dll, files, **kwargs))
+  return importlib.import_module(f"extra.nv_pma.cupti.{name}")
+
+def __getattr__(nm):
+  match nm:
+    case "cupti":
+      return load("cupti", "'/usr/local/cuda/targets/x86_64-linux/lib/libcupti.so'", [
+        "/usr/local/cuda/include/cupti_result.h", "/usr/local/cuda/include/cupti_activity.h",
+        "/usr/local/cuda/include/cupti_callbacks.h", "/usr/local/cuda/include/cupti_events.h",
+        "/usr/local/cuda/include/cupti_metrics.h", "/usr/local/cuda/include/cupti_driver_cbid.h",
+        "/usr/local/cuda/include/cupti_runtime_cbid.h", "/usr/local/cuda/include/cupti_profiler_target.h",
+        "/usr/local/cuda/include/cupti_profiler_host.h", "/usr/local/cuda/include/cupti_pmsampling.h",
+        "/usr/local/cuda/include/generated_cuda_meta.h", "/usr/local/cuda/include/generated_cuda_runtime_api_meta.h"
+      ], args=["-D__CUDA_API_VERSION_INTERNAL", "-I/usr/local/cuda/include"], parse_macros=False)
+    case _: raise AttributeError(f"no such autogen: {nm}")

extra/nv_pma/cupti/cu_prof_ext.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+import ctypes
+from tinygrad.helpers import DEBUG, getenv
+from extra.nv_pma.cupti import cupti
+
+def stall_reason_name(reason: int) -> str:
+  name = cupti.CUpti_ActivityPCSamplingStallReason.get(reason)
+  return name.replace("CUPTI_ACTIVITY_PC_SAMPLING_STALL_", "").lower() if name else str(reason)
+
+class CUPTIProfiler:
+  def __init__(self):
+    self.initialized = False
+    self.pc_sampling_enabled = False
+    self.buffers: list[ctypes.Array] = []
+    self.kernel_stalls: dict[int, dict[int, int]] = {}
+    self.raw_buffers: list[bytes] = []
+    self.pc_samples: list[dict] = []
+
+  def _check_cupti(self, status, soft=False):
+    if status != cupti.CUPTI_SUCCESS:
+      if soft: return False
+      raise RuntimeError(f"CUPTI Error {status}")
+    return True
+
+  def init(self, ctx, device_id: int = 0, profile_level: int = 2):
+    if self.initialized: return
+
+    # Initialize profiler API
+    init_params = cupti.CUpti_Profiler_Initialize_Params()
+    init_params.structSize = 16
+    cupti.cuptiProfilerInitialize(ctypes.byref(init_params))
+
+    # Register buffer callbacks for Activity API
+    self._buf_req_cb = cupti.CUpti_BuffersCallbackRequestFunc(self._buffer_requested)
+    self._buf_comp_cb = cupti.CUpti_BuffersCallbackCompleteFunc(self._buffer_completed)
+    self._check_cupti(cupti.cuptiActivityRegisterCallbacks(self._buf_req_cb, self._buf_comp_cb))
+
+    # PROFILE=1: kernel timing, PROFILE=2: PC sampling with stall reasons
+    if profile_level >= 2:
+      # PC sampling for stall analysis (requires elevated privileges)
+      if DEBUG >= 1: print("  CUPTI: PC sampling mode (before)")
+      pc_status = cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING)
+      if pc_status == cupti.CUPTI_SUCCESS:
+        config = cupti.CUpti_ActivityPCSamplingConfig()
+        config.size, config.samplingPeriod = 16, cupti.CUPTI_ACTIVITY_PC_SAMPLING_PERIOD_MIN
+        cfg_status = cupti.dll.cuptiActivityConfigurePCSampling(ctx, ctypes.byref(config))
+        if cfg_status == cupti.CUPTI_SUCCESS:
+          if DEBUG >= 1: print("  CUPTI: PC sampling mode (before stall analysis)")
+          cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING_RECORD_INFO)
+          self.pc_sampling_enabled = True
+          if DEBUG >= 1: print("  CUPTI: PC sampling mode (stall analysis)")
+        elif cfg_status == 35:
+          if DEBUG >= 1: print("  CUPTI: PC sampling needs: echo 'options nvidia NVreg_RestrictProfilingToAdminUsers=0'|sudo tee /etc/modprobe.d/nvidia.conf && sudo reboot")
+      # Fall back to kernel timing if PC sampling setup failed
+      if not self.pc_sampling_enabled:
+        self._check_cupti(cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_KERNEL))
+    else:
+      # Kernel activity tracing for timing
+      self._check_cupti(cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_KERNEL))
+
+    self.initialized = True
+
+  def _buffer_requested(self, buffer, size, max_num_records):
+    buf = (ctypes.c_uint8 * 1024 * 1024)()  # 1MB buffer
+    self.buffers.append(buf)
+    buffer[0] = ctypes.cast(buf, ctypes.POINTER(ctypes.c_uint8))
+    size[0] = ctypes.sizeof(buf)
+    max_num_records[0] = 0
+
+  def _buffer_completed(self, ctx, stream_id, buffer, size, valid_size):
+    if valid_size > 0:
+      record = ctypes.POINTER(cupti.CUpti_Activity)()
+      while cupti.cuptiActivityGetNextRecord(buffer, valid_size, ctypes.byref(record)) == cupti.CUPTI_SUCCESS:
+        kind = record.contents.kind
+        if kind == cupti.CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL:
+          kernel = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityKernel9)).contents
+          name = ctypes.string_at(kernel.name).decode() if kernel.name else "unknown"
+          duration_us = (kernel.end - kernel.start) / 1000.0
+          grid, block = (kernel.gridX, kernel.gridY, kernel.gridZ), (kernel.blockX, kernel.blockY, kernel.blockZ)
+          print(f"  CUPTI: {name[:40]:40s} | {duration_us:10.2f} us | grid={grid} block={block} | regs={kernel.registersPerThread:3d} smem={kernel.staticSharedMemory + kernel.dynamicSharedMemory:6d}B")
+        elif kind == cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING:
+          pc = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityPCSampling3)).contents
+          cid = pc.correlationId
+          if cid not in self.kernel_stalls: self.kernel_stalls[cid] = {}
+          self.kernel_stalls[cid][pc.stallReason] = self.kernel_stalls[cid].get(pc.stallReason, 0) + pc.samples
+          self.pc_samples.append({
+            'correlationId': pc.correlationId, 'pcOffset': pc.pcOffset, 'stallReason': pc.stallReason,
+            'samples': pc.samples, 'latencySamples': pc.latencySamples, 'functionId': pc.functionId, 'sourceLocatorId': pc.sourceLocatorId
+          })
+          if DEBUG >= 3:
+            print(f"    PC {pc.pcOffset:#x} stall={stall_reason_name(pc.stallReason)} samples={pc.samples} latency={pc.latencySamples} func={pc.functionId} src={pc.sourceLocatorId}")
+        elif kind == cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING_RECORD_INFO:
+          info = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityPCSamplingRecordInfo)).contents
+          cid = info.correlationId
+          if cid in self.kernel_stalls:
+            stalls = self.kernel_stalls[cid]
+            total = sum(stalls.values())
+            if total > 0:
+              top = sorted(stalls.items(), key=lambda x: -x[1])[:5]
+              stall_str = " ".join(f"{stall_reason_name(r)}:{100*c//total}%" for r,c in top if c > 0)
+              print(f"  CUPTI stalls (corr={cid}): {total} samples | {stall_str}")
+            del self.kernel_stalls[cid]
+        else: print(f"  CUPTI: Unhandled activity kind {kind}")
+
+  def flush(self):
+    if not self.initialized: return
+    self._check_cupti(cupti.cuptiActivityFlushAll(0))
+
+# Module-level profiler instance
+_profiler: CUPTIProfiler | None = None
+
+def get_profiler() -> CUPTIProfiler | None:
+  return _profiler
+
+def get_cupti_raw_buffers() -> list[bytes]:
+  return _profiler.raw_buffers if _profiler else []
+
+def clear_cupti_raw_buffers():
+  if _profiler: _profiler.raw_buffers.clear()
+
+def get_cupti_pc_samples() -> list[dict]:
+  return _profiler.pc_samples if _profiler else []
+
+def clear_cupti_pc_samples():
+  if _profiler: _profiler.pc_samples.clear()
+
+# Raw PMA buffer access (from ioctl interception)
+def get_pma_raw_dumps() -> list[bytes]:
+  try:
+    from extra.nv_gpu_driver.nv_ioctl import get_pma_raw_dumps as _get
+    return _get()
+  except ImportError: return []
+
+def clear_pma_raw_dumps():
+  try:
+    from extra.nv_gpu_driver.nv_ioctl import clear_pma_raw_dumps as _clear
+    _clear()
+  except ImportError: pass
+
+def enable(profile_level:int=2):
+  global _profiler
+  if _profiler is not None: return
+
+  _profiler = CUPTIProfiler()
+
+  # Patch CUDADevice to initialize CUPTI profiler
+  from tinygrad.runtime.ops_cuda import CUDADevice
+  _orig_init = CUDADevice.__init__
+  _orig_sync = CUDADevice.synchronize
+
+  def _patched_init(self, device: str):
+    _orig_init(self, device)
+    device_id = int(device.split(":")[1]) if ":" in device else 0
+    _profiler.init(self.context, device_id, profile_level)
+
+  def _patched_sync(self):
+    _orig_sync(self)
+    if _profiler: _profiler.flush()
+
+  CUDADevice.__init__ = _patched_init
+  CUDADevice.synchronize = _patched_sync
+
+def enable_auto():
+  if (profile_level:=getenv("PROFILE", 0)) > 0: enable(profile_level)

extra/nv_pma/decode.py

@@ -0,0 +1,189 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+import enum, collections
+from typing import Iterator
+from tinygrad.helpers import colored
+from tinygrad.renderer.amd.sqtt import PacketType, bits
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# STALL REASONS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class StallReason(enum.IntEnum):
+  # Based on CUpti_ActivityPCSamplingStallReason
+  INVALID = 0
+  NONE = 1              # selected, selected_not_issued
+  INST_FETCH = 2        # branch_resolving, no_instructions
+  EXEC_DEPENDENCY = 3   # short_scoreboard, wait
+  MEMORY_DEPENDENCY = 4 # long_scoreboard
+  TEXTURE = 5           # tex_throttle
+  SYNC = 6              # barrier, membar
+  CONSTANT_MEMORY = 7   # imc_miss
+  PIPE_BUSY = 8         # mio_throttle, math_pipe_throttle
+  MEMORY_THROTTLE = 9   # drain, lg_throttle
+  NOT_SELECTED = 10     # not_selected
+  OTHER = 11            # misc, dispatch_stall
+  SLEEPING = 12         # sleeping
+
+STALL_KEY_MAP_AMPERE: dict[int, StallReason] = {
+  1: StallReason.MEMORY_THROTTLE, 15: StallReason.MEMORY_THROTTLE,
+  2: StallReason.CONSTANT_MEMORY,
+  3: StallReason.SYNC,
+  6: StallReason.INST_FETCH, 11: StallReason.INST_FETCH,
+  7: StallReason.EXEC_DEPENDENCY, 10: StallReason.EXEC_DEPENDENCY,
+  9: StallReason.MEMORY_DEPENDENCY,
+  12: StallReason.PIPE_BUSY,
+  17: StallReason.OTHER, 20: StallReason.OTHER,
+  18: StallReason.NONE,
+}
+
+STALL_KEY_MAP_BLACKWELL: dict[int, StallReason] = {
+  0x01: StallReason.MEMORY_THROTTLE, 0x0e: StallReason.MEMORY_THROTTLE,
+  0x02: StallReason.SYNC,
+  0x05: StallReason.INST_FETCH, 0x0a: StallReason.INST_FETCH,
+  0x06: StallReason.EXEC_DEPENDENCY, 0x09: StallReason.EXEC_DEPENDENCY,
+  0x08: StallReason.MEMORY_DEPENDENCY,
+  0x0b: StallReason.PIPE_BUSY, 0x0f: StallReason.PIPE_BUSY,
+  0x10: StallReason.OTHER, 0x13: StallReason.OTHER,
+  0x11: StallReason.NONE,
+}
+
+# Lookup table for extracting sample bytes from 32-byte packet (bytes 0-3, 8-31, skipping header at 4-7)
+LOOKUP_28B = [0, 1, 2, 3, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PACKET HEADER
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class PMAHeader(PacketType):
+  num_bytes  = bits[4:0]    # number of sample bytes in this packet
+  tpc_id_lo  = bits[15:8]   # TPC identifier low 8 bits
+  tpc_id_hi  = bits[27:25]  # TPC identifier high 3 bits
+  dropped    = bits[28:28]  # dropped flag (resets byte accumulator)
+  @property
+  def tpc_id(self) -> int: return self.tpc_id_lo | (self.tpc_id_hi << 8)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# 8-BYTE SAMPLE FORMAT (Ampere/Ada/Hopper)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class PMASampleAmpere8B(PacketType):
+  pc_raw     = bits[44:0]   # raw PC value (pc_offset = pc_raw << 4)
+  stall_key  = bits[49:45]  # stall reason key
+  wave_id    = bits[55:50]  # warp/wave identifier
+  active     = bits[62:62]  # 1 if warp was executing, 0 if scheduled but not issued
+  @property
+  def pc_offset(self) -> int: return self.pc_raw << 4
+  @property
+  def stall_reason(self) -> StallReason: return STALL_KEY_MAP_AMPERE.get(self.stall_key, StallReason.OTHER)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# 9-BYTE SAMPLE FORMAT (Blackwell+)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class PMASampleBlackwell9B(PacketType):
+  stall_key  = bits[5:0]    # stall reason key
+  pc_raw     = bits[60:8]   # raw PC value (pc_offset = pc_raw << 4)
+  wave_hi    = bits[7:6]    # wave_id high 2 bits
+  wave_lo    = bits[71:68]  # wave_id low 4 bits
+  active     = bits[67:67]  # 1 if warp was executing, 0 if scheduled but not issued
+  @property
+  def pc_offset(self) -> int: return self.pc_raw << 4
+  @property
+  def stall_reason(self) -> StallReason: return STALL_KEY_MAP_BLACKWELL.get(self.stall_key, StallReason.OTHER)
+  @property
+  def wave_id(self) -> int: return (self.wave_hi << 4) | self.wave_lo
+
+PMASample = PMASampleAmpere8B|PMASampleBlackwell9B
+
+def decode(data: bytes, sm_version: int = 0x800) -> Iterator[tuple[PMASample, int]]:
+  use_9byte = sm_version >= 0xa04
+  record_size = 9 if use_9byte else 8
+  sample_cls = PMASampleBlackwell9B if use_9byte else PMASampleAmpere8B
+
+  tpc_state: dict[int, list[int]] = collections.defaultdict(list)
+  for pkt_idx in range(len(data) // 32):
+    pkt = data[pkt_idx * 32:(pkt_idx + 1) * 32]
+    hdr = PMAHeader.from_raw(int.from_bytes(pkt[4:8], 'little'))
+
+    if hdr.dropped: tpc_state[hdr.tpc_id].clear()
+
+    for i in range(hdr.num_bytes):
+      tpc_state[hdr.tpc_id].append(pkt[LOOKUP_28B[i]])
+
+    while len(tpc_state[hdr.tpc_id]) >= record_size:
+      yield sample_cls.from_raw(int.from_bytes(bytes(tpc_state[hdr.tpc_id][:record_size]), 'little')), hdr.tpc_id
+      del tpc_state[hdr.tpc_id][:record_size]
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# CLI
+# ═══════════════════════════════════════════════════════════════════════════════
+
+STALL_COLORS = {
+  StallReason.NONE: "green", StallReason.INST_FETCH: "yellow", StallReason.EXEC_DEPENDENCY: "cyan",
+  StallReason.MEMORY_DEPENDENCY: "red", StallReason.SYNC: "magenta", StallReason.CONSTANT_MEMORY: "blue",
+  StallReason.PIPE_BUSY: "yellow", StallReason.MEMORY_THROTTLE: "RED", StallReason.OTHER: "white",
+}
+
+def decode_tpc_id(tpc_id:int) -> tuple[int, int, int]:
+  # NOTE: valid only for ops_nv, cuda encoding is different
+  return (tpc_id >> 5, (tpc_id >> 1) & 0xf, tpc_id & 1)
+
+def print_packets(data:bytes, sm_version:int=0x800) -> None:
+  record_size = 9 if sm_version >= 0x890 else 8
+  tpc_state: dict[int, list[int]] = collections.defaultdict(list)
+  for i in range(len(data) // 32):
+    pkt = data[i * 32:(i + 1) * 32]
+    hdr = PMAHeader.from_raw(int.from_bytes(pkt[4:8], 'little'))
+    if hdr.dropped: tpc_state[hdr.tpc_id].clear()
+    for j in range(hdr.num_bytes): tpc_state[hdr.tpc_id].append(pkt[LOOKUP_28B[j]])
+    # Show complete records extracted from this packet
+    records = []
+    while len(tpc_state[hdr.tpc_id]) >= record_size:
+      records.append(bytes(tpc_state[hdr.tpc_id][:record_size]).hex())
+      del tpc_state[hdr.tpc_id][:record_size]
+    leftover = len(tpc_state[hdr.tpc_id])
+    print(f"Pkt {i:3d}: tpc={hdr.tpc_id:4d} n={hdr.num_bytes:2d} drop={hdr.dropped} left={leftover} | {' '.join(records)}")
+
+def print_aggregated(samples:list[tuple[PMASample, int]]) -> None:
+  if not samples: return
+  base_pc = min(s.pc_offset for s, _ in samples)
+  counter: collections.Counter[tuple[int, StallReason]] = collections.Counter((s.pc_offset, s.stall_reason) for s, _ in samples)
+  print(f"\nAggregated samples (base_pc=0x{base_pc:x}):")
+  for (pc, reason), cnt in sorted(counter.items()):
+    stall_str = colored(f"{reason.name:17}", STALL_COLORS.get(reason, "white"))
+    print(f"  pc=0x{pc - base_pc:06x} {stall_str} samples={cnt:4d}")
+
+if __name__ == "__main__":
+  import sys, pickle
+
+  if len(sys.argv) < 2:
+    print("Usage: python decode.py <pkl_file> [--raw] [--sm=0xNNN]")
+    sys.exit(1)
+
+  with open(sys.argv[1], "rb") as f:
+    data = pickle.load(f)
+
+  if isinstance(data, dict):
+    sm_version = 0x800  # default to Ampere
+    for arg in sys.argv:
+      if arg.startswith("--sm="): sm_version = int(arg[5:], 0)
+    dumps = [(i, x, sm_version) for i, x in enumerate(data["pma_raw_dumps"])]
+  else:
+    devs = {e.device: e for e in data if type(e).__name__ == "ProfileDeviceEvent"}
+    dumps = []
+    for i, e in enumerate(e for e in data if type(e).__name__ == "ProfilePMAEvent"):
+      dumps.append((i, e.blob, devs[e.device].props.get('sm_version', 0x800)))
+
+  for dump_idx, raw, sm_ver in dumps:
+    print(f"\n{'='*60}\nDump {dump_idx} ({len(raw)} bytes, {len(raw)//32} packets)\n{'='*60}")
+    if "--raw" in sys.argv: print_packets(raw, sm_ver)
+    else:
+      samples = []
+      for s, tpc_id in decode(raw, sm_ver):
+        gpc, tpc, sm = decode_tpc_id(tpc_id)
+        stall_str = colored(f"{s.stall_reason.name:17}", STALL_COLORS.get(s.stall_reason, "white"))
+        print(f"pc=0x{s.pc_offset:06x} {stall_str} ev={s.stall_key:2d} active={s.active} wave={s.wave_id:2d} gpc={gpc} tpc={tpc} sm={sm}")
+        samples.append((s, tpc_id))
+      print(f"\nDecoded {len(samples)} samples:")
+      print_aggregated(samples)