mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-20 16:26:07 +00:00
Compare commits
238
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a5afdcc79b | ||
|
|
9f7f2f0e0c | ||
|
|
50554115ee | ||
|
|
ab244c7f81 | ||
|
|
a879b54234 | ||
|
|
a8ae9757dd | ||
|
|
5aa71f437b | ||
|
|
eafcd44d95 | ||
|
|
3960e2758c | ||
|
|
9302ab003a | ||
|
|
551454f476 | ||
|
|
8a2549d42b | ||
|
|
7d1d9d4568 | ||
|
|
dc4ae7dd08 | ||
|
|
0a14e1fcd4 | ||
|
|
fc10470883 | ||
|
|
6790165ef8 | ||
|
|
e855ec8ee3 | ||
|
|
bbc55962ee | ||
|
|
3751b29a3d | ||
|
|
7c1f1cb2bc | ||
|
|
1696991988 | ||
|
|
a021b84604 | ||
|
|
fa5475307c | ||
|
|
5abc262e22 | ||
|
|
f9ca072b61 | ||
|
|
14e9a71a41 | ||
|
|
a0dd9d2146 | ||
|
|
32e1c267ee | ||
|
|
0cb024a5bb | ||
|
|
255e0573b1 | ||
|
|
164bc678a6 | ||
|
|
b46da603fe | ||
|
|
fd60626ea1 | ||
|
|
add7da268f | ||
|
|
e9ce12028e | ||
|
|
7e5687f6a3 | ||
|
|
1381daac06 | ||
|
|
8c55ef4f01 | ||
|
|
899a56446e | ||
|
|
986e865830 | ||
|
|
434dbafab5 | ||
|
|
76b577ee76 | ||
|
|
e5500ae4ad | ||
|
|
86708ccac5 | ||
|
|
f9147422a3 | ||
|
|
62c1a014a6 | ||
|
|
e0eea0d833 | ||
|
|
2a2c1eacf6 | ||
|
|
a92778aa0c | ||
|
|
2ab18ea7e3 | ||
|
|
002ea39da7 | ||
|
|
fe00682502 | ||
|
|
84b88a0a31 | ||
|
|
e610821c52 | ||
|
|
176a934ddd | ||
|
|
2a217ba206 | ||
|
|
79d00521f8 | ||
|
|
7fe91e5db9 | ||
|
|
1364449cab | ||
|
|
a28c8105a5 | ||
|
|
6cd318e377 | ||
|
|
fd10fd245a | ||
|
|
05fcb57696 | ||
|
|
7c967399a4 | ||
|
|
330a0b686e | ||
|
|
a8c821f45e | ||
|
|
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 |
@@ -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,35 @@ 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
|
||||
find tinygrad/runtime/autogen -type f -name "*.py" -not -name "__init__.py" -not -name "comgr_3.py" -not -name "metal.py" -not -name "libclang.py" -delete
|
||||
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
|
||||
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
|
||||
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
|
||||
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd"
|
||||
python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v14_0_2"
|
||||
python3 -c "from tinygrad.runtime.autogen import 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
|
||||
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
|
||||
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
|
||||
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
|
||||
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
|
||||
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
|
||||
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 +81,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 +117,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
|
||||
|
||||
+122
-226
@@ -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
|
||||
|
||||
@@ -389,10 +329,12 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
- name: Remove amdgpu
|
||||
run: sudo rmmod amdgpu || true
|
||||
- name: Cleanup running AM processes
|
||||
run: python extra/amdpci/am_smi.py --pids --kill
|
||||
- name: Setcap to python
|
||||
run: ./extra/amdpci/setup_python_cap.sh
|
||||
- name: Remove amd modules
|
||||
run: ./extra/hcq/hcq_smi.py amd rmmod
|
||||
- name: Kill stale pids
|
||||
run: ./extra/hcq/hcq_smi.py amd kill_pids
|
||||
#- name: Insert amdgpu
|
||||
# run: sudo modprobe amdgpu
|
||||
- name: Symlink models and datasets
|
||||
@@ -426,7 +368,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 +379,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 +394,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
|
||||
|
||||
@@ -521,10 +441,12 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
- name: Remove amdgpu
|
||||
run: sudo rmmod amdgpu || true
|
||||
- name: Cleanup running AM processes
|
||||
run: python extra/amdpci/am_smi.py --pids --kill
|
||||
- name: Setcap to python
|
||||
run: ./extra/amdpci/setup_python_cap.sh
|
||||
- name: Remove amd modules
|
||||
run: ./extra/hcq/hcq_smi.py amd rmmod
|
||||
- name: Kill stale pids
|
||||
run: ./extra/hcq/hcq_smi.py amd kill_pids
|
||||
- name: Symlink models and datasets
|
||||
run: |
|
||||
mkdir -p weights
|
||||
@@ -543,31 +465,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
|
||||
|
||||
@@ -582,10 +493,12 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
- name: Remove amdgpu
|
||||
run: sudo rmmod amdgpu || true
|
||||
- name: Cleanup running AM processes
|
||||
run: python extra/amdpci/am_smi.py --pids --kill
|
||||
- name: Setcap to python
|
||||
run: ./extra/amdpci/setup_python_cap.sh
|
||||
- name: Remove amd modules
|
||||
run: ./extra/hcq/hcq_smi.py amd rmmod
|
||||
- name: Kill stale pids
|
||||
run: ./extra/hcq/hcq_smi.py amd kill_pids
|
||||
- name: Symlink models and datasets
|
||||
run: |
|
||||
mkdir -p weights
|
||||
@@ -606,19 +519,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
|
||||
|
||||
@@ -641,17 +547,19 @@ jobs:
|
||||
- name: reset process replay
|
||||
run: test/external/process_replay/reset.py
|
||||
- name: openpilot compile3 0.10.0 driving_policy
|
||||
run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
|
||||
run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
|
||||
- name: openpilot compile3 0.10.0 dmonitoring
|
||||
run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
|
||||
- name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
|
||||
run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
|
||||
run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
|
||||
run: BENCHMARK_LOG=image_1_openpilot_0_10_1_vision PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: openpilot compile3 0.10.1 driving_vision
|
||||
run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
|
||||
- name: openpilot compile3 0.10.1 driving_policy
|
||||
run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
|
||||
run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
|
||||
- name: openpilot compile3 0.10.1 dmonitoring
|
||||
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
|
||||
- name: benchmark MobileNetV2 on DSP
|
||||
@@ -676,6 +584,8 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
- name: Setcap to python
|
||||
run: ./extra/amdpci/setup_python_cap.sh
|
||||
- name: Remove amd modules
|
||||
run: ./extra/hcq/hcq_smi.py amd rmmod
|
||||
- name: Kill stale pids
|
||||
@@ -708,7 +618,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 +628,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
|
||||
|
||||
@@ -746,6 +648,8 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
- name: Setcap to python
|
||||
run: ./extra/amdpci/setup_python_cap.sh
|
||||
- name: Remove nv modules
|
||||
run: ./extra/hcq/hcq_smi.py nv rmmod
|
||||
- name: Kill stale pids
|
||||
@@ -778,21 +682,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
|
||||
|
||||
+28
-34
@@ -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,16 +105,12 @@ 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
|
||||
run: python3 extra/torch_backend/test.py
|
||||
- name: custom tests
|
||||
run: python3 -m pytest -n auto extra/torch_backend/test.py --durations=20
|
||||
- name: Test one op in torch tests
|
||||
run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
|
||||
- name: Test Ops with TINY_BACKEND
|
||||
@@ -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
|
||||
|
||||
@@ -667,6 +667,11 @@ jobs:
|
||||
key: rdna3-emu
|
||||
deps: testing_minimal
|
||||
amd: 'true'
|
||||
python-version: '3.13'
|
||||
- name: Verify AMD autogen is up to date
|
||||
run: |
|
||||
python -m extra.assembly.amd.amdxml
|
||||
git diff --exit-code extra/assembly/amd/autogen/
|
||||
- name: Install LLVM 21
|
||||
run: |
|
||||
wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
|
||||
@@ -675,32 +680,19 @@ jobs:
|
||||
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: PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
|
||||
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: PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
|
||||
|
||||
testamdautogen:
|
||||
name: AMD autogen
|
||||
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-autogen
|
||||
pydeps: "pdfplumber"
|
||||
- name: Verify AMD autogen is up to date
|
||||
run: |
|
||||
python -m extra.assembly.amd.dsl --arch all
|
||||
python -m extra.assembly.amd.pcode --arch all
|
||||
git diff --exit-code extra/assembly/amd/autogen/
|
||||
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:
|
||||
@@ -789,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
|
||||
|
||||
@@ -192,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
|
||||
@@ -209,6 +212,14 @@ Key patterns to watch (from ResNet50 benchmark):
|
||||
|
||||
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.
|
||||
|
||||
+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.
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
|
||||
::: tinygrad.Tensor.neg
|
||||
::: tinygrad.Tensor.log
|
||||
::: tinygrad.Tensor.log2
|
||||
::: tinygrad.Tensor.log10
|
||||
::: tinygrad.Tensor.exp
|
||||
::: tinygrad.Tensor.exp2
|
||||
::: tinygrad.Tensor.sqrt
|
||||
@@ -87,4 +88,8 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
|
||||
::: tinygrad.Tensor.float
|
||||
::: tinygrad.Tensor.half
|
||||
::: tinygrad.Tensor.int
|
||||
::: tinygrad.Tensor.bool
|
||||
::: tinygrad.Tensor.bool
|
||||
::: tinygrad.Tensor.bfloat16
|
||||
::: tinygrad.Tensor.double
|
||||
::: tinygrad.Tensor.long
|
||||
::: tinygrad.Tensor.short
|
||||
@@ -27,5 +27,6 @@
|
||||
::: tinygrad.Tensor.flatten
|
||||
::: tinygrad.Tensor.unflatten
|
||||
::: tinygrad.Tensor.diag
|
||||
::: tinygrad.Tensor.diagonal
|
||||
::: tinygrad.Tensor.roll
|
||||
::: tinygrad.Tensor.rearrange
|
||||
@@ -7,6 +7,7 @@
|
||||
::: tinygrad.Tensor.any
|
||||
::: tinygrad.Tensor.all
|
||||
::: tinygrad.Tensor.isclose
|
||||
::: tinygrad.Tensor.allclose
|
||||
::: tinygrad.Tensor.mean
|
||||
::: tinygrad.Tensor.var
|
||||
::: tinygrad.Tensor.var_mean
|
||||
@@ -30,7 +31,9 @@
|
||||
::: tinygrad.Tensor.matmul
|
||||
::: tinygrad.Tensor.einsum
|
||||
::: tinygrad.Tensor.cumsum
|
||||
::: tinygrad.Tensor.cumprod
|
||||
::: tinygrad.Tensor.cummax
|
||||
::: tinygrad.Tensor.cummin
|
||||
::: tinygrad.Tensor.triu
|
||||
::: tinygrad.Tensor.tril
|
||||
::: tinygrad.Tensor.interpolate
|
||||
@@ -38,7 +41,9 @@
|
||||
::: tinygrad.Tensor.scatter_reduce
|
||||
::: tinygrad.Tensor.masked_select
|
||||
::: tinygrad.Tensor.masked_fill
|
||||
::: tinygrad.Tensor.nonzero
|
||||
::: tinygrad.Tensor.sort
|
||||
::: tinygrad.Tensor.argsort
|
||||
::: tinygrad.Tensor.topk
|
||||
::: tinygrad.Tensor.multinomial
|
||||
|
||||
@@ -56,3 +61,8 @@
|
||||
::: tinygrad.Tensor.sparse_categorical_crossentropy
|
||||
::: tinygrad.Tensor.cross_entropy
|
||||
::: tinygrad.Tensor.nll_loss
|
||||
|
||||
## Linear Algebra
|
||||
|
||||
::: tinygrad.Tensor.qr
|
||||
::: tinygrad.Tensor.svd
|
||||
|
||||
@@ -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
|
||||
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." AMD=1 DEBUG=0 JIT=1 FLASH_ATTENTION=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=0 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 WANDB=1 PARALLEL=0
|
||||
|
||||
RUNMLPERF=1 python3 examples/mlperf/model_train.py
|
||||
+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,345 @@
|
||||
# AMD ISA code generator - generates enum.py, ins.py, operands.py, str_pcode.py
|
||||
# Sources: XML from https://gpuopen.com/download/machine-readable-isa/latest/
|
||||
# PDF manuals from AMD documentation
|
||||
import re, zlib, xml.etree.ElementTree as ET, zipfile
|
||||
from tinygrad.helpers import fetch
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# Configuration
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
ARCHS = {
|
||||
"rdna3": {"xml": "amdgpu_isa_rdna3_5.xml", "pdf": "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content"},
|
||||
"rdna4": {"xml": "amdgpu_isa_rdna4.xml", "pdf": "https://docs.amd.com/api/khub/documents/uQpkEvk3pv~kfAb2x~j4uw/content"},
|
||||
"cdna": {"xml": "amdgpu_isa_cdna4.xml", "pdf": "https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-cdna4-instruction-set-architecture.pdf"},
|
||||
}
|
||||
XML_URL = "https://gpuopen.com/download/machine-readable-isa/latest/"
|
||||
# Map XML encoding names to codebase names (arch-specific overrides in ARCH_NAME_MAP)
|
||||
NAME_MAP = {"VOP3_SDST_ENC": "VOP3SD", "VOPDXY": "VOPD", "VDS": "DS"}
|
||||
ARCH_NAME_MAP = {"cdna": {"VOP3": "VOP3A", "VOP3_SDST_ENC": "VOP3B"}}
|
||||
# Instructions missing from XML but present in PDF
|
||||
FIXES = {"rdna3": {"SOPK": {22: "S_SUBVECTOR_LOOP_BEGIN", 23: "S_SUBVECTOR_LOOP_END"}, "FLAT": {55: "FLAT_ATOMIC_CSUB_U32"}},
|
||||
"rdna4": {"SOP1": {80: "S_GET_BARRIER_STATE", 81: "S_BARRIER_INIT", 82: "S_BARRIER_JOIN"}, "SOPP": {9: "S_WAITCNT", 21: "S_BARRIER_LEAVE"}}}
|
||||
# Encoding suffixes to strip (variants we don't generate separate classes for)
|
||||
_ENC_SUFFIXES = ("_INST_LITERAL", "_VOP_DPP16", "_VOP_DPP8", "_VOP_DPP", "_VOP_SDWA", "_NSA1", "_MFMA")
|
||||
# Field name normalization
|
||||
_FIELD_RENAMES = {"opsel_hi_2": "opsel_hi2", "op_sel_hi_2": "opsel_hi2", "op_sel": "opsel", "bound_ctrl": "bc",
|
||||
"tgt": "target", "row_en": "row", "unorm": "unrm", "clamp": "clmp", "wait_exp": "waitexp"}
|
||||
# Encoding variants to skip entirely
|
||||
_SKIP_ENCODINGS = ("LITERAL", "NSA", "DPP", "SDWA", "MFMA")
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# XML parsing helpers
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _strip_enc(name: str) -> str:
|
||||
"""Strip ENC_ prefix and encoding variant suffixes."""
|
||||
name = name.removeprefix("ENC_")
|
||||
for sfx in _ENC_SUFFIXES: name = name.replace(sfx, "")
|
||||
return name
|
||||
|
||||
def _norm_field(name: str) -> str:
|
||||
"""Normalize field name to match expected names."""
|
||||
for old, new in _FIELD_RENAMES.items(): name = name.replace(old, new)
|
||||
return name
|
||||
|
||||
def _map_flat(enc_name: str, instr_name: str) -> str:
|
||||
"""Map FLAT/GLOBAL/SCRATCH encoding to correct enum based on instruction prefix."""
|
||||
if enc_name in ("FLAT_GLBL", "FLAT_GLOBAL"): return "GLOBAL"
|
||||
if enc_name == "FLAT_SCRATCH": return "SCRATCH"
|
||||
if enc_name in ("FLAT", "VFLAT", "VGLOBAL", "VSCRATCH"):
|
||||
v = "V" if enc_name.startswith("V") else ""
|
||||
if instr_name.startswith("GLOBAL_"): return f"{v}GLOBAL"
|
||||
if instr_name.startswith("SCRATCH_"): return f"{v}SCRATCH"
|
||||
return f"{v}FLAT"
|
||||
return enc_name
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# XML parsing
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def parse_xml(filename: str, arch: str):
|
||||
root = ET.fromstring(zipfile.ZipFile(fetch(XML_URL)).read(filename))
|
||||
name_map = {**NAME_MAP, **ARCH_NAME_MAP.get(arch, {})}
|
||||
encodings, enums, types, fmts, op_types_set = {}, {}, {}, {}, set()
|
||||
# Extract DataFormats with BitCount
|
||||
for df in root.findall("ISA/DataFormats/DataFormat"):
|
||||
name, bits = df.findtext("DataFormatName"), df.findtext("BitCount")
|
||||
if name and bits: fmts[name] = int(bits)
|
||||
# Extract encoding definitions
|
||||
for enc in root.findall("ISA/Encodings/Encoding"):
|
||||
name = enc.findtext("EncodingName")
|
||||
if not name.startswith("ENC_") and name not in ("VOP3_SDST_ENC", "VOPDXY"): continue
|
||||
if any(s in name for s in _SKIP_ENCODINGS): continue
|
||||
fields = [(_norm_field(f.findtext("FieldName").lower()), int(f.find("BitLayout/Range").findtext("BitOffset") or 0) + int(f.find("BitLayout/Range").findtext("BitCount") or 0) - 1,
|
||||
int(f.find("BitLayout/Range").findtext("BitOffset") or 0))
|
||||
for f in enc.findall(".//MicrocodeFormat/BitMap/Field") if f.find("BitLayout/Range") is not None]
|
||||
ident = (enc.findall("EncodingIdentifiers/EncodingIdentifier") or [None])[0]
|
||||
enc_field = next((f for f in fields if f[0] == "encoding"), None)
|
||||
enc_bits = "".join(ident.text[len(ident.text)-1-b] for b in range(enc_field[1], enc_field[2]-1, -1)) if ident is not None and enc_field else None
|
||||
base_name = _strip_enc(name)
|
||||
encodings[name_map.get(base_name, base_name)] = (fields, enc_bits)
|
||||
# Extract instruction opcodes and operand info
|
||||
for instr in root.findall("ISA/Instructions/Instruction"):
|
||||
name = instr.findtext("InstructionName")
|
||||
for enc in instr.findall("InstructionEncodings/InstructionEncoding"):
|
||||
if enc.findtext("EncodingCondition") != "default": continue
|
||||
base, opcode = _map_flat(_strip_enc(enc.findtext("EncodingName")), name), int(enc.findtext("Opcode") or 0)
|
||||
enc_name = name_map.get(base, base)
|
||||
# ADDTID instructions go in both FLAT and GLOBAL enums (pcode uses FLATOp for these)
|
||||
if "ADDTID" in name:
|
||||
if base == "GLOBAL": enums.setdefault("FLAT", {})[opcode] = name
|
||||
elif base == "VGLOBAL": enums.setdefault("VFLAT", {})[opcode] = name
|
||||
enums.setdefault(enc_name, {})[opcode] = name
|
||||
# Extract operand info
|
||||
op_info = {op.findtext("FieldName").lower(): (op.findtext("DataFormatName"), int(op.findtext("OperandSize") or 0), op.findtext("OperandType"))
|
||||
for op in enc.findall("Operands/Operand") if op.findtext("FieldName")}
|
||||
for fmt, _, otype in op_info.values():
|
||||
if fmt and fmt not in fmts: fmts[fmt] = 0
|
||||
if otype: op_types_set.add(otype)
|
||||
if op_info: types[(name, enc_name)] = op_info
|
||||
return encodings, enums, types, fmts, op_types_set
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PDF parsing
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def extract_pdf_text(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
|
||||
|
||||
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_pcode(pages: list[list[tuple[float, float, str, str]]], name_to_op: dict[str, int]) -> dict[tuple[str, int], str]:
|
||||
"""Extract pseudocode for instructions. Returns {(name, opcode): pseudocode}."""
|
||||
# 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):
|
||||
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:
|
||||
sorted_lines = sorted(lines, key=lambda x: (x[0], -x[1]))
|
||||
# Stop at large Y gaps (>30) - indicates section break
|
||||
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
|
||||
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
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# Code generation
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def write_common(all_fmts, all_op_types, path):
|
||||
lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
|
||||
lines.append("class Fmt(Enum):")
|
||||
for fmt in sorted(all_fmts.keys()): lines.append(f" {fmt} = auto()")
|
||||
lines.append("")
|
||||
lines.append("FMT_BITS = {")
|
||||
for fmt, bits in sorted(all_fmts.items()): lines.append(f" Fmt.{fmt}: {bits},")
|
||||
lines.append("}")
|
||||
lines.append("")
|
||||
lines.append("class OpType(Enum):")
|
||||
for ot in sorted(all_op_types): lines.append(f" {ot} = auto()")
|
||||
with open(path, "w") as f: f.write("\n".join(lines))
|
||||
|
||||
def write_enum(enums, path):
|
||||
lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum", "from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401", ""]
|
||||
for name, ops in sorted(enums.items()):
|
||||
if not ops: continue
|
||||
suffix = "_E32" if name in ("VOP1", "VOP2", "VOPC") else "_E64" if name == "VOP3" else ""
|
||||
lines.append(f"class {name}Op(Enum):")
|
||||
aliases = []
|
||||
for op, mem in sorted(ops.items()):
|
||||
msuf = suffix if name != "VOP3" or op < 512 else ""
|
||||
lines.append(f" {mem}{msuf} = {op}")
|
||||
if msuf: aliases.append((mem, msuf))
|
||||
for mem, msuf in aliases: lines.append(f" {mem} = {mem}{msuf}")
|
||||
lines.append("")
|
||||
with open(path, "w") as f: f.write("\n".join(lines))
|
||||
|
||||
def write_ins(encodings, enums, arch, path):
|
||||
_VGPR_FIELDS = {"vdst", "vdstx", "vsrc0", "vsrc1", "vsrc2", "vsrc3", "vsrcx1", "vsrcy1", "vaddr", "vdata", "data", "data0", "data1", "addr"}
|
||||
def field_def(name, hi, lo, fmt, enc_bits=None):
|
||||
bits = hi - lo + 1
|
||||
if name == "encoding" and enc_bits: return f"FixedBitField({hi}, {lo}, 0b{enc_bits})"
|
||||
if name == "op" and fmt not in ("DPP", "SDWA"): return f"EnumBitField({hi}, {lo}, {fmt}Op)"
|
||||
if name in ("opx", "opy"): return f"EnumBitField({hi}, {lo}, VOPDOp)"
|
||||
if name == "vdsty": return f"VDSTYField({hi}, {lo})"
|
||||
if name in _VGPR_FIELDS and bits == 8: return f"VGPRField({hi}, {lo})"
|
||||
if name == "sbase" and bits == 6: return f"SBaseField({hi}, {lo})"
|
||||
if name in ("srsrc", "ssamp") and bits == 5: return f"SRsrcField({hi}, {lo})"
|
||||
if name in ("sdst", "sdata") and bits == 7: return f"SGPRField({hi}, {lo})"
|
||||
if name in ("soffset", "saddr") and bits == 7: return f"SGPRField({hi}, {lo}, default=NULL)"
|
||||
if name.startswith("ssrc") and bits == 8: return f"SSrcField({hi}, {lo})"
|
||||
if name in ("saddr", "soffset") and bits == 8: return f"SSrcField({hi}, {lo}, default=NULL)"
|
||||
if name.startswith("src") and bits == 9: return f"SrcField({hi}, {lo})"
|
||||
if fmt == "VOP3P" and name == "opsel_hi": return f"BitField({hi}, {lo}, default=3)"
|
||||
if fmt == "VOP3P" and name == "opsel_hi2": return f"BitField({hi}, {lo}, default=1)"
|
||||
return f"BitField({hi}, {lo})"
|
||||
ORDER = ['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']
|
||||
sort_fields = lambda fields: sorted(fields, key=lambda f: (ORDER.index(f[0]) if f[0] in ORDER else 999, f[2]))
|
||||
|
||||
lines = ["# autogenerated from AMD ISA XML - do not edit", "# ruff: noqa: F401,F403",
|
||||
"from extra.assembly.amd.dsl import *", f"from extra.assembly.amd.autogen.{arch}.enum import *", "import functools", ""]
|
||||
for enc_name, (fields, enc_bits) in sorted(encodings.items()):
|
||||
if enc_name in ("FLAT", "VFLAT"):
|
||||
prefix = "V" if enc_name == "VFLAT" else ""
|
||||
for cls, seg, op_enum in [(f"{prefix}FLAT", 0, f"{prefix}FLATOp"), (f"{prefix}GLOBAL", 2, f"{prefix}GLOBALOp"), (f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]:
|
||||
lines.append(f"class {cls}(Inst):")
|
||||
for fn, hi, lo in sort_fields(fields):
|
||||
if fn == "seg": lines.append(f" seg = FixedBitField({hi}, {lo}, {seg})")
|
||||
elif fn == "op": lines.append(f" op = EnumBitField({hi}, {lo}, {op_enum})")
|
||||
else: lines.append(f" {fn} = {field_def(fn, hi, lo, cls, enc_bits)}")
|
||||
lines.append("")
|
||||
elif enc_name not in ("FLAT_GLOBAL", "FLAT_SCRATCH", "FLAT_GLBL", "VGLOBAL", "VSCRATCH", "DPP", "SDWA"):
|
||||
lines.append(f"class {enc_name}(Inst):")
|
||||
for fn, hi, lo in sort_fields(fields):
|
||||
lines.append(f" {fn} = {field_def(fn, hi, lo, enc_name, enc_bits if fn == 'encoding' else None)}")
|
||||
lines.append("")
|
||||
# SDST variants
|
||||
for base, field in [("VOP1", "vdst = SSrcField(24, 17)"), ("VOP3", "vdst = SSrcField(7, 0)")]:
|
||||
if base in encodings: lines += [f"class {base}_SDST({base}):", f" {field}", ""]
|
||||
# Instruction helpers
|
||||
lines.append("# instruction helpers")
|
||||
SDST_OPS = {"V_READFIRSTLANE_B32", "V_READLANE_B32"}
|
||||
for fmt, ops in sorted(enums.items()):
|
||||
if fmt not in encodings and fmt not in ("GLOBAL", "SCRATCH", "VGLOBAL", "VSCRATCH"): continue
|
||||
suffix = "_E32" if fmt in ("VOP1", "VOP2", "VOPC") else "_E64" if fmt == "VOP3" else ""
|
||||
for op, name in sorted(ops.items()):
|
||||
msuf = suffix if fmt != "VOP3" or op < 512 else ""
|
||||
cls = "VOP1_SDST" if fmt == "VOP1" and name in SDST_OPS else "VOP3_SDST" if fmt == "VOP3" and (name in SDST_OPS or op < 256) else fmt
|
||||
lines.append(f"{name.lower()}{msuf.lower()} = functools.partial({cls}, {fmt}Op.{name}{msuf})")
|
||||
with open(path, "w") as f: f.write("\n".join(lines))
|
||||
|
||||
def write_operands(types, enums, arch, path):
|
||||
valid = {(name, fmt) for fmt, ops in enums.items() for name in ops.values()}
|
||||
lines = ["# autogenerated from AMD ISA XML - do not edit",
|
||||
"from extra.assembly.amd.autogen.common import Fmt, OpType",
|
||||
f"from extra.assembly.amd.autogen.{arch}.enum import *", ""]
|
||||
lines.append("# instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}")
|
||||
lines.append("OPERANDS = {")
|
||||
def fmt_val(v):
|
||||
fmt, size, otype = v
|
||||
return f"({f'Fmt.{fmt}' if fmt else 'None'}, {size}, {f'OpType.{otype}' if otype else 'None'})"
|
||||
for (name, enc_base), fields in sorted(types.items()):
|
||||
if (name, enc_base) not in valid: continue
|
||||
fstr = ", ".join(f'"{k}": {fmt_val(v)}' for k, v in sorted(fields.items()))
|
||||
lines.append(f' {enc_base}Op.{name}: {{{fstr}}},')
|
||||
lines.append("}")
|
||||
with open(path, "w") as f: f.write("\n".join(lines))
|
||||
|
||||
def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, str]], arch: str, path: str):
|
||||
"""Write str_pcode.py file from extracted pseudocode."""
|
||||
entries: list[tuple[str, str, int, str]] = []
|
||||
for fmt_name, ops in enums.items():
|
||||
member_suffix = "_E32" if fmt_name in ("VOP1", "VOP2", "VOPC") else "_E64" if fmt_name == "VOP3" else ""
|
||||
for opcode, name in ops.items():
|
||||
if (name, opcode) in pcode:
|
||||
msuf = member_suffix if fmt_name != "VOP3" or opcode < 512 else ""
|
||||
entries.append((f"{fmt_name}Op", f"{name}{msuf}", opcode, pcode[(name, opcode)]))
|
||||
enum_names = sorted(set(e[0] for e in entries))
|
||||
lines = ["# autogenerated from AMD ISA PDF - do not edit", "# ruff: noqa: E501",
|
||||
f"from extra.assembly.amd.autogen.{arch}.enum import {', '.join(enum_names)}", "", "PCODE = {"]
|
||||
for enum_name, name, opcode, code in sorted(entries, key=lambda x: (x[0], x[2])):
|
||||
lines.append(f" {enum_name}.{name}: {code!r},")
|
||||
lines.append("}")
|
||||
with open(path, "w") as f: f.write("\n".join(lines))
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# Main
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
if __name__ == "__main__":
|
||||
import pathlib
|
||||
all_fmts, all_op_types, arch_data = {}, set(), {}
|
||||
# First pass: parse XML for all architectures
|
||||
for arch, cfg in ARCHS.items():
|
||||
print(f"Parsing XML: {cfg['xml']} -> {arch}")
|
||||
encodings, enums, types, fmts, op_types_set = parse_xml(cfg["xml"], arch)
|
||||
for fmt, ops in FIXES.get(arch, {}).items(): enums.setdefault(fmt, {}).update(ops)
|
||||
arch_data[arch] = {"encodings": encodings, "enums": enums, "types": types}
|
||||
for fmt, bits in fmts.items():
|
||||
assert fmt not in all_fmts or all_fmts[fmt] == bits, f"FMT_BITS mismatch for {fmt}: {all_fmts[fmt]} vs {bits}"
|
||||
all_fmts[fmt] = bits
|
||||
all_op_types.update(op_types_set)
|
||||
# Write common.py
|
||||
common_path = pathlib.Path(__file__).parent / "autogen" / "common.py"
|
||||
write_common(all_fmts, all_op_types, common_path)
|
||||
print(f"Wrote common.py: {len(all_fmts)} formats, {len(all_op_types)} op types")
|
||||
# Write per-arch files from XML
|
||||
for arch, data in arch_data.items():
|
||||
base = pathlib.Path(__file__).parent / "autogen" / arch
|
||||
write_enum(data["enums"], base / "enum.py")
|
||||
write_ins(data["encodings"], data["enums"], arch, base / "ins.py")
|
||||
write_operands(data["types"], data["enums"], arch, base / "operands.py")
|
||||
print(f" {arch}: {len(data['encodings'])} encodings, {sum(len(v) for v in data['enums'].values())} instructions")
|
||||
# Second pass: parse PDFs and write pcode
|
||||
for arch, cfg in ARCHS.items():
|
||||
print(f"Parsing PDF: {arch}...")
|
||||
pages = extract_pdf_text(cfg["pdf"])
|
||||
name_to_op = {name: op for ops in arch_data[arch]["enums"].values() for op, name in ops.items()}
|
||||
pcode = extract_pcode(pages, name_to_op)
|
||||
base = pathlib.Path(__file__).parent / "autogen" / arch
|
||||
write_pcode(pcode, arch_data[arch]["enums"], arch, base / "str_pcode.py")
|
||||
print(f" {arch}: {len(pcode)} pcode entries")
|
||||
@@ -1,667 +0,0 @@
|
||||
# RDNA3 assembler and disassembler
|
||||
from __future__ import annotations
|
||||
import re
|
||||
from extra.assembly.amd.dsl import Inst, RawImm, Reg, SrcMod, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory, FLOAT_ENC, SRC_FIELDS, unwrap
|
||||
from extra.assembly.amd.dsl import VCC_LO, VCC_HI, VCC, EXEC_LO, EXEC_HI, EXEC, SCC, M0, NULL, OFF
|
||||
from extra.assembly.amd.autogen.rdna3 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.rdna3 import VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, VOPDOp, VINTERPOp
|
||||
from extra.assembly.amd.autogen.rdna3 import SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, DSOp, FLATOp, MUBUFOp, MTBUFOp, MIMGOp
|
||||
|
||||
# VOP3SD opcodes that share VOP3 encoding
|
||||
VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
|
||||
|
||||
def detect_format(data: bytes) -> 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')
|
||||
hi2 = (word >> 30) & 0x3
|
||||
if hi2 == 0b11:
|
||||
enc = (word >> 26) & 0xf
|
||||
if enc == 0b0010: return VOPD
|
||||
if enc == 0b0011: return VOP3P
|
||||
if enc == 0b0100: return VINTERP
|
||||
if enc == 0b0101: return VOP3SD if ((word >> 16) & 0x3ff) in VOP3SD_OPS else VOP3
|
||||
if enc == 0b0110: return DS
|
||||
if enc == 0b0111: return FLAT
|
||||
if enc == 0b1000: return MUBUF
|
||||
if enc == 0b1010: return MTBUF
|
||||
if enc == 0b1100 or enc == 0b1111: return MIMG
|
||||
if enc == 0b1101: return SMEM
|
||||
if enc == 0b1110: return EXP
|
||||
raise ValueError(f"unknown 64-bit format enc={enc:#06b} word={word:#010x}")
|
||||
if hi2 == 0b10:
|
||||
enc = (word >> 23) & 0x7f
|
||||
if enc == 0b1111101: return SOP1
|
||||
if enc == 0b1111110: return SOPC
|
||||
if enc == 0b1111111: return SOPP
|
||||
return SOPK if ((word >> 28) & 0xf) == 0b1011 else SOP2
|
||||
# hi2 == 0b00 or 0b01: VOP1/VOP2/VOPC (bit 31 = 0)
|
||||
assert (word >> 31) == 0, f"expected bit 31 = 0 for VOP, got word={word:#010x}"
|
||||
enc = (word >> 25) & 0x7f
|
||||
if enc == 0b0111110: return VOPC
|
||||
if enc == 0b0111111: return VOP1
|
||||
if enc <= 0b0111101: return VOP2
|
||||
raise ValueError(f"unknown VOP format enc={enc:#09b} word={word:#010x}")
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# CONSTANTS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
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"}
|
||||
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_IDS = {v.lower(): k for k, v in HWREG.items()}
|
||||
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'}
|
||||
VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# HELPERS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
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(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 _hl(v: int, hi_thresh: int = 128) -> str: return 'h' if v >= hi_thresh else 'l'
|
||||
|
||||
def _fmt_sdst(v: int, n: int = 1) -> str:
|
||||
if v == 124: return "null"
|
||||
if 108 <= v <= 123: return _reg("ttmp", v - 108, n)
|
||||
if n > 1: return SPECIAL_PAIRS.get(v) or _sreg(v, n)
|
||||
return {126: "exec_lo", 127: "exec_hi", 106: "vcc_lo", 107: "vcc_hi", 125: "m0"}.get(v, f"s{v}")
|
||||
|
||||
def _fmt_src(v: int, n: int = 1) -> str:
|
||||
if n == 1: return decode_src(v)
|
||||
if v >= 256: return _vreg(v - 256, n)
|
||||
if v <= 105: return _sreg(v, n)
|
||||
if n == 2 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
|
||||
if 108 <= v <= 123: return _reg("ttmp", v - 108, n)
|
||||
return decode_src(v)
|
||||
|
||||
def _fmt_v16(v: int, base: int = 256, hi_thresh: int = 384) -> str:
|
||||
return f"v{(v - base) & 0x7f}.{_hl(v, hi_thresh)}"
|
||||
|
||||
def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
|
||||
return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
|
||||
|
||||
def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
|
||||
def _is16(op: str) -> bool: return _has(op, 'f16', 'i16', 'u16', 'b16') and not _has(op, '_f32', '_i32')
|
||||
def _is64(op: str) -> bool: return _has(op, 'f64', 'i64', 'u64', 'b64')
|
||||
def _omod(v: int) -> str: return {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(v, "")
|
||||
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, any_hi: bool) -> str:
|
||||
"""Format VOP3 source operand with modifiers."""
|
||||
if 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" if any_hi else inst.lit(v))
|
||||
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 ""
|
||||
if is16_d and (opsel & 8): return f" op_sel:[1,1,1{',1' if n == 3 else ''}]"
|
||||
if n == 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}]"
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# DISASSEMBLER
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _disasm_vop1(inst: VOP1) -> str:
|
||||
op = VOP1Op(inst.op)
|
||||
if op in (VOP1Op.V_NOP, VOP1Op.V_PIPEFLUSH): return op.name.lower()
|
||||
F64_OPS = {VOP1Op.V_CEIL_F64, VOP1Op.V_FLOOR_F64, VOP1Op.V_FRACT_F64, VOP1Op.V_FREXP_MANT_F64, VOP1Op.V_RCP_F64, VOP1Op.V_RNDNE_F64, VOP1Op.V_RSQ_F64, VOP1Op.V_SQRT_F64, VOP1Op.V_TRUNC_F64}
|
||||
is_f64_d = op in F64_OPS or op in (VOP1Op.V_CVT_F64_F32, VOP1Op.V_CVT_F64_I32, VOP1Op.V_CVT_F64_U32)
|
||||
is_f64_s = op in F64_OPS or op in (VOP1Op.V_CVT_F32_F64, VOP1Op.V_CVT_I32_F64, VOP1Op.V_CVT_U32_F64, VOP1Op.V_FREXP_EXP_I32_F64)
|
||||
name = op.name.lower()
|
||||
parts = name.split('_')
|
||||
is_16d = 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)
|
||||
is_16s = parts[-1] in ('f16','i16','u16','b16') and 'sat_pk' not in name
|
||||
if op == VOP1Op.V_READFIRSTLANE_B32: return f"v_readfirstlane_b32 {decode_src(inst.vdst)}, v{inst.src0 - 256 if inst.src0 >= 256 else inst.src0}"
|
||||
dst = _vreg(inst.vdst, 2) if is_f64_d else _fmt_v16(inst.vdst, 0, 128) if is_16d else f"v{inst.vdst}"
|
||||
src = _fmt_src(inst.src0, 2) if is_f64_s else _fmt_v16(inst.src0) if is_16s and inst.src0 >= 256 else inst.lit(inst.src0)
|
||||
return f"{name}_e32 {dst}, {src}"
|
||||
|
||||
def _disasm_vop2(inst: VOP2) -> str:
|
||||
op = VOP2Op(inst.op)
|
||||
name = op.name.lower()
|
||||
suf = "" if op == VOP2Op.V_DOT2ACC_F32_F16 else "_e32"
|
||||
is16 = _is16(name) and 'pk_' not in name
|
||||
# fmaak: dst = src0 * vsrc1 + K, fmamk: dst = src0 * K + vsrc1
|
||||
if op in (VOP2Op.V_FMAAK_F32, VOP2Op.V_FMAAK_F16): return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}, 0x{inst._literal:x}"
|
||||
if op in (VOP2Op.V_FMAMK_F32, VOP2Op.V_FMAMK_F16): return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, 0x{inst._literal:x}, v{inst.vsrc1}"
|
||||
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_fmt_v16(inst.src0) if inst.src0 >= 256 else inst.lit(inst.src0)}, {_fmt_v16(inst.vsrc1, 0, 128)}"
|
||||
return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}" + (", vcc_lo" if op == VOP2Op.V_CNDMASK_B32 else "")
|
||||
|
||||
VOPC_CLASS = {VOPCOp.V_CMP_CLASS_F16, VOPCOp.V_CMP_CLASS_F32, VOPCOp.V_CMP_CLASS_F64,
|
||||
VOPCOp.V_CMPX_CLASS_F16, VOPCOp.V_CMPX_CLASS_F32, VOPCOp.V_CMPX_CLASS_F64}
|
||||
|
||||
def _disasm_vopc(inst: VOPC) -> str:
|
||||
op = VOPCOp(inst.op)
|
||||
name = op.name.lower()
|
||||
is64, is16 = _is64(name), _is16(name)
|
||||
s0 = _fmt_src(inst.src0, 2) if is64 else _fmt_v16(inst.src0) if is16 and inst.src0 >= 256 else inst.lit(inst.src0)
|
||||
s1 = _vreg(inst.vsrc1, 2) if is64 and op not in VOPC_CLASS else _fmt_v16(inst.vsrc1, 0, 128) if is16 else f"v{inst.vsrc1}"
|
||||
return f"{name}_e32 {s0}, {s1}" if op.value >= 128 else f"{name}_e32 vcc_lo, {s0}, {s1}"
|
||||
|
||||
NO_ARG_SOPP = {SOPPOp.S_ENDPGM, 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}
|
||||
|
||||
def _disasm_sopp(inst: SOPP) -> str:
|
||||
op, name = SOPPOp(inst.op), SOPPOp(inst.op).name.lower()
|
||||
if op in NO_ARG_SOPP: return name
|
||||
if 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 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:
|
||||
op = SMEMOp(inst.op)
|
||||
name = op.name.lower()
|
||||
if op in (SMEMOp.S_GL1_INV, SMEMOp.S_DCACHE_INV): return name
|
||||
off_s = f"{decode_src(inst.soffset)} offset:0x{inst.offset:x}" if inst.offset and inst.soffset != 124 else f"0x{inst.offset:x}" if inst.offset else decode_src(inst.soffset)
|
||||
sbase_idx, sbase_count = inst.sbase * 2, 4 if (8 <= inst.op <= 12 or name == 's_atc_probe_buffer') else 2
|
||||
sbase_str = _fmt_src(sbase_idx, sbase_count) 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}"
|
||||
width = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(inst.op, 1)
|
||||
return f"{name} {_fmt_sdst(inst.sdata, width)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (inst.dlc, " dlc"))
|
||||
|
||||
def _disasm_flat(inst: FLAT) -> str:
|
||||
name = FLATOp(inst.op).name.lower()
|
||||
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)
|
||||
suffix = name.split('_')[-1]
|
||||
w = {'b32':1,'b64':2,'b96':3,'b128':4,'u8':1,'i8':1,'u16':1,'i16':1,'u32':1,'i32':1,'u64':2,'i64':2,'f32':1,'f64':2}.get(suffix, 1)
|
||||
if 'cmpswap' in name: w *= 2
|
||||
if name.endswith('_x2') or 'x2' in suffix: w = max(w, 2)
|
||||
mods = f"{f' offset:{off_val}' if off_val else ''}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
|
||||
# saddr
|
||||
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)}"
|
||||
elif inst.saddr in SPECIAL_PAIRS: saddr_s = f", {SPECIAL_PAIRS[inst.saddr]}"
|
||||
elif 108 <= inst.saddr <= 123: saddr_s = f", {_reg('ttmp', inst.saddr - 108, 2)}"
|
||||
else: saddr_s = f", {_sreg(inst.saddr, 2) if inst.saddr < 106 else decode_src(inst.saddr)}"
|
||||
# addtid: no addr
|
||||
if 'addtid' in name: return f"{instr} v{inst.data if 'store' in name else inst.vdst}{saddr_s}{mods}"
|
||||
# addr width
|
||||
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, 1 if seg == 'scratch' or (inst.saddr not in (0x7F, 124)) else 2)
|
||||
data_s, vdst_s = _vreg(inst.data, w), _vreg(inst.vdst, w // 2 if 'cmpswap' in name else w)
|
||||
if 'atomic' in name:
|
||||
return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if inst.glc 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} {_vreg(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
|
||||
|
||||
def _disasm_ds(inst: DS) -> str:
|
||||
op, name = DSOp(inst.op), DSOp(inst.op).name.lower()
|
||||
gds = " gds" if inst.gds else ""
|
||||
off = f" offset:{inst.offset0 | (inst.offset1 << 8)}" if inst.offset0 or inst.offset1 else ""
|
||||
off2 = f" offset0:{inst.offset0} offset1:{inst.offset1}" if inst.offset0 or inst.offset1 else ""
|
||||
w = 4 if '128' in name else 3 if '96' in name else 2 if (name.endswith('64') or 'gs_reg' in name) else 1
|
||||
d0, d1, dst, addr = _vreg(inst.data0, w), _vreg(inst.data1, w), _vreg(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 '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} v{inst.vdst}{off}{gds}"
|
||||
if 'gs_reg' in name: return f"{name} {_vreg(inst.vdst, 2)}, v{inst.data0}{off}{gds}"
|
||||
if '2addr' in name:
|
||||
if 'load' in name: return f"{name} {_vreg(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} {_vreg(inst.vdst, w*2)}, {addr}, {d0}, {d1}{off2}{gds}"
|
||||
if 'load' in name: return f"{name} v{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} v{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} v{inst.vdst}, {addr}{off}{gds}"
|
||||
if 'permute' in name: return f"{name} v{inst.vdst}, {addr}, v{inst.data0}{off}{gds}"
|
||||
if 'condxchg' in name: return f"{name} {_vreg(inst.vdst, 2)}, {addr}, {_vreg(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 = VOP3SDOp(inst.op) if inst.op in VOP3SD_OPS else VOP3Op(inst.op)
|
||||
name = op.name.lower()
|
||||
|
||||
# VOP3SD (shared encoding)
|
||||
if inst.op in VOP3SD_OPS:
|
||||
sdst = (inst.clmp << 7) | (inst.opsel << 3) | inst.abs
|
||||
is64, mad64 = 'f64' in name, _has(name, 'mad_i64_i32', 'mad_u64_u32')
|
||||
def src(v, neg, ext=False): s = _fmt_src(v, 2) if ext or is64 else inst.lit(v); return f"-{s}" if neg else s
|
||||
s0, s1, s2 = src(inst.src0, inst.neg & 1), src(inst.src1, inst.neg & 2), src(inst.src2, inst.neg & 4, mad64)
|
||||
dst = _vreg(inst.vdst, 2) if is64 or mad64 else f"v{inst.vdst}"
|
||||
if op in (VOP3SDOp.V_ADD_CO_U32, VOP3SDOp.V_SUB_CO_U32, VOP3SDOp.V_SUBREV_CO_U32): return f"{name} {dst}, {_fmt_sdst(sdst, 1)}, {s0}, {s1}"
|
||||
if op in (VOP3SDOp.V_ADD_CO_CI_U32, VOP3SDOp.V_SUB_CO_CI_U32, VOP3SDOp.V_SUBREV_CO_CI_U32): return f"{name} {dst}, {_fmt_sdst(sdst, 1)}, {s0}, {s1}, {s2}"
|
||||
return f"{name} {dst}, {_fmt_sdst(sdst, 1)}, {s0}, {s1}, {s2}" + _omod(inst.omod)
|
||||
|
||||
# Detect operand sizes
|
||||
is64 = _is64(name)
|
||||
is64_src, is64_dst = False, False
|
||||
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')
|
||||
is64_src, is64_dst = '64' in m.group(2), '64' in m.group(1)
|
||||
is16_s2, is64 = is16_s, False
|
||||
elif re.match(r'v_mad_[iu]32_[iu]16', name): is16_s = True
|
||||
elif 'pack_b32' in name: is16_s = is16_s2 = True
|
||||
else: is16_d = is16_s = is16_s2 = _is16(name) and not _has(name, 'dot2', 'pk_', 'sad', 'msad', 'qsad', 'mqsad')
|
||||
|
||||
# Source counts
|
||||
shift64 = 'rev' in name and '64' in name and name.startswith('v_')
|
||||
ldexp64 = op == VOP3Op.V_LDEXP_F64
|
||||
trig = op == VOP3Op.V_TRIG_PREOP_F64
|
||||
sad64, mqsad = _has(name, 'qsad_pk', 'mqsad_pk'), 'mqsad_u32' in name
|
||||
s0n = 2 if ((is64 and not shift64) or sad64 or mqsad or is64_src) else 1
|
||||
s1n = 2 if (is64 and not _has(name, 'class') and not ldexp64 and not trig) else 1
|
||||
s2n = 4 if mqsad else 2 if (is64 or sad64) else 1
|
||||
|
||||
any_hi = inst.opsel != 0
|
||||
s0 = _vop3_src(inst, inst.src0, inst.neg&1, inst.abs&1, inst.opsel&1, s0n, is16_s, any_hi)
|
||||
s1 = _vop3_src(inst, inst.src1, inst.neg&2, inst.abs&2, inst.opsel&2, s1n, is16_s, any_hi)
|
||||
s2 = _vop3_src(inst, inst.src2, inst.neg&4, inst.abs&4, inst.opsel&4, s2n, is16_s2, any_hi)
|
||||
|
||||
# Destination
|
||||
dn = 4 if mqsad else 2 if (is64 or sad64 or is64_dst) else 1
|
||||
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" if any_hi else f"v{inst.vdst}"
|
||||
else: dst = f"v{inst.vdst}"
|
||||
|
||||
cl, om = " clamp" if inst.clmp 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}" if name.startswith('v_cmpx') else f"{name}_e64 {_fmt_sdst(inst.vdst, 1)}, {s0}, {s1}"
|
||||
if inst.op < 384: # VOP2
|
||||
os = _opsel_str(inst.opsel, 3, need_opsel, is16_d) if 'cndmask' in name else _opsel_str(inst.opsel, 2, need_opsel, is16_d)
|
||||
return f"{name}_e64 {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if 'cndmask' in name else f"{name}_e64 {dst}, {s0}, {s1}{os}{cl}{om}"
|
||||
if inst.op < 512: # VOP1
|
||||
return f"{name}_e64" if op in (VOP3Op.V_NOP, VOP3Op.V_PIPEFLUSH) else f"{name}_e64 {dst}, {s0}{_opsel_str(inst.opsel, 1, need_opsel, is16_d)}{cl}{om}"
|
||||
# Native VOP3
|
||||
is3 = _has(name, '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')
|
||||
os = _opsel_str(inst.opsel, 3 if is3 else 2, need_opsel, is16_d)
|
||||
return f"{name} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if is3 else f"{name} {dst}, {s0}, {s1}{os}{cl}{om}"
|
||||
|
||||
def _disasm_vop3sd(inst: VOP3SD) -> str:
|
||||
op, name = VOP3SDOp(inst.op), VOP3SDOp(inst.op).name.lower()
|
||||
is64, mad64 = 'f64' in name, _has(name, 'mad_i64_i32', 'mad_u64_u32')
|
||||
def src(v, neg, ext=False): s = _fmt_src(v, 2) if ext or is64 else inst.lit(v); return f"-{s}" if neg else s
|
||||
s0, s1, s2 = src(inst.src0, inst.neg & 1), src(inst.src1, inst.neg & 2), src(inst.src2, inst.neg & 4, mad64)
|
||||
dst, is2src = _vreg(inst.vdst, 2) if is64 or mad64 else f"v{inst.vdst}", op in (VOP3SDOp.V_ADD_CO_U32, VOP3SDOp.V_SUB_CO_U32, VOP3SDOp.V_SUBREV_CO_U32)
|
||||
suffix = "_e64" if name.startswith('v_') and 'co_' in name else ""
|
||||
return f"{name}{suffix} {dst}, {_fmt_sdst(inst.sdst, 1)}, {s0}, {s1}{'' if is2src else f', {s2}'}{' clamp' if inst.clmp else ''}{_omod(inst.omod)}"
|
||||
|
||||
def _disasm_vopd(inst: VOPD) -> str:
|
||||
lit = inst._literal or inst.literal
|
||||
vdst_y, nx, ny = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), VOPDOp(inst.opx).name.lower(), VOPDOp(inst.opy).name.lower()
|
||||
def half(n, vd, s0, vs1): return f"{n} v{vd}, {inst.lit(s0)}{f', 0x{lit:x}' if lit and _has(n, 'fmaak', 'fmamk') else ''}" if 'mov' in n else f"{n} v{vd}, {inst.lit(s0)}, v{vs1}{f', 0x{lit:x}' if lit and _has(n, 'fmaak', 'fmamk') else ''}"
|
||||
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
|
||||
|
||||
def _disasm_vop3p(inst: VOP3P) -> str:
|
||||
name = VOP3POp(inst.op).name.lower()
|
||||
is_wmma, is_3src, is_fma_mix = 'wmma' in name, _has(name, 'fma', 'mad', 'dot', 'wmma'), 'fma_mix' in name
|
||||
if is_wmma:
|
||||
sc = 2 if 'iu4' in name else 4 if 'iu8' in name else 8
|
||||
src0, src1, src2, dst = _fmt_src(inst.src0, sc), _fmt_src(inst.src1, sc), _fmt_src(inst.src2, 8), _vreg(inst.vdst, 8)
|
||||
else: src0, src1, src2, dst = _fmt_src(inst.src0, 1), _fmt_src(inst.src1, 1), _fmt_src(inst.src2, 1), f"v{inst.vdst}"
|
||||
n, opsel_hi = 3 if is_3src else 2, inst.opsel_hi | (inst.opsel_hi2 << 2)
|
||||
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 inst.clmp else [])
|
||||
else:
|
||||
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 is_3src 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 is_3src else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
|
||||
|
||||
def _disasm_buf(inst: MUBUF | MTBUF) -> str:
|
||||
op = MTBUFOp(inst.op) if isinstance(inst, MTBUF) else MUBUFOp(inst.op)
|
||||
name = op.name.lower()
|
||||
if 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 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 = _reg("ttmp", inst.srsrc*4 - 108, 4) if 108 <= inst.srsrc*4 <= 123 else _sreg(inst.srsrc*4, 4)
|
||||
mods = ([f"format:{inst.format}"] if isinstance(inst, MTBUF) else []) + [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]
|
||||
return f"{name} {_vreg(inst.vdata, w)}, {vaddr}, {srsrc}, {decode_src(inst.soffset)}{' ' + ' '.join(mods) if mods else ''}"
|
||||
|
||||
def _mimg_vaddr_width(name: str, dim: int, a16: bool) -> int:
|
||||
"""Calculate vaddr register count for MIMG sample/gather operations."""
|
||||
# 1d,2d,3d,cube,1d_arr,2d_arr,2d_msaa,2d_msaa_arr
|
||||
base = [1, 2, 3, 3, 2, 3, 3, 4][dim] # address coords
|
||||
grad = [1, 2, 3, 2, 1, 2, 2, 2][dim] # gradient coords (for derivatives)
|
||||
if 'get_resinfo' in name: return 1 # only mip level
|
||||
packed, unpacked = 0, 0
|
||||
if '_mip' in name: packed += 1
|
||||
elif 'sample' in name or 'gather' in name:
|
||||
if '_o' in name: unpacked += 1 # offset
|
||||
if re.search(r'_c(_|$)', name): unpacked += 1 # compare (not _cl)
|
||||
if '_d' in name: unpacked += (grad + 1) & ~1 if '_g16' in name else grad*2 # derivatives
|
||||
if '_b' in name: unpacked += 1 # bias
|
||||
if '_l' in name and '_cl' not in name and '_lz' not in name: packed += 1 # LOD
|
||||
if '_cl' in name: packed += 1 # clamp
|
||||
return (base + packed + 1) // 2 + unpacked if a16 else base + packed + unpacked
|
||||
|
||||
def _disasm_mimg(inst: MIMG) -> str:
|
||||
name = MIMGOp(inst.op).name.lower()
|
||||
srsrc_base = inst.srsrc * 4
|
||||
srsrc_str = _reg("ttmp", srsrc_base - 108, 8) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 8)
|
||||
# BVH intersect ray: special case with 4 SGPR srsrc
|
||||
if 'bvh' in name:
|
||||
vaddr = (9 if '64' in name else 8) if inst.a16 else (12 if '64' in name else 11)
|
||||
srsrc = _reg("ttmp", srsrc_base - 108, 4) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 4)
|
||||
return f"{name} {_vreg(inst.vdata, 4)}, {_vreg(inst.vaddr, vaddr)}, {srsrc}{' a16' if inst.a16 else ''}"
|
||||
# vdata width from dmask (gather4/msaa_load always 4), d16 packs, tfe adds 1
|
||||
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
|
||||
# vaddr width
|
||||
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)
|
||||
# modifiers
|
||||
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 for sample/gather/get_lod
|
||||
ssamp_str = ""
|
||||
if 'sample' in name or 'gather' in name or 'get_lod' in name:
|
||||
ssamp_base = inst.ssamp * 4
|
||||
ssamp_str = ", " + (_reg("ttmp", ssamp_base - 108, 4) if 108 <= ssamp_base <= 123 else _sreg(ssamp_base, 4))
|
||||
return f"{name} {_vreg(inst.vdata, vdata)}, {vaddr_str}, {srsrc_str}{ssamp_str} {' '.join(mods)}"
|
||||
|
||||
def _sop_widths(name: str) -> tuple[int, int, int]:
|
||||
"""Return (dst_width, src0_width, src1_width) in register count for SOP instructions."""
|
||||
if name in ('s_bitset0_b64', 's_bitset1_b64', 's_bfm_b64'): return 2, 1, 1
|
||||
if name in ('s_lshl_b64', 's_lshr_b64', 's_ashr_i64', 's_bfe_u64', 's_bfe_i64'): return 2, 2, 1
|
||||
if 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)$', name): return 2 if m.group(2) == '64' else 1, 2 if m.group(4) == '64' else 1, 1
|
||||
if m := re.search(r'_(b|i|u)(32|64)$', name): sz = 2 if m.group(2) == '64' else 1; return sz, sz, sz
|
||||
return 1, 1, 1
|
||||
|
||||
def _disasm_sop1(inst: SOP1) -> str:
|
||||
op, name = SOP1Op(inst.op), SOP1Op(inst.op).name.lower()
|
||||
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} {_fmt_src(inst.ssrc0, 2)}"
|
||||
if op == SOP1Op.S_SWAPPC_B64: return f"{name} {_fmt_sdst(inst.sdst, 2)}, {_fmt_src(inst.ssrc0, 2)}"
|
||||
if op in (SOP1Op.S_SENDMSG_RTN_B32, SOP1Op.S_SENDMSG_RTN_B64): return f"{name} {_fmt_sdst(inst.sdst, 2 if 'b64' in name else 1)}, sendmsg({MSG.get(inst.ssrc0, str(inst.ssrc0))})"
|
||||
dn, s0n, _ = _sop_widths(name)
|
||||
return f"{name} {_fmt_sdst(inst.sdst, dn)}, {inst.lit(inst.ssrc0) if s0n == 1 else _fmt_src(inst.ssrc0, s0n)}"
|
||||
|
||||
def _disasm_sop2(inst: SOP2) -> str:
|
||||
name = SOP2Op(inst.op).name.lower()
|
||||
dn, s0n, s1n = _sop_widths(name)
|
||||
return f"{name} {_fmt_sdst(inst.sdst, dn)}, {inst.lit(inst.ssrc0) if inst.ssrc0 == 255 else _fmt_src(inst.ssrc0, s0n)}, {inst.lit(inst.ssrc1) if inst.ssrc1 == 255 else _fmt_src(inst.ssrc1, s1n)}"
|
||||
|
||||
def _disasm_sopc(inst: SOPC) -> str:
|
||||
name = SOPCOp(inst.op).name.lower()
|
||||
_, s0n, s1n = _sop_widths(name)
|
||||
return f"{name} {_fmt_src(inst.ssrc0, s0n)}, {_fmt_src(inst.ssrc1, s1n)}"
|
||||
|
||||
def _disasm_sopk(inst: SOPK) -> str:
|
||||
op, name = SOPKOp(inst.op), SOPKOp(inst.op).name.lower()
|
||||
if op == SOPKOp.S_VERSION: return f"{name} 0x{inst.simm16:x}"
|
||||
if op in (SOPKOp.S_SETREG_B32, SOPKOp.S_GETREG_B32):
|
||||
hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
|
||||
hs = f"0x{inst.simm16:x}" if hid in (16, 17) else f"hwreg({HWREG.get(hid, str(hid))}, {hoff}, {hsz})"
|
||||
return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1)}" if op == SOPKOp.S_SETREG_B32 else f"{name} {_fmt_sdst(inst.sdst, 1)}, {hs}"
|
||||
dn, _, _ = _sop_widths(name)
|
||||
return f"{name} {_fmt_sdst(inst.sdst, dn)}, 0x{inst.simm16:x}"
|
||||
|
||||
def _disasm_vinterp(inst: VINTERP) -> str:
|
||||
name = VINTERPOp(inst.op).name.lower()
|
||||
src0 = f"-{inst.lit(inst.src0)}" if inst.neg & 1 else inst.lit(inst.src0)
|
||||
src1 = f"-{inst.lit(inst.src1)}" if inst.neg & 2 else inst.lit(inst.src1)
|
||||
src2 = f"-{inst.lit(inst.src2)}" if inst.neg & 4 else inst.lit(inst.src2)
|
||||
mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
|
||||
return f"{name} v{inst.vdst}, {src0}, {src1}, {src2}" + (" " + mods if mods else "")
|
||||
|
||||
def _disasm_generic(inst: Inst) -> str:
|
||||
name = f"op_{inst.op}"
|
||||
def format_field(field_name, val):
|
||||
val = unwrap(val)
|
||||
if field_name in SRC_FIELDS: return inst.lit(val) if val != 255 else "0xff"
|
||||
return f"{'s' if field_name == 'sdst' else 'v'}{val}" if field_name in ('sdst', 'vdst') else f"v{val}" if field_name == 'vsrc1' else f"0x{val:x}" if field_name == 'simm16' else str(val)
|
||||
operands = [format_field(field_name, inst._values.get(field_name, 0)) for field_name in inst._fields if field_name not in ('encoding', 'op')]
|
||||
return f"{name} {', '.join(operands)}" if operands else name
|
||||
|
||||
DISASM_HANDLERS = {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}
|
||||
|
||||
def disasm(inst: Inst) -> str: return DISASM_HANDLERS.get(type(inst), _disasm_generic)(inst)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# 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 = {'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}
|
||||
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'}
|
||||
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:]
|
||||
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
|
||||
|
||||
def get_dsl(text: str) -> str:
|
||||
text, kw = 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
|
||||
if (m := _extract(text, r'\s+clamp(?:\s|$)'))[0]: kw.append('clmp=1'); 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'))
|
||||
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))
|
||||
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+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
|
||||
if waitexp: kw.append(f'waitexp={waitexp}')
|
||||
|
||||
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}")
|
||||
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]})"
|
||||
|
||||
# SMEM
|
||||
if mn in SMEM_OPS:
|
||||
gs, ds = ", glc=1" if glc else "", ", dlc=1" if dlc else ""
|
||||
if len(ops) >= 3 and re.match(r'^-?[0-9]|^-?0x', ops[2].strip().lower()):
|
||||
return f"{mn}(sdata={args[0]}, sbase={args[1]}, offset={args[2]}, soffset=RawImm(124){gs}{ds})"
|
||||
if off_val and len(ops) >= 3: return f"{mn}(sdata={args[0]}, sbase={args[1]}, offset={off_val}, soffset={args[2]}{gs}{ds})"
|
||||
if len(ops) >= 3: return f"{mn}(sdata={args[0]}, sbase={args[1]}, soffset={args[2]}{gs}{ds})"
|
||||
|
||||
# Buffer
|
||||
if mn.startswith('buffer_') and len(ops) >= 2 and ops[1].strip().lower() == 'off':
|
||||
return f"{mn}(vdata={args[0]}, vaddr=0, srsrc={args[2]}, soffset={f'RawImm({args[3].strip()})' if len(args) > 3 else 'RawImm(0)'})"
|
||||
|
||||
# FLAT/GLOBAL/SCRATCH load/store/atomic - saddr needs RawImm(124) for off/null
|
||||
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_'):
|
||||
off0, off1 = (str(int(off_val, 0) & 0xff), str((int(off_val, 0) >> 8) & 0xff)) if off_val else ("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]]
|
||||
|
||||
# VCC ops cleanup
|
||||
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: mn, args = mn.replace('_e32', '') + '_e32', [args[0], args[2], args[3]]
|
||||
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
|
||||
|
||||
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'])
|
||||
|
||||
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 asm(text: str) -> Inst:
|
||||
from extra.assembly.amd.autogen import rdna3 as ag
|
||||
dsl = get_dsl(text)
|
||||
ns = {n: getattr(ag, n) for n in dir(ag) if not n.startswith('_')}
|
||||
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})
|
||||
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
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,320 @@
|
||||
# autogenerated from AMD ISA XML - do not edit
|
||||
from enum import Enum, auto
|
||||
|
||||
class Fmt(Enum):
|
||||
FMT_ANY = auto()
|
||||
FMT_BUF = auto()
|
||||
FMT_IMG = auto()
|
||||
FMT_IMG_BVH = auto()
|
||||
FMT_NUM_B1 = auto()
|
||||
FMT_NUM_B1024 = auto()
|
||||
FMT_NUM_B128 = auto()
|
||||
FMT_NUM_B16 = auto()
|
||||
FMT_NUM_B256 = auto()
|
||||
FMT_NUM_B32 = auto()
|
||||
FMT_NUM_B512 = auto()
|
||||
FMT_NUM_B64 = auto()
|
||||
FMT_NUM_B8 = auto()
|
||||
FMT_NUM_B96 = auto()
|
||||
FMT_NUM_BF16 = auto()
|
||||
FMT_NUM_BF6 = auto()
|
||||
FMT_NUM_BF8 = auto()
|
||||
FMT_NUM_F16 = auto()
|
||||
FMT_NUM_F32 = auto()
|
||||
FMT_NUM_F64 = auto()
|
||||
FMT_NUM_FP4 = auto()
|
||||
FMT_NUM_FP6 = auto()
|
||||
FMT_NUM_FP8 = auto()
|
||||
FMT_NUM_I16 = auto()
|
||||
FMT_NUM_I24 = auto()
|
||||
FMT_NUM_I32 = auto()
|
||||
FMT_NUM_I4 = auto()
|
||||
FMT_NUM_I64 = auto()
|
||||
FMT_NUM_I8 = auto()
|
||||
FMT_NUM_IU4 = auto()
|
||||
FMT_NUM_IU8 = auto()
|
||||
FMT_NUM_M64 = auto()
|
||||
FMT_NUM_PK16_BF16 = auto()
|
||||
FMT_NUM_PK16_BF8 = auto()
|
||||
FMT_NUM_PK16_F16 = auto()
|
||||
FMT_NUM_PK16_F32 = auto()
|
||||
FMT_NUM_PK16_FP8 = auto()
|
||||
FMT_NUM_PK16_I32 = auto()
|
||||
FMT_NUM_PK16_I8 = auto()
|
||||
FMT_NUM_PK2_B16 = auto()
|
||||
FMT_NUM_PK2_B32 = auto()
|
||||
FMT_NUM_PK2_B64 = auto()
|
||||
FMT_NUM_PK2_BF16 = auto()
|
||||
FMT_NUM_PK2_BF8 = auto()
|
||||
FMT_NUM_PK2_F16 = auto()
|
||||
FMT_NUM_PK2_F32 = auto()
|
||||
FMT_NUM_PK2_FP4 = auto()
|
||||
FMT_NUM_PK2_FP8 = auto()
|
||||
FMT_NUM_PK2_I16 = auto()
|
||||
FMT_NUM_PK2_I8 = auto()
|
||||
FMT_NUM_PK2_U16 = auto()
|
||||
FMT_NUM_PK2_U8 = auto()
|
||||
FMT_NUM_PK32_BF16 = auto()
|
||||
FMT_NUM_PK32_BF6 = auto()
|
||||
FMT_NUM_PK32_BF8 = auto()
|
||||
FMT_NUM_PK32_F16 = auto()
|
||||
FMT_NUM_PK32_F32 = auto()
|
||||
FMT_NUM_PK32_FP6 = auto()
|
||||
FMT_NUM_PK32_FP8 = auto()
|
||||
FMT_NUM_PK32_I32 = auto()
|
||||
FMT_NUM_PK32_I8 = auto()
|
||||
FMT_NUM_PK4_B8 = auto()
|
||||
FMT_NUM_PK4_BF16 = auto()
|
||||
FMT_NUM_PK4_BF8 = auto()
|
||||
FMT_NUM_PK4_F16 = auto()
|
||||
FMT_NUM_PK4_F32 = auto()
|
||||
FMT_NUM_PK4_F64 = auto()
|
||||
FMT_NUM_PK4_FP8 = auto()
|
||||
FMT_NUM_PK4_I32 = auto()
|
||||
FMT_NUM_PK4_I8 = auto()
|
||||
FMT_NUM_PK4_IU8 = auto()
|
||||
FMT_NUM_PK4_U8 = auto()
|
||||
FMT_NUM_PK8_B32 = auto()
|
||||
FMT_NUM_PK8_BF16 = auto()
|
||||
FMT_NUM_PK8_BF8 = auto()
|
||||
FMT_NUM_PK8_F16 = auto()
|
||||
FMT_NUM_PK8_FP8 = auto()
|
||||
FMT_NUM_PK8_I4 = auto()
|
||||
FMT_NUM_PK8_I8 = auto()
|
||||
FMT_NUM_PK8_IU4 = auto()
|
||||
FMT_NUM_PK8_U4 = auto()
|
||||
FMT_NUM_PK8_U8 = auto()
|
||||
FMT_NUM_PK_F16 = auto()
|
||||
FMT_NUM_PK_I16 = auto()
|
||||
FMT_NUM_PK_I8 = auto()
|
||||
FMT_NUM_PK_U16 = auto()
|
||||
FMT_NUM_PK_U8 = auto()
|
||||
FMT_NUM_U16 = auto()
|
||||
FMT_NUM_U24 = auto()
|
||||
FMT_NUM_U32 = auto()
|
||||
FMT_NUM_U4 = auto()
|
||||
FMT_NUM_U64 = auto()
|
||||
FMT_NUM_U8 = auto()
|
||||
FMT_RSRC = auto()
|
||||
FMT_RSRC_SCALAR = auto()
|
||||
FMT_RSRC_SCRATCH = auto()
|
||||
FMT_RSRC_SCRATCH_BYTE = auto()
|
||||
FMT_RSRC_SCRATCH_STRIDE = auto()
|
||||
FMT_RSRC_TYPED = auto()
|
||||
FMT_RSRC_TYPED_BYTE = auto()
|
||||
FMT_RSRC_TYPED_SCRATCH = auto()
|
||||
FMT_RSRC_TYPED_STRIDE = auto()
|
||||
FMT_RSRC_VECTOR = auto()
|
||||
FMT_RSRC_VECTOR_BYTE = auto()
|
||||
FMT_RSRC_VECTOR_STRIDE = auto()
|
||||
FMT_SAMP = auto()
|
||||
FMT_WMMA_AB_16X16_BF16 = auto()
|
||||
FMT_WMMA_AB_16X16_BF8 = auto()
|
||||
FMT_WMMA_AB_16X16_F16 = auto()
|
||||
FMT_WMMA_AB_16X16_FP8 = auto()
|
||||
FMT_WMMA_AB_16X16_IU4 = auto()
|
||||
FMT_WMMA_AB_16X16_IU8 = auto()
|
||||
FMT_WMMA_AB_16X32_BF16 = auto()
|
||||
FMT_WMMA_AB_16X32_BF8 = auto()
|
||||
FMT_WMMA_AB_16X32_F16 = auto()
|
||||
FMT_WMMA_AB_16X32_FP8 = auto()
|
||||
FMT_WMMA_AB_16X32_IU4 = auto()
|
||||
FMT_WMMA_AB_16X32_IU8 = auto()
|
||||
FMT_WMMA_AB_16X64_IU4 = auto()
|
||||
FMT_WMMA_AB_BF16 = auto()
|
||||
FMT_WMMA_AB_F16 = auto()
|
||||
FMT_WMMA_AB_IU4 = auto()
|
||||
FMT_WMMA_AB_IU8 = auto()
|
||||
FMT_WMMA_DC_16X16_BF16 = auto()
|
||||
FMT_WMMA_DC_16X16_F16 = auto()
|
||||
FMT_WMMA_DC_16X16_F32 = auto()
|
||||
FMT_WMMA_DC_16X16_I32 = auto()
|
||||
FMT_WMMA_DC_BF16 = auto()
|
||||
FMT_WMMA_DC_F16 = auto()
|
||||
FMT_WMMA_DC_F32 = auto()
|
||||
FMT_WMMA_DC_I32 = auto()
|
||||
FMT_WMMA_INDEX_SET = auto()
|
||||
|
||||
FMT_BITS = {
|
||||
Fmt.FMT_ANY: 1,
|
||||
Fmt.FMT_BUF: 64,
|
||||
Fmt.FMT_IMG: 256,
|
||||
Fmt.FMT_IMG_BVH: 128,
|
||||
Fmt.FMT_NUM_B1: 1,
|
||||
Fmt.FMT_NUM_B1024: 1024,
|
||||
Fmt.FMT_NUM_B128: 128,
|
||||
Fmt.FMT_NUM_B16: 16,
|
||||
Fmt.FMT_NUM_B256: 256,
|
||||
Fmt.FMT_NUM_B32: 32,
|
||||
Fmt.FMT_NUM_B512: 512,
|
||||
Fmt.FMT_NUM_B64: 64,
|
||||
Fmt.FMT_NUM_B8: 8,
|
||||
Fmt.FMT_NUM_B96: 96,
|
||||
Fmt.FMT_NUM_BF16: 16,
|
||||
Fmt.FMT_NUM_BF6: 6,
|
||||
Fmt.FMT_NUM_BF8: 8,
|
||||
Fmt.FMT_NUM_F16: 16,
|
||||
Fmt.FMT_NUM_F32: 32,
|
||||
Fmt.FMT_NUM_F64: 64,
|
||||
Fmt.FMT_NUM_FP4: 4,
|
||||
Fmt.FMT_NUM_FP6: 6,
|
||||
Fmt.FMT_NUM_FP8: 8,
|
||||
Fmt.FMT_NUM_I16: 16,
|
||||
Fmt.FMT_NUM_I24: 24,
|
||||
Fmt.FMT_NUM_I32: 32,
|
||||
Fmt.FMT_NUM_I4: 4,
|
||||
Fmt.FMT_NUM_I64: 64,
|
||||
Fmt.FMT_NUM_I8: 8,
|
||||
Fmt.FMT_NUM_IU4: 4,
|
||||
Fmt.FMT_NUM_IU8: 8,
|
||||
Fmt.FMT_NUM_M64: 64,
|
||||
Fmt.FMT_NUM_PK16_BF16: 256,
|
||||
Fmt.FMT_NUM_PK16_BF8: 128,
|
||||
Fmt.FMT_NUM_PK16_F16: 256,
|
||||
Fmt.FMT_NUM_PK16_F32: 512,
|
||||
Fmt.FMT_NUM_PK16_FP8: 128,
|
||||
Fmt.FMT_NUM_PK16_I32: 512,
|
||||
Fmt.FMT_NUM_PK16_I8: 128,
|
||||
Fmt.FMT_NUM_PK2_B16: 32,
|
||||
Fmt.FMT_NUM_PK2_B32: 64,
|
||||
Fmt.FMT_NUM_PK2_B64: 128,
|
||||
Fmt.FMT_NUM_PK2_BF16: 32,
|
||||
Fmt.FMT_NUM_PK2_BF8: 16,
|
||||
Fmt.FMT_NUM_PK2_F16: 32,
|
||||
Fmt.FMT_NUM_PK2_F32: 64,
|
||||
Fmt.FMT_NUM_PK2_FP4: 8,
|
||||
Fmt.FMT_NUM_PK2_FP8: 16,
|
||||
Fmt.FMT_NUM_PK2_I16: 32,
|
||||
Fmt.FMT_NUM_PK2_I8: 16,
|
||||
Fmt.FMT_NUM_PK2_U16: 32,
|
||||
Fmt.FMT_NUM_PK2_U8: 16,
|
||||
Fmt.FMT_NUM_PK32_BF16: 512,
|
||||
Fmt.FMT_NUM_PK32_BF6: 192,
|
||||
Fmt.FMT_NUM_PK32_BF8: 256,
|
||||
Fmt.FMT_NUM_PK32_F16: 512,
|
||||
Fmt.FMT_NUM_PK32_F32: 1024,
|
||||
Fmt.FMT_NUM_PK32_FP6: 192,
|
||||
Fmt.FMT_NUM_PK32_FP8: 256,
|
||||
Fmt.FMT_NUM_PK32_I32: 1024,
|
||||
Fmt.FMT_NUM_PK32_I8: 256,
|
||||
Fmt.FMT_NUM_PK4_B8: 32,
|
||||
Fmt.FMT_NUM_PK4_BF16: 64,
|
||||
Fmt.FMT_NUM_PK4_BF8: 32,
|
||||
Fmt.FMT_NUM_PK4_F16: 64,
|
||||
Fmt.FMT_NUM_PK4_F32: 128,
|
||||
Fmt.FMT_NUM_PK4_F64: 256,
|
||||
Fmt.FMT_NUM_PK4_FP8: 32,
|
||||
Fmt.FMT_NUM_PK4_I32: 128,
|
||||
Fmt.FMT_NUM_PK4_I8: 32,
|
||||
Fmt.FMT_NUM_PK4_IU8: 32,
|
||||
Fmt.FMT_NUM_PK4_U8: 32,
|
||||
Fmt.FMT_NUM_PK8_B32: 256,
|
||||
Fmt.FMT_NUM_PK8_BF16: 128,
|
||||
Fmt.FMT_NUM_PK8_BF8: 64,
|
||||
Fmt.FMT_NUM_PK8_F16: 128,
|
||||
Fmt.FMT_NUM_PK8_FP8: 64,
|
||||
Fmt.FMT_NUM_PK8_I4: 32,
|
||||
Fmt.FMT_NUM_PK8_I8: 64,
|
||||
Fmt.FMT_NUM_PK8_IU4: 32,
|
||||
Fmt.FMT_NUM_PK8_U4: 32,
|
||||
Fmt.FMT_NUM_PK8_U8: 64,
|
||||
Fmt.FMT_NUM_PK_F16: 32,
|
||||
Fmt.FMT_NUM_PK_I16: 32,
|
||||
Fmt.FMT_NUM_PK_I8: 32,
|
||||
Fmt.FMT_NUM_PK_U16: 32,
|
||||
Fmt.FMT_NUM_PK_U8: 32,
|
||||
Fmt.FMT_NUM_U16: 16,
|
||||
Fmt.FMT_NUM_U24: 24,
|
||||
Fmt.FMT_NUM_U32: 32,
|
||||
Fmt.FMT_NUM_U4: 4,
|
||||
Fmt.FMT_NUM_U64: 64,
|
||||
Fmt.FMT_NUM_U8: 8,
|
||||
Fmt.FMT_RSRC: 128,
|
||||
Fmt.FMT_RSRC_SCALAR: 128,
|
||||
Fmt.FMT_RSRC_SCRATCH: 128,
|
||||
Fmt.FMT_RSRC_SCRATCH_BYTE: 128,
|
||||
Fmt.FMT_RSRC_SCRATCH_STRIDE: 128,
|
||||
Fmt.FMT_RSRC_TYPED: 128,
|
||||
Fmt.FMT_RSRC_TYPED_BYTE: 128,
|
||||
Fmt.FMT_RSRC_TYPED_SCRATCH: 128,
|
||||
Fmt.FMT_RSRC_TYPED_STRIDE: 128,
|
||||
Fmt.FMT_RSRC_VECTOR: 128,
|
||||
Fmt.FMT_RSRC_VECTOR_BYTE: 128,
|
||||
Fmt.FMT_RSRC_VECTOR_STRIDE: 128,
|
||||
Fmt.FMT_SAMP: 128,
|
||||
Fmt.FMT_WMMA_AB_16X16_BF16: 128,
|
||||
Fmt.FMT_WMMA_AB_16X16_BF8: 64,
|
||||
Fmt.FMT_WMMA_AB_16X16_F16: 128,
|
||||
Fmt.FMT_WMMA_AB_16X16_FP8: 64,
|
||||
Fmt.FMT_WMMA_AB_16X16_IU4: 32,
|
||||
Fmt.FMT_WMMA_AB_16X16_IU8: 64,
|
||||
Fmt.FMT_WMMA_AB_16X32_BF16: 256,
|
||||
Fmt.FMT_WMMA_AB_16X32_BF8: 128,
|
||||
Fmt.FMT_WMMA_AB_16X32_F16: 256,
|
||||
Fmt.FMT_WMMA_AB_16X32_FP8: 128,
|
||||
Fmt.FMT_WMMA_AB_16X32_IU4: 64,
|
||||
Fmt.FMT_WMMA_AB_16X32_IU8: 128,
|
||||
Fmt.FMT_WMMA_AB_16X64_IU4: 128,
|
||||
Fmt.FMT_WMMA_AB_BF16: 256,
|
||||
Fmt.FMT_WMMA_AB_F16: 256,
|
||||
Fmt.FMT_WMMA_AB_IU4: 64,
|
||||
Fmt.FMT_WMMA_AB_IU8: 128,
|
||||
Fmt.FMT_WMMA_DC_16X16_BF16: 128,
|
||||
Fmt.FMT_WMMA_DC_16X16_F16: 128,
|
||||
Fmt.FMT_WMMA_DC_16X16_F32: 256,
|
||||
Fmt.FMT_WMMA_DC_16X16_I32: 256,
|
||||
Fmt.FMT_WMMA_DC_BF16: 256,
|
||||
Fmt.FMT_WMMA_DC_F16: 256,
|
||||
Fmt.FMT_WMMA_DC_F32: 256,
|
||||
Fmt.FMT_WMMA_DC_I32: 256,
|
||||
Fmt.FMT_WMMA_INDEX_SET: 32,
|
||||
}
|
||||
|
||||
class OpType(Enum):
|
||||
OPR_ACCVGPR = auto()
|
||||
OPR_ATTR = auto()
|
||||
OPR_CLAUSE = auto()
|
||||
OPR_DELAY = auto()
|
||||
OPR_EXEC = auto()
|
||||
OPR_HWREG = auto()
|
||||
OPR_LABEL = auto()
|
||||
OPR_SDST = auto()
|
||||
OPR_SDST_NULL = auto()
|
||||
OPR_SENDMSG = auto()
|
||||
OPR_SENDMSG_RTN = auto()
|
||||
OPR_SIMM16 = auto()
|
||||
OPR_SIMM24 = auto()
|
||||
OPR_SIMM4 = auto()
|
||||
OPR_SIMM5 = auto()
|
||||
OPR_SIMM8 = auto()
|
||||
OPR_SLEEP = auto()
|
||||
OPR_SMEM_OFFSET = auto()
|
||||
OPR_SMEM_OFFSET_NOK = auto()
|
||||
OPR_SRC = auto()
|
||||
OPR_SRC_ACCVGPR = auto()
|
||||
OPR_SRC_NOLDS = auto()
|
||||
OPR_SRC_NOLIT = auto()
|
||||
OPR_SRC_SIMPLE = auto()
|
||||
OPR_SRC_VGPR = auto()
|
||||
OPR_SRC_VGPR_OR_ACCVGPR = auto()
|
||||
OPR_SRC_VGPR_OR_ACCVGPR_OR_CONST = auto()
|
||||
OPR_SRC_VGPR_OR_INLINE = auto()
|
||||
OPR_SREG = auto()
|
||||
OPR_SREG_LITERAL = auto()
|
||||
OPR_SREG_M0 = auto()
|
||||
OPR_SREG_M0_INL = auto()
|
||||
OPR_SREG_NOVCC = auto()
|
||||
OPR_SSRC = auto()
|
||||
OPR_SSRC_BARRIER_ID = auto()
|
||||
OPR_SSRC_LANESEL = auto()
|
||||
OPR_SSRC_NOLIT = auto()
|
||||
OPR_TGT = auto()
|
||||
OPR_VERSION = auto()
|
||||
OPR_VGPR = auto()
|
||||
OPR_VGPR_OR_ACCVGPR = auto()
|
||||
OPR_VGPR_OR_LDS = auto()
|
||||
OPR_WAITCNT = auto()
|
||||
OPR_WAITCNT_DEPCTR = auto()
|
||||
OPR_WAIT_ALU = auto()
|
||||
OPR_WAIT_EVENT = auto()
|
||||
OPR_WAIT_MEM_DS = auto()
|
||||
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
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
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,83 @@
|
||||
# Instruction format detection and decoding
|
||||
from __future__ import annotations
|
||||
from extra.assembly.amd.dsl import Inst, FixedBitField
|
||||
from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP2, VOP3, VOP3_SDST, VOP3SD, VOP3P, VOPC, VOPD, VINTERP,
|
||||
SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT)
|
||||
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP2 as R4_VOP2,
|
||||
VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
|
||||
VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOPC as R4_SOPC, SOPK as R4_SOPK, SOPP as R4_SOPP,
|
||||
SMEM as R4_SMEM, DS as R4_DS)
|
||||
from extra.assembly.amd.autogen.cdna.ins import (VOP1 as C_VOP1, VOP2 as C_VOP2, VOPC as C_VOPC, VOP3A, VOP3B, VOP3P as C_VOP3P,
|
||||
SOP1 as C_SOP1, SOP2 as C_SOP2, SOPC as C_SOPC, SOPK as C_SOPK, SOPP as C_SOPP, SMEM as C_SMEM, DS as C_DS, FLAT as C_FLAT)
|
||||
|
||||
def _matches_encoding(word: int, cls: type[Inst]) -> bool:
|
||||
"""Check if word matches the encoding pattern of an instruction class."""
|
||||
enc = next(((n, f) for n, f in cls._fields if isinstance(f, FixedBitField) and n == 'encoding'), None)
|
||||
if enc is None: return False
|
||||
bf = enc[1]
|
||||
return ((word >> bf.lo) & bf.mask) == bf.default
|
||||
|
||||
# Order matters: more specific encodings first, VOP2 last (it's a catch-all for bit31=0)
|
||||
_RDNA_FORMATS_64 = [VOPD, VOP3P, VINTERP, VOP3, DS, FLAT, SMEM]
|
||||
_RDNA_FORMATS_32 = [SOP1, SOPC, SOPP, SOPK, VOPC, VOP1, SOP2, VOP2] # SOP2/VOP2 are catch-alls
|
||||
_CDNA_FORMATS_64 = [C_VOP3P, VOP3A, C_DS, C_FLAT, C_SMEM]
|
||||
_CDNA_FORMATS_32 = [C_SOP1, C_SOPC, C_SOPP, C_SOPK, C_VOPC, C_VOP1, C_SOP2, C_VOP2]
|
||||
_CDNA_VOP3B_OPS = {281, 282, 283, 284, 285, 286, 480, 481, 488, 489} # VOP3B opcodes
|
||||
_RDNA4_FORMATS_64 = [R4_VOPD, R4_VOP3P, R4_VINTERP, R4_VOP3, R4_DS, R4_SMEM]
|
||||
_RDNA4_FORMATS_32 = [R4_SOP1, R4_SOPC, R4_SOPP, R4_SOPK, R4_VOPC, R4_VOP1, R4_SOP2, R4_VOP2]
|
||||
_RDNA4_VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
|
||||
_RDNA3_VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
|
||||
# Instructions with SGPR destination (READLANE, READFIRSTLANE, and VOP3-encoded VOPC)
|
||||
_VOP1_SDST_OPS = {2} # V_READFIRSTLANE_B32_E32
|
||||
_VOP3_SDST_OPS = {386, 864} # V_READFIRSTLANE_B32_E64, V_READLANE_B32 (V_WRITELANE_B32=865 writes to VGPR)
|
||||
# VOP3-encoded VOPC instructions (opcodes < 256) also have SGPR destination
|
||||
|
||||
def detect_format(data: bytes, arch: str = "rdna3") -> type[Inst]:
|
||||
"""Detect instruction format from machine code bytes."""
|
||||
assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
|
||||
word = int.from_bytes(data[:4], 'little')
|
||||
if arch == "cdna":
|
||||
# SDWA (0xf9) and DPP (0xfa) use special src0 encodings - not supported
|
||||
if (word & 0x1ff) in (0xf9, 0xfa): raise ValueError(f"CDNA SDWA/DPP not supported word={word:#010x}")
|
||||
if (word >> 30) == 0b11:
|
||||
for cls in _CDNA_FORMATS_64:
|
||||
if _matches_encoding(word, cls):
|
||||
return VOP3B if cls is VOP3A and ((word >> 16) & 0x3ff) in _CDNA_VOP3B_OPS else cls
|
||||
raise ValueError(f"unknown CDNA 64-bit format word={word:#010x}")
|
||||
for cls in _CDNA_FORMATS_32:
|
||||
if _matches_encoding(word, cls): return cls
|
||||
raise ValueError(f"unknown CDNA 32-bit format word={word:#010x}")
|
||||
if arch == "rdna4":
|
||||
if (word >> 30) == 0b11:
|
||||
for cls in _RDNA4_FORMATS_64:
|
||||
if _matches_encoding(word, cls):
|
||||
if cls is R4_VOP3:
|
||||
opcode = (word >> 16) & 0x3ff
|
||||
if opcode in _RDNA4_VOP3SD_OPS: return R4_VOP3SD
|
||||
if opcode in _VOP3_SDST_OPS or opcode < 256: return R4_VOP3_SDST # VOP3-encoded VOPC (op < 256) writes to SGPR
|
||||
return cls
|
||||
raise ValueError(f"unknown RDNA4 64-bit format word={word:#010x}")
|
||||
for cls in _RDNA4_FORMATS_32:
|
||||
if _matches_encoding(word, cls):
|
||||
if cls is R4_VOP1 and ((word >> 9) & 0xff) in _VOP1_SDST_OPS: return R4_VOP1_SDST
|
||||
return cls
|
||||
raise ValueError(f"unknown RDNA4 32-bit format word={word:#010x}")
|
||||
# RDNA3 (default)
|
||||
if (word >> 30) == 0b11:
|
||||
for cls in _RDNA_FORMATS_64:
|
||||
if _matches_encoding(word, cls):
|
||||
if cls is VOP3:
|
||||
opcode = (word >> 16) & 0x3ff
|
||||
if opcode in _RDNA3_VOP3SD_OPS: return VOP3SD
|
||||
if opcode in _VOP3_SDST_OPS or opcode < 256: return VOP3_SDST # VOP3-encoded VOPC (op < 256) writes to SGPR
|
||||
return cls
|
||||
raise ValueError(f"unknown 64-bit format word={word:#010x}")
|
||||
for cls in _RDNA_FORMATS_32:
|
||||
if _matches_encoding(word, cls):
|
||||
if cls is VOP1 and ((word >> 9) & 0xff) in _VOP1_SDST_OPS: return VOP1_SDST
|
||||
return cls
|
||||
raise ValueError(f"unknown 32-bit format word={word:#010x}")
|
||||
|
||||
def decode_inst(data: bytes, arch: str = "rdna3") -> Inst:
|
||||
"""Decode machine code bytes into an instruction."""
|
||||
return detect_format(data, arch).from_bytes(data)
|
||||
@@ -0,0 +1,671 @@
|
||||
# RDNA3/RDNA4/CDNA disassembler
|
||||
from __future__ import annotations
|
||||
import re, struct
|
||||
from extra.assembly.amd.dsl import Inst, Reg
|
||||
|
||||
# Special register mappings for disassembly
|
||||
SPECIAL_GPRS = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'null', 125: 'm0', 126: 'exec_lo', 127: 'exec_hi',
|
||||
128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc'}
|
||||
SPECIAL_GPRS_CDNA = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'null', 125: 'm0', 126: 'exec_lo', 127: 'exec_hi',
|
||||
128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc',
|
||||
102: 'flat_scratch_lo', 103: 'flat_scratch_hi', 104: 'xnack_mask_lo', 105: 'xnack_mask_hi'}
|
||||
SPECIAL_PAIRS = {106: 'vcc', 126: 'exec'}
|
||||
SPECIAL_PAIRS_CDNA = {106: 'vcc', 126: 'exec', 102: 'flat_scratch', 104: 'xnack_mask'}
|
||||
|
||||
def decode_src(v, cdna: bool = False) -> str:
|
||||
"""Decode a source operand encoding to its string representation."""
|
||||
v = _unwrap(v)
|
||||
gprs = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
|
||||
if v in gprs: return gprs[v]
|
||||
if v < 106: return f's{v}'
|
||||
if 108 <= v < 124: return f'ttmp{v - 108}'
|
||||
if 129 <= v <= 192: return str(v - 128) # positive integers 1-64
|
||||
if 193 <= v <= 208: return str(-(v - 192)) # negative integers -1 to -16
|
||||
if v >= 256: return f'v{v - 256}'
|
||||
return f's{v}'
|
||||
|
||||
def _unwrap(v) -> int:
|
||||
"""Unwrap Reg to int offset, or return int as-is."""
|
||||
return v.offset if isinstance(v, Reg) else v
|
||||
|
||||
def _vi(v) -> int:
|
||||
"""Get VGPR index from Reg or int (for v[N] fields that encode as 256+N)."""
|
||||
off = _unwrap(v)
|
||||
return off - 256 if off >= 256 else off
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# LITERAL FORMATTING
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
_FLOAT_DEC = {240: 0.5, 241: -0.5, 242: 1.0, 243: -1.0, 244: 2.0, 245: -2.0, 246: 4.0, 247: -4.0}
|
||||
|
||||
def _lit(inst, v, neg=0) -> str:
|
||||
"""Format literal/inline constant value."""
|
||||
v = _unwrap(v)
|
||||
if v == 255:
|
||||
lit = inst._literal
|
||||
if lit is None: return "0"
|
||||
s = f"0x{lit:x}"
|
||||
elif v in _FLOAT_DEC: s = str(_FLOAT_DEC[v])
|
||||
elif 128 <= v <= 192: s = str(v - 128)
|
||||
elif 193 <= v <= 208: s = str(-(v - 192))
|
||||
elif v < 128: s = decode_src(v)
|
||||
elif v >= 256: s = f"v{v - 256}"
|
||||
else: s = decode_src(v)
|
||||
return f"-{s}" if neg else s
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# INSTRUCTION METADATA - fallback functions when inst.num_srcs()/inst.operands unavailable
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _num_srcs(inst) -> int:
|
||||
"""Fallback: get number of source operands from instruction name."""
|
||||
name = getattr(inst, 'op_name', '') or ''
|
||||
n = name.upper()
|
||||
# FMAC/MAC ops are 2-source (dst is implicit accumulator), but FMA/MAD ops are 3-source
|
||||
if 'FMAC' in n or 'V_MAC_' in n: return 2
|
||||
if any(x in n for x in ('FMA', 'MAD', 'CNDMASK', 'BFE', 'BFI', 'LERP', 'MED3', 'SAD', 'DIV_FMAS', 'DIV_FIXUP', 'DIV_SCALE', 'CUBE')): return 3
|
||||
# PERMLANE_VAR ops are 2-source, but PERMLANE (non-VAR) are 3-source
|
||||
if 'PERMLANE' in n and '_VAR' not in n: return 3
|
||||
if any(x in n for x in ('_ADD3', '_LSHL_ADD', '_ADD_LSHL', '_LSHL_OR', '_AND_OR', 'OR3_B32', 'AND_OR_B32', 'ALIGNBIT', 'ALIGNBYTE', 'V_PERM_', 'XOR3', 'XAD', 'MULLIT', 'MINMAX', 'MAXMIN', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM', 'MINIMUM3', 'MAXIMUM3', 'MIN3', 'MAX3', 'DOT2', 'CVT_PK_U8_F32', 'DOT4', 'DOT8', 'WMMA', 'SWMMAC')): return 3
|
||||
return 2
|
||||
|
||||
# SWMMAC register counts: (dst, src0, src1, src2)
|
||||
def _swmmac_regs(name: str) -> tuple[int, int, int, int]:
|
||||
"""Return (dst, src0, src1, src2) register counts for SWMMAC instructions."""
|
||||
if 'f16_16x16x32' in name or 'bf16_16x16x32' in name: return (4, 4, 8, 1)
|
||||
if 'f32_16x16x32_f16' in name or 'f32_16x16x32_bf16' in name: return (8, 4, 8, 1)
|
||||
if 'i32_16x16x32_iu4' in name: return (8, 1, 2, 1)
|
||||
if 'i32_16x16x64_iu4' in name: return (8, 2, 4, 1)
|
||||
if 'i32_16x16x32_iu8' in name or 'f32_16x16x32_fp8' in name or 'f32_16x16x32_bf8' in name: return (8, 2, 4, 1)
|
||||
return (8, 8, 8, 8)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# IMPORTS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP2, VOP3, VOP3_SDST, VOP3SD, VOP3P, VOPC, VOPD, VINTERP, SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH,
|
||||
VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOPDOp, SOP1Op, SOPKOp, SOPPOp, SMEMOp, DSOp)
|
||||
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP2 as R4_VOP2, VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
|
||||
VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOPC as R4_SOPC, SOPK as R4_SOPK, SOPP as R4_SOPP,
|
||||
SMEM as R4_SMEM, DS as R4_DS, VOPDOp as R4_VOPDOp)
|
||||
from extra.assembly.amd.autogen.cdna.ins import FLAT as C_FLAT
|
||||
|
||||
def _is_cdna(inst: Inst) -> bool: return 'cdna' in inst.__class__.__module__
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# CONSTANTS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
HWREG = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
|
||||
6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
|
||||
19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI', 22: 'HW_REG_XNACK_MASK',
|
||||
23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
|
||||
HWREG_RDNA4 = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 4: 'HW_REG_STATE_PRIV', 5: 'HW_REG_GPR_ALLOC',
|
||||
6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 10: 'HW_REG_PERF_SNAPSHOT_DATA', 11: 'HW_REG_PERF_SNAPSHOT_PC_LO',
|
||||
12: 'HW_REG_PERF_SNAPSHOT_PC_HI', 15: 'HW_REG_PERF_SNAPSHOT_DATA1', 16: 'HW_REG_PERF_SNAPSHOT_DATA2',
|
||||
17: 'HW_REG_EXCP_FLAG_PRIV', 18: 'HW_REG_EXCP_FLAG_USER', 19: 'HW_REG_TRAP_CTRL',
|
||||
20: 'HW_REG_SCRATCH_BASE_LO', 21: 'HW_REG_SCRATCH_BASE_HI', 23: 'HW_REG_HW_ID1',
|
||||
24: 'HW_REG_HW_ID2', 26: 'HW_REG_SCHED_MODE', 29: 'HW_REG_SHADER_CYCLES_LO',
|
||||
30: 'HW_REG_SHADER_CYCLES_HI', 31: 'HW_REG_DVGPR_ALLOC_LO', 32: 'HW_REG_DVGPR_ALLOC_HI'}
|
||||
MSG = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
|
||||
131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA',
|
||||
134: 'MSG_RTN_GET_TBA_TO_PC', 135: 'MSG_RTN_GET_SE_AID_ID'}
|
||||
# CDNA opcode name aliases for disasm (new name -> old name expected by tests)
|
||||
_CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32', 'v_fmamk_f32': 'v_madmk_f32', 'v_fmaak_f32': 'v_madak_f32'}
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# HELPERS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{_unwrap(b)}" if n == 1 else f"{p}[{_unwrap(b)}:{_unwrap(b)+n-1}]"
|
||||
def _sreg(b: int, n: int = 1) -> str: return _reg("s", _unwrap(b), n)
|
||||
def _vreg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("v", b - 256 if b >= 256 else b, n)
|
||||
def _areg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("a", b - 256 if b >= 256 else b, n) # accumulator registers for GFX90a
|
||||
def _ttmp(b, n: int = 1) -> str: b = _unwrap(b); return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
|
||||
def _sreg_or_ttmp(b, n: int = 1) -> str: return _ttmp(b, n) or _sreg(b, n)
|
||||
|
||||
def _fmt_sdst(v, n: int = 1, cdna: bool = False) -> str:
|
||||
v = _unwrap(v)
|
||||
if t := _ttmp(v, n): return t
|
||||
pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
|
||||
gprs = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
|
||||
if n > 1: return pairs.get(v) or gprs.get(v) or _sreg(v, n) # also check gprs for null/m0
|
||||
return gprs.get(v, f"s{v}")
|
||||
|
||||
def _fmt_src(v, n: int = 1, cdna: bool = False) -> str:
|
||||
v = _unwrap(v)
|
||||
if v == 253: return "src_scc" # SCC as source operand
|
||||
if n == 1: return decode_src(v, cdna)
|
||||
if v >= 256: return _vreg(v, n)
|
||||
if v <= 101: return _sreg(v, n) # s0-s101 can be pairs, but 102+ are special on CDNA
|
||||
pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
|
||||
if n == 2 and v in pairs: return pairs[v]
|
||||
if v <= 105: return _sreg(v, n) # s102-s105 regular pairs for RDNA
|
||||
if t := _ttmp(v, n): return t
|
||||
return decode_src(v, cdna)
|
||||
|
||||
def _fmt_v16(v, base: int = 256, hi_thresh: int = 384) -> str:
|
||||
v = _unwrap(v)
|
||||
return f"v{(v - base) & 0x7f}.{'h' if v >= hi_thresh else 'l'}"
|
||||
|
||||
def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
|
||||
def _omod(v: int) -> str: return {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(v, "")
|
||||
def _src16(inst, v: int) -> str: v = _unwrap(v); return _fmt_v16(v) if v >= 256 else _lit(inst, v) # format 16-bit src: vgpr.h/l or literal
|
||||
def _mods(*pairs) -> str: return " ".join(m for c, m in pairs if c)
|
||||
def _fmt_bits(label: str, val: int, count: int) -> str: return f"{label}:[{','.join(str((val >> i) & 1) for i in range(count))}]"
|
||||
|
||||
def _vop3_src(inst, v: int, neg: int, abs_: int, hi: int, n: int, f16: bool) -> str:
|
||||
"""Format VOP3 source operand with modifiers."""
|
||||
v = _unwrap(v)
|
||||
if v == 255: s = _lit(inst, v) # literal constant takes priority
|
||||
elif n > 1: s = _fmt_src(v, n)
|
||||
elif f16 and v >= 256: s = f"v{v - 256}.h" if hi else f"v{v - 256}.l"
|
||||
elif v == 253: s = "src_scc" # VOP3 sources use src_scc not scc
|
||||
else: s = _lit(inst, v)
|
||||
if abs_: s = f"|{s}|"
|
||||
return f"-{s}" if neg else s
|
||||
|
||||
def _opsel_str(opsel: int, n: int, need: bool, is16_d: bool) -> str:
|
||||
"""Format op_sel modifier string."""
|
||||
if not need: return ""
|
||||
dst_hi = (opsel >> 3) & 1
|
||||
if n == 1: return f" op_sel:[{opsel & 1},{dst_hi}]"
|
||||
if n == 2: return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{dst_hi}]"
|
||||
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{dst_hi}]"
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# DISASSEMBLER
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _disasm_vop1(inst: VOP1) -> str:
|
||||
name, cdna = inst.op_name.lower() or f'vop1_op_{inst.op}', _is_cdna(inst)
|
||||
name = name.replace('_e32', '') # Strip _e32 suffix
|
||||
if any(x in name for x in ('v_nop', 'v_pipeflush', 'v_clrexcp')): return name # no operands
|
||||
if 'readfirstlane' in name:
|
||||
src = inst.src0.fmt() if inst.src0.offset >= 256 else decode_src(inst.src0.offset, cdna)
|
||||
vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
|
||||
return f"{name} {_fmt_sdst(vdst_off, 1, cdna)}, {src}"
|
||||
bits = inst.canonical_op_bits
|
||||
is16_dst, is16_src = not cdna and bits['d'] == 16, not cdna and bits['s0'] == 16
|
||||
# v_cvt_pk_f32_fp8/bf8: pcode has None dst type but outputs 2 VGPRs
|
||||
if 'cvt_pk_f32_fp8' in name or 'cvt_pk_f32_bf8' in name: is16_src = True
|
||||
# Format dst
|
||||
if is16_dst: dst = _fmt_v16(inst.vdst)
|
||||
else: dst = inst.vdst.fmt()
|
||||
# Format src
|
||||
if inst.src0.offset == 255: src = _lit(inst, inst.src0)
|
||||
elif is16_src and inst.src0.offset >= 256: src = _fmt_v16(inst.src0)
|
||||
elif inst.src0.sz > 1: src = _fmt_src(inst.src0, inst.src0.sz, cdna)
|
||||
else: src = _lit(inst, inst.src0)
|
||||
return f"{name} {dst}, {src}"
|
||||
|
||||
_VOP2_CARRY_OUT = {'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'} # carry out only
|
||||
_VOP2_CARRY_INOUT = {'v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'} # carry in and out (CDNA)
|
||||
_VOP2_CARRY_INOUT_RDNA = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'} # carry in and out (RDNA)
|
||||
def _disasm_vop2(inst: VOP2) -> str:
|
||||
name, cdna = inst.op_name.lower(), _is_cdna(inst)
|
||||
if cdna: name = _CDNA_DISASM_ALIASES.get(name, name) # apply CDNA aliases
|
||||
suf = "" if cdna or name.endswith('_e32') or (not cdna and inst.op == VOP2Op.V_DOT2ACC_F32_F16_E32) else "_e32"
|
||||
lit = getattr(inst, '_literal', None)
|
||||
is16 = not cdna and inst.canonical_op_bits['d'] == 16
|
||||
# fmaak/madak: dst = src0 * vsrc1 + K, fmamk/madmk: dst = src0 * K + vsrc1
|
||||
if 'fmaak' in name or 'madak' in name or (not cdna and inst.op in (VOP2Op.V_FMAAK_F32_E32, VOP2Op.V_FMAAK_F16_E32)):
|
||||
if lit is None: return f"op_{inst.op.value if hasattr(inst.op, 'value') else inst.op}"
|
||||
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1)}, 0x{lit:x}"
|
||||
return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, 0x{lit:x}"
|
||||
if 'fmamk' in name or 'madmk' in name or (not cdna and inst.op in (VOP2Op.V_FMAMK_F32_E32, VOP2Op.V_FMAMK_F16_E32)):
|
||||
if lit is None: return f"op_{inst.op.value if hasattr(inst.op, 'value') else inst.op}"
|
||||
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, 0x{lit:x}, {_fmt_v16(inst.vsrc1)}"
|
||||
return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, 0x{lit:x}, {inst.vsrc1.fmt()}"
|
||||
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1)}"
|
||||
vcc = "vcc" if cdna else "vcc_lo"
|
||||
if cdna and name in _VOP2_CARRY_OUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}"
|
||||
if cdna and name in _VOP2_CARRY_INOUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
|
||||
if not cdna and name in _VOP2_CARRY_INOUT_RDNA: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
|
||||
sn0 = inst.canonical_op_regs.get('s0', 1)
|
||||
if inst.vdst.sz > 1 or sn0 > 1 or inst.vsrc1.sz > 1:
|
||||
src0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else _fmt_src(inst.src0, sn0, cdna)
|
||||
return f"{name.replace('_e32', '')} {inst.vdst.fmt()}, {src0}, {inst.vsrc1.fmt()}"
|
||||
return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}" + (f", {vcc}" if name == 'v_cndmask_b32' else "")
|
||||
|
||||
def _disasm_vopc(inst: VOPC) -> str:
|
||||
name, cdna = inst.op_name.lower(), _is_cdna(inst)
|
||||
bits = inst.canonical_op_bits
|
||||
is16 = bits['s0'] == 16
|
||||
if cdna:
|
||||
s0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else _fmt_src(inst.src0, inst.src0.sz, cdna)
|
||||
return f"{name} vcc, {s0}, {inst.vsrc1.fmt()}" # CDNA VOPC always outputs vcc
|
||||
# RDNA: v_cmpx_* writes to exec (no vcc), v_cmp_* writes to vcc_lo
|
||||
has_vcc = 'cmpx' not in name
|
||||
s0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else inst.src0.fmt() if inst.src0.sz > 1 else _src16(inst, inst.src0.offset) if is16 else _lit(inst, inst.src0)
|
||||
s1 = inst.vsrc1.fmt() if inst.vsrc1.sz > 1 else _fmt_v16(inst.vsrc1) if is16 else inst.vsrc1.fmt()
|
||||
suf = "" if name.endswith('_e32') else "_e32"
|
||||
return f"{name}{suf} vcc_lo, {s0}, {s1}" if has_vcc else f"{name}{suf} {s0}, {s1}"
|
||||
|
||||
NO_ARG_SOPP = {SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
|
||||
SOPPOp.S_WAIT_IDLE, SOPPOp.S_ENDPGM_SAVED, SOPPOp.S_CODE_END, SOPPOp.S_ENDPGM_ORDERED_PS_DONE, SOPPOp.S_TTRACEDATA}
|
||||
_CDNA_NO_ARG_SOPP = {'s_endpgm', 's_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_nop', 's_sethalt', 's_sleep',
|
||||
's_setprio', 's_trap', 's_incperflevel', 's_decperflevel', 's_sendmsg', 's_sendmsghalt'}
|
||||
|
||||
def _disasm_sopp(inst: SOPP) -> str:
|
||||
name, cdna = inst.op_name.lower(), _is_cdna(inst)
|
||||
is_rdna4 = 'rdna4' in inst.__class__.__module__
|
||||
# Ops that have no argument when simm16 == 0
|
||||
no_arg_zero = {'s_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_wait_idle', 's_endpgm_saved',
|
||||
's_endpgm_ordered_ps_done', 's_code_end'}
|
||||
if name in no_arg_zero: return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
|
||||
if name == 's_endpgm': return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
|
||||
if cdna:
|
||||
if name == 's_waitcnt':
|
||||
vm, lgkm, exp = inst.simm16 & 0xf, (inst.simm16 >> 8) & 0x3f, (inst.simm16 >> 4) & 0x7
|
||||
p = [f"vmcnt({vm})" if vm != 0xf else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
|
||||
return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
|
||||
if name.startswith(('s_cbranch', 's_branch')): return f"{name} {inst.simm16}"
|
||||
return f"{name} 0x{inst.simm16:x}" if inst.simm16 else name
|
||||
# RDNA (use name-based checks instead of enum-based for cross-arch compatibility)
|
||||
if name == 's_waitcnt':
|
||||
if is_rdna4:
|
||||
return f"{name} {inst.simm16}" if inst.simm16 else f"{name} 0"
|
||||
vm, exp, lgkm = (inst.simm16 >> 10) & 0x3f, inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x3f
|
||||
p = [f"vmcnt({vm})" if vm != 0x3f else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
|
||||
return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
|
||||
if name == 's_delay_alu':
|
||||
deps = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
|
||||
skips = ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
|
||||
id0, skip, id1 = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 7) & 0xf
|
||||
dep = lambda v: deps[v-1] if 0 < v <= len(deps) else str(v)
|
||||
p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
|
||||
return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
|
||||
if name.startswith(('s_cbranch', 's_branch')): return f"{name} 0x{inst.simm16:x}"
|
||||
return f"{name} 0x{inst.simm16:x}"
|
||||
|
||||
def _disasm_smem(inst: SMEM) -> str:
|
||||
name, cdna = inst.op_name.lower(), _is_cdna(inst)
|
||||
if name in ('s_gl1_inv', 's_dcache_inv'): return name
|
||||
soe, imm = getattr(inst, 'soe', 0), getattr(inst, 'imm', 1)
|
||||
is_rdna4 = 'rdna4' in inst.__class__.__module__
|
||||
offset = inst.ioffset if is_rdna4 else getattr(inst, 'offset', 0)
|
||||
if cdna:
|
||||
if soe and imm: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
|
||||
elif imm: off_s = f"0x{offset:x}"
|
||||
elif offset < 256: off_s = decode_src(offset, cdna)
|
||||
else: off_s = decode_src(inst.soffset, cdna)
|
||||
elif offset and inst.soffset != 124: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
|
||||
elif offset: off_s = f"0x{offset:x}"
|
||||
else: off_s = decode_src(inst.soffset, cdna)
|
||||
is_buffer = 'buffer' in name or 's_atc_probe_buffer' == name
|
||||
sbase_idx, sbase_count = _unwrap(inst.sbase), 4 if is_buffer else 2
|
||||
sbase_str = _fmt_src(sbase_idx, sbase_count, cdna) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
|
||||
if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {_unwrap(inst.sdata)}, {sbase_str}, {off_s}"
|
||||
if 'prefetch' in name:
|
||||
off = getattr(inst, 'ioffset', getattr(inst, 'offset', 0))
|
||||
if off >= 0x800000: off = off - 0x1000000
|
||||
off_s = f"0x{off:x}" if off > 255 else str(off)
|
||||
soff_s = decode_src(inst.soffset, cdna) if inst.soffset != 124 else "null"
|
||||
if 'pc_rel' in name: return f"{name} {off_s}, {soff_s}, {_unwrap(inst.sdata)}"
|
||||
return f"{name} {sbase_str}, {off_s}, {soff_s}, {_unwrap(inst.sdata)}"
|
||||
# Use get_field_bits for register count
|
||||
dst_n = inst.canonical_op_regs.get('d', 1)
|
||||
th, scope = getattr(inst, 'th', 0), getattr(inst, 'scope', 0)
|
||||
if is_rdna4: # RDNA4 uses th/scope instead of glc/dlc
|
||||
th_names = ['TH_LOAD_RT', 'TH_LOAD_NT', 'TH_LOAD_HT', 'TH_LOAD_LU']
|
||||
scope_names = ['SCOPE_CU', 'SCOPE_SE', 'SCOPE_DEV', 'SCOPE_SYS']
|
||||
mods = (f" th:{th_names[th]}" if th else "") + (f" scope:{scope_names[scope]}" if scope else "")
|
||||
return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}{mods}"
|
||||
if th or scope:
|
||||
th_names = ['TH_LOAD_RT', 'TH_LOAD_NT', 'TH_LOAD_HT', 'TH_LOAD_LU']
|
||||
scope_names = ['SCOPE_CU', 'SCOPE_SE', 'SCOPE_DEV', 'SCOPE_SYS']
|
||||
mods = (f" th:{th_names[th]}" if th else "") + (f" scope:{scope_names[scope]}" if scope else "")
|
||||
return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}{mods}"
|
||||
return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
|
||||
|
||||
def _disasm_flat(inst: FLAT) -> str:
|
||||
name, cdna = inst.op_name.lower(), _is_cdna(inst)
|
||||
acc = getattr(inst, 'acc', 0)
|
||||
reg_fn = _areg if acc else _vreg
|
||||
seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
|
||||
instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
|
||||
off_val = inst.offset if seg == 'flat' else (inst.offset if inst.offset < 4096 else inst.offset - 8192)
|
||||
# Use get_field_bits: data for stores/atomics, d for loads
|
||||
regs = inst.canonical_op_regs
|
||||
w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
|
||||
off_s = f" offset:{off_val}" if off_val else ""
|
||||
if cdna: mods = f"{off_s}{' glc' if inst.sc0 else ''}{' slc' if inst.nt else ''}"
|
||||
else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
|
||||
if seg == 'flat' or _unwrap(inst.saddr) == 0x7F: saddr_s = ""
|
||||
elif _unwrap(inst.saddr) == 124: saddr_s = ", off"
|
||||
elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr, cdna)}"
|
||||
elif _unwrap(inst.saddr) in (SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS): saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
|
||||
elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
|
||||
else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
|
||||
if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
|
||||
if cdna: addr_w = 1 if seg == 'scratch' else 2
|
||||
else: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
|
||||
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
|
||||
data_s, vdst_s = reg_fn(inst.data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
|
||||
glc_or_sc0 = inst.sc0 if cdna else inst.glc
|
||||
if 'atomic' in name:
|
||||
return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}"
|
||||
if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
|
||||
return f"{instr} {reg_fn(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
|
||||
|
||||
def _disasm_ds(inst: DS) -> str:
|
||||
op, name = inst.op, inst.op_name.lower()
|
||||
acc = getattr(inst, 'acc', 0)
|
||||
reg_fn = _areg if acc else _vreg
|
||||
gds = " gds" if getattr(inst, 'gds', 0) else ""
|
||||
off = f" offset:{inst.offset0 | (inst.offset1 << 8)}" if inst.offset0 or inst.offset1 else ""
|
||||
off2 = (" offset0:" + str(inst.offset0) if inst.offset0 else "") + (" offset1:" + str(inst.offset1) if inst.offset1 else "")
|
||||
# Use get_field_bits: data for stores/writes/atomics, d for loads
|
||||
regs = inst.canonical_op_regs
|
||||
w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'write' in name or ('load' not in name and 'read' not in name) else regs.get('d', 1)
|
||||
d0, d1, dst, addr = reg_fn(inst.data0, w), reg_fn(inst.data1, w), reg_fn(inst.vdst, w), _vreg(inst.addr)
|
||||
|
||||
if name == 'ds_nop': return name
|
||||
if name == 'ds_bvh_stack_rtn_b32': return f"{name} {_vreg(inst.vdst)}, {addr}, {_vreg(inst.data0)}, {_vreg(inst.data1, 4)}{off}{gds}"
|
||||
if 'bvh_stack_push' in name:
|
||||
d1_regs = 8 if 'push8' in name else 4
|
||||
vdst_regs = 2 if 'pop2' in name else 1
|
||||
vdst_s = _vreg(inst.vdst, vdst_regs) if vdst_regs > 1 else _vreg(inst.vdst)
|
||||
return f"{name} {vdst_s}, {addr}, {_vreg(inst.data0)}, {_vreg(inst.data1, d1_regs)}{off}{gds}"
|
||||
if 'gws_sema' in name and 'sema_br' not in name: return f"{name}{off}{gds}"
|
||||
if 'gws_' in name: return f"{name} {addr}{off}{gds}"
|
||||
if name in ('ds_consume', 'ds_append'): return f"{name} {reg_fn(inst.vdst)}{off}{gds}"
|
||||
if 'gs_reg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {reg_fn(inst.data0)}{off}{gds}"
|
||||
if '2addr' in name:
|
||||
if 'load' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}{off2}{gds}"
|
||||
if 'store' in name and 'xchg' not in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
|
||||
return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}, {d0}, {d1}{off2}{gds}"
|
||||
if 'write2' in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
|
||||
if 'read2' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}{off2}{gds}"
|
||||
if 'load' in name: return f"{name} {reg_fn(inst.vdst)}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
|
||||
if 'store' in name and not _has(name, 'cmp', 'xchg'):
|
||||
return f"{name} {reg_fn(inst.data0)}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
|
||||
if 'swizzle' in name or name == 'ds_ordered_count': return f"{name} {reg_fn(inst.vdst)}, {addr}{off}{gds}"
|
||||
if 'permute' in name: return f"{name} {reg_fn(inst.vdst)}, {addr}, {reg_fn(inst.data0)}{off}{gds}"
|
||||
if 'condxchg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {addr}, {reg_fn(inst.data0, 2)}{off}{gds}"
|
||||
if _has(name, 'cmpstore', 'mskor', 'wrap'):
|
||||
return f"{name} {dst}, {addr}, {d0}, {d1}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}, {d1}{off}{gds}"
|
||||
return f"{name} {dst}, {addr}, {d0}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}{off}{gds}"
|
||||
|
||||
def _disasm_vop3(inst: VOP3) -> str:
|
||||
op, name = inst.op, inst.op_name.lower()
|
||||
n_up = name.upper()
|
||||
bits = inst.canonical_op_bits
|
||||
|
||||
# RDNA4 v_s_* scalar VOP3 instructions - vdst is SGPR (VGPRField adds 256)
|
||||
if name.startswith('v_s_'):
|
||||
src = _lit(inst, inst.src0) if _unwrap(inst.src0) == 255 else ("src_scc" if _unwrap(inst.src0) == 253 else _fmt_src(inst.src0, max(1, bits['s0'] // 32)))
|
||||
if inst.neg & 1: src = f"-{src}"
|
||||
if inst.abs & 1: src = f"|{src}|"
|
||||
clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
|
||||
vdst_raw = _unwrap(inst.vdst)
|
||||
return f"{name} s{vdst_raw - 256 if vdst_raw >= 256 else vdst_raw}, {src}" + (" clamp" if clamp else "") + _omod(inst.omod)
|
||||
|
||||
# Use get_field_bits for register sizes and 16-bit detection
|
||||
r0, r1, r2 = max(1, bits['s0'] // 32), max(1, bits['s1'] // 32), max(1, bits['s2'] // 32)
|
||||
dn = max(1, bits['d'] // 32)
|
||||
is16_d, is16_s, is16_s2 = bits['d'] == 16, bits['s0'] == 16, bits['s2'] == 16
|
||||
|
||||
s0 = _vop3_src(inst, inst.src0, inst.neg&1, inst.abs&1, inst.opsel&1, r0, is16_s)
|
||||
s1 = _vop3_src(inst, inst.src1, inst.neg&2, inst.abs&2, inst.opsel&2, r1, is16_s)
|
||||
s2 = _vop3_src(inst, inst.src2, inst.neg&4, inst.abs&4, inst.opsel&4, r2, is16_s2)
|
||||
|
||||
# Format destination
|
||||
if 'readlane' in name:
|
||||
vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
|
||||
dst = _fmt_sdst(vdst_off, 1)
|
||||
elif is16_d: dst = f"{inst.vdst.fmt()}.h" if (inst.opsel & 8) else f"{inst.vdst.fmt()}.l"
|
||||
else: dst = inst.vdst.fmt()
|
||||
|
||||
clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
|
||||
cl, om = " clamp" if clamp else "", _omod(inst.omod)
|
||||
nonvgpr_opsel = (inst.src0.offset < 256 and (inst.opsel & 1)) or (inst.src1.offset < 256 and (inst.opsel & 2)) or (inst.src2.offset < 256 and (inst.opsel & 4))
|
||||
need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)
|
||||
|
||||
op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
|
||||
e64 = "" if name.endswith('_e64') else "_e64"
|
||||
if op_val < 256: # VOPC
|
||||
vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
|
||||
return f"{name}{e64} {s0}, {s1}{cl}" if name.startswith('v_cmpx') else f"{name}{e64} {_fmt_sdst(vdst_off, 1)}, {s0}, {s1}{cl}"
|
||||
if op_val < 384: # VOP2
|
||||
n = inst.num_srcs() or 2
|
||||
os = _opsel_str(inst.opsel, n, need_opsel, is16_d)
|
||||
return f"{name}{e64} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name}{e64} {dst}, {s0}, {s1}{os}{cl}{om}"
|
||||
if op_val < 512: # VOP1
|
||||
if re.match(r'v_cvt_f32_(bf|fp)8', name) and inst.opsel:
|
||||
os = f" byte_sel:{((inst.opsel & 1) << 1) | ((inst.opsel >> 1) & 1)}"
|
||||
else:
|
||||
os = _opsel_str(inst.opsel, 1, need_opsel, is16_d)
|
||||
if 'v_nop' in name or 'v_pipeflush' in name: return f"{name}{e64}"
|
||||
return f"{name}{e64} {dst}, {s0}{os}{cl}{om}"
|
||||
# Native VOP3
|
||||
n = inst.num_srcs() or 2
|
||||
os = f" byte_sel:{inst.opsel >> 2}" if 'cvt_sr' in name and inst.opsel else _opsel_str(inst.opsel, n, need_opsel, is16_d)
|
||||
return f"{name} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{os}{cl}{om}"
|
||||
|
||||
def _disasm_vop3sd(inst: VOP3SD) -> str:
|
||||
name = inst.op_name.lower()
|
||||
def src(reg, neg):
|
||||
s = _lit(inst, reg.offset) if reg.offset == 255 else ("src_scc" if reg.offset == 253 else (reg.fmt() if reg.sz > 1 else _lit(inst, reg.offset)))
|
||||
return f"neg({s})" if neg and reg.offset == 255 else (f"-{s}" if neg else s)
|
||||
s0, s1, s2 = src(inst.src0, inst.neg & 1), src(inst.src1, inst.neg & 2), src(inst.src2, inst.neg & 4)
|
||||
# VOP3SD: _co_ ops (add/sub) without _ci_ have only 2 sources, all others (mad, div_scale, _co_ci_) have 3 sources
|
||||
has_only_two_srcs = '_co_' in name and '_ci_' not in name and 'mad' not in name
|
||||
srcs = f"{s0}, {s1}" if has_only_two_srcs else f"{s0}, {s1}, {s2}"
|
||||
clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
|
||||
return f"{name} {inst.vdst.fmt()}, {_fmt_sdst(inst.sdst, 1)}, {srcs}{' clamp' if clamp else ''}{_omod(inst.omod)}"
|
||||
|
||||
def _disasm_vopd(inst: VOPD) -> str:
|
||||
lit = inst._literal or getattr(inst, 'literal', None)
|
||||
is_rdna4 = 'rdna4' in inst.__class__.__module__
|
||||
op_enum = R4_VOPDOp if is_rdna4 else VOPDOp
|
||||
vdst_y, nx, ny = (_unwrap(inst.vdsty) << 1) | ((_unwrap(inst.vdstx) & 1) ^ 1), op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
|
||||
def half(n, vd, s0, vs1):
|
||||
vd, vs1 = _vi(vd), _vi(vs1)
|
||||
if 'mov' in n: return f"{n} v{vd}, {_lit(inst, s0)}"
|
||||
if 'fmamk' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, 0x{lit:x}, v{vs1}"
|
||||
if 'fmaak' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}, 0x{lit:x}"
|
||||
return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}"
|
||||
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
|
||||
|
||||
def _disasm_vop3p(inst: VOP3P) -> str:
|
||||
name = inst.op_name.lower()
|
||||
is_wmma, is_swmmac, n, is_fma_mix = 'wmma' in name, 'swmmac' in name, inst.num_srcs() or 2, 'fma_mix' in name
|
||||
def get_src(reg):
|
||||
return _lit(inst, reg.offset) if reg.offset == 255 else reg.fmt()
|
||||
src0, src1, src2, dst = get_src(inst.src0), get_src(inst.src1), get_src(inst.src2), inst.vdst.fmt()
|
||||
opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
|
||||
clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
|
||||
if is_fma_mix:
|
||||
def m(s, neg, abs_): return f"-{f'|{s}|' if abs_ else s}" if neg else (f"|{s}|" if abs_ else s)
|
||||
src0, src1, src2 = m(src0, inst.neg & 1, inst.neg_hi & 1), m(src1, inst.neg & 2, inst.neg_hi & 2), m(src2, inst.neg & 4, inst.neg_hi & 4)
|
||||
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if clamp else [])
|
||||
elif is_swmmac:
|
||||
mods = ([f"index_key:{inst.opsel}"] if inst.opsel else []) + ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + \
|
||||
([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
|
||||
else:
|
||||
opsel_hi_default = 7 if n == 3 else 3
|
||||
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else []) + \
|
||||
([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
|
||||
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
|
||||
|
||||
def _disasm_sop1(inst: SOP1) -> str:
|
||||
op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
|
||||
# Use get_field_bits for register sizes
|
||||
regs = inst.canonical_op_regs
|
||||
dst_regs, src_regs = regs.get('d', 1), regs.get('s0', 1)
|
||||
src = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, src_regs, cdna)
|
||||
if not cdna:
|
||||
if 'getpc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2)}"
|
||||
if 'setpc_b64' in name or 'rfe_b64' in name: return f"{name} {src}"
|
||||
if 'swappc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2)}, {src}"
|
||||
if 'sendmsg_rtn' in name:
|
||||
v = _unwrap(inst.ssrc0)
|
||||
msg_str = MSG.get(v)
|
||||
return f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, sendmsg({msg_str})" if msg_str else f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, 0x{v:x}"
|
||||
sop1_src_only = ('S_ALLOC_VGPR', 'S_SLEEP_VAR', 'S_BARRIER_SIGNAL', 'S_BARRIER_SIGNAL_ISFIRST', 'S_BARRIER_INIT', 'S_BARRIER_JOIN')
|
||||
if inst.op_name in sop1_src_only: return f"{name} {src}"
|
||||
return f"{name} {_fmt_sdst(inst.sdst, dst_regs, cdna)}, {src}"
|
||||
|
||||
def _disasm_sop2(inst: SOP2) -> str:
|
||||
cdna, name = _is_cdna(inst), inst.op_name.lower()
|
||||
lit = getattr(inst, '_literal', None)
|
||||
# Use get_field_bits for register sizes
|
||||
regs = inst.canonical_op_regs
|
||||
dn, s0n, s1n = regs['d'], regs['s0'], regs['s1']
|
||||
s0 = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, s0n, cdna)
|
||||
s1 = _lit(inst, inst.ssrc1) if _unwrap(inst.ssrc1) == 255 else _fmt_src(inst.ssrc1, s1n, cdna)
|
||||
dst = _fmt_sdst(inst.sdst, dn, cdna)
|
||||
if 'fmamk' in name and lit is not None: return f"{name} {dst}, {s0}, 0x{lit:x}, {s1}"
|
||||
if 'fmaak' in name and lit is not None: return f"{name} {dst}, {s0}, {s1}, 0x{lit:x}"
|
||||
return f"{name} {dst}, {s0}, {s1}"
|
||||
|
||||
def _disasm_sopc(inst: SOPC) -> str:
|
||||
cdna, regs = _is_cdna(inst), inst.canonical_op_regs
|
||||
s0 = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, regs['s0'], cdna)
|
||||
s1 = _lit(inst, inst.ssrc1) if _unwrap(inst.ssrc1) == 255 else _fmt_src(inst.ssrc1, regs['s1'], cdna)
|
||||
return f"{inst.op_name.lower()} {s0}, {s1}"
|
||||
|
||||
def _disasm_sopk(inst: SOPK) -> str:
|
||||
op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
|
||||
is_rdna4 = 'rdna4' in inst.__class__.__module__
|
||||
hw = HWREG_RDNA4 if is_rdna4 else HWREG
|
||||
def fmt_hwreg(hid, hoff, hsz):
|
||||
if hid not in hw: return f"0x{inst.simm16:x}"
|
||||
hr_name = hw[hid]
|
||||
return f"hwreg({hr_name})" if hoff == 0 and hsz == 32 else f"hwreg({hr_name}, {hoff}, {hsz})"
|
||||
if name == 's_setreg_imm32_b32':
|
||||
hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
|
||||
return f"{name} {fmt_hwreg(hid, hoff, hsz)}, 0x{inst._literal:x}"
|
||||
if name == 's_version': return f"{name} 0x{inst.simm16:x}"
|
||||
if name in ('s_setreg_b32', 's_getreg_b32'):
|
||||
hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
|
||||
hs = fmt_hwreg(hid, hoff, hsz)
|
||||
return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1, cdna)}" if 'setreg' in name else f"{name} {_fmt_sdst(inst.sdst, 1, cdna)}, {hs}"
|
||||
if name in ('s_subvector_loop_begin', 's_subvector_loop_end'):
|
||||
return f"{name} {_fmt_sdst(inst.sdst, 1)}, 0x{inst.simm16:x}"
|
||||
return f"{name} {_fmt_sdst(inst.sdst, inst.canonical_op_regs['d'], cdna)}, 0x{inst.simm16:x}"
|
||||
|
||||
def _disasm_vinterp(inst: VINTERP) -> str:
|
||||
mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
|
||||
return f"{inst.op_name.lower()} {inst.vdst.fmt()}, {_lit(inst, inst.src0, inst.neg & 1)}, {_lit(inst, inst.src1, inst.neg & 2)}, {_lit(inst, inst.src2, inst.neg & 4)}" + (" " + mods if mods else "")
|
||||
|
||||
DISASM_HANDLERS: dict[type, callable] = {
|
||||
VOP1: _disasm_vop1, VOP1_SDST: _disasm_vop1, VOP2: _disasm_vop2, VOPC: _disasm_vopc, VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOPD: _disasm_vopd, VOP3P: _disasm_vop3p,
|
||||
VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, GLOBAL: _disasm_flat, SCRATCH: _disasm_flat,
|
||||
SOP1: _disasm_sop1, SOP2: _disasm_sop2, SOPC: _disasm_sopc, SOPK: _disasm_sopk,
|
||||
# RDNA4
|
||||
R4_VOP1: _disasm_vop1, R4_VOP1_SDST: _disasm_vop1, R4_VOP2: _disasm_vop2, R4_VOPC: _disasm_vopc, R4_VOP3: _disasm_vop3, R4_VOP3_SDST: _disasm_vop3, R4_VOP3SD: _disasm_vop3sd,
|
||||
R4_VOPD: _disasm_vopd, R4_VOP3P: _disasm_vop3p, R4_VINTERP: _disasm_vinterp, R4_SOPP: _disasm_sopp, R4_SMEM: _disasm_smem,
|
||||
R4_DS: _disasm_ds, R4_SOP1: _disasm_sop1, R4_SOP2: _disasm_sop2, R4_SOPC: _disasm_sopc, R4_SOPK: _disasm_sopk}
|
||||
|
||||
def disasm(inst: Inst) -> str: return DISASM_HANDLERS[type(inst)](inst)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# CDNA DISASSEMBLER SUPPORT
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
try:
|
||||
from extra.assembly.amd.autogen.cdna.ins import (VOP1 as CDNA_VOP1, VOP2 as CDNA_VOP2, VOPC as CDNA_VOPC, VOP3A, VOP3B, VOP3P as CDNA_VOP3P,
|
||||
SOP1 as CDNA_SOP1, SOP2 as CDNA_SOP2, SOPC as CDNA_SOPC, SOPK as CDNA_SOPK, SOPP as CDNA_SOPP, SMEM as CDNA_SMEM, DS as CDNA_DS,
|
||||
FLAT as CDNA_FLAT, VOP1Op as CDNA_VOP1Op, VOP2Op as CDNA_VOP2Op, VOPCOp as CDNA_VOPCOp)
|
||||
|
||||
def _cdna_src(inst, v, neg, abs_=0, n=1):
|
||||
s = _lit(inst, v) if v == 255 else _fmt_src(v, n, cdna=True)
|
||||
if abs_: s = f"|{s}|"
|
||||
return f"neg({s})" if neg and v == 255 else (f"-{s}" if neg else s)
|
||||
|
||||
_CDNA_VOP3_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32'}
|
||||
|
||||
def _disasm_vop3a(inst) -> str:
|
||||
op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
|
||||
name = inst.op_name.lower() or f'vop3a_op_{op_val}'
|
||||
n = inst.num_srcs() or _num_srcs(inst)
|
||||
cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
|
||||
orig_name = name
|
||||
name = _CDNA_VOP3_ALIASES.get(name, name)
|
||||
if name != orig_name:
|
||||
s0, s1 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, 1), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, 1)
|
||||
s2 = ""
|
||||
dst = _vreg(inst.vdst)
|
||||
else:
|
||||
regs = inst.canonical_op_regs
|
||||
dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
|
||||
s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, r0), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, r1), _cdna_src(inst, inst.src2, inst.neg&4, inst.abs&4, r2)
|
||||
dst = _vreg(inst.vdst, dregs) if dregs > 1 else _vreg(inst.vdst)
|
||||
if op_val >= 512:
|
||||
return f"{name} {dst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{cl}{om}"
|
||||
if op_val < 256:
|
||||
sdst = _fmt_sdst(inst.vdst, 2, cdna=True)
|
||||
return f"{name}_e64 {sdst}, {s0}, {s1}{cl}"
|
||||
if 320 <= op_val < 512:
|
||||
if name in ('v_nop', 'v_clrexcp'): return f"{name}_e64"
|
||||
return f"{name}_e64 {dst}, {s0}{cl}{om}"
|
||||
if name == 'v_cndmask_b32':
|
||||
s2 = _fmt_src(inst.src2, 2, cdna=True)
|
||||
return f"{name}_e64 {dst}, {s0}, {s1}, {s2}{cl}{om}"
|
||||
if name in ('v_mul_legacy_f32', 'v_mac_f32'):
|
||||
return f"{name}_e64 {dst}, {s0}, {s1}{cl}{om}"
|
||||
suf = "_e64" if op_val < 512 else ""
|
||||
return f"{name}{suf} {dst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name}{suf} {dst}, {s0}, {s1}{cl}{om}"
|
||||
|
||||
def _disasm_vop3b(inst) -> str:
|
||||
op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
|
||||
name = inst.op_name.lower() or f'vop3b_op_{op_val}'
|
||||
n = inst.num_srcs() or _num_srcs(inst)
|
||||
regs = inst.canonical_op_regs
|
||||
dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
|
||||
s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, n=r0), _cdna_src(inst, inst.src1, inst.neg&2, n=r1), _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
|
||||
dst = _vreg(inst.vdst, dregs) if dregs > 1 else _vreg(inst.vdst)
|
||||
sdst = _fmt_sdst(inst.sdst, 2, cdna=True)
|
||||
cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
|
||||
if name in ('v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'):
|
||||
s2 = _fmt_src(inst.src2, 2, cdna=True)
|
||||
return f"{name}_e64 {dst}, {sdst}, {s0}, {s1}, {s2}{cl}{om}"
|
||||
suf = "_e64" if 'co_' in name else ""
|
||||
return f"{name}{suf} {dst}, {sdst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name}{suf} {dst}, {sdst}, {s0}, {s1}{cl}{om}"
|
||||
|
||||
def _disasm_cdna_vop3p(inst) -> str:
|
||||
name, n = inst.op_name.lower(), inst.num_srcs() or 2
|
||||
is_mfma = 'mfma' in name or 'smfmac' in name
|
||||
is_accvgpr = 'accvgpr' in name
|
||||
get_src = lambda v, sc: _lit(inst, v) if v == 255 else _fmt_src(v, sc, cdna=True)
|
||||
|
||||
# Handle accvgpr read/write (accumulator register operations)
|
||||
if is_accvgpr:
|
||||
src0_off = _unwrap(inst.src0)
|
||||
vdst_off = _vi(inst.vdst)
|
||||
if 'read' in name:
|
||||
# v_accvgpr_read_b32 vN, aM - reads from accumulator to VGPR
|
||||
return f"{name}_b32 v{vdst_off}, a{src0_off - 256 if src0_off >= 256 else src0_off}"
|
||||
if 'write' in name:
|
||||
# v_accvgpr_write_b32 aM, src - writes to accumulator from source
|
||||
src = _lit(inst, inst.src0) if src0_off == 255 else (f"v{src0_off - 256}" if src0_off >= 256 else decode_src(src0_off, cdna=True))
|
||||
return f"{name}_b32 a{vdst_off}, {src}"
|
||||
|
||||
# Handle MFMA instructions with accumulator destinations
|
||||
if is_mfma:
|
||||
sc = 2 if 'iu4' in name else 4 if 'iu8' in name or 'i4' in name else 8 if 'f16' in name or 'bf16' in name else 4
|
||||
src0, src1, src2 = get_src(inst.src0, sc), get_src(inst.src1, sc), get_src(inst.src2, 16)
|
||||
dst = _areg(inst.vdst, 16) # MFMA uses accumulator registers
|
||||
opsel_hi = inst.opsel_hi
|
||||
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != 3 else []) + \
|
||||
([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
|
||||
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}"
|
||||
|
||||
# Standard VOP3P instructions
|
||||
src0, src1, src2, dst = get_src(inst.src0, 1), get_src(inst.src1, 1), get_src(inst.src2, 1), _vreg(inst.vdst)
|
||||
opsel_hi = inst.opsel_hi # CDNA VOP3P only has 2 bits for opsel_hi (no opsel_hi2)
|
||||
opsel_hi_default = 3 # CDNA default is 0b11 (2 bits), not 0b111 like RDNA
|
||||
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else []) + \
|
||||
([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
|
||||
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
|
||||
|
||||
DISASM_HANDLERS.update({CDNA_VOP1: _disasm_vop1, CDNA_VOP2: _disasm_vop2, CDNA_VOPC: _disasm_vopc,
|
||||
CDNA_SOP1: _disasm_sop1, CDNA_SOP2: _disasm_sop2, CDNA_SOPC: _disasm_sopc, CDNA_SOPK: _disasm_sopk, CDNA_SOPP: _disasm_sopp,
|
||||
CDNA_SMEM: _disasm_smem, CDNA_DS: _disasm_ds, CDNA_FLAT: _disasm_flat,
|
||||
VOP3A: _disasm_vop3a, VOP3B: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p})
|
||||
except ImportError:
|
||||
pass
|
||||
+340
-650
File diff suppressed because it is too large
Load Diff
+394
-664
File diff suppressed because it is too large
Load Diff
+498
-728
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,386 @@
|
||||
"""SQTT (SQ Thread Trace) packet encoder and decoder for AMD GPUs.
|
||||
|
||||
This module provides encoding and decoding of raw SQTT byte streams.
|
||||
The format is nibble-based with variable-width packets determined by a state machine.
|
||||
Uses BitField infrastructure from dsl.py, similar to GPU instruction encoding.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
from typing import Iterator
|
||||
from enum import Enum
|
||||
from extra.assembly.amd.dsl import BitField, FixedBitField, bits
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# FIELD ENUMS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class MemSrc(Enum):
|
||||
LDS = 0
|
||||
LDS_ALT = 1
|
||||
VMEM = 2
|
||||
VMEM_ALT = 3
|
||||
|
||||
class AluSrc(Enum):
|
||||
NONE = 0
|
||||
SALU = 1
|
||||
VALU = 2
|
||||
VALU_SALU = 3
|
||||
|
||||
class InstOp(Enum):
|
||||
"""SQTT instruction operation types.
|
||||
|
||||
Memory ops appear in two ranges depending on which SIMD executes them:
|
||||
- 0x1x-0x2x range: ops on traced SIMD
|
||||
- 0x5x range: ops on other SIMD (OTHER_ prefix)
|
||||
|
||||
GLOBAL memory ops encoding depends on addressing mode AND size:
|
||||
- Loads: 0x21 (saddr=SGPR) or 0x22 (saddr=NULL), all sizes same
|
||||
- Stores: base + size_offset, where VADDR is shifted +1 from SADDR
|
||||
SADDR: 0x24(32) 0x25(64) 0x26(96) 0x27(128)
|
||||
VADDR: 0x25(32) 0x26(64) 0x27(96) 0x28(128)
|
||||
|
||||
OTHER_ range follows same pattern but values overlap differently.
|
||||
"""
|
||||
SALU = 0x0
|
||||
SMEM = 0x1
|
||||
JUMP = 0x3 # branch taken
|
||||
JUMP_NO = 0x4 # branch not taken
|
||||
MESSAGE = 0x9
|
||||
VALU_TRANS = 0xb # transcendental: exp, log, rcp, sqrt, sin, cos
|
||||
VALU_64_SHIFT = 0xd # 64-bit shifts: lshl, lshr, ashr
|
||||
VALU_MAD64 = 0xe # 64-bit multiply-add
|
||||
VALU_64 = 0xf # 64-bit: add, mul, fma, rcp, sqrt, rounding, frexp, div helpers
|
||||
VINTERP = 0x12 # interpolation: v_interp_p10_f32, v_interp_p2_f32
|
||||
BARRIER = 0x13
|
||||
|
||||
# FLAT memory ops on traced SIMD (0x1x range)
|
||||
FLAT_LOAD = 0x1c
|
||||
FLAT_STORE = 0x1d
|
||||
FLAT_STORE_64 = 0x1e
|
||||
FLAT_STORE_96 = 0x1f
|
||||
FLAT_STORE_128 = 0x20
|
||||
|
||||
# GLOBAL memory ops on traced SIMD (0x2x range)
|
||||
GLOBAL_LOAD = 0x21 # saddr=SGPR, all sizes
|
||||
GLOBAL_LOAD_VADDR = 0x22 # saddr=NULL, all sizes
|
||||
GLOBAL_STORE = 0x24 # saddr=SGPR, 32-bit
|
||||
GLOBAL_STORE_64 = 0x25 # saddr=SGPR 64 or saddr=NULL 32
|
||||
GLOBAL_STORE_96 = 0x26 # saddr=SGPR 96 or saddr=NULL 64
|
||||
GLOBAL_STORE_128 = 0x27 # saddr=SGPR 128 or saddr=NULL 96
|
||||
GLOBAL_STORE_VADDR_128 = 0x28 # saddr=NULL, 128-bit
|
||||
|
||||
# LDS ops on traced SIMD
|
||||
LDS_LOAD = 0x29
|
||||
LDS_STORE = 0x2b
|
||||
LDS_STORE_64 = 0x2c
|
||||
LDS_STORE_128 = 0x2e
|
||||
|
||||
# Memory ops on other SIMD (0x5x range)
|
||||
OTHER_LDS_LOAD = 0x50
|
||||
OTHER_LDS_STORE = 0x51
|
||||
OTHER_LDS_STORE_64 = 0x52
|
||||
OTHER_LDS_STORE_128 = 0x54
|
||||
OTHER_FLAT_LOAD = 0x55
|
||||
OTHER_FLAT_STORE = 0x56
|
||||
OTHER_FLAT_STORE_64 = 0x57
|
||||
OTHER_FLAT_STORE_96 = 0x58
|
||||
OTHER_FLAT_STORE_128 = 0x59
|
||||
OTHER_GLOBAL_LOAD = 0x5a # saddr=SGPR, all sizes
|
||||
OTHER_GLOBAL_LOAD_VADDR = 0x5b # saddr=NULL or saddr=SGPR store 32
|
||||
OTHER_GLOBAL_STORE_64 = 0x5c # saddr=SGPR 64 or saddr=NULL 32
|
||||
OTHER_GLOBAL_STORE_96 = 0x5d # saddr=SGPR 96 or saddr=NULL 64
|
||||
OTHER_GLOBAL_STORE_128 = 0x5e # saddr=SGPR 128 or saddr=NULL 96
|
||||
OTHER_GLOBAL_STORE_VADDR_128 = 0x5f # saddr=NULL, 128-bit
|
||||
|
||||
# EXEC-modifying ops (0x7x range)
|
||||
SALU_SAVEEXEC = 0x72 # s_*_saveexec_b32/b64
|
||||
VALU_CMPX = 0x73 # v_cmpx_*
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PACKET TYPE BASE CLASS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class PacketType:
|
||||
"""Base class for SQTT packet types."""
|
||||
encoding: FixedBitField
|
||||
_raw: int
|
||||
_time: int
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
cls._fields = {k: v for k, v in cls.__dict__.items() if isinstance(v, BitField)}
|
||||
cls._size_nibbles = ((max((f.hi for f in cls._fields.values()), default=0) + 4) // 4)
|
||||
|
||||
@classmethod
|
||||
def from_raw(cls, raw: int, time: int = 0):
|
||||
inst = object.__new__(cls)
|
||||
inst._raw, inst._time = raw, time
|
||||
return inst
|
||||
|
||||
def __repr__(self) -> str:
|
||||
fields_str = ", ".join(f"{k}={getattr(self, k)}" for k in self._fields if not k.startswith('_'))
|
||||
return f"{self.__class__.__name__}({fields_str})"
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PACKET TYPE DEFINITIONS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class VALUINST(PacketType): # exclude: 1 << 2
|
||||
encoding = bits[2:0] == 0b011
|
||||
delta = bits[5:3]
|
||||
flag = bits[6:6]
|
||||
wave = bits[11:7]
|
||||
|
||||
class VMEMEXEC(PacketType): # exclude: 1 << 0
|
||||
encoding = bits[3:0] == 0b1111
|
||||
delta = bits[5:4]
|
||||
src = bits[7:6].enum(MemSrc)
|
||||
|
||||
class ALUEXEC(PacketType): # exclude: 1 << 1
|
||||
encoding = bits[3:0] == 0b1110
|
||||
delta = bits[5:4]
|
||||
src = bits[7:6].enum(AluSrc)
|
||||
|
||||
class IMMEDIATE(PacketType): # exclude: 1 << 5
|
||||
encoding = bits[3:0] == 0b1101
|
||||
delta = bits[6:4]
|
||||
wave = bits[11:7]
|
||||
|
||||
class IMMEDIATE_MASK(PacketType): # exclude: 1 << 5
|
||||
encoding = bits[4:0] == 0b00100
|
||||
delta = bits[7:5]
|
||||
mask = bits[23:8]
|
||||
|
||||
class WAVERDY(PacketType): # exclude: 1 << 3
|
||||
encoding = bits[4:0] == 0b10100
|
||||
delta = bits[7:5]
|
||||
mask = bits[23:8]
|
||||
|
||||
class TS_DELTA_S8_W3(PacketType):
|
||||
encoding = bits[6:0] == 0b0100001
|
||||
delta = bits[10:8]
|
||||
_padding = bits[63:11]
|
||||
|
||||
class WAVEEND(PacketType): # exclude: 1 << 4
|
||||
encoding = bits[4:0] == 0b10101
|
||||
delta = bits[7:5]
|
||||
flag7 = bits[8:8]
|
||||
simd = bits[10:9]
|
||||
cu_lo = bits[13:11]
|
||||
wave = bits[19:15]
|
||||
@property
|
||||
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
|
||||
|
||||
class WAVESTART(PacketType): # exclude: 1 << 4
|
||||
encoding = bits[4:0] == 0b01100
|
||||
delta = bits[6:5]
|
||||
flag7 = bits[7:7]
|
||||
simd = bits[9:8]
|
||||
cu_lo = bits[12:10]
|
||||
wave = bits[17:13]
|
||||
id7 = bits[31:18]
|
||||
@property
|
||||
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
|
||||
|
||||
class TS_DELTA_S5_W2(PacketType):
|
||||
encoding = bits[4:0] == 0b11100
|
||||
delta = bits[6:5]
|
||||
_padding = bits[47:7]
|
||||
|
||||
class WAVEALLOC(PacketType): # exclude: 1 << 10
|
||||
encoding = bits[4:0] == 0b00101
|
||||
delta = bits[7:5]
|
||||
_padding = bits[19:8]
|
||||
|
||||
class TS_DELTA_S5_W3(PacketType):
|
||||
encoding = bits[4:0] == 0b00110
|
||||
delta = bits[7:5]
|
||||
_padding = bits[51:8]
|
||||
|
||||
class PERF(PacketType): # exclude: 1 << 11
|
||||
encoding = bits[4:0] == 0b10110
|
||||
delta = bits[7:5]
|
||||
arg = bits[27:8]
|
||||
|
||||
class TS_DELTA_SHORT(PacketType):
|
||||
encoding = bits[3:0] == 0b1000
|
||||
delta = bits[7:4]
|
||||
|
||||
class NOP(PacketType):
|
||||
encoding = bits[3:0] == 0b0000
|
||||
delta = None # type: ignore
|
||||
_padding = bits[3:0]
|
||||
|
||||
class TS_WAVE_STATE(PacketType):
|
||||
encoding = bits[6:0] == 0b1010001
|
||||
delta = bits[15:7]
|
||||
coarse = bits[23:16]
|
||||
@property
|
||||
def wave_interest(self) -> bool: return bool(self.coarse & 1)
|
||||
@property
|
||||
def terminate_all(self) -> bool: return bool(self.coarse & 8)
|
||||
|
||||
class EVENT(PacketType): # exclude: 1 << 7
|
||||
encoding = bits[7:0] == 0b01100001
|
||||
delta = bits[10:8]
|
||||
event = bits[23:11]
|
||||
|
||||
class EVENT_BIG(PacketType):
|
||||
encoding = bits[7:0] == 0b11100001
|
||||
delta = bits[10:8]
|
||||
event = bits[31:11]
|
||||
|
||||
class REG(PacketType):
|
||||
encoding = bits[3:0] == 0b1001
|
||||
delta = bits[6:4]
|
||||
slot = bits[9:7]
|
||||
hi_byte = bits[15:8]
|
||||
subop = bits[31:16]
|
||||
val32 = bits[63:32]
|
||||
@property
|
||||
def is_config(self) -> bool: return bool(self.hi_byte & 0x80)
|
||||
|
||||
class SNAPSHOT(PacketType):
|
||||
encoding = bits[6:0] == 0b1110001
|
||||
delta = bits[9:7]
|
||||
snap = bits[63:10]
|
||||
|
||||
class TS_DELTA_OR_MARK(PacketType):
|
||||
encoding = bits[6:0] == 0b0000001
|
||||
delta = bits[47:12]
|
||||
bit8 = bits[8:8]
|
||||
bit9 = bits[9:9]
|
||||
@property
|
||||
def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
|
||||
|
||||
class LAYOUT_HEADER(PacketType):
|
||||
encoding = bits[6:0] == 0b0010001
|
||||
delta = None # type: ignore
|
||||
layout = bits[12:7]
|
||||
simd = bits[14:13]
|
||||
group = bits[17:15]
|
||||
sel_a = bits[31:28]
|
||||
sel_b = bits[36:33]
|
||||
flag4 = bits[59:59]
|
||||
_padding = bits[63:60]
|
||||
|
||||
class INST(PacketType):
|
||||
encoding = bits[2:0] == 0b010
|
||||
delta = bits[6:4]
|
||||
flag1 = bits[3:3]
|
||||
flag2 = bits[7:7]
|
||||
wave = bits[12:8]
|
||||
op = bits[19:13].enum(InstOp)
|
||||
|
||||
class UTILCTR(PacketType):
|
||||
encoding = bits[6:0] == 0b0110001
|
||||
delta = bits[8:7]
|
||||
ctr = bits[47:9]
|
||||
|
||||
# All packet types in encoding priority order (more specific masks first, NOP last as fallback)
|
||||
PACKET_TYPES: list[type[PacketType]] = [
|
||||
EVENT, EVENT_BIG,
|
||||
TS_DELTA_S8_W3, TS_WAVE_STATE, SNAPSHOT, TS_DELTA_OR_MARK, LAYOUT_HEADER, UTILCTR,
|
||||
IMMEDIATE_MASK, WAVERDY, WAVEEND, WAVESTART, TS_DELTA_S5_W2, WAVEALLOC, TS_DELTA_S5_W3, PERF,
|
||||
VMEMEXEC, ALUEXEC, IMMEDIATE, TS_DELTA_SHORT, REG,
|
||||
VALUINST, INST,
|
||||
NOP,
|
||||
]
|
||||
|
||||
def _build_state_table() -> tuple[bytes, dict[int, type[PacketType]]]:
|
||||
table = [len(PACKET_TYPES) - 1] * 256 # default to NOP
|
||||
opcode_to_class: dict[int, type[PacketType]] = {i: cls for i, cls in enumerate(PACKET_TYPES)}
|
||||
|
||||
for byte_val in range(256):
|
||||
for opcode, pkt_cls in enumerate(PACKET_TYPES):
|
||||
if (byte_val & pkt_cls.encoding.mask) == pkt_cls.encoding.default:
|
||||
table[byte_val] = opcode
|
||||
break
|
||||
|
||||
return bytes(table), opcode_to_class
|
||||
|
||||
STATE_TO_OPCODE, OPCODE_TO_CLASS = _build_state_table()
|
||||
|
||||
# Precompute special case opcodes
|
||||
_TS_DELTA_OR_MARK_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_OR_MARK)
|
||||
_TS_DELTA_SHORT_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_SHORT)
|
||||
|
||||
# Combined lookup: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
|
||||
# special_case: 0=none, 1=TS_DELTA_OR_MARK, 2=TS_DELTA_SHORT
|
||||
_DECODE_INFO: dict[int, tuple] = {}
|
||||
for _opcode, _pkt_cls in OPCODE_TO_CLASS.items():
|
||||
_delta_field = getattr(_pkt_cls, 'delta', None)
|
||||
_delta_lo = _delta_field.lo if _delta_field else 0
|
||||
_delta_mask = _delta_field.mask if _delta_field else 0
|
||||
_special = 1 if _opcode == _TS_DELTA_OR_MARK_OPCODE else (2 if _opcode == _TS_DELTA_SHORT_OPCODE else 0)
|
||||
_DECODE_INFO[_opcode] = (_pkt_cls, _pkt_cls._size_nibbles, _delta_lo, _delta_mask, _special)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# DECODER
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def decode(data: bytes) -> Iterator[PacketType]:
|
||||
"""Decode raw SQTT blob, yielding packet instances."""
|
||||
n, reg, pos, nib_off, nib_count, time = len(data), 0, 0, 0, 16, 0
|
||||
|
||||
while pos + ((nib_count + nib_off + 1) >> 1) <= n:
|
||||
need = nib_count - nib_off
|
||||
# 1. if unaligned, read high nibble to align
|
||||
if nib_off: reg, pos = (reg >> 4) | ((data[pos] >> 4) << 60), pos + 1
|
||||
# 2. read all full bytes at once
|
||||
if (byte_count := need >> 1):
|
||||
chunk = int.from_bytes(data[pos:pos + byte_count], 'little')
|
||||
reg, pos = (reg >> (byte_count * 8)) | (chunk << (64 - byte_count * 8)), pos + byte_count
|
||||
# 3. if odd, read low nibble
|
||||
if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)
|
||||
|
||||
opcode = STATE_TO_OPCODE[reg & 0xFF]
|
||||
pkt_cls, nib_count, delta_lo, delta_mask, special = _DECODE_INFO[opcode]
|
||||
delta = (reg >> delta_lo) & delta_mask
|
||||
if special == 1 and (reg >> 9) & 1 and not (reg >> 8) & 1: delta = 0 # TS_DELTA_OR_MARK marker
|
||||
elif special == 2: delta += 8 # TS_DELTA_SHORT
|
||||
time += delta
|
||||
yield pkt_cls.from_raw(reg, time)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PRINTER
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
PACKET_COLORS = {
|
||||
"INST": "WHITE", "VALUINST": "BLACK", "VMEMEXEC": "yellow", "ALUEXEC": "yellow",
|
||||
"IMMEDIATE": "YELLOW", "IMMEDIATE_MASK": "YELLOW", "WAVERDY": "cyan", "WAVEALLOC": "cyan",
|
||||
"WAVEEND": "blue", "WAVESTART": "blue", "PERF": "magenta", "EVENT": "red", "EVENT_BIG": "red",
|
||||
"REG": "green", "LAYOUT_HEADER": "white", "SNAPSHOT": "white", "UTILCTR": "green",
|
||||
}
|
||||
|
||||
def format_packet(p) -> str:
|
||||
from tinygrad.helpers import colored
|
||||
name = type(p).__name__
|
||||
if isinstance(p, INST):
|
||||
op_name = p.op.name if isinstance(p.op, InstOp) else f"0x{p.op:02x}"
|
||||
fields = f"wave={p.wave} op={op_name}" + (" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "")
|
||||
elif isinstance(p, VALUINST): fields = f"wave={p.wave}" + (" flag" if p.flag else "")
|
||||
elif isinstance(p, ALUEXEC): fields = f"src={p.src.name if isinstance(p.src, AluSrc) else p.src}"
|
||||
elif isinstance(p, VMEMEXEC): fields = f"src={p.src.name if isinstance(p.src, MemSrc) else p.src}"
|
||||
elif isinstance(p, (WAVESTART, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
|
||||
elif hasattr(p, '_fields'):
|
||||
fields = " ".join(f"{k}=0x{getattr(p, k):x}" if k in {'snap', 'val32'} else f"{k}={getattr(p, k)}"
|
||||
for k in p._fields if not k.startswith('_') and k not in {'delta', 'encoding'})
|
||||
else: fields = ""
|
||||
return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name, 'white'))} {fields}"
|
||||
|
||||
def print_packets(packets) -> None:
|
||||
skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"}
|
||||
for p in packets:
|
||||
if type(p).__name__ not in skip: print(format_packet(p))
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys, pickle
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python sqtt.py <pkl_file>")
|
||||
sys.exit(1)
|
||||
with open(sys.argv[1], "rb") as f:
|
||||
data = pickle.load(f)
|
||||
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
|
||||
for i, event in enumerate(sqtt_events):
|
||||
print(f"\n=== event {i} ===")
|
||||
print_packets(decode(event.blob))
|
||||
@@ -1,13 +1,12 @@
|
||||
#!/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
|
||||
"""Benchmark comparing Python vs Rust RDNA3 emulators on real tinygrad kernels."""
|
||||
import ctypes, time, os
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
# 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, step_wave, WaveState, WAVE_SIZE
|
||||
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():
|
||||
@@ -42,7 +41,7 @@ def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = Non
|
||||
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):
|
||||
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
|
||||
@@ -50,13 +49,13 @@ def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size
|
||||
lib_ptr = ctypes.addressof(kernel_buf)
|
||||
|
||||
# Warmup
|
||||
run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
|
||||
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)
|
||||
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}")
|
||||
@@ -65,27 +64,12 @@ def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size
|
||||
|
||||
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)."""
|
||||
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)
|
||||
|
||||
@@ -112,7 +96,9 @@ def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], d
|
||||
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]
|
||||
@@ -122,67 +108,22 @@ def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], d
|
||||
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)
|
||||
# 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
|
||||
|
||||
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()
|
||||
|
||||
@@ -197,98 +138,55 @@ def main():
|
||||
|
||||
results = []
|
||||
|
||||
# Synthetic workloads
|
||||
if not args.tinygrad_only:
|
||||
print("\n[SYNTHETIC WORKLOADS]")
|
||||
print("-" * 90)
|
||||
print("\n[TINYGRAD KERNELS]")
|
||||
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
|
||||
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
|
||||
|
||||
print(f"\n{name}: {n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
|
||||
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
|
||||
|
||||
buf_sizes = [4096]
|
||||
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes)
|
||||
set_valid_mem_ranges(ranges)
|
||||
print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
|
||||
|
||||
# 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
|
||||
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
|
||||
set_valid_mem_ranges(ranges)
|
||||
|
||||
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]")
|
||||
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
|
||||
|
||||
results.append(("synthetic", name, n_insts, n_workgroups, py_time, rust_time))
|
||||
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]")
|
||||
|
||||
# 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}")
|
||||
results.append((op_name, n_insts, n_workgroups, py_time, rust_time))
|
||||
|
||||
# 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(f"{'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:
|
||||
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"{test_type:<10} {name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
|
||||
|
||||
|
||||
print(f"{name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
# Currently many of these tests fail - they document desired behavior
|
||||
|
||||
import unittest
|
||||
from extra.assembly.amd.autogen.rdna3 import *
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
from extra.assembly.amd.dsl import Inst, RawImm, SGPR, VGPR
|
||||
|
||||
class TestRegisterSliceSyntax(unittest.TestCase):
|
||||
|
||||
@@ -22,3 +22,45 @@ def get_llvm_objdump():
|
||||
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,221 @@
|
||||
"""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.emu import WaveState, run_asm, set_valid_mem_ranges
|
||||
from extra.assembly.amd.dsl import NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, M0
|
||||
from extra.assembly.amd.pcode import _i32, _f32
|
||||
|
||||
# For backwards compatibility with tests using SrcEnum.NULL etc.
|
||||
class SrcEnum:
|
||||
NULL = NULL
|
||||
VCC_LO = VCC_LO
|
||||
VCC_HI = VCC_HI
|
||||
EXEC_LO = EXEC_LO
|
||||
EXEC_HI = EXEC_HI
|
||||
SCC = SCC
|
||||
M0 = M0
|
||||
POS_HALF = 0.5
|
||||
NEG_HALF = -0.5
|
||||
POS_ONE = 1.0
|
||||
NEG_ONE = -1.0
|
||||
POS_TWO = 2.0
|
||||
NEG_TWO = -2.0
|
||||
POS_FOUR = 4.0
|
||||
NEG_FOUR = -4.0
|
||||
|
||||
VCC = VCC_LO # For VOP3SD sdst field (VCC_LO is exported from dsl)
|
||||
USE_HW = os.environ.get("USE_HW", "0") == "1"
|
||||
FLOAT_TOLERANCE = 1e-5
|
||||
|
||||
# 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(VCC_LO, 0))
|
||||
|
||||
epilogue = [
|
||||
s_mov_b32(s[90], VCC_LO),
|
||||
s_cselect_b32(s[91], 1, 0),
|
||||
s_load_b64(s[92:93], s[80:81], 0, soffset=NULL),
|
||||
s_waitcnt(0), # simm16=0 waits for all
|
||||
v_lshlrev_b32_e32(v[240], 2, v[255]),
|
||||
]
|
||||
for i in range(N_VGPRS):
|
||||
epilogue.append(global_store_b32(addr=v[240], data=v[i], saddr=s[92:93], offset=i * WAVE_SIZE * 4))
|
||||
epilogue.append(v_mov_b32_e32(v[241], 0))
|
||||
epilogue.append(v_cmp_eq_u32_e32(v[255], v[241]))
|
||||
epilogue.append(s_and_saveexec_b32(s[94], VCC_LO))
|
||||
epilogue.append(v_mov_b32_e32(v[240], 0))
|
||||
for i in range(N_SGPRS):
|
||||
epilogue.append(v_mov_b32_e32(v[243], s[i]))
|
||||
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + i * 4))
|
||||
epilogue.append(v_mov_b32_e32(v[243], s[90]))
|
||||
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + SGPR_BYTES))
|
||||
epilogue.append(v_mov_b32_e32(v[243], s[91]))
|
||||
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=VGPR_BYTES + SGPR_BYTES + 4))
|
||||
epilogue.append(s_mov_b32(EXEC_LO, s[94]))
|
||||
epilogue.append(s_endpgm())
|
||||
return prologue, epilogue
|
||||
|
||||
def parse_output(out_buf: bytes, n_lanes: int) -> WaveState:
|
||||
"""Parse output buffer into WaveState."""
|
||||
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:3], 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:1], data1=v[2:3], vdst=v[0], offset0=0, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4:7], 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:3], 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:7], 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:3], 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:1], 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:3], 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:3], 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:1], data1=v[2:3], vdst=v[0], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[4:7], 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:5], offset0=0, offset1=1),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[6:7], 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:1], 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:3], vdst=v[4:5], offset0=0),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_B64, addr=v[10], vdst=v[6:7], 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:1], data1=v[0:1], vdst=v[0], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[2:5], 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:5], offset0=1, offset1=2),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[6:7], 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:1], data1=v[0:1], 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:7], data1=v[6:7], vdst=v[8:11], 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:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
] + 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], 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:1], data=v[2:3], 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:1], data=v[4:5], vdst=v[6:7], 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:1], data=v[2:3], 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:1], data=v[4:5], vdst=v[6:7], 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:1], data=v[2:3], 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:1], data=v[4:5], vdst=v[6:7], 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:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
s_mov_b32(s[0], 0xDEADBEEF),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
FLAT(FLATOp.FLAT_LOAD_B32, addr=v[0:1], 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:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
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:1], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
FLAT(FLATOp.FLAT_LOAD_B64, addr=v[0:1], vdst=v[4: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][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:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
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:1], data=v[2:4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
FLAT(FLATOp.FLAT_LOAD_B96, addr=v[0:1], vdst=v[5:7], 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:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
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:1], data=v[2:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
FLAT(FLATOp.FLAT_LOAD_B128, addr=v[0:1], vdst=v[6:9], 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:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
] + 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:1], 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 = GLOBAL(GLOBALOp.GLOBAL_ATOMIC_ADD_U32, addr=v[0:1], 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_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:1], data=v[2:3], 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 = GLOBAL(GLOBALOp.GLOBAL_ATOMIC_ADD_U64, addr=v[0:1], data=v[4:5], vdst=v[6:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
|
||||
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:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
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:1], data=v[2:4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B96, addr=v[0:1], vdst=v[5:7], 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], 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:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
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:1], data=v[2:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B128, addr=v[0:1], vdst=v[6:9], 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], 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:81], 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:3], saddr=s[2:3], offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B64, addr=v[0], vdst=v[4:5], data=v[4:5], saddr=s[2:3], offset=TEST_OFFSET),
|
||||
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:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
s_mov_b32(s[4], 0xCAFE),
|
||||
v_mov_b32_e32(v[2], s[4]),
|
||||
global_store_b16(addr=v[0:1], 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]),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[0:1], vdst=v[3], data=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
|
||||
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:81], 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:3], 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
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[0], saddr=s[2:3], offset=TEST_OFFSET),
|
||||
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:81], 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:3], 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),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_U8, addr=v[3], vdst=v[5], data=v[4], saddr=s[2:3], offset=TEST_OFFSET),
|
||||
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:81], 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:3], offset=TEST_OFFSET),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[1], saddr=s[2:3], offset=TEST_OFFSET),
|
||||
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:81], 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:3], offset=TEST_OFFSET),
|
||||
global_store_b32(addr=v[3], data=v[10], saddr=s[2:3], 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
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[0], saddr=s[2:3], offset=TEST_OFFSET+6),
|
||||
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:81], 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:3], 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),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_I8, addr=v[3], vdst=v[5], data=v[4], saddr=s[2:3], offset=TEST_OFFSET),
|
||||
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:81], 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:11], saddr=s[2:3], offset=TEST_OFFSET),
|
||||
global_store_b8(addr=v[0], data=v[12], saddr=s[2:3], offset=TEST_OFFSET+8),
|
||||
s_waitcnt(vmcnt=0),
|
||||
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_U16, addr=v[2], vdst=v[0], data=v[0], saddr=s[2:3], offset=TEST_OFFSET+3),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[1], saddr=s[2:3], offset=TEST_OFFSET+6),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_U8, addr=v[2], vdst=v[3], data=v[3], saddr=s[2:3], offset=TEST_OFFSET),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_U8, addr=v[2], vdst=v[4], data=v[4], saddr=s[2:3], offset=TEST_OFFSET+8),
|
||||
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:1], saddr=s[2:3], offset=TEST_OFFSET+16),
|
||||
s_waitcnt(vmcnt=0),
|
||||
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B64, addr=v[2], vdst=v[6:7], data=v[6:7], saddr=s[2:3], offset=TEST_OFFSET+16),
|
||||
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,381 @@
|
||||
"""Tests for SOP instructions - scalar operations.
|
||||
|
||||
Includes: s_add_u32, s_mov_b32, s_and_b32, s_or_b32, s_quadmask_b32, s_wqm_b32,
|
||||
s_cbranch_vccnz, s_cbranch_vccz
|
||||
"""
|
||||
import unittest
|
||||
from extra.assembly.amd.test.hw.helpers import *
|
||||
|
||||
class TestBasicScalar(unittest.TestCase):
|
||||
"""Tests for basic scalar operations."""
|
||||
|
||||
def test_s_add_u32(self):
|
||||
"""S_ADD_U32 adds two scalar values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 100),
|
||||
s_mov_b32(s[1], 200),
|
||||
s_add_u32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 300)
|
||||
|
||||
def test_s_add_u32_carry(self):
|
||||
"""S_ADD_U32 sets SCC on overflow."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 64),
|
||||
s_not_b32(s[0], s[0]), # ~64 = 0xffffffbf
|
||||
s_mov_b32(s[1], 64),
|
||||
s_add_u32(s[2], s[0], s[1]), # 0xffffffbf + 64 = 0xffffffff
|
||||
s_mov_b32(s[3], 1),
|
||||
s_add_u32(s[4], s[2], s[3]), # 0xffffffff + 1 = overflow
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[4], 0)
|
||||
self.assertEqual(st.scc, 1)
|
||||
|
||||
def test_s_brev_b32(self):
|
||||
"""S_BREV_B32 reverses bits of a 32-bit value."""
|
||||
# 10 = 0b00000000000000000000000000001010
|
||||
# reversed = 0b01010000000000000000000000000000 = 0x50000000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 10),
|
||||
s_brev_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0x50000000)
|
||||
|
||||
def test_s_brev_b32_all_ones(self):
|
||||
"""S_BREV_B32 with all ones stays all ones."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFFFFFFFF),
|
||||
s_brev_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
|
||||
|
||||
def test_s_brev_b32_single_bit(self):
|
||||
"""S_BREV_B32 with bit 0 set becomes bit 31."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 1),
|
||||
s_brev_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0x80000000)
|
||||
|
||||
|
||||
class TestQuadmaskWqm(unittest.TestCase):
|
||||
"""Tests for S_QUADMASK_B32 and S_WQM_B32."""
|
||||
|
||||
def test_s_quadmask_b32_all_quads_active(self):
|
||||
"""S_QUADMASK_B32 with all quads active."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFFFFFFFF), # All lanes active
|
||||
s_quadmask_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Each quad (4 lanes) with any bit set -> 1 bit in result
|
||||
# 32 lanes = 8 quads, all active -> 0xFF
|
||||
self.assertEqual(st.sgpr[1], 0xFF)
|
||||
|
||||
def test_s_quadmask_b32_alternating_quads(self):
|
||||
"""S_QUADMASK_B32 with alternating quads active."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x0F0F0F0F), # Quads 0,2,4,6 active
|
||||
s_quadmask_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Quads 0,2,4,6 have at least one bit -> 0b01010101 = 0x55
|
||||
self.assertEqual(st.sgpr[1], 0x55)
|
||||
|
||||
def test_s_quadmask_b32_no_quads_active(self):
|
||||
"""S_QUADMASK_B32 with no quads active."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0),
|
||||
s_quadmask_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0)
|
||||
|
||||
def test_s_quadmask_b32_single_lane_per_quad(self):
|
||||
"""S_QUADMASK_B32 with single lane active in each quad."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x11111111), # Bit 0 of each nibble
|
||||
s_quadmask_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# All 8 quads have at least one lane -> 0xFF
|
||||
self.assertEqual(st.sgpr[1], 0xFF)
|
||||
|
||||
def test_s_wqm_b32_all_active(self):
|
||||
"""S_WQM_B32 with all lanes active returns all 1s."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFFFFFFFF),
|
||||
s_wqm_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
|
||||
|
||||
def test_s_wqm_b32_alternating_quads(self):
|
||||
"""S_WQM_B32 with single lane per quad expands to full quads."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x11111111), # One lane per quad
|
||||
s_wqm_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Each quad with any bit expands to all 4 bits
|
||||
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
|
||||
|
||||
def test_s_wqm_b32_zero(self):
|
||||
"""S_WQM_B32 with zero input returns zero."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0),
|
||||
s_wqm_b32(s[1], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0)
|
||||
|
||||
|
||||
class TestBranch(unittest.TestCase):
|
||||
"""Tests for branch instructions."""
|
||||
|
||||
def test_cbranch_vccnz_ignores_vcc_hi(self):
|
||||
"""S_CBRANCH_VCCNZ should only check VCC_LO in wave32."""
|
||||
instructions = [
|
||||
# Set VCC_LO = 0, VCC_HI = 1
|
||||
s_mov_b32(VCC_LO, 0),
|
||||
s_mov_b32(VCC_HI, 1),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
# If VCC_HI is incorrectly used, branch will be taken
|
||||
s_cbranch_vccnz(1), # Skip next instruction if VCC != 0
|
||||
v_mov_b32_e32(v[0], 42), # This should execute
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 42, "Branch should NOT be taken (VCC_LO is 0)")
|
||||
|
||||
def test_cbranch_vccz_ignores_vcc_hi(self):
|
||||
"""S_CBRANCH_VCCZ should only check VCC_LO in wave32."""
|
||||
instructions = [
|
||||
# Set VCC_LO = 1, VCC_HI = 0
|
||||
s_mov_b32(VCC_LO, 1),
|
||||
s_mov_b32(VCC_HI, 0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
# If VCC_HI is incorrectly used, branch will be taken
|
||||
s_cbranch_vccz(1), # Skip next instruction if VCC == 0
|
||||
v_mov_b32_e32(v[0], 42), # This should execute
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 42, "Branch should NOT be taken (VCC_LO is 1)")
|
||||
|
||||
def test_cbranch_vccnz_branches_on_vcc_lo(self):
|
||||
"""S_CBRANCH_VCCNZ branches when VCC_LO is non-zero."""
|
||||
instructions = [
|
||||
s_mov_b32(VCC_LO, 1),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_cbranch_vccnz(1), # Skip next instruction if VCC != 0
|
||||
v_mov_b32_e32(v[0], 42), # This should be skipped
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0, "Branch should be taken (VCC_LO is 1)")
|
||||
|
||||
|
||||
class Test64BitLiterals(unittest.TestCase):
|
||||
"""Tests for 64-bit literal encoding in instructions."""
|
||||
|
||||
def test_64bit_literal_negative_encoding(self):
|
||||
"""64-bit literal -2^32 encodes correctly."""
|
||||
lit = -4294967296.0 # -2^32
|
||||
lit_bits = f2i64(lit)
|
||||
instructions = [
|
||||
s_mov_b32(s[0], lit_bits & 0xffffffff),
|
||||
s_mov_b32(s[1], lit_bits >> 32),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = i642f(st.vgpr[0][0] | (st.vgpr[0][1] << 32))
|
||||
self.assertAlmostEqual(result, -4294967296.0, places=5)
|
||||
|
||||
class TestSCCBehavior(unittest.TestCase):
|
||||
"""Tests for SCC condition code behavior."""
|
||||
|
||||
def test_scc_from_s_cmp(self):
|
||||
"""SCC should be set by scalar compare."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 10),
|
||||
s_cmp_eq_u32(s[0], 10),
|
||||
s_cselect_b32(s[1], 1, 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 1, "SCC should be true")
|
||||
self.assertEqual(st.scc, 1)
|
||||
|
||||
def test_scc_clear(self):
|
||||
"""SCC should be cleared by failing compare."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 10),
|
||||
s_cmp_eq_u32(s[0], 20),
|
||||
s_cselect_b32(s[1], 1, 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0, "SCC should be false")
|
||||
self.assertEqual(st.scc, 0)
|
||||
|
||||
|
||||
class TestSignedArithmetic(unittest.TestCase):
|
||||
"""Tests for S_ADD_I32, S_SUB_I32 and their SCC overflow behavior."""
|
||||
|
||||
def test_s_add_i32_no_overflow(self):
|
||||
"""S_ADD_I32: 1 + 1 = 2, no overflow, SCC=0."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 1),
|
||||
s_add_i32(s[1], s[0], 1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 2)
|
||||
self.assertEqual(st.scc, 0, "No overflow, SCC should be 0")
|
||||
|
||||
def test_s_add_i32_positive_overflow(self):
|
||||
"""S_ADD_I32: MAX_INT + 1 overflows, SCC=1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7FFFFFFF), # MAX_INT
|
||||
s_add_i32(s[1], s[0], 1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[1], 0x80000000) # Wraps to MIN_INT
|
||||
self.assertEqual(st.scc, 1, "Overflow, SCC should be 1")
|
||||
|
||||
def test_s_add_i32_negative_no_overflow(self):
|
||||
"""S_ADD_I32: -10 + 20 = 10, no overflow."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFFFFFFF6), # -10 in two's complement
|
||||
s_mov_b32(s[1], 20),
|
||||
s_add_i32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 10)
|
||||
self.assertEqual(st.scc, 0)
|
||||
|
||||
def test_s_add_i32_negative_overflow(self):
|
||||
"""S_ADD_I32: MIN_INT + (-1) underflows, SCC=1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x80000000), # MIN_INT
|
||||
s_mov_b32(s[1], 0xFFFFFFFF), # -1
|
||||
s_add_i32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 0x7FFFFFFF) # Wraps to MAX_INT
|
||||
self.assertEqual(st.scc, 1, "Underflow, SCC should be 1")
|
||||
|
||||
def test_s_sub_i32_no_overflow(self):
|
||||
"""S_SUB_I32: 10 - 5 = 5, no overflow."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 10),
|
||||
s_mov_b32(s[1], 5),
|
||||
s_sub_i32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 5)
|
||||
self.assertEqual(st.scc, 0)
|
||||
|
||||
def test_s_sub_i32_overflow(self):
|
||||
"""S_SUB_I32: MAX_INT - (-1) overflows, SCC=1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7FFFFFFF), # MAX_INT
|
||||
s_mov_b32(s[1], 0xFFFFFFFF), # -1
|
||||
s_sub_i32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 0x80000000) # Wraps to MIN_INT
|
||||
self.assertEqual(st.scc, 1, "Overflow, SCC should be 1")
|
||||
|
||||
def test_s_mul_hi_u32(self):
|
||||
"""S_MUL_HI_U32: high 32 bits of u32 * u32."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x80000000), # 2^31
|
||||
s_mov_b32(s[1], 4),
|
||||
s_mul_hi_u32(s[2], s[0], s[1]), # (2^31 * 4) >> 32 = 2
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 2)
|
||||
|
||||
def test_s_mul_i32(self):
|
||||
"""S_MUL_I32: signed multiply low 32 bits."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFFFFFFFF), # -1
|
||||
s_mov_b32(s[1], 10),
|
||||
s_mul_i32(s[2], s[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[2], 0xFFFFFFF6) # -10
|
||||
|
||||
def test_division_sequence_from_llvm(self):
|
||||
"""Test the division sequence pattern from LLVM-generated code."""
|
||||
# This sequence is from the sin kernel and computes integer division
|
||||
# s10 = dividend, s18 = divisor, result in s6/s14
|
||||
dividend = 0x28BE60DB # Some value from the sin kernel
|
||||
divisor = 3 # Simplified divisor
|
||||
instructions = [
|
||||
s_mov_b32(s[10], dividend),
|
||||
s_mov_b32(s[18], divisor),
|
||||
# Compute reciprocal approximation: s6 = ~0 / divisor (approx)
|
||||
s_mov_b32(s[11], 0),
|
||||
s_sub_i32(s[11], s[11], s[18]), # s11 = -divisor
|
||||
# For testing, just verify basic arithmetic works
|
||||
s_mul_i32(s[6], s[10], 2),
|
||||
s_add_i32(s[7], s[6], 1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[6], (dividend * 2) & 0xFFFFFFFF)
|
||||
self.assertEqual(st.sgpr[7], ((dividend * 2) + 1) & 0xFFFFFFFF)
|
||||
|
||||
|
||||
class Test64BitCompare(unittest.TestCase):
|
||||
"""Tests for 64-bit scalar compare instructions."""
|
||||
|
||||
def test_s_cmp_eq_u64_equal(self):
|
||||
"""S_CMP_EQ_U64: comparing equal 64-bit values sets SCC=1."""
|
||||
val = 0x123456789ABCDEF0
|
||||
instructions = [
|
||||
s_mov_b32(s[0], val & 0xFFFFFFFF),
|
||||
s_mov_b32(s[1], val >> 32),
|
||||
s_mov_b32(s[2], val & 0xFFFFFFFF),
|
||||
s_mov_b32(s[3], val >> 32),
|
||||
s_cmp_eq_u64(s[0:1], s[2:3]),
|
||||
s_cselect_b32(s[4], 1, 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.scc, 1)
|
||||
self.assertEqual(st.sgpr[4], 1)
|
||||
|
||||
def test_s_cmp_eq_u64_different_upper_bits(self):
|
||||
"""S_CMP_EQ_U64: values differing only in upper 32 bits are not equal."""
|
||||
# This is the bug case - if only lower 32 bits are compared, these would be equal
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0), # lower 32 bits of value 0
|
||||
s_mov_b32(s[1], 0), # upper 32 bits of value 0
|
||||
s_mov_b32(s[2], 0), # lower 32 bits of 0x100000000
|
||||
s_mov_b32(s[3], 1), # upper 32 bits of 0x100000000
|
||||
s_cmp_eq_u64(s[0:1], s[2:3]),
|
||||
s_cselect_b32(s[4], 1, 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.scc, 0, "0 != 0x100000000, SCC should be 0")
|
||||
self.assertEqual(st.sgpr[4], 0)
|
||||
|
||||
def test_s_cmp_lg_u64_different(self):
|
||||
"""S_CMP_LG_U64: different 64-bit values sets SCC=1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0),
|
||||
s_mov_b32(s[1], 0), # s[0:1] = 0
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 1), # s[2:3] = 0x100000000
|
||||
s_cmp_lg_u64(s[0:1], s[2:3]),
|
||||
s_cselect_b32(s[4], 1, 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.scc, 1, "0 != 0x100000000, SCC should be 1")
|
||||
self.assertEqual(st.sgpr[4], 1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
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], literal=0x3f800000),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
|
||||
|
||||
def test_v_fmamk_f32_basic(self):
|
||||
"""V_FMAMK_F32: d = a * K + b."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 2.0),
|
||||
v_mov_b32_e32(v[1], 1.0),
|
||||
v_fmamk_f32_e32(v[2], v[0], v[1], literal=0x40800000),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
|
||||
|
||||
def test_v_fmamk_f32_small_constant(self):
|
||||
"""V_FMAMK_F32 with small constant."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 4.0),
|
||||
v_mov_b32_e32(v[1], 1.0),
|
||||
v_fmamk_f32_e32(v[2], v[0], v[1], literal=f2i(0.5)),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 3.0, places=5)
|
||||
|
||||
|
||||
class TestBitManipulation(unittest.TestCase):
|
||||
"""Tests for bit manipulation VOP2 instructions."""
|
||||
|
||||
def test_v_and_b32(self):
|
||||
"""V_AND_B32 bitwise and."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xff),
|
||||
s_mov_b32(s[1], 0x0f),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_and_b32_e32(v[1], s[1], v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0x0f)
|
||||
|
||||
def test_v_and_b32_quadrant(self):
|
||||
"""V_AND_B32 for quadrant extraction (n & 3)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 15915),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_and_b32_e32(v[1], 3, v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 15915 & 3)
|
||||
|
||||
def test_v_lshrrev_b32(self):
|
||||
"""V_LSHRREV_B32 logical shift right."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xff00),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_lshrrev_b32_e32(v[1], 8, v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0xff)
|
||||
|
||||
def test_v_lshlrev_b32(self):
|
||||
"""V_LSHLREV_B32 logical shift left."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xff),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_lshlrev_b32_e32(v[1], 8, v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0xff00)
|
||||
|
||||
def test_v_xor_b32(self):
|
||||
"""V_XOR_B32 bitwise xor (used in sin for sign)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x80000000),
|
||||
s_mov_b32(s[1], f2i(1.0)),
|
||||
v_mov_b32_e32(v[0], s[1]),
|
||||
v_xor_b32_e32(v[1], s[0], v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][1]), -1.0, places=5)
|
||||
|
||||
def test_v_xor_b32_sign_flip(self):
|
||||
"""V_XOR_B32 for sign flip pattern."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x80000000),
|
||||
v_mov_b32_e32(v[0], -2.0),
|
||||
v_xor_b32_e32(v[1], s[0], v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 2.0, places=5)
|
||||
|
||||
|
||||
class TestSpecialValues(unittest.TestCase):
|
||||
"""Tests for special float values - inf, nan, zero handling."""
|
||||
|
||||
def test_v_mul_f32_zero_times_inf(self):
|
||||
"""V_MUL_F32: 0 * inf = NaN."""
|
||||
import math
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
s_mov_b32(s[0], 0x7f800000),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
v_mul_f32_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
|
||||
|
||||
def test_v_add_f32_inf_minus_inf(self):
|
||||
"""V_ADD_F32: inf + (-inf) = NaN."""
|
||||
import math
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7f800000),
|
||||
s_mov_b32(s[1], 0xff800000),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_add_f32_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
|
||||
|
||||
|
||||
class TestF16Ops(unittest.TestCase):
|
||||
"""Tests for 16-bit VOP2 operations."""
|
||||
|
||||
def test_v_add_f16_basic(self):
|
||||
"""V_ADD_F16 adds two f16 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_add_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x4200, f"Expected 0x4200 (f16 3.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_add_f16_negative(self):
|
||||
"""V_ADD_F16 with negative values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0xc000), # f16 -2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_add_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0xbc00, f"Expected 0xbc00 (f16 -1.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_mul_f16_basic(self):
|
||||
"""V_MUL_F16 multiplies two f16 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x4000), # f16 2.0
|
||||
s_mov_b32(s[1], 0x4200), # f16 3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mul_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x4600, f"Expected 0x4600 (f16 6.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_mul_f16_by_zero(self):
|
||||
"""V_MUL_F16 by zero."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
v_mul_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x0000, f"Expected 0x0000 (f16 0.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_fmac_f16_basic(self):
|
||||
"""V_FMAC_F16: d = d + a*b."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x4000), # f16 2.0
|
||||
s_mov_b32(s[1], 0x4200), # f16 3.0
|
||||
s_mov_b32(s[2], 0x3c00), # f16 1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
v_fmac_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
|
||||
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_fmaak_f16_basic(self):
|
||||
"""V_FMAAK_F16: d = a * b + K."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x4000), # f16 2.0
|
||||
s_mov_b32(s[1], 0x4200), # f16 3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_fmaak_f16_e32(v[2], v[0], v[1], literal=0x3c00), # + f16 1.0
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
|
||||
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
|
||||
|
||||
|
||||
class TestHiHalfOps(unittest.TestCase):
|
||||
"""Tests for VOP2 16-bit operations with hi-half operands."""
|
||||
|
||||
def test_v_add_f16_src0_hi_fold(self):
|
||||
"""V_ADD_F16 with src0 hi-half fold (same register, different halves)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # lo=f16(1.0), hi=f16(2.0)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
VOP3(VOP3Op.V_ADD_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0001),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][1] & 0xffff
|
||||
self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
|
||||
|
||||
def test_v_add_f16_src0_hi_different_reg(self):
|
||||
"""V_ADD_F16 with src0 hi-half from different register."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40000000), # hi=f16(2.0), lo=0
|
||||
s_mov_b32(s[1], 0x00003c00), # hi=0, lo=f16(1.0)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
VOP3(VOP3Op.V_ADD_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
|
||||
|
||||
def test_v_mul_f16_src0_hi(self):
|
||||
"""V_MUL_F16 with src0 from high half."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40000000), # hi=f16(2.0), lo=0
|
||||
s_mov_b32(s[1], 0x00004200), # hi=0, lo=f16(3.0)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
VOP3(VOP3Op.V_MUL_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x4600, f"Expected f16(6.0)=0x4600, got 0x{result:04x}")
|
||||
|
||||
def test_v_mul_f16_hi_half(self):
|
||||
"""V_MUL_F16 reading from high half."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # lo=1.0, hi=2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
VOP3(VOP3Op.V_MUL_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0011),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][1] & 0xffff
|
||||
self.assertEqual(result, 0x4400, f"Expected f16(4.0)=0x4400, got 0x{result:04x}")
|
||||
|
||||
def test_v_fma_f16_hi_dest(self):
|
||||
"""V_FMA_F16 writing to high half with opsel.
|
||||
|
||||
Uses V_FMA_F16 (not V_FMAC_F16) because it has explicit src2 operand
|
||||
which makes opsel handling clearer.
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c000000), # hi=f16(1.0), lo=0
|
||||
s_mov_b32(s[1], 0x4000), # f16(2.0) in lo
|
||||
s_mov_b32(s[2], 0x4200), # f16(3.0) in lo
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
# V_FMA_F16: dst = src0 * src1 + src2
|
||||
# opsel=0b1100: bit2=src2 hi, bit3=dst hi
|
||||
# So: v[0].hi = v[1].lo * v[2].lo + v[0].hi = 2.0 * 3.0 + 1.0 = 7.0
|
||||
VOP3(VOP3Op.V_FMA_F16, vdst=v[0], src0=v[1], src1=v[2], src2=v[0], opsel=0b1100),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
hi = (st.vgpr[0][0] >> 16) & 0xffff
|
||||
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
|
||||
self.assertEqual(hi, 0x4700, f"Expected f16(7.0)=0x4700 in hi, got 0x{hi:04x}")
|
||||
|
||||
def test_v_add_f16_multilane(self):
|
||||
"""V_ADD_F16 with multiple lanes."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_add_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=4)
|
||||
for lane in range(4):
|
||||
result = st.vgpr[lane][2] & 0xffff
|
||||
self.assertEqual(result, 0x4200, f"Lane {lane}: expected 0x4200, got 0x{result:04x}")
|
||||
|
||||
|
||||
class TestCndmask(unittest.TestCase):
|
||||
"""Tests for V_CNDMASK_B32 and V_CNDMASK_B16."""
|
||||
|
||||
def test_v_cndmask_b16_select_src0(self):
|
||||
"""V_CNDMASK_B16 selects src0 when VCC bit is 0."""
|
||||
instructions = [
|
||||
s_mov_b32(VCC_LO, 0), # VCC = 0
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cndmask_b16(v[2], v[0], v[1], VCC),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x3c00, f"Expected src0=0x3c00, got 0x{result:04x}")
|
||||
|
||||
def test_v_cndmask_b16_select_src1(self):
|
||||
"""V_CNDMASK_B16 selects src1 when VCC bit is 1."""
|
||||
instructions = [
|
||||
s_mov_b32(VCC_LO, 1), # VCC = 1
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cndmask_b16(v[2], v[0], v[1], VCC),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x4000, f"Expected src1=0x4000, got 0x{result:04x}")
|
||||
|
||||
def test_v_cndmask_b16_write_hi(self):
|
||||
"""V_CNDMASK_B16 can write to high 16 bits with opsel."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c003800), # src0: hi=1.0, lo=0.5
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], 0x4000c000), # src1: hi=2.0, lo=-2.0
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], 0xDEAD0000), # v2 initial: hi=0xDEAD, lo=0
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
s_mov_b32(VCC_LO, 0), # vcc = 0, select src0
|
||||
# opsel=0b1011: bit0=src0 hi, bit1=src1 hi, bit3=dst hi
|
||||
VOP3(VOP3Op.V_CNDMASK_B16, vdst=v[2], src0=v[0], src1=v[1], src2=SrcEnum.VCC_LO, opsel=0b1011),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
hi = (st.vgpr[0][2] >> 16) & 0xffff
|
||||
lo = st.vgpr[0][2] & 0xffff
|
||||
# vcc=0 selects src0.h = 1.0 = 0x3c00, writes to hi
|
||||
self.assertEqual(hi, 0x3c00, f"Expected hi=0x3c00 (1.0), got 0x{hi:04x}")
|
||||
self.assertEqual(lo, 0x0000, f"Expected lo preserved as 0, got 0x{lo:04x}")
|
||||
|
||||
|
||||
class TestSpecialFloatValues(unittest.TestCase):
|
||||
"""Tests for special float value handling in VOP2 instructions."""
|
||||
|
||||
def test_neg_zero_add(self):
|
||||
"""-0.0 + 0.0 = +0.0 (IEEE 754)."""
|
||||
neg_zero = 0x80000000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], neg_zero),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_add_f32_e32(v[1], 0.0, v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
|
||||
|
||||
def test_neg_zero_mul(self):
|
||||
"""-0.0 * -1.0 = +0.0."""
|
||||
neg_zero = 0x80000000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], neg_zero),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mul_f32_e32(v[1], -1.0, v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
|
||||
|
||||
def test_inf_minus_inf(self):
|
||||
"""+inf - inf = NaN."""
|
||||
import math
|
||||
pos_inf = 0x7f800000
|
||||
neg_inf = 0xff800000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], pos_inf),
|
||||
s_mov_b32(s[1], neg_inf),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_sub_f32_e32(v[2], v[0], v[1]), # inf - (-inf) = inf
|
||||
v_add_f32_e32(v[3], v[0], v[1]), # inf + (-inf) = NaN
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], pos_inf, "inf - (-inf) = inf")
|
||||
self.assertTrue(math.isnan(i2f(st.vgpr[0][3])), "inf + (-inf) = NaN")
|
||||
|
||||
def test_denormal_f32_mul_ftz(self):
|
||||
"""Denormal * normal - RDNA3 flushes denormals to zero (FTZ mode)."""
|
||||
smallest_denorm = 0x00000001 # Smallest positive denormal
|
||||
instructions = [
|
||||
s_mov_b32(s[0], smallest_denorm),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mul_f32_e32(v[1], 2.0, v[0]), # Denormal input gets flushed to 0
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][1], 0x00000000)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,532 @@
|
||||
"""Tests for VOP3P instructions - packed 16-bit vector operations.
|
||||
|
||||
Includes: v_pk_add_f16, v_pk_mul_f16, v_pk_fma_f16, v_pack_b32_f16, v_wmma_*, v_dot2_*
|
||||
"""
|
||||
import unittest
|
||||
from extra.assembly.amd.test.hw.helpers import *
|
||||
|
||||
class TestPackInstructions(unittest.TestCase):
|
||||
"""Tests for pack instructions."""
|
||||
|
||||
def test_v_pack_b32_f16(self):
|
||||
"""V_PACK_B32_F16 packs two f16 values into one 32-bit register."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pack_b32_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x40003c00, f"Expected 0x40003c00, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_opsel_hi_hi(self):
|
||||
"""V_PACK_B32_F16 with opsel to read high halves."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
|
||||
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pack_b32_f16(v[2], v[0], v[1], opsel=0b0011),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x44004000, f"Expected 0x44004000, got 0x{result:08x}")
|
||||
|
||||
|
||||
class TestPackMore(unittest.TestCase):
|
||||
"""Additional pack instruction tests."""
|
||||
|
||||
def test_v_pack_b32_f16_basic(self):
|
||||
"""V_PACK_B32_F16 packs two f16 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pack_b32_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x40003c00, f"Expected 0x40003c00, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_with_cvt(self):
|
||||
"""V_PACK_B32_F16 after V_CVT_F16_F32 conversions."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3f800000),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
v_cvt_f16_f32_e32(v[2], v[0]),
|
||||
v_cvt_f16_f32_e32(v[3], v[1]),
|
||||
v_pack_b32_f16(v[4], v[2], v[3]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][4]
|
||||
self.assertEqual(result, 0x3c003c00, f"Expected 0x3c003c00, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_packed_sources(self):
|
||||
"""V_PACK_B32_F16 with packed f16 sources (reads lo halves)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
|
||||
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pack_b32_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
# Expected: hi=v1.lo=0x4200 (3.0), lo=v0.lo=0x3c00 (1.0) -> 0x42003c00
|
||||
self.assertEqual(result, 0x42003c00, f"Expected 0x42003c00, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_opsel_lo_hi(self):
|
||||
"""V_PACK_B32_F16 with opsel=0b0010 to read lo from src0, hi from src1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00),
|
||||
s_mov_b32(s[1], 0x44004200),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pack_b32_f16(v[2], v[0], v[1], opsel=0b0010),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x44003c00, f"Expected 0x44003c00, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_opsel_hi_lo(self):
|
||||
"""V_PACK_B32_F16 with opsel=0b0001 to read hi from src0, lo from src1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00),
|
||||
s_mov_b32(s[1], 0x44004200),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pack_b32_f16(v[2], v[0], v[1], opsel=0b0001),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x42004000, f"Expected 0x42004000, got 0x{result:08x}")
|
||||
|
||||
def test_v_pack_b32_f16_zeros(self):
|
||||
"""V_PACK_B32_F16 with zero values."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
v_mov_b32_e32(v[1], 0),
|
||||
v_pack_b32_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0)
|
||||
|
||||
def test_v_pack_b32_f16_both_positive(self):
|
||||
"""V_PACK_B32_F16 with positive f16 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x4200), # f16 3.0
|
||||
s_mov_b32(s[1], 0x4400), # f16 4.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pack_b32_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x44004200, f"Expected 0x44004200, got 0x{result:08x}")
|
||||
|
||||
|
||||
class TestFmaMix(unittest.TestCase):
|
||||
"""Tests for V_FMA_MIX_F32 and V_FMA_MIXLO_F16."""
|
||||
|
||||
def test_v_fma_mix_f32_all_f32_sources(self):
|
||||
"""V_FMA_MIX_F32 with all f32 sources."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(2.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(3.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], f2i(1.0)),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = i2f(st.vgpr[0][3])
|
||||
self.assertAlmostEqual(result, 7.0, places=5)
|
||||
|
||||
def test_v_fma_mix_f32_src2_f16_lo(self):
|
||||
"""V_FMA_MIX_F32 with src2 as f16 from lo bits."""
|
||||
from extra.assembly.amd.pcode import f32_to_f16
|
||||
f16_2 = f32_to_f16(2.0)
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(1.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(3.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], f16_2),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = i2f(st.vgpr[0][3])
|
||||
self.assertAlmostEqual(result, 5.0, places=5)
|
||||
|
||||
def test_v_fma_mix_f32_src2_f16_hi(self):
|
||||
"""V_FMA_MIX_F32 with src2 as f16 from hi bits."""
|
||||
from extra.assembly.amd.pcode import f32_to_f16
|
||||
f16_2 = f32_to_f16(2.0)
|
||||
val = (f16_2 << 16) | 0
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(1.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(3.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], val),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=4, opsel_hi=0, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = i2f(st.vgpr[0][3])
|
||||
self.assertAlmostEqual(result, 5.0, places=5)
|
||||
|
||||
def test_v_fma_mix_f32_with_abs(self):
|
||||
"""V_FMA_MIX_F32 with abs modifier on src2."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(2.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(3.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], f2i(-1.0)),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0, neg_hi=4),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = i2f(st.vgpr[0][3])
|
||||
self.assertAlmostEqual(result, 7.0, places=5)
|
||||
|
||||
def test_v_fma_mixlo_f16(self):
|
||||
"""V_FMA_MIXLO_F16 writes to low 16 bits of destination."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(2.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(3.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], f2i(1.0)),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
s_mov_b32(s[3], 0xdead0000),
|
||||
v_mov_b32_e32(v[3], s[3]),
|
||||
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
lo = _f16(st.vgpr[0][3] & 0xffff)
|
||||
hi = (st.vgpr[0][3] >> 16) & 0xffff
|
||||
self.assertAlmostEqual(lo, 7.0, places=1)
|
||||
self.assertEqual(hi, 0xdead, f"hi should be preserved, got 0x{hi:04x}")
|
||||
|
||||
def test_v_fma_mixlo_f16_all_f32_sources(self):
|
||||
"""V_FMA_MIXLO_F16 with all f32 sources."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], f2i(1.0)),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], f2i(2.0)),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
s_mov_b32(s[2], f2i(3.0)),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
v_mov_b32_e32(v[3], 0),
|
||||
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
lo = _f16(st.vgpr[0][3] & 0xffff)
|
||||
# 1*2+3 = 5
|
||||
self.assertAlmostEqual(lo, 5.0, places=1)
|
||||
|
||||
def test_v_fma_mixlo_f16_sin_case(self):
|
||||
"""V_FMA_MIXLO_F16 case from sin kernel."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3f800000), # f32 1.0
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
s_mov_b32(s[1], 0xaf05a309), # f32 tiny negative
|
||||
s_mov_b32(s[6], s[1]),
|
||||
s_mov_b32(s[2], 0xc0490fdb), # f32 -π
|
||||
v_mov_b32_e32(v[5], s[2]),
|
||||
s_mov_b32(s[3], 0x3f800000),
|
||||
v_mov_b32_e32(v[3], s[3]),
|
||||
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[3], src1=s[6], src2=v[5], opsel=0, opsel_hi=0, opsel_hi2=0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
lo = _f16(st.vgpr[0][3] & 0xffff)
|
||||
self.assertAlmostEqual(lo, -3.14159, delta=0.01)
|
||||
|
||||
|
||||
class TestVOP3P(unittest.TestCase):
|
||||
"""Tests for VOP3P packed 16-bit operations."""
|
||||
|
||||
def test_v_pk_add_f16_basic(self):
|
||||
"""V_PK_ADD_F16 adds two packed f16 values."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
|
||||
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pk_add_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
self.assertAlmostEqual(lo, 4.0, places=2)
|
||||
self.assertAlmostEqual(hi, 6.0, places=2)
|
||||
|
||||
def test_v_pk_mul_f16_basic(self):
|
||||
"""V_PK_MUL_F16 multiplies two packed f16 values."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x42004000), # hi=3.0, lo=2.0
|
||||
s_mov_b32(s[1], 0x45004400), # hi=5.0, lo=4.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pk_mul_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
self.assertAlmostEqual(lo, 8.0, places=1)
|
||||
self.assertAlmostEqual(hi, 15.0, places=1)
|
||||
|
||||
def test_v_pk_fma_f16_basic(self):
|
||||
"""V_PK_FMA_F16: D = A * B + C for packed f16."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x42004000), # A: hi=3.0, lo=2.0
|
||||
s_mov_b32(s[1], 0x45004400), # B: hi=5.0, lo=4.0
|
||||
s_mov_b32(s[2], 0x3c003c00), # C: hi=1.0, lo=1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], s[2]),
|
||||
v_pk_fma_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][3]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
self.assertAlmostEqual(lo, 9.0, places=1) # 2*4+1
|
||||
self.assertAlmostEqual(hi, 16.0, places=0) # 3*5+1
|
||||
|
||||
def test_v_pk_add_f16_with_inline_constant(self):
|
||||
"""V_PK_ADD_F16 with inline constant POS_ONE (1.0).
|
||||
Inline constants for VOP3P are f16 values in the low 16 bits only.
|
||||
hi half of inline constant is 0, so hi result = v0.hi + 0 = 1.0.
|
||||
"""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c003c00), # packed f16: hi=1.0, lo=1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_pk_add_f16(v[1], v[0], SrcEnum.POS_ONE, opsel_hi=3, opsel_hi2=1), # Add inline constant 1.0
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][1]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
# lo = 1.0 + 1.0 = 2.0, hi = 1.0 + 0.0 = 1.0 (inline const hi half is 0)
|
||||
self.assertAlmostEqual(lo, 2.0, places=2)
|
||||
self.assertAlmostEqual(hi, 1.0, places=2)
|
||||
|
||||
def test_v_pk_mul_f16_with_inline_constant(self):
|
||||
"""V_PK_MUL_F16 with inline constant POS_TWO (2.0).
|
||||
Inline constant has value only in low 16 bits, hi is 0.
|
||||
"""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
# v0 = packed (3.0, 4.0), multiply by POS_TWO
|
||||
# lo = 3.0 * 2.0 = 6.0, hi = 4.0 * 0.0 = 0.0 (inline const hi is 0)
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x44004200), # packed f16: hi=4.0, lo=3.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_pk_mul_f16(v[1], v[0], SrcEnum.POS_TWO, opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][1]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
self.assertAlmostEqual(lo, 6.0, places=1)
|
||||
self.assertAlmostEqual(hi, 0.0, places=1)
|
||||
|
||||
|
||||
class TestWMMA(unittest.TestCase):
|
||||
"""Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions."""
|
||||
|
||||
def test_v_wmma_f32_16x16x16_f16_all_ones(self):
|
||||
"""V_WMMA_F32_16X16X16_F16 with all ones produces 16.0."""
|
||||
instructions = []
|
||||
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
|
||||
for i in range(16, 32):
|
||||
instructions.append(v_mov_b32_e32(v[i], s[0]))
|
||||
for i in range(8):
|
||||
instructions.append(v_mov_b32_e32(v[i], 0))
|
||||
instructions.append(v_wmma_f32_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
|
||||
st = run_program(instructions, n_lanes=32)
|
||||
expected = f2i(16.0)
|
||||
for lane in range(32):
|
||||
for reg in range(8):
|
||||
result = st.vgpr[lane][reg]
|
||||
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
|
||||
|
||||
def test_v_wmma_f32_16x16x16_f16_with_accumulator(self):
|
||||
"""V_WMMA_F32_16X16X16_F16 with non-zero accumulator."""
|
||||
instructions = []
|
||||
instructions.append(s_mov_b32(s[0], 0x3c003c00))
|
||||
instructions.append(s_mov_b32(s[1], f2i(5.0)))
|
||||
for i in range(16, 32):
|
||||
instructions.append(v_mov_b32_e32(v[i], s[0]))
|
||||
for i in range(8):
|
||||
instructions.append(v_mov_b32_e32(v[i], s[1]))
|
||||
instructions.append(v_wmma_f32_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
|
||||
st = run_program(instructions, n_lanes=32)
|
||||
expected = f2i(21.0) # 16 + 5
|
||||
for lane in range(32):
|
||||
for reg in range(8):
|
||||
result = st.vgpr[lane][reg]
|
||||
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
|
||||
|
||||
|
||||
class TestSpecialOps(unittest.TestCase):
|
||||
"""Tests for special operations (SAD, PERM, DOT2)."""
|
||||
|
||||
def test_v_sad_u8_basic(self):
|
||||
"""V_SAD_U8 computes sum of absolute differences."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x04030201), # bytes: 1, 2, 3, 4
|
||||
s_mov_b32(s[1], 0x05040302), # bytes: 2, 3, 4, 5
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_sad_u8(v[3], v[0], v[1], v[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# |1-2| + |2-3| + |3-4| + |4-5| = 1 + 1 + 1 + 1 = 4
|
||||
self.assertEqual(st.vgpr[0][3], 4)
|
||||
|
||||
def test_v_sad_u8_identical_bytes(self):
|
||||
"""V_SAD_U8 with identical inputs returns accumulator."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x04030201),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], 10),
|
||||
v_mov_b32_e32(v[2], s[1]),
|
||||
v_sad_u8(v[3], v[0], v[0], v[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# Same inputs -> SAD = 0, result = accumulator = 10
|
||||
self.assertEqual(st.vgpr[0][3], 10)
|
||||
|
||||
def test_v_sad_u16_basic(self):
|
||||
"""V_SAD_U16 computes sum of absolute differences of u16 pairs."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x00030001), # hi=3, lo=1
|
||||
s_mov_b32(s[1], 0x00050002), # hi=5, lo=2
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_sad_u16(v[3], v[0], v[1], v[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# |1-2| + |3-5| = 1 + 2 = 3
|
||||
self.assertEqual(st.vgpr[0][3], 3)
|
||||
|
||||
def test_v_sad_u32_basic(self):
|
||||
"""V_SAD_U32 computes absolute difference of u32 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 100),
|
||||
s_mov_b32(s[1], 70),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_sad_u32(v[3], v[0], v[1], v[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# |100-70| = 30
|
||||
self.assertEqual(st.vgpr[0][3], 30)
|
||||
|
||||
def test_v_msad_u8_masked(self):
|
||||
"""V_MSAD_U8 masked SAD operation."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x04030201),
|
||||
s_mov_b32(s[1], 0x05040302),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_msad_u8(v[3], v[0], v[1], v[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# V_MSAD_U8 skips bytes where src0 is 0
|
||||
# Since no bytes are 0, result same as V_SAD_U8 = 4
|
||||
self.assertEqual(st.vgpr[0][3], 4)
|
||||
|
||||
def test_v_perm_b32_select_bytes(self):
|
||||
"""V_PERM_B32 selects bytes from two sources.
|
||||
|
||||
V_PERM_B32 concatenates {S1, S0} as a 64-bit value with S1 in low 32 bits.
|
||||
Selector byte values 0-3 select from S1, values 4-7 select from S0.
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x44332211), # src0: bytes 4-7 in 64-bit view
|
||||
s_mov_b32(s[1], 0x88776655), # src1: bytes 0-3 in 64-bit view
|
||||
s_mov_b32(s[2], 0x07060504), # select bytes 4,5,6,7 (from src0)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_perm_b32(v[2], v[0], v[1], s[2]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 0x44332211)
|
||||
|
||||
def test_v_dot2_f32_bf16_basic(self):
|
||||
"""V_DOT2_F32_BF16 computes dot product of bf16 pairs."""
|
||||
# bf16 1.0 = 0x3f80, bf16 2.0 = 0x4000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3f803f80), # packed bf16: lo=1.0, hi=1.0
|
||||
s_mov_b32(s[1], 0x40003f80), # packed bf16: lo=1.0, hi=2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_mov_b32_e32(v[2], 0),
|
||||
v_dot2_f32_bf16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
# 1.0*1.0 + 1.0*2.0 + 0 = 3.0
|
||||
result = i2f(st.vgpr[0][3])
|
||||
self.assertAlmostEqual(result, 3.0, places=4)
|
||||
|
||||
|
||||
class TestPackedMixedSigns(unittest.TestCase):
|
||||
"""Tests for packed operations with mixed sign values."""
|
||||
|
||||
def test_pk_add_f16_mixed_signs(self):
|
||||
"""V_PK_ADD_F16 with mixed positive/negative values."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xc0003c00), # packed: hi=-2.0, lo=1.0
|
||||
s_mov_b32(s[1], 0x3c003c00), # packed: hi=1.0, lo=1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pk_add_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
lo = _f16(result & 0xffff)
|
||||
hi = _f16((result >> 16) & 0xffff)
|
||||
self.assertAlmostEqual(lo, 2.0, places=2) # 1.0 + 1.0
|
||||
self.assertAlmostEqual(hi, -1.0, places=2) # -2.0 + 1.0
|
||||
|
||||
def test_pk_mul_f16_zero(self):
|
||||
"""V_PK_MUL_F16 with zero."""
|
||||
from extra.assembly.amd.pcode import _f16
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40004000), # packed: 2.0, 2.0
|
||||
s_mov_b32(s[1], 0x00000000), # packed: 0.0, 0.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_pk_mul_f16(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2]
|
||||
self.assertEqual(result, 0x00000000, "2.0 * 0.0 should be 0.0")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,476 @@
|
||||
"""Tests for VOPC instructions - vector compare operations.
|
||||
|
||||
Includes: v_cmp_class_f32, v_cmp_class_f16, v_cmp_eq_*, v_cmp_lt_*, v_cmp_gt_*
|
||||
"""
|
||||
import unittest
|
||||
from extra.assembly.amd.test.hw.helpers import *
|
||||
|
||||
VCC = 106 # SGPR index for VCC_LO
|
||||
|
||||
class TestCmpClass(unittest.TestCase):
|
||||
"""Tests for V_CMP_CLASS_F32 float classification."""
|
||||
|
||||
def test_cmp_class_quiet_nan(self):
|
||||
"""V_CMP_CLASS_F32 detects quiet NaN."""
|
||||
quiet_nan = 0x7fc00000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], quiet_nan),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0b0000000010), # bit 1 = quiet NaN
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
|
||||
|
||||
def test_cmp_class_signaling_nan(self):
|
||||
"""V_CMP_CLASS_F32 detects signaling NaN."""
|
||||
signal_nan = 0x7f800001
|
||||
instructions = [
|
||||
s_mov_b32(s[0], signal_nan),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0b0000000001), # bit 0 = signaling NaN
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
|
||||
|
||||
def test_cmp_class_positive_inf(self):
|
||||
"""V_CMP_CLASS_F32 detects +inf."""
|
||||
pos_inf = 0x7f800000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], pos_inf),
|
||||
s_mov_b32(s[1], 0b1000000000), # bit 9 = +inf
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect +inf")
|
||||
|
||||
def test_cmp_class_negative_inf(self):
|
||||
"""V_CMP_CLASS_F32 detects -inf."""
|
||||
neg_inf = 0xff800000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], neg_inf),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0b0000000100), # bit 2 = -inf
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect -inf")
|
||||
|
||||
def test_cmp_class_normal_positive(self):
|
||||
"""V_CMP_CLASS_F32 detects positive normal."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 1.0),
|
||||
s_mov_b32(s[1], 0b0100000000), # bit 8 = positive normal
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
|
||||
|
||||
def test_cmp_class_normal_negative(self):
|
||||
"""V_CMP_CLASS_F32 detects negative normal."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], -1.0),
|
||||
v_mov_b32_e32(v[1], 0b0000001000), # bit 3 = negative normal
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
|
||||
|
||||
def test_cmp_class_quiet_nan_not_signaling(self):
|
||||
"""Quiet NaN does not match signaling NaN mask."""
|
||||
quiet_nan = 0x7fc00000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], quiet_nan),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0b0000000001), # bit 0 = signaling NaN only
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "Quiet NaN should not match signaling mask")
|
||||
|
||||
def test_cmp_class_signaling_nan_not_quiet(self):
|
||||
"""Signaling NaN does not match quiet NaN mask."""
|
||||
signal_nan = 0x7f800001
|
||||
instructions = [
|
||||
s_mov_b32(s[0], signal_nan),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0b0000000010), # bit 1 = quiet NaN only
|
||||
v_cmp_class_f32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "Signaling NaN should not match quiet mask")
|
||||
|
||||
def test_v_cmp_sets_vcc_bits(self):
|
||||
"""V_CMP_EQ sets VCC bits based on per-lane comparison."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 5),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
v_cmp_eq_u32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=4)
|
||||
self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
|
||||
|
||||
|
||||
class TestCmpClassF16(unittest.TestCase):
|
||||
"""Tests for V_CMP_CLASS_F16 float classification.
|
||||
|
||||
Class bit mapping:
|
||||
bit 0 = signaling NaN
|
||||
bit 1 = quiet NaN
|
||||
bit 2 = -infinity
|
||||
bit 3 = -normal
|
||||
bit 4 = -denormal
|
||||
bit 5 = -zero
|
||||
bit 6 = +zero
|
||||
bit 7 = +denormal
|
||||
bit 8 = +normal
|
||||
bit 9 = +infinity
|
||||
"""
|
||||
|
||||
def test_cmp_class_f16_positive_zero(self):
|
||||
"""V_CMP_CLASS_F16: +zero matches bit 6."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0x0000), # f16 +0.0
|
||||
v_mov_b32_e32(v[1], 0x40), # bit 6 = +zero
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect positive zero")
|
||||
|
||||
def test_cmp_class_f16_negative_zero(self):
|
||||
"""V_CMP_CLASS_F16: -zero matches bit 5."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x8000), # f16 -0.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x20), # bit 5 = -zero
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect negative zero")
|
||||
|
||||
def test_cmp_class_f16_positive_normal(self):
|
||||
"""V_CMP_CLASS_F16: +1.0 (normal) matches bit 8."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 +1.0
|
||||
s_mov_b32(s[1], 0x100), # bit 8 = +normal
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
|
||||
|
||||
def test_cmp_class_f16_negative_normal(self):
|
||||
"""V_CMP_CLASS_F16: -1.0 (normal) matches bit 3."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xbc00), # f16 -1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x08), # bit 3 = -normal
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
|
||||
|
||||
def test_cmp_class_f16_positive_infinity(self):
|
||||
"""V_CMP_CLASS_F16: +inf matches bit 9."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7c00), # f16 +inf
|
||||
s_mov_b32(s[1], 0x200), # bit 9 = +inf
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect positive infinity")
|
||||
|
||||
def test_cmp_class_f16_negative_infinity(self):
|
||||
"""V_CMP_CLASS_F16: -inf matches bit 2."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xfc00), # f16 -inf
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x04), # bit 2 = -inf
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect negative infinity")
|
||||
|
||||
def test_cmp_class_f16_quiet_nan(self):
|
||||
"""V_CMP_CLASS_F16: quiet NaN matches bit 1."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7e00), # f16 quiet NaN
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x02), # bit 1 = quiet NaN
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
|
||||
|
||||
def test_cmp_class_f16_signaling_nan(self):
|
||||
"""V_CMP_CLASS_F16: signaling NaN matches bit 0."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7c01), # f16 signaling NaN
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x01), # bit 0 = signaling NaN
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
|
||||
|
||||
def test_cmp_class_f16_positive_denormal(self):
|
||||
"""V_CMP_CLASS_F16: positive denormal matches bit 7."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 1), # f16 +denormal (0x0001)
|
||||
v_mov_b32_e32(v[1], 0x80), # bit 7 = +denormal
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect positive denormal")
|
||||
|
||||
def test_cmp_class_f16_negative_denormal(self):
|
||||
"""V_CMP_CLASS_F16: negative denormal matches bit 4."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x8001), # f16 -denormal
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], 0x10), # bit 4 = -denormal
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Should detect negative denormal")
|
||||
|
||||
def test_cmp_class_f16_combined_mask_zeros(self):
|
||||
"""V_CMP_CLASS_F16: mask 0x60 covers both +zero and -zero."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0), # f16 +0.0
|
||||
v_mov_b32_e32(v[1], 0x60), # bits 5 and 6 (+-zero)
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x60")
|
||||
|
||||
def test_cmp_class_f16_combined_mask_1f8(self):
|
||||
"""V_CMP_CLASS_F16: mask 0x1f8 covers -normal,-denorm,-zero,+zero,+denorm,+normal.
|
||||
|
||||
This is the exact mask used in the f16 sin kernel at PC=46.
|
||||
"""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0), # f16 +0.0
|
||||
s_mov_b32(s[0], 0x1f8),
|
||||
v_mov_b32_e32(v[1], s[0]), # mask 0x1f8
|
||||
v_cmp_class_f16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x1f8")
|
||||
|
||||
def test_cmp_class_f16_vop3_encoding(self):
|
||||
"""V_CMP_CLASS_F16 in VOP3 encoding (v_cmp_class_f16_e64)."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 0), # f16 +0.0
|
||||
s_mov_b32(s[0], 0x1f8), # class mask
|
||||
v_cmp_class_f16_e64(VCC_LO, v[0], s[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with VOP3 encoding")
|
||||
|
||||
def test_cmp_class_f16_vop3_normal_positive(self):
|
||||
"""V_CMP_CLASS_F16 VOP3 encoding with +1.0 (normal)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 +1.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], 0x1f8), # class mask
|
||||
v_cmp_class_f16_e64(VCC_LO, v[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +1.0 (normal) with mask 0x1f8")
|
||||
|
||||
def test_cmp_class_f16_vop3_nan_fails_mask(self):
|
||||
"""V_CMP_CLASS_F16 VOP3: NaN should NOT match mask 0x1f8 (no NaN bits set)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7e00), # f16 quiet NaN
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], 0x1f8), # class mask
|
||||
v_cmp_class_f16_e64(VCC_LO, v[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for NaN with mask 0x1f8 (no NaN bits)")
|
||||
|
||||
def test_cmp_class_f16_vop3_inf_fails_mask(self):
|
||||
"""V_CMP_CLASS_F16 VOP3: +inf should NOT match mask 0x1f8 (no inf bits set)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x7c00), # f16 +inf
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
s_mov_b32(s[1], 0x1f8), # class mask
|
||||
v_cmp_class_f16_e64(VCC_LO, v[0], s[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for +inf with mask 0x1f8 (no inf bits)")
|
||||
|
||||
|
||||
class TestCmpInt(unittest.TestCase):
|
||||
"""Tests for integer comparison operations."""
|
||||
|
||||
def test_v_cmp_eq_u32(self):
|
||||
"""V_CMP_EQ_U32 sets VCC bits based on per-lane comparison."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 5),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
v_cmp_eq_u32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=4)
|
||||
self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
|
||||
|
||||
def test_cmp_eq_u16_opsel_lo_lo(self):
|
||||
"""V_CMP_EQ_U16 comparing lo halves."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x12340005), # lo=5, hi=0x1234
|
||||
s_mov_b32(s[1], 0xABCD0005), # lo=5, hi=0xABCD
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_cmp_eq_u16_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Lo halves should be equal")
|
||||
|
||||
def test_cmp_eq_u16_opsel_hi_hi(self):
|
||||
"""V_CMP_EQ_U16 comparing hi halves with VOP3 opsel."""
|
||||
instructions = [
|
||||
s_mov_b32(s[2], 0x00051234), # hi=5, lo=0x1234
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0x0005ABCD), # hi=5, lo=0xABCD
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
v_cmp_eq_u16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[0] & 1, 1, "Hi halves should be equal: 5==5")
|
||||
|
||||
def test_cmp_eq_u16_opsel_hi_hi_equal(self):
|
||||
"""V_CMP_EQ_U16 VOP3 with opsel=3 compares hi halves (equal case)."""
|
||||
instructions = [
|
||||
s_mov_b32(s[2], 0x12340005), # lo=5, hi=0x1234
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0x12340009), # lo=9, hi=0x1234
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
v_cmp_eq_u16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[0] & 1, 1, "hi==hi should be true: 0x1234==0x1234")
|
||||
|
||||
def test_cmp_gt_u16_opsel_hi(self):
|
||||
"""V_CMP_GT_U16 VOP3 with opsel=3 compares hi halves."""
|
||||
instructions = [
|
||||
s_mov_b32(s[2], 0x99990005), # lo=5, hi=0x9999
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0x12340005), # lo=5, hi=0x1234
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
v_cmp_gt_u16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[0] & 1, 1, "hi>hi should be true: 0x9999>0x1234")
|
||||
|
||||
|
||||
class TestCmpFloat(unittest.TestCase):
|
||||
"""Tests for float comparison operations."""
|
||||
|
||||
def test_v_cmp_lt_f16_vsrc1_hi(self):
|
||||
"""V_CMP_LT_F16 with both operands from high half using VOP3 opsel."""
|
||||
instructions = [
|
||||
s_mov_b32(s[2], 0x3c000000), # hi=1.0 (f16), lo=0
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0x40000000), # hi=2.0 (f16), lo=0
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
v_cmp_lt_f16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[0] & 1, 1, "1.0 < 2.0 should be true")
|
||||
|
||||
def test_v_cmp_gt_f16_vsrc1_hi(self):
|
||||
"""V_CMP_GT_F16 with both operands from high half using VOP3 opsel."""
|
||||
instructions = [
|
||||
s_mov_b32(s[2], 0x40000000), # hi=2.0 (f16), lo=0
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
s_mov_b32(s[2], 0x3c000000), # hi=1.0 (f16), lo=0
|
||||
v_mov_b32_e32(v[1], s[2]),
|
||||
v_cmp_gt_f16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.sgpr[0] & 1, 1, "2.0 > 1.0 should be true")
|
||||
|
||||
def test_v_cmp_eq_f16_vsrc1_hi_equal(self):
|
||||
"""v_cmp_eq_f16 with equal low and high halves."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x42004200), # hi=3.0 (0x4200), lo=3.0 (0x4200)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_cmp_eq_f16_e32(v[0], v[0].h),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (3.0 == 3.0)")
|
||||
|
||||
def test_v_cmp_neq_f16_vsrc1_hi(self):
|
||||
"""v_cmp_neq_f16 with different low and high halves."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x40003c00), # hi=2.0 (0x4000), lo=1.0 (0x3c00)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_cmp_lg_f16_e32(v[0], v[0].h),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (1.0 != 2.0)")
|
||||
|
||||
def test_v_cmp_nge_f16_inf_self(self):
|
||||
"""v_cmp_nge_f16 comparing -inf with itself (unordered less than).
|
||||
|
||||
Regression test: -inf < -inf should be false (IEEE 754).
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0xFC00FC00), # both halves = -inf (0xFC00)
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_cmp_nge_f16_e32(v[0], v[0].h),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "Expected vcc=0 (-inf >= -inf)")
|
||||
|
||||
def test_v_cmp_f16_multilane(self):
|
||||
"""v_cmp_lt_f16 with vsrc1=v128 across multiple lanes."""
|
||||
instructions = [
|
||||
# Lane 0: v0 = 0x40003c00 (hi=2.0, lo=1.0) -> 1.0 < 2.0 = true
|
||||
# Lane 1: v0 = 0x3c004000 (hi=1.0, lo=2.0) -> 2.0 < 1.0 = false
|
||||
v_mov_b32_e32(v[0], 0x40003c00), # default
|
||||
v_cmp_eq_u32_e32(1, v[255]), # vcc = (lane == 1)
|
||||
v_cndmask_b32_e64(v[0], v[0], 0x3c004000, SrcEnum.VCC_LO),
|
||||
v_cmp_lt_f16_e32(v[0], v[0].h),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=2)
|
||||
self.assertEqual(st.vcc & 1, 1, "Lane 0: expected vcc=1 (1.0 < 2.0)")
|
||||
self.assertEqual((st.vcc >> 1) & 1, 0, "Lane 1: expected vcc=0 (2.0 < 1.0)")
|
||||
|
||||
|
||||
class TestVCCBehavior(unittest.TestCase):
|
||||
"""Tests for VCC condition code behavior."""
|
||||
|
||||
def test_vcc_all_lanes_true(self):
|
||||
"""VCC should have all bits set when all lanes compare true."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 5),
|
||||
v_mov_b32_e32(v[1], 5),
|
||||
v_cmp_eq_u32_e32(v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=32)
|
||||
self.assertEqual(st.vcc, 0xFFFFFFFF, "All 32 lanes should be true")
|
||||
|
||||
def test_vcc_lane_dependent(self):
|
||||
"""VCC should differ per lane based on lane_id comparison."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], 16),
|
||||
v_cmp_lt_u32_e32(v[255], v[0]), # lanes 0-15 are < 16
|
||||
]
|
||||
st = run_program(instructions, n_lanes=32)
|
||||
self.assertEqual(st.vcc & 0xFFFF, 0xFFFF, "Lanes 0-15 should be true")
|
||||
self.assertEqual(st.vcc >> 16, 0x0000, "Lanes 16-31 should be false")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -9,10 +9,9 @@ os.environ["AMD"] = "1"
|
||||
os.environ["MOCKGPU"] = "1"
|
||||
os.environ["PYTHON_REMU"] = "1"
|
||||
|
||||
from extra.assembly.amd.emu import WaveState, decode_program, step_wave, WAVE_SIZE, set_valid_mem_ranges
|
||||
from extra.assembly.amd.emu import WaveState, decode_program, WAVE_SIZE, set_valid_mem_ranges, LDSMem
|
||||
from extra.assembly.amd.test.helpers import KernelInfo
|
||||
|
||||
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
|
||||
from extra.assembly.amd.test.bench_emu import REMU_PATH
|
||||
|
||||
def _is_f32_nan(bits: int) -> bool:
|
||||
"""Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
|
||||
@@ -92,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
|
||||
@@ -163,8 +158,9 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
|
||||
# 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'))
|
||||
'v_cvt_f16_f32', 's_add_i32', 's_sub_i32'))
|
||||
diffs = rust_before.diff(python_before, n_lanes)
|
||||
if diffs:
|
||||
trace_lines = []
|
||||
@@ -397,6 +393,9 @@ class TestTinygradKernels(unittest.TestCase):
|
||||
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
|
||||
self._test_kernel(lambda T: T([2.0], dtype=dtypes.float64).sin())
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
import unittest
|
||||
import functools
|
||||
from tinygrad import Tensor, Device, dtypes
|
||||
from tinygrad.uop.ops import UOp, Ops, KernelInfo
|
||||
from tinygrad.renderer import Estimates
|
||||
from tinygrad.runtime.support.compiler_amd import HIPCompiler
|
||||
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
from extra.assembly.amd.dsl import s, v, Inst
|
||||
|
||||
def assemble_insts(insts:list[Inst], name:str, arch:str, kernarg_size:int=8) -> tuple[UOp, UOp]:
|
||||
kd = {"kernarg_size":kernarg_size, "user_sgpr_kernarg_segment_ptr":1, "next_free_vgpr":8, "next_free_sgpr":8, "wavefront_size32":1}
|
||||
disasm = "\n".join([inst.disasm() for inst in insts])
|
||||
hsasrc = f".text\n.globl {name}\n.p2align 8\n.type fn_name,@function\n{name}:\n{disasm}\ns_code_end\n"
|
||||
hsasrc += f".rodata\n.p2align 6\n.amdhsa_kernel {name}\n"+"\n".join([f".amdhsa_{k} {v}" for k,v in kd.items()])+"\n.end_amdhsa_kernel"
|
||||
binary = HIPCompiler(arch).compile(hsasrc)
|
||||
return UOp(Ops.SOURCE, arg=disasm), UOp(Ops.BINARY, arg=binary)
|
||||
|
||||
def custom_add_one(A:UOp, arch:str) -> UOp:
|
||||
A = A.flatten()
|
||||
assert dtypes.is_float(A.dtype.base), f"buffer dtype must be float32, got {A.dtype}"
|
||||
threads = UOp.special(A.size, "lidx0")
|
||||
insts = [
|
||||
s_load_b64(s[0:1], s[0:1], soffset=NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_lshlrev_b32_e32(v[0], 2, v[0]), # element offset
|
||||
global_load_b32(v[1], v[0], saddr=s[0:1]),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[2], 1.0),
|
||||
v_add_f32_e32(v[1], v[1], v[2]),
|
||||
global_store_b32(addr=v[0], data=v[1], saddr=s[0:1]),
|
||||
s_endpgm(),
|
||||
]
|
||||
sink = UOp.sink(A.base, threads, arg=KernelInfo(name:=f"custom_add_one_{A.size}", estimates=Estimates(ops=A.size, mem=A.size*4*2)))
|
||||
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=(*sink.src, sink)), *assemble_insts(insts, name, arch)))
|
||||
|
||||
class TestCustomKernel(unittest.TestCase):
|
||||
def test_simple(self):
|
||||
a = Tensor.full((16, 16), 1.).contiguous().realize()
|
||||
a = Tensor.custom_kernel(a, fxn=functools.partial(custom_add_one, arch=Device[Device.DEFAULT].arch))[0]
|
||||
ei = a.schedule()[-1].lower()
|
||||
self.assertEqual(ei.prg.estimates.ops, a.numel())
|
||||
self.assertEqual(ei.prg.estimates.mem, a.nbytes()*2)
|
||||
ei.run()
|
||||
self.assertTrue((a.numpy() == 2.).all())
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,170 @@
|
||||
import unittest
|
||||
from extra.assembly.amd.dsl import *
|
||||
from extra.assembly.amd.dsl import VDSTYField
|
||||
from extra.assembly.amd.autogen.rdna3.enum import VOP1Op, VOP2Op
|
||||
from extra.assembly.amd.autogen.rdna3.ins import VOP1
|
||||
|
||||
class TestRegisters(unittest.TestCase):
|
||||
def test_vgpr_single(self):
|
||||
self.assertEqual(repr(v[5]), "v[5]")
|
||||
self.assertEqual(v[5].offset, 261) # 256 + 5
|
||||
self.assertEqual(v[5].sz, 1)
|
||||
|
||||
def test_sgpr_single(self):
|
||||
self.assertEqual(repr(s[10]), "s[10]")
|
||||
self.assertEqual(s[10].offset, 10)
|
||||
|
||||
def test_vgpr_range(self):
|
||||
self.assertEqual(repr(v[0:3]), "v[0:3]")
|
||||
self.assertEqual(v[0:3].offset, 256)
|
||||
self.assertEqual(v[0:3].sz, 4)
|
||||
|
||||
def test_sgpr_range(self):
|
||||
self.assertEqual(repr(s[4:5]), "s[4:5]")
|
||||
self.assertEqual(s[4:5].sz, 2)
|
||||
|
||||
def test_ttmp_reslice(self):
|
||||
# ttmp is src[108:123], so ttmp[0] should be src[108]
|
||||
self.assertEqual(ttmp[0].offset, 108)
|
||||
self.assertEqual(ttmp[1].offset, 109)
|
||||
# ttmp[0:1] is 2 elements (inclusive slicing)
|
||||
self.assertEqual(ttmp[0:1].offset, 108)
|
||||
self.assertEqual(ttmp[0:1].sz, 2)
|
||||
# ttmp[0:1][0] should be src[108]
|
||||
self.assertEqual(ttmp[0:1][0].offset, 108)
|
||||
|
||||
def test_special_regs(self):
|
||||
self.assertEqual(NULL.offset, 124)
|
||||
self.assertEqual(M0.offset, 125)
|
||||
self.assertEqual(EXEC_LO.offset, 126)
|
||||
self.assertEqual(EXEC_HI.offset, 127)
|
||||
# Check repr round-trips
|
||||
self.assertEqual(repr(NULL), "NULL")
|
||||
self.assertEqual(repr(M0), "M0")
|
||||
self.assertEqual(repr(EXEC_LO), "EXEC_LO")
|
||||
self.assertEqual(repr(EXEC), "EXEC")
|
||||
|
||||
def test_vcc(self):
|
||||
self.assertEqual(VCC.offset, 106)
|
||||
self.assertEqual(VCC.sz, 2)
|
||||
self.assertEqual(VCC_LO.offset, 106)
|
||||
self.assertEqual(VCC_HI.offset, 107)
|
||||
# Check repr round-trips
|
||||
self.assertEqual(repr(VCC_LO), "VCC_LO")
|
||||
self.assertEqual(repr(VCC_HI), "VCC_HI")
|
||||
self.assertEqual(repr(VCC), "VCC")
|
||||
|
||||
def test_float_constants(self):
|
||||
self.assertEqual(src[240].offset, 240)
|
||||
self.assertEqual(repr(src[240]), "0.5")
|
||||
self.assertEqual(repr(src[242]), "1.0")
|
||||
self.assertEqual(repr(src[243]), "-1.0")
|
||||
|
||||
def test_int_constants(self):
|
||||
self.assertEqual(repr(src[128]), "0")
|
||||
self.assertEqual(repr(src[129]), "1")
|
||||
self.assertEqual(repr(src[192]), "64")
|
||||
self.assertEqual(repr(src[193]), "-1")
|
||||
self.assertEqual(repr(src[208]), "-16")
|
||||
|
||||
class TestEnumBitField(unittest.TestCase):
|
||||
def test_enum_name(self):
|
||||
self.assertEqual(VOP1Op.V_MOV_B32_E32.name, "V_MOV_B32_E32")
|
||||
|
||||
def test_enum_value(self):
|
||||
self.assertEqual(VOP1Op.V_MOV_B32_E32.value, 1)
|
||||
|
||||
def test_enum_comparison(self):
|
||||
self.assertEqual(VOP1Op.V_MOV_B32_E32, VOP1Op.V_MOV_B32_E32)
|
||||
self.assertNotEqual(VOP1Op.V_NOP_E32, VOP1Op.V_MOV_B32_E32)
|
||||
|
||||
def test_enum_different_types(self):
|
||||
# VOP1Op and VOP2Op are different enums, even if same value
|
||||
self.assertNotEqual(VOP1Op.V_MOV_B32_E32, VOP2Op.V_CNDMASK_B32_E32)
|
||||
|
||||
def test_wrong_enum_type_raises(self):
|
||||
# Passing VOP2Op to VOP1 should raise
|
||||
with self.assertRaises(RuntimeError):
|
||||
VOP1(VOP2Op.V_CNDMASK_B32_E32, v[5], v[6])
|
||||
|
||||
class TestVOP1(unittest.TestCase):
|
||||
def test_class_setup(self):
|
||||
self.assertEqual(VOP1._size(), 4)
|
||||
field_names = [n for n, _ in VOP1._fields]
|
||||
self.assertIn('encoding', field_names)
|
||||
self.assertIn('op', field_names)
|
||||
self.assertIn('vdst', field_names)
|
||||
self.assertIn('src0', field_names)
|
||||
|
||||
def test_encoding_vgpr_vgpr(self):
|
||||
i = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
|
||||
raw = i._raw
|
||||
# Check each field
|
||||
self.assertEqual((raw >> 25) & 0x7f, 0b0111111) # encoding
|
||||
self.assertEqual((raw >> 17) & 0xff, 5) # vdst (just VGPR index)
|
||||
self.assertEqual((raw >> 9) & 0xff, 1) # op
|
||||
self.assertEqual(raw & 0x1ff, 262) # src0 (256 + 6)
|
||||
|
||||
def test_encoding_vgpr_sgpr(self):
|
||||
i = VOP1(VOP1Op.V_MOV_B32_E32, v[5], s[10])
|
||||
raw = i._raw
|
||||
self.assertEqual((raw >> 17) & 0xff, 5) # vdst (just VGPR index)
|
||||
self.assertEqual(raw & 0x1ff, 10) # src0 (SGPR encoded)
|
||||
|
||||
def test_to_bytes(self):
|
||||
i = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
|
||||
b = i.to_bytes()
|
||||
self.assertEqual(len(b), 4)
|
||||
self.assertEqual(int.from_bytes(b, 'little'), i._raw)
|
||||
|
||||
def test_from_bytes(self):
|
||||
i1 = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
|
||||
i2 = VOP1.from_bytes(i1.to_bytes())
|
||||
self.assertEqual(i1._raw, i2._raw)
|
||||
|
||||
def test_repr(self):
|
||||
i = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
|
||||
self.assertEqual(repr(i), "v_mov_b32_e32(v[5], v[6])")
|
||||
|
||||
def test_repr_sgpr_src(self):
|
||||
i = VOP1(VOP1Op.V_MOV_B32_E32, v[5], s[10])
|
||||
self.assertEqual(repr(i), "v_mov_b32_e32(v[5], s[10])")
|
||||
|
||||
def test_kwargs(self):
|
||||
i1 = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
|
||||
i2 = VOP1(op=VOP1Op.V_MOV_B32_E32, vdst=v[5], src0=v[6])
|
||||
self.assertEqual(i1._raw, i2._raw)
|
||||
|
||||
def test_kwargs_partial(self):
|
||||
i1 = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
|
||||
i2 = VOP1(VOP1Op.V_MOV_B32_E32, src0=v[6], vdst=v[5])
|
||||
self.assertEqual(i1._raw, i2._raw)
|
||||
|
||||
class TestVDSTYField(unittest.TestCase):
|
||||
def test_encode_even_vgpr(self):
|
||||
f = VDSTYField(6, 0) # 7-bit field
|
||||
self.assertEqual(f.encode(v[0]), 0)
|
||||
self.assertEqual(f.encode(v[2]), 1)
|
||||
self.assertEqual(f.encode(v[4]), 2)
|
||||
self.assertEqual(f.encode(v[254]), 127)
|
||||
|
||||
def test_encode_non_vgpr_raises(self):
|
||||
f = VDSTYField(6, 0)
|
||||
with self.assertRaises(ValueError) as ctx:
|
||||
f.encode(s[0])
|
||||
self.assertIn("VGPR", str(ctx.exception))
|
||||
|
||||
def test_encode_non_reg_raises(self):
|
||||
f = VDSTYField(6, 0)
|
||||
with self.assertRaises(TypeError) as ctx:
|
||||
f.encode(42)
|
||||
self.assertIn("Reg", str(ctx.exception))
|
||||
|
||||
def test_decode_returns_raw(self):
|
||||
f = VDSTYField(6, 0)
|
||||
# decode returns raw value, actual vdsty computed with vdstx context
|
||||
self.assertEqual(f.decode(0), 0)
|
||||
self.assertEqual(f.decode(127), 127)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,175 +1,13 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
|
||||
"""Test DS and other compute-relevant instruction formats.
|
||||
|
||||
Note: Graphics-only formats (EXP, MUBUF, MTBUF, MIMG) are not supported - use GLOBAL/FLAT for memory access in compute.
|
||||
"""
|
||||
import unittest
|
||||
from extra.assembly.amd.autogen.rdna3 import *
|
||||
from extra.assembly.amd.dsl import encode_src
|
||||
|
||||
class TestMUBUF(unittest.TestCase):
|
||||
"""Test MUBUF (buffer) instructions."""
|
||||
|
||||
def test_buffer_load_b32_basic(self):
|
||||
# buffer_load_b32 v5, off, s[8:11], s3 offset:4095
|
||||
# GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]
|
||||
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]))
|
||||
|
||||
def test_buffer_load_b32_idxen(self):
|
||||
# buffer_load_b32 v5, v0, s[8:11], s3 idxen offset:4095
|
||||
# GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]
|
||||
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, idxen=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]))
|
||||
|
||||
def test_buffer_load_b32_offen(self):
|
||||
# buffer_load_b32 v5, v0, s[8:11], s3 offen offset:4095
|
||||
# GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]
|
||||
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, offen=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]))
|
||||
|
||||
def test_buffer_load_b32_glc(self):
|
||||
# buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc
|
||||
# GFX11: encoding: [0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]
|
||||
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]))
|
||||
|
||||
def test_buffer_load_b32_slc(self):
|
||||
# buffer_load_b32 v5, off, s[8:11], s3 offset:4095 slc
|
||||
# GFX11: encoding: [0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]
|
||||
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, slc=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]))
|
||||
|
||||
def test_buffer_load_b32_dlc(self):
|
||||
# buffer_load_b32 v5, off, s[8:11], s3 offset:4095 dlc
|
||||
# GFX11: encoding: [0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]
|
||||
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, dlc=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]))
|
||||
|
||||
def test_buffer_load_b32_all_flags(self):
|
||||
# buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc slc dlc
|
||||
# GFX11: encoding: [0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]
|
||||
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1, slc=1, dlc=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]))
|
||||
|
||||
def test_buffer_store_b32(self):
|
||||
# buffer_store_b32 v1, off, s[12:15], s4 offset:4095
|
||||
# GFX11: encoding: [0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]
|
||||
inst = buffer_store_b32(vdata=v[1], vaddr=v[0], srsrc=s[12:16], soffset=s[4], offset=4095)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]))
|
||||
|
||||
def test_buffer_load_b64(self):
|
||||
# buffer_load_b64 v[5:6], off, s[8:11], s3 offset:4095
|
||||
# GFX11: encoding: [0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]
|
||||
inst = buffer_load_b64(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]))
|
||||
|
||||
def test_buffer_load_soffset_m0(self):
|
||||
# buffer_load_b32 v5, off, s[8:11], m0 offset:4095
|
||||
# GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]
|
||||
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=M0, offset=4095)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]))
|
||||
|
||||
def test_buffer_load_soffset_inline_const(self):
|
||||
# buffer_load_b32 v5, off, s[8:11], 0 offset:4095
|
||||
# GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]
|
||||
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=0, offset=4095)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]))
|
||||
|
||||
def test_buffer_disasm_roundtrip(self):
|
||||
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
|
||||
decoded = MUBUF.from_bytes(inst.to_bytes())
|
||||
self.assertEqual(decoded.to_bytes(), inst.to_bytes())
|
||||
|
||||
|
||||
class TestMTBUF(unittest.TestCase):
|
||||
"""Test MTBUF (typed buffer) instructions."""
|
||||
|
||||
def test_tbuffer_load_format_x(self):
|
||||
# tbuffer_load_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
|
||||
# BUF_FMT_32_FLOAT = 22
|
||||
# GFX11: encoding: [0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]
|
||||
inst = tbuffer_load_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]))
|
||||
|
||||
def test_tbuffer_store_format_x(self):
|
||||
# tbuffer_store_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
|
||||
# BUF_FMT_32_FLOAT = 22
|
||||
# GFX11: encoding: [0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]
|
||||
inst = tbuffer_store_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]))
|
||||
|
||||
def test_tbuffer_load_format_xy(self):
|
||||
# tbuffer_load_format_xy v[5:6], off, s[8:11], s3 format:[BUF_FMT_32_32_FLOAT] offset:4095
|
||||
# BUF_FMT_32_32_FLOAT = 50
|
||||
# GFX11: encoding: [0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]
|
||||
inst = tbuffer_load_format_xy(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=50)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]))
|
||||
|
||||
|
||||
class TestMIMG(unittest.TestCase):
|
||||
"""Test MIMG (image) instructions."""
|
||||
|
||||
def test_image_load_2d(self):
|
||||
# image_load v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
|
||||
# GFX11: encoding: [0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
|
||||
inst = image_load(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1) # dim=1 is SQ_RSRC_IMG_2D
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
|
||||
|
||||
def test_image_store_2d(self):
|
||||
# image_store v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
|
||||
# GFX11: encoding: [0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]
|
||||
inst = image_store(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]))
|
||||
|
||||
def test_image_load_1d(self):
|
||||
# image_load v[0:3], v4, s[0:7] dmask:0xf dim:SQ_RSRC_IMG_1D
|
||||
# GFX11: encoding: [0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
|
||||
inst = image_load(vdata=v[0:4], vaddr=v[4], srsrc=s[0:8], dmask=0xf, dim=0) # dim=0 is SQ_RSRC_IMG_1D
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
|
||||
|
||||
def test_image_sample(self):
|
||||
# image_sample v[0:3], v[4:5], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
|
||||
# GFX11: encoding: [0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]
|
||||
inst = image_sample(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], ssamp=s[8:12], dmask=0xf, dim=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]))
|
||||
|
||||
def test_image_load_d16(self):
|
||||
# image_load v[0:1], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D d16
|
||||
# GFX11: encoding: [0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]
|
||||
inst = image_load(vdata=v[0:2], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1, d16=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]))
|
||||
|
||||
|
||||
class TestEXP(unittest.TestCase):
|
||||
"""Test EXP (export) instructions."""
|
||||
|
||||
def test_exp_mrt0(self):
|
||||
# exp mrt0 v0, v1, v2, v3
|
||||
# GFX11: encoding: [0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]
|
||||
inst = EXP(en=0xf, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[2], vsrc3=v[3])
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]))
|
||||
|
||||
def test_exp_mrtz(self):
|
||||
# exp mrtz v4, v3, v2, v1
|
||||
# GFX11: encoding: [0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]
|
||||
inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1])
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]))
|
||||
|
||||
def test_exp_mrtz_done(self):
|
||||
# exp mrtz v4, v3, v2, v1 done
|
||||
# GFX11: encoding: [0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x01]
|
||||
inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[3], done=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x03]))
|
||||
|
||||
def test_exp_partial_mask(self):
|
||||
# exp mrt0 v0, v1, off, off (en=0x3, only first two components)
|
||||
# GFX11: encoding: [0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]
|
||||
inst = EXP(en=0x3, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[0], vsrc3=v[0])
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]))
|
||||
|
||||
def test_exp_row_en(self):
|
||||
# exp mrtz v4, v3, v2, v1 row_en
|
||||
# GFX11: encoding: [0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]
|
||||
inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1], row=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]))
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
from extra.assembly.amd.dsl import VCC_HI, EXEC_LO, NULL
|
||||
OFF = NULL # OFF is alias for NULL
|
||||
from extra.assembly.amd.decode import detect_format
|
||||
|
||||
|
||||
class TestDS(unittest.TestCase):
|
||||
@@ -196,7 +34,7 @@ class TestDS(unittest.TestCase):
|
||||
def test_ds_load_b64(self):
|
||||
# ds_load_b64 v[0:1], v2
|
||||
# GFX11: encoding: [0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]
|
||||
inst = ds_load_b64(vdst=v[0:2], addr=v[2])
|
||||
inst = ds_load_b64(vdst=v[0:1], addr=v[2])
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]))
|
||||
|
||||
def test_ds_add_u32(self):
|
||||
@@ -240,31 +78,31 @@ class TestFLAT(unittest.TestCase):
|
||||
def test_global_load_b32(self):
|
||||
# global_load_b32 v0, v[1:2], off (seg=2 for global)
|
||||
# GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]
|
||||
inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF)
|
||||
inst = global_load_b32(vdst=v[0], addr=v[1:2], saddr=OFF)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]))
|
||||
|
||||
def test_global_store_b32(self):
|
||||
# global_store_b32 v[0:1], v2, off (seg=2 for global)
|
||||
# GFX11: encoding: [0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]
|
||||
inst = global_store_b32(addr=v[0:2], data=v[2], saddr=OFF)
|
||||
inst = global_store_b32(addr=v[0:1], data=v[2], saddr=OFF)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]))
|
||||
|
||||
def test_global_load_b32_saddr(self):
|
||||
# global_load_b32 v0, v1, s[0:1] (seg=2 for global)
|
||||
# GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]
|
||||
inst = global_load_b32(vdst=v[0], addr=v[1], saddr=s[0:2])
|
||||
inst = global_load_b32(vdst=v[0], addr=v[1], saddr=s[0:1])
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]))
|
||||
|
||||
def test_global_load_b32_offset(self):
|
||||
# global_load_b32 v0, v[1:2], off offset:256 (seg=2 for global)
|
||||
# GFX11: encoding: [0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]
|
||||
inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF, offset=256)
|
||||
inst = global_load_b32(vdst=v[0], addr=v[1:2], saddr=OFF, offset=256)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]))
|
||||
|
||||
def test_global_load_b64(self):
|
||||
# global_load_b64 v[0:1], v[2:3], off (seg=2 for global)
|
||||
# GFX11: encoding: [0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]
|
||||
inst = global_load_b64(vdst=v[0:2], addr=v[2:4], saddr=OFF)
|
||||
inst = global_load_b64(vdst=v[0:1], addr=v[2:3], saddr=OFF)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]))
|
||||
|
||||
|
||||
@@ -274,19 +112,19 @@ class TestSMEM(unittest.TestCase):
|
||||
def test_smem_dlc_bit_position(self):
|
||||
# s_load_b32 s5, s[2:3], s0 dlc - tests that DLC is at bit 13 (not bit 14)
|
||||
# GFX11: encoding: [0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]
|
||||
inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], dlc=1)
|
||||
inst = s_load_b32(sdata=s[5], sbase=s[2:3], soffset=s[0], dlc=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]))
|
||||
|
||||
def test_smem_glc_bit_position(self):
|
||||
# s_load_b32 s5, s[2:3], s0 glc - tests that GLC is at bit 14 (not bit 16)
|
||||
# GFX11: encoding: [0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]
|
||||
inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1)
|
||||
inst = s_load_b32(sdata=s[5], sbase=s[2:3], soffset=s[0], glc=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]))
|
||||
|
||||
def test_smem_glc_dlc_combined(self):
|
||||
# s_load_b32 s5, s[2:3], s0 glc dlc - tests both flags together
|
||||
# GFX11: encoding: [0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]
|
||||
inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1, dlc=1)
|
||||
inst = s_load_b32(sdata=s[5], sbase=s[2:3], soffset=s[0], glc=1, dlc=1)
|
||||
self.assertEqual(inst.to_bytes(), bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]))
|
||||
|
||||
def test_smem_disasm_roundtrip_dlc(self):
|
||||
@@ -308,25 +146,84 @@ 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.amd.dsl import RawImm
|
||||
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=RawImm(107), src1=0xaf123456, src2=v[255])
|
||||
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=VCC_HI, 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)
|
||||
|
||||
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.amd.dsl import RawImm
|
||||
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=NULL, src1=RawImm(126), src2=0xaf123456)
|
||||
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=NULL, src1=EXEC_LO, 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.amd.dsl import RawImm
|
||||
inst = VOP3P(VOP3POp.V_PK_ADD_F16, vdst=v[5], src0=RawImm(240), src1=0x12345678, src2=v[0])
|
||||
inst = VOP3P(VOP3POp.V_PK_ADD_F16, vdst=v[5], src0=0.5, 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(sdata=s[0], sbase=s[2:3], offset=0).to_bytes()), SMEM)
|
||||
self.assertEqual(detect_format(s_load_b64(sdata=s[0:1], sbase=s[2:3], soffset=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(vdst=v[0], addr=v[1:2], saddr=NULL).to_bytes()), FLAT)
|
||||
self.assertEqual(detect_format(global_store_b32(addr=v[0:1], data=v[2], saddr=NULL).to_bytes()), FLAT)
|
||||
|
||||
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()
|
||||
|
||||
@@ -2,9 +2,8 @@
|
||||
# the Inst constructor should be looking at the types of the fields to correctly set the value
|
||||
|
||||
import unittest, struct
|
||||
from extra.assembly.amd.autogen.rdna3 import *
|
||||
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):
|
||||
@@ -13,20 +12,20 @@ class TestIntegration(unittest.TestCase):
|
||||
if not hasattr(self, 'inst'): return
|
||||
b = self.inst.to_bytes()
|
||||
st = self.inst.disasm()
|
||||
reasm = asm(st)
|
||||
desc = f"{st:25s} {self.inst} {b!r} {reasm}"
|
||||
# Test that the instruction can be compiled by LLVM and produces the same bytes
|
||||
desc = f"{st:25s} {self.inst} {b!r}"
|
||||
self.assertEqual(b, compile_asm(st), desc)
|
||||
# TODO: this compare should work for valid things
|
||||
#self.assertEqual(self.inst, reasm)
|
||||
self.assertEqual(repr(self.inst), repr(reasm))
|
||||
print(desc)
|
||||
|
||||
def test_wmma(self):
|
||||
self.inst = v_wmma_f32_16x16x16_f16(v[0:7], v[184:191], v[136:143], v[0:7])
|
||||
|
||||
def test_load_b128(self):
|
||||
self.inst = s_load_b128(s[4:7], s[0:1], NULL, 0)
|
||||
|
||||
def test_load_b128_wrong_size(self):
|
||||
# this should have to be 4 regs on the loaded to
|
||||
with self.assertRaises(Exception):
|
||||
with self.assertRaises(TypeError):
|
||||
self.inst = s_load_b128(s[4:6], s[0:1], NULL, 0)
|
||||
|
||||
def test_mov_b32(self):
|
||||
@@ -140,7 +139,7 @@ class TestRegisterSliceSyntax(unittest.TestCase):
|
||||
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)")
|
||||
self.assertEqual(reg.sz, 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
|
||||
@@ -151,7 +150,7 @@ class TestRegisterSliceSyntax(unittest.TestCase):
|
||||
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")
|
||||
self.assertEqual(reg_from_disasm.sz, 4, "s[4:7] from disasm should give 4 registers")
|
||||
|
||||
class TestInstructionEquality(unittest.TestCase):
|
||||
"""
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
|
||||
import unittest, re, io, sys, subprocess
|
||||
from extra.assembly.amd.autogen.rdna3 import *
|
||||
from extra.assembly.amd.asm import waitcnt, asm
|
||||
from extra.assembly.amd.test.helpers import get_llvm_mc
|
||||
import unittest, io, sys
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
|
||||
def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
|
||||
return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
|
||||
|
||||
def disassemble(lib: bytes, arch: str = "gfx1100") -> str:
|
||||
"""Disassemble ELF binary using tinygrad's compiler, return raw output."""
|
||||
@@ -40,7 +41,7 @@ def assemble_and_disassemble(instructions: list, arch: str = "gfx1100") -> list[
|
||||
return parse_disassembly(disassemble(lib, arch))
|
||||
|
||||
class TestIntegration(unittest.TestCase):
|
||||
"""Test our assembler output matches LLVM disassembly."""
|
||||
"""Test our DSL output matches LLVM disassembly."""
|
||||
|
||||
def test_simple_sop1(self):
|
||||
"""Test SOP1 instructions round-trip."""
|
||||
@@ -90,9 +91,9 @@ class TestIntegration(unittest.TestCase):
|
||||
def test_memory_ops(self):
|
||||
"""Test memory instructions."""
|
||||
instructions = [
|
||||
s_load_b32(s[0], s[0:2], NULL),
|
||||
s_load_b32(s[0], s[0:1], NULL),
|
||||
s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
|
||||
global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
|
||||
global_store_b32(addr=v[0:1], data=v[2], saddr=OFF),
|
||||
s_endpgm(),
|
||||
]
|
||||
disasm = assemble_and_disassemble(instructions)
|
||||
@@ -108,12 +109,12 @@ class TestIntegration(unittest.TestCase):
|
||||
v_mov_b32_e32(v[0], s[0]), # base addr low
|
||||
v_mov_b32_e32(v[1], s[1]), # base addr high
|
||||
# Load value
|
||||
global_load_b32(vdst=v[2], addr=v[0:2], saddr=OFF),
|
||||
global_load_b32(vdst=v[2], addr=v[0:1], saddr=OFF),
|
||||
s_waitcnt(simm16=waitcnt(vmcnt=0)),
|
||||
# Add 1.0
|
||||
v_add_f32_e32(v[2], 1.0, v[2]),
|
||||
# Store result
|
||||
global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
|
||||
global_store_b32(addr=v[0:1], data=v[2], saddr=OFF),
|
||||
s_endpgm(),
|
||||
]
|
||||
disasm = assemble_and_disassemble(instructions)
|
||||
@@ -148,79 +149,6 @@ class TestIntegration(unittest.TestCase):
|
||||
return
|
||||
self.fail("Could not find s_mov_b32 in disassembly")
|
||||
|
||||
class TestAsm(unittest.TestCase):
|
||||
"""Test asm() string parsing."""
|
||||
|
||||
def test_asm_basic(self):
|
||||
"""Test basic instruction parsing."""
|
||||
inst = asm('s_mov_b32 s0, s1')
|
||||
self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], s[1]).to_bytes())
|
||||
|
||||
def test_asm_with_immediates(self):
|
||||
"""Test parsing with immediate values."""
|
||||
inst = asm('s_add_u32 s0, s1, 10')
|
||||
self.assertEqual(inst.to_bytes(), s_add_u32(s[0], s[1], 10).to_bytes())
|
||||
|
||||
def test_asm_float_const(self):
|
||||
"""Test parsing float constants."""
|
||||
inst = asm('v_mul_f32_e32 v0, 1.0, v1')
|
||||
self.assertEqual(inst.to_bytes(), v_mul_f32_e32(v[0], 1.0, v[1]).to_bytes())
|
||||
|
||||
def test_asm_hex_immediate(self):
|
||||
"""Test parsing hex immediates."""
|
||||
inst = asm('s_waitcnt 0xfc07')
|
||||
self.assertEqual(inst.to_bytes(), s_waitcnt(simm16=0xfc07).to_bytes())
|
||||
|
||||
def test_asm_special_regs(self):
|
||||
"""Test parsing special registers."""
|
||||
inst = asm('s_mov_b32 s0, vcc_lo')
|
||||
self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], VCC_LO).to_bytes())
|
||||
|
||||
def test_asm_register_range(self):
|
||||
"""Test parsing register ranges."""
|
||||
inst = asm('s_load_b128 s[4:7], s[0:1], null')
|
||||
self.assertEqual(inst.to_bytes(), s_load_b128(s[4:7], s[0:1], NULL).to_bytes())
|
||||
|
||||
def test_asm_matches_llvm(self):
|
||||
"""Test asm() output matches LLVM assembler."""
|
||||
from tinygrad.runtime.support.compiler_amd import HIPCompiler
|
||||
compiler = HIPCompiler('gfx1100')
|
||||
|
||||
def get_llvm_bytes(instr: str) -> bytes:
|
||||
src = f'.text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n{instr}\n'
|
||||
lib = compiler.compile(src)
|
||||
raw = disassemble(lib)
|
||||
for line in raw.splitlines():
|
||||
if instr.split()[0] in line and '//' in line:
|
||||
hex_str = line.split('//')[1].strip().split(':')[1].strip()
|
||||
return bytes.fromhex(hex_str)[::-1]
|
||||
return b''
|
||||
|
||||
tests = ['s_mov_b32 s0, s1', 's_endpgm', 'v_add_f32_e32 v0, v1, v2']
|
||||
for t in tests:
|
||||
self.assertEqual(asm(t).to_bytes(), get_llvm_bytes(t), f"mismatch for: {t}")
|
||||
|
||||
def test_asm_vop3_modifiers(self):
|
||||
"""Test asm() with VOP3 modifiers (neg, abs, clamp)."""
|
||||
def get_llvm_encoding(instr: str) -> str:
|
||||
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(' ','')
|
||||
return ''
|
||||
|
||||
tests = [
|
||||
'v_fma_f32 v0, -v1, v2, v3', # neg on src0
|
||||
'v_fma_f32 v0, v1, |v2|, v3', # abs on src1
|
||||
'v_fma_f32 v0, v1, v2, v3 clamp', # clamp
|
||||
'v_fma_f32 v0, -v1, |v2|, v3 clamp', # all modifiers
|
||||
'v_fma_f32 v0, -|v1|, v2, v3', # neg+abs on same operand
|
||||
]
|
||||
for t in tests:
|
||||
our_hex = asm(t).to_bytes().hex()
|
||||
llvm_hex = get_llvm_encoding(t)
|
||||
self.assertEqual(our_hex, llvm_hex, f"mismatch for: {t}")
|
||||
|
||||
class TestTinygradIntegration(unittest.TestCase):
|
||||
"""Test that we can parse disassembled tinygrad kernels."""
|
||||
|
||||
|
||||
@@ -1,201 +1,127 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test RDNA3 assembler/disassembler against LLVM test vectors."""
|
||||
import unittest, re, subprocess
|
||||
"""Test AMD assembler/disassembler against LLVM test vectors.
|
||||
|
||||
Only compute-relevant instruction formats are tested. Graphics-only formats not supported:
|
||||
- MUBUF/MTBUF: buffer instructions with resource descriptors (use GLOBAL/FLAT instead)
|
||||
- MIMG: image/texture instructions
|
||||
- EXP/VEXPORT: export instructions for pixel/vertex output
|
||||
- VIMAGE/VSAMPLE: image sampling instructions (RDNA4)
|
||||
- VBUFFER: buffer instructions (RDNA4)
|
||||
"""
|
||||
import unittest, re, subprocess, functools
|
||||
from tinygrad.helpers import fetch
|
||||
from extra.assembly.amd.autogen.rdna3 import *
|
||||
from extra.assembly.amd.asm import asm
|
||||
from extra.assembly.amd.disasm import disasm
|
||||
from extra.assembly.amd.decode import decode_inst, detect_format
|
||||
from extra.assembly.amd.test.helpers import get_llvm_mc
|
||||
|
||||
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/main/llvm/test/MC/AMDGPU"
|
||||
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-21.1.0/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),
|
||||
}
|
||||
# RDNA3 (gfx11) test files for compute instructions
|
||||
# Excluded: gfx11_asm_mubuf.s, gfx11_asm_mtbuf.s, gfx11_asm_mimg.s, gfx11_asm_mubuf_alias.s, gfx11_asm_mtbuf_alias.s (graphics-only)
|
||||
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_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']
|
||||
# CDNA (gfx9/gfx90a/gfx942) test files for compute instructions
|
||||
# Excluded: gfx9_asm_mubuf.s, gfx9_asm_mtbuf.s, gfx90a_ldst_acc.s (has MIMG mixed in)
|
||||
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',
|
||||
'gfx90a_asm_features.s', 'flat-scratch-gfx942.s', 'gfx942_asm_features.s',
|
||||
'mai-gfx90a.s', 'mai-gfx942.s']
|
||||
# RDNA4 (gfx12) test files for compute instructions
|
||||
# Excluded: gfx12_asm_vbuffer_mubuf.s, gfx12_asm_vbuffer_mtbuf.s, gfx12_asm_exp.s (graphics-only)
|
||||
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_wmma_w32.s']
|
||||
|
||||
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)
|
||||
# Format 1: "// GFX11: v_foo ... ; encoding: [0x01,0x02,...]"
|
||||
# Format 2: "// GFX11: [0x01,0x02,...]" (used by DS, older files)
|
||||
if m := re.search(r'(?:GFX11|W32|W64)[^:]*:.*?encoding:\s*\[(.*?)\]', lines[j]):
|
||||
hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
|
||||
elif m := re.search(r'(?:GFX11|W32|W64)[^:]*:\s*\[(0x[0-9a-fA-F,x\s]+)\]', lines[j]):
|
||||
hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
|
||||
else:
|
||||
continue
|
||||
if hex_bytes:
|
||||
try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
|
||||
except ValueError: pass
|
||||
break
|
||||
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
|
||||
|
||||
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
|
||||
@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":
|
||||
# Match GFX11 and W32 only (wavefront32 mode)
|
||||
tests = _parse_llvm_tests(text, r'(?:GFX11|W32)')
|
||||
elif arch == "rdna4":
|
||||
# Match GFX12 (but not GFX1250) and W32 only (wavefront32 mode)
|
||||
tests = _parse_llvm_tests(text, r'(?:GFX12(?!50)|W32)')
|
||||
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 tests
|
||||
|
||||
def compile_asm_batch(instrs: list[str]) -> list[bytes]:
|
||||
"""Compile multiple instructions with a single llvm-mc call."""
|
||||
def _compile_asm_batch(instrs: list[str], arch: str = "rdna3") -> list[bytes]:
|
||||
if not instrs: return []
|
||||
asm_text = ".text\n" + "\n".join(instrs) + "\n"
|
||||
result = subprocess.run(
|
||||
[get_llvm_mc(), '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
|
||||
input=asm_text, capture_output=True, text=True, timeout=30)
|
||||
if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
|
||||
# Parse all encodings from output
|
||||
results = []
|
||||
for line in result.stdout.split('\n'):
|
||||
if 'encoding:' not in line: continue
|
||||
enc = line.split('encoding:')[1].strip()
|
||||
if enc.startswith('[') and enc.endswith(']'):
|
||||
results.append(bytes.fromhex(enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')))
|
||||
if len(results) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(results)}")
|
||||
return results
|
||||
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]
|
||||
|
||||
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 _make_test(f: str, arch: str, test_type: str):
|
||||
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)
|
||||
tests = _get_tests(f, arch)
|
||||
name = f"{arch}_{test_type}_{f}"
|
||||
if test_type == "roundtrip":
|
||||
passed, skipped = 0, 0
|
||||
for _, data in tests:
|
||||
try:
|
||||
decoded = detect_format(data, arch).from_bytes(data)
|
||||
self.assertEqual(decoded.to_bytes()[:len(data)], data)
|
||||
passed += 1
|
||||
except ValueError: skipped += 1 # skip invalid opcodes not in enum
|
||||
print(f"{name}: {passed} passed, {skipped} skipped")
|
||||
if arch in ("rdna3", "rdna4"):
|
||||
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
|
||||
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
|
||||
|
||||
def _make_disasm_test(name):
|
||||
def test(self):
|
||||
_, fmt_cls, op_enum = LLVM_TEST_FILES[name]
|
||||
# VOP3SD opcodes that share encoding with VOP3 (only for vop3sd test, not vopc promotions)
|
||||
vop3sd_opcodes = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
|
||||
is_vopc_promotion = name in ('vop3_from_vopc', 'vop3_from_vopcx')
|
||||
undocumented = {'smem': {34, 35}, 'sopk': {22, 23}, 'sopp': {8, 58, 59}}
|
||||
class TestLLVM(unittest.TestCase): pass
|
||||
|
||||
# First pass: decode all instructions and collect disasm strings
|
||||
to_test: list[tuple[str, bytes, str | None, str | None]] = [] # (asm_text, data, disasm_str, error)
|
||||
skipped = 0
|
||||
for asm_text, data in self.tests.get(name, []):
|
||||
if len(data) > fmt_cls._size(): continue
|
||||
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
|
||||
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}
|
||||
if temp_op in sopp_no_imm and simm16 != 0: skipped += 1; continue
|
||||
try:
|
||||
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)
|
||||
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)
|
||||
if decoded.to_bytes()[:len(data)] != data:
|
||||
to_test.append((asm_text, data, None, "decode roundtrip failed"))
|
||||
continue
|
||||
to_test.append((asm_text, data, decoded.disasm(), None))
|
||||
except Exception as e:
|
||||
to_test.append((asm_text, data, None, f"exception: {e}"))
|
||||
|
||||
# Batch compile all disasm strings with single llvm-mc call
|
||||
disasm_strs = [(i, t[2]) for i, t in enumerate(to_test) if t[2] is not None]
|
||||
llvm_results = compile_asm_batch([s for _, s in disasm_strs]) if disasm_strs else []
|
||||
llvm_map = {i: llvm_results[j] for j, (i, _) in enumerate(disasm_strs)}
|
||||
|
||||
# Match results back
|
||||
passed, failed = 0, 0
|
||||
failures: list[str] = []
|
||||
for idx, (asm_text, data, disasm_str, error) in enumerate(to_test):
|
||||
if error:
|
||||
failed += 1; failures.append(f"{error} for {data.hex()}")
|
||||
elif disasm_str is not None and idx in llvm_map:
|
||||
llvm_bytes = llvm_map[idx]
|
||||
if llvm_bytes is not None and llvm_bytes == data: passed += 1
|
||||
elif llvm_bytes is not None: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
|
||||
|
||||
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))
|
||||
for f in RDNA_FILES:
|
||||
setattr(TestLLVM, f"test_rdna3_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "roundtrip"))
|
||||
setattr(TestLLVM, f"test_rdna3_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "disasm"))
|
||||
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_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "disasm"))
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -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,8 +47,7 @@ dev.synchronize()
|
||||
elapsed = time.perf_counter() - st
|
||||
|
||||
self.assertNotEqual(result.returncode, 0, "should have raised")
|
||||
self.assertTrue("NotImplementedError" in result.stderr or "ValueError" in result.stderr,
|
||||
f"expected NotImplementedError or ValueError in 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, 9.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
|
||||
|
||||
|
||||
@@ -1,12 +1,16 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Tests for the RDNA3 pseudocode DSL."""
|
||||
import unittest
|
||||
from extra.assembly.amd.pcode import (Reg, TypedView, SliceProxy, ExecContext, compile_pseudocode, _expr, 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)
|
||||
from extra.assembly.amd.autogen.rdna3.gen_pcode import VOP3SDOp_FUNCTIONS, VOPCOp_FUNCTIONS
|
||||
from extra.assembly.amd.autogen.rdna3 import VOP3SDOp, VOPCOp
|
||||
from extra.assembly.amd.emu import _run_pcode
|
||||
from extra.assembly.amd.pcode import (Reg, TypedView, TypedView, MASK32, MASK64,
|
||||
_f32, _i32, _f16, _i16, f32_to_f16, isNAN, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
|
||||
BYTE_PERMUTE, v_sad_u8, v_msad_u8, _compile_pseudocode, _expr, compile_pseudocode)
|
||||
from extra.assembly.amd.test.helpers import ExecContext
|
||||
from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
|
||||
from extra.assembly.amd.autogen.rdna3.enum import VOP3SDOp, VOPCOp
|
||||
|
||||
# Compile pseudocode functions on demand for regression tests
|
||||
_VOP3SDOp_V_DIV_SCALE_F32 = compile_pseudocode('VOP3SDOp', 'V_DIV_SCALE_F32', PCODE[VOP3SDOp.V_DIV_SCALE_F32])
|
||||
_VOPCOp_V_CMP_CLASS_F32 = compile_pseudocode('VOPCOp', 'V_CMP_CLASS_F32', PCODE[VOPCOp.V_CMP_CLASS_F32_E32])
|
||||
|
||||
class TestReg(unittest.TestCase):
|
||||
def test_u32_read(self):
|
||||
@@ -42,7 +46,7 @@ class TestReg(unittest.TestCase):
|
||||
class TestTypedView(unittest.TestCase):
|
||||
def test_bit_slice(self):
|
||||
r = Reg(0xDEADBEEF)
|
||||
# Slices return SliceProxy which supports .u32, .u16 etc (matching pseudocode like S1.u32[1:0].u32)
|
||||
# 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)
|
||||
@@ -67,7 +71,7 @@ class TestTypedView(unittest.TestCase):
|
||||
# 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] # SliceProxy, int value = 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)
|
||||
@@ -85,7 +89,7 @@ class TestTypedView(unittest.TestCase):
|
||||
self.assertFalse(r1.u32 < r2.u32)
|
||||
self.assertTrue(r1.u32 != r2.u32)
|
||||
|
||||
class TestSliceProxy(unittest.TestCase):
|
||||
class TestTypedView(unittest.TestCase):
|
||||
def test_slice_read(self):
|
||||
r = Reg(0x56781234)
|
||||
self.assertEqual(r[15:0].u16, 0x1234)
|
||||
@@ -154,19 +158,19 @@ class TestExecContext(unittest.TestCase):
|
||||
self.assertEqual(ctx.SCC._val, 0)
|
||||
|
||||
def test_ternary(self):
|
||||
code = compile_pseudocode("D0.u32 = S0.u32 > S1.u32 ? 1'1U : 1'0U")
|
||||
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 }")
|
||||
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
|
||||
code = _compile_pseudocode("""tmp = S0.u32 + S1.u32
|
||||
D0.u32 = tmp.u32""")
|
||||
ctx = ExecContext(s0=100, s1=200)
|
||||
ctx.run(code)
|
||||
@@ -174,7 +178,7 @@ D0.u32 = tmp.u32""")
|
||||
|
||||
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)
|
||||
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
|
||||
@@ -184,7 +188,7 @@ D0.u32 = tmp.u32""")
|
||||
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)
|
||||
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)
|
||||
@@ -206,7 +210,7 @@ D0.u32 = tmp.u32""")
|
||||
|
||||
def test_for_loop(self):
|
||||
# CTZ pattern - find first set bit
|
||||
code = compile_pseudocode("""tmp = -1
|
||||
code = _compile_pseudocode("""tmp = -1
|
||||
for i in 0 : 31 do
|
||||
if S0.u32[i] == 1 then
|
||||
tmp = i
|
||||
@@ -229,49 +233,47 @@ 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_lane, even when VCC=0 (no scaling needed).
|
||||
Bug: when VCC._val == vcc (both 0), vcc_lane wasn't returned, so VCC bits weren't written.
|
||||
"""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)
|
||||
fn = VOP3SDOp_FUNCTIONS[VOP3SDOp.V_DIV_SCALE_F32]
|
||||
result = _run_pcode(fn, VOP3SDOp, VOP3SDOp.V_DIV_SCALE_F32, s0, s1, s2, 0, 0, 0, 0, 0xffffffff, 0)
|
||||
# Must always have vcc_lane in result
|
||||
self.assertIn('vcc_lane', result, "V_DIV_SCALE_F32 must always return vcc_lane")
|
||||
self.assertEqual(result['vcc_lane'], 0, "vcc_lane should be 0 when no scaling needed")
|
||||
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
|
||||
fn = VOPCOp_FUNCTIONS[VOPCOp.V_CMP_CLASS_F32]
|
||||
# Test quiet NaN detection (bit 1 in mask)
|
||||
s1_quiet = 0b0000000010 # bit 1 = quiet NaN
|
||||
result = _run_pcode(fn, VOPCOp, VOPCOp.V_CMP_CLASS_F32, quiet_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0)
|
||||
self.assertEqual(result['vcc_lane'], 1, "Should detect quiet NaN with quiet NaN mask")
|
||||
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 = _run_pcode(fn, VOPCOp, VOPCOp.V_CMP_CLASS_F32, signal_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0)
|
||||
self.assertEqual(result['vcc_lane'], 1, "Should detect signaling NaN with signaling NaN mask")
|
||||
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 = _run_pcode(fn, VOPCOp, VOPCOp.V_CMP_CLASS_F32, quiet_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0)
|
||||
self.assertEqual(result['vcc_lane'], 0, "Quiet NaN should not 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 = _run_pcode(fn, VOPCOp, VOPCOp.V_CMP_CLASS_F32, signal_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0)
|
||||
self.assertEqual(result['vcc_lane'], 0, "Signaling NaN should not 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 test_isnan_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."""
|
||||
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")
|
||||
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."""
|
||||
@@ -310,7 +312,7 @@ class TestBF16(unittest.TestCase):
|
||||
self.assertAlmostEqual(float(r.bf16), 3.0, places=1)
|
||||
|
||||
def test_bf16_slice_property(self):
|
||||
"""Test SliceProxy.bf16 property."""
|
||||
"""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)
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test PDF pseudocode extraction from amdxml.py."""
|
||||
import unittest
|
||||
from extra.assembly.amd.amdxml import extract_pdf_text, extract_pcode, parse_xml, ARCHS, FIXES
|
||||
|
||||
EXPECTED_PAGES = {"rdna3": 655, "rdna4": 711, "cdna": 610}
|
||||
|
||||
class TestPcodePDF(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.pages = {arch: extract_pdf_text(cfg["pdf"]) for arch, cfg in ARCHS.items()}
|
||||
cls.enums = {}
|
||||
for arch, cfg in ARCHS.items():
|
||||
_, enums, _, _, _ = parse_xml(cfg["xml"], arch)
|
||||
for fmt, ops in FIXES.get(arch, {}).items(): enums.setdefault(fmt, {}).update(ops)
|
||||
cls.enums[arch] = enums
|
||||
cls.pcode = {arch: extract_pcode(cls.pages[arch], {n: op for ops in cls.enums[arch].values() for op, n in ops.items()}) for arch in ARCHS}
|
||||
|
||||
def test_page_counts(self):
|
||||
for name, exp in EXPECTED_PAGES.items():
|
||||
self.assertEqual(len(self.pages[name]), exp, f"{name} page count")
|
||||
|
||||
def test_pcode_extracted(self):
|
||||
"""Check we extracted a reasonable number of pcode entries."""
|
||||
for name in ARCHS:
|
||||
self.assertGreater(len(self.pcode[name]), 500, f"{name} pcode count too low")
|
||||
|
||||
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 ARCHS:
|
||||
for (op_name, opcode), code in self.pcode[name].items():
|
||||
self.assertNotIn('=>', code, f"{name} {op_name} contains example line with '=>'")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,153 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test that PDF parser correctly extracts format fields."""
|
||||
import unittest, os
|
||||
from extra.assembly.amd.autogen.rdna3 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),
|
||||
}
|
||||
|
||||
# Skip PDF parsing tests by default - only run with TEST_PDF_PARSER=1
|
||||
# These are slow (~5s) and only needed when regenerating autogen/
|
||||
@unittest.skipUnless(os.environ.get("TEST_PDF_PARSER"), "set TEST_PDF_PARSER=1 to run PDF parser tests")
|
||||
class TestPDFParserGenerate(unittest.TestCase):
|
||||
"""Test the PDF parser by running generate() and checking results."""
|
||||
|
||||
def test_pdf_parser(self):
|
||||
"""Single test that validates all PDF parser outputs."""
|
||||
from extra.assembly.amd.dsl import generate
|
||||
result = generate()
|
||||
|
||||
# test_all_formats_present
|
||||
for fmt_name in EXPECTED_FORMATS:
|
||||
self.assertIn(fmt_name, result["formats"], f"missing format {fmt_name}")
|
||||
|
||||
# test_format_count
|
||||
self.assertEqual(len(result["formats"]), 23)
|
||||
|
||||
# test_no_duplicate_fields
|
||||
for fmt_name, fields in 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}")
|
||||
|
||||
# test_expected_fields
|
||||
for fmt_name, (expected_fields, has_encoding) in EXPECTED_FORMATS.items():
|
||||
fields = {f[0] for f in 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")
|
||||
|
||||
# test_vopd_no_dpp16_fields
|
||||
vopd_fields = {f[0] for f in result["formats"].get("VOPD", [])}
|
||||
for field in ['DPP_CTRL', 'BANK_MASK', 'ROW_MASK']:
|
||||
self.assertNotIn(field, vopd_fields, f"VOPD should not have {field}")
|
||||
|
||||
# test_dpp16_no_vinterp_fields
|
||||
dpp16_fields = {f[0] for f in result["formats"].get("DPP16", [])}
|
||||
for field in ['VDST', 'WAITEXP']:
|
||||
self.assertNotIn(field, dpp16_fields, f"DPP16 should not have {field}")
|
||||
|
||||
# test_sopp_no_smem_fields
|
||||
sopp_fields = {f[0] for f in 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,6 +1,6 @@
|
||||
#!/usr/bin/env python3
|
||||
import unittest, subprocess
|
||||
from extra.assembly.amd.autogen.rdna3 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:
|
||||
@@ -37,10 +37,10 @@ class TestRDNA3Asm(unittest.TestCase):
|
||||
v_add_co_u32(v[2], VCC_LO, s[6], v[2]),
|
||||
v_add_co_ci_u32_e32(v[3], s[7], v[3]),
|
||||
v_add_co_u32(v[0], VCC_LO, s[4], v[0]),
|
||||
global_load_b32(vdst=v[2], addr=v[2], saddr=OFF),
|
||||
global_load_b32(vdst=v[2], addr=v[2:3], saddr=OFF),
|
||||
v_add_co_ci_u32_e32(v[1], s[5], v[1]),
|
||||
s_waitcnt(0x03f7), # vmcnt(0)
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=OFF),
|
||||
global_store_b32(addr=v[0:1], data=v[2], saddr=OFF),
|
||||
s_endpgm(),
|
||||
]
|
||||
|
||||
|
||||
@@ -1,12 +1,17 @@
|
||||
#!/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.autogen.rdna3 import *
|
||||
from extra.assembly.amd.dsl import Inst
|
||||
from extra.assembly.amd.asm import asm
|
||||
from extra.assembly.amd.asm import detect_format
|
||||
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
|
||||
@@ -30,61 +35,40 @@ def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
|
||||
continue
|
||||
return results
|
||||
|
||||
def compile_asm(instr: str, compiler=None) -> bytes:
|
||||
"""Compile a single instruction with llvm-mc and return the machine code bytes."""
|
||||
llvm_mc = get_llvm_mc()
|
||||
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: raise RuntimeError(f"llvm-mc failed for '{instr}': {result.stderr.strip()}")
|
||||
# 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)
|
||||
raise RuntimeError(f"no encoding found in llvm-mc output for: {instr}")
|
||||
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]) -> list[bytes]:
|
||||
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 []
|
||||
llvm_mc = get_llvm_mc()
|
||||
src = ".text\n" + "\n".join(instrs) + "\n"
|
||||
result = subprocess.run(
|
||||
[llvm_mc, '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
|
||||
input=src, capture_output=True, text=True)
|
||||
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()}")
|
||||
# Parse all encodings in order
|
||||
encodings = []
|
||||
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(' ', '')
|
||||
encodings.append(bytes.fromhex(hex_vals))
|
||||
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], compiler) -> list[str]:
|
||||
def compile_and_disasm_batch(instrs: list[str], arch: str = 'rdna3') -> list[str]:
|
||||
"""Compile instructions with LLVM and get LLVM's disassembly."""
|
||||
import tempfile, os
|
||||
import tempfile
|
||||
if not instrs: return []
|
||||
# Build assembly source with all instructions
|
||||
src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n"
|
||||
src += "\n".join(f" {instr}" for instr in instrs) + "\n"
|
||||
# Use llvm-mc to assemble to object file
|
||||
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', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-filetype=obj', '-o', obj_path],
|
||||
input=src, capture_output=True, text=True)
|
||||
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()}")
|
||||
# Disassemble with llvm-objdump
|
||||
result = subprocess.run([get_llvm_objdump(), '-d', '--mcpu=gfx1100', obj_path], capture_output=True, text=True)
|
||||
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()}")
|
||||
# Parse disassembly output
|
||||
results: list[str] = []
|
||||
for line in result.stdout.splitlines():
|
||||
if '//' not in line: continue
|
||||
@@ -96,18 +80,22 @@ def compile_and_disasm_batch(instrs: list[str], compiler) -> list[str]:
|
||||
|
||||
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
|
||||
2. disasm() -> LLVM asm -> bytes matches original (validates disasm correctness)
|
||||
3. our disasm() matches LLVM's disassembly string (informational)
|
||||
"""
|
||||
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('gfx1100')
|
||||
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)
|
||||
@@ -115,7 +103,7 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
|
||||
offset = 0
|
||||
while offset < len(kernel.code):
|
||||
remaining = kernel.code[offset:]
|
||||
fmt = detect_format(remaining)
|
||||
fmt = detect_format(remaining, arch)
|
||||
if fmt is None:
|
||||
decoded_instrs.append((ki, offset, None, None, None, False, "no format"))
|
||||
offset += 4
|
||||
@@ -141,22 +129,22 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
|
||||
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
|
||||
asm_test_instrs: list[tuple[int, str, bytes]] = [] # (idx, our_disasm, orig_bytes) 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
|
||||
# Skip unknown opcodes and malformed instructions
|
||||
if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm): continue
|
||||
asm_test_instrs.append((idx, our_disasm))
|
||||
asm_test_instrs.append((idx, our_disasm, orig_bytes))
|
||||
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])
|
||||
asm_llvm_map = {idx: result for (idx, _), result in zip(asm_test_instrs, asm_llvm_results)}
|
||||
# Batch compile for asm test (our disasm -> LLVM asm -> bytes)
|
||||
asm_llvm_results = compile_asm_batch([d for _, d, _ in asm_test_instrs], arch)
|
||||
asm_llvm_map = {idx: (result, orig) for (idx, _, orig), 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], compiler)
|
||||
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
|
||||
@@ -177,20 +165,16 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
|
||||
decode_failed += 1
|
||||
decode_failures.append(f"K{ki}@{offset}: {our_disasm}: {decode_err}")
|
||||
|
||||
# Asm test
|
||||
# Asm test: our disasm -> LLVM asm -> compare bytes with original
|
||||
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
|
||||
llvm_bytes, orig = asm_llvm_map[idx]
|
||||
if llvm_bytes == orig[:len(llvm_bytes)]:
|
||||
asm_passed += 1
|
||||
else:
|
||||
asm_failed += 1
|
||||
asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': llvm={llvm_bytes.hex()} orig={orig[:len(llvm_bytes)].hex()}")
|
||||
else:
|
||||
asm_skipped += 1
|
||||
|
||||
@@ -207,9 +191,9 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
|
||||
else:
|
||||
disasm_skipped += 1
|
||||
|
||||
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")
|
||||
print(f"[{arch}] decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
|
||||
print(f"[{arch}] asm via llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
|
||||
print(f"[{arch}] disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
|
||||
self.assertEqual(decode_failed, 0, 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
|
||||
@@ -259,5 +243,11 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
|
||||
# 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("RDNA4 decode roundtrip not yet supported")
|
||||
class TestTinygradKernelRoundtripRDNA4(TestTinygradKernelRoundtrip): arch = 'rdna4'
|
||||
|
||||
@unittest.skip("CDNA decode roundtrip not yet supported")
|
||||
class TestTinygradKernelRoundtripCDNA(TestTinygradKernelRoundtrip): arch = 'cdna'
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,199 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Tests for SQTT packet decoding using real captured examples."""
|
||||
import pickle, unittest, ctypes, threading
|
||||
from pathlib import Path
|
||||
from tinygrad.helpers import DEBUG
|
||||
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, print_packets)
|
||||
|
||||
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}
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# 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(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
|
||||
|
||||
exc = None
|
||||
def worker():
|
||||
nonlocal exc
|
||||
try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
|
||||
except Exception as e: exc = e
|
||||
(t:=threading.Thread(target=worker, daemon=True)).start()
|
||||
t.join(timeout=1)
|
||||
if exc is not None: raise exc
|
||||
if t.is_alive(): raise RuntimeError("rocprof decoder timeout")
|
||||
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 = list(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_packet_counts(self):
|
||||
expected = {
|
||||
"profile_empty_run_0": [559, 600],
|
||||
"profile_empty_run_1": [517, 570],
|
||||
"profile_gemm_run_0": [1489, 604, 1789, 466, 17570, 407],
|
||||
"profile_gemm_run_1": [1453, 604, 1871, 493, 17827, 460],
|
||||
"profile_plus_run_0": [695, 668],
|
||||
"profile_plus_run_1": [663, 593],
|
||||
}
|
||||
for name, (events, *_) in self.examples.items():
|
||||
with self.subTest(example=name):
|
||||
counts = [len(list(decode(e.blob))) for e in events]
|
||||
self.assertEqual(counts, expected[name], f"packet count mismatch 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:
|
||||
wave_starts: dict[tuple[int, int, int], int] = {}
|
||||
for p in decode(event.blob):
|
||||
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()
|
||||
@@ -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,684 @@
|
||||
# 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.uop.ops import UOp, Ops, KernelInfo
|
||||
from tinygrad.helpers import getenv, colored
|
||||
from tinygrad.engine.realize import Estimates
|
||||
from extra.assembly.amd.dsl import s, v, VCC_LO, NULL
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
|
||||
# =============================================================================
|
||||
# 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 = 182 # lane_id & 7 (column within 8-wide tile chunk)
|
||||
V_OUTPUT_ROW = 171 # output row coordinate
|
||||
V_LANE_MOD8_X4 = 174 # V_LANE_ID_MOD8 << 2 (byte offset)
|
||||
V_LANE_DIV8_X4 = 175 # (lane_id >> 3) << 2
|
||||
V_ADDR_HI_ZERO = 188 # always 0 (for 64-bit address high bits)
|
||||
V_LDS_A_BASE = 186 # LDS A-tile base address for inner loop
|
||||
V_LDS_B_BASE = 170 # LDS B-tile base address for inner loop
|
||||
V_GLOBAL_A_ADDR = 171 # global memory A prefetch address (reuses V_OUTPUT_ROW slot during main loop)
|
||||
V_GLOBAL_B_ADDR = 178 # global memory B prefetch address
|
||||
|
||||
# LDS tile register destinations - SEPARATE from DATA to avoid overlap
|
||||
# A on banks 2-3, B on banks 0-1 to avoid bank conflicts in VOPD
|
||||
# Double-buffered: buffer 0 (v130-v161), buffer 1 (v194-v225)
|
||||
V_A_TILE_REGS = [[130, 134, 138, 142], [194, 198, 202, 206]] # A tile: banks 2,2,2,2 (mod 4 = 2)
|
||||
V_B_TILE_REGS = [[132, 136, 140, 144, 148, 152, 156, 160], [192, 196, 200, 204, 208, 212, 216, 220]] # B tile: banks 0,0,0,0 (mod 4 = 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, b_tile_regs):
|
||||
"""Generate 64 dual FMAC ops from accumulator grid with optimized iteration order."""
|
||||
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 for each buffer
|
||||
FMAC_PATTERN = [derive_fmac_pattern(ACC_GRID, V_A_TILE_REGS[i], V_B_TILE_REGS[i]) for i in range(2)]
|
||||
|
||||
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 = 189 # single base register for A stores (use +512 offsets)
|
||||
V_LDS_A_DATA = [155, 172, 173, 154, 159, 176, 177, 158] # 8 data registers for A prefetch (mod 4: 3,0,1,2,3,0,1,2)
|
||||
V_LDS_B_ADDR = 190 # single base register for B stores (use 16-bit offsets)
|
||||
V_LDS_B_DATA = [163, 180, 181, 162, 167, 184, 185, 166] # 8 data registers for B prefetch (mod 4: 3,0,1,2,3,0,1,2)
|
||||
|
||||
# Global memory prefetch schedule: (vdst1, vdst2, addr_vreg, saddr_lo1, saddr_lo2)
|
||||
# First 2 pairs from B prefetch pointers (s[32:39]), next 4 pairs from A prefetch pointers (s[40:55])
|
||||
PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR, S_PREFETCH_B+8+4*i, S_PREFETCH_B+10+4*i) for i in range(2)] + \
|
||||
[(V_LDS_B_DATA[2*(i-2)], V_LDS_B_DATA[2*(i-2)+1], V_GLOBAL_A_ADDR, S_PREFETCH_A+4*(i-2), S_PREFETCH_A+2+4*(i-2)) for i in range(2, 6)]
|
||||
|
||||
# Initial tile prefetch: (vdst, saddr_lo) - load into A data regs using B prefetch pointers (s[24:31])
|
||||
INIT_PREFETCH = [(V_LDS_A_DATA[i], S_PREFETCH_B+2*i) for i in range(4)]
|
||||
|
||||
# Initial tile loads: (vdst, addr_lo) pairs - use temp regs in accumulator gaps
|
||||
INIT_TILE_LOADS = [(23,5),(24,9),(25,7),(26,2),(27,11),(28,13),(29,6),(30,8),(31,10),(12,12),(13,14),(3,2),(4,4),(5,8),(6,6),(7,10)]
|
||||
|
||||
# A matrix row offset registers (scattered to avoid accumulator conflicts)
|
||||
ROW_REGS = [165, 146, 147, 164, 169, 150, 151, 168] # mod 4: 1,2,3,0,1,2,3,0
|
||||
|
||||
# =============================================================================
|
||||
# Kernel class
|
||||
# =============================================================================
|
||||
|
||||
class Kernel:
|
||||
def __init__(self, arch='gfx1100'):
|
||||
self.instructions, self.labels, self.branch_targets, self.arch = [], {}, {}, arch
|
||||
|
||||
def 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 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."""
|
||||
vmcnt, lgkmcnt, expcnt = vm if vm is not None else 63, lgkm if lgkm is not None else 63, 7
|
||||
waitcnt = (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
|
||||
self.emit(s_waitcnt(simm16=waitcnt))
|
||||
|
||||
def to_asm(self):
|
||||
import re
|
||||
# Instruction stream with labels
|
||||
label_at = {pos: name for name, pos in self.labels.items()}
|
||||
body = []
|
||||
for i, inst in enumerate(self.instructions):
|
||||
if i in label_at: body.append(f'.{label_at[i]}:')
|
||||
asm = inst.disasm()
|
||||
if i in self.branch_targets:
|
||||
asm = re.sub(r'(s_cbranch_\w+|s_branch)\s+\S+', rf'\1 .{self.branch_targets[i]}', asm)
|
||||
body.append('\t' + asm)
|
||||
|
||||
# limit wave occupancy by using more LDS
|
||||
lds_size = max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536))
|
||||
|
||||
# HSA kernel descriptor attributes (zeros included for compatibility)
|
||||
hsa = [
|
||||
('group_segment_fixed_size', lds_size), ('private_segment_fixed_size', 0), ('kernarg_size', 36),
|
||||
('user_sgpr_count', 14), ('user_sgpr_dispatch_ptr', 0), ('user_sgpr_queue_ptr', 0),
|
||||
('user_sgpr_kernarg_segment_ptr', 1), ('user_sgpr_dispatch_id', 0), ('user_sgpr_private_segment_size', 0),
|
||||
('wavefront_size32', 1), ('uses_dynamic_stack', 0), ('enable_private_segment', 0),
|
||||
('system_sgpr_workgroup_id_x', 1), ('system_sgpr_workgroup_id_y', 1), ('system_sgpr_workgroup_id_z', 0),
|
||||
('system_sgpr_workgroup_info', 0), ('system_vgpr_workitem_id', 0), ('next_free_vgpr', 222),
|
||||
('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: 222', ' .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=NULL))
|
||||
k.emit(s_load_b64(sdata=s[S_KERNARG_OUT[0]:S_KERNARG_OUT[1]], sbase=s[0:1], offset=0x10, soffset=NULL))
|
||||
k.emit(s_mov_b32(s[S_DIM_N], MATRIX_DIM))
|
||||
k.emit(s_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
|
||||
k.emit(s_mov_b64(s[S_PREFETCH_B:S_PREFETCH_B+1], s[S_KERNARG_B[0]:S_KERNARG_B[1]])) # B[0]: no offset
|
||||
for i in range(1, 8): # B: 16KB apart
|
||||
k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_KERNARG_B[0]], i * 0x4000))
|
||||
k.emit(s_addc_u32(s[S_PREFETCH_B+i*2+1], s[S_KERNARG_B[1]], 0))
|
||||
k.emit(s_mov_b64(s[S_PREFETCH_A:S_PREFETCH_A+1], s[S_KERNARG_A[0]:S_KERNARG_A[1]])) # A[0]: no offset
|
||||
for i in range(1, 8): # A: 256KB apart
|
||||
k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_KERNARG_A[0]], i * 0x40000))
|
||||
k.emit(s_addc_u32(s[S_PREFETCH_A+i*2+1], s[S_KERNARG_A[1]], 0))
|
||||
|
||||
# Global prefetch addresses: B = (tile_x + lane_id) * 4, A = ((tile_y << 12) + (lane_id/8)*4K + lane_id%8) * 4
|
||||
k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], s[S_TILE_X], v[0]))
|
||||
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.emit(global_load_b32(vdst=v[vdst], addr=v[addr:addr+1], saddr=NULL))
|
||||
|
||||
# Batch 2: A addresses (rows 0-4) and loads
|
||||
for dst, ri in [(6,0), (8,1), (10,2), (12,3), (14,4)]:
|
||||
k.emit(v_add_nc_u32_e32(v[dst], v[ROW_REGS[ri]], v[V_LANE_ID_MOD8]))
|
||||
addr64(dst, S_A_PTR[0])
|
||||
k.emit(s_clause(simm16=4)) # 5 consecutive global loads
|
||||
for vdst, addr in INIT_TILE_LOADS[6:11]: k.emit(global_load_b32(vdst=v[vdst], addr=v[addr:addr+1], saddr=NULL))
|
||||
|
||||
# Batch 3: B cols 6-7, A rows 5-7, and loads
|
||||
for dst, mult, tmp in [(2,6,15), (4,7,4)]: b_addr(dst, mult, tmp)
|
||||
for dst, ri, tmp in [(8,5,16), (6,6,18), (10,7,20)]: a_addr(dst, ROW_REGS[ri], tmp)
|
||||
k.emit(s_clause(simm16=4)) # 5 consecutive global loads
|
||||
for vdst, addr in INIT_TILE_LOADS[11:]: k.emit(global_load_b32(vdst=v[vdst], addr=v[addr:addr+1], saddr=NULL))
|
||||
|
||||
# ===========================================================================
|
||||
# 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))
|
||||
|
||||
# ===========================================================================
|
||||
# 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_ALU, NO_DS, NO_GLOBAL = getenv("NO_ALU", 0), 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
|
||||
|
||||
if not NO_GLOBAL:
|
||||
# Advance prefetch pointers (VGPR)
|
||||
k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], 0x20000, v[V_GLOBAL_B_ADDR]))
|
||||
k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], 0x20, v[V_GLOBAL_A_ADDR]))
|
||||
|
||||
"""
|
||||
# Advance prefetch pointers (SGPRs, 64-bit adds)
|
||||
k.emit(s_clause(simm16=31))
|
||||
for i in range(8):
|
||||
k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_PREFETCH_B+i*2], 0x20000))
|
||||
k.emit(s_addc_u32(s[S_PREFETCH_B+i*2+1], s[S_PREFETCH_B+i*2+1], 0))
|
||||
for i in range(8):
|
||||
k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_PREFETCH_A+i*2], 0x20))
|
||||
k.emit(s_addc_u32(s[S_PREFETCH_A+i*2+1], s[S_PREFETCH_A+i*2+1], 0))
|
||||
"""
|
||||
|
||||
# do the fetch
|
||||
for vdst, saddr_lo in INIT_PREFETCH:
|
||||
k.emit(global_load_b32(vdst=v[vdst], addr=v[V_GLOBAL_B_ADDR], saddr=s[saddr_lo:saddr_lo+1]))
|
||||
|
||||
k.label('SKIP_PREFETCH')
|
||||
|
||||
# wait for local stores to finish (either initial or loop)
|
||||
# then sync the warp so it's safe to load local
|
||||
k.waitcnt(lgkm=0)
|
||||
k.emit(s_barrier())
|
||||
|
||||
# Load initial tiles for iter=0 into buffer 0
|
||||
if not NO_DS:
|
||||
a_tile_regs = V_A_TILE_REGS[0]
|
||||
b_tile_regs = V_B_TILE_REGS[0]
|
||||
k.emit(s_clause(simm16=len(a_tile_regs) + len(b_tile_regs) - 1))
|
||||
for i, vdst in enumerate(a_tile_regs):
|
||||
a_off = (i & 1) * 8 + (i >> 1) * 64 # iter=0
|
||||
k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_A_BASE], offset0=a_off & 0xFF, offset1=a_off >> 8))
|
||||
for i, vdst in enumerate(b_tile_regs):
|
||||
b_off = (i & 1) * 8 + (i & 2) * 64 + (i >> 2) * 256 # iter=0
|
||||
k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_B_BASE], offset0=b_off & 0xFF, offset1=b_off >> 8))
|
||||
|
||||
# 8 inner loop iterations
|
||||
# Double-buffered inner loop: load next iteration's tiles while computing current
|
||||
# Buffer 0 used for even iterations, buffer 1 for odd iterations
|
||||
for iter in range(8):
|
||||
buf = iter & 1 # current compute buffer
|
||||
next_buf = 1 - buf # next load buffer
|
||||
|
||||
# Load tiles for NEXT iteration into next_buf (except on last iteration)
|
||||
if not NO_DS and iter < 7:
|
||||
next_iter = iter + 1
|
||||
a_tile_regs = V_A_TILE_REGS[next_buf]
|
||||
b_tile_regs = V_B_TILE_REGS[next_buf]
|
||||
k.emit(s_clause(simm16=len(a_tile_regs) + len(b_tile_regs) - 1)) # 12 loads total: 4 A + 8 B
|
||||
# A tile: 4 ds_load_b64
|
||||
for i, vdst in enumerate(a_tile_regs):
|
||||
a_off = (i & 1) * 8 + (i >> 1) * 64 + next_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(b_tile_regs):
|
||||
b_off = (i & 1) * 8 + (i & 2) * 64 + (i >> 2) * 256 + next_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))
|
||||
|
||||
# Issue global prefetch (first 6 iterations only)
|
||||
if iter < 6 and not NO_GLOBAL:
|
||||
vdst1, vdst2, addr, slo1, slo2 = PREFETCH_LOADS[iter]
|
||||
k.emit(global_load_b32(vdst=v[vdst1], addr=v[addr], saddr=s[slo1:slo1+1]))
|
||||
k.emit(global_load_b32(vdst=v[vdst2], addr=v[addr], saddr=s[slo2:slo2+1]))
|
||||
|
||||
# Wait for current buffer's loads to complete
|
||||
# iter 0-6: 12 loads for next iteration are in flight, wait for the other 12 (lgkm=12)
|
||||
# iter 7: no loads in flight, wait for all (lgkm=0)
|
||||
k.waitcnt(lgkm=0 if iter == 7 else 12)
|
||||
|
||||
# 64 dual FMACs using current buffer
|
||||
if not NO_ALU:
|
||||
fmac_pattern = FMAC_PATTERN[buf]
|
||||
k.emit(s_clause(simm16=len(fmac_pattern)-1))
|
||||
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]))
|
||||
|
||||
# wait for all global stores to finish
|
||||
# then sync the warp so it's safe to store local
|
||||
k.waitcnt(vm=0)
|
||||
k.emit(s_barrier())
|
||||
|
||||
# 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
|
||||
if not NO_DS:
|
||||
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.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=NULL))
|
||||
|
||||
k.emit(s_sendmsg(simm16=3)) # DEALLOC_VGPRS
|
||||
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)
|
||||
|
||||
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}")
|
||||
|
||||
dname:str = Device.DEFAULT
|
||||
def asm_kernel(A:UOp, B:UOp, C:UOp) -> UOp:
|
||||
gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
|
||||
lidxs = [UOp.special(n, f"lidx{i}") for i,n in enumerate(local)]
|
||||
sink = UOp.sink(A.base, B.base, C.base, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
|
||||
estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
|
||||
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=asm),
|
||||
UOp(Ops.BINARY, arg=binary)))
|
||||
c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
|
||||
ei = c.schedule()[0].lower()
|
||||
|
||||
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 != err or err > 1e-06: raise RuntimeError("matmul is wrong!")
|
||||
|
||||
def run_sqtt():
|
||||
"""Run with SQTT profiling and write trace files."""
|
||||
import subprocess, os
|
||||
|
||||
# Run test_matmul in a subprocess with SQTT enabled from the start (no verify)
|
||||
env = {**os.environ, "AMD": "1", "SQTT": "1", "CNT": "1", "PROFILE": "1", "PYTHONPATH": ".", "VERIFY": "0"}
|
||||
result = subprocess.run(
|
||||
["python", "-c", "from extra.gemm.amd_asm_matmul import test_matmul; test_matmul()"],
|
||||
capture_output=True, text=True, env=env, timeout=120
|
||||
)
|
||||
print(result.stdout)
|
||||
|
||||
# Run roc.py to extract trace data
|
||||
result = subprocess.run(
|
||||
["python", "extra/sqtt/roc.py", "--profile", "/tmp/profile.pkl.tiny", "--kernel", "kernel"],
|
||||
capture_output=True, text=True, env={**os.environ, "DEBUG": "5"}, timeout=60
|
||||
)
|
||||
output = result.stdout + result.stderr
|
||||
|
||||
# Write full output to trace file
|
||||
with open("/tmp/sqtt_trace.txt", "w") as f:
|
||||
f.write(output)
|
||||
print(f"Wrote {len(output)} bytes to /tmp/sqtt_trace.txt")
|
||||
|
||||
if __name__ == "__main__":
|
||||
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))
|
||||
|
||||
@@ -1,16 +1,9 @@
|
||||
.text
|
||||
.section .text.
|
||||
.global gemm
|
||||
.p2align 8
|
||||
.type gemm,@function
|
||||
|
||||
gemm:
|
||||
// ** global buffers
|
||||
s_load_dwordx2 s[28:29], s[0:1], 0x0 // C
|
||||
s_load_dwordx4 s[32:35], s[0:1], 0x8 // A, B
|
||||
s_load_dwordx2 s[34:35], s[0:1], 0x08 // A
|
||||
s_load_dwordx2 s[32:33], s[0:1], 0x10 // B
|
||||
// ** others kernel args
|
||||
s_load_dword s24, s[0:1], 0x18 // N
|
||||
s_load_dword s54, s[0:1], 0x1C // num work groups
|
||||
s_waitcnt lgkmcnt(0)
|
||||
// "info"
|
||||
s_mov_b32 s51, 1 // gemm_info = 1
|
||||
@@ -221,11 +214,11 @@ gemm:
|
||||
s_mul_hi_u32 s87, s86, s40 // 000000003288: 96572856
|
||||
s_mul_i32 s86, s86, s40 // 00000000328C: 92562856
|
||||
s_and_b32 s84, s50, 0x8000 // 000000003290: 8654FF32 00008000
|
||||
s_cbranch_scc1 label_GSUC_A // 000000003298: BF850003
|
||||
s_cbranch_scc1 skip_offset_A // 000000003298: BF850003
|
||||
s_mul_hi_u32 s85, 64, s6 // 00000000329C: 965506C0
|
||||
s_mul_i32 s84, 64, s6 // 0000000032A0: 925406C0
|
||||
|
||||
label_GSUC_A:
|
||||
skip_offset_A:
|
||||
s_add_u32 s86, s86, s84 // 000000003330: 80565456
|
||||
s_addc_u32 s87, s87, s85 // 000000003334: 82575557
|
||||
s_mov_b64 s[60:61], 1 // 000000003338: BEBC0181
|
||||
@@ -259,11 +252,11 @@ label_GSUC_A:
|
||||
s_mul_hi_u32 s87, s86, s42 // 0000000033B4: 96572A56
|
||||
s_mul_i32 s86, s86, s42 // 0000000033B8: 92562A56
|
||||
s_and_b32 s84, s50, 0x8000 // 0000000033BC: 8654FF32 00008000
|
||||
s_cbranch_scc1 label_GSUC_B // 0000000033C4: BF850003
|
||||
s_cbranch_scc1 skip_offset_B // 0000000033C4: BF850003
|
||||
s_mul_hi_u32 s85, 64, s6 // 0000000033C8: 965506C0
|
||||
s_mul_i32 s84, 64, s6 // 0000000033CC: 925406C0
|
||||
|
||||
label_GSUC_B:
|
||||
skip_offset_B:
|
||||
s_add_u32 s86, s86, s84 // 00000000345C: 80565456
|
||||
s_addc_u32 s87, s87, s85 // 000000003460: 82575557
|
||||
s_mov_b64 s[62:63], 1 // 000000003464: BEBE0181
|
||||
@@ -308,8 +301,6 @@ label_GSUC_B:
|
||||
s_and_b32 s87, s10, 0xe000 // 0000000035A4: 8657FF0A 0000E000
|
||||
s_and_b32 s10, s10, 0xff // 0000000035AC: 860AFF0A 000000FF
|
||||
s_mov_b32 s84, s10 // 0000000035B4: BED4000A
|
||||
|
||||
label_beginStaggerUIter:
|
||||
s_lshl_b32 s85, s84, s86 // 0000000035B8: 8E555654
|
||||
s_cmp_ge_u32 s13, s85 // 0000000035BC: BF09550D
|
||||
s_sub_u32 s85, s84, 1 // 0000000035CC: 80D58154
|
||||
@@ -344,7 +335,7 @@ label_beginStaggerUIter:
|
||||
s_cselect_b32 s58, s62, -1 // 0000000036A4: 853AC13E
|
||||
s_add_u32 s51, s51, 2 // 0000000036A8: 80338233
|
||||
s_cmp_eq_u32 s12, 0 // 0000000036AC: BF06800C
|
||||
s_cbranch_scc1 label_ShadowInitStart // 0000000036B0: BF850092
|
||||
s_cbranch_scc1 init_output_buffers // 0000000036B0: BF850092
|
||||
s_mov_b32 m0, s46 // 0000000036B4: BEFC002E
|
||||
buffer_load_dwordx4 v0, s[52:55], 0 offen lds // 0000000036B8: E05D1000 800D0000
|
||||
s_add_u32 m0, m0, 0x1040 // 0000000036C0: 807CFF7C 00001040
|
||||
@@ -431,7 +422,7 @@ label_beginStaggerUIter:
|
||||
s_cmp_eq_u32 s63, 0 // 0000000038F4: BF06803F
|
||||
s_cselect_b32 s58, s62, -1 // 0000000038F8: 853AC13E
|
||||
|
||||
label_ShadowInitStart:
|
||||
init_output_buffers:
|
||||
s_mov_b64 s[16:17], s[28:29] // 0000000038FC: BE90011C
|
||||
s_mov_b32 s18, 0x80000000 // 000000003900: BE9200FF 80000000
|
||||
s_mov_b32 s19, 0x20000 // 000000003908: BE9300FF 00020000
|
||||
@@ -476,12 +467,10 @@ label_ShadowInitStart:
|
||||
s_lshl_b64 s[84:85], s[84:85], 2 // 0000000039C4: 8ED48254
|
||||
s_add_u32 s16, s16, s84 // 0000000039C8: 80105410
|
||||
s_addc_u32 s17, s17, s85 // 0000000039CC: 82115511
|
||||
|
||||
label_NoBranch_T8JHFHKM7BO5OHXW:
|
||||
s_xor_b32 s46, s48, s46 // 0000000039F0: 882E2E30
|
||||
s_xor_b32 s47, s49, s47 // 0000000039F4: 882F2F31
|
||||
s_cmp_eq_u32 s12, 1 // 0000000039F8: BF06810C
|
||||
s_cbranch_scc1 label_skipPGR2 // 0000000039FC: BF850040
|
||||
s_cbranch_scc1 after_prefetch // 0000000039FC: BF850040
|
||||
s_mov_b32 m0, s46 // 000000003A00: BEFC002E
|
||||
buffer_load_dwordx4 v0, s[52:55], 0 offen lds // 000000003A04: E05D1000 800D0000
|
||||
s_add_u32 m0, m0, 0x1040 // 000000003A0C: 807CFF7C 00001040
|
||||
@@ -517,7 +506,7 @@ label_NoBranch_T8JHFHKM7BO5OHXW:
|
||||
s_xor_b32 s46, s48, s46 // 000000003AF8: 882E2E30
|
||||
s_xor_b32 s47, s49, s47 // 000000003AFC: 882F2F31
|
||||
|
||||
label_skipPGR2:
|
||||
after_prefetch:
|
||||
s_waitcnt vmcnt(24) // 000000003B00: BF8C4F78
|
||||
s_barrier // 000000003B04: BF8A0000
|
||||
ds_read_b128 v[4:7], v2 // 000000003B08: D9FE0000 04000002
|
||||
@@ -539,14 +528,12 @@ label_skipPGR2:
|
||||
ds_read_b128 v[92:95], v3 offset:768 // 000000003B80: D9FE0300 5C000003
|
||||
ds_read_b128 v[96:99], v3 offset:896 // 000000003B88: D9FE0380 60000003
|
||||
s_waitcnt lgkmcnt(0) // 000000003B90: BF8CC07F
|
||||
|
||||
label_openLoopL:
|
||||
s_cmp_eq_u32 s12, 1 // 000000003B94: BF06810C
|
||||
s_cbranch_scc1 label_toPGR1 // 000000003B98: BF8502E5
|
||||
s_cbranch_scc1 final_compute // 000000003B98: BF8502E5
|
||||
s_cmp_le_u32 s12, 2 // 000000003B9C: BF0B820C
|
||||
s_cbranch_scc1 label_LoopEndL // 000000003BA0: BF85019E
|
||||
s_cbranch_scc1 loop_epilogue // 000000003BA0: BF85019E
|
||||
|
||||
label_LoopBeginL:
|
||||
main_loop:
|
||||
v_mfma_f32_16x16x32_bf16 a[0:3], v[68:71], v[4:7], a[0:3] // 000000003BA4: D3B58000 04020944
|
||||
ds_read_b128 v[36:39], v2 offset:64 // 000000003BAC: D9FE0040 24000002
|
||||
v_mfma_f32_16x16x32_bf16 a[4:7], v[68:71], v[8:11], a[4:7] // 000000003BB4: D3B58004 04121144
|
||||
@@ -770,9 +757,9 @@ label_LoopBeginL:
|
||||
s_cmp_eq_i32 s12, 2 // 000000004208: BF00820C
|
||||
s_waitcnt lgkmcnt(0) // 00000000420C: BF8CC07F
|
||||
v_mfma_f32_16x16x32_bf16 a[252:255], v[128:131], v[64:67], a[252:255]// 000000004210: D3B580FC 07F28180
|
||||
s_cbranch_scc0 label_LoopBeginL // 000000004218: BF84FE62
|
||||
s_cbranch_scc0 main_loop // 000000004218: BF84FE62
|
||||
|
||||
label_LoopEndL:
|
||||
loop_epilogue:
|
||||
v_mfma_f32_16x16x32_bf16 a[0:3], v[68:71], v[4:7], a[0:3] // 00000000421C: D3B58000 04020944
|
||||
ds_read_b128 v[36:39], v2 offset:64 // 000000004224: D9FE0040 24000002
|
||||
v_mfma_f32_16x16x32_bf16 a[4:7], v[68:71], v[8:11], a[4:7] // 00000000422C: D3B58004 04121144
|
||||
@@ -939,7 +926,7 @@ label_LoopEndL:
|
||||
v_mfma_f32_16x16x32_bf16 a[248:251], v[128:131], v[60:63], a[248:251]// 000000004720: D3B580F8 07E27980
|
||||
v_mfma_f32_16x16x32_bf16 a[252:255], v[128:131], v[64:67], a[252:255]// 000000004728: D3B580FC 07F28180
|
||||
|
||||
label_toPGR1:
|
||||
final_compute:
|
||||
s_and_b32 s8, s50, 0x3fff // 000000004730: 8608FF32 00003FFF
|
||||
s_and_b32 s84, 0xff, s24 // 000000004750: 865418FF 000000FF
|
||||
s_add_u32 s85, -1, s14 // 000000004758: 80550EC1
|
||||
@@ -1095,7 +1082,6 @@ label_toPGR1:
|
||||
v_mfma_f32_16x16x32_bf16 a[248:251], v[128:131], v[60:63], a[248:251]// 000000004BFC: D3B580F8 07E27980
|
||||
v_mfma_f32_16x16x32_bf16 a[252:255], v[128:131], v[64:67], a[252:255]// 000000004C04: D3B580FC 07F28180
|
||||
|
||||
label_toPGR1end_OptNLL:
|
||||
v_lshrrev_b32_e32 v4, 6, v134 // 000000004C0C: 20090C86
|
||||
v_lshrrev_b32_e32 v5, 1, v4 // 000000004C10: 200A0881
|
||||
v_mul_lo_u32 v5, 16, v5 // 000000004C14: D2850005 00020A90
|
||||
@@ -1114,7 +1100,6 @@ label_toPGR1end_OptNLL:
|
||||
s_mul_i32 s8, 0x100, s3 // 000000004C64: 920803FF 00000100
|
||||
v_add_u32_e32 v1, s8, v1 // 000000004C6C: 68020208
|
||||
|
||||
label_GW_B0_E0:
|
||||
v_add_lshl_u32 v11, v3, v0, 1 // 000000004C70: D1FE000B 02060103
|
||||
v_accvgpr_read_b32 v16, a0 // 000000004C78: D3D84010 18000100
|
||||
v_accvgpr_read_b32 v17, a4 // 000000004C80: D3D84011 18000104
|
||||
@@ -1633,86 +1618,4 @@ label_GW_B0_E0:
|
||||
s_addc_u32 s17, s17, 0 // 000000005B14: 82118011
|
||||
buffer_store_dwordx4 v[40:43], v11, s[16:19], 0 offen nt // 000000005B18: E07E1000 8004280B
|
||||
s_nop 0 // 000000005B20: BF800000
|
||||
|
||||
end:
|
||||
s_endpgm // 00000001F5D0: BF810000
|
||||
|
||||
.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 32
|
||||
|
||||
# ---- 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 SGPR enables (descriptor bits >448) ----
|
||||
.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 specific (RSRC3) ----
|
||||
.amdhsa_accum_offset 248
|
||||
.amdhsa_uses_dynamic_stack 0
|
||||
.amdhsa_tg_split 0
|
||||
|
||||
.end_amdhsa_kernel
|
||||
|
||||
.amdgpu_metadata
|
||||
---
|
||||
amdhsa.kernels:
|
||||
- .args:
|
||||
- .address_space: global
|
||||
.name: C
|
||||
.offset: 0
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: bf16
|
||||
- .address_space: global
|
||||
.name: B
|
||||
.offset: 8
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: bf16
|
||||
- .address_space: global
|
||||
.name: A
|
||||
.offset: 16
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: bf16
|
||||
- .name: sz
|
||||
.offset: 24
|
||||
.size: 4
|
||||
.value_kind: by_value
|
||||
.value_type: u32
|
||||
- .name: num_wg
|
||||
.offset: 28
|
||||
.size: 4
|
||||
.value_kind: by_value
|
||||
.value_type: u32
|
||||
.group_segment_fixed_size: 133120
|
||||
.kernarg_segment_align: 8
|
||||
.kernarg_segment_size: 32
|
||||
.max_flat_workgroup_size: 256
|
||||
.name: gemm
|
||||
.private_segment_fixed_size: 0
|
||||
.sgpr_count: 88
|
||||
.sgpr_spill_count: 0
|
||||
.symbol: gemm.kd
|
||||
.vgpr_count: 248
|
||||
.vgpr_spill_count: 0
|
||||
.wavefront_size: 64
|
||||
amdhsa.version:
|
||||
- 1
|
||||
- 0
|
||||
...
|
||||
.end_amdgpu_metadata
|
||||
@@ -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
|
||||
@@ -2,10 +2,8 @@
|
||||
# VIZ=2 to profile
|
||||
import pathlib
|
||||
from tinygrad import Tensor, Device, dtypes, Context
|
||||
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, getenv
|
||||
from tinygrad.uop.ops import UOp, Ops, KernelInfo
|
||||
from tinygrad.helpers import getenv
|
||||
|
||||
fp = pathlib.Path(__file__).parent/"gemm.s"
|
||||
|
||||
@@ -23,8 +21,8 @@ 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
|
||||
Bt = B.t().contiguous() # transpose B for the asm gemm
|
||||
C_torch = A@B
|
||||
|
||||
# ** copy buffers to AMD
|
||||
|
||||
@@ -33,31 +31,32 @@ C_torch = A@Bt
|
||||
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_tiny = from_torch(A) @ from_torch(B)
|
||||
C_asm = Tensor.empty_like(C_tiny)
|
||||
C_asm.uop.buffer.allocate()
|
||||
|
||||
# ** assembly custom kernel
|
||||
|
||||
def custom_asm_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
|
||||
lidx = UOp.special(THREADS_PER_WG, "lidx0")
|
||||
gidx = UOp.special(NUM_WG, "gidx0")
|
||||
|
||||
src = (pathlib.Path(__file__).parent/"template.s").read_text().replace("INSTRUCTIONS", fp.read_text())
|
||||
|
||||
sz = UOp.variable("SZ", 256, 8192)
|
||||
|
||||
sink = UOp.sink(C.base, A.base, B.base, sz, lidx, gidx, arg=KernelInfo(name="gemm"))
|
||||
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
|
||||
|
||||
C_asm = Tensor.custom_kernel(C_asm, from_torch(A), from_torch(Bt), fxn=custom_asm_gemm)[0]
|
||||
|
||||
# ** run gemms
|
||||
|
||||
# baseline tinygrad
|
||||
sched = C_tiny.schedule()
|
||||
assert len(sched) == 1
|
||||
eis:list[ExecItem] = [sched[-1].lower()]
|
||||
ast = sched[-1].ast
|
||||
|
||||
# assembly gemm
|
||||
@track_rewrites(name=lambda ret: TracingKey(ret.name, (ret.function_name,), ret))
|
||||
def get_asm_prg() -> ProgramSpec:
|
||||
src = fp.read_text()
|
||||
lib = Device[Device.DEFAULT].compiler.compile(src)
|
||||
return ProgramSpec("gemm", src, Device.DEFAULT, ast, lib=lib, global_size=[NUM_WG, 1, 1], local_size=[THREADS_PER_WG, 1, 1],
|
||||
globals=[0, 1, 2], vars=[UOp.variable("SZ", 256, 8192), UOp.variable("NUM_WG", 1, 1024)])
|
||||
eis.append(ExecItem(ast, [C_asm.uop.buffer, from_torch(B).uop.buffer, from_torch(A).uop.buffer], fixedvars={"SZ":N, "NUM_WG":NUM_WG},
|
||||
prg=CompiledRunner(get_asm_prg())))
|
||||
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(wait=True)
|
||||
et = ei.run({"SZ":N}, wait=True)
|
||||
print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
|
||||
|
||||
# ** correctness
|
||||
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)))
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_matmul(asm_kernel(), dtype=dtypes.half, N=N)
|
||||
@@ -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()
|
||||
|
||||
@@ -46,6 +46,22 @@ def install_hook(c_function, python_function):
|
||||
# *** ioctl lib end ***
|
||||
|
||||
import tinygrad.runtime.autogen.kfd as kfd_ioctl
|
||||
import tinygrad.runtime.autogen.hsa as hsa
|
||||
|
||||
def print_aql_queue(read_pointer_address):
|
||||
rptr_offset = getattr(hsa.amd_queue_v2_t, 'read_dispatch_id').offset
|
||||
queue_base = read_pointer_address - rptr_offset
|
||||
queue = hsa.amd_queue_v2_t.from_address(queue_base)
|
||||
print(f" AQL Queue @ 0x{queue_base:X}:")
|
||||
for field_name, field_type in hsa.amd_queue_v2_t._fields_:
|
||||
val = getattr(queue, field_name)
|
||||
if isinstance(val, int): print(f" {field_name}: 0x{val:X}")
|
||||
elif hasattr(val, '_length_'):
|
||||
arr_vals = [f"{format_struct(v)}" if hasattr(v, '_fields_') else f"{v:#X}" for v in val]
|
||||
print(f" {field_name}: [{', '.join(arr_vals)}]")
|
||||
elif hasattr(val, '_fields_'): print(f" {field_name}: {format_struct(val)}")
|
||||
else: print(f" {field_name}: {val}")
|
||||
|
||||
def ioctls_from_header():
|
||||
hdr = (pathlib.Path(__file__).parent / "kfd_ioctl.h").read_text().replace("\\\n", "")
|
||||
pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_IOW?R?\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
|
||||
@@ -67,6 +83,7 @@ def ioctl(fd, request, argp):
|
||||
if name == "AMDKFD_IOC_SVM":
|
||||
out = ctypes.cast(s.attrs, ctypes.POINTER(kfd_ioctl.struct_kfd_ioctl_svm_attribute))
|
||||
for i in range(s.nattr): print(f"{i}: {kfd_ioctl.enum_kfd_ioctl_svm_attr_type.get(out[i].type):40s}: {out[i].value:#x}")
|
||||
if name == "AMDKFD_IOC_CREATE_QUEUE" and s.queue_type == kfd_ioctl.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL: print_aql_queue(s.read_pointer_address)
|
||||
else:
|
||||
print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : ioctl",
|
||||
f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", os.readlink(f"/proc/self/fd/{fd}") if fd >= 0 else "")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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.amd.autogen.rdna3 import *
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
from extra.assembly.amd.asm import waitcnt
|
||||
from test.testextra.test_cfg_viz import template
|
||||
|
||||
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.
|
||||
|
||||
|
||||
@@ -145,7 +145,7 @@ OPNAME = {
|
||||
ALUSRC = {
|
||||
1: "SALU",
|
||||
2: "VALU",
|
||||
3: "VALU_ALT",
|
||||
3: "VALU_SALU",
|
||||
}
|
||||
|
||||
MEMSRC = {
|
||||
|
||||
@@ -275,7 +275,7 @@ class RGP:
|
||||
),
|
||||
shader_engine_index=sqtt_event.se,
|
||||
sqtt_version={11: sqtt.SQTT_VERSION_3_2, 12: sqtt.SQTT_VERSION_3_3}.get(gfx_ver),
|
||||
v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_0_v1(
|
||||
v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_v1(
|
||||
instrumentation_spec_version=1,
|
||||
instrumentation_api_version=0,
|
||||
compute_unit_index=0,
|
||||
|
||||
+7
-4
@@ -22,9 +22,7 @@ class WaveSlot:
|
||||
@property
|
||||
def cu_loc(self) -> str: return f"SE:{self.se} CU:{self.cu}"
|
||||
@property
|
||||
def simd_loc(self) -> str: return f"{self.cu_loc} SIMD:{self.simd}"
|
||||
@property
|
||||
def wave_loc(self) -> str: return f"{self.simd_loc} W:{self.wave_id}"
|
||||
def wave_loc(self) -> str: return f"{self.cu_loc} SIMD:{self.simd} W:{self.wave_id}"
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class WaveExec(WaveSlot):
|
||||
@@ -113,12 +111,17 @@ 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
|
||||
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_data(data:dict) -> None:
|
||||
|
||||
+120
-34
@@ -11,6 +11,15 @@ NUM_WORKERS = 1
|
||||
Q_BLOCK_SIZE = 16
|
||||
KV_BLOCK_SIZE = 16
|
||||
|
||||
def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None) -> Tensor:
|
||||
if not isinstance(ref.device, tuple): return Tensor.empty(*shape, device=ref.device)
|
||||
shape = tuple(s // len(ref.device) if i == ref.uop.axis else s for i, s in enumerate(shape))
|
||||
axis = ref.uop.axis if axis is None else axis
|
||||
return Tensor(Tensor.empty(*shape, device=ref.device).uop.multi(axis), device=ref.device)
|
||||
|
||||
def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
|
||||
return _sharded_empty(ref.shape, ref, axis)
|
||||
|
||||
def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False):
|
||||
if len(xq.shape) == 3: xq, xk, xv = xq.unsqueeze(0), xk.unsqueeze(0), xv.unsqueeze(0)
|
||||
|
||||
@@ -45,7 +54,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
|
||||
k_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
|
||||
v_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
|
||||
|
||||
q_reg_fl = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
|
||||
q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
|
||||
q_reg_transposed = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
|
||||
k_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
|
||||
@@ -69,9 +77,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
|
||||
scale_vec = warp.ones(scale_vec)
|
||||
|
||||
# load q tile
|
||||
q_reg_fl = warp.load(q_reg_fl, q, (), (batch, q_seq, head, 0), axis=1)
|
||||
q_reg_fl *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
|
||||
q_reg = warp.copy(q_reg, q_reg_fl)
|
||||
q_reg = warp.load(q_reg, q, (), (batch, q_seq, head, 0), axis=1)
|
||||
q_reg_transposed = warp.transpose(q_reg_transposed, q_reg)
|
||||
|
||||
for kv_idx in ker.range(N // KV_BLOCK_SIZE):
|
||||
@@ -85,6 +91,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
|
||||
att_block = warp.zero(att_block.after(kv_idx))
|
||||
k_reg_transposed = warp.transpose(k_reg_transposed, k_reg)
|
||||
att_block = warp.mma_AtB(att_block, k_reg_transposed, q_reg_transposed)
|
||||
att_block *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
|
||||
|
||||
# apply attention mask
|
||||
mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
|
||||
@@ -217,11 +224,11 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
|
||||
|
||||
return ker.finish()
|
||||
|
||||
def custom_backward_kv(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
|
||||
with Kernel("fa_custom_backward_kv", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
|
||||
def custom_backward_k(dku:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
|
||||
with Kernel("fa_custom_backward_k", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
|
||||
warp = ker.warp
|
||||
|
||||
dk, dv, do, q, k, v, mask = GL(dku, ker), GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
|
||||
dk, do, q, k, v, mask = GL(dku, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
|
||||
l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)
|
||||
|
||||
head_kv = ker.blockIdx_x
|
||||
@@ -242,7 +249,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
|
||||
mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
|
||||
|
||||
dk_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.float32, TileLayout.COL)
|
||||
dv_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.float32, TileLayout.COL)
|
||||
do_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
|
||||
do_reg_col = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
|
||||
|
||||
@@ -256,7 +262,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
|
||||
delta_vec_reg = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
|
||||
|
||||
dk_reg = warp.zero(dk_reg)
|
||||
dv_reg = warp.zero(dv_reg)
|
||||
|
||||
# load kv tile
|
||||
k_reg = warp.load(k_reg, k, (), (batch, kv_seq, head_kv, 0), axis=1)
|
||||
@@ -295,13 +300,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
|
||||
att_block -= l_vec_reg
|
||||
att_block = att_block.exp2()
|
||||
|
||||
att_block_mma = warp.copy(att_block_mma, att_block)
|
||||
att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
|
||||
att_smem = warp.store(att_smem, att_block_transposed)
|
||||
att_block_row = warp.load(att_block_row, att_smem)
|
||||
dv_reg_ = warp.mma_AB(dv_reg, att_block_row, do_reg_col)
|
||||
|
||||
dp_block = warp.zero(dp_block.after(g, q_idx, dv_reg_))
|
||||
dp_block = warp.zero(dp_block.after(g, q_idx))
|
||||
dp_block = warp.mma_ABt(dp_block, v_reg, do_reg)
|
||||
dp_block -= delta_vec_reg
|
||||
att_block *= dp_block
|
||||
@@ -312,41 +311,128 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
|
||||
att_block_row = warp.load(att_block_row, att_smem)
|
||||
dk_reg = warp.mma_AB(dk_reg, att_block_row, q_reg_col)
|
||||
dk_reg = ker.endrange(2)
|
||||
dv_reg = dv_reg.after(dk_reg)
|
||||
|
||||
dk_reg *= 1.0 / math.sqrt(D)
|
||||
|
||||
dk = warp.store(dk, dk_reg, (batch, kv_seq, head_kv, 0), axis=1)
|
||||
|
||||
return ker.finish()
|
||||
|
||||
def custom_backward_v(dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp) -> UOp:
|
||||
with Kernel("fa_custom_backward_v", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
|
||||
warp = ker.warp
|
||||
|
||||
dv, do, q, k, v, mask = GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
|
||||
l_vec = GL(l_vecu, ker)
|
||||
|
||||
head_kv = ker.blockIdx_x
|
||||
batch = ker.blockIdx_z
|
||||
kv_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
|
||||
|
||||
q_smem = ker.st((Q_BLOCK_SIZE, D), dtypes.bfloat16)
|
||||
do_smem = ker.st((Q_BLOCK_SIZE, D), dtypes.bfloat16)
|
||||
att_smem = ker.st((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16)
|
||||
|
||||
q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
|
||||
q_reg_t = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
|
||||
q_reg_col = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
|
||||
k_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
|
||||
k_reg_t = ker.rt((D, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
|
||||
v_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
|
||||
mask_reg = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.float32)
|
||||
mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
|
||||
|
||||
dv_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.float32, TileLayout.COL)
|
||||
do_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
|
||||
do_reg_col = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
|
||||
|
||||
att_block = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
|
||||
att_block_mma = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
|
||||
att_block_transposed = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
|
||||
att_block_row = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16)
|
||||
|
||||
l_vec_reg = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
|
||||
|
||||
dv_reg = warp.zero(dv_reg)
|
||||
|
||||
# load kv tile
|
||||
k_reg = warp.load(k_reg, k, (), (batch, kv_seq, head_kv, 0), axis=1)
|
||||
k_reg_t = warp.transpose(k_reg_t, k_reg)
|
||||
v_reg = warp.load(v_reg, v, (), (batch, kv_seq, head_kv, 0), axis=1)
|
||||
|
||||
for q_idx in ker.range(N // Q_BLOCK_SIZE):
|
||||
for g in ker.range(GROUP_SIZE):
|
||||
head_q = head_kv * GROUP_SIZE + g
|
||||
|
||||
# load q and do
|
||||
q_smem = warp.load(q_smem, q, (), (batch, q_idx, head_q, 0), axis=1)
|
||||
do_smem = warp.load(do_smem, do, (), (batch, q_idx, head_q, 0), axis=1)
|
||||
|
||||
q_reg = warp.load(q_reg, q_smem)
|
||||
q_reg_t = warp.transpose(q_reg_t, q_reg)
|
||||
q_reg_col = warp.load(q_reg_col, q_smem)
|
||||
do_reg = warp.load(do_reg, do_smem)
|
||||
do_reg_col = warp.load(do_reg_col, do_smem)
|
||||
|
||||
# load l_vec
|
||||
l_vec_reg = warp.load(l_vec_reg, l_vec, (), (batch, head_q, 0, q_idx), axis=2)
|
||||
l_vec_reg *= 1.0 / math.log(2)
|
||||
|
||||
# mma qk^t
|
||||
att_block = warp.zero(att_block.after(g))
|
||||
att_block = warp.mma_AtB(att_block, k_reg_t, q_reg_t)
|
||||
|
||||
# apply attention mask
|
||||
mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_idx, kv_seq), axis=2)
|
||||
mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
|
||||
att_block += mask_reg_transposed
|
||||
|
||||
att_block *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
|
||||
att_block -= l_vec_reg
|
||||
att_block = att_block.exp2()
|
||||
|
||||
att_block_mma = warp.copy(att_block_mma, att_block)
|
||||
att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
|
||||
att_smem = warp.store(att_smem, att_block_transposed)
|
||||
att_block_row = warp.load(att_block_row, att_smem)
|
||||
dv_reg = warp.mma_AB(dv_reg, att_block_row, do_reg_col)
|
||||
dv_reg = ker.endrange(2)
|
||||
|
||||
dv_reg = warp.map(dv_reg, lambda x, idx: x + v_reg[*idx].cast(dtypes.float32) * 1e-30)
|
||||
|
||||
dv = warp.store(dv, dv_reg, (batch, kv_seq, head_kv, 0), axis=1)
|
||||
|
||||
return ker.finish(2)
|
||||
return ker.finish()
|
||||
|
||||
single_device = xq.device[0] if isinstance(xq.device, tuple) else xq.device
|
||||
|
||||
if is_causal:
|
||||
if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
|
||||
attn_mask = Tensor.ones((B, 1, N, N), requires_grad=False, device=xq.device, dtype=dtypes.bool).tril()
|
||||
attn_mask = Tensor.ones((B, 1, N, N), requires_grad=False, device=single_device, dtype=dtypes.bool).tril()
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
|
||||
else:
|
||||
attn_mask = Tensor.zeros((B, 1, N, N), requires_grad=False, device=xq.device, dtype=dtypes.float32)
|
||||
attn_mask = Tensor.zeros((B, 1, N, N), requires_grad=False, device=single_device, dtype=dtypes.float32)
|
||||
if isinstance(xq.device, tuple) and not isinstance(attn_mask.device, tuple):
|
||||
attn_mask = attn_mask.shard(xq.device, axis=0)
|
||||
|
||||
attn = Tensor.empty_like(xq)
|
||||
l_vec = Tensor.empty(B, H, 1, N, requires_grad=False, device=xq.device, dtype=dtypes.float32).detach()
|
||||
attn = _sharded_empty_like(xq, axis=0)
|
||||
l_vec = _sharded_empty((B, H, 1, N), xq, axis=0)
|
||||
|
||||
def grad(gradu:UOp, kernel:UOp) -> tuple[None, None, UOp, UOp, UOp, None]:
|
||||
grad = Tensor(gradu)
|
||||
grad_q = Tensor.empty_like(q := Tensor(kernel.src[2]))
|
||||
grad_k = Tensor.empty_like(k := Tensor(kernel.src[3]))
|
||||
grad_v = Tensor.empty_like(v := Tensor(kernel.src[4]))
|
||||
mask = Tensor(kernel.src[5])
|
||||
def grad(gradu:UOp, _) -> tuple[None, None, UOp, UOp, UOp, None]:
|
||||
grad = Tensor(gradu, device=gradu.device)
|
||||
grad_q = _sharded_empty_like(xq, axis=0)
|
||||
grad_k = _sharded_empty_like(xk, axis=0)
|
||||
grad_v = _sharded_empty_like(xv, axis=0)
|
||||
|
||||
delta_vec = (grad * attn).sum(-1).transpose(1, 2).unsqueeze(-2).detach()
|
||||
print(l_vec.shape, delta_vec.shape, grad.shape, attn.shape, grad_q.shape, grad_k.shape, grad_v.shape)
|
||||
delta_vec = (grad * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
|
||||
|
||||
grad_q = Tensor.custom_kernel(grad_q, grad, q, k, v, mask, l_vec, delta_vec, fxn=custom_backward_q)[0]
|
||||
grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, q, k, v, mask, l_vec, delta_vec, fxn=custom_backward_kv)[:2]
|
||||
grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_q)[0]
|
||||
grad_k = Tensor.custom_kernel(grad_k, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_k)[0]
|
||||
grad_v = Tensor.custom_kernel(grad_v, grad, xq, xk, xv, attn_mask, l_vec, fxn=custom_backward_v)[0]
|
||||
return (None, None, grad_q.uop, grad_k.uop, grad_v.uop, None)
|
||||
|
||||
attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward, grad_fxn=grad)[:2]
|
||||
attn = attn[:, :N_, :, :D_]
|
||||
attn_ = attn[:, :N_, :, :D_]
|
||||
|
||||
return attn.transpose(1, 2).cast(odtype)
|
||||
return attn_.transpose(1, 2).cast(odtype)
|
||||
|
||||
@@ -24,13 +24,11 @@ class Group:
|
||||
|
||||
# ops that only work on a single warp
|
||||
|
||||
clear_rid = 1000
|
||||
def clear(self, reg:ALL_TILES, value:float=0):
|
||||
reg = cast(UOp, reg)
|
||||
assert self.warps == 1
|
||||
|
||||
rngs_for_shape = tuple(UOp.range(dim, Group.clear_rid + i) for i, dim in enumerate(reg.shape))
|
||||
Group.clear_rid += len(reg.shape)
|
||||
rngs_for_shape = tuple(self.ker.raw_range(dim) for dim in reg.shape)
|
||||
|
||||
reg_store = reg[*rngs_for_shape].store(value).end(*rngs_for_shape)
|
||||
|
||||
@@ -41,14 +39,12 @@ class Group:
|
||||
def ones(self, reg:ALL_TILES): return self.clear(reg, 1)
|
||||
def neg_inf(self, reg:ALL_TILES): return self.clear(reg, -math.inf)
|
||||
|
||||
copy_rid = 300
|
||||
def copy(self, dst:ALL_TILES, src:ALL_TILES):
|
||||
dst, src = cast(UOp, dst), cast(UOp, src)
|
||||
assert self.warps == 1
|
||||
assert dst.shape == src.shape
|
||||
|
||||
rngs_for_shape = tuple(UOp.range(dim, Group.copy_rid + i) for i, dim in enumerate(dst.shape))
|
||||
Group.copy_rid += len(dst.shape)
|
||||
rngs_for_shape = tuple(self.ker.raw_range(dim) for dim in dst.shape)
|
||||
|
||||
src_load = src[*rngs_for_shape]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
@@ -219,8 +215,7 @@ class Group:
|
||||
red_reg = self.ker.alloc((1,), src.dtype.base, AddrSpace.REG)
|
||||
|
||||
for height in self.ker.range(src.shape[-3], track=False):
|
||||
i = UOp.range(red_reg.size, Group.clear_rid)
|
||||
Group.clear_rid += 1
|
||||
i = self.ker.raw_range(red_reg.size)
|
||||
red_reg = red_reg.after(height, *[tkr._rng for tkr in self.ker.range_stack])
|
||||
reg_store = red_reg.flatten()[i].store(init_value).end(i)
|
||||
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
|
||||
@@ -254,8 +249,7 @@ class Group:
|
||||
red_reg = self.ker.alloc((1,), src.dtype.base, AddrSpace.REG)
|
||||
|
||||
for width in self.ker.range(src.shape[-2], track=False):
|
||||
i = UOp.range(red_reg.size, Group.clear_rid)
|
||||
Group.clear_rid += 1
|
||||
i = self.ker.raw_range(red_reg.size)
|
||||
red_reg = red_reg.after(width, *[tkr._rng for tkr in self.ker.range_stack])
|
||||
reg_store = red_reg.flatten()[i].store(init_value).end(i)
|
||||
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
|
||||
@@ -302,9 +296,20 @@ class Group:
|
||||
row = laneid % rt.base_shape.rows
|
||||
col = rt.base_shape.stride * (laneid // rt.base_shape.rows) + inner
|
||||
|
||||
sheight = height
|
||||
swidth = width
|
||||
if len(idxs) == 2:
|
||||
row_idx = idxs[0] * dst.shape[-3] * rt.base_shape.rows
|
||||
col_idx = idxs[1] * dst.shape[-2] * rt.base_shape.cols
|
||||
|
||||
row += row_idx % st.base_shape.rows
|
||||
col += col_idx % st.base_shape.cols
|
||||
sheight += row_idx // st.base_shape.rows
|
||||
swidth += col_idx // st.base_shape.cols
|
||||
|
||||
srow, scol = cast(ST, src).swizzle(row, col)
|
||||
|
||||
src_load = src[*idxs[:-2], height, width, srow, scol]
|
||||
src_load = src[*idxs[:-2], sheight, swidth, srow, scol]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dst[*dst_idxs, height, width, inner].store(src_load)
|
||||
@@ -318,28 +323,31 @@ class Group:
|
||||
idxs = tuple(idx * st.cols if i == 3 else idx for i, idx in enumerate(idxs))
|
||||
src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
|
||||
|
||||
for height in self.ker.range(dst.shape[-4], track=False):
|
||||
for width in self.ker.range(dst.shape[-3], track=False):
|
||||
elements_per_thread = st.base_shape.elements_per_thread
|
||||
memcpy_per_row = st.base_shape.cols // elements_per_thread
|
||||
total_calls = st.base_shape.num_elements // (self.group_threads * elements_per_thread)
|
||||
elements_per_thread = st.base_shape.elements_per_thread
|
||||
memcpy_per_row = st.cols // elements_per_thread
|
||||
total_calls = (dst.shape[-4] * dst.shape[-3] * st.base_shape.num_elements) // (self.group_threads * elements_per_thread)
|
||||
|
||||
for outer in self.ker.range(total_calls, track=False):
|
||||
for inner in self.ker.range(elements_per_thread, axis_type=AxisType.UPCAST, track=False):
|
||||
load_idx = outer * self.group_threads + self.laneid
|
||||
row = load_idx // memcpy_per_row
|
||||
col = (load_idx * elements_per_thread) % st.base_shape.cols + inner
|
||||
for outer in self.ker.range(total_calls, track=False):
|
||||
for inner in self.ker.range(elements_per_thread, axis_type=AxisType.UPCAST, track=False):
|
||||
load_idx = outer * self.group_threads + self.laneid
|
||||
row = load_idx // memcpy_per_row
|
||||
col = (load_idx * elements_per_thread) % st.cols + inner
|
||||
height = row // st.base_shape.rows
|
||||
width = col // st.base_shape.cols
|
||||
|
||||
srow, scol = cast(ST, dst).swizzle(row, col)
|
||||
row = row % st.base_shape.rows
|
||||
col = col % st.base_shape.cols
|
||||
|
||||
src_i += height * st.base_shape.rows * row_stride + width * st.base_shape.cols
|
||||
src_i += row * row_stride + col
|
||||
srow, scol = cast(ST, dst).swizzle(row, col)
|
||||
|
||||
src_load = srcf[src_i]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
|
||||
dst_store = dst_store.end(height, width, outer, inner).barrier()
|
||||
src_i += height * st.base_shape.rows * row_stride + width * st.base_shape.cols
|
||||
src_i += row * row_stride + col
|
||||
|
||||
src_load = srcf[src_i]
|
||||
if src.dtype.base != dst.dtype.base:
|
||||
src_load = src_load.cast(dst.dtype.base)
|
||||
dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
|
||||
dst_store = dst_store.end(height, width, outer, inner).barrier()
|
||||
elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
|
||||
srcf = src.flatten()
|
||||
row_stride = prod(src.shape[axis+1:])
|
||||
@@ -471,7 +479,7 @@ class Group:
|
||||
idxs = tuple(idx * rv.length if i == 3 else idx for i, idx in enumerate(idxs))
|
||||
dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
|
||||
|
||||
for outer in self.ker.range(src.shape[-2]):
|
||||
for outer in self.ker.range(src.shape[-2], track=False):
|
||||
dst_i += outer * reductions + (laneid % reductions)
|
||||
|
||||
src_load = src[outer, 0]
|
||||
|
||||
@@ -55,6 +55,11 @@ class Kernel(AbstractContextManager):
|
||||
if track: self.range_stack.append(rng)
|
||||
return rng
|
||||
|
||||
def raw_range(self, end:int=0, axis_type:AxisType=AxisType.LOOP):
|
||||
rng = UOp.range(end, self.range_id, axis_type=axis_type)
|
||||
self.range_id += 1
|
||||
return rng
|
||||
|
||||
def alloc(self, shape, dtype, addrspace:AddrSpace, name:str|None=None):
|
||||
match addrspace:
|
||||
case AddrSpace.GLOBAL:
|
||||
|
||||
@@ -7,7 +7,7 @@ from tinygrad.uop.ops import Ops
|
||||
from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
|
||||
import torch.lib
|
||||
TORCH_DEBUG = getenv("TORCH_DEBUG")
|
||||
import torch, pathlib, math, operator, functools, weakref
|
||||
import torch, pathlib, operator, functools, weakref
|
||||
torch.autograd.grad_mode.set_multithreading_enabled(False)
|
||||
from tinygrad.dtype import _from_torch_dtype, _to_torch_dtype
|
||||
|
||||
@@ -74,6 +74,7 @@ view_ops = {
|
||||
"aten.select.int": lambda self, dim, idx: self[(slice(None),) * (dim%self.ndim) + (idx,)],
|
||||
"aten.permute": Tensor.permute,
|
||||
"aten.alias": lambda self: self,
|
||||
"aten.diagonal": Tensor.diagonal,
|
||||
}
|
||||
|
||||
# torch 2.10 handles this natively
|
||||
@@ -154,7 +155,9 @@ def index_put(self, indices, values, accumulate=False):
|
||||
return aten.index_put(self.cpu(), [z.cpu() if isinstance(z, torch.Tensor) else None for z in indices], values.clone().cpu(), accumulate).tiny()
|
||||
|
||||
@torch.library.impl("aten::isin.Tensor_Tensor_out", "privateuseone")
|
||||
def isin_tensor_tensor_out(x, y, *, assume_unique=False, invert=False, out=None): return out.copy_(aten.isin(x.cpu(), y.cpu(), assume_unique=assume_unique, invert=invert).tiny())
|
||||
def isin_tensor_tensor_out(x, y, *, assume_unique=False, invert=False, out=None):
|
||||
result = (unwrap(x).unsqueeze(-1) == unwrap(y).flatten()).any(-1)
|
||||
return out.copy_(wrap(~result if invert else result))
|
||||
|
||||
@torch.library.impl("aten::randperm.generator_out", "privateuseone")
|
||||
def randperm_generator(n, generator=None, out=None):
|
||||
@@ -162,13 +165,16 @@ def randperm_generator(n, generator=None, out=None):
|
||||
|
||||
@torch.library.impl("aten::cummax", "privateuseone")
|
||||
def cummax(self, dim):
|
||||
# TODO: support cummax with indices to match torch
|
||||
cummax, indices = aten.cummax(self.cpu(), dim)
|
||||
return (cummax.tiny(), indices.tiny())
|
||||
values, indices = unwrap(self).cummax(dim)
|
||||
return (wrap(values), wrap(indices.cast(dtypes.int64)))
|
||||
|
||||
@torch.library.impl("aten::cummin", "privateuseone")
|
||||
def cummin(self, dim):
|
||||
values, indices = unwrap(self).cummin(dim)
|
||||
return (wrap(values), wrap(indices.cast(dtypes.int64)))
|
||||
|
||||
@torch.library.impl("aten::nonzero", "privateuseone")
|
||||
# TODO: move to tinygrad
|
||||
def nonzero(self): return aten.nonzero(self.cpu()).tiny()
|
||||
def nonzero(self): return wrap(unwrap(self).nonzero())
|
||||
|
||||
@torch.library.impl("aten::_linalg_eigh", "privateuseone")
|
||||
# TODO: move to tinygrad
|
||||
@@ -472,7 +478,7 @@ for k,v in get_decompositions(decomps).items():
|
||||
# the goal is to make as much as we can this
|
||||
simple_tensor_methods = [
|
||||
# unary (ish)
|
||||
"log", "log2", "sqrt", "rsqrt", "sign", "silu", "hardsigmoid", "exp", "exp2", "neg", "reciprocal", "bitwise_not",
|
||||
"log", "log2", "log10", "sqrt", "rsqrt", "sign", "silu", "hardsigmoid", "exp", "exp2", "neg", "reciprocal", "bitwise_not",
|
||||
"sigmoid", "clamp", "mish", "erf", "leaky_relu",
|
||||
# trig
|
||||
"acos", "acosh", "cos", "cosh", "asin", "asinh", "sin", "sinh", "atan", "atanh", "tan", "tanh",
|
||||
@@ -516,7 +522,6 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
|
||||
"aten.bitwise_left_shift.Tensor_out": lambda x,y: x*(2**y),
|
||||
"aten.bitwise_right_shift.Tensor_out": lambda x,y: x//(2**y),
|
||||
# not in tinygrad. are there decomps for these?
|
||||
"aten.log10.out": lambda self: self.log2() * (math.log(2) / math.log(10)),
|
||||
"aten.log1p.out": lambda self: (self+1).log(),
|
||||
"aten.expm1.out": lambda self: self.exp() - 1,
|
||||
"aten.fmax.out": lambda input,other: Tensor.where(input.isnan() & ~other.isnan(), other, Tensor.where(~input.isnan() & other.isnan(), input, Tensor.maximum(input, other))),
|
||||
@@ -625,7 +630,8 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
|
||||
Tensor.linspace(start, stop, steps, **({"dtype": _from_torch_dtype(dtype)} if dtype is not None else {})),
|
||||
"aten.topk": Tensor.topk,
|
||||
"aten.constant_pad_nd": lambda self, padding, value=0.0: self.pad(padding, mode="constant", value=value).contiguous(),
|
||||
"aten.cumsum": lambda self, dim: self.cumsum(dim).contiguous(), # TODO: fix test_simple_cumsum, fails without contiguous for shapes >512
|
||||
# TODO: input contiguous is needed to prevent CFGContext circular dependency assertion for shapes >512 (see test_cumsum_arange_large)
|
||||
"aten.cumsum": lambda self, dim: self.contiguous().cumsum(dim),
|
||||
"aten.logsumexp": lambda self, axis, keepdim=False: self.logsumexp(axis[0], keepdim=keepdim),
|
||||
"aten.roll": Tensor.roll,
|
||||
"aten.logcumsumexp": Tensor.logcumsumexp,
|
||||
@@ -751,16 +757,3 @@ def _pad_circular(self, padding): return _PadCircular.apply(self, padding)
|
||||
|
||||
@torch.library.impl("aten::_pad_circular", "AutogradPrivateUse1")
|
||||
def _pad_circular_autograd(self, padding): return _PadCircular.apply(self, padding)
|
||||
|
||||
# only needed for test_diag_backward_gradient_values
|
||||
# was going through torch before, but now we are using tinygrad directly and tracking views
|
||||
# Tensor.diagonal does not support all cases tests in the tests
|
||||
@torch.library.impl("aten::diagonal", "privateuseone")
|
||||
@wrap_view_op
|
||||
def diagonal(self, offset=0, dim1=0, dim2=1):
|
||||
if offset != 0: raise NotImplementedError(f"diagonal with {offset=} not implemented")
|
||||
dim1, dim2 = dim1 % self.ndim, dim2 % self.ndim
|
||||
if dim1 != self.ndim - 2 or dim2 != self.ndim - 1: raise NotImplementedError(f"diagonal with {dim1=}, {dim2=} not implemented, only last two dims supported")
|
||||
batch_shape, m, n = self.shape[:-2], self.shape[-2], self.shape[-1]
|
||||
diag_len = min(m, n)
|
||||
return self.reshape(*batch_shape, m*n).pad(tuple((0,0) for _ in batch_shape) + ((0, diag_len),)).reshape(*batch_shape, diag_len, n+1)[..., :, 0]
|
||||
|
||||
@@ -191,6 +191,20 @@ class TestTorchBackend(unittest.TestCase):
|
||||
assert torch.equal(tensor_a, tensor_b)
|
||||
assert not torch.equal(tensor_a, tensor_c)
|
||||
|
||||
@unittest.skip("# TODO: this test is slow")
|
||||
def test_linalg_svd(self):
|
||||
A = torch.randn(5, 5, device=device)
|
||||
U, S, Vh = torch.linalg.svd(A)
|
||||
np.testing.assert_equal(U.shape, (5,5))
|
||||
np.testing.assert_equal(Vh.shape, (5,5))
|
||||
np.testing.assert_allclose(torch.dist(A, U @ torch.diag(S) @ Vh).cpu().numpy(), 0, atol=1e-5)
|
||||
|
||||
A = torch.randn(5, 3, device=device)
|
||||
U, S, Vh = torch.linalg.svd(A, full_matrices=False)
|
||||
np.testing.assert_equal(U.shape, (5,3))
|
||||
np.testing.assert_equal(Vh.shape, (3,3))
|
||||
np.testing.assert_allclose(torch.dist(A, U @ torch.diag(S) @ Vh).cpu().numpy(), 0, atol=1e-5)
|
||||
|
||||
def test_linalg_eigh(self):
|
||||
a = torch.tensor([[1, 2], [2, 1]], dtype=torch.float32, device=device)
|
||||
w, v = torch.linalg.eigh(a)
|
||||
@@ -686,6 +700,15 @@ class TestTorchBackend(unittest.TestCase):
|
||||
expected = np.array([4.0, 3.0, 2.0, 1.0])
|
||||
np.testing.assert_allclose(a.grad.cpu().numpy(), expected, rtol=1e-5)
|
||||
|
||||
def test_cumsum_arange_large(self):
|
||||
# Tests cumsum with an unrealized arange input with size > 512 (the split threshold)
|
||||
# This exercises the _split_cumalu path which uses a two-stage algorithm
|
||||
for size in [513, 1022]:
|
||||
a = torch.arange(size, dtype=torch.float32, device=device)
|
||||
result = torch.cumsum(a, dim=0)
|
||||
expected = torch.arange(size, dtype=torch.float32).cumsum(dim=0)
|
||||
np.testing.assert_allclose(result.cpu().numpy(), expected.numpy(), rtol=1e-5)
|
||||
|
||||
def test_diag_1d_to_2d(self):
|
||||
a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device=device, requires_grad=True)
|
||||
b = torch.diag(a)
|
||||
|
||||
@@ -1,14 +1,7 @@
|
||||
#include "TinyGPUDriver.h"
|
||||
#include "TinyGPUDriverUserClient.h"
|
||||
#include <AudioDriverKit/AudioDriverKit.h>
|
||||
#include <DriverKit/IOUserServer.h>
|
||||
#include <DriverKit/IOLib.h>
|
||||
#include <DriverKit/OSString.h>
|
||||
#include <DriverKit/IOMemoryMap.h>
|
||||
#include <DriverKit/IODMACommand.h>
|
||||
#include <DriverKit/IODispatchQueue.h>
|
||||
#include <PCIDriverKit/PCIDriverKit.h>
|
||||
#include <DriverKit/OSAction.h>
|
||||
|
||||
struct TinyGPUDriver_IVars
|
||||
{
|
||||
@@ -34,9 +27,6 @@ bool TinyGPUDriver::init()
|
||||
|
||||
void TinyGPUDriver::free()
|
||||
{
|
||||
if (ivars != nullptr) {
|
||||
|
||||
}
|
||||
IOSafeDeleteNULL(ivars, TinyGPUDriver_IVars, 1);
|
||||
super::free();
|
||||
}
|
||||
@@ -64,32 +54,6 @@ kern_return_t TinyGPUDriver::Start_Impl(IOService* in_provider)
|
||||
ivars->pci->ConfigurationRead16(kIOPCIConfigurationOffsetDeviceID, &dev);
|
||||
os_log(OS_LOG_DEFAULT, "tinygpu: opened device ven=0x%04x dev=0x%04x", ven, dev);
|
||||
|
||||
#if 0
|
||||
uint32_t off = 0x100;
|
||||
while (off) {
|
||||
uint32_t hdr = 0, next = 0, cap_id = 0;
|
||||
ivars->pci->ConfigurationRead32(off, &hdr);
|
||||
cap_id = hdr & 0xFFFFu;
|
||||
next = (hdr >> 20) & 0xFFCu;
|
||||
os_log(OS_LOG_DEFAULT, "tinygpu: cap: %u", cap_id);
|
||||
if (cap_id == 0x15) {
|
||||
uint32_t cap = 0, ctrl = 0;
|
||||
ivars->pci->ConfigurationRead32(off+0x4, &cap);
|
||||
ivars->pci->ConfigurationRead32(off+0x8, &ctrl);
|
||||
|
||||
uint32_t new_bar_size = 31 - __builtin_clz(cap >> 4);
|
||||
uint32_t new_ctrl = (ctrl & ~0x1f00) | (new_bar_size << 8);
|
||||
ivars->pci->ConfigurationWrite32(off+0x8, new_ctrl);
|
||||
|
||||
os_log(OS_LOG_DEFAULT, "tinygpu: rebar: cap=%u ctrl=%u new_bar_size=%u new_ctrl=%u", cap, ctrl, new_bar_size, new_ctrl);
|
||||
ivars->pci->Reset(0);
|
||||
break;
|
||||
}
|
||||
off = next;
|
||||
}
|
||||
ivars->pci->Reset(kIOPCIDeviceResetTypeHotReset);
|
||||
#endif
|
||||
|
||||
uint16_t commandRegister;
|
||||
ivars->pci->ConfigurationRead16(kIOPCIConfigurationOffsetCommand, &commandRegister);
|
||||
commandRegister |= (kIOPCICommandIOSpace | kIOPCICommandBusMaster | kIOPCICommandMemorySpace);
|
||||
@@ -130,96 +94,66 @@ error:
|
||||
|
||||
kern_return_t TinyGPUDriver::MapBar(uint32_t bar, IOMemoryDescriptor** memory)
|
||||
{
|
||||
kern_return_t err = 0;
|
||||
uint8_t barMemoryIndex, barMemoryType;
|
||||
uint64_t barMemorySize;
|
||||
err = ivars->pci->GetBARInfo(bar, &barMemoryIndex, &barMemorySize, &barMemoryType);
|
||||
kern_return_t err = ivars->pci->GetBARInfo(bar, &barMemoryIndex, &barMemorySize, &barMemoryType);
|
||||
if (err) return err;
|
||||
os_log(OS_LOG_DEFAULT, "tinygpu: bar mapping %d idx=%d", bar, barMemoryIndex);
|
||||
return ivars->pci->_CopyDeviceMemoryWithIndex(barMemoryIndex, memory, this);
|
||||
}
|
||||
|
||||
os_log(OS_LOG_DEFAULT, "tinygpu: requested bar mapping %d, %d", bar, (uint32_t)barMemoryIndex);
|
||||
err = ivars->pci->_CopyDeviceMemoryWithIndex(barMemoryIndex, memory, this);
|
||||
return err;
|
||||
static kern_return_t WriteDMASegments(IOMemoryDescriptor* mem, IOAddressSegment* segments, uint32_t segCount,
|
||||
uint64_t mapOffset = 0, uint64_t mapSize = 0)
|
||||
{
|
||||
// write dma segments to mapped memory as [addr0, len0, addr1, len1, ..., 0, 0]
|
||||
|
||||
IOMemoryMap* map = nullptr;
|
||||
kern_return_t err = mem->CreateMapping(0, 0, 0, mapOffset, mapSize, &map);
|
||||
if (err || !map) return err ?: kIOReturnError;
|
||||
|
||||
uint64_t* out = (uint64_t*)map->GetAddress();
|
||||
for (uint32_t i = 0; i < segCount; i++) { out[i * 2] = segments[i].address; out[i * 2 + 1] = segments[i].length; }
|
||||
out[segCount * 2] = 0; out[segCount * 2 + 1] = 0;
|
||||
map->release();
|
||||
return 0;
|
||||
}
|
||||
|
||||
kern_return_t TinyGPUDriver::SetupDMA(IOMemoryDescriptor* memory, uint64_t size, IODMACommand** outCmd,
|
||||
IOAddressSegment* segments, uint32_t* segCount)
|
||||
{
|
||||
IODMACommandSpecification dmaSpec = {.options = 0, .maxAddressBits = 40};
|
||||
IODMACommand* dmaCmd = nullptr;
|
||||
|
||||
kern_return_t err = IODMACommand::Create(ivars->pci, kIODMACommandCreateNoOptions, &dmaSpec, &dmaCmd);
|
||||
if (err) { os_log(OS_LOG_DEFAULT, "tinygpu: DMA create failed err=%d", err); return err; }
|
||||
|
||||
uint64_t flags = kIOMemoryDirectionInOut;
|
||||
err = dmaCmd->PrepareForDMA(kIODMACommandPrepareForDMANoOptions, memory, 0, size, &flags, segCount, segments);
|
||||
if (err) { os_log(OS_LOG_DEFAULT, "tinygpu: PrepareForDMA failed err=%d", err); dmaCmd->release(); return err; }
|
||||
|
||||
*outCmd = dmaCmd;
|
||||
return 0;
|
||||
}
|
||||
|
||||
kern_return_t TinyGPUDriver::CreateDMA(size_t size, TinyGPUCreateDMAResp* dmaDesc)
|
||||
{
|
||||
kern_return_t err = 0;
|
||||
IOMemoryMap* memoryMap = nullptr;
|
||||
IOBufferMemoryDescriptor* sharedBuf = nullptr;
|
||||
kern_return_t err = IOBufferMemoryDescriptor::Create(kIOMemoryDirectionInOut, size, IOVMPageSize, &sharedBuf);
|
||||
if (err) { os_log(OS_LOG_DEFAULT, "tinygpu: alloc failed err=%d", err); return err; }
|
||||
|
||||
IODMACommand* dmaCmd = nullptr;
|
||||
uint64_t flags = kIOMemoryDirectionInOut;
|
||||
uint32_t segCount = 32;
|
||||
IOAddressSegment segments[32];
|
||||
IODMACommandSpecification dmaSpec = {
|
||||
.options = 0,
|
||||
.maxAddressBits = 40,
|
||||
};
|
||||
uint32_t segCount = 32;
|
||||
err = SetupDMA(sharedBuf, size, &dmaCmd, segments, &segCount);
|
||||
if (err) { sharedBuf->release(); return err; }
|
||||
|
||||
err = IOBufferMemoryDescriptor::Create(kIOMemoryDirectionInOut, size, IOVMPageSize, &sharedBuf);
|
||||
if (err) {
|
||||
os_log(OS_LOG_DEFAULT, "tinygpu: failed to alloc user buffer, err=%d", err);
|
||||
goto error;
|
||||
}
|
||||
|
||||
err = IODMACommand::Create(ivars->pci, kIODMACommandCreateNoOptions, &dmaSpec, &dmaCmd);
|
||||
if (err) {
|
||||
os_log(OS_LOG_DEFAULT, "tinygpu: failed to create dma command, err=%d", err);
|
||||
goto error;
|
||||
}
|
||||
|
||||
err = dmaCmd->PrepareForDMA(kIODMACommandPrepareForDMANoOptions, sharedBuf, 0, size,
|
||||
&flags, &segCount, segments);
|
||||
if (err) {
|
||||
os_log(OS_LOG_DEFAULT, "tinygpu: failed to prepare for dma, err=%d", err);
|
||||
goto error;
|
||||
}
|
||||
|
||||
// pass addresses to userland
|
||||
{
|
||||
// debug
|
||||
for (int i = 0; i < segCount; i++) {
|
||||
os_log(OS_LOG_DEFAULT, "tinygpu: new dma mapping (sz=0x%zx) %d 0x%llx 0x%llx", size, i, segments[i].address, segments[i].length);
|
||||
}
|
||||
|
||||
err = sharedBuf->CreateMapping(0, 0, 0, IOVMPageSize, IOVMPageSize, &memoryMap); // one page should be fine
|
||||
if (err) {
|
||||
os_log(OS_LOG_DEFAULT, "tinygpu: failed to map memory, err=%d", err);
|
||||
goto error;
|
||||
}
|
||||
|
||||
// Send back gpu addresses
|
||||
uint64_t* addr = (uint64_t*)memoryMap->GetAddress();
|
||||
for (int i = 0; i < segCount; i++) {
|
||||
addr[i * 2] = segments[i].address;
|
||||
addr[i * 2 + 1] = segments[i].length;
|
||||
}
|
||||
addr[segCount * 2] = 0;
|
||||
addr[segCount * 2 + 1] = 0;
|
||||
|
||||
// free memoryMap
|
||||
memoryMap->release();
|
||||
memoryMap = nullptr;
|
||||
}
|
||||
err = WriteDMASegments(sharedBuf, segments, segCount, IOVMPageSize, IOVMPageSize);
|
||||
if (err) { dmaCmd->CompleteDMA(kIODMACommandCompleteDMANoOptions); dmaCmd->release(); sharedBuf->release(); return err; }
|
||||
|
||||
dmaDesc->sharedBuf = sharedBuf;
|
||||
dmaDesc->dmaCmd = dmaCmd;
|
||||
os_log(OS_LOG_DEFAULT, "tinygpu: CreateDMA size=0x%zx segs=%u", size, segCount);
|
||||
return 0;
|
||||
|
||||
error:
|
||||
if (memoryMap) {
|
||||
memoryMap->release();
|
||||
memoryMap = nullptr;
|
||||
}
|
||||
if (dmaCmd) {
|
||||
dmaCmd->CompleteDMA(kIODMACommandCompleteDMANoOptions);
|
||||
dmaCmd->release();
|
||||
dmaCmd = nullptr;
|
||||
}
|
||||
if (sharedBuf) {
|
||||
sharedBuf->release();
|
||||
sharedBuf = nullptr;
|
||||
}
|
||||
return err;
|
||||
}
|
||||
|
||||
kern_return_t TinyGPUDriver::CfgRead(uint32_t off, uint32_t size, uint32_t* outVal)
|
||||
@@ -257,3 +191,8 @@ kern_return_t TinyGPUDriver::ResetDevice()
|
||||
ivars->pci->Reset(kIOPCIDeviceResetTypeFunctionReset);
|
||||
return 0;
|
||||
}
|
||||
|
||||
IOPCIDevice* TinyGPUDriver::GetPCI()
|
||||
{
|
||||
return ivars->pci;
|
||||
}
|
||||
|
||||
@@ -2,20 +2,16 @@
|
||||
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
|
||||
<plist version="1.0">
|
||||
<dict>
|
||||
<key>com.apple.developer.driverkit</key>
|
||||
<true/>
|
||||
<key>com.apple.developer.driverkit.transport.pci</key>
|
||||
<array>
|
||||
<dict>
|
||||
<key>IOPCIMatch</key>
|
||||
<string>0x00001002&0x0000FFFF</string>
|
||||
</dict>
|
||||
<dict>
|
||||
<key>IOPCIMatch</key>
|
||||
<string>0x000010de&0x0000FFFF</string>
|
||||
<key>IOPCIPrimaryMatch</key>
|
||||
<string>0xFFFFFFFF&0x00000000</string>
|
||||
</dict>
|
||||
</array>
|
||||
<key>com.apple.developer.driverkit.allow-any-userclient-access</key>
|
||||
<true/>
|
||||
<key>com.apple.developer.driverkit</key>
|
||||
<true/>
|
||||
</dict>
|
||||
</plist>
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user