failing test from application

.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:
@@ -66,14 +62,14 @@ runs:
       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 }}
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
     - name: Cache Python packages
       if: github.event_name != 'pull_request'
       id: restore-venv
       uses: actions/cache@v4
       with:
         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 ****
 
@@ -149,7 +145,7 @@ runs:
       run: |
         wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
         sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
-        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/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
 
@@ -199,13 +195,13 @@ runs:
       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 }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
     - name: Cache apt
       if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name != 'pull_request'
       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 +233,7 @@ runs:
       shell: bash
       run: |
         sudo mkdir -p /usr/local/lib
-        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/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
@@ -330,9 +326,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'
@@ -44,11 +43,11 @@ jobs:
       run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev
     - name: Regenerate autogen files
       run: |
-        find tinygrad/runtime/autogen -type f -name "*.py" -not -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 import comgr, 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.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"
@@ -60,9 +59,8 @@ jobs:
     - 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
@@ -82,7 +80,6 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
-        key: 'autogen-mac'
         llvm: 'true'
     - name: Regenerate autogen files
       run: |
@@ -91,9 +88,8 @@ jobs:
     - 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
@@ -103,8 +99,8 @@ jobs:
         name: autogen-macos-patch
         path: autogen-macos.patch
 
-  autogen-comgr-2:
+  autogen-comgr-3:
-    name: In-tree Autogen (comgr 2)
+    name: In-tree Autogen (comgr 3)
     runs-on: ubuntu-24.04
     timeout-minutes: 15
     steps:
@@ -112,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
+        rm tinygrad/runtime/autogen/comgr_3.py
-        python3 -c "from tinygrad.runtime.autogen import comgr"
+        python3 -c "from tinygrad.runtime.autogen import comgr_3"
     - name: Check for differences
       run: |
         if ! git diff --quiet; then
-          git diff
+          git diff > autogen-comgr3.patch
-          git diff > autogen-comgr2.patch
+          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
-          echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
           exit 1
         fi
     - name: Upload patch artifact
       if: failure()
       uses: actions/upload-artifact@v4
       with:
-        name: autogen-comgr2-patch
+        name: autogen-comgr3-patch
-        path: autogen-comgr2.patch
+        path: autogen-comgr3.patch

.github/workflows/benchmark.yml

@@ -16,43 +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=2 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
-      run: FLOAT16=1 CL=1 IMAGE=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-
   testmacbenchmark:
     name: Mac Benchmark
     env:
@@ -332,13 +295,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
@@ -510,7 +473,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
@@ -520,9 +483,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
 
@@ -617,27 +577,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="." DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot load_pickle 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
-
   testreddriverbenchmark:
     name: AM Benchmark
     runs-on: [self-hosted, Linux, tinyboxrandom]

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 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 }}
@@ -29,9 +29,9 @@ jobs:
         deps: testing_unit
         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
@@ -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
@@ -158,25 +158,25 @@ jobs:
         key: be-minimal
         deps: testing_unit
     - 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
+      run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
     - name: Test IMAGE=2 support
       run: |
-        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
+        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
+        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 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/backend/test_ops.py TestOps.test_gemm
+        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/backend/test_ops.py TestOps.test_gemm_fp16
+        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=METAL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/null/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/null/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/null/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/null/test_uops_stats.py TestUOpsStats.test_simple_matmul
+        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul
 
   linter:
     name: Linters
@@ -244,37 +244,6 @@ jobs:
     - 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
-
   unittest:
     name: Unit Tests
     runs-on: ubuntu-latest
@@ -297,8 +266,20 @@ jobs:
       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 test/unit/test_device.py TestRunAsModule.test_module_runs
-        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+        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 +293,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
+    - name: Repo line count < 20000 lines
-      run: MAX_LINE_COUNT=24000 python sz.py
+      run: MAX_LINE_COUNT=20000 python sz.py
 
   spec:
     strategy:
@@ -333,7 +314,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
@@ -371,7 +352,7 @@ jobs:
           opencl: 'true'
       - name: Test CL IMAGE=2 ops
         run: |
-          CL=1 IMAGE=2 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=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
       - name: Run process replay tests
@@ -395,7 +376,7 @@ jobs:
       - name: Run Kernel Count Test
         run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
       - name: Run fused optimizer tests
-        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/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:
@@ -454,7 +435,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
 
@@ -484,11 +465,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
 
@@ -568,11 +549,11 @@ jobs:
         pydeps: "pillow"
         llvm: "true"
     - name: Test LLVM=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/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)
@@ -586,7 +567,7 @@ jobs:
       with:
         key: dsp-minimal
         deps: testing_unit
-        pydeps: "onnx==1.18.0 onnxruntime ml_dtypes"
+        pydeps: "onnx==1.18.0 onnxruntime"
         llvm: "true"
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v3
@@ -598,15 +579,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)
@@ -625,13 +606,61 @@ jobs:
     - 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
 
+  testamd:
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [amd, amdllvm]
+
+    name: Linux (${{ matrix.backend }})
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    env:
+      AMD: 1
+      MOCKGPU: 1
+      FORWARD_ONLY: 1
+      AMD_LLVM: ${{ matrix.backend == 'amdllvm' && '1' || matrix.backend != 'amdllvm' && '0' }}
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: ${{ matrix.backend }}-minimal
+          deps: testing_unit
+          amd: 'true'
+          llvm: ${{ matrix.backend == 'amdllvm' && 'true' }}
+      - name: Check Device.DEFAULT and print some source
+        run: |
+          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+      - name: Run LLVM test
+        if: matrix.backend=='amdllvm'
+        run: python test/device/test_amd_llvm.py
+      - name: Run pytest (amd)
+        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py --durations=20
+      - name: Run pytest (amd)
+        run: python -m pytest test/external/external_test_am.py --durations=20
+      - name: Run TRANSCENDENTAL math
+        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
+      - name: Run TestOps.test_add with SQTT
+        run: |
+          VIZ=-2 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 AMD emulated mmapeak on NULL backend
+        env:
+          AMD: 0
+        run: PYTHONPATH=. NULL=1 EMULATE=AMD python extra/mmapeak/mmapeak.py
+      - name: Run process replay tests
+        uses: ./.github/actions/process-replay
+
   testamdasm:
     name: AMD ASM IDE
     runs-on: ubuntu-24.04
@@ -652,94 +681,29 @@ jobs:
           python-version: '3.14'
       - name: Verify AMD autogen is up to date
         run: |
-          python -m tinygrad.renderer.amd.generate
+          python -m extra.assembly.amd.generate
-          git diff --exit-code tinygrad/runtime/autogen/amd/
+          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_rocprof_decoder.py
+        run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
-      - name: Run AMD renderer tests
+      - name: Run RDNA3 emulator tests
-        run: AMD_LLVM=0 python -m pytest -n=auto test/amd/ --durations 20
+        run: AMD_LLVM=0 python -m pytest -n=auto extra/assembly/amd/ --durations 20
-      - name: Run AMD renderer tests (AMD_LLVM=1)
+      - name: Run RDNA3 emulator tests (AMD_LLVM=1)
-        run: AMD_LLVM=1 python -m pytest -n=auto test/amd/ --durations 20
+        run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
-      - name: Run SQTT profiling tests
+      - name: Run RDNA3 dtype tests
-        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py --durations 20
-      - name: Run AMD emulated tests on NULL backend
+      - name: Run RDNA3 dtype tests (AMD_LLVM=1)
-        env:
+        run: AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py --durations 20
-          AMD: 0
+      # TODO: run all once emulator is faster
-        run: |
+      - name: Run RDNA3 ops tests
-          PYTHONPATH=. NULL=1 EMULATE=AMD python extra/mmapeak/mmapeak.py
+        run: SKIP_SLOW_TEST=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril or test_nonzero or test_softmax_argmax" --durations 20
-          PYTHONPATH=. NULL=1 EMULATE=AMD_CDNA4 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
-      - name: Run ASM matmul on MOCKGPU
-        run: PYTHONPATH="." AMD=1 MOCKGPU=1 N=256 python3 extra/gemm/amd_asm_matmul.py
-      - name: Run LLVM test
-        run: AMD_LLVM=1 python test/device/test_amd_llvm.py
-
-  testmockam:
-    name: Linux (am)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    env:
-      AMD: 1
-      MOCKGPU: 1
-      AMD_IFACE: PCI
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: mockam
-          deps: testing_unit
-          amd: 'true'
-      - name: Run test_tiny on MOCKAM
-        run: python test/test_tiny.py
-      - name: Run test_tiny on MOCKAM USB
-        run: AMD_IFACE=USB python test/test_tiny.py
-      - name: Run test_hcq on MOCKAM
-        run: python -m pytest test/device/test_hcq.py
-
-  testamd:
-    strategy:
-      fail-fast: false
-      matrix:
-        backend: [amd, amdllvm]
-        arch: [rdna3, rdna4]
-        #arch: [rdna3, rdna4, cdna4]
-
-    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: ubuntu-22.04
-    timeout-minutes: 15
-    env:
-      AMD: 1
-      MOCKGPU: 1
-      MOCKGPU_ARCH: ${{ matrix.arch }}
-      SKIP_SLOW_TEST: 1
-      AMD_LLVM: ${{ matrix.backend == 'amdllvm' && '1' || matrix.backend != 'amdllvm' && '0' }}
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: ${{ matrix.backend }}-minimal
-          deps: testing_unit
-          amd: 'true'
-          llvm: ${{ matrix.backend == 'amdllvm' && 'true' }}
-      - name: Check Device.DEFAULT and print some source
-        run: |
-          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
-          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
-      - 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/testextra/test_cfg_viz.py test/external/external_test_am.py --durations=20
-      - name: Run TRANSCENDENTAL math
-        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
-      - name: Run process replay tests
-        uses: ./.github/actions/process-replay
 
   testnvidia:
     strategy:
@@ -768,12 +732,12 @@ jobs:
       - name: Check Device.DEFAULT and print some source
         run: |
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
-          DEBUG=5 FORWARD_ONLY=1 python3 test/test_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
+        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 TestOps.test_add with PMA
-        run: VIZ=-1 PMA=1 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
+        run: VIZ=-1 PMA=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -802,11 +766,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
 
@@ -831,22 +795,18 @@ jobs:
         llvm: 'true'
     - name: Run unit tests
       run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
-    - name: Run NULL backend tests
-      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
     - name: Run ONNX
       run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
     - name: Test tensor core ops (fake)
-      run: METAL=1 DEBUG=3 TC=2 python test/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
@@ -869,8 +829,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
@@ -889,14 +847,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
+    - name: Clean npm cache
-      run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m pytest -n=auto test/backend --durations=20
+      run: npm cache clean --force
-    #- name: Clean npm cache
+    - name: Install Puppeteer
-    #  run: npm cache clean --force
+      run: npm install puppeteer
-    #- name: Install Puppeteer
-    #  run: npm install puppeteer
     # this is also flaky
     #- name: Run WEBGPU Efficientnet
     #  run: node test/web/test_webgpu.js
@@ -938,7 +896,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
@@ -971,16 +929,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
+        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 NULL backend tests
-        if: matrix.backend=='llvm'
-        shell: bash
-        run: CPU=0 CPU_LLVM=0 NULL=1 python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
       - name: Run pytest (${{ matrix.backend }})
         shell: bash
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
-          python -m pytest -n=auto test/test_tiny.py test/backend/test_ops.py --durations=20
+          python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
 
 # ****** Compile-only Tests ******
 
@@ -1001,36 +955,13 @@ jobs:
           key: compile-${{ matrix.backend }}
           deps: testing_unit
           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
+          DEBUG=4 python3 test/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          python -m pytest -n=auto test/test_ops.py --durations=20
-  qcomclcompiletests:
-    name: Compile-only (QCOM CL)
-    runs-on: ubuntu-24.04-arm
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: compile-qcomcl
-          deps: testing_unit
-          tinydreno: 'true'
-          python-version: '3.12'
-      - name: Set env
-        shell: bash
-        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\nNULL_QCOMCL=1" >> $GITHUB_ENV
-      - name: Run test_ops
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20

.gitignore

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

@@ -72,7 +72,7 @@ vliw_prepare = PatternMatcher([
   # 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])),
+  (UPat(Ops.DEFINE_GLOBAL, 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

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 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)=}"
+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())
+    X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
-    Y[idx].flatten().assign(y.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())
+      boxes[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
-      labels[idx].flatten().assign(clipped_labels.tobytes())
+      labels[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
-      matches[idx].flatten().assign(match_idxs.tobytes())
+      matches[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
-      anchors[idx].flatten().assign(anchor.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.stack(batch, dim=0)
-    yield Tensor(stacked, device="NPY")
 
 def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
   return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)

examples/mlperf/model_train.py

@@ -3,7 +3,7 @@ from pathlib import Path
 import multiprocessing
 
 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
-from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker, DEBUG
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker
 from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
 
@@ -1282,10 +1282,9 @@ def train_bert():
         previous_step = i
 
 def train_llama3():
-  from examples.mlperf.models.llama import Transformer
+  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,7 +1294,6 @@ def train_llama3():
   grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
   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)
@@ -1335,16 +1333,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
+  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
-  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
-  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
-
-  model = Transformer(**model_params, max_context=SEQLEN)
-
   params = get_parameters(model)
   # weights are all bfloat16 for now
   assert params and all(p.dtype == dtypes.bfloat16 for p in params)
@@ -1358,8 +1350,6 @@ def train_llama3():
     for v in get_parameters(model):
       v.shard_(device, axis=None)
 
-    vocab_mask.shard_(device, axis=None)
-
   if (MP := getenv("MP", 1)) > 1:
     device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
     for k,v in get_state_dict(model).items():
@@ -1367,7 +1357,6 @@ def train_llama3():
       elif '.attention.wq' in k: v.shard_(device, axis=0)
       elif '.attention.wk' in k: v.shard_(device, axis=0)
       elif '.attention.wv' in k: v.shard_(device, axis=0)
-      elif '.attention.wqkv' in k: v.shard_(device, axis=0)
       elif '.attention.wo' in k: v.shard_(device, axis=1)
       elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
       elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
@@ -1380,22 +1369,13 @@ def train_llama3():
       # prevents memory spike on device 0
       v.realize()
 
-    vocab_mask.shard_(device, axis=2).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)
-  is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
-  is_fake_offload = Device.DEFAULT == "NULL"
-  optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
-  optim = GradAccClipAdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
-                           eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
 
   # init grads
-  if is_offload_optim:
+  for p in optim.params:
-    for p in optim.params:
+    p.grad = p.zeros_like().contiguous().realize()
-      p.grad = Tensor.zeros(p.shape, dtype=p.dtype, device=optim_device, requires_grad=False).contiguous().realize()
+  grads = [p.grad for p in optim.params]
-  else:
-    for p in optim.params:
-      p.grad = p.zeros_like().contiguous().realize()
-  grads: list[Tensor] = [p.grad for p in optim.params]
 
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
 
@@ -1412,60 +1392,68 @@ def train_llama3():
   def minibatch(tokens:Tensor):
     if (DP := getenv("DP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.to(None).shard(device, 0)
+      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)
-    if DP == 1 and MP == 1: tokens = tokens.to(None)
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    logits:Tensor = model(tokens[:, :-1])
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
     loss.backward()
     assert all(p.grad is g for p,g in zip(optim.params, grads))
-    loss_cpu = loss.flatten().float().to("CPU")
+    Tensor.realize(loss, *grads)
-    Tensor.realize(loss_cpu, *grads)
+    return loss
-    return loss_cpu
 
   @TinyJit
   def optim_step():
-    grad_norm = optim.fstep(grads)
+    for p in optim.params:
+      p.grad.assign(p.grad / grad_acc)
+
+    # 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 g in grads:
+        total_norm += g.float().square().sum()
+      total_norm = total_norm.sqrt().contiguous().realize()
+      for g in grads:
+        g.assign((g * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype)).realize()
+
+    optim.step()
     scheduler.step()
 
     for g in grads:
-      g.assign(g.zeros_like())
+      g.assign(g.zeros_like().contiguous()).realize()
 
-    lr_cpu = optim.lr.float().to("CPU")
+    lr = optim.lr
-    grad_norm_cpu = grad_norm.float().to("CPU")
+    Tensor.realize(lr, *grads)
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads)
 
-    return lr_cpu, grad_norm_cpu
+    return lr
 
   @TinyJit
   @Tensor.train(False)
   def eval_step(tokens:Tensor):
     if (DP := getenv("DP", 1)) > 1:
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.to(None).shard(device, 0)
+      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)
-    if DP == 1 and MP == 1: tokens = tokens.to(None)
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    logits:Tensor = model(tokens[:, :-1])
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    return loss.flatten().float()
-    return loss.flatten().float().to("CPU")
 
   # ** 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.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
-      yield Tensor(fake_data_np, device="NPY")
 
   def get_train_iter():
     if getenv("FAKEDATA", 0):
       return fake_data(BS, SAMPLES)
     else:
       from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=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
@@ -1485,53 +1473,49 @@ def train_llama3():
   step_times = []
   while i < MAX_STEPS:
     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):
-      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()
+        dt = time.perf_counter()
-        data_time += mst - ist
+        loss = minibatch(tokens)
-        losses.append(minibatch(tokens).item())
-        dev_time += time.perf_counter() - mst
       if stopped: break
 
       gt = time.perf_counter()
-      ret = optim_step()
+      lr = optim_step()
-      lr, grad_norm = ret[0].item(), ret[1].item()
+      ot = time.perf_counter()
-      et = time.perf_counter()
 
-      loss = sum(losses) / len(losses)
+      loss = loss.float().item()
-      optim_time = et - gt
+      lr = lr.item()
-      dev_time += optim_time
+
+      et = time.perf_counter()
       step_time = et - st
       gbs_time = gt - st
+      optim_time = ot - gt
+      data_time = dt - ist
+      dev_time = step_time - data_time * grad_acc
       if BENCHMARK: step_times.append(step_time)
 
       i += 1
-      sequences_seen += actual_gbs
+      sequences_seen += GBS
 
       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
       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")
+          f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
-      if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
 
       if WANDB:
         wandb.log({
-          "train/loss": loss,
+          "lr": lr, "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,
@@ -1560,7 +1544,7 @@ 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 (sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
       if EVAL_BS == 0: return
       tqdm.write(f"evaluating after {sequences_seen} sequences")
       profile_marker(f"eval @ {i}")
@@ -1568,7 +1552,7 @@ def train_llama3():
       # 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()

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/optim.py

@@ -1,57 +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]]:
-    for i in range(len(grads)):
-      if grads[i].device != self.m[i].device: grads[i].assign(grads[i].to(self.m[i].device))
-
-    if self.fused:
-      grads[0].assign(grads[0] / self.grad_acc)
-      total_norm = grads[0].float().square().sum().sqrt()
-      grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
-    else:
-      for i in range(len(grads)):
-        grads[i].assign(grads[i] / self.grad_acc)
-      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
-      for i in range(len(grads)):
-        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype))
-
-    ret = []
-    self.b1_t *= self.b1
-    self.b2_t *= self.b2
-    for i, g in enumerate(grads):
-      self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
-      self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
-      m_hat = (self.m[i] / (1.0 - self.b1_t)).cast(self.m[i].dtype)
-      v_hat = (self.v[i] / (1.0 - self.b2_t)).cast(self.v[i].dtype)
-      up = m_hat / (v_hat.sqrt() + self.eps)
-      ret.append((self.lr * up).cast(g.dtype))
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
-
-  def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
-    wd = self.wd if t.ndim >= 2 else 0.0
-    up = up.shard_like(t) + self.lr.to(t.device) * wd * t.detach()
-    return t.detach() - up.cast(t.dtype)

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

@@ -1,37 +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:-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

@@ -2,20 +2,15 @@
 
 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 FLASH_ATTENTION=${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 DP=8 BS=8 EVAL_BS=8 GRADIENT_ACC_STEPS=2
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
@@ -23,20 +18,15 @@ 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 LR="2.5e-4" END_LR="2.5e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
 
-export SEED=${SEED:-5760}
+export SEED=5760
-export DATA_SEED=${DATA_SEED:-5760}
 
-export JITBEAM=${JITBEAM:-3}
+export JITBEAM=3
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 
-export FAKEDATA=1 BENCHMARK=10
+export FAKEDATA=1 BENCHMARK=10 LLAMA_LAYERS=2
-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

@@ -5,17 +5,15 @@ 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 FLASH_ATTENTION=${FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-0}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
@@ -23,15 +21,13 @@ 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 LR="2.5e-4" END_LR="2.5e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
 
-export SEED=${SEED:-$RANDOM}
+export 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 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,6 +0,0 @@
-#!/bin/bash
-export BENCHMARK=5
-export EVAL_BS=0
-export VIZ=${VIZ:--1}
-examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
-extra/viz/cli.py --profile --device "AMD" --top 20

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

@@ -2,9 +2,8 @@
 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
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

examples/openpilot/compile3.py

@@ -31,7 +31,7 @@ def compile(onnx_file):
   for i in range(3):
     GlobalCounters.reset()
     print(f"run {i}")
-    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1), 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

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,6)|(13,0,12): table_t = dev.smu.smu_mod.MetricsTableX_t
-        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTable_t
         case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
       tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
     return tables
@@ -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,6)|(13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
-      case (13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.SocketPowerLimit)
       case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
 
   def get_mem_usage(self, dev):
+    return 0
+
     usage = 0
     pt_stack = [dev.mm.root_page_table]
     while len(pt_stack) > 0:
@@ -244,8 +244,8 @@ class SMICtx:
         entry = pt.entries[i]
 
         if (entry & am.AMDGPU_PTE_VALID) == 0: continue
-        if pt.lv < am.AMDGPU_VM_PDB0 and not dev.gmc.is_pte_huge_page(pt.lv, entry):
+        if pt.lv!=am.AMDGPU_VM_PTB and not dev.gmc.is_pte_huge_page(pt.lv, entry):
-          pt_stack.append(AMPageTableEntry(dev, dev.xgmi2paddr(entry & 0x0000FFFFFFFFF000), lv=pt.lv+1))
+          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, PCIDevImplBase
-from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev
 
 if __name__ == "__main__":
   gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
-  for gpu in gpus:
+  pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
-    drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
-    if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
-      raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
-  pcidevs = [PCIDevice("AM", gpu, bars=[0, 2, 5]) for gpu in gpus]
   amdevs = []
   with Context(DEBUG=2):
     for pcidev in pcidevs:

extra/amdpci/proclogs.py

@@ -7,8 +7,8 @@ class GFXFake:
   def __init__(self): self.xccs = 8
 
 class AMDFake(AMDev):
-  def __init__(self, pci_dev):
+  def __init__(self, pci_dev, dma_regions=None):
-    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
+    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/`
+Test with `PYTHONPATH="." pytest -n12 extra/assembly/amd/`
-`AMD_LLVM=1 pytest -n12 test/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`
+* pcode.py -- pseudocode execution environment. pseudocode should be transformed as little as possible.
-* generate.py -- extract assembly format + instruction pseudocode from AMD XML + PDF
+* asm.py -- an asm/disasm function to transform to and from AMD assembly syntax
-* test/mockgpu/amd/pcode.py -- pseudocode to UOp transformation
+* emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
-* sqtt.py -- SQTT parser
 
-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`
+`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_dtype_alu.py test/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=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 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/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=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.
+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 tinygrad/renderer/amd/*.py`
+Get line count with `cloc --by-file extra/assembly/amd/*.py`
 

extra/assembly/amd/__init__.py

@@ -0,0 +1,67 @@
+# 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, 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, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
+  VOPC as R4_VOPC, VOPD as R4_VOPD, 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, 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, VOP3P_MFMA as C_VOP3P_MFMA, 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_SOP1_LIT, R4_SOPC, R4_SOPC_LIT, R4_SOPP, R4_SOPK, R4_SOPK_LIT, R4_VOPC, R4_VOP1_SDST, R4_VOP1, R4_VOP1_LIT,
+            R4_SOP2, R4_SOP2_LIT, R4_VOP2, R4_VOP2_LIT],
+  "cdna": [C_VOP3PX2, C_VOP3P_MFMA, 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/autogen/cdna/enum.py

@@ -1,5 +1,5 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401
+from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401
 
 class DSOp(ReprEnum):
   DS_ADD_U32 = 0

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

@@ -1,7 +1,7 @@
 # autogenerated from AMD ISA XML - do not edit
-# ruff: noqa: E501,F401
+# ruff: noqa: F401,F403
-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 extra.assembly.amd.dsl import *
-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
+from extra.assembly.amd.autogen.cdna.enum import *
 import functools
 
 class DS(Inst):

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 extra.assembly.amd.autogen.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.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/rdna3/enum.py

@@ -1,5 +1,5 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401
+from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401
 
 class DSOp(ReprEnum):
   DS_ADD_U32 = 0

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

@@ -1,7 +1,7 @@
 # autogenerated from AMD ISA XML - do not edit
-# ruff: noqa: E501,F401
+# ruff: noqa: F401,F403
-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 extra.assembly.amd.dsl import *
-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
+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 extra.assembly.amd.autogen.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.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,5 +1,5 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401
+from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401
 
 class DSOp(ReprEnum):
   DS_ADD_U32 = 0

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

@@ -1,7 +1,7 @@
 # autogenerated from AMD ISA XML - do not edit
-# ruff: noqa: E501,F401
+# ruff: noqa: F401,F403
-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 extra.assembly.amd.dsl import *
-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
+from extra.assembly.amd.autogen.rdna4.enum import *
 import functools
 
 class DS(Inst):
@@ -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 extra.assembly.amd.autogen.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.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/disasm.py

@@ -1,16 +1,14 @@
 # RDNA3/RDNA4/CDNA disassembler
 from __future__ import annotations
-import re
+import re, struct
 from typing import Callable
-from tinygrad.renderer.amd.dsl import Inst, Reg
+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',
+                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'}
-                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',
+                     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',
-                     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,26 +70,23 @@ 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',
+  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
-    '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,
+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,
-  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
+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__
@@ -105,15 +100,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:
+def _vreg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("v", b - 256 if b >= 256 else b, n)
-  b = _unwrap(b)
+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
-  return _reg("v", b - 256 if b >= 256 else b, n)
+def _ttmp(b, n: int = 1) -> str | None: b = _unwrap(b); return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
-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 _fmt_sdst(v, n: int = 1, cdna: bool = False) -> str:
   v = _unwrap(v)
@@ -141,9 +130,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:
+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
-  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 +201,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:
+  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}"
-    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 +217,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)
+  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)
-  elif inst.src0.sz > 1: s0 = inst.src0.fmt()
-  elif is16: s0 = _src16(inst, inst.src0.offset)
-  else: s0 = _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}"
@@ -270,11 +253,10 @@ 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',
+    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']
-            '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}"
@@ -285,7 +267,7 @@ def _disasm_smem(inst: SMEM) -> str:
   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]
+  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 +278,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)
+  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)
-  elif sbase_idx <= 105: sbase_str = _sreg(sbase_idx, sbase_count)
-  else: sbase_str = _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,12 +304,6 @@ 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)
   acc = getattr(inst, 'acc', 0)
@@ -337,10 +311,9 @@ def _disasm_flat(inst: FLAT) -> str:
   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'
   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)
+  # Global/scratch uses 13-bit signed offset
-  offset = inst.ioffset if r4 else inst.offset  # type: ignore[attr-defined]
+  offset = inst.ioffset if r4 else inst.offset
-  if r4: off_val = offset if offset < (1 << 23) else offset - (1 << 24)  # sign extend 24-bit
+  if seg != 'flat':
-  elif seg != 'flat':
     if cdna:
       # CDNA: bit 12 is sign bit but not in offset field
       raw = int.from_bytes(inst.to_bytes(), 'little')
@@ -354,22 +327,19 @@ def _disasm_flat(inst: FLAT) -> str:
   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]
+  if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"
-  elif r4:
+  elif r4: mods = f"{off_s}{' scope' if inst.scope else ''}{' th' if inst.th else ''}"
-    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 "")
   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):
+  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)]}"
-    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}"
+  if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
-   # RDNA4: vaddr instead of addr, vsrc instead of data
+  # RDNA4: vaddr instead of addr, vsrc instead of data
-  addr = inst.vaddr if r4 else inst.addr  # type: ignore[attr-defined]
+  addr = inst.vaddr if r4 else inst.addr
-  data = inst.vsrc if r4 else inst.data  # type: ignore[attr-defined]
+  data = inst.vsrc if r4 else inst.data
   # 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
@@ -381,14 +351,13 @@ def _disasm_flat(inst: FLAT) -> str:
   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)
   if 'atomic' in name:
-    glc_or_sc0 = inst.sc0 if cdna else (inst.th & 1 if r4 else inst.glc)  # type: ignore[attr-defined]
+    glc_or_sc0 = inst.sc0 if cdna else inst.glc
-    sfx = f"{saddr_s if seg != 'flat' else ''}{mods}"
+    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}"
-    return f"{instr} {vdst_s}, {addr_s}, {data_s}{sfx}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{sfx}"
   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 +386,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):
