Merge branch 'master' into call_inline

inline all these calls
2026-02-10 14:12:08 +08:00 · 2026-02-10 14:02:43 +08:00 · 2026-02-10 13:25:29 +08:00 · 2026-02-10 12:19:51 +08:00
420 changed files with 17943 additions and 28724 deletions
@@ -145,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

@@ -233,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
@@ -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.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v13_0_12, smu_v14_0_2"
+        python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, 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_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:
-    name: In-tree Autogen (comgr 2)
+  autogen-comgr-3:
+    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
-        python3 -c "from tinygrad.runtime.autogen import comgr"
+        rm tinygrad/runtime/autogen/comgr_3.py
+        python3 -c "from tinygrad.runtime.autogen import comgr_3"
    - name: Check for differences
      run: |
        if ! git diff --quiet; then
-          git diff
-          git diff > autogen-comgr2.patch
-          echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          git diff > autogen-comgr3.patch
+          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
          exit 1
        fi
    - name: Upload patch artifact
      if: failure()
      uses: actions/upload-artifact@v4
      with:
-        name: autogen-comgr2-patch
-        path: autogen-comgr2.patch
+        name: autogen-comgr3-patch
+        path: autogen-comgr3.patch
@@ -21,9 +21,6 @@ jobs:
  # 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:
@@ -44,14 +41,22 @@ jobs:
      run: |
        echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
        rm -f /tmp/pytest-db-ci*
+    # TODO: remove this step once all old caches are migrated
+    - name: Migrate old huggingface cache (symlinks break onnxruntime 1.24+)
+      run: |
+        cd ~/Library/Caches/tinygrad/downloads/models 2>/dev/null || exit 0
+        for old_dir in models--*; do
+          [ -d "$old_dir" ] || continue
+          repo_id=$(echo "$old_dir" | sed 's/models--//; s/--/\//g')
+          snapshot=$(ls -1 "$old_dir/snapshots" 2>/dev/null | head -1)
+          [ -n "$snapshot" ] || continue
+          mkdir -p "$repo_id"
+          cp -RLn "$old_dir/snapshots/$snapshot/"* "$repo_id/" 2>/dev/null || true
+        done
    - 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
@@ -332,13 +337,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 +515,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 +525,8 @@ 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: 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 +621,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]
@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '17'
+  CACHE_VERSION: '16'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
@@ -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/backend/test_ops.py TestOps.test_simple_conv2d
+        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_simple_conv2d
    - name: Test emulated METAL tensor cores
      run: |
-        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_big_gemm
+        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated AMX tensor cores
-      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
+      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
    - name: Test emulated AMD tensor cores
      run: |
        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
@@ -197,9 +197,9 @@ jobs:
        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated CUDA tensor cores
      run: |
-        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
-        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
        DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated INTEL OpenCL tensor cores
      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
@@ -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
@@ -299,6 +268,20 @@ jobs:
      run: |
        CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
+    - name: Run targetted tests on NULL backend
+      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
+    # 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
    - name: Run GC tests
      run: python test/external/external_uop_gc.py
    - name: External Benchmark Schedule
@@ -312,8 +295,8 @@ jobs:
        python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
        #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 24000 lines
-      run: MAX_LINE_COUNT=24000 python sz.py
+    - name: Repo line count < 20000 lines
+      run: MAX_LINE_COUNT=20000 python sz.py

  spec:
    strategy:
@@ -333,7 +316,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/null --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}

  fuzzing:
    name: Fuzzing
@@ -371,7 +354,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 +378,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 +437,7 @@ jobs:
      - name: Test Additional ONNX Ops (CPU)
        run: CPU=1 CPU_LLVM=0 python3 test/external/external_test_onnx_ops.py
      - name: Test Quantize ONNX
-        run: CPU=1 CPU_LLVM=0 python3 test/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

@@ -568,11 +551,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)
@@ -604,9 +587,9 @@ jobs:
    - 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 +608,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 --ignore=test/null --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 +683,31 @@ jobs:
          python-version: '3.14'
      - name: Verify AMD autogen is up to date
        run: |
