mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-20 12:26:08 +00:00
Compare commits
207
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6b0a9f5ee6 | ||
|
|
cad7feec02 | ||
|
|
635ed2df9d | ||
|
|
6c0f0e29ff | ||
|
|
9347757c2d | ||
|
|
3a92df66ea | ||
|
|
7c234a9c7c | ||
|
|
91bde927ef | ||
|
|
44135e2e84 | ||
|
|
8b1b15aec0 | ||
|
|
8b5ff403fa | ||
|
|
d8aba24967 | ||
|
|
9973a81356 | ||
|
|
c5492f8f75 | ||
|
|
d5f954858d | ||
|
|
3e2c05ee9f | ||
|
|
35c9701df0 | ||
|
|
92246ea731 | ||
|
|
c34c6d9468 | ||
|
|
eacccc5ace | ||
|
|
ed295e74dc | ||
|
|
cff33c8d78 | ||
|
|
74fa3c7d09 | ||
|
|
0fbc551622 | ||
|
|
ba209d6305 | ||
|
|
6b308b89b7 | ||
|
|
40f9fa2db4 | ||
|
|
2917ed1616 | ||
|
|
f3f4d9b387 | ||
|
|
c70c112254 | ||
|
|
13e5d00d0e | ||
|
|
a071adffc0 | ||
|
|
b878f9d5a4 | ||
|
|
efcb32f6a9 | ||
|
|
b34c637767 | ||
|
|
16b652302e | ||
|
|
3f61a96d79 | ||
|
|
e7b5d8a434 | ||
|
|
e372c841ba | ||
|
|
1732a4ec4b | ||
|
|
f3aceaa08b | ||
|
|
309197bca5 | ||
|
|
15a056715d | ||
|
|
027b935269 | ||
|
|
2db04d0696 | ||
|
|
cb500466c2 | ||
|
|
3caa1e2c98 | ||
|
|
5f1ede7f7e | ||
|
|
5bd4593eda | ||
|
|
241f0402b4 | ||
|
|
25c82dd242 | ||
|
|
35900290b2 | ||
|
|
87f4bc5446 | ||
|
|
2833c5a54b | ||
|
|
72a3f78d19 | ||
|
|
c714881832 | ||
|
|
a8896f28e1 | ||
|
|
325f4006ff | ||
|
|
7fb18f7e47 | ||
|
|
4491ec0c9e | ||
|
|
6ddddc68af | ||
|
|
b699b9f763 | ||
|
|
02084f5376 | ||
|
|
2b3e01e79c | ||
|
|
45f7fd073d | ||
|
|
21d0f6bb76 | ||
|
|
3170365a5b | ||
|
|
0120d69caa | ||
|
|
20653d2996 | ||
|
|
ea7b149ca5 | ||
|
|
f86c728440 | ||
|
|
eda6a73897 | ||
|
|
ce464b147a | ||
|
|
83063cc3e4 | ||
|
|
9497ec00f2 | ||
|
|
5cff5698f7 | ||
|
|
7a81a3cb98 | ||
|
|
34fe105386 | ||
|
|
4f2f38bf64 | ||
|
|
70405b4f3c | ||
|
|
b2a0b9c551 | ||
|
|
aae08b20e0 | ||
|
|
785d04d127 | ||
|
|
f6a78a29e0 | ||
|
|
404eed6172 | ||
|
|
f550f9204c | ||
|
|
7abf4591ba | ||
|
|
cfb8bf5814 | ||
|
|
7ebda28692 | ||
|
|
ad041416ca | ||
|
|
bf356ae996 | ||
|
|
606786e152 | ||
|
|
34ea053b26 | ||
|
|
280790e438 | ||
|
|
9a9564118c | ||
|
|
8328511808 | ||
|
|
bdb421f13e | ||
|
|
66caa9fe1d | ||
|
|
8003db2a28 | ||
|
|
c1b8644a3f | ||
|
|
35c2870b1f | ||
|
|
a49924a0e9 | ||
|
|
3b354bc11f | ||
|
|
efb2ae87c6 | ||
|
|
bd55507ee4 | ||
|
|
6242a9d151 | ||
|
|
9f082e8e25 | ||
|
|
2cc64d71b0 | ||
|
|
7cbafb2ef1 | ||
|
|
9dc524536f | ||
|
|
61dc70f1a8 | ||
|
|
0e282025ff | ||
|
|
2e2b5fed12 | ||
|
|
f49e4714af | ||
|
|
a78fcc55a4 | ||
|
|
fcbb896e05 | ||
|
|
ff7853a65a | ||
|
|
42abb0586c | ||
|
|
ebbaad6bfd | ||
|
|
5f52266225 | ||
|
|
5a1a561e0f | ||
|
|
b27527f05a | ||
|
|
ecbac8a338 | ||
|
|
af0392efea | ||
|
|
e036d6df89 | ||
|
|
dfb813b760 | ||
|
|
cb7c76a3bd | ||
|
|
51398edf9c | ||
|
|
8e416df438 | ||
|
|
b8ea0d779c | ||
|
|
4d5c4d256d | ||
|
|
ed222070f7 | ||
|
|
ce84a23142 | ||
|
|
24723327ac | ||
|
|
9726500de8 | ||
|
|
c0f52c9dcb | ||
|
|
c69470be52 | ||
|
|
b91b46091c | ||
|
|
17ef4af72c | ||
|
|
6a5430ab00 | ||
|
|
baff10d32c | ||
|
|
1c5ed8e8b5 | ||
|
|
526fd4ec71 | ||
|
|
20777f30b9 | ||
|
|
0ed58c1fcd | ||
|
|
e2987001ee | ||
|
|
8bf7c9c1d2 | ||
|
|
2bb07d4824 | ||
|
|
52acadc160 | ||
|
|
c0c1c1c8c8 | ||
|
|
b6d08f247d | ||
|
|
f14428090f | ||
|
|
13973e4dea | ||
|
|
051fe6c8bc | ||
|
|
a9a7b33404 | ||
|
|
29402034a1 | ||
|
|
ba9aa5cd6f | ||
|
|
4968060ad4 | ||
|
|
35bd39e4ba | ||
|
|
b998a80b5d | ||
|
|
404755bafd | ||
|
|
25440f0f72 | ||
|
|
f7ee644950 | ||
|
|
b063518ea7 | ||
|
|
b23f4517ab | ||
|
|
3f3786ded9 | ||
|
|
a14896fff2 | ||
|
|
c475c3a6d7 | ||
|
|
0221b96761 | ||
|
|
dc27eb48ac | ||
|
|
efc99d0c55 | ||
|
|
49d1bf93d6 | ||
|
|
04c79505ec | ||
|
|
39f99b207a | ||
|
|
7e14cdcb06 | ||
|
|
69cdc8066d | ||
|
|
9c89be5235 | ||
|
|
2b838dc1d8 | ||
|
|
a19d21ea9c | ||
|
|
b557c46233 | ||
|
|
d7e1f26e3d | ||
|
|
ab58926b00 | ||
|
|
0497387e45 | ||
|
|
fc4faed0b2 | ||
|
|
94bca91f3e | ||
|
|
7322d9ec4a | ||
|
|
0d326f5b9b | ||
|
|
9c6850fc01 | ||
|
|
9d8397be11 | ||
|
|
72236bbd3d | ||
|
|
81cf9ea0ab | ||
|
|
37f0fa11b6 | ||
|
|
35db73b231 | ||
|
|
d178235309 | ||
|
|
ff856a74cb | ||
|
|
39923203ba | ||
|
|
63a1bb8507 | ||
|
|
0a98fd38b3 | ||
|
|
0e409ff5ce | ||
|
|
f1471a3b99 | ||
|
|
37720fd6c0 | ||
|
|
25ef866e89 | ||
|
|
88eb230326 | ||
|
|
f541540129 | ||
|
|
c6769badc2 | ||
|
|
fc5278746f | ||
|
|
f07c39cfa4 |
@@ -14,10 +14,12 @@ on:
|
||||
paths:
|
||||
- 'tinygrad/runtime/autogen/**/*'
|
||||
- 'tinygrad/runtime/support/autogen.py'
|
||||
- '.github/workflows/autogen.yml'
|
||||
workflow_dispatch:
|
||||
paths:
|
||||
- 'tinygrad/runtime/autogen/**/*'
|
||||
- 'tinygrad/runtime/support/autogen.py'
|
||||
- '.github/workflows/autogen.yml'
|
||||
|
||||
jobs:
|
||||
autogen:
|
||||
@@ -39,102 +41,45 @@ jobs:
|
||||
pydeps: 'pyyaml mako'
|
||||
- name: Install autogen support packages
|
||||
run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
|
||||
- name: Verify OpenCL autogen
|
||||
- name: Regenerate autogen files
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
|
||||
rm tinygrad/runtime/autogen/opencl.py
|
||||
python3 -c "from tinygrad.runtime.autogen import opencl"
|
||||
diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
|
||||
- name: Verify CUDA autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
|
||||
mv tinygrad/runtime/autogen/nvrtc.py /tmp/nvrtc.py.bak
|
||||
mv tinygrad/runtime/autogen/nvjitlink.py /tmp/nvjitlink.py.bak
|
||||
mv tinygrad/runtime/autogen/nv_570.py /tmp/nv_570.py.bak
|
||||
mv tinygrad/runtime/autogen/nv.py /tmp/nv.py.bak
|
||||
rm tinygrad/runtime/autogen/{cuda,nvrtc,nvjitlink,nv_570,nv}.py
|
||||
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv"
|
||||
diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
|
||||
diff /tmp/nvrtc.py.bak tinygrad/runtime/autogen/nvrtc.py
|
||||
diff /tmp/nvjitlink.py.bak tinygrad/runtime/autogen/nvjitlink.py
|
||||
diff /tmp/nv_570.py.bak tinygrad/runtime/autogen/nv_570.py
|
||||
diff /tmp/nv.py.bak tinygrad/runtime/autogen/nv.py
|
||||
- name: Verify AMD autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
|
||||
mv tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
|
||||
mv tinygrad/runtime/autogen/hip.py /tmp/hip.py.bak
|
||||
mv tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
|
||||
mv tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
|
||||
mv tinygrad/runtime/autogen/rocprof.py /tmp/rocprof.py.bak
|
||||
mv tinygrad/runtime/autogen/am/am.py /tmp/am_am.py.bak
|
||||
mv tinygrad/runtime/autogen/am/pm4_soc15.py /tmp/am_pm4_soc15.py.bak
|
||||
mv tinygrad/runtime/autogen/am/pm4_nv.py /tmp/am_pm4_nv.py.bak
|
||||
mv tinygrad/runtime/autogen/am/sdma_4_0_0.py /tmp/am_sdma_4_0_0.py.bak
|
||||
mv tinygrad/runtime/autogen/am/sdma_5_0_0.py /tmp/am_sdma_5_0_0.py.bak
|
||||
mv tinygrad/runtime/autogen/am/sdma_6_0_0.py /tmp/am_sdma_6_0_0.py.bak
|
||||
mv tinygrad/runtime/autogen/am/smu_v13_0_0.py /tmp/am_smu_v13_0_0.py.bak
|
||||
mv tinygrad/runtime/autogen/am/smu_v14_0_2.py /tmp/am_smu_v14_0_2.py.bak
|
||||
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof; from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v14_0_2"
|
||||
diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
|
||||
diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
|
||||
diff /tmp/hip.py.bak tinygrad/runtime/autogen/hip.py
|
||||
diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
|
||||
diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
|
||||
diff /tmp/rocprof.py.bak tinygrad/runtime/autogen/rocprof.py
|
||||
diff /tmp/am_am.py.bak tinygrad/runtime/autogen/am/am.py
|
||||
diff /tmp/am_pm4_soc15.py.bak tinygrad/runtime/autogen/am/pm4_soc15.py
|
||||
diff /tmp/am_pm4_nv.py.bak tinygrad/runtime/autogen/am/pm4_nv.py
|
||||
diff /tmp/am_sdma_4_0_0.py.bak tinygrad/runtime/autogen/am/sdma_4_0_0.py
|
||||
diff /tmp/am_sdma_5_0_0.py.bak tinygrad/runtime/autogen/am/sdma_5_0_0.py
|
||||
diff /tmp/am_sdma_6_0_0.py.bak tinygrad/runtime/autogen/am/sdma_6_0_0.py
|
||||
diff /tmp/am_smu_v13_0_0.py.bak tinygrad/runtime/autogen/am/smu_v13_0_0.py
|
||||
diff /tmp/am_smu_v14_0_2.py.bak tinygrad/runtime/autogen/am/smu_v14_0_2.py
|
||||
- name: Verify Linux autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
|
||||
mv tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
|
||||
mv tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
|
||||
mv tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
|
||||
mv tinygrad/runtime/autogen/pci.py /tmp/pci.py.bak
|
||||
mv tinygrad/runtime/autogen/vfio.py /tmp/vfio.py.bak
|
||||
rm tinygrad/runtime/autogen/{comgr,hsa,hip,amd_gpu,sqtt,rocprof}.py
|
||||
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof"
|
||||
rm tinygrad/runtime/autogen/am/{am,pm4_soc15,pm4_nv,sdma_4_0_0,sdma_5_0_0,sdma_6_0_0,smu_v13_0_0,smu_v14_0_2}.py
|
||||
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_v14_0_2"
|
||||
rm tinygrad/runtime/autogen/{libc,kfd,io_uring,ib,pci,vfio}.py
|
||||
python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
|
||||
diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
|
||||
diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
|
||||
diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
|
||||
diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
|
||||
diff /tmp/pci.py.bak tinygrad/runtime/autogen/pci.py
|
||||
diff /tmp/vfio.py.bak tinygrad/runtime/autogen/vfio.py
|
||||
- name: Verify LLVM autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
|
||||
rm tinygrad/runtime/autogen/llvm.py
|
||||
python3 -c "from tinygrad.runtime.autogen import llvm"
|
||||
diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
|
||||
- name: Verify WebGPU autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
|
||||
rm tinygrad/runtime/autogen/webgpu.py
|
||||
python3 -c "from tinygrad.runtime.autogen import webgpu"
|
||||
diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
|
||||
- name: Verify Qualcomm autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/kgsl.py /tmp/kgsl.py.bak
|
||||
mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
|
||||
rm tinygrad/runtime/autogen/{kgsl,qcom_dsp}.py
|
||||
python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
|
||||
diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
|
||||
diff /tmp/qcom_dsp.py.bak tinygrad/runtime/autogen/qcom_dsp.py
|
||||
- name: Verify libusb autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/libusb.py /tmp/libusb.py.bak
|
||||
rm tinygrad/runtime/autogen/libusb.py
|
||||
python3 -c "from tinygrad.runtime.autogen import libusb"
|
||||
diff /tmp/libusb.py.bak tinygrad/runtime/autogen/libusb.py
|
||||
- name: Verify mesa autogen
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
|
||||
rm tinygrad/runtime/autogen/mesa.py
|
||||
python3 -c "from tinygrad.runtime.autogen import mesa"
|
||||
diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
|
||||
- name: Verify libclang autogen
|
||||
run: |
|
||||
cp tinygrad/runtime/autogen/libclang.py /tmp/libclang.py.bak
|
||||
rm tinygrad/runtime/autogen/avcodec.py
|
||||
python3 -c "from tinygrad.runtime.autogen import avcodec"
|
||||
REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
|
||||
diff /tmp/libclang.py.bak tinygrad/runtime/autogen/libclang.py
|
||||
- name: Check for differences
|
||||
run: |
|
||||
if ! git diff --quiet; then
|
||||
git diff > autogen-ubuntu.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-ubuntu-patch
|
||||
path: autogen-ubuntu.patch
|
||||
|
||||
autogen-mac:
|
||||
name: In-tree Autogen (macos)
|
||||
runs-on: macos-14
|
||||
@@ -146,11 +91,24 @@ jobs:
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
llvm: 'true'
|
||||
- name: Verify macos autogen
|
||||
- name: Regenerate autogen files
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/metal.py /tmp/metal.py.bak
|
||||
rm tinygrad/runtime/autogen/metal.py
|
||||
LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
|
||||
diff /tmp/metal.py.bak tinygrad/runtime/autogen/metal.py
|
||||
- name: Check for differences
|
||||
run: |
|
||||
if ! git diff --quiet; then
|
||||
git diff > autogen-macos.patch
|
||||
echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
|
||||
exit 1
|
||||
fi
|
||||
- name: Upload patch artifact
|
||||
if: failure()
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: autogen-macos-patch
|
||||
path: autogen-macos.patch
|
||||
|
||||
autogen-comgr-3:
|
||||
name: In-tree Autogen (comgr 3)
|
||||
runs-on: ubuntu-24.04
|
||||
@@ -169,8 +127,20 @@ jobs:
|
||||
echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
|
||||
sudo apt -qq update || true
|
||||
sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
|
||||
- name: Verify comgr (3) autogen
|
||||
- name: Regenerate autogen files
|
||||
run: |
|
||||
mv tinygrad/runtime/autogen/comgr_3.py /tmp/comgr_3.py.bak
|
||||
rm tinygrad/runtime/autogen/comgr_3.py
|
||||
python3 -c "from tinygrad.runtime.autogen import comgr_3"
|
||||
diff /tmp/comgr_3.py.bak tinygrad/runtime/autogen/comgr_3.py
|
||||
- name: Check for differences
|
||||
run: |
|
||||
if ! git diff --quiet; then
|
||||
git diff > autogen-comgr3.patch
|
||||
echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
|
||||
exit 1
|
||||
fi
|
||||
- name: Upload patch artifact
|
||||
if: failure()
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: autogen-comgr3-patch
|
||||
path: autogen-comgr3.patch
|
||||
|
||||
+98
-212
@@ -49,19 +49,19 @@ jobs:
|
||||
- name: Print macOS version
|
||||
run: sw_vers
|
||||
- name: Run Stable Diffusion
|
||||
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
|
||||
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
|
||||
- name: Run Stable Diffusion without fp16
|
||||
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
|
||||
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing
|
||||
- name: Run Stable Diffusion v2
|
||||
# TODO: very slow step time
|
||||
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
|
||||
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing
|
||||
# process replay can't capture this, the graph is too large
|
||||
- name: Run SDXL
|
||||
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
|
||||
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing
|
||||
- name: Run model inference benchmark
|
||||
run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
|
||||
- name: Test speed vs torch
|
||||
run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
|
||||
run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py
|
||||
- name: Test tensor cores
|
||||
run: METAL=1 python3.11 test/opt/test_tensor_cores.py
|
||||
- name: Test AMX tensor cores
|
||||
@@ -71,84 +71,59 @@ jobs:
|
||||
DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
|
||||
DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
|
||||
- name: Run Tensor Core GEMM (float)
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py
|
||||
- name: Run Tensor Core GEMM (half)
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py
|
||||
- name: Run Tensor Core GEMM (bfloat16)
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
|
||||
run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py
|
||||
- name: Fuzz Padded Tensor Core GEMM
|
||||
run: METAL=1 M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3.11 ./extra/gemm/fuzz_matmul.py
|
||||
- name: Run LLaMA
|
||||
run: |
|
||||
BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
|
||||
BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
|
||||
BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run LLaMA with BEAM
|
||||
run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
|
||||
run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run quantized LLaMA
|
||||
run: |
|
||||
BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8 | tee llama_int8.txt
|
||||
BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4 | tee llama_nf4.txt
|
||||
BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8
|
||||
BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4
|
||||
- name: Run quantized LLaMA3
|
||||
run: |
|
||||
BENCHMARK_LOG=llama3_int8 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8 | tee llama3_int8.txt
|
||||
BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4 | tee llama3_nf4.txt
|
||||
BENCHMARK_LOG=llama3_int8 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8
|
||||
BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4
|
||||
#- name: Run LLaMA 7B on 4 (virtual) GPUs
|
||||
# run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_four_gpu.txt
|
||||
# run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run GPT2
|
||||
run: |
|
||||
BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
|
||||
BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
|
||||
BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run GPT2 w HALF
|
||||
run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
|
||||
run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
|
||||
- name: Run GPT2 w HALF/BEAM
|
||||
run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
|
||||
run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
|
||||
- name: Run OLMoE
|
||||
run: BENCHMARK_LOG=olmoe python3.11 examples/olmoe.py
|
||||
- name: Train MNIST
|
||||
run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
|
||||
run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py
|
||||
|
||||
# NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
|
||||
# the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
|
||||
#- name: Run 10 CIFAR training steps
|
||||
# run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
|
||||
# run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py
|
||||
#- name: Run 10 CIFAR training steps w HALF
|
||||
# run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
|
||||
# run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py
|
||||
|
||||
#- name: Run 10 CIFAR training steps w BF16
|
||||
# run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
|
||||
# run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py
|
||||
# TODO: too slow
|
||||
# - name: Run 10 CIFAR training steps w winograd
|
||||
# run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
|
||||
# run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (Mac)
|
||||
path: |
|
||||
onnx_inference_speed.csv
|
||||
torch_speed.txt
|
||||
llama_unjitted.txt
|
||||
llama_jitted.txt
|
||||
llama_beam.txt
|
||||
llama_int8.txt
|
||||
llama_nf4.txt
|
||||
llama3_int8.txt
|
||||
llama3_nf4.txt
|
||||
llama_four_gpu.txt
|
||||
gpt2_unjitted.txt
|
||||
gpt2_jitted.txt
|
||||
gpt2_half.txt
|
||||
gpt2_half_beam.txt
|
||||
matmul.txt
|
||||
matmul_half.txt
|
||||
matmul_bfloat16.txt
|
||||
sd.txt
|
||||
sd_no_fp16.txt
|
||||
sdv2.txt
|
||||
sdxl.txt
|
||||
beautiful_mnist.txt
|
||||
train_cifar.txt
|
||||
train_cifar_half.txt
|
||||
train_cifar_bf16.txt
|
||||
train_cifar_wino.txt
|
||||
- name: Run process replay tests
|
||||
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3.11 process_replay.py
|
||||
|
||||
@@ -215,7 +190,7 @@ jobs:
|
||||
- name: Run model inference benchmark
|
||||
run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
|
||||
- name: Test speed vs torch
|
||||
run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
|
||||
run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
|
||||
- name: Test speed vs theoretical
|
||||
run: NV=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
|
||||
- name: Test benchmark allreduce
|
||||
@@ -226,79 +201,58 @@ jobs:
|
||||
NV=1 NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
|
||||
- name: Run Tensor Core GEMM (CUDA)
|
||||
run: |
|
||||
CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
|
||||
CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
|
||||
CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
|
||||
CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_fp8.txt
|
||||
CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
|
||||
CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
- name: Run Tensor Core GEMM (PTX)
|
||||
run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
|
||||
run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
- name: Run Tensor Core GEMM (NV)
|
||||
run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
|
||||
run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
- name: Test NV=1
|
||||
run: DEBUG=2 NV=1 python -m pytest -rA test/test_tiny.py
|
||||
- name: Test CUDA=1
|
||||
run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
|
||||
- name: Run Stable Diffusion
|
||||
run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
|
||||
run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
|
||||
# TODO: too slow
|
||||
# - name: Run SDXL
|
||||
# run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
|
||||
# run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
|
||||
- name: Run LLaMA
|
||||
run: |
|
||||
BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
|
||||
BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
|
||||
BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run LLaMA with BEAM
|
||||
run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
|
||||
run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
# - name: Run LLaMA 7B on 4 GPUs
|
||||
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_four_gpu.txt
|
||||
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
# - name: Run LLaMA 7B on 6 GPUs
|
||||
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_six_gpu.txt
|
||||
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run LLaMA-3 8B BEAM
|
||||
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
|
||||
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
|
||||
- name: Run LLaMA-3 8B on 4 GPUs with BEAM
|
||||
run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
|
||||
run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
|
||||
- name: Run quantized LLaMA3
|
||||
run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8 | tee llama3_fp8.txt
|
||||
run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8
|
||||
# - name: Run LLaMA-3 8B on 6 GPUs
|
||||
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
|
||||
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
|
||||
# - name: Run LLaMA-2 70B
|
||||
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_2_70B.txt
|
||||
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run Mixtral 8x7B
|
||||
run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
|
||||
run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing
|
||||
- name: Run GPT2
|
||||
run: |
|
||||
BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
|
||||
BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
|
||||
BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run GPT2 w HALF
|
||||
run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
|
||||
run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing
|
||||
- name: Run GPT2 w HALF/BEAM
|
||||
run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
|
||||
run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (NVIDIA)
|
||||
path: |
|
||||
onnx_inference_speed.csv
|
||||
torch_speed.txt
|
||||
matmul.txt
|
||||
matmul_bfloat16.txt
|
||||
matmul_tf32.txt
|
||||
matmul_ptx.txt
|
||||
matmul_nv.txt
|
||||
sd.txt
|
||||
sdxl.txt
|
||||
llama_unjitted.txt
|
||||
llama_jitted.txt
|
||||
llama_beam.txt
|
||||
llama3_beam.txt
|
||||
llama3_four_gpu.txt
|
||||
llama3_six_gpu.txt
|
||||
llama3_fp8.txt
|
||||
llama_2_70B.txt
|
||||
mixtral.txt
|
||||
gpt2_unjitted.txt
|
||||
gpt2_jitted.txt
|
||||
gpt2_half.txt
|
||||
gpt2_half_beam.txt
|
||||
- name: Run process replay tests
|
||||
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
|
||||
|
||||
@@ -337,44 +291,30 @@ jobs:
|
||||
# - name: Fuzz Padded Tensor Core GEMM (PTX)
|
||||
# run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
|
||||
- name: HEVC Decode Benchmark
|
||||
run: VALIDATE=1 MAX_FRAMES=100 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
|
||||
run: VALIDATE=1 MAX_FRAMES=100 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
|
||||
- name: Train MNIST
|
||||
run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
|
||||
run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
|
||||
- name: Run 10 CIFAR training steps
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py
|
||||
- name: Run 10 CIFAR training steps w HALF
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=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 | tee train_cifar_bf16.txt
|
||||
run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
|
||||
# - name: Run 10 CIFAR training steps w winograd
|
||||
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
|
||||
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
|
||||
- name: Run full CIFAR training w 1 GPU
|
||||
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
|
||||
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
- name: Run full CIFAR training steps w 6 GPUS
|
||||
run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
|
||||
run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
- name: Run MLPerf resnet eval on training data
|
||||
run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
|
||||
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
|
||||
run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
|
||||
run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
- name: Run 10 MLPerf Bert training steps (6 gpu)
|
||||
# TODO: remove BERT_LAYERS once scheduler is fast
|
||||
run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (NVIDIA Training)
|
||||
path: |
|
||||
beautiful_mnist.txt
|
||||
train_cifar.txt
|
||||
train_cifar_half.txt
|
||||
train_cifar_bf16.txt
|
||||
train_cifar_wino.txt
|
||||
train_cifar_one_gpu.txt
|
||||
train_cifar_six_gpu.txt
|
||||
train_resnet.txt
|
||||
train_resnet_one_gpu.txt
|
||||
train_bert.txt
|
||||
run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
|
||||
- name: Run process replay tests
|
||||
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
|
||||
|
||||
@@ -426,7 +366,7 @@ jobs:
|
||||
#- name: Test speed vs torch
|
||||
# run: |
|
||||
# python3 -c "import torch; print(torch.__version__)"
|
||||
# LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
|
||||
# LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
|
||||
- name: Test speed vs theoretical
|
||||
run: AMD=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
|
||||
- name: Test tensor cores AMD_LLVM=0
|
||||
@@ -437,7 +377,7 @@ jobs:
|
||||
- name: Run Tensor Core GEMM (AMD)
|
||||
run: |
|
||||
AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
|
||||
AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
|
||||
- name: Test AMD=1
|
||||
run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
|
||||
#- name: Test HIP=1
|
||||
@@ -452,61 +392,39 @@ jobs:
|
||||
- name: Test AM warm start time
|
||||
run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
|
||||
- name: Run Stable Diffusion
|
||||
run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
|
||||
run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
|
||||
- name: Run SDXL
|
||||
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
|
||||
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing
|
||||
- name: Run LLaMA 7B
|
||||
run: |
|
||||
BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
|
||||
BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
|
||||
BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run LLaMA 7B with BEAM
|
||||
run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
|
||||
run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
# - name: Run LLaMA 7B on 4 GPUs
|
||||
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_four_gpu.txt
|
||||
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
# - name: Run LLaMA 7B on 6 GPUs
|
||||
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_six_gpu.txt
|
||||
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run LLaMA-3 8B BEAM
|
||||
run: BENCHMARK_LOG=llama3_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
|
||||
run: BENCHMARK_LOG=llama3_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
|
||||
- name: Run LLaMA-3 8B on 4 GPUs with BEAM
|
||||
run: BENCHMARK_LOG=llama3_beam_4gpu AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
|
||||
run: BENCHMARK_LOG=llama3_beam_4gpu AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
|
||||
# - name: Run LLaMA-3 8B on 6 GPUs
|
||||
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
|
||||
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
|
||||
#- name: Restore amdgpu
|
||||
# run: sudo modprobe amdgpu
|
||||
# - name: Run LLaMA-2 70B
|
||||
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_2_70B.txt
|
||||
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run Mixtral 8x7B
|
||||
run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
|
||||
run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing
|
||||
- name: Run GPT2
|
||||
run: |
|
||||
BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
|
||||
BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
|
||||
BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
|
||||
- name: Run GPT2 w HALF
|
||||
run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
|
||||
run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing
|
||||
- name: Run GPT2 w HALF/BEAM
|
||||
run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (AMD)
|
||||
path: |
|
||||
onnx_inference_speed.csv
|
||||
torch_speed.txt
|
||||
llama_unjitted.txt
|
||||
llama_jitted.txt
|
||||
llama_beam.txt
|
||||
llama3_beam.txt
|
||||
llama3_four_gpu.txt
|
||||
llama3_six_gpu.txt
|
||||
llama_2_70B.txt
|
||||
gpt2_unjitted.txt
|
||||
gpt2_jitted.txt
|
||||
gpt2_half.txt
|
||||
gpt2_half_beam.txt
|
||||
matmul.txt
|
||||
matmul_amd.txt
|
||||
sd.txt
|
||||
sdxl.txt
|
||||
mixtral.txt
|
||||
run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
|
||||
- name: Run process replay tests
|
||||
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
|
||||
|
||||
@@ -543,31 +461,20 @@ jobs:
|
||||
- name: reset process replay
|
||||
run: test/external/process_replay/reset.py
|
||||
- name: Train MNIST
|
||||
run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
|
||||
run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
|
||||
- name: Run 10 CIFAR training steps
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py
|
||||
- name: Run 10 CIFAR training steps w HALF
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=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 | tee train_cifar_bf16.txt
|
||||
# run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
|
||||
# TODO: too slow
|
||||
# - name: Run 10 CIFAR training steps w winograd
|
||||
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
|
||||
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
|
||||
- name: Run full CIFAR training w 1 GPU
|
||||
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
|
||||
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
- name: Run full CIFAR training steps w 6 GPUS
|
||||
run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (AMD Training)
|
||||
path: |
|
||||
beautiful_mnist.txt
|
||||
train_cifar.txt
|
||||
train_cifar_half.txt
|
||||
train_cifar_bf16.txt
|
||||
train_cifar_wino.txt
|
||||
train_cifar_one_gpu.txt
|
||||
train_cifar_six_gpu.txt
|
||||
run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
- 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
|
||||
|
||||
@@ -606,19 +513,12 @@ jobs:
|
||||
- name: Run MLPerf resnet eval
|
||||
run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
|
||||
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
|
||||
run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
|
||||
run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
- name: Run 10 MLPerf Bert training steps (6 gpu)
|
||||
# TODO: remove BERT_LAYERS once scheduler is fast
|
||||
run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (AMD MLPerf)
|
||||
path: |
|
||||
train_resnet.txt
|
||||
train_resnet_one_gpu.txt
|
||||
train_bert.txt
|
||||
run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
|
||||
- name: Run process replay tests
|
||||
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
|
||||
|
||||
@@ -648,6 +548,8 @@ jobs:
|
||||
run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
|
||||
run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
|
||||
run: BENCHMARK_LOG=image_1_openpilot_0_10_1_vision PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: openpilot compile3 0.10.1 driving_vision
|
||||
run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: openpilot compile3 0.10.1 driving_policy
|
||||
@@ -708,7 +610,7 @@ jobs:
|
||||
# AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
|
||||
# AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
|
||||
- name: Run Tensor Core GEMM (AMD)
|
||||
run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee am_matmul_amd.txt
|
||||
run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
|
||||
- name: Test AMD=1
|
||||
run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
|
||||
- name: Test DISK copy time
|
||||
@@ -718,20 +620,12 @@ jobs:
|
||||
AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
|
||||
AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
|
||||
- name: Run full CIFAR training w 1 GPU
|
||||
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
|
||||
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
# - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
# run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
|
||||
# run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
- name: Run 10 MLPerf Bert training steps (1 gpu)
|
||||
# TODO: remove BERT_LAYERS once scheduler is fast
|
||||
run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (AM Driver)
|
||||
path: |
|
||||
am_matmul_amd.txt
|
||||
am_train_cifar_one_gpu.txt
|
||||
am_train_resnet_one_gpu.txt
|
||||
am_train_bert_one_gpu.txt
|
||||
run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
|
||||
- name: Run process replay tests
|
||||
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
|
||||
|
||||
@@ -778,21 +672,13 @@ jobs:
|
||||
NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
|
||||
NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
|
||||
- name: Test LLAMA-3
|
||||
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
|
||||
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0
|
||||
- name: Run full CIFAR training w 1 GPU
|
||||
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
|
||||
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
|
||||
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
|
||||
run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
|
||||
- name: Run 10 MLPerf Bert training steps (1 gpu)
|
||||
# TODO: remove BERT_LAYERS once scheduler is fast
|
||||
run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (NV Driver)
|
||||
path: |
|
||||
nv_llama3_beam.txt
|
||||
nv_train_cifar_one_gpu.txt
|
||||
nv_train_resnet_one_gpu.txt
|
||||
nv_train_bert_one_gpu.txt
|
||||
run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
|
||||
- name: Run process replay tests
|
||||
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
|
||||
|
||||
+55
-13
@@ -5,6 +5,7 @@ env:
|
||||
CAPTURE_PROCESS_REPLAY: 1
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
PYTHONPATH: ${{ github.workspace }}
|
||||
IGNORE_OOB: 0
|
||||
|
||||
on:
|
||||
push:
|
||||
@@ -36,6 +37,8 @@ jobs:
|
||||
name: Docs
|
||||
runs-on: ubuntu-22.04
|
||||
timeout-minutes: 10
|
||||
env:
|
||||
IGNORE_OOB: 1
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
@@ -102,15 +105,11 @@ jobs:
|
||||
run: |
|
||||
sudo apt update || true
|
||||
sudo apt install -y --no-install-recommends ninja-build
|
||||
- name: Lint with ruff
|
||||
run: |
|
||||
pip3 install --upgrade --force-reinstall ruff==0.11.0
|
||||
python3 -m ruff check extra/torch_backend/backend.py
|
||||
- name: Test one op
|
||||
run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
|
||||
- name: Test ResNet-18
|
||||
run: DEBUG=2 python3 extra/torch_backend/example.py
|
||||
- name: My (custom) tests
|
||||
- name: custom tests
|
||||
run: python3 extra/torch_backend/test.py
|
||||
- name: Test one op in torch tests
|
||||
run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
|
||||
@@ -219,7 +218,6 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
|
||||
# TODO: run the pre-commit hook to replace a lot of this
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
@@ -233,13 +231,14 @@ jobs:
|
||||
run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string=' ' --recursive=y .
|
||||
- name: Lint with ruff
|
||||
run: |
|
||||
pip3 install --upgrade --force-reinstall ruff==0.11.0
|
||||
python3 -m ruff check .
|
||||
pip3 install --upgrade --force-reinstall ruff==0.14.10
|
||||
pre-commit run ruff --all-files
|
||||
python3 -m ruff check examples/mlperf/ --ignore E501
|
||||
python3 -m ruff check extra/thunder/tiny/ --ignore E501 --ignore F841 --ignore E722
|
||||
python3 -m ruff check extra/torch_backend/backend.py
|
||||
- name: Run mypy
|
||||
run: |
|
||||
python -m mypy --strict-equality --lineprecision-report .
|
||||
python -m mypy --lineprecision-report .
|
||||
cat lineprecision.txt
|
||||
- name: Run TYPED=1
|
||||
run: TYPED=1 python -c "import tinygrad"
|
||||
@@ -258,6 +257,7 @@ jobs:
|
||||
key: unittest-12
|
||||
pydeps: "pillow numpy ftfy regex"
|
||||
deps: testing_unit
|
||||
llvm: 'true'
|
||||
- name: Check Device.DEFAULT
|
||||
run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
|
||||
- name: Run unit tests
|
||||
@@ -310,7 +310,7 @@ jobs:
|
||||
deps: testing_unit
|
||||
python-version: '3.14'
|
||||
- name: Test SPEC=2
|
||||
run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
|
||||
run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --ignore test/unit/test_autogen.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
|
||||
|
||||
fuzzing:
|
||||
name: Fuzzing
|
||||
@@ -473,6 +473,8 @@ jobs:
|
||||
name: Test LLM
|
||||
runs-on: ubuntu-24.04
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
IGNORE_OOB: 1
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
@@ -603,9 +605,7 @@ jobs:
|
||||
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/ --ignore=test/models --ignore=test/unit \
|
||||
--ignore=test/test_copy_speed.py --ignore=test/test_rearrange_einops.py \
|
||||
--ignore=test/test_fuzz_shape_ops.py --durations=20
|
||||
WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
@@ -654,6 +654,46 @@ jobs:
|
||||
- name: Run process replay tests
|
||||
uses: ./.github/actions/process-replay
|
||||
|
||||
testamdasm:
|
||||
name: AMD ASM IDE
|
||||
runs-on: ubuntu-24.04
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
- name: Setup Environment
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: rdna3-emu
|
||||
deps: testing_minimal
|
||||
amd: 'true'
|
||||
python-version: '3.13'
|
||||
- name: Verify AMD autogen is up to date
|
||||
run: |
|
||||
python -m extra.assembly.amd.pdf
|
||||
git diff --exit-code extra/assembly/amd/autogen/
|
||||
- name: Install LLVM 21
|
||||
run: |
|
||||
wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
|
||||
echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
|
||||
sudo apt-get update
|
||||
sudo apt-get install llvm-21 llvm-21-tools cloc
|
||||
- name: RDNA3 Line Count
|
||||
run: cloc --by-file extra/assembly/amd/*.py
|
||||
- name: Install rocprof-trace-decoder
|
||||
run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
|
||||
- name: Run RDNA3 emulator tests
|
||||
run: python -m pytest -n=auto extra/assembly/amd/ --durations 20
|
||||
- name: Run RDNA3 emulator tests (AMD_LLVM=1)
|
||||
run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
|
||||
- name: Run RDNA3 dtype tests
|
||||
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
|
||||
- name: Run RDNA3 dtype tests (AMD_LLVM=1)
|
||||
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
|
||||
# TODO: run all once emulator is faster
|
||||
- name: Run RDNA3 ops tests
|
||||
run: SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril"
|
||||
|
||||
testnvidia:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
@@ -741,6 +781,8 @@ jobs:
|
||||
ocelot: 'true'
|
||||
llvm: 'true'
|
||||
- name: Run unit tests
|
||||
env:
|
||||
LIBCLANG_PATH: '/opt/homebrew/opt/llvm@20/lib/libclang.dylib'
|
||||
run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
|
||||
- name: Run ONNX
|
||||
run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
|
||||
|
||||
@@ -16,7 +16,7 @@ repos:
|
||||
pass_filenames: false
|
||||
- id: mypy
|
||||
name: mypy
|
||||
entry: python3 -m mypy tinygrad/ --strict-equality
|
||||
entry: python3 -m mypy
|
||||
language: system
|
||||
always_run: true
|
||||
pass_filenames: false
|
||||
|
||||
@@ -34,33 +34,6 @@ result = graph_rewrite(uop, pm)
|
||||
### Schedule Cache
|
||||
Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.
|
||||
|
||||
## Directory Structure
|
||||
|
||||
```
|
||||
tinygrad/
|
||||
├── tensor.py # Tensor class, user API
|
||||
├── device.py # Buffer, device management
|
||||
├── dtype.py # Data types
|
||||
├── helpers.py # Utilities, environment vars
|
||||
├── uop/
|
||||
│ ├── ops.py # UOp class, Ops enum, PatternMatcher
|
||||
│ ├── spec.py # UOp type verification
|
||||
│ └── symbolic.py # Symbolic math simplification
|
||||
├── engine/
|
||||
│ ├── schedule.py # Schedule creation, caching
|
||||
│ ├── realize.py # Tensor realization
|
||||
│ ├── jit.py # JIT compilation
|
||||
│ └── memory.py # Memory planning
|
||||
├── schedule/
|
||||
│ ├── rangeify.py # Convert movements to ranges
|
||||
│ └── indexing.py # Index calculations
|
||||
├── codegen/
|
||||
│ ├── kernel.py # Kernel optimization
|
||||
│ └── uopgraph.py # UOp graph transformations
|
||||
├── renderer/ # Code generation (CUDA, Metal, etc.)
|
||||
└── runtime/ # Device backends
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
```bash
|
||||
@@ -79,7 +52,7 @@ VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
|
||||
|
||||
## Common Environment Variables
|
||||
|
||||
- `DEBUG=1-4` - Increasing verbosity
|
||||
- `DEBUG=1-7` - Increasing verbosity (7 shows assembly output)
|
||||
- `VIZ=1` - Enable graph visualization
|
||||
- `SPEC=1` - Enable UOp spec verification
|
||||
- `NOOPT=1` - Disable optimizations
|
||||
@@ -100,6 +73,16 @@ VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
|
||||
- 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
|
||||
@@ -209,9 +192,12 @@ When optimizing tinygrad internals:
|
||||
|
||||
9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.
|
||||
|
||||
## Pattern Matching Profiling
|
||||
## Pattern Matching Analysis
|
||||
|
||||
Use `TRACK_MATCH_STATS=2` to identify expensive patterns:
|
||||
**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
|
||||
@@ -225,3 +211,17 @@ Key patterns to watch (from ResNet50 benchmark):
|
||||
- `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.
|
||||
|
||||
+1
-1
@@ -70,7 +70,7 @@ AMD backend supports several interfaces for communicating with devices:
|
||||
|
||||
* `KFD`: uses the amdgpu driver
|
||||
* `PCI`: uses the [AM driver](developer/am.md)
|
||||
* `USB`: USB3 interafce for asm24xx chips.
|
||||
* `USB`: USB3 interface for asm24xx chips.
|
||||
|
||||
You can force an interface by setting `AMD_IFACE` to one of these values. In the case of `AMD_IFACE=PCI`, this may unbind your GPU from the amdgpu driver.
|
||||
|
||||
|
||||
@@ -213,12 +213,13 @@ class InterleavedDataset:
|
||||
self.queues[queue_index].queue.extend(load_file(file))
|
||||
|
||||
# Reference: https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/run_pretraining.py, Line 394
|
||||
def batch_load_train_bert(BS:int):
|
||||
def batch_load_train_bert(BS:int, seed:int|None=None):
|
||||
from extra.datasets.wikipedia import get_wiki_train_files
|
||||
rng = random.Random(seed)
|
||||
fs = sorted(get_wiki_train_files())
|
||||
train_files = []
|
||||
while fs: # TF shuffle
|
||||
random.shuffle(fs)
|
||||
rng.shuffle(fs)
|
||||
train_files.append(fs.pop(0))
|
||||
|
||||
cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
|
||||
|
||||
@@ -219,7 +219,18 @@ def get_mlperf_bert_model():
|
||||
config = get_mlperf_bert_config()
|
||||
if getenv("DISABLE_DROPOUT", 0):
|
||||
config["hidden_dropout_prob"] = config["attention_probs_dropout_prob"] = 0.0
|
||||
return BertForPretraining(**config)
|
||||
model = BertForPretraining(**config)
|
||||
if getenv("FP8_TRAIN"):
|
||||
from extra.fp8.fp8_linear import convert_to_float8_training
|
||||
def module_filter_fn(mod, fqn):
|
||||
if isinstance(mod, LinearBert):
|
||||
skip_layers = [] if (ln:=config["num_hidden_layers"]) <= 2 else ["bert.encoder.layer.0.", f"bert.encoder.layer.{ln-1}"]
|
||||
if mod.weight.shape[-1] >= 1024 and "encoder" in fqn and not any(name in fqn for name in skip_layers):
|
||||
print(f"replacing linear with fp8: {fqn} {mod.weight.shape}")
|
||||
return True
|
||||
return False
|
||||
convert_to_float8_training(model, module_filter_fn)
|
||||
return model
|
||||
|
||||
def get_fake_data_bert(BS:int):
|
||||
return {
|
||||
|
||||
@@ -1008,6 +1008,7 @@ def train_bert():
|
||||
config["DISABLE_DROPOUT"] = getenv("DISABLE_DROPOUT", 0)
|
||||
config["TRAIN_BEAM"] = TRAIN_BEAM = getenv("TRAIN_BEAM", BEAM.value)
|
||||
config["EVAL_BEAM"] = EVAL_BEAM = getenv("EVAL_BEAM", BEAM.value)
|
||||
config["FP8_TRAIN"] = getenv("FP8_TRAIN", 0)
|
||||
|
||||
Tensor.manual_seed(seed) # seed for weight initialization
|
||||
|
||||
@@ -1085,7 +1086,7 @@ def train_bert():
|
||||
if RUNMLPERF:
|
||||
# only load real data with RUNMLPERF
|
||||
eval_it = iter(batch_load_val_bert(EVAL_BS))
|
||||
train_it = iter(tqdm(batch_load_train_bert(BS), total=train_steps, disable=BENCHMARK))
|
||||
train_it = iter(tqdm(batch_load_train_bert(BS, seed=seed), total=train_steps, disable=BENCHMARK))
|
||||
for _ in range(start_step): next(train_it) # Fast forward
|
||||
else:
|
||||
# repeat fake data
|
||||
@@ -1147,7 +1148,7 @@ def train_bert():
|
||||
|
||||
device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
|
||||
loss = loss.item()
|
||||
assert not math.isnan(loss)
|
||||
if not getenv("FP8_TRAIN"): assert not math.isnan(loss)
|
||||
lr = lr.item()
|
||||
|
||||
cl = time.perf_counter()
|
||||
@@ -1160,7 +1161,7 @@ def train_bert():
|
||||
if WANDB:
|
||||
wandb.log({"lr": lr, "train/loss": loss, "train/global_norm": global_norm.item(), "train/step_time": cl - st,
|
||||
"train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
|
||||
"train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
|
||||
"train/mem":GlobalCounters.mem_used / 1e9, "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
|
||||
|
||||
train_data, next_data = next_data, None
|
||||
i += 1
|
||||
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." AMD=1
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
|
||||
# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
|
||||
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
|
||||
export TRAIN_STEPS=3900
|
||||
|
||||
export IGNORE_OOB=1
|
||||
export REWRITE_STACK_LIMIT=5000000
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
export BEAM_TIMEOUT_SEC=15
|
||||
export FP8_TRAIN=1
|
||||
# search
|
||||
IGNORE_BEAM_CACHE=1 BENCHMARK=10 BERT_LAYERS=2 RUNMLPERF=0 python3 examples/mlperf/model_train.py
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
RUNMLPERF=1 python3 examples/mlperf/model_train.py
|
||||
@@ -7,7 +7,7 @@ if __name__ == "__main__":
|
||||
with open(fetch(sys.argv[1]), "rb") as f:
|
||||
run_onnx_jit = pickle.load(f)
|
||||
input_name = run_onnx_jit.captured.expected_names[0]
|
||||
device = run_onnx_jit.captured.expected_st_vars_dtype_device[0][-1]
|
||||
device = run_onnx_jit.captured.expected_input_info[0][-1]
|
||||
print(f"input goes into {input_name=} on {device=}")
|
||||
hit = 0
|
||||
for i,(img,y) in enumerate(imagenet_dataloader(cnt=getenv("CNT", 100))):
|
||||
|
||||
+10
-11
@@ -153,8 +153,7 @@ class SMICtx:
|
||||
tables = {}
|
||||
for dev in self.devs:
|
||||
match dev.ip_ver[am.MP1_HWIP]:
|
||||
case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableX_t
|
||||
case (13,0,12): table_t = dev.smu.smu_mod.MetricsTableV2_t
|
||||
case (13,0,6)|(13,0,12): table_t = dev.smu.smu_mod.MetricsTableX_t
|
||||
case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
|
||||
tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
|
||||
return tables
|
||||
@@ -165,17 +164,17 @@ class SMICtx:
|
||||
|
||||
def get_gfx_activity(self, dev, metrics):
|
||||
match dev.ip_ver[am.MP1_HWIP]:
|
||||
case (13,0,6): return max(0, min(100, self._smuq10_round(metrics.SocketGfxBusy)))
|
||||
case (13,0,6)|(13,0,12): return max(0, min(100, self._smuq10_round(metrics.SocketGfxBusy)))
|
||||
case _: return metrics.SmuMetrics.AverageGfxActivity
|
||||
|
||||
def get_mem_activity(self, dev, metrics):
|
||||
match dev.ip_ver[am.MP1_HWIP]:
|
||||
case (13,0,6): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
|
||||
case (13,0,6)|(13,0,12): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
|
||||
case _: return metrics.SmuMetrics.AverageUclkActivity
|
||||
|
||||
def get_temps(self, dev, metrics, compact=False):
|
||||
match dev.ip_ver[am.MP1_HWIP]:
|
||||
case (13,0,6):
|
||||
case (13,0,6)|(13,0,12):
|
||||
temps = {
|
||||
"Hotspot": self._smuq10_round(metrics.MaxSocketTemperature),
|
||||
"HBM": self._smuq10_round(metrics.MaxHbmTemperature),
|
||||
@@ -191,7 +190,7 @@ class SMICtx:
|
||||
|
||||
def get_voltage(self, dev, metrics, compact=False):
|
||||
match dev.ip_ver[am.MP1_HWIP]:
|
||||
case (13,0,6): return {}
|
||||
case (13,0,6)|(13,0,12): return {}
|
||||
case _:
|
||||
voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.SVI_PLANE_e.items()
|
||||
if k < dev.smu.smu_mod.SVI_PLANE_COUNT and metrics.SmuMetrics.AvgVoltage[k] != 0]
|
||||
@@ -205,33 +204,33 @@ class SMICtx:
|
||||
def get_gfx_freq(self, dev, metrics):
|
||||
if metrics is None: return 0
|
||||
match dev.ip_ver[am.MP1_HWIP]:
|
||||
case (13,0,6): return self._smuq10_round(metrics.GfxclkFrequency[0])
|
||||
case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.GfxclkFrequency[0])
|
||||
case _:
|
||||
return metrics.SmuMetrics.AverageGfxclkFrequencyPostDs if self.get_gfx_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
|
||||
metrics.SmuMetrics.AverageGfxclkFrequencyPreDs
|
||||
|
||||
def get_mem_freq(self, dev, metrics):
|
||||
match dev.ip_ver[am.MP1_HWIP]:
|
||||
case (13,0,6): return self._smuq10_round(metrics.UclkFrequency)
|
||||
case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.UclkFrequency)
|
||||
case _:
|
||||
return metrics.SmuMetrics.AverageMemclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
|
||||
metrics.SmuMetrics.AverageMemclkFrequencyPreDs
|
||||
|
||||
def get_fckl_freq(self, dev, metrics):
|
||||
match dev.ip_ver[am.MP1_HWIP]:
|
||||
case (13,0,6): return self._smuq10_round(metrics.FclkFrequency)
|
||||
case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.FclkFrequency)
|
||||
case _:
|
||||
return metrics.SmuMetrics.AverageFclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
|
||||
metrics.SmuMetrics.AverageFclkFrequencyPreDs
|
||||
|
||||
def get_fan_rpm_pwm(self, dev, metrics):
|
||||
match dev.ip_ver[am.MP1_HWIP]:
|
||||
case (13,0,6): return None, None
|
||||
case (13,0,6)|(13,0,12): return None, None
|
||||
case _: return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm
|
||||
|
||||
def get_power(self, dev, metrics):
|
||||
match dev.ip_ver[am.MP1_HWIP]:
|
||||
case (13,0,6): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
|
||||
case (13,0,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):
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
An integrated environment for AMD GPU assembly and emulation
|
||||
|
||||
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
|
||||
* 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. asm and emu shouldn't be required for dsl.
|
||||
|
||||
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.
|
||||
|
||||
test_llvm.py tests asm/disasm on the LLVM tests, confirming it behaves the same as LLVM.
|
||||
|
||||
tinygrad's dtype tests should pass with and without LLVM. they run in about 12 seconds.
|
||||
|
||||
`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 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 `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 ~2000 lines.
|
||||
Get line count with `cloc --by-file extra/assembly/amd/*.py`
|
||||
|
||||
@@ -0,0 +1,444 @@
|
||||
# RDNA3/RDNA4/CDNA assembler
|
||||
from __future__ import annotations
|
||||
import re
|
||||
from extra.assembly.amd.dsl import RawImm, SrcMod, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory
|
||||
from extra.assembly.amd.dsl import VCC_LO, VCC_HI, VCC, EXEC_LO, EXEC_HI, EXEC, SCC, M0, NULL, OFF
|
||||
from extra.assembly.amd.dsl import FLOAT_ENC
|
||||
from extra.assembly.amd.autogen.rdna3 import ins
|
||||
from extra.assembly.amd.autogen.rdna3.ins import VOP2Op, VOPDOp, SOPKOp
|
||||
from extra.assembly.amd.autogen.rdna3.enum import BufFmt
|
||||
from extra.assembly.amd.autogen.rdna4 import ins as rdna4_ins
|
||||
|
||||
# Re-export disasm for backwards compatibility
|
||||
from extra.assembly.amd.disasm import disasm, HWREG, HWREG_RDNA4
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# CONSTANTS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
# RDNA unified buffer format
|
||||
BUF_FMT = {e.name: e.value for e in BufFmt}
|
||||
_BUF_FMT_EXT = {'BUF_FMT_32_32_32_32_SINT': 62, 'BUF_FMT_32_32_32_32_FLOAT': 63, 'BUF_FMT_8_FLOAT': 108}
|
||||
BUF_FMT.update(_BUF_FMT_EXT)
|
||||
def _parse_buf_fmt_combo(s: str) -> int:
|
||||
parts = [p.strip().replace('BUF_DATA_FORMAT_', '').replace('BUF_NUM_FORMAT_', '') for p in s.split(',')]
|
||||
return BUF_FMT.get(f'BUF_FMT_{parts[0]}_{parts[1]}') if len(parts) == 2 else None
|
||||
|
||||
def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
|
||||
return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# ASSEMBLER
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
SPEC_REGS = {'vcc_lo': RawImm(106), 'vcc_hi': RawImm(107), 'vcc': RawImm(106), 'null': RawImm(124), 'off': RawImm(124), 'm0': RawImm(125),
|
||||
'exec_lo': RawImm(126), 'exec_hi': RawImm(127), 'exec': RawImm(126), 'scc': RawImm(253), 'src_scc': RawImm(253)}
|
||||
FLOATS = {str(k): k for k in FLOAT_ENC} # Valid float literal strings: '0.5', '-0.5', '1.0', etc.
|
||||
REG_MAP: dict[str, _RegFactory] = {'s': s, 'v': v, 't': ttmp, 'ttmp': ttmp}
|
||||
SMEM_OPS = {'s_load_b32', 's_load_b64', 's_load_b96', 's_load_b128', 's_load_b256', 's_load_b512',
|
||||
's_load_i8', 's_load_u8', 's_load_i16', 's_load_u16',
|
||||
's_buffer_load_b32', 's_buffer_load_b64', 's_buffer_load_b96', 's_buffer_load_b128', 's_buffer_load_b256', 's_buffer_load_b512',
|
||||
's_buffer_load_i8', 's_buffer_load_u8', 's_buffer_load_i16', 's_buffer_load_u16',
|
||||
's_atc_probe', 's_atc_probe_buffer'}
|
||||
SPEC_DSL = {'vcc_lo': 'VCC_LO', 'vcc_hi': 'VCC_HI', 'vcc': 'VCC_LO', 'null': 'NULL', 'off': 'OFF', 'm0': 'M0',
|
||||
'exec_lo': 'EXEC_LO', 'exec_hi': 'EXEC_HI', 'exec': 'EXEC_LO', 'scc': 'SCC', 'src_scc': 'SCC'}
|
||||
|
||||
def _op2dsl(op: str) -> str:
|
||||
op = op.strip()
|
||||
neg = op.startswith('-') and not (op[1:2].isdigit() or (len(op) > 2 and op[1] == '0' and op[2] in 'xX'))
|
||||
if neg: op = op[1:]
|
||||
if op.startswith('neg(') and op.endswith(')'): neg = True; op = op[4:-1]
|
||||
abs_ = (op.startswith('|') and op.endswith('|')) or (op.startswith('abs(') and op.endswith(')'))
|
||||
if abs_: op = op[1:-1] if op.startswith('|') else op[4:-1]
|
||||
hi = ".h" if op.endswith('.h') else ".l" if op.endswith('.l') else ""
|
||||
if hi: op = op[:-2]
|
||||
lo = op.lower()
|
||||
def wrap(b): return f"{'-' if neg else ''}abs({b}){hi}" if abs_ else f"-{b}{hi}" if neg else f"{b}{hi}"
|
||||
if lo in SPEC_DSL: return wrap(SPEC_DSL[lo])
|
||||
if op in FLOATS: return wrap(op)
|
||||
rp = {'s': 's', 'v': 'v', 't': 'ttmp', 'ttmp': 'ttmp'}
|
||||
if m := re.match(r'^([svt](?:tmp)?)\[(\d+):(\d+)\]$', lo): return wrap(f"{rp[m.group(1)]}[{m.group(2)}:{m.group(3)}]")
|
||||
if m := re.match(r'^([svt](?:tmp)?)(\d+)$', lo): return wrap(f"{rp[m.group(1)]}[{m.group(2)}]")
|
||||
if re.match(r'^-?\d+$|^-?0x[0-9a-fA-F]+$', op): return f"SrcMod({op}, neg={neg}, abs_={abs_})" if neg or abs_ else op
|
||||
return wrap(op)
|
||||
|
||||
def _parse_ops(s: str) -> list[str]:
|
||||
ops, cur, depth, pipe = [], "", 0, False
|
||||
for c in s:
|
||||
if c in '[(': depth += 1
|
||||
elif c in '])': depth -= 1
|
||||
elif c == '|': pipe = not pipe
|
||||
if c == ',' and depth == 0 and not pipe: ops.append(cur.strip()); cur = ""
|
||||
else: cur += c
|
||||
if cur.strip(): ops.append(cur.strip())
|
||||
return ops
|
||||
|
||||
def _extract(text: str, pat: str, flags=re.I):
|
||||
if m := re.search(pat, text, flags): return m, text[:m.start()] + ' ' + text[m.end():]
|
||||
return None, text
|
||||
|
||||
# Instruction aliases: LLVM uses different names for some instructions
|
||||
_ALIASES = {
|
||||
'v_cmp_tru_f16': 'v_cmp_t_f16', 'v_cmp_tru_f32': 'v_cmp_t_f32', 'v_cmp_tru_f64': 'v_cmp_t_f64',
|
||||
'v_cmpx_tru_f16': 'v_cmpx_t_f16', 'v_cmpx_tru_f32': 'v_cmpx_t_f32', 'v_cmpx_tru_f64': 'v_cmpx_t_f64',
|
||||
'v_cvt_flr_i32_f32': 'v_cvt_floor_i32_f32', 'v_cvt_rpi_i32_f32': 'v_cvt_nearest_i32_f32',
|
||||
'v_ffbh_i32': 'v_cls_i32', 'v_ffbh_u32': 'v_clz_i32_u32', 'v_ffbl_b32': 'v_ctz_i32_b32',
|
||||
'v_cvt_pkrtz_f16_f32': 'v_cvt_pk_rtz_f16_f32', 'v_fmac_legacy_f32': 'v_fmac_dx9_zero_f32', 'v_mul_legacy_f32': 'v_mul_dx9_zero_f32',
|
||||
's_load_dword': 's_load_b32', 's_load_dwordx2': 's_load_b64', 's_load_dwordx4': 's_load_b128',
|
||||
's_load_dwordx8': 's_load_b256', 's_load_dwordx16': 's_load_b512',
|
||||
's_buffer_load_dword': 's_buffer_load_b32', 's_buffer_load_dwordx2': 's_buffer_load_b64',
|
||||
's_buffer_load_dwordx4': 's_buffer_load_b128', 's_buffer_load_dwordx8': 's_buffer_load_b256',
|
||||
's_buffer_load_dwordx16': 's_buffer_load_b512',
|
||||
'v_cvt_pknorm_i16_f16': 'v_cvt_pk_norm_i16_f16', 'v_cvt_pknorm_u16_f16': 'v_cvt_pk_norm_u16_f16',
|
||||
'v_add3_nc_u32': 'v_add3_u32', 'v_xor_add_u32': 'v_xad_u32',
|
||||
'v_interp_p2_new_f32': 'v_interp_p2_f32',
|
||||
's_ff1_i32_b32': 's_ctz_i32_b32', 's_ff1_i32_b64': 's_ctz_i32_b64',
|
||||
's_flbit_i32_b32': 's_clz_i32_u32', 's_flbit_i32_b64': 's_clz_i32_u64', 's_flbit_i32': 's_cls_i32', 's_flbit_i32_i64': 's_cls_i32_i64',
|
||||
's_andn1_saveexec_b32': 's_and_not0_saveexec_b32', 's_andn1_saveexec_b64': 's_and_not0_saveexec_b64',
|
||||
's_andn1_wrexec_b32': 's_and_not0_wrexec_b32', 's_andn1_wrexec_b64': 's_and_not0_wrexec_b64',
|
||||
's_andn2_saveexec_b32': 's_and_not1_saveexec_b32', 's_andn2_saveexec_b64': 's_and_not1_saveexec_b64',
|
||||
's_andn2_wrexec_b32': 's_and_not1_wrexec_b32', 's_andn2_wrexec_b64': 's_and_not1_wrexec_b64',
|
||||
's_orn1_saveexec_b32': 's_or_not0_saveexec_b32', 's_orn1_saveexec_b64': 's_or_not0_saveexec_b64',
|
||||
's_orn2_saveexec_b32': 's_or_not1_saveexec_b32', 's_orn2_saveexec_b64': 's_or_not1_saveexec_b64',
|
||||
's_andn2_b32': 's_and_not1_b32', 's_andn2_b64': 's_and_not1_b64',
|
||||
's_orn2_b32': 's_or_not1_b32', 's_orn2_b64': 's_or_not1_b64',
|
||||
'v_dot2c_f32_f16': 'v_dot2acc_f32_f16',
|
||||
'v_fma_legacy_f32': 'v_fma_dx9_zero_f32',
|
||||
'ds_read_b32': 'ds_load_b32', 'ds_read_b64': 'ds_load_b64', 'ds_read_b96': 'ds_load_b96', 'ds_read_b128': 'ds_load_b128',
|
||||
'ds_read_i8': 'ds_load_i8', 'ds_read_u8': 'ds_load_u8', 'ds_read_i16': 'ds_load_i16', 'ds_read_u16': 'ds_load_u16',
|
||||
'ds_read_i8_d16': 'ds_load_i8_d16', 'ds_read_u8_d16': 'ds_load_u8_d16', 'ds_read_i8_d16_hi': 'ds_load_i8_d16_hi', 'ds_read_u8_d16_hi': 'ds_load_u8_d16_hi',
|
||||
'ds_read_u16_d16': 'ds_load_u16_d16', 'ds_read_u16_d16_hi': 'ds_load_u16_d16_hi',
|
||||
'ds_read2_b32': 'ds_load_2addr_b32', 'ds_read2_b64': 'ds_load_2addr_b64',
|
||||
'ds_read2st64_b32': 'ds_load_2addr_stride64_b32', 'ds_read2st64_b64': 'ds_load_2addr_stride64_b64',
|
||||
'ds_read_addtid_b32': 'ds_load_addtid_b32', 'ds_write_addtid_b32': 'ds_store_addtid_b32',
|
||||
'ds_write_b32': 'ds_store_b32', 'ds_write_b64': 'ds_store_b64', 'ds_write_b96': 'ds_store_b96', 'ds_write_b128': 'ds_store_b128',
|
||||
'ds_write_b8': 'ds_store_b8', 'ds_write_b16': 'ds_store_b16',
|
||||
'ds_write_b8_d16_hi': 'ds_store_b8_d16_hi', 'ds_write_b16_d16_hi': 'ds_store_b16_d16_hi',
|
||||
'ds_write2_b32': 'ds_store_2addr_b32', 'ds_write2_b64': 'ds_store_2addr_b64',
|
||||
'ds_write2st64_b32': 'ds_store_2addr_stride64_b32', 'ds_write2st64_b64': 'ds_store_2addr_stride64_b64',
|
||||
'ds_wrxchg_rtn_b32': 'ds_storexchg_rtn_b32', 'ds_wrxchg_rtn_b64': 'ds_storexchg_rtn_b64',
|
||||
'ds_wrxchg2_rtn_b32': 'ds_storexchg_2addr_rtn_b32', 'ds_wrxchg2_rtn_b64': 'ds_storexchg_2addr_rtn_b64',
|
||||
'ds_wrxchg2st64_rtn_b32': 'ds_storexchg_2addr_stride64_rtn_b32', 'ds_wrxchg2st64_rtn_b64': 'ds_storexchg_2addr_stride64_rtn_b64',
|
||||
}
|
||||
|
||||
def _apply_alias(text: str) -> str:
|
||||
mn = text.split()[0].lower() if ' ' in text else text.lower().rstrip('_')
|
||||
for m in (mn, mn.removesuffix('_e32'), mn.removesuffix('_e64')):
|
||||
if m in _ALIASES: return _ALIASES[m] + text[len(m):]
|
||||
return text
|
||||
|
||||
def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
|
||||
|
||||
def get_dsl(text: str, arch: str = "rdna3") -> str:
|
||||
text, kw = _apply_alias(text.strip()), []
|
||||
# Extract modifiers
|
||||
for pat, val in [(r'\s+mul:2(?:\s|$)', 1), (r'\s+mul:4(?:\s|$)', 2), (r'\s+div:2(?:\s|$)', 3)]:
|
||||
if (m := _extract(text, pat))[0]: kw.append(f'omod={val}'); text = m[1]; break
|
||||
clamp_found = False
|
||||
if (m := _extract(text, r'\s+clamp(?:\s|$)'))[0]: clamp_found = True; text = m[1]
|
||||
opsel, m, text = None, *_extract(text, r'\s+op_sel:\[([^\]]+)\]')
|
||||
if m:
|
||||
bits, mn = [int(x.strip()) for x in m.group(1).split(',')], text.split()[0].lower()
|
||||
is3p = mn.startswith(('v_pk_', 'v_wmma_', 'v_dot', 'v_fma_mix'))
|
||||
opsel = (bits[0] | (bits[1] << 1) | (bits[2] << 2)) if len(bits) == 3 and is3p else \
|
||||
(bits[0] | (bits[1] << 1) | (bits[2] << 3)) if len(bits) == 3 else sum(b << i for i, b in enumerate(bits))
|
||||
opsel_hi_val, m, text = None, *_extract(text, r'\s+op_sel_hi:\[([^\]]+)\]')
|
||||
if m: opsel_hi_val = [int(x.strip()) for x in m.group(1).split(',')]
|
||||
m, text = _extract(text, r'\s+wait_exp:(\d+)'); waitexp = m.group(1) if m else None
|
||||
m, text = _extract(text, r'\s+offset:(0x[0-9a-fA-F]+|-?\d+)'); off_val = m.group(1) if m else None
|
||||
m, text = _extract(text, r'\s+dlc(?:\s|$)'); dlc = 1 if m else None
|
||||
m, text = _extract(text, r'\s+glc(?:\s|$)'); glc = 1 if m else None
|
||||
m, text = _extract(text, r'\s+slc(?:\s|$)'); slc = 1 if m else None
|
||||
m, text = _extract(text, r'\s+tfe(?:\s|$)'); tfe = 1 if m else None
|
||||
m, text = _extract(text, r'\s+offen(?:\s|$)'); offen = 1 if m else None
|
||||
m, text = _extract(text, r'\s+idxen(?:\s|$)'); idxen = 1 if m else None
|
||||
m, text = _extract(text, r'\s+format:\[([^\]]+)\]'); fmt_val = m.group(1) if m else None
|
||||
m, text = _extract(text, r'\s+format:(\d+)'); fmt_val = m.group(1) if m and not fmt_val else fmt_val
|
||||
m, text = _extract(text, r'\s+neg_lo:\[([^\]]+)\]'); neg_lo = sum(int(x.strip()) << i for i, x in enumerate(m.group(1).split(','))) if m else None
|
||||
m, text = _extract(text, r'\s+neg_hi:\[([^\]]+)\]'); neg_hi = sum(int(x.strip()) << i for i, x in enumerate(m.group(1).split(','))) if m else None
|
||||
m, text = _extract(text, r'\s+byte_sel:(\d+)'); byte_sel = int(m.group(1)) if m else None
|
||||
m, text = _extract(text, r'\s+offset0:(\d+)'); ds_off0 = int(m.group(1)) if m else None
|
||||
m, text = _extract(text, r'\s+offset1:(\d+)'); ds_off1 = int(m.group(1)) if m else None
|
||||
m, text = _extract(text, r'\s+index_key:(\d+)'); index_key = int(m.group(1)) if m else None
|
||||
if waitexp: kw.append(f'waitexp={waitexp}')
|
||||
if byte_sel is not None:
|
||||
if opsel is None: opsel = 0
|
||||
opsel |= (byte_sel << 2)
|
||||
if ds_off0 is not None: kw.append(f'offset0={ds_off0}')
|
||||
if ds_off1 is not None: kw.append(f'offset1={ds_off1}')
|
||||
if index_key is not None: kw.append(f'opsel={index_key}')
|
||||
|
||||
parts = text.replace(',', ' ').split()
|
||||
if not parts: raise ValueError("empty instruction")
|
||||
mn, op_str = parts[0].lower(), text[len(parts[0]):].strip()
|
||||
ops, args = _parse_ops(op_str), [_op2dsl(o) for o in _parse_ops(op_str)]
|
||||
|
||||
# s_waitcnt
|
||||
if mn == 's_waitcnt':
|
||||
vm, exp, lgkm = 0x3f, 0x7, 0x3f
|
||||
for p in op_str.replace(',', ' ').split():
|
||||
if m := re.match(r'vmcnt\((\d+)\)', p): vm = int(m.group(1))
|
||||
elif m := re.match(r'expcnt\((\d+)\)', p): exp = int(m.group(1))
|
||||
elif m := re.match(r'lgkmcnt\((\d+)\)', p): lgkm = int(m.group(1))
|
||||
elif re.match(r'^0x[0-9a-f]+$|^\d+$', p): return f"s_waitcnt(simm16={int(p, 0)})"
|
||||
return f"s_waitcnt(simm16={waitcnt(vm, exp, lgkm)})"
|
||||
|
||||
# VOPD
|
||||
if '::' in text:
|
||||
xp, yp = text.split('::')
|
||||
xps, yps = xp.strip().replace(',', ' ').split(), yp.strip().replace(',', ' ').split()
|
||||
xo, yo = [_op2dsl(p) for p in xps[1:]], [_op2dsl(p) for p in yps[1:]]
|
||||
vdx, sx0, vsx1 = xo[0], xo[1] if len(xo) > 1 else '0', xo[2] if len(xo) > 2 else 'v[0]'
|
||||
vdy, sy0, vsy1 = yo[0], yo[1] if len(yo) > 1 else '0', yo[2] if len(yo) > 2 else 'v[0]'
|
||||
lit = xo[3] if 'fmaak' in xps[0].lower() and len(xo) > 3 else yo[3] if 'fmaak' in yps[0].lower() and len(yo) > 3 else None
|
||||
if 'fmamk' in xps[0].lower() and len(xo) > 3: lit, vsx1 = xo[2], xo[3]
|
||||
elif 'fmamk' in yps[0].lower() and len(yo) > 3: lit, vsy1 = yo[2], yo[3]
|
||||
return f"VOPD(VOPDOp.{xps[0].upper()}, VOPDOp.{yps[0].upper()}, vdstx={vdx}, vdsty={vdy}, srcx0={sx0}, vsrcx1={vsx1}, srcy0={sy0}, vsrcy1={vsy1}{f', literal={lit}' if lit else ''})"
|
||||
|
||||
# Special instructions
|
||||
if mn == 's_setreg_imm32_b32': raise ValueError(f"unsupported: {mn}")
|
||||
sop1_no_dest = ('s_alloc_vgpr', 's_barrier_init', 's_barrier_join', 's_barrier_signal', 's_barrier_signal_isfirst', 's_sleep_var')
|
||||
if mn in sop1_no_dest:
|
||||
return f"{mn}(sdst=RawImm(128), ssrc0={args[0]})"
|
||||
if mn in ('s_setpc_b64', 's_rfe_b64'): return f"{mn}(ssrc0={args[0]})"
|
||||
if mn in ('s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'): return f"{mn}(sdst={args[0]}, ssrc0=RawImm({args[1].strip()}))"
|
||||
if mn == 's_version': return f"{mn}(simm16={args[0]})"
|
||||
if mn == 's_setreg_b32': return f"{mn}(simm16={args[0]}, sdst={args[1]})"
|
||||
|
||||
# Export instructions (RDNA4 VEXPORT)
|
||||
if mn == 'export':
|
||||
target_map = {**{f'mrt{i}': i for i in range(8)}, 'mrtz': 8, **{f'pos{i}': 12+i for i in range(4)}}
|
||||
m, exp_str = _extract(op_str, r'\s+done(?:\s|$)')
|
||||
done_val = 1 if m else 0
|
||||
exp_parts = exp_str.replace(',', ' ').split()
|
||||
target_name = exp_parts[0].lower().strip()
|
||||
target = target_map.get(target_name, 0)
|
||||
vsrcs, en = [], 0
|
||||
for i, o in enumerate(exp_parts[1:5]):
|
||||
o = o.strip().lower()
|
||||
if o == 'off': vsrcs.append('v[0]')
|
||||
else: vsrcs.append(_op2dsl(o)); en |= (1 << i)
|
||||
return f"VEXPORT(target={target}, en={en}, vsrc0={vsrcs[0]}, vsrc1={vsrcs[1]}, vsrc2={vsrcs[2]}, vsrc3={vsrcs[3]}, done={done_val})"
|
||||
|
||||
# SMEM
|
||||
if mn in SMEM_OPS:
|
||||
gs, ds = ", glc=1" if glc else "", ", dlc=1" if dlc else ""
|
||||
off_field = "ioffset" if arch == "rdna4" else "offset"
|
||||
th_s, scope_s, smem_str = "", "", op_str
|
||||
if arch == "rdna4":
|
||||
m, smem_str = _extract(op_str, r'\s+th:TH_(\w+)')
|
||||
th_val = {'LOAD_RT': 0, 'LOAD_NT': 1, 'LOAD_HT': 2, 'LOAD_LU': 3, 'STORE_RT': 0, 'STORE_NT': 1, 'STORE_HT': 2, 'STORE_LU': 3}.get(m.group(1), 0) if m else None
|
||||
m, smem_str = _extract(smem_str, r'\s+scope:SCOPE_(\w+)')
|
||||
scope_val = {'CU': 0, 'SE': 1, 'DEV': 2, 'SYS': 3}.get(m.group(1), 0) if m else None
|
||||
if scope_val is None:
|
||||
m, smem_str = _extract(smem_str, r'\s+scope:(0?x?[0-9a-fA-F]+)')
|
||||
scope_val = int(m.group(1), 0) if m else None
|
||||
th_s = f", th={th_val}" if th_val else ""
|
||||
scope_s = f", scope={scope_val}" if scope_val else ""
|
||||
smem_ops = _parse_ops(smem_str)
|
||||
smem_args = [_op2dsl(o) for o in smem_ops]
|
||||
if len(smem_ops) >= 3 and re.match(r'^-?[0-9]|^-?0x', smem_ops[2].strip().lower()):
|
||||
return f"{mn}(sdata={smem_args[0]}, sbase={smem_args[1]}, {off_field}={smem_ops[2].strip()}, soffset=RawImm(124){gs}{ds}{th_s}{scope_s})"
|
||||
if off_val and len(smem_ops) >= 3: return f"{mn}(sdata={smem_args[0]}, sbase={smem_args[1]}, {off_field}={off_val}, soffset={smem_args[2]}{gs}{ds}{th_s}{scope_s})"
|
||||
if len(smem_ops) >= 3: return f"{mn}(sdata={smem_args[0]}, sbase={smem_args[1]}, soffset={smem_args[2]}{gs}{ds}{th_s}{scope_s})"
|
||||
|
||||
# Buffer (MUBUF/MTBUF/VBUFFER) instructions
|
||||
if mn.startswith(('buffer_', 'tbuffer_')):
|
||||
is_tbuf = mn.startswith('tbuffer_')
|
||||
fmt_num = None
|
||||
if fmt_val is not None:
|
||||
if fmt_val.isdigit(): fmt_num = int(fmt_val)
|
||||
else: fmt_num = BUF_FMT.get(fmt_val.replace(' ', '')) or _parse_buf_fmt_combo(fmt_val)
|
||||
if mn in ('buffer_gl0_inv', 'buffer_gl1_inv', 'buffer_wbl2', 'buffer_inv'): return f"{mn}()"
|
||||
if arch == "rdna4":
|
||||
m, buf_text = _extract(op_str, r'\s+th:TH_(\w+)')
|
||||
th_val = {'LOAD_RT': 0, 'LOAD_NT': 1, 'LOAD_HT': 2, 'LOAD_BYPASS': 3, 'LOAD_LU': 4, 'LOAD_RT_NT': 5, 'LOAD_NT_HT': 6, 'LOAD_RT_WB': 7,
|
||||
'STORE_RT': 0, 'STORE_NT': 1, 'STORE_HT': 2, 'STORE_BYPASS': 3, 'STORE_LU': 4, 'STORE_RT_NT': 5, 'STORE_NT_HT': 6,
|
||||
'ATOMIC_RT': 0, 'ATOMIC_NT': 1, 'ATOMIC_RETURN': 1, 'ATOMIC_RT_RETURN': 1, 'ATOMIC_NT_RETURN': 3, 'ATOMIC_CASCADE_RT': 6, 'ATOMIC_CASCADE_NT': 6}.get(m.group(1), 0) if m else 0
|
||||
m, buf_text = _extract(buf_text, r'\s+scope:SCOPE_(\w+)')
|
||||
scope_val = {'CU': 0, 'SE': 1, 'DEV': 2, 'SYS': 3}.get(m.group(1), 0) if m else 0
|
||||
buf_ops = _parse_ops(buf_text)
|
||||
buf_args = [_op2dsl(o) for o in buf_ops]
|
||||
vbuf_mods = "".join([f", ioffset={off_val}" if off_val else "", ", offen=1" if offen else "", ", idxen=1" if idxen else "",
|
||||
f", th={th_val}" if th_val else "", f", scope={scope_val}" if scope_val else "",
|
||||
", tfe=1" if tfe else ""])
|
||||
if is_tbuf and fmt_num is not None: vbuf_mods = f", format={fmt_num}" + vbuf_mods
|
||||
elif is_tbuf: vbuf_mods = ", format=1" + vbuf_mods
|
||||
else: vbuf_mods = ", format=1" + vbuf_mods
|
||||
vaddr_idx = 1
|
||||
if len(buf_ops) > vaddr_idx and buf_ops[vaddr_idx].strip().lower() == 'off': vaddr_val = "v[0]"
|
||||
else: vaddr_val = buf_args[vaddr_idx] if len(buf_args) > vaddr_idx else "v[0]"
|
||||
rsrc_idx, soff_idx = (2, 3) if len(buf_ops) > 1 else (1, 2)
|
||||
rsrc_raw = buf_ops[rsrc_idx].strip() if len(buf_ops) > rsrc_idx else "s[0:3]"
|
||||
if m := re.match(r's\[(\d+):\d+\]', rsrc_raw.lower()): rsrc_val = m.group(1)
|
||||
elif m := re.match(r's(\d+)', rsrc_raw.lower()): rsrc_val = m.group(1)
|
||||
elif m := re.match(r'ttmp\[(\d+):\d+\]', rsrc_raw.lower()): rsrc_val = str(108 + int(m.group(1)))
|
||||
elif m := re.match(r'ttmp(\d+)', rsrc_raw.lower()): rsrc_val = str(108 + int(m.group(1)))
|
||||
else: rsrc_val = "0"
|
||||
soff_raw = buf_ops[soff_idx].strip() if len(buf_ops) > soff_idx else "0"
|
||||
soff_lower = soff_raw.lower()
|
||||
if soff_lower == 'm0': soff_val = "RawImm(125)"
|
||||
elif soff_lower in ('null', 'off'): soff_val = "RawImm(124)"
|
||||
elif m := re.match(r's(\d+)', soff_lower): soff_val = f"RawImm({m.group(1)})"
|
||||
else: soff_val = f"RawImm({soff_raw})"
|
||||
return f"{mn}(vdata={buf_args[0]}, vaddr={vaddr_val}, rsrc={rsrc_val}, soffset={soff_val}{vbuf_mods})"
|
||||
buf_mods = "".join([f", offset={off_val}" if off_val else "", ", glc=1" if glc else "", ", dlc=1" if dlc else "",
|
||||
", slc=1" if slc else "", ", tfe=1" if tfe else "", ", offen=1" if offen else "", ", idxen=1" if idxen else ""])
|
||||
if is_tbuf and fmt_num is not None: buf_mods = f", format={fmt_num}" + buf_mods
|
||||
vaddr_idx = 1
|
||||
if len(ops) > vaddr_idx and ops[vaddr_idx].strip().lower() == 'off': vaddr_val = "v[0]"
|
||||
else: vaddr_val = args[vaddr_idx] if len(args) > vaddr_idx else "v[0]"
|
||||
srsrc_idx, soff_idx = (2, 3) if len(ops) > 1 else (1, 2)
|
||||
srsrc_val = args[srsrc_idx] if len(args) > srsrc_idx else "s[0:3]"
|
||||
soff_val = args[soff_idx] if len(args) > soff_idx else "0"
|
||||
return f"{mn}(vdata={args[0]}, vaddr={vaddr_val}, srsrc={srsrc_val}, soffset={soff_val}{buf_mods})"
|
||||
|
||||
# FLAT/GLOBAL/SCRATCH load/store/atomic
|
||||
def _saddr(a): return 'RawImm(124)' if a in ('OFF', 'NULL') else a
|
||||
flat_mods = f"{f', offset={off_val}' if off_val else ''}{', glc=1' if glc else ''}{', slc=1' if slc else ''}{', dlc=1' if dlc else ''}"
|
||||
for pre, flds in [('flat_load','vdst,addr,saddr'), ('global_load','vdst,addr,saddr'), ('scratch_load','vdst,addr,saddr'),
|
||||
('flat_store','addr,data,saddr'), ('global_store','addr,data,saddr'), ('scratch_store','addr,data,saddr')]:
|
||||
if mn.startswith(pre) and len(args) >= 2:
|
||||
f0, f1, f2 = flds.split(',')
|
||||
return f"{mn}({f0}={args[0]}, {f1}={args[1]}{f', {f2}={_saddr(args[2])}' if len(args) >= 3 else ', saddr=RawImm(124)'}{flat_mods})"
|
||||
for pre in ('flat_atomic', 'global_atomic', 'scratch_atomic'):
|
||||
if mn.startswith(pre):
|
||||
if glc and len(args) >= 3: return f"{mn}(vdst={args[0]}, addr={args[1]}, data={args[2]}{f', saddr={_saddr(args[3])}' if len(args) >= 4 else ', saddr=RawImm(124)'}{flat_mods})"
|
||||
if len(args) >= 2: return f"{mn}(addr={args[0]}, data={args[1]}{f', saddr={_saddr(args[2])}' if len(args) >= 3 else ', saddr=RawImm(124)'}{flat_mods})"
|
||||
|
||||
# DS instructions
|
||||
if mn.startswith('ds_'):
|
||||
if ds_off0 is not None or ds_off1 is not None:
|
||||
off0, off1 = str(ds_off0 or 0), str(ds_off1 or 0)
|
||||
elif off_val:
|
||||
off0, off1 = str(int(off_val, 0) & 0xff), str((int(off_val, 0) >> 8) & 0xff)
|
||||
else:
|
||||
off0, off1 = "0", "0"
|
||||
gds_s = ", gds=1" if 'gds' in text.lower().split()[-1:] else ""
|
||||
off_kw = f", offset0={off0}, offset1={off1}{gds_s}"
|
||||
if mn == 'ds_nop' or mn in ('ds_gws_sema_v', 'ds_gws_sema_p', 'ds_gws_sema_release_all'): return f"{mn}({off_kw.lstrip(', ')})"
|
||||
if 'gws_' in mn: return f"{mn}(addr={args[0]}{off_kw})"
|
||||
if 'consume' in mn or 'append' in mn: return f"{mn}(vdst={args[0]}{off_kw})"
|
||||
if 'gs_reg' in mn: return f"{mn}(vdst={args[0]}, data0={args[1]}{off_kw})"
|
||||
if '2addr' in mn:
|
||||
if 'load' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}{off_kw})"
|
||||
if 'store' in mn and 'xchg' not in mn: return f"{mn}(addr={args[0]}, data0={args[1]}, data1={args[2]}{off_kw})"
|
||||
return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}, data1={args[3]}{off_kw})"
|
||||
if 'load' in mn: return f"{mn}(vdst={args[0]}{off_kw})" if 'addtid' in mn else f"{mn}(vdst={args[0]}, addr={args[1]}{off_kw})"
|
||||
if 'store' in mn and not _has(mn, 'cmp', 'xchg'):
|
||||
return f"{mn}(data0={args[0]}{off_kw})" if 'addtid' in mn else f"{mn}(addr={args[0]}, data0={args[1]}{off_kw})"
|
||||
if 'swizzle' in mn or 'ordered_count' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}{off_kw})"
|
||||
if 'permute' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}{off_kw})"
|
||||
if 'bvh' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}, data1={args[3]}{off_kw})"
|
||||
if 'condxchg' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}{off_kw})"
|
||||
if _has(mn, 'cmpstore', 'mskor', 'wrap'):
|
||||
return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}, data1={args[3]}{off_kw})" if '_rtn' in mn else f"{mn}(addr={args[0]}, data0={args[1]}, data1={args[2]}{off_kw})"
|
||||
return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}{off_kw})" if '_rtn' in mn else f"{mn}(addr={args[0]}, data0={args[1]}{off_kw})"
|
||||
|
||||
# v_fmaak/v_fmamk literal extraction
|
||||
lit_s = ""
|
||||
if mn in ('v_fmaak_f32', 'v_fmaak_f16') and len(args) == 4: lit_s, args = f", literal={args[3].strip()}", args[:3]
|
||||
elif mn in ('v_fmamk_f32', 'v_fmamk_f16') and len(args) == 4: lit_s, args = f", literal={args[2].strip()}", [args[0], args[1], args[3]]
|
||||
elif mn in ('s_fmaak_f32',) and len(args) == 4: lit_s, args = f", literal={args[3].strip()}", args[:3]
|
||||
elif mn in ('s_fmamk_f32',) and len(args) == 4: lit_s, args = f", literal={args[2].strip()}", [args[0], args[1], args[3]]
|
||||
elif mn in ('v_cndmask_b32', 'v_cndmask_b32_e32') and len(args) == 4 and ops[3].strip().lower() in ('vcc_lo', 'vcc'):
|
||||
mn, args = 'v_cndmask_b32_e32', args[:3]
|
||||
|
||||
_SGPR_NAMES = {'vcc_lo': 106, 'vcc_hi': 107, 'vcc': 106, 'null': 124, 'm0': 125, 'exec_lo': 126, 'exec_hi': 127}
|
||||
vcc_ops = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'}
|
||||
if mn.replace('_e32', '') in vcc_ops and len(args) >= 5:
|
||||
carry_in = ops[4].strip().lower() if len(ops) > 4 else 'vcc_lo'
|
||||
carry_out = ops[1].strip().lower() if len(ops) > 1 else 'vcc_lo'
|
||||
if carry_in in ('vcc_lo', 'vcc') and carry_out in ('vcc_lo', 'vcc'):
|
||||
mn, args = mn.replace('_e32', '') + '_e32', [args[0], args[2], args[3]]
|
||||
else:
|
||||
mn_base = mn.replace('_e32', '').replace('_e64', '')
|
||||
sdst = _SGPR_NAMES.get(carry_out, 124) if carry_out in _SGPR_NAMES else (int(carry_out[1:]) if carry_out.startswith('s') and carry_out[1:].isdigit() else 124)
|
||||
src2 = _SGPR_NAMES.get(carry_in, 0) if carry_in in _SGPR_NAMES else (int(carry_in[1:]) if carry_in.startswith('s') and carry_in[1:].isdigit() else 0)
|
||||
return f"{mn_base}(vdst={args[0]}, sdst=RawImm({sdst}), src0={args[2]}, src1={args[3]}, src2=RawImm({src2}))"
|
||||
if mn.replace('_e64', '') in vcc_ops and mn.endswith('_e64'): mn = mn.replace('_e64', '')
|
||||
if mn.startswith('v_cmp') and not mn.endswith('_e64') and len(args) >= 3 and ops[0].strip().lower() in ('vcc_lo', 'vcc_hi', 'vcc'): args = args[1:]
|
||||
if 'cmpx' in mn and mn.endswith('_e64') and len(args) == 2: args = ['RawImm(126)'] + args
|
||||
if ((mn.startswith('v_cmp') and 'cmpx' not in mn and mn.endswith('_e64')) or mn.startswith('v_s_') or mn in ('v_readlane_b32', 'v_readfirstlane_b32')) and len(args) >= 1:
|
||||
dst = ops[0].strip().lower()
|
||||
if dst.startswith('s') and dst[1:].isdigit(): args[0] = f'RawImm({int(dst[1:])})'
|
||||
elif dst.startswith('s[') and ':' in dst: args[0] = f'RawImm({int(dst[2:].split(":")[0])})'
|
||||
elif dst.startswith('ttmp') and dst[4:].isdigit(): args[0] = f'RawImm({108 + int(dst[4:])})'
|
||||
elif dst.startswith('ttmp[') and ':' in dst: args[0] = f'RawImm({108 + int(dst[5:].split(":")[0])})'
|
||||
elif dst in _SGPR_NAMES: args[0] = f'RawImm({_SGPR_NAMES[dst]})'
|
||||
|
||||
fn = mn.replace('.', '_')
|
||||
if opsel is not None: args = [re.sub(r'\.[hl]$', '', a) for a in args]
|
||||
|
||||
# v_fma_mix*: extract inline neg/abs modifiers
|
||||
if 'fma_mix' in mn and neg_lo is None and neg_hi is None:
|
||||
inline_neg, inline_abs, clean_args = 0, 0, [args[0]]
|
||||
for i, op in enumerate(ops[1:4]):
|
||||
op = op.strip()
|
||||
neg = op.startswith('-') and not (op[1:2].isdigit() or (len(op) > 2 and op[1] == '0' and op[2] in 'xX'))
|
||||
if neg: op = op[1:]
|
||||
abs_ = op.startswith('|') and op.endswith('|')
|
||||
if abs_: op = op[1:-1]
|
||||
if neg: inline_neg |= (1 << i)
|
||||
if abs_: inline_abs |= (1 << i)
|
||||
clean_args.append(_op2dsl(op))
|
||||
args = clean_args + args[4:]
|
||||
if inline_neg: neg_lo = inline_neg
|
||||
if inline_abs: neg_hi = inline_abs
|
||||
|
||||
all_kw = list(kw)
|
||||
if lit_s: all_kw.append(lit_s.lstrip(', '))
|
||||
if opsel is not None: all_kw.append(f'opsel={opsel}')
|
||||
if neg_lo is not None: all_kw.append(f'neg={neg_lo}')
|
||||
if neg_hi is not None: all_kw.append(f'neg_hi={neg_hi}')
|
||||
if 'bvh' in mn and 'intersect_ray' in mn: all_kw.extend(['dmask=15', 'unrm=1', 'r128=1'])
|
||||
vop3p_ops = {'v_pk_', 'v_dot2', 'v_dot4', 'v_dot8', 'v_wmma', 'v_swmmac'}
|
||||
is_vop3p = any(mn.startswith(p) for p in vop3p_ops)
|
||||
is_fma_mix = 'fma_mix' in mn
|
||||
if opsel_hi_val is not None:
|
||||
opsel_hi_enc = opsel_hi_val[0] | (opsel_hi_val[1] << 1) if len(opsel_hi_val) >= 2 else opsel_hi_val[0]
|
||||
opsel_hi2_enc = opsel_hi_val[2] if len(opsel_hi_val) >= 3 else (0 if is_fma_mix else 1)
|
||||
all_kw.extend([f'opsel_hi={opsel_hi_enc}', f'opsel_hi2={opsel_hi2_enc}'])
|
||||
elif is_vop3p and not is_fma_mix:
|
||||
all_kw.extend(['opsel_hi=3', 'opsel_hi2=1'])
|
||||
if clamp_found:
|
||||
if arch == 'rdna4': all_kw.append('cm=1')
|
||||
else: all_kw.append('clmp=1')
|
||||
|
||||
a_str, kw_str = ', '.join(args), ', '.join(all_kw)
|
||||
return f"{fn}({a_str}, {kw_str})" if kw_str and a_str else f"{fn}({kw_str})" if kw_str else f"{fn}({a_str})"
|
||||
|
||||
def _hwreg(id_, offset=0, size=32): return id_ | (offset << 6) | ((size - 1) << 11)
|
||||
def _sendmsg(id_, op=0, stream=0): return id_ | (op << 4) | (stream << 8)
|
||||
|
||||
_HWREG_NAMES = {'HW_REG_MODE': 1, 'HW_REG_STATUS': 2, 'HW_REG_TRAPSTS': 3, 'HW_REG_HW_ID': 4, 'HW_REG_GPR_ALLOC': 5,
|
||||
'HW_REG_LDS_ALLOC': 6, 'HW_REG_IB_STS': 7, 'HW_REG_PC_LO': 8, 'HW_REG_PC_HI': 9, 'HW_REG_INST_DW0': 10, 'HW_REG_INST_DW1': 11,
|
||||
'HW_REG_IB_DBG0': 12, 'HW_REG_IB_DBG1': 13, 'HW_REG_FLUSH_IB': 14, 'HW_REG_SH_MEM_BASES': 15, 'HW_REG_SQ_SHADER_TBA_LO': 16,
|
||||
'HW_REG_SQ_SHADER_TBA_HI': 17, 'HW_REG_SQ_SHADER_TMA_LO': 18, 'HW_REG_SQ_SHADER_TMA_HI': 19, 'HW_REG_FLAT_SCR_LO': 20,
|
||||
'HW_REG_FLAT_SCR_HI': 21, 'HW_REG_XNACK_MASK': 22, 'HW_REG_HW_ID1': 23, 'HW_REG_HW_ID2': 24, 'HW_REG_POPS_PACKER': 25,
|
||||
'HW_REG_PERF_SNAPSHOT_DATA': 26, 'HW_REG_PERF_SNAPSHOT_PC_LO': 27, 'HW_REG_PERF_SNAPSHOT_PC_HI': 28, 'HW_REG_SHADER_CYCLES': 29,
|
||||
'HW_REG_SHADER_CYCLES_HI': 30, 'HW_REG_WAVE_MODE': 31, 'HW_REG_WAVE_SCRATCH_BASE': 32}
|
||||
_HWREG_NAMES_RDNA4 = {v: k for k, v in HWREG_RDNA4.items()}
|
||||
_SENDMSG_NAMES = {'MSG_INTERRUPT': 1, 'MSG_GS': 2, 'MSG_GS_DONE': 3, 'MSG_SAVEWAVE': 4, 'MSG_STALL_WAVE_GEN': 5,
|
||||
'MSG_HALT_WAVES': 6, 'MSG_ORDERED_PS_DONE': 7, 'MSG_EARLY_PRIM_DEALLOC': 8, 'MSG_GS_ALLOC_REQ': 9, 'MSG_GET_DOORBELL': 10,
|
||||
'MSG_GET_DDID': 11, 'MSG_HS_TESSFACTOR': 2, 'MSG_DEALLOC_VGPRS': 10, 'MSG_RTN_GET_DOORBELL': 128, 'MSG_RTN_GET_DDID': 129,
|
||||
'MSG_RTN_GET_TMA': 130, 'MSG_RTN_GET_REALTIME': 131, 'MSG_RTN_SAVE_WAVE': 132, 'MSG_RTN_GET_TBA': 133,
|
||||
'MSG_RTN_GET_TBA_TO_PC': 134, 'MSG_RTN_GET_SE_AID_ID': 135}
|
||||
|
||||
def asm(text: str, arch: str = "rdna3") -> Inst:
|
||||
dsl = get_dsl(text, arch)
|
||||
if arch == "rdna4":
|
||||
ns = {n: getattr(rdna4_ins, n) for n in dir(rdna4_ins) if not n.startswith('_')}
|
||||
hwreg_names = _HWREG_NAMES_RDNA4
|
||||
else:
|
||||
ns = {n: getattr(ins, n) for n in dir(ins) if not n.startswith('_')}
|
||||
hwreg_names = _HWREG_NAMES
|
||||
def hwreg(id_, offset=0, size=32): return _hwreg(hwreg_names.get(id_, id_) if isinstance(id_, str) else id_, offset, size)
|
||||
def sendmsg(id_, op=0, stream=0): return _sendmsg(_SENDMSG_NAMES.get(id_, id_) if isinstance(id_, str) else id_, op, stream)
|
||||
ns.update({'s': s, 'v': v, 'ttmp': ttmp, 'abs': abs, 'RawImm': RawImm, 'SrcMod': SrcMod, 'VGPR': VGPR, 'SGPR': SGPR, 'TTMP': TTMP,
|
||||
'VCC_LO': VCC_LO, 'VCC_HI': VCC_HI, 'VCC': VCC, 'EXEC_LO': EXEC_LO, 'EXEC_HI': EXEC_HI, 'EXEC': EXEC, 'SCC': SCC, 'M0': M0, 'NULL': NULL, 'OFF': OFF,
|
||||
'hwreg': hwreg, 'sendmsg': sendmsg, **{k: k for k in hwreg_names}, **{k: k for k in _SENDMSG_NAMES}})
|
||||
try: return eval(dsl, ns)
|
||||
except NameError:
|
||||
if m := re.match(r'^(v_\w+)(\(.*\))$', dsl): return eval(f"{m.group(1)}_e32{m.group(2)}", ns)
|
||||
raise
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
+86
-1676
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,62 @@
|
||||
# Instruction format detection and decoding
|
||||
from __future__ import annotations
|
||||
from extra.assembly.amd.dsl import Inst
|
||||
from extra.assembly.amd.autogen.rdna3.ins import VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, VOPD, VINTERP, SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT, MUBUF, MTBUF, MIMG, EXP
|
||||
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP2 as R4_VOP2, VOP3 as R4_VOP3, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
|
||||
VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOPC as R4_SOPC, SOPK as R4_SOPK, SOPP as R4_SOPP,
|
||||
SMEM as R4_SMEM, DS as R4_DS, VBUFFER as R4_VBUFFER, VEXPORT as R4_VEXPORT)
|
||||
from extra.assembly.amd.autogen.cdna.ins import (VOP1 as C_VOP1, VOP2 as C_VOP2, VOPC as C_VOPC, VOP3A, VOP3B, VOP3P as C_VOP3P,
|
||||
SOP1 as C_SOP1, SOP2 as C_SOP2, SOPC as C_SOPC, SOPK as C_SOPK, SOPP as C_SOPP, SMEM as C_SMEM, DS as C_DS,
|
||||
FLAT as C_FLAT, MUBUF as C_MUBUF, MTBUF as C_MTBUF, SDWA, DPP)
|
||||
|
||||
def _matches_encoding(word: int, cls: type[Inst]) -> bool:
|
||||
"""Check if word matches the encoding pattern of an instruction class."""
|
||||
if cls._encoding is None: return False
|
||||
bf, val = cls._encoding
|
||||
return ((word >> bf.lo) & bf.mask()) == val
|
||||
|
||||
# Order matters: more specific encodings first, VOP2 last (it's a catch-all for bit31=0)
|
||||
_RDNA_FORMATS_64 = [VOPD, VOP3P, VINTERP, VOP3, DS, FLAT, MUBUF, MTBUF, MIMG, SMEM, EXP]
|
||||
_RDNA_FORMATS_32 = [SOP1, SOPC, SOPP, SOPK, VOPC, VOP1, SOP2, VOP2] # SOP2/VOP2 are catch-alls
|
||||
_CDNA_FORMATS_64 = [C_VOP3P, VOP3A, C_DS, C_FLAT, C_MUBUF, C_MTBUF, C_SMEM]
|
||||
_CDNA_FORMATS_32 = [SDWA, DPP, C_SOP1, C_SOPC, C_SOPP, C_SOPK, C_VOPC, C_VOP1, C_SOP2, C_VOP2]
|
||||
_CDNA_VOP3B_OPS = {281, 282, 283, 284, 285, 286, 480, 481, 488, 489} # VOP3B opcodes
|
||||
_RDNA4_FORMATS_64 = [R4_VOPD, R4_VOP3P, R4_VINTERP, R4_VOP3, R4_DS, R4_VBUFFER, R4_SMEM, R4_VEXPORT]
|
||||
_RDNA4_FORMATS_32 = [R4_SOP1, R4_SOPC, R4_SOPP, R4_SOPK, R4_VOPC, R4_VOP1, R4_SOP2, R4_VOP2]
|
||||
_RDNA4_VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
|
||||
|
||||
def detect_format(data: bytes, arch: str = "rdna3") -> type[Inst]:
|
||||
"""Detect instruction format from machine code bytes."""
|
||||
assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
|
||||
word = int.from_bytes(data[:4], 'little')
|
||||
if arch == "cdna":
|
||||
if (word >> 30) == 0b11:
|
||||
for cls in _CDNA_FORMATS_64:
|
||||
if _matches_encoding(word, cls):
|
||||
return VOP3B if cls is VOP3A and ((word >> 16) & 0x3ff) in _CDNA_VOP3B_OPS else cls
|
||||
raise ValueError(f"unknown CDNA 64-bit format word={word:#010x}")
|
||||
for cls in _CDNA_FORMATS_32:
|
||||
if _matches_encoding(word, cls): return cls
|
||||
raise ValueError(f"unknown CDNA 32-bit format word={word:#010x}")
|
||||
if arch == "rdna4":
|
||||
if (word >> 30) == 0b11:
|
||||
for cls in _RDNA4_FORMATS_64:
|
||||
if _matches_encoding(word, cls):
|
||||
return R4_VOP3SD if cls is R4_VOP3 and ((word >> 16) & 0x3ff) in _RDNA4_VOP3SD_OPS else cls
|
||||
raise ValueError(f"unknown RDNA4 64-bit format word={word:#010x}")
|
||||
for cls in _RDNA4_FORMATS_32:
|
||||
if _matches_encoding(word, cls): return cls
|
||||
raise ValueError(f"unknown RDNA4 32-bit format word={word:#010x}")
|
||||
# RDNA3 (default)
|
||||
if (word >> 30) == 0b11:
|
||||
for cls in _RDNA_FORMATS_64:
|
||||
if _matches_encoding(word, cls):
|
||||
return VOP3SD if cls is VOP3 and ((word >> 16) & 0x3ff) in Inst._VOP3SD_OPS else cls
|
||||
raise ValueError(f"unknown 64-bit format word={word:#010x}")
|
||||
for cls in _RDNA_FORMATS_32:
|
||||
if _matches_encoding(word, cls): return cls
|
||||
raise ValueError(f"unknown 32-bit format word={word:#010x}")
|
||||
|
||||
def decode_inst(data: bytes, arch: str = "rdna3") -> Inst:
|
||||
"""Decode machine code bytes into an instruction."""
|
||||
return detect_format(data, arch).from_bytes(data)
|
||||
@@ -0,0 +1,745 @@
|
||||
# RDNA3/RDNA4/CDNA disassembler
|
||||
from __future__ import annotations
|
||||
import re
|
||||
from extra.assembly.amd.dsl import Inst, decode_src, SPECIAL_GPRS, SPECIAL_GPRS_CDNA, SPECIAL_PAIRS, SPECIAL_PAIRS_CDNA
|
||||
from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, VOPD, VINTERP, SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT, MUBUF, MTBUF, MIMG, EXP,
|
||||
VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOPDOp, SOP1Op, SOPKOp, SOPPOp, SMEMOp, DSOp, MUBUFOp)
|
||||
from extra.assembly.amd.autogen.rdna3.enum import BufFmt
|
||||
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP2 as R4_VOP2, VOP3 as R4_VOP3, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
|
||||
VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOPC as R4_SOPC, SOPK as R4_SOPK, SOPP as R4_SOPP,
|
||||
SMEM as R4_SMEM, DS as R4_DS, VBUFFER as R4_VBUFFER, VEXPORT as R4_VEXPORT, VOPDOp as R4_VOPDOp)
|
||||
from extra.assembly.amd.autogen.cdna.ins import FLAT as C_FLAT, MUBUF as C_MUBUF, MTBUF as C_MTBUF
|
||||
|
||||
def _is_cdna(inst: Inst) -> bool: return 'cdna' in inst.__class__.__module__
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# CONSTANTS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
HWREG = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
|
||||
6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
|
||||
19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI', 22: 'HW_REG_XNACK_MASK',
|
||||
23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
|
||||
HWREG_RDNA4 = {1: 'HW_REG_WAVE_MODE', 2: 'HW_REG_WAVE_STATUS', 4: 'HW_REG_WAVE_STATE_PRIV', 5: 'HW_REG_WAVE_GPR_ALLOC',
|
||||
6: 'HW_REG_WAVE_LDS_ALLOC', 7: 'HW_REG_IB_STS', 10: 'HW_REG_PERF_SNAPSHOT_DATA', 11: 'HW_REG_PERF_SNAPSHOT_PC_LO',
|
||||
12: 'HW_REG_PERF_SNAPSHOT_PC_HI', 15: 'HW_REG_PERF_SNAPSHOT_DATA1', 16: 'HW_REG_PERF_SNAPSHOT_DATA2',
|
||||
17: 'HW_REG_WAVE_EXCP_FLAG_PRIV', 18: 'HW_REG_WAVE_EXCP_FLAG_USER', 19: 'HW_REG_WAVE_TRAP_CTRL',
|
||||
20: 'HW_REG_WAVE_SCRATCH_BASE_LO', 21: 'HW_REG_WAVE_SCRATCH_BASE_HI', 23: 'HW_REG_WAVE_HW_ID1',
|
||||
24: 'HW_REG_WAVE_HW_ID2', 26: 'HW_REG_WAVE_SCHED_MODE', 29: 'HW_REG_SHADER_CYCLES_LO',
|
||||
30: 'HW_REG_SHADER_CYCLES_HI', 31: 'HW_REG_WAVE_DVGPR_ALLOC_LO', 32: 'HW_REG_WAVE_DVGPR_ALLOC_HI'}
|
||||
MSG = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
|
||||
131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA',
|
||||
134: 'MSG_RTN_GET_TBA_TO_PC', 135: 'MSG_RTN_GET_SE_AID_ID'}
|
||||
# CDNA opcode name aliases for disasm (new name -> old name expected by tests)
|
||||
_CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32', 'v_fmamk_f32': 'v_madmk_f32', 'v_fmaak_f32': 'v_madak_f32'}
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# HELPERS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{b}" if n == 1 else f"{p}[{b}:{b+n-1}]"
|
||||
def _sreg(b: int, n: int = 1) -> str: return _reg("s", b, n)
|
||||
def _vreg(b: int, n: int = 1) -> str: return _reg("v", b, n)
|
||||
def _areg(b: int, n: int = 1) -> str: return _reg("a", b, n) # accumulator registers for GFX90a
|
||||
def _ttmp(b: int, n: int = 1) -> str: return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
|
||||
def _sreg_or_ttmp(b: int, n: int = 1) -> str: return _ttmp(b, n) or _sreg(b, n)
|
||||
|
||||
def _fmt_sdst(v: int, n: int = 1, cdna: bool = False) -> str:
|
||||
if t := _ttmp(v, n): return t
|
||||
pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
|
||||
gprs = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
|
||||
if n > 1: return pairs.get(v) or gprs.get(v) or _sreg(v, n) # also check gprs for null/m0
|
||||
return gprs.get(v, f"s{v}")
|
||||
|
||||
def _fmt_src(v: int, n: int = 1, cdna: bool = False) -> str:
|
||||
if n == 1: return decode_src(v, cdna)
|
||||
if v >= 256: return _vreg(v - 256, n)
|
||||
if v <= 101: return _sreg(v, n) # s0-s101 can be pairs, but 102+ are special on CDNA
|
||||
pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
|
||||
if n == 2 and v in pairs: return pairs[v]
|
||||
if v <= 105: return _sreg(v, n) # s102-s105 regular pairs for RDNA
|
||||
if t := _ttmp(v, n): return t
|
||||
return decode_src(v, cdna)
|
||||
|
||||
def _fmt_v16(v: int, base: int = 256, hi_thresh: int = 384) -> str:
|
||||
return f"v{(v - base) & 0x7f}.{'h' if v >= hi_thresh else 'l'}"
|
||||
|
||||
def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
|
||||
def _omod(v: int) -> str: return {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(v, "")
|
||||
def _src16(inst, v: int) -> str: return _fmt_v16(v) if v >= 256 else inst.lit(v) # format 16-bit src: vgpr.h/l or literal
|
||||
def _mods(*pairs) -> str: return " ".join(m for c, m in pairs if c)
|
||||
def _fmt_bits(label: str, val: int, count: int) -> str: return f"{label}:[{','.join(str((val >> i) & 1) for i in range(count))}]"
|
||||
|
||||
def _vop3_src(inst, v: int, neg: int, abs_: int, hi: int, n: int, f16: bool) -> str:
|
||||
"""Format VOP3 source operand with modifiers."""
|
||||
if v == 255: s = inst.lit(v) # literal constant takes priority
|
||||
elif n > 1: s = _fmt_src(v, n)
|
||||
elif f16 and v >= 256: s = f"v{v - 256}.h" if hi else f"v{v - 256}.l"
|
||||
elif v == 253: s = "src_scc" # VOP3 sources use src_scc not scc
|
||||
else: s = inst.lit(v)
|
||||
if abs_: s = f"|{s}|"
|
||||
return f"-{s}" if neg else s
|
||||
|
||||
def _opsel_str(opsel: int, n: int, need: bool, is16_d: bool) -> str:
|
||||
"""Format op_sel modifier string."""
|
||||
if not need: return ""
|
||||
dst_hi = (opsel >> 3) & 1
|
||||
if n == 1: return f" op_sel:[{opsel & 1},{dst_hi}]"
|
||||
if n == 2: return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{dst_hi}]"
|
||||
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{dst_hi}]"
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# DISASSEMBLER
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _disasm_vop1(inst: VOP1) -> str:
|
||||
name, cdna = inst.op_name.lower() or f'vop1_op_{inst.op}', _is_cdna(inst)
|
||||
suf = "" if cdna else "_e32"
|
||||
if name in ('v_nop', 'v_pipeflush', 'v_clrexcp'): return name # no operands
|
||||
if 'readfirstlane' in name:
|
||||
src = f"v{inst.src0 - 256}" if inst.src0 >= 256 else decode_src(inst.src0, cdna)
|
||||
return f"{name} {_fmt_sdst(inst.vdst, 1, cdna)}, {src}"
|
||||
# 16-bit dst: uses .h/.l suffix for RDNA (CDNA uses plain vN)
|
||||
parts = name.split('_')
|
||||
is_16d = not cdna and (any(p in ('f16','i16','u16','b16') for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in ('f16','i16','u16','b16') and 'cvt' not in name))
|
||||
# v_cvt_pk_f32_fp8 and v_cvt_pk_f32_bf8 output to 2 VGPRs, and take 16-bit src
|
||||
is_pk_fp8 = 'cvt_pk_f32_fp8' in name or 'cvt_pk_f32_bf8' in name
|
||||
dregs = 2 if is_pk_fp8 else inst.dst_regs()
|
||||
dst = _vreg(inst.vdst, dregs) if dregs > 1 else _fmt_v16(inst.vdst, 0, 128) if is_16d else f"v{inst.vdst}"
|
||||
src = inst.lit(inst.src0) if inst.src0 == 255 else _fmt_src(inst.src0, inst.src_regs(0), cdna) if inst.src_regs(0) > 1 else _src16(inst, inst.src0) if not cdna and (inst.is_src_16(0) or is_pk_fp8) and 'sat_pk' not in name else inst.lit(inst.src0)
|
||||
return f"{name}{suf} {dst}, {src}"
|
||||
|
||||
_VOP2_CARRY_OUT = {'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'} # carry out only
|
||||
_VOP2_CARRY_INOUT = {'v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'} # carry in and out (CDNA)
|
||||
_VOP2_CARRY_INOUT_RDNA = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'} # carry in and out (RDNA)
|
||||
def _disasm_vop2(inst: VOP2) -> str:
|
||||
name, cdna = inst.op_name.lower(), _is_cdna(inst)
|
||||
if cdna: name = _CDNA_DISASM_ALIASES.get(name, name) # apply CDNA aliases
|
||||
suf = "" if cdna or (not cdna and inst.op == VOP2Op.V_DOT2ACC_F32_F16) else "_e32"
|
||||
lit = getattr(inst, '_literal', None)
|
||||
is16 = not cdna and inst.is_16bit()
|
||||
# fmaak/madak: dst = src0 * vsrc1 + K, fmamk/madmk: dst = src0 * K + vsrc1
|
||||
if 'fmaak' in name or 'madak' in name or (not cdna and inst.op in (VOP2Op.V_FMAAK_F32, VOP2Op.V_FMAAK_F16)):
|
||||
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1, 0, 128)}, 0x{lit:x}"
|
||||
return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}, 0x{lit:x}"
|
||||
if 'fmamk' in name or 'madmk' in name or (not cdna and inst.op in (VOP2Op.V_FMAMK_F32, VOP2Op.V_FMAMK_F16)):
|
||||
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_src16(inst, inst.src0)}, 0x{lit:x}, {_fmt_v16(inst.vsrc1, 0, 128)}"
|
||||
return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, 0x{lit:x}, v{inst.vsrc1}"
|
||||
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1, 0, 128)}"
|
||||
vcc = "vcc" if cdna else "vcc_lo"
|
||||
# CDNA carry ops output vcc after vdst
|
||||
if cdna and name in _VOP2_CARRY_OUT: return f"{name}{suf} v{inst.vdst}, {vcc}, {inst.lit(inst.src0)}, v{inst.vsrc1}"
|
||||
if cdna and name in _VOP2_CARRY_INOUT: return f"{name}{suf} v{inst.vdst}, {vcc}, {inst.lit(inst.src0)}, v{inst.vsrc1}, {vcc}"
|
||||
# RDNA carry-in/out ops: v_add_co_ci_u32, etc. - format: vdst, vcc_lo, src0, vsrc1, vcc_lo
|
||||
if not cdna and name in _VOP2_CARRY_INOUT_RDNA: return f"{name}{suf} v{inst.vdst}, {vcc}, {inst.lit(inst.src0)}, v{inst.vsrc1}, {vcc}"
|
||||
# Handle 64-bit register operands (v_add_f64, v_mul_f64, etc.)
|
||||
dn, sn0, sn1 = inst.dst_regs(), inst.src_regs(0), inst.src_regs(1)
|
||||
if dn > 1 or sn0 > 1 or sn1 > 1:
|
||||
dst = _vreg(inst.vdst, dn) if dn > 1 else f"v{inst.vdst}"
|
||||
src0 = inst.lit(inst.src0) if inst.src0 == 255 else _fmt_src(inst.src0, sn0, cdna)
|
||||
src1 = _vreg(inst.vsrc1, sn1) if sn1 > 1 else f"v{inst.vsrc1}"
|
||||
return f"{name} {dst}, {src0}, {src1}"
|
||||
return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}" + (f", {vcc}" if name == 'v_cndmask_b32' else "")
|
||||
|
||||
def _disasm_vopc(inst: VOPC) -> str:
|
||||
name, cdna = inst.op_name.lower(), _is_cdna(inst)
|
||||
if cdna:
|
||||
s0 = inst.lit(inst.src0) if inst.src0 == 255 else _fmt_src(inst.src0, inst.src_regs(0), cdna)
|
||||
s1 = _vreg(inst.vsrc1, inst.src_regs(1)) if inst.src_regs(1) > 1 else f"v{inst.vsrc1}"
|
||||
return f"{name} vcc, {s0}, {s1}" # CDNA VOPC always outputs vcc
|
||||
# RDNA: v_cmpx_* writes to exec (no vcc), v_cmp_* writes to vcc_lo
|
||||
has_vcc = 'cmpx' not in name
|
||||
s0 = inst.lit(inst.src0) if inst.src0 == 255 else _fmt_src(inst.src0, inst.src_regs(0)) if inst.src_regs(0) > 1 else _src16(inst, inst.src0) if inst.is_16bit() else inst.lit(inst.src0)
|
||||
s1 = _vreg(inst.vsrc1, inst.src_regs(1)) if inst.src_regs(1) > 1 else _fmt_v16(inst.vsrc1, 0, 128) if inst.is_16bit() else f"v{inst.vsrc1}"
|
||||
return f"{name}_e32 vcc_lo, {s0}, {s1}" if has_vcc else f"{name}_e32 {s0}, {s1}"
|
||||
|
||||
NO_ARG_SOPP = {SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
|
||||
SOPPOp.S_WAIT_IDLE, SOPPOp.S_ENDPGM_SAVED, SOPPOp.S_CODE_END, SOPPOp.S_ENDPGM_ORDERED_PS_DONE, SOPPOp.S_TTRACEDATA}
|
||||
_CDNA_NO_ARG_SOPP = {'s_endpgm', 's_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_nop', 's_sethalt', 's_sleep',
|
||||
's_setprio', 's_trap', 's_incperflevel', 's_decperflevel', 's_sendmsg', 's_sendmsghalt'}
|
||||
|
||||
def _disasm_sopp(inst: SOPP) -> str:
|
||||
name, cdna = inst.op_name.lower(), _is_cdna(inst)
|
||||
if cdna:
|
||||
if name == 's_endpgm': return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
|
||||
if name in ('s_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata'): return name
|
||||
if name == 's_waitcnt':
|
||||
vm, lgkm, exp = inst.simm16 & 0xf, (inst.simm16 >> 8) & 0x3f, (inst.simm16 >> 4) & 0x7
|
||||
p = [f"vmcnt({vm})" if vm != 0xf else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
|
||||
return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
|
||||
if name.startswith(('s_cbranch', 's_branch')): return f"{name} {inst.simm16}"
|
||||
return f"{name} 0x{inst.simm16:x}" if inst.simm16 else name
|
||||
# RDNA
|
||||
if inst.op in NO_ARG_SOPP: return name
|
||||
if inst.op == SOPPOp.S_ENDPGM: return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
|
||||
if inst.op == SOPPOp.S_WAITCNT:
|
||||
vm, exp, lgkm = (inst.simm16 >> 10) & 0x3f, inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x3f
|
||||
p = [f"vmcnt({vm})" if vm != 0x3f else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
|
||||
return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
|
||||
if inst.op == SOPPOp.S_DELAY_ALU:
|
||||
deps, skips = ['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'], ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
|
||||
id0, skip, id1 = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 7) & 0xf
|
||||
dep = lambda v: deps[v-1] if 0 < v <= len(deps) else str(v)
|
||||
p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
|
||||
return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
|
||||
return f"{name} {inst.simm16}" if name.startswith(('s_cbranch', 's_branch')) else f"{name} 0x{inst.simm16:x}"
|
||||
|
||||
def _disasm_smem(inst: SMEM) -> str:
|
||||
name, cdna = inst.op_name.lower(), _is_cdna(inst)
|
||||
if inst.op in (SMEMOp.S_GL1_INV, SMEMOp.S_DCACHE_INV): return name
|
||||
soe, imm = getattr(inst, 'soe', 0), getattr(inst, 'imm', 1)
|
||||
is_rdna4 = 'rdna4' in inst.__class__.__module__
|
||||
offset = inst.ioffset if is_rdna4 else getattr(inst, 'offset', 0)
|
||||
if cdna:
|
||||
if soe and imm: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
|
||||
elif imm: off_s = f"0x{offset:x}"
|
||||
elif offset < 256: off_s = decode_src(offset, cdna)
|
||||
else: off_s = decode_src(inst.soffset, cdna)
|
||||
elif offset and inst.soffset != 124: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
|
||||
elif offset: off_s = f"0x{offset:x}"
|
||||
else: off_s = decode_src(inst.soffset, cdna)
|
||||
is_buffer = 'buffer' in name or 's_atc_probe_buffer' == name
|
||||
sbase_idx, sbase_count = inst.sbase * 2, 4 if is_buffer else 2
|
||||
sbase_str = _fmt_src(sbase_idx, sbase_count, cdna) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
|
||||
if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {inst.sdata}, {sbase_str}, {off_s}"
|
||||
if 'prefetch' in name:
|
||||
off = getattr(inst, 'ioffset', inst.offset)
|
||||
if off >= 0x800000: off = off - 0x1000000
|
||||
off_s = f"0x{off:x}" if off > 255 else str(off)
|
||||
soff_s = decode_src(inst.soffset, cdna) if inst.soffset != 124 else "null"
|
||||
if 'pc_rel' in name: return f"{name} {off_s}, {soff_s}, {inst.sdata}"
|
||||
return f"{name} {sbase_str}, {off_s}, {soff_s}, {inst.sdata}"
|
||||
th, scope = getattr(inst, 'th', 0), getattr(inst, 'scope', 0)
|
||||
if th or scope:
|
||||
th_names = ['TH_LOAD_RT', 'TH_LOAD_NT', 'TH_LOAD_HT', 'TH_LOAD_LU']
|
||||
scope_names = ['SCOPE_CU', 'SCOPE_SE', 'SCOPE_DEV', 'SCOPE_SYS']
|
||||
mods = (f" th:{th_names[th]}" if th else "") + (f" scope:{scope_names[scope]}" if scope else "")
|
||||
return f"{name} {_fmt_sdst(inst.sdata, inst.dst_regs(), cdna)}, {sbase_str}, {off_s}{mods}"
|
||||
return f"{name} {_fmt_sdst(inst.sdata, inst.dst_regs(), cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
|
||||
|
||||
def _disasm_flat(inst: FLAT) -> str:
|
||||
name, cdna = inst.op_name.lower(), _is_cdna(inst)
|
||||
acc = getattr(inst, 'acc', 0)
|
||||
reg_fn = _areg if acc else _vreg
|
||||
seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
|
||||
instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
|
||||
off_val = inst.offset if seg == 'flat' else (inst.offset if inst.offset < 4096 else inst.offset - 8192)
|
||||
w = inst.dst_regs() * (2 if '_x2' in name else 1) * (2 if 'cmpswap' in name else 1)
|
||||
off_s = f" offset:{off_val}" if off_val else ""
|
||||
if cdna: mods = f"{off_s}{' glc' if inst.sc0 else ''}{' slc' if inst.nt else ''}"
|
||||
else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
|
||||
if seg == 'flat' or inst.saddr == 0x7F: saddr_s = ""
|
||||
elif inst.saddr == 124: saddr_s = ", off"
|
||||
elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr, cdna)}"
|
||||
elif inst.saddr in (SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS): saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[inst.saddr]}"
|
||||
elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
|
||||
else: saddr_s = f", {_sreg(inst.saddr, 2) if inst.saddr < 106 else decode_src(inst.saddr, cdna)}"
|
||||
if 'addtid' in name: return f"{instr} {'a' if acc else 'v'}{inst.data if 'store' in name else inst.vdst}{saddr_s}{mods}"
|
||||
if cdna: addr_w = 1 if seg == 'scratch' else 2
|
||||
else: addr_w = 1 if seg == 'scratch' or (inst.saddr not in (0x7F, 124)) else 2
|
||||
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
|
||||
data_s, vdst_s = reg_fn(inst.data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
|
||||
glc_or_sc0 = inst.sc0 if cdna else inst.glc
|
||||
if 'atomic' in name:
|
||||
return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}"
|
||||
if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
|
||||
return f"{instr} {reg_fn(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
|
||||
|
||||
def _disasm_ds(inst: DS) -> str:
|
||||
op, name = inst.op, inst.op_name.lower()
|
||||
acc = getattr(inst, 'acc', 0)
|
||||
reg_fn = _areg if acc else _vreg
|
||||
rp = 'a' if acc else 'v'
|
||||
gds = " gds" if inst.gds else ""
|
||||
off = f" offset:{inst.offset0 | (inst.offset1 << 8)}" if inst.offset0 or inst.offset1 else ""
|
||||
off2 = (" offset0:" + str(inst.offset0) if inst.offset0 else "") + (" offset1:" + str(inst.offset1) if inst.offset1 else "")
|
||||
w = inst.dst_regs()
|
||||
d0, d1, dst, addr = reg_fn(inst.data0, w), reg_fn(inst.data1, w), reg_fn(inst.vdst, w), f"v{inst.addr}"
|
||||
|
||||
if op == DSOp.DS_NOP: return name
|
||||
if op == DSOp.DS_BVH_STACK_RTN_B32: return f"{name} v{inst.vdst}, {addr}, v{inst.data0}, {_vreg(inst.data1, 4)}{off}{gds}"
|
||||
if 'bvh_stack_push' in name:
|
||||
d1_regs = 8 if 'push8' in name else 4
|
||||
vdst_regs = 2 if 'pop2' in name else 1
|
||||
vdst_s = _vreg(inst.vdst, vdst_regs) if vdst_regs > 1 else f"v{inst.vdst}"
|
||||
return f"{name} {vdst_s}, {addr}, v{inst.data0}, {_vreg(inst.data1, d1_regs)}{off}{gds}"
|
||||
if 'gws_sema' in name and op != DSOp.DS_GWS_SEMA_BR: return f"{name}{off}{gds}"
|
||||
if 'gws_' in name: return f"{name} {addr}{off}{gds}"
|
||||
if op in (DSOp.DS_CONSUME, DSOp.DS_APPEND): return f"{name} {rp}{inst.vdst}{off}{gds}"
|
||||
if 'gs_reg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {rp}{inst.data0}{off}{gds}"
|
||||
if '2addr' in name:
|
||||
if 'load' in name: return f"{name} {reg_fn(inst.vdst, w*2)}, {addr}{off2}{gds}"
|
||||
if 'store' in name and 'xchg' not in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
|
||||
return f"{name} {reg_fn(inst.vdst, w*2)}, {addr}, {d0}, {d1}{off2}{gds}"
|
||||
if 'write2' in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
|
||||
if 'read2' in name: return f"{name} {reg_fn(inst.vdst, w*2)}, {addr}{off2}{gds}"
|
||||
if 'load' in name: return f"{name} {rp}{inst.vdst}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
|
||||
if 'store' in name and not _has(name, 'cmp', 'xchg'):
|
||||
return f"{name} {rp}{inst.data0}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
|
||||
if 'swizzle' in name or op == DSOp.DS_ORDERED_COUNT: return f"{name} {rp}{inst.vdst}, {addr}{off}{gds}"
|
||||
if 'permute' in name: return f"{name} {rp}{inst.vdst}, {addr}, {rp}{inst.data0}{off}{gds}"
|
||||
if 'condxchg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {addr}, {reg_fn(inst.data0, 2)}{off}{gds}"
|
||||
if _has(name, 'cmpstore', 'mskor', 'wrap'):
|
||||
return f"{name} {dst}, {addr}, {d0}, {d1}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}, {d1}{off}{gds}"
|
||||
return f"{name} {dst}, {addr}, {d0}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}{off}{gds}"
|
||||
|
||||
def _disasm_vop3(inst: VOP3) -> str:
|
||||
op, name = inst.op, inst.op_name.lower()
|
||||
|
||||
# RDNA4 v_s_* scalar VOP3 instructions
|
||||
if name.startswith('v_s_'):
|
||||
src = inst.lit(inst.src0) if inst.src0 == 255 else ("src_scc" if inst.src0 == 253 else _fmt_src(inst.src0, inst.src_regs(0)))
|
||||
if inst.neg & 1: src = f"-{src}"
|
||||
if inst.abs & 1: src = f"|{src}|"
|
||||
clamp = inst.cm if 'cm' in inst._fields else getattr(inst, 'clmp', 0)
|
||||
return f"{name} s{inst.vdst}, {src}" + (" clamp" if clamp else "") + _omod(inst.omod)
|
||||
|
||||
# VOP3SD (shared encoding)
|
||||
if isinstance(op, VOP3SDOp):
|
||||
sdst = (inst.clmp << 7) | (inst.opsel << 3) | inst.abs
|
||||
def src(v, neg, n):
|
||||
s = inst.lit(v) if v == 255 else ("src_scc" if v == 253 else (_fmt_src(v, n) if n > 1 else inst.lit(v)))
|
||||
return f"neg({s})" if neg and v == 255 else (f"-{s}" if neg else s)
|
||||
s0, s1, s2 = src(inst.src0, inst.neg & 1, inst.src_regs(0)), src(inst.src1, inst.neg & 2, inst.src_regs(1)), src(inst.src2, inst.neg & 4, inst.src_regs(2))
|
||||
dst = _vreg(inst.vdst, inst.dst_regs()) if inst.dst_regs() > 1 else f"v{inst.vdst}"
|
||||
srcs = f"{s0}, {s1}, {s2}" if inst.num_srcs() == 3 else f"{s0}, {s1}"
|
||||
return f"{name} {dst}, {_fmt_sdst(sdst, 1)}, {srcs}" + _omod(inst.omod)
|
||||
|
||||
# Detect 16-bit operand sizes
|
||||
is16_d = is16_s = is16_s2 = False
|
||||
if 'cvt_pk' in name: is16_s = name.endswith('16')
|
||||
elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', name):
|
||||
is16_d, is16_s = _has(m.group(1), 'f16','i16','u16','b16'), _has(m.group(2), 'f16','i16','u16','b16')
|
||||
is16_s2 = is16_s
|
||||
elif re.match(r'v_mad_[iu]32_[iu]16', name): is16_s = True
|
||||
elif 'pack_b32' in name: is16_s = is16_s2 = True
|
||||
elif 'sat_pk' in name: is16_d = True
|
||||
else: is16_d = is16_s = is16_s2 = inst.is_16bit()
|
||||
|
||||
s0 = _vop3_src(inst, inst.src0, inst.neg&1, inst.abs&1, inst.opsel&1, inst.src_regs(0), is16_s)
|
||||
s1 = _vop3_src(inst, inst.src1, inst.neg&2, inst.abs&2, inst.opsel&2, inst.src_regs(1), is16_s)
|
||||
s2 = _vop3_src(inst, inst.src2, inst.neg&4, inst.abs&4, inst.opsel&4, inst.src_regs(2), is16_s2)
|
||||
|
||||
# Destination
|
||||
dn = inst.dst_regs()
|
||||
if op == VOP3Op.V_READLANE_B32: dst = _fmt_sdst(inst.vdst, 1)
|
||||
elif dn > 1: dst = _vreg(inst.vdst, dn)
|
||||
elif is16_d: dst = f"v{inst.vdst}.h" if (inst.opsel & 8) else f"v{inst.vdst}.l"
|
||||
else: dst = f"v{inst.vdst}"
|
||||
|
||||
clamp = inst.cm if 'cm' in inst._fields else getattr(inst, 'clmp', 0)
|
||||
cl, om = " clamp" if clamp else "", _omod(inst.omod)
|
||||
nonvgpr_opsel = (inst.src0 < 256 and (inst.opsel & 1)) or (inst.src1 < 256 and (inst.opsel & 2)) or (inst.src2 < 256 and (inst.opsel & 4))
|
||||
need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)
|
||||
|
||||
if inst.op < 256: # VOPC
|
||||
return f"{name}_e64 {s0}, {s1}{cl}" if name.startswith('v_cmpx') else f"{name}_e64 {_fmt_sdst(inst.vdst, 1)}, {s0}, {s1}{cl}"
|
||||
if inst.op < 384: # VOP2
|
||||
n = inst.num_srcs()
|
||||
os = _opsel_str(inst.opsel, n, need_opsel, is16_d)
|
||||
return f"{name}_e64 {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name}_e64 {dst}, {s0}, {s1}{os}{cl}{om}"
|
||||
if inst.op < 512: # VOP1
|
||||
if re.match(r'v_cvt_f32_(bf|fp)8', name) and inst.opsel:
|
||||
os = f" byte_sel:{((inst.opsel & 1) << 1) | ((inst.opsel >> 1) & 1)}"
|
||||
else:
|
||||
os = _opsel_str(inst.opsel, 1, need_opsel, is16_d)
|
||||
return f"{name}_e64" if op in (VOP3Op.V_NOP, VOP3Op.V_PIPEFLUSH) else f"{name}_e64 {dst}, {s0}{os}{cl}{om}"
|
||||
# Native VOP3
|
||||
n = inst.num_srcs()
|
||||
if 'cvt_sr' in name and inst.opsel:
|
||||
os = f" byte_sel:{inst.opsel >> 2}"
|
||||
else:
|
||||
os = _opsel_str(inst.opsel, n, need_opsel, is16_d)
|
||||
return f"{name} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{os}{cl}{om}"
|
||||
|
||||
def _disasm_vop3sd(inst: VOP3SD) -> str:
|
||||
name = inst.op_name.lower()
|
||||
src2_n = 2 if '_co_' in name and '64' in name else inst.src_regs(2)
|
||||
def src(v, neg, n):
|
||||
s = inst.lit(v) if v == 255 else ("src_scc" if v == 253 else (_fmt_src(v, n) if n > 1 else inst.lit(v)))
|
||||
return f"neg({s})" if neg and v == 255 else (f"-{s}" if neg else s)
|
||||
s0, s1, s2 = src(inst.src0, inst.neg & 1, inst.src_regs(0)), src(inst.src1, inst.neg & 2, inst.src_regs(1)), src(inst.src2, inst.neg & 4, src2_n)
|
||||
dst = _vreg(inst.vdst, inst.dst_regs()) if inst.dst_regs() > 1 else f"v{inst.vdst}"
|
||||
srcs = f"{s0}, {s1}, {s2}" if inst.num_srcs() == 3 else f"{s0}, {s1}"
|
||||
clamp = inst.cm if 'cm' in inst._fields else getattr(inst, 'clmp', 0)
|
||||
return f"{name} {dst}, {_fmt_sdst(inst.sdst, 1)}, {srcs}{' clamp' if clamp else ''}{_omod(inst.omod)}"
|
||||
|
||||
def _disasm_vopd(inst: VOPD) -> str:
|
||||
lit = inst._literal or inst.literal
|
||||
is_rdna4 = 'rdna4' in inst.__class__.__module__
|
||||
op_enum = R4_VOPDOp if is_rdna4 else VOPDOp
|
||||
vdst_y, nx, ny = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
|
||||
def half(n, vd, s0, vs1):
|
||||
if 'mov' in n: return f"{n} v{vd}, {inst.lit(s0)}"
|
||||
if 'fmamk' in n and lit: return f"{n} v{vd}, {inst.lit(s0)}, 0x{lit:x}, v{vs1}"
|
||||
if 'fmaak' in n and lit: return f"{n} v{vd}, {inst.lit(s0)}, v{vs1}, 0x{lit:x}"
|
||||
return f"{n} v{vd}, {inst.lit(s0)}, v{vs1}"
|
||||
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
|
||||
|
||||
def _swmmac_regs(name: str) -> tuple[int, int, int, int]:
|
||||
"""Return (dst, src0, src1, src2) register counts for SWMMAC instructions."""
|
||||
if 'f16_16x16x32' in name or 'bf16_16x16x32' in name: return (4, 4, 8, 1)
|
||||
if 'f32_16x16x32_f16' in name or 'f32_16x16x32_bf16' in name: return (8, 4, 8, 1)
|
||||
if 'i32_16x16x32_iu4' in name: return (8, 1, 2, 1)
|
||||
if 'i32_16x16x64_iu4' in name: return (8, 2, 4, 1)
|
||||
if 'i32_16x16x32_iu8' in name or 'f32_16x16x32_fp8' in name or 'f32_16x16x32_bf8' in name: return (8, 2, 4, 1)
|
||||
return (8, 8, 8, 8)
|
||||
|
||||
def _disasm_vop3p(inst: VOP3P) -> str:
|
||||
name = inst.op_name.lower()
|
||||
is_wmma, is_swmmac, n, is_fma_mix = 'wmma' in name, 'swmmac' in name, inst.num_srcs(), 'fma_mix' in name
|
||||
def get_src(v, sc): return inst.lit(v) if v == 255 else _fmt_src(v, sc)
|
||||
if is_swmmac:
|
||||
dn, s0n, s1n, s2n = _swmmac_regs(name)
|
||||
src0, src1, src2, dst = get_src(inst.src0, s0n), get_src(inst.src1, s1n), get_src(inst.src2, s2n), _vreg(inst.vdst, dn)
|
||||
elif is_wmma:
|
||||
is_rdna4_wmma = 'rdna4' in inst.__class__.__module__
|
||||
sc = 1 if '16x16x16_iu4' in name else 2 if ('iu4' in name or 'iu8' in name or 'fp8' in name or 'bf8' in name) else 4
|
||||
if not is_rdna4_wmma: sc *= 2
|
||||
dc = 8 if not is_rdna4_wmma else (4 if ('f16_16x16' in name or 'bf16_16x16' in name) and 'f32' not in name else 8)
|
||||
src0, src1, src2, dst = get_src(inst.src0, sc), get_src(inst.src1, sc), get_src(inst.src2, dc), _vreg(inst.vdst, dc)
|
||||
else: src0, src1, src2, dst = get_src(inst.src0, 1), get_src(inst.src1, 1), get_src(inst.src2, 1), f"v{inst.vdst}"
|
||||
opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
|
||||
clamp = inst.cm if 'cm' in inst._fields else getattr(inst, 'clmp', 0)
|
||||
if is_fma_mix:
|
||||
def m(s, neg, abs_): return f"-{f'|{s}|' if abs_ else s}" if neg else (f"|{s}|" if abs_ else s)
|
||||
src0, src1, src2 = m(src0, inst.neg & 1, inst.neg_hi & 1), m(src1, inst.neg & 2, inst.neg_hi & 2), m(src2, inst.neg & 4, inst.neg_hi & 4)
|
||||
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if clamp else [])
|
||||
elif is_swmmac:
|
||||
mods = ([f"index_key:{inst.opsel}"] if inst.opsel else []) + ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + \
|
||||
([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
|
||||
else:
|
||||
opsel_hi_default = 7 if n == 3 else 3
|
||||
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else []) + \
|
||||
([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
|
||||
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
|
||||
|
||||
def _disasm_buf(inst: MUBUF | MTBUF) -> str:
|
||||
name, cdna = inst.op_name.lower(), _is_cdna(inst)
|
||||
acc = getattr(inst, 'acc', 0)
|
||||
reg_fn = _areg if acc else _vreg
|
||||
if cdna and name in ('buffer_wbl2', 'buffer_inv'): return name
|
||||
if not cdna and inst.op in (MUBUFOp.BUFFER_GL0_INV, MUBUFOp.BUFFER_GL1_INV): return name
|
||||
w = (2 if _has(name, 'xyz', 'xyzw') else 1) if 'd16' in name else \
|
||||
((2 if _has(name, 'b64', 'u64', 'i64') else 1) * (2 if 'cmpswap' in name else 1)) if 'atomic' in name else \
|
||||
{'b32':1,'b64':2,'b96':3,'b128':4,'b16':1,'x':1,'xy':2,'xyz':3,'xyzw':4}.get(name.split('_')[-1], 1)
|
||||
if hasattr(inst, 'tfe') and inst.tfe: w += 1
|
||||
vaddr = _vreg(inst.vaddr, 2) if inst.offen and inst.idxen else f"v{inst.vaddr}" if inst.offen or inst.idxen else "off"
|
||||
srsrc = _sreg_or_ttmp(inst.srsrc*4, 4)
|
||||
is_mtbuf = isinstance(inst, MTBUF) or isinstance(inst, C_MTBUF)
|
||||
if is_mtbuf:
|
||||
dfmt, nfmt = inst.format & 0xf, (inst.format >> 4) & 0x7
|
||||
if acc: fmt_s = f" dfmt:{dfmt}, nfmt:{nfmt},"
|
||||
elif not cdna: fmt_s = f" format:{inst.format}" if inst.format else ""
|
||||
else:
|
||||
dfmt_names = ['INVALID', '8', '16', '8_8', '32', '16_16', '10_11_11', '11_11_10', '10_10_10_2', '2_10_10_10', '8_8_8_8', '32_32', '16_16_16_16', '32_32_32', '32_32_32_32', 'RESERVED_15']
|
||||
nfmt_names = ['UNORM', 'SNORM', 'USCALED', 'SSCALED', 'UINT', 'SINT', 'RESERVED_6', 'FLOAT']
|
||||
if dfmt == 1 and nfmt == 0: fmt_s = ""
|
||||
elif nfmt == 0: fmt_s = f" format:[BUF_DATA_FORMAT_{dfmt_names[dfmt]}]"
|
||||
elif dfmt == 1: fmt_s = f" format:[BUF_NUM_FORMAT_{nfmt_names[nfmt]}]"
|
||||
else: fmt_s = f" format:[BUF_DATA_FORMAT_{dfmt_names[dfmt]},BUF_NUM_FORMAT_{nfmt_names[nfmt]}]"
|
||||
else: fmt_s = ""
|
||||
if cdna: mods = [m for c, m in [(inst.idxen,"idxen"),(inst.offen,"offen"),(inst.offset,f"offset:{inst.offset}"),(inst.sc0,"glc"),(inst.nt,"slc"),(inst.sc1,"sc1")] if c]
|
||||
else: mods = [m for c, m in [(inst.idxen,"idxen"),(inst.offen,"offen"),(inst.offset,f"offset:{inst.offset}"),(inst.glc,"glc"),(inst.dlc,"dlc"),(inst.slc,"slc"),(inst.tfe,"tfe")] if c]
|
||||
soffset_s = decode_src(inst.soffset, cdna)
|
||||
if cdna and not acc and is_mtbuf: return f"{name} {reg_fn(inst.vdata, w)}, {vaddr}, {srsrc}, {soffset_s}{fmt_s}{' ' + ' '.join(mods) if mods else ''}"
|
||||
return f"{name} {reg_fn(inst.vdata, w)}, {vaddr}, {srsrc},{fmt_s} {soffset_s}{' ' + ' '.join(mods) if mods else ''}"
|
||||
|
||||
def _mimg_vaddr_width(name: str, dim: int, a16: bool) -> int:
|
||||
base = [1, 2, 3, 3, 2, 3, 3, 4][dim]
|
||||
grad = [1, 2, 3, 2, 1, 2, 2, 2][dim]
|
||||
if 'get_resinfo' in name: return 1
|
||||
packed, unpacked = 0, 0
|
||||
if '_mip' in name: packed += 1
|
||||
elif 'sample' in name or 'gather' in name:
|
||||
if '_o' in name: unpacked += 1
|
||||
if re.search(r'_c(_|$)', name): unpacked += 1
|
||||
if '_d' in name: unpacked += (grad + 1) & ~1 if '_g16' in name else grad*2
|
||||
if '_b' in name: unpacked += 1
|
||||
if '_l' in name and '_cl' not in name and '_lz' not in name: packed += 1
|
||||
if '_cl' in name: packed += 1
|
||||
return (base + packed + 1) // 2 + unpacked if a16 else base + packed + unpacked
|
||||
|
||||
def _disasm_mimg(inst: MIMG) -> str:
|
||||
name = inst.op_name.lower()
|
||||
srsrc_base = inst.srsrc * 4
|
||||
srsrc_str = _sreg_or_ttmp(srsrc_base, 8)
|
||||
if 'bvh' in name:
|
||||
vaddr = (9 if '64' in name else 8) if inst.a16 else (12 if '64' in name else 11)
|
||||
return f"{name} {_vreg(inst.vdata, 4)}, {_vreg(inst.vaddr, vaddr)}, {_sreg_or_ttmp(srsrc_base, 4)}{' a16' if inst.a16 else ''}"
|
||||
vdata = 4 if 'gather4' in name or 'msaa_load' in name else (bin(inst.dmask).count('1') or 1)
|
||||
if inst.d16: vdata = (vdata + 1) // 2
|
||||
if inst.tfe: vdata += 1
|
||||
dim_names = ['1d', '2d', '3d', 'cube', '1d_array', '2d_array', '2d_msaa', '2d_msaa_array']
|
||||
dim = dim_names[inst.dim] if inst.dim < len(dim_names) else f"dim_{inst.dim}"
|
||||
vaddr = _mimg_vaddr_width(name, inst.dim, inst.a16)
|
||||
vaddr_str = f"v{inst.vaddr}" if vaddr == 1 else _vreg(inst.vaddr, vaddr)
|
||||
mods = [f"dmask:0x{inst.dmask:x}"] if inst.dmask and (inst.dmask != 15 or 'atomic' in name) else []
|
||||
mods.append(f"dim:SQ_RSRC_IMG_{dim.upper()}")
|
||||
for flag, mod in [(inst.unrm,"unorm"),(inst.glc,"glc"),(inst.slc,"slc"),(inst.dlc,"dlc"),(inst.r128,"r128"),
|
||||
(inst.a16,"a16"),(inst.tfe,"tfe"),(inst.lwe,"lwe"),(inst.d16,"d16")]:
|
||||
if flag: mods.append(mod)
|
||||
ssamp_str = ""
|
||||
if 'sample' in name or 'gather' in name or 'get_lod' in name:
|
||||
ssamp_str = ", " + _sreg_or_ttmp(inst.ssamp * 4, 4)
|
||||
return f"{name} {_vreg(inst.vdata, vdata)}, {vaddr_str}, {srsrc_str}{ssamp_str} {' '.join(mods)}"
|
||||
|
||||
def _disasm_sop1(inst: SOP1) -> str:
|
||||
op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
|
||||
src = inst.lit(inst.ssrc0) if inst.ssrc0 == 255 else _fmt_src(inst.ssrc0, inst.src_regs(0), cdna)
|
||||
if not cdna:
|
||||
if op == SOP1Op.S_GETPC_B64: return f"{name} {_fmt_sdst(inst.sdst, 2)}"
|
||||
if op in (SOP1Op.S_SETPC_B64, SOP1Op.S_RFE_B64): return f"{name} {src}"
|
||||
if op == SOP1Op.S_SWAPPC_B64: return f"{name} {_fmt_sdst(inst.sdst, 2)}, {src}"
|
||||
if op in (SOP1Op.S_SENDMSG_RTN_B32, SOP1Op.S_SENDMSG_RTN_B64): return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs())}, sendmsg({MSG.get(inst.ssrc0, str(inst.ssrc0))})"
|
||||
sop1_src_only = ('S_ALLOC_VGPR', 'S_SLEEP_VAR', 'S_BARRIER_SIGNAL', 'S_BARRIER_SIGNAL_ISFIRST', 'S_BARRIER_INIT', 'S_BARRIER_JOIN')
|
||||
if inst.op_name in sop1_src_only: return f"{name} {src}"
|
||||
return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs(), cdna)}, {src}"
|
||||
|
||||
def _disasm_sop2(inst: SOP2) -> str:
|
||||
cdna, name = _is_cdna(inst), inst.op_name.lower()
|
||||
lit = getattr(inst, '_literal', None)
|
||||
s0 = inst.lit(inst.ssrc0) if inst.ssrc0 == 255 else _fmt_src(inst.ssrc0, inst.src_regs(0), cdna)
|
||||
s1 = inst.lit(inst.ssrc1) if inst.ssrc1 == 255 else _fmt_src(inst.ssrc1, inst.src_regs(1), cdna)
|
||||
dst = _fmt_sdst(inst.sdst, inst.dst_regs(), cdna)
|
||||
if 'fmamk' in name and lit is not None: return f"{name} {dst}, {s0}, 0x{lit:x}, {s1}"
|
||||
if 'fmaak' in name and lit is not None: return f"{name} {dst}, {s0}, {s1}, 0x{lit:x}"
|
||||
return f"{name} {dst}, {s0}, {s1}"
|
||||
|
||||
def _disasm_sopc(inst: SOPC) -> str:
|
||||
cdna = _is_cdna(inst)
|
||||
s0 = inst.lit(inst.ssrc0) if inst.ssrc0 == 255 else _fmt_src(inst.ssrc0, inst.src_regs(0), cdna)
|
||||
s1 = inst.lit(inst.ssrc1) if inst.ssrc1 == 255 else _fmt_src(inst.ssrc1, inst.src_regs(1), cdna)
|
||||
return f"{inst.op_name.lower()} {s0}, {s1}"
|
||||
|
||||
def _disasm_sopk(inst: SOPK) -> str:
|
||||
op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
|
||||
is_rdna4 = 'rdna4' in inst.__class__.__module__
|
||||
hw = HWREG
|
||||
def fmt_hwreg(hid, hoff, hsz):
|
||||
if hid not in hw: return f"0x{inst.simm16:x}"
|
||||
hr_name = str(hid) if is_rdna4 else hw[hid]
|
||||
return f"hwreg({hr_name})" if hoff == 0 and hsz == 32 else f"hwreg({hr_name}, {hoff}, {hsz})"
|
||||
if name == 's_setreg_imm32_b32' or (not cdna and op == SOPKOp.S_SETREG_IMM32_B32):
|
||||
hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
|
||||
return f"{name} {fmt_hwreg(hid, hoff, hsz)}, 0x{inst._literal:x}"
|
||||
if not cdna and op == SOPKOp.S_VERSION: return f"{name} 0x{inst.simm16:x}"
|
||||
if (not cdna and op in (SOPKOp.S_SETREG_B32, SOPKOp.S_GETREG_B32)) or (cdna and name in ('s_setreg_b32', 's_getreg_b32')):
|
||||
hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
|
||||
hs = fmt_hwreg(hid, hoff, hsz)
|
||||
return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1, cdna)}" if 'setreg' in name else f"{name} {_fmt_sdst(inst.sdst, 1, cdna)}, {hs}"
|
||||
if not cdna and op in (SOPKOp.S_SUBVECTOR_LOOP_BEGIN, SOPKOp.S_SUBVECTOR_LOOP_END):
|
||||
return f"{name} {_fmt_sdst(inst.sdst, 1)}, 0x{inst.simm16:x}"
|
||||
return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs(), cdna)}, 0x{inst.simm16:x}"
|
||||
|
||||
def _disasm_vinterp(inst: VINTERP) -> str:
|
||||
mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
|
||||
return f"{inst.op_name.lower()} v{inst.vdst}, {inst.lit(inst.src0, inst.neg & 1)}, {inst.lit(inst.src1, inst.neg & 2)}, {inst.lit(inst.src2, inst.neg & 4)}" + (" " + mods if mods else "")
|
||||
|
||||
EXP_TARGETS = {0: 'mrt0', 1: 'mrt1', 2: 'mrt2', 3: 'mrt3', 4: 'mrt4', 5: 'mrt5', 6: 'mrt6', 7: 'mrt7',
|
||||
8: 'mrtz', 9: 'null', 12: 'pos0', 13: 'pos1', 14: 'pos2', 15: 'pos3', 16: 'pos4',
|
||||
32: 'param0', 33: 'param1', 34: 'param2', 35: 'param3', 36: 'param4', 37: 'param5'}
|
||||
def _disasm_vexport(inst) -> str:
|
||||
tgt = EXP_TARGETS.get(inst.target, f'{inst.target}')
|
||||
srcs = [f'v{getattr(inst, f"vsrc{i}")}' if inst.en & (1 << i) else 'off' for i in range(4)]
|
||||
mods = _mods((inst.done, "done"), (inst.row, "row_en"))
|
||||
return f"export {tgt} {', '.join(srcs)}" + (" " + mods if mods else "")
|
||||
|
||||
def _disasm_vbuffer(inst) -> str:
|
||||
name = inst.op_name.lower().replace('buffer_', 'buffer_').replace('tbuffer_', 'tbuffer_')
|
||||
w = (2 if _has(name, 'xyz', 'xyzw') else 1) if 'd16' in name else \
|
||||
((2 if _has(name, 'b64', 'u64', 'i64') else 1) * (2 if 'cmpswap' in name else 1)) if 'atomic' in name else \
|
||||
{'b32':1,'b64':2,'b96':3,'b128':4,'b16':1,'x':1,'xy':2,'xyz':3,'xyzw':4}.get(name.split('_')[-1], inst.dst_regs())
|
||||
if getattr(inst, 'tfe', 0): w += 1
|
||||
vdata = _vreg(inst.vdata, w) if w else f'v{inst.vdata}'
|
||||
vaddr = _vreg(inst.vaddr, 2) if inst.offen and inst.idxen else (f'v{inst.vaddr}' if inst.offen or inst.idxen else 'off')
|
||||
srsrc = f'ttmp[{inst.rsrc - 108}:{inst.rsrc - 108 + 3}]' if inst.rsrc >= 108 else f's[{inst.rsrc}:{inst.rsrc + 3}]'
|
||||
soff = decode_src(inst.soffset) if inst.soffset >= 106 else f's{inst.soffset}'
|
||||
fmt = getattr(inst, 'format', 0)
|
||||
fmt_names = {e.value: e.name for e in BufFmt}
|
||||
fmt_s = f" format:[{fmt_names[fmt]}]" if fmt > 1 and fmt in fmt_names else (f" format:{fmt}" if fmt > 1 else "")
|
||||
if 'atomic' in name: th_names = {1: 'TH_ATOMIC_RETURN', 6: 'TH_ATOMIC_CASCADE_NT'}
|
||||
elif 'store' in name: th_names = {3: 'TH_STORE_BYPASS', 6: 'TH_STORE_NT_HT'}
|
||||
else: th_names = {3: 'TH_LOAD_BYPASS', 6: 'TH_LOAD_NT_HT'}
|
||||
scope_names = {1: 'SCOPE_SE', 2: 'SCOPE_DEV', 3: 'SCOPE_SYS'}
|
||||
mods = _mods((inst.idxen, "idxen"), (inst.offen, "offen"), (inst.ioffset, f"offset:{inst.ioffset}"),
|
||||
(inst.th in th_names, f"th:{th_names.get(inst.th, '')}"), (inst.scope in scope_names, f"scope:{scope_names.get(inst.scope, '')}"))
|
||||
return f"{name} {vdata}, {vaddr}, {srsrc}, {soff}{fmt_s}" + (" " + mods if mods else "")
|
||||
|
||||
DISASM_HANDLERS: dict[type, callable] = {
|
||||
VOP1: _disasm_vop1, VOP2: _disasm_vop2, VOPC: _disasm_vopc, VOP3: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOPD: _disasm_vopd, VOP3P: _disasm_vop3p,
|
||||
VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, MUBUF: _disasm_buf, MTBUF: _disasm_buf,
|
||||
MIMG: _disasm_mimg, SOP1: _disasm_sop1, SOP2: _disasm_sop2, SOPC: _disasm_sopc, SOPK: _disasm_sopk,
|
||||
# RDNA4
|
||||
R4_VOP1: _disasm_vop1, R4_VOP2: _disasm_vop2, R4_VOPC: _disasm_vopc, R4_VOP3: _disasm_vop3, R4_VOP3SD: _disasm_vop3sd,
|
||||
R4_VOPD: _disasm_vopd, R4_VOP3P: _disasm_vop3p, R4_VINTERP: _disasm_vinterp, R4_SOPP: _disasm_sopp, R4_SMEM: _disasm_smem,
|
||||
R4_DS: _disasm_ds, R4_SOP1: _disasm_sop1, R4_SOP2: _disasm_sop2, R4_SOPC: _disasm_sopc, R4_SOPK: _disasm_sopk,
|
||||
R4_VEXPORT: _disasm_vexport, R4_VBUFFER: _disasm_vbuffer}
|
||||
|
||||
def disasm(inst: Inst) -> str: return DISASM_HANDLERS[type(inst)](inst)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# CDNA DISASSEMBLER SUPPORT
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
try:
|
||||
from extra.assembly.amd.autogen.cdna.ins import (VOP1 as CDNA_VOP1, VOP2 as CDNA_VOP2, VOPC as CDNA_VOPC, VOP3A, VOP3B, VOP3P as CDNA_VOP3P,
|
||||
SOP1 as CDNA_SOP1, SOP2 as CDNA_SOP2, SOPC as CDNA_SOPC, SOPK as CDNA_SOPK, SOPP as CDNA_SOPP, SMEM as CDNA_SMEM, DS as CDNA_DS,
|
||||
FLAT as CDNA_FLAT, MUBUF as CDNA_MUBUF, MTBUF as CDNA_MTBUF, SDWA, DPP, VOP1Op as CDNA_VOP1Op, VOP2Op as CDNA_VOP2Op, VOPCOp as CDNA_VOPCOp)
|
||||
|
||||
def _cdna_src(inst, v, neg, abs_=0, n=1):
|
||||
s = inst.lit(v) if v == 255 else _fmt_src(v, n, cdna=True)
|
||||
if abs_: s = f"|{s}|"
|
||||
return f"neg({s})" if neg and v == 255 else (f"-{s}" if neg else s)
|
||||
|
||||
_CDNA_VOP3_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32'}
|
||||
|
||||
def _disasm_vop3a(inst) -> str:
|
||||
op_val = inst._values.get('op', 0)
|
||||
if hasattr(op_val, 'value'): op_val = op_val.value
|
||||
name = inst.op_name.lower() or f'vop3a_op_{op_val}'
|
||||
from extra.assembly.amd.dsl import spec_num_srcs, spec_regs
|
||||
n = spec_num_srcs(name) if name else inst.num_srcs()
|
||||
cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
|
||||
orig_name = name
|
||||
name = _CDNA_VOP3_ALIASES.get(name, name)
|
||||
if name != orig_name:
|
||||
s0, s1 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, 1), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, 1)
|
||||
s2 = ""
|
||||
dst = f"v{inst.vdst}"
|
||||
else:
|
||||
dregs, r0, r1, r2 = spec_regs(name) if name else (inst.dst_regs(), inst.src_regs(0), inst.src_regs(1), inst.src_regs(2))
|
||||
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 f"v{inst.vdst}"
|
||||
if op_val >= 512:
|
||||
return f"{name} {dst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{cl}{om}"
|
||||
if op_val < 256:
|
||||
sdst = _fmt_sdst(inst.vdst, 2, cdna=True)
|
||||
return f"{name}_e64 {sdst}, {s0}, {s1}{cl}"
|
||||
if 320 <= op_val < 512:
|
||||
if name in ('v_nop', 'v_clrexcp'): return f"{name}_e64"
|
||||
return f"{name}_e64 {dst}, {s0}{cl}{om}"
|
||||
if name == 'v_cndmask_b32':
|
||||
s2 = _fmt_src(inst.src2, 2, cdna=True)
|
||||
return f"{name}_e64 {dst}, {s0}, {s1}, {s2}{cl}{om}"
|
||||
if name in ('v_mul_legacy_f32', 'v_mac_f32'):
|
||||
return f"{name}_e64 {dst}, {s0}, {s1}{cl}{om}"
|
||||
suf = "_e64" if op_val < 512 else ""
|
||||
return f"{name}{suf} {dst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name}{suf} {dst}, {s0}, {s1}{cl}{om}"
|
||||
|
||||
def _disasm_vop3b(inst) -> str:
|
||||
op_val = inst._values.get('op', 0)
|
||||
if hasattr(op_val, 'value'): op_val = op_val.value
|
||||
name = inst.op_name.lower() or f'vop3b_op_{op_val}'
|
||||
from extra.assembly.amd.dsl import spec_num_srcs, spec_regs
|
||||
n = spec_num_srcs(name) if name else inst.num_srcs()
|
||||
dregs, r0, r1, r2 = spec_regs(name) if name else (inst.dst_regs(), inst.src_regs(0), inst.src_regs(1), inst.src_regs(2))
|
||||
s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, n=r0), _cdna_src(inst, inst.src1, inst.neg&2, n=r1), _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
|
||||
dst = _vreg(inst.vdst, dregs) if dregs > 1 else f"v{inst.vdst}"
|
||||
sdst = _fmt_sdst(inst.sdst, 2, cdna=True)
|
||||
cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
|
||||
if name in ('v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'):
|
||||
s2 = _fmt_src(inst.src2, 2, cdna=True)
|
||||
return f"{name}_e64 {dst}, {sdst}, {s0}, {s1}, {s2}{cl}{om}"
|
||||
suf = "_e64" if 'co_' in name else ""
|
||||
return f"{name}{suf} {dst}, {sdst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name}{suf} {dst}, {sdst}, {s0}, {s1}{cl}{om}"
|
||||
|
||||
def _disasm_cdna_vop3p(inst) -> str:
|
||||
name, n, is_mfma = inst.op_name.lower(), inst.num_srcs(), 'mfma' in inst.op_name.lower() or 'smfmac' in inst.op_name.lower()
|
||||
get_src = lambda v, sc: inst.lit(v) if v == 255 else _fmt_src(v, sc, cdna=True)
|
||||
if is_mfma: sc = 2 if 'iu4' in name else 4 if 'iu8' in name or 'i4' in name else 8 if 'f16' in name or 'bf16' in name else 4; src0, src1, src2, dst = get_src(inst.src0, sc), get_src(inst.src1, sc), get_src(inst.src2, 16), _vreg(inst.vdst, 16)
|
||||
else: src0, src1, src2, dst = get_src(inst.src0, 1), get_src(inst.src1, 1), get_src(inst.src2, 1), f"v{inst.vdst}"
|
||||
opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
|
||||
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != (7 if n == 3 else 3) else []) + \
|
||||
([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
|
||||
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
|
||||
|
||||
_SEL = {0: 'BYTE_0', 1: 'BYTE_1', 2: 'BYTE_2', 3: 'BYTE_3', 4: 'WORD_0', 5: 'WORD_1', 6: 'DWORD'}
|
||||
_UNUSED = {0: 'UNUSED_PAD', 1: 'UNUSED_SEXT', 2: 'UNUSED_PRESERVE'}
|
||||
_DPP = {0x130: "wave_shl:1", 0x134: "wave_rol:1", 0x138: "wave_shr:1", 0x13c: "wave_ror:1", 0x140: "row_mirror", 0x141: "row_half_mirror", 0x142: "row_bcast:15", 0x143: "row_bcast:31"}
|
||||
|
||||
def _sdwa_src0(v, is_sgpr, sext=0, neg=0, abs_=0):
|
||||
s = decode_src(v, cdna=True) if is_sgpr else f"v{v}"
|
||||
if sext: s = f"sext({s})"
|
||||
if abs_: s = f"|{s}|"
|
||||
return f"-{s}" if neg else s
|
||||
|
||||
def _sdwa_vsrc1(v, sext=0, neg=0, abs_=0):
|
||||
s = f"v{v}"
|
||||
if sext: s = f"sext({s})"
|
||||
if abs_: s = f"|{s}|"
|
||||
return f"-{s}" if neg else s
|
||||
|
||||
_OMOD_SDWA = {0: "", 1: " mul:2", 2: " mul:4", 3: " div:2"}
|
||||
|
||||
def _disasm_sdwa(inst) -> str:
|
||||
vop2_op = inst.vop2_op
|
||||
src0 = _sdwa_src0(inst.src0, inst.s0, inst.src0_sext, inst.src0_neg, inst.src0_abs)
|
||||
clamp = " clamp" if inst.clmp else ""
|
||||
omod = _OMOD_SDWA.get(inst.omod, "")
|
||||
if vop2_op == 63:
|
||||
try: name = CDNA_VOP1Op(inst.vop_op).name.lower()
|
||||
except ValueError: name = f"vop1_op_{inst.vop_op}"
|
||||
dst = f"v{inst.vdst}"
|
||||
mods = [f"dst_sel:{_SEL[inst.dst_sel]}", f"dst_unused:{_UNUSED[inst.dst_u]}", f"src0_sel:{_SEL[inst.src0_sel]}"]
|
||||
return f"{name}_sdwa {dst}, {src0}{clamp}{omod} " + " ".join(mods)
|
||||
elif vop2_op == 62:
|
||||
try: name = CDNA_VOPCOp(inst.vdst).name.lower()
|
||||
except ValueError: name = f"vopc_op_{inst.vdst}"
|
||||
src1 = _sdwa_vsrc1(inst.vop_op, inst.src1_sext, inst.src1_neg, inst.src1_abs)
|
||||
sdst_enc = inst.dst_sel | (inst.dst_u << 3) | (inst.clmp << 5) | (inst.omod << 6)
|
||||
if sdst_enc == 0: sdst = "vcc"
|
||||
else:
|
||||
sdst_val = sdst_enc - 128 if sdst_enc >= 128 else sdst_enc
|
||||
sdst = _fmt_sdst(sdst_val, 2, cdna=True)
|
||||
mods = [f"src0_sel:{_SEL[inst.src0_sel]}", f"src1_sel:{_SEL[inst.src1_sel]}"]
|
||||
return f"{name}_sdwa {sdst}, {src0}, {src1} " + " ".join(mods)
|
||||
else:
|
||||
try: name = CDNA_VOP2Op(vop2_op).name.lower()
|
||||
except ValueError: name = f"vop2_op_{vop2_op}"
|
||||
name = _CDNA_DISASM_ALIASES.get(name, name)
|
||||
dst = f"v{inst.vdst}"
|
||||
src1 = _sdwa_vsrc1(inst.vop_op, inst.src1_sext, inst.src1_neg, inst.src1_abs)
|
||||
mods = [f"dst_sel:{_SEL[inst.dst_sel]}", f"dst_unused:{_UNUSED[inst.dst_u]}", f"src0_sel:{_SEL[inst.src0_sel]}", f"src1_sel:{_SEL[inst.src1_sel]}"]
|
||||
if name == 'v_cndmask_b32':
|
||||
return f"{name}_sdwa {dst}, {src0}, {src1}, vcc{clamp}{omod} " + " ".join(mods)
|
||||
if name in ('v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'):
|
||||
return f"{name}_sdwa {dst}, vcc, {src0}, {src1}, vcc{clamp}{omod} " + " ".join(mods)
|
||||
if '_co_' in name:
|
||||
return f"{name}_sdwa {dst}, vcc, {src0}, {src1}{clamp}{omod} " + " ".join(mods)
|
||||
return f"{name}_sdwa {dst}, {src0}, {src1}{clamp}{omod} " + " ".join(mods)
|
||||
|
||||
def _dpp_src(v, neg=0, abs_=0):
|
||||
s = f"v{v}" if v < 256 else f"v{v - 256}"
|
||||
if abs_: s = f"|{s}|"
|
||||
return f"-{s}" if neg else s
|
||||
|
||||
def _disasm_dpp(inst) -> str:
|
||||
vop2_op = inst.vop2_op
|
||||
ctrl = inst.dpp_ctrl
|
||||
dpp = f"quad_perm:[{ctrl&3},{(ctrl>>2)&3},{(ctrl>>4)&3},{(ctrl>>6)&3}]" if ctrl < 0x100 else f"row_shl:{ctrl&0xf}" if ctrl < 0x110 else f"row_shr:{ctrl&0xf}" if ctrl < 0x120 else f"row_ror:{ctrl&0xf}" if ctrl < 0x130 else _DPP.get(ctrl, f"dpp_ctrl:0x{ctrl:x}")
|
||||
src0 = _dpp_src(inst.src0, inst.src0_neg, inst.src0_abs)
|
||||
mods = [dpp, f"row_mask:0x{inst.row_mask:x}", f"bank_mask:0x{inst.bank_mask:x}"] + (["bound_ctrl:0"] if inst.bound_ctrl else [])
|
||||
if vop2_op == 63:
|
||||
try: name = CDNA_VOP1Op(inst.vop_op).name.lower()
|
||||
except ValueError: name = f"vop1_op_{inst.vop_op}"
|
||||
return f"{name}_dpp v{inst.vdst}, {src0} " + " ".join(mods)
|
||||
else:
|
||||
try: name = CDNA_VOP2Op(vop2_op).name.lower()
|
||||
except ValueError: name = f"vop2_op_{vop2_op}"
|
||||
name = _CDNA_DISASM_ALIASES.get(name, name)
|
||||
src1 = _dpp_src(inst.vop_op, inst.src1_neg, inst.src1_abs)
|
||||
if name == 'v_cndmask_b32':
|
||||
return f"{name}_dpp v{inst.vdst}, {src0}, {src1}, vcc " + " ".join(mods)
|
||||
if name in ('v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'):
|
||||
return f"{name}_dpp v{inst.vdst}, vcc, {src0}, {src1}, vcc " + " ".join(mods)
|
||||
if '_co_' in name:
|
||||
return f"{name}_dpp v{inst.vdst}, vcc, {src0}, {src1} " + " ".join(mods)
|
||||
return f"{name}_dpp v{inst.vdst}, {src0}, {src1} " + " ".join(mods)
|
||||
|
||||
DISASM_HANDLERS.update({CDNA_VOP1: _disasm_vop1, CDNA_VOP2: _disasm_vop2, CDNA_VOPC: _disasm_vopc,
|
||||
CDNA_SOP1: _disasm_sop1, CDNA_SOP2: _disasm_sop2, CDNA_SOPC: _disasm_sopc, CDNA_SOPK: _disasm_sopk, CDNA_SOPP: _disasm_sopp,
|
||||
CDNA_SMEM: _disasm_smem, CDNA_DS: _disasm_ds, CDNA_FLAT: _disasm_flat, CDNA_MUBUF: _disasm_buf, CDNA_MTBUF: _disasm_buf,
|
||||
VOP3A: _disasm_vop3a, VOP3B: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p, SDWA: _disasm_sdwa, DPP: _disasm_dpp})
|
||||
except ImportError:
|
||||
pass
|
||||
@@ -0,0 +1,613 @@
|
||||
# library for RDNA3 assembly DSL
|
||||
# mypy: ignore-errors
|
||||
from __future__ import annotations
|
||||
import re, struct
|
||||
from enum import IntEnum
|
||||
from functools import cache
|
||||
from typing import overload, Annotated, TypeVar, Generic
|
||||
from extra.assembly.amd.autogen.rdna3.enum import (VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, VOPDOp, SOP1Op, SOP2Op,
|
||||
SOPCOp, SOPKOp, SOPPOp, SMEMOp, DSOp, FLATOp, MUBUFOp, MTBUFOp, MIMGOp, VINTERPOp)
|
||||
from extra.assembly.amd.autogen.cdna.enum import VOP1Op as CDNA_VOP1Op, VOP2Op as CDNA_VOP2Op
|
||||
from extra.assembly.amd.autogen.rdna4.enum import (VOP1Op as RDNA4_VOP1Op, VOP2Op as RDNA4_VOP2Op, VOP3Op as RDNA4_VOP3Op,
|
||||
VOP3SDOp as RDNA4_VOP3SDOp, VOP3POp as RDNA4_VOP3POp, VOPCOp as RDNA4_VOPCOp, VOPDOp as RDNA4_VOPDOp,
|
||||
SOP1Op as RDNA4_SOP1Op, SOP2Op as RDNA4_SOP2Op, SOPCOp as RDNA4_SOPCOp, SOPKOp as RDNA4_SOPKOp, SOPPOp as RDNA4_SOPPOp,
|
||||
SMEMOp as RDNA4_SMEMOp, DSOp as RDNA4_DSOp, VBUFFEROp as RDNA4_VBUFFEROp, VINTERPOp as RDNA4_VINTERPOp)
|
||||
|
||||
# Source operand encoding - constant across all AMD ISAs
|
||||
class SrcEnum(IntEnum):
|
||||
VCC_LO=106; VCC_HI=107; NULL=124; M0=125; EXEC_LO=126; EXEC_HI=127; ZERO=128
|
||||
DPP8=233; DPP8FI=234; SHARED_BASE=235; SHARED_LIMIT=236; PRIVATE_BASE=237; PRIVATE_LIMIT=238
|
||||
POS_HALF=240; NEG_HALF=241; POS_ONE=242; NEG_ONE=243; POS_TWO=244; NEG_TWO=245
|
||||
POS_FOUR=246; NEG_FOUR=247; INV_2PI=248; DPP16=250; VCCZ=251; EXECZ=252; SCC=253; LDS_DIRECT=254
|
||||
VCC_LO, VCC_HI, NULL, M0, EXEC_LO, EXEC_HI, ZERO = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.M0, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.ZERO
|
||||
DPP8FI, SHARED_BASE, SHARED_LIMIT, PRIVATE_BASE, PRIVATE_LIMIT = SrcEnum.DPP8FI, SrcEnum.SHARED_BASE, SrcEnum.SHARED_LIMIT, SrcEnum.PRIVATE_BASE, SrcEnum.PRIVATE_LIMIT
|
||||
POS_HALF, NEG_HALF, POS_ONE, NEG_ONE, POS_TWO, NEG_TWO = SrcEnum.POS_HALF, SrcEnum.NEG_HALF, SrcEnum.POS_ONE, SrcEnum.NEG_ONE, SrcEnum.POS_TWO, SrcEnum.NEG_TWO
|
||||
POS_FOUR, NEG_FOUR, INV_2PI, VCCZ, EXECZ, SCC, LDS_DIRECT = SrcEnum.POS_FOUR, SrcEnum.NEG_FOUR, SrcEnum.INV_2PI, SrcEnum.VCCZ, SrcEnum.EXECZ, SrcEnum.SCC, SrcEnum.LDS_DIRECT
|
||||
OFF = NULL
|
||||
|
||||
# Common masks
|
||||
MASK32, MASK64, MASK128 = 0xffffffff, 0xffffffffffffffff, (1 << 128) - 1
|
||||
|
||||
# Float/int bit conversion (simple versions for literal encoding)
|
||||
def _i32(f: float) -> int: return struct.unpack("<I", struct.pack("<f", f))[0]
|
||||
def _i64(f: float) -> int: return struct.unpack("<Q", struct.pack("<d", f))[0]
|
||||
|
||||
# Instruction spec - register counts and dtypes derived from instruction names
|
||||
_REGS = {'B32': 1, 'B64': 2, 'B96': 3, 'B128': 4, 'B256': 8, 'B512': 16,
|
||||
'F32': 1, 'I32': 1, 'U32': 1, 'F64': 2, 'I64': 2, 'U64': 2,
|
||||
'F16': 1, 'I16': 1, 'U16': 1, 'B16': 1, 'I8': 1, 'U8': 1, 'B8': 1,
|
||||
'DWORD': 1, 'DWORDX2': 2, 'DWORDX3': 3, 'DWORDX4': 4, 'DWORDX8': 8, 'DWORDX16': 16,
|
||||
'BYTE': 1, 'SHORT': 1, 'UBYTE': 1, 'SBYTE': 1, 'USHORT': 1, 'SSHORT': 1}
|
||||
_CVT_RE = re.compile(r'CVT_([FIUB]\d+)_([FIUB]\d+)$')
|
||||
_MAD_MUL_RE = re.compile(r'(?:MAD|MUL)_([IU]\d+)_([IU]\d+)$')
|
||||
_PACK_RE = re.compile(r'PACK_([FIUB]\d+)_([FIUB]\d+)$')
|
||||
_DST_SRC_RE = re.compile(r'_([FIUB]\d+)_([FIUB]\d+)$')
|
||||
_SINGLE_RE = re.compile(r'_([FIUB](?:32|64|16|8|96|128|256|512)|DWORD(?:X(?:2|3|4|8|16))?|[US]?BYTE|[US]?SHORT)$')
|
||||
@cache
|
||||
def _suffix(name: str) -> tuple[str | None, str | None]:
|
||||
name = name.upper()
|
||||
if m := _CVT_RE.search(name): return m.group(1), m.group(2)
|
||||
if m := _MAD_MUL_RE.search(name): return m.group(1), m.group(2)
|
||||
if m := _PACK_RE.search(name): return m.group(1), m.group(2)
|
||||
if m := _DST_SRC_RE.search(name): return m.group(1), m.group(2)
|
||||
if m := _SINGLE_RE.search(name): return m.group(1), m.group(1)
|
||||
return None, None
|
||||
_SPECIAL_REGS = {
|
||||
'V_LSHLREV_B64': (2, 1, 2, 1), 'V_LSHRREV_B64': (2, 1, 2, 1), 'V_ASHRREV_I64': (2, 1, 2, 1),
|
||||
'S_LSHL_B64': (2, 2, 1, 1), 'S_LSHR_B64': (2, 2, 1, 1), 'S_ASHR_I64': (2, 2, 1, 1),
|
||||
'S_BFE_U64': (2, 2, 1, 1), 'S_BFE_I64': (2, 2, 1, 1), 'S_BFM_B64': (2, 1, 1, 1),
|
||||
'S_BITSET0_B64': (2, 1, 1, 1), 'S_BITSET1_B64': (2, 1, 1, 1),
|
||||
'S_BITCMP0_B64': (1, 2, 1, 1), 'S_BITCMP1_B64': (1, 2, 1, 1),
|
||||
'V_LDEXP_F64': (2, 2, 1, 1), 'V_TRIG_PREOP_F64': (2, 2, 1, 1),
|
||||
'V_CMP_CLASS_F64': (1, 2, 1, 1), 'V_CMPX_CLASS_F64': (1, 2, 1, 1),
|
||||
'V_CMP_CLASS_F32': (1, 1, 1, 1), 'V_CMPX_CLASS_F32': (1, 1, 1, 1),
|
||||
'V_CMP_CLASS_F16': (1, 1, 1, 1), 'V_CMPX_CLASS_F16': (1, 1, 1, 1),
|
||||
'V_MAD_U64_U32': (2, 1, 1, 2), 'V_MAD_I64_I32': (2, 1, 1, 2),
|
||||
'V_QSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_U32_U8': (4, 2, 1, 4),
|
||||
'V_CVT_PK_F32_BF8': (2, 1, 1, 1), 'V_CVT_PK_F32_FP8': (2, 1, 1, 1),
|
||||
}
|
||||
_SPECIAL_DTYPE = {
|
||||
'V_LSHLREV_B64': ('B64', 'U32', 'B64', None), 'V_LSHRREV_B64': ('B64', 'U32', 'B64', None), 'V_ASHRREV_I64': ('I64', 'U32', 'I64', None),
|
||||
'S_LSHL_B64': ('B64', 'B64', 'U32', None), 'S_LSHR_B64': ('B64', 'B64', 'U32', None), 'S_ASHR_I64': ('I64', 'I64', 'U32', None),
|
||||
'S_BFE_U64': ('U64', 'U64', 'U32', None), 'S_BFE_I64': ('I64', 'I64', 'U32', None),
|
||||
'S_BFM_B64': ('B64', 'U32', 'U32', None), 'S_BITSET0_B64': ('B64', 'U32', None, None), 'S_BITSET1_B64': ('B64', 'U32', None, None),
|
||||
'S_BITCMP0_B64': ('SCC', 'B64', 'U32', None), 'S_BITCMP1_B64': ('SCC', 'B64', 'U32', None),
|
||||
'V_LDEXP_F64': ('F64', 'F64', 'I32', None), 'V_TRIG_PREOP_F64': ('F64', 'F64', 'U32', None),
|
||||
'V_CMP_CLASS_F64': ('VCC', 'F64', 'U32', None), 'V_CMPX_CLASS_F64': ('EXEC', 'F64', 'U32', None),
|
||||
'V_CMP_CLASS_F32': ('VCC', 'F32', 'U32', None), 'V_CMPX_CLASS_F32': ('EXEC', 'F32', 'U32', None),
|
||||
'V_CMP_CLASS_F16': ('VCC', 'F16', 'U32', None), 'V_CMPX_CLASS_F16': ('EXEC', 'F16', 'U32', None),
|
||||
'V_MAD_U64_U32': ('U64', 'U32', 'U32', 'U64'), 'V_MAD_I64_I32': ('I64', 'I32', 'I32', 'I64'),
|
||||
'V_QSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'), 'V_MQSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'),
|
||||
'V_MQSAD_U32_U8': ('B128', 'B64', 'B64', 'B128'),
|
||||
}
|
||||
@cache
|
||||
def spec_regs(name: str) -> tuple[int, int, int, int]:
|
||||
uname = name.upper()
|
||||
if uname in _SPECIAL_REGS: return _SPECIAL_REGS[uname]
|
||||
if 'SAD' in uname and 'U8' in uname and 'QSAD' not in uname and 'MQSAD' not in uname: return 1, 1, 1, 1
|
||||
dst_suf, src_suf = _suffix(name)
|
||||
return _REGS.get(dst_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1)
|
||||
@cache
|
||||
def spec_dtype(name: str) -> tuple[str | None, str | None, str | None, str | None]:
|
||||
uname = name.upper()
|
||||
if uname in _SPECIAL_DTYPE: return _SPECIAL_DTYPE[uname]
|
||||
if 'SAD' in uname and ('U8' in uname or 'U16' in uname) and 'QSAD' not in uname and 'MQSAD' not in uname: return 'U32', 'U32', 'U32', 'U32'
|
||||
if '_CMP_' in uname or '_CMPX_' in uname:
|
||||
dst_suf, src_suf = _suffix(name)
|
||||
return 'EXEC' if '_CMPX_' in uname else 'VCC', src_suf, src_suf, None
|
||||
dst_suf, src_suf = _suffix(name)
|
||||
return dst_suf, src_suf, src_suf, src_suf
|
||||
_F16_RE = re.compile(r'_[FIUB]16(?:_|$)')
|
||||
_F64_RE = re.compile(r'_[FIUB]64(?:_|$)')
|
||||
@cache
|
||||
def spec_is_16bit(name: str) -> bool:
|
||||
uname = name.upper()
|
||||
if 'SAD' in uname or 'PACK' in uname or '_PK_' in uname or 'SAT_PK' in uname or 'DOT2' in uname: return False
|
||||
if '_F32' in uname or '_I32' in uname or '_U32' in uname or '_B32' in uname: return False
|
||||
return bool(_F16_RE.search(uname))
|
||||
@cache
|
||||
def spec_is_64bit(name: str) -> bool: return bool(_F64_RE.search(name.upper()))
|
||||
_3SRC = {'FMA', 'MAD', 'MIN3', 'MAX3', 'MED3', 'DIV_FIX', 'DIV_FMAS', 'DIV_SCALE', 'SAD', 'LERP', 'ALIGN', 'CUBE', 'BFE', 'BFI',
|
||||
'PERM_B32', 'PERMLANE', 'CNDMASK', 'XOR3', 'OR3', 'ADD3', 'LSHL_OR', 'AND_OR', 'LSHL_ADD', 'ADD_LSHL', 'XAD', 'MAXMIN',
|
||||
'MINMAX', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM', 'MINIMUM3', 'MAXIMUM3', 'DOT2', 'DOT4', 'DOT8', 'WMMA', 'CVT_PK_U8', 'MULLIT', 'CO_CI'}
|
||||
_2SRC = {'FMAC', 'PERMLANE16_VAR', 'PERMLANEX16_VAR'} # FMAC uses dst as implicit accumulator, _VAR permlane only 2 sources
|
||||
def spec_num_srcs(name: str) -> int:
|
||||
name = name.upper()
|
||||
if any(k in name for k in _2SRC): return 2
|
||||
return 3 if any(k in name for k in _3SRC) else 2
|
||||
def is_dtype_16(dt: str | None) -> bool: return dt is not None and '16' in dt
|
||||
def is_dtype_64(dt: str | None) -> bool: return dt is not None and '64' in dt
|
||||
|
||||
# Bit field DSL
|
||||
class BitField:
|
||||
def __init__(self, hi: int, lo: int, name: str | None = None): self.hi, self.lo, self.name, self._marker = hi, lo, name, None
|
||||
def __set_name__(self, owner, name):
|
||||
import typing
|
||||
self.name, self._owner = name, owner
|
||||
# Cache marker at class definition time
|
||||
hints = typing.get_type_hints(owner, include_extras=True)
|
||||
if name in hints:
|
||||
hint = hints[name]
|
||||
if typing.get_origin(hint) is Annotated:
|
||||
args = typing.get_args(hint)
|
||||
self._marker = args[1] if len(args) > 1 else None
|
||||
def __eq__(self, val: int) -> tuple[BitField, int]: return (self, val) # type: ignore
|
||||
def mask(self) -> int: return (1 << (self.hi - self.lo + 1)) - 1
|
||||
@property
|
||||
def marker(self) -> type | None: return self._marker
|
||||
@overload
|
||||
def __get__(self, obj: None, objtype: type) -> BitField: ...
|
||||
@overload
|
||||
def __get__(self, obj: object, objtype: type | None = None) -> int: ...
|
||||
def __get__(self, obj, objtype=None):
|
||||
if obj is None: return self
|
||||
val = unwrap(obj._values.get(self.name, 0))
|
||||
# Convert to IntEnum if marker is an IntEnum subclass
|
||||
if self.marker and isinstance(self.marker, type) and issubclass(self.marker, IntEnum):
|
||||
# VOP3 with VOPC opcodes (0-255) -> VOPCOp, VOP3SD opcodes -> VOP3SDOp
|
||||
if self.marker is VOP3Op:
|
||||
if val < 256: return VOPCOp(val)
|
||||
if val in Inst._VOP3SD_OPS: return VOP3SDOp(val)
|
||||
try: return self.marker(val)
|
||||
except ValueError: pass
|
||||
return val
|
||||
|
||||
class _Bits:
|
||||
def __getitem__(self, key) -> BitField: return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
|
||||
bits = _Bits()
|
||||
|
||||
# Source operand with modifiers - base class for anything that can be a src with neg/abs
|
||||
class SrcMod:
|
||||
__slots__ = ('val', 'neg', 'abs_')
|
||||
def __init__(self, val: int, neg: bool = False, abs_: bool = False): self.val, self.neg, self.abs_ = val, neg, abs_
|
||||
def __repr__(self): return f"{'-' if self.neg else ''}{'|' if self.abs_ else ''}{self.val}{'|' if self.abs_ else ''}"
|
||||
def __neg__(self): return SrcMod(self.val, not self.neg, self.abs_)
|
||||
def __abs__(self): return SrcMod(self.val, self.neg, True)
|
||||
|
||||
# Register types
|
||||
class Reg(SrcMod):
|
||||
__slots__ = ('idx', 'count', 'hi')
|
||||
def __init__(self, idx: int, count: int = 1, hi: bool = False, neg: bool = False, abs_: bool = False):
|
||||
self.idx, self.count, self.hi = idx, count, hi
|
||||
super().__init__(idx, neg, abs_)
|
||||
def __repr__(self): return f"{self.__class__.__name__.lower()[0]}[{self.idx}]" if self.count == 1 else f"{self.__class__.__name__.lower()[0]}[{self.idx}:{self.idx + self.count}]"
|
||||
def __neg__(self): return self.__class__(self.idx, self.count, self.hi, not self.neg, self.abs_)
|
||||
def __abs__(self): return self.__class__(self.idx, self.count, self.hi, self.neg, True)
|
||||
@property
|
||||
def l(self): return self.__class__(self.idx, self.count, False, self.neg, self.abs_)
|
||||
@property
|
||||
def h(self): return self.__class__(self.idx, self.count, True, self.neg, self.abs_)
|
||||
|
||||
T = TypeVar('T', bound=Reg)
|
||||
class _RegFactory(Generic[T]):
|
||||
def __init__(self, cls: type[T], name: str): self._cls, self._name = cls, name
|
||||
@overload
|
||||
def __getitem__(self, key: int) -> Reg: ...
|
||||
@overload
|
||||
def __getitem__(self, key: slice) -> Reg: ...
|
||||
def __getitem__(self, key: int | slice) -> Reg:
|
||||
return self._cls(key.start, key.stop - key.start + 1) if isinstance(key, slice) else self._cls(key)
|
||||
def __repr__(self): return f"<{self._name} factory>"
|
||||
|
||||
class SGPR(Reg): pass
|
||||
class VGPR(Reg): pass
|
||||
class TTMP(Reg): pass
|
||||
s: _RegFactory[SGPR] = _RegFactory(SGPR, "SGPR")
|
||||
v: _RegFactory[VGPR] = _RegFactory(VGPR, "VGPR")
|
||||
ttmp: _RegFactory[TTMP] = _RegFactory(TTMP, "TTMP")
|
||||
|
||||
# Special registers as SrcMod objects (support -VCC_LO, abs(EXEC_LO), etc.)
|
||||
VCC_LO, VCC_HI, VCC = SrcMod(106), SrcMod(107), SrcMod(106)
|
||||
EXEC_LO, EXEC_HI, EXEC = SrcMod(126), SrcMod(127), SrcMod(126)
|
||||
SCC, M0, NULL, OFF = SrcMod(253), SrcMod(125), SrcMod(124), SrcMod(124)
|
||||
|
||||
# Field type markers (runtime classes for validation)
|
||||
class _SSrc: pass
|
||||
class _Src: pass
|
||||
class _Imm: pass
|
||||
class _SImm: pass
|
||||
class _VDSTYEnc: pass # VOPD vdsty: encoded = actual >> 1, actual = (encoded << 1) | ((vdstx & 1) ^ 1)
|
||||
class _SGPRField: pass
|
||||
class _VGPRField: pass
|
||||
|
||||
# Type aliases for annotations - tells mypy it's a BitField while preserving marker info
|
||||
SSrc = Annotated[BitField, _SSrc]
|
||||
Src = Annotated[BitField, _Src]
|
||||
Imm = Annotated[BitField, _Imm]
|
||||
SImm = Annotated[BitField, _SImm]
|
||||
VDSTYEnc = Annotated[BitField, _VDSTYEnc]
|
||||
SGPRField = Annotated[BitField, _SGPRField]
|
||||
VGPRField = Annotated[BitField, _VGPRField]
|
||||
class RawImm:
|
||||
def __init__(self, val: int): self.val = val
|
||||
def __repr__(self): return f"RawImm({self.val})"
|
||||
def __eq__(self, other): return isinstance(other, RawImm) and self.val == other.val
|
||||
|
||||
def unwrap(val) -> int:
|
||||
if isinstance(val, RawImm): return val.val
|
||||
if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val # Special registers like VCC_LO, NULL
|
||||
if hasattr(val, 'value'): return val.value # IntEnum
|
||||
if hasattr(val, 'idx'): return val.idx # Reg
|
||||
return val
|
||||
|
||||
# Encoding/decoding constants
|
||||
FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
|
||||
FLOAT_DEC = {v: str(k) for k, v in FLOAT_ENC.items()}
|
||||
SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
|
||||
SPECIAL_GPRS_CDNA = {102: "flat_scratch_lo", 103: "flat_scratch_hi", 104: "xnack_mask_lo", 105: "xnack_mask_hi",
|
||||
106: "vcc_lo", 107: "vcc_hi", 124: "m0", 126: "exec_lo", 127: "exec_hi",
|
||||
251: "src_vccz", 252: "src_execz", 253: "src_scc", 254: "src_lds_direct"}
|
||||
SPECIAL_PAIRS = {106: "vcc", 126: "exec"}
|
||||
SPECIAL_PAIRS_CDNA = {102: "flat_scratch", 104: "xnack_mask", 106: "vcc", 126: "exec"}
|
||||
SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
|
||||
RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata', 'vsrc1'}
|
||||
|
||||
def _encode_reg(val: Reg) -> int: return (108 if isinstance(val, TTMP) else 0) + val.idx
|
||||
|
||||
def _is_encoded_src(v: int) -> bool: return 106 <= v <= 127 or 128 <= v <= 208 or 240 <= v <= 255 # Special regs (106-127) or inline const
|
||||
|
||||
def encode_src(val) -> int:
|
||||
if isinstance(val, VGPR): return 256 + _encode_reg(val)
|
||||
if isinstance(val, Reg): return _encode_reg(val)
|
||||
if isinstance(val, SrcMod) and not isinstance(val, Reg):
|
||||
v = val.val
|
||||
if _is_encoded_src(v): return v # Already encoded (special reg 106-127 or inline const 128-208 or float 240-255)
|
||||
if isinstance(v, int) and 0 <= v <= 64: return 128 + v # Encode as inline constant
|
||||
if isinstance(v, int) and -16 <= v <= -1: return 192 - v
|
||||
return 255 # Literal
|
||||
if hasattr(val, 'value'): return val.value # IntEnum
|
||||
if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
|
||||
if isinstance(val, int): return 128 + val if 0 <= val <= 64 else 192 - val if -16 <= val <= -1 else 255
|
||||
return 255
|
||||
|
||||
def decode_src(val: int, cdna: bool = False) -> str:
|
||||
special = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
|
||||
if val in special: return special[val]
|
||||
if val <= 105: return f"s{val}"
|
||||
if val in FLOAT_DEC: return FLOAT_DEC[val]
|
||||
if 108 <= val <= 123: return f"ttmp{val - 108}"
|
||||
if 128 <= val <= 192: return str(val - 128)
|
||||
if 193 <= val <= 208: return str(-(val - 192))
|
||||
if 256 <= val <= 511: return f"v{val - 256}"
|
||||
return "lit" if val == 255 else f"?{val}"
|
||||
|
||||
# Instruction base class
|
||||
class Inst:
|
||||
_fields: dict[str, BitField]
|
||||
_encoding: tuple[BitField, int] | None = None
|
||||
_defaults: dict[str, int] = {}
|
||||
_values: dict[str, int | RawImm]
|
||||
_words: int # size in 32-bit words, set by decode_program
|
||||
_literal: int | None
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
# Merge fields from parent classes
|
||||
cls._fields = {}
|
||||
for base in reversed(cls.__mro__):
|
||||
if base is Inst or not hasattr(base, '_fields'): continue
|
||||
cls._fields.update(base._fields)
|
||||
# Add this class's own fields (overrides parents)
|
||||
cls._fields.update({n: v[0] if isinstance(v, tuple) else v for n, v in cls.__dict__.items() if isinstance(v, BitField) or (isinstance(v, tuple) and len(v) == 2 and isinstance(v[0], BitField))})
|
||||
# Compute size from max bit (exclude optional MIMG NSA fields: addr1/addr2 at bits 64+)
|
||||
optional_nsa = {'addr1', 'addr2'}
|
||||
max_bit = max((bf.hi for n, bf in cls._fields.items() if n not in optional_nsa), default=0) if cls._fields else 0
|
||||
cls._sz = 12 if max_bit > 63 else 8 if max_bit > 31 else 4
|
||||
if 'encoding' in cls._fields and isinstance(cls.__dict__.get('encoding'), tuple): cls._encoding = cls.__dict__['encoding']
|
||||
|
||||
def _or_field(self, name: str, bit: int):
|
||||
cur = self._values.get(name, 0)
|
||||
self._values[name] = (cur.val if isinstance(cur, RawImm) else cur) | bit
|
||||
|
||||
def _encode_src(self, name: str, val):
|
||||
"""Encode a source field, handling modifiers and literals."""
|
||||
encoded = encode_src(val)
|
||||
has_opsel = 'opsel' in self._fields
|
||||
if isinstance(val, Reg) and val.hi and not has_opsel: encoded |= 0x80 # hi bit in src for VOP1/2/C
|
||||
self._values[name] = RawImm(encoded)
|
||||
# Handle neg/abs/opsel modifiers
|
||||
if isinstance(val, SrcMod):
|
||||
mod_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
|
||||
if val.neg and 'neg' in self._fields: self._or_field('neg', mod_bit)
|
||||
if val.abs_ and 'abs' in self._fields: self._or_field('abs', mod_bit)
|
||||
if isinstance(val, Reg) and val.hi and has_opsel:
|
||||
self._or_field('opsel', {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0))
|
||||
# Track literal value if needed
|
||||
if encoded == 255 and self._literal is None:
|
||||
import struct
|
||||
# Check if THIS source uses 64-bit encoding (not just src0)
|
||||
src_idx = {'src0': 0, 'src1': 1, 'src2': 2, 'ssrc0': 0, 'ssrc1': 1}.get(name, 0)
|
||||
src_regs = self.src_regs(src_idx)
|
||||
is_64 = src_regs == 2
|
||||
if isinstance(val, SrcMod) and not isinstance(val, Reg): lit32 = val.val & MASK32
|
||||
elif isinstance(val, int) and not isinstance(val, IntEnum): lit32 = val & MASK32
|
||||
elif isinstance(val, float): lit32 = (_i64(val) >> 32) if is_64 else _i32(val) # f64: high 32 bits of f64 repr
|
||||
else: return
|
||||
self._literal = (lit32 << 32) if is_64 else lit32
|
||||
|
||||
def _encode_raw(self, name: str, val):
|
||||
"""Encode a raw register field (vdst, vdata, etc.)."""
|
||||
if isinstance(val, Reg):
|
||||
encoded = _encode_reg(val)
|
||||
if val.hi and 'opsel' not in self._fields: encoded |= 0x80
|
||||
self._values[name] = encoded
|
||||
if name == 'vdst' and val.hi and 'opsel' in self._fields: self._or_field('opsel', 8)
|
||||
elif hasattr(val, 'value'): self._values[name] = val.value
|
||||
|
||||
def _validate(self, orig_args: dict):
|
||||
"""Format-specific validation. Override in subclass or check by class name."""
|
||||
cls_name, op = self.__class__.__name__, orig_args.get('op')
|
||||
op_val = op.value if hasattr(op, 'value') else op
|
||||
op_name = op.name if hasattr(op, 'name') else None
|
||||
# SMEM: register count must match opcode (derive from name: b32=1, b64=2, b96=3, b128=4, b256=8, b512=16, i8/u8/i16/u16=1)
|
||||
if cls_name == 'SMEM' and op_name:
|
||||
expected = {'B32': 1, 'B64': 2, 'B96': 3, 'B128': 4, 'B256': 8, 'B512': 16, 'I8': 1, 'U8': 1, 'I16': 1, 'U16': 1}.get(op_name.split('_')[-1])
|
||||
sdata = orig_args.get('sdata')
|
||||
if expected and isinstance(sdata, Reg) and sdata.count != expected:
|
||||
raise ValueError(f"SMEM op {op_name} expects {expected} registers, got {sdata.count}")
|
||||
# SOP1: derive expected register sizes from op name (e.g., S_MOV_B64 -> dst=2, src=2; S_CTZ_I32_B64 -> dst=1, src=2)
|
||||
if cls_name == 'SOP1' and hasattr(orig_args.get('op'), 'name'):
|
||||
op_name = orig_args['op'].name
|
||||
# Special cases: BITSET takes bit index (1 reg) regardless of dst size
|
||||
if 'BITSET' in op_name:
|
||||
dst_size = 2 if op_name.endswith('_B64') else 1
|
||||
src_size = 1 # bit index is always 1 reg
|
||||
else:
|
||||
# Extract sizes from name: last suffix is src type, second-to-last (if exists) is dst type
|
||||
sizes = {'B32': 1, 'I32': 1, 'U32': 1, 'B64': 2, 'I64': 2, 'U64': 2, 'B128': 4, 'B256': 8, 'B512': 16}
|
||||
parts = op_name.split('_')
|
||||
src_size = sizes.get(parts[-1], 1) if parts[-1] in sizes else 1
|
||||
dst_size = sizes.get(parts[-2], src_size) if len(parts) >= 2 and parts[-2] in sizes else src_size
|
||||
for fld, expected in [('sdst', dst_size), ('ssrc0', src_size)]:
|
||||
if isinstance(orig_args.get(fld), Reg) and orig_args[fld].count != expected:
|
||||
raise ValueError(f"SOP1 {op_name} expects {expected} register(s) for {fld}, got {orig_args[fld].count}")
|
||||
|
||||
def __init__(self, *args, literal: int | None = None, **kwargs):
|
||||
self._values, self._literal = dict(self._defaults), None
|
||||
field_names = [n for n in self._fields if n != 'encoding']
|
||||
# Map Python-friendly names to actual field names (abs_ -> abs for Python reserved word)
|
||||
if 'abs_' in kwargs: kwargs['abs'] = kwargs.pop('abs_')
|
||||
# If more args than fields, treat extra arg as literal (for FMAAK/FMAMK style instructions)
|
||||
# FMAMK has K in middle (vdst, src0, K, vsrc1), FMAAK has K at end (vdst, src0, vsrc1, K)
|
||||
args = list(args)
|
||||
if len(args) > len(field_names) and literal is None:
|
||||
for i, a in enumerate(args):
|
||||
if isinstance(a, int) and not isinstance(a, SrcEnum) and i < len(field_names) and field_names[i] in ('vsrc1',):
|
||||
literal = args.pop(i)
|
||||
break
|
||||
else:
|
||||
literal = args.pop() # fallback: last arg is literal
|
||||
orig_args = dict(zip(field_names, args)) | kwargs
|
||||
self._values.update(orig_args)
|
||||
self._precompute()
|
||||
self._validate(orig_args)
|
||||
# Pre-shift literal for 64-bit sources (literal param is always raw 32-bit value from user)
|
||||
if literal is not None:
|
||||
# Find which source uses the literal (255) and check its register count
|
||||
for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
|
||||
v = orig_args.get(n)
|
||||
if (isinstance(v, RawImm) and v.val == 255) or (isinstance(v, int) and v == 255):
|
||||
self._literal = (literal << 32) if self.src_regs(idx) == 2 else literal
|
||||
break
|
||||
else:
|
||||
self._literal = literal # fallback if no literal source found
|
||||
cls_name = self.__class__.__name__
|
||||
|
||||
# Format-specific setup
|
||||
if cls_name == 'FLAT' and 'sve' in self._fields:
|
||||
seg = self._values.get('seg', 0)
|
||||
if (seg.val if isinstance(seg, RawImm) else seg) == 1 and isinstance(orig_args.get('addr'), VGPR): self._values['sve'] = 1
|
||||
if cls_name == 'VOP3P':
|
||||
op = orig_args.get('op')
|
||||
if hasattr(op, 'value'): op = op.value
|
||||
# fma_mix ops (32-34) default to opsel_hi=0, WMMA ops (64-69) default to opsel_hi=7 to match LLVM
|
||||
if op in (32, 33, 34) and 'opsel_hi' not in orig_args: self._values['opsel_hi'] = self._values['opsel_hi2'] = 0
|
||||
if op in range(64, 70) and 'opsel_hi' not in orig_args: self._values['opsel_hi'], self._values['opsel_hi2'] = 3, 1
|
||||
|
||||
# Encode all fields
|
||||
for name, val in list(self._values.items()):
|
||||
if name == 'encoding': continue
|
||||
if isinstance(val, RawImm):
|
||||
if name in RAW_FIELDS: self._values[name] = val.val
|
||||
continue
|
||||
field = self._fields.get(name)
|
||||
marker = field.marker if field else None
|
||||
# Type validation
|
||||
if marker is _SGPRField and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
|
||||
if marker is _VGPRField and not isinstance(val, VGPR): raise TypeError(f"field '{name}' requires VGPR, got {type(val).__name__}")
|
||||
if marker is _SSrc and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires scalar source, got VGPR")
|
||||
# Encode by field type
|
||||
if name in SRC_FIELDS: self._encode_src(name, val)
|
||||
elif name in RAW_FIELDS: self._encode_raw(name, val)
|
||||
elif name == 'sbase': self._values[name] = (val.idx if isinstance(val, Reg) else val.val if isinstance(val, SrcMod) else val * 2) // 2
|
||||
elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg): self._values[name] = _encode_reg(val) // 4
|
||||
elif marker is _VDSTYEnc and isinstance(val, VGPR): self._values[name] = val.idx >> 1
|
||||
self._precompute_fields()
|
||||
|
||||
def _encode_field(self, name: str, val) -> int:
|
||||
if isinstance(val, RawImm): return val.val
|
||||
if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val # Special regs like VCC_LO
|
||||
if name in {'srsrc', 'ssamp'}: return _encode_reg(val) // 4 if isinstance(val, Reg) else val
|
||||
if name == 'sbase': return val.idx // 2 if isinstance(val, Reg) else val.val // 2 if isinstance(val, SrcMod) else val
|
||||
if name in RAW_FIELDS: return _encode_reg(val) if isinstance(val, Reg) else val
|
||||
if isinstance(val, Reg) or name in SRC_FIELDS: return encode_src(val)
|
||||
return val.value if hasattr(val, 'value') else val
|
||||
|
||||
def to_int(self) -> int:
|
||||
word = (self._encoding[1] & self._encoding[0].mask()) << self._encoding[0].lo if self._encoding else 0
|
||||
for n, bf in self._fields.items():
|
||||
if n != 'encoding' and n in self._values: word |= (self._encode_field(n, self._values[n]) & bf.mask()) << bf.lo
|
||||
return word
|
||||
|
||||
def _get_literal(self) -> int | None:
|
||||
for n in SRC_FIELDS:
|
||||
if n in self._values and not isinstance(v := self._values[n], RawImm) and isinstance(v, int) and not isinstance(v, IntEnum) and not (0 <= v <= 64 or -16 <= v <= -1): return v
|
||||
return None
|
||||
|
||||
def _is_64bit_op(self) -> bool:
|
||||
"""Check if this instruction uses 64-bit operands (and thus 64-bit literals)."""
|
||||
op = self._values.get('op')
|
||||
if op is None: return False
|
||||
op_name = op.name if hasattr(op, 'name') else None
|
||||
# Look up op name from int if needed (happens in from_bytes path)
|
||||
if op_name is None and self.__class__.__name__ == 'VOP3':
|
||||
try: op_name = VOP3Op(op).name
|
||||
except ValueError: pass
|
||||
if op_name is None and self.__class__.__name__ == 'VOPC':
|
||||
try: op_name = VOPCOp(op).name
|
||||
except ValueError: pass
|
||||
if op_name is None: return False
|
||||
# V_LDEXP_F64 has 32-bit integer src1, so literal is 32-bit
|
||||
return op_name != 'V_LDEXP_F64' and op_name.endswith(('_F64', '_B64', '_I64', '_U64'))
|
||||
|
||||
def to_bytes(self) -> bytes:
|
||||
result = self.to_int().to_bytes(self._size(), 'little')
|
||||
lit = self._get_literal() or getattr(self, '_literal', None)
|
||||
if lit is None: return result
|
||||
# For 64-bit sources, literal is stored in high 32 bits internally, but encoded as 4 bytes
|
||||
# Find which source uses the literal (255) and check its register count
|
||||
lit_src_is_64 = False
|
||||
for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
|
||||
if n not in self._values: continue
|
||||
v = self._values[n]
|
||||
if (isinstance(v, RawImm) and v.val == 255) or (isinstance(v, int) and v == 255):
|
||||
lit_src_is_64 = self.is_src_64(idx)
|
||||
break
|
||||
lit32 = (lit >> 32) if lit_src_is_64 else lit
|
||||
return result + (lit32 & MASK32).to_bytes(4, 'little')
|
||||
|
||||
@classmethod
|
||||
def _size(cls) -> int: return cls._sz
|
||||
def size(self) -> int:
|
||||
# Literal is always 4 bytes in the binary (for 64-bit ops, it's in high 32 bits)
|
||||
return self._size() + (4 if self._literal is not None else 0)
|
||||
|
||||
@classmethod
|
||||
def from_int(cls, word: int):
|
||||
inst = object.__new__(cls)
|
||||
inst._values = {n: RawImm(v) if n in SRC_FIELDS else v for n, bf in cls._fields.items() if n != 'encoding' for v in [(word >> bf.lo) & bf.mask()]}
|
||||
inst._literal = None
|
||||
inst._precompute()
|
||||
inst._precompute_fields()
|
||||
return inst
|
||||
|
||||
@classmethod
|
||||
def from_bytes(cls, data: bytes):
|
||||
import typing
|
||||
inst = cls.from_int(int.from_bytes(data[:cls._size()], 'little'))
|
||||
op_val = inst._values.get('op', 0)
|
||||
# Check for instructions that always have a literal constant (FMAMK/FMAAK/MADMK/MADAK, SETREG_IMM32)
|
||||
op_name = ''
|
||||
if cls.__name__ in ('VOP2', 'SOP2', 'SOPK') and 'op' in (hints := typing.get_type_hints(cls, include_extras=True)):
|
||||
if typing.get_origin(hints['op']) is typing.Annotated:
|
||||
try: op_name = typing.get_args(hints['op'])[1](op_val).name
|
||||
except (ValueError, TypeError): pass
|
||||
has_literal = any(x in op_name for x in ('FMAMK', 'FMAAK', 'MADMK', 'MADAK', 'SETREG_IMM32'))
|
||||
# VOPD fmaak/fmamk always have a literal (opx/opy value 1 or 2)
|
||||
opx, opy = inst._values.get('opx', 0), inst._values.get('opy', 0)
|
||||
has_literal = has_literal or (cls.__name__ == 'VOPD' and (opx in (1, 2) or opy in (1, 2)))
|
||||
for n in SRC_FIELDS:
|
||||
if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255: has_literal = True
|
||||
if has_literal:
|
||||
# For 64-bit ops, the literal is 32 bits placed in the HIGH 32 bits of the 64-bit value
|
||||
# (low 32 bits are zero). This is how AMD hardware interprets 32-bit literals for 64-bit ops.
|
||||
# Check which source uses the literal and whether THAT source is 64-bit
|
||||
if len(data) >= cls._size() + 4:
|
||||
lit32 = int.from_bytes(data[cls._size():cls._size()+4], 'little')
|
||||
# Find which source has literal (255) and check its register count
|
||||
lit_src_is_64 = False
|
||||
for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
|
||||
if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255:
|
||||
lit_src_is_64 = inst.src_regs(idx) == 2
|
||||
break
|
||||
inst._literal = (lit32 << 32) if lit_src_is_64 else lit32
|
||||
return inst
|
||||
|
||||
def __repr__(self):
|
||||
# Use _fields order and exclude fields that are 0/default (for consistent repr after roundtrip)
|
||||
def is_zero(v): return (isinstance(v, int) and v == 0) or (isinstance(v, VGPR) and v.idx == 0 and v.count == 1)
|
||||
items = [(k, self._values[k]) for k in self._fields if k in self._values and k != 'encoding'
|
||||
and not (is_zero(self._values[k]) and k not in {'op'})]
|
||||
lit = f", literal={hex(self._literal)}" if self._literal is not None else ""
|
||||
return f"{self.__class__.__name__}({', '.join(f'{k}={v}' for k, v in items)}{lit})"
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
if name.startswith('_'): raise AttributeError(name)
|
||||
return unwrap(self._values.get(name, 0))
|
||||
|
||||
def lit(self, v: int, neg: bool = False) -> str:
|
||||
if v == 255 and self._literal is not None:
|
||||
# For 64-bit sources, literal is stored shifted - extract the 32-bit value
|
||||
lit32 = (self._literal >> 32) if self._literal > 0xffffffff else self._literal
|
||||
s = f"0x{lit32:x}"
|
||||
else:
|
||||
s = decode_src(v, 'cdna' in self.__class__.__module__)
|
||||
return f"-{s}" if neg else s
|
||||
|
||||
def __eq__(self, other):
|
||||
if not isinstance(other, Inst): return NotImplemented
|
||||
return self.__class__ == other.__class__ and self._values == other._values and self._literal == other._literal
|
||||
|
||||
def __hash__(self): return hash((self.__class__.__name__, tuple(sorted((k, repr(v)) for k, v in self._values.items())), self._literal))
|
||||
|
||||
def disasm(self) -> str:
|
||||
from extra.assembly.amd.asm import disasm
|
||||
return disasm(self)
|
||||
|
||||
_enum_map = {'VOP1': VOP1Op, 'VOP2': VOP2Op, 'VOP3': VOP3Op, 'VOP3SD': VOP3SDOp, 'VOP3P': VOP3POp, 'VOPC': VOPCOp,
|
||||
'SOP1': SOP1Op, 'SOP2': SOP2Op, 'SOPC': SOPCOp, 'SOPK': SOPKOp, 'SOPP': SOPPOp,
|
||||
'SMEM': SMEMOp, 'DS': DSOp, 'FLAT': FLATOp, 'MUBUF': MUBUFOp, 'MTBUF': MTBUFOp, 'MIMG': MIMGOp,
|
||||
'VOPD': VOPDOp, 'VINTERP': VINTERPOp}
|
||||
_VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
|
||||
|
||||
def _precompute(self):
|
||||
"""Precompute op, op_name, _spec_regs, _spec_dtype for fast access."""
|
||||
val = self._values.get('op')
|
||||
if val is None: self.op = None
|
||||
elif hasattr(val, 'name'): self.op = val
|
||||
else:
|
||||
cls_name = self.__class__.__name__
|
||||
is_cdna = cls_name in ('VOP3A', 'VOP3B')
|
||||
# Try marker enum first (VOP3AOp, VOP3BOp, etc.)
|
||||
marker = self._fields['op'].marker if 'op' in self._fields else None
|
||||
if marker and issubclass(marker, IntEnum):
|
||||
try: self.op = marker(val)
|
||||
except ValueError: self.op = val
|
||||
elif cls_name in self._enum_map:
|
||||
try: self.op = self._enum_map[cls_name](val)
|
||||
except ValueError: self.op = val
|
||||
else: self.op = val
|
||||
# Fallback for promoted instructions when marker lookup failed
|
||||
if not hasattr(self.op, 'name') and cls_name in ('VOP3', 'VOP3A', 'VOP3B') and isinstance(val, int):
|
||||
if val < 256:
|
||||
try: self.op = VOPCOp(val)
|
||||
except ValueError: pass
|
||||
elif is_cdna and 256 <= val < 512:
|
||||
try: self.op = (CDNA_VOP1Op(val - 320) if val >= 320 else CDNA_VOP2Op(val - 256))
|
||||
except ValueError: pass
|
||||
elif val in self._VOP3SD_OPS and not is_cdna:
|
||||
try: self.op = VOP3SDOp(val)
|
||||
except ValueError: pass
|
||||
elif 256 <= val < 512 and not is_cdna:
|
||||
try: self.op = VOP1Op(val - 384) if val >= 384 else VOP2Op(val - 256)
|
||||
except ValueError: pass
|
||||
self.op_name = self.op.name if hasattr(self.op, 'name') else ''
|
||||
self._spec_regs = spec_regs(self.op_name)
|
||||
self._spec_dtype = spec_dtype(self.op_name)
|
||||
|
||||
def _precompute_fields(self):
|
||||
"""Unwrap all field values as direct attributes for fast access."""
|
||||
for name, val in self._values.items():
|
||||
if name != 'op': setattr(self, name, unwrap(val))
|
||||
def dst_regs(self) -> int: return self._spec_regs[0]
|
||||
def src_regs(self, n: int) -> int: return self._spec_regs[n + 1]
|
||||
def num_srcs(self) -> int: return spec_num_srcs(self.op_name)
|
||||
def dst_dtype(self) -> str | None: return self._spec_dtype[0]
|
||||
def src_dtype(self, n: int) -> str | None: return self._spec_dtype[n + 1]
|
||||
def is_src_16(self, n: int) -> bool: return self._spec_regs[n + 1] == 1 and is_dtype_16(self._spec_dtype[n + 1])
|
||||
def is_src_64(self, n: int) -> bool: return self._spec_regs[n + 1] == 2
|
||||
def is_16bit(self) -> bool: return spec_is_16bit(self.op_name)
|
||||
def is_64bit(self) -> bool: return spec_is_64bit(self.op_name)
|
||||
def is_dst_16(self) -> bool: return self._spec_regs[0] == 1 and is_dtype_16(self._spec_dtype[0])
|
||||
|
||||
|
||||
@@ -0,0 +1,433 @@
|
||||
# RDNA3 emulator - executes compiled pseudocode from AMD ISA PDF
|
||||
# mypy: ignore-errors
|
||||
from __future__ import annotations
|
||||
import ctypes, functools
|
||||
from tinygrad.runtime.autogen import hsa
|
||||
from extra.assembly.amd.dsl import Inst, unwrap, FLOAT_ENC, MASK32, MASK64
|
||||
from extra.assembly.amd.pcode import _f32, _i32, _sext, _f16, _i16, _f64, _i64
|
||||
from extra.assembly.amd.decode import decode_inst
|
||||
from extra.assembly.amd.pcode import compile_pseudocode
|
||||
from extra.assembly.amd.autogen.rdna3.str_pcode import PSEUDOCODE_STRINGS
|
||||
from extra.assembly.amd.dsl import SrcEnum
|
||||
from extra.assembly.amd.autogen.rdna3.ins import (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD,
|
||||
SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, SCRATCHOp, VOPDOp)
|
||||
|
||||
WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
|
||||
VCC_LO, VCC_HI, NULL, EXEC_LO, EXEC_HI, SCC = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.SCC
|
||||
|
||||
# Inline constants for src operands 128-254. Build tables for f32, f16, and f64 formats.
|
||||
_FLOAT_CONSTS = {v: k for k, v in FLOAT_ENC.items()} | {248: 0.15915494309189535} # INV_2PI
|
||||
def _build_inline_consts(mask, to_bits):
|
||||
tbl = list(range(65)) + [((-i) & mask) for i in range(1, 17)] + [0] * (127 - 81)
|
||||
for k, v in _FLOAT_CONSTS.items(): tbl[k - 128] = to_bits(v)
|
||||
return tbl
|
||||
_INLINE_CONSTS = _build_inline_consts(MASK32, _i32)
|
||||
_INLINE_CONSTS_F16 = _build_inline_consts(0xffff, _i16)
|
||||
_INLINE_CONSTS_F64 = _build_inline_consts(MASK64, _i64)
|
||||
|
||||
# Helper: extract/write 16-bit half from/to 32-bit value
|
||||
def _src16(raw: int, is_hi: bool) -> int: return ((raw >> 16) & 0xffff) if is_hi else (raw & 0xffff)
|
||||
def _dst16(cur: int, val: int, is_hi: bool) -> int: return (cur & 0x0000ffff) | ((val & 0xffff) << 16) if is_hi else (cur & 0xffff0000) | (val & 0xffff)
|
||||
def _vgpr_hi(src: int) -> bool: return src >= 256 and ((src - 256) & 0x80) != 0
|
||||
def _vgpr_masked(src: int) -> int: return ((src - 256) & 0x7f) + 256 if src >= 256 else src
|
||||
|
||||
# VOP3 source modifier: apply abs/neg to value
|
||||
def _mod_src(val: int, idx: int, neg: int, abs_: int, is64: bool = False) -> int:
|
||||
to_f, to_i = (_f64, _i64) if is64 else (_f32, _i32)
|
||||
if (abs_ >> idx) & 1: val = to_i(abs(to_f(val)))
|
||||
if (neg >> idx) & 1: val = to_i(-to_f(val))
|
||||
return val
|
||||
|
||||
# Read source operand with VOP3 modifiers
|
||||
def _read_src(st, inst, src, idx: int, lane: int, neg: int, abs_: int, opsel: int) -> int:
|
||||
if src is None: return 0
|
||||
literal, regs, is_src_16 = inst._literal, inst.src_regs(idx), inst.is_src_16(idx)
|
||||
if regs == 2: return _mod_src(st.rsrc64(src, lane, literal), idx, neg, abs_, is64=True)
|
||||
if isinstance(inst, VOP3P):
|
||||
opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
|
||||
if 'FMA_MIX' in inst.op_name:
|
||||
raw = st.rsrc(src, lane, literal)
|
||||
sign_bit = (15 if not (opsel & (1 << idx)) else 31) if (opsel_hi >> idx) & 1 else 31
|
||||
if inst.neg_hi & (1 << idx): raw &= ~(1 << sign_bit)
|
||||
if neg & (1 << idx): raw ^= (1 << sign_bit)
|
||||
return raw
|
||||
raw = st.rsrc_f16(src, lane, literal)
|
||||
hi = _src16(raw, opsel_hi & (1 << idx)) ^ (0x8000 if inst.neg_hi & (1 << idx) else 0)
|
||||
lo = _src16(raw, opsel & (1 << idx)) ^ (0x8000 if neg & (1 << idx) else 0)
|
||||
return (hi << 16) | lo
|
||||
if is_src_16 and isinstance(inst, VOP3):
|
||||
raw = st.rsrc_f16(src, lane, literal) if 128 <= src < 255 else st.rsrc(src, lane, literal)
|
||||
val = _src16(raw, bool(opsel & (1 << idx)))
|
||||
if abs_ & (1 << idx): val &= 0x7fff
|
||||
if neg & (1 << idx): val ^= 0x8000
|
||||
return val
|
||||
if is_src_16 and isinstance(inst, (VOP1, VOP2, VOPC)):
|
||||
if src >= 256: return _src16(_mod_src(st.rsrc(_vgpr_masked(src), lane, literal), idx, neg, abs_), _vgpr_hi(src))
|
||||
return _mod_src(st.rsrc_f16(src, lane, literal), idx, neg, abs_) & 0xffff
|
||||
return _mod_src(st.rsrc(src, lane, literal), idx, neg, abs_)
|
||||
|
||||
# Helper: get number of dwords from memory op name
|
||||
def _op_ndwords(name: str) -> int:
|
||||
if '_B128' in name: return 4
|
||||
if '_B96' in name: return 3
|
||||
if any(s in name for s in ('_B64', '_U64', '_I64', '_F64')): return 2
|
||||
return 1
|
||||
|
||||
# Helper: build multi-dword int from consecutive VGPRs
|
||||
def _vgpr_read(V: list, base: int, ndwords: int) -> int: return sum(V[base + i] << (32 * i) for i in range(ndwords))
|
||||
|
||||
# Helper: write multi-dword value to consecutive VGPRs
|
||||
def _vgpr_write(V: list, base: int, val: int, ndwords: int):
|
||||
for i in range(ndwords): V[base + i] = (val >> (32 * i)) & MASK32
|
||||
|
||||
# Memory access
|
||||
_valid_mem_ranges: list[tuple[int, int]] = []
|
||||
def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: _valid_mem_ranges.clear(); _valid_mem_ranges.extend(ranges)
|
||||
def _mem_valid(addr: int, size: int) -> bool:
|
||||
return not _valid_mem_ranges or any(s <= addr and addr + size <= s + z for s, z in _valid_mem_ranges)
|
||||
def _ctypes_at(addr: int, size: int): return (ctypes.c_uint8 if size == 1 else ctypes.c_uint16 if size == 2 else ctypes.c_uint64 if size == 8 else ctypes.c_uint32).from_address(addr)
|
||||
def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
|
||||
def mem_write(addr: int, size: int, val: int) -> None:
|
||||
if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
|
||||
|
||||
def _make_mem_accessor(read_fn, write_fn):
|
||||
"""Create a memory accessor class with the given read/write functions."""
|
||||
class _MemAccessor:
|
||||
__slots__ = ('_addr',)
|
||||
def __init__(self, addr: int): self._addr = int(addr)
|
||||
u8 = property(lambda s: read_fn(s._addr, 1), lambda s, v: write_fn(s._addr, 1, int(v)))
|
||||
u16 = property(lambda s: read_fn(s._addr, 2), lambda s, v: write_fn(s._addr, 2, int(v)))
|
||||
u32 = property(lambda s: read_fn(s._addr, 4), lambda s, v: write_fn(s._addr, 4, int(v)))
|
||||
u64 = property(lambda s: read_fn(s._addr, 8), lambda s, v: write_fn(s._addr, 8, int(v)))
|
||||
i8 = property(lambda s: _sext(read_fn(s._addr, 1), 8), lambda s, v: write_fn(s._addr, 1, int(v)))
|
||||
i16 = property(lambda s: _sext(read_fn(s._addr, 2), 16), lambda s, v: write_fn(s._addr, 2, int(v)))
|
||||
i32 = property(lambda s: _sext(read_fn(s._addr, 4), 32), lambda s, v: write_fn(s._addr, 4, int(v)))
|
||||
i64 = property(lambda s: _sext(read_fn(s._addr, 8), 64), lambda s, v: write_fn(s._addr, 8, int(v)))
|
||||
b8, b16, b32, b64 = u8, u16, u32, u64
|
||||
return _MemAccessor
|
||||
|
||||
_GlobalMemAccessor = _make_mem_accessor(mem_read, mem_write)
|
||||
|
||||
class _GlobalMem:
|
||||
"""Global memory wrapper that supports MEM[addr].u32 style access."""
|
||||
def __getitem__(self, addr) -> _GlobalMemAccessor: return _GlobalMemAccessor(addr)
|
||||
GlobalMem = _GlobalMem()
|
||||
|
||||
class LDSMem:
|
||||
"""LDS memory wrapper that supports MEM[addr].u32 style access."""
|
||||
__slots__ = ('_lds',)
|
||||
def __init__(self, lds: bytearray): self._lds = lds
|
||||
def _read(self, addr: int, size: int) -> int:
|
||||
addr = addr & 0xffff
|
||||
return int.from_bytes(self._lds[addr:addr+size], 'little') if addr + size <= len(self._lds) else 0
|
||||
def _write(self, addr: int, size: int, val: int):
|
||||
addr = addr & 0xffff
|
||||
if addr + size <= len(self._lds): self._lds[addr:addr+size] = (int(val) & ((1 << (size*8)) - 1)).to_bytes(size, 'little')
|
||||
def __getitem__(self, addr): return _make_mem_accessor(self._read, self._write)(addr)
|
||||
|
||||
# SMEM dst register count (for writing result back to SGPRs)
|
||||
SMEM_DST_COUNT = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}
|
||||
|
||||
# VOPD op -> VOP3 op mapping (VOPD is dual-issue of VOP1/VOP2 ops, use VOP3 enums for pseudocode lookup)
|
||||
_VOPD_TO_VOP = {
|
||||
VOPDOp.V_DUAL_FMAC_F32: VOP3Op.V_FMAC_F32, VOPDOp.V_DUAL_FMAAK_F32: VOP2Op.V_FMAAK_F32, VOPDOp.V_DUAL_FMAMK_F32: VOP2Op.V_FMAMK_F32,
|
||||
VOPDOp.V_DUAL_MUL_F32: VOP3Op.V_MUL_F32, VOPDOp.V_DUAL_ADD_F32: VOP3Op.V_ADD_F32, VOPDOp.V_DUAL_SUB_F32: VOP3Op.V_SUB_F32,
|
||||
VOPDOp.V_DUAL_SUBREV_F32: VOP3Op.V_SUBREV_F32, VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: VOP3Op.V_MUL_DX9_ZERO_F32,
|
||||
VOPDOp.V_DUAL_MOV_B32: VOP3Op.V_MOV_B32, VOPDOp.V_DUAL_CNDMASK_B32: VOP3Op.V_CNDMASK_B32,
|
||||
VOPDOp.V_DUAL_MAX_F32: VOP3Op.V_MAX_F32, VOPDOp.V_DUAL_MIN_F32: VOP3Op.V_MIN_F32,
|
||||
VOPDOp.V_DUAL_ADD_NC_U32: VOP3Op.V_ADD_NC_U32, VOPDOp.V_DUAL_LSHLREV_B32: VOP3Op.V_LSHLREV_B32, VOPDOp.V_DUAL_AND_B32: VOP3Op.V_AND_B32,
|
||||
}
|
||||
|
||||
|
||||
class WaveState:
|
||||
__slots__ = ('sgpr', 'vgpr', 'scc', 'pc', '_pend_sgpr', 'lds', 'n_lanes')
|
||||
def __init__(self, lds: LDSMem | None = None, n_lanes: int = WAVE_SIZE):
|
||||
self.sgpr, self.vgpr = [0] * SGPR_COUNT, [[0] * VGPR_COUNT for _ in range(WAVE_SIZE)]
|
||||
self.sgpr[EXEC_LO], self.scc, self.pc, self._pend_sgpr, self.lds, self.n_lanes = 0xffffffff, 0, 0, {}, lds, n_lanes
|
||||
|
||||
@property
|
||||
def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
|
||||
@vcc.setter
|
||||
def vcc(self, v: int): self.sgpr[VCC_LO], self.sgpr[VCC_HI] = v & MASK32, (v >> 32) & MASK32
|
||||
@property
|
||||
def exec_mask(self) -> int: return self.sgpr[EXEC_LO] | (self.sgpr[EXEC_HI] << 32)
|
||||
@exec_mask.setter
|
||||
def exec_mask(self, v: int): self.sgpr[EXEC_LO], self.sgpr[EXEC_HI] = v & MASK32, (v >> 32) & MASK32
|
||||
|
||||
def rsgpr(self, i: int) -> int: return 0 if i == NULL else self.scc if i == SCC else self.sgpr[i] if i < SGPR_COUNT else 0
|
||||
def wsgpr(self, i: int, v: int):
|
||||
if i < SGPR_COUNT and i != NULL: self.sgpr[i] = v & MASK32
|
||||
def rsgpr64(self, i: int) -> int: return self.rsgpr(i) | (self.rsgpr(i+1) << 32)
|
||||
def wsgpr64(self, i: int, v: int): self.wsgpr(i, v & MASK32); self.wsgpr(i+1, (v >> 32) & MASK32)
|
||||
|
||||
def _rsrc_base(self, v: int, lane: int, consts, literal: int):
|
||||
if v < SGPR_COUNT: return self.sgpr[v]
|
||||
if v == SCC: return self.scc
|
||||
if v < 255: return consts[v - 128]
|
||||
if v == 255: return literal
|
||||
return self.vgpr[lane][v - 256] if v <= 511 else 0
|
||||
def rsrc(self, v: int, lane: int, literal: int = 0) -> int: return self._rsrc_base(v, lane, _INLINE_CONSTS, literal)
|
||||
def rsrc_f16(self, v: int, lane: int, literal: int = 0) -> int: return self._rsrc_base(v, lane, _INLINE_CONSTS_F16, literal)
|
||||
def rsrc64(self, v: int, lane: int, literal: int = 0) -> int:
|
||||
if 128 <= v < 255: return _INLINE_CONSTS_F64[v - 128]
|
||||
if v == 255: return literal # literal is already shifted in from_bytes for 64-bit ops
|
||||
return self.rsrc(v, lane, literal) | ((self.rsrc(v+1, lane, literal) if v < VCC_LO or 256 <= v <= 511 else 0) << 32)
|
||||
|
||||
def pend_sgpr_lane(self, reg: int, lane: int, val: int):
|
||||
if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
|
||||
if val: self._pend_sgpr[reg] |= (1 << lane)
|
||||
def commit_pends(self):
|
||||
for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
|
||||
self._pend_sgpr.clear()
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# EXECUTION - All ops use pseudocode from PDF
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def exec_scalar(st: WaveState, inst: Inst):
|
||||
"""Execute scalar instruction. Returns 0 to continue execution."""
|
||||
# Get op enum and lookup compiled function
|
||||
if isinstance(inst, SMEM): ssrc0, sdst = None, None
|
||||
elif isinstance(inst, SOP1): ssrc0, sdst = inst.ssrc0, inst.sdst
|
||||
elif isinstance(inst, SOP2): ssrc0, sdst = inst.ssrc0, inst.sdst
|
||||
elif isinstance(inst, SOPC): ssrc0, sdst = inst.ssrc0, None
|
||||
elif isinstance(inst, SOPK): ssrc0, sdst = inst.sdst, inst.sdst # sdst is both src and dst
|
||||
elif isinstance(inst, SOPP): ssrc0, sdst = None, None
|
||||
else: raise NotImplementedError(f"Unknown scalar type {type(inst)}")
|
||||
|
||||
# SMEM: memory loads
|
||||
if isinstance(inst, SMEM):
|
||||
addr = st.rsgpr64(inst.sbase * 2) + _sext(inst.offset, 21)
|
||||
if inst.soffset not in (NULL, 0x7f): addr += st.rsrc(inst.soffset, 0, inst._literal)
|
||||
result = inst._fn(GlobalMem, addr & MASK64)
|
||||
if 'SDATA' in result:
|
||||
sdata = result['SDATA']
|
||||
for i in range(SMEM_DST_COUNT.get(inst.op, 1)): st.wsgpr(inst.sdata + i, (sdata >> (i * 32)) & MASK32)
|
||||
st.pc += inst._words
|
||||
return 0
|
||||
|
||||
# Build context - use inst methods to determine operand sizes
|
||||
literal = inst._literal
|
||||
s0 = st.rsrc64(ssrc0, 0, literal) if inst.is_src_64(0) else (st.rsrc(ssrc0, 0, literal) if not isinstance(inst, (SOPK, SOPP)) else (st.rsgpr(inst.sdst) if isinstance(inst, SOPK) else 0))
|
||||
s1 = st.rsrc64(inst.ssrc1, 0, literal) if inst.is_src_64(1) else (st.rsrc(inst.ssrc1, 0, literal) if isinstance(inst, (SOP2, SOPC)) else inst.simm16 if isinstance(inst, SOPK) else 0)
|
||||
d0 = st.rsgpr64(sdst) if inst.dst_regs() == 2 and sdst is not None else (st.rsgpr(sdst) if sdst is not None else 0)
|
||||
literal = inst.simm16 if isinstance(inst, (SOPK, SOPP)) else inst._literal
|
||||
|
||||
# Call compiled function with int parameters
|
||||
result = inst._fn(s0, s1, 0, d0, st.scc, st.vcc & MASK32, 0, st.exec_mask & MASK32, literal, None, pc=st.pc * 4)
|
||||
|
||||
# Apply results (already int values)
|
||||
if sdst is not None and 'D0' in result:
|
||||
(st.wsgpr64 if inst.dst_regs() == 2 else st.wsgpr)(sdst, result['D0'])
|
||||
if 'SCC' in result: st.scc = result['SCC'] & 1
|
||||
if 'EXEC' in result: st.exec_mask = result['EXEC']
|
||||
if 'PC' in result:
|
||||
# Convert absolute byte address to word offset
|
||||
pc_val = result['PC']
|
||||
new_pc = pc_val if pc_val < 0x8000000000000000 else pc_val - 0x10000000000000000
|
||||
st.pc = new_pc // 4
|
||||
else:
|
||||
st.pc += inst._words
|
||||
return 0
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# VECTOR INSTRUCTIONS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def exec_vopd(st: WaveState, inst, V: list, lane: int) -> None:
|
||||
"""VOPD: dual-issue, execute two ops simultaneously (read all inputs before writes)."""
|
||||
literal, vdstx, vdsty = inst._literal, inst.vdstx, (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1)
|
||||
sx0, sx1, dx, sy0, sy1, dy = st.rsrc(inst.srcx0, lane, literal), V[inst.vsrcx1], V[vdstx], st.rsrc(inst.srcy0, lane, literal), V[inst.vsrcy1], V[vdsty]
|
||||
V[vdstx] = inst._fnx(sx0, sx1, 0, dx, st.scc, st.vcc, lane, st.exec_mask, literal, None)['D0']
|
||||
V[vdsty] = inst._fny(sy0, sy1, 0, dy, st.scc, st.vcc, lane, st.exec_mask, literal, None)['D0']
|
||||
|
||||
def exec_flat(st: WaveState, inst, V: list, lane: int) -> None:
|
||||
"""FLAT/GLOBAL/SCRATCH memory ops."""
|
||||
ndwords = _op_ndwords(inst.op_name)
|
||||
addr = V[inst.addr] | (V[inst.addr + 1] << 32)
|
||||
ADDR = (st.rsgpr64(inst.saddr) + V[inst.addr] + _sext(inst.offset, 13)) & MASK64 if inst.saddr not in (NULL, 0x7f) else (addr + _sext(inst.offset, 13)) & MASK64
|
||||
vdata_src = inst.vdst if 'LOAD' in inst.op_name else inst.data
|
||||
result = inst._fn(GlobalMem, ADDR, _vgpr_read(V, vdata_src, ndwords), V[inst.vdst])
|
||||
if 'VDATA' in result: _vgpr_write(V, inst.vdst, result['VDATA'], ndwords)
|
||||
if 'RETURN_DATA' in result: _vgpr_write(V, inst.vdst, result['RETURN_DATA'], ndwords)
|
||||
|
||||
def exec_ds(st: WaveState, inst, V: list, lane: int) -> None:
|
||||
"""DS (LDS) memory ops."""
|
||||
ndwords = _op_ndwords(inst.op_name)
|
||||
data0, data1 = _vgpr_read(V, inst.data0, ndwords), _vgpr_read(V, inst.data1, ndwords) if inst.data1 is not None else 0
|
||||
result = inst._fn(st.lds, V[inst.addr], data0, data1, inst.offset0, inst.offset1)
|
||||
if 'RETURN_DATA' in result and ('_RTN' in inst.op_name or '_LOAD' in inst.op_name):
|
||||
_vgpr_write(V, inst.vdst, result['RETURN_DATA'], ndwords * 2 if '_2ADDR_' in inst.op_name else ndwords)
|
||||
|
||||
def exec_vop(st: WaveState, inst: Inst, V: list, lane: int) -> None:
|
||||
"""VOP1/VOP2/VOP3/VOP3SD/VOP3P/VOPC: standard ALU ops."""
|
||||
if isinstance(inst, VOP3P):
|
||||
src0, src1, src2, vdst, dst_hi = inst.src0, inst.src1, inst.src2, inst.vdst, False
|
||||
neg, abs_, opsel = inst.neg, 0, inst.opsel
|
||||
elif isinstance(inst, VOP1):
|
||||
src0, src1, src2, vdst = inst.src0, None, None, inst.vdst & 0x7f if inst.is_dst_16() else inst.vdst
|
||||
neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst & 0x80) != 0 and inst.is_dst_16()
|
||||
elif isinstance(inst, VOP2):
|
||||
src0, src1, src2, vdst = inst.src0, inst.vsrc1 + 256, None, inst.vdst & 0x7f if inst.is_dst_16() else inst.vdst
|
||||
neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst & 0x80) != 0 and inst.is_dst_16()
|
||||
elif isinstance(inst, (VOP3, VOP3SD)):
|
||||
src0, src1, src2, vdst = inst.src0, inst.src1, (None if isinstance(inst, VOP3) and inst.op.value < 256 else inst.src2), inst.vdst
|
||||
neg, abs_, opsel, dst_hi = (inst.neg, inst.abs, inst.opsel, False) if isinstance(inst, VOP3) else (0, 0, 0, False)
|
||||
elif isinstance(inst, VOPC):
|
||||
src0, src1, src2, vdst, neg, abs_, opsel, dst_hi = inst.src0, inst.vsrc1 + 256, None, VCC_LO, 0, 0, 0, False
|
||||
else:
|
||||
raise NotImplementedError(f"exec_vop: unhandled instruction type {type(inst).__name__}")
|
||||
|
||||
s0 = _read_src(st, inst, src0, 0, lane, neg, abs_, opsel)
|
||||
s1 = _read_src(st, inst, src1, 1, lane, neg, abs_, opsel)
|
||||
s2 = _read_src(st, inst, src2, 2, lane, neg, abs_, opsel)
|
||||
if isinstance(inst, VOP2) and inst.is_16bit(): d0 = _src16(V[vdst], dst_hi)
|
||||
elif inst.dst_regs() == 2: d0 = V[vdst] | (V[vdst + 1] << 32)
|
||||
else: d0 = V[vdst]
|
||||
|
||||
if isinstance(inst, VOP3SD) and 'CO_CI' in inst.op_name: vcc_for_fn = st.rsgpr64(inst.src2)
|
||||
elif isinstance(inst, VOP3) and inst.op in (VOP3Op.V_CNDMASK_B32, VOP3Op.V_CNDMASK_B16) and src2 is not None and src2 < 256: vcc_for_fn = st.rsgpr64(src2)
|
||||
else: vcc_for_fn = st.vcc
|
||||
src0_idx = (src0 - 256) if src0 is not None and src0 >= 256 else (src0 if src0 is not None else 0)
|
||||
extra_kwargs = {'opsel': opsel, 'opsel_hi': inst.opsel_hi | (inst.opsel_hi2 << 2)} if isinstance(inst, VOP3P) and 'FMA_MIX' in inst.op_name else {}
|
||||
result = inst._fn(s0, s1, s2, d0, st.scc, vcc_for_fn, lane, st.exec_mask, inst._literal, st.vgpr, src0_idx, vdst, **extra_kwargs)
|
||||
|
||||
# Check if this is a VOPC instruction (either standalone VOPC or VOP3 with VOPC opcode)
|
||||
is_vopc = isinstance(inst.op, VOPCOp) or (isinstance(inst, VOP3) and inst.op.value < 256)
|
||||
if 'VCC' in result:
|
||||
if isinstance(inst, VOP3SD): st.pend_sgpr_lane(inst.sdst, lane, (result['VCC'] >> lane) & 1)
|
||||
else: st.pend_sgpr_lane(VCC_LO if isinstance(inst, VOP2) and 'CO_CI' in inst.op_name else vdst, lane, (result['VCC'] >> lane) & 1)
|
||||
if 'EXEC' in result:
|
||||
st.pend_sgpr_lane(EXEC_LO, lane, (result['EXEC'] >> lane) & 1)
|
||||
elif is_vopc:
|
||||
st.pend_sgpr_lane(vdst, lane, (result['D0'] >> lane) & 1)
|
||||
if not is_vopc:
|
||||
d0_val = result['D0']
|
||||
if inst.dst_regs() == 2: V[vdst], V[vdst + 1] = d0_val & MASK32, (d0_val >> 32) & MASK32
|
||||
elif not isinstance(inst, VOP3P) and inst.is_dst_16(): V[vdst] = _dst16(V[vdst], d0_val, bool(opsel & 8) if isinstance(inst, VOP3) else dst_hi)
|
||||
else: V[vdst] = d0_val & MASK32
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# WMMA (Wave Matrix Multiply-Accumulate)
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def exec_wmma(st: WaveState, inst, op: VOP3POp) -> None:
|
||||
"""Execute WMMA instruction - 16x16x16 matrix multiply across the wave."""
|
||||
src0, src1, src2, vdst = inst.src0, inst.src1, inst.src2, inst.vdst
|
||||
# Read 16x16 f16 matrix from 16 lanes × 8 VGPRs (2 f16 per VGPR)
|
||||
def read_f16_mat(src):
|
||||
return [f for l in range(16) for r in range(8) for v in [st.vgpr[l][src-256+r] if src >= 256 else st.rsgpr(src+r)] for f in [_f16(v&0xffff), _f16((v>>16)&0xffff)]]
|
||||
mat_a, mat_b = read_f16_mat(src0), read_f16_mat(src1)
|
||||
# Read matrix C (16x16 f32) from lanes 0-31, VGPRs src2 to src2+7
|
||||
mat_c = [_f32(st.vgpr[i % 32][src2 - 256 + i // 32] if src2 >= 256 else st.rsgpr(src2 + i // 32)) for i in range(256)]
|
||||
# Compute D = A × B + C (16x16 matrix multiply)
|
||||
mat_d = [sum(mat_a[row*16+k] * mat_b[col*16+k] for k in range(16)) + mat_c[row*16+col] for row in range(16) for col in range(16)]
|
||||
# Write result - f16 packed or f32
|
||||
if op == VOP3POp.V_WMMA_F16_16X16X16_F16:
|
||||
for i in range(0, 256, 2):
|
||||
st.vgpr[(i//2) % 32][vdst + (i//2)//32] = ((_i16(mat_d[i+1]) & 0xffff) << 16) | (_i16(mat_d[i]) & 0xffff)
|
||||
else:
|
||||
for i in range(256): st.vgpr[i % 32][vdst + i//32] = _i32(mat_d[i])
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PROGRAM DECODE
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
# Wave-level dispatch functions: (st, inst) -> return_code (0 = continue, -1 = end, -2 = barrier)
|
||||
def dispatch_endpgm(st, inst): return -1
|
||||
def dispatch_barrier(st, inst): st.pc += inst._words; return -2
|
||||
def dispatch_nop(st, inst): st.pc += inst._words; return 0
|
||||
def dispatch_wmma(st, inst): exec_wmma(st, inst, inst.op); st.pc += inst._words; return 0
|
||||
def dispatch_writelane(st, inst): st.vgpr[st.rsrc(inst.src1, 0, inst._literal) & 0x1f][inst.vdst] = st.rsrc(inst.src0, 0, inst._literal) & MASK32; st.pc += inst._words; return 0
|
||||
def dispatch_readlane(st, inst):
|
||||
src0_idx = (inst.src0 - 256) if inst.src0 >= 256 else inst.src0
|
||||
s1 = st.rsrc(inst.src1, 0, inst._literal) if getattr(inst, 'src1', None) is not None else 0
|
||||
result = inst._fn(0, s1, 0, 0, st.scc, st.vcc, 0, st.exec_mask, inst._literal, st.vgpr, src0_idx, inst.vdst)
|
||||
st.wsgpr(inst.vdst, result['D0'])
|
||||
st.pc += inst._words; return 0
|
||||
|
||||
# Per-lane dispatch wrapper: wraps per-lane exec functions into wave-level dispatch
|
||||
@functools.cache
|
||||
def dispatch_lane(exec_fn):
|
||||
def dispatch(st, inst):
|
||||
exec_mask, vgpr, n_lanes = st.exec_mask, st.vgpr, st.n_lanes
|
||||
for lane in range(n_lanes):
|
||||
if exec_mask >> lane & 1: exec_fn(st, inst, vgpr[lane], lane)
|
||||
st.commit_pends()
|
||||
st.pc += inst._words
|
||||
return 0
|
||||
return dispatch
|
||||
|
||||
def decode_program(data: bytes) -> dict[int, Inst]:
|
||||
result: dict[int, Inst] = {}
|
||||
i = 0
|
||||
while i < len(data):
|
||||
inst = decode_inst(data[i:])
|
||||
inst._words = inst.size() // 4
|
||||
|
||||
# Determine dispatch function and pcode function
|
||||
if isinstance(inst, SOPP) and inst.op == SOPPOp.S_CODE_END: break
|
||||
elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_ENDPGM: inst._dispatch = dispatch_endpgm
|
||||
elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER: inst._dispatch = dispatch_barrier
|
||||
elif isinstance(inst, SOPP) and inst.op in (SOPPOp.S_CLAUSE, SOPPOp.S_WAITCNT, SOPPOp.S_WAITCNT_DEPCTR, SOPPOp.S_SENDMSG, SOPPOp.S_SET_INST_PREFETCH_DISTANCE, SOPPOp.S_DELAY_ALU): inst._dispatch = dispatch_nop
|
||||
elif isinstance(inst, (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM)): inst._dispatch = exec_scalar
|
||||
elif isinstance(inst, VOP1) and inst.op == VOP1Op.V_NOP: inst._dispatch = dispatch_nop
|
||||
elif isinstance(inst, VOP3P) and 'WMMA' in inst.op_name: inst._dispatch = dispatch_wmma
|
||||
elif isinstance(inst, VOP3) and inst.op == VOP3Op.V_WRITELANE_B32: inst._dispatch = dispatch_writelane
|
||||
elif isinstance(inst, (VOP1, VOP3)) and inst.op in (VOP1Op.V_READFIRSTLANE_B32, VOP3Op.V_READFIRSTLANE_B32, VOP3Op.V_READLANE_B32): inst._dispatch = dispatch_readlane
|
||||
elif isinstance(inst, VOPD): inst._dispatch = dispatch_lane(exec_vopd)
|
||||
elif isinstance(inst, FLAT): inst._dispatch = dispatch_lane(exec_flat)
|
||||
elif isinstance(inst, DS): inst._dispatch = dispatch_lane(exec_ds)
|
||||
else: inst._dispatch = dispatch_lane(exec_vop)
|
||||
|
||||
# Compile pcode for instructions that use it (not VOPD which has _fnx/_fny, not special dispatches)
|
||||
# VOPD needs separate functions for X and Y ops
|
||||
if isinstance(inst, VOPD):
|
||||
def _compile_vopd_op(op): return compile_pseudocode(type(op).__name__, op.name, PSEUDOCODE_STRINGS[type(op)][op])
|
||||
inst._fnx, inst._fny = _compile_vopd_op(_VOPD_TO_VOP[inst.opx]), _compile_vopd_op(_VOPD_TO_VOP[inst.opy])
|
||||
elif inst._dispatch not in (dispatch_endpgm, dispatch_barrier, dispatch_nop, dispatch_wmma, dispatch_writelane):
|
||||
assert type(inst.op) != int, f"inst op of {inst} is int"
|
||||
inst._fn = compile_pseudocode(type(inst.op).__name__, inst.op.name, PSEUDOCODE_STRINGS[type(inst.op)][inst.op])
|
||||
result[i // 4] = inst
|
||||
i += inst._words * 4
|
||||
return result
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# MAIN EXECUTION LOOP
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def exec_wave(program: dict[int, Inst], st: WaveState) -> int:
|
||||
while (inst := program.get(st.pc)) and (result := inst._dispatch(st, inst)) == 0: pass
|
||||
return result
|
||||
|
||||
def exec_workgroup(program: dict[int, Inst], workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int, rsrc2: int) -> None:
|
||||
lx, ly, lz = local_size
|
||||
total_threads = lx * ly * lz
|
||||
# GRANULATED_LDS_SIZE is in 512-byte units (see ops_amd.py: lds_size = ((group_segment_size + 511) // 512))
|
||||
lds_size = ((rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE_SHIFT) * 512
|
||||
lds = LDSMem(bytearray(lds_size)) if lds_size else None
|
||||
waves: list[WaveState] = []
|
||||
for wave_start in range(0, total_threads, WAVE_SIZE):
|
||||
n_lanes = min(WAVE_SIZE, total_threads - wave_start)
|
||||
st = WaveState(lds, n_lanes)
|
||||
st.exec_mask = (1 << n_lanes) - 1
|
||||
st.wsgpr64(0, args_ptr) # s[0:1] = kernel arguments pointer
|
||||
# COMPUTE_PGM_RSRC2: USER_SGPR_COUNT is where workgroup IDs start, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z control which are passed
|
||||
sgpr_idx = (rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT_SHIFT
|
||||
if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X: st.sgpr[sgpr_idx] = workgroup_id[0]; sgpr_idx += 1
|
||||
if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Y: st.sgpr[sgpr_idx] = workgroup_id[1]; sgpr_idx += 1
|
||||
if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Z: st.sgpr[sgpr_idx] = workgroup_id[2]
|
||||
# VGPR0 = packed workitem IDs: (Z << 20) | (Y << 10) | X
|
||||
for tid in range(wave_start, wave_start + n_lanes):
|
||||
st.vgpr[tid - wave_start][0] = ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
|
||||
waves.append(st)
|
||||
while waves:
|
||||
waves = [st for st in waves if exec_wave(program, st) != -1]
|
||||
|
||||
def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int, rsrc2: int = 0x19c) -> int:
|
||||
program = decode_program((ctypes.c_char * lib_sz).from_address(lib).raw)
|
||||
for gidz in range(gz):
|
||||
for gidy in range(gy):
|
||||
for gidx in range(gx): exec_workgroup(program, (gidx, gidy, gidz), (lx, ly, lz), args_ptr, rsrc2)
|
||||
return 0
|
||||
@@ -0,0 +1,821 @@
|
||||
# DSL for RDNA3 pseudocode - makes pseudocode expressions work directly as Python
|
||||
import struct, math, re, functools
|
||||
from extra.assembly.amd.dsl import MASK32, MASK64
|
||||
|
||||
# Float/int bit conversion functions
|
||||
_struct_f, _struct_I = struct.Struct("<f"), struct.Struct("<I")
|
||||
_struct_e, _struct_H = struct.Struct("<e"), struct.Struct("<H")
|
||||
_struct_d, _struct_Q = struct.Struct("<d"), struct.Struct("<Q")
|
||||
def _f32(i):
|
||||
i = i & MASK32
|
||||
# RDNA3 default mode: flush f32 denormals to zero (FTZ)
|
||||
# Denormal: exponent=0 (bits 23-30) and mantissa!=0 (bits 0-22)
|
||||
if (i & 0x7f800000) == 0 and (i & 0x007fffff) != 0: return 0.0
|
||||
return _struct_f.unpack(_struct_I.pack(i))[0]
|
||||
def _i32(f):
|
||||
if isinstance(f, int): f = float(f)
|
||||
if math.isnan(f): return 0xffc00000 if math.copysign(1.0, f) < 0 else 0x7fc00000
|
||||
if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
|
||||
try:
|
||||
bits = _struct_I.unpack(_struct_f.pack(f))[0]
|
||||
# RDNA3 default mode: flush f32 denormals to zero (FTZ)
|
||||
if (bits & 0x7f800000) == 0 and (bits & 0x007fffff) != 0: return 0x80000000 if bits & 0x80000000 else 0
|
||||
return bits
|
||||
except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
|
||||
def _sext(v, b): return v - (1 << b) if v & (1 << (b - 1)) else v
|
||||
def _f16(i): return _struct_e.unpack(_struct_H.pack(i & 0xffff))[0]
|
||||
def _i16(f):
|
||||
if math.isnan(f): return 0x7e00
|
||||
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
|
||||
try: return _struct_H.unpack(_struct_e.pack(f))[0]
|
||||
except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
|
||||
def _f64(i): return _struct_d.unpack(_struct_Q.pack(i & MASK64))[0]
|
||||
def _i64(f):
|
||||
if math.isnan(f): return 0x7ff8000000000000
|
||||
if math.isinf(f): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
|
||||
try: return _struct_Q.unpack(_struct_d.pack(f))[0]
|
||||
except (OverflowError, struct.error): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# INTERNAL HELPERS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _div(a, b):
|
||||
try: return a / b
|
||||
except ZeroDivisionError:
|
||||
if a == 0.0 or math.isnan(a): return float("nan")
|
||||
return math.copysign(float("inf"), a * b) if b == 0.0 else float("inf") if a > 0 else float("-inf")
|
||||
def _check_nan_type(x, quiet_bit_expected, default):
|
||||
try:
|
||||
if not math.isnan(float(x)): return False
|
||||
if hasattr(x, '_reg') and hasattr(x, '_bits'):
|
||||
bits = x._reg._val & ((1 << x._bits) - 1)
|
||||
exp_bits, quiet_pos, mant_mask = {16: (0x1f, 9, 0x3ff), 32: (0xff, 22, 0x7fffff), 64: (0x7ff, 51, 0xfffffffffffff)}.get(x._bits, (0,0,0))
|
||||
exp_shift = {16: 10, 32: 23, 64: 52}.get(x._bits, 0)
|
||||
if exp_bits and ((bits >> exp_shift) & exp_bits) == exp_bits and (bits & mant_mask) != 0:
|
||||
return ((bits >> quiet_pos) & 1) == quiet_bit_expected
|
||||
return default
|
||||
except (TypeError, ValueError): return False
|
||||
def _gt_neg_zero(a, b): return (a > b) or (a == 0 and b == 0 and not math.copysign(1, a) < 0 and math.copysign(1, b) < 0)
|
||||
def _lt_neg_zero(a, b): return (a < b) or (a == 0 and b == 0 and math.copysign(1, a) < 0 and not math.copysign(1, b) < 0)
|
||||
def _fpop(fn):
|
||||
def wrapper(x):
|
||||
x = float(x)
|
||||
if math.isnan(x) or math.isinf(x): return x
|
||||
result = float(fn(x))
|
||||
return math.copysign(0.0, x) if result == 0.0 else result
|
||||
return wrapper
|
||||
def _f_to_int(f, lo, hi): f = float(f); return 0 if math.isnan(f) else (hi if f >= hi else lo if f <= lo else int(f))
|
||||
def _f16_to_f32_bits(bits): return struct.unpack("<e", struct.pack("<H", int(bits) & 0xffff))[0]
|
||||
def _brev(v, bits): return int(bin(v & ((1 << bits) - 1))[2:].zfill(bits)[::-1], 2)
|
||||
def _ctz(v, bits):
|
||||
v, n = int(v) & ((1 << bits) - 1), 0
|
||||
if v == 0: return bits
|
||||
while (v & 1) == 0: v >>= 1; n += 1
|
||||
return n
|
||||
|
||||
def _bf16(i):
|
||||
"""Convert bf16 bits to float. BF16 is just the top 16 bits of f32."""
|
||||
return struct.unpack("<f", struct.pack("<I", (i & 0xffff) << 16))[0]
|
||||
def _ibf16(f):
|
||||
"""Convert float to bf16 bits (truncate to top 16 bits of f32)."""
|
||||
if math.isnan(f): return 0x7fc0 # bf16 quiet NaN
|
||||
if math.isinf(f): return 0x7f80 if f > 0 else 0xff80 # bf16 ±infinity
|
||||
try: return (struct.unpack("<I", struct.pack("<f", float(f)))[0] >> 16) & 0xffff
|
||||
except (OverflowError, struct.error): return 0x7f80 if f > 0 else 0xff80
|
||||
def _trig(fn, x):
|
||||
# V_SIN/COS_F32: hardware does frac on input cycles before computing
|
||||
if math.isinf(x) or math.isnan(x): return float("nan")
|
||||
frac_cycles = fract(x / (2 * math.pi))
|
||||
result = fn(frac_cycles * 2 * math.pi)
|
||||
# Hardware returns exactly 0 for cos(π/2), sin(π), etc. due to lookup table
|
||||
# Round very small results (below f32 precision) to exactly 0
|
||||
if abs(result) < 1e-7: return 0.0
|
||||
return result
|
||||
|
||||
class _SafeFloat(float):
|
||||
"""Float subclass that uses _div for division to handle 0/inf correctly."""
|
||||
def __truediv__(self, o): return _div(float(self), float(o))
|
||||
def __rtruediv__(self, o): return _div(float(o), float(self))
|
||||
|
||||
class _Inf:
|
||||
f16 = f32 = f64 = float('inf')
|
||||
def __neg__(self): return _NegInf()
|
||||
def __pos__(self): return self
|
||||
def __float__(self): return float('inf')
|
||||
def __eq__(self, other): return float(other) == float('inf') if not isinstance(other, _NegInf) else False
|
||||
def __req__(self, other): return self.__eq__(other)
|
||||
class _NegInf:
|
||||
f16 = f32 = f64 = float('-inf')
|
||||
def __neg__(self): return _Inf()
|
||||
def __pos__(self): return self
|
||||
def __float__(self): return float('-inf')
|
||||
def __eq__(self, other): return float(other) == float('-inf') if not isinstance(other, _Inf) else False
|
||||
def __req__(self, other): return self.__eq__(other)
|
||||
|
||||
class _RoundMode:
|
||||
NEAREST_EVEN = 0
|
||||
|
||||
class _WaveMode:
|
||||
IEEE = False
|
||||
|
||||
class _DenormChecker:
|
||||
"""Comparator for denormalized floats. x == DENORM.f32 checks if x is denormalized."""
|
||||
def __init__(self, bits): self._bits = bits
|
||||
def _check(self, other):
|
||||
f = float(other)
|
||||
if math.isinf(f) or math.isnan(f) or f == 0.0: return False
|
||||
if self._bits == 64:
|
||||
bits = struct.unpack("<Q", struct.pack("<d", f))[0]
|
||||
return (bits >> 52) & 0x7ff == 0
|
||||
bits = struct.unpack("<I", struct.pack("<f", f))[0]
|
||||
return (bits >> 23) & 0xff == 0
|
||||
def __eq__(self, other): return self._check(other)
|
||||
def __req__(self, other): return self._check(other)
|
||||
def __ne__(self, other): return not self._check(other)
|
||||
|
||||
class _Denorm:
|
||||
f32 = _DenormChecker(32)
|
||||
f64 = _DenormChecker(64)
|
||||
|
||||
_pack = lambda hi, lo: ((int(hi) & 0xffff) << 16) | (int(lo) & 0xffff)
|
||||
_pack32 = lambda hi, lo: ((int(hi) & 0xffffffff) << 32) | (int(lo) & 0xffffffff)
|
||||
|
||||
class TypedView:
|
||||
"""View into a Reg with typed access. Used for both full-width (Reg.u32) and slices (Reg[31:16])."""
|
||||
__slots__ = ('_reg', '_high', '_low', '_signed', '_float', '_bf16', '_reversed')
|
||||
def __init__(self, reg, high, low=0, signed=False, is_float=False, is_bf16=False):
|
||||
# Handle reversed slices like [0:31] which means bit-reverse
|
||||
if high < low: high, low, reversed = low, high, True
|
||||
else: reversed = False
|
||||
self._reg, self._high, self._low, self._reversed = reg, high, low, reversed
|
||||
self._signed, self._float, self._bf16 = signed, is_float, is_bf16
|
||||
|
||||
def _nbits(self): return self._high - self._low + 1
|
||||
def _mask(self): return (1 << self._nbits()) - 1
|
||||
def _get(self):
|
||||
v = (self._reg._val >> self._low) & self._mask()
|
||||
return _brev(v, self._nbits()) if self._reversed else v
|
||||
def _set(self, v):
|
||||
v = int(v)
|
||||
if self._reversed: v = _brev(v, self._nbits())
|
||||
self._reg._val = (self._reg._val & ~(self._mask() << self._low)) | ((v & self._mask()) << self._low)
|
||||
|
||||
@property
|
||||
def _val(self): return self._get()
|
||||
@property
|
||||
def _bits(self): return self._nbits()
|
||||
|
||||
# Type accessors for slices (e.g., D0[31:16].f16)
|
||||
u8 = property(lambda s: s._get() & 0xff)
|
||||
u16 = property(lambda s: s._get() & 0xffff, lambda s, v: s._set(v))
|
||||
u32 = property(lambda s: s._get() & MASK32, lambda s, v: s._set(v))
|
||||
i16 = property(lambda s: _sext(s._get() & 0xffff, 16), lambda s, v: s._set(v))
|
||||
i32 = property(lambda s: _sext(s._get() & MASK32, 32), lambda s, v: s._set(v))
|
||||
f16 = property(lambda s: _f16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _i16(float(v))))
|
||||
f32 = property(lambda s: _f32(s._get()), lambda s, v: s._set(_i32(float(v))))
|
||||
bf16 = property(lambda s: _bf16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _ibf16(float(v))))
|
||||
b16, b32 = u16, u32
|
||||
|
||||
# Chained type access (e.g., jump_addr.i64 when jump_addr is already TypedView)
|
||||
@property
|
||||
def i64(s): return s if s._nbits() == 64 and s._signed else int(s)
|
||||
@property
|
||||
def u64(s): return s if s._nbits() == 64 and not s._signed else int(s) & MASK64
|
||||
|
||||
def __getitem__(self, key):
|
||||
if isinstance(key, slice):
|
||||
high, low = int(key.start), int(key.stop)
|
||||
return TypedView(self._reg, high, low)
|
||||
return (self._get() >> int(key)) & 1
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
if isinstance(key, slice):
|
||||
high, low = int(key.start), int(key.stop)
|
||||
if high < low: high, low, value = low, high, _brev(int(value), low - high + 1)
|
||||
mask = (1 << (high - low + 1)) - 1
|
||||
self._reg._val = (self._reg._val & ~(mask << low)) | ((int(value) & mask) << low)
|
||||
elif value: self._reg._val |= (1 << int(key))
|
||||
else: self._reg._val &= ~(1 << int(key))
|
||||
|
||||
def __int__(self): return _sext(self._get(), self._nbits()) if self._signed else self._get()
|
||||
def __index__(self): return int(self)
|
||||
def __trunc__(self): return int(float(self)) if self._float else int(self)
|
||||
def __float__(self):
|
||||
if self._float:
|
||||
if self._bf16: return _bf16(self._get())
|
||||
bits = self._nbits()
|
||||
return _f16(self._get()) if bits == 16 else _f32(self._get()) if bits == 32 else _f64(self._get())
|
||||
return float(int(self))
|
||||
def __bool__(s): return bool(int(s))
|
||||
|
||||
# Arithmetic - floats use float(), ints use int()
|
||||
def __add__(s, o): return float(s) + float(o) if s._float else int(s) + int(o)
|
||||
def __radd__(s, o): return float(o) + float(s) if s._float else int(o) + int(s)
|
||||
def __sub__(s, o): return float(s) - float(o) if s._float else int(s) - int(o)
|
||||
def __rsub__(s, o): return float(o) - float(s) if s._float else int(o) - int(s)
|
||||
def __mul__(s, o): return float(s) * float(o) if s._float else int(s) * int(o)
|
||||
def __rmul__(s, o): return float(o) * float(s) if s._float else int(o) * int(s)
|
||||
def __truediv__(s, o): return _div(float(s), float(o)) if s._float else _div(int(s), int(o))
|
||||
def __rtruediv__(s, o): return _div(float(o), float(s)) if s._float else _div(int(o), int(s))
|
||||
def __pow__(s, o): return float(s) ** float(o) if s._float else int(s) ** int(o)
|
||||
def __rpow__(s, o): return float(o) ** float(s) if s._float else int(o) ** int(s)
|
||||
def __neg__(s): return -float(s) if s._float else -int(s)
|
||||
def __abs__(s): return abs(float(s)) if s._float else abs(int(s))
|
||||
|
||||
# Bitwise - GPU shifts mask the shift amount to valid range
|
||||
def __and__(s, o): return int(s) & int(o)
|
||||
def __or__(s, o): return int(s) | int(o)
|
||||
def __xor__(s, o): return int(s) ^ int(o)
|
||||
def __invert__(s): return ~int(s)
|
||||
def __lshift__(s, o): n = int(o); return int(s) << n if 0 <= n < 64 or s._nbits() > 64 else 0
|
||||
def __rshift__(s, o): n = int(o); return int(s) >> n if 0 <= n < 64 or s._nbits() > 64 else 0
|
||||
def __rand__(s, o): return int(o) & int(s)
|
||||
def __ror__(s, o): return int(o) | int(s)
|
||||
def __rxor__(s, o): return int(o) ^ int(s)
|
||||
def __rlshift__(s, o): n = int(s); return int(o) << n if 0 <= n < 64 else 0
|
||||
def __rrshift__(s, o): n = int(s); return int(o) >> n if 0 <= n < 64 else 0
|
||||
|
||||
# Comparison - handle _DenormChecker specially
|
||||
def __eq__(s, o):
|
||||
if isinstance(o, _DenormChecker): return o._check(s)
|
||||
return float(s) == float(o) if s._float else int(s) == int(o)
|
||||
def __ne__(s, o):
|
||||
if isinstance(o, _DenormChecker): return not o._check(s)
|
||||
return float(s) != float(o) if s._float else int(s) != int(o)
|
||||
def __lt__(s, o): return float(s) < float(o) if s._float else int(s) < int(o)
|
||||
def __le__(s, o): return float(s) <= float(o) if s._float else int(s) <= int(o)
|
||||
def __gt__(s, o): return float(s) > float(o) if s._float else int(s) > int(o)
|
||||
def __ge__(s, o): return float(s) >= float(o) if s._float else int(s) >= int(o)
|
||||
|
||||
class Reg:
|
||||
"""GPU register: D0.f32 = S0.f32 + S1.f32 just works. Supports up to 128 bits for DS_LOAD_B128."""
|
||||
__slots__ = ('_val',)
|
||||
def __init__(self, val=0): self._val = int(val)
|
||||
|
||||
# Typed views - TypedView(reg, high, signed, is_float, is_bf16)
|
||||
u64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
|
||||
i64 = property(lambda s: TypedView(s, 63, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK64))
|
||||
b64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
|
||||
f64 = property(lambda s: TypedView(s, 63, is_float=True), lambda s, v: setattr(s, '_val', v if isinstance(v, int) else _i64(float(v))))
|
||||
u32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
|
||||
i32 = property(lambda s: TypedView(s, 31, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK32))
|
||||
b32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
|
||||
f32 = property(lambda s: TypedView(s, 31, is_float=True), lambda s, v: setattr(s, '_val', _i32(float(v))))
|
||||
u24 = property(lambda s: TypedView(s, 23))
|
||||
i24 = property(lambda s: TypedView(s, 23, signed=True))
|
||||
u16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
|
||||
i16 = property(lambda s: TypedView(s, 15, signed=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
|
||||
b16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
|
||||
f16 = property(lambda s: TypedView(s, 15, is_float=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _i16(float(v))) & 0xffff)))
|
||||
bf16 = property(lambda s: TypedView(s, 15, is_float=True, is_bf16=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _ibf16(float(v))) & 0xffff)))
|
||||
u8 = property(lambda s: TypedView(s, 7))
|
||||
i8 = property(lambda s: TypedView(s, 7, signed=True))
|
||||
u3 = property(lambda s: TypedView(s, 2)) # 3-bit for opsel fields
|
||||
u1 = property(lambda s: TypedView(s, 0)) # single bit
|
||||
|
||||
def __getitem__(s, key):
|
||||
if isinstance(key, slice): return TypedView(s, int(key.start), int(key.stop))
|
||||
return (s._val >> int(key)) & 1
|
||||
|
||||
def __setitem__(s, key, value):
|
||||
if isinstance(key, slice):
|
||||
high, low = int(key.start), int(key.stop)
|
||||
if high < low: high, low = low, high
|
||||
mask = (1 << (high - low + 1)) - 1
|
||||
s._val = (s._val & ~(mask << low)) | ((int(value) & mask) << low)
|
||||
elif value: s._val |= (1 << int(key))
|
||||
else: s._val &= ~(1 << int(key))
|
||||
|
||||
def __int__(s): return s._val
|
||||
def __index__(s): return s._val
|
||||
def __bool__(s): return bool(s._val)
|
||||
|
||||
# Arithmetic (for tmp = tmp + 1 patterns). Float operands trigger f32 interpretation.
|
||||
def __add__(s, o): return (_f32(s._val) + float(o)) if isinstance(o, float) else s._val + int(o)
|
||||
def __radd__(s, o): return (float(o) + _f32(s._val)) if isinstance(o, float) else int(o) + s._val
|
||||
def __sub__(s, o): return (_f32(s._val) - float(o)) if isinstance(o, float) else s._val - int(o)
|
||||
def __rsub__(s, o): return (float(o) - _f32(s._val)) if isinstance(o, float) else int(o) - s._val
|
||||
def __mul__(s, o): return (_f32(s._val) * float(o)) if isinstance(o, float) else s._val * int(o)
|
||||
def __rmul__(s, o): return (float(o) * _f32(s._val)) if isinstance(o, float) else int(o) * s._val
|
||||
def __and__(s, o): return s._val & int(o)
|
||||
def __rand__(s, o): return int(o) & s._val
|
||||
def __or__(s, o): return s._val | int(o)
|
||||
def __ror__(s, o): return int(o) | s._val
|
||||
def __xor__(s, o): return s._val ^ int(o)
|
||||
def __rxor__(s, o): return int(o) ^ s._val
|
||||
def __lshift__(s, o): n = int(o); return s._val << n if 0 <= n < 64 else 0
|
||||
def __rshift__(s, o): n = int(o); return s._val >> n if 0 <= n < 64 else 0
|
||||
def __invert__(s): return ~s._val
|
||||
|
||||
# Comparison (for tmp >= 0x100000000 patterns)
|
||||
def __lt__(s, o): return s._val < int(o)
|
||||
def __le__(s, o): return s._val <= int(o)
|
||||
def __gt__(s, o): return s._val > int(o)
|
||||
def __ge__(s, o): return s._val >= int(o)
|
||||
def __eq__(s, o): return s._val == int(o)
|
||||
def __ne__(s, o): return s._val != int(o)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PSEUDOCODE API - Functions and constants from AMD ISA pseudocode
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
# Rounding and float operations
|
||||
trunc, floor, ceil = _fpop(math.trunc), _fpop(math.floor), _fpop(math.ceil)
|
||||
def sqrt(x): return _SafeFloat(math.sqrt(x)) if x >= 0 else _SafeFloat(float("nan"))
|
||||
def log2(x): return math.log2(x) if x > 0 else (float("-inf") if x == 0 else float("nan"))
|
||||
def fract(x): return x - math.floor(x)
|
||||
def sin(x): return _trig(math.sin, x)
|
||||
def cos(x): return _trig(math.cos, x)
|
||||
def pow(a, b):
|
||||
try: return a ** b
|
||||
except OverflowError: return float("inf") if b > 0 else 0.0
|
||||
def isEven(x):
|
||||
x = float(x)
|
||||
if math.isinf(x) or math.isnan(x): return False
|
||||
return int(x) % 2 == 0
|
||||
def mantissa(f):
|
||||
if f == 0.0 or math.isinf(f) or math.isnan(f): return f
|
||||
m, _ = math.frexp(f)
|
||||
return m # AMD V_FREXP_MANT returns mantissa in [0.5, 1.0) range
|
||||
def signext_from_bit(val, bit):
|
||||
bit = int(bit)
|
||||
if bit == 0: return 0
|
||||
mask = (1 << bit) - 1
|
||||
val = int(val) & mask
|
||||
if val & (1 << (bit - 1)): return val - (1 << bit)
|
||||
return val
|
||||
|
||||
# Type conversions
|
||||
i32_to_f32 = u32_to_f32 = i32_to_f64 = u32_to_f64 = f32_to_f64 = f64_to_f32 = float
|
||||
def f32_to_i32(f): return _f_to_int(f, -2147483648, 2147483647)
|
||||
def f32_to_u32(f): return _f_to_int(f, 0, 4294967295)
|
||||
f64_to_i32, f64_to_u32 = f32_to_i32, f32_to_u32
|
||||
def f32_to_f16(f):
|
||||
f = float(f)
|
||||
if math.isnan(f): return 0x7e00 # f16 NaN
|
||||
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00 # f16 ±infinity
|
||||
try: return struct.unpack("<H", struct.pack("<e", f))[0]
|
||||
except OverflowError: return 0x7c00 if f > 0 else 0xfc00 # overflow -> ±infinity
|
||||
def f16_to_f32(v): return v if isinstance(v, float) else _f16_to_f32_bits(v)
|
||||
def i16_to_f16(v): return f32_to_f16(float(_sext(int(v) & 0xffff, 16)))
|
||||
def u16_to_f16(v): return f32_to_f16(float(int(v) & 0xffff))
|
||||
def f16_to_i16(bits): f = _f16_to_f32_bits(bits); return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
|
||||
def f16_to_u16(bits): f = _f16_to_f32_bits(bits); return max(0, min(65535, int(f))) if not math.isnan(f) else 0
|
||||
def bf16_to_f32(v): return _bf16(v) if isinstance(v, int) else float(v)
|
||||
def f32_to_bf16(f): return _ibf16(f)
|
||||
def u8_to_u32(v): return int(v) & 0xff
|
||||
def u4_to_u32(v): return int(v) & 0xf
|
||||
def u32_to_u16(u): return int(u) & 0xffff
|
||||
def i32_to_i16(i): return ((int(i) + 32768) & 0xffff) - 32768
|
||||
def f16_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
|
||||
def f16_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
|
||||
def f32_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
|
||||
def f32_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
|
||||
def v_cvt_i16_f32(f): return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
|
||||
def v_cvt_u16_f32(f): return max(0, min(65535, int(f))) if not math.isnan(f) else 0
|
||||
def SAT8(v): return max(0, min(255, int(v)))
|
||||
def f32_to_u8(f): return max(0, min(255, int(f))) if not math.isnan(f) else 0
|
||||
|
||||
# Min/max operations
|
||||
def v_min_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _lt_neg_zero(a, b) else b)
|
||||
def v_max_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _gt_neg_zero(a, b) else b)
|
||||
v_min_f16, v_max_f16 = v_min_f32, v_max_f32
|
||||
v_min_i32, v_max_i32 = min, max
|
||||
v_min_i16, v_max_i16 = min, max
|
||||
def v_min_u32(a, b): return min(a & MASK32, b & MASK32)
|
||||
def v_max_u32(a, b): return max(a & MASK32, b & MASK32)
|
||||
def v_min_u16(a, b): return min(a & 0xffff, b & 0xffff)
|
||||
def v_max_u16(a, b): return max(a & 0xffff, b & 0xffff)
|
||||
def v_min3_f32(a, b, c): return v_min_f32(v_min_f32(a, b), c)
|
||||
def v_max3_f32(a, b, c): return v_max_f32(v_max_f32(a, b), c)
|
||||
v_min3_f16, v_max3_f16 = v_min3_f32, v_max3_f32
|
||||
v_min3_i32, v_max3_i32, v_min3_i16, v_max3_i16 = min, max, min, max
|
||||
def v_min3_u32(a, b, c): return min(a & MASK32, b & MASK32, c & MASK32)
|
||||
def v_max3_u32(a, b, c): return max(a & MASK32, b & MASK32, c & MASK32)
|
||||
def v_min3_u16(a, b, c): return min(a & 0xffff, b & 0xffff, c & 0xffff)
|
||||
def v_max3_u16(a, b, c): return max(a & 0xffff, b & 0xffff, c & 0xffff)
|
||||
|
||||
# SAD/MSAD operations
|
||||
def ABSDIFF(a, b): return abs(int(a) - int(b))
|
||||
def v_sad_u8(s0, s1, s2):
|
||||
"""V_SAD_U8: Sum of absolute differences of 4 byte pairs plus accumulator."""
|
||||
s0, s1, s2 = int(s0), int(s1), int(s2)
|
||||
result = s2
|
||||
for i in range(4):
|
||||
a = (s0 >> (i * 8)) & 0xff
|
||||
b = (s1 >> (i * 8)) & 0xff
|
||||
result += abs(a - b)
|
||||
return result & 0xffffffff
|
||||
def v_msad_u8(s0, s1, s2):
|
||||
"""V_MSAD_U8: Masked sum of absolute differences (skip if reference byte is 0)."""
|
||||
s0, s1, s2 = int(s0), int(s1), int(s2)
|
||||
result = s2
|
||||
for i in range(4):
|
||||
a = (s0 >> (i * 8)) & 0xff
|
||||
b = (s1 >> (i * 8)) & 0xff
|
||||
if b != 0: # Only add diff if reference (s1) byte is non-zero
|
||||
result += abs(a - b)
|
||||
return result & 0xffffffff
|
||||
|
||||
def BYTE_PERMUTE(data, sel):
|
||||
"""Select a byte from 64-bit data based on selector value."""
|
||||
sel = int(sel) & 0xff
|
||||
if sel <= 7: return (int(data) >> (sel * 8)) & 0xff
|
||||
if sel == 8: return 0xff if ((int(data) >> 15) & 1) else 0x00
|
||||
if sel == 9: return 0xff if ((int(data) >> 31) & 1) else 0x00
|
||||
if sel == 10: return 0xff if ((int(data) >> 47) & 1) else 0x00
|
||||
if sel == 11: return 0xff if ((int(data) >> 63) & 1) else 0x00
|
||||
if sel == 12: return 0x00
|
||||
return 0xff
|
||||
|
||||
# Pseudocode functions
|
||||
def s_ff1_i32_b32(v): return _ctz(v, 32)
|
||||
def s_ff1_i32_b64(v): return _ctz(v, 64)
|
||||
GT_NEG_ZERO, LT_NEG_ZERO = _gt_neg_zero, _lt_neg_zero
|
||||
def isNAN(x):
|
||||
try: return math.isnan(float(x))
|
||||
except (TypeError, ValueError): return False
|
||||
def isQuietNAN(x): return _check_nan_type(x, 1, True)
|
||||
def isSignalNAN(x): return _check_nan_type(x, 0, False)
|
||||
def fma(a, b, c):
|
||||
try: return math.fma(a, b, c)
|
||||
except ValueError: return float('nan')
|
||||
def ldexp(m, e): return math.ldexp(m, e)
|
||||
def sign(f): return 1 if math.copysign(1.0, f) < 0 else 0
|
||||
def exponent(f):
|
||||
if hasattr(f, '_bits') and hasattr(f, '_float') and f._float:
|
||||
raw = f._val
|
||||
if f._bits == 16: return (raw >> 10) & 0x1f
|
||||
if f._bits == 32: return (raw >> 23) & 0xff
|
||||
if f._bits == 64: return (raw >> 52) & 0x7ff
|
||||
f = float(f)
|
||||
if math.isinf(f) or math.isnan(f): return 255
|
||||
if f == 0.0: return 0
|
||||
try: bits = struct.unpack("<I", struct.pack("<f", f))[0]; return (bits >> 23) & 0xff
|
||||
except: return 0
|
||||
def signext(x): return int(x)
|
||||
def cvtToQuietNAN(x): return float('nan')
|
||||
|
||||
def F(x):
|
||||
"""32'F(x) or 64'F(x) - interpret x as float. If x is int, treat as bit pattern."""
|
||||
if isinstance(x, int): return _f32(x)
|
||||
if isinstance(x, TypedView): return x
|
||||
return float(x)
|
||||
|
||||
# Constants
|
||||
PI = math.pi
|
||||
WAVE32, WAVE64 = True, False
|
||||
OVERFLOW_F32, UNDERFLOW_F32 = float('inf'), 0.0
|
||||
OVERFLOW_F64, UNDERFLOW_F64 = float('inf'), 0.0
|
||||
MAX_FLOAT_F32 = 3.4028235e+38
|
||||
INF = _Inf()
|
||||
ROUND_MODE = _RoundMode()
|
||||
WAVE_MODE = _WaveMode()
|
||||
DENORM = _Denorm()
|
||||
|
||||
# 2/PI with 1201 bits of precision for V_TRIG_PREOP_F64
|
||||
TWO_OVER_PI_1201 = Reg(0x0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# COMPILER: pseudocode -> Python (minimal transforms)
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _filter_pseudocode(pseudocode: str) -> str:
|
||||
"""Filter raw PDF pseudocode to only include actual code lines."""
|
||||
pcode_lines, in_lambda, depth = [], 0, 0
|
||||
for line in pseudocode.split('\n'):
|
||||
s = line.strip()
|
||||
if not s: continue
|
||||
if '=>' in s or re.match(r'^[A-Z_]+\(', s): continue # Skip example lines
|
||||
if '= lambda(' in s: in_lambda += 1; continue # Skip lambda definitions
|
||||
if in_lambda > 0:
|
||||
if s.endswith(');'): in_lambda -= 1
|
||||
continue
|
||||
# Only include lines that look like pseudocode
|
||||
is_code = (any(p in s for p in ['D0.', 'D1.', 'S0.', 'S1.', 'S2.', 'SCC =', 'SCC ?', 'VCC', 'EXEC', 'tmp =', 'tmp[', 'lane =', 'PC =',
|
||||
'D0[', 'D1[', 'S0[', 'S1[', 'S2[', 'MEM[', 'RETURN_DATA', 'VADDR', 'VDATA', 'VDST', 'SADDR', 'OFFSET']) or
|
||||
s.startswith(('if ', 'else', 'elsif', 'endif', 'declare ', 'for ', 'endfor', '//')) or
|
||||
re.match(r'^[a-z_]+\s*=', s) or re.match(r'^[a-z_]+\[', s) or (depth > 0 and '=' in s))
|
||||
if s.startswith('if '): depth += 1
|
||||
elif s.startswith('endif'): depth = max(0, depth - 1)
|
||||
if is_code: pcode_lines.append(s)
|
||||
return '\n'.join(pcode_lines)
|
||||
|
||||
def _compile_pseudocode(pseudocode: str) -> str:
|
||||
"""Compile pseudocode to Python. Transforms are minimal - most syntax just works."""
|
||||
pseudocode = re.sub(r'\bpass\b', 'pass_', pseudocode) # 'pass' is Python keyword
|
||||
raw_lines = pseudocode.strip().split('\n')
|
||||
joined_lines: list[str] = []
|
||||
for line in raw_lines:
|
||||
line = line.strip()
|
||||
if joined_lines and (joined_lines[-1].rstrip().endswith(('||', '&&', '(', ',')) or
|
||||
(joined_lines[-1].count('(') > joined_lines[-1].count(')'))):
|
||||
joined_lines[-1] = joined_lines[-1].rstrip() + ' ' + line
|
||||
else:
|
||||
joined_lines.append(line)
|
||||
|
||||
lines = []
|
||||
indent, need_pass, in_first_match_loop = 0, False, False
|
||||
for line in joined_lines:
|
||||
line = line.split('//')[0].strip() # Strip C-style comments
|
||||
if not line: continue
|
||||
if line.startswith('if '):
|
||||
lines.append(' ' * indent + f"if {_expr(line[3:].rstrip(' then'))}:")
|
||||
indent += 1
|
||||
need_pass = True
|
||||
elif line.startswith('elsif '):
|
||||
if need_pass: lines.append(' ' * indent + "pass")
|
||||
indent -= 1
|
||||
lines.append(' ' * indent + f"elif {_expr(line[6:].rstrip(' then'))}:")
|
||||
indent += 1
|
||||
need_pass = True
|
||||
elif line == 'else':
|
||||
if need_pass: lines.append(' ' * indent + "pass")
|
||||
indent -= 1
|
||||
lines.append(' ' * indent + "else:")
|
||||
indent += 1
|
||||
need_pass = True
|
||||
elif line.startswith('endif'):
|
||||
if need_pass: lines.append(' ' * indent + "pass")
|
||||
indent -= 1
|
||||
need_pass = False
|
||||
elif line.startswith('endfor'):
|
||||
if need_pass: lines.append(' ' * indent + "pass")
|
||||
indent -= 1
|
||||
need_pass, in_first_match_loop = False, False
|
||||
elif line.startswith('declare '):
|
||||
pass
|
||||
elif m := re.match(r'for (\w+) in (.+?)\s*:\s*(.+?) do', line):
|
||||
start, end = _expr(m[2].strip()), _expr(m[3].strip())
|
||||
lines.append(' ' * indent + f"for {m[1]} in range({start}, int({end})+1):")
|
||||
indent += 1
|
||||
need_pass, in_first_match_loop = True, True
|
||||
elif '=' in line and not line.startswith('=='):
|
||||
need_pass = False
|
||||
line = line.rstrip(';')
|
||||
if m := re.match(r'\{\s*D1\.[ui]1\s*,\s*D0\.[ui]64\s*\}\s*=\s*(.+)', line):
|
||||
rhs = _expr(m[1])
|
||||
lines.append(' ' * indent + f"_full = {rhs}")
|
||||
lines.append(' ' * indent + f"D0.u64 = int(_full) & 0xffffffffffffffff")
|
||||
lines.append(' ' * indent + f"D1 = Reg((int(_full) >> 64) & 1)")
|
||||
elif any(op in line for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^=')):
|
||||
for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^='):
|
||||
if op in line:
|
||||
lhs, rhs = line.split(op, 1)
|
||||
lines.append(' ' * indent + f"{lhs.strip()} {op} {_expr(rhs.strip())}")
|
||||
break
|
||||
else:
|
||||
lhs, rhs = line.split('=', 1)
|
||||
lhs_s, rhs_s = _expr(lhs.strip()), rhs.strip()
|
||||
stmt = _assign(lhs_s, _expr(rhs_s))
|
||||
if in_first_match_loop and rhs_s == 'i' and (lhs_s == 'tmp' or lhs_s == 'D0.i32'):
|
||||
stmt += "; break"
|
||||
lines.append(' ' * indent + stmt)
|
||||
if need_pass: lines.append(' ' * indent + "pass")
|
||||
return '\n'.join(lines)
|
||||
|
||||
def _assign(lhs: str, rhs: str) -> str:
|
||||
if lhs in ('tmp', 'SCC', 'VCC', 'EXEC', 'D0', 'D1', 'saveexec', 'PC'):
|
||||
return f"{lhs} = Reg({rhs})"
|
||||
return f"{lhs} = {rhs}"
|
||||
|
||||
def _expr(e: str) -> str:
|
||||
e = e.strip()
|
||||
e = e.replace('&&', ' and ').replace('||', ' or ').replace('<>', ' != ')
|
||||
e = re.sub(r'!([^=])', r' not \1', e)
|
||||
e = re.sub(r'\{\s*(\w+\.u32)\s*,\s*(\w+\.u32)\s*\}', r'_pack32(\1, \2)', e)
|
||||
def pack(m):
|
||||
hi, lo = _expr(m[1].strip()), _expr(m[2].strip())
|
||||
return f'_pack({hi}, {lo})'
|
||||
e = re.sub(r'\{\s*([^,{}]+)\s*,\s*([^,{}]+)\s*\}', pack, e)
|
||||
e = re.sub(r"1201'B\(2\.0\s*/\s*PI\)", "TWO_OVER_PI_1201", e)
|
||||
e = re.sub(r"\d+'([0-9a-fA-Fx]+)[UuFf]*", r'\1', e)
|
||||
e = re.sub(r"\d+'[FIBU]\(", "(", e)
|
||||
e = re.sub(r'\bB\(', '(', e)
|
||||
e = re.sub(r'([0-9a-fA-Fx])ULL\b', r'\1', e)
|
||||
e = re.sub(r'([0-9a-fA-Fx])LL\b', r'\1', e)
|
||||
e = re.sub(r'([0-9a-fA-Fx])U\b', r'\1', e)
|
||||
e = re.sub(r'(\d\.?\d*)F\b', r'\1', e)
|
||||
e = re.sub(r'(\[laneId\])\.[uib]\d+', r'\1', e)
|
||||
e = e.replace('+INF', 'INF').replace('-INF', '(-INF)')
|
||||
e = re.sub(r'NAN\.f\d+', 'float("nan")', e)
|
||||
def convert_verilog_slice(m):
|
||||
start, width = m.group(1).strip(), m.group(2).strip()
|
||||
return f'[({start}) + ({width}) - 1 : ({start})]'
|
||||
e = re.sub(r'\[([^:\[\]]+)\s*\+:\s*([^:\[\]]+)\]', convert_verilog_slice, e)
|
||||
def process_brackets(s):
|
||||
result, i = [], 0
|
||||
while i < len(s):
|
||||
if s[i] == '[':
|
||||
depth, start = 1, i + 1
|
||||
j = start
|
||||
while j < len(s) and depth > 0:
|
||||
if s[j] == '[': depth += 1
|
||||
elif s[j] == ']': depth -= 1
|
||||
j += 1
|
||||
inner = _expr(s[start:j-1])
|
||||
result.append('[' + inner + ']')
|
||||
i = j
|
||||
else:
|
||||
result.append(s[i])
|
||||
i += 1
|
||||
return ''.join(result)
|
||||
e = process_brackets(e)
|
||||
while '?' in e:
|
||||
depth, bracket, q = 0, 0, -1
|
||||
for i, c in enumerate(e):
|
||||
if c == '(': depth += 1
|
||||
elif c == ')': depth -= 1
|
||||
elif c == '[': bracket += 1
|
||||
elif c == ']': bracket -= 1
|
||||
elif c == '?' and depth == 0 and bracket == 0: q = i; break
|
||||
if q < 0: break
|
||||
depth, bracket, col = 0, 0, -1
|
||||
for i in range(q + 1, len(e)):
|
||||
if e[i] == '(': depth += 1
|
||||
elif e[i] == ')': depth -= 1
|
||||
elif e[i] == '[': bracket += 1
|
||||
elif e[i] == ']': bracket -= 1
|
||||
elif e[i] == ':' and depth == 0 and bracket == 0: col = i; break
|
||||
if col < 0: break
|
||||
cond, t, f = e[:q].strip(), e[q+1:col].strip(), e[col+1:].strip()
|
||||
e = f'(({t}) if ({cond}) else ({f}))'
|
||||
return e
|
||||
|
||||
def _apply_pseudocode_fixes(op_name: str, code: str) -> str:
|
||||
"""Apply known fixes for PDF pseudocode bugs."""
|
||||
if op_name == 'V_DIV_FMAS_F32':
|
||||
code = code.replace('D0.f32 = 2.0 ** 32 * fma(S0.f32, S1.f32, S2.f32)',
|
||||
'D0.f32 = (2.0 ** 64 if exponent(S2.f32) > 127 else 2.0 ** -64) * fma(S0.f32, S1.f32, S2.f32)')
|
||||
if op_name == 'V_DIV_FMAS_F64':
|
||||
code = code.replace('D0.f64 = 2.0 ** 64 * fma(S0.f64, S1.f64, S2.f64)',
|
||||
'D0.f64 = (2.0 ** 128 if exponent(S2.f64) > 1023 else 2.0 ** -128) * fma(S0.f64, S1.f64, S2.f64)')
|
||||
if op_name == 'V_DIV_SCALE_F32':
|
||||
code = code.replace('D0.f32 = float("nan")', 'VCC = Reg(0x1); D0.f32 = float("nan")')
|
||||
code = code.replace('elif S1.f32 == DENORM.f32:\n D0.f32 = ldexp(S0.f32, 64)', 'elif False:\n pass')
|
||||
code += '\nif S1.f32 == DENORM.f32:\n D0.f32 = float("nan")'
|
||||
code = code.replace('elif exponent(S2.f32) <= 23:\n D0.f32 = ldexp(S0.f32, 64)', 'elif exponent(S2.f32) <= 23:\n VCC = Reg(0x1); D0.f32 = ldexp(S0.f32, 64)')
|
||||
code = code.replace('elif S2.f32 / S1.f32 == DENORM.f32:\n VCC = Reg(0x1)\n if S0.f32 == S2.f32:\n D0.f32 = ldexp(S0.f32, 64)', 'elif S2.f32 / S1.f32 == DENORM.f32:\n VCC = Reg(0x1)')
|
||||
if op_name == 'V_DIV_SCALE_F64':
|
||||
code = code.replace('D0.f64 = float("nan")', 'VCC = Reg(0x1); D0.f64 = float("nan")')
|
||||
code = code.replace('elif S1.f64 == DENORM.f64:\n D0.f64 = ldexp(S0.f64, 128)', 'elif False:\n pass')
|
||||
code += '\nif S1.f64 == DENORM.f64:\n D0.f64 = float("nan")'
|
||||
code = code.replace('elif exponent(S2.f64) <= 52:\n D0.f64 = ldexp(S0.f64, 128)', 'elif exponent(S2.f64) <= 52:\n VCC = Reg(0x1); D0.f64 = ldexp(S0.f64, 128)')
|
||||
code = code.replace('elif S2.f64 / S1.f64 == DENORM.f64:\n VCC = Reg(0x1)\n if S0.f64 == S2.f64:\n D0.f64 = ldexp(S0.f64, 128)', 'elif S2.f64 / S1.f64 == DENORM.f64:\n VCC = Reg(0x1)')
|
||||
if op_name == 'V_DIV_FIXUP_F32':
|
||||
code = code.replace('D0.f32 = ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))',
|
||||
'D0.f32 = ((-OVERFLOW_F32) if (sign_out) else (OVERFLOW_F32)) if isNAN(S0.f32) else ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))')
|
||||
if op_name == 'V_DIV_FIXUP_F64':
|
||||
code = code.replace('D0.f64 = ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))',
|
||||
'D0.f64 = ((-OVERFLOW_F64) if (sign_out) else (OVERFLOW_F64)) if isNAN(S0.f64) else ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))')
|
||||
if op_name == 'V_TRIG_PREOP_F64':
|
||||
code = code.replace('result = F((TWO_OVER_PI_1201[1200 : 0] << shift.u32) & 0x1fffffffffffff)',
|
||||
'result = float(((TWO_OVER_PI_1201[1200 : 0] << int(shift)) >> (1201 - 53)) & 0x1fffffffffffff)')
|
||||
return code
|
||||
|
||||
def _generate_function(cls_name: str, op_name: str, pc: str, code: str) -> str:
|
||||
"""Generate a single compiled pseudocode function.
|
||||
Functions take int parameters and return dict of int values.
|
||||
Reg wrapping happens inside the function, only for registers actually used."""
|
||||
has_d1 = '{ D1' in pc
|
||||
is_cmpx = (cls_name in ('VOPCOp', 'VOP3Op')) and 'EXEC.u64[laneId]' in pc
|
||||
is_div_scale = 'DIV_SCALE' in op_name
|
||||
has_sdst = cls_name == 'VOP3SDOp' and ('VCC.u64[laneId]' in pc or is_div_scale)
|
||||
is_ds = cls_name == 'DSOp'
|
||||
is_flat = cls_name in ('FLATOp', 'GLOBALOp', 'SCRATCHOp')
|
||||
is_smem = cls_name == 'SMEMOp'
|
||||
has_s_array = 'S[i]' in pc # FMA_MIX style: S[0], S[1], S[2] array access
|
||||
combined = code + pc
|
||||
|
||||
fn_name = f"_{cls_name}_{op_name}"
|
||||
|
||||
# Detect which registers are used/modified
|
||||
def needs_init(name): return name in combined and not re.search(rf'^\s*{name}\s*=\s*Reg\(', code, re.MULTILINE)
|
||||
modifies_d0 = is_div_scale or bool(re.search(r'\bD0\b[.\[]', combined))
|
||||
modifies_exec = is_cmpx or bool(re.search(r'EXEC\.(u32|u64|b32|b64)\s*=', combined))
|
||||
modifies_vcc = has_sdst or bool(re.search(r'VCC\.(u32|u64|b32|b64)\s*=|VCC\.u64\[laneId\]\s*=', combined))
|
||||
modifies_scc = bool(re.search(r'\bSCC\s*=', combined))
|
||||
modifies_pc = bool(re.search(r'\bPC\s*=', combined))
|
||||
|
||||
# Build function signature and Reg init lines
|
||||
if is_smem:
|
||||
lines = [f"def {fn_name}(MEM, addr):"]
|
||||
reg_inits = ["ADDR=Reg(addr)", "SDATA=Reg(0)"]
|
||||
special_regs = []
|
||||
elif is_ds:
|
||||
lines = [f"def {fn_name}(MEM, addr, data0, data1, offset0, offset1):"]
|
||||
reg_inits = ["ADDR=Reg(addr)", "DATA0=Reg(data0)", "DATA1=Reg(data1)", "OFFSET0=Reg(offset0)", "OFFSET1=Reg(offset1)", "RETURN_DATA=Reg(0)"]
|
||||
special_regs = [('DATA', 'DATA0'), ('DATA2', 'DATA1'), ('OFFSET', 'OFFSET0'), ('ADDR_BASE', 'ADDR')]
|
||||
elif is_flat:
|
||||
lines = [f"def {fn_name}(MEM, addr, vdata, vdst):"]
|
||||
reg_inits = ["ADDR=addr", "VDATA=Reg(vdata)", "VDST=Reg(vdst)", "RETURN_DATA=Reg(0)"]
|
||||
special_regs = [('DATA', 'VDATA')]
|
||||
elif has_s_array:
|
||||
# FMA_MIX style: needs S[i] array, opsel, opsel_hi for source selection (neg/neg_hi applied in emu.py before call)
|
||||
lines = [f"def {fn_name}(s0, s1, s2, d0, scc, vcc, laneId, exec_mask, literal, VGPR, src0_idx=0, vdst_idx=0, pc=None, opsel=0, opsel_hi=0):"]
|
||||
reg_inits = ["S0=Reg(s0)", "S1=Reg(s1)", "S2=Reg(s2)", "S=[S0,S1,S2]", "D0=Reg(d0)", "OPSEL=Reg(opsel)", "OPSEL_HI=Reg(opsel_hi)"]
|
||||
special_regs = []
|
||||
# Detect array declarations like "declare in : 32'F[3]" and create them (rename 'in' to 'ins' since 'in' is a keyword)
|
||||
if "in[" in combined:
|
||||
reg_inits.append("ins=[Reg(0),Reg(0),Reg(0)]")
|
||||
code = code.replace("in[", "ins[")
|
||||
else:
|
||||
lines = [f"def {fn_name}(s0, s1, s2, d0, scc, vcc, laneId, exec_mask, literal, VGPR, src0_idx=0, vdst_idx=0, pc=None):"]
|
||||
# Only create Regs for registers actually used in the pseudocode
|
||||
reg_inits = []
|
||||
if 'S0' in combined: reg_inits.append("S0=Reg(s0)")
|
||||
if 'S1' in combined: reg_inits.append("S1=Reg(s1)")
|
||||
if 'S2' in combined: reg_inits.append("S2=Reg(s2)")
|
||||
if modifies_d0 or 'D0' in combined: reg_inits.append("D0=Reg(s0)" if is_div_scale else "D0=Reg(d0)")
|
||||
if modifies_scc or 'SCC' in combined: reg_inits.append("SCC=Reg(scc)")
|
||||
if modifies_vcc or 'VCC' in combined: reg_inits.append("VCC=Reg(vcc)")
|
||||
if modifies_exec or 'EXEC' in combined: reg_inits.append("EXEC=Reg(exec_mask)")
|
||||
if modifies_pc or 'PC' in combined: reg_inits.append("PC=Reg(pc) if pc is not None else None")
|
||||
special_regs = [('D1', 'Reg(0)'), ('SIMM16', 'Reg(literal)'), ('SIMM32', 'Reg(literal)'),
|
||||
('SRC0', 'Reg(src0_idx)'), ('VDST', 'Reg(vdst_idx)')]
|
||||
if needs_init('tmp'): special_regs.insert(0, ('tmp', 'Reg(0)'))
|
||||
if needs_init('saveexec'): special_regs.insert(0, ('saveexec', 'Reg(EXEC._val)'))
|
||||
|
||||
# Build init code
|
||||
init_parts = reg_inits.copy()
|
||||
for name, init in special_regs:
|
||||
if name in combined: init_parts.append(f"{name}={init}")
|
||||
if 'EXEC_LO' in code: init_parts.append("EXEC_LO=TypedView(EXEC, 31, 0)")
|
||||
if 'EXEC_HI' in code: init_parts.append("EXEC_HI=TypedView(EXEC, 63, 32)")
|
||||
if 'VCCZ' in code and not re.search(r'^\s*VCCZ\s*=', code, re.MULTILINE): init_parts.append("VCCZ=Reg(1 if VCC._val == 0 else 0)")
|
||||
if 'EXECZ' in code and not re.search(r'^\s*EXECZ\s*=', code, re.MULTILINE): init_parts.append("EXECZ=Reg(1 if EXEC._val == 0 else 0)")
|
||||
|
||||
# Add init line and separator
|
||||
if init_parts: lines.append(f" {'; '.join(init_parts)}")
|
||||
|
||||
# Add compiled pseudocode
|
||||
for line in code.split('\n'):
|
||||
if line.strip(): lines.append(f" {line}")
|
||||
|
||||
# Build result dict
|
||||
result_items = []
|
||||
if modifies_d0: result_items.append("'D0': D0._val")
|
||||
if modifies_scc: result_items.append("'SCC': SCC._val")
|
||||
if modifies_vcc: result_items.append("'VCC': VCC._val")
|
||||
if modifies_exec: result_items.append("'EXEC': EXEC._val")
|
||||
if has_d1: result_items.append("'D1': D1._val")
|
||||
if modifies_pc: result_items.append("'PC': PC._val")
|
||||
if is_smem and 'SDATA' in combined and re.search(r'^\s*SDATA[\.\[].*=', code, re.MULTILINE):
|
||||
result_items.append("'SDATA': SDATA._val")
|
||||
if is_ds and 'RETURN_DATA' in combined and re.search(r'^\s*RETURN_DATA[\.\[].*=', code, re.MULTILINE):
|
||||
result_items.append("'RETURN_DATA': RETURN_DATA._val")
|
||||
if is_flat:
|
||||
if 'RETURN_DATA' in combined and re.search(r'^\s*RETURN_DATA[\.\[].*=', code, re.MULTILINE):
|
||||
result_items.append("'RETURN_DATA': RETURN_DATA._val")
|
||||
if re.search(r'^\s*VDATA[\.\[].*=', code, re.MULTILINE):
|
||||
result_items.append("'VDATA': VDATA._val")
|
||||
lines.append(f" return {{{', '.join(result_items)}}}")
|
||||
return '\n'.join(lines)
|
||||
|
||||
# Build the globals dict for exec() - includes all pcode symbols
|
||||
_PCODE_GLOBALS = {
|
||||
'Reg': Reg, 'TypedView': TypedView, '_pack': _pack, '_pack32': _pack32,
|
||||
'ABSDIFF': ABSDIFF, 'BYTE_PERMUTE': BYTE_PERMUTE, 'DENORM': DENORM, 'F': F,
|
||||
'GT_NEG_ZERO': GT_NEG_ZERO, 'LT_NEG_ZERO': LT_NEG_ZERO, 'INF': INF,
|
||||
'MAX_FLOAT_F32': MAX_FLOAT_F32, 'OVERFLOW_F32': OVERFLOW_F32, 'OVERFLOW_F64': OVERFLOW_F64,
|
||||
'UNDERFLOW_F32': UNDERFLOW_F32, 'UNDERFLOW_F64': UNDERFLOW_F64,
|
||||
'PI': PI, 'ROUND_MODE': ROUND_MODE, 'WAVE_MODE': WAVE_MODE,
|
||||
'WAVE32': WAVE32, 'WAVE64': WAVE64, 'TWO_OVER_PI_1201': TWO_OVER_PI_1201,
|
||||
'SAT8': SAT8, 'trunc': trunc, 'floor': floor, 'ceil': ceil, 'sqrt': sqrt,
|
||||
'log2': log2, 'fract': fract, 'sin': sin, 'cos': cos, 'pow': pow,
|
||||
'isEven': isEven, 'mantissa': mantissa, 'signext_from_bit': signext_from_bit,
|
||||
'i32_to_f32': i32_to_f32, 'u32_to_f32': u32_to_f32, 'i32_to_f64': i32_to_f64,
|
||||
'u32_to_f64': u32_to_f64, 'f32_to_f64': f32_to_f64, 'f64_to_f32': f64_to_f32,
|
||||
'f32_to_i32': f32_to_i32, 'f32_to_u32': f32_to_u32, 'f64_to_i32': f64_to_i32,
|
||||
'f64_to_u32': f64_to_u32, 'f32_to_f16': f32_to_f16, 'f16_to_f32': f16_to_f32,
|
||||
'i16_to_f16': i16_to_f16, 'u16_to_f16': u16_to_f16, 'f16_to_i16': f16_to_i16,
|
||||
'f16_to_u16': f16_to_u16, 'bf16_to_f32': bf16_to_f32, 'f32_to_bf16': f32_to_bf16,
|
||||
'u8_to_u32': u8_to_u32, 'u4_to_u32': u4_to_u32, 'u32_to_u16': u32_to_u16,
|
||||
'i32_to_i16': i32_to_i16, 'f16_to_snorm': f16_to_snorm, 'f16_to_unorm': f16_to_unorm,
|
||||
'f32_to_snorm': f32_to_snorm, 'f32_to_unorm': f32_to_unorm,
|
||||
'v_cvt_i16_f32': v_cvt_i16_f32, 'v_cvt_u16_f32': v_cvt_u16_f32, 'f32_to_u8': f32_to_u8,
|
||||
'v_min_f32': v_min_f32, 'v_max_f32': v_max_f32, 'v_min_f16': v_min_f16, 'v_max_f16': v_max_f16,
|
||||
'v_min_i32': v_min_i32, 'v_max_i32': v_max_i32, 'v_min_i16': v_min_i16, 'v_max_i16': v_max_i16,
|
||||
'v_min_u32': v_min_u32, 'v_max_u32': v_max_u32, 'v_min_u16': v_min_u16, 'v_max_u16': v_max_u16,
|
||||
'v_min3_f32': v_min3_f32, 'v_max3_f32': v_max3_f32, 'v_min3_f16': v_min3_f16, 'v_max3_f16': v_max3_f16,
|
||||
'v_min3_i32': v_min3_i32, 'v_max3_i32': v_max3_i32, 'v_min3_i16': v_min3_i16, 'v_max3_i16': v_max3_i16,
|
||||
'v_min3_u32': v_min3_u32, 'v_max3_u32': v_max3_u32, 'v_min3_u16': v_min3_u16, 'v_max3_u16': v_max3_u16,
|
||||
'v_sad_u8': v_sad_u8, 'v_msad_u8': v_msad_u8,
|
||||
's_ff1_i32_b32': s_ff1_i32_b32, 's_ff1_i32_b64': s_ff1_i32_b64,
|
||||
'isNAN': isNAN, 'isQuietNAN': isQuietNAN, 'isSignalNAN': isSignalNAN,
|
||||
'fma': fma, 'ldexp': ldexp, 'sign': sign, 'exponent': exponent,
|
||||
'signext': signext, 'cvtToQuietNAN': cvtToQuietNAN,
|
||||
}
|
||||
|
||||
@functools.cache
|
||||
def compile_pseudocode(cls_name: str, op_name: str, pseudocode: str):
|
||||
"""Compile pseudocode string to executable function. Cached for performance."""
|
||||
filtered = _filter_pseudocode(pseudocode)
|
||||
code = _compile_pseudocode(filtered)
|
||||
code = _apply_pseudocode_fixes(op_name, code)
|
||||
fn_code = _generate_function(cls_name, op_name, filtered, code)
|
||||
fn_name = f"_{cls_name}_{op_name}"
|
||||
local_ns = {}
|
||||
exec(fn_code, _PCODE_GLOBALS, local_ns)
|
||||
return local_ns[fn_name]
|
||||
@@ -0,0 +1,319 @@
|
||||
# Generic PDF text extractor - no external dependencies
|
||||
import re, zlib
|
||||
from tinygrad.helpers import fetch, merge_dicts
|
||||
|
||||
PDF_URLS = {
|
||||
"rdna3": "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content",
|
||||
"rdna4": "https://docs.amd.com/api/khub/documents/uQpkEvk3pv~kfAb2x~j4uw/content",
|
||||
"cdna": "https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-cdna4-instruction-set-architecture.pdf",
|
||||
}
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# Generic PDF extraction tools
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def extract(url: str) -> list[list[tuple[float, float, str, str]]]:
|
||||
"""Extract positioned text from PDF. Returns list of text elements (x, y, text, font) per page."""
|
||||
data = fetch(url).read_bytes()
|
||||
|
||||
# Parse xref table to locate objects
|
||||
xref: dict[int, int] = {}
|
||||
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)
|
||||
start_obj, count = map(int, data[pos:line_end].split()[:2])
|
||||
pos = line_end + 1
|
||||
for i in range(count):
|
||||
if data[pos+17:pos+18] == b'n' and (off := int(data[pos:pos+10])) > 0: xref[start_obj + i] = off
|
||||
pos += 20
|
||||
|
||||
def get_stream(n: int) -> bytes:
|
||||
obj = data[xref[n]:data.find(b'endobj', xref[n])]
|
||||
raw = obj[obj.find(b'stream\n') + 7:obj.find(b'\nendstream')]
|
||||
return zlib.decompress(raw) if b'/FlateDecode' in obj else raw
|
||||
|
||||
# Find page content streams and extract text
|
||||
pages = []
|
||||
for n in sorted(xref):
|
||||
if b'/Type /Page' not in data[xref[n]:xref[n]+500]: continue
|
||||
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 = [], ''
|
||||
for bt in re.finditer(r'BT(.*?)ET', stream, re.S):
|
||||
x, y = 0.0, 0.0
|
||||
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
|
||||
|
||||
def extract_tables(pages: list[list[tuple[float, float, str, str]]]) -> dict[int, tuple[str, list[list[str]]]]:
|
||||
"""Extract numbered tables from PDF pages. Returns {table_num: (title, rows)} where rows is list of cells per row."""
|
||||
def group_by_y(texts, key=lambda y: round(y)):
|
||||
by_y: dict[int, list[tuple[float, float, str]]] = {}
|
||||
for x, y, t, _ in texts:
|
||||
by_y.setdefault(key(y), []).append((x, y, t))
|
||||
return by_y
|
||||
|
||||
# Find all table headers by merging text on same line
|
||||
table_positions = []
|
||||
for page_idx, texts in enumerate(pages):
|
||||
for items in group_by_y(texts).values():
|
||||
line = ''.join(t for _, t in sorted((x, t) for x, _, t in items))
|
||||
if m := re.search(r'Table (\d+)\. (.+)', line):
|
||||
table_positions.append((int(m.group(1)), m.group(2).strip(), page_idx, items[0][1]))
|
||||
table_positions.sort(key=lambda t: (t[2], -t[3]))
|
||||
|
||||
# For each table, find rows with matching X positions
|
||||
result: dict[int, tuple[str, list[list[str]]]] = {}
|
||||
for num, title, start_page, header_y in table_positions:
|
||||
rows, col_xs = [], None
|
||||
for page_idx in range(start_page, len(pages)):
|
||||
page_texts = [(x, y, t) for x, y, t, _ in pages[page_idx] if 30 < y < 760 and (page_idx > start_page or y < header_y)]
|
||||
for items in sorted(group_by_y([(x, y, t, '') for x, y, t in page_texts], key=lambda y: round(y / 5)).values(), key=lambda items: -items[0][1]):
|
||||
xs = tuple(sorted(round(x) for x, _, _ in items))
|
||||
if col_xs is None:
|
||||
if len(xs) < 2: continue # Skip single-column rows before table starts
|
||||
col_xs = xs
|
||||
elif len(xs) == 1 and xs[0] in col_xs: continue # Skip continuation rows at known column positions
|
||||
elif not any(c in xs for c in col_xs[:2]): break # Row missing first columns = end of table
|
||||
rows.append([t for _, t in sorted((x, t) for x, _, t in items)])
|
||||
else: continue
|
||||
break
|
||||
if rows: result[num] = (title, rows)
|
||||
return result
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# AMD specific extraction
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def extract_enums(tables: dict[int, tuple[str, list[list[str]]]]) -> dict[str, dict[int, str]]:
|
||||
"""Extract all enums from tables. Returns {enum_name: {value: name}}."""
|
||||
enums: dict[str, dict[int, str]] = {}
|
||||
for num, (title, rows) in tables.items():
|
||||
# Opcode enums from "XXX Opcodes" tables
|
||||
if m := re.match(r'(\w+) (?:Y-)?Opcodes', title):
|
||||
fmt_name = 'VOPD' if 'Y-Opcodes' in title else m.group(1)
|
||||
ops: dict[int, str] = {}
|
||||
for row in rows:
|
||||
for i in range(0, len(row) - 1, 2):
|
||||
if row[i].isdigit() and re.match(r'^[A-Z][A-Z0-9_]+$', row[i + 1]):
|
||||
ops[int(row[i])] = row[i + 1]
|
||||
if ops: enums[fmt_name] = ops
|
||||
# BufFmt from "Data Format" tables
|
||||
if 'Data Format' in title:
|
||||
for row in rows:
|
||||
for i in range(0, len(row) - 1, 2):
|
||||
if row[i].isdigit() and re.match(r'^[\dA-Z_]+$', row[i + 1]) and 'INVALID' not in row[i + 1]:
|
||||
enums.setdefault('BufFmt', {})[int(row[i])] = row[i + 1]
|
||||
return enums
|
||||
|
||||
def extract_ins(tables: dict[int, tuple[str, list[list[str]]]]) -> tuple[dict[str, list[tuple[str, int, int]]], dict[str, str]]:
|
||||
"""Extract formats and encodings from 'XXX Fields' tables. Returns (formats, encodings)."""
|
||||
formats: dict[str, list[tuple[str, int, int]]] = {}
|
||||
encodings: dict[str, str] = {}
|
||||
for num, (title, rows) in tables.items():
|
||||
if not (m := re.match(r'(\w+) Fields$', title)): continue
|
||||
fmt_name = m.group(1)
|
||||
fields = []
|
||||
for row in rows:
|
||||
if len(row) < 2: continue
|
||||
if (bits := re.match(r'\[?(\d+):(\d+)\]?$', row[1])) or (bits := re.match(r'\[(\d+)\]$', row[1])):
|
||||
field_name = row[0].lower()
|
||||
hi, lo = int(bits.group(1)), int(bits.group(2)) if bits.lastindex >= 2 else int(bits.group(1))
|
||||
if field_name == 'encoding' and len(row) >= 3:
|
||||
enc_bits = None
|
||||
if "'b" in row[2]: enc_bits = row[2].split("'b")[-1].replace('_', '')
|
||||
elif (enc := re.search(r':\s*([01_]+)', row[2])): enc_bits = enc.group(1).replace('_', '')
|
||||
if enc_bits:
|
||||
# If encoding bits exceed field width, extend field to match (AMD docs sometimes have this)
|
||||
declared_width, actual_width = hi - lo + 1, len(enc_bits)
|
||||
if actual_width > declared_width: lo = hi - actual_width + 1
|
||||
encodings[fmt_name] = enc_bits
|
||||
fields.append((field_name, hi, lo))
|
||||
if fields: formats[fmt_name] = fields
|
||||
return formats, encodings
|
||||
|
||||
def extract_pcode(pages: list[list[tuple[float, float, str, str]]], enums: dict[str, dict[int, str]]) -> dict[tuple[str, int], str]:
|
||||
"""Extract pseudocode for instructions. Returns {(name, opcode): pseudocode}."""
|
||||
# Build lookup from instruction name to opcode
|
||||
name_to_op = {name: op for ops in enums.values() for op, name in ops.items()}
|
||||
|
||||
# First pass: find all instruction headers across all pages
|
||||
all_instructions: list[tuple[int, float, str, int]] = [] # (page_idx, y, name, opcode)
|
||||
for page_idx, page in enumerate(pages):
|
||||
by_y: dict[int, list[tuple[float, str]]] = {}
|
||||
for x, y, t, _ in page:
|
||||
by_y.setdefault(round(y), []).append((x, t))
|
||||
for y, items in sorted(by_y.items(), reverse=True):
|
||||
left = [(x, t) for x, t in items if 55 < x < 65]
|
||||
right = [(x, t) for x, t in items if 535 < x < 550]
|
||||
if left and right and left[0][1] in name_to_op and right[0][1].isdigit():
|
||||
all_instructions.append((page_idx, y, left[0][1], int(right[0][1])))
|
||||
|
||||
# Second pass: extract pseudocode between consecutive instructions
|
||||
pcode: dict[tuple[str, int], str] = {}
|
||||
for i, (page_idx, y, name, opcode) in enumerate(all_instructions):
|
||||
# Get end boundary from next instruction
|
||||
if i + 1 < len(all_instructions):
|
||||
next_page, next_y = all_instructions[i + 1][0], all_instructions[i + 1][1]
|
||||
else:
|
||||
next_page, next_y = page_idx, 0
|
||||
# Collect F6 text from current position to next instruction (pseudocode is at x ≈ 69)
|
||||
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
|
||||
lines.extend((p, y2, t) for x, y2, t, f in pages[p] if f in ('/F6.0', '/F7.0') and end_y < y2 < start_y and 60 < x < 80)
|
||||
if lines:
|
||||
# Sort by page first, then by y descending within each page (higher y = earlier text in PDF)
|
||||
sorted_lines = sorted(lines, key=lambda x: (x[0], -x[1]))
|
||||
# Stop at large Y gaps (>30) - indicates section break (Notes, examples, etc)
|
||||
filtered = [sorted_lines[0]]
|
||||
for j in range(1, len(sorted_lines)):
|
||||
prev_page, prev_y, _ = sorted_lines[j-1]
|
||||
curr_page, curr_y, _ = sorted_lines[j]
|
||||
if curr_page == prev_page and prev_y - curr_y > 30: break
|
||||
if curr_page != prev_page and prev_y > 60 and curr_y < 730: break # examples spilled to next page (not at very top)
|
||||
filtered.append(sorted_lines[j])
|
||||
pcode_lines = [t.replace('Ê', '').strip() for _, _, t in filtered]
|
||||
if pcode_lines: pcode[(name, opcode)] = '\n'.join(pcode_lines)
|
||||
return pcode
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# Write autogen files
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def write_enums(enums: dict[str, dict[int, str]], arch: str, path: str):
|
||||
"""Write enum.py file from extracted enums."""
|
||||
lines = ["# autogenerated from AMD ISA PDF by pdf.py - do not edit", "from enum import IntEnum", ""]
|
||||
for name, values in sorted(enums.items()):
|
||||
suffix = "Op" if name not in ('Src', 'BufFmt') else ("Enum" if name == 'Src' else "")
|
||||
prefix = "BUF_FMT_" if name == 'BufFmt' else ""
|
||||
lines.append(f"class {name}{suffix}(IntEnum):")
|
||||
for val, member in sorted(values.items()):
|
||||
lines.append(f" {prefix}{member} = {val}")
|
||||
lines.append("")
|
||||
with open(path, "w") as f:
|
||||
f.write("\n".join(lines))
|
||||
|
||||
def write_ins(formats: dict[str, list[tuple[str, int, int]]], encodings: dict[str, str], enums: dict[str, dict[int, str]], arch: str, path: str):
|
||||
"""Write ins.py file from extracted formats and enums."""
|
||||
# Field types and ordering
|
||||
def field_type(name, fmt):
|
||||
if name == 'op' and fmt in enums: return f'Annotated[BitField, {fmt}Op]'
|
||||
if name in ('opx', 'opy'): return 'Annotated[BitField, VOPDOp]'
|
||||
if name == 'vdsty': return 'VDSTYEnc'
|
||||
if name in ('vdst', 'vsrc1', 'vaddr', 'vdata', 'data', 'data0', 'data1', 'addr', 'vsrc0', 'vsrc2', 'vsrc3'): return 'VGPRField'
|
||||
if name in ('sdst', 'sbase', 'sdata', 'srsrc', 'ssamp'): return 'SGPRField'
|
||||
if name.startswith('ssrc') or name in ('saddr', 'soffset'): return 'SSrc'
|
||||
if name in ('src0', 'srcx0', 'srcy0') or name.startswith('src') and name[3:].isdigit(): return 'Src'
|
||||
if name.startswith('simm'): return 'SImm'
|
||||
if name == 'offset' or name.startswith('imm'): return 'Imm'
|
||||
return None
|
||||
field_priority = ['encoding', 'op', 'opx', 'opy', 'vdst', 'vdstx', 'vdsty', 'sdst', 'vdata', 'sdata', 'addr', 'vaddr', 'data', 'data0', 'data1',
|
||||
'src0', 'srcx0', 'srcy0', 'vsrc0', 'ssrc0', 'src1', 'vsrc1', 'vsrcx1', 'vsrcy1', 'ssrc1', 'src2', 'vsrc2', 'src3', 'vsrc3',
|
||||
'saddr', 'sbase', 'srsrc', 'ssamp', 'soffset', 'offset', 'simm16', 'en', 'target', 'attr', 'attr_chan',
|
||||
'omod', 'neg', 'neg_hi', 'abs', 'clmp', 'opsel', 'opsel_hi', 'waitexp', 'wait_va',
|
||||
'dmask', 'dim', 'seg', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe', 'unrm', 'done', 'row']
|
||||
def sort_fields(fields):
|
||||
order = {name: i for i, name in enumerate(field_priority)}
|
||||
return sorted(fields, key=lambda f: (order.get(f[0], 1000), f[2]))
|
||||
|
||||
# Generate format classes
|
||||
lines = ["# autogenerated from AMD ISA PDF by pdf.py - do not edit", "# ruff: noqa: F401,F403",
|
||||
"from typing import Annotated",
|
||||
"from extra.assembly.amd.dsl import *",
|
||||
f"from extra.assembly.amd.autogen.{arch}.enum import *", "import functools", ""]
|
||||
for fmt_name, fields in sorted(formats.items()):
|
||||
lines.append(f"class {fmt_name}(Inst):")
|
||||
for name, hi, lo in sort_fields(fields):
|
||||
bits_str = f"bits[{hi}:{lo}]" if hi != lo else f"bits[{hi}]"
|
||||
if name == 'encoding' and fmt_name in encodings: lines.append(f" encoding = {bits_str} == 0b{encodings[fmt_name]}")
|
||||
else:
|
||||
ftype = field_type(name, fmt_name)
|
||||
lines.append(f" {name}{f':{ftype}' if ftype else ''} = {bits_str}")
|
||||
lines.append("")
|
||||
|
||||
# Generate instruction helpers
|
||||
lines.append("# instruction helpers")
|
||||
for fmt_name, ops in sorted(enums.items()):
|
||||
seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=1"}.get(fmt_name, "")
|
||||
tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt_name, f"{fmt_name}, {fmt_name}Op")
|
||||
suffix = "_e32" if fmt_name in ("VOP1", "VOP2", "VOPC") else "_e64" if fmt_name == "VOP3" and len(ops) > 0 else ""
|
||||
if fmt_name in formats or fmt_name in ("GLOBAL", "SCRATCH"):
|
||||
for op_val, name in sorted(ops.items()):
|
||||
fn_suffix = suffix if fmt_name != "VOP3" or op_val < 512 else ""
|
||||
lines.append(f"{name.lower()}{fn_suffix} = functools.partial({tgt}.{name}{seg})")
|
||||
|
||||
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: str):
|
||||
"""Write str_pcode.py file from extracted pseudocode."""
|
||||
# Group pseudocode by enum class
|
||||
by_enum: dict[str, list[tuple[str, int, str]]] = {}
|
||||
for fmt_name, ops in enums.items():
|
||||
for opcode, name in ops.items():
|
||||
if (name, opcode) in pcode: by_enum.setdefault(f"{fmt_name}Op", []).append((name, opcode, pcode[(name, opcode)]))
|
||||
# Generate file
|
||||
enum_names = sorted(by_enum.keys())
|
||||
lines = [f"# autogenerated by pdf.py - do not edit", f"# to regenerate: python -m extra.assembly.amd.pdf",
|
||||
"# ruff: noqa: E501", f"from extra.assembly.amd.autogen.{arch}.enum import {', '.join(enum_names)}", ""]
|
||||
for enum_name in enum_names:
|
||||
lines.append(f"{enum_name}_PCODE = {{")
|
||||
for name, opcode, code in sorted(by_enum[enum_name], key=lambda x: x[1]):
|
||||
lines.append(f" {enum_name}.{name}: {code!r},")
|
||||
lines.append("}\n")
|
||||
lines.append(f"PSEUDOCODE_STRINGS = {{{', '.join(f'{e}: {e}_PCODE' for e in enum_names)}}}")
|
||||
with open(path, "w") as f:
|
||||
f.write("\n".join(lines))
|
||||
|
||||
if __name__ == "__main__":
|
||||
import pathlib
|
||||
for arch, url in PDF_URLS.items():
|
||||
print(f"Processing {arch}...")
|
||||
pages = extract(url)
|
||||
tables = extract_tables(pages)
|
||||
enums = extract_enums(tables)
|
||||
formats, encodings = extract_ins(tables)
|
||||
pcode = extract_pcode(pages, enums)
|
||||
# Fix known PDF errors
|
||||
if arch == 'rdna3':
|
||||
fixes = {'SOPP': {8: 'S_WAITCNT_DEPCTR', 58: 'S_TTRACEDATA', 59: 'S_TTRACEDATA_IMM'},
|
||||
'SOPK': {22: 'S_SUBVECTOR_LOOP_BEGIN', 23: 'S_SUBVECTOR_LOOP_END'},
|
||||
'SMEM': {34: 'S_ATC_PROBE', 35: 'S_ATC_PROBE_BUFFER'},
|
||||
'DS': {24: 'DS_GWS_SEMA_RELEASE_ALL', 25: 'DS_GWS_INIT', 26: 'DS_GWS_SEMA_V', 27: 'DS_GWS_SEMA_BR', 28: 'DS_GWS_SEMA_P', 29: 'DS_GWS_BARRIER'},
|
||||
'FLAT': {40: 'GLOBAL_LOAD_ADDTID_B32', 41: 'GLOBAL_STORE_ADDTID_B32', 55: 'FLAT_ATOMIC_CSUB_U32'}}
|
||||
for fmt, ops in fixes.items(): enums[fmt] = merge_dicts([enums[fmt], ops])
|
||||
if arch == 'rdna4':
|
||||
fixes = {'SMEM': {34: 'S_ATC_PROBE', 35: 'S_ATC_PROBE_BUFFER'},
|
||||
'SOP1': {81: 'S_BARRIER_INIT', 82: 'S_BARRIER_JOIN'},
|
||||
'SOPP': {21: 'S_BARRIER_LEAVE', 58: 'S_TTRACEDATA', 59: 'S_TTRACEDATA_IMM'}}
|
||||
for fmt, ops in fixes.items(): enums[fmt] = merge_dicts([enums[fmt], ops])
|
||||
if arch in ('rdna3', 'rdna4'):
|
||||
# RDNA SMEM: PDF says DLC=[14], GLC=[16] but hardware uses DLC=[13], GLC=[14]
|
||||
if 'SMEM' in formats:
|
||||
formats['SMEM'] = [(n, 13 if n == 'dlc' else 14 if n == 'glc' else h, 13 if n == 'dlc' else 14 if n == 'glc' else l)
|
||||
for n, h, l in formats['SMEM']]
|
||||
if arch == 'cdna':
|
||||
# CDNA DS: PDF is missing the GDS field (bit 16)
|
||||
if 'DS' in formats and not any(n == 'gds' for n, _, _ in formats['DS']):
|
||||
formats['DS'].append(('gds', 16, 16))
|
||||
# CDNA DPP/SDWA: PDF only documents modifier fields (bits[63:32]), need to add VOP overlay fields (bits[31:0])
|
||||
vop_overlay = [('encoding', 8, 0), ('vop_op', 16, 9), ('vdst', 24, 17), ('vop2_op', 31, 25)]
|
||||
if 'DPP' in formats and not any(n == 'encoding' for n, _, _ in formats['DPP']):
|
||||
formats['DPP'] = vop_overlay + [('bc' if n == 'bound_ctrl' else n, h, l) for n, h, l in formats['DPP']]
|
||||
encodings['DPP'] = '11111010'
|
||||
if 'SDWA' in formats and not any(n == 'encoding' for n, _, _ in formats['SDWA']):
|
||||
formats['SDWA'] = vop_overlay + [(n, h, l) for n, h, l in formats['SDWA']]
|
||||
encodings['SDWA'] = '11111001'
|
||||
base = pathlib.Path(__file__).parent / "autogen" / arch
|
||||
write_enums(enums, arch, base / "enum.py")
|
||||
write_ins(formats, encodings, enums, arch, base / "ins.py")
|
||||
write_pcode(pcode, enums, arch, base / "str_pcode.py")
|
||||
print(f" {len(tables)} tables, {len(pcode)} pcode -> {base}")
|
||||
@@ -0,0 +1,381 @@
|
||||
"""SQTT (SQ Thread Trace) packet encoder and decoder for AMD GPUs.
|
||||
|
||||
This module provides encoding and decoding of raw SQTT byte streams.
|
||||
The format is nibble-based with variable-width packets determined by a state machine.
|
||||
Uses BitField infrastructure from dsl.py, similar to GPU instruction encoding.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
from enum import IntEnum
|
||||
from typing import get_type_hints
|
||||
from extra.assembly.amd.dsl import BitField, bits
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# FIELD ENUMS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class MemSrc(IntEnum):
|
||||
LDS = 0
|
||||
LDS_ALT = 1
|
||||
VMEM = 2
|
||||
VMEM_ALT = 3
|
||||
|
||||
class AluSrc(IntEnum):
|
||||
NONE = 0
|
||||
SALU = 1
|
||||
VALU = 2
|
||||
VALU_ALT = 3
|
||||
|
||||
class InstOp(IntEnum):
|
||||
"""SQTT instruction operation types.
|
||||
|
||||
Memory ops appear in two ranges depending on which SIMD executes them:
|
||||
- 0x1x-0x2x range: ops on traced SIMD
|
||||
- 0x5x range: ops on other SIMD (OTHER_ prefix)
|
||||
|
||||
GLOBAL memory ops encoding depends on addressing mode AND size:
|
||||
- Loads: 0x21 (saddr=SGPR) or 0x22 (saddr=NULL), all sizes same
|
||||
- Stores: base + size_offset, where VADDR is shifted +1 from SADDR
|
||||
SADDR: 0x24(32) 0x25(64) 0x26(96) 0x27(128)
|
||||
VADDR: 0x25(32) 0x26(64) 0x27(96) 0x28(128)
|
||||
|
||||
OTHER_ range follows same pattern but values overlap differently.
|
||||
"""
|
||||
SALU = 0x0
|
||||
SMEM = 0x1
|
||||
JUMP = 0x3 # branch taken
|
||||
JUMP_NO = 0x4 # branch not taken
|
||||
MESSAGE = 0x9
|
||||
VALU_TRANS = 0xb # transcendental: exp, log, rcp, sqrt, sin, cos
|
||||
VALU_64_SHIFT = 0xd # 64-bit shifts: lshl, lshr, ashr
|
||||
VALU_MAD64 = 0xe # 64-bit multiply-add
|
||||
VALU_64 = 0xf # 64-bit: add, mul, fma, rcp, sqrt, rounding, frexp, div helpers
|
||||
VINTERP = 0x12 # interpolation: v_interp_p10_f32, v_interp_p2_f32
|
||||
BARRIER = 0x13
|
||||
|
||||
# FLAT memory ops on traced SIMD (0x1x range)
|
||||
FLAT_LOAD = 0x1c
|
||||
FLAT_STORE = 0x1d
|
||||
FLAT_STORE_64 = 0x1e
|
||||
FLAT_STORE_96 = 0x1f
|
||||
FLAT_STORE_128 = 0x20
|
||||
|
||||
# GLOBAL memory ops on traced SIMD (0x2x range)
|
||||
GLOBAL_LOAD = 0x21 # saddr=SGPR, all sizes
|
||||
GLOBAL_LOAD_VADDR = 0x22 # saddr=NULL, all sizes
|
||||
GLOBAL_STORE = 0x24 # saddr=SGPR, 32-bit
|
||||
GLOBAL_STORE_64 = 0x25 # saddr=SGPR 64 or saddr=NULL 32
|
||||
GLOBAL_STORE_96 = 0x26 # saddr=SGPR 96 or saddr=NULL 64
|
||||
GLOBAL_STORE_128 = 0x27 # saddr=SGPR 128 or saddr=NULL 96
|
||||
GLOBAL_STORE_VADDR_128 = 0x28 # saddr=NULL, 128-bit
|
||||
|
||||
# LDS ops on traced SIMD
|
||||
LDS_LOAD = 0x29
|
||||
LDS_STORE = 0x2b
|
||||
LDS_STORE_64 = 0x2c
|
||||
LDS_STORE_128 = 0x2e
|
||||
|
||||
# Memory ops on other SIMD (0x5x range)
|
||||
OTHER_LDS_LOAD = 0x50
|
||||
OTHER_LDS_STORE = 0x51
|
||||
OTHER_LDS_STORE_64 = 0x52
|
||||
OTHER_LDS_STORE_128 = 0x54
|
||||
OTHER_FLAT_LOAD = 0x55
|
||||
OTHER_FLAT_STORE = 0x56
|
||||
OTHER_FLAT_STORE_64 = 0x57
|
||||
OTHER_FLAT_STORE_96 = 0x58
|
||||
OTHER_FLAT_STORE_128 = 0x59
|
||||
OTHER_GLOBAL_LOAD = 0x5a # saddr=SGPR, all sizes
|
||||
OTHER_GLOBAL_LOAD_VADDR = 0x5b # saddr=NULL or saddr=SGPR store 32
|
||||
OTHER_GLOBAL_STORE_64 = 0x5c # saddr=SGPR 64 or saddr=NULL 32
|
||||
OTHER_GLOBAL_STORE_96 = 0x5d # saddr=SGPR 96 or saddr=NULL 64
|
||||
OTHER_GLOBAL_STORE_128 = 0x5e # saddr=SGPR 128 or saddr=NULL 96
|
||||
OTHER_GLOBAL_STORE_VADDR_128 = 0x5f # saddr=NULL, 128-bit
|
||||
|
||||
# EXEC-modifying ops (0x7x range)
|
||||
SALU_SAVEEXEC = 0x72 # s_*_saveexec_b32/b64
|
||||
VALU_CMPX = 0x73 # v_cmpx_*
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PACKET TYPE BASE CLASS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class PacketType:
|
||||
"""Base class for SQTT packet types."""
|
||||
_encoding: tuple[BitField, int] | None = None
|
||||
_field_types: dict[str, type] = {}
|
||||
_values: dict[str, int]
|
||||
_raw: int
|
||||
_time: int
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
if 'encoding' in cls.__dict__ and isinstance(cls.__dict__['encoding'], tuple):
|
||||
cls._encoding = cls.__dict__['encoding']
|
||||
# Cache field type annotations for enum conversion
|
||||
try: cls._field_types = {k: v for k, v in get_type_hints(cls).items() if isinstance(v, type) and issubclass(v, IntEnum)}
|
||||
except Exception: cls._field_types = {}
|
||||
# Cache fields and precompute extraction info: (name, lo, mask, enum_type)
|
||||
cls._fields = {k: v for k, v in cls.__dict__.items() if isinstance(v, BitField) and k != 'encoding'}
|
||||
cls._extract_info = [(name, bf.lo, bf.mask(), cls._field_types.get(name)) for name, bf in cls._fields.items()]
|
||||
cls._size_nibbles = ((max((f.hi for f in cls._fields.values()), default=0) + 4) // 4)
|
||||
|
||||
@classmethod
|
||||
def from_raw(cls, raw: int, time: int = 0):
|
||||
inst = object.__new__(cls)
|
||||
inst._raw, inst._time, inst._values = raw, time, {}
|
||||
for name, lo, mask, enum_type in cls._extract_info:
|
||||
val = (raw >> lo) & mask
|
||||
if enum_type is not None:
|
||||
try: val = enum_type(val)
|
||||
except ValueError: pass
|
||||
inst._values[name] = val
|
||||
return inst
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
if name.startswith('_'): raise AttributeError(name)
|
||||
return self._values.get(name, 0)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
fields_str = ", ".join(f"{k}={v}" for k, v in self._values.items() if not k.startswith('_'))
|
||||
return f"{self.__class__.__name__}({fields_str})"
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PACKET TYPE DEFINITIONS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class VALUINST(PacketType): # exclude: 1 << 2
|
||||
encoding = bits[2:0] == 0b011
|
||||
delta = bits[5:3]
|
||||
flag = bits[6:6]
|
||||
wave = bits[11:7]
|
||||
|
||||
class VMEMEXEC(PacketType): # exclude: 1 << 0
|
||||
encoding = bits[3:0] == 0b1111
|
||||
delta = bits[5:4]
|
||||
src: MemSrc = bits[7:6]
|
||||
|
||||
class ALUEXEC(PacketType): # exclude: 1 << 1
|
||||
encoding = bits[3:0] == 0b1110
|
||||
delta = bits[5:4]
|
||||
src: AluSrc = bits[7:6]
|
||||
|
||||
class IMMEDIATE(PacketType): # exclude: 1 << 5
|
||||
encoding = bits[3:0] == 0b1101
|
||||
delta = bits[6:4]
|
||||
wave = bits[11:7]
|
||||
|
||||
class IMMEDIATE_MASK(PacketType): # exclude: 1 << 5
|
||||
encoding = bits[4:0] == 0b00100
|
||||
delta = bits[7:5]
|
||||
mask = bits[23:8]
|
||||
|
||||
class WAVERDY(PacketType): # exclude: 1 << 3
|
||||
encoding = bits[4:0] == 0b10100
|
||||
delta = bits[7:5]
|
||||
mask = bits[23:8]
|
||||
|
||||
class TS_DELTA_S8_W3(PacketType):
|
||||
encoding = bits[6:0] == 0b0100001
|
||||
delta = bits[10:8]
|
||||
_padding = bits[63:11]
|
||||
|
||||
class WAVEEND(PacketType): # exclude: 1 << 4
|
||||
encoding = bits[4:0] == 0b10101
|
||||
delta = bits[7:5]
|
||||
flag7 = bits[8:8]
|
||||
simd = bits[10:9]
|
||||
cu_lo = bits[13:11]
|
||||
wave = bits[19:15]
|
||||
@property
|
||||
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
|
||||
|
||||
class WAVESTART(PacketType): # exclude: 1 << 4
|
||||
encoding = bits[4:0] == 0b01100
|
||||
delta = bits[6:5]
|
||||
flag7 = bits[7:7]
|
||||
simd = bits[9:8]
|
||||
cu_lo = bits[12:10]
|
||||
wave = bits[17:13]
|
||||
id7 = bits[31:18]
|
||||
@property
|
||||
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
|
||||
|
||||
class TS_DELTA_S5_W2(PacketType):
|
||||
encoding = bits[4:0] == 0b11100
|
||||
delta = bits[6:5]
|
||||
_padding = bits[47:7]
|
||||
|
||||
class WAVEALLOC(PacketType): # exclude: 1 << 10
|
||||
encoding = bits[4:0] == 0b00101
|
||||
delta = bits[7:5]
|
||||
_padding = bits[19:8]
|
||||
|
||||
class TS_DELTA_S5_W3(PacketType):
|
||||
encoding = bits[4:0] == 0b00110
|
||||
delta = bits[7:5]
|
||||
_padding = bits[51:8]
|
||||
|
||||
class PERF(PacketType): # exclude: 1 << 11
|
||||
encoding = bits[4:0] == 0b10110
|
||||
delta = bits[7:5]
|
||||
arg = bits[27:8]
|
||||
|
||||
class TS_DELTA_SHORT(PacketType):
|
||||
encoding = bits[3:0] == 0b1000
|
||||
delta = bits[7:4]
|
||||
|
||||
class NOP(PacketType):
|
||||
encoding = bits[3:0] == 0b0000
|
||||
delta = None # type: ignore
|
||||
_padding = bits[3:0]
|
||||
|
||||
class TS_WAVE_STATE(PacketType):
|
||||
encoding = bits[6:0] == 0b1010001
|
||||
delta = bits[15:7]
|
||||
coarse = bits[23:16]
|
||||
@property
|
||||
def wave_interest(self) -> bool: return bool(self.coarse & 1)
|
||||
@property
|
||||
def terminate_all(self) -> bool: return bool(self.coarse & 8)
|
||||
|
||||
class EVENT(PacketType): # exclude: 1 << 7
|
||||
encoding = bits[7:0] == 0b01100001
|
||||
delta = bits[10:8]
|
||||
event = bits[23:11]
|
||||
|
||||
class EVENT_BIG(PacketType):
|
||||
encoding = bits[7:0] == 0b11100001
|
||||
delta = bits[10:8]
|
||||
event = bits[31:11]
|
||||
|
||||
class REG(PacketType):
|
||||
encoding = bits[3:0] == 0b1001
|
||||
delta = bits[6:4]
|
||||
slot = bits[9:7]
|
||||
hi_byte = bits[15:8]
|
||||
subop = bits[31:16]
|
||||
val32 = bits[63:32]
|
||||
@property
|
||||
def is_config(self) -> bool: return bool(self.hi_byte & 0x80)
|
||||
|
||||
class SNAPSHOT(PacketType):
|
||||
encoding = bits[6:0] == 0b1110001
|
||||
delta = bits[9:7]
|
||||
snap = bits[63:10]
|
||||
|
||||
class TS_DELTA_OR_MARK(PacketType):
|
||||
encoding = bits[6:0] == 0b0000001
|
||||
delta = bits[47:12]
|
||||
bit8 = bits[8:8]
|
||||
bit9 = bits[9:9]
|
||||
@property
|
||||
def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
|
||||
|
||||
class LAYOUT_HEADER(PacketType):
|
||||
encoding = bits[6:0] == 0b0010001
|
||||
delta = None # type: ignore
|
||||
layout = bits[12:7]
|
||||
simd = bits[14:13]
|
||||
group = bits[17:15]
|
||||
sel_a = bits[31:28]
|
||||
sel_b = bits[36:33]
|
||||
flag4 = bits[59:59]
|
||||
_padding = bits[63:60]
|
||||
|
||||
class INST(PacketType):
|
||||
encoding = bits[2:0] == 0b010
|
||||
delta = bits[6:4]
|
||||
flag1 = bits[3:3]
|
||||
flag2 = bits[7:7]
|
||||
wave = bits[12:8]
|
||||
op: InstOp = bits[19:13]
|
||||
|
||||
class UTILCTR(PacketType):
|
||||
encoding = bits[6:0] == 0b0110001
|
||||
delta = bits[8:7]
|
||||
ctr = bits[47:9]
|
||||
|
||||
# All packet types in encoding priority order (more specific masks first, NOP last as fallback)
|
||||
PACKET_TYPES: list[type[PacketType]] = [
|
||||
EVENT, EVENT_BIG,
|
||||
TS_DELTA_S8_W3, TS_WAVE_STATE, SNAPSHOT, TS_DELTA_OR_MARK, LAYOUT_HEADER, UTILCTR,
|
||||
IMMEDIATE_MASK, WAVERDY, WAVEEND, WAVESTART, TS_DELTA_S5_W2, WAVEALLOC, TS_DELTA_S5_W3, PERF,
|
||||
VMEMEXEC, ALUEXEC, IMMEDIATE, TS_DELTA_SHORT, REG,
|
||||
VALUINST, INST,
|
||||
NOP,
|
||||
]
|
||||
|
||||
def _build_state_table() -> tuple[bytes, dict[int, type[PacketType]]]:
|
||||
table = [len(PACKET_TYPES) - 1] * 256 # default to NOP
|
||||
opcode_to_class: dict[int, type[PacketType]] = {i: cls for i, cls in enumerate(PACKET_TYPES)}
|
||||
|
||||
for byte_val in range(256):
|
||||
for opcode, pkt_cls in enumerate(PACKET_TYPES):
|
||||
if pkt_cls._encoding is None: continue
|
||||
mask_bf, pattern = pkt_cls._encoding
|
||||
if (byte_val & mask_bf.mask()) == pattern:
|
||||
table[byte_val] = opcode
|
||||
break
|
||||
|
||||
return bytes(table), opcode_to_class
|
||||
|
||||
STATE_TO_OPCODE, OPCODE_TO_CLASS = _build_state_table()
|
||||
|
||||
# Precompute special case opcodes
|
||||
_TS_DELTA_OR_MARK_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_OR_MARK)
|
||||
_TS_DELTA_SHORT_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_SHORT)
|
||||
_TS_DELTA_OR_MARK_BIT8 = (TS_DELTA_OR_MARK.bit8.lo, TS_DELTA_OR_MARK.bit8.mask())
|
||||
_TS_DELTA_OR_MARK_BIT9 = (TS_DELTA_OR_MARK.bit9.lo, TS_DELTA_OR_MARK.bit9.mask())
|
||||
|
||||
# Combined lookup: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
|
||||
# special_case: 0=none, 1=TS_DELTA_OR_MARK, 2=TS_DELTA_SHORT
|
||||
_DECODE_INFO: dict[int, tuple] = {}
|
||||
for _opcode, _pkt_cls in OPCODE_TO_CLASS.items():
|
||||
_delta_field = getattr(_pkt_cls, 'delta', None)
|
||||
_delta_lo = _delta_field.lo if _delta_field else 0
|
||||
_delta_mask = _delta_field.mask() if _delta_field else 0
|
||||
_special = 1 if _opcode == _TS_DELTA_OR_MARK_OPCODE else (2 if _opcode == _TS_DELTA_SHORT_OPCODE else 0)
|
||||
_DECODE_INFO[_opcode] = (_pkt_cls, _pkt_cls._size_nibbles, _delta_lo, _delta_mask, _special)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# DECODER
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def decode(data: bytes) -> list[PacketType]:
|
||||
"""Decode raw SQTT blob into list of packet instances."""
|
||||
packets: list[PacketType] = []
|
||||
packets_append = packets.append
|
||||
n = len(data)
|
||||
reg = 0
|
||||
offset = 0
|
||||
nib_count = 16
|
||||
time = 0
|
||||
state_to_opcode = STATE_TO_OPCODE
|
||||
decode_info = _DECODE_INFO
|
||||
mask64 = (1 << 64) - 1
|
||||
|
||||
while (offset >> 3) < n:
|
||||
target = offset + nib_count * 4
|
||||
while offset < target and (offset >> 3) < n:
|
||||
byte = data[offset >> 3]
|
||||
nib = (byte >> (offset & 4)) & 0xF
|
||||
reg = ((reg >> 4) | (nib << 60)) & mask64
|
||||
offset += 4
|
||||
if offset < target: break
|
||||
|
||||
opcode = state_to_opcode[reg & 0xFF]
|
||||
pkt_cls, nib_count, delta_lo, delta_mask, special = decode_info[opcode]
|
||||
|
||||
delta = (reg >> delta_lo) & delta_mask
|
||||
|
||||
if special == 1: # TS_DELTA_OR_MARK
|
||||
bit8 = (reg >> _TS_DELTA_OR_MARK_BIT8[0]) & _TS_DELTA_OR_MARK_BIT8[1]
|
||||
bit9 = (reg >> _TS_DELTA_OR_MARK_BIT9[0]) & _TS_DELTA_OR_MARK_BIT9[1]
|
||||
if bit9 and not bit8: delta = 0
|
||||
elif special == 2: # TS_DELTA_SHORT
|
||||
delta = delta + 8
|
||||
|
||||
time += delta
|
||||
packets_append(pkt_cls.from_raw(reg, time))
|
||||
|
||||
return packets
|
||||
@@ -0,0 +1,192 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Benchmark comparing Python vs Rust RDNA3 emulators on real tinygrad kernels."""
|
||||
import ctypes, time, os
|
||||
from pathlib import Path
|
||||
|
||||
# Set AMD=1 before importing tinygrad
|
||||
os.environ["AMD"] = "1"
|
||||
|
||||
from extra.assembly.amd.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program
|
||||
|
||||
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
|
||||
if not REMU_PATH.exists():
|
||||
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.dylib"
|
||||
|
||||
def get_rust_remu():
|
||||
"""Load the Rust libremu shared library."""
|
||||
if not REMU_PATH.exists(): return None
|
||||
remu = ctypes.CDLL(str(REMU_PATH))
|
||||
remu.run_asm.restype = ctypes.c_int32
|
||||
remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
|
||||
ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
|
||||
return remu
|
||||
|
||||
def count_instructions(kernel: bytes) -> int:
|
||||
"""Count instructions in a kernel."""
|
||||
return len(decode_program(kernel))
|
||||
|
||||
def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = None):
|
||||
"""Allocate buffers and return args pointer + valid ranges."""
|
||||
if init_data is None: init_data = {}
|
||||
buffers = []
|
||||
for i, size in enumerate(buf_sizes):
|
||||
padded = ((size + 15) // 16) * 16 + 16
|
||||
data = init_data.get(i, b'\x00' * padded)
|
||||
data_list = list(data) + [0] * (padded - len(data))
|
||||
buf = (ctypes.c_uint8 * padded)(*data_list[:padded])
|
||||
buffers.append(buf)
|
||||
args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
|
||||
args_ptr = ctypes.addressof(args)
|
||||
ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
|
||||
ranges.add((args_ptr, ctypes.sizeof(args)))
|
||||
return buffers, args, args_ptr, ranges
|
||||
|
||||
def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, rsrc2: int, iterations: int = 5):
|
||||
"""Benchmark an emulator and return average time."""
|
||||
gx, gy, gz = global_size
|
||||
lx, ly, lz = local_size
|
||||
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
|
||||
lib_ptr = ctypes.addressof(kernel_buf)
|
||||
|
||||
# Warmup
|
||||
run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
|
||||
|
||||
# Timed runs
|
||||
times = []
|
||||
for _ in range(iterations):
|
||||
start = time.perf_counter()
|
||||
result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
|
||||
end = time.perf_counter()
|
||||
if result != 0:
|
||||
print(f" {name} returned error: {result}")
|
||||
return None
|
||||
times.append(end - start)
|
||||
|
||||
return sum(times) / len(times)
|
||||
|
||||
def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes], int] | None:
|
||||
"""Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data, rsrc2)."""
|
||||
try:
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.runtime.support.elf import elf_loader
|
||||
from tinygrad.runtime.autogen import hsa
|
||||
import numpy as np
|
||||
np.random.seed(42)
|
||||
|
||||
ops = {
|
||||
"add": lambda: Tensor.empty(1024) + Tensor.empty(1024),
|
||||
"mul": lambda: Tensor.empty(1024) * Tensor.empty(1024),
|
||||
"matmul_small": lambda: Tensor.empty(16, 16) @ Tensor.empty(16, 16),
|
||||
"matmul_medium": lambda: Tensor.empty(64, 64) @ Tensor.empty(64, 64),
|
||||
"reduce_sum": lambda: Tensor.empty(4096).sum(),
|
||||
"reduce_max": lambda: Tensor.empty(4096).max(),
|
||||
"softmax": lambda: Tensor.empty(256).softmax(),
|
||||
"layernorm": lambda: Tensor.empty(32, 64).layernorm(),
|
||||
"conv2d": lambda: Tensor.empty(1, 4, 16, 16).conv2d(Tensor.empty(4, 4, 3, 3)),
|
||||
"gelu": lambda: Tensor.empty(1024).gelu(),
|
||||
"exp": lambda: Tensor.empty(1024).exp(),
|
||||
"sin": lambda: Tensor.empty(1024).sin(),
|
||||
}
|
||||
|
||||
if op_name not in ops: return None
|
||||
out = ops[op_name]()
|
||||
sched = out.schedule()
|
||||
|
||||
for ei in sched:
|
||||
lowered = ei.lower()
|
||||
if ei.ast.op.name == 'SINK' and lowered.prg and lowered.prg.p.lib:
|
||||
lib = bytes(lowered.prg.p.lib)
|
||||
image = memoryview(bytearray(lib))
|
||||
_, sections, _ = elf_loader(lib)
|
||||
rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
|
||||
for sec in sections:
|
||||
if sec.name == '.text':
|
||||
buf_sizes = [b.nbytes for b in lowered.bufs]
|
||||
# Get initial data from numpy arrays if available
|
||||
buf_data = {}
|
||||
for i, buf in enumerate(lowered.bufs):
|
||||
if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
|
||||
try: buf_data[i] = bytes(buf.base._buf)
|
||||
except: pass
|
||||
# Extract rsrc2 from ELF (same as ops_amd.py)
|
||||
group_segment_size = image[rodata_entry:rodata_entry+4].cast("I")[0]
|
||||
lds_size = ((group_segment_size + 511) // 512) & 0x1FF
|
||||
code = hsa.amd_kernel_code_t.from_buffer_copy(bytes(image[rodata_entry:rodata_entry+256]) + b'\x00'*256)
|
||||
rsrc2 = code.compute_pgm_rsrc2 | (lds_size << 15)
|
||||
return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data, rsrc2)
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f" Error getting kernel: {e}")
|
||||
return None
|
||||
|
||||
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
|
||||
|
||||
def main():
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
|
||||
parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
|
||||
args = parser.parse_args()
|
||||
|
||||
rust_remu = get_rust_remu()
|
||||
if rust_remu is None:
|
||||
print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
|
||||
print("Running Python-only benchmarks...\n")
|
||||
|
||||
print("=" * 90)
|
||||
print("RDNA3 Emulator Benchmark: Python vs Rust")
|
||||
print("=" * 90)
|
||||
|
||||
results = []
|
||||
|
||||
print("\n[TINYGRAD KERNELS]")
|
||||
print("-" * 90)
|
||||
|
||||
for op_name in TINYGRAD_TESTS:
|
||||
print(f"\n{op_name}:", end=" ", flush=True)
|
||||
kernel_info = get_tinygrad_kernel(op_name)
|
||||
if kernel_info is None:
|
||||
print("failed to compile")
|
||||
continue
|
||||
|
||||
kernel, global_size, local_size, buf_sizes, buf_data, rsrc2 = kernel_info
|
||||
n_insts = count_instructions(kernel)
|
||||
n_workgroups = global_size[0] * global_size[1] * global_size[2]
|
||||
n_threads = local_size[0] * local_size[1] * local_size[2]
|
||||
total_work = n_insts * n_workgroups * n_threads
|
||||
|
||||
print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
|
||||
|
||||
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
|
||||
set_valid_mem_ranges(ranges)
|
||||
|
||||
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
|
||||
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations) if rust_remu else None
|
||||
|
||||
if py_time:
|
||||
py_rate = total_work / py_time / 1e6
|
||||
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
|
||||
if rust_time:
|
||||
rust_rate = total_work / rust_time / 1e6
|
||||
speedup = py_time / rust_time if py_time else 0
|
||||
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
|
||||
|
||||
results.append((op_name, n_insts, n_workgroups, py_time, rust_time))
|
||||
|
||||
# Summary table
|
||||
print("\n" + "=" * 90)
|
||||
print("SUMMARY")
|
||||
print("=" * 90)
|
||||
print(f"{'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
|
||||
print("-" * 90)
|
||||
|
||||
for name, n_insts, n_wgs, py_time, rust_time in results:
|
||||
py_ms = f"{py_time*1000:.3f}" if py_time else "error"
|
||||
if rust_time:
|
||||
rust_ms = f"{rust_time*1000:.3f}"
|
||||
speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
|
||||
else:
|
||||
rust_ms, speedup = "N/A", "N/A"
|
||||
print(f"{name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,196 @@
|
||||
# Usability tests for the RDNA3 ASM DSL
|
||||
# These tests demonstrate how the DSL *should* work for a good user experience
|
||||
# Currently many of these tests fail - they document desired behavior
|
||||
|
||||
import unittest
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
from extra.assembly.amd.dsl import Inst, RawImm, SGPR, VGPR
|
||||
|
||||
class TestRegisterSliceSyntax(unittest.TestCase):
|
||||
"""
|
||||
Issue: Register slice syntax should use AMD assembly convention (inclusive end).
|
||||
|
||||
In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
|
||||
The DSL should match this convention so that:
|
||||
- s[4:7] gives 4 registers
|
||||
- Disassembler output can be copied directly back into DSL code
|
||||
|
||||
Fix: Change _RegFactory.__getitem__ to use inclusive end:
|
||||
key.stop - key.start + 1 (instead of key.stop - key.start)
|
||||
"""
|
||||
def test_register_slice_count(self):
|
||||
# s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
|
||||
reg = s[4:7]
|
||||
self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
|
||||
|
||||
def test_register_slice_roundtrip(self):
|
||||
# Round-trip: DSL -> disasm -> DSL should preserve register count
|
||||
reg = s[4:7] # 4 registers in AMD convention
|
||||
inst = s_load_b128(reg, s[0:1], NULL, 0)
|
||||
disasm = inst.disasm()
|
||||
# Disasm shows s[4:7] - user should be able to copy this back
|
||||
self.assertIn("s[4:7]", disasm)
|
||||
# And s[4:7] in DSL should give the same 4 registers
|
||||
reg_from_disasm = s[4:7]
|
||||
self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
|
||||
|
||||
|
||||
class TestReprReadability(unittest.TestCase):
|
||||
"""
|
||||
Issue: repr() leaks internal RawImm type and omits zero-valued fields.
|
||||
|
||||
When you create v_mov_b32_e32(v[0], v[1]), the repr shows:
|
||||
VOP1(op=1, src0=RawImm(257))
|
||||
|
||||
Problems:
|
||||
1. vdst=v[0] is omitted because 0 is treated as "default"
|
||||
2. src0 shows RawImm(257) instead of v[1]
|
||||
3. User sees encoded values (257 = 256 + 1) instead of register names
|
||||
|
||||
Expected repr: VOP1(op=1, vdst=v[0], src0=v[1])
|
||||
"""
|
||||
def test_repr_shows_registers_not_raw_imm(self):
|
||||
inst = v_mov_b32_e32(v[0], v[1])
|
||||
# Should show v[1], not RawImm(257)
|
||||
self.assertNotIn("RawImm", repr(inst), "repr should not expose RawImm internal type")
|
||||
self.assertIn("v[1]", repr(inst), "repr should show register name")
|
||||
|
||||
def test_repr_includes_zero_dst(self):
|
||||
inst = v_mov_b32_e32(v[0], v[1])
|
||||
# v[0] is a valid destination register, should be shown
|
||||
self.assertIn("vdst", repr(inst), "repr should include vdst even when 0")
|
||||
|
||||
def test_repr_roundtrip(self):
|
||||
# repr should produce something that can be eval'd back
|
||||
inst = v_mov_b32_e32(v[0], v[1])
|
||||
# This would require repr to output valid Python, e.g.:
|
||||
# "VOP1(op=VOP1Op.V_MOV_B32, vdst=v[0], src0=v[1])"
|
||||
r = repr(inst)
|
||||
# At minimum, it should be human-readable
|
||||
self.assertIn("v[", r, "repr should show register syntax")
|
||||
|
||||
|
||||
class TestInstructionEquality(unittest.TestCase):
|
||||
"""
|
||||
Issue: No __eq__ method - instruction comparison requires repr() workaround.
|
||||
|
||||
Two identical instructions should compare equal with ==, but currently:
|
||||
inst1 == inst2 returns False
|
||||
|
||||
The test_handwritten.py works around this with:
|
||||
self.assertEqual(repr(self.inst), repr(reasm))
|
||||
"""
|
||||
def test_identical_instructions_equal(self):
|
||||
inst1 = v_mov_b32_e32(v[0], v[1])
|
||||
inst2 = v_mov_b32_e32(v[0], v[1])
|
||||
self.assertEqual(inst1, inst2, "identical instructions should be equal")
|
||||
|
||||
def test_different_instructions_not_equal(self):
|
||||
inst1 = v_mov_b32_e32(v[0], v[1])
|
||||
inst2 = v_mov_b32_e32(v[0], v[2])
|
||||
self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
|
||||
|
||||
|
||||
class TestVOPDHelperSignature(unittest.TestCase):
|
||||
"""
|
||||
Issue: VOPD helper functions have confusing semantics.
|
||||
|
||||
v_dual_mul_f32 is defined as:
|
||||
v_dual_mul_f32 = functools.partial(VOPD, VOPDOp.V_DUAL_MUL_F32)
|
||||
|
||||
This binds VOPDOp.V_DUAL_MUL_F32 to the FIRST positional arg of VOPD.__init__,
|
||||
which is 'opx'. So v_dual_mul_f32 sets the X operation.
|
||||
|
||||
But then test_dual_mul in test_handwritten.py does:
|
||||
v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], ...)
|
||||
|
||||
This passes V_DUAL_MUL_F32 as the SECOND positional arg (opy), making both
|
||||
X and Y operations the same. This is confusing because:
|
||||
1. The function name suggests it handles the X operation
|
||||
2. But you still pass an opcode as the first arg (which becomes opy)
|
||||
|
||||
Expected: Either make the helper fully specify both ops, or make the
|
||||
signature clearer about what the positional arg means.
|
||||
"""
|
||||
def test_vopd_helper_opy_should_be_required(self):
|
||||
# Using only keyword args "works" but opy silently defaults to 0
|
||||
inst = v_dual_mul_f32(vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
|
||||
self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32)
|
||||
# Bug: opy defaults to 0 (V_DUAL_FMAC_F32) silently - should require explicit opy
|
||||
# This test documents the bug - it should fail once fixed
|
||||
self.assertNotEqual(inst.opy, VOPDOp.V_DUAL_FMAC_F32, "opy should not silently default to FMAC")
|
||||
|
||||
def test_vopd_helper_positional_arg_is_opy(self):
|
||||
# The first positional arg after the partial becomes opy, not a second opx
|
||||
inst = v_dual_mul_f32(VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
|
||||
self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32) # From partial
|
||||
self.assertEqual(inst.opy, VOPDOp.V_DUAL_MOV_B32) # From first positional arg
|
||||
|
||||
|
||||
class TestFieldAccessPreservesType(unittest.TestCase):
|
||||
"""
|
||||
Issue: Field access loses type information.
|
||||
|
||||
After creating an instruction, accessing fields returns encoded int values:
|
||||
inst = v_mov_b32_e32(v[0], v[1])
|
||||
inst.vdst # returns 0, not VGPR(0)
|
||||
|
||||
This makes it impossible to round-trip register types through field access.
|
||||
"""
|
||||
def test_vdst_returns_register(self):
|
||||
inst = v_mov_b32_e32(v[5], v[1])
|
||||
vdst = inst.vdst
|
||||
# Should return a VGPR, not an int
|
||||
self.assertIsInstance(vdst, (VGPR, int), "vdst should return VGPR or at least be usable")
|
||||
# Ideally: self.assertIsInstance(vdst, VGPR)
|
||||
|
||||
def test_src_returns_register_for_vgpr_source(self):
|
||||
inst = v_mov_b32_e32(v[0], v[1])
|
||||
# src0 is encoded as 257 (256 + 1 for v1)
|
||||
# Ideally it should decode back to v[1]
|
||||
src0_raw = inst._values.get('src0')
|
||||
# Currently returns RawImm(257), should return VGPR(1) or similar
|
||||
self.assertNotIsInstance(src0_raw, RawImm, "source should not be RawImm internally")
|
||||
|
||||
|
||||
class TestArgumentDiscoverability(unittest.TestCase):
|
||||
"""
|
||||
Issue: No clear signature for positional arguments.
|
||||
|
||||
inspect.signature(s_load_b128) shows: (*args, literal=None, **kwargs)
|
||||
|
||||
Users have no way to know the argument order without reading source code.
|
||||
The order is implicitly defined by the class field definition order.
|
||||
|
||||
Possible fixes:
|
||||
1. Add explicit parameter names to functools.partial
|
||||
2. Generate type stubs with proper signatures
|
||||
3. Add docstrings listing the expected arguments
|
||||
"""
|
||||
def test_signature_has_named_params(self):
|
||||
import inspect
|
||||
sig = inspect.signature(s_load_b128)
|
||||
params = list(sig.parameters.keys())
|
||||
# Currently: ['args', 'literal', 'kwargs'] (from *args, literal=None, **kwargs)
|
||||
# Expected: something like ['sdata', 'sbase', 'soffset', 'offset', 'literal']
|
||||
self.assertIn('sdata', params, "signature should show field names")
|
||||
|
||||
|
||||
class TestSpecialConstants(unittest.TestCase):
|
||||
"""
|
||||
Issue: NULL and other constants are IntEnum values that might be confusing.
|
||||
|
||||
NULL = SrcEnum.NULL = 124, but users might expect NULL to be a special object
|
||||
that clearly represents "no register" rather than a magic number.
|
||||
"""
|
||||
def test_null_has_clear_repr(self):
|
||||
# NULL should have a clear string representation
|
||||
self.assertIn("NULL", str(NULL) or repr(NULL), "NULL should be clearly identifiable")
|
||||
|
||||
def test_null_is_distinguishable_from_int(self):
|
||||
# NULL should be distinguishable from the raw integer 124
|
||||
self.assertNotEqual(type(NULL), int, "NULL should not be plain int")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,66 @@
|
||||
"""Shared test helpers for RDNA3 tests."""
|
||||
import shutil
|
||||
from dataclasses import dataclass
|
||||
|
||||
@dataclass
|
||||
class KernelInfo:
|
||||
code: bytes
|
||||
global_size: tuple[int, int, int]
|
||||
local_size: tuple[int, int, int]
|
||||
buf_idxs: list[int] # indices into shared buffer pool
|
||||
buf_sizes: list[int] # sizes for each buffer index
|
||||
|
||||
# LLVM tool detection (shared across test files)
|
||||
def get_llvm_mc():
|
||||
"""Find llvm-mc executable, preferring newer versions."""
|
||||
for p in ['llvm-mc', 'llvm-mc-21', 'llvm-mc-20']:
|
||||
if shutil.which(p): return p
|
||||
raise FileNotFoundError("llvm-mc not found")
|
||||
|
||||
def get_llvm_objdump():
|
||||
"""Find llvm-objdump executable, preferring newer versions."""
|
||||
for p in ['llvm-objdump', 'llvm-objdump-21', 'llvm-objdump-20']:
|
||||
if shutil.which(p): return p
|
||||
raise FileNotFoundError("llvm-objdump not found")
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# EXECUTION CONTEXT (for testing compiled pseudocode)
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class ExecContext:
|
||||
"""Context for running compiled pseudocode in tests."""
|
||||
def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=0xffffffff, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
|
||||
from extra.assembly.amd.pcode import Reg, MASK32, MASK64, TypedView
|
||||
self._Reg, self._MASK64, self._TypedView = Reg, MASK64, TypedView
|
||||
self.S0, self.S1, self.S2 = Reg(s0), Reg(s1), Reg(s2)
|
||||
self.D0, self.D1 = Reg(d0), Reg(0)
|
||||
self.SCC, self.VCC, self.EXEC = Reg(scc), Reg(vcc), Reg(exec_mask)
|
||||
self.tmp, self.saveexec = Reg(0), Reg(exec_mask)
|
||||
self.lane, self.laneId, self.literal = lane, lane, literal
|
||||
self.SIMM16, self.SIMM32 = Reg(literal), Reg(literal)
|
||||
self.VGPR = vgprs if vgprs is not None else {}
|
||||
self.SRC0, self.VDST = Reg(src0_idx), Reg(vdst_idx)
|
||||
|
||||
def run(self, code: str):
|
||||
"""Execute compiled code."""
|
||||
import extra.assembly.amd.pcode as pcode
|
||||
ns = {k: getattr(pcode, k) for k in dir(pcode) if not k.startswith('_')}
|
||||
# Also include underscore-prefixed helpers that compiled pseudocode uses
|
||||
for k in ['_pack', '_pack32']:
|
||||
if hasattr(pcode, k): ns[k] = getattr(pcode, k)
|
||||
ns.update({
|
||||
'S0': self.S0, 'S1': self.S1, 'S2': self.S2, 'D0': self.D0, 'D1': self.D1,
|
||||
'SCC': self.SCC, 'VCC': self.VCC, 'EXEC': self.EXEC,
|
||||
'EXEC_LO': self._TypedView(self.EXEC, 31, 0), 'EXEC_HI': self._TypedView(self.EXEC, 63, 32),
|
||||
'tmp': self.tmp, 'saveexec': self.saveexec,
|
||||
'lane': self.lane, 'laneId': self.laneId, 'literal': self.literal,
|
||||
'SIMM16': self.SIMM16, 'SIMM32': self.SIMM32, 'VGPR': self.VGPR, 'SRC0': self.SRC0, 'VDST': self.VDST,
|
||||
})
|
||||
exec(code, ns)
|
||||
def _sync(ctx_reg, ns_val):
|
||||
if isinstance(ns_val, self._Reg): ctx_reg._val = ns_val._val
|
||||
else: ctx_reg._val = int(ns_val) & self._MASK64
|
||||
for name in ('SCC', 'VCC', 'EXEC', 'D0', 'D1', 'tmp', 'saveexec'):
|
||||
if ns.get(name) is not getattr(self, name): _sync(getattr(self, name), ns[name])
|
||||
|
||||
def result(self) -> dict: return {"d0": self.D0._val, "scc": self.SCC._val & 1}
|
||||
@@ -0,0 +1 @@
|
||||
"""Hardware-validated emulator tests for RDNA3 instructions."""
|
||||
@@ -0,0 +1,202 @@
|
||||
"""Test infrastructure for hardware-validated RDNA3 emulator tests.
|
||||
|
||||
Uses run_asm() with memory output, so tests can run on both emulator and real hardware.
|
||||
Set USE_HW=1 to run on both emulator and real hardware, comparing results.
|
||||
"""
|
||||
import ctypes, os, struct
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
from extra.assembly.amd.dsl import RawImm
|
||||
from extra.assembly.amd.emu import WaveState, run_asm, set_valid_mem_ranges
|
||||
from extra.assembly.amd.pcode import _i32, _f32
|
||||
|
||||
VCC = SrcEnum.VCC_LO # For VOP3SD sdst field
|
||||
USE_HW = os.environ.get("USE_HW", "0") == "1"
|
||||
FLOAT_TOLERANCE = 1e-5
|
||||
|
||||
# Output buffer layout: vgpr[16][32], sgpr[16], vcc, scc
|
||||
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
|
||||
OUT_BYTES = VGPR_BYTES + SGPR_BYTES + 8 # + vcc + scc
|
||||
|
||||
# Float conversion helpers
|
||||
def f2i(f: float) -> int: return _i32(f)
|
||||
def i2f(i: int) -> float: return _f32(i)
|
||||
def f2i64(f: float) -> int: return struct.unpack('<Q', struct.pack('<d', f))[0]
|
||||
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)
|
||||
|
||||
def get_prologue_epilogue(n_lanes: int) -> tuple[list, list]:
|
||||
"""Generate prologue and epilogue instructions for state capture."""
|
||||
prologue = [
|
||||
s_mov_b32(s[80], s[0]),
|
||||
s_mov_b32(s[81], s[1]),
|
||||
v_mov_b32_e32(v[255], v[0]),
|
||||
]
|
||||
for i in range(N_VGPRS):
|
||||
prologue.append(v_mov_b32_e32(v[i], 0))
|
||||
for i in range(N_SGPRS):
|
||||
prologue.append(s_mov_b32(s[i], 0))
|
||||
prologue.append(s_mov_b32(s[SrcEnum.VCC_LO - 128], 0))
|
||||
|
||||
epilogue = [
|
||||
s_mov_b32(s[90], SrcEnum.VCC_LO),
|
||||
s_cselect_b32(s[91], 1, 0),
|
||||
s_load_b64(s[92:93], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_lshlrev_b32_e32(v[240], 2, v[255]),
|
||||
]
|
||||
for i in range(N_VGPRS):
|
||||
epilogue.append(global_store_b32(addr=v[240], data=v[i], saddr=s[92], 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], SrcEnum.VCC_LO))
|
||||
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], 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], 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], offset=VGPR_BYTES + SGPR_BYTES + 4))
|
||||
epilogue.append(s_mov_b32(s[SrcEnum.EXEC_LO - 128], s[94]))
|
||||
epilogue.append(s_endpgm())
|
||||
return prologue, epilogue
|
||||
|
||||
def parse_output(out_buf: bytes, n_lanes: int) -> WaveState:
|
||||
"""Parse output buffer into WaveState."""
|
||||
st = WaveState()
|
||||
for i in range(N_VGPRS):
|
||||
for lane in range(n_lanes):
|
||||
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]
|
||||
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."""
|
||||
out_buf = (ctypes.c_uint8 * OUT_BYTES)(*([0] * OUT_BYTES))
|
||||
out_addr = ctypes.addressof(out_buf)
|
||||
|
||||
prologue, epilogue = get_prologue_epilogue(n_lanes)
|
||||
code = assemble(prologue + instructions + epilogue)
|
||||
|
||||
args = (ctypes.c_uint64 * 1)(out_addr)
|
||||
args_ptr = ctypes.addressof(args)
|
||||
kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
|
||||
lib_ptr = ctypes.addressof(kernel_buf)
|
||||
|
||||
set_valid_mem_ranges({(out_addr, OUT_BYTES), (args_ptr, 8)})
|
||||
# rsrc2: USER_SGPR_COUNT=2, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z=1, LDS_SIZE=128 (64KB)
|
||||
rsrc2 = 0x19c | (128 << 15)
|
||||
result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2)
|
||||
assert result == 0, f"run_asm failed with {result}"
|
||||
|
||||
return parse_output(bytes(out_buf), n_lanes)
|
||||
|
||||
def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
|
||||
"""Run instructions on real AMD hardware via HIPCompiler and AMDProgram."""
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.runtime.ops_amd import AMDProgram
|
||||
from tinygrad.runtime.support.compiler_amd import HIPCompiler
|
||||
from tinygrad.helpers import flat_mv
|
||||
|
||||
dev = Device["AMD"]
|
||||
compiler = HIPCompiler(dev.arch)
|
||||
|
||||
prologue, epilogue = get_prologue_epilogue(n_lanes)
|
||||
code = assemble(prologue + instructions + epilogue)
|
||||
|
||||
byte_str = ', '.join(f'0x{b:02x}' for b in code)
|
||||
asm_src = f""".text
|
||||
.globl test
|
||||
.p2align 8
|
||||
.type test,@function
|
||||
test:
|
||||
.byte {byte_str}
|
||||
|
||||
.rodata
|
||||
.p2align 6
|
||||
.amdhsa_kernel test
|
||||
.amdhsa_next_free_vgpr 256
|
||||
.amdhsa_next_free_sgpr 96
|
||||
.amdhsa_wavefront_size32 1
|
||||
.amdhsa_user_sgpr_kernarg_segment_ptr 1
|
||||
.amdhsa_kernarg_size 8
|
||||
.amdhsa_group_segment_fixed_size 65536
|
||||
.end_amdhsa_kernel
|
||||
|
||||
.amdgpu_metadata
|
||||
---
|
||||
amdhsa.version:
|
||||
- 1
|
||||
- 0
|
||||
amdhsa.kernels:
|
||||
- .name: test
|
||||
.symbol: test.kd
|
||||
.kernarg_segment_size: 8
|
||||
.group_segment_fixed_size: 65536
|
||||
.private_segment_fixed_size: 0
|
||||
.kernarg_segment_align: 8
|
||||
.wavefront_size: 32
|
||||
.sgpr_count: 96
|
||||
.vgpr_count: 256
|
||||
.max_flat_workgroup_size: 1024
|
||||
...
|
||||
.end_amdgpu_metadata
|
||||
"""
|
||||
|
||||
lib = compiler.compile(asm_src)
|
||||
prg = AMDProgram(dev, "test", lib)
|
||||
|
||||
out_gpu = dev.allocator.alloc(OUT_BYTES)
|
||||
prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
|
||||
|
||||
out_buf = bytearray(OUT_BYTES)
|
||||
dev.allocator._copyout(flat_mv(memoryview(out_buf)), out_gpu)
|
||||
|
||||
return parse_output(bytes(out_buf), n_lanes)
|
||||
|
||||
def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS) -> list[str]:
|
||||
"""Compare two WaveStates and return list of differences."""
|
||||
import math
|
||||
diffs = []
|
||||
for i in range(n_vgprs):
|
||||
for lane in range(n_lanes):
|
||||
emu_val = emu_st.vgpr[lane][i]
|
||||
hw_val = hw_st.vgpr[lane][i]
|
||||
if emu_val != hw_val:
|
||||
emu_f, hw_f = _f32(emu_val), _f32(hw_val)
|
||||
if math.isnan(emu_f) and math.isnan(hw_f):
|
||||
continue
|
||||
diffs.append(f"v[{i}] lane {lane}: emu=0x{emu_val:08x} ({emu_f:.6g}) hw=0x{hw_val:08x} ({hw_f:.6g})")
|
||||
for i in range(N_SGPRS):
|
||||
emu_val = emu_st.sgpr[i]
|
||||
hw_val = hw_st.sgpr[i]
|
||||
if emu_val != hw_val:
|
||||
diffs.append(f"s[{i}]: emu=0x{emu_val:08x} hw=0x{hw_val:08x}")
|
||||
if emu_st.vcc != hw_st.vcc:
|
||||
diffs.append(f"vcc: emu=0x{emu_st.vcc:08x} hw=0x{hw_st.vcc:08x}")
|
||||
if emu_st.scc != hw_st.scc:
|
||||
diffs.append(f"scc: emu={emu_st.scc} hw={hw_st.scc}")
|
||||
return diffs
|
||||
|
||||
def run_program(instructions: list, n_lanes: int = 1) -> WaveState:
|
||||
"""Run instructions and return WaveState.
|
||||
|
||||
If USE_HW=1, runs on both emulator and hardware, compares results, and raises if they differ.
|
||||
Otherwise, runs only on emulator.
|
||||
"""
|
||||
emu_st = run_program_emu(instructions, n_lanes)
|
||||
if USE_HW:
|
||||
hw_st = run_program_hw(instructions, n_lanes)
|
||||
diffs = compare_wave_states(emu_st, hw_st, n_lanes)
|
||||
if diffs:
|
||||
raise AssertionError(f"Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
|
||||
return hw_st
|
||||
return emu_st
|
||||
@@ -0,0 +1,629 @@
|
||||
"""Tests for DS instructions - data share (LDS) operations.
|
||||
|
||||
Includes: ds_store_b32, ds_load_b32, ds_store_2addr_*, ds_load_2addr_*,
|
||||
ds_add_*, ds_max_*, ds_min_*, ds_and_*, ds_or_*, ds_xor_*,
|
||||
ds_inc_*, ds_dec_*, ds_cmpstore_*, ds_storexchg_*
|
||||
"""
|
||||
import unittest
|
||||
from extra.assembly.amd.test.hw.helpers import *
|
||||
|
||||
class TestDS2Addr(unittest.TestCase):
|
||||
"""Tests for DS_*_2ADDR instructions."""
|
||||
|
||||
def test_ds_store_load_2addr_b32(self):
|
||||
"""DS_STORE_2ADDR_B32 and DS_LOAD_2ADDR_B32 with offset * 4."""
|
||||
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]),
|
||||
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2], offset0=0, offset1=1),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
|
||||
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
|
||||
|
||||
def test_ds_store_load_2addr_b64(self):
|
||||
"""DS_STORE_2ADDR_B64 and DS_LOAD_2ADDR_B64."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[0], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[0], 0xCAFEBABE),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
s_mov_b32(s[0], 0x12345678),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0x9ABCDEF0),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
DS(DSOp.DS_STORE_2ADDR_B64, addr=v[10], data0=v[0], data1=v[2], vdst=v[0], offset0=0, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4], offset0=0, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
|
||||
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
|
||||
self.assertEqual(st.vgpr[0][6], 0x12345678)
|
||||
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
|
||||
|
||||
|
||||
class TestDS2AddrMore(unittest.TestCase):
|
||||
"""Additional DS_*_2ADDR tests."""
|
||||
|
||||
def test_ds_store_load_2addr_b32_nonzero_offsets(self):
|
||||
"""DS_STORE_2ADDR_B32 with non-zero offsets (offset*4 scaling)."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 0x11111111),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0x22222222),
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=2, offset1=5),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2], offset0=2, offset1=5),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have value from offset 8 (2*4)")
|
||||
self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have value from offset 20 (5*4)")
|
||||
|
||||
def test_ds_2addr_b64_no_overlap(self):
|
||||
"""DS_LOAD_2ADDR_B64 with adjacent offsets should not overlap."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 0x11111111),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_mov_b32(s[2], 0x22222222),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=4),
|
||||
s_mov_b32(s[2], 0x33333333),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=8),
|
||||
s_mov_b32(s[2], 0x44444444),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=12),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4], offset0=0, offset1=1),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should be 0x11111111")
|
||||
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should be 0x22222222")
|
||||
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should be 0x33333333")
|
||||
self.assertEqual(st.vgpr[0][7], 0x44444444, "v7 should be 0x44444444")
|
||||
|
||||
def test_ds_load_2addr_b32_no_overwrite(self):
|
||||
"""DS_LOAD_2ADDR_B32 should only write 2 VGPRs."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[4], s[2]), # Sentinel
|
||||
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2], offset0=0, offset1=1),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
|
||||
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
|
||||
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should be untouched")
|
||||
|
||||
def test_ds_load_b64_no_overwrite(self):
|
||||
"""DS_LOAD_B64 should only write 2 VGPRs."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0xCAFEBABE),
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_store_b64(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 0x12345678),
|
||||
v_mov_b32_e32(v[4], s[2]), # Sentinel
|
||||
ds_load_b64(addr=v[10], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
|
||||
self.assertEqual(st.vgpr[0][3], 0xCAFEBABE)
|
||||
self.assertEqual(st.vgpr[0][4], 0x12345678, "v4 should be untouched")
|
||||
|
||||
|
||||
class TestDSAtomic(unittest.TestCase):
|
||||
"""Tests for DS atomic operations."""
|
||||
|
||||
def test_ds_max_rtn_u32(self):
|
||||
"""DS_MAX_RTN_U32: atomically store max and return old value."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 100),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 200),
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_max_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
|
||||
self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
|
||||
|
||||
def test_ds_min_rtn_u32(self):
|
||||
"""DS_MIN_RTN_U32: atomically store min and return old value."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 200),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 100),
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_min_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 200)
|
||||
self.assertEqual(st.vgpr[0][3], 100)
|
||||
|
||||
def test_ds_and_rtn_b32(self):
|
||||
"""DS_AND_RTN_B32: atomically AND and return old value."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 0xFF00FF00),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 0xFFFF0000),
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_and_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0xFF00FF00)
|
||||
self.assertEqual(st.vgpr[0][3], 0xFF000000)
|
||||
|
||||
def test_ds_or_rtn_b32(self):
|
||||
"""DS_OR_RTN_B32: atomically OR and return old value."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 0x00FF0000),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 0x000000FF),
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_or_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0x00FF0000)
|
||||
self.assertEqual(st.vgpr[0][3], 0x00FF00FF)
|
||||
|
||||
def test_ds_xor_rtn_b32(self):
|
||||
"""DS_XOR_RTN_B32: atomically XOR and return old value."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 0xFFFFFFFF),
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_xor_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
|
||||
self.assertEqual(st.vgpr[0][3], 0x55555555)
|
||||
|
||||
def test_ds_inc_rtn_u32(self):
|
||||
"""DS_INC_RTN_U32: increment with wrap."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 5),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 10), # limit
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_inc_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 5)
|
||||
self.assertEqual(st.vgpr[0][3], 6)
|
||||
|
||||
def test_ds_dec_rtn_u32(self):
|
||||
"""DS_DEC_RTN_U32: decrement with wrap."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 5),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 10), # limit
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 5)
|
||||
self.assertEqual(st.vgpr[0][3], 4)
|
||||
|
||||
def test_ds_cmpstore_b32_match(self):
|
||||
"""DS_CMPSTORE_B32: conditional store when compare matches."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 100),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 200),
|
||||
v_mov_b32_e32(v[1], s[2]), # new value
|
||||
s_mov_b32(s[2], 100),
|
||||
v_mov_b32_e32(v[2], s[2]), # compare = 100 (matches)
|
||||
ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][4], 200)
|
||||
|
||||
def test_ds_cmpstore_b32_no_match(self):
|
||||
"""DS_CMPSTORE_B32: no store when compare doesn't match."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 100),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 200),
|
||||
v_mov_b32_e32(v[1], s[2]), # new value
|
||||
s_mov_b32(s[2], 50),
|
||||
v_mov_b32_e32(v[2], s[2]), # compare = 50 (doesn't match)
|
||||
ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][4], 100)
|
||||
|
||||
def test_ds_max_u32_no_rtn(self):
|
||||
"""DS_MAX_U32 (no RTN): atomically store max, no return value."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 100),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 200),
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_max_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
|
||||
|
||||
def test_ds_add_u32_no_rtn_preserves_vdst(self):
|
||||
"""DS_ADD_U32 (no RTN) should NOT write to vdst."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[2], s[2]), # sentinel
|
||||
s_mov_b32(s[2], 100),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 50),
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_add_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF, "v2 should preserve sentinel")
|
||||
self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
|
||||
|
||||
def test_ds_add_rtn_u32_writes_vdst(self):
|
||||
"""DS_ADD_RTN_U32 should write old value to vdst."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[2], s[2]), # sentinel
|
||||
s_mov_b32(s[2], 100),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 50),
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_add_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
|
||||
self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
|
||||
|
||||
def test_ds_dec_rtn_u32_wrap(self):
|
||||
"""DS_DEC_RTN_U32: decrement wraps when value is 0 or > limit."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[2], 0), # Start at 0
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[2], 10), # limit
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0, "v2 should have old value (0)")
|
||||
# When mem == 0 or mem > limit, result = limit
|
||||
self.assertEqual(st.vgpr[0][3], 10, "v3 should wrap to limit (10)")
|
||||
|
||||
|
||||
class TestDSStorexchg(unittest.TestCase):
|
||||
"""Tests for DS_STOREXCHG instructions."""
|
||||
|
||||
def test_ds_storexchg_rtn_b32(self):
|
||||
"""DS_STOREXCHG_RTN_B32: exchange value and return old."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[0], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[0], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
DS(DSOp.DS_STOREXCHG_RTN_B32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
|
||||
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
|
||||
|
||||
|
||||
class TestDSRegisterWidth(unittest.TestCase):
|
||||
"""Regression tests: DS loads should only write correct number of VGPRs."""
|
||||
|
||||
def test_ds_load_b32_no_overwrite(self):
|
||||
"""DS_LOAD_B32 should only write 1 VGPR."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_mov_b32(s[0], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
s_mov_b32(s[0], 0x11111111),
|
||||
v_mov_b32_e32(v[2], s[0]), # sentinel
|
||||
ds_store_b32(addr=v[0], data0=v[1], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
ds_load_b32(addr=v[0], vdst=v[1], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
|
||||
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should be untouched")
|
||||
|
||||
|
||||
class TestDS2AddrStride64(unittest.TestCase):
|
||||
"""Tests for DS_*_2ADDR_STRIDE64 (offset * 256 for B32, offset * 512 for B64)."""
|
||||
|
||||
def test_ds_store_load_2addr_stride64_b32(self):
|
||||
"""DS_STORE_2ADDR_STRIDE64_B32: stores at ADDR + offset*256."""
|
||||
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]),
|
||||
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[2], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 from addr 256")
|
||||
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 from addr 512")
|
||||
|
||||
def test_ds_store_load_2addr_stride64_b64(self):
|
||||
"""DS_STORE_2ADDR_STRIDE64_B64: stores at ADDR + offset*512."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[0], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[0], 0xCAFEBABE),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
s_mov_b32(s[0], 0x12345678),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0x9ABCDEF0),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0], data1=v[2], vdst=v[0], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[4], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
|
||||
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
|
||||
self.assertEqual(st.vgpr[0][6], 0x12345678)
|
||||
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
|
||||
|
||||
def test_ds_storexchg_2addr_rtn_b32(self):
|
||||
"""DS_STOREXCHG_2ADDR_RTN_B32: exchange at two addresses."""
|
||||
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(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[0], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
DS(DSOp.DS_STOREXCHG_2ADDR_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4], offset0=0, offset1=1),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[6], offset0=0, offset1=1),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][4], 0x11111111, "old val 0")
|
||||
self.assertEqual(st.vgpr[0][5], 0x22222222, "old val 1")
|
||||
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 = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
s_mov_b32(s[0], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[0], s[0]), # initial low
|
||||
s_mov_b32(s[0], 0xCAFEBABE),
|
||||
v_mov_b32_e32(v[1], s[0]), # initial high
|
||||
DS(DSOp.DS_STORE_B64, addr=v[10], data0=v[0], vdst=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[0], 0x12345678),
|
||||
v_mov_b32_e32(v[2], s[0]), # new low
|
||||
s_mov_b32(s[0], 0x9ABCDEF0),
|
||||
v_mov_b32_e32(v[3], s[0]), # new high
|
||||
DS(DSOp.DS_STOREXCHG_RTN_B64, addr=v[10], data0=v[2], vdst=v[4], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_B64, addr=v[10], vdst=v[6], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should have old low dword")
|
||||
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE, "v5 should have old high dword")
|
||||
self.assertEqual(st.vgpr[0][6], 0x12345678, "v6 should have new low dword")
|
||||
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0, "v7 should have new high dword")
|
||||
|
||||
def test_ds_store_load_2addr_stride64_b64_roundtrip(self):
|
||||
"""DS_STORE_2ADDR_STRIDE64_B64 followed by DS_LOAD_2ADDR_STRIDE64_B64 works correctly."""
|
||||
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(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0], data1=v[0], vdst=v[0], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[2], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have val1 low")
|
||||
self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have val1 high")
|
||||
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have val2 low")
|
||||
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have val2 high")
|
||||
|
||||
def test_ds_storexchg_2addr_stride64_rtn_b32(self):
|
||||
"""DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32: exchange at two addresses (offset*256)."""
|
||||
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(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[0], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[6], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have old value")
|
||||
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have old value")
|
||||
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "v6 should have new value")
|
||||
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "v7 should have new value")
|
||||
|
||||
def test_ds_storexchg_2addr_stride64_rtn_b64_returns_old(self):
|
||||
"""DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64: returns old values correctly."""
|
||||
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(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0], data1=v[0], vdst=v[0], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[0], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[6], s[0]),
|
||||
s_mov_b32(s[0], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[7], s[0]),
|
||||
DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64, addr=v[10], data0=v[6], data1=v[6], vdst=v[8], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][8], 0x11111111, "v8 should have old val1 low")
|
||||
self.assertEqual(st.vgpr[0][9], 0x22222222, "v9 should have old val1 high")
|
||||
self.assertEqual(st.vgpr[0][10], 0x11111111, "v10 should have old val2 low")
|
||||
self.assertEqual(st.vgpr[0][11], 0x22222222, "v11 should have old val2 high")
|
||||
|
||||
|
||||
class TestAtomicOrdering(unittest.TestCase):
|
||||
"""Tests for atomic operation return values and ordering."""
|
||||
|
||||
def test_ds_add_rtn_sequence(self):
|
||||
"""DS_ADD_RTN returns correct old values in sequence."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[10], 0),
|
||||
v_mov_b32_e32(v[0], 100),
|
||||
DS(DSOp.DS_STORE_B32, addr=v[10], data0=v[0], vdst=v[0], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[1], 25),
|
||||
DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[3], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_B32, addr=v[10], vdst=v[4], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 100, "First add should return 100")
|
||||
self.assertEqual(st.vgpr[0][3], 125, "Second add should return 125")
|
||||
self.assertEqual(st.vgpr[0][4], 150, "Final value should be 150")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,363 @@
|
||||
"""Tests for FLAT instructions - flat memory operations.
|
||||
|
||||
Includes: flat_load_*, flat_store_*, flat_atomic_*
|
||||
"""
|
||||
import unittest
|
||||
from extra.assembly.amd.test.hw.helpers import *
|
||||
|
||||
class TestFlatAtomic(unittest.TestCase):
|
||||
"""Tests for FLAT atomic instructions."""
|
||||
|
||||
def _make_test(self, setup_instrs, atomic_instr, check_fn, test_offset=2000):
|
||||
"""Helper to create atomic test instructions."""
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
] + setup_instrs + [atomic_instr, s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
check_fn(st)
|
||||
|
||||
def test_flat_atomic_add_u32(self):
|
||||
"""FLAT_ATOMIC_ADD_U32 adds to memory and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 100),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 50),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_ADD_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][4], 100)
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_swap_b32(self):
|
||||
"""FLAT_ATOMIC_SWAP_B32 swaps memory value and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_SWAP_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_and_b32(self):
|
||||
"""FLAT_ATOMIC_AND_B32 ANDs with memory and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 0xFF00FF00),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 0xFFFF0000),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_AND_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][4], 0xFF00FF00)
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_or_b32(self):
|
||||
"""FLAT_ATOMIC_OR_B32 ORs with memory and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 0x00FF0000),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 0x0000FF00),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_OR_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][4], 0x00FF0000)
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_inc_u32(self):
|
||||
"""FLAT_ATOMIC_INC_U32 increments and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 10),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 100), # threshold
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_INC_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][4], 10)
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_dec_u32(self):
|
||||
"""FLAT_ATOMIC_DEC_U32 decrements and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 10),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 100),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_DEC_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][4], 10)
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_sub_u32(self):
|
||||
"""FLAT_ATOMIC_SUB_U32 subtracts from memory and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 100),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 30),
|
||||
v_mov_b32_e32(v[3], s[0]), # sub 30
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_SUB_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][4], 100, "v4 should have old value (100)")
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_xor_b32(self):
|
||||
"""FLAT_ATOMIC_XOR_B32 XORs with memory and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 0xFFFFFFFF),
|
||||
v_mov_b32_e32(v[3], s[0]), # XOR mask
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_XOR_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA, "v4 should have old value")
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_min_u32(self):
|
||||
"""FLAT_ATOMIC_MIN_U32 stores min and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 100),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 50),
|
||||
v_mov_b32_e32(v[3], s[0]), # compare value (smaller)
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_MIN_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][4], 100, "v4 should have old value (100)")
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_max_u32(self):
|
||||
"""FLAT_ATOMIC_MAX_U32 stores max and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 50),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 100),
|
||||
v_mov_b32_e32(v[3], s[0]), # compare value (larger)
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_MAX_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][4], 50, "v4 should have old value (50)")
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_inc_u64_returns_old_value(self):
|
||||
"""FLAT_ATOMIC_INC_U64 should return full 64-bit old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
# Store initial 64-bit value: 0xCAFEBABE_DEADBEEF
|
||||
s_mov_b32(s[0], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0xCAFEBABE),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
# Threshold: 0xFFFFFFFF_FFFFFFFF
|
||||
s_mov_b32(s[0], 0xFFFFFFFF),
|
||||
v_mov_b32_e32(v[4], s[0]),
|
||||
v_mov_b32_e32(v[5], s[0]),
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_INC_U64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][6], 0xDEADBEEF, "v6 should have old value low dword")
|
||||
self.assertEqual(st.vgpr[0][7], 0xCAFEBABE, "v7 should have old value high dword")
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_add_u64(self):
|
||||
"""FLAT_ATOMIC_ADD_U64 adds 64-bit value and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 0x11111111),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0x22222222),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 0x00000001), # add 1
|
||||
v_mov_b32_e32(v[4], s[0]),
|
||||
s_mov_b32(s[0], 0x00000000),
|
||||
v_mov_b32_e32(v[5], s[0]),
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_ADD_U64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][6], 0x11111111, "v6 should have old value low")
|
||||
self.assertEqual(st.vgpr[0][7], 0x22222222, "v7 should have old value high")
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_flat_atomic_swap_b64(self):
|
||||
"""FLAT_ATOMIC_SWAP_B64 swaps 64-bit value and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 0xCCCCCCCC),
|
||||
v_mov_b32_e32(v[4], s[0]),
|
||||
s_mov_b32(s[0], 0xDDDDDDDD),
|
||||
v_mov_b32_e32(v[5], s[0]),
|
||||
]
|
||||
atomic = FLAT(FLATOp.FLAT_ATOMIC_SWAP_B64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "v6 should have old value low")
|
||||
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "v7 should have old value high")
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
|
||||
class TestFlatLoad(unittest.TestCase):
|
||||
"""Tests for FLAT load instructions."""
|
||||
|
||||
def test_flat_load_b32(self):
|
||||
"""FLAT_LOAD_B32 loads 32-bit value correctly."""
|
||||
TEST_OFFSET = 2000
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
s_mov_b32(s[0], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
FLAT(FLATOp.FLAT_LOAD_B32, addr=v[0], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
|
||||
|
||||
def test_flat_load_b64(self):
|
||||
"""FLAT_LOAD_B64 loads 64-bit value correctly."""
|
||||
TEST_OFFSET = 2000
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
s_mov_b32(s[0], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0xCAFEBABE),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
FLAT(FLATOp.FLAT_LOAD_B64, addr=v[0], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
|
||||
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
|
||||
|
||||
def test_flat_load_b96(self):
|
||||
"""FLAT_LOAD_B96 loads 96-bit (3 dword) value correctly."""
|
||||
TEST_OFFSET = 2000
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
s_mov_b32(s[0], 0x11111111),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0x22222222),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
s_mov_b32(s[0], 0x33333333),
|
||||
v_mov_b32_e32(v[4], s[0]),
|
||||
global_store_b96(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
FLAT(FLATOp.FLAT_LOAD_B96, addr=v[0], vdst=v[5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][5], 0x11111111)
|
||||
self.assertEqual(st.vgpr[0][6], 0x22222222)
|
||||
self.assertEqual(st.vgpr[0][7], 0x33333333)
|
||||
|
||||
def test_flat_load_b128(self):
|
||||
"""FLAT_LOAD_B128 loads 128-bit value correctly."""
|
||||
TEST_OFFSET = 2000
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
s_mov_b32(s[0], 0x11111111),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0x22222222),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
s_mov_b32(s[0], 0x33333333),
|
||||
v_mov_b32_e32(v[4], s[0]),
|
||||
s_mov_b32(s[0], 0x44444444),
|
||||
v_mov_b32_e32(v[5], s[0]),
|
||||
global_store_b128(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
FLAT(FLATOp.FLAT_LOAD_B128, addr=v[0], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][6], 0x11111111)
|
||||
self.assertEqual(st.vgpr[0][7], 0x22222222)
|
||||
self.assertEqual(st.vgpr[0][8], 0x33333333)
|
||||
self.assertEqual(st.vgpr[0][9], 0x44444444)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,364 @@
|
||||
"""Tests for GLOBAL instructions - global memory operations.
|
||||
|
||||
Includes: global_load_*, global_store_*, global_atomic_*, global_load_d16_*
|
||||
"""
|
||||
import unittest
|
||||
from extra.assembly.amd.test.hw.helpers import *
|
||||
|
||||
class TestGlobalAtomic(unittest.TestCase):
|
||||
"""Tests for GLOBAL atomic instructions."""
|
||||
|
||||
def _make_test(self, setup_instrs, atomic_instr, check_fn, test_offset=2000):
|
||||
"""Helper to create atomic test instructions."""
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
] + setup_instrs + [atomic_instr, s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
check_fn(st)
|
||||
|
||||
def test_global_atomic_add_u32(self):
|
||||
"""GLOBAL_ATOMIC_ADD_U32 adds to memory and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 100),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 50),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
]
|
||||
atomic = FLAT(GLOBALOp.GLOBAL_ATOMIC_ADD_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1, seg=2)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][4], 100)
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
def test_global_atomic_add_u64(self):
|
||||
"""GLOBAL_ATOMIC_ADD_U64 adds 64-bit value and returns old value."""
|
||||
TEST_OFFSET = 2000
|
||||
setup = [
|
||||
s_mov_b32(s[0], 0xFFFFFFFF),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0x00000000),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[0], 0x00000001),
|
||||
v_mov_b32_e32(v[4], s[0]),
|
||||
s_mov_b32(s[0], 0x00000000),
|
||||
v_mov_b32_e32(v[5], s[0]),
|
||||
]
|
||||
atomic = FLAT(GLOBALOp.GLOBAL_ATOMIC_ADD_U64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1, seg=2)
|
||||
def check(st):
|
||||
self.assertEqual(st.vgpr[0][6], 0xFFFFFFFF)
|
||||
self.assertEqual(st.vgpr[0][7], 0x00000000)
|
||||
self._make_test(setup, atomic, check, TEST_OFFSET)
|
||||
|
||||
|
||||
class TestGlobalLoad(unittest.TestCase):
|
||||
"""Tests for GLOBAL load instructions."""
|
||||
|
||||
def test_global_load_b96(self):
|
||||
"""GLOBAL_LOAD_B96 loads 96-bit value correctly."""
|
||||
TEST_OFFSET = 2000
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
s_mov_b32(s[0], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
s_mov_b32(s[0], 0xCCCCCCCC),
|
||||
v_mov_b32_e32(v[4], s[0]),
|
||||
global_store_b96(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_B96, addr=v[0], vdst=v[5], saddr=SrcEnum.NULL, offset=TEST_OFFSET, seg=2),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][5], 0xAAAAAAAA)
|
||||
self.assertEqual(st.vgpr[0][6], 0xBBBBBBBB)
|
||||
self.assertEqual(st.vgpr[0][7], 0xCCCCCCCC)
|
||||
|
||||
def test_global_load_b128(self):
|
||||
"""GLOBAL_LOAD_B128 loads 128-bit value correctly."""
|
||||
TEST_OFFSET = 2000
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
s_mov_b32(s[0], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
s_mov_b32(s[0], 0xCAFEBABE),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
s_mov_b32(s[0], 0x12345678),
|
||||
v_mov_b32_e32(v[4], s[0]),
|
||||
s_mov_b32(s[0], 0x9ABCDEF0),
|
||||
v_mov_b32_e32(v[5], s[0]),
|
||||
global_store_b128(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_B128, addr=v[0], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, seg=2),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][6], 0xDEADBEEF)
|
||||
self.assertEqual(st.vgpr[0][7], 0xCAFEBABE)
|
||||
self.assertEqual(st.vgpr[0][8], 0x12345678)
|
||||
self.assertEqual(st.vgpr[0][9], 0x9ABCDEF0)
|
||||
|
||||
|
||||
class TestGlobalStore(unittest.TestCase):
|
||||
"""Tests for GLOBAL store instructions."""
|
||||
|
||||
def test_global_store_b64_basic(self):
|
||||
"""GLOBAL_STORE_B64 stores 8 bytes from v[n:n+1] to memory."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[4], 0xDEADBEEF),
|
||||
s_mov_b32(s[5], 0xCAFEBABE),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
v_mov_b32_e32(v[3], s[5]),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
global_store_b64(addr=v[0], data=v[2], saddr=s[2], offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_B64, addr=v[0], vdst=v[4], data=v[4], saddr=s[2], offset=TEST_OFFSET, seg=2),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[4]),
|
||||
v_mov_b32_e32(v[1], v[5]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
|
||||
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
|
||||
|
||||
|
||||
class TestD16HiLoads(unittest.TestCase):
|
||||
"""Tests for D16_HI load instructions that load into high 16 bits."""
|
||||
|
||||
def test_global_load_d16_hi_b16_preserves_low_bits(self):
|
||||
"""GLOBAL_LOAD_D16_HI_B16 must preserve low 16 bits of destination."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
s_mov_b32(s[4], 0xCAFE),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
global_store_b16(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[4], 0x0000BEEF),
|
||||
v_mov_b32_e32(v[3], s[4]),
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[0], vdst=v[3], data=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET, seg=2),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][0]
|
||||
self.assertEqual(result, 0xCAFEBEEF, f"Expected 0xCAFEBEEF, got 0x{result:08x}")
|
||||
|
||||
def test_global_load_d16_hi_b16_data_differs_from_vdst(self):
|
||||
"""GLOBAL_LOAD_D16_HI_B16 where data field differs from vdst."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[4], 0xCAFE),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
v_mov_b32_e32(v[3], 0),
|
||||
global_store_b16(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[4], 0x0000DEAD),
|
||||
v_mov_b32_e32(v[0], s[4]), # data field - should NOT affect result
|
||||
v_mov_b32_e32(v[1], 0), # vdst - low bits should be preserved
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[0], saddr=s[2], offset=TEST_OFFSET, seg=2),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[1]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][0]
|
||||
self.assertEqual(result, 0xCAFE0000, f"Expected 0xCAFE0000, got 0x{result:08x}")
|
||||
|
||||
def test_global_load_d16_hi_u8_data_differs_from_vdst(self):
|
||||
"""GLOBAL_LOAD_D16_HI_U8 where data field differs from vdst."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[4], 0xAB),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
v_mov_b32_e32(v[3], 0),
|
||||
global_store_b8(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[4], 0x0000DEAD),
|
||||
v_mov_b32_e32(v[4], s[4]), # data field
|
||||
s_mov_b32(s[4], 0x0000BEEF),
|
||||
v_mov_b32_e32(v[5], s[4]), # vdst
|
||||
v_mov_b32_e32(v[3], 0),
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_U8, addr=v[3], vdst=v[5], data=v[4], saddr=s[2], offset=TEST_OFFSET, seg=2),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[5]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][0]
|
||||
self.assertEqual(result, 0x00ABBEEF, f"Expected 0x00ABBEEF, got 0x{result:08x}")
|
||||
|
||||
def test_global_load_d16_hi_b16_same_addr_and_dst_zero_addr(self):
|
||||
"""GLOBAL_LOAD_D16_HI_B16 with same register for addr and vdst, addr value=0."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[4], 0xCAFE),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
v_mov_b32_e32(v[3], 0),
|
||||
global_store_b16(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[1], saddr=s[2], offset=TEST_OFFSET, seg=2),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[1]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][0]
|
||||
self.assertEqual(result, 0xCAFE0000, f"Expected 0xCAFE0000, got 0x{result:08x}")
|
||||
|
||||
def test_global_load_d16_hi_b16_tril_exact_pattern(self):
|
||||
"""Exact pattern from tril() failure: data=v0 differs from vdst=v1."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[4], 0x01010101),
|
||||
v_mov_b32_e32(v[10], s[4]),
|
||||
v_mov_b32_e32(v[3], 0),
|
||||
global_store_b32(addr=v[3], data=v[10], saddr=s[2], offset=TEST_OFFSET),
|
||||
global_store_b32(addr=v[3], data=v[10], saddr=s[2], offset=TEST_OFFSET+4),
|
||||
s_waitcnt(vmcnt=0),
|
||||
# Set v[0] to 0x0101 (simulating prior u16 load result)
|
||||
s_mov_b32(s[4], 0x0101),
|
||||
v_mov_b32_e32(v[0], s[4]),
|
||||
# Set v[1] to 0
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
# Load using v[1] as addr AND vdst, but v[0] as data
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[0], saddr=s[2], offset=TEST_OFFSET+6, seg=2),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[1]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][0]
|
||||
# Expected: hi=0x0101 (loaded), lo=0x0000 (from v1) -> 0x01010000
|
||||
self.assertEqual(result, 0x01010000, f"Expected 0x01010000, got 0x{result:08x}")
|
||||
|
||||
def test_global_load_d16_hi_i8_data_differs_from_vdst(self):
|
||||
"""GLOBAL_LOAD_D16_HI_I8 where data field differs from vdst."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[4], 0x80), # negative signed byte = -128
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
v_mov_b32_e32(v[3], 0),
|
||||
global_store_b8(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
s_mov_b32(s[4], 0x0000DEAD),
|
||||
v_mov_b32_e32(v[4], s[4]), # data field
|
||||
s_mov_b32(s[4], 0x0000BEEF),
|
||||
v_mov_b32_e32(v[5], s[4]), # vdst
|
||||
v_mov_b32_e32(v[3], 0),
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_I8, addr=v[3], vdst=v[5], data=v[4], saddr=s[2], offset=TEST_OFFSET, seg=2),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[5]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][0]
|
||||
# 0x80 sign-extended = 0xFF80, lo=0xBEEF -> 0xFF80BEEF
|
||||
self.assertEqual(result, 0xFF80BEEF, f"Expected 0xFF80BEEF, got 0x{result:08x}")
|
||||
|
||||
def test_global_store_b64_tril_pattern(self):
|
||||
"""Test the exact pattern from tril() kernel that was failing."""
|
||||
TEST_OFFSET = 256
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
s_mov_b32(s[4], 0x01010101),
|
||||
v_mov_b32_e32(v[10], s[4]),
|
||||
v_mov_b32_e32(v[11], s[4]),
|
||||
s_mov_b32(s[4], 0x01),
|
||||
v_mov_b32_e32(v[12], s[4]),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
global_store_b64(addr=v[0], data=v[10], saddr=s[2], offset=TEST_OFFSET),
|
||||
global_store_b8(addr=v[0], data=v[12], saddr=s[2], offset=TEST_OFFSET+8),
|
||||
s_waitcnt(vmcnt=0),
|
||||
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_U16, addr=v[2], vdst=v[0], data=v[0], saddr=s[2], offset=TEST_OFFSET+3, seg=2),
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[1], saddr=s[2], offset=TEST_OFFSET+6, seg=2),
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_U8, addr=v[2], vdst=v[3], data=v[3], saddr=s[2], offset=TEST_OFFSET, seg=2),
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_U8, addr=v[2], vdst=v[4], data=v[4], saddr=s[2], offset=TEST_OFFSET+8, seg=2),
|
||||
s_waitcnt(vmcnt=0),
|
||||
|
||||
v_and_b32_e32(v[5], 0xffff, v[0]),
|
||||
v_lshlrev_b32_e32(v[0], 24, v[0]),
|
||||
v_lshrrev_b32_e32(v[5], 8, v[5]),
|
||||
v_or_b32_e32(v[0], v[3], v[0]),
|
||||
v_or_b32_e32(v[1], v[5], v[1]),
|
||||
|
||||
global_store_b64(addr=v[2], data=v[0], saddr=s[2], offset=TEST_OFFSET+16),
|
||||
s_waitcnt(vmcnt=0),
|
||||
|
||||
FLAT(GLOBALOp.GLOBAL_LOAD_B64, addr=v[2], vdst=v[6], data=v[6], saddr=s[2], offset=TEST_OFFSET+16, seg=2),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[6]),
|
||||
v_mov_b32_e32(v[1], v[7]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
|
||||
v0 = st.vgpr[0][0]
|
||||
v1 = st.vgpr[0][1]
|
||||
self.assertEqual(v0, 0x01000001, f"v0: expected 0x01000001, got 0x{v0:08x}")
|
||||
self.assertEqual(v1, 0x01010001, f"v1: expected 0x01010001, got 0x{v1:08x}")
|
||||
|
||||
byte5 = (v1 >> 8) & 0xff
|
||||
self.assertEqual(byte5, 0x00, f"byte5: expected 0x00, got 0x{byte5:02x}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,342 @@
|
||||
"""Tests for SOP instructions - scalar operations.
|
||||
|
||||
Includes: s_add_u32, s_mov_b32, s_and_b32, s_or_b32, s_quadmask_b32, s_wqm_b32,
|
||||
s_cbranch_vccnz, s_cbranch_vccz
|
||||
"""
|
||||
import unittest
|
||||
from extra.assembly.amd.test.hw.helpers import *
|
||||
|
||||
class TestBasicScalar(unittest.TestCase):
|
||||
"""Tests for basic scalar operations."""
|
||||
|
||||
def test_s_add_u32(self):
|
||||
"""S_ADD_U32 adds two scalar values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 100),
|
||||
s_mov_b32(s[1], 200),
|
||||
s_add_u32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 300)
|
||||
|
||||
def test_s_add_u32_carry(self):
|
||||
"""S_ADD_U32 sets SCC on overflow."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 64),
|
||||
s_not_b32(s[0], s[0]), # ~64 = 0xffffffbf
|
||||
s_mov_b32(s[1], 64),
|
||||
s_add_u32(s[2], s[0], s[1]), # 0xffffffbf + 64 = 0xffffffff
|
||||
s_mov_b32(s[3], 1),
|
||||
s_add_u32(s[4], s[2], s[3]), # 0xffffffff + 1 = overflow
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[4], 0)
|
||||
self.assertEqual(st.scc, 1)
|
||||
|
||||
def test_s_brev_b32(self):
|
||||
"""S_BREV_B32 reverses bits of a 32-bit value."""
|
||||
# 10 = 0b00000000000000000000000000001010
|
||||
# reversed = 0b01010000000000000000000000000000 = 0x50000000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 10),
|
||||
s_brev_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0x50000000)
|
||||
|
||||
def test_s_brev_b32_all_ones(self):
|
||||
"""S_BREV_B32 with all ones stays all ones."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFFFFFFFF),
|
||||
s_brev_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
|
||||
|
||||
def test_s_brev_b32_single_bit(self):
|
||||
"""S_BREV_B32 with bit 0 set becomes bit 31."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 1),
|
||||
s_brev_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0x80000000)
|
||||
|
||||
|
||||
class TestQuadmaskWqm(unittest.TestCase):
|
||||
"""Tests for S_QUADMASK_B32 and S_WQM_B32."""
|
||||
|
||||
def test_s_quadmask_b32_all_quads_active(self):
|
||||
"""S_QUADMASK_B32 with all quads active."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFFFFFFFF), # All lanes active
|
||||
s_quadmask_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Each quad (4 lanes) with any bit set -> 1 bit in result
|
||||
# 32 lanes = 8 quads, all active -> 0xFF
|
||||
self.assertEqual(st.sgpr[1], 0xFF)
|
||||
|
||||
def test_s_quadmask_b32_alternating_quads(self):
|
||||
"""S_QUADMASK_B32 with alternating quads active."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x0F0F0F0F), # Quads 0,2,4,6 active
|
||||
s_quadmask_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Quads 0,2,4,6 have at least one bit -> 0b01010101 = 0x55
|
||||
self.assertEqual(st.sgpr[1], 0x55)
|
||||
|
||||
def test_s_quadmask_b32_no_quads_active(self):
|
||||
"""S_QUADMASK_B32 with no quads active."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0),
|
||||
s_quadmask_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0)
|
||||
|
||||
def test_s_quadmask_b32_single_lane_per_quad(self):
|
||||
"""S_QUADMASK_B32 with single lane active in each quad."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x11111111), # Bit 0 of each nibble
|
||||
s_quadmask_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# All 8 quads have at least one lane -> 0xFF
|
||||
self.assertEqual(st.sgpr[1], 0xFF)
|
||||
|
||||
def test_s_wqm_b32_all_active(self):
|
||||
"""S_WQM_B32 with all lanes active returns all 1s."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFFFFFFFF),
|
||||
s_wqm_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
|
||||
|
||||
def test_s_wqm_b32_alternating_quads(self):
|
||||
"""S_WQM_B32 with single lane per quad expands to full quads."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x11111111), # One lane per quad
|
||||
s_wqm_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Each quad with any bit expands to all 4 bits
|
||||
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
|
||||
|
||||
def test_s_wqm_b32_zero(self):
|
||||
"""S_WQM_B32 with zero input returns zero."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0),
|
||||
s_wqm_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0)
|
||||
|
||||
|
||||
class TestBranch(unittest.TestCase):
|
||||
"""Tests for branch instructions."""
|
||||
|
||||
def test_cbranch_vccnz_ignores_vcc_hi(self):
|
||||
"""S_CBRANCH_VCCNZ should only check VCC_LO in wave32."""
|
||||
instructions = [
|
||||
# Set VCC_LO = 0, VCC_HI = 1
|
||||
s_mov_b32(s[SrcEnum.VCC_LO - 128], 0),
|
||||
s_mov_b32(s[SrcEnum.VCC_HI - 128], 1),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
# If VCC_HI is incorrectly used, branch will be taken
|
||||
s_cbranch_vccnz(1), # Skip next instruction if VCC != 0
|
||||
v_mov_b32_e32(v[0], 42), # This should execute
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 42, "Branch should NOT be taken (VCC_LO is 0)")
|
||||
|
||||
def test_cbranch_vccz_ignores_vcc_hi(self):
|
||||
"""S_CBRANCH_VCCZ should only check VCC_LO in wave32."""
|
||||
instructions = [
|
||||
# Set VCC_LO = 1, VCC_HI = 0
|
||||
s_mov_b32(s[SrcEnum.VCC_LO - 128], 1),
|
||||
s_mov_b32(s[SrcEnum.VCC_HI - 128], 0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
# If VCC_HI is incorrectly used, branch will be taken
|
||||
s_cbranch_vccz(1), # Skip next instruction if VCC == 0
|
||||
v_mov_b32_e32(v[0], 42), # This should execute
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 42, "Branch should NOT be taken (VCC_LO is 1)")
|
||||
|
||||
def test_cbranch_vccnz_branches_on_vcc_lo(self):
|
||||
"""S_CBRANCH_VCCNZ branches when VCC_LO is non-zero."""
|
||||
instructions = [
|
||||
s_mov_b32(s[SrcEnum.VCC_LO - 128], 1),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_cbranch_vccnz(1), # Skip next instruction if VCC != 0
|
||||
v_mov_b32_e32(v[0], 42), # This should be skipped
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0, "Branch should be taken (VCC_LO is 1)")
|
||||
|
||||
|
||||
class Test64BitLiterals(unittest.TestCase):
|
||||
"""Tests for 64-bit literal encoding in instructions."""
|
||||
|
||||
def test_64bit_literal_negative_encoding(self):
|
||||
"""64-bit literal -2^32 encodes correctly."""
|
||||
lit = -4294967296.0 # -2^32
|
||||
lit_bits = f2i64(lit)
|
||||
instructions = [
|
||||
s_mov_b32(s[0], lit_bits & 0xffffffff),
|
||||
s_mov_b32(s[1], lit_bits >> 32),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = i642f(st.vgpr[0][0] | (st.vgpr[0][1] << 32))
|
||||
self.assertAlmostEqual(result, -4294967296.0, places=5)
|
||||
|
||||
def test_64bit_literal_positive_encoding(self):
|
||||
"""64-bit instruction encodes large positive literals correctly."""
|
||||
large_val = 0x12345678
|
||||
inst = v_add_f64(v[2], v[0], large_val)
|
||||
self.assertIsNotNone(inst._literal, "Literal should be set")
|
||||
actual_lit = (inst._literal >> 32) & 0xffffffff
|
||||
self.assertEqual(actual_lit, large_val, f"Literal should be {large_val:#x}, got {actual_lit:#x}")
|
||||
|
||||
|
||||
class TestSCCBehavior(unittest.TestCase):
|
||||
"""Tests for SCC condition code behavior."""
|
||||
|
||||
def test_scc_from_s_cmp(self):
|
||||
"""SCC should be set by scalar compare."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 10),
|
||||
s_cmp_eq_u32(s[0], 10),
|
||||
s_cselect_b32(s[1], 1, 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 1, "SCC should be true")
|
||||
self.assertEqual(st.scc, 1)
|
||||
|
||||
def test_scc_clear(self):
|
||||
"""SCC should be cleared by failing compare."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 10),
|
||||
s_cmp_eq_u32(s[0], 20),
|
||||
s_cselect_b32(s[1], 1, 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0, "SCC should be false")
|
||||
self.assertEqual(st.scc, 0)
|
||||
|
||||
|
||||
class TestSignedArithmetic(unittest.TestCase):
|
||||
"""Tests for S_ADD_I32, S_SUB_I32 and their SCC overflow behavior."""
|
||||
|
||||
def test_s_add_i32_no_overflow(self):
|
||||
"""S_ADD_I32: 1 + 1 = 2, no overflow, SCC=0."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 1),
|
||||
s_add_i32(s[1], s[0], 1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 2)
|
||||
self.assertEqual(st.scc, 0, "No overflow, SCC should be 0")
|
||||
|
||||
def test_s_add_i32_positive_overflow(self):
|
||||
"""S_ADD_I32: MAX_INT + 1 overflows, SCC=1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7FFFFFFF), # MAX_INT
|
||||
s_add_i32(s[1], s[0], 1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0x80000000) # Wraps to MIN_INT
|
||||
self.assertEqual(st.scc, 1, "Overflow, SCC should be 1")
|
||||
|
||||
def test_s_add_i32_negative_no_overflow(self):
|
||||
"""S_ADD_I32: -10 + 20 = 10, no overflow."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFFFFFFF6), # -10 in two's complement
|
||||
s_mov_b32(s[1], 20),
|
||||
s_add_i32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 10)
|
||||
self.assertEqual(st.scc, 0)
|
||||
|
||||
def test_s_add_i32_negative_overflow(self):
|
||||
"""S_ADD_I32: MIN_INT + (-1) underflows, SCC=1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x80000000), # MIN_INT
|
||||
s_mov_b32(s[1], 0xFFFFFFFF), # -1
|
||||
s_add_i32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 0x7FFFFFFF) # Wraps to MAX_INT
|
||||
self.assertEqual(st.scc, 1, "Underflow, SCC should be 1")
|
||||
|
||||
def test_s_sub_i32_no_overflow(self):
|
||||
"""S_SUB_I32: 10 - 5 = 5, no overflow."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 10),
|
||||
s_mov_b32(s[1], 5),
|
||||
s_sub_i32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 5)
|
||||
self.assertEqual(st.scc, 0)
|
||||
|
||||
def test_s_sub_i32_overflow(self):
|
||||
"""S_SUB_I32: MAX_INT - (-1) overflows, SCC=1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7FFFFFFF), # MAX_INT
|
||||
s_mov_b32(s[1], 0xFFFFFFFF), # -1
|
||||
s_sub_i32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 0x80000000) # Wraps to MIN_INT
|
||||
self.assertEqual(st.scc, 1, "Overflow, SCC should be 1")
|
||||
|
||||
def test_s_mul_hi_u32(self):
|
||||
"""S_MUL_HI_U32: high 32 bits of u32 * u32."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x80000000), # 2^31
|
||||
s_mov_b32(s[1], 4),
|
||||
s_mul_hi_u32(s[2], s[0], s[1]), # (2^31 * 4) >> 32 = 2
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 2)
|
||||
|
||||
def test_s_mul_i32(self):
|
||||
"""S_MUL_I32: signed multiply low 32 bits."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFFFFFFFF), # -1
|
||||
s_mov_b32(s[1], 10),
|
||||
s_mul_i32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 0xFFFFFFF6) # -10
|
||||
|
||||
def test_division_sequence_from_llvm(self):
|
||||
"""Test the division sequence pattern from LLVM-generated code."""
|
||||
# This sequence is from the sin kernel and computes integer division
|
||||
# s10 = dividend, s18 = divisor, result in s6/s14
|
||||
dividend = 0x28BE60DB # Some value from the sin kernel
|
||||
divisor = 3 # Simplified divisor
|
||||
instructions = [
|
||||
s_mov_b32(s[10], dividend),
|
||||
s_mov_b32(s[18], divisor),
|
||||
# Compute reciprocal approximation: s6 = ~0 / divisor (approx)
|
||||
s_mov_b32(s[11], 0),
|
||||
s_sub_i32(s[11], s[11], s[18]), # s11 = -divisor
|
||||
# For testing, just verify basic arithmetic works
|
||||
s_mul_i32(s[6], s[10], 2),
|
||||
s_add_i32(s[7], s[6], 1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[6], (dividend * 2) & 0xFFFFFFFF)
|
||||
self.assertEqual(st.sgpr[7], ((dividend * 2) + 1) & 0xFFFFFFFF)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,451 @@
|
||||
"""Tests for VOP2 instructions - two operand vector operations.
|
||||
|
||||
Includes: v_add_f32, v_mul_f32, v_and_b32, v_or_b32, v_xor_b32,
|
||||
v_lshrrev_b32, v_lshlrev_b32, v_fmac_f32, v_fmaak_f32, v_fmamk_f32,
|
||||
v_add_nc_u32, v_cndmask_b32, v_add_f16, v_mul_f16
|
||||
"""
|
||||
import unittest
|
||||
from extra.assembly.amd.test.hw.helpers import *
|
||||
|
||||
class TestBasicArithmetic(unittest.TestCase):
|
||||
"""Tests for basic arithmetic VOP2 instructions."""
|
||||
|
||||
def test_v_add_f32(self):
|
||||
"""V_ADD_F32 adds two floats."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 1.0),
|
||||
v_mov_b32_e32(v[1], 2.0),
|
||||
v_add_f32_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 3.0, places=5)
|
||||
|
||||
def test_v_mul_f32(self):
|
||||
"""V_MUL_F32 multiplies two floats."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 2.0),
|
||||
v_mov_b32_e32(v[1], 4.0),
|
||||
v_mul_f32_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 8.0, places=5)
|
||||
|
||||
def test_v_fmac_f32(self):
|
||||
"""V_FMAC_F32: d = d + a*b using inline constants."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 2.0),
|
||||
v_mov_b32_e32(v[1], 4.0),
|
||||
v_mov_b32_e32(v[2], 1.0),
|
||||
v_fmac_f32_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
|
||||
|
||||
def test_v_fmaak_f32(self):
|
||||
"""V_FMAAK_F32: d = a * b + K using inline constants."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 2.0),
|
||||
v_mov_b32_e32(v[1], 4.0),
|
||||
v_fmaak_f32_e32(v[2], v[0], v[1], 0x3f800000),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
|
||||
|
||||
def test_v_fmamk_f32_basic(self):
|
||||
"""V_FMAMK_F32: d = a * K + b."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 2.0),
|
||||
v_mov_b32_e32(v[1], 1.0),
|
||||
v_fmamk_f32_e32(v[2], v[0], 0x40800000, v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
|
||||
|
||||
def test_v_fmamk_f32_small_constant(self):
|
||||
"""V_FMAMK_F32 with small constant."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 4.0),
|
||||
v_mov_b32_e32(v[1], 1.0),
|
||||
v_fmamk_f32_e32(v[2], v[0], f2i(0.5), v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 3.0, places=5)
|
||||
|
||||
|
||||
class TestBitManipulation(unittest.TestCase):
|
||||
"""Tests for bit manipulation VOP2 instructions."""
|
||||
|
||||
def test_v_and_b32(self):
|
||||
"""V_AND_B32 bitwise and."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xff),
|
||||
s_mov_b32(s[1], 0x0f),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_and_b32_e32(v[1], s[1], v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0x0f)
|
||||
|
||||
def test_v_and_b32_quadrant(self):
|
||||
"""V_AND_B32 for quadrant extraction (n & 3)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 15915),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_and_b32_e32(v[1], 3, v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 15915 & 3)
|
||||
|
||||
def test_v_lshrrev_b32(self):
|
||||
"""V_LSHRREV_B32 logical shift right."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xff00),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_lshrrev_b32_e32(v[1], 8, v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0xff)
|
||||
|
||||
def test_v_lshlrev_b32(self):
|
||||
"""V_LSHLREV_B32 logical shift left."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xff),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_lshlrev_b32_e32(v[1], 8, v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0xff00)
|
||||
|
||||
def test_v_xor_b32(self):
|
||||
"""V_XOR_B32 bitwise xor (used in sin for sign)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x80000000),
|
||||
s_mov_b32(s[1], f2i(1.0)),
|
||||
v_mov_b32_e32(v[0], s[1]),
|
||||
v_xor_b32_e32(v[1], s[0], v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][1]), -1.0, places=5)
|
||||
|
||||
def test_v_xor_b32_sign_flip(self):
|
||||
"""V_XOR_B32 for sign flip pattern."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x80000000),
|
||||
v_mov_b32_e32(v[0], -2.0),
|
||||
v_xor_b32_e32(v[1], s[0], v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 2.0, places=5)
|
||||
|
||||
|
||||
class TestSpecialValues(unittest.TestCase):
|
||||
"""Tests for special float values - inf, nan, zero handling."""
|
||||
|
||||
def test_v_mul_f32_zero_times_inf(self):
|
||||
"""V_MUL_F32: 0 * inf = NaN."""
|
||||
import math
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_mov_b32(s[0], 0x7f800000),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
v_mul_f32_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
|
||||
|
||||
def test_v_add_f32_inf_minus_inf(self):
|
||||
"""V_ADD_F32: inf + (-inf) = NaN."""
|
||||
import math
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7f800000),
|
||||
s_mov_b32(s[1], 0xff800000),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_add_f32_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
|
||||
|
||||
|
||||
class TestF16Ops(unittest.TestCase):
|
||||
"""Tests for 16-bit VOP2 operations."""
|
||||
|
||||
def test_v_add_f16_basic(self):
|
||||
"""V_ADD_F16 adds two f16 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_add_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x4200, f"Expected 0x4200 (f16 3.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_add_f16_negative(self):
|
||||
"""V_ADD_F16 with negative values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0xc000), # f16 -2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_add_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0xbc00, f"Expected 0xbc00 (f16 -1.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_mul_f16_basic(self):
|
||||
"""V_MUL_F16 multiplies two f16 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x4000), # f16 2.0
|
||||
s_mov_b32(s[1], 0x4200), # f16 3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mul_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x4600, f"Expected 0x4600 (f16 6.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_mul_f16_by_zero(self):
|
||||
"""V_MUL_F16 by zero."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
v_mul_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x0000, f"Expected 0x0000 (f16 0.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_fmac_f16_basic(self):
|
||||
"""V_FMAC_F16: d = d + a*b."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x4000), # f16 2.0
|
||||
s_mov_b32(s[1], 0x4200), # f16 3.0
|
||||
s_mov_b32(s[2], 0x3c00), # f16 1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
v_fmac_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
|
||||
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_fmaak_f16_basic(self):
|
||||
"""V_FMAAK_F16: d = a * b + K."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x4000), # f16 2.0
|
||||
s_mov_b32(s[1], 0x4200), # f16 3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_fmaak_f16_e32(v[2], v[0], v[1], 0x3c00), # + f16 1.0
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
|
||||
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
|
||||
|
||||
|
||||
class TestHiHalfOps(unittest.TestCase):
|
||||
"""Tests for VOP2 16-bit operations with hi-half operands."""
|
||||
|
||||
def test_v_add_f16_src0_hi_fold(self):
|
||||
"""V_ADD_F16 with src0 hi-half fold (same register, different halves)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # lo=f16(1.0), hi=f16(2.0)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
VOP3(VOP3Op.V_ADD_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0001),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][1] & 0xffff
|
||||
self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
|
||||
|
||||
def test_v_add_f16_src0_hi_different_reg(self):
|
||||
"""V_ADD_F16 with src0 hi-half from different register."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40000000), # hi=f16(2.0), lo=0
|
||||
s_mov_b32(s[1], 0x00003c00), # hi=0, lo=f16(1.0)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
VOP3(VOP3Op.V_ADD_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
|
||||
|
||||
def test_v_mul_f16_src0_hi(self):
|
||||
"""V_MUL_F16 with src0 from high half."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40000000), # hi=f16(2.0), lo=0
|
||||
s_mov_b32(s[1], 0x00004200), # hi=0, lo=f16(3.0)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
VOP3(VOP3Op.V_MUL_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x4600, f"Expected f16(6.0)=0x4600, got 0x{result:04x}")
|
||||
|
||||
def test_v_mul_f16_hi_half(self):
|
||||
"""V_MUL_F16 reading from high half."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # lo=1.0, hi=2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
VOP3(VOP3Op.V_MUL_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0011),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][1] & 0xffff
|
||||
self.assertEqual(result, 0x4400, f"Expected f16(4.0)=0x4400, got 0x{result:04x}")
|
||||
|
||||
def test_v_fma_f16_hi_dest(self):
|
||||
"""V_FMA_F16 writing to high half with opsel.
|
||||
|
||||
Uses V_FMA_F16 (not V_FMAC_F16) because it has explicit src2 operand
|
||||
which makes opsel handling clearer.
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c000000), # hi=f16(1.0), lo=0
|
||||
s_mov_b32(s[1], 0x4000), # f16(2.0) in lo
|
||||
s_mov_b32(s[2], 0x4200), # f16(3.0) in lo
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
# V_FMA_F16: dst = src0 * src1 + src2
|
||||
# opsel=0b1100: bit2=src2 hi, bit3=dst hi
|
||||
# So: v[0].hi = v[1].lo * v[2].lo + v[0].hi = 2.0 * 3.0 + 1.0 = 7.0
|
||||
VOP3(VOP3Op.V_FMA_F16, vdst=v[0], src0=v[1], src1=v[2], src2=v[0], opsel=0b1100),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
hi = (st.vgpr[0][0] >> 16) & 0xffff
|
||||
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
|
||||
self.assertEqual(hi, 0x4700, f"Expected f16(7.0)=0x4700 in hi, got 0x{hi:04x}")
|
||||
|
||||
def test_v_add_f16_multilane(self):
|
||||
"""V_ADD_F16 with multiple lanes."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_add_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=4)
|
||||
for lane in range(4):
|
||||
result = st.vgpr[lane][2] & 0xffff
|
||||
self.assertEqual(result, 0x4200, f"Lane {lane}: expected 0x4200, got 0x{result:04x}")
|
||||
|
||||
|
||||
class TestCndmask(unittest.TestCase):
|
||||
"""Tests for V_CNDMASK_B32 and V_CNDMASK_B16."""
|
||||
|
||||
def test_v_cndmask_b16_select_src0(self):
|
||||
"""V_CNDMASK_B16 selects src0 when VCC bit is 0."""
|
||||
instructions = [
|
||||
s_mov_b32(s[SrcEnum.VCC_LO - 128], 0), # VCC = 0
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cndmask_b16(v[2], v[0], v[1], VCC),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x3c00, f"Expected src0=0x3c00, got 0x{result:04x}")
|
||||
|
||||
def test_v_cndmask_b16_select_src1(self):
|
||||
"""V_CNDMASK_B16 selects src1 when VCC bit is 1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[SrcEnum.VCC_LO - 128], 1), # VCC = 1
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cndmask_b16(v[2], v[0], v[1], VCC),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x4000, f"Expected src1=0x4000, got 0x{result:04x}")
|
||||
|
||||
def test_v_cndmask_b16_write_hi(self):
|
||||
"""V_CNDMASK_B16 can write to high 16 bits with opsel."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c003800), # src0: hi=1.0, lo=0.5
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], 0x4000c000), # src1: hi=2.0, lo=-2.0
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], 0xDEAD0000), # v2 initial: hi=0xDEAD, lo=0
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
s_mov_b32(s[SrcEnum.VCC_LO - 128], 0), # vcc = 0, select src0
|
||||
# opsel=0b1011: bit0=src0 hi, bit1=src1 hi, bit3=dst hi
|
||||
VOP3(VOP3Op.V_CNDMASK_B16, vdst=v[2], src0=v[0], src1=v[1], src2=SrcEnum.VCC_LO, opsel=0b1011),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
hi = (st.vgpr[0][2] >> 16) & 0xffff
|
||||
lo = st.vgpr[0][2] & 0xffff
|
||||
# vcc=0 selects src0.h = 1.0 = 0x3c00, writes to hi
|
||||
self.assertEqual(hi, 0x3c00, f"Expected hi=0x3c00 (1.0), got 0x{hi:04x}")
|
||||
self.assertEqual(lo, 0x0000, f"Expected lo preserved as 0, got 0x{lo:04x}")
|
||||
|
||||
|
||||
class TestSpecialFloatValues(unittest.TestCase):
|
||||
"""Tests for special float value handling in VOP2 instructions."""
|
||||
|
||||
def test_neg_zero_add(self):
|
||||
"""-0.0 + 0.0 = +0.0 (IEEE 754)."""
|
||||
neg_zero = 0x80000000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], neg_zero),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_add_f32_e32(v[1], 0.0, v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
|
||||
|
||||
def test_neg_zero_mul(self):
|
||||
"""-0.0 * -1.0 = +0.0."""
|
||||
neg_zero = 0x80000000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], neg_zero),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mul_f32_e32(v[1], -1.0, v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
|
||||
|
||||
def test_inf_minus_inf(self):
|
||||
"""+inf - inf = NaN."""
|
||||
import math
|
||||
pos_inf = 0x7f800000
|
||||
neg_inf = 0xff800000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], pos_inf),
|
||||
s_mov_b32(s[1], neg_inf),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_sub_f32_e32(v[2], v[0], v[1]), # inf - (-inf) = inf
|
||||
v_add_f32_e32(v[3], v[0], v[1]), # inf + (-inf) = NaN
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], pos_inf, "inf - (-inf) = inf")
|
||||
self.assertTrue(math.isnan(i2f(st.vgpr[0][3])), "inf + (-inf) = NaN")
|
||||
|
||||
def test_denormal_f32_mul_ftz(self):
|
||||
"""Denormal * normal - RDNA3 flushes denormals to zero (FTZ mode)."""
|
||||
smallest_denorm = 0x00000001 # Smallest positive denormal
|
||||
instructions = [
|
||||
s_mov_b32(s[0], smallest_denorm),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mul_f32_e32(v[1], 2.0, v[0]), # Denormal input gets flushed to 0
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0x00000000)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,538 @@
|
||||
"""Tests for VOP3P instructions - packed 16-bit vector operations.
|
||||
|
||||
Includes: v_pk_add_f16, v_pk_mul_f16, v_pk_fma_f16, v_pack_b32_f16, v_wmma_*, v_dot2_*
|
||||
"""
|
||||
import unittest
|
||||
from extra.assembly.amd.test.hw.helpers import *
|
||||
|
||||
class TestPackInstructions(unittest.TestCase):
|
||||
"""Tests for pack instructions."""
|
||||
|
||||
def test_v_pack_b32_f16(self):
|
||||
"""V_PACK_B32_F16 packs two f16 values into one 32-bit register."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pack_b32_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x40003c00, f"Expected 0x40003c00, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_opsel_hi_hi(self):
|
||||
"""V_PACK_B32_F16 with opsel to read high halves."""
|
||||
inst = v_pack_b32_f16(v[2], v[0], v[1])
|
||||
inst._values['opsel'] = 0b0011
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
|
||||
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
inst,
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x44004000, f"Expected 0x44004000, got 0x{result:08x}")
|
||||
|
||||
|
||||
class TestPackMore(unittest.TestCase):
|
||||
"""Additional pack instruction tests."""
|
||||
|
||||
def test_v_pack_b32_f16_basic(self):
|
||||
"""V_PACK_B32_F16 packs two f16 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pack_b32_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x40003c00, f"Expected 0x40003c00, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_with_cvt(self):
|
||||
"""V_PACK_B32_F16 after V_CVT_F16_F32 conversions."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3f800000),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
v_cvt_f16_f32_e32(v[2], v[0]),
|
||||
v_cvt_f16_f32_e32(v[3], v[1]),
|
||||
v_pack_b32_f16(v[4], v[2], v[3]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][4]
|
||||
self.assertEqual(result, 0x3c003c00, f"Expected 0x3c003c00, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_packed_sources(self):
|
||||
"""V_PACK_B32_F16 with packed f16 sources (reads lo halves)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
|
||||
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pack_b32_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
# Expected: hi=v1.lo=0x4200 (3.0), lo=v0.lo=0x3c00 (1.0) -> 0x42003c00
|
||||
self.assertEqual(result, 0x42003c00, f"Expected 0x42003c00, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_opsel_lo_hi(self):
|
||||
"""V_PACK_B32_F16 with opsel=0b0010 to read lo from src0, hi from src1."""
|
||||
inst = v_pack_b32_f16(v[2], v[0], v[1])
|
||||
inst._values['opsel'] = 0b0010
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00),
|
||||
s_mov_b32(s[1], 0x44004200),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
inst,
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x44003c00, f"Expected 0x44003c00, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_opsel_hi_lo(self):
|
||||
"""V_PACK_B32_F16 with opsel=0b0001 to read hi from src0, lo from src1."""
|
||||
inst = v_pack_b32_f16(v[2], v[0], v[1])
|
||||
inst._values['opsel'] = 0b0001
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00),
|
||||
s_mov_b32(s[1], 0x44004200),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
inst,
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x42004000, f"Expected 0x42004000, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_zeros(self):
|
||||
"""V_PACK_B32_F16 with zero values."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
v_pack_b32_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0)
|
||||
|
||||
def test_v_pack_b32_f16_both_positive(self):
|
||||
"""V_PACK_B32_F16 with positive f16 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x4200), # f16 3.0
|
||||
s_mov_b32(s[1], 0x4400), # f16 4.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pack_b32_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x44004200, f"Expected 0x44004200, got 0x{result:08x}")
|
||||
|
||||
|
||||
class TestFmaMix(unittest.TestCase):
|
||||
"""Tests for V_FMA_MIX_F32 and V_FMA_MIXLO_F16."""
|
||||
|
||||
def test_v_fma_mix_f32_all_f32_sources(self):
|
||||
"""V_FMA_MIX_F32 with all f32 sources."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(2.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(3.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], f2i(1.0)),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = i2f(st.vgpr[0][3])
|
||||
self.assertAlmostEqual(result, 7.0, places=5)
|
||||
|
||||
def test_v_fma_mix_f32_src2_f16_lo(self):
|
||||
"""V_FMA_MIX_F32 with src2 as f16 from lo bits."""
|
||||
from extra.assembly.amd.pcode import f32_to_f16
|
||||
f16_2 = f32_to_f16(2.0)
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(1.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(3.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], f16_2),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = i2f(st.vgpr[0][3])
|
||||
self.assertAlmostEqual(result, 5.0, places=5)
|
||||
|
||||
def test_v_fma_mix_f32_src2_f16_hi(self):
|
||||
"""V_FMA_MIX_F32 with src2 as f16 from hi bits."""
|
||||
from extra.assembly.amd.pcode import f32_to_f16
|
||||
f16_2 = f32_to_f16(2.0)
|
||||
val = (f16_2 << 16) | 0
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(1.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(3.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], val),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=4, opsel_hi=0, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = i2f(st.vgpr[0][3])
|
||||
self.assertAlmostEqual(result, 5.0, places=5)
|
||||
|
||||
def test_v_fma_mix_f32_with_abs(self):
|
||||
"""V_FMA_MIX_F32 with abs modifier on src2."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(2.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(3.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], f2i(-1.0)),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0, neg_hi=4),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = i2f(st.vgpr[0][3])
|
||||
self.assertAlmostEqual(result, 7.0, places=5)
|
||||
|
||||
def test_v_fma_mixlo_f16(self):
|
||||
"""V_FMA_MIXLO_F16 writes to low 16 bits of destination."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(2.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(3.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], f2i(1.0)),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
s_mov_b32(s[3], 0xdead0000),
|
||||
v_mov_b32_e32(v[3], s[3]),
|
||||
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
lo = _f16(st.vgpr[0][3] & 0xffff)
|
||||
hi = (st.vgpr[0][3] >> 16) & 0xffff
|
||||
self.assertAlmostEqual(lo, 7.0, places=1)
|
||||
self.assertEqual(hi, 0xdead, f"hi should be preserved, got 0x{hi:04x}")
|
||||
|
||||
def test_v_fma_mixlo_f16_all_f32_sources(self):
|
||||
"""V_FMA_MIXLO_F16 with all f32 sources."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(1.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(2.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], f2i(3.0)),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
v_mov_b32_e32(v[3], 0),
|
||||
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
lo = _f16(st.vgpr[0][3] & 0xffff)
|
||||
# 1*2+3 = 5
|
||||
self.assertAlmostEqual(lo, 5.0, places=1)
|
||||
|
||||
def test_v_fma_mixlo_f16_sin_case(self):
|
||||
"""V_FMA_MIXLO_F16 case from sin kernel."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3f800000), # f32 1.0
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
s_mov_b32(s[1], 0xaf05a309), # f32 tiny negative
|
||||
s_mov_b32(s[6], s[1]),
|
||||
s_mov_b32(s[2], 0xc0490fdb), # f32 -π
|
||||
v_mov_b32_e32(v[5], s[2]),
|
||||
s_mov_b32(s[3], 0x3f800000),
|
||||
v_mov_b32_e32(v[3], s[3]),
|
||||
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[3], src1=s[6], src2=v[5], opsel=0, opsel_hi=0, opsel_hi2=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
lo = _f16(st.vgpr[0][3] & 0xffff)
|
||||
self.assertAlmostEqual(lo, -3.14159, delta=0.01)
|
||||
|
||||
|
||||
class TestVOP3P(unittest.TestCase):
|
||||
"""Tests for VOP3P packed 16-bit operations."""
|
||||
|
||||
def test_v_pk_add_f16_basic(self):
|
||||
"""V_PK_ADD_F16 adds two packed f16 values."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
|
||||
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pk_add_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
self.assertAlmostEqual(lo, 4.0, places=2)
|
||||
self.assertAlmostEqual(hi, 6.0, places=2)
|
||||
|
||||
def test_v_pk_mul_f16_basic(self):
|
||||
"""V_PK_MUL_F16 multiplies two packed f16 values."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x42004000), # hi=3.0, lo=2.0
|
||||
s_mov_b32(s[1], 0x45004400), # hi=5.0, lo=4.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pk_mul_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
self.assertAlmostEqual(lo, 8.0, places=1)
|
||||
self.assertAlmostEqual(hi, 15.0, places=1)
|
||||
|
||||
def test_v_pk_fma_f16_basic(self):
|
||||
"""V_PK_FMA_F16: D = A * B + C for packed f16."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x42004000), # A: hi=3.0, lo=2.0
|
||||
s_mov_b32(s[1], 0x45004400), # B: hi=5.0, lo=4.0
|
||||
s_mov_b32(s[2], 0x3c003c00), # C: hi=1.0, lo=1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
v_pk_fma_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][3]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
self.assertAlmostEqual(lo, 9.0, places=1) # 2*4+1
|
||||
self.assertAlmostEqual(hi, 16.0, places=0) # 3*5+1
|
||||
|
||||
def test_v_pk_add_f16_with_inline_constant(self):
|
||||
"""V_PK_ADD_F16 with inline constant POS_ONE (1.0).
|
||||
Inline constants for VOP3P are f16 values in the low 16 bits only.
|
||||
hi half of inline constant is 0, so hi result = v0.hi + 0 = 1.0.
|
||||
"""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c003c00), # packed f16: hi=1.0, lo=1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_pk_add_f16(v[1], v[0], SrcEnum.POS_ONE, opsel_hi=3, opsel_hi2=1), # Add inline constant 1.0
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][1]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
# lo = 1.0 + 1.0 = 2.0, hi = 1.0 + 0.0 = 1.0 (inline const hi half is 0)
|
||||
self.assertAlmostEqual(lo, 2.0, places=2)
|
||||
self.assertAlmostEqual(hi, 1.0, places=2)
|
||||
|
||||
def test_v_pk_mul_f16_with_inline_constant(self):
|
||||
"""V_PK_MUL_F16 with inline constant POS_TWO (2.0).
|
||||
Inline constant has value only in low 16 bits, hi is 0.
|
||||
"""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
# v0 = packed (3.0, 4.0), multiply by POS_TWO
|
||||
# lo = 3.0 * 2.0 = 6.0, hi = 4.0 * 0.0 = 0.0 (inline const hi is 0)
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x44004200), # packed f16: hi=4.0, lo=3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_pk_mul_f16(v[1], v[0], SrcEnum.POS_TWO, opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][1]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
self.assertAlmostEqual(lo, 6.0, places=1)
|
||||
self.assertAlmostEqual(hi, 0.0, places=1)
|
||||
|
||||
|
||||
class TestWMMA(unittest.TestCase):
|
||||
"""Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions."""
|
||||
|
||||
def test_v_wmma_f32_16x16x16_f16_all_ones(self):
|
||||
"""V_WMMA_F32_16X16X16_F16 with all ones produces 16.0."""
|
||||
instructions = []
|
||||
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
|
||||
for i in range(16, 32):
|
||||
instructions.append(v_mov_b32_e32(v[i], s[0]))
|
||||
for i in range(8):
|
||||
instructions.append(v_mov_b32_e32(v[i], 0))
|
||||
instructions.append(v_wmma_f32_16x16x16_f16(v[0], v[16], v[24], v[0]))
|
||||
st = run_program(instructions, n_lanes=32)
|
||||
expected = f2i(16.0)
|
||||
for lane in range(32):
|
||||
for reg in range(8):
|
||||
result = st.vgpr[lane][reg]
|
||||
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
|
||||
|
||||
def test_v_wmma_f32_16x16x16_f16_with_accumulator(self):
|
||||
"""V_WMMA_F32_16X16X16_F16 with non-zero accumulator."""
|
||||
instructions = []
|
||||
instructions.append(s_mov_b32(s[0], 0x3c003c00))
|
||||
instructions.append(s_mov_b32(s[1], f2i(5.0)))
|
||||
for i in range(16, 32):
|
||||
instructions.append(v_mov_b32_e32(v[i], s[0]))
|
||||
for i in range(8):
|
||||
instructions.append(v_mov_b32_e32(v[i], s[1]))
|
||||
instructions.append(v_wmma_f32_16x16x16_f16(v[0], v[16], v[24], v[0]))
|
||||
st = run_program(instructions, n_lanes=32)
|
||||
expected = f2i(21.0) # 16 + 5
|
||||
for lane in range(32):
|
||||
for reg in range(8):
|
||||
result = st.vgpr[lane][reg]
|
||||
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
|
||||
|
||||
|
||||
class TestSpecialOps(unittest.TestCase):
|
||||
"""Tests for special operations (SAD, PERM, DOT2)."""
|
||||
|
||||
def test_v_sad_u8_basic(self):
|
||||
"""V_SAD_U8 computes sum of absolute differences."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x04030201), # bytes: 1, 2, 3, 4
|
||||
s_mov_b32(s[1], 0x05040302), # bytes: 2, 3, 4, 5
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_sad_u8(v[3], v[0], v[1], v[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# |1-2| + |2-3| + |3-4| + |4-5| = 1 + 1 + 1 + 1 = 4
|
||||
self.assertEqual(st.vgpr[0][3], 4)
|
||||
|
||||
def test_v_sad_u8_identical_bytes(self):
|
||||
"""V_SAD_U8 with identical inputs returns accumulator."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x04030201),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], 10),
|
||||
v_mov_b32_e32(v[2], s[1]),
|
||||
v_sad_u8(v[3], v[0], v[0], v[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Same inputs -> SAD = 0, result = accumulator = 10
|
||||
self.assertEqual(st.vgpr[0][3], 10)
|
||||
|
||||
def test_v_sad_u16_basic(self):
|
||||
"""V_SAD_U16 computes sum of absolute differences of u16 pairs."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x00030001), # hi=3, lo=1
|
||||
s_mov_b32(s[1], 0x00050002), # hi=5, lo=2
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_sad_u16(v[3], v[0], v[1], v[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# |1-2| + |3-5| = 1 + 2 = 3
|
||||
self.assertEqual(st.vgpr[0][3], 3)
|
||||
|
||||
def test_v_sad_u32_basic(self):
|
||||
"""V_SAD_U32 computes absolute difference of u32 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 100),
|
||||
s_mov_b32(s[1], 70),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_sad_u32(v[3], v[0], v[1], v[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# |100-70| = 30
|
||||
self.assertEqual(st.vgpr[0][3], 30)
|
||||
|
||||
def test_v_msad_u8_masked(self):
|
||||
"""V_MSAD_U8 masked SAD operation."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x04030201),
|
||||
s_mov_b32(s[1], 0x05040302),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_msad_u8(v[3], v[0], v[1], v[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# V_MSAD_U8 skips bytes where src0 is 0
|
||||
# Since no bytes are 0, result same as V_SAD_U8 = 4
|
||||
self.assertEqual(st.vgpr[0][3], 4)
|
||||
|
||||
def test_v_perm_b32_select_bytes(self):
|
||||
"""V_PERM_B32 selects bytes from two sources.
|
||||
|
||||
V_PERM_B32 concatenates {S1, S0} as a 64-bit value with S1 in low 32 bits.
|
||||
Selector byte values 0-3 select from S1, values 4-7 select from S0.
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x44332211), # src0: bytes 4-7 in 64-bit view
|
||||
s_mov_b32(s[1], 0x88776655), # src1: bytes 0-3 in 64-bit view
|
||||
s_mov_b32(s[2], 0x07060504), # select bytes 4,5,6,7 (from src0)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_perm_b32(v[2], v[0], v[1], s[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0x44332211)
|
||||
|
||||
def test_v_dot2_f32_bf16_basic(self):
|
||||
"""V_DOT2_F32_BF16 computes dot product of bf16 pairs."""
|
||||
# bf16 1.0 = 0x3f80, bf16 2.0 = 0x4000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3f803f80), # packed bf16: lo=1.0, hi=1.0
|
||||
s_mov_b32(s[1], 0x40003f80), # packed bf16: lo=1.0, hi=2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_dot2_f32_bf16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# 1.0*1.0 + 1.0*2.0 + 0 = 3.0
|
||||
result = i2f(st.vgpr[0][3])
|
||||
self.assertAlmostEqual(result, 3.0, places=4)
|
||||
|
||||
|
||||
class TestPackedMixedSigns(unittest.TestCase):
|
||||
"""Tests for packed operations with mixed sign values."""
|
||||
|
||||
def test_pk_add_f16_mixed_signs(self):
|
||||
"""V_PK_ADD_F16 with mixed positive/negative values."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xc0003c00), # packed: hi=-2.0, lo=1.0
|
||||
s_mov_b32(s[1], 0x3c003c00), # packed: hi=1.0, lo=1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pk_add_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
self.assertAlmostEqual(lo, 2.0, places=2) # 1.0 + 1.0
|
||||
self.assertAlmostEqual(hi, -1.0, places=2) # -2.0 + 1.0
|
||||
|
||||
def test_pk_mul_f16_zero(self):
|
||||
"""V_PK_MUL_F16 with zero."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40004000), # packed: 2.0, 2.0
|
||||
s_mov_b32(s[1], 0x00000000), # packed: 0.0, 0.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pk_mul_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x00000000, "2.0 * 0.0 should be 0.0")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,486 @@
|
||||
"""Tests for VOPC instructions - vector compare operations.
|
||||
|
||||
Includes: v_cmp_class_f32, v_cmp_class_f16, v_cmp_eq_*, v_cmp_lt_*, v_cmp_gt_*
|
||||
"""
|
||||
import unittest
|
||||
from extra.assembly.amd.test.hw.helpers import *
|
||||
|
||||
VCC = 106 # SGPR index for VCC_LO
|
||||
|
||||
class TestCmpClass(unittest.TestCase):
|
||||
"""Tests for V_CMP_CLASS_F32 float classification."""
|
||||
|
||||
def test_cmp_class_quiet_nan(self):
|
||||
"""V_CMP_CLASS_F32 detects quiet NaN."""
|
||||
quiet_nan = 0x7fc00000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], quiet_nan),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0b0000000010), # bit 1 = quiet NaN
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
|
||||
|
||||
def test_cmp_class_signaling_nan(self):
|
||||
"""V_CMP_CLASS_F32 detects signaling NaN."""
|
||||
signal_nan = 0x7f800001
|
||||
instructions = [
|
||||
s_mov_b32(s[0], signal_nan),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0b0000000001), # bit 0 = signaling NaN
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
|
||||
|
||||
def test_cmp_class_positive_inf(self):
|
||||
"""V_CMP_CLASS_F32 detects +inf."""
|
||||
pos_inf = 0x7f800000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], pos_inf),
|
||||
s_mov_b32(s[1], 0b1000000000), # bit 9 = +inf
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect +inf")
|
||||
|
||||
def test_cmp_class_negative_inf(self):
|
||||
"""V_CMP_CLASS_F32 detects -inf."""
|
||||
neg_inf = 0xff800000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], neg_inf),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0b0000000100), # bit 2 = -inf
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect -inf")
|
||||
|
||||
def test_cmp_class_normal_positive(self):
|
||||
"""V_CMP_CLASS_F32 detects positive normal."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 1.0),
|
||||
s_mov_b32(s[1], 0b0100000000), # bit 8 = positive normal
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
|
||||
|
||||
def test_cmp_class_normal_negative(self):
|
||||
"""V_CMP_CLASS_F32 detects negative normal."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], -1.0),
|
||||
v_mov_b32_e32(v[1], 0b0000001000), # bit 3 = negative normal
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
|
||||
|
||||
def test_cmp_class_quiet_nan_not_signaling(self):
|
||||
"""Quiet NaN does not match signaling NaN mask."""
|
||||
quiet_nan = 0x7fc00000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], quiet_nan),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0b0000000001), # bit 0 = signaling NaN only
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "Quiet NaN should not match signaling mask")
|
||||
|
||||
def test_cmp_class_signaling_nan_not_quiet(self):
|
||||
"""Signaling NaN does not match quiet NaN mask."""
|
||||
signal_nan = 0x7f800001
|
||||
instructions = [
|
||||
s_mov_b32(s[0], signal_nan),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0b0000000010), # bit 1 = quiet NaN only
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "Signaling NaN should not match quiet mask")
|
||||
|
||||
def test_v_cmp_sets_vcc_bits(self):
|
||||
"""V_CMP_EQ sets VCC bits based on per-lane comparison."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 5),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
v_cmp_eq_u32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=4)
|
||||
self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
|
||||
|
||||
|
||||
class TestCmpClassF16(unittest.TestCase):
|
||||
"""Tests for V_CMP_CLASS_F16 float classification.
|
||||
|
||||
Class bit mapping:
|
||||
bit 0 = signaling NaN
|
||||
bit 1 = quiet NaN
|
||||
bit 2 = -infinity
|
||||
bit 3 = -normal
|
||||
bit 4 = -denormal
|
||||
bit 5 = -zero
|
||||
bit 6 = +zero
|
||||
bit 7 = +denormal
|
||||
bit 8 = +normal
|
||||
bit 9 = +infinity
|
||||
"""
|
||||
|
||||
def test_cmp_class_f16_positive_zero(self):
|
||||
"""V_CMP_CLASS_F16: +zero matches bit 6."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0x0000), # f16 +0.0
|
||||
v_mov_b32_e32(v[1], 0x40), # bit 6 = +zero
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect positive zero")
|
||||
|
||||
def test_cmp_class_f16_negative_zero(self):
|
||||
"""V_CMP_CLASS_F16: -zero matches bit 5."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x8000), # f16 -0.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x20), # bit 5 = -zero
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect negative zero")
|
||||
|
||||
def test_cmp_class_f16_positive_normal(self):
|
||||
"""V_CMP_CLASS_F16: +1.0 (normal) matches bit 8."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 +1.0
|
||||
s_mov_b32(s[1], 0x100), # bit 8 = +normal
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
|
||||
|
||||
def test_cmp_class_f16_negative_normal(self):
|
||||
"""V_CMP_CLASS_F16: -1.0 (normal) matches bit 3."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xbc00), # f16 -1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x08), # bit 3 = -normal
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
|
||||
|
||||
def test_cmp_class_f16_positive_infinity(self):
|
||||
"""V_CMP_CLASS_F16: +inf matches bit 9."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7c00), # f16 +inf
|
||||
s_mov_b32(s[1], 0x200), # bit 9 = +inf
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect positive infinity")
|
||||
|
||||
def test_cmp_class_f16_negative_infinity(self):
|
||||
"""V_CMP_CLASS_F16: -inf matches bit 2."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xfc00), # f16 -inf
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x04), # bit 2 = -inf
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect negative infinity")
|
||||
|
||||
def test_cmp_class_f16_quiet_nan(self):
|
||||
"""V_CMP_CLASS_F16: quiet NaN matches bit 1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7e00), # f16 quiet NaN
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x02), # bit 1 = quiet NaN
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
|
||||
|
||||
def test_cmp_class_f16_signaling_nan(self):
|
||||
"""V_CMP_CLASS_F16: signaling NaN matches bit 0."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7c01), # f16 signaling NaN
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x01), # bit 0 = signaling NaN
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
|
||||
|
||||
def test_cmp_class_f16_positive_denormal(self):
|
||||
"""V_CMP_CLASS_F16: positive denormal matches bit 7."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 1), # f16 +denormal (0x0001)
|
||||
v_mov_b32_e32(v[1], 0x80), # bit 7 = +denormal
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect positive denormal")
|
||||
|
||||
def test_cmp_class_f16_negative_denormal(self):
|
||||
"""V_CMP_CLASS_F16: negative denormal matches bit 4."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x8001), # f16 -denormal
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x10), # bit 4 = -denormal
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect negative denormal")
|
||||
|
||||
def test_cmp_class_f16_combined_mask_zeros(self):
|
||||
"""V_CMP_CLASS_F16: mask 0x60 covers both +zero and -zero."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0), # f16 +0.0
|
||||
v_mov_b32_e32(v[1], 0x60), # bits 5 and 6 (+-zero)
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x60")
|
||||
|
||||
def test_cmp_class_f16_combined_mask_1f8(self):
|
||||
"""V_CMP_CLASS_F16: mask 0x1f8 covers -normal,-denorm,-zero,+zero,+denorm,+normal.
|
||||
|
||||
This is the exact mask used in the f16 sin kernel at PC=46.
|
||||
"""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0), # f16 +0.0
|
||||
s_mov_b32(s[0], 0x1f8),
|
||||
v_mov_b32_e32(v[1], s[0]), # mask 0x1f8
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x1f8")
|
||||
|
||||
def test_cmp_class_f16_vop3_encoding(self):
|
||||
"""V_CMP_CLASS_F16 in VOP3 encoding (v_cmp_class_f16_e64)."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0), # f16 +0.0
|
||||
s_mov_b32(s[0], 0x1f8), # class mask
|
||||
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with VOP3 encoding")
|
||||
|
||||
def test_cmp_class_f16_vop3_normal_positive(self):
|
||||
"""V_CMP_CLASS_F16 VOP3 encoding with +1.0 (normal)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 +1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], 0x1f8), # class mask
|
||||
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +1.0 (normal) with mask 0x1f8")
|
||||
|
||||
def test_cmp_class_f16_vop3_nan_fails_mask(self):
|
||||
"""V_CMP_CLASS_F16 VOP3: NaN should NOT match mask 0x1f8 (no NaN bits set)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7e00), # f16 quiet NaN
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], 0x1f8), # class mask
|
||||
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for NaN with mask 0x1f8 (no NaN bits)")
|
||||
|
||||
def test_cmp_class_f16_vop3_inf_fails_mask(self):
|
||||
"""V_CMP_CLASS_F16 VOP3: +inf should NOT match mask 0x1f8 (no inf bits set)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7c00), # f16 +inf
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], 0x1f8), # class mask
|
||||
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for +inf with mask 0x1f8 (no inf bits)")
|
||||
|
||||
|
||||
class TestCmpInt(unittest.TestCase):
|
||||
"""Tests for integer comparison operations."""
|
||||
|
||||
def test_v_cmp_eq_u32(self):
|
||||
"""V_CMP_EQ_U32 sets VCC bits based on per-lane comparison."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 5),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
v_cmp_eq_u32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=4)
|
||||
self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
|
||||
|
||||
def test_cmp_eq_u16_opsel_lo_lo(self):
|
||||
"""V_CMP_EQ_U16 comparing lo halves."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x12340005), # lo=5, hi=0x1234
|
||||
s_mov_b32(s[1], 0xABCD0005), # lo=5, hi=0xABCD
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cmp_eq_u16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Lo halves should be equal")
|
||||
|
||||
def test_cmp_eq_u16_opsel_hi_hi(self):
|
||||
"""V_CMP_EQ_U16 comparing hi halves with VOP3 opsel.
|
||||
|
||||
VOPC doesn't have opsel, so we use VOP3 form for hi-half comparisons.
|
||||
VOP3 compares write result to SGPR via vdst field.
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(s[2], 0x00051234), # hi=5, lo=0x1234
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0x0005ABCD), # hi=5, lo=0xABCD
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
# opsel=3 means compare hi halves, vdst=v[0] actually writes to s[0]
|
||||
VOP3(VOP3Op.V_CMP_EQ_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Result is in sgpr[0], not vcc
|
||||
self.assertEqual(st.sgpr[0] & 1, 1, "Hi halves should be equal: 5==5")
|
||||
|
||||
def test_cmp_eq_u16_opsel_hi_hi_equal(self):
|
||||
"""V_CMP_EQ_U16 VOP3 with opsel=3 compares hi halves (equal case)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[2], 0x12340005), # lo=5, hi=0x1234
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0x12340009), # lo=9, hi=0x1234
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
VOP3(VOP3Op.V_CMP_EQ_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[0] & 1, 1, "hi==hi should be true: 0x1234==0x1234")
|
||||
|
||||
def test_cmp_gt_u16_opsel_hi(self):
|
||||
"""V_CMP_GT_U16 VOP3 with opsel=3 compares hi halves."""
|
||||
instructions = [
|
||||
s_mov_b32(s[2], 0x99990005), # lo=5, hi=0x9999
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0x12340005), # lo=5, hi=0x1234
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
VOP3(VOP3Op.V_CMP_GT_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[0] & 1, 1, "hi>hi should be true: 0x9999>0x1234")
|
||||
|
||||
|
||||
class TestCmpFloat(unittest.TestCase):
|
||||
"""Tests for float comparison operations."""
|
||||
|
||||
def test_v_cmp_lt_f16_vsrc1_hi(self):
|
||||
"""V_CMP_LT_F16 with both operands from high half using VOP3 opsel."""
|
||||
instructions = [
|
||||
s_mov_b32(s[2], 0x3c000000), # hi=1.0 (f16), lo=0
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0x40000000), # hi=2.0 (f16), lo=0
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
# opsel=3 means read hi halves for both src0 and src1
|
||||
VOP3(VOP3Op.V_CMP_LT_F16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Result is in sgpr[0]
|
||||
self.assertEqual(st.sgpr[0] & 1, 1, "1.0 < 2.0 should be true")
|
||||
|
||||
def test_v_cmp_gt_f16_vsrc1_hi(self):
|
||||
"""V_CMP_GT_F16 with both operands from high half using VOP3 opsel."""
|
||||
instructions = [
|
||||
s_mov_b32(s[2], 0x40000000), # hi=2.0 (f16), lo=0
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0x3c000000), # hi=1.0 (f16), lo=0
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
# opsel=3 means read hi halves for both src0 and src1
|
||||
VOP3(VOP3Op.V_CMP_GT_F16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Result is in sgpr[0]
|
||||
self.assertEqual(st.sgpr[0] & 1, 1, "2.0 > 1.0 should be true")
|
||||
|
||||
def test_v_cmp_eq_f16_vsrc1_hi_equal(self):
|
||||
"""v_cmp_eq_f16 with equal low and high halves."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x42004200), # hi=3.0 (0x4200), lo=3.0 (0x4200)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_cmp_eq_f16_e32(v[0], v[0].h),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (3.0 == 3.0)")
|
||||
|
||||
def test_v_cmp_neq_f16_vsrc1_hi(self):
|
||||
"""v_cmp_neq_f16 with different low and high halves."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # hi=2.0 (0x4000), lo=1.0 (0x3c00)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_cmp_lg_f16_e32(v[0], v[0].h),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (1.0 != 2.0)")
|
||||
|
||||
def test_v_cmp_nge_f16_inf_self(self):
|
||||
"""v_cmp_nge_f16 comparing -inf with itself (unordered less than).
|
||||
|
||||
Regression test: -inf < -inf should be false (IEEE 754).
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFC00FC00), # both halves = -inf (0xFC00)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_cmp_nge_f16_e32(v[0], v[0].h),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "Expected vcc=0 (-inf >= -inf)")
|
||||
|
||||
def test_v_cmp_f16_multilane(self):
|
||||
"""v_cmp_lt_f16 with vsrc1=v128 across multiple lanes."""
|
||||
instructions = [
|
||||
# Lane 0: v0 = 0x40003c00 (hi=2.0, lo=1.0) -> 1.0 < 2.0 = true
|
||||
# Lane 1: v0 = 0x3c004000 (hi=1.0, lo=2.0) -> 2.0 < 1.0 = false
|
||||
v_mov_b32_e32(v[0], 0x40003c00), # default
|
||||
v_cmp_eq_u32_e32(1, v[255]), # vcc = (lane == 1)
|
||||
v_cndmask_b32_e64(v[0], v[0], 0x3c004000, SrcEnum.VCC_LO),
|
||||
v_cmp_lt_f16_e32(v[0], v[0].h),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=2)
|
||||
self.assertEqual(st.vcc & 1, 1, "Lane 0: expected vcc=1 (1.0 < 2.0)")
|
||||
self.assertEqual((st.vcc >> 1) & 1, 0, "Lane 1: expected vcc=0 (2.0 < 1.0)")
|
||||
|
||||
|
||||
class TestVCCBehavior(unittest.TestCase):
|
||||
"""Tests for VCC condition code behavior."""
|
||||
|
||||
def test_vcc_all_lanes_true(self):
|
||||
"""VCC should have all bits set when all lanes compare true."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 5),
|
||||
v_mov_b32_e32(v[1], 5),
|
||||
v_cmp_eq_u32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=32)
|
||||
self.assertEqual(st.vcc, 0xFFFFFFFF, "All 32 lanes should be true")
|
||||
|
||||
def test_vcc_lane_dependent(self):
|
||||
"""VCC should differ per lane based on lane_id comparison."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 16),
|
||||
v_cmp_lt_u32_e32(v[255], v[0]), # lanes 0-15 are < 16
|
||||
]
|
||||
st = run_program(instructions, n_lanes=32)
|
||||
self.assertEqual(st.vcc & 0xFFFF, 0xFFFF, "Lanes 0-15 should be true")
|
||||
self.assertEqual(st.vcc >> 16, 0x0000, "Lanes 16-31 should be false")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
+87
-152
@@ -7,18 +7,20 @@ from pathlib import Path
|
||||
# This allows generating AMD GPU kernels without requiring real hardware
|
||||
os.environ["AMD"] = "1"
|
||||
os.environ["MOCKGPU"] = "1"
|
||||
os.environ["PYTHON_REMU"] = "1"
|
||||
|
||||
from extra.assembly.rdna3.emu import WaveState, decode_program, step_wave, WAVE_SIZE
|
||||
from extra.assembly.amd.emu import WaveState, decode_program, WAVE_SIZE, set_valid_mem_ranges, LDSMem
|
||||
from extra.assembly.amd.test.helpers import KernelInfo
|
||||
from extra.assembly.amd.test.bench_emu import REMU_PATH
|
||||
|
||||
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
|
||||
def _is_f32_nan(bits: int) -> bool:
|
||||
"""Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
|
||||
return (bits & 0x7f800000) == 0x7f800000 and (bits & 0x007fffff) != 0
|
||||
|
||||
@dataclass
|
||||
class KernelInfo:
|
||||
code: bytes
|
||||
global_size: tuple[int, int, int]
|
||||
local_size: tuple[int, int, int]
|
||||
buf_idxs: list[int] # indices into shared buffer pool
|
||||
buf_sizes: list[int] # sizes for each buffer index
|
||||
def _vals_equal(a: int, b: int) -> bool:
|
||||
"""Compare two 32-bit values, treating all NaN bit patterns as equal."""
|
||||
if a == b: return True
|
||||
return _is_f32_nan(a) and _is_f32_nan(b)
|
||||
|
||||
@dataclass
|
||||
class StateSnapshot:
|
||||
@@ -29,20 +31,20 @@ class StateSnapshot:
|
||||
sgpr: list[int]
|
||||
vgpr: list[list[int]]
|
||||
|
||||
def diff(self, other: 'StateSnapshot', n_lanes: int) -> list[str]:
|
||||
def diff(self, other: 'StateSnapshot', n_lanes: int, arrow: str = " vs ") -> list[str]:
|
||||
"""Return list of differences between two states."""
|
||||
diffs = []
|
||||
if self.pc != other.pc: diffs.append(f"pc: {self.pc} vs {other.pc}")
|
||||
if self.scc != other.scc: diffs.append(f"scc: {self.scc} vs {other.scc}")
|
||||
if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x} vs 0x{other.vcc:08x}")
|
||||
if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x} vs 0x{other.exec_mask:08x}")
|
||||
if self.pc != other.pc: diffs.append(f"pc: {self.pc}{arrow}{other.pc}")
|
||||
if self.scc != other.scc: diffs.append(f"scc: {self.scc}{arrow}{other.scc}")
|
||||
if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x}{arrow}0x{other.vcc:08x}")
|
||||
if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x}{arrow}0x{other.exec_mask:08x}")
|
||||
for i, (a, b) in enumerate(zip(self.sgpr, other.sgpr)):
|
||||
# Skip VCC_LO/HI (106/107) and EXEC_LO/HI (126/127) as they alias vcc/exec_mask which are compared separately
|
||||
if i in (106, 107, 126, 127): continue
|
||||
if a != b: diffs.append(f"sgpr[{i}]: 0x{a:08x} vs 0x{b:08x}")
|
||||
if not _vals_equal(a, b): diffs.append(f"sgpr[{i}]: 0x{a:08x}{arrow}0x{b:08x}")
|
||||
for lane in range(n_lanes):
|
||||
for i, (a, b) in enumerate(zip(self.vgpr[lane], other.vgpr[lane])):
|
||||
if a != b: diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x} vs 0x{b:08x}")
|
||||
if not _vals_equal(a, b): diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x}{arrow}0x{b:08x}")
|
||||
return diffs
|
||||
|
||||
class CStateSnapshot(ctypes.Structure):
|
||||
@@ -89,19 +91,15 @@ class PythonEmulator:
|
||||
def __init__(self):
|
||||
self.state: WaveState | None = None
|
||||
self.program: dict | None = None
|
||||
self.lds: bytearray | None = None
|
||||
self.n_lanes = 0
|
||||
|
||||
def create(self, kernel: bytes, n_lanes: int):
|
||||
self.program = decode_program(kernel)
|
||||
self.state = WaveState()
|
||||
self.state = WaveState(LDSMem(bytearray(65536)), n_lanes)
|
||||
self.state.exec_mask = (1 << n_lanes) - 1
|
||||
self.lds = bytearray(65536)
|
||||
self.n_lanes = n_lanes
|
||||
|
||||
def step(self) -> int:
|
||||
assert self.program is not None and self.state is not None and self.lds is not None
|
||||
return step_wave(self.program, self.state, self.lds, self.n_lanes)
|
||||
assert self.program is not None and self.state is not None
|
||||
return self.program[self.state.pc]._dispatch(self.state, self.program[self.state.pc])
|
||||
def set_sgpr(self, idx: int, val: int):
|
||||
assert self.state is not None
|
||||
self.state.sgpr[idx] = val & 0xffffffff
|
||||
@@ -117,7 +115,7 @@ class PythonEmulator:
|
||||
|
||||
def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
|
||||
program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
|
||||
max_workgroups: int = 64) -> tuple[bool, str, int]:
|
||||
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
|
||||
total_steps = 0
|
||||
@@ -157,25 +155,53 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
|
||||
|
||||
if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
|
||||
|
||||
# Instructions with known Rust emulator bugs - sync Python to Rust after execution
|
||||
# v_div_scale/v_div_fixup: Rust has different VCC handling
|
||||
# v_cvt_f16_f32: Rust clears high 16 bits, but hardware (and Python) preserves them
|
||||
# s_add_i32/s_sub_i32: Rust has incorrect SCC overflow detection
|
||||
sync_after = any(x in inst_str for x in ('v_div_scale_f32', 'v_div_scale_f64', 'v_div_fixup_f32', 'v_div_fixup_f64',
|
||||
'v_cvt_f16_f32', 's_add_i32', 's_sub_i32'))
|
||||
diffs = rust_before.diff(python_before, n_lanes)
|
||||
if diffs:
|
||||
trace_lines = []
|
||||
for s, pc, d, rb, pb in trace[:-1]:
|
||||
for idx, (s, pc, d, rb, pb) in enumerate(trace):
|
||||
trace_lines.append(f" step {s}: PC={pc:3d} {d}")
|
||||
if trace.index((s, pc, d, rb, pb)) < len(trace) - 2:
|
||||
next_rb, next_pb = trace[trace.index((s, pc, d, rb, pb)) + 1][3:5]
|
||||
inst_diffs = rb.diff(next_rb, n_lanes)
|
||||
if inst_diffs: trace_lines.append(f" rust changes: {', '.join(inst_diffs[:3])}")
|
||||
if idx < len(trace) - 1:
|
||||
next_rb, next_pb = trace[idx + 1][3:5]
|
||||
rust_diffs = rb.diff(next_rb, n_lanes, "->")
|
||||
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(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(f" python: (no changes)")
|
||||
trace_str = "\n".join(trace_lines)
|
||||
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ:\n " + "\n ".join(diffs[:10]) + f"\n Recent instructions:\n{trace_str}", total_steps
|
||||
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()
|
||||
|
||||
if rust_result != python_result:
|
||||
# Rust returns 1 for unsupported instructions - skip test
|
||||
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)
|
||||
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:
|
||||
rust_after = rust.get_snapshot()
|
||||
for i in range(128): python.set_sgpr(i, rust_after.sgpr[i])
|
||||
for lane in range(n_lanes):
|
||||
for i in range(256): python.set_vgpr(lane, i, rust_after.vgpr[lane][i])
|
||||
assert python.state is not None
|
||||
python.state.pc, python.state.scc, python.state.vcc, python.state.exec_mask = rust_after.pc, rust_after.scc, rust_after.vcc, rust_after.exec_mask
|
||||
|
||||
if rust_result == -1:
|
||||
total_steps += step + 1
|
||||
break
|
||||
@@ -196,7 +222,6 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
|
||||
def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int, int], max_steps: int = 1000,
|
||||
debug: bool = False, trace_len: int = 10, buf_data: dict[int, bytes] | None = None) -> tuple[bool, str]:
|
||||
"""Run all kernels through both emulators with shared buffer pool."""
|
||||
from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
|
||||
if buf_data is None: buf_data = {}
|
||||
|
||||
# Allocate shared buffer pool with padding for over-reads (GPU loads up to 16 bytes at once)
|
||||
@@ -240,8 +265,6 @@ def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int
|
||||
def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list, max_steps: int = 1000, debug: bool = False,
|
||||
global_size: tuple[int, int, int] = (1, 1, 1), trace_len: int = 10) -> tuple[bool, str]:
|
||||
"""Run both emulators with memory set up for tinygrad kernels, executing all workgroups. Legacy wrapper."""
|
||||
from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
|
||||
|
||||
# Allocate buffers
|
||||
buffers = []
|
||||
for size in buf_sizes:
|
||||
@@ -315,7 +338,6 @@ def get_kernel_from_tinygrad(op_fn) -> tuple[bytes, tuple[int, int, int], tuple[
|
||||
k = kernels[-1]
|
||||
return k.code, k.global_size, k.local_size, k.buf_sizes
|
||||
|
||||
@unittest.skipUnless(REMU_PATH.exists(), "libremu.so not found")
|
||||
class TestTinygradKernels(unittest.TestCase):
|
||||
"""Compare emulators on real tinygrad-compiled kernels."""
|
||||
|
||||
@@ -324,143 +346,56 @@ class TestTinygradKernels(unittest.TestCase):
|
||||
ok, msg = compare_emulators_multi_kernel(kernels, buf_pool, max_steps=max_steps, buf_data=buf_data)
|
||||
self.assertTrue(ok, msg)
|
||||
|
||||
# Basic unary ops
|
||||
def test_neg(self): self._test_kernel(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
|
||||
def test_relu(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
|
||||
def test_exp(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).exp())
|
||||
def test_log(self): self._test_kernel(lambda T: T([1.0, 2.0, 3.0]).log())
|
||||
def test_sin(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).sin())
|
||||
def test_sqrt(self): self._test_kernel(lambda T: T([1.0, 4.0, 9.0]).sqrt())
|
||||
def test_recip(self): self._test_kernel(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
|
||||
|
||||
# Binary ops
|
||||
def test_add(self): self._test_kernel(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
|
||||
def test_sub(self): self._test_kernel(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
|
||||
def test_mul(self): self._test_kernel(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
|
||||
def test_div(self): self._test_kernel(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
|
||||
def test_max_binary(self): self._test_kernel(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
|
||||
# Basic ops - consolidated tests covering key instruction patterns
|
||||
def test_unary_ops(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu().exp().log().sqrt().reciprocal())
|
||||
def test_binary_ops(self): self._test_kernel(lambda T: (T([1.0, 2.0]) + T([3.0, 4.0])) * T([0.5, 0.5]) - T([1.0, 1.0]))
|
||||
def test_trig(self): self._test_kernel(lambda T: T([0.1, 1.0, 3.14, -1.0]*8).sin() + T([0.1, 1.0, 3.14, -1.0]*8).cos())
|
||||
def test_compare(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
|
||||
def test_bitwise(self): self._test_kernel(lambda T: (T([0xF0, 0x0F, 0xFF]*11).int() & T([0x0F, 0x0F, 0x00]*11).int()) | T([1]*33).int())
|
||||
def test_int_ops(self): self._test_kernel(lambda T: ((T.empty(64).int() + T.empty(64).int()) * T.empty(64).int()).float())
|
||||
|
||||
# Reductions
|
||||
def test_sum_reduce(self): self._test_kernel(lambda T: T.empty(64).sum())
|
||||
def test_max_reduce(self): self._test_kernel(lambda T: T.empty(64).max())
|
||||
def test_mean_reduce(self): self._test_kernel(lambda T: T.empty(32).mean())
|
||||
def test_reduce(self): self._test_kernel(lambda T: T.empty(64).sum() + T.empty(64).max())
|
||||
def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
|
||||
|
||||
# Matmul - various sizes
|
||||
def test_gemm_4x4(self): self._test_kernel(lambda T: T.empty(4, 4) @ T.empty(4, 4), max_steps=100000)
|
||||
def test_gemm_8x8(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=200000)
|
||||
@unittest.skip("too slow")
|
||||
def test_gemm_16x16(self): self._test_kernel(lambda T: T.empty(16, 16) @ T.empty(16, 16), max_steps=500000)
|
||||
def test_gemv(self): self._test_kernel(lambda T: T.empty(1, 16) @ T.empty(16, 16), max_steps=100000)
|
||||
# Matmul
|
||||
def test_gemm(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=100000)
|
||||
@unittest.skip("Rust emulator crashes on this kernel (assertion failure in thread.rs)")
|
||||
def test_gemm_fp16(self): self._test_kernel(lambda T: T.empty(16, 16).half() @ T.empty(16, 16).half(), max_steps=100000)
|
||||
|
||||
# Complex ops
|
||||
def test_softmax(self): self._test_kernel(lambda T: T.empty(16).softmax())
|
||||
def test_layernorm(self): self._test_kernel(lambda T: T.empty(8, 8).layernorm())
|
||||
|
||||
# Memory patterns
|
||||
def test_contiguous(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
|
||||
def test_reshape(self): self._test_kernel(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
|
||||
def test_expand(self): self._test_kernel(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
|
||||
def test_memory(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous() + T.empty(4, 1).expand(4, 4))
|
||||
|
||||
# Cast ops
|
||||
def test_cast_int(self): self._test_kernel(lambda T: T.empty(16).int().float())
|
||||
def test_cast_half(self): self._test_kernel(lambda T: T.empty(16).half().float())
|
||||
def test_cast(self): self._test_kernel(lambda T: T.empty(32).half().float() + T.empty(32).int().float())
|
||||
|
||||
# Min/max (uses comparison internally)
|
||||
def test_min_binary(self): self._test_kernel(lambda T: T([1.0, 5.0, 3.0]).minimum(T([3.0, 2.0, 4.0])))
|
||||
# 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)))
|
||||
|
||||
# Comparison ops (test VOPC instructions) - use 32+ elements to force vector instructions
|
||||
def test_cmp_lt(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
|
||||
def test_cmp_eq(self): self._test_kernel(lambda T: (T.empty(64) == T.empty(64)).where(T.empty(64), T.empty(64)))
|
||||
def test_where(self): self._test_kernel(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
|
||||
# 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)
|
||||
|
||||
# Bitwise ops
|
||||
def test_bitwise_and(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0xFF]).int() & T([0x0F, 0x0F, 0x00]).int())
|
||||
def test_bitwise_or(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0x00]).int() | T([0x0F, 0x0F, 0xFF]).int())
|
||||
def test_bitwise_xor(self): self._test_kernel(lambda T: T([0xFF, 0x0F, 0xF0]).int() ^ T([0x0F, 0xF0, 0xF0]).int())
|
||||
|
||||
# Integer ops - use 32+ elements to force vector instructions
|
||||
def test_int_add(self): self._test_kernel(lambda T: (T.empty(64).int() + T.empty(64).int()).float())
|
||||
def test_int_mul(self): self._test_kernel(lambda T: (T.empty(64).int() * T.empty(64).int()).float())
|
||||
def test_int_mod(self): self._test_kernel(lambda T: (T.empty(64).int().abs() % (T.empty(64).int().abs() + 1)).float())
|
||||
|
||||
# More math ops - use 32+ elements to force vector instructions
|
||||
def test_abs(self): self._test_kernel(lambda T: T.empty(64).abs())
|
||||
def test_floor(self): self._test_kernel(lambda T: T.empty(64).floor())
|
||||
def test_ceil(self): self._test_kernel(lambda T: T.empty(64).ceil())
|
||||
def test_trunc(self): self._test_kernel(lambda T: T.empty(64).trunc())
|
||||
|
||||
# Fused ops
|
||||
def test_fma(self): self._test_kernel(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
|
||||
|
||||
# Argmax/argmin (tests different reduction pattern) - use 32+ elements to force vector instructions
|
||||
def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
|
||||
def test_argmin(self): self._test_kernel(lambda T: T.empty(64).argmin())
|
||||
|
||||
# Exact value tests - use 32+ elements to force vector instructions (small tensors use scalar ops which Rust emu doesn't fully support)
|
||||
def test_abs_exact(self): self._test_kernel(lambda T: T([-1., 0., 1.]*11).abs()) # 33 elements
|
||||
def test_neg_exact(self): self._test_kernel(lambda T: -T([-1., 0., 1.]*11))
|
||||
def test_log_special(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).log())
|
||||
def test_exp_exact(self): self._test_kernel(lambda T: T([0., 1., -1.]*11).exp())
|
||||
def test_reciprocal_exact(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).reciprocal())
|
||||
|
||||
# Integer division and mod - use 32+ elements
|
||||
def test_int_div(self): self._test_kernel(lambda T: (T([10, 20, 30]*11).int() // T([3, 4, 5]*11).int()).float())
|
||||
def test_int_neg(self): self._test_kernel(lambda T: (-T([1, -2, 3]*11).int()).float())
|
||||
|
||||
# Mixed precision - use 32+ elements
|
||||
def test_half_add(self): self._test_kernel(lambda T: (T([1., 2.]*16).half() + T([3., 4.]*16).half()).float())
|
||||
def test_half_mul(self): self._test_kernel(lambda T: (T([2., 3.]*16).half() * T([4., 5.]*16).half()).float())
|
||||
|
||||
# Matrix ops - patterns from test_ops.py failures
|
||||
def test_cat(self): self._test_kernel(lambda T: T.empty(32, 64).cat(T.empty(32, 64), dim=1))
|
||||
def test_gather(self): self._test_kernel(lambda T: T.empty(64).gather(0, T.arange(32).int()))
|
||||
|
||||
# Tests from test_ops.py that are failing
|
||||
def test_permute(self): self._test_kernel(lambda T: T.empty(3, 4, 5, 6).permute((3, 2, 1, 0)).contiguous())
|
||||
def test_cat_large(self): self._test_kernel(lambda T: T.empty(45, 65, 9).cat(T.empty(45, 65, 9), T.empty(45, 65, 9), dim=1))
|
||||
def test_gather_small(self): self._test_kernel(lambda T: T.empty(10).gather(0, T.arange(5).int()))
|
||||
@unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
|
||||
def test_cross_entropy(self): self._test_kernel(lambda T: T.randn(32, 10).softmax().log().sum())
|
||||
def test_cross_entropy_class(self):
|
||||
import numpy as np
|
||||
np.random.seed(0)
|
||||
classes = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
|
||||
x_np = np.random.randn(32, 10).astype(np.float32)
|
||||
self._test_kernel(lambda T: (T(x_np.tolist()).reshape(32,10) + 0).cross_entropy((T(classes).int().reshape(32) + 0)))
|
||||
|
||||
# Regression tests for BFE operations with width=0 (walrus operator bug)
|
||||
# Regression tests
|
||||
def test_topk(self): self._test_kernel(lambda T: T.empty(64).topk(3)[0])
|
||||
def test_interpolate_uint8(self): self._test_kernel(lambda T: T.empty(2,3,64,64).relu().cast('uint8').interpolate((10,10), mode="linear"))
|
||||
|
||||
# Regression test for 64-bit comparison (V_CMP_GT_I64, V_CMP_LT_U64, etc.) with rsrc64
|
||||
def test_interpolate(self): self._test_kernel(lambda T: T.empty(1,2,16,16).relu().cast('uint8').interpolate((8,8), mode="linear"))
|
||||
def test_index_int64(self):
|
||||
from tinygrad import dtypes
|
||||
self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
|
||||
|
||||
@unittest.skip("only works with mock GPU")
|
||||
def test_index_int64_2d(self):
|
||||
def test_gelu(self): self._test_kernel(lambda T: T.empty(32, 32).gelu())
|
||||
def test_cross_entropy(self):
|
||||
import numpy as np
|
||||
np.random.seed(0)
|
||||
classes = np.random.randint(0, 10, (16,), dtype=np.int32).tolist()
|
||||
x_np = np.random.randn(16, 10).astype(np.float32)
|
||||
self._test_kernel(lambda T: (T(x_np.tolist()).reshape(16,10) + 0).cross_entropy((T(classes).int().reshape(16) + 0)))
|
||||
def test_isinf(self): self._test_kernel(lambda T: T([float('-inf'), 0., float('inf'), 1.1]*8).isinf())
|
||||
def test_sin_f64(self):
|
||||
from tinygrad import dtypes
|
||||
# Tests 64-bit compare with inline constants (comparing against 0)
|
||||
self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), T.arange(4).cast(dtypes.int64)])
|
||||
|
||||
# Pooling operations - regression test for VCC wave32 mode (S_CBRANCH_VCCZ should only check VCC_LO)
|
||||
def test_avg_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4), stride=2))
|
||||
@unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
|
||||
def test_avg_pool3d(self):
|
||||
import numpy as np
|
||||
np.random.seed(0)
|
||||
self._test_kernel(lambda T: T(np.random.randn(1, 1, 16, 16, 16).astype(np.float32).tolist()).avg_pool2d(kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False))
|
||||
def test_max_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4), stride=2))
|
||||
|
||||
# Convolution operations - multi-kernel tests
|
||||
def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv2d(T.empty(4, 4, 3, 3)), max_steps=100000)
|
||||
def test_conv_transpose2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv_transpose2d(T.empty(4, 4, 3, 3)), max_steps=200000)
|
||||
@unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
|
||||
def test_conv_transpose3d(self):
|
||||
import numpy as np
|
||||
np.random.seed(0)
|
||||
self._test_kernel(lambda T: T(np.random.randn(2, 4, 9, 9, 9).astype(np.float32).tolist()).conv_transpose2d(
|
||||
T(np.random.randn(4, 4, 3, 3, 3).astype(np.float32).tolist())), max_steps=500000)
|
||||
self._test_kernel(lambda T: T([2.0], dtype=dtypes.float64).sin())
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+80
-5
@@ -1,8 +1,9 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
|
||||
import unittest
|
||||
from extra.assembly.rdna3.autogen import *
|
||||
from extra.assembly.rdna3.lib import encode_src
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
from extra.assembly.amd.dsl import encode_src, RawImm
|
||||
from extra.assembly.amd.decode import detect_format
|
||||
|
||||
class TestMUBUF(unittest.TestCase):
|
||||
"""Test MUBUF (buffer) instructions."""
|
||||
@@ -308,7 +309,7 @@ class TestVOP3Literal(unittest.TestCase):
|
||||
def test_vop3_with_literal(self):
|
||||
# v_add3_u32 v5, vcc_hi, 0xaf123456, v255
|
||||
# GFX11: encoding: [0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf]
|
||||
from extra.assembly.rdna3.lib import RawImm
|
||||
from extra.assembly.amd.dsl import RawImm
|
||||
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=RawImm(107), src1=0xaf123456, src2=v[255])
|
||||
expected = bytes([0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf])
|
||||
self.assertEqual(inst.to_bytes(), expected)
|
||||
@@ -316,17 +317,91 @@ class TestVOP3Literal(unittest.TestCase):
|
||||
def test_vop3_literal_null_operand(self):
|
||||
# v_add3_u32 v5, null, exec_lo, 0xaf123456
|
||||
# GFX11: encoding: [0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf]
|
||||
from extra.assembly.rdna3.lib import RawImm
|
||||
from extra.assembly.amd.dsl import RawImm
|
||||
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=NULL, src1=RawImm(126), src2=0xaf123456)
|
||||
expected = bytes([0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf])
|
||||
self.assertEqual(inst.to_bytes(), expected)
|
||||
|
||||
def test_vop3p_with_literal(self):
|
||||
# Test VOP3P literal encoding (also uses Inst64)
|
||||
from extra.assembly.rdna3.lib import RawImm
|
||||
from extra.assembly.amd.dsl import RawImm
|
||||
inst = VOP3P(VOP3POp.V_PK_ADD_F16, vdst=v[5], src0=RawImm(240), src1=0x12345678, src2=v[0])
|
||||
self.assertEqual(len(inst.to_bytes()), 12) # 8 bytes + 4 byte literal
|
||||
|
||||
|
||||
class TestDetectFormat(unittest.TestCase):
|
||||
"""Test detect_format uses encoding from autogen classes."""
|
||||
|
||||
def test_detect_sopp(self):
|
||||
self.assertEqual(detect_format(s_endpgm().to_bytes()), SOPP)
|
||||
self.assertEqual(detect_format(s_nop(0).to_bytes()), SOPP)
|
||||
self.assertEqual(detect_format(s_barrier().to_bytes()), SOPP)
|
||||
|
||||
def test_detect_sop1(self):
|
||||
self.assertEqual(detect_format(s_mov_b32(s[0], 0).to_bytes()), SOP1)
|
||||
self.assertEqual(detect_format(s_mov_b64(s[0:1], 0).to_bytes()), SOP1)
|
||||
|
||||
def test_detect_sop2(self):
|
||||
self.assertEqual(detect_format(s_add_u32(s[0], s[1], s[2]).to_bytes()), SOP2)
|
||||
self.assertEqual(detect_format(s_mul_i32(s[0], s[1], s[2]).to_bytes()), SOP2)
|
||||
|
||||
def test_detect_sopc(self):
|
||||
self.assertEqual(detect_format(s_cmp_eq_i32(s[0], s[1]).to_bytes()), SOPC)
|
||||
|
||||
def test_detect_sopk(self):
|
||||
self.assertEqual(detect_format(s_movk_i32(s[0], 0x1234).to_bytes()), SOPK)
|
||||
|
||||
def test_detect_vop1(self):
|
||||
self.assertEqual(detect_format(v_mov_b32_e32(v[0], 0).to_bytes()), VOP1)
|
||||
self.assertEqual(detect_format(v_rcp_f32_e32(v[0], v[1]).to_bytes()), VOP1)
|
||||
|
||||
def test_detect_vop2(self):
|
||||
self.assertEqual(detect_format(v_add_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
|
||||
self.assertEqual(detect_format(v_mul_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
|
||||
|
||||
def test_detect_vopc(self):
|
||||
self.assertEqual(detect_format(v_cmp_eq_f32_e32(v[0], v[1]).to_bytes()), VOPC)
|
||||
self.assertEqual(detect_format(v_cmp_lt_i32_e32(v[0], v[1]).to_bytes()), VOPC)
|
||||
|
||||
def test_detect_vop3(self):
|
||||
self.assertEqual(detect_format(v_add_f32_e64(v[0], v[1], v[2]).to_bytes()), VOP3)
|
||||
self.assertEqual(detect_format(v_fma_f32(v[0], v[1], v[2], v[3]).to_bytes()), VOP3)
|
||||
|
||||
def test_detect_vop3p(self):
|
||||
self.assertEqual(detect_format(VOP3P(VOP3POp.V_PK_ADD_F16, v[0], v[1], v[2], v[3]).to_bytes()), VOP3P)
|
||||
|
||||
def test_detect_smem(self):
|
||||
self.assertEqual(detect_format(s_load_b32(s[0], s[2:3], 0).to_bytes()), SMEM)
|
||||
self.assertEqual(detect_format(s_load_b64(s[0:1], s[2:3], s[5]).to_bytes()), SMEM)
|
||||
|
||||
def test_detect_ds(self):
|
||||
self.assertEqual(detect_format(ds_load_b32(v[0], v[1]).to_bytes()), DS)
|
||||
self.assertEqual(detect_format(ds_store_b32(v[0], v[1]).to_bytes()), DS)
|
||||
|
||||
def test_detect_flat(self):
|
||||
self.assertEqual(detect_format(global_load_b32(v[0], v[1:3], RawImm(124)).to_bytes()), FLAT)
|
||||
self.assertEqual(detect_format(global_store_b32(v[0:2], v[2], RawImm(124)).to_bytes()), FLAT)
|
||||
|
||||
def test_detect_mubuf(self):
|
||||
self.assertEqual(detect_format(buffer_load_b32(v[0], v[1], s[0:4], s[5]).to_bytes()), MUBUF)
|
||||
|
||||
def test_detect_mtbuf(self):
|
||||
self.assertEqual(detect_format(tbuffer_load_format_x(v[0], v[1], s[0:4], s[5], format=22).to_bytes()), MTBUF)
|
||||
|
||||
def test_detect_mimg(self):
|
||||
self.assertEqual(detect_format(image_load(v[0:4], v[4:6], s[0:8], dmask=0xf, dim=1).to_bytes()), MIMG)
|
||||
|
||||
def test_detect_exp(self):
|
||||
self.assertEqual(detect_format(EXP(en=0xf, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[2], vsrc3=v[3]).to_bytes()), EXP)
|
||||
|
||||
def test_detect_vopd(self):
|
||||
inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=0, srcy0=0)
|
||||
self.assertEqual(detect_format(inst.to_bytes()), VOPD)
|
||||
|
||||
def test_detect_vinterp(self):
|
||||
inst = VINTERP(VINTERPOp.V_INTERP_P10_F32, vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
|
||||
self.assertEqual(detect_format(inst.to_bytes()), VINTERP)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+7
-4
@@ -2,10 +2,10 @@
|
||||
# the Inst constructor should be looking at the types of the fields to correctly set the value
|
||||
|
||||
import unittest, struct
|
||||
from extra.assembly.rdna3.autogen import *
|
||||
from extra.assembly.rdna3.lib import Inst
|
||||
from extra.assembly.rdna3.asm import asm
|
||||
from extra.assembly.rdna3.test.test_roundtrip import compile_asm
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
from extra.assembly.amd.dsl import Inst
|
||||
from extra.assembly.amd.asm import asm
|
||||
from extra.assembly.amd.test.test_roundtrip import compile_asm
|
||||
|
||||
class TestIntegration(unittest.TestCase):
|
||||
inst: Inst
|
||||
@@ -21,6 +21,9 @@ class TestIntegration(unittest.TestCase):
|
||||
self.assertEqual(repr(self.inst), repr(reasm))
|
||||
print(desc)
|
||||
|
||||
def test_wmma(self):
|
||||
self.inst = v_wmma_f32_16x16x16_f16(v[0:7], v[189:192], v[140:143], v[0:7])
|
||||
|
||||
def test_load_b128(self):
|
||||
self.inst = s_load_b128(s[4:7], s[0:1], NULL, 0)
|
||||
|
||||
+5
-18
@@ -1,17 +1,9 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
|
||||
import unittest, re, io, sys
|
||||
from extra.assembly.rdna3.autogen import *
|
||||
from extra.assembly.rdna3.asm import waitcnt, asm
|
||||
|
||||
def get_amd_toolchain():
|
||||
"""Check if AMD toolchain is available."""
|
||||
try:
|
||||
from tinygrad.runtime.support.compiler_amd import HIPCompiler
|
||||
HIPCompiler("gfx1100").compile(".text\ns_endpgm")
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
import unittest, re, io, sys, subprocess
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
from extra.assembly.amd.asm import waitcnt, asm
|
||||
from extra.assembly.amd.test.helpers import get_llvm_mc
|
||||
|
||||
def disassemble(lib: bytes, arch: str = "gfx1100") -> str:
|
||||
"""Disassemble ELF binary using tinygrad's compiler, return raw output."""
|
||||
@@ -47,7 +39,6 @@ def assemble_and_disassemble(instructions: list, arch: str = "gfx1100") -> list[
|
||||
lib = HIPCompiler(arch).compile(asm_src)
|
||||
return parse_disassembly(disassemble(lib, arch))
|
||||
|
||||
@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
|
||||
class TestIntegration(unittest.TestCase):
|
||||
"""Test our assembler output matches LLVM disassembly."""
|
||||
|
||||
@@ -157,7 +148,6 @@ class TestIntegration(unittest.TestCase):
|
||||
return
|
||||
self.fail("Could not find s_mov_b32 in disassembly")
|
||||
|
||||
@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
|
||||
class TestAsm(unittest.TestCase):
|
||||
"""Test asm() string parsing."""
|
||||
|
||||
@@ -212,10 +202,8 @@ class TestAsm(unittest.TestCase):
|
||||
|
||||
def test_asm_vop3_modifiers(self):
|
||||
"""Test asm() with VOP3 modifiers (neg, abs, clamp)."""
|
||||
import subprocess, re
|
||||
|
||||
def get_llvm_encoding(instr: str) -> str:
|
||||
result = subprocess.run(['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-show-encoding'],
|
||||
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', '-mcpu=gfx1100', '-show-encoding'],
|
||||
input=instr, capture_output=True, text=True)
|
||||
if m := re.search(r'encoding:\s*\[(.*?)\]', result.stdout):
|
||||
return m.group(1).replace('0x','').replace(',','').replace(' ','')
|
||||
@@ -233,7 +221,6 @@ class TestAsm(unittest.TestCase):
|
||||
llvm_hex = get_llvm_encoding(t)
|
||||
self.assertEqual(our_hex, llvm_hex, f"mismatch for: {t}")
|
||||
|
||||
@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
|
||||
class TestTinygradIntegration(unittest.TestCase):
|
||||
"""Test that we can parse disassembled tinygrad kernels."""
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test AMD assembler/disassembler against LLVM test vectors."""
|
||||
import unittest, re, subprocess, functools
|
||||
from tinygrad.helpers import fetch
|
||||
from extra.assembly.amd.asm import asm, disasm
|
||||
from extra.assembly.amd.decode import decode_inst, detect_format
|
||||
from extra.assembly.amd.test.helpers import get_llvm_mc
|
||||
|
||||
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-21.1.0/llvm/test/MC/AMDGPU"
|
||||
|
||||
RDNA_FILES = ['gfx11_asm_sop1.s', 'gfx11_asm_sop2.s', 'gfx11_asm_sopp.s', 'gfx11_asm_sopk.s', 'gfx11_asm_sopc.s',
|
||||
'gfx11_asm_vop1.s', 'gfx11_asm_vop2.s', 'gfx11_asm_vopc.s', 'gfx11_asm_vop3.s', 'gfx11_asm_vop3p.s', 'gfx11_asm_vinterp.s',
|
||||
'gfx11_asm_vopd.s', 'gfx11_asm_vopcx.s', 'gfx11_asm_vop3_from_vop1.s', 'gfx11_asm_vop3_from_vop2.s', 'gfx11_asm_vop3_from_vopc.s',
|
||||
'gfx11_asm_vop3_from_vopcx.s', 'gfx11_asm_ds.s', 'gfx11_asm_smem.s', 'gfx11_asm_flat.s', 'gfx11_asm_mubuf.s', 'gfx11_asm_mtbuf.s',
|
||||
'gfx11_asm_mimg.s', 'gfx11_asm_wmma.s', 'gfx11_asm_vop3_features.s', 'gfx11_asm_vop3p_features.s', 'gfx11_asm_vopd_features.s',
|
||||
'gfx11_asm_vop3_alias.s', 'gfx11_asm_vop3p_alias.s', 'gfx11_asm_vopc_alias.s', 'gfx11_asm_vopcx_alias.s', 'gfx11_asm_vinterp_alias.s',
|
||||
'gfx11_asm_smem_alias.s', 'gfx11_asm_mubuf_alias.s', 'gfx11_asm_mtbuf_alias.s']
|
||||
# CDNA test files - includes gfx9 files for shared instructions, plus gfx90a/gfx942 specific files
|
||||
# gfx90a_ldst_acc.s has MIMG mixed in, filtered via is_mimg check
|
||||
CDNA_FILES = ['gfx9_asm_sop1.s', 'gfx9_asm_sop2.s', 'gfx9_asm_sopp.s', 'gfx9_asm_sopk.s', 'gfx9_asm_sopc.s',
|
||||
'gfx9_asm_vop1.s', 'gfx9_asm_vop2.s', 'gfx9_asm_vopc.s', 'gfx9_asm_vop3.s', 'gfx9_asm_vop3p.s',
|
||||
'gfx9_asm_ds.s', 'gfx9_asm_flat.s', 'gfx9_asm_smem.s', 'gfx9_asm_mubuf.s', 'gfx9_asm_mtbuf.s',
|
||||
'gfx90a_ldst_acc.s', 'gfx90a_asm_features.s', 'flat-scratch-gfx942.s', 'gfx942_asm_features.s',
|
||||
'mai-gfx90a.s', 'mai-gfx942.s']
|
||||
# RDNA4 (gfx12) test files - excludes alias/err/fake16/dpp files, and vimage/vsample (not supported)
|
||||
# NOTE: vflat/vdsdir excluded - not implemented; features.s has mixed formats
|
||||
RDNA4_FILES = ['gfx12_asm_sop1.s', 'gfx12_asm_sop2.s', 'gfx12_asm_sopp.s', 'gfx12_asm_sopk.s', 'gfx12_asm_sopc.s',
|
||||
'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_vbuffer_mubuf.s', 'gfx12_asm_vbuffer_mtbuf.s', 'gfx12_asm_wmma_w32.s', 'gfx12_asm_exp.s']
|
||||
|
||||
def _is_mimg(data: bytes) -> bool: return (int.from_bytes(data[:4], 'little') >> 26) & 0x3f == 0b111100
|
||||
|
||||
def _parse_llvm_tests(text: str, pattern: str) -> list[tuple[str, bytes]]:
|
||||
tests = []
|
||||
for block in text.split('\n\n'):
|
||||
asm_text, encoding = None, None
|
||||
for line in block.split('\n'):
|
||||
line = line.strip()
|
||||
if not line or line.startswith(('.', ';')): continue
|
||||
if not line.startswith('//'):
|
||||
asm_text = line.split('//')[0].strip() or asm_text
|
||||
if m := re.search(pattern + r'[^:]*:.*?(?:encoding:\s*)?\[(0x[0-9a-f,x\s]+)\]', line, re.I):
|
||||
encoding = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
|
||||
if asm_text and encoding:
|
||||
try: tests.append((asm_text, bytes.fromhex(encoding)))
|
||||
except ValueError: pass
|
||||
return tests
|
||||
|
||||
@functools.cache
|
||||
def _get_tests(f: str, arch: str) -> list[tuple[str, bytes]]:
|
||||
text = fetch(f"{LLVM_BASE}/{f}").read_bytes().decode('utf-8', errors='ignore')
|
||||
if arch == "rdna3":
|
||||
tests = _parse_llvm_tests(text, r'(?:GFX11|W32|W64)')
|
||||
elif arch == "rdna4":
|
||||
# Match GFX12 but not GFX1250 (which has different lit64 encoding)
|
||||
tests = _parse_llvm_tests(text, r'(?:GFX12(?!50)|W32|W64)')
|
||||
elif 'gfx90a' in f or 'gfx942' in f:
|
||||
tests = _parse_llvm_tests(text, r'(?:GFX90A|GFX942)')
|
||||
else:
|
||||
tests = _parse_llvm_tests(text, r'(?:VI9|GFX9|CHECK)')
|
||||
return [(a, d) for a, d in tests if not _is_mimg(d)] if arch == "cdna" else tests
|
||||
|
||||
def _compile_asm_batch(instrs: list[str], arch: str = "rdna3") -> list[bytes]:
|
||||
if not instrs: return []
|
||||
mcpu = {'rdna3': 'gfx1100', 'rdna4': 'gfx1200'}.get(arch, 'gfx1100')
|
||||
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', '-mattr=+real-true16,+wavefrontsize32', '-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 _make_test(f: str, arch: str, test_type: str):
|
||||
def test(self):
|
||||
tests = _get_tests(f, arch)
|
||||
name = f"{arch}_{test_type}_{f}"
|
||||
if test_type == "roundtrip":
|
||||
for _, data in tests:
|
||||
decoded = detect_format(data, arch).from_bytes(data)
|
||||
self.assertEqual(decoded.to_bytes()[:len(data)], data)
|
||||
print(f"{name}: {len(tests)} passed")
|
||||
elif test_type == "asm":
|
||||
passed, skipped = 0, 0
|
||||
for asm_text, expected in tests:
|
||||
try:
|
||||
self.assertEqual(asm(asm_text, arch).to_bytes(), expected)
|
||||
passed += 1
|
||||
except: skipped += 1
|
||||
print(f"{name}: {passed} passed, {skipped} skipped")
|
||||
elif test_type == "disasm":
|
||||
to_test = []
|
||||
for _, data in tests:
|
||||
try:
|
||||
decoded = decode_inst(data, arch)
|
||||
# Skip if roundtrip fails, disasm fails, or op_name is missing (disasm starts with space)
|
||||
if decoded.to_bytes()[:len(data)] == data and (d := disasm(decoded)) and not d.startswith(' '): to_test.append((data, d))
|
||||
except: pass
|
||||
skipped = len(tests) - len(to_test)
|
||||
print(f"{name}: {len(to_test)} passed, {skipped} skipped")
|
||||
if arch in ("rdna3", "rdna4"):
|
||||
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
|
||||
for (data, _), llvm in zip(to_test, _compile_asm_batch([t[1] for t in to_test], arch)): self.assertEqual(llvm, data)
|
||||
return test
|
||||
|
||||
class TestLLVM(unittest.TestCase): pass
|
||||
|
||||
for f in RDNA_FILES:
|
||||
setattr(TestLLVM, f"test_rdna3_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "roundtrip"))
|
||||
setattr(TestLLVM, f"test_rdna3_asm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "asm"))
|
||||
setattr(TestLLVM, f"test_rdna3_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "disasm"))
|
||||
for f in CDNA_FILES:
|
||||
setattr(TestLLVM, f"test_cdna_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "roundtrip"))
|
||||
setattr(TestLLVM, f"test_cdna_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "disasm"))
|
||||
for f in RDNA4_FILES:
|
||||
setattr(TestLLVM, f"test_rdna4_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "roundtrip"))
|
||||
setattr(TestLLVM, f"test_rdna4_asm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "asm"))
|
||||
setattr(TestLLVM, f"test_rdna4_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "disasm"))
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+5
-4
@@ -20,11 +20,12 @@ runner = get_runner(dev.device, si.ast)
|
||||
prg = runner._prg
|
||||
lib = bytearray(prg.lib)
|
||||
|
||||
# Find s_endpgm (0xBFB00000) and replace with invalid SOPP op=127 (0xBFFF0000)
|
||||
# Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
|
||||
# VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400
|
||||
found = False
|
||||
for i in range(0, len(lib) - 4, 4):
|
||||
if struct.unpack("<I", lib[i:i+4])[0] == 0xBFB00000:
|
||||
lib[i:i+4] = struct.pack("<I", 0xBFFF0000)
|
||||
lib[i:i+4] = struct.pack("<I", 0x7E008400)
|
||||
found = True
|
||||
break
|
||||
assert found, "s_endpgm not found"
|
||||
@@ -46,9 +47,9 @@ dev.synchronize()
|
||||
elapsed = time.perf_counter() - st
|
||||
|
||||
self.assertNotEqual(result.returncode, 0, "should have raised")
|
||||
self.assertIn("NotImplementedError", result.stderr)
|
||||
self.assertTrue("Error" in result.stderr, f"expected an error in stderr, got: {result.stderr[:500]}")
|
||||
# Should exit immediately, not wait for the full timeout
|
||||
self.assertLess(elapsed, 5.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
|
||||
self.assertLess(elapsed, 9.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,403 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Tests for the RDNA3 pseudocode DSL."""
|
||||
import unittest
|
||||
from extra.assembly.amd.pcode import (Reg, TypedView, TypedView, MASK32, MASK64,
|
||||
_f32, _i32, _f16, _i16, f32_to_f16, isNAN, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
|
||||
BYTE_PERMUTE, v_sad_u8, v_msad_u8, _compile_pseudocode, _expr, compile_pseudocode)
|
||||
from extra.assembly.amd.test.helpers import ExecContext
|
||||
from extra.assembly.amd.autogen.rdna3.str_pcode import VOP3SDOp_PCODE, VOPCOp_PCODE
|
||||
from extra.assembly.amd.autogen.rdna3.enum import VOP3SDOp, VOPCOp
|
||||
|
||||
# Compile pseudocode functions on demand for regression tests
|
||||
_VOP3SDOp_V_DIV_SCALE_F32 = compile_pseudocode('VOP3SDOp', 'V_DIV_SCALE_F32', VOP3SDOp_PCODE[VOP3SDOp.V_DIV_SCALE_F32])
|
||||
_VOPCOp_V_CMP_CLASS_F32 = compile_pseudocode('VOPCOp', 'V_CMP_CLASS_F32', VOPCOp_PCODE[VOPCOp.V_CMP_CLASS_F32])
|
||||
|
||||
class TestReg(unittest.TestCase):
|
||||
def test_u32_read(self):
|
||||
r = Reg(0xDEADBEEF)
|
||||
self.assertEqual(int(r.u32), 0xDEADBEEF)
|
||||
|
||||
def test_u32_write(self):
|
||||
r = Reg(0)
|
||||
r.u32 = 0x12345678
|
||||
self.assertEqual(r._val, 0x12345678)
|
||||
|
||||
def test_f32_read(self):
|
||||
r = Reg(0x40400000) # 3.0f
|
||||
self.assertAlmostEqual(float(r.f32), 3.0)
|
||||
|
||||
def test_f32_write(self):
|
||||
r = Reg(0)
|
||||
r.f32 = 3.0
|
||||
self.assertEqual(r._val, 0x40400000)
|
||||
|
||||
def test_i32_signed(self):
|
||||
r = Reg(0xFFFFFFFF) # -1 as signed
|
||||
self.assertEqual(int(r.i32), -1)
|
||||
|
||||
def test_u64(self):
|
||||
r = Reg(0xDEADBEEFCAFEBABE)
|
||||
self.assertEqual(int(r.u64), 0xDEADBEEFCAFEBABE)
|
||||
|
||||
def test_f64(self):
|
||||
r = Reg(0x4008000000000000) # 3.0 as f64
|
||||
self.assertAlmostEqual(float(r.f64), 3.0)
|
||||
|
||||
class TestTypedView(unittest.TestCase):
|
||||
def test_bit_slice(self):
|
||||
r = Reg(0xDEADBEEF)
|
||||
# Slices return TypedView which supports .u32, .u16 etc (matching pseudocode like S1.u32[1:0].u32)
|
||||
self.assertEqual(r.u32[7:0].u32, 0xEF)
|
||||
self.assertEqual(r.u32[15:8].u32, 0xBE)
|
||||
self.assertEqual(r.u32[23:16].u32, 0xAD)
|
||||
self.assertEqual(r.u32[31:24].u32, 0xDE)
|
||||
# Also works with int() for arithmetic
|
||||
self.assertEqual(int(r.u32[7:0]), 0xEF)
|
||||
|
||||
def test_single_bit_read(self):
|
||||
r = Reg(0b11010101)
|
||||
self.assertEqual(r.u32[0], 1)
|
||||
self.assertEqual(r.u32[1], 0)
|
||||
self.assertEqual(r.u32[2], 1)
|
||||
self.assertEqual(r.u32[3], 0)
|
||||
|
||||
def test_single_bit_write(self):
|
||||
r = Reg(0)
|
||||
r.u32[5] = 1
|
||||
r.u32[3] = 1
|
||||
self.assertEqual(r._val, 0b00101000)
|
||||
|
||||
def test_nested_bit_access(self):
|
||||
# S0.u32[S1.u32[4:0]] - access bit at position from another register
|
||||
s0 = Reg(0b11010101)
|
||||
s1 = Reg(3)
|
||||
bit_pos = s1.u32[4:0] # TypedView, int value = 3
|
||||
bit_val = s0.u32[int(bit_pos)] # bit 3 of s0 = 0
|
||||
self.assertEqual(int(bit_pos), 3)
|
||||
self.assertEqual(bit_val, 0)
|
||||
|
||||
def test_arithmetic(self):
|
||||
r1 = Reg(0x40400000) # 3.0f
|
||||
r2 = Reg(0x40800000) # 4.0f
|
||||
result = r1.f32 + r2.f32
|
||||
self.assertAlmostEqual(result, 7.0)
|
||||
|
||||
def test_comparison(self):
|
||||
r1 = Reg(5)
|
||||
r2 = Reg(3)
|
||||
self.assertTrue(r1.u32 > r2.u32)
|
||||
self.assertFalse(r1.u32 < r2.u32)
|
||||
self.assertTrue(r1.u32 != r2.u32)
|
||||
|
||||
class TestTypedView(unittest.TestCase):
|
||||
def test_slice_read(self):
|
||||
r = Reg(0x56781234)
|
||||
self.assertEqual(r[15:0].u16, 0x1234)
|
||||
self.assertEqual(r[31:16].u16, 0x5678)
|
||||
|
||||
def test_slice_write(self):
|
||||
r = Reg(0)
|
||||
r[15:0].u16 = 0x1234
|
||||
r[31:16].u16 = 0x5678
|
||||
self.assertEqual(r._val, 0x56781234)
|
||||
|
||||
def test_slice_f16(self):
|
||||
r = Reg(0)
|
||||
r[15:0].f16 = 3.0
|
||||
self.assertAlmostEqual(_f16(r._val & 0xffff), 3.0, places=2)
|
||||
|
||||
class TestCompiler(unittest.TestCase):
|
||||
def test_ternary(self):
|
||||
result = _expr("a > b ? 1 : 0")
|
||||
self.assertIn("if", result)
|
||||
self.assertIn("else", result)
|
||||
|
||||
def test_type_prefix_strip(self):
|
||||
self.assertEqual(_expr("1'0U"), "0")
|
||||
self.assertEqual(_expr("32'1"), "1")
|
||||
self.assertEqual(_expr("16'0xFFFF"), "0xFFFF")
|
||||
|
||||
def test_suffix_strip(self):
|
||||
self.assertEqual(_expr("0ULL"), "0")
|
||||
self.assertEqual(_expr("1LL"), "1")
|
||||
self.assertEqual(_expr("5U"), "5")
|
||||
self.assertEqual(_expr("3.14F"), "3.14")
|
||||
|
||||
def test_boolean_ops(self):
|
||||
self.assertIn("and", _expr("a && b"))
|
||||
self.assertIn("or", _expr("a || b"))
|
||||
self.assertIn("!=", _expr("a <> b"))
|
||||
|
||||
def test_pack16(self):
|
||||
result = _expr("{ a, b }")
|
||||
self.assertIn("_pack", result)
|
||||
|
||||
def test_type_cast_strip(self):
|
||||
self.assertEqual(_expr("64'U(x)"), "(x)")
|
||||
self.assertEqual(_expr("32'I(y)"), "(y)")
|
||||
|
||||
class TestExecContext(unittest.TestCase):
|
||||
def test_float_add(self):
|
||||
ctx = ExecContext(s0=0x40400000, s1=0x40800000) # 3.0f, 4.0f
|
||||
ctx.D0.f32 = ctx.S0.f32 + ctx.S1.f32
|
||||
self.assertAlmostEqual(_f32(ctx.D0._val), 7.0)
|
||||
|
||||
def test_float_mul(self):
|
||||
ctx = ExecContext(s0=0x40400000, s1=0x40800000) # 3.0f, 4.0f
|
||||
ctx.run("D0.f32 = S0.f32 * S1.f32")
|
||||
self.assertAlmostEqual(_f32(ctx.D0._val), 12.0)
|
||||
|
||||
def test_scc_comparison(self):
|
||||
ctx = ExecContext(s0=42, s1=42)
|
||||
ctx.run("SCC = S0.u32 == S1.u32")
|
||||
self.assertEqual(ctx.SCC._val, 1)
|
||||
|
||||
def test_scc_comparison_false(self):
|
||||
ctx = ExecContext(s0=42, s1=43)
|
||||
ctx.run("SCC = S0.u32 == S1.u32")
|
||||
self.assertEqual(ctx.SCC._val, 0)
|
||||
|
||||
def test_ternary(self):
|
||||
code = _compile_pseudocode("D0.u32 = S0.u32 > S1.u32 ? 1'1U : 1'0U")
|
||||
ctx = ExecContext(s0=5, s1=3)
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val, 1)
|
||||
|
||||
def test_pack(self):
|
||||
code = _compile_pseudocode("D0 = { S1[15:0].u16, S0[15:0].u16 }")
|
||||
ctx = ExecContext(s0=0x1234, s1=0x5678)
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val, 0x56781234)
|
||||
|
||||
def test_tmp_with_typed_access(self):
|
||||
code = _compile_pseudocode("""tmp = S0.u32 + S1.u32
|
||||
D0.u32 = tmp.u32""")
|
||||
ctx = ExecContext(s0=100, s1=200)
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val, 300)
|
||||
|
||||
def test_s_add_u32_pattern(self):
|
||||
# Real pseudocode pattern from S_ADD_U32
|
||||
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
|
||||
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
|
||||
D0.u32 = tmp.u32""")
|
||||
# Test overflow case
|
||||
ctx = ExecContext(s0=0xFFFFFFFF, s1=0x00000001)
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val, 0) # Wraps to 0
|
||||
self.assertEqual(ctx.SCC._val, 1) # Carry set
|
||||
|
||||
def test_s_add_u32_no_overflow(self):
|
||||
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
|
||||
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
|
||||
D0.u32 = tmp.u32""")
|
||||
ctx = ExecContext(s0=100, s1=200)
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val, 300)
|
||||
self.assertEqual(ctx.SCC._val, 0) # No carry
|
||||
|
||||
def test_vcc_lane_read(self):
|
||||
ctx = ExecContext(vcc=0b1010, lane=1)
|
||||
# Lane 1 is set
|
||||
self.assertEqual(ctx.VCC.u64[1], 1)
|
||||
self.assertEqual(ctx.VCC.u64[2], 0)
|
||||
|
||||
def test_vcc_lane_write(self):
|
||||
ctx = ExecContext(vcc=0, lane=0)
|
||||
ctx.VCC.u64[3] = 1
|
||||
ctx.VCC.u64[1] = 1
|
||||
self.assertEqual(ctx.VCC._val, 0b1010)
|
||||
|
||||
def test_for_loop(self):
|
||||
# CTZ pattern - find first set bit
|
||||
code = _compile_pseudocode("""tmp = -1
|
||||
for i in 0 : 31 do
|
||||
if S0.u32[i] == 1 then
|
||||
tmp = i
|
||||
endif
|
||||
endfor
|
||||
D0.i32 = tmp""")
|
||||
ctx = ExecContext(s0=0b1000) # Bit 3 is set
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val & MASK32, 3)
|
||||
|
||||
def test_result_dict(self):
|
||||
ctx = ExecContext(s0=5, s1=3)
|
||||
ctx.D0.u32 = 42
|
||||
ctx.SCC._val = 1
|
||||
result = ctx.result()
|
||||
self.assertEqual(result['d0'], 42)
|
||||
self.assertEqual(result['scc'], 1)
|
||||
|
||||
class TestPseudocodeRegressions(unittest.TestCase):
|
||||
"""Regression tests for pseudocode instruction emulation bugs."""
|
||||
|
||||
def test_v_div_scale_f32_vcc_always_returned(self):
|
||||
"""V_DIV_SCALE_F32 must always return VCC, even when VCC=0 (no scaling needed).
|
||||
Bug: when VCC._val == vcc (both 0), VCC wasn't returned, so VCC bits weren't written.
|
||||
This caused division to produce wrong results for multiple lanes."""
|
||||
# Normal case: 1.0 / 3.0, no scaling needed, VCC should be 0
|
||||
s0 = 0x3f800000 # 1.0
|
||||
s1 = 0x40400000 # 3.0
|
||||
s2 = 0x3f800000 # 1.0 (numerator)
|
||||
result = _VOP3SDOp_V_DIV_SCALE_F32(s0, s1, s2, 0, 0, 0, 0, 0xffffffff, 0, None)
|
||||
# Must always have VCC in result
|
||||
self.assertIn('VCC', result, "V_DIV_SCALE_F32 must always return VCC")
|
||||
self.assertEqual(result['VCC'] & 1, 0, "VCC lane 0 should be 0 when no scaling needed")
|
||||
|
||||
def test_v_cmp_class_f32_detects_quiet_nan(self):
|
||||
"""V_CMP_CLASS_F32 must correctly identify quiet NaN vs signaling NaN.
|
||||
Bug: isQuietNAN and isSignalNAN both used math.isnan which can't distinguish them."""
|
||||
quiet_nan = 0x7fc00000 # quiet NaN: exponent=255, bit22=1
|
||||
signal_nan = 0x7f800001 # signaling NaN: exponent=255, bit22=0
|
||||
# Test quiet NaN detection (bit 1 in mask)
|
||||
s1_quiet = 0b0000000010 # bit 1 = quiet NaN
|
||||
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
|
||||
self.assertEqual(result['D0'] & 1, 1, "Should detect quiet NaN with quiet NaN mask")
|
||||
# Test signaling NaN detection (bit 0 in mask)
|
||||
s1_signal = 0b0000000001 # bit 0 = signaling NaN
|
||||
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
|
||||
self.assertEqual(result['D0'] & 1, 1, "Should detect signaling NaN with signaling NaN mask")
|
||||
# Test that quiet NaN doesn't match signaling NaN mask
|
||||
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
|
||||
self.assertEqual(result['D0'] & 1, 0, "Quiet NaN should not match signaling NaN mask")
|
||||
# Test that signaling NaN doesn't match quiet NaN mask
|
||||
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
|
||||
self.assertEqual(result['D0'] & 1, 0, "Signaling NaN should not match quiet NaN mask")
|
||||
|
||||
def testisNAN_with_typed_view(self):
|
||||
"""isNAN must work with TypedView objects, not just Python floats.
|
||||
Bug: isNAN checked isinstance(x, float) which returned False for TypedView."""
|
||||
nan_reg = Reg(0x7fc00000) # quiet NaN
|
||||
normal_reg = Reg(0x3f800000) # 1.0
|
||||
inf_reg = Reg(0x7f800000) # +inf
|
||||
self.assertTrue(isNAN(nan_reg.f32), "isNAN should return True for NaN TypedView")
|
||||
self.assertFalse(isNAN(normal_reg.f32), "isNAN should return False for normal TypedView")
|
||||
self.assertFalse(isNAN(inf_reg.f32), "isNAN should return False for inf TypedView")
|
||||
|
||||
class TestBF16(unittest.TestCase):
|
||||
"""Tests for BF16 (bfloat16) support."""
|
||||
|
||||
def test_bf16_conversion(self):
|
||||
"""Test bf16 <-> f32 conversion."""
|
||||
# bf16 is just the top 16 bits of f32
|
||||
# 1.0f = 0x3f800000, bf16 = 0x3f80
|
||||
self.assertAlmostEqual(_bf16(0x3f80), 1.0, places=2)
|
||||
self.assertEqual(_ibf16(1.0), 0x3f80)
|
||||
# 2.0f = 0x40000000, bf16 = 0x4000
|
||||
self.assertAlmostEqual(_bf16(0x4000), 2.0, places=2)
|
||||
self.assertEqual(_ibf16(2.0), 0x4000)
|
||||
# -1.0f = 0xbf800000, bf16 = 0xbf80
|
||||
self.assertAlmostEqual(_bf16(0xbf80), -1.0, places=2)
|
||||
self.assertEqual(_ibf16(-1.0), 0xbf80)
|
||||
|
||||
def test_bf16_special_values(self):
|
||||
"""Test bf16 special values (inf, nan)."""
|
||||
import math
|
||||
# +inf: f32 = 0x7f800000, bf16 = 0x7f80
|
||||
self.assertTrue(math.isinf(_bf16(0x7f80)))
|
||||
self.assertEqual(_ibf16(float('inf')), 0x7f80)
|
||||
# -inf: f32 = 0xff800000, bf16 = 0xff80
|
||||
self.assertTrue(math.isinf(_bf16(0xff80)))
|
||||
self.assertEqual(_ibf16(float('-inf')), 0xff80)
|
||||
# NaN: quiet NaN bf16 = 0x7fc0
|
||||
self.assertTrue(math.isnan(_bf16(0x7fc0)))
|
||||
self.assertEqual(_ibf16(float('nan')), 0x7fc0)
|
||||
|
||||
def test_bf16_register_property(self):
|
||||
"""Test Reg.bf16 property."""
|
||||
r = Reg(0)
|
||||
r.bf16 = 3.0 # 3.0f = 0x40400000, bf16 = 0x4040
|
||||
self.assertEqual(r._val & 0xffff, 0x4040)
|
||||
self.assertAlmostEqual(float(r.bf16), 3.0, places=1)
|
||||
|
||||
def test_bf16_slice_property(self):
|
||||
"""Test TypedView.bf16 property."""
|
||||
r = Reg(0x40404040) # Two bf16 3.0 values
|
||||
self.assertAlmostEqual(r[15:0].bf16, 3.0, places=1)
|
||||
self.assertAlmostEqual(r[31:16].bf16, 3.0, places=1)
|
||||
|
||||
class TestBytePermute(unittest.TestCase):
|
||||
"""Tests for BYTE_PERMUTE helper function (V_PERM_B32)."""
|
||||
|
||||
def test_byte_select_0_to_7(self):
|
||||
"""Test selecting bytes 0-7 from 64-bit data."""
|
||||
# data = {s0, s1} where s0 is bytes 0-3, s1 is bytes 4-7
|
||||
# Combined: 0x0706050403020100 (byte 0 = 0x00, byte 7 = 0x07)
|
||||
data = 0x0706050403020100
|
||||
for i in range(8):
|
||||
self.assertEqual(BYTE_PERMUTE(data, i), i, f"byte {i} should be {i}")
|
||||
|
||||
def test_sign_extend_bytes(self):
|
||||
"""Test sign extension selectors 8-11."""
|
||||
# sel 8: sign of byte 1 (bits 15:8)
|
||||
# sel 9: sign of byte 3 (bits 31:24)
|
||||
# sel 10: sign of byte 5 (bits 47:40)
|
||||
# sel 11: sign of byte 7 (bits 63:56)
|
||||
data = 0x8000800080008000 # All relevant bytes have sign bit set
|
||||
self.assertEqual(BYTE_PERMUTE(data, 8), 0xff)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 9), 0xff)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 10), 0xff)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 11), 0xff)
|
||||
data = 0x7f007f007f007f00 # No sign bits set
|
||||
self.assertEqual(BYTE_PERMUTE(data, 8), 0x00)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 9), 0x00)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 10), 0x00)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 11), 0x00)
|
||||
|
||||
def test_constant_zero(self):
|
||||
"""Test selector 12 returns 0x00."""
|
||||
self.assertEqual(BYTE_PERMUTE(0xffffffffffffffff, 12), 0x00)
|
||||
|
||||
def test_constant_ff(self):
|
||||
"""Test selectors >= 13 return 0xFF."""
|
||||
for sel in [13, 14, 15, 255]:
|
||||
self.assertEqual(BYTE_PERMUTE(0, sel), 0xff, f"sel {sel} should be 0xff")
|
||||
|
||||
class TestSADHelpers(unittest.TestCase):
|
||||
"""Tests for V_SAD_U8 and V_MSAD_U8 helper functions."""
|
||||
|
||||
def test_v_sad_u8_basic(self):
|
||||
"""Test v_sad_u8 with simple values."""
|
||||
# s0 = 0x04030201, s1 = 0x04030201 -> diff = 0 for all bytes
|
||||
result = v_sad_u8(0x04030201, 0x04030201, 0)
|
||||
self.assertEqual(result, 0)
|
||||
# s0 = 0x05040302, s1 = 0x04030201 -> diff = 1+1+1+1 = 4
|
||||
result = v_sad_u8(0x05040302, 0x04030201, 0)
|
||||
self.assertEqual(result, 4)
|
||||
|
||||
def test_v_sad_u8_with_accumulator(self):
|
||||
"""Test v_sad_u8 with non-zero accumulator."""
|
||||
# s0 = 0x05040302, s1 = 0x04030201, s2 = 100 -> 4 + 100 = 104
|
||||
result = v_sad_u8(0x05040302, 0x04030201, 100)
|
||||
self.assertEqual(result, 104)
|
||||
|
||||
def test_v_sad_u8_large_diff(self):
|
||||
"""Test v_sad_u8 with maximum byte differences."""
|
||||
# s0 = 0xffffffff, s1 = 0x00000000 -> diff = 255*4 = 1020
|
||||
result = v_sad_u8(0xffffffff, 0x00000000, 0)
|
||||
self.assertEqual(result, 1020)
|
||||
|
||||
def test_v_msad_u8_basic(self):
|
||||
"""Test v_msad_u8 masks when reference byte is 0."""
|
||||
# s0 = 0x10101010, s1 = 0x00000000 -> all masked, result = 0
|
||||
result = v_msad_u8(0x10101010, 0x00000000, 0)
|
||||
self.assertEqual(result, 0)
|
||||
# s0 = 0x10101010, s1 = 0x01010101 -> diff = |0x10-0x01|*4 = 15*4 = 60
|
||||
result = v_msad_u8(0x10101010, 0x01010101, 0)
|
||||
self.assertEqual(result, 60)
|
||||
|
||||
def test_v_msad_u8_partial_mask(self):
|
||||
"""Test v_msad_u8 with partial masking."""
|
||||
# s0 = 0x10101010, s1 = 0x00010001 -> bytes 1 and 3 masked
|
||||
# diff = |0x10-0x01| + |0x10-0x01| = 15 + 15 = 30
|
||||
result = v_msad_u8(0x10101010, 0x00010001, 0)
|
||||
self.assertEqual(result, 30)
|
||||
|
||||
def test_v_msad_u8_with_accumulator(self):
|
||||
"""Test v_msad_u8 with non-zero accumulator."""
|
||||
result = v_msad_u8(0x10101010, 0x01010101, 50)
|
||||
self.assertEqual(result, 110) # 60 + 50
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,71 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test pdf.py PDF parser and enum generation."""
|
||||
import unittest, tempfile, importlib.util
|
||||
from extra.assembly.amd.pdf import extract, extract_tables, extract_enums, extract_pcode, write_enums, PDF_URLS
|
||||
|
||||
EXPECTED = {
|
||||
"rdna3": {"pages": 655, "tables": 115, "sop2_ops": 67, "sop2_first": "S_ADD_U32"},
|
||||
"rdna4": {"pages": 711, "tables": 125, "sop2_ops": 74, "sop2_first": "S_ADD_CO_U32"},
|
||||
"cdna": {"pages": 610, "tables": 104, "sop2_ops": 52, "sop2_first": "S_ADD_U32"},
|
||||
}
|
||||
|
||||
class TestPDF2(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.data = {name: extract(url) for name, url in PDF_URLS.items()}
|
||||
cls.tables = {name: extract_tables(pages) for name, pages in cls.data.items()}
|
||||
cls.enums = {name: extract_enums(cls.tables[name]) for name in PDF_URLS}
|
||||
cls.pcode = {name: extract_pcode(cls.data[name], cls.enums[name]) for name in PDF_URLS}
|
||||
|
||||
def test_page_counts(self):
|
||||
for name, exp in EXPECTED.items():
|
||||
self.assertEqual(len(self.data[name]), exp["pages"], f"{name} page count")
|
||||
|
||||
def test_table_counts(self):
|
||||
for name, exp in EXPECTED.items():
|
||||
self.assertEqual(len(self.tables[name]), exp["tables"], f"{name} table count")
|
||||
|
||||
def test_tables_sequential(self):
|
||||
for name in PDF_URLS:
|
||||
nums = sorted(self.tables[name].keys())
|
||||
missing = set(range(1, max(nums) + 1)) - set(nums)
|
||||
self.assertEqual(missing, set(), f"{name} missing tables: {missing}")
|
||||
|
||||
def test_generate_enums(self):
|
||||
for name, exp in EXPECTED.items():
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
|
||||
write_enums(self.enums[name], name, f.name)
|
||||
spec = importlib.util.spec_from_file_location("enum", f.name)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
# Check SOP2Op
|
||||
self.assertTrue(hasattr(mod, 'SOP2Op'), f"{name} missing SOP2Op")
|
||||
self.assertEqual(len(mod.SOP2Op), exp["sop2_ops"], f"{name} SOP2Op count")
|
||||
self.assertEqual(mod.SOP2Op(0).name, exp["sop2_first"], f"{name} SOP2Op first")
|
||||
# Check all enums have at least 2 ops
|
||||
for attr in dir(mod):
|
||||
if attr.endswith('Op'):
|
||||
self.assertGreaterEqual(len(getattr(mod, attr)), 2, f"{name} {attr} has too few ops")
|
||||
|
||||
def test_pcode_rdna3_tricky(self):
|
||||
"""Test specific pseudocode patterns that are tricky to extract correctly."""
|
||||
pcode = self.pcode['rdna3']
|
||||
# BUFFER_ATOMIC_MAX_U64: should have 4 statements (not truncated)
|
||||
self.assertEqual(pcode[('BUFFER_ATOMIC_MAX_U64', 72)],
|
||||
'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];\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")
|
||||
|
||||
def test_pcode_no_examples(self):
|
||||
"""Pseudocode should not contain example lines with '=>'."""
|
||||
for name in PDF_URLS:
|
||||
for (op_name, opcode), code in self.pcode[name].items():
|
||||
# No example lines (test vectors like "S_CTZ_I32_B32(0xaaaaaaaa) => 1")
|
||||
self.assertNotIn('=>', code, f"{name} {op_name} contains example line with '=>'")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+3
-2
@@ -1,11 +1,12 @@
|
||||
#!/usr/bin/env python3
|
||||
import unittest, subprocess
|
||||
from extra.assembly.rdna3.autogen import *
|
||||
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(
|
||||
["llvm-mc", "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
|
||||
[get_llvm_mc(), "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
|
||||
input=asm, capture_output=True, text=True
|
||||
)
|
||||
out = b''
|
||||
@@ -0,0 +1,258 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
|
||||
import unittest, io, sys, re, subprocess, os
|
||||
from extra.assembly.amd.dsl import Inst
|
||||
from extra.assembly.amd.asm import asm
|
||||
from extra.assembly.amd.decode import decode_inst, detect_format
|
||||
from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump
|
||||
|
||||
# arch: (mcpu, mattr)
|
||||
ARCH_CONFIG = {
|
||||
'rdna3': ('gfx1100', '+real-true16,+wavefrontsize32'),
|
||||
'rdna4': ('gfx1200', '+real-true16,+wavefrontsize32'),
|
||||
'cdna': ('gfx942', '+wavefrontsize64'),
|
||||
}
|
||||
|
||||
def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
|
||||
"""Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
|
||||
old_stdout = sys.stdout
|
||||
sys.stdout = io.StringIO()
|
||||
compiler.disassemble(lib)
|
||||
output = sys.stdout.getvalue()
|
||||
sys.stdout = old_stdout
|
||||
|
||||
results = []
|
||||
for line in output.splitlines():
|
||||
if '//' not in line: continue
|
||||
instr = line.split('//')[0].strip()
|
||||
if not instr: continue
|
||||
comment = line.split('//')[1].strip()
|
||||
if ':' not in comment: continue
|
||||
hex_str = comment.split(':')[1].strip().split()[0]
|
||||
try:
|
||||
machine_bytes = bytes.fromhex(hex_str)[::-1] # big-endian to little-endian
|
||||
results.append((instr, machine_bytes))
|
||||
except ValueError:
|
||||
continue
|
||||
return results
|
||||
|
||||
def compile_asm(instr: str, arch: str = 'rdna3') -> bytes:
|
||||
"""Compile a single instruction using LLVM."""
|
||||
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-mc call."""
|
||||
if not instrs: return []
|
||||
mcpu, mattr = ARCH_CONFIG[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)
|
||||
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 = ARCH_CONFIG[arch]
|
||||
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)
|
||||
|
||||
class TestTinygradKernelRoundtrip(unittest.TestCase):
|
||||
"""Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
|
||||
arch = 'rdna3'
|
||||
|
||||
def _test_kernel_roundtrip(self, op_fn):
|
||||
"""Generate kernel from op_fn, test:
|
||||
1. decode -> reencode matches original bytes
|
||||
2. asm(disasm()) matches LLVM output
|
||||
3. our disasm() matches LLVM's disassembly string exactly
|
||||
"""
|
||||
arch = self.arch
|
||||
mcpu, mattr = ARCH_CONFIG[arch]
|
||||
|
||||
from extra.assembly.amd.test.test_compare_emulators import get_kernels_from_tinygrad
|
||||
from tinygrad.runtime.support.compiler_amd import HIPCompiler
|
||||
|
||||
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
|
||||
compiler = HIPCompiler(mcpu)
|
||||
|
||||
# First pass: decode all instructions and collect info
|
||||
decoded_instrs: list[tuple] = [] # list of (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err)
|
||||
for ki, kernel in enumerate(kernels):
|
||||
offset = 0
|
||||
while offset < len(kernel.code):
|
||||
remaining = kernel.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
|
||||
|
||||
try:
|
||||
decoded = fmt.from_bytes(remaining) # pass all remaining bytes so from_bytes can read literal
|
||||
size = decoded.size() # actual size including literal
|
||||
orig_bytes = remaining[:size]
|
||||
reencoded = decoded.to_bytes()
|
||||
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))
|
||||
except Exception as e:
|
||||
decoded_instrs.append((ki, offset, remaining[:base_size], None, None, False, str(e)))
|
||||
size = base_size
|
||||
|
||||
offset += size
|
||||
|
||||
# Collect disasm strings for batched LLVM calls - skip unknown opcodes (op_X) that LLVM can't compile
|
||||
asm_test_instrs: list[tuple[int, str]] = [] # (idx, our_disasm) for asm test
|
||||
disasm_test_instrs: list[tuple[int, str]] = [] # (idx, our_disasm) for disasm comparison test
|
||||
|
||||
for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
|
||||
if our_disasm is None: continue
|
||||
# Skip unknown opcodes and malformed instructions for both tests
|
||||
if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm): continue
|
||||
asm_test_instrs.append((idx, our_disasm))
|
||||
disasm_test_instrs.append((idx, our_disasm))
|
||||
|
||||
# Batch compile for asm test
|
||||
asm_llvm_results = compile_asm_batch([d for _, d in asm_test_instrs], arch)
|
||||
asm_llvm_map = {idx: result for (idx, _), result in zip(asm_test_instrs, asm_llvm_results)}
|
||||
|
||||
# Batch compile+disasm for disasm comparison test
|
||||
disasm_llvm_results = compile_and_disasm_batch([d for _, d in disasm_test_instrs], arch)
|
||||
disasm_llvm_map = {idx: result for (idx, _), result in zip(disasm_test_instrs, disasm_llvm_results)}
|
||||
|
||||
# Now evaluate results
|
||||
decode_passed, decode_failed, decode_skipped = 0, 0, 0
|
||||
asm_passed, asm_failed, asm_skipped = 0, 0, 0
|
||||
disasm_passed, disasm_failed, disasm_skipped = 0, 0, 0
|
||||
decode_failures: list[str] = []
|
||||
asm_failures: list[str] = []
|
||||
disasm_failures: list[str] = []
|
||||
|
||||
for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
|
||||
# Decode test
|
||||
if decode_ok:
|
||||
decode_passed += 1
|
||||
elif decode_err == "no format":
|
||||
decode_skipped += 1
|
||||
else:
|
||||
decode_failed += 1
|
||||
decode_failures.append(f"K{ki}@{offset}: {our_disasm}: {decode_err}")
|
||||
|
||||
# Asm test
|
||||
if our_disasm is None:
|
||||
asm_skipped += 1
|
||||
elif idx in asm_llvm_map:
|
||||
llvm_bytes = asm_llvm_map[idx]
|
||||
try:
|
||||
our_bytes = asm(our_disasm).to_bytes()
|
||||
if our_bytes[:len(llvm_bytes)] == llvm_bytes:
|
||||
asm_passed += 1
|
||||
else:
|
||||
asm_failed += 1
|
||||
asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': ours={our_bytes[:len(llvm_bytes)].hex()} llvm={llvm_bytes.hex()}")
|
||||
except Exception:
|
||||
asm_skipped += 1
|
||||
else:
|
||||
asm_skipped += 1
|
||||
|
||||
# Disasm comparison test
|
||||
if our_disasm is None:
|
||||
disasm_skipped += 1
|
||||
elif idx in disasm_llvm_map:
|
||||
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}'")
|
||||
else:
|
||||
disasm_skipped += 1
|
||||
|
||||
print(f"[{arch}] decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
|
||||
print(f"[{arch}] asm vs 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, 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
|
||||
def test_neg(self): self._test_kernel_roundtrip(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
|
||||
def test_relu(self): self._test_kernel_roundtrip(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
|
||||
def test_exp(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).exp())
|
||||
def test_log(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 3.0]).log())
|
||||
def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
|
||||
def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
|
||||
def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
|
||||
|
||||
# Binary ops
|
||||
def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
|
||||
def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
|
||||
def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
|
||||
def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
|
||||
def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
|
||||
|
||||
# Reductions
|
||||
def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
|
||||
def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
|
||||
def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
|
||||
|
||||
# Matmul
|
||||
def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
|
||||
def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
|
||||
|
||||
# Complex ops
|
||||
def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
|
||||
def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
|
||||
|
||||
# Memory patterns
|
||||
def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
|
||||
def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
|
||||
def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
|
||||
|
||||
# Cast ops
|
||||
def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
|
||||
def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
|
||||
|
||||
# Comparison ops
|
||||
def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
|
||||
def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
|
||||
|
||||
# Fused ops
|
||||
def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
|
||||
|
||||
@unittest.skip("no asm support for RDNA4")
|
||||
class TestTinygradKernelRoundtripRDNA4(TestTinygradKernelRoundtrip): arch = 'rdna4'
|
||||
|
||||
@unittest.skip("no asm support for CDNA")
|
||||
class TestTinygradKernelRoundtripCDNA(TestTinygradKernelRoundtrip): arch = 'cdna'
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,206 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Tests for SQTT packet decoding using real captured examples."""
|
||||
import pickle, unittest, ctypes
|
||||
from pathlib import Path
|
||||
from tinygrad.helpers import DEBUG, colored
|
||||
from tinygrad.runtime.autogen import rocprof
|
||||
from tinygrad.runtime.support.elf import elf_loader
|
||||
from extra.assembly.amd.decode 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, WAVEEND, INST, VALUINST, IMMEDIATE, IMMEDIATE_MASK,
|
||||
ALUEXEC, VMEMEXEC, PACKET_TYPES, InstOp, AluSrc, MemSrc)
|
||||
|
||||
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}
|
||||
|
||||
PACKET_COLORS = {
|
||||
"INST": "WHITE", "VALUINST": "BLACK", "VMEMEXEC": "yellow", "ALUEXEC": "yellow",
|
||||
"IMMEDIATE": "YELLOW", "IMMEDIATE_MASK": "YELLOW", "WAVERDY": "cyan", "WAVEALLOC": "cyan",
|
||||
"WAVEEND": "blue", "WAVESTART": "blue", "PERF": "magenta", "EVENT": "red", "EVENT_BIG": "red",
|
||||
"REG": "green", "LAYOUT_HEADER": "white", "SNAPSHOT": "white", "UTILCTR": "green",
|
||||
}
|
||||
|
||||
def format_packet(p, time_offset: int = 0) -> str:
|
||||
name, cycle = type(p).__name__, p._time - time_offset
|
||||
if isinstance(p, INST):
|
||||
op_name = p.op.name if isinstance(p.op, InstOp) else f"0x{p.op:02x}"
|
||||
fields = f"wave={p.wave} op={op_name}" + (" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "")
|
||||
elif isinstance(p, VALUINST): fields = f"wave={p.wave}" + (" flag" if p.flag else "")
|
||||
elif isinstance(p, ALUEXEC): fields = f"src={p.src.name if isinstance(p.src, AluSrc) else p.src}"
|
||||
elif isinstance(p, VMEMEXEC): fields = f"src={p.src.name if isinstance(p.src, MemSrc) else p.src}"
|
||||
elif isinstance(p, (WAVESTART, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
|
||||
elif hasattr(p, '_values'):
|
||||
fields = " ".join(f"{k}=0x{v:x}" if k in {'snap', 'val32'} else f"{k}={v}"
|
||||
for k, v in p._values.items() if not k.startswith('_') and k != 'delta')
|
||||
else: fields = ""
|
||||
return f"{cycle:8}: {colored(f'{name:18}', PACKET_COLORS.get(name, 'white'))} {fields}"
|
||||
|
||||
def print_packets(packets: list) -> None:
|
||||
skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"}
|
||||
time_offset = packets[0]._time if packets else 0
|
||||
for p in packets:
|
||||
if type(p).__name__ not in skip: print(format_packet(p, time_offset))
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# ROCPROF DECODER
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int):
|
||||
"""Run rocprof decoder on SQTT blobs, returning raw occupancy and instruction records."""
|
||||
image, sections, _ = elf_loader(lib)
|
||||
text = next((sh for sh in sections if sh.name == ".text"), None)
|
||||
assert text is not None, "no .text section found"
|
||||
text_off, text_size = text.header.sh_addr, text.header.sh_size
|
||||
|
||||
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, _):
|
||||
blob = next(blob_iter, None)
|
||||
if blob is None: return 0
|
||||
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, _):
|
||||
if record_type == rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
|
||||
for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr):
|
||||
occupancy_records.append((ev.wave_id, ev.simd, ev.cu, ev.time, ev.start))
|
||||
elif record_type == rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
|
||||
for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr):
|
||||
if ev.instructions_size > 0:
|
||||
sz = ev.instructions_size * ctypes.sizeof(rocprof.rocprofiler_thread_trace_decoder_inst_t)
|
||||
insts_blob = bytearray(sz)
|
||||
ctypes.memmove((ctypes.c_char * sz).from_buffer(insts_blob), ev.instructions_array, sz)
|
||||
insts = list((rocprof.rocprofiler_thread_trace_decoder_inst_t * ev.instructions_size).from_buffer(insts_blob))
|
||||
wave_insts.append([(inst.time, inst.stall) for inst in insts])
|
||||
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
|
||||
|
||||
@rocprof.rocprof_trace_decoder_isa_callback_t
|
||||
def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
|
||||
offset = pc.address - base
|
||||
if offset < text_off or offset >= text_off + text_size:
|
||||
mem_size_ptr[0] = 0
|
||||
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
|
||||
try:
|
||||
inst = decode_inst(data := image[offset:])
|
||||
mem_size_ptr[0] = inst._size()
|
||||
except (ValueError, AssertionError):
|
||||
mem_size_ptr[0] = 0
|
||||
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
|
||||
if isinstance(inst, SOPP) and inst.op == SOPPOp.S_ENDPGM: mem_size_ptr[0] = 0
|
||||
# rocprof parses instruction string to determine type; v_nop works for all
|
||||
if (max_sz := size_ptr[0]) == 0: return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES
|
||||
ctypes.memmove(instr_ptr, b"v_nop", min(5, max_sz - 1))
|
||||
size_ptr[0] = min(5, max_sz - 1)
|
||||
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
|
||||
|
||||
rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
|
||||
return occupancy_records, wave_insts
|
||||
|
||||
class TestSQTTExamples(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.examples = {}
|
||||
for pkl_path in sorted(EXAMPLES_DIR.glob("*.pkl")):
|
||||
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 sqtt_events and prg:
|
||||
cls.examples[pkl_path.stem] = (sqtt_events, prg.lib, prg.base)
|
||||
|
||||
def test_examples_loaded(self):
|
||||
self.assertGreater(len(self.examples), 0, "no example files found")
|
||||
|
||||
def test_decode_all_examples(self):
|
||||
for name, (events, *_) in self.examples.items():
|
||||
for i, event in enumerate(events):
|
||||
with self.subTest(example=name, event=i):
|
||||
packets = decode(event.blob)
|
||||
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):
|
||||
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):
|
||||
self.assertIn(type(pkt), PACKET_TYPES, f"unknown packet type {type(pkt)} in {name}")
|
||||
|
||||
def test_wave_lifecycle(self):
|
||||
for name, (events, *_) in self.examples.items():
|
||||
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)]), 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):
|
||||
for name, (events, *_) in self.examples.items():
|
||||
for i, event in enumerate(events):
|
||||
with self.subTest(example=name, event=i):
|
||||
times = [p._time for p in decode(event.blob)]
|
||||
self.assertEqual(times, sorted(times), f"timestamps not monotonic in {name}")
|
||||
|
||||
def test_gemm_has_instructions(self):
|
||||
for name, (events, *_) in self.examples.items():
|
||||
if "gemm" not 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, INST)]), 0, f"no INST packets in {name}")
|
||||
|
||||
def test_rocprof_wave_times_match(self):
|
||||
"""Wave start/end times must match rocprof exactly."""
|
||||
for name, (events, lib, base) in self.examples.items():
|
||||
with self.subTest(example=name):
|
||||
occupancy, _ = run_rocprof_decoder([e.blob for e in events], lib, base)
|
||||
# extract from rocprof occupancy records
|
||||
roc_starts: dict[tuple[int, int, int], int] = {}
|
||||
roc_waves: list[tuple[int, int]] = []
|
||||
for wave_id, simd, cu, time, is_start in occupancy:
|
||||
key = (wave_id, simd, cu)
|
||||
if is_start: roc_starts[key] = time
|
||||
elif key in roc_starts: roc_waves.append((roc_starts.pop(key), time))
|
||||
# extract from our decoder
|
||||
our_waves: list[tuple[int, int]] = []
|
||||
for event in events:
|
||||
packets = decode(event.blob)
|
||||
wave_starts: dict[tuple[int, int, int], int] = {}
|
||||
for p in packets:
|
||||
if isinstance(p, WAVESTART): 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}")
|
||||
|
||||
def test_rocprof_inst_times_match(self):
|
||||
"""Instruction times must match rocprof exactly (excluding s_endpgm)."""
|
||||
for name, (events, lib, base) in self.examples.items():
|
||||
with self.subTest(example=name):
|
||||
_, wave_insts = run_rocprof_decoder([e.blob for e in events], lib, base)
|
||||
# skip last inst per wave (s_endpgm) - it needs special handling (time + duration instead of time + stall)
|
||||
roc_insts = [time + stall for insts in wave_insts for time, stall in insts[:-1]]
|
||||
# extract from our decoder
|
||||
our_insts: list[int] = []
|
||||
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, VALUINST): our_insts.append(p._time)
|
||||
elif isinstance(p, IMMEDIATE): our_insts.append(p._time)
|
||||
elif isinstance(p, IMMEDIATE_MASK):
|
||||
for _ in range(bin(p.mask).count('1')): our_insts.append(p._time)
|
||||
self.assertEqual(sorted(our_insts), sorted(roc_insts), f"instruction times mismatch in {name}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,254 +0,0 @@
|
||||
# Pure combinational ALU functions for RDNA3 emulation
|
||||
from __future__ import annotations
|
||||
import struct, math
|
||||
from typing import Callable
|
||||
from extra.assembly.rdna3.autogen import SOP1Op, SOP2Op, SOPCOp, SOPKOp, VOP1Op, VOP2Op, VOP3Op
|
||||
|
||||
# Format base offsets for unified opcode space
|
||||
SOP2_BASE, SOP1_BASE, SOPC_BASE, SOPK_BASE = 0x000, 0x100, 0x200, 0x300
|
||||
VOP2_BASE, VOP1_BASE = 0x100, 0x180
|
||||
|
||||
# Float conversion helpers
|
||||
_I, _f, _H, _e = struct.Struct('<I'), struct.Struct('<f'), struct.Struct('<H'), struct.Struct('<e')
|
||||
def f32(i: int) -> float: return _f.unpack(_I.pack(i & 0xffffffff))[0]
|
||||
def i32(f: float) -> int:
|
||||
if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
|
||||
try: return _I.unpack(_f.pack(f))[0]
|
||||
except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
|
||||
def f16(i: int) -> float: return _e.unpack(_H.pack(i & 0xffff))[0]
|
||||
def i16(f: float) -> int:
|
||||
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
|
||||
try: return _H.unpack(_e.pack(f))[0]
|
||||
except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
|
||||
def sext(v: int, b: int) -> int: return v - (1 << b) if v & (1 << (b-1)) else v
|
||||
def clz(x: int) -> int: return 32 - x.bit_length() if x else 32
|
||||
def cls(x: int) -> int: x &= 0xffffffff; return 31 if x in (0, 0xffffffff) else clz(~x & 0xffffffff if x >> 31 else x) - 1
|
||||
def _cvt_i32_f32(v): return (0x7fffffff if v > 0 else 0x80000000) if math.isinf(v) else (0 if math.isnan(v) else max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff)
|
||||
def _cvt_u32_f32(v): return (0xffffffff if v > 0 else 0) if math.isinf(v) else (0 if math.isnan(v) or v < 0 else min(0xffffffff, int(v)))
|
||||
|
||||
# SALU: op -> fn(s0, s1, scc_in) -> (result, scc_out)
|
||||
SALU: dict[int, Callable] = {
|
||||
# SOP2
|
||||
SOP2_BASE + SOP2Op.S_ADD_U32: lambda a, b, scc: ((a + b) & 0xffffffff, int((a + b) >= 0x100000000)),
|
||||
SOP2_BASE + SOP2Op.S_SUB_U32: lambda a, b, scc: ((a - b) & 0xffffffff, int(b > a)),
|
||||
SOP2_BASE + SOP2Op.S_ADDC_U32: lambda a, b, scc: ((r := a + b + scc) & 0xffffffff, int(r >= 0x100000000)),
|
||||
SOP2_BASE + SOP2Op.S_SUBB_U32: lambda a, b, scc: ((a - b - scc) & 0xffffffff, int((b + scc) > a)),
|
||||
SOP2_BASE + SOP2Op.S_ADD_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 32)) & 0xffffffff, int(((a >> 31) == (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
|
||||
SOP2_BASE + SOP2Op.S_SUB_I32: lambda a, b, scc: ((r := sext(a, 32) - sext(b, 32)) & 0xffffffff, int(((a >> 31) != (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
|
||||
SOP2_BASE + SOP2Op.S_AND_B32: lambda a, b, scc: ((r := a & b), int(r != 0)),
|
||||
SOP2_BASE + SOP2Op.S_OR_B32: lambda a, b, scc: ((r := a | b), int(r != 0)),
|
||||
SOP2_BASE + SOP2Op.S_XOR_B32: lambda a, b, scc: ((r := a ^ b), int(r != 0)),
|
||||
SOP2_BASE + SOP2Op.S_AND_NOT1_B32: lambda a, b, scc: ((r := a & (~b & 0xffffffff)), int(r != 0)),
|
||||
SOP2_BASE + SOP2Op.S_OR_NOT1_B32: lambda a, b, scc: ((r := a | (~b & 0xffffffff)), int(r != 0)),
|
||||
SOP2_BASE + SOP2Op.S_LSHL_B32: lambda a, b, scc: ((r := (a << (b & 0x1f)) & 0xffffffff), int(r != 0)),
|
||||
SOP2_BASE + SOP2Op.S_LSHR_B32: lambda a, b, scc: ((r := a >> (b & 0x1f)), int(r != 0)),
|
||||
SOP2_BASE + SOP2Op.S_ASHR_I32: lambda a, b, scc: ((r := sext(a, 32) >> (b & 0x1f)) & 0xffffffff, int(r != 0)),
|
||||
SOP2_BASE + SOP2Op.S_MUL_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 32)) & 0xffffffff, scc),
|
||||
SOP2_BASE + SOP2Op.S_MUL_HI_U32: lambda a, b, scc: (((a * b) >> 32) & 0xffffffff, scc),
|
||||
SOP2_BASE + SOP2Op.S_MUL_HI_I32: lambda a, b, scc: (((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff, scc),
|
||||
SOP2_BASE + SOP2Op.S_MIN_I32: lambda a, b, scc: (a, 1) if sext(a, 32) < sext(b, 32) else (b, 0),
|
||||
SOP2_BASE + SOP2Op.S_MIN_U32: lambda a, b, scc: (a, 1) if a < b else (b, 0),
|
||||
SOP2_BASE + SOP2Op.S_MAX_I32: lambda a, b, scc: (a, 1) if sext(a, 32) > sext(b, 32) else (b, 0),
|
||||
SOP2_BASE + SOP2Op.S_MAX_U32: lambda a, b, scc: (a, 1) if a > b else (b, 0),
|
||||
SOP2_BASE + SOP2Op.S_CSELECT_B32: lambda a, b, scc: (a if scc else b, scc),
|
||||
SOP2_BASE + SOP2Op.S_BFE_U32: lambda a, b, scc: ((r := ((a >> (b & 0x1f)) & ((1 << ((b >> 16) & 0x7f)) - 1)) if (b >> 16) & 0x7f else 0), int(r != 0)),
|
||||
SOP2_BASE + SOP2Op.S_BFE_I32: lambda a, b, scc: ((r := sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := (b >> 16) & 0x7f) else 0), int(r != 0)),
|
||||
SOP2_BASE + SOP2Op.S_PACK_LL_B32_B16: lambda a, b, scc: ((a & 0xffff) | ((b & 0xffff) << 16), scc),
|
||||
SOP2_BASE + SOP2Op.S_PACK_LH_B32_B16: lambda a, b, scc: ((a & 0xffff) | (b & 0xffff0000), scc),
|
||||
SOP2_BASE + SOP2Op.S_PACK_HH_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | (b & 0xffff0000), scc),
|
||||
SOP2_BASE + SOP2Op.S_PACK_HL_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | ((b & 0xffff) << 16), scc),
|
||||
SOP2_BASE + SOP2Op.S_ADD_F32: lambda a, b, scc: (i32(f32(a) + f32(b)), scc),
|
||||
SOP2_BASE + SOP2Op.S_SUB_F32: lambda a, b, scc: (i32(f32(a) - f32(b)), scc),
|
||||
SOP2_BASE + SOP2Op.S_MUL_F32: lambda a, b, scc: (i32(f32(a) * f32(b)), scc),
|
||||
# SOP1
|
||||
SOP1_BASE + SOP1Op.S_MOV_B32: lambda a, b, scc: (a, scc),
|
||||
SOP1_BASE + SOP1Op.S_NOT_B32: lambda a, b, scc: ((r := (~a) & 0xffffffff), int(r != 0)),
|
||||
SOP1_BASE + SOP1Op.S_BREV_B32: lambda a, b, scc: (int(f'{a & 0xffffffff:032b}'[::-1], 2), scc),
|
||||
SOP1_BASE + SOP1Op.S_CLZ_I32_U32: lambda a, b, scc: (clz(a), scc),
|
||||
SOP1_BASE + SOP1Op.S_CLS_I32: lambda a, b, scc: (cls(a), scc),
|
||||
SOP1_BASE + SOP1Op.S_SEXT_I32_I8: lambda a, b, scc: (sext(a & 0xff, 8) & 0xffffffff, scc),
|
||||
SOP1_BASE + SOP1Op.S_SEXT_I32_I16: lambda a, b, scc: (sext(a & 0xffff, 16) & 0xffffffff, scc),
|
||||
SOP1_BASE + SOP1Op.S_ABS_I32: lambda a, b, scc: ((r := abs(sext(a, 32)) & 0xffffffff), int(r != 0)),
|
||||
SOP1_BASE + SOP1Op.S_CVT_F32_I32: lambda a, b, scc: (i32(float(sext(a, 32))), scc),
|
||||
SOP1_BASE + SOP1Op.S_CVT_F32_U32: lambda a, b, scc: (i32(float(a)), scc),
|
||||
SOP1_BASE + SOP1Op.S_CVT_I32_F32: lambda a, b, scc: (_cvt_i32_f32(f32(a)), scc),
|
||||
SOP1_BASE + SOP1Op.S_CVT_U32_F32: lambda a, b, scc: (_cvt_u32_f32(f32(a)), scc),
|
||||
SOP1_BASE + SOP1Op.S_CEIL_F32: lambda a, b, scc: (i32(math.ceil(f32(a))), scc),
|
||||
SOP1_BASE + SOP1Op.S_FLOOR_F32: lambda a, b, scc: (i32(math.floor(f32(a))), scc),
|
||||
SOP1_BASE + SOP1Op.S_TRUNC_F32: lambda a, b, scc: (i32(math.trunc(f32(a))), scc),
|
||||
SOP1_BASE + SOP1Op.S_RNDNE_F32: lambda a, b, scc: (i32(round(f32(a))), scc),
|
||||
SOP1_BASE + SOP1Op.S_CVT_F16_F32: lambda a, b, scc: (i16(f32(a)), scc),
|
||||
SOP1_BASE + SOP1Op.S_CVT_F32_F16: lambda a, b, scc: (i32(f16(a)), scc),
|
||||
# SOPC
|
||||
SOPC_BASE + SOPCOp.S_CMP_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 32))),
|
||||
SOPC_BASE + SOPCOp.S_CMP_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 32))),
|
||||
SOPC_BASE + SOPCOp.S_CMP_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 32))),
|
||||
SOPC_BASE + SOPCOp.S_CMP_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 32))),
|
||||
SOPC_BASE + SOPCOp.S_CMP_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 32))),
|
||||
SOPC_BASE + SOPCOp.S_CMP_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 32))),
|
||||
SOPC_BASE + SOPCOp.S_CMP_EQ_U32: lambda a, b, scc: (0, int(a == b)),
|
||||
SOPC_BASE + SOPCOp.S_CMP_LG_U32: lambda a, b, scc: (0, int(a != b)),
|
||||
SOPC_BASE + SOPCOp.S_CMP_GT_U32: lambda a, b, scc: (0, int(a > b)),
|
||||
SOPC_BASE + SOPCOp.S_CMP_GE_U32: lambda a, b, scc: (0, int(a >= b)),
|
||||
SOPC_BASE + SOPCOp.S_CMP_LT_U32: lambda a, b, scc: (0, int(a < b)),
|
||||
SOPC_BASE + SOPCOp.S_CMP_LE_U32: lambda a, b, scc: (0, int(a <= b)),
|
||||
SOPC_BASE + SOPCOp.S_BITCMP0_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) == 0)),
|
||||
SOPC_BASE + SOPCOp.S_BITCMP1_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) != 0)),
|
||||
# SOPK
|
||||
SOPK_BASE + SOPKOp.S_MOVK_I32: lambda a, b, scc: (sext(b, 16) & 0xffffffff, scc),
|
||||
SOPK_BASE + SOPKOp.S_CMOVK_I32: lambda a, b, scc: ((sext(b, 16) & 0xffffffff) if scc else a, scc),
|
||||
SOPK_BASE + SOPKOp.S_ADDK_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 16)) & 0xffffffff, int(((a >> 31) == ((b >> 15) & 1)) and ((a >> 31) != ((r >> 31) & 1)))),
|
||||
SOPK_BASE + SOPKOp.S_MULK_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 16)) & 0xffffffff, scc),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 16))),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 16))),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 16))),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 16))),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 16))),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 16))),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_EQ_U32: lambda a, b, scc: (0, int(a == (b & 0xffff))),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_LG_U32: lambda a, b, scc: (0, int(a != (b & 0xffff))),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_GT_U32: lambda a, b, scc: (0, int(a > (b & 0xffff))),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_GE_U32: lambda a, b, scc: (0, int(a >= (b & 0xffff))),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_LT_U32: lambda a, b, scc: (0, int(a < (b & 0xffff))),
|
||||
SOPK_BASE + SOPKOp.S_CMPK_LE_U32: lambda a, b, scc: (0, int(a <= (b & 0xffff))),
|
||||
}
|
||||
|
||||
# VALU: op -> fn(s0, s1, s2) -> result
|
||||
VALU: dict[int, Callable] = {
|
||||
# VOP2
|
||||
VOP2_BASE + VOP2Op.V_ADD_F32: lambda a, b, c: i32(f32(a) + f32(b)),
|
||||
VOP2_BASE + VOP2Op.V_SUB_F32: lambda a, b, c: i32(f32(a) - f32(b)),
|
||||
VOP2_BASE + VOP2Op.V_SUBREV_F32: lambda a, b, c: i32(f32(b) - f32(a)),
|
||||
VOP2_BASE + VOP2Op.V_MUL_F32: lambda a, b, c: i32(f32(a) * f32(b)),
|
||||
VOP2_BASE + VOP2Op.V_MIN_F32: lambda a, b, c: i32(min(f32(a), f32(b))),
|
||||
VOP2_BASE + VOP2Op.V_MAX_F32: lambda a, b, c: i32(max(f32(a), f32(b))),
|
||||
VOP2_BASE + VOP2Op.V_ADD_NC_U32: lambda a, b, c: (a + b) & 0xffffffff,
|
||||
VOP2_BASE + VOP2Op.V_SUB_NC_U32: lambda a, b, c: (a - b) & 0xffffffff,
|
||||
VOP2_BASE + VOP2Op.V_SUBREV_NC_U32: lambda a, b, c: (b - a) & 0xffffffff,
|
||||
VOP2_BASE + VOP2Op.V_AND_B32: lambda a, b, c: a & b,
|
||||
VOP2_BASE + VOP2Op.V_OR_B32: lambda a, b, c: a | b,
|
||||
VOP2_BASE + VOP2Op.V_XOR_B32: lambda a, b, c: a ^ b,
|
||||
VOP2_BASE + VOP2Op.V_XNOR_B32: lambda a, b, c: (~(a ^ b)) & 0xffffffff,
|
||||
VOP2_BASE + VOP2Op.V_LSHLREV_B32: lambda a, b, c: (b << (a & 0x1f)) & 0xffffffff,
|
||||
VOP2_BASE + VOP2Op.V_LSHRREV_B32: lambda a, b, c: b >> (a & 0x1f),
|
||||
VOP2_BASE + VOP2Op.V_ASHRREV_I32: lambda a, b, c: (sext(b, 32) >> (a & 0x1f)) & 0xffffffff,
|
||||
VOP2_BASE + VOP2Op.V_MIN_I32: lambda a, b, c: a if sext(a, 32) < sext(b, 32) else b,
|
||||
VOP2_BASE + VOP2Op.V_MAX_I32: lambda a, b, c: a if sext(a, 32) > sext(b, 32) else b,
|
||||
VOP2_BASE + VOP2Op.V_MIN_U32: lambda a, b, c: min(a, b),
|
||||
VOP2_BASE + VOP2Op.V_MAX_U32: lambda a, b, c: max(a, b),
|
||||
VOP2_BASE + VOP2Op.V_MUL_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) & 0xffffffff,
|
||||
VOP2_BASE + VOP2Op.V_MUL_HI_I32_I24: lambda a, b, c: ((sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) >> 32) & 0xffffffff,
|
||||
VOP2_BASE + VOP2Op.V_MUL_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff)) & 0xffffffff,
|
||||
VOP2_BASE + VOP2Op.V_MUL_HI_U32_U24: lambda a, b, c: (((a & 0xffffff) * (b & 0xffffff)) >> 32) & 0xffffffff,
|
||||
VOP2_BASE + VOP2Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
|
||||
VOP2_BASE + VOP2Op.V_LDEXP_F16: lambda a, b, c: i16(math.ldexp(f16(a), sext(b, 32))),
|
||||
VOP2_BASE + VOP2Op.V_ADD_F16: lambda a, b, c: i16(f16(a) + f16(b)),
|
||||
VOP2_BASE + VOP2Op.V_SUB_F16: lambda a, b, c: i16(f16(a) - f16(b)),
|
||||
VOP2_BASE + VOP2Op.V_MUL_F16: lambda a, b, c: i16(f16(a) * f16(b)),
|
||||
VOP2_BASE + VOP2Op.V_MIN_F16: lambda a, b, c: i16(min(f16(a), f16(b))),
|
||||
VOP2_BASE + VOP2Op.V_MAX_F16: lambda a, b, c: i16(max(f16(a), f16(b))),
|
||||
# VOP1
|
||||
VOP1_BASE + VOP1Op.V_MOV_B32: lambda a, b, c: a,
|
||||
VOP1_BASE + VOP1Op.V_NOT_B32: lambda a, b, c: (~a) & 0xffffffff,
|
||||
VOP1_BASE + VOP1Op.V_BFREV_B32: lambda a, b, c: int(f'{a & 0xffffffff:032b}'[::-1], 2),
|
||||
VOP1_BASE + VOP1Op.V_CLZ_I32_U32: lambda a, b, c: clz(a),
|
||||
VOP1_BASE + VOP1Op.V_CLS_I32: lambda a, b, c: cls(a),
|
||||
VOP1_BASE + VOP1Op.V_CVT_F32_I32: lambda a, b, c: i32(float(sext(a, 32))),
|
||||
VOP1_BASE + VOP1Op.V_CVT_F32_U32: lambda a, b, c: i32(float(a)),
|
||||
VOP1_BASE + VOP1Op.V_CVT_I32_F32: lambda a, b, c: _cvt_i32_f32(f32(a)),
|
||||
VOP1_BASE + VOP1Op.V_CVT_U32_F32: lambda a, b, c: _cvt_u32_f32(f32(a)),
|
||||
VOP1_BASE + VOP1Op.V_CVT_F16_F32: lambda a, b, c: i16(f32(a)),
|
||||
VOP1_BASE + VOP1Op.V_CVT_F32_F16: lambda a, b, c: i32(f16(a)),
|
||||
VOP1_BASE + VOP1Op.V_RCP_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
|
||||
VOP1_BASE + VOP1Op.V_RCP_IFLAG_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
|
||||
VOP1_BASE + VOP1Op.V_RSQ_F32: lambda a, b, c: i32(1.0 / math.sqrt(f32(a)) if f32(a) > 0 else (float('nan') if f32(a) < 0 else float('inf'))),
|
||||
VOP1_BASE + VOP1Op.V_SQRT_F32: lambda a, b, c: i32(math.sqrt(f32(a)) if f32(a) >= 0 else float('nan')),
|
||||
VOP1_BASE + VOP1Op.V_LOG_F32: lambda a, b, c: i32(math.log2(f32(a)) if f32(a) > 0 else (float('-inf') if f32(a) == 0 else float('nan'))),
|
||||
VOP1_BASE + VOP1Op.V_EXP_F32: lambda a, b, c: i32(float('inf') if f32(a) > 128 else (0.0 if f32(a) < -150 else math.pow(2.0, f32(a)))),
|
||||
VOP1_BASE + VOP1Op.V_SIN_F32: lambda a, b, c: i32(math.sin(f32(a) * 2 * math.pi)),
|
||||
VOP1_BASE + VOP1Op.V_COS_F32: lambda a, b, c: i32(math.cos(f32(a) * 2 * math.pi)),
|
||||
VOP1_BASE + VOP1Op.V_FLOOR_F32: lambda a, b, c: i32(math.floor(f32(a))),
|
||||
VOP1_BASE + VOP1Op.V_CEIL_F32: lambda a, b, c: i32(math.ceil(f32(a))),
|
||||
VOP1_BASE + VOP1Op.V_TRUNC_F32: lambda a, b, c: i32(math.trunc(f32(a))),
|
||||
VOP1_BASE + VOP1Op.V_RNDNE_F32: lambda a, b, c: i32(round(f32(a))),
|
||||
VOP1_BASE + VOP1Op.V_FRACT_F32: lambda a, b, c: i32((v := f32(a)) - math.floor(v)),
|
||||
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE0: lambda a, b, c: i32(float(a & 0xff)),
|
||||
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE1: lambda a, b, c: i32(float((a >> 8) & 0xff)),
|
||||
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE2: lambda a, b, c: i32(float((a >> 16) & 0xff)),
|
||||
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE3: lambda a, b, c: i32(float((a >> 24) & 0xff)),
|
||||
VOP1_BASE + VOP1Op.V_FREXP_MANT_F32: lambda a, b, c: i32(math.frexp(v)[0] if (v := f32(a)) != 0 else 0.0),
|
||||
VOP1_BASE + VOP1Op.V_FREXP_EXP_I32_F32: lambda a, b, c: (math.frexp(v)[1] if (v := f32(a)) != 0 else 0) & 0xffffffff,
|
||||
# VOP3
|
||||
VOP3Op.V_FMA_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
|
||||
VOP3Op.V_DIV_FMAS_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
|
||||
VOP3Op.V_ADD3_U32: lambda a, b, c: (a + b + c) & 0xffffffff,
|
||||
VOP3Op.V_LSHL_ADD_U32: lambda a, b, c: ((a << (b & 0x1f)) + c) & 0xffffffff,
|
||||
VOP3Op.V_ADD_LSHL_U32: lambda a, b, c: ((a + b) << (c & 0x1f)) & 0xffffffff,
|
||||
VOP3Op.V_XOR3_B32: lambda a, b, c: a ^ b ^ c,
|
||||
VOP3Op.V_OR3_B32: lambda a, b, c: a | b | c,
|
||||
VOP3Op.V_AND_OR_B32: lambda a, b, c: (a & b) | c,
|
||||
VOP3Op.V_LSHL_OR_B32: lambda a, b, c: ((a << (b & 0x1f)) | c) & 0xffffffff,
|
||||
VOP3Op.V_XAD_U32: lambda a, b, c: ((a ^ b) + c) & 0xffffffff,
|
||||
VOP3Op.V_MAD_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff) + c) & 0xffffffff,
|
||||
VOP3Op.V_MAD_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24) + sext(c, 32)) & 0xffffffff,
|
||||
VOP3Op.V_BFE_U32: lambda a, b, c: (a >> (b & 0x1f)) & ((1 << (c & 0x1f)) - 1) if c & 0x1f else 0,
|
||||
VOP3Op.V_BFE_I32: lambda a, b, c: sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := c & 0x1f) else 0,
|
||||
VOP3Op.V_ALIGNBIT_B32: lambda a, b, c: (((a << 32) | b) >> (c & 0x1f)) & 0xffffffff,
|
||||
VOP3Op.V_MUL_LO_U32: lambda a, b, c: (a * b) & 0xffffffff,
|
||||
VOP3Op.V_MUL_HI_U32: lambda a, b, c: ((a * b) >> 32) & 0xffffffff,
|
||||
VOP3Op.V_MUL_HI_I32: lambda a, b, c: ((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff,
|
||||
VOP3Op.V_LDEXP_F32: lambda a, b, c: i32(math.ldexp(f32(a), sext(b, 32))),
|
||||
VOP3Op.V_DIV_FIXUP_F32: lambda a, b, c: i32(math.copysign(float('inf'), f32(c)) if f32(b) == 0.0 else f32(c) / f32(b)),
|
||||
VOP3Op.V_PACK_B32_F16: lambda a, b, c: (a & 0xffff) | ((b & 0xffff) << 16),
|
||||
VOP3Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
|
||||
VOP3Op.V_LSHLREV_B16: lambda a, b, c: ((b & 0xffff) << (a & 0xf)) & 0xffff,
|
||||
VOP3Op.V_LSHRREV_B16: lambda a, b, c: (b & 0xffff) >> (a & 0xf),
|
||||
VOP3Op.V_ASHRREV_I16: lambda a, b, c: (sext(b & 0xffff, 16) >> (a & 0xf)) & 0xffff,
|
||||
VOP3Op.V_ADD_NC_U16: lambda a, b, c: ((a & 0xffff) + (b & 0xffff)) & 0xffff,
|
||||
VOP3Op.V_SUB_NC_U16: lambda a, b, c: ((a & 0xffff) - (b & 0xffff)) & 0xffff,
|
||||
VOP3Op.V_MUL_LO_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff)) & 0xffff,
|
||||
VOP3Op.V_MIN_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff),
|
||||
VOP3Op.V_MAX_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff),
|
||||
VOP3Op.V_MIN_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) < sext(b & 0xffff, 16) else (b & 0xffff),
|
||||
VOP3Op.V_MAX_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) > sext(b & 0xffff, 16) else (b & 0xffff),
|
||||
VOP3Op.V_MAD_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff) + (c & 0xffff)) & 0xffff,
|
||||
VOP3Op.V_MAD_I16: lambda a, b, c: (sext(a & 0xffff, 16) * sext(b & 0xffff, 16) + sext(c & 0xffff, 16)) & 0xffff,
|
||||
VOP3Op.V_FMA_F16: lambda a, b, c: i16(f16(a) * f16(b) + f16(c)),
|
||||
VOP3Op.V_MIN3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[0] & 0xffffffff,
|
||||
VOP3Op.V_MAX3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[2] & 0xffffffff,
|
||||
VOP3Op.V_MED3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[1] & 0xffffffff,
|
||||
VOP3Op.V_MIN3_F16: lambda a, b, c: i16(min(f16(a), f16(b), f16(c))),
|
||||
VOP3Op.V_MAX3_F16: lambda a, b, c: i16(max(f16(a), f16(b), f16(c))),
|
||||
VOP3Op.V_MED3_F16: lambda a, b, c: i16(sorted([f16(a), f16(b), f16(c)])[1]),
|
||||
VOP3Op.V_MIN3_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff, c & 0xffff),
|
||||
VOP3Op.V_MAX3_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff, c & 0xffff),
|
||||
VOP3Op.V_MED3_U16: lambda a, b, c: sorted([a & 0xffff, b & 0xffff, c & 0xffff])[1],
|
||||
VOP3Op.V_MIN3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[0] & 0xffff,
|
||||
VOP3Op.V_MAX3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[2] & 0xffff,
|
||||
VOP3Op.V_MED3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[1] & 0xffff,
|
||||
}
|
||||
|
||||
def _cmp8(a, b): return [False, a < b, a == b, a <= b, a > b, a != b, a >= b, True]
|
||||
def _cmp6(a, b): return [a < b, a == b, a <= b, a > b, a != b, a >= b]
|
||||
|
||||
def vopc(op: int, s0: int, s1: int, s0_hi: int = 0, s1_hi: int = 0) -> int:
|
||||
base = op & 0x7f
|
||||
if 16 <= base <= 31: # F32
|
||||
f0, f1, cmp, nan = f32(s0), f32(s1), base - 16, math.isnan(f32(s0)) or math.isnan(f32(s1))
|
||||
return int([False, f0<f1, f0==f1, f0<=f1, f0>f1, f0!=f1, f0>=f1, not nan, nan, f0<f1 or nan, f0==f1 or nan, f0<=f1 or nan, f0>f1 or nan, f0!=f1 or nan, f0>=f1 or nan, True][cmp])
|
||||
if 49 <= base <= 54: return int(_cmp6(sext(s0 & 0xffff, 16), sext(s1 & 0xffff, 16))[base - 49]) # I16
|
||||
if 57 <= base <= 62: return int(_cmp6(s0 & 0xffff, s1 & 0xffff)[base - 57]) # U16
|
||||
if 64 <= base <= 79: # I32/U32
|
||||
cmp = (base - 64) % 8
|
||||
return int(_cmp8(sext(s0, 32), sext(s1, 32))[cmp] if base < 72 else _cmp8(s0, s1)[cmp])
|
||||
if 80 <= base <= 95: # I64/U64
|
||||
s0_64, s1_64 = s0 | (s0_hi << 32), s1 | (s1_hi << 32)
|
||||
return int(_cmp8(sext(s0_64, 64), sext(s1_64, 64))[(base - 80) % 8] if base < 88 else _cmp8(s0_64, s1_64)[(base - 80) % 8])
|
||||
if base == 126: # CLASS_F32
|
||||
f, mask = f32(s0), s1
|
||||
if math.isnan(f): return int(bool(mask & 0x3))
|
||||
if math.isinf(f): return int(bool(mask & (0x4 if f < 0 else 0x200)))
|
||||
if f == 0.0: return int(bool(mask & (0x20 if (s0 >> 31) & 1 else 0x40)))
|
||||
exp, sign = (s0 >> 23) & 0xff, (s0 >> 31) & 1
|
||||
return int(bool(mask & ((0x10 if sign else 0x80) if exp == 0 else (0x8 if sign else 0x100))))
|
||||
raise NotImplementedError(f"VOPC op {op} (base {base})")
|
||||
@@ -1,602 +0,0 @@
|
||||
# RDNA3 assembler and disassembler
|
||||
from __future__ import annotations
|
||||
import re
|
||||
from extra.assembly.rdna3.lib import Inst, RawImm, Reg, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory, FLOAT_ENC, SRC_FIELDS, unwrap
|
||||
|
||||
# Decoding helpers
|
||||
SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
|
||||
SPECIAL_DEC = {**SPECIAL_GPRS, **{v: str(k) for k, v in FLOAT_ENC.items()}}
|
||||
SPECIAL_PAIRS = {106: "vcc", 126: "exec"} # Special register pairs (for 64-bit ops)
|
||||
# GFX11 hwreg names (IDs 16-17 are TBA - not supported, IDs 18-19 are PERF_SNAPSHOT)
|
||||
HWREG_NAMES = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
|
||||
6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
|
||||
19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI',
|
||||
22: 'HW_REG_XNACK_MASK', 23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
|
||||
HWREG_IDS = {v.lower(): k for k, v in HWREG_NAMES.items()} # Reverse map for assembler
|
||||
MSG_NAMES = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
|
||||
131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA'}
|
||||
_16BIT_TYPES = ('f16', 'i16', 'u16', 'b16')
|
||||
def _is_16bit(s: str) -> bool: return any(s.endswith(x) for x in _16BIT_TYPES)
|
||||
|
||||
def decode_src(val: int) -> str:
|
||||
if val <= 105: return f"s{val}"
|
||||
if val in SPECIAL_DEC: return SPECIAL_DEC[val]
|
||||
if 108 <= val <= 123: return f"ttmp{val - 108}"
|
||||
if 128 <= val <= 192: return str(val - 128)
|
||||
if 193 <= val <= 208: return str(-(val - 192))
|
||||
if 256 <= val <= 511: return f"v{val - 256}"
|
||||
return "lit" if val == 255 else f"?{val}"
|
||||
|
||||
def _reg(prefix: str, base: int, cnt: int = 1) -> str: return f"{prefix}{base}" if cnt == 1 else f"{prefix}[{base}:{base+cnt-1}]"
|
||||
def _sreg(base: int, cnt: int = 1) -> str: return _reg("s", base, cnt)
|
||||
def _vreg(base: int, cnt: int = 1) -> str: return _reg("v", base, cnt)
|
||||
|
||||
def _fmt_sdst(v: int, cnt: int = 1) -> str:
|
||||
"""Format SGPR destination with special register names."""
|
||||
if v == 124: return "null"
|
||||
if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
|
||||
if cnt > 1 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
|
||||
if cnt > 1: return _sreg(v, cnt)
|
||||
return {126: "exec_lo", 127: "exec_hi", 106: "vcc_lo", 107: "vcc_hi", 125: "m0"}.get(v, f"s{v}")
|
||||
|
||||
def _fmt_ssrc(v: int, cnt: int = 1) -> str:
|
||||
"""Format SGPR source with special register names and pairs."""
|
||||
if cnt == 2:
|
||||
if v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
|
||||
if v <= 105: return _sreg(v, 2)
|
||||
if 108 <= v <= 123: return _reg("ttmp", v - 108, 2)
|
||||
return decode_src(v)
|
||||
|
||||
def _fmt_src_n(v: int, cnt: int) -> str:
|
||||
"""Format source with given register count (1, 2, or 4)."""
|
||||
if cnt == 1: return decode_src(v)
|
||||
if v >= 256: return _vreg(v - 256, cnt)
|
||||
if v <= 105: return _sreg(v, cnt)
|
||||
if cnt == 2 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
|
||||
if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
|
||||
return decode_src(v)
|
||||
|
||||
def _fmt_src64(v: int) -> str:
|
||||
"""Format 64-bit source (VGPR pair, SGPR pair, or special pair)."""
|
||||
return _fmt_src_n(v, 2)
|
||||
|
||||
def _parse_sop_sizes(op_name: str) -> tuple[int, ...]:
|
||||
"""Parse dst and src sizes from SOP instruction name. Returns (dst_cnt, src0_cnt) or (dst_cnt, src0_cnt, src1_cnt)."""
|
||||
if op_name in ('s_bitset0_b64', 's_bitset1_b64'): return (2, 1)
|
||||
if op_name in ('s_lshl_b64', 's_lshr_b64', 's_ashr_i64', 's_bfe_u64', 's_bfe_i64'): return (2, 2, 1)
|
||||
if op_name in ('s_bfm_b64',): return (2, 1, 1)
|
||||
# SOPC: s_bitcmp0_b64, s_bitcmp1_b64 - 64-bit src0, 32-bit src1 (bit index)
|
||||
if op_name in ('s_bitcmp0_b64', 's_bitcmp1_b64'): return (1, 2, 1)
|
||||
if m := re.search(r'_(b|i|u)(32|64)_(b|i|u)(32|64)$', op_name):
|
||||
return (2 if m.group(2) == '64' else 1, 2 if m.group(4) == '64' else 1)
|
||||
if m := re.search(r'_(b|i|u)(32|64)$', op_name):
|
||||
sz = 2 if m.group(2) == '64' else 1
|
||||
return (sz, sz)
|
||||
return (1, 1)
|
||||
|
||||
# Waitcnt helpers (RDNA3 format: bits 15:10=vmcnt, bits 9:4=lgkmcnt, bits 3:0=expcnt)
|
||||
def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
|
||||
return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
|
||||
def decode_waitcnt(val: int) -> tuple[int, int, int]:
|
||||
return (val >> 10) & 0x3f, val & 0xf, (val >> 4) & 0x3f # vmcnt, expcnt, lgkmcnt
|
||||
|
||||
# VOP3SD opcodes (shared encoding with VOP3 but different field layout)
|
||||
# Note: opcodes 0-255 are VOPC promoted to VOP3 - never treat as VOP3SD
|
||||
VOP3SD_OPCODES = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
|
||||
|
||||
# Disassembler
|
||||
def disasm(inst: Inst) -> str:
|
||||
op_val = unwrap(inst._values.get('op', 0))
|
||||
cls_name = inst.__class__.__name__
|
||||
# VOP3 and VOP3SD share encoding - check opcode to determine which
|
||||
is_vop3sd = cls_name == 'VOP3' and op_val in VOP3SD_OPCODES
|
||||
try:
|
||||
from extra.assembly.rdna3 import autogen
|
||||
if is_vop3sd:
|
||||
op_name = autogen.VOP3SDOp(op_val).name.lower()
|
||||
else:
|
||||
op_name = getattr(autogen, f"{cls_name}Op")(op_val).name.lower() if hasattr(autogen, f"{cls_name}Op") else f"op_{op_val}"
|
||||
except (ValueError, KeyError): op_name = f"op_{op_val}"
|
||||
def fmt_src(v): return f"0x{inst._literal:x}" if v == 255 and getattr(inst, '_literal', None) else decode_src(v)
|
||||
|
||||
# VOP1
|
||||
if cls_name == 'VOP1':
|
||||
vdst, src0 = unwrap(inst._values['vdst']), unwrap(inst._values['src0'])
|
||||
if op_name == 'v_nop': return 'v_nop'
|
||||
if op_name == 'v_pipeflush': return 'v_pipeflush'
|
||||
parts = op_name.split('_')
|
||||
is_16bit_dst = any(p in _16BIT_TYPES for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in _16BIT_TYPES and 'cvt' not in op_name)
|
||||
is_16bit_src = parts[-1] in _16BIT_TYPES and 'sat_pk' not in op_name
|
||||
_F64_OPS = ('v_ceil_f64', 'v_floor_f64', 'v_fract_f64', 'v_frexp_mant_f64', 'v_rcp_f64', 'v_rndne_f64', 'v_rsq_f64', 'v_sqrt_f64', 'v_trunc_f64')
|
||||
is_f64_dst = op_name in _F64_OPS or op_name in ('v_cvt_f64_f32', 'v_cvt_f64_i32', 'v_cvt_f64_u32')
|
||||
is_f64_src = op_name in _F64_OPS or op_name in ('v_cvt_f32_f64', 'v_cvt_i32_f64', 'v_cvt_u32_f64', 'v_frexp_exp_i32_f64')
|
||||
if op_name == 'v_readfirstlane_b32':
|
||||
return f"v_readfirstlane_b32 {decode_src(vdst)}, v{src0 - 256 if src0 >= 256 else src0}"
|
||||
dst_str = _vreg(vdst, 2) if is_f64_dst else f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}" if is_16bit_dst else f"v{vdst}"
|
||||
src_str = _fmt_src64(src0) if is_f64_src else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit_src and src0 >= 256 else fmt_src(src0)
|
||||
return f"{op_name}_e32 {dst_str}, {src_str}"
|
||||
|
||||
# VOP2
|
||||
if cls_name == 'VOP2':
|
||||
vdst, src0_raw, vsrc1 = unwrap(inst._values['vdst']), unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
|
||||
suffix = "" if op_name == "v_dot2acc_f32_f16" else "_e32"
|
||||
is_16bit_op = ('_f16' in op_name or '_i16' in op_name or '_u16' in op_name) and '_f32' not in op_name and '_i32' not in op_name and 'pk_' not in op_name
|
||||
if is_16bit_op:
|
||||
dst_str = f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}"
|
||||
src0_str = f"v{(src0_raw - 256) & 0x7f}.{'h' if src0_raw >= 384 else 'l'}" if src0_raw >= 256 else fmt_src(src0_raw)
|
||||
vsrc1_str = f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}"
|
||||
else:
|
||||
dst_str, src0_str, vsrc1_str = f"v{vdst}", fmt_src(src0_raw), f"v{vsrc1}"
|
||||
return f"{op_name}{suffix} {dst_str}, {src0_str}, {vsrc1_str}" + (", vcc_lo" if op_name == "v_cndmask_b32" else "")
|
||||
|
||||
# VOPC
|
||||
if cls_name == 'VOPC':
|
||||
src0, vsrc1 = unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
|
||||
is_64bit = any(x in op_name for x in ('f64', 'i64', 'u64'))
|
||||
is_64bit_vsrc1 = is_64bit and 'class' not in op_name
|
||||
is_16bit = any(x in op_name for x in ('_f16', '_i16', '_u16')) and 'f32' not in op_name
|
||||
is_cmpx = op_name.startswith('v_cmpx') # VOPCX writes to exec, no vcc destination
|
||||
src0_str = _fmt_src64(src0) if is_64bit else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit and src0 >= 256 else fmt_src(src0)
|
||||
vsrc1_str = _vreg(vsrc1, 2) if is_64bit_vsrc1 else f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}" if is_16bit else f"v{vsrc1}"
|
||||
return f"{op_name}_e32 {src0_str}, {vsrc1_str}" if is_cmpx else f"{op_name}_e32 vcc_lo, {src0_str}, {vsrc1_str}"
|
||||
|
||||
# SOPP
|
||||
if cls_name == 'SOPP':
|
||||
simm16 = unwrap(inst._values.get('simm16', 0))
|
||||
# No-operand instructions (simm16 is ignored)
|
||||
no_imm_ops = ('s_endpgm', 's_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_ttracedata_imm',
|
||||
's_wait_idle', 's_endpgm_saved', 's_code_end', 's_endpgm_ordered_ps_done')
|
||||
if op_name in no_imm_ops: return op_name
|
||||
if op_name == 's_waitcnt':
|
||||
vmcnt, expcnt, lgkmcnt = decode_waitcnt(simm16)
|
||||
parts = []
|
||||
if vmcnt != 0x3f: parts.append(f"vmcnt({vmcnt})")
|
||||
if expcnt != 0x7: parts.append(f"expcnt({expcnt})")
|
||||
if lgkmcnt != 0x3f: parts.append(f"lgkmcnt({lgkmcnt})")
|
||||
return f"s_waitcnt {' '.join(parts)}" if parts else "s_waitcnt 0"
|
||||
if op_name == 's_delay_alu':
|
||||
dep_names = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
|
||||
skip_names = ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
|
||||
id0, skip, id1 = simm16 & 0xf, (simm16 >> 4) & 0x7, (simm16 >> 7) & 0xf
|
||||
def dep_name(v): return dep_names[v-1] if 0 < v <= len(dep_names) else str(v)
|
||||
parts = [f"instid0({dep_name(id0)})"] if id0 else []
|
||||
if skip: parts.append(f"instskip({skip_names[skip]})")
|
||||
if id1: parts.append(f"instid1({dep_name(id1)})")
|
||||
return f"s_delay_alu {' | '.join(p for p in parts if p)}" if parts else "s_delay_alu 0"
|
||||
if op_name.startswith('s_cbranch') or op_name.startswith('s_branch'):
|
||||
return f"{op_name} {simm16}"
|
||||
# Most SOPP ops require immediate (s_nop, s_setkill, s_sethalt, s_sleep, s_setprio, s_sendmsg*, etc.)
|
||||
return f"{op_name} 0x{simm16:x}"
|
||||
|
||||
# SMEM
|
||||
if cls_name == 'SMEM':
|
||||
if op_name in ('s_gl1_inv', 's_dcache_inv'): return op_name
|
||||
sdata, sbase, soffset, offset = unwrap(inst._values['sdata']), unwrap(inst._values['sbase']), unwrap(inst._values['soffset']), unwrap(inst._values.get('offset', 0))
|
||||
glc, dlc = unwrap(inst._values.get('glc', 0)), unwrap(inst._values.get('dlc', 0))
|
||||
# Format offset: "soffset offset:X" if both, "0x{offset:x}" if only imm, or decode_src(soffset)
|
||||
off_str = f"{decode_src(soffset)} offset:0x{offset:x}" if offset and soffset != 124 else f"0x{offset:x}" if offset else decode_src(soffset)
|
||||
sbase_idx, sbase_cnt = sbase * 2, 4 if (8 <= op_val <= 12 or op_name == 's_atc_probe_buffer') else 2
|
||||
sbase_str = _fmt_ssrc(sbase_idx, sbase_cnt) if sbase_cnt == 2 else _sreg(sbase_idx, sbase_cnt) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_cnt)
|
||||
if op_name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{op_name} {sdata}, {sbase_str}, {off_str}"
|
||||
width = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val, 1)
|
||||
mods = [m for m in ["glc" if glc else "", "dlc" if dlc else ""] if m]
|
||||
return f"{op_name} {_fmt_sdst(sdata, width)}, {sbase_str}, {off_str}" + (" " + " ".join(mods) if mods else "")
|
||||
|
||||
# FLAT
|
||||
if cls_name == 'FLAT':
|
||||
vdst, addr, data, saddr, offset, seg = [unwrap(inst._values.get(f, 0)) for f in ['vdst', 'addr', 'data', 'saddr', 'offset', 'seg']]
|
||||
instr = f"{['flat', 'scratch', 'global'][seg] if seg < 3 else 'flat'}_{op_name.split('_', 1)[1] if '_' in op_name else op_name}"
|
||||
width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'u8':1, 'i8':1, 'u16':1, 'i16':1}.get(op_name.split('_')[-1], 1)
|
||||
addr_str = _vreg(addr, 2) if saddr == 0x7F else _vreg(addr)
|
||||
saddr_str = "" if saddr == 0x7F else f", {_sreg(saddr, 2)}" if saddr < 106 else ", off" if saddr == 124 else f", {decode_src(saddr)}"
|
||||
off_str = f" offset:{offset}" if offset else ""
|
||||
vdata_str = _vreg(data if 'store' in op_name else vdst, width)
|
||||
return f"{instr} {addr_str}, {vdata_str}{saddr_str}{off_str}" if 'store' in op_name else f"{instr} {vdata_str}, {addr_str}{saddr_str}{off_str}"
|
||||
|
||||
# VOP3: vector ops with modifiers (can be 1, 2, or 3 sources depending on opcode range)
|
||||
if cls_name == 'VOP3':
|
||||
# Handle VOP3SD opcodes (same encoding, different field layout)
|
||||
if is_vop3sd:
|
||||
vdst = unwrap(inst._values.get('vdst', 0))
|
||||
# VOP3SD: sdst is at bits [14:8], but VOP3 decodes opsel at [14:11], abs at [10:8], clmp at [15]
|
||||
# We need to reconstruct sdst from these fields
|
||||
opsel_raw = unwrap(inst._values.get('opsel', 0))
|
||||
abs_raw = unwrap(inst._values.get('abs', 0))
|
||||
clmp_raw = unwrap(inst._values.get('clmp', 0))
|
||||
sdst = (clmp_raw << 7) | (opsel_raw << 3) | abs_raw
|
||||
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
|
||||
neg = unwrap(inst._values.get('neg', 0))
|
||||
omod = unwrap(inst._values.get('omod', 0))
|
||||
omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
|
||||
is_f64 = 'f64' in op_name
|
||||
# v_mad_i64_i32/v_mad_u64_u32: 64-bit dst and src2, 32-bit src0/src1
|
||||
is_mad64 = 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
|
||||
def fmt_sd_src(v, neg_bit, is_64bit=False):
|
||||
s = _fmt_src64(v) if (is_64bit or is_f64) else fmt_src(v)
|
||||
return f"-{s}" if neg_bit else s
|
||||
src0_str, src1_str = fmt_sd_src(src0, neg & 1), fmt_sd_src(src1, neg & 2)
|
||||
src2_str = fmt_sd_src(src2, neg & 4, is_mad64)
|
||||
dst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}"
|
||||
sdst_str = _fmt_sdst(sdst, 1)
|
||||
# v_add_co_u32, v_sub_co_u32, v_subrev_co_u32, v_add_co_ci_u32, etc. only use 2 sources
|
||||
if op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32', 'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'):
|
||||
return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}"
|
||||
# v_div_scale uses 3 sources
|
||||
return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}, {src2_str}" + omod_str
|
||||
|
||||
vdst = unwrap(inst._values.get('vdst', 0))
|
||||
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
|
||||
neg, abs_, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('abs', 0)), unwrap(inst._values.get('clmp', 0))
|
||||
opsel = unwrap(inst._values.get('opsel', 0))
|
||||
# Check if 64-bit op (needs register pairs)
|
||||
is_f64 = 'f64' in op_name or 'i64' in op_name or 'u64' in op_name or 'b64' in op_name
|
||||
# v_cmp_class_* has 64-bit src0 but 32-bit src1 (class mask)
|
||||
is_class = 'class' in op_name
|
||||
# Shift ops: v_*rev_*64 have 32-bit shift amount (src0), 64-bit value (src1)
|
||||
is_shift64 = 'rev' in op_name and '64' in op_name and op_name.startswith('v_')
|
||||
# v_ldexp_f64: 64-bit src0 (mantissa), 32-bit src1 (exponent)
|
||||
is_ldexp64 = op_name == 'v_ldexp_f64'
|
||||
# v_trig_preop_f64: 64-bit dst/src0, 32-bit src1 (exponent/scale)
|
||||
is_trig_preop = op_name == 'v_trig_preop_f64'
|
||||
# v_readlane_b32: destination is SGPR (despite vdst field)
|
||||
is_readlane = op_name == 'v_readlane_b32'
|
||||
# SAD/QSAD/MQSAD instructions have mixed sizes
|
||||
# v_qsad_pk_u16_u8, v_mqsad_pk_u16_u8: 64-bit dst/src0/src2, 32-bit src1
|
||||
# v_mqsad_u32_u8: 128-bit (4 reg) dst/src2, 64-bit src0, 32-bit src1
|
||||
is_sad64 = any(x in op_name for x in ('qsad_pk', 'mqsad_pk'))
|
||||
is_mqsad_u32 = 'mqsad_u32' in op_name
|
||||
# Detect 16-bit and 64-bit operand sizes for various instruction patterns
|
||||
if 'cvt_pk' in op_name:
|
||||
is_f16_dst, is_f16_src, is_f16_src2 = False, op_name.endswith('16'), False
|
||||
elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', op_name):
|
||||
dst_type, src_type = m.group(1), m.group(2)
|
||||
is_f16_dst, is_f16_src, is_f16_src2 = _is_16bit(dst_type), _is_16bit(src_type), _is_16bit(src_type)
|
||||
is_f64_dst, is_f64_src, is_f64 = '64' in dst_type, '64' in src_type, False
|
||||
elif re.match(r'v_mad_[iu]32_[iu]16', op_name):
|
||||
is_f16_dst, is_f16_src, is_f16_src2 = False, True, False # 32-bit dst, 16-bit src0/src1, 32-bit src2
|
||||
elif 'pack_b32' in op_name:
|
||||
is_f16_dst, is_f16_src, is_f16_src2 = False, True, True # 32-bit dst, 16-bit sources
|
||||
else:
|
||||
is_16bit_op = any(x in op_name for x in _16BIT_TYPES) and not any(x in op_name for x in ('dot2', 'pk_', 'sad', 'msad', 'qsad', 'mqsad'))
|
||||
is_f16_dst = is_f16_src = is_f16_src2 = is_16bit_op
|
||||
# Check if any opsel bit is set (any operand uses .h) - if so, we need explicit .l for low-half
|
||||
any_hi = opsel != 0
|
||||
def fmt_vop3_src(v, neg_bit, abs_bit, hi_bit=False, reg_cnt=1, is_16=False):
|
||||
s = _fmt_src_n(v, reg_cnt) if reg_cnt > 1 else f"v{v - 256}.h" if is_16 and v >= 256 and hi_bit else f"v{v - 256}.l" if is_16 and v >= 256 and any_hi else fmt_src(v)
|
||||
if abs_bit: s = f"|{s}|"
|
||||
return f"-{s}" if neg_bit else s
|
||||
# Determine register count for each source (check for cvt-specific 64-bit flags first)
|
||||
is_src0_64 = locals().get('is_f64_src', is_f64 and not is_shift64) or is_sad64 or is_mqsad_u32
|
||||
is_src1_64 = is_f64 and not is_class and not is_ldexp64 and not is_trig_preop
|
||||
src0_cnt = 2 if is_src0_64 else 1
|
||||
src1_cnt = 2 if is_src1_64 else 1
|
||||
src2_cnt = 4 if is_mqsad_u32 else 2 if (is_f64 or is_sad64) else 1
|
||||
src0_str = fmt_vop3_src(src0, neg & 1, abs_ & 1, opsel & 1, src0_cnt, is_f16_src)
|
||||
src1_str = fmt_vop3_src(src1, neg & 2, abs_ & 2, opsel & 2, src1_cnt, is_f16_src)
|
||||
src2_str = fmt_vop3_src(src2, neg & 4, abs_ & 4, opsel & 4, src2_cnt, is_f16_src2)
|
||||
# Format destination - for 16-bit ops, use .h/.l suffix; readlane uses SGPR dest
|
||||
is_dst_64 = locals().get('is_f64_dst', is_f64) or is_sad64
|
||||
dst_cnt = 4 if is_mqsad_u32 else 2 if is_dst_64 else 1
|
||||
if is_readlane:
|
||||
dst_str = _fmt_sdst(vdst, 1)
|
||||
elif dst_cnt > 1:
|
||||
dst_str = _vreg(vdst, dst_cnt)
|
||||
elif is_f16_dst:
|
||||
dst_str = f"v{vdst}.h" if (opsel & 8) else f"v{vdst}.l" if any_hi else f"v{vdst}"
|
||||
else:
|
||||
dst_str = f"v{vdst}"
|
||||
clamp_str = " clamp" if clmp else ""
|
||||
omod = unwrap(inst._values.get('omod', 0))
|
||||
omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
|
||||
# op_sel for non-VGPR sources (when opsel bits are set but source is not a VGPR)
|
||||
# For 16-bit ops with VGPR sources, opsel is encoded in .h/.l suffix
|
||||
# For non-VGPR sources or non-16-bit ops, we need explicit op_sel
|
||||
has_nonvgpr_opsel = (src0 < 256 and (opsel & 1)) or (src1 < 256 and (opsel & 2)) or (src2 < 256 and (opsel & 4))
|
||||
need_opsel = has_nonvgpr_opsel or (opsel and not is_f16_src)
|
||||
# Helper to format opsel string based on source count
|
||||
def fmt_opsel(num_src):
|
||||
if not need_opsel: return ""
|
||||
# When dst is .h (for 16-bit ops) and non-VGPR sources have opsel, use all 1s
|
||||
if is_f16_dst and (opsel & 8): # dst is .h
|
||||
return f" op_sel:[1,1,1{',1' if num_src == 3 else ''}]"
|
||||
# Otherwise output actual opsel values
|
||||
if num_src == 3:
|
||||
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{(opsel >> 3) & 1}]"
|
||||
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1}]"
|
||||
# Determine number of sources based on opcode range:
|
||||
# 0-255: VOPC promoted (comparison, 2 src, sdst)
|
||||
# 256-383: VOP2 promoted (2 src)
|
||||
# 384-511: VOP1 promoted (1 src)
|
||||
# 512+: Native VOP3 (2 or 3 src depending on instruction)
|
||||
if op_val < 256: # VOPC promoted
|
||||
# VOPCX (v_cmpx_*) writes to exec, no explicit destination
|
||||
if op_name.startswith('v_cmpx'):
|
||||
return f"{op_name}_e64 {src0_str}, {src1_str}"
|
||||
return f"{op_name}_e64 {_fmt_sdst(vdst, 1)}, {src0_str}, {src1_str}"
|
||||
elif op_val < 384: # VOP2 promoted
|
||||
# v_cndmask_b32 in VOP3 format has 3 sources (src2 is mask selector)
|
||||
if 'cndmask' in op_name:
|
||||
return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
|
||||
return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
|
||||
elif op_val < 512: # VOP1 promoted
|
||||
if op_name in ('v_nop', 'v_pipeflush'): return f"{op_name}_e64"
|
||||
return f"{op_name}_e64 {dst_str}, {src0_str}" + fmt_opsel(1) + clamp_str + omod_str
|
||||
else: # Native VOP3 - determine 2 vs 3 sources based on instruction name
|
||||
# 3-source ops: fma, mad, min3, max3, med3, div_fixup, div_fmas, sad, msad, qsad, mqsad, lerp, alignbit/byte, cubeid/sc/tc/ma, bfe, bfi, perm_b32, permlane, cndmask
|
||||
# Note: v_writelane_b32 is 2-src (src0, src1 with vdst as 3rd operand - read-modify-write)
|
||||
is_3src = any(x in op_name for x in ('fma', 'mad', 'min3', 'max3', 'med3', 'div_fix', 'div_fmas', 'sad', 'lerp', 'align', 'cube',
|
||||
'bfe', 'bfi', 'perm_b32', 'permlane', 'cndmask', 'xor3', 'or3', 'add3', 'lshl_or', 'and_or', 'lshl_add',
|
||||
'add_lshl', 'xad', 'maxmin', 'minmax', 'dot2', 'cvt_pk_u8', 'mullit'))
|
||||
if is_3src:
|
||||
return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
|
||||
return f"{op_name} {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
|
||||
|
||||
# VOP3SD: 3-source with scalar destination (v_div_scale_*, v_add_co_u32, v_mad_*64_*32, etc.)
|
||||
if cls_name == 'VOP3SD':
|
||||
vdst, sdst = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('sdst', 0))
|
||||
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
|
||||
neg, omod, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('omod', 0)), unwrap(inst._values.get('clmp', 0))
|
||||
is_f64, is_mad64 = 'f64' in op_name, 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
|
||||
def fmt_neg(v, neg_bit, is_64=False): return f"-{_fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)}" if neg_bit else _fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)
|
||||
srcs = [fmt_neg(src0, neg & 1), fmt_neg(src1, neg & 2), fmt_neg(src2, neg & 4, is_mad64)]
|
||||
dst_str, sdst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}", _fmt_sdst(sdst, 1)
|
||||
clamp_str, omod_str = " clamp" if clmp else "", {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
|
||||
is_2src = op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32')
|
||||
suffix = "_e64" if op_name.startswith('v_') and 'co_' in op_name else ""
|
||||
return f"{op_name}{suffix} {dst_str}, {sdst_str}, {', '.join(srcs[:2] if is_2src else srcs)}" + clamp_str + omod_str
|
||||
|
||||
# VOPD: dual-issue instructions
|
||||
if cls_name == 'VOPD':
|
||||
from extra.assembly.rdna3 import autogen
|
||||
opx, opy, vdstx, vdsty_enc = [unwrap(inst._values.get(f, 0)) for f in ('opx', 'opy', 'vdstx', 'vdsty')]
|
||||
srcx0, vsrcx1, srcy0, vsrcy1 = [unwrap(inst._values.get(f, 0)) for f in ('srcx0', 'vsrcx1', 'srcy0', 'vsrcy1')]
|
||||
vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1) # Decode vdsty
|
||||
def fmt_vopd(op, vdst, src0, vsrc1):
|
||||
try: name = autogen.VOPDOp(op).name.lower()
|
||||
except (ValueError, KeyError): name = f"op_{op}"
|
||||
return f"{name} v{vdst}, {fmt_src(src0)}" if 'mov' in name else f"{name} v{vdst}, {fmt_src(src0)}, v{vsrc1}"
|
||||
return f"{fmt_vopd(opx, vdstx, srcx0, vsrcx1)} :: {fmt_vopd(opy, vdsty, srcy0, vsrcy1)}"
|
||||
|
||||
# VOP3P: packed vector ops
|
||||
if cls_name == 'VOP3P':
|
||||
vdst, clmp = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('clmp', 0))
|
||||
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
|
||||
neg, neg_hi = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('neg_hi', 0))
|
||||
opsel, opsel_hi, opsel_hi2 = unwrap(inst._values.get('opsel', 0)), unwrap(inst._values.get('opsel_hi', 0)), unwrap(inst._values.get('opsel_hi2', 0))
|
||||
is_wmma, is_3src = 'wmma' in op_name, any(x in op_name for x in ('fma', 'mad', 'dot', 'wmma'))
|
||||
def fmt_bits(name, val, n): return f"{name}:[{','.join(str((val >> i) & 1) for i in range(n))}]"
|
||||
# WMMA: f16/bf16 use 8-reg sources, iu8 uses 4-reg, iu4 uses 2-reg; all have 8-reg dst
|
||||
if is_wmma:
|
||||
src_cnt = 2 if 'iu4' in op_name else 4 if 'iu8' in op_name else 8
|
||||
src0_str, src1_str, src2_str = _fmt_src_n(src0, src_cnt), _fmt_src_n(src1, src_cnt), _fmt_src_n(src2, 8)
|
||||
dst_str = _vreg(vdst, 8)
|
||||
else:
|
||||
src0_str, src1_str, src2_str = _fmt_src_n(src0, 1), _fmt_src_n(src1, 1), _fmt_src_n(src2, 1)
|
||||
dst_str = f"v{vdst}"
|
||||
n = 3 if is_3src else 2
|
||||
full_opsel_hi = opsel_hi | (opsel_hi2 << 2)
|
||||
mods = [fmt_bits("op_sel", opsel, n)] if opsel else []
|
||||
if full_opsel_hi != (0b111 if is_3src else 0b11): mods.append(fmt_bits("op_sel_hi", full_opsel_hi, n))
|
||||
if neg: mods.append(fmt_bits("neg_lo", neg, n))
|
||||
if neg_hi: mods.append(fmt_bits("neg_hi", neg_hi, n))
|
||||
if clmp: mods.append("clamp")
|
||||
mod_str = " " + " ".join(mods) if mods else ""
|
||||
return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}{mod_str}" if is_3src else f"{op_name} {dst_str}, {src0_str}, {src1_str}{mod_str}"
|
||||
|
||||
# VINTERP: interpolation instructions
|
||||
if cls_name == 'VINTERP':
|
||||
vdst = unwrap(inst._values.get('vdst', 0))
|
||||
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
|
||||
neg, waitexp, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('waitexp', 0)), unwrap(inst._values.get('clmp', 0))
|
||||
def fmt_neg_vi(v, neg_bit): return f"-{v}" if neg_bit else v
|
||||
srcs = [fmt_neg_vi(f"v{s - 256}" if s >= 256 else fmt_src(s), neg & (1 << i)) for i, s in enumerate([src0, src1, src2])]
|
||||
mods = [m for m in [f"wait_exp:{waitexp}" if waitexp else "", "clamp" if clmp else ""] if m]
|
||||
return f"{op_name} v{vdst}, {', '.join(srcs)}" + (" " + " ".join(mods) if mods else "")
|
||||
|
||||
# MUBUF/MTBUF helpers
|
||||
def _buf_vaddr(vaddr, offen, idxen): return _vreg(vaddr, 2) if offen and idxen else f"v{vaddr}" if offen or idxen else "off"
|
||||
def _buf_srsrc(srsrc): srsrc_base = srsrc * 4; return _reg("ttmp", srsrc_base - 108, 4) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 4)
|
||||
|
||||
# MUBUF: buffer load/store
|
||||
if cls_name == 'MUBUF':
|
||||
vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
|
||||
offset, offen, idxen = unwrap(inst._values.get('offset', 0)), unwrap(inst._values.get('offen', 0)), unwrap(inst._values.get('idxen', 0))
|
||||
glc, dlc, slc, tfe = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc', 'tfe')]
|
||||
if op_name in ('buffer_gl0_inv', 'buffer_gl1_inv'): return op_name
|
||||
# Determine data width from op name
|
||||
if 'd16' in op_name: width = 2 if any(x in op_name for x in ('xyz', 'xyzw')) else 1
|
||||
elif 'atomic' in op_name:
|
||||
base_width = 2 if any(x in op_name for x in ('b64', 'u64', 'i64')) else 1
|
||||
width = base_width * 2 if 'cmpswap' in op_name else base_width
|
||||
else: width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'b16':1, 'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
|
||||
if tfe: width += 1
|
||||
mods = [m for m in ["offen" if offen else "", "idxen" if idxen else "", f"offset:{offset}" if offset else "",
|
||||
"glc" if glc else "", "dlc" if dlc else "", "slc" if slc else "", "tfe" if tfe else ""] if m]
|
||||
return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)}" + (" " + " ".join(mods) if mods else "")
|
||||
|
||||
# MTBUF: typed buffer load/store
|
||||
if cls_name == 'MTBUF':
|
||||
vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
|
||||
offset, tbuf_fmt, offen, idxen = [unwrap(inst._values.get(f, 0)) for f in ('offset', 'format', 'offen', 'idxen')]
|
||||
glc, dlc, slc = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc')]
|
||||
mods = [f"format:{tbuf_fmt}"] + [m for m in ["idxen" if idxen else "", "offen" if offen else "", f"offset:{offset}" if offset else "",
|
||||
"glc" if glc else "", "dlc" if dlc else "", "slc" if slc else ""] if m]
|
||||
width = 2 if 'd16' in op_name and any(x in op_name for x in ('xyz', 'xyzw')) else 1 if 'd16' in op_name else {'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
|
||||
return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)} {' '.join(mods)}"
|
||||
|
||||
# SOP1/SOP2/SOPC/SOPK
|
||||
if cls_name in ('SOP1', 'SOP2', 'SOPC', 'SOPK'):
|
||||
sizes = _parse_sop_sizes(op_name)
|
||||
dst_cnt, src0_cnt = sizes[0], sizes[1]
|
||||
src1_cnt = sizes[2] if len(sizes) > 2 else src0_cnt
|
||||
if cls_name == 'SOP1':
|
||||
sdst, ssrc0 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('ssrc0', 0))
|
||||
if op_name == 's_getpc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}"
|
||||
if op_name in ('s_setpc_b64', 's_rfe_b64'): return f"{op_name} {_fmt_ssrc(ssrc0, 2)}"
|
||||
if op_name == 's_swappc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}, {_fmt_ssrc(ssrc0, 2)}"
|
||||
if op_name in ('s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'):
|
||||
return f"{op_name} {_fmt_sdst(sdst, 2 if 'b64' in op_name else 1)}, sendmsg({MSG_NAMES.get(ssrc0, str(ssrc0))})"
|
||||
ssrc0_str = fmt_src(ssrc0) if src0_cnt == 1 else _fmt_ssrc(ssrc0, src0_cnt)
|
||||
return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {ssrc0_str}"
|
||||
if cls_name == 'SOP2':
|
||||
sdst, ssrc0, ssrc1 = [unwrap(inst._values.get(f, 0)) for f in ('sdst', 'ssrc0', 'ssrc1')]
|
||||
return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {_fmt_ssrc(ssrc0, src0_cnt)}, {_fmt_ssrc(ssrc1, src1_cnt)}"
|
||||
if cls_name == 'SOPC':
|
||||
return f"{op_name} {_fmt_ssrc(unwrap(inst._values.get('ssrc0', 0)), src0_cnt)}, {_fmt_ssrc(unwrap(inst._values.get('ssrc1', 0)), src1_cnt)}"
|
||||
if cls_name == 'SOPK':
|
||||
sdst, simm16 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('simm16', 0))
|
||||
if op_name == 's_version': return f"{op_name} 0x{simm16:x}"
|
||||
if op_name in ('s_setreg_b32', 's_getreg_b32'):
|
||||
hwreg_id, hwreg_offset, hwreg_size = simm16 & 0x3f, (simm16 >> 6) & 0x1f, ((simm16 >> 11) & 0x1f) + 1
|
||||
hwreg_str = f"0x{simm16:x}" if hwreg_id in (16, 17) else f"hwreg({HWREG_NAMES.get(hwreg_id, str(hwreg_id))}, {hwreg_offset}, {hwreg_size})"
|
||||
return f"{op_name} {hwreg_str}, {_fmt_sdst(sdst, 1)}" if op_name == 's_setreg_b32' else f"{op_name} {_fmt_sdst(sdst, 1)}, {hwreg_str}"
|
||||
return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, 0x{simm16:x}"
|
||||
|
||||
# Generic fallback
|
||||
def fmt_field(n, v):
|
||||
v = unwrap(v)
|
||||
if n in SRC_FIELDS: return fmt_src(v) if v != 255 else "0xff"
|
||||
if n in ('sdst', 'vdst'): return f"{'s' if n == 'sdst' else 'v'}{v}"
|
||||
return f"v{v}" if n == 'vsrc1' else f"0x{v:x}" if n == 'simm16' else str(v)
|
||||
ops = [fmt_field(n, inst._values.get(n, 0)) for n in inst._fields if n not in ('encoding', 'op')]
|
||||
return f"{op_name} {', '.join(ops)}" if ops else op_name
|
||||
|
||||
# Assembler
|
||||
SPECIAL_REGS = {'vcc_lo': RawImm(106), 'vcc_hi': RawImm(107), 'null': RawImm(124), 'off': RawImm(124), 'm0': RawImm(125), 'exec_lo': RawImm(126), 'exec_hi': RawImm(127), 'scc': RawImm(253)}
|
||||
FLOAT_CONSTS = {'0.5': 0.5, '-0.5': -0.5, '1.0': 1.0, '-1.0': -1.0, '2.0': 2.0, '-2.0': -2.0, '4.0': 4.0, '-4.0': -4.0}
|
||||
REG_MAP: dict[str, _RegFactory] = {'s': s, 'v': v, 't': ttmp, 'ttmp': ttmp}
|
||||
|
||||
def parse_operand(op: str) -> tuple:
|
||||
op = op.strip().lower()
|
||||
neg = op.startswith('-') and not op[1:2].isdigit(); op = op[1:] if neg else op
|
||||
abs_ = op.startswith('|') and op.endswith('|') or op.startswith('abs(') and op.endswith(')')
|
||||
op = op[1:-1] if op.startswith('|') else op[4:-1] if op.startswith('abs(') else op
|
||||
hi_half = op.endswith('.h')
|
||||
op = re.sub(r'\.[lh]$', '', op)
|
||||
if op in FLOAT_CONSTS: return (FLOAT_CONSTS[op], neg, abs_, hi_half)
|
||||
if re.match(r'^-?\d+$', op): return (int(op), neg, abs_, hi_half)
|
||||
if m := re.match(r'^-?0x([0-9a-f]+)$', op):
|
||||
v = -int(m.group(1), 16) if op.startswith('-') else int(m.group(1), 16)
|
||||
return (v, neg, abs_, hi_half)
|
||||
if op in SPECIAL_REGS: return (SPECIAL_REGS[op], neg, abs_, hi_half)
|
||||
if op == 'lit': return (RawImm(255), neg, abs_, hi_half) # literal marker (actual value comes from literal word)
|
||||
if m := re.match(r'^([svt](?:tmp)?)\[(\d+):(\d+)\]$', op): return (REG_MAP[m.group(1)][int(m.group(2)):int(m.group(3))], neg, abs_, hi_half)
|
||||
if m := re.match(r'^([svt](?:tmp)?)(\d+)$', op):
|
||||
reg = REG_MAP[m.group(1)][int(m.group(2))]
|
||||
reg.hi = hi_half
|
||||
return (reg, neg, abs_, hi_half)
|
||||
# hwreg(name, offset, size) or hwreg(name) -> simm16 encoding
|
||||
if m := re.match(r'^hwreg\((\w+)(?:,\s*(\d+),\s*(\d+))?\)$', op):
|
||||
name_str = m.group(1).lower()
|
||||
hwreg_id = HWREG_IDS.get(name_str, int(name_str) if name_str.isdigit() else None)
|
||||
if hwreg_id is None: raise ValueError(f"unknown hwreg name: {name_str}")
|
||||
offset, size = int(m.group(2)) if m.group(2) else 0, int(m.group(3)) if m.group(3) else 32
|
||||
return (((size - 1) << 11) | (offset << 6) | hwreg_id, neg, abs_, hi_half)
|
||||
raise ValueError(f"cannot parse operand: {op}")
|
||||
|
||||
SMEM_OPS = {'s_load_b32', 's_load_b64', 's_load_b128', 's_load_b256', 's_load_b512',
|
||||
's_buffer_load_b32', 's_buffer_load_b64', 's_buffer_load_b128', 's_buffer_load_b256', 's_buffer_load_b512'}
|
||||
SOP1_SRC_ONLY = {'s_setpc_b64', 's_rfe_b64'}
|
||||
SOP1_MSG_IMM = {'s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'}
|
||||
SOPK_IMM_ONLY = {'s_version'}
|
||||
SOPK_IMM_FIRST = {'s_setreg_b32'}
|
||||
SOPK_UNSUPPORTED = {'s_setreg_imm32_b32'}
|
||||
|
||||
def asm(text: str) -> Inst:
|
||||
from extra.assembly.rdna3 import autogen
|
||||
text = text.strip()
|
||||
clamp = 'clamp' in text.lower()
|
||||
if clamp: text = re.sub(r'\s+clamp\s*$', '', text, flags=re.I)
|
||||
modifiers = {}
|
||||
if m := re.search(r'\s+wait_exp:(\d+)', text, re.I): modifiers['waitexp'] = int(m.group(1)); text = text[:m.start()] + text[m.end():]
|
||||
parts = text.replace(',', ' ').split()
|
||||
if not parts: raise ValueError("empty instruction")
|
||||
mnemonic, op_str = parts[0].lower(), text[len(parts[0]):].strip()
|
||||
# Handle s_waitcnt specially before operand parsing
|
||||
if mnemonic == 's_waitcnt':
|
||||
vmcnt, expcnt, lgkmcnt = 0x3f, 0x7, 0x3f
|
||||
for part in op_str.replace(',', ' ').split():
|
||||
if m := re.match(r'vmcnt\((\d+)\)', part): vmcnt = int(m.group(1))
|
||||
elif m := re.match(r'expcnt\((\d+)\)', part): expcnt = int(m.group(1))
|
||||
elif m := re.match(r'lgkmcnt\((\d+)\)', part): lgkmcnt = int(m.group(1))
|
||||
elif re.match(r'^0x[0-9a-f]+$|^\d+$', part): return autogen.s_waitcnt(simm16=int(part, 0))
|
||||
return autogen.s_waitcnt(simm16=waitcnt(vmcnt, expcnt, lgkmcnt))
|
||||
# Handle VOPD dual-issue instructions: opx dst, src :: opy dst, src
|
||||
if '::' in text:
|
||||
x_part, y_part = text.split('::')
|
||||
x_parts, y_parts = x_part.strip().replace(',', ' ').split(), y_part.strip().replace(',', ' ').split()
|
||||
opx_name, opy_name = x_parts[0].upper(), y_parts[0].upper()
|
||||
opx, opy = autogen.VOPDOp[opx_name], autogen.VOPDOp[opy_name]
|
||||
x_ops, y_ops = [parse_operand(p)[0] for p in x_parts[1:]], [parse_operand(p)[0] for p in y_parts[1:]]
|
||||
vdstx, srcx0 = x_ops[0], x_ops[1] if len(x_ops) > 1 else 0
|
||||
vsrcx1 = x_ops[2] if len(x_ops) > 2 else VGPR(0)
|
||||
vdsty, srcy0 = y_ops[0], y_ops[1] if len(y_ops) > 1 else 0
|
||||
vsrcy1 = y_ops[2] if len(y_ops) > 2 else VGPR(0)
|
||||
# Handle fmaak/fmamk literals (4th operand on x or y side)
|
||||
lit = None
|
||||
if 'fmaak' in opx_name.lower() and len(x_ops) > 3: lit = unwrap(x_ops[3])
|
||||
elif 'fmamk' in opx_name.lower() and len(x_ops) > 3: lit, vsrcx1 = unwrap(x_ops[2]), x_ops[3]
|
||||
elif 'fmaak' in opy_name.lower() and len(y_ops) > 3: lit = unwrap(y_ops[3])
|
||||
elif 'fmamk' in opy_name.lower() and len(y_ops) > 3: lit, vsrcy1 = unwrap(y_ops[2]), y_ops[3]
|
||||
return autogen.VOPD(opx, opy, vdstx=vdstx, vdsty=vdsty, srcx0=srcx0, vsrcx1=vsrcx1, srcy0=srcy0, vsrcy1=vsrcy1, literal=lit)
|
||||
operands, current, depth, in_pipe = [], "", 0, False
|
||||
for ch in op_str:
|
||||
if ch in '[(': depth += 1
|
||||
elif ch in '])': depth -= 1
|
||||
elif ch == '|': in_pipe = not in_pipe
|
||||
if ch == ',' and depth == 0 and not in_pipe: operands.append(current.strip()); current = ""
|
||||
else: current += ch
|
||||
if current.strip(): operands.append(current.strip())
|
||||
parsed = [parse_operand(op) for op in operands]
|
||||
values = [p[0] for p in parsed]
|
||||
neg_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[1])
|
||||
abs_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[2])
|
||||
opsel_bits = (8 if len(parsed) > 0 and parsed[0][3] else 0) | sum((1 << i) for i, p in enumerate(parsed[1:4]) if p[3])
|
||||
lit = None
|
||||
if mnemonic in ('v_fmaak_f32', 'v_fmaak_f16') and len(values) == 4: lit, values = unwrap(values[3]), values[:3]
|
||||
elif mnemonic in ('v_fmamk_f32', 'v_fmamk_f16') and len(values) == 4: lit, values = unwrap(values[2]), [values[0], values[1], values[3]]
|
||||
vcc_ops = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32', 'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'}
|
||||
if mnemonic.replace('_e32', '') in vcc_ops and len(values) >= 5: values = [values[0], values[2], values[3]]
|
||||
if mnemonic.startswith('v_cmp') and len(values) >= 3 and operands[0].strip().lower() in ('vcc_lo', 'vcc_hi', 'vcc'):
|
||||
values = values[1:]
|
||||
# CMPX instructions with _e64 suffix: prepend implicit EXEC_LO destination (vdst=126)
|
||||
if 'cmpx' in mnemonic and mnemonic.endswith('_e64') and len(values) == 2:
|
||||
values = [VGPR(126, 1)] + values
|
||||
# Recalculate modifiers: parsed[0]=src0, parsed[1]=src1 (no vdst in user input)
|
||||
neg_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[1])
|
||||
abs_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[2])
|
||||
opsel_bits = sum((1 << i) for i, p in enumerate(parsed[:2]) if p[3])
|
||||
vop3sd_ops = {'v_div_scale_f32', 'v_div_scale_f64'}
|
||||
if mnemonic in vop3sd_ops and len(parsed) >= 5:
|
||||
neg_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[1])
|
||||
abs_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[2])
|
||||
if mnemonic in SOPK_UNSUPPORTED: raise ValueError(f"unsupported instruction: {mnemonic}")
|
||||
elif mnemonic in SOP1_SRC_ONLY:
|
||||
return getattr(autogen, mnemonic)(ssrc0=values[0])
|
||||
elif mnemonic in SOP1_MSG_IMM:
|
||||
return getattr(autogen, mnemonic)(sdst=values[0], ssrc0=RawImm(unwrap(values[1])))
|
||||
elif mnemonic in SOPK_IMM_ONLY:
|
||||
return getattr(autogen, mnemonic)(simm16=values[0])
|
||||
elif mnemonic in SOPK_IMM_FIRST:
|
||||
return getattr(autogen, mnemonic)(simm16=values[0], sdst=values[1])
|
||||
elif mnemonic in SMEM_OPS and len(operands) >= 3 and re.match(r'^-?[0-9]|^-?0x', operands[2].strip().lower()):
|
||||
return getattr(autogen, mnemonic)(sdata=values[0], sbase=values[1], offset=values[2], soffset=RawImm(124))
|
||||
elif mnemonic.startswith('buffer_') and len(operands) >= 2 and operands[1].strip().lower() == 'off':
|
||||
return getattr(autogen, mnemonic)(vdata=values[0], vaddr=0, srsrc=values[2], soffset=RawImm(unwrap(values[3])) if len(values) > 3 else RawImm(0))
|
||||
elif (mnemonic.startswith('flat_load') or mnemonic.startswith('global_load') or mnemonic.startswith('scratch_load')) and len(values) >= 3:
|
||||
offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
|
||||
return getattr(autogen, mnemonic)(vdst=values[0], addr=values[1], saddr=values[2], offset=offset)
|
||||
elif (mnemonic.startswith('flat_store') or mnemonic.startswith('global_store') or mnemonic.startswith('scratch_store')) and len(values) >= 3:
|
||||
offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
|
||||
return getattr(autogen, mnemonic)(addr=values[0], data=values[1], saddr=values[2], offset=offset)
|
||||
for suffix in (['_e32', ''] if not (neg_bits or abs_bits or clamp) else ['', '_e32']):
|
||||
if hasattr(autogen, name := mnemonic.replace('.', '_') + suffix):
|
||||
use_opsel = 'opsel' in getattr(autogen, name).func._fields
|
||||
vals = [type(v)(v.idx, v.count, False) if isinstance(v, Reg) and v.hi and use_opsel else v for v in values]
|
||||
inst = getattr(autogen, name)(*vals, literal=lit, **modifiers)
|
||||
if neg_bits and 'neg' in inst._fields: inst._values['neg'] = neg_bits
|
||||
if opsel_bits and use_opsel: inst._values['opsel'] = opsel_bits
|
||||
if abs_bits and 'abs' in inst._fields: inst._values['abs'] = abs_bits
|
||||
if clamp and 'clmp' in inst._fields: inst._values['clmp'] = 1
|
||||
return inst
|
||||
raise ValueError(f"unknown instruction: {mnemonic}")
|
||||
@@ -1,505 +0,0 @@
|
||||
# RDNA3 emulator - pure Python implementation for testing
|
||||
from __future__ import annotations
|
||||
import ctypes, struct, math
|
||||
from typing import Callable
|
||||
from extra.assembly.rdna3.lib import Inst, Inst32, Inst64, RawImm
|
||||
|
||||
Program = dict[int, Inst] # pc (word offset) -> instruction
|
||||
from extra.assembly.rdna3.autogen import (
|
||||
SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD, SrcEnum,
|
||||
SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, VOPDOp
|
||||
)
|
||||
from extra.assembly.rdna3.alu import (
|
||||
f32, i32, f16, i16, sext, vopc, SALU, VALU,
|
||||
SOP1_BASE, SOP2_BASE, SOPC_BASE, SOPK_BASE, VOP1_BASE, VOP2_BASE
|
||||
)
|
||||
|
||||
WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
|
||||
VCC_LO, VCC_HI, NULL, M0, EXEC_LO, EXEC_HI, SCC = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.M0, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.SCC
|
||||
# Pre-computed inline constant table for src operands 128-254 (index = src - 128)
|
||||
_INLINE_CONSTS = [0] * 127
|
||||
for _i in range(65): _INLINE_CONSTS[_i] = _i # 128-192 -> 0-64
|
||||
for _i in range(1, 17): _INLINE_CONSTS[64 + _i] = ((-_i) & 0xffffffff) # 193-208 -> -1 to -16
|
||||
for _k, _v in {SrcEnum.POS_HALF: 0x3f000000, SrcEnum.NEG_HALF: 0xbf000000, SrcEnum.POS_ONE: 0x3f800000, SrcEnum.NEG_ONE: 0xbf800000,
|
||||
SrcEnum.POS_TWO: 0x40000000, SrcEnum.NEG_TWO: 0xc0000000, SrcEnum.POS_FOUR: 0x40800000, SrcEnum.NEG_FOUR: 0xc0800000,
|
||||
SrcEnum.INV_2PI: 0x3e22f983}.items(): _INLINE_CONSTS[_k - 128] = _v
|
||||
|
||||
_valid_mem_ranges: list[tuple[int, int]] = []
|
||||
def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: global _valid_mem_ranges; _valid_mem_ranges = list(ranges)
|
||||
def _mem_valid(addr: int, size: int) -> bool:
|
||||
for s, z in _valid_mem_ranges:
|
||||
if s <= addr and addr + size <= s + z: return True
|
||||
return not _valid_mem_ranges
|
||||
def _ctypes_at(addr: int, size: int): return (ctypes.c_uint8 if size == 1 else ctypes.c_uint16 if size == 2 else ctypes.c_uint32).from_address(addr)
|
||||
def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
|
||||
def mem_write(addr: int, size: int, val: int) -> None:
|
||||
if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
|
||||
|
||||
# Memory op tables - (cnt, sz, sign) for loads, (cnt, sz) for stores
|
||||
def _mem_ops(ops, suffix_map):
|
||||
return {getattr(e, f"{p}_{s}"): v for e in ops for s, v in suffix_map.items() for p in [e.__name__.replace("Op", "")]}
|
||||
_LOAD_MAP = {'LOAD_B32': (1,4,0), 'LOAD_B64': (2,4,0), 'LOAD_B96': (3,4,0), 'LOAD_B128': (4,4,0), 'LOAD_U8': (1,1,0), 'LOAD_I8': (1,1,1), 'LOAD_U16': (1,2,0), 'LOAD_I16': (1,2,1)}
|
||||
_STORE_MAP = {'STORE_B32': (1,4), 'STORE_B64': (2,4), 'STORE_B96': (3,4), 'STORE_B128': (4,4), 'STORE_B8': (1,1), 'STORE_B16': (1,2)}
|
||||
FLAT_LOAD = _mem_ops([GLOBALOp, FLATOp], _LOAD_MAP)
|
||||
FLAT_STORE = _mem_ops([GLOBALOp, FLATOp], _STORE_MAP)
|
||||
DS_LOAD: dict[int, tuple[int,int,int]] = {DSOp.DS_LOAD_B32: (1,4,0), DSOp.DS_LOAD_B64: (2,4,0), DSOp.DS_LOAD_B128: (4,4,0), DSOp.DS_LOAD_U8: (1,1,0), DSOp.DS_LOAD_I8: (1,1,1), DSOp.DS_LOAD_U16: (1,2,0), DSOp.DS_LOAD_I16: (1,2,1)}
|
||||
DS_STORE: dict[int, tuple[int,int]] = {DSOp.DS_STORE_B32: (1,4), DSOp.DS_STORE_B64: (2,4), DSOp.DS_STORE_B128: (4,4), DSOp.DS_STORE_B8: (1,1), DSOp.DS_STORE_B16: (1,2)}
|
||||
FLAT_D16_LO = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_U8', (1, 0)), ('LOAD_D16_I8', (1, 1)), ('LOAD_D16_B16', (2, 0))]}
|
||||
FLAT_D16_HI = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_HI_U8', (1, 0)), ('LOAD_D16_HI_I8', (1, 1)), ('LOAD_D16_HI_B16', (2, 0))]}
|
||||
FLAT_D16_STORE = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('STORE_D16_HI_B8', 1), ('STORE_D16_HI_B16', 2)]}
|
||||
SMEM_LOAD: dict[int, int] = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}
|
||||
SOPK_WAIT = {SOPKOp.S_WAITCNT_VSCNT, SOPKOp.S_WAITCNT_VMCNT, SOPKOp.S_WAITCNT_EXPCNT, SOPKOp.S_WAITCNT_LGKMCNT}
|
||||
|
||||
class WaveState:
|
||||
__slots__ = ('sgpr', 'vgpr', 'scc', 'pc', 'literal', '_pend_sgpr')
|
||||
def __init__(self):
|
||||
self.sgpr, self.vgpr = [0] * SGPR_COUNT, [[0] * VGPR_COUNT for _ in range(WAVE_SIZE)]
|
||||
self.sgpr[EXEC_LO] = 0xffffffff # wave32: all lanes active
|
||||
self.scc = self.pc = self.literal = 0
|
||||
self._pend_sgpr = {}
|
||||
|
||||
@property
|
||||
def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
|
||||
@vcc.setter
|
||||
def vcc(self, v: int) -> None: self.sgpr[VCC_LO] = v & 0xffffffff; self.sgpr[VCC_HI] = (v >> 32) & 0xffffffff
|
||||
@property
|
||||
def exec_mask(self) -> int: return self.sgpr[EXEC_LO] | (self.sgpr[EXEC_HI] << 32)
|
||||
@exec_mask.setter
|
||||
def exec_mask(self, v: int) -> None: self.sgpr[EXEC_LO] = v & 0xffffffff; self.sgpr[EXEC_HI] = (v >> 32) & 0xffffffff
|
||||
|
||||
def rsgpr(self, i: int) -> int:
|
||||
if i == NULL: return 0
|
||||
if i == SCC: return self.scc
|
||||
return self.sgpr[i] if i < SGPR_COUNT else 0
|
||||
def wsgpr(self, i: int, v: int) -> None:
|
||||
if i < SGPR_COUNT and i != NULL: self.sgpr[i] = v & 0xffffffff
|
||||
def rsgpr64(self, i: int) -> int: return self.rsgpr(i) | (self.rsgpr(i+1) << 32)
|
||||
def wsgpr64(self, i: int, v: int) -> None: self.wsgpr(i, v & 0xffffffff); self.wsgpr(i+1, (v >> 32) & 0xffffffff)
|
||||
|
||||
def rsrc(self, v: int, lane: int) -> int:
|
||||
if v < SGPR_COUNT: return self.sgpr[v]
|
||||
if v == SCC: return self.scc
|
||||
if v < 255: return _INLINE_CONSTS[v - 128]
|
||||
if v == 255: return self.literal
|
||||
return self.vgpr[lane][v - 256] if v <= 511 else 0
|
||||
|
||||
def rsrc64(self, v: int, lane: int) -> int:
|
||||
return self.rsrc(v, lane) | ((self.rsrc(v+1, lane) if v < VCC_LO or 256 <= v <= 511 else 0) << 32)
|
||||
|
||||
def pend_sgpr_lane(self, reg: int, lane: int, val: int) -> None:
|
||||
if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
|
||||
if val: self._pend_sgpr[reg] |= (1 << lane)
|
||||
|
||||
def commit_pends(self) -> None:
|
||||
for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
|
||||
self._pend_sgpr.clear()
|
||||
|
||||
def decode_format(word: int) -> tuple[type[Inst] | None, bool]:
|
||||
hi2 = (word >> 30) & 0x3
|
||||
if hi2 == 0b11:
|
||||
enc = (word >> 26) & 0xf
|
||||
if enc == 0b1101: return SMEM, True
|
||||
if enc == 0b0101:
|
||||
op = (word >> 16) & 0x3ff
|
||||
return (VOP3SD, True) if op in (288, 289, 290, 764, 765, 766, 767, 768, 769, 770) else (VOP3, True)
|
||||
return {0b0011: (VOP3P, True), 0b0110: (DS, True), 0b0111: (FLAT, True), 0b0010: (VOPD, True)}.get(enc, (None, True))
|
||||
if hi2 == 0b10:
|
||||
enc = (word >> 23) & 0x7f
|
||||
return {0b1111101: (SOP1, False), 0b1111110: (SOPC, False), 0b1111111: (SOPP, False)}.get(enc, (SOPK, False) if ((word >> 28) & 0xf) == 0b1011 else (SOP2, False))
|
||||
enc = (word >> 25) & 0x7f
|
||||
return (VOPC, False) if enc == 0b0111110 else (VOP1, False) if enc == 0b0111111 else (VOP2, False)
|
||||
|
||||
def _unwrap(v) -> int: return v.val if isinstance(v, RawImm) else v.value if hasattr(v, 'value') else v
|
||||
|
||||
def decode_program(data: bytes) -> Program:
|
||||
result: Program = {}
|
||||
i = 0
|
||||
while i < len(data):
|
||||
word = int.from_bytes(data[i:i+4], 'little')
|
||||
inst_class, is_64 = decode_format(word)
|
||||
if inst_class is None: i += 4; continue
|
||||
base_size = 8 if is_64 else 4
|
||||
inst = inst_class.from_bytes(data[i:i+base_size])
|
||||
for name, val in inst._values.items(): setattr(inst, name, _unwrap(val))
|
||||
has_literal = any(getattr(inst, fld, None) == 255 for fld in ('src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'srcx0', 'srcy0'))
|
||||
if inst_class == VOP2 and inst.op in (44, 45, 55, 56): has_literal = True
|
||||
if inst_class == VOPD and (inst.opx in (1, 2) or inst.opy in (1, 2)): has_literal = True
|
||||
if inst_class == SOP2 and inst.op in (69, 70): has_literal = True
|
||||
if has_literal: inst._literal = int.from_bytes(data[i+base_size:i+base_size+4], 'little')
|
||||
inst._words = inst.size() // 4 # cache size for step_wave
|
||||
result[i // 4] = inst
|
||||
i += inst._words * 4
|
||||
return result
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# SCALAR EXECUTION
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
def exec_sop1(st: WaveState, inst: SOP1) -> int:
|
||||
s0, op = st.rsrc(inst.ssrc0, 0), inst.op
|
||||
# 64-bit and special ops handled inline
|
||||
if op == SOP1Op.S_MOV_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0)); return 0
|
||||
if op == SOP1Op.S_NOT_B64: r = (~st.rsrc64(inst.ssrc0, 0)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
|
||||
if op == SOP1Op.S_BITSET0_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) & ~(1 << (s0 & 0x1f))); return 0
|
||||
if op == SOP1Op.S_BITSET1_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) | (1 << (s0 & 0x1f))); return 0
|
||||
if op == SOP1Op.S_AND_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
|
||||
if op == SOP1Op.S_OR_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 | old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
|
||||
if op == SOP1Op.S_AND_NOT1_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & (~old & 0xffffffff); st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
|
||||
if op == SOP1Op.S_GETPC_B64: return -3
|
||||
if op == SOP1Op.S_SETPC_B64: return -4
|
||||
if op == SOP1Op.S_SWAPPC_B64: return -5
|
||||
if (fn := SALU.get(SOP1_BASE + op)) is None: raise NotImplementedError(f"SOP1 op {op}")
|
||||
r, scc = fn(s0, 0, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
|
||||
|
||||
_SOP2_64: dict[int, Callable[[int, int], int]] = {SOP2Op.S_AND_B64: lambda a, b: a & b, SOP2Op.S_OR_B64: lambda a, b: a | b, SOP2Op.S_XOR_B64: lambda a, b: a ^ b}
|
||||
def exec_sop2(st: WaveState, inst: SOP2) -> int:
|
||||
s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
|
||||
# 64-bit ops handled inline
|
||||
if op == SOP2Op.S_LSHL_B64: r = (st.rsrc64(inst.ssrc0, 0) << (s1 & 0x3f)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
|
||||
if op == SOP2Op.S_LSHR_B64: r = st.rsrc64(inst.ssrc0, 0) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
|
||||
if op == SOP2Op.S_ASHR_I64: r = sext(st.rsrc64(inst.ssrc0, 0), 64) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r & 0xffffffffffffffff); st.scc = int(r != 0); return 0
|
||||
if (fn := _SOP2_64.get(op)): r = fn(st.rsrc64(inst.ssrc0, 0), st.rsrc64(inst.ssrc1, 0)); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
|
||||
if op == SOP2Op.S_CSELECT_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0) if st.scc else st.rsrc64(inst.ssrc1, 0)); return 0
|
||||
if op == SOP2Op.S_FMAC_F32: st.wsgpr(inst.sdst, i32(f32(st.rsgpr(inst.sdst)) + f32(s0) * f32(s1))); return 0
|
||||
if op == SOP2Op.S_FMAAK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(s1) + f32(inst._literal or 0))); return 0
|
||||
if op == SOP2Op.S_FMAMK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(inst._literal or 0) + f32(s1))); return 0
|
||||
if (fn := SALU.get(SOP2_BASE + op)) is None: raise NotImplementedError(f"SOP2 op {op}")
|
||||
r, scc = fn(s0, s1, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
|
||||
|
||||
def exec_sopc(st: WaveState, inst: SOPC) -> int:
|
||||
s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
|
||||
if op == SOPCOp.S_CMP_EQ_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) == st.rsrc64(inst.ssrc1, 0)); return 0
|
||||
if op == SOPCOp.S_CMP_LG_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) != st.rsrc64(inst.ssrc1, 0)); return 0
|
||||
if (fn := SALU.get(SOPC_BASE + op)) is None: raise NotImplementedError(f"SOPC op {op}")
|
||||
st.scc = fn(s0, s1, st.scc)[1]; return 0
|
||||
|
||||
_SOPK_CMP = frozenset((SOPKOp.S_CMPK_EQ_I32, SOPKOp.S_CMPK_LG_I32, SOPKOp.S_CMPK_GT_I32, SOPKOp.S_CMPK_GE_I32,
|
||||
SOPKOp.S_CMPK_LT_I32, SOPKOp.S_CMPK_LE_I32, SOPKOp.S_CMPK_EQ_U32, SOPKOp.S_CMPK_LG_U32,
|
||||
SOPKOp.S_CMPK_GT_U32, SOPKOp.S_CMPK_GE_U32, SOPKOp.S_CMPK_LT_U32, SOPKOp.S_CMPK_LE_U32))
|
||||
def exec_sopk(st: WaveState, inst: SOPK) -> int:
|
||||
simm, s0, op = inst.simm16, st.rsgpr(inst.sdst), inst.op
|
||||
if op in SOPK_WAIT: return 0
|
||||
if (fn := SALU.get(SOPK_BASE + op)) is None: raise NotImplementedError(f"SOPK op {op}")
|
||||
r, scc = fn(s0, simm, st.scc)
|
||||
if op not in _SOPK_CMP: st.wsgpr(inst.sdst, r)
|
||||
st.scc = scc; return 0
|
||||
|
||||
def exec_sopp(st: WaveState, inst: SOPP) -> int:
|
||||
if inst.op == SOPPOp.S_ENDPGM: return -1
|
||||
if inst.op == SOPPOp.S_BARRIER: return -2
|
||||
if inst.op == SOPPOp.S_BRANCH: return sext(inst.simm16, 16)
|
||||
if inst.op == SOPPOp.S_CBRANCH_SCC0: return sext(inst.simm16, 16) if st.scc == 0 else 0
|
||||
if inst.op == SOPPOp.S_CBRANCH_SCC1: return sext(inst.simm16, 16) if st.scc == 1 else 0
|
||||
# In wave32 mode, only VCC_LO is used for lane masks; VCC_HI is a free SGPR
|
||||
if inst.op == SOPPOp.S_CBRANCH_VCCZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) == 0 else 0
|
||||
if inst.op == SOPPOp.S_CBRANCH_VCCNZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) != 0 else 0
|
||||
if inst.op == SOPPOp.S_CBRANCH_EXECZ: return sext(inst.simm16, 16) if st.exec_mask == 0 else 0
|
||||
if inst.op == SOPPOp.S_CBRANCH_EXECNZ: return sext(inst.simm16, 16) if st.exec_mask != 0 else 0
|
||||
# Scheduling hints and wait instructions are no-ops in emulation
|
||||
if inst.op <= 31: return 0 # S_NOP, S_CLAUSE, S_DELAY_ALU, S_WAITCNT, etc.
|
||||
# S_WAKEUP(52), S_SETPRIO(53), S_SENDMSG(54), S_SENDMSGHALT(55), perf counters, S_ICACHE_INV(60) are no-ops
|
||||
if inst.op in (52, 53, 54, 55, 56, 57, 60): return 0
|
||||
raise NotImplementedError(f"SOPP op {inst.op}")
|
||||
|
||||
def exec_smem(st: WaveState, inst: SMEM) -> int:
|
||||
addr = st.rsgpr64(inst.sbase * 2) + sext(inst.offset, 21)
|
||||
if inst.soffset not in (NULL, 0x7f): addr += st.rsrc(inst.soffset, 0)
|
||||
if (cnt := SMEM_LOAD.get(inst.op)) is None: raise NotImplementedError(f"SMEM op {inst.op}")
|
||||
for i in range(cnt): st.wsgpr(inst.sdata + i, mem_read((addr + i * 4) & 0xffffffffffffffff, 4))
|
||||
return 0
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# VECTOR EXECUTION
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
def f64(hi: int, lo: int) -> float: return struct.unpack('<d', struct.pack('<Q', (hi << 32) | lo))[0]
|
||||
def i64_parts(f: float) -> tuple[int, int]:
|
||||
if math.isnan(f): val = 0x7ff8000000000000
|
||||
elif math.isinf(f): val = 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
|
||||
else: val = struct.unpack('<Q', struct.pack('<d', f))[0]
|
||||
return val & 0xffffffff, (val >> 32) & 0xffffffff
|
||||
|
||||
def exec_vop1(st: WaveState, inst: VOP1, lane: int) -> None:
|
||||
if inst.op == VOP1Op.V_NOP: return
|
||||
V, s0 = st.vgpr[lane], st.rsrc(inst.src0, lane)
|
||||
if inst.op == VOP1Op.V_READFIRSTLANE_B32:
|
||||
first = (st.exec_mask & -st.exec_mask).bit_length() - 1 if st.exec_mask else 0
|
||||
st.wsgpr(inst.vdst, st.rsrc(inst.src0, first) if inst.src0 >= 256 else s0); return
|
||||
# F64 ops handled inline
|
||||
if inst.op == VOP1Op.V_CVT_F64_F32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(f32(s0))); return
|
||||
if inst.op == VOP1Op.V_CVT_F64_I32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(sext(s0, 32))); return
|
||||
if inst.op == VOP1Op.V_CVT_F64_U32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(s0)); return
|
||||
if inst.op in (VOP1Op.V_CVT_F32_F64, VOP1Op.V_CVT_I32_F64, VOP1Op.V_CVT_U32_F64):
|
||||
src = inst.src0 - 256 if inst.src0 >= 256 else inst.src0
|
||||
lo, hi = (V[src], V[src+1]) if inst.src0 >= 256 else (st.sgpr[src], st.sgpr[src+1])
|
||||
v = f64(hi, lo)
|
||||
if inst.op == VOP1Op.V_CVT_F32_F64: V[inst.vdst] = i32(v)
|
||||
elif inst.op == VOP1Op.V_CVT_I32_F64: V[inst.vdst] = (max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff) if math.isfinite(v) else 0
|
||||
else: V[inst.vdst] = max(0, min(0xffffffff, int(v))) if math.isfinite(v) and v == v else 0
|
||||
return
|
||||
if (fn := VALU.get(VOP1_BASE + inst.op)): V[inst.vdst] = fn(s0, 0, 0); return
|
||||
raise NotImplementedError(f"VOP1 op {inst.op}")
|
||||
|
||||
def exec_vop2(st: WaveState, inst: VOP2, lane: int) -> None:
|
||||
V, s0, s1, op = st.vgpr[lane], st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1], inst.op
|
||||
if op == VOP2Op.V_CNDMASK_B32: V[inst.vdst] = s1 if (st.vcc >> lane) & 1 else s0; return
|
||||
if op == VOP2Op.V_FMAC_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(V[inst.vdst])); return
|
||||
if op == VOP2Op.V_FMAMK_F32: V[inst.vdst] = i32(f32(s0)*f32(st.literal)+f32(s1)); return
|
||||
if op == VOP2Op.V_FMAAK_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(st.literal)); return
|
||||
if op == VOP2Op.V_FMAC_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(V[inst.vdst])); return
|
||||
if op == VOP2Op.V_FMAMK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(st.literal)+f16(s1)); return
|
||||
if op == VOP2Op.V_FMAAK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(st.literal)); return
|
||||
if op == VOP2Op.V_PK_FMAC_F16:
|
||||
lo = i16(f16(s0 & 0xffff) * f16(s1 & 0xffff) + f16(V[inst.vdst] & 0xffff))
|
||||
hi = i16(f16((s0 >> 16) & 0xffff) * f16((s1 >> 16) & 0xffff) + f16((V[inst.vdst] >> 16) & 0xffff))
|
||||
V[inst.vdst] = lo | (hi << 16); return
|
||||
if op == VOP2Op.V_ADD_CO_CI_U32: r = s0+s1+((st.vcc>>lane)&1); st.pend_sgpr_lane(VCC_LO, lane, r >= 0x100000000); V[inst.vdst] = r & 0xffffffff; return
|
||||
if op == VOP2Op.V_SUB_CO_CI_U32: b = (st.vcc>>lane)&1; st.pend_sgpr_lane(VCC_LO, lane, s1+b > s0); V[inst.vdst] = (s0-s1-b) & 0xffffffff; return
|
||||
if (fn := VALU.get(VOP2_BASE + op)): V[inst.vdst] = fn(s0, s1, 0); return
|
||||
raise NotImplementedError(f"VOP2 op {op}")
|
||||
|
||||
def vop3_mod(val: int, neg: int, abs_: int, idx: int) -> int:
|
||||
if (abs_ >> idx) & 1: val = i32(abs(f32(val)))
|
||||
if (neg >> idx) & 1: val = i32(-f32(val))
|
||||
return val
|
||||
|
||||
def exec_vop3(st: WaveState, inst: VOP3, lane: int) -> None:
|
||||
op, src0, src1, src2, vdst, neg, abs_ = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.neg, getattr(inst, 'abs', 0)
|
||||
V = st.vgpr[lane]
|
||||
# VOPC encoded in VOP3 (0-255)
|
||||
if 0 <= op <= 255:
|
||||
base = op & 0x7f
|
||||
# For 64-bit comparisons (I64: 80-87, U64: 88-95), read raw 64-bit values (no float modifiers)
|
||||
if 80 <= base <= 95:
|
||||
s0_64, s1_64 = st.rsrc64(src0, lane), st.rsrc64(src1, lane)
|
||||
result = vopc(op, s0_64 & 0xffffffff, s1_64 & 0xffffffff, (s0_64 >> 32) & 0xffffffff, (s1_64 >> 32) & 0xffffffff)
|
||||
else:
|
||||
s0, s1 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1)
|
||||
result = vopc(op, s0, s1)
|
||||
is_cmpx = op >= 128
|
||||
st.pend_sgpr_lane(vdst, lane, result)
|
||||
if is_cmpx: st.pend_sgpr_lane(EXEC_LO, lane, result)
|
||||
return
|
||||
s0, s1, s2 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1), vop3_mod(st.rsrc(src2, lane), neg, abs_, 2)
|
||||
# Special ops
|
||||
if op == VOP3Op.V_FMAC_F32: V[vdst] = i32(f32(s0)*f32(s1)+f32(V[vdst])); return
|
||||
if op == VOP3Op.V_READLANE_B32: st.wsgpr(vdst, st.vgpr[s1 & 0x1f][src0 - 256] if src0 >= 256 else s0); return
|
||||
if op == VOP3Op.V_WRITELANE_B32: st.vgpr[s1 & 0x1f][vdst] = s0; return
|
||||
if op == VOP3Op.V_CNDMASK_B32:
|
||||
mask = st.rsgpr(src2) if src2 < 256 else st.vcc
|
||||
V[vdst] = s1 if (mask >> lane) & 1 else s0; return
|
||||
if op in (VOP3Op.V_LSHLREV_B64, VOP3Op.V_LSHRREV_B64, VOP3Op.V_ASHRREV_I64):
|
||||
v64 = st.rsrc64(src1, lane)
|
||||
r = ((v64 << (s0 & 0x3f)) & 0xffffffffffffffff if op == VOP3Op.V_LSHLREV_B64 else
|
||||
v64 >> (s0 & 0x3f) if op == VOP3Op.V_LSHRREV_B64 else sext(v64, 64) >> (s0 & 0x3f))
|
||||
V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff; return
|
||||
if op in (VOP3Op.V_ADD_F64, VOP3Op.V_MUL_F64, VOP3Op.V_FMA_F64, VOP3Op.V_MAX_F64, VOP3Op.V_MIN_F64):
|
||||
a, b = f64(st.rsrc(src0+1, lane), s0), f64(st.rsrc(src1+1, lane), s1)
|
||||
c = f64(st.rsrc(src2+1, lane), s2) if op == VOP3Op.V_FMA_F64 else 0.0
|
||||
rf = a + b if op == VOP3Op.V_ADD_F64 else a * b if op == VOP3Op.V_MUL_F64 else a * b + c if op == VOP3Op.V_FMA_F64 else max(a, b) if op == VOP3Op.V_MAX_F64 else min(a, b)
|
||||
V[vdst], V[vdst+1] = i64_parts(rf); return
|
||||
if (fn := VALU.get(op)): V[vdst] = fn(s0, s1, s2); return
|
||||
raise NotImplementedError(f"VOP3 op {op}")
|
||||
|
||||
def exec_vopc(st: WaveState, inst: VOPC, lane: int) -> None:
|
||||
result, is_cmpx = vopc(inst.op, st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1]), inst.op >= 128
|
||||
st.pend_sgpr_lane(EXEC_LO if is_cmpx else VCC_LO, lane, result)
|
||||
|
||||
def exec_vop3sd(st: WaveState, inst: VOP3SD, lane: int) -> None:
|
||||
op, src0, src1, src2, vdst, sdst, neg = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.sdst, inst.neg
|
||||
s0, s1, s2 = st.rsrc(src0, lane), st.rsrc(src1, lane), st.rsrc(src2, lane)
|
||||
if (neg >> 0) & 1: s0 = i32(-f32(s0))
|
||||
if (neg >> 1) & 1: s1 = i32(-f32(s1))
|
||||
if (neg >> 2) & 1: s2 = i32(-f32(s2))
|
||||
V = st.vgpr[lane]
|
||||
if op == VOP3SDOp.V_ADD_CO_U32: r = s0 + s1; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
|
||||
elif op == VOP3SDOp.V_SUB_CO_U32: V[vdst] = (s0 - s1) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 > s0)
|
||||
elif op == VOP3SDOp.V_SUBREV_CO_U32: V[vdst] = (s1 - s0) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s0 > s1)
|
||||
elif op == VOP3SDOp.V_ADD_CO_CI_U32:
|
||||
cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
|
||||
r = s0 + s1 + cin; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
|
||||
elif op == VOP3SDOp.V_SUB_CO_CI_U32:
|
||||
cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
|
||||
V[vdst] = (s0 - s1 - cin) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 + cin > s0)
|
||||
elif op == VOP3SDOp.V_MAD_U64_U32:
|
||||
s2_64 = s2 | (st.rsrc(src2+1, lane) << 32); r = s0 * s1 + s2_64
|
||||
V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
|
||||
elif op == VOP3SDOp.V_MAD_I64_I32:
|
||||
s2_64 = sext(s2 | (st.rsrc(src2+1, lane) << 32), 64)
|
||||
r = (sext(s0, 32) * sext(s1, 32) + s2_64) & 0xffffffffffffffff
|
||||
V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
|
||||
elif op == VOP3SDOp.V_DIV_SCALE_F32: V[vdst] = 0; st.pend_sgpr_lane(sdst, lane, False)
|
||||
elif op == VOP3SDOp.V_DIV_SCALE_F64: V[vdst], V[vdst+1] = s0, st.rsrc(src0+1, lane); st.pend_sgpr_lane(VCC_LO, lane, s0 == s2)
|
||||
else: raise NotImplementedError(f"VOP3SD op {op}")
|
||||
|
||||
def exec_flat(st: WaveState, inst: FLAT, lane: int) -> None:
|
||||
op, addr_reg, data_reg, vdst, offset, saddr, V = inst.op, inst.addr, inst.data, inst.vdst, sext(inst.offset, 13), inst.saddr, st.vgpr[lane]
|
||||
addr = V[addr_reg] | (V[addr_reg+1] << 32)
|
||||
addr = (st.rsgpr64(saddr) + V[addr_reg] + offset) & 0xffffffffffffffff if saddr not in (NULL, 0x7f) else (addr + offset) & 0xffffffffffffffff
|
||||
if op in FLAT_LOAD:
|
||||
cnt, sz, sign = FLAT_LOAD[op]
|
||||
for i in range(cnt): val = mem_read(addr + i * sz, sz); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
|
||||
elif op in FLAT_STORE:
|
||||
cnt, sz = FLAT_STORE[op]
|
||||
for i in range(cnt): mem_write(addr + i * sz, sz, V[data_reg + i] & ((1 << (sz * 8)) - 1))
|
||||
elif op in FLAT_D16_LO: sz, sign = FLAT_D16_LO[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0xffff0000) | ((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff))
|
||||
elif op in FLAT_D16_HI: sz, sign = FLAT_D16_HI[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0x0000ffff) | (((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff)) << 16)
|
||||
elif op in FLAT_D16_STORE: mem_write(addr, FLAT_D16_STORE[op], (V[data_reg] >> 16) & ((1 << (FLAT_D16_STORE[op] * 8)) - 1))
|
||||
else: raise NotImplementedError(f"FLAT op {op}")
|
||||
|
||||
def exec_ds(st: WaveState, inst: DS, lane: int, lds: bytearray) -> None:
|
||||
op, addr, vdst, V = inst.op, (st.vgpr[lane][inst.addr] + inst.offset0) & 0xffff, inst.vdst, st.vgpr[lane]
|
||||
if op in DS_LOAD:
|
||||
cnt, sz, sign = DS_LOAD[op]
|
||||
for i in range(cnt): val = int.from_bytes(lds[addr+i*sz:addr+i*sz+sz], 'little'); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
|
||||
elif op in DS_STORE:
|
||||
cnt, sz = DS_STORE[op]
|
||||
for i in range(cnt): lds[addr+i*sz:addr+i*sz+sz] = (V[inst.data0 + i] & ((1 << (sz * 8)) - 1)).to_bytes(sz, 'little')
|
||||
else: raise NotImplementedError(f"DS op {op}")
|
||||
|
||||
VOPD_OPS: dict[int, Callable[[int, int, int, int, int], int]] = {
|
||||
VOPDOp.V_DUAL_MUL_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)), VOPDOp.V_DUAL_ADD_F32: lambda a, b, d, l, lit: i32(f32(a)+f32(b)),
|
||||
VOPDOp.V_DUAL_SUB_F32: lambda a, b, d, l, lit: i32(f32(a)-f32(b)), VOPDOp.V_DUAL_SUBREV_F32: lambda a, b, d, l, lit: i32(f32(b)-f32(a)),
|
||||
VOPDOp.V_DUAL_MAX_F32: lambda a, b, d, l, lit: i32(max(f32(a), f32(b))), VOPDOp.V_DUAL_MIN_F32: lambda a, b, d, l, lit: i32(min(f32(a), f32(b))),
|
||||
VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: lambda a, b, d, l, lit: i32(0.0 if f32(a) == 0.0 or f32(b) == 0.0 else f32(a)*f32(b)),
|
||||
VOPDOp.V_DUAL_MOV_B32: lambda a, b, d, l, lit: a, VOPDOp.V_DUAL_ADD_NC_U32: lambda a, b, d, l, lit: (a + b) & 0xffffffff,
|
||||
VOPDOp.V_DUAL_LSHLREV_B32: lambda a, b, d, l, lit: (b << (a & 0x1f)) & 0xffffffff, VOPDOp.V_DUAL_AND_B32: lambda a, b, d, l, lit: a & b,
|
||||
VOPDOp.V_DUAL_FMAC_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(d)), VOPDOp.V_DUAL_FMAAK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(lit)),
|
||||
VOPDOp.V_DUAL_FMAMK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(lit)+f32(b)), VOPDOp.V_DUAL_CNDMASK_B32: lambda a, b, d, l, lit: b if l else a,
|
||||
}
|
||||
def exec_vopd(st: WaveState, inst: VOPD, lane: int) -> None:
|
||||
V, vdsty, vcc_lane = st.vgpr[lane], (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), (st.vcc >> lane) & 1
|
||||
sx0, sx1, sy0, sy1, dstx = st.rsrc(inst.srcx0, lane), V[inst.vsrcx1], st.rsrc(inst.srcy0, lane), V[inst.vsrcy1], inst.vdstx
|
||||
if (fn := VOPD_OPS.get(inst.opx)): V[dstx] = fn(sx0, sx1, V[dstx], vcc_lane, st.literal)
|
||||
else: raise NotImplementedError(f"VOPD opx {inst.opx}")
|
||||
if (fn := VOPD_OPS.get(inst.opy)): V[vdsty] = fn(sy0, sy1, V[vdsty], vcc_lane, st.literal)
|
||||
else: raise NotImplementedError(f"VOPD opy {inst.opy}")
|
||||
|
||||
def exec_vop3p(st: WaveState, inst: VOP3P, lane: int) -> None:
|
||||
op, vdst, V = inst.op, inst.vdst, st.vgpr[lane]
|
||||
s0, s1, s2 = st.rsrc(inst.src0, lane), st.rsrc(inst.src1, lane), st.rsrc(inst.src2, lane)
|
||||
opsel, opsel_hi = [(inst.opsel >> i) & 1 for i in range(3)], [(inst.opsel_hi >> i) & 1 for i in range(2)] + [inst.opsel_hi2]
|
||||
neg, neg_hi = inst.neg, inst.neg_hi
|
||||
def get_src(src: int, idx: int, for_mix: bool = False) -> float:
|
||||
if for_mix:
|
||||
if not opsel_hi[idx]: return abs(f32(src)) if (neg_hi >> idx) & 1 else f32(src)
|
||||
return float(f16((src >> 16) & 0xffff) if opsel[idx] else f16(src & 0xffff))
|
||||
use_hi = opsel[idx]
|
||||
val = ((src >> 16) & 0xffff) if use_hi else (src & 0xffff)
|
||||
f = f16(val)
|
||||
if use_hi and (neg >> idx) & 1: f = -f
|
||||
elif not use_hi and (neg_hi >> idx) & 1: f = -f
|
||||
return f
|
||||
if op == VOP3POp.V_FMA_MIX_F32: V[vdst] = i32(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
|
||||
elif op == VOP3POp.V_FMA_MIXLO_F16: V[vdst] = (V[vdst] & 0xffff0000) | i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
|
||||
elif op == VOP3POp.V_FMA_MIXHI_F16: V[vdst] = (V[vdst] & 0x0000ffff) | (i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True)) << 16)
|
||||
else: raise NotImplementedError(f"VOP3P op {op}")
|
||||
|
||||
def exec_wmma_f32_16x16x16_f16(st: WaveState, inst: VOP3P, n_lanes: int) -> None:
|
||||
src0_base, src1_base, src2_base = (inst.src0 - 256) if inst.src0 >= 256 else inst.src0, (inst.src1 - 256) if inst.src1 >= 256 else inst.src1, (inst.src2 - 256) if inst.src2 >= 256 else inst.src2
|
||||
src0_is_vgpr, src1_is_vgpr, src2_is_vgpr, vdst = inst.src0 >= 256, inst.src1 >= 256, inst.src2 >= 256, inst.vdst
|
||||
A, B, C = [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)]
|
||||
for lane in range(min(n_lanes, 16)):
|
||||
V = st.vgpr[lane]
|
||||
for reg in range(8):
|
||||
val = V[src0_base + reg] if src0_is_vgpr else st.sgpr[src0_base + reg]
|
||||
A[lane][reg * 2], A[lane][reg * 2 + 1] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
|
||||
val = V[src1_base + reg] if src1_is_vgpr else st.sgpr[src1_base + reg]
|
||||
B[reg * 2][lane], B[reg * 2 + 1][lane] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
|
||||
for row in range(16):
|
||||
for col in range(16):
|
||||
idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
|
||||
if lane_idx < n_lanes:
|
||||
val = st.vgpr[lane_idx][src2_base + reg] if src2_is_vgpr else st.sgpr[src2_base + reg]
|
||||
C[row][col] = f32(val)
|
||||
for row in range(16):
|
||||
for col in range(16):
|
||||
for k in range(16): C[row][col] += A[row][k] * B[k][col]
|
||||
for row in range(16):
|
||||
for col in range(16):
|
||||
idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
|
||||
if lane_idx < n_lanes and (st.exec_mask & (1 << lane_idx)): st.vgpr[lane_idx][vdst + reg] = i32(C[row][col])
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# MAIN EXECUTION LOOP
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
SCALAR: dict[type, Callable[..., int]] = {SOP1: exec_sop1, SOP2: exec_sop2, SOPC: exec_sopc, SOPK: exec_sopk, SOPP: exec_sopp, SMEM: exec_smem}
|
||||
VECTOR: dict[type, Callable[..., None]] = {VOP1: exec_vop1, VOP2: exec_vop2, VOP3: exec_vop3, VOP3SD: exec_vop3sd, VOPC: exec_vopc, FLAT: exec_flat, DS: exec_ds, VOPD: exec_vopd, VOP3P: exec_vop3p}
|
||||
|
||||
_WMMA_OPS = frozenset((VOP3POp.V_WMMA_F32_16X16X16_F16, VOP3POp.V_WMMA_F32_16X16X16_BF16, VOP3POp.V_WMMA_F16_16X16X16_F16,
|
||||
VOP3POp.V_WMMA_BF16_16X16X16_BF16, VOP3POp.V_WMMA_I32_16X16X16_IU8, VOP3POp.V_WMMA_I32_16X16X16_IU4))
|
||||
|
||||
def step_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> int:
|
||||
inst = program.get(st.pc)
|
||||
if inst is None: return 1
|
||||
inst_words, st.literal, inst_type = inst._words, inst._literal or 0, type(inst)
|
||||
if (handler := SCALAR.get(inst_type)) is not None:
|
||||
delta = handler(st, inst)
|
||||
if delta == -1: return -1
|
||||
if delta == -2: st.pc += inst_words; return -2
|
||||
if delta == -3: # S_GETPC_B64
|
||||
sop1 = inst if isinstance(inst, SOP1) else None
|
||||
assert sop1 is not None
|
||||
next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc += inst_words; return 0
|
||||
if delta == -4: # S_SETPC_B64
|
||||
sop1 = inst if isinstance(inst, SOP1) else None
|
||||
assert sop1 is not None
|
||||
st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
|
||||
if delta == -5: # S_SWAPPC_B64
|
||||
sop1 = inst if isinstance(inst, SOP1) else None
|
||||
assert sop1 is not None
|
||||
next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
|
||||
st.pc += inst_words + delta
|
||||
else:
|
||||
vec_handler, exec_mask = VECTOR[inst_type], st.exec_mask
|
||||
if inst_type is DS:
|
||||
for lane in range(n_lanes):
|
||||
if exec_mask & (1 << lane): vec_handler(st, inst, lane, lds)
|
||||
elif inst_type is VOP3P:
|
||||
vop3p = inst if isinstance(inst, VOP3P) else None
|
||||
assert vop3p is not None
|
||||
if vop3p.op in _WMMA_OPS:
|
||||
exec_wmma_f32_16x16x16_f16(st, vop3p, n_lanes)
|
||||
else:
|
||||
for lane in range(n_lanes):
|
||||
if exec_mask & (1 << lane): vec_handler(st, vop3p, lane)
|
||||
else:
|
||||
for lane in range(n_lanes):
|
||||
if exec_mask & (1 << lane): vec_handler(st, inst, lane)
|
||||
st.commit_pends(); st.pc += inst_words
|
||||
return 0
|
||||
|
||||
def exec_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int, wg_id: tuple[int,int,int]=(0,0,0), local_size: tuple[int,int,int]=(1,1,1), wave_start: int=0) -> int:
|
||||
while st.pc in program:
|
||||
result = step_wave(program, st, lds, n_lanes)
|
||||
if result == -1: return 0
|
||||
if result == -2: return -2
|
||||
return 0
|
||||
|
||||
def exec_workgroup(program: Program, workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int, dispatch_dim: int) -> None:
|
||||
lx, ly, lz = local_size
|
||||
total_threads, lds = lx * ly * lz, bytearray(65536)
|
||||
waves: list[tuple[WaveState, int, int]] = []
|
||||
for wave_start in range(0, total_threads, WAVE_SIZE):
|
||||
n_lanes, st = min(WAVE_SIZE, total_threads - wave_start), WaveState()
|
||||
st.exec_mask = (1 << n_lanes) - 1
|
||||
st.wsgpr64(0, args_ptr)
|
||||
gx, gy, gz = workgroup_id
|
||||
if dispatch_dim >= 3: st.sgpr[13], st.sgpr[14], st.sgpr[15] = gx, gy, gz
|
||||
elif dispatch_dim == 2: st.sgpr[14], st.sgpr[15] = gx, gy
|
||||
else: st.sgpr[15] = gx
|
||||
for i in range(n_lanes):
|
||||
tid = wave_start + i
|
||||
st.vgpr[i][0] = tid if local_size == (lx, 1, 1) else ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
|
||||
waves.append((st, n_lanes, wave_start))
|
||||
has_barrier = any(isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER for inst in program.values())
|
||||
for _ in range(2 if has_barrier else 1):
|
||||
for st, n_lanes, wave_start in waves: exec_wave(program, st, lds, n_lanes, workgroup_id, local_size, wave_start)
|
||||
|
||||
def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
|
||||
data = (ctypes.c_char * lib_sz).from_address(lib).raw
|
||||
program = decode_program(data)
|
||||
if not program: return -1
|
||||
dispatch_dim = 3 if gz > 1 else (2 if gy > 1 else 1)
|
||||
for gidz in range(gz):
|
||||
for gidy in range(gy):
|
||||
for gidx in range(gx): exec_workgroup(program, (gidx, gidy, gidz), (lx, ly, lz), args_ptr, dispatch_dim)
|
||||
return 0
|
||||
@@ -1,191 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
# generates autogen/__init__.py by parsing the AMD RDNA3.5 ISA PDF
|
||||
import re, pdfplumber, pathlib
|
||||
from tinygrad.helpers import fetch
|
||||
|
||||
PDF_URL = "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content"
|
||||
FIELD_TYPES = {'SSRC0': 'SSrc', 'SSRC1': 'SSrc', 'SOFFSET': 'SSrc', 'SADDR': 'SSrc', 'SRC0': 'Src', 'SRC1': 'Src', 'SRC2': 'Src',
|
||||
'SDST': 'SGPRField', 'SBASE': 'SGPRField', 'SDATA': 'SGPRField', 'SRSRC': 'SGPRField', 'VDST': 'VGPRField', 'VSRC1': 'VGPRField', 'VDATA': 'VGPRField',
|
||||
'VADDR': 'VGPRField', 'ADDR': 'VGPRField', 'DATA': 'VGPRField', 'DATA0': 'VGPRField', 'DATA1': 'VGPRField', 'SIMM16': 'SImm', 'OFFSET': 'Imm',
|
||||
'OPX': 'VOPDOp', 'OPY': 'VOPDOp', 'SRCX0': 'Src', 'SRCY0': 'Src', 'VSRCX1': 'VGPRField', 'VSRCY1': 'VGPRField', 'VDSTX': 'VGPRField', 'VDSTY': 'VDSTYEnc'}
|
||||
FIELD_ORDER = {
|
||||
'SOP2': ['op', 'sdst', 'ssrc0', 'ssrc1'], 'SOP1': ['op', 'sdst', 'ssrc0'], 'SOPC': ['op', 'ssrc0', 'ssrc1'],
|
||||
'SOPK': ['op', 'sdst', 'simm16'], 'SOPP': ['op', 'simm16'], 'VOP1': ['op', 'vdst', 'src0'], 'VOPC': ['op', 'src0', 'vsrc1'],
|
||||
'VOP2': ['op', 'vdst', 'src0', 'vsrc1'], 'VOP3SD': ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2', 'clmp'],
|
||||
'SMEM': ['op', 'sdata', 'sbase', 'soffset', 'offset', 'glc', 'dlc'], 'DS': ['op', 'vdst', 'addr', 'data0', 'data1'],
|
||||
'VOP3': ['op', 'vdst', 'src0', 'src1', 'src2', 'omod', 'neg', 'abs', 'clmp', 'opsel'],
|
||||
'VOP3P': ['op', 'vdst', 'src0', 'src1', 'src2', 'neg', 'neg_hi', 'opsel', 'opsel_hi', 'clmp'],
|
||||
'FLAT': ['op', 'vdst', 'addr', 'data', 'saddr', 'offset', 'seg', 'dlc', 'glc', 'slc'],
|
||||
'MUBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
|
||||
'MTBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
|
||||
'MIMG': ['op', 'vdata', 'vaddr', 'srsrc', 'ssamp', 'dmask', 'dim', 'unrm', 'dlc', 'glc', 'slc'],
|
||||
'EXP': ['en', 'target', 'vsrc0', 'vsrc1', 'vsrc2', 'vsrc3', 'done', 'row'],
|
||||
'VINTERP': ['op', 'vdst', 'src0', 'src1', 'src2', 'waitexp', 'clmp', 'opsel', 'neg'],
|
||||
'VOPD': ['opx', 'opy', 'vdstx', 'vdsty', 'srcx0', 'vsrcx1', 'srcy0', 'vsrcy1'],
|
||||
'LDSDIR': ['op', 'vdst', 'attr', 'attr_chan', 'wait_va']}
|
||||
SRC_EXTRAS = {233: 'DPP8', 234: 'DPP8FI', 250: 'DPP16', 251: 'VCCZ', 252: 'EXECZ', 254: 'LDS_DIRECT'}
|
||||
FLOAT_MAP = {'0.5': 'POS_HALF', '-0.5': 'NEG_HALF', '1.0': 'POS_ONE', '-1.0': 'NEG_ONE', '2.0': 'POS_TWO', '-2.0': 'NEG_TWO',
|
||||
'4.0': 'POS_FOUR', '-4.0': 'NEG_FOUR', '1/(2*PI)': 'INV_2PI', '0': 'ZERO'}
|
||||
|
||||
def parse_bits(s: str) -> tuple[int, int] | None:
|
||||
return (int(m.group(1)), int(m.group(2) or m.group(1))) if (m := re.match(r'\[(\d+)(?::(\d+))?\]', s)) else None
|
||||
|
||||
def parse_fields_table(table: list, fmt: str, enums: set[str]) -> list[tuple]:
|
||||
fields = []
|
||||
for row in table[1:]:
|
||||
if not row or not row[0]: continue
|
||||
name, bits_str = row[0].split('\n')[0].strip(), (row[1] or '').split('\n')[0].strip()
|
||||
if not (bits := parse_bits(bits_str)): continue
|
||||
enc_val, hi, lo = None, bits[0], bits[1]
|
||||
if name == 'ENCODING' and row[2] and (m := re.search(r"'b([01_]+)", row[2])):
|
||||
enc_bits = m.group(1).replace('_', '')
|
||||
enc_val = int(enc_bits, 2)
|
||||
declared_width, actual_width = hi - lo + 1, len(enc_bits)
|
||||
if actual_width > declared_width: lo = hi - actual_width + 1
|
||||
ftype = f"{fmt}Op" if name == 'OP' and f"{fmt}Op" in enums else FIELD_TYPES.get(name.upper())
|
||||
fields.append((name, hi, lo, enc_val, ftype))
|
||||
return fields
|
||||
|
||||
def generate(output_path: pathlib.Path|str|None = None) -> dict:
|
||||
"""Generate RDNA3.5 instruction definitions from the AMD ISA PDF. Returns dict with formats for testing."""
|
||||
pdf = pdfplumber.open(fetch(PDF_URL))
|
||||
pages = pdf.pages[150:200]
|
||||
page_texts = [p.extract_text() or '' for p in pages]
|
||||
page_tables = [[t.extract() for t in p.find_tables()] for p in pages]
|
||||
full_text = '\n'.join(page_texts)
|
||||
|
||||
# parse SSRC encoding from first page with VCC_LO
|
||||
src_enum = dict(SRC_EXTRAS)
|
||||
for text in page_texts[:10]:
|
||||
if 'SSRC0' in text and 'VCC_LO' in text:
|
||||
for m in re.finditer(r'^(\d+)\s+(\S+)', text, re.M):
|
||||
val, name = int(m.group(1)), m.group(2).rstrip('.:')
|
||||
if name in FLOAT_MAP: src_enum[val] = FLOAT_MAP[name]
|
||||
elif re.match(r'^[A-Z][A-Z0-9_]*$', name): src_enum[val] = name
|
||||
break
|
||||
|
||||
# parse opcode tables
|
||||
enums: dict[str, dict[int, str]] = {}
|
||||
for m in re.finditer(r'Table \d+\. (\w+) Opcodes(.*?)(?=Table \d+\.|\n\d+\.\d+\.\d+\.\s+\w+\s*\nDescription|$)', full_text, re.S):
|
||||
if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+([A-Z][A-Z0-9_]+)', m.group(2))}:
|
||||
enums[m.group(1) + "Op"] = ops
|
||||
if vopd_m := re.search(r'Table \d+\. VOPD Y-Opcodes\n(.*?)(?=Table \d+\.|15\.\d)', full_text, re.S):
|
||||
if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+(V_DUAL_\w+)', vopd_m.group(1))}:
|
||||
enums["VOPDOp"] = ops
|
||||
enum_names = set(enums.keys())
|
||||
|
||||
def is_fields_table(t) -> bool: return t and len(t) > 1 and t[0] and 'Field' in str(t[0][0] or '')
|
||||
def has_encoding(fields) -> bool: return any(f[0] == 'ENCODING' for f in fields)
|
||||
def has_header_before_fields(text) -> bool:
|
||||
return (pos := text.find('Field Name')) != -1 and bool(re.search(r'\d+\.\d+\.\d+\.\s+\w+\s*\n', text[:pos]))
|
||||
|
||||
# find format headers with their page indices
|
||||
format_headers = [] # (fmt_name, page_idx)
|
||||
for i, text in enumerate(page_texts):
|
||||
for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n?Description', text): format_headers.append((m.group(1), i, m.start()))
|
||||
for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n', text):
|
||||
if m.start() > len(text) - 200 and 'Description' not in text[m.end():] and i + 1 < len(page_texts):
|
||||
next_text = page_texts[i + 1].lstrip()
|
||||
if next_text.startswith('Description') or (next_text.startswith('"RDNA') and 'Description' in next_text[:200]):
|
||||
format_headers.append((m.group(1), i, m.start()))
|
||||
|
||||
# parse instruction formats
|
||||
formats: dict[str, list] = {}
|
||||
for fmt_name, page_idx, header_pos in format_headers:
|
||||
if fmt_name in formats: continue
|
||||
text, tables = page_texts[page_idx], page_tables[page_idx]
|
||||
field_pos = text.find('Field Name', header_pos)
|
||||
|
||||
# find fields table with ENCODING (same page or up to 2 pages ahead)
|
||||
fields = None
|
||||
for offset in range(3):
|
||||
if page_idx + offset >= len(pages): break
|
||||
if offset > 0 and has_header_before_fields(page_texts[page_idx + offset]): break
|
||||
for t in page_tables[page_idx + offset] if offset > 0 or field_pos > header_pos else []:
|
||||
if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)) and has_encoding(f):
|
||||
fields = f
|
||||
break
|
||||
if fields: break
|
||||
|
||||
# for modifier formats (no ENCODING), accept first fields table on same page
|
||||
if not fields and field_pos > header_pos:
|
||||
for t in tables:
|
||||
if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)):
|
||||
fields = f
|
||||
break
|
||||
|
||||
if not fields: continue
|
||||
field_names = {f[0] for f in fields}
|
||||
|
||||
# check next pages for continuation fields (tables without ENCODING)
|
||||
for pg_offset in range(1, 3):
|
||||
if page_idx + pg_offset >= len(pages) or has_header_before_fields(page_texts[page_idx + pg_offset]): break
|
||||
for t in page_tables[page_idx + pg_offset]:
|
||||
if is_fields_table(t) and (extra := parse_fields_table(t, fmt_name, enum_names)) and not has_encoding(extra):
|
||||
for ef in extra:
|
||||
if ef[0] not in field_names:
|
||||
fields.append(ef)
|
||||
field_names.add(ef[0])
|
||||
break
|
||||
formats[fmt_name] = fields
|
||||
|
||||
# fix known PDF errors (verified against LLVM test vectors)
|
||||
# SMEM: PDF says DLC=bit14, GLC=bit16 but actual encoding is DLC=bit13, GLC=bit14
|
||||
if 'SMEM' in formats:
|
||||
formats['SMEM'] = [(n, 13 if n == 'DLC' else 14 if n == 'GLC' else h, 13 if n == 'DLC' else 14 if n == 'GLC' else l, e, t)
|
||||
for n, h, l, e, t in formats['SMEM']]
|
||||
|
||||
# generate output
|
||||
def enum_lines(name, items):
|
||||
return [f"class {name}(IntEnum):"] + [f" {n} = {v}" for v, n in sorted(items.items())] + [""]
|
||||
def field_key(f): return order.index(f[0].lower()) if f[0].lower() in order else 1000
|
||||
lines = ["# autogenerated from AMD RDNA3.5 ISA PDF by gen.py - do not edit", "from enum import IntEnum",
|
||||
"from typing import Annotated",
|
||||
"from extra.assembly.rdna3.lib import bits, BitField, Inst32, Inst64, SGPR, VGPR, TTMP as TTMP, s as s, v as v, ttmp as ttmp, SSrc, Src, SImm, Imm, VDSTYEnc, SGPRField, VGPRField",
|
||||
"import functools", ""]
|
||||
lines += enum_lines("SrcEnum", src_enum) + sum([enum_lines(n, ops) for n, ops in sorted(enums.items())], [])
|
||||
# Format-specific field defaults (verified against LLVM test vectors)
|
||||
format_defaults = {'VOP3P': {'opsel_hi': 3, 'opsel_hi2': 1}}
|
||||
lines.append("# instruction formats")
|
||||
for fmt_name, fields in sorted(formats.items()):
|
||||
base = "Inst64" if max(f[1] for f in fields) > 31 or fmt_name == 'VOP3SD' else "Inst32"
|
||||
order = FIELD_ORDER.get(fmt_name, [])
|
||||
lines.append(f"class {fmt_name}({base}):")
|
||||
if enc := next((f for f in fields if f[0] == 'ENCODING'), None):
|
||||
enc_str = f"bits[{enc[1]}:{enc[2]}] == 0b{enc[3]:b}" if enc[1] != enc[2] else f"bits[{enc[1]}] == {enc[3]}"
|
||||
lines.append(f" encoding = {enc_str}")
|
||||
if defaults := format_defaults.get(fmt_name):
|
||||
lines.append(f" _defaults = {defaults}")
|
||||
for name, hi, lo, _, ftype in sorted([f for f in fields if f[0] != 'ENCODING'], key=field_key):
|
||||
# Wrap IntEnum types (ending in Op) with Annotated[BitField, ...] for correct typing
|
||||
if ftype and ftype.endswith('Op'):
|
||||
ann = f":Annotated[BitField, {ftype}]"
|
||||
else:
|
||||
ann = f":{ftype}" if ftype else ""
|
||||
lines.append(f" {name.lower()}{ann} = bits[{hi}]" if hi == lo else f" {name.lower()}{ann} = bits[{hi}:{lo}]")
|
||||
lines.append("")
|
||||
lines.append("# instruction helpers")
|
||||
for cls_name, ops in sorted(enums.items()):
|
||||
fmt = cls_name[:-2]
|
||||
for op_val, name in sorted(ops.items()):
|
||||
seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=2"}.get(fmt, "")
|
||||
tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt, f"{fmt}, {cls_name}")
|
||||
if fmt in formats or fmt in ("GLOBAL", "SCRATCH"):
|
||||
# VOP1/VOP2/VOPC get _e32 suffix, VOP3 promoted ops (< 512) get _e64 suffix
|
||||
if fmt in ("VOP1", "VOP2", "VOPC"):
|
||||
suffix = "_e32"
|
||||
elif fmt == "VOP3" and op_val < 512:
|
||||
suffix = "_e64"
|
||||
else:
|
||||
suffix = ""
|
||||
lines.append(f"{name.lower()}{suffix} = functools.partial({tgt}.{name}{seg})")
|
||||
# export SrcEnum values, but skip DPP8/DPP16 which conflict with class names
|
||||
skip_exports = {'DPP8', 'DPP16'}
|
||||
lines += [""] + [f"{name} = SrcEnum.{name}" for _, name in sorted(src_enum.items()) if name not in skip_exports] + ["OFF = NULL\n"]
|
||||
|
||||
if output_path is not None: pathlib.Path(output_path).write_text('\n'.join(lines))
|
||||
return {"formats": formats, "enums": enums, "src_enum": src_enum}
|
||||
|
||||
if __name__ == "__main__":
|
||||
result = generate("extra/assembly/rdna3/autogen/__init__.py")
|
||||
print(f"generated SrcEnum ({len(result['src_enum'])}) + {len(result['enums'])} opcode enums + {len(result['formats'])} format classes")
|
||||
@@ -1,254 +0,0 @@
|
||||
# library for RDNA3 assembly DSL
|
||||
from __future__ import annotations
|
||||
from enum import IntEnum
|
||||
from typing import overload, Annotated, TypeVar, Generic
|
||||
|
||||
# Bit field DSL
|
||||
class BitField:
|
||||
def __init__(self, hi: int, lo: int, name: str | None = None): self.hi, self.lo, self.name = hi, lo, name
|
||||
def __set_name__(self, owner, name): self.name, self._owner = name, owner
|
||||
def __eq__(self, val: int) -> tuple[BitField, int]: return (self, val) # type: ignore
|
||||
def mask(self) -> int: return (1 << (self.hi - self.lo + 1)) - 1
|
||||
@property
|
||||
def marker(self) -> type | None:
|
||||
# Get marker from Annotated type hint if present
|
||||
import typing
|
||||
if hasattr(self, '_owner') and self.name:
|
||||
hints = typing.get_type_hints(self._owner, include_extras=True)
|
||||
if self.name in hints:
|
||||
hint = hints[self.name]
|
||||
if typing.get_origin(hint) is Annotated:
|
||||
args = typing.get_args(hint)
|
||||
return args[1] if len(args) > 1 else None
|
||||
return None
|
||||
@overload
|
||||
def __get__(self, obj: None, objtype: type) -> BitField: ...
|
||||
@overload
|
||||
def __get__(self, obj: object, objtype: type | None = None) -> int: ...
|
||||
def __get__(self, obj, objtype=None):
|
||||
if obj is None: return self
|
||||
val = unwrap(obj._values.get(self.name, 0))
|
||||
# Convert to IntEnum if marker is an IntEnum subclass
|
||||
if self.marker and isinstance(self.marker, type) and issubclass(self.marker, IntEnum):
|
||||
try: return self.marker(val)
|
||||
except ValueError: pass
|
||||
return val
|
||||
|
||||
class _Bits:
|
||||
def __getitem__(self, key) -> BitField: return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
|
||||
bits = _Bits()
|
||||
|
||||
# Register types
|
||||
class Reg:
|
||||
def __init__(self, idx: int, count: int = 1, hi: bool = False, neg: bool = False): self.idx, self.count, self.hi, self.neg = idx, count, hi, neg
|
||||
def __repr__(self): return f"{self.__class__.__name__.lower()[0]}[{self.idx}]" if self.count == 1 else f"{self.__class__.__name__.lower()[0]}[{self.idx}:{self.idx + self.count}]"
|
||||
def __neg__(self): return self.__class__(self.idx, self.count, self.hi, neg=not self.neg)
|
||||
|
||||
T = TypeVar('T', bound=Reg)
|
||||
class _RegFactory(Generic[T]):
|
||||
def __init__(self, cls: type[T], name: str): self._cls, self._name = cls, name
|
||||
@overload
|
||||
def __getitem__(self, key: int) -> Reg: ...
|
||||
@overload
|
||||
def __getitem__(self, key: slice) -> Reg: ...
|
||||
def __getitem__(self, key: int | slice) -> Reg:
|
||||
return self._cls(key.start, key.stop - key.start + 1) if isinstance(key, slice) else self._cls(key)
|
||||
def __repr__(self): return f"<{self._name} factory>"
|
||||
|
||||
class SGPR(Reg): pass
|
||||
class VGPR(Reg): pass
|
||||
class TTMP(Reg): pass
|
||||
s: _RegFactory[SGPR] = _RegFactory(SGPR, "SGPR")
|
||||
v: _RegFactory[VGPR] = _RegFactory(VGPR, "VGPR")
|
||||
ttmp: _RegFactory[TTMP] = _RegFactory(TTMP, "TTMP")
|
||||
|
||||
# Field type markers (runtime classes for validation)
|
||||
class _SSrc: pass
|
||||
class _Src: pass
|
||||
class _Imm: pass
|
||||
class _SImm: pass
|
||||
class _VDSTYEnc: pass # VOPD vdsty: encoded = actual >> 1, actual = (encoded << 1) | ((vdstx & 1) ^ 1)
|
||||
class _SGPRField: pass
|
||||
class _VGPRField: pass
|
||||
|
||||
# Type aliases for annotations - tells mypy it's a BitField while preserving marker info
|
||||
SSrc = Annotated[BitField, _SSrc]
|
||||
Src = Annotated[BitField, _Src]
|
||||
Imm = Annotated[BitField, _Imm]
|
||||
SImm = Annotated[BitField, _SImm]
|
||||
VDSTYEnc = Annotated[BitField, _VDSTYEnc]
|
||||
SGPRField = Annotated[BitField, _SGPRField]
|
||||
VGPRField = Annotated[BitField, _VGPRField]
|
||||
class RawImm:
|
||||
def __init__(self, val: int): self.val = val
|
||||
def __repr__(self): return f"RawImm({self.val})"
|
||||
def __eq__(self, other): return isinstance(other, RawImm) and self.val == other.val
|
||||
|
||||
def unwrap(val) -> int:
|
||||
return val.val if isinstance(val, RawImm) else val.value if hasattr(val, 'value') else val.idx if hasattr(val, 'idx') else val
|
||||
|
||||
# Encoding helpers
|
||||
FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
|
||||
SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
|
||||
RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata'}
|
||||
|
||||
def _encode_reg(val) -> int:
|
||||
if isinstance(val, TTMP): return 108 + val.idx
|
||||
return val.idx | (0x80 if val.hi else 0)
|
||||
|
||||
def encode_src(val) -> int:
|
||||
if isinstance(val, VGPR): return 256 + _encode_reg(val)
|
||||
if isinstance(val, Reg): return _encode_reg(val)
|
||||
if hasattr(val, 'value'): return val.value
|
||||
if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
|
||||
return 128 + val if isinstance(val, int) and 0 <= val <= 64 else 192 + (-val) if isinstance(val, int) and -16 <= val <= -1 else 255
|
||||
|
||||
# Instruction base class
|
||||
class Inst:
|
||||
_fields: dict[str, BitField]
|
||||
_encoding: tuple[BitField, int] | None = None
|
||||
_defaults: dict[str, int] = {}
|
||||
_values: dict[str, int | RawImm]
|
||||
_words: int # size in 32-bit words, set by decode_program
|
||||
_literal: int | None
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
cls._fields = {n: v[0] if isinstance(v, tuple) else v for n, v in cls.__dict__.items() if isinstance(v, BitField) or (isinstance(v, tuple) and len(v) == 2 and isinstance(v[0], BitField))}
|
||||
if 'encoding' in cls._fields and isinstance(cls.__dict__.get('encoding'), tuple): cls._encoding = cls.__dict__['encoding']
|
||||
|
||||
def __init__(self, *args, literal: int | None = None, **kwargs):
|
||||
self._values, self._literal = dict(self._defaults), literal
|
||||
# Map positional args to field names
|
||||
field_names = [n for n in self._fields if n != 'encoding']
|
||||
orig_args = dict(zip(field_names, args))
|
||||
orig_args.update(kwargs)
|
||||
self._values.update(orig_args)
|
||||
# Validate register counts for SMEM instructions (before encoding)
|
||||
if self.__class__.__name__ == 'SMEM':
|
||||
op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
|
||||
if op_val is not None:
|
||||
if hasattr(op_val, 'value'): op_val = op_val.value
|
||||
expected_cnt = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val)
|
||||
sdata_val = orig_args.get('sdata')
|
||||
if expected_cnt is not None and isinstance(sdata_val, Reg) and sdata_val.count != expected_cnt:
|
||||
raise ValueError(f"SMEM op {op_val} expects {expected_cnt} registers, got {sdata_val.count}")
|
||||
# Validate register counts for SOP1 instructions (b32 = 1 reg, b64 = 2 regs)
|
||||
if self.__class__.__name__ == 'SOP1':
|
||||
op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
|
||||
if op_val is not None and hasattr(op_val, 'name'):
|
||||
expected = 2 if op_val.name.endswith('_B64') else 1
|
||||
sdst_val, ssrc0_val = orig_args.get('sdst'), orig_args.get('ssrc0')
|
||||
if isinstance(sdst_val, Reg) and sdst_val.count != expected:
|
||||
raise ValueError(f"SOP1 {op_val.name} expects {expected} destination register(s), got {sdst_val.count}")
|
||||
if isinstance(ssrc0_val, Reg) and ssrc0_val.count != expected:
|
||||
raise ValueError(f"SOP1 {op_val.name} expects {expected} source register(s), got {ssrc0_val.count}")
|
||||
# Type check and encode values
|
||||
for name, val in list(self._values.items()):
|
||||
if name == 'encoding': continue
|
||||
# For RawImm, only process RAW_FIELDS to unwrap to int
|
||||
if isinstance(val, RawImm):
|
||||
if name in RAW_FIELDS: self._values[name] = val.val
|
||||
continue
|
||||
field = self._fields.get(name)
|
||||
marker = field.marker if field else None
|
||||
# Type validation
|
||||
if marker is _SGPRField:
|
||||
if isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
|
||||
if not isinstance(val, (SGPR, TTMP, int, RawImm)): raise TypeError(f"field '{name}' requires SGPR, got {type(val).__name__}")
|
||||
if marker is _VGPRField:
|
||||
if not isinstance(val, VGPR): raise TypeError(f"field '{name}' requires VGPR, got {type(val).__name__}")
|
||||
if marker is _SSrc and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires scalar source, got VGPR")
|
||||
# Encode source fields as RawImm for consistent disassembly
|
||||
if name in SRC_FIELDS:
|
||||
encoded = encode_src(val)
|
||||
self._values[name] = RawImm(encoded)
|
||||
# Handle negation modifier for VOP3 instructions
|
||||
if isinstance(val, Reg) and val.neg and 'neg' in self._fields:
|
||||
neg_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
|
||||
cur_neg = self._values.get('neg', 0)
|
||||
self._values['neg'] = (cur_neg.val if isinstance(cur_neg, RawImm) else cur_neg) | neg_bit
|
||||
# Track literal value if needed (encoded as 255)
|
||||
if encoded == 255 and self._literal is None and isinstance(val, int) and not isinstance(val, IntEnum):
|
||||
self._literal = val
|
||||
elif encoded == 255 and self._literal is None and isinstance(val, float):
|
||||
import struct
|
||||
self._literal = struct.unpack('<I', struct.pack('<f', val))[0]
|
||||
# Encode raw register fields for consistent repr
|
||||
elif name in RAW_FIELDS:
|
||||
if isinstance(val, Reg): self._values[name] = _encode_reg(val)
|
||||
elif hasattr(val, 'value'): self._values[name] = val.value # IntEnum like SrcEnum.NULL
|
||||
# Encode sbase (divided by 2) and srsrc/ssamp (divided by 4)
|
||||
elif name == 'sbase' and isinstance(val, Reg):
|
||||
self._values[name] = val.idx // 2
|
||||
elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg):
|
||||
self._values[name] = val.idx // 4
|
||||
# VOPD vdsty: encode as actual >> 1 (constraint: vdsty parity must be opposite of vdstx)
|
||||
elif marker is _VDSTYEnc and isinstance(val, VGPR):
|
||||
self._values[name] = val.idx >> 1
|
||||
|
||||
def _encode_field(self, name: str, val) -> int:
|
||||
if isinstance(val, RawImm): return val.val
|
||||
if name in {'srsrc', 'ssamp'}: return val.idx // 4 if isinstance(val, Reg) else val
|
||||
if name == 'sbase': return val.idx // 2 if isinstance(val, Reg) else val
|
||||
if name in RAW_FIELDS: return _encode_reg(val) if isinstance(val, Reg) else val
|
||||
if isinstance(val, Reg) or name in SRC_FIELDS: return encode_src(val)
|
||||
return val.value if hasattr(val, 'value') else val
|
||||
|
||||
def to_int(self) -> int:
|
||||
word = (self._encoding[1] & self._encoding[0].mask()) << self._encoding[0].lo if self._encoding else 0
|
||||
for n, bf in self._fields.items():
|
||||
if n != 'encoding' and n in self._values: word |= (self._encode_field(n, self._values[n]) & bf.mask()) << bf.lo
|
||||
return word
|
||||
|
||||
def _get_literal(self) -> int | None:
|
||||
for n in SRC_FIELDS:
|
||||
if n in self._values and not isinstance(v := self._values[n], RawImm) and isinstance(v, int) and not isinstance(v, IntEnum) and not (0 <= v <= 64 or -16 <= v <= -1): return v
|
||||
return None
|
||||
|
||||
def to_bytes(self) -> bytes:
|
||||
result = self.to_int().to_bytes(self._size(), 'little')
|
||||
return result + (lit & 0xffffffff).to_bytes(4, 'little') if (lit := self._get_literal() or getattr(self, '_literal', None)) else result
|
||||
|
||||
@classmethod
|
||||
def _size(cls) -> int: return 4 if issubclass(cls, Inst32) else 8
|
||||
def size(self) -> int: return self._size() + (4 if self._literal is not None else 0)
|
||||
|
||||
@classmethod
|
||||
def from_int(cls, word: int):
|
||||
inst = object.__new__(cls)
|
||||
inst._values = {n: RawImm(v) if n in SRC_FIELDS else v for n, bf in cls._fields.items() if n != 'encoding' for v in [(word >> bf.lo) & bf.mask()]}
|
||||
inst._literal = None
|
||||
return inst
|
||||
|
||||
@classmethod
|
||||
def from_bytes(cls, data: bytes):
|
||||
inst = cls.from_int(int.from_bytes(data[:cls._size()], 'little'))
|
||||
op_val = inst._values.get('op', 0)
|
||||
has_literal = cls.__name__ == 'VOP2' and op_val in (44, 45, 55, 56)
|
||||
has_literal = has_literal or (cls.__name__ == 'SOP2' and op_val in (69, 70))
|
||||
for n in SRC_FIELDS:
|
||||
if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255: has_literal = True
|
||||
if has_literal and len(data) >= cls._size() + 4: inst._literal = int.from_bytes(data[cls._size():cls._size()+4], 'little')
|
||||
return inst
|
||||
|
||||
def __repr__(self):
|
||||
# Use _fields order and exclude fields that are 0/default (for consistent repr after roundtrip)
|
||||
def is_zero(v): return (isinstance(v, int) and v == 0) or (isinstance(v, VGPR) and v.idx == 0 and v.count == 1)
|
||||
items = [(k, self._values[k]) for k in self._fields if k in self._values and k != 'encoding'
|
||||
and not (is_zero(self._values[k]) and k not in {'op'})]
|
||||
lit = f", literal={hex(self._literal)}" if self._literal is not None else ""
|
||||
return f"{self.__class__.__name__}({', '.join(f'{k}={v}' for k, v in items)}{lit})"
|
||||
|
||||
def __eq__(self, other):
|
||||
if not isinstance(other, Inst): return NotImplemented
|
||||
return self.__class__ == other.__class__ and self._values == other._values and self._literal == other._literal
|
||||
|
||||
def __hash__(self): return hash((self.__class__.__name__, tuple(sorted((k, repr(v)) for k, v in self._values.items())), self._literal))
|
||||
|
||||
def disasm(self) -> str:
|
||||
from extra.assembly.rdna3.asm import disasm
|
||||
return disasm(self)
|
||||
|
||||
class Inst32(Inst): pass
|
||||
class Inst64(Inst): pass
|
||||
@@ -1,294 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Benchmark comparing Python vs Rust RDNA3 emulators on synthetic and real tinygrad kernels."""
|
||||
import ctypes, time, os, struct, cProfile, pstats, io
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
# Set AMD=1 before importing tinygrad
|
||||
os.environ["AMD"] = "1"
|
||||
|
||||
from extra.assembly.rdna3.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program, step_wave, WaveState, WAVE_SIZE
|
||||
|
||||
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
|
||||
if not REMU_PATH.exists():
|
||||
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.dylib"
|
||||
|
||||
def get_rust_remu():
|
||||
"""Load the Rust libremu shared library."""
|
||||
if not REMU_PATH.exists(): return None
|
||||
remu = ctypes.CDLL(str(REMU_PATH))
|
||||
remu.run_asm.restype = ctypes.c_int32
|
||||
remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
|
||||
ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
|
||||
return remu
|
||||
|
||||
def count_instructions(kernel: bytes) -> int:
|
||||
"""Count instructions in a kernel."""
|
||||
return len(decode_program(kernel))
|
||||
|
||||
def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = None):
|
||||
"""Allocate buffers and return args pointer + valid ranges."""
|
||||
if init_data is None: init_data = {}
|
||||
buffers = []
|
||||
for i, size in enumerate(buf_sizes):
|
||||
padded = ((size + 15) // 16) * 16 + 16
|
||||
data = init_data.get(i, b'\x00' * padded)
|
||||
data_list = list(data) + [0] * (padded - len(data))
|
||||
buf = (ctypes.c_uint8 * padded)(*data_list[:padded])
|
||||
buffers.append(buf)
|
||||
args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
|
||||
args_ptr = ctypes.addressof(args)
|
||||
ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
|
||||
ranges.add((args_ptr, ctypes.sizeof(args)))
|
||||
return buffers, args, args_ptr, ranges
|
||||
|
||||
def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, iterations: int = 5):
|
||||
"""Benchmark an emulator and return average time."""
|
||||
gx, gy, gz = global_size
|
||||
lx, ly, lz = local_size
|
||||
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
|
||||
lib_ptr = ctypes.addressof(kernel_buf)
|
||||
|
||||
# Warmup
|
||||
run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
|
||||
|
||||
# Timed runs
|
||||
times = []
|
||||
for _ in range(iterations):
|
||||
start = time.perf_counter()
|
||||
result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
|
||||
end = time.perf_counter()
|
||||
if result != 0:
|
||||
print(f" {name} returned error: {result}")
|
||||
return None
|
||||
times.append(end - start)
|
||||
|
||||
return sum(times) / len(times)
|
||||
|
||||
def create_synthetic_kernel(n_ops: int) -> bytes:
|
||||
"""Create a synthetic kernel with n_ops vector operations."""
|
||||
instructions = []
|
||||
# VOP2 instructions: v_add_f32, v_mul_f32, v_max_f32, v_min_f32
|
||||
ops = [
|
||||
(0b0000011 << 25) | (1 << 17) | (0 << 9) | 256, # v_add_f32 v0, v0, v1
|
||||
(0b0001000 << 25) | (1 << 17) | (0 << 9) | 256, # v_mul_f32 v0, v0, v1
|
||||
(0b0010000 << 25) | (1 << 17) | (0 << 9) | 256, # v_max_f32 v0, v0, v1
|
||||
(0b0001111 << 25) | (1 << 17) | (0 << 9) | 256, # v_min_f32 v0, v0, v1
|
||||
]
|
||||
for i in range(n_ops):
|
||||
instructions.append(ops[i % len(ops)])
|
||||
# S_ENDPGM
|
||||
instructions.append((0b101111111 << 23) | (48 << 16) | 0)
|
||||
return b''.join(struct.pack('<I', inst) for inst in instructions)
|
||||
|
||||
def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes]] | None:
|
||||
"""Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data)."""
|
||||
try:
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.runtime.support.elf import elf_loader
|
||||
import numpy as np
|
||||
np.random.seed(42)
|
||||
|
||||
ops = {
|
||||
"add": lambda: Tensor.empty(1024) + Tensor.empty(1024),
|
||||
"mul": lambda: Tensor.empty(1024) * Tensor.empty(1024),
|
||||
"matmul_small": lambda: Tensor.empty(16, 16) @ Tensor.empty(16, 16),
|
||||
"matmul_medium": lambda: Tensor.empty(64, 64) @ Tensor.empty(64, 64),
|
||||
"reduce_sum": lambda: Tensor.empty(4096).sum(),
|
||||
"reduce_max": lambda: Tensor.empty(4096).max(),
|
||||
"softmax": lambda: Tensor.empty(256).softmax(),
|
||||
"layernorm": lambda: Tensor.empty(32, 64).layernorm(),
|
||||
"conv2d": lambda: Tensor.empty(1, 4, 16, 16).conv2d(Tensor.empty(4, 4, 3, 3)),
|
||||
"gelu": lambda: Tensor.empty(1024).gelu(),
|
||||
"exp": lambda: Tensor.empty(1024).exp(),
|
||||
"sin": lambda: Tensor.empty(1024).sin(),
|
||||
}
|
||||
|
||||
if op_name not in ops: return None
|
||||
out = ops[op_name]()
|
||||
sched = out.schedule()
|
||||
|
||||
for ei in sched:
|
||||
lowered = ei.lower()
|
||||
if ei.ast.op.name == 'SINK' and lowered.prg and lowered.prg.p.lib:
|
||||
lib = bytes(lowered.prg.p.lib)
|
||||
_, sections, _ = elf_loader(lib)
|
||||
for sec in sections:
|
||||
if sec.name == '.text':
|
||||
buf_sizes = [b.nbytes for b in lowered.bufs]
|
||||
# Get initial data from numpy arrays if available
|
||||
buf_data = {}
|
||||
for i, buf in enumerate(lowered.bufs):
|
||||
if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
|
||||
try: buf_data[i] = bytes(buf.base._buf)
|
||||
except: pass
|
||||
return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data)
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f" Error getting kernel: {e}")
|
||||
return None
|
||||
|
||||
def profile_python_emu(kernel: bytes, global_size, local_size, args_ptr, n_runs: int = 1):
|
||||
"""Profile the Python emulator to find bottlenecks."""
|
||||
gx, gy, gz = global_size
|
||||
lx, ly, lz = local_size
|
||||
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
|
||||
lib_ptr = ctypes.addressof(kernel_buf)
|
||||
|
||||
pr = cProfile.Profile()
|
||||
pr.enable()
|
||||
for _ in range(n_runs):
|
||||
python_run_asm(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
|
||||
pr.disable()
|
||||
|
||||
s = io.StringIO()
|
||||
ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
|
||||
ps.print_stats(20)
|
||||
return s.getvalue()
|
||||
|
||||
def measure_step_rate(kernel: bytes, n_steps: int = 10000) -> float:
|
||||
"""Measure raw step_wave() performance (steps per second)."""
|
||||
program = decode_program(kernel)
|
||||
if not program: return 0.0
|
||||
|
||||
st = WaveState()
|
||||
st.exec_mask = 0xffffffff
|
||||
lds = bytearray(65536)
|
||||
n_lanes = 32
|
||||
|
||||
# Reset PC for each measurement
|
||||
start = time.perf_counter()
|
||||
for _ in range(n_steps):
|
||||
st.pc = 0
|
||||
while st.pc in program:
|
||||
result = step_wave(program, st, lds, n_lanes)
|
||||
if result == -1: break
|
||||
elapsed = time.perf_counter() - start
|
||||
return n_steps / elapsed if elapsed > 0 else 0
|
||||
|
||||
# Test configurations
|
||||
SYNTHETIC_TESTS = [
|
||||
("synthetic_10ops", 10, (1, 1, 1), (32, 1, 1)),
|
||||
("synthetic_100ops", 100, (1, 1, 1), (32, 1, 1)),
|
||||
("synthetic_500ops", 500, (1, 1, 1), (32, 1, 1)),
|
||||
("synthetic_100ops_4wg", 100, (4, 1, 1), (32, 1, 1)),
|
||||
("synthetic_100ops_16wg", 100, (16, 1, 1), (32, 1, 1)),
|
||||
]
|
||||
|
||||
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
|
||||
|
||||
def main():
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
|
||||
parser.add_argument("--profile", action="store_true", help="Profile Python emulator")
|
||||
parser.add_argument("--synthetic-only", action="store_true", help="Only run synthetic tests")
|
||||
parser.add_argument("--tinygrad-only", action="store_true", help="Only run tinygrad tests")
|
||||
parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
|
||||
args = parser.parse_args()
|
||||
|
||||
rust_remu = get_rust_remu()
|
||||
if rust_remu is None:
|
||||
print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
|
||||
print("Running Python-only benchmarks...\n")
|
||||
|
||||
print("=" * 90)
|
||||
print("RDNA3 Emulator Benchmark: Python vs Rust")
|
||||
print("=" * 90)
|
||||
|
||||
results = []
|
||||
|
||||
# Synthetic workloads
|
||||
if not args.tinygrad_only:
|
||||
print("\n[SYNTHETIC WORKLOADS]")
|
||||
print("-" * 90)
|
||||
|
||||
for name, n_ops, global_size, local_size in SYNTHETIC_TESTS:
|
||||
kernel = create_synthetic_kernel(n_ops)
|
||||
n_insts = count_instructions(kernel)
|
||||
n_workgroups = global_size[0] * global_size[1] * global_size[2]
|
||||
n_threads = local_size[0] * local_size[1] * local_size[2]
|
||||
total_work = n_insts * n_workgroups * n_threads
|
||||
|
||||
print(f"\n{name}: {n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
|
||||
|
||||
buf_sizes = [4096]
|
||||
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes)
|
||||
set_valid_mem_ranges(ranges)
|
||||
|
||||
# Benchmark
|
||||
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
|
||||
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
|
||||
|
||||
if py_time:
|
||||
py_rate = total_work / py_time / 1e6
|
||||
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
|
||||
if rust_time:
|
||||
rust_rate = total_work / rust_time / 1e6
|
||||
speedup = py_time / rust_time if py_time else 0
|
||||
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
|
||||
|
||||
results.append(("synthetic", name, n_insts, n_workgroups, py_time, rust_time))
|
||||
|
||||
# Tinygrad kernels
|
||||
if not args.synthetic_only:
|
||||
print("\n[TINYGRAD KERNELS]")
|
||||
print("-" * 90)
|
||||
|
||||
for op_name in TINYGRAD_TESTS:
|
||||
print(f"\n{op_name}:", end=" ", flush=True)
|
||||
kernel_info = get_tinygrad_kernel(op_name)
|
||||
if kernel_info is None:
|
||||
print("failed to compile")
|
||||
continue
|
||||
|
||||
kernel, global_size, local_size, buf_sizes, buf_data = kernel_info
|
||||
n_insts = count_instructions(kernel)
|
||||
n_workgroups = global_size[0] * global_size[1] * global_size[2]
|
||||
n_threads = local_size[0] * local_size[1] * local_size[2]
|
||||
total_work = n_insts * n_workgroups * n_threads
|
||||
|
||||
print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
|
||||
|
||||
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
|
||||
set_valid_mem_ranges(ranges)
|
||||
|
||||
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
|
||||
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
|
||||
|
||||
if py_time:
|
||||
py_rate = total_work / py_time / 1e6
|
||||
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
|
||||
if rust_time:
|
||||
rust_rate = total_work / rust_time / 1e6
|
||||
speedup = py_time / rust_time if py_time else 0
|
||||
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
|
||||
|
||||
results.append(("tinygrad", op_name, n_insts, n_workgroups, py_time, rust_time))
|
||||
|
||||
# Optional profiling
|
||||
if args.profile and py_time:
|
||||
print("\n [PROFILE - Top 10 functions]")
|
||||
profile_output = profile_python_emu(kernel, global_size, local_size, args_ptr)
|
||||
for line in profile_output.split('\n')[5:15]:
|
||||
if line.strip(): print(f" {line}")
|
||||
|
||||
# Summary table
|
||||
print("\n" + "=" * 90)
|
||||
print("SUMMARY")
|
||||
print("=" * 90)
|
||||
print(f"{'Type':<10} {'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
|
||||
print("-" * 90)
|
||||
|
||||
for test_type, name, n_insts, n_wgs, py_time, rust_time in results:
|
||||
py_ms = f"{py_time*1000:.3f}" if py_time else "error"
|
||||
if rust_time:
|
||||
rust_ms = f"{rust_time*1000:.3f}"
|
||||
speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
|
||||
else:
|
||||
rust_ms, speedup = "N/A", "N/A"
|
||||
print(f"{test_type:<10} {name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,845 +0,0 @@
|
||||
# Unit tests for RDNA3 Python emulator
|
||||
import unittest
|
||||
import ctypes
|
||||
import struct
|
||||
import math
|
||||
from extra.assembly.rdna3.emu import (
|
||||
WaveState, decode_program, exec_wave, exec_workgroup, run_asm,
|
||||
i32, f32, sext, WAVE_SIZE, set_valid_mem_ranges
|
||||
)
|
||||
from extra.assembly.rdna3.autogen import *
|
||||
from extra.assembly.rdna3.lib import RawImm
|
||||
|
||||
def run_kernel(kernel: bytes, n_threads: int = 1, n_outputs: int = 1) -> list[int]:
|
||||
"""Helper to run a kernel and return output values."""
|
||||
output = (ctypes.c_uint32 * (n_threads * n_outputs))(*[0xdead] * (n_threads * n_outputs))
|
||||
output_ptr = ctypes.addressof(output)
|
||||
args = (ctypes.c_uint64 * 1)(output_ptr)
|
||||
args_ptr = ctypes.addressof(args)
|
||||
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
|
||||
kernel_ptr = ctypes.addressof(kernel_buf)
|
||||
# Register valid memory ranges for bounds checking
|
||||
set_valid_mem_ranges({
|
||||
(output_ptr, ctypes.sizeof(output)),
|
||||
(args_ptr, ctypes.sizeof(args)),
|
||||
(kernel_ptr, len(kernel)),
|
||||
})
|
||||
result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
|
||||
assert result == 0, f"run_asm failed with {result}"
|
||||
return [output[i] for i in range(n_threads * n_outputs)]
|
||||
|
||||
def make_store_kernel(setup_instrs: list, store_vreg: int = 1) -> bytes:
|
||||
"""Create a kernel that runs setup instructions then stores v[store_vreg] to output[tid]."""
|
||||
kernel = b''
|
||||
# Load output pointer
|
||||
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
|
||||
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
|
||||
# Run setup instructions
|
||||
for instr in setup_instrs:
|
||||
kernel += instr.to_bytes()
|
||||
# Compute offset: v3 = tid * 4
|
||||
kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
|
||||
# Store result
|
||||
kernel += global_store_b32(addr=v[3], data=v[store_vreg], saddr=s[2]).to_bytes()
|
||||
kernel += s_endpgm().to_bytes()
|
||||
return kernel
|
||||
|
||||
class TestScalarOps(unittest.TestCase):
|
||||
def test_s_mov_b32(self):
|
||||
state = WaveState()
|
||||
kernel = s_mov_b32(s[5], 42).to_bytes() + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[5], 42)
|
||||
|
||||
def test_s_add_u32(self):
|
||||
state = WaveState()
|
||||
state.sgpr[0], state.sgpr[1] = 100, 50
|
||||
kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[2], 150)
|
||||
self.assertEqual(state.scc, 0) # no carry
|
||||
|
||||
def test_s_add_u32_carry(self):
|
||||
state = WaveState()
|
||||
state.sgpr[0], state.sgpr[1] = 0xffffffff, 1
|
||||
kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[2], 0)
|
||||
self.assertEqual(state.scc, 1) # carry
|
||||
|
||||
def test_s_sub_u32(self):
|
||||
state = WaveState()
|
||||
state.sgpr[0], state.sgpr[1] = 100, 30
|
||||
kernel = s_sub_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[2], 70)
|
||||
self.assertEqual(state.scc, 0) # no borrow
|
||||
|
||||
def test_s_and_b32(self):
|
||||
state = WaveState()
|
||||
state.sgpr[0], state.sgpr[1] = 0xff00, 0x0ff0
|
||||
kernel = s_and_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[2], 0x0f00)
|
||||
|
||||
def test_s_or_b32(self):
|
||||
state = WaveState()
|
||||
state.sgpr[0], state.sgpr[1] = 0xff00, 0x00ff
|
||||
kernel = s_or_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[2], 0xffff)
|
||||
|
||||
def test_s_lshl_b32(self):
|
||||
state = WaveState()
|
||||
state.sgpr[0], state.sgpr[1] = 1, 4
|
||||
kernel = s_lshl_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[2], 16)
|
||||
|
||||
def test_s_lshr_b32(self):
|
||||
state = WaveState()
|
||||
state.sgpr[0], state.sgpr[1] = 256, 4
|
||||
kernel = s_lshr_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[2], 16)
|
||||
|
||||
def test_s_mul_i32(self):
|
||||
state = WaveState()
|
||||
state.sgpr[0], state.sgpr[1] = 7, 6
|
||||
kernel = s_mul_i32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[2], 42)
|
||||
|
||||
def test_s_cmp_eq_u32(self):
|
||||
state = WaveState()
|
||||
state.sgpr[0], state.sgpr[1] = 42, 42
|
||||
kernel = s_cmp_eq_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.scc, 1)
|
||||
|
||||
def test_s_cmp_lg_u32(self):
|
||||
state = WaveState()
|
||||
state.sgpr[0], state.sgpr[1] = 42, 43
|
||||
kernel = s_cmp_lg_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.scc, 1)
|
||||
|
||||
class TestVectorOps(unittest.TestCase):
|
||||
def test_v_mov_b32(self):
|
||||
kernel = make_store_kernel([v_mov_b32_e32(v[1], 42)])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out, [42])
|
||||
|
||||
def test_v_add_nc_u32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 10),
|
||||
v_mov_b32_e32(v[2], 32),
|
||||
v_add_nc_u32_e32(v[1], v[1], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out, [42])
|
||||
|
||||
def test_v_sub_nc_u32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 50),
|
||||
v_mov_b32_e32(v[2], 8),
|
||||
v_sub_nc_u32_e32(v[1], v[1], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out, [42])
|
||||
|
||||
def test_v_mul_lo_u32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 6),
|
||||
v_mov_b32_e32(v[2], 7),
|
||||
v_mul_lo_u32(v[1], v[1], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out, [42])
|
||||
|
||||
def test_v_and_b32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 0xff0f),
|
||||
v_mov_b32_e32(v[2], 0x0fff),
|
||||
v_and_b32_e32(v[1], v[1], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out, [0x0f0f])
|
||||
|
||||
def test_v_or_b32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 0xf000),
|
||||
v_mov_b32_e32(v[2], 0x000f),
|
||||
v_or_b32_e32(v[1], v[1], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out, [0xf00f])
|
||||
|
||||
def test_v_lshlrev_b32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 1),
|
||||
v_lshlrev_b32_e32(v[1], 5, v[1]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out, [32])
|
||||
|
||||
def test_v_lshrrev_b32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 128),
|
||||
v_lshrrev_b32_e32(v[1], 3, v[1]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out, [16])
|
||||
|
||||
def test_v_add_f32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(1.5)),
|
||||
v_mov_b32_e32(v[2], i32(2.5)),
|
||||
v_add_f32_e32(v[1], v[1], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(f32(out[0]), 4.0)
|
||||
|
||||
def test_v_mul_f32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(3.0)),
|
||||
v_mov_b32_e32(v[2], i32(4.0)),
|
||||
v_mul_f32_e32(v[1], v[1], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(f32(out[0]), 12.0)
|
||||
|
||||
def test_v_max_f32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(3.0)),
|
||||
v_mov_b32_e32(v[2], i32(5.0)),
|
||||
v_max_f32_e32(v[1], v[1], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(f32(out[0]), 5.0)
|
||||
|
||||
def test_v_min_f32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(3.0)),
|
||||
v_mov_b32_e32(v[2], i32(5.0)),
|
||||
v_min_f32_e32(v[1], v[1], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(f32(out[0]), 3.0)
|
||||
|
||||
class TestThreading(unittest.TestCase):
|
||||
def test_thread_id(self):
|
||||
"""Each thread should get its own thread ID in v0."""
|
||||
kernel = make_store_kernel([v_mov_b32_e32(v[1], v[0])], store_vreg=1)
|
||||
out = run_kernel(kernel, n_threads=4)
|
||||
self.assertEqual(out, [0, 1, 2, 3])
|
||||
|
||||
def test_thread_local_ops(self):
|
||||
"""Each thread computes tid * 10."""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[2], 10),
|
||||
v_mul_lo_u32(v[1], v[0], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=4)
|
||||
self.assertEqual(out, [0, 10, 20, 30])
|
||||
|
||||
def test_exec_mask(self):
|
||||
"""Test that exec mask controls which lanes execute."""
|
||||
kernel = b''
|
||||
kernel += s_load_b64(s[2:3], s[0:1], 0, soffset=NULL).to_bytes()
|
||||
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
|
||||
kernel += v_mov_b32_e32(v[1], 100).to_bytes() # default value
|
||||
kernel += s_mov_b32(EXEC_LO, 0b0101).to_bytes() # only lanes 0 and 2
|
||||
kernel += v_mov_b32_e32(v[1], 42).to_bytes() # only for active lanes
|
||||
kernel += s_mov_b32(EXEC_LO, 0xf).to_bytes() # restore all lanes
|
||||
kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
|
||||
kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[2]).to_bytes()
|
||||
kernel += s_endpgm().to_bytes()
|
||||
out = run_kernel(kernel, n_threads=4)
|
||||
self.assertEqual(out, [42, 100, 42, 100])
|
||||
|
||||
class TestBranching(unittest.TestCase):
|
||||
def test_s_branch(self):
|
||||
"""Test unconditional branch."""
|
||||
state = WaveState()
|
||||
kernel = b''
|
||||
kernel += s_mov_b32(s[0], 1).to_bytes()
|
||||
kernel += s_branch(1).to_bytes() # skip next instruction
|
||||
kernel += s_mov_b32(s[0], 2).to_bytes() # should be skipped
|
||||
kernel += s_mov_b32(s[1], 3).to_bytes()
|
||||
kernel += s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[0], 1) # not overwritten
|
||||
self.assertEqual(state.sgpr[1], 3)
|
||||
|
||||
def test_s_cbranch_scc0(self):
|
||||
"""Test conditional branch on SCC=0."""
|
||||
state = WaveState()
|
||||
state.scc = 0
|
||||
kernel = b''
|
||||
kernel += s_mov_b32(s[0], 1).to_bytes()
|
||||
kernel += s_cbranch_scc0(1).to_bytes() # branch if scc=0
|
||||
kernel += s_mov_b32(s[0], 2).to_bytes() # should be skipped
|
||||
kernel += s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[0], 1)
|
||||
|
||||
def test_s_cbranch_scc1(self):
|
||||
"""Test conditional branch on SCC=1."""
|
||||
state = WaveState()
|
||||
state.scc = 1
|
||||
kernel = b''
|
||||
kernel += s_mov_b32(s[0], 1).to_bytes()
|
||||
kernel += s_cbranch_scc1(1).to_bytes() # branch if scc=1
|
||||
kernel += s_mov_b32(s[0], 2).to_bytes() # should be skipped
|
||||
kernel += s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.sgpr[0], 1)
|
||||
|
||||
def test_unknown_sopp_opcode(self):
|
||||
"""Regression test: unknown SOPP opcodes should be ignored, not crash."""
|
||||
state = WaveState()
|
||||
# Create a raw SOPP instruction with opcode 8 (undefined in our enum)
|
||||
# SOPP format: bits[31:23] = 0b101111111, bits[22:16] = op, bits[15:0] = simm16
|
||||
unknown_sopp = (0b101111111 << 23) | (8 << 16) | 0 # op=8, simm16=0
|
||||
kernel = unknown_sopp.to_bytes(4, 'little') + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
# Should not raise an exception
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
|
||||
class TestMemory(unittest.TestCase):
|
||||
def test_global_load_store(self):
|
||||
"""Test global load followed by store."""
|
||||
# Create input buffer
|
||||
input_buf = (ctypes.c_uint32 * 4)(10, 20, 30, 40)
|
||||
input_ptr = ctypes.addressof(input_buf)
|
||||
output_buf = (ctypes.c_uint32 * 4)(*[0]*4)
|
||||
output_ptr = ctypes.addressof(output_buf)
|
||||
args = (ctypes.c_uint64 * 2)(output_ptr, input_ptr)
|
||||
args_ptr = ctypes.addressof(args)
|
||||
|
||||
# Kernel: load from input[tid], add 1, store to output[tid]
|
||||
kernel = b''
|
||||
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes() # output ptr
|
||||
kernel += s_load_b64(s[4:5], s[0:1], soffset=NULL, offset=8).to_bytes() # input ptr
|
||||
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
|
||||
kernel += v_lshlrev_b32_e32(v[2], 2, v[0]).to_bytes() # offset = tid * 4
|
||||
kernel += global_load_b32(vdst=v[1], addr=v[2], saddr=s[4]).to_bytes()
|
||||
kernel += s_waitcnt(vmcnt=0).to_bytes()
|
||||
kernel += v_add_nc_u32_e32(v[1], 1, v[1]).to_bytes() # add 1
|
||||
kernel += global_store_b32(addr=v[2], data=v[1], saddr=s[2]).to_bytes()
|
||||
kernel += s_endpgm().to_bytes()
|
||||
|
||||
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
|
||||
kernel_ptr = ctypes.addressof(kernel_buf)
|
||||
set_valid_mem_ranges({
|
||||
(input_ptr, ctypes.sizeof(input_buf)),
|
||||
(output_ptr, ctypes.sizeof(output_buf)),
|
||||
(args_ptr, ctypes.sizeof(args)),
|
||||
(kernel_ptr, len(kernel)),
|
||||
})
|
||||
result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, 4, 1, 1, args_ptr)
|
||||
self.assertEqual(result, 0)
|
||||
self.assertEqual([output_buf[i] for i in range(4)], [11, 21, 31, 41])
|
||||
|
||||
class TestFloatOps(unittest.TestCase):
|
||||
def test_v_rcp_f32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(4.0)),
|
||||
v_rcp_f32_e32(v[1], v[1]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertAlmostEqual(f32(out[0]), 0.25, places=5)
|
||||
|
||||
def test_v_sqrt_f32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(16.0)),
|
||||
v_sqrt_f32_e32(v[1], v[1]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertAlmostEqual(f32(out[0]), 4.0, places=5)
|
||||
|
||||
def test_v_floor_f32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(3.7)),
|
||||
v_floor_f32_e32(v[1], v[1]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(f32(out[0]), 3.0)
|
||||
|
||||
def test_v_ceil_f32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(3.2)),
|
||||
v_ceil_f32_e32(v[1], v[1]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(f32(out[0]), 4.0)
|
||||
|
||||
def test_v_cvt_f32_i32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 42),
|
||||
v_cvt_f32_i32_e32(v[1], v[1]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(f32(out[0]), 42.0)
|
||||
|
||||
def test_v_cvt_i32_f32(self):
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(42.9)),
|
||||
v_cvt_i32_f32_e32(v[1], v[1]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out[0], 42)
|
||||
|
||||
class TestVOP3(unittest.TestCase):
|
||||
def test_v_fma_f32(self):
|
||||
"""Test fused multiply-add: a*b + c"""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(2.0)),
|
||||
v_mov_b32_e32(v[2], i32(3.0)),
|
||||
v_mov_b32_e32(v[4], i32(4.0)),
|
||||
v_fma_f32(v[1], v[1], v[2], v[4]), # 2*3+4 = 10
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(f32(out[0]), 10.0)
|
||||
|
||||
def test_v_add3_u32(self):
|
||||
"""Test 3-operand add."""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 10),
|
||||
v_mov_b32_e32(v[2], 20),
|
||||
v_mov_b32_e32(v[4], 12),
|
||||
v_add3_u32(v[1], v[1], v[2], v[4]), # 10+20+12 = 42
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out[0], 42)
|
||||
|
||||
def test_v_neg_modifier(self):
|
||||
"""Test VOP3 negation modifier."""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(5.0)),
|
||||
v_mov_b32_e32(v[2], i32(3.0)),
|
||||
# v_add_f32 with neg on src1: 5 + (-3) = 2
|
||||
v_add_f32_e64(v[1], v[1], v[2], neg=0b010),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(f32(out[0]), 2.0)
|
||||
|
||||
def test_v_ldexp_f32(self):
|
||||
"""Regression test: V_LDEXP_F32 used by exp()."""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(1.5)),
|
||||
v_mov_b32_e32(v[2], 3), # exponent
|
||||
v_ldexp_f32(v[1], v[1], v[2]), # 1.5 * 2^3 = 12.0
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(f32(out[0]), 12.0)
|
||||
|
||||
def test_v_xad_u32(self):
|
||||
"""Regression test: V_XAD_U32 (xor-add) used by random number generation."""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 3),
|
||||
v_mov_b32_e32(v[2], 4),
|
||||
v_mov_b32_e32(v[4], 5),
|
||||
v_xad_u32(v[1], v[1], v[2], v[4]), # (3^4)+5 = 7+5 = 12
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out[0], 12)
|
||||
|
||||
def test_v_lshl_or_b32(self):
|
||||
"""Regression test: V_LSHL_OR_B32 operand order is (s0 << s1) | s2, not (s0 << s2) | s1."""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 5), # s0 = value to shift
|
||||
v_mov_b32_e32(v[2], 2), # s1 = shift amount
|
||||
v_mov_b32_e32(v[4], 3), # s2 = value to OR
|
||||
v_lshl_or_b32(v[1], v[1], v[2], v[4]), # (5 << 2) | 3 = 20 | 3 = 23
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out[0], 23)
|
||||
|
||||
def test_v_sqrt_f32_negative(self):
|
||||
"""Regression test: V_SQRT_F32 should return NaN for negative inputs, not 0."""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(-1.0)),
|
||||
v_sqrt_f32_e32(v[1], v[1]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertTrue(math.isnan(f32(out[0])))
|
||||
|
||||
def test_v_rsq_f32_negative(self):
|
||||
"""Regression test: V_RSQ_F32 should return NaN for negative inputs, not inf."""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i32(-1.0)),
|
||||
v_rsq_f32_e32(v[1], v[1]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertTrue(math.isnan(f32(out[0])))
|
||||
|
||||
class TestVOPD(unittest.TestCase):
|
||||
def test_vopd_add_nc_u32(self):
|
||||
"""Test VOPD V_DUAL_ADD_NC_U32."""
|
||||
state = WaveState()
|
||||
state.vgpr[0][1] = 100
|
||||
state.vgpr[0][2] = 50
|
||||
# vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
|
||||
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
|
||||
opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
|
||||
kernel += s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.vgpr[0][3], 100) # MOV result
|
||||
self.assertEqual(state.vgpr[0][4], 150) # 100 + 50
|
||||
|
||||
def test_vopd_lshlrev(self):
|
||||
"""Test VOPD V_DUAL_LSHLREV_B32."""
|
||||
state = WaveState()
|
||||
state.vgpr[0][1] = 0x10
|
||||
state.vgpr[0][2] = 0
|
||||
# vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
|
||||
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
|
||||
opy=VOPDOp.V_DUAL_LSHLREV_B32, srcy0=4, vsrcy1=VGPR(1), vdsty=VGPR(4)).to_bytes() # V4 = V1 << 4
|
||||
kernel += s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.vgpr[0][3], 0x10) # MOV result
|
||||
self.assertEqual(state.vgpr[0][4], 0x100) # 0x10 << 4 = 0x100
|
||||
|
||||
def test_vopd_and(self):
|
||||
"""Test VOPD V_DUAL_AND_B32."""
|
||||
state = WaveState()
|
||||
state.vgpr[0][1] = 0xff
|
||||
state.vgpr[0][2] = 0x0f
|
||||
# vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
|
||||
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
|
||||
opy=VOPDOp.V_DUAL_AND_B32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
|
||||
kernel += s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.vgpr[0][3], 0xff)
|
||||
self.assertEqual(state.vgpr[0][4], 0x0f) # 0xff & 0x0f = 0x0f
|
||||
|
||||
def test_vopd_parallel_read(self):
|
||||
"""Regression: VOPD must read all inputs before writing - Y op reads register that X op writes."""
|
||||
state = WaveState()
|
||||
state.vgpr[0][4] = 0
|
||||
state.vgpr[0][7] = 5 # Y op reads v7 as vsrcy1, X op writes to v7
|
||||
# X: MOV v7, v0 (v0=0, so v7 becomes 0)
|
||||
# Y: ADD v6, v4, v7 (should use original v7=5, not the overwritten 0)
|
||||
# vdsty_enc=3 with vdstx=7 (odd) -> vdsty = (3 << 1) | (7&1)^1 = 6 | 0 = 6
|
||||
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[0], vsrcx1=VGPR(0), vdstx=VGPR(7),
|
||||
opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[4], vsrcy1=VGPR(7), vdsty=VGPR(6)).to_bytes()
|
||||
kernel += s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.vgpr[0][7], 0) # X op: v7 = v0 = 0
|
||||
self.assertEqual(state.vgpr[0][6], 5) # Y op: v6 = v4 + v7 = 0 + 5 (original v7)
|
||||
|
||||
class TestDecoder(unittest.TestCase):
|
||||
def test_vopd_literal_handling(self):
|
||||
"""Regression test: VOPD srcx0/srcy0 with literal (255) wasn't consuming the literal dword."""
|
||||
state = WaveState()
|
||||
# Create VOPD with srcx0=255 (literal), followed by literal value 0x12345678
|
||||
vopd_bytes = VOPD(opx=8, srcx0=RawImm(255), vsrcx1=VGPR(0), vdstx=VGPR(1), # MOV: V1 = literal
|
||||
opy=8, srcy0=RawImm(128), vsrcy1=VGPR(0), vdsty=VGPR(2)).to_bytes() # MOV: V2 = 0
|
||||
literal_bytes = (0x12345678).to_bytes(4, 'little')
|
||||
kernel = vopd_bytes + literal_bytes + s_endpgm().to_bytes()
|
||||
prog = decode_program(kernel)
|
||||
# Should decode as 3 instructions: VOPD (with literal), then S_ENDPGM
|
||||
# The literal should NOT be decoded as a separate instruction
|
||||
self.assertEqual(len(prog), 2) # VOPD + S_ENDPGM
|
||||
exec_wave(prog, state, bytearray(65536), 1)
|
||||
self.assertEqual(state.vgpr[0][1], 0x12345678)
|
||||
|
||||
def test_s_endpgm_stops_decode(self):
|
||||
"""Regression test: decoder should stop at S_ENDPGM, not read past into metadata."""
|
||||
# Create a kernel followed by garbage that looks like an invalid instruction
|
||||
kernel = s_mov_b32(s[0], 42).to_bytes() + s_endpgm().to_bytes()
|
||||
garbage = bytes([0xff] * 16) # garbage after kernel
|
||||
prog = decode_program(kernel + garbage)
|
||||
# Should only have 2 instructions (s_mov_b32 and s_endpgm)
|
||||
self.assertEqual(len(prog), 2)
|
||||
|
||||
class TestFloatConversion(unittest.TestCase):
|
||||
"""Unit tests for i32/i16/f32/f16 float conversion functions."""
|
||||
|
||||
def test_i32_preserves_nan_sign(self):
|
||||
"""NaN sign bit should be preserved when converting float to int bits."""
|
||||
from extra.assembly.rdna3.emu import i32, f32
|
||||
# 0 * -inf produces a negative NaN
|
||||
neg_nan = 0.0 * float('-inf')
|
||||
bits = i32(neg_nan)
|
||||
# Should have sign bit set (0xffc00000), not canonical positive NaN (0x7fc00000)
|
||||
self.assertEqual(bits & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{bits:08x}")
|
||||
self.assertTrue(math.isnan(f32(bits)))
|
||||
|
||||
def test_i32_preserves_positive_nan(self):
|
||||
"""Positive NaN should remain positive."""
|
||||
from extra.assembly.rdna3.emu import i32, f32
|
||||
pos_nan = float('nan')
|
||||
bits = i32(pos_nan)
|
||||
# Standard Python NaN is positive (0x7fc00000)
|
||||
self.assertEqual(bits & 0x80000000, 0, f"Expected positive NaN, got 0x{bits:08x}")
|
||||
self.assertTrue(math.isnan(f32(bits)))
|
||||
|
||||
def test_i32_overflow_to_inf(self):
|
||||
"""Values too large for f32 should become inf."""
|
||||
from extra.assembly.rdna3.emu import i32, f32
|
||||
big = 2.0 ** 200
|
||||
self.assertEqual(i32(big), 0x7f800000) # +inf
|
||||
self.assertEqual(i32(-big), 0xff800000) # -inf
|
||||
|
||||
def test_i32_inf(self):
|
||||
"""Infinity should be preserved."""
|
||||
from extra.assembly.rdna3.emu import i32
|
||||
self.assertEqual(i32(float('inf')), 0x7f800000)
|
||||
self.assertEqual(i32(float('-inf')), 0xff800000)
|
||||
|
||||
def test_i32_normal_values(self):
|
||||
"""Normal float values should round-trip correctly (within f32 precision)."""
|
||||
from extra.assembly.rdna3.emu import i32, f32
|
||||
# Use values exactly representable in float32
|
||||
for val in [0.0, 1.0, -1.0, 0.5, -0.5, 100.0, -100.0, 1e10]:
|
||||
bits = i32(val)
|
||||
self.assertAlmostEqual(f32(bits), val, places=5)
|
||||
|
||||
def test_i16_overflow_to_inf(self):
|
||||
"""Values too large for f16 should become inf."""
|
||||
from extra.assembly.rdna3.emu import i16
|
||||
big = 100000.0 # way larger than f16 max (65504)
|
||||
self.assertEqual(i16(big), 0x7c00) # +inf
|
||||
self.assertEqual(i16(-big), 0xfc00) # -inf
|
||||
|
||||
def test_i16_inf(self):
|
||||
"""Infinity should be preserved."""
|
||||
from extra.assembly.rdna3.emu import i16
|
||||
self.assertEqual(i16(float('inf')), 0x7c00)
|
||||
self.assertEqual(i16(float('-inf')), 0xfc00)
|
||||
|
||||
def test_fma_nan_sign_preserved(self):
|
||||
"""FMA producing NaN should preserve the correct sign bit."""
|
||||
from extra.assembly.rdna3.emu import i32, f32
|
||||
# 0 * (-inf) + 1.0 = NaN (from 0 * -inf)
|
||||
a, b, c = 0.0, float('-inf'), 1.0
|
||||
result = i32(a * b + c)
|
||||
# The NaN should be negative since 0 * -inf produces negative NaN
|
||||
self.assertEqual(result & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{result:08x}")
|
||||
|
||||
class TestMultiWave(unittest.TestCase):
|
||||
def test_all_waves_execute(self):
|
||||
"""Regression test: all waves in a workgroup must execute, not just the first."""
|
||||
n_threads = 64 # 2 waves of 32 threads each
|
||||
output = (ctypes.c_uint32 * n_threads)(*[0xdead] * n_threads)
|
||||
output_ptr = ctypes.addressof(output)
|
||||
args = (ctypes.c_uint64 * 1)(output_ptr)
|
||||
args_ptr = ctypes.addressof(args)
|
||||
|
||||
# Simple kernel: store tid to output[tid]
|
||||
kernel = b''
|
||||
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
|
||||
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
|
||||
kernel += v_lshlrev_b32_e32(v[1], 2, v[0]).to_bytes() # offset = tid * 4
|
||||
kernel += global_store_b32(addr=v[1], data=v[0], saddr=s[2]).to_bytes()
|
||||
kernel += s_endpgm().to_bytes()
|
||||
|
||||
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
|
||||
kernel_ptr = ctypes.addressof(kernel_buf)
|
||||
set_valid_mem_ranges({
|
||||
(output_ptr, ctypes.sizeof(output)),
|
||||
(args_ptr, ctypes.sizeof(args)),
|
||||
(kernel_ptr, len(kernel)),
|
||||
})
|
||||
result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
|
||||
self.assertEqual(result, 0)
|
||||
# All threads should have written their tid
|
||||
for i in range(n_threads):
|
||||
self.assertEqual(output[i], i, f"Thread {i} didn't execute")
|
||||
|
||||
class TestRegressions(unittest.TestCase):
|
||||
"""Regression tests for bugs fixed in the emulator."""
|
||||
|
||||
def test_v_fmac_f16(self):
|
||||
"""V_FMAC_F16: fused multiply-add for FP16. Regression for VOP2 op 54."""
|
||||
from extra.assembly.rdna3.emu import i16, f16
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], i16(2.0)), # v1.lo = 2.0 (fp16)
|
||||
v_mov_b32_e32(v[2], i16(3.0)), # v2.lo = 3.0 (fp16)
|
||||
# v1 = v1 * v2 + v1 = 2.0 * 3.0 + 2.0 = 8.0
|
||||
VOP2(VOP2Op.V_FMAC_F16, v[1], v[1], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertAlmostEqual(f16(out[0] & 0xffff), 8.0, places=2)
|
||||
|
||||
def test_v_cvt_f64_f32(self):
|
||||
"""V_CVT_F64_F32: convert float32 to float64. Regression for VOP1 op 16."""
|
||||
kernel = b''
|
||||
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
|
||||
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
|
||||
kernel += v_mov_b32_e32(v[1], i32(3.14159)).to_bytes()
|
||||
kernel += VOP1(VOP1Op.V_CVT_F64_F32, v[4], v[1]).to_bytes() # v4:v5 = f64(v1)
|
||||
kernel += v_lshlrev_b32_e32(v[3], 3, v[0]).to_bytes() # offset = tid * 8
|
||||
kernel += global_store_b64(addr=v[3], data=v[4], saddr=s[2]).to_bytes()
|
||||
kernel += s_endpgm().to_bytes()
|
||||
output = (ctypes.c_double * 1)(0.0)
|
||||
output_ptr = ctypes.addressof(output)
|
||||
args = (ctypes.c_uint64 * 1)(output_ptr)
|
||||
args_ptr = ctypes.addressof(args)
|
||||
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
|
||||
kernel_ptr = ctypes.addressof(kernel_buf)
|
||||
set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
|
||||
run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
|
||||
self.assertAlmostEqual(output[0], 3.14159, places=4)
|
||||
|
||||
def test_v_add_f64(self):
|
||||
"""V_ADD_F64: add two float64 values. Regression for VOP3 op 807."""
|
||||
from extra.assembly.rdna3.emu import i64_parts
|
||||
kernel = b''
|
||||
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
|
||||
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
|
||||
# Load 1.5 into v1:v2
|
||||
lo, hi = i64_parts(1.5)
|
||||
kernel += v_mov_b32_e32(v[1], lo).to_bytes()
|
||||
kernel += v_mov_b32_e32(v[2], hi).to_bytes()
|
||||
# Load 2.5 into v3:v4
|
||||
lo, hi = i64_parts(2.5)
|
||||
kernel += v_mov_b32_e32(v[3], lo).to_bytes()
|
||||
kernel += v_mov_b32_e32(v[4], hi).to_bytes()
|
||||
# v5:v6 = v1:v2 + v3:v4 = 1.5 + 2.5 = 4.0
|
||||
kernel += VOP3(VOP3Op.V_ADD_F64, v[5], v[1], v[3]).to_bytes()
|
||||
kernel += v_lshlrev_b32_e32(v[7], 3, v[0]).to_bytes()
|
||||
kernel += global_store_b64(addr=v[7], data=v[5], saddr=s[2]).to_bytes()
|
||||
kernel += s_endpgm().to_bytes()
|
||||
output = (ctypes.c_double * 1)(0.0)
|
||||
output_ptr = ctypes.addressof(output)
|
||||
args = (ctypes.c_uint64 * 1)(output_ptr)
|
||||
args_ptr = ctypes.addressof(args)
|
||||
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
|
||||
kernel_ptr = ctypes.addressof(kernel_buf)
|
||||
set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
|
||||
run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
|
||||
self.assertAlmostEqual(output[0], 4.0, places=10)
|
||||
|
||||
def test_flat_load_d16_hi_b16(self):
|
||||
"""FLAT_LOAD_D16_HI_B16: load 16-bit to high half. Regression for FLAT op 35."""
|
||||
from extra.assembly.rdna3.emu import i16
|
||||
# Create a buffer with test data
|
||||
src_data = (ctypes.c_uint16 * 1)(0x1234)
|
||||
src_ptr = ctypes.addressof(src_data)
|
||||
output = (ctypes.c_uint32 * 1)(0xABCD0000) # preset low bits
|
||||
output_ptr = ctypes.addressof(output)
|
||||
args = (ctypes.c_uint64 * 2)(output_ptr, src_ptr)
|
||||
args_ptr = ctypes.addressof(args)
|
||||
|
||||
kernel = b''
|
||||
kernel += s_load_b128(s[0:3], s[0:1], soffset=NULL, offset=0).to_bytes()
|
||||
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
|
||||
kernel += v_mov_b32_e32(v[1], 0xDEAD).to_bytes() # initial value with low bits set
|
||||
kernel += v_mov_b32_e32(v[2], 0).to_bytes() # offset = 0
|
||||
kernel += FLAT(FLATOp.FLAT_LOAD_D16_HI_B16, v[1], v[2], saddr=s[2], offset=0).to_bytes()
|
||||
kernel += s_waitcnt(vmcnt=0).to_bytes()
|
||||
kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
|
||||
kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[0]).to_bytes()
|
||||
kernel += s_endpgm().to_bytes()
|
||||
|
||||
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
|
||||
kernel_ptr = ctypes.addressof(kernel_buf)
|
||||
set_valid_mem_ranges({(output_ptr, 4), (src_ptr, 2), (args_ptr, 16), (kernel_ptr, len(kernel))})
|
||||
run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
|
||||
# High 16 bits should be 0x1234, low 16 bits preserved as 0xDEAD
|
||||
self.assertEqual(output[0], 0x1234DEAD)
|
||||
|
||||
def test_v_mad_u16(self):
|
||||
"""V_MAD_U16: multiply-add unsigned 16-bit. Regression for VOP3 op 577."""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 10), # a = 10
|
||||
v_mov_b32_e32(v[2], 20), # b = 20
|
||||
v_mov_b32_e32(v[4], 5), # c = 5
|
||||
VOP3(VOP3Op.V_MAD_U16, v[1], v[1], v[2], v[4]), # v1 = 10*20+5 = 205
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out[0] & 0xffff, 205)
|
||||
|
||||
def test_v_lshrrev_b16(self):
|
||||
"""V_LSHRREV_B16: logical shift right 16-bit. Regression for VOP3 op 825."""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 0x8000), # value to shift
|
||||
v_mov_b32_e32(v[2], 4), # shift amount
|
||||
VOP3(VOP3Op.V_LSHRREV_B16, v[1], v[2], v[1]), # v1 = 0x8000 >> 4 = 0x0800
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out[0] & 0xffff, 0x0800)
|
||||
|
||||
def test_v_min_u16(self):
|
||||
"""V_MIN_U16: minimum of two unsigned 16-bit values. Regression for VOP3 op 779."""
|
||||
kernel = make_store_kernel([
|
||||
v_mov_b32_e32(v[1], 100),
|
||||
v_mov_b32_e32(v[2], 50),
|
||||
VOP3(VOP3Op.V_MIN_U16, v[1], v[1], v[2]),
|
||||
])
|
||||
out = run_kernel(kernel, n_threads=1)
|
||||
self.assertEqual(out[0] & 0xffff, 50)
|
||||
|
||||
class TestWMMA(unittest.TestCase):
|
||||
"""Tests for WMMA (Wave Matrix Multiply Accumulate) instructions."""
|
||||
|
||||
def test_wmma_f32_16x16x16_f16_identity(self):
|
||||
"""V_WMMA_F32_16X16X16_F16 with identity matrix. Regression for VOP3P op 64."""
|
||||
from extra.assembly.rdna3.emu import i16, f16, exec_wmma_f32_16x16x16_f16, WaveState
|
||||
# Test using direct emulator call rather than full kernel to simplify
|
||||
st = WaveState()
|
||||
st.exec_mask = 0xffffffff # all 32 lanes active
|
||||
|
||||
# Set up A as identity matrix: A[i][i] = 1.0, rest = 0.0
|
||||
# Lane i holds row i of A in 8 regs (2 fp16 per reg)
|
||||
for lane in range(16):
|
||||
for reg in range(8):
|
||||
col0, col1 = reg * 2, reg * 2 + 1
|
||||
val0 = i16(1.0) if col0 == lane else 0
|
||||
val1 = i16(1.0) if col1 == lane else 0
|
||||
st.vgpr[lane][0 + reg] = val0 | (val1 << 16) # src0 = v0:v7
|
||||
|
||||
# Set up B as identity matrix: lane i holds column i of B
|
||||
for lane in range(16):
|
||||
for reg in range(8):
|
||||
row0, row1 = reg * 2, reg * 2 + 1
|
||||
val0 = i16(1.0) if row0 == lane else 0
|
||||
val1 = i16(1.0) if row1 == lane else 0
|
||||
st.vgpr[lane][8 + reg] = val0 | (val1 << 16) # src1 = v8:v15
|
||||
|
||||
# Set up C as zeros
|
||||
for lane in range(32):
|
||||
for reg in range(8):
|
||||
st.vgpr[lane][16 + reg] = 0 # src2 = v16:v23
|
||||
|
||||
# Create a fake VOP3P instruction
|
||||
inst = VOP3P(VOP3POp.V_WMMA_F32_16X16X16_F16, v[24], src0=VGPR(0), src1=VGPR(8), src2=VGPR(16))
|
||||
|
||||
# Execute WMMA
|
||||
exec_wmma_f32_16x16x16_f16(st, inst, 32)
|
||||
|
||||
# Check result: C should be identity (since A @ B where both are identity)
|
||||
# Output i = row*16+col goes to lane (i%32), reg (i//32)
|
||||
for row in range(16):
|
||||
for col in range(16):
|
||||
idx = row * 16 + col
|
||||
lane, reg = idx % 32, idx // 32
|
||||
result = st.vgpr[lane][24 + reg]
|
||||
expected = 1.0 if row == col else 0.0
|
||||
self.assertAlmostEqual(f32(result), expected, places=3,
|
||||
msg=f"C[{row},{col}] = {f32(result)}, expected {expected}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,173 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test RDNA3 assembler/disassembler against LLVM test vectors."""
|
||||
import unittest, re
|
||||
from tinygrad.helpers import fetch
|
||||
from extra.assembly.rdna3.autogen import *
|
||||
from extra.assembly.rdna3.asm import asm
|
||||
from extra.assembly.rdna3.test.test_roundtrip import compile_asm, disassemble_lib
|
||||
|
||||
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/main/llvm/test/MC/AMDGPU"
|
||||
|
||||
# Format info: (filename, format_class, op_enum)
|
||||
LLVM_TEST_FILES = {
|
||||
# Scalar ALU
|
||||
'sop1': ('gfx11_asm_sop1.s', SOP1, SOP1Op),
|
||||
'sop2': ('gfx11_asm_sop2.s', SOP2, SOP2Op),
|
||||
'sopp': ('gfx11_asm_sopp.s', SOPP, SOPPOp),
|
||||
'sopk': ('gfx11_asm_sopk.s', SOPK, SOPKOp),
|
||||
'sopc': ('gfx11_asm_sopc.s', SOPC, SOPCOp),
|
||||
# Vector ALU
|
||||
'vop1': ('gfx11_asm_vop1.s', VOP1, VOP1Op),
|
||||
'vop2': ('gfx11_asm_vop2.s', VOP2, VOP2Op),
|
||||
'vopc': ('gfx11_asm_vopc.s', VOPC, VOPCOp),
|
||||
'vop3': ('gfx11_asm_vop3.s', VOP3, VOP3Op),
|
||||
'vop3p': ('gfx11_asm_vop3p.s', VOP3P, VOP3POp),
|
||||
'vop3sd': ('gfx11_asm_vop3.s', VOP3SD, VOP3SDOp), # VOP3SD shares file with VOP3
|
||||
'vinterp': ('gfx11_asm_vinterp.s', VINTERP, VINTERPOp),
|
||||
'vopd': ('gfx11_asm_vopd.s', VOPD, VOPDOp),
|
||||
'vopcx': ('gfx11_asm_vopcx.s', VOPC, VOPCOp), # VOPCX uses VOPC format
|
||||
# VOP3 promotions (VOP1/VOP2/VOPC promoted to VOP3 encoding)
|
||||
'vop3_from_vop1': ('gfx11_asm_vop3_from_vop1.s', VOP3, VOP3Op),
|
||||
'vop3_from_vop2': ('gfx11_asm_vop3_from_vop2.s', VOP3, VOP3Op),
|
||||
'vop3_from_vopc': ('gfx11_asm_vop3_from_vopc.s', VOP3, VOP3Op),
|
||||
'vop3_from_vopcx': ('gfx11_asm_vop3_from_vopcx.s', VOP3, VOP3Op),
|
||||
# Memory
|
||||
'ds': ('gfx11_asm_ds.s', DS, DSOp),
|
||||
'smem': ('gfx11_asm_smem.s', SMEM, SMEMOp),
|
||||
'flat': ('gfx11_asm_flat.s', FLAT, FLATOp),
|
||||
'mubuf': ('gfx11_asm_mubuf.s', MUBUF, MUBUFOp),
|
||||
'mtbuf': ('gfx11_asm_mtbuf.s', MTBUF, MTBUFOp),
|
||||
'mimg': ('gfx11_asm_mimg.s', MIMG, MIMGOp),
|
||||
# WMMA (matrix multiply)
|
||||
'wmma': ('gfx11_asm_wmma.s', VOP3P, VOP3POp),
|
||||
# Additional features
|
||||
'vop3_features': ('gfx11_asm_vop3_features.s', VOP3, VOP3Op),
|
||||
'vop3p_features': ('gfx11_asm_vop3p_features.s', VOP3P, VOP3POp),
|
||||
'vopd_features': ('gfx11_asm_vopd_features.s', VOPD, VOPDOp),
|
||||
# Alias files (alternative mnemonics)
|
||||
'vop3_alias': ('gfx11_asm_vop3_alias.s', VOP3, VOP3Op),
|
||||
'vop3p_alias': ('gfx11_asm_vop3p_alias.s', VOP3P, VOP3POp),
|
||||
'vopc_alias': ('gfx11_asm_vopc_alias.s', VOPC, VOPCOp),
|
||||
'vopcx_alias': ('gfx11_asm_vopcx_alias.s', VOPC, VOPCOp),
|
||||
'vinterp_alias': ('gfx11_asm_vinterp_alias.s', VINTERP, VINTERPOp),
|
||||
'smem_alias': ('gfx11_asm_smem_alias.s', SMEM, SMEMOp),
|
||||
'mubuf_alias': ('gfx11_asm_mubuf_alias.s', MUBUF, MUBUFOp),
|
||||
'mtbuf_alias': ('gfx11_asm_mtbuf_alias.s', MTBUF, MTBUFOp),
|
||||
}
|
||||
|
||||
def parse_llvm_tests(text: str) -> list[tuple[str, bytes]]:
|
||||
"""Parse LLVM test format into (asm, expected_bytes) pairs."""
|
||||
tests, lines = [], text.split('\n')
|
||||
for i, line in enumerate(lines):
|
||||
line = line.strip()
|
||||
if not line or line.startswith(('//', '.', ';')): continue
|
||||
asm_text = line.split('//')[0].strip()
|
||||
if not asm_text: continue
|
||||
for j in range(i, min(i + 3, len(lines))):
|
||||
# Match GFX11, W32, or W64 encodings (all valid for gfx11)
|
||||
if m := re.search(r'(?:GFX11|W32|W64)[^:]*:.*?encoding:\s*\[(.*?)\]', lines[j]):
|
||||
hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
|
||||
if hex_bytes:
|
||||
try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
|
||||
except ValueError: pass
|
||||
break
|
||||
return tests
|
||||
|
||||
def try_assemble(text: str):
|
||||
"""Try to assemble instruction text, return bytes or None on failure."""
|
||||
try: return asm(text).to_bytes()
|
||||
except: return None
|
||||
|
||||
class TestLLVM(unittest.TestCase):
|
||||
"""Test assembler and disassembler against all LLVM test vectors."""
|
||||
tests: dict[str, list[tuple[str, bytes]]] = {}
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
for name, (filename, _, _) in LLVM_TEST_FILES.items():
|
||||
try:
|
||||
data = fetch(f"{LLVM_BASE}/{filename}").read_bytes()
|
||||
cls.tests[name] = parse_llvm_tests(data.decode('utf-8', errors='ignore'))
|
||||
except Exception as e:
|
||||
print(f"Warning: couldn't fetch {filename}: {e}")
|
||||
cls.tests[name] = []
|
||||
|
||||
# Generate test methods dynamically for each format
|
||||
def _make_asm_test(name):
|
||||
def test(self):
|
||||
passed, failed, skipped = 0, 0, 0
|
||||
for asm_text, expected in self.tests.get(name, []):
|
||||
result = try_assemble(asm_text)
|
||||
if result is None: skipped += 1
|
||||
elif result == expected: passed += 1
|
||||
else: failed += 1
|
||||
print(f"{name.upper()} asm: {passed} passed, {failed} failed, {skipped} skipped")
|
||||
self.assertEqual(failed, 0)
|
||||
return test
|
||||
|
||||
def _make_disasm_test(name):
|
||||
def test(self):
|
||||
from tinygrad.runtime.support.compiler_amd import HIPCompiler
|
||||
compiler = HIPCompiler('gfx1100')
|
||||
_, fmt_cls, op_enum = LLVM_TEST_FILES[name]
|
||||
passed, failed, skipped, failures = 0, 0, 0, []
|
||||
# VOP3SD opcodes that share encoding with VOP3 (only for vop3sd test, not vopc promotions)
|
||||
# Note: opcodes 0-255 are VOPC promoted to VOP3, never VOP3SD
|
||||
vop3sd_opcodes = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
|
||||
# vop3_from_vopc/vopcx tests have VOPC opcodes 0-255, not VOP3SD - don't detect as VOP3SD
|
||||
is_vopc_promotion = name in ('vop3_from_vopc', 'vop3_from_vopcx')
|
||||
# Undocumented opcodes not in AMD ISA PDF - skip these
|
||||
undocumented = {'smem': {34, 35}, 'sopk': {22, 23}, 'sopp': {8, 58, 59}} # s_atc_probe*, s_subvector_loop*, s_waitcnt_depctr, unknown
|
||||
for asm_text, data in self.tests.get(name, []):
|
||||
if len(data) > fmt_cls._size(): continue # skip literals (need different handling)
|
||||
# Skip undocumented opcodes
|
||||
temp_inst = fmt_cls.from_bytes(data)
|
||||
temp_op = temp_inst._values.get('op', 0)
|
||||
temp_op = temp_op.val if hasattr(temp_op, 'val') else temp_op
|
||||
if temp_op in undocumented.get(name, set()): skipped += 1; continue
|
||||
# Skip SOPP no-imm instructions with non-zero simm16 (can't roundtrip through LLVM)
|
||||
if name == 'sopp':
|
||||
simm16 = temp_inst._values.get('simm16', 0)
|
||||
simm16 = simm16.val if hasattr(simm16, 'val') else simm16
|
||||
sopp_no_imm = {48, 54, 53, 55, 60, 61, 62} # s_endpgm, s_barrier, s_wakeup, s_icache_inv, s_wait_idle, s_endpgm_saved, s_code_end
|
||||
if temp_op in sopp_no_imm and simm16 != 0: skipped += 1; continue
|
||||
try:
|
||||
# VOP3 and VOP3SD share encoding - peek at opcode to determine which class to use
|
||||
if fmt_cls.__name__ in ('VOP3', 'VOP3SD'):
|
||||
temp = VOP3.from_bytes(data)
|
||||
op_val = temp._values.get('op', 0)
|
||||
op_val = op_val.val if hasattr(op_val, 'val') else op_val
|
||||
is_vop3sd = (op_val in vop3sd_opcodes) and not is_vopc_promotion
|
||||
decoded = VOP3SD.from_bytes(data) if is_vop3sd else VOP3.from_bytes(data)
|
||||
# Validate opcode with appropriate enum
|
||||
if is_vop3sd:
|
||||
VOP3SDOp(op_val)
|
||||
else:
|
||||
VOP3Op(op_val)
|
||||
else:
|
||||
decoded = fmt_cls.from_bytes(data)
|
||||
op_val = decoded._values.get('op', 0)
|
||||
op_val = op_val.val if hasattr(op_val, 'val') else op_val
|
||||
op_enum(op_val) # validate opcode
|
||||
if decoded.to_bytes()[:len(data)] != data:
|
||||
failed += 1; failures.append(f"decode roundtrip failed for {data.hex()}"); continue
|
||||
disasm_str = decoded.disasm()
|
||||
# Test: LLVM should assemble our disasm output to the same bytes
|
||||
llvm_bytes = compile_asm(disasm_str, compiler)
|
||||
if llvm_bytes is None:
|
||||
failed += 1; failures.append(f"LLVM failed to assemble: '{disasm_str}' (from '{asm_text}')")
|
||||
elif llvm_bytes == data: passed += 1
|
||||
else: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
|
||||
except Exception as e:
|
||||
failed += 1; failures.append(f"exception for {data.hex()}: {e}")
|
||||
print(f"{name.upper()} disasm: {passed} passed, {failed} failed" + (f", {skipped} skipped" if skipped else ""))
|
||||
if failures[:10]: print(" " + "\n ".join(failures[:10]))
|
||||
self.assertEqual(failed, 0)
|
||||
return test
|
||||
|
||||
for name in LLVM_TEST_FILES:
|
||||
setattr(TestLLVM, f'test_{name}_asm', _make_asm_test(name))
|
||||
setattr(TestLLVM, f'test_{name}_disasm', _make_disasm_test(name))
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,158 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test that PDF parser correctly extracts format fields."""
|
||||
import unittest
|
||||
from extra.assembly.rdna3.autogen import (
|
||||
SOP1, SOP2, SOPK, SOPP, VOP1, VOP2, VOP3SD, VOPC, FLAT, VOPD,
|
||||
SOP1Op, SOP2Op, VOP1Op, VOP3Op
|
||||
)
|
||||
|
||||
# expected formats with key fields and whether they have ENCODING
|
||||
EXPECTED_FORMATS = {
|
||||
'DPP16': (['SRC0', 'DPP_CTRL', 'BANK_MASK', 'ROW_MASK'], False),
|
||||
'DPP8': (['SRC0', 'LANE_SEL0', 'LANE_SEL7'], False),
|
||||
'DS': (['OP', 'ADDR', 'DATA0', 'DATA1', 'VDST'], True),
|
||||
'EXP': (['EN', 'TARGET', 'VSRC0', 'VSRC1', 'VSRC2', 'VSRC3'], True),
|
||||
'FLAT': (['OP', 'ADDR', 'DATA', 'SADDR', 'VDST', 'OFFSET'], True),
|
||||
'LDSDIR': (['VDST', 'OP'], True),
|
||||
'MIMG': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'DMASK'], True),
|
||||
'MTBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'FORMAT', 'SOFFSET'], True),
|
||||
'MUBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'SOFFSET'], True),
|
||||
'SMEM': (['OP', 'SBASE', 'SDATA', 'OFFSET', 'SOFFSET'], True),
|
||||
'SOP1': (['OP', 'SDST', 'SSRC0'], True),
|
||||
'SOP2': (['OP', 'SDST', 'SSRC0', 'SSRC1'], True),
|
||||
'SOPC': (['OP', 'SSRC0', 'SSRC1'], True),
|
||||
'SOPK': (['OP', 'SDST', 'SIMM16'], True),
|
||||
'SOPP': (['OP', 'SIMM16'], True),
|
||||
'VINTERP': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
|
||||
'VOP1': (['OP', 'VDST', 'SRC0'], True),
|
||||
'VOP2': (['OP', 'VDST', 'SRC0', 'VSRC1'], True),
|
||||
'VOP3': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
|
||||
'VOP3P': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
|
||||
'VOP3SD': (['OP', 'VDST', 'SDST', 'SRC0', 'SRC1', 'SRC2'], True),
|
||||
'VOPC': (['OP', 'SRC0', 'VSRC1'], True),
|
||||
'VOPD': (['OPX', 'OPY', 'SRCX0', 'SRCY0', 'VDSTX', 'VDSTY'], True),
|
||||
}
|
||||
|
||||
class TestPDFParserGenerate(unittest.TestCase):
|
||||
"""Test the PDF parser by running generate() and checking results."""
|
||||
result: dict
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
from extra.assembly.rdna3.gen import generate
|
||||
cls.result = generate()
|
||||
|
||||
def test_all_formats_present(self):
|
||||
"""All expected formats should be parsed."""
|
||||
for fmt_name in EXPECTED_FORMATS:
|
||||
self.assertIn(fmt_name, self.result["formats"], f"missing format {fmt_name}")
|
||||
|
||||
def test_format_count(self):
|
||||
"""Should have exactly 23 formats."""
|
||||
self.assertEqual(len(self.result["formats"]), 23)
|
||||
|
||||
def test_no_duplicate_fields(self):
|
||||
"""No format should have duplicate field names."""
|
||||
for fmt_name, fields in self.result["formats"].items():
|
||||
field_names = [f[0] for f in fields]
|
||||
self.assertEqual(len(field_names), len(set(field_names)), f"{fmt_name} has duplicate fields: {field_names}")
|
||||
|
||||
def test_expected_fields(self):
|
||||
"""Each format should have its expected key fields."""
|
||||
for fmt_name, (expected_fields, has_encoding) in EXPECTED_FORMATS.items():
|
||||
fields = {f[0] for f in self.result["formats"].get(fmt_name, [])}
|
||||
for field in expected_fields:
|
||||
self.assertIn(field, fields, f"{fmt_name} missing {field}")
|
||||
if has_encoding:
|
||||
self.assertIn("ENCODING", fields, f"{fmt_name} should have ENCODING")
|
||||
else:
|
||||
self.assertNotIn("ENCODING", fields, f"{fmt_name} should not have ENCODING")
|
||||
|
||||
def test_vopd_no_dpp16_fields(self):
|
||||
"""VOPD should not have DPP16-specific fields (parser boundary bug)."""
|
||||
vopd_fields = {f[0] for f in self.result["formats"].get("VOPD", [])}
|
||||
for field in ['DPP_CTRL', 'BANK_MASK', 'ROW_MASK']:
|
||||
self.assertNotIn(field, vopd_fields, f"VOPD should not have {field}")
|
||||
|
||||
def test_dpp16_no_vinterp_fields(self):
|
||||
"""DPP16 should not have VINTERP-specific fields."""
|
||||
dpp16_fields = {f[0] for f in self.result["formats"].get("DPP16", [])}
|
||||
for field in ['VDST', 'WAITEXP']:
|
||||
self.assertNotIn(field, dpp16_fields, f"DPP16 should not have {field}")
|
||||
|
||||
def test_sopp_no_smem_fields(self):
|
||||
"""SOPP should not have SMEM fields (page break bug)."""
|
||||
sopp_fields = {f[0] for f in self.result["formats"].get("SOPP", [])}
|
||||
for field in ['SBASE', 'SDATA']:
|
||||
self.assertNotIn(field, sopp_fields, f"SOPP should not have {field}")
|
||||
|
||||
class TestPDFParser(unittest.TestCase):
|
||||
"""Verify format classes have correct fields from PDF parsing."""
|
||||
|
||||
def test_sop2_fields(self):
|
||||
"""SOP2 should have op, sdst, ssrc0, ssrc1."""
|
||||
for field in ['op', 'sdst', 'ssrc0', 'ssrc1']:
|
||||
self.assertIn(field, SOP2._fields)
|
||||
self.assertEqual(SOP2._fields['op'].hi, 29)
|
||||
self.assertEqual(SOP2._fields['op'].lo, 23)
|
||||
|
||||
def test_sop1_fields(self):
|
||||
"""SOP1 should have op, sdst, ssrc0 with correct bit positions."""
|
||||
for field in ['op', 'sdst', 'ssrc0']:
|
||||
self.assertIn(field, SOP1._fields)
|
||||
self.assertNotIn('simm16', SOP1._fields)
|
||||
self.assertEqual(SOP1._fields['ssrc0'].hi, 7)
|
||||
self.assertEqual(SOP1._fields['ssrc0'].lo, 0)
|
||||
assert SOP1._encoding is not None
|
||||
self.assertEqual(SOP1._encoding[0].hi, 31)
|
||||
self.assertEqual(SOP1._encoding[1], 0b101111101)
|
||||
|
||||
def test_vop3sd_fields(self):
|
||||
"""VOP3SD should have all fields including src0/src1/src2 from page continuation."""
|
||||
for field in ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2']:
|
||||
self.assertIn(field, VOP3SD._fields)
|
||||
self.assertEqual(VOP3SD._fields['src0'].hi, 40)
|
||||
self.assertEqual(VOP3SD._fields['src0'].lo, 32)
|
||||
self.assertEqual(VOP3SD._size(), 8)
|
||||
|
||||
def test_flat_has_vdst(self):
|
||||
"""FLAT should have vdst field."""
|
||||
self.assertIn('vdst', FLAT._fields)
|
||||
self.assertEqual(FLAT._fields['vdst'].hi, 63)
|
||||
self.assertEqual(FLAT._fields['vdst'].lo, 56)
|
||||
|
||||
def test_encoding_bits(self):
|
||||
"""Verify encoding bits are correct for major formats."""
|
||||
tests = [
|
||||
(SOP2, 31, 30, 0b10),
|
||||
(SOPK, 31, 28, 0b1011),
|
||||
(SOPP, 31, 23, 0b101111111),
|
||||
(VOP1, 31, 25, 0b0111111),
|
||||
(VOP2, 31, 31, 0b0),
|
||||
(VOPC, 31, 25, 0b0111110),
|
||||
(FLAT, 31, 26, 0b110111),
|
||||
]
|
||||
for cls, hi, lo, val in tests:
|
||||
assert cls._encoding is not None
|
||||
self.assertEqual(cls._encoding[0].hi, hi, f"{cls.__name__} encoding hi")
|
||||
self.assertEqual(cls._encoding[0].lo, lo, f"{cls.__name__} encoding lo")
|
||||
self.assertEqual(cls._encoding[1], val, f"{cls.__name__} encoding val")
|
||||
|
||||
def test_opcode_enums_exist(self):
|
||||
"""Verify opcode enums are generated with expected counts."""
|
||||
self.assertGreater(len(SOP1Op), 50)
|
||||
self.assertGreater(len(SOP2Op), 50)
|
||||
self.assertGreater(len(VOP1Op), 50)
|
||||
self.assertGreater(len(VOP3Op), 200)
|
||||
|
||||
def test_vopd_no_duplicate_fields(self):
|
||||
"""VOPD should not have duplicate fields and should not include DPP16 fields."""
|
||||
field_names = list(VOPD._fields.keys())
|
||||
self.assertEqual(len(field_names), len(set(field_names)))
|
||||
for field in ['srcx0', 'srcy0', 'opx', 'opy']:
|
||||
self.assertIn(field, VOPD._fields)
|
||||
for field in ['dpp_ctrl', 'bank_mask', 'row_mask']:
|
||||
self.assertNotIn(field, VOPD._fields)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,234 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
|
||||
import unittest, io, sys, re
|
||||
from extra.assembly.rdna3.autogen import *
|
||||
from extra.assembly.rdna3.lib import Inst
|
||||
from extra.assembly.rdna3.asm import asm
|
||||
|
||||
# Instruction format detection based on encoding bits
|
||||
def detect_format(data: bytes) -> type[Inst] | None:
|
||||
"""Detect instruction format from machine code bytes."""
|
||||
if len(data) < 4: return None
|
||||
word = int.from_bytes(data[:4], 'little')
|
||||
enc_9bit = (word >> 23) & 0x1FF # 9-bit encoding for SOP1/SOPC/SOPP
|
||||
enc_8bit = (word >> 24) & 0xFF
|
||||
|
||||
# Check 9-bit encodings first (most specific)
|
||||
if enc_9bit == 0x17D: return SOP1 # bits 31:23 = 101111101
|
||||
if enc_9bit == 0x17E: return SOPC # bits 31:23 = 101111110
|
||||
if enc_9bit == 0x17F: return SOPP # bits 31:23 = 101111111
|
||||
# SOPK: bits 31:28 = 1011, bits 27:23 = opcode (check after SOP1/SOPC/SOPP)
|
||||
if enc_8bit in range(0xB0, 0xC0): return SOPK
|
||||
# SOP2: bits 31:23 in range 0x100-0x17C (0x80-0xBE in bits 31:24, but not SOPK)
|
||||
if 0x80 <= enc_8bit <= 0x9F: return SOP2
|
||||
# VOP1: bits 31:25 = 0111111 (0x3F)
|
||||
if (word >> 25) == 0x3F: return VOP1
|
||||
# VOPC: bits 31:25 = 0111110 (0x3E)
|
||||
if (word >> 25) == 0x3E: return VOPC
|
||||
# VOP2: bits 31:30 = 00
|
||||
if (word >> 30) == 0: return VOP2
|
||||
|
||||
# Check 64-bit formats
|
||||
if len(data) >= 8:
|
||||
if enc_8bit in (0xD4, 0xD5, 0xD7): return VOP3
|
||||
if enc_8bit == 0xD6: return VOP3SD
|
||||
if enc_8bit == 0xCC: return VOP3P
|
||||
if enc_8bit == 0xCD: return VINTERP
|
||||
if enc_8bit in (0xC8, 0xC9): return VOPD
|
||||
if enc_8bit == 0xF4: return SMEM
|
||||
if enc_8bit == 0xD8: return DS
|
||||
if enc_8bit in (0xDC, 0xDD, 0xDE, 0xDF): return FLAT
|
||||
if enc_8bit in (0xE0, 0xE1, 0xE2, 0xE3): return MUBUF
|
||||
if enc_8bit in (0xE8, 0xE9, 0xEA, 0xEB): return MTBUF
|
||||
|
||||
return None
|
||||
|
||||
def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
|
||||
"""Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
|
||||
old_stdout = sys.stdout
|
||||
sys.stdout = io.StringIO()
|
||||
compiler.disassemble(lib)
|
||||
output = sys.stdout.getvalue()
|
||||
sys.stdout = old_stdout
|
||||
|
||||
results = []
|
||||
for line in output.splitlines():
|
||||
if '//' not in line: continue
|
||||
instr = line.split('//')[0].strip()
|
||||
if not instr: continue
|
||||
comment = line.split('//')[1].strip()
|
||||
if ':' not in comment: continue
|
||||
hex_str = comment.split(':')[1].strip().split()[0]
|
||||
try:
|
||||
machine_bytes = bytes.fromhex(hex_str)[::-1] # big-endian to little-endian
|
||||
results.append((instr, machine_bytes))
|
||||
except ValueError:
|
||||
continue
|
||||
return results
|
||||
|
||||
def compile_asm(instr: str, compiler=None) -> bytes | None:
|
||||
"""Compile a single instruction with llvm-mc and return the machine code bytes."""
|
||||
import subprocess
|
||||
try:
|
||||
result = subprocess.run(
|
||||
['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
|
||||
input=f".text\n{instr}\n", capture_output=True, text=True)
|
||||
if result.returncode != 0: return None
|
||||
# Parse encoding: [0x01,0x39,0x0a,0x7e]
|
||||
for line in result.stdout.split('\n'):
|
||||
if 'encoding:' in line:
|
||||
enc = line.split('encoding:')[1].strip()
|
||||
if enc.startswith('[') and enc.endswith(']'):
|
||||
hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
|
||||
return bytes.fromhex(hex_vals)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
class TestTinygradKernelRoundtrip(unittest.TestCase):
|
||||
"""Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
|
||||
|
||||
def _test_kernel_roundtrip(self, op_fn):
|
||||
"""Generate kernel from op_fn, test:
|
||||
1. decode -> reencode matches original bytes
|
||||
2. asm(disasm()) matches LLVM output
|
||||
3. our disasm() matches LLVM's disassembly string exactly
|
||||
"""
|
||||
from extra.assembly.rdna3.test.test_compare_emulators import get_kernels_from_tinygrad
|
||||
from tinygrad.runtime.support.compiler_amd import HIPCompiler
|
||||
|
||||
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
|
||||
compiler = HIPCompiler('gfx1100')
|
||||
|
||||
decode_passed, decode_failed, decode_skipped = 0, 0, 0
|
||||
asm_passed, asm_failed, asm_skipped = 0, 0, 0
|
||||
disasm_passed, disasm_failed, disasm_skipped = 0, 0, 0
|
||||
decode_failures, asm_failures, disasm_failures = [], [], []
|
||||
|
||||
for ki, kernel in enumerate(kernels):
|
||||
offset = 0
|
||||
while offset < len(kernel.code):
|
||||
remaining = kernel.code[offset:]
|
||||
fmt = detect_format(remaining)
|
||||
if fmt is None:
|
||||
decode_skipped += 1
|
||||
asm_skipped += 1
|
||||
disasm_skipped += 1
|
||||
offset += 4
|
||||
continue
|
||||
|
||||
size = fmt._size()
|
||||
if len(remaining) < size:
|
||||
break
|
||||
|
||||
orig_bytes = remaining[:size]
|
||||
|
||||
# Test 1: decode -> reencode roundtrip
|
||||
try:
|
||||
decoded = fmt.from_bytes(orig_bytes)
|
||||
reencoded = decoded.to_bytes()
|
||||
if reencoded[:size] == orig_bytes:
|
||||
decode_passed += 1
|
||||
else:
|
||||
decode_failed += 1
|
||||
decode_failures.append(f"K{ki}@{offset}: {decoded.disasm()}: orig={orig_bytes.hex()} reenc={reencoded[:size].hex()}")
|
||||
|
||||
our_disasm = decoded.disasm()
|
||||
|
||||
# Test 2: asm(disasm()) matches LLVM output
|
||||
try:
|
||||
our_bytes = asm(our_disasm).to_bytes()
|
||||
llvm_bytes = compile_asm(our_disasm, compiler)
|
||||
if llvm_bytes is None:
|
||||
asm_skipped += 1
|
||||
elif our_bytes[:len(llvm_bytes)] == llvm_bytes:
|
||||
asm_passed += 1
|
||||
else:
|
||||
asm_failed += 1
|
||||
asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': ours={our_bytes[:len(llvm_bytes)].hex()} llvm={llvm_bytes.hex()}")
|
||||
except Exception:
|
||||
asm_skipped += 1
|
||||
|
||||
# Test 3: our disasm() matches LLVM's disassembly string exactly
|
||||
# Skip if instruction uses op_XX (unknown opcode) or looks malformed (many raw field values)
|
||||
if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm):
|
||||
disasm_skipped += 1
|
||||
else:
|
||||
try:
|
||||
# Get LLVM's disassembly of our instruction
|
||||
src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n {our_disasm}\n"
|
||||
lib = compiler.compile(src)
|
||||
llvm_instrs = disassemble_lib(lib, compiler)
|
||||
if llvm_instrs:
|
||||
llvm_disasm = llvm_instrs[0][0]
|
||||
if our_disasm == llvm_disasm:
|
||||
disasm_passed += 1
|
||||
else:
|
||||
disasm_failed += 1
|
||||
disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
|
||||
else:
|
||||
disasm_skipped += 1
|
||||
except Exception:
|
||||
disasm_skipped += 1
|
||||
|
||||
except Exception:
|
||||
decode_skipped += 1
|
||||
asm_skipped += 1
|
||||
disasm_skipped += 1
|
||||
|
||||
offset += size
|
||||
|
||||
print(f"decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
|
||||
print(f"asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
|
||||
print(f"disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
|
||||
self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
|
||||
self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
|
||||
self.assertEqual(disasm_failed, 0, f"Disasm failures:\n" + "\n".join(disasm_failures[:20]))
|
||||
|
||||
# Basic unary ops
|
||||
def test_neg(self): self._test_kernel_roundtrip(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
|
||||
def test_relu(self): self._test_kernel_roundtrip(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
|
||||
def test_exp(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).exp())
|
||||
def test_log(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 3.0]).log())
|
||||
def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
|
||||
def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
|
||||
def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
|
||||
|
||||
# Binary ops
|
||||
def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
|
||||
def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
|
||||
def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
|
||||
def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
|
||||
def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
|
||||
|
||||
# Reductions
|
||||
def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
|
||||
def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
|
||||
def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
|
||||
|
||||
# Matmul
|
||||
def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
|
||||
def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
|
||||
|
||||
# Complex ops
|
||||
def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
|
||||
def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
|
||||
|
||||
# Memory patterns
|
||||
def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
|
||||
def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
|
||||
def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
|
||||
|
||||
# Cast ops
|
||||
def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
|
||||
def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
|
||||
|
||||
# Comparison ops
|
||||
def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
|
||||
def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
|
||||
|
||||
# Fused ops
|
||||
def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,4 +0,0 @@
|
||||
*.deb
|
||||
build
|
||||
src
|
||||
sniffer/sniff.so
|
||||
@@ -1,20 +0,0 @@
|
||||
Built ROCT-Thunk-Interface (hsakmt)
|
||||
hsakmt-roct-dev_5.4.4.99999-local_amd64.deb
|
||||
note: installs to /opt/rocm
|
||||
Built ROCm-Device-Libs
|
||||
Works with ROCM_PATH=/home/tiny/build/ROCm-Device-Libs/build/dist
|
||||
rocm-device-libs_1.0.0.99999-local_amd64.deb
|
||||
Built ROCm-CompilerSupport (amd_comgr)
|
||||
no deb, sudo make install to /usr/local
|
||||
Built ROCR-Runtime
|
||||
hsa-rocr_1.8.0-local_amd64.deb
|
||||
hsa-rocr-dev_1.8.0-local_amd64.deb
|
||||
Built ROCm-OpenCL-Runtime
|
||||
rocm-ocl-icd_2.0.0-local_amd64.deb
|
||||
ISSUE: these depend on "comgr"
|
||||
rocm-opencl_2.0.0-local_amd64.deb
|
||||
rocm-opencl-dev_2.0.0-local_amd64.deb
|
||||
Did sudo make install
|
||||
|
||||
|
||||
|
||||
@@ -1,41 +0,0 @@
|
||||
# run two "rocm-bandwidth-test" in a loop
|
||||
# amdgpu-6.0.5-1581431.20.04
|
||||
# fixed in kernel 6.2.14
|
||||
|
||||
[ 72.153646] RIP: 0010:pm_send_runlist+0x4a/0x630 [amdgpu]
|
||||
[ 72.153815] Code: 30 65 48 8b 04 25 28 00 00 00 48 89 45 d0 31 c0 80 fb 01 0f 87 aa 9d 49 00 83 e3 01 0f 85 1c 05 00 00 49 8b 3f b8 01 00 00 00 <48> 8b 97 30 01 00 00 44 8b b7 6c 01 00 00 8b 9f 70 01 00 00 8b 8a
|
||||
[ 72.153900] RSP: 0018:ffffb48445c03c30 EFLAGS: 00010246
|
||||
[ 72.153928] RAX: 0000000000000001 RBX: 0000000000000000 RCX: 0000000000000000
|
||||
[ 72.153962] RDX: 000000000000007b RSI: ffff9395e1562558 RDI: 0000000000000000
|
||||
[ 72.153996] RBP: ffffb48445c03cb8 R08: 0000000000000000 R09: 0000000000000001
|
||||
[ 72.154030] R10: ffff9395c900d840 R11: 0000000000000000 R12: 0000000000000000
|
||||
[ 72.154065] R13: ffff9395c9e00400 R14: 0000000000000001 R15: ffff9395e15624e0
|
||||
[ 72.154099] FS: 00007f345c6463c0(0000) GS:ffff93a4aee80000(0000) knlGS:0000000000000000
|
||||
[ 72.154137] CS: 0010 DS: 0000 ES: 0000 CR0: 0000000080050033
|
||||
[ 72.154165] CR2: 0000000000000130 CR3: 0000000112840000 CR4: 0000000000750ee0
|
||||
[ 72.154201] PKRU: 55555554
|
||||
[ 72.154215] Call Trace:
|
||||
[ 72.154230] <TASK>
|
||||
[ 72.154244] map_queues_cpsch+0x75/0xc0 [amdgpu]
|
||||
[ 72.154365] debug_map_and_unlock+0x51/0x90 [amdgpu]
|
||||
[ 72.154480] debug_refresh_runlist+0x1f/0x30 [amdgpu]
|
||||
[ 72.154591] kfd_dbg_runtime_disable+0x13c/0x240 [amdgpu]
|
||||
[ 72.154705] kfd_ioctl_dbg_set_debug_trap+0x69d/0x8b0 [amdgpu]
|
||||
[ 72.154820] kfd_ioctl+0x24a/0x5b0 [amdgpu]
|
||||
[ 72.154925] ? kfd_ioctl_create_queue+0x770/0x770 [amdgpu]
|
||||
[ 72.155035] ? syscall_exit_to_user_mode+0x27/0x50
|
||||
[ 72.155061] ? exit_to_user_mode_prepare+0x3d/0x1c0
|
||||
[ 72.155088] __x64_sys_ioctl+0x95/0xd0
|
||||
[ 72.155109] do_syscall_64+0x5c/0xc0
|
||||
[ 72.155128] ? syscall_exit_to_user_mode+0x27/0x50
|
||||
[ 72.155151] ? do_syscall_64+0x69/0xc0
|
||||
[ 72.155172] entry_SYSCALL_64_after_hwframe+0x61/0xcb
|
||||
[ 72.155198] RIP: 0033:0x7f345c7f63ab
|
||||
[ 72.155218] Code: 0f 1e fa 48 8b 05 e5 7a 0d 00 64 c7 00 26 00 00 00 48 c7 c0 ff ff ff ff c3 66 0f 1f 44 00 00 f3 0f 1e fa b8 10 00 00 00 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 8b 0d b5 7a 0d 00 f7 d8 64 89 01 48
|
||||
[ 72.155301] RSP: 002b:00007ffc97cc89f8 EFLAGS: 00000246 ORIG_RAX: 0000000000000010
|
||||
[ 72.155339] RAX: ffffffffffffffda RBX: 00007ffc97cc8a30 RCX: 00007f345c7f63ab
|
||||
[ 72.155375] RDX: 00007ffc97cc8a30 RSI: 00000000c0284b82 RDI: 0000000000000003
|
||||
[ 72.155411] RBP: 00000000c0284b82 R08: 0000000000000000 R09: 0000000000000000
|
||||
[ 72.155447] R10: 00007f345cd4ddb0 R11: 0000000000000246 R12: 00007ffc97cc8a30
|
||||
[ 72.155481] R13: 0000000000000003 R14: 00007ffc97cc8d20 R15: 0000000000000000
|
||||
[ 72.155517] </TASK>
|
||||
@@ -1,41 +0,0 @@
|
||||
# run two tinygrad matrix example in a loop
|
||||
# amdgpu-6.0.5-1581431.20.04
|
||||
# NOT fixed in kernel 6.2.14
|
||||
|
||||
[ 553.016624] gmc_v11_0_process_interrupt: 30 callbacks suppressed
|
||||
[ 553.016631] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:9 pasid:32770, for process python3 pid 10001 thread python3 pid 10001)
|
||||
[ 553.016790] amdgpu 0000:0b:00.0: amdgpu: in page starting at address 0x00007f0000000000 from client 10
|
||||
[ 553.016892] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00901A30
|
||||
[ 553.016974] amdgpu 0000:0b:00.0: amdgpu: Faulty UTCL2 client ID: SDMA0 (0xd)
|
||||
[ 553.017051] amdgpu 0000:0b:00.0: amdgpu: MORE_FAULTS: 0x0
|
||||
[ 553.017111] amdgpu 0000:0b:00.0: amdgpu: WALKER_ERROR: 0x0
|
||||
[ 553.017173] amdgpu 0000:0b:00.0: amdgpu: PERMISSION_FAULTS: 0x3
|
||||
[ 553.017238] amdgpu 0000:0b:00.0: amdgpu: MAPPING_ERROR: 0x0
|
||||
[ 553.017300] amdgpu 0000:0b:00.0: amdgpu: RW: 0x0
|
||||
[ 553.123921] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
|
||||
[ 553.124153] amdgpu: failed to add hardware queue to MES, doorbell=0x1a16
|
||||
[ 553.124195] amdgpu: MES might be in unrecoverable state, issue a GPU reset
|
||||
[ 553.124237] amdgpu: Failed to restore queue 2
|
||||
[ 553.124266] amdgpu: Failed to restore process queues
|
||||
[ 553.124270] amdgpu: Failed to evict queue 3
|
||||
[ 553.124297] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
|
||||
|
||||
# alternative crash in kernel 6.2.14
|
||||
|
||||
[ 151.097948] gmc_v11_0_process_interrupt: 30 callbacks suppressed
|
||||
[ 151.097953] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:8 pasid:32771, for process python3 pid 7525 thread python3 pid 7525)
|
||||
[ 151.097993] amdgpu 0000:0b:00.0: amdgpu: in page starting at address 0x00007f0000000000 from client 10
|
||||
[ 151.098008] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00801A30
|
||||
[ 151.098020] amdgpu 0000:0b:00.0: amdgpu: Faulty UTCL2 client ID: SDMA0 (0xd)
|
||||
[ 151.098032] amdgpu 0000:0b:00.0: amdgpu: MORE_FAULTS: 0x0
|
||||
[ 151.098042] amdgpu 0000:0b:00.0: amdgpu: WALKER_ERROR: 0x0
|
||||
[ 151.098052] amdgpu 0000:0b:00.0: amdgpu: PERMISSION_FAULTS: 0x3
|
||||
[ 151.098062] amdgpu 0000:0b:00.0: amdgpu: MAPPING_ERROR: 0x0
|
||||
[ 151.098071] amdgpu 0000:0b:00.0: amdgpu: RW: 0x0
|
||||
[ 151.209517] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
|
||||
[ 151.209724] amdgpu: failed to add hardware queue to MES, doorbell=0x1002
|
||||
[ 151.209734] amdgpu: MES might be in unrecoverable state, issue a GPU reset
|
||||
[ 151.209743] amdgpu: Failed to restore queue 1
|
||||
[ 151.209751] amdgpu: Failed to restore process queues
|
||||
[ 151.209759] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
|
||||
[ 151.209858] amdgpu 0000:0b:00.0: amdgpu: GPU reset begin!
|
||||
@@ -1,20 +0,0 @@
|
||||
# two tinygrad + two bandwidth test
|
||||
# RDNA2, driver 6.0.5
|
||||
# recovered from this!
|
||||
|
||||
[ 136.971209] gmc_v10_0_process_interrupt: 39 callbacks suppressed
|
||||
[ 136.971218] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:11 pasid:32773, for process rocm-bandwidth- pid 20281 thread rocm-bandwidth- pid 20281)
|
||||
[ 136.971228] amdgpu 0000:0b:00.0: amdgpu: in page starting at address 0x00007f5c2b800000 from client 0x1b (UTCL2)
|
||||
[ 136.971232] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00B01A31
|
||||
[ 136.971233] amdgpu 0000:0b:00.0: amdgpu: Faulty UTCL2 client ID: SDMA0 (0xd)
|
||||
[ 136.971235] amdgpu 0000:0b:00.0: amdgpu: MORE_FAULTS: 0x1
|
||||
[ 136.971236] amdgpu 0000:0b:00.0: amdgpu: WALKER_ERROR: 0x0
|
||||
[ 136.971236] amdgpu 0000:0b:00.0: amdgpu: PERMISSION_FAULTS: 0x3
|
||||
[ 136.971237] amdgpu 0000:0b:00.0: amdgpu: MAPPING_ERROR: 0x0
|
||||
[ 136.971238] amdgpu 0000:0b:00.0: amdgpu: RW: 0x0
|
||||
...
|
||||
[ 136.993979] amdgpu 0000:0b:00.0: amdgpu: IH ring buffer overflow (0x000BE5A0, 0x0003C480, 0x0003E5C0)
|
||||
[ 138.209072] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004000 flags=0x0000]
|
||||
[ 138.209078] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004d80 flags=0x0000]
|
||||
[ 138.209081] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005000 flags=0x0000]
|
||||
[ 138.209084] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005d80 flags=0x0000]
|
||||
@@ -1,33 +0,0 @@
|
||||
# ROCK-Kernel-Driver 0b579de9622f5c93021dcb7927d13926313740a2
|
||||
# non fatal "crash"
|
||||
|
||||
[ 127.418045] ------------[ cut here ]------------
|
||||
[ 127.418046] User pages unexpectedly invalid
|
||||
[ 127.418056] WARNING: CPU: 16 PID: 260 at drivers/gpu/drm/amd/amdgpu/amdgpu_amdkfd_gpuvm.c:3000 amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
|
||||
[ 127.418235] Modules linked in: rfcomm cmac algif_hash algif_skcipher af_alg bnep nls_iso8859_1 iwlmvm mac80211 intel_rapl_msr intel_rapl_common edac_mce_amd snd_hda_codec_realtek snd_hda_codec_generic snd_hda_codec_hdmi kvm_amd binfmt_misc snd_hda_intel snd_intel_dspcfg kvm libarc4 snd_intel_sdw_acpi snd_hda_codec btusb iwlwifi btrtl snd_hda_core btbcm btintel irqbypass btmtk snd_hwdep crct10dif_pclmul snd_pcm polyval_clmulni bluetooth snd_seq_midi snd_seq_midi_event snd_rawmidi snd_seq polyval_generic cfg80211 ghash_clmulni_intel eeepc_wmi snd_seq_device snd_timer aesni_intel asus_wmi ecdh_generic snd platform_profile crypto_simd ledtrig_audio cryptd ecc ccp soundcore sparse_keymap rapl k10temp wmi_bmof mac_hid sch_fq_codel msr parport_pc ppdev lp parport ramoops pstore_blk efi_pstore reed_solomon pstore_zone ip_tables x_tables autofs4 amdgpu hid_generic usbhid hid i2c_algo_bit drm_ttm_helper ttm video iommu_v2 drm_buddy gpu_sched drm_display_helper drm_kms_helper syscopyarea
|
||||
[ 127.418276] sysfillrect sysimgblt fb_sys_fops drm nvme nvme_core cec r8169 ahci crc32_pclmul rc_core i2c_piix4 xhci_pci libahci nvme_common xhci_pci_renesas realtek wmi
|
||||
[ 127.418284] CPU: 16 PID: 260 Comm: kworker/16:1 Tainted: G W 6.0.0 #4
|
||||
[ 127.418286] Hardware name: System manufacturer System Product Name/TUF GAMING X570-PLUS (WI-FI), BIOS 3603 03/20/2021
|
||||
[ 127.418287] Workqueue: events amdgpu_amdkfd_restore_userptr_worker [amdgpu]
|
||||
[ 127.418455] RIP: 0010:amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
|
||||
[ 127.418601] Code: ff e8 2b 8a 96 d1 e9 66 fe ff ff 48 c7 c7 40 4f f5 c0 e8 56 7b 8a d1 0f 0b e9 2e ff ff ff 48 c7 c7 d8 d0 ed c0 e8 43 7b 8a d1 <0f> 0b e9 0a fe ff ff 4c 89 ef e8 f8 89 96 d1 e9 cb fd ff ff e8 ce
|
||||
[ 127.418603] RSP: 0018:ffffb36740a83dc8 EFLAGS: 00010282
|
||||
[ 127.418604] RAX: 0000000000000000 RBX: ffff9d159ee9df30 RCX: 0000000000000027
|
||||
[ 127.418605] RDX: 0000000000000027 RSI: ffffb36740a83c88 RDI: ffff9d242a220568
|
||||
[ 127.418606] RBP: ffffb36740a83e58 R08: ffff9d242a220560 R09: 0000000000000001
|
||||
[ 127.418607] R10: 0000000000000001 R11: 0000000000000020 R12: ffff9d159ee9df98
|
||||
[ 127.418607] R13: ffff9d159ee9df70 R14: ffff9d159ee9dee0 R15: ffff9d159ee9dee0
|
||||
[ 127.418608] FS: 0000000000000000(0000) GS:ffff9d242a200000(0000) knlGS:0000000000000000
|
||||
[ 127.418609] CS: 0010 DS: 0000 ES: 0000 CR0: 0000000080050033
|
||||
[ 127.418610] CR2: 00007fd5d4715000 CR3: 0000000120ffe000 CR4: 0000000000750ee0
|
||||
[ 127.418611] PKRU: 55555554
|
||||
[ 127.418611] Call Trace:
|
||||
[ 127.418612] <TASK>
|
||||
[ 127.418613] process_one_work+0x21f/0x3f0
|
||||
[ 127.418615] worker_thread+0x4a/0x3c0
|
||||
[ 127.418617] ? process_one_work+0x3f0/0x3f0
|
||||
[ 127.418618] kthread+0xf0/0x120
|
||||
[ 127.418619] ? kthread_complete_and_exit+0x20/0x20
|
||||
[ 127.418620] ret_from_fork+0x22/0x30
|
||||
[ 127.418622] </TASK>
|
||||
[ 127.418623] ---[ end trace 0000000000000000 ]---
|
||||
@@ -1,80 +0,0 @@
|
||||
import numpy as np
|
||||
import pathlib
|
||||
from hexdump import hexdump
|
||||
from tinygrad.helpers import colored
|
||||
from extra.helpers import enable_early_exec
|
||||
early_exec = enable_early_exec()
|
||||
|
||||
from tinygrad.runtime.ops_cl import CLProgram, CLBuffer, ROCM_LLVM_PATH
|
||||
|
||||
ENABLE_NON_ASM = False
|
||||
|
||||
WMMA = True
|
||||
DUAL_ALU = True
|
||||
F32 = True
|
||||
|
||||
if ENABLE_NON_ASM:
|
||||
buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
|
||||
prg_empty = CLProgram("code", "__kernel void code(__global float *a) { a[0] = 1; }")
|
||||
asm_real = prg_empty.binary()
|
||||
with open("/tmp/cc.elf", "wb") as f:
|
||||
f.write(asm_real)
|
||||
prg_empty([1], [1], buf, wait=True)
|
||||
print(buf.toCPU())
|
||||
|
||||
print(colored("creating CLBuffer", "green"))
|
||||
buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
|
||||
code = open(pathlib.Path(__file__).parent / "prog.s", "r").read()
|
||||
|
||||
gen = []
|
||||
FLOPS = 0
|
||||
MAX_REG = 251
|
||||
for j in range(1):
|
||||
if WMMA:
|
||||
KY, KX = 4, 4
|
||||
for y in range(KY):
|
||||
for x in range(KX):
|
||||
c = (y*KX+x)*8
|
||||
a = (KY*KX*8) + y*8
|
||||
b = (KY*KX*8) + (KY*8) + x*8
|
||||
gen.append(f"v_wmma_f32_16x16x16_f16 v[{c}:{c+7}], v[{a}:{a+7}], v[{b}:{b+7}], v[{c}:{c+7}]")
|
||||
FLOPS += 16*8*2
|
||||
else:
|
||||
for i in range(0, MAX_REG, 6):
|
||||
if DUAL_ALU:
|
||||
if F32:
|
||||
gen.append(f"v_dual_fmac_f32 v{i+0}, v{i+1}, v{i+2} :: v_dual_fmac_f32 v{i+3}, v{i+4}, v{i+5}")
|
||||
FLOPS += 4
|
||||
else:
|
||||
gen.append(f"v_dual_dot2acc_f32_f16 v{i+0}, v{i+1}, v{i+2} :: v_dual_dot2acc_f32_f16 v{i+3}, v{i+4}, v{i+5}")
|
||||
FLOPS += 8
|
||||
else:
|
||||
assert F32
|
||||
gen.append(f"v_fmac_f32 v{i+0}, v{i+1}, v{i+2}")
|
||||
gen.append(f"v_fmac_f32 v{i+3}, v{i+4}, v{i+5}")
|
||||
code = code.replace("// FLOPS", '\n'.join(gen))
|
||||
print(code)
|
||||
|
||||
|
||||
# fix: COMGR failed to get code object ISA name. set triple to 'amdgcn-amd-amdhsa'
|
||||
|
||||
object = early_exec(([ROCM_LLVM_PATH / "llvm-mc", '--arch=amdgcn', '--mcpu=gfx1100', '--triple=amdgcn-amd-amdhsa', '--filetype=obj', '-'], code.encode("utf-8")))
|
||||
asm = early_exec(([ROCM_LLVM_PATH / "ld.lld", "/dev/stdin", "-o", "/dev/stdout", "--pie"], object))
|
||||
|
||||
with open("/tmp/cc2.o", "wb") as f:
|
||||
f.write(object)
|
||||
with open("/tmp/cc2.elf", "wb") as f:
|
||||
f.write(asm)
|
||||
|
||||
print(colored("creating CLProgram", "green"))
|
||||
prg = CLProgram("code", asm)
|
||||
|
||||
print(colored("running program", "green"))
|
||||
G = 512
|
||||
FLOPS *= 100000*G*G # loop * global_size
|
||||
for i in range(3):
|
||||
tm = prg(buf, global_size=[G//256, G, 1], local_size=[256, 1, 1], wait=True)
|
||||
print(f"ran in {tm*1e3:.2f} ms, {FLOPS/(tm*1e9):.2f} GFLOPS")
|
||||
|
||||
print(colored("transferring buffer", "green"))
|
||||
print(buf.toCPU())
|
||||
@@ -1,80 +0,0 @@
|
||||
.global _start
|
||||
_start:
|
||||
.rodata
|
||||
.align 0x10
|
||||
.global code.kd
|
||||
.type code.kd,STT_OBJECT
|
||||
# amd_kernel_code_t (must be at 0x440 for kernel_code_entry_byte_offset to be right)
|
||||
code.kd:
|
||||
# amd_kernel_..., amd_machine_...
|
||||
.long 0,0,0,0
|
||||
# kernel_code_entry_byte_offset, kernel_code_prefetch_byte_offset
|
||||
.long 0x00000bc0,0x00000000,0x00000000,0x00000000
|
||||
# kernel_code_prefetch_byte_size, max_scratch_backing_memory_byte_size
|
||||
.long 0,0,0,0
|
||||
# compute_pgm_rsrc1, compute_pgm_rsrc2, kernel_code_properties, workitem_private_segment_byte_size
|
||||
.long 0x60af0000,0x0000009e,0x00000408,0x00000000
|
||||
# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_32 | AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_16_64
|
||||
# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_DX10_CLAMP | AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_IEEE_MODE
|
||||
# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT = 0xF
|
||||
# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X
|
||||
# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_KERNARG_SEGMENT_PTR = 1
|
||||
# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_RESERVED1 = 1
|
||||
.text
|
||||
.global code
|
||||
.type code,STT_FUNC
|
||||
code:
|
||||
# https://llvm.org/docs/AMDGPUUsage.html#initial-kernel-execution-state
|
||||
# s[0:1] contains the kernarg_address
|
||||
# TODO: can we use s[2:3] if this was really a wave since we only alloced 2 SGPRs?
|
||||
s_load_b64 s[2:3], s[0:1], null
|
||||
|
||||
s_mov_b32 s8, 0
|
||||
loop:
|
||||
s_addk_i32 s8, 1
|
||||
s_cmp_eq_u32 s8, 100000
|
||||
// FLOPS
|
||||
s_cbranch_scc0 loop
|
||||
|
||||
# wait for the s_load_b64
|
||||
s_waitcnt lgkmcnt(0)
|
||||
|
||||
v_dual_mov_b32 v0, 4 :: v_dual_mov_b32 v1, 2.0
|
||||
global_store_b32 v0, v1, s[2:3]
|
||||
|
||||
# Deallocate all VGPRs for this wave. Use only when next instruction is S_ENDPGM.
|
||||
s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
|
||||
s_endpgm
|
||||
s_code_end
|
||||
|
||||
.amdgpu_metadata
|
||||
amdhsa.kernels:
|
||||
- .args:
|
||||
- .address_space: global
|
||||
.name: a
|
||||
.offset: 0
|
||||
.size: 8
|
||||
.type_name: 'float*'
|
||||
.value_kind: global_buffer
|
||||
.group_segment_fixed_size: 0
|
||||
.kernarg_segment_align: 8
|
||||
.kernarg_segment_size: 8
|
||||
.language: OpenCL C
|
||||
.language_version:
|
||||
- 1
|
||||
- 2
|
||||
.max_flat_workgroup_size: 256
|
||||
.name: code
|
||||
.private_segment_fixed_size: 0
|
||||
.sgpr_count: 2
|
||||
.sgpr_spill_count: 0
|
||||
.symbol: code.kd
|
||||
.uses_dynamic_stack: false
|
||||
.vgpr_count: 256
|
||||
.vgpr_spill_count: 0
|
||||
.wavefront_size: 32
|
||||
amdhsa.target: amdgcn-amd-amdhsa--gfx1100
|
||||
amdhsa.version:
|
||||
- 1
|
||||
- 2
|
||||
.end_amdgpu_metadata
|
||||
@@ -1,11 +0,0 @@
|
||||
#!/bin/bash
|
||||
mkdir -p src
|
||||
cd src
|
||||
git clone https://github.com/RadeonOpenCompute/ROCT-Thunk-Interface.git -b rocm-5.5.0
|
||||
git clone https://github.com/RadeonOpenCompute/ROCm-Device-Libs.git -b rocm-5.5.0
|
||||
git clone https://github.com/RadeonOpenCompute/llvm-project.git -b rocm-5.5.0 --depth 1
|
||||
git clone https://github.com/RadeonOpenCompute/ROCR-Runtime.git -b rocm-5.5.0
|
||||
git clone https://github.com/ROCm-Developer-Tools/ROCclr.git -b rocm-5.5.0
|
||||
git clone https://github.com/RadeonOpenCompute/ROCm-CompilerSupport.git -b rocm-5.5.0
|
||||
git clone https://github.com/RadeonOpenCompute/ROCm-OpenCL-Runtime.git -b rocm-5.5.0
|
||||
cd ../
|
||||
@@ -1,69 +0,0 @@
|
||||
#!/bin/bash
|
||||
mkdir -p build/debs
|
||||
cd build
|
||||
|
||||
# ROCT-Thunk-Interface (hsakmt)
|
||||
if [ ! -f debs/hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ]
|
||||
then
|
||||
mkdir -p ROCT-Thunk-Interface
|
||||
cd ROCT-Thunk-Interface
|
||||
cmake ../../src/ROCT-Thunk-Interface
|
||||
make -j32 package
|
||||
cp hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ../debs
|
||||
cd ../
|
||||
fi
|
||||
|
||||
|
||||
# build custom LLVM
|
||||
if [ ! -f llvm-project/bin/clang ]
|
||||
then
|
||||
mkdir -p llvm-project
|
||||
cd llvm-project
|
||||
cmake -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_PROJECTS="llvm;clang;lld" -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" ../../src/llvm-project/llvm
|
||||
make -j32
|
||||
cd ..
|
||||
fi
|
||||
|
||||
# use custom LLVM
|
||||
export PATH="$PWD/llvm-project/bin:$PATH"
|
||||
|
||||
# ROCm-Device-Libs
|
||||
if [ ! -f debs/rocm-device-libs_1.0.0.99999-local_amd64.deb ]
|
||||
then
|
||||
mkdir -p ROCm-Device-Libs
|
||||
cd ROCm-Device-Libs
|
||||
cmake ../../src/ROCm-Device-Libs
|
||||
make -j32 package
|
||||
cp rocm-device-libs_1.0.0.99999-local_amd64.deb ../debs
|
||||
cd ../
|
||||
fi
|
||||
|
||||
# ROCR-Runtime
|
||||
if [ ! -f debs/hsa-rocr_1.8.0-local_amd64.deb ]
|
||||
then
|
||||
mkdir -p ROCR-Runtime
|
||||
cd ROCR-Runtime
|
||||
cmake ../../src/ROCR-Runtime/src
|
||||
make -j32 package
|
||||
cp hsa-rocr_1.8.0-local_amd64.deb ../debs
|
||||
cp hsa-rocr-dev_1.8.0-local_amd64.deb ../debs
|
||||
cd ../
|
||||
fi
|
||||
|
||||
# ROCm-OpenCL-Runtime (needs ROCclr)
|
||||
if [ ! -f debs/rocm-opencl_2.0.0-local_amd64.deb ]
|
||||
then
|
||||
mkdir -p ROCm-OpenCL-Runtime
|
||||
cd ROCm-OpenCL-Runtime
|
||||
cmake ../../src/ROCm-OpenCL-Runtime
|
||||
make -j32 package
|
||||
cp rocm-opencl_2.0.0-local_amd64.deb ../debs
|
||||
cp rocm-opencl-dev_2.0.0-local_amd64.deb ../debs
|
||||
cp rocm-ocl-icd_2.0.0-local_amd64.deb ../debs
|
||||
fi
|
||||
|
||||
# ROCm-CompilerSupport (broken)
|
||||
#mkdir -p ROCm-CompilerSupport
|
||||
#cd ROCm-CompilerSupport
|
||||
#cmake ../../src/ROCm-CompilerSupport/lib/comgr
|
||||
#make -j32
|
||||
@@ -1,14 +0,0 @@
|
||||
#!/bin/bash
|
||||
rm amdgpu-install_5.5.50500-1_all.deb
|
||||
wget https://repo.radeon.com/amdgpu-install/5.5/ubuntu/$(lsb_release -cs)/amdgpu-install_5.5.50500-1_all.deb
|
||||
sudo dpkg -i amdgpu-install_5.5.50500-1_all.deb
|
||||
sudo apt-get update
|
||||
|
||||
# kernel driver
|
||||
sudo apt-get install amdgpu-dkms
|
||||
|
||||
# for opencl
|
||||
sudo apt-get install rocm-opencl-runtime
|
||||
|
||||
# for HIP
|
||||
sudo apt-get install hip-runtime-amd rocm-device-libs hip-dev
|
||||
@@ -1,11 +0,0 @@
|
||||
#!/bin/bash -e
|
||||
clang sniff.cc -Werror -shared -fPIC -I../src/ -I../src/ROCT-Thunk-Interface/include -I../src/ROCm-Device-Libs/ockl/inc -o sniff.so -lstdc++
|
||||
#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so /home/tiny/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
|
||||
#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so $HOME/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
|
||||
#AMD_LOG_LEVEL=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
|
||||
DEBUG=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
|
||||
#AMD_LOG_LEVEL=5 HSAKMT_DEBUG_LEVEL=7 DEBUG=5 LD_PRELOAD=$PWD/sniff.so strace -F python3 ../rdna3/asm.py
|
||||
#LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
|
||||
#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so FORWARD_ONLY=1 DEBUG=2 python3 ../../../test/test_ops.py TestOps.test_add
|
||||
#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 0 -d 1 -m 1
|
||||
#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 1 -d 2 -m 1
|
||||
@@ -1,282 +0,0 @@
|
||||
// template copied from https://github.com/geohot/cuda_ioctl_sniffer/blob/master/sniff.cc
|
||||
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <dlfcn.h>
|
||||
#include <signal.h>
|
||||
#include <ucontext.h>
|
||||
|
||||
#include <sys/mman.h>
|
||||
|
||||
// includes from the ROCm sources
|
||||
#include <linux/kfd_ioctl.h>
|
||||
#include <hsa.h>
|
||||
#include <amd_hsa_kernel_code.h>
|
||||
#include <ROCR-Runtime/src/core/inc/sdma_registers.h>
|
||||
using namespace rocr::AMD;
|
||||
|
||||
#include <string>
|
||||
#include <map>
|
||||
std::map<int, std::string> files;
|
||||
std::map<uint64_t, uint64_t> ring_base_addresses;
|
||||
|
||||
#define D(args...) fprintf(stderr, args)
|
||||
|
||||
uint64_t doorbell_offset = -1;
|
||||
std::map<uint64_t, int> queue_types;
|
||||
|
||||
void hexdump(void *d, int l) {
|
||||
for (int i = 0; i < l; i++) {
|
||||
if (i%0x10 == 0 && i != 0) printf("\n");
|
||||
if (i%0x10 == 8) printf(" ");
|
||||
if (i%0x10 == 0) printf("%8X: ", i);
|
||||
printf("%2.2X ", ((uint8_t*)d)[i]);
|
||||
}
|
||||
printf("\n");
|
||||
}
|
||||
|
||||
extern "C" {
|
||||
|
||||
// https://defuse.ca/online-x86-assembler.htm#disassembly2
|
||||
static void handler(int sig, siginfo_t *si, void *unused) {
|
||||
ucontext_t *u = (ucontext_t *)unused;
|
||||
uint8_t *rip = (uint8_t*)u->uc_mcontext.gregs[REG_RIP];
|
||||
|
||||
int store_size = 0;
|
||||
uint64_t value;
|
||||
if (rip[0] == 0x48 && rip[1] == 0x89 && rip[2] == 0x30) {
|
||||
// 0: 48 89 30 mov QWORD PTR [rax],rsi
|
||||
store_size = 8;
|
||||
value = u->uc_mcontext.gregs[REG_RSI];
|
||||
u->uc_mcontext.gregs[REG_RIP] += 3;
|
||||
} else if (rip[0] == 0x4c && rip[1] == 0x89 && rip[2] == 0x28) {
|
||||
// 0: 4c 89 28 mov QWORD PTR [rax],r13
|
||||
store_size = 8;
|
||||
value = u->uc_mcontext.gregs[REG_R13];
|
||||
u->uc_mcontext.gregs[REG_RIP] += 3;
|
||||
} else {
|
||||
D("segfault %02X %02X %02X %02X %02X %02X %02X %02X rip: %p addr: %p\n", rip[0], rip[1], rip[2], rip[3], rip[4], rip[5], rip[6], rip[7], rip, si->si_addr);
|
||||
D("rax: %llx rcx: %llx rdx: %llx rsi: %llx rbx: %llx\n", u->uc_mcontext.gregs[REG_RAX], u->uc_mcontext.gregs[REG_RCX], u->uc_mcontext.gregs[REG_RDX], u->uc_mcontext.gregs[REG_RSI], u->uc_mcontext.gregs[REG_RBX]);
|
||||
exit(-1);
|
||||
}
|
||||
|
||||
uint64_t ring_base_address = ring_base_addresses[((uint64_t)si->si_addr)&0xFFF];
|
||||
int queue_type = queue_types[((uint64_t)si->si_addr)&0xFFF];
|
||||
D("%16p: \u001b[31mDING DONG\u001b[0m (queue_type %d) store(%d): 0x%8lx -> %p ring_base_address:0x%lx\n", rip, queue_type, store_size, value, si->si_addr, ring_base_address);
|
||||
|
||||
if (queue_type == KFD_IOC_QUEUE_TYPE_SDMA) {
|
||||
uint8_t *sdma_ptr = (uint8_t*)(ring_base_address);
|
||||
while (sdma_ptr < ((uint8_t*)(ring_base_address)+value)) {
|
||||
D("0x%3lx: ", sdma_ptr-(uint8_t*)(ring_base_address));
|
||||
if (sdma_ptr[0] == SDMA_OP_TIMESTAMP) {
|
||||
D("SDMA_PKT_TIMESTAMP\n");
|
||||
sdma_ptr += sizeof(SDMA_PKT_TIMESTAMP);
|
||||
} else if (sdma_ptr[0] == SDMA_OP_GCR) {
|
||||
D("SDMA_PKT_GCR\n");
|
||||
sdma_ptr += sizeof(SDMA_PKT_GCR);
|
||||
} else if (sdma_ptr[0] == SDMA_OP_ATOMIC) {
|
||||
D("SDMA_PKT_ATOMIC\n");
|
||||
sdma_ptr += sizeof(SDMA_PKT_ATOMIC);
|
||||
} else if (sdma_ptr[0] == SDMA_OP_FENCE) {
|
||||
D("SDMA_PKT_FENCE\n");
|
||||
sdma_ptr += sizeof(SDMA_PKT_FENCE);
|
||||
} else if (sdma_ptr[0] == SDMA_OP_TRAP) {
|
||||
D("SDMA_PKT_TRAP\n");
|
||||
sdma_ptr += sizeof(SDMA_PKT_TRAP);
|
||||
} else if (sdma_ptr[0] == SDMA_OP_COPY && sdma_ptr[1] == SDMA_SUBOP_COPY_LINEAR) {
|
||||
SDMA_PKT_COPY_LINEAR *pkt = (SDMA_PKT_COPY_LINEAR *)sdma_ptr;
|
||||
D("SDMA_PKT_COPY_LINEAR: count:0x%x src:0x%lx dst:0x%lx\n", pkt->COUNT_UNION.count+1,
|
||||
(uint64_t)pkt->SRC_ADDR_LO_UNION.src_addr_31_0 | ((uint64_t)pkt->SRC_ADDR_HI_UNION.src_addr_63_32 << 32),
|
||||
(uint64_t)pkt->DST_ADDR_LO_UNION.dst_addr_31_0 | ((uint64_t)pkt->DST_ADDR_HI_UNION.dst_addr_63_32 << 32)
|
||||
);
|
||||
sdma_ptr += sizeof(SDMA_PKT_COPY_LINEAR);
|
||||
} else {
|
||||
D("unhandled packet type %d %d, exiting\n", sdma_ptr[0], sdma_ptr[1]);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
//hexdump((void*)(ring_base_address), 0x100);
|
||||
} else if (queue_type == KFD_IOC_QUEUE_TYPE_COMPUTE_AQL) {
|
||||
hsa_kernel_dispatch_packet_t *pkt = (hsa_kernel_dispatch_packet_t *)(ring_base_address+value*0x40);
|
||||
if ((pkt->header&0xFF) == HSA_PACKET_TYPE_KERNEL_DISPATCH) {
|
||||
D("HSA_PACKET_TYPE_KERNEL_DISPATCH -- setup:%d workgroup[%d, %d, %d] grid[%d, %d, %d] kernel_object:0x%lx kernarg_address:%p\n", pkt->setup, pkt->workgroup_size_x, pkt->workgroup_size_y, pkt->workgroup_size_z, pkt->grid_size_x, pkt->grid_size_y, pkt->grid_size_z, pkt->kernel_object, pkt->kernarg_address);
|
||||
amd_kernel_code_t *code = (amd_kernel_code_t *)pkt->kernel_object;
|
||||
D("kernel_code_entry_byte_offset:%lx\n", code->kernel_code_entry_byte_offset);
|
||||
uint32_t *kernel_code = (uint32_t*)(pkt->kernel_object + code->kernel_code_entry_byte_offset);
|
||||
int code_len = 0;
|
||||
while (kernel_code[code_len] != 0xbf9f0000 && kernel_code[code_len] != 0) code_len++;
|
||||
hexdump(kernel_code, code_len*4);
|
||||
/*FILE *f = fopen("/tmp/kernel_code", "wb");
|
||||
fwrite(kernel_code, 4, code_len, f);
|
||||
fclose(f);
|
||||
system("python -c 'print(\" \".join([(\"0x%02X\"%x) for x in open(\"/tmp/kernel_code\", \"rb\").read()]))' | ../build/llvm-project/bin/llvm-mc --disassemble --arch=amdgcn --mcpu=gfx1100 --show-encoding");*/
|
||||
D("kernargs (kernarg_segment_byte_size:0x%lx)\n", code->kernarg_segment_byte_size);
|
||||
// get length
|
||||
int i;
|
||||
for (i = 0; i < 0x400; i+=0x10) {
|
||||
if (memcmp((void*)((uint64_t)pkt->kernarg_address+i), "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", 0x10) == 0) break;
|
||||
}
|
||||
hexdump((void*)pkt->kernarg_address, i+0x10);
|
||||
} else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_BARRIER_AND) {
|
||||
hsa_barrier_and_packet_t *pkt_and = (hsa_barrier_and_packet_t *)(ring_base_address+value*0x40);
|
||||
D("HSA_PACKET_TYPE_BARRIER_AND completion_signal:0x%lx\n", pkt_and->completion_signal.handle);
|
||||
//hexdump((void*)(ring_base_address+value*0x40), 0x40);
|
||||
} else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_VENDOR_SPECIFIC) {
|
||||
D("HSA_PACKET_TYPE_VENDOR_SPECIFIC\n");
|
||||
hexdump((void*)(ring_base_address+value*0x40), 0x40);
|
||||
} else {
|
||||
hexdump((void*)(ring_base_address+value*0x40), 0x40);
|
||||
}
|
||||
}
|
||||
|
||||
mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_READ | PROT_WRITE);
|
||||
if (store_size == 8) {
|
||||
*(volatile uint64_t*)(si->si_addr) = value;
|
||||
} else if (store_size == 4) {
|
||||
*(volatile uint32_t*)(si->si_addr) = value;
|
||||
} else if (store_size == 2) {
|
||||
*(volatile uint16_t*)(si->si_addr) = value;
|
||||
} else {
|
||||
D("store size not supported\n");
|
||||
exit(-1);
|
||||
}
|
||||
mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_NONE);
|
||||
}
|
||||
|
||||
void register_sigsegv_handler() {
|
||||
struct sigaction sa = {0};
|
||||
sa.sa_flags = SA_SIGINFO;
|
||||
sigemptyset(&sa.sa_mask);
|
||||
sa.sa_sigaction = handler;
|
||||
if (sigaction(SIGSEGV, &sa, NULL) == -1) {
|
||||
D("ERROR: failed to register sigsegv handler");
|
||||
exit(-1);
|
||||
}
|
||||
// NOTE: python (or ocl runtime?) blocks the SIGSEGV signal
|
||||
sigset_t x;
|
||||
sigemptyset(&x);
|
||||
sigaddset(&x, SIGSEGV);
|
||||
sigprocmask(SIG_UNBLOCK, &x, NULL);
|
||||
}
|
||||
|
||||
int (*my_open)(const char *pathname, int flags, mode_t mode);
|
||||
#undef open
|
||||
int open(const char *pathname, int flags, mode_t mode) {
|
||||
if (my_open == NULL) my_open = reinterpret_cast<decltype(my_open)>(dlsym(RTLD_NEXT, "open"));
|
||||
int ret = my_open(pathname, flags, mode);
|
||||
//D("open %s (0o%o) = %d\n", pathname, flags, ret);
|
||||
files[ret] = pathname;
|
||||
return ret;
|
||||
}
|
||||
|
||||
|
||||
int (*my_open64)(const char *pathname, int flags, mode_t mode);
|
||||
#undef open
|
||||
int open64(const char *pathname, int flags, mode_t mode) {
|
||||
if (my_open64 == NULL) my_open64 = reinterpret_cast<decltype(my_open64)>(dlsym(RTLD_NEXT, "open64"));
|
||||
int ret = my_open64(pathname, flags, mode);
|
||||
//D("open %s (0o%o) = %d\n", pathname, flags, ret);
|
||||
files[ret] = pathname;
|
||||
return ret;
|
||||
}
|
||||
|
||||
void *(*my_mmap)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
|
||||
#undef mmap
|
||||
void *mmap(void *addr, size_t length, int prot, int flags, int fd, off_t offset) {
|
||||
if (my_mmap == NULL) my_mmap = reinterpret_cast<decltype(my_mmap)>(dlsym(RTLD_NEXT, "mmap"));
|
||||
void *ret = my_mmap(addr, length, prot, flags, fd, offset);
|
||||
|
||||
if (doorbell_offset != -1 && offset == doorbell_offset) {
|
||||
D("HIDDEN DOORBELL %p, handled by %p\n", addr, handler);
|
||||
register_sigsegv_handler();
|
||||
mprotect(addr, length, PROT_NONE);
|
||||
}
|
||||
|
||||
if (fd != -1) D("mmapped %p (target %p) with flags 0x%x length 0x%zx fd %d %s offset 0x%lx\n", ret, addr, flags, length, fd, files[fd].c_str(), offset);
|
||||
return ret;
|
||||
}
|
||||
|
||||
void *(*my_mmap64)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
|
||||
#undef mmap64
|
||||
void *mmap64(void *addr, size_t length, int prot, int flags, int fd, off_t offset) { return mmap(addr, length, prot, flags, fd, offset); }
|
||||
|
||||
int ioctl_num = 1;
|
||||
int (*my_ioctl)(int filedes, unsigned long request, void *argp) = NULL;
|
||||
#undef ioctl
|
||||
int ioctl(int filedes, unsigned long request, void *argp) {
|
||||
if (my_ioctl == NULL) my_ioctl = reinterpret_cast<decltype(my_ioctl)>(dlsym(RTLD_NEXT, "ioctl"));
|
||||
int ret = 0;
|
||||
ret = my_ioctl(filedes, request, argp);
|
||||
if (!files.count(filedes)) return ret;
|
||||
|
||||
uint8_t type = (request >> 8) & 0xFF;
|
||||
uint8_t nr = (request >> 0) & 0xFF;
|
||||
uint16_t size = (request >> 16) & 0xFFF;
|
||||
|
||||
D("%3d: %d = %3d(%20s) 0x%3x ", ioctl_num, ret, filedes, files[filedes].c_str(), size);
|
||||
|
||||
if (request == AMDKFD_IOC_SET_EVENT) {
|
||||
kfd_ioctl_set_event_args *args = (kfd_ioctl_set_event_args *)argp;
|
||||
D("AMDKFD_IOC_SET_EVENT event_id:%d", args->event_id);
|
||||
} else if (request == AMDKFD_IOC_ALLOC_MEMORY_OF_GPU) {
|
||||
kfd_ioctl_alloc_memory_of_gpu_args *args = (kfd_ioctl_alloc_memory_of_gpu_args *)argp;
|
||||
D("AMDKFD_IOC_ALLOC_MEMORY_OF_GPU va_addr:0x%llx size:0x%llx handle:%llX gpu_id:0x%x", args->va_addr, args->size, args->handle, args->gpu_id);
|
||||
} else if (request == AMDKFD_IOC_MAP_MEMORY_TO_GPU) {
|
||||
kfd_ioctl_map_memory_to_gpu_args *args = (kfd_ioctl_map_memory_to_gpu_args *)argp;
|
||||
D("AMDKFD_IOC_MAP_MEMORY_TO_GPU handle:%llX", args->handle);
|
||||
} else if (request == AMDKFD_IOC_CREATE_EVENT) {
|
||||
kfd_ioctl_create_event_args *args = (kfd_ioctl_create_event_args *)argp;
|
||||
D("AMDKFD_IOC_CREATE_EVENT event_page_offset:0x%llx event_type:%d event_id:%d", args->event_page_offset, args->event_type, args->event_id);
|
||||
} else if (request == AMDKFD_IOC_WAIT_EVENTS) {
|
||||
D("AMDKFD_IOC_WAIT_EVENTS");
|
||||
} else if (request == AMDKFD_IOC_SET_XNACK_MODE) {
|
||||
D("AMDKFD_IOC_SET_XNACK_MODE");
|
||||
} else if (request == AMDKFD_IOC_SVM || (type == 0x4b && nr == 0x20)) {
|
||||
// NOTE: this one is variable length
|
||||
kfd_ioctl_svm_args *args = (kfd_ioctl_svm_args *)argp;
|
||||
D("AMDKFD_IOC_SVM start_addr:0x%llx size:0x%llx op:%d", args->start_addr, args->size, args->op);
|
||||
} else if (request == AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU) {
|
||||
kfd_ioctl_unmap_memory_from_gpu_args *args = (kfd_ioctl_unmap_memory_from_gpu_args *)argp;
|
||||
D("AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU handle:%llX", args->handle);
|
||||
} else if (request == AMDKFD_IOC_FREE_MEMORY_OF_GPU) {
|
||||
D("AMDKFD_IOC_FREE_MEMORY_OF_GPU");
|
||||
} else if (request == AMDKFD_IOC_SET_SCRATCH_BACKING_VA) {
|
||||
D("AMDKFD_IOC_SET_SCRATCH_BACKING_VA");
|
||||
} else if (request == AMDKFD_IOC_GET_TILE_CONFIG) {
|
||||
D("AMDKFD_IOC_GET_TILE_CONFIG");
|
||||
} else if (request == AMDKFD_IOC_SET_TRAP_HANDLER) {
|
||||
D("AMDKFD_IOC_SET_TRAP_HANDLER");
|
||||
} else if (request == AMDKFD_IOC_GET_VERSION) {
|
||||
kfd_ioctl_get_version_args *args = (kfd_ioctl_get_version_args *)argp;
|
||||
D("AMDKFD_IOC_GET_VERSION major_version:%d minor_version:%d", args->major_version, args->minor_version);
|
||||
} else if (request == AMDKFD_IOC_GET_PROCESS_APERTURES_NEW) {
|
||||
D("AMDKFD_IOC_GET_PROCESS_APERTURES_NEW");
|
||||
} else if (request == AMDKFD_IOC_ACQUIRE_VM) {
|
||||
D("AMDKFD_IOC_ACQUIRE_VM");
|
||||
} else if (request == AMDKFD_IOC_SET_MEMORY_POLICY) {
|
||||
D("AMDKFD_IOC_SET_MEMORY_POLICY");
|
||||
} else if (request == AMDKFD_IOC_GET_CLOCK_COUNTERS) {
|
||||
D("AMDKFD_IOC_GET_CLOCK_COUNTERS");
|
||||
} else if (request == AMDKFD_IOC_CREATE_QUEUE) {
|
||||
kfd_ioctl_create_queue_args *args = (kfd_ioctl_create_queue_args *)argp;
|
||||
D("AMDKFD_IOC_CREATE_QUEUE\n");
|
||||
D("queue_type:%d ring_base_address:0x%llx\n", args->queue_type, args->ring_base_address);
|
||||
D("eop_buffer_address:0x%llx ctx_save_restore_address:0x%llx\n", args->eop_buffer_address, args->ctx_save_restore_address);
|
||||
D("ring_size:0x%x queue_priority:%d\n", args->ring_size, args->queue_priority);
|
||||
D("RETURNS write_pointer_address:0x%llx read_pointer_address:0x%llx doorbell_offset:0x%llx queue_id:%d\n", args->write_pointer_address, args->read_pointer_address, args->doorbell_offset, args->queue_id);
|
||||
//D("RETURNS *write_pointer_address:0x%llx *read_pointer_address:0x%llx\n", *(uint64_t*)args->write_pointer_address, *(uint64_t*)args->read_pointer_address);
|
||||
ring_base_addresses[args->doorbell_offset&0xFFF] = args->ring_base_address;
|
||||
queue_types[args->doorbell_offset&0xFFF] = args->queue_type;
|
||||
doorbell_offset = args->doorbell_offset&~0xFFF;
|
||||
} else {
|
||||
D("type:0x%x nr:0x%x size:0x%x", type, nr, size);
|
||||
}
|
||||
|
||||
D("\n");
|
||||
ioctl_num++;
|
||||
return ret;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
from typing import Callable, Any
|
||||
from tinygrad import Tensor, dtypes, nn, UOp
|
||||
from tinygrad.uop.ops import KernelInfo, AxisType, Ops
|
||||
|
||||
def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
|
||||
fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
|
||||
fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
|
||||
x_abs_max = x.abs().max().detach()
|
||||
scale = fp8_max / (x_abs_max + 1e-8)
|
||||
x_scaled = x * scale
|
||||
x_det = x_scaled.detach()
|
||||
x_clamped = x_det.clamp(fp8_min, fp8_max)
|
||||
x_clamped_ste = x_scaled + (x_clamped - x_det)
|
||||
res = x_clamped_ste.cast(dtype)
|
||||
return res, scale.float().reciprocal()
|
||||
|
||||
def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
|
||||
SEQ = inp.shape[1]
|
||||
OUT = weight.shape[0]
|
||||
IN = weight.shape[-1]
|
||||
seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
|
||||
out_idx = UOp.range(OUT, 3, AxisType.LOOP)
|
||||
batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
|
||||
reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
|
||||
product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
|
||||
reduced = product.reduce(reduce_idx, arg=Ops.ADD)
|
||||
store_op = output.index((seq_idx*OUT+out_idx+batch_idx*OUT*SEQ), ptr=True).store(reduced).end(batch_idx, seq_idx, out_idx)
|
||||
return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))
|
||||
|
||||
def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
|
||||
_, input_uop, weight_uop = kernel.src
|
||||
input_tensor = Tensor(input_uop, device=input_uop.device)
|
||||
grad_tensor = Tensor(gradient, device=gradient.device)
|
||||
weight_tensor = Tensor(weight_uop, device=weight_uop.device)
|
||||
grad_quantized, scale = quantize_to_fp8(grad_tensor)
|
||||
scale_scalar = scale.reshape(())
|
||||
grad_weight = Tensor.einsum("bso,bsi->oi", grad_quantized, input_tensor, dtype=dtypes.float)
|
||||
grad_weight = grad_weight * scale_scalar
|
||||
grad_2d = grad_quantized.reshape(grad_tensor.shape[0] * grad_tensor.shape[1], grad_tensor.shape[-1])
|
||||
grad_input = (grad_2d.dot(weight_tensor, dtype=dtypes.float)).contiguous().reshape(input_tensor.shape) * scale
|
||||
return (None, grad_input.uop, grad_weight.uop)
|
||||
|
||||
class FP8Linear:
|
||||
def __init__(self, in_features:int, out_features:int, bias:bool=True):
|
||||
self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.float32)
|
||||
self.bias = Tensor.empty(out_features, dtype=dtypes.float32) if bias else None
|
||||
|
||||
def __call__(self, x: Tensor) -> Tensor:
|
||||
original_ndim = len(x.shape)
|
||||
if original_ndim == 2: x = x.reshape(x.shape[0], 1, x.shape[1])
|
||||
batch, seq, _ = x.shape
|
||||
w_fp8, w_scale = quantize_to_fp8(self.weight)
|
||||
x_fp8, x_scale = quantize_to_fp8(x)
|
||||
GPUS = self.weight.device
|
||||
if isinstance(GPUS, tuple) and len(GPUS) > 1:
|
||||
y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.multi(0), device=GPUS)
|
||||
else:
|
||||
y = Tensor.empty((batch, seq, self.weight.shape[0]), dtype=dtypes.float)
|
||||
y = Tensor.custom_kernel(y, x_fp8, w_fp8, fxn=custom_matmul, grad_fxn=custom_matmul_backward)[0]
|
||||
y = y * w_scale * x_scale
|
||||
if self.bias is not None: y = y + self.bias
|
||||
if original_ndim == 2: y = y.reshape(batch, self.weight.shape[0])
|
||||
return y.cast(x.dtype)
|
||||
|
||||
def _replace_linear(layer: nn.Linear):
|
||||
fp8_linear = FP8Linear(layer.weight.shape[1], layer.weight.shape[0], layer.bias is not None)
|
||||
fp8_linear.weight = layer.weight
|
||||
if layer.bias is not None: fp8_linear.bias = layer.bias
|
||||
return fp8_linear
|
||||
|
||||
def _swap_linear_with_fp8(model, module_filter_fn:Callable[[Any, str],bool]|None=None, fqn:str="", parent:Any|None=None,
|
||||
attr_name:str="", visited:set|None=None):
|
||||
if visited is None: visited = set()
|
||||
if id(model) in visited: return
|
||||
visited.add(id(model))
|
||||
if isinstance(model, (str, int, float, bool, type(None), Tensor, UOp)): return
|
||||
elif isinstance(model, nn.Linear):
|
||||
if module_filter_fn is not None and not module_filter_fn(model, fqn): return
|
||||
fp8_linear = _replace_linear(model)
|
||||
if parent is not None and attr_name:
|
||||
setattr(parent, attr_name, fp8_linear)
|
||||
elif isinstance(model, list):
|
||||
for i, item in enumerate(model):
|
||||
child_fqn = f"{fqn}.{i}" if fqn else str(i)
|
||||
if isinstance(item, nn.Linear) and (module_filter_fn is None or module_filter_fn(item, child_fqn)): model[i] = _replace_linear(item)
|
||||
else: _swap_linear_with_fp8(item, module_filter_fn, child_fqn, None, "", visited)
|
||||
elif isinstance(model, dict):
|
||||
for key, item in list(model.items()):
|
||||
child_fqn = f"{fqn}.{key}" if fqn else str(key)
|
||||
if isinstance(item, nn.Linear) and (module_filter_fn is None or module_filter_fn(item, child_fqn)): model[key] = _replace_linear(item)
|
||||
else: _swap_linear_with_fp8(item, module_filter_fn, child_fqn, None, "", visited)
|
||||
elif hasattr(model, "__dict__"):
|
||||
for attr_key in list(vars(model).keys()):
|
||||
try: attr = getattr(model, attr_key)
|
||||
except Exception: continue
|
||||
child_fqn = f"{fqn}.{attr_key}" if fqn else attr_key
|
||||
_swap_linear_with_fp8(attr, module_filter_fn, child_fqn, model, attr_key, visited)
|
||||
|
||||
def convert_to_float8_training(model, module_filter_fn:Callable[[Any,str],bool]|None=None):
|
||||
_swap_linear_with_fp8(model, module_filter_fn, "", None, "")
|
||||
return model
|
||||
@@ -0,0 +1,655 @@
|
||||
# RDNA3 128x128 tiled GEMM kernel - DSL version
|
||||
# Computes C = A @ B for 4096x4096 float32 matrices using 128x128 tiles
|
||||
#
|
||||
# Architecture: RDNA3 (gfx1100)
|
||||
# Tile size: 128x128 (each workgroup computes one tile of C)
|
||||
# Workgroup: 128 threads (arranged as 32x4 for coalesced memory access)
|
||||
# Inner loop: 8 iterations per K-block, processing 8 columns of A and 8 rows of B
|
||||
#
|
||||
# Accumulators: 128 vgprs (v[2-117], v[120-124], v[126-129], v[131-133])
|
||||
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from tinygrad import Tensor, Device, Context, GlobalCounters
|
||||
from tinygrad.helpers import getenv, colored
|
||||
from tinygrad.engine.realize import Runner, Estimates, ExecItem
|
||||
from extra.assembly.amd.dsl import s, v, VCC_LO, RawImm, EXEC_LO
|
||||
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
|
||||
ADDR_MASK = 0x3fffff80 # Address alignment mask
|
||||
|
||||
# s_waitcnt encodings: wait for memory operations to complete
|
||||
WAIT_LGKM = 64519 # wait for LDS/GDS/KMEM (lgkm_cnt=0)
|
||||
WAIT_ALL = 0 # wait for everything
|
||||
WAIT_VMEM = 1015 # wait for VMEM only (vm_cnt=0, lgkm_cnt=63)
|
||||
|
||||
# =============================================================================
|
||||
# Named register assignments (VGPRs) - COMPACT LAYOUT
|
||||
# =============================================================================
|
||||
V_LANE_ID_MOD8 = 214 # lane_id & 7 (column within 8-wide tile chunk)
|
||||
V_OUTPUT_ROW = 131 # output row coordinate
|
||||
V_LANE_MOD8_X4 = 134 # V_LANE_ID_MOD8 << 2 (byte offset)
|
||||
V_LANE_DIV8_X4 = 135 # (lane_id >> 3) << 2
|
||||
V_ADDR_HI_ZERO = 136 # always 0 (for 64-bit address high bits)
|
||||
V_LDS_A_BASE = 133 # LDS A-tile base address for inner loop (in ACC_RESERVED gap)
|
||||
V_LDS_B_BASE = 130 # LDS B-tile base address for inner loop (in ACC_RESERVED gap)
|
||||
V_GLOBAL_A_ADDR = 131 # global memory A prefetch address (reuses V_OUTPUT_ROW slot during main loop)
|
||||
V_GLOBAL_B_ADDR = 154 # global memory B prefetch address
|
||||
|
||||
# LDS tile register destinations - SEPARATE from DATA to avoid overlap
|
||||
# DATA regs (v155-170) receive global prefetch
|
||||
# A on banks 2-3, B on banks 0-1 to avoid bank conflicts in VOPD
|
||||
# This layout matches kernel8's optimization for VGPR cache utilization
|
||||
V_A_TILE_REGS = [186, 190, 194, 198] # A tile: banks 2,2,2,2 (186%4=2, 190%4=2, etc.)
|
||||
V_B_TILE_REGS = [184, 188, 192, 196, 200, 204, 208, 212] # B tile: banks 0,0,0,0,0,0,0,0
|
||||
|
||||
# =============================================================================
|
||||
# Named register assignments (SGPRs)
|
||||
# =============================================================================
|
||||
S_OUT_PTR = (0, 1) # output C matrix base pointer
|
||||
S_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_A_PTR = (8, 9) # A matrix base pointer
|
||||
S_B_PTR = (10, 11) # B matrix base pointer
|
||||
S_LOOP_CTR = 12 # loop counter (increments by 8)
|
||||
S_PREFETCH_FLAG = 13 # prefetch condition flag / row stride in epilogue
|
||||
S_WORKGROUP_X = 14 # workgroup_id_x
|
||||
S_WORKGROUP_Y = 15 # workgroup_id_y
|
||||
# Kernarg load destinations (before copy to working regs)
|
||||
S_KERNARG_OUT = (16, 17) # output pointer from kernarg
|
||||
S_KERNARG_A = (20, 21) # A pointer from kernarg
|
||||
S_KERNARG_B = (22, 23) # B pointer from kernarg
|
||||
# Prefetch base pointers (8 pairs each, 16KB/256KB apart)
|
||||
S_PREFETCH_B = 24 # s[24:39] - 8 B tile pointers
|
||||
S_PREFETCH_A = 40 # s[40:55] - 8 A tile pointers
|
||||
|
||||
# =============================================================================
|
||||
# Data tables
|
||||
# =============================================================================
|
||||
|
||||
# Accumulator grid: ACC_GRID[a_idx][b_idx] = vgpr for C[a,b]
|
||||
# a_idx: which A value (0-7), b_idx: which B value (0-15)
|
||||
# Scattered due to VOPD bank constraints (vdst_x % 4 != vdst_y % 4)
|
||||
# Range is from v2 - v129
|
||||
ACC_GRID = [
|
||||
[ 5, 3, 9, 8, 37, 35, 41, 40, 69, 67, 73, 72, 101, 99,105,104], # a0
|
||||
[ 4, 2, 7, 6, 36, 34, 39, 38, 68, 66, 71, 70, 100, 98,103,102], # a1
|
||||
[ 17, 16, 13, 11, 49, 48, 45, 43, 81, 80, 77, 75, 113,112,109,107], # a2
|
||||
[ 15, 14, 12, 10, 47, 46, 44, 42, 79, 78, 76, 74, 111,110,108,106], # a3
|
||||
[ 21, 19, 25, 24, 53, 51, 57, 56, 85, 83, 89, 88, 117,115,121,120], # a4
|
||||
[ 20, 18, 23, 22, 52, 50, 55, 54, 84, 82, 87, 86, 116,114,123,122], # a5
|
||||
[125,128, 29, 27, 33, 32, 61, 59, 65, 64, 93, 91, 97, 96,129,127], # a6
|
||||
[119,118, 28, 26, 31, 30, 60, 58, 63, 62, 92, 90, 95, 94,124,126], # a7
|
||||
]
|
||||
|
||||
# Optimized (a_pair, b_pair) iteration order for better GPU scheduling
|
||||
# Interleaves A and B pairs to maximize instruction-level parallelism
|
||||
FMAC_PAIR_ORDER = [
|
||||
(0,0),(0,1),(1,1),(1,0), (2,0),(2,1),(3,1),(3,2), (0,2),(0,3),(1,3),(1,2), (2,2),(2,3),(3,3),(3,4),
|
||||
(0,4),(0,5),(1,5),(1,4), (2,4),(2,5),(3,5),(3,6), (0,6),(0,7),(1,7),(1,6), (2,6),(2,7),(3,7),(3,0),
|
||||
]
|
||||
|
||||
def derive_fmac_pattern(acc_grid, a_tile_regs=None, b_tile_regs=None):
|
||||
"""Generate 64 dual FMAC ops from accumulator grid with optimized iteration order."""
|
||||
if a_tile_regs is None: a_tile_regs = V_A_TILE_REGS
|
||||
if b_tile_regs is None: b_tile_regs = V_B_TILE_REGS
|
||||
pattern = []
|
||||
for idx, (a_pair, b_pair) in enumerate(FMAC_PAIR_ORDER):
|
||||
a_even, a_odd = a_pair * 2, a_pair * 2 + 1
|
||||
b_even, b_odd = b_pair * 2, b_pair * 2 + 1
|
||||
a_base, b_base = a_tile_regs[a_pair], b_tile_regs[b_pair]
|
||||
# Op 1: normal order -> C[a_even, b_even] + C[a_odd, b_odd]
|
||||
pattern.append((acc_grid[a_even][b_even], acc_grid[a_odd][b_odd],
|
||||
a_base, b_base, a_base+1, b_base+1))
|
||||
# Op 2: alternate swapping A vs B to vary register banks
|
||||
if idx % 2 == 0: # swap B
|
||||
pattern.append((acc_grid[a_even][b_odd], acc_grid[a_odd][b_even],
|
||||
a_base, b_base+1, a_base+1, b_base))
|
||||
else: # swap A
|
||||
pattern.append((acc_grid[a_odd][b_even], acc_grid[a_even][b_odd],
|
||||
a_base+1, b_base, a_base, b_base+1))
|
||||
return pattern
|
||||
|
||||
# Derived: 64 dual FMAC operations
|
||||
FMAC_PATTERN = derive_fmac_pattern(ACC_GRID)
|
||||
|
||||
def derive_permute_swaps(acc_grid, out_regs):
|
||||
"""Derive swap sequence to permute accumulators from FMAC layout to output order.
|
||||
|
||||
After FMAC loop: acc_grid[a][b] holds C[a,b]
|
||||
Output order: for row_half in 0,1; col_group in 0-3; row_in_group in 0-3; b_off in 0-3
|
||||
-> need C[row_half*4 + row_in_group, col_group*4 + b_off] in descending reg order
|
||||
"""
|
||||
def target_ab(i):
|
||||
row_half, col_group = i // 64, (i // 16) % 4
|
||||
row_in_group, b_off = (i // 4) % 4, i % 4
|
||||
return (row_half * 4 + row_in_group, col_group * 4 + b_off)
|
||||
|
||||
reg_contents = {acc_grid[a][b]: (a, b) for a in range(8) for b in range(16)}
|
||||
ab_location = {ab: r for r, ab in reg_contents.items()}
|
||||
|
||||
swaps = []
|
||||
for i in range(128):
|
||||
target_reg, needed_ab = out_regs[i], target_ab(i)
|
||||
current_reg = ab_location[needed_ab]
|
||||
if current_reg != target_reg:
|
||||
swaps.append((current_reg, target_reg))
|
||||
ab_at_target = reg_contents.get(target_reg)
|
||||
reg_contents[target_reg], ab_location[needed_ab] = needed_ab, target_reg
|
||||
if ab_at_target is not None:
|
||||
reg_contents[current_reg], ab_location[ab_at_target] = ab_at_target, current_reg
|
||||
return swaps
|
||||
|
||||
# Derived: swap sequence to arrange accumulators for output
|
||||
OUT_REGS = list(range(129, 1, -1))
|
||||
PERMUTE_SWAPS = derive_permute_swaps(ACC_GRID, OUT_REGS)
|
||||
|
||||
# =============================================================================
|
||||
# LDS tile staging registers - COMPACT LAYOUT
|
||||
# =============================================================================
|
||||
# DATA regs receive contiguous global prefetch, then write to LDS
|
||||
# TILE regs receive scattered LDS loads (ds_load_b64 pairs), then feed FMACs
|
||||
# These are SEPARATE - DATA lives during prefetch/store, TILE lives during inner loop
|
||||
V_LDS_A_ADDR = 153 # single base register for A stores (use +512 offsets)
|
||||
V_LDS_A_DATA = list(range(155, 163)) # 8 data registers for A prefetch (v155-162)
|
||||
V_LDS_B_ADDR = 145 # single base register for B stores (use 16-bit offsets)
|
||||
V_LDS_B_DATA = list(range(163, 171)) # 8 data registers for B prefetch (v163-170)
|
||||
|
||||
# Global memory prefetch schedule: (vdst1, vdst2, addr_vreg, saddr_lo1, saddr_lo2)
|
||||
# First 2 pairs from B prefetch pointers (s[32:39]), next 4 pairs from A prefetch pointers (s[40:55])
|
||||
PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR, S_PREFETCH_B+8+4*i, S_PREFETCH_B+10+4*i) for i in range(2)] + \
|
||||
[(V_LDS_B_DATA[2*(i-2)], V_LDS_B_DATA[2*(i-2)+1], V_GLOBAL_A_ADDR, S_PREFETCH_A+4*(i-2), S_PREFETCH_A+2+4*(i-2)) for i in range(2, 6)]
|
||||
|
||||
# Initial tile prefetch: (vdst, saddr_lo) - load into A data regs using B prefetch pointers (s[24:31])
|
||||
INIT_PREFETCH = [(V_LDS_A_DATA[i], S_PREFETCH_B+2*i) for i in range(4)]
|
||||
|
||||
# Initial tile loads: (vdst, addr_lo) pairs - use temp regs in accumulator gaps
|
||||
INIT_TILE_LOADS = [(23,5),(24,9),(25,7),(26,2),(27,11),(28,13),(29,6),(30,8),(31,10),(12,12),(13,14),(3,2),(4,4),(5,8),(6,6),(7,10)]
|
||||
|
||||
# A matrix row offset registers (scattered to avoid accumulator conflicts)
|
||||
ROW_REGS = list(range(137, 145)) # v137-v144 (8 regs)
|
||||
|
||||
# =============================================================================
|
||||
# Kernel class
|
||||
# =============================================================================
|
||||
|
||||
class Kernel:
|
||||
def __init__(self, arch='gfx1100'):
|
||||
self.instructions, self.labels, self.branch_targets, self.arch = [], {}, {}, arch
|
||||
|
||||
def emit(self, inst): self.instructions.append(inst); return inst
|
||||
def label(self, name): self.labels[name] = len(self.instructions)
|
||||
def branch_to(self, label): self.branch_targets[len(self.instructions) - 1] = label
|
||||
|
||||
def add64(self, dst_lo, dst_hi, src_lo, src_hi, off):
|
||||
"""s[dst_lo:dst_hi] = s[src_lo:src_hi] + off"""
|
||||
if off: self.emit(s_add_u32(s[dst_lo], s[src_lo], off)); self.emit(s_addc_u32(s[dst_hi], s[src_hi], 0))
|
||||
elif dst_lo != src_lo: self.emit(s_mov_b64(s[dst_lo:dst_hi], s[src_lo:src_hi]))
|
||||
|
||||
def global_load(self, vdst, addr, saddr=None):
|
||||
"""Global load b32"""
|
||||
self.emit(global_load_b32(vdst=v[vdst], addr=v[addr:addr+1],
|
||||
saddr=s[saddr:saddr+2] if saddr else RawImm(124)))
|
||||
|
||||
def waitcnt(self, lgkm=None, vm=None):
|
||||
"""Wait for memory operations. lgkm=N waits until N lgkm ops remain, vm=N waits until N vmem ops remain."""
|
||||
from extra.assembly.amd.asm import waitcnt as encode_waitcnt
|
||||
if lgkm == 0 and vm is None: self.emit(s_waitcnt(simm16=WAIT_LGKM))
|
||||
elif vm == 0 and lgkm is None: self.emit(s_waitcnt(simm16=WAIT_VMEM))
|
||||
elif lgkm == 0 and vm == 0: self.emit(s_waitcnt(simm16=WAIT_ALL))
|
||||
elif vm is not None and lgkm is None:
|
||||
self.emit(s_waitcnt(simm16=encode_waitcnt(vmcnt=vm, expcnt=7, lgkmcnt=63)))
|
||||
elif lgkm is not None and vm is None:
|
||||
self.emit(s_waitcnt(simm16=encode_waitcnt(vmcnt=63, expcnt=7, lgkmcnt=lgkm)))
|
||||
else: raise ValueError(f"unsupported waitcnt: lgkm={lgkm}, vm={vm}")
|
||||
|
||||
def barrier(self): self.emit(s_barrier())
|
||||
|
||||
def to_asm(self):
|
||||
import re
|
||||
# Instruction stream with labels
|
||||
label_at = {pos: name for name, pos in self.labels.items()}
|
||||
body = []
|
||||
for i, inst in enumerate(self.instructions):
|
||||
if i in label_at: body.append(f'.{label_at[i]}:')
|
||||
asm = inst.disasm()
|
||||
if i in self.branch_targets:
|
||||
asm = re.sub(r'(s_cbranch_\w+|s_branch)\s+\S+', rf'\1 .{self.branch_targets[i]}', asm)
|
||||
body.append('\t' + asm)
|
||||
|
||||
# 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', 214),
|
||||
('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: 214', ' .wavefront_size: 32', f'amdhsa.target: amdgcn-amd-amdhsa--{self.arch}',
|
||||
'amdhsa.version:', ' - 1', ' - 2', '...', '\t.end_amdgpu_metadata'])
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Kernel builder
|
||||
# =============================================================================
|
||||
|
||||
def build_kernel(arch='gfx1100'):
|
||||
k = Kernel(arch)
|
||||
|
||||
# ===========================================================================
|
||||
# PROLOGUE: Load kernel arguments, compute tile coordinates and addresses
|
||||
# ===========================================================================
|
||||
k.emit(s_load_b128(sdata=s[S_KERNARG_A[0]:S_KERNARG_B[1]], sbase=s[0:1], offset=0x0, soffset=RawImm(124)))
|
||||
k.emit(s_load_b64(sdata=s[S_KERNARG_OUT[0]:S_KERNARG_OUT[1]], sbase=s[0:1], offset=0x10, soffset=RawImm(124)))
|
||||
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))
|
||||
|
||||
# Lane-derived values
|
||||
k.emit(v_and_b32_e32(v[V_LANE_ID_MOD8], 7, v[0]))
|
||||
k.emit(v_lshrrev_b32_e32(v[4], 3, v[0]))
|
||||
k.emit(v_or_b32_e32(v[1], s[S_TILE_X], v[0]))
|
||||
k.emit(v_or_b32_e32(v[22], s[S_TILE_Y], v[4]))
|
||||
k.emit(v_lshlrev_b32_e32(v[V_LANE_MOD8_X4], 2, v[V_LANE_ID_MOD8]))
|
||||
k.emit(v_mov_b32_e32(v[2], 0)) # v[1] always positive, sign extension is 0
|
||||
k.emit(v_lshlrev_b64(v[5:6], 2, v[1:2]))
|
||||
k.waitcnt(lgkm=0)
|
||||
|
||||
# Copy pointers to working registers
|
||||
k.emit(s_mov_b64(s[S_OUT_PTR[0]:S_OUT_PTR[1]], s[S_KERNARG_OUT[0]:S_KERNARG_OUT[1]]))
|
||||
k.emit(s_mov_b64(s[S_A_PTR[0]:S_A_PTR[1]], s[S_KERNARG_A[0]:S_KERNARG_A[1]]))
|
||||
k.emit(s_mov_b64(s[S_B_PTR[0]:S_B_PTR[1]], s[S_KERNARG_B[0]:S_KERNARG_B[1]]))
|
||||
|
||||
# Compute 8 A and B matrix tile base pointers for prefetch
|
||||
for i in range(8): k.add64(S_PREFETCH_B + i*2, S_PREFETCH_B + i*2 + 1, S_KERNARG_B[0], S_KERNARG_B[1], i * 0x4000) # B: 16KB apart
|
||||
for i in range(8): k.add64(S_PREFETCH_A + i*2, S_PREFETCH_A + i*2 + 1, S_KERNARG_A[0], S_KERNARG_A[1], i * 0x40000) # A: 256KB apart
|
||||
|
||||
# Global prefetch addresses: B = (tile_x + lane_id) * 4, A = ((tile_y << 12) + (lane_id/8)*4K + lane_id%8) * 4
|
||||
k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], s[S_TILE_X], v[0]))
|
||||
k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_B_ADDR], 2, v[V_GLOBAL_B_ADDR]))
|
||||
k.emit(s_lshl_b32(s[19], s[S_TILE_Y], 12))
|
||||
k.emit(v_lshl_add_u32(v[V_GLOBAL_A_ADDR], v[4], 12, v[V_LANE_ID_MOD8])) # (lane_id/8)*4K + lane_id%8
|
||||
k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], s[19], v[V_GLOBAL_A_ADDR]))
|
||||
k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_A_ADDR], 2, v[V_GLOBAL_A_ADDR]))
|
||||
|
||||
# ===========================================================================
|
||||
# Tile address computation for initial A/B matrix loads
|
||||
# ===========================================================================
|
||||
k.emit(s_lshl_b32(s[S_LOOP_BOUND], s[S_DIM_N], 4)) # row stride = 16*N
|
||||
k.emit(v_mul_lo_u32(v[ROW_REGS[0]], v[22], s[S_DIM_N])) # A matrix row offsets
|
||||
for i in range(1, 8): k.emit(v_add_nc_u32_e32(v[ROW_REGS[i]], s[S_LOOP_BOUND], v[ROW_REGS[i-1]]))
|
||||
|
||||
def addr64(dst, base_s): # 64-bit address: v[dst:dst+1] = s[base_s:base_s+1] + v[dst]*4
|
||||
k.emit(v_mov_b32_e32(v[dst+1], 0)) # offset always positive, sign ext = 0
|
||||
k.emit(v_lshlrev_b64(v[dst:dst+1], 2, v[dst:dst+1]))
|
||||
k.emit(v_add_co_u32(v[dst], VCC_LO, s[base_s], v[dst]))
|
||||
k.emit(v_add_co_ci_u32_e32(v[dst+1], s[base_s+1], v[dst+1]))
|
||||
|
||||
def b_addr(dst, mult, tmp=None): # B address for col + mult*N
|
||||
tmp = tmp if tmp is not None else dst
|
||||
k.emit(v_mad_u32_u24(v[tmp], s[S_DIM_N], mult, v[1]))
|
||||
if tmp != dst:
|
||||
k.emit(v_mov_b32_e32(v[tmp+1], 0)) # offset always positive
|
||||
k.emit(v_lshlrev_b64(v[dst:dst+1], 2, v[tmp:tmp+1]))
|
||||
k.emit(v_add_co_u32(v[dst], VCC_LO, s[S_B_PTR[0]], v[dst]))
|
||||
k.emit(v_add_co_ci_u32_e32(v[dst+1], s[S_B_PTR[1]], v[dst+1]))
|
||||
else: addr64(dst, S_B_PTR[0])
|
||||
|
||||
def a_addr(dst, row_reg, tmp): # A address for row_reg + lane_id_mod8
|
||||
k.emit(v_add_nc_u32_e32(v[tmp], v[row_reg], v[V_LANE_ID_MOD8]))
|
||||
k.emit(v_mov_b32_e32(v[tmp+1], 0)) # offset always positive
|
||||
k.emit(v_lshlrev_b64(v[dst:dst+1], 2, v[tmp:tmp+1]))
|
||||
k.emit(v_add_co_u32(v[dst], VCC_LO, s[S_A_PTR[0]], v[dst]))
|
||||
k.emit(v_add_co_ci_u32_e32(v[dst+1], s[S_A_PTR[1]], v[dst+1]))
|
||||
|
||||
# Batch 1: B addresses (cols 0-5) and loads
|
||||
k.emit(v_add_co_u32(v[5], VCC_LO, s[S_B_PTR[0]], v[5]))
|
||||
k.emit(v_add_co_ci_u32_e32(v[6], s[S_B_PTR[1]], v[6]))
|
||||
for dst, mult in [(9,1), (7,2), (2,3), (11,4), (13,5)]: b_addr(dst, mult)
|
||||
k.emit(s_clause(simm16=5)) # 6 consecutive global loads
|
||||
for vdst, addr in INIT_TILE_LOADS[:6]: k.global_load(vdst, addr)
|
||||
|
||||
# Batch 2: A addresses (rows 0-4) and loads
|
||||
for dst, ri in [(6,0), (8,1), (10,2), (12,3), (14,4)]:
|
||||
k.emit(v_add_nc_u32_e32(v[dst], v[ROW_REGS[ri]], v[V_LANE_ID_MOD8]))
|
||||
addr64(dst, S_A_PTR[0])
|
||||
k.emit(s_clause(simm16=4)) # 5 consecutive global loads
|
||||
for vdst, addr in INIT_TILE_LOADS[6:11]: k.global_load(vdst, addr)
|
||||
|
||||
# Batch 3: B cols 6-7, A rows 5-7, and loads
|
||||
for dst, mult, tmp in [(2,6,15), (4,7,4)]: b_addr(dst, mult, tmp)
|
||||
for dst, ri, tmp in [(8,5,16), (6,6,18), (10,7,20)]: a_addr(dst, ROW_REGS[ri], tmp)
|
||||
k.emit(s_clause(simm16=4)) # 5 consecutive global loads
|
||||
for vdst, addr in INIT_TILE_LOADS[11:]: k.global_load(vdst, addr)
|
||||
|
||||
# ===========================================================================
|
||||
# LDS store address computation (bank-conflict-avoiding swizzle)
|
||||
# ===========================================================================
|
||||
# This section computes LDS store addresses with a swizzle pattern to avoid bank conflicts.
|
||||
# Key outputs:
|
||||
# v[8]: A-tile initial store base (used only for initial stores with stride64)
|
||||
# V_LDS_B_ADDR (v145): B-tile store base (used for both initial and main loop)
|
||||
# V_LANE_DIV8_X4 (v135): (lane_id >> 3) << 2 for epilogue
|
||||
#
|
||||
# The swizzle ensures that threads in the same wavefront write to different LDS banks.
|
||||
# Formula: swizzled_addr = base + (lane_id & 7) * LDS_A_STRIDE + swizzle_offset
|
||||
# where swizzle_offset depends on (lane_id >> 3) to distribute across banks.
|
||||
|
||||
# v[22] = tile_y | (lane_id >> 3) from prologue, used as base for row offsets
|
||||
# Compute 7 row offsets for B-tile rows 1-7 (row 0 computed separately in v[9])
|
||||
k.emit(v_add_nc_u32_e32(v[9], s[S_LOOP_CTR], v[22])) # row 0 base (S_LOOP_CTR=0)
|
||||
for i in range(7): k.emit(v_or_b32_e32(v[10 + i if i < 2 else 12 + i], 16 * (i + 1), v[22])) # rows 1-7
|
||||
|
||||
# Extract sign bit of workgroup_x (always 0 for valid workgroups, used for masking)
|
||||
k.emit(s_bfe_i32(s[S_LOOP_BOUND], s[S_WORKGROUP_X], 0x10018))
|
||||
k.emit(v_and_b32_e32(v[9], ADDR_MASK, v[9]))
|
||||
k.emit(s_lshr_b32(s[S_LOOP_BOUND], s[S_LOOP_BOUND], 25))
|
||||
|
||||
# Compute masked row offsets for bank conflict avoidance pattern
|
||||
# Pattern: v[row] = row_val - (row_val & ADDR_MASK) extracts lower bits
|
||||
k.emit(v_add_nc_u32_e32(v[19], s[S_LOOP_CTR], v[10]))
|
||||
k.emit(v_add_nc_u32_e32(v[8], s[S_LOOP_BOUND], v[1])) # A-tile base computation
|
||||
for d, r in zip([20, 21, 32, 33, 34, 35], [11, 14, 15, 16, 17, 18]):
|
||||
k.emit(v_add_nc_u32_e32(v[d], s[S_LOOP_CTR], v[r]))
|
||||
k.emit(v_and_b32_e32(v[8], ADDR_MASK, v[8]))
|
||||
k.emit(v_sub_nc_u32_e32(v[9], v[22], v[9])) # row 0 swizzle offset
|
||||
for d, s_ in zip([19, 20, 21, 22, 32, 33, 34], [20, 21, 22, 32, 33, 34, 35]):
|
||||
k.emit(v_and_b32_e32(v[d], ADDR_MASK, v[s_]))
|
||||
k.emit(v_sub_nc_u32_e32(v[8], v[1], v[8])) # A-tile swizzle
|
||||
|
||||
# Apply swizzle offsets and scale to byte offsets
|
||||
k.emit(v_lshlrev_b32_e32(v[9], 2, v[9])) # row 0 offset * 4
|
||||
for r, t in zip([10, 11, 14, 15, 16, 17, 18], [19, 20, 21, 22, 32, 33, 34]):
|
||||
k.emit(v_sub_nc_u32_e32(v[r], v[r], v[t])) # rows 1-7 swizzle
|
||||
k.emit(v_bfe_u32(v[2], v[0], 3, 2)) # v[2] = (lane_id >> 3) & 3
|
||||
k.emit(v_lshlrev_b32_e32(v[8], 2, v[8])) # A-tile base * 4
|
||||
|
||||
# Compute B-tile base address: LDS_A_STRIDE * (lane_id % 8) + row0_offset
|
||||
k.emit(v_mad_u32_u24(v[V_LDS_B_ADDR], LDS_A_STRIDE, v[V_LANE_ID_MOD8], v[9]))
|
||||
# Scale row offsets 1-7 to byte offsets (row 0 already in v[9])
|
||||
for d, r in zip([9, 10, 11, 14, 15, 16, 17], [10, 11, 14, 15, 16, 17, 18]):
|
||||
k.emit(v_lshlrev_b32_e32(v[d], 2, v[r]))
|
||||
k.emit(v_lshlrev_b32_e32(v[V_LANE_DIV8_X4], 2, v[2]))
|
||||
k.emit(v_add_nc_u32_e32(v[8], 0x80, v[8])) # A-tile initial store base + 128
|
||||
|
||||
# Store initial tile data to LDS
|
||||
k.waitcnt(vm=0)
|
||||
for i, (d0, d1) in enumerate([(0,1), (2,3), (4,5), (11,12)]):
|
||||
k.emit(ds_store_2addr_stride64_b32(addr=v[8], data0=v[INIT_TILE_LOADS[d0][0]], data1=v[INIT_TILE_LOADS[d1][0]], offset0=16+i*4, offset1=18+i*4))
|
||||
# B stores: single base with offsets 0,64,128,192,256,320,384,448
|
||||
for i, idx in enumerate([6,7,8,9,10,13,14,15]):
|
||||
offset = i * 64
|
||||
k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[INIT_TILE_LOADS[idx][0]], offset0=offset & 0xFF, offset1=offset >> 8))
|
||||
k.waitcnt(lgkm=0)
|
||||
k.barrier()
|
||||
|
||||
# ===========================================================================
|
||||
# INIT: Compute LDS base addresses, then zero accumulators
|
||||
# ===========================================================================
|
||||
# v[3] = v[1] & 0x7F (lower 7 bits) since S_LOOP_BOUND=0 for valid workgroups
|
||||
k.emit(v_lshlrev_b32_e32(v[2], 4, v[2]))
|
||||
k.emit(v_add_nc_u32_e32(v[3], s[S_LOOP_BOUND], v[1]))
|
||||
k.emit(v_and_b32_e32(v[3], ADDR_MASK, v[3]))
|
||||
k.emit(v_sub_nc_u32_e32(v[3], v[1], v[3]))
|
||||
k.emit(v_lshl_or_b32(v[V_LDS_B_BASE], v[V_LANE_ID_MOD8], 4, LDS_BASE_OFFSET))
|
||||
k.emit(v_lshl_add_u32(v[V_LDS_A_ADDR], v[3], 2, LDS_BASE_OFFSET))
|
||||
k.emit(v_lshlrev_b32_e32(v[3], 2, v[0]))
|
||||
k.emit(v_and_or_b32(v[V_LDS_A_BASE], 0x180, v[3], v[2]))
|
||||
|
||||
# Zero all 128 accumulators using VOPD dual moves (64 instructions instead of 128)
|
||||
for i in range(0, len(OUT_REGS), 2):
|
||||
k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[OUT_REGS[i]], vdsty=v[OUT_REGS[i+1]], srcx0=0, srcy0=0))
|
||||
|
||||
k.emit(s_add_i32(s[S_LOOP_BOUND], s[S_DIM_N], -8))
|
||||
k.emit(s_add_u32(s[S_A_PTR[0]], s[S_A_PTR[0]], 32))
|
||||
k.emit(s_addc_u32(s[S_A_PTR[1]], s[S_A_PTR[1]], 0))
|
||||
# S_LOOP_CTR is already 0 from prologue initialization
|
||||
k.emit(s_branch(simm16=0)); k.branch_to('LOOP_ENTRY')
|
||||
|
||||
# ===========================================================================
|
||||
# MAIN GEMM LOOP
|
||||
# ===========================================================================
|
||||
|
||||
NO_DS, NO_GLOBAL = getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
|
||||
|
||||
k.label('LOOP_INC')
|
||||
k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
|
||||
k.emit(s_cmp_ge_i32(s[S_LOOP_CTR], s[S_DIM_N]))
|
||||
k.emit(s_cbranch_scc1(simm16=0)); k.branch_to('EPILOGUE')
|
||||
|
||||
k.label('LOOP_ENTRY')
|
||||
k.emit(s_cmp_lt_i32(s[S_LOOP_CTR], s[S_LOOP_BOUND]))
|
||||
k.emit(s_cselect_b32(s[S_PREFETCH_FLAG], -1, 0)) # s_cselect doesn't modify SCC
|
||||
k.emit(s_cbranch_scc0(simm16=0)); k.branch_to('SKIP_PREFETCH') # branch if loop_ctr >= loop_bound
|
||||
|
||||
# Advance prefetch pointers
|
||||
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]))
|
||||
|
||||
if not NO_GLOBAL:
|
||||
for vdst, saddr_lo in INIT_PREFETCH:
|
||||
k.global_load(vdst, V_GLOBAL_B_ADDR, saddr_lo)
|
||||
|
||||
k.label('SKIP_PREFETCH')
|
||||
|
||||
# 8 inner loop iterations
|
||||
for iter in range(8):
|
||||
# Load A tile (4 pairs) and B tile (8 pairs) from LDS
|
||||
if not NO_DS:
|
||||
k.emit(s_clause(simm16=11)) # 12 loads total: 4 A + 8 B
|
||||
# A tile: 4 ds_load_b64
|
||||
for i, vdst in enumerate(V_A_TILE_REGS):
|
||||
a_off = (i & 1) * 8 + (i >> 1) * 64 + iter * LDS_A_STRIDE
|
||||
k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_A_BASE], offset0=a_off & 0xFF, offset1=a_off >> 8))
|
||||
# B tile: 8 ds_load_b64
|
||||
for i, vdst in enumerate(V_B_TILE_REGS):
|
||||
b_off = (i & 1) * 8 + (i & 2) * 64 + (i >> 2) * 256 + iter * LDS_B_STRIDE
|
||||
k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_B_BASE], offset0=b_off & 0xFF, offset1=b_off >> 8))
|
||||
k.waitcnt(lgkm=0)
|
||||
|
||||
# 64 dual FMACs
|
||||
k.emit(s_clause(simm16=63))
|
||||
for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
|
||||
k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
|
||||
vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))
|
||||
|
||||
# Issue global prefetch AFTER FMACs (first 6 iterations only)
|
||||
if iter < 6 and not NO_GLOBAL:
|
||||
vdst1, vdst2, addr, slo1, slo2 = PREFETCH_LOADS[iter]
|
||||
k.global_load(vdst1, addr, slo1)
|
||||
k.global_load(vdst2, addr, slo2)
|
||||
|
||||
k.emit(s_and_not1_b32(VCC_LO, EXEC_LO, s[S_PREFETCH_FLAG]))
|
||||
k.waitcnt(vm=0)
|
||||
k.barrier()
|
||||
k.emit(s_cbranch_vccnz(simm16=0)); k.branch_to('LOOP_INC')
|
||||
|
||||
# Store prefetched data to LDS
|
||||
# NOTE: Register naming reflects LDS tile organization, not source matrix:
|
||||
# V_LDS_A_DATA (v155-162) holds data that goes to LDS A-tile region
|
||||
# V_LDS_B_DATA (v163-170) holds data that goes to LDS B-tile region
|
||||
# The data sources are swapped: A-tile receives B matrix rows, B-tile receives A matrix columns
|
||||
for i in range(4): # A tile: 8 values via 4 stride64 stores
|
||||
k.emit(ds_store_2addr_stride64_b32(addr=v[V_LDS_A_ADDR], data0=v[V_LDS_A_DATA[i*2]], data1=v[V_LDS_A_DATA[i*2+1]], offset0=i*4, offset1=i*4+2))
|
||||
for i in range(8): # B tile: 8 values via 8 scalar stores with 64-byte spacing
|
||||
offset = i * 64
|
||||
k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[V_LDS_B_DATA[i]], offset0=offset & 0xFF, offset1=offset >> 8))
|
||||
|
||||
k.waitcnt(lgkm=0)
|
||||
k.barrier()
|
||||
k.emit(s_branch(simm16=0)); k.branch_to('LOOP_INC')
|
||||
|
||||
# ===========================================================================
|
||||
# EPILOGUE: Permute and store results
|
||||
# ===========================================================================
|
||||
k.label('EPILOGUE')
|
||||
|
||||
# Rearrange accumulators from FMAC layout to contiguous output order
|
||||
for a, b in PERMUTE_SWAPS:
|
||||
k.emit(v_swap_b32_e32(v[a], v[b]))
|
||||
|
||||
# Compute output coordinates: v[V_LANE_ID_MOD8] = col, v[V_OUTPUT_ROW] = row
|
||||
k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32,
|
||||
vdstx=v[149], vdsty=v[150], srcx0=v[V_LANE_MOD8_X4], vsrcx1=v[0], srcy0=v[V_LANE_DIV8_X4], vsrcy1=v[0]))
|
||||
k.emit(v_and_b32_e32(v[0], 0x60, v[0]))
|
||||
k.emit(v_or_b32_e32(v[V_LANE_ID_MOD8], s[S_TILE_X], v[149]))
|
||||
k.emit(v_add_nc_u32_e32(v[0], s[S_TILE_Y], v[0]))
|
||||
k.emit(v_or_b32_e32(v[V_OUTPUT_ROW], v[0], v[150]))
|
||||
|
||||
# Precompute row offsets: v[144-147] for rows 0-3, v[148-151] for rows 16-19
|
||||
for base, row_off in [(144, 0), (148, 16)]:
|
||||
if row_off: k.emit(v_or_b32_e32(v[1], row_off, v[V_OUTPUT_ROW]))
|
||||
k.emit(v_mul_lo_u32(v[base], v[1] if row_off else v[V_OUTPUT_ROW], s[S_DIM_N]))
|
||||
for i in range(3): k.emit(v_add_nc_u32_e32(v[base + 1 + i], s[S_DIM_N], v[base + i]))
|
||||
|
||||
k.emit(v_mov_b32_e32(v[V_ADDR_HI_ZERO], 0))
|
||||
k.emit(s_lshl_b32(s[S_PREFETCH_FLAG], s[S_DIM_N], 2)) # row stride in bytes
|
||||
|
||||
# Store 128 output values as 32 groups of 4 (128-bit stores)
|
||||
# Layout: 2 row halves (0-3, 16-19) x 4 col groups x 4 rows = 32 stores of 4 floats
|
||||
epilogue_reserved = {V_LANE_ID_MOD8, V_OUTPUT_ROW, V_LANE_MOD8_X4, V_LANE_DIV8_X4, V_ADDR_HI_ZERO}
|
||||
|
||||
for i, (row_half, col_off, row_in_group) in enumerate([(rh, co, ri)
|
||||
for rh in range(2) for co in [0, 32, 64, 96] for ri in range(4)]):
|
||||
row = row_half * 16 + row_in_group
|
||||
srcs = OUT_REGS[i*4:(i+1)*4]
|
||||
|
||||
# Find temp register for scaled values (must not conflict with reserved regs)
|
||||
tmp = max(srcs) + 5
|
||||
while any(r in epilogue_reserved for r in range(tmp, tmp + 4)): tmp += 1
|
||||
|
||||
# Copy values to temp regs for output (alpha=1.0 hardcoded, so just move)
|
||||
for j, src in enumerate(srcs):
|
||||
k.emit(v_mov_b32_e32(v[tmp + j], v[src]))
|
||||
|
||||
# Compute output address
|
||||
if row_in_group == 0: # first row: compute base address for this column group
|
||||
if col_off == 0: k.emit(v_mov_b32_e32(v[0], v[V_LANE_ID_MOD8]))
|
||||
else: k.emit(v_add_nc_u32_e32(v[0], col_off, v[V_LANE_ID_MOD8]))
|
||||
row_base = 144 + row if row < 4 else 148 + row - 16
|
||||
k.emit(v_add_nc_u32_e32(v[0], v[row_base], v[0]))
|
||||
k.emit(v_lshlrev_b32_e32(v[0], 2, v[0]))
|
||||
k.emit(v_add_co_u32(v[0], VCC_LO, s[S_OUT_PTR[0]], v[0]))
|
||||
k.emit(v_add_co_ci_u32_e32(v[1], s[S_OUT_PTR[1]], v[V_ADDR_HI_ZERO]))
|
||||
else: # subsequent rows: just add stride
|
||||
k.emit(v_add_co_u32(v[0], VCC_LO, s[S_PREFETCH_FLAG], v[0]))
|
||||
k.emit(v_add_co_ci_u32_e32(v[1], v[1], v[V_ADDR_HI_ZERO]))
|
||||
|
||||
k.emit(global_store_b128(addr=v[0:1], data=v[tmp:tmp+3], saddr=RawImm(124)))
|
||||
|
||||
k.emit(s_sendmsg(simm16=3))
|
||||
k.emit(s_endpgm())
|
||||
|
||||
return k.to_asm()
|
||||
|
||||
# =============================================================================
|
||||
# Test harness
|
||||
# =============================================================================
|
||||
|
||||
N = getenv("N", 4096)
|
||||
BLOCK_M, BLOCK_N = 128, 128
|
||||
THREADS = 128
|
||||
|
||||
def test_matmul():
|
||||
dev = Device[Device.DEFAULT]
|
||||
print(f"Device arch: {dev.arch}")
|
||||
|
||||
if getenv("STOCK", 0):
|
||||
# Load the stock kernel from amd_seb/kernel8_batched_gmem.s
|
||||
stock_path = Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s"
|
||||
asm = stock_path.read_text()
|
||||
print(f"Loaded stock kernel from {stock_path}")
|
||||
else:
|
||||
asm = build_kernel(dev.arch)
|
||||
if getenv("PRINT_ASM", 0): print(asm)
|
||||
|
||||
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)
|
||||
b = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
|
||||
c = Tensor.empty(N, N)
|
||||
Tensor.realize(a, b, c)
|
||||
|
||||
grid, local = (N // BLOCK_N, N // BLOCK_M, 1), (THREADS, 1, 1)
|
||||
print(f"Grid: {grid}, Local: {local}")
|
||||
|
||||
_prg = dev.runtime("kernel", binary)
|
||||
class AsmRunner(Runner):
|
||||
def __init__(self):
|
||||
super().__init__(colored("kernel", "cyan"), Device.DEFAULT, Estimates(ops=N*N*N*2, mem=N*N*4*3))
|
||||
def __call__(self, rawbufs, var_vals, wait=False):
|
||||
c_buf, a_buf, b_buf = [x.ensure_allocated()._buf for x in rawbufs]
|
||||
return _prg(a_buf, b_buf, c_buf, global_size=grid, local_size=local, wait=wait)
|
||||
|
||||
ei = ExecItem(None, [c.uop.buffer, a.uop.buffer, b.uop.buffer], prg=AsmRunner())
|
||||
|
||||
ets = []
|
||||
with Context(DEBUG=2):
|
||||
for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True))
|
||||
print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
|
||||
|
||||
if getenv("VERIFY", 1):
|
||||
GlobalCounters.reset()
|
||||
with Context(DEBUG=2): tc = (a @ b).realize()
|
||||
with Context(DEBUG=0): err = (c - tc).square().mean().item()
|
||||
print(f"mean squared error {err}")
|
||||
if err > 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__":
|
||||
if getenv("ASM", 0): print(build_kernel(Device[Device.DEFAULT].arch))
|
||||
elif getenv("SQTT", 0): run_sqtt()
|
||||
else: test_matmul()
|
||||
@@ -140,11 +140,11 @@ def hand_spec_kernel3():
|
||||
|
||||
return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
|
||||
|
||||
def test_matmul(sink:UOp, N=N):
|
||||
def test_matmul(sink:UOp, dtype=dtypes.float32, N=N):
|
||||
rng = np.random.default_rng()
|
||||
a = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
|
||||
b = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
|
||||
hc = Tensor.empty(N, N)
|
||||
a = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
|
||||
b = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
|
||||
hc = Tensor.empty(N, N, dtype=dtype)
|
||||
Tensor.realize(a, b, hc)
|
||||
|
||||
ei = ExecItem(sink, [t.uop.buffer for t in [hc, a, b]], prg=get_runner(Device.DEFAULT, sink))
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,78 @@
|
||||
.text
|
||||
.section .text.
|
||||
.global gemm
|
||||
.p2align 8
|
||||
.type gemm,@function
|
||||
|
||||
gemm:
|
||||
INSTRUCTIONS
|
||||
|
||||
.section .rodata,"a",@progbits
|
||||
.p2align 6, 0x0
|
||||
.amdhsa_kernel gemm
|
||||
# basic memory requirements
|
||||
.amdhsa_group_segment_fixed_size 133120
|
||||
.amdhsa_private_segment_fixed_size 0
|
||||
.amdhsa_kernarg_size 28
|
||||
# register usage (RSRC1)
|
||||
.amdhsa_next_free_vgpr 504
|
||||
.amdhsa_next_free_sgpr 96
|
||||
# workgroup / workitem IDs (RSRC2)
|
||||
.amdhsa_system_sgpr_workgroup_id_x 1
|
||||
.amdhsa_system_sgpr_workgroup_id_y 1
|
||||
.amdhsa_system_sgpr_workgroup_id_z 1
|
||||
# user SGPRs, we only specify the kernel args ptr in s[0:1]
|
||||
.amdhsa_user_sgpr_kernarg_segment_ptr 1
|
||||
.amdhsa_user_sgpr_count 2
|
||||
.amdhsa_user_sgpr_kernarg_preload_length 0
|
||||
.amdhsa_user_sgpr_kernarg_preload_offset 0
|
||||
# gfx90a / gfx940 specifics (RSRC3)
|
||||
.amdhsa_accum_offset 248
|
||||
.amdhsa_uses_dynamic_stack 0
|
||||
.amdhsa_tg_split 0
|
||||
.end_amdhsa_kernel
|
||||
|
||||
.amdgpu_metadata
|
||||
---
|
||||
amdhsa.kernels:
|
||||
- .name: gemm
|
||||
.symbol: gemm.kd
|
||||
.args:
|
||||
- .name: C
|
||||
.address_space: global
|
||||
.offset: 0
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: bf16
|
||||
- .name: B
|
||||
.address_space: global
|
||||
.offset: 8
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: bf16
|
||||
- .name: A
|
||||
.address_space: global
|
||||
.offset: 16
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: bf16
|
||||
- .name: sz
|
||||
.offset: 24
|
||||
.size: 4
|
||||
.value_kind: by_value
|
||||
.value_type: u32
|
||||
.group_segment_fixed_size: 133120
|
||||
.private_segment_fixed_size: 0
|
||||
.kernarg_segment_align: 8
|
||||
.kernarg_segment_size: 28
|
||||
.max_flat_workgroup_size: 256
|
||||
.sgpr_count: 88
|
||||
.sgpr_spill_count: 0
|
||||
.vgpr_count: 248
|
||||
.vgpr_spill_count: 0
|
||||
.wavefront_size: 64
|
||||
amdhsa.version:
|
||||
- 1
|
||||
- 0
|
||||
...
|
||||
.end_amdgpu_metadata
|
||||
@@ -0,0 +1,72 @@
|
||||
# Run assembly on the AMD runtime and check correctness
|
||||
# VIZ=2 to profile
|
||||
import pathlib
|
||||
from tinygrad import Tensor, Device, dtypes, Context
|
||||
from tinygrad.uop.ops import UOp, Ops, KernelInfo
|
||||
from tinygrad.helpers import getenv
|
||||
|
||||
fp = pathlib.Path(__file__).parent/"gemm.s"
|
||||
|
||||
N = getenv("N", 8192)
|
||||
THREADS_PER_WG = 256
|
||||
NUM_WG = N//THREADS_PER_WG * N//THREADS_PER_WG
|
||||
|
||||
assert N % THREADS_PER_WG == 0, "N must be divisible by THREADS_PER_WG"
|
||||
|
||||
# ** generate inputs on CPU
|
||||
|
||||
scale = 10.0
|
||||
|
||||
import torch
|
||||
torch.manual_seed(0)
|
||||
A = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
|
||||
B = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
|
||||
Bt = B.t().contiguous() # transpose B for the asm gemm
|
||||
C_torch = A@B
|
||||
|
||||
# ** copy buffers to AMD
|
||||
|
||||
# input creation and validation run on the copy engine for simpler tracing
|
||||
|
||||
def from_torch(t:torch.Tensor) -> Tensor:
|
||||
return Tensor.from_blob(t.data_ptr(), t.shape, dtype=dtypes.bfloat16, device="cpu").to(Device.DEFAULT).realize()
|
||||
|
||||
C_tiny = from_torch(A) @ from_torch(B)
|
||||
C_asm = Tensor.empty_like(C_tiny)
|
||||
|
||||
# ** assembly custom kernel
|
||||
|
||||
def custom_asm_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
|
||||
lidx = UOp.special(THREADS_PER_WG, "lidx0")
|
||||
gidx = UOp.special(NUM_WG, "gidx0")
|
||||
|
||||
src = (pathlib.Path(__file__).parent/"template.s").read_text().replace("INSTRUCTIONS", fp.read_text())
|
||||
|
||||
sz = UOp.variable("SZ", 256, 8192)
|
||||
|
||||
sink = UOp.sink(C.base, A.base, B.base, sz, lidx, gidx, arg=KernelInfo(name="gemm"))
|
||||
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)), arg=())
|
||||
|
||||
C_asm = Tensor.custom_kernel(C_asm, from_torch(A), from_torch(Bt), fxn=custom_asm_gemm)[0]
|
||||
|
||||
# ** run gemms
|
||||
|
||||
sched = Tensor.schedule(C_tiny, C_asm)
|
||||
eis = [si.lower() for si in sched]
|
||||
|
||||
with Context(DEBUG=2):
|
||||
for ei in eis:
|
||||
et = ei.run({"SZ":N}, wait=True)
|
||||
print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
|
||||
|
||||
# ** correctness
|
||||
|
||||
import ctypes
|
||||
|
||||
def torch_bf16(t:Tensor) -> torch.tensor:
|
||||
asm_out = t.to("cpu").realize().uop.buffer._buf
|
||||
buf = (ctypes.c_uint16*C_asm.uop.size).from_address(asm_out.va_addr)
|
||||
return torch.frombuffer(buf, dtype=torch.bfloat16, count=C_asm.uop.size).reshape(C_asm.shape)
|
||||
|
||||
assert torch.allclose(torch_bf16(C_asm), C_torch, rtol=1e-2, atol=1e-3)
|
||||
assert torch.allclose(torch_bf16(C_tiny), C_torch, rtol=1e-2, atol=1e-3)
|
||||
-15338
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,76 @@
|
||||
.text
|
||||
.section .text.
|
||||
.global gemm
|
||||
.p2align 8
|
||||
.type gemm,@function
|
||||
|
||||
gemm:
|
||||
INSTRUCTIONS
|
||||
|
||||
.section .rodata,"a",@progbits
|
||||
.p2align 6, 0x0
|
||||
.amdhsa_kernel gemm
|
||||
# basic memory requirements
|
||||
.amdhsa_group_segment_fixed_size 30336
|
||||
.amdhsa_private_segment_fixed_size 0
|
||||
.amdhsa_kernarg_size 32
|
||||
# register usage (RSRC1)
|
||||
.amdhsa_next_free_vgpr 256
|
||||
.amdhsa_next_free_sgpr 100
|
||||
# workgroup / workitem IDs (RSRC2)
|
||||
.amdhsa_system_sgpr_workgroup_id_x 1
|
||||
.amdhsa_system_sgpr_workgroup_id_y 1
|
||||
.amdhsa_system_sgpr_workgroup_id_z 1
|
||||
# user SGPRs: kernarg ptr in s[0:1]
|
||||
.amdhsa_user_sgpr_kernarg_segment_ptr 1
|
||||
.amdhsa_user_sgpr_count 2
|
||||
# gfx10+ / gfx11 specifics (RSRC1[29..31])
|
||||
.amdhsa_wavefront_size32 1
|
||||
.amdhsa_workgroup_processor_mode 1
|
||||
.amdhsa_memory_ordered 1
|
||||
.amdhsa_forward_progress 1
|
||||
# misc for gfx11
|
||||
.amdhsa_dx10_clamp 1
|
||||
.amdhsa_ieee_mode 1
|
||||
.amdhsa_uses_dynamic_stack 0
|
||||
.end_amdhsa_kernel
|
||||
|
||||
.amdgpu_metadata
|
||||
---
|
||||
amdhsa.kernels:
|
||||
- .args:
|
||||
- .address_space: generic
|
||||
.name: C
|
||||
.offset: 0
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: f16
|
||||
- .address_space: generic
|
||||
.name: A
|
||||
.offset: 8
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: f16
|
||||
- .address_space: generic
|
||||
.name: B
|
||||
.offset: 16
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: f16
|
||||
.group_segment_fixed_size: 30336
|
||||
.kernarg_segment_align: 8
|
||||
.kernarg_segment_size: 32
|
||||
.max_flat_workgroup_size: 128
|
||||
.name: gemm
|
||||
.private_segment_fixed_size: 0
|
||||
.sgpr_count: 70
|
||||
.sgpr_spill_count: 0
|
||||
.symbol: gemm.kd
|
||||
.vgpr_count: 256
|
||||
.vgpr_spill_count: 0
|
||||
.wavefront_size: 32
|
||||
amdhsa.version:
|
||||
- 1
|
||||
- 1
|
||||
...
|
||||
.end_amdgpu_metadata
|
||||
@@ -0,0 +1,30 @@
|
||||
import math, pathlib
|
||||
|
||||
from tinygrad import Device, dtypes
|
||||
from tinygrad.uop.ops import UOp, Ops, KernelInfo
|
||||
|
||||
from extra.gemm.amd_uop_matmul import test_matmul
|
||||
|
||||
N = 4096
|
||||
TN = 96
|
||||
THREADS_PER_WG = 128
|
||||
NUM_WG = math.ceil(N / TN) * math.ceil(N / TN)
|
||||
|
||||
dname:str = Device.DEFAULT
|
||||
template:str = (pathlib.Path(__file__).parent/"template.s").read_text()
|
||||
|
||||
def asm_kernel() -> UOp:
|
||||
lidx = UOp.special(THREADS_PER_WG, "lidx0")
|
||||
gidx = UOp.special(NUM_WG, "gidx0")
|
||||
|
||||
a = UOp.placeholder((N*N,), dtypes.half, slot=1)
|
||||
b = UOp.placeholder((N*N,), dtypes.half, slot=2)
|
||||
c = UOp.placeholder((N*N,), dtypes.half, slot=0)
|
||||
|
||||
src = template.replace("INSTRUCTIONS", (pathlib.Path(__file__).parent/"gemm.s").read_text())
|
||||
|
||||
sink = UOp.sink(a, b, c, lidx, gidx, arg=KernelInfo(name="gemm"))
|
||||
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)), arg=())
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_matmul(asm_kernel(), dtype=dtypes.half, N=N)
|
||||
@@ -1,65 +0,0 @@
|
||||
# Run assembly on the AMD runtime and check correctness
|
||||
# VIZ=2 to profile
|
||||
import pathlib
|
||||
from tinygrad import Tensor, Device, dtypes
|
||||
from tinygrad.engine.realize import ExecItem, CompiledRunner
|
||||
from tinygrad.renderer import ProgramSpec
|
||||
from tinygrad.uop.ops import track_rewrites, UOp
|
||||
from tinygrad.helpers import TracingKey
|
||||
|
||||
fp = pathlib.Path(__file__).parent/"gemm.s"
|
||||
|
||||
# ** generate inputs on CPU
|
||||
|
||||
N = 8192
|
||||
scale = 10.0
|
||||
|
||||
import torch
|
||||
torch.manual_seed(0)
|
||||
A = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
|
||||
B = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
|
||||
Bt = B.t().contiguous() # transpose B for the baseline gemm
|
||||
C_torch = A@Bt
|
||||
|
||||
# ** copy buffers to AMD
|
||||
|
||||
# input creation and validation run on the copy engine for simpler tracing
|
||||
|
||||
def from_torch(t:torch.Tensor) -> Tensor:
|
||||
return Tensor.from_blob(t.data_ptr(), t.shape, dtype=dtypes.bfloat16, device="cpu").to(Device.DEFAULT).realize()
|
||||
|
||||
C_tiny = Tensor.matmul(from_torch(A), from_torch(Bt), dtype=dtypes.float32).cast(dtypes.bfloat16)
|
||||
C_asm = Tensor.empty_like(C_tiny)
|
||||
C_asm.uop.buffer.allocate()
|
||||
|
||||
# ** run gemms
|
||||
|
||||
@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name,), ret=ret))
|
||||
def get_asm_gemm(ast:UOp, fp:pathlib.Path) -> ProgramSpec:
|
||||
src = fp.read_text()
|
||||
lib = Device[Device.DEFAULT].compiler.compile(src)
|
||||
return ProgramSpec("gemm", src, Device.DEFAULT, ast, lib=lib, global_size=[1024, 1, 1], local_size=[256, 1, 1], globals=[0, 1, 2])
|
||||
|
||||
sched = C_tiny.schedule()
|
||||
assert len(sched) == 1
|
||||
eis:list[ExecItem] = [sched[-1].lower()]
|
||||
ast = eis[0].ast
|
||||
prg = get_asm_gemm(ast, fp)
|
||||
eis.append(ExecItem(ast, [C_asm.uop.buffer, from_torch(B).uop.buffer, from_torch(A).uop.buffer], prg=CompiledRunner(prg)))
|
||||
|
||||
|
||||
for ei in eis:
|
||||
et = ei.run(wait=True)
|
||||
print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
|
||||
|
||||
# ** correctness
|
||||
|
||||
import ctypes
|
||||
|
||||
def torch_bf16(t:Tensor) -> torch.tensor:
|
||||
asm_out = t.to("cpu").realize().uop.buffer._buf
|
||||
buf = (ctypes.c_uint16*C_asm.uop.size).from_address(asm_out.va_addr)
|
||||
return torch.frombuffer(buf, dtype=torch.bfloat16, count=C_asm.uop.size).reshape(C_asm.shape)
|
||||
|
||||
assert torch.allclose(torch_bf16(C_asm), C_torch, rtol=1e-2, atol=1e-3)
|
||||
assert torch.allclose(torch_bf16(C_tiny), C_torch, rtol=1e-2, atol=1e-3)
|
||||
@@ -1,12 +1,12 @@
|
||||
# unpack the complete kernel descriptor of an amdgpu ELF of for gfx950
|
||||
# unpack the complete kernel descriptor of an amdgpu ELF
|
||||
# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
|
||||
import struct, pathlib
|
||||
import struct, pathlib, sys
|
||||
from tinygrad.runtime.support.elf import elf_loader
|
||||
|
||||
def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
|
||||
def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
|
||||
|
||||
with open(fp:=pathlib.Path(__file__).parent/"lib", "rb") as f:
|
||||
with open(sys.argv[1], "rb") as f:
|
||||
lib = f.read()
|
||||
|
||||
image, sections, relocs = elf_loader(lib)
|
||||
@@ -49,7 +49,7 @@ print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
|
||||
print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
|
||||
print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
|
||||
|
||||
# rsrc 3
|
||||
# rsrc 3 (gfx950)
|
||||
|
||||
accum_offset_raw = bits(pgm_rsrc3, 0, 5)
|
||||
assert_zero(pgm_rsrc3, 6, 15)
|
||||
@@ -169,10 +169,10 @@ assert_zero(desc, 458, 459)
|
||||
uses_dynamic_stack = bits(desc, 459, 460)
|
||||
print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
|
||||
|
||||
# gfx950 only
|
||||
assert_zero(desc, 460, 463)
|
||||
kernarg_preload_spec_length = bits(desc, 464, 470)
|
||||
print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
|
||||
|
||||
kernarg_preload_spec_offset = bits(desc, 471, 479)
|
||||
print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
|
||||
|
||||
|
||||
+47
-30
@@ -1,8 +1,42 @@
|
||||
import argparse, os, hashlib
|
||||
from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch
|
||||
import argparse, os, hashlib, functools
|
||||
from typing import Iterator, Callable
|
||||
from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch, ceildiv
|
||||
from extra.hevc.hevc import parse_hevc_file_headers, untile_nv12, to_bgr, nv_gpu
|
||||
from tinygrad import Tensor, dtypes, Device, Variable, TinyJit
|
||||
|
||||
# rounds up hevc input data to 32 bytes, so more optimal kernels can be generated
|
||||
HEVC_ROUNDUP = getenv("DATA_ROUNDUP", 32)
|
||||
|
||||
@functools.cache
|
||||
def _hevc_jitted_decoder(out_image_size:tuple[int, int], max_hist:int, inplace:bool):
|
||||
def hevc_decode_frame(pos:Variable, hevc_tensor:Tensor, offset:Variable, sz:Variable, opaque:Tensor, i:Variable, *hist:Tensor, outbuf:Tensor|None=None):
|
||||
x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist)
|
||||
if outbuf is not None: outbuf.assign(x).realize()
|
||||
return x.realize()
|
||||
return TinyJit(hevc_decode_frame)
|
||||
|
||||
def hevc_decode(hevc_tensor:Tensor, opaque:Tensor, frame_info:list, luma_h:int, luma_w:int,
|
||||
history:list[Tensor]|None=None, preallocated_outputs:list[Tensor]|None=None, warmup=False) -> Iterator[Tensor]:
|
||||
out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
|
||||
max_hist = max((hs for _, _, _, hs, _ in frame_info), default=0)
|
||||
|
||||
v_pos = Variable("pos", 0, max_hist + 1)
|
||||
v_offset = Variable("offset", 0, hevc_tensor.numel()-1)
|
||||
v_sz = Variable("sz", 1, ceildiv(hevc_tensor.numel(), HEVC_ROUNDUP))
|
||||
v_i = Variable("i", 0, len(frame_info)-1)
|
||||
|
||||
decode_jit = _hevc_jitted_decoder(out_image_size, max_hist, preallocated_outputs is not None)
|
||||
history = history or [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(max_hist)]
|
||||
assert len(history) == max_hist, f"history length {len(history)} does not match max_hist {max_hist}"
|
||||
|
||||
for i, (offset, sz, frame_pos, _, is_hist) in enumerate(frame_info):
|
||||
history = history[-max_hist:] if max_hist > 0 else []
|
||||
img = decode_jit(v_pos.bind(frame_pos), hevc_tensor, v_offset.bind(offset), v_sz.bind(ceildiv(sz, HEVC_ROUNDUP)),
|
||||
opaque, v_i.bind(i), *history, outbuf=preallocated_outputs[i] if preallocated_outputs else None)
|
||||
res = preallocated_outputs[i] if preallocated_outputs else img.clone().realize()
|
||||
if is_hist: history.append(res)
|
||||
yield res
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input_file", type=str, default="")
|
||||
@@ -19,49 +53,32 @@ if __name__ == "__main__":
|
||||
dat_hash = hashlib.md5(dat).hexdigest()
|
||||
|
||||
with Timing("prep infos: "):
|
||||
dat_nv = hevc_tensor.to("NV")
|
||||
opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat)
|
||||
|
||||
frame_info = frame_info[:getenv("MAX_FRAMES", len(frame_info))]
|
||||
|
||||
# move all needed data to gpu
|
||||
#all_slices = []
|
||||
with Timing("copy to gpu: "):
|
||||
opaque_nv = opaque.to("NV").contiguous().realize()
|
||||
hevc_tensor = hevc_tensor.to("NV")
|
||||
|
||||
out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
|
||||
max_hist = max(history_sz for _, _, _, history_sz, _ in frame_info)
|
||||
|
||||
# define variables
|
||||
v_pos = Variable("pos", 0, max_hist + 1)
|
||||
v_offset = Variable("offset", 0, hevc_tensor.numel()-1)
|
||||
v_sz = Variable("sz", 0, hevc_tensor.numel())
|
||||
v_i = Variable("i", 0, len(frame_info)-1)
|
||||
# preallocate output/hist buffers
|
||||
max_hist = max((hs for _, _, _, hs, _ in frame_info), default=0)
|
||||
hist = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(max_hist)]
|
||||
out_images = [Tensor.zeros(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(len(frame_info))]
|
||||
|
||||
@TinyJit
|
||||
def decode_jit(pos:Variable, src:Tensor, data:Tensor, *hist:Tensor):
|
||||
return src.decode_hevc_frame(pos, out_image_size, data, hist).realize()
|
||||
# warmup decode
|
||||
_ = list(hevc_decode(hevc_tensor, opaque_nv, frame_info[:3], luma_h, luma_w, history=hist, preallocated_outputs=out_images))
|
||||
Device.default.synchronize()
|
||||
|
||||
# warm up
|
||||
history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)]
|
||||
for i in range(3):
|
||||
hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(frame_info[0][0])), bound_offset+v_sz.bind(frame_info[0][1])),))
|
||||
decode_jit(v_pos.bind(0), hevc_frame, opaque_nv[v_i.bind(0)], *history)
|
||||
|
||||
out_images = []
|
||||
# decode all frames using the iterator
|
||||
with Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps")):
|
||||
for i, (offset, sz, frame_pos, history_sz, is_hist) in enumerate(frame_info):
|
||||
history = history[-max_hist:] if max_hist > 0 else []
|
||||
# TODO: this shrink should work as a slice
|
||||
hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(offset)), bound_offset+v_sz.bind(sz)),))
|
||||
|
||||
outimg = decode_jit(v_pos.bind(frame_pos), hevc_frame, opaque_nv[v_i.bind(i)], *history).clone()
|
||||
out_images.append(outimg)
|
||||
if is_hist: history.append(outimg)
|
||||
|
||||
images = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w, history=hist, preallocated_outputs=out_images))
|
||||
Device.default.synchronize()
|
||||
|
||||
# validation
|
||||
if getenv("VALIDATE", 0):
|
||||
import pickle
|
||||
if dat_hash == "b813bfdbec194fd17fdf0e3ceb8cea1c":
|
||||
@@ -70,7 +87,7 @@ if __name__ == "__main__":
|
||||
else: decoded_frames = pickle.load(open(f"extra/hevc/decoded_frames_{dat_hash}.pkl", "rb"))
|
||||
else: import cv2
|
||||
|
||||
for i, img in tqdm(enumerate(out_images)):
|
||||
for i, img in tqdm(enumerate(images)):
|
||||
if getenv("VALIDATE", 0):
|
||||
if i < len(decoded_frames) and len(decoded_frames[i]) > 0:
|
||||
img = untile_nv12(img, h, w, luma_w, chroma_off).realize()
|
||||
|
||||
@@ -4,7 +4,8 @@ import os, pathlib
|
||||
os.environ["AMD_AQL"] = "1"
|
||||
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.runtime.ops_amd import AMDProgram, HIPCompiler
|
||||
from tinygrad.runtime.support.compiler_amd import HIPCompiler
|
||||
from tinygrad.runtime.ops_amd import AMDProgram
|
||||
|
||||
NUM_WORKGROUPS = 96
|
||||
WAVE_SIZE = 32
|
||||
|
||||
@@ -129,7 +129,7 @@ class LSTM:
|
||||
return self.do_step(x_, hc_)
|
||||
|
||||
if hc is None:
|
||||
hc = Tensor.zeros(self.layers, 2 * x.shape[1], self.hidden_size, requires_grad=False)
|
||||
hc = Tensor.zeros(self.layers, 2 * x.shape[1], self.hidden_size, requires_grad=False).contiguous().realize()
|
||||
|
||||
output = None
|
||||
for t in range(x.shape[0]):
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
import unittest
|
||||
import numpy as np
|
||||
|
||||
from tinygrad.helpers import BEAM, Timing, CI, Context
|
||||
from tinygrad import Variable, Tensor
|
||||
from tinygrad.helpers import BEAM, Timing, CI, prod
|
||||
from tinygrad import Variable, Device, Tensor
|
||||
from tinygrad.nn import Conv2d
|
||||
from tinygrad.uop.ops import AxisType
|
||||
from tinygrad.codegen.opt import Opt, OptOps
|
||||
from tinygrad.codegen.opt.postrange import Scheduler
|
||||
from tinygrad.codegen.opt.search import get_kernel_actions
|
||||
|
||||
def rand(*shape):
|
||||
return Tensor(np.random.rand(*shape).astype(np.float32))
|
||||
@@ -75,5 +79,27 @@ class TestBeamSearch(unittest.TestCase):
|
||||
a = (a + a) * a
|
||||
a.realize()
|
||||
|
||||
@unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
|
||||
def test_tc_up(self):
|
||||
tc = Device[Device.DEFAULT].renderer.tensor_cores[0]
|
||||
size = max(tc.dims[0], tc.dims[1]) * 8
|
||||
a, b = Tensor.rand(size, size, dtype=tc.dtype_in), Tensor.rand(size, size, dtype=tc.dtype_in)
|
||||
ast = a.matmul(b, dtype=tc.dtype_out).schedule()[-1].ast
|
||||
s = Scheduler(ast, Device[Device.DEFAULT].renderer)
|
||||
s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1)))
|
||||
up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)])
|
||||
actions = get_kernel_actions(s, include_0=False, max_up=int(up))
|
||||
upcasted = [s for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]
|
||||
assert len(upcasted) > 0, f"expected upcast/unroll actions after TC with max_up={up}, but got none"
|
||||
|
||||
def test_max_up(self):
|
||||
a = Tensor.rand(16, 16)
|
||||
ast = a.schedule()[-1].ast
|
||||
s = Scheduler(ast, Device[Device.DEFAULT].renderer)
|
||||
for max_up in (2, 4):
|
||||
actions = get_kernel_actions(s, include_0=False, max_up=max_up)
|
||||
for up_opts in [s.applied_opts for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]:
|
||||
assert len([opt for opt in up_opts if opt.arg > max_up]) == 0 and len([op for op in up_opts if op.arg <= max_up]) > 0
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
+10
-16
@@ -3,21 +3,18 @@
|
||||
|
||||
import numpy as np
|
||||
import unittest
|
||||
import subprocess, struct, math, textwrap
|
||||
from tinygrad import Tensor, dtypes, Device, UOp
|
||||
from tinygrad.uop.ops import Ops
|
||||
import subprocess, struct, math, functools
|
||||
from tinygrad import Tensor, dtypes, Device
|
||||
from tinygrad.helpers import getenv
|
||||
from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
|
||||
from tinygrad.renderer import ProgramSpec
|
||||
from tinygrad.engine.realize import CompiledRunner
|
||||
|
||||
from extra.assembly.rdna3.autogen import *
|
||||
from extra.assembly.rdna3.asm import waitcnt
|
||||
from test.testextra.test_cfg_viz import template
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
from extra.assembly.amd.asm import waitcnt
|
||||
|
||||
from test.testextra.test_cfg_viz import asm_kernel
|
||||
|
||||
def get_output(asm:list, n_threads:int=1, vdst:VGPR=v[1]):
|
||||
out = Tensor([0]*n_threads, dtype=dtypes.uint32).realize()
|
||||
src = "\n".join(inst.disasm() for inst in [
|
||||
insts = [
|
||||
s_load_b64(s[0:1], s[0:1], NULL),
|
||||
*asm,
|
||||
v_lshlrev_b32_e32(v[0], 2, v[0]),
|
||||
@@ -25,12 +22,9 @@ def get_output(asm:list, n_threads:int=1, vdst:VGPR=v[1]):
|
||||
#global_store_b32(v[0], v[1], s[0:1]),
|
||||
global_store_b32(addr=v[0], data=vdst, saddr=s[0:1]),
|
||||
s_endpgm()
|
||||
])
|
||||
prg = ProgramSpec("test", template.replace("fn_name", "test").replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK),
|
||||
global_size=[1, 1, 1], local_size=[n_threads, 1, 1], globals=[0])
|
||||
car = CompiledRunner(prg)
|
||||
if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
|
||||
car([out.uop.buffer], {}, wait=True)
|
||||
]
|
||||
out = Tensor.custom_kernel(out, fxn=functools.partial(asm_kernel, name="test", insts=insts, device=out.device, n_threads=n_threads))[0]
|
||||
out.realize()
|
||||
return out.tolist()
|
||||
|
||||
def f16_to_bits(x:float) -> int: return struct.unpack('<H', struct.pack('<e', x))[0]
|
||||
|
||||
@@ -8,7 +8,7 @@ SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get p
|
||||
|
||||
`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 shdaer engine with instruction tracing enabled will not have instruction timings.
|
||||
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.
|
||||
|
||||
|
||||
Regular → Executable
+47
-18
@@ -1,11 +1,9 @@
|
||||
import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools, threading
|
||||
#!/usr/bin/env python3
|
||||
import ctypes, pathlib, argparse, pickle, dataclasses, threading
|
||||
from typing import Generator
|
||||
from tinygrad.helpers import temp, unwrap, DEBUG
|
||||
from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
|
||||
from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
|
||||
from tinygrad.runtime.autogen import llvm, rocprof
|
||||
from tinygrad.runtime.support.elf import elf_loader
|
||||
from tinygrad.viz.serve import llvm_disasm
|
||||
from tinygrad.runtime.ops_amd import ProfileSQTTEvent
|
||||
from tinygrad.runtime.autogen import rocprof
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class InstExec:
|
||||
@@ -115,28 +113,59 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, tuple[st
|
||||
|
||||
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
|
||||
|
||||
exc:Exception|None = None
|
||||
def worker():
|
||||
nonlocal exc
|
||||
try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
|
||||
except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
|
||||
except AttributeError as e:
|
||||
exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install")
|
||||
exc.__cause__ = e
|
||||
(t:=threading.Thread(target=worker, daemon=True)).start()
|
||||
t.join()
|
||||
if exc is not None:
|
||||
raise exc
|
||||
return ROCParseCtx
|
||||
|
||||
def print_pmc(events:list[ProfilePMCEvent]) -> None:
|
||||
from tinygrad.viz.serve import unpack_pmc
|
||||
def print_data(data:dict) -> None:
|
||||
from tabulate import tabulate
|
||||
for e in events:
|
||||
print("**", e.kern)
|
||||
data = unpack_pmc(e)
|
||||
print(tabulate([r[:-1] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
|
||||
# plaintext
|
||||
if "src" in data: print(data["src"])
|
||||
# table format
|
||||
elif "cols" in data:
|
||||
print(tabulate([r[:len(data["cols"])] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
|
||||
|
||||
def main() -> None:
|
||||
import tinygrad.viz.serve as viz
|
||||
viz.ctxs = []
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
|
||||
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)')
|
||||
parser.add_argument('-n', type=int, default=3, metavar="NUM", help='Max traces to print (optional number, default: 3 traces)')
|
||||
args = parser.parse_args()
|
||||
|
||||
with args.profile.open("rb") as f: profile = pickle.load(f)
|
||||
#rctx = decode(profile, disasm)
|
||||
#print('SQTT:', rctx.inst_execs.keys())
|
||||
|
||||
print_pmc([ev for ev in profile if isinstance(ev, ProfilePMCEvent)])
|
||||
viz.get_profile(profile)
|
||||
|
||||
# List all kernels
|
||||
if args.kernel is None:
|
||||
for c in viz.ctxs:
|
||||
print(c["name"])
|
||||
for s in c["steps"]: print(" "+s["name"])
|
||||
return None
|
||||
|
||||
# Find kernel trace
|
||||
trace = next((c for c in viz.ctxs if c["name"] == f"Exec {args.kernel}"), None)
|
||||
if not trace: raise RuntimeError(f"no matching trace for {args.kernel}")
|
||||
n = 0
|
||||
for s in trace["steps"]:
|
||||
print(s["name"])
|
||||
data = viz.get_render(s["query"])
|
||||
print_data(data)
|
||||
n += 1
|
||||
if n > args.n: break
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user