no long

.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 && CHECK_OOB=0 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 && IGNORE_OOB=1 PYTHONPATH=. python3 process_replay.py
         git checkout $CURRENT_HEAD  # restore to branch

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

@@ -45,10 +45,6 @@ inputs:
     description: "Install mesa"
     required: false
     default: 'false'
-  tinydreno:
-    description: "Install tinydreno"
-    required: false
-    default: 'false'
 runs:
   using: "composite"
   steps:
@@ -60,40 +56,32 @@ 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 }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
 
     # **** Caching downloads ****
 
-    - name: Cache downloads (PR)
-      if: inputs.key != '' && github.event_name == 'pull_request'
-      uses: actions/cache/restore@v4
-      with:
-        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
-        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
-    - name: Cache downloads
-      if: inputs.key != '' && github.event_name != 'pull_request'
+    - name: Cache downloads (Linux)
+      if: inputs.key != '' && runner.os == 'Linux'
       uses: actions/cache@v4
       with:
-        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
+        path: ~/.cache/tinygrad/downloads/
+        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
+    - name: Cache downloads (macOS)
+      if: inputs.key != '' && runner.os == 'macOS'
+      uses: actions/cache@v4
+      with:
+        path: ~/Library/Caches/tinygrad/downloads/
         key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
 
     # **** Python deps ****
 
     - name: Install dependencies in venv (with extra)
-      if: inputs.deps != '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
+      if: inputs.deps != '' && steps.restore-venv.outputs.cache-hit != 'true'
       shell: bash
       run: |
         python -m venv .venv
@@ -104,7 +92,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-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
+      if: inputs.deps == '' && steps.restore-venv.outputs.cache-hit != 'true'
       shell: bash
       run: |
         python -m venv .venv
@@ -149,7 +137,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/7.1 $(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
 
@@ -194,18 +182,12 @@ 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') && github.event_name != 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
       uses: actions/cache@v4
       with:
         path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
 
     - name: Run apt Update + Install
       if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
@@ -237,7 +219,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/tinygrad/amdcomgr_dylib/releases/latest | \
+        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
           jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
           sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
         cargo build --release --manifest-path ./extra/remu/Cargo.toml
@@ -257,17 +239,8 @@ 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' && github.event_name != 'pull_request'
+      if: inputs.ocelot == 'true'
       id: cache-build
       uses: actions/cache@v4
       env:
@@ -276,7 +249,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-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true'
+      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
       shell: bash
       run: |
         git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
@@ -287,7 +260,6 @@ runs:
 
         CMAKE_ARGS="-Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5"
         if [[ "${{ runner.os }}" == "macOS" ]]; then
-          sudo xcode-select -s /Applications/Xcode_16.2.app/Contents/Developer
           CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
         fi
 
@@ -331,9 +303,3 @@ 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,7 +32,6 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
-        key: 'autogen'
         opencl: 'true'
         amd: 'true'
         cuda: 'true'
@@ -41,14 +40,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 libdrm-dev
+      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
     - name: Regenerate autogen files
       run: |
-        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
+        find tinygrad/runtime/autogen -type f -name "*.py" -not -name "__init__.py" -not -name "comgr_3.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
         python3 -c "from tinygrad.runtime.autogen import opencl"
         python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
-        python3 -c "from tinygrad.runtime.autogen import comgr_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 comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd"
+        python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v14_0_2"
         python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
         python3 -c "from tinygrad.runtime.autogen import llvm"
         python3 -c "from tinygrad.runtime.autogen import webgpu"
@@ -56,14 +55,12 @@ jobs:
         python3 -c "from tinygrad.runtime.autogen import libusb"
         python3 -c "from tinygrad.runtime.autogen import mesa"
         python3 -c "from tinygrad.runtime.autogen import avcodec"
-        python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
         REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
     - name: Check for differences
       run: |
         if ! git diff --quiet; then
-          git diff
           git diff > autogen-ubuntu.patch
-          echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
           exit 1
         fi
     - name: Upload patch artifact
@@ -83,18 +80,16 @@ 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 tinygrad/runtime/autogen/iokit.py tinygrad/runtime/autogen/corefoundation.py
-        python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
+        LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
     - name: Check for differences
       run: |
         if ! git diff --quiet; then
-          git diff
           git diff > autogen-macos.patch
-          echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
           exit 1
         fi
     - name: Upload patch artifact
@@ -104,8 +99,8 @@ jobs:
         name: autogen-macos-patch
         path: autogen-macos.patch
 
-  autogen-comgr-2:
-    name: In-tree Autogen (comgr 2)
+  autogen-comgr-3:
+    name: In-tree Autogen (comgr 3)
     runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
@@ -113,32 +108,29 @@ 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.2 $(lsb_release -cs) main
+        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.4 $(lsb_release -cs) main
         EOF
         echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
         sudo apt -qq update || true
         sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
     - name: Regenerate autogen files
       run: |
-        rm tinygrad/runtime/autogen/comgr.py
-        python3 -c "from tinygrad.runtime.autogen import comgr"
+        rm tinygrad/runtime/autogen/comgr_3.py
+        python3 -c "from tinygrad.runtime.autogen import comgr_3"
     - name: Check for differences
       run: |
         if ! git diff --quiet; then
-          git diff
-          git diff > autogen-comgr2.patch
-          echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          git diff > autogen-comgr3.patch
+          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
           exit 1
         fi
     - name: Upload patch artifact
       if: failure()
       uses: actions/upload-artifact@v4
       with:
-        name: autogen-comgr2-patch
-        path: autogen-comgr2.patch
+        name: autogen-comgr3-patch
+        path: autogen-comgr3.patch

.github/workflows/benchmark.yml

@@ -16,41 +16,6 @@ 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=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:
@@ -180,18 +145,14 @@ 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=. GMMU=0 DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
     - name: UsbGPU tiny tests
-      run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
     - name: UsbGPU copy speeds
-      run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
     #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
     - name: UsbGPU (USB4/TB) boot time
       run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
     - name: UsbGPU (USB4/TB) tiny tests
@@ -330,13 +291,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 ASSERT_FPS=1400 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
+      run: VALIDATE=1 MAX_FRAMES=100 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
     - name: Train MNIST
       run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
     - 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=120 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
     - name: Run 10 CIFAR training steps w BF16
       run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
     # - name: Run 10 CIFAR training steps w winograd
@@ -371,9 +332,9 @@ jobs:
     - name: Setcap to python
       run: ./extra/amdpci/setup_python_cap.sh
     - name: Remove amd modules
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+      run: ./extra/hcq/hcq_smi.py amd rmmod
     - name: Kill stale pids
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+      run: ./extra/hcq/hcq_smi.py amd kill_pids
     #- name: Insert amdgpu
     #  run: sudo modprobe amdgpu
     - name: Symlink models and datasets
@@ -483,9 +444,9 @@ jobs:
     - name: Setcap to python
       run: ./extra/amdpci/setup_python_cap.sh
     - name: Remove amd modules
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+      run: ./extra/hcq/hcq_smi.py amd rmmod
     - name: Kill stale pids
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+      run: ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -508,7 +469,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=230 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
     # - name: Run 10 CIFAR training steps w BF16
     #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
     # TODO: too slow
@@ -518,9 +479,6 @@ 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
 
@@ -538,9 +496,9 @@ jobs:
     - name: Setcap to python
       run: ./extra/amdpci/setup_python_cap.sh
     - name: Remove amd modules
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+      run: ./extra/hcq/hcq_smi.py amd rmmod
     - name: Kill stale pids
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+      run: ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -588,22 +546,22 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: test/external/process_replay/reset.py
-    - name: openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-    - name: IR3 openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM QCOM_IR3=1 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.11.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.11.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.10.0 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.10.0 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
+      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
       run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=image_1_openpilot_0_10_1_vision PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.10.1 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: benchmark MobileNetV2 on DSP
       run: |
         # generate quantized weights
@@ -615,27 +573,6 @@ jobs:
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 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="." GMMU=0 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="." GMMU=0 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]
@@ -650,9 +587,9 @@ jobs:
     - name: Setcap to python
       run: ./extra/amdpci/setup_python_cap.sh
     - name: Remove amd modules
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+      run: ./extra/hcq/hcq_smi.py amd rmmod
     - name: Kill stale pids
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+      run: ./extra/hcq/hcq_smi.py amd kill_pids
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -697,14 +634,6 @@ jobs:
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Remote
-      run: |
-        pkill -f 'extra/remote/serve.py' || true
-        PYTHONPATH=. python3 extra/remote/serve.py 6482 &
-        sleep 1
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD=1 AMD_IFACE=PCI python3 test/test_tiny.py
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD=1 AMD_AQL=1 AMD_IFACE=PCI python3 test/test_tiny.py
-        pkill -f 'extra/remote/serve.py' || true
     - 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
 
@@ -722,9 +651,9 @@ jobs:
     - name: Setcap to python
       run: ./extra/amdpci/setup_python_cap.sh
     - name: Remove nv modules
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv rmmod
+      run: ./extra/hcq/hcq_smi.py nv rmmod
     - name: Kill stale pids
-      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
+      run: ./extra/hcq/hcq_smi.py nv kill_pids
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -761,12 +690,5 @@ jobs:
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Remote
-      run: |
-        pkill -f 'extra/remote/serve.py' || true
-        PYTHONPATH=. python3 extra/remote/serve.py 6483 &
-        sleep 1
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6483 NV=1 python3 test/test_tiny.py
-        pkill -f 'extra/remote/serve.py' || true
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

.github/workflows/test.yml

@@ -1,11 +1,11 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '18'
+  CACHE_VERSION: '15'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
-  CHECK_OOB: 1
+  IGNORE_OOB: 0
 
 on:
   push:
@@ -26,19 +26,19 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: llvm-speed
-        deps: testing_unit
+        deps: testing_minimal
         llvm: 'true'
     - name: Speed Test
-      run: CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+      run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
     - name: Speed Test (BEAM=2)
-      run: BEAM=2 CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
 
   docs:
     name: Docs
     runs-on: ubuntu-22.04
     timeout-minutes: 10
     env:
-      CHECK_OOB: 0
+      IGNORE_OOB: 1
     steps:
     - name: Checkout Code
       uses: actions/checkout@v4
@@ -70,7 +70,7 @@ jobs:
         source venv/bin/activate
         pip install $GITHUB_WORKSPACE
         cp $GITHUB_WORKSPACE/examples/beautiful_mnist.py .
-        BS=2 STEPS=10 MAX_BUFFER_SIZE=0 python beautiful_mnist.py
+        BS=2 STEPS=10 python beautiful_mnist.py
     - name: Test Docs Build
       run: python -m mkdocs build --strict
     - name: Test Docs
@@ -98,7 +98,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_unit
+        deps: testing_minimal
         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_tiny.py TestTiny.test_plus
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
     - name: Test ResNet-18
       run: DEBUG=2 python3 extra/torch_backend/example.py
     - name: custom tests
@@ -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/backend/test_ops.py --durations=20
+      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
     - name: Test in-place operations on views
       run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
     - name: Test multi-gpu
@@ -134,14 +134,14 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_unit
+        deps: testing_minimal
         llvm: 'true'
     - name: Install ninja
       run: |
         sudo apt update || true
         sudo apt install -y --no-install-recommends ninja-build
     - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+      run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
     - name: Test some torch tests (expect failure)
       run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
 
@@ -156,27 +156,27 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: be-minimal
-        deps: testing_unit
+        deps: testing_minimal
     - name: Test dtype with Python emulator
-      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/backend/test_dtype.py test/backend/test_dtype_alu.py
+      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/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/backend/test_ops.py --durations=20
+      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20
     - name: Test uops with Python emulator
-      run: PYTHON=1 python3 -m pytest test/backend/test_uops.py --durations=20
+      run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
     - name: Test symbolic with Python emulator
-      run: PYTHON=1 python3 test/backend/test_symbolic_ops.py
+      run: PYTHON=1 python3 test/test_symbolic_ops.py
     - name: test_renderer_failures with Python emulator
-      run: PYTHON=1 python3 -m pytest -rA test/backend/test_renderer_failures.py::TestRendererFailures
-    - name: Test IMAGE support
+      run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
+    - name: Test IMAGE=2 support
       run: |
-        IMAGE=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
-        IMAGE=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
+        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
     - name: Test emulated METAL tensor cores
       run: |
-        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/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/backend/test_ops.py TestOps.test_gemm
+      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/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/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 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
         DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
     - 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/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
+        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
 
   linter:
     name: Linters
@@ -239,41 +239,9 @@ jobs:
     - 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: 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
+      run: TYPED=1 python -c "import tinygrad"
 
   unittest:
     name: Unit Tests
@@ -287,18 +255,29 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: unittest-13
-        pydeps: "pillow ftfy regex pre-commit"
+        pydeps: "pillow numpy 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/null/test_device.py TestRunAsModule.test_module_runs
-        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+        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
     - name: Run GC tests
       run: python test/external/external_uop_gc.py
     - name: External Benchmark Schedule
@@ -312,8 +291,8 @@ jobs:
         python extra/optimization/extract_dataset.py
         gzip -c /tmp/sops > extra/datasets/sops.gz
         #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 24000 lines
-      run: MAX_LINE_COUNT=24000 python sz.py
+    - name: Repo line count < 20000 lines
+      run: MAX_LINE_COUNT=20000 python sz.py
 
   spec:
     strategy:
@@ -333,7 +312,7 @@ jobs:
         deps: testing_unit
         python-version: '3.14'
     - name: Test SPEC=2
-      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 }}
+      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --ignore test/unit/test_autogen.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
 
   fuzzing:
     name: Fuzzing
@@ -367,13 +346,13 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: gpu-image
-          deps: testing_unit
+          deps: testing_minimal
           opencl: 'true'
-      - name: Test CL IMAGE=1 ops
+      - name: Test CL IMAGE=2 ops
         run: |
-          CL=1 IMAGE=1 python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
         # TODO: training is broken
-        # CL=1 IMAGE=1 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
+        # CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -388,14 +367,14 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: gen-dataset
-          deps: testing
+          deps: testing_minimal
           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/backend/test_optim.py -k "not muon"
+        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/test_optim.py -k "not muon"
       - name: Upload artifact
         uses: actions/upload-artifact@v4
         with:
@@ -418,13 +397,13 @@ jobs:
           llvm: 'true'
       - name: Test openpilot model kernel count and gate usage
         run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=17 FLOAT16=1 CL=1 IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
       - name: Test openpilot CL compile fp16
-        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
       - name: Test openpilot CL compile fp32 (test correctness)
-        run: CL=1 IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
+        run: DEBUGCL=1 CL=1 IMAGE=2 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
       - name: Test openpilot LLVM compile fp16
-        run: IMAGE=1 FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+        run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -443,7 +422,7 @@ jobs:
         with:
           key: onnxoptc
           deps: testing
-          python-version: '3.12'
+          python-version: '3.11'
           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
@@ -454,7 +433,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/backend/test_quantize_onnx.py
+        run: CPU=1 CPU_LLVM=0 python3 test/test_quantize_onnx.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -471,7 +450,7 @@ jobs:
           key: onnxoptl
           deps: testing
           pydeps: "tensorflow==2.19"
-          python-version: '3.12'
+          python-version: '3.11'
           opencl: 'true'
       - name: Test ONNX (CL)
         run: CL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -484,11 +463,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 NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py
+        run: NULL=1 python examples/beautiful_mnist_multigpu.py
       - name: Test Bert training
-        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
+        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
       - name: Test llama 3 training
-        run: NULL=1 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
+        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -497,7 +476,7 @@ jobs:
     runs-on: ubuntu-24.04
     timeout-minutes: 15
     env:
-      CHECK_OOB: 0
+      IGNORE_OOB: 1
     steps:
       - name: Checkout Code
         uses: actions/checkout@v4
@@ -505,13 +484,8 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: apps_llm
-      - name: Test 1B LLM (llama)
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster
-      - name: Test 1B LLM (llama q4)
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster
-      - name: Test 1B LLM (qwen)
-        # NOTE: qwen is dumb and only knows about female chickens
-        run: echo "What's a female chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen
+      - name: Test 1B LLM
+        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm | grep -i rooster
 
 # ****** Models Tests ******
 
@@ -550,7 +524,7 @@ jobs:
       with:
         key: metal
         deps: testing
-        python-version: '3.12'
+        python-version: '3.11'
     - name: Test models (Metal)
       run: METAL=1 python -m pytest -n=auto test/models --durations=20
     - name: Test LLaMA compile speed
@@ -569,15 +543,15 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: devectorize-minimal
-        deps: testing_unit
+        deps: testing_minimal
         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/backend/test_ops.py
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/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/backend/test_ops.py
+      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
 
   testdsp:
     name: Linux (DSP)
@@ -590,8 +564,8 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: dsp-minimal
-        deps: testing_unit
-        pydeps: "onnx==1.18.0 onnxruntime ml_dtypes"
+        deps: testing_minimal
+        pydeps: "onnx==1.18.0 onnxruntime pillow"
         llvm: "true"
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v3
@@ -603,15 +577,15 @@ jobs:
         load: true
         tags: qemu-hexagon:latest
         cache-from: type=gha
-        cache-to: ${{ github.event_name != 'pull_request' && 'type=gha,mode=min' || '' }}
+        cache-to: 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/backend/test_transcendental.py TestTranscendentalVectorized
+      run: CC=clang-20 DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
     - name: Test quantize onnx
-      run: DEBUG=2 DSP=1 python3 test/backend/test_quantize_onnx.py
+      run: DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
 
   testwebgpu:
     name: Linux (WebGPU)
@@ -624,105 +598,32 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: webgpu-minimal
-        deps: testing_unit
-        python-version: '3.12'
+        deps: testing_minimal
+        python-version: '3.11'
         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_tiny.py TestTiny.test_plus
+        WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
     - name: Run selected webgpu tests
       run: |
-          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/backend --durations=20
+          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
-  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: GMMU=0 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, cdna4]
 
-    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
+    name: Linux (${{ matrix.backend }})
     runs-on: ubuntu-22.04
-    timeout-minutes: 15
+    timeout-minutes: 20
     env:
       AMD: 1
       MOCKGPU: 1
-      MOCKGPU_ARCH: ${{ matrix.arch }}
-      SKIP_SLOW_TEST: 1
+      FORWARD_ONLY: 1
       AMD_LLVM: ${{ matrix.backend == 'amdllvm' && '1' || matrix.backend != 'amdllvm' && '0' }}
     steps:
       - name: Checkout Code
@@ -731,20 +632,70 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: ${{ matrix.backend }}-minimal
-          deps: testing_unit
+          deps: testing_minimal
           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_tiny.py TestTiny.test_plus
+          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
       - name: Run pytest (amd)
-        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/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
+        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py --durations=20
+      - name: Run pytest (amd)
+        run: python -m pytest test/external/external_test_am.py --durations=20
       - name: Run TRANSCENDENTAL math
-        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
+        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
+      - name: Run TestOps.test_add with SQTT
+        run: |
+          VIZ=1 PMC=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
+          VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
+          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
+  testamdasm:
+    name: AMD ASM IDE
+    runs-on: ubuntu-24.04
+    timeout-minutes: 10
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: rdna3-emu
+          deps: testing_minimal
+          amd: 'true'
+          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
@@ -764,7 +715,7 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: ${{ matrix.backend }}-minimal
-          deps: testing_unit
+          deps: testing_minimal
           cuda: 'true'
           ocelot: 'true'
       - name: Set env
@@ -772,12 +723,10 @@ 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_tiny.py TestTiny.test_plus
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
       - name: Run pytest (cuda)
         # skip multitensor because it's slow
-        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
+        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
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -797,7 +746,7 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: ${{ matrix.backend }}-minimal
-          deps: testing_unit
+          deps: testing_minimal
           opencl: ${{ matrix.backend == 'opencl' && 'true' }}
           llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
           mesa: ${{ matrix.backend == 'lvp' && 'true' }}
@@ -806,11 +755,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_tiny.py TestTiny.test_plus
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
       - name: Run pytest (${{ matrix.backend }})
-        run: python -m pytest -n=auto test/backend --durations=20
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
       - name: Run TRANSCENDENTAL math
-        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
+        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 process replay tests
         uses: ./.github/actions/process-replay
 
@@ -828,29 +777,27 @@ jobs:
       with:
         key: metal
         deps: testing
-        python-version: '3.12'
+        python-version: '3.11'
         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/backend/test_ops.py TestOps.test_gemm
+      run: METAL=1 DEBUG=3 TC=2 python test/test_ops.py TestOps.test_gemm
     - name: Test tensor core ops (real)
-      run: METAL=1 DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
+      run: METAL=1 DEBUG=3 python test/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/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
+      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
     - name: Run pytest (amd)
       env:
         MOCKGPU: 1
@@ -873,8 +820,6 @@ 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
@@ -893,14 +838,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: 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
+    - 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
@@ -932,7 +877,8 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: macos-${{ matrix.backend }}-minimal
-          deps: testing_unit
+          deps: testing_minimal
+          pydeps: "capstone"
           llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
           mesa: ${{ matrix.backend == 'lvp' && 'true' }}
       - name: Set env
@@ -942,7 +888,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/backend --durations=20
+        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
       - name: Run macOS-specific unit test
@@ -975,16 +921,12 @@ 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_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
+        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
       - 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/backend/test_ops.py --durations=20
+          python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
 
 # ****** Compile-only Tests ******
 
@@ -1003,38 +945,15 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: compile-${{ matrix.backend }}
-          deps: testing_unit
+          deps: testing_minimal
           mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
-          python-version: '3.12'
+          python-version: '3.14'
       - name: Set env
         shell: bash
-        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
+        run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
       - name: Run test_ops
         shell: bash
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/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
+          DEBUG=4 python3 test/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/test_ops.py --durations=20

.gitignore

@@ -66,5 +66,3 @@ target
 .mypy_cache
 mutants
 .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/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
+        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/unit/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
         language: system
         always_run: true
         pass_filenames: false

AGENTS.md

@@ -0,0 +1,17 @@
+# 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

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

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/backend/test_ops.py        # just the ops tests
+python3 test/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_kernel_graph
+::: tinygrad.schedule.rangeify.get_rangeify_map
     options:
         members: false
         show_labels: false

examples/anthropic_challenge.py

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

@@ -1,6 +1,6 @@
 # model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
 from typing import Callable
-from tinygrad import Tensor, TinyJit, nn, GlobalCounters, function
+from tinygrad import Tensor, TinyJit, nn, GlobalCounters
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist
 
@@ -15,31 +15,30 @@ class Model:
       nn.BatchNorm(64), Tensor.max_pool2d,
       lambda x: x.flatten(1), nn.Linear(576, 10)]
 
-  @function
   def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
 
-  @TinyJit
-  @Tensor.train()
-  def train_step(self, X_train:Tensor, Y_train:Tensor) -> Tensor:
-    opt.zero_grad()
-    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
-    loss = self(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
-    return loss.realize(*opt.schedule_step())
-
-  @TinyJit
-  def get_test_acc(self, X_test:Tensor, Y_test:Tensor) -> Tensor: return (self(X_test).argmax(axis=1) == Y_test).mean()*100
-
 if __name__ == "__main__":
   X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))
 
   model = Model()
   opt = (nn.optim.Muon if getenv("MUON") else nn.optim.SGD if getenv("SGD") else nn.optim.Adam)(nn.state.get_parameters(model))
 
+  @TinyJit
+  @Tensor.train()
+  def train_step() -> Tensor:
+    opt.zero_grad()
+    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
+    loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
+    return loss.realize(*opt.schedule_step())
+
+  @TinyJit
+  def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
+
   test_acc = float('nan')
   for i in (t:=trange(getenv("STEPS", 70))):
     GlobalCounters.reset()   # NOTE: this makes it nice for DEBUG=2 timing
-    loss = model.train_step(X_train, Y_train)
-    if i%10 == 9: test_acc = model.get_test_acc(X_test, Y_test).item()
+    loss = train_step()
+    if i%10 == 9: test_acc = get_test_acc().item()
     t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
 
   # verify eval acc

examples/benchmark_onnx.py

@@ -5,7 +5,7 @@ from extra.onnx_helpers import get_example_inputs, validate
 
 def load_onnx_model(onnx_file):
   run_onnx = OnnxRunner(onnx_file)
-  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True)
+  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True)
   return run_onnx_jit, run_onnx.graph_inputs
 
 if __name__ == "__main__":

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)=}"
-assert EVAL_BS % len(GPUS) == 0, f"{EVAL_BS=} is not a multiple of {len(GPUS)=}"
+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"
 
 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,7 +65,17 @@ 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))
-      X[idx].flatten().assign(img.tobytes())
+
+      # 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!
       q_out.put(idx)
     q_out.put(None)
 
@@ -254,8 +264,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].flatten().assign(x.tobytes())
-    Y[idx].flatten().assign(y.tobytes())
+    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()
 
     queue_out.put(idx)
   queue_out.put(None)
@@ -369,12 +379,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].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())
+      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()
 
-    imgs[idx].flatten().assign(img.tobytes())
+    imgs[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
 
     queue_out.put(idx)
   queue_out.put(None)
@@ -396,7 +406,6 @@ 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}")
@@ -543,7 +552,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: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.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.count, = struct.unpack("<Q", self.idx.read(8))
     doc_count, = struct.unpack("<Q", self.idx.read(8))
 
@@ -560,7 +569,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")).numpy()
+    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
 
   def _index(self, idx) -> tuple[int, int]:
     return int(self.pointers[idx]), int(self.sizes[idx])
@@ -569,7 +578,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].view(self.dtype)
+    return self.bin_t[ptr:ptr+length*self.dtype.itemsize].bitcast(self.dtype).to(None)
 
 # https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
 class GPTDataset:
@@ -628,7 +637,7 @@ class GPTDataset:
         sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
 
     # concat all parts
-    text = np.concatenate(sample_parts, axis=0)
+    text = Tensor.cat(*sample_parts)
 
     return text
 
@@ -771,8 +780,7 @@ 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)]
-    stacked = np.stack(batch, axis=0)
-    yield Tensor(stacked, device="NPY")
+    yield Tensor.stack(batch, dim=0)
 
 def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
   return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)

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, profile_marker, DEBUG
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
 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
 
@@ -13,8 +13,6 @@ from extra.bench_log import BenchEvent, WallTimeEvent
 # TODO: fix benchmark logging and use tinygrad tqdm
 from tqdm import tqdm
 
-from tinygrad.uop.ops import UOp
-
 def train_resnet():
   from extra.models import resnet
   from examples.mlperf.dataloader import batch_load_resnet
@@ -1284,10 +1282,9 @@ def train_bert():
         previous_step = i
 
 def train_llama3():
-  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad
+  from extra.models.llama import Transformer
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW
 
   BENCHMARK = getenv("BENCHMARK")
 
@@ -1295,18 +1292,13 @@ def train_llama3():
   BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
   BS                 = config["BS"]                     = getenv("BS", 16)
   grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
+  assert grad_acc == 1, f"{grad_acc=} is not supported"
   GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
   SEED               = config["SEED"]                   = getenv("SEED", 5760)
-  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)
@@ -1320,12 +1312,10 @@ def train_llama3():
   opt_adamw_weight_decay = 0.1
 
   opt_gradient_clip_norm = 1.0
-  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
+  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)
 
   # ** init wandb **
   WANDB = getenv("WANDB")
@@ -1337,16 +1327,10 @@ 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
   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 = FlatTransformer(**model_params, max_context=SEQLEN)
-
+  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
   params = get_parameters(model)
   # weights are all bfloat16 for now
   assert params and all(p.dtype == dtypes.bfloat16 for p in params)
@@ -1355,26 +1339,32 @@ def train_llama3():
     for v in get_parameters(model):
       v = v.assign(Tensor.empty(v.shape))
 
-  is_dp = (DP := getenv("DP", 1)) > 1
-  is_mp = (MP := getenv("MP", 1)) > 1
-  is_sharding = is_dp or is_mp
-  device_count = max(DP, MP)
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
+  if (DP := getenv("DP", 1)) > 1:
+    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+    for v in get_parameters(model):
+      v.shard_(device, axis=None)
 
-  model.shard(device, is_mp)
-
-  if is_dp: vocab_mask.shard_(device, axis=None).realize()
-  if is_mp: 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(params, 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
-  grads = [Tensor.zeros_like(p).contiguous() for p in optim.params]
+  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():
+      if 'scale' in k: v.shard_(device, axis=None)  # from quantized
+      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.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)
+      elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
+      elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
+      elif 'output.weight' in k: v.shard_(device, axis=0)
+      else:
+        # attention_norm, ffn_norm, norm
+        v.shard_(device, axis=None)
+      # 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)
   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"):
@@ -1387,131 +1377,101 @@ def train_llama3():
     load_state_dict(scheduler, safe_load(fn), realize=False)
 
   @TinyJit
-  def minibatch(tokens:Tensor):
-    if is_dp: tokens = tokens.to(None).shard(device, 0)
-    if is_mp: tokens = tokens.shard(device)
-    if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1])
-    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+  @Tensor.train()
+  def train_step(model, tokens:Tensor):
+    optim.zero_grad()
+    if (DP := getenv("DP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+      tokens = tokens.shard(device, 0)
+    if (MP := getenv("MP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+      tokens = tokens.shard(device)
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    loss.backward()
+    # 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)
 
-    for i,(t,g) in enumerate(zip(optim.params, loss.gradient(*optim.params))):
-      grads[i].replace(Tensor(grads[i].uop.after(UOp.group(*apply_grad(grads[i].uop, g.uop))), device=t.device))
-
-    loss_cpu = loss.flatten().float().to("CPU")
-    return loss_cpu.realize(*grads)
-
-  @TinyJit
-  def optim_step():
-    grad_norm = optim.fstep(grads)
+    optim.step()
     scheduler.step()
 
-    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
+    lr = optim.lr
+    loss.realize(lr)
+    return loss, lr
 
   @TinyJit
   @Tensor.train(False)
-  def eval_step(tokens:Tensor):
-    if is_dp: tokens = tokens.to(None).shard(device, 0)
-    if is_mp: tokens = tokens.shard(device)
-    if not is_sharding: 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")
+  def eval_step(model, tokens:Tensor):
+    if (DP := getenv("DP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+      tokens = tokens.shard(device, 0)
+    if (MP := getenv("MP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+      tokens = tokens.shard(device)
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    return loss.flatten().float()
 
   # ** data iters **
   def fake_data(bs, samples):
-    import numpy as np
     for _ in range(samples // bs):
-      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")
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
 
   def get_train_iter():
     if getenv("FAKEDATA", 0):
       return fake_data(BS, SAMPLES)
     else:
       from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=DATA_SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
+      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
 
   if getenv("FAKEDATA", 0):
     eval_dataset = None
   else:
     from examples.mlperf.dataloader import get_llama3_dataset
-    eval_dataset = get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
+    eval_dataset = get_llama3_dataset(1024 if SMALL else 5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
 
   def get_eval_iter():
     if eval_dataset is None:
-      return fake_data(EVAL_BS, EVAL_SAMPLES)
+      return fake_data(EVAL_BS, 5760)
     from examples.mlperf.dataloader import iterate_llama3_dataset
     return iterate_llama3_dataset(eval_dataset, EVAL_BS)
 
-  num_params = sum(p.numel() for p in params) - model_params["vocab_size"]*model_params["dim"]
-  train_iter = get_train_iter()
+  iter = get_train_iter()
   i, sequences_seen = resume_ckpt, 0
   step_times = []
-  while i < MAX_STEPS:
+  for tokens in tqdm(iter, total=SAMPLES//GBS):
     GlobalCounters.reset()
-    actual_gbs = GBS if i >= 2 else BS
     if getenv("TRAIN", 1):
-      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)
+      t = time.perf_counter()
+      loss, lr = train_step(model, tokens)
+      loss = loss.float().item()
+      lr = lr.item()
 
       i += 1
-      sequences_seen += actual_gbs
+      sequences_seen += tokens.shape[0]
+
+      sec = time.perf_counter()-t
+      if BENCHMARK: step_times.append(sec)
 
       mem_gb = GlobalCounters.mem_used / 1e9
-      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
+      gflops = GlobalCounters.global_ops / 1e9 / sec
       tqdm.write(
-          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())))
+        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")
 
       if WANDB:
-        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
-        })
+        wandb.log({"lr": lr, "train/loss": loss, "train/step_time": sec, "train/GFLOPS": gflops, "train/sequences_seen": sequences_seen})
 
       if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
         tqdm.write("saving checkpoint")
@@ -1530,20 +1490,18 @@ def train_llama3():
         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
+    if (sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
       tqdm.write(f"evaluating after {sequences_seen} sequences")
-      profile_marker(f"eval @ {i}")
 
       # run eval
       eval_losses = []
       eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences")
+      tqdm.write(f"evaluating {5760//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):
+      for j,tokens in tqdm(enumerate(eval_iter), total=5760//EVAL_BS):
+        eval_losses += eval_step(model, tokens).tolist()
+        
+        if BENCHMARK and (j+1) == min(BENCHMARK, 5760//EVAL_BS):
           return
 
       log_perplexity = Tensor(eval_losses).mean().float().item()
@@ -1631,7 +1589,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():
@@ -1670,7 +1628,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/flat_llama.py

@@ -1,193 +0,0 @@
-import math, os
-if __name__ == "__main__":
-  os.environ["DEFAULT_FLOAT"] = "bfloat16"
-  os.environ["OPTIM_DTYPE"] = "bfloat16"
-  if "DEV" not in os.environ: os.environ["DEV"] = "NULL"
-  # CDNA
-  os.environ["EMULATE"] = "AMD_CDNA4"
-  os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
-  os.environ["ALL2ALL"] = "1"
-  os.environ["USE_ATOMICS"] = "1"
-  if "HK_FLASH_ATTENTION" not in os.environ:
-    os.environ["HK_FLASH_ATTENTION"] = "1"
-    if "ASM_GEMM" not in os.environ:
-      os.environ["ASM_GEMM"] = "1"
-from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
-from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
-from tinygrad.uop.ops import Ops, UOp
-from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
-
-FP8 = getenv("FP8", 0)
-
-FP8_DTYPE = dtypes.fp8e4m3
-FP8_MAX = 448.0
-
-def quantize_fp8(x:Tensor):
-  scale = FP8_MAX / (x.abs().max().detach() + 1e-8)
-  x_scaled = x * scale
-  x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
-  return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal()
-
-def matmul(x:Tensor, w:Tensor) -> Tensor:
-  if not FP8: return x @ w.T
-  # weights are already FP8, just quantize activations
-  x_fp8, x_scale = quantize_fp8(x)
-  return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale
-
-def rmsnorm(x_in:Tensor, eps:float):
-  x = x_in.float()
-  x = x * (x.square().mean(-1, keepdim=True) + eps).rsqrt()
-  return x.cast(x_in.dtype)
-
-class FlatTransformer:
-  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):
-    self.vocab_size = vocab_size
-    self.n_layers = n_layers
-    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
-
-    # Attention
-    self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
-    self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim)
-
-    # FeedForward
-    self.w1 = self.lin_per_layer(dim, hidden_dim)
-    self.w2 = self.lin_per_layer(hidden_dim, dim)
-    self.w3 = self.lin_per_layer(dim, hidden_dim)
-
-    self.norm_eps = norm_eps
-    self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
-    self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
-
-    # output
-    self.norm = nn.RMSNorm(dim, norm_eps)
-    self.tok_embeddings = nn.Embedding(vocab_size, dim)
-    self.output = nn.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 lin_per_layer(self, in_features:int, out_features:int):
-    bound = 1 / math.sqrt(in_features)
-    dt = FP8_DTYPE if FP8 else None
-    if getenv("ZEROS"): return Tensor.zeros(self.n_layers, out_features, in_features, dtype=dt)
-    return Tensor.uniform(self.n_layers, out_features, in_features, low=-bound, high=bound, dtype=dt)
-
-  def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor):
-    x = rmsnorm(x, self.norm_eps) * attention_norm
-    xqkv = matmul(x, wqkv)
-
-    bsz, seqlen, _ = xqkv.shape
-    # interleaved layout: each kv group has [n_rep q heads, 1 k head, 1 v head] for clean MP sharding
-    xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
-    xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
-    xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-    xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    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 matmul(attn, wo)
-
-  def feed_forward(self, x:Tensor, ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
-    x = rmsnorm(x, self.norm_eps) * ffn_norm
-    x_w1 = matmul(x, w1).silu()
-    x_w3 = matmul(x.contiguous_backward(), w3)
-    return matmul(x_w1 * x_w3, w2)
-
-  @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor,
-                attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
-    h = x + self.attention(x, freqs_cis, attention_norm, wqkv, wo)
-    return h + self.feed_forward(h, ffn_norm, w1, w2, w3)
-
-  def shard(self, device:tuple[str, ...], mp:bool=False):
-    from tinygrad.nn.state import get_parameters
-    if not mp:
-      for v in get_parameters(self): v.shard_(device, axis=None)
-    else:
-      # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
-      self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
-      self.wo.shard_(device, axis=2).realize()             # (n_layers, dim, in) shard in
-      self.w1.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
-      self.w2.shard_(device, axis=2).realize()             # (n_layers, dim, hidden) shard in
-      self.w3.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
-      self.attention_norm.shard_(device, axis=None).realize()
-      self.ffn_norm.shard_(device, axis=None).realize()
-      self.norm.weight.shard_(device, axis=None).realize()
-      self.tok_embeddings.weight.shard_(device, axis=0).realize()
-      self.output.weight.shard_(device, axis=0).realize()
-      self.freqs_cis.shard_(device, axis=None).realize()
-
-  def __call__(self, tokens:Tensor):
-    h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
-    for i in range(self.n_layers):
-      h = self.run_layer(h, freqs_cis,
-                         self.attention_norm[i], self.wqkv[i], self.wo[i],
-                         self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i])
-    logits = self.output(self.norm(h))
-    return logits
-
-# TODO: this shouldn't be needed, but it prevents a copy of the grads. CAT can help
-def apply_grad(old_grad:UOp, new_grad:UOp) -> list[UOp]:
-  if new_grad.op == Ops.ADD:
-    return apply_grad(old_grad, new_grad.src[0])+apply_grad(old_grad, new_grad.src[1])
-  elif new_grad.op == Ops.PAD:
-    grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(new_grad.src[0].shape, new_grad.marg)])
-    return apply_grad(old_grad.shrink(grad_shrink), new_grad.src[0])
-  else:
-    return [old_grad.store(old_grad + new_grad)]
-
-if __name__ == "__main__":
-  config = {}
-  BS                 = config["BS"]                     = getenv("BS", 16)
-  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
-
-  from examples.llama3 import MODEL_PARAMS
-  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
-  model = FlatTransformer(**model_params, max_context=SEQLEN)
-  state = nn.state.get_state_dict(model)
-  print("tensor count:", len(state))
-
-  # shard the model
-  from tinygrad import Device
-  if (DP := getenv("DP", 1)) > 1:
-    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)))
-  if (MP := getenv("MP", 1)) > 1:
-    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)), mp=True)
-
-  # preallocate all the grad buffers and zero them out
-  grads = {x:Tensor.zeros_like(x).contiguous() for x in state.values() if x.requires_grad is None}
-
-  # print model size
-  sz = 0
-  for k,v in state.items():
-    print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device}  {v.nbytes()/1e9:.2f} GB")
-    sz += v.nbytes()
-  print(f"total sz: {sz/1e9:.2f} GB")
-
-  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=model.vocab_size, dtype=dtypes.int)
-  with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
-  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
-  if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
-  if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))
-
-  @TinyJit
-  def jit_step(tokens:Tensor):
-    with Timing("python forward: "): loss = model(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
-    with Timing("python backward: "):
-      for t,g in zip(grads, loss.gradient(*grads)):
-        grads[t] = Tensor(grads[t].uop.after(UOp.group(*apply_grad(grads[t].uop, g.uop))), device=t.device)
-    with Timing("run step: "): loss.realize(*grads.values())
-
-  for i in range(6):
-    GlobalCounters.reset()
-    profile_marker(f"step {i}")
-    with Timing(colored(f"*** step {i}: ", "red")):
-      jit_step(tokens)
-  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

examples/mlperf/models/llama.py

@@ -1,80 +0,0 @@
-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/models/test_flat_llama.py

@@ -1,140 +0,0 @@
-import os
-os.environ["WQKV"] = "1"
-import unittest
-import numpy as np
-from tinygrad import Tensor, nn, dtypes
-from tinygrad.nn.state import get_parameters
-from tinygrad.device import is_dtype_supported
-from examples.mlperf.models.llama import Transformer
-from examples.mlperf.models.flat_llama import FlatTransformer
-
-def copy_weights(flat:FlatTransformer, ref:Transformer):
-  n_layers = flat.n_layers
-  Tensor.realize(*nn.state.get_state_dict(ref).values())
-  flat.wqkv.assign(Tensor(np.stack([ref.layers[i].attention.wqkv.weight.numpy() for i in range(n_layers)])))
-  flat.wo.assign(Tensor(np.stack([ref.layers[i].attention.wo.weight.numpy() for i in range(n_layers)])))
-  flat.w1.assign(Tensor(np.stack([ref.layers[i].feed_forward.w1.weight.numpy() for i in range(n_layers)])))
-  flat.w2.assign(Tensor(np.stack([ref.layers[i].feed_forward.w2.weight.numpy() for i in range(n_layers)])))
-  flat.w3.assign(Tensor(np.stack([ref.layers[i].feed_forward.w3.weight.numpy() for i in range(n_layers)])))
-  flat.attention_norm.assign(Tensor(np.stack([ref.layers[i].attention_norm.weight.numpy() for i in range(n_layers)])))
-  flat.ffn_norm.assign(Tensor(np.stack([ref.layers[i].ffn_norm.weight.numpy() for i in range(n_layers)])))
-  flat.norm.weight.assign(Tensor(ref.norm.weight.numpy()))
-  flat.tok_embeddings.weight.assign(Tensor(ref.tok_embeddings.weight.numpy()))
-  flat.output.weight.assign(Tensor(ref.output.weight.numpy()))
-
-class TestFlatLlama(unittest.TestCase):
-  def test_forward_match(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    ref = Transformer(**params)
-    flat = FlatTransformer(**params)
-    copy_weights(flat, ref)
-    Tensor.realize(*nn.state.get_state_dict(flat).values())
-
-    tokens = Tensor([[1, 50, 100, 999, 2]])
-    ref_logits = ref(tokens).realize()
-    flat_logits = flat(tokens).realize()
-    self.assertEqual(ref_logits.shape, flat_logits.shape)
-    diff = (ref_logits - flat_logits).abs().max().item()
-    self.assertLess(diff, 1e-5, f"forward mismatch: max abs diff {diff}")
-
-  def test_backward_match(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    ref = Transformer(**params)
-    flat = FlatTransformer(**params)
-    copy_weights(flat, ref)
-
-    for p in get_parameters(ref): p.requires_grad_(True)
-    for p in get_parameters(flat): p.requires_grad_(True)
-    Tensor.realize(*nn.state.get_state_dict(flat).values())
-
-    tokens = Tensor([[1, 50, 100, 999, 2, 10]])
-
-    ref_loss = ref(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
-    ref_loss.backward()
-    ref_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(ref).items() if v.grad is not None}
-
-    flat_loss = flat(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
-    flat_loss.backward()
-    flat_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(flat).items() if v.grad is not None}
-
-    # check loss matches
-    self.assertAlmostEqual(ref_loss.item(), flat_loss.item(), places=4)
-
-    # check output weight grad matches
-    diff = abs(ref_grads["output.weight"] - flat_grads["output.weight"]).max()
-    self.assertLess(diff, 1e-4, f"output.weight grad mismatch: max abs diff {diff}")
-
-    # check per-layer weight grads match
-    for i in range(params["n_layers"]):
-      for flat_key, ref_key in [
-        ("wqkv", f"layers.{i}.attention.wqkv.weight"),
-        ("wo", f"layers.{i}.attention.wo.weight"),
-        ("w1", f"layers.{i}.feed_forward.w1.weight"),
-        ("w2", f"layers.{i}.feed_forward.w2.weight"),
-        ("w3", f"layers.{i}.feed_forward.w3.weight"),
-      ]:
-        diff = abs(ref_grads[ref_key] - flat_grads[flat_key][i]).max()
-        self.assertLess(diff, 1e-4, f"layer {i} {flat_key} grad mismatch: max abs diff {diff}")
-
-  @unittest.skipUnless(os.getenv("CPU", "") == "1", "multi-device CPU test")
-  def test_forward_match_mp(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    from tinygrad import Device
-    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
-    ref = Transformer(**params)
-    flat = FlatTransformer(**params)
-    copy_weights(flat, ref)
-    Tensor.realize(*nn.state.get_state_dict(flat).values())
-    flat.shard(devices, mp=True)
-
-    tokens = Tensor([[1, 50, 100, 999, 2]], device=devices[0])
-    ref_logits = ref(tokens.to(devices[0])).numpy()
-    flat_logits = flat(tokens.shard(devices)).numpy()
-    self.assertEqual(ref_logits.shape, flat_logits.shape)
-    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
-
-  @unittest.skipUnless(os.getenv("CPU", "") == "1", "multi-device CPU test")
-  def test_forward_match_dp(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    from tinygrad import Device
-    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
-    ref = Transformer(**params)
-    flat = FlatTransformer(**params)
-    copy_weights(flat, ref)
-    Tensor.realize(*nn.state.get_state_dict(flat).values())
-    flat.shard(devices)
-
-    tokens = Tensor([[1, 50, 100, 999, 2], [2, 100, 50, 1, 999]], device=devices[0])
-    ref_logits = ref(tokens.to(devices[0])).numpy()
-    flat_logits = flat(tokens.shard(devices, axis=0)).numpy()
-    self.assertEqual(ref_logits.shape, flat_logits.shape)
-    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), "fp8 not supported on this device")
-  def test_forward_fp8(self):
-    import examples.mlperf.models.flat_llama as flat_llama_mod
-    old_fp8 = flat_llama_mod.FP8
-    try:
-      flat_llama_mod.FP8 = 1
-      Tensor.manual_seed(42)
-      params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-      ref = Transformer(**params)
-      flat = FlatTransformer(**params)
-      copy_weights(flat, ref)
-      Tensor.realize(*nn.state.get_state_dict(flat).values())
-
-      tokens = Tensor([[1, 50, 100, 999, 2]])
-      ref_logits = ref(tokens).numpy()
-      flat_logits = flat(tokens).numpy()
-      self.assertEqual(ref_logits.shape, flat_logits.shape)
-      # FP8 has lower precision, allow larger tolerance
-      np.testing.assert_allclose(flat_logits, ref_logits, atol=1.0, rtol=0.1)
-    finally:
-      flat_llama_mod.FP8 = old_fp8
-
-if __name__ == "__main__":
-  unittest.main()

examples/mlperf/optim.py

@@ -1,59 +0,0 @@
-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]]:
-    grads = list(grads)
-
-    for i in range(len(grads)):
-      if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
-
-    if self.fused:
-      grads[0].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 >= 3 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 CHECK_OOB=0
+export IGNORE_OOB=1
 
 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 CHECK_OOB=0
+export IGNORE_OOB=1
 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 CHECK_OOB=0
+export IGNORE_OOB=1
 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 CHECK_OOB=0
+export IGNORE_OOB=1
 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 CHECK_OOB=0
+export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=5000000
 
 export BEAM=0 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

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 CHECK_OOB=0
+export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=5000000
 
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

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 CHECK_OOB=0
+export IGNORE_OOB=1
 export REWRITE_STACK_LIMIT=5000000
 
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

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 CHECK_OOB=0
+export IGNORE_OOB=1
 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 CHECK_OOB=0
+export IGNORE_OOB=1
 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 CHECK_OOB=0
+export IGNORE_OOB=1
 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 CHECK_OOB=0
+export IGNORE_OOB=1
 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 CHECK_OOB=0
+export IGNORE_OOB=1
 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 CHECK_OOB=0
+export IGNORE_OOB=1
 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

@@ -1,38 +0,0 @@
-#!/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}
-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

@@ -1,32 +0,0 @@
-#!/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

@@ -1,42 +0,0 @@
-#!/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

@@ -1,43 +0,0 @@
-#!/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

@@ -1,37 +1,26 @@
 #!/usr/bin/env bash
 
-export PYTHONPATH="."
-export DEV=${DEV:-AMD}
-export EMULATE="AMD_CDNA4"
-export CHECK_OOB=0
+export PYTHONPATH="." AMD=1
+export IGNORE_OOB=1
 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 FLASH_ATTENTION=1
 
 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 DP=8 BS=8 EVAL_BS=8
 
 export MODEL="llama3"
 export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
-export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export 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 LR="1e-3" END_LR="1e-4" WARMUP_STEPS=1024 MAX_STEPS=1200000
+export SAMPLES=$((MAX_STEPS * BS))
 
-export SEED=${SEED:-$RANDOM}
-export DATA_SEED=${DATA_SEED:-5760}
+export 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 JITBEAM=3
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 
 python3 examples/mlperf/model_train.py

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

@@ -1,38 +0,0 @@
-#!/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

@@ -1,5 +0,0 @@
-#!/bin/bash
-export BENCHMARK=5
-export EVAL_BS=0
-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" --limit 20

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

@@ -1,10 +0,0 @@
-#!/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/mnist_gan.py

@@ -1,11 +1,12 @@
 from pathlib import Path
+import numpy as np
 import torch
 from torchvision.utils import make_grid, save_image
 from tinygrad.nn.state import get_parameters
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import trange
 from tinygrad.nn import optim
-from tinygrad.nn.datasets import mnist
+from extra.datasets import fetch_mnist
 
 class LinearGen:
   def __init__(self):
@@ -37,14 +38,14 @@ class LinearDisc:
     return x
 
 def make_batch(images):
-  sample = Tensor.randint(batch_size, low=0, high=images.shape[0])
-  return images[sample].reshape(batch_size, 28*28).cast('float').div(127.5).sub(1.0)
+  sample = np.random.randint(0, len(images), size=(batch_size))
+  image_b = images[sample].reshape(-1, 28*28).astype(np.float32) / 127.5 - 1.0
+  return Tensor(image_b)
 
 def make_labels(bs, col, val=-2.0):
-  y = Tensor.zeros(bs, 2)
-  if col == 0: y = y + Tensor([val, 0.0])
-  else: y = y + Tensor([0.0, val])
-  return y
+  y = np.zeros((bs, 2), np.float32)
+  y[range(bs), [col] * bs] = val  # Can we do label smoothing? i.e -2.0 changed to -1.98789.
+  return Tensor(y)
 
 def train_discriminator(optimizer, data_real, data_fake):
   real_labels = make_labels(batch_size, 1)
@@ -70,12 +71,12 @@ def train_generator(optimizer, data_fake):
 
 if __name__ == "__main__":
   # data for training and validation
-  X_train, _, _, _ = mnist()
+  images_real = np.vstack(fetch_mnist()[::2])
   ds_noise = Tensor.randn(64, 128, requires_grad=False)
   # parameters
   epochs, batch_size, k = 300, 512, 1
   sample_interval = epochs // 10
-  n_steps = X_train.shape[0] // batch_size
+  n_steps = len(images_real) // batch_size
   # models and optimizer
   generator = LinearGen()
   discriminator = LinearDisc()
@@ -83,24 +84,24 @@ if __name__ == "__main__":
   output_dir = Path(".").resolve() / "outputs"
   output_dir.mkdir(exist_ok=True)
   # optimizers
-  optim_g = optim.Adam(get_parameters(generator), lr=0.0002, b1=0.5)  # 0.0002 for equilibrium!
-  optim_d = optim.Adam(get_parameters(discriminator), lr=0.0002, b1=0.5)
+  optim_g = optim.Adam(get_parameters(generator),lr=0.0002, b1=0.5)  # 0.0002 for equilibrium!
+  optim_d = optim.Adam(get_parameters(discriminator),lr=0.0002, b1=0.5)
   # training loop
-  with Tensor.train():
-    for epoch in (t := trange(epochs)):
-      loss_g, loss_d = 0.0, 0.0
-      for _ in range(n_steps):
-        data_real = make_batch(X_train)
-        for step in range(k):  # Try with k = 5 or 7.
-          noise = Tensor.randn(batch_size, 128)
-          data_fake = generator.forward(noise).detach()
-          loss_d += train_discriminator(optim_d, data_real, data_fake)
+  Tensor.training = True
+  for epoch in (t := trange(epochs)):
+    loss_g, loss_d = 0.0, 0.0
+    for _ in range(n_steps):
+      data_real = make_batch(images_real)
+      for step in range(k):  # Try with k = 5 or 7.
         noise = Tensor.randn(batch_size, 128)
-        data_fake = generator.forward(noise)
-        loss_g += train_generator(optim_g, data_fake)
-      if (epoch + 1) % sample_interval == 0:
-        fake_images = generator.forward(ds_noise).detach().numpy()
-        fake_images = (fake_images.reshape(-1, 1, 28, 28) + 1) / 2  # 0 - 1 range.
-        save_image(make_grid(torch.tensor(fake_images)), output_dir / f"image_{epoch+1}.jpg")
-      t.set_description(f"Generator loss: {loss_g/n_steps}, Discriminator loss: {loss_d/n_steps}")
+        data_fake = generator.forward(noise).detach()
+        loss_d += train_discriminator(optim_d, data_real, data_fake)
+      noise = Tensor.randn(batch_size, 128)
+      data_fake = generator.forward(noise)
+      loss_g += train_generator(optim_g, data_fake)
+    if (epoch + 1) % sample_interval == 0:
+      fake_images = generator.forward(ds_noise).detach().numpy()
+      fake_images = (fake_images.reshape(-1, 1, 28, 28) + 1) / 2  # 0 - 1 range.
+      save_image(make_grid(torch.tensor(fake_images)), output_dir / f"image_{epoch+1}.jpg")
+    t.set_description(f"Generator loss: {loss_g/n_steps}, Discriminator loss: {loss_d/n_steps}")
   print("Training Completed!")

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), OPENPILOT_HACKS=1):
+    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
       ret = run_onnx_jit(**inputs).numpy()
     # copy i == 1 so use of JITBEAM is okay
     if i == 1: test_val = np.copy(ret)

examples/openpilot/load_pickle.py

@@ -1,16 +0,0 @@
-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,6 +6,7 @@ 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
@@ -335,7 +336,6 @@ 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_memoryview())
+        data = bytes(state_dict[name].uop.base.realized.as_buffer())
         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.randint(1, 77, low=0, high=49408, dtype=dtypes.int32)], forward = model.cond_stage_model.transformer.text_model),
+    Step(name = "textModel", input = [Tensor.randn(1, 77)], 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)

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 = vram_bar, doorbell_bar, mmio_bar
+    self.vram, self.doorbell64, self.mmio, self.dma_regions = vram_bar, doorbell_bar, mmio_bar, None
     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:", "Processing accelerators:"]
+    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
     lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
     self.lspci = {l.split()[0]: l.split(" ", 1)[1] for l in lspci}
     for k,v in self.lspci.items():
@@ -153,8 +153,7 @@ class SMICtx:
     tables = {}
     for dev in self.devs:
       match dev.ip_ver[am.MP1_HWIP]:
-        case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableV0_t
-        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTable_t
+        case (13,0,6)|(13,0,12): table_t = dev.smu.smu_mod.MetricsTableX_t
         case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
       tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
     return tables
@@ -231,11 +230,12 @@ class SMICtx:
 
   def get_power(self, dev, metrics):
     match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
-      case (13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.SocketPowerLimit)
+      case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
       case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
 
   def get_mem_usage(self, dev):
+    return 0
+
     usage = 0
     pt_stack = [dev.mm.root_page_table]
     while len(pt_stack) > 0:
@@ -244,8 +244,8 @@ class SMICtx:
         entry = pt.entries[i]
 
         if (entry & am.AMDGPU_PTE_VALID) == 0: continue
-        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))
+        if pt.lv!=am.AMDGPU_VM_PTB and not dev.gmc.is_pte_huge_page(pt.lv, entry):
+          pt_stack.append(AMPageTableEntry(dev, entry & 0x0000FFFFFFFFF000, lv=pt.lv+1))
           continue
         if (entry & am.AMDGPU_PTE_SYSTEM) != 0: continue
         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, activity_line_width, opt_text=mem_fmt)}"] \
+                    + [f"MEM Usage    {draw_bar((mem_used / mem_total) / 100, 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,18 +1,12 @@
 #!/usr/bin/env python3
 
-import os
 from tinygrad.helpers import Context
-from tinygrad.runtime.support.system import System, PCIDevice
-from tinygrad.runtime.support.hcq import FileIOInterface
+from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
 from tinygrad.runtime.support.am.amdev import AMDev
 
 if __name__ == "__main__":
   gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
-  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) for gpu in gpus]
+  pcidevs = [PCIDevice(f"reset:{gpu}", 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):
-    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
+  def __init__(self, pci_dev, dma_regions=None):
+    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
     self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
     self._run_discovery()
     self._build_regs()
@@ -19,9 +19,8 @@ amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
 amdev.AMDev = AMDFake
 from tinygrad.runtime.ops_amd import PCIIface
 
-def parse_amdgpu_logs(log_content, register_names=None, register_objects=None, *, only_xcc0: bool = False):
+def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
   register_map = register_names or {}
-  register_objs = register_objects or {}
 
   def replace_register(match):
     reg = match.group(1)
@@ -38,28 +37,6 @@ def parse_amdgpu_logs(log_content, register_names=None, register_objects=None, *
   # 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 = []
@@ -73,18 +50,16 @@ 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, reg_objs, only_xcc0=only_xcc0)
+  processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
 
   with open(sys.argv[2], 'w') as f:
     f.write(processed_log)

extra/assembly/amd/README

@@ -1,17 +1,17 @@
 An integrated environment for AMD GPU assembly and emulation
 
-Test with `pytest -n12 test/amd/`
-`AMD_LLVM=1 pytest -n12 test/amd/`
+Test with `PYTHONPATH="." pytest -n12 extra/assembly/amd/`
+`AMD_LLVM=1 PYTHONPATH="." pytest -n12 extra/assembly/amd/`
 
+* pdf.py -- extract assembly format + instruction pseudocode from AMD PDF
 * dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
-* test/mockgpu/amd/emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
-* generate.py -- extract assembly format + instruction pseudocode from AMD XML + PDF
-* test/mockgpu/amd/pcode.py -- pseudocode to UOp transformation
-* sqtt.py -- SQTT parser
+* pcode.py -- pseudocode execution environment. pseudocode should be transformed as little as possible.
+* asm.py -- an asm/disasm function to transform to and from AMD assembly syntax
+* emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
 
-The code should be as readable and deduplicated as possible. emu (in test/mockgpu/amd/) shouldn't be required for dsl.
+The code should be as readable and deduplicated as possible. asm and emu shouldn't be required for dsl.
 
-The autogen folder is autogenerated from the AMD PDFs with `python3 -m tinygrad.renderer.amd.pdf --arch all`
+The autogen folder is autogenerated from the AMD PDFs with `python3 -m extra.assembly.amd.pdf --arch all`
 
 test_emu.py has a good set of instruction tests for the emulation, with USE_HW=1 it will compare to real hardware.
 Whenever an instruction is fixed, regression tests should be added here and confirmed with real hardware.
@@ -20,20 +20,20 @@ test_llvm.py tests asm/disasm on the LLVM tests, confirming it behaves the same
 
 tinygrad's dtype tests should pass with and without LLVM. they run in about 12 seconds.
 
-`AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
-`AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
+`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
+`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
 
 The ops tests also pass, but they are very slow, so you should run them one at a time.
 
-`SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/backend/test_ops.py`
-`SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/backend/test_ops.py`
+`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_ops.py`
+`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_ops.py`
 
-When something is caught by main tinygrad tests, a local regression test should be added to `test/amd`.
+When something is caught by main tinygrad tests, a local regression test should be added to `extra/assembly/amd/test`.
 While working with tinygrad, you can dump the assembly with `DEBUG=7`. These tests all pass on real hardware
 If a test is failing with `AMD=1 PYTHON_REMU=1 MOCKGPU=1` it's because an instruction is emulated incorrectly.
 You can test without `MOCKGPU=1` to test on real hardware, if it works on real hardware there's a bug in the emulator.
 IMPORTANT: if a test is failing in the emulator, it's an instruction bug. Use DEBUG=7, get the instructions, and debug.
 
-Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~3000 lines.
-Get line count with `cloc --by-file tinygrad/renderer/amd/*.py`
+Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~2000 lines.
+Get line count with `cloc --by-file extra/assembly/amd/*.py`
 

extra/assembly/amd/amdxml.py

@@ -1,7 +1,7 @@
 # AMD ISA code generator - generates enum.py, ins.py, operands.py, str_pcode.py
 # Sources: XML from https://gpuopen.com/download/machine-readable-isa/latest/
 #          PDF manuals from AMD documentation
-import re, zlib, xml.etree.ElementTree as ET, zipfile, pathlib
+import re, zlib, xml.etree.ElementTree as ET, zipfile
 from tinygrad.helpers import fetch
 
 # ═══════════════════════════════════════════════════════════════════════════════
@@ -22,8 +22,6 @@ FIXES = {"rdna3": {"SOPK": {22: "S_SUBVECTOR_LOOP_BEGIN", 23: "S_SUBVECTOR_LOOP_
          "rdna4": {"SOP1": {80: "S_GET_BARRIER_STATE", 81: "S_BARRIER_INIT", 82: "S_BARRIER_JOIN"}, "SOPP": {9: "S_WAITCNT", 21: "S_BARRIER_LEAVE"}},
          "cdna": {"DS": {152: "DS_GWS_SEMA_RELEASE_ALL", 154: "DS_GWS_SEMA_V", 156: "DS_GWS_SEMA_P"},
                   "VOP3P": {44: "V_MFMA_LD_SCALE_B32", 62: "V_MFMA_F32_16X16X8_XF32", 63: "V_MFMA_F32_32X32X4_XF32"}}}
-# Fields missing from XML but present in hardware (format: {arch: {encoding: [(name, hi, lo), ...]}})
-FIELD_FIXES = {"cdna": {"VOP3P": [("opsel_hi2", 14, 14)]}}
 # Encoding suffixes to strip (variants we don't generate separate classes for)
 _ENC_SUFFIXES = ("_NSA1",)
 # Encoding suffix to class suffix mapping (for variants we DO generate)
@@ -77,13 +75,8 @@ def parse_xml(filename: str):
   for ot in root.findall(".//OperandTypes/OperandType"):
     ot_name = ot.findtext("OperandTypeName")
     for field in ot.findall(".//Field"):
-      key = (ot_name, field.findtext("FieldName"))
-      if (enum_name := op_enum_map.get(key)):  # type: ignore[arg-type]
-        def _pv_val(pv: ET.Element) -> tuple[int, str]:
-          v, n = pv.findtext("Value"), pv.findtext("Name")
-          assert v is not None and n is not None
-          return int(v), n.upper()
-        enums[enum_name] = dict(_pv_val(pv) for pv in field.findall(".//PredefinedValue"))
+      if (enum_name := op_enum_map.get((ot_name, field.findtext("FieldName")))):
+        enums[enum_name] = {int(pv.findtext("Value")): pv.findtext("Name").upper() for pv in field.findall(".//PredefinedValue")}
   # Extract DataFormats with BitCount
   for df in root.findall("ISA/DataFormats/DataFormat"):
     name, bits = df.findtext("DataFormatName"), df.findtext("BitCount")
@@ -91,26 +84,17 @@ def parse_xml(filename: str):
   # Extract encoding definitions
   for enc in root.findall("ISA/Encodings/Encoding"):
     name = enc.findtext("EncodingName")
-    assert name is not None
     is_base = name.startswith("ENC_") or name in ("VOP3_SDST_ENC", "VOPDXY")
     is_variant = any(sfx in name for sfx in _ENC_SUFFIX_MAP)
     if not is_base and not is_variant: continue
     if any(s in name for s in _SKIP_ENCODINGS): continue
-    fields: list[tuple[str, int, int]] = []
-    for f in enc.findall(".//MicrocodeFormat/BitMap/Field"):
-      br = f.find("BitLayout/Range")
-      if br is None: continue
-      fn = f.findtext("FieldName")
-      assert fn is not None
-      fields.append((_norm_field(fn.lower()),
-        int(br.findtext("BitOffset") or 0) + int(br.findtext("BitCount") or 0) - 1, int(br.findtext("BitOffset") or 0)))
-    ident_list = enc.findall("EncodingIdentifiers/EncodingIdentifier")
-    ident = ident_list[0] if ident_list else None
+    fields = [(_norm_field(f.findtext("FieldName").lower()), int(f.find("BitLayout/Range").findtext("BitOffset") or 0) + int(f.find("BitLayout/Range").findtext("BitCount") or 0) - 1,
+               int(f.find("BitLayout/Range").findtext("BitOffset") or 0))
+              for f in enc.findall(".//MicrocodeFormat/BitMap/Field") if f.find("BitLayout/Range") is not None]
+    ident = (enc.findall("EncodingIdentifiers/EncodingIdentifier") or [None])[0]
     enc_field = next((f for f in fields if f[0] == "encoding"), None)
-    # For multi-dword formats, encoding field may be in higher dword but identifier is always in dword0; use % 32
-    enc_bits: str | None = None
-    if ident is not None and ident.text is not None and enc_field:
-      enc_bits = "".join(ident.text[len(ident.text)-1-b] for b in range(enc_field[1] % 32, (enc_field[2] % 32)-1, -1))
+    # For multi-dword formats, encoding field may be in higher dword but identifier pattern is always in dword0; use % 32
+    enc_bits = "".join(ident.text[len(ident.text)-1-b] for b in range(enc_field[1] % 32, (enc_field[2] % 32)-1, -1)) if ident is not None and enc_field else None
     base_name = _strip_enc(name)
     encodings[NAME_MAP.get(base_name, base_name)] = (fields, enc_bits)
   # Extract instruction opcodes and operand info
@@ -118,12 +102,9 @@ def parse_xml(filename: str):
   opcode_encs: dict[str, dict[int, set[str]]] = {}  # {base_fmt: {opcode: {enc_names}}}
   for instr in root.findall("ISA/Instructions/Instruction"):
     name = instr.findtext("InstructionName")
-    assert name is not None
     for enc in instr.findall("InstructionEncodings/InstructionEncoding"):
       if enc.findtext("EncodingCondition") != "default": continue
-      enc_enc_name = enc.findtext("EncodingName")
-      assert enc_enc_name is not None
-      base, opcode = _map_flat(_strip_enc(enc_enc_name), name), int(enc.findtext("Opcode") or 0)
+      base, opcode = _map_flat(_strip_enc(enc.findtext("EncodingName")), name), int(enc.findtext("Opcode") or 0)
       enc_name = NAME_MAP.get(base, base)
       # Encoding variants use the same Op enum as the base format
       base_enum = enc_name
@@ -137,21 +118,19 @@ def parse_xml(filename: str):
         elif base == "VGLOBAL": enums.setdefault("VFLAT", {})[opcode] = name
       enums.setdefault(base_enum, {})[opcode] = name
       # Extract operand info
-      op_info: dict[str, tuple[str | None, int, str | None]] = {}
-      for op in enc.findall("Operands/Operand"):
-        fn = op.findtext("FieldName")
-        if fn: op_info[fn.lower()] = (op.findtext("DataFormatName"), int(op.findtext("OperandSize") or 0), op.findtext("OperandType"))
+      op_info = {op.findtext("FieldName").lower(): (op.findtext("DataFormatName"), int(op.findtext("OperandSize") or 0), op.findtext("OperandType"))
+                 for op in enc.findall("Operands/Operand") if op.findtext("FieldName")}
       for fmt, _, otype in op_info.values():
         if fmt and fmt not in fmts: fmts[fmt] = 0
         if otype: op_types_set.add(otype)
       if op_info: types[(name, base_enum)] = op_info
-  # Find opcodes that only exist in a specific variant encoding (no base format version)
-  suffix_only_ops: dict[str, dict[str, set[int]]] = {}  # {suffix: {base_fmt: {opcodes}}}
+  # Find opcodes that only exist in _LIT encoding (no base format version)
+  lit_only_ops: dict[str, set[int]] = {}
   for base_fmt, opcodes in opcode_encs.items():
     for opcode, encs in opcodes.items():
-      suffix = next((s for s in _ENC_SUFFIX_MAP.values() if all(s in e for e in encs)), None)
-      if suffix is not None: suffix_only_ops.setdefault(suffix, {}).setdefault(base_fmt, set()).add(opcode)
-  return encodings, enums, types, fmts, op_types_set, suffix_only_ops
+      if all("_LIT" in e for e in encs):
+        lit_only_ops.setdefault(base_fmt, set()).add(opcode)
+  return encodings, enums, types, fmts, op_types_set, lit_only_ops
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # PDF parsing
@@ -162,9 +141,7 @@ def extract_pdf_text(url: str) -> list[list[tuple[float, float, str, str]]]:
   data = fetch(url).read_bytes()
   # Parse xref table to locate objects
   xref: dict[int, int] = {}
-  xref_match = re.search(rb'startxref\s+(\d+)', data)
-  assert xref_match is not None
-  pos = int(xref_match.group(1)) + 4
+  pos = int(re.search(rb'startxref\s+(\d+)', data).group(1)) + 4
   while data[pos:pos+7] != b'trailer':
     while data[pos:pos+1] in b' \r\n': pos += 1
     line_end = data.find(b'\n', pos)
@@ -185,19 +162,14 @@ def extract_pdf_text(url: str) -> list[list[tuple[float, float, str, str]]]:
     if not (m := re.search(rb'/Contents (\d+) 0 R', data[xref[n]:xref[n]+500])): continue
     stream = get_stream(int(m.group(1))).decode('latin-1')
     elements, font = [], ''
-    _RE_BT = (r'(/F[\d.]+) [\d.]+ Tf|([\d.+-]+) ([\d.+-]+) Td|[\d.+-]+ [\d.+-]+ [\d.+-]+ [\d.+-]+ ([\d.+-]+) ([\d.+-]+) Tm'
-              r'|<([0-9A-Fa-f]+)>.*?Tj|\[([^\]]+)\] TJ')
     for bt in re.finditer(r'BT(.*?)ET', stream, re.S):
       x, y = 0.0, 0.0
-      for sm in re.finditer(_RE_BT, bt.group(1)):
-        if sm.group(1): font = sm.group(1)
-        elif sm.group(2): x, y = x + float(sm.group(2)), y + float(sm.group(3))
-        elif sm.group(4): x, y = float(sm.group(4)), float(sm.group(5))
-        elif sm.group(6) and (t := bytes.fromhex(sm.group(6)).decode('latin-1')).strip():
-          elements.append((x, y, t, font))
-        elif sm.group(7):
-          t = ''.join(bytes.fromhex(h).decode('latin-1') for h in re.findall(r'<([0-9A-Fa-f]+)>', sm.group(7)))
-          if t.strip(): elements.append((x, y, t, font))
+      for m in re.finditer(r'(/F[\d.]+) [\d.]+ Tf|([\d.+-]+) ([\d.+-]+) Td|[\d.+-]+ [\d.+-]+ [\d.+-]+ [\d.+-]+ ([\d.+-]+) ([\d.+-]+) Tm|<([0-9A-Fa-f]+)>.*?Tj|\[([^\]]+)\] TJ', bt.group(1)):
+        if m.group(1): font = m.group(1)
+        elif m.group(2): x, y = x + float(m.group(2)), y + float(m.group(3))
+        elif m.group(4): x, y = float(m.group(4)), float(m.group(5))
+        elif m.group(6) and (t := bytes.fromhex(m.group(6)).decode('latin-1')).strip(): elements.append((x, y, t, font))
+        elif m.group(7) and (t := ''.join(bytes.fromhex(h).decode('latin-1') for h in re.findall(r'<([0-9A-Fa-f]+)>', m.group(7)))).strip(): elements.append((x, y, t, font))
     pages.append(sorted(elements, key=lambda e: (-e[1], e[0])))
   return pages
 
@@ -223,7 +195,7 @@ def extract_pcode(pages: list[list[tuple[float, float, str, str]]], name_to_op:
     else:
       next_page, next_y = page_idx, 0
     # Collect F6 text from current position to next instruction (pseudocode is at x ≈ 69)
-    lines: list[tuple[int, float, str]] = []
+    lines = []
     for p in range(page_idx, next_page + 1):
       start_y = y if p == page_idx else 800
       end_y = next_y if p == next_page else 0
@@ -246,12 +218,8 @@ def extract_pcode(pages: list[list[tuple[float, float, str, str]]], name_to_op:
 # Code generation
 # ═══════════════════════════════════════════════════════════════════════════════
 
-def write_common(all_fmts: dict[str, int], all_op_types: set[str], path: pathlib.Path) -> None:
-  lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
-  lines.append("class ReprEnum(Enum):")
-  lines.append('  """Enum with clean repr that roundtrips with eval()."""')
-  lines.append('  def __repr__(self): return f"{type(self).__name__}.{self.name}"')
-  lines.append("")
+def write_common(all_fmts, all_op_types, path):
+  lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
   lines.append("class Fmt(Enum):")
   for fmt in sorted(all_fmts.keys()): lines.append(f"  {fmt} = auto()")
   lines.append("")
@@ -264,12 +232,11 @@ def write_common(all_fmts: dict[str, int], all_op_types: set[str], path: pathlib
   with open(path, "w") as f: f.write("\n".join(lines))
 
 def write_enum(enums, path):
-  lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit",
-    "from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401", ""]
+  lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum", "from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType  # noqa: F401", ""]
   for name, ops in sorted(enums.items()):
     if not ops: continue
     suffix = "_E32" if name in ("VOP1", "VOP2", "VOPC") else "_E64" if name == "VOP3" else ""
-    lines.append(f"class {name}(ReprEnum):" if name in ("HWREG", "MSG") else f"class {name}Op(ReprEnum):")
+    lines.append(f"class {name}(Enum):" if name in ("HWREG", "MSG") else f"class {name}Op(Enum):")
     aliases = []
     for op, mem in sorted(ops.items()):
       msuf = suffix if name != "VOP3" or op < 512 else ""
@@ -279,8 +246,8 @@ def write_enum(enums, path):
     lines.append("")
   with open(path, "w") as f: f.write("\n".join(lines))
 
-def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
-  _VGPR_FIELDS = {"vdst", "vdstx", "vsrc0", "vsrc1", "vsrc2", "vsrc3", "vsrcx1", "vsrcy1", "vaddr", "vdata", "data", "data0", "data1", "addr", "vsrc"}
+def write_ins(encodings, enums, lit_only_ops, types, arch, path):
+  _VGPR_FIELDS = {"vdst", "vdstx", "vsrc0", "vsrc1", "vsrc2", "vsrc3", "vsrcx1", "vsrcy1", "vaddr", "vdata", "data", "data0", "data1", "addr"}
   _VARIANT_SUFFIXES = ("_LIT", "_DPP16", "_DPP8", "_SDWA_SDST", "_SDWA", "_MFMA")
   def get_base_fmt(fmt):
     for sfx in _VARIANT_SUFFIXES: fmt = fmt.replace(sfx, "")
@@ -300,8 +267,6 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
     if name.startswith("ssrc") and bits == 8: return f"SSrcField({hi}, {lo})"
     if name in ("saddr", "soffset") and bits == 8: return f"SSrcField({hi}, {lo}, default=NULL)"
     if name.startswith("src") and bits == 9: return f"SrcField({hi}, {lo})"
-    # GLOBAL/SCRATCH: offset is 13-bit signed [12:0], FLAT: 12-bit unsigned (XML has 12-bit for all)
-    if name == "offset" and base_fmt in ("GLOBAL", "SCRATCH"): return f"BitField(12, {lo})"
     if base_fmt == "VOP3P" and name == "opsel_hi": return f"BitField({hi}, {lo}, default=3)"
     if base_fmt == "VOP3P" and name == "opsel_hi2": return f"BitField({hi}, {lo}, default=1)"
     return f"BitField({hi}, {lo})"
@@ -313,7 +278,7 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
            'dpp', 'fi', 'bc', 'row_mask', 'bank_mask', 'src0_neg', 'src0_abs', 'src1_neg', 'src1_abs',
            'cbsz', 'abid', 'acc_cd', 'acc', 'blgp', 'lane_sel_0', 'lane_sel_1', 'lane_sel_2', 'lane_sel_3',
            'lane_sel_4', 'lane_sel_5', 'lane_sel_6', 'lane_sel_7', 'dst_sel', 'dst_unused', 'src0_sel', 'src1_sel']
-  def sort_fields(fields): return sorted(fields, key=lambda f: (ORDER.index(f[0]) if f[0] in ORDER else 999, f[2]))
+  sort_fields = lambda fields: sorted(fields, key=lambda f: (ORDER.index(f[0]) if f[0] in ORDER else 999, f[2]))
 
   # Separate base encodings from variants
   base_encodings, variant_encodings = {}, {}
@@ -323,29 +288,15 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
     else: variant_encodings[enc_name] = data
 
   # Build sets of ops by their vdst type from operand metadata
-  sdst_opcodes: dict[str, set[int]] = {}  # ops where vdst is OPR_SREG (writes to SGPR)
+  sdst_opcodes = {}  # ops where vdst is OPR_SREG (writes to SGPR)
   for fmt, ops in enums.items():
     for op, name in ops.items():
       op_types = types.get((name, fmt), {})
       vdst_type = op_types.get("vdst", (None, None, None))[2]
       if vdst_type == "OPR_SREG": sdst_opcodes.setdefault(fmt, set()).add(op)
 
-  # collect only the XxxOp enums that are actually referenced in this arch's instruction definitions
-  enum_names = sorted(f"{k}Op" for k in enums if enums[k] and k not in ("HWREG", "MSG"))
-  # also re-export HWREG/MSG enums (plain enums, not instruction format ops)
-  enum_names += sorted(k for k in enums if k in ("HWREG", "MSG") and enums[k])
-  # collect DSL field types actually used by scanning generated field definitions
-  all_field_defs = " ".join(field_def(fn, hi, lo, enc, eb) for enc, (flds, eb) in encodings.items() for fn, hi, lo in flds)
-  _ALL_DSL = ["BitField", "EnumBitField", "FixedBitField", "NULL", "SBaseField", "SGPRField", "SRsrcField",
-              "SSrcField", "SrcField", "VDSTYField", "VGPRField"]
-  dsl_names = ["Inst"] + [n for n in _ALL_DSL if n in all_field_defs]
-  # also re-export register names so `from ins import *` still provides them to downstream users
-  _DSL_REGS = ["s", "v", "src", "VCC_LO", "VCC_HI", "VCC", "EXEC_LO", "EXEC_HI", "EXEC", "NULL", "OFF", "M0",
-               "SCC", "VCCZ", "EXECZ", "ttmp", "INV_2PI", "SDWA", "DPP", "DPP16", "LIT", "SRC_LDS_DIRECT"]
-  dsl_reexport = sorted(set(dsl_names + _DSL_REGS))
-  lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit", "# ruff: noqa: E501,F401",
-    f"from tinygrad.renderer.amd.dsl import {', '.join(dsl_reexport)}",
-    f"from tinygrad.runtime.autogen.amd.{arch}.enum import {', '.join(enum_names)}", "import functools", ""]
+  lines = ["# autogenerated from AMD ISA XML - do not edit", "# ruff: noqa: F401,F403",
+           "from extra.assembly.amd.dsl import *", f"from extra.assembly.amd.autogen.{arch}.enum import *", "import functools", ""]
 
   def fmt_allowed(op_enum: str, ops: set[int]) -> str:
     """Format allowed ops as {EnumName.MEMBER, ...}."""
@@ -354,19 +305,14 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
 
   # Generate base classes first
   for enc_name, (fields, enc_bits) in sorted(base_encodings.items()):
+    # Get lit-only ops for this format (these can't be used in base class)
+    base_lit_ops = lit_only_ops.get(enc_name, set())
     all_ops = set(enums.get(enc_name, {}).keys())
-    # Get suffix-only ops for this format (these can't be used in base class)
-    base_suffix_ops = set().union(*(d.get(enc_name, set()) for d in suffix_only_ops.values()))
     # Exclude SDST ops from base class (they need VOP1_SDST/VOP3_SDST/VOP3B)
-    base_allowed = all_ops - base_suffix_ops - sdst_opcodes.get(enc_name, set())
-    # RDNA3 FLAT/GLOBAL/SCRATCH share encoding bits, differentiated by seg field
-    # RDNA4 VFLAT/VGLOBAL/VSCRATCH have distinct encoding bits, no seg field needed
-    has_seg_field = any(fn == "seg" for fn, _, _ in fields)
-    if enc_name in ("FLAT", "VFLAT") and has_seg_field:
+    base_allowed = all_ops - base_lit_ops - sdst_opcodes.get(enc_name, set())
+    if enc_name in ("FLAT", "VFLAT"):
       prefix = "V" if enc_name == "VFLAT" else ""
-      flat_variants = [(f"{prefix}FLAT", 0, f"{prefix}FLATOp"), (f"{prefix}GLOBAL", 2, f"{prefix}GLOBALOp"),
-                       (f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]
-      for cls, seg, op_enum in flat_variants:
+      for cls, seg, op_enum in [(f"{prefix}FLAT", 0, f"{prefix}FLATOp"), (f"{prefix}GLOBAL", 2, f"{prefix}GLOBALOp"), (f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]:
         cls_ops = set(enums.get(cls, {}).keys())
         lines.append(f"class {cls}(Inst):")
         for fn, hi, lo in sort_fields(fields):
@@ -374,7 +320,7 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
           elif fn == "op": lines.append(f"  op = EnumBitField({hi}, {lo}, {op_enum}, {fmt_allowed(op_enum, cls_ops)})")
           else: lines.append(f"  {fn} = {field_def(fn, hi, lo, cls, enc_bits)}")
         lines.append("")
-    elif enc_name not in ("FLAT_GLOBAL", "FLAT_SCRATCH", "FLAT_GLBL", "DPP", "SDWA"):
+    elif enc_name not in ("FLAT_GLOBAL", "FLAT_SCRATCH", "FLAT_GLBL", "VGLOBAL", "VSCRATCH", "DPP", "SDWA"):
       lines.append(f"class {enc_name}(Inst):")
       for fn, hi, lo in sort_fields(fields):
         if fn == "op":
@@ -390,19 +336,15 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
     if base not in base_encodings: continue  # skip if no base class
     base_fields = {f[0] for f in base_encodings[base][0]}
     extra_fields = [(fn, hi, lo) for fn, hi, lo in fields if fn not in base_fields]
-    # Check if this is a suffix-only variant
-    variant_suffix = next((sfx for sfx in _VARIANT_SUFFIXES if enc_name.endswith(sfx)), None)
-    is_suffix_variant = variant_suffix in suffix_only_ops
+    is_lit = enc_name.endswith("_LIT")
     all_ops = set(enums.get(base, {}).keys())
-    if extra_fields or is_suffix_variant:
+    if extra_fields or is_lit:
       lines.append(f"class {enc_name}({base}):")
       op_field = next((f for f in base_encodings[base][0] if f[0] == "op"), None)
       # _LIT classes: override op to allow all opcodes (base excludes lit-only ops)
-      # other classes override op to only suffix-only opcodes
-      if op_field and is_suffix_variant:
+      if op_field and is_lit:
         _, hi, lo = op_field
-        allowed_ops = all_ops if variant_suffix == "_LIT" else suffix_only_ops[variant_suffix][base]
-        lines.append(f"  op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', allowed_ops)})")
+        lines.append(f"  op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', all_ops)})")
       for fn, hi, lo in sort_fields(extra_fields):
         lines.append(f"  {fn} = {field_def(fn, hi, lo, enc_name)}")
       lines.append("")
@@ -436,28 +378,23 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
   for fmt, ops in sorted(enums.items()):
     if fmt not in base_encodings and fmt not in ("GLOBAL", "SCRATCH", "VGLOBAL", "VSCRATCH"): continue
     suffix = "_E32" if fmt in ("VOP1", "VOP2", "VOPC") else "_E64" if fmt == "VOP3" else ""
-    op_to_suffix = {op:suffix for suffix,ops in suffix_only_ops.items() for op in ops.get(fmt, set())}
+    lit_ops = lit_only_ops.get(fmt, set())
     fmt_sdst_ops = sdst_opcodes.get(fmt, set())
     for op, name in sorted(ops.items()):
-      # ADDTID ops are in both FLAT and GLOBAL enums (for pcode); only generate helper for GLOBAL/VGLOBAL
-      if "ADDTID" in name and fmt in ("FLAT", "VFLAT"): continue
       msuf = suffix if fmt != "VOP3" or op < 512 else ""
-      # Determine class: SDST variants, suffix-specific variants (e.g., _MFMA, _LIT), or base
+      # Determine class: SDST variants, LIT-only instructions, or base
       if fmt == "VOP1" and op in fmt_sdst_ops: cls = "VOP1_SDST"
       elif fmt == "VOP3" and (op in fmt_sdst_ops or op < 256): cls = "VOP3_SDST"
-      elif op_to_suffix.get(op): cls = f"{fmt}{op_to_suffix[op]}"
+      elif op in lit_ops: cls = f"{fmt}_LIT"
       else: cls = fmt
       lines.append(f"{name.lower()}{msuf.lower()} = functools.partial({cls}, {fmt}Op.{name}{msuf})")
   with open(path, "w") as f: f.write("\n".join(lines))
 
-def write_operands(types: dict, enums: dict, arch: str, path: pathlib.Path) -> None:
+def write_operands(types, enums, arch, path):
   valid = {(name, fmt) for fmt, ops in enums.items() for name in ops.values()}
-  # only import enums that are actually used as keys in OPERANDS
-  used_bases = {eb for (nm, eb) in types if (nm, eb) in valid}
-  enum_names = sorted(f"{k}Op" for k in used_bases)
-  lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit",
-    "from tinygrad.runtime.autogen.amd.common import Fmt, OpType",
-    f"from tinygrad.runtime.autogen.amd.{arch}.enum import {', '.join(enum_names)}", ""]
+  lines = ["# autogenerated from AMD ISA XML - do not edit",
+           "from extra.assembly.amd.autogen.common import Fmt, OpType",
+           f"from extra.assembly.amd.autogen.{arch}.enum import *", ""]
   lines.append("# instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}")
   lines.append("OPERANDS = {")
   def fmt_val(v):
@@ -470,7 +407,7 @@ def write_operands(types: dict, enums: dict, arch: str, path: pathlib.Path) -> N
   lines.append("}")
   with open(path, "w") as f: f.write("\n".join(lines))
 
-def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, str]], arch: str, path: pathlib.Path) -> None:
+def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, str]], arch: str, path: str):
   """Write str_pcode.py file from extracted pseudocode."""
   entries: list[tuple[str, str, int, str]] = []
   for fmt_name, ops in enums.items():
@@ -481,7 +418,7 @@ def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, st
         entries.append((f"{fmt_name}Op", f"{name}{msuf}", opcode, pcode[(name, opcode)]))
   enum_names = sorted(set(e[0] for e in entries))
   lines = ["# autogenerated from AMD ISA PDF - do not edit", "# ruff: noqa: E501",
-           f"from tinygrad.runtime.autogen.amd.{arch}.enum import {', '.join(enum_names)}", "", "PCODE = {"]
+           f"from extra.assembly.amd.autogen.{arch}.enum import {', '.join(enum_names)}", "", "PCODE = {"]
   for enum_name, name, opcode, code in sorted(entries, key=lambda x: (x[0], x[2])):
     lines.append(f"  {enum_name}.{name}: {code!r},")
   lines.append("}")
@@ -492,31 +429,27 @@ def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, st
 # ═══════════════════════════════════════════════════════════════════════════════
 
 if __name__ == "__main__":
-  all_fmts: dict[str, int] = {}
-  all_op_types: set[str] = set()
-  arch_data: dict[str, dict] = {}
+  import pathlib
+  all_fmts, all_op_types, arch_data = {}, set(), {}
   # First pass: parse XML for all architectures
   for arch, cfg in ARCHS.items():
     print(f"Parsing XML: {cfg['xml']} -> {arch}")
-    encodings, enums, types, fmts, op_types_set, suffix_only_ops = parse_xml(cfg["xml"])
+    encodings, enums, types, fmts, op_types_set, lit_only_ops = parse_xml(cfg["xml"])
     for fmt, ops in FIXES.get(arch, {}).items(): enums.setdefault(fmt, {}).update(ops)
-    for fmt, fields in FIELD_FIXES.get(arch, {}).items():
-      if fmt in encodings: encodings[fmt] = (encodings[fmt][0] + fields, encodings[fmt][1])
-    arch_data[arch] = {"encodings": encodings, "enums": enums, "types": types, "suffix_only_ops": suffix_only_ops}
+    arch_data[arch] = {"encodings": encodings, "enums": enums, "types": types, "lit_only_ops": lit_only_ops}
     for fmt, bits in fmts.items():
       assert fmt not in all_fmts or all_fmts[fmt] == bits, f"FMT_BITS mismatch for {fmt}: {all_fmts[fmt]} vs {bits}"
       all_fmts[fmt] = bits
     all_op_types.update(op_types_set)
   # Write common.py
-  autogen_base = pathlib.Path(__file__).parents[2] / "runtime" / "autogen" / "amd"
-  common_path = autogen_base / "common.py"
+  common_path = pathlib.Path(__file__).parent / "autogen" / "common.py"
   write_common(all_fmts, all_op_types, common_path)
   print(f"Wrote common.py: {len(all_fmts)} formats, {len(all_op_types)} op types")
   # Write per-arch files from XML
   for arch, data in arch_data.items():
-    base = autogen_base / arch
+    base = pathlib.Path(__file__).parent / "autogen" / arch
     write_enum(data["enums"], base / "enum.py")
-    write_ins(data["encodings"], data["enums"], data["suffix_only_ops"], data["types"], arch, base / "ins.py")
+    write_ins(data["encodings"], data["enums"], data["lit_only_ops"], data["types"], arch, base / "ins.py")
     write_operands(data["types"], data["enums"], arch, base / "operands.py")
     print(f"  {arch}: {len(data['encodings'])} encodings, {sum(len(v) for v in data['enums'].values())} instructions")
   # Second pass: parse PDFs and write pcode
@@ -525,6 +458,6 @@ if __name__ == "__main__":
     pages = extract_pdf_text(cfg["pdf"])
     name_to_op = {name: op for ops in arch_data[arch]["enums"].values() for op, name in ops.items()}
     pcode = extract_pcode(pages, name_to_op)
-    base = autogen_base / arch
+    base = pathlib.Path(__file__).parent / "autogen" / arch
     write_pcode(pcode, arch_data[arch]["enums"], arch, base / "str_pcode.py")
     print(f"  {arch}: {len(pcode)} pcode entries")

extra/assembly/amd/autogen/cdna/enum.py

@@ -1,7 +1,8 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401
+from enum import Enum
+from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType  # noqa: F401
 
-class DSOp(ReprEnum):
+class DSOp(Enum):
   DS_ADD_U32 = 0
   DS_SUB_U32 = 1
   DS_RSUB_U32 = 2
@@ -132,7 +133,7 @@ class DSOp(ReprEnum):
   DS_READ_B96 = 254
   DS_READ_B128 = 255
 
-class FLATOp(ReprEnum):
+class FLATOp(Enum):
   FLAT_LOAD_UBYTE = 16
   FLAT_LOAD_SBYTE = 17
   FLAT_LOAD_USHORT = 18
@@ -188,7 +189,7 @@ class FLATOp(ReprEnum):
   FLAT_ATOMIC_INC_X2 = 107
   FLAT_ATOMIC_DEC_X2 = 108
 
-class GLOBALOp(ReprEnum):
+class GLOBALOp(Enum):
   GLOBAL_LOAD_UBYTE = 16
   GLOBAL_LOAD_SBYTE = 17
   GLOBAL_LOAD_USHORT = 18
@@ -251,7 +252,7 @@ class GLOBALOp(ReprEnum):
   GLOBAL_LOAD_LDS_DWORDX4 = 125
   GLOBAL_LOAD_LDS_DWORDX3 = 126
 
-class HWREG(ReprEnum):
+class HWREG(Enum):
   HW_REG_MODE = 1
   HW_REG_STATUS = 2
   HW_REG_TRAPSTS = 3
@@ -277,7 +278,7 @@ class HWREG(ReprEnum):
   HW_REG_SQ_PERF_SNAPSHOT_PC_LO = 23
   HW_REG_SQ_PERF_SNAPSHOT_PC_HI = 24
 
-class MTBUFOp(ReprEnum):
+class MTBUFOp(Enum):
   TBUFFER_LOAD_FORMAT_X = 0
   TBUFFER_LOAD_FORMAT_XY = 1
   TBUFFER_LOAD_FORMAT_XYZ = 2
@@ -295,7 +296,7 @@ class MTBUFOp(ReprEnum):
   TBUFFER_STORE_FORMAT_D16_XYZ = 14
   TBUFFER_STORE_FORMAT_D16_XYZW = 15
 
-class MUBUFOp(ReprEnum):
+class MUBUFOp(Enum):
   BUFFER_LOAD_FORMAT_X = 0
   BUFFER_LOAD_FORMAT_XY = 1
   BUFFER_LOAD_FORMAT_XYZ = 2
@@ -371,7 +372,7 @@ class MUBUFOp(ReprEnum):
   BUFFER_ATOMIC_INC_X2 = 107
   BUFFER_ATOMIC_DEC_X2 = 108
 
-class SCRATCHOp(ReprEnum):
+class SCRATCHOp(Enum):
   SCRATCH_LOAD_UBYTE = 16
   SCRATCH_LOAD_SBYTE = 17
   SCRATCH_LOAD_USHORT = 18
@@ -400,7 +401,7 @@ class SCRATCHOp(ReprEnum):
   SCRATCH_LOAD_LDS_SSHORT = 41
   SCRATCH_LOAD_LDS_DWORD = 42
 
-class SMEMOp(ReprEnum):
+class SMEMOp(Enum):
   S_LOAD_DWORD = 0
   S_LOAD_DWORDX2 = 1
   S_LOAD_DWORDX4 = 2
@@ -486,7 +487,7 @@ class SMEMOp(ReprEnum):
   S_ATOMIC_INC_X2 = 171
   S_ATOMIC_DEC_X2 = 172
 
-class SOP1Op(ReprEnum):
+class SOP1Op(Enum):
   S_MOV_B32 = 0
   S_MOV_B64 = 1
   S_CMOV_B32 = 2
@@ -542,7 +543,7 @@ class SOP1Op(ReprEnum):
   S_ANDN2_WREXEC_B64 = 54
   S_BITREPLICATE_B64_B32 = 55
 
-class SOP2Op(ReprEnum):
+class SOP2Op(Enum):
   S_ADD_U32 = 0
   S_SUB_U32 = 1
   S_ADD_I32 = 2
@@ -597,7 +598,7 @@ class SOP2Op(ReprEnum):
   S_PACK_LH_B32_B16 = 51
   S_PACK_HH_B32_B16 = 52
 
-class SOPCOp(ReprEnum):
+class SOPCOp(Enum):
   S_CMP_EQ_I32 = 0
   S_CMP_LG_I32 = 1
   S_CMP_GT_I32 = 2
@@ -619,7 +620,7 @@ class SOPCOp(ReprEnum):
   S_CMP_EQ_U64 = 18
   S_CMP_LG_U64 = 19
 
-class SOPKOp(ReprEnum):
+class SOPKOp(Enum):
   S_MOVK_I32 = 0
   S_CMOVK_I32 = 1
   S_CMPK_EQ_I32 = 2
@@ -642,7 +643,7 @@ class SOPKOp(ReprEnum):
   S_SETREG_IMM32_B32 = 20
   S_CALL_B64 = 21
 
-class SOPPOp(ReprEnum):
+class SOPPOp(Enum):
   S_NOP = 0
   S_ENDPGM = 1
   S_BRANCH = 2
@@ -676,7 +677,7 @@ class SOPPOp(ReprEnum):
   S_ENDPGM_ORDERED_PS_DONE = 30
   S_SET_VALU_COEXEC_MODE = 31
 
-class VOP1Op(ReprEnum):
+class VOP1Op(Enum):
   V_NOP_E32 = 0
   V_MOV_B32_E32 = 1
   V_READFIRSTLANE_B32_E32 = 2
@@ -854,7 +855,7 @@ class VOP1Op(ReprEnum):
   V_PERMLANE32_SWAP_B32 = V_PERMLANE32_SWAP_B32_E32
   V_CVT_F32_BF16 = V_CVT_F32_BF16_E32
 
-class VOP2Op(ReprEnum):
+class VOP2Op(Enum):
   V_CNDMASK_B32_E32 = 0
   V_ADD_F32_E32 = 1
   V_SUB_F32_E32 = 2
@@ -980,7 +981,7 @@ class VOP2Op(ReprEnum):
   V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
   V_XNOR_B32 = V_XNOR_B32_E32
 
-class VOP3Op(ReprEnum):
+class VOP3Op(Enum):
   V_CMP_CLASS_F32_E64 = 16
   V_CMPX_CLASS_F32_E64 = 17
   V_CMP_CLASS_F64_E64 = 18
@@ -1878,7 +1879,7 @@ class VOP3Op(ReprEnum):
   V_ADD_LSHL_U32 = V_ADD_LSHL_U32_E64
   V_ADD3_U32 = V_ADD3_U32_E64
 
-class VOP3POp(ReprEnum):
+class VOP3POp(Enum):
   V_PK_MAD_I16 = 0
   V_PK_MUL_LO_U16 = 1
   V_PK_ADD_I16 = 2
@@ -1987,11 +1988,11 @@ class VOP3POp(ReprEnum):
   V_SMFMAC_F32_32X32X32_FP8_BF8 = 126
   V_SMFMAC_F32_32X32X32_FP8_FP8 = 127
 
-class VOP3PX2Op(ReprEnum):
+class VOP3PX2Op(Enum):
   V_MFMA_SCALE_F32_16X16X128_F8F6F4 = 45
   V_MFMA_SCALE_F32_32X32X64_F8F6F4 = 46
 
-class VOP3SDOp(ReprEnum):
+class VOP3SDOp(Enum):
   V_ADD_CO_U32 = 281
   V_SUB_CO_U32 = 282
   V_SUBREV_CO_U32 = 283
@@ -2003,7 +2004,7 @@ class VOP3SDOp(ReprEnum):
   V_MAD_U64_U32 = 488
   V_MAD_I64_I32 = 489
 
-class VOPCOp(ReprEnum):
+class VOPCOp(Enum):
   V_CMP_CLASS_F32_E32 = 16
   V_CMPX_CLASS_F32_E32 = 17
   V_CMP_CLASS_F64_E32 = 18

extra/assembly/amd/autogen/cdna/ins.py

@@ -1,7 +1,7 @@
 # autogenerated from AMD ISA XML - do not edit
-# ruff: noqa: E501,F401
-from tinygrad.renderer.amd.dsl import BitField, DPP, DPP16, EXEC, EXECZ, EXEC_HI, EXEC_LO, EnumBitField, FixedBitField, INV_2PI, Inst, LIT, M0, NULL, OFF, SBaseField, SCC, SDWA, SGPRField, SRC_LDS_DIRECT, SRsrcField, SSrcField, SrcField, VCC, VCCZ, VCC_HI, VCC_LO, VGPRField, s, src, ttmp, v
-from tinygrad.runtime.autogen.amd.cdna.enum import DSOp, FLATOp, GLOBALOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3PX2Op, VOP3SDOp, VOPCOp, HWREG
+# ruff: noqa: F401,F403
+from extra.assembly.amd.dsl import *
+from extra.assembly.amd.autogen.cdna.enum import *
 import functools
 
 class DS(Inst):
@@ -38,7 +38,7 @@ class GLOBAL(Inst):
   addr = VGPRField(39, 32)
   data = VGPRField(47, 40)
   saddr = SGPRField(54, 48, default=NULL)
-  offset = BitField(12, 0)
+  offset = BitField(11, 0)
   seg = FixedBitField(15, 14, 2)
   acc = BitField(55, 55)
   sve = BitField(13, 13)
@@ -53,7 +53,7 @@ class SCRATCH(Inst):
   addr = VGPRField(39, 32)
   data = VGPRField(47, 40)
   saddr = SGPRField(54, 48, default=NULL)
-  offset = BitField(12, 0)
+  offset = BitField(11, 0)
   seg = FixedBitField(15, 14, 1)
   acc = BitField(55, 55)
   sve = BitField(13, 13)
@@ -164,7 +164,7 @@ class VOP3(Inst):
 
 class VOP3P(Inst):
   encoding = FixedBitField(31, 23, 0b110100111)
-  op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_PK_MAD_I16, VOP3POp.V_PK_MUL_LO_U16, VOP3POp.V_PK_ADD_I16, VOP3POp.V_PK_SUB_I16, VOP3POp.V_PK_LSHLREV_B16, VOP3POp.V_PK_LSHRREV_B16, VOP3POp.V_PK_ASHRREV_I16, VOP3POp.V_PK_MAX_I16, VOP3POp.V_PK_MIN_I16, VOP3POp.V_PK_MAD_U16, VOP3POp.V_PK_ADD_U16, VOP3POp.V_PK_SUB_U16, VOP3POp.V_PK_MAX_U16, VOP3POp.V_PK_MIN_U16, VOP3POp.V_PK_FMA_F16, VOP3POp.V_PK_ADD_F16, VOP3POp.V_PK_MUL_F16, VOP3POp.V_PK_MIN_F16, VOP3POp.V_PK_MAX_F16, VOP3POp.V_DOT2_F32_BF16, VOP3POp.V_PK_MINIMUM3_F16, VOP3POp.V_PK_MAXIMUM3_F16, VOP3POp.V_MAD_MIX_F32, VOP3POp.V_MAD_MIXLO_F16, VOP3POp.V_MAD_MIXHI_F16, VOP3POp.V_DOT2_F32_F16, VOP3POp.V_DOT2_I32_I16, VOP3POp.V_DOT2_U32_U16, VOP3POp.V_DOT4_I32_I8, VOP3POp.V_DOT4_U32_U8, VOP3POp.V_DOT8_I32_I4, VOP3POp.V_DOT8_U32_U4, VOP3POp.V_MFMA_LD_SCALE_B32, VOP3POp.V_PK_FMA_F32, VOP3POp.V_PK_MUL_F32, VOP3POp.V_PK_ADD_F32, VOP3POp.V_PK_MOV_B32, VOP3POp.V_MFMA_F32_16X16X8_XF32, VOP3POp.V_MFMA_F32_32X32X4_XF32, VOP3POp.V_ACCVGPR_READ, VOP3POp.V_ACCVGPR_WRITE})
+  op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_PK_MAD_I16, VOP3POp.V_PK_MUL_LO_U16, VOP3POp.V_PK_ADD_I16, VOP3POp.V_PK_SUB_I16, VOP3POp.V_PK_LSHLREV_B16, VOP3POp.V_PK_LSHRREV_B16, VOP3POp.V_PK_ASHRREV_I16, VOP3POp.V_PK_MAX_I16, VOP3POp.V_PK_MIN_I16, VOP3POp.V_PK_MAD_U16, VOP3POp.V_PK_ADD_U16, VOP3POp.V_PK_SUB_U16, VOP3POp.V_PK_MAX_U16, VOP3POp.V_PK_MIN_U16, VOP3POp.V_PK_FMA_F16, VOP3POp.V_PK_ADD_F16, VOP3POp.V_PK_MUL_F16, VOP3POp.V_PK_MIN_F16, VOP3POp.V_PK_MAX_F16, VOP3POp.V_DOT2_F32_BF16, VOP3POp.V_PK_MINIMUM3_F16, VOP3POp.V_PK_MAXIMUM3_F16, VOP3POp.V_MAD_MIX_F32, VOP3POp.V_MAD_MIXLO_F16, VOP3POp.V_MAD_MIXHI_F16, VOP3POp.V_DOT2_F32_F16, VOP3POp.V_DOT2_I32_I16, VOP3POp.V_DOT2_U32_U16, VOP3POp.V_DOT4_I32_I8, VOP3POp.V_DOT4_U32_U8, VOP3POp.V_DOT8_I32_I4, VOP3POp.V_DOT8_U32_U4, VOP3POp.V_MFMA_LD_SCALE_B32, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4, VOP3POp.V_PK_FMA_F32, VOP3POp.V_PK_MUL_F32, VOP3POp.V_PK_ADD_F32, VOP3POp.V_PK_MOV_B32, VOP3POp.V_MFMA_F32_16X16X32_BF16, VOP3POp.V_MFMA_I32_16X16X64_I8, VOP3POp.V_MFMA_F32_32X32X16_BF16, VOP3POp.V_MFMA_I32_32X32X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_BF16, VOP3POp.V_SMFMAC_I32_16X16X128_I8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X8_XF32, VOP3POp.V_MFMA_F32_32X32X4_XF32, VOP3POp.V_MFMA_F32_32X32X1_2B_F32, VOP3POp.V_MFMA_F32_16X16X1_4B_F32, VOP3POp.V_MFMA_F32_4X4X1_16B_F32, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X2_F32, VOP3POp.V_MFMA_F32_16X16X4_F32, VOP3POp.V_SMFMAC_F32_32X32X32_BF16, VOP3POp.V_SMFMAC_I32_32X32X64_I8, VOP3POp.V_MFMA_F32_32X32X4_2B_F16, VOP3POp.V_MFMA_F32_16X16X4_4B_F16, VOP3POp.V_MFMA_F32_4X4X4_16B_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X8_F16, VOP3POp.V_MFMA_F32_16X16X16_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8, VOP3POp.V_MFMA_I32_32X32X4_2B_I8, VOP3POp.V_MFMA_I32_16X16X4_4B_I8, VOP3POp.V_MFMA_I32_4X4X4_16B_I8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8, VOP3POp.V_MFMA_F32_16X16X32_F16, VOP3POp.V_MFMA_F32_32X32X16_F16, VOP3POp.V_MFMA_I32_32X32X16_I8, VOP3POp.V_MFMA_I32_16X16X32_I8, VOP3POp.V_ACCVGPR_READ, VOP3POp.V_ACCVGPR_WRITE, VOP3POp.V_SMFMAC_F32_16X16X64_F16, VOP3POp.V_SMFMAC_F32_32X32X32_F16, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16, VOP3POp.V_MFMA_F32_32X32X8_BF16, VOP3POp.V_MFMA_F32_16X16X16_BF16, VOP3POp.V_SMFMAC_F32_16X16X32_F16, VOP3POp.V_SMFMAC_F32_32X32X16_F16, VOP3POp.V_SMFMAC_F32_16X16X32_BF16, VOP3POp.V_SMFMAC_F32_32X32X16_BF16, VOP3POp.V_SMFMAC_I32_16X16X64_I8, VOP3POp.V_SMFMAC_I32_32X32X32_I8, VOP3POp.V_MFMA_F64_16X16X4_F64, VOP3POp.V_MFMA_F64_4X4X4_4B_F64, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8})
   vdst = VGPRField(7, 0)
   src0 = SrcField(40, 32)
   src1 = SrcField(49, 41)
@@ -174,7 +174,6 @@ class VOP3P(Inst):
   clmp = BitField(15, 15)
   opsel = BitField(13, 11)
   opsel_hi = BitField(60, 59, default=3)
-  opsel_hi2 = BitField(14, 14, default=1)
 
 class VOP3PX2(Inst):
   encoding = FixedBitField(95, 87, 0b110100111)
@@ -310,7 +309,6 @@ class VOP2_SDWA_SDST(VOP2):
   s1 = BitField(63, 63)
 
 class VOP3P_MFMA(VOP3P):
-  op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_MFMA_F32_16X16X128_F8F6F4, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4, VOP3POp.V_MFMA_F32_16X16X32_BF16, VOP3POp.V_MFMA_I32_16X16X64_I8, VOP3POp.V_MFMA_F32_32X32X16_BF16, VOP3POp.V_MFMA_I32_32X32X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_BF16, VOP3POp.V_SMFMAC_I32_16X16X128_I8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X1_2B_F32, VOP3POp.V_MFMA_F32_16X16X1_4B_F32, VOP3POp.V_MFMA_F32_4X4X1_16B_F32, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X2_F32, VOP3POp.V_MFMA_F32_16X16X4_F32, VOP3POp.V_SMFMAC_F32_32X32X32_BF16, VOP3POp.V_SMFMAC_I32_32X32X64_I8, VOP3POp.V_MFMA_F32_32X32X4_2B_F16, VOP3POp.V_MFMA_F32_16X16X4_4B_F16, VOP3POp.V_MFMA_F32_4X4X4_16B_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X8_F16, VOP3POp.V_MFMA_F32_16X16X16_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8, VOP3POp.V_MFMA_I32_32X32X4_2B_I8, VOP3POp.V_MFMA_I32_16X16X4_4B_I8, VOP3POp.V_MFMA_I32_4X4X4_16B_I8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8, VOP3POp.V_MFMA_F32_16X16X32_F16, VOP3POp.V_MFMA_F32_32X32X16_F16, VOP3POp.V_MFMA_I32_32X32X16_I8, VOP3POp.V_MFMA_I32_16X16X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_F16, VOP3POp.V_SMFMAC_F32_32X32X32_F16, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16, VOP3POp.V_MFMA_F32_32X32X8_BF16, VOP3POp.V_MFMA_F32_16X16X16_BF16, VOP3POp.V_SMFMAC_F32_16X16X32_F16, VOP3POp.V_SMFMAC_F32_32X32X16_F16, VOP3POp.V_SMFMAC_F32_16X16X32_BF16, VOP3POp.V_SMFMAC_F32_32X32X16_BF16, VOP3POp.V_SMFMAC_I32_16X16X64_I8, VOP3POp.V_SMFMAC_I32_32X32X32_I8, VOP3POp.V_MFMA_F64_16X16X4_F64, VOP3POp.V_MFMA_F64_4X4X4_4B_F64, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8})
   cbsz = BitField(10, 8)
   abid = BitField(14, 11)
   acc_cd = BitField(15, 15)
@@ -1649,80 +1647,80 @@ v_dot4_u32_u8 = functools.partial(VOP3P, VOP3POp.V_DOT4_U32_U8)
 v_dot8_i32_i4 = functools.partial(VOP3P, VOP3POp.V_DOT8_I32_I4)
 v_dot8_u32_u4 = functools.partial(VOP3P, VOP3POp.V_DOT8_U32_U4)
 v_mfma_ld_scale_b32 = functools.partial(VOP3P, VOP3POp.V_MFMA_LD_SCALE_B32)
-v_mfma_f32_16x16x128_f8f6f4 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4)
-v_mfma_f32_32x32x64_f8f6f4 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4)
+v_mfma_f32_16x16x128_f8f6f4 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4)
+v_mfma_f32_32x32x64_f8f6f4 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4)
 v_pk_fma_f32 = functools.partial(VOP3P, VOP3POp.V_PK_FMA_F32)
 v_pk_mul_f32 = functools.partial(VOP3P, VOP3POp.V_PK_MUL_F32)
 v_pk_add_f32 = functools.partial(VOP3P, VOP3POp.V_PK_ADD_F32)
 v_pk_mov_b32 = functools.partial(VOP3P, VOP3POp.V_PK_MOV_B32)
-v_mfma_f32_16x16x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF16)
-v_mfma_i32_16x16x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X64_I8)
-v_mfma_f32_32x32x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF16)
-v_mfma_i32_32x32x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X32_I8)
-v_smfmac_f32_16x16x64_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF16)
-v_smfmac_i32_16x16x128_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_16X16X128_I8)
-v_smfmac_f32_16x16x128_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8)
-v_smfmac_f32_16x16x128_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8)
-v_smfmac_f32_16x16x128_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8)
+v_mfma_f32_16x16x32_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF16)
+v_mfma_i32_16x16x64_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X64_I8)
+v_mfma_f32_32x32x16_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF16)
+v_mfma_i32_32x32x32_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X32_I8)
+v_smfmac_f32_16x16x64_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF16)
+v_smfmac_i32_16x16x128_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_16X16X128_I8)
+v_smfmac_f32_16x16x128_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8)
+v_smfmac_f32_16x16x128_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8)
+v_smfmac_f32_16x16x128_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8)
 v_mfma_f32_16x16x8_xf32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X8_XF32)
 v_mfma_f32_32x32x4_xf32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_XF32)
-v_mfma_f32_32x32x1_2b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X1_2B_F32)
-v_mfma_f32_16x16x1_4b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X1_4B_F32)
-v_mfma_f32_4x4x1_16b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X1_16B_F32)
-v_smfmac_f32_16x16x128_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8)
-v_mfma_f32_32x32x2_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X2_F32)
-v_mfma_f32_16x16x4_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_F32)
-v_smfmac_f32_32x32x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF16)
-v_smfmac_i32_32x32x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_32X32X64_I8)
-v_mfma_f32_32x32x4_2b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X4_2B_F16)
-v_mfma_f32_16x16x4_4b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_4B_F16)
-v_mfma_f32_4x4x4_16b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X4_16B_F16)
-v_smfmac_f32_32x32x64_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8)
-v_mfma_f32_32x32x8_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X8_F16)
-v_mfma_f32_16x16x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X16_F16)
-v_smfmac_f32_32x32x64_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8)
-v_smfmac_f32_32x32x64_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8)
-v_mfma_i32_32x32x4_2b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X4_2B_I8)
-v_mfma_i32_16x16x4_4b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X4_4B_I8)
-v_mfma_i32_4x4x4_16b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_4X4X4_16B_I8)
-v_smfmac_f32_32x32x64_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8)
-v_mfma_f32_16x16x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_F16)
-v_mfma_f32_32x32x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_F16)
-v_mfma_i32_32x32x16_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X16_I8)
-v_mfma_i32_16x16x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X32_I8)
+v_mfma_f32_32x32x1_2b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X1_2B_F32)
+v_mfma_f32_16x16x1_4b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X1_4B_F32)
+v_mfma_f32_4x4x1_16b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X1_16B_F32)
+v_smfmac_f32_16x16x128_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8)
+v_mfma_f32_32x32x2_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X2_F32)
+v_mfma_f32_16x16x4_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_F32)
+v_smfmac_f32_32x32x32_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF16)
+v_smfmac_i32_32x32x64_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_32X32X64_I8)
+v_mfma_f32_32x32x4_2b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_2B_F16)
+v_mfma_f32_16x16x4_4b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_4B_F16)
+v_mfma_f32_4x4x4_16b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X4_16B_F16)
+v_smfmac_f32_32x32x64_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8)
+v_mfma_f32_32x32x8_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X8_F16)
+v_mfma_f32_16x16x16_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X16_F16)
+v_smfmac_f32_32x32x64_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8)
+v_smfmac_f32_32x32x64_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8)
+v_mfma_i32_32x32x4_2b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X4_2B_I8)
+v_mfma_i32_16x16x4_4b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X4_4B_I8)
+v_mfma_i32_4x4x4_16b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_4X4X4_16B_I8)
+v_smfmac_f32_32x32x64_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8)
+v_mfma_f32_16x16x32_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_F16)
+v_mfma_f32_32x32x16_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_F16)
+v_mfma_i32_32x32x16_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X16_I8)
+v_mfma_i32_16x16x32_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X32_I8)
 v_accvgpr_read = functools.partial(VOP3P, VOP3POp.V_ACCVGPR_READ)
 v_accvgpr_write = functools.partial(VOP3P, VOP3POp.V_ACCVGPR_WRITE)
-v_smfmac_f32_16x16x64_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_F16)
-v_smfmac_f32_32x32x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_F16)
-v_mfma_f32_32x32x4_2b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16)
-v_mfma_f32_16x16x4_4b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16)
-v_mfma_f32_4x4x4_16b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16)
-v_mfma_f32_32x32x8_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X8_BF16)
-v_mfma_f32_16x16x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X16_BF16)
-v_smfmac_f32_16x16x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X32_F16)
-v_smfmac_f32_32x32x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X16_F16)
-v_smfmac_f32_16x16x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X32_BF16)
-v_smfmac_f32_32x32x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X16_BF16)
-v_smfmac_i32_16x16x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_16X16X64_I8)
-v_smfmac_i32_32x32x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_32X32X32_I8)
-v_mfma_f64_16x16x4_f64 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F64_16X16X4_F64)
-v_mfma_f64_4x4x4_4b_f64 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F64_4X4X4_4B_F64)
-v_mfma_f32_16x16x32_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8)
-v_mfma_f32_16x16x32_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8)
-v_mfma_f32_16x16x32_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8)
-v_mfma_f32_16x16x32_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8)
-v_mfma_f32_32x32x16_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8)
-v_mfma_f32_32x32x16_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8)
-v_mfma_f32_32x32x16_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8)
-v_mfma_f32_32x32x16_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8)
-v_smfmac_f32_16x16x64_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8)
-v_smfmac_f32_16x16x64_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8)
-v_smfmac_f32_16x16x64_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8)
-v_smfmac_f32_16x16x64_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8)
-v_smfmac_f32_32x32x32_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8)
-v_smfmac_f32_32x32x32_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8)
-v_smfmac_f32_32x32x32_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8)
-v_smfmac_f32_32x32x32_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8)
+v_smfmac_f32_16x16x64_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_F16)
+v_smfmac_f32_32x32x32_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_F16)
+v_mfma_f32_32x32x4_2b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16)
+v_mfma_f32_16x16x4_4b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16)
+v_mfma_f32_4x4x4_16b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16)
+v_mfma_f32_32x32x8_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X8_BF16)
+v_mfma_f32_16x16x16_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X16_BF16)
+v_smfmac_f32_16x16x32_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X32_F16)
+v_smfmac_f32_32x32x16_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X16_F16)
+v_smfmac_f32_16x16x32_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X32_BF16)
+v_smfmac_f32_32x32x16_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X16_BF16)
+v_smfmac_i32_16x16x64_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_16X16X64_I8)
+v_smfmac_i32_32x32x32_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_32X32X32_I8)
+v_mfma_f64_16x16x4_f64 = functools.partial(VOP3P, VOP3POp.V_MFMA_F64_16X16X4_F64)
+v_mfma_f64_4x4x4_4b_f64 = functools.partial(VOP3P, VOP3POp.V_MFMA_F64_4X4X4_4B_F64)
+v_mfma_f32_16x16x32_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8)
+v_mfma_f32_16x16x32_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8)
+v_mfma_f32_16x16x32_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8)
+v_mfma_f32_16x16x32_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8)
+v_mfma_f32_32x32x16_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8)
+v_mfma_f32_32x32x16_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8)
+v_mfma_f32_32x32x16_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8)
+v_mfma_f32_32x32x16_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8)
+v_smfmac_f32_16x16x64_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8)
+v_smfmac_f32_16x16x64_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8)
+v_smfmac_f32_16x16x64_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8)
+v_smfmac_f32_16x16x64_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8)
+v_smfmac_f32_32x32x32_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8)
+v_smfmac_f32_32x32x32_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8)
+v_smfmac_f32_32x32x32_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8)
+v_smfmac_f32_32x32x32_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8)
 v_mfma_scale_f32_16x16x128_f8f6f4 = functools.partial(VOP3PX2, VOP3PX2Op.V_MFMA_SCALE_F32_16X16X128_F8F6F4)
 v_mfma_scale_f32_32x32x64_f8f6f4 = functools.partial(VOP3PX2, VOP3PX2Op.V_MFMA_SCALE_F32_32X32X64_F8F6F4)
 v_add_co_u32 = functools.partial(VOP3SD, VOP3SDOp.V_ADD_CO_U32)

extra/assembly/amd/autogen/cdna/operands.py

@@ -1,6 +1,6 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import Fmt, OpType
-from tinygrad.runtime.autogen.amd.cdna.enum import DSOp, FLATOp, GLOBALOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3PX2Op, VOP3SDOp, VOPCOp
+from extra.assembly.amd.autogen.common import Fmt, OpType
+from extra.assembly.amd.autogen.cdna.enum import *
 
 # instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}
 OPERANDS = {

extra/assembly/amd/autogen/cdna/str_pcode.py

@@ -1,6 +1,6 @@
 # autogenerated from AMD ISA PDF - do not edit
 # ruff: noqa: E501
-from tinygrad.runtime.autogen.amd.cdna.enum import DSOp, FLATOp, GLOBALOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp
+from extra.assembly.amd.autogen.cdna.enum import DSOp, FLATOp, GLOBALOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp
 
 PCODE = {
   DSOp.DS_ADD_U32: 'addr = CalcDsAddr(ADDR.b32, OFFSET0.b32, OFFSET1.b32);\ntmp = MEM[addr].u32;\nMEM[addr].u32 += DATA.u32;\nRETURN_DATA.u32 = tmp',

extra/assembly/amd/autogen/common.py

@@ -1,10 +1,6 @@
 # autogenerated from AMD ISA XML - do not edit
 from enum import Enum, auto
 
-class ReprEnum(Enum):
-  """Enum with clean repr that roundtrips with eval()."""
-  def __repr__(self): return f"{type(self).__name__}.{self.name}"
-
 class Fmt(Enum):
   FMT_ANY = auto()
   FMT_BUF = auto()

extra/assembly/amd/autogen/rdna3/enum.py

@@ -1,7 +1,8 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401
+from enum import Enum
+from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType  # noqa: F401
 
-class DSOp(ReprEnum):
+class DSOp(Enum):
   DS_ADD_U32 = 0
   DS_SUB_U32 = 1
   DS_RSUB_U32 = 2
@@ -129,10 +130,10 @@ class DSOp(ReprEnum):
   DS_LOAD_B96 = 254
   DS_LOAD_B128 = 255
 
-class EXPOp(ReprEnum):
+class EXPOp(Enum):
   EXP = 0
 
-class FLATOp(ReprEnum):
+class FLATOp(Enum):
   FLAT_LOAD_U8 = 16
   FLAT_LOAD_I8 = 17
   FLAT_LOAD_U16 = 18
@@ -190,7 +191,7 @@ class FLATOp(ReprEnum):
   FLAT_ATOMIC_MAX_F32 = 82
   FLAT_ATOMIC_ADD_F32 = 86
 
-class GLOBALOp(ReprEnum):
+class GLOBALOp(Enum):
   GLOBAL_LOAD_U8 = 16
   GLOBAL_LOAD_I8 = 17
   GLOBAL_LOAD_U16 = 18
@@ -253,7 +254,7 @@ class GLOBALOp(ReprEnum):
   GLOBAL_ATOMIC_MAX_F32 = 82
   GLOBAL_ATOMIC_ADD_F32 = 86
 
-class HWREG(ReprEnum):
+class HWREG(Enum):
   HW_REG_MODE = 1
   HW_REG_STATUS = 2
   HW_REG_TRAPSTS = 3
@@ -279,11 +280,11 @@ class HWREG(ReprEnum):
   HW_REG_IB_STS2 = 28
   HW_REG_SHADER_CYCLES = 29
 
-class LDSDIROp(ReprEnum):
+class LDSDIROp(Enum):
   LDS_PARAM_LOAD = 0
   LDS_DIRECT_LOAD = 1
 
-class MIMGOp(ReprEnum):
+class MIMGOp(Enum):
   IMAGE_LOAD = 0
   IMAGE_LOAD_MIP = 1
   IMAGE_LOAD_PCK = 2
@@ -369,7 +370,7 @@ class MIMGOp(ReprEnum):
   IMAGE_GATHER4_C_B_CL = 101
   IMAGE_GATHER4H = 144
 
-class MSG(ReprEnum):
+class MSG(Enum):
   MSG_RTN_GET_DOORBELL = 128
   MSG_RTN_GET_DDID = 129
   MSG_RTN_GET_TMA = 130
@@ -379,7 +380,7 @@ class MSG(ReprEnum):
   MSG_RTN_GET_TBA_TO_PC = 134
   MSG_RTN_ILLEGAL_MSG = 255
 
-class MTBUFOp(ReprEnum):
+class MTBUFOp(Enum):
   TBUFFER_LOAD_FORMAT_X = 0
   TBUFFER_LOAD_FORMAT_XY = 1
   TBUFFER_LOAD_FORMAT_XYZ = 2
@@ -397,7 +398,7 @@ class MTBUFOp(ReprEnum):
   TBUFFER_STORE_D16_FORMAT_XYZ = 14
   TBUFFER_STORE_D16_FORMAT_XYZW = 15
 
-class MUBUFOp(ReprEnum):
+class MUBUFOp(Enum):
   BUFFER_LOAD_FORMAT_X = 0
   BUFFER_LOAD_FORMAT_XY = 1
   BUFFER_LOAD_FORMAT_XYZ = 2
@@ -478,7 +479,7 @@ class MUBUFOp(ReprEnum):
   BUFFER_ATOMIC_MAX_F32 = 82
   BUFFER_ATOMIC_ADD_F32 = 86
 
-class SCRATCHOp(ReprEnum):
+class SCRATCHOp(Enum):
   SCRATCH_LOAD_U8 = 16
   SCRATCH_LOAD_I8 = 17
   SCRATCH_LOAD_U16 = 18
@@ -507,7 +508,7 @@ class SCRATCHOp(ReprEnum):
   SCRATCH_LOAD_LDS_I16 = 48
   SCRATCH_LOAD_LDS_B32 = 49
 
-class SMEMOp(ReprEnum):
+class SMEMOp(Enum):
   S_LOAD_B32 = 0
   S_LOAD_B64 = 1
   S_LOAD_B128 = 2
@@ -523,7 +524,7 @@ class SMEMOp(ReprEnum):
   S_ATC_PROBE = 34
   S_ATC_PROBE_BUFFER = 35
 
-class SOP1Op(ReprEnum):
+class SOP1Op(Enum):
   S_MOV_B32 = 0
   S_MOV_B64 = 1
   S_CMOV_B32 = 2
@@ -605,7 +606,7 @@ class SOP1Op(ReprEnum):
   S_TRUNC_F16 = 109
   S_RNDNE_F16 = 110
 
-class SOP2Op(ReprEnum):
+class SOP2Op(Enum):
   S_ADD_U32 = 0
   S_SUB_U32 = 1
   S_ADD_I32 = 2
@@ -674,7 +675,7 @@ class SOP2Op(ReprEnum):
   S_MUL_F16 = 77
   S_FMAC_F16 = 78
 
-class SOPCOp(ReprEnum):
+class SOPCOp(Enum):
   S_CMP_EQ_I32 = 0
   S_CMP_LG_I32 = 1
   S_CMP_GT_I32 = 2
@@ -722,7 +723,7 @@ class SOPCOp(ReprEnum):
   S_CMP_NEQ_F16 = 93
   S_CMP_NLT_F16 = 94
 
-class SOPKOp(ReprEnum):
+class SOPKOp(Enum):
   S_MOVK_I32 = 0
   S_VERSION = 1
   S_CMOVK_I32 = 2
@@ -751,7 +752,7 @@ class SOPKOp(ReprEnum):
   S_WAITCNT_EXPCNT = 26
   S_WAITCNT_LGKMCNT = 27
 
-class SOPPOp(ReprEnum):
+class SOPPOp(Enum):
   S_NOP = 0
   S_SETKILL = 1
   S_SETHALT = 2
@@ -792,7 +793,7 @@ class SOPPOp(ReprEnum):
   S_ICACHE_INV = 60
   S_BARRIER = 61
 
-class VINTERPOp(ReprEnum):
+class VINTERPOp(Enum):
   V_INTERP_P10_F32 = 0
   V_INTERP_P2_F32 = 1
   V_INTERP_P10_F16_F32 = 2
@@ -800,7 +801,7 @@ class VINTERPOp(ReprEnum):
   V_INTERP_P10_RTZ_F16_F32 = 4
   V_INTERP_P2_RTZ_F16_F32 = 5
 
-class VOP1Op(ReprEnum):
+class VOP1Op(Enum):
   V_NOP_E32 = 0
   V_MOV_B32_E32 = 1
   V_READFIRSTLANE_B32_E32 = 2
@@ -974,7 +975,7 @@ class VOP1Op(ReprEnum):
   V_CVT_I32_I16 = V_CVT_I32_I16_E32
   V_CVT_U32_U16 = V_CVT_U32_U16_E32
 
-class VOP2Op(ReprEnum):
+class VOP2Op(Enum):
   V_CNDMASK_B32_E32 = 1
   V_DOT2ACC_F32_F16_E32 = 2
   V_ADD_F32_E32 = 3
@@ -1068,7 +1069,7 @@ class VOP2Op(ReprEnum):
   V_LDEXP_F16 = V_LDEXP_F16_E32
   V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
 
-class VOP3Op(ReprEnum):
+class VOP3Op(Enum):
   V_CMP_F_F16_E64 = 0
   V_CMP_LT_F16_E64 = 1
   V_CMP_EQ_F16_E64 = 2
@@ -1808,7 +1809,7 @@ class VOP3Op(ReprEnum):
   V_CVT_I32_I16 = V_CVT_I32_I16_E64
   V_CVT_U32_U16 = V_CVT_U32_U16_E64
 
-class VOP3POp(ReprEnum):
+class VOP3POp(Enum):
   V_PK_MAD_I16 = 0
   V_PK_MUL_LO_U16 = 1
   V_PK_ADD_I16 = 2
@@ -1844,7 +1845,7 @@ class VOP3POp(ReprEnum):
   V_WMMA_I32_16X16X16_IU8 = 68
   V_WMMA_I32_16X16X16_IU4 = 69
 
-class VOP3SDOp(ReprEnum):
+class VOP3SDOp(Enum):
   V_ADD_CO_CI_U32 = 288
   V_SUB_CO_CI_U32 = 289
   V_SUBREV_CO_CI_U32 = 290
@@ -1856,7 +1857,7 @@ class VOP3SDOp(ReprEnum):
   V_SUB_CO_U32 = 769
   V_SUBREV_CO_U32 = 770
 
-class VOPCOp(ReprEnum):
+class VOPCOp(Enum):
   V_CMP_F_F16_E32 = 0
   V_CMP_LT_F16_E32 = 1
   V_CMP_EQ_F16_E32 = 2
@@ -2238,7 +2239,7 @@ class VOPCOp(ReprEnum):
   V_CMPX_CLASS_F32 = V_CMPX_CLASS_F32_E32
   V_CMPX_CLASS_F64 = V_CMPX_CLASS_F64_E32
 
-class VOPDOp(ReprEnum):
+class VOPDOp(Enum):
   V_DUAL_FMAC_F32 = 0
   V_DUAL_FMAAK_F32 = 1
   V_DUAL_FMAMK_F32 = 2

extra/assembly/amd/autogen/rdna3/ins.py

@@ -1,7 +1,7 @@
 # autogenerated from AMD ISA XML - do not edit
-# ruff: noqa: E501,F401
-from tinygrad.renderer.amd.dsl import BitField, DPP, DPP16, EXEC, EXECZ, EXEC_HI, EXEC_LO, EnumBitField, FixedBitField, INV_2PI, Inst, LIT, M0, NULL, OFF, SBaseField, SCC, SDWA, SGPRField, SRC_LDS_DIRECT, SRsrcField, SSrcField, SrcField, VCC, VCCZ, VCC_HI, VCC_LO, VDSTYField, VGPRField, s, src, ttmp, v
-from tinygrad.runtime.autogen.amd.rdna3.enum import DSOp, EXPOp, FLATOp, GLOBALOp, LDSDIROp, MIMGOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, HWREG, MSG
+# ruff: noqa: F401,F403
+from extra.assembly.amd.dsl import *
+from extra.assembly.amd.autogen.rdna3.enum import *
 import functools
 
 class DS(Inst):
@@ -593,6 +593,9 @@ flat_load_d16_hi_i8 = functools.partial(FLAT, FLATOp.FLAT_LOAD_D16_HI_I8)
 flat_load_d16_hi_b16 = functools.partial(FLAT, FLATOp.FLAT_LOAD_D16_HI_B16)
 flat_store_d16_hi_b8 = functools.partial(FLAT, FLATOp.FLAT_STORE_D16_HI_B8)
 flat_store_d16_hi_b16 = functools.partial(FLAT, FLATOp.FLAT_STORE_D16_HI_B16)
+global_load_addtid_b32 = functools.partial(FLAT, FLATOp.GLOBAL_LOAD_ADDTID_B32)
+global_store_addtid_b32 = functools.partial(FLAT, FLATOp.GLOBAL_STORE_ADDTID_B32)
+global_load_lds_addtid_b32 = functools.partial(FLAT, FLATOp.GLOBAL_LOAD_LDS_ADDTID_B32)
 flat_atomic_swap_b32 = functools.partial(FLAT, FLATOp.FLAT_ATOMIC_SWAP_B32)
 flat_atomic_cmpswap_b32 = functools.partial(FLAT, FLATOp.FLAT_ATOMIC_CMPSWAP_B32)
 flat_atomic_add_u32 = functools.partial(FLAT, FLATOp.FLAT_ATOMIC_ADD_U32)

extra/assembly/amd/autogen/rdna3/operands.py

@@ -1,6 +1,6 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import Fmt, OpType
-from tinygrad.runtime.autogen.amd.rdna3.enum import DSOp, EXPOp, FLATOp, GLOBALOp, LDSDIROp, MIMGOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp
+from extra.assembly.amd.autogen.common import Fmt, OpType
+from extra.assembly.amd.autogen.rdna3.enum import *
 
 # instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}
 OPERANDS = {

extra/assembly/amd/autogen/rdna3/str_pcode.py

@@ -1,6 +1,6 @@
 # autogenerated from AMD ISA PDF - do not edit
 # ruff: noqa: E501
-from tinygrad.runtime.autogen.amd.rdna3.enum import DSOp, FLATOp, GLOBALOp, LDSDIROp, MIMGOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp
+from extra.assembly.amd.autogen.rdna3.enum import DSOp, FLATOp, GLOBALOp, LDSDIROp, MIMGOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp
 
 PCODE = {
   DSOp.DS_ADD_U32: 'tmp = MEM[ADDR].u32;\nMEM[ADDR].u32 += DATA.u32;\nRETURN_DATA.u32 = tmp',

extra/assembly/amd/autogen/rdna4/enum.py

@@ -1,7 +1,8 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401
+from enum import Enum
+from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType  # noqa: F401
 
-class DSOp(ReprEnum):
+class DSOp(Enum):
   DS_ADD_U32 = 0
   DS_SUB_U32 = 1
   DS_RSUB_U32 = 2
@@ -126,7 +127,7 @@ class DSOp(ReprEnum):
   DS_LOAD_B96 = 254
   DS_LOAD_B128 = 255
 
-class HWREG(ReprEnum):
+class HWREG(Enum):
   HW_REG_WAVE_MODE = 1
   HW_REG_WAVE_STATUS = 2
   HW_REG_WAVE_STATE_PRIV = 4
@@ -147,7 +148,7 @@ class HWREG(ReprEnum):
   HW_REG_SHADER_CYCLES_LO = 29
   HW_REG_SHADER_CYCLES_HI = 30
 
-class MSG(ReprEnum):
+class MSG(Enum):
   MSG_RTN_GET_DOORBELL = 128
   MSG_RTN_GET_DDID = 129
   MSG_RTN_GET_TMA = 130
@@ -158,7 +159,7 @@ class MSG(ReprEnum):
   MSG_RTN_GET_SE_HW_ID = 135
   MSG_RTN_ILLEGAL_MSG = 255
 
-class SMEMOp(ReprEnum):
+class SMEMOp(Enum):
   S_LOAD_B32 = 0
   S_LOAD_B64 = 1
   S_LOAD_B128 = 2
@@ -188,7 +189,7 @@ class SMEMOp(ReprEnum):
   S_BUFFER_PREFETCH_DATA = 39
   S_PREFETCH_DATA_PC_REL = 40
 
-class SOP1Op(ReprEnum):
+class SOP1Op(Enum):
   S_MOV_B32 = 0
   S_MOV_B64 = 1
   S_CMOV_B32 = 2
@@ -277,7 +278,7 @@ class SOP1Op(ReprEnum):
   S_TRUNC_F16 = 109
   S_RNDNE_F16 = 110
 
-class SOP2Op(ReprEnum):
+class SOP2Op(Enum):
   S_ADD_CO_U32 = 0
   S_SUB_CO_U32 = 1
   S_ADD_CO_I32 = 2
@@ -353,7 +354,7 @@ class SOP2Op(ReprEnum):
   S_SUB_NC_U64 = 84
   S_MUL_U64 = 85
 
-class SOPCOp(ReprEnum):
+class SOPCOp(Enum):
   S_CMP_EQ_I32 = 0
   S_CMP_LG_I32 = 1
   S_CMP_GT_I32 = 2
@@ -401,7 +402,7 @@ class SOPCOp(ReprEnum):
   S_CMP_NEQ_F16 = 93
   S_CMP_NLT_F16 = 94
 
-class SOPKOp(ReprEnum):
+class SOPKOp(Enum):
   S_MOVK_I32 = 0
   S_VERSION = 1
   S_CMOVK_I32 = 2
@@ -412,7 +413,7 @@ class SOPKOp(ReprEnum):
   S_SETREG_IMM32_B32 = 19
   S_CALL_B64 = 20
 
-class SOPPOp(ReprEnum):
+class SOPPOp(Enum):
   S_NOP = 0
   S_SETKILL = 1
   S_SETHALT = 2
@@ -457,7 +458,7 @@ class SOPPOp(ReprEnum):
   S_WAIT_LOADCNT_DSCNT = 72
   S_WAIT_STORECNT_DSCNT = 73
 
-class VBUFFEROp(ReprEnum):
+class VBUFFEROp(Enum):
   BUFFER_LOAD_FORMAT_X = 0
   BUFFER_LOAD_FORMAT_XY = 1
   BUFFER_LOAD_FORMAT_XYZ = 2
@@ -548,14 +549,14 @@ class VBUFFEROp(ReprEnum):
   TBUFFER_STORE_D16_FORMAT_XYZ = 142
   TBUFFER_STORE_D16_FORMAT_XYZW = 143
 
-class VDSDIROp(ReprEnum):
+class VDSDIROp(Enum):
   DS_PARAM_LOAD = 0
   DS_DIRECT_LOAD = 1
 
-class VEXPORTOp(ReprEnum):
+class VEXPORTOp(Enum):
   EXPORT = 0
 
-class VFLATOp(ReprEnum):
+class VFLATOp(Enum):
   FLAT_LOAD_U8 = 16
   FLAT_LOAD_I8 = 17
   FLAT_LOAD_U16 = 18
@@ -614,7 +615,7 @@ class VFLATOp(ReprEnum):
   FLAT_ATOMIC_PK_ADD_F16 = 89
   FLAT_ATOMIC_PK_ADD_BF16 = 90
 
-class VGLOBALOp(ReprEnum):
+class VGLOBALOp(Enum):
   GLOBAL_LOAD_U8 = 16
   GLOBAL_LOAD_I8 = 17
   GLOBAL_LOAD_U16 = 18
@@ -681,7 +682,7 @@ class VGLOBALOp(ReprEnum):
   GLOBAL_ATOMIC_PK_ADD_BF16 = 90
   GLOBAL_ATOMIC_ORDERED_ADD_B64 = 115
 
-class VIMAGEOp(ReprEnum):
+class VIMAGEOp(Enum):
   IMAGE_LOAD = 0
   IMAGE_LOAD_MIP = 1
   IMAGE_LOAD_PCK = 2
@@ -716,7 +717,7 @@ class VIMAGEOp(ReprEnum):
   IMAGE_ATOMIC_PK_ADD_F16 = 134
   IMAGE_ATOMIC_PK_ADD_BF16 = 135
 
-class VINTERPOp(ReprEnum):
+class VINTERPOp(Enum):
   V_INTERP_P10_F32 = 0
   V_INTERP_P2_F32 = 1
   V_INTERP_P10_F16_F32 = 2
@@ -724,7 +725,7 @@ class VINTERPOp(ReprEnum):
   V_INTERP_P10_RTZ_F16_F32 = 4
   V_INTERP_P2_RTZ_F16_F32 = 5
 
-class VOP1Op(ReprEnum):
+class VOP1Op(Enum):
   V_NOP_E32 = 0
   V_MOV_B32_E32 = 1
   V_READFIRSTLANE_B32_E32 = 2
@@ -906,7 +907,7 @@ class VOP1Op(ReprEnum):
   V_CVT_PK_F32_FP8 = V_CVT_PK_F32_FP8_E32
   V_CVT_PK_F32_BF8 = V_CVT_PK_F32_BF8_E32
 
-class VOP2Op(ReprEnum):
+class VOP2Op(Enum):
   V_CNDMASK_B32_E32 = 1
   V_ADD_F64_E32 = 2
   V_ADD_F32_E32 = 3
@@ -1006,7 +1007,7 @@ class VOP2Op(ReprEnum):
   V_LDEXP_F16 = V_LDEXP_F16_E32
   V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
 
-class VOP3Op(ReprEnum):
+class VOP3Op(Enum):
   V_CMP_LT_F16_E64 = 1
   V_CMP_EQ_F16_E64 = 2
   V_CMP_LE_F16_E64 = 3
@@ -1731,7 +1732,7 @@ class VOP3Op(ReprEnum):
   V_CVT_PK_F32_FP8 = V_CVT_PK_F32_FP8_E64
   V_CVT_PK_F32_BF8 = V_CVT_PK_F32_BF8_E64
 
-class VOP3POp(ReprEnum):
+class VOP3POp(Enum):
   V_PK_MAD_I16 = 0
   V_PK_MUL_LO_U16 = 1
   V_PK_ADD_I16 = 2
@@ -1789,7 +1790,7 @@ class VOP3POp(ReprEnum):
   V_SWMMAC_F32_16X16X32_BF8_FP8 = 89
   V_SWMMAC_F32_16X16X32_BF8_BF8 = 90
 
-class VOP3SDOp(ReprEnum):
+class VOP3SDOp(Enum):
   V_ADD_CO_CI_U32 = 288
   V_SUB_CO_CI_U32 = 289
   V_SUBREV_CO_CI_U32 = 290
@@ -1801,7 +1802,7 @@ class VOP3SDOp(ReprEnum):
   V_SUB_CO_U32 = 769
   V_SUBREV_CO_U32 = 770
 
-class VOPCOp(ReprEnum):
+class VOPCOp(Enum):
   V_CMP_LT_F16_E32 = 1
   V_CMP_EQ_F16_E32 = 2
   V_CMP_LE_F16_E32 = 3
@@ -2127,7 +2128,7 @@ class VOPCOp(ReprEnum):
   V_CMPX_CLASS_F32 = V_CMPX_CLASS_F32_E32
   V_CMPX_CLASS_F64 = V_CMPX_CLASS_F64_E32
 
-class VOPDOp(ReprEnum):
+class VOPDOp(Enum):
   V_DUAL_FMAC_F32 = 0
   V_DUAL_FMAAK_F32 = 1
   V_DUAL_FMAMK_F32 = 2
@@ -2146,7 +2147,7 @@ class VOPDOp(ReprEnum):
   V_DUAL_LSHLREV_B32 = 17
   V_DUAL_AND_B32 = 18
 
-class VSAMPLEOp(ReprEnum):
+class VSAMPLEOp(Enum):
   IMAGE_MSAA_LOAD = 24
   IMAGE_SAMPLE = 27
   IMAGE_SAMPLE_D = 28
@@ -2206,7 +2207,7 @@ class VSAMPLEOp(ReprEnum):
   IMAGE_GATHER4_C_B_CL = 101
   IMAGE_GATHER4H = 144
 
-class VSCRATCHOp(ReprEnum):
+class VSCRATCHOp(Enum):
   SCRATCH_LOAD_U8 = 16
   SCRATCH_LOAD_I8 = 17
   SCRATCH_LOAD_U16 = 18

extra/assembly/amd/autogen/rdna4/ins.py

@@ -1,7 +1,7 @@
 # autogenerated from AMD ISA XML - do not edit
-# ruff: noqa: E501,F401
-from tinygrad.renderer.amd.dsl import BitField, DPP, DPP16, EXEC, EXECZ, EXEC_HI, EXEC_LO, EnumBitField, FixedBitField, INV_2PI, Inst, LIT, M0, NULL, OFF, SBaseField, SCC, SDWA, SGPRField, SRC_LDS_DIRECT, SSrcField, SrcField, VCC, VCCZ, VCC_HI, VCC_LO, VDSTYField, VGPRField, s, src, ttmp, v
-from tinygrad.runtime.autogen.amd.rdna4.enum import DSOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VBUFFEROp, VDSDIROp, VEXPORTOp, VFLATOp, VGLOBALOp, VIMAGEOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, VSAMPLEOp, VSCRATCHOp, HWREG, MSG
+# ruff: noqa: F401,F403
+from extra.assembly.amd.dsl import *
+from extra.assembly.amd.autogen.rdna4.enum import *
 import functools
 
 class DS(Inst):
@@ -101,11 +101,11 @@ class VFLAT(Inst):
   sve = BitField(49, 49)
   scope = BitField(51, 50)
   th = BitField(54, 52)
-  vsrc = VGPRField(62, 55)
+  vsrc = BitField(62, 55)
   ioffset = BitField(95, 72)
 
 class VGLOBAL(Inst):
-  encoding = FixedBitField(31, 24, 0b11101110)
+  encoding = FixedBitField(31, 24, 0b11101100)
   op = EnumBitField(21, 14, VGLOBALOp, {VGLOBALOp.GLOBAL_LOAD_U8, VGLOBALOp.GLOBAL_LOAD_I8, VGLOBALOp.GLOBAL_LOAD_U16, VGLOBALOp.GLOBAL_LOAD_I16, VGLOBALOp.GLOBAL_LOAD_B32, VGLOBALOp.GLOBAL_LOAD_B64, VGLOBALOp.GLOBAL_LOAD_B96, VGLOBALOp.GLOBAL_LOAD_B128, VGLOBALOp.GLOBAL_STORE_B8, VGLOBALOp.GLOBAL_STORE_B16, VGLOBALOp.GLOBAL_STORE_B32, VGLOBALOp.GLOBAL_STORE_B64, VGLOBALOp.GLOBAL_STORE_B96, VGLOBALOp.GLOBAL_STORE_B128, VGLOBALOp.GLOBAL_LOAD_D16_U8, VGLOBALOp.GLOBAL_LOAD_D16_I8, VGLOBALOp.GLOBAL_LOAD_D16_B16, VGLOBALOp.GLOBAL_LOAD_D16_HI_U8, VGLOBALOp.GLOBAL_LOAD_D16_HI_I8, VGLOBALOp.GLOBAL_LOAD_D16_HI_B16, VGLOBALOp.GLOBAL_STORE_D16_HI_B8, VGLOBALOp.GLOBAL_STORE_D16_HI_B16, VGLOBALOp.GLOBAL_LOAD_ADDTID_B32, VGLOBALOp.GLOBAL_STORE_ADDTID_B32, VGLOBALOp.GLOBAL_INV, VGLOBALOp.GLOBAL_WB, VGLOBALOp.GLOBAL_ATOMIC_SWAP_B32, VGLOBALOp.GLOBAL_ATOMIC_CMPSWAP_B32, VGLOBALOp.GLOBAL_ATOMIC_ADD_U32, VGLOBALOp.GLOBAL_ATOMIC_SUB_U32, VGLOBALOp.GLOBAL_ATOMIC_SUB_CLAMP_U32, VGLOBALOp.GLOBAL_ATOMIC_MIN_I32, VGLOBALOp.GLOBAL_ATOMIC_MIN_U32, VGLOBALOp.GLOBAL_ATOMIC_MAX_I32, VGLOBALOp.GLOBAL_ATOMIC_MAX_U32, VGLOBALOp.GLOBAL_ATOMIC_AND_B32, VGLOBALOp.GLOBAL_ATOMIC_OR_B32, VGLOBALOp.GLOBAL_ATOMIC_XOR_B32, VGLOBALOp.GLOBAL_ATOMIC_INC_U32, VGLOBALOp.GLOBAL_ATOMIC_DEC_U32, VGLOBALOp.GLOBAL_ATOMIC_SWAP_B64, VGLOBALOp.GLOBAL_ATOMIC_CMPSWAP_B64, VGLOBALOp.GLOBAL_ATOMIC_ADD_U64, VGLOBALOp.GLOBAL_ATOMIC_SUB_U64, VGLOBALOp.GLOBAL_ATOMIC_MIN_I64, VGLOBALOp.GLOBAL_ATOMIC_MIN_U64, VGLOBALOp.GLOBAL_ATOMIC_MAX_I64, VGLOBALOp.GLOBAL_ATOMIC_MAX_U64, VGLOBALOp.GLOBAL_ATOMIC_AND_B64, VGLOBALOp.GLOBAL_ATOMIC_OR_B64, VGLOBALOp.GLOBAL_ATOMIC_XOR_B64, VGLOBALOp.GLOBAL_ATOMIC_INC_U64, VGLOBALOp.GLOBAL_ATOMIC_DEC_U64, VGLOBALOp.GLOBAL_WBINV, VGLOBALOp.GLOBAL_ATOMIC_COND_SUB_U32, VGLOBALOp.GLOBAL_ATOMIC_MIN_NUM_F32, VGLOBALOp.GLOBAL_ATOMIC_MAX_NUM_F32, VGLOBALOp.GLOBAL_LOAD_BLOCK, VGLOBALOp.GLOBAL_STORE_BLOCK, VGLOBALOp.GLOBAL_ATOMIC_ADD_F32, VGLOBALOp.GLOBAL_LOAD_TR_B128, VGLOBALOp.GLOBAL_LOAD_TR_B64, VGLOBALOp.GLOBAL_ATOMIC_PK_ADD_F16, VGLOBALOp.GLOBAL_ATOMIC_PK_ADD_BF16, VGLOBALOp.GLOBAL_ATOMIC_ORDERED_ADD_B64})
   vdst = VGPRField(39, 32)
   vaddr = VGPRField(71, 64)
@@ -114,7 +114,20 @@ class VGLOBAL(Inst):
   sve = BitField(49, 49)
   scope = BitField(51, 50)
   th = BitField(54, 52)
-  vsrc = VGPRField(62, 55)
+  vsrc = BitField(62, 55)
+  ioffset = BitField(95, 72)
+
+class VSCRATCH(Inst):
+  encoding = FixedBitField(31, 24, 0b11101100)
+  op = EnumBitField(21, 14, VSCRATCHOp, {VSCRATCHOp.SCRATCH_LOAD_U8, VSCRATCHOp.SCRATCH_LOAD_I8, VSCRATCHOp.SCRATCH_LOAD_U16, VSCRATCHOp.SCRATCH_LOAD_I16, VSCRATCHOp.SCRATCH_LOAD_B32, VSCRATCHOp.SCRATCH_LOAD_B64, VSCRATCHOp.SCRATCH_LOAD_B96, VSCRATCHOp.SCRATCH_LOAD_B128, VSCRATCHOp.SCRATCH_STORE_B8, VSCRATCHOp.SCRATCH_STORE_B16, VSCRATCHOp.SCRATCH_STORE_B32, VSCRATCHOp.SCRATCH_STORE_B64, VSCRATCHOp.SCRATCH_STORE_B96, VSCRATCHOp.SCRATCH_STORE_B128, VSCRATCHOp.SCRATCH_LOAD_D16_U8, VSCRATCHOp.SCRATCH_LOAD_D16_I8, VSCRATCHOp.SCRATCH_LOAD_D16_B16, VSCRATCHOp.SCRATCH_LOAD_D16_HI_U8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_I8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_B16, VSCRATCHOp.SCRATCH_STORE_D16_HI_B8, VSCRATCHOp.SCRATCH_STORE_D16_HI_B16, VSCRATCHOp.SCRATCH_LOAD_BLOCK, VSCRATCHOp.SCRATCH_STORE_BLOCK})
+  vdst = VGPRField(39, 32)
+  vaddr = VGPRField(71, 64)
+  saddr = SGPRField(6, 0, default=NULL)
+  nv = BitField(7, 7)
+  sve = BitField(49, 49)
+  scope = BitField(51, 50)
+  th = BitField(54, 52)
+  vsrc = BitField(62, 55)
   ioffset = BitField(95, 72)
 
 class VIMAGE(Inst):
@@ -240,19 +253,6 @@ class VSAMPLE(Inst):
   vaddr2 = BitField(87, 80)
   vaddr3 = BitField(95, 88)
 
-class VSCRATCH(Inst):
-  encoding = FixedBitField(31, 24, 0b11101101)
-  op = EnumBitField(21, 14, VSCRATCHOp, {VSCRATCHOp.SCRATCH_LOAD_U8, VSCRATCHOp.SCRATCH_LOAD_I8, VSCRATCHOp.SCRATCH_LOAD_U16, VSCRATCHOp.SCRATCH_LOAD_I16, VSCRATCHOp.SCRATCH_LOAD_B32, VSCRATCHOp.SCRATCH_LOAD_B64, VSCRATCHOp.SCRATCH_LOAD_B96, VSCRATCHOp.SCRATCH_LOAD_B128, VSCRATCHOp.SCRATCH_STORE_B8, VSCRATCHOp.SCRATCH_STORE_B16, VSCRATCHOp.SCRATCH_STORE_B32, VSCRATCHOp.SCRATCH_STORE_B64, VSCRATCHOp.SCRATCH_STORE_B96, VSCRATCHOp.SCRATCH_STORE_B128, VSCRATCHOp.SCRATCH_LOAD_D16_U8, VSCRATCHOp.SCRATCH_LOAD_D16_I8, VSCRATCHOp.SCRATCH_LOAD_D16_B16, VSCRATCHOp.SCRATCH_LOAD_D16_HI_U8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_I8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_B16, VSCRATCHOp.SCRATCH_STORE_D16_HI_B8, VSCRATCHOp.SCRATCH_STORE_D16_HI_B16, VSCRATCHOp.SCRATCH_LOAD_BLOCK, VSCRATCHOp.SCRATCH_STORE_BLOCK})
-  vdst = VGPRField(39, 32)
-  vaddr = VGPRField(71, 64)
-  saddr = SGPRField(6, 0, default=NULL)
-  nv = BitField(7, 7)
-  sve = BitField(49, 49)
-  scope = BitField(51, 50)
-  th = BitField(54, 52)
-  vsrc = VGPRField(62, 55)
-  ioffset = BitField(95, 72)
-
 class SOP1_LIT(SOP1):
   op = EnumBitField(15, 8, SOP1Op, {SOP1Op.S_MOV_B32, SOP1Op.S_MOV_B64, SOP1Op.S_CMOV_B32, SOP1Op.S_CMOV_B64, SOP1Op.S_BREV_B32, SOP1Op.S_BREV_B64, SOP1Op.S_CTZ_I32_B32, SOP1Op.S_CTZ_I32_B64, SOP1Op.S_CLZ_I32_U32, SOP1Op.S_CLZ_I32_U64, SOP1Op.S_CLS_I32, SOP1Op.S_CLS_I32_I64, SOP1Op.S_SEXT_I32_I8, SOP1Op.S_SEXT_I32_I16, SOP1Op.S_BITSET0_B32, SOP1Op.S_BITSET0_B64, SOP1Op.S_BITSET1_B32, SOP1Op.S_BITSET1_B64, SOP1Op.S_BITREPLICATE_B64_B32, SOP1Op.S_ABS_I32, SOP1Op.S_BCNT0_I32_B32, SOP1Op.S_BCNT0_I32_B64, SOP1Op.S_BCNT1_I32_B32, SOP1Op.S_BCNT1_I32_B64, SOP1Op.S_QUADMASK_B32, SOP1Op.S_QUADMASK_B64, SOP1Op.S_WQM_B32, SOP1Op.S_WQM_B64, SOP1Op.S_NOT_B32, SOP1Op.S_NOT_B64, SOP1Op.S_AND_SAVEEXEC_B32, SOP1Op.S_AND_SAVEEXEC_B64, SOP1Op.S_OR_SAVEEXEC_B32, SOP1Op.S_OR_SAVEEXEC_B64, SOP1Op.S_XOR_SAVEEXEC_B32, SOP1Op.S_XOR_SAVEEXEC_B64, SOP1Op.S_NAND_SAVEEXEC_B32, SOP1Op.S_NAND_SAVEEXEC_B64, SOP1Op.S_NOR_SAVEEXEC_B32, SOP1Op.S_NOR_SAVEEXEC_B64, SOP1Op.S_XNOR_SAVEEXEC_B32, SOP1Op.S_XNOR_SAVEEXEC_B64, SOP1Op.S_AND_NOT0_SAVEEXEC_B32, SOP1Op.S_AND_NOT0_SAVEEXEC_B64, SOP1Op.S_OR_NOT0_SAVEEXEC_B32, SOP1Op.S_OR_NOT0_SAVEEXEC_B64, SOP1Op.S_AND_NOT1_SAVEEXEC_B32, SOP1Op.S_AND_NOT1_SAVEEXEC_B64, SOP1Op.S_OR_NOT1_SAVEEXEC_B32, SOP1Op.S_OR_NOT1_SAVEEXEC_B64, SOP1Op.S_AND_NOT0_WREXEC_B32, SOP1Op.S_AND_NOT0_WREXEC_B64, SOP1Op.S_AND_NOT1_WREXEC_B32, SOP1Op.S_AND_NOT1_WREXEC_B64, SOP1Op.S_MOVRELS_B32, SOP1Op.S_MOVRELS_B64, SOP1Op.S_MOVRELD_B32, SOP1Op.S_MOVRELD_B64, SOP1Op.S_MOVRELSD_2_B32, SOP1Op.S_GETPC_B64, SOP1Op.S_SETPC_B64, SOP1Op.S_SWAPPC_B64, SOP1Op.S_RFE_B64, SOP1Op.S_SENDMSG_RTN_B32, SOP1Op.S_SENDMSG_RTN_B64, SOP1Op.S_BARRIER_SIGNAL, SOP1Op.S_BARRIER_SIGNAL_ISFIRST, SOP1Op.S_GET_BARRIER_STATE, SOP1Op.S_BARRIER_INIT, SOP1Op.S_BARRIER_JOIN, SOP1Op.S_ALLOC_VGPR, SOP1Op.S_SLEEP_VAR, SOP1Op.S_CEIL_F32, SOP1Op.S_FLOOR_F32, SOP1Op.S_TRUNC_F32, SOP1Op.S_RNDNE_F32, SOP1Op.S_CVT_F32_I32, SOP1Op.S_CVT_F32_U32, SOP1Op.S_CVT_I32_F32, SOP1Op.S_CVT_U32_F32, SOP1Op.S_CVT_F16_F32, SOP1Op.S_CVT_F32_F16, SOP1Op.S_CVT_HI_F32_F16, SOP1Op.S_CEIL_F16, SOP1Op.S_FLOOR_F16, SOP1Op.S_TRUNC_F16, SOP1Op.S_RNDNE_F16})
   literal = BitField(63, 32)
@@ -973,6 +973,8 @@ flat_load_d16_hi_i8 = functools.partial(VFLAT, VFLATOp.FLAT_LOAD_D16_HI_I8)
 flat_load_d16_hi_b16 = functools.partial(VFLAT, VFLATOp.FLAT_LOAD_D16_HI_B16)
 flat_store_d16_hi_b8 = functools.partial(VFLAT, VFLATOp.FLAT_STORE_D16_HI_B8)
 flat_store_d16_hi_b16 = functools.partial(VFLAT, VFLATOp.FLAT_STORE_D16_HI_B16)
+global_load_addtid_b32 = functools.partial(VFLAT, VFLATOp.GLOBAL_LOAD_ADDTID_B32)
+global_store_addtid_b32 = functools.partial(VFLAT, VFLATOp.GLOBAL_STORE_ADDTID_B32)
 flat_atomic_swap_b32 = functools.partial(VFLAT, VFLATOp.FLAT_ATOMIC_SWAP_B32)
 flat_atomic_cmpswap_b32 = functools.partial(VFLAT, VFLATOp.FLAT_ATOMIC_CMPSWAP_B32)
 flat_atomic_add_u32 = functools.partial(VFLAT, VFLATOp.FLAT_ATOMIC_ADD_U32)

extra/assembly/amd/autogen/rdna4/operands.py

@@ -1,6 +1,6 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import Fmt, OpType
-from tinygrad.runtime.autogen.amd.rdna4.enum import DSOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VBUFFEROp, VDSDIROp, VEXPORTOp, VFLATOp, VGLOBALOp, VIMAGEOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, VSAMPLEOp, VSCRATCHOp
+from extra.assembly.amd.autogen.common import Fmt, OpType
+from extra.assembly.amd.autogen.rdna4.enum import *
 
 # instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}
 OPERANDS = {

extra/assembly/amd/autogen/rdna4/str_pcode.py

@@ -1,6 +1,6 @@
 # autogenerated from AMD ISA PDF - do not edit
 # ruff: noqa: E501
-from tinygrad.runtime.autogen.amd.rdna4.enum import DSOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VBUFFEROp, VFLATOp, VGLOBALOp, VIMAGEOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, VSAMPLEOp, VSCRATCHOp
+from extra.assembly.amd.autogen.rdna4.enum import DSOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VBUFFEROp, VFLATOp, VGLOBALOp, VIMAGEOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, VSAMPLEOp, VSCRATCHOp
 
 PCODE = {
   DSOp.DS_ADD_U32: 'addr = CalcDsAddr(vgpr_a.b32, offset.b32);\ntmp = MEM[addr].u32;\nMEM[addr].u32 += DATA.u32;\nRETURN_DATA.u32 = tmp',

extra/assembly/amd/decode.py

@@ -0,0 +1,65 @@
+# Instruction format detection and decoding
+from __future__ import annotations
+from extra.assembly.amd.dsl import Inst, FixedBitField, EnumBitField
+
+# SDWA/DPP variant detection: src0 field (bits 0-8) encodes the variant
+# 0xf9 (249) = SDWA, 0xfa (250) = DPP16 for CDNA (GFX9)
+_VARIANT_SRC0 = {"_SDWA_SDST": 0xf9, "_SDWA": 0xf9, "_DPP16": 0xfa}
+
+def _matches(data: bytes, cls: type[Inst]) -> bool:
+  """Check if data matches all FixedBitFields and op is in allowed."""
+  for _, field in cls._fields:
+    dword_idx = field.lo // 32
+    if len(data) < (dword_idx + 1) * 4: return False
+    word = int.from_bytes(data[dword_idx*4:(dword_idx+1)*4], 'little')
+    field_lo = field.lo % 32
+    if isinstance(field, FixedBitField):
+      if ((word >> field_lo) & field.mask) != field.default: return False
+    if isinstance(field, EnumBitField) and field.allowed is not None:
+      try: opcode = field.decode((word >> field_lo) & field.mask)
+      except ValueError: return False  # opcode not in enum
+      if opcode not in field.allowed: return False
+  # Check SDWA/DPP variant based on src0 field (bits 0-8) - only for variant classes
+  name = cls.__name__
+  word = int.from_bytes(data[:4], 'little')
+  for suffix, expected_src0 in _VARIANT_SRC0.items():
+    if name.endswith(suffix): return (word & 0x1ff) == expected_src0
+  return True
+
+# Import instruction classes for each architecture
+from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3SD, VOP3P, VOPC, VOPD, VINTERP,
+  SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH)
+from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT,
+  VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
+  VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
+  SOPC as R4_SOPC, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP,
+  SMEM as R4_SMEM, DS as R4_DS, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
+from extra.assembly.amd.autogen.cdna.ins import (VOP1 as C_VOP1, VOP1_SDWA as C_VOP1_SDWA, VOP1_DPP16 as C_VOP1_DPP16,
+  VOP2 as C_VOP2, VOP2_LIT as C_VOP2_LIT, VOP2_SDWA as C_VOP2_SDWA, VOP2_DPP16 as C_VOP2_DPP16,
+  VOPC as C_VOPC, VOPC_SDWA_SDST as C_VOPC_SDWA_SDST,
+  VOP3 as C_VOP3, VOP3_SDST as C_VOP3_SDST, VOP3SD as C_VOP3SD, VOP3P as C_VOP3P, VOP3PX2 as C_VOP3PX2,
+  SOP1 as C_SOP1, SOP2 as C_SOP2, SOPC as C_SOPC, SOPK as C_SOPK, SOPK_LIT as C_SOPK_LIT, SOPP as C_SOPP, SMEM as C_SMEM, DS as C_DS,
+  FLAT as C_FLAT, GLOBAL as C_GLOBAL, SCRATCH as C_SCRATCH, MUBUF as C_MUBUF)
+
+# Order matters: more specific encodings first, catch-alls (SOP2, VOP2) last
+# Order: base before _LIT (base matches regular ops, _LIT catches lit-only ops excluded from base)
+_FORMATS = {
+  "rdna3": [VOPD, VOP3P, VINTERP, VOP3SD, VOP3_SDST, VOP3, DS, GLOBAL, SCRATCH, FLAT, SMEM,
+            SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, VOPC, VOP1_SDST, VOP1, VOP1_LIT, VOP2, VOP2_LIT],
+  "rdna4": [R4_VOPD, R4_VOP3P, R4_VINTERP, R4_VOP3SD, R4_VOP3_SDST, R4_VOP3, R4_DS, R4_GLOBAL, R4_SCRATCH, R4_FLAT, R4_SMEM,
+            R4_SOP1, R4_SOPC, R4_SOPP, R4_SOPK, R4_SOPK_LIT, R4_VOPC, R4_VOP1_SDST, R4_VOP1, R4_SOP2, R4_SOP2_LIT, R4_VOP2, R4_VOP2_LIT],
+  "cdna": [C_VOP3PX2, C_VOP3P, C_VOP3SD, C_VOP3_SDST, C_VOP3, C_DS, C_GLOBAL, C_SCRATCH, C_FLAT, C_MUBUF, C_SMEM,
+            C_SOP1, C_SOPC, C_SOPP, C_SOPK, C_SOPK_LIT, C_VOPC_SDWA_SDST, C_VOPC,
+            C_VOP1_DPP16, C_VOP1_SDWA, C_VOP1, C_VOP2_DPP16, C_VOP2_SDWA, C_SOP2, C_VOP2, C_VOP2_LIT],
+}
+
+def detect_format(data: bytes, arch: str = "rdna3") -> type[Inst]:
+  """Detect instruction format from machine code bytes."""
+  assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
+  for cls in _FORMATS[arch]:
+    if _matches(data, cls): return cls
+  raise ValueError(f"unknown {arch} format word={int.from_bytes(data[:4], 'little'):#010x}")
+
+def decode_inst(data: bytes, arch: str = "rdna3") -> Inst:
+  """Decode machine code bytes into an instruction."""
+  return detect_format(data, arch).from_bytes(data)

extra/assembly/amd/disasm.py

@@ -1,16 +1,13 @@
 # RDNA3/RDNA4/CDNA disassembler
 from __future__ import annotations
-import re
-from typing import Callable
-from tinygrad.renderer.amd.dsl import Inst, Reg
+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'}
+                128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc'}
 SPECIAL_GPRS_CDNA = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'm0', 126: 'exec_lo', 127: 'exec_hi',
-                     128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0',
-                     246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc',
+                     128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc',
                      102: 'flat_scratch_lo', 103: 'flat_scratch_hi', 104: 'xnack_mask_lo', 105: 'xnack_mask_hi',
                      251: 'src_vccz', 252: 'src_execz'}
 SPECIAL_PAIRS = {106: 'vcc', 126: 'exec'}
@@ -72,29 +69,25 @@ def _num_srcs(inst) -> int:
   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
+  if any(x in n for x in ('_ADD3', '_LSHL_ADD', '_ADD_LSHL', '_LSHL_OR', '_AND_OR', 'OR3_B32', 'AND_OR_B32', 'ALIGNBIT', 'ALIGNBYTE', 'V_PERM_', 'XOR3', 'XAD', 'MULLIT', 'MINMAX', 'MAXMIN', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM', 'MINIMUM3', 'MAXIMUM3', 'MIN3', 'MAX3', 'DOT2', 'CVT_PK_U8_F32', 'DOT4', 'DOT8', 'WMMA', 'SWMMAC')): return 3
   return 2
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # IMPORTS
 # ═══════════════════════════════════════════════════════════════════════════════
 
-from tinygrad.runtime.autogen.amd.rdna3.ins import (VOP1, VOP1_SDST, VOP1_SDST_LIT, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3_SDST_LIT,
+from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP1_SDST_LIT, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3_SDST_LIT,
   VOP3_LIT, VOP3SD, VOP3SD_LIT, VOP3P, VOP3P_LIT, VOPC, VOPC_LIT, VOPD, VOPD_LIT, VINTERP, SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPC_LIT,
   SOPK, SOPK_LIT, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH, VOP2Op, VOPDOp, SOPPOp, HWREG, MSG)
-from tinygrad.runtime.autogen.amd.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST,
-  VOP1_SDST_LIT as R4_VOP1_SDST_LIT, VOP1_LIT as R4_VOP1_LIT,
+from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP1_SDST_LIT as R4_VOP1_SDST_LIT, VOP1_LIT as R4_VOP1_LIT,
   VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT, VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3_SDST_LIT as R4_VOP3_SDST_LIT, VOP3_LIT as R4_VOP3_LIT,
   VOP3SD as R4_VOP3SD, VOP3SD_LIT as R4_VOP3SD_LIT, VOP3P as R4_VOP3P, VOP3P_LIT as R4_VOP3P_LIT, VOPC as R4_VOPC, VOPC_LIT as R4_VOPC_LIT,
   VOPD as R4_VOPD, VOPD_LIT as R4_VOPD_LIT, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP1_LIT as R4_SOP1_LIT, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
   SOPC as R4_SOPC, SOPC_LIT as R4_SOPC_LIT, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP, SMEM as R4_SMEM, DS as R4_DS,
-  VOPDOp as R4_VOPDOp, HWREG as HWREG_RDNA4, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
-from tinygrad.runtime.autogen.amd.cdna.ins import HWREG as HWREG_CDNA
+  VOPDOp as R4_VOPDOp, HWREG as HWREG_RDNA4)
+from extra.assembly.amd.autogen.cdna.ins import FLAT as C_FLAT, HWREG as HWREG_CDNA
 
 def _is_cdna(inst: Inst) -> bool: return 'cdna' in inst.__class__.__module__
-def _is_r4(inst: Inst) -> bool: return 'rdna4' in inst.__class__.__module__
 
 # CDNA opcode name aliases for disasm (new name -> old name expected by tests)
 _CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32', 'v_fmamk_f32': 'v_madmk_f32', 'v_fmaak_f32': 'v_madak_f32'}
@@ -105,15 +98,9 @@ _CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v
 
 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 | None:
-  b = _unwrap(b)
-  return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
+def _vreg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("v", b - 256 if b >= 256 else b, n)
+def _areg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("a", b - 256 if b >= 256 else b, n)  # accumulator registers for GFX90a
+def _ttmp(b, n: int = 1) -> str: b = _unwrap(b); return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
 
 def _fmt_sdst(v, n: int = 1, cdna: bool = False) -> str:
   v = _unwrap(v)
@@ -141,9 +128,7 @@ def _fmt_v16(v, base: int = 256, hi_thresh: int = 384) -> str:
 
 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 _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))}]"
 
@@ -214,8 +199,7 @@ def _disasm_vop2(inst: VOP2) -> str:
   basename = name.replace('_e32', '')
   if cdna and basename in _VOP2_CARRY_OUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}"
   if cdna and basename in _VOP2_CARRY_INOUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
-  if not cdna and basename in _VOP2_CARRY_INOUT_RDNA:
-    return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
+  if not cdna and basename in _VOP2_CARRY_INOUT_RDNA: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
   sn0 = inst.canonical_op_regs.get('s0', 1)
   if inst.vdst.sz > 1 or sn0 > 1 or inst.vsrc1.sz > 1:
     src0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else _fmt_src(inst.src0, sn0, cdna)
@@ -231,10 +215,7 @@ def _disasm_vopc(inst: VOPC) -> str:
     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
-  if inst.src0.offset == 255: s0 = _lit(inst, inst.src0)
-  elif inst.src0.sz > 1: s0 = inst.src0.fmt()
-  elif is16: s0 = _src16(inst, inst.src0.offset)
-  else: s0 = _lit(inst, inst.src0)
+  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}"
@@ -244,7 +225,7 @@ NO_ARG_SOPP = {SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
 
 def _disasm_sopp(inst: SOPP) -> str:
   name, cdna = inst.op_name.lower(), _is_cdna(inst)
-  is_rdna4 = _is_r4(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'}
@@ -270,22 +251,21 @@ def _disasm_sopp(inst: SOPP) -> str:
     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']
+    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
-    def dep(v): return deps[v-1] if 0 < v <= len(deps) else str(v)
+    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} {inst.simm16}"
+  if name.startswith(('s_cbranch', 's_branch')): return f"{name} 0x{inst.simm16:x}"
   return f"{name} 0x{inst.simm16:x}"
 
 def _disasm_smem(inst: SMEM) -> str:
   name, cdna = inst.op_name.lower(), _is_cdna(inst)
   if name in ('s_gl1_inv', 's_dcache_inv', 's_dcache_inv_vol', 's_dcache_wb', 's_dcache_wb_vol', 's_icache_inv'): return name
   soe, imm = getattr(inst, 'soe', 0) or getattr(inst, 'soffset_en', 0), getattr(inst, 'imm', 1)
-  is_rdna4 = _is_r4(inst)
-  offset = inst.ioffset if is_rdna4 else getattr(inst, 'offset', 0)  # type: ignore[attr-defined]
+  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}"
@@ -296,9 +276,7 @@ def _disasm_smem(inst: SMEM) -> str:
   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
-  if sbase_count == 2: sbase_str = _fmt_src(sbase_idx, sbase_count, cdna)
-  elif sbase_idx <= 105: sbase_str = _sreg(sbase_idx, sbase_count)
-  else: sbase_str = _reg("ttmp", sbase_idx - 108, sbase_count)
+  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))
@@ -324,71 +302,54 @@ def _disasm_smem(inst: SMEM) -> str:
   if name in ('s_memrealtime', 's_memtime'): return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}"
   return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
 
-R4_TH_LOAD = {1: 'TH_LOAD_NT', 2: 'TH_LOAD_HT', 3: 'TH_LOAD_LU', 4: 'TH_LOAD_RT_WB', 5: 'TH_LOAD_NT_WB'}
-R4_TH_STORE = {1: 'TH_STORE_NT', 2: 'TH_STORE_HT', 3: 'TH_STORE_ST', 4: 'TH_STORE_RT_WB', 5: 'TH_STORE_NT_WB'}
-R4_TH_ATOMIC = {1: 'TH_ATOMIC_RETURN', 2: 'TH_ATOMIC_NT', 3: 'TH_ATOMIC_RETURN_NT',
-                4: 'TH_ATOMIC_CASCADE_RT', 5: 'TH_ATOMIC_CASCADE_RETURN', 6: 'TH_ATOMIC_CASCADE_NT', 7: 'TH_ATOMIC_CASCADE_RETURN_NT'}
-R4_SCOPE = {1: 'SCOPE_SE', 2: 'SCOPE_DEV', 3: 'SCOPE_SYS'}
-
 def _disasm_flat(inst: FLAT) -> str:
-  name, cdna, r4 = inst.op_name.lower(), _is_cdna(inst), _is_r4(inst)
+  name, cdna = inst.op_name.lower(), _is_cdna(inst)
   acc = getattr(inst, 'acc', 0)
   reg_fn = _areg if acc else _vreg
-  if r4: seg = 'flat' if (cls_name:=inst.__class__.__name__) == 'VFLAT' else ('global' if cls_name == 'VGLOBAL' else 'scratch')
-  else: seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
+  seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
   instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
-  # Global/scratch uses 13-bit signed offset (RDNA3/CDNA), 24-bit signed offset (RDNA4)
-  offset = inst.ioffset if r4 else inst.offset  # type: ignore[attr-defined]
-  if r4: off_val = offset if offset < (1 << 23) else offset - (1 << 24)  # sign extend 24-bit
-  elif seg != 'flat':
+  # Global/scratch uses 13-bit signed offset
+  if seg != 'flat':
     if cdna:
       # CDNA: bit 12 is sign bit but not in offset field
       raw = int.from_bytes(inst.to_bytes(), 'little')
-      off_val = offset | ((raw >> 12) & 1) << 12  # get bit 12
+      off_val = inst.offset | ((raw >> 12) & 1) << 12  # get bit 12
     else:
-      off_val = offset
+      off_val = inst.offset
     off_val = off_val if off_val < 4096 else off_val - 8192  # sign extend 13-bit
   else:
-    off_val = offset
+    off_val = inst.offset
   # Use get_field_bits: data for stores/atomics, d for loads
   regs = inst.canonical_op_regs
   w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
   off_s = f" offset:{off_val}" if off_val else ""
-  if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"  # type: ignore[attr-defined]
-  elif r4:
-    th_names = R4_TH_ATOMIC if 'atomic' in name else (R4_TH_STORE if 'store' in name else R4_TH_LOAD)
-    mods = off_s + (f" th:{th_names[inst.th]}" if inst.th in th_names else "") + (f" scope:{R4_SCOPE[inst.scope]}" if inst.scope in R4_SCOPE else "")
+  if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"
   else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
   if seg == 'flat': saddr_s = ""
   elif _unwrap(inst.saddr) in (0x7F, 124): saddr_s = ", off"
   elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr, cdna)}"
-  elif _unwrap(inst.saddr) in (SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS):
-    saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
+  elif _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.vsrc if r4 else inst.data) if 'store' in name else inst.vdst)}{saddr_s}{mods}"
-   # RDNA4: vaddr instead of addr, vsrc instead of data
-  addr = inst.vaddr if r4 else inst.addr  # type: ignore[attr-defined]
-  data = inst.vsrc if r4 else inst.data  # type: ignore[attr-defined]
+  if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
   # load_lds_* instructions: vaddr, saddr (no vdst, data goes to LDS)
   if 'load_lds' in name:
     addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
-    addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
+    addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
     return f"{instr} {addr_s}{saddr_s}{mods}"
   if seg == 'flat': addr_w = 2  # flat always uses 64-bit vaddr
   elif cdna: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
   else: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
-  addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
-  data_s, vdst_s = reg_fn(data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
+  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:
-    glc_or_sc0 = inst.sc0 if cdna else (inst.th & 1 if r4 else inst.glc)  # type: ignore[attr-defined]
-    sfx = f"{saddr_s if seg != 'flat' else ''}{mods}"
-    return f"{instr} {vdst_s}, {addr_s}, {data_s}{sfx}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{sfx}"
+    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:
-  name = inst.op_name.lower()
+  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 ""
@@ -417,8 +378,7 @@ def _disasm_ds(inst: DS) -> str:
   if 'write2' in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
   if 'read2' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}{off2}{gds}"
   if 'xchg2' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}, {d0}, {d1}{off2}{gds}"
-  if 'load' in name or ('read' in name and 'read2' not in name):
-    return f"{name} {reg_fn(inst.vdst)}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
+  if 'load' in name or ('read' in name and 'read2' not in name): return f"{name} {reg_fn(inst.vdst)}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
   if ('store' in name or 'write' in name) and not _has(name, 'cmp', 'xchg', 'write2'):
     return f"{name} {reg_fn(inst.data0)}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
   if 'swizzle' in name or name == 'ds_ordered_count': return f"{name} {reg_fn(inst.vdst)}, {addr}{off}{gds}"
@@ -429,15 +389,13 @@ def _disasm_ds(inst: DS) -> str:
   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:
-  name = inst.op_name.lower()
+  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_'):
-    s0v = _unwrap(inst.src0)
-    if s0v == 255: src = _lit(inst, inst.src0)
-    elif s0v == 253: src = "src_scc"
-    else: src = _fmt_src(inst.src0, max(1, bits['s0'] // 32))
+    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)
@@ -446,6 +404,7 @@ def _disasm_vop3(inst: VOP3) -> str:
 
   # 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)
@@ -461,8 +420,7 @@ def _disasm_vop3(inst: VOP3) -> str:
 
   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)))
+  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
@@ -500,19 +458,20 @@ def _disasm_vop3sd(inst: VOP3SD) -> str:
 
 def _disasm_vopd(inst: VOPD) -> str:
   lit = inst._literal
-  op_enum = R4_VOPDOp if _is_r4(inst) else VOPDOp
-  nx, ny = op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
+  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, inst.vdsty, inst.srcy0, inst.vsrcy1)}"
+  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_swmmac, n, is_fma_mix = 'swmmac' in name, inst.num_srcs() or 2, 'fma_mix' in name
+  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()
@@ -521,22 +480,18 @@ def _disasm_vop3p(inst: VOP3P) -> str:
   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 []))
+    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 []))
-  mod_s = ' ' + ' '.join(mods) if mods else ''
-  return f"{name} {dst}, {src0}, {src1}, {src2}{mod_s}" if n == 3 else f"{name} {dst}, {src0}, {src1}{mod_s}"
+    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:
-  name, cdna = inst.op_name.lower(), _is_cdna(inst)
+  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)
@@ -550,8 +505,8 @@ def _disasm_sop1(inst: SOP1) -> str:
       try: msg_str = MSG(v).name if v != 255 else None  # MSG_RTN_ILLEGAL_MSG (255) not supported by LLVM
       except ValueError: msg_str = None
       return f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, sendmsg({msg_str})" if msg_str else f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, 0x{v:x}"
-  sop1_src_only = ('S_ALLOC_VGPR', 'S_SLEEP_VAR', 'S_BARRIER_SIGNAL', 'S_BARRIER_SIGNAL_ISFIRST',
-                   'S_BARRIER_INIT', 'S_BARRIER_JOIN', 'S_SET_GPR_IDX_IDX', 'S_CBRANCH_JOIN')
+  sop1_src_only = ('S_ALLOC_VGPR', 'S_SLEEP_VAR', 'S_BARRIER_SIGNAL', 'S_BARRIER_SIGNAL_ISFIRST', 'S_BARRIER_INIT', 'S_BARRIER_JOIN', 'S_SET_GPR_IDX_IDX',
+                   'S_CBRANCH_JOIN')
   if inst.op_name in sop1_src_only: return f"{name} {src}"
   if cdna:
     if 'getpc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2, cdna)}"
@@ -589,8 +544,8 @@ _HWREG_BLACKLIST_CDNA = {'HW_REG_PC_LO', 'HW_REG_PC_HI', 'HW_REG_IB_DBG1', 'HW_R
                          'HW_REG_SQ_SHADER_TMA_LO', 'HW_REG_SQ_SHADER_TMA_HI', 'HW_REG_SQ_PERF_SNAPSHOT_DATA', 'HW_REG_SQ_PERF_SNAPSHOT_DATA1',
                          'HW_REG_SQ_PERF_SNAPSHOT_PC_LO', 'HW_REG_SQ_PERF_SNAPSHOT_PC_HI', 'HW_REG_XCC_ID'}
 def _disasm_sopk(inst: SOPK) -> str:
-  name, cdna = inst.op_name.lower(), _is_cdna(inst)
-  is_rdna4 = _is_r4(inst)
+  op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
+  is_rdna4 = 'rdna4' in inst.__class__.__module__
   hw = HWREG_CDNA if cdna else (HWREG_RDNA4 if is_rdna4 else HWREG)
   blacklist = _HWREG_BLACKLIST_CDNA if cdna else _HWREG_BLACKLIST
   def fmt_hwreg(hid, hoff, hsz):
@@ -612,14 +567,12 @@ def _disasm_sopk(inst: SOPK) -> str:
 
 def _disasm_vinterp(inst: VINTERP) -> str:
   mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
-  s0, s1, s2 = _lit(inst, inst.src0, inst.neg & 1), _lit(inst, inst.src1, inst.neg & 2), _lit(inst, inst.src2, inst.neg & 4)
-  return f"{inst.op_name.lower()} {inst.vdst.fmt()}, {s0}, {s1}, {s2}" + (" " + mods if mods else "")
+  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[..., str]] = {
+DISASM_HANDLERS: dict[type, callable] = {
   VOP1: _disasm_vop1, VOP1_SDST: _disasm_vop1, VOP1_SDST_LIT: _disasm_vop1, VOP1_LIT: _disasm_vop1,
   VOP2: _disasm_vop2, VOP2_LIT: _disasm_vop2, VOPC: _disasm_vopc, VOPC_LIT: _disasm_vopc,
-  VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3_SDST_LIT: _disasm_vop3, VOP3_LIT: _disasm_vop3,
-  VOP3SD: _disasm_vop3sd, VOP3SD_LIT: _disasm_vop3sd,
+  VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3_SDST_LIT: _disasm_vop3, VOP3_LIT: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOP3SD_LIT: _disasm_vop3sd,
   VOPD: _disasm_vopd, VOPD_LIT: _disasm_vopd, VOP3P: _disasm_vop3p, VOP3P_LIT: _disasm_vop3p,
   VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, GLOBAL: _disasm_flat, SCRATCH: _disasm_flat,
   SOP1: _disasm_sop1, SOP1_LIT: _disasm_sop1, SOP2: _disasm_sop2, SOP2_LIT: _disasm_sop2,
@@ -629,7 +582,6 @@ DISASM_HANDLERS: dict[type, Callable[..., str]] = {
   R4_VOP2: _disasm_vop2, R4_VOP2_LIT: _disasm_vop2, R4_VOPC: _disasm_vopc, R4_VOPC_LIT: _disasm_vopc,
   R4_VOP3: _disasm_vop3, R4_VOP3_SDST: _disasm_vop3, R4_VOP3_SDST_LIT: _disasm_vop3, R4_VOP3_LIT: _disasm_vop3,
   R4_VOP3SD: _disasm_vop3sd, R4_VOP3SD_LIT: _disasm_vop3sd, R4_VOP3P: _disasm_vop3p, R4_VOP3P_LIT: _disasm_vop3p,
-  R4_FLAT: _disasm_flat, R4_GLOBAL: _disasm_flat, R4_SCRATCH: _disasm_flat,
   R4_VOPD: _disasm_vopd, R4_VOPD_LIT: _disasm_vopd, R4_VINTERP: _disasm_vinterp, R4_SOPP: _disasm_sopp, R4_SMEM: _disasm_smem, R4_DS: _disasm_ds,
   R4_SOP1: _disasm_sop1, R4_SOP1_LIT: _disasm_sop1, R4_SOP2: _disasm_sop2, R4_SOP2_LIT: _disasm_sop2,
   R4_SOPC: _disasm_sopc, R4_SOPC_LIT: _disasm_sopc, R4_SOPK: _disasm_sopk, R4_SOPK_LIT: _disasm_sopk}
@@ -640,11 +592,11 @@ def disasm(inst: Inst) -> str: return DISASM_HANDLERS[type(inst)](inst)
 # CDNA DISASSEMBLER SUPPORT
 # ═══════════════════════════════════════════════════════════════════════════════
 
-from tinygrad.runtime.autogen.amd.cdna.ins import (VOP1 as CDNA_VOP1, VOP1_LIT as CDNA_VOP1_LIT,
+from extra.assembly.amd.autogen.cdna.ins import (VOP1 as CDNA_VOP1, VOP1_LIT as CDNA_VOP1_LIT,
   VOP1_SDWA as CDNA_VOP1_SDWA, VOP1_DPP16 as CDNA_VOP1_DPP16,
   VOP2 as CDNA_VOP2, VOP2_LIT as CDNA_VOP2_LIT, VOP2_SDWA as CDNA_VOP2_SDWA, VOP2_DPP16 as CDNA_VOP2_DPP16,
   VOPC as CDNA_VOPC, VOPC_LIT as CDNA_VOPC_LIT, VOPC_SDWA_SDST as CDNA_VOPC_SDWA_SDST,
-  VOP3 as CDNA_VOP3, VOP3_SDST as CDNA_VOP3_SDST, VOP3SD as CDNA_VOP3SD, VOP3P as CDNA_VOP3P, VOP3P_MFMA as CDNA_VOP3P_MFMA, VOP3PX2 as CDNA_VOP3PX2,
+  VOP3 as CDNA_VOP3, VOP3_SDST as CDNA_VOP3_SDST, VOP3SD as CDNA_VOP3SD, VOP3P as CDNA_VOP3P, VOP3PX2 as CDNA_VOP3PX2,
   SOP1 as CDNA_SOP1, SOP1_LIT as CDNA_SOP1_LIT, SOP2 as CDNA_SOP2, SOP2_LIT as CDNA_SOP2_LIT,
   SOPC as CDNA_SOPC, SOPC_LIT as CDNA_SOPC_LIT, SOPK as CDNA_SOPK, SOPK_LIT as CDNA_SOPK_LIT,
   SOPP as CDNA_SOPP, SMEM as CDNA_SMEM, DS as CDNA_DS,
@@ -674,9 +626,7 @@ def _disasm_vop3a(inst) -> str:
   else:
     regs = inst.canonical_op_regs
     dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
-    s0 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, r0)
-    s1 = _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, r1)
-    s2 = _cdna_src(inst, inst.src2, inst.neg&4, inst.abs&4, r2)
+    s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, r0), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, r1), _cdna_src(inst, inst.src2, inst.neg&4, inst.abs&4, r2)
     dst = _vreg(inst.vdst, dregs) if dregs > 1 else _vreg(inst.vdst)
   if op_val >= 512:
     return f"{name} {dst}, {s0}, {s1}, {s2}{opsel}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{opsel}{cl}{om}"
@@ -700,9 +650,7 @@ def _disasm_vop3b(inst) -> str:
   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 = _cdna_src(inst, inst.src0, inst.neg&1, n=r0)
-  s1 = _cdna_src(inst, inst.src1, inst.neg&2, n=r1)
-  s2 = _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
+  s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, n=r0), _cdna_src(inst, inst.src1, inst.neg&2, n=r1), _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
   # CDNA VOP3_SDST uses vdst field for sdst (but vdst adds 256), RDNA uses separate sdst field
   sdst_val = getattr(inst, 'sdst', None)
   if sdst_val is None and hasattr(inst, 'vdst'):
@@ -724,7 +672,7 @@ 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
-  def get_src(v, sc): return _lit(inst, v) if v == 255 else _fmt_src(v, sc, cdna=True)
+  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:
@@ -786,12 +734,9 @@ def _disasm_cdna_vop3p(inst) -> str:
   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 []))
-  mod_s = ' ' + ' '.join(mods) if mods else ''
-  return f"{name} {dst}, {src0}, {src1}, {src2}{mod_s}" if n == 3 else f"{name} {dst}, {src0}, {src1}{mod_s}"
+  mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else []) + \
+          ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
+  return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
 
 def _disasm_mubuf(inst) -> str:
   name = inst.op_name.lower()
@@ -950,6 +895,5 @@ DISASM_HANDLERS.update({CDNA_VOP1: _disasm_vop1, CDNA_VOP1_LIT: _disasm_vop1,
   CDNA_SOP1: _disasm_sop1, CDNA_SOP1_LIT: _disasm_sop1, CDNA_SOP2: _disasm_sop2, CDNA_SOP2_LIT: _disasm_sop2,
   CDNA_SOPC: _disasm_sopc, CDNA_SOPC_LIT: _disasm_sopc, CDNA_SOPK: _disasm_sopk, CDNA_SOPK_LIT: _disasm_sopk, CDNA_SOPP: _disasm_sopp,
   CDNA_SMEM: _disasm_smem, CDNA_DS: _disasm_ds, CDNA_FLAT: _disasm_flat, CDNA_GLOBAL: _disasm_flat, CDNA_SCRATCH: _disasm_flat,
-  CDNA_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b,
-  CDNA_VOP3P: _disasm_cdna_vop3p, CDNA_VOP3P_MFMA: _disasm_cdna_vop3p,
+  CDNA_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p,
   CDNA_MUBUF: _disasm_mubuf, CDNA_VOP3PX2: _disasm_vop3px2})

extra/assembly/amd/dsl.py

@@ -1,24 +1,27 @@
 # dsl.py - clean DSL for AMD assembly
-from typing import Any
 
 # ══════════════════════════════════════════════════════════════
 # Registers - unified src encoding space (0-511)
 # ══════════════════════════════════════════════════════════════
 
+def _reg_size(t: str | None) -> int: return {'b64': 2, 'f64': 2, 'u64': 2, 'i64': 2, 'b128': 4}.get(t, 1)
+
 class Reg:
-  # Register names vary by arch: RDNA has NULL@124/M0@125, CDNA has M0@124/reserved@125
-  # RDNA4 has DPP8@233, CDNA has SDWA@249/DPP@250/VCCZ@251/EXECZ@252
-  _NAMES = {102: "FLAT_SCRATCH_LO", 103: "FLAT_SCRATCH_HI", 104: "XNACK_MASK_LO", 105: "XNACK_MASK_HI",
-            106: "VCC_LO", 107: "VCC_HI", 124: "NULL", 125: "M0", 126: "EXEC_LO", 127: "EXEC_HI",
-            233: "DPP8", 234: "DPP8FI", 235: "SHARED_BASE", 236: "SHARED_LIMIT", 237: "PRIVATE_BASE", 238: "PRIVATE_LIMIT",
+  _NAMES = {106: "VCC_LO", 107: "VCC_HI", 124: "NULL", 125: "M0", 126: "EXEC_LO", 127: "EXEC_HI",
             240: "0.5", 241: "-0.5", 242: "1.0", 243: "-1.0", 244: "2.0", 245: "-2.0", 246: "4.0", 247: "-4.0",
-            248: "INV_2PI", 249: "SDWA", 250: "DPP", 251: "VCCZ", 252: "EXECZ", 253: "SCC", 254: "SRC_LDS_DIRECT", 255: "LIT"}
+            248: "INV_2PI", 250: "DPP16", 253: "SCC", 255: "LIT"}
   _PAIRS = {106: "VCC", 126: "EXEC"}
 
   def __init__(self, offset: int = 0, sz: int = 512, *, neg: bool = False, abs_: bool = False, hi: bool = False):
     self.offset, self.sz = offset, sz
     self.neg, self.abs_, self.hi = neg, abs_, hi
 
+  # TODO: remove these legacy aliases
+  @property
+  def count(self): return self.sz
+  @property
+  def idx(self): return self.offset
+
   def __hash__(self): return hash((self.offset, self.sz, self.neg, self.abs_, self.hi))
   def __getitem__(self, key):
     if isinstance(key, slice):
@@ -44,15 +47,11 @@ class Reg:
   def fmt(self, sz=None, parens=False, upper=False) -> str:
     o, sz = self.offset, sz or self.sz
     l, r = ("[", "]") if parens or sz > 1 else ("", "")  # brackets for multi-reg or when parens=True
-    if 256 <= o < 512:
-      idx = o - 256
-      base = f"v{l}{idx}{r}" if sz == 1 else f"v[{idx}:{idx + sz - 1}]"
+    if 256 <= o < 512: idx = o - 256; base = f"v{l}{idx}{r}" if sz == 1 else f"v[{idx}:{idx + sz - 1}]"
     elif o < 106: base = f"s{l}{o}{r}" if sz == 1 else f"s[{o}:{o + sz - 1}]"
     elif sz == 2 and o in self._PAIRS: base = self._PAIRS[o] if upper else self._PAIRS[o].lower()
     elif o in self._NAMES: base = self._NAMES[o] if upper else self._NAMES[o].lower()  # special regs (any sz)
-    elif 108 <= o < 124:
-      idx = o - 108
-      base = f"ttmp{l}{idx}{r}" if sz == 1 else f"ttmp[{idx}:{idx + sz - 1}]"
+    elif 108 <= o < 124: idx = o - 108; base = f"ttmp{l}{idx}{r}" if sz == 1 else f"ttmp[{idx}:{idx + sz - 1}]"
     elif 128 <= o <= 192: base = str(o - 128)  # inline int constants (0-64)
     elif 193 <= o <= 208: base = str(-(o - 192))  # inline negative int constants (-1 to -16)
     else: raise RuntimeError(f"unknown register: offset={o}, sz={sz}")
@@ -79,13 +78,9 @@ EXEC = src[126:127]
 # 128: 0, 129-192: integers 1-64, 193-208: integers -1 to -16
 # 240-248: float constants (0.5, -0.5, 1.0, -1.0, 2.0, -2.0, 4.0, -4.0, 1/(2*PI))
 INV_2PI = src[248]
-SDWA = src[249]
-DPP = DPP16 = src[250]
-VCCZ = src[251]
-EXECZ = src[252]
+DPP16 = src[250]
 SCC = src[253]
-SRC_LDS_DIRECT = src[254]
-LIT = src[255]           # literal constant marker
+# 255: literal constant
 v = src[256:511]         # VGPR0-255
 
 # ══════════════════════════════════════════════════════════════
@@ -98,13 +93,12 @@ class _Bits:
 bits = _Bits()
 
 class BitField:
-  name: str | None
-  def __init__(self, hi: int, lo: int, default = 0):
+  def __init__(self, hi: int, lo: int, default: int = 0):
     self.hi, self.lo, self.default, self.name, self.mask = hi, lo, default, None, (1 << (hi - lo + 1)) - 1
-  def __set_name__(self, owner, name: str): self.name = name
-  def __eq__(self, other) -> 'FixedBitField':  # type: ignore[override]
+  def __set_name__(self, owner, name): self.name = name
+  def __eq__(self, other) -> 'FixedBitField':
     if isinstance(other, int): return FixedBitField(self.hi, self.lo, other)
-    raise TypeError(f"BitField.__eq__ expects int, got {type(other).__name__}")
+    return NotImplemented
   def enum(self, enum_cls) -> 'EnumBitField': return EnumBitField(self.hi, self.lo, enum_cls)
   def encode(self, val) -> int:
     assert isinstance(val, int), f"BitField.encode expects int, got {type(val).__name__}"
@@ -113,14 +107,11 @@ class BitField:
   def set(self, raw: int, val) -> int:
     if val is None: val = self.default
     encoded = self.encode(val)
-    # Handle signed values: convert negative to 2's complement
-    if encoded < 0: encoded = encoded & self.mask
     if encoded < 0 or encoded > self.mask: raise RuntimeError(f"field '{self.name}': value {encoded} doesn't fit in {self.hi - self.lo + 1} bits")
     return (raw & ~(self.mask << self.lo)) | (encoded << self.lo)
   def __get__(self, obj, objtype=None):
     if obj is None: return self
     return self.decode((obj._raw >> self.lo) & self.mask)
-  def __set__(self, obj, val): obj._raw = self.set(obj._raw, val)
 
 class FixedBitField(BitField):
   def set(self, raw: int, val=None) -> int:
@@ -155,8 +146,7 @@ class SrcField(BitField):
     expected_size = self._valid_range[1] - self._valid_range[0] + 1
     actual_size = 1 << (hi - lo + 1)
     if actual_size != expected_size:
-      raise RuntimeError(f"{self.__class__.__name__}: field size {hi - lo + 1} bits ({actual_size}) "
-                         f"doesn't match range {self._valid_range} ({expected_size})")
+      raise RuntimeError(f"{self.__class__.__name__}: field size {hi - lo + 1} bits ({actual_size}) doesn't match range {self._valid_range} ({expected_size})")
 
   def encode(self, val) -> int:
     """Encode value. Returns 255 (literal marker) for out-of-range values."""
@@ -178,10 +168,7 @@ class SrcField(BitField):
     # Resize register based on operand info (skip non-resizable special registers)
     # VCC/EXEC pairs (106, 126), NULL (124), M0 (125), float constants (240-255)
     if reg.offset not in (124, 125) and not 240 <= reg.offset <= 255:
-      # Map variant field names (vsrc0->src0, vsrc1->src1, etc.) for DPP/SDWA classes
-      assert self.name is not None
-      name = self.name[1:] if self.name.startswith('v') and self.name[1:] in obj.op_regs else self.name
-      if sz := obj.op_regs.get(name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
+      if sz := obj.op_regs.get(self.name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
     return reg
 
 class VGPRField(SrcField):
@@ -224,21 +211,16 @@ class VDSTYField(BitField):
     if not isinstance(val, Reg): raise TypeError(f"VDSTYField requires Reg, got {type(val).__name__}")
     if not (256 <= val.offset < 512): raise ValueError(f"VDSTYField requires VGPR, got offset {val.offset}")
     return (val.offset - 256) >> 1
-  def __get__(self, obj, objtype=None):
-    if obj is None: return self
-    raw = (obj._raw >> self.lo) & self.mask
-    vdstx_bit0 = (obj.vdstx.offset - 256) & 1
-    vgpr_idx = (raw << 1) | (vdstx_bit0 ^ 1)
-    return Reg(256 + vgpr_idx, 1)
+  def decode(self, raw): return raw  # raw value, actual vdsty = (raw << 1) | ((vdstx & 1) ^ 1)
 
 # ══════════════════════════════════════════════════════════════
 # Operand info from XML
 # ══════════════════════════════════════════════════════════════
 
 import functools
-from tinygrad.runtime.autogen.amd.rdna3.operands import OPERANDS as OPERANDS_RDNA3
-from tinygrad.runtime.autogen.amd.rdna4.operands import OPERANDS as OPERANDS_RDNA4
-from tinygrad.runtime.autogen.amd.cdna.operands import OPERANDS as OPERANDS_CDNA
+from extra.assembly.amd.autogen.rdna3.operands import OPERANDS as OPERANDS_RDNA3
+from extra.assembly.amd.autogen.rdna4.operands import OPERANDS as OPERANDS_RDNA4
+from extra.assembly.amd.autogen.cdna.operands import OPERANDS as OPERANDS_CDNA
 OPERANDS = {**OPERANDS_CDNA, **OPERANDS_RDNA3, **OPERANDS_RDNA4}
 
 # ══════════════════════════════════════════════════════════════
@@ -258,15 +240,6 @@ def _get_variant(cls, suffix: str):
   module = sys.modules.get(cls.__module__)
   return getattr(module, f"{cls.__name__}{suffix}", None) if module else None
 
-def _canonical_name(name: str) -> str | None:
-  """Map operand name to canonical name."""
-  if name in ('src0', 'vsrc0', 'ssrc0'): return 's0'
-  if name in ('src1', 'vsrc1', 'ssrc1'): return 's1'
-  if name == 'src2': return 's2'
-  if name in ('vdst', 'sdst', 'sdata'): return 'd'
-  if name in ('data', 'vdata', 'data0', 'vsrc'): return 'data'
-  return None
-
 class Inst:
   _fields: list[tuple[str, BitField]]
   _base_size: int
@@ -276,36 +249,37 @@ class Inst:
     inherited = {}
     for base in reversed(cls.__mro__[1:]):
       if hasattr(base, '_fields'):
-        inherited.update(dict(base._fields))
+        inherited.update({name: field for name, field in base._fields})
     inherited.update({name: val for name, val in cls.__dict__.items() if isinstance(val, BitField)})
     cls._fields = list(inherited.items())
     cls._base_size = (max(f.hi for _, f in cls._fields) + 8) // 8
 
   def __new__(cls, *args, **kwargs):
-    # Auto-upgrade to variant if needed (only for base classes, not variants)
+    # Auto-upgrade to _LIT variant if needed (only for base classes, not variants)
     if not any(cls.__name__.endswith(sfx) for sfx in ('_LIT', '_DPP16', '_DPP8', '_SDWA', '_SDWA_SDST', '_MFMA')):
-      args_iter = iter(args)
-      for name, field in cls._fields:
-        if isinstance(field, FixedBitField): continue
-        val = kwargs.get(name) if name in kwargs else next(args_iter, None)
-        if not isinstance(field, SrcField): continue
-        if isinstance(val, Reg) and val.offset == 255 and (lit_cls := _get_variant(cls, '_LIT')): return lit_cls(*args, **kwargs)
-        if isinstance(val, Reg) and val.offset == 249:
-          if (sdwa_cls := _get_variant(cls, '_SDWA') or _get_variant(cls, '_SDWA_SDST')): return sdwa_cls(*args, **kwargs)
-        if isinstance(val, Reg) and val.offset == 250 and (dpp_cls := _get_variant(cls, '_DPP16')): return dpp_cls(*args, **kwargs)
-        if _needs_literal(val) and (lit_cls := _get_variant(cls, '_LIT')): return lit_cls(*args, **kwargs)
+      lit_cls = _get_variant(cls, '_LIT')
+      if lit_cls is not None:
+        # Check if any src field needs a literal
+        # Map positional args to field names to find src values
+        args_iter = iter(args)
+        for name, field in cls._fields:
+          if isinstance(field, FixedBitField): continue
+          val = kwargs.get(name) if name in kwargs else next(args_iter, None)
+          if isinstance(field, SrcField) and _needs_literal(val):
+            return lit_cls(*args, **kwargs)
     return object.__new__(cls)
 
   def __init__(self, *args, **kwargs):
     self._raw = 0
     # Map positional args to field names (skip FixedBitFields)
     args_iter = iter(args)
-    vals: dict[str, Any] = {}
+    vals = {}
     for name, field in self._fields:
       if isinstance(field, FixedBitField): vals[name] = None
       elif name in kwargs: vals[name] = kwargs[name]
       else: vals[name] = next(args_iter, None)
-    assert not (remaining := list(args_iter)), f"too many positional args: {remaining}"
+    remaining = list(args_iter)
+    assert not remaining, f"too many positional args: {remaining}"
     # Extract modifiers from Reg objects and merge into neg/abs/opsel
     neg_bits, abs_bits, opsel_bits = 0, 0, 0
     for name, bit in [('src0', 0), ('src1', 1), ('src2', 2)]:
@@ -330,16 +304,16 @@ class Inst:
     # Set all field values
     for name, field in self._fields:
       self._raw = field.set(self._raw, vals[name])
-    # Validate register sizes against operand info (skip special registers like NULL, VCC, EXEC, SDWA/DPP markers)
+    # Validate register sizes against operand info (skip special registers like NULL, VCC, EXEC)
     for name, expected in self.op_regs.items():
       if (val := vals.get(name)) is None: continue
-      if isinstance(val, Reg) and val.sz != expected and not (106 <= val.offset <= 127 or 249 <= val.offset <= 255):
+      if isinstance(val, Reg) and val.sz != expected and not (106 <= val.offset <= 127 or val.offset == 253):
         raise TypeError(f"{name} expects {expected} register(s), got {val.sz}")
 
   @property
-  def op_name(self) -> str: return getattr(self, 'op').name
+  def op_name(self) -> str: return self.op.name
   @property
-  def operands(self) -> dict: return OPERANDS.get(getattr(self, 'op'), {}) if hasattr(self, 'op') else {}
+  def operands(self) -> dict: return OPERANDS.get(self.op, {}) if hasattr(self, 'op') else {}
   def _is_cdna(self) -> bool: return 'cdna' in type(self).__module__
 
   @functools.cached_property
@@ -351,9 +325,9 @@ class Inst:
     if not self._is_cdna():
       name = self.op_name.lower()
       if 'cndmask' in name and 'src2' in bits: bits['src2'] = 32
-      if '_co_ci_' in name and 'src2' in bits: bits['src2'] = 32  # carry-in source
-      # VOP3SD: sdst is always wavefront-size dependent (carry-out or condition mask)
-      if 'VOP3SD' in type(self).__name__ and 'sdst' in bits: bits['sdst'] = 32
+      if '_co_ci_' in name:
+        if 'src2' in bits: bits['src2'] = 32
+        if 'sdst' in bits: bits['sdst'] = 32
       if 'cmp' in name and 'vdst' in bits: bits['vdst'] = 32
     # GLOBAL/FLAT: addr is 32-bit if saddr is valid SGPR, 64-bit if saddr is NULL
     # SCRATCH: addr is always 32-bit (offset from scratch base, not absolute address)
@@ -367,8 +341,8 @@ class Inst:
     # VGPRs: FP8/BF8(0,1)=8, FP6/BF6(2,3)=6, FP4(4)=4
     if 'f8f6f4' in getattr(self, 'op_name', '').lower():
       # Use explicit fields if available (VOP3PX2), else extract from VOP3P-MAI bit positions
-      cbsz = getattr(self, 'cbsz') if hasattr(type(self), 'cbsz') else (self._raw >> 8) & 0x7
-      blgp = getattr(self, 'blgp') if hasattr(type(self), 'blgp') else (self._raw >> 61) & 0x7
+      cbsz = self.cbsz if hasattr(type(self), 'cbsz') else (self._raw >> 8) & 0x7
+      blgp = self.blgp if hasattr(type(self), 'blgp') else (self._raw >> 61) & 0x7
       vgprs = {0: 8, 1: 8, 2: 6, 3: 6, 4: 4}
       bits['src0'], bits['src1'] = vgprs.get(cbsz, 8) * 32, vgprs.get(blgp, 8) * 32
     return bits
@@ -382,17 +356,12 @@ class Inst:
     """Get bit widths with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
     bits = {'d': 32, 's0': 32, 's1': 32, 's2': 32, 'data': 32}
     for name, val in self.op_bits.items():
-      if (cn := _canonical_name(name)): bits[cn] = val
+      if name in ('src0', 'vsrc0', 'ssrc0'): bits['s0'] = val
+      elif name in ('src1', 'vsrc1', 'ssrc1'): bits['s1'] = val
+      elif name == 'src2': bits['s2'] = val
+      elif name in ('vdst', 'sdst', 'sdata'): bits['d'] = val
+      elif name in ('data', 'vdata', 'data0'): bits['data'] = val
     return bits
-
-  @functools.cached_property
-  def canonical_operands(self) -> dict:
-    """Get operands with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
-    result = {}
-    for name, val in self.operands.items():
-      if (cn := _canonical_name(name)): result[cn] = val
-    return result
-
   @property
   def canonical_op_regs(self) -> dict[str, int]:
     """Get register counts with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
@@ -408,7 +377,9 @@ class Inst:
   @classmethod
   def _size(cls) -> int: return cls._base_size
   def size(self) -> int: return self._base_size
-  def disasm(self) -> str: raise NotImplementedError("disasm is no longer supported")
+  def disasm(self) -> str:
+    from extra.assembly.amd.disasm import disasm
+    return disasm(self)
 
   def to_bytes(self) -> bytes: return self._raw.to_bytes(self._base_size, 'little')
 

extra/assembly/amd/emu.py

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

extra/assembly/amd/pcode.py

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

extra/assembly/amd/sqtt.py

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

extra/assembly/amd/sqttmap.py

@@ -0,0 +1,132 @@
+# maps SQTT trace packets to instructions.
+from dataclasses import dataclass
+from typing import Iterator
+
+from tinygrad.runtime.support.elf import elf_loader
+
+from extra.assembly.amd.sqtt import decode, print_packets, INST, VALUINST, IMMEDIATE, WAVESTART, WAVEEND, InstOp, PacketType, IMMEDIATE_MASK
+from extra.assembly.amd.dsl import Inst
+from extra.assembly.amd.decode import decode_inst
+from extra.assembly.amd.autogen.rdna3.ins import SOPP, s_endpgm
+from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
+
+@dataclass(frozen=True)
+class InstructionInfo:
+  pc: int
+  wave: int
+  inst: Inst
+
+def map_insts(data:bytes, lib:bytes) -> Iterator[tuple[PacketType, InstructionInfo|None]]:
+  """maps SQTT packets to instructions, yields (packet, instruction_info or None)"""
+  # map pcs to insts
+  pc_map:dict[int, Inst] = {}
+  image, sections, _ = elf_loader(lib)
+  text = next((sh for sh in sections if sh.name == ".text"), None)
+  assert text is not None, "no .text section found"
+  text_off, text_size = text.header.sh_addr, text.header.sh_size
+  offset = text_off
+  while offset < text_off + text_size:
+    inst = decode_inst(image[offset:])
+    pc_map[offset-text_off] = inst
+    offset += inst.size()
+
+  wave_pc:dict[int, int] = {}
+  # only processing packets on one [CU, SIMD] unit
+  def simd_select(p) -> bool: return getattr(p, "cu", 0) == 0 and getattr(p, "simd", 0) == 0
+  for p in decode(data):
+    if not simd_select(p): continue
+    if isinstance(p, WAVESTART):
+      assert p.wave not in wave_pc, "only one inflight wave per unit"
+      wave_pc[p.wave] = 0
+      continue
+    if isinstance(p, WAVEEND):
+      pc = wave_pc.pop(p.wave)
+      yield (p, InstructionInfo(pc, p.wave, s_endpgm()))
+      continue
+    # skip OTHER_ instructions, they don't belong to this unit
+    if isinstance(p, INST) and p.op.name.startswith("OTHER_"): continue
+    if isinstance(p, IMMEDIATE_MASK):
+      # immediate mask may yield multiple times per packet
+      for wave in range(16):
+        if p.mask & (1 << wave):
+          inst = pc_map[pc:=wave_pc[wave]]
+          # can this assert be more strict?
+          assert isinstance(inst, SOPP), f"IMMEDIATE_MASK packet must map to SOPP, got {inst}"
+          wave_pc[wave] += inst.size()
+          yield (p, InstructionInfo(pc, wave, inst))
+      continue
+    if isinstance(p, (VALUINST, INST, IMMEDIATE)):
+      inst = pc_map[pc:=wave_pc[p.wave]]
+      # s_delay_alu doesn't get a packet?
+      if isinstance(inst, SOPP) and inst.op in {SOPPOp.S_DELAY_ALU}:
+        wave_pc[p.wave] += inst.size()
+        inst = pc_map[pc:=wave_pc[p.wave]]
+      # identify a branch instruction, only used for asserts
+      is_branch = isinstance(inst, SOPP) and "BRANCH" in inst.op_name
+      if is_branch: assert isinstance(p, INST) and p.op in {InstOp.JUMP_NO, InstOp.JUMP}, f"branch can only be folowed by jump packets, got {p}"
+      # JUMP handling
+      if isinstance(p, INST) and p.op is InstOp.JUMP:
+        assert is_branch, f"JUMP packet must map to a branch instruction, got {inst}"
+        x = inst.simm16 & 0xffff
+        wave_pc[p.wave] += inst.size() + (x - 0x10000 if x & 0x8000 else x)*4
+      else:
+        if is_branch: assert inst.op != SOPPOp.S_BRANCH, f"S_BRANCH must have a JUMP packet, got {p}"
+        wave_pc[p.wave] += inst.size()
+      yield (p, InstructionInfo(pc, p.wave, inst))
+      continue
+    # for all other packets (VMEMEXEC, ALUEXEC, etc.), yield with None
+    yield (p, None)
+
+# test to compare every packet with the rocprof decoder
+
+def test_rocprof_inst_traces_match(sqtt, prg, target):
+  from tinygrad.viz.serve import llvm_disasm
+  from extra.sqtt.roc import decode as roc_decode, InstExec
+  disasm = {addr+prg.base:inst_disasm for addr, inst_disasm in llvm_disasm(target, prg.lib).items()}
+  rctx = roc_decode([sqtt], {prg.name:disasm})
+  rwaves = rctx.inst_execs[(sqtt.kern, sqtt.exec_tag)]
+  rwaves_iter:dict[int, list[Iterator[InstExec]]] = {} # wave unit (0-15) -> list of inst trace iterators for all executions on that unit
+  for w in rwaves: rwaves_iter.setdefault(w.wave_id, []).append(w.unpack_insts())
+  rwaves_base = next(iter(disasm)) # base program counter
+
+  passed_insts = 0
+  for pkt, info in map_insts(sqtt.blob, prg.lib):
+    if DEBUG >= 2: print_packets([pkt])
+    if info is None: continue
+    if DEBUG >= 2: print(f"{' '*29}{info.inst.disasm()}")
+    rocprof_inst = next(rwaves_iter[info.wave][0])
+    ref_pc = rocprof_inst.pc-rwaves_base
+    # always check pc matches
+    assert ref_pc == info.pc, f"pc mismatch {ref_pc}:{disasm[rocprof_inst.pc][0]} != {info.pc}:{info.inst.disasm()}"
+    # special handling for s_endpgm, it marks the wave completion.
+    if info.inst == s_endpgm():
+      completed_wave = list(rwaves_iter[info.wave].pop(0))
+      assert len(completed_wave) == 0, f"incomplete instructions in wave {info.wave}"
+    # otherwise the packet timestamp is time + "stall"
+    else:
+      assert pkt._time == rocprof_inst.time+rocprof_inst.stall
+    passed_insts += 1
+
+  for k,v in rwaves_iter.items():
+    assert len(v) == 0, f"incomplete wave {k}"
+
+  print(f"passed for {passed_insts} instructions across {len(rwaves)} waves scheduled on {len(rwaves_iter)} wave units")
+
+if __name__ == "__main__":
+  import argparse, pickle, pathlib
+  from tinygrad.helpers import temp, DEBUG
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--profile', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
+                      default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  parser.add_argument('--kernel', type=str, default=None, metavar="NAME", help='Kernel to focus on (optional name, default: all kernels)')
+  args = parser.parse_args()
+  with open(args.profile, "rb") as f:
+    data = pickle.load(f)
+  sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
+  kern_events = {e.name:e for e in data if type(e).__name__ == "ProfileProgramEvent"}
+  target = next((e for e in data if type(e).__name__ == "ProfileDeviceEvent" and e.device.startswith("AMD"))).props["gfx_target_version"]
+  for e in sqtt_events:
+    if args.kernel is not None and args.kernel != e.kern: continue
+    if not e.itrace: continue
+    print(f"==== {e.kern}")
+    test_rocprof_inst_traces_match(e, kern_events[e.kern], target)

extra/assembly/amd/test/bench_emu.py

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

extra/assembly/amd/test/external_test_usability.py

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

extra/assembly/amd/test/helpers.py

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

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

@@ -4,16 +4,16 @@ Uses run_asm() with memory output, so tests can run on both emulator and real ha
 Set USE_HW=1 to run on both emulator and hardware, comparing results.
 """
 import ctypes, math, os, struct
-from tinygrad.runtime.autogen.amd.rdna3.ins import *
+from extra.assembly.amd.autogen.rdna3.ins import *
 
-from test.mockgpu.amd.emu import run_asm
-from tinygrad.renderer.amd.dsl import NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, M0
+from extra.assembly.amd.emu import run_asm
+from extra.assembly.amd.dsl import NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, M0
 
 def _i32(f: float) -> int: return struct.unpack('<I', struct.pack('<f', f))[0]
 def _f32(i: int) -> float: return struct.unpack('<f', struct.pack('<I', i & 0xFFFFFFFF))[0]
 
 # f16 conversion helpers
-def f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
+def _f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
 def f32_to_f16(f: float) -> int:
   f = float(f)
   if math.isnan(f): return 0x7e00
@@ -43,29 +43,11 @@ VCC = VCC_LO  # For VOP3SD sdst field (VCC_LO is exported from dsl)
 USE_HW = os.environ.get("USE_HW", "0") == "1"
 FLOAT_TOLERANCE = 1e-5
 
-def get_gpu_target() -> tuple[int, int, int]:
-  """Get the GPU target as (major, minor, stepping) tuple."""
-  if not USE_HW: return (0, 0, 0)
-  from tinygrad.device import Device
-  return Device["AMD"].target  # type: ignore[attr-defined]
-
-def skip_unless_gfx(min_major: int, min_minor: int = 0, reason: str = ""):
-  """Skip test if GPU target is below the minimum required version."""
-  import unittest
-  def decorator(test_func):
-    if not USE_HW: return test_func
-    target = get_gpu_target()
-    if target[0] < min_major or (target[0] == min_major and target[1] < min_minor):
-      return unittest.skip(reason or f"requires gfx{min_major}{min_minor}0+")(test_func)
-    return test_func
-  return decorator
-
-# Output buffer layout: vgpr[N_VGPRS][n_lanes], sgpr[N_SGPRS], vcc, scc, exec
+# Output buffer layout: vgpr[16][32], sgpr[16], vcc, scc, exec
 N_VGPRS, N_SGPRS, WAVE_SIZE = 16, 16, 32
+VGPR_BYTES = N_VGPRS * WAVE_SIZE * 4  # 16 regs * 32 lanes * 4 bytes = 2048
 SGPR_BYTES = N_SGPRS * 4  # 16 regs * 4 bytes = 64
-_VGPR_REGION = N_VGPRS * WAVE_SIZE * 4  # minimum vgpr region size (tests may use as scratch)
-def _out_bytes(n_lanes: int) -> int: return max(N_VGPRS * n_lanes * 4, _VGPR_REGION) + SGPR_BYTES + 12
-OUT_BYTES = _out_bytes(WAVE_SIZE)  # default for single-wave (backward compat)
+OUT_BYTES = VGPR_BYTES + SGPR_BYTES + 12  # + vcc + scc + exec
 
 # Float conversion helpers
 def f2i(f: float) -> int: return _i32(f)
@@ -76,10 +58,10 @@ def i642f(i: int) -> float: return struct.unpack('<d', struct.pack('<Q', i))[0]
 def assemble(instructions: list) -> bytes:
   return b''.join(inst.to_bytes() for inst in instructions)
 
-# Simple WaveState class for test output parsing (mirrors test/mockgpu/amd/emu.py interface for tests)
+# Simple WaveState class for test output parsing (mirrors emu.py interface for tests)
 class WaveState:
-  def __init__(self, n_lanes: int = 32):
-    self.vgpr = [[0] * 256 for _ in range(n_lanes)]  # vgpr[lane][reg]
+  def __init__(self):
+    self.vgpr = [[0] * 256 for _ in range(32)]  # vgpr[lane][reg]
     self.sgpr = [0] * 128
     self.vcc = 0
     self.scc = 0
@@ -103,53 +85,49 @@ def get_prologue_epilogue(n_lanes: int) -> tuple[list, list]:
     # Save EXEC early (before we modify it for VGPR stores)
     s_mov_b32(s[95], EXEC_LO),
     # Restore EXEC to all active lanes for VGPR stores (test may have modified EXEC)
-    s_mov_b32(EXEC_LO, (1 << min(n_lanes, WAVE_SIZE)) - 1),
+    s_mov_b32(EXEC_LO, (1 << n_lanes) - 1),
     s_load_b64(s[92:93], s[80:81], 0, soffset=NULL),
     s_waitcnt(0),  # simm16=0 waits for all
     v_lshlrev_b32_e32(v[240], 2, v[255]),
   ]
-  vgpr_bytes = N_VGPRS * n_lanes * 4
   for i in range(N_VGPRS):
-    epilogue.append(global_store_b32(addr=v[240], data=v[i], saddr=s[92:93], offset=i * n_lanes * 4))
+    epilogue.append(global_store_b32(addr=v[240], data=v[i], saddr=s[92:93], offset=i * WAVE_SIZE * 4))
   epilogue.append(v_mov_b32_e32(v[241], 0))
   epilogue.append(v_cmp_eq_u32_e32(v[255], v[241]))
   epilogue.append(s_and_saveexec_b32(s[94], VCC_LO))
-  # Scalar stores: only thread 0. Use v[240]=vgpr_bytes as base offset so immediate offsets stay small.
-  epilogue.append(v_mov_b32_e32(v[240], vgpr_bytes))
+  epilogue.append(v_mov_b32_e32(v[240], 0))
   for i in range(N_SGPRS):
     epilogue.append(v_mov_b32_e32(v[243], s[i]))
-    epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=i * 4))
+    epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + i * 4))
   epilogue.append(v_mov_b32_e32(v[243], s[90]))
-  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=SGPR_BYTES))
+  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + SGPR_BYTES))
   epilogue.append(v_mov_b32_e32(v[243], s[91]))
-  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=SGPR_BYTES + 4))
+  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + SGPR_BYTES + 4))
   # Store EXEC (saved earlier in s[95])
   epilogue.append(v_mov_b32_e32(v[243], s[95]))
-  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=SGPR_BYTES + 8))
+  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + SGPR_BYTES + 8))
   epilogue.append(s_mov_b32(EXEC_LO, s[94]))
   epilogue.append(s_endpgm())
   return prologue, epilogue
 
 def parse_output(out_buf: bytes, n_lanes: int) -> WaveState:
   """Parse output buffer into WaveState."""
-  vgpr_bytes = N_VGPRS * n_lanes * 4
-  st = WaveState(n_lanes)
+  st = WaveState()
   for i in range(N_VGPRS):
     for lane in range(n_lanes):
-      off = i * n_lanes * 4 + lane * 4
+      off = i * WAVE_SIZE * 4 + lane * 4
       st.vgpr[lane][i] = struct.unpack_from('<I', out_buf, off)[0]
   for i in range(N_SGPRS):
-    st.sgpr[i] = struct.unpack_from('<I', out_buf, vgpr_bytes + i * 4)[0]
-  st.vcc = struct.unpack_from('<I', out_buf, vgpr_bytes + SGPR_BYTES)[0]
-  st.scc = struct.unpack_from('<I', out_buf, vgpr_bytes + SGPR_BYTES + 4)[0]
+    st.sgpr[i] = struct.unpack_from('<I', out_buf, VGPR_BYTES + i * 4)[0]
+  st.vcc = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES)[0]
+  st.scc = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES + 4)[0]
   # Store EXEC in its proper location (index 126)
-  st.sgpr[EXEC_LO.offset] = struct.unpack_from('<I', out_buf, vgpr_bytes + SGPR_BYTES + 8)[0]
+  st.sgpr[EXEC_LO.offset] = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES + 8)[0]
   return st
 
 def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
   """Run instructions via emulator run_asm, dump state to memory, return WaveState."""
-  buf_sz = _out_bytes(n_lanes)
-  out_buf = (ctypes.c_uint8 * buf_sz)(*([0] * buf_sz))
+  out_buf = (ctypes.c_uint8 * OUT_BYTES)(*([0] * OUT_BYTES))
   out_addr = ctypes.addressof(out_buf)
 
   prologue, epilogue = get_prologue_epilogue(n_lanes)
@@ -163,7 +141,7 @@ def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
   # rsrc2: USER_SGPR_COUNT=2, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z=1, LDS_SIZE=128 (64KB)
   rsrc2 = 0x19c | (128 << 15)
   scratch_size = 0x10000  # 64KB per lane, matches .amdhsa_private_segment_fixed_size in run_program_hw
-  result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2, scratch_size)
+  result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2)
   assert result == 0, f"run_asm failed with {result}"
 
   return parse_output(bytes(out_buf), n_lanes)
@@ -176,7 +154,7 @@ def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
   from tinygrad.helpers import flat_mv
 
   dev = Device["AMD"]
-  compiler = HIPCompiler(dev.arch)  # type: ignore[attr-defined]
+  compiler = HIPCompiler(dev.arch)
 
   prologue, epilogue = get_prologue_epilogue(n_lanes)
   code = assemble(prologue + instructions + epilogue)
@@ -223,24 +201,18 @@ amdhsa.kernels:
 """
 
   lib = compiler.compile(asm_src)
-  prg = AMDProgram(dev, "test", lib)  # type: ignore[arg-type]
+  prg = AMDProgram(dev, "test", lib)
 
-  buf_sz = _out_bytes(n_lanes)
-  out_gpu = dev.allocator.alloc(buf_sz)
-  assert out_gpu.va_addr % 16 == 0, f"buffer not 16-byte aligned: 0x{out_gpu.va_addr:x}"
+  out_gpu = dev.allocator.alloc(OUT_BYTES)
   prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
 
-  out_buf = bytearray(buf_sz)
+  out_buf = bytearray(OUT_BYTES)
   dev.allocator._copyout(flat_mv(memoryview(out_buf)), out_gpu)
 
   return parse_output(bytes(out_buf), n_lanes)
 
-def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS, ulp_tolerance: int = 0) -> list[str]:
-  """Compare two WaveStates and return list of differences.
-
-  Args:
-    ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
-  """
+def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS) -> list[str]:
+  """Compare two WaveStates and return list of differences."""
   import math
   diffs = []
   for i in range(n_vgprs):
@@ -251,11 +223,6 @@ def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgp
         emu_f, hw_f = _f32(emu_val), _f32(hw_val)
         if math.isnan(emu_f) and math.isnan(hw_f):
           continue
-        # Check ULP difference for floats (only for same-sign values)
-        if ulp_tolerance > 0 and (emu_val < 0x80000000) == (hw_val < 0x80000000):
-          ulp_diff = abs(int(emu_val) - int(hw_val))
-          if ulp_diff <= ulp_tolerance:
-            continue
         diffs.append(f"v[{i}] lane {lane}: emu=0x{emu_val:08x} ({emu_f:.6g}) hw=0x{hw_val:08x} ({hw_f:.6g})")
   for i in range(N_SGPRS):
     emu_val = emu_st.sgpr[i]
@@ -268,20 +235,17 @@ def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgp
     diffs.append(f"scc: emu={emu_st.scc} hw={hw_st.scc}")
   return diffs
 
-def run_program(instructions: list, n_lanes: int = 1, ulp_tolerance: int = 0) -> WaveState:
+def run_program(instructions: list, n_lanes: int = 1) -> WaveState:
   """Run instructions and return WaveState.
 
   If USE_HW=1, runs on both emulator and hardware, compares results, and raises if they differ.
   Otherwise, runs only on emulator.
-
-  Args:
-    ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
   """
   emu_st = run_program_emu(instructions, n_lanes)
   if USE_HW:
     hw_st = run_program_hw(instructions, n_lanes)
-    diffs = compare_wave_states(emu_st, hw_st, n_lanes, ulp_tolerance=ulp_tolerance)
+    diffs = compare_wave_states(emu_st, hw_st, n_lanes)
     if diffs:
-      raise AssertionError("Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
+      raise AssertionError(f"Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
     return hw_st
   return emu_st

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

@@ -5,7 +5,7 @@ Includes: ds_store_b32, ds_load_b32, ds_store_2addr_*, ds_load_2addr_*,
           ds_inc_*, ds_dec_*, ds_cmpstore_*, ds_storexchg_*
 """
 import unittest
-from test.amd.hw.helpers import *
+from extra.assembly.amd.test.hw.helpers import *
 
 class TestDS2Addr(unittest.TestCase):
   """Tests for DS_*_2ADDR instructions."""
@@ -117,58 +117,6 @@ class TestDS2AddrMore(unittest.TestCase):
     self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
     self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should be untouched")
 
-  def test_ds_load_2addr_b64_addr_overlaps_vdst(self):
-    """DS_LOAD_2ADDR_B64 where addr register overlaps vdst range.
-
-    Hardware reads the address before writing any results, so addr=v[4]
-    with vdst=v[4:7] must load all 4 dwords using the original v[4] value.
-    """
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[2], 0xAAAAAAAA),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
-      s_mov_b32(s[2], 0xBBBBBBBB),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=4),
-      s_mov_b32(s[2], 0xCCCCCCCC),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=8),
-      s_mov_b32(s[2], 0xDDDDDDDD),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=12),
-      s_waitcnt(lgkmcnt=0),
-      # addr=v[4] overlaps vdst=v[4:7]
-      v_mov_b32_e32(v[4], 0),
-      DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[4], vdst=v[4:7], offset0=0, offset1=1),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA, "v4 = LDS[0:4]")
-    self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB, "v5 = LDS[4:8]")
-    self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC, "v6 = LDS[8:12]")
-    self.assertEqual(st.vgpr[0][7], 0xDDDDDDDD, "v7 = LDS[12:16]")
-
-  def test_ds_load_2addr_b32_addr_overlaps_vdst(self):
-    """DS_LOAD_2ADDR_B32 where addr register overlaps vdst range."""
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[2], 0xAAAAAAAA),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
-      s_mov_b32(s[2], 0xBBBBBBBB),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=4),
-      s_waitcnt(lgkmcnt=0),
-      # addr=v[2] overlaps vdst=v[2:3]
-      v_mov_b32_e32(v[2], 0),
-      DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[2], vdst=v[2:3], offset0=0, offset1=1),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 = LDS[0:4]")
-    self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 = LDS[4:8]")
-
   def test_ds_load_b64_no_overwrite(self):
     """DS_LOAD_B64 should only write 2 VGPRs."""
     instructions = [
@@ -190,50 +138,6 @@ class TestDS2AddrMore(unittest.TestCase):
     self.assertEqual(st.vgpr[0][4], 0x12345678, "v4 should be untouched")
 
 
-class TestDSB96(unittest.TestCase):
-  """Tests for DS_STORE_B96 and DS_LOAD_B96 (96-bit / 3 dwords)."""
-
-  def test_ds_store_load_b96(self):
-    """DS_STORE_B96 stores 3 VGPRs, DS_LOAD_B96 loads them back."""
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[0], 0x11111111),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[0], 0x22222222),
-      v_mov_b32_e32(v[1], s[0]),
-      s_mov_b32(s[0], 0x33333333),
-      v_mov_b32_e32(v[2], s[0]),
-      ds_store_b96(addr=v[10], data0=v[0:2]),
-      s_waitcnt(lgkmcnt=0),
-      ds_load_b96(addr=v[10], vdst=v[4:6]),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have first dword")
-    self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have second dword")
-    self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should have third dword")
-
-  def test_ds_store_b96_with_offset(self):
-    """DS_STORE_B96 with non-zero offset."""
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[0], 0xAAAAAAAA),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[0], 0xBBBBBBBB),
-      v_mov_b32_e32(v[1], s[0]),
-      s_mov_b32(s[0], 0xCCCCCCCC),
-      v_mov_b32_e32(v[2], s[0]),
-      DS(DSOp.DS_STORE_B96, addr=v[10], data0=v[0:2], offset0=12),
-      s_waitcnt(lgkmcnt=0),
-      DS(DSOp.DS_LOAD_B96, addr=v[10], vdst=v[4:6], offset0=12),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
-    self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB)
-    self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC)
-
-
 class TestDSB128(unittest.TestCase):
   """Tests for DS_STORE_B128 and DS_LOAD_B128 (128-bit / 4 dwords)."""
 
@@ -653,6 +557,7 @@ class TestDS2AddrStride64(unittest.TestCase):
     self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "new val 0")
     self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "new val 1")
 
+
   def test_ds_storexchg_rtn_b64(self):
     """DS_STOREXCHG_RTN_B64: exchange 64-bit value and return old."""
     instructions = [
@@ -770,193 +675,5 @@ class TestAtomicOrdering(unittest.TestCase):
     self.assertEqual(st.vgpr[0][4], 150, "Final value should be 150")
 
 
-class TestDsPermute(unittest.TestCase):
-  """Tests for DS_PERMUTE_B32 and DS_BPERMUTE_B32 instructions."""
-
-  def test_ds_permute_b32_identity(self):
-    """DS_PERMUTE_B32 with identity permutation (lane 0 sends to lane 0)."""
-    # For simplicity, test with single lane
-    instructions = [
-      v_mov_b32_e32(v[0], 0),  # addr = 0 (lane 0)
-      v_mov_b32_e32(v[1], 0xDEADBEEF),  # data
-      ds_permute_b32(v[2], v[0], v[1]),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # Lane 0 sends to lane 0, so lane 0 gets 0xDEADBEEF
-    self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
-
-  def test_ds_bpermute_b32_identity(self):
-    """DS_BPERMUTE_B32 with identity permutation (each lane reads from itself)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0),  # addr = 0 (read from lane 0)
-      v_mov_b32_e32(v[1], 0xCAFEBABE),  # data in lane 0
-      ds_bpermute_b32(v[2], v[0], v[1]),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # Lane 0 reads from lane 0's v[1]
-    self.assertEqual(st.vgpr[0][2], 0xCAFEBABE)
-
-  def test_ds_permute_b32_broadcast(self):
-    """DS_PERMUTE_B32 broadcast - all lanes send to lane 0."""
-    # With 4 lanes, all sending to lane 0, highest lane wins
-    instructions = [
-      v_mov_b32_e32(v[0], 0),  # All lanes send to addr 0 (lane 0)
-      v_mov_b32_e32(v[1], 0x11111111),  # All lanes send same data
-      ds_permute_b32(v[2], v[0], v[1]),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=4)
-    # Lane 0 receives data (highest numbered active lane wins)
-    self.assertEqual(st.vgpr[0][2], 0x11111111)
-
-  def test_ds_bpermute_b32_xor_swap(self):
-    """DS_BPERMUTE_B32 with XOR-1 pattern — each lane reads from lane^1.
-
-    This is the pattern used by warp_shfl_xor in flash attention for reduce_max/reduce_sum.
-    Each lane has a unique value (lane_id + 100), and reads from the adjacent lane.
-    """
-    instructions = [
-      # v[0] = (lane_id ^ 1) * 4  (byte offset for bpermute)
-      v_xor_b32_e32(v[0], 1, v[255]),
-      v_lshlrev_b32_e32(v[0], 2, v[0]),
-      # v[1] = lane_id + 100 (unique per-lane value)
-      s_mov_b32(s[0], 100),
-      v_add_nc_u32_e32(v[1], s[0], v[255]),
-      # ds_bpermute: v[2] = v[1] from lane (lane_id ^ 1)
-      ds_bpermute_b32(vdst=v[2], addr=v[0], data0=v[1]),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=32)
-    for lane in range(32):
-      src_lane = lane ^ 1
-      expected = src_lane + 100
-      self.assertEqual(st.vgpr[lane][2], expected, f"lane {lane}: expected v[1] from lane {src_lane} = {expected}, got {st.vgpr[lane][2]}")
-
-
-class TestDSSubDword(unittest.TestCase):
-  """Tests for sub-dword DS operations (ds_store_b16, ds_store_b16_d16_hi)."""
-
-  def test_ds_store_b16_and_d16_hi(self):
-    """DS_STORE_B16 stores low 16 bits, DS_STORE_B16_D16_HI stores high 16 bits to adjacent LDS half-words."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0),
-      v_mov_b32_e32(v[1], 0xBEEF1234),
-      DS(DSOp.DS_STORE_B16, addr=v[0], data0=v[1], offset0=0),
-      DS(DSOp.DS_STORE_B16_D16_HI, addr=v[0], data0=v[1], offset0=2),
-      s_waitcnt(lgkmcnt=0),
-      ds_load_b32(vdst=v[2], addr=v[0], offset0=0),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xBEEF1234, "lo=0x1234 at byte 0, hi=0xBEEF at byte 2")
-
-
-class TestDSLargeOffset(unittest.TestCase):
-  """Tests for DS instructions with offsets > 255 (offset1 > 0).
-
-  The DS offset is a 16-bit value encoded as (offset1 << 8) | offset0.
-  These tests verify that offset1 is used correctly, not just offset0.
-  """
-
-  def test_ds_store_load_b32_offset_256(self):
-    """DS_STORE_B32/DS_LOAD_B32 with offset=256 (offset0=0, offset1=1)."""
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[0], 0xDEADBEEF),
-      v_mov_b32_e32(v[0], s[0]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=0, offset1=1),  # offset = 256
-      s_waitcnt(lgkmcnt=0),
-      ds_load_b32(addr=v[10], vdst=v[1], offset0=0, offset1=1),  # offset = 256
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
-
-  def test_ds_store_load_b32_offset_300(self):
-    """DS_STORE_B32/DS_LOAD_B32 with offset=300 (offset0=44, offset1=1)."""
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[0], 0xCAFEBABE),
-      v_mov_b32_e32(v[0], s[0]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=44, offset1=1),  # offset = 300
-      s_waitcnt(lgkmcnt=0),
-      ds_load_b32(addr=v[10], vdst=v[1], offset0=44, offset1=1),  # offset = 300
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
-
-  def test_ds_store_load_b64_offset_512(self):
-    """DS_STORE_B64/DS_LOAD_B64 with offset=512 (offset0=0, offset1=2)."""
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[0], 0x11111111),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[0], 0x22222222),
-      v_mov_b32_e32(v[1], s[0]),
-      ds_store_b64(addr=v[10], data0=v[0:1], offset0=0, offset1=2),  # offset = 512
-      s_waitcnt(lgkmcnt=0),
-      ds_load_b64(addr=v[10], vdst=v[2:3], offset0=0, offset1=2),  # offset = 512
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x11111111)
-    self.assertEqual(st.vgpr[0][3], 0x22222222)
-
-  def test_ds_large_offset_distinct_from_small(self):
-    """Verify offset=256 and offset=0 address different LDS locations."""
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[0], 0xAAAAAAAA),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[0], 0xBBBBBBBB),
-      v_mov_b32_e32(v[1], s[0]),
-      # Store 0xAAAAAAAA at offset=0, 0xBBBBBBBB at offset=256
-      ds_store_b32(addr=v[10], data0=v[0], offset0=0, offset1=0),  # offset = 0
-      ds_store_b32(addr=v[10], data0=v[1], offset0=0, offset1=1),  # offset = 256
-      s_waitcnt(lgkmcnt=0),
-      # Read back both
-      ds_load_b32(addr=v[10], vdst=v[2], offset0=0, offset1=0),  # offset = 0
-      ds_load_b32(addr=v[10], vdst=v[3], offset0=0, offset1=1),  # offset = 256
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "offset=0 should read 0xAAAAAAAA")
-    self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "offset=256 should read 0xBBBBBBBB")
-
-  def test_ds_store_load_b32_offset_448(self):
-    """DS_STORE_B32/DS_LOAD_B32 with offset=448 (offset0=192, offset1=1) - matches matmul B tile."""
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[0], 0x12345678),
-      v_mov_b32_e32(v[0], s[0]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=192, offset1=1),  # offset = 448
-      s_waitcnt(lgkmcnt=0),
-      ds_load_b32(addr=v[10], vdst=v[1], offset0=192, offset1=1),  # offset = 448
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][1], 0x12345678)
-
-  def test_ds_load_b64_offset_392(self):
-    """DS_LOAD_B64 with offset=392 (offset0=136, offset1=1) - matches matmul B tile load."""
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[0], 0xAABBCCDD),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[0], 0x11223344),
-      v_mov_b32_e32(v[1], s[0]),
-      ds_store_b64(addr=v[10], data0=v[0:1], offset0=136, offset1=1),  # offset = 392
-      s_waitcnt(lgkmcnt=0),
-      ds_load_b64(addr=v[10], vdst=v[2:3], offset0=136, offset1=1),  # offset = 392
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xAABBCCDD)
-    self.assertEqual(st.vgpr[0][3], 0x11223344)
-
-
 if __name__ == '__main__':
   unittest.main()

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

@@ -3,7 +3,7 @@
 Includes: flat_load_*, flat_store_*, flat_atomic_*
 """
 import unittest
-from test.amd.hw.helpers import *
+from extra.assembly.amd.test.hw.helpers import *
 
 class TestFlatAtomic(unittest.TestCase):
   """Tests for FLAT atomic instructions."""

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

@@ -3,7 +3,7 @@
 Includes: global_load_*, global_store_*, global_atomic_*, global_load_d16_*
 """
 import unittest
-from test.amd.hw.helpers import *
+from extra.assembly.amd.test.hw.helpers import *
 
 class TestGlobalAtomic(unittest.TestCase):
   """Tests for GLOBAL atomic instructions."""
@@ -523,157 +523,5 @@ class TestD16HiLoads(unittest.TestCase):
     self.assertEqual(byte5, 0x00, f"byte5: expected 0x00, got 0x{byte5:02x}")
 
 
-class TestGlobalOffset(unittest.TestCase):
-  """Tests for GLOBAL instructions with different offsets.
-
-  These tests verify that instruction deduplication correctly handles different offset values.
-  If offset is made dynamic incorrectly, instructions with different offsets may load/store wrong data.
-  """
-
-  def test_global_load_different_offsets(self):
-    """Load from two different offsets and verify correct values."""
-    instructions = [
-      s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
-      s_waitcnt(lgkmcnt=0),
-      v_mov_b32_e32(v[0], s[2]),
-      v_mov_b32_e32(v[1], s[3]),
-      # Store 0xAAAAAAAA at offset 100
-      s_mov_b32(s[0], 0xAAAAAAAA),
-      v_mov_b32_e32(v[2], s[0]),
-      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
-      # Store 0xBBBBBBBB at offset 200
-      s_mov_b32(s[0], 0xBBBBBBBB),
-      v_mov_b32_e32(v[2], s[0]),
-      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
-      s_waitcnt(vmcnt=0),
-      # Load from offset 100 -> should get 0xAAAAAAAA
-      GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[3], saddr=SrcEnum.NULL, offset=100),
-      # Load from offset 200 -> should get 0xBBBBBBBB
-      GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=200),
-      s_waitcnt(vmcnt=0),
-      v_mov_b32_e32(v[0], v[3]),
-      v_mov_b32_e32(v[1], v[4]),
-      s_mov_b32(s[2], 0),
-      s_mov_b32(s[3], 0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
-    self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
-
-  def test_global_store_different_offsets(self):
-    """Store to two different offsets and verify correct values."""
-    instructions = [
-      s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
-      s_waitcnt(lgkmcnt=0),
-      v_mov_b32_e32(v[0], s[2]),
-      v_mov_b32_e32(v[1], s[3]),
-      # Store 0x11111111 at offset 300
-      s_mov_b32(s[0], 0x11111111),
-      v_mov_b32_e32(v[2], s[0]),
-      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=300),
-      # Store 0x22222222 at offset 400
-      s_mov_b32(s[0], 0x22222222),
-      v_mov_b32_e32(v[3], s[0]),
-      global_store_b32(addr=v[0:1], data=v[3], saddr=SrcEnum.NULL, offset=400),
-      s_waitcnt(vmcnt=0),
-      # Load back to verify
-      GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=300),
-      GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[5], saddr=SrcEnum.NULL, offset=400),
-      s_waitcnt(vmcnt=0),
-      v_mov_b32_e32(v[0], v[4]),
-      v_mov_b32_e32(v[1], v[5]),
-      s_mov_b32(s[2], 0),
-      s_mov_b32(s[3], 0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0x11111111, f"offset 300: expected 0x11111111, got 0x{st.vgpr[0][0]:08x}")
-    self.assertEqual(st.vgpr[0][1], 0x22222222, f"offset 400: expected 0x22222222, got 0x{st.vgpr[0][1]:08x}")
-
-  def test_global_negative_offset_no_saddr(self):
-    """Test negative offset without saddr (VGPR pair for address).
-    Store 0xAAAA at offset 100, 0xBBBB at offset 200.
-    Load with offset -100 from vaddr pointing to base+200 -> should get 0xAAAA (at 100).
-    Load with offset -100 from vaddr pointing to base+300 -> should get 0xBBBB (at 200)."""
-    instructions = [
-      s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
-      s_waitcnt(lgkmcnt=0),
-      v_mov_b32_e32(v[0], s[2]),
-      v_mov_b32_e32(v[1], s[3]),
-      # Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
-      s_mov_b32(s[0], 0xAAAAAAAA),
-      v_mov_b32_e32(v[2], s[0]),
-      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
-      s_mov_b32(s[0], 0xBBBBBBBB),
-      v_mov_b32_e32(v[2], s[0]),
-      global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
-      s_waitcnt(vmcnt=0),
-      # vaddr = base+200, load with offset -100 -> should get value at 100
-      s_add_u32(s[4], s[2], 200),
-      s_addc_u32(s[5], s[3], 0),
-      v_mov_b32_e32(v[4], s[4]),
-      v_mov_b32_e32(v[5], s[5]),
-      GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[6], saddr=SrcEnum.NULL, offset=-100),
-      # vaddr = base+300, load with offset -100 -> should get value at 200
-      s_add_u32(s[4], s[2], 300),
-      s_addc_u32(s[5], s[3], 0),
-      v_mov_b32_e32(v[4], s[4]),
-      v_mov_b32_e32(v[5], s[5]),
-      GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[7], saddr=SrcEnum.NULL, offset=-100),
-      s_waitcnt(vmcnt=0),
-      v_mov_b32_e32(v[0], v[6]),
-      v_mov_b32_e32(v[1], v[7]),
-      v_mov_b32_e32(v[4], 0),
-      v_mov_b32_e32(v[5], 0),
-      v_mov_b32_e32(v[6], 0),
-      v_mov_b32_e32(v[7], 0),
-      s_mov_b32(s[2], 0),
-      s_mov_b32(s[3], 0),
-      s_mov_b32(s[4], 0),
-      s_mov_b32(s[5], 0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
-    self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
-
-  def test_global_negative_offset_with_saddr(self):
-    """Test negative offset with saddr (SGPR pair for base address).
-    Store 0xAAAA at offset 100, 0xBBBB at offset 200.
-    Load with offset -100 from saddr pointing to base+200 -> should get 0xAAAA (at 100).
-    Load with offset -100 from saddr pointing to base+300 -> should get 0xBBBB (at 200)."""
-    instructions = [
-      s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
-      s_waitcnt(lgkmcnt=0),
-      v_mov_b32_e32(v[0], 0),
-      # Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
-      s_mov_b32(s[0], 0xAAAAAAAA),
-      v_mov_b32_e32(v[2], s[0]),
-      global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
-      s_mov_b32(s[0], 0xBBBBBBBB),
-      v_mov_b32_e32(v[2], s[0]),
-      global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
-      s_waitcnt(vmcnt=0),
-      # saddr = base+200, load with offset -100 -> should get value at 100
-      s_add_u32(s[4], s[2], 200),
-      s_addc_u32(s[5], s[3], 0),
-      GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[6], saddr=s[4:5], offset=-100),
-      # saddr = base+300, load with offset -100 -> should get value at 200
-      s_add_u32(s[4], s[2], 300),
-      s_addc_u32(s[5], s[3], 0),
-      GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[7], saddr=s[4:5], offset=-100),
-      s_waitcnt(vmcnt=0),
-      v_mov_b32_e32(v[0], v[6]),
-      v_mov_b32_e32(v[1], v[7]),
-      v_mov_b32_e32(v[6], 0),
-      v_mov_b32_e32(v[7], 0),
-      s_mov_b32(s[2], 0),
-      s_mov_b32(s[3], 0),
-      s_mov_b32(s[4], 0),
-      s_mov_b32(s[5], 0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
-    self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
-
-
 if __name__ == '__main__':
   unittest.main()

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

@@ -4,7 +4,7 @@ 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 test.amd.hw.helpers import *
+from extra.assembly.amd.test.hw.helpers import *
 
 class TestBasicScalar(unittest.TestCase):
   """Tests for basic scalar operations."""
@@ -62,7 +62,6 @@ class TestBasicScalar(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     self.assertEqual(st.sgpr[1], 0x80000000)
 
-  @skip_unless_gfx(11, 5, "SALU FP ops require gfx1150+")
   def test_s_fmamk_f32(self):
     """S_FMAMK_F32: D = S0 * literal + S1."""
     # 2.0 * 3.0 + 1.0 = 7.0
@@ -74,7 +73,6 @@ class TestBasicScalar(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     self.assertEqual(st.sgpr[2], f2i(7.0))
 
-  @skip_unless_gfx(11, 5, "SALU FP ops require gfx1150+")
   def test_s_fmamk_f32_negative(self):
     """S_FMAMK_F32 with negative values."""
     # -2.0 * 4.0 + 10.0 = 2.0
@@ -87,50 +85,6 @@ class TestBasicScalar(unittest.TestCase):
     self.assertEqual(st.sgpr[2], f2i(2.0))
 
 
-class TestPack(unittest.TestCase):
-  """Tests for S_PACK instructions."""
-
-  def test_s_pack_ll_b32_b16(self):
-    """S_PACK_LL_B32_B16 packs low 16 bits of two sources into one 32-bit result."""
-    instructions = [
-      s_mov_b32(s[0], 0xDEADAAAA),
-      s_mov_b32(s[1], 0xDEADBBBB),
-      s_pack_ll_b32_b16(s[2], s[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 0xBBBBAAAA)
-
-  def test_s_pack_lh_b32_b16(self):
-    """S_PACK_LH_B32_B16: D0 = { S1[31:16], S0[15:0] }."""
-    instructions = [
-      s_mov_b32(s[0], 0xDEADAAAA),
-      s_mov_b32(s[1], 0xDEADBBBB),
-      s_pack_lh_b32_b16(s[2], s[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 0xDEADAAAA)
-
-  def test_s_pack_hh_b32_b16(self):
-    """S_PACK_HH_B32_B16: D0 = { S1[31:16], S0[31:16] }."""
-    instructions = [
-      s_mov_b32(s[0], 0xDEADAAAA),
-      s_mov_b32(s[1], 0xDEADBBBB),
-      s_pack_hh_b32_b16(s[2], s[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 0xDEADDEAD)
-
-  def test_s_pack_hl_b32_b16(self):
-    """S_PACK_HL_B32_B16: D0 = { S1[15:0], S0[31:16] }."""
-    instructions = [
-      s_mov_b32(s[0], 0xDEADAAAA),
-      s_mov_b32(s[1], 0xDEADBBBB),
-      s_pack_hl_b32_b16(s[2], s[0], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 0xBBBBDEAD)
-
-
 class TestQuadmaskWqm(unittest.TestCase):
   """Tests for S_QUADMASK_B32 and S_WQM_B32."""
 
@@ -665,343 +619,5 @@ class Test64BitCompare(unittest.TestCase):
     self.assertEqual(st.sgpr[4], 1)
 
 
-class TestSOPPNop(unittest.TestCase):
-  """Tests for S_NOP and other SOPP instructions with expression-based for loops.
-
-  S_NOP's pcode uses 'for i in 0U : SIMM16.u16[3 : 0].u32 do' which requires
-  the parser to handle non-constant loop bounds.
-  """
-
-  def test_s_nop_basic(self):
-    """S_NOP executes without side effects."""
-    # S_NOP with immediate 0 should just do nothing
-    instructions = [
-      s_mov_b32(s[0], 42),
-      s_nop(0),  # nop with simm16=0
-      s_mov_b32(s[1], 100),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[0], 42)
-    self.assertEqual(st.sgpr[1], 100)
-
-  def test_s_nop_with_count(self):
-    """S_NOP with count parameter executes multiple nops."""
-    # S_NOP with immediate 3 should execute 4 nops (0:3 inclusive)
-    instructions = [
-      s_mov_b32(s[0], 1),
-      s_nop(3),  # nop with simm16=3 -> 4 iterations
-      s_add_u32(s[0], s[0], 1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[0], 2)
-
-
-class TestNullRegister(unittest.TestCase):
-  """Tests for NULL register (124) behavior - writes should be discarded, reads return 0."""
-
-  def test_s_mov_b32_from_null(self):
-    """S_MOV_B32 from NULL should read as 0."""
-    instructions = [
-      s_mov_b32(s[0], 0xDEADBEEF),  # Set s[0] to sentinel
-      s_mov_b32(s[0], NULL),  # Read from NULL - should be 0
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[0], 0)
-
-  def test_s_add_u32_with_null_src(self):
-    """S_ADD_U32 with NULL as source should use 0."""
-    instructions = [
-      s_mov_b32(s[0], 100),
-      s_add_u32(s[1], s[0], NULL),  # 100 + 0 = 100
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[1], 100)
-
-  def test_s_mov_b32_to_null(self):
-    """S_MOV_B32 to NULL (sdst=124) should discard the write."""
-    instructions = [
-      s_mov_b32(s[0], 0xDEADBEEF),  # Set s[0] to sentinel
-      s_mov_b32(NULL, 42),  # Write to NULL - should be discarded
-      # s[0] should still be 0xDEADBEEF since NULL write doesn't affect it
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[0], 0xDEADBEEF)
-
-  def test_s_add_u32_to_null(self):
-    """S_ADD_U32 with sdst=NULL should discard result but still set SCC."""
-    instructions = [
-      s_mov_b32(s[0], 0xFFFFFFFF),
-      s_mov_b32(s[1], 1),
-      s_add_u32(NULL, s[0], s[1]),  # overflow, write to NULL
-      s_cselect_b32(s[2], 1, 0),  # capture SCC
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # SCC should still be set from overflow even though result was discarded
-    self.assertEqual(st.sgpr[2], 1)
-    self.assertEqual(st.scc, 1)
-
-  def test_s_and_b32_to_null(self):
-    """S_AND_B32 with sdst=NULL should discard result but still set SCC."""
-    instructions = [
-      s_mov_b32(s[0], 0xFF00FF00),
-      s_mov_b32(s[1], 0x0F0F0F0F),
-      s_and_b32(NULL, s[0], s[1]),  # result=0x0F000F00, non-zero so SCC=1
-      s_cselect_b32(s[2], 1, 0),  # capture SCC
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 1)  # SCC=1 because result was non-zero
-    self.assertEqual(st.scc, 1)
-
-  def test_s_or_b32_to_null_zero_result(self):
-    """S_OR_B32 with sdst=NULL and zero result should set SCC=0."""
-    instructions = [
-      s_mov_b32(s[0], 0),
-      s_mov_b32(s[1], 0),
-      s_or_b32(NULL, s[0], s[1]),  # result=0, so SCC=0
-      s_cselect_b32(s[2], 1, 0),  # capture SCC
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[2], 0)  # SCC=0 because result was zero
-    self.assertEqual(st.scc, 0)
-
-
-class Test64BitSOP1InlineConstants(unittest.TestCase):
-  """Tests for 64-bit SOP1 instructions with inline constants.
-
-  Regression tests for bug where rsrc_dyn didn't properly handle 64-bit
-  inline constants, incorrectly duplicating lo bits to hi instead of
-  zero/sign-extending.
-  """
-
-  def test_s_mov_b64_inline_0(self):
-    """S_MOV_B64 with inline constant 0."""
-    instructions = [
-      s_mov_b64(s[0:1], 0),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0)
-    self.assertEqual(st.vgpr[0][1], 0)
-
-  def test_s_mov_b64_inline_16(self):
-    """S_MOV_B64 with inline constant 16 should set lo=16, hi=0."""
-    instructions = [
-      s_mov_b64(s[0:1], 16),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 16)
-    self.assertEqual(st.vgpr[0][1], 0)
-
-  def test_s_mov_b64_inline_64(self):
-    """S_MOV_B64 with inline constant 64 (max positive)."""
-    instructions = [
-      s_mov_b64(s[0:1], 64),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 64)
-    self.assertEqual(st.vgpr[0][1], 0)
-
-  def test_s_mov_b64_inline_neg1(self):
-    """S_MOV_B64 with inline constant -1 should sign-extend."""
-    instructions = [
-      s_mov_b64(s[0:1], -1),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0xFFFFFFFF)
-    self.assertEqual(st.vgpr[0][1], 0xFFFFFFFF)
-
-  def test_s_mov_b64_inline_neg16(self):
-    """S_MOV_B64 with inline constant -16 should sign-extend."""
-    instructions = [
-      s_mov_b64(s[0:1], -16),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0xFFFFFFF0)
-    self.assertEqual(st.vgpr[0][1], 0xFFFFFFFF)
-
-  def test_s_mov_b64_float_const_1_0(self):
-    """S_MOV_B64 with float inline constant 1.0 - casts F32 to F64."""
-    instructions = [
-      s_mov_b64(s[0:1], 1.0),  # inline constant 242 (1.0f)
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # Hardware casts F32 to F64: 1.0f64 = 0x3FF0000000000000
-    self.assertEqual(st.vgpr[0][0], 0x00000000)  # lo
-    self.assertEqual(st.vgpr[0][1], 0x3FF00000)  # hi
-
-  def test_s_or_b64_inline_constant(self):
-    """S_OR_B64 with 64-bit inline constant."""
-    instructions = [
-      s_mov_b64(s[0:1], 0),
-      s_or_b64(s[2:3], s[0:1], 16),
-      v_mov_b32_e32(v[0], s[2]),
-      v_mov_b32_e32(v[1], s[3]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 16)
-    self.assertEqual(st.vgpr[0][1], 0)
-
-  def test_s_and_b64_inline_constant(self):
-    """S_AND_B64 with 64-bit inline constant."""
-    instructions = [
-      s_mov_b32(s[0], 0xFFFFFFFF),
-      s_mov_b32(s[1], 0xFFFFFFFF),
-      s_and_b64(s[2:3], s[0:1], 16),
-      v_mov_b32_e32(v[0], s[2]),
-      v_mov_b32_e32(v[1], s[3]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 16)
-    self.assertEqual(st.vgpr[0][1], 0)
-
-
-class Test64BitSOPLiterals(unittest.TestCase):
-  """Tests for 64-bit SOP instructions with 32-bit literals.
-
-  Tests the behavior when a 64-bit SOP instruction uses a 32-bit literal
-  (offset 255 in instruction encoding). The literal is zero-extended to 64 bits.
-  """
-
-  def test_s_mov_b64_literal(self):
-    """S_MOV_B64 with 32-bit literal value - zero-extended to 64 bits."""
-    instructions = [
-      s_mov_b64(s[0:1], 0x12345678),  # literal > 64, uses literal encoding
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0x12345678)
-    self.assertEqual(st.vgpr[0][1], 0)
-
-  def test_s_or_b64_literal(self):
-    """S_OR_B64 with 32-bit literal value - zero-extended to 64 bits."""
-    instructions = [
-      s_mov_b64(s[0:1], 0),
-      s_or_b64(s[2:3], s[0:1], 0x12345678),  # literal
-      v_mov_b32_e32(v[0], s[2]),
-      v_mov_b32_e32(v[1], s[3]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0x12345678)
-    self.assertEqual(st.vgpr[0][1], 0)
-
-  def test_s_and_b64_literal(self):
-    """S_AND_B64 with 32-bit literal value - zero-extended to 64 bits."""
-    instructions = [
-      s_mov_b32(s[0], 0xFFFFFFFF),
-      s_mov_b32(s[1], 0xFFFFFFFF),
-      s_and_b64(s[2:3], s[0:1], 0x12345678),  # literal
-      v_mov_b32_e32(v[0], s[2]),
-      v_mov_b32_e32(v[1], s[3]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0x12345678)
-    self.assertEqual(st.vgpr[0][1], 0)
-
-  def test_s_mov_b64_literal_negative(self):
-    """S_MOV_B64 with 0xFFFFFFFF literal - zero-extended (not sign-extended)."""
-    instructions = [
-      s_mov_b64(s[0:1], 0xFFFFFFFF),  # -1 as 32-bit, but zero-extended to 64-bit
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0xFFFFFFFF)
-    self.assertEqual(st.vgpr[0][1], 0)  # zero-extended, not sign-extended
-
-  def test_s_mov_b64_literal_high_bit(self):
-    """S_MOV_B64 with 0x80000000 literal - zero-extended (not sign-extended)."""
-    instructions = [
-      s_mov_b64(s[0:1], 0x80000000),  # high bit set, but zero-extended
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0x80000000)
-    self.assertEqual(st.vgpr[0][1], 0)  # zero-extended, not sign-extended
-
-
-class TestBarrier(unittest.TestCase):
-  """Tests for s_barrier — workgroup synchronization across wavefronts."""
-
-  def test_barrier_cross_wave_lds(self):
-    """Wave 0 writes to LDS, s_barrier, wave 1 reads — verifies cross-wave synchronization.
-
-    64 threads (2 waves of 32). Each thread writes (tid+1) to LDS[tid*4], then after
-    s_barrier, reads LDS[(tid^32)*4] — the value written by the other wave. Without barrier
-    support, wave 1 would read stale/zero LDS values.
-    """
-    instructions = [
-      # v[255] = tid (saved by prologue), copy to v[1]
-      v_mov_b32_e32(v[1], v[255]),
-      # v[2] = tid + 1
-      v_add_nc_u32_e32(v[2], 1, v[1]),
-      # v[3] = tid * 4
-      v_lshlrev_b32_e32(v[3], 2, v[1]),
-      # Store (tid+1) to LDS[tid*4]
-      ds_store_b32(addr=v[3], data0=v[2]),
-      s_waitcnt(lgkmcnt=0),
-      s_barrier(),
-      # Read from the other wave's slot: LDS[(tid^32)*4]
-      v_xor_b32_e32(v[4], 32, v[1]),
-      v_lshlrev_b32_e32(v[5], 2, v[4]),
-      ds_load_b32(addr=v[5], vdst=v[0]),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=64)
-    for tid in range(64):
-      self.assertEqual(st.vgpr[tid][0], (tid ^ 32) + 1, f"tid={tid}")
-
-  def test_barrier_two_phases(self):
-    """Two barriers with three phases — tests multiple barriers in sequence.
-
-    Phase 1: all threads write (tid+100) to LDS[tid*4], barrier.
-    Phase 2: all threads read other wave's value, add 1000, write to LDS[(tid+64)*4], barrier.
-    Phase 3: all threads read the other wave's phase-2 output into v[0].
-    """
-    instructions = [
-      # v[255] = tid (saved by prologue), copy to v[1]
-      v_mov_b32_e32(v[1], v[255]),
-      # v[2] = tid + 100
-      v_add_nc_u32_e32(v[2], 100, v[1]),
-      # v[3] = tid * 4
-      v_lshlrev_b32_e32(v[3], 2, v[1]),
-      # Phase 1: write (tid+100) to LDS[tid*4]
-      ds_store_b32(addr=v[3], data0=v[2]),
-      s_waitcnt(lgkmcnt=0),
-      s_barrier(),
-      # Phase 2: read from other wave, add 1000, write to separate LDS region
-      v_xor_b32_e32(v[4], 32, v[1]),
-      v_lshlrev_b32_e32(v[5], 2, v[4]),
-      ds_load_b32(addr=v[5], vdst=v[6]),
-      s_waitcnt(lgkmcnt=0),
-      v_add_nc_u32_e32(v[7], 0x3e8, v[6]),
-      v_add_nc_u32_e32(v[8], 64, v[1]),
-      v_lshlrev_b32_e32(v[9], 2, v[8]),
-      ds_store_b32(addr=v[9], data0=v[7]),
-      s_waitcnt(lgkmcnt=0),
-      s_barrier(),
-      # Phase 3: read other wave's phase-2 output into v[0]
-      v_add_nc_u32_e32(v[10], 64, v[4]),
-      v_lshlrev_b32_e32(v[11], 2, v[10]),
-      ds_load_b32(addr=v[11], vdst=v[0]),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=64)
-    for tid in range(64):
-      self.assertEqual(st.vgpr[tid][0], tid + 100 + 1000, f"tid={tid}")
-
-
 if __name__ == '__main__':
   unittest.main()

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

@@ -5,7 +5,7 @@ Includes: v_mov_b32, v_cvt_*, v_sin_f32, v_rcp_f32, v_exp_f32, v_rndne_f32,
           v_readfirstlane_b32
 """
 import unittest
-from test.amd.hw.helpers import *
+from extra.assembly.amd.test.hw.helpers import *
 
 class TestMov(unittest.TestCase):
   """Tests for V_MOV_B32."""
@@ -255,6 +255,7 @@ class TestF16Conversions(unittest.TestCase):
 
   def test_v_cvt_f16_f32_small(self):
     """V_CVT_F16_F32 converts small f32 value."""
+    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     instructions = [
       v_mov_b32_e32(v[0], 0.5),
       v_cvt_f16_f32_e32(v[1], v[0]),
@@ -292,6 +293,7 @@ class TestF16Conversions(unittest.TestCase):
 
   def test_v_cvt_f16_f32_reads_full_32bit_source(self):
     """V_CVT_F16_F32 must read full 32-bit f32 source."""
+    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], 0x3fc00000),  # f32 1.5
       v_mov_b32_e32(v[0], s[0]),
@@ -300,7 +302,7 @@ class TestF16Conversions(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][1]
     lo_bits = result & 0xffff
-    self.assertEqual(lo_bits, 0x3e00, f"Expected f16(1.5)=0x3e00, got 0x{lo_bits:04x} ({f16(lo_bits)})")
+    self.assertEqual(lo_bits, 0x3e00, f"Expected f16(1.5)=0x3e00, got 0x{lo_bits:04x} ({_f16(lo_bits)})")
 
   def test_v_cvt_i16_f16_zero(self):
     """V_CVT_I16_F16 converts f16 zero to i16 zero."""
@@ -373,6 +375,7 @@ class TestF64Conversions(unittest.TestCase):
 
   def test_v_cvt_f64_f32_pi(self):
     """V_CVT_F64_F32 converts f32 pi to f64."""
+    import math
     instructions = [
       s_mov_b32(s[0], f2i(3.14159265)),
       v_mov_b32_e32(v[0], s[0]),
@@ -693,6 +696,7 @@ class TestCvtF16Modifiers(unittest.TestCase):
 
   def test_v_cvt_f32_f16_abs_negative(self):
     """V_CVT_F32_F16 with |abs| on negative value."""
+    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_neg1 = f32_to_f16(-1.0)  # 0xbc00
     instructions = [
       s_mov_b32(s[0], f16_neg1),
@@ -705,6 +709,7 @@ class TestCvtF16Modifiers(unittest.TestCase):
 
   def test_v_cvt_f32_f16_abs_positive(self):
     """V_CVT_F32_F16 with |abs| on positive value (should stay positive)."""
+    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_2 = f32_to_f16(2.0)  # 0x4000
     instructions = [
       s_mov_b32(s[0], f16_2),
@@ -717,6 +722,7 @@ class TestCvtF16Modifiers(unittest.TestCase):
 
   def test_v_cvt_f32_f16_neg_positive(self):
     """V_CVT_F32_F16 with neg on positive value."""
+    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_2 = f32_to_f16(2.0)  # 0x4000
     instructions = [
       s_mov_b32(s[0], f16_2),
@@ -729,6 +735,7 @@ class TestCvtF16Modifiers(unittest.TestCase):
 
   def test_v_cvt_f32_f16_neg_negative(self):
     """V_CVT_F32_F16 with neg on negative value (double negative)."""
+    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_neg2 = f32_to_f16(-2.0)  # 0xc000
     instructions = [
       s_mov_b32(s[0], f16_neg2),
@@ -741,6 +748,7 @@ class TestCvtF16Modifiers(unittest.TestCase):
 
   def test_v_cvt_f16_f32_then_pack_for_wmma(self):
     """CVT F32->F16 followed by pack (common WMMA pattern)."""
+    from extra.assembly.amd.test.hw.helpers import _f16
     f32_val = 3.5
     instructions = [
       s_mov_b32(s[0], f2i(f32_val)),
@@ -749,8 +757,8 @@ class TestCvtF16Modifiers(unittest.TestCase):
       v_pack_b32_f16(v[2], v[1], v[1]),  # Pack same value
     ]
     st = run_program(instructions, n_lanes=1)
-    lo = f16(st.vgpr[0][2] & 0xffff)
-    hi = f16((st.vgpr[0][2] >> 16) & 0xffff)
+    lo = _f16(st.vgpr[0][2] & 0xffff)
+    hi = _f16((st.vgpr[0][2] >> 16) & 0xffff)
     self.assertAlmostEqual(lo, f32_val, places=1)
     self.assertAlmostEqual(hi, f32_val, places=1)
 
@@ -796,6 +804,7 @@ class TestConversionRounding(unittest.TestCase):
 
   def test_f16_to_f32_precision(self):
     """F16 to F32 conversion precision."""
+    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_val = f32_to_f16(1.5)
     instructions = [
       s_mov_b32(s[0], f16_val),
@@ -807,6 +816,7 @@ class TestConversionRounding(unittest.TestCase):
 
   def test_f16_denormal_to_f32(self):
     """F16 denormal converts to small positive f32."""
+    from extra.assembly.amd.test.hw.helpers import _f16
     f16_denorm = 0x0001  # Smallest positive f16 denormal
     instructions = [
       v_mov_b32_e32(v[0], f16_denorm),
@@ -1502,132 +1512,5 @@ class TestReciprocalF16(unittest.TestCase):
     self.assertAlmostEqual(result, 0.25, places=2, msg="1/4.0 should be 0.25")
 
 
-class TestCvtNormF16(unittest.TestCase):
-  """Tests for V_CVT_NORM_I16_F16 and V_CVT_NORM_U16_F16."""
-
-  def test_cvt_norm_i16_f16_positive(self):
-    """V_CVT_NORM_I16_F16: f16 1.0 -> i16 max (32767)."""
-    instructions = [
-      s_mov_b32(s[0], f32_to_f16(1.0)),
-      v_mov_b32_e32(v[0], s[0]),
-      v_cvt_norm_i16_f16_e32(v[1], v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][1] & 0xffff
-    self.assertEqual(result, 32767)
-
-  def test_cvt_norm_i16_f16_negative(self):
-    """V_CVT_NORM_I16_F16: f16 -1.0 -> i16 -32767 (0x8001)."""
-    instructions = [
-      s_mov_b32(s[0], f32_to_f16(-1.0)),
-      v_mov_b32_e32(v[0], s[0]),
-      v_cvt_norm_i16_f16_e32(v[1], v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][1] & 0xffff
-    self.assertEqual(result, 0x8001)  # -32767, hardware uses symmetric range
-
-  def test_cvt_norm_i16_f16_zero(self):
-    """V_CVT_NORM_I16_F16: f16 0.0 -> i16 0."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0),
-      v_cvt_norm_i16_f16_e32(v[1], v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][1] & 0xffff
-    self.assertEqual(result, 0)
-
-  def test_cvt_norm_u16_f16_one(self):
-    """V_CVT_NORM_U16_F16: f16 1.0 -> u16 max (65535)."""
-    instructions = [
-      s_mov_b32(s[0], f32_to_f16(1.0)),
-      v_mov_b32_e32(v[0], s[0]),
-      v_cvt_norm_u16_f16_e32(v[1], v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][1] & 0xffff
-    self.assertEqual(result, 65535)
-
-  def test_cvt_norm_u16_f16_half(self):
-    """V_CVT_NORM_U16_F16: f16 0.5 -> u16 ~32768."""
-    instructions = [
-      s_mov_b32(s[0], f32_to_f16(0.5)),
-      v_mov_b32_e32(v[0], s[0]),
-      v_cvt_norm_u16_f16_e32(v[1], v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][1] & 0xffff
-    self.assertAlmostEqual(result, 32768, delta=1)
-
-
-class TestPermlane64(unittest.TestCase):
-  """Tests for V_PERMLANE64_B32 instruction (wave64 cross-half swap)."""
-
-  def test_v_permlane64_b32_is_nop_in_wave32(self):
-    """V_PERMLANE64_B32 is a NOP in wave32 mode.
-
-    Per AMD pcode: "if WAVE32 then s_nop(...) else ... endif"
-    The emulator runs in wave32 mode, so this instruction should not modify registers.
-    """
-    instructions = [
-      v_mov_b32_e32(v[0], 0xCAFEBABE),  # source
-      v_mov_b32_e32(v[1], 0x12345678),  # dest (should be preserved)
-      v_permlane64_b32_e32(v[1], v[0]),  # NOP in wave32
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # Dest register should be unchanged (NOP behavior in wave32)
-    self.assertEqual(st.vgpr[0][1], 0x12345678)
-
-
-class TestSwap(unittest.TestCase):
-  """Tests for V_SWAP_B32 - swap two VGPRs."""
-
-  def test_v_swap_b32_basic(self):
-    """V_SWAP_B32 swaps two VGPR values."""
-    instructions = [
-      v_mov_b32_e32(v[0], 42),
-      v_mov_b32_e32(v[1], 99),
-      v_swap_b32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 99)
-    self.assertEqual(st.vgpr[0][1], 42)
-
-  def test_v_swap_b32_same_reg(self):
-    """V_SWAP_B32 with same src and dst is a no-op."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0xDEADBEEF),
-      v_swap_b32_e32(v[0], v[0]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
-
-  def test_v_swap_b32_multi_lane(self):
-    """V_SWAP_B32 swaps per-lane values independently."""
-    instructions = [
-      # v[0] = lane_id * 10, v[1] = lane_id * 100
-      v_lshlrev_b32_e32(v[0], 1, v[255]),       # v[0] = lane_id * 2
-      v_add_nc_u32_e32(v[0], v[0], v[255]),      # v[0] = lane_id * 3
-      v_mul_u32_u24_e32(v[1], 100, v[255]),      # v[1] = lane_id * 100
-      v_swap_b32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=4)
-    for lane in range(4):
-      self.assertEqual(st.vgpr[lane][0], lane * 100)
-      self.assertEqual(st.vgpr[lane][1], lane * 3)
-
-  def test_v_swap_b32_chain(self):
-    """Two swaps in sequence restore original values."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0xAAAAAAAA),
-      v_mov_b32_e32(v[1], 0x55555555),
-      v_swap_b32_e32(v[0], v[1]),
-      v_swap_b32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA)
-    self.assertEqual(st.vgpr[0][1], 0x55555555)
-
-
 if __name__ == '__main__':
   unittest.main()

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

@@ -5,7 +5,7 @@ Includes: v_add_f32, v_mul_f32, v_and_b32, v_or_b32, v_xor_b32,
           v_add_nc_u32, v_cndmask_b32, v_add_f16, v_mul_f16
 """
 import unittest
-from test.amd.hw.helpers import *
+from extra.assembly.amd.test.hw.helpers import *
 
 class TestBasicArithmetic(unittest.TestCase):
   """Tests for basic arithmetic VOP2 instructions."""
@@ -237,32 +237,6 @@ class TestF16Ops(unittest.TestCase):
     # 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_max_f16_basic(self):
-    """V_MAX_F16 returns the maximum of 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_max_f16_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    self.assertEqual(result, 0x4000, f"Expected 0x4000 (f16 2.0), got 0x{result:04x}")
-
-  def test_v_min_f16_basic(self):
-    """V_MIN_F16 returns the minimum of 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_min_f16_e32(v[2], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][2] & 0xffff
-    self.assertEqual(result, 0x3c00, f"Expected 0x3c00 (f16 1.0), got 0x{result:04x}")
-
   def test_v_fmaak_f16_basic(self):
     """V_FMAAK_F16: d = a * b + K."""
     instructions = [
@@ -836,81 +810,6 @@ class TestCarryOps(unittest.TestCase):
     self.assertEqual(st.vgpr[0][2], 0)  # Overflowed to 0
     self.assertEqual(st.vcc, 1)  # Carry out
 
-  def test_v_add_co_ci_u32_clears_carry(self):
-    """V_ADD_CO_CI_U32: VCC must be updated even when no carry is generated.
-
-    This tests the case where VCC=1 going in (carry-in consumed) but the addition
-    does not overflow, so VCC must be cleared to 0.
-
-    Regression test for: VCC not being written by v_add_co_ci_u32_e32.
-    """
-    instructions = [
-      s_mov_b32(VCC_LO, 1),  # VCC = 1 (carry in)
-      v_mov_b32_e32(v[0], 1),  # S0 = 1
-      v_mov_b32_e32(v[1], 1),  # S1 = 1
-      v_add_co_ci_u32_e32(v[2], v[0], v[1]),  # D0 = 1 + 1 + 1 = 3 (no overflow)
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 3)  # 1 + 1 + 1 = 3
-    self.assertEqual(st.vcc, 0)  # No carry out - VCC must be cleared
-
-  def test_v_add_co_ci_u32_multilane_clears_vcc(self):
-    """V_ADD_CO_CI_U32 with multiple lanes: VCC bits must be updated per-lane.
-
-    When VCC has multiple bits set (one per active lane), and the addition doesn't
-    overflow for any lane, all VCC bits must be cleared.
-
-    Regression test for: VCC not being written by v_add_co_ci_u32_e32 in multi-lane case.
-    """
-    instructions = [
-      s_mov_b32(VCC_LO, 0b11),  # VCC = 0b11 (lanes 0,1 have carry-in)
-      v_mov_b32_e32(v[0], 1),  # S0 = 1 for all lanes
-      v_mov_b32_e32(v[1], 1),  # S1 = 1 for all lanes
-      v_add_co_ci_u32_e32(v[2], v[0], v[1]),  # D0 = 1 + 1 + 1 = 3 (no overflow)
-    ]
-    st = run_program(instructions, n_lanes=2)
-    self.assertEqual(st.vgpr[0][2], 3)  # lane 0: 1 + 1 + 1 = 3
-    self.assertEqual(st.vgpr[1][2], 3)  # lane 1: 1 + 1 + 1 = 3
-    self.assertEqual(st.vcc, 0)  # No carry out for any lane - all VCC bits must be cleared
-
-  def test_v_add_co_ci_u32_preserves_inactive_vcc_bits(self):
-    """V_ADD_CO_CI_U32: VCC carry-out overwrites entire VCC register.
-
-    VOP2 carry instructions write ALL VCC bits based on carry-out, clearing
-    bits for lanes that don't overflow regardless of EXEC mask.
-
-    Note: This differs from VOPC which only writes active lane bits.
-    """
-    instructions = [
-      s_mov_b32(VCC_LO, 0x00010000),  # VCC bit 16 set
-      v_mov_b32_e32(v[0], 1),  # S0 = 1
-      v_mov_b32_e32(v[1], 1),  # S1 = 1
-      v_add_co_ci_u32_e32(v[2], v[0], v[1]),  # D0 = 1 + 1 + 0 = 2 (no carry)
-    ]
-    st = run_program(instructions, n_lanes=4)
-    self.assertEqual(st.vgpr[0][2], 2)  # lane 0: 1 + 1 + 0 = 2
-    # VCC should be completely cleared (all lanes have no carry-out)
-    self.assertEqual(st.vcc, 0)
-
-  def test_v_add_co_ci_u32_all_lanes_same_result(self):
-    """V_ADD_CO_CI_U32: all active lanes should produce the same result.
-
-    When the same constant inputs are used across all lanes, each lane should
-    compute the same result and write to its own VGPR slot.
-
-    Regression test for: VGPR writes not happening for all lanes.
-    """
-    instructions = [
-      s_mov_b32(VCC_LO, 0),  # No carry-in
-      v_mov_b32_e32(v[0], 3),  # inline constant 3
-      v_mov_b32_e32(v[1], 5),  # value 5
-      v_add_co_ci_u32_e32(v[1], 3, v[1]),  # v[1] = 3 + v[1] + 0 = 3 + 5 = 8
-    ]
-    st = run_program(instructions, n_lanes=4)
-    # All 4 lanes should have v[1] = 8
-    for lane in range(4):
-      self.assertEqual(st.vgpr[lane][1], 8, f"lane {lane} should have v[1]=8")
-
   def test_v_sub_co_ci_u32_no_borrow(self):
     """V_SUB_CO_CI_U32: D0 = S0 - S1 - VCC_IN, when VCC_IN=0."""
     instructions = [
@@ -961,23 +860,6 @@ class TestCarryOps(unittest.TestCase):
     self.assertEqual(st.vgpr[0][0], 16)
     self.assertEqual(st.sgpr[10], 0)  # No carry out
 
-  def test_v_add_co_ci_u32_vop3sd_null_sdst(self):
-    """VOP3SD V_ADD_CO_CI_U32 with sdst=NULL: carry output is discarded.
-
-    When sdst=NULL (register 124), the carry-out should NOT be written anywhere.
-    We verify this by checking that VCC (which we set to a sentinel value) is unchanged.
-    """
-    instructions = [
-      s_mov_b32(VCC_LO, 0xDEADBEEF),  # Sentinel value in VCC
-      s_mov_b32(s[6], 0),  # carry-in = 0
-      # VOP3SD with NULL sdst: carry-out should be discarded
-      # Uses 0xFFFFFFFF + 1 + 0 = 0 with carry-out=1, but carry should not be written
-      v_add_co_ci_u32(v[0], NULL, 0xFFFFFFFF, 1, s[6]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0)  # 0xFFFFFFFF + 1 + 0 = 0 (overflow)
-    self.assertEqual(st.vcc, 0xDEADBEEF)  # VCC unchanged - carry was discarded
-
 
 if __name__ == '__main__':
   unittest.main()

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

@@ -3,7 +3,7 @@
 Includes: v_pk_add_f16, v_pk_mul_f16, v_pk_fma_f16, v_pack_b32_f16, v_wmma_*, v_dot2_*
 """
 import unittest
-from test.amd.hw.helpers import *
+from extra.assembly.amd.test.hw.helpers import *
 
 class TestPackInstructions(unittest.TestCase):
   """Tests for pack instructions."""
@@ -149,6 +149,7 @@ class TestFmaMix(unittest.TestCase):
 
   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.test.hw.helpers import f32_to_f16
     f16_2 = f32_to_f16(2.0)
     instructions = [
       s_mov_b32(s[0], f2i(1.0)),
@@ -165,6 +166,7 @@ class TestFmaMix(unittest.TestCase):
 
   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.test.hw.helpers import f32_to_f16
     f16_2 = f32_to_f16(2.0)
     val = (f16_2 << 16) | 0
     instructions = [
@@ -197,6 +199,7 @@ class TestFmaMix(unittest.TestCase):
 
   def test_v_fma_mix_f32_with_abs_f16_src2_lo(self):
     """V_FMA_MIX_F32 with abs modifier on f16 src2 (lo half). Regression test for sin(1.0) bug."""
+    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_neg1 = f32_to_f16(-1.0)  # 0xbc00
     instructions = [
       s_mov_b32(s[0], f2i(0.0)),  # src0 = 0.0 (f32)
@@ -214,6 +217,7 @@ class TestFmaMix(unittest.TestCase):
 
   def test_v_fma_mix_f32_with_neg_f16_src2_lo(self):
     """V_FMA_MIX_F32 with neg modifier on f16 src2 (lo half)."""
+    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_1 = f32_to_f16(1.0)  # 0x3c00
     instructions = [
       s_mov_b32(s[0], f2i(0.0)),  # src0 = 0.0 (f32)
@@ -231,6 +235,7 @@ class TestFmaMix(unittest.TestCase):
 
   def test_v_fma_mix_f32_with_abs_f16_src2_hi(self):
     """V_FMA_MIX_F32 with abs modifier on f16 src2 (hi half)."""
+    from extra.assembly.amd.test.hw.helpers import f32_to_f16
     f16_neg1 = f32_to_f16(-1.0)  # 0xbc00
     val = (f16_neg1 << 16) | 0  # -1.0 in hi, 0 in lo
     instructions = [
@@ -249,6 +254,7 @@ class TestFmaMix(unittest.TestCase):
 
   def test_v_fma_mixlo_f16(self):
     """V_FMA_MIXLO_F16 writes to low 16 bits of destination."""
+    from extra.assembly.amd.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], f2i(2.0)),
       v_mov_b32_e32(v[0], s[0]),
@@ -261,13 +267,14 @@ class TestFmaMix(unittest.TestCase):
       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)
+    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.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], f2i(1.0)),
       v_mov_b32_e32(v[0], s[0]),
@@ -279,12 +286,13 @@ class TestFmaMix(unittest.TestCase):
       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)
+    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.test.hw.helpers import _f16
     instructions = [
       s_mov_b32(s[0], 0x3f800000),  # f32 1.0
       v_mov_b32_e32(v[3], s[0]),
@@ -297,7 +305,7 @@ class TestFmaMix(unittest.TestCase):
       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)
+    lo = _f16(st.vgpr[0][3] & 0xffff)
     self.assertAlmostEqual(lo, -3.14159, delta=0.01)
 
 
@@ -306,6 +314,7 @@ class TestVOP3P(unittest.TestCase):
 
   def test_v_pk_add_f16_basic(self):
     """V_PK_ADD_F16 adds two packed f16 values."""
+    from extra.assembly.amd.test.hw.helpers 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
@@ -315,13 +324,14 @@ class TestVOP3P(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][2]
-    lo = f16(result & 0xffff)
-    hi = f16((result >> 16) & 0xffff)
+    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.test.hw.helpers 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
@@ -331,13 +341,14 @@ class TestVOP3P(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][2]
-    lo = f16(result & 0xffff)
-    hi = f16((result >> 16) & 0xffff)
+    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.test.hw.helpers 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
@@ -349,8 +360,8 @@ class TestVOP3P(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][3]
-    lo = f16(result & 0xffff)
-    hi = f16((result >> 16) & 0xffff)
+    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
 
@@ -359,6 +370,7 @@ class TestVOP3P(unittest.TestCase):
     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.test.hw.helpers 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]),
@@ -366,8 +378,8 @@ class TestVOP3P(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][1]
-    lo = f16(result & 0xffff)
-    hi = f16((result >> 16) & 0xffff)
+    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)
@@ -376,6 +388,7 @@ class TestVOP3P(unittest.TestCase):
     """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.test.hw.helpers 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 = [
@@ -385,141 +398,18 @@ class TestVOP3P(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][1]
-    lo = f16(result & 0xffff)
-    hi = f16((result >> 16) & 0xffff)
+    lo = _f16(result & 0xffff)
+    hi = _f16((result >> 16) & 0xffff)
     self.assertAlmostEqual(lo, 6.0, places=1)
     self.assertAlmostEqual(hi, 0.0, places=1)
 
-  def test_v_pk_add_u16_float_inline_const_opsel(self):
-    """V_PK_ADD_U16 with float inline constant 2.0
-    Regression test: for integer packed ops, do not perform the f32->f16 conversion.
-    """
-    # src1 = inline float constant 2.0
-    instructions = [
-      s_mov_b32(s[0], 0x00030005),  # packed u16: hi=3, lo=5
-      v_mov_b32_e32(v[0], s[0]),
-      v_pk_add_u16(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 = result & 0xffff
-    hi = (result >> 16) & 0xffff
-    # lo = 5 + 0x0000 = 0x0005, hi = 3 + 0x4000 = 0x4003
-    self.assertEqual(lo, 0x0005, f"lo: expected 0x0005, got 0x{lo:04x}")
-    self.assertEqual(hi, 0x4003, f"hi: expected 0x4003, got 0x{hi:04x}")
-
-  def test_v_pk_add_u16_literal_constant(self):
-    """V_PK_ADD_U16 with a literal constant (value > 64, requires VOP3P_LIT encoding).
-    Regression test: VOP3P literal constants were not passed to rsrc_dyn, so literal src read as 0.
-    """
-    instructions = [
-      s_mov_b32(s[0], 0x1C001C00),  # packed u16: hi=0x1C00, lo=0x1C00 (f16 for 2^-8)
-      v_mov_b32_e32(v[0], s[0]),
-      v_pk_add_u16(v[1], 0x2000, v[0], opsel_hi=2, opsel_hi2=1),  # add 0x2000 bias to both halves
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][1]
-    lo = result & 0xffff
-    hi = (result >> 16) & 0xffff
-    # lo = 0x1C00 + 0x2000 = 0x3C00 (f16 1.0), hi = 0x1C00 + 0x2000 = 0x3C00 (f16 1.0)
-    self.assertEqual(lo, 0x3C00, f"lo: expected 0x3C00, got 0x{lo:04x}")
-    self.assertEqual(hi, 0x3C00, f"hi: expected 0x3C00, got 0x{hi:04x}")
-
-
-class TestWMMAF16(unittest.TestCase):
-  """Tests for WMMA F16 output variant (V_WMMA_F16_16X16X16_F16).
-
-  Note: RDNA3 WMMA F16 uses 8 VGPRs for accumulator/output (same as F32 variant),
-  but values are packed as f16. This differs from RDNA4 which uses 4 VGPRs.
-  """
-
-  def test_v_wmma_f16_16x16x16_f16_all_ones(self):
-    """V_WMMA_F16_16X16X16_F16 with all ones produces 16.0 in f16."""
-    instructions: list[Inst] = []
-    instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
-    # Initialize A matrix in v[16:23] (8 regs)
-    for i in range(16, 24):
-      instructions.append(v_mov_b32_e32(v[i], s[0]))
-    # Initialize B matrix in v[24:31] (8 regs)
-    for i in range(24, 32):
-      instructions.append(v_mov_b32_e32(v[i], s[0]))
-    # Initialize C (accumulator) in v[0:7] to zero (8 regs for RDNA3 WMMA F16)
-    for i in range(8):
-      instructions.append(v_mov_b32_e32(v[i], 0))
-    # WMMA F16: D = A @ B + C
-    instructions.append(v_wmma_f16_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
-    st = run_program(instructions, n_lanes=32)
-    # Result should be 16.0 in f16, stored in lo 16 bits of each VGPR (hi bits are 0)
-    for lane in range(32):
-      for reg in range(8):
-        result = st.vgpr[lane][reg]
-        lo = f16(result & 0xffff)
-        self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
-        self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
-
-  def test_v_wmma_f16_16x16x16_f16_with_accumulator(self):
-    """V_WMMA_F16_16X16X16_F16 with non-zero accumulator."""
-    instructions: list[Inst] = []
-    instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
-    instructions.append(s_mov_b32(s[1], 0x4500))  # f16 5.0 in lo bits only
-    # Initialize A matrix in v[16:23] (8 regs)
-    for i in range(16, 24):
-      instructions.append(v_mov_b32_e32(v[i], s[0]))
-    # Initialize B matrix in v[24:31] (8 regs)
-    for i in range(24, 32):
-      instructions.append(v_mov_b32_e32(v[i], s[0]))
-    # Initialize C (accumulator) in v[0:7] to 5.0 in lo bits (8 regs for RDNA3 WMMA F16)
-    for i in range(8):
-      instructions.append(v_mov_b32_e32(v[i], s[1]))
-    # WMMA F16: D = A @ B + C
-    instructions.append(v_wmma_f16_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
-    st = run_program(instructions, n_lanes=32)
-    # Result should be 16.0 + 5.0 = 21.0 in f16, stored in lo 16 bits (hi bits are 0)
-    for lane in range(32):
-      for reg in range(8):
-        result = st.vgpr[lane][reg]
-        lo = f16(result & 0xffff)
-        self.assertAlmostEqual(lo, 21.0, places=0, msg=f"v[{reg}] lane {lane}: expected 21.0, got {lo}")
-        self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
-
-  def test_v_wmma_f16_16x16x16_f16_high_registers(self):
-    """V_WMMA_F16_16X16X16_F16 with high register indices.
-
-    Regression test: WMMA was using static register indices instead of dynamic.
-    This test uses v[64:71] for A, v[80:87] for B, v[96:103] for C/D.
-    """
-    instructions: list[Inst] = []
-    instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
-    # Initialize A matrix in v[64:71] (8 regs)
-    for i in range(64, 72):
-      instructions.append(v_mov_b32_e32(v[i], s[0]))
-    # Initialize B matrix in v[80:87] (8 regs)
-    for i in range(80, 88):
-      instructions.append(v_mov_b32_e32(v[i], s[0]))
-    # Initialize C (accumulator) in v[96:103] to zero (8 regs for RDNA3 WMMA F16)
-    for i in range(96, 104):
-      instructions.append(v_mov_b32_e32(v[i], 0))
-    # WMMA F16: D = A @ B + C, result in v[96:103]
-    instructions.append(v_wmma_f16_16x16x16_f16(v[96:103], v[64:71], v[80:87], v[96:103]))
-    # Copy results to v[0:7] for checking
-    for i in range(8):
-      instructions.append(v_mov_b32_e32(v[i], v[96+i]))
-    st = run_program(instructions, n_lanes=32)
-    # Result should be 16.0 in f16, stored in lo 16 bits (hi bits are 0)
-    for lane in range(32):
-      for reg in range(8):
-        result = st.vgpr[lane][reg]
-        lo = f16(result & 0xffff)
-        self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
-        self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
-
 
 class TestWMMA(unittest.TestCase):
-  """Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions with F32 output."""
+  """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: list[Inst] = []
+    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]))
@@ -535,7 +425,7 @@ class TestWMMA(unittest.TestCase):
 
   def test_v_wmma_f32_16x16x16_f16_with_accumulator(self):
     """V_WMMA_F32_16X16X16_F16 with non-zero accumulator."""
-    instructions: list[Inst] = []
+    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):
@@ -550,75 +440,6 @@ class TestWMMA(unittest.TestCase):
         result = st.vgpr[lane][reg]
         self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
 
-  def test_v_wmma_f32_16x16x16_f16_high_registers(self):
-    """V_WMMA_F32_16X16X16_F16 with high register indices.
-
-    Regression test: WMMA was using static register indices instead of dynamic,
-    causing incorrect results when registers weren't at the default positions.
-    This test uses v[64:71] for A, v[80:87] for B, v[96:103] for C/D.
-    """
-    instructions: list[Inst] = []
-    instructions.append(s_mov_b32(s[0], 0x3c003c00))  # packed f16 1.0
-    # Initialize A matrix in v[64:71]
-    for i in range(64, 72):
-      instructions.append(v_mov_b32_e32(v[i], s[0]))
-    # Initialize B matrix in v[80:87]
-    for i in range(80, 88):
-      instructions.append(v_mov_b32_e32(v[i], s[0]))
-    # Initialize C (accumulator) in v[96:103] to zero
-    for i in range(96, 104):
-      instructions.append(v_mov_b32_e32(v[i], 0))
-    # WMMA: D = A @ B + C, result in v[96:103]
-    instructions.append(v_wmma_f32_16x16x16_f16(v[96:103], v[64:71], v[80:87], v[96:103]))
-    # Copy results to v[0:7] for checking
-    for i in range(8):
-      instructions.append(v_mov_b32_e32(v[i], v[96+i]))
-    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)}")
-
-
-class TestWMMABF16(unittest.TestCase):
-  """Tests for WMMA BF16 instructions."""
-
-  def test_v_wmma_f32_16x16x16_bf16_all_ones(self):
-    """V_WMMA_F32_16X16X16_BF16 with all ones produces 16.0."""
-    instructions: list[Inst] = []
-    # BF16 1.0 = 0x3f80, packed = 0x3f803f80
-    instructions.append(s_mov_b32(s[0], 0x3f803f80))
-    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_bf16(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_bf16_with_accumulator(self):
-    """V_WMMA_F32_16X16X16_BF16 with non-zero accumulator."""
-    instructions: list[Inst] = []
-    # BF16 1.0 = 0x3f80, packed = 0x3f803f80
-    instructions.append(s_mov_b32(s[0], 0x3f803f80))
-    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_bf16(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)."""
@@ -732,6 +553,7 @@ class TestPackedMixedSigns(unittest.TestCase):
 
   def test_pk_add_f16_mixed_signs(self):
     """V_PK_ADD_F16 with mixed positive/negative values."""
+    from extra.assembly.amd.test.hw.helpers 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
@@ -741,13 +563,14 @@ class TestPackedMixedSigns(unittest.TestCase):
     ]
     st = run_program(instructions, n_lanes=1)
     result = st.vgpr[0][2]
-    lo = f16(result & 0xffff)
-    hi = f16((result >> 16) & 0xffff)
+    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.test.hw.helpers 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
@@ -760,277 +583,5 @@ class TestPackedMixedSigns(unittest.TestCase):
     self.assertEqual(result, 0x00000000, "2.0 * 0.0 should be 0.0")
 
 
-class TestDot2F32F16(unittest.TestCase):
-  """Tests for V_DOT2_F32_F16 - dot product of f16 pairs producing f32."""
-
-  def test_v_dot2_f32_f16_basic(self):
-    """V_DOT2_F32_F16: dot product of two packed f16 pairs -> f32."""
-    # src0 = {hi=2.0, lo=1.0}, src1 = {hi=4.0, lo=3.0}
-    # result = 1.0*3.0 + 2.0*4.0 + 0 = 3 + 8 = 11.0
-    src0 = (f32_to_f16(2.0) << 16) | f32_to_f16(1.0)
-    src1 = (f32_to_f16(4.0) << 16) | f32_to_f16(3.0)
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], src1),
-      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_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = i2f(st.vgpr[0][3])
-    self.assertAlmostEqual(result, 11.0, places=2)
-
-  def test_v_dot2_f32_f16_with_accumulator(self):
-    """V_DOT2_F32_F16 with non-zero f32 accumulator."""
-    # src0 = {hi=1.0, lo=1.0}, src1 = {hi=1.0, lo=1.0}, acc = 5.0
-    # result = 1.0*1.0 + 1.0*1.0 + 5.0 = 7.0
-    src0 = (f32_to_f16(1.0) << 16) | f32_to_f16(1.0)
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], f2i(5.0)),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[0]),  # same as src0
-      v_mov_b32_e32(v[2], s[1]),
-      v_dot2_f32_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = i2f(st.vgpr[0][3])
-    self.assertAlmostEqual(result, 7.0, places=2)
-
-  def test_v_dot2_f32_f16_negative_values(self):
-    """V_DOT2_F32_F16 with negative f16 values."""
-    # src0 = {hi=-2.0, lo=3.0}, src1 = {hi=1.0, lo=2.0}
-    # result = 3.0*2.0 + (-2.0)*1.0 + 0 = 6 - 2 = 4.0
-    # NOTE: Hardware DOT2 may have up to 1 ULP difference due to internal implementation
-    src0 = (f32_to_f16(-2.0) << 16) | f32_to_f16(3.0)
-    src1 = (f32_to_f16(1.0) << 16) | f32_to_f16(2.0)
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], src1),
-      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_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
-    ]
-    st = run_program(instructions, n_lanes=1, ulp_tolerance=1)
-    result = i2f(st.vgpr[0][3])
-    self.assertAlmostEqual(result, 4.0, places=2)
-
-
-class TestDot2F16F16(unittest.TestCase):
-  """Tests for V_DOT2_F16_F16 - dot product of f16 pairs producing f16."""
-
-  def test_v_dot2_f16_f16_basic(self):
-    """V_DOT2_F16_F16: dot product of two packed f16 pairs -> f16."""
-    # src0 = {hi=2.0, lo=1.0}, src1 = {hi=3.0, lo=2.0}
-    # result = 1.0*2.0 + 2.0*3.0 + 0 = 2 + 6 = 8.0 (f16)
-    src0 = (f32_to_f16(2.0) << 16) | f32_to_f16(1.0)
-    src1 = (f32_to_f16(3.0) << 16) | f32_to_f16(2.0)
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], src1),
-      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_f16_f16(v[3], v[0], v[1], v[2]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = f16(st.vgpr[0][3] & 0xffff)
-    self.assertAlmostEqual(result, 8.0, places=1)
-
-  def test_v_dot2_f16_f16_with_accumulator(self):
-    """V_DOT2_F16_F16 with non-zero f16 accumulator."""
-    # src0 = {hi=1.0, lo=1.0}, src1 = {hi=1.0, lo=1.0}, acc = 3.0 (f16)
-    # result = 1.0*1.0 + 1.0*1.0 + 3.0 = 5.0 (f16)
-    src0 = (f32_to_f16(1.0) << 16) | f32_to_f16(1.0)
-    acc = f32_to_f16(3.0)
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[2], acc),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[0]),  # same as src0
-      v_mov_b32_e32(v[2], s[2]),
-      v_dot2_f16_f16(v[3], v[0], v[1], v[2]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = f16(st.vgpr[0][3] & 0xffff)
-    self.assertAlmostEqual(result, 5.0, places=1)
-
-
-class TestSignedDotProducts(unittest.TestCase):
-  """Tests for V_DOT4_I32_IU8 and V_DOT8_I32_IU4 with signed inputs."""
-
-  def test_v_dot4_i32_iu8_signed_both(self):
-    """V_DOT4_I32_IU8 with both inputs signed (neg=0b011)."""
-    # src0 = {-1, -2, 3, 4} as i8 = {0xff, 0xfe, 0x03, 0x04}
-    # src1 = {1, 1, 1, 1} as i8
-    # result = (-1)*1 + (-2)*1 + 3*1 + 4*1 = -1 - 2 + 3 + 4 = 4
-    src0 = (0xff << 24) | (0xfe << 16) | (0x03 << 8) | 0x04  # -1, -2, 3, 4
-    src1 = 0x01010101  # 1, 1, 1, 1
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], src1),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], 0),
-      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b011),  # both signed
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][3]
-    # Result is i32, interpret as signed
-    if result >= 0x80000000:
-      result = result - 0x100000000
-    self.assertEqual(result, 4)
-
-  def test_v_dot4_i32_iu8_src0_signed(self):
-    """V_DOT4_I32_IU8 with only src0 signed (neg=0b001)."""
-    # src0 = {-1, -1, -1, -1} as i8 = {0xff, 0xff, 0xff, 0xff}
-    # src1 = {2, 2, 2, 2} as u8
-    # result = (-1)*2 + (-1)*2 + (-1)*2 + (-1)*2 = -8
-    src0 = 0xffffffff  # -1, -1, -1, -1 (as i8)
-    src1 = 0x02020202  # 2, 2, 2, 2 (as u8)
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], src1),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], 0),
-      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b001),  # src0 signed
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][3]
-    if result >= 0x80000000:
-      result = result - 0x100000000
-    self.assertEqual(result, -8)
-
-  def test_v_dot4_i32_iu8_src1_signed(self):
-    """V_DOT4_I32_IU8 with only src1 signed (neg=0b010)."""
-    # src0 = {2, 2, 2, 2} as u8
-    # src1 = {-1, -1, -1, -1} as i8 = {0xff, 0xff, 0xff, 0xff}
-    # result = 2*(-1) + 2*(-1) + 2*(-1) + 2*(-1) = -8
-    src0 = 0x02020202  # 2, 2, 2, 2 (as u8)
-    src1 = 0xffffffff  # -1, -1, -1, -1 (as i8)
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], src1),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], 0),
-      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b010),  # src1 signed
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][3]
-    if result >= 0x80000000:
-      result = result - 0x100000000
-    self.assertEqual(result, -8)
-
-  def test_v_dot4_i32_iu8_unsigned_as_reference(self):
-    """V_DOT4_I32_IU8 with both unsigned (neg=0) - same as V_DOT4_U32_U8."""
-    # src0 = {0xff, 0xff, 0xff, 0xff} = 255 each as u8
-    # src1 = {1, 1, 1, 1}
-    # result = 255*1 + 255*1 + 255*1 + 255*1 = 1020
-    src0 = 0xffffffff
-    src1 = 0x01010101
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], src1),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], 0),
-      v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0),  # both unsigned
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][3], 1020)
-
-  def test_v_dot8_i32_iu4_signed_both(self):
-    """V_DOT8_I32_IU4 with both inputs signed (neg=0b011)."""
-    # src0 = 8 nibbles: {-1, -2, 3, 4, -1, -2, 3, 4} as i4
-    # i4 -1 = 0xf, -2 = 0xe, 3 = 0x3, 4 = 0x4
-    # src0 = 0xfe34fe34
-    # src1 = {1, 1, 1, 1, 1, 1, 1, 1} as i4 = 0x11111111
-    # result = 2 * ((-1)*1 + (-2)*1 + 3*1 + 4*1) = 2 * 4 = 8
-    src0 = 0xfe34fe34
-    src1 = 0x11111111
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], src1),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], 0),
-      v_dot8_i32_iu4(v[3], v[0], v[1], v[2], neg=0b011),  # both signed
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][3]
-    if result >= 0x80000000:
-      result = result - 0x100000000
-    self.assertEqual(result, 8)
-
-  def test_v_dot8_i32_iu4_all_negative(self):
-    """V_DOT8_I32_IU4 with all negative signed values."""
-    # src0 = 8 nibbles all -1 (0xf) = 0xffffffff
-    # src1 = 8 nibbles all 1 = 0x11111111
-    # result = 8 * ((-1)*1) = -8
-    src0 = 0xffffffff  # all -1 as i4
-    src1 = 0x11111111  # all 1
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], src1),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_mov_b32_e32(v[2], 0),
-      v_dot8_i32_iu4(v[3], v[0], v[1], v[2], neg=0b011),  # both signed
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][3]
-    if result >= 0x80000000:
-      result = result - 0x100000000
-    self.assertEqual(result, -8)
-
-
-class TestPkMinMaxF16(unittest.TestCase):
-  """Tests for V_PK_MIN_F16 and V_PK_MAX_F16."""
-
-  def test_v_pk_min_f16_basic(self):
-    """V_PK_MIN_F16: packed min of two f16 pairs."""
-    # src0 = {hi=3.0, lo=1.0}, src1 = {hi=2.0, lo=4.0}
-    # result = {min(3,2)=2, min(1,4)=1}
-    src0 = (f32_to_f16(3.0) << 16) | f32_to_f16(1.0)
-    src1 = (f32_to_f16(2.0) << 16) | f32_to_f16(4.0)
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], src1),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pk_min_f16(v[2], v[0], v[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, 1.0, delta=0.01)
-    self.assertAlmostEqual(hi, 2.0, delta=0.01)
-
-  def test_v_pk_max_f16_basic(self):
-    """V_PK_MAX_F16: packed max of two f16 pairs."""
-    # src0 = {hi=3.0, lo=1.0}, src1 = {hi=2.0, lo=4.0}
-    # result = {max(3,2)=3, max(1,4)=4}
-    src0 = (f32_to_f16(3.0) << 16) | f32_to_f16(1.0)
-    src1 = (f32_to_f16(2.0) << 16) | f32_to_f16(4.0)
-    instructions = [
-      s_mov_b32(s[0], src0),
-      s_mov_b32(s[1], src1),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], s[1]),
-      v_pk_max_f16(v[2], v[0], v[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, delta=0.01)
-    self.assertAlmostEqual(hi, 3.0, delta=0.01)
-
-
 if __name__ == '__main__':
   unittest.main()

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

@@ -3,7 +3,7 @@
 Includes: v_cmp_class_f32, v_cmp_class_f16, v_cmp_eq_*, v_cmp_lt_*, v_cmp_gt_*
 """
 import unittest
-from test.amd.hw.helpers import *
+from extra.assembly.amd.test.hw.helpers import *
 
 VCC = 106  # SGPR index for VCC_LO
 
@@ -104,34 +104,6 @@ class TestCmpClass(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     self.assertEqual(st.vcc & 1, 0, "Signaling NaN should not match quiet mask")
 
-  def test_v_cmp_lg_f32_nan(self):
-    """v_cmp_lg_f32 is ordered not-equal (<>): NaN <> x should be False per IEEE 754."""
-    quiet_nan = 0x7fc00000
-    one_f32 = 0x3f800000  # 1.0f
-    instructions = [
-      s_mov_b32(s[0], quiet_nan),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], one_f32),
-      v_mov_b32_e32(v[1], s[1]),
-      v_cmp_lg_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 0, "v_cmp_lg_f32(NaN, 1.0) should be 0")
-
-  def test_v_cmp_neq_f32_nan(self):
-    """v_cmp_neq_f32 is unordered not-equal (!=): NaN != x should be True per IEEE 754."""
-    quiet_nan = 0x7fc00000
-    one_f32 = 0x3f800000  # 1.0f
-    instructions = [
-      s_mov_b32(s[0], quiet_nan),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], one_f32),
-      v_mov_b32_e32(v[1], s[1]),
-      v_cmp_neq_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "v_cmp_neq_f32(NaN, 1.0) should be 1")
-
   def test_v_cmp_sets_vcc_bits(self):
     """V_CMP_EQ sets VCC bits based on per-lane comparison."""
     instructions = [
@@ -759,111 +731,6 @@ class TestVCCBehavior(unittest.TestCase):
     self.assertEqual(st.vcc >> 16, 0x0000, "Lanes 16-31 should be false")
 
 
-class TestCmpNge(unittest.TestCase):
-  """Tests for V_CMP_NGE (not-greater-or-equal) with NaN semantics.
-
-  NGE = !(a >= b). With NaN inputs:
-  - If either input is NaN, a >= b is false, so !(false) = true
-  - This differs from a < b which returns false for NaN inputs
-  """
-
-  def test_v_cmp_nge_f32_normal_values(self):
-    """v_cmp_nge_f32: basic comparison with normal floats."""
-    instructions = [
-      v_mov_b32_e32(v[0], f2i(1.0)),
-      v_mov_b32_e32(v[1], f2i(2.0)),
-      v_cmp_nge_f32_e32(v[0], v[1]),  # !(1.0 >= 2.0) = !(false) = true
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "!(1.0 >= 2.0) should be true")
-
-  def test_v_cmp_nge_f32_equal_values(self):
-    """v_cmp_nge_f32: equal values should return false."""
-    instructions = [
-      v_mov_b32_e32(v[0], f2i(1.0)),
-      v_mov_b32_e32(v[1], f2i(1.0)),
-      v_cmp_nge_f32_e32(v[0], v[1]),  # !(1.0 >= 1.0) = !(true) = false
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 0, "!(1.0 >= 1.0) should be false")
-
-  def test_v_cmp_nge_f32_greater_value(self):
-    """v_cmp_nge_f32: greater value should return false."""
-    instructions = [
-      v_mov_b32_e32(v[0], f2i(2.0)),
-      v_mov_b32_e32(v[1], f2i(1.0)),
-      v_cmp_nge_f32_e32(v[0], v[1]),  # !(2.0 >= 1.0) = !(true) = false
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 0, "!(2.0 >= 1.0) should be false")
-
-  def test_v_cmp_nge_f32_neg_inf(self):
-    """v_cmp_nge_f32: -inf compared to normal value."""
-    neg_inf = 0xff800000  # -inf
-    instructions = [
-      s_mov_b32(s[0], neg_inf),
-      v_mov_b32_e32(v[0], s[0]),
-      v_mov_b32_e32(v[1], f2i(1.0)),
-      v_cmp_nge_f32_e32(v[0], v[1]),  # !(-inf >= 1.0) = !(false) = true
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "!(-inf >= 1.0) should be true")
-
-  def test_v_cmp_nge_f32_clears_inactive_vcc_bits(self):
-    """v_cmp_nge_f32 with partial EXEC clears inactive VCC bits (hardware behavior)."""
-    neg_inf = 0xff800000  # -inf
-    instructions = [
-      # Set VCC to all 1s first
-      s_mov_b32(VCC_LO, 0xFFFFFFFF),
-      # Set EXEC to only lane 0
-      s_mov_b32(EXEC_LO, 0x00000001),
-      # v0 = 1.0 for lane 0
-      v_mov_b32_e32(v[0], f2i(1.0)),
-      # Compare: !(-inf >= 1.0) = true for lane 0
-      v_cmp_nge_f32_e32(neg_inf, v[0]),
-    ]
-    st = run_program(instructions, n_lanes=16)
-    # Hardware clears inactive lane bits, only active lane results remain
-    # Lane 0 result = 1 (true), lanes 1-15 = 0 (cleared)
-    self.assertEqual(st.vcc, 0x00000001, "VCC should only have active lane results")
-
-  def test_v_cmp_nge_f32_nan_src0(self):
-    """v_cmp_nge_f32: NaN in src0 should return true (NaN >= x is false)."""
-    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], f2i(1.0)),
-      v_cmp_nge_f32_e32(v[0], v[1]),  # !(NaN >= 1.0) = !(false) = true
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "!(NaN >= 1.0) should be true")
-
-  def test_v_cmp_nge_f32_nan_src1(self):
-    """v_cmp_nge_f32: NaN in src1 should return true (x >= NaN is false)."""
-    quiet_nan = 0x7fc00000
-    instructions = [
-      s_mov_b32(s[0], quiet_nan),
-      v_mov_b32_e32(v[0], f2i(1.0)),
-      v_mov_b32_e32(v[1], s[0]),
-      v_cmp_nge_f32_e32(v[0], v[1]),  # !(1.0 >= NaN) = !(false) = true
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "!(1.0 >= NaN) should be true")
-
-  def test_v_cmp_nge_f32_both_nan(self):
-    """v_cmp_nge_f32: both NaN should return true."""
-    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], s[0]),
-      v_cmp_nge_f32_e32(v[0], v[1]),  # !(NaN >= NaN) = !(false) = true
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "!(NaN >= NaN) should be true")
-
-
 class TestCmpxPartialWavefront(unittest.TestCase):
   """Tests for V_CMPX with partial wavefronts (fewer than 32 active lanes).
 

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

@@ -7,8 +7,9 @@ VOPD executes two operations simultaneously. Key behavior:
 - Op Y can use ops 0-18 (includes ADD_NC_U32, LSHLREV, AND)
 """
 import unittest
-from test.amd.hw.helpers import run_program, v, v_mov_b32_e32
-from tinygrad.runtime.autogen.amd.rdna3.ins import VOPD, VOPD_LIT, VOPDOp
+from extra.assembly.amd.test.hw.helpers import run_program, run_program_emu, run_program_hw, compare_wave_states, \
+  v, s, v_mov_b32_e32, s_mov_b32
+from extra.assembly.amd.autogen.rdna3.ins import VOPD, VOPD_LIT, VOPDOp
 
 class TestVOPDBasic(unittest.TestCase):
   """Basic VOPD functionality tests."""
@@ -108,7 +109,7 @@ class TestVOPDLiterals(unittest.TestCase):
     Tests that the 32-bit literal (SIMM32) is correctly passed to the instruction.
     fma(2.0, 3.0, 10.0) = 2*3 + 10 = 16.0
     """
-    from test.amd.hw.helpers import f2i, i2f
+    from extra.assembly.amd.test.hw.helpers import f2i, i2f
     instructions = [
       v_mov_b32_e32(v[0], f2i(2.0)),  # v[0] = 2.0
       v_mov_b32_e32(v[1], f2i(3.0)),  # v[1] = 3.0
@@ -126,7 +127,7 @@ class TestVOPDLiterals(unittest.TestCase):
     Tests that the 32-bit literal (SIMM32) is correctly used as the multiplier.
     fma(2.0, 5.0, 3.0) = 2*5 + 3 = 13.0
     """
-    from test.amd.hw.helpers import f2i, i2f
+    from extra.assembly.amd.test.hw.helpers import f2i, i2f
     instructions = [
       v_mov_b32_e32(v[0], f2i(2.0)),  # v[0] = 2.0
       v_mov_b32_e32(v[1], f2i(3.0)),  # v[1] = 3.0
@@ -138,47 +139,6 @@ class TestVOPDLiterals(unittest.TestCase):
     self.assertAlmostEqual(i2f(st.vgpr[0][2]), 13.0, places=5, msg="fma(2.0, 5.0, 3.0) should be 13.0")
 
 
-class TestVOPDDot2Acc(unittest.TestCase):
-  """Tests for V_DUAL_DOT2ACC_F32_F16 - packed f16 dot product accumulate."""
-
-  def test_vopd_dot2acc_f32_f16_basic(self):
-    """V_DUAL_DOT2ACC_F32_F16: D += lo(S0)*lo(S1) + hi(S0)*hi(S1).
-
-    S0 = pack(1.0h, 2.0h), S1 = pack(3.0h, 4.0h), D = 10.0f
-    result = 10.0 + 1.0*3.0 + 2.0*4.0 = 10.0 + 3.0 + 8.0 = 21.0
-    """
-    from test.amd.hw.helpers import f2i, i2f, f32_to_f16
-    pk_s0 = f32_to_f16(1.0) | (f32_to_f16(2.0) << 16)  # lo=1.0h, hi=2.0h
-    pk_s1 = f32_to_f16(3.0) | (f32_to_f16(4.0) << 16)  # lo=3.0h, hi=4.0h
-    instructions = [
-      v_mov_b32_e32(v[0], pk_s0),
-      v_mov_b32_e32(v[1], pk_s1),
-      v_mov_b32_e32(v[3], f2i(10.0)),  # accumulator in v[3] (vdsty with vdstx=v[4])
-      # X: v[4] = MOV v[0] (don't care), Y: v[3] += dot2(v[0], v[1])
-      VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_DOT2ACC_F32_F16, v[4], v[3], v[0], v[0], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][3]), 21.0, places=2, msg="10.0 + 1.0*3.0 + 2.0*4.0 = 21.0")
-
-  def test_vopd_dot2acc_f32_f16_zero_accum(self):
-    """V_DUAL_DOT2ACC_F32_F16 with zero accumulator — pure dot product.
-
-    S0 = pack(0.5h, -1.0h), S1 = pack(2.0h, 3.0h), D = 0.0f
-    result = 0.0 + 0.5*2.0 + (-1.0)*3.0 = 1.0 - 3.0 = -2.0
-    """
-    from test.amd.hw.helpers import f2i, i2f, f32_to_f16
-    pk_s0 = f32_to_f16(0.5) | (f32_to_f16(-1.0) << 16)
-    pk_s1 = f32_to_f16(2.0) | (f32_to_f16(3.0) << 16)
-    instructions = [
-      v_mov_b32_e32(v[0], pk_s0),
-      v_mov_b32_e32(v[1], pk_s1),
-      v_mov_b32_e32(v[3], f2i(0.0)),  # zero accumulator in v[3]
-      VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_DOT2ACC_F32_F16, v[4], v[3], v[0], v[0], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][3]), -2.0, places=2, msg="0.5*2.0 + (-1.0)*3.0 = -2.0")
-
-
 class TestVOPDMultilane(unittest.TestCase):
   """Tests for VOPD with multiple lanes."""
 

extra/assembly/amd/test/test_compare_emulators.py

@@ -1,16 +1,17 @@
 # Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
-import unittest, ctypes
+import unittest, ctypes, os
 from dataclasses import dataclass
 from pathlib import Path
-from tinygrad import Device
 
-from test.mockgpu.amd.emu import WaveState, _decode_at, WAVE_SIZE, VCC_LO, EXEC_LO, SCC
-from tinygrad.renderer.amd import decode_inst
-import tinygrad
-REMU_PATH = Path(tinygrad.__file__).parent.parent / "extra/remu/target/release/libremu.so"
-if not REMU_PATH.exists(): REMU_PATH = Path(tinygrad.__file__).parent.parent / "extra/remu/target/release/libremu.dylib"
+# Set environment before any tinygrad imports to use MOCKGPU
+# This allows generating AMD GPU kernels without requiring real hardware
+os.environ["AMD"] = "1"
+os.environ["MOCKGPU"] = "1"
+os.environ["PYTHON_REMU"] = "1"
 
-def set_valid_mem_ranges(ranges): pass  # emu2 doesn't need this
+from extra.assembly.amd.emu import WaveState, decode_program, WAVE_SIZE, set_valid_mem_ranges, LDSMem
+from extra.assembly.amd.test.helpers import KernelInfo
+from extra.assembly.amd.test.bench_emu import REMU_PATH
 
 def _is_f32_nan(bits: int) -> bool:
   """Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
@@ -21,15 +22,6 @@ def _vals_equal(a: int, b: int) -> bool:
   if a == b: return True
   return _is_f32_nan(a) and _is_f32_nan(b)
 
-@dataclass
-class KernelSnapshot:
-  code: bytes
-  src: str
-  global_size: tuple[int, int, int]
-  local_size: tuple[int, int, int]
-  buf_idxs: list[int]  # indices into shared buffer pool
-  buf_sizes: list[int]  # sizes for each buffer index
-
 @dataclass
 class StateSnapshot:
   pc: int
@@ -93,71 +85,39 @@ class RustEmulator:
     return snap.to_snapshot()
 
   def free(self):
-    if self.ctx:
-      self.lib.wave_free(self.ctx)
-      self.ctx = None
+    if self.ctx: self.lib.wave_free(self.ctx); self.ctx = None
 
 class PythonEmulator:
   def __init__(self):
     self.state: WaveState | None = None
-    self.program: dict[int, tuple] = {}  # lazily populated: pc -> (name, fxn, globals)
-    self.vmem_buf = None
-    self.lds_buf = None
-    self.kernel_buf = None  # Keep kernel bytes alive
-    self.lib_addr = 0  # Base address of kernel code
+    self.program: dict | None = None
 
   def create(self, kernel: bytes, n_lanes: int):
-    import ctypes
-    from tinygrad.device import Buffer, BufferSpec
-    from tinygrad.dtype import dtypes
-    # Store kernel in a ctypes buffer so _decode_at can read from memory at actual PC address
-    self.kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    self.lib_addr = ctypes.addressof(self.kernel_buf)
-    self.program = {}
-    self.state = WaveState(n_lanes)
-    self.state.pc = self.lib_addr  # Set PC to code base address
-    self.vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
-    self.lds_buf = Buffer('CPU', 65536 // 4, dtypes.uint32).ensure_allocated()
-
-  def _ensure_decoded(self, pc: int):
-    if pc not in self.program:
-      runner, _ = _decode_at(pc, "rdna3")
-      self.program[pc] = (runner.p.function_name, runner._prg.fxn, runner.p.globals)
+    self.program = decode_program(kernel)
+    self.state = WaveState(LDSMem(bytearray(65536)), n_lanes)
+    self.state.exec_mask = (1 << n_lanes) - 1
 
   def step(self) -> int:
-    import ctypes
-    assert self.state is not None
-    pc = self.state.pc
-    if pc == 0xFFFFFFFFFFFFFFFF: return -1
-    self._ensure_decoded(pc)
-    name, fxn, globals_list = self.program[pc]
-    buf_addrs = {0: self.state.sgpr_buf._buf.va_addr, 1: self.state.vgpr_buf._buf.va_addr,  # type: ignore[union-attr]
-                 2: self.vmem_buf._buf.va_addr, 3: self.lds_buf._buf.va_addr}  # type: ignore[union-attr]
-    fxn(*[ctypes.c_uint64(buf_addrs[g]) for g in globals_list], ctypes.c_int32(0))
-    return -1 if self.state.pc == 0xFFFFFFFFFFFFFFFF else 0
-
+    assert self.program is not None and self.state is not None
+    return self.program[self.state.pc]._dispatch(self.state, self.program[self.state.pc])
   def set_sgpr(self, idx: int, val: int):
     assert self.state is not None
-    self.state._write_sgpr(idx, val)
+    self.state.sgpr[idx] = val & 0xffffffff
   def set_vgpr(self, lane: int, idx: int, val: int):
     assert self.state is not None
-    self.state._write_vgpr(idx, lane, val)
+    self.state.vgpr[lane][idx] = val & 0xffffffff
 
   def get_snapshot(self) -> StateSnapshot:
     assert self.state is not None
-    sgpr = [self.state._read_sgpr(i) for i in range(128)]
-    vgpr = [[self.state._read_vgpr(reg, lane) for reg in range(256)] for lane in range(WAVE_SIZE)]
-    # Convert actual PC address to word offset for comparison with Rust emulator
-    pc_offset = (self.state.pc - self.lib_addr) // 4 if self.state.pc != 0xFFFFFFFFFFFFFFFF else 0xFFFFFFFFFFFFFFFF
-    return StateSnapshot(pc=pc_offset, scc=self.state._read_sgpr(SCC.offset), vcc=sgpr[VCC_LO.offset],
-                         exec_mask=sgpr[EXEC_LO.offset], sgpr=sgpr, vgpr=vgpr)
+    return StateSnapshot(pc=self.state.pc, scc=self.state.scc, vcc=self.state.vcc & 0xffffffff,
+                         exec_mask=self.state.exec_mask & 0xffffffff, sgpr=list(self.state.sgpr),
+                         vgpr=[list(self.state.vgpr[i]) for i in range(WAVE_SIZE)])
 
 def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
-                      local_size: tuple[int, int, int], max_steps: int, debug: bool, trace_len: int,
-                      kernel_idx: int = 0, max_workgroups: int = 8) -> tuple[bool, str, int]:
+                      program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
+                      max_workgroups: int = 8) -> tuple[bool, str, int]:
   """Run a single kernel through both emulators. Returns (success, message, total_steps)."""
   gx, gy, gz = global_size
-  lx, ly, lz = local_size
   total_steps = 0
   wg_count = 0
 
@@ -180,53 +140,28 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
           emu.set_sgpr(13, gidx)
           emu.set_sgpr(14, gidy)
           emu.set_sgpr(15, gidz)
-          # Initialize v[0] with packed workitem IDs for each lane
-          for lane in range(n_lanes):
-            tid = lane
-            z, y, x = tid // (lx * ly), (tid // lx) % ly, tid % lx
-            emu.set_vgpr(lane, 0, (z << 20) | (y << 10) | x)
 
         step = 0
         trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
-        prev_sync_after = False  # Track if previous instruction had known Rust bugs
         try:
           while step < max_steps:
             rust_before = rust.get_snapshot()
             python_before = python.get_snapshot()
 
-            pc_addr = python.lib_addr + python_before.pc * 4  # Convert word offset to actual address
-            python._ensure_decoded(pc_addr)
-            inst_hex_name = python.program[pc_addr][0]
-            # Decode the instruction to get mnemonic for sync_after checks
-            try:
-              # Format is mnemonic_hexbytes, e.g. v_exp_f32_e32_014b027e -> hex is 014b027e
-              parts = inst_hex_name.rsplit('_', 1)
-              inst_bytes_hex = parts[1] if len(parts) == 2 else ""
-              inst_bytes = bytes.fromhex(inst_bytes_hex) if inst_bytes_hex else b''
-              decoded = decode_inst(inst_bytes) if inst_bytes else None
-              inst_mnemonic = repr(decoded).split('(')[0] if decoded else ""
-            except Exception:
-              inst_mnemonic = ""
-            # For generic instructions, use function name for sync_after check
-            if not inst_mnemonic: inst_mnemonic = inst_hex_name
-            inst_str = inst_hex_name
+            inst = program.get(python_before.pc)
+            inst_str = inst.disasm() if inst else f"unknown at PC={python_before.pc}"
             trace.append((step, python_before.pc, inst_str, rust_before, python_before))
             if len(trace) > trace_len: trace.pop(0)
 
             if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
 
-            # Instructions with known Rust emulator bugs or precision differences - sync Python to Rust after execution
+            # Instructions with known Rust emulator bugs - sync Python to Rust after execution
             # v_div_scale/v_div_fixup: Rust has different VCC handling
             # v_cvt_f16_f32: Rust clears high 16 bits, but hardware (and Python) preserves them
             # s_add_i32/s_sub_i32: Rust has incorrect SCC overflow detection
-            # v_exp_f32/v_log_f32/v_ldexp_f32: precision differences in transcendental functions
-            # s_delay_alu: Rust handles differently
-            # v_add_co_ci_u32/v_sub_co_ci_u32/v_subrev_co_ci_u32: Rust preserves inactive VCC bits, but hardware clears all bits
-            sync_after = any(x in inst_mnemonic.lower() for x in ('v_div_scale', 'v_div_fixup', 'v_cvt_f16_f32', 's_add_i32', 's_sub_i32',
-                                                                   'v_exp_f32', 'v_log_f32', 'v_ldexp_f32', 's_delay_alu',
-                                                                   'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'))
-            # Skip comparison if previous instruction had known Rust bugs (states were synced but may still differ slightly)
-            diffs = rust_before.diff(python_before, n_lanes) if not prev_sync_after else []
+            sync_after = any(x in inst_str for x in ('v_div_scale_f32', 'v_div_scale_f64', 'v_div_fixup_f32', 'v_div_fixup_f64',
+                                                      'v_cvt_f16_f32', 's_add_i32', 's_sub_i32'))
+            diffs = rust_before.diff(python_before, n_lanes)
             if diffs:
               trace_lines = []
               for idx, (s, pc, d, rb, pb) in enumerate(trace):
@@ -237,18 +172,16 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
                   python_diffs = pb.diff(next_pb, n_lanes, "->")
                   if rust_diffs: trace_lines.append(f"             rust:   {', '.join(rust_diffs[:5])}")
                   if python_diffs: trace_lines.append(f"             python: {', '.join(python_diffs[:5])}")
-                  elif rust_diffs: trace_lines.append("             python: (no changes)")
+                  elif rust_diffs: trace_lines.append(f"             python: (no changes)")
                 else:
                   # Last traced instruction - compare with current state
                   rust_diffs = rb.diff(rust_before, n_lanes, "->")
                   python_diffs = pb.diff(python_before, n_lanes, "->")
                   if rust_diffs: trace_lines.append(f"             rust:   {', '.join(rust_diffs[:5])}")
                   if python_diffs: trace_lines.append(f"             python: {', '.join(python_diffs[:5])}")
-                  elif rust_diffs: trace_lines.append("             python: (no changes)")
+                  elif rust_diffs: trace_lines.append(f"             python: (no changes)")
               trace_str = "\n".join(trace_lines)
-              msg = f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ (rust vs python):\n  "
-              msg += "\n  ".join(diffs[:10]) + f"\n  Recent instructions:\n{trace_str}"
-              return False, msg, total_steps
+              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ (rust vs python):\n  " + "\n  ".join(diffs[:10]) + f"\n  Recent instructions:\n{trace_str}", total_steps
 
             rust_result = rust.step()
             python_result = python.step()
@@ -258,9 +191,7 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
               if rust_result == 1 and python_result == 0:
                 raise unittest.SkipTest(f"Rust emulator doesn't support instruction: {inst_str}")
               trace_str = "\n".join(f"    step {s}: PC={pc:3d} {d}" for s, pc, d, _, _ in trace)
-              msg = (f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: "
-                     f"rust={rust_result}, python={python_result}, inst={inst_str}\n  Recent instructions:\n{trace_str}")
-              return False, msg, total_steps
+              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: rust={rust_result}, python={python_result}, inst={inst_str}\n  Recent instructions:\n{trace_str}", total_steps
 
             # Sync Python state to Rust after instructions with known Rust emulator differences
             if sync_after:
@@ -269,12 +200,7 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
               for lane in range(n_lanes):
                 for i in range(256): python.set_vgpr(lane, i, rust_after.vgpr[lane][i])
               assert python.state is not None
-              # Convert Rust's word-based PC to Python's actual address
-              python.state.pc = python.lib_addr + rust_after.pc * 4
-              python.state._write_sgpr(SCC.offset, rust_after.scc)
-              python.state._write_sgpr(VCC_LO.offset, rust_after.vcc)
-              python.state._write_sgpr(EXEC_LO.offset, rust_after.exec_mask)
-            prev_sync_after = sync_after
+              python.state.pc, python.state.scc, python.state.vcc, python.state.exec_mask = rust_after.pc, rust_after.scc, rust_after.vcc, rust_after.exec_mask
 
             if rust_result == -1:
               total_steps += step + 1
@@ -293,7 +219,7 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
 
   return True, f"Completed {gx*gy*gz} workgroups", total_steps
 
-def compare_emulators_multi_kernel(kernels: list[KernelSnapshot], buf_pool: dict[int, int], max_steps: int = 1000,
+def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int, int], max_steps: int = 1000,
                                     debug: bool = False, trace_len: int = 10, buf_data: dict[int, bytes] | None = None) -> tuple[bool, str]:
   """Run all kernels through both emulators with shared buffer pool."""
   if buf_data is None: buf_data = {}
@@ -323,11 +249,12 @@ def compare_emulators_multi_kernel(kernels: list[KernelSnapshot], buf_pool: dict
     kernel_ranges = ranges | {(args_ptr, ctypes.sizeof(args))}
     set_valid_mem_ranges(kernel_ranges)
 
+    program = decode_program(kernel.code)
     n_lanes = kernel.local_size[0] * kernel.local_size[1] * kernel.local_size[2]
 
     ok, msg, steps = run_single_kernel(
       kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
-      kernel.local_size, max_steps, debug, trace_len, ki
+      program, max_steps, debug, trace_len, ki
     )
     total_steps += steps
     if not ok:
@@ -353,11 +280,11 @@ def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list,
   ranges.add((args_ptr, ctypes.sizeof(args)))
   set_valid_mem_ranges(ranges)
 
-  # Legacy wrapper assumes local_size = (n_lanes, 1, 1)
-  ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, (n_lanes, 1, 1), max_steps, debug, trace_len)
+  program = decode_program(kernel)
+  ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, program, max_steps, debug, trace_len)
   return ok, msg
 
-def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, int], dict[int, bytes]]:
+def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int], dict[int, bytes]]:
   """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
   from tinygrad import Tensor
   from tinygrad.runtime.support.elf import elf_loader
@@ -395,9 +322,8 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, in
                 buf_pool[buf_id] = b.nbytes
               buf_idxs.append(buf_id)
               buf_sizes.append(b.nbytes)
-            kernels.append(KernelSnapshot(
+            kernels.append(KernelInfo(
               code=bytes(sec.content),
-              src=lowered.prg.p.src,
               global_size=tuple(lowered.prg.p.global_size),
               local_size=tuple(lowered.prg.p.local_size),
               buf_idxs=buf_idxs,
@@ -412,7 +338,6 @@ def get_kernel_from_tinygrad(op_fn) -> tuple[bytes, tuple[int, int, int], tuple[
   k = kernels[-1]
   return k.code, k.global_size, k.local_size, k.buf_sizes
 
-@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestTinygradKernels(unittest.TestCase):
   """Compare emulators on real tinygrad-compiled kernels."""
 
@@ -449,8 +374,7 @@ class TestTinygradKernels(unittest.TestCase):
   def test_cast(self): self._test_kernel(lambda T: T.empty(32).half().float() + T.empty(32).int().float())
 
   # Pooling - regression for VCC wave32 mode
-  def test_pool2d(self):
-    self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4)) + T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4)))
+  def test_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4)) + T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4)))
 
   # Convolution
   def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 2, 8, 8).conv2d(T.empty(2, 2, 3, 3)), max_steps=50000)
@@ -462,7 +386,6 @@ class TestTinygradKernels(unittest.TestCase):
     from tinygrad import dtypes
     self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
   def test_gelu(self): self._test_kernel(lambda T: T.empty(32, 32).gelu())
-  def test_exp(self): self._test_kernel(lambda T: T.empty(1024).exp())
   def test_cross_entropy(self):
     import numpy as np
     np.random.seed(0)
@@ -474,55 +397,5 @@ class TestTinygradKernels(unittest.TestCase):
     from tinygrad import dtypes
     self._test_kernel(lambda T: T([2.0], dtype=dtypes.float64).sin())
 
-  def test_sin_large_f32(self):
-    """Test sin with large values that trigger Payne-Hanek range reduction."""
-    # Values around 859240 trigger the Payne-Hanek algorithm
-    # This tests the integer multiply-high instructions used in range reduction
-    self._test_kernel(lambda T: T([859240.0, 1000000.0, 100594688.0]).sin())
-
-  def test_clip_zero_one(self):
-    """Test clip(0, 1) - regression for binary_crossentropy failure."""
-    import numpy as np
-    np.random.seed(0)
-    x_np = np.random.uniform(-2, 2, (32, 10)).astype(np.float32).tolist()
-    self._test_kernel(lambda T: T(x_np).clip(0, 1))
-
-  def test_mod_int64(self):
-    """Test int64 modulo, especially edge cases like 1 % -1."""
-    from tinygrad import dtypes
-    self._test_kernel(lambda T: T([1, 10, -10, 7], dtype=dtypes.int64) % T([-1, 3, 3, -3], dtype=dtypes.int64))
-
-  def test_expand_flatten_sum(self):
-    """Test flatten of expanded tensor followed by sum.
-
-    Bug: flatten() of an expanded tensor produces wrong results for certain sizes.
-    Sizes that are multiples of 32 work (32, 48, 64), but sizes like 33, 49, 50 fail.
-    This breaks masked_select and nonzero operations.
-    """
-    import numpy as np
-    np.random.seed(0)
-    x_np = np.random.uniform(-2, 2, (33,)).astype(np.float32)
-    self._test_kernel(lambda T: (T(x_np.tolist()) > 0.5).unsqueeze(-1).expand(33, 3).flatten().sum())
-
-  @unittest.skip("slow and broken with AMD_LLVM=1")
-  def test_nonzero(self):
-    """Test nonzero operation - counts and gathers indices of non-zero elements."""
-    import numpy as np
-    np.random.seed(42)
-    x_np = np.random.rand(10, 5, 3).astype(np.float32)
-    self._test_kernel(lambda T: (T(x_np.tolist()) > 0.5).nonzero())
-
-  @unittest.skip("Precision differences in v_exp/v_log accumulate across kernels, causing memory divergence")
-  def test_softmax_argmax_fused(self):
-    """Test fused softmax+argmax - tracks exp2 precision issue.
-
-    The fused kernel recomputes softmax inline and Python emulator's exp2 polynomial
-    has up to 1 ULP error vs native exp2f, causing accumulated differences.
-    """
-    import torch
-    torch.manual_seed(0)
-    x_np = torch.rand(4, 10).numpy()
-    self._test_kernel(lambda T: T(x_np.tolist()).softmax(1).argmax())
-
 if __name__ == "__main__":
   unittest.main()

extra/assembly/amd/test/test_custom_kernel.py

@@ -0,0 +1,48 @@
+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_dsl2.py

@@ -1,8 +1,8 @@
 import unittest
-from tinygrad.renderer.amd.dsl import *
-from tinygrad.renderer.amd.dsl import VDSTYField
-from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op
-from tinygrad.runtime.autogen.amd.rdna3.ins import VOP1
+from extra.assembly.amd.dsl import *
+from extra.assembly.amd.dsl import VDSTYField
+from extra.assembly.amd.autogen.rdna3.enum import VOP1Op, VOP2Op
+from extra.assembly.amd.autogen.rdna3.ins import VOP1
 
 class TestRegisters(unittest.TestCase):
   def test_vgpr_single(self):

extra/assembly/amd/test/test_formats.py

@@ -4,10 +4,10 @@
 Note: Graphics-only formats (EXP, MUBUF, MTBUF, MIMG) are not supported - use GLOBAL/FLAT for memory access in compute.
 """
 import unittest
-from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.renderer.amd.dsl import VCC_HI, EXEC_LO, NULL
+from extra.assembly.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.dsl import VCC_HI, EXEC_LO, NULL
 OFF = NULL  # OFF is alias for NULL
-from tinygrad.renderer.amd import detect_format
+from extra.assembly.amd.decode import detect_format
 
 
 class TestDS(unittest.TestCase):

extra/assembly/amd/test/test_handwritten.py

@@ -2,26 +2,21 @@
 # the Inst constructor should be looking at the types of the fields to correctly set the value
 
 import unittest, struct
-from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.renderer.amd.dsl import Inst
-from test.amd.test_roundtrip import compile_asm
-from test.amd.disasm import disasm
+from extra.assembly.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.dsl import Inst
+from extra.assembly.amd.test.test_roundtrip import compile_asm
 
-class IntegrationTestBase(unittest.TestCase):
+class TestIntegration(unittest.TestCase):
   inst: Inst
-  arch: str
   def tearDown(self):
     if not hasattr(self, 'inst'): return
     b = self.inst.to_bytes()
-    st = disasm(self.inst)
+    st = self.inst.disasm()
     # Test that the instruction can be compiled by LLVM and produces the same bytes
     desc = f"{st:25s} {self.inst} {b!r}"
-    self.assertEqual(b, compile_asm(st, arch=self.arch), desc)
+    self.assertEqual(b, compile_asm(st), desc)
     print(desc)
 
-class TestIntegration(IntegrationTestBase):
-  arch: str = "rdna3"
-
   def test_wmma(self):
     self.inst = v_wmma_f32_16x16x16_f16(v[0:7], v[184:191], v[136:143], v[0:7])
 
@@ -129,17 +124,6 @@ class TestIntegration(IntegrationTestBase):
     int_inst = s_mov_b32(s[0], struct.unpack("I", struct.pack("f", 1337.0))[0])
     self.assertEqual(self.inst, int_inst)
 
-class TestIntegrationCDNA(IntegrationTestBase):
-  arch = "cdna"
-
-  def test_mfma(self):
-    from tinygrad.runtime.autogen.amd.cdna.ins import v_mfma_f32_16x16x16_f16
-    self.inst = v_mfma_f32_16x16x16_f16(v[0:3], v[0:1], v[0:1], 0)
-
-  def test_mfma_fp8(self):
-    from tinygrad.runtime.autogen.amd.cdna.ins import v_mfma_f32_16x16x128_f8f6f4
-    self.inst = v_mfma_f32_16x16x128_f8f6f4(v[0:3], v[0:5], v[0:5], 1, cbsz=2, blgp=2)
-
 class TestRegisterSliceSyntax(unittest.TestCase):
   """
   Issue: Register slice syntax should use AMD assembly convention (inclusive end).
@@ -161,9 +145,9 @@ class TestRegisterSliceSyntax(unittest.TestCase):
     # 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)
-    d = disasm(inst)
+    disasm = inst.disasm()
     # Disasm shows s[4:7] - user should be able to copy this back
-    self.assertIn("s[4:7]", d)
+    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.sz, 4, "s[4:7] from disasm should give 4 registers")

extra/assembly/amd/test/test_integration.py

@@ -0,0 +1,258 @@
+#!/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_llvm.py

@@ -8,11 +8,11 @@ Only compute-relevant instruction formats are tested. Graphics-only formats not
   - VIMAGE/VSAMPLE: image sampling instructions (RDNA4)
   - VBUFFER: buffer instructions (RDNA4)
 """
-import unittest, re, functools
+import unittest, re, subprocess, functools
 from tinygrad.helpers import fetch
-from test.amd.disasm import disasm
-from tinygrad.renderer.amd import decode_inst, detect_format
-from test.amd.helpers import llvm_assemble, llvm_filter_valid_asm, get_target, get_mattr
+from extra.assembly.amd.disasm import disasm
+from extra.assembly.amd.decode import decode_inst, detect_format
+from extra.assembly.amd.test.helpers import get_llvm_mc, get_target, get_mattr
 
 LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-21.1.0/llvm/test/MC/AMDGPU"
 
@@ -40,7 +40,7 @@ RDNA4_FILES = ['gfx12_asm_sop1.s', 'gfx12_asm_sop2.s', 'gfx12_asm_sopp.s', 'gfx1
   'gfx12_asm_vop1.s', 'gfx12_asm_vop2.s', 'gfx12_asm_vopc.s', 'gfx12_asm_vopcx.s', 'gfx12_asm_vop3.s', 'gfx12_asm_vop3c.s',
   'gfx12_asm_vop3cx.s', 'gfx12_asm_vop3p.s', 'gfx12_asm_vop3_from_vop1.s', 'gfx12_asm_vop3_from_vop2.s',
   'gfx12_asm_vop3p_features.s', 'gfx12_asm_vopd.s', 'gfx12_asm_vopd_features.s',
-  'gfx12_asm_ds.s', 'gfx12_asm_smem.s', 'gfx12_asm_vflat.s',
+  'gfx12_asm_ds.s', 'gfx12_asm_smem.s',
   'gfx12_asm_wmma_w32.s']
 
 def _parse_llvm_tests(text: str, pattern: str) -> list[tuple[str, bytes]]:
@@ -74,13 +74,42 @@ def _get_tests_uncached(f: str, arch: str) -> list[tuple[str, bytes]]:
   # Exclude v_interp_* (graphics-only, not on CDNA)
   if arch == "cdna": tests = [(asm, data) for asm, data in tests if not asm.startswith('v_interp_')]
   # Filter out tests where original ASM isn't valid on target (e.g., gfx9 tests with gfx942/gfx950 constraints)
-  if arch == "cdna" and not ('gfx942' in f or 'gfx950' in f or 'gfx90a' in f):
-    tests = llvm_filter_valid_asm(tests, get_target(arch), get_mattr(arch))
+  if arch == "cdna" and not ('gfx942' in f or 'gfx950' in f or 'gfx90a' in f): tests = _filter_valid_asm(tests, arch)
   return tests
 
 @functools.cache
 def _get_tests(f: str, arch: str) -> list[tuple[str, bytes]]: return _get_tests_uncached(f, arch)
 
+def _compile_asm_batch(instrs: list[str], arch: str = "rdna3", mcpu: str|None = None) -> list[bytes]:
+  if not instrs: return []
+  mcpu, mattr = mcpu or get_target(arch), get_mattr(arch)
+  result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', f'-mattr={mattr}', '-show-encoding'],
+    input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True, timeout=30)
+  if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
+  return [bytes.fromhex(line.split('encoding:')[1].strip()[1:-1].replace('0x', '').replace(',', '').replace(' ', ''))
+          for line in result.stdout.split('\n') if 'encoding:' in line]
+
+def _filter_valid_asm(tests: list[tuple[str, bytes]], arch: str) -> list[tuple[str, bytes]]:
+  """Filter out tests where the original ASM isn't valid on the target (e.g., gfx9 tests with gfx942/gfx950 constraints)."""
+  if not tests: return []
+  mcpu = get_target(arch)
+  # Batch assemble all instructions, parse stderr to find which lines failed
+  instrs = [asm for asm, _ in tests]
+  result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', '-show-encoding'],
+    input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True, timeout=30)
+  # Parse error lines from stderr (format: "<stdin>:N:..." where N is 1-indexed, line 1 is ".text")
+  failed_lines = set()
+  for line in result.stderr.split('\n'):
+    if m := re.match(r'<stdin>:(\d+):', line): failed_lines.add(int(m.group(1)) - 1)  # -1 for .text, so line 2 -> index 1 -> tests[0]
+  # Also filter out tests where LLVM roundtrip doesn't match original (reserved bits set in original)
+  valid = [(asm, data) for i, (asm, data) in enumerate(tests) if (i + 1) not in failed_lines]
+  if not valid: return []
+  llvm_result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', '-show-encoding'],
+    input=".text\n" + "\n".join(asm for asm, _ in valid) + "\n", capture_output=True, text=True, timeout=30)
+  llvm_bytes = [bytes.fromhex(line.split('encoding:')[1].strip()[1:-1].replace('0x', '').replace(',', '').replace(' ', ''))
+                for line in llvm_result.stdout.split('\n') if 'encoding:' in line]
+  return [(asm, data) for (asm, data), lb in zip(valid, llvm_bytes) if lb == data]
+
 def _make_test(f: str, arch: str, test_type: str):
   def test(self):
     tests = _get_tests(f, arch)
@@ -96,28 +125,6 @@ def _make_test(f: str, arch: str, test_type: str):
         except ValueError: skipped += 1  # skip invalid opcodes not in enum
       print(f"{name}: {passed} passed, {skipped} skipped")
       self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
-    elif test_type == "repr":
-      # Test that eval(repr(inst)) reproduces the instruction
-      if arch == "rdna3": import tinygrad.runtime.autogen.amd.rdna3.ins as ins  # type: ignore[no-redef]
-      elif arch == "rdna4": import tinygrad.runtime.autogen.amd.rdna4.ins as ins  # type: ignore[no-redef]
-      elif arch == "cdna": import tinygrad.runtime.autogen.amd.cdna.ins as ins  # type: ignore[no-redef]
-      ns = {k: getattr(ins, k) for k in dir(ins) if not k.startswith('_')}
-      passed, skipped = 0, 0
-      for _, data in tests:
-        try:
-          decoded = detect_format(data, arch).from_bytes(data)
-          if decoded.to_bytes()[:len(data)] != data:
-            skipped += 1
-            continue  # skip if binary roundtrip fails
-          r = repr(decoded)
-          try:
-            decoded2 = eval(r, ns)  # noqa: S307
-            if decoded == decoded2: passed += 1
-            else: skipped += 1
-          except Exception: skipped += 1
-        except ValueError: skipped += 1
-      print(f"{name}: {passed} passed, {skipped} skipped")
-      self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
     elif test_type == "disasm":
       to_test = []
       for _, data in tests:
@@ -126,12 +133,12 @@ def _make_test(f: str, arch: str, test_type: str):
           enc = decoded.to_bytes()[:len(data)]
           # Skip if roundtrip fails, disasm fails, or op_name is missing (disasm starts with space)
           if enc == data and (d := disasm(decoded)) and not d.startswith(' '): to_test.append((enc, d))
-        except Exception: pass
+        except: pass
       skipped = len(tests) - len(to_test)
       print(f"{name}: {len(to_test)} passed, {skipped} skipped")
       self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
       # Compare disasm->reassemble with original encoding (filter reserved bit cases where LLVM can't reproduce)
-      llvm_bytes = llvm_assemble([t[1] for t in to_test], mcpu, get_mattr(arch))
+      llvm_bytes = _compile_asm_batch([t[1] for t in to_test], arch, mcpu)
       valid = [(enc, d, llvm) for (enc, d), llvm in zip(to_test, llvm_bytes) if llvm == enc]
       print(f"{name}: {len(valid)}/{len(to_test)} matched LLVM encoding")
       for enc, _, llvm in valid: self.assertEqual(llvm, enc)
@@ -142,15 +149,12 @@ class TestLLVM(unittest.TestCase): pass
 for f in RDNA_FILES:
   setattr(TestLLVM, f"test_rdna3_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "roundtrip"))
   setattr(TestLLVM, f"test_rdna3_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "disasm"))
-  setattr(TestLLVM, f"test_rdna3_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "repr"))
 for f in CDNA_FILES:
   setattr(TestLLVM, f"test_cdna_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "roundtrip"))
   setattr(TestLLVM, f"test_cdna_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "disasm"))
-  setattr(TestLLVM, f"test_cdna_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "repr"))
 for f in RDNA4_FILES:
   setattr(TestLLVM, f"test_rdna4_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "roundtrip"))
   setattr(TestLLVM, f"test_rdna4_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "disasm"))
-  setattr(TestLLVM, f"test_rdna4_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "repr"))
 
 if __name__ == "__main__":
   unittest.main()