-          python -m tinygrad.renderer.amd.generate
-          git diff --exit-code tinygrad/runtime/autogen/amd/
+          python -m extra.assembly.amd.generate
+          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
-      - name: Run AMD renderer tests
-        run: AMD_LLVM=0 python -m pytest -n=auto test/amd/ --durations 20
-      - name: Run AMD renderer tests (AMD_LLVM=1)
-        run: AMD_LLVM=1 python -m pytest -n=auto test/amd/ --durations 20
-      - name: Run SQTT profiling tests
-        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
-      - name: Run AMD emulated tests on NULL backend
-        env:
-          AMD: 0
-        run: |
-          PYTHONPATH=. NULL=1 EMULATE=AMD python extra/mmapeak/mmapeak.py
-          PYTHONPATH=. NULL=1 EMULATE=AMD_CDNA4 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
-      - name: Run ASM matmul on MOCKGPU
-        run: PYTHONPATH="." AMD=1 MOCKGPU=1 N=256 python3 extra/gemm/amd_asm_matmul.py
-      - name: Run LLVM test
-        run: AMD_LLVM=1 python test/device/test_amd_llvm.py
-
-  testmockam:
-    name: Linux (am)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    env:
-      AMD: 1
-      MOCKGPU: 1
-      AMD_IFACE: PCI
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: mockam
-          deps: testing_unit
-          amd: 'true'
-      - name: Run test_tiny on MOCKAM
-        run: python test/test_tiny.py
-      - name: Run test_tiny on MOCKAM USB
-        run: AMD_IFACE=USB python test/test_tiny.py
-      - name: Run test_hcq on MOCKAM
-        run: python -m pytest test/device/test_hcq.py
-
-  testamd:
-    strategy:
-      fail-fast: false
-      matrix:
-        backend: [amd, amdllvm]
-        arch: [rdna3, rdna4]
-        #arch: [rdna3, rdna4, cdna4]
-
-    name: Linux (${{ matrix.backend }} ${{ 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
+        run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
+      - name: Run RDNA3 emulator tests
+        run: AMD_LLVM=0 python -m pytest -n=auto extra/assembly/amd/ --durations 20
+      - name: Run RDNA3 emulator tests (AMD_LLVM=1)
+        run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
+      - name: Run RDNA3 dtype tests
+        run: AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py --durations 20
+      - name: Run RDNA3 dtype tests (AMD_LLVM=1)
+        run: AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py --durations 20
+      # TODO: run all once emulator is faster
+      - name: Run RDNA3 ops tests
+        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
+      - name: Run RDNA4 emulator tests
+        run: MOCKGPU_ARCH=rdna4 python -m pytest test/test_tiny.py -v --durations 20

  testnvidia:
    strategy:
@@ -768,12 +736,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/null --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 +770,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 --ignore=test/null --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

@@ -836,17 +804,15 @@ jobs:
    - 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 +835,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 +853,14 @@ jobs:
        key: osx-webgpu
        deps: testing
        webgpu: 'true'
+    - name: Test infinity math in WGSL
+      run: WEBGPU=1 python -m pytest -n=auto test/test_renderer_failures.py::TestWGSLFailures::test_multiply_infinity --durations=20
    - name: Build WEBGPU Efficientnet
      run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m examples.compile_efficientnet
-    - name: Run selected webgpu tests
-      run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m pytest -n=auto test/backend --durations=20
-    #- name: Clean npm cache
-    #  run: npm cache clean --force
-    #- name: Install Puppeteer
-    #  run: npm install puppeteer
+    - name: Clean npm cache
+      run: npm cache clean --force
+    - name: Install Puppeteer
+      run: npm install puppeteer
    # this is also flaky
    #- name: Run WEBGPU Efficientnet
    #  run: node test/web/test_webgpu.js
@@ -938,7 +902,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 --ignore=test/null --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay
      - name: Run macOS-specific unit test
@@ -980,7 +944,7 @@ jobs:
        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 ******

@@ -1009,5 +973,5 @@ jobs:
        shell: bash
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          DEBUG=4 python3 test/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/test_ops.py --durations=20
@@ -66,5 +66,3 @@ target
 .mypy_cache
 mutants
 .mutmut-cache
-dagre/
-graphlib/
@@ -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/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
        language: system
        always_run: true
        pass_filenames: false
@@ -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.
@@ -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.
@@ -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
 ```

@@ -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
@@ -19,8 +19,8 @@ cifar_std = [0.24703225141799082, 0.24348516474564, 0.26158783926049628]
 BS, STEPS = getenv("BS", 512), getenv("STEPS", 1000)
 EVAL_BS = getenv("EVAL_BS", BS)
 GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
-assert BS % len(GPUS) == 0, f"{BS=} is not a multiple of {len(GPUS)=}"
-assert EVAL_BS % len(GPUS) == 0, f"{EVAL_BS=} is not a multiple of {len(GPUS)=}"
+assert BS % len(GPUS) == 0, f"{BS=} is not a multiple of {len(GPUS)=}, uneven multi GPU is slow"
+assert EVAL_BS % len(GPUS) == 0, f"{EVAL_BS=} is not a multiple of {len(GPUS)=}, uneven multi GPU is slow"

 class UnsyncedBatchNorm:
  def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1, num_devices=len(GPUS)):
@@ -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_memoryview(force_zero_copy=True)[:] = img.tobytes()
+
+      # ideal
+      #X[idx].assign(img.tobytes())   # NOTE: this is slow!
      q_out.put(idx)
    q_out.put(None)

@@ -254,8 +264,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
      x = random_brightness_augmentation(x)
      x = gaussian_noise(x)

-    X[idx].flatten().assign(x.tobytes())
-    Y[idx].flatten().assign(y.tobytes())
+    X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = x.tobytes()
+    Y[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = y.tobytes()

    queue_out.put(idx)
  queue_out.put(None)
@@ -369,12 +379,12 @@ def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue
      clipped_match_idxs = np.clip(match_idxs, 0, None)
      clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]

-      boxes[idx].flatten().assign(clipped_boxes.tobytes())
-      labels[idx].flatten().assign(clipped_labels.tobytes())
-      matches[idx].flatten().assign(match_idxs.tobytes())
-      anchors[idx].flatten().assign(anchor.tobytes())
+      boxes[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_boxes.tobytes()
+      labels[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_labels.tobytes()
+      matches[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = match_idxs.tobytes()
+      anchors[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = anchor.tobytes()

-    imgs[idx].flatten().assign(img.tobytes())
+    imgs[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()

    queue_out.put(idx)
  queue_out.put(None)
@@ -396,7 +406,6 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
      queue_in.put((idx, img, tgt))

  def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
-    shm_name = f"{shm_name}_{os.getpid()}"
    if os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
    shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(size))
    shm_tensor = Tensor.empty(*size, dtype=dtype, device=f"disk:/dev/shm/{shm_name}")
@@ -543,7 +552,7 @@ class BinIdxDataset:
    version, = struct.unpack("<Q", self.idx.read(8))
    assert version == 1, "unsupported index version"
    dtype_code, = struct.unpack("<B", self.idx.read(1))
-    self.dtype = {1:np.dtype(np.uint8), 2:np.dtype(np.int8), 3:np.dtype(np.int16), 4:np.dtype(np.int32), 5:np.dtype(np.int64), 6:np.dtype(np.float64), 7:np.dtype(np.double), 8:np.dtype(np.uint16)}[dtype_code]
+    self.dtype = {1:dtypes.uint8, 2:dtypes.int8, 3:dtypes.int16, 4:dtypes.int32, 5:dtypes.int64, 6:dtypes.float64, 7:dtypes.double, 8:dtypes.uint16}[dtype_code]
    self.count, = struct.unpack("<Q", self.idx.read(8))
    doc_count, = struct.unpack("<Q", self.idx.read(8))

@@ -560,7 +569,7 @@ class BinIdxDataset:
    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()

    # bin file
-    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin")).numpy()
+    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))

  def _index(self, idx) -> tuple[int, int]:
    return int(self.pointers[idx]), int(self.sizes[idx])
@@ -569,7 +578,7 @@ class BinIdxDataset:
    ptr, size = self._index(idx)
    if length is None: length = size - offset
    ptr += offset * self.dtype.itemsize
-    return self.bin_t[ptr:ptr+length*self.dtype.itemsize].view(self.dtype)
+    return self.bin_t[ptr:ptr+length*self.dtype.itemsize].bitcast(self.dtype).to(None)

 # https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
 class GPTDataset:
@@ -628,7 +637,7 @@ class GPTDataset:
        sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))

    # concat all parts
-    text = np.concatenate(sample_parts, axis=0)
+    text = Tensor.cat(*sample_parts)

    return text

@@ -771,8 +780,7 @@ def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:b
 def iterate_llama3_dataset(dataset:BlendedGPTDataset, bs:int):
  for b in range(math.ceil(dataset.samples / bs)):
    batch = [dataset.get(b * bs + i) for i in range(bs)]
-    stacked = np.stack(batch, axis=0)
-    yield Tensor(stacked, device="NPY")
+    yield Tensor.stack(batch, dim=0)

 def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
  return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)
@@ -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

@@ -1285,7 +1285,6 @@ def train_llama3():
  from extra.models.llama import Transformer
  from examples.llama3 import MODEL_PARAMS
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW

  BENCHMARK = getenv("BENCHMARK")

@@ -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,14 +1333,9 @@ def train_llama3():
  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
  # vocab_size from the mixtral tokenizer
  if not SMALL: model_params |= {"vocab_size": 32000}
-  real_vocab_size = model_params['vocab_size']
  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
  print(f"model parameters: {model_params}")

-  # pad vocab
-  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
-  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
-
  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
  params = get_parameters(model)
  # weights are all bfloat16 for now
@@ -1357,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():
@@ -1366,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)
@@ -1379,18 +1369,13 @@ def train_llama3():
      # prevents memory spike on device 0
      v.realize()

-    vocab_mask.shard_(device, axis=2).realize()
-
-  optim_device = "CPU" if getenv("OFFLOAD_OPTIM") else None
-  optim = GradAccClipAdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
-                           eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
+  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)

  # init grads
  for p in optim.params:
-    p.grad = p.empty_like().realize()
-  grads: list[Tensor] = [p.grad for p in optim.params]
-  for p in optim.params:
-    p.grad.assign(p.grad.zeros_like()).realize()
+    p.grad = p.zeros_like().contiguous().realize()
+  grads = [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)

@@ -1407,58 +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)
-    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
    loss.backward()
    assert all(p.grad is g for p,g in zip(optim.params, grads))
    Tensor.realize(loss, *grads)
-    return loss.flatten().float().to("CPU")
+    return loss

  @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()).realize()
+      g.assign(g.zeros_like().contiguous()).realize()

    lr = optim.lr
    Tensor.realize(lr, *grads)

-    return lr.float().to("CPU"), grad_norm.float().to("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)
-    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten().float().to("CPU")
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    return loss.flatten().float()

  # ** data iters **
  def fake_data(bs, samples):
-    import numpy as np
    for _ in range(samples // bs):
-      fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
-      yield Tensor(fake_data_np, device="NPY")
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)

  def get_train_iter():
    if getenv("FAKEDATA", 0):
      return fake_data(BS, SAMPLES)
    else:
      from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=DATA_SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
+      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))

  if getenv("FAKEDATA", 0):
    eval_dataset = None
@@ -1478,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 if i >= 2 else 1):
+      for _ in range(grad_acc):
        ist = time.perf_counter()
        try: tokens = next(train_iter)
        except StopIteration:
          stopped = True
          break
-        mst = time.perf_counter()
-        data_time += mst - ist
-        losses.append(minibatch(tokens).item())
-        dev_time += time.perf_counter() - mst
+        dt = time.perf_counter()
+        loss = minibatch(tokens)
      if stopped: break

      gt = time.perf_counter()
-      ret = optim_step()
-      lr, grad_norm = ret[0].item(), ret[1].item()
-      et = time.perf_counter()
+      lr = optim_step()
+      ot = time.perf_counter()

-      loss = sum(losses) / len(losses)
-      optim_time = et - gt
-      dev_time += optim_time
+      loss = loss.float().item()
+      lr = lr.item()
+
+      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")
-      if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
+          f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")

      if WANDB:
        wandb.log({
-          "train/loss": loss,
-          "train/lr": lr,
-          "train/grad_norm": grad_norm,
+          "lr": lr, "train/loss": loss,
          "train/step_time": step_time,
          "train/gbs_time": gbs_time,
          "train/optim_time": optim_time,
@@ -1553,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}")
@@ -1561,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()
@@ -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).realize()
-      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous().realize()
-      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)).realize()
-
-    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)
-      v_hat = self.v[i] / (1.0 - self.b2_t)
-      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)
@@ -1,36 +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 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
@@ -1,31 +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 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
@@ -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:-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=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

 export MODEL="llama3"
@@ -23,20 +21,16 @@ 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="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}

 export SEED=${SEED:-5760}
-export DATA_SEED=${DATA_SEED:-5760}

 export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5

-export FAKEDATA=1 BENCHMARK=10
-if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=2
-fi
+export FAKEDATA=1 BENCHMARK=10 LLAMA_LAYERS=2

 python3 examples/mlperf/model_train.py
@@ -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
@@ -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:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

 export MODEL="llama3"
@@ -23,15 +21,14 @@ 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="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}

-export SEED=${SEED:-$RANDOM}
-export DATA_SEED=${DATA_SEED:-5760}
+export SEED=${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
@@ -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
@@ -3,4 +3,4 @@ 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
+PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
@@ -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)
@@ -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"
@@ -154,7 +154,7 @@ class SMICtx:
    for dev in self.devs:
      match dev.ip_ver[am.MP1_HWIP]:
        case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableV0_t
-        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTable_t
+        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTableV2_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 +231,12 @@ class SMICtx:

  def get_power(self, dev, metrics):
    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
-      case (13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.SocketPowerLimit)
+      case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
      case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX

  def get_mem_usage(self, dev):
+    return 0
+
    usage = 0
    pt_stack = [dev.mm.root_page_table]
    while len(pt_stack) > 0:
@@ -244,8 +245,8 @@ class SMICtx:
        entry = pt.entries[i]

        if (entry & am.AMDGPU_PTE_VALID) == 0: continue
-        if pt.lv < am.AMDGPU_VM_PDB0 and not dev.gmc.is_pte_huge_page(pt.lv, entry):
-          pt_stack.append(AMPageTableEntry(dev, dev.xgmi2paddr(entry & 0x0000FFFFFFFFF000), lv=pt.lv+1))
+        if pt.lv!=am.AMDGPU_VM_PTB and not dev.gmc.is_pte_huge_page(pt.lv, entry):
+          pt_stack.append(AMPageTableEntry(dev, entry & 0x0000FFFFFFFFF000, lv=pt.lv+1))
          continue
        if (entry & am.AMDGPU_PTE_SYSTEM) != 0: continue
        usage += (1 << ((9 * (3-pt.lv)) + 12))
@@ -1,17 +1,17 @@
 An integrated environment for AMD GPU assembly and emulation

-Test with `pytest -n12 test/amd/`
-`AMD_LLVM=1 pytest -n12 test/amd/`
+Test with `PYTHONPATH="." pytest -n12 extra/assembly/amd/`
+`AMD_LLVM=1 PYTHONPATH="." pytest -n12 extra/assembly/amd/`

+* pdf.py -- extract assembly format + instruction pseudocode from AMD PDF
 * dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
-* test/mockgpu/amd/emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
-* generate.py -- extract assembly format + instruction pseudocode from AMD XML + PDF
-* test/mockgpu/amd/pcode.py -- pseudocode to UOp transformation
-* sqtt.py -- SQTT parser
+* pcode.py -- pseudocode execution environment. pseudocode should be transformed as little as possible.
+* asm.py -- an asm/disasm function to transform to and from AMD assembly syntax
+* emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`

-The code should be as readable and deduplicated as possible. emu (in test/mockgpu/amd/) shouldn't be required for dsl.
+The code should be as readable and deduplicated as possible. asm and emu shouldn't be required for dsl.

-The autogen folder is autogenerated from the AMD PDFs with `python3 -m tinygrad.renderer.amd.pdf --arch all`
+The autogen folder is autogenerated from the AMD PDFs with `python3 -m extra.assembly.amd.pdf --arch all`

 test_emu.py has a good set of instruction tests for the emulation, with USE_HW=1 it will compare to real hardware.
 Whenever an instruction is fixed, regression tests should be added here and confirmed with real hardware.
@@ -20,20 +20,20 @@ test_llvm.py tests asm/disasm on the LLVM tests, confirming it behaves the same

 tinygrad's dtype tests should pass with and without LLVM. they run in about 12 seconds.

-`AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
-`AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
+`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
+`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`

 The ops tests also pass, but they are very slow, so you should run them one at a time.

-`SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/backend/test_ops.py`
-`SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/backend/test_ops.py`
+`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_ops.py`
+`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_ops.py`

-When something is caught by main tinygrad tests, a local regression test should be added to `test/amd`.
+When something is caught by main tinygrad tests, a local regression test should be added to `extra/assembly/amd/test`.
 While working with tinygrad, you can dump the assembly with `DEBUG=7`. These tests all pass on real hardware
 If a test is failing with `AMD=1 PYTHON_REMU=1 MOCKGPU=1` it's because an instruction is emulated incorrectly.
 You can test without `MOCKGPU=1` to test on real hardware, if it works on real hardware there's a bug in the emulator.
 IMPORTANT: if a test is failing in the emulator, it's an instruction bug. Use DEBUG=7, get the instructions, and debug.

-Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~3000 lines.
-Get line count with `cloc --by-file tinygrad/renderer/amd/*.py`
+Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~2000 lines.
+Get line count with `cloc --by-file extra/assembly/amd/*.py`

@@ -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)
@@ -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
@@ -1,7 +1,7 @@
 # autogenerated from AMD ISA XML - do not edit
-# ruff: noqa: E501,F401
-from tinygrad.renderer.amd.dsl import BitField, DPP, DPP16, EXEC, EXECZ, EXEC_HI, EXEC_LO, EnumBitField, FixedBitField, INV_2PI, Inst, LIT, M0, NULL, OFF, SBaseField, SCC, SDWA, SGPRField, SRC_LDS_DIRECT, SRsrcField, SSrcField, SrcField, VCC, VCCZ, VCC_HI, VCC_LO, VGPRField, s, src, ttmp, v
-from tinygrad.runtime.autogen.amd.cdna.enum import DSOp, FLATOp, GLOBALOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3PX2Op, VOP3SDOp, VOPCOp, HWREG
+# ruff: noqa: F401,F403
+from extra.assembly.amd.dsl import *
+from extra.assembly.amd.autogen.cdna.enum import *
 import functools

 class DS(Inst):
@@ -1,6 +1,6 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import Fmt, OpType
-from tinygrad.runtime.autogen.amd.cdna.enum import DSOp, FLATOp, GLOBALOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3PX2Op, VOP3SDOp, VOPCOp
+from extra.assembly.amd.autogen.common import Fmt, OpType
+from extra.assembly.amd.autogen.cdna.enum import *

 # instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}
 OPERANDS = {
@@ -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',
@@ -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
@@ -1,7 +1,7 @@
 # autogenerated from AMD ISA XML - do not edit
-# ruff: noqa: E501,F401
-from tinygrad.renderer.amd.dsl import BitField, DPP, DPP16, EXEC, EXECZ, EXEC_HI, EXEC_LO, EnumBitField, FixedBitField, INV_2PI, Inst, LIT, M0, NULL, OFF, SBaseField, SCC, SDWA, SGPRField, SRC_LDS_DIRECT, SRsrcField, SSrcField, SrcField, VCC, VCCZ, VCC_HI, VCC_LO, VDSTYField, VGPRField, s, src, ttmp, v
-from tinygrad.runtime.autogen.amd.rdna3.enum import DSOp, EXPOp, FLATOp, GLOBALOp, LDSDIROp, MIMGOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, HWREG, MSG
+# ruff: noqa: F401,F403
+from extra.assembly.amd.dsl import *
+from extra.assembly.amd.autogen.rdna3.enum import *
 import functools

 class DS(Inst):
@@ -593,6 +593,9 @@ flat_load_d16_hi_i8 = functools.partial(FLAT, FLATOp.FLAT_LOAD_D16_HI_I8)
 flat_load_d16_hi_b16 = functools.partial(FLAT, FLATOp.FLAT_LOAD_D16_HI_B16)
 flat_store_d16_hi_b8 = functools.partial(FLAT, FLATOp.FLAT_STORE_D16_HI_B8)
 flat_store_d16_hi_b16 = functools.partial(FLAT, FLATOp.FLAT_STORE_D16_HI_B16)
+global_load_addtid_b32 = functools.partial(FLAT, FLATOp.GLOBAL_LOAD_ADDTID_B32)
+global_store_addtid_b32 = functools.partial(FLAT, FLATOp.GLOBAL_STORE_ADDTID_B32)
+global_load_lds_addtid_b32 = functools.partial(FLAT, FLATOp.GLOBAL_LOAD_LDS_ADDTID_B32)
 flat_atomic_swap_b32 = functools.partial(FLAT, FLATOp.FLAT_ATOMIC_SWAP_B32)
 flat_atomic_cmpswap_b32 = functools.partial(FLAT, FLATOp.FLAT_ATOMIC_CMPSWAP_B32)
 flat_atomic_add_u32 = functools.partial(FLAT, FLATOp.FLAT_ATOMIC_ADD_U32)
@@ -1,6 +1,6 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import Fmt, OpType
-from tinygrad.runtime.autogen.amd.rdna3.enum import DSOp, EXPOp, FLATOp, GLOBALOp, LDSDIROp, MIMGOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp
+from extra.assembly.amd.autogen.common import Fmt, OpType
+from extra.assembly.amd.autogen.rdna3.enum import *

 # instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}
 OPERANDS = {
@@ -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',
@@ -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
@@ -1,7 +1,7 @@
 # autogenerated from AMD ISA XML - do not edit
-# ruff: noqa: E501,F401
-from tinygrad.renderer.amd.dsl import BitField, DPP, DPP16, EXEC, EXECZ, EXEC_HI, EXEC_LO, EnumBitField, FixedBitField, INV_2PI, Inst, LIT, M0, NULL, OFF, SBaseField, SCC, SDWA, SGPRField, SRC_LDS_DIRECT, SSrcField, SrcField, VCC, VCCZ, VCC_HI, VCC_LO, VDSTYField, VGPRField, s, src, ttmp, v
-from tinygrad.runtime.autogen.amd.rdna4.enum import DSOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VBUFFEROp, VDSDIROp, VEXPORTOp, VFLATOp, VGLOBALOp, VIMAGEOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, VSAMPLEOp, VSCRATCHOp, HWREG, MSG
+# ruff: noqa: F401,F403
+from extra.assembly.amd.dsl import *
+from extra.assembly.amd.autogen.rdna4.enum import *
 import functools

 class DS(Inst):
@@ -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)
@@ -1,6 +1,6 @@
 # autogenerated from AMD ISA XML - do not edit
-from tinygrad.runtime.autogen.amd.common import Fmt, OpType
-from tinygrad.runtime.autogen.amd.rdna4.enum import DSOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VBUFFEROp, VDSDIROp, VEXPORTOp, VFLATOp, VGLOBALOp, VIMAGEOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, VSAMPLEOp, VSCRATCHOp
+from extra.assembly.amd.autogen.common import Fmt, OpType
+from extra.assembly.amd.autogen.rdna4.enum import *

 # instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}
 OPERANDS = {
@@ -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',
@@ -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',
-                246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc'}
+                128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc'}
 SPECIAL_GPRS_CDNA = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'm0', 126: 'exec_lo', 127: 'exec_hi',
-                     128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0',
-                     246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc',
+                     128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc',
                     102: 'flat_scratch_lo', 103: 'flat_scratch_hi', 104: 'xnack_mask_lo', 105: 'xnack_mask_hi',
                     251: 'src_vccz', 252: 'src_execz'}
 SPECIAL_PAIRS = {106: 'vcc', 126: 'exec'}
@@ -72,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',
-    'ALIGNBYTE', 'V_PERM_', 'XOR3', 'XAD', 'MULLIT', 'MINMAX', 'MAXMIN', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM',
-    'MINIMUM3', 'MAXIMUM3', 'MIN3', 'MAX3', 'DOT2', 'CVT_PK_U8_F32', 'DOT4', 'DOT8', 'WMMA', 'SWMMAC')): return 3
+  if any(x in n for x in ('_ADD3', '_LSHL_ADD', '_ADD_LSHL', '_LSHL_OR', '_AND_OR', 'OR3_B32', 'AND_OR_B32', 'ALIGNBIT', 'ALIGNBYTE', 'V_PERM_', 'XOR3', 'XAD', 'MULLIT', 'MINMAX', 'MAXMIN', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM', 'MINIMUM3', 'MAXIMUM3', 'MIN3', 'MAX3', 'DOT2', 'CVT_PK_U8_F32', 'DOT4', 'DOT8', 'WMMA', 'SWMMAC')): return 3
  return 2

 # ═══════════════════════════════════════════════════════════════════════════════
 # IMPORTS
 # ═══════════════════════════════════════════════════════════════════════════════

-from tinygrad.runtime.autogen.amd.rdna3.ins import (VOP1, VOP1_SDST, VOP1_SDST_LIT, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3_SDST_LIT,
+from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP1_SDST_LIT, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3_SDST_LIT,
  VOP3_LIT, VOP3SD, VOP3SD_LIT, VOP3P, VOP3P_LIT, VOPC, VOPC_LIT, VOPD, VOPD_LIT, VINTERP, SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPC_LIT,
  SOPK, SOPK_LIT, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH, VOP2Op, VOPDOp, SOPPOp, HWREG, MSG)
-from tinygrad.runtime.autogen.amd.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST,
-  VOP1_SDST_LIT as R4_VOP1_SDST_LIT, VOP1_LIT as R4_VOP1_LIT,
+from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP1_SDST_LIT as R4_VOP1_SDST_LIT, VOP1_LIT as R4_VOP1_LIT,
  VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT, VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3_SDST_LIT as R4_VOP3_SDST_LIT, VOP3_LIT as R4_VOP3_LIT,
  VOP3SD as R4_VOP3SD, VOP3SD_LIT as R4_VOP3SD_LIT, VOP3P as R4_VOP3P, VOP3P_LIT as R4_VOP3P_LIT, VOPC as R4_VOPC, VOPC_LIT as R4_VOPC_LIT,
  VOPD as R4_VOPD, VOPD_LIT as R4_VOPD_LIT, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP1_LIT as R4_SOP1_LIT, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
  SOPC as R4_SOPC, SOPC_LIT as R4_SOPC_LIT, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP, SMEM as R4_SMEM, DS as R4_DS,
  VOPDOp as R4_VOPDOp, HWREG as HWREG_RDNA4, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
-from tinygrad.runtime.autogen.amd.cdna.ins import HWREG as HWREG_CDNA
+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:
-  b = _unwrap(b)
-  return _reg("v", b - 256 if b >= 256 else b, n)
-def _areg(b: int, n: int = 1) -> str:
-  b = _unwrap(b)
-  return _reg("a", b - 256 if b >= 256 else b, n)  # accumulator registers for GFX90a
-def _ttmp(b, n: int = 1) -> str | None:
-  b = _unwrap(b)
-  return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
+def _vreg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("v", b - 256 if b >= 256 else b, n)
+def _areg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("a", b - 256 if b >= 256 else b, n)  # accumulator registers for GFX90a
+def _ttmp(b, n: int = 1) -> str | 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:
-  v = _unwrap(v)
-  return _fmt_v16(v) if v >= 256 else _lit(inst, v)  # format 16-bit src: vgpr.h/l or literal
+def _src16(inst, v: int) -> str: v = _unwrap(v); return _fmt_v16(v) if v >= 256 else _lit(inst, v)  # format 16-bit src: vgpr.h/l or literal
 def _mods(*pairs) -> str: return " ".join(m for c, m in pairs if c)
 def _fmt_bits(label: str, val: int, count: int) -> str: return f"{label}:[{','.join(str((val >> i) & 1) for i in range(count))}]"

@@ -214,8 +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:
-    return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
+  if not cdna and basename in _VOP2_CARRY_INOUT_RDNA: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
  sn0 = inst.canonical_op_regs.get('s0', 1)
  if inst.vdst.sz > 1 or sn0 > 1 or inst.vsrc1.sz > 1:
    src0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else _fmt_src(inst.src0, sn0, cdna)
@@ -231,10 +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)
-  elif inst.src0.sz > 1: s0 = inst.src0.fmt()
-  elif is16: s0 = _src16(inst, inst.src0.offset)
-  else: s0 = _lit(inst, inst.src0)
+  s0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else inst.src0.fmt() if inst.src0.sz > 1 else _src16(inst, inst.src0.offset) if is16 else _lit(inst, inst.src0)
  s1 = inst.vsrc1.fmt() if inst.vsrc1.sz > 1 else _fmt_v16(inst.vsrc1) if is16 else inst.vsrc1.fmt()
  suf = "" if name.endswith('_e32') else "_e32"
  return f"{name}{suf} vcc_lo, {s0}, {s1}" if has_vcc else f"{name}{suf} {s0}, {s1}"
@@ -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',
-            'TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
+    deps = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
    skips = ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
    id0, skip, id1 = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 7) & 0xf
-    def dep(v): return deps[v-1] if 0 < v <= len(deps) else str(v)
+    dep = lambda v: deps[v-1] if 0 < v <= len(deps) else str(v)
    p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
    return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
  if name.startswith(('s_cbranch', 's_branch')): return f"{name} {inst.simm16}"
@@ -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)
-  elif sbase_idx <= 105: sbase_str = _sreg(sbase_idx, sbase_count)
-  else: sbase_str = _reg("ttmp", sbase_idx - 108, sbase_count)
+  sbase_str = _fmt_src(sbase_idx, sbase_count, cdna) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
  if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {_unwrap(inst.sdata)}, {sbase_str}, {off_s}"
  if 'prefetch' in name:
    off = getattr(inst, 'ioffset', getattr(inst, 'offset', 0))
@@ -324,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)
-  offset = inst.ioffset if r4 else inst.offset  # type: ignore[attr-defined]
-  if r4: off_val = offset if offset < (1 << 23) else offset - (1 << 24)  # sign extend 24-bit
-  elif seg != 'flat':
+  # Global/scratch uses 13-bit signed offset
+  offset = inst.ioffset if r4 else inst.offset
+  if seg != 'flat':
    if cdna:
      # CDNA: bit 12 is sign bit but not in offset field
      raw = int.from_bytes(inst.to_bytes(), 'little')
@@ -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]
-  elif r4:
-    th_names = R4_TH_ATOMIC if 'atomic' in name else (R4_TH_STORE if 'store' in name else R4_TH_LOAD)
-    mods = off_s + (f" th:{th_names[inst.th]}" if inst.th in th_names else "") + (f" scope:{R4_SCOPE[inst.scope]}" if inst.scope in R4_SCOPE else "")
+  if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"
+  elif r4: mods = f"{off_s}{' scope' if inst.scope else ''}{' th' if inst.th else ''}"
  else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
  if seg == 'flat': saddr_s = ""
  elif _unwrap(inst.saddr) in (0x7F, 124): saddr_s = ", off"
  elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr, cdna)}"
-  elif _unwrap(inst.saddr) in (SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS):
-    saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
+  elif _unwrap(inst.saddr) in (SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS): saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
  elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
  else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
-  if 'addtid' in name: return f"{instr} {reg_fn((inst.vsrc if r4 else inst.data) if 'store' in name else inst.vdst)}{saddr_s}{mods}"
-   # RDNA4: vaddr instead of addr, vsrc instead of data
-  addr = inst.vaddr if r4 else inst.addr  # type: ignore[attr-defined]
-  data = inst.vsrc if r4 else inst.data  # type: ignore[attr-defined]
+  if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
+  # RDNA4: vaddr instead of addr, vsrc instead of data
+  addr = inst.vaddr if r4 else inst.addr
+  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]
-    sfx = f"{saddr_s if seg != 'flat' else ''}{mods}"
-    return f"{instr} {vdst_s}, {addr_s}, {data_s}{sfx}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{sfx}"
+    glc_or_sc0 = inst.sc0 if cdna else inst.glc
+    return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}"
  if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
  return f"{instr} {reg_fn(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"

 def _disasm_ds(inst: DS) -> str:
-  name = inst.op_name.lower()
+  op, name = inst.op, inst.op_name.lower()
  acc = getattr(inst, 'acc', 0)
  reg_fn = _areg if acc else _vreg
  gds = " gds" if getattr(inst, 'gds', 0) else ""
@@ -417,8 +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):
-    return f"{name} {reg_fn(inst.vdst)}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
+  if 'load' in name or ('read' in name and 'read2' not in name): return f"{name} {reg_fn(inst.vdst)}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
  if ('store' in name or 'write' in name) and not _has(name, 'cmp', 'xchg', 'write2'):
    return f"{name} {reg_fn(inst.data0)}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
  if 'swizzle' in name or name == 'ds_ordered_count': return f"{name} {reg_fn(inst.vdst)}, {addr}{off}{gds}"
@@ -429,15 +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)
-    if s0v == 255: src = _lit(inst, inst.src0)
-    elif s0v == 253: src = "src_scc"
-    else: src = _fmt_src(inst.src0, max(1, bits['s0'] // 32))
+    src = _lit(inst, inst.src0) if _unwrap(inst.src0) == 255 else ("src_scc" if _unwrap(inst.src0) == 253 else _fmt_src(inst.src0, max(1, bits['s0'] // 32)))
    if inst.neg & 1: src = f"-{src}"
    if inst.abs & 1: src = f"|{src}|"
    clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
@@ -446,6 +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))
-                   or (inst.src2.offset < 256 and (inst.opsel & 4)))
+  nonvgpr_opsel = (inst.src0.offset < 256 and (inst.opsel & 1)) or (inst.src1.offset < 256 and (inst.opsel & 2)) or (inst.src2.offset < 256 and (inst.opsel & 4))
  need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)

  op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
@@ -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 [])
-            + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if clamp else []))
+    mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if clamp else [])
  elif is_swmmac:
    mods = ([f"index_key:{inst.opsel}"] if inst.opsel else []) + ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + \
           ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
  else:
    opsel_hi_default = 7 if n == 3 else 3
-    mods = (([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else [])
-            + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else [])
-            + ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else [])
-            + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else []))
-  mod_s = ' ' + ' '.join(mods) if mods else ''
-  return f"{name} {dst}, {src0}, {src1}, {src2}{mod_s}" if n == 3 else f"{name} {dst}, {src0}, {src1}{mod_s}"
+    mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else []) + \
+           ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
+  return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"

 def _disasm_sop1(inst: SOP1) -> str:
-  name, cdna = inst.op_name.lower(), _is_cdna(inst)
+  op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
  # Use get_field_bits for register sizes
  regs = inst.canonical_op_regs
  dst_regs, src_regs = regs.get('d', 1), regs.get('s0', 1)
@@ -550,8 +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',
-                   'S_BARRIER_INIT', 'S_BARRIER_JOIN', 'S_SET_GPR_IDX_IDX', 'S_CBRANCH_JOIN')
+  sop1_src_only = ('S_ALLOC_VGPR', 'S_SLEEP_VAR', 'S_BARRIER_SIGNAL', 'S_BARRIER_SIGNAL_ISFIRST', 'S_BARRIER_INIT', 'S_BARRIER_JOIN', 'S_SET_GPR_IDX_IDX',
+                   'S_CBRANCH_JOIN')
  if inst.op_name in sop1_src_only: return f"{name} {src}"
  if cdna:
    if 'getpc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2, cdna)}"
@@ -589,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()}, {s0}, {s1}, {s2}" + (" " + mods if mods else "")
+  return f"{inst.op_name.lower()} {inst.vdst.fmt()}, {_lit(inst, inst.src0, inst.neg & 1)}, {_lit(inst, inst.src1, inst.neg & 2)}, {_lit(inst, inst.src2, inst.neg & 4)}" + (" " + mods if mods else "")

 DISASM_HANDLERS: dict[type, Callable[..., str]] = {
  VOP1: _disasm_vop1, VOP1_SDST: _disasm_vop1, VOP1_SDST_LIT: _disasm_vop1, VOP1_LIT: _disasm_vop1,
  VOP2: _disasm_vop2, VOP2_LIT: _disasm_vop2, VOPC: _disasm_vopc, VOPC_LIT: _disasm_vopc,
-  VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3_SDST_LIT: _disasm_vop3, VOP3_LIT: _disasm_vop3,
-  VOP3SD: _disasm_vop3sd, VOP3SD_LIT: _disasm_vop3sd,
+  VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3_SDST_LIT: _disasm_vop3, VOP3_LIT: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOP3SD_LIT: _disasm_vop3sd,
  VOPD: _disasm_vopd, VOPD_LIT: _disasm_vopd, VOP3P: _disasm_vop3p, VOP3P_LIT: _disasm_vop3p,
  VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, GLOBAL: _disasm_flat, SCRATCH: _disasm_flat,
  SOP1: _disasm_sop1, SOP1_LIT: _disasm_sop1, SOP2: _disasm_sop2, SOP2_LIT: _disasm_sop2,
@@ -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)
-    s1 = _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, r1)
-    s2 = _cdna_src(inst, inst.src2, inst.neg&4, inst.abs&4, r2)
+    s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, r0), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, r1), _cdna_src(inst, inst.src2, inst.neg&4, inst.abs&4, r2)
    dst = _vreg(inst.vdst, dregs) if dregs > 1 else _vreg(inst.vdst)
  if op_val >= 512:
    return f"{name} {dst}, {s0}, {s1}, {s2}{opsel}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{opsel}{cl}{om}"
@@ -700,9 +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)
-  s1 = _cdna_src(inst, inst.src1, inst.neg&2, n=r1)
-  s2 = _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
+  s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, n=r0), _cdna_src(inst, inst.src1, inst.neg&2, n=r1), _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
  # CDNA VOP3_SDST uses vdst field for sdst (but vdst adds 256), RDNA uses separate sdst field
  sdst_val = getattr(inst, 'sdst', None)
  if sdst_val is None and hasattr(inst, 'vdst'):
@@ -724,7 +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 [])
-          + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else [])
-          + ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else [])
-          + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else []))
-  mod_s = ' ' + ' '.join(mods) if mods else ''
-  return f"{name} {dst}, {src0}, {src1}, {src2}{mod_s}" if n == 3 else f"{name} {dst}, {src0}, {src1}{mod_s}"
+  mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else []) + \
+          ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
+  return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"

 def _disasm_mubuf(inst) -> str:
  name = inst.op_name.lower()
@@ -950,6 +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_VOP3P: _disasm_cdna_vop3p, CDNA_VOP3P_MFMA: _disasm_cdna_vop3p,
+  CDNA_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p, CDNA_VOP3P_MFMA: _disasm_cdna_vop3p,
  CDNA_MUBUF: _disasm_mubuf, CDNA_VOP3PX2: _disasm_vop3px2})
@@ -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:
-      idx = o - 256
-      base = f"v{l}{idx}{r}" if sz == 1 else f"v[{idx}:{idx + sz - 1}]"
+    if 256 <= o < 512: idx = o - 256; base = f"v{l}{idx}{r}" if sz == 1 else f"v[{idx}:{idx + sz - 1}]"
    elif o < 106: base = f"s{l}{o}{r}" if sz == 1 else f"s[{o}:{o + sz - 1}]"
    elif sz == 2 and o in self._PAIRS: base = self._PAIRS[o] if upper else self._PAIRS[o].lower()
    elif o in self._NAMES: base = self._NAMES[o] if upper else self._NAMES[o].lower()  # special regs (any sz)
-    elif 108 <= o < 124:
-      idx = o - 108
-      base = f"ttmp{l}{idx}{r}" if sz == 1 else f"ttmp[{idx}:{idx + sz - 1}]"
+    elif 108 <= o < 124: idx = o - 108; base = f"ttmp{l}{idx}{r}" if sz == 1 else f"ttmp[{idx}:{idx + sz - 1}]"
    elif 128 <= o <= 192: base = str(o - 128)  # inline int constants (0-64)
    elif 193 <= o <= 208: base = str(-(o - 192))  # inline negative int constants (-1 to -16)
    else: raise RuntimeError(f"unknown register: offset={o}, sz={sz}")
@@ -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}) "
-                         f"doesn't match range {self._valid_range} ({expected_size})")
+      raise RuntimeError(f"{self.__class__.__name__}: field size {hi - lo + 1} bits ({actual_size}) doesn't match range {self._valid_range} ({expected_size})")

  def encode(self, val) -> int:
    """Encode value. Returns 255 (literal marker) for out-of-range values."""
@@ -236,9 +231,9 @@ class VDSTYField(BitField):
 # ══════════════════════════════════════════════════════════════

 import functools
-from tinygrad.runtime.autogen.amd.rdna3.operands import OPERANDS as OPERANDS_RDNA3
-from tinygrad.runtime.autogen.amd.rdna4.operands import OPERANDS as OPERANDS_RDNA4
-from tinygrad.runtime.autogen.amd.cdna.operands import OPERANDS as OPERANDS_CDNA
+from extra.assembly.amd.autogen.rdna3.operands import OPERANDS as OPERANDS_RDNA3
+from extra.assembly.amd.autogen.rdna4.operands import OPERANDS as OPERANDS_RDNA4
+from extra.assembly.amd.autogen.cdna.operands import OPERANDS as OPERANDS_CDNA
 OPERANDS = {**OPERANDS_CDNA, **OPERANDS_RDNA3, **OPERANDS_RDNA4}

 # ══════════════════════════════════════════════════════════════
@@ -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')

@@ -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(key)):  # type: ignore[arg-type]
-        def _pv_val(pv: ET.Element) -> tuple[int, str]:
-          v, n = pv.findtext("Value"), pv.findtext("Name")
-          assert v is not None and n is not None
-          return int(v), n.upper()
-        enums[enum_name] = dict(_pv_val(pv) for pv in field.findall(".//PredefinedValue"))
+      if (enum_name := op_enum_map.get((ot_name, field.findtext("FieldName")))):
+        enums[enum_name] = {int(pv.findtext("Value")): pv.findtext("Name").upper() for pv in field.findall(".//PredefinedValue")}
  # Extract DataFormats with BitCount
  for df in root.findall("ISA/DataFormats/DataFormat"):
    name, bits = df.findtext("DataFormatName"), df.findtext("BitCount")
@@ -91,26 +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]] = []
-    for f in enc.findall(".//MicrocodeFormat/BitMap/Field"):
-      br = f.find("BitLayout/Range")
-      if br is None: continue
-      fn = f.findtext("FieldName")
-      assert fn is not None
-      fields.append((_norm_field(fn.lower()),
-        int(br.findtext("BitOffset") or 0) + int(br.findtext("BitCount") or 0) - 1, int(br.findtext("BitOffset") or 0)))
-    ident_list = enc.findall("EncodingIdentifiers/EncodingIdentifier")
-    ident = ident_list[0] if ident_list else None
+    fields = [(_norm_field(f.findtext("FieldName").lower()), int(f.find("BitLayout/Range").findtext("BitOffset") or 0) + int(f.find("BitLayout/Range").findtext("BitCount") or 0) - 1,
+               int(f.find("BitLayout/Range").findtext("BitOffset") or 0))
+              for f in enc.findall(".//MicrocodeFormat/BitMap/Field") if f.find("BitLayout/Range") is not None]
+    ident = (enc.findall("EncodingIdentifiers/EncodingIdentifier") or [None])[0]
    enc_field = next((f for f in fields if f[0] == "encoding"), None)
-    # For multi-dword formats, encoding field may be in higher dword but identifier is always in dword0; use % 32
-    enc_bits: str | None = None
-    if ident is not None and ident.text is not None and enc_field:
-      enc_bits = "".join(ident.text[len(ident.text)-1-b] for b in range(enc_field[1] % 32, (enc_field[2] % 32)-1, -1))
+    # For multi-dword formats, encoding field may be in higher dword but identifier pattern is always in dword0; use % 32
+    enc_bits = "".join(ident.text[len(ident.text)-1-b] for b in range(enc_field[1] % 32, (enc_field[2] % 32)-1, -1)) if ident is not None and enc_field else None
    base_name = _strip_enc(name)
    encodings[NAME_MAP.get(base_name, base_name)] = (fields, enc_bits)
  # Extract instruction opcodes and operand info
@@ -118,12 +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")
-      assert enc_enc_name is not None
-      base, opcode = _map_flat(_strip_enc(enc_enc_name), name), int(enc.findtext("Opcode") or 0)
+      base, opcode = _map_flat(_strip_enc(enc.findtext("EncodingName")), name), int(enc.findtext("Opcode") or 0)
      enc_name = NAME_MAP.get(base, base)
      # Encoding variants use the same Op enum as the base format
      base_enum = enc_name
@@ -137,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]] = {}
-      for op in enc.findall("Operands/Operand"):
-        fn = op.findtext("FieldName")
-        if fn: op_info[fn.lower()] = (op.findtext("DataFormatName"), int(op.findtext("OperandSize") or 0), op.findtext("OperandType"))
+      op_info = {op.findtext("FieldName").lower(): (op.findtext("DataFormatName"), int(op.findtext("OperandSize") or 0), op.findtext("OperandType"))
+                 for op in enc.findall("Operands/Operand") if op.findtext("FieldName")}
      for fmt, _, otype in op_info.values():
        if fmt and fmt not in fmts: fmts[fmt] = 0
        if otype: op_types_set.add(otype)
@@ -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)
-  assert xref_match is not None
-  pos = int(xref_match.group(1)) + 4
+  pos = int(re.search(rb'startxref\s+(\d+)', data).group(1)) + 4
  while data[pos:pos+7] != b'trailer':
    while data[pos:pos+1] in b' \r\n': pos += 1
    line_end = data.find(b'\n', pos)
@@ -185,19 +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)):
-        if sm.group(1): font = sm.group(1)
-        elif sm.group(2): x, y = x + float(sm.group(2)), y + float(sm.group(3))
-        elif sm.group(4): x, y = float(sm.group(4)), float(sm.group(5))
-        elif sm.group(6) and (t := bytes.fromhex(sm.group(6)).decode('latin-1')).strip():
-          elements.append((x, y, t, font))
-        elif sm.group(7):
-          t = ''.join(bytes.fromhex(h).decode('latin-1') for h in re.findall(r'<([0-9A-Fa-f]+)>', sm.group(7)))
-          if t.strip(): elements.append((x, y, t, font))
+      for m in re.finditer(r'(/F[\d.]+) [\d.]+ Tf|([\d.+-]+) ([\d.+-]+) Td|[\d.+-]+ [\d.+-]+ [\d.+-]+ [\d.+-]+ ([\d.+-]+) ([\d.+-]+) Tm|<([0-9A-Fa-f]+)>.*?Tj|\[([^\]]+)\] TJ', bt.group(1)):
+        if m.group(1): font = m.group(1)
+        elif m.group(2): x, y = x + float(m.group(2)), y + float(m.group(3))
+        elif m.group(4): x, y = float(m.group(4)), float(m.group(5))
+        elif m.group(6) and (t := bytes.fromhex(m.group(6)).decode('latin-1')).strip(): elements.append((x, y, t, font))
+        elif m.group(7) and (t := ''.join(bytes.fromhex(h).decode('latin-1') for h in re.findall(r'<([0-9A-Fa-f]+)>', m.group(7)))).strip(): elements.append((x, y, t, font))
    pages.append(sorted(elements, key=lambda e: (-e[1], e[0])))
  return pages

@@ -223,7 +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:
-  lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
+def write_common(all_fmts, all_op_types, path):
+  lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
  lines.append("class 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",
-    "from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType  # noqa: F401", ""]
+  lines = ["# autogenerated from AMD ISA XML - do not edit", "from extra.assembly.amd.autogen.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
-  enum_names = sorted(f"{k}Op" for k in enums if enums[k] and k not in ("HWREG", "MSG"))
-  # also re-export HWREG/MSG enums (plain enums, not instruction format ops)
-  enum_names += sorted(k for k in enums if k in ("HWREG", "MSG") and enums[k])
-  # collect DSL field types actually used by scanning generated field definitions
-  all_field_defs = " ".join(field_def(fn, hi, lo, enc, eb) for enc, (flds, eb) in encodings.items() for fn, hi, lo in flds)
-  _ALL_DSL = ["BitField", "EnumBitField", "FixedBitField", "NULL", "SBaseField", "SGPRField", "SRsrcField",
-              "SSrcField", "SrcField", "VDSTYField", "VGPRField"]
-  dsl_names = ["Inst"] + [n for n in _ALL_DSL if n in all_field_defs]
-  # also re-export register names so `from ins import *` still provides them to downstream users
-  _DSL_REGS = ["s", "v", "src", "VCC_LO", "VCC_HI", "VCC", "EXEC_LO", "EXEC_HI", "EXEC", "NULL", "OFF", "M0",
-               "SCC", "VCCZ", "EXECZ", "ttmp", "INV_2PI", "SDWA", "DPP", "DPP16", "LIT", "SRC_LDS_DIRECT"]
-  dsl_reexport = sorted(set(dsl_names + _DSL_REGS))
-  lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit", "# ruff: noqa: E501,F401",
-    f"from tinygrad.renderer.amd.dsl import {', '.join(dsl_reexport)}",
-    f"from tinygrad.runtime.autogen.amd.{arch}.enum import {', '.join(enum_names)}", "import functools", ""]
+  lines = ["# autogenerated from AMD ISA XML - do not edit", "# ruff: noqa: F401,F403",
+           "from extra.assembly.amd.dsl import *", f"from extra.assembly.amd.autogen.{arch}.enum import *", "import functools", ""]

  def fmt_allowed(op_enum: str, ops: set[int]) -> str:
    """Format allowed ops as {EnumName.MEMBER, ...}."""
@@ -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"),
-                       (f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]
-      for cls, seg, op_enum in flat_variants:
+      for cls, seg, op_enum in [(f"{prefix}FLAT", 0, f"{prefix}FLATOp"), (f"{prefix}GLOBAL", 2, f"{prefix}GLOBALOp"), (f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]:
        cls_ops = set(enums.get(cls, {}).keys())
        lines.append(f"class {cls}(Inst):")
        for fn, hi, lo in sort_fields(fields):
@@ -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
-  used_bases = {eb for (nm, eb) in types if (nm, eb) in valid}
-  enum_names = sorted(f"{k}Op" for k in used_bases)
-  lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit",
-    "from tinygrad.runtime.autogen.amd.common import Fmt, OpType",
-    f"from tinygrad.runtime.autogen.amd.{arch}.enum import {', '.join(enum_names)}", ""]
+  lines = ["# autogenerated from AMD ISA XML - do not edit",
+           "from extra.assembly.amd.autogen.common import Fmt, OpType",
+           f"from extra.assembly.amd.autogen.{arch}.enum import *", ""]
  lines.append("# instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}")
  lines.append("OPERANDS = {")
  def fmt_val(v):
@@ -470,7 +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] = {}
-  all_op_types: set[str] = set()
-  arch_data: dict[str, dict] = {}
+  import pathlib
+  all_fmts, all_op_types, arch_data = {}, set(), {}
  # First pass: parse XML for all architectures
  for arch, cfg in ARCHS.items():
    print(f"Parsing XML: {cfg['xml']} -> {arch}")
@@ -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 = autogen_base / "common.py"
+  common_path = pathlib.Path(__file__).parent / "autogen" / "common.py"
  write_common(all_fmts, all_op_types, common_path)
  print(f"Wrote common.py: {len(all_fmts)} formats, {len(all_op_types)} op types")
  # Write per-arch files from XML
  for arch, data in arch_data.items():
-    base = autogen_base / arch
+    base = pathlib.Path(__file__).parent / "autogen" / arch
    write_enum(data["enums"], base / "enum.py")
    write_ins(data["encodings"], data["enums"], data["suffix_only_ops"], data["types"], arch, base / "ins.py")
    write_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")
@@ -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)
-  # 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
+  return ((val >> _const(dt, lo)) if lo > 0 else val) & _const(val.dtype, (1 << (hi - lo + 1)) - 1)

 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):
-  t = UOp(Ops.TRUNC, x.dtype, (x,))
-  return ((x < _const(x.dtype, 0)) & x.ne(t)).where(t - _const(x.dtype, 1), t)
+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)
 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 \
-    (dtypes.uint16, dtypes.half) if val.dtype == dtypes.half else (dtypes.uint32, dtypes.float32)
+  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)
  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)) |
-    _const(dtypes.uint64, 0x3fe0000000000000)).bitcast(dtypes.float64)
+  return ((val.bitcast(dtypes.uint64) & _const(dtypes.uint64, 0x800FFFFFFFFFFFFF)) | _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(
-  "0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd"
-  "63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414"
-  "da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6", 16)
+TWO_OVER_PI = 0x0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6
 # 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_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),
-  '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),
+  '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_to_unorm': lambda a: _floor(_f16_extract(a).cast(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),
@@ -288,7 +274,6 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
  # 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 +310,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():
-      i += 1
-      continue
+    if c.isspace(): i += 1; continue
    if i + 1 < n and s[i:i+2] in ('+=', '-='):
-      tokens.append(Token('ASSIGN_OP', s[i:i+2]))
-      i += 2
-      continue
+      tokens.append(Token('ASSIGN_OP', s[i:i+2])); i += 2; continue
    if i + 1 < n and s[i:i+2] in ('||', '&&', '>=', '<=', '==', '!=', '<>', '>>', '<<', '**', '+:', '-:'):
-      tokens.append(Token('OP', s[i:i+2]))
-      i += 2
-      continue
-    if c in '|^&><+-*/~!%':
-      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
+      tokens.append(Token('OP', s[i:i+2])); i += 2; continue
+    if c in '|^&><+-*/~!%': 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 +337,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:
-          i += len(sfx)
-          break
-      tokens.append(Token('NUM', s[start:i]))
-      continue
+        if s[i:i+len(sfx)] == sfx: i += len(sfx); 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]))
-      continue
+      tokens.append(Token('IDENT', s[start:i])); 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):
-    self.tokens, self.vars, self.funcs, self.pos = tokens, env, funcs if funcs is not None else _FUNCS, 0
+  def __init__(self, tokens: list[Token], vars: dict, funcs: dict | None = None):
+    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:
-    tok = self.tokens[self.pos]
-    self.pos += 1
-    return tok
-  def eat(self, kind: str) -> Token:
-    if self.peek().type != kind: raise RuntimeError(f"expected {kind}, got {self.peek()}")
+  def _advance(self) -> Token: tok = self.tokens[self.pos]; self.pos += 1; return tok
+  def eat(self, type: str) -> Token:
+    if self.peek().type != type: raise RuntimeError(f"expected {type}, 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_val(self, val: str, kind: str) -> Token | None:
-    return self._advance() if self.peek().type == kind and self.peek().val == val else None
-  def eat_val(self, val: str, kind: str) -> Token:
-    if self.peek().type != kind or self.peek().val != val: raise RuntimeError(f"expected {kind}:{val}, got {self.peek()}")
+  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, type: str) -> Token | None:
+    return self._advance() if self.peek().type == type and self.peek().val == val else None
+  def eat_val(self, val: str, type: str) -> Token:
+    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 +380,9 @@ class Parser:
      case '||' | '|': return left | right
      case '&&' | '&': return left & right
      case '^': return left ^ right