+  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}"
-    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 +397,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)
+    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 s0v == 255: src = _lit(inst, inst.src0)
-    elif s0v == 253: src = "src_scc"
-    else: src = _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 +412,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 +428,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))
+  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))
-                   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
@@ -512,7 +478,7 @@ def _disasm_vopd(inst: VOPD) -> str:
 
 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 +487,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 [])
+    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 [])
-            + ([_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 [])
+    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("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 [])
-            + ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg 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 ''}"
-            + ([_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}"
 
 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 +512,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',
+  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_BARRIER_INIT', 'S_BARRIER_JOIN', 'S_SET_GPR_IDX_IDX', 'S_CBRANCH_JOIN')
+                   '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,7 +551,7 @@ _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)
+  op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
   is_rdna4 = _is_r4(inst)
   hw = HWREG_CDNA if cdna else (HWREG_RDNA4 if is_rdna4 else HWREG)
   blacklist = _HWREG_BLACKLIST_CDNA if cdna else _HWREG_BLACKLIST
@@ -612,14 +574,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()}, {_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 "")
-  return f"{inst.op_name.lower()} {inst.vdst.fmt()}, {s0}, {s1}, {s2}" + (" " + mods if mods else "")
 
 DISASM_HANDLERS: dict[type, Callable[..., str]] = {
   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,
+  VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3_SDST_LIT: _disasm_vop3, VOP3_LIT: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOP3SD_LIT: _disasm_vop3sd,
-  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,
@@ -640,7 +600,7 @@ 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,
@@ -674,9 +634,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)
+    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)
-    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)
     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 +658,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)
+  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)
-  s1 = _cdna_src(inst, inst.src1, inst.neg&2, n=r1)
-  s2 = _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 +680,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 +742,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 [])
+  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("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 [])
-          + ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg 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 ''}"
-          + ([_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}"
 
 def _disasm_mubuf(inst) -> str:
   name = inst.op_name.lower()
@@ -950,6 +903,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_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p, CDNA_VOP3P_MFMA: _disasm_cdna_vop3p,
-  CDNA_VOP3P: _disasm_cdna_vop3p, CDNA_VOP3P_MFMA: _disasm_cdna_vop3p,
   CDNA_MUBUF: _disasm_mubuf, CDNA_VOP3PX2: _disasm_vop3px2})

extra/assembly/amd/dsl.py

@@ -44,15 +44,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:
+    if 256 <= o < 512: idx = o - 256; base = f"v{l}{idx}{r}" if sz == 1 else f"v[{idx}:{idx + sz - 1}]"
-      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:
+    elif 108 <= o < 124: idx = o - 108; base = f"ttmp{l}{idx}{r}" if sz == 1 else f"ttmp[{idx}:{idx + sz - 1}]"
-      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}")
@@ -99,7 +95,7 @@ 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]
@@ -155,8 +151,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}) "
+      raise RuntimeError(f"{self.__class__.__name__}: field size {hi - lo + 1} bits ({actual_size}) doesn't match range {self._valid_range} ({expected_size})")
-                         f"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."""
@@ -236,9 +231,9 @@ class VDSTYField(BitField):
 # ══════════════════════════════════════════════════════════════
 
 import functools
-from tinygrad.runtime.autogen.amd.rdna3.operands import OPERANDS as OPERANDS_RDNA3
+from extra.assembly.amd.autogen.rdna3.operands import OPERANDS as OPERANDS_RDNA3
-from tinygrad.runtime.autogen.amd.rdna4.operands import OPERANDS as OPERANDS_RDNA4
+from extra.assembly.amd.autogen.rdna4.operands import OPERANDS as OPERANDS_RDNA4
-from tinygrad.runtime.autogen.amd.cdna.operands import OPERANDS as OPERANDS_CDNA
+from extra.assembly.amd.autogen.cdna.operands import OPERANDS as OPERANDS_CDNA
 OPERANDS = {**OPERANDS_CDNA, **OPERANDS_RDNA3, **OPERANDS_RDNA4}
 
 # ══════════════════════════════════════════════════════════════
@@ -276,7 +271,7 @@ 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
@@ -408,7 +403,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/generate.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
 
 # ═══════════════════════════════════════════════════════════════════════════════
@@ -77,13 +77,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((ot_name, field.findtext("FieldName")))):
-      if (enum_name := op_enum_map.get(key)):  # type: ignore[arg-type]
+        enums[enum_name] = {int(pv.findtext("Value")): pv.findtext("Name").upper() for pv in field.findall(".//PredefinedValue")}
-        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"))
   # Extract DataFormats with BitCount
   for df in root.findall("ISA/DataFormats/DataFormat"):
     name, bits = df.findtext("DataFormatName"), df.findtext("BitCount")
@@ -91,26 +86,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]] = []
+    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,
-    for f in enc.findall(".//MicrocodeFormat/BitMap/Field"):
+               int(f.find("BitLayout/Range").findtext("BitOffset") or 0))
-      br = f.find("BitLayout/Range")
+              for f in enc.findall(".//MicrocodeFormat/BitMap/Field") if f.find("BitLayout/Range") is not None]
-      if br is None: continue
+    ident = (enc.findall("EncodingIdentifiers/EncodingIdentifier") or [None])[0]
-      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
     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
+    # For multi-dword formats, encoding field may be in higher dword but identifier pattern is always in dword0; use % 32
-    enc_bits: str | None = None
+    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
-    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))
     base_name = _strip_enc(name)
     encodings[NAME_MAP.get(base_name, base_name)] = (fields, enc_bits)
   # Extract instruction opcodes and operand info
@@ -118,12 +104,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")
+      base, opcode = _map_flat(_strip_enc(enc.findtext("EncodingName")), name), int(enc.findtext("Opcode") or 0)
-      assert enc_enc_name is not None
-      base, opcode = _map_flat(_strip_enc(enc_enc_name), 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,10 +120,8 @@ 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]] = {}
+      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"):
+                 for op in enc.findall("Operands/Operand") if op.findtext("FieldName")}
-        fn = op.findtext("FieldName")
-        if fn: op_info[fn.lower()] = (op.findtext("DataFormatName"), int(op.findtext("OperandSize") or 0), op.findtext("OperandType"))
       for fmt, _, otype in op_info.values():
         if fmt and fmt not in fmts: fmts[fmt] = 0
         if otype: op_types_set.add(otype)
@@ -162,9 +143,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)
+  pos = int(re.search(rb'startxref\s+(\d+)', data).group(1)) + 4
-  assert xref_match is not None
-  pos = int(xref_match.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 +164,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)):
+      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 sm.group(1): font = sm.group(1)
+        if m.group(1): font = m.group(1)
-        elif sm.group(2): x, y = x + float(sm.group(2)), y + float(sm.group(3))
+        elif m.group(2): x, y = x + float(m.group(2)), y + float(m.group(3))
-        elif sm.group(4): x, y = float(sm.group(4)), float(sm.group(5))
+        elif m.group(4): x, y = float(m.group(4)), float(m.group(5))
-        elif sm.group(6) and (t := bytes.fromhex(sm.group(6)).decode('latin-1')).strip():
+        elif m.group(6) and (t := bytes.fromhex(m.group(6)).decode('latin-1')).strip(): elements.append((x, y, t, font))
-          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))
-        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))
     pages.append(sorted(elements, key=lambda e: (-e[1], e[0])))
   return pages
 