-      case '==': return left.eq(right)
-      case '!=': return left.ne(right)
-      case '>=' | '<=' | '>' | '<' | '<>':
-        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.eq(right) if op == '==' else left.ne(right)
+      case '!=' : return left.ne(right)
+      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])
      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)
@@ -558,9 +520,7 @@ class Parser:
      self.eat('LBRACKET')
      self.eat_val('laneId', 'IDENT')
      self.eat('RBRACKET')
-      lane = self.vars['laneId']
-      shift = lane.cast(base.dtype) if base.dtype != dtypes.uint32 else _to_u32(lane)
-      result = (base >> shift) & _const(base.dtype, 1)
+      result = (base >> _to_u32(self.vars['laneId'])) & _u32(1)
      if self.try_eat('DOT'):
        dt_name = self.eat('IDENT').val
        return result.cast(DTYPES.get(dt_name, dtypes.uint32))
@@ -569,8 +529,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).bitcast(dt)
+      return (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16) if dt == dtypes.uint16 else (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16).bitcast(dt)
    if field == 'i4': return _signext_4bit(base)
    return _cast_to(base, dt)

@@ -580,7 +539,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 +625,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,
-            ('B',32): dtypes.uint32, ('B',64): dtypes.uint64}.get((type_char, bits), dtypes.uint64 if bits > 32 else dtypes.uint32)
+            ('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)
      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 +686,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 +712,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)
-        hi = mem.index(idx + _const(dtypes.int, 1), *gate, ptr=True).load().cast(dtypes.uint32)
-        val = (lo | (hi << _u32(8))).cast(dt)
+        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)
      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 +723,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):
-        val = (val >> (((addr >> _const(adt, 1)) & _const(adt, 1)).cast(dtypes.uint32) * _u32(16))) & _u32(0xFFFF)
+      elif dt in (dtypes.uint16, dtypes.int16): 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 +756,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:
-  return Parser(toks, env, funcs).parse()
+def parse_tokens(toks: list[Token], vars: dict[str, VarVal], funcs: dict | None = None) -> UOp:
+  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 +767,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 +781,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 +792,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':
-      i += 1
-      continue
+    if toks[0].type != 'IDENT' and toks[0].type != 'LBRACE': i += 1; continue
    first = toks[0].val.lower() if toks[0].type == 'IDENT' else '{'

    # Block terminators
@@ -855,19 +801,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':
-          p.eat('NUM')
-          p.eat('QUOTE')
+        if p.at('NUM') and p.peek(1).type == '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 +833,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(
-                  val, old.cast(val.dtype) if val.dtype != old.dtype and val.dtype.itemsize == old.dtype.itemsize else old)
+                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)
          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))
-                block_assigns[found_var] = env[found_var] = not_found.where(cond, found)
+                cond = _to_bool(parse_expr(cond_str, vars, funcs))
+                block_assigns[found_var] = vars[found_var] = not_found.where(cond, found)
                break
        else:
-          block_assigns.update(iter_assigns)
-          env.update(iter_assigns)
+          block_assigns.update(iter_assigns); vars.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))
-      i += 1
-      continue
+        vars.setdefault(toks[1].val, _u32(0))
+      i += 1; 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 +886,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:
-                  body_lines_lst.append(lines[i][:j])
-                  break
+                if depth == 0: 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':
-        compound_op = toks[j].val
-        j += 1
+      if j < len(toks) and toks[j].type == 'ASSIGN_OP': 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 +914,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
-      continue
+      i += 1; continue

    # VGPR assignment: VGPR[lane][reg] = value
    if first == 'vgpr' and toks[1].type == 'LBRACKET':
@@ -987,12 +923,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)
-        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)))
-        i += 1
-        continue
+        ln, rg, val = parse_tokens(lane_toks, vars, funcs), parse_tokens(reg_toks, vars, funcs), parse_tokens(toks[j:], vars, 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)))
+        i += 1; continue

    # Compound destination: {hi.type, lo.type} = value
    if first == '{':
@@ -1006,20 +939,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[hi_var] = env[hi_var] = hi_val
+          block_assigns[lo_var] = vars[lo_var] = lo_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
-          continue
+          i += 1; 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 \
-        (toks[1].type == 'LBRACKET' or (toks[1].type == 'DOT' and toks[3].type == 'LBRACKET')):
+    if len(toks) >= 5 and toks[0].type == 'IDENT' and (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 +967,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)
-          old = block_assigns.get(var, env.get(var))
-          assert isinstance(old, UOp)
-          block_assigns[var] = env[var] = _set_bits(old, _val_to_bits(val), hi - lo + 1, lo)
-          i += 1
-          continue
-        except Exception: pass
+          if var not in vars: vars[var] = _const(dtypes.uint64 if hi >= 32 else dtypes.uint32, 0)
+          old = block_assigns.get(var, vars.get(var))
+          block_assigns[var] = vars[var] = _set_bits(old, _val_to_bits(val), hi - lo + 1, lo)
+          i += 1; continue
+        except: pass
      elif toks[1].type == 'LBRACKET':  # bit index: var[expr] (only for var[...], not var.type[...])
-        existing = block_assigns.get(var, env.get(var))
-        if existing is not None and isinstance(existing, UOp) and \
-            not any(f'{var}{k}' in env or f'{var}{k}' in block_assigns for k in range(8)):
+        existing = block_assigns.get(var, vars.get(var))
+        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)):
          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(
-              existing, _to_u32(parse_tokens(bit_toks, env, funcs)), parse_tokens(toks[j+1:], env, funcs))
-            i += 1
-            continue
+            block_assigns[var] = vars[var] = _set_bit(existing, _to_u32(parse_tokens(bit_toks, vars, funcs)), parse_tokens(toks[j+1:], vars, funcs))
+            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 +993,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)
-          existing = block_assigns.get(var, env.get(var))
+          val = parse_tokens(toks[j+1:], vars, funcs)
+          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
-          i += 1
-          continue
+            block_assigns[f'{var}@{idx}'] = vars[f'{var}@{idx}'] = val
+          i += 1; 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)
-          val = parse_tokens(toks[j+1:], env, funcs)
+          idx_expr = parse_tokens(idx_toks, vars, 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
-          i += 1
-          continue
+            block_assigns[f'{var}@{elem_idx}'] = vars[f'{var}@{elem_idx}'] = new_val
+          i += 1; 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)))
-      rhs = parse_tokens(toks[assign_op+1:], env, funcs)
+      old = block_assigns.get(var, vars.get(var, _u32(0)))
+      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)
-      i += 1
-      continue
+      block_assigns[var] = vars[var] = (old + rhs) if toks[assign_op].val == '+=' else (old - rhs)
+      i += 1; continue

    # Typed element: var.type[idx] = value
-    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':
+    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':
      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
-        continue
+        i += 1; continue

    # Dynamic bit: var.type[expr_with_brackets] = value
-    if len(toks) >= 5 and toks[0].type == 'IDENT' and toks[1].type == 'DOT' and \
-        toks[2].type == 'IDENT' and toks[3].type == 'LBRACKET':
+    if len(toks) >= 5 and toks[0].type == 'IDENT' and toks[1].type == 'DOT' and 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':
-          depth += 1
-          has_inner = True
+        if toks[j].type == 'LBRACKET': 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)
-          old = block_assigns.get(var, env.get(var, _u32(0)))
-          block_assigns[var] = env[var] = _set_bit(old, bit_pos, val)
-          i += 1
-          continue
+          val = parse_tokens(toks[j+1:], vars, funcs)
+          old = block_assigns.get(var, vars.get(var, _u32(0)))
+          block_assigns[var] = vars[var] = _set_bit(old, bit_pos, val)
+          i += 1; 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()
-      env.update(env_snap)
+      vars.clear(); vars.update(vars_snap)
      while i < len(lines):
        ltoks = tokenize(lines[i])
        if ltoks[0].type != 'IDENT': break
@@ -1158,22 +1074,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, branch, ret = parse_block(lines, i, env, funcs, assigns if take else None)
+          i += 1; i, branch, ret = parse_block(lines, i, vars, 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()
-          env.update(env_snap)
+          vars.clear(); vars.update(vars_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()
-          env.update(env_snap)
-        elif lf == 'endif':
-          i += 1
-          break
+          vars.clear(); vars.update(vars_snap)
+        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 +1097,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)
-        env.update(ba)
+        block_assigns.update(ba); vars.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))))
-          for cond, ba in reversed(conditions):  # type: ignore[assignment]
+          res: Any = else_assigns.get(var, block_assigns.get(var, vars.get(var, _u32(0))))
+          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 +1116,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)
-        i += 1
-        break
+        block_assigns[base_var] = vars[base_var] = parse_tokens(toks[j+1:], vars, funcs)
+        i += 1; break
    else: i += 1
  return i, block_assigns, None

-def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
-  return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
+def parse_expr(expr: str, vars: dict[str, VarVal], funcs: dict | None = None) -> UOp:
+  return parse_tokens(tokenize(expr.strip().rstrip(';')), vars, funcs)

@@ -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 tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm # same encoding as RDNA4
+from extra.assembly.amd.dsl import BitField, FixedBitField, bits

 # ═══════════════════════════════════════════════════════════════════════════════
 # 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,41 +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
-  JUMP_UNCOND = 0x5
-  MESSAGE = 0x9
-  VALU_TRANS = 0xb
-  VALU_B2 = 0xd
-  VALU_B4 = 0xe
+  UNK_06 = 0x6
+  VMEM = 0x10
+  UNK_11 = 0x11
  VINTERP = 0x12
-  VMEM_RD_1 = 0x21
-  VMEM_RD_2 = 0x22
-  VMEM_WR_1 = 0x23
-  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
-  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
+  UNK_14 = 0x14
+  OTHER_VMEM = 0x5e
+  UNK_60 = 0x60

 # ═══════════════════════════════════════════════════════════════════════════════
 # PACKET TYPE BASE CLASS
@@ -148,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._size_nibbles = ((max((f.hi for f in cls._fields.values()), default=0) + 4) // 4)  # 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)

  @classmethod
  def from_raw(cls, raw: int, time: int = 0):
@@ -158,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})"

 # ═══════════════════════════════════════════════════════════════════════════════
@@ -170,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]
@@ -192,7 +171,7 @@ class TS_DELTA_OR_MARK(PacketType):
  @property
  def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)

-class TS_DELTA_OR_MARK_RDNA4(PacketType):  # Layout 4: 48->64 bits
+class TS_DELTA_OR_MARK_L4(PacketType):  # Layout 4: 48->64 bits
  encoding = bits[6:0] == 0b0000001
  delta = bits[63:12]
  bit7 = bits[7:7]
@@ -206,7 +185,7 @@ class TS_DELTA_S5_W2(PacketType):
  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]
-  wave = bits[11:7]
-  op = bits[19:12].enum(InstOpRDNA4)
+  flag1 = bits[6:6]
+  flag2 = bits[7:7]
+  wave = bits[12:8]
+  op = bits[19:13].enum(InstOpL4)

 class UTILCTR(PacketType):
  encoding = bits[6:0] == 0b0110001
@@ -369,179 +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_RDNA3,
-  9: WAVESTART_RDNA4, 10: TS_DELTA_S5_W2_RDNA4, 11: WAVEALLOC_RDNA4,
-  12: TS_DELTA_S5_W3_RDNA4, 13: PERF_RDNA4, 22: TS_DELTA_OR_MARK_RDNA4, 24: INST_RDNA4,
+PACKET_TYPES_L4: dict[int, type[PacketType]] = {
+  **PACKET_TYPES_L3,
+  7: TS_DELTA_S8_W3_L4, 9: WAVESTART_L4, 10: TS_DELTA_S5_W2_L4, 11: WAVEALLOC_L4,
+  12: TS_DELTA_S5_W3_L4, 13: PERF_L4, 22: TS_DELTA_OR_MARK_L4, 24: INST_L4,
 }
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# CDNA PACKET TYPE DEFINITIONS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-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_7(PacketType):
-  """pkt_fmt=7: 16-bit packet"""
-  encoding = bits[3:0] == 7
-  unk_padding = bits[15:4]
-
-class CDNA_PKT_8(PacketType):
-  """pkt_fmt=8: 16-bit packet"""
-  encoding = bits[3:0] == 8
-  unk_padding = bits[15:4]
-
-class CDNA_PKT_9(PacketType):
-  """pkt_fmt=9: 16-bit packet"""
-  encoding = bits[3:0] == 9
-  unk_padding = bits[15:4]
-
-class CDNA_PKT_12(PacketType):
-  """pkt_fmt=12: 48-bit packet"""
-  encoding = bits[3:0] == 12
-  unk_padding = bits[47:4]
-
-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]
-
-PACKET_TYPES_CDNA: 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,
-  7: CDNA_PKT_7, 8: CDNA_PKT_8, 9: CDNA_PKT_9, 10: CDNA_EXEC, 11: CDNA_PKT_11, 12: CDNA_PKT_12,
-  13: CDNA_INST, 14: CDNA_PKT_14, 15: CDNA_PKT_15,
-}
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# 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_DELTA (*4), 4=CDNA_TIMESTAMP (absolute)
-  _special = {TS_DELTA_OR_MARK: 1, TS_DELTA_OR_MARK_RDNA4: 1, TS_DELTA_SHORT: 2, CDNA_DELTA: 3, CDNA_TIMESTAMP: 4}
+  # special_case: 0=none, 1=TS_DELTA_OR_MARK (check is_marker), 2=TS_DELTA_SHORT (add 8)
  decode_info = {}
  for opcode, pkt_cls in packet_types.items():
    delta_field = getattr(pkt_cls, 'delta', None)
-    special = _special.get(pkt_cls, 0)
-    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]
+    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)
  return decode_info, state_table

-_DECODE_INFO_RDNA3, _STATE_TABLE_RDNA3 = _build_decode_tables(PACKET_TYPES_RDNA3)
-_DECODE_INFO_RDNA4, _STATE_TABLE_RDNA4 = _build_decode_tables(PACKET_TYPES_RDNA4)
-_DECODE_INFO_CDNA, _STATE_TABLE_CDNA = _build_decode_tables(PACKET_TYPES_CDNA)
+_DECODE_INFO_L3, _STATE_TABLE_L3 = _build_decode_tables(PACKET_TYPES_L3)
+_DECODE_INFO_L4, _STATE_TABLE_L4 = _build_decode_tables(PACKET_TYPES_L4)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# DECODER
+# ═══════════════════════════════════════════════════════════════════════════════

 def decode(data: bytes) -> Iterator[PacketType]:
-  """Decode raw SQTT blob, yielding packet instances. Auto-detects RDNA (layout 3/4) vs CDNA."""
-  n, reg, pos, nib_off, nib_count, time, ts_offset = len(data), 0, 0, 0, 16, 0, None
-  decode_info, state_table = _DECODE_INFO_RDNA3, _STATE_TABLE_RDNA3  # start RDNA3, auto-detect switches if needed
+  """Decode raw SQTT blob, yielding packet instances. Auto-detects layout from LAYOUT_HEADER."""
+  n, reg, pos, nib_off, nib_count, time = len(data), 0, 0, 0, 16, 0
+  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
@@ -549,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 + read_bytes], 'little')
-      reg, pos = (reg >> (read_bytes * 8)) | (chunk << (64 - read_bytes * 8)), pos + byte_count
+      chunk = int.from_bytes(data[pos:pos + byte_count], 'little')
+      reg, pos = (reg >> (byte_count * 8)) | (chunk << (64 - byte_count * 8)), pos + byte_count
    # 3. if odd, read low nibble
    if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)

@@ -562,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)
-    if pkt_cls is LAYOUT_HEADER:
-      if pkt.layout == 4: decode_info, state_table = _DECODE_INFO_RDNA4, _STATE_TABLE_RDNA4
-      elif pkt.layout != 3:  # not a real LAYOUT_HEADER — switch to CDNA and re-decode first packet
-        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)
+    # detect layout from first LAYOUT_HEADER and switch decode tables if needed
+    # NOTE: CDNA uses a completely different 16-bit header format, not nibbles - not supported here
+    if pkt_cls is LAYOUT_HEADER and pkt.layout == 4:
+      decode_info, state_table = _DECODE_INFO_L4, _STATE_TABLE_L4
    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
 # ═══════════════════════════════════════════════════════════════════════════════
@@ -653,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)):
-    op_name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4)) 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 "")
+  if isinstance(p, (INST, INST_L4)):
+    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 "")
  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:
-    p, inst = data if isinstance(data, tuple) else (data, None)
-    if type(p).__name__.replace("_RDNA4", "") not in skip: print(format_packet(p), f"inst={inst.inst}" if inst is not None else '')
+  for p in packets:
+    if type(p).__name__.replace("_L4", "") not in skip: print(format_packet(p))

 if __name__ == "__main__":
  import sys, pickle
-  from tinygrad.helpers import temp
-  with open(temp("profile.pkl", append_user=True) if len(sys.argv) < 2 else sys.argv[1], "rb") as f:
+  if len(sys.argv) < 2:
+    print("Usage: python sqtt.py <pkl_file>")
+    sys.exit(1)
+  with open(sys.argv[1], "rb") as f:
    data = pickle.load(f)
-  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} {prg.name if prg is not None else ''} ===")
-    print_packets(map_insts(event.blob, prg.lib, dev_targets[prg.device]) if prg is not None else decode(event.blob))
+    print(f"\n=== event {i} ===")
+    print_packets(decode(event.blob))
@@ -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}")
@@ -0,0 +1,122 @@
+# maps SQTT trace packets to instructions.
+from dataclasses import dataclass
+from typing import Iterator
+
+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.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, target:int) -> 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):
+      assert p.wave not in wave_pc, "only one inflight wave per unit"
+      wave_pc[p.wave] = next(iter(pc_map))
+      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 amd_decode
+  from extra.sqtt.roc import decode as roc_decode, InstExec
+  addr_table = amd_decode(prg.lib, target)
+  disasm = {addr+prg.base:(inst.disasm(), inst.size()) for addr,inst in addr_table.items()}
+  rctx = roc_decode([sqtt], {prg.tag:disasm})
+  rwaves = rctx.inst_execs.get((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())
+
+  passed_insts = 0
+  for pkt, info in map_insts(sqtt.blob, prg.lib, target):
+    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-prg.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}"
+
+  if len(rwaves):
+    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.tag:e for e in data if type(e).__name__ == "ProfileProgramEvent"}
+  target = next((e for e in data if type(e).__name__ == "ProfileDeviceEvent" and e.device.startswith("AMD"))).props["gfx_target_version"]
+  for e in sqtt_events:
+    if args.kernel is not None and args.kernel != e.kern: continue
+    if not e.itrace: continue
+    print(f"==== {e.kern}")
+    test_rocprof_inst_traces_match(e, kern_events[e.kern], target)
@@ -0,0 +1,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()
@@ -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]
@@ -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 tinygrad.renderer.amd.dsl import NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, M0
+from extra.assembly.amd.emu import run_asm
+from extra.assembly.amd.dsl import NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, M0

 def _i32(f: float) -> int: return struct.unpack('<I', struct.pack('<f', f))[0]
 def _f32(i: int) -> float: return struct.unpack('<f', struct.pack('<I', i & 0xFFFFFFFF))[0]
@@ -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)
-def _out_bytes(n_lanes: int) -> int: return max(N_VGPRS * n_lanes * 4, _VGPR_REGION) + SGPR_BYTES + 12
-OUT_BYTES = _out_bytes(WAVE_SIZE)  # default for single-wave (backward compat)
+OUT_BYTES = VGPR_BYTES + SGPR_BYTES + 12  # + vcc + scc + exec

 # Float conversion helpers
 def f2i(f: float) -> int: return _i32(f)
@@ -76,10 +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):
-    self.vgpr = [[0] * 256 for _ in range(n_lanes)]  # vgpr[lane][reg]
+  def __init__(self):
+    self.vgpr = [[0] * 256 for _ in range(32)]  # vgpr[lane][reg]
    self.sgpr = [0] * 128
    self.vcc = 0
    self.scc = 0
@@ -103,53 +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], vgpr_bytes))
+  epilogue.append(v_mov_b32_e32(v[240], 0))
  for i in range(N_SGPRS):
    epilogue.append(v_mov_b32_e32(v[243], s[i]))
-    epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=i * 4))
+    epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + i * 4))
  epilogue.append(v_mov_b32_e32(v[243], s[90]))
-  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=SGPR_BYTES))
+  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + SGPR_BYTES))
  epilogue.append(v_mov_b32_e32(v[243], s[91]))
-  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=SGPR_BYTES + 4))
+  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + SGPR_BYTES + 4))
  # Store EXEC (saved earlier in s[95])
  epilogue.append(v_mov_b32_e32(v[243], s[95]))
-  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=SGPR_BYTES + 8))
+  epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + SGPR_BYTES + 8))
  epilogue.append(s_mov_b32(EXEC_LO, s[94]))
  epilogue.append(s_endpgm())
  return prologue, epilogue

 def parse_output(out_buf: bytes, n_lanes: int) -> WaveState:
  """Parse output buffer into WaveState."""
-  vgpr_bytes = N_VGPRS * n_lanes * 4
-  st = WaveState(n_lanes)
+  st = WaveState()
  for i in range(N_VGPRS):
    for lane in range(n_lanes):
-      off = i * n_lanes * 4 + lane * 4
+      off = i * WAVE_SIZE * 4 + lane * 4
      st.vgpr[lane][i] = struct.unpack_from('<I', out_buf, off)[0]
  for i in range(N_SGPRS):
-    st.sgpr[i] = struct.unpack_from('<I', out_buf, vgpr_bytes + i * 4)[0]
-  st.vcc = struct.unpack_from('<I', out_buf, vgpr_bytes + SGPR_BYTES)[0]
-  st.scc = struct.unpack_from('<I', out_buf, vgpr_bytes + SGPR_BYTES + 4)[0]
+    st.sgpr[i] = struct.unpack_from('<I', out_buf, VGPR_BYTES + i * 4)[0]
+  st.vcc = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES)[0]
+  st.scc = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES + 4)[0]
  # Store EXEC in its proper location (index 126)
-  st.sgpr[EXEC_LO.offset] = struct.unpack_from('<I', out_buf, vgpr_bytes + SGPR_BYTES + 8)[0]
+  st.sgpr[EXEC_LO.offset] = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES + 8)[0]
  return st

 def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
  """Run instructions via emulator run_asm, dump state to memory, return WaveState."""
-  buf_sz = _out_bytes(n_lanes)
-  out_buf = (ctypes.c_uint8 * buf_sz)(*([0] * buf_sz))
+  out_buf = (ctypes.c_uint8 * OUT_BYTES)(*([0] * OUT_BYTES))
  out_addr = ctypes.addressof(out_buf)

  prologue, epilogue = get_prologue_epilogue(n_lanes)
@@ -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(buf_sz)
+  out_gpu = dev.allocator.alloc(OUT_BYTES)
  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