@@ -223,7 +197,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,8 +220,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:
+def write_common(all_fmts, all_op_types, path):
-  lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
+  lines = ["# 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}"')
@@ -264,8 +238,7 @@ 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",
+  lines = ["# autogenerated from AMD ISA XML - do not edit", "from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401", ""]
-    "from tinygrad.runtime.autogen.amd.common import ReprEnum, 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 ""
@@ -313,7 +286,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 +296,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
+  lines = ["# autogenerated from AMD ISA XML - do not edit", "# ruff: noqa: F401,F403",
-  enum_names = sorted(f"{k}Op" for k in enums if enums[k] and k not in ("HWREG", "MSG"))
+           "from extra.assembly.amd.dsl import *", f"from extra.assembly.amd.autogen.{arch}.enum import *", "import functools", ""]
-  # 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", ""]
 
   def fmt_allowed(op_enum: str, ops: set[int]) -> str:
     """Format allowed ops as {EnumName.MEMBER, ...}."""
@@ -364,9 +323,7 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
     has_seg_field = any(fn == "seg" for fn, _, _ in fields)
     if enc_name in ("FLAT", "VFLAT") and has_seg_field:
       prefix = "V" if enc_name == "VFLAT" else ""
-      flat_variants = [(f"{prefix}FLAT", 0, f"{prefix}FLATOp"), (f"{prefix}GLOBAL", 2, f"{prefix}GLOBALOp"),
+      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")]:
-                       (f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]
-      for cls, seg, op_enum in flat_variants:
         cls_ops = set(enums.get(cls, {}).keys())
         lines.append(f"class {cls}(Inst):")
         for fn, hi, lo in sort_fields(fields):
@@ -439,8 +396,6 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
     op_to_suffix = {op:suffix for suffix,ops in suffix_only_ops.items() for op in 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
       if fmt == "VOP1" and op in fmt_sdst_ops: cls = "VOP1_SDST"
@@ -450,14 +405,11 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
       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
+  lines = ["# autogenerated from AMD ISA XML - do not edit",
-  used_bases = {eb for (nm, eb) in types if (nm, eb) in valid}
+           "from extra.assembly.amd.autogen.common import Fmt, OpType",
-  enum_names = sorted(f"{k}Op" for k in used_bases)
+           f"from extra.assembly.amd.autogen.{arch}.enum import *", ""]
-  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.append("# instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}")
   lines.append("OPERANDS = {")
   def fmt_val(v):
@@ -470,7 +422,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 +433,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,9 +444,8 @@ def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, st
 # ═══════════════════════════════════════════════════════════════════════════════
 
 if __name__ == "__main__":
-  all_fmts: dict[str, int] = {}
+  import pathlib
-  all_op_types: set[str] = set()
+  all_fmts, all_op_types, arch_data = {}, set(), {}
-  arch_data: dict[str, dict] = {}
   # First pass: parse XML for all architectures
   for arch, cfg in ARCHS.items():
     print(f"Parsing XML: {cfg['xml']} -> {arch}")
@@ -508,13 +459,12 @@ if __name__ == "__main__":
       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 = pathlib.Path(__file__).parent / "autogen" / "common.py"
-  common_path = autogen_base / "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_operands(data["types"], data["enums"], arch, base / "operands.py")
@@ -525,6 +475,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/pcode.py

@@ -40,10 +40,7 @@ def _bitreverse(v: UOp, bits: int) -> UOp:
 
 def _extract_bits(val: UOp, hi: int, lo: int) -> UOp:
   dt = dtypes.uint64 if val.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
-  result = ((val >> _const(dt, lo)) if lo > 0 else val) & _const(val.dtype, (1 << (hi - lo + 1)) - 1)
+  return ((val >> _const(dt, lo)) if lo > 0 else val) & _const(val.dtype, (1 << (hi - lo + 1)) - 1)
-  # Downcast to uint32 when extracting <=32 bits from a 64-bit value, so .f32 bitcast works correctly
-  if dt == dtypes.uint64 and (hi - lo + 1) <= 32: result = result.cast(dtypes.uint32)
-  return result
 
 def _set_bit(old, pos, val):
   mask = _u32(1) << pos
@@ -55,9 +52,7 @@ def _val_to_bits(val):
   if val.dtype == dtypes.float64: return val.bitcast(dtypes.uint64)
   return val if val.dtype == dtypes.uint32 else val.cast(dtypes.uint32)
 
-def _floor(x):
+def _floor(x): t = UOp(Ops.TRUNC, x.dtype, (x,)); return ((x < _const(x.dtype, 0)) & x.ne(t)).where(t - _const(x.dtype, 1), t)
-  t = UOp(Ops.TRUNC, x.dtype, (x,))
-  return ((x < _const(x.dtype, 0)) & x.ne(t)).where(t - _const(x.dtype, 1), t)
 def _f16_extract(v): return (v & _u32(0xFFFF)).cast(dtypes.uint16).bitcast(dtypes.half) if v.dtype == dtypes.uint32 else v
 
 def _check_nan(v: UOp, quiet: bool) -> UOp:
@@ -123,8 +118,7 @@ def _f_to_u(f, dt): return UOp(Ops.TRUNC, f.dtype, ((f < _const(f.dtype, 0.0)).w
 
 def _cvt_quiet(val: UOp) -> UOp:
   bits, _, _, qb, _ = _float_info(val)
-  bt, ft = (dtypes.uint64, dtypes.float64) if val.dtype == dtypes.float64 else \
+  bt, ft = (dtypes.uint64, dtypes.float64) if val.dtype == dtypes.float64 else (dtypes.uint16, dtypes.half) if val.dtype == dtypes.half else (dtypes.uint32, dtypes.float32)
-    (dtypes.uint16, dtypes.half) if val.dtype == dtypes.half else (dtypes.uint32, dtypes.float32)
   return (val.bitcast(bt) | qb).bitcast(ft)
 
 def _is_denorm(val: UOp) -> UOp:
@@ -169,18 +163,14 @@ def _ldexp(val: UOp, exp: UOp) -> UOp:
 def _frexp_mant(val: UOp) -> UOp:
   val = val.bitcast(dtypes.float32) if val.dtype == dtypes.uint32 else val.bitcast(dtypes.float64) if val.dtype == dtypes.uint64 else val
   if val.dtype == dtypes.float32: return ((val.bitcast(dtypes.uint32) & _u32(0x807FFFFF)) | _u32(0x3f000000)).bitcast(dtypes.float32)
-  return ((val.bitcast(dtypes.uint64) & _const(dtypes.uint64, 0x800FFFFFFFFFFFFF)) |
+  return ((val.bitcast(dtypes.uint64) & _const(dtypes.uint64, 0x800FFFFFFFFFFFFF)) | _const(dtypes.uint64, 0x3fe0000000000000)).bitcast(dtypes.float64)
-    _const(dtypes.uint64, 0x3fe0000000000000)).bitcast(dtypes.float64)
 
 def _frexp_exp(val: UOp) -> UOp:
   val = val.bitcast(dtypes.float32) if val.dtype == dtypes.uint32 else val.bitcast(dtypes.float64) if val.dtype == dtypes.uint64 else val
   if val.dtype == dtypes.float32: return ((val.bitcast(dtypes.uint32) >> _u32(23)) & _u32(0xFF)).cast(dtypes.int) - _const(dtypes.int, 126)
   return ((val.bitcast(dtypes.uint64) >> _const(dtypes.uint64, 52)) & _const(dtypes.uint64, 0x7FF)).cast(dtypes.int) - _const(dtypes.int, 1022)
 
-TWO_OVER_PI = int(
+TWO_OVER_PI = 0x0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6
-  "0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd"
-  "63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414"
-  "da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6", 16)
 # TWO_OVER_PI as 19 u64 words for trig_preop_result (word[0] = bits 0-63, word[18] = bits 1152-1200)
 _PREOP_WORDS = tuple((TWO_OVER_PI >> (64 * i)) & 0xFFFFFFFFFFFFFFFF for i in range(19))
 def _trig_preop(val: UOp) -> UOp:
@@ -257,14 +247,10 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
   # Normalization conversions: map [-1,1] or [0,1] to integer range
   # Use floor(x + 0.5) for round-to-nearest
   # SNORM: round(value * 32767), range is [-32767, 32767] (hardware behavior)
-  'f16_to_snorm': lambda a: _floor(
+  'f16_to_snorm': lambda a: _floor(_f16_extract(a).cast(dtypes.float32) * _const(dtypes.float32, 32767) + _const(dtypes.float32, 0.5)).cast(dtypes.int).cast(dtypes.int16),
-    _f16_extract(a).cast(dtypes.float32) * _const(dtypes.float32, 32767) + _const(dtypes.float32, 0.5)).cast(dtypes.int).cast(dtypes.int16),
+  'f16_to_unorm': lambda a: _floor(_f16_extract(a).cast(dtypes.float32) * _const(dtypes.float32, 65535) + _const(dtypes.float32, 0.5)).cast(dtypes.uint16),
-  'f16_to_unorm': lambda a: _floor(
+  'f32_to_snorm': lambda a: _floor(a.bitcast(dtypes.float32) * _const(dtypes.float32, 32767) + _const(dtypes.float32, 0.5)).cast(dtypes.int).cast(dtypes.int16),
-    _f16_extract(a).cast(dtypes.float32) * _const(dtypes.float32, 65535) + _const(dtypes.float32, 0.5)).cast(dtypes.uint16),
+  'f32_to_unorm': lambda a: _floor(a.bitcast(dtypes.float32) * _const(dtypes.float32, 65535) + _const(dtypes.float32, 0.5)).cast(dtypes.uint16),
-  'f32_to_snorm': lambda a: _floor(
-    a.bitcast(dtypes.float32) * _const(dtypes.float32, 32767) + _const(dtypes.float32, 0.5)).cast(dtypes.int).cast(dtypes.int16),
-  'f32_to_unorm': lambda a: _floor(
-    a.bitcast(dtypes.float32) * _const(dtypes.float32, 65535) + _const(dtypes.float32, 0.5)).cast(dtypes.uint16),
   'f32_to_u8': lambda a: _f_to_u(a.bitcast(dtypes.float32), dtypes.uint8),
   # Integer truncation conversions
   'i32_to_i16': lambda a: a.cast(dtypes.int).cast(dtypes.int16),
@@ -285,10 +271,6 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
   # System NOPs - these are scheduling hints, no effect on emulation
   'MIN': lambda a, b: (a < b).where(a, b),
   's_nop': lambda a: _u32(0),
-  # Address calculation for memory operations
-  'CalcDsAddr': lambda a, o, *r: a.cast(dtypes.uint32) + o.cast(dtypes.uint32),
-  'CalcGlobalAddr': lambda v, s, *r: v.cast(dtypes.uint64) + s.cast(dtypes.uint64),
-  'CalcScratchAddr': lambda v, s, *r: v.cast(dtypes.uint64) + s.cast(dtypes.uint64),
 }
 for is_max, name in [(False, 'min'), (True, 'max')]:
   for dt, sfx in [(dtypes.float32, 'f32'), (dtypes.int, 'i32'), (dtypes.uint32, 'u32'), (dtypes.int16, 'i16'), (dtypes.uint16, 'u16')]:
@@ -325,35 +307,21 @@ _SINGLE_CHAR = {'(': 'LPAREN', ')': 'RPAREN', '[': 'LBRACKET', ']': 'RBRACKET',
 
 class Token:
   __slots__ = ('type', 'val')
-  def __init__(self, kind: str, val: str): self.type, self.val = kind, val
+  def __init__(self, type: str, val: str): self.type, self.val = type, val
   def __repr__(self): return f'{self.type}:{self.val}'
 
 def tokenize(s: str) -> list[Token]:
   tokens, i, n = [], 0, len(s)
   while i < n:
     c = s[i]
-    if c.isspace():
+    if c.isspace(): i += 1; continue
-      i += 1
-      continue
     if i + 1 < n and s[i:i+2] in ('+=', '-='):
-      tokens.append(Token('ASSIGN_OP', s[i:i+2]))
+      tokens.append(Token('ASSIGN_OP', s[i:i+2])); i += 2; continue
-      i += 2
-      continue
     if i + 1 < n and s[i:i+2] in ('||', '&&', '>=', '<=', '==', '!=', '<>', '>>', '<<', '**', '+:', '-:'):
-      tokens.append(Token('OP', s[i:i+2]))
+      tokens.append(Token('OP', s[i:i+2])); i += 2; continue
-      i += 2
+    if c in '|^&><+-*/~!%': tokens.append(Token('OP', c)); i += 1; continue
-      continue
+    if (t := _SINGLE_CHAR.get(c)): tokens.append(Token(t, c)); i += 1; continue
-    if c in '|^&><+-*/~!%':
+    if c == ';': i += 1; continue
-      tokens.append(Token('OP', c))
-      i += 1
-      continue
-    if (t := _SINGLE_CHAR.get(c)):
-      tokens.append(Token(t, c))
-      i += 1
-      continue
-    if c == ';':
-      i += 1
-      continue
     if c.isdigit() or (c == '-' and i + 1 < n and s[i+1].isdigit()):
       start = i
       if c == '-': i += 1
@@ -366,38 +334,31 @@ def tokenize(s: str) -> list[Token]:
           i += 1
           while i < n and s[i].isdigit(): i += 1
       for sfx in ('ULL', 'LL', 'UL', 'U', 'L', 'F', 'f'):
-        if s[i:i+len(sfx)] == sfx:
+        if s[i:i+len(sfx)] == sfx: i += len(sfx); break
-          i += len(sfx)
+      tokens.append(Token('NUM', s[start:i])); continue
-          break
-      tokens.append(Token('NUM', s[start:i]))
-      continue
     if c.isalpha() or c == '_':
       start = i
       while i < n and (s[i].isalnum() or s[i] == '_'): i += 1
-      tokens.append(Token('IDENT', s[start:i]))
+      tokens.append(Token('IDENT', s[start:i])); continue
-      continue
     raise RuntimeError(f"unexpected char '{c}' at pos {i} in: {s}")
   tokens.append(Token('EOF', ''))
   return tokens
 
 class Parser:
-  def __init__(self, tokens: list[Token], env: dict, funcs: dict | None = None):
+  def __init__(self, tokens: list[Token], vars: dict, funcs: dict | None = None):
-    self.tokens, self.vars, self.funcs, self.pos = tokens, env, funcs if funcs is not None else _FUNCS, 0
+    self.tokens, self.vars, self.funcs, self.pos = tokens, vars, funcs if funcs is not None else _FUNCS, 0
 
   def peek(self, offset=0) -> Token: return self.tokens[min(self.pos + offset, len(self.tokens) - 1)]
   def at(self, *types) -> bool: return self.peek().type in types
-  def _advance(self) -> Token:
+  def _advance(self) -> Token: tok = self.tokens[self.pos]; self.pos += 1; return tok
-    tok = self.tokens[self.pos]
+  def eat(self, type: str) -> Token:
-    self.pos += 1
+    if self.peek().type != type: raise RuntimeError(f"expected {type}, got {self.peek()}")
-    return tok
-  def eat(self, kind: str) -> Token:
-    if self.peek().type != kind: raise RuntimeError(f"expected {kind}, got {self.peek()}")
     return self._advance()
-  def try_eat(self, kind: str) -> Token | None: return self._advance() if self.peek().type == kind else None
+  def try_eat(self, type: str) -> Token | None: return self._advance() if self.peek().type == type else None
-  def try_eat_val(self, val: str, kind: str) -> Token | None:
+  def try_eat_val(self, val: str, type: str) -> Token | None:
-    return self._advance() if self.peek().type == kind and self.peek().val == val else None
+    return self._advance() if self.peek().type == type and self.peek().val == val else None
-  def eat_val(self, val: str, kind: str) -> Token:
+  def eat_val(self, val: str, type: str) -> Token:
-    if self.peek().type != kind or self.peek().val != val: raise RuntimeError(f"expected {kind}:{val}, got {self.peek()}")
+    if self.peek().type != type or self.peek().val != val: raise RuntimeError(f"expected {type}:{val}, got {self.peek()}")
     return self._advance()
 
   def parse(self) -> UOp:
@@ -416,11 +377,9 @@ class Parser:
       case '||' | '|': return left | right
       case '&&' | '&': return left & right
       case '^': return left ^ right
-      case '==': return left.eq(right)
+      case '==' | '<>': return left.eq(right) if op == '==' else left.ne(right)
-      case '!=': return left.ne(right)
+      case '!=' : return left.ne(right)
-      case '>=' | '<=' | '>' | '<' | '<>':
+      case '>=' | '<=' | '>' | '<': return self._cmp_nan(left, right, {'>=':(lambda a,b:a>=b),'<=':(lambda a,b:a<=b),'>':(lambda a,b:a>b),'<':(lambda a,b:a<b)}[op])
-        ops = {'>=':(lambda a,b:a>=b),'<=':(lambda a,b:a<=b),'>':(lambda a,b:a>b),'<':(lambda a,b:a<b),'<>':(lambda a,b:a.ne(b))}
-        return self._cmp_nan(left, right, ops[op])
       case '>>' | '<<': return (left >> right) if op == '>>' else (left << right)
       case '+' | '-':
         if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.arg - right.arg)
@@ -478,7 +437,7 @@ class Parser:
       self.eat('COMMA')
       lo = self.parse()
       self.eat('RBRACE')
-      return (hi.cast(dt:=_BITS_DT.get((s:=lo.dtype.bitsize) * 2, dtypes.uint64)) << _const(dt, s)) | lo.cast(dt)
+      return (hi.cast(dtypes.uint64) << _u64(32)) | lo.cast(dtypes.uint64)
     if self.at('NUM'):
       num = self.eat('NUM').val
       if self.try_eat('QUOTE'):
@@ -558,9 +517,7 @@ class Parser:
       self.eat('LBRACKET')
       self.eat_val('laneId', 'IDENT')
       self.eat('RBRACKET')
-      lane = self.vars['laneId']
+      result = (base >> _to_u32(self.vars['laneId'])) & _u32(1)
-      shift = lane.cast(base.dtype) if base.dtype != dtypes.uint32 else _to_u32(lane)
-      result = (base >> shift) & _const(base.dtype, 1)
       if self.try_eat('DOT'):
         dt_name = self.eat('IDENT').val
         return result.cast(DTYPES.get(dt_name, dtypes.uint32))
@@ -569,8 +526,7 @@ class Parser:
     if dt is None: return base
     if dt == base.dtype: return base
     if dt.itemsize == 2 and base.dtype.itemsize == 4:
-      if dt == dtypes.uint16: return (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16)
+      return (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16) if dt == dtypes.uint16 else (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16).bitcast(dt)
-      return (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16).bitcast(dt)
     if field == 'i4': return _signext_4bit(base)
     return _cast_to(base, dt)
 
@@ -580,7 +536,7 @@ class Parser:
 
   def _handle_bracket_rest(self, first: UOp, base: UOp, var_name: str | None = None) -> UOp:
     if self.at('OP') and self.peek().val in ('+:', '-:'):
-      self.eat('OP')
+      op = self.eat('OP').val
       width = self.parse()
       self.eat('RBRACKET')
       if width.op == Ops.CONST:
@@ -666,8 +622,7 @@ class Parser:
       inner = self.parse()
       self.eat('RPAREN')
       dt = {('U',32): dtypes.uint32, ('U',64): dtypes.uint64, ('I',32): dtypes.int, ('I',64): dtypes.int64,
-            ('F',16): dtypes.half, ('F',32): dtypes.float32, ('F',64): dtypes.float64,
+            ('F',16): dtypes.half, ('F',32): dtypes.float32, ('F',64): dtypes.float64, ('B',32): dtypes.uint32, ('B',64): dtypes.uint64}.get((type_char, bits), dtypes.uint64 if bits > 32 else dtypes.uint32)
-            ('B',32): dtypes.uint32, ('B',64): dtypes.uint64}.get((type_char, bits), dtypes.uint64 if bits > 32 else dtypes.uint32)
       if type_char == 'F' and inner.dtype in (dtypes.uint32, dtypes.uint64, dtypes.ulong, dtypes.int, dtypes.int64):
         if inner.dtype.itemsize != dt.itemsize: inner = inner.cast(dtypes.uint32 if dt.itemsize == 4 else dtypes.uint64)
         return inner.bitcast(dt)
@@ -728,7 +683,7 @@ class Parser:
   def _call_func(self, name: str, args: list[UOp]) -> UOp:
     if name in self.vars and isinstance(self.vars[name], tuple) and self.vars[name][0] == 'lambda':
       _, params, body = self.vars[name]
-      lv = {**self.vars, **dict(zip(params, args))}
+      lv = {**self.vars, **{p: a for p, a in zip(params, args)}}
       if ';' in body or '\n' in body or 'return' in body.lower():
         lines = [l.strip() for l in body.replace(';', '\n').split('\n') if l.strip() and not l.strip().startswith('//')]
         _, _, result = parse_block(lines, 0, lv, self.funcs)
@@ -754,9 +709,7 @@ class Parser:
       elif dt in (dtypes.uint8, dtypes.int8):
         val = mem.index(idx, *gate, ptr=True).load().cast(dt)
       elif dt in (dtypes.uint16, dtypes.int16, dtypes.short):
-        lo = mem.index(idx, *gate, ptr=True).load().cast(dtypes.uint32)
+        val = (mem.index(idx, *gate, ptr=True).load().cast(dtypes.uint32) | (mem.index(idx + _const(dtypes.int, 1), *gate, ptr=True).load().cast(dtypes.uint32) << _u32(8))).cast(dt)
-        hi = mem.index(idx + _const(dtypes.int, 1), *gate, ptr=True).load().cast(dtypes.uint32)
-        val = (lo | (hi << _u32(8))).cast(dt)
       else:
         val = _u32(0)
         for i in range(4): val = val | (mem.index(idx + _const(dtypes.int, i), *gate, ptr=True).load().cast(dtypes.uint32) << _u32(i * 8))
@@ -767,8 +720,7 @@ class Parser:
         idx2 = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int)
         val = val.cast(dtypes.uint64) | (mem.index(idx2, *gate).cast(dtypes.uint64) << _u64(32))
       elif dt in (dtypes.uint8, dtypes.int8): val = (val >> ((addr & _const(adt, 3)).cast(dtypes.uint32) * _u32(8))) & _u32(0xFF)
-      elif dt in (dtypes.uint16, dtypes.int16):
+      elif dt in (dtypes.uint16, dtypes.int16): val = (val >> (((addr >> _const(adt, 1)) & _const(adt, 1)).cast(dtypes.uint32) * _u32(16))) & _u32(0xFFFF)
-        val = (val >> (((addr >> _const(adt, 1)) & _const(adt, 1)).cast(dtypes.uint32) * _u32(16))) & _u32(0xFFFF)
     return val
 
   def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
@@ -801,8 +753,8 @@ def _match_bracket(toks: list[Token], start: int) -> tuple[int, list[Token]]:
   return j, [t for t in toks[start+1:j-1] if t.type != 'EOF']
 
 def _tok_str(toks: list[Token]) -> str: return ' '.join(t.val for t in toks if t.type != 'EOF')
-def parse_tokens(toks: list[Token], env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
+def parse_tokens(toks: list[Token], vars: dict[str, VarVal], funcs: dict | None = None) -> UOp:
-  return Parser(toks, env, funcs).parse()
+  return Parser(toks, vars, funcs).parse()
 
 # Unified block parser for pcode
 def _subst_loop_var(line: str, loop_var: str, val: int) -> str:
@@ -812,13 +764,6 @@ def _subst_loop_var(line: str, loop_var: str, val: int) -> str:
 
 def _set_bits(old: UOp, val: UOp, width: int, offset: int) -> UOp:
   """Set bits [offset:offset+width) in old to val, masking and shifting appropriately."""
-  if old.dtype in (dtypes.half, dtypes.float32): old = _val_to_bits(old)
-  is64 = old.dtype in (dtypes.uint64, dtypes.int64) or offset + width > 32
-  if is64:
-    old = old.cast(dtypes.uint64) if old.dtype != dtypes.uint64 else old
-    mask = _u64(((1 << width) - 1) << offset)
-    v = (val.cast(dtypes.uint64) if val.dtype != dtypes.uint64 else val) & _u64((1 << width) - 1)
-    return (old & (mask ^ _u64(0xFFFFFFFFFFFFFFFF))) | (v << _u64(offset))
   mask = _u32(((1 << width) - 1) << offset)
   v = (val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val) & _u32((1 << width) - 1)
   return (old & (mask ^ _u32(0xFFFFFFFF))) | (v << _u32(offset))
@@ -833,7 +778,7 @@ def _find_paren_end(s: str, start: int = 0, open_ch: str = '(', close_ch: str =
       if depth == 0: return j
   return len(s)
 
-def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dict | None = None,
+def parse_block(lines: list[str], start: int, vars: dict[str, VarVal], funcs: dict | None = None,
                 assigns: list | None = None) -> tuple[int, dict[str, VarVal], UOp | None]:
   """Parse a block of pcode. Returns (next_line, block_assigns, return_value).
   If assigns list is provided, side effects (MEM/VGPR writes) are appended to it."""
@@ -844,9 +789,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
   while i < len(lines):
     line = lines[i]
     toks = tokenize(line)
-    if toks[0].type != 'IDENT' and toks[0].type != 'LBRACE':
+    if toks[0].type != 'IDENT' and toks[0].type != 'LBRACE': i += 1; continue
-      i += 1
-      continue
     first = toks[0].val.lower() if toks[0].type == 'IDENT' else '{'
 
     # Block terminators
@@ -855,19 +798,17 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
     # return expr (lambda bodies)
     if first == 'return':
       rest = line[line.lower().find('return') + 6:].strip()
-      return i + 1, block_assigns, parse_expr(rest, env, funcs)
+      return i + 1, block_assigns, parse_expr(rest, vars, funcs)
 
     # for loop
     if first == 'for':
       # Parse: for VAR in [SIZE']START : [SIZE']END do
-      p = Parser(toks, env, funcs)
+      p = Parser(toks, vars, funcs)
       p.eat_val('for', 'IDENT')
       loop_var = p.eat('IDENT').val
       p.eat_val('in', 'IDENT')
       def parse_bound():
-        if p.at('NUM') and p.peek(1).type == 'QUOTE':
+        if p.at('NUM') and p.peek(1).type == 'QUOTE': p.eat('NUM'); p.eat('QUOTE')
-          p.eat('NUM')
-          p.eat('QUOTE')
         if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
         expr = p.parse().simplify()
         assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}"
@@ -889,41 +830,38 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       # Execute loop with break support
       has_break = any('break' in bl.lower() for bl in body_lines)
       found_var = f'_found_{id(body_lines)}' if has_break else None
-      if found_var: env[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
+      if found_var: vars[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
       for loop_i in range(start_val, end_val + 1):
         subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
-        _, iter_assigns, _ = parse_block(subst_lines, 0, {**env, **block_assigns}, funcs, assigns)
+        _, iter_assigns, _ = parse_block(subst_lines, 0, {**vars, **block_assigns}, funcs, assigns)
         if has_break:
           assert found_var is not None
-          found = block_assigns.get(found_var, env.get(found_var))
+          found = block_assigns.get(found_var, vars.get(found_var))
           assert isinstance(found, UOp)
           not_found = found.eq(_const(dtypes.bool, False))
           for var, val in iter_assigns.items():
             if var != found_var and isinstance(val, UOp):
-              old = block_assigns.get(var, env.get(var, _u32(0)))
+              old = block_assigns.get(var, vars.get(var, _u32(0)))
               if isinstance(old, UOp):
-                block_assigns[var] = env[var] = not_found.where(
+                block_assigns[var] = vars[var] = not_found.where(val, old.cast(val.dtype) if val.dtype != old.dtype and val.dtype.itemsize == old.dtype.itemsize else old)
-                  val, old.cast(val.dtype) if val.dtype != old.dtype and val.dtype.itemsize == old.dtype.itemsize else old)
           for j, bl in enumerate(body_lines):
             bl_l = bl.strip().lower()
             if bl_l.startswith('if ') and bl_l.endswith(' then'):
               if any(body_lines[k].strip().lower() == 'break' for k in range(j+1, len(body_lines))):
                 cond_str = _subst_loop_var(bl.strip()[3:-5].strip(), loop_var, loop_i)
-                cond = _to_bool(parse_expr(cond_str, env, funcs))
+                cond = _to_bool(parse_expr(cond_str, vars, funcs))
-                block_assigns[found_var] = env[found_var] = not_found.where(cond, found)
+                block_assigns[found_var] = vars[found_var] = not_found.where(cond, found)
                 break
         else:
-          block_assigns.update(iter_assigns)
+          block_assigns.update(iter_assigns); vars.update(iter_assigns)
-          env.update(iter_assigns)
       continue
 
     # declare
     if first == 'declare':
-      # Initialize scalar declarations (skip arrays and env already passed as srcs)
+      # Initialize scalar declarations (skip arrays and vars already passed as srcs)
       if '[' not in line and len(toks) >= 2 and toks[1].type == 'IDENT':
-        env.setdefault(toks[1].val, _u32(0))
+        vars.setdefault(toks[1].val, _u32(0))
-      i += 1
+      i += 1; continue
-      continue
 
     # lambda definition
     if first != '{' and '=' in line and 'lambda' in line and any(t.type == 'IDENT' and t.val == 'lambda' for t in toks):
@@ -945,30 +883,26 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
               if ch == '(': depth += 1
               elif ch == ')':
                 depth -= 1
-                if depth == 0:
+                if depth == 0: body_lines_lst.append(lines[i][:j]); break
-                  body_lines_lst.append(lines[i][:j])
-                  break
             else: body_lines_lst.append(lines[i])
             i += 1
           body = '\n'.join(body_lines_lst).strip()
-        env[name] = ('lambda', params, body)
+        vars[name] = ('lambda', params, body)
         continue
 
     # MEM assignment: MEM[addr].type (+|-)?= value
     if first == 'mem' and toks[1].type == 'LBRACKET':
       j, addr_toks = _match_bracket(toks, 1)
-      addr = parse_tokens(addr_toks, env, funcs)
+      addr = parse_tokens(addr_toks, vars, funcs)
       if j < len(toks) and toks[j].type == 'DOT': j += 1
       dt_name = toks[j].val if j < len(toks) and toks[j].type == 'IDENT' else 'u32'
       dt, j = DTYPES.get(dt_name, dtypes.uint32), j + 1
       compound_op = None
-      if j < len(toks) and toks[j].type == 'ASSIGN_OP':
+      if j < len(toks) and toks[j].type == 'ASSIGN_OP': compound_op = toks[j].val; j += 1
-        compound_op = toks[j].val
-        j += 1
       elif j < len(toks) and toks[j].type == 'EQUALS': j += 1
-      rhs = parse_tokens(toks[j:], env, funcs)
+      rhs = parse_tokens(toks[j:], vars, funcs)
       if compound_op:
-        mem = env.get('_vmem') if '_vmem' in env else env.get('_lds')
+        mem = vars.get('_vmem') if '_vmem' in vars else vars.get('_lds')
         if isinstance(mem, UOp):
           adt = dtypes.uint64 if addr.dtype == dtypes.uint64 else dtypes.uint32
           idx = (addr >> _const(adt, 2)).cast(dtypes.int)
@@ -977,8 +911,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
             old = old.cast(dtypes.uint64) | (mem.index(((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int)).cast(dtypes.uint64) << _u64(32))
           rhs = (old + rhs) if compound_op == '+=' else (old - rhs)
       if assigns is not None: assigns.append((f'MEM[{_tok_str(addr_toks)}].{dt_name}', (addr, rhs)))
-      i += 1
+      i += 1; continue
-      continue
 
     # VGPR assignment: VGPR[lane][reg] = value
     if first == 'vgpr' and toks[1].type == 'LBRACKET':
@@ -987,12 +920,9 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         j, reg_toks = _match_bracket(toks, j)
         if j < len(toks) and toks[j].type == 'DOT': j += 2  # skip .type suffix
         if j < len(toks) and toks[j].type == 'EQUALS': j += 1
-        ln = parse_tokens(lane_toks, env, funcs)
+        ln, rg, val = parse_tokens(lane_toks, vars, funcs), parse_tokens(reg_toks, vars, funcs), parse_tokens(toks[j:], vars, funcs)
-        rg, val = parse_tokens(reg_toks, env, funcs), parse_tokens(toks[j:], env, funcs)
+        if assigns is not None: assigns.append((f'VGPR[{_tok_str(lane_toks)}][{_tok_str(reg_toks)}]', (_to_u32(rg) * _u32(32) + _to_u32(ln), val)))
-        if assigns is not None:
+        i += 1; continue
-          assigns.append((f'VGPR[{_tok_str(lane_toks)}][{_tok_str(reg_toks)}]', (_to_u32(rg) * _u32(32) + _to_u32(ln), val)))
-        i += 1
-        continue
 
     # Compound destination: {hi.type, lo.type} = value
     if first == '{':
@@ -1006,20 +936,18 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
           j += 3
           if j < len(toks) and toks[j].type == 'RBRACE': j += 1
           if j < len(toks) and toks[j].type == 'EQUALS': j += 1
-          val = parse_tokens(toks[j:], env, funcs)
+          val = parse_tokens(toks[j:], vars, funcs)
           lo_dt, hi_dt = DTYPES.get(lo_type, dtypes.uint64), DTYPES.get(hi_type, dtypes.uint32)
           lo_bits = 64 if lo_dt in (dtypes.uint64, dtypes.int64) else 32
           lo_val = val.cast(lo_dt) if val.dtype.itemsize * 8 <= lo_bits else (val & _const(val.dtype, (1 << lo_bits) - 1)).cast(lo_dt)
           hi_val = (val >> _const(val.dtype, lo_bits)).cast(hi_dt)
-          block_assigns[lo_var] = env[lo_var] = lo_val
+          block_assigns[lo_var] = vars[lo_var] = lo_val
-          block_assigns[hi_var] = env[hi_var] = hi_val
+          block_assigns[hi_var] = vars[hi_var] = hi_val
           if assigns is not None: assigns.extend([(f'{lo_var}.{lo_type}', lo_val), (f'{hi_var}.{hi_type}', hi_val)])
-          i += 1
+          i += 1; continue
-          continue
 
     # Bit slice/index: var[hi:lo] = value, var.type[hi:lo] = value, or var[expr] = value
-    if len(toks) >= 5 and toks[0].type == 'IDENT' and \
+    if len(toks) >= 5 and toks[0].type == 'IDENT' and (toks[1].type == 'LBRACKET' or (toks[1].type == 'DOT' and toks[3].type == 'LBRACKET')):
-        (toks[1].type == 'LBRACKET' or (toks[1].type == 'DOT' and toks[3].type == 'LBRACKET')):
       bracket_start = 2 if toks[1].type == 'LBRACKET' else 4
       j = bracket_start
       colon_pos = None
@@ -1036,28 +964,23 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
           j += 1
           if j < len(toks) and toks[j].type == 'DOT': j += 2
           if j < len(toks) and toks[j].type == 'EQUALS': j += 1
-          val = parse_tokens(toks[j:], env, funcs)
+          val = parse_tokens(toks[j:], vars, funcs)
           dt_suffix = toks[2].val if toks[1].type == 'DOT' else None
           if assigns is not None: assigns.append((f'{var}[{hi}:{lo}]' + (f'.{dt_suffix}' if dt_suffix else ''), val))
-          if var not in env: env[var] = _const(dtypes.uint64 if hi >= 32 else dtypes.uint32, 0)
+          if var not in vars: vars[var] = _const(dtypes.uint64 if hi >= 32 else dtypes.uint32, 0)
-          old = block_assigns.get(var, env.get(var))
+          old = block_assigns.get(var, vars.get(var))
-          assert isinstance(old, UOp)
+          block_assigns[var] = vars[var] = _set_bits(old, _val_to_bits(val), hi - lo + 1, lo)
-          block_assigns[var] = env[var] = _set_bits(old, _val_to_bits(val), hi - lo + 1, lo)
+          i += 1; continue
-          i += 1
+        except: pass
-          continue
-        except Exception: pass
       elif toks[1].type == 'LBRACKET':  # bit index: var[expr] (only for var[...], not var.type[...])
-        existing = block_assigns.get(var, env.get(var))
+        existing = block_assigns.get(var, vars.get(var))
-        if existing is not None and isinstance(existing, UOp) and \
+        if existing is not None and isinstance(existing, UOp) and not any(f'{var}{k}' in vars or f'{var}{k}' in block_assigns for k in range(8)):
-            not any(f'{var}{k}' in env or f'{var}{k}' in block_assigns for k in range(8)):
           bit_toks = toks[2:j]
           j += 1
           while j < len(toks) and toks[j].type != 'EQUALS': j += 1
           if j < len(toks):
-            block_assigns[var] = env[var] = _set_bit(
+            block_assigns[var] = vars[var] = _set_bit(existing, _to_u32(parse_tokens(bit_toks, vars, funcs)), parse_tokens(toks[j+1:], vars, funcs))
-              existing, _to_u32(parse_tokens(bit_toks, env, funcs)), parse_tokens(toks[j+1:], env, funcs))
+            i += 1; continue
-            i += 1
-            continue
 
     # Array element: var[idx] = value (static index) or var[expr] = value (dynamic)
     if len(toks) >= 4 and toks[0].type == 'IDENT' and toks[1].type == 'LBRACKET':
@@ -1067,90 +990,80 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         # Static index: var[NUM] = value
         if len(idx_toks) == 1 and idx_toks[0].type == 'NUM':
           idx = int(idx_toks[0].val.rstrip('UuLl'))
-          val = parse_tokens(toks[j+1:], env, funcs)
+          val = parse_tokens(toks[j+1:], vars, funcs)
-          existing = block_assigns.get(var, env.get(var))
+          existing = block_assigns.get(var, vars.get(var))
           if existing is not None and isinstance(existing, UOp):
-            block_assigns[var] = env[var] = _set_bit(existing, _u32(idx), val)
+            block_assigns[var] = vars[var] = _set_bit(existing, _u32(idx), val)
           else:
-            block_assigns[f'{var}@{idx}'] = env[f'{var}@{idx}'] = val
+            block_assigns[f'{var}@{idx}'] = vars[f'{var}@{idx}'] = val
-          i += 1
+          i += 1; continue
-          continue
         # Dynamic index: var[expr] = value where var has @-elements
-        elems = [(k.split('@')[1], v) for k, v in {**env, **block_assigns}.items() if k.startswith(f'{var}@') and isinstance(v, UOp)]
+        elems = [(k.split('@')[1], v) for k, v in {**vars, **block_assigns}.items() if k.startswith(f'{var}@') and isinstance(v, UOp)]
         if elems:
-          idx_expr = parse_tokens(idx_toks, env, funcs)
+          idx_expr = parse_tokens(idx_toks, vars, funcs)
-          val = parse_tokens(toks[j+1:], env, funcs)
+          val = parse_tokens(toks[j+1:], vars, funcs)
           for elem_idx_str, old_elem in elems:
             elem_idx = int(elem_idx_str)
             cond = _to_u32(idx_expr).eq(_u32(elem_idx))
             new_val = cond.where(val.cast(old_elem.dtype) if val.dtype != old_elem.dtype else val, old_elem)
-            block_assigns[f'{var}@{elem_idx}'] = env[f'{var}@{elem_idx}'] = new_val
+            block_assigns[f'{var}@{elem_idx}'] = vars[f'{var}@{elem_idx}'] = new_val
-          i += 1
+          i += 1; continue
-          continue
 
     # Compound assignment: var += or var -=
     assign_op = next((j for j, t in enumerate(toks) if t.type == 'ASSIGN_OP'), None)
     if assign_op is not None:
       var = toks[0].val
-      old = block_assigns.get(var, env.get(var, _u32(0)))
+      old = block_assigns.get(var, vars.get(var, _u32(0)))
-      rhs = parse_tokens(toks[assign_op+1:], env, funcs)
+      rhs = parse_tokens(toks[assign_op+1:], vars, funcs)
       if rhs.dtype != old.dtype: rhs = rhs.cast(old.dtype)
-      block_assigns[var] = env[var] = (old + rhs) if toks[assign_op].val == '+=' else (old - rhs)
+      block_assigns[var] = vars[var] = (old + rhs) if toks[assign_op].val == '+=' else (old - rhs)
-      i += 1
+      i += 1; continue
-      continue
 
     # Typed element: var.type[idx] = value
-    if len(toks) >= 7 and toks[0].type == 'IDENT' and toks[1].type == 'DOT' and \
+    if len(toks) >= 7 and toks[0].type == 'IDENT' and toks[1].type == 'DOT' and toks[2].type == 'IDENT' and toks[3].type == 'LBRACKET' and toks[4].type == 'NUM':
-        toks[2].type == 'IDENT' and toks[3].type == 'LBRACKET' and toks[4].type == 'NUM':
       var, dt_name, idx = toks[0].val, toks[2].val, int(toks[4].val)
       dt = DTYPES.get(dt_name, dtypes.uint32)
       j = 6
       while j < len(toks) and toks[j].type != 'EQUALS': j += 1
       if j < len(toks):
-        val, old = parse_tokens(toks[j+1:], env, funcs), block_assigns.get(var, env.get(var, _u32(0)))
+        val, old = parse_tokens(toks[j+1:], vars, funcs), block_assigns.get(var, vars.get(var, _u32(0)))
         bw = dt.itemsize * 8
-        block_assigns[var] = env[var] = _set_bits(old, val, bw, idx * bw)
+        block_assigns[var] = vars[var] = _set_bits(old, val, bw, idx * bw)
         if assigns is not None: assigns.append((f'{var}.{dt_name}[{idx}]', val))
-        i += 1
+        i += 1; continue
-        continue
 
     # Dynamic bit: var.type[expr_with_brackets] = value
-    if len(toks) >= 5 and toks[0].type == 'IDENT' and toks[1].type == 'DOT' and \
+    if len(toks) >= 5 and toks[0].type == 'IDENT' and toks[1].type == 'DOT' and toks[2].type == 'IDENT' and toks[3].type == 'LBRACKET':
-        toks[2].type == 'IDENT' and toks[3].type == 'LBRACKET':
       j, depth, has_inner = 4, 1, False
       while j < len(toks) and depth > 0:
-        if toks[j].type == 'LBRACKET':
+        if toks[j].type == 'LBRACKET': depth += 1; has_inner = True
-          depth += 1
-          has_inner = True
         elif toks[j].type == 'RBRACKET': depth -= 1
         j += 1
       if has_inner:
         var = toks[0].val
-        bit_pos = _to_u32(parse_tokens(toks[4:j-1], env, funcs))
+        bit_pos = _to_u32(parse_tokens(toks[4:j-1], vars, funcs))
         while j < len(toks) and toks[j].type != 'EQUALS': j += 1
         if j < len(toks):
-          val = parse_tokens(toks[j+1:], env, funcs)
+          val = parse_tokens(toks[j+1:], vars, funcs)
-          old = block_assigns.get(var, env.get(var, _u32(0)))
+          old = block_assigns.get(var, vars.get(var, _u32(0)))
-          block_assigns[var] = env[var] = _set_bit(old, bit_pos, val)
+          block_assigns[var] = vars[var] = _set_bit(old, bit_pos, val)
-          i += 1
+          i += 1; continue
-          continue
 
     # If/elsif/else - skip branches with statically false conditions (WAVE32/WAVE64)
     if first == 'if':
       def parse_cond(s, kw):
         ll = s.lower()
-        return _to_bool(parse_expr(s[ll.find(kw) + len(kw):ll.rfind('then')].strip(), env, funcs))
+        return _to_bool(parse_expr(s[ll.find(kw) + len(kw):ll.rfind('then')].strip(), vars, funcs))
       def is_const(c, v): return c.op == Ops.CONST and c.arg is v
       cond = parse_cond(line, 'if')
       conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else []
       else_branch: tuple[UOp | None, dict[str, VarVal]] = (None, {})
-      env_snap = dict(env)
+      vars_snap = dict(vars)
       static_true = is_const(cond, True)  # track if any condition is statically true
       i += 1
-      i, branch, ret = parse_block(lines, i, env, funcs, assigns if not is_const(cond, False) else None)
+      i, branch, ret = parse_block(lines, i, vars, funcs, assigns if not is_const(cond, False) else None)
       if conditions: conditions[0] = (cond, ret if ret is not None else branch)
-      env.clear()
+      vars.clear(); vars.update(vars_snap)
-      env.update(env_snap)
       while i < len(lines):
         ltoks = tokenize(lines[i])
         if ltoks[0].type != 'IDENT': break
@@ -1158,22 +1071,17 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         if lf == 'elsif':
           c = parse_cond(lines[i], 'elsif')
           take = not static_true and not is_const(c, False)
-          i += 1
+          i += 1; i, branch, ret = parse_block(lines, i, vars, funcs, assigns if take else None)
-          i, branch, ret = parse_block(lines, i, env, funcs, assigns if take else None)
           if take:
             conditions.append((c, ret if ret is not None else branch))
             if is_const(c, True): static_true = True
-          env.clear()
+          vars.clear(); vars.update(vars_snap)
-          env.update(env_snap)
         elif lf == 'else':
           i += 1
-          i, branch, ret = parse_block(lines, i, env, funcs, assigns if not static_true else None)
+          i, branch, ret = parse_block(lines, i, vars, funcs, assigns if not static_true else None)
           if not static_true: else_branch = (ret, branch)
-          env.clear()
+          vars.clear(); vars.update(vars_snap)
-          env.update(env_snap)
+        elif lf == 'endif': i += 1; break
-        elif lf == 'endif':
-          i += 1
-          break
         else: break
       # Check if any branch returned a value (lambda-style)
       if any(isinstance(br, UOp) for _, br in conditions):
@@ -1186,19 +1094,18 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       # If statically true, use that branch directly; otherwise merge with WHERE
       if static_true:
         ba = next((b for c, b in conditions if is_const(c, True) and isinstance(b, dict)), {})
-        block_assigns.update(ba)
+        block_assigns.update(ba); vars.update(ba)
-        env.update(ba)
       else:
         else_assigns = else_branch[1]
         all_vars = set().union(*[ba.keys() for _, ba in conditions if isinstance(ba, dict)], else_assigns.keys())
         for var in all_vars:
-          res: Any = else_assigns.get(var, block_assigns.get(var, env.get(var, _u32(0))))
+          res: Any = else_assigns.get(var, block_assigns.get(var, vars.get(var, _u32(0))))
-          for cond, ba in reversed(conditions):  # type: ignore[assignment]
+          for cond, ba in reversed(conditions):
             if isinstance(ba, dict) and var in ba:
               tv = ba[var]
               if isinstance(tv, UOp) and isinstance(res, UOp):
                 res = cond.where(tv, res.cast(tv.dtype) if tv.dtype != res.dtype and tv.dtype.itemsize == res.dtype.itemsize else res)
-          block_assigns[var] = env[var] = res
+          block_assigns[var] = vars[var] = res
       continue
 
     # Regular assignment: var = value
@@ -1206,12 +1113,11 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       if t.type == 'EQUALS':
         if any(toks[k].type == 'OP' and toks[k].val in ('<', '>', '!', '=') for k in range(j)): break
         base_var = toks[0].val
-        block_assigns[base_var] = env[base_var] = parse_tokens(toks[j+1:], env, funcs)
+        block_assigns[base_var] = vars[base_var] = parse_tokens(toks[j+1:], vars, funcs)
-        i += 1
+        i += 1; break
-        break
     else: i += 1
   return i, block_assigns, None
 
-def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
+def parse_expr(expr: str, vars: dict[str, VarVal], funcs: dict | None = None) -> UOp:
-  return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
+  return parse_tokens(tokenize(expr.strip().rstrip(';')), vars, funcs)
 

extra/assembly/amd/sqtt.py

@@ -5,11 +5,9 @@ The format is nibble-based with variable-width packets determined by a state mac
 Uses BitField infrastructure from dsl.py, similar to GPU instruction encoding.
 """
 from __future__ import annotations
-from dataclasses import dataclass
 from typing import Iterator
 from enum import Enum
-from tinygrad.renderer.amd.dsl import BitField, FixedBitField, Inst, bits
+from extra.assembly.amd.dsl import BitField, FixedBitField, bits
-from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm # same encoding as RDNA4
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # FIELD ENUMS
@@ -46,7 +44,6 @@ class InstOp(Enum):
   SMEM = 0x1
   JUMP = 0x3              # branch taken
   JUMP_NO = 0x4           # branch not taken
-  CALL = 0x5              # s_call_b64
   MESSAGE = 0x9
   VALU_TRANS = 0xb        # transcendental: exp, log, rcp, sqrt, sin, cos
   VALU_64_SHIFT = 0xd     # 64-bit shifts: lshl, lshr, ashr
@@ -73,10 +70,8 @@ class InstOp(Enum):
 
   # LDS ops on traced SIMD
   LDS_LOAD = 0x29
-  LDS_ATOMIC = 0x2a        # ds_append, ds_consume, ds_store_addtid_b32
   LDS_STORE = 0x2b
   LDS_STORE_64 = 0x2c
-  LDS_STORE_96 = 0x2d
   LDS_STORE_128 = 0x2e
 
   # Memory ops on other SIMD (0x5x range)
@@ -100,44 +95,20 @@ class InstOp(Enum):
   SALU_SAVEEXEC = 0x72    # s_*_saveexec_b32/b64
   VALU_CMPX = 0x73        # v_cmpx_*
 
-class InstOpRDNA4(Enum):
+class InstOpL4(Enum):
   """SQTT instruction operation types for RDNA4 (gfx1200). Different encoding from RDNA3."""
+  # TODO: we need to do discovery of all of these from instructions
   SALU = 0x0
   SMEM = 0x1
-  JUMP = 0x3
+  UNK_02 = 0x2
   JUMP_NO = 0x4
-  CALL = 0x5
+  UNK_06 = 0x6
-  MESSAGE = 0x9
+  VMEM = 0x10
-  VALU_TRANS = 0xb
+  UNK_11 = 0x11
-  VALU_B2 = 0xd
-  VALU_B4 = 0xe
   VINTERP = 0x12
-  VMEM_RD_1 = 0x21
+  UNK_14 = 0x14
-  VMEM_RD_2 = 0x22
+  OTHER_VMEM = 0x5e
-  VMEM_WR_1 = 0x23
+  UNK_60 = 0x60
-  VMEM_WR_2 = 0x24
-  VMEM_WR_3 = 0x25
-  VMEM_WR_4 = 0x26
-  VMEM_WR_5 = 0x27
-  VMEM_WR_6 = 0x28
-  LDS_RD = 0x29
-  LDS_WR_1 = 0x2a
-  LDS_WR_2 = 0x2b
-  LDS_WR_3 = 0x2c
-  LDS_WR_4 = 0x2d
-  LDS_WR_5 = 0x2e
-  OTHER_LDS_1 = 0x50
-  OTHER_LDS_2 = 0x51
-  BARRIER_SIGNAL = 0x7a
-  WMMA_8 = 0x8c
-  WMMA_16 = 0x8d
-  VALU_DPFP = 0x92
-  SALU_FLOAT3 = 0x98
-  VALU_SCL_TRANS = 0x99
-  SALU_2 = 0x9b
-  SALU_5 = 0x9c
-  OTHER_VMEM = 0xbd
-  OTHER_VMEM_5 = 0xc1
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # PACKET TYPE BASE CLASS
@@ -151,8 +122,8 @@ class PacketType:
 
   def __init_subclass__(cls, **kwargs):
     super().__init_subclass__(**kwargs)
-    cls._fields = {k: v for k, v in cls.__dict__.items() if isinstance(v, BitField)}  # type: ignore[attr-defined]
+    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)  # type: ignore[attr-defined]
+    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):
@@ -161,7 +132,7 @@ class PacketType:
     return inst
 
   def __repr__(self) -> str:
-    fields_str = ", ".join(f"{k}={getattr(self, k)}" for k in self._fields if not k.startswith('_') and k != 'encoding')  # type: ignore[attr-defined]
+    fields_str = ", ".join(f"{k}={getattr(self, k)}" for k in self._fields if not k.startswith('_') and k != 'encoding')
     return f"{self.__class__.__name__}({fields_str})"
 
 # ═══════════════════════════════════════════════════════════════════════════════
@@ -173,12 +144,17 @@ class TS_DELTA_S8_W3(PacketType):
   delta = bits[10:8]
   _padding = bits[63:11]
 
+class TS_DELTA_S8_W3_L4(PacketType):  # Layout 4: 64->72 bits
+  encoding = bits[6:0] == 0b0100001
+  delta = bits[10:8]
+  _padding = bits[71:11]
+
 class TS_DELTA_S5_W3(PacketType):
   encoding = bits[4:0] == 0b00110
   delta = bits[7:5]
   _padding = bits[51:8]
 
-class TS_DELTA_S5_W3_RDNA4(PacketType):  # Layout 4: 52->56 bits
+class TS_DELTA_S5_W3_L4(PacketType):  # Layout 4: 52->56 bits
   encoding = bits[4:0] == 0b00110
   delta = bits[9:7]
   _padding = bits[55:10]
@@ -190,23 +166,26 @@ class TS_DELTA_SHORT(PacketType):
 class TS_DELTA_OR_MARK(PacketType):
   encoding = bits[6:0] == 0b0000001
   delta = bits[47:12]
-  pl = bits[8:8]
+  bit8 = bits[8:8]
-  rt = bits[9:9]
+  bit9 = bits[9:9]
   @property
-  def is_marker(self) -> bool: return bool(self.rt and not self.pl)
+  def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
 
-class TS_DELTA_OR_MARK_RDNA4(TS_DELTA_OR_MARK):
+class TS_DELTA_OR_MARK_L4(PacketType):  # Layout 4: 48->64 bits
+  encoding = bits[6:0] == 0b0000001
   delta = bits[63:12]
-  rt = bits[7:7]
+  bit7 = bits[7:7]
-  pl = bits[8:8]
+  bit8 = bits[8:8]
-  tl = bits[9:9]
+  bit9 = bits[9:9]
+  @property
+  def is_marker(self) -> bool: return bool((self.bit9 and not self.bit8) or self.bit7)
 
 class TS_DELTA_S5_W2(PacketType):
   encoding = bits[4:0] == 0b11100
   delta = bits[6:5]
   _padding = bits[47:7]
 
-class TS_DELTA_S5_W2_RDNA4(PacketType):  # Layout 4: 48->40 bits
+class TS_DELTA_S5_W2_L4(PacketType):  # Layout 4: 48->40 bits
   encoding = bits[4:0] == 0b11100
   delta = bits[6:5]
   _padding = bits[39:7]
@@ -267,7 +246,7 @@ class WAVESTART(PacketType):  # exclude: 1 << 4
   @property
   def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
 
-class WAVESTART_RDNA4(PacketType):  # Layout 4 has wave field at different position
+class WAVESTART_L4(PacketType):  # Layout 4 has wave field at different position
   encoding = bits[4:0] == 0b01100
   delta = bits[6:5]
   flag7 = bits[7:7]
@@ -283,7 +262,7 @@ class WAVEALLOC(PacketType):  # exclude: 1 << 10
   delta = bits[7:5]
   _padding = bits[19:8]
 
-class WAVEALLOC_RDNA4(PacketType):  # Layout 4: 20->24 bits
+class WAVEALLOC_L4(PacketType):  # Layout 4: 20->24 bits
   encoding = bits[4:0] == 0b00101
   delta = bits[7:5]
   _padding = bits[23:8]
@@ -293,7 +272,7 @@ class PERF(PacketType):  # exclude: 1 << 11
   delta = bits[7:5]
   arg = bits[27:8]
 
-class PERF_RDNA4(PacketType):  # Layout 4: 28->32 bits
+class PERF_L4(PacketType):  # Layout 4: 28->32 bits
   encoding = bits[4:0] == 0b10110
   delta = bits[9:7]
   arg = bits[31:10]
@@ -356,12 +335,13 @@ class INST(PacketType):
   wave = bits[12:8]
   op = bits[19:13].enum(InstOp)
 
-class INST_RDNA4(PacketType):  # Layout 4: different delta position and InstOp encoding
+class INST_L4(PacketType):  # Layout 4: different delta position and InstOp encoding
   encoding = bits[2:0] == 0b010
   delta = bits[5:3]
-  w64h = bits[6:6]
+  flag1 = bits[6:6]
-  wave = bits[11:7]
+  flag2 = bits[7:7]
-  op = bits[19:12].enum(InstOpRDNA4)
+  wave = bits[12:8]
+  op = bits[19:13].enum(InstOpL4)
 
 class UTILCTR(PacketType):
   encoding = bits[6:0] == 0b0110001
@@ -369,188 +349,40 @@ class UTILCTR(PacketType):
   ctr = bits[47:9]
 
 # Packet types with rocprof type IDs as keys
-PACKET_TYPES_RDNA3: dict[int, type[PacketType]] = {
+PACKET_TYPES_L3: dict[int, type[PacketType]] = {
   1: VALUINST, 2: VMEMEXEC, 3: ALUEXEC, 4: IMMEDIATE, 5: IMMEDIATE_MASK, 6: WAVERDY, 7: TS_DELTA_S8_W3, 8: WAVEEND,
   9: WAVESTART, 10: TS_DELTA_S5_W2, 11: WAVEALLOC, 12: TS_DELTA_S5_W3, 13: PERF, 14: UTILCTR, 15: TS_DELTA_SHORT,
   16: NOP, 17: TS_WAVE_STATE, 18: EVENT, 19: EVENT_BIG, 20: REG, 21: SNAPSHOT, 22: TS_DELTA_OR_MARK, 23: LAYOUT_HEADER, 24: INST,
 }
-PACKET_TYPES_RDNA4: dict[int, type[PacketType]] = {
+PACKET_TYPES_L4: dict[int, type[PacketType]] = {
-  **PACKET_TYPES_RDNA3,
+  **PACKET_TYPES_L3,
-  9: WAVESTART_RDNA4, 10: TS_DELTA_S5_W2_RDNA4, 11: WAVEALLOC_RDNA4,
+  7: TS_DELTA_S8_W3_L4, 9: WAVESTART_L4, 10: TS_DELTA_S5_W2_L4, 11: WAVEALLOC_L4,
-  12: TS_DELTA_S5_W3_RDNA4, 13: PERF_RDNA4, 22: TS_DELTA_OR_MARK_RDNA4, 24: INST_RDNA4,
+  12: TS_DELTA_S5_W3_L4, 13: PERF_L4, 22: TS_DELTA_OR_MARK_L4, 24: INST_L4,
 }
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# CDNA PACKET TYPE DEFINITIONS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-class CDNA_MISC(PacketType):
-  """pkt_fmt=0: 16-bit (Misc)"""
-  encoding = bits[3:0] == 0
-  delta = bits[11:4]
-  sh = bits[12:12]
-  misc_type = bits[15:13]
-
-class CDNA_TIMESTAMP(PacketType):
-  """pkt_fmt=1: 64-bit timestamp packet (case 0x0)"""
-  encoding = bits[3:0] == 1
-  _reserved = bits[15:4]
-  timestamp = bits[63:16]   # stored as (data_word >> 0x10) in low 46 bits of local_58
-
-class CDNA_REG(PacketType):
-  """pkt_fmt=2: 64-bit (Reg)"""
-  encoding = bits[3:0] == 2
-  pipe = bits[6:5]
-  _me_raw = bits[8:7]
-  _reserved = bits[15:9]
-  regaddr = bits[31:16]
-  regdata = bits[63:32]
-
-class CDNA_WAVESTART(PacketType):
-  """type 3: 32-bit wave start (Wave/group_id)"""
-  encoding = bits[3:0] == 3
-  sh = bits[5:5]
-  cu = bits[9:6]
-  wave = bits[13:10]
-  simd = bits[15:14]
-  pipe = bits[17:16]
-  me = bits[19:18]
-  _reserved = bits[21:20]
-  count = bits[28:22]
-  _padding = bits[31:29]
-
-class CDNA_WAVEALLOC(PacketType):
-  """pkt_fmt=4: 16-bit (Wave)"""
-  encoding = bits[3:0] == 4
-  sh = bits[5:5]
-  cu = bits[9:6]
-  wave = bits[13:10]
-  simd = bits[15:14]
-
-class CDNA_REG_CS(PacketType):
-  """type 5: 48-bit register CS write (RegCs)"""
-  encoding = bits[3:0] == 5
-  pipe = bits[6:5]
-  _me_raw = bits[8:7]
-  regaddr = bits[15:9]
-  regdata = bits[47:16]
-
-class CDNA_WAVEEND(PacketType):
-  """type 6: 16-bit wave end (group_id)"""
-  encoding = bits[3:0] == 6
-  sh = bits[5:5]
-  cu = bits[9:6]
-  wave = bits[13:10]
-  simd = bits[15:14]
-
-class CDNA_INST(PacketType):
-  """pkt_fmt=10: 16-bit (MsgInst)"""
-  encoding = bits[3:0] == 10
-  wave = bits[8:5]
-  simd = bits[10:9]
-  inst_type = bits[15:11]
-
-class CDNA_INST_PC(PacketType):
-  """pkt_fmt=11: 64-bit (MsgInstPc)"""
-  encoding = bits[3:0] == 11
-  wave = bits[8:5]
-  simd = bits[10:9]
-  _reserved = bits[14:11]
-  err = bits[15:15]
-  pc = bits[63:16]
-
-class CDNA_ISSUE(PacketType):
-  """pkt_fmt=13: 32-bit (Issue)"""
-  encoding = bits[3:0] == 13
-  simd = bits[6:5]
-  _gap = bits[7:7]
-  inst0 = bits[9:8]
-  inst1 = bits[11:10]
-  inst2 = bits[13:12]
-  inst3 = bits[15:14]
-  inst4 = bits[17:16]
-  inst5 = bits[19:18]
-  inst6 = bits[21:20]
-  inst7 = bits[23:22]
-  inst8 = bits[25:24]
-  inst9 = bits[27:26]
-  _padding = bits[31:28]
-
-class CDNA_PERF(PacketType):
-  """pkt_fmt=14: 64-bit (MsgPerf)"""
-  encoding = bits[3:0] == 14
-  sh = bits[5:5]
-  cu = bits[9:6]
-  cntr_bank = bits[11:10]
-  cntr0 = bits[24:12]
-  cntr1 = bits[37:25]
-  cntr2 = bits[50:38]
-  cntr3 = bits[63:51]
-
-class CDNA_EVENT(PacketType):
-  """pkt_fmt=7: 16-bit"""
-  encoding = bits[3:0] == 7
-  _reserved = bits[15:4]
-
-class CDNA_EVENT_CS(PacketType):
-  """pkt_fmt=8: 16-bit"""
-  encoding = bits[3:0] == 8
-  _reserved = bits[15:4]
-
-class CDNA_EVENT_GFX1(PacketType):
-  """pkt_fmt=9: 16-bit"""
-  encoding = bits[3:0] == 9
-  _reserved = bits[15:4]
-
-class CDNA_USERDATA(PacketType):
-  """pkt_fmt=12: 48-bit (UserData)"""
-  encoding = bits[3:0] == 12
-  sh = bits[5:5]
-  cu = bits[9:6]
-  wave = bits[13:10]
-  simd = bits[15:14]
-  data = bits[47:16]
-
-class CDNA_REG_CS_PRIV(PacketType):
-  """pkt_fmt=15: 48-bit (RegCs)"""
-  encoding = bits[3:0] == 15
-  pipe = bits[6:5]
-  _me_raw = bits[8:7]
-  regaddr = bits[15:9]
-  regdata = bits[47:16]
-
-PACKET_TYPES_CDNA: dict[int, type[PacketType]] = {
-  0: CDNA_MISC, 1: CDNA_TIMESTAMP, 2: CDNA_REG, 3: CDNA_WAVESTART, 4: CDNA_WAVEALLOC, 5: CDNA_REG_CS, 6: CDNA_WAVEEND,
-  7: CDNA_EVENT, 8: CDNA_EVENT_CS, 9: CDNA_EVENT_GFX1, 10: CDNA_INST, 11: CDNA_INST_PC, 12: CDNA_USERDATA,
-  13: CDNA_ISSUE, 14: CDNA_PERF, 15: CDNA_REG_CS_PRIV, 16: LAYOUT_HEADER,
-}
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# DECODER
-# ═══════════════════════════════════════════════════════════════════════════════
-
 def _build_decode_tables(packet_types: dict[int, type[PacketType]]) -> tuple[dict[int, tuple], bytes]:
   # Build state table: byte -> opcode. Sort by mask specificity (more bits first), NOP last
   sorted_types = sorted(packet_types.items(), key=lambda x: (-bin(x[1].encoding.mask).count('1'), x[0] == 16))
   state_table = bytes(next((op for op, cls in sorted_types if (b & cls.encoding.mask) == cls.encoding.default), 16) for b in range(256))
   # Build decode info: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
-  # special_case: 0=none, 1=TS_DELTA_OR_MARK (check is_marker), 2=TS_DELTA_SHORT (add 8), 3=CDNA_MISC (*4), 4=CDNA_TIMESTAMP (absolute)
+  # special_case: 0=none, 1=TS_DELTA_OR_MARK (check is_marker), 2=TS_DELTA_SHORT (add 8)
-  _special = {TS_DELTA_OR_MARK: 1, TS_DELTA_OR_MARK_RDNA4: 1, TS_DELTA_SHORT: 2, CDNA_MISC: 3, CDNA_TIMESTAMP: 4}
   decode_info = {}
   for opcode, pkt_cls in packet_types.items():
     delta_field = getattr(pkt_cls, 'delta', None)
-    special = _special.get(pkt_cls, 0)
+    special = {22: 1, 15: 2}.get(opcode, 0)  # TS_DELTA_OR_MARK=22, TS_DELTA_SHORT=15
-    decode_info[opcode] = (pkt_cls, pkt_cls._size_nibbles, delta_field.lo if delta_field else 0, delta_field.mask if delta_field else 0, special)  # type: ignore[attr-defined]
+    decode_info[opcode] = (pkt_cls, pkt_cls._size_nibbles, delta_field.lo if delta_field else 0, delta_field.mask if delta_field else 0, special)
   return decode_info, state_table
 
-_DECODE_INFO_RDNA3, _STATE_TABLE_RDNA3 = _build_decode_tables(PACKET_TYPES_RDNA3)
+_DECODE_INFO_L3, _STATE_TABLE_L3 = _build_decode_tables(PACKET_TYPES_L3)
-_DECODE_INFO_RDNA4, _STATE_TABLE_RDNA4 = _build_decode_tables(PACKET_TYPES_RDNA4)
+_DECODE_INFO_L4, _STATE_TABLE_L4 = _build_decode_tables(PACKET_TYPES_L4)
-_DECODE_INFO_CDNA, _STATE_TABLE_CDNA = _build_decode_tables(PACKET_TYPES_CDNA)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# DECODER
+# ═══════════════════════════════════════════════════════════════════════════════
 
 def decode(data: bytes) -> Iterator[PacketType]:
-  """Decode raw SQTT blob, yielding packet instances. Auto-detects RDNA (layout 3/4) vs CDNA."""
+  """Decode raw SQTT blob, yielding packet instances. Auto-detects layout from LAYOUT_HEADER."""
-  n, reg, pos, nib_off, nib_count, time, ts_offset = len(data), 0, 0, 0, 16, 0, None
+  n, reg, pos, nib_off, nib_count, time = len(data), 0, 0, 0, 16, 0
-  decode_info, state_table = _DECODE_INFO_RDNA3, _STATE_TABLE_RDNA3  # start RDNA3, auto-detect switches if needed
+  decode_info, state_table = _DECODE_INFO_L3, _STATE_TABLE_L3  # default to layout 3, will update after seeing LAYOUT_HEADER
 
   while pos + ((nib_count + nib_off + 1) >> 1) <= n:
     need = nib_count - nib_off
@@ -558,9 +390,8 @@ def decode(data: bytes) -> Iterator[PacketType]:
     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):
-      read_bytes = min(byte_count, 8)
+      chunk = int.from_bytes(data[pos:pos + byte_count], 'little')
-      chunk = int.from_bytes(data[pos:pos + read_bytes], 'little')
+      reg, pos = (reg >> (byte_count * 8)) | (chunk << (64 - byte_count * 8)), pos + byte_count
-      reg, pos = (reg >> (read_bytes * 8)) | (chunk << (64 - read_bytes * 8)), pos + byte_count
     # 3. if odd, read low nibble
     if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)
 
@@ -571,83 +402,14 @@ def decode(data: bytes) -> Iterator[PacketType]:
       pkt = pkt_cls.from_raw(reg, 0)  # create packet to check is_marker
       if pkt.is_marker: delta = 0
     elif special == 2: delta += 8  # TS_DELTA_SHORT
-    elif special == 3: delta *= 4  # CDNA_DELTA
-    elif special == 4:  # CDNA_TIMESTAMP (absolute timestamp anchoring)
-      if (reg >> 4) & 0xfff == 0:  # unk_0 == 0 means absolute timestamp
-        abs_ts = reg >> 16
-        if ts_offset is None: ts_offset = abs_ts - time
-        else: time = ((abs_ts - ts_offset) & ~3) - 4
-      delta = 0
     time += delta
     pkt = pkt_cls.from_raw(reg, time)
-    # auto-detect: first packet is always LAYOUT_HEADER (RDNA layout 3/4) or misdetected (CDNA)
+    # detect layout from first LAYOUT_HEADER and switch decode tables if needed
-    if pkt_cls is LAYOUT_HEADER:
+    # NOTE: CDNA uses a completely different 16-bit header format, not nibbles - not supported here
-      if pkt.layout == 4: decode_info, state_table = _DECODE_INFO_RDNA4, _STATE_TABLE_RDNA4
+    if pkt_cls is LAYOUT_HEADER and pkt.layout == 4:
-      elif pkt.layout != 3:  # not a real LAYOUT_HEADER — switch to CDNA and re-decode first packet
+      decode_info, state_table = _DECODE_INFO_L4, _STATE_TABLE_L4
-        decode_info, state_table = _DECODE_INFO_CDNA, _STATE_TABLE_CDNA
-        opcode = state_table[reg & 0xFF]
-        pkt_cls, nib_count, delta_lo, delta_mask, special = decode_info[opcode]
-        if special == 4 and (reg >> 4) & 0xfff == 0:  # CDNA_TIMESTAMP absolute
-          ts_offset = (reg >> 16) - time
-        pkt = pkt_cls.from_raw(reg, time)
     yield pkt
 
-# ═══════════════════════════════════════════════════════════════════════════════
-# MAPPER
-# ═══════════════════════════════════════════════════════════════════════════════
-
-@dataclass(frozen=True)
-class InstructionInfo:
-  pc: int
-  wave: int
-  inst: Inst
-
-def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, InstructionInfo|None]]:
-  """maps SQTT packets to instructions, yields (packet, instruction_info or None)"""
-  # map pcs to insts
-  from tinygrad.viz.serve import amd_decode
-  pc_map = amd_decode(lib, target)
-  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, WAVESTART_RDNA4)):
-      assert p.wave not in wave_pc, "only one inflight wave per unit"
-      wave_pc[p.wave] = next(iter(pc_map))
-    elif isinstance(p, WAVEEND):
-      pc = wave_pc.pop(p.wave)
-      yield (p, InstructionInfo(pc, p.wave, s_endpgm()))
-    # skip OTHER_ instructions, they don't belong to this unit
-    elif isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("OTHER_"): pass
-    elif 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 type(inst).__name__ == "SOPP", f"IMMEDIATE_MASK packet must map to SOPP, got {inst}"
-          wave_pc[wave] += inst.size()
-          yield (p, InstructionInfo(pc, wave, inst))
-    elif isinstance(p, (VALUINST, INST, INST_RDNA4, IMMEDIATE)):
-      inst = pc_map[pc:=wave_pc[p.wave]]
-      # s_delay_alu and s_wait_alu instructions are skipped
-      while (inst_op:=getattr(inst, 'op_name', '')) in {"S_DELAY_ALU", "S_WAIT_ALU"}:
-        wave_pc[p.wave] += inst.size()
-        inst = pc_map[pc:=wave_pc[p.wave]]
-      # assert branch always has a JUMP packet
-      if "BRANCH" in inst_op and not (isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("JUMP")):
-        raise AssertionError(f"{inst_op} can only be followed by JUMP, got {p}")
-      # JUMP handling
-      if isinstance(p, (INST, INST_RDNA4)) and p.op in {InstOp.JUMP, InstOpRDNA4.JUMP}:
-        x = getattr(inst, 'simm16') & 0xffff
-        wave_pc[p.wave] += inst.size() + (x - 0x10000 if x & 0x8000 else x)*4
-      else:
-        wave_pc[p.wave] += inst.size()
-      yield (p, InstructionInfo(pc, p.wave, inst))
-    # for all other packets (VMEMEXEC, ALUEXEC, etc.), yield with None
-    else: yield (p, None)
-
 # ═══════════════════════════════════════════════════════════════════════════════
 # PRINTER
 # ═══════════════════════════════════════════════════════════════════════════════
@@ -662,37 +424,35 @@ PACKET_COLORS = {
 def format_packet(p) -> str:
   from tinygrad.helpers import colored
   name = type(p).__name__
-  if isinstance(p, (INST, INST_RDNA4)):
+  if isinstance(p, (INST, INST_L4)):
-    op_name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4)) else f"0x{p.op:02x}"
+    op_name = p.op.name if isinstance(p.op, (InstOp, InstOpL4)) else f"0x{p.op:02x}"
-    fields = f"wave={p.wave} op={op_name}" + ((" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "") if isinstance(p, INST) else "")
+    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, WAVESTART_RDNA4, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
+  elif isinstance(p, (WAVESTART, WAVESTART_L4, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
   elif hasattr(p, '_fields'):
-    filt = {'delta', 'encoding'} if not isinstance(p, (TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4)) else {'encoding'}
+    filt = {'delta', 'encoding'} if not isinstance(p, (TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_L4)) else {'encoding'}
     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 filt)
   else: fields = ""
-  return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name.replace('_RDNA4', ''), 'white'))} {fields}"
+  return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name.replace('_L4', ''), 'white'))} {fields}"
 
 def print_packets(packets) -> None:
   from tinygrad.helpers import getenv
   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"} if not getenv("NOSKIP") else {"NOP"}
-  for data in packets:
+  for p in packets:
-    p, inst = data if isinstance(data, tuple) else (data, None)
+    if type(p).__name__.replace("_L4", "") not in skip: print(format_packet(p))
-    if type(p).__name__.replace("_RDNA4", "") not in skip: print(format_packet(p), f"inst={inst.inst}" if inst is not None else '')
 
 if __name__ == "__main__":
   import sys, pickle
-  from tinygrad.helpers import temp
+  if len(sys.argv) < 2:
-  with open(temp("profile.pkl", append_user=True) if len(sys.argv) < 2 else sys.argv[1], "rb") as f:
+    print("Usage: python sqtt.py <pkl_file>")
+    sys.exit(1)
+  with open(sys.argv[1], "rb") as f:
     data = pickle.load(f)
-  prg_events = {e.tag: e for e in data if type(e).__name__ == "ProfileProgramEvent" and e.tag is not None}
   sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
-  dev_targets = {e.device:f"gfx{e.props['gfx_target_version']//1000}" for e in data if type(e).__name__ == "ProfileDeviceEvent" and e.props}
   for i, event in enumerate(sqtt_events):
-    prg = prg_events.get(event.kern)
+    print(f"\n=== event {i} ===")
-    print(f"\n=== event {i} {prg.name if prg is not None else ''} ===")
+    print_packets(decode(event.blob))
-    print_packets(map_insts(event.blob, prg.lib, dev_targets[prg.device]) if prg is not None else decode(event.blob))

extra/assembly/amd/sqtt_cdna.py

@@ -0,0 +1,161 @@
+"""SQTT (SQ Thread Trace) packet decoder for CDNA/MI300 GPUs.
+
+CDNA uses a completely different 16-bit header format from RDNA's nibble-based encoding.
+"""
+from __future__ import annotations
+from typing import Iterator
+from extra.assembly.amd.dsl import bits
+from extra.assembly.amd.sqtt import PacketType
+
+# CDNA pkt_fmt -> size in bytes (extracted from rocprof hash table)
+CDNA_PKT_SIZES = {0: 2, 1: 8, 2: 8, 3: 4, 4: 2, 5: 6, 6: 2, 7: 2, 8: 2, 9: 2, 10: 2, 11: 8, 12: 6, 13: 4, 14: 8, 15: 6}
+
+class CDNA_DELTA(PacketType):
+  """pkt_fmt=0: 16-bit timestamp delta packet"""
+  encoding = bits[3:0] == 0
+  delta = bits[11:4]      # (data >> 4) & 0xff
+  unk_0 = bits[12:12]     # (data >> 0xc) & 1
+  unk_1 = bits[15:13]     # (data >> 0xd)
+
+class CDNA_TIMESTAMP(PacketType):
+  """pkt_fmt=1: 64-bit timestamp packet (case 0x0)"""
+  encoding = bits[3:0] == 1
+  unk_0 = bits[15:4]
+  timestamp = bits[63:16]   # stored as (data_word >> 0x10) in low 46 bits of local_58
+
+class CDNA_PKT_2(PacketType):
+  """pkt_fmt=2: 64-bit packet (case 0x4)"""
+  encoding = bits[3:0] == 2
+  unk_0 = bits[6:5]       # (data >> 5) & 3
+  unk_1 = bits[7:7]       # (data >> 7) + 1 & 1
+  unk_padding = bits[63:8]
+
+class CDNA_WAVESTART(PacketType):
+  """pkt_fmt=3: 32-bit WAVESTART packet (case 0x8)"""
+  encoding = bits[3:0] == 3
+  unk_0 = bits[5:5]       # (data >> 5) & 1
+  unk_1 = bits[9:6]       # (data >> 6) & 0xf
+  wave = bits[13:10]      # (data >> 10) & 0xf
+  simd = bits[15:14]      # (data >> 0xe) & 3
+  cu = bits[17:16]        # (data >> 0x10) & 3
+  unk_5 = bits[19:18]     # (data >> 0x12) & 3
+  unk_6 = bits[28:22]     # (data >> 0x16) & 0x7f
+  unk_padding = bits[31:29]
+
+class CDNA_PKT_4(PacketType):
+  """pkt_fmt=4: 16-bit packet (case 0xc, same as 0x8/0x14)"""
+  encoding = bits[3:0] == 4
+  unk_0 = bits[5:5]       # (data_word >> 5) & 1
+  unk_1 = bits[9:6]       # (data_word >> 6) & 0xf
+  unk_2 = bits[13:10]     # (data_word >> 10) & 0xf
+  unk_3 = bits[15:14]     # (data_word >> 0xe)
+
+class CDNA_PKT_5(PacketType):
+  """pkt_fmt=5: 48-bit packet (case 0x10)"""
+  encoding = bits[3:0] == 5
+  unk_0 = bits[6:5]       # (data >> 5) & 3
+  unk_1 = bits[7:7]       # (data >> 7) + 1 & 1
+  unk_2 = bits[15:9]      # (data >> 9) & 0x7f
+  unk_padding = bits[47:16]
+
+class CDNA_WAVEEND(PacketType):
+  """pkt_fmt=6: 16-bit WAVEEND packet (case 0x14, same as 0x8/0xc)"""
+  encoding = bits[3:0] == 6
+  unk_0 = bits[5:5]       # (data_word >> 5) & 1
+  unk_1 = bits[9:6]       # (data_word >> 6) & 0xf
+  wave = bits[13:10]      # (data_word >> 10) & 0xf
+  simd = bits[15:14]      # (data_word >> 0xe)
+
+class CDNA_EXEC(PacketType):
+  """pkt_fmt=10: 16-bit EXEC packet (case 0x24)"""
+  encoding = bits[3:0] == 10
+  unk_0 = bits[8:5]       # (data_word >> 5) & 0xf
+  unk_1 = bits[10:9]      # (data_word >> 9) & 3
+  unk_2 = bits[15:11]     # (data_word >> 0xb)
+
+class CDNA_PKT_11(PacketType):
+  """pkt_fmt=11: 64-bit packet (case 0x28)"""
+  encoding = bits[3:0] == 11
+  unk_0 = bits[8:5]       # (data_word >> 5) & 0xf
+  unk_1 = bits[10:9]      # (data_word >> 9) & 3
+  unk_2 = bits[15:15]     # (data_word >> 0xf) & 1
+  unk_padding = bits[63:16]
+
+class CDNA_INST(PacketType):
+  """pkt_fmt=13: 32-bit INST packet (case 0x30)"""
+  encoding = bits[3:0] == 13
+  unk_0 = bits[6:5]       # (data >> 5) & 3
+  unk_1 = bits[9:8]       # (data >> 8) & 3
+  unk_2 = bits[11:10]     # (data >> 10) & 3
+  unk_3 = bits[13:12]     # (data >> 0xc) & 3
+  unk_4 = bits[15:14]     # (data >> 0xe) & 3
+  unk_5 = bits[19:18]     # (data >> 0x12) & 3
+  unk_6 = bits[21:20]     # (data >> 0x14) & 3
+  unk_7 = bits[23:22]     # (data >> 0x16) & 3
+  unk_8 = bits[25:24]     # (data >> 0x18) & 3
+  unk_9 = bits[27:26]     # (data >> 0x1a) & 3
+  unk_padding = bits[31:28]
+
+class CDNA_PKT_14(PacketType):
+  """pkt_fmt=14: 64-bit packet (case 0x34)"""
+  encoding = bits[3:0] == 14
+  unk_0 = bits[5:5]       # (data >> 5) & 1
+  unk_1 = bits[9:6]       # (data >> 6) & 0xf
+  unk_2 = bits[11:10]     # (data >> 10) & 3
+  unk_3 = bits[24:12]     # (data >> 0xc) & 0x1fff
+  unk_4 = bits[37:25]     # (data >> 0x19) & 0x1fff
+  unk_5 = bits[50:38]     # (data >> 0x26) & 0x1fff
+  unk_6 = bits[51:51]     # (data >> 0x33) & 1
+  unk_padding = bits[63:52]
+
+class CDNA_PKT_15(PacketType):
+  """pkt_fmt=15: 48-bit packet (case 0x38, same as 0x10)"""
+  encoding = bits[3:0] == 15
+  unk_0 = bits[6:5]       # (data >> 5) & 3
+  unk_1 = bits[7:7]       # (data >> 7) + 1 & 1
+  unk_2 = bits[15:9]      # (data >> 9) & 0x7f
+  unk_padding = bits[47:16]
+
+CDNA_PKT_TYPES: dict[int, type[PacketType]] = {
+  0: CDNA_DELTA, 1: CDNA_TIMESTAMP, 2: CDNA_PKT_2, 3: CDNA_WAVESTART, 4: CDNA_PKT_4,
+  5: CDNA_PKT_5, 6: CDNA_WAVEEND, 10: CDNA_EXEC, 11: CDNA_PKT_11, 13: CDNA_INST, 14: CDNA_PKT_14, 15: CDNA_PKT_15,
+}
+
+# Validate CDNA packet definitions
+for pkt_fmt, pkt_cls in CDNA_PKT_TYPES.items():
+  assert pkt_cls.encoding.default == pkt_fmt, f"{pkt_cls.__name__} encoding {pkt_cls.encoding.default} != pkt_fmt {pkt_fmt}"
+  assert CDNA_PKT_SIZES[pkt_fmt] * 2 == pkt_cls._size_nibbles, f"{pkt_cls.__name__} size {pkt_cls._size_nibbles//2} != {CDNA_PKT_SIZES[pkt_fmt]}"
+
+def decode(data: bytes) -> Iterator[PacketType]:
+  """Decode CDNA SQTT blob using 16-bit header format."""
+  pos, time, ts_offset = 0, 0, None
+  while pos + 2 <= len(data):
+    header = int.from_bytes(data[pos:pos+2], 'little')
+    pkt_fmt = header & 0xf
+    pkt_size = CDNA_PKT_SIZES[pkt_fmt]
+    if pos + pkt_size > len(data): break
+
+    raw = int.from_bytes(data[pos:pos+pkt_size], 'little')
+    # pkt_fmt=0 has delta in bits[11:4], accumulate it
+    if pkt_fmt == 0: time += ((raw >> 4) & 0xff) * 4
+    # pkt_fmt=1 with unk_0=0 is absolute timestamp - use it to anchor time
+    if pkt_fmt == 1 and ((raw >> 4) & 0xfff) == 0:
+      abs_ts = raw >> 16
+      if ts_offset is None: ts_offset = abs_ts - time  # first timestamp: save offset
+      else: time = ((abs_ts - ts_offset) & ~3) - 4     # subsequent: compute time, align to 4, subtract 4
+    pkt_cls = CDNA_PKT_TYPES[pkt_fmt]
+    yield pkt_cls.from_raw(raw, time)
+    pos += pkt_size
+
+if __name__ == "__main__":
+  import sys, pickle
+  if len(sys.argv) < 2:
+    print("Usage: python sqtt_cdna.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} ===")
+    for pkt in decode(event.blob):
+      print(f"{pkt._time:8}: {pkt}")

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 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,266 @@
+#!/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, decode_program
+from extra.assembly.amd import decode_inst
+from extra.assembly.amd.autogen.rdna3.ins import SOPP, SOPPOp
+
+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 profile_instructions(kernel: bytes):
+  """Profile individual instruction compile times."""
+  from extra.assembly.amd.emu import _get_runner, _canonical_runner_cache
+  from tinygrad.helpers import Context
+  _get_runner.cache_clear()
+  _canonical_runner_cache.clear()
+
+  results = []
+  i = 0
+  while i < len(kernel):
+    inst = decode_inst(kernel[i:])
+    if isinstance(inst, SOPP) and inst.op == SOPPOp.S_CODE_END: break
+    inst_bytes = bytes(kernel[i:i + inst.size() + 4])
+    try: inst_str = repr(inst)
+    except Exception: inst_str = f"<{type(inst).__name__}>"
+
+    # Time the full compile (sink + render + compile)
+    start = time.perf_counter()
+    with Context(CCACHE=0):
+      runner, is_new = _get_runner(inst_bytes)
+    compile_time = time.perf_counter() - start
+
+    results.append({
+      'inst_str': inst_str + ('' if is_new else ' [CACHED]'),
+      'compile_ms': compile_time * 1000 if is_new else 0,
+    })
+    i += inst.size()
+
+  return sorted(results, key=lambda x: x['compile_ms'], reverse=True)
+
+def benchmark_python_split(kernel: bytes, global_size, local_size, args_ptr, rsrc2: int, iterations: int = 5):
+  """Benchmark Python emulator with compile and execution times."""
+  from extra.assembly.amd.emu import _get_runner, _canonical_runner_cache
+  from tinygrad.helpers import Context
+  _get_runner.cache_clear()
+  _canonical_runner_cache.clear()
+  decode_program.cache_clear()
+
+  # Measure compile time (decode_program builds sinks, renders, and compiles)
+  compile_start = time.perf_counter()
+  with Context(CCACHE=0):
+    program = decode_program(kernel)
+  compile_time = time.perf_counter() - compile_start
+  n_compiled = len(_canonical_runner_cache)
+
+  # Execution time
+  exec_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, iterations)
+  return compile_time, exec_time, len(program), n_compiled
+
+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", "sin", "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")
+  parser.add_argument("--profile", type=str, default=None, help="Profile instructions for a specific kernel (e.g. 'sin')")
+  parser.add_argument("--top", type=int, default=20, help="Number of top instructions to show in profile")
+  args = parser.parse_args()
+
+  # Profile mode: show individual instruction timing
+  if args.profile:
+    kernel_info = get_tinygrad_kernel(args.profile)
+    if kernel_info is None:
+      print(f"Failed to get kernel for '{args.profile}'")
+      return
+    kernel = kernel_info[0]
+    print(f"Profiling instructions for '{args.profile}' kernel...")
+    print("=" * 110)
+    results = profile_instructions(kernel)
+    print(f"{'Instruction':<90} {'Compile(ms)':>12}")
+    print("-" * 110)
+    for r in results[:args.top]:
+      inst = r['inst_str'][:87] + "..." if len(r['inst_str']) > 90 else r['inst_str']
+      print(f"{inst:<90} {r['compile_ms']:>12.3f}")
+    print("-" * 110)
+    total = sum(r['compile_ms'] for r in results)
+    print(f"{'TOTAL':<90} {total:>12.3f}")
+    return
+
+  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
+    buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
+
+    # Benchmark Python emulator (must be first to measure compile time before cache is populated)
+    py_compile, py_exec, n_insts, n_compiled = benchmark_python_split(kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
+
+    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_compiled} unique) × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
+    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_compile is not None:
+      py_exec_rate = total_work / py_exec / 1e6
+      print(f"  Compile:        {py_compile*1000:8.3f} ms  ({n_compiled} unique)")
+      print(f"  Exec:           {py_exec*1000:8.3f} ms  ({py_exec_rate:7.2f} M ops/s)")
+    if rust_time:
+      rust_rate = total_work / rust_time / 1e6
+      speedup = py_exec / rust_time if py_exec 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_compiled, n_workgroups, py_compile, py_exec, rust_time))
+
+  # Summary table
+  print("\n" + "=" * 110)
+  print("SUMMARY")
+  print("=" * 110)
+  print(f"{'Name':<16} {'Insts':<6} {'Unique':<6} {'WGs':<5} {'Compile (ms)':<14} {'Exec (ms)':<12} {'Rust (ms)':<12} {'Speedup':<10}")
+  print("-" * 110)
+
+  for name, n_insts, n_compiled, n_wgs, py_compile, py_exec, rust_time in results:
+    compile_ms = f"{py_compile*1000:.3f}" if py_compile else "error"
+    exec_ms = f"{py_exec*1000:.3f}" if py_exec else "error"
+    if rust_time:
+      rust_ms = f"{rust_time*1000:.3f}"
+      speedup = f"{py_exec/rust_time:.1f}x" if py_exec else "N/A"
+    else:
+      rust_ms, speedup = "N/A", "N/A"
+    print(f"{name:<16} {n_insts:<6} {n_compiled:<6} {n_wgs:<5} {compile_ms:<14} {exec_ms:<12} {rust_ms:<12} {speedup:<10}")
+
+if __name__ == "__main__":
+  main()

extra/assembly/amd/test/helpers.py

@@ -0,0 +1,38 @@
+"""Shared test helpers for RDNA3 tests."""
+import shutil
+from dataclasses import dataclass
+
+@dataclass
+class KernelInfo:
+  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
+
+# 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]

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

@@ -4,10 +4,10 @@ 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 extra.assembly.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.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]
@@ -47,7 +47,7 @@ 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]
+  return Device["AMD"].target
 
 def skip_unless_gfx(min_major: int, min_minor: int = 0, reason: str = ""):
   """Skip test if GPU target is below the minimum required version."""
@@ -60,12 +60,11 @@ def skip_unless_gfx(min_major: int, min_minor: int = 0, reason: str = ""):
     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)
+OUT_BYTES = VGPR_BYTES + SGPR_BYTES + 12  # + vcc + scc + exec
-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)
 
 # Float conversion helpers
 def f2i(f: float) -> int: return _i32(f)
@@ -76,10 +75,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):
+  def __init__(self):
-    self.vgpr = [[0] * 256 for _ in range(n_lanes)]  # vgpr[lane][reg]
+    self.vgpr = [[0] * 256 for _ in range(32)]  # vgpr[lane][reg]
     self.sgpr = [0] * 128
     self.vcc = 0
     self.scc = 0
@@ -103,53 +102,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], 0))
-  epilogue.append(v_mov_b32_e32(v[240], vgpr_bytes))
   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()
-  st = WaveState(n_lanes)
   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.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.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.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 * OUT_BYTES)(*([0] * OUT_BYTES))
-  out_buf = (ctypes.c_uint8 * buf_sz)(*([0] * buf_sz))
   out_addr = ctypes.addressof(out_buf)
 
   prologue, epilogue = get_prologue_epilogue(n_lanes)
@@ -176,7 +171,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,14 +218,13 @@ 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(OUT_BYTES)
-  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}"
   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)
@@ -282,6 +276,6 @@ def run_program(instructions: list, n_lanes: int = 1, ulp_tolerance: int = 0) ->
     hw_st = run_program_hw(instructions, n_lanes)
     diffs = compare_wave_states(emu_st, hw_st, n_lanes, ulp_tolerance=ulp_tolerance)
     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."""
@@ -601,6 +601,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 = [
@@ -760,110 +761,5 @@ class TestDsPermute(unittest.TestCase):
     self.assertEqual(st.vgpr[0][2], 0x11111111)
 
 
-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."""

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

@@ -3,7 +3,7 @@
 Includes: scratch_load_*, scratch_store_*
 """
 import unittest
-from test.amd.hw.helpers import *
+from extra.assembly.amd.test.hw.helpers import *
 
 class TestScratchStore(unittest.TestCase):
   """Tests for SCRATCH store instructions."""

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

@@ -4,7 +4,7 @@ Includes: s_load_b32, s_load_b64, s_load_b128, s_load_b256, s_load_b512
 Tests both immediate and register offset addressing modes.
 """
 import unittest
-from test.amd.hw.helpers import *
+from extra.assembly.amd.test.hw.helpers import *
 
 # Use offset into output buffer for test data (output buffer is 2124 bytes)
 TEST_OFFSET = 2000

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."""
@@ -87,50 +87,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."""
 
@@ -932,76 +888,5 @@ class Test64BitSOPLiterals(unittest.TestCase):
     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."""
@@ -373,6 +373,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]),
@@ -1579,55 +1580,5 @@ class TestPermlane64(unittest.TestCase):
     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."""

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

@@ -4,7 +4,7 @@ Includes: v_fma_f32, v_div_scale_f32, v_div_fmas_f32, v_div_fixup_f32,
           v_alignbit_b32, v_bfe_i32, v_mad_u64_u32, v_readlane_b32, v_writelane_b32
 """
 import unittest
-from test.amd.hw.helpers import *
+from extra.assembly.amd.test.hw.helpers import *
 
 class TestFMA(unittest.TestCase):
   """Tests for FMA instructions."""
@@ -725,7 +725,7 @@ class TestLaneOps(unittest.TestCase):
     # v[5] should have the value only in lane 1
     for lane in range(4):
       if lane == 1:
-        self.assertEqual(st.vgpr[lane][5], 0x12345678, "v[5] lane 1 should have 0x12345678")
+        self.assertEqual(st.vgpr[lane][5], 0x12345678, f"v[5] lane 1 should have 0x12345678")
       else:
         self.assertEqual(st.vgpr[lane][5], 0, f"v[5] lane {lane} should be 0")
 
@@ -1082,6 +1082,7 @@ class TestF64Ops(unittest.TestCase):
     """Full f64->i64 conversion sequence with negative value."""
     import struct
     val = f2i64(-8.0)
+    lit = 0xC1F00000  # high 32 bits of f64 -2^32
     instructions = [
       s_mov_b32(s[0], val & 0xffffffff),
       s_mov_b32(s[1], (val >> 32) & 0xffffffff),
@@ -1137,6 +1138,7 @@ class TestF64Ops(unittest.TestCase):
     # v_fma_f64 v[7:8], v[17:18], v[7:8], v[15:16]
     # We need to capture the exact input values and verify output matches hardware
     # v[7:8] before = 0x3f80fdf3_d69db28f (0.008296875941334462)
+    v78 = 0x3f80fdf3d69db28f
     # For the FMA to produce 0xbf457ef0_ab8c254d, we need v[17:18] and v[15:16]
     # Let's test with known precision-sensitive values
     a = 1.0000000001
@@ -1393,7 +1395,7 @@ class TestWMMAMore(unittest.TestCase):
 
   def test_v_wmma_f32_16x16x16_f16_basic(self):
     """V_WMMA_F32_16X16X16_F16 basic test - verify output is non-zero."""
-    instructions: list[Inst] = []
+    instructions = []
     instructions.append(s_mov_b32(s[0], 0x3c003c00))
     for i in range(16, 32):
       instructions.append(v_mov_b32_e32(v[i], s[0]))
@@ -1849,6 +1851,7 @@ class TestMed3(unittest.TestCase):
 
   def test_v_med3_f32_with_nan(self):
     """V_MED3_F32: NaN handling - returns min of non-NaN values."""
+    import math
     instructions = [
       s_mov_b32(s[0], 0x7fc00000),  # NaN
       v_mov_b32_e32(v[0], s[0]),
@@ -2487,6 +2490,7 @@ class TestDivScaleF64(unittest.TestCase):
     independently. This catches the bug where the emulator was setting VCC
     for all lanes to the same value.
     """
+    import math
     # Use lane-varying input: lane 0 gets 2.0, lane 1 gets 3.0, etc.
     # All normal values should result in VCC=0 for each lane
     instructions = [
@@ -2717,6 +2721,7 @@ class TestDivScaleFmasF64Integration(unittest.TestCase):
     This is the exact bug scenario: tan([2.0, 3.0, 4.0]) was failing because
     VCC from DIV_SCALE was being set incorrectly for all lanes.
     """
+    import math
     # Set up values like tan() would: different values per lane
     instructions = [
       # Create per-lane values: 2.0, 3.0, 4.0, 5.0
@@ -2754,7 +2759,7 @@ class TestVOP3VOPC(unittest.TestCase):
 
   def test_v_cmp_ge_f32_e64_nan(self):
     """V_CMP_GE_F32_E64: |NaN| >= |0.0| should be FALSE (NaN comparisons always false)."""
-    from tinygrad.runtime.autogen.amd.rdna3.ins import VOP3_SDST
+    from extra.assembly.amd.autogen.rdna3.ins import VOP3_SDST
     instructions = [
       s_mov_b32(s[0], 0xffc00000),  # NaN
       s_mov_b32(s[1], 0x00000000),  # 0.0
@@ -2944,306 +2949,6 @@ class TestVOP3Clamp(unittest.TestCase):
     self.assertAlmostEqual(i2f(st.vgpr[3][1]), 1.0, places=5, msg="lane 3: 2.5 should clamp to 1.0")
 
 
-class TestVOP3ClampUint32(unittest.TestCase):
-  """Tests for VOP3 clamp modifier on unsigned 32-bit integer operations."""
-
-  def test_v_sub_nc_u32_e64_clamp_underflow(self):
-    """V_SUB_NC_U32_E64 with clamp: 0 - 1 should saturate to 0."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0),
-      v_mov_b32_e32(v[1], 1),
-      v_sub_nc_u32_e64(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_sub_nc_u32_e64_clamp_no_underflow(self):
-    """V_SUB_NC_U32_E64 with clamp: 100 - 50 = 50 (no saturation needed)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 100),
-      v_mov_b32_e32(v[1], 50),
-      v_sub_nc_u32_e64(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 50, f"expected 50, got {st.vgpr[0][2]}")
-
-  def test_v_add_nc_u32_e64_clamp_overflow(self):
-    """V_ADD_NC_U32_E64 with clamp: 0xFFFFFFFF + 1 should saturate to 0xFFFFFFFF."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0xFFFFFFFF),
-      v_mov_b32_e32(v[1], 1),
-      v_add_nc_u32_e64(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_add_nc_u32_e64_clamp_no_overflow(self):
-    """V_ADD_NC_U32_E64 with clamp: 100 + 50 = 150 (no saturation needed)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 100),
-      v_mov_b32_e32(v[1], 50),
-      v_add_nc_u32_e64(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 150, f"expected 150, got {st.vgpr[0][2]}")
-
-
-class TestVOP3ClampUint16(unittest.TestCase):
-  """Tests for VOP3 clamp modifier on unsigned 16-bit integer operations."""
-
-  def test_v_sub_nc_u16_clamp_underflow(self):
-    """V_SUB_NC_U16 with clamp: 0 - 1 should saturate to 0."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0),
-      v_mov_b32_e32(v[1], 1),
-      v_sub_nc_u16(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 0, f"expected 0, got 0x{st.vgpr[0][2] & 0xFFFF:04x}")
-
-  def test_v_sub_nc_u16_clamp_no_underflow(self):
-    """V_SUB_NC_U16 with clamp: 100 - 50 = 50 (no saturation needed)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 100),
-      v_mov_b32_e32(v[1], 50),
-      v_sub_nc_u16(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 50, f"expected 50, got {st.vgpr[0][2] & 0xFFFF}")
-
-  def test_v_add_nc_u16_clamp_overflow(self):
-    """V_ADD_NC_U16 with clamp: 0xFFFF + 1 should saturate to 0xFFFF."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0xFFFF),
-      v_mov_b32_e32(v[1], 1),
-      v_add_nc_u16(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][2] & 0xFFFF:04x}")
-
-  def test_v_add_nc_u16_clamp_no_overflow(self):
-    """V_ADD_NC_U16 with clamp: 100 + 50 = 150 (no saturation needed)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 100),
-      v_mov_b32_e32(v[1], 50),
-      v_add_nc_u16(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2] & 0xFFFF, 150, f"expected 150, got {st.vgpr[0][2] & 0xFFFF}")
-
-
-class TestVOP3ClampInt32(unittest.TestCase):
-  """Tests for VOP3 clamp modifier on signed 32-bit integer operations."""
-
-  def test_v_add_nc_i32_clamp_overflow(self):
-    """V_ADD_NC_I32 with clamp: INT_MAX + 1 should saturate to INT_MAX."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0x7FFFFFFF),  # S0 = INT_MAX
-      v_mov_b32_e32(v[1], 1),            # S1 = 1
-      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF, f"expected 0x7FFFFFFF, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_add_nc_i32_clamp_underflow(self):
-    """V_ADD_NC_I32 with clamp: INT_MIN + (-1) should saturate to INT_MIN."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0x80000000),  # S0 = INT_MIN
-      v_mov_b32_e32(v[1], 0xFFFFFFFF),  # S1 = -1
-      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x80000000, f"expected 0x80000000, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_sub_nc_i32_clamp_underflow(self):
-    """V_SUB_NC_I32 with clamp: INT_MIN - 1 should saturate to INT_MIN."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0x80000000),  # S0 = INT_MIN
-      v_mov_b32_e32(v[1], 1),            # S1 = 1
-      v_sub_nc_i32(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x80000000, f"expected 0x80000000, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_sub_nc_i32_clamp_overflow(self):
-    """V_SUB_NC_I32 with clamp: INT_MAX - (-1) should saturate to INT_MAX."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0x7FFFFFFF),  # S0 = INT_MAX
-      v_mov_b32_e32(v[1], 0xFFFFFFFF),  # S1 = -1
-      v_sub_nc_i32(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF, f"expected 0x7FFFFFFF, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_add_nc_i32_no_saturation_positive(self):
-    """V_ADD_NC_I32 with clamp: 100 + 200 = 300 (no saturation needed)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 100),
-      v_mov_b32_e32(v[1], 200),
-      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 300, f"expected 300, got {st.vgpr[0][2]}")
-
-  def test_v_add_nc_i32_no_saturation_negative(self):
-    """V_ADD_NC_I32 with clamp: -100 + -200 = -300 (no saturation needed)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0xFFFFFF9C),  # -100
-      v_mov_b32_e32(v[1], 0xFFFFFF38),  # -200
-      v_add_nc_i32(v[2], v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    expected = 0xFFFFFED4  # -300
-    self.assertEqual(st.vgpr[0][2], expected, f"expected 0x{expected:08x}, got 0x{st.vgpr[0][2]:08x}")
-
-
-class TestVOP3ClampCarry(unittest.TestCase):
-  """Tests for VOP3 clamp modifier on carry operations (VOP3SD)."""
-
-  def test_v_add_co_u32_clamp_overflow(self):
-    """V_ADD_CO_U32 with clamp: 0xFFFFFFFF + 1 should saturate to 0xFFFFFFFF."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0xFFFFFFFF),
-      v_mov_b32_e32(v[1], 1),
-      v_add_co_u32(v[2], VCC, v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_add_co_u32_clamp_no_overflow(self):
-    """V_ADD_CO_U32 with clamp: 100 + 200 = 300 (no saturation)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 100),
-      v_mov_b32_e32(v[1], 200),
-      v_add_co_u32(v[2], VCC, v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 300, f"expected 300, got {st.vgpr[0][2]}")
-
-  def test_v_sub_co_u32_clamp_underflow(self):
-    """V_SUB_CO_U32 with clamp: 0 - 1 should saturate to 0."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0),
-      v_mov_b32_e32(v[1], 1),
-      v_sub_co_u32(v[2], VCC, v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_sub_co_u32_clamp_no_underflow(self):
-    """V_SUB_CO_U32 with clamp: 300 - 100 = 200 (no saturation)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 300),
-      v_mov_b32_e32(v[1], 100),
-      v_sub_co_u32(v[2], VCC, v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 200, f"expected 200, got {st.vgpr[0][2]}")
-
-  def test_v_subrev_co_u32_clamp_underflow(self):
-    """V_SUBREV_CO_U32 with clamp: 1 - 0 reversed = 0 - 1 should saturate to 0."""
-    instructions = [
-      v_mov_b32_e32(v[0], 1),  # This becomes the subtrahend
-      v_mov_b32_e32(v[1], 0),  # This becomes the minuend (0 - 1)
-      v_subrev_co_u32(v[2], VCC, v[0], v[1], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_add_co_ci_u32_clamp_overflow(self):
-    """V_ADD_CO_CI_U32 with clamp: 0xFFFFFFFF + 1 + 0 should saturate to 0xFFFFFFFF."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0xFFFFFFFF),
-      v_mov_b32_e32(v[1], 1),
-      s_mov_b64(VCC, 0),  # No carry in
-      v_add_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_add_co_ci_u32_clamp_overflow_with_carry(self):
-    """V_ADD_CO_CI_U32 with clamp: 0xFFFFFFFE + 1 + 1 should saturate to 0xFFFFFFFF."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0xFFFFFFFE),
-      v_mov_b32_e32(v[1], 1),
-      s_mov_b64(VCC, 1),  # Carry in = 1
-      v_add_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_sub_co_ci_u32_clamp_underflow(self):
-    """V_SUB_CO_CI_U32 with clamp: 0 - 1 - 0 should saturate to 0."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0),
-      v_mov_b32_e32(v[1], 1),
-      s_mov_b64(VCC, 0),  # No borrow in
-      v_sub_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
-
-  def test_v_subrev_co_ci_u32_clamp_underflow(self):
-    """V_SUBREV_CO_CI_U32 with clamp: reversed 1 - 0 - 0 = 0 - 1 should saturate to 0."""
-    instructions = [
-      v_mov_b32_e32(v[0], 1),
-      v_mov_b32_e32(v[1], 0),
-      s_mov_b64(VCC, 0),
-      v_subrev_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
-
-
-class TestVOP3ClampMAD(unittest.TestCase):
-  """Tests for VOP3 clamp modifier on MAD (multiply-add) operations."""
-
-  def test_v_mad_u16_clamp_overflow(self):
-    """V_MAD_U16 with clamp: 0xFFFF * 2 + 0 should saturate to 0xFFFF."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0xFFFF),
-      v_mov_b32_e32(v[1], 2),
-      v_mov_b32_e32(v[2], 0),
-      v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
-
-  def test_v_mad_u16_clamp_overflow_with_add(self):
-    """V_MAD_U16 with clamp: 0x8000 * 2 + 0x1000 should saturate to 0xFFFF."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0x8000),  # 32768
-      v_mov_b32_e32(v[1], 2),        # * 2 = 65536
-      v_mov_b32_e32(v[2], 0x1000),   # + 4096 = 69632 > 0xFFFF
-      v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
-
-  def test_v_mad_u16_no_overflow(self):
-    """V_MAD_U16 with clamp: 100 * 100 + 50 = 10050 (no saturation)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 100),
-      v_mov_b32_e32(v[1], 100),
-      v_mov_b32_e32(v[2], 50),
-      v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 10050, f"expected 10050, got {st.vgpr[0][3] & 0xFFFF}")
-
-  def test_v_mad_u16_no_clamp(self):
-    """V_MAD_U16 without clamp: 0xFFFF * 2 + 0 should wrap to 0xFFFE."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0xFFFF),
-      v_mov_b32_e32(v[1], 2),
-      v_mov_b32_e32(v[2], 0),
-      v_mad_u16(v[3], v[0], v[1], v[2], clmp=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    # 0xFFFF * 2 = 0x1FFFE, low 16 bits = 0xFFFE
-    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFE, f"expected 0xFFFE, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
-
-
 class TestCvtPkF16(unittest.TestCase):
   """Tests for V_CVT_PK_RTZ_F16_F32 - pack two f32 to f16 with round toward zero."""
 

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."""
@@ -390,24 +390,6 @@ class TestVOP3P(unittest.TestCase):
     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}")
-
 
 class TestWMMAF16(unittest.TestCase):
   """Tests for WMMA F16 output variant (V_WMMA_F16_16X16X16_F16).
@@ -418,7 +400,7 @@ class TestWMMAF16(unittest.TestCase):
 
   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 = []
     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):
@@ -442,7 +424,7 @@ class TestWMMAF16(unittest.TestCase):
 
   def test_v_wmma_f16_16x16x16_f16_with_accumulator(self):
     """V_WMMA_F16_16X16X16_F16 with non-zero accumulator."""
-    instructions: list[Inst] = []
+    instructions = []
     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)
@@ -471,7 +453,7 @@ class TestWMMAF16(unittest.TestCase):
     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 = []
     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):
@@ -502,7 +484,7 @@ class TestWMMA(unittest.TestCase):
 
   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]))
@@ -518,7 +500,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):
@@ -540,7 +522,7 @@ class TestWMMA(unittest.TestCase):
     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 = []
     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):
@@ -569,7 +551,7 @@ class TestWMMABF16(unittest.TestCase):
 
   def test_v_wmma_f32_16x16x16_bf16_all_ones(self):
     """V_WMMA_F32_16X16X16_BF16 with all ones produces 16.0."""
-    instructions: list[Inst] = []
+    instructions = []
     # BF16 1.0 = 0x3f80, packed = 0x3f803f80
     instructions.append(s_mov_b32(s[0], 0x3f803f80))
     for i in range(16, 32):
@@ -586,7 +568,7 @@ class TestWMMABF16(unittest.TestCase):
 
   def test_v_wmma_f32_16x16x16_bf16_with_accumulator(self):
     """V_WMMA_F32_16X16X16_BF16 with non-zero accumulator."""
-    instructions: list[Inst] = []
+    instructions = []
     # BF16 1.0 = 0x3f80, packed = 0x3f803f80
     instructions.append(s_mov_b32(s[0], 0x3f803f80))
     instructions.append(s_mov_b32(s[1], f2i(5.0)))

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 = [

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 extra.assembly.amd.test.hw.helpers import run_program, run_program_emu, run_program_hw, compare_wave_states, \
-from tinygrad.runtime.autogen.amd.rdna3.ins import VOPD, VOPD_LIT, VOPDOp
+  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

extra/assembly/amd/test/test_compare_emulators.py

@@ -1,14 +1,11 @@
 # Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
 import unittest, ctypes
 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 extra.assembly.amd.emu import WaveState, decode_program, WAVE_SIZE, VCC_LO, EXEC_LO, SCC
-from tinygrad.renderer.amd import decode_inst
+from extra.assembly.amd import decode_inst
-import tinygrad
+from extra.assembly.amd.test.helpers import KernelInfo
-REMU_PATH = Path(tinygrad.__file__).parent.parent / "extra/remu/target/release/libremu.so"
+from extra.assembly.amd.test.bench_emu import REMU_PATH
-if not REMU_PATH.exists(): REMU_PATH = Path(tinygrad.__file__).parent.parent / "extra/remu/target/release/libremu.dylib"
 
 def set_valid_mem_ranges(ranges): pass  # emu2 doesn't need this
 
@@ -21,15 +18,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,14 +81,12 @@ class RustEmulator:
     return snap.to_snapshot()
 
   def free(self):
-    if self.ctx:
+    if self.ctx: self.lib.wave_free(self.ctx); self.ctx = None
-      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.program: dict | None = None
     self.vmem_buf = None
     self.lds_buf = None
     self.kernel_buf = None  # Keep kernel bytes alive
@@ -110,29 +96,27 @@ class PythonEmulator:
     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
+    # Store kernel in a ctypes buffer so generic instructions can read from vmem 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 = {}
+    # Remap program dict to use actual addresses (like run_asm does)
+    program_raw = decode_program(kernel)
+    self.program = {self.lib_addr + offset: val for offset, val in program_raw.items()}
     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)
-
   def step(self) -> int:
     import ctypes
-    assert self.state is not None
+    assert self.program is not None and self.state is not None
     pc = self.state.pc
-    if pc == 0xFFFFFFFFFFFFFFFF: return -1
+    if pc == 0xFFFFFFFFFFFFFFFF or pc not in self.program: return -1
-    self._ensure_decoded(pc)
+    name, fxn, globals_list, _runner = self.program[pc]
-    name, fxn, globals_list = self.program[pc]
+    if fxn is None: return 1  # unsupported instruction
-    buf_addrs = {0: self.state.sgpr_buf._buf.va_addr, 1: self.state.vgpr_buf._buf.va_addr,  # type: ignore[union-attr]
+    buf_addrs = {0: self.state.sgpr_buf._buf.va_addr, 1: self.state.vgpr_buf._buf.va_addr,
-                 2: self.vmem_buf._buf.va_addr, 3: self.lds_buf._buf.va_addr}  # type: ignore[union-attr]
+                 2: self.vmem_buf._buf.va_addr, 3: self.lds_buf._buf.va_addr}
+    # Direct ctypes call - bypasses HCQ overhead
     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
 
@@ -153,7 +137,7 @@ class PythonEmulator:
                          exec_mask=sgpr[EXEC_LO.offset], sgpr=sgpr, vgpr=vgpr)
 
 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,
+                      local_size: tuple[int, int, int], program, max_steps: int, debug: bool, trace_len: int,
                       kernel_idx: int = 0, max_workgroups: int = 8) -> tuple[bool, str, int]:
   """Run a single kernel through both emulators. Returns (success, message, total_steps)."""
   gx, gy, gz = global_size
@@ -194,9 +178,8 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
             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
+            inst_info = python.program.get(python.lib_addr + python_before.pc * 4)  # Convert word offset to actual address
-            python._ensure_decoded(pc_addr)
+            inst_hex_name = inst_info[0] if inst_info else f"unknown at PC={python_before.pc}"
-            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
@@ -205,7 +188,7 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
               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:
+            except:
               inst_mnemonic = ""
             # For generic instructions, use function name for sync_after check
             if not inst_mnemonic: inst_mnemonic = inst_hex_name
@@ -237,18 +220,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  "
+              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
-              msg += "\n  ".join(diffs[:10]) + f"\n  Recent instructions:\n{trace_str}"
-              return False, msg, total_steps
 
             rust_result = rust.step()
             python_result = python.step()
@@ -258,9 +239,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: "
+              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
-                     f"rust={rust_result}, python={python_result}, inst={inst_str}\n  Recent instructions:\n{trace_str}")
-              return False, msg, total_steps
 
             # Sync Python state to Rust after instructions with known Rust emulator differences
             if sync_after:
@@ -293,7 +272,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 +302,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
+      kernel.local_size, program, max_steps, debug, trace_len, ki
     )
     total_steps += steps
     if not ok:
@@ -353,11 +333,12 @@ 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)
 
+  program = decode_program(kernel)
   # 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)
+  ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, (n_lanes, 1, 1), 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,7 +376,7 @@ 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),
@@ -412,7 +393,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 +429,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):
+  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)))
-    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)

extra/assembly/amd/test/test_custom_kernel.py

@@ -3,12 +3,20 @@ import functools
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-from tinygrad.runtime.autogen.amd.rdna3.ins import *
+from tinygrad.runtime.support.compiler_amd import HIPCompiler
-from tinygrad.runtime.autogen.amd.rdna4.ins import s_barrier_wait, s_barrier_signal
-from tinygrad.renderer.amd.dsl import s, v
-from test.amd.helpers import TARGET_TO_ARCH
 
-def custom_add_one(A:UOp) -> UOp:
+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")
@@ -23,10 +31,10 @@ def custom_add_one(A:UOp) -> UOp:
     global_store_b32(addr=v[0], data=v[1], saddr=s[0:1]),
     s_endpgm(),
   ]
-  sink = UOp.sink(A.base, threads, arg=KernelInfo(f"custom_add_one_{A.size}", estimates=Estimates(ops=A.size, mem=A.size*4*2)))
+  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=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=(*sink.src, sink)), *assemble_insts(insts, name, arch)))
 
-def custom_add_var(A:UOp, B:UOp) -> UOp:
+def custom_add_var(A:UOp, B:UOp, arch:str) -> UOp:
   A,B = A.flatten(), B.flatten()
   assert A.dtype.base == dtypes.uint32, f"buffer dtype must be uint32, got {A.dtype}"
   threads = UOp.special(A.size, "lidx0")
@@ -42,31 +50,14 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
     global_store_b32(addr=v[0], data=v[1], saddr=s[4:5]),
     s_endpgm(),
   ]
-  sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.size}"))
+  sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(name:=f"custom_add_one_{A.size}"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               *assemble_insts(insts, name, arch, kernarg_size=16)))
 
-def custom_wave_sync(A:UOp, arch:str) -> UOp:
-  # 4 waves across 1024 WG — enough to saturate a SIMD with many concurrent WGs
-  # s_sleep yields the SIMD so waves from different WGs interleave, causing barrier packet reordering
-  threads = UOp.special(128, "lidx0")
-  wg = UOp.special(1024, "gidx0")
-  insts = []
-  for _ in range(4):
-    insts.append(s_sleep(4))
-    insts += [s_barrier()] if arch == "rdna3" else [s_barrier_signal(), s_barrier_wait()]
-    insts += [s_nop(0)]*4
-  insts.append(s_endpgm())
-  sink = UOp.sink(A.base, threads, wg, arg=KernelInfo("custom_wave_sync"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-
-@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestCustomKernel(unittest.TestCase):
-  def setUp(self): self.arch = TARGET_TO_ARCH[Device["AMD"].arch]
-
   def test_simple(self):
-    if self.arch != "rdna3": self.skipTest("only rdna3")
     a = Tensor.full((16, 16), 1.).contiguous().realize()
-    a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
+    a = Tensor.custom_kernel(a, fxn=functools.partial(custom_add_one, arch=Device[Device.DEFAULT].renderer.arch))[0]
     ei = a.schedule()[-1].lower()
     self.assertEqual(ei.prg.estimates.ops, a.numel())
     self.assertEqual(ei.prg.estimates.mem, a.nbytes()*2)
@@ -74,18 +65,13 @@ class TestCustomKernel(unittest.TestCase):
     self.assertTrue((a.numpy() == 2.).all())
 
   def test_variable(self):
-    if self.arch != "rdna3": self.skipTest("only rdna3")
     b = Tensor.full((16, 16), 1, dtype=dtypes.uint32).contiguous().realize()
     a = Tensor.zeros_like(b).contiguous().realize()
-    a = Tensor.custom_kernel(a, b, fxn=custom_add_var)[0]
+    a = Tensor.custom_kernel(a, b, fxn=functools.partial(custom_add_var, arch=Device[Device.DEFAULT].renderer.arch))[0]
     ei = a.schedule()[-1].lower()
     for i in range(4):
       ei.run({"var":i})
       self.assertTrue((a.numpy() == 1+i).all())
 
-  def test_wave_sync(self):
-    if self.arch not in {"rdna3", "rdna4"}: self.skipTest("only rdna3 or rdna4")
-    Tensor.empty(1).custom_kernel(fxn=functools.partial(custom_wave_sync, arch=self.arch))[0].realize()
-
 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 extra.assembly.amd.dsl import *
-from tinygrad.renderer.amd.dsl import VDSTYField
+from extra.assembly.amd.dsl import VDSTYField
-from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op
+from extra.assembly.amd.autogen.rdna3.enum import VOP1Op, VOP2Op
-from tinygrad.runtime.autogen.amd.rdna3.ins import VOP1
+from extra.assembly.amd.autogen.rdna3.ins import VOP1
 
 class TestRegisters(unittest.TestCase):
   def test_vgpr_single(self):

extra/assembly/amd/test/test_emu2_pcode.py

@@ -4,14 +4,14 @@ from collections import defaultdict
 from tinygrad.helpers import DEBUG
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops
-from test.mockgpu.amd.emu import parse_pcode
+from extra.assembly.amd.emu import parse_pcode
-from test.mockgpu.amd.pcode import parse_expr
+from extra.assembly.amd.pcode import parse_expr
-from tinygrad.runtime.autogen.amd.rdna3.str_pcode import PCODE
+from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
-from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op, SOP2Op, DSOp
+from extra.assembly.amd.autogen.rdna3.enum import VOP1Op, VOP2Op, VOP3Op, SOP1Op, SOP2Op, DSOp
 
 def _srcs():
   """Create minimal source variables for pcode parsing."""
-  def u32(v=0): return UOp.const(dtypes.uint32, v)
+  u32 = lambda v=0: UOp.const(dtypes.uint32, v)
   return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(dtypes.uint64, 0), 'laneId': u32()}
 
 class TestBasicParsing(unittest.TestCase):
@@ -90,16 +90,16 @@ class TestParseExpr(unittest.TestCase):
 
   def test_variable_lookup(self):
     """Test variable lookup in parse_expr."""
-    vrs = {'x': UOp.const(dtypes.uint32, 42)}
+    vars = {'x': UOp.const(dtypes.uint32, 42)}
-    result = parse_expr('x', vrs)
+    result = parse_expr('x', vars)
     self.assertEqual(result.arg, 42)
 
   def test_binary_ops(self):
     """Test parsing binary operations."""
-    vrs = {'a': UOp.const(dtypes.uint32, 10), 'b': UOp.const(dtypes.uint32, 5)}
+    vars = {'a': UOp.const(dtypes.uint32, 10), 'b': UOp.const(dtypes.uint32, 5)}
 
     # Addition
-    result = parse_expr('a + b', vrs)
+    result = parse_expr('a + b', vars)
     self.assertEqual(result.op, Ops.ADD)
 
     # Subtraction with constant folding
@@ -109,8 +109,8 @@ class TestParseExpr(unittest.TestCase):
 
   def test_ternary(self):
     """Test parsing ternary expressions."""
-    vrs = {'cond': UOp.const(dtypes.bool, True), 'a': UOp.const(dtypes.uint32, 1), 'b': UOp.const(dtypes.uint32, 0)}
+    vars = {'cond': UOp.const(dtypes.bool, True), 'a': UOp.const(dtypes.uint32, 1), 'b': UOp.const(dtypes.uint32, 0)}
-    result = parse_expr('cond ? a : b', vrs)
+    result = parse_expr('cond ? a : b', vars)
     self.assertEqual(result.op, Ops.WHERE)
 
 class TestForLoopParsing(unittest.TestCase):
@@ -120,14 +120,13 @@ class TestForLoopParsing(unittest.TestCase):
     """Verify CLZ pcode is available."""
     pcode = PCODE.get(VOP1Op.V_CLZ_I32_U32_E32)
     self.assertIsNotNone(pcode)
-    assert pcode is not None
     self.assertIn('for', pcode.lower())
 
   def test_clz_parsing(self):
     """Test CLZ pcode parsing produces correct structure."""
     pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
     S0 = UOp.const(dtypes.uint32, 0xFFFFFFFF)  # All ones - CLZ should be 0
-    _vrs, assigns = parse_pcode(pcode, {'S0': S0})
+    vars, assigns = parse_pcode(pcode, {'S0': S0})
 
     self.assertEqual(len(assigns), 1)
     dest, val = assigns[0]
@@ -139,7 +138,7 @@ class TestForLoopParsing(unittest.TestCase):
     """Test CLZ with input 0 - should return -1."""
     pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
     S0 = UOp.const(dtypes.uint32, 0)
-    _vrs, assigns = parse_pcode(pcode, {'S0': S0})
+    vars, assigns = parse_pcode(pcode, {'S0': S0})
 
     # Check that the innermost value (default) is -1 (may be wrapped in CAST)
     val = assigns[0][1]
@@ -158,7 +157,7 @@ class TestForLoopParsing(unittest.TestCase):
       self.skipTest("V_CTZ_I32_B32_E32 pcode not available")
 
     S0 = UOp.const(dtypes.uint32, 1)  # LSB set - CTZ should be 0
-    _vrs, assigns = parse_pcode(pcode, {'S0': S0})
+    vars, assigns = parse_pcode(pcode, {'S0': S0})
     self.assertEqual(len(assigns), 1)
 
 class TestDSPcodePatterns(unittest.TestCase):
@@ -168,7 +167,6 @@ class TestDSPcodePatterns(unittest.TestCase):
     """Test DS_LOAD_B32 pcode is parseable."""
     pcode = PCODE.get(DSOp.DS_LOAD_B32)
     self.assertIsNotNone(pcode)
-    assert pcode is not None
     self.assertIn('RETURN_DATA', pcode)
     self.assertIn('MEM[', pcode)
 
@@ -176,18 +174,17 @@ class TestDSPcodePatterns(unittest.TestCase):
     """Test DS_STORE_B32 pcode is parseable."""
     pcode = PCODE.get(DSOp.DS_STORE_B32)
     self.assertIsNotNone(pcode)
-    assert pcode is not None
     self.assertIn('MEM[', pcode)
     self.assertIn('DATA', pcode)
 
   def test_mem_read_parsing(self):
     """Test MEM[addr].type read expression parsing."""
     # Create a mock LDS buffer
-    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
     addr = UOp.const(dtypes.uint32, 0)
-    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}
+    vars = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}
 
-    result = parse_expr('MEM[ADDR + OFFSET].b32', vrs)
+    result = parse_expr('MEM[ADDR + OFFSET].b32', vars)
     # Should be an INDEX operation into LDS
     self.assertIsNotNone(result)
 
@@ -195,7 +192,6 @@ class TestDSPcodePatterns(unittest.TestCase):
     """Test DS_STORE_2ADDR_B32 pcode parsing produces MEM writes."""
     pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
     self.assertIsNotNone(pcode)
-    assert pcode is not None
     srcs = {
       'ADDR': UOp.const(dtypes.uint32, 0),
       'OFFSET0': UOp.const(dtypes.uint32, 0),
@@ -211,14 +207,13 @@ class TestDSPcodePatterns(unittest.TestCase):
       self.assertTrue(dest.startswith('MEM['))
       # val should be (addr, write_val) tuple
       self.assertIsInstance(val, tuple)
-      self.assertEqual(len(val), 2)  # type: ignore[arg-type]
+      self.assertEqual(len(val), 2)
 
   def test_ds_load_2addr_b32_parsing(self):
     """Test DS_LOAD_2ADDR_B32 pcode parsing produces RETURN_DATA assignments."""
     pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
     self.assertIsNotNone(pcode)
-    assert pcode is not None
+    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
-    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
     srcs = {
       'ADDR': UOp.const(dtypes.uint32, 0),
       'OFFSET0': UOp.const(dtypes.uint32, 0),
@@ -235,7 +230,6 @@ class TestDSPcodePatterns(unittest.TestCase):
   def test_ds_store_address_calculation(self):
     """Test DS_STORE_2ADDR_B32 calculates correct addresses (offset * 4)."""
     pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
-    assert pcode is not None
     srcs = {
       'ADDR': UOp.const(dtypes.uint32, 100),
       'OFFSET0': UOp.const(dtypes.uint32, 2),
@@ -246,14 +240,14 @@ class TestDSPcodePatterns(unittest.TestCase):
     srcs['laneId'] = UOp.const(dtypes.uint32, 0)
     _, assigns = parse_pcode(pcode, srcs)
     # Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
-    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
+    addr0, _ = assigns[0][1]
-    self.assertEqual(assigns[0][1][0].simplify().arg, 108)  # type: ignore[index]
+    addr1, _ = assigns[1][1]
-    self.assertEqual(assigns[1][1][0].simplify().arg, 120)  # type: ignore[index]
+    self.assertEqual(addr0.simplify().arg, 108)
+    self.assertEqual(addr1.simplify().arg, 120)
 
   def test_ds_store_data_values(self):
     """Test DS_STORE_2ADDR_B32 uses correct data values."""
     pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
-    assert pcode is not None
     srcs = {
       'ADDR': UOp.const(dtypes.uint32, 0),
       'OFFSET0': UOp.const(dtypes.uint32, 0),
@@ -263,10 +257,11 @@ class TestDSPcodePatterns(unittest.TestCase):
     }
     srcs['laneId'] = UOp.const(dtypes.uint32, 0)
     _, assigns = parse_pcode(pcode, srcs)
-    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
+    _, val0 = assigns[0][1]
+    _, val1 = assigns[1][1]
     # DATA[31:0] should preserve the value
-    self.assertEqual(assigns[0][1][1].simplify().arg, 0xAAAAAAAA)  # type: ignore[index]
+    self.assertEqual(val0.simplify().arg, 0xAAAAAAAA)
-    self.assertEqual(assigns[1][1][1].simplify().arg, 0xBBBBBBBB)  # type: ignore[index]
+    self.assertEqual(val1.simplify().arg, 0xBBBBBBBB)
 
 class TestConditionalParsing(unittest.TestCase):
   """Test conditional (if/elsif/else) pcode parsing."""
@@ -278,7 +273,7 @@ class TestConditionalParsing(unittest.TestCase):
     s0 = UOp.const(dtypes.uint32, 10)
     s1 = UOp.const(dtypes.uint32, 20)
     scc = UOp.const(dtypes.uint32, 1)
-    _vrs, assigns = parse_pcode(pcode, {'S0': s0, 'S1': s1, 'SCC': scc})
+    vars, assigns = parse_pcode(pcode, {'S0': s0, 'S1': s1, 'SCC': scc})
     self.assertEqual(len(assigns), 1)
     dest, val = assigns[0]
     self.assertTrue(dest.startswith('D0'))
@@ -291,7 +286,7 @@ class TestAllPcode(unittest.TestCase):
   def _make_srcs(self):
     """Create dummy source variables for pcode parsing."""
     u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
-    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
+    lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
     return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
             'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
             'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),
@@ -299,8 +294,7 @@ class TestAllPcode(unittest.TestCase):
             'ADDR': u32(), 'ADDR_BASE': u32(), 'TADDR': u32(), 'DATA': u32(), 'DATA0': u32(), 'DATA1': u32(), 'DATA2': u32(),
             'VDATA': u32(), 'VDATA0': u32(), 'VDATA1': u32(), 'VDATA2': u32(), 'VDATA3': u32(),
             'OPSEL': u32(), 'OPSEL_HI': u32(), 'NEG': u32(), 'NEG_HI': u32(), 'CLAMP': u32(),
-            'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'ROUND_TOWARD_ZERO': u32(),
+            'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'ROUND_TOWARD_ZERO': u32(), 'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(),
-            'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(),
             'MAX_FLOAT_F32': u32(0x7f7fffff), 'Unsigned': u32(1), 'clampedLOD': u32(),
             '_lds': lds, '_vmem': lds, '_active': UOp.const(dtypes.bool, True)}
 
@@ -312,9 +306,7 @@ class TestAllPcode(unittest.TestCase):
       try:
         parse_pcode(pcode, srcs)
         passed += 1
-      except RuntimeError as e:
+      except RuntimeError as e: skipped += 1; errors[str(e)].append(op.name)
-        skipped += 1
-        errors[str(e)].append(op.name)
       except Exception as e: self.fail(f"[{arch}] {op.name}: {e}\nPcode: {pcode[:200]}")
     total = len(pcode_dict)
     pct = 100 * passed / total
@@ -325,15 +317,15 @@ class TestAllPcode(unittest.TestCase):
     self.assertGreaterEqual(pct, min_pct, f"[{arch}] {pct:.1f}% < {min_pct}% threshold")
 
   def test_parse_all_cdna_pcode(self):
-    from tinygrad.runtime.autogen.amd.cdna.str_pcode import PCODE as CDNA_PCODE
+    from extra.assembly.amd.autogen.cdna.str_pcode import PCODE as CDNA_PCODE
     self._parse_all_pcode(CDNA_PCODE, "CDNA", min_pct=60)
 
   def test_parse_all_rdna3_pcode(self):
-    from tinygrad.runtime.autogen.amd.rdna3.str_pcode import PCODE as RDNA3_PCODE
+    from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE as RDNA3_PCODE
     self._parse_all_pcode(RDNA3_PCODE, "RDNA3", min_pct=90)
 
   def test_parse_all_rdna4_pcode(self):
-    from tinygrad.runtime.autogen.amd.rdna4.str_pcode import PCODE as RDNA4_PCODE
+    from extra.assembly.amd.autogen.rdna4.str_pcode import PCODE as RDNA4_PCODE
     self._parse_all_pcode(RDNA4_PCODE, "RDNA4", min_pct=65)
 
 if __name__ == "__main__":

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 extra.assembly.amd.autogen.rdna3.ins import *
-from tinygrad.renderer.amd.dsl import VCC_HI, EXEC_LO, NULL
+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 import detect_format
 
 
 class TestDS(unittest.TestCase):

extra/assembly/amd/test/test_handwritten.py

@@ -2,10 +2,9 @@
 # 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 extra.assembly.amd.autogen.rdna3.ins import *
-from tinygrad.renderer.amd.dsl import Inst
+from extra.assembly.amd.dsl import Inst
-from test.amd.test_roundtrip import compile_asm
+from extra.assembly.amd.test.test_roundtrip import compile_asm
-from test.amd.disasm import disasm
 
 class IntegrationTestBase(unittest.TestCase):
   inst: Inst
@@ -13,7 +12,7 @@ class IntegrationTestBase(unittest.TestCase):
   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)
@@ -133,11 +132,11 @@ class TestIntegrationCDNA(IntegrationTestBase):
   arch = "cdna"
 
   def test_mfma(self):
-    from tinygrad.runtime.autogen.amd.cdna.ins import v_mfma_f32_16x16x16_f16
+    from extra.assembly.amd.autogen.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
+    from extra.assembly.amd.autogen.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):
@@ -161,9 +160,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 extra.assembly.amd.disasm import disasm
-from tinygrad.renderer.amd import decode_inst, detect_format
+from extra.assembly.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.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):
+  if arch == "cdna" and not ('gfx942' in f or 'gfx950' in f or 'gfx90a' in f): tests = _filter_valid_asm(tests, arch)
-    tests = llvm_filter_valid_asm(tests, get_target(arch), get_mattr(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)
@@ -98,17 +127,15 @@ def _make_test(f: str, arch: str, test_type: str):
       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]
+      if arch == "rdna3": import extra.assembly.amd.autogen.rdna3.ins as ins
-      elif arch == "rdna4": import tinygrad.runtime.autogen.amd.rdna4.ins as ins  # type: ignore[no-redef]
+      elif arch == "rdna4": import extra.assembly.amd.autogen.rdna4.ins as ins
-      elif arch == "cdna": import tinygrad.runtime.autogen.amd.cdna.ins as ins  # type: ignore[no-redef]
+      elif arch == "cdna": import extra.assembly.amd.autogen.cdna.ins as ins
       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:
+          if decoded.to_bytes()[:len(data)] != data: skipped += 1; continue  # skip if binary roundtrip fails
-            skipped += 1
-            continue  # skip if binary roundtrip fails
           r = repr(decoded)
           try:
             decoded2 = eval(r, ns)  # noqa: S307
@@ -126,12 +153,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)

extra/assembly/amd/test/test_mockgpu_invalid.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """Test that invalid instructions raise exceptions through the mock GPU stack."""
-import unittest, subprocess, os, sys, time
+import unittest, subprocess, os, time
 
 class TestMockGPUInvalidInstruction(unittest.TestCase):
   def test_unsupported_instruction_raises(self):
@@ -43,7 +43,7 @@ dev.synchronize()
     env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
 
     st = time.perf_counter()
-    result = subprocess.run([sys.executable, "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
+    result = subprocess.run(["python", "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
     elapsed = time.perf_counter() - st
 
     self.assertNotEqual(result.returncode, 0, "should have raised")

extra/assembly/amd/test/test_pdf.py

@@ -1,15 +1,11 @@
 #!/usr/bin/env python3
 """Test PDF pseudocode extraction from generate.py."""
 import unittest
-from tinygrad.renderer.amd.generate import extract_pdf_text, extract_pcode, parse_xml, ARCHS, FIXES
+from extra.assembly.amd.generate import extract_pdf_text, extract_pcode, parse_xml, ARCHS, FIXES
 
 EXPECTED_PAGES = {"rdna3": 655, "rdna4": 711, "cdna": 610}
 
 class TestPcodePDF(unittest.TestCase):
-  pages: dict
-  enums: dict
-  pcode: dict
-
   @classmethod
   def setUpClass(cls):
     cls.pages = {arch: extract_pdf_text(cfg["pdf"]) for arch, cfg in ARCHS.items()}
@@ -37,8 +33,7 @@ class TestPcodePDF(unittest.TestCase):
       'tmp = MEM[ADDR].u64;\nsrc = DATA.u64;\nMEM[ADDR].u64 = src >= tmp ? src : tmp;\nRETURN_DATA.u64 = tmp')
     # GLOBAL_STORE_B128: should have 4 MEM stores (not truncated)
     self.assertEqual(pcode[('GLOBAL_STORE_B128', 29)],
-      'MEM[ADDR].b32 = VDATA[31 : 0];\nMEM[ADDR + 4U].b32 = VDATA[63 : 32];\n'
+      'MEM[ADDR].b32 = VDATA[31 : 0];\nMEM[ADDR + 4U].b32 = VDATA[63 : 32];\nMEM[ADDR + 8U].b32 = VDATA[95 : 64];\nMEM[ADDR + 12U].b32 = VDATA[127 : 96]')
-      'MEM[ADDR + 8U].b32 = VDATA[95 : 64];\nMEM[ADDR + 12U].b32 = VDATA[127 : 96]')
     # S_CMOVK_I32: should have full if/endif block
     self.assertEqual(pcode[('S_CMOVK_I32', 2)],
       "if SCC then\nD0.i32 = 32'I(signext(SIMM16.i16))\nendif")

extra/assembly/amd/test/test_rdna3_asm.py

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

extra/assembly/amd/test/test_roundtrip.py

@@ -1,10 +1,9 @@
 #!/usr/bin/env python3
 """Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
-import unittest, io, sys, re
+import unittest, io, sys, re, subprocess, os
-from tinygrad import Device
+from extra.assembly.amd.dsl import Inst
-from tinygrad.renderer.amd import detect_format
+from extra.assembly.amd import decode_inst, detect_format
-from test.amd.helpers import llvm_assemble, llvm_disasm, get_target, get_mattr
+from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump, get_target, get_mattr
-from test.amd.disasm import disasm
 
 def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
   """Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
@@ -31,20 +30,46 @@ def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
 
 def compile_asm(instr: str, arch: str = 'rdna3') -> bytes:
   """Compile a single instruction using LLVM."""
-  return llvm_assemble([instr], get_target(arch), get_mattr(arch))[0]
+  return compile_asm_batch([instr], arch)[0]
 
 def compile_asm_batch(instrs: list[str], arch: str = 'rdna3') -> list[bytes]:
-  """Compile multiple instructions with a single LLVM emission."""
+  """Compile multiple instructions with a single llvm-mc call."""
-  return llvm_assemble(instrs, get_target(arch), get_mattr(arch))
+  if not instrs: return []
+  result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={get_target(arch)}', f'-mattr={get_mattr(arch)}', '-show-encoding'],
+                          input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True)
+  if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
+  encodings = []
+  for line in result.stdout.split('\n'):
+    if 'encoding:' in line:
+      enc = line.split('encoding:')[1].strip()
+      if enc.startswith('[') and enc.endswith(']'):
+        encodings.append(bytes.fromhex(enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')))
+  if len(encodings) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(encodings)}")
+  return encodings
 
 def compile_and_disasm_batch(instrs: list[str], arch: str = 'rdna3') -> list[str]:
   """Compile instructions with LLVM and get LLVM's disassembly."""
+  import tempfile
   if not instrs: return []
   mcpu, mattr = get_target(arch), get_mattr(arch)
-  code = b''.join(llvm_assemble(instrs, mcpu, mattr))
+  src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n" + "\n".join(f"  {instr}" for instr in instrs) + "\n"
-  return llvm_disasm(code, mcpu, mattr)[:len(instrs)]
+  with tempfile.NamedTemporaryFile(suffix='.o', delete=False) as f:
+    obj_path = f.name
+  try:
+    result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', f'-mattr={mattr}', '-filetype=obj', '-o', obj_path],
+                            input=src, capture_output=True, text=True)
+    if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
+    result = subprocess.run([get_llvm_objdump(), '-d', f'--mcpu={mcpu}', obj_path], capture_output=True, text=True)
+    if result.returncode != 0: raise RuntimeError(f"llvm-objdump failed: {result.stderr.strip()}")
+    results: list[str] = []
+    for line in result.stdout.splitlines():
+      if '//' not in line: continue
+      instr = line.split('//')[0].strip()
+      if instr: results.append(instr)
+    return results[:len(instrs)]
+  finally:
+    os.unlink(obj_path)
 
-@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestTinygradKernelRoundtrip(unittest.TestCase):
   """Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
   arch = 'rdna3'
@@ -57,7 +82,7 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
     """
     arch = self.arch
 
-    from test.amd.test_compare_emulators import get_kernels_from_tinygrad
+    from extra.assembly.amd.test.test_compare_emulators import get_kernels_from_tinygrad
     from tinygrad.runtime.support.elf import elf_loader
     from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
     from tinygrad.helpers import AMD_LLVM
@@ -74,6 +99,11 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
       while offset < len(code):
         remaining = code[offset:]
         fmt = detect_format(remaining, arch)
+        if fmt is None:
+          decoded_instrs.append((ki, offset, None, None, None, False, "no format"))
+          offset += 4
+          continue
+
         base_size = fmt._size()
         if len(remaining) < base_size:
           break
@@ -83,7 +113,7 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
           size = decoded.size()  # actual size including literal
           orig_bytes = remaining[:size]
           reencoded = decoded.to_bytes()
-          our_disasm = disasm(decoded)
+          our_disasm = decoded.disasm()
           decode_ok = reencoded == orig_bytes
           decode_err: str | None = None if decode_ok else f"orig={orig_bytes.hex()} reenc={reencoded.hex()}"
           decoded_instrs.append((ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err))
@@ -147,20 +177,20 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
       if our_disasm is None:
         disasm_skipped += 1
       elif idx in disasm_llvm_map:
-        llvm_disasm_str = disasm_llvm_map[idx]
+        llvm_disasm = disasm_llvm_map[idx]
-        if our_disasm == llvm_disasm_str:
+        if our_disasm == llvm_disasm:
           disasm_passed += 1
         else:
           disasm_failed += 1
-          disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm_str}'")
+          disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
       else:
         disasm_skipped += 1
 
     print(f"[{arch}] decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
     print(f"[{arch}] asm via llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
     print(f"[{arch}] disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
-    self.assertEqual(decode_failed, 0, "Decode failures:\n" + "\n".join(decode_failures[:20]))
+    self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
-    self.assertEqual(asm_failed, 0, "Asm failures:\n" + "\n".join(asm_failures[:20]))
+    self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
     # Note: disasm string comparison is informational only - formatting differences between LLVM versions are expected
 
   # Basic unary ops

extra/assembly/amd/test/test_sqtt_examples.py

@@ -5,23 +5,21 @@ from pathlib import Path
 from tinygrad.helpers import DEBUG
 from tinygrad.runtime.autogen import rocprof
 from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.renderer.amd import decode_inst
+from extra.assembly.amd import decode_inst
-from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
+from extra.assembly.amd.autogen.rdna3.ins import SOPP
-from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
+from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
-from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
+from extra.assembly.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_L4, WAVEEND, INST, INST_L4, VALUINST, IMMEDIATE, IMMEDIATE_MASK,
-                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART,
+                                     ALUEXEC, VMEMEXEC, PACKET_TYPES_L3, PACKET_TYPES_L4, InstOp, InstOpL4, print_packets)
-                                     InstOp, InstOpRDNA4, print_packets, CDNA_WAVEEND, CDNA_INST)
+from extra.assembly.amd.test.helpers import TARGET_TO_ARCH
-from test.amd.helpers import TARGET_TO_ARCH
 
-import tinygrad
+EXAMPLES_DIR = Path(__file__).parent.parent.parent.parent / "sqtt/examples"
-EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
 # INST ops for non-traced SIMDs (excluded from instruction count)
 OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LDS_STORE_64, InstOp.OTHER_LDS_STORE_128,
                   InstOp.OTHER_FLAT_LOAD, InstOp.OTHER_FLAT_STORE, InstOp.OTHER_FLAT_STORE_64, InstOp.OTHER_FLAT_STORE_96,
                   InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
                   InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
                   InstOp.OTHER_GLOBAL_STORE_VADDR_128}
-OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_5, InstOpRDNA4.OTHER_LDS_1, InstOpRDNA4.OTHER_LDS_2}
+OTHER_SIMD_OPS_L4 = {InstOpL4.OTHER_VMEM, InstOpL4.UNK_60}
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # ROCPROF DECODER
@@ -34,18 +32,18 @@ def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
   assert text is not None, "no .text section found"
   text_off, text_size = text.header.sh_addr, text.header.sh_size
 
-  blob_iter, current_blob = iter(blobs), [None]  # type: ignore[var-annotated]
+  blob_iter, current_blob = iter(blobs), [None]
   occupancy_records: list[tuple[int, int, int, int, bool]] = []  # (wave_id, simd, cu, time, is_start)
   wave_insts: list[list[tuple[int, int]]] = []  # per-wave list of (time, stall)
 
   @rocprof.rocprof_trace_decoder_se_data_callback_t
-  def copy_cb(buf, buf_size, _):  # type: ignore[no-untyped-def]
+  def copy_cb(buf, buf_size, _):
     blob = next(blob_iter, None)
     if blob is None: return 0
-    current_blob[0] = (ctypes.c_ubyte * len(blob)).from_buffer_copy(blob)  # type: ignore[call-overload]
+    current_blob[0] = (ctypes.c_ubyte * len(blob)).from_buffer_copy(blob)
-    buf[0] = ctypes.cast(current_blob[0], ctypes.POINTER(ctypes.c_ubyte))  # type: ignore[arg-type]
+    buf[0] = ctypes.cast(current_blob[0], ctypes.POINTER(ctypes.c_ubyte))
-    buf_size[0] = len(current_blob[0])  # type: ignore[arg-type]
+    buf_size[0] = len(current_blob[0])
-    return len(current_blob[0])  # type: ignore[arg-type]
+    return len(current_blob[0])
 
   @rocprof.rocprof_trace_decoder_trace_callback_t
   def trace_cb(record_type, events_ptr, n, _):