@@ -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()
@@ -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."""
@@ -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."""
@@ -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."""
@@ -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
@@ -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."""
@@ -932,76 +932,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()
@@ -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()
@@ -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."""
@@ -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
@@ -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."""
@@ -418,7 +418,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 +442,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 +471,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 +502,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 +518,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 +540,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 +569,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 +586,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)))
@@ -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 = [
@@ -7,8 +7,9 @@ VOPD executes two operations simultaneously. Key behavior:
 - Op Y can use ops 0-18 (includes ADD_NC_U32, LSHLREV, AND)
 """
 import unittest
-from test.amd.hw.helpers import run_program, v, v_mov_b32_e32
-from tinygrad.runtime.autogen.amd.rdna3.ins import VOPD, VOPD_LIT, VOPDOp
+from extra.assembly.amd.test.hw.helpers import run_program, run_program_emu, run_program_hw, compare_wave_states, \
+  v, s, v_mov_b32_e32, s_mov_b32
+from extra.assembly.amd.autogen.rdna3.ins import VOPD, VOPD_LIT, VOPDOp

 class TestVOPDBasic(unittest.TestCase):
  """Basic VOPD functionality tests."""
@@ -108,7 +109,7 @@ class TestVOPDLiterals(unittest.TestCase):
    Tests that the 32-bit literal (SIMM32) is correctly passed to the instruction.
    fma(2.0, 3.0, 10.0) = 2*3 + 10 = 16.0
    """
-    from test.amd.hw.helpers import f2i, i2f
+    from extra.assembly.amd.test.hw.helpers import f2i, i2f
    instructions = [
      v_mov_b32_e32(v[0], f2i(2.0)),  # v[0] = 2.0
      v_mov_b32_e32(v[1], f2i(3.0)),  # v[1] = 3.0
@@ -126,7 +127,7 @@ class TestVOPDLiterals(unittest.TestCase):
    Tests that the 32-bit literal (SIMM32) is correctly used as the multiplier.
    fma(2.0, 5.0, 3.0) = 2*5 + 3 = 13.0
    """
-    from test.amd.hw.helpers import f2i, i2f
+    from extra.assembly.amd.test.hw.helpers import f2i, i2f
    instructions = [
      v_mov_b32_e32(v[0], f2i(2.0)),  # v[0] = 2.0
      v_mov_b32_e32(v[1], f2i(3.0)),  # v[1] = 3.0
@@ -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 tinygrad.renderer.amd import decode_inst
-import tinygrad
-REMU_PATH = Path(tinygrad.__file__).parent.parent / "extra/remu/target/release/libremu.so"
-if not REMU_PATH.exists(): REMU_PATH = Path(tinygrad.__file__).parent.parent / "extra/remu/target/release/libremu.dylib"
+from extra.assembly.amd.emu import WaveState, decode_program, WAVE_SIZE, VCC_LO, EXEC_LO, SCC
+from extra.assembly.amd import decode_inst
+from extra.assembly.amd.test.helpers import KernelInfo
+from extra.assembly.amd.test.bench_emu import REMU_PATH

 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:
-      self.lib.wave_free(self.ctx)
-      self.ctx = None
+    if self.ctx: self.lib.wave_free(self.ctx); self.ctx = None

 class PythonEmulator:
  def __init__(self):
    self.state: WaveState | None = None
-    self.program: dict[int, tuple] = {}  # lazily populated: pc -> (name, fxn, globals)
+    self.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
-    self._ensure_decoded(pc)
-    name, fxn, globals_list = self.program[pc]
-    buf_addrs = {0: self.state.sgpr_buf._buf.va_addr, 1: self.state.vgpr_buf._buf.va_addr,  # type: ignore[union-attr]
-                 2: self.vmem_buf._buf.va_addr, 3: self.lds_buf._buf.va_addr}  # type: ignore[union-attr]
+    if pc == 0xFFFFFFFFFFFFFFFF or pc not in self.program: return -1
+    name, fxn, globals_list, _runner = 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,
+                 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
-            python._ensure_decoded(pc_addr)
-            inst_hex_name = python.program[pc_addr][0]
+            inst_info = python.program.get(python.lib_addr + python_before.pc * 4)  # Convert word offset to actual address
+            inst_hex_name = inst_info[0] if inst_info else f"unknown at PC={python_before.pc}"
            # 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  "
-              msg += "\n  ".join(diffs[:10]) + f"\n  Recent instructions:\n{trace_str}"
-              return False, msg, total_steps
+              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ (rust vs python):\n  " + "\n  ".join(diffs[:10]) + f"\n  Recent instructions:\n{trace_str}", total_steps

            rust_result = rust.step()
            python_result = python.step()
@@ -258,9 +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: "
-                     f"rust={rust_result}, python={python_result}, inst={inst_str}\n  Recent instructions:\n{trace_str}")
-              return False, msg, total_steps
+              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: rust={rust_result}, python={python_result}, inst={inst_str}\n  Recent instructions:\n{trace_str}", total_steps

            # Sync Python state to Rust after instructions with known Rust emulator differences
            if sync_after:
@@ -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):
-    self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4)) + T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4)))
+  def test_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4)) + T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4)))

  # Convolution
  def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 2, 8, 8).conv2d(T.empty(2, 2, 3, 3)), max_steps=50000)
@@ -1,12 +1,22 @@
 import unittest
+import functools
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
+from tinygrad.runtime.support.compiler_amd import HIPCompiler

-from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.renderer.amd.dsl import s, v
+from extra.assembly.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.dsl import s, v, Inst

-def custom_add_one(A:UOp) -> UOp:
+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")
@@ -21,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)))
-  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]))))
+  sink = UOp.sink(A.base, threads, arg=KernelInfo(name:=f"custom_add_one_{A.size}", estimates=Estimates(ops=A.size, mem=A.size*4*2)))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=(*sink.src, sink)), *assemble_insts(insts, name, arch)))

-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")
@@ -40,14 +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}"))
-  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]))))
+  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=(*sink.src, sink)),
+                               *assemble_insts(insts, name, arch, kernarg_size=16)))

-@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestCustomKernel(unittest.TestCase):
  def test_simple(self):
    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)
@@ -57,7 +67,7 @@ class TestCustomKernel(unittest.TestCase):
  def test_variable(self):
    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})
@@ -1,8 +1,8 @@
 import unittest
-from tinygrad.renderer.amd.dsl import *
-from tinygrad.renderer.amd.dsl import VDSTYField
-from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op
-from tinygrad.runtime.autogen.amd.rdna3.ins import VOP1
+from extra.assembly.amd.dsl import *
+from extra.assembly.amd.dsl import VDSTYField
+from extra.assembly.amd.autogen.rdna3.enum import VOP1Op, VOP2Op
+from extra.assembly.amd.autogen.rdna3.ins import VOP1

 class TestRegisters(unittest.TestCase):
  def test_vgpr_single(self):
@@ -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 test.mockgpu.amd.pcode import parse_expr
-from tinygrad.runtime.autogen.amd.rdna3.str_pcode import PCODE
-from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op, SOP2Op, DSOp
+from extra.assembly.amd.emu import parse_pcode
+from extra.assembly.amd.pcode import parse_expr
+from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
+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)}
-    result = parse_expr('x', vrs)
+    vars = {'x': UOp.const(dtypes.uint32, 42)}
+    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)}
-    result = parse_expr('cond ? a : b', vrs)
+    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', 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,7 +174,6 @@ 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)

@@ -185,9 +182,9 @@ class TestDSPcodePatterns(unittest.TestCase):
    # Create a mock LDS buffer
    lds = UOp(Ops.PARAM, 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,13 +207,12 @@ 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.PARAM, dtypes.uint32.ptr(16384), arg=3)
    srcs = {
      'ADDR': 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
-    self.assertEqual(assigns[0][1][0].simplify().arg, 108)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][0].simplify().arg, 120)  # type: ignore[index]
+    addr0, _ = assigns[0][1]
+    addr1, _ = assigns[1][1]
+    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(assigns[1][1][1].simplify().arg, 0xBBBBBBBB)  # type: ignore[index]
+    self.assertEqual(val0.simplify().arg, 0xAAAAAAAA)
+    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'))
@@ -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(),
-            'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(),
+            'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'ROUND_TOWARD_ZERO': 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:
-        skipped += 1
-        errors[str(e)].append(op.name)
+      except RuntimeError as e: 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__":
@@ -4,10 +4,10 @@
 Note: Graphics-only formats (EXP, MUBUF, MTBUF, MIMG) are not supported - use GLOBAL/FLAT for memory access in compute.
 """
 import unittest
-from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.renderer.amd.dsl import VCC_HI, EXEC_LO, NULL
+from extra.assembly.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.dsl import VCC_HI, EXEC_LO, NULL
 OFF = NULL  # OFF is alias for NULL
-from tinygrad.renderer.amd import detect_format
+from extra.assembly.amd import detect_format


 class TestDS(unittest.TestCase):
@@ -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 tinygrad.renderer.amd.dsl import Inst
-from test.amd.test_roundtrip import compile_asm
-from test.amd.disasm import disasm
+from extra.assembly.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.dsl import Inst
+from extra.assembly.amd.test.test_roundtrip import compile_asm

 class IntegrationTestBase(unittest.TestCase):
  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")
@@ -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()
@@ -8,11 +8,11 @@ Only compute-relevant instruction formats are tested. Graphics-only formats not
  - VIMAGE/VSAMPLE: image sampling instructions (RDNA4)
  - VBUFFER: buffer instructions (RDNA4)
 """
-import unittest, re, functools
+import unittest, re, subprocess, functools
 from tinygrad.helpers import fetch
-from test.amd.disasm import disasm
-from tinygrad.renderer.amd import decode_inst, detect_format
-from test.amd.helpers import llvm_assemble, llvm_filter_valid_asm, get_target, get_mattr
+from extra.assembly.amd.disasm import disasm
+from extra.assembly.amd import decode_inst, detect_format
+from extra.assembly.amd.test.helpers import get_llvm_mc, get_target, get_mattr

 LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-21.1.0/llvm/test/MC/AMDGPU"

@@ -40,7 +40,7 @@ RDNA4_FILES = ['gfx12_asm_sop1.s', 'gfx12_asm_sop2.s', 'gfx12_asm_sopp.s', 'gfx1
  'gfx12_asm_vop1.s', 'gfx12_asm_vop2.s', 'gfx12_asm_vopc.s', 'gfx12_asm_vopcx.s', 'gfx12_asm_vop3.s', 'gfx12_asm_vop3c.s',
  'gfx12_asm_vop3cx.s', 'gfx12_asm_vop3p.s', 'gfx12_asm_vop3_from_vop1.s', 'gfx12_asm_vop3_from_vop2.s',
  'gfx12_asm_vop3p_features.s', 'gfx12_asm_vopd.s', 'gfx12_asm_vopd_features.s',
-  'gfx12_asm_ds.s', 'gfx12_asm_smem.s', 'gfx12_asm_vflat.s',
+  'gfx12_asm_ds.s', 'gfx12_asm_smem.s',
  'gfx12_asm_wmma_w32.s']

 def _parse_llvm_tests(text: str, pattern: str) -> list[tuple[str, bytes]]:
@@ -74,13 +74,42 @@ def _get_tests_uncached(f: str, arch: str) -> list[tuple[str, bytes]]:
  # Exclude v_interp_* (graphics-only, not on CDNA)
  if arch == "cdna": tests = [(asm, data) for asm, data in tests if not asm.startswith('v_interp_')]
  # Filter out tests where original ASM isn't valid on target (e.g., gfx9 tests with gfx942/gfx950 constraints)
-  if arch == "cdna" and not ('gfx942' in f or 'gfx950' in f or 'gfx90a' in f):
-    tests = llvm_filter_valid_asm(tests, get_target(arch), get_mattr(arch))
+  if arch == "cdna" and not ('gfx942' in f or 'gfx950' in f or 'gfx90a' in f): tests = _filter_valid_asm(tests, arch)
  return tests

@functools.cache
 def _get_tests(f: str, arch: str) -> list[tuple[str, bytes]]: return _get_tests_uncached(f, arch)

+def _compile_asm_batch(instrs: list[str], arch: str = "rdna3", mcpu: str|None = None) -> list[bytes]:
+  if not instrs: return []
+  mcpu, mattr = mcpu or get_target(arch), get_mattr(arch)
+  result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', f'-mattr={mattr}', '-show-encoding'],
+    input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True, timeout=30)
+  if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
+  return [bytes.fromhex(line.split('encoding:')[1].strip()[1:-1].replace('0x', '').replace(',', '').replace(' ', ''))
+          for line in result.stdout.split('\n') if 'encoding:' in line]
+
+def _filter_valid_asm(tests: list[tuple[str, bytes]], arch: str) -> list[tuple[str, bytes]]:
+  """Filter out tests where the original ASM isn't valid on the target (e.g., gfx9 tests with gfx942/gfx950 constraints)."""
+  if not tests: return []
+  mcpu = get_target(arch)
+  # Batch assemble all instructions, parse stderr to find which lines failed
+  instrs = [asm for asm, _ in tests]
+  result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', '-show-encoding'],
+    input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True, timeout=30)
+  # Parse error lines from stderr (format: "<stdin>:N:..." where N is 1-indexed, line 1 is ".text")
+  failed_lines = set()
+  for line in result.stderr.split('\n'):
+    if m := re.match(r'<stdin>:(\d+):', line): failed_lines.add(int(m.group(1)) - 1)  # -1 for .text, so line 2 -> index 1 -> tests[0]
+  # Also filter out tests where LLVM roundtrip doesn't match original (reserved bits set in original)
+  valid = [(asm, data) for i, (asm, data) in enumerate(tests) if (i + 1) not in failed_lines]
+  if not valid: return []
+  llvm_result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', '-show-encoding'],
+    input=".text\n" + "\n".join(asm for asm, _ in valid) + "\n", capture_output=True, text=True, timeout=30)
+  llvm_bytes = [bytes.fromhex(line.split('encoding:')[1].strip()[1:-1].replace('0x', '').replace(',', '').replace(' ', ''))
+                for line in llvm_result.stdout.split('\n') if 'encoding:' in line]
+  return [(asm, data) for (asm, data), lb in zip(valid, llvm_bytes) if lb == data]
+
 def _make_test(f: str, arch: str, test_type: str):
  def test(self):
    tests = _get_tests(f, arch)
@@ -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]
-      elif arch == "rdna4": import tinygrad.runtime.autogen.amd.rdna4.ins as ins  # type: ignore[no-redef]
-      elif arch == "cdna": import tinygrad.runtime.autogen.amd.cdna.ins as ins  # type: ignore[no-redef]
+      if arch == "rdna3": import extra.assembly.amd.autogen.rdna3.ins as ins
+      elif arch == "rdna4": import extra.assembly.amd.autogen.rdna4.ins as ins
+      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:
-            skipped += 1
-            continue  # skip if binary roundtrip fails
+          if decoded.to_bytes()[:len(data)] != data: skipped += 1; continue  # skip if binary roundtrip fails
          r = repr(decoded)
          try:
            decoded2 = eval(r, ns)  # noqa: S307
@@ -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)
@@ -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")
@@ -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 + 8U].b32 = VDATA[95 : 64];\nMEM[ADDR + 12U].b32 = VDATA[127 : 96]')
+      '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]')
    # 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")
@@ -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()
@@ -0,0 +1,98 @@
+import unittest, ctypes
+from extra.assembly.amd.autogen.rdna4 import ins as ir4
+from extra.assembly.amd.dsl import v, s
+from extra.assembly.amd.emu import WaveState, decode_program
+from tinygrad.device import Buffer, BufferSpec
+from tinygrad.dtype import dtypes
+
+class TestRDNA4Emu(unittest.TestCase):
+  def _run(self, insts: list, sgprs: dict[int, int] = None, vgprs: dict[tuple[int, int], int] = None) -> WaveState:
+    """Run instructions and return final WaveState."""
+    # Add S_ENDPGM if not present
+    if not any(isinstance(i, ir4.SOPP) and i.op == ir4.SOPPOp.S_ENDPGM for i in insts):
+      insts = list(insts) + [ir4.SOPP(ir4.SOPPOp.S_ENDPGM, simm=0)]
+
+    # Assemble and decode
+    code = b''.join(i.to_bytes() for i in insts)
+    code_buf = (ctypes.c_uint8 * len(code)).from_buffer_copy(code)
+    code_addr = ctypes.addressof(code_buf)
+    program_raw = decode_program(code, "rdna4")
+    program = {code_addr + offset: val for offset, val in program_raw.items()}
+
+    # Setup wave state
+    st = WaveState(n_lanes=1)
+    st.pc = code_addr
+    if sgprs:
+      for idx, val in sgprs.items(): st._write_sgpr(idx, val)
+    if vgprs:
+      for (reg, lane), val in vgprs.items(): st._write_vgpr(reg, lane, val)
+
+    # Setup vmem buffer with external_ptr=0 (maps to address 0, allows any pointer access)
+    vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
+
+    # Execute
+    c_bufs = [ctypes.c_uint64(st.sgpr_buf._buf.va_addr), ctypes.c_uint64(st.vgpr_buf._buf.va_addr),
+              ctypes.c_uint64(vmem_buf._buf.va_addr), ctypes.c_uint64(0), ctypes.c_uint64(0)]
+    for _ in range(100):
+      if (pc := st.pc) == 0xFFFFFFFFFFFFFFFF or pc not in program: break
+      _, fxn, globals_list, _ = program[pc]
+      fxn(*[c_bufs[g] for g in globals_list])
+    return st
+
+  def test_vopd_dual_mov(self):
+    """Test VOPD with two V_DUAL_MOV_B32 operations: v[1]=s[1], v[2]=s[2]."""
+    insts = [ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+                      vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0])]
+    st = self._run(insts, sgprs={1: 0x40e00000, 2: 0x41100000})  # 7.0f, 9.0f
+    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
+    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
+
+  def test_vopd_dual_mov_after_other_vopd(self):
+    """Test VOPD reuse: first VOPD(v[3]=0, v[0]=?), then VOPD(v[1]=s[1], v[2]=s[2])."""
+    # This matches the BEAM kernel sequence that fails
+    insts = [
+      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+               vdstx=v[3], vdsty=v[0], srcx0=0, srcy0=s[0], vsrcx1=v[0], vsrcy1=v[0]),  # v[3]=0, v[0]=s[0]
+      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+               vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]),  # v[1]=s[1], v[2]=s[2]
+    ]
+    st = self._run(insts, sgprs={0: 0x40a00000, 1: 0x40e00000, 2: 0x41100000})  # 5.0f, 7.0f, 9.0f
+    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
+    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
+
+  def test_vopd_with_s_add_f32_sequence(self):
+    """Test full BEAM kernel sequence: s_add_f32 then VOPD."""
+    # This is the exact sequence from the failing BEAM kernel
+    insts = [
+      ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[0], ssrc0=s[0], ssrc1=s[8]),   # s[0] = s[0] + s[8]
+      ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[1], ssrc0=s[1], ssrc1=s[9]),   # s[1] = s[1] + s[9]
+      ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[2], ssrc0=s[2], ssrc1=s[10]),  # s[2] = s[2] + s[10]
+      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+               vdstx=v[3], vdsty=v[0], srcx0=0, srcy0=s[0], vsrcx1=v[0], vsrcy1=v[0]),
+      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+               vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]),
+    ]
+    # Input: s[0:2] = [1,2,3], s[8:10] = [4,5,6]
+    # After s_add_f32: s[0:2] = [5,7,9]
+    st = self._run(insts, sgprs={0: 0x3f800000, 1: 0x40000000, 2: 0x40400000,  # 1.0, 2.0, 3.0
+                                  8: 0x40800000, 9: 0x40a00000, 10: 0x40c00000})  # 4.0, 5.0, 6.0
+    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
+    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
+
+  def test_s_mov_b32_then_vopd(self):
+    """Test s_mov_b32 followed by VOPD - simulates BEAM kernel sequence."""
+    # Use s_mov_b32 with SGPR source (copy from pre-initialized SGPRs)
+    # s[10:12] will have values set by test harness, copy to s[0:2], then VOPD to VGPRs
+    insts = [
+      ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[0], ssrc0=s[10]),  # s[0] = s[10]
+      ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[1], ssrc0=s[11]),  # s[1] = s[11]
+      ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[2], ssrc0=s[12]),  # s[2] = s[12]
+      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
+               vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]),
+    ]
+    st = self._run(insts, sgprs={10: 0x40a00000, 11: 0x40e00000, 12: 0x41100000})  # 5.0, 7.0, 9.0
+    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
+    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
+
+if __name__ == '__main__':
+  unittest.main()
@@ -1,10 +1,9 @@
 #!/usr/bin/env python3
 """Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
-import unittest, io, sys, re
-from tinygrad import Device
-from tinygrad.renderer.amd import detect_format
-from test.amd.helpers import llvm_assemble, llvm_disasm, get_target, get_mattr
-from test.amd.disasm import disasm
+import unittest, io, sys, re, subprocess, os
+from extra.assembly.amd.dsl import Inst
+from extra.assembly.amd import decode_inst, detect_format
+from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump, get_target, get_mattr

 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."""
-  return llvm_assemble(instrs, get_target(arch), get_mattr(arch))
+  """Compile multiple instructions with a single llvm-mc call."""
+  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))
-  return llvm_disasm(code, mcpu, mattr)[:len(instrs)]
+  src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n" + "\n".join(f"  {instr}" for instr in instrs) + "\n"
+  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]
-        if our_disasm == llvm_disasm_str:
+        llvm_disasm = disasm_llvm_map[idx]
+        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(asm_failed, 0, "Asm failures:\n" + "\n".join(asm_failures[:20]))
+    self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
+    self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
    # Note: disasm string comparison is informational only - formatting differences between LLVM versions are expected

  # Basic unary ops
@@ -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 tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
-from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
-from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
-                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4,
-                                     InstOp, InstOpRDNA4, print_packets)
-from test.amd.helpers import TARGET_TO_ARCH
+from extra.assembly.amd import decode_inst
+from extra.assembly.amd.autogen.rdna3.ins import SOPP
+from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
+from extra.assembly.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_L4, WAVEEND, INST, INST_L4, VALUINST, IMMEDIATE, IMMEDIATE_MASK,
+                                     ALUEXEC, VMEMEXEC, PACKET_TYPES_L3, PACKET_TYPES_L4, InstOp, InstOpL4, print_packets)
+from extra.assembly.amd.test.helpers import TARGET_TO_ARCH

-import tinygrad
-EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
+EXAMPLES_DIR = Path(__file__).parent.parent.parent.parent / "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}
+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]
-    buf[0] = ctypes.cast(current_blob[0], ctypes.POINTER(ctypes.c_ubyte))  # type: ignore[arg-type]
-    buf_size[0] = len(current_blob[0])  # type: ignore[arg-type]
-    return len(current_blob[0])  # type: ignore[arg-type]
+    current_blob[0] = (ctypes.c_ubyte * len(blob)).from_buffer_copy(blob)
+    buf[0] = ctypes.cast(current_blob[0], ctypes.POINTER(ctypes.c_ubyte))
+    buf_size[0] = len(current_blob[0])
+    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:
-            print(f"\n=== {name} event {i} ===")
-            print_packets(packets)
+          if DEBUG >= 2: 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())
+    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,7 +133,7 @@ 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))]), 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)]), 0, f"no WAVEEND in {name}")

  def test_time_monotonic(self):
@@ -153,9 +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_names = [p.op.name for p in all_packets if isinstance(p, (INST, INST_RDNA4))]
-        self.assertGreater(len(inst_names), 0, f"no INST packets in {name}")
-        self.assertGreater(len([n for n in inst_names if n.startswith("JUMP")]), 0, f"no JUMP packets in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (INST, INST_L4))]), 0, f"no INST packets in {name}")

  expected: dict[str, list[int]] = {}  # override in subclasses
  def test_packet_counts(self):
@@ -183,7 +176,7 @@ 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, 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) 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}")
@@ -200,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):
@@ -210,19 +203,22 @@ class SQTTExamplesTestBase(unittest.TestCase):
 class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
  target = "gfx1100"
  expected = {
-    "profile_empty_run_0": [1974, 1961, 2014, 2065, 2092, 1998],
-    "profile_empty_run_1": [1979, 1972, 2019, 2070, 2097, 2003],
-    "profile_gemm_run_0": [2038, 11076, 2324, 2129, 2156, 2062],
-    "profile_gemm_run_1": [2038, 11037, 2318, 2129, 2156, 2062],
-    "profile_ops_run_0": [2038, 5070, 2078, 2129, 2156, 2062],
-    "profile_ops_run_1": [2038, 5007, 2078, 2129, 2156, 2062],
-    "profile_plus_run_0": [1979, 1979, 2030, 2070, 2097, 2003],
-    "profile_plus_run_1": [1979, 2043, 2030, 2070, 2097, 2003],
+    "profile_empty_run_0": [1844, 1885, 1905, 1956, 1983, 1889],
+    "profile_empty_run_1": [1780, 1885, 1905, 1956, 1983, 1889],
+    "profile_gemm_run_0": [2656, 2025, 2045, 2096, 2123, 2029, 3183, 2019, 2039, 2090, 2117, 2023, 19119, 2013, 2033, 2084, 2111, 2017],
+    "profile_gemm_run_1": [2662, 2025, 2045, 2096, 2123, 2029, 3179, 2019, 2039, 2090, 2117, 2023, 19113, 2071, 2091, 2142, 2169, 2075],
+    "profile_plus_run_0": [1886, 2013, 2033, 2084, 2111, 2017],
+    "profile_plus_run_1": [1988, 2071, 2091, 2142, 2169, 2075],
  }

 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