@@ -89,14 +87,13 @@ def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
     try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
     except Exception as e: exc = e
   (t:=threading.Thread(target=worker, daemon=True)).start()
-  t.join(timeout=5)
+  t.join(timeout=1)
   if exc is not None: raise exc
   if t.is_alive(): raise RuntimeError("rocprof decoder timeout")
   return occupancy_records, wave_insts
 
 class SQTTExamplesTestBase(unittest.TestCase):
   target: str
-  examples: dict
 
   @classmethod
   def setUpClass(cls):
@@ -118,19 +115,17 @@ class SQTTExamplesTestBase(unittest.TestCase):
       for i, event in enumerate(events):
         with self.subTest(example=name, event=i):
           packets = list(decode(event.blob))
-          if DEBUG >= 2:
+          if DEBUG >= 2: print(f"\n=== {name} event {i} ==="); print_packets(packets)
-            print(f"\n=== {name} event {i} ===")
-            print_packets(packets)
           self.assertGreater(len(packets), 0, f"no packets decoded from {name} event {i}")
           self.assertIsInstance(packets[0], LAYOUT_HEADER, f"first packet should be LAYOUT_HEADER in {name}")
 
   def test_packet_types_valid(self):
-    all_classes = set(PACKET_TYPES_RDNA3.values()) | set(PACKET_TYPES_RDNA4.values()) | set(PACKET_TYPES_CDNA.values())
+    all_classes = set(PACKET_TYPES_L3.values()) | set(PACKET_TYPES_L4.values())
     for name, (events, *_) in self.examples.items():
       for i, event in enumerate(events):
         with self.subTest(example=name, event=i):
           for pkt in decode(event.blob):