-@unittest.skip("TODO: fix CDNA")
-class TestSQTTExamplesCDNA(SQTTExamplesTestBase): target = "gfx950"
+# 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"

 if __name__ == "__main__":
  unittest.main()
@@ -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()
@@ -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(self.time)
+    _events[self.event]["wall"].append(time.monotonic() - self.start)
    return False

 class KernelTimeEvent:
@@ -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 tinygrad.runtime.autogen.amd.rdna3.ins import *
+from extra.assembly.amd.dsl import s, v, VCC_LO, NULL
+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_WORKGROUP_Y = 3         # workgroup_id_y (system SGPR)
+S_TILE_X = 2              # workgroup_x << 7
+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_TILE_Y = 15             # workgroup_y << 7
+S_WORKGROUP_X = 14        # workgroup_id_x
+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):
-    """Patch branch offsets and return the finalized instruction list."""
+  def to_asm(self):
+    # 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'):
-  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}"
+def build_kernel(arch='gfx1100'):
  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)
-    k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_KERNARG_B[0]], i * N * 4))
+  for i in range(1, 8):  # B: 16KB apart
+    k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_KERNARG_B[0]], i * 0x4000))
    k.emit(s_addc_u32(s[S_PREFETCH_B+i*2+1], s[S_KERNARG_B[1]], 0))
  k.emit(s_mov_b64(s[S_PREFETCH_A:S_PREFETCH_A+1], s[S_KERNARG_A[0]:S_KERNARG_A[1]]))  # A[0]: no offset
-  for i in range(1, 8):  # A: each pointer 16 rows of A apart (16*N*4 bytes)
-    k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_KERNARG_A[0]], i * N * 64))
+  for i in range(1, 8):  # A: 256KB apart
+    k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_KERNARG_A[0]], i * 0x40000))
    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(v_mul_lo_u32(v[V_GLOBAL_A_ADDR], v[4], N))                          # (lane_id/8)*N
-  k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], v[V_LANE_ID_MOD8], v[V_GLOBAL_A_ADDR]))  # + lane_id%8
+  k.emit(s_lshl_b32(s[19], s[S_TILE_Y], 12))
+  k.emit(v_lshl_add_u32(v[V_GLOBAL_A_ADDR], v[4], 12, v[V_LANE_ID_MOD8]))  # (lane_id/8)*4K + lane_id%8
  k.emit(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, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
-                                                                                  estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    sink = UOp.sink(A.base, B.base, C.base, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
+                                                                           estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=asm),
+                                 UOp(Ops.BINARY, arg=binary)))
  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:
-      c_np, tc_np = c.numpy(), tc.numpy()
-      for bi in range(N // 128):
-        for bj in range(N // 128):
-          blk_c = c_np[bi*128:(bi+1)*128, bj*128:(bj+1)*128]
-          blk_ref = tc_np[bi*128:(bi+1)*128, bj*128:(bj+1)*128]
-          blk_diff = blk_c - blk_ref
-          zero_rows = [i for i in range(128) if np.all(np.abs(blk_c[i,:]) < 1e-10)]
-          nz_rows = [i for i in range(128) if i not in zero_rows]
-          nz_mse = float(np.mean(blk_diff[nz_rows,:]**2)) if nz_rows else 0
-          print(f"Block ({bi},{bj}): zero_rows={zero_rows}, nz_rows_mse={nz_mse:.2e}")
-          # show first few non-zero row comparisons
-          if nz_rows and nz_mse > 1e-6:
-            for r in nz_rows[:3]:
-              print(f"  row {r} asm[0:8]:  {blk_c[r,:8]}")
-              print(f"  row {r} ref[0:8]:  {blk_ref[r,:8]}")
-      raise RuntimeError("matmul is wrong!")
+    if err != err or err > 1e-06: raise RuntimeError("matmul is wrong!")
+
+def run_sqtt():
+  """Run with SQTT profiling and write trace files."""
+  import subprocess, os
+
+  # Run test_matmul in a subprocess with SQTT enabled from the start (no verify)
+  env = {**os.environ, "AMD": "1", "SQTT": "1", "CNT": "1", "PROFILE": "1", "PYTHONPATH": ".", "VERIFY": "0"}
+  result = subprocess.run(
+    ["python", "-c", "from extra.gemm.amd_asm_matmul import test_matmul; test_matmul()"],
+    capture_output=True, text=True, env=env, timeout=120
+  )
+  print(result.stdout)
+
+  # Run roc.py to extract trace data
+  result = subprocess.run(
+    ["python", "extra/sqtt/roc.py", "--profile", "/tmp/profile.pkl.tiny", "--kernel", "kernel"],
+    capture_output=True, text=True, env={**os.environ, "DEBUG": "5"}, timeout=60
+  )
+  output = result.stdout + result.stderr
+
+  # Write full output to trace file
+  with open("/tmp/sqtt_trace.txt", "w") as f:
+    f.write(output)
+  print(f"Wrote {len(output)} bytes to /tmp/sqtt_trace.txt")

 if __name__ == "__main__":
-  test_matmul()
+  if getenv("ASM", 0): print(build_kernel(Device[Device.DEFAULT].arch))
+  elif getenv("SQTT", 0): run_sqtt()
+  else: test_matmul()
@@ -1,57 +1,47 @@
 import atexit, functools
+from tinygrad.runtime.support.compiler_amd import HIPCompiler
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, DEBUG
-from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG
+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

-@functools.cache
-def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
+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(NUM_WG, "gidx0")
-  insts = build_kernel(batch, M, N, K, A.dtype.base)
-  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
-  sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
-                  arg=KernelInfo(name=f"gemm_{batch}_{M}_{N}_{K}", estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname),
-                                UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  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
-def _asm_gemm_report():
-  print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used')
-  if DEBUG >= 2 and counters["todos"]:
-    from collections import Counter
-    for msg, cnt in Counter(counters["todos"]).most_common(): print(f'  {cnt:3d}x {msg}')
-atexit.register(_asm_gemm_report)
+atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))

 def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
  N = b.shape[1]
+  # only sharding on the batch or K is tested, others might work too
  if isinstance(a.device, tuple):
-    if a.ndim == 2 and a.uop.axis == 0 and b.uop.axis is None: M //= len(a.device)
-    elif a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
-    elif a.ndim == 2 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
+    if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis == 2 and b.uop.axis == 0: K //= len(a.device)
    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
    dname = a.device[0]
  else: dname = a.device
  arch = getattr(Device[dname].renderer, "arch", "")
  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
-  if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
-    return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
+  if (key:=(M, N, K)) not in GEMM_ARGS and arch == "gfx950": return todo(f"GEMM shape not supported {key} on {arch}")
  return True

 # ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
@@ -86,10 +76,6 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
 def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
  counters["used"] += 1
-  unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
-  if unfold_batch:
-    orig_batch = a.shape[0]
-    a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
  squeeze = a.ndim == 2
  if squeeze: a = a.unsqueeze(0)

@@ -97,26 +83,18 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
  N = b.shape[1]
  is_multi = isinstance(a.device, tuple)
  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
-  if (m_sharded:=is_multi and a.uop.axis == 1): M //= len(a.device)
-  n_sharded = is_multi and b.uop.axis == 1

  if is_multi:
-    if n_sharded:
-      out = Tensor(Tensor.empty(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
-    elif m_sharded:
-      out = Tensor(Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
-    else:
-      out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+    out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
  else:
    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)

-  renderer = Device[a.device[0] if is_multi else a.device].renderer
-  dname, arch = renderer.device, getattr(renderer, "arch", "")
+  dname = a.device[0] if is_multi else a.device
+  arch = getattr(Device[dname].renderer, "arch", "")
  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
+    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]
  if k_sharded: out = out.sum(0)
-  out = out.squeeze(0) if squeeze else out
-  if unfold_batch: out = out.reshape(orig_batch, -1, out.shape[-1])
-  return out
+  return out.squeeze(0) if squeeze else out
@@ -10,9 +10,9 @@ HEVC_ROUNDUP = getenv("DATA_ROUNDUP", 32)
@functools.cache
 def _hevc_jitted_decoder(out_image_size:tuple[int, int], max_hist:int, inplace:bool):
  def hevc_decode_frame(pos:Variable, hevc_tensor:Tensor, offset:Variable, sz:Variable, opaque:Tensor, i:Variable, *hist:Tensor, outbuf:Tensor|None=None):
-    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist).realize()
+    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist)
    if outbuf is not None: outbuf.assign(x).realize()
-    return x
+    return x.realize()
  return TinyJit(hevc_decode_frame)

 def hevc_decode(hevc_tensor:Tensor, opaque:Tensor, frame_info:list, luma_h:int, luma_w:int,
@@ -74,14 +74,10 @@ if __name__ == "__main__":
  Device.default.synchronize()

  # decode all frames using the iterator
-  tm = Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps"))
-  with tm:
+  with Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps")):
    images = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w, history=hist, preallocated_outputs=out_images))
    Device.default.synchronize()

-  fps = len(frame_info)/(tm.et/1e9)
-  assert fps >= getenv("ASSERT_FPS", 0), f"HEVC decode too slow: {fps:.2f} fps"
-
  # validation
  if getenv("VALIDATE", 0):
    import pickle
@@ -1,27 +1,29 @@
-import os
+import os, pathlib

 # TODO: there is a timing bug without this
 os.environ["AMD_AQL"] = "1"

-from tinygrad import Tensor, Device
-from tinygrad.helpers import getenv
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from tinygrad.renderer.amd.dsl import Reg, Inst, s, v
+from tinygrad.device import Device
+from tinygrad.runtime.support.compiler_amd import HIPCompiler
+from extra.assembly.amd.dsl import Reg, Inst, s, v

 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
-NUM_WAVES = 4
+NUM_WAVES = 2
 FLOPS_PER_MATMUL = 16*16*16*2
-INTERNAL_LOOP = getenv("LOOP", 10_000)
+INTERNAL_LOOP = 1_000_00
 INSTRUCTIONS_PER_LOOP = 200
+DIRECTIVE = ".amdhsa_wavefront_size32 1"

-def repeat(insts:list[Inst], n:int, counter_sreg:Reg) -> list[Inst]:
+assemblyTemplate = (pathlib.Path(__file__).parent / "template.s").read_text()
+
+def repeat(insts:list[Inst], n:int, counter_sreg:Reg) -> bytes:
+  preamble = s_mov_b32(counter_sreg, n).to_bytes()
  insts_bytes = b"".join([inst.to_bytes() for inst in insts])
  sub_inst, cmp_inst = s_sub_u32(counter_sreg, counter_sreg, 1), s_cmp_lg_i32(counter_sreg, 0)
  loop_sz = len(insts_bytes) + sub_inst.size() + cmp_inst.size()
  branch_inst = s_cbranch_scc1(simm16=-((loop_sz // 4) + 1) & 0xFFFF)
-  return [s_mov_b32(counter_sreg, n)] + insts + [sub_inst, cmp_inst, branch_inst, s_endpgm()]
+  return preamble + insts_bytes + sub_inst.to_bytes() + cmp_inst.to_bytes() + branch_inst.to_bytes() + s_endpgm().to_bytes()

 def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs):
  if accum:
@@ -30,17 +32,16 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs)
    inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[1]:vgprIndices[2]], 1)
  else:
    inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[3]:vgprIndices[4]], v[vgprIndices[5]])
-  insts = repeat([inst for _ in range(INSTRUCTIONS_PER_LOOP)], n=INTERNAL_LOOP, counter_sreg=s[1])
-  def fxn(A:UOp) -> UOp:
-    threads = UOp.special(WAVE_SIZE * NUM_WAVES, "lidx0")
-    gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
-    FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
-    sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-  dummy = Tensor.zeros(1).contiguous().realize()
-  out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
-  ei = out.schedule()[-1].lower()
-  elapsed = min([ei.run(wait=True) for _ in range(2)])
+  vgprs:set = set()
+  for n,_ in inst._fields:
+    if isinstance(val:=getattr(inst, n), Reg) and val.offset >= v.offset: vgprs |= {val.offset+i for i in range(val.sz)}
+  inst_bytes = repeat([inst for _ in range(INSTRUCTIONS_PER_LOOP)], n=INTERNAL_LOOP, counter_sreg=s[1])
+  inst_hex = "\n".join("  .byte " + ",".join(f"0x{b:02x}" for b in inst_bytes[i:i+16]) for i in range(0, len(inst_bytes), 16)) + "\n"
+  src = assemblyTemplate.replace("INTERNAL_LOOP", str(INTERNAL_LOOP)).replace("INSTRUCTION", inst_hex).replace("VGPR_COUNT", str(len(vgprs)))
+  src = src.replace("DIRECTIVE", DIRECTIVE)
+  lib = COMPILER.compile(src)
+  fxn = DEV.runtime("matmul", lib)
+  elapsed = min([fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) for _ in range(2)])
  FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
  print(f"{inst.op_name.lower():<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")

@@ -48,8 +49,9 @@ if __name__=="__main__":
  DEV = Device[Device.DEFAULT]
  arch = DEV.renderer.arch

+  COMPILER = HIPCompiler(arch)
  if arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
-    from tinygrad.runtime.autogen.amd.rdna3.ins import *
+    from extra.assembly.amd.autogen.rdna3.ins import *
    if arch == 'gfx1103': NUM_WORKGROUPS = 8
    if arch == 'gfx1151': NUM_WORKGROUPS = 32
    launchBenchmark(v_wmma_bf16_16x16x16_bf16, (7,8,15))
@@ -59,7 +61,7 @@ if __name__=="__main__":
    launchBenchmark(v_wmma_i32_16x16x16_iu4, (7,8,9))
    launchBenchmark(v_wmma_i32_16x16x16_iu8, (7,8,11))
  elif arch in {'gfx1200', 'gfx1201'}:
-    from tinygrad.runtime.autogen.amd.rdna4.ins import *
+    from extra.assembly.amd.autogen.rdna4.ins import *
    # this instruction does not exist in the rdna4 isa, use the co version
    s_sub_u32 = s_sub_co_u32
    NUM_WORKGROUPS = 64
@@ -88,7 +90,8 @@ if __name__=="__main__":
    FLOPS_PER_MATMUL = 16*16*64*2
    launchBenchmark(v_swmmac_i32_16x16x64_iu4, (7,8,9,10,13,14), False)
  elif arch == 'gfx950':
-    from tinygrad.runtime.autogen.amd.cdna.ins import *
+    from extra.assembly.amd.autogen.cdna.ins import *
+    DIRECTIVE = ".amdhsa_accum_offset 4"
    NUM_WORKGROUPS = 256
    WAVE_SIZE = 64
    NUM_WAVES = 4
@@ -0,0 +1,33 @@
+    .text
+    .globl matmul
+    .p2align 8
+    .type matmul,@function
+matmul:
+    INSTRUCTION
+
+.rodata
+.p2align 6
+.amdhsa_kernel matmul
+  .amdhsa_next_free_vgpr VGPR_COUNT
+  .amdhsa_next_free_sgpr 3
+  DIRECTIVE
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.version:
+  - 1
+  - 0
+amdhsa.kernels:
+  - .name: matmul
+    .symbol: matmul.kd
+    .kernarg_segment_size: 0
+    .group_segment_fixed_size: 0
+    .private_segment_fixed_size: 0
+    .kernarg_segment_align: 4
+    .wavefront_size: 32
+    .sgpr_count: 8
+    .vgpr_count: 32
+    .max_flat_workgroup_size: 1024
+...
+.end_amdgpu_metadata
@@ -41,13 +41,9 @@ class Attention:
    self.n_rep = self.n_heads // self.n_kv_heads
    self.max_context = max_context

-    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.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
+    self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+    self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
    self.wo = linear(self.n_heads * self.head_dim, dim, bias=False)

    self.q_norm = nn.RMSNorm(dim, qk_norm) if qk_norm is not None else None
@@ -55,11 +51,9 @@ class Attention:

  def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]=None) -> 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)
+      if not hasattr(self, 'wqkv'): self.wqkv = Tensor.cat(self.wq.weight, self.wk.weight, self.wv.weight)
+      xqkv = x @ self.wqkv.T
+      xq, xk, xv = xqkv.split([self.wq.weight.shape[0], self.wk.weight.shape[0], self.wv.weight.shape[0]], dim=2)
    else:
      xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)

@@ -206,14 +200,14 @@ class Transformer:

  def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):
    _bsz, seqlen = tokens.shape
-    h = self.tok_embeddings(tokens).contiguous()
+    h = self.tok_embeddings(tokens)
    freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :]

    if self.max_context != 0 and seqlen > 1:
      mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1)
    else: mask = None
    for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
-    logits = self.output(self.norm(h).contiguous().contiguous_backward()).contiguous_backward()
+    logits = self.output(self.norm(h))
    if math.isnan(temperature): return logits

    return sample(logits[:, -1, :].flatten(), temperature, top_k, top_p, alpha_f, alpha_p)
@@ -150,7 +150,7 @@ class ResNet:
        continue # Skip FC if transfer learning

      if 'bn' not in k and 'downsample' not in k: assert obj.shape == dat.shape, (k, obj.shape, dat.shape)
-      obj.assign(dat.to(obj.device).cast(obj.dtype).reshape(obj.shape))
+      obj.assign(dat.to(obj.device).reshape(obj.shape))

 ResNet18 = lambda num_classes=1000: ResNet(18, num_classes=num_classes)
 ResNet34 = lambda num_classes=1000: ResNet(34, num_classes=num_classes)
@@ -3,7 +3,7 @@ from __future__ import annotations
 import enum, collections
 from typing import Iterator
 from tinygrad.helpers import colored
-from tinygrad.renderer.amd.sqtt import PacketType, bits
+from extra.assembly.amd.sqtt import PacketType, bits

 # ═══════════════════════════════════════════════════════════════════════════════
 # STALL REASONS
@@ -129,6 +129,14 @@ def decode_tpc_id(tpc_id:int) -> tuple[int, int, int]:
  # NOTE: valid only for ops_nv, cuda encoding is different
  return (tpc_id >> 5, (tpc_id >> 1) & 0xf, tpc_id & 1)

+def print_samples(samples:list[tuple[PMASample, int]]) -> None:
+  if not samples: return
+  base_pc = min(s.pc_offset for s, _ in samples)
+  for s, tpc_id in samples:
+    gpc, tpc, sm = decode_tpc_id(tpc_id)
+    stall_str = colored(f"{s.stall_reason.name:17}", STALL_COLORS.get(s.stall_reason, "white"))
+    print(f"pc=0x{s.pc_offset - base_pc:06x} {stall_str} ev={s.stall_key:2d} active={s.active} wave={s.wave_id:2d} gpc={gpc} tpc={tpc} sm={sm}")
+
 def print_packets(data:bytes, sm_version:int=0x800) -> None:
  record_size = 9 if sm_version >= 0x890 else 8
  tpc_state: dict[int, list[int]] = collections.defaultdict(list)
@@ -179,11 +187,7 @@ if __name__ == "__main__":
    print(f"\n{'='*60}\nDump {dump_idx} ({len(raw)} bytes, {len(raw)//32} packets)\n{'='*60}")
    if "--raw" in sys.argv: print_packets(raw, sm_ver)
    else:
-      samples = []
-      for s, tpc_id in decode(raw, sm_ver):
-        gpc, tpc, sm = decode_tpc_id(tpc_id)
-        stall_str = colored(f"{s.stall_reason.name:17}", STALL_COLORS.get(s.stall_reason, "white"))
-        print(f"pc=0x{s.pc_offset:06x} {stall_str} ev={s.stall_key:2d} active={s.active} wave={s.wave_id:2d} gpc={gpc} tpc={tpc} sm={sm}")
-        samples.append((s, tpc_id))
+      samples = list(decode(raw, sm_ver))
      print(f"\nDecoded {len(samples)} samples:")
+      print_samples(samples)
      print_aggregated(samples)
@@ -7,7 +7,7 @@ export CAPTURE_PROCESS_REPLAY=1
 rm "$LOGOPS" 2>/dev/null || true
 test/external/process_replay/reset.py

-CI=1 python3 -m pytest -n=auto test/backend/test_ops.py test/backend/test_nn.py test/unit/test_winograd.py test/null/test_real_world.py --durations=20
+CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/unit/test_winograd.py test/null/test_real_world.py --durations=20
 CL=1 python3 -m pytest test/test_tiny.py

 # extract, sort and uniq
@@ -7,8 +7,8 @@ import subprocess, struct, math, functools
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.helpers import getenv

-from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.renderer.amd.asm import waitcnt
+from extra.assembly.amd.autogen.rdna3.ins import *
+from extra.assembly.amd.asm import waitcnt

 from test.testextra.test_cfg_viz import asm_kernel

@@ -1,24 +0,0 @@
-#!/bin/sh
-install_loc="$HOME/.local/bin"
-docker build --platform=linux/amd64 -t cuda-nvcc:12.8 - <<'EOF'
-FROM ubuntu:22.04
-RUN apt-get update && apt-get install -y --no-install-recommends wget ca-certificates && \
-  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb && \
-  dpkg -i cuda-keyring_1.1-1_all.deb && \
-  apt-get update && apt-get install -y --no-install-recommends cuda-nvcc-12-8 cuda-nvdisasm-12-8 cuda-cuobjdump-12-8 && rm -rf /var/lib/apt/lists/*
-ENV PATH=/usr/local/cuda/bin:$PATH
-EOF
-
-mkdir -p "$install_loc"
-tee "$install_loc/nvccshim" >/dev/null <<'EOF'
-#!/bin/sh
-set -eu
-# assume the final arg is the input path
-# mount it so that container can read it
-dir=$(dirname "${@: -1}")
-exec docker run --rm --platform=linux/amd64 -v "$dir":"$dir" cuda-nvcc:12.8 "$(basename "$0")" "$@"
-EOF
-chmod +x "$install_loc/nvccshim"
-for t in nvcc nvdisasm; do
-  ln -sf "$install_loc/nvccshim" "$install_loc/$t"
-done
@@ -2,13 +2,30 @@

 ## Getting SQ Thread Trace

-`VIZ=2` to enable SQTT profiling.
-
-`SQTT_ITRACE_SE_MASK=X` to select shader engines for instruction tracing, -1 = all, 0 = disabled, >0 = SE bitmask, default 0b11.
+SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.

 `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.

-## Viewing the traces
+`SQTT_ITRACE_SE_MASK=X` to select for which shader engines instruction tracing will be enabled, -1 is all, 0 is none (instruction tracing disabled), >0 is
+bitfield/mask for SEs to enable instruction tracing on. Masking shader engines will give smaller file sizes at a cost of less hits and kernels that
+don't have any wavefront on first simd of shader engine with instruction tracing enabled will not have instruction timings.
+The default is 2 (second shader engine only), only one for file size reasons, second instead of first because dispatch starts from it so there is
+greater chance that kernels with small global size will have instruction tracing data.
+ 
+Note that instruction tracing might not be available for kernels with small global dims, this is not a bug, but it can be improved with various hacks
+to the point where it can reliably trace a kernel consisting of a single wavefront (am only, not quite reliable under amdgpu due to waves sometimes
+being dispatched starting from different simds). More info in comments in ops_amd.py

- Web UI: `tinygrad/viz/serve.py`
- Command line: `python -m tinygrad.renderer.amd.sqtt`
+## Converting pickled profile with SQTT data into RGP file
+
+```bash
+extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
+```
+
+Then load gpu0.rgp into Radeon GPU Profiler. It works just fine both in wine (macos, native version available for linux) and via ssh X forwarding
+
+If multiple gpus are used you can select which one to export with `-d` like this:
+
+```bash
+extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -d 'AMD:5' -o /tmp/gpu5.rgp
+```
@@ -0,0 +1,152 @@
+import os
+os.environ["PYTHONPATH"] = "."
+os.environ["SQTT"] = "1"
+if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
+os.environ["PROFILE"] = "1"
+os.environ["AMD_LLVM"] = "0"
+
+from dataclasses import replace
+import atexit, contextlib
+from tinygrad import Tensor
+from tinygrad.helpers import system, OSX
+from tinygrad.runtime.ops_amd import AMDProgram
+from extra.sqtt.roc import decode, WaveExec, ProfileSQTTEvent
+from tinygrad.device import Device
+
+from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
+
+dev = Device["AMD"]
+
+@contextlib.contextmanager
+def save_sqtt():
+  # clear the old traces
+  dev.profile_events.clear()
+  sqtt:dict[str, list[WaveExec]] = {}
+  yield sqtt
+  events = dev.profile_events
+
+  #rctx = decode(events)
+  #assert len(rctx.inst_execs) > 0, "empty sqtt output"
+  #sqtt.update(rctx.inst_execs)
+
+  for e in events:
+    if isinstance(e, ProfileSQTTEvent):
+      print(replace(e, blob=b''))
+      if e.se == 0:
+        parse_sqtt_print_packets(e.blob)
+
+template = """.text
+.globl matmul
+.p2align 8
+.type matmul,@function
+matmul:
+  INSTRUCTION
+
+.rodata
+.p2align 6
+.amdhsa_kernel matmul
+  .amdhsa_kernarg_size 8
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
+  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
+  .amdhsa_wavefront_size32 1
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.version:
+  - 1
+  - 0
+amdhsa.kernels:
+  - .name: matmul
+    .symbol: matmul.kd
+    .group_segment_fixed_size: 0
+    .private_segment_fixed_size: 0
+    .wavefront_size: 32
+    .sgpr_count: 8
+    .vgpr_count: 8
+    .max_flat_workgroup_size: 1024
+    .kernarg_segment_align: 8
+    .kernarg_segment_size: 8
+    .args:
+      - .address_space:  global
+        .name:           a
+        .offset:         0
+        .size:           8
+        .type_name:      'float*'
+        .value_kind:     global_buffer
+...
+.end_amdgpu_metadata
+"""
+
+def run_asm(src, num_workgroups=1, num_waves=1):
+  WAVE_SIZE = 32
+  t = Tensor.empty(0x1000).realize()
+  buf = t.uop.buffer.ensure_allocated()
+  lib = dev.compiler.compile(template.replace("INSTRUCTION", '\n'.join(src)))
+  dev.compiler.disassemble(lib)
+  fxn = AMDProgram(dev, "matmul", lib)
+  fxn(buf._buf, global_size=(num_workgroups,1,1), local_size=(WAVE_SIZE*num_waves,1,1), wait=True)
+
+if __name__ == "__main__":
+  with save_sqtt() as sqtt:
+    run_asm([
+      "s_nop 100",
+      "s_nop 100",
+      "s_load_b64 s[0:1], s[0:1], null",
+      "s_waitcnt lgkmcnt(0)",
+      "s_nop 100",
+      "s_nop 100",
+      "s_add_i32 s2, s2, 10",
+      "s_add_i32 s2, s2, 10",
+      "s_nop 100",
+      "s_nop 100",
+      "v_mov_b32_e32 v0, 0",
+      "v_mov_b32_e32 v0, 0",
+      "s_nop 100",
+      "s_nop 100",
+      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
+      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
+      "s_nop 100",
+      "s_nop 100",
+      "global_load_b128 v[2:5], v0, s[0:1]",
+      "global_load_b128 v[2:5], v0, s[0:1]",
+      "s_nop 100",
+      "s_nop 100",
+      "s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)",
+      "s_endpgm",
+    ], num_workgroups=1, num_waves=1)
+  exit(0)
+
+  with save_sqtt() as sqtt:
+    #(Tensor.empty(16,16) @ Tensor.empty(16,16)).elu().realize()
+    #Tensor.empty(1, 64).sum(axis=1).realize()
+    Tensor.empty(1).log2().realize()
+  exit(0)
+
+  with save_sqtt() as sqtt:
+    # what's in v0?
+    run_asm([
+      "v_mov_b32_e32 v0, 0",
+      "v_mov_b32_e32 v1, 0",
+      "s_clause 0x1",
+      "s_load_b64 s[0:1], s[0:1], null",
+      "s_waitcnt lgkmcnt(0)",
+    ]+[
+      "global_load_b32 v1, v0, s[0:1]",
+    ]*10+[
+      "global_load_b32 v10, v1, s[0:1]",
+      "s_waitcnt vmcnt(0)",
+
+      #"v_rcp_f32 v1, v0"
+      #"v_add_f32_e32 v1 v0 v0",
+      #"v_add_f32_e32 v5 v4 v4",
+      #"v_add_f32_e32 v7 v6 v6",
+      #"v_add_f32_e32 v1 v0 v0",
+      #"v_add_f32_e32 v2 v1 v1",
+      #"s_nop 1"
+    ]*5+[
+      "v_add_f32_e32 v3 v2 v2",
+    ]*5+[
+      "v_mul_f32_e32 v3 v2 v2",
+    ]*7)
@@ -0,0 +1,548 @@
+import pickle, sys
+from tinygrad.helpers import getenv, Timing, colored
+from extra.sqtt.roc import decode, ProfileSQTTEvent
+
+# do these enums match fields in the packets?
+#from tinygrad.runtime.support.amd import import_soc
+#soc = import_soc([11])
+#perf_sel = {getattr(soc, k):k for k in dir(soc) if k.startswith("SQ_PERF_")}
+
+# Instruction packets (one per ISA op)
+# NOTE: these are bad guesses and may be wrong! feel free to update if you know better
+# some names were taken from SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT
+
+# we see 18 opcodes
+# opcodes(18):  1  2  3  4  5  6  8  9  F 10 11 12 14 15 16 17 18 19
+# if you exclude everything, you are left with 6
+# opcodes( 6): 10 11 14 15 16 17
+# sometimes we see a lot of B, but not repeatable
+
+# not seen
+# 7 A C
+
+# NOTE: INST runs before EXEC
+
+OPCODE_COLORS = {
+  # dispatches are BLACK
+  0x1: "BLACK",
+  0x18: "BLACK",
+
+  # execs are yellow
+  0x2: "yellow",
+  0x3: "yellow",
+  0x4: "YELLOW",
+  0x5: "YELLOW",
+
+  # waves are blue
+  0x8: "blue",
+  0x9: "blue",
+  0x6: "cyan",
+  0xb: "cyan",
+}
+
+OPCODE_NAMES = {
+  # gated by SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT (but others must be enabled for it to show)
+  0x01: "VALUINST",
+  # gated by SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT
+  0x02: "VMEMEXEC",
+  # gated by SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT
+  0x03: "ALUEXEC",
+  # gated by SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT
+  0x04: "IMMEDIATE",
+  0x05: "IMMEDIATE_MASK",
+
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT
+  0x06: "WAVERDY",
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT
+  0x08: "WAVEEND",
+  0x09: "WAVESTART",
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVEALLOC_SHIFT
+  0x0B: "WAVEALLOC",  # FFF00
+
+  # gated by NOT SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT
+  0x0D: "PERF",
+  # gated by SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT
+  0x12: "EVENT",
+  0x13: "EVENT_BIG",  # FFFFF800
+  # some gated by SQ_TT_TOKEN_EXCLUDE_REG_SHIFT, some always there. something is broken with the timing on this
+  0x14: "REG",
+  # gated by SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
+  0x18: "INST",
+  # gated by SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT
+  0x19: "UTILCTR",
+
+  # this is the first (8 byte) packet in the bitstream
+  0x17: "LAYOUT_HEADER",       # layout/mode/group + selectors A/B (reversed)
+
+  # pure time (no extra bits)
+  0x0F: "TS_DELTA_SHORT",
+  0x10: "NOP",
+  0x11: "TS_WAVE_STATE",     # almost pure time, has a small flag
+
+  # not a good name, but seen and understood mostly
+  0x15: "SNAPSHOT",          # small delta + 50-ish bits of snapshot
+  0x16: "TS_DELTA_OR_MARK",  # 36-bit long delta or 36-bit marker
+
+  # packets we haven't seen / rarely see 0x0b
+  0x07: "TS_DELTA_S8_W3_7",         # shift=8,  width=3  (small delta)
+  0x0A: "TS_DELTA_S5_W2_A",         # shift=5,  width=2
+  0x0C: "TS_DELTA_S5_W3_B",         # shift=5,  width=3 (different consumer)
+}
+
+#  SALU    =  0x0 / s_mov_b32
+#  SMEM    =  0x1 / s_load_b*
+#  JUMP    =  0x3 / s_cbranch_scc0
+#  NEXT    =  0x4 / s_cbranch_execz
+#  MESSAGE =  0x9 / s_sendmsg
+#  VALU    =  0xb / v_(exp,log)_f32_e32
+#  VALU    =  0xd / v_lshlrev_b64
+#  VALU    =  0xe / v_mad_u64_u32
+#  VMEM    = 0x21 / global_load_b32
+#  VMEM    = 0x22 / global_load_b32
+#  VMEM    = 0x24 / global_store_b32
+#  VMEM    = 0x25 / global_store_b64
+#  VMEM    = 0x27 / global_store
+#  VMEM    = 0x28 / global_store_b64
+#  LDS     = 0x29 / ds_load_b128
+#  LDS     = 0x2b / ds_store_b32
+#  LDS     = 0x2e / ds_store_b128
+#  ????    = 0x5a / hidden global_load  instruction
+#  ????    = 0x5b / hidden global_load  instruction
+#  ????    = 0x5c / hidden global_store instruction
+#  VALU    = 0x73 / v_cmpx_eq_u32_e32 (not normal VALUINST)
+OPNAME = {
+  0x0: "SALU",
+  0x1: "SMEM",
+  0x3: "JUMP",
+  0x4: "NEXT",
+  0x9: "MESSAGE",
+  0xb: "VALU",
+  0xd: "VALU",
+  0xe: "VALU",
+  0x21: "VMEM_LOAD",
+  0x22: "VMEM_LOAD",
+  0x24: "VMEM_STORE",
+  0x25: "VMEM_STORE",
+  0x26: "VMEM_STORE",
+  0x27: "VMEM_STORE",
+  0x28: "VMEM_STORE",
+  0x29: "LDS_LOAD",
+  0x2b: "LDS_STORE",
+  0x2e: "LDS_STORE",
+  0x50: "__SIMD_LDS_LOAD",
+  0x51: "__SIMD_LDS_LOAD",
+  0x54: "__SIMD_LDS_STORE",
+  0x5a: "__SIMD_VMEM_LOAD",
+  0x5b: "__SIMD_VMEM_LOAD",
+  0x5c: "__SIMD_VMEM_STORE",
+  0x5d: "__SIMD_VMEM_STORE",
+  0x5e: "__SIMD_VMEM_STORE",
+  0x5f: "__SIMD_VMEM_STORE",
+  0x72: "SALU_OR",
+  0x73: "VALU_CMPX",
+}
+
+ALUSRC = {
+  1: "SALU",
+  2: "VALU",
+  3: "VALU_SALU",
+}
+
+MEMSRC = {
+  0: "LDS",
+  1: "__LDS",
+  2: "VMEM",
+  3: "__VMEM",
+}
+
+
+# these tables are from rocprof trace decoder
+# rocprof_trace_decoder_parse_data-0x11c6a0
+# parse_sqtt_180 = b *rocprof_trace_decoder_parse_data-0x11c6a0+0x110040
+
+# ---------- 1. local_138: 256-byte state->opcode table ----------
+
+STATE_TO_OPCODE: bytes = bytes([
+  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x12, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x13, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+])
+
+# opcode mask (the bits used to determine the opcode, worked out by looking at the repeats in STATE_TO_OPCODE)
+
+opcode_mask = {
+  0x10: 0b1111,
+
+  0x16: 0b1111111,
+  0x17: 0b1111111,
+  0x07: 0b1111111,
+  0x19: 0b1111111,
+  0x11: 0b1111111,
+  0x12: 0b11111111,
+  0x13: 0b11111111,
+  0x15: 0b1111111,
+
+  0x18: 0b111,
+  0x1: 0b111,
+
+  0x5: 0b11111,
+  0x6: 0b11111,
+  0xb: 0b11111,
+  0x8: 0b11111,
+  0xc: 0b11111,
+  0xd: 0b11111,
+
+  0xf: 0b1111,
+  0x14: 0b1111,
+
+  0x9: 0b11111,
+  0xa: 0b11111,
+
+  0x4: 0b1111,
+  0x3: 0b1111,
+  0x2: 0b1111,
+}
+
+# ---------- 2. DAT_0012e280: nibble budget per opcode&0x1F ----------
+
+NIBBLE_BUDGET = [
+  0x08, 0x0C, 0x08, 0x08, 0x0C, 0x18, 0x18, 0x40, 0x14, 0x20, 0x30, 0x14, 0x34, 0x1C, 0x30, 0x08,
+  0x04, 0x18, 0x18, 0x20, 0x40, 0x40, 0x30, 0x40, 0x14, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+]
+
+# ---------- 3. delta_map from your hash nodes ----------
+
+# opcode -> (shift, width)
+DELTA_MAP_DEFAULT = {
+  0x01: (3,  3),   # shift=3,  end=6
+  0x02: (4,  2),   # shift=4,  end=6
+  0x03: (4,  2),   # shift=4,  end=6
+  0x04: (4,  3),   # shift=4,  end=7
+  0x05: (5,  3),   # shift=5,  end=8
+  0x06: (5,  3),   # shift=5,  end=8
+  0x07: (8,  3),   # shift=8,  end=11
+  0x08: (5,  3),   # shift=5,  end=8
+  0x09: (5,  2),   # shift=5,  end=7
+  0x0A: (5,  2),   # shift=5,  end=7
+  0x0B: (5,  3),   # shift=5,  end=8
+  0x0C: (5,  3),   # shift=5,  end=8
+  0x0D: (5,  3),   # shift=5,  end=8
+  # NOTE: 0x0e can never be decoded, it's not in the STATE_TO_OPCODE table
+  #0x0E: (7,  2),   # shift=7,  end=9
+  0x0F: (4,  4),   # shift=4,  end=8
+  0x10: (0,  0),   # shift=0,  end=0  (no delta)
+  0x11: (7,  9),   # shift=7,  end=16
+  0x12: (8,  3),   # shift=8,  end=11
+  0x13: (8,  3),   # shift=8,  end=11
+  0x14: (4,  3),   # shift=4,  end=7
+  0x15: (7,  3),   # shift=7,  end=10
+  0x16: (12, 36),  # shift=12, end=48 (36-bit field, matches the 0x16 special-case)
+  0x17: (0,  0),   # shift=0,  end=0  (no delta)
+  0x18: (4,  3),   # shift=4,  end=7
+  0x19: (7,  2),   # shift=7,  end=9
+}
+
+# ---------- 4. One-line-per-packet parser ----------
+
+def reg_mask(opcode):
+  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]
+  shift, width = DELTA_MAP_DEFAULT[opcode]
+  delta_mask = ((1 << width) - 1) << shift
+  assert delta_mask & opcode_mask[opcode] == 0, "masks shouldn't overlap"
+  return ((1 << nb_bits) - 1) & ~(delta_mask | opcode_mask[opcode])
+
+def decode_packet_fields(opcode: int, reg: int) -> str:
+  """
+  Decode packet payloads conservatively, using:
+    - NIBBLE_BUDGET[opcode & 0x1F] to mask reg down to true width.
+    - DELTA_MAP_DEFAULT[opcode] to expose the "primary" field (often delta).
+    - Per-opcode layouts derived from rocprof's decompiled consumers.
+  """
+  # --- 0. Restrict to real packet bits not used in delta ---------------------------------
+  pkt = reg & reg_mask(opcode)
+  fields: list[str] = []
+
+  match opcode:
+    case 0x01: # VALUINST
+      # 6 bit field
+      flag = (pkt >> 6) & 1
+      wave = pkt >> 7
+      fields.append(f"wave={wave:x}")
+      if flag: fields.append("flag")
+    case 0x02: # VMEMEXEC
+      # 2 bit field (pipe is a guess)
+      src = pkt>>6
+      fields.append(f"src={src} [{MEMSRC.get(src, '')}]")
+    case 0x03: # ALUEXEC
+      # 2 bit field
+      src = pkt>>6
+      fields.append(f"src={src} [{ALUSRC.get(src, '')}]")
+    case 0x04: # IMMEDIATE_4
+      # 5 bit field (actually 4)
+      wave = pkt >> 7
+      fields.append(f"wave={wave:x}")
+    case 0x05: # IMMEDIATE_5
+      # 16 bit field
+      # 1 bit per wave
+      fields.append(f"mask={pkt>>8:016b}")
+    case 0x6:
+      # wave ready FFFF00
+      # 16 bit field
+      # 1 bit per wave
+      fields.append(f"mask={pkt>>8:016b}")
+    case 0x0d:
+      # 20 bit field
+      fields.append(f"arg = {pkt>>8:X}")
+    case 0x12:
+      fields.append(f"event = {pkt>>11:X}")
+    case 0x15:
+      fields.append(f"snap = {pkt>>10:X}")
+    case 0x19:
+      # wave end
+      fields.append(f"ctr = {pkt>>9:X}")
+    case 0xf:
+      extracted_delta = (reg >> 4) & 0xF
+      fields.append(f"strange_delta=0x{extracted_delta:x}")
+    case 0x11:
+      # DELTA_MAP_DEFAULT: shift=7, width=9 -> small delta.
+      # FF0000 is the mask
+      coarse    = pkt >> 16
+      fields.append(f"coarse=0x{coarse:02x}")
+      # From decomp:
+      #  - when layout<3 and coarse&1, it sets a "has interesting wave" flag
+      #  - when coarse&8, it marks all live waves as "terminated"
+      if coarse & 0x01:
+        fields.append("flag_wave_interest=1")
+      if coarse & 0x08:
+        fields.append("flag_terminate_all=1")
+    case 0x8:
+      # wave end, this is 20 bits (FFF00)
+      flag7   = (pkt >> 8) & 1
+      simd    = (pkt >> 9) & 3
+      cu      = ((pkt >> 11) & 0x7) | (flag7 << 3)
+      wave    = (pkt >> 15) & 0x1f
+      fields.append(f"wave={wave:x}")
+      fields.append(f"simd={simd}")
+      fields.append(f"cu={cu}")
+    case 0x9:
+      # From case 9 (WAVESTART) in multiple consumers:
+      #   flag7  = (w >> 7) & 1        (low bit of uVar41)
+      #   cls2   = (w >> 8) & 3        (class / group)
+      #   slot4  = (w >> 10) & 0xf     (slot / group index)
+      #   idx_lo = (w >> 0xd) & 0x1f   (low index, layout<4 path)
+      #   idx_hi = (w >> 0xf) & 0x1f   (high index, layout>=4 path)
+      #   id7    = (w >> 0x19) & 0x7f  (7-bit id)
+      flag7   = (pkt >> 7) & 1
+      simd    = (pkt >> 8) & 3
+      cu      = ((pkt >> 10) & 0x7) | (flag7 << 3)
+      wave    = (pkt >> 13) & 0x1F
+      id7     = (pkt >> 17)
+      fields.append(f"wave={wave:x}")
+      fields.append(f"simd={simd}")
+      fields.append(f"cu={cu}")
+      fields.append(f"id7=0x{id7:x}")
+    case 0x18:
+      # FFF88 is the mask
+      # From case 0x18:
+      #   low3   = w & 7
+      #   grp3   = (w >> 3) or (w >> 4) & 7   (layout-dependent)
+      #   flags  = bits 6 (B6) and 7 (B7)
+      #   hi8    = (w >> 0xc) & 0xff   (layout 4 path)
+      #   hi7    = (w >> 0xd) & 0x7f   (other layouts)
+      #   idx5   = (w >> 7) or (w >> 8) & 0x1f, used as wave index
+      flag1 = (pkt >> 3) & 1
+      flag2 = (pkt >> 7) & 1
+      wave = (pkt >> 8) & 0x1F
+      op = (pkt >> 13)
+      fields.append(f"wave={wave:x}")
+      fields.append(f"op=0x{op:02x} [{OPNAME.get(op, '')}]")
+      if flag1: fields.append("flag1")
+      if flag2: fields.append("flag2")
+    case 0x14:
+      subop   = (pkt >> 16) & 0xFFFF       # (short)(w >> 0x10)
+      val32   = (pkt >> 32) & 0xFFFFFFFF   # (uint)(w >> 0x20)
+      slot    = (pkt >> 7) & 0x7           # index in local_168[...] tables
+      hi_byte = (pkt >> 8) & 0xFF          # determines config vs marker
+
+      fields.append(f"subop=0x{subop:04x}")
+      fields.append(f"slot={slot}")
+      fields.append(f"val32=0x{val32:08x}")
+
+      if hi_byte & 0x80:
+        # Config flavour: writes config words into per-slot state arrays.
+        fields.append("kind=config")
+        if subop == 0x000C:
+          fields.append("slot=lo")
+        elif subop == 0x000D:
+          fields.append("slot=hi")
+      else:
+        # COR marker: subop 0xC342, payload "COR\0" → start of a COR region.
+        if subop == 0xC342:
+          fields.append("kind=cor_stream")
+          if val32 == 0x434F5200:
+            fields.append("cor_magic='COR\\0'")
+    case 0x16:
+      # Bits:
+      #   bit8  -> 0x100
+      #   bit9  -> 0x200
+      #   bits 12..47 -> 36-bit field used as delta or marker
+      bit8 = bool(pkt & 0x100)
+      bit9 = bool(pkt & 0x200)
+      if not bit9:
+        mode = "delta"
+      elif not bit8:
+        mode = "marker"
+      else:
+        mode = "other"
+      # need to use reg here
+      val36 = (reg >> 12) & ((1 << 36) - 1)
+      fields.append(f"mode={mode}")
+      if mode != "delta":
+        fields.append(f"val36=0x{val36:x}")
+    case 0x17:
+      # From decomp (two sites with identical logic):
+      #   layout = (w >> 7) & 0x3f
+      #   mode   = (w >> 0xd) & 3
+      #   group  = (w >> 0xf) & 7
+      #   sel_a  = (w >> 0x1c) & 0xf
+      #   sel_b  = (w >> 0x21) & 7
+      #   flag4  = (w >> 0x3b) & 1  (only meaningful when layout == 4)
+      layout = (pkt >> 7)  & 0x3F
+      simd   = (pkt >> 13) & 0x3  # you can change this by changing traced simd
+      group  = (pkt >> 15) & 0x7
+      sel_a  = (pkt >> 0x1C) & 0xF
+      sel_b  = (pkt >> 0x21) & 0x7
+      flag4  = (pkt >> 0x3B) & 0x1
+
+      fields.append(f"layout={layout}")
+      fields.append(f"group={group}")
+      fields.append(f"simd={simd}")
+      fields.append(f"sel_a={sel_a}")
+      fields.append(f"sel_b={sel_b}")
+      if layout == 4:
+        fields.append(f"layout4_flag={flag4}")
+    case _:
+      fields.append(f"{pkt:X} & {reg_mask(opcode):X}")
+  return ",".join(fields)
+
+FILTER_LEVEL = getenv("FILTER", 1)
+
+DEFAULT_FILTER: tuple[int, ...] = tuple()
+# NOP + pure time + "sample"
+if FILTER_LEVEL >= 0: DEFAULT_FILTER += (0x10, 0xf, 0x11)
+# reg + event + sample + marker
+# TODO: events are probably good
+if FILTER_LEVEL >= 1: DEFAULT_FILTER += (0x14, 0x12, 0x16)
+# instruction runs + valuinst
+if FILTER_LEVEL >= 2: DEFAULT_FILTER += (0x01, 0x02, 0x03)
+# instructions dispatch (inst, immed)
+if FILTER_LEVEL >= 3: DEFAULT_FILTER += (0x4, 0x5, 0x18)
+# waves
+if FILTER_LEVEL >= 4: DEFAULT_FILTER += (0x6, 0x8, 0x9)
+
+def parse_sqtt_print_packets(data: bytes, filter=DEFAULT_FILTER, verbose=True) -> None:
+  """
+  Minimal debug: print ONE LINE per decoded token (packet).
+
+  Now prints only the actual nibbles that belong to each packet, instead of
+  the full 64-bit shift register.
+  """
+  n = len(data)
+  time = 0
+  last_printed_time = 0
+  reg = 0          # shift register
+  offset = 0       # bit offset, in steps of 4 (one nibble)
+  nib_budget = 0x40
+  flags = 0
+  token_index = 0
+  opcodes_seen = set()
+
+  while (offset >> 3) < n:
+    # 1) Fill register with nibbles according to nib_budget
+    if nib_budget != 0:
+      target = offset + 4 + ((nib_budget - 1) & ~3)
+      while offset != target and (offset >> 3) < n:
+        byte = data[offset >> 3]
+        nib = (byte >> (offset & 4)) & 0xF
+        reg = ((reg >> 4) | (nib << 60)) & ((1 << 64) - 1)
+        offset += 4
+      if offset != target: break  # don't parse past the end
+
+    # 2) Decode token from low 8 bits
+    opcode = STATE_TO_OPCODE[reg & 0xFF]
+    opcodes_seen.add(opcode)
+
+    # 4) Set next nibble budget based on opcode
+    nib_budget = NIBBLE_BUDGET[opcode & 0x1F]
+
+    # 5) Get delta
+    shift, width = DELTA_MAP_DEFAULT[opcode]
+    delta = (reg >> shift) & ((1 << width) - 1)
+
+    # 6) Update time and handle special opcodes 0xF/0x16
+    if opcode == 0x16:
+      two_bits = (reg >> 8) & 0x3
+      if two_bits == 1:
+        flags |= 0x01
+
+      # Common 36-bit field at bits [12..47]
+      if (reg & 0x200) == 0:
+        # delta mode: add 36-bit delta to time
+        pass
+      elif (reg & 0x100) == 0:
+        # marker / other modes: no time advance
+        # real marker: bit9=1, bit8=0, non-zero payload
+        # "other" 0x16 variants, ignored for timing
+        delta = 0
+      else:
+        raise RuntimeError("unknown 0x16 delta")
+    elif opcode == 0x0F:
+      # opcode 0x0F has an offset of 4 to the delta
+      # update: it's actually computed to be 8 to match WAVESTART
+      delta = delta + 8
+
+    # Append extra decoded fields into the note string
+    note = decode_packet_fields(opcode, reg)
+
+    # this delta happens before the instruction
+    time += delta
+    token_index += 1
+
+    if verbose and (filter is None or opcode not in filter):
+      print(f"{time:8d} +{time-last_printed_time:8d} : "+colored(f"{OPCODE_NAMES[opcode]:18s} ", OPCODE_COLORS.get(opcode, "white"))+f"{note}")
+      last_printed_time = time
+
+  # Optional summary at the end
+  print(f"# done: tokens={token_index:_}, final_time={time}, flags=0x{flags:02x}")
+  if verbose:
+    print(f"opcodes({len(opcodes_seen):2d}):",
+          ' '.join([colored(f"{op:2X}", "WHITE" if op in opcodes_seen else "BLACK") for op in sorted(opcode_mask)]))
+
+
+def parse(fn:str):
+  with Timing(f"unpickle {fn}: "): dat = pickle.load(open(fn, "rb"))
+  #if getenv("ROCM", 0):
+  #  with Timing(f"decode {fn}: "): ctx = decode(dat)
+  dat_sqtt = [x for x in dat if isinstance(x, ProfileSQTTEvent)]
+  print(f"got {len(dat_sqtt)} SQTT events in {fn}")
+  return dat_sqtt
+
+if __name__ == "__main__":
+  fn = "extra/sqtt/examples/profile_gemm_run_0.pkl"
+  dat_sqtt = parse(sys.argv[1] if len(sys.argv) > 1 else fn)
+  for i,dat in enumerate(dat_sqtt):
+    with Timing(f"decode pkt {i} with len {len(dat.blob):_}: "):
+      parse_sqtt_print_packets(dat.blob, verbose=getenv("V", 1))
@@ -1,148 +0,0 @@
-#!/usr/bin/env python3
-# Run all ALU and memory instructions in the ISA
-import functools, inspect
-from enum import Enum
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AddrSpace
-from tinygrad.renderer.amd.dsl import Inst, Reg, OPERANDS, SrcField, VGPRField, SGPRField, SSrcField, SBaseField, AlignedSGPRField, BitField
-from tinygrad.renderer.amd.dsl import FixedBitField, EnumBitField, s, v, NULL, VCC_LO
-from extra.gemm.amd_asm_matmul import Kernel
-
-# skip instructions that mutate wave state (PC, EXEC, allocations, signals)
-SKIP = {"S_SETPC_B64", "S_SWAPPC_B64", "S_RFE_B64", "S_BARRIER_SIGNAL_ISFIRST", "S_GET_BARRIER_STATE", "S_ALLOC_VGPR", "S_SLEEP_VAR", "S_GETPC_B64",
-        "S_SENDMSG_RTN_B32", "S_SENDMSG_RTN_B64"}
-# skip barriers, s_waits, wrap level atomics, and ray tracing (bvh)
-SKIP_SUBSTR = ["SAVEEXEC", "CMPX", "WREXEC", "MOVREL", "ATOMIC", "S_BUFFER_", "S_ATC_PROBE", "BARRIER", "S_WAITCNT", "BVH",
-               "DS_CMPSTORE_RTN", "DS_WRAP_RTN_B32", "DS_ORDERED_COUNT", "DS_GWS", "GS_REG", "GLOBAL_LOAD_LDS", "GLOBAL_STORE_BLOCK"]
-
-ALU_FORMATS = {"VOP1", "VOP1_LIT", "VOP1_SDST", "VOP2", "VOP2_LIT", "VOP3", "VOP3_SDST", "VOP3SD", "VOP3P", "VOP3P_MFMA", "VOP3PX2",
-               "VOPC", "SOP1", "SOP1_LIT", "SOP2", "SOP2_LIT", "SOPC", "SOPC_LIT", "SOPK", "SOPK_LIT", "VINTERP"}
-# intentionally not testing scratch memory ops
-MEM_FORMATS = {"VGLOBAL", "GLOBAL", "SMEM", "DS"}
-
-def should_skip(op:Enum) -> bool: return (name:=op.name) in SKIP or any(sub in name for sub in SKIP_SUBSTR)
-
-# ** named register assignments
-
-# ALU operands
-ALU_VGPR_STRIDE = 16          # v[0], v[16], v[32], ... per ALU operand slot
-ALU_SGPR_STRIDE = 4           # s[0], s[4], s[8], ... per ALU operand slot
-
-# memory address registers
-S_KERNARG_PTR = (0, 1)
-S_BUF_PTR = (2, 3)
-V_VADDR = (0, 1)
-V_DS_ADDR = 0
-
-# memory data registers
-MEM_VGPR_BASE = 32            # v[32], v[48], ... for vdst/vdata/vsrc
-MEM_VGPR_STRIDE = 16          # spacing between memory data vgpr slots
-MEM_SGPR_BASE = 8             # s[8], s[10], ... for SMEM sdata
-MEM_SGPR_STRIDE = 2           # spacing between memory data sgpr slots
-
-# ** create an ALU instruction based on the operands
-
-def create_alu_inst(op:Enum, builder:functools.partial[Inst]) -> Inst:
-  inst_cls, operands, slot = builder.func, OPERANDS[op], 0
-  kwargs:dict[str, Reg|int] = {}
-  for name, field in inst_cls._fields:
-    if isinstance(field, (FixedBitField, EnumBitField)): continue
-    nregs = max(1, operands[name][1] // 32) if name in operands else 1
-    is_sreg = name in operands and "SREG" in str(operands[name][2])
-    base_v, base_s = slot * ALU_VGPR_STRIDE, slot * ALU_SGPR_STRIDE
-    if name == "sdst" and isinstance(field, SGPRField): reg = VCC_LO
-    elif is_sreg and not isinstance(field, VGPRField): reg = VCC_LO
-    elif isinstance(field, VGPRField): reg = v[base_v:base_v+nregs-1] if nregs > 1 else v[base_v]
-    elif isinstance(field, SSrcField): reg = VCC_LO if nregs <= 2 else s[base_s:base_s+nregs-1] if nregs > 1 else s[base_s]
-    elif isinstance(field, SGPRField): reg = s[base_s:base_s+nregs-1] if nregs > 1 else s[base_s]
-    elif isinstance(field, SrcField): reg = v[base_v:base_v+nregs-1] if nregs > 1 else v[base_v]
-    else: reg = None
-    if reg is not None: kwargs[name] = reg; slot += 1
-    elif isinstance(field, BitField): kwargs[name] = field.default
-  return builder(**kwargs)
-
-# ** create a memory instruction with pre set address registers
-
-MEM_PRESET_REGS:dict[str, dict[str, Reg]] = {
-  "VGLOBAL":{"saddr":s[S_BUF_PTR[0]:S_BUF_PTR[1]], "vaddr":v[V_VADDR[0]:V_VADDR[1]]},
-  "GLOBAL":{"saddr":s[S_BUF_PTR[0]:S_BUF_PTR[1]], "addr":v[V_DS_ADDR]},  # addr is 32-bit offset when saddr is valid SGPR
-  "DS":{"addr":v[V_DS_ADDR]},
-  "SMEM":{"sbase":s[S_KERNARG_PTR[0]:S_KERNARG_PTR[1]], "soffset":NULL},
-}
-
-def create_mem_inst(op:Enum, builder:functools.partial[Inst]) -> Inst:
-  inst_cls, operands, field_map = builder.func, OPERANDS.get(op, {}), MEM_PRESET_REGS.get(builder.func.__name__, {})
-  kwargs:dict[str, Reg|int] = {}
-  vslot, sslot = 0, 0
-  for name, field in inst_cls._fields:
-    if isinstance(field, (FixedBitField, EnumBitField)): continue
-    if name in field_map:
-      kwargs[name] = field_map[name]
-      continue
-    nregs = max(1, operands[name][1] // 32) if name in operands else 1
-    if isinstance(field, VGPRField):
-      vi = MEM_VGPR_BASE + vslot * MEM_VGPR_STRIDE
-      kwargs[name] = v[vi:vi+nregs-1] if nregs > 1 else v[vi]
-      vslot += 1
-    elif isinstance(field, (SGPRField, AlignedSGPRField, SBaseField)):
-      si = MEM_SGPR_BASE + sslot * MEM_SGPR_STRIDE
-      kwargs[name] = s[si:si+nregs-1] if nregs > 1 else s[si]
-      sslot += 1
-    elif isinstance(field, BitField): kwargs[name] = field.default
-  return builder(**kwargs)
-
-# ** collect all memory and ALU instructions from the ISA autogen
-
-def collect_instructions() -> tuple[list[Inst], list[Inst], list[str]]:
-  op_map:dict[Enum, functools.partial[Inst]] = {}
-  for name, obj in inspect.getmembers(all_insts):
-    if isinstance(obj, functools.partial) and len(obj.args) == 1: op_map[obj.args[0]] = obj
-  alu_insts:list[Inst] = []
-  mem_insts:list[Inst] = []
-  skipped:list[str] = []
-  for op_enum, builder in op_map.items():
-    if should_skip(op_enum) or op_enum not in OPERANDS: skipped.append(op_enum.name); continue
-    fmt = builder.func.__name__
-    if fmt in ALU_FORMATS: alu_insts.append(create_alu_inst(op_enum, builder))
-    elif fmt in MEM_FORMATS: mem_insts.append(create_mem_inst(op_enum, builder))
-  return alu_insts, mem_insts, skipped
-
-def exec_insts(insts:list):
-  k = Kernel(arch)
-  # ** prologue for global memory
-  k.emit(s_load_b64(sdata=s[S_BUF_PTR[0]:S_BUF_PTR[1]], sbase=s[S_KERNARG_PTR[0]:S_KERNARG_PTR[1]], soffset=NULL))
-  k.waitcnt(lgkm=0)
-  k.emit(v_mov_b32_e32(v[V_VADDR[0]], 0))
-  k.emit(v_mov_b32_e32(v[V_VADDR[1]], 0))
-  # ** emit
-  for inst in insts: k.emit(inst)
-  k.emit(s_endpgm())
-  # ** run
-  NUM_THREADS, NUM_GRIDS, BUF_SIZE = 32, 1, 1024*1024
-  def fxn(A:UOp, B:UOp, C:UOp) -> UOp:
-    lidx, gidx = UOp.special(NUM_THREADS, "lidx0"), UOp.special(NUM_GRIDS, "gidx0")
-    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=BUF_SIZE, addrspace=AddrSpace.LOCAL), (), "lds")
-    sink = UOp.sink(A.base, B.base, C.base, lds, lidx, gidx, arg=KernelInfo(name="discover_ops"))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in k.finalize()))))
-  A = Tensor.empty(BUF_SIZE, dtype=dtypes.uint8)
-  B = Tensor.empty(1, dtype=dtypes.uint8)
-  C = Tensor.empty(1, dtype=dtypes.uint8)
-  Tensor.custom_kernel(A, B, C, fxn=fxn)[0].realize()
-
-if __name__ == "__main__":
-  import sys
-  arch = Device[Device.DEFAULT].renderer.arch
-  if arch.startswith("gfx12"):
-    from tinygrad.runtime.autogen.amd.rdna4.ins import *
-    import tinygrad.runtime.autogen.amd.rdna4.ins as all_insts
-  elif arch.startswith("gfx11"):
-    from tinygrad.runtime.autogen.amd.rdna3.ins import *
-    import tinygrad.runtime.autogen.amd.rdna3.ins as all_insts
-    # these don"t exist in RDNA3, only RDNA3.5 and above
-    SKIP.update(["S_FMAAK_F32", "S_FMAMK_F32"])
-  else:
-    print(f"{arch} not supported yet")
-    sys.exit(0)
-  alu_insts, mem_insts, skipped = collect_instructions()
-  print(f"collected {len(alu_insts)} ALU + {len(mem_insts)} memory instructions ({len(skipped)} skipped)")
-  exec_insts(mem_insts+alu_insts)
@@ -1,25 +1,23 @@
-import os, subprocess, sys, shlex
+import os, subprocess
 from pathlib import Path
 from tinygrad.helpers import temp

 EXAMPLES_DIR = Path(__file__).parent
 PROFILE_PATH = Path(temp("profile.pkl", append_user=True))

-EXAMPLES = {
-  "empty":"test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
-  "plus":"test/test_tiny.py TestTiny.test_plus",
-  "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=64, N)@Tensor.empty(N, N)).realize()\"",
-  "ops":"extra/sqtt/examples/discover_ops.py"
-}
+EXAMPLES = [
+  "test.test_custom_kernel.TestCustomKernel.test_empty",
+  "test.test_tiny.TestTiny.test_plus",
+  "test.test_tiny.TestTiny.test_gemm",
+]

 if __name__ == "__main__":
  arch = subprocess.check_output(["python", "-c", "from tinygrad import Device; print(Device['AMD'].arch)"], text=True,
                                 env={**os.environ, "DEBUG":"0"}).rstrip()
  (EXAMPLES_DIR/arch).mkdir(exist_ok=True)
-  for name,test in EXAMPLES.items():
+  for test in EXAMPLES:
    for i in range(2):
-      # AM_RESET=1 gets a clear trace, does not work on mi300 machines
-      subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
-                     env={**os.environ, "AMD":"1", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
-      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
+      subprocess.run(["python", "-m", "unittest", test], cwd=EXAMPLES_DIR.parent.parent.parent,
+                     env={**os.environ, "AMD":"1", "SQTT_LIMIT_SE":"-1", "VIZ":"-2"}, check=True)
+      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{test.split('.')[-1].replace('test_', '')}_run_{i}.pkl")
      print(f"saved SQTT trace to {dest}")
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
George HotzandGitHub	2fe45b0660	Merge branch 'master' into call_inline	2026-02-10 14:12:08 +08:00
George HotzandGitHub	1d88723aa0	Merge branch 'master' into call_inline	2026-02-10 14:02:43 +08:00
geohot	b0dd3af093	inline all these calls	2026-02-10 13:25:29 +08:00
geohot	e89221e9aa	add inline flag for call	2026-02-10 12:19:51 +08:00