-            # Use isinstance to handle layout-specific subclasses (e.g., WAVESTART_RDNA4)
+            # Use isinstance to handle layout-specific subclasses (e.g., WAVESTART_L4)
             self.assertTrue(any(isinstance(pkt, cls) for cls in all_classes), f"unknown packet type {type(pkt)} in {name}")
 
   def test_wave_lifecycle(self):
@@ -138,8 +133,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
       if "empty" in name: continue
       with self.subTest(example=name):
         all_packets = [p for e in events for p in decode(e.blob)]
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART))]), 0, f"no WAVESTART in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_L4))]), 0, f"no WAVESTART in {name}")
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, WAVEEND)]), 0, f"no WAVEEND in {name}")
 
   def test_time_monotonic(self):
     for name, (events, *_) in self.examples.items():
@@ -153,10 +148,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
       if "gemm" not in name: continue
       with self.subTest(example=name):
         all_packets = [p for e in events for p in decode(e.blob)]
-        inst_packets = [p for p in all_packets if isinstance(p, (INST, INST_RDNA4, CDNA_INST))]
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (INST, INST_L4))]), 0, f"no INST packets in {name}")
-        self.assertGreater(len(inst_packets), 0, f"no INST packets in {name}")
-        if isinstance(inst_packets[0], (INST, INST_RDNA4)):
-          self.assertGreater(len([p for p in inst_packets if p.op.name.startswith("JUMP")]), 0, f"no JUMP packets in {name}")
 
   expected: dict[str, list[int]] = {}  # override in subclasses
   def test_packet_counts(self):
@@ -184,8 +176,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
         for event in events:
           wave_starts: dict[tuple[int, int, int], int] = {}
           for p in decode(event.blob):
-            if isinstance(p, (WAVESTART, CDNA_WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
+            if isinstance(p, (WAVESTART, WAVESTART_L4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
-            elif isinstance(p, (WAVEEND, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
+            elif isinstance(p, WAVEEND) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
               our_waves.append((wave_starts[key], p._time))
         self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
 
@@ -201,7 +193,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
         for event in events:
           for p in decode(event.blob):
             if isinstance(p, INST) and p.op not in OTHER_SIMD_OPS: our_insts.append(p._time)
-            elif isinstance(p, INST_RDNA4) and p.op not in OTHER_SIMD_OPS_RDNA4: our_insts.append(p._time)
+            elif isinstance(p, INST_L4) and p.op not in OTHER_SIMD_OPS_L4: our_insts.append(p._time)
             elif isinstance(p, VALUINST): our_insts.append(p._time)
             elif isinstance(p, IMMEDIATE): our_insts.append(p._time)
             elif isinstance(p, IMMEDIATE_MASK):
@@ -211,22 +203,22 @@ class SQTTExamplesTestBase(unittest.TestCase):
 class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
   target = "gfx1100"
   expected = {
-    "profile_empty_run_0": [1880, 1867, 1920, 1971, 1998, 1904],
+    "profile_empty_run_0": [1803, 1908, 1928, 1979, 2006, 1912],
-    "profile_empty_run_1": [1880, 1867, 1920, 1971, 1998, 1904],
+    "profile_empty_run_1": [1803, 1908, 1928, 1979, 2006, 1912],
-    "profile_gemm_run_0": [3275, 3278, 2426, 2475, 2511, 2431],
+    "profile_gemm_run_0": [2531, 1844, 1864, 1915, 1942, 1848, 3074, 1919, 1939, 1990, 2017, 1923, 19026, 1919, 1939, 1990, 2017, 1929],
-    "profile_gemm_run_1": [3264, 3268, 2420, 2469, 2504, 2401],
+    "profile_gemm_run_1": [2554, 1844, 1864, 1915, 1942, 1848, 3084, 1919, 1939, 1990, 2017, 1923, 19010, 1919, 1939, 1990, 2017, 1923],
-    "profile_ops_run_0": [1944, 4903, 1984, 2035, 2062, 1968],
+    "profile_plus_run_0": [1900, 1908, 1928, 1979, 2006, 1912],
-    "profile_ops_run_1": [1944, 4918, 1984, 2035, 2062, 1968],
+    "profile_plus_run_1": [1856, 1908, 1928, 1979, 2006, 1912],
-    "profile_plus_run_0": [1938, 1932, 1978, 2029, 2056, 1962],
-    "profile_plus_run_1": [1891, 1874, 1931, 1982, 2009, 1915],
   }
 
 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
-
+# CDNA/MI300 (gfx950) uses a completely different 16-bit header packet format, not the nibble-based format.
+# See decode_tt_header_stream in ghidra/librocprof-trace-decoder.c - it reads 16-bit headers and uses
+# pkt_fmt = header & 0xf to look up packet_class (0x10=2bytes, 0x20=4bytes, 0x30=6bytes, 0x40=8bytes).
+# This is NOT implemented yet - the nibble decoder produces garbage for CDNA data.
+@unittest.skip("CDNA/MI300 uses 16-bit header format, not nibble-based - decoder not implemented")
 class TestSQTTExamplesCDNA(SQTTExamplesTestBase):
   target = "gfx950"
-  def test_rocprof_wave_times_match(self): self.skipTest("TODO: requires timestamp patching")
-  def test_rocprof_inst_times_match(self): self.skipTest("TODO: requires timestamp patching")
 
 if __name__ == "__main__":
   unittest.main()

extra/assembly/amd/test/test_sqtt_tables.py

@@ -0,0 +1,170 @@
+"""Tests comparing sqtt.py PACKET_TYPES_L3/L4 against AMD's rocprof-trace-decoder binary."""
+import unittest, struct, ctypes, pickle
+from pathlib import Path
+
+ROCPROF_LIB = Path("/usr/lib/librocprof-trace-decoder.so")
+EXAMPLES_DIR = Path(__file__).parent.parent.parent.parent / "sqtt/examples"
+
+def _find_segment(perms: str):
+  """Find a segment of the loaded library with given permissions (e.g. 'rw-p', 'r--p')."""
+  with open('/proc/self/maps', 'r') as f:
+    for line in f:
+      if 'librocprof-trace-decoder.so' in line and f' {perms} ' in line:
+        parts = line.split()
+        return int(parts[0].split('-')[0], 16), int(parts[2], 16)
+  return None, None
+
+def _read_array(file_offset: int, count: int):
+  """Read an array of uint8 at file_offset from the loaded library."""
+  base, seg_offset = _find_segment('rw-p')
+  if base is None: return None
+  return list((ctypes.c_uint8 * count).from_address(base + (file_offset - seg_offset)))
+
+def _load_lib():
+  if not ROCPROF_LIB.exists(): return False
+  ctypes.CDLL(str(ROCPROF_LIB))
+  return True
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# RDNA EXTRACTION (nibble-based format)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def extract_bit_tables():
+  """Extract bit budget tables. Returns (layout2, layout3, layout4) or None."""
+  if not _load_lib(): return None
+  return _read_array(0x2d220, 32), _read_array(0x2d280, 32), _read_array(0x2d2c0, 32)
+
+def extract_delta_fields():
+  """Extract delta bitfield tables. Returns (layout2, layout3, layout4) dicts mapping type_id -> (lo, hi)."""
+  if not _load_lib(): return None
+  ro_base, ro_offset = _find_segment('r--p')
+  if ro_base is None: return None
+
+  def read_table(file_offset, num_entries):
+    addr = ro_base + (file_offset - ro_offset)
+    data = bytes((ctypes.c_uint8 * (num_entries * 12)).from_address(addr))
+    return {type_id: (lo, hi) for j in range(0, len(data), 12)
+            for type_id, lo, hi in [struct.unpack('<III', data[j:j+12])] if type_id < 32}
+
+  return read_table(0x26800, 24), read_table(0x26dc0, 25), read_table(0x27300, 27)
+
+def extract_packet_encodings():
+  """Extract packet encodings. Returns (L2, L3, L4) dicts mapping type_id -> (mask, value)."""
+  if not _load_lib(): return None
+  rw_base, rw_offset = _find_segment('rw-p')
+  if rw_base is None: return None
+
+  # Read base encodings from registration vector at 0x2d340
+  vec_start = ctypes.c_void_p.from_address(rw_base + (0x2d340 - rw_offset)).value
+  vec_end = ctypes.c_void_p.from_address(rw_base + (0x2d348 - rw_offset)).value
+  base = {}
+  if vec_start and vec_end:
+    for i in range((vec_end - vec_start) // 32):
+      addr = vec_start + i * 32
+      type_id = ctypes.c_uint8.from_address(addr).value
+      pat_start = ctypes.c_void_p.from_address(addr + 8).value
+      pat_end = ctypes.c_void_p.from_address(addr + 16).value
+      if pat_start and pat_end and 0 < (n := pat_end - pat_start) <= 8:
+        pat = list((ctypes.c_uint8 * n).from_address(pat_start))
+        base[type_id] = (sum(1 << j for j in range(n)), sum(b << j for j, b in enumerate(pat)))
+
+  return {**base, 17: (0x7f, 0x51), 25: (0x7f, 0x31)}, base, {**base}  # L2 has overrides
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# CDNA EXTRACTION (16-bit header format)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def extract_cdna_packet_sizes():
+  """Extract CDNA pkt_fmt -> size mapping by running rocprof decoder to populate its hash table."""
+  from extra.assembly.amd.test.test_sqtt_examples import run_rocprof_decoder
+
+  if not (pkl_path := next((EXAMPLES_DIR / "gfx950").glob("*.pkl"), None)): return None
+  with open(pkl_path, "rb") as f: data = pickle.load(f)
+  sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
+  prg = next((e for e in data if type(e).__name__ == "ProfileProgramEvent"), None)
+  if not sqtt_events or not prg: return None
+
+  # Run decoder to trigger hash table initialization
+  run_rocprof_decoder([e.blob for e in sqtt_events], prg.lib, prg.base, "gfx950")
+
+  # Extract hash table: head at 0x2d4f0, nodes are 16 bytes (next[8], key[4], value[4])
+  rw_base, rw_offset = _find_segment('rw-p')
+  if not (head := ctypes.c_void_p.from_address(rw_base + (0x2d4f0 - rw_offset)).value if rw_base else None): return None
+
+  pkt_sizes, node, seen = {}, head, set()
+  while node and node not in seen and len(pkt_sizes) < 20:
+    seen.add(node)
+    key, val = ctypes.c_uint32.from_address(node + 8).value, ctypes.c_uint32.from_address(node + 12).value
+    if key < 16 and val in (0x10, 0x20, 0x30, 0x40): pkt_sizes[key] = {0x10: 2, 0x20: 4, 0x30: 6, 0x40: 8}[val]
+    node = ctypes.c_void_p.from_address(node).value
+  return pkt_sizes if len(pkt_sizes) == 16 else None
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# TESTS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class TestSQTTMatchesBinary(unittest.TestCase):
+  def test_bit_counts_match_layout3(self): self._test_bit_counts(3)
+  def test_bit_counts_match_layout4(self): self._test_bit_counts(4)
+  def test_encodings_match_layout3(self): self._test_encodings(3)
+  def test_encodings_match_layout4(self): self._test_encodings(4)
+  def test_delta_fields_match_layout3(self): self._test_delta_fields(3)
+  def test_delta_fields_match_layout4(self): self._test_delta_fields(4)
+
+  def test_cdna_packet_sizes(self):
+    """Extract and verify CDNA pkt_fmt -> size mapping from rocprof's hash table."""
+    if not (EXAMPLES_DIR / "gfx950").exists(): self.skipTest("no CDNA examples")
+    pkt_sizes = extract_cdna_packet_sizes()
+    self.assertIsNotNone(pkt_sizes, "failed to extract CDNA packet sizes")
+    from extra.assembly.amd.sqtt_cdna import CDNA_PKT_SIZES
+    for pkt_fmt, size in CDNA_PKT_SIZES.items():
+      with self.subTest(pkt_fmt=pkt_fmt): self.assertEqual(pkt_sizes.get(pkt_fmt), size)
+
+  def _test_bit_counts(self, layout: int):
+    if not (tables := extract_bit_tables()): self.skipTest("rocprof-trace-decoder not installed")
+    from extra.assembly.amd.sqtt import PACKET_TYPES_L3, PACKET_TYPES_L4
+    for type_id, pkt_cls in {3: PACKET_TYPES_L3, 4: PACKET_TYPES_L4}[layout].items():
+      with self.subTest(packet=pkt_cls.__name__):
+        self.assertEqual(pkt_cls._size_nibbles * 4, tables[layout - 2][type_id])
+
+  def _test_encodings(self, layout: int):
+    if not (encodings := extract_packet_encodings()): self.skipTest("rocprof-trace-decoder not installed")
+    from extra.assembly.amd.sqtt import PACKET_TYPES_L3, PACKET_TYPES_L4
+    for type_id, pkt_cls in {3: PACKET_TYPES_L3, 4: PACKET_TYPES_L4}[layout].items():
+      with self.subTest(packet=pkt_cls.__name__):
+        self.assertEqual((pkt_cls.encoding.mask, pkt_cls.encoding.default), encodings[layout - 2][type_id])
+
+  def _test_delta_fields(self, layout: int):
+    if not (deltas := extract_delta_fields()): self.skipTest("rocprof-trace-decoder not installed")
+    from extra.assembly.amd.sqtt import PACKET_TYPES_L3, PACKET_TYPES_L4
+    for type_id, pkt_cls in {3: PACKET_TYPES_L3, 4: PACKET_TYPES_L4}[layout].items():
+      if type_id not in deltas[layout - 2]: continue
+      delta = getattr(pkt_cls, 'delta', None)
+      actual = (0, 0) if delta is None else (delta.lo, delta.hi + 1)
+      with self.subTest(packet=pkt_cls.__name__): self.assertEqual(actual, deltas[layout - 2][type_id])
+
+if __name__ == "__main__":
+  tables = extract_bit_tables()
+  encodings = extract_packet_encodings()
+  deltas = extract_delta_fields()
+
+  TYPE_NAMES = {1: 'VALUINST', 2: 'VMEMEXEC', 3: 'ALUEXEC', 4: 'IMMEDIATE', 5: 'IMMEDIATE_MASK', 6: 'WAVERDY',
+    7: 'TS_DELTA_S8_W3', 8: 'WAVEEND', 9: 'WAVESTART', 10: 'TS_DELTA_S5_W2', 11: 'WAVEALLOC', 12: 'TS_DELTA_S5_W3',
+    13: 'PERF', 14: 'UTILCTR', 15: 'TS_DELTA_SHORT', 16: 'NOP', 17: 'TS_WAVE_STATE', 18: 'EVENT', 19: 'EVENT_BIG',
+    20: 'REG', 21: 'SNAPSHOT', 22: 'TS_DELTA_OR_MARK', 23: 'LAYOUT_HEADER', 24: 'INST', 25: 'UNK_25'}
+
+  print("L2:", tables[0], "\nL3:", tables[1], "\nL4:", tables[2])
+  if encodings and tables:
+    print(f"\n{'TypeID':>6} {'Name':>18} {'L2 enc':>12} {'L3 enc':>12} {'L4 enc':>12} {'L2':>4} {'L3':>4} {'L4':>4} {'L2 delta':>12} {'L3 delta':>12} {'L4 delta':>12}")
+    print("-" * 140)
+    for type_id in sorted(set(encodings[0]) | set(encodings[1]) | set(encodings[2])):
+      name = TYPE_NAMES.get(type_id, f'UNK_{type_id}')
+      bits = [tables[i][type_id] if type_id < len(tables[i]) else 0 for i in range(3)]
+      enc_strs = [f"0x{encodings[i][type_id][0]:02x}/0x{encodings[i][type_id][1]:02x}" if type_id in encodings[i] else "-" for i in range(3)]
+      delta_strs = [f"[{d[1]-1}:{d[0]}]" if (d := deltas[i].get(type_id, (0, 0)))[1] > d[0] else "-" for i in range(3)]
+      print(f"{type_id:6d} {name:>18} {enc_strs[0]:>12} {enc_strs[1]:>12} {enc_strs[2]:>12} {bits[0]:4d} {bits[1]:4d} {bits[2]:4d} {delta_strs[0]:>12} {delta_strs[1]:>12} {delta_strs[2]:>12}")
+
+  cdna = extract_cdna_packet_sizes()
+  if cdna: print(f"\nCDNA packet sizes: {cdna}")
+
+  unittest.main()

extra/bench_log.py

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

extra/fp8/fp8_linear.py

@@ -28,7 +28,7 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
   return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))
 
 def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
-  _, input_uop, weight_uop = kernel.src[1:]
+  _, input_uop, weight_uop = kernel.src
   input_tensor = Tensor(input_uop, device=input_uop.device)
   grad_tensor = Tensor(gradient, device=gradient.device)
   weight_tensor = Tensor(weight_uop, device=weight_uop.device)

extra/gemm/amd_asm_matmul.py

@@ -1,5 +1,5 @@
 # RDNA3 128x128 tiled GEMM kernel - DSL version
-# Computes C = A @ B for NxN float32 matrices using 128x128 tiles
+# Computes C = A @ B for 4096x4096 float32 matrices using 128x128 tiles
 #
 # Architecture: RDNA3 (gfx1100)
 # Tile size: 128x128 (each workgroup computes one tile of C)
@@ -9,18 +9,19 @@
 # Accumulators: 128 vgprs (v[2-129])
 
 import numpy as np
+from pathlib import Path
 from tinygrad import Tensor, Device, Context, GlobalCounters
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.helpers import getenv, colored
-from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.engine.realize import Estimates
-from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL
+from extra.assembly.amd.dsl import s, v, VCC_LO, NULL
-from tinygrad.runtime.autogen.amd.rdna3.ins import *
+from extra.assembly.amd.autogen.rdna3.ins import *
 
 # =============================================================================
 # Kernel constants
 # =============================================================================
 LDS_SIZE = 8320       # Local data share size in bytes
+MATRIX_DIM = 4096     # Matrix dimension N (assumes square NxN matrices)
 LDS_A_STRIDE = 0x210  # LDS stride for A tile (528 bytes)
 LDS_B_STRIDE = 0x200  # LDS stride for B tile (512 bytes)
 LDS_BASE_OFFSET = 0x1080  # Base LDS offset for tiles
@@ -50,18 +51,18 @@ V_B_TILE_REGS = [132, 136, 140, 144, 148, 152, 156, 160]  # B tile: banks 0,0,0,
 # Named register assignments (SGPRs)
 # =============================================================================
 S_OUT_PTR = (0, 1)        # output C matrix base pointer
-S_WORKGROUP_X = 2         # workgroup_id_x (system SGPR, follows user SGPRs)
+S_TILE_X = 2              # workgroup_x << 7
-S_WORKGROUP_Y = 3         # workgroup_id_y (system SGPR)
+S_TILE_Y = 3              # workgroup_y << 7
 S_DIM_N = 4               # matrix dimension N
 S_LOOP_BOUND = 7          # K-8 (loop termination bound)
 S_LOOP_CTR = 12           # loop counter (increments by 8)
 S_PREFETCH_FLAG = 13      # prefetch condition flag / row stride in epilogue
-S_TILE_X = 14             # workgroup_x << 7
+S_WORKGROUP_X = 14        # workgroup_id_x
-S_TILE_Y = 15             # workgroup_y << 7
+S_WORKGROUP_Y = 15        # workgroup_id_y
 # Kernarg load destinations
 S_KERNARG_A = (20, 21)    # A pointer from kernarg
 S_KERNARG_B = (22, 23)    # B pointer from kernarg
-# Prefetch base pointers (8 pairs each, B: N*4 bytes apart, A: N*64 bytes apart)
+# Prefetch base pointers (8 pairs each, 16KB/256KB apart)
 S_PREFETCH_B = 24         # s[24:39] - 8 B tile pointers
 S_PREFETCH_A = 40         # s[40:55] - 8 A tile pointers
 
@@ -182,23 +183,54 @@ class Kernel:
     waitcnt = (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
     self.emit(s_waitcnt(simm16=waitcnt))
 
-  def finalize(self):
+  def to_asm(self):
-    """Patch branch offsets and return the finalized instruction list."""
+    # Patch branch offsets: simm16 = (target_pos - branch_end_pos) / 4
     for inst in self.instructions:
       if inst._target is None: continue
       offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
       if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
       inst.simm16 = offset_dwords
-    return self.instructions
+
+    # TODO: replace this with direct ELF
+    body = ['\t' + inst.disasm() for inst in self.instructions]
+
+    # limit wave occupancy by using more LDS
+    lds_size = max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536))
+
+    # HSA kernel descriptor attributes (zeros included for compatibility)
+    hsa = [
+      ('group_segment_fixed_size', lds_size), ('private_segment_fixed_size', 0), ('kernarg_size', 36),
+      ('user_sgpr_count', 14), ('user_sgpr_dispatch_ptr', 0), ('user_sgpr_queue_ptr', 0),
+      ('user_sgpr_kernarg_segment_ptr', 1), ('user_sgpr_dispatch_id', 0), ('user_sgpr_private_segment_size', 0),
+      ('wavefront_size32', 1), ('uses_dynamic_stack', 0), ('enable_private_segment', 0),
+      ('system_sgpr_workgroup_id_x', 1), ('system_sgpr_workgroup_id_y', 1), ('system_sgpr_workgroup_id_z', 0),
+      ('system_sgpr_workgroup_info', 0), ('system_vgpr_workitem_id', 0), ('next_free_vgpr', 179),
+      ('next_free_sgpr', 16), ('float_round_mode_32', 0), ('float_round_mode_16_64', 0),
+      ('float_denorm_mode_32', 3), ('float_denorm_mode_16_64', 3), ('dx10_clamp', 1), ('ieee_mode', 1),
+      ('fp16_overflow', 0), ('workgroup_processor_mode', 0), ('memory_ordered', 1), ('forward_progress', 0),
+      ('shared_vgpr_count', 0)]
+
+    return '\n'.join([
+      '\t.text', f'\t.amdgcn_target "amdgcn-amd-amdhsa--{self.arch}"',
+      '\t.protected\tkernel', '\t.globl\tkernel', '\t.p2align\t8', '\t.type\tkernel,@function', 'kernel:',
+      *body,
+      '\t.section\t.rodata,"a",@progbits', '\t.p2align\t6, 0x0', '\t.amdhsa_kernel kernel',
+      *[f'\t\t.amdhsa_{k} {v}' for k, v in hsa],
+      '\t.end_amdhsa_kernel', '\t.text', '.Lfunc_end0:', '\t.size\tkernel, .Lfunc_end0-kernel',
+      '\t.amdgpu_metadata', '---', 'amdhsa.kernels:', '  - .args:',
+      *[f'      - .address_space: global\n        .offset: {i*8}\n        .size: 8\n        .value_kind: global_buffer' for i in range(3)],
+      f'    .group_segment_fixed_size: {lds_size}', '    .kernarg_segment_align: 8',
+      '    .kernarg_segment_size: 24', '    .max_flat_workgroup_size: 128', '    .name: kernel',
+      '    .private_segment_fixed_size: 0', '    .sgpr_count: 60', '    .symbol: kernel.kd',
+      '    .vgpr_count: 179', '    .wavefront_size: 32', f'amdhsa.target: amdgcn-amd-amdhsa--{self.arch}',
+      'amdhsa.version:', '  - 1', '  - 2', '...', '\t.end_amdgpu_metadata'])
 
 
 # =============================================================================
 # Kernel builder
 # =============================================================================
 
-def build_kernel(N, arch='gfx1100'):
+def build_kernel(arch='gfx1100'):
-  assert N % 128 == 0, f"N must be a multiple of 128 (tile size), got {N}"
-  assert N >= 256, f"N must be >= 256 (prefetch pipeline requires at least 2 K-blocks), got {N}"
   k = Kernel(arch)
 
   # ===========================================================================
@@ -206,7 +238,7 @@ def build_kernel(N, arch='gfx1100'):
   # ===========================================================================
   k.emit(s_load_b128(sdata=s[S_KERNARG_A[0]:S_KERNARG_B[1]], sbase=s[0:1], offset=0x0, soffset=NULL))
   k.emit(s_load_b64(sdata=s[S_OUT_PTR[0]:S_OUT_PTR[1]], sbase=s[0:1], offset=0x10, soffset=NULL))
-  k.emit(s_mov_b32(s[S_DIM_N], N))
+  k.emit(s_mov_b32(s[S_DIM_N], MATRIX_DIM))
   k.emit(s_mov_b32(s[S_LOOP_CTR], 0))  # used by LDS swizzle, always 0 for valid workgroups
   k.emit(s_lshl_b32(s[S_TILE_X], s[S_WORKGROUP_X], 7))
   k.emit(s_lshl_b32(s[S_TILE_Y], s[S_WORKGROUP_Y], 7))
@@ -221,20 +253,19 @@ def build_kernel(N, arch='gfx1100'):
 
   # Compute 8 A and B matrix tile base pointers for prefetch
   k.emit(s_mov_b64(s[S_PREFETCH_B:S_PREFETCH_B+1], s[S_KERNARG_B[0]:S_KERNARG_B[1]]))  # B[0]: no offset
-  for i in range(1, 8):  # B: each pointer 1 row of B apart (N*4 bytes)
+  for i in range(1, 8):  # B: 16KB apart
-    k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_KERNARG_B[0]], i * N * 4))
+    k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_KERNARG_B[0]], i * 0x4000))
     k.emit(s_addc_u32(s[S_PREFETCH_B+i*2+1], s[S_KERNARG_B[1]], 0))
   k.emit(s_mov_b64(s[S_PREFETCH_A:S_PREFETCH_A+1], s[S_KERNARG_A[0]:S_KERNARG_A[1]]))  # A[0]: no offset
-  for i in range(1, 8):  # A: each pointer 16 rows of A apart (16*N*4 bytes)
+  for i in range(1, 8):  # A: 256KB apart
-    k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_KERNARG_A[0]], i * N * 64))
+    k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_KERNARG_A[0]], i * 0x40000))
     k.emit(s_addc_u32(s[S_PREFETCH_A+i*2+1], s[S_KERNARG_A[1]], 0))
 
-  # Global prefetch addresses: B = (tile_x + lane_id) * 4, A = (tile_y*N + (lane_id/8)*N + lane_id%8) * 4
+  # Global prefetch addresses: B = (tile_x + lane_id) * 4, A = ((tile_y << 12) + (lane_id/8)*4K + lane_id%8) * 4
   k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], s[S_TILE_X], v[V_LANE_ID]))
   k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_B_ADDR], 2, v[V_GLOBAL_B_ADDR]))
-  k.emit(s_mul_i32(s[19], s[S_TILE_Y], N))
+  k.emit(s_lshl_b32(s[19], s[S_TILE_Y], 12))
-  k.emit(v_mul_lo_u32(v[V_GLOBAL_A_ADDR], v[4], N))                          # (lane_id/8)*N
+  k.emit(v_lshl_add_u32(v[V_GLOBAL_A_ADDR], v[4], 12, v[V_LANE_ID_MOD8]))  # (lane_id/8)*4K + lane_id%8
-  k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], v[V_LANE_ID_MOD8], v[V_GLOBAL_A_ADDR]))  # + lane_id%8
   k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], s[19], v[V_GLOBAL_A_ADDR]))
   k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_A_ADDR], 2, v[V_GLOBAL_A_ADDR]))
 
@@ -305,13 +336,13 @@ def build_kernel(N, arch='gfx1100'):
 
   if not NO_GLOBAL:
     # Advance prefetch pointers (VGPR)
-    #k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], N * 32, v[V_GLOBAL_B_ADDR]))
+    #k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], 0x20000, v[V_GLOBAL_B_ADDR]))
     #k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], 0x20, v[V_GLOBAL_A_ADDR]))
 
-    # Advance prefetch pointers (64-bit adds): B advances 8 rows (8*N*4 bytes), A advances 8 cols (8*4 bytes)
+    # Advance prefetch pointers (64-bit adds)
     k.emit(s_clause(simm16=31))
     for i in range(8):
-      k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_PREFETCH_B+i*2], N * 32))
+      k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_PREFETCH_B+i*2], 0x20000))
       k.emit(s_addc_u32(s[S_PREFETCH_B+i*2+1], s[S_PREFETCH_B+i*2+1], 0))
     for i in range(8):
       k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_PREFETCH_A+i*2], 0x20))
@@ -428,7 +459,7 @@ def build_kernel(N, arch='gfx1100'):
   k.emit(s_sendmsg(simm16=3))  # DEALLOC_VGPRS
   k.emit(s_endpgm())
 
-  return k.finalize()
+  return k.to_asm()
 
 # =============================================================================
 # Test harness
@@ -442,7 +473,16 @@ def test_matmul():
   dev = Device[Device.DEFAULT]
   print(f"Device arch: {dev.renderer.arch}")
 
-  insts = build_kernel(N, dev.renderer.arch)
+  if getenv("STOCK", 0):
+    # Load the stock kernel from amd_seb/kernel8_batched_gmem.s
+    stock_path = Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s"
+    asm = stock_path.read_text()
+    print(f"Loaded stock kernel from {stock_path}")
+  else:
+    asm = build_kernel(dev.renderer.arch)
+
+  binary = dev.compiler.compile(asm)
+  print(f"Compiled! Binary size: {len(binary)} bytes")
 
   rng = np.random.default_rng(42)
   a = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
@@ -457,10 +497,10 @@ def test_matmul():
   def asm_kernel(A:UOp, B:UOp, C:UOp) -> UOp:
     gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
     lidxs = [UOp.special(n, f"lidx{i}") for i,n in enumerate(local)]
-    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536)), addrspace=AddrSpace.LOCAL), (), 'lds')
+    sink = UOp.sink(A.base, B.base, C.base, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
-    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
+                                                                           estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
-                                                                                  estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=asm),
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+                                 UOp(Ops.BINARY, arg=binary)))
   c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
   ei = c.schedule()[0].lower()
 
@@ -474,23 +514,33 @@ def test_matmul():
     with Context(DEBUG=2): tc = (a @ b).realize()
     with Context(DEBUG=0): err = (c - tc).square().mean().item()
     print(f"mean squared error {err}")
-    if err != err or err > 1e-06:
+    if err != err or err > 1e-06: raise RuntimeError("matmul is wrong!")
-      c_np, tc_np = c.numpy(), tc.numpy()
+
-      for bi in range(N // 128):
+def run_sqtt():
-        for bj in range(N // 128):
+  """Run with SQTT profiling and write trace files."""
-          blk_c = c_np[bi*128:(bi+1)*128, bj*128:(bj+1)*128]
+  import subprocess, os
-          blk_ref = tc_np[bi*128:(bi+1)*128, bj*128:(bj+1)*128]
+
-          blk_diff = blk_c - blk_ref
+  # Run test_matmul in a subprocess with SQTT enabled from the start (no verify)
-          zero_rows = [i for i in range(128) if np.all(np.abs(blk_c[i,:]) < 1e-10)]
+  env = {**os.environ, "AMD": "1", "SQTT": "1", "CNT": "1", "PROFILE": "1", "PYTHONPATH": ".", "VERIFY": "0"}
-          nz_rows = [i for i in range(128) if i not in zero_rows]
+  result = subprocess.run(
-          nz_mse = float(np.mean(blk_diff[nz_rows,:]**2)) if nz_rows else 0
+    ["python", "-c", "from extra.gemm.amd_asm_matmul import test_matmul; test_matmul()"],
-          print(f"Block ({bi},{bj}): zero_rows={zero_rows}, nz_rows_mse={nz_mse:.2e}")
+    capture_output=True, text=True, env=env, timeout=120
-          # show first few non-zero row comparisons
+  )
-          if nz_rows and nz_mse > 1e-6:
+  print(result.stdout)
-            for r in nz_rows[:3]:
+
-              print(f"  row {r} asm[0:8]:  {blk_c[r,:8]}")
+  # Run roc.py to extract trace data
-              print(f"  row {r} ref[0:8]:  {blk_ref[r,:8]}")
+  result = subprocess.run(
-      raise RuntimeError("matmul is wrong!")
+    ["python", "extra/sqtt/roc.py", "--profile", "/tmp/profile.pkl.tiny", "--kernel", "kernel"],
+    capture_output=True, text=True, env={**os.environ, "DEBUG": "5"}, timeout=60
+  )
+  output = result.stdout + result.stderr
+
+  # Write full output to trace file
+  with open("/tmp/sqtt_trace.txt", "w") as f:
+    f.write(output)
+  print(f"Wrote {len(output)} bytes to /tmp/sqtt_trace.txt")
 
 if __name__ == "__main__":
-  test_matmul()
+  if getenv("ASM", 0): print(build_kernel(Device[Device.DEFAULT].arch))
+  elif getenv("SQTT", 0): run_sqtt()
+  else: test_matmul()

extra/gemm/amd_copy_matmul.py

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

extra/gemm/amd_uop_matmul.py

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

extra/gemm/asm/cdna/gemm.py

@@ -0,0 +1,95 @@
+import atexit, functools
+from tinygrad.runtime.support.compiler_amd import HIPCompiler
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.renderer import Estimates
+from tinygrad.helpers import getenv, all_same, dedup
+from extra.gemm.asm.cdna.asm import build_kernel, GEMM_ARGS
+
+# ** CDNA4 assembly gemm
+
+WORKGROUP_SIZE = 256
+
+def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str, arch:str, wg:int) -> UOp:
+  batch, M, K = A.shape
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
+  gidx = UOp.special(wg, "gidx0")
+  k = build_kernel(batch, M, N, K, A.dtype.base)
+  sink = UOp.sink(C.base, A.base, B.base, lidx, gidx,
+                  arg=KernelInfo(name=k.name, estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
+  binary = HIPCompiler(arch).compile(k.to_asm())
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=k.to_text()), UOp(Ops.BINARY, arg=binary)))
+
+counters = {"used":0, "todos":[]}
+def todo(msg:str) -> bool: counters["todos"].append(msg); return False
+atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))
+
+def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
+  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
+  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
+  # only sharding on the batch is tested, others might work too
+  if isinstance(a.device, tuple) and not (a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None):
+    return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
+  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
+  N = b.shape[1]
+  if isinstance(a.device, tuple): batch //= len(a.device)
+  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
+  if (key:=(M, N, K)) not in GEMM_ARGS: return todo(f"GEMM shape not supported {key}")
+  return True
+
+# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
+# note: this can be removed after we have GEMM on mixins
+
+def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
+  M, K = A.shape[0]*A.shape[1], A.shape[2]
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  m = UOp.range(M, 1, AxisType.LOOP)
+  n = UOp.range(N, 2, AxisType.LOOP)
+  k = UOp.range(K, 0, AxisType.REDUCE)
+  mul = (A.index((m*UOp.const(dtypes.index, K)+k))*B.index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
+  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
+  store = C.index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
+  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
+
+# ** backward gemm, might use the asm gemm
+
+def custom_gemm_bw(gradient:UOp, kernel:UOp):
+  out, a, b = kernel.src
+  assert all_same([gradient.device, a.device, b.device, out.device])
+  a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
+  grad_a = (g_t @ b_t.T).uop
+  a_T = a_t.transpose(-2, -1)
+  a_T = a_T.reshape(*a_T.shape[:-1], 1, a_T.shape[-1])
+  g_r = g_t.reshape(*g_t.shape[:-2], 1, *g_t.shape[-2:]).transpose(-1, -2)
+  grad_b = (a_T * g_r).sum((-1, 0)).uop
+  return (None, grad_a, grad_b)
+
+# ** main gemm function
+
+def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
+  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
+  counters["used"] += 1
+  squeeze = a.ndim == 2
+  if squeeze: a = a.unsqueeze(0)
+
+  batch, M, K = a.shape
+  N = b.shape[1]
+  is_multi = isinstance(a.device, tuple)
+
+  if is_multi:
+    out = Tensor(Tensor.empty(batch//len(a.device), M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+  else:
+    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
+
+  dname = a.device[0] if is_multi else a.device
+  arch = getattr(Device[dname].renderer, "arch", None)
+  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
+    numWG = GEMM_ARGS[(M, N, K)][0]
+    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=numWG, arch=arch), grad_fxn=custom_gemm_bw)[0]
+  else:
+    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
+  return out.squeeze(0) if squeeze else out

extra/gemm/asm/cdna/test_asm_gemm.py

@@ -0,0 +1,46 @@
+import unittest
+from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad.helpers import getenv
+from extra.gemm.asm.cdna.gemm import asm_gemm
+
+def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, multi=False) -> None:
+  Tensor.manual_seed(0)
+  a_rand = Tensor.randn((batch, M, K), dtype=dtypes.float).sub(0.5).cast(dtype)
+  b_rand = Tensor.randn((K, N), dtype=dtypes.float).sub(0.5).cast(dtype)
+  with Context(DEBUG=0):
+    Tensor.realize(a_rand, b_rand)
+
+  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8)) if multi else None
+
+  a, b = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
+  if multi: a, b = a.shard(devs, axis=0), b.shard(devs, axis=None)
+  tst = asm_gemm(a, b)
+  tst.sum().backward()
+  Tensor.realize(tst, a.grad, b.grad)
+
+  a_ref, b_ref = Tensor(a_rand.numpy(), requires_grad=True).cast(dtype), Tensor(b_rand.numpy(), requires_grad=True).cast(dtype)
+  if multi: a_ref, b_ref = a_ref.shard(devs, axis=0), b_ref.shard(devs, axis=None)
+  with Context(ASM_GEMM=0): ref = a_ref @ b_ref
+  ref.sum().backward()
+  Tensor.realize(ref, a_ref.grad, b_ref.grad)
+
+  with Context(DEBUG=0):
+    assert (tst - ref).square().max().float().item() < 1e-6, "forward mismatch"
+    assert (a.grad - a_ref.grad).square().max().float().item() < 1e-3, "grad_a mismatch"
+    assert (b.grad - b_ref.grad).square().max().float().item() < 1e-3, "grad_b mismatch"
+
+class TestGemm(unittest.TestCase):
+  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=dtypes.half)
+
+  def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, multi=True)
+  def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, multi=True)
+  def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, multi=True)
+  def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, multi=True)
+  def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, multi=True)
+  def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, multi=True)
+  def test_gemm_unsupported(self):
+    with self.assertRaisesRegex(AssertionError, "shape not supported"):
+      verify_asm_gemm(8, 8192, 1024, 4096, multi=True)
+
+if __name__ == "__main__":
+  unittest.main()

extra/gemm/asm/rdna3/template.s

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

extra/gemm/asm/rdna3/test.py

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

extra/gemm/asm/unpack_kd.py

@@ -0,0 +1,179 @@
+# unpack the complete kernel descriptor of an amdgpu ELF
+# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
+import struct, pathlib, sys
+from tinygrad.runtime.support.elf import elf_loader
+
+def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
+def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
+
+with open(sys.argv[1], "rb") as f:
+  lib = f.read()
+
+image, sections, relocs = elf_loader(lib)
+rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"))
+
+# rodata is exactly 64 bytes
+kd = image[rodata_entry:rodata_entry+64]
+desc = int.from_bytes(kd, byteorder="little")
+
+group_segment_fixed_size = bits(desc, 0, 31)
+private_segment_fixed_size = bits(desc, 32, 63)
+kernarg_size = bits(desc, 64, 95)
+reserved_127_96 = bits(desc, 96, 127)
+assert reserved_127_96 == 0
+
+print("GROUP_SEGMENT_FIXED_SIZE:", group_segment_fixed_size)
+print("PRIVATE_SEGMENT_FIXED_SIZE:", private_segment_fixed_size)
+print("KERNARG_SIZE:", kernarg_size)
+print("RESERVED 127:96:", reserved_127_96)
+
+entry_off = bits(desc, 128, 191)
+
+# sign-extend manually if needed
+if entry_off & (1 << 63):
+  entry_off -= 1 << 64
+
+print("KERNEL_CODE_ENTRY_BYTE_OFFSET:", entry_off)
+
+kd_addr = 0x1840
+entry_addr = kd_addr + entry_off
+
+print("Computed entry address: 0x%016x" % entry_addr)
+print("256B aligned:", entry_addr % 256 == 0)
+
+pgm_rsrc3 = bits(desc, 352, 383)
+pgm_rsrc1 = bits(desc, 384, 415)
+pgm_rsrc2 = bits(desc, 416, 447)
+
+print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
+print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
+print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
+
+# rsrc 3 (gfx950)
+
+accum_offset_raw = bits(pgm_rsrc3, 0, 5)
+assert_zero(pgm_rsrc3, 6, 15)
+tg_split = bits(pgm_rsrc3, 16, 16)
+accum_offset_vgprs = (accum_offset_raw + 1) * 4
+print("RSRC3.ACCUM_OFFSET (AccVGPR index):", accum_offset_vgprs)
+print("RSRC3.TG_SPLIT:", tg_split)
+
+# rsrc 1
+
+vgpr_gran = bits(pgm_rsrc1, 0, 5)
+sgpr_gran = bits(pgm_rsrc1, 6, 9)
+assert_zero(pgm_rsrc1, 27, 28)
+
+# NOTE: this is vgprs + agprs
+vgprs_used = (vgpr_gran + 1) * 8
+assert 0 <= vgprs_used <= 512
+
+k = sgpr_gran // 2
+sgprs_used = (k + 1) * 16
+
+print("RSRC1.VGPRS:", vgprs_used)
+print("RSRC1.SGPRS:", sgprs_used)
+
+assert_zero(pgm_rsrc1, 10, 11)
+
+float_round_mode_32 = bits(pgm_rsrc1, 12, 13)
+float_round_mode_16_64 = bits(pgm_rsrc1, 15, 14)
+float_denorm_mode_32 = bits(pgm_rsrc1, 16, 17)
+float_denorm_mode_16_64 = bits(pgm_rsrc1, 18, 19)
+
+priv = bits(pgm_rsrc1, 20, 20)
+assert priv == 0
+enable_dx10_clamp_wg_rr_en = bits(pgm_rsrc1, 21, 21)
+debug_mode = bits(pgm_rsrc1, 22, 22)
+enable_ieee_mode = bits(pgm_rsrc1, 23, 23)
+bulky = bits(pgm_rsrc1, 24, 24)
+assert bulky == 0
+cdbg_user = bits(pgm_rsrc1, 25, 25)
+assert cdbg_user == 0
+fp16_ovfl = bits(pgm_rsrc1, 26, 26)
+assert_zero(pgm_rsrc1, 27, 28)  # reserved
+assert_zero(pgm_rsrc1, 29, 29)  # WGP_MODE (reserved on gfx9)
+assert_zero(pgm_rsrc1, 30, 30)  # MEM_ORDERED (reserved on gfx9)
+assert_zero(pgm_rsrc1, 31, 31)  # FWD_PROGRESS (reserved on gfx9)
+
+# rsrc 2
+
+enable_private_segment = bits(pgm_rsrc2, 0, 0)  # SCRATCH_EN
+user_sgpr_count = bits(pgm_rsrc2, 1, 5)         # USER_SGPR
+enable_trap_handler = bits(pgm_rsrc2, 6, 6)     # TRAP_PRESENT (must be 0 here)
+assert enable_trap_handler == 0
+
+enable_sgpr_workgroup_id_x = bits(pgm_rsrc2, 7, 7)
+enable_sgpr_workgroup_id_y = bits(pgm_rsrc2, 8, 8)
+enable_sgpr_workgroup_id_z = bits(pgm_rsrc2, 9, 9)
+enable_sgpr_workgroup_info = bits(pgm_rsrc2, 10, 10)
+
+enable_vgpr_workitem_id = bits(pgm_rsrc2, 11, 12)  # TIDIG_CMP_CNT enum (0..3)
+
+enable_exception_address_watch = bits(pgm_rsrc2, 13, 13)
+assert enable_exception_address_watch == 0
+enable_exception_memory = bits(pgm_rsrc2, 14, 14)
+assert enable_exception_memory == 0
+
+granulated_lds_size = bits(pgm_rsrc2, 15, 23)
+assert granulated_lds_size == 0  # spec: must be 0; CP uses dispatch packet rounding
+
+enable_exception_fp_invalid = bits(pgm_rsrc2, 24, 24)
+enable_exception_fp_denorm_src = bits(pgm_rsrc2, 25, 25)
+enable_exception_fp_div0 = bits(pgm_rsrc2, 26, 26)
+enable_exception_fp_overflow = bits(pgm_rsrc2, 27, 27)
+enable_exception_fp_underflow = bits(pgm_rsrc2, 28, 28)
+enable_exception_fp_inexact = bits(pgm_rsrc2, 29, 29)
+enable_exception_int_div0 = bits(pgm_rsrc2, 30, 30)
+
+assert_zero(pgm_rsrc2, 31, 31)
+
+print("RSRC2.ENABLE_PRIVATE_SEGMENT:", enable_private_segment)
+print("RSRC2.USER_SGPR_COUNT:", user_sgpr_count)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_X:", enable_sgpr_workgroup_id_x)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Y:", enable_sgpr_workgroup_id_y)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Z:", enable_sgpr_workgroup_id_z)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_INFO:", enable_sgpr_workgroup_info)
+print("RSRC2.ENABLE_VGPR_WORKITEM_ID (enum):", enable_vgpr_workitem_id)
+
+print("RSRC2.EXC_FP_INVALID:", enable_exception_fp_invalid)
+print("RSRC2.EXC_FP_DENORM_SRC:", enable_exception_fp_denorm_src)
+print("RSRC2.EXC_FP_DIV0:", enable_exception_fp_div0)
+print("RSRC2.EXC_FP_OVERFLOW:", enable_exception_fp_overflow)
+print("RSRC2.EXC_FP_UNDERFLOW:", enable_exception_fp_underflow)
+print("RSRC2.EXC_FP_INEXACT:", enable_exception_fp_inexact)
+print("RSRC2.EXC_INT_DIV0:", enable_exception_int_div0)
+
+# user sgprs
+
+enable_sgpr_private_segment_buffer = bits(desc, 448, 448)
+enable_sgpr_dispatch_ptr = bits(desc, 449, 449)
+enable_sgpr_queue_ptr = bits(desc, 450, 450)
+enable_sgpr_kernarg_segment_ptr = bits(desc, 451, 451)
+enable_sgpr_dispatch_id = bits(desc, 452, 452)
+enable_sgpr_flat_scratch_init = bits(desc, 453, 453)
+enable_sgpr_private_segment_size = bits(desc, 454, 454)
+
+assert_zero(desc, 455, 457)
+
+print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER:", enable_sgpr_private_segment_buffer)
+print("DESC.ENABLE_SGPR_DISPATCH_PTR:", enable_sgpr_dispatch_ptr)
+print("DESC.ENABLE_SGPR_QUEUE_PTR:", enable_sgpr_queue_ptr)
+print("DESC.ENABLE_SGPR_KERNARG_SEGMENT_PTR:", enable_sgpr_kernarg_segment_ptr)
+print("DESC.ENABLE_SGPR_DISPATCH_ID:", enable_sgpr_dispatch_id)
+print("DESC.ENABLE_SGPR_FLAT_SCRATCH_INIT:", enable_sgpr_flat_scratch_init)
+print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_SIZE:", enable_sgpr_private_segment_size)
+
+assert_zero(desc, 458, 459)
+
+uses_dynamic_stack = bits(desc, 459, 460)
+print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
+
+# gfx950 only
+assert_zero(desc, 460, 463)
+kernarg_preload_spec_length = bits(desc, 464, 470)
+print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
+kernarg_preload_spec_offset = bits(desc, 471, 479)
+print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
+
+assert_zero(desc, 480, 511)