Compare commits

..
Author SHA1 Message Date
George HotzandGitHub 66676a1ac0 Merge branch 'master' into llama_trainer 2025-11-27 09:37:44 -08:00
geohot ee3ed9e646 wip junk 2025-11-24 20:03:12 -08:00
geohot 62f98ef817 fused optim 2025-11-24 19:51:07 -08:00
geohot f6dcb9a777 why did fakedata have fakeweights? 2025-11-24 19:29:42 -08:00
geohot dfaaeb0720 improve llama trainer 2025-11-24 19:11:14 -08:00
454 changed files with 29771 additions and 67453 deletions
+3 -3
View File
@@ -70,13 +70,13 @@ runs:
uses: actions/cache@v4
with:
path: ~/.cache/tinygrad/downloads/
key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
key: downloads-cache-${{ inputs.key }}-${{ env.CACHE_VERSION }}
- name: Cache downloads (macOS)
if: inputs.key != '' && runner.os == 'macOS'
uses: actions/cache@v4
with:
path: ~/Library/Caches/tinygrad/downloads/
key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
key: osx-downloads-cache-${{ inputs.key }}-${{ env.CACHE_VERSION }}
# **** Python deps ****
@@ -298,7 +298,7 @@ runs:
- name: Install mesa (linux)
if: inputs.mesa == 'true' && runner.os == 'Linux'
shell: bash
run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa_cpu-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
- name: Install mesa (macOS)
if: inputs.mesa == 'true' && runner.os == 'macOS'
shell: bash
+101 -69
View File
@@ -14,12 +14,10 @@ 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:
@@ -41,45 +39,104 @@ 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: Regenerate autogen files
- name: Verify OpenCL autogen
run: |
rm tinygrad/runtime/autogen/opencl.py
mv tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
python3 -c "from tinygrad.runtime.autogen import opencl"
rm tinygrad/runtime/autogen/{cuda,nvrtc,nvjitlink,nv_570,nv}.py
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv"
rm tinygrad/runtime/autogen/{comgr,hsa,hip,amd_gpu,sqtt,rocprof}.py
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof"
rm tinygrad/runtime/autogen/am/{am,pm4_soc15,pm4_nv,sdma_4_0_0,sdma_5_0_0,sdma_6_0_0,smu_v13_0_0,smu_v14_0_2}.py
python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v14_0_2"
rm tinygrad/runtime/autogen/{libc,kfd,io_uring,ib,pci,vfio}.py
python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
rm tinygrad/runtime/autogen/llvm.py
python3 -c "from tinygrad.runtime.autogen import llvm"
rm tinygrad/runtime/autogen/webgpu.py
python3 -c "from tinygrad.runtime.autogen import webgpu"
rm tinygrad/runtime/autogen/{kgsl,qcom_dsp}.py
python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
rm tinygrad/runtime/autogen/libusb.py
python3 -c "from tinygrad.runtime.autogen import libusb"
rm tinygrad/runtime/autogen/mesa.py
python3 -c "from tinygrad.runtime.autogen import mesa"
rm tinygrad/runtime/autogen/avcodec.py
python3 -c "from tinygrad.runtime.autogen import avcodec"
REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
- name: Check for differences
diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
- name: Verify CUDA autogen
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
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 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/adreno.py /tmp/adreno.py.bak
mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
python3 -c "from tinygrad.runtime.autogen import kgsl, adreno, qcom_dsp"
diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
diff /tmp/adreno.py.bak tinygrad/runtime/autogen/adreno.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
REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
diff /tmp/libclang.py.bak tinygrad/runtime/autogen/libclang.py
autogen-mac:
name: In-tree Autogen (macos)
runs-on: macos-14
@@ -91,24 +148,11 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
llvm: 'true'
- name: Regenerate autogen files
- name: Verify macos autogen
run: |
rm tinygrad/runtime/autogen/metal.py
mv tinygrad/runtime/autogen/metal.py /tmp/metal.py.bak
LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
- 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
diff /tmp/metal.py.bak tinygrad/runtime/autogen/metal.py
autogen-comgr-3:
name: In-tree Autogen (comgr 3)
runs-on: ubuntu-24.04
@@ -127,20 +171,8 @@ 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: Regenerate autogen files
- name: Verify comgr (3) autogen
run: |
rm tinygrad/runtime/autogen/comgr_3.py
mv tinygrad/runtime/autogen/comgr_3.py /tmp/comgr_3.py.bak
python3 -c "from tinygrad.runtime.autogen import comgr_3"
- 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
diff /tmp/comgr_3.py.bak tinygrad/runtime/autogen/comgr_3.py
+260 -154
View File
@@ -14,6 +14,12 @@ on:
- update_benchmark
- update_benchmark_staging
workflow_dispatch:
inputs:
run_process_replay:
description: "Run process replay tests"
required: false
default: false
type: boolean
jobs:
testmacbenchmark:
@@ -33,7 +39,6 @@ jobs:
- name: Symlink models and datasets
run: |
mkdir -p weights
mkdir -p extra/disassemblers
ln -s ~/tinygrad/extra/disassemblers/applegpu extra/disassemblers/applegpu
ln -s ~/tinygrad/weights/sd-v1-4.ckpt weights/sd-v1-4.ckpt
ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
@@ -49,19 +54,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
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
- 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
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
- 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
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
# 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
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
- name: Run model inference benchmark
run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
run: METAL=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
run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
- name: Test tensor cores
run: METAL=1 python3.11 test/opt/test_tensor_cores.py
- name: Test AMX tensor cores
@@ -71,80 +76,54 @@ 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
run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
- name: Run Tensor Core GEMM (half)
run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py
run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
- name: Run Tensor Core GEMM (bfloat16)
run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py
run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
- 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
BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
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
- 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
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
- 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
BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4
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
- name: Run quantized LLaMA3
run: |
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
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
#- 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
# run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_four_gpu.txt
- name: Run GPT2
run: |
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
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
- name: Run GPT2 w HALF
run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
- 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
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
- 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
run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
# 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
# run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
#- 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
# 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
#- name: Run 10 CIFAR training steps w BF16
# run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py
# run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
# 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
- uses: actions/upload-artifact@v4
with:
name: Speed (Mac)
path: |
onnx_inference_speed.csv
- 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
testusbgpu:
name: UsbGPU Benchmark
env:
PYTHONPYCACHEPREFIX: /tmp/tiny_python_pycache
runs-on: [self-hosted, macOS]
timeout-minutes: 10
defaults:
run:
shell: bash -e -o pipefail {0}
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
# 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
- name: UsbGPU boot time
run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
- name: UsbGPU tiny tests
@@ -153,10 +132,38 @@ jobs:
run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
#- name: UsbGPU openpilot test
# run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
- name: UsbGPU (USB4/TB) boot time
run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
- name: UsbGPU (USB4/TB) tiny tests
run: PYTHONPATH=. NV=1 NV_IFACE=PCI NV_NAK=1 python3.11 test/test_tiny.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
testnvidiabenchmark:
name: tinybox green Benchmark
@@ -190,7 +197,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
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
- 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
@@ -201,58 +208,79 @@ 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
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
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
- 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
run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
- name: Run Tensor Core GEMM (NV)
run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
- 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
run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
# 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
# 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
- 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
BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
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
- 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
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
# - 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
# 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
# - 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
# 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
- 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
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
- 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
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
- 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
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
# - 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
# 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
# - 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
# 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
- name: Run Mixtral 8x7B
run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing
run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
- name: Run GPT2
run: |
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
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
- 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
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
- 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
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
- 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
@@ -290,31 +318,45 @@ jobs:
# TODO: too slow
# - 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 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
run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
# TODO: too slow
- name: Run 10 CIFAR training steps
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py
- name: Run 10 CIFAR training steps w HALF
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=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
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=1300 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
# - name: Run 10 CIFAR training steps w HALF
# run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
# - name: Run 10 CIFAR training steps w BF16
# run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
# TODO: too slow
# - name: Run 10 CIFAR training steps w winograd
# 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
- 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
# 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
# - 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
# - 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
- 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
- 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
#- 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
#- 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
- 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
run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 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
- 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
@@ -366,18 +408,16 @@ 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
# 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
- 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
run: AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
# TODO: this is flaky
# - name: Test tensor cores AMD_LLVM=1
# run: AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
- name: Run Tensor Core GEMM (AMD)
- name: Test tensor cores
run: |
AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
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
- 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 matmul_amd.txt
- name: Test AMD=1
run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
#- name: Test HIP=1
@@ -392,39 +432,62 @@ 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
- 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
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
# TODO: too slow
# - 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
- 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
BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
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
- 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
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
# - 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
# 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
# - 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
# 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
- 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
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
- 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
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
# - 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
# 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
#- 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
# 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
- name: Run Mixtral 8x7B
run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing
run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
- name: Run GPT2
run: |
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
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
- 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
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
- 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
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
- 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
@@ -461,20 +524,35 @@ 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
run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
# TODO: too slow
- name: Run 10 CIFAR training steps
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py
- name: Run 10 CIFAR training steps w HALF
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=2000 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
# - name: Run 10 CIFAR training steps w HALF
# run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=390 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
# - 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
# 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
# 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
- 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
- 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
# 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
# - 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
#- 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
#- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
# run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD 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_remote.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
train_cifar_six_gpu_remote.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
@@ -512,13 +590,20 @@ jobs:
run: test/external/process_replay/reset.py
- 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
- 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
#- 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
#- 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
- 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
run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 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
- 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
@@ -540,16 +625,18 @@ jobs:
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
- name: reset process replay
run: test/external/process_replay/reset.py
# - name: openpilot compile3 0.9.9 driving_vision
# run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
# - name: openpilot compile3 0.9.9 driving_policy
# run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
# - name: openpilot compile3 0.9.9 dmonitoring
# run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
- name: openpilot compile3 0.10.0 driving_policy
run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=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
- 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
@@ -610,7 +697,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
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
- name: Test AMD=1
run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
- name: Test DISK copy time
@@ -619,13 +706,23 @@ jobs:
run: |
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
# TODO: too slow
# - 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
# TODO: enable
# - 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
# 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
- 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
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
- 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
@@ -672,13 +769,22 @@ 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
- 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
- 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
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
# TODO: too slow
# - 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
#- 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
- 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
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
- 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
+115 -98
View File
@@ -1,11 +1,10 @@
name: Unit Tests
env:
# increment this when downloads substantially change to avoid the internet
CACHE_VERSION: '15'
CACHE_VERSION: '13'
CAPTURE_PROCESS_REPLAY: 1
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PYTHONPATH: ${{ github.workspace }}
IGNORE_OOB: 0
on:
push:
@@ -37,8 +36,6 @@ jobs:
name: Docs
runs-on: ubuntu-22.04
timeout-minutes: 10
env:
IGNORE_OOB: 1
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -74,7 +71,9 @@ jobs:
- name: Test Docs Build
run: python -m mkdocs build --strict
- name: Test Docs
run: python docs/abstractions3.py
run: |
python docs/abstractions2.py
python docs/abstractions3.py
- name: Test README
run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py && python README.py
- name: Test Quickstart
@@ -105,11 +104,15 @@ 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: custom tests
- name: My (custom) tests
run: python3 extra/torch_backend/test.py
- name: Test one op in torch tests
run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
@@ -218,6 +221,7 @@ 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
@@ -231,17 +235,16 @@ 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.14.10
pre-commit run ruff --all-files
pip3 install --upgrade --force-reinstall ruff==0.11.0
python3 -m ruff check .
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 --lineprecision-report .
python -m mypy --strict-equality --lineprecision-report .
cat lineprecision.txt
- name: Run TYPED=1
run: TYPED=1 python -c "import tinygrad"
# broken because of UPatAny
#- name: Run TYPED=1
# run: TYPED=1 python -c "import tinygrad"
unittest:
name: Unit Tests
@@ -257,13 +260,10 @@ 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
run: |
CPU=1 python test/unit/test_device.py TestRunAsModule.test_module_runs
CPU=1 python -m pytest -n=auto test/unit/ --durations=20 --deselect=test/unit/test_device.py::TestRunAsModule::test_module_runs
run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
- name: Run targetted tests on NULL backend
run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
# TODO: too slow
@@ -289,8 +289,8 @@ jobs:
python extra/optimization/extract_dataset.py
gzip -c /tmp/sops > extra/datasets/sops.gz
#DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
- name: Repo line count < 20000 lines
run: MAX_LINE_COUNT=20000 python sz.py
- name: Repo line count < 19000 lines
run: MAX_LINE_COUNT=19000 python sz.py
spec:
strategy:
@@ -310,7 +310,7 @@ jobs:
deps: testing_unit
python-version: '3.14'
- name: Test SPEC=2
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 }}
run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
fuzzing:
name: Fuzzing
@@ -447,7 +447,7 @@ jobs:
with:
key: onnxoptl
deps: testing
pydeps: "tensorflow==2.19"
pydeps: "tensorflow==2.15.1 tensorflow_addons"
python-version: '3.11'
opencl: 'true'
- name: Test ONNX (CL)
@@ -465,7 +465,7 @@ jobs:
- name: Test Bert training
run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
- name: Test llama 3 training
run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -473,8 +473,6 @@ jobs:
name: Test LLM
runs-on: ubuntu-24.04
timeout-minutes: 15
env:
IGNORE_OOB: 1
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -605,7 +603,9 @@ 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 --durations=20
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
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -641,7 +641,7 @@ jobs:
if: matrix.backend=='amdllvm'
run: python test/device/test_amd_llvm.py
- name: Run pytest (amd)
run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py --durations=20
run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py --durations=20
- name: Run pytest (amd)
run: python -m pytest test/external/external_test_am.py --durations=20
- name: Run TRANSCENDENTAL math
@@ -654,46 +654,6 @@ jobs:
- name: Run process replay tests
uses: ./.github/actions/process-replay
testamdasm:
name: AMD ASM IDE
runs-on: ubuntu-24.04
timeout-minutes: 10
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: rdna3-emu
deps: testing_minimal
amd: 'true'
python-version: '3.13'
- name: Verify AMD autogen is up to date
run: |
python -m extra.assembly.amd.pdf
git diff --exit-code extra/assembly/amd/autogen/
- name: Install LLVM 21
run: |
wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
sudo apt-get update
sudo apt-get install llvm-21 llvm-21-tools cloc
- name: RDNA3 Line Count
run: cloc --by-file extra/assembly/amd/*.py
- name: Install rocprof-trace-decoder
run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
- name: Run RDNA3 emulator tests
run: python -m pytest -n=auto extra/assembly/amd/ --durations 20
- name: Run RDNA3 emulator tests (AMD_LLVM=1)
run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
- name: Run RDNA3 dtype tests
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
- name: Run RDNA3 dtype tests (AMD_LLVM=1)
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
# TODO: run all once emulator is faster
- name: Run RDNA3 ops tests
run: SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril"
testnvidia:
strategy:
fail-fast: false
@@ -761,6 +721,71 @@ jobs:
- name: Run process replay tests
uses: ./.github/actions/process-replay
amdremote:
name: Linux (remote)
runs-on: ubuntu-22.04
timeout-minutes: 20
env:
REMOTE: 1
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: linux-remote
deps: testing_minimal
amd: 'true'
llvm: 'true'
opencl: 'true'
- name: Start remote server
run: |
start_server() {
systemd-run --user \
--unit="$1" \
--setenv=REMOTEDEV="$2" \
--setenv=MOCKGPU=1 \
--setenv=PYTHONPATH=. \
--setenv=PORT="$3" \
--working-directory="$(pwd)" \
python tinygrad/runtime/ops_remote.py
}
start_server "remote-server-amd-1" "AMD" 6667
start_server "remote-server-amd-2" "AMD" 6668
start_server "remote-server-gpu" "CL" 7667
start_server "remote-server-cpu" "CPU" 8667
- name: Check Device.DEFAULT and print some source
env:
HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
- name: Run REMOTE=1 Test (AMD)
env:
HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
run: |
python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py --durations 20
- name: Run REMOTE=1 Test (CL)
env:
HOST: 127.0.0.1:7667*6
run: |
python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py --durations 20
IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
- name: Run REMOTE=1 Test (CPU)
env:
HOST: 127.0.0.1:8667*6
run: |
python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py --durations 20
- name: Show remote server logs
if: always()
run: |
journalctl --user -u remote-server-amd-1 --no-pager
journalctl --user -u remote-server-amd-2 --no-pager
journalctl --user -u remote-server-gpu --no-pager
journalctl --user -u remote-server-cpu --no-pager
# ****** OSX Tests ******
testmetal:
@@ -781,8 +806,6 @@ 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
@@ -860,6 +883,30 @@ jobs:
- name: Test ONNX Runner (WEBGPU)
run: WEBGPU=1 python3 test/external/external_test_onnx_runner.py
osxremote:
name: MacOS (remote metal)
runs-on: macos-15
timeout-minutes: 10
env:
REMOTE: 1
REMOTEDEV: METAL
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: macos-remote
deps: testing_minimal
- name: Check Device.DEFAULT and print some source
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
- name: Run REMOTE=1 Test
run: |
python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
osxtests:
strategy:
fail-fast: false
@@ -925,33 +972,3 @@ jobs:
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
# ****** Compile-only Tests ******
compiletests:
strategy:
fail-fast: false
matrix:
backend: [ir3, nak]
name: Compile-only (${{ matrix.backend }})
runs-on: ubuntu-24.04
timeout-minutes: 15
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: compile-${{ matrix.backend }}
deps: testing_minimal
mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
python-version: '3.14'
- name: Set env
shell: bash
run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
- name: Run test_ops
shell: bash
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
DEBUG=4 python3 test/test_ops.py TestOps.test_add
python -m pytest -n=auto test/test_ops.py --durations=20
+3 -3
View File
@@ -16,7 +16,7 @@ repos:
pass_filenames: false
- id: mypy
name: mypy
entry: python3 -m mypy
entry: python3 -m mypy tinygrad/ --strict-equality
language: system
always_run: true
pass_filenames: false
@@ -27,8 +27,8 @@ repos:
always_run: true
pass_filenames: false
- id: tests
name: comprehensive test suite
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
name: subset of tests
entry: env OMP_NUM_THREADS=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
language: system
always_run: true
pass_filenames: false
-227
View File
@@ -1,227 +0,0 @@
# Claude Code Guide for tinygrad
## Architecture Overview
tinygrad compiles tensor operations into optimized kernels. The pipeline:
1. **Tensor** (`tensor.py`) - User-facing API, creates UOp graph
2. **UOp** (`uop/ops.py`) - Unified IR for all operations (both tensor and kernel level)
3. **Schedule** (`engine/schedule.py`, `schedule/`) - Converts tensor UOps to kernel UOps
4. **Codegen** (`codegen/`) - Converts kernel UOps to device code
5. **Runtime** (`runtime/`) - Device-specific execution
## Key Concepts
### UOp (Universal Operation)
Everything is a UOp - tensors, operations, buffers, kernels. Key properties:
- `op`: The operation type (Ops enum)
- `dtype`: Data type
- `src`: Tuple of source UOps
- `arg`: Operation-specific argument
- `tag`: Optional tag for graph transformations
UOps are **immutable and cached** - creating the same UOp twice returns the same object (ucache).
### PatternMatcher
Used extensively for graph transformations:
```python
pm = PatternMatcher([
(UPat(Ops.ADD, src=(UPat.cvar("x"), UPat.cvar("x"))), lambda x: x * 2),
])
result = graph_rewrite(uop, pm)
```
### Schedule Cache
Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.
## Testing
```bash
# Run specific test
python -m pytest test/unit/test_schedule_cache.py -xvs
# Run with timeout
python -m pytest test/test_symbolic_ops.py -x --timeout=60
# Debug with print
DEBUG=2 python -m pytest test/test_schedule.py::test_name -xvs
# Visualize UOp graphs
VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
```
## Common Environment Variables
- `DEBUG=1-7` - Increasing verbosity (7 shows assembly output)
- `VIZ=1` - Enable graph visualization
- `SPEC=1` - Enable UOp spec verification
- `NOOPT=1` - Disable optimizations
- `DEVICE=CPU/CUDA/AMD/METAL` - Set default device
## Debugging Tips
1. **Print UOp graphs**: `print(tensor.uop)` or `print(tensor.uop.sink())`
2. **Check schedule**: `tensor.schedule()` returns list of ExecItems
3. **Trace graph rewrites**: Use `VIZ=1` or add print in PatternMatcher callbacks
4. **Find UOps by type**: `[u for u in uop.toposort() if u.op is Ops.SOMETHING]`
## Workflow Rules
- **NEVER commit without explicit user approval** - always show the diff and wait for approval
- **NEVER amend commits** - always create a new commit instead
- Run `pre-commit run --all-files` before committing to catch linting/type errors
- Run tests before proposing commits
- Test with `SPEC=2` when modifying UOp-related code
## Auto-generated Files (DO NOT EDIT)
The following files are auto-generated and should never be edited manually:
- `extra/assembly/amd/autogen/{arch}/__init__.py` - Generated by `python -m extra.assembly.amd.dsl --arch {arch}`
- `extra/assembly/amd/autogen/{arch}/gen_pcode.py` - Generated by `python -m extra.assembly.amd.pcode --arch {arch}`
Where `{arch}` is one of: `rdna3`, `rdna4`, `cdna`
To add missing instruction implementations, add them to `extra/assembly/amd/emu.py` instead.
## Style Notes
- 2-space indentation, 150 char line limit
- PatternMatchers should be defined at module level (slow to construct)
- Prefer `graph_rewrite` over manual graph traversal
- UOp methods like `.replace()` preserve tags unless explicitly changed
- Use `.rtag(value)` to add tags to UOps
## Lessons Learned
### UOp ucache Behavior
UOps are cached by their contents - creating a UOp with identical (op, dtype, src, arg) returns the **same object**. This means:
- `uop.replace(tag=None)` on a tagged UOp returns the original untagged UOp if it exists in cache
- Two UOps with same structure are identical (`is` comparison works)
### Spec Validation
When adding new UOp patterns, update `tinygrad/uop/spec.py`. Test with:
```bash
SPEC=2 python3 test/unit/test_something.py
```
Spec issues appear as `RuntimeError: SPEC ISSUE None: UOp(...)`.
### Schedule Cache Key Normalization
The schedule cache strips values from BIND nodes so different bound values (e.g., KV cache positions) hit the same cache entry:
- `pm_pre_sched_cache`: BIND(DEFINE_VAR, CONST) → BIND(DEFINE_VAR) for cache key
- `pm_post_sched_cache`: restores original BIND from context
- When accessing `bind.src[1]`, check `len(bind.src) > 1` first (might be stripped)
- Extract var_vals from `input_buffers` dict after graph_rewrite (avoids extra toposort)
### Avoiding Extra Work
- Use ctx dict from graph_rewrite to collect info during traversal instead of separate toposort
- Only extract var_vals when schedule is non-empty (no kernels = no vars needed)
- PatternMatchers are slow to construct - define at module level, not in functions
### Readability Over Speed
Don't add complexity for marginal performance gains. Simpler code that's slightly slower is often better:
```python
# BAD: "optimized" with extra complexity
if has_afters: # skip toposort if no AFTERs
after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
# GOOD: simple, always works
after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
```
The conditional check adds complexity, potential bugs, and often negligible speedup. Only optimize when profiling shows a real bottleneck.
### Testing LLM Changes
```bash
# Quick smoke test
echo "Hello" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b"
# Check cache hits (should see "cache hit" after warmup)
echo "Hello world" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b" 2>&1 | grep cache
# Test with beam search
echo "Hello" | BEAM=2 python tinygrad/apps/llm.py --model "llama3.2:1b"
```
## Common Patterns
### Graph Transformation
```python
def my_transform(ctx, x):
# Return new UOp or None to skip
return x.replace(arg=new_arg)
pm = PatternMatcher([
(UPat(Ops.SOMETHING, name="x"), my_transform),
])
result = graph_rewrite(input_uop, pm, ctx={})
```
### Finding Variables
```python
# Get all variables in a UOp graph
variables = uop.variables()
# Get bound variable values
var, val = bind_uop.unbind()
```
### Shape Handling
```python
# Shapes can be symbolic (contain UOps)
shape = tensor.shape # tuple[sint, ...] where sint = int | UOp
```
## Performance Optimization
When optimizing tinygrad internals:
1. **Measure wall time, not just call counts** - Reducing `graph_rewrite` calls doesn't always improve wall time. The overhead of conditional checks can exceed the cost of the operation being skipped.
2. **Profile each optimization individually** - Run benchmarks with and without each change to measure actual impact. Use `test/external/external_benchmark_schedule.py` for schedule/rewrite timing.
3. **Early exits in hot paths are effective** - Simple checks like `if self.op is Ops.CONST: return self` in `simplify()` can eliminate many unnecessary `graph_rewrite` calls.
4. **`graph_rewrite` is expensive** - Each call has overhead even for small graphs. Avoid calling it when the result is trivially known (e.g., simplifying a CONST returns itself).
5. **Beware iterator overhead** - Checks like `all(x.op is Ops.CONST for x in self.src)` can be slower than just running the operation, especially for small sequences.
6. **Verify cache hit rates before adding/keeping caches** - Measure actual hit rates with real workloads. A cache with 0% hit rate is pure overhead (e.g., `pm_cache` was removed because the algorithm guarantees each UOp is only passed to `pm_rewrite` once).
7. **Use `TRACK_MATCH_STATS=2` to profile pattern matching** - This shows match rates and time per pattern. Look for patterns with 0% match rate that still cost significant time - these are pure overhead for that workload.
8. **Cached properties beat manual traversal** - `backward_slice` uses `@functools.cached_property`. A DFS with early-exit sounds faster but is actually slower because it doesn't benefit from caching. The cache hit benefit often outweighs algorithmic improvements.
9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.
## Pattern Matching Analysis
**Use the right tool:**
- `TRACK_MATCH_STATS=2` - **Profiling**: identify expensive patterns
- `VIZ=-1` - **Inspection**: see all transformations, what every match pattern does, the before/after diffs
```bash
TRACK_MATCH_STATS=2 PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
```
Output format: `matches / attempts -- match_time / total_time ms -- location`
Key patterns to watch (from ResNet50 benchmark):
- `split_load_store`: ~146ms, 31% match rate - does real work
- `simplify_valid`: ~75ms, 0% match rate in this workload - checks AND ops for INDEX in backward slice
- `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches
Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.
```bash
# Save the trace
VIZ=-1 python test/test_tiny.py TestTiny.test_gemm
# Explore it
./extra/viz/cli.py --help
```
## AMD Performance Counter Profiling
Set VIZ to `-2` to save performance counters traces for the AMD backend.
Use the CLI in `./extra/sqtt/roc.py` to explore the trace.
+135
View File
@@ -0,0 +1,135 @@
# tinygrad is a tensor library, and as a tensor library it has multiple parts
# 1. a "runtime". this allows buffer management, compilation, and running programs
# 2. a "Device" that uses the runtime but specifies compute in an abstract way for all
# 3. a "UOp" that fuses the compute into kernels, using memory only when needed
# 4. a "Tensor" that provides an easy to use frontend with autograd ".backward()"
print("******** first, the runtime ***********")
from tinygrad.runtime.ops_cpu import ClangJITCompiler, CPUDevice, CPUProgram
cpu = CPUDevice()
# allocate some buffers
out = cpu.allocator.alloc(4)
a = cpu.allocator.alloc(4)
b = cpu.allocator.alloc(4)
# load in some values (little endian)
cpu.allocator._copyin(a, memoryview(bytearray([2,0,0,0])))
cpu.allocator._copyin(b, memoryview(bytearray([3,0,0,0])))
# compile a program to a binary
lib = ClangJITCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")
# create a runtime for the program
fxn = cpu.runtime("add", lib)
# run the program
fxn(out, a, b)
# check the data out
print(val := cpu.allocator._as_buffer(out).cast("I").tolist()[0])
assert val == 5
print("******** second, the Device ***********")
DEVICE = "CPU" # NOTE: you can change this!
import struct
from tinygrad.dtype import dtypes
from tinygrad.device import Buffer, Device
from tinygrad.uop.ops import UOp, Ops
# allocate some buffers + load in values
out = Buffer(DEVICE, 1, dtypes.int32).allocate()
a = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
# NOTE: a._buf is the same as the return from cpu.allocator.alloc
# describe the computation
idx = UOp.const(dtypes.index, 0)
buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
alu = buf_1.index(idx) + buf_2.index(idx)
output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.index(idx), alu))
s = UOp(Ops.SINK, dtypes.void, (st_0,))
# convert the computation to a "linearized" format (print the format)
from tinygrad.engine.realize import get_program, CompiledRunner
program = get_program(s, Device[DEVICE].renderer)
# compile a program (and print the source)
fxn = CompiledRunner(program)
print(fxn.p.src)
# NOTE: fxn.clprg is the CPUProgram
# run the program
fxn.exec([out, a, b])
# check the data out
assert out.as_buffer().cast('I')[0] == 5
print("******** third, the UOp ***********")
from tinygrad.engine.realize import run_schedule
from tinygrad.engine.schedule import create_schedule_with_vars
from tinygrad.schedule.rangeify import get_rangeify_map
# allocate some values + load in values
a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
b = UOp.new_buffer(DEVICE, 1, dtypes.int32)
a.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
b.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
# describe the computation
out = a + b
s = UOp(Ops.SINK, dtypes.void, (out,))
# group the computation into kernels
becomes_map = get_rangeify_map(s)
# the compute maps to an assign
assign = becomes_map[a+b].base
# the first source is the output buffer (data)
assert assign.src[0].op is Ops.BUFFER
# the second source is the kernel (compute)
assert assign.src[1].op is Ops.KERNEL
# schedule the kernel graph in a linear list
s = UOp(Ops.SINK, dtypes.void, (assign,))
sched, _ = create_schedule_with_vars(s)
assert len(sched) == 1
# DEBUGGING: print the compute ast
print(sched[-1].ast)
# NOTE: sched[-1].ast is the same as st_0 above
# the output will be stored in a new buffer
out = assign.buf_uop
assert out.op is Ops.BUFFER and not out.buffer.is_allocated()
print(out)
# run that schedule
run_schedule(sched)
# check the data out
assert out.is_realized and out.buffer.as_buffer().cast('I')[0] == 5
print("******** fourth, the Tensor ***********")
from tinygrad import Tensor
a = Tensor([2], dtype=dtypes.int32, device=DEVICE)
b = Tensor([3], dtype=dtypes.int32, device=DEVICE)
out = a + b
# check the data out
print(val:=out.item())
assert val == 5
+11 -5
View File
@@ -38,19 +38,25 @@ optim.schedule_step() # this will step the optimizer without running realize
# The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
# l1.uop and l2.uop define a computation graph
from tinygrad.engine.schedule import ExecItem
schedule: List[ExecItem] = Tensor.schedule(l1, l2)
from tinygrad.engine.schedule import ScheduleItem
schedule: List[ScheduleItem] = Tensor.schedule(l1, l2)
print(f"The schedule contains {len(schedule)} items.")
for si in schedule: print(str(si)[:80])
# *****
# 4. Lower and run the schedule.
# 4. Lower a schedule.
for si in tqdm(schedule): si.run()
from tinygrad.engine.realize import lower_schedule_item, ExecItem
lowered: List[ExecItem] = [lower_schedule_item(si) for si in tqdm(schedule)]
# *****
# 5. Print the weight change
# 5. Run the schedule
for ei in tqdm(lowered): ei.run()
# *****
# 6. Print the weight change
print("first weight change\n", l1.numpy()-l1n)
print("second weight change\n", l2.numpy()-l2n)
+5 -5
View File
@@ -13,19 +13,19 @@ There's also a [doc describing speed](../developer/speed.md)
Everything in [Tensor](../tensor/index.md) is syntactic sugar around constructing a graph of [UOps](../developer/uop.md).
The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view. Inputs to a base can be either base or view, inputs to a view can only be a single base.
The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view (specified by a ShapeTracker). Inputs to a base can be either base or view, inputs to a view can only be a single base.
## Scheduling
The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of UOps into a list of `ExecItem`. One `ExecItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of UOps into a list of `ScheduleItem`. One `ScheduleItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
::: tinygrad.engine.schedule.ExecItem
::: tinygrad.engine.schedule.ScheduleItem
## Lowering
The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers `ExecItem` by populating its `prg` field with
The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers `ScheduleItem` to `ExecItem` with
::: tinygrad.engine.realize.run_schedule
::: tinygrad.engine.realize.lower_schedule
There's a ton of complexity hidden behind this, see the `codegen/` directory.
+2 -2
View File
@@ -26,9 +26,9 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi
## tinygrad/codegen
Transform the optimized ast into a linearized and rendered program.
Transform the optimized ast into a linearized list of UOps.
::: tinygrad.codegen.get_program
::: tinygrad.codegen.full_rewrite
options:
members: false
show_labels: false
+1 -1
View File
@@ -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 interface for asm24xx chips.
* `USB`: USB3 interafce 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.
+9
View File
@@ -0,0 +1,9 @@
import globals from "globals";
import pluginJs from "@eslint/js";
import pluginHtml from "eslint-plugin-html";
export default [
{files: ["**/*.html"], plugins: {html: pluginHtml}, rules:{"max-len": ["error", {"code": 150}]}},
{languageOptions: {globals: globals.browser}},
pluginJs.configs.recommended,
];
+1 -1
View File
@@ -21,7 +21,7 @@ if __name__ == "__main__":
X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))
model = Model()
opt = (nn.optim.Muon if getenv("MUON") else nn.optim.SGD if getenv("SGD") else nn.optim.Adam)(nn.state.get_parameters(model))
opt = (nn.optim.Adam if not getenv("MUON") else nn.optim.Muon)(nn.state.get_parameters(model))
@TinyJit
@Tensor.train()
+93
View File
@@ -0,0 +1,93 @@
#!/usr/bin/env python3
import os, sys, traceback
sys.path.append(os.getcwd())
from io import StringIO
from contextlib import redirect_stdout
from tinygrad import Tensor, nn
from tinygrad.helpers import Timing, colored, getenv, fetch
from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
from sentencepiece import SentencePieceProcessor
def create_fixed_tokenizer(output_file):
print("creating fixed tokenizer")
import extra.junk.sentencepiece_model_pb2 as spb2
mp = spb2.ModelProto()
mp.ParseFromString(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/tokenizer.model?download=true").read_bytes())
mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_end|>", score=0))
mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_start|>", score=0))
with open(output_file, "wb") as f:
f.write(mp.SerializeToString())
# example:
# echo -en "write 2+2\nwrite hello world\ny\n" | TEMP=0 python3 examples/coder.py
if __name__ == "__main__":
# https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/config.json
with Timing("create model: "):
model = Transformer(4096, 14336, n_heads=32, n_layers=32, norm_eps=1e-5, vocab_size=32002, n_kv_heads=8, max_context=4096, jit=getenv("JIT", 1))
with Timing("download weights: "):
part1 = nn.state.torch_load(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/pytorch_model-00001-of-00002.bin?download=true"))
part2 = nn.state.torch_load(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/pytorch_model-00002-of-00002.bin?download=true"))
with Timing("weights -> model: "):
nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part1, 32, 32, 8)), strict=False)
nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part2, 32, 32, 8)), strict=False)
if not os.path.isfile("/tmp/tokenizer.model"): create_fixed_tokenizer("/tmp/tokenizer.model")
spp = SentencePieceProcessor(model_file="/tmp/tokenizer.model")
# https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/tokenizer_config.json
# "chat_template": "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
IM_END = 32000
IM_START = 32001
def encode_prompt(k, v): return [IM_START]+spp.encode(f"{k}\n{v}")+[IM_END]+spp.encode("\n")
def start_prompt(k): return [IM_START]+spp.encode(f"{k}\n")
def output(outputted, toks, color):
cur = spp.decode(toks)[len(outputted):]
sys.stdout.write(colored(cur, color))
sys.stdout.flush()
outputted += cur
return outputted
# *** app below this line ***
toks = [spp.bos_id()] + encode_prompt("system", "You are Quentin. Quentin is a useful assistant who writes Python code to answer questions. He keeps the code as short as possible and doesn't read from user input")
PROMPT = getenv("PROMPT", 1)
temperature = getenv("TEMP", 0.7)
start_pos = 0
outputted = output("", toks, "green")
turn = True
while 1:
if PROMPT:
toks += encode_prompt("user", input("Q: ")) + start_prompt("assistant")
else:
toks += start_prompt("user" if turn else "assistant")
turn = not turn
old_output_len = len(outputted)
while 1:
tok = model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
start_pos = len(toks)
toks.append(tok)
outputted = output(outputted, toks, "blue" if not turn else "cyan")
if tok == IM_END: break
if tok == spp.eos_id(): break
new_output = outputted[old_output_len:]
if new_output.endswith("```") and '```python\n' in new_output:
python_code = new_output.split('```python\n')[1].split("```")[0]
# AI safety. Warning to user. Do not press y if the AI is trying to do unsafe things.
if input(colored(f" <-- PYTHON DETECTED, RUN IT? ", "red")).lower() == 'y':
my_stdout = StringIO()
try:
with redirect_stdout(my_stdout): exec(python_code)
result = my_stdout.getvalue()
except Exception as e:
result = ''.join(traceback.format_exception_only(e))
toks += spp.encode(f"\nOutput:\n```\n{result}```")
outputted = output(outputted, toks, "yellow")
old_output_len = len(outputted)
print("")
+341
View File
@@ -0,0 +1,341 @@
import argparse
import multiprocessing as mp
import os
import re
import sys
import time
from contextlib import contextmanager
from pathlib import Path
import numpy as np
import pyaudio
import yaml
from llama import LLaMa
from vits import MODELS as VITS_MODELS
from vits import Y_LENGTH_ESTIMATE_SCALARS, HParams, Synthesizer, TextMapper, get_hparams_from_file, load_model
from whisper import init_whisper, transcribe_waveform
from sentencepiece import SentencePieceProcessor
from tinygrad.helpers import Timing, fetch
from tinygrad import Tensor, dtypes
# Whisper constants
RATE = 16000
CHUNK = 1600
# LLaMa constants
IM_START = 32001
IM_END = 32002
# Functions for encoding prompts to chatml md
def encode_prompt(spp, k, v): return [IM_START]+spp.encode(f"{k}\n{v}")+[IM_END]+spp.encode("\n")
def start_prompt(spp, k): return [IM_START]+spp.encode(f"{k}\n")
def chunks(lst, n):
for i in range(0, len(lst), n): yield lst[i:i + n]
def create_fixed_tokenizer():
"""Function needed for extending tokenizer with additional chat tokens"""
import extra.junk.sentencepiece_model_pb2 as spb2
tokenizer_path = fetch("https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/resolve/main/tokenizer.model")
if SentencePieceProcessor(model_file=str(tokenizer_path)).vocab_size() != 32003:
print("creating fixed tokenizer")
mp = spb2.ModelProto()
mp.ParseFromString(tokenizer_path.read_bytes())
# https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/blob/main/added_tokens.json
mp.pieces.append(spb2.ModelProto.SentencePiece(piece="[PAD]", score=0))
mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_start|>", score=0))
mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_end|>", score=0))
tokenizer_path.write_bytes(mp.SerializeToString())
return tokenizer_path
def llama_prepare(llama: LLaMa, temperature: float, pre_prompt_path: Path) -> tuple[list[int], str, str, str]:
"""Prepares a llama model from a specified pre-prompt file"""
with open(str(pre_prompt_path)) as f:
config = yaml.safe_load(f.read())
toks = [llama.tokenizer.bos_id()] + encode_prompt(llama.tokenizer, "system", config["pre_prompt"].replace("\n", " "))
for i in config["examples"]:
toks += encode_prompt(llama.tokenizer, config["user_delim"], i["user_prompt"])
toks += encode_prompt(llama.tokenizer, config["resp_delim"], i["resp_prompt"])
llama.model(Tensor([toks]), 0, temperature).realize() # NOTE: outputs are not used
return toks, config["user_delim"], config["resp_delim"], len(toks), llama.tokenizer.decode(toks)
def llama_generate(
llama: LLaMa,
toks: list[int],
outputted: str,
prompt: str,
start_pos: int,
user_delim: str,
resp_delim: str,
temperature=0.7,
max_tokens=1000
):
"""Generates an output for the specified prompt"""
toks += encode_prompt(llama.tokenizer, user_delim, prompt)
toks += start_prompt(llama.tokenizer, resp_delim)
outputted = llama.tokenizer.decode(toks)
init_length = len(outputted)
for _ in range(max_tokens):
token = llama.model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
start_pos = len(toks)
toks.append(token)
cur = llama.tokenizer.decode(toks)
# Print is just for debugging
sys.stdout.write(cur[len(outputted):])
sys.stdout.flush()
outputted = cur
if toks[-1] == IM_END: break
else:
toks.append(IM_END)
print() # because the output is flushed
return outputted, start_pos, outputted[init_length:].replace("<|im_end|>", "")
def tts(
text_to_synthesize: str,
synth: Synthesizer,
hps: HParams,
emotion_embedding: Path,
speaker_id: int,
model_to_use: str,
noise_scale: float,
noise_scale_w: float,
length_scale: float,
estimate_max_y_length: bool,
text_mapper: TextMapper,
model_has_multiple_speakers: bool,
pad_length=600,
vits_pad_length=1000
):
if model_to_use == "mmts-tts": text_to_synthesize = text_mapper.filter_oov(text_to_synthesize.lower())
# Convert the input text to a tensor.
stn_tst = text_mapper.get_text(text_to_synthesize, hps.data.add_blank, hps.data.text_cleaners)
init_shape = stn_tst.shape
assert init_shape[0] < pad_length, "text is too long"
x_tst, x_tst_lengths = stn_tst.pad(((0, pad_length - init_shape[0]),), value=1).unsqueeze(0), Tensor([init_shape[0]], dtype=dtypes.int64)
sid = Tensor([speaker_id], dtype=dtypes.int64) if model_has_multiple_speakers else None
# Perform inference.
audio_tensor = synth.infer(x_tst, x_tst_lengths, sid, noise_scale, length_scale, noise_scale_w, emotion_embedding=emotion_embedding,
max_y_length_estimate_scale=Y_LENGTH_ESTIMATE_SCALARS[model_to_use] if estimate_max_y_length else None, pad_length=vits_pad_length)[0, 0]
# Save the audio output.
audio_data = (np.clip(audio_tensor.numpy(), -1.0, 1.0) * 32767).astype(np.int16)
return audio_data
def init_vits(
model_to_use: str,
emotion_path: Path,
speaker_id: int,
seed: int,
):
model_config = VITS_MODELS[model_to_use]
# Load the hyperparameters from the config file.
hps = get_hparams_from_file(fetch(model_config[0]))
# If model has multiple speakers, validate speaker id and retrieve name if available.
model_has_multiple_speakers = hps.data.n_speakers > 0
if model_has_multiple_speakers:
if speaker_id >= hps.data.n_speakers: raise ValueError(f"Speaker ID {speaker_id} is invalid for this model.")
if hps.__contains__("speakers"): # maps speaker ids to names
speakers = hps.speakers
if isinstance(speakers, list): speakers = {speaker: i for i, speaker in enumerate(speakers)}
# Load emotions if any. TODO: find an english model with emotions, this is untested atm.
emotion_embedding = None
if emotion_path is not None:
if emotion_path.endswith(".npy"): emotion_embedding = Tensor(np.load(emotion_path), dtype=dtypes.int64).unsqueeze(0)
else: raise ValueError("Emotion path must be a .npy file.")
# Load symbols, instantiate TextMapper and clean the text.
if hps.__contains__("symbols"): symbols = hps.symbols
elif model_to_use == "mmts-tts": symbols = [x.replace("\n", "") for x in fetch("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/vocab.txt").open(encoding="utf-8").readlines()]
else: symbols = ['_'] + list(';:,.!?¡¿—…"«»“” ') + list('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz') + list("ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'")
text_mapper = TextMapper(apply_cleaners=True, symbols=symbols)
# Load the model.
if seed is not None:
Tensor.manual_seed(seed)
np.random.seed(seed)
net_g = load_model(text_mapper.symbols, hps, model_config)
return net_g, emotion_embedding, text_mapper, hps, model_has_multiple_speakers
@contextmanager
def output_stream(num_channels: int, sample_rate: int):
try:
p = pyaudio.PyAudio()
stream = p.open(format=pyaudio.paInt16, channels=num_channels, rate=sample_rate, output=True)
yield stream
except KeyboardInterrupt: pass
finally:
stream.stop_stream()
stream.close()
p.terminate()
@contextmanager
def log_writer():
try:
logs = []
yield logs
finally:
sep = "="*os.get_terminal_size()[1]
print(f"{sep[:-1]}\nCHAT LOG")
print(*logs, sep="\n")
print(sep)
def listener(q: mp.Queue, event: mp.Event):
try:
p = pyaudio.PyAudio()
stream = p.open(format=pyaudio.paInt16, channels=1, rate=RATE, input=True, frames_per_buffer=CHUNK)
did_print = False
while True:
data = stream.read(CHUNK) # read data to avoid overflow
if event.is_set():
if not did_print:
print("listening")
did_print = True
q.put(((np.frombuffer(data, np.int16)/32768).astype(np.float32)*3))
else:
did_print = False
finally:
stream.stop_stream()
stream.close()
p.terminate()
def mp_output_stream(q: mp.Queue, counter: mp.Value, num_channels: int, sample_rate: int):
with output_stream(num_channels, sample_rate) as stream:
while True:
try:
stream.write(q.get())
counter.value += 1
except KeyboardInterrupt:
break
if __name__ == "__main__":
import nltk
nltk.download("punkt")
# Parse CLI arguments
parser = argparse.ArgumentParser("Have a tiny conversation with tinygrad")
# Whisper args
parser.add_argument("--whisper_model_name", type=str, default="tiny.en")
# LLAMA args
parser.add_argument("--llama_pre_prompt_path", type=Path, default=Path(__file__).parent / "conversation_data" / "pre_prompt_stacy.yaml", help="Path to yaml file which contains all pre-prompt data needed. ")
parser.add_argument("--llama_count", type=int, default=1000, help="Max number of tokens to generate")
parser.add_argument("--llama_temperature", type=float, default=0.7, help="Temperature in the softmax")
parser.add_argument("--llama_quantize", type=str, default=None, help="Quantize the weights to int8 or nf4 in memory")
parser.add_argument("--llama_model", type=Path, default=None, help="Folder with the original weights to load, or single .index.json, .safetensors or .bin file")
parser.add_argument("--llama_gen", type=str, default="tiny", required=False, help="Generation of the model to use")
parser.add_argument("--llama_size", type=str, default="1B-Chat", required=False, help="Size of model to use")
parser.add_argument("--llama_tokenizer", type=Path, default=None, required=False, help="Path to llama tokenizer.model")
# vits args
parser.add_argument("--vits_model_to_use", default="vctk", help="Specify the model to use. Default is 'vctk'.")
parser.add_argument("--vits_speaker_id", type=int, default=12, help="Specify the speaker ID. Default is 6.")
parser.add_argument("--vits_noise_scale", type=float, default=0.667, help="Specify the noise scale. Default is 0.667.")
parser.add_argument("--vits_noise_scale_w", type=float, default=0.8, help="Specify the noise scale w. Default is 0.8.")
parser.add_argument("--vits_length_scale", type=float, default=1, help="Specify the length scale. Default is 1.")
parser.add_argument("--vits_seed", type=int, default=None, help="Specify the seed (set to None if no seed). Default is 1337.")
parser.add_argument("--vits_num_channels", type=int, default=1, help="Specify the number of audio output channels. Default is 1.")
parser.add_argument("--vits_sample_width", type=int, default=2, help="Specify the number of bytes per sample, adjust if necessary. Default is 2.")
parser.add_argument("--vits_emotion_path", type=Path, default=None, help="Specify the path to emotion reference.")
parser.add_argument("--vits_estimate_max_y_length", type=str, default=False, help="If true, overestimate the output length and then trim it to the correct length, to prevent premature realization, much more performant for larger inputs, for smaller inputs not so much. Default is False.")
parser.add_argument("--vits_vocab_path", type=Path, default=None, help="Path to the TTS vocabulary.")
# conversation args
parser.add_argument("--max_sentence_length", type=int, default=20, help="Max words in one sentence to pass to vits")
args = parser.parse_args()
# Init models
model, enc = init_whisper(args.whisper_model_name)
synth, emotion_embedding, text_mapper, hps, model_has_multiple_speakers = init_vits(args.vits_model_to_use, args.vits_emotion_path, args.vits_speaker_id, args.vits_seed)
# Download tinyllama chat as a default model
if args.llama_model is None:
args.llama_model = fetch("https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/resolve/main/model.safetensors", "tinyllamachat.safetensors")
args.llama_gen = "tiny"
args.llama_size = "1B-Chat"
# Add 3 more tokens to the tokenizer
if args.llama_gen == "tiny" and args.llama_size.endswith("Chat"): args.llama_tokenizer = create_fixed_tokenizer()
tokenizer_path = args.llama_tokenizer or args.llama_model.parent / "tokenizer.model"
llama = LLaMa.build(args.llama_model, tokenizer_path, args.llama_gen, args.llama_size, args.llama_quantize)
toks, user_delim, resp_delim, start_pos, outputted = llama_prepare(llama, args.llama_temperature, args.llama_pre_prompt_path)
# Start child process for mic input
q = mp.Queue()
is_listening_event = mp.Event()
p = mp.Process(target=listener, args=(q, is_listening_event,))
p.daemon = True
p.start()
# Start child process for speaker output
out_q = mp.Queue()
out_counter = mp.Value("i", 0)
out_p = mp.Process(target=mp_output_stream, args=(out_q, out_counter, args.vits_num_channels, hps.data.sampling_rate,))
out_p.daemon = True
out_p.start()
# JIT tts
for i in ["Hello, I'm a chat bot", "I am capable of doing a lot of things"]:
tts(
i, synth, hps, emotion_embedding,
args.vits_speaker_id, args.vits_model_to_use, args.vits_noise_scale,
args.vits_noise_scale_w, args.vits_length_scale,
args.vits_estimate_max_y_length, text_mapper, model_has_multiple_speakers
)
# Start the pipeline
with log_writer() as log:
while True:
tokens = [enc._special_tokens["<|startoftranscript|>"], enc._special_tokens["<|notimestamps|>"]]
total = np.array([])
out_counter.value = 0
s = time.perf_counter()
is_listening_event.set()
prev_text = None
while True:
for _ in range(RATE // CHUNK): total = np.concatenate([total, q.get()])
txt = transcribe_waveform(model, enc, [total], truncate=True)
print(txt, end="\r")
if txt == "[BLANK_AUDIO]" or re.match(r"^\([\w+ ]+\)$", txt.strip()): continue
if prev_text is not None and prev_text == txt:
is_listening_event.clear()
break
prev_text = txt
print() # to avoid llama printing on the same line
log.append(f"{user_delim.capitalize()}: {txt}")
# Generate with llama
with Timing("llama generation: "):
outputted, start_pos, response = llama_generate(
llama, toks, outputted, txt, start_pos,
user_delim=user_delim, resp_delim=resp_delim, temperature=args.llama_temperature,
max_tokens=args.llama_count
)
log.append(f"{resp_delim.capitalize()}: {response}")
# Convert to voice
with Timing("tts: "):
sentences = nltk.sent_tokenize(response.replace('"', ""))
for i in sentences:
total = np.array([], dtype=np.int16)
for j in chunks(i.split(), args.max_sentence_length):
audio_data = tts(
" ".join(j), synth, hps, emotion_embedding,
args.vits_speaker_id, args.vits_model_to_use, args.vits_noise_scale,
args.vits_noise_scale_w, args.vits_length_scale,
args.vits_estimate_max_y_length, text_mapper, model_has_multiple_speakers
)
total = np.concatenate([total, audio_data])
out_q.put(total.tobytes())
while out_counter.value < len(sentences): continue
log.append(f"Total: {time.perf_counter() - s}")
+89
View File
@@ -0,0 +1,89 @@
# load weights from
# https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth
# a rough copy of
# https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/model.py
import sys
import ast
import time
import numpy as np
from PIL import Image
from tinygrad.tensor import Tensor
from tinygrad.helpers import getenv, fetch, Timing
from tinygrad.engine.jit import TinyJit
from extra.models.efficientnet import EfficientNet
np.set_printoptions(suppress=True)
# TODO: you should be able to put these in the jitted function
bias = Tensor([0.485, 0.456, 0.406])
scale = Tensor([0.229, 0.224, 0.225])
@TinyJit
def _infer(model, img):
img = img.permute((2,0,1))
img = img / 255.0
img = img - bias.reshape((1,-1,1,1))
img = img / scale.reshape((1,-1,1,1))
return model.forward(img).realize()
def infer(model, img):
# preprocess image
aspect_ratio = img.size[0] / img.size[1]
img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0))))
img = np.array(img)
y0,x0=(np.asarray(img.shape)[:2]-224)//2
retimg = img = img[y0:y0+224, x0:x0+224]
# if you want to look at the image
"""
import matplotlib.pyplot as plt
plt.imshow(img)
plt.show()
"""
# run the net
out = _infer(model, Tensor(img.astype("float32"))).numpy()
# if you want to look at the outputs
"""
import matplotlib.pyplot as plt
plt.plot(out[0])
plt.show()
"""
return out, retimg
if __name__ == "__main__":
# instantiate my net
model = EfficientNet(getenv("NUM", 0))
model.load_from_pretrained()
# category labels
lbls = ast.literal_eval(fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt").read_text())
# load image and preprocess
url = sys.argv[1] if len(sys.argv) >= 2 else "https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/showcase/stable_diffusion_by_tinygrad.jpg"
if url == 'webcam':
import cv2
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
while 1:
_ = cap.grab() # discard one frame to circumvent capture buffering
ret, frame = cap.read()
img = Image.fromarray(frame[:, :, [2,1,0]])
lt = time.monotonic_ns()
out, retimg = infer(model, img)
print(f"{(time.monotonic_ns()-lt)*1e-6:7.2f} ms", np.argmax(out), np.max(out), lbls[np.argmax(out)])
SCALE = 3
simg = cv2.resize(retimg, (224*SCALE, 224*SCALE))
retimg = cv2.cvtColor(simg, cv2.COLOR_RGB2BGR)
cv2.imshow('capture', retimg)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
else:
img = Image.open(fetch(url))
for i in range(getenv("CNT", 1)):
with Timing("did inference in "):
out, _ = infer(model, img)
print(np.argmax(out), np.max(out), lbls[np.argmax(out)])
+498
View File
@@ -0,0 +1,498 @@
# pip3 install sentencepiece
# This file incorporates code from the following:
# Github Name | License | Link
# black-forest-labs/flux | Apache | https://github.com/black-forest-labs/flux/tree/main/model_licenses
from tinygrad import Tensor, nn, dtypes, TinyJit
from tinygrad.nn.state import safe_load, load_state_dict
from tinygrad.helpers import fetch, tqdm, colored
from sdxl import FirstStage
from extra.models.clip import FrozenClosedClipEmbedder
from extra.models.t5 import T5Embedder
import numpy as np
import math, time, argparse, tempfile
from typing import List, Dict, Optional, Union, Tuple, Callable
from dataclasses import dataclass
from pathlib import Path
from PIL import Image
urls:dict = {
"flux-schnell": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors",
"flux-dev": "https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/flux1-dev.sft",
"ae": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors",
"T5_1_of_2": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/text_encoder_2/model-00001-of-00002.safetensors",
"T5_2_of_2": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/text_encoder_2/model-00002-of-00002.safetensors",
"T5_tokenizer": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/tokenizer_2/spiece.model",
"clip": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/text_encoder/model.safetensors"
}
def tensor_identity(x:Tensor) -> Tensor: return x
class AutoEncoder:
def __init__(self, scale_factor:float, shift_factor:float):
self.decoder = FirstStage.Decoder(128, 3, 3, 16, [1, 2, 4, 4], 2, 256)
self.scale_factor = scale_factor
self.shift_factor = shift_factor
def decode(self, z:Tensor) -> Tensor:
z = z / self.scale_factor + self.shift_factor
return self.decoder(z)
# Conditioner
class ClipEmbedder(FrozenClosedClipEmbedder):
def __call__(self, texts:Union[str, List[str], Tensor]) -> Tensor:
if isinstance(texts, str): texts = [texts]
assert isinstance(texts, (list,tuple)), f"expected list of strings, got {type(texts).__name__}"
tokens = Tensor.cat(*[Tensor(self.tokenizer.encode(text)) for text in texts], dim=0)
return self.transformer.text_model(tokens.reshape(len(texts),-1))[:, tokens.argmax(-1)]
# https://github.com/black-forest-labs/flux/blob/main/src/flux/math.py
def attention(q:Tensor, k:Tensor, v:Tensor, pe:Tensor) -> Tensor:
q, k = apply_rope(q, k, pe)
x = Tensor.scaled_dot_product_attention(q, k, v)
return x.rearrange("B H L D -> B L (H D)")
def rope(pos:Tensor, dim:int, theta:int) -> Tensor:
assert dim % 2 == 0
scale = Tensor.arange(0, dim, 2, dtype=dtypes.float32, device=pos.device) / dim # NOTE: this is torch.float64 in reference implementation
omega = 1.0 / (theta**scale)
out = Tensor.einsum("...n,d->...nd", pos, omega)
out = Tensor.stack(Tensor.cos(out), -Tensor.sin(out), Tensor.sin(out), Tensor.cos(out), dim=-1)
out = out.rearrange("b n d (i j) -> b n d i j", i=2, j=2)
return out.float()
def apply_rope(xq:Tensor, xk:Tensor, freqs_cis:Tensor) -> Tuple[Tensor, Tensor]:
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
return xq_out.reshape(*xq.shape).cast(xq.dtype), xk_out.reshape(*xk.shape).cast(xk.dtype)
# https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py
class EmbedND:
def __init__(self, dim:int, theta:int, axes_dim:List[int]):
self.dim = dim
self.theta = theta
self.axes_dim = axes_dim
def __call__(self, ids:Tensor) -> Tensor:
n_axes = ids.shape[-1]
emb = Tensor.cat(*[rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], dim=-3)
return emb.unsqueeze(1)
class MLPEmbedder:
def __init__(self, in_dim:int, hidden_dim:int):
self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
def __call__(self, x:Tensor) -> Tensor:
return self.out_layer(self.in_layer(x).silu())
class QKNorm:
def __init__(self, dim:int):
self.query_norm = nn.RMSNorm(dim)
self.key_norm = nn.RMSNorm(dim)
def __call__(self, q:Tensor, k:Tensor) -> Tuple[Tensor, Tensor]:
return self.query_norm(q), self.key_norm(k)
class SelfAttention:
def __init__(self, dim:int, num_heads:int = 8, qkv_bias:bool = False):
self.num_heads = num_heads
head_dim = dim // num_heads
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.norm = QKNorm(head_dim)
self.proj = nn.Linear(dim, dim)
def __call__(self, x:Tensor, pe:Tensor) -> Tensor:
qkv = self.qkv(x)
q, k, v = qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
q, k = self.norm(q, k)
x = attention(q, k, v, pe=pe)
return self.proj(x)
@dataclass
class ModulationOut:
shift:Tensor
scale:Tensor
gate:Tensor
class Modulation:
def __init__(self, dim:int, double:bool):
self.is_double = double
self.multiplier = 6 if double else 3
self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
def __call__(self, vec:Tensor) -> Tuple[ModulationOut, Optional[ModulationOut]]:
out = self.lin(vec.silu())[:, None, :].chunk(self.multiplier, dim=-1)
return ModulationOut(*out[:3]), ModulationOut(*out[3:]) if self.is_double else None
class DoubleStreamBlock:
def __init__(self, hidden_size:int, num_heads:int, mlp_ratio:float, qkv_bias:bool = False):
mlp_hidden_dim = int(hidden_size * mlp_ratio)
self.num_heads = num_heads
self.hidden_size = hidden_size
self.img_mod = Modulation(hidden_size, double=True)
self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.img_mlp = [nn.Linear(hidden_size, mlp_hidden_dim, bias=True), Tensor.gelu, nn.Linear(mlp_hidden_dim, hidden_size, bias=True)]
self.txt_mod = Modulation(hidden_size, double=True)
self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.txt_mlp = [nn.Linear(hidden_size, mlp_hidden_dim, bias=True), Tensor.gelu, nn.Linear(mlp_hidden_dim, hidden_size, bias=True)]
def __call__(self, img:Tensor, txt:Tensor, vec:Tensor, pe:Tensor) -> tuple[Tensor, Tensor]:
img_mod1, img_mod2 = self.img_mod(vec)
txt_mod1, txt_mod2 = self.txt_mod(vec)
assert img_mod2 is not None and txt_mod2 is not None
# prepare image for attention
img_modulated = self.img_norm1(img)
img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
img_qkv = self.img_attn.qkv(img_modulated)
img_q, img_k, img_v = img_qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
img_q, img_k = self.img_attn.norm(img_q, img_k)
# prepare txt for attention
txt_modulated = self.txt_norm1(txt)
txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
txt_qkv = self.txt_attn.qkv(txt_modulated)
txt_q, txt_k, txt_v = txt_qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k)
# run actual attention
q = Tensor.cat(txt_q, img_q, dim=2)
k = Tensor.cat(txt_k, img_k, dim=2)
v = Tensor.cat(txt_v, img_v, dim=2)
attn = attention(q, k, v, pe=pe)
txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
# calculate the img bloks
img = img + img_mod1.gate * self.img_attn.proj(img_attn)
img = img + img_mod2.gate * ((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift).sequential(self.img_mlp)
# calculate the txt bloks
txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
txt = txt + txt_mod2.gate * ((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift).sequential(self.txt_mlp)
return img, txt
class SingleStreamBlock:
"""
A DiT block with parallel linear layers as described in
https://arxiv.org/abs/2302.05442 and adapted modulation interface.
"""
def __init__(self,hidden_size:int, num_heads:int, mlp_ratio:float=4.0, qk_scale:Optional[float]=None):
self.hidden_dim = hidden_size
self.num_heads = num_heads
head_dim = hidden_size // num_heads
self.scale = qk_scale or head_dim**-0.5
self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
# qkv and mlp_in
self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
# proj and mlp_out
self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
self.norm = QKNorm(head_dim)
self.hidden_size = hidden_size
self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.mlp_act = Tensor.gelu
self.modulation = Modulation(hidden_size, double=False)
def __call__(self, x:Tensor, vec:Tensor, pe:Tensor) -> Tensor:
mod, _ = self.modulation(vec)
x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
qkv, mlp = Tensor.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
q, k, v = qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
q, k = self.norm(q, k)
# compute attention
attn = attention(q, k, v, pe=pe)
# compute activation in mlp stream, cat again and run second linear layer
output = self.linear2(Tensor.cat(attn, self.mlp_act(mlp), dim=2))
return x + mod.gate * output
class LastLayer:
def __init__(self, hidden_size:int, patch_size:int, out_channels:int):
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
self.adaLN_modulation:List[Callable[[Tensor], Tensor]] = [Tensor.silu, nn.Linear(hidden_size, 2 * hidden_size, bias=True)]
def __call__(self, x:Tensor, vec:Tensor) -> Tensor:
shift, scale = vec.sequential(self.adaLN_modulation).chunk(2, dim=1)
x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
return self.linear(x)
def timestep_embedding(t:Tensor, dim:int, max_period:int=10000, time_factor:float=1000.0) -> Tensor:
"""
Create sinusoidal timestep embeddings.
:param t: a 1-D Tensor of N indices, one per batch element.
These may be fractional.
:param dim: the dimension of the output.
:param max_period: controls the minimum frequency of the embeddings.
:return: an (N, D) Tensor of positional embeddings.
"""
t = time_factor * t
half = dim // 2
freqs = Tensor.exp(-math.log(max_period) * Tensor.arange(0, stop=half, dtype=dtypes.float32) / half).to(t.device)
args = t[:, None].float() * freqs[None]
embedding = Tensor.cat(Tensor.cos(args), Tensor.sin(args), dim=-1)
if dim % 2: embedding = Tensor.cat(*[embedding, Tensor.zeros_like(embedding[:, :1])], dim=-1)
if Tensor.is_floating_point(t): embedding = embedding.cast(t.dtype)
return embedding
# https://github.com/black-forest-labs/flux/blob/main/src/flux/model.py
class Flux:
"""
Transformer model for flow matching on sequences.
"""
def __init__(
self,
guidance_embed:bool,
in_channels:int = 64,
vec_in_dim:int = 768,
context_in_dim:int = 4096,
hidden_size:int = 3072,
mlp_ratio:float = 4.0,
num_heads:int = 24,
depth:int = 19,
depth_single_blocks:int = 38,
axes_dim:Optional[List[int]] = None,
theta:int = 10_000,
qkv_bias:bool = True,
):
axes_dim = axes_dim or [16, 56, 56]
self.guidance_embed = guidance_embed
self.in_channels = in_channels
self.out_channels = self.in_channels
if hidden_size % num_heads != 0:
raise ValueError(f"Hidden size {hidden_size} must be divisible by num_heads {num_heads}")
pe_dim = hidden_size // num_heads
if sum(axes_dim) != pe_dim:
raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}")
self.hidden_size = hidden_size
self.num_heads = num_heads
self.pe_embedder = EmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim)
self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
self.vector_in = MLPEmbedder(vec_in_dim, self.hidden_size)
self.guidance_in:Callable[[Tensor], Tensor] = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if guidance_embed else tensor_identity
self.txt_in = nn.Linear(context_in_dim, self.hidden_size)
self.double_blocks = [DoubleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias) for _ in range(depth)]
self.single_blocks = [SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio) for _ in range(depth_single_blocks)]
self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
def __call__(self, img:Tensor, img_ids:Tensor, txt:Tensor, txt_ids:Tensor, timesteps:Tensor, y:Tensor, guidance:Optional[Tensor] = None) -> Tensor:
if img.ndim != 3 or txt.ndim != 3:
raise ValueError("Input img and txt tensors must have 3 dimensions.")
# running on sequences img
img = self.img_in(img)
vec = self.time_in(timestep_embedding(timesteps, 256))
if self.guidance_embed:
if guidance is None:
raise ValueError("Didn't get guidance strength for guidance distilled model.")
vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
vec = vec + self.vector_in(y)
txt = self.txt_in(txt)
ids = Tensor.cat(txt_ids, img_ids, dim=1)
pe = self.pe_embedder(ids)
for double_block in self.double_blocks:
img, txt = double_block(img=img, txt=txt, vec=vec, pe=pe)
img = Tensor.cat(txt, img, dim=1)
for single_block in self.single_blocks:
img = single_block(img, vec=vec, pe=pe)
img = img[:, txt.shape[1] :, ...]
return self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
# https://github.com/black-forest-labs/flux/blob/main/src/flux/util.py
def load_flow_model(name:str, model_path:str):
# Loading Flux
print("Init model")
model = Flux(guidance_embed=(name != "flux-schnell"))
if not model_path: model_path = fetch(urls[name])
state_dict = {k.replace("scale", "weight"): v for k, v in safe_load(model_path).items()}
load_state_dict(model, state_dict)
return model
def load_T5(max_length:int=512):
# max length 64, 128, 256 and 512 should work (if your sequence is short enough)
print("Init T5")
T5 = T5Embedder(max_length, fetch(urls["T5_tokenizer"]))
pt_1 = fetch(urls["T5_1_of_2"])
pt_2 = fetch(urls["T5_2_of_2"])
load_state_dict(T5.encoder, safe_load(pt_1) | safe_load(pt_2), strict=False)
return T5
def load_clip():
print("Init Clip")
clip = ClipEmbedder()
load_state_dict(clip.transformer, safe_load(fetch(urls["clip"])))
return clip
def load_ae() -> AutoEncoder:
# Loading the autoencoder
print("Init AE")
ae = AutoEncoder(0.3611, 0.1159)
load_state_dict(ae, safe_load(fetch(urls["ae"])))
return ae
# https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py
def prepare(T5:T5Embedder, clip:ClipEmbedder, img:Tensor, prompt:Union[str, List[str]]) -> Dict[str, Tensor]:
bs, _, h, w = img.shape
if bs == 1 and not isinstance(prompt, str):
bs = len(prompt)
img = img.rearrange("b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
if img.shape[0] == 1 and bs > 1:
img = img.expand((bs, *img.shape[1:]))
img_ids = Tensor.zeros(h // 2, w // 2, 3).contiguous()
img_ids[..., 1] = img_ids[..., 1] + Tensor.arange(h // 2)[:, None]
img_ids[..., 2] = img_ids[..., 2] + Tensor.arange(w // 2)[None, :]
img_ids = img_ids.rearrange("h w c -> 1 (h w) c")
img_ids = img_ids.expand((bs, *img_ids.shape[1:]))
if isinstance(prompt, str):
prompt = [prompt]
txt = T5(prompt).realize()
if txt.shape[0] == 1 and bs > 1:
txt = txt.expand((bs, *txt.shape[1:]))
txt_ids = Tensor.zeros(bs, txt.shape[1], 3)
vec = clip(prompt).realize()
if vec.shape[0] == 1 and bs > 1:
vec = vec.expand((bs, *vec.shape[1:]))
return {"img": img, "img_ids": img_ids.to(img.device), "txt": txt.to(img.device), "txt_ids": txt_ids.to(img.device), "vec": vec.to(img.device)}
def get_schedule(num_steps:int, image_seq_len:int, base_shift:float=0.5, max_shift:float=1.15, shift:bool=True) -> List[float]:
# extra step for zero
step_size = -1.0 / num_steps
timesteps = Tensor.arange(1, 0 + step_size, step_size)
# shifting the schedule to favor high timesteps for higher signal images
if shift:
# estimate mu based on linear estimation between two points
mu = 0.5 + (max_shift - base_shift) * (image_seq_len - 256) / (4096 - 256)
timesteps = math.exp(mu) / (math.exp(mu) + (1 / timesteps - 1))
return timesteps.tolist()
@TinyJit
def run(model, *args): return model(*args).realize()
def denoise(model, img:Tensor, img_ids:Tensor, txt:Tensor, txt_ids:Tensor, vec:Tensor, timesteps:List[float], guidance:float=4.0) -> Tensor:
# this is ignored for schnell
guidance_vec = Tensor((guidance,), device=img.device, dtype=img.dtype).expand((img.shape[0],))
for t_curr, t_prev in tqdm(list(zip(timesteps[:-1], timesteps[1:])), "Denoising"):
t_vec = Tensor((t_curr,), device=img.device, dtype=img.dtype).expand((img.shape[0],))
pred = run(model, img, img_ids, txt, txt_ids, t_vec, vec, guidance_vec)
img = img + (t_prev - t_curr) * pred
return img
def unpack(x:Tensor, height:int, width:int) -> Tensor:
return x.rearrange("b (h w) (c ph pw) -> b c (h ph) (w pw)", h=math.ceil(height / 16), w=math.ceil(width / 16), ph=2, pw=2)
# https://github.com/black-forest-labs/flux/blob/main/src/flux/cli.py
if __name__ == "__main__":
default_prompt = "bananas and a can of coke"
parser = argparse.ArgumentParser(description="Run Flux.1", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--name", type=str, default="flux-schnell", help="Name of the model to load")
parser.add_argument("--model_path", type=str, default="", help="path of the model file")
parser.add_argument("--width", type=int, default=512, help="width of the sample in pixels (should be a multiple of 16)")
parser.add_argument("--height", type=int, default=512, help="height of the sample in pixels (should be a multiple of 16)")
parser.add_argument("--seed", type=int, default=None, help="Set a seed for sampling")
parser.add_argument("--prompt", type=str, default=default_prompt, help="Prompt used for sampling")
parser.add_argument('--out', type=str, default=Path(tempfile.gettempdir()) / "rendered.png", help="Output filename")
parser.add_argument("--num_steps", type=int, default=None, help="number of sampling steps (default 4 for schnell, 50 for guidance distilled)") #noqa:E501
parser.add_argument("--guidance", type=float, default=3.5, help="guidance value used for guidance distillation")
parser.add_argument("--output_dir", type=str, default="output", help="output directory")
args = parser.parse_args()
if args.name not in ["flux-schnell", "flux-dev"]:
raise ValueError(f"Got unknown model name: {args.name}, chose from flux-schnell and flux-dev")
if args.num_steps is None:
args.num_steps = 4 if args.name == "flux-schnell" else 50
# allow for packing and conversion to latent space
height = 16 * (args.height // 16)
width = 16 * (args.width // 16)
if args.seed is None: args.seed = Tensor._seed
else: Tensor.manual_seed(args.seed)
print(f"Generating with seed {args.seed}:\n{args.prompt}")
t0 = time.perf_counter()
# prepare input noise
x = Tensor.randn(1, 16, 2 * math.ceil(height / 16), 2 * math.ceil(width / 16), dtype="bfloat16")
# load text embedders
T5 = load_T5(max_length=256 if args.name == "flux-schnell" else 512)
clip = load_clip()
# embed text to get inputs for model
inp = prepare(T5, clip, x, prompt=args.prompt)
timesteps = get_schedule(args.num_steps, inp["img"].shape[1], shift=(args.name != "flux-schnell"))
# done with text embedders
del T5, clip
# load model
model = load_flow_model(args.name, args.model_path)
# denoise initial noise
x = denoise(model, **inp, timesteps=timesteps, guidance=args.guidance)
# done with model
del model, run
# load autoencoder
ae = load_ae()
# decode latents to pixel space
x = unpack(x.float(), height, width)
x = ae.decode(x).realize()
t1 = time.perf_counter()
print(f"Done in {t1 - t0:.1f}s. Saving {args.out}")
# bring into PIL format and save
x = x.clamp(-1, 1)
x = x[0].rearrange("c h w -> h w c")
x = (127.5 * (x + 1.0)).cast("uint8")
img = Image.fromarray(x.numpy())
img.save(args.out)
# validation!
if args.prompt == default_prompt and args.name=="flux-schnell" and args.seed == 0 and args.width == args.height == 512:
ref_image = Tensor(np.array(Image.open("examples/flux1_seed0.png")))
distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item()
assert distance < 4e-3, colored(f"validation failed with {distance=}", "red")
print(colored(f"output validated with {distance=}", "green"))
Binary file not shown.

After

Width:  |  Height:  |  Size: 286 KiB

-108
View File
@@ -1,108 +0,0 @@
import itertools
from typing import Callable
from tinygrad import nn, Tensor, dtypes, Device, TinyJit
from tinygrad.helpers import getenv, trange, partition
class Model:
def __init__(self):
self.layers: list[Callable[[Tensor], Tensor]] = [
nn.Conv2d(1, 32, 5), Tensor.relu,
nn.Conv2d(32, 32, 5), Tensor.relu,
nn.BatchNorm(32), Tensor.max_pool2d,
nn.Conv2d(32, 64, 3), Tensor.relu,
nn.Conv2d(64, 64, 3), Tensor.relu,
nn.BatchNorm(64), Tensor.max_pool2d,
lambda x: x.flatten(1), nn.Linear(576, 10)]
def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
# TODO: refactor this into optim/onnx
def functional_adam(g:Tensor, m:Tensor, v:Tensor, b1_t:Tensor, b2_t:Tensor, lr=0.001, b1=0.9, b2=0.999, eps=1e-6) -> Tensor:
b1_t *= b1
b2_t *= b2
m.assign(b1 * m + (1.0 - b1) * g)
v.assign(b2 * v + (1.0 - b2) * (g * g))
m_hat = m / (1.0 - b1_t)
v_hat = v / (1.0 - b2_t)
return lr * (m_hat / (v_hat.sqrt() + eps))
if __name__ == "__main__":
BS = getenv("BS", 512)
ACC_STEPS = getenv("ACC_STEPS", 8)
X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
model = Model()
params = nn.state.get_parameters(model)
# init params, set requires grad on the ones we need gradients of
for x in params:
if x.requires_grad is None: x.requires_grad_()
x.replace(x.contiguous())
Tensor.realize(*params)
# split params (with grads) and buffers (without)
params, buffers = partition(params, lambda x: x.requires_grad)
print(f"params: {len(params)} buffers: {len(buffers)}")
# optim params
pos_params = list(itertools.accumulate(params, lambda x,y: x+y.numel(), initial=0))
adam_m = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
adam_v = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
adam_params = [adam_m, adam_v, adam_b1_t, adam_b2_t]
# create loss and grads. init all state so the JIT works on microbatch
for x in params: x.assign(x.detach())
loss = Tensor.zeros(tuple()).contiguous()
grads = Tensor.zeros(pos_params[-1]).contiguous()
Tensor.realize(*params, *buffers, *adam_params, loss, grads)
@TinyJit
@Tensor.train()
def microbatch():
samples = Tensor.randint(BS // ACC_STEPS, high=X_train.shape[0])
for t in params: t.grad = None
# divide by ACC_STEPS at the loss
uloss = (model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]) / ACC_STEPS).backward()
ugrads = Tensor.cat(*[t.grad.contiguous().flatten() for t in params], dim=0)
for t in params: t.grad = None
# concat the grads and assign them
loss.assign(loss + uloss)
grads.assign(grads + ugrads)
Tensor.realize(*params, *buffers, loss, grads)
@TinyJit
def optimizer():
# run optimizer (on CPU, where adam params live)
delta = functional_adam(grads.to("CPU"), adam_m, adam_v, adam_b1_t, adam_b2_t)
# update the params, copying back the delta one at a time to avoid OOM
# NOTE: the scheduler is ordering things poorly, all the copies are happening before the adds
for j,tt in enumerate(params):
tt.assign(tt.detach() - delta[pos_params[j]:pos_params[j+1]].reshape(tt.shape).to(Device.DEFAULT))
# realize everything, zero out loss and grads
loss.assign(Tensor.zeros_like(loss))
grads.assign(Tensor.zeros_like(grads))
Tensor.realize(*params, *adam_params, loss, grads)
@TinyJit
def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
test_acc = float('nan')
for i in (t:=trange(getenv("STEPS", 70))):
# microbatch sets the gradients
for _ in range(ACC_STEPS): microbatch()
# get the loss before the optimizer clears it
# this is already realized so this isn't a schedule
loss_item = loss.item()
# run the optimizer
optimizer()
# eval
if i%10 == 9: test_acc = get_test_acc().item()
t.set_description(f"loss: {loss_item:6.2f} test_accuracy: {test_acc:5.2f}%")
+299
View File
@@ -0,0 +1,299 @@
from extra.models.mask_rcnn import MaskRCNN
from extra.models.resnet import ResNet
from extra.models.mask_rcnn import BoxList
from torch.nn import functional as F
from torchvision import transforms as T
from torchvision.transforms import functional as Ft
import random
from tinygrad.tensor import Tensor
from PIL import Image
import numpy as np
import torch
import argparse
import cv2
class Resize:
def __init__(self, min_size, max_size):
if not isinstance(min_size, (list, tuple)):
min_size = (min_size,)
self.min_size = min_size
self.max_size = max_size
# modified from torchvision to add support for max size
def get_size(self, image_size):
w, h = image_size
size = random.choice(self.min_size)
max_size = self.max_size
if max_size is not None:
min_original_size = float(min((w, h)))
max_original_size = float(max((w, h)))
if max_original_size / min_original_size * size > max_size:
size = int(round(max_size * min_original_size / max_original_size))
if (w <= h and w == size) or (h <= w and h == size):
return (h, w)
if w < h:
ow = size
oh = int(size * h / w)
else:
oh = size
ow = int(size * w / h)
return (oh, ow)
def __call__(self, image):
size = self.get_size(image.size)
image = Ft.resize(image, size)
return image
class Normalize:
def __init__(self, mean, std, to_bgr255=True):
self.mean = mean
self.std = std
self.to_bgr255 = to_bgr255
def __call__(self, image):
if self.to_bgr255:
image = image[[2, 1, 0]] * 255
else:
image = image[[0, 1, 2]] * 255
image = Ft.normalize(image, mean=self.mean, std=self.std)
return image
transforms = lambda size_scale: T.Compose(
[
Resize(int(800*size_scale), int(1333*size_scale)),
T.ToTensor(),
Normalize(
mean=[102.9801, 115.9465, 122.7717], std=[1., 1., 1.], to_bgr255=True
),
]
)
def expand_boxes(boxes, scale):
w_half = (boxes[:, 2] - boxes[:, 0]) * .5
h_half = (boxes[:, 3] - boxes[:, 1]) * .5
x_c = (boxes[:, 2] + boxes[:, 0]) * .5
y_c = (boxes[:, 3] + boxes[:, 1]) * .5
w_half *= scale
h_half *= scale
boxes_exp = torch.zeros_like(boxes)
boxes_exp[:, 0] = x_c - w_half
boxes_exp[:, 2] = x_c + w_half
boxes_exp[:, 1] = y_c - h_half
boxes_exp[:, 3] = y_c + h_half
return boxes_exp
def expand_masks(mask, padding):
N = mask.shape[0]
M = mask.shape[-1]
pad2 = 2 * padding
scale = float(M + pad2) / M
padded_mask = mask.new_zeros((N, 1, M + pad2, M + pad2))
padded_mask[:, :, padding:-padding, padding:-padding] = mask
return padded_mask, scale
def paste_mask_in_image(mask, box, im_h, im_w, thresh=0.5, padding=1):
# TODO: remove torch
mask = torch.tensor(mask.numpy())
box = torch.tensor(box.numpy())
padded_mask, scale = expand_masks(mask[None], padding=padding)
mask = padded_mask[0, 0]
box = expand_boxes(box[None], scale)[0]
box = box.to(dtype=torch.int32)
TO_REMOVE = 1
w = int(box[2] - box[0] + TO_REMOVE)
h = int(box[3] - box[1] + TO_REMOVE)
w = max(w, 1)
h = max(h, 1)
mask = mask.expand((1, 1, -1, -1))
mask = mask.to(torch.float32)
mask = F.interpolate(mask, size=(h, w), mode='bilinear', align_corners=False)
mask = mask[0][0]
if thresh >= 0:
mask = mask > thresh
else:
mask = (mask * 255).to(torch.uint8)
im_mask = torch.zeros((im_h, im_w), dtype=torch.uint8)
x_0 = max(box[0], 0)
x_1 = min(box[2] + 1, im_w)
y_0 = max(box[1], 0)
y_1 = min(box[3] + 1, im_h)
im_mask[y_0:y_1, x_0:x_1] = mask[
(y_0 - box[1]): (y_1 - box[1]), (x_0 - box[0]): (x_1 - box[0])
]
return im_mask
class Masker:
def __init__(self, threshold=0.5, padding=1):
self.threshold = threshold
self.padding = padding
def forward_single_image(self, masks, boxes):
boxes = boxes.convert("xyxy")
im_w, im_h = boxes.size
res = [
paste_mask_in_image(mask[0], box, im_h, im_w, self.threshold, self.padding)
for mask, box in zip(masks, boxes.bbox)
]
if len(res) > 0:
res = torch.stack(*res, dim=0)[:, None]
else:
res = masks.new_empty((0, 1, masks.shape[-2], masks.shape[-1]))
return Tensor(res.numpy())
def __call__(self, masks, boxes):
if isinstance(boxes, BoxList):
boxes = [boxes]
results = []
for mask, box in zip(masks, boxes):
result = self.forward_single_image(mask, box)
results.append(result)
return results
masker = Masker(threshold=0.5, padding=1)
def select_top_predictions(predictions, confidence_threshold=0.9):
scores = predictions.get_field("scores").numpy()
keep = [idx for idx, score in enumerate(scores) if score > confidence_threshold]
return predictions[keep]
def compute_prediction(original_image, model, confidence_threshold, size_scale=1.0):
image = transforms(size_scale)(original_image).numpy()
image = Tensor(image, requires_grad=False)
predictions = model(image)
prediction = predictions[0]
prediction = select_top_predictions(prediction, confidence_threshold)
width, height = original_image.size
prediction = prediction.resize((width, height))
if prediction.has_field("mask"):
masks = prediction.get_field("mask")
masks = masker([masks], [prediction])[0]
prediction.add_field("mask", masks)
return prediction
def compute_prediction_batched(batch, model, size_scale=1.0):
imgs = []
for img in batch:
imgs.append(transforms(size_scale)(img).numpy())
image = [Tensor(image, requires_grad=False) for image in imgs]
predictions = model(image)
del image
return predictions
palette = np.array([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
def findContours(*args, **kwargs):
if cv2.__version__.startswith('4'):
contours, hierarchy = cv2.findContours(*args, **kwargs)
elif cv2.__version__.startswith('3'):
_, contours, hierarchy = cv2.findContours(*args, **kwargs)
return contours, hierarchy
def compute_colors_for_labels(labels):
l = labels[:, None]
colors = l * palette
colors = (colors % 255).astype("uint8")
return colors
def overlay_mask(image, predictions):
image = np.asarray(image)
masks = predictions.get_field("mask").numpy()
labels = predictions.get_field("labels").numpy()
colors = compute_colors_for_labels(labels).tolist()
for mask, color in zip(masks, colors):
thresh = mask[0, :, :, None]
contours, hierarchy = findContours(
thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
)
image = cv2.drawContours(image, contours, -1, color, 3)
composite = image
return composite
CATEGORIES = [
"__background", "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant",
"bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard",
"sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle",
"wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
"carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table",
"toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster",
"sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush",
]
def overlay_boxes(image, predictions):
labels = predictions.get_field("labels").numpy()
boxes = predictions.bbox
image = np.asarray(image)
colors = compute_colors_for_labels(labels).tolist()
for box, color in zip(boxes, colors):
box = torch.tensor(box.numpy())
box = box.to(torch.int64)
top_left, bottom_right = box[:2].tolist(), box[2:].tolist()
image = cv2.rectangle(
image, tuple(top_left), tuple(bottom_right), tuple(color), 1
)
return image
def overlay_class_names(image, predictions):
scores = predictions.get_field("scores").numpy().tolist()
labels = predictions.get_field("labels").numpy().tolist()
labels = [CATEGORIES[int(i)] for i in labels]
boxes = predictions.bbox.numpy()
image = np.asarray(image)
template = "{}: {:.2f}"
for box, score, label in zip(boxes, scores, labels):
x, y = box[:2]
s = template.format(label, score)
x, y = int(x), int(y)
cv2.putText(
image, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, .5, (255, 255, 255), 1
)
return image
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Run MaskRCNN', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--image', type=str, help="Path of the image to run")
parser.add_argument('--threshold', type=float, default=0.7, help="Detector threshold")
parser.add_argument('--size_scale', type=float, default=1.0, help="Image resize multiplier")
parser.add_argument('--out', type=str, default="/tmp/rendered.png", help="Output filename")
args = parser.parse_args()
resnet = ResNet(50, num_classes=None, stride_in_1x1=True)
model_tiny = MaskRCNN(resnet)
model_tiny.load_from_pretrained()
img = Image.open(args.image)
top_result_tiny = compute_prediction(img, model_tiny, confidence_threshold=args.threshold, size_scale=args.size_scale)
bbox_image = overlay_boxes(img, top_result_tiny)
mask_image = overlay_mask(bbox_image, top_result_tiny)
final_image = overlay_class_names(mask_image, top_result_tiny)
im = Image.fromarray(final_image)
print(f"saving {args.out}")
im.save(args.out)
im.show()
+40 -19
View File
@@ -213,13 +213,12 @@ 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, seed:int|None=None):
def batch_load_train_bert(BS:int):
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
rng.shuffle(fs)
random.shuffle(fs)
train_files.append(fs.pop(0))
cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
@@ -764,26 +763,48 @@ class BlendedGPTDataset:
return dataset_idx, dataset_sample_idx
def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False) -> BlendedGPTDataset:
if small:
if val:
return BlendedGPTDataset(
[base_dir / "c4-validation-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
return BlendedGPTDataset(
[base_dir / "c4-train.en_6_text_document"], [1.0], samples, seqlen, seed, shuffle=True)
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
if val:
return BlendedGPTDataset(
[base_dir / "validation" / "c4-validationn-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
return BlendedGPTDataset(
[base_dir / "c4-train.en_6_text_document", base_dir / "c4-train.en_7_text_document"], [1.0, 1.0], samples, seqlen, seed, shuffle=True)
dataset = BlendedGPTDataset([
base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
], [
1.0
], samples, seqlen, seed, False)
else:
dataset = BlendedGPTDataset([
base_dir / "c4-train.en_6_text_document",
base_dir / "c4-train.en_7_text_document",
], [
1.0, 1.0
], samples, seqlen, seed, True)
def iterate_llama3_dataset(dataset:BlendedGPTDataset, bs:int):
for b in range(math.ceil(dataset.samples / bs)):
batch = [dataset.get(b * bs + i) for i in range(bs)]
for b in range(math.ceil(samples / bs)):
batch = []
for i in range(bs):
tokens = dataset.get(b * bs + i)
batch.append(tokens)
yield Tensor.stack(batch, dim=0)
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)
def batch_load_llama3_small(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
if val:
dataset = BlendedGPTDataset([
base_dir / "c4-validation-91205-samples.en_text_document",
], [
1.0
], samples, seqlen, seed, False)
else:
dataset = BlendedGPTDataset([
base_dir / "c4-train.en_6_text_document",
], [
1.0
], samples, seqlen, seed, True)
for b in range(math.ceil(samples / bs)):
batch = []
for i in range(bs):
tokens = dataset.get(b * bs + i)
batch.append(tokens)
yield Tensor.stack(batch, dim=0)
if __name__ == "__main__":
def load_unet3d(val):
+8 -19
View File
@@ -219,28 +219,17 @@ 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
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
return BertForPretraining(**config)
def get_fake_data_bert(BS:int):
return {
"input_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
"input_mask": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
"segment_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
"masked_lm_positions": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
"masked_lm_ids": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
"masked_lm_weights": Tensor.zeros((BS, 76), dtype=dtypes.float32, device="CPU").contiguous(),
"next_sentence_labels": Tensor.zeros((BS, 1), dtype=dtypes.int32, device="CPU").contiguous(),
"input_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
"input_mask": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
"segment_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
"masked_lm_positions": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
"masked_lm_ids": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
"masked_lm_weights": Tensor.empty((BS, 76), dtype=dtypes.float32, device="CPU"),
"next_sentence_labels": Tensor.empty((BS, 1), dtype=dtypes.int32, device="CPU"),
}
def find_matches(match_quality_matrix:np.ndarray, high_threshold:float=0.5, low_threshold:float=0.4, allow_low_quality_matches:bool=False) -> np.ndarray:
+3 -1
View File
@@ -59,7 +59,9 @@ class EmbeddingBert(nn.Embedding):
arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).reshape(arange_shp)
arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype)
# TODO: contiguous() here because the embedding dropout creates different asts on each device, and search becomes very slow.
# Should fix with fixing random ast on multi device, and fuse arange to make embedding fast.
return (arange == idx).mul(vals).sum(2, dtype=vals.dtype).contiguous()
class LayerNormBert:
def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
+93
View File
@@ -0,0 +1,93 @@
import math
from pathlib import Path
from tinygrad import Device, nn, Tensor, TinyJit
from tinygrad.helpers import getenv, profile_marker
from extra.models.llama import Transformer
from examples.llama3 import MODEL_PARAMS
from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
config = {}
BASEDIR = config["BASEDIR"] = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
BS = config["BS"] = getenv("BS", 16)
grad_acc = config["GRADIENT_ACC_STEPS"] = getenv("GRADIENT_ACC_STEPS", 1)
GBS = config["GLOBAL_BATCH_SIZE"] = BS * grad_acc
SEED = config["SEED"] = getenv("SEED", 5760)
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
TRAIN_ON_VAL = config["TRAIN_ON_VAL"] = getenv("TRAIN_ON_VAL", 0)
SMALL = config["SMALL"] = getenv("SMALL", 0)
SAMPLES = config["SAMPLES"] = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
EVAL_FREQ = config["EVAL_FREQ"] = getenv("EVAL_FREQ", 46080)
EVAL_BS = config["EVAL_BS"] = getenv("EVAL_BS", 16)
EVAL_TARGET = config["EVAL_TARGET"] = getenv("EVAL_TARGET", 5.6)
opt_adamw_beta_1 = 0.9
opt_adamw_beta_2 = 0.95
opt_adamw_epsilon = 1e-5
opt_adamw_weight_decay = 0.1
opt_gradient_clip_norm = 1.0
opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
opt_learning_rate_decay_steps = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS)) - opt_learning_rate_warmup_steps
opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152) # NOTE: cannot change for benchmark
opt_end_learning_rate = getenv("END_LR", 8e-7)
# TODO: confirm weights are in bf16
# vocab_size from the mixtral tokenizer
params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
params = params | {"vocab_size": 32000} if not SMALL else params
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
if __name__ == "__main__":
profile_marker("create model")
model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
# shard the model, either data parallel (DP) or model parallel (MP)
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
for v in nn.state.get_parameters(model):
v.shard_(device, axis=None)
if (MP := getenv("MP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
for k,v in nn.state.get_state_dict(model).items():
if 'scale' in k: v.shard_(device, axis=None) # from quantized
elif '.attention.wq' in k: v.shard_(device, axis=0)
elif '.attention.wk' in k: v.shard_(device, axis=0)
elif '.attention.wv' in k: v.shard_(device, axis=0)
elif '.attention.wo' in k: v.shard_(device, axis=1)
elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
elif 'output.weight' in k: v.shard_(device, axis=0)
else:
# attention_norm, ffn_norm, norm
v.shard_(device, axis=None)
# prevents memory spike on device 0
v.realize()
profile_marker("create optim")
optim = nn.optim.AdamW(nn.state.get_parameters(model), lr=0.0,
b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, fused=True)
scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate,
opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
profile_marker("init params")
optim.lr.realize(*[p.replace(p.contiguous()) for p in optim.params])
# TODO: make this work with multigpu
cat_params = Tensor.cat(*[t.flatten() for t in optim.params], dim=0)
cat_grads = Tensor.zeros_like(cat_params)
@profile_marker("microbatch")
@TinyJit
@Tensor.train()
def microbatch(batch:Tensor):
logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
loss = logits.sparse_categorical_crossentropy(batch[:, 1:]).backward()
return loss.realize(cat_grads)
+44 -4
View File
@@ -204,6 +204,43 @@ def eval_bert():
st = time.perf_counter()
def eval_mrcnn():
from tqdm import tqdm
from extra.models.mask_rcnn import MaskRCNN
from extra.models.resnet import ResNet
from extra.datasets.coco import BASEDIR, images, convert_prediction_to_coco_bbox, convert_prediction_to_coco_mask, accumulate_predictions_for_coco, evaluate_predictions_on_coco, iterate
from examples.mask_rcnn import compute_prediction_batched, Image
mdl = MaskRCNN(ResNet(50, num_classes=None, stride_in_1x1=True))
mdl.load_from_pretrained()
bbox_output = '/tmp/results_bbox.json'
mask_output = '/tmp/results_mask.json'
accumulate_predictions_for_coco([], bbox_output, rm=True)
accumulate_predictions_for_coco([], mask_output, rm=True)
#TODO: bs > 1 not as accurate
bs = 1
for batch in tqdm(iterate(images, bs=bs), total=len(images)//bs):
batch_imgs = []
for image_row in batch:
image_name = image_row['file_name']
img = Image.open(BASEDIR/f'val2017/{image_name}').convert("RGB")
batch_imgs.append(img)
batch_result = compute_prediction_batched(batch_imgs, mdl)
for image_row, result in zip(batch, batch_result):
image_name = image_row['file_name']
box_pred = convert_prediction_to_coco_bbox(image_name, result)
mask_pred = convert_prediction_to_coco_mask(image_name, result)
accumulate_predictions_for_coco(box_pred, bbox_output)
accumulate_predictions_for_coco(mask_pred, mask_output)
del batch_imgs
del batch_result
evaluate_predictions_on_coco(bbox_output, iou_type='bbox')
evaluate_predictions_on_coco(mask_output, iou_type='segm')
def eval_llama3():
from extra.models.llama import Transformer
from examples.llama3 import MODEL_PARAMS, load, convert_from_huggingface
@@ -234,9 +271,12 @@ def eval_llama3():
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
return loss.flatten().float()
from examples.mlperf.dataloader import get_llama3_dataset, iterate_llama3_dataset
eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
iter = iterate_llama3_dataset(eval_dataset, BS)
if SMALL:
from examples.mlperf.dataloader import batch_load_llama3_small
iter = batch_load_llama3_small(BS, 5760, SEQLEN, BASEDIR, val=True)
else:
from examples.mlperf.dataloader import batch_load_llama3
iter = batch_load_llama3(BS, 5760, SEQLEN, BASEDIR, val=True)
losses = []
for tokens in tqdm(iter, total=5760//BS):
@@ -501,7 +541,7 @@ if __name__ == "__main__":
# inference only
Tensor.training = False
models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(",")
for m in models:
nm = f"eval_{m}"
if nm in globals():
+134 -131
View File
@@ -3,7 +3,7 @@ from pathlib import Path
import multiprocessing
from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker
from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
@@ -918,6 +918,40 @@ def train_rnnt():
# TODO: RNN-T
pass
@TinyJit
def train_step_bert(model, optimizer, scheduler, loss_scaler:float, GPUS, grad_acc:int, **kwargs):
optimizer.zero_grad()
for i in range(grad_acc):
input_ids, segment_ids = kwargs[f"input_ids{i}"], kwargs[f"segment_ids{i}"]
# NOTE: these two have different names
attention_mask, masked_positions = kwargs[f"input_mask{i}"], kwargs[f"masked_lm_positions{i}"]
masked_lm_ids, masked_lm_weights, next_sentence_labels = kwargs[f"masked_lm_ids{i}"], kwargs[f"masked_lm_weights{i}"], kwargs[f"next_sentence_labels{i}"]
for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
if len(GPUS) > 1: t.shard_(GPUS, axis=0)
else: t.to_(GPUS[0])
lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
(loss * loss_scaler).backward()
# TODO: OOM without this realize with large grad_acc
Tensor.realize(*[p.grad for p in optimizer.params])
global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer[0].device)
for p in optimizer.params:
p.grad = p.grad / loss_scaler
global_norm += p.grad.float().square().sum()
global_norm = global_norm.sqrt().contiguous()
for p in optimizer.params:
p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)
optimizer.step()
scheduler.step()
# TODO: no to("CPU") here because it blocks and messes the python time
Tensor.realize(loss, global_norm, optimizer.optimizers[0].lr)
return loss, global_norm, optimizer.optimizers[0].lr
@TinyJit
def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
@@ -980,8 +1014,7 @@ def train_bert():
# ** hyperparameters **
BS = config["BS"] = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
grad_acc = config["GRADIENT_ACC_STEPS"] = getenv("GRADIENT_ACC_STEPS", 1)
# TODO: implement grad accumulation + mlperf logging
assert grad_acc == 1
# TODO: mlperf logging
GBS = config["GLOBAL_BATCH_SIZE"] = BS * grad_acc
EVAL_BS = config["EVAL_BS"] = getenv("EVAL_BS", 1 * len(GPUS))
max_lr = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(GBS/96))
@@ -1008,7 +1041,6 @@ 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
@@ -1041,8 +1073,8 @@ def train_bert():
# ** Optimizer **
parameters_no_wd = [v for k, v in get_state_dict(model).items() if "bias" in k or "LayerNorm" in k]
parameters_wd = [x for x in parameters if x not in set(parameters_no_wd)]
optimizer_wd = LAMB(parameters_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
parameters = [x for x in parameters if x not in set(parameters_no_wd)]
optimizer_wd = LAMB(parameters, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
optimizer_no_wd = LAMB(parameters_no_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=0.0, adam=False)
optimizer_group = OptimizerGroup(optimizer_wd, optimizer_no_wd)
@@ -1086,7 +1118,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, seed=seed), total=train_steps, disable=BENCHMARK))
train_it = iter(tqdm(batch_load_train_bert(BS), total=train_steps, disable=BENCHMARK))
for _ in range(start_step): next(train_it) # Fast forward
else:
# repeat fake data
@@ -1099,38 +1131,12 @@ def train_bert():
# ** train loop **
wc_start = time.perf_counter()
i, train_data = start_step, next(train_it)
i, train_data = start_step, [next(train_it) for _ in range(grad_acc)]
if RUNMLPERF:
if MLLOGGER:
MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})
@TinyJit
def train_step_bert(input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
if len(GPUS) > 1: t.shard_(GPUS, axis=0)
else: t.to_(GPUS[0])
optimizer_group.zero_grad()
lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
(loss * loss_scaler).backward()
global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer_group[0].device)
for p in optimizer_group.params:
p.grad = p.grad / loss_scaler
global_norm += p.grad.float().square().sum()
global_norm = global_norm.sqrt().contiguous()
for p in optimizer_group.params:
p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)
optimizer_group.step()
scheduler_group.step()
# TODO: no to("CPU") here because it blocks and messes the python time
Tensor.realize(loss, global_norm, optimizer_group.optimizers[0].lr)
return loss, global_norm, optimizer_group.optimizers[0].lr
while train_data is not None and i < train_steps and not achieved:
if getenv("TRAIN", 1):
Tensor.training = True
@@ -1138,17 +1144,21 @@ def train_bert():
st = time.perf_counter()
GlobalCounters.reset()
with WallTimeEvent(BenchEvent.STEP):
loss, global_norm, lr = train_step_bert(
train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"])
data = {f"{k}{i}":v for i,d in enumerate(train_data) for k,v in d.items()}
loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler, GPUS, grad_acc, **data)
pt = time.perf_counter()
next_data = next(train_it)
try:
next_data = [next(train_it) for _ in range(grad_acc)]
except StopIteration:
next_data = None
dt = time.perf_counter()
device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
loss = loss.item()
if not getenv("FP8_TRAIN"): assert not math.isnan(loss)
assert not math.isnan(loss)
lr = lr.item()
cl = time.perf_counter()
@@ -1161,7 +1171,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/mem":GlobalCounters.mem_used / 1e9, "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
"train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
train_data, next_data = next_data, None
i += 1
@@ -1178,8 +1188,8 @@ def train_bert():
if MLLOGGER and RUNMLPERF:
MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
if getenv("RESET_STEP"): train_step_bert.reset()
elif getenv("FREE_INTERMEDIATE") and train_step_bert.captured is not None:
# TODO: this hangs on tiny green after 90 minutes of training
elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None:
# TODO: FREE_INTERMEDIATE nan'ed after jit step 2
train_step_bert.captured.free_intermediates()
eval_lm_losses = []
eval_clsf_losses = []
@@ -1214,7 +1224,7 @@ def train_bert():
return
if getenv("RESET_STEP"): eval_step_bert.reset()
elif getenv("FREE_INTERMEDIATE") and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
del eval_data
avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
@@ -1290,7 +1300,6 @@ def train_llama3():
BASEDIR = config["BASEDIR"] = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
BS = config["BS"] = getenv("BS", 16)
grad_acc = config["GRADIENT_ACC_STEPS"] = getenv("GRADIENT_ACC_STEPS", 1)
assert grad_acc == 1, f"{grad_acc=} is not supported"
GBS = config["GLOBAL_BATCH_SIZE"] = BS * grad_acc
SEED = config["SEED"] = getenv("SEED", 5760)
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
@@ -1315,25 +1324,12 @@ def train_llama3():
opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152) # NOTE: cannot change for benchmark
opt_end_learning_rate = getenv("END_LR", 8e-7)
# ** init wandb **
WANDB = getenv("WANDB")
if WANDB:
import wandb
wandb_args = {"id": wandb_id, "resume": "must"} if (wandb_id := getenv("WANDB_RESUME", "")) else {}
wandb.init(config=config, **wandb_args, project="MLPerf-LLaMA3")
model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
# TODO: confirm weights are in bf16
# vocab_size from the mixtral tokenizer
if not SMALL: model_params |= {"vocab_size": 32000}
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
params = get_parameters(model)
# weights are all bfloat16 for now
assert params and all(p.dtype == dtypes.bfloat16 for p in params)
if getenv("FAKEDATA"):
for v in get_parameters(model):
v = v.assign(Tensor.empty(v.shape))
params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
params = params | {"vocab_size": 32000} if not SMALL else params
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
@@ -1360,9 +1356,13 @@ def train_llama3():
v.realize()
optim = AdamW(get_parameters(model), lr=0.0,
b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, fused=True)
scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
# init tensors
profile_marker("init tensors")
optim.lr.realize(*[p.replace(p.contiguous()) for p in optim.params])
if resume_ckpt := getenv("RESUME_CKPT"):
fn = f"./ckpts/llama3_{resume_ckpt}.safe"
print(f"loading initial checkpoint from {fn}")
@@ -1374,38 +1374,50 @@ def train_llama3():
@TinyJit
@Tensor.train()
def train_step(model, tokens:Tensor):
def train_step(tokens:Tensor, grad_acc:int):
optim.zero_grad()
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
tokens = tokens.shard(device, 0)
if (MP := getenv("MP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
tokens = tokens.shard(device)
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
loss.backward()
# grad acc. NOTE: this has to become multidevice aware, this cat should be per device
cat_params = Tensor.cat(*[t.flatten() for t in optim.params], dim=0)
cat_grads = Tensor.zeros_like(cat_params)
total_loss = Tensor(0, dtype=dtypes.float)
for batch in tokens.split(tokens.shape[0]//grad_acc):
profile_marker("grads")
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
batch = batch.shard(device, 0)
if (MP := getenv("MP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
batch = batch.shard(device)
logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
loss = logits.sparse_categorical_crossentropy(batch[:, 1:])
loss.backward()
total_loss += loss/grad_acc
cat_grads += Tensor.cat(*[t.grad.flatten() for t in optim.params], dim=0)
total_loss.realize(cat_grads)
# L2 norm grad clip
# https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
# https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
profile_marker("optimizer")
if not getenv("DISABLE_GRAD_CLIP_NORM"):
total_norm = Tensor(0.0, dtype=dtypes.float32, device=optim.params[0].device)
for p in optim.params:
total_norm += p.grad.float().square().sum()
total_norm = total_norm.sqrt().contiguous()
for p in optim.params:
p.grad = p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
total_norm = cat_grads.float().square().sum().sqrt().contiguous()
cat_grads = cat_grads * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
optim.step()
scheduler.step()
# run the optimizer
# NOTE: this is copied from _schedule_step
out, extra = optim._step([cat_params], [cat_grads]) # this will go on CPU
lr = optim.lr
loss.realize(lr)
return loss, lr
# update the parameters
updated_params = [out[0][optim.pos_params[i]:optim.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(optim.params)]
for i, tt in enumerate(optim.params): tt.assign(updated_params[i])
Tensor.realize(*optim.params, *extra, *optim.buffers, *scheduler.schedule_step())
return total_loss
@TinyJit
@Tensor.train(False)
def eval_step(model, tokens:Tensor):
def eval_step(tokens:Tensor):
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
tokens = tokens.shard(device, 0)
@@ -1419,62 +1431,56 @@ def train_llama3():
# ** data iters **
def fake_data(bs, samples):
for _ in range(samples // bs):
yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
yield Tensor.randint(bs, SEQLEN + 1, low=0, high=params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
def get_train_iter():
if getenv("FAKEDATA", 0):
return fake_data(BS, SAMPLES)
return fake_data(GBS, SAMPLES)
else:
from examples.mlperf.dataloader import batch_load_llama3
return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
if getenv("FAKEDATA", 0):
eval_dataset = None
else:
from examples.mlperf.dataloader import get_llama3_dataset
eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
if SMALL:
from examples.mlperf.dataloader import batch_load_llama3_small
return batch_load_llama3_small(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
else:
from examples.mlperf.dataloader import batch_load_llama3
return batch_load_llama3(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
def get_eval_iter():
if eval_dataset is None:
if getenv("FAKEDATA", 0):
return fake_data(EVAL_BS, 5760)
from examples.mlperf.dataloader import iterate_llama3_dataset
return iterate_llama3_dataset(eval_dataset, EVAL_BS)
else:
if SMALL:
from examples.mlperf.dataloader import batch_load_llama3_small
return batch_load_llama3_small(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
else:
from examples.mlperf.dataloader import batch_load_llama3
return batch_load_llama3(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
iter = get_train_iter()
i, sequences_seen = resume_ckpt, 0
for tokens in tqdm(iter, total=SAMPLES//GBS):
profile_marker(f"train step {i}")
t = time.perf_counter()
GlobalCounters.reset()
if getenv("TRAIN", 1):
t = time.perf_counter()
loss, lr = train_step(model, tokens)
loss = loss.float().item()
lr = lr.item()
loss = train_step(tokens, grad_acc)
loss, lr = loss.float().item(), optim.lr.item()
i += 1
sequences_seen += tokens.shape[0]
i += 1
sequences_seen += tokens.shape[0]
sec = time.perf_counter()-t
mem_gb = GlobalCounters.mem_used / 1e9
gflops = GlobalCounters.global_ops / 1e9 / sec
tqdm.write(
f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS")
tqdm.write(f"{loss:.4f} loss, {lr:.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
if (fname:=getenv("LOSS_FILE", "")):
with open(fname, "a") as f:
f.write(f"{i} {loss:.4f} {lr:.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
if (fname:=getenv("LOSS_FILE", "")):
with open(fname, "a") as f:
f.write(f"{i} {loss:.4f} {lr:.12f} {mem_gb:.2f}\n")
if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
tqdm.write("saving checkpoint")
if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
fn = f"{ckpt_dir}/llama3_{i}.safe"
safe_save(get_state_dict(model), fn)
if WANDB:
wandb.log({"lr": lr, "train/loss": loss, "train/step_time": sec, "train/GFLOPS": gflops, "train/sequences_seen": sequences_seen})
if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
tqdm.write("saving checkpoint")
if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
fn = f"{ckpt_dir}/llama3_{i}.safe"
safe_save(get_state_dict(model), fn)
tqdm.write("saving optim checkpoint")
fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
safe_save(get_state_dict(scheduler), fn)
tqdm.write("saving optim checkpoint")
fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
safe_save(get_state_dict(scheduler), fn)
if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
tqdm.write(f"evaluating after {sequences_seen} sequences")
@@ -1485,14 +1491,11 @@ def train_llama3():
tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
for tokens in tqdm(eval_iter, total=5760//EVAL_BS):
eval_losses += eval_step(model, tokens).tolist()
eval_losses += eval_step(tokens).tolist()
log_perplexity = Tensor(eval_losses).mean().float().item()
tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
if WANDB:
wandb.log({"eval/log_perplexity": log_perplexity, "eval/sequences_seen": sequences_seen})
if log_perplexity < EVAL_TARGET:
tqdm.write(f"target achieved after {sequences_seen} sequences")
if getenv("CKPT"):
@@ -1571,7 +1574,7 @@ def train_stable_diffusion():
loss, out_lr = loss.detach().to("CPU"), optimizer.lr.to("CPU")
Tensor.realize(loss, out_lr)
return loss, out_lr
# checkpointing takes ~9 minutes without this, and ~1 minute with this
@TinyJit
def ckpt_to_cpu():
@@ -1610,7 +1613,7 @@ def train_stable_diffusion():
if i == 3:
for _ in range(3): ckpt_to_cpu() # do this at the beginning of run to prevent OOM surprises when checkpointing
print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
total_train_time = time.perf_counter() - train_start_time
if WANDB:
wandb.log({"train/loss": loss_item, "train/lr": lr_item, "train/loop_time_prev": loop_time, "train/dl_time": dl_time, "train/step": i,
@@ -1,24 +0,0 @@
#!/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
@@ -1,31 +0,0 @@
#!/bin/bash
set -e # Exit on any error
set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." AMD=1
export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_8xMI350X"
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"
# pip install -e ".[mlperf]"
export LOGMLPERF=1
export SEED=$RANDOM
DATETIME=$(date "+%m%d%H%M")
LOGFILE="bert_8xMI350x_${DATETIME}_${SEED}.log"
BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
# run
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -2,7 +2,7 @@
export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
@@ -2,7 +2,7 @@
export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
@@ -5,7 +5,7 @@ set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." NV=1
export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_green"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
+118
View File
@@ -0,0 +1,118 @@
import json, pprint
from tinygrad import fetch, nn, Tensor
from tinygrad.helpers import DEBUG
class FeedForward:
def __init__(self, model_dim, intermediate_dim):
self.proj_1 = nn.Linear(model_dim, 2*intermediate_dim, bias=False)
self.proj_2 = nn.Linear(intermediate_dim, model_dim, bias=False)
def __call__(self, x):
y_12 = self.proj_1(x)
y_1, y_2 = y_12.chunk(2, dim=-1)
return self.proj_2(y_1.silu() * y_2)
# NOTE: this RoPE doesn't match LLaMA's?
def _rotate_half(x: Tensor) -> Tensor:
x1, x2 = x.chunk(2, dim=-1)
return Tensor.cat(-x2, x1, dim=-1)
def _apply_rotary_pos_emb(x: Tensor, pos_sin: Tensor, pos_cos: Tensor) -> Tensor:
return (x * pos_cos) + (_rotate_half(x) * pos_sin)
class Attention:
def __init__(self, model_dim, num_query_heads, num_kv_heads, head_dim):
self.qkv_proj = nn.Linear(model_dim, (num_query_heads + num_kv_heads*2) * head_dim, bias=False)
self.num_query_heads, self.num_kv_heads = num_query_heads, num_kv_heads
self.head_dim = head_dim
self.q_norm = nn.RMSNorm(head_dim)
self.k_norm = nn.RMSNorm(head_dim)
self.out_proj = nn.Linear(num_query_heads * head_dim, model_dim, bias=False)
def __call__(self, x:Tensor) -> Tensor:
batch_size, seq_len, embed_dim = x.shape
qkv = self.qkv_proj(x)
qkv = qkv.reshape(batch_size, seq_len, self.num_query_heads+self.num_kv_heads*2, self.head_dim).transpose(1, 2)
xq,xk,xv = qkv.split([self.num_query_heads, self.num_kv_heads, self.num_kv_heads], dim=1)
xq = self.q_norm(xq)
xk = self.k_norm(xk)
# add positional embedding (how many kernels is this?)
freq_constant = 10000
inv_freq = 1.0 / (freq_constant ** (Tensor.arange(0, self.head_dim, 2) / self.head_dim))
pos_index_theta = Tensor.einsum("i,j->ij", Tensor.arange(seq_len), inv_freq)
emb = Tensor.cat(pos_index_theta, pos_index_theta, dim=-1)
cos_emb, sin_emb = emb.cos()[None, None, :, :], emb.sin()[None, None, :, :]
xq = _apply_rotary_pos_emb(xq, sin_emb, cos_emb)
xk = _apply_rotary_pos_emb(xk, sin_emb, cos_emb)
# grouped-query attention
num_groups = self.num_query_heads // self.num_kv_heads
xk = xk.repeat_interleave(num_groups, dim=1)
xv = xv.repeat_interleave(num_groups, dim=1)
# masked attention
#start_pos = 0
#mask = Tensor.full((1, 1, seq_len, start_pos+seq_len), float("-inf"), dtype=xq.dtype, device=xq.device).triu(start_pos+1)
#attn_output = xq.scaled_dot_product_attention(xk, xv, mask).transpose(1, 2)
# causal is fine, no mask needed
attn_output = xq.scaled_dot_product_attention(xk, xv, is_causal=True).transpose(1, 2)
return self.out_proj(attn_output.reshape(batch_size, seq_len, self.num_query_heads * self.head_dim))
class Layer:
def __init__(self, model_dim, intermediate_dim, num_query_heads, num_kv_heads, head_dim):
self.ffn = FeedForward(model_dim, intermediate_dim)
self.attn = Attention(model_dim, num_query_heads, num_kv_heads, head_dim)
self.ffn_norm = nn.RMSNorm(model_dim)
self.attn_norm = nn.RMSNorm(model_dim)
def __call__(self, x:Tensor) -> Tensor: # (batch, seq_len, embed_dim)
x = x + self.attn(self.attn_norm(x))
x = x + self.ffn(self.ffn_norm(x))
return x
# stupidly complex
def make_divisible(v, divisor):
new_v = max(divisor, int(v + divisor / 2) // divisor * divisor)
if new_v < 0.9 * v: new_v += divisor
return new_v
class Transformer:
def __init__(self, cfg):
if DEBUG >= 3: pprint.pp(cfg)
self.layers = [Layer(cfg['model_dim'], make_divisible(int(cfg["model_dim"] * cfg['ffn_multipliers'][i]), cfg['ffn_dim_divisor']),
cfg['num_query_heads'][i], cfg['num_kv_heads'][i], cfg['head_dim']) for i in range(cfg['num_transformer_layers'])]
self.norm = nn.RMSNorm(cfg['model_dim'])
self.token_embeddings = nn.Embedding(cfg['vocab_size'], cfg['model_dim'])
def __call__(self, tokens:Tensor):
# _bsz, seqlen = tokens.shape
x = self.token_embeddings(tokens)
for l in self.layers: x = l(x)
return self.norm(x) @ self.token_embeddings.weight.T
if __name__ == "__main__":
#model_name = "OpenELM-270M-Instruct"
model_name = "OpenELM-270M" # this is fp32
model = Transformer(json.loads(fetch(f"https://huggingface.co/apple/{model_name}/resolve/main/config.json?download=true").read_bytes()))
weights = nn.state.safe_load(fetch(f"https://huggingface.co/apple/{model_name}/resolve/main/model.safetensors?download=true"))
if DEBUG >= 3:
for k, v in weights.items(): print(k, v.shape)
nn.state.load_state_dict(model, {k.removeprefix("transformer."):v for k,v in weights.items()})
from sentencepiece import SentencePieceProcessor
tokenizer = SentencePieceProcessor(fetch("https://github.com/karpathy/llama2.c/raw/master/tokenizer.model").as_posix())
toks = [tokenizer.bos_id()] + tokenizer.encode("Some car brands include")
for i in range(100):
ttoks = Tensor([toks])
out = model(ttoks).realize()
t0 = out[0].argmax(axis=-1).tolist()
toks.append(t0[-1])
# hmmm...passthrough still doesn't match (it shouldn't, it outputs the most likely)
print(tokenizer.decode(toks))
#print(toks)
#print(tokenizer.decode(t0))
#print(t0)
@@ -0,0 +1,55 @@
from tinygrad.helpers import trange
from tinygrad.nn.datasets import mnist
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from functools import partial
class Model(nn.Module):
def __init__(self):
super().__init__()
self.c1 = nn.Conv2d(1, 32, 5)
self.c2 = nn.Conv2d(32, 32, 5)
self.bn1 = nn.BatchNorm(32)
self.m1 = nn.MaxPool2d(2)
self.c3 = nn.Conv2d(32, 64, 3)
self.c4 = nn.Conv2d(64, 64, 3)
self.bn2 = nn.BatchNorm(64)
self.m2 = nn.MaxPool2d(2)
self.lin = nn.Linear(576, 10)
def __call__(self, x):
x = mx.maximum(self.c1(x), 0)
x = mx.maximum(self.c2(x), 0)
x = self.m1(self.bn1(x))
x = mx.maximum(self.c3(x), 0)
x = mx.maximum(self.c4(x), 0)
x = self.m2(self.bn2(x))
return self.lin(mx.flatten(x, 1))
if __name__ == "__main__":
X_train, Y_train, X_test, Y_test = mnist()
X_train = mx.array(X_train.float().permute((0,2,3,1)).numpy())
Y_train = mx.array(Y_train.numpy())
X_test = mx.array(X_test.float().permute((0,2,3,1)).numpy())
Y_test = mx.array(Y_test.numpy())
model = Model()
optimizer = optim.Adam(1e-3)
def loss_fn(model, x, y): return nn.losses.cross_entropy(model(x), y).mean()
state = [model.state, optimizer.state]
@partial(mx.compile, inputs=state, outputs=state)
def step(samples):
# Compiled functions will also treat any inputs not in the parameter list as constants.
X,Y = X_train[samples], Y_train[samples]
loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
loss, grads = loss_and_grad_fn(model, X, Y)
optimizer.update(model, grads)
return loss
test_acc = float('nan')
for i in (t:=trange(70)):
samples = mx.random.randint(0, X_train.shape[0], (512,)) # putting this in JIT didn't work well
loss = step(samples)
if i%10 == 9: test_acc = ((model(X_test).argmax(axis=-1) == Y_test).sum() * 100 / X_test.shape[0]).item()
t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
+45
View File
@@ -0,0 +1,45 @@
import gymnasium as gym
import numpy as np
from gymnasium.envs.registration import register
# a very simple game
# one of <size> lights will light up
# take the action of the lit up light
# in <hard_mode>, you act differently based on the step number and need to track this
class PressTheLightUpButton(gym.Env):
metadata = {"render_modes": []}
def __init__(self, render_mode=None, size=2, game_length=10, hard_mode=False):
self.size, self.game_length = size, game_length
self.observation_space = gym.spaces.Box(0, 1, shape=(self.size,), dtype=np.float32)
self.action_space = gym.spaces.Discrete(self.size)
self.step_num = 0
self.done = True
self.hard_mode = hard_mode
def _get_obs(self):
obs = [0]*self.size
if self.step_num < len(self.state):
obs[self.state[self.step_num]] = 1
return np.array(obs, dtype=np.float32)
def reset(self, seed=None, options=None):
super().reset(seed=seed)
self.state = np.random.randint(0, self.size, size=self.game_length)
self.step_num = 0
self.done = False
return self._get_obs(), {}
def step(self, action):
target = ((action + self.step_num) % self.size) if self.hard_mode else action
reward = int(target == self.state[self.step_num])
self.step_num += 1
if not reward:
self.done = True
return self._get_obs(), reward, self.done, self.step_num >= self.game_length, {}
register(
id="PressTheLightUpButton-v0",
entry_point="examples.rl.lightupbutton:PressTheLightUpButton",
max_episode_steps=None,
)
+1 -1
View File
@@ -115,7 +115,7 @@ if __name__ == "__main__":
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
if not args.fakeweights:
default_weights_url = 'https://huggingface.co/sd2-community/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
weights_fn = args.weights_fn
if not weights_fn:
weights_url = args.weights_url if args.weights_url else default_weights_url
+136
View File
@@ -0,0 +1,136 @@
#!/usr/bin/env python
#inspired by https://github.com/Matuzas77/MNIST-0.17/blob/master/MNIST_final_solution.ipynb
import sys
import numpy as np
from tinygrad.nn.state import get_parameters
from tinygrad.tensor import Tensor
from tinygrad.nn import BatchNorm2d, optim
from tinygrad.helpers import getenv
from extra.datasets import fetch_mnist
from extra.augment import augment_img
from extra.training import train, evaluate
GPU = getenv("GPU")
QUICK = getenv("QUICK")
DEBUG = getenv("DEBUG")
class SqueezeExciteBlock2D:
def __init__(self, filters):
self.filters = filters
self.weight1 = Tensor.scaled_uniform(self.filters, self.filters//32)
self.bias1 = Tensor.scaled_uniform(1,self.filters//32)
self.weight2 = Tensor.scaled_uniform(self.filters//32, self.filters)
self.bias2 = Tensor.scaled_uniform(1, self.filters)
def __call__(self, input):
se = input.avg_pool2d(kernel_size=(input.shape[2], input.shape[3])) #GlobalAveragePool2D
se = se.reshape(shape=(-1, self.filters))
se = se.dot(self.weight1) + self.bias1
se = se.relu()
se = se.dot(self.weight2) + self.bias2
se = se.sigmoid().reshape(shape=(-1,self.filters,1,1)) #for broadcasting
se = input.mul(se)
return se
class ConvBlock:
def __init__(self, h, w, inp, filters=128, conv=3):
self.h, self.w = h, w
self.inp = inp
#init weights
self.cweights = [Tensor.scaled_uniform(filters, inp if i==0 else filters, conv, conv) for i in range(3)]
self.cbiases = [Tensor.scaled_uniform(1, filters, 1, 1) for i in range(3)]
#init layers
self._bn = BatchNorm2d(128)
self._seb = SqueezeExciteBlock2D(filters)
def __call__(self, input):
x = input.reshape(shape=(-1, self.inp, self.w, self.h))
for cweight, cbias in zip(self.cweights, self.cbiases):
x = x.pad(padding=[1,1,1,1]).conv2d(cweight).add(cbias).relu()
x = self._bn(x)
x = self._seb(x)
return x
class BigConvNet:
def __init__(self):
self.conv = [ConvBlock(28,28,1), ConvBlock(28,28,128), ConvBlock(14,14,128)]
self.weight1 = Tensor.scaled_uniform(128,10)
self.weight2 = Tensor.scaled_uniform(128,10)
def parameters(self):
if DEBUG: #keeping this for a moment
pars = [par for par in get_parameters(self) if par.requires_grad]
no_pars = 0
for par in pars:
print(par.shape)
no_pars += np.prod(par.shape)
print('no of parameters', no_pars)
return pars
else:
return get_parameters(self)
def save(self, filename):
with open(filename+'.npy', 'wb') as f:
for par in get_parameters(self):
#if par.requires_grad:
np.save(f, par.numpy())
def load(self, filename):
with open(filename+'.npy', 'rb') as f:
for par in get_parameters(self):
#if par.requires_grad:
try:
par.numpy()[:] = np.load(f)
if GPU:
par.gpu()
except:
print('Could not load parameter')
def forward(self, x):
x = self.conv[0](x)
x = self.conv[1](x)
x = x.avg_pool2d(kernel_size=(2,2))
x = self.conv[2](x)
x1 = x.avg_pool2d(kernel_size=(14,14)).reshape(shape=(-1,128)) #global
x2 = x.max_pool2d(kernel_size=(14,14)).reshape(shape=(-1,128)) #global
xo = x1.dot(self.weight1) + x2.dot(self.weight2)
return xo
if __name__ == "__main__":
lrs = [1e-4, 1e-5] if QUICK else [1e-3, 1e-4, 1e-5, 1e-5]
epochss = [2, 1] if QUICK else [13, 3, 3, 1]
BS = 32
lmbd = 0.00025
lossfn = lambda out,y: out.sparse_categorical_crossentropy(y) + lmbd*(model.weight1.abs() + model.weight2.abs()).sum()
X_train, Y_train, X_test, Y_test = fetch_mnist()
X_train = X_train.reshape(-1, 28, 28).astype(np.uint8)
X_test = X_test.reshape(-1, 28, 28).astype(np.uint8)
steps = len(X_train)//BS
np.random.seed(1337)
if QUICK:
steps = 1
X_test, Y_test = X_test[:BS], Y_test[:BS]
model = BigConvNet()
if len(sys.argv) > 1:
try:
model.load(sys.argv[1])
print('Loaded weights "'+sys.argv[1]+'", evaluating...')
evaluate(model, X_test, Y_test, BS=BS)
except:
print('could not load weights "'+sys.argv[1]+'".')
if GPU:
params = get_parameters(model)
[x.gpu_() for x in params]
for lr, epochs in zip(lrs, epochss):
optimizer = optim.Adam(model.parameters(), lr=lr)
for epoch in range(1,epochs+1):
#first epoch without augmentation
X_aug = X_train if epoch == 1 else augment_img(X_train)
train(model, X_aug, Y_train, optimizer, steps=steps, lossfn=lossfn, BS=BS)
accuracy = evaluate(model, X_test, Y_test, BS=BS)
model.save(f'examples/checkpoint{accuracy * 1e6:.0f}')
+17
View File
@@ -0,0 +1,17 @@
from tinygrad.tensor import Tensor
from tinygrad.nn import Conv2d, BatchNorm2d
from tinygrad.nn.state import get_parameters
if __name__ == "__main__":
with Tensor.train():
BS, C1, H, W = 4, 16, 224, 224
C2, K, S, P = 64, 7, 2, 1
x = Tensor.uniform(BS, C1, H, W)
conv = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
bn = BatchNorm2d(C2, track_running_stats=False)
for t in get_parameters([x, conv, bn]): t.realize()
print("running network")
x.sequential([conv, bn]).numpy()
+669
View File
@@ -0,0 +1,669 @@
# original implementation: https://github.com/svc-develop-team/so-vits-svc
from __future__ import annotations
import sys, logging, time, io, math, argparse, operator, numpy as np
from functools import partial, reduce
from pathlib import Path
from typing import Tuple, Optional, Type
from tinygrad import nn, dtypes, Tensor
from tinygrad.helpers import getenv, fetch
from tinygrad.nn.state import torch_load
from examples.vits import ResidualCouplingBlock, PosteriorEncoder, Encoder, ResBlock1, ResBlock2, LRELU_SLOPE, sequence_mask, split, get_hparams_from_file, load_checkpoint, weight_norm, HParams
from examples.sovits_helpers import preprocess
import soundfile
DEBUG = getenv("DEBUG")
F0_BIN = 256
F0_MAX = 1100.0
F0_MIN = 50.0
F0_MEL_MIN = 1127 * np.log(1 + F0_MIN / 700)
F0_MEL_MAX = 1127 * np.log(1 + F0_MAX / 700)
class SpeechEncoder:
def __init__(self, hidden_dim, model:ContentVec): self.hidden_dim, self.model = hidden_dim, model
def encode(self, ): raise NotImplementedError("implement me")
@classmethod
def load_from_pretrained(cls, checkpoint_path:str, checkpoint_url:str) -> ContentVec:
contentvec = ContentVec.load_from_pretrained(checkpoint_path, checkpoint_url)
return cls(contentvec)
class ContentVec256L9(SpeechEncoder):
def __init__(self, model:ContentVec): super().__init__(hidden_dim=256, model=model)
def encode(self, wav: Tensor):
feats = wav
if len(feats.shape) == 2: # double channels
feats = feats.mean(-1)
assert len(feats.shape) == 1, feats.dim()
feats = feats.reshape(1, -1)
padding_mask = Tensor.zeros_like(feats).cast(dtypes.bool)
logits = self.model.extract_features(feats.to(wav.device), padding_mask=padding_mask.to(wav.device), output_layer=9)
feats = self.model.final_proj(logits[0])
return feats.transpose(1,2)
class ContentVec768L12(SpeechEncoder):
def __init__(self, model:ContentVec): super().__init__(hidden_dim=768, model=model)
def encode(self, wav: Tensor):
feats = wav
if len(feats.shape) == 2: # double channels
feats = feats.mean(-1)
assert len(feats.shape) == 1, feats.dim()
feats = feats.reshape(1, -1)
padding_mask = Tensor.zeros_like(feats).cast(dtypes.bool)
logits = self.model.extract_features(feats.to(wav.device), padding_mask=padding_mask.to(wav.device), output_layer=12)
return logits[0].transpose(1,2)
# original code for contentvec: https://github.com/auspicious3000/contentvec/
class ContentVec:
# self.final_proj dims are hardcoded and depend on fairseq.data.dictionary Dictionary in the checkpoint. This param can't yet be loaded since there is no pickle for it. See with DEBUG=2.
# This means that the ContentVec only works with the hubert weights used in all SVC models
def __init__(self, cfg: HParams):
self.feature_grad_mult, self.untie_final_proj = cfg.feature_grad_mult, cfg.untie_final_proj
feature_enc_layers = eval(cfg.conv_feature_layers)
self.embed = feature_enc_layers[-1][0]
final_dim = cfg.final_dim if cfg.final_dim > 0 else cfg.encoder_embed_dim
self.feature_extractor = ConvFeatureExtractionModel(conv_layers=feature_enc_layers, dropout=0.0, mode=cfg.extractor_mode, conv_bias=cfg.conv_bias)
self.post_extract_proj = nn.Linear(self.embed, cfg.encoder_embed_dim) if self.embed != cfg.encoder_embed_dim else None
self.encoder = TransformerEncoder(cfg)
self.layer_norm = nn.LayerNorm(self.embed)
self.final_proj = nn.Linear(cfg.encoder_embed_dim, final_dim * 1) if self.untie_final_proj else nn.Linear(cfg.encoder_embed_dim, final_dim)
self.mask_emb = Tensor.uniform(cfg.encoder_embed_dim, dtype=dtypes.float32)
self.label_embs_concat = Tensor.uniform(504, final_dim, dtype=dtypes.float32)
def forward_features(self, source, padding_mask):
if self.feature_grad_mult > 0:
features = self.feature_extractor(source, padding_mask)
if self.feature_grad_mult != 1.0: pass # training: GradMultiply.forward(features, self.feature_grad_mult)
else:
features = self.feature_extractor(source, padding_mask)
return features
def forward_padding_mask(self, features, padding_mask): # replaces original forward_padding_mask for batch inference
lengths_org = tilde(padding_mask.cast(dtypes.bool)).cast(dtypes.int64).sum(1) # ensure its bool for tilde
lengths = (lengths_org - 400).float().div(320).floor().cast(dtypes.int64) + 1 # intermediate float to divide
padding_mask = lengths_to_padding_mask(lengths)
return padding_mask
def extract_features(self, source: Tensor, spk_emb:Tensor=None, padding_mask=None, ret_conv=False, output_layer=None, tap=False):
features = self.forward_features(source, padding_mask)
if padding_mask is not None:
padding_mask = self.forward_padding_mask(features, padding_mask)
features = features.transpose(1, 2)
features = self.layer_norm(features)
if self.post_extract_proj is not None:
features = self.post_extract_proj(features)
x, _ = self.encoder(features, spk_emb, padding_mask=padding_mask, layer=(None if output_layer is None else output_layer - 1), tap=tap)
res = features if ret_conv else x
return res, padding_mask
@classmethod
def load_from_pretrained(cls, checkpoint_path:str, checkpoint_url:str) -> ContentVec:
fetch(checkpoint_url, checkpoint_path)
cfg = load_fairseq_cfg(checkpoint_path)
enc = cls(cfg.model)
_ = load_checkpoint_enc(checkpoint_path, enc, None)
logging.debug(f"{cls.__name__}: Loaded model with cfg={cfg}")
return enc
class TransformerEncoder:
def __init__(self, cfg: HParams):
def make_conv() -> nn.Conv1d:
layer = nn.Conv1d(self.embedding_dim, self.embedding_dim, kernel_size=cfg.conv_pos, padding=cfg.conv_pos // 2, groups=cfg.conv_pos_groups)
std = std = math.sqrt(4 / (cfg.conv_pos * self.embedding_dim))
layer.weight, layer.bias = (Tensor.normal(*layer.weight.shape, std=std)), (Tensor.zeros(*layer.bias.shape))
# for training: layer.weights need to be weight_normed
return layer
self.dropout, self.embedding_dim, self.layer_norm_first, self.layerdrop, self.num_layers, self.num_layers_1 = cfg.dropout, cfg.encoder_embed_dim, cfg.layer_norm_first, cfg.encoder_layerdrop, cfg.encoder_layers, cfg.encoder_layers_1
self.pos_conv, self.pos_conv_remove = [make_conv()], (1 if cfg.conv_pos % 2 == 0 else 0)
self.layers = [
TransformerEncoderLayer(self.embedding_dim, cfg.encoder_ffn_embed_dim, cfg.encoder_attention_heads, self.dropout, cfg.attention_dropout, cfg.activation_dropout, cfg.activation_fn, self.layer_norm_first, cond_layer_norm=(i >= cfg.encoder_layers))
for i in range(cfg.encoder_layers + cfg.encoder_layers_1)
]
self.layer_norm = nn.LayerNorm(self.embedding_dim)
self.cond_layer_norm = CondLayerNorm(self.embedding_dim) if cfg.encoder_layers_1 > 0 else None
# training: apply init_bert_params
def __call__(self, x, spk_emb, padding_mask=None, layer=None, tap=False):
x, layer_results = self.extract_features(x, spk_emb, padding_mask, layer, tap)
if self.layer_norm_first and layer is None:
x = self.cond_layer_norm(x, spk_emb) if (self.num_layers_1 > 0) else self.layer_norm(x)
return x, layer_results
def extract_features(self, x: Tensor, spk_emb: Tensor, padding_mask=None, tgt_layer=None, tap=False):
if tgt_layer is not None: # and not self.training
assert tgt_layer >= 0 and tgt_layer < len(self.layers)
if padding_mask is not None:
# x[padding_mask] = 0
assert padding_mask.shape == x.shape[:len(padding_mask.shape)] # first few dims of x must match padding_mask
tmp_mask = padding_mask.unsqueeze(-1).repeat((1, 1, x.shape[-1]))
tmp_mask = tilde(tmp_mask.cast(dtypes.bool))
x = tmp_mask.where(x, 0)
x_conv = self.pos_conv[0](x.transpose(1,2))
if self.pos_conv_remove > 0: x_conv = x_conv[:, :, : -self.pos_conv_remove]
x_conv = x_conv.gelu().transpose(1, 2)
x = (x + x_conv).transpose(0, 1) # B x T x C -> T x B x C
if not self.layer_norm_first: x = self.layer_norm(x)
x = x.dropout(p=self.dropout)
layer_results = []
r = None
for i, layer in enumerate(self.layers):
if i < self.num_layers: # if (not self.training or (dropout_probability > self.layerdrop)) and (i < self.num_layers):
assert layer.cond_layer_norm == False
x = layer(x, self_attn_padding_mask=padding_mask, need_weights=False)
if tgt_layer is not None or tap:
layer_results.append(x.transpose(0, 1))
if i>= self.num_layers:
assert layer.cond_layer_norm == True
x = layer(x, emb=spk_emb, self_attn_padding_mask=padding_mask, need_weights=False)
if i == tgt_layer:
r = x
break
if r is not None:
x = r
x = x.transpose(0, 1) # T x B x C -> B x T x C
return x, layer_results
class TransformerEncoderLayer:
def __init__(self, embedding_dim=768.0, ffn_embedding_dim=3072.0, num_attention_heads=8.0, dropout=0.1, attention_dropout=0.1, activation_dropout=0.1, activation_fn="relu", layer_norm_first=False, cond_layer_norm=False):
def get_activation_fn(activation):
if activation == "relu": return Tensor.relu
if activation == "gelu": return Tensor.gelu
else: raise RuntimeError(f"activation function={activation} is not forseen")
self.embedding_dim, self.dropout, self.activation_dropout, self.layer_norm_first, self.num_attention_heads, self.cond_layer_norm, self.activation_fn = embedding_dim, dropout, activation_dropout, layer_norm_first, num_attention_heads, cond_layer_norm, get_activation_fn(activation_fn)
self.self_attn = MultiHeadAttention(self.embedding_dim, self.num_attention_heads)
self.self_attn_layer_norm = nn.LayerNorm(self.embedding_dim) if not cond_layer_norm else CondLayerNorm(self.embedding_dim)
self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)
self.final_layer_norm = nn.LayerNorm(self.embedding_dim) if not cond_layer_norm else CondLayerNorm(self.embedding_dim)
def __call__(self, x:Tensor, self_attn_mask:Tensor=None, self_attn_padding_mask:Tensor=None, emb:Tensor=None, need_weights=False):
#self_attn_padding_mask = self_attn_padding_mask.reshape(x.shape[0], 1, 1, self_attn_padding_mask.shape[1]).expand(-1, self.num_attention_heads, -1, -1).reshape(x.shape[0] * self.num_attention_heads, 1, self_attn_padding_mask.shape[1]) if self_attn_padding_mask is not None else None
assert self_attn_mask is None and self_attn_padding_mask is not None
residual = x
if self.layer_norm_first:
x = self.self_attn_layer_norm(x) if not self.cond_layer_norm else self.self_attn_layer_norm(x, emb)
x = self.self_attn(x=x, mask=self_attn_padding_mask)
x = x.dropout(self.dropout)
x = residual + x
x = self.final_layer_norm(x) if not self.cond_layer_norm else self.final_layer_norm(x, emb)
x = self.activation_fn(self.fc1(x))
x = x.dropout(self.activation_dropout)
x = self.fc2(x)
x = x.dropout(self.dropout)
x = residual + x
else:
x = self.self_attn(x=x, mask=self_attn_padding_mask)
x = x.dropout(self.dropout)
x = residual + x
x = self.self_attn_layer_norm(x) if not self.cond_layer_norm else self.self_attn_layer_norm(x, emb)
residual = x
x = self.activation_fn(self.fc1(x))
x = x.dropout(self.activation_dropout)
x = self.fc2(x)
x = x.dropout(self.dropout)
x = residual + x
x = self.final_layer_norm(x) if not self.cond_layer_norm else self.final_layer_norm(x, emb)
return x
class MultiHeadAttention:
def __init__(self, n_state, n_head):
self.n_state, self.n_head = n_state, n_head
self.q_proj, self.k_proj, self.v_proj, self.out_proj = [nn.Linear(n_state, n_state) for _ in range(4)]
def __call__(self, x:Tensor, xa:Optional[Tensor]=None, mask:Optional[Tensor]=None):
x = x.transpose(0,1) # TxBxC -> BxTxC
q, k, v = self.q_proj(x), self.k_proj(xa or x), self.v_proj(xa or x)
q, k, v = [x.reshape(*q.shape[:2], self.n_head, -1) for x in (q, k, v)]
wv = Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), None).transpose(1, 2).reshape(*x.shape[:2], -1)
ret = self.out_proj(wv).transpose(0,1) # BxTxC -> TxBxC
return ret
class ConvFeatureExtractionModel:
def __init__(self, conv_layers, dropout=.0, mode="default", conv_bias=False):
assert mode in {"default", "group_norm_masked", "layer_norm"}
def block(n_in, n_out, k, stride, is_layer_norm=False, is_group_norm=False, conv_bias=False):
def make_conv():
conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias)
conv.weight = Tensor.kaiming_normal(*conv.weight.shape)
return conv
assert (is_layer_norm and is_group_norm) == False, "layer norm and group norm are exclusive"
if is_layer_norm:
return [make_conv(), partial(Tensor.dropout, p=dropout),[partial(Tensor.transpose, dim0=-2, dim1=-1), nn.LayerNorm(dim, elementwise_affine=True), partial(Tensor.transpose, dim0=-2, dim1=-1)], Tensor.gelu]
elif is_group_norm and mode == "default":
return [make_conv(), partial(Tensor.dropout, p=dropout), nn.GroupNorm(dim, dim, affine=True), Tensor.gelu]
elif is_group_norm and mode == "group_norm_masked":
return [make_conv(), partial(Tensor.dropout, p=dropout), GroupNormMasked(dim, dim, affine=True), Tensor.gelu]
else:
return [make_conv(), partial(Tensor.dropout, p=dropout), Tensor.gelu]
in_d, self.conv_layers, self.mode = 1, [], mode
for i, cl in enumerate(conv_layers):
assert len(cl) == 3, "invalid conv definition: " + str(cl)
(dim, k, stride) = cl
if i == 0: self.cl = cl
self.conv_layers.append(block(in_d, dim, k, stride, is_layer_norm=(mode == "layer_norm"), is_group_norm=((mode == "default" or mode == "group_norm_masked") and i == 0), conv_bias=conv_bias))
in_d = dim
def __call__(self, x:Tensor, padding_mask:Tensor):
x = x.unsqueeze(1) # BxT -> BxCxT
if self.mode == "group_norm_masked":
if padding_mask is not None:
_, k, stride = self.cl
lengths_org = tilde(padding_mask.cast(dtypes.bool)).cast(dtypes.int64).sum(1) # ensure padding_mask is bool for tilde
lengths = (((lengths_org - k) / stride) + 1).floor().cast(dtypes.int64)
padding_mask = tilde(lengths_to_padding_mask(lengths)).cast(dtypes.int64) # lengths_to_padding_mask returns bool tensor
x = self.conv_layers[0][0](x) # padding_mask is numeric
x = self.conv_layers[0][1](x)
x = self.conv_layers[0][2](x, padding_mask)
x = self.conv_layers[0][3](x)
else:
x = x.sequential(self.conv_layers[0]) # default
for _, conv in enumerate(self.conv_layers[1:], start=1):
conv = reduce(lambda a,b: operator.iconcat(a,b if isinstance(b, list) else [b]), conv, []) # flatten
x = x.sequential(conv)
return x
class CondLayerNorm: # https://github.com/auspicious3000/contentvec/blob/main/contentvec/modules/cond_layer_norm.py#L10
def __init__(self, dim_last, eps=1e-5, dim_spk=256, elementwise_affine=True):
self.dim_last, self.eps, self.dim_spk, self.elementwise_affine = dim_last, eps, dim_spk, elementwise_affine
if self.elementwise_affine:
self.weight_ln = nn.Linear(self.dim_spk, self.dim_last, bias=False)
self.bias_ln = nn.Linear(self.dim_spk, self.dim_last, bias=False)
self.weight_ln.weight, self.bias_ln.weight = (Tensor.ones(*self.weight_ln.weight.shape)), (Tensor.zeros(*self.bias_ln.weight.shape))
def __call__(self, x: Tensor, spk_emb: Tensor):
axis = tuple(-1-i for i in range(len(x.shape[1:])))
x = x.layernorm(axis=axis, eps=self.eps)
if not self.elementwise_affine: return x
weights, bias = self.weight_ln(spk_emb), self.bias_ln(spk_emb)
return weights * x + bias
class GroupNormMasked: # https://github.com/auspicious3000/contentvec/blob/d746688a32940f4bee410ed7c87ec9cf8ff04f74/contentvec/modules/fp32_group_norm.py#L16
def __init__(self, num_groups, num_channels, eps=1e-5, affine=True):
self.num_groups, self.num_channels, self.eps, self.affine = num_groups, num_channels, eps, affine
self.weight, self.bias = (Tensor.ones(num_channels)), (Tensor.zeros(num_channels)) if self.affine else (None, None)
def __call__(self, x:Tensor, mask:Tensor):
bsz, n_c, length = x.shape
assert n_c % self.num_groups == 0
x = x.reshape(bsz, self.num_groups, n_c // self.num_groups, length)
if mask is None: mask = Tensor.ones_like(x)
else: mask = mask.reshape(bsz, 1, 1, length)
x = x * mask
lengths = mask.sum(axis=3, keepdim=True)
assert x.shape[2] == 1
mean_ = x.mean(dim=3, keepdim=True)
mean = mean_ * length / lengths
var = (((x.std(axis=3, keepdim=True) ** 2) + mean_**2) * length / lengths - mean**2) + self.eps
return x.add(-mean).div(var.sqrt()).reshape(bsz, n_c, length).mul(self.weight.reshape(1,-1,1)).add(self.bias.reshape(1,-1,1))
class Synthesizer:
def __init__(self, spec_channels, segment_size, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels, ssl_dim, n_speakers, sampling_rate=44100, vol_embedding=False, n_flow_layer=4, **kwargs):
self.spec_channels, self.inter_channels, self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout, self.resblock, self.resblock_kernel_sizes, self.resblock_dilation_sizes, self.upsample_rates, self.upsample_initial_channel, self.upsample_kernel_sizes, self.segment_size, self.n_speakers, self.gin_channels, self.vol_embedding = spec_channels, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, segment_size, n_speakers, gin_channels, vol_embedding
self.emb_g = nn.Embedding(n_speakers, gin_channels)
if vol_embedding: self.emb_vol = nn.Linear(1, hidden_channels)
self.pre = nn.Conv1d(ssl_dim, hidden_channels, kernel_size=5, padding=2)
self.enc_p = TextEncoder(inter_channels, hidden_channels, kernel_size, n_layers, filter_channels=filter_channels, n_heads=n_heads, p_dropout=p_dropout)
self.dec = Generator(sampling_rate, inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels)
self.enc_q = PosteriorEncoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, n_flow_layer, gin_channels=gin_channels)
self.emb_uv = nn.Embedding(vocab_size=2, embed_size=hidden_channels)
def infer(self, c:Tensor, f0:Tensor, uv:Tensor, g:Tensor=None, noise_scale=0.35, seed=52468, vol=None) -> Tuple[Tensor, Tensor]:
Tensor.manual_seed(getenv('SEED', seed))
c_lengths = (Tensor.ones([c.shape[0]]) * c.shape[-1]).to(c.device)
if len(g.shape) == 1: g = g.unsqueeze(0)
g = self.emb_g(g).transpose(1, 2)
x_mask = sequence_mask(c_lengths, c.shape[2]).unsqueeze(1).cast(c.dtype)
vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol is not None and self.vol_embedding else 0
x = self.pre(c) * x_mask + self.emb_uv(uv.cast(dtypes.int64)).transpose(1, 2) + vol
z_p, _, _, c_mask = self.enc_p.forward(x, x_mask, f0=self._f0_to_coarse(f0), noise_scale=noise_scale)
z = self.flow.forward(z_p, c_mask, g=g, reverse=True)
o = self.dec.forward(z * c_mask, g=g, f0=f0)
return o,f0
def _f0_to_coarse(self, f0 : Tensor):
f0_mel = 1127 * (1 + f0 / 700).log()
a = (F0_BIN - 2) / (F0_MEL_MAX - F0_MEL_MIN)
b = F0_MEL_MIN * a - 1.
f0_mel = (f0_mel > 0).where(f0_mel * a - b, f0_mel)
f0_coarse = f0_mel.ceil().cast(dtype=dtypes.int64)
f0_coarse = f0_coarse * (f0_coarse > 0)
f0_coarse = f0_coarse + ((f0_coarse < 1) * 1)
f0_coarse = f0_coarse * (f0_coarse < F0_BIN)
f0_coarse = f0_coarse + ((f0_coarse >= F0_BIN) * (F0_BIN - 1))
return f0_coarse
@classmethod
def load_from_pretrained(cls, config_path:str, config_url:str, weights_path:str, weights_url:str) -> Synthesizer:
fetch(config_url, config_path)
hps = get_hparams_from_file(config_path)
fetch(weights_url, weights_path)
net_g = cls(hps.data.filter_length // 2 + 1, hps.train.segment_size // hps.data.hop_length, **hps.model)
_ = load_checkpoint(weights_path, net_g, None, skip_list=["f0_decoder"])
logging.debug(f"{cls.__name__}:Loaded model with hps: {hps}")
return net_g, hps
class TextEncoder:
def __init__(self, out_channels, hidden_channels, kernel_size, n_layers, gin_channels=0, filter_channels=None, n_heads=None, p_dropout=None):
self.out_channels, self.hidden_channels, self.kernel_size, self.n_layers, self.gin_channels = out_channels, hidden_channels, kernel_size, n_layers, gin_channels
self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
self.f0_emb = nn.Embedding(256, hidden_channels) # n_vocab = 256
self.enc_ = Encoder(hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout)
def forward(self, x, x_mask, f0=None, noise_scale=1):
x = x + self.f0_emb(f0).transpose(1, 2)
x = self.enc_.forward(x * x_mask, x_mask)
stats = self.proj(x) * x_mask
m, logs = split(stats, self.out_channels, dim=1)
z = (m + randn_like(m) * logs.exp() * noise_scale) * x_mask
return z, m, logs, x_mask
class Upsample:
def __init__(self, scale_factor):
assert scale_factor % 1 == 0, "Only integer scale factor allowed."
self.scale = int(scale_factor)
def forward(self, x:Tensor):
repeats = tuple([1] * len(x.shape) + [self.scale])
new_shape = (*x.shape[:-1], x.shape[-1] * self.scale)
return x.unsqueeze(-1).repeat(repeats).reshape(new_shape)
class SineGen:
def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voice_threshold=0, flag_for_pulse=False):
self.sine_amp, self.noise_std, self.harmonic_num, self.sampling_rate, self.voiced_threshold, self.flag_for_pulse = sine_amp, noise_std, harmonic_num, samp_rate, voice_threshold, flag_for_pulse
self.dim = self.harmonic_num + 1
def _f02uv(self, f0): return (f0 > self.voiced_threshold).float() #generate uv signal
def _f02sine(self, f0_values):
def padDiff(x : Tensor): return (x.pad((0,0,-1,1)) - x).pad((0,0,0,-1))
def mod(x: Tensor, n: int) -> Tensor: return x - n * x.div(n).floor() # this is what the % operator does in pytorch.
rad_values = mod((f0_values / self.sampling_rate) , 1) # convert to F0 in rad
rand_ini = Tensor.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device) # initial phase noise
#rand_ini[:, 0] = 0
m = Tensor.ones(f0_values.shape[0]).unsqueeze(1).pad((0,f0_values.shape[2]-1,0,0)).cast(dtypes.bool)
m = tilde(m)
rand_ini = m.where(rand_ini, 0)
#rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
tmp = rad_values[:, 0, :] + rand_ini
m = Tensor.ones(tmp.shape).pad((0,0,0,rad_values.shape[1]-1,0)).cast(dtypes.bool)
m = tilde(m)
tmp = tmp.unsqueeze(1).pad((0,0,0,rad_values.shape[1]-1,0))
rad_values = m.where(rad_values, tmp)
tmp_over_one = mod(rad_values.cumsum(1), 1)
tmp_over_one_idx = padDiff(tmp_over_one) < 0
cumsum_shift = Tensor.zeros_like(rad_values)
#cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
tmp_over_one_idx = (tmp_over_one_idx * -1.0).pad((0,0,1,0))
cumsum_shift = tmp_over_one_idx
sines = ((rad_values + cumsum_shift).cumsum(1) * 2 * np.pi).sin()
return sines
def forward(self, f0, upp=None):
fn = f0.mul(Tensor([[range(1, self.harmonic_num + 2)]], dtype=dtypes.float32).to(f0.device))
sine_waves = self._f02sine(fn) * self.sine_amp #generate sine waveforms
uv = self._f02uv(f0) # generate uv signal
noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
noise = noise_amp * randn_like(sine_waves)
sine_waves = sine_waves * uv + noise
return sine_waves, uv, noise
class SourceHnNSF:
def __init__(self, sampling_rate, harmonic_num=0, sine_amp=0.1, add_noise_std=0.003, voiced_threshold=0):
self.sine_amp, self.noise_std = sine_amp, add_noise_std
self.l_sin_gen = SineGen(sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold)
self.l_linear = nn.Linear(harmonic_num + 1, 1)
def forward(self, x, upp=None):
sine_waves, uv, _ = self.l_sin_gen.forward(x, upp)
sine_merge = self.l_linear(sine_waves.cast(self.l_linear.weight.dtype)).tanh()
noise = randn_like(uv) * self.sine_amp / 3
return sine_merge, noise, uv
# most of the hifigan in standard vits is reused here, but need to upsample and construct harmonic source from f0
class Generator:
def __init__(self, sampling_rate, inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels):
self.sampling_rate, self.inter_channels, self.resblock, self.resblock_kernel_sizes, self.resblock_dilation_sizes, self.upsample_rates, self.upsample_initial_channel, self.upsample_kernel_sizes, self.gin_channels = sampling_rate, inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels
self.num_kernels, self.num_upsamples = len(resblock_kernel_sizes), len(upsample_rates)
self.conv_pre = nn.Conv1d(inter_channels, upsample_initial_channel, 7, 1, padding=3)
self.f0_upsamp = Upsample(scale_factor=np.prod(upsample_rates))
self.m_source = SourceHnNSF(sampling_rate, harmonic_num=8)
resblock = ResBlock1 if resblock == '1' else ResBlock2
self.ups, self.noise_convs, self.resblocks = [], [], []
for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
c_cur = upsample_initial_channel//(2**(i+1))
self.ups.append(nn.ConvTranspose1d(upsample_initial_channel//(2**i), c_cur, k, u, padding=(k-u)//2))
stride_f0 = int(np.prod(upsample_rates[i + 1:]))
self.noise_convs.append(nn.Conv1d(1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2) if (i + 1 < len(upsample_rates)) else nn.Conv1d(1, c_cur, kernel_size=1))
for i in range(len(self.ups)):
ch = upsample_initial_channel // (2 ** (i + 1))
for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
self.resblocks.append(resblock(ch, k, d))
self.conv_post = nn.Conv1d(ch, 1, 7, 1, padding=3)
if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
self.upp = np.prod(upsample_rates)
def forward(self, x, f0, g=None):
f0 = self.f0_upsamp.forward(f0[:, None]).transpose(1, 2) # bs,n,t
har_source, _, _ = self.m_source.forward(f0, self.upp)
har_source = har_source.transpose(1, 2)
x = self.conv_pre(x)
if g is not None: x = x + self.cond(g)
for i in range(self.num_upsamples):
x, xs = self.ups[i](x.leaky_relu(LRELU_SLOPE)), None
x_source = self.noise_convs[i](har_source)
x = x + x_source
for j in range(self.num_kernels):
if xs is None: xs = self.resblocks[i * self.num_kernels + j].forward(x)
else: xs += self.resblocks[i * self.num_kernels + j].forward(x)
x = xs / self.num_kernels
return self.conv_post(x.leaky_relu()).tanh()
# **** helpers ****
def randn_like(x:Tensor) -> Tensor: return Tensor.randn(*x.shape, dtype=x.dtype).to(device=x.device)
def tilde(x: Tensor) -> Tensor:
if x.dtype == dtypes.bool: return (1 - x).cast(dtypes.bool)
return (x + 1) * -1 # this seems to be what the ~ operator does in pytorch for non bool
def lengths_to_padding_mask(lens:Tensor) -> Tensor:
bsz, max_lens = lens.shape[0], lens.max().numpy().item()
mask = Tensor.arange(max_lens).to(lens.device).reshape(1, max_lens)
mask = mask.expand(bsz, -1) >= lens.reshape(bsz, 1).expand(-1, max_lens)
return mask.cast(dtypes.bool)
def repeat_expand_2d_left(content, target_len): # content : [h, t]
src_len = content.shape[-1]
temp = np.arange(src_len+1) * target_len / src_len
current_pos, cols = 0, []
for i in range(target_len):
if i >= temp[current_pos+1]:
current_pos += 1
cols.append(content[:, current_pos])
return Tensor.stack(*cols).transpose(0, 1)
def load_fairseq_cfg(checkpoint_path):
assert Path(checkpoint_path).is_file()
state = torch_load(checkpoint_path)
cfg = state["cfg"] if ("cfg" in state and state["cfg"] is not None) else None
if cfg is None: raise RuntimeError(f"No cfg exist in state keys = {state.keys()}")
return HParams(**cfg)
def load_checkpoint_enc(checkpoint_path, model: ContentVec, optimizer=None, skip_list=[]):
assert Path(checkpoint_path).is_file()
start_time = time.time()
checkpoint_dict = torch_load(checkpoint_path)
saved_state_dict = checkpoint_dict['model']
weight_g, weight_v, parent = None, None, None
for key, v in saved_state_dict.items():
if any(layer in key for layer in skip_list): continue
try:
obj, skip = model, False
for k in key.split('.'):
if k.isnumeric(): obj = obj[int(k)]
elif isinstance(obj, dict): obj = obj[k]
else:
if k in ["weight_g", "weight_v"]:
parent, skip = obj, True
if k == "weight_g": weight_g = v
else: weight_v = v
if not skip:
parent = obj
obj = getattr(obj, k)
if weight_g and weight_v:
setattr(obj, "weight_g", weight_g.numpy())
setattr(obj, "weight_v", weight_v.numpy())
obj, v = getattr(parent, "weight"), weight_norm(weight_v, weight_g, 0)
weight_g, weight_v, parent, skip = None, None, None, False
if not skip and obj.shape == v.shape:
if "feature_extractor" in key and (isinstance(parent, (nn.GroupNorm, nn.LayerNorm))): # cast
obj.assign(v.to(obj.device).float())
else:
obj.assign(v.to(obj.device))
elif not skip: logging.error(f"MISMATCH SHAPE IN {key}, {obj.shape} {v.shape}")
except Exception as e: raise e
logging.info(f"Loaded checkpoint '{checkpoint_path}' in {time.time() - start_time:.4f}s")
return model, optimizer
def pad_array(arr, target_length):
current_length = arr.shape[0]
if current_length >= target_length: return arr
pad_width = target_length - current_length
pad_left = pad_width // 2
pad_right = pad_width - pad_left
padded_arr = np.pad(arr, (pad_left, pad_right), 'constant', constant_values=(0, 0))
return padded_arr
def split_list_by_n(list_collection, n, pre=0):
for i in range(0, len(list_collection), n):
yield list_collection[i-pre if i-pre>=0 else i: i + n]
def get_sid(spk2id:HParams, speaker:str) -> Tensor:
speaker_id = spk2id[speaker]
if not speaker_id and type(speaker) is int:
if len(spk2id.__dict__) >= speaker: speaker_id = speaker
if speaker_id is None: raise RuntimeError(f"speaker={speaker} not in the speaker list")
return Tensor([int(speaker_id)], dtype=dtypes.int64).unsqueeze(0)
def get_encoder(ssl_dim) -> Type[SpeechEncoder]:
if ssl_dim == 256: return ContentVec256L9
if ssl_dim == 768: return ContentVec768L12
#########################################################################################
# CODE: https://github.com/svc-develop-team/so-vits-svc
#########################################################################################
# CONTENTVEC:
# CODE: https://github.com/auspicious3000/contentvec
# PAPER: https://arxiv.org/abs/2204.09224
#########################################################################################
# INSTALLATION: dependencies are for preprocessing and loading/saving audio.
# pip3 install soundfile librosa praat-parselmouth
#########################################################################################
# EXAMPLE USAGE:
# python3 examples/so_vits_svc.py --model tf2spy --file ~/recording.wav
#########################################################################################
# DEMO USAGE (uses audio sample from LJ-Speech):
# python3 examples/so_vits_svc.py --model saul_goodman
#########################################################################################
SO_VITS_SVC_PATH = Path(__file__).parents[1] / "weights/So-VITS-SVC"
VITS_MODELS = { # config_path, weights_path, config_url, weights_url
"saul_goodman" : (SO_VITS_SVC_PATH / "config_saul_gman.json", SO_VITS_SVC_PATH / "pretrained_saul_gman.pth", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/Saul_Goodman_80000/config.json", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/Saul_Goodman_80000/G_80000.pth"),
"drake" : (SO_VITS_SVC_PATH / "config_drake.json", SO_VITS_SVC_PATH / "pretrained_drake.pth", "https://huggingface.co/jaspa/so-vits-svc/resolve/main/aubrey/config_aubrey.json", "https://huggingface.co/jaspa/so-vits-svc/resolve/main/aubrey/pretrained_aubrey.pth"),
"cartman" : (SO_VITS_SVC_PATH / "config_cartman.json", SO_VITS_SVC_PATH / "pretrained_cartman.pth", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/EricCartman/config.json", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/EricCartman/G_10200.pth"),
"tf2spy" : (SO_VITS_SVC_PATH / "config_tf2spy.json", SO_VITS_SVC_PATH / "pretrained_tf2spy.pth", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_spy_60k/config.json", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_spy_60k/G_60000.pth"),
"tf2heavy" : (SO_VITS_SVC_PATH / "config_tf2heavy.json", SO_VITS_SVC_PATH / "pretrained_tf2heavy.pth", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_heavy_100k/config.json", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_heavy_100k/G_100000.pth"),
"lady_gaga" : (SO_VITS_SVC_PATH / "config_gaga.json", SO_VITS_SVC_PATH / "pretrained_gaga.pth", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/LadyGaga/config.json", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/LadyGaga/G_14400.pth")
}
ENCODER_MODELS = { # weights_path, weights_url
"contentvec": (SO_VITS_SVC_PATH / "contentvec_checkpoint.pt", "https://huggingface.co/lj1995/VoiceConversionWebUI/resolve/main/hubert_base.pt")
}
ENCODER_MODEL = "contentvec"
DEMO_PATH, DEMO_URL = Path(__file__).parents[1] / "temp/LJ037-0171.wav", "https://keithito.com/LJ-Speech-Dataset/LJ037-0171.wav"
if __name__=="__main__":
logging.basicConfig(stream=sys.stdout, level=(logging.INFO if DEBUG < 1 else logging.DEBUG))
parser = argparse.ArgumentParser()
parser.add_argument("-m", "--model", default=None, help=f"Specify the model to use. All supported models: {VITS_MODELS.keys()}", required=True)
parser.add_argument("-f", "--file", default=DEMO_PATH, help=f"Specify the path of the input file")
parser.add_argument("--out_dir", default=str(Path(__file__).parents[1] / "temp"), help="Specify the output path.")
parser.add_argument("--out_path", default=None, help="Specify the full output path. Overrides the --out_dir and --name parameter.")
parser.add_argument("--base_name", default="test", help="Specify the base of the output file name. Default is 'test'.")
parser.add_argument("--speaker", default=None, help="If not specified, the first available speaker is chosen. Usually there is only one speaker per model.")
parser.add_argument("--noise_scale", default=0.4)
parser.add_argument("--tran", default=0.0, help="Pitch shift, supports positive and negative (semitone) values. Default 0.0")
parser.add_argument("--pad_seconds", default=0.5)
parser.add_argument("--lg_num", default=0.0)
parser.add_argument("--clip_seconds", default=0.0)
parser.add_argument("--slice_db", default=-40)
args = parser.parse_args()
vits_model = args.model
encoder_location, vits_location = ENCODER_MODELS[ENCODER_MODEL], VITS_MODELS[vits_model]
Tensor.training = False
# Get Synthesizer and ContentVec
net_g, hps = Synthesizer.load_from_pretrained(vits_location[0], vits_location[2], vits_location[1], vits_location[3])
Encoder = get_encoder(hps.model.ssl_dim)
encoder = Encoder.load_from_pretrained(encoder_location[0], encoder_location[1])
# model config args
target_sample, spk2id, hop_length, target_sample = hps.data.sampling_rate, hps.spk, hps.data.hop_length, hps.data.sampling_rate
vol_embedding = hps.model.vol_embedding if hasattr(hps.data, "vol_embedding") and hps.model.vol_embedding is not None else False
# args
slice_db, clip_seconds, lg_num, pad_seconds, tran, noise_scale, audio_path = args.slice_db, args.clip_seconds, args.lg_num, args.pad_seconds, args.tran, args.noise_scale, args.file
speaker = args.speaker if args.speaker is not None else list(hps.spk.__dict__.keys())[0]
### Loading audio and slicing ###
if audio_path == DEMO_PATH: fetch(DEMO_URL, DEMO_PATH)
assert Path(audio_path).is_file() and Path(audio_path).suffix == ".wav"
chunks = preprocess.cut(audio_path, db_thresh=slice_db)
audio_data, audio_sr = preprocess.chunks2audio(audio_path, chunks)
per_size = int(clip_seconds * audio_sr)
lg_size = int(lg_num * audio_sr)
### Infer per slice ###
global_frame = 0
audio = []
for (slice_tag, data) in audio_data:
print(f"\n====segment start, {round(len(data) / audio_sr, 3)}s====")
length = int(np.ceil(len(data) / audio_sr * target_sample))
if slice_tag:
print("empty segment")
_audio = np.zeros(length)
audio.extend(list(pad_array(_audio, length)))
global_frame += length // hop_length
continue
datas = [data] if per_size == 0 else split_list_by_n(data, per_size, lg_size)
for k, dat in enumerate(datas):
per_length = int(np.ceil(len(dat) / audio_sr * target_sample)) if clip_seconds!=0 else length
pad_len = int(audio_sr * pad_seconds)
dat = np.concatenate([np.zeros([pad_len]), dat, np.zeros([pad_len])])
raw_path = io.BytesIO()
soundfile.write(raw_path, dat, audio_sr, format="wav")
raw_path.seek(0)
### Infer START ###
wav, sr = preprocess.load_audiofile(raw_path)
wav = preprocess.sinc_interp_resample(wav, sr, target_sample)[0]
wav16k, f0, uv = preprocess.get_unit_f0(wav, tran, hop_length, target_sample)
sid = get_sid(spk2id, speaker)
n_frames = f0.shape[1]
# ContentVec infer
start = time.time()
c = encoder.encode(wav16k)
c = repeat_expand_2d_left(c.squeeze(0).realize(), f0.shape[1]) # interpolate speech encoding to match f0
c = c.unsqueeze(0).realize()
enc_time = time.time() - start
# VITS infer
vits_start = time.time()
out_audio, f0 = net_g.infer(c, f0=f0, uv=uv, g=sid, noise_scale=noise_scale, vol=None)
out_audio = out_audio[0,0].float().realize()
vits_time = time.time() - vits_start
infer_time = time.time() - start
logging.info("total infer time:{:.2f}s, speech_enc time:{:.2f}s, vits time:{:.2f}s".format(infer_time, enc_time, vits_time))
### Infer END ###
out_sr, out_frame = out_audio.shape[-1], n_frames
global_frame += out_frame
_audio = out_audio.numpy()
pad_len = int(target_sample * pad_seconds)
_audio = _audio[pad_len:-pad_len]
_audio = pad_array(_audio, per_length)
audio.extend(list(_audio))
audio = np.array(audio)
out_path = Path(args.out_path or Path(args.out_dir)/f"{args.model}{f'_spk_{speaker}'}_{args.base_name}.wav")
out_path.parent.mkdir(parents=True, exist_ok=True)
soundfile.write(out_path, audio, target_sample, format="flac")
logging.info(f"Saved audio output to {out_path}")
+204
View File
@@ -0,0 +1,204 @@
import math
from typing import Optional, Tuple
from tinygrad import Tensor, dtypes
import librosa
import soundfile
import numpy as np
import parselmouth
class PMF0Predictor: # from https://github.com/svc-develop-team/so-vits-svc/
def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
self.hop_length, self.f0_min, self.f0_max, self.sampling_rate, self.name = hop_length, f0_min, f0_max, sampling_rate, "pm"
def interpolate_f0(self,f0):
vuv_vector = np.zeros_like(f0, dtype=np.float32)
vuv_vector[f0 > 0.0] = 1.0
vuv_vector[f0 <= 0.0] = 0.0
nzindex = np.nonzero(f0)[0]
data = f0[nzindex]
nzindex = nzindex.astype(np.float32)
time_org = self.hop_length / self.sampling_rate * nzindex
time_frame = np.arange(f0.shape[0]) * self.hop_length / self.sampling_rate
if data.shape[0] <= 0: return np.zeros(f0.shape[0], dtype=np.float32),vuv_vector
if data.shape[0] == 1: return np.ones(f0.shape[0], dtype=np.float32) * f0[0],vuv_vector
f0 = np.interp(time_frame, time_org, data, left=data[0], right=data[-1])
return f0,vuv_vector
def compute_f0(self,wav,p_len=None):
x = wav
if p_len is None: p_len = x.shape[0]//self.hop_length
else: assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
time_step = self.hop_length / self.sampling_rate * 1000
f0 = parselmouth.Sound(x, self.sampling_rate) \
.to_pitch_ac(time_step=time_step / 1000, voicing_threshold=0.6,pitch_floor=self.f0_min, pitch_ceiling=self.f0_max) \
.selected_array['frequency']
pad_size=(p_len - len(f0) + 1) // 2
if(pad_size>0 or p_len - len(f0) - pad_size>0):
f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
f0,uv = self.interpolate_f0(f0)
return f0
def compute_f0_uv(self,wav,p_len=None):
x = wav
if p_len is None: p_len = x.shape[0]//self.hop_length
else: assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
time_step = self.hop_length / self.sampling_rate * 1000
f0 = parselmouth.Sound(x, self.sampling_rate).to_pitch_ac(
time_step=time_step / 1000, voicing_threshold=0.6,
pitch_floor=self.f0_min, pitch_ceiling=self.f0_max).selected_array['frequency']
pad_size=(p_len - len(f0) + 1) // 2
if(pad_size>0 or p_len - len(f0) - pad_size>0):
f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
f0,uv = self.interpolate_f0(f0)
return f0,uv
class Slicer: # from https://github.com/svc-develop-team/so-vits-svc/
def __init__(self, sr: int, threshold: float = -40., min_length: int = 5000, min_interval: int = 300, hop_size: int = 20, max_sil_kept: int = 5000):
if not min_length >= min_interval >= hop_size:
raise ValueError('The following condition must be satisfied: min_length >= min_interval >= hop_size')
if not max_sil_kept >= hop_size:
raise ValueError('The following condition must be satisfied: max_sil_kept >= hop_size')
min_interval = sr * min_interval / 1000
self.threshold = 10 ** (threshold / 20.)
self.hop_size = round(sr * hop_size / 1000)
self.win_size = min(round(min_interval), 4 * self.hop_size)
self.min_length = round(sr * min_length / 1000 / self.hop_size)
self.min_interval = round(min_interval / self.hop_size)
self.max_sil_kept = round(sr * max_sil_kept / 1000 / self.hop_size)
def _apply_slice(self, waveform, begin, end):
if len(waveform.shape) > 1: return waveform[:, begin * self.hop_size: min(waveform.shape[1], end * self.hop_size)]
else: return waveform[begin * self.hop_size: min(waveform.shape[0], end * self.hop_size)]
def slice(self, waveform):
samples = librosa.to_mono(waveform) if len(waveform.shape) > 1 else waveform
if samples.shape[0] <= self.min_length: return {"0": {"slice": False, "split_time": f"0,{len(waveform)}"}}
rms_list = librosa.feature.rms(y=samples, frame_length=self.win_size, hop_length=self.hop_size).squeeze(0)
sil_tags, silence_start, clip_start = [], None, 0
for i, rms in enumerate(rms_list):
if rms < self.threshold: # Keep looping while frame is silent.
if silence_start is None: # Record start of silent frames.
silence_start = i
continue
if silence_start is None: continue # Keep looping while frame is not silent and silence start has not been recorded.
# Clear recorded silence start if interval is not enough or clip is too short
is_leading_silence = silence_start == 0 and i > self.max_sil_kept
need_slice_middle = i - silence_start >= self.min_interval and i - clip_start >= self.min_length
if not is_leading_silence and not need_slice_middle:
silence_start = None
continue
if i - silence_start <= self.max_sil_kept: # Need slicing. Record the range of silent frames to be removed.
pos = rms_list[silence_start: i + 1].argmin() + silence_start
sil_tags.append((0, pos) if silence_start == 0 else (pos, pos))
clip_start = pos
elif i - silence_start <= self.max_sil_kept * 2:
pos = rms_list[i - self.max_sil_kept: silence_start + self.max_sil_kept + 1].argmin()
pos += i - self.max_sil_kept
pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
if silence_start == 0:
sil_tags.append((0, pos_r))
clip_start = pos_r
else:
sil_tags.append((min(pos_l, pos), max(pos_r, pos)))
clip_start = max(pos_r, pos)
else:
pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
sil_tags.append((0, pos_r) if silence_start == 0 else (pos_l, pos_r))
clip_start = pos_r
silence_start = None
total_frames = rms_list.shape[0]
if silence_start is not None and total_frames - silence_start >= self.min_interval: # Deal with trailing silence.
silence_end = min(total_frames, silence_start + self.max_sil_kept)
pos = rms_list[silence_start: silence_end + 1].argmin() + silence_start
sil_tags.append((pos, total_frames + 1))
if len(sil_tags) == 0: return {"0": {"slice": False, "split_time": f"0,{len(waveform)}"}} # Apply and return slices.
chunks = []
if sil_tags[0][0]:
chunks.append({"slice": False, "split_time": f"0,{min(waveform.shape[0], sil_tags[0][0] * self.hop_size)}"})
for i in range(0, len(sil_tags)):
if i: chunks.append({"slice": False, "split_time": f"{sil_tags[i - 1][1] * self.hop_size},{min(waveform.shape[0], sil_tags[i][0] * self.hop_size)}"})
chunks.append({"slice": True, "split_time": f"{sil_tags[i][0] * self.hop_size},{min(waveform.shape[0], sil_tags[i][1] * self.hop_size)}"})
if sil_tags[-1][1] * self.hop_size < len(waveform):
chunks.append({"slice": False, "split_time": f"{sil_tags[-1][1] * self.hop_size},{len(waveform)}"})
chunk_dict = {}
for i in range(len(chunks)): chunk_dict[str(i)] = chunks[i]
return chunk_dict
# sinc_interp_hann audio resampling
class Resample:
def __init__(self, orig_freq:int=16000, new_freq:int=16000, lowpass_filter_width:int=6, rolloff:float=0.99, beta:Optional[float]=None, dtype:Optional[dtypes]=None):
self.orig_freq, self.new_freq, self.lowpass_filter_width, self.rolloff, self.beta = orig_freq, new_freq, lowpass_filter_width, rolloff, beta
self.gcd = math.gcd(int(self.orig_freq), int(self.new_freq))
self.kernel, self.width = self._get_sinc_resample_kernel(dtype) if self.orig_freq != self.new_freq else (None, None)
def __call__(self, waveform:Tensor) -> Tensor:
if self.orig_freq == self.new_freq: return waveform
return self._apply_sinc_resample_kernel(waveform)
def _apply_sinc_resample_kernel(self, waveform:Tensor):
if not waveform.is_floating_point(): raise TypeError(f"Waveform tensor expected to be of type float, but received {waveform.dtype}.")
orig_freq, new_freq = (int(self.orig_freq) // self.gcd), (int(self.new_freq) // self.gcd)
shape = waveform.shape
waveform = waveform.reshape(-1, shape[-1]) # pack batch
num_wavs, length = waveform.shape
target_length = int(math.ceil(new_freq * length / orig_freq))
waveform = waveform.pad((self.width, self.width + orig_freq))
resampled = waveform[:, None].conv2d(self.kernel, stride=orig_freq)
resampled = resampled.transpose(1, 2).reshape(num_wavs, -1)
resampled = resampled[..., :target_length]
resampled = resampled.reshape(shape[:-1] + resampled.shape[-1:]) # unpack batch
return resampled
def _get_sinc_resample_kernel(self, dtype=None):
orig_freq, new_freq = (int(self.orig_freq) // self.gcd), (int(self.new_freq) // self.gcd)
if self.lowpass_filter_width <= 0: raise ValueError("Low pass filter width should be positive.")
base_freq = min(orig_freq, new_freq)
base_freq *= self.rolloff
width = math.ceil(self.lowpass_filter_width * orig_freq / base_freq)
idx = Tensor.arange(-width, width + orig_freq, dtype=(dtype if dtype is not None else dtypes.float32))[None, None] / orig_freq
t = Tensor.arange(0, -new_freq, -1, dtype=dtype)[:, None, None] / new_freq + idx
t *= base_freq
t = t.clip(-self.lowpass_filter_width, self.lowpass_filter_width)
window = (t * math.pi / self.lowpass_filter_width / 2).cos() ** 2
t *= math.pi
scale = base_freq / orig_freq
kernels = Tensor.where(t == 0, Tensor(1.0, dtype=t.dtype).to(t.device), t.sin() / t)
kernels *= window * scale
if dtype is None: kernels = kernels.cast(dtype=dtypes.float32)
return kernels, width
def sinc_interp_resample(x:Tensor, orig_freq:int=16000, new_freq:int=1600, lowpass_filter_width:int=6, rolloff:float=0.99, beta:Optional[float]=None):
resamp = Resample(orig_freq, new_freq, lowpass_filter_width, rolloff, beta, x.dtype)
return resamp(x)
def cut(audio_path, db_thresh=-30, min_len=5000):
audio, sr = librosa.load(audio_path, sr=None)
slicer = Slicer(sr=sr, threshold=db_thresh, min_length=min_len)
chunks = slicer.slice(audio)
return chunks
def chunks2audio(audio_path, chunks):
chunks = dict(chunks)
audio, sr = load_audiofile(audio_path)
if len(audio.shape) == 2 and audio.shape[1] >= 2:
audio = audio.mean(0).unsqueeze(0)
audio = audio.numpy()[0]
result = []
for k, v in chunks.items():
tag = v["split_time"].split(",")
if tag[0] != tag[1]:
result.append((v["slice"], audio[int(tag[0]):int(tag[1])]))
return result, sr
def load_audiofile(filepath:str, frame_offset:int=0, num_frames:int=-1, channels_first:bool=True):
with soundfile.SoundFile(filepath, "r") as file_:
frames = file_._prepare_read(frame_offset, None, num_frames)
waveform = file_.read(frames, "float32", always_2d=True)
sample_rate = file_.samplerate
waveform = Tensor(waveform)
if channels_first: waveform = waveform.transpose(0, 1)
return waveform, sample_rate
def get_unit_f0(wav:Tensor, tran, hop_length, target_sample, f0_filter=False) -> Tuple[Tensor,Tensor,Tensor]:
f0_predictor = PMF0Predictor(hop_length, sampling_rate=target_sample)
f0, uv = f0_predictor.compute_f0_uv(wav.numpy())
if f0_filter and sum(f0) == 0: raise RuntimeError("No voice detected")
f0 = Tensor(f0.astype(np.float32)).float()
f0 = (f0 * 2 ** (tran / 12)).unsqueeze(0)
uv = Tensor(uv.astype(np.float32)).float().unsqueeze(0)
wav16k = sinc_interp_resample(wav[None,:], target_sample, 16000)[0]
return wav16k.realize(), f0.realize(), uv.realize()
+1 -1
View File
@@ -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_input_info[0][-1]
device = run_onnx_jit.captured.expected_st_vars_dtype_device[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))):
+104
View File
@@ -0,0 +1,104 @@
import traceback
import time
from multiprocessing import Process, Queue
import numpy as np
from tinygrad.nn.state import get_parameters
from tinygrad.nn import optim
from tinygrad.helpers import getenv, trange
from tinygrad.tensor import Tensor
from extra.datasets import fetch_cifar
from extra.models.efficientnet import EfficientNet
class TinyConvNet:
def __init__(self, classes=10):
conv = 3
inter_chan, out_chan = 8, 16 # for speed
self.c1 = Tensor.uniform(inter_chan,3,conv,conv)
self.c2 = Tensor.uniform(out_chan,inter_chan,conv,conv)
self.l1 = Tensor.uniform(out_chan*6*6, classes)
def forward(self, x):
x = x.conv2d(self.c1).relu().max_pool2d()
x = x.conv2d(self.c2).relu().max_pool2d()
x = x.reshape(shape=[x.shape[0], -1])
return x.dot(self.l1)
if __name__ == "__main__":
IMAGENET = getenv("IMAGENET")
classes = 1000 if IMAGENET else 10
TINY = getenv("TINY")
TRANSFER = getenv("TRANSFER")
if TINY:
model = TinyConvNet(classes)
elif TRANSFER:
model = EfficientNet(getenv("NUM", 0), classes, has_se=True)
model.load_from_pretrained()
else:
model = EfficientNet(getenv("NUM", 0), classes, has_se=False)
parameters = get_parameters(model)
print("parameter count", len(parameters))
optimizer = optim.Adam(parameters, lr=0.001)
BS, steps = getenv("BS", 64 if TINY else 16), getenv("STEPS", 2048)
print(f"training with batch size {BS} for {steps} steps")
if IMAGENET:
from extra.datasets.imagenet import fetch_batch
def loader(q):
while 1:
try:
q.put(fetch_batch(BS))
except Exception:
traceback.print_exc()
q = Queue(16)
for i in range(2):
p = Process(target=loader, args=(q,))
p.daemon = True
p.start()
else:
X_train, Y_train, _, _ = fetch_cifar()
X_train = X_train.reshape((-1, 3, 32, 32))
Y_train = Y_train.reshape((-1,))
with Tensor.train():
for i in (t := trange(steps)):
if IMAGENET:
X, Y = q.get(True)
else:
samp = np.random.randint(0, X_train.shape[0], size=(BS))
X, Y = X_train.numpy()[samp], Y_train.numpy()[samp]
st = time.time()
out = model.forward(Tensor(X.astype(np.float32), requires_grad=False))
fp_time = (time.time()-st)*1000.0
y = np.zeros((BS,classes), np.float32)
y[range(y.shape[0]),Y] = -classes
y = Tensor(y, requires_grad=False)
loss = out.log_softmax().mul(y).mean()
optimizer.zero_grad()
st = time.time()
loss.backward()
bp_time = (time.time()-st)*1000.0
st = time.time()
optimizer.step()
opt_time = (time.time()-st)*1000.0
st = time.time()
loss = loss.numpy()
cat = out.argmax(axis=1).numpy()
accuracy = (cat == Y).mean()
finish_time = (time.time()-st)*1000.0
# printing
t.set_description("loss %.2f accuracy %.2f -- %.2f + %.2f + %.2f + %.2f = %.2f" %
(loss, accuracy,
fp_time, bp_time, opt_time, finish_time,
fp_time + bp_time + opt_time + finish_time))
del out, y, loss
+46
View File
@@ -0,0 +1,46 @@
import ast
import numpy as np
from PIL import Image
from tinygrad.tensor import Tensor
from tinygrad.helpers import getenv, fetch
from extra.models.vit import ViT
"""
fn = "gs://vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz"
import tensorflow as tf
with tf.io.gfile.GFile(fn, "rb") as f:
dat = f.read()
with open("cache/"+ fn.rsplit("/", 1)[1], "wb") as g:
g.write(dat)
"""
Tensor.training = False
if getenv("LARGE", 0) == 1:
m = ViT(embed_dim=768, num_heads=12)
else:
# tiny
m = ViT(embed_dim=192, num_heads=3)
m.load_from_pretrained()
# category labels
lbls = ast.literal_eval(fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt").read_text())
#url = "https://upload.wikimedia.org/wikipedia/commons/4/41/Chicken.jpg"
url = "https://repository-images.githubusercontent.com/296744635/39ba6700-082d-11eb-98b8-cb29fb7369c0"
# junk
img = Image.open(fetch(url))
aspect_ratio = img.size[0] / img.size[1]
img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0))))
img = np.array(img)
y0,x0=(np.asarray(img.shape)[:2]-224)//2
img = img[y0:y0+224, x0:x0+224]
img = np.moveaxis(img, [2,0,1], [0,1,2])
img = img.astype(np.float32)[:3].reshape(1,3,224,224)
img /= 255.0
img -= 0.5
img /= 0.5
out = m.forward(Tensor(img))
outnp = out.numpy().ravel()
choice = outnp.argmax()
print(out.shape, choice, outnp[choice], lbls[choice])
+740
View File
@@ -0,0 +1,740 @@
import json, logging, math, re, sys, time, wave, argparse, numpy as np
from phonemizer.phonemize import default_separator, _phonemize
from phonemizer.backend import EspeakBackend
from phonemizer.punctuation import Punctuation
from functools import reduce
from pathlib import Path
from typing import List
from tinygrad import nn, dtypes
from tinygrad.helpers import fetch
from tinygrad.nn.state import torch_load
from tinygrad.tensor import Tensor
from tinygrad.engine.jit import TinyJit
from unidecode import unidecode
LRELU_SLOPE = 0.1
class Synthesizer:
def __init__(self, n_vocab, spec_channels, segment_size, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, n_speakers=0, gin_channels=0, use_sdp=True, emotion_embedding=False, **kwargs):
self.n_vocab, self.spec_channels, self.inter_channels, self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout, self.resblock, self.resblock_kernel_sizes, self.resblock_dilation_sizes, self.upsample_rates, self.upsample_initial_channel, self.upsample_kernel_sizes, self.segment_size, self.n_speakers, self.gin_channels, self.use_sdp = n_vocab, spec_channels, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, segment_size, n_speakers, gin_channels, use_sdp
self.enc_p = TextEncoder(n_vocab, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, emotion_embedding)
self.dec = Generator(inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=gin_channels)
self.enc_q = PosteriorEncoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
self.dp = StochasticDurationPredictor(hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels) if use_sdp else DurationPredictor(hidden_channels, 256, 3, 0.5, gin_channels=gin_channels)
if n_speakers > 1: self.emb_g = nn.Embedding(n_speakers, gin_channels)
def infer(self, x, x_lengths, sid=None, noise_scale=1.0, length_scale=1, noise_scale_w=1., max_len=None, emotion_embedding=None, max_y_length_estimate_scale=None, pad_length=-1):
x, m_p, logs_p, x_mask = self.enc_p.forward(x.realize(), x_lengths.realize(), emotion_embedding.realize() if emotion_embedding is not None else emotion_embedding)
g = self.emb_g(sid.reshape(1, 1)).squeeze(1).unsqueeze(-1) if self.n_speakers > 0 else None
logw = self.dp.forward(x, x_mask.realize(), g=g.realize(), reverse=self.use_sdp, noise_scale=noise_scale_w if self.use_sdp else 1.0)
w_ceil = Tensor.ceil(logw.exp() * x_mask * length_scale)
y_lengths = Tensor.maximum(w_ceil.sum([1, 2]), 1).cast(dtypes.int64)
return self.generate(g, logs_p, m_p, max_len, max_y_length_estimate_scale, noise_scale, w_ceil, x, x_mask, y_lengths, pad_length)
def generate(self, g, logs_p, m_p, max_len, max_y_length_estimate_scale, noise_scale, w_ceil, x, x_mask, y_lengths, pad_length):
max_y_length = y_lengths.max().item() if max_y_length_estimate_scale is None else max(15, x.shape[-1]) * max_y_length_estimate_scale
y_mask = sequence_mask(y_lengths, max_y_length).unsqueeze(1).cast(x_mask.dtype)
attn_mask = x_mask.unsqueeze(2) * y_mask.unsqueeze(-1)
attn = generate_path(w_ceil, attn_mask)
m_p_2 = attn.squeeze(1).matmul(m_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']
logs_p_2 = attn.squeeze(1).matmul(logs_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']
z_p = m_p_2 + Tensor.randn(*m_p_2.shape, dtype=m_p_2.dtype) * logs_p_2.exp() * noise_scale
row_len = y_mask.shape[2]
if pad_length > -1:
# Pad flow forward inputs to enable JIT
assert pad_length > row_len, "pad length is too small"
y_mask = y_mask.pad(((0, 0), (0, 0), (0, pad_length - row_len))).cast(z_p.dtype)
# New y_mask tensor to remove sts mask
y_mask = Tensor(y_mask.numpy(), device=y_mask.device, dtype=y_mask.dtype, requires_grad=y_mask.requires_grad)
z_p = z_p.squeeze(0).pad(((0, 0), (0, pad_length - z_p.shape[2])), value=1).unsqueeze(0)
z = self.flow.forward(z_p.realize(), y_mask.realize(), g=g.realize(), reverse=True)
result_length = reduce(lambda x, y: x * y, self.dec.upsample_rates, row_len)
o = self.dec.forward((z * y_mask)[:, :, :max_len], g=g)[:, :, :result_length]
if max_y_length_estimate_scale is not None:
length_scaler = o.shape[-1] / max_y_length
o.realize()
real_max_y_length = y_lengths.max().numpy()
if real_max_y_length > max_y_length:
logging.warning(f"Underestimated max length by {(((real_max_y_length / max_y_length) * 100) - 100):.2f}%, recomputing inference without estimate...")
return self.generate(g, logs_p, m_p, max_len, None, noise_scale, w_ceil, x, x_mask, y_lengths)
if real_max_y_length < max_y_length:
overestimation = ((max_y_length / real_max_y_length) * 100) - 100
logging.info(f"Overestimated max length by {overestimation:.2f}%")
if overestimation > 10: logging.warning("Warning: max length overestimated by more than 10%")
o = o[:, :, :(real_max_y_length * length_scaler).astype(np.int32)]
return o
class StochasticDurationPredictor:
def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, n_flows=4, gin_channels=0):
filter_channels = in_channels # it needs to be removed from future version.
self.in_channels, self.filter_channels, self.kernel_size, self.p_dropout, self.n_flows, self.gin_channels = in_channels, filter_channels, kernel_size, p_dropout, n_flows, gin_channels
self.log_flow, self.flows = Log(), [ElementwiseAffine(2)]
for _ in range(n_flows):
self.flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3))
self.flows.append(Flip())
self.post_pre, self.post_proj = nn.Conv1d(1, filter_channels, 1), nn.Conv1d(filter_channels, filter_channels, 1)
self.post_convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
self.post_flows = [ElementwiseAffine(2)]
for _ in range(4):
self.post_flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3))
self.post_flows.append(Flip())
self.pre, self.proj = nn.Conv1d(in_channels, filter_channels, 1), nn.Conv1d(filter_channels, filter_channels, 1)
self.convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
@TinyJit
def forward(self, x: Tensor, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
x = self.pre(x.detach())
if g is not None: x = x + self.cond(g.detach())
x = self.convs.forward(x, x_mask)
x = self.proj(x) * x_mask
if not reverse:
flows = self.flows
assert w is not None
log_det_tot_q = 0
h_w = self.post_proj(self.post_convs.forward(self.post_pre(w), x_mask)) * x_mask
e_q = Tensor.randn(w.size(0), 2, w.size(2), dtype=x.dtype).to(device=x.device) * x_mask
z_q = e_q
for flow in self.post_flows:
z_q, log_det_q = flow.forward(z_q, x_mask, g=(x + h_w))
log_det_tot_q += log_det_q
z_u, z1 = z_q.split([1, 1], 1)
u = z_u.sigmoid() * x_mask
z0 = (w - u) * x_mask
log_det_tot_q += Tensor.sum((z_u.logsigmoid() + (-z_u).logsigmoid()) * x_mask, [1,2])
log_q = Tensor.sum(-0.5 * (math.log(2*math.pi) + (e_q**2)) * x_mask, [1,2]) - log_det_tot_q
log_det_tot = 0
z0, log_det = self.log_flow.forward(z0, x_mask)
log_det_tot += log_det
z = z0.cat(z1, 1)
for flow in flows:
z, log_det = flow.forward(z, x_mask, g=x, reverse=reverse)
log_det_tot = log_det_tot + log_det
nll = Tensor.sum(0.5 * (math.log(2*math.pi) + (z**2)) * x_mask, [1,2]) - log_det_tot
return (nll + log_q).realize() # [b]
flows = list(reversed(self.flows))
flows = flows[:-2] + [flows[-1]] # remove a useless vflow
z = Tensor.randn(x.shape[0], 2, x.shape[2], dtype=x.dtype).to(device=x.device) * noise_scale
for flow in flows: z = flow.forward(z, x_mask, g=x, reverse=reverse)
z0, z1 = z.split([1, 1], 1)
return z0.realize()
class DurationPredictor:
def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0):
self.in_channels, self.filter_channels, self.kernel_size, self.p_dropout, self.gin_channels = in_channels, filter_channels, kernel_size, p_dropout, gin_channels
self.conv_1, self.norm_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size//2), LayerNorm(filter_channels)
self.conv_2, self.norm_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size//2), LayerNorm(filter_channels)
self.proj = nn.Conv1d(filter_channels, 1, 1)
if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, in_channels, 1)
def forward(self, x: Tensor, x_mask, g=None):
x = x.detach()
if g is not None: x = x + self.cond(g.detach())
x = self.conv_1(x * x_mask).relu()
x = self.norm_1(x).dropout(self.p_dropout)
x = self.conv_2(x * x_mask).relu(x)
x = self.norm_2(x).dropout(self.p_dropout)
return self.proj(x * x_mask) * x_mask
class TextEncoder:
def __init__(self, n_vocab, out_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, emotion_embedding):
self.n_vocab, self.out_channels, self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout = n_vocab, out_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout
if n_vocab!=0:self.emb = nn.Embedding(n_vocab, hidden_channels)
if emotion_embedding: self.emo_proj = nn.Linear(1024, hidden_channels)
self.encoder = Encoder(hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout)
self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
@TinyJit
def forward(self, x: Tensor, x_lengths: Tensor, emotion_embedding=None):
if self.n_vocab!=0: x = (self.emb(x) * math.sqrt(self.hidden_channels))
if emotion_embedding: x = x + self.emo_proj(emotion_embedding).unsqueeze(1)
x = x.transpose(1, -1) # [b, t, h] -transpose-> [b, h, t]
x_mask = sequence_mask(x_lengths, x.shape[2]).unsqueeze(1).cast(x.dtype)
x = self.encoder.forward(x * x_mask, x_mask)
m, logs = (self.proj(x) * x_mask).split(self.out_channels, dim=1)
return x.realize(), m.realize(), logs.realize(), x_mask.realize()
class ResidualCouplingBlock:
def __init__(self, channels, hidden_channels, kernel_size, dilation_rate, n_layers, n_flows=4, gin_channels=0):
self.channels, self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.n_flows, self.gin_channels = channels, hidden_channels, kernel_size, dilation_rate, n_layers, n_flows, gin_channels
self.flows = []
for _ in range(n_flows):
self.flows.append(ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels, mean_only=True))
self.flows.append(Flip())
@TinyJit
def forward(self, x, x_mask, g=None, reverse=False):
for flow in reversed(self.flows) if reverse else self.flows: x = flow.forward(x, x_mask, g=g, reverse=reverse)
return x.realize()
class PosteriorEncoder:
def __init__(self, in_channels, out_channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0):
self.in_channels, self.out_channels, self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.gin_channels = in_channels, out_channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels
self.pre, self.proj = nn.Conv1d(in_channels, hidden_channels, 1), nn.Conv1d(hidden_channels, out_channels * 2, 1)
self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels)
def forward(self, x, x_lengths, g=None):
x_mask = sequence_mask(x_lengths, x.size(2)).unsqueeze(1).cast(x.dtype)
stats = self.proj(self.enc.forward(self.pre(x) * x_mask, x_mask, g=g)) * x_mask
m, logs = stats.split(self.out_channels, dim=1)
z = (m + Tensor.randn(m.shape, m.dtype) * logs.exp()) * x_mask
return z, m, logs, x_mask
class Generator:
def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=0):
self.num_kernels, self.num_upsamples = len(resblock_kernel_sizes), len(upsample_rates)
self.conv_pre = nn.Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)
resblock = ResBlock1 if resblock == '1' else ResBlock2
self.ups = [nn.ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)), k, u, padding=(k-u)//2) for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes))]
self.resblocks = []
self.upsample_rates = upsample_rates
for i in range(len(self.ups)):
ch = upsample_initial_channel // (2 ** (i + 1))
for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
self.resblocks.append(resblock(ch, k, d))
self.conv_post = nn.Conv1d(ch, 1, 7, 1, padding=3, bias=False)
if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
@TinyJit
def forward(self, x: Tensor, g=None):
x = self.conv_pre(x)
if g is not None: x = x + self.cond(g)
for i in range(self.num_upsamples):
x = self.ups[i](x.leaky_relu(LRELU_SLOPE))
xs = sum(self.resblocks[i * self.num_kernels + j].forward(x) for j in range(self.num_kernels))
x = (xs / self.num_kernels).realize()
res = self.conv_post(x.leaky_relu()).tanh().realize()
return res
class LayerNorm(nn.LayerNorm):
def __init__(self, channels, eps=1e-5): super().__init__(channels, eps, elementwise_affine=True)
def forward(self, x: Tensor): return self.__call__(x.transpose(1, -1)).transpose(1, -1)
class WN:
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
assert (kernel_size % 2 == 1)
self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.gin_channels, self.p_dropout = hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels, p_dropout
self.in_layers, self.res_skip_layers = [], []
if gin_channels != 0: self.cond_layer = nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers, 1)
for i in range(n_layers):
dilation = dilation_rate ** i
self.in_layers.append(nn.Conv1d(hidden_channels, 2 * hidden_channels, kernel_size, dilation=dilation, padding=int((kernel_size * dilation - dilation) / 2)))
self.res_skip_layers.append(nn.Conv1d(hidden_channels, 2 * hidden_channels if i < n_layers - 1 else hidden_channels, 1))
def forward(self, x, x_mask, g=None, **kwargs):
output = Tensor.zeros_like(x)
if g is not None: g = self.cond_layer(g)
for i in range(self.n_layers):
x_in = self.in_layers[i](x)
if g is not None:
cond_offset = i * 2 * self.hidden_channels
g_l = g[:, cond_offset:cond_offset + 2 * self.hidden_channels, :]
else:
g_l = Tensor.zeros_like(x_in)
acts = fused_add_tanh_sigmoid_multiply(x_in, g_l, self.hidden_channels)
res_skip_acts = self.res_skip_layers[i](acts)
if i < self.n_layers - 1:
x = (x + res_skip_acts[:, :self.hidden_channels, :]) * x_mask
output = output + res_skip_acts[:, self.hidden_channels:, :]
else:
output = output + res_skip_acts
return output * x_mask
class ResBlock1:
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
self.convs1 = [nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[i], padding=get_padding(kernel_size, dilation[i])) for i in range(3)]
self.convs2 = [nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)) for _ in range(3)]
def forward(self, x: Tensor, x_mask=None):
for c1, c2 in zip(self.convs1, self.convs2):
xt = x.leaky_relu(LRELU_SLOPE)
xt = c1(xt if x_mask is None else xt * x_mask).leaky_relu(LRELU_SLOPE)
x = c2(xt if x_mask is None else xt * x_mask) + x
return x if x_mask is None else x * x_mask
class ResBlock2:
def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
self.convs = [nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[i], padding=get_padding(kernel_size, dilation[i])) for i in range(2)]
def forward(self, x, x_mask=None):
for c in self.convs:
xt = x.leaky_relu(LRELU_SLOPE)
xt = c(xt if x_mask is None else xt * x_mask)
x = xt + x
return x if x_mask is None else x * x_mask
class DDSConv: # Dilated and Depth-Separable Convolution
def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):
self.channels, self.kernel_size, self.n_layers, self.p_dropout = channels, kernel_size, n_layers, p_dropout
self.convs_sep, self.convs_1x1, self.norms_1, self.norms_2 = [], [], [], []
for i in range(n_layers):
dilation = kernel_size ** i
padding = (kernel_size * dilation - dilation) // 2
self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size, groups=channels, dilation=dilation, padding=padding))
self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
self.norms_1.append(LayerNorm(channels))
self.norms_2.append(LayerNorm(channels))
def forward(self, x, x_mask, g=None):
if g is not None: x = x + g
for i in range(self.n_layers):
y = self.convs_sep[i](x * x_mask)
y = self.norms_1[i].forward(y).gelu()
y = self.convs_1x1[i](y)
y = self.norms_2[i].forward(y).gelu()
x = x + y.dropout(self.p_dropout)
return x * x_mask
class ConvFlow:
def __init__(self, in_channels, filter_channels, kernel_size, n_layers, num_bins=10, tail_bound=5.0):
self.in_channels, self.filter_channels, self.kernel_size, self.n_layers, self.num_bins, self.tail_bound = in_channels, filter_channels, kernel_size, n_layers, num_bins, tail_bound
self.half_channels = in_channels // 2
self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.)
self.proj = nn.Conv1d(filter_channels, self.half_channels * (num_bins * 3 - 1), 1)
def forward(self, x, x_mask, g=None, reverse=False):
x0, x1 = x.split([self.half_channels] * 2, 1)
h = self.proj(self.convs.forward(self.pre(x0), x_mask, g=g)) * x_mask
b, c, t = x0.shape
h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) # [b, cx?, t] -> [b, c, t, ?]
un_normalized_widths = h[..., :self.num_bins] / math.sqrt(self.filter_channels)
un_normalized_heights = h[..., self.num_bins:2*self.num_bins] / math.sqrt(self.filter_channels)
un_normalized_derivatives = h[..., 2 * self.num_bins:]
x1, log_abs_det = piecewise_rational_quadratic_transform(x1, un_normalized_widths, un_normalized_heights, un_normalized_derivatives, inverse=reverse, tails='linear', tail_bound=self.tail_bound)
x = x0.cat(x1, dim=1) * x_mask
return x if reverse else (x, Tensor.sum(log_abs_det * x_mask, [1,2]))
class ResidualCouplingLayer:
def __init__(self, channels, hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=0, gin_channels=0, mean_only=False):
assert channels % 2 == 0, "channels should be divisible by 2"
self.channels, self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.mean_only = channels, hidden_channels, kernel_size, dilation_rate, n_layers, mean_only
self.half_channels = channels // 2
self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
def forward(self, x, x_mask, g=None, reverse=False):
x0, x1 = x.split([self.half_channels] * 2, 1)
stats = self.post(self.enc.forward(self.pre(x0) * x_mask, x_mask, g=g)) * x_mask
if not self.mean_only:
m, logs = stats.split([self.half_channels] * 2, 1)
else:
m = stats
logs = Tensor.zeros_like(m)
if not reverse: return x0.cat((m + x1 * logs.exp() * x_mask), dim=1)
return x0.cat(((x1 - m) * (-logs).exp() * x_mask), dim=1)
class Log:
def forward(self, x : Tensor, x_mask, reverse=False):
if not reverse:
y = x.maximum(1e-5).log() * x_mask
return y, (-y).sum([1, 2])
return x.exp() * x_mask
class Flip:
def forward(self, x: Tensor, *args, reverse=False, **kwargs):
return x.flip([1]) if reverse else (x.flip([1]), Tensor.zeros(x.shape[0], dtype=x.dtype).to(device=x.device))
class ElementwiseAffine:
def __init__(self, channels): self.m, self.logs = Tensor.zeros(channels, 1), Tensor.zeros(channels, 1)
def forward(self, x, x_mask, reverse=False, **kwargs): # x if reverse else y, logdet
return (x - self.m) * Tensor.exp(-self.logs) * x_mask if reverse \
else ((self.m + Tensor.exp(self.logs) * x) * x_mask, Tensor.sum(self.logs * x_mask, [1, 2]))
class MultiHeadAttention:
def __init__(self, channels, out_channels, n_heads, p_dropout=0., window_size=None, heads_share=True, block_length=None, proximal_bias=False, proximal_init=False):
assert channels % n_heads == 0
self.channels, self.out_channels, self.n_heads, self.p_dropout, self.window_size, self.heads_share, self.block_length, self.proximal_bias, self.proximal_init = channels, out_channels, n_heads, p_dropout, window_size, heads_share, block_length, proximal_bias, proximal_init
self.attn, self.k_channels = None, channels // n_heads
self.conv_q, self.conv_k, self.conv_v = [nn.Conv1d(channels, channels, 1) for _ in range(3)]
self.conv_o = nn.Conv1d(channels, out_channels, 1)
if window_size is not None: self.emb_rel_k, self.emb_rel_v = [Tensor.randn(1 if heads_share else n_heads, window_size * 2 + 1, self.k_channels) * (self.k_channels ** -0.5) for _ in range(2)]
def forward(self, x, c, attn_mask=None):
q, k, v = self.conv_q(x), self.conv_k(c), self.conv_v(c)
x, self.attn = self.attention(q, k, v, mask=attn_mask)
return self.conv_o(x)
def attention(self, query: Tensor, key: Tensor, value: Tensor, mask=None):# reshape [b, d, t] -> [b, n_h, t, d_k]
b, d, t_s, t_t = key.shape[0], key.shape[1], key.shape[2], query.shape[2]
query = query.reshape(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
key = key.reshape(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
value = value.reshape(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
scores = (query / math.sqrt(self.k_channels)) @ key.transpose(-2, -1)
if self.window_size is not None:
assert t_s == t_t, "Relative attention is only available for self-attention."
key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
rel_logits = self._matmul_with_relative_keys(query / math.sqrt(self.k_channels), key_relative_embeddings)
scores = scores + self._relative_position_to_absolute_position(rel_logits)
if mask is not None:
scores = Tensor.where(mask, scores, -1e4)
if self.block_length is not None:
assert t_s == t_t, "Local attention is only available for self-attention."
scores = Tensor.where(Tensor.ones_like(scores).triu(-self.block_length).tril(self.block_length), scores, -1e4)
p_attn = scores.softmax(axis=-1) # [b, n_h, t_t, t_s]
output = p_attn.matmul(value)
if self.window_size is not None:
relative_weights = self._absolute_position_to_relative_position(p_attn)
value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s)
output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings)
output = output.transpose(2, 3).contiguous().reshape(b, d, t_t) # [b, n_h, t_t, d_k] -> [b, d, t_t]
return output, p_attn
def _matmul_with_relative_values(self, x, y): return x.matmul(y.unsqueeze(0)) # x: [b, h, l, m], y: [h or 1, m, d], ret: [b, h, l, d]
def _matmul_with_relative_keys(self, x, y): return x.matmul(y.unsqueeze(0).transpose(-2, -1)) # x: [b, h, l, d], y: [h or 1, m, d], re, : [b, h, l, m]
def _get_relative_embeddings(self, relative_embeddings, length):
pad_length, slice_start_position = max(length - (self.window_size + 1), 0), max((self.window_size + 1) - length, 0)
padded_relative_embeddings = relative_embeddings if pad_length <= 0\
else relative_embeddings.pad(convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]))
return padded_relative_embeddings[:, slice_start_position:(slice_start_position + 2 * length - 1)] #used_relative_embeddings
def _relative_position_to_absolute_position(self, x: Tensor): # x: [b, h, l, 2*l-1] -> [b, h, l, l]
batch, heads, length, _ = x.shape
x = x.pad(convert_pad_shape([[0,0],[0,0],[0,0],[0,1]]))
x_flat = x.reshape([batch, heads, length * 2 * length]).pad(convert_pad_shape([[0,0],[0,0],[0,length-1]]))
return x_flat.reshape([batch, heads, length+1, 2*length-1])[:, :, :length, length-1:]
def _absolute_position_to_relative_position(self, x: Tensor): # x: [b, h, l, l] -> [b, h, l, 2*l-1]
batch, heads, length, _ = x.shape
x = x.pad(convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length-1]]))
x_flat = x.reshape([batch, heads, length**2 + length*(length -1)]).pad(convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
return x_flat.reshape([batch, heads, length, 2*length])[:,:,:,1:]
class FFN:
def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., activation=None, causal=False):
self.in_channels, self.out_channels, self.filter_channels, self.kernel_size, self.p_dropout, self.activation, self.causal = in_channels, out_channels, filter_channels, kernel_size, p_dropout, activation, causal
self.padding = self._causal_padding if causal else self._same_padding
self.conv_1, self.conv_2 = nn.Conv1d(in_channels, filter_channels, kernel_size), nn.Conv1d(filter_channels, out_channels, kernel_size)
def forward(self, x, x_mask):
x = self.conv_1(self.padding(x * x_mask))
x = x * (1.702 * x).sigmoid() if self.activation == "gelu" else x.relu()
return self.conv_2(self.padding(x.dropout(self.p_dropout) * x_mask)) * x_mask
def _causal_padding(self, x):return x if self.kernel_size == 1 else x.pad(convert_pad_shape([[0, 0], [0, 0], [self.kernel_size - 1, 0]]))
def _same_padding(self, x): return x if self.kernel_size == 1 else x.pad(convert_pad_shape([[0, 0], [0, 0], [(self.kernel_size - 1) // 2, self.kernel_size // 2]]))
class Encoder:
def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=4, **kwargs):
self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout, self.window_size = hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, window_size
self.attn_layers, self.norm_layers_1, self.ffn_layers, self.norm_layers_2 = [], [], [], []
for _ in range(n_layers):
self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, window_size=window_size))
self.norm_layers_1.append(LayerNorm(hidden_channels))
self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout))
self.norm_layers_2.append(LayerNorm(hidden_channels))
def forward(self, x, x_mask):
attn_mask, x = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1), x * x_mask
for i in range(self.n_layers):
y = self.attn_layers[i].forward(x, x, attn_mask).dropout(self.p_dropout)
x = self.norm_layers_1[i].forward(x + y)
y = self.ffn_layers[i].forward(x, x_mask).dropout(self.p_dropout)
x = self.norm_layers_2[i].forward(x + y)
return x * x_mask
DEFAULT_MIN_BIN_WIDTH, DEFAULT_MIN_BIN_HEIGHT, DEFAULT_MIN_DERIVATIVE = 1e-3, 1e-3, 1e-3
def piecewise_rational_quadratic_transform(inputs, un_normalized_widths, un_normalized_heights, un_normalized_derivatives, inverse=False, tails=None, tail_bound=1., min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE):
if tails is None: spline_fn, spline_kwargs = rational_quadratic_spline, {}
else: spline_fn, spline_kwargs = unconstrained_rational_quadratic_spline, {'tails': tails, 'tail_bound': tail_bound}
return spline_fn(inputs=inputs, un_normalized_widths=un_normalized_widths, un_normalized_heights=un_normalized_heights, un_normalized_derivatives=un_normalized_derivatives, inverse=inverse, min_bin_width=min_bin_width, min_bin_height=min_bin_height, min_derivative=min_derivative, **spline_kwargs)
def unconstrained_rational_quadratic_spline(inputs, un_normalized_widths, un_normalized_heights, un_normalized_derivatives, inverse=False, tails='linear', tail_bound=1., min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE):
if not tails == 'linear': raise RuntimeError('{} tails are not implemented.'.format(tails))
constant = np.log(np.exp(1 - min_derivative) - 1).item()
un_normalized_derivatives = cat_lr(un_normalized_derivatives, constant, constant)
output, log_abs_det = rational_quadratic_spline(inputs=inputs.squeeze(dim=0).squeeze(dim=0), unnormalized_widths=un_normalized_widths.squeeze(dim=0).squeeze(dim=0), unnormalized_heights=un_normalized_heights.squeeze(dim=0).squeeze(dim=0), unnormalized_derivatives=un_normalized_derivatives.squeeze(dim=0).squeeze(dim=0), inverse=inverse, left=-tail_bound, right=tail_bound, bottom=-tail_bound, top=tail_bound, min_bin_width=min_bin_width, min_bin_height=min_bin_height, min_derivative=min_derivative)
return output.unsqueeze(dim=0).unsqueeze(dim=0), log_abs_det.unsqueeze(dim=0).unsqueeze(dim=0)
def rational_quadratic_spline(inputs: Tensor, unnormalized_widths: Tensor, unnormalized_heights: Tensor, unnormalized_derivatives: Tensor, inverse=False, left=0., right=1., bottom=0., top=1., min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE):
num_bins = unnormalized_widths.shape[-1]
if min_bin_width * num_bins > 1.0: raise ValueError('Minimal bin width too large for the number of bins')
if min_bin_height * num_bins > 1.0: raise ValueError('Minimal bin height too large for the number of bins')
widths = min_bin_width + (1 - min_bin_width * num_bins) * unnormalized_widths.softmax(axis=-1)
cum_widths = cat_lr(((right - left) * widths[..., :-1].cumsum(axis=1) + left), left, right + 1e-6 if not inverse else right)
widths = cum_widths[..., 1:] - cum_widths[..., :-1]
derivatives = min_derivative + (unnormalized_derivatives.exp()+1).log()
heights = min_bin_height + (1 - min_bin_height * num_bins) * unnormalized_heights.softmax(axis=-1)
cum_heights = cat_lr(((top - bottom) * heights[..., :-1].cumsum(axis=1) + bottom), bottom, top + 1e-6 if inverse else top)
heights = cum_heights[..., 1:] - cum_heights[..., :-1]
bin_idx = ((inputs[..., None] >= (cum_heights if inverse else cum_widths)).sum(axis=-1) - 1)[..., None]
input_cum_widths = gather(cum_widths, bin_idx, axis=-1)[..., 0]
input_bin_widths = gather(widths, bin_idx, axis=-1)[..., 0]
input_cum_heights = gather(cum_heights, bin_idx, axis=-1)[..., 0]
input_delta = gather(heights / widths, bin_idx, axis=-1)[..., 0]
input_derivatives = gather(derivatives, bin_idx, axis=-1)[..., 0]
input_derivatives_plus_one = gather(derivatives[..., 1:], bin_idx, axis=-1)[..., 0]
input_heights = gather(heights, bin_idx, axis=-1)[..., 0]
if inverse:
a = ((inputs - input_cum_heights) * (input_derivatives + input_derivatives_plus_one - 2 * input_delta) + input_heights * (input_delta - input_derivatives))
b = (input_heights * input_derivatives - (inputs - input_cum_heights) * (input_derivatives + input_derivatives_plus_one - 2 * input_delta))
c = - input_delta * (inputs - input_cum_heights)
discriminant = b.square() - 4 * a * c
# assert (discriminant.numpy() >= 0).all()
root = (2 * c) / (-b - discriminant.sqrt())
theta_one_minus_theta = root * (1 - root)
denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta) * theta_one_minus_theta)
derivative_numerator = input_delta.square() * (input_derivatives_plus_one * root.square() + 2 * input_delta * theta_one_minus_theta + input_derivatives * (1 - root).square())
return root * input_bin_widths + input_cum_widths, -(derivative_numerator.log() - 2 * denominator.log())
theta = (inputs - input_cum_widths) / input_bin_widths
theta_one_minus_theta = theta * (1 - theta)
numerator = input_heights * (input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta)
denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta) * theta_one_minus_theta)
derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2) + 2 * input_delta * theta_one_minus_theta + input_derivatives * (1 - theta).pow(2))
return input_cum_heights + numerator / denominator, derivative_numerator.log() - 2 * denominator.log()
def sequence_mask(length: Tensor, max_length): return Tensor.arange(max_length, dtype=length.dtype, device=length.device).unsqueeze(0) < length.unsqueeze(1)
def generate_path(duration: Tensor, mask: Tensor): # duration: [b, 1, t_x], mask: [b, 1, t_y, t_x]
b, _, t_y, t_x = mask.shape
path = sequence_mask(duration.cumsum(axis=2).reshape(b * t_x), t_y).cast(mask.dtype).reshape(b, t_x, t_y)
path = path - path.pad(convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
return path.unsqueeze(1).transpose(2, 3) * mask
def fused_add_tanh_sigmoid_multiply(input_a: Tensor, input_b: Tensor, n_channels: int):
n_channels_int, in_act = n_channels, input_a + input_b
t_act, s_act = in_act[:, :n_channels_int, :].tanh(), in_act[:, n_channels_int:, :].sigmoid()
return t_act * s_act
def cat_lr(t, left, right): return Tensor.full(get_shape(t), left).cat(t, dim=-1).cat(Tensor.full(get_shape(t), right), dim=-1)
def get_shape(tensor):
(shape := list(tensor.shape))[-1] = 1
return tuple(shape)
def convert_pad_shape(pad_shape): return tuple(tuple(x) for x in pad_shape)
def get_padding(kernel_size, dilation=1): return int((kernel_size*dilation - dilation)/2)
def gather(x, indices, axis):
indices = (indices < 0).where(indices + x.shape[axis], indices).transpose(0, axis)
permute_args = list(range(x.ndim))
permute_args[0], permute_args[axis] = permute_args[axis], permute_args[0]
permute_args.append(permute_args.pop(0))
x = x.permute(*permute_args)
reshape_arg = [1] * x.ndim + [x.shape[-1]]
return ((indices.unsqueeze(indices.ndim).expand(*indices.shape, x.shape[-1]) ==
Tensor.arange(x.shape[-1]).reshape(*reshape_arg).expand(*indices.shape, x.shape[-1])) * x).sum(indices.ndim).transpose(0, axis)
def norm_except_dim(v, dim):
if dim == -1: return np.linalg.norm(v)
if dim == 0:
(output_shape := [1] * v.ndim)[0] = v.shape[0]
return np.linalg.norm(v.reshape(v.shape[0], -1), axis=1).reshape(output_shape)
if dim == v.ndim - 1:
(output_shape := [1] * v.ndim)[-1] = v.shape[-1]
return np.linalg.norm(v.reshape(-1, v.shape[-1]), axis=0).reshape(output_shape)
transposed_v = np.transpose(v, (dim,) + tuple(i for i in range(v.ndim) if i != dim))
return np.transpose(norm_except_dim(transposed_v, 0), (dim,) + tuple(i for i in range(v.ndim) if i != dim))
def weight_norm(v: Tensor, g: Tensor, dim):
v, g = v.numpy(), g.numpy()
return Tensor(v * (g / norm_except_dim(v, dim)))
# HPARAMS LOADING
def get_hparams_from_file(path):
with open(path, "r") as f:
data = f.read()
return HParams(**json.loads(data))
class HParams:
def __init__(self, **kwargs):
for k, v in kwargs.items(): self[k] = v if type(v) != dict else HParams(**v)
def keys(self): return self.__dict__.keys()
def items(self): return self.__dict__.items()
def values(self): return self.__dict__.values()
def __len__(self): return len(self.__dict__)
def __getitem__(self, key): return getattr(self, key)
def __setitem__(self, key, value): return setattr(self, key, value)
def __contains__(self, key): return key in self.__dict__
def __repr__(self): return self.__dict__.__repr__()
# MODEL LOADING
def load_model(symbols, hps, model) -> Synthesizer:
net_g = Synthesizer(len(symbols), hps.data.filter_length // 2 + 1, hps.train.segment_size // hps.data.hop_length, n_speakers = hps.data.n_speakers, **hps.model)
_ = load_checkpoint(fetch(model[1]), net_g, None)
return net_g
def load_checkpoint(checkpoint_path, model: Synthesizer, optimizer=None, skip_list=[]):
assert Path(checkpoint_path).is_file()
start_time = time.time()
checkpoint_dict = torch_load(checkpoint_path)
iteration, learning_rate = checkpoint_dict['iteration'], checkpoint_dict['learning_rate']
if optimizer: optimizer.load_state_dict(checkpoint_dict['optimizer'])
saved_state_dict = checkpoint_dict['model']
weight_g, weight_v, parent = None, None, None
for key, v in saved_state_dict.items():
if any(layer in key for layer in skip_list): continue
try:
obj, skip = model, False
for k in key.split('.'):
if k.isnumeric(): obj = obj[int(k)]
elif isinstance(obj, dict): obj = obj[k]
else:
if isinstance(obj, (LayerNorm, nn.LayerNorm)) and k in ["gamma", "beta"]:
k = "weight" if k == "gamma" else "bias"
elif k in ["weight_g", "weight_v"]:
parent, skip = obj, True
if k == "weight_g": weight_g = v
else: weight_v = v
if not skip: obj = getattr(obj, k)
if weight_g is not None and weight_v is not None:
setattr(obj, "weight_g", weight_g.numpy())
setattr(obj, "weight_v", weight_v.numpy())
obj, v = getattr(parent, "weight"), weight_norm(weight_v, weight_g, 0)
weight_g, weight_v, parent, skip = None, None, None, False
if not skip and obj.shape == v.shape: obj.assign(v.to(obj.device))
elif not skip: logging.error(f"MISMATCH SHAPE IN {key}, {obj.shape} {v.shape}")
except Exception as e: raise e
logging.info(f"Loaded checkpoint '{checkpoint_path}' (iteration {iteration}) in {time.time() - start_time:.4f}s")
return model, optimizer, learning_rate, iteration
# Used for cleaning input text and mapping to symbols
class TextMapper: # Based on https://github.com/keithito/tacotron
def __init__(self, symbols, apply_cleaners=True):
self.apply_cleaners, self.symbols, self._inflect = apply_cleaners, symbols, None
self._symbol_to_id, _id_to_symbol = {s: i for i, s in enumerate(symbols)}, {i: s for i, s in enumerate(symbols)}
self._whitespace_re, self._abbreviations = re.compile(r'\s+'), [(re.compile('\\b%s\\.' % x[0], re.IGNORECASE), x[1]) for x in [('mrs', 'misess'), ('mr', 'mister'), ('dr', 'doctor'), ('st', 'saint'), ('co', 'company'), ('jr', 'junior'), ('maj', 'major'), ('gen', 'general'), ('drs', 'doctors'), ('rev', 'reverend'), ('lt', 'lieutenant'), ('hon', 'honorable'), ('sgt', 'sergeant'), ('capt', 'captain'), ('esq', 'esquire'), ('ltd', 'limited'), ('col', 'colonel'), ('ft', 'fort'), ]]
self.phonemizer = EspeakBackend(
language="en-us", punctuation_marks=Punctuation.default_marks(), preserve_punctuation=True, with_stress=True,
)
def text_to_sequence(self, text, cleaner_names):
if self.apply_cleaners:
for name in cleaner_names:
cleaner = getattr(self, name)
if not cleaner: raise ModuleNotFoundError('Unknown cleaner: %s' % name)
text = cleaner(text)
else: text = text.strip()
return [self._symbol_to_id[symbol] for symbol in text]
def get_text(self, text, add_blank=False, cleaners=('english_cleaners2',)):
text_norm = self.text_to_sequence(text, cleaners)
return Tensor(self.intersperse(text_norm, 0) if add_blank else text_norm, dtype=dtypes.int64)
def intersperse(self, lst, item):
(result := [item] * (len(lst) * 2 + 1))[1::2] = lst
return result
def phonemize(self, text, strip=True): return _phonemize(self.phonemizer, text, default_separator, strip, 1, False, False)
def filter_oov(self, text): return "".join(list(filter(lambda x: x in self._symbol_to_id, text)))
def base_english_cleaners(self, text): return self.collapse_whitespace(self.phonemize(self.expand_abbreviations(unidecode(text.lower()))))
def english_cleaners2(self, text): return self.base_english_cleaners(text)
def transliteration_cleaners(self, text): return self.collapse_whitespace(unidecode(text.lower()))
def cjke_cleaners(self, text): return re.sub(r'([^\.,!\?\-…~])$', r'\1.', re.sub(r'\s+$', '', self.english_to_ipa2(text).replace('ɑ', 'a').replace('ɔ', 'o').replace('ɛ', 'e').replace('ɪ', 'i').replace('ʊ', 'u')))
def cjke_cleaners2(self, text): return re.sub(r'([^\.,!\?\-…~])$', r'\1.', re.sub(r'\s+$', '', self.english_to_ipa2(text)))
def cjks_cleaners(self, text): return re.sub(r'([^\.,!\?\-…~])$', r'\1.', re.sub(r'\s+$', '', self.english_to_lazy_ipa(text)))
def english_to_ipa2(self, text):
_ipa_to_ipa2 = [(re.compile('%s' % x[0]), x[1]) for x in [ ('r', 'ɹ'), ('ʤ', ''), ('ʧ', '')]]
return reduce(lambda t, rx: re.sub(rx[0], rx[1], t), _ipa_to_ipa2, self.mark_dark_l(self.english_to_ipa(text))).replace('...', '')
def mark_dark_l(self, text): return re.sub(r'l([^aeiouæɑɔəɛɪʊ ]*(?: |$))', lambda x: 'ɫ' + x.group(1), text)
def english_to_ipa(self, text):
import eng_to_ipa as ipa
return self.collapse_whitespace(ipa.convert(self.normalize_numbers(self.expand_abbreviations(unidecode(text).lower()))))
def english_to_lazy_ipa(self, text):
_lazy_ipa = [(re.compile('%s' % x[0]), x[1]) for x in [('r', 'ɹ'), ('æ', 'e'), ('ɑ', 'a'), ('ɔ', 'o'), ('ð', 'z'), ('θ', 's'), ('ɛ', 'e'), ('ɪ', 'i'), ('ʊ', 'u'), ('ʒ', 'ʥ'), ('ʤ', 'ʥ'), ('ˈ', '')]]
return reduce(lambda t, rx: re.sub(rx[0], rx[1], t), _lazy_ipa, self.english_to_ipa(text))
def expand_abbreviations(self, text): return reduce(lambda t, abbr: re.sub(abbr[0], abbr[1], t), self._abbreviations, text)
def collapse_whitespace(self, text): return re.sub(self._whitespace_re, ' ', text)
def normalize_numbers(self, text):
import inflect
self._inflect = inflect.engine()
text = re.sub(re.compile(r'([0-9][0-9\,]+[0-9])'), self._remove_commas, text)
text = re.sub(re.compile(r'£([0-9\,]*[0-9]+)'), r'\1 pounds', text)
text = re.sub(re.compile(r'\$([0-9\.\,]*[0-9]+)'), self._expand_dollars, text)
text = re.sub(re.compile(r'([0-9]+\.[0-9]+)'), self._expand_decimal_point, text)
text = re.sub(re.compile(r'[0-9]+(st|nd|rd|th)'), self._expand_ordinal, text)
text = re.sub(re.compile(r'[0-9]+'), self._expand_number, text)
return text
def _remove_commas(self, m): return m.group(1).replace(',', '') # george won't like this
def _expand_dollars(self, m):
match = m.group(1)
parts = match.split('.')
if len(parts) > 2: return match + ' dollars' # Unexpected format
dollars, cents = int(parts[0]) if parts[0] else 0, int(parts[1]) if len(parts) > 1 and parts[1] else 0
if dollars and cents: return '%s %s, %s %s' % (dollars, 'dollar' if dollars == 1 else 'dollars', cents, 'cent' if cents == 1 else 'cents')
if dollars: return '%s %s' % (dollars, 'dollar' if dollars == 1 else 'dollars')
if cents: return '%s %s' % (cents, 'cent' if cents == 1 else 'cents')
return 'zero dollars'
def _expand_decimal_point(self, m): return m.group(1).replace('.', ' point ')
def _expand_ordinal(self, m): return self._inflect.number_to_words(m.group(0))
def _expand_number(self, _inflect, m):
num = int(m.group(0))
if 1000 < num < 3000:
if num == 2000: return 'two thousand'
if 2000 < num < 2010: return 'two thousand ' + self._inflect.number_to_words(num % 100)
if num % 100 == 0: return self._inflect.number_to_words(num // 100) + ' hundred'
return _inflect.number_to_words(num, andword='', zero='oh', group=2).replace(', ', ' ')
return self._inflect.number_to_words(num, andword='')
#########################################################################################
# PAPER: https://arxiv.org/abs/2106.06103
# CODE: https://github.com/jaywalnut310/vits/tree/main
#########################################################################################
# INSTALLATION: this is based on default config, dependencies are for preprocessing.
# vctk, ljs | pip3 install unidecode phonemizer | phonemizer requires [eSpeak](https://espeak.sourceforge.net) backend to be installed on your system
# mmts-tts | pip3 install unidecode |
# uma_trilingual, cjks, voistock | pip3 install unidecode inflect eng_to_ipa |
#########################################################################################
# Some good speakers to try out, there may be much better ones, I only tried out a few:
# male vctk 1 | --model_to_use vctk --speaker_id 2
# male vctk 2 | --model_to_use vctk --speaker_id 6
# anime lady 1 | --model_to_use uma_trilingual --speaker_id 36
# anime lady 2 | --model_to_use uma_trilingual --speaker_id 121
#########################################################################################
VITS_PATH = Path(__file__).parents[1] / "weights/VITS/"
MODELS = { # config_url, weights_url
"ljs": ("https://raw.githubusercontent.com/jaywalnut310/vits/main/configs/ljs_base.json", "https://drive.google.com/uc?export=download&id=1q86w74Ygw2hNzYP9cWkeClGT5X25PvBT&confirm=t"),
"vctk": ("https://huggingface.co/csukuangfj/vits-vctk/resolve/main/vctk_base.json", "https://huggingface.co/csukuangfj/vits-vctk/resolve/main/pretrained_vctk.pth"),
"mmts-tts": ("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/config.json", "https://huggingface.co/facebook/mms-tts/resolve/main/full_models/eng/G_100000.pth"),
"uma_trilingual": ("https://huggingface.co/spaces/Plachta/VITS-Umamusume-voice-synthesizer/raw/main/configs/uma_trilingual.json", "https://huggingface.co/spaces/Plachta/VITS-Umamusume-voice-synthesizer/resolve/main/pretrained_models/G_trilingual.pth"),
"cjks": ("https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/14/config.json", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/14/model.pth"),
"voistock": ("https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/15/config.json", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/15/model.pth"),
}
Y_LENGTH_ESTIMATE_SCALARS = {"ljs": 2.8, "vctk": 1.74, "mmts-tts": 1.9, "uma_trilingual": 2.3, "cjks": 3.3, "voistock": 3.1}
if __name__ == '__main__':
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
parser = argparse.ArgumentParser()
parser.add_argument("--model_to_use", default="vctk", help="Specify the model to use. Default is 'vctk'.")
parser.add_argument("--speaker_id", type=int, default=6, help="Specify the speaker ID. Default is 6.")
parser.add_argument("--out_path", default=None, help="Specify the full output path. Overrides the --out_dir and --name parameter.")
parser.add_argument("--out_dir", default=str(Path(__file__).parents[1] / "temp"), help="Specify the output path.")
parser.add_argument("--base_name", default="test", help="Specify the base of the output file name. Default is 'test'.")
parser.add_argument("--text_to_synthesize", default="""Hello person. If the code you are contributing isn't some of the highest quality code you've written in your life, either put in the effort to make it great, or don't bother.""", help="Specify the text to synthesize. Default is a greeting message.")
parser.add_argument("--noise_scale", type=float, default=0.667, help="Specify the noise scale. Default is 0.667.")
parser.add_argument("--noise_scale_w", type=float, default=0.8, help="Specify the noise scale w. Default is 0.8.")
parser.add_argument("--length_scale", type=float, default=1, help="Specify the length scale. Default is 1.")
parser.add_argument("--seed", type=int, default=1337, help="Specify the seed (set to None if no seed). Default is 1337.")
parser.add_argument("--num_channels", type=int, default=1, help="Specify the number of audio output channels. Default is 1.")
parser.add_argument("--sample_width", type=int, default=2, help="Specify the number of bytes per sample, adjust if necessary. Default is 2.")
parser.add_argument("--emotion_path", type=str, default=None, help="Specify the path to emotion reference.")
parser.add_argument("--estimate_max_y_length", type=str, default=False, help="If true, overestimate the output length and then trim it to the correct length, to prevent premature realization, much more performant for larger inputs, for smaller inputs not so much. Default is False.")
args = parser.parse_args()
model_config = MODELS[args.model_to_use]
# Load the hyperparameters from the config file.
hps = get_hparams_from_file(fetch(model_config[0]))
# If model has multiple speakers, validate speaker id and retrieve name if available.
model_has_multiple_speakers = hps.data.n_speakers > 0
if model_has_multiple_speakers:
logging.info(f"Model has {hps.data.n_speakers} speakers")
if args.speaker_id >= hps.data.n_speakers: raise ValueError(f"Speaker ID {args.speaker_id} is invalid for this model.")
speaker_name = "?"
if hps.__contains__("speakers"): # maps speaker ids to names
speakers = hps.speakers
if isinstance(speakers, List): speakers = {speaker: i for i, speaker in enumerate(speakers)}
speaker_name = next((key for key, value in speakers.items() if value == args.speaker_id), None)
logging.info(f"You selected speaker {args.speaker_id} (name: {speaker_name})")
# Load emotions if any. TODO: find an english model with emotions, this is untested atm.
emotion_embedding = None
if args.emotion_path is not None:
if args.emotion_path.endswith(".npy"): emotion_embedding = Tensor(np.load(args.emotion_path), dtype=dtypes.int64).unsqueeze(0)
else: raise ValueError("Emotion path must be a .npy file.")
# Load symbols, instantiate TextMapper and clean the text.
if hps.__contains__("symbols"): symbols = hps.symbols
elif args.model_to_use == "mmts-tts": symbols = [x.replace("\n", "") for x in fetch("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/vocab.txt").open(encoding="utf-8").readlines()]
else: symbols = ['_'] + list(';:,.!?¡¿—…"«»“” ') + list('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz') + list("ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'")
text_mapper = TextMapper(apply_cleaners=True, symbols=symbols)
# Load the model.
if args.seed is not None:
Tensor.manual_seed(args.seed)
np.random.seed(args.seed)
net_g = load_model(text_mapper.symbols, hps, model_config)
logging.debug(f"Loaded model with hps: {hps}")
# Convert the input text to a tensor.
text_to_synthesize = args.text_to_synthesize
if args.model_to_use == "mmts-tts": text_to_synthesize = text_mapper.filter_oov(text_to_synthesize.lower())
stn_tst = text_mapper.get_text(text_to_synthesize, hps.data.add_blank, hps.data.text_cleaners)
logging.debug(f"Converted input text to tensor \"{text_to_synthesize}\" -> Tensor({stn_tst.shape}): {stn_tst.numpy()}")
x_tst, x_tst_lengths = stn_tst.unsqueeze(0), Tensor([stn_tst.shape[0]], dtype=dtypes.int64)
sid = Tensor([args.speaker_id], dtype=dtypes.int64) if model_has_multiple_speakers else None
# Perform inference.
start_time = time.time()
audio_tensor = net_g.infer(x_tst, x_tst_lengths, sid, args.noise_scale, args.length_scale, args.noise_scale_w, emotion_embedding=emotion_embedding,
max_y_length_estimate_scale=Y_LENGTH_ESTIMATE_SCALARS[args.model_to_use] if args.estimate_max_y_length else None)[0, 0].realize()
logging.info(f"Inference took {(time.time() - start_time):.2f}s")
# Save the audio output.
audio_data = (np.clip(audio_tensor.numpy(), -1.0, 1.0) * 32767).astype(np.int16)
out_path = Path(args.out_path or Path(args.out_dir)/f"{args.model_to_use}{f'_sid_{args.speaker_id}' if model_has_multiple_speakers else ''}_{args.base_name}.wav")
out_path.parent.mkdir(parents=True, exist_ok=True)
with wave.open(str(out_path), 'wb') as wav_file:
wav_file.setnchannels(args.num_channels)
wav_file.setsampwidth(args.sample_width)
wav_file.setframerate(hps.data.sampling_rate)
wav_file.setnframes(len(audio_data))
wav_file.writeframes(audio_data.tobytes())
logging.info(f"Saved audio output to {out_path}")
+26 -99
View File
@@ -26,13 +26,11 @@ def color_temp(temp):
def color_voltage(voltage): return colored(f"{voltage/1000:>5.3f}V", "cyan")
def draw_bar(percentage, width=40, fill='|', empty=' ', opt_text='', color='cyan'):
percentage = 0.0 if percentage != percentage else percentage # NaN guard
percentage = max(0.0, min(1.0, float(percentage)))
filled_width = int(width * percentage)
if not opt_text: opt_text = f'{percentage*100:.1f}%'
bar = fill * filled_width + empty * (width - filled_width)
if opt_text and len(opt_text) <= len(bar): bar = (bar[:-len(opt_text)] + opt_text)
bar = (bar[:-len(opt_text)] + opt_text) if opt_text else bar
bar = colored(bar[:filled_width], color) + bar[filled_width:]
return f'[{bar}]'
@@ -90,7 +88,6 @@ class SMICtx:
self.opened_pci_resources = {}
self.prev_lines_cnt = 0
self.prev_terminal_width = 0
self.prev_terminal_height = 0
remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
@@ -98,20 +95,6 @@ class SMICtx:
for k,v in self.lspci.items():
for part in remove_parts: self.lspci[k] = self.lspci[k].replace(part, "").strip().rstrip()
def _smuq10_round(self, v:int) -> int:
v = int(v)
return (v + 512) >> 10 # SMUQ10_ROUND
def _fmt_kb(self, kb:int) -> str:
kb = int(kb)
if kb < 1024: return f"{kb}KB"
mb = kb / 1024.0
if mb < 1024: return f"{mb:.1f}MB"
gb = mb / 1024.0
if gb < 1024: return f"{gb:.2f}GB"
tb = gb / 1024.0
return f"{tb:.2f}TB"
def _open_am_device(self, pcibus):
if pcibus not in self.opened_pci_resources:
bar_fds = {bar: os.open(f"/sys/bus/pci/devices/{pcibus}/resource{bar}", os.O_RDWR | os.O_SYNC) for bar in [0, 2, 5]}
@@ -133,7 +116,6 @@ class SMICtx:
def rescan_devs(self):
pattern = os.path.join('/tmp', 'am_*.lock')
for d in [f[8:-5] for f in glob.glob(pattern)]:
if d.startswith("usb"): continue
if d not in self.opened_pcidevs:
self._open_am_device(d)
@@ -149,52 +131,21 @@ class SMICtx:
os.system('clear')
if DEBUG >= 2: print(f"Removed AM device {d.pcibus}")
def collect(self):
tables = {}
for dev in self.devs:
match dev.ip_ver[am.MP1_HWIP]:
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
def collect(self): return {d: d.smu.read_metrics() if d.pci_state == "D0" else None for d in self.devs}
def _pick_nonzero_avg(self, vals) -> int:
xs = [x for x in vals if x > 0]
return int(sum(xs) / len(xs)) if xs else 0
def get_gfx_activity(self, dev, metrics):
match dev.ip_ver[am.MP1_HWIP]:
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)|(13,0,12): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
case _: return metrics.SmuMetrics.AverageUclkActivity
def get_gfx_activity(self, dev, metrics): return metrics.SmuMetrics.AverageGfxActivity
def get_mem_activity(self, dev, metrics): return metrics.SmuMetrics.AverageUclkActivity
def get_temps(self, dev, metrics, compact=False):
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6)|(13,0,12):
temps = {
"Hotspot": self._smuq10_round(metrics.MaxSocketTemperature),
"HBM": self._smuq10_round(metrics.MaxHbmTemperature),
"VR": self._smuq10_round(metrics.MaxVrTemperature),
}
if compact: return {k: temps[k] for k in ("Hotspot", "HBM") if temps.get(k, 0) != 0}
return {k: v for k, v in temps.items() if v != 0}
case _:
temps_keys = [(k, name) for k, name in dev.smu.smu_mod.TEMP_e.items()
if k < dev.smu.smu_mod.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
if compact: temps_keys = [(k, name) for k, name in temps_keys if k in (dev.smu.smu_mod.TEMP_HOTSPOT, dev.smu.smu_mod.TEMP_MEM)]
return {name: metrics.SmuMetrics.AvgTemperature[k] for k, name in temps_keys}
temps_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_TEMP_e__enumvalues.items()
if k < dev.smu.smu_mod.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
if compact: temps_keys = [(k, name) for k, name in temps_keys if k in (dev.smu.smu_mod.TEMP_HOTSPOT, dev.smu.smu_mod.TEMP_MEM)]
return {name: metrics.SmuMetrics.AvgTemperature[k] for k, name in temps_keys}
def get_voltage(self, dev, metrics, compact=False):
match dev.ip_ver[am.MP1_HWIP]:
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()
voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_SVI_PLANE_e__enumvalues.items()
if k < dev.smu.smu_mod.SVI_PLANE_COUNT and metrics.SmuMetrics.AvgVoltage[k] != 0]
return {name: metrics.SmuMetrics.AvgVoltage[k] for k, name in voltage_keys}
return {name: metrics.SmuMetrics.AvgVoltage[k] for k, name in voltage_keys}
def get_busy_threshold(self, dev):
match dev.ip_ver[am.MP1_HWIP]:
@@ -202,40 +153,22 @@ class SMICtx:
case _: return 15
def get_gfx_freq(self, dev, metrics):
if metrics is None: return 0
match dev.ip_ver[am.MP1_HWIP]:
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
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)|(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
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)|(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
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)|(13,0,12): return None, None
case _: return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm
def get_fan_rpm_pwm(self, dev, metrics): return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm
def get_power(self, dev, metrics):
match dev.ip_ver[am.MP1_HWIP]:
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_power(self, dev, metrics): return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
def get_mem_usage(self, dev):
return 0
usage = 0
pt_stack = [dev.mm.root_page_table]
while len(pt_stack) > 0:
@@ -244,7 +177,7 @@ class SMICtx:
entry = pt.entries[i]
if (entry & am.AMDGPU_PTE_VALID) == 0: continue
if pt.lv!=am.AMDGPU_VM_PTB and not dev.gmc.is_pte_huge_page(pt.lv, entry):
if pt.lv!=am.AMDGPU_VM_PTB and not dev.gmc.is_pte_huge_page(entry):
pt_stack.append(AMPageTableEntry(dev, entry & 0x0000FFFFFFFFF000, lv=pt.lv+1))
continue
if (entry & am.AMDGPU_PTE_SYSTEM) != 0: continue
@@ -286,28 +219,23 @@ class SMICtx:
temps_table_compact = ["Temps (°C):" + '/'.join([f"{color_temp(val)} {name}" for name, val in temps_data_compact.items()])]
fan_rpm, fan_pwm = self.get_fan_rpm_pwm(dev, metrics)
power_table = ["=== Power ==="]
power_table += ["Fan: N/A"] if fan_rpm is None or fan_pwm is None else [f"Fan Speed: {fan_rpm} RPM", f"Fan Power: {fan_pwm}%"]
power_table = ["=== Power ==="] + [f"Fan Speed: {fan_rpm} RPM"] + [f"Fan Power: {fan_pwm}%"]
total_power, max_power = self.get_power(dev, metrics)
if max_power > 0:
power_line = [f"Power: " + draw_bar(total_power / max_power, 16, opt_text=f"{total_power}/{max_power}W")]
power_line_compact = [f"Power: " + draw_bar(total_power / max_power, activity_line_width, opt_text=f"{total_power}/{max_power}W")]
else:
power_line = ["Power: N/A"]
power_line_compact = ["Power: N/A"]
power_line = [f"Power: " + draw_bar(total_power / max_power, 16, opt_text=f"{total_power}/{max_power}W")]
power_line_compact = [f"Power: " + draw_bar(total_power / max_power, activity_line_width, opt_text=f"{total_power}/{max_power}W")]
voltage_data = self.get_voltage(dev, metrics)
voltage_table = None if not voltage_data else (["=== Voltages ==="] + [f"{name:<20}: {color_voltage(voltage)}" for name, voltage in voltage_data.items()])
voltage_table = ["=== Voltages ==="] + [f"{name:<20}: {color_voltage(voltage)}" for name, voltage in voltage_data.items()]
gfx_freq = self.get_gfx_freq(dev, metrics)
mclk_freq = self.get_mem_freq(dev, metrics)
fclk_freq = self.get_fckl_freq(dev, metrics)
frequency_table = ["=== Frequencies ===", f"GFXCLK: {gfx_freq:>4} MHz", f"FCLK : {fclk_freq:>4} MHz", f"MCLK : {mclk_freq:>4} MHz"]
if self.prev_terminal_width >= 231:
power_table += power_line
if voltage_table is not None: power_table += [""] + voltage_table
power_table += power_line + [""] + voltage_table
activity_line += [""]
elif self.prev_terminal_width >= 171:
power_table += power_line + [""] + frequency_table
@@ -379,5 +307,4 @@ if __name__ == "__main__":
smi_ctx.draw(args.list)
if args.list: break
time.sleep(1)
except KeyboardInterrupt:
print("Exiting...")
except KeyboardInterrupt: print("Exiting...")
-14
View File
@@ -1,14 +0,0 @@
#!/usr/bin/env python3
from tinygrad.helpers import Context
from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
from tinygrad.runtime.support.am.amdev import AMDev
if __name__ == "__main__":
gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
amdevs = []
with Context(DEBUG=2):
for pcidev in pcidevs:
amdevs.append(AMDev(pcidev, reset_mode=True))
for amdev in amdevs: amdev.smu.mode1_reset()
+20 -36
View File
@@ -1,65 +1,48 @@
import re, ctypes, sys, importlib
from tinygrad.helpers import getenv
from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
class GFXFake:
def __init__(self): self.xccs = 8
class AMDFake(AMDev):
def __init__(self, pci_dev, dma_regions=None):
self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
def __init__(self, devfmt, vram, doorbell, mmio, dma_regions=None):
self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
self._run_discovery()
self._build_regs()
self.gfx = GFXFake()
amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
amdev.AMDev = AMDFake
from tinygrad.runtime.ops_amd import PCIIface
def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
register_map = register_names or {}
def parse_amdgpu_logs(log_content, register_names=None):
register_map = register_names
final = ""
def replace_register(match):
reg = match.group(1)
return f"Reading register {register_map.get(int(reg, 16), reg)}"
register = match.group(1)
return f"Reading register {register_map.get(int(register, base=16), register)}"
processed_log = re.sub(r'Reading register (0x[0-9a-fA-F]+)', replace_register, log_content)
pattern = r'Reading register (0x[0-9a-fA-F]+)'
processed_log = re.sub(pattern, replace_register, log_content)
def replace_register_2(match):
reg = match.group(1)
return f"Writing register {register_map.get(int(reg, 16), reg)}"
processed_log = re.sub(r'Writing register (0x[0-9a-fA-F]+)', replace_register_2, processed_log)
# remove timing prefix
processed_log = re.sub(r'^\[\s*\d+(?:\.\d+)?\]\s*', '', processed_log, flags=re.MULTILINE)
# keep only xcc=0 lines (but keep lines with no xcc at all)
if only_xcc0:
kept = []
for line in processed_log.splitlines(True):
if "xcc=" not in line or re.search(r'\bxcc=0\b', line): kept.append(line)
processed_log = "".join(kept)
register = match.group(1)
return f"Writing register {register_map.get(int(register, base=16), register)}"
pattern = r'Writing register (0x[0-9a-fA-F]+)'
processed_log = re.sub(pattern, replace_register_2, processed_log)
return processed_log
def main():
only_xcc0 = bool(getenv("ONLY_XCC0", 0))
reg_names = {}
dev = PCIIface(None, 0)
for x, y in dev.dev_impl.__dict__.items():
if isinstance(y, AMRegister):
for xcc, addr in y.addr.items():
reg_names[addr] = f"{x}, xcc={xcc}"
for inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"
with open(sys.argv[1], 'r') as f:
log_content = f.read()
log_content = log_content_them = f.read()
processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
processed_log = parse_amdgpu_logs(log_content, reg_names)
with open(sys.argv[2], 'w') as f:
f.write(processed_log)
@@ -68,4 +51,5 @@ if __name__ == '__main__':
if len(sys.argv) != 3:
print("Usage: <input_file_path> <output_file_path>")
sys.exit(1)
main()
main()
-39
View File
@@ -1,39 +0,0 @@
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`
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
-641
View File
@@ -1,641 +0,0 @@
# library for RDNA3 assembly DSL
# mypy: ignore-errors
from __future__ import annotations
import struct, math, re
from enum import IntEnum
from functools import cache
from typing import overload, Annotated, TypeVar, Generic
from extra.assembly.amd.autogen.rdna3.enum import (VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, VOPDOp, SOP1Op, SOP2Op,
SOPCOp, SOPKOp, SOPPOp, SMEMOp, DSOp, FLATOp, MUBUFOp, MTBUFOp, MIMGOp, VINTERPOp)
from extra.assembly.amd.autogen.cdna.enum import VOP1Op as CDNA_VOP1Op, VOP2Op as CDNA_VOP2Op
from extra.assembly.amd.autogen.rdna4.enum import (VOP1Op as RDNA4_VOP1Op, VOP2Op as RDNA4_VOP2Op, VOP3Op as RDNA4_VOP3Op,
VOP3SDOp as RDNA4_VOP3SDOp, VOP3POp as RDNA4_VOP3POp, VOPCOp as RDNA4_VOPCOp, VOPDOp as RDNA4_VOPDOp,
SOP1Op as RDNA4_SOP1Op, SOP2Op as RDNA4_SOP2Op, SOPCOp as RDNA4_SOPCOp, SOPKOp as RDNA4_SOPKOp, SOPPOp as RDNA4_SOPPOp,
SMEMOp as RDNA4_SMEMOp, DSOp as RDNA4_DSOp, VBUFFEROp as RDNA4_VBUFFEROp, VINTERPOp as RDNA4_VINTERPOp)
# Source operand encoding - constant across all AMD ISAs
class SrcEnum(IntEnum):
VCC_LO=106; VCC_HI=107; NULL=124; M0=125; EXEC_LO=126; EXEC_HI=127; ZERO=128
DPP8=233; DPP8FI=234; SHARED_BASE=235; SHARED_LIMIT=236; PRIVATE_BASE=237; PRIVATE_LIMIT=238
POS_HALF=240; NEG_HALF=241; POS_ONE=242; NEG_ONE=243; POS_TWO=244; NEG_TWO=245
POS_FOUR=246; NEG_FOUR=247; INV_2PI=248; DPP16=250; VCCZ=251; EXECZ=252; SCC=253; LDS_DIRECT=254
VCC_LO, VCC_HI, NULL, M0, EXEC_LO, EXEC_HI, ZERO = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.M0, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.ZERO
DPP8FI, SHARED_BASE, SHARED_LIMIT, PRIVATE_BASE, PRIVATE_LIMIT = SrcEnum.DPP8FI, SrcEnum.SHARED_BASE, SrcEnum.SHARED_LIMIT, SrcEnum.PRIVATE_BASE, SrcEnum.PRIVATE_LIMIT
POS_HALF, NEG_HALF, POS_ONE, NEG_ONE, POS_TWO, NEG_TWO = SrcEnum.POS_HALF, SrcEnum.NEG_HALF, SrcEnum.POS_ONE, SrcEnum.NEG_ONE, SrcEnum.POS_TWO, SrcEnum.NEG_TWO
POS_FOUR, NEG_FOUR, INV_2PI, VCCZ, EXECZ, SCC, LDS_DIRECT = SrcEnum.POS_FOUR, SrcEnum.NEG_FOUR, SrcEnum.INV_2PI, SrcEnum.VCCZ, SrcEnum.EXECZ, SrcEnum.SCC, SrcEnum.LDS_DIRECT
OFF = NULL
# Common masks and bit conversion functions
MASK32, MASK64, MASK128 = 0xffffffff, 0xffffffffffffffff, (1 << 128) - 1
_struct_f, _struct_I = struct.Struct("<f"), struct.Struct("<I")
_struct_e, _struct_H = struct.Struct("<e"), struct.Struct("<H")
_struct_d, _struct_Q = struct.Struct("<d"), struct.Struct("<Q")
def _f32(i):
i = i & MASK32
# RDNA3 default mode: flush f32 denormals to zero (FTZ)
# Denormal: exponent=0 (bits 23-30) and mantissa!=0 (bits 0-22)
if (i & 0x7f800000) == 0 and (i & 0x007fffff) != 0: return 0.0
return _struct_f.unpack(_struct_I.pack(i))[0]
def _i32(f):
if isinstance(f, int): f = float(f)
if math.isnan(f): return 0xffc00000 if math.copysign(1.0, f) < 0 else 0x7fc00000
if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
try:
bits = _struct_I.unpack(_struct_f.pack(f))[0]
# RDNA3 default mode: flush f32 denormals to zero (FTZ)
if (bits & 0x7f800000) == 0 and (bits & 0x007fffff) != 0: return 0x80000000 if bits & 0x80000000 else 0
return bits
except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
def _sext(v, b): return v - (1 << b) if v & (1 << (b - 1)) else v
def _f16(i): return _struct_e.unpack(_struct_H.pack(i & 0xffff))[0]
def _i16(f):
if math.isnan(f): return 0x7e00
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
try: return _struct_H.unpack(_struct_e.pack(f))[0]
except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
def _f64(i): return _struct_d.unpack(_struct_Q.pack(i & MASK64))[0]
def _i64(f):
if math.isnan(f): return 0x7ff8000000000000
if math.isinf(f): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
try: return _struct_Q.unpack(_struct_d.pack(f))[0]
except (OverflowError, struct.error): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
# Instruction spec - register counts and dtypes derived from instruction names
_REGS = {'B32': 1, 'B64': 2, 'B96': 3, 'B128': 4, 'B256': 8, 'B512': 16,
'F32': 1, 'I32': 1, 'U32': 1, 'F64': 2, 'I64': 2, 'U64': 2,
'F16': 1, 'I16': 1, 'U16': 1, 'B16': 1, 'I8': 1, 'U8': 1, 'B8': 1,
'DWORD': 1, 'DWORDX2': 2, 'DWORDX3': 3, 'DWORDX4': 4, 'DWORDX8': 8, 'DWORDX16': 16,
'BYTE': 1, 'SHORT': 1, 'UBYTE': 1, 'SBYTE': 1, 'USHORT': 1, 'SSHORT': 1}
_CVT_RE = re.compile(r'CVT_([FIUB]\d+)_([FIUB]\d+)$')
_MAD_MUL_RE = re.compile(r'(?:MAD|MUL)_([IU]\d+)_([IU]\d+)$')
_PACK_RE = re.compile(r'PACK_([FIUB]\d+)_([FIUB]\d+)$')
_DST_SRC_RE = re.compile(r'_([FIUB]\d+)_([FIUB]\d+)$')
_SINGLE_RE = re.compile(r'_([FIUB](?:32|64|16|8|96|128|256|512)|DWORD(?:X(?:2|3|4|8|16))?|[US]?BYTE|[US]?SHORT)$')
@cache
def _suffix(name: str) -> tuple[str | None, str | None]:
name = name.upper()
if m := _CVT_RE.search(name): return m.group(1), m.group(2)
if m := _MAD_MUL_RE.search(name): return m.group(1), m.group(2)
if m := _PACK_RE.search(name): return m.group(1), m.group(2)
if m := _DST_SRC_RE.search(name): return m.group(1), m.group(2)
if m := _SINGLE_RE.search(name): return m.group(1), m.group(1)
return None, None
_SPECIAL_REGS = {
'V_LSHLREV_B64': (2, 1, 2, 1), 'V_LSHRREV_B64': (2, 1, 2, 1), 'V_ASHRREV_I64': (2, 1, 2, 1),
'S_LSHL_B64': (2, 2, 1, 1), 'S_LSHR_B64': (2, 2, 1, 1), 'S_ASHR_I64': (2, 2, 1, 1),
'S_BFE_U64': (2, 2, 1, 1), 'S_BFE_I64': (2, 2, 1, 1), 'S_BFM_B64': (2, 1, 1, 1),
'S_BITSET0_B64': (2, 1, 1, 1), 'S_BITSET1_B64': (2, 1, 1, 1),
'S_BITCMP0_B64': (1, 2, 1, 1), 'S_BITCMP1_B64': (1, 2, 1, 1),
'V_LDEXP_F64': (2, 2, 1, 1), 'V_TRIG_PREOP_F64': (2, 2, 1, 1),
'V_CMP_CLASS_F64': (1, 2, 1, 1), 'V_CMPX_CLASS_F64': (1, 2, 1, 1),
'V_CMP_CLASS_F32': (1, 1, 1, 1), 'V_CMPX_CLASS_F32': (1, 1, 1, 1),
'V_CMP_CLASS_F16': (1, 1, 1, 1), 'V_CMPX_CLASS_F16': (1, 1, 1, 1),
'V_MAD_U64_U32': (2, 1, 1, 2), 'V_MAD_I64_I32': (2, 1, 1, 2),
'V_QSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_U32_U8': (4, 2, 1, 4),
'V_CVT_PK_F32_BF8': (2, 1, 1, 1), 'V_CVT_PK_F32_FP8': (2, 1, 1, 1),
}
_SPECIAL_DTYPE = {
'V_LSHLREV_B64': ('B64', 'U32', 'B64', None), 'V_LSHRREV_B64': ('B64', 'U32', 'B64', None), 'V_ASHRREV_I64': ('I64', 'U32', 'I64', None),
'S_LSHL_B64': ('B64', 'B64', 'U32', None), 'S_LSHR_B64': ('B64', 'B64', 'U32', None), 'S_ASHR_I64': ('I64', 'I64', 'U32', None),
'S_BFE_U64': ('U64', 'U64', 'U32', None), 'S_BFE_I64': ('I64', 'I64', 'U32', None),
'S_BFM_B64': ('B64', 'U32', 'U32', None), 'S_BITSET0_B64': ('B64', 'U32', None, None), 'S_BITSET1_B64': ('B64', 'U32', None, None),
'S_BITCMP0_B64': ('SCC', 'B64', 'U32', None), 'S_BITCMP1_B64': ('SCC', 'B64', 'U32', None),
'V_LDEXP_F64': ('F64', 'F64', 'I32', None), 'V_TRIG_PREOP_F64': ('F64', 'F64', 'U32', None),
'V_CMP_CLASS_F64': ('VCC', 'F64', 'U32', None), 'V_CMPX_CLASS_F64': ('EXEC', 'F64', 'U32', None),
'V_CMP_CLASS_F32': ('VCC', 'F32', 'U32', None), 'V_CMPX_CLASS_F32': ('EXEC', 'F32', 'U32', None),
'V_CMP_CLASS_F16': ('VCC', 'F16', 'U32', None), 'V_CMPX_CLASS_F16': ('EXEC', 'F16', 'U32', None),
'V_MAD_U64_U32': ('U64', 'U32', 'U32', 'U64'), 'V_MAD_I64_I32': ('I64', 'I32', 'I32', 'I64'),
'V_QSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'), 'V_MQSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'),
'V_MQSAD_U32_U8': ('B128', 'B64', 'B64', 'B128'),
}
@cache
def spec_regs(name: str) -> tuple[int, int, int, int]:
uname = name.upper()
if uname in _SPECIAL_REGS: return _SPECIAL_REGS[uname]
if 'SAD' in uname and 'U8' in uname and 'QSAD' not in uname and 'MQSAD' not in uname: return 1, 1, 1, 1
dst_suf, src_suf = _suffix(name)
return _REGS.get(dst_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1)
@cache
def spec_dtype(name: str) -> tuple[str | None, str | None, str | None, str | None]:
uname = name.upper()
if uname in _SPECIAL_DTYPE: return _SPECIAL_DTYPE[uname]
if 'SAD' in uname and ('U8' in uname or 'U16' in uname) and 'QSAD' not in uname and 'MQSAD' not in uname: return 'U32', 'U32', 'U32', 'U32'
if '_CMP_' in uname or '_CMPX_' in uname:
dst_suf, src_suf = _suffix(name)
return 'EXEC' if '_CMPX_' in uname else 'VCC', src_suf, src_suf, None
dst_suf, src_suf = _suffix(name)
return dst_suf, src_suf, src_suf, src_suf
_F16_RE = re.compile(r'_[FIUB]16(?:_|$)')
_F64_RE = re.compile(r'_[FIUB]64(?:_|$)')
@cache
def spec_is_16bit(name: str) -> bool:
uname = name.upper()
if 'SAD' in uname or 'PACK' in uname or '_PK_' in uname or 'SAT_PK' in uname or 'DOT2' in uname: return False
if '_F32' in uname or '_I32' in uname or '_U32' in uname or '_B32' in uname: return False
return bool(_F16_RE.search(uname))
@cache
def spec_is_64bit(name: str) -> bool: return bool(_F64_RE.search(name.upper()))
_3SRC = {'FMA', 'MAD', 'MIN3', 'MAX3', 'MED3', 'DIV_FIX', 'DIV_FMAS', 'DIV_SCALE', 'SAD', 'LERP', 'ALIGN', 'CUBE', 'BFE', 'BFI',
'PERM_B32', 'PERMLANE', 'CNDMASK', 'XOR3', 'OR3', 'ADD3', 'LSHL_OR', 'AND_OR', 'LSHL_ADD', 'ADD_LSHL', 'XAD', 'MAXMIN',
'MINMAX', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM', 'MINIMUM3', 'MAXIMUM3', 'DOT2', 'DOT4', 'DOT8', 'WMMA', 'CVT_PK_U8', 'MULLIT', 'CO_CI'}
_2SRC = {'FMAC', 'PERMLANE16_VAR', 'PERMLANEX16_VAR'} # FMAC uses dst as implicit accumulator, _VAR permlane only 2 sources
def spec_num_srcs(name: str) -> int:
name = name.upper()
if any(k in name for k in _2SRC): return 2
return 3 if any(k in name for k in _3SRC) else 2
def is_dtype_16(dt: str | None) -> bool: return dt is not None and '16' in dt
def is_dtype_64(dt: str | None) -> bool: return dt is not None and '64' in dt
# Bit field DSL
class BitField:
def __init__(self, hi: int, lo: int, name: str | None = None): self.hi, self.lo, self.name, self._marker = hi, lo, name, None
def __set_name__(self, owner, name):
import typing
self.name, self._owner = name, owner
# Cache marker at class definition time
hints = typing.get_type_hints(owner, include_extras=True)
if name in hints:
hint = hints[name]
if typing.get_origin(hint) is Annotated:
args = typing.get_args(hint)
self._marker = args[1] if len(args) > 1 else None
def __eq__(self, val: int) -> tuple[BitField, int]: return (self, val) # type: ignore
def mask(self) -> int: return (1 << (self.hi - self.lo + 1)) - 1
@property
def marker(self) -> type | None: return self._marker
@overload
def __get__(self, obj: None, objtype: type) -> BitField: ...
@overload
def __get__(self, obj: object, objtype: type | None = None) -> int: ...
def __get__(self, obj, objtype=None):
if obj is None: return self
val = unwrap(obj._values.get(self.name, 0))
# Convert to IntEnum if marker is an IntEnum subclass
if self.marker and isinstance(self.marker, type) and issubclass(self.marker, IntEnum):
# VOP3 with VOPC opcodes (0-255) -> VOPCOp, VOP3SD opcodes -> VOP3SDOp
if self.marker is VOP3Op:
if val < 256: return VOPCOp(val)
if val in Inst._VOP3SD_OPS: return VOP3SDOp(val)
try: return self.marker(val)
except ValueError: pass
return val
class _Bits:
def __getitem__(self, key) -> BitField: return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
bits = _Bits()
# Source operand with modifiers - base class for anything that can be a src with neg/abs
class SrcMod:
__slots__ = ('val', 'neg', 'abs_')
def __init__(self, val: int, neg: bool = False, abs_: bool = False): self.val, self.neg, self.abs_ = val, neg, abs_
def __repr__(self): return f"{'-' if self.neg else ''}{'|' if self.abs_ else ''}{self.val}{'|' if self.abs_ else ''}"
def __neg__(self): return SrcMod(self.val, not self.neg, self.abs_)
def __abs__(self): return SrcMod(self.val, self.neg, True)
# Register types
class Reg(SrcMod):
__slots__ = ('idx', 'count', 'hi')
def __init__(self, idx: int, count: int = 1, hi: bool = False, neg: bool = False, abs_: bool = False):
self.idx, self.count, self.hi = idx, count, hi
super().__init__(idx, neg, abs_)
def __repr__(self): return f"{self.__class__.__name__.lower()[0]}[{self.idx}]" if self.count == 1 else f"{self.__class__.__name__.lower()[0]}[{self.idx}:{self.idx + self.count}]"
def __neg__(self): return self.__class__(self.idx, self.count, self.hi, not self.neg, self.abs_)
def __abs__(self): return self.__class__(self.idx, self.count, self.hi, self.neg, True)
@property
def l(self): return self.__class__(self.idx, self.count, False, self.neg, self.abs_)
@property
def h(self): return self.__class__(self.idx, self.count, True, self.neg, self.abs_)
T = TypeVar('T', bound=Reg)
class _RegFactory(Generic[T]):
def __init__(self, cls: type[T], name: str): self._cls, self._name = cls, name
@overload
def __getitem__(self, key: int) -> Reg: ...
@overload
def __getitem__(self, key: slice) -> Reg: ...
def __getitem__(self, key: int | slice) -> Reg:
return self._cls(key.start, key.stop - key.start + 1) if isinstance(key, slice) else self._cls(key)
def __repr__(self): return f"<{self._name} factory>"
class SGPR(Reg): pass
class VGPR(Reg): pass
class TTMP(Reg): pass
s: _RegFactory[SGPR] = _RegFactory(SGPR, "SGPR")
v: _RegFactory[VGPR] = _RegFactory(VGPR, "VGPR")
ttmp: _RegFactory[TTMP] = _RegFactory(TTMP, "TTMP")
# Special registers as SrcMod objects (support -VCC_LO, abs(EXEC_LO), etc.)
VCC_LO, VCC_HI, VCC = SrcMod(106), SrcMod(107), SrcMod(106)
EXEC_LO, EXEC_HI, EXEC = SrcMod(126), SrcMod(127), SrcMod(126)
SCC, M0, NULL, OFF = SrcMod(253), SrcMod(125), SrcMod(124), SrcMod(124)
# Field type markers (runtime classes for validation)
class _SSrc: pass
class _Src: pass
class _Imm: pass
class _SImm: pass
class _VDSTYEnc: pass # VOPD vdsty: encoded = actual >> 1, actual = (encoded << 1) | ((vdstx & 1) ^ 1)
class _SGPRField: pass
class _VGPRField: pass
# Type aliases for annotations - tells mypy it's a BitField while preserving marker info
SSrc = Annotated[BitField, _SSrc]
Src = Annotated[BitField, _Src]
Imm = Annotated[BitField, _Imm]
SImm = Annotated[BitField, _SImm]
VDSTYEnc = Annotated[BitField, _VDSTYEnc]
SGPRField = Annotated[BitField, _SGPRField]
VGPRField = Annotated[BitField, _VGPRField]
class RawImm:
def __init__(self, val: int): self.val = val
def __repr__(self): return f"RawImm({self.val})"
def __eq__(self, other): return isinstance(other, RawImm) and self.val == other.val
def unwrap(val) -> int:
if isinstance(val, RawImm): return val.val
if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val # Special registers like VCC_LO, NULL
if hasattr(val, 'value'): return val.value # IntEnum
if hasattr(val, 'idx'): return val.idx # Reg
return val
# Encoding/decoding constants
FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
FLOAT_DEC = {v: str(k) for k, v in FLOAT_ENC.items()}
SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
SPECIAL_GPRS_CDNA = {102: "flat_scratch_lo", 103: "flat_scratch_hi", 104: "xnack_mask_lo", 105: "xnack_mask_hi",
106: "vcc_lo", 107: "vcc_hi", 124: "m0", 126: "exec_lo", 127: "exec_hi",
251: "src_vccz", 252: "src_execz", 253: "src_scc", 254: "src_lds_direct"}
SPECIAL_PAIRS = {106: "vcc", 126: "exec"}
SPECIAL_PAIRS_CDNA = {102: "flat_scratch", 104: "xnack_mask", 106: "vcc", 126: "exec"}
SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata', 'vsrc1'}
def _encode_reg(val: Reg) -> int: return (108 if isinstance(val, TTMP) else 0) + val.idx
def _is_encoded_src(v: int) -> bool: return 106 <= v <= 127 or 128 <= v <= 208 or 240 <= v <= 255 # Special regs (106-127) or inline const
def encode_src(val) -> int:
if isinstance(val, VGPR): return 256 + _encode_reg(val)
if isinstance(val, Reg): return _encode_reg(val)
if isinstance(val, SrcMod) and not isinstance(val, Reg):
v = val.val
if _is_encoded_src(v): return v # Already encoded (special reg 106-127 or inline const 128-208 or float 240-255)
if isinstance(v, int) and 0 <= v <= 64: return 128 + v # Encode as inline constant
if isinstance(v, int) and -16 <= v <= -1: return 192 - v
return 255 # Literal
if hasattr(val, 'value'): return val.value # IntEnum
if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
if isinstance(val, int): return 128 + val if 0 <= val <= 64 else 192 - val if -16 <= val <= -1 else 255
return 255
def decode_src(val: int, cdna: bool = False) -> str:
special = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
if val in special: return special[val]
if val <= 105: return f"s{val}"
if val in FLOAT_DEC: return FLOAT_DEC[val]
if 108 <= val <= 123: return f"ttmp{val - 108}"
if 128 <= val <= 192: return str(val - 128)
if 193 <= val <= 208: return str(-(val - 192))
if 256 <= val <= 511: return f"v{val - 256}"
return "lit" if val == 255 else f"?{val}"
# Instruction base class
class Inst:
_fields: dict[str, BitField]
_encoding: tuple[BitField, int] | None = None
_defaults: dict[str, int] = {}
_values: dict[str, int | RawImm]
_words: int # size in 32-bit words, set by decode_program
_literal: int | None
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
# Merge fields from parent classes
cls._fields = {}
for base in reversed(cls.__mro__):
if base is Inst or not hasattr(base, '_fields'): continue
cls._fields.update(base._fields)
# Add this class's own fields (overrides parents)
cls._fields.update({n: v[0] if isinstance(v, tuple) else v for n, v in cls.__dict__.items() if isinstance(v, BitField) or (isinstance(v, tuple) and len(v) == 2 and isinstance(v[0], BitField))})
# Compute size from max bit (exclude optional MIMG NSA fields: addr1/addr2 at bits 64+)
optional_nsa = {'addr1', 'addr2'}
max_bit = max((bf.hi for n, bf in cls._fields.items() if n not in optional_nsa), default=0) if cls._fields else 0
cls._sz = 12 if max_bit > 63 else 8 if max_bit > 31 else 4
if 'encoding' in cls._fields and isinstance(cls.__dict__.get('encoding'), tuple): cls._encoding = cls.__dict__['encoding']
def _or_field(self, name: str, bit: int):
cur = self._values.get(name, 0)
self._values[name] = (cur.val if isinstance(cur, RawImm) else cur) | bit
def _encode_src(self, name: str, val):
"""Encode a source field, handling modifiers and literals."""
encoded = encode_src(val)
has_opsel = 'opsel' in self._fields
if isinstance(val, Reg) and val.hi and not has_opsel: encoded |= 0x80 # hi bit in src for VOP1/2/C
self._values[name] = RawImm(encoded)
# Handle neg/abs/opsel modifiers
if isinstance(val, SrcMod):
mod_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
if val.neg and 'neg' in self._fields: self._or_field('neg', mod_bit)
if val.abs_ and 'abs' in self._fields: self._or_field('abs', mod_bit)
if isinstance(val, Reg) and val.hi and has_opsel:
self._or_field('opsel', {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0))
# Track literal value if needed
if encoded == 255 and self._literal is None:
import struct
# Check if THIS source uses 64-bit encoding (not just src0)
src_idx = {'src0': 0, 'src1': 1, 'src2': 2, 'ssrc0': 0, 'ssrc1': 1}.get(name, 0)
src_regs = self.src_regs(src_idx)
is_64 = src_regs == 2
if isinstance(val, SrcMod) and not isinstance(val, Reg): lit32 = val.val & MASK32
elif isinstance(val, int) and not isinstance(val, IntEnum): lit32 = val & MASK32
elif isinstance(val, float): lit32 = (_i64(val) >> 32) if is_64 else _i32(val) # f64: high 32 bits of f64 repr
else: return
self._literal = (lit32 << 32) if is_64 else lit32
def _encode_raw(self, name: str, val):
"""Encode a raw register field (vdst, vdata, etc.)."""
if isinstance(val, Reg):
encoded = _encode_reg(val)
if val.hi and 'opsel' not in self._fields: encoded |= 0x80
self._values[name] = encoded
if name == 'vdst' and val.hi and 'opsel' in self._fields: self._or_field('opsel', 8)
elif hasattr(val, 'value'): self._values[name] = val.value
def _validate(self, orig_args: dict):
"""Format-specific validation. Override in subclass or check by class name."""
cls_name, op = self.__class__.__name__, orig_args.get('op')
op_val = op.value if hasattr(op, 'value') else op
op_name = op.name if hasattr(op, 'name') else None
# SMEM: register count must match opcode (derive from name: b32=1, b64=2, b96=3, b128=4, b256=8, b512=16, i8/u8/i16/u16=1)
if cls_name == 'SMEM' and op_name:
expected = {'B32': 1, 'B64': 2, 'B96': 3, 'B128': 4, 'B256': 8, 'B512': 16, 'I8': 1, 'U8': 1, 'I16': 1, 'U16': 1}.get(op_name.split('_')[-1])
sdata = orig_args.get('sdata')
if expected and isinstance(sdata, Reg) and sdata.count != expected:
raise ValueError(f"SMEM op {op_name} expects {expected} registers, got {sdata.count}")
# SOP1: derive expected register sizes from op name (e.g., S_MOV_B64 -> dst=2, src=2; S_CTZ_I32_B64 -> dst=1, src=2)
if cls_name == 'SOP1' and hasattr(orig_args.get('op'), 'name'):
op_name = orig_args['op'].name
# Special cases: BITSET takes bit index (1 reg) regardless of dst size
if 'BITSET' in op_name:
dst_size = 2 if op_name.endswith('_B64') else 1
src_size = 1 # bit index is always 1 reg
else:
# Extract sizes from name: last suffix is src type, second-to-last (if exists) is dst type
sizes = {'B32': 1, 'I32': 1, 'U32': 1, 'B64': 2, 'I64': 2, 'U64': 2, 'B128': 4, 'B256': 8, 'B512': 16}
parts = op_name.split('_')
src_size = sizes.get(parts[-1], 1) if parts[-1] in sizes else 1
dst_size = sizes.get(parts[-2], src_size) if len(parts) >= 2 and parts[-2] in sizes else src_size
for fld, expected in [('sdst', dst_size), ('ssrc0', src_size)]:
if isinstance(orig_args.get(fld), Reg) and orig_args[fld].count != expected:
raise ValueError(f"SOP1 {op_name} expects {expected} register(s) for {fld}, got {orig_args[fld].count}")
def __init__(self, *args, literal: int | None = None, **kwargs):
self._values, self._literal = dict(self._defaults), None
field_names = [n for n in self._fields if n != 'encoding']
# Map Python-friendly names to actual field names (abs_ -> abs for Python reserved word)
if 'abs_' in kwargs: kwargs['abs'] = kwargs.pop('abs_')
# If more args than fields, treat extra arg as literal (for FMAAK/FMAMK style instructions)
# FMAMK has K in middle (vdst, src0, K, vsrc1), FMAAK has K at end (vdst, src0, vsrc1, K)
args = list(args)
if len(args) > len(field_names) and literal is None:
for i, a in enumerate(args):
if isinstance(a, int) and not isinstance(a, SrcEnum) and i < len(field_names) and field_names[i] in ('vsrc1',):
literal = args.pop(i)
break
else:
literal = args.pop() # fallback: last arg is literal
orig_args = dict(zip(field_names, args)) | kwargs
self._values.update(orig_args)
self._precompute()
self._validate(orig_args)
# Pre-shift literal for 64-bit sources (literal param is always raw 32-bit value from user)
if literal is not None:
# Find which source uses the literal (255) and check its register count
for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
v = orig_args.get(n)
if (isinstance(v, RawImm) and v.val == 255) or (isinstance(v, int) and v == 255):
self._literal = (literal << 32) if self.src_regs(idx) == 2 else literal
break
else:
self._literal = literal # fallback if no literal source found
cls_name = self.__class__.__name__
# Format-specific setup
if cls_name == 'FLAT' and 'sve' in self._fields:
seg = self._values.get('seg', 0)
if (seg.val if isinstance(seg, RawImm) else seg) == 1 and isinstance(orig_args.get('addr'), VGPR): self._values['sve'] = 1
if cls_name == 'VOP3P':
op = orig_args.get('op')
if hasattr(op, 'value'): op = op.value
# fma_mix ops (32-34) default to opsel_hi=0, WMMA ops (64-69) default to opsel_hi=7 to match LLVM
if op in (32, 33, 34) and 'opsel_hi' not in orig_args: self._values['opsel_hi'] = self._values['opsel_hi2'] = 0
if op in range(64, 70) and 'opsel_hi' not in orig_args: self._values['opsel_hi'], self._values['opsel_hi2'] = 3, 1
# Encode all fields
for name, val in list(self._values.items()):
if name == 'encoding': continue
if isinstance(val, RawImm):
if name in RAW_FIELDS: self._values[name] = val.val
continue
field = self._fields.get(name)
marker = field.marker if field else None
# Type validation
if marker is _SGPRField and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
if marker is _VGPRField and not isinstance(val, VGPR): raise TypeError(f"field '{name}' requires VGPR, got {type(val).__name__}")
if marker is _SSrc and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires scalar source, got VGPR")
# Encode by field type
if name in SRC_FIELDS: self._encode_src(name, val)
elif name in RAW_FIELDS: self._encode_raw(name, val)
elif name == 'sbase': self._values[name] = (val.idx if isinstance(val, Reg) else val.val if isinstance(val, SrcMod) else val * 2) // 2
elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg): self._values[name] = _encode_reg(val) // 4
elif marker is _VDSTYEnc and isinstance(val, VGPR): self._values[name] = val.idx >> 1
self._precompute_fields()
def _encode_field(self, name: str, val) -> int:
if isinstance(val, RawImm): return val.val
if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val # Special regs like VCC_LO
if name in {'srsrc', 'ssamp'}: return _encode_reg(val) // 4 if isinstance(val, Reg) else val
if name == 'sbase': return val.idx // 2 if isinstance(val, Reg) else val.val // 2 if isinstance(val, SrcMod) else val
if name in RAW_FIELDS: return _encode_reg(val) if isinstance(val, Reg) else val
if isinstance(val, Reg) or name in SRC_FIELDS: return encode_src(val)
return val.value if hasattr(val, 'value') else val
def to_int(self) -> int:
word = (self._encoding[1] & self._encoding[0].mask()) << self._encoding[0].lo if self._encoding else 0
for n, bf in self._fields.items():
if n != 'encoding' and n in self._values: word |= (self._encode_field(n, self._values[n]) & bf.mask()) << bf.lo
return word
def _get_literal(self) -> int | None:
for n in SRC_FIELDS:
if n in self._values and not isinstance(v := self._values[n], RawImm) and isinstance(v, int) and not isinstance(v, IntEnum) and not (0 <= v <= 64 or -16 <= v <= -1): return v
return None
def _is_64bit_op(self) -> bool:
"""Check if this instruction uses 64-bit operands (and thus 64-bit literals)."""
op = self._values.get('op')
if op is None: return False
op_name = op.name if hasattr(op, 'name') else None
# Look up op name from int if needed (happens in from_bytes path)
if op_name is None and self.__class__.__name__ == 'VOP3':
try: op_name = VOP3Op(op).name
except ValueError: pass
if op_name is None and self.__class__.__name__ == 'VOPC':
try: op_name = VOPCOp(op).name
except ValueError: pass
if op_name is None: return False
# V_LDEXP_F64 has 32-bit integer src1, so literal is 32-bit
return op_name != 'V_LDEXP_F64' and op_name.endswith(('_F64', '_B64', '_I64', '_U64'))
def to_bytes(self) -> bytes:
result = self.to_int().to_bytes(self._size(), 'little')
lit = self._get_literal() or getattr(self, '_literal', None)
if lit is None: return result
# For 64-bit sources, literal is stored in high 32 bits internally, but encoded as 4 bytes
# Find which source uses the literal (255) and check its register count
lit_src_is_64 = False
for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
if n not in self._values: continue
v = self._values[n]
if (isinstance(v, RawImm) and v.val == 255) or (isinstance(v, int) and v == 255):
lit_src_is_64 = self.is_src_64(idx)
break
lit32 = (lit >> 32) if lit_src_is_64 else lit
return result + (lit32 & MASK32).to_bytes(4, 'little')
@classmethod
def _size(cls) -> int: return cls._sz
def size(self) -> int:
# Literal is always 4 bytes in the binary (for 64-bit ops, it's in high 32 bits)
return self._size() + (4 if self._literal is not None else 0)
@classmethod
def from_int(cls, word: int):
inst = object.__new__(cls)
inst._values = {n: RawImm(v) if n in SRC_FIELDS else v for n, bf in cls._fields.items() if n != 'encoding' for v in [(word >> bf.lo) & bf.mask()]}
inst._literal = None
inst._precompute()
inst._precompute_fields()
return inst
@classmethod
def from_bytes(cls, data: bytes):
import typing
inst = cls.from_int(int.from_bytes(data[:cls._size()], 'little'))
op_val = inst._values.get('op', 0)
# Check for instructions that always have a literal constant (FMAMK/FMAAK/MADMK/MADAK, SETREG_IMM32)
op_name = ''
if cls.__name__ in ('VOP2', 'SOP2', 'SOPK') and 'op' in (hints := typing.get_type_hints(cls, include_extras=True)):
if typing.get_origin(hints['op']) is typing.Annotated:
try: op_name = typing.get_args(hints['op'])[1](op_val).name
except (ValueError, TypeError): pass
has_literal = any(x in op_name for x in ('FMAMK', 'FMAAK', 'MADMK', 'MADAK', 'SETREG_IMM32'))
# VOPD fmaak/fmamk always have a literal (opx/opy value 1 or 2)
opx, opy = inst._values.get('opx', 0), inst._values.get('opy', 0)
has_literal = has_literal or (cls.__name__ == 'VOPD' and (opx in (1, 2) or opy in (1, 2)))
for n in SRC_FIELDS:
if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255: has_literal = True
if has_literal:
# For 64-bit ops, the literal is 32 bits placed in the HIGH 32 bits of the 64-bit value
# (low 32 bits are zero). This is how AMD hardware interprets 32-bit literals for 64-bit ops.
# Check which source uses the literal and whether THAT source is 64-bit
if len(data) >= cls._size() + 4:
lit32 = int.from_bytes(data[cls._size():cls._size()+4], 'little')
# Find which source has literal (255) and check its register count
lit_src_is_64 = False
for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255:
lit_src_is_64 = inst.src_regs(idx) == 2
break
inst._literal = (lit32 << 32) if lit_src_is_64 else lit32
return inst
def __repr__(self):
# Use _fields order and exclude fields that are 0/default (for consistent repr after roundtrip)
def is_zero(v): return (isinstance(v, int) and v == 0) or (isinstance(v, VGPR) and v.idx == 0 and v.count == 1)
items = [(k, self._values[k]) for k in self._fields if k in self._values and k != 'encoding'
and not (is_zero(self._values[k]) and k not in {'op'})]
lit = f", literal={hex(self._literal)}" if self._literal is not None else ""
return f"{self.__class__.__name__}({', '.join(f'{k}={v}' for k, v in items)}{lit})"
def __getattr__(self, name: str):
if name.startswith('_'): raise AttributeError(name)
return unwrap(self._values.get(name, 0))
def lit(self, v: int, neg: bool = False) -> str:
if v == 255 and self._literal is not None:
# For 64-bit sources, literal is stored shifted - extract the 32-bit value
lit32 = (self._literal >> 32) if self._literal > 0xffffffff else self._literal
s = f"0x{lit32:x}"
else:
s = decode_src(v, 'cdna' in self.__class__.__module__)
return f"-{s}" if neg else s
def __eq__(self, other):
if not isinstance(other, Inst): return NotImplemented
return self.__class__ == other.__class__ and self._values == other._values and self._literal == other._literal
def __hash__(self): return hash((self.__class__.__name__, tuple(sorted((k, repr(v)) for k, v in self._values.items())), self._literal))
def disasm(self) -> str:
from extra.assembly.amd.asm import disasm
return disasm(self)
_enum_map = {'VOP1': VOP1Op, 'VOP2': VOP2Op, 'VOP3': VOP3Op, 'VOP3SD': VOP3SDOp, 'VOP3P': VOP3POp, 'VOPC': VOPCOp,
'SOP1': SOP1Op, 'SOP2': SOP2Op, 'SOPC': SOPCOp, 'SOPK': SOPKOp, 'SOPP': SOPPOp,
'SMEM': SMEMOp, 'DS': DSOp, 'FLAT': FLATOp, 'MUBUF': MUBUFOp, 'MTBUF': MTBUFOp, 'MIMG': MIMGOp,
'VOPD': VOPDOp, 'VINTERP': VINTERPOp}
_VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
def _precompute(self):
"""Precompute op, op_name, _spec_regs, _spec_dtype for fast access."""
val = self._values.get('op')
if val is None: self.op = None
elif hasattr(val, 'name'): self.op = val
else:
cls_name = self.__class__.__name__
is_cdna = cls_name in ('VOP3A', 'VOP3B')
# Try marker enum first (VOP3AOp, VOP3BOp, etc.)
marker = self._fields['op'].marker if 'op' in self._fields else None
if marker and issubclass(marker, IntEnum):
try: self.op = marker(val)
except ValueError: self.op = val
elif cls_name in self._enum_map:
try: self.op = self._enum_map[cls_name](val)
except ValueError: self.op = val
else: self.op = val
# Fallback for promoted instructions when marker lookup failed
if not hasattr(self.op, 'name') and cls_name in ('VOP3', 'VOP3A', 'VOP3B') and isinstance(val, int):
if val < 256:
try: self.op = VOPCOp(val)
except ValueError: pass
elif is_cdna and 256 <= val < 512:
try: self.op = (CDNA_VOP1Op(val - 320) if val >= 320 else CDNA_VOP2Op(val - 256))
except ValueError: pass
elif val in self._VOP3SD_OPS and not is_cdna:
try: self.op = VOP3SDOp(val)
except ValueError: pass
elif 256 <= val < 512 and not is_cdna:
try: self.op = VOP1Op(val - 384) if val >= 384 else VOP2Op(val - 256)
except ValueError: pass
self.op_name = self.op.name if hasattr(self.op, 'name') else ''
self._spec_regs = spec_regs(self.op_name)
self._spec_dtype = spec_dtype(self.op_name)
def _precompute_fields(self):
"""Unwrap all field values as direct attributes for fast access."""
for name, val in self._values.items():
if name != 'op': setattr(self, name, unwrap(val))
def dst_regs(self) -> int: return self._spec_regs[0]
def src_regs(self, n: int) -> int: return self._spec_regs[n + 1]
def num_srcs(self) -> int: return spec_num_srcs(self.op_name)
def dst_dtype(self) -> str | None: return self._spec_dtype[0]
def src_dtype(self, n: int) -> str | None: return self._spec_dtype[n + 1]
def is_src_16(self, n: int) -> bool: return self._spec_regs[n + 1] == 1 and is_dtype_16(self._spec_dtype[n + 1])
def is_src_64(self, n: int) -> bool: return self._spec_regs[n + 1] == 2
def is_16bit(self) -> bool: return spec_is_16bit(self.op_name)
def is_64bit(self) -> bool: return spec_is_64bit(self.op_name)
def is_dst_16(self) -> bool: return self._spec_regs[0] == 1 and is_dtype_16(self._spec_dtype[0])
-720
View File
@@ -1,720 +0,0 @@
# RDNA3 emulator - executes compiled pseudocode from AMD ISA PDF
# mypy: ignore-errors
from __future__ import annotations
import ctypes, functools
from tinygrad.helpers import DEBUG, colored, ansilen
from tinygrad.runtime.autogen import hsa
from extra.assembly.amd.dsl import Inst, unwrap, FLOAT_ENC, MASK32, MASK64, _f32, _i32, _sext, _f16, _i16, _f64, _i64, SrcEnum
from extra.assembly.amd.pcode import Reg, compile_pseudocode
from extra.assembly.amd.asm import detect_format, disasm
from extra.assembly.amd.autogen.rdna3.str_pcode import PSEUDOCODE_STRINGS
from extra.assembly.amd.autogen.rdna3.ins import (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD,
SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, SCRATCHOp, VOPDOp)
WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
VCC_LO, VCC_HI, NULL, EXEC_LO, EXEC_HI, SCC = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.SCC
# Inline constants for src operands 128-254. Build tables for f32, f16, and f64 formats.
_FLOAT_CONSTS = {v: k for k, v in FLOAT_ENC.items()} | {248: 0.15915494309189535} # INV_2PI
def _build_inline_consts(mask, to_bits):
tbl = list(range(65)) + [((-i) & mask) for i in range(1, 17)] + [0] * (127 - 81)
for k, v in _FLOAT_CONSTS.items(): tbl[k - 128] = to_bits(v)
return tbl
_INLINE_CONSTS = _build_inline_consts(MASK32, _i32)
_INLINE_CONSTS_F16 = _build_inline_consts(0xffff, _i16)
_INLINE_CONSTS_F64 = _build_inline_consts(MASK64, _i64)
# Helper: extract/write 16-bit half from/to 32-bit value
def _src16(raw: int, is_hi: bool) -> int: return ((raw >> 16) & 0xffff) if is_hi else (raw & 0xffff)
def _dst16(cur: int, val: int, is_hi: bool) -> int: return (cur & 0x0000ffff) | ((val & 0xffff) << 16) if is_hi else (cur & 0xffff0000) | (val & 0xffff)
def _vgpr_hi(src: int) -> bool: return src >= 256 and ((src - 256) & 0x80) != 0
def _vgpr_masked(src: int) -> int: return ((src - 256) & 0x7f) + 256 if src >= 256 else src
# VOP3 source modifier: apply abs/neg to value
def _mod_src(val: int, idx: int, neg: int, abs_: int, is64: bool = False) -> int:
to_f, to_i = (_f64, _i64) if is64 else (_f32, _i32)
if (abs_ >> idx) & 1: val = to_i(abs(to_f(val)))
if (neg >> idx) & 1: val = to_i(-to_f(val))
return val
# Read source operand with VOP3 modifiers
def _read_src(st, inst, src, idx: int, lane: int, neg: int, abs_: int, opsel: int) -> int:
if src is None: return 0
literal, regs, is_src_16 = inst._literal, inst.src_regs(idx), inst.is_src_16(idx)
if regs == 2: return _mod_src(st.rsrc64(src, lane, literal), idx, neg, abs_, is64=True)
if isinstance(inst, VOP3P):
opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
if 'FMA_MIX' in inst.op_name:
raw = st.rsrc(src, lane, literal)
sign_bit = (15 if not (opsel & (1 << idx)) else 31) if (opsel_hi >> idx) & 1 else 31
if inst.neg_hi & (1 << idx): raw &= ~(1 << sign_bit)
if neg & (1 << idx): raw ^= (1 << sign_bit)
return raw
raw = st.rsrc_f16(src, lane, literal)
hi = _src16(raw, opsel_hi & (1 << idx)) ^ (0x8000 if inst.neg_hi & (1 << idx) else 0)
lo = _src16(raw, opsel & (1 << idx)) ^ (0x8000 if neg & (1 << idx) else 0)
return (hi << 16) | lo
if is_src_16 and isinstance(inst, VOP3):
raw = st.rsrc_f16(src, lane, literal) if 128 <= src < 255 else st.rsrc(src, lane, literal)
val = _src16(raw, bool(opsel & (1 << idx)))
if abs_ & (1 << idx): val &= 0x7fff
if neg & (1 << idx): val ^= 0x8000
return val
if is_src_16 and isinstance(inst, (VOP1, VOP2, VOPC)):
if src >= 256: return _src16(_mod_src(st.rsrc(_vgpr_masked(src), lane, literal), idx, neg, abs_), _vgpr_hi(src))
return _mod_src(st.rsrc_f16(src, lane, literal), idx, neg, abs_) & 0xffff
return _mod_src(st.rsrc(src, lane, literal), idx, neg, abs_)
# Helper: get number of dwords from memory op name
def _op_ndwords(name: str) -> int:
if '_B128' in name: return 4
if '_B96' in name: return 3
if any(s in name for s in ('_B64', '_U64', '_I64', '_F64')): return 2
return 1
# Helper: build multi-dword int from consecutive VGPRs
def _vgpr_read(V: list, base: int, ndwords: int) -> int: return sum(V[base + i] << (32 * i) for i in range(ndwords))
# Helper: write multi-dword value to consecutive VGPRs
def _vgpr_write(V: list, base: int, val: int, ndwords: int):
for i in range(ndwords): V[base + i] = (val >> (32 * i)) & MASK32
# Memory access
_valid_mem_ranges: list[tuple[int, int]] = []
def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: _valid_mem_ranges.clear(); _valid_mem_ranges.extend(ranges)
def _mem_valid(addr: int, size: int) -> bool:
return not _valid_mem_ranges or any(s <= addr and addr + size <= s + z for s, z in _valid_mem_ranges)
def _ctypes_at(addr: int, size: int): return (ctypes.c_uint8 if size == 1 else ctypes.c_uint16 if size == 2 else ctypes.c_uint64 if size == 8 else ctypes.c_uint32).from_address(addr)
def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
def mem_write(addr: int, size: int, val: int) -> None:
if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
def _make_mem_accessor(read_fn, write_fn):
"""Create a memory accessor class with the given read/write functions."""
class _MemAccessor:
__slots__ = ('_addr',)
def __init__(self, addr: int): self._addr = int(addr)
u8 = property(lambda s: read_fn(s._addr, 1), lambda s, v: write_fn(s._addr, 1, int(v)))
u16 = property(lambda s: read_fn(s._addr, 2), lambda s, v: write_fn(s._addr, 2, int(v)))
u32 = property(lambda s: read_fn(s._addr, 4), lambda s, v: write_fn(s._addr, 4, int(v)))
u64 = property(lambda s: read_fn(s._addr, 8), lambda s, v: write_fn(s._addr, 8, int(v)))
i8 = property(lambda s: _sext(read_fn(s._addr, 1), 8), lambda s, v: write_fn(s._addr, 1, int(v)))
i16 = property(lambda s: _sext(read_fn(s._addr, 2), 16), lambda s, v: write_fn(s._addr, 2, int(v)))
i32 = property(lambda s: _sext(read_fn(s._addr, 4), 32), lambda s, v: write_fn(s._addr, 4, int(v)))
i64 = property(lambda s: _sext(read_fn(s._addr, 8), 64), lambda s, v: write_fn(s._addr, 8, int(v)))
b8, b16, b32, b64 = u8, u16, u32, u64
return _MemAccessor
_GlobalMemAccessor = _make_mem_accessor(mem_read, mem_write)
class _GlobalMem:
"""Global memory wrapper that supports MEM[addr].u32 style access."""
def __getitem__(self, addr) -> _GlobalMemAccessor: return _GlobalMemAccessor(addr)
GlobalMem = _GlobalMem()
class LDSMem:
"""LDS memory wrapper that supports MEM[addr].u32 style access."""
__slots__ = ('_lds',)
def __init__(self, lds: bytearray): self._lds = lds
def _read(self, addr: int, size: int) -> int:
addr = addr & 0xffff
return int.from_bytes(self._lds[addr:addr+size], 'little') if addr + size <= len(self._lds) else 0
def _write(self, addr: int, size: int, val: int):
addr = addr & 0xffff
if addr + size <= len(self._lds): self._lds[addr:addr+size] = (int(val) & ((1 << (size*8)) - 1)).to_bytes(size, 'little')
def __getitem__(self, addr): return _make_mem_accessor(self._read, self._write)(addr)
# SMEM dst register count (for writing result back to SGPRs)
SMEM_DST_COUNT = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}
# VOPD op -> VOP3 op mapping (VOPD is dual-issue of VOP1/VOP2 ops, use VOP3 enums for pseudocode lookup)
_VOPD_TO_VOP = {
VOPDOp.V_DUAL_FMAC_F32: VOP3Op.V_FMAC_F32, VOPDOp.V_DUAL_FMAAK_F32: VOP2Op.V_FMAAK_F32, VOPDOp.V_DUAL_FMAMK_F32: VOP2Op.V_FMAMK_F32,
VOPDOp.V_DUAL_MUL_F32: VOP3Op.V_MUL_F32, VOPDOp.V_DUAL_ADD_F32: VOP3Op.V_ADD_F32, VOPDOp.V_DUAL_SUB_F32: VOP3Op.V_SUB_F32,
VOPDOp.V_DUAL_SUBREV_F32: VOP3Op.V_SUBREV_F32, VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: VOP3Op.V_MUL_DX9_ZERO_F32,
VOPDOp.V_DUAL_MOV_B32: VOP3Op.V_MOV_B32, VOPDOp.V_DUAL_CNDMASK_B32: VOP3Op.V_CNDMASK_B32,
VOPDOp.V_DUAL_MAX_F32: VOP3Op.V_MAX_F32, VOPDOp.V_DUAL_MIN_F32: VOP3Op.V_MIN_F32,
VOPDOp.V_DUAL_ADD_NC_U32: VOP3Op.V_ADD_NC_U32, VOPDOp.V_DUAL_LSHLREV_B32: VOP3Op.V_LSHLREV_B32, VOPDOp.V_DUAL_AND_B32: VOP3Op.V_AND_B32,
}
class WaveState:
__slots__ = ('sgpr', 'vgpr', 'scc', 'pc', '_pend_sgpr', 'lds', 'n_lanes')
def __init__(self, lds: LDSMem | None = None, n_lanes: int = WAVE_SIZE):
self.sgpr, self.vgpr = [0] * SGPR_COUNT, [[0] * VGPR_COUNT for _ in range(WAVE_SIZE)]
self.sgpr[EXEC_LO], self.scc, self.pc, self._pend_sgpr, self.lds, self.n_lanes = 0xffffffff, 0, 0, {}, lds, n_lanes
@property
def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
@vcc.setter
def vcc(self, v: int): self.sgpr[VCC_LO], self.sgpr[VCC_HI] = v & MASK32, (v >> 32) & MASK32
@property
def exec_mask(self) -> int: return self.sgpr[EXEC_LO] | (self.sgpr[EXEC_HI] << 32)
@exec_mask.setter
def exec_mask(self, v: int): self.sgpr[EXEC_LO], self.sgpr[EXEC_HI] = v & MASK32, (v >> 32) & MASK32
def rsgpr(self, i: int) -> int: return 0 if i == NULL else self.scc if i == SCC else self.sgpr[i] if i < SGPR_COUNT else 0
def wsgpr(self, i: int, v: int):
if i < SGPR_COUNT and i != NULL: self.sgpr[i] = v & MASK32
def rsgpr64(self, i: int) -> int: return self.rsgpr(i) | (self.rsgpr(i+1) << 32)
def wsgpr64(self, i: int, v: int): self.wsgpr(i, v & MASK32); self.wsgpr(i+1, (v >> 32) & MASK32)
def _rsrc_base(self, v: int, lane: int, consts, literal: int):
if v < SGPR_COUNT: return self.sgpr[v]
if v == SCC: return self.scc
if v < 255: return consts[v - 128]
if v == 255: return literal
return self.vgpr[lane][v - 256] if v <= 511 else 0
def rsrc(self, v: int, lane: int, literal: int = 0) -> int: return self._rsrc_base(v, lane, _INLINE_CONSTS, literal)
def rsrc_f16(self, v: int, lane: int, literal: int = 0) -> int: return self._rsrc_base(v, lane, _INLINE_CONSTS_F16, literal)
def rsrc64(self, v: int, lane: int, literal: int = 0) -> int:
if 128 <= v < 255: return _INLINE_CONSTS_F64[v - 128]
if v == 255: return literal # literal is already shifted in from_bytes for 64-bit ops
return self.rsrc(v, lane, literal) | ((self.rsrc(v+1, lane, literal) if v < VCC_LO or 256 <= v <= 511 else 0) << 32)
def pend_sgpr_lane(self, reg: int, lane: int, val: int):
if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
if val: self._pend_sgpr[reg] |= (1 << lane)
def commit_pends(self):
for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
self._pend_sgpr.clear()
# ═══════════════════════════════════════════════════════════════════════════════
# EXECUTION - All ops use pseudocode from PDF
# ═══════════════════════════════════════════════════════════════════════════════
def exec_scalar(st: WaveState, inst: Inst):
"""Execute scalar instruction. Returns 0 to continue execution."""
# Get op enum and lookup compiled function
if isinstance(inst, SMEM): ssrc0, sdst = None, None
elif isinstance(inst, SOP1): ssrc0, sdst = inst.ssrc0, inst.sdst
elif isinstance(inst, SOP2): ssrc0, sdst = inst.ssrc0, inst.sdst
elif isinstance(inst, SOPC): ssrc0, sdst = inst.ssrc0, None
elif isinstance(inst, SOPK): ssrc0, sdst = inst.sdst, inst.sdst # sdst is both src and dst
elif isinstance(inst, SOPP): ssrc0, sdst = None, None
else: raise NotImplementedError(f"Unknown scalar type {type(inst)}")
# SMEM: memory loads
if isinstance(inst, SMEM):
addr = st.rsgpr64(inst.sbase * 2) + _sext(inst.offset, 21)
if inst.soffset not in (NULL, 0x7f): addr += st.rsrc(inst.soffset, 0, inst._literal)
result = inst._fn(GlobalMem, addr & MASK64)
if 'SDATA' in result:
sdata = result['SDATA']
for i in range(SMEM_DST_COUNT.get(inst.op, 1)): st.wsgpr(inst.sdata + i, (sdata >> (i * 32)) & MASK32)
st.pc += inst._words
return 0
# Build context - use inst methods to determine operand sizes
literal = inst._literal
s0 = st.rsrc64(ssrc0, 0, literal) if inst.is_src_64(0) else (st.rsrc(ssrc0, 0, literal) if not isinstance(inst, (SOPK, SOPP)) else (st.rsgpr(inst.sdst) if isinstance(inst, SOPK) else 0))
s1 = st.rsrc64(inst.ssrc1, 0, literal) if inst.is_src_64(1) else (st.rsrc(inst.ssrc1, 0, literal) if isinstance(inst, (SOP2, SOPC)) else inst.simm16 if isinstance(inst, SOPK) else 0)
d0 = st.rsgpr64(sdst) if inst.dst_regs() == 2 and sdst is not None else (st.rsgpr(sdst) if sdst is not None else 0)
literal = inst.simm16 if isinstance(inst, (SOPK, SOPP)) else inst._literal
# Call compiled function with int parameters
result = inst._fn(s0, s1, 0, d0, st.scc, st.vcc & MASK32, 0, st.exec_mask & MASK32, literal, None, pc=st.pc * 4)
# Apply results (already int values)
if sdst is not None and 'D0' in result:
(st.wsgpr64 if inst.dst_regs() == 2 else st.wsgpr)(sdst, result['D0'])
if 'SCC' in result: st.scc = result['SCC'] & 1
if 'EXEC' in result: st.exec_mask = result['EXEC']
if 'PC' in result:
# Convert absolute byte address to word offset
pc_val = result['PC']
new_pc = pc_val if pc_val < 0x8000000000000000 else pc_val - 0x10000000000000000
st.pc = new_pc // 4
else:
st.pc += inst._words
return 0
# ═══════════════════════════════════════════════════════════════════════════════
# VECTOR INSTRUCTIONS
# ═══════════════════════════════════════════════════════════════════════════════
def exec_vopd(st: WaveState, inst, V: list, lane: int) -> None:
"""VOPD: dual-issue, execute two ops simultaneously (read all inputs before writes)."""
literal, vdstx, vdsty = inst._literal, inst.vdstx, (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1)
sx0, sx1, dx, sy0, sy1, dy = st.rsrc(inst.srcx0, lane, literal), V[inst.vsrcx1], V[vdstx], st.rsrc(inst.srcy0, lane, literal), V[inst.vsrcy1], V[vdsty]
V[vdstx] = inst._fnx(sx0, sx1, 0, dx, st.scc, st.vcc, lane, st.exec_mask, literal, None)['D0']
V[vdsty] = inst._fny(sy0, sy1, 0, dy, st.scc, st.vcc, lane, st.exec_mask, literal, None)['D0']
def exec_flat(st: WaveState, inst, V: list, lane: int) -> None:
"""FLAT/GLOBAL/SCRATCH memory ops."""
ndwords = _op_ndwords(inst.op_name)
addr = V[inst.addr] | (V[inst.addr + 1] << 32)
ADDR = (st.rsgpr64(inst.saddr) + V[inst.addr] + _sext(inst.offset, 13)) & MASK64 if inst.saddr not in (NULL, 0x7f) else (addr + _sext(inst.offset, 13)) & MASK64
vdata_src = inst.vdst if 'LOAD' in inst.op_name else inst.data
result = inst._fn(GlobalMem, ADDR, _vgpr_read(V, vdata_src, ndwords), V[inst.vdst])
if 'VDATA' in result: _vgpr_write(V, inst.vdst, result['VDATA'], ndwords)
if 'RETURN_DATA' in result: _vgpr_write(V, inst.vdst, result['RETURN_DATA'], ndwords)
def exec_ds(st: WaveState, inst, V: list, lane: int) -> None:
"""DS (LDS) memory ops."""
ndwords = _op_ndwords(inst.op_name)
data0, data1 = _vgpr_read(V, inst.data0, ndwords), _vgpr_read(V, inst.data1, ndwords) if inst.data1 is not None else 0
result = inst._fn(st.lds, V[inst.addr], data0, data1, inst.offset0, inst.offset1)
if 'RETURN_DATA' in result and ('_RTN' in inst.op_name or '_LOAD' in inst.op_name):
_vgpr_write(V, inst.vdst, result['RETURN_DATA'], ndwords * 2 if '_2ADDR_' in inst.op_name else ndwords)
def exec_vop(st: WaveState, inst: Inst, V: list, lane: int) -> None:
"""VOP1/VOP2/VOP3/VOP3SD/VOP3P/VOPC: standard ALU ops."""
if isinstance(inst, VOP3P):
src0, src1, src2, vdst, dst_hi = inst.src0, inst.src1, inst.src2, inst.vdst, False
neg, abs_, opsel = inst.neg, 0, inst.opsel
elif isinstance(inst, VOP1):
src0, src1, src2, vdst = inst.src0, None, None, inst.vdst & 0x7f if inst.is_dst_16() else inst.vdst
neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst & 0x80) != 0 and inst.is_dst_16()
elif isinstance(inst, VOP2):
src0, src1, src2, vdst = inst.src0, inst.vsrc1 + 256, None, inst.vdst & 0x7f if inst.is_dst_16() else inst.vdst
neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst & 0x80) != 0 and inst.is_dst_16()
elif isinstance(inst, (VOP3, VOP3SD)):
src0, src1, src2, vdst = inst.src0, inst.src1, (None if isinstance(inst, VOP3) and inst.op.value < 256 else inst.src2), inst.vdst
neg, abs_, opsel, dst_hi = (inst.neg, inst.abs, inst.opsel, False) if isinstance(inst, VOP3) else (0, 0, 0, False)
elif isinstance(inst, VOPC):
src0, src1, src2, vdst, neg, abs_, opsel, dst_hi = inst.src0, inst.vsrc1 + 256, None, VCC_LO, 0, 0, 0, False
else:
raise NotImplementedError(f"exec_vop: unhandled instruction type {type(inst).__name__}")
s0 = _read_src(st, inst, src0, 0, lane, neg, abs_, opsel)
s1 = _read_src(st, inst, src1, 1, lane, neg, abs_, opsel)
s2 = _read_src(st, inst, src2, 2, lane, neg, abs_, opsel)
if isinstance(inst, VOP2) and inst.is_16bit(): d0 = _src16(V[vdst], dst_hi)
elif inst.dst_regs() == 2: d0 = V[vdst] | (V[vdst + 1] << 32)
else: d0 = V[vdst]
if isinstance(inst, VOP3SD) and 'CO_CI' in inst.op_name: vcc_for_fn = st.rsgpr64(inst.src2)
elif isinstance(inst, VOP3) and inst.op in (VOP3Op.V_CNDMASK_B32, VOP3Op.V_CNDMASK_B16) and src2 is not None and src2 < 256: vcc_for_fn = st.rsgpr64(src2)
else: vcc_for_fn = st.vcc
src0_idx = (src0 - 256) if src0 is not None and src0 >= 256 else (src0 if src0 is not None else 0)
extra_kwargs = {'opsel': opsel, 'opsel_hi': inst.opsel_hi | (inst.opsel_hi2 << 2)} if isinstance(inst, VOP3P) and 'FMA_MIX' in inst.op_name else {}
result = inst._fn(s0, s1, s2, d0, st.scc, vcc_for_fn, lane, st.exec_mask, inst._literal, st.vgpr, src0_idx, vdst, **extra_kwargs)
# Check if this is a VOPC instruction (either standalone VOPC or VOP3 with VOPC opcode)
is_vopc = isinstance(inst.op, VOPCOp) or (isinstance(inst, VOP3) and inst.op.value < 256)
if 'VCC' in result:
if isinstance(inst, VOP3SD): st.pend_sgpr_lane(inst.sdst, lane, (result['VCC'] >> lane) & 1)
else: st.pend_sgpr_lane(VCC_LO if isinstance(inst, VOP2) and 'CO_CI' in inst.op_name else vdst, lane, (result['VCC'] >> lane) & 1)
if 'EXEC' in result:
st.pend_sgpr_lane(EXEC_LO, lane, (result['EXEC'] >> lane) & 1)
elif is_vopc:
st.pend_sgpr_lane(vdst, lane, (result['D0'] >> lane) & 1)
if not is_vopc:
d0_val = result['D0']
if inst.dst_regs() == 2: V[vdst], V[vdst + 1] = d0_val & MASK32, (d0_val >> 32) & MASK32
elif not isinstance(inst, VOP3P) and inst.is_dst_16(): V[vdst] = _dst16(V[vdst], d0_val, bool(opsel & 8) if isinstance(inst, VOP3) else dst_hi)
else: V[vdst] = d0_val & MASK32
# ═══════════════════════════════════════════════════════════════════════════════
# WMMA (Wave Matrix Multiply-Accumulate)
# ═══════════════════════════════════════════════════════════════════════════════
def exec_wmma(st: WaveState, inst, op: VOP3POp) -> None:
"""Execute WMMA instruction - 16x16x16 matrix multiply across the wave."""
src0, src1, src2, vdst = inst.src0, inst.src1, inst.src2, inst.vdst
# Read 16x16 f16 matrix from 16 lanes × 8 VGPRs (2 f16 per VGPR)
def read_f16_mat(src):
return [f for l in range(16) for r in range(8) for v in [st.vgpr[l][src-256+r] if src >= 256 else st.rsgpr(src+r)] for f in [_f16(v&0xffff), _f16((v>>16)&0xffff)]]
mat_a, mat_b = read_f16_mat(src0), read_f16_mat(src1)
# Read matrix C (16x16 f32) from lanes 0-31, VGPRs src2 to src2+7
mat_c = [_f32(st.vgpr[i % 32][src2 - 256 + i // 32] if src2 >= 256 else st.rsgpr(src2 + i // 32)) for i in range(256)]
# Compute D = A × B + C (16x16 matrix multiply)
mat_d = [sum(mat_a[row*16+k] * mat_b[col*16+k] for k in range(16)) + mat_c[row*16+col] for row in range(16) for col in range(16)]
# Write result - f16 packed or f32
if op == VOP3POp.V_WMMA_F16_16X16X16_F16:
for i in range(0, 256, 2):
st.vgpr[(i//2) % 32][vdst + (i//2)//32] = ((_i16(mat_d[i+1]) & 0xffff) << 16) | (_i16(mat_d[i]) & 0xffff)
else:
for i in range(256): st.vgpr[i % 32][vdst + i//32] = _i32(mat_d[i])
# SQTT TRACING
# ═══════════════════════════════════════════════════════════════════════════════
WAVESTART_TO_INST_CYCLES = 32
SNOP_EXTRA_DELAY_MIN, SNOP_EXTRA_DELAY_MAX = 11, 22 # s_nop(11-22) has +4 penalty
SNOP_EXTRA_DELAY_CYCLES = 4
from extra.assembly.amd.sqtt import WAVESTART, WAVEEND, IMMEDIATE, VALUINST, ALUEXEC, AluSrc
def _get_src_vgprs(inst: Inst) -> list[int]:
if isinstance(inst, VOP1): return [inst.src0 - 256] if inst.src0 >= 256 else []
if isinstance(inst, VOP2): return ([inst.src0 - 256] if inst.src0 >= 256 else []) + [inst.vsrc1]
if isinstance(inst, VOP3): return [s - 256 for s in [inst.src0, inst.src1, getattr(inst, 'src2', None)] if s is not None and s >= 256]
return []
class SQTTState:
"""SQTT tracing with cycle-accurate RDNA3 VALU pipeline model.
NOTE: This is a hardware-plausible model derived from observed SQTT timing patterns.
The model should be verified by tests against real hardware traces, not by fitting
formulas to expected outputs. If tests fail, the model needs to be understood and
fixed, not hacked with magic constants.
Physical model:
- alu[4]: 4-stage ALU pipeline, each slot holds dest_vgpr or None
- in_flight: up to 12 in-flight instructions (issued but not yet completed)
- issue_queue: instructions waiting to enter ALU (sources not ready)
- fwd_slots: 4 forwarding slots, reserved at issue, freed when consumer forwards
- completed: vgprs with results ready (exited ALU)
Forwarding model (4 slots):
- Slot reserved at ISSUE time if available (len(fwd_slots) < 4)
- Slot freed when a consumer uses the result for forwarding
- Consumer can forward if: has a slot AND producer is completed
- If no slot at issue, instruction uses regfile path (+4 cycle penalty)
"""
def __init__(self, wave_id: int = 0, simd: int = 0, cu: int = 0):
self.wave_id, self.simd, self.cu = wave_id, simd, cu
self.cycle = 0
self.packets = []
# 4-stage ALU pipeline: each slot holds dest_vgpr or None
self.alu = [None, None, None, None]
# In-flight instructions: max 12 at a time, each is (dest_vgpr, srcs, has_fwd_slot)
self.in_flight: list[tuple[int, list[int], bool]] = []
# Issue queue: list of (dest_vgpr, srcs, ready_at, has_fwd_slot, was_warm) waiting for deps
# ready_at: cycle when this instruction can enter ALU (0 = no restriction)
# has_fwd_slot: True if this instruction reserved a forwarding slot at issue time
# was_warm: True if forwarding path was warm when this instruction was issued
self.issue_queue: list[tuple[int, list[int], int, bool, bool]] = []
# 4 forwarding slots: consumer adds producer at issue, freed when consumer forwards
self.fwd_slots: list[int] = [] # producer vgprs reserved for forwarding
# VGPRs that had a dependent try to add them to fwd_slots (successful or not)
self.had_dependent: set[int] = set()
# VGPRs that were issued after forwarding chain broke (can't forward)
self.fwd_chain_broken: set[int] = set()
# Set of completed vgprs (results ready, exited ALU)
self.completed: set[int] = set()
# Cold start: first forwarding use has +1 cycle penalty
self.forward_warm = False
self.cold_used = False # True if cold start penalty was applied
def emit(self, pkt_class, **kwargs):
self.packets.append(pkt_class(_time=self.cycle, **kwargs))
def _fmt_alu(self) -> str:
# Fixed width: each slot 3 chars, total ALU[xxx,xxx,xxx,xxx] = 20 chars
slots = [f'v{v}' if v is not None else '-' for v in self.alu]
return 'ALU[' + ','.join(f'{s:>3}' for s in slots) + ']'
def _fmt_fwd(self) -> str:
items = [f'v{v}' for v in self.fwd_slots]
content = 'FWD[' + ','.join(items) + ']' if items else 'FWD[]'
padded = f'{content:<24}'
return colored(padded, 'yellow') if items else padded
def _fmt_iq(self) -> str:
def fmt_item(d, r, fwd):
s = f'v{d}'
if r != 0: s += f'@{abs(r)}'
if not fwd: s += 'R'
return s
items = [fmt_item(d, r, fwd) for d, _, r, fwd, _ in self.issue_queue]
return 'IQ[' + ','.join(items) + ']' if items else 'IQ[]'
def _debug_line(self, events: list[str] | None = None):
if DEBUG < 3: return
# Skip empty cycles (nothing in ALU, no events, no IQ)
has_alu = any(s is not None for s in self.alu)
if not has_alu and not events and not self.issue_queue: return
cycle = colored(f'C{self.cycle:>3}:', 'cyan')
alu = self._fmt_alu()
fwd = self._fmt_fwd()
iq = f'{self._fmt_iq():<28}'
ev_str = ' '.join(events) if events else ''
ev_padded = f'{ev_str:<20}' if ev_str else ' ' * 20
print(f"{cycle} {alu} {fwd} {iq} {ev_padded}")
def _can_issue(self) -> bool:
return len(self.in_flight) < 12
def _has_pending_write(self, vgpr: int) -> bool:
"""Check if there's a pending write to this VGPR (in ALU, in-flight, or issue queue)."""
if any(slot == vgpr for slot in self.alu if slot is not None): return True
if any(d == vgpr for d, _, _ in self.in_flight): return True
if any(d == vgpr for d, _, _, _, _ in self.issue_queue): return True
return False
def _all_srcs_ready(self, srcs: list[int]) -> bool:
"""Returns True if all sources are ready (completed or no pending write)."""
for src in srcs:
if src in self.completed: continue
if not self._has_pending_write(src): continue # initial value
return False
return True
def tick(self):
self.cycle += 1
if self.cycle > 10000: raise RuntimeError("cycle limit exceeded")
events = []
# 1. ALU[3] exits - capture but don't add to completed yet
exiting = self.alu[3]
if exiting is not None:
self.emit(ALUEXEC, src=AluSrc.VALU)
events.append(colored(f"EXEC v{exiting}", 'red'))
# 2. Slide ALU pipeline
self.alu[3] = self.alu[2]
self.alu[2] = self.alu[1]
self.alu[1] = self.alu[0]
self.alu[0] = None
# 3. Try to promote from issue_queue to ALU[0] (before adding exiting to completed)
if self.alu[0] is None and self.issue_queue:
for i, (dest, srcs, ready_at, has_fwd_slot, was_warm) in enumerate(self.issue_queue):
# Check if instruction has a minimum ready cycle
if ready_at > 0 and self.cycle < ready_at:
continue
# Check if sources are ready
ready = self._all_srcs_ready(srcs)
has_deps = len(srcs) > 0
if not ready:
continue
# Cold start penalty: first dependent instruction has +1 cycle delay (delta=6 vs delta=5)
# Only applies if forwarding path wasn't warm when this instruction was issued
if has_deps and not was_warm and not self.cold_used:
self.cold_used = True
self.issue_queue[i] = (dest, srcs, self.cycle + 1, has_fwd_slot, was_warm)
continue
# Forwarding: consumer can forward if:
# 1. Not in fwd_chain_broken (chain must be intact), AND
# 2. Producer has a slot (source is in fwd_slots), AND
# 3. Either activated by dependent OR successfully added producer at issue
# Note: if issued cold with no slot, activation only counts if the activator also has a dependent
chain_intact = dest not in self.fwd_chain_broken
producer_has_slot = has_deps and any(src in self.fwd_slots for src in srcs)
# Check activation validity
if dest in self.had_dependent:
if was_warm or has_fwd_slot:
activated_by_dependent = True
else:
# Cold + no slot: activation only counts if activator itself has a dependent
# This handles the chain_6 vs chain_7 difference (chain_7 has v6 which activates v5)
activated_by_dependent = (dest + 1) in self.had_dependent # activator is dest+1 in a chain
else:
activated_by_dependent = False
can_forward = chain_intact and producer_has_slot and (activated_by_dependent or has_fwd_slot)
# Regfile path: has dependencies but can't forward
must_use_regfile = has_deps and not can_forward
# Regfile penalty: add +4 cycles latency (only apply once)
if must_use_regfile and ready_at == 0:
self.issue_queue[i] = (dest, srcs, self.cycle + 4, has_fwd_slot, was_warm)
continue
# Enter ALU
self.alu[0] = dest
self.issue_queue.pop(i)
# Free producer's forwarding slot when consumer dispatches (regardless of fwd/rf)
for src in srcs:
if src in self.fwd_slots:
self.fwd_slots.remove(src)
break
events.append(colored(f"v{dest}->ALU" + ("(fwd)" if can_forward else "(rf)" if must_use_regfile else ""), 'green'))
break
# 4. Now add exiting instruction to completed (after promotion decision)
if exiting is not None:
self.completed.add(exiting)
# Remove from in_flight - any VALU completing warms up the forward path
for idx, (d, _, _) in enumerate(self.in_flight):
if d == exiting:
self.forward_warm = True
self.in_flight.pop(idx)
break
self._debug_line(events)
def _pipeline_empty(self) -> bool:
if any(s is not None for s in self.alu): return False
if self.issue_queue: return False
if self.in_flight: return False
return True
def process_instruction(self, inst: Inst):
if isinstance(inst, SOPP) and inst.op == SOPPOp.S_DELAY_ALU:
# TODO: implement s_delay_alu properly
return
elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_NOP:
# s_nop(N) delays N+1 cycles, plus extra penalty for s_nop(11-22)
cycles = inst.simm16 + 1
if SNOP_EXTRA_DELAY_MIN <= inst.simm16 <= SNOP_EXTRA_DELAY_MAX:
cycles += SNOP_EXTRA_DELAY_CYCLES
if DEBUG >= 3:
cycle = colored(f'C{self.cycle:>3}:', 'cyan')
# 20 (ALU) + 1 + 24 (FWD) + 1 + 28 (IQ) + 1 + 20 (events) = 95 padding after cycle
print(f"{cycle} {' ' * 95} {disasm(inst)}")
for _ in range(cycles): self.tick()
self.emit(IMMEDIATE, wave=self.wave_id)
elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_ENDPGM:
# Drain pipeline before ending
while not self._pipeline_empty(): self.tick()
self.emit(WAVEEND, wave=self.wave_id, simd=self.simd, cu_lo=self.cu & 0x7, flag7=self.cu >> 3)
elif isinstance(inst, (VOP1, VOP2, VOP3)):
# Check for issue stall (no free in-flight slots)
while not self._can_issue():
self.tick()
# Issue: add to in_flight and issue_queue
srcs = _get_src_vgprs(inst)
dest = inst.vdst
# Clear stale state for this dest (WAW hazard)
self.completed.discard(dest)
if dest in self.fwd_slots: self.fwd_slots.remove(dest)
# Consumer adds producer to fwd_slots (if room and has dependency)
# If producer is in fwd_chain_broken, or we can't add, the chain breaks for this instruction too
has_fwd_slot = False
if srcs:
producer = srcs[0]
self.had_dependent.add(producer) # record that producer has a dependent
# Check if producer's forwarding chain is already broken
if producer in self.fwd_chain_broken:
# Chain is broken, this instruction also can't forward
self.fwd_chain_broken.add(dest)
elif len(self.fwd_slots) >= 4:
# Can't add producer, chain breaks
self.fwd_chain_broken.add(dest)
else:
# Can add producer
if producer not in self.fwd_slots:
self.fwd_slots.append(producer)
has_fwd_slot = len(self.fwd_slots) < 4
# Record if forwarding path was warm at issue time
was_warm = self.forward_warm
self.in_flight.append((dest, srcs, has_fwd_slot))
self.issue_queue.append((dest, srcs, 0, has_fwd_slot, was_warm))
self.emit(VALUINST, wave=self.wave_id)
if DEBUG >= 3:
cycle = colored(f'C{self.cycle:>3}:', 'cyan')
slot_info = "" if has_fwd_slot else colored(" NO_SLOT", 'red')
issue = colored(f'ISSUE v{dest}', 'magenta') + slot_info
padding = 95 - ansilen(issue)
print(f"{cycle} {issue}{' ' * padding} {disasm(inst)}")
# One cycle per instruction issued, then try to enter ALU
self.tick()
return
# One cycle per instruction issued (for non-VALU)
self.tick()
def emit_wavestart(self):
self.emit(WAVESTART, wave=self.wave_id, simd=self.simd, cu_lo=self.cu & 0x7, flag7=self.cu >> 3)
for _ in range(WAVESTART_TO_INST_CYCLES): self.tick()
# ═══════════════════════════════════════════════════════════════════════════════
# PROGRAM DECODE
# ═══════════════════════════════════════════════════════════════════════════════
# Wave-level dispatch functions: (st, inst) -> return_code (0 = continue, -1 = end, -2 = barrier)
def dispatch_endpgm(st, inst): return -1
def dispatch_barrier(st, inst): st.pc += inst._words; return -2
def dispatch_nop(st, inst): st.pc += inst._words; return 0
def dispatch_wmma(st, inst): exec_wmma(st, inst, inst.op); st.pc += inst._words; return 0
def dispatch_writelane(st, inst): st.vgpr[st.rsrc(inst.src1, 0, inst._literal) & 0x1f][inst.vdst] = st.rsrc(inst.src0, 0, inst._literal) & MASK32; st.pc += inst._words; return 0
def dispatch_readlane(st, inst):
src0_idx = (inst.src0 - 256) if inst.src0 >= 256 else inst.src0
s1 = st.rsrc(inst.src1, 0, inst._literal) if getattr(inst, 'src1', None) is not None else 0
result = inst._fn(0, s1, 0, 0, st.scc, st.vcc, 0, st.exec_mask, inst._literal, st.vgpr, src0_idx, inst.vdst)
st.wsgpr(inst.vdst, result['D0'])
st.pc += inst._words; return 0
# Per-lane dispatch wrapper: wraps per-lane exec functions into wave-level dispatch
@functools.cache
def dispatch_lane(exec_fn):
def dispatch(st, inst):
exec_mask, vgpr, n_lanes = st.exec_mask, st.vgpr, st.n_lanes
for lane in range(n_lanes):
if exec_mask >> lane & 1: exec_fn(st, inst, vgpr[lane], lane)
st.commit_pends()
st.pc += inst._words
return 0
return dispatch
def decode_program(data: bytes) -> dict[int, Inst]:
result: dict[int, Inst] = {}
i = 0
while i < len(data):
inst = detect_format(data[i:]).from_bytes(data[i:])
inst._words = inst.size() // 4
# Determine dispatch function and pcode function
if isinstance(inst, SOPP) and inst.op == SOPPOp.S_CODE_END: break
elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_ENDPGM: inst._dispatch = dispatch_endpgm
elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER: inst._dispatch = dispatch_barrier
elif isinstance(inst, SOPP) and inst.op in (SOPPOp.S_CLAUSE, SOPPOp.S_WAITCNT, SOPPOp.S_WAITCNT_DEPCTR, SOPPOp.S_SENDMSG, SOPPOp.S_SET_INST_PREFETCH_DISTANCE, SOPPOp.S_DELAY_ALU): inst._dispatch = dispatch_nop
elif isinstance(inst, (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM)): inst._dispatch = exec_scalar
elif isinstance(inst, VOP1) and inst.op == VOP1Op.V_NOP: inst._dispatch = dispatch_nop
elif isinstance(inst, VOP3P) and 'WMMA' in inst.op_name: inst._dispatch = dispatch_wmma
elif isinstance(inst, VOP3) and inst.op == VOP3Op.V_WRITELANE_B32: inst._dispatch = dispatch_writelane
elif isinstance(inst, (VOP1, VOP3)) and inst.op in (VOP1Op.V_READFIRSTLANE_B32, VOP3Op.V_READFIRSTLANE_B32, VOP3Op.V_READLANE_B32): inst._dispatch = dispatch_readlane
elif isinstance(inst, VOPD): inst._dispatch = dispatch_lane(exec_vopd)
elif isinstance(inst, FLAT): inst._dispatch = dispatch_lane(exec_flat)
elif isinstance(inst, DS): inst._dispatch = dispatch_lane(exec_ds)
else: inst._dispatch = dispatch_lane(exec_vop)
# Compile pcode for instructions that use it (not VOPD which has _fnx/_fny, not special dispatches)
# VOPD needs separate functions for X and Y ops
if isinstance(inst, VOPD):
def _compile_vopd_op(op): return compile_pseudocode(type(op).__name__, op.name, PSEUDOCODE_STRINGS[type(op)][op])
inst._fnx, inst._fny = _compile_vopd_op(_VOPD_TO_VOP[inst.opx]), _compile_vopd_op(_VOPD_TO_VOP[inst.opy])
elif inst._dispatch not in (dispatch_endpgm, dispatch_barrier, dispatch_nop, dispatch_wmma, dispatch_writelane):
assert type(inst.op) != int, f"inst op of {inst} is int"
inst._fn = compile_pseudocode(type(inst.op).__name__, inst.op.name, PSEUDOCODE_STRINGS[type(inst.op)][inst.op])
result[i // 4] = inst
i += inst._words * 4
return result
# ═══════════════════════════════════════════════════════════════════════════════
# MAIN EXECUTION LOOP
# ═══════════════════════════════════════════════════════════════════════════════
def exec_wave(program: dict[int, Inst], st: WaveState) -> int:
while (inst := program.get(st.pc)) and (result := inst._dispatch(st, inst)) == 0: pass
return result
def exec_workgroup(program: dict[int, Inst], workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int, rsrc2: int) -> None:
lx, ly, lz = local_size
total_threads = lx * ly * lz
# GRANULATED_LDS_SIZE is in 512-byte units (see ops_amd.py: lds_size = ((group_segment_size + 511) // 512))
lds_size = ((rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE_SHIFT) * 512
lds = LDSMem(bytearray(lds_size)) if lds_size else None
waves: list[WaveState] = []
for wave_start in range(0, total_threads, WAVE_SIZE):
n_lanes = min(WAVE_SIZE, total_threads - wave_start)
st = WaveState(lds, n_lanes)
st.exec_mask = (1 << n_lanes) - 1
st.wsgpr64(0, args_ptr) # s[0:1] = kernel arguments pointer
# COMPUTE_PGM_RSRC2: USER_SGPR_COUNT is where workgroup IDs start, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z control which are passed
sgpr_idx = (rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT_SHIFT
if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X: st.sgpr[sgpr_idx] = workgroup_id[0]; sgpr_idx += 1
if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Y: st.sgpr[sgpr_idx] = workgroup_id[1]; sgpr_idx += 1
if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Z: st.sgpr[sgpr_idx] = workgroup_id[2]
# VGPR0 = packed workitem IDs: (Z << 20) | (Y << 10) | X
for tid in range(wave_start, wave_start + n_lanes):
st.vgpr[tid - wave_start][0] = ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
waves.append(st)
while waves:
waves = [st for st in waves if exec_wave(program, st) != -1]
def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int, rsrc2: int = 0x19c) -> int:
program = decode_program((ctypes.c_char * lib_sz).from_address(lib).raw)
for gidz in range(gz):
for gidy in range(gy):
for gidx in range(gx): exec_workgroup(program, (gidx, gidy, gidz), (lx, ly, lz), args_ptr, rsrc2)
return 0
-787
View File
@@ -1,787 +0,0 @@
# DSL for RDNA3 pseudocode - makes pseudocode expressions work directly as Python
import struct, math, re, functools
from extra.assembly.amd.dsl import MASK32, MASK64, _f32, _i32, _sext, _f16, _i16, _f64, _i64
# ═══════════════════════════════════════════════════════════════════════════════
# INTERNAL HELPERS
# ═══════════════════════════════════════════════════════════════════════════════
def _div(a, b):
try: return a / b
except ZeroDivisionError:
if a == 0.0 or math.isnan(a): return float("nan")
return math.copysign(float("inf"), a * b) if b == 0.0 else float("inf") if a > 0 else float("-inf")
def _check_nan_type(x, quiet_bit_expected, default):
try:
if not math.isnan(float(x)): return False
if hasattr(x, '_reg') and hasattr(x, '_bits'):
bits = x._reg._val & ((1 << x._bits) - 1)
exp_bits, quiet_pos, mant_mask = {16: (0x1f, 9, 0x3ff), 32: (0xff, 22, 0x7fffff), 64: (0x7ff, 51, 0xfffffffffffff)}.get(x._bits, (0,0,0))
exp_shift = {16: 10, 32: 23, 64: 52}.get(x._bits, 0)
if exp_bits and ((bits >> exp_shift) & exp_bits) == exp_bits and (bits & mant_mask) != 0:
return ((bits >> quiet_pos) & 1) == quiet_bit_expected
return default
except (TypeError, ValueError): return False
def _gt_neg_zero(a, b): return (a > b) or (a == 0 and b == 0 and not math.copysign(1, a) < 0 and math.copysign(1, b) < 0)
def _lt_neg_zero(a, b): return (a < b) or (a == 0 and b == 0 and math.copysign(1, a) < 0 and not math.copysign(1, b) < 0)
def _fpop(fn):
def wrapper(x):
x = float(x)
if math.isnan(x) or math.isinf(x): return x
result = float(fn(x))
return math.copysign(0.0, x) if result == 0.0 else result
return wrapper
def _f_to_int(f, lo, hi): f = float(f); return 0 if math.isnan(f) else (hi if f >= hi else lo if f <= lo else int(f))
def _f16_to_f32_bits(bits): return struct.unpack("<e", struct.pack("<H", int(bits) & 0xffff))[0]
def _brev(v, bits): return int(bin(v & ((1 << bits) - 1))[2:].zfill(bits)[::-1], 2)
def _ctz(v, bits):
v, n = int(v) & ((1 << bits) - 1), 0
if v == 0: return bits
while (v & 1) == 0: v >>= 1; n += 1
return n
def _bf16(i):
"""Convert bf16 bits to float. BF16 is just the top 16 bits of f32."""
return struct.unpack("<f", struct.pack("<I", (i & 0xffff) << 16))[0]
def _ibf16(f):
"""Convert float to bf16 bits (truncate to top 16 bits of f32)."""
if math.isnan(f): return 0x7fc0 # bf16 quiet NaN
if math.isinf(f): return 0x7f80 if f > 0 else 0xff80 # bf16 ±infinity
try: return (struct.unpack("<I", struct.pack("<f", float(f)))[0] >> 16) & 0xffff
except (OverflowError, struct.error): return 0x7f80 if f > 0 else 0xff80
def _trig(fn, x):
# V_SIN/COS_F32: hardware does frac on input cycles before computing
if math.isinf(x) or math.isnan(x): return float("nan")
frac_cycles = fract(x / (2 * math.pi))
result = fn(frac_cycles * 2 * math.pi)
# Hardware returns exactly 0 for cos(π/2), sin(π), etc. due to lookup table
# Round very small results (below f32 precision) to exactly 0
if abs(result) < 1e-7: return 0.0
return result
class _SafeFloat(float):
"""Float subclass that uses _div for division to handle 0/inf correctly."""
def __truediv__(self, o): return _div(float(self), float(o))
def __rtruediv__(self, o): return _div(float(o), float(self))
class _Inf:
f16 = f32 = f64 = float('inf')
def __neg__(self): return _NegInf()
def __pos__(self): return self
def __float__(self): return float('inf')
def __eq__(self, other): return float(other) == float('inf') if not isinstance(other, _NegInf) else False
def __req__(self, other): return self.__eq__(other)
class _NegInf:
f16 = f32 = f64 = float('-inf')
def __neg__(self): return _Inf()
def __pos__(self): return self
def __float__(self): return float('-inf')
def __eq__(self, other): return float(other) == float('-inf') if not isinstance(other, _Inf) else False
def __req__(self, other): return self.__eq__(other)
class _RoundMode:
NEAREST_EVEN = 0
class _WaveMode:
IEEE = False
class _DenormChecker:
"""Comparator for denormalized floats. x == DENORM.f32 checks if x is denormalized."""
def __init__(self, bits): self._bits = bits
def _check(self, other):
f = float(other)
if math.isinf(f) or math.isnan(f) or f == 0.0: return False
if self._bits == 64:
bits = struct.unpack("<Q", struct.pack("<d", f))[0]
return (bits >> 52) & 0x7ff == 0
bits = struct.unpack("<I", struct.pack("<f", f))[0]
return (bits >> 23) & 0xff == 0
def __eq__(self, other): return self._check(other)
def __req__(self, other): return self._check(other)
def __ne__(self, other): return not self._check(other)
class _Denorm:
f32 = _DenormChecker(32)
f64 = _DenormChecker(64)
_pack = lambda hi, lo: ((int(hi) & 0xffff) << 16) | (int(lo) & 0xffff)
_pack32 = lambda hi, lo: ((int(hi) & 0xffffffff) << 32) | (int(lo) & 0xffffffff)
class TypedView:
"""View into a Reg with typed access. Used for both full-width (Reg.u32) and slices (Reg[31:16])."""
__slots__ = ('_reg', '_high', '_low', '_signed', '_float', '_bf16', '_reversed')
def __init__(self, reg, high, low=0, signed=False, is_float=False, is_bf16=False):
# Handle reversed slices like [0:31] which means bit-reverse
if high < low: high, low, reversed = low, high, True
else: reversed = False
self._reg, self._high, self._low, self._reversed = reg, high, low, reversed
self._signed, self._float, self._bf16 = signed, is_float, is_bf16
def _nbits(self): return self._high - self._low + 1
def _mask(self): return (1 << self._nbits()) - 1
def _get(self):
v = (self._reg._val >> self._low) & self._mask()
return _brev(v, self._nbits()) if self._reversed else v
def _set(self, v):
v = int(v)
if self._reversed: v = _brev(v, self._nbits())
self._reg._val = (self._reg._val & ~(self._mask() << self._low)) | ((v & self._mask()) << self._low)
@property
def _val(self): return self._get()
@property
def _bits(self): return self._nbits()
# Type accessors for slices (e.g., D0[31:16].f16)
u8 = property(lambda s: s._get() & 0xff)
u16 = property(lambda s: s._get() & 0xffff, lambda s, v: s._set(v))
u32 = property(lambda s: s._get() & MASK32, lambda s, v: s._set(v))
i16 = property(lambda s: _sext(s._get() & 0xffff, 16), lambda s, v: s._set(v))
i32 = property(lambda s: _sext(s._get() & MASK32, 32), lambda s, v: s._set(v))
f16 = property(lambda s: _f16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _i16(float(v))))
f32 = property(lambda s: _f32(s._get()), lambda s, v: s._set(_i32(float(v))))
bf16 = property(lambda s: _bf16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _ibf16(float(v))))
b16, b32 = u16, u32
# Chained type access (e.g., jump_addr.i64 when jump_addr is already TypedView)
@property
def i64(s): return s if s._nbits() == 64 and s._signed else int(s)
@property
def u64(s): return s if s._nbits() == 64 and not s._signed else int(s) & MASK64
def __getitem__(self, key):
if isinstance(key, slice):
high, low = int(key.start), int(key.stop)
return TypedView(self._reg, high, low)
return (self._get() >> int(key)) & 1
def __setitem__(self, key, value):
if isinstance(key, slice):
high, low = int(key.start), int(key.stop)
if high < low: high, low, value = low, high, _brev(int(value), low - high + 1)
mask = (1 << (high - low + 1)) - 1
self._reg._val = (self._reg._val & ~(mask << low)) | ((int(value) & mask) << low)
elif value: self._reg._val |= (1 << int(key))
else: self._reg._val &= ~(1 << int(key))
def __int__(self): return _sext(self._get(), self._nbits()) if self._signed else self._get()
def __index__(self): return int(self)
def __trunc__(self): return int(float(self)) if self._float else int(self)
def __float__(self):
if self._float:
if self._bf16: return _bf16(self._get())
bits = self._nbits()
return _f16(self._get()) if bits == 16 else _f32(self._get()) if bits == 32 else _f64(self._get())
return float(int(self))
def __bool__(s): return bool(int(s))
# Arithmetic - floats use float(), ints use int()
def __add__(s, o): return float(s) + float(o) if s._float else int(s) + int(o)
def __radd__(s, o): return float(o) + float(s) if s._float else int(o) + int(s)
def __sub__(s, o): return float(s) - float(o) if s._float else int(s) - int(o)
def __rsub__(s, o): return float(o) - float(s) if s._float else int(o) - int(s)
def __mul__(s, o): return float(s) * float(o) if s._float else int(s) * int(o)
def __rmul__(s, o): return float(o) * float(s) if s._float else int(o) * int(s)
def __truediv__(s, o): return _div(float(s), float(o)) if s._float else _div(int(s), int(o))
def __rtruediv__(s, o): return _div(float(o), float(s)) if s._float else _div(int(o), int(s))
def __pow__(s, o): return float(s) ** float(o) if s._float else int(s) ** int(o)
def __rpow__(s, o): return float(o) ** float(s) if s._float else int(o) ** int(s)
def __neg__(s): return -float(s) if s._float else -int(s)
def __abs__(s): return abs(float(s)) if s._float else abs(int(s))
# Bitwise - GPU shifts mask the shift amount to valid range
def __and__(s, o): return int(s) & int(o)
def __or__(s, o): return int(s) | int(o)
def __xor__(s, o): return int(s) ^ int(o)
def __invert__(s): return ~int(s)
def __lshift__(s, o): n = int(o); return int(s) << n if 0 <= n < 64 or s._nbits() > 64 else 0
def __rshift__(s, o): n = int(o); return int(s) >> n if 0 <= n < 64 or s._nbits() > 64 else 0
def __rand__(s, o): return int(o) & int(s)
def __ror__(s, o): return int(o) | int(s)
def __rxor__(s, o): return int(o) ^ int(s)
def __rlshift__(s, o): n = int(s); return int(o) << n if 0 <= n < 64 else 0
def __rrshift__(s, o): n = int(s); return int(o) >> n if 0 <= n < 64 else 0
# Comparison - handle _DenormChecker specially
def __eq__(s, o):
if isinstance(o, _DenormChecker): return o._check(s)
return float(s) == float(o) if s._float else int(s) == int(o)
def __ne__(s, o):
if isinstance(o, _DenormChecker): return not o._check(s)
return float(s) != float(o) if s._float else int(s) != int(o)
def __lt__(s, o): return float(s) < float(o) if s._float else int(s) < int(o)
def __le__(s, o): return float(s) <= float(o) if s._float else int(s) <= int(o)
def __gt__(s, o): return float(s) > float(o) if s._float else int(s) > int(o)
def __ge__(s, o): return float(s) >= float(o) if s._float else int(s) >= int(o)
class Reg:
"""GPU register: D0.f32 = S0.f32 + S1.f32 just works. Supports up to 128 bits for DS_LOAD_B128."""
__slots__ = ('_val',)
def __init__(self, val=0): self._val = int(val)
# Typed views - TypedView(reg, high, signed, is_float, is_bf16)
u64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
i64 = property(lambda s: TypedView(s, 63, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK64))
b64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
f64 = property(lambda s: TypedView(s, 63, is_float=True), lambda s, v: setattr(s, '_val', v if isinstance(v, int) else _i64(float(v))))
u32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
i32 = property(lambda s: TypedView(s, 31, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK32))
b32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
f32 = property(lambda s: TypedView(s, 31, is_float=True), lambda s, v: setattr(s, '_val', _i32(float(v))))
u24 = property(lambda s: TypedView(s, 23))
i24 = property(lambda s: TypedView(s, 23, signed=True))
u16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
i16 = property(lambda s: TypedView(s, 15, signed=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
b16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
f16 = property(lambda s: TypedView(s, 15, is_float=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _i16(float(v))) & 0xffff)))
bf16 = property(lambda s: TypedView(s, 15, is_float=True, is_bf16=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _ibf16(float(v))) & 0xffff)))
u8 = property(lambda s: TypedView(s, 7))
i8 = property(lambda s: TypedView(s, 7, signed=True))
u3 = property(lambda s: TypedView(s, 2)) # 3-bit for opsel fields
u1 = property(lambda s: TypedView(s, 0)) # single bit
def __getitem__(s, key):
if isinstance(key, slice): return TypedView(s, int(key.start), int(key.stop))
return (s._val >> int(key)) & 1
def __setitem__(s, key, value):
if isinstance(key, slice):
high, low = int(key.start), int(key.stop)
if high < low: high, low = low, high
mask = (1 << (high - low + 1)) - 1
s._val = (s._val & ~(mask << low)) | ((int(value) & mask) << low)
elif value: s._val |= (1 << int(key))
else: s._val &= ~(1 << int(key))
def __int__(s): return s._val
def __index__(s): return s._val
def __bool__(s): return bool(s._val)
# Arithmetic (for tmp = tmp + 1 patterns). Float operands trigger f32 interpretation.
def __add__(s, o): return (_f32(s._val) + float(o)) if isinstance(o, float) else s._val + int(o)
def __radd__(s, o): return (float(o) + _f32(s._val)) if isinstance(o, float) else int(o) + s._val
def __sub__(s, o): return (_f32(s._val) - float(o)) if isinstance(o, float) else s._val - int(o)
def __rsub__(s, o): return (float(o) - _f32(s._val)) if isinstance(o, float) else int(o) - s._val
def __mul__(s, o): return (_f32(s._val) * float(o)) if isinstance(o, float) else s._val * int(o)
def __rmul__(s, o): return (float(o) * _f32(s._val)) if isinstance(o, float) else int(o) * s._val
def __and__(s, o): return s._val & int(o)
def __rand__(s, o): return int(o) & s._val
def __or__(s, o): return s._val | int(o)
def __ror__(s, o): return int(o) | s._val
def __xor__(s, o): return s._val ^ int(o)
def __rxor__(s, o): return int(o) ^ s._val
def __lshift__(s, o): n = int(o); return s._val << n if 0 <= n < 64 else 0
def __rshift__(s, o): n = int(o); return s._val >> n if 0 <= n < 64 else 0
def __invert__(s): return ~s._val
# Comparison (for tmp >= 0x100000000 patterns)
def __lt__(s, o): return s._val < int(o)
def __le__(s, o): return s._val <= int(o)
def __gt__(s, o): return s._val > int(o)
def __ge__(s, o): return s._val >= int(o)
def __eq__(s, o): return s._val == int(o)
def __ne__(s, o): return s._val != int(o)
# ═══════════════════════════════════════════════════════════════════════════════
# PSEUDOCODE API - Functions and constants from AMD ISA pseudocode
# ═══════════════════════════════════════════════════════════════════════════════
# Rounding and float operations
trunc, floor, ceil = _fpop(math.trunc), _fpop(math.floor), _fpop(math.ceil)
def sqrt(x): return _SafeFloat(math.sqrt(x)) if x >= 0 else _SafeFloat(float("nan"))
def log2(x): return math.log2(x) if x > 0 else (float("-inf") if x == 0 else float("nan"))
def fract(x): return x - math.floor(x)
def sin(x): return _trig(math.sin, x)
def cos(x): return _trig(math.cos, x)
def pow(a, b):
try: return a ** b
except OverflowError: return float("inf") if b > 0 else 0.0
def isEven(x):
x = float(x)
if math.isinf(x) or math.isnan(x): return False
return int(x) % 2 == 0
def mantissa(f):
if f == 0.0 or math.isinf(f) or math.isnan(f): return f
m, _ = math.frexp(f)
return m # AMD V_FREXP_MANT returns mantissa in [0.5, 1.0) range
def signext_from_bit(val, bit):
bit = int(bit)
if bit == 0: return 0
mask = (1 << bit) - 1
val = int(val) & mask
if val & (1 << (bit - 1)): return val - (1 << bit)
return val
# Type conversions
i32_to_f32 = u32_to_f32 = i32_to_f64 = u32_to_f64 = f32_to_f64 = f64_to_f32 = float
def f32_to_i32(f): return _f_to_int(f, -2147483648, 2147483647)
def f32_to_u32(f): return _f_to_int(f, 0, 4294967295)
f64_to_i32, f64_to_u32 = f32_to_i32, f32_to_u32
def f32_to_f16(f):
f = float(f)
if math.isnan(f): return 0x7e00 # f16 NaN
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00 # f16 ±infinity
try: return struct.unpack("<H", struct.pack("<e", f))[0]
except OverflowError: return 0x7c00 if f > 0 else 0xfc00 # overflow -> ±infinity
def f16_to_f32(v): return v if isinstance(v, float) else _f16_to_f32_bits(v)
def i16_to_f16(v): return f32_to_f16(float(_sext(int(v) & 0xffff, 16)))
def u16_to_f16(v): return f32_to_f16(float(int(v) & 0xffff))
def f16_to_i16(bits): f = _f16_to_f32_bits(bits); return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
def f16_to_u16(bits): f = _f16_to_f32_bits(bits); return max(0, min(65535, int(f))) if not math.isnan(f) else 0
def bf16_to_f32(v): return _bf16(v) if isinstance(v, int) else float(v)
def f32_to_bf16(f): return _ibf16(f)
def u8_to_u32(v): return int(v) & 0xff
def u4_to_u32(v): return int(v) & 0xf
def u32_to_u16(u): return int(u) & 0xffff
def i32_to_i16(i): return ((int(i) + 32768) & 0xffff) - 32768
def f16_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
def f16_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
def f32_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
def f32_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
def v_cvt_i16_f32(f): return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
def v_cvt_u16_f32(f): return max(0, min(65535, int(f))) if not math.isnan(f) else 0
def SAT8(v): return max(0, min(255, int(v)))
def f32_to_u8(f): return max(0, min(255, int(f))) if not math.isnan(f) else 0
# Min/max operations
def v_min_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _lt_neg_zero(a, b) else b)
def v_max_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _gt_neg_zero(a, b) else b)
v_min_f16, v_max_f16 = v_min_f32, v_max_f32
v_min_i32, v_max_i32 = min, max
v_min_i16, v_max_i16 = min, max
def v_min_u32(a, b): return min(a & MASK32, b & MASK32)
def v_max_u32(a, b): return max(a & MASK32, b & MASK32)
def v_min_u16(a, b): return min(a & 0xffff, b & 0xffff)
def v_max_u16(a, b): return max(a & 0xffff, b & 0xffff)
def v_min3_f32(a, b, c): return v_min_f32(v_min_f32(a, b), c)
def v_max3_f32(a, b, c): return v_max_f32(v_max_f32(a, b), c)
v_min3_f16, v_max3_f16 = v_min3_f32, v_max3_f32
v_min3_i32, v_max3_i32, v_min3_i16, v_max3_i16 = min, max, min, max
def v_min3_u32(a, b, c): return min(a & MASK32, b & MASK32, c & MASK32)
def v_max3_u32(a, b, c): return max(a & MASK32, b & MASK32, c & MASK32)
def v_min3_u16(a, b, c): return min(a & 0xffff, b & 0xffff, c & 0xffff)
def v_max3_u16(a, b, c): return max(a & 0xffff, b & 0xffff, c & 0xffff)
# SAD/MSAD operations
def ABSDIFF(a, b): return abs(int(a) - int(b))
def v_sad_u8(s0, s1, s2):
"""V_SAD_U8: Sum of absolute differences of 4 byte pairs plus accumulator."""
s0, s1, s2 = int(s0), int(s1), int(s2)
result = s2
for i in range(4):
a = (s0 >> (i * 8)) & 0xff
b = (s1 >> (i * 8)) & 0xff
result += abs(a - b)
return result & 0xffffffff
def v_msad_u8(s0, s1, s2):
"""V_MSAD_U8: Masked sum of absolute differences (skip if reference byte is 0)."""
s0, s1, s2 = int(s0), int(s1), int(s2)
result = s2
for i in range(4):
a = (s0 >> (i * 8)) & 0xff
b = (s1 >> (i * 8)) & 0xff
if b != 0: # Only add diff if reference (s1) byte is non-zero
result += abs(a - b)
return result & 0xffffffff
def BYTE_PERMUTE(data, sel):
"""Select a byte from 64-bit data based on selector value."""
sel = int(sel) & 0xff
if sel <= 7: return (int(data) >> (sel * 8)) & 0xff
if sel == 8: return 0xff if ((int(data) >> 15) & 1) else 0x00
if sel == 9: return 0xff if ((int(data) >> 31) & 1) else 0x00
if sel == 10: return 0xff if ((int(data) >> 47) & 1) else 0x00
if sel == 11: return 0xff if ((int(data) >> 63) & 1) else 0x00
if sel == 12: return 0x00
return 0xff
# Pseudocode functions
def s_ff1_i32_b32(v): return _ctz(v, 32)
def s_ff1_i32_b64(v): return _ctz(v, 64)
GT_NEG_ZERO, LT_NEG_ZERO = _gt_neg_zero, _lt_neg_zero
def isNAN(x):
try: return math.isnan(float(x))
except (TypeError, ValueError): return False
def isQuietNAN(x): return _check_nan_type(x, 1, True)
def isSignalNAN(x): return _check_nan_type(x, 0, False)
def fma(a, b, c):
try: return math.fma(a, b, c)
except ValueError: return float('nan')
def ldexp(m, e): return math.ldexp(m, e)
def sign(f): return 1 if math.copysign(1.0, f) < 0 else 0
def exponent(f):
if hasattr(f, '_bits') and hasattr(f, '_float') and f._float:
raw = f._val
if f._bits == 16: return (raw >> 10) & 0x1f
if f._bits == 32: return (raw >> 23) & 0xff
if f._bits == 64: return (raw >> 52) & 0x7ff
f = float(f)
if math.isinf(f) or math.isnan(f): return 255
if f == 0.0: return 0
try: bits = struct.unpack("<I", struct.pack("<f", f))[0]; return (bits >> 23) & 0xff
except: return 0
def signext(x): return int(x)
def cvtToQuietNAN(x): return float('nan')
def F(x):
"""32'F(x) or 64'F(x) - interpret x as float. If x is int, treat as bit pattern."""
if isinstance(x, int): return _f32(x)
if isinstance(x, TypedView): return x
return float(x)
# Constants
PI = math.pi
WAVE32, WAVE64 = True, False
OVERFLOW_F32, UNDERFLOW_F32 = float('inf'), 0.0
OVERFLOW_F64, UNDERFLOW_F64 = float('inf'), 0.0
MAX_FLOAT_F32 = 3.4028235e+38
INF = _Inf()
ROUND_MODE = _RoundMode()
WAVE_MODE = _WaveMode()
DENORM = _Denorm()
# 2/PI with 1201 bits of precision for V_TRIG_PREOP_F64
TWO_OVER_PI_1201 = Reg(0x0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6)
# ═══════════════════════════════════════════════════════════════════════════════
# COMPILER: pseudocode -> Python (minimal transforms)
# ═══════════════════════════════════════════════════════════════════════════════
def _filter_pseudocode(pseudocode: str) -> str:
"""Filter raw PDF pseudocode to only include actual code lines."""
pcode_lines, in_lambda, depth = [], 0, 0
for line in pseudocode.split('\n'):
s = line.strip()
if not s: continue
if '=>' in s or re.match(r'^[A-Z_]+\(', s): continue # Skip example lines
if '= lambda(' in s: in_lambda += 1; continue # Skip lambda definitions
if in_lambda > 0:
if s.endswith(');'): in_lambda -= 1
continue
# Only include lines that look like pseudocode
is_code = (any(p in s for p in ['D0.', 'D1.', 'S0.', 'S1.', 'S2.', 'SCC =', 'SCC ?', 'VCC', 'EXEC', 'tmp =', 'tmp[', 'lane =', 'PC =',
'D0[', 'D1[', 'S0[', 'S1[', 'S2[', 'MEM[', 'RETURN_DATA', 'VADDR', 'VDATA', 'VDST', 'SADDR', 'OFFSET']) or
s.startswith(('if ', 'else', 'elsif', 'endif', 'declare ', 'for ', 'endfor', '//')) or
re.match(r'^[a-z_]+\s*=', s) or re.match(r'^[a-z_]+\[', s) or (depth > 0 and '=' in s))
if s.startswith('if '): depth += 1
elif s.startswith('endif'): depth = max(0, depth - 1)
if is_code: pcode_lines.append(s)
return '\n'.join(pcode_lines)
def _compile_pseudocode(pseudocode: str) -> str:
"""Compile pseudocode to Python. Transforms are minimal - most syntax just works."""
pseudocode = re.sub(r'\bpass\b', 'pass_', pseudocode) # 'pass' is Python keyword
raw_lines = pseudocode.strip().split('\n')
joined_lines: list[str] = []
for line in raw_lines:
line = line.strip()
if joined_lines and (joined_lines[-1].rstrip().endswith(('||', '&&', '(', ',')) or
(joined_lines[-1].count('(') > joined_lines[-1].count(')'))):
joined_lines[-1] = joined_lines[-1].rstrip() + ' ' + line
else:
joined_lines.append(line)
lines = []
indent, need_pass, in_first_match_loop = 0, False, False
for line in joined_lines:
line = line.split('//')[0].strip() # Strip C-style comments
if not line: continue
if line.startswith('if '):
lines.append(' ' * indent + f"if {_expr(line[3:].rstrip(' then'))}:")
indent += 1
need_pass = True
elif line.startswith('elsif '):
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
lines.append(' ' * indent + f"elif {_expr(line[6:].rstrip(' then'))}:")
indent += 1
need_pass = True
elif line == 'else':
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
lines.append(' ' * indent + "else:")
indent += 1
need_pass = True
elif line.startswith('endif'):
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
need_pass = False
elif line.startswith('endfor'):
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
need_pass, in_first_match_loop = False, False
elif line.startswith('declare '):
pass
elif m := re.match(r'for (\w+) in (.+?)\s*:\s*(.+?) do', line):
start, end = _expr(m[2].strip()), _expr(m[3].strip())
lines.append(' ' * indent + f"for {m[1]} in range({start}, int({end})+1):")
indent += 1
need_pass, in_first_match_loop = True, True
elif '=' in line and not line.startswith('=='):
need_pass = False
line = line.rstrip(';')
if m := re.match(r'\{\s*D1\.[ui]1\s*,\s*D0\.[ui]64\s*\}\s*=\s*(.+)', line):
rhs = _expr(m[1])
lines.append(' ' * indent + f"_full = {rhs}")
lines.append(' ' * indent + f"D0.u64 = int(_full) & 0xffffffffffffffff")
lines.append(' ' * indent + f"D1 = Reg((int(_full) >> 64) & 1)")
elif any(op in line for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^=')):
for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^='):
if op in line:
lhs, rhs = line.split(op, 1)
lines.append(' ' * indent + f"{lhs.strip()} {op} {_expr(rhs.strip())}")
break
else:
lhs, rhs = line.split('=', 1)
lhs_s, rhs_s = _expr(lhs.strip()), rhs.strip()
stmt = _assign(lhs_s, _expr(rhs_s))
if in_first_match_loop and rhs_s == 'i' and (lhs_s == 'tmp' or lhs_s == 'D0.i32'):
stmt += "; break"
lines.append(' ' * indent + stmt)
if need_pass: lines.append(' ' * indent + "pass")
return '\n'.join(lines)
def _assign(lhs: str, rhs: str) -> str:
if lhs in ('tmp', 'SCC', 'VCC', 'EXEC', 'D0', 'D1', 'saveexec', 'PC'):
return f"{lhs} = Reg({rhs})"
return f"{lhs} = {rhs}"
def _expr(e: str) -> str:
e = e.strip()
e = e.replace('&&', ' and ').replace('||', ' or ').replace('<>', ' != ')
e = re.sub(r'!([^=])', r' not \1', e)
e = re.sub(r'\{\s*(\w+\.u32)\s*,\s*(\w+\.u32)\s*\}', r'_pack32(\1, \2)', e)
def pack(m):
hi, lo = _expr(m[1].strip()), _expr(m[2].strip())
return f'_pack({hi}, {lo})'
e = re.sub(r'\{\s*([^,{}]+)\s*,\s*([^,{}]+)\s*\}', pack, e)
e = re.sub(r"1201'B\(2\.0\s*/\s*PI\)", "TWO_OVER_PI_1201", e)
e = re.sub(r"\d+'([0-9a-fA-Fx]+)[UuFf]*", r'\1', e)
e = re.sub(r"\d+'[FIBU]\(", "(", e)
e = re.sub(r'\bB\(', '(', e)
e = re.sub(r'([0-9a-fA-Fx])ULL\b', r'\1', e)
e = re.sub(r'([0-9a-fA-Fx])LL\b', r'\1', e)
e = re.sub(r'([0-9a-fA-Fx])U\b', r'\1', e)
e = re.sub(r'(\d\.?\d*)F\b', r'\1', e)
e = re.sub(r'(\[laneId\])\.[uib]\d+', r'\1', e)
e = e.replace('+INF', 'INF').replace('-INF', '(-INF)')
e = re.sub(r'NAN\.f\d+', 'float("nan")', e)
def convert_verilog_slice(m):
start, width = m.group(1).strip(), m.group(2).strip()
return f'[({start}) + ({width}) - 1 : ({start})]'
e = re.sub(r'\[([^:\[\]]+)\s*\+:\s*([^:\[\]]+)\]', convert_verilog_slice, e)
def process_brackets(s):
result, i = [], 0
while i < len(s):
if s[i] == '[':
depth, start = 1, i + 1
j = start
while j < len(s) and depth > 0:
if s[j] == '[': depth += 1
elif s[j] == ']': depth -= 1
j += 1
inner = _expr(s[start:j-1])
result.append('[' + inner + ']')
i = j
else:
result.append(s[i])
i += 1
return ''.join(result)
e = process_brackets(e)
while '?' in e:
depth, bracket, q = 0, 0, -1
for i, c in enumerate(e):
if c == '(': depth += 1
elif c == ')': depth -= 1
elif c == '[': bracket += 1
elif c == ']': bracket -= 1
elif c == '?' and depth == 0 and bracket == 0: q = i; break
if q < 0: break
depth, bracket, col = 0, 0, -1
for i in range(q + 1, len(e)):
if e[i] == '(': depth += 1
elif e[i] == ')': depth -= 1
elif e[i] == '[': bracket += 1
elif e[i] == ']': bracket -= 1
elif e[i] == ':' and depth == 0 and bracket == 0: col = i; break
if col < 0: break
cond, t, f = e[:q].strip(), e[q+1:col].strip(), e[col+1:].strip()
e = f'(({t}) if ({cond}) else ({f}))'
return e
def _apply_pseudocode_fixes(op_name: str, code: str) -> str:
"""Apply known fixes for PDF pseudocode bugs."""
if op_name == 'V_DIV_FMAS_F32':
code = code.replace('D0.f32 = 2.0 ** 32 * fma(S0.f32, S1.f32, S2.f32)',
'D0.f32 = (2.0 ** 64 if exponent(S2.f32) > 127 else 2.0 ** -64) * fma(S0.f32, S1.f32, S2.f32)')
if op_name == 'V_DIV_FMAS_F64':
code = code.replace('D0.f64 = 2.0 ** 64 * fma(S0.f64, S1.f64, S2.f64)',
'D0.f64 = (2.0 ** 128 if exponent(S2.f64) > 1023 else 2.0 ** -128) * fma(S0.f64, S1.f64, S2.f64)')
if op_name == 'V_DIV_SCALE_F32':
code = code.replace('D0.f32 = float("nan")', 'VCC = Reg(0x1); D0.f32 = float("nan")')
code = code.replace('elif S1.f32 == DENORM.f32:\n D0.f32 = ldexp(S0.f32, 64)', 'elif False:\n pass')
code += '\nif S1.f32 == DENORM.f32:\n D0.f32 = float("nan")'
code = code.replace('elif exponent(S2.f32) <= 23:\n D0.f32 = ldexp(S0.f32, 64)', 'elif exponent(S2.f32) <= 23:\n VCC = Reg(0x1); D0.f32 = ldexp(S0.f32, 64)')
code = code.replace('elif S2.f32 / S1.f32 == DENORM.f32:\n VCC = Reg(0x1)\n if S0.f32 == S2.f32:\n D0.f32 = ldexp(S0.f32, 64)', 'elif S2.f32 / S1.f32 == DENORM.f32:\n VCC = Reg(0x1)')
if op_name == 'V_DIV_SCALE_F64':
code = code.replace('D0.f64 = float("nan")', 'VCC = Reg(0x1); D0.f64 = float("nan")')
code = code.replace('elif S1.f64 == DENORM.f64:\n D0.f64 = ldexp(S0.f64, 128)', 'elif False:\n pass')
code += '\nif S1.f64 == DENORM.f64:\n D0.f64 = float("nan")'
code = code.replace('elif exponent(S2.f64) <= 52:\n D0.f64 = ldexp(S0.f64, 128)', 'elif exponent(S2.f64) <= 52:\n VCC = Reg(0x1); D0.f64 = ldexp(S0.f64, 128)')
code = code.replace('elif S2.f64 / S1.f64 == DENORM.f64:\n VCC = Reg(0x1)\n if S0.f64 == S2.f64:\n D0.f64 = ldexp(S0.f64, 128)', 'elif S2.f64 / S1.f64 == DENORM.f64:\n VCC = Reg(0x1)')
if op_name == 'V_DIV_FIXUP_F32':
code = code.replace('D0.f32 = ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))',
'D0.f32 = ((-OVERFLOW_F32) if (sign_out) else (OVERFLOW_F32)) if isNAN(S0.f32) else ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))')
if op_name == 'V_DIV_FIXUP_F64':
code = code.replace('D0.f64 = ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))',
'D0.f64 = ((-OVERFLOW_F64) if (sign_out) else (OVERFLOW_F64)) if isNAN(S0.f64) else ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))')
if op_name == 'V_TRIG_PREOP_F64':
code = code.replace('result = F((TWO_OVER_PI_1201[1200 : 0] << shift.u32) & 0x1fffffffffffff)',
'result = float(((TWO_OVER_PI_1201[1200 : 0] << int(shift)) >> (1201 - 53)) & 0x1fffffffffffff)')
return code
def _generate_function(cls_name: str, op_name: str, pc: str, code: str) -> str:
"""Generate a single compiled pseudocode function.
Functions take int parameters and return dict of int values.
Reg wrapping happens inside the function, only for registers actually used."""
has_d1 = '{ D1' in pc
is_cmpx = (cls_name in ('VOPCOp', 'VOP3Op')) and 'EXEC.u64[laneId]' in pc
is_div_scale = 'DIV_SCALE' in op_name
has_sdst = cls_name == 'VOP3SDOp' and ('VCC.u64[laneId]' in pc or is_div_scale)
is_ds = cls_name == 'DSOp'
is_flat = cls_name in ('FLATOp', 'GLOBALOp', 'SCRATCHOp')
is_smem = cls_name == 'SMEMOp'
has_s_array = 'S[i]' in pc # FMA_MIX style: S[0], S[1], S[2] array access
combined = code + pc
fn_name = f"_{cls_name}_{op_name}"
# Detect which registers are used/modified
def needs_init(name): return name in combined and not re.search(rf'^\s*{name}\s*=\s*Reg\(', code, re.MULTILINE)
modifies_d0 = is_div_scale or bool(re.search(r'\bD0\b[.\[]', combined))
modifies_exec = is_cmpx or bool(re.search(r'EXEC\.(u32|u64|b32|b64)\s*=', combined))
modifies_vcc = has_sdst or bool(re.search(r'VCC\.(u32|u64|b32|b64)\s*=|VCC\.u64\[laneId\]\s*=', combined))
modifies_scc = bool(re.search(r'\bSCC\s*=', combined))
modifies_pc = bool(re.search(r'\bPC\s*=', combined))
# Build function signature and Reg init lines
if is_smem:
lines = [f"def {fn_name}(MEM, addr):"]
reg_inits = ["ADDR=Reg(addr)", "SDATA=Reg(0)"]
special_regs = []
elif is_ds:
lines = [f"def {fn_name}(MEM, addr, data0, data1, offset0, offset1):"]
reg_inits = ["ADDR=Reg(addr)", "DATA0=Reg(data0)", "DATA1=Reg(data1)", "OFFSET0=Reg(offset0)", "OFFSET1=Reg(offset1)", "RETURN_DATA=Reg(0)"]
special_regs = [('DATA', 'DATA0'), ('DATA2', 'DATA1'), ('OFFSET', 'OFFSET0'), ('ADDR_BASE', 'ADDR')]
elif is_flat:
lines = [f"def {fn_name}(MEM, addr, vdata, vdst):"]
reg_inits = ["ADDR=addr", "VDATA=Reg(vdata)", "VDST=Reg(vdst)", "RETURN_DATA=Reg(0)"]
special_regs = [('DATA', 'VDATA')]
elif has_s_array:
# FMA_MIX style: needs S[i] array, opsel, opsel_hi for source selection (neg/neg_hi applied in emu.py before call)
lines = [f"def {fn_name}(s0, s1, s2, d0, scc, vcc, laneId, exec_mask, literal, VGPR, src0_idx=0, vdst_idx=0, pc=None, opsel=0, opsel_hi=0):"]
reg_inits = ["S0=Reg(s0)", "S1=Reg(s1)", "S2=Reg(s2)", "S=[S0,S1,S2]", "D0=Reg(d0)", "OPSEL=Reg(opsel)", "OPSEL_HI=Reg(opsel_hi)"]
special_regs = []
# Detect array declarations like "declare in : 32'F[3]" and create them (rename 'in' to 'ins' since 'in' is a keyword)
if "in[" in combined:
reg_inits.append("ins=[Reg(0),Reg(0),Reg(0)]")
code = code.replace("in[", "ins[")
else:
lines = [f"def {fn_name}(s0, s1, s2, d0, scc, vcc, laneId, exec_mask, literal, VGPR, src0_idx=0, vdst_idx=0, pc=None):"]
# Only create Regs for registers actually used in the pseudocode
reg_inits = []
if 'S0' in combined: reg_inits.append("S0=Reg(s0)")
if 'S1' in combined: reg_inits.append("S1=Reg(s1)")
if 'S2' in combined: reg_inits.append("S2=Reg(s2)")
if modifies_d0 or 'D0' in combined: reg_inits.append("D0=Reg(s0)" if is_div_scale else "D0=Reg(d0)")
if modifies_scc or 'SCC' in combined: reg_inits.append("SCC=Reg(scc)")
if modifies_vcc or 'VCC' in combined: reg_inits.append("VCC=Reg(vcc)")
if modifies_exec or 'EXEC' in combined: reg_inits.append("EXEC=Reg(exec_mask)")
if modifies_pc or 'PC' in combined: reg_inits.append("PC=Reg(pc) if pc is not None else None")
special_regs = [('D1', 'Reg(0)'), ('SIMM16', 'Reg(literal)'), ('SIMM32', 'Reg(literal)'),
('SRC0', 'Reg(src0_idx)'), ('VDST', 'Reg(vdst_idx)')]
if needs_init('tmp'): special_regs.insert(0, ('tmp', 'Reg(0)'))
if needs_init('saveexec'): special_regs.insert(0, ('saveexec', 'Reg(EXEC._val)'))
# Build init code
init_parts = reg_inits.copy()
for name, init in special_regs:
if name in combined: init_parts.append(f"{name}={init}")
if 'EXEC_LO' in code: init_parts.append("EXEC_LO=TypedView(EXEC, 31, 0)")
if 'EXEC_HI' in code: init_parts.append("EXEC_HI=TypedView(EXEC, 63, 32)")
if 'VCCZ' in code and not re.search(r'^\s*VCCZ\s*=', code, re.MULTILINE): init_parts.append("VCCZ=Reg(1 if VCC._val == 0 else 0)")
if 'EXECZ' in code and not re.search(r'^\s*EXECZ\s*=', code, re.MULTILINE): init_parts.append("EXECZ=Reg(1 if EXEC._val == 0 else 0)")
# Add init line and separator
if init_parts: lines.append(f" {'; '.join(init_parts)}")
# Add compiled pseudocode
for line in code.split('\n'):
if line.strip(): lines.append(f" {line}")
# Build result dict
result_items = []
if modifies_d0: result_items.append("'D0': D0._val")
if modifies_scc: result_items.append("'SCC': SCC._val")
if modifies_vcc: result_items.append("'VCC': VCC._val")
if modifies_exec: result_items.append("'EXEC': EXEC._val")
if has_d1: result_items.append("'D1': D1._val")
if modifies_pc: result_items.append("'PC': PC._val")
if is_smem and 'SDATA' in combined and re.search(r'^\s*SDATA[\.\[].*=', code, re.MULTILINE):
result_items.append("'SDATA': SDATA._val")
if is_ds and 'RETURN_DATA' in combined and re.search(r'^\s*RETURN_DATA[\.\[].*=', code, re.MULTILINE):
result_items.append("'RETURN_DATA': RETURN_DATA._val")
if is_flat:
if 'RETURN_DATA' in combined and re.search(r'^\s*RETURN_DATA[\.\[].*=', code, re.MULTILINE):
result_items.append("'RETURN_DATA': RETURN_DATA._val")
if re.search(r'^\s*VDATA[\.\[].*=', code, re.MULTILINE):
result_items.append("'VDATA': VDATA._val")
lines.append(f" return {{{', '.join(result_items)}}}")
return '\n'.join(lines)
# Build the globals dict for exec() - includes all pcode symbols
_PCODE_GLOBALS = {
'Reg': Reg, 'TypedView': TypedView, '_pack': _pack, '_pack32': _pack32,
'ABSDIFF': ABSDIFF, 'BYTE_PERMUTE': BYTE_PERMUTE, 'DENORM': DENORM, 'F': F,
'GT_NEG_ZERO': GT_NEG_ZERO, 'LT_NEG_ZERO': LT_NEG_ZERO, 'INF': INF,
'MAX_FLOAT_F32': MAX_FLOAT_F32, 'OVERFLOW_F32': OVERFLOW_F32, 'OVERFLOW_F64': OVERFLOW_F64,
'UNDERFLOW_F32': UNDERFLOW_F32, 'UNDERFLOW_F64': UNDERFLOW_F64,
'PI': PI, 'ROUND_MODE': ROUND_MODE, 'WAVE_MODE': WAVE_MODE,
'WAVE32': WAVE32, 'WAVE64': WAVE64, 'TWO_OVER_PI_1201': TWO_OVER_PI_1201,
'SAT8': SAT8, 'trunc': trunc, 'floor': floor, 'ceil': ceil, 'sqrt': sqrt,
'log2': log2, 'fract': fract, 'sin': sin, 'cos': cos, 'pow': pow,
'isEven': isEven, 'mantissa': mantissa, 'signext_from_bit': signext_from_bit,
'i32_to_f32': i32_to_f32, 'u32_to_f32': u32_to_f32, 'i32_to_f64': i32_to_f64,
'u32_to_f64': u32_to_f64, 'f32_to_f64': f32_to_f64, 'f64_to_f32': f64_to_f32,
'f32_to_i32': f32_to_i32, 'f32_to_u32': f32_to_u32, 'f64_to_i32': f64_to_i32,
'f64_to_u32': f64_to_u32, 'f32_to_f16': f32_to_f16, 'f16_to_f32': f16_to_f32,
'i16_to_f16': i16_to_f16, 'u16_to_f16': u16_to_f16, 'f16_to_i16': f16_to_i16,
'f16_to_u16': f16_to_u16, 'bf16_to_f32': bf16_to_f32, 'f32_to_bf16': f32_to_bf16,
'u8_to_u32': u8_to_u32, 'u4_to_u32': u4_to_u32, 'u32_to_u16': u32_to_u16,
'i32_to_i16': i32_to_i16, 'f16_to_snorm': f16_to_snorm, 'f16_to_unorm': f16_to_unorm,
'f32_to_snorm': f32_to_snorm, 'f32_to_unorm': f32_to_unorm,
'v_cvt_i16_f32': v_cvt_i16_f32, 'v_cvt_u16_f32': v_cvt_u16_f32, 'f32_to_u8': f32_to_u8,
'v_min_f32': v_min_f32, 'v_max_f32': v_max_f32, 'v_min_f16': v_min_f16, 'v_max_f16': v_max_f16,
'v_min_i32': v_min_i32, 'v_max_i32': v_max_i32, 'v_min_i16': v_min_i16, 'v_max_i16': v_max_i16,
'v_min_u32': v_min_u32, 'v_max_u32': v_max_u32, 'v_min_u16': v_min_u16, 'v_max_u16': v_max_u16,
'v_min3_f32': v_min3_f32, 'v_max3_f32': v_max3_f32, 'v_min3_f16': v_min3_f16, 'v_max3_f16': v_max3_f16,
'v_min3_i32': v_min3_i32, 'v_max3_i32': v_max3_i32, 'v_min3_i16': v_min3_i16, 'v_max3_i16': v_max3_i16,
'v_min3_u32': v_min3_u32, 'v_max3_u32': v_max3_u32, 'v_min3_u16': v_min3_u16, 'v_max3_u16': v_max3_u16,
'v_sad_u8': v_sad_u8, 'v_msad_u8': v_msad_u8,
's_ff1_i32_b32': s_ff1_i32_b32, 's_ff1_i32_b64': s_ff1_i32_b64,
'isNAN': isNAN, 'isQuietNAN': isQuietNAN, 'isSignalNAN': isSignalNAN,
'fma': fma, 'ldexp': ldexp, 'sign': sign, 'exponent': exponent,
'signext': signext, 'cvtToQuietNAN': cvtToQuietNAN,
}
@functools.cache
def compile_pseudocode(cls_name: str, op_name: str, pseudocode: str):
"""Compile pseudocode string to executable function. Cached for performance."""
filtered = _filter_pseudocode(pseudocode)
code = _compile_pseudocode(filtered)
code = _apply_pseudocode_fixes(op_name, code)
fn_code = _generate_function(cls_name, op_name, filtered, code)
fn_name = f"_{cls_name}_{op_name}"
local_ns = {}
exec(fn_code, _PCODE_GLOBALS, local_ns)
return local_ns[fn_name]
-319
View File
@@ -1,319 +0,0 @@
# Generic PDF text extractor - no external dependencies
import re, zlib
from tinygrad.helpers import fetch, merge_dicts
PDF_URLS = {
"rdna3": "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content",
"rdna4": "https://docs.amd.com/api/khub/documents/uQpkEvk3pv~kfAb2x~j4uw/content",
"cdna": "https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-cdna4-instruction-set-architecture.pdf",
}
# ═══════════════════════════════════════════════════════════════════════════════
# Generic PDF extraction tools
# ═══════════════════════════════════════════════════════════════════════════════
def extract(url: str) -> list[list[tuple[float, float, str, str]]]:
"""Extract positioned text from PDF. Returns list of text elements (x, y, text, font) per page."""
data = fetch(url).read_bytes()
# Parse xref table to locate objects
xref: dict[int, int] = {}
pos = int(re.search(rb'startxref\s+(\d+)', data).group(1)) + 4
while data[pos:pos+7] != b'trailer':
while data[pos:pos+1] in b' \r\n': pos += 1
line_end = data.find(b'\n', pos)
start_obj, count = map(int, data[pos:line_end].split()[:2])
pos = line_end + 1
for i in range(count):
if data[pos+17:pos+18] == b'n' and (off := int(data[pos:pos+10])) > 0: xref[start_obj + i] = off
pos += 20
def get_stream(n: int) -> bytes:
obj = data[xref[n]:data.find(b'endobj', xref[n])]
raw = obj[obj.find(b'stream\n') + 7:obj.find(b'\nendstream')]
return zlib.decompress(raw) if b'/FlateDecode' in obj else raw
# Find page content streams and extract text
pages = []
for n in sorted(xref):
if b'/Type /Page' not in data[xref[n]:xref[n]+500]: continue
if not (m := re.search(rb'/Contents (\d+) 0 R', data[xref[n]:xref[n]+500])): continue
stream = get_stream(int(m.group(1))).decode('latin-1')
elements, font = [], ''
for bt in re.finditer(r'BT(.*?)ET', stream, re.S):
x, y = 0.0, 0.0
for m in re.finditer(r'(/F[\d.]+) [\d.]+ Tf|([\d.+-]+) ([\d.+-]+) Td|[\d.+-]+ [\d.+-]+ [\d.+-]+ [\d.+-]+ ([\d.+-]+) ([\d.+-]+) Tm|<([0-9A-Fa-f]+)>.*?Tj|\[([^\]]+)\] TJ', bt.group(1)):
if m.group(1): font = m.group(1)
elif m.group(2): x, y = x + float(m.group(2)), y + float(m.group(3))
elif m.group(4): x, y = float(m.group(4)), float(m.group(5))
elif m.group(6) and (t := bytes.fromhex(m.group(6)).decode('latin-1')).strip(): elements.append((x, y, t, font))
elif m.group(7) and (t := ''.join(bytes.fromhex(h).decode('latin-1') for h in re.findall(r'<([0-9A-Fa-f]+)>', m.group(7)))).strip(): elements.append((x, y, t, font))
pages.append(sorted(elements, key=lambda e: (-e[1], e[0])))
return pages
def extract_tables(pages: list[list[tuple[float, float, str, str]]]) -> dict[int, tuple[str, list[list[str]]]]:
"""Extract numbered tables from PDF pages. Returns {table_num: (title, rows)} where rows is list of cells per row."""
def group_by_y(texts, key=lambda y: round(y)):
by_y: dict[int, list[tuple[float, float, str]]] = {}
for x, y, t, _ in texts:
by_y.setdefault(key(y), []).append((x, y, t))
return by_y
# Find all table headers by merging text on same line
table_positions = []
for page_idx, texts in enumerate(pages):
for items in group_by_y(texts).values():
line = ''.join(t for _, t in sorted((x, t) for x, _, t in items))
if m := re.search(r'Table (\d+)\. (.+)', line):
table_positions.append((int(m.group(1)), m.group(2).strip(), page_idx, items[0][1]))
table_positions.sort(key=lambda t: (t[2], -t[3]))
# For each table, find rows with matching X positions
result: dict[int, tuple[str, list[list[str]]]] = {}
for num, title, start_page, header_y in table_positions:
rows, col_xs = [], None
for page_idx in range(start_page, len(pages)):
page_texts = [(x, y, t) for x, y, t, _ in pages[page_idx] if 30 < y < 760 and (page_idx > start_page or y < header_y)]
for items in sorted(group_by_y([(x, y, t, '') for x, y, t in page_texts], key=lambda y: round(y / 5)).values(), key=lambda items: -items[0][1]):
xs = tuple(sorted(round(x) for x, _, _ in items))
if col_xs is None:
if len(xs) < 2: continue # Skip single-column rows before table starts
col_xs = xs
elif len(xs) == 1 and xs[0] in col_xs: continue # Skip continuation rows at known column positions
elif not any(c in xs for c in col_xs[:2]): break # Row missing first columns = end of table
rows.append([t for _, t in sorted((x, t) for x, _, t in items)])
else: continue
break
if rows: result[num] = (title, rows)
return result
# ═══════════════════════════════════════════════════════════════════════════════
# AMD specific extraction
# ═══════════════════════════════════════════════════════════════════════════════
def extract_enums(tables: dict[int, tuple[str, list[list[str]]]]) -> dict[str, dict[int, str]]:
"""Extract all enums from tables. Returns {enum_name: {value: name}}."""
enums: dict[str, dict[int, str]] = {}
for num, (title, rows) in tables.items():
# Opcode enums from "XXX Opcodes" tables
if m := re.match(r'(\w+) (?:Y-)?Opcodes', title):
fmt_name = 'VOPD' if 'Y-Opcodes' in title else m.group(1)
ops: dict[int, str] = {}
for row in rows:
for i in range(0, len(row) - 1, 2):
if row[i].isdigit() and re.match(r'^[A-Z][A-Z0-9_]+$', row[i + 1]):
ops[int(row[i])] = row[i + 1]
if ops: enums[fmt_name] = ops
# BufFmt from "Data Format" tables
if 'Data Format' in title:
for row in rows:
for i in range(0, len(row) - 1, 2):
if row[i].isdigit() and re.match(r'^[\dA-Z_]+$', row[i + 1]) and 'INVALID' not in row[i + 1]:
enums.setdefault('BufFmt', {})[int(row[i])] = row[i + 1]
return enums
def extract_ins(tables: dict[int, tuple[str, list[list[str]]]]) -> tuple[dict[str, list[tuple[str, int, int]]], dict[str, str]]:
"""Extract formats and encodings from 'XXX Fields' tables. Returns (formats, encodings)."""
formats: dict[str, list[tuple[str, int, int]]] = {}
encodings: dict[str, str] = {}
for num, (title, rows) in tables.items():
if not (m := re.match(r'(\w+) Fields$', title)): continue
fmt_name = m.group(1)
fields = []
for row in rows:
if len(row) < 2: continue
if (bits := re.match(r'\[?(\d+):(\d+)\]?$', row[1])) or (bits := re.match(r'\[(\d+)\]$', row[1])):
field_name = row[0].lower()
hi, lo = int(bits.group(1)), int(bits.group(2)) if bits.lastindex >= 2 else int(bits.group(1))
if field_name == 'encoding' and len(row) >= 3:
enc_bits = None
if "'b" in row[2]: enc_bits = row[2].split("'b")[-1].replace('_', '')
elif (enc := re.search(r':\s*([01_]+)', row[2])): enc_bits = enc.group(1).replace('_', '')
if enc_bits:
# If encoding bits exceed field width, extend field to match (AMD docs sometimes have this)
declared_width, actual_width = hi - lo + 1, len(enc_bits)
if actual_width > declared_width: lo = hi - actual_width + 1
encodings[fmt_name] = enc_bits
fields.append((field_name, hi, lo))
if fields: formats[fmt_name] = fields
return formats, encodings
def extract_pcode(pages: list[list[tuple[float, float, str, str]]], enums: dict[str, dict[int, str]]) -> dict[tuple[str, int], str]:
"""Extract pseudocode for instructions. Returns {(name, opcode): pseudocode}."""
# Build lookup from instruction name to opcode
name_to_op = {name: op for ops in enums.values() for op, name in ops.items()}
# First pass: find all instruction headers across all pages
all_instructions: list[tuple[int, float, str, int]] = [] # (page_idx, y, name, opcode)
for page_idx, page in enumerate(pages):
by_y: dict[int, list[tuple[float, str]]] = {}
for x, y, t, _ in page:
by_y.setdefault(round(y), []).append((x, t))
for y, items in sorted(by_y.items(), reverse=True):
left = [(x, t) for x, t in items if 55 < x < 65]
right = [(x, t) for x, t in items if 535 < x < 550]
if left and right and left[0][1] in name_to_op and right[0][1].isdigit():
all_instructions.append((page_idx, y, left[0][1], int(right[0][1])))
# Second pass: extract pseudocode between consecutive instructions
pcode: dict[tuple[str, int], str] = {}
for i, (page_idx, y, name, opcode) in enumerate(all_instructions):
# Get end boundary from next instruction
if i + 1 < len(all_instructions):
next_page, next_y = all_instructions[i + 1][0], all_instructions[i + 1][1]
else:
next_page, next_y = page_idx, 0
# Collect F6 text from current position to next instruction (pseudocode is at x ≈ 69)
lines = []
for p in range(page_idx, next_page + 1):
start_y = y if p == page_idx else 800
end_y = next_y if p == next_page else 0
lines.extend((p, y2, t) for x, y2, t, f in pages[p] if f in ('/F6.0', '/F7.0') and end_y < y2 < start_y and 60 < x < 80)
if lines:
# Sort by page first, then by y descending within each page (higher y = earlier text in PDF)
sorted_lines = sorted(lines, key=lambda x: (x[0], -x[1]))
# Stop at large Y gaps (>30) - indicates section break (Notes, examples, etc)
filtered = [sorted_lines[0]]
for j in range(1, len(sorted_lines)):
prev_page, prev_y, _ = sorted_lines[j-1]
curr_page, curr_y, _ = sorted_lines[j]
if curr_page == prev_page and prev_y - curr_y > 30: break
if curr_page != prev_page and prev_y > 60 and curr_y < 730: break # examples spilled to next page (not at very top)
filtered.append(sorted_lines[j])
pcode_lines = [t.replace('Ê', '').strip() for _, _, t in filtered]
if pcode_lines: pcode[(name, opcode)] = '\n'.join(pcode_lines)
return pcode
# ═══════════════════════════════════════════════════════════════════════════════
# Write autogen files
# ═══════════════════════════════════════════════════════════════════════════════
def write_enums(enums: dict[str, dict[int, str]], arch: str, path: str):
"""Write enum.py file from extracted enums."""
lines = ["# autogenerated from AMD ISA PDF by pdf.py - do not edit", "from enum import IntEnum", ""]
for name, values in sorted(enums.items()):
suffix = "Op" if name not in ('Src', 'BufFmt') else ("Enum" if name == 'Src' else "")
prefix = "BUF_FMT_" if name == 'BufFmt' else ""
lines.append(f"class {name}{suffix}(IntEnum):")
for val, member in sorted(values.items()):
lines.append(f" {prefix}{member} = {val}")
lines.append("")
with open(path, "w") as f:
f.write("\n".join(lines))
def write_ins(formats: dict[str, list[tuple[str, int, int]]], encodings: dict[str, str], enums: dict[str, dict[int, str]], arch: str, path: str):
"""Write ins.py file from extracted formats and enums."""
# Field types and ordering
def field_type(name, fmt):
if name == 'op' and fmt in enums: return f'Annotated[BitField, {fmt}Op]'
if name in ('opx', 'opy'): return 'Annotated[BitField, VOPDOp]'
if name == 'vdsty': return 'VDSTYEnc'
if name in ('vdst', 'vsrc1', 'vaddr', 'vdata', 'data', 'data0', 'data1', 'addr', 'vsrc0', 'vsrc2', 'vsrc3'): return 'VGPRField'
if name in ('sdst', 'sbase', 'sdata', 'srsrc', 'ssamp'): return 'SGPRField'
if name.startswith('ssrc') or name in ('saddr', 'soffset'): return 'SSrc'
if name in ('src0', 'srcx0', 'srcy0') or name.startswith('src') and name[3:].isdigit(): return 'Src'
if name.startswith('simm'): return 'SImm'
if name == 'offset' or name.startswith('imm'): return 'Imm'
return None
field_priority = ['encoding', 'op', 'opx', 'opy', 'vdst', 'vdstx', 'vdsty', 'sdst', 'vdata', 'sdata', 'addr', 'vaddr', 'data', 'data0', 'data1',
'src0', 'srcx0', 'srcy0', 'vsrc0', 'ssrc0', 'src1', 'vsrc1', 'vsrcx1', 'vsrcy1', 'ssrc1', 'src2', 'vsrc2', 'src3', 'vsrc3',
'saddr', 'sbase', 'srsrc', 'ssamp', 'soffset', 'offset', 'simm16', 'en', 'target', 'attr', 'attr_chan',
'omod', 'neg', 'neg_hi', 'abs', 'clmp', 'opsel', 'opsel_hi', 'waitexp', 'wait_va',
'dmask', 'dim', 'seg', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe', 'unrm', 'done', 'row']
def sort_fields(fields):
order = {name: i for i, name in enumerate(field_priority)}
return sorted(fields, key=lambda f: (order.get(f[0], 1000), f[2]))
# Generate format classes
lines = ["# autogenerated from AMD ISA PDF by pdf.py - do not edit", "# ruff: noqa: F401,F403",
"from typing import Annotated",
"from extra.assembly.amd.dsl import *",
f"from extra.assembly.amd.autogen.{arch}.enum import *", "import functools", ""]
for fmt_name, fields in sorted(formats.items()):
lines.append(f"class {fmt_name}(Inst):")
for name, hi, lo in sort_fields(fields):
bits_str = f"bits[{hi}:{lo}]" if hi != lo else f"bits[{hi}]"
if name == 'encoding' and fmt_name in encodings: lines.append(f" encoding = {bits_str} == 0b{encodings[fmt_name]}")
else:
ftype = field_type(name, fmt_name)
lines.append(f" {name}{f':{ftype}' if ftype else ''} = {bits_str}")
lines.append("")
# Generate instruction helpers
lines.append("# instruction helpers")
for fmt_name, ops in sorted(enums.items()):
seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=1"}.get(fmt_name, "")
tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt_name, f"{fmt_name}, {fmt_name}Op")
suffix = "_e32" if fmt_name in ("VOP1", "VOP2", "VOPC") else "_e64" if fmt_name == "VOP3" and len(ops) > 0 else ""
if fmt_name in formats or fmt_name in ("GLOBAL", "SCRATCH"):
for op_val, name in sorted(ops.items()):
fn_suffix = suffix if fmt_name != "VOP3" or op_val < 512 else ""
lines.append(f"{name.lower()}{fn_suffix} = functools.partial({tgt}.{name}{seg})")
with open(path, "w") as f:
f.write("\n".join(lines))
def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, str]], arch: str, path: str):
"""Write str_pcode.py file from extracted pseudocode."""
# Group pseudocode by enum class
by_enum: dict[str, list[tuple[str, int, str]]] = {}
for fmt_name, ops in enums.items():
for opcode, name in ops.items():
if (name, opcode) in pcode: by_enum.setdefault(f"{fmt_name}Op", []).append((name, opcode, pcode[(name, opcode)]))
# Generate file
enum_names = sorted(by_enum.keys())
lines = [f"# autogenerated by pdf.py - do not edit", f"# to regenerate: python -m extra.assembly.amd.pdf",
"# ruff: noqa: E501", f"from extra.assembly.amd.autogen.{arch}.enum import {', '.join(enum_names)}", ""]
for enum_name in enum_names:
lines.append(f"{enum_name}_PCODE = {{")
for name, opcode, code in sorted(by_enum[enum_name], key=lambda x: x[1]):
lines.append(f" {enum_name}.{name}: {code!r},")
lines.append("}\n")
lines.append(f"PSEUDOCODE_STRINGS = {{{', '.join(f'{e}: {e}_PCODE' for e in enum_names)}}}")
with open(path, "w") as f:
f.write("\n".join(lines))
if __name__ == "__main__":
import pathlib
for arch, url in PDF_URLS.items():
print(f"Processing {arch}...")
pages = extract(url)
tables = extract_tables(pages)
enums = extract_enums(tables)
formats, encodings = extract_ins(tables)
pcode = extract_pcode(pages, enums)
# Fix known PDF errors
if arch == 'rdna3':
fixes = {'SOPP': {8: 'S_WAITCNT_DEPCTR', 58: 'S_TTRACEDATA', 59: 'S_TTRACEDATA_IMM'},
'SOPK': {22: 'S_SUBVECTOR_LOOP_BEGIN', 23: 'S_SUBVECTOR_LOOP_END'},
'SMEM': {34: 'S_ATC_PROBE', 35: 'S_ATC_PROBE_BUFFER'},
'DS': {24: 'DS_GWS_SEMA_RELEASE_ALL', 25: 'DS_GWS_INIT', 26: 'DS_GWS_SEMA_V', 27: 'DS_GWS_SEMA_BR', 28: 'DS_GWS_SEMA_P', 29: 'DS_GWS_BARRIER'},
'FLAT': {40: 'GLOBAL_LOAD_ADDTID_B32', 41: 'GLOBAL_STORE_ADDTID_B32', 55: 'FLAT_ATOMIC_CSUB_U32'}}
for fmt, ops in fixes.items(): enums[fmt] = merge_dicts([enums[fmt], ops])
if arch == 'rdna4':
fixes = {'SMEM': {34: 'S_ATC_PROBE', 35: 'S_ATC_PROBE_BUFFER'},
'SOP1': {81: 'S_BARRIER_INIT', 82: 'S_BARRIER_JOIN'},
'SOPP': {21: 'S_BARRIER_LEAVE', 58: 'S_TTRACEDATA', 59: 'S_TTRACEDATA_IMM'}}
for fmt, ops in fixes.items(): enums[fmt] = merge_dicts([enums[fmt], ops])
if arch in ('rdna3', 'rdna4'):
# RDNA SMEM: PDF says DLC=[14], GLC=[16] but hardware uses DLC=[13], GLC=[14]
if 'SMEM' in formats:
formats['SMEM'] = [(n, 13 if n == 'dlc' else 14 if n == 'glc' else h, 13 if n == 'dlc' else 14 if n == 'glc' else l)
for n, h, l in formats['SMEM']]
if arch == 'cdna':
# CDNA DS: PDF is missing the GDS field (bit 16)
if 'DS' in formats and not any(n == 'gds' for n, _, _ in formats['DS']):
formats['DS'].append(('gds', 16, 16))
# CDNA DPP/SDWA: PDF only documents modifier fields (bits[63:32]), need to add VOP overlay fields (bits[31:0])
vop_overlay = [('encoding', 8, 0), ('vop_op', 16, 9), ('vdst', 24, 17), ('vop2_op', 31, 25)]
if 'DPP' in formats and not any(n == 'encoding' for n, _, _ in formats['DPP']):
formats['DPP'] = vop_overlay + [('bc' if n == 'bound_ctrl' else n, h, l) for n, h, l in formats['DPP']]
encodings['DPP'] = '11111010'
if 'SDWA' in formats and not any(n == 'encoding' for n, _, _ in formats['SDWA']):
formats['SDWA'] = vop_overlay + [(n, h, l) for n, h, l in formats['SDWA']]
encodings['SDWA'] = '11111001'
base = pathlib.Path(__file__).parent / "autogen" / arch
write_enums(enums, arch, base / "enum.py")
write_ins(formats, encodings, enums, arch, base / "ins.py")
write_pcode(pcode, enums, arch, base / "str_pcode.py")
print(f" {len(tables)} tables, {len(pcode)} pcode -> {base}")
-450
View File
@@ -1,450 +0,0 @@
"""SQTT (SQ Thread Trace) packet encoder and decoder for AMD GPUs.
This module provides encoding and decoding of raw SQTT byte streams.
The format is nibble-based with variable-width packets determined by a state machine.
Uses BitField infrastructure from dsl.py, similar to GPU instruction encoding.
"""
from __future__ import annotations
from enum import IntEnum
from typing import get_type_hints
from extra.assembly.amd.dsl import BitField, bits
# ═══════════════════════════════════════════════════════════════════════════════
# FIELD ENUMS
# ═══════════════════════════════════════════════════════════════════════════════
class MemSrc(IntEnum):
LDS = 0
LDS_ALT = 1
VMEM = 2
VMEM_ALT = 3
class AluSrc(IntEnum):
NONE = 0
SALU = 1
VALU = 2
VALU_ALT = 3
class InstOp(IntEnum):
"""SQTT instruction operation types.
Memory ops appear in two ranges depending on which SIMD executes them:
- 0x1x-0x2x range: ops on traced SIMD
- 0x5x range: ops on other SIMD (OTHER_ prefix)
GLOBAL memory ops encoding depends on addressing mode AND size:
- Loads: 0x21 (saddr=SGPR) or 0x22 (saddr=NULL), all sizes same
- Stores: base + size_offset, where VADDR is shifted +1 from SADDR
SADDR: 0x24(32) 0x25(64) 0x26(96) 0x27(128)
VADDR: 0x25(32) 0x26(64) 0x27(96) 0x28(128)
OTHER_ range follows same pattern but values overlap differently.
"""
SALU = 0x0
SMEM = 0x1
JUMP = 0x3 # branch taken
JUMP_NO = 0x4 # branch not taken
MESSAGE = 0x9
VALU_TRANS = 0xb # transcendental: exp, log, rcp, sqrt, sin, cos
VALU_64_SHIFT = 0xd # 64-bit shifts: lshl, lshr, ashr
VALU_MAD64 = 0xe # 64-bit multiply-add
VALU_64 = 0xf # 64-bit: add, mul, fma, rcp, sqrt, rounding, frexp, div helpers
VINTERP = 0x12 # interpolation: v_interp_p10_f32, v_interp_p2_f32
BARRIER = 0x13
# FLAT memory ops on traced SIMD (0x1x range)
FLAT_LOAD = 0x1c
FLAT_STORE = 0x1d
FLAT_STORE_64 = 0x1e
FLAT_STORE_96 = 0x1f
FLAT_STORE_128 = 0x20
# GLOBAL memory ops on traced SIMD (0x2x range)
GLOBAL_LOAD = 0x21 # saddr=SGPR, all sizes
GLOBAL_LOAD_VADDR = 0x22 # saddr=NULL, all sizes
GLOBAL_STORE = 0x24 # saddr=SGPR, 32-bit
GLOBAL_STORE_64 = 0x25 # saddr=SGPR 64 or saddr=NULL 32
GLOBAL_STORE_96 = 0x26 # saddr=SGPR 96 or saddr=NULL 64
GLOBAL_STORE_128 = 0x27 # saddr=SGPR 128 or saddr=NULL 96
GLOBAL_STORE_VADDR_128 = 0x28 # saddr=NULL, 128-bit
# LDS ops on traced SIMD
LDS_LOAD = 0x29
LDS_STORE = 0x2b
LDS_STORE_64 = 0x2c
LDS_STORE_128 = 0x2e
# Memory ops on other SIMD (0x5x range)
OTHER_LDS_LOAD = 0x50
OTHER_LDS_STORE = 0x51
OTHER_LDS_STORE_64 = 0x52
OTHER_LDS_STORE_128 = 0x54
OTHER_FLAT_LOAD = 0x55
OTHER_FLAT_STORE = 0x56
OTHER_FLAT_STORE_64 = 0x57
OTHER_FLAT_STORE_96 = 0x58
OTHER_FLAT_STORE_128 = 0x59
OTHER_GLOBAL_LOAD = 0x5a # saddr=SGPR, all sizes
OTHER_GLOBAL_LOAD_VADDR = 0x5b # saddr=NULL or saddr=SGPR store 32
OTHER_GLOBAL_STORE_64 = 0x5c # saddr=SGPR 64 or saddr=NULL 32
OTHER_GLOBAL_STORE_96 = 0x5d # saddr=SGPR 96 or saddr=NULL 64
OTHER_GLOBAL_STORE_128 = 0x5e # saddr=SGPR 128 or saddr=NULL 96
OTHER_GLOBAL_STORE_VADDR_128 = 0x5f # saddr=NULL, 128-bit
# EXEC-modifying ops (0x7x range)
SALU_SAVEEXEC = 0x72 # s_*_saveexec_b32/b64
VALU_CMPX = 0x73 # v_cmpx_*
# ═══════════════════════════════════════════════════════════════════════════════
# PACKET TYPE BASE CLASS
# ═══════════════════════════════════════════════════════════════════════════════
class PacketType:
"""Base class for SQTT packet types."""
_encoding: tuple[BitField, int] | None = None
_field_types: dict[str, type] = {}
_values: dict[str, int]
_raw: int
_time: int
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
if 'encoding' in cls.__dict__ and isinstance(cls.__dict__['encoding'], tuple):
cls._encoding = cls.__dict__['encoding']
# Cache field type annotations for enum conversion
try: cls._field_types = {k: v for k, v in get_type_hints(cls).items() if isinstance(v, type) and issubclass(v, IntEnum)}
except Exception: cls._field_types = {}
# Cache fields and precompute extraction info: (name, lo, mask, enum_type)
cls._fields = {k: v for k, v in cls.__dict__.items() if isinstance(v, BitField) and k != 'encoding'}
cls._extract_info = [(name, bf.lo, bf.mask(), cls._field_types.get(name)) for name, bf in cls._fields.items()]
cls._size_nibbles = ((max((f.hi for f in cls._fields.values()), default=0) + 4) // 4)
def __init__(self, _time: int = 0, **kwargs):
"""Construct packet from named fields (like assembly instructions)."""
raw = 0
if self._encoding:
bf, pattern = self._encoding
raw |= pattern << bf.lo
for name, bf in self._fields.items():
val = kwargs.get(name, 0)
if isinstance(val, IntEnum): val = val.value
raw |= (val & bf.mask()) << bf.lo
self._raw, self._time, self._values = raw, _time, {}
for name, lo, mask, enum_type in self._extract_info:
val = (raw >> lo) & mask
if enum_type is not None:
try: val = enum_type(val)
except ValueError: pass
self._values[name] = val
@classmethod
def from_raw(cls, raw: int, time: int = 0):
inst = object.__new__(cls)
inst._raw, inst._time, inst._values = raw, time, {}
for name, lo, mask, enum_type in cls._extract_info:
val = (raw >> lo) & mask
if enum_type is not None:
try: val = enum_type(val)
except ValueError: pass
inst._values[name] = val
return inst
def __getattr__(self, name: str):
if name.startswith('_'): raise AttributeError(name)
return self._values.get(name, 0)
def __repr__(self) -> str:
fields_str = ", ".join(f"{k}={v}" for k, v in self._values.items() if not k.startswith('_'))
return f"{self.__class__.__name__}({fields_str})"
# ═══════════════════════════════════════════════════════════════════════════════
# PACKET TYPE DEFINITIONS
# ═══════════════════════════════════════════════════════════════════════════════
class VALUINST(PacketType): # exclude: 1 << 2
encoding = bits[2:0] == 0b011
delta = bits[5:3]
flag = bits[6:6]
wave = bits[11:7]
class VMEMEXEC(PacketType): # exclude: 1 << 0
encoding = bits[3:0] == 0b1111
delta = bits[5:4]
src: MemSrc = bits[7:6]
class ALUEXEC(PacketType): # exclude: 1 << 1
encoding = bits[3:0] == 0b1110
delta = bits[5:4]
src: AluSrc = bits[7:6]
class IMMEDIATE(PacketType): # exclude: 1 << 5
encoding = bits[3:0] == 0b1101
delta = bits[6:4]
wave = bits[11:7]
class IMMEDIATE_MASK(PacketType): # exclude: 1 << 5
encoding = bits[4:0] == 0b00100
delta = bits[7:5]
mask = bits[23:8]
class WAVERDY(PacketType): # exclude: 1 << 3
encoding = bits[4:0] == 0b10100
delta = bits[7:5]
mask = bits[23:8]
class TS_DELTA_S8_W3(PacketType):
encoding = bits[6:0] == 0b0100001
delta = bits[10:8]
_padding = bits[63:11]
class WAVEEND(PacketType): # exclude: 1 << 4
encoding = bits[4:0] == 0b10101
delta = bits[7:5]
flag7 = bits[8:8]
simd = bits[10:9]
cu_lo = bits[13:11]
wave = bits[19:15]
@property
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
class WAVESTART(PacketType): # exclude: 1 << 4
encoding = bits[4:0] == 0b01100
delta = bits[6:5]
flag7 = bits[7:7]
simd = bits[9:8]
cu_lo = bits[12:10]
wave = bits[17:13]
id7 = bits[31:18]
@property
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
class TS_DELTA_S5_W2(PacketType):
encoding = bits[4:0] == 0b11100
delta = bits[6:5]
_padding = bits[47:7]
class WAVEALLOC(PacketType): # exclude: 1 << 10
encoding = bits[4:0] == 0b00101
delta = bits[7:5]
_padding = bits[19:8]
class TS_DELTA_S5_W3(PacketType):
encoding = bits[4:0] == 0b00110
delta = bits[7:5]
_padding = bits[51:8]
class PERF(PacketType): # exclude: 1 << 11
encoding = bits[4:0] == 0b10110
delta = bits[7:5]
arg = bits[27:8]
class TS_DELTA_SHORT(PacketType):
encoding = bits[3:0] == 0b1000
delta = bits[7:4]
class NOP(PacketType):
encoding = bits[3:0] == 0b0000
delta = None # type: ignore
_padding = bits[3:0]
class TS_WAVE_STATE(PacketType):
encoding = bits[6:0] == 0b1010001
delta = bits[15:7]
coarse = bits[23:16]
@property
def wave_interest(self) -> bool: return bool(self.coarse & 1)
@property
def terminate_all(self) -> bool: return bool(self.coarse & 8)
class EVENT(PacketType): # exclude: 1 << 7
encoding = bits[7:0] == 0b01100001
delta = bits[10:8]
event = bits[23:11]
class EVENT_BIG(PacketType):
encoding = bits[7:0] == 0b11100001
delta = bits[10:8]
event = bits[31:11]
class REG(PacketType):
encoding = bits[3:0] == 0b1001
delta = bits[6:4]
slot = bits[9:7]
hi_byte = bits[15:8]
subop = bits[31:16]
val32 = bits[63:32]
@property
def is_config(self) -> bool: return bool(self.hi_byte & 0x80)
class SNAPSHOT(PacketType):
encoding = bits[6:0] == 0b1110001
delta = bits[9:7]
snap = bits[63:10]
class TS_DELTA_OR_MARK(PacketType):
encoding = bits[6:0] == 0b0000001
delta = bits[47:12]
bit8 = bits[8:8]
bit9 = bits[9:9]
@property
def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
class LAYOUT_HEADER(PacketType):
encoding = bits[6:0] == 0b0010001
delta = None # type: ignore
layout = bits[12:7]
simd = bits[14:13]
group = bits[17:15]
sel_a = bits[31:28]
sel_b = bits[36:33]
flag4 = bits[59:59]
_padding = bits[63:60]
class INST(PacketType):
encoding = bits[2:0] == 0b010
delta = bits[6:4]
flag1 = bits[3:3]
flag2 = bits[7:7]
wave = bits[12:8]
op: InstOp = bits[19:13]
class UTILCTR(PacketType):
encoding = bits[6:0] == 0b0110001
delta = bits[8:7]
ctr = bits[47:9]
# All packet types in encoding priority order (more specific masks first, NOP last as fallback)
PACKET_TYPES: list[type[PacketType]] = [
EVENT, EVENT_BIG,
TS_DELTA_S8_W3, TS_WAVE_STATE, SNAPSHOT, TS_DELTA_OR_MARK, LAYOUT_HEADER, UTILCTR,
IMMEDIATE_MASK, WAVERDY, WAVEEND, WAVESTART, TS_DELTA_S5_W2, WAVEALLOC, TS_DELTA_S5_W3, PERF,
VMEMEXEC, ALUEXEC, IMMEDIATE, TS_DELTA_SHORT, REG,
VALUINST, INST,
NOP,
]
PACKET_BY_NAME: dict[str, type[PacketType]] = {cls.__name__: cls for cls in PACKET_TYPES}
def _build_state_table() -> tuple[bytes, dict[int, type[PacketType]]]:
table = [len(PACKET_TYPES) - 1] * 256 # default to NOP
opcode_to_class: dict[int, type[PacketType]] = {i: cls for i, cls in enumerate(PACKET_TYPES)}
for byte_val in range(256):
for opcode, pkt_cls in enumerate(PACKET_TYPES):
if pkt_cls._encoding is None: continue
mask_bf, pattern = pkt_cls._encoding
if (byte_val & mask_bf.mask()) == pattern:
table[byte_val] = opcode
break
return bytes(table), opcode_to_class
STATE_TO_OPCODE, OPCODE_TO_CLASS = _build_state_table()
OPCODE_TO_BYTES: dict[int, list[int]] = {}
for _byte_val, _opcode in enumerate(STATE_TO_OPCODE):
if _opcode not in OPCODE_TO_BYTES: OPCODE_TO_BYTES[_opcode] = []
OPCODE_TO_BYTES[_opcode].append(_byte_val)
# Precompute special case opcodes
_TS_DELTA_OR_MARK_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_OR_MARK)
_TS_DELTA_SHORT_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_SHORT)
_TS_DELTA_OR_MARK_BIT8 = (TS_DELTA_OR_MARK.bit8.lo, TS_DELTA_OR_MARK.bit8.mask())
_TS_DELTA_OR_MARK_BIT9 = (TS_DELTA_OR_MARK.bit9.lo, TS_DELTA_OR_MARK.bit9.mask())
# Combined lookup: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
# special_case: 0=none, 1=TS_DELTA_OR_MARK, 2=TS_DELTA_SHORT
_DECODE_INFO: dict[int, tuple] = {}
for _opcode, _pkt_cls in OPCODE_TO_CLASS.items():
_delta_field = getattr(_pkt_cls, 'delta', None)
_delta_lo = _delta_field.lo if _delta_field else 0
_delta_mask = _delta_field.mask() if _delta_field else 0
_special = 1 if _opcode == _TS_DELTA_OR_MARK_OPCODE else (2 if _opcode == _TS_DELTA_SHORT_OPCODE else 0)
_DECODE_INFO[_opcode] = (_pkt_cls, _pkt_cls._size_nibbles, _delta_lo, _delta_mask, _special)
# ═══════════════════════════════════════════════════════════════════════════════
# DECODER
# ═══════════════════════════════════════════════════════════════════════════════
def decode(data: bytes) -> list[PacketType]:
"""Decode raw SQTT blob into list of packet instances."""
packets: list[PacketType] = []
packets_append = packets.append
n = len(data)
reg = 0
offset = 0
nib_count = 16
time = 0
state_to_opcode = STATE_TO_OPCODE
decode_info = _DECODE_INFO
mask64 = (1 << 64) - 1
while (offset >> 3) < n:
target = offset + nib_count * 4
while offset < target and (offset >> 3) < n:
byte = data[offset >> 3]
nib = (byte >> (offset & 4)) & 0xF
reg = ((reg >> 4) | (nib << 60)) & mask64
offset += 4
if offset < target: break
opcode = state_to_opcode[reg & 0xFF]
pkt_cls, nib_count, delta_lo, delta_mask, special = decode_info[opcode]
delta = (reg >> delta_lo) & delta_mask
if special == 1: # TS_DELTA_OR_MARK
bit8 = (reg >> _TS_DELTA_OR_MARK_BIT8[0]) & _TS_DELTA_OR_MARK_BIT8[1]
bit9 = (reg >> _TS_DELTA_OR_MARK_BIT9[0]) & _TS_DELTA_OR_MARK_BIT9[1]
if bit9 and not bit8: delta = 0
elif special == 2: # TS_DELTA_SHORT
delta = delta + 8
time += delta
packets_append(pkt_cls.from_raw(reg, time))
return packets
# ═══════════════════════════════════════════════════════════════════════════════
# ENCODER
# ═══════════════════════════════════════════════════════════════════════════════
def encode(packets: list[PacketType]) -> bytes:
"""Encode a list of packet instances into raw SQTT blob."""
if not packets: return b''
read_lengths = [16]
for p in packets[:-1]:
read_lengths.append(type(p)._size_nibbles)
total_nibbles = sum(read_lengths)
bits_arr = [0] * (total_nibbles * 4)
cumulative = 0
for i, p in enumerate(packets):
cumulative += read_lengths[i]
pkt_cls = type(p)
opcode = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is pkt_cls)
byte_vals = OPCODE_TO_BYTES.get(opcode)
if not byte_vals: raise ValueError(f"No encoding for {pkt_cls.__name__}")
opcode_byte = byte_vals[0]
delta_field = getattr(pkt_cls, 'delta', None)
if delta_field is not None and delta_field.hi < 8:
delta = p._values.get('delta', 0)
if isinstance(delta, IntEnum): delta = delta.value
if pkt_cls is TS_DELTA_SHORT: delta = max(0, delta - 8)
delta = delta & delta_field.mask()
opcode_byte = (opcode_byte & ~(delta_field.mask() << delta_field.lo)) | (delta << delta_field.lo)
opcode_nibble_pos = max(0, cumulative - 16)
opcode_bit_pos = opcode_nibble_pos * 4
for b in range(8):
if opcode_bit_pos + b < len(bits_arr):
bits_arr[opcode_bit_pos + b] = (opcode_byte >> b) & 1
nibbles = [sum(bits_arr[i + j] << j for j in range(4) if i + j < len(bits_arr)) for i in range(0, len(bits_arr), 4)]
while len(nibbles) % 2: nibbles.append(0)
return bytes(nibbles[i] | (nibbles[i + 1] << 4) for i in range(0, len(nibbles), 2))
-192
View File
@@ -1,192 +0,0 @@
#!/usr/bin/env python3
"""Benchmark comparing Python vs Rust RDNA3 emulators on real tinygrad kernels."""
import ctypes, time, os
from pathlib import Path
# Set AMD=1 before importing tinygrad
os.environ["AMD"] = "1"
from extra.assembly.amd.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
if not REMU_PATH.exists():
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.dylib"
def get_rust_remu():
"""Load the Rust libremu shared library."""
if not REMU_PATH.exists(): return None
remu = ctypes.CDLL(str(REMU_PATH))
remu.run_asm.restype = ctypes.c_int32
remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
return remu
def count_instructions(kernel: bytes) -> int:
"""Count instructions in a kernel."""
return len(decode_program(kernel))
def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = None):
"""Allocate buffers and return args pointer + valid ranges."""
if init_data is None: init_data = {}
buffers = []
for i, size in enumerate(buf_sizes):
padded = ((size + 15) // 16) * 16 + 16
data = init_data.get(i, b'\x00' * padded)
data_list = list(data) + [0] * (padded - len(data))
buf = (ctypes.c_uint8 * padded)(*data_list[:padded])
buffers.append(buf)
args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
args_ptr = ctypes.addressof(args)
ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
ranges.add((args_ptr, ctypes.sizeof(args)))
return buffers, args, args_ptr, ranges
def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, rsrc2: int, iterations: int = 5):
"""Benchmark an emulator and return average time."""
gx, gy, gz = global_size
lx, ly, lz = local_size
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
lib_ptr = ctypes.addressof(kernel_buf)
# Warmup
run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
# Timed runs
times = []
for _ in range(iterations):
start = time.perf_counter()
result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
end = time.perf_counter()
if result != 0:
print(f" {name} returned error: {result}")
return None
times.append(end - start)
return sum(times) / len(times)
def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes], int] | None:
"""Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data, rsrc2)."""
try:
from tinygrad import Tensor
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.autogen import hsa
import numpy as np
np.random.seed(42)
ops = {
"add": lambda: Tensor.empty(1024) + Tensor.empty(1024),
"mul": lambda: Tensor.empty(1024) * Tensor.empty(1024),
"matmul_small": lambda: Tensor.empty(16, 16) @ Tensor.empty(16, 16),
"matmul_medium": lambda: Tensor.empty(64, 64) @ Tensor.empty(64, 64),
"reduce_sum": lambda: Tensor.empty(4096).sum(),
"reduce_max": lambda: Tensor.empty(4096).max(),
"softmax": lambda: Tensor.empty(256).softmax(),
"layernorm": lambda: Tensor.empty(32, 64).layernorm(),
"conv2d": lambda: Tensor.empty(1, 4, 16, 16).conv2d(Tensor.empty(4, 4, 3, 3)),
"gelu": lambda: Tensor.empty(1024).gelu(),
"exp": lambda: Tensor.empty(1024).exp(),
"sin": lambda: Tensor.empty(1024).sin(),
}
if op_name not in ops: return None
out = ops[op_name]()
sched = out.schedule()
for ei in sched:
lowered = ei.lower()
if ei.ast.op.name == 'SINK' and lowered.prg and lowered.prg.p.lib:
lib = bytes(lowered.prg.p.lib)
image = memoryview(bytearray(lib))
_, sections, _ = elf_loader(lib)
rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
for sec in sections:
if sec.name == '.text':
buf_sizes = [b.nbytes for b in lowered.bufs]
# Get initial data from numpy arrays if available
buf_data = {}
for i, buf in enumerate(lowered.bufs):
if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
try: buf_data[i] = bytes(buf.base._buf)
except: pass
# Extract rsrc2 from ELF (same as ops_amd.py)
group_segment_size = image[rodata_entry:rodata_entry+4].cast("I")[0]
lds_size = ((group_segment_size + 511) // 512) & 0x1FF
code = hsa.amd_kernel_code_t.from_buffer_copy(bytes(image[rodata_entry:rodata_entry+256]) + b'\x00'*256)
rsrc2 = code.compute_pgm_rsrc2 | (lds_size << 15)
return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data, rsrc2)
return None
except Exception as e:
print(f" Error getting kernel: {e}")
return None
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
def main():
import argparse
parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
args = parser.parse_args()
rust_remu = get_rust_remu()
if rust_remu is None:
print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
print("Running Python-only benchmarks...\n")
print("=" * 90)
print("RDNA3 Emulator Benchmark: Python vs Rust")
print("=" * 90)
results = []
print("\n[TINYGRAD KERNELS]")
print("-" * 90)
for op_name in TINYGRAD_TESTS:
print(f"\n{op_name}:", end=" ", flush=True)
kernel_info = get_tinygrad_kernel(op_name)
if kernel_info is None:
print("failed to compile")
continue
kernel, global_size, local_size, buf_sizes, buf_data, rsrc2 = kernel_info
n_insts = count_instructions(kernel)
n_workgroups = global_size[0] * global_size[1] * global_size[2]
n_threads = local_size[0] * local_size[1] * local_size[2]
total_work = n_insts * n_workgroups * n_threads
print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
set_valid_mem_ranges(ranges)
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations) if rust_remu else None
if py_time:
py_rate = total_work / py_time / 1e6
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
if rust_time:
rust_rate = total_work / rust_time / 1e6
speedup = py_time / rust_time if py_time else 0
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
results.append((op_name, n_insts, n_workgroups, py_time, rust_time))
# Summary table
print("\n" + "=" * 90)
print("SUMMARY")
print("=" * 90)
print(f"{'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
print("-" * 90)
for name, n_insts, n_wgs, py_time, rust_time in results:
py_ms = f"{py_time*1000:.3f}" if py_time else "error"
if rust_time:
rust_ms = f"{rust_time*1000:.3f}"
speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
else:
rust_ms, speedup = "N/A", "N/A"
print(f"{name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
if __name__ == "__main__":
main()
-855
View File
@@ -1,855 +0,0 @@
#!/usr/bin/env python3
"""SQTT InstOp discovery tool - finds instruction opcodes by running different instructions.
Requires profiling enabled:
echo 'profile_standard' | sudo tee /sys/class/drm/card1/device/power_dpm_force_performance_level
Run with: DEBUG=1 python extra/assembly/amd/test/discover_instops.py
For full traces: DEBUG=2 python extra/assembly/amd/test/discover_instops.py
"""
import os
os.environ["SQTT"] = "1"
os.environ["PROFILE"] = "1"
os.environ["SQTT_LIMIT_SE"] = "2" # Force work to traced SE only
os.environ["SQTT_TOKEN_EXCLUDE"] = "3784" # Exclude WAVERDY, REG, EVENT, UTILCTR, WAVEALLOC, PERF
from tinygrad.helpers import DEBUG, colored
from tinygrad.runtime.ops_amd import SQTT_SIMD_SEL
from extra.assembly.amd.autogen.rdna3.ins import (
# VALU - basic (these are safe, just register ops)
v_mov_b32_e32, v_add_f32_e32, v_mul_f32_e32,
v_and_b32_e32, v_or_b32_e32, v_xor_b32_e32,
v_lshlrev_b32_e32, v_lshrrev_b32_e32,
# VALU - transcendental
v_exp_f32_e32, v_log_f32_e32, v_rcp_f32_e32, v_sqrt_f32_e32,
v_sin_f32_e32, v_cos_f32_e32,
# VALU - 64-bit
v_lshlrev_b64, v_lshrrev_b64, v_ashrrev_i64,
v_add_f64, v_mul_f64, v_max_f64, v_min_f64,
v_fma_f64,
# VALU - 64-bit transcendental
v_rcp_f64_e32, v_rsq_f64_e32, v_sqrt_f64_e32,
v_trunc_f64_e32, v_ceil_f64_e32, v_floor_f64_e32, v_fract_f64_e32,
v_frexp_exp_i32_f64_e32, v_frexp_mant_f64_e32,
# VALU - div helpers
v_div_fixup_f32, v_div_fixup_f64, v_div_fmas_f32, v_div_fmas_f64, v_div_scale_f32,
# VALU - MAD64
v_mad_u64_u32, v_mad_i64_i32,
# VALU - compare (writes to VCC, safe)
v_cmp_eq_u32_e32,
# VALU - cmpx (modifies EXEC) - various types
v_cmpx_eq_u32_e32, v_cmpx_lt_u32_e32, v_cmpx_gt_u32_e32,
v_cmpx_eq_f32_e32, v_cmpx_lt_f32_e32,
v_cmpx_eq_i32_e32,
v_cmpx_class_f32_e32,
# VALU - readlane/writelane
v_readlane_b32, v_writelane_b32,
v_readfirstlane_b32_e32,
# SALU - basic (safe, just register ops)
s_mov_b32, s_add_u32, s_and_b32, s_or_b32,
s_lshl_b32, s_lshr_b32,
s_nop, s_endpgm, s_waitcnt,
# SALU - float
s_ceil_f32, s_floor_f32, s_trunc_f32,
# SALU - branch (safe if offset is 0 = next instruction)
s_branch, s_cbranch_scc0, s_cbranch_execz, s_cbranch_execnz,
# SALU - message
s_sendmsg,
# SALU - bit manipulation
s_brev_b32, s_bcnt1_i32_b32, s_ctz_i32_b32, s_clz_i32_u32,
# SALU - saveexec (modifies EXEC)
s_and_saveexec_b32, s_or_saveexec_b32, s_xor_saveexec_b32,
# SMEM - scalar memory (load from kernarg pointer in s[0:1])
s_load_b32, s_load_b64,
# GLOBAL - global memory (load/store) - various widths
global_load_u8, global_load_u16, global_load_b32, global_load_b64, global_load_b96, global_load_b128,
global_store_b8, global_store_b16, global_store_b32, global_store_b64, global_store_b96, global_store_b128,
# GLOBAL - atomics
global_atomic_add_u32, global_atomic_add_u64,
# FLAT - flat memory access
flat_load_b32, flat_load_b64, flat_load_b96, flat_load_b128,
flat_store_b8, flat_store_b16, flat_store_b32, flat_store_b64, flat_store_b96, flat_store_b128,
# LDS - local data share - various widths
ds_load_b32, ds_load_b64, ds_load_b128,
ds_store_b32, ds_store_b64, ds_store_b128,
# LDS - atomics
ds_add_u32, ds_max_u32, ds_min_u32,
# VOP3P - packed
v_pk_add_f16, v_pk_mul_f16, v_pk_fma_f16, v_pk_add_i16,
# VOP3 - misc
v_bfe_u32, v_bfi_b32, v_alignbit_b32, v_fma_f32,
v_add3_u32, v_xad_u32, v_lshl_or_b32, v_add_nc_u32_e32,
# VOP3 - carry-out
v_add_co_u32, v_add_co_ci_u32_e32,
# VOPD - dual issue
v_dual_add_f32, v_dual_mul_f32,
# VOP2 - fmac
v_fmac_f32_e32,
# DOT
v_dot2_f16_f16,
# WMMA
v_wmma_f32_16x16x16_f16, v_wmma_f16_16x16x16_f16, v_wmma_i32_16x16x16_iu8,
# Permlane ops
v_permlane64_b32_e32, v_permlane16_b32, v_permlanex16_b32,
# Interpolation
v_interp_p10_f32, v_interp_p2_f32,
# Barrier
s_barrier,
# SrcEnum for NULL soffset
SrcEnum,
)
from extra.assembly.amd.dsl import v, s
from extra.assembly.amd.sqtt import InstOp, INST, WAVESTART, WAVEEND, ALUEXEC, VMEMEXEC
from extra.assembly.amd.test.test_sqtt_hw import (
run_asm_sqtt, decode_all_blobs, get_inst_ops, print_blobs, get_wave_packets, format_packet, PACKET_COLORS, count_valuinst
)
# ═══════════════════════════════════════════════════════════════════════════════
# INSTRUCTION TEST CASES - only safe instructions that don't access memory
# ═══════════════════════════════════════════════════════════════════════════════
# Helper: load buffer address from kernarg (s[0:1] -> s[2:3])
# The runtime passes kernarg pointer in s[0:1], kernarg contains buffer address
def _load_buf_addr():
return [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL), # load buf addr from kernarg
s_waitcnt(lgkmcnt=0), # wait for SMEM load
]
INSTRUCTION_TESTS: dict[str, tuple[str, list]] = {
# SALU (0x0) - scalar ALU, just register operations
"SALU_mov": ("s_mov_b32", [s_mov_b32(s[4], 0), s_mov_b32(s[5], 1)]),
"SALU_add": ("s_add_u32", [s_mov_b32(s[4], 1), s_mov_b32(s[5], 2), s_add_u32(s[6], s[4], s[5])]),
"SALU_logic": ("s_and/or", [s_and_b32(s[6], s[4], s[5]), s_or_b32(s[7], s[4], s[5])]),
"SALU_shift": ("s_lshl/lshr", [s_lshl_b32(s[6], s[4], 1), s_lshr_b32(s[7], s[4], 1)]),
"SALU_nop": ("s_nop", [s_nop(0)]),
# JUMP (0x3) - branch taken
"JUMP_branch": ("s_branch", [s_branch(0)]),
"JUMP_cbranch_execnz": ("s_cbranch_execnz", [s_cbranch_execnz(0)]), # EXEC != 0, branch taken
# JUMP_NO (0x4) - branch not taken
"JUMP_NO_cbranch_execz": ("s_cbranch_execz", [s_cbranch_execz(0)]), # EXEC != 0, branch not taken
# VALU (0xb) - vector ALU, just register operations
"VALU_mov": ("v_mov_b32", [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 1.0)]),
"VALU_add": ("v_add_f32", [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])]),
"VALU_mul": ("v_mul_f32", [v_mul_f32_e32(v[2], v[0], v[1])]),
"VALU_logic": ("v_and/or/xor", [v_and_b32_e32(v[2], v[0], v[1]), v_or_b32_e32(v[3], v[0], v[1]), v_xor_b32_e32(v[4], v[0], v[1])]),
"VALU_shift": ("v_lshl/lshr", [v_lshlrev_b32_e32(v[2], 1, v[0]), v_lshrrev_b32_e32(v[3], 1, v[0])]),
# VALU transcendental - still just register ops
"VALU_exp": ("v_exp_f32", [v_mov_b32_e32(v[0], 1.0), v_exp_f32_e32(v[1], v[0])]),
"VALU_log": ("v_log_f32", [v_mov_b32_e32(v[0], 1.0), v_log_f32_e32(v[1], v[0])]),
"VALU_rcp": ("v_rcp_f32", [v_mov_b32_e32(v[0], 1.0), v_rcp_f32_e32(v[1], v[0])]),
"VALU_sqrt": ("v_sqrt_f32", [v_mov_b32_e32(v[0], 1.0), v_sqrt_f32_e32(v[1], v[0])]),
# VALU 64-bit shift (0xd)
"VALU64_lshl": ("v_lshlrev_b64", [v_lshlrev_b64(v[0:1], 1, v[2:3])]),
"VALU64_lshr": ("v_lshrrev_b64", [v_lshrrev_b64(v[0:1], 1, v[2:3])]),
"VALU64_ashr": ("v_ashrrev_i64", [v_ashrrev_i64(v[0:1], 1, v[2:3])]),
# VALU 64-bit arithmetic
"VALU64_add": ("v_add_f64", [v_add_f64(v[0:1], v[2:3], v[4:5])]),
"VALU64_mul": ("v_mul_f64", [v_mul_f64(v[0:1], v[2:3], v[4:5])]),
"VALU64_max": ("v_max_f64", [v_max_f64(v[0:1], v[2:3], v[4:5])]),
"VALU64_min": ("v_min_f64", [v_min_f64(v[0:1], v[2:3], v[4:5])]),
"VALU64_fma": ("v_fma_f64", [v_fma_f64(v[0:1], v[2:3], v[4:5], v[6:7])]),
# VALU 64-bit transcendental
"VALU64_rcp": ("v_rcp_f64", [v_rcp_f64_e32(v[0:1], v[2:3])]),
"VALU64_rsq": ("v_rsq_f64", [v_rsq_f64_e32(v[0:1], v[2:3])]),
"VALU64_sqrt": ("v_sqrt_f64", [v_sqrt_f64_e32(v[0:1], v[2:3])]),
# VALU 64-bit rounding
"VALU64_trunc": ("v_trunc_f64", [v_trunc_f64_e32(v[0:1], v[2:3])]),
"VALU64_ceil": ("v_ceil_f64", [v_ceil_f64_e32(v[0:1], v[2:3])]),
"VALU64_floor": ("v_floor_f64", [v_floor_f64_e32(v[0:1], v[2:3])]),
"VALU64_fract": ("v_fract_f64", [v_fract_f64_e32(v[0:1], v[2:3])]),
# VALU 64-bit frexp
"VALU64_frexp_exp": ("v_frexp_exp_i32_f64", [v_frexp_exp_i32_f64_e32(v[0], v[2:3])]),
"VALU64_frexp_mant": ("v_frexp_mant_f64", [v_frexp_mant_f64_e32(v[0:1], v[2:3])]),
# VALU 64-bit div helpers
"VALU64_div_fixup": ("v_div_fixup_f64", [v_div_fixup_f64(v[0:1], v[2:3], v[4:5], v[6:7])]),
"VALU64_div_fmas": ("v_div_fmas_f64", [v_div_fmas_f64(v[0:1], v[2:3], v[4:5], v[6:7])]),
# VALU 32-bit div helpers
"VALU_div_fixup": ("v_div_fixup_f32", [v_div_fixup_f32(v[0], v[1], v[2], v[3])]),
"VALU_div_fmas": ("v_div_fmas_f32", [v_div_fmas_f32(v[0], v[1], v[2], v[3])]),
"VALU_div_scale": ("v_div_scale_f32", [v_div_scale_f32(v[0], SrcEnum.VCC_LO, v[1], v[2], v[3])]),
# VALU MAD64 (0xe)
"VALU_mad64u": ("v_mad_u64_u32", [
v_mov_b32_e32(v[2], 2),
v_mov_b32_e32(v[3], 3),
v_mov_b32_e32(v[4], 0),
v_mov_b32_e32(v[5], 0),
v_mad_u64_u32(v[0:1], SrcEnum.NULL, v[2], v[3], v[4:5]),
]),
"VALU_mad64i": ("v_mad_i64_i32", [
v_mov_b32_e32(v[2], 2),
v_mov_b32_e32(v[3], 3),
v_mov_b32_e32(v[4], 0),
v_mov_b32_e32(v[5], 0),
v_mad_i64_i32(v[0:1], SrcEnum.NULL, v[2], v[3], v[4:5]),
]),
# VALU compare - writes to VCC
"VALU_cmp": ("v_cmp_eq_u32", [v_cmp_eq_u32_e32(v[0], v[1])]),
# VALU CMPX (0x73) - modifies EXEC
"VALU_cmpx_eq_u32": ("v_cmpx_eq_u32", [v_cmpx_eq_u32_e32(v[0], v[1])]),
# SALU saveexec (0x72) - modifies EXEC safely by ANDing with all-ones mask
"SALU_saveexec": ("s_and_saveexec_b32", [
s_mov_b32(s[5], 0xFFFFFFFF), # all lanes mask
s_and_saveexec_b32(s[4], s[5]), # EXEC = EXEC & 0xFFFFFFFF = EXEC (unchanged)
]),
# SALU float ops
"SALU_ceil": ("s_ceil_f32", [s_ceil_f32(s[4], s[5])]),
"SALU_floor": ("s_floor_f32", [s_floor_f32(s[4], s[5])]),
"SALU_trunc": ("s_trunc_f32", [s_trunc_f32(s[4], s[5])]),
# SALU bit ops
"SALU_brev": ("s_brev_b32", [s_brev_b32(s[4], s[5])]),
"SALU_bcnt1": ("s_bcnt1_i32_b32", [s_bcnt1_i32_b32(s[4], s[5])]),
"SALU_ctz": ("s_ctz_i32_b32", [s_ctz_i32_b32(s[4], s[5])]),
"SALU_clz": ("s_clz_i32_u32", [s_clz_i32_u32(s[4], s[5])]),
# VALU sin/cos
"VALU_sin": ("v_sin_f32", [v_sin_f32_e32(v[0], v[1])]),
"VALU_cos": ("v_cos_f32", [v_cos_f32_e32(v[0], v[1])]),
# VOP3P - packed operations
"VALU_pk_add_f16": ("v_pk_add_f16", [v_pk_add_f16(v[0], v[1], v[2])]),
"VALU_pk_mul_f16": ("v_pk_mul_f16", [v_pk_mul_f16(v[0], v[1], v[2])]),
"VALU_pk_fma_f16": ("v_pk_fma_f16", [v_pk_fma_f16(v[0], v[1], v[2], v[3])]),
"VALU_pk_add_i16": ("v_pk_add_i16", [v_pk_add_i16(v[0], v[1], v[2])]),
# VOP3 - misc
"VALU_bfe_u32": ("v_bfe_u32", [v_bfe_u32(v[0], v[1], 0, 8)]),
"VALU_bfi_b32": ("v_bfi_b32", [v_bfi_b32(v[0], v[1], v[2], v[3])]),
"VALU_alignbit": ("v_alignbit_b32", [v_alignbit_b32(v[0], v[1], v[2], 4)]),
"VALU_fma_f32": ("v_fma_f32", [v_fma_f32(v[0], v[1], v[2], v[3])]),
# VOP3 - integer add variants (used by tinygrad kernels)
"VALU_add3": ("v_add3_u32", [v_add3_u32(v[0], v[1], v[2], v[3])]),
"VALU_xad": ("v_xad_u32", [v_xad_u32(v[0], v[1], v[2], v[3])]),
"VALU_lshl_or": ("v_lshl_or_b32", [v_lshl_or_b32(v[0], v[1], 4, v[2])]),
"VALU_add_nc": ("v_add_nc_u32", [v_add_nc_u32_e32(v[0], v[1], v[2])]),
# VOP3 - carry-out adds (used for 64-bit address calculation)
"VALU_add_co": ("v_add_co_u32", [v_add_co_u32(v[0], SrcEnum.VCC_LO, v[1], v[2])]),
"VALU_add_co_ci": ("v_add_co_ci_u32", [v_add_co_ci_u32_e32(v[0], v[1], v[2])]),
# VOPD - dual issue (used by tinygrad kernels)
"VALU_dual_add": ("v_dual_add_f32", [v_dual_add_f32(v[0], v[1], v[2], v[3], v[4], v[5])]),
"VALU_dual_mul": ("v_dual_mul_f32", [v_dual_mul_f32(v[0], v[1], v[2], v[3], v[4], v[5])]),
# VOP2 - fmac
"VALU_fmac": ("v_fmac_f32", [v_fmac_f32_e32(v[0], v[1], v[0])]),
# DOT products
"VALU_dot2": ("v_dot2_f16_f16", [v_dot2_f16_f16(v[0], v[1], v[2], v[3])]),
# WMMA - wave matrix multiply accumulate
"VALU_wmma_f32_f16": ("v_wmma_f32_16x16x16_f16", [v_wmma_f32_16x16x16_f16(v[0:7], v[8:15], v[16:23], v[0:7])]),
"VALU_wmma_f16_f16": ("v_wmma_f16_16x16x16_f16", [v_wmma_f16_16x16x16_f16(v[0:7], v[8:15], v[16:23], v[0:7])]),
"VALU_wmma_i32_iu8": ("v_wmma_i32_16x16x16_iu8", [v_wmma_i32_16x16x16_iu8(v[0:7], v[8:11], v[12:15], v[0:7])]),
# Permlane operations - cross-lane data movement
# NOTE: permlane64 produces NO SQTT packets in wave32 mode (it's for wave64 pairs)
# NOTE: permlane16/x16 produce VALUINST packets (no specific InstOp)
"VALU_permlane16": ("v_permlane16_b32", [v_permlane16_b32(v[0], v[1], s[2], s[3])]),
"VALU_permlanex16": ("v_permlanex16_b32", [v_permlanex16_b32(v[0], v[1], s[2], s[3])]),
# Interpolation - used in graphics shaders (produces InstOp 0x12 VINTERP)
"VINTERP_p10": ("v_interp_p10_f32", [v_interp_p10_f32(v[0], v[1], v[2], v[3])]),
"VINTERP_p2": ("v_interp_p2_f32", [v_interp_p2_f32(v[0], v[1], v[2], v[3])]),
# Barrier - wave synchronization
# NOTE: s_barrier produces NO SQTT instruction packets (with 1 wave, it's essentially a no-op)
"SALU_barrier": ("s_barrier", [s_barrier()]),
# LDS atomics
"LDS_atomic_add": ("ds_add_u32", [
v_mov_b32_e32(v[0], 0), # LDS address
v_mov_b32_e32(v[1], 1), # data to add
ds_add_u32(addr=v[0], data0=v[1]),
s_waitcnt(lgkmcnt=0),
]),
# ═══════════════════════════════════════════════════════════════════════════════
# GLOBAL ATOMICS - access real buffer passed via kernarg
# ═══════════════════════════════════════════════════════════════════════════════
# GLOBAL atomic add 32-bit (0x28 GLOBAL_ATOMIC)
"GLOBAL_atomic_add": ("global_atomic_add_u32", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL), # load buf addr from kernarg
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0), # offset = 0
v_mov_b32_e32(v[1], 1), # data to add
global_atomic_add_u32(addr=v[0], data=v[1], saddr=s[2]),
s_waitcnt(vmcnt=0),
]),
# GLOBAL atomic add 64-bit
"GLOBAL_atomic_add64": ("global_atomic_add_u64", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[2], 1),
v_mov_b32_e32(v[3], 0),
global_atomic_add_u64(addr=v[0], data=v[2:3], saddr=s[2]),
s_waitcnt(vmcnt=0),
]),
# ═══════════════════════════════════════════════════════════════════════════════
# MEMORY INSTRUCTIONS - access real buffer passed via kernarg
# ═══════════════════════════════════════════════════════════════════════════════
# SMEM (0x1) - scalar memory load from buffer
"SMEM_load": ("s_load_b32", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL), # load buf addr from kernarg
s_waitcnt(lgkmcnt=0),
s_load_b32(s[4], s[2], 0, soffset=SrcEnum.NULL), # load from buffer
s_waitcnt(lgkmcnt=0),
]),
# GLOBAL load (0x21 GLOBAL_LOAD) - global memory load
"GLOBAL_load": ("global_load_b32", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL), # load buf addr from kernarg
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0), # offset = 0
global_load_b32(v[1], addr=v[0], saddr=s[2], offset=0), # load from buffer
s_waitcnt(vmcnt=0),
]),
# GLOBAL store (0x24 GLOBAL_STORE) - global memory store
"GLOBAL_store": ("global_store_b32", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL), # load buf addr from kernarg
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0), # offset = 0
v_mov_b32_e32(v[1], 42), # data to store
global_store_b32(addr=v[0], data=v[1], saddr=s[2], offset=0), # store to buffer
s_waitcnt(vmcnt=0),
]),
# GLOBAL 8-bit load/store
"GLOBAL_load8": ("global_load_u8", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
global_load_u8(v[1], addr=v[0], saddr=s[2], offset=0),
s_waitcnt(vmcnt=0),
]),
"GLOBAL_store8": ("global_store_b8", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 42),
global_store_b8(addr=v[0], data=v[1], saddr=s[2], offset=0),
s_waitcnt(vmcnt=0),
]),
# GLOBAL 16-bit load/store
"GLOBAL_load16": ("global_load_u16", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
global_load_u16(v[1], addr=v[0], saddr=s[2], offset=0),
s_waitcnt(vmcnt=0),
]),
"GLOBAL_store16": ("global_store_b16", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 42),
global_store_b16(addr=v[0], data=v[1], saddr=s[2], offset=0),
s_waitcnt(vmcnt=0),
]),
# LDS load (0x29 LDS_LOAD) - local data share read
"LDS_load": ("ds_load_b32", [
v_mov_b32_e32(v[0], 0), # LDS address = 0
ds_load_b32(v[1], v[0], offset=0), # read from LDS
s_waitcnt(lgkmcnt=0),
]),
# LDS store (0x2b LDS_STORE) - local data share write
"LDS_store": ("ds_store_b32", [
v_mov_b32_e32(v[0], 0), # LDS address = 0
v_mov_b32_e32(v[1], 42), # data to store
ds_store_b32(v[0], v[1], offset=0), # write to LDS
s_waitcnt(lgkmcnt=0),
]),
# ═══════════════════════════════════════════════════════════════════════════════
# WIDER MEMORY OPERATIONS - to discover more InstOp variants
# ═══════════════════════════════════════════════════════════════════════════════
# GLOBAL 64-bit load
"GLOBAL_load64": ("global_load_b64", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
global_load_b64(v[2:3], addr=v[0], saddr=s[2], offset=0),
s_waitcnt(vmcnt=0),
]),
# GLOBAL 96-bit load
"GLOBAL_load96": ("global_load_b96", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
global_load_b96(v[4:6], addr=v[0], saddr=s[2], offset=0),
s_waitcnt(vmcnt=0),
]),
# GLOBAL 128-bit load
"GLOBAL_load128": ("global_load_b128", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
global_load_b128(v[4:7], addr=v[0], saddr=s[2], offset=0),
s_waitcnt(vmcnt=0),
]),
# GLOBAL 64-bit store
"GLOBAL_store64": ("global_store_b64", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[2], 42),
v_mov_b32_e32(v[3], 43),
global_store_b64(addr=v[0], data=v[2:3], saddr=s[2], offset=0),
s_waitcnt(vmcnt=0),
]),
# GLOBAL 96-bit store
"GLOBAL_store96": ("global_store_b96", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[4], 42),
v_mov_b32_e32(v[5], 43),
v_mov_b32_e32(v[6], 44),
global_store_b96(addr=v[0], data=v[4:6], saddr=s[2], offset=0),
s_waitcnt(vmcnt=0),
]),
# GLOBAL 128-bit store
"GLOBAL_store128": ("global_store_b128", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[4], 42),
v_mov_b32_e32(v[5], 43),
v_mov_b32_e32(v[6], 44),
v_mov_b32_e32(v[7], 45),
global_store_b128(addr=v[0], data=v[4:7], saddr=s[2], offset=0),
s_waitcnt(vmcnt=0),
]),
# ═══════════════════════════════════════════════════════════════════════════════
# GLOBAL VADDR (vector-only addressing, saddr=NULL) - used by tinygrad kernels
# ═══════════════════════════════════════════════════════════════════════════════
# GLOBAL VADDR load (all sizes use same opcode 0x22)
"GLOBAL_VADDR_load": ("global_load_b32 vaddr", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
global_load_b32(v[4], addr=v[0:1], saddr=SrcEnum.NULL, offset=0),
s_waitcnt(vmcnt=0),
]),
"GLOBAL_VADDR_load128": ("global_load_b128 vaddr", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
global_load_b128(v[4:7], addr=v[0:1], saddr=SrcEnum.NULL, offset=0),
s_waitcnt(vmcnt=0),
]),
# GLOBAL VADDR stores (size encoded: 32->0x25, 64->0x26, 96->0x27, 128->0x28)
"GLOBAL_VADDR_store": ("global_store_b32 vaddr", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
v_mov_b32_e32(v[4], 42),
global_store_b32(addr=v[0:1], data=v[4], saddr=SrcEnum.NULL, offset=0),
s_waitcnt(vmcnt=0),
]),
"GLOBAL_VADDR_store64": ("global_store_b64 vaddr", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
v_mov_b32_e32(v[4], 42),
v_mov_b32_e32(v[5], 43),
global_store_b64(addr=v[0:1], data=v[4:5], saddr=SrcEnum.NULL, offset=0),
s_waitcnt(vmcnt=0),
]),
"GLOBAL_VADDR_store96": ("global_store_b96 vaddr", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
v_mov_b32_e32(v[4], 42),
v_mov_b32_e32(v[5], 43),
v_mov_b32_e32(v[6], 44),
global_store_b96(addr=v[0:1], data=v[4:6], saddr=SrcEnum.NULL, offset=0),
s_waitcnt(vmcnt=0),
]),
"GLOBAL_VADDR_store128": ("global_store_b128 vaddr", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
v_mov_b32_e32(v[4], 42),
v_mov_b32_e32(v[5], 43),
v_mov_b32_e32(v[6], 44),
v_mov_b32_e32(v[7], 45),
global_store_b128(addr=v[0:1], data=v[4:7], saddr=SrcEnum.NULL, offset=0),
s_waitcnt(vmcnt=0),
]),
# LDS 64-bit load
"LDS_load64": ("ds_load_b64", [
v_mov_b32_e32(v[0], 0),
ds_load_b64(v[2:3], v[0], offset=0),
s_waitcnt(lgkmcnt=0),
]),
# LDS 128-bit load
"LDS_load128": ("ds_load_b128", [
v_mov_b32_e32(v[0], 0),
ds_load_b128(v[4:7], v[0], offset=0),
s_waitcnt(lgkmcnt=0),
]),
# LDS 64-bit store
"LDS_store64": ("ds_store_b64", [
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[2], 42),
v_mov_b32_e32(v[3], 43),
ds_store_b64(v[0], v[2:3], offset=0),
s_waitcnt(lgkmcnt=0),
]),
# LDS 128-bit store
"LDS_store128": ("ds_store_b128", [
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[4], 42),
v_mov_b32_e32(v[5], 43),
v_mov_b32_e32(v[6], 44),
v_mov_b32_e32(v[7], 45),
ds_store_b128(v[0], v[4:7], offset=0),
s_waitcnt(lgkmcnt=0),
]),
# MESSAGE (0x9) - s_sendmsg
"MESSAGE": ("s_sendmsg", [
s_sendmsg(0), # send message 0 (NOP message)
]),
# ═══════════════════════════════════════════════════════════════════════════════
# FLAT MEMORY - uses 64-bit virtual address in VGPRs
# ═══════════════════════════════════════════════════════════════════════════════
# FLAT load - load using 64-bit address from buffer
"FLAT_load": ("flat_load_b32", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL), # load buf addr from kernarg
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]), # addr lo
v_mov_b32_e32(v[1], s[3]), # addr hi
flat_load_b32(v[2], addr=v[0:1]),
s_waitcnt(vmcnt=0, lgkmcnt=0),
]),
# FLAT store
"FLAT_store": ("flat_store_b32", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
v_mov_b32_e32(v[2], 42),
flat_store_b32(addr=v[0:1], data=v[2]),
s_waitcnt(vmcnt=0, lgkmcnt=0),
]),
# FLAT 64-bit
"FLAT_load64": ("flat_load_b64", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
flat_load_b64(v[2:3], addr=v[0:1]),
s_waitcnt(vmcnt=0, lgkmcnt=0),
]),
"FLAT_store64": ("flat_store_b64", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
v_mov_b32_e32(v[4], 42),
v_mov_b32_e32(v[5], 43),
flat_store_b64(addr=v[0:1], data=v[4:5]),
s_waitcnt(vmcnt=0, lgkmcnt=0),
]),
# FLAT 96-bit
"FLAT_load96": ("flat_load_b96", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
flat_load_b96(v[4:6], addr=v[0:1]),
s_waitcnt(vmcnt=0, lgkmcnt=0),
]),
"FLAT_store96": ("flat_store_b96", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
v_mov_b32_e32(v[4], 42),
v_mov_b32_e32(v[5], 43),
v_mov_b32_e32(v[6], 44),
flat_store_b96(addr=v[0:1], data=v[4:6]),
s_waitcnt(vmcnt=0, lgkmcnt=0),
]),
# FLAT 128-bit
"FLAT_load128": ("flat_load_b128", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
flat_load_b128(v[4:7], addr=v[0:1]),
s_waitcnt(vmcnt=0, lgkmcnt=0),
]),
"FLAT_store128": ("flat_store_b128", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
v_mov_b32_e32(v[4], 42),
v_mov_b32_e32(v[5], 43),
v_mov_b32_e32(v[6], 44),
v_mov_b32_e32(v[7], 45),
flat_store_b128(addr=v[0:1], data=v[4:7]),
s_waitcnt(vmcnt=0, lgkmcnt=0),
]),
# FLAT 8/16-bit stores
"FLAT_store8": ("flat_store_b8", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
v_mov_b32_e32(v[2], 42),
flat_store_b8(addr=v[0:1], data=v[2]),
s_waitcnt(vmcnt=0, lgkmcnt=0),
]),
"FLAT_store16": ("flat_store_b16", [
s_load_b64(s[2:3], s[0], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
v_mov_b32_e32(v[2], 42),
flat_store_b16(addr=v[0:1], data=v[2]),
s_waitcnt(vmcnt=0, lgkmcnt=0),
]),
}
def run_with_retry(instructions: list, max_attempts: int = 20) -> tuple[list[tuple[int, list[bytes]]], list[list], set, int]:
"""Run instructions multiple times to collect InstOp variants.
Memory ops produce different InstOp values (0x2x vs 0x5x) depending on which SIMD executes them:
- 0x2x range: wave ran on traced SIMD (matched)
- 0x5x range: wave ran on other SIMD (not matched)
Returns list of (traced_simd, blobs) tuples, all_packets, all_ops, max_valuinst_count.
"""
all_ops = set()
all_runs: list[tuple[int, list[bytes]]] = []
all_packets = []
max_valuinst = 0
SQTT_SIMD_SEL.value = 0 # only trace SIMD 0
for _ in range(max_attempts):
blobs = run_asm_sqtt(instructions)
packets = decode_all_blobs(blobs)
# get ops and valuinst from all SIMDs
ops = set()
valuinst_count = 0
for simd in [0, 1, 2, 3]:
ops.update(get_inst_ops(packets, traced_simd=simd))
valuinst_count = max(valuinst_count, count_valuinst(packets, traced_simd=simd))
all_runs.append((0, blobs))
all_packets.append(packets)
all_ops.update(ops)
max_valuinst = max(max_valuinst, valuinst_count)
return all_runs, all_packets, all_ops, max_valuinst
def discover_all_instops() -> tuple[dict[int, set[str]], dict[str, Exception], dict[str, int]]:
"""Run all instruction tests and collect InstOp values."""
discovered: dict[int, set[str]] = {}
failures: dict[str, Exception] = {}
valuinst_tests: dict[str, int] = {} # tests that produced VALUINST packets
for test_name, (instr_name, instructions) in INSTRUCTION_TESTS.items():
try:
all_runs, _, ops, valuinst_count = run_with_retry(instructions)
for op in ops:
if op not in discovered:
discovered[op] = set()
discovered[op].add(f"{test_name}")
if valuinst_count > 0:
valuinst_tests[test_name] = valuinst_count
if DEBUG >= 2:
print(f"\n{''*60}")
print(f"{test_name} ({instr_name}): ops={[hex(op) for op in sorted(ops)]}")
# collect wave patterns from traced SIMD runs (group by exact timing)
patterns: dict[tuple, list] = {} # pattern (types + timing) -> list of (wave_packets, t0)
for traced_simd, blobs in all_runs:
for blob in blobs:
packets = decode_all_blobs([blob])
wave_packets = get_wave_packets(packets)
# only include runs where wave ran on traced SIMD
ws = next((p for p in wave_packets if isinstance(p, WAVESTART)), None)
if ws and ws.simd == traced_simd and wave_packets:
t0 = wave_packets[0]._time
# pattern includes types AND normalized timing
pattern = tuple((type(p).__name__, p._time - t0) for p in wave_packets)
if pattern not in patterns:
patterns[pattern] = []
patterns[pattern].append((wave_packets, t0))
if patterns:
counts = {p: len(runs) for p, runs in patterns.items()}
most_common = max(counts, key=counts.get)
count = counts[most_common]
total = sum(counts.values())
print(f"\n=== most common pattern ({count}/{total} runs) ===")
wave_packets, t0 = patterns[most_common][0]
last_time = t0
for p in wave_packets:
print(format_packet(p, last_time, t0))
last_time = p._time
if len(patterns) > 1:
print(f"\n variations: {len(patterns)} unique timing patterns")
if DEBUG >= 3:
for traced_simd, blobs in all_runs:
print(f"\n=== traced simd={traced_simd} ===")
print_blobs(blobs, wave_only=False)
if DEBUG >= 1:
status = colored("", "green") if ops else (colored("V", "cyan") if valuinst_count > 0 else colored("", "yellow"))
ops_str = ", ".join(hex(op) for op in sorted(ops)) if ops else "none"
valuinst_str = f" valuinst={valuinst_count}" if valuinst_count > 0 and not ops else ""
print(f" {status} {test_name:25s} ops=[{ops_str}]{valuinst_str}")
except Exception as e:
failures[test_name] = e
if DEBUG >= 1:
print(f" {colored('', 'red')} {test_name:25s} FAILED: {e}")
return discovered, failures, valuinst_tests
def print_summary(discovered: dict[int, set[str]], failures: dict[str, Exception], valuinst_tests: dict[str, int]) -> None:
"""Print discovery summary."""
known_ops = {e.value for e in InstOp}
discovered_ops = set(discovered.keys())
print("\n" + "=" * 60)
print("DISCOVERED INSTOP VALUES")
print("=" * 60)
for op in sorted(discovered_ops):
try:
name = InstOp(op).name
status = colored("known", "green")
except ValueError:
name = f"UNKNOWN"
status = colored("NEW!", "yellow")
sources = ", ".join(sorted(discovered[op]))
print(f" 0x{op:02x} {name:20s} ({status}) <- {sources}")
# VALUINST tests (instructions that only produce VALUINST, not INST packets)
valuinst_only = {k: v for k, v in valuinst_tests.items() if not any(k in tests for tests in discovered.values())}
if valuinst_only:
print("\n" + "=" * 60)
print(colored("VALUINST-ONLY INSTRUCTIONS (no InstOp, use VALUINST packet)", "cyan"))
print("=" * 60)
for test_name, count in sorted(valuinst_only.items()):
print(f" {test_name}: {count} VALUINST packets")
# Missing from enum
missing = known_ops - discovered_ops
if missing:
print("\n" + "=" * 60)
print("ENUM VALUES NOT DISCOVERED")
print("=" * 60)
print("(need memory ops: SMEM, VMEM, LDS)")
for op in sorted(missing):
print(f" 0x{op:02x} {InstOp(op).name}")
# New values to add
new_ops = discovered_ops - known_ops
if new_ops:
print("\n" + "=" * 60)
print(colored("NEW INSTOP VALUES TO ADD TO ENUM", "yellow"))
print("=" * 60)
for op in sorted(new_ops):
sources = ", ".join(sorted(discovered[op]))
print(f" {op:#04x}: \"{sources}\",")
# Stats
print("\n" + "=" * 60)
print("STATISTICS")
print("=" * 60)
print(f" Tests run: {len(INSTRUCTION_TESTS)}")
print(f" Tests passed: {len(INSTRUCTION_TESTS) - len(failures)}")
print(f" Tests failed: {len(failures)}")
print(f" Known ops: {len(known_ops)}")
print(f" Discovered: {len(discovered_ops)}")
if known_ops:
print(f" Coverage: {len(discovered_ops & known_ops)}/{len(known_ops)} ({100*len(discovered_ops & known_ops)//len(known_ops)}%)")
print(f" New ops found: {len(new_ops)}")
print(f" VALUINST-only: {len(valuinst_only)}")
if __name__ == "__main__":
print("=" * 60)
print("SQTT InstOp Discovery Tool")
print("=" * 60)
print(f"Testing {len(INSTRUCTION_TESTS)} instruction categories...\n")
discovered, failures, valuinst_tests = discover_all_instops()
print_summary(discovered, failures, valuinst_tests)
@@ -1,289 +0,0 @@
#!/usr/bin/env python3
"""SQTT InstOp discovery from tinygrad-generated kernels.
Runs various tinygrad operations and captures SQTT traces to find new InstOp values.
Requires profiling enabled:
echo 'profile_standard' | sudo tee /sys/class/drm/card1/device/power_dpm_force_performance_level
Run with: DEBUG=1 python extra/assembly/amd/test/discover_instops_tensor.py
For full traces: DEBUG=2 python extra/assembly/amd/test/discover_instops_tensor.py
"""
import os
os.environ["SQTT"] = "1"
os.environ["PROFILE"] = "1"
os.environ["SQTT_LIMIT_SE"] = "2" # Force work to traced SE only
os.environ["SQTT_TOKEN_EXCLUDE"] = "3784" # Exclude noisy packet types
from tinygrad import Tensor, dtypes, Device
from tinygrad.helpers import DEBUG, colored
from tinygrad.runtime.ops_amd import ProfileSQTTEvent, SQTT_SIMD_SEL
from extra.assembly.amd.sqtt import InstOp, decode, INST, WAVESTART, WAVEEND
# ═══════════════════════════════════════════════════════════════════════════════
# HELPERS
# ═══════════════════════════════════════════════════════════════════════════════
def get_inst_ops_from_blobs(blobs: list[bytes]) -> set[int]:
"""Extract all InstOp values from SQTT blobs."""
ops = set()
for blob in blobs:
packets = decode(blob)
in_wave = False
for p in packets:
if isinstance(p, WAVESTART):
in_wave = True
if in_wave and isinstance(p, INST):
ops.add(p.op if isinstance(p.op, int) else p.op.value)
if isinstance(p, WAVEEND):
in_wave = False
return ops
def run_and_capture(fn, attempts: int = 5) -> tuple[set[int], list[bytes]]:
"""Run a function multiple times and collect SQTT traces."""
dev = Device["AMD"]
all_ops = set()
all_blobs = []
SQTT_SIMD_SEL.value = 0
for _ in range(attempts):
dev.profile_events.clear()
fn()
blobs = [ev.blob for ev in dev.profile_events if isinstance(ev, ProfileSQTTEvent)]
ops = get_inst_ops_from_blobs(blobs)
all_ops.update(ops)
all_blobs.extend(blobs)
return all_ops, all_blobs
# ═══════════════════════════════════════════════════════════════════════════════
# TENSOR OPERATIONS TO TEST
# ═══════════════════════════════════════════════════════════════════════════════
TENSOR_TESTS: dict[str, tuple[str, callable]] = {
# Basic arithmetic
"add_f32": ("tensor add float32", lambda: (Tensor.rand(1024) + Tensor.rand(1024)).realize()),
"mul_f32": ("tensor mul float32", lambda: (Tensor.rand(1024) * Tensor.rand(1024)).realize()),
"sub_f32": ("tensor sub float32", lambda: (Tensor.rand(1024) - Tensor.rand(1024)).realize()),
"div_f32": ("tensor div float32", lambda: (Tensor.rand(1024) / (Tensor.rand(1024) + 0.1)).realize()),
# Transcendental
"exp_f32": ("tensor exp float32", lambda: Tensor.rand(1024).exp().realize()),
"log_f32": ("tensor log float32", lambda: (Tensor.rand(1024) + 0.1).log().realize()),
"sqrt_f32": ("tensor sqrt float32", lambda: Tensor.rand(1024).sqrt().realize()),
"sin_f32": ("tensor sin float32", lambda: Tensor.rand(1024).sin().realize()),
"cos_f32": ("tensor cos float32", lambda: Tensor.rand(1024).cos().realize()),
"tanh_f32": ("tensor tanh float32", lambda: Tensor.rand(1024).tanh().realize()),
"sigmoid_f32": ("tensor sigmoid float32", lambda: Tensor.rand(1024).sigmoid().realize()),
# Reductions
"sum_f32": ("tensor sum float32", lambda: Tensor.rand(1024).sum().realize()),
"max_f32": ("tensor max float32", lambda: Tensor.rand(1024).max().realize()),
"mean_f32": ("tensor mean float32", lambda: Tensor.rand(1024).mean().realize()),
# Matmul - small
"matmul_small": ("matmul 32x32", lambda: (Tensor.rand(32, 32) @ Tensor.rand(32, 32)).realize()),
# Matmul - medium (might use WMMA)
"matmul_medium": ("matmul 128x128", lambda: (Tensor.rand(128, 128) @ Tensor.rand(128, 128)).realize()),
# Matmul - larger (more likely to use WMMA)
"matmul_large": ("matmul 256x256", lambda: (Tensor.rand(256, 256) @ Tensor.rand(256, 256)).realize()),
# Different dtypes
"add_f16": ("tensor add float16", lambda: (Tensor.rand(1024, dtype=dtypes.float16) + Tensor.rand(1024, dtype=dtypes.float16)).realize()),
"mul_f16": ("tensor mul float16", lambda: (Tensor.rand(1024, dtype=dtypes.float16) * Tensor.rand(1024, dtype=dtypes.float16)).realize()),
"matmul_f16": ("matmul float16 128x128", lambda: (Tensor.rand(128, 128, dtype=dtypes.float16) @ Tensor.rand(128, 128, dtype=dtypes.float16)).realize()),
# Integer ops
"add_i32": ("tensor add int32", lambda: (Tensor.randint(1024, high=1000) + Tensor.randint(1024, high=1000)).realize()),
"mul_i32": ("tensor mul int32", lambda: (Tensor.randint(1024, high=100) * Tensor.randint(1024, high=100)).realize()),
# Bitwise
"and_i32": ("tensor bitwise and", lambda: (Tensor.randint(1024, high=1000) & Tensor.randint(1024, high=1000)).realize()),
"or_i32": ("tensor bitwise or", lambda: (Tensor.randint(1024, high=1000) | Tensor.randint(1024, high=1000)).realize()),
"xor_i32": ("tensor bitwise xor", lambda: (Tensor.randint(1024, high=1000) ^ Tensor.randint(1024, high=1000)).realize()),
"lshift_i32": ("tensor left shift", lambda: (Tensor.randint(1024, high=1000) << 2).realize()),
"rshift_i32": ("tensor right shift", lambda: (Tensor.randint(1024, high=1000) >> 2).realize()),
# Comparisons
"cmp_eq": ("tensor compare eq", lambda: (Tensor.rand(1024) == 0.5).realize()),
"cmp_lt": ("tensor compare lt", lambda: (Tensor.rand(1024) < 0.5).realize()),
"cmp_gt": ("tensor compare gt", lambda: (Tensor.rand(1024) > 0.5).realize()),
# Where/select
"where": ("tensor where", lambda: Tensor.rand(1024).where(Tensor.rand(1024), Tensor.rand(1024)).realize()),
# Reshaping/movement (may not generate interesting ops but let's check)
"reshape": ("tensor reshape", lambda: Tensor.rand(32, 32).reshape(16, 64).realize()),
"permute": ("tensor permute", lambda: Tensor.rand(32, 32).permute(1, 0).contiguous().realize()),
"expand": ("tensor expand", lambda: Tensor.rand(1, 32).expand(32, 32).contiguous().realize()),
# Pad
"pad": ("tensor pad", lambda: Tensor.rand(30, 30).pad(((1, 1), (1, 1))).realize()),
# Conv2D - small
"conv2d_small": ("conv2d 3x3", lambda: Tensor.rand(1, 3, 32, 32).conv2d(Tensor.rand(8, 3, 3, 3)).realize()),
# Conv2D - larger
"conv2d_medium": ("conv2d 3x3 64ch", lambda: Tensor.rand(1, 64, 32, 32).conv2d(Tensor.rand(64, 64, 3, 3)).realize()),
# Pooling
"maxpool": ("max pool 2x2", lambda: Tensor.rand(1, 3, 32, 32).max_pool2d((2, 2)).realize()),
"avgpool": ("avg pool 2x2", lambda: Tensor.rand(1, 3, 32, 32).avg_pool2d((2, 2)).realize()),
# Softmax
"softmax": ("softmax", lambda: Tensor.rand(32, 128).softmax().realize()),
# LayerNorm-like
"layernorm": ("layer norm pattern", lambda: _layernorm(Tensor.rand(32, 128))),
# BatchNorm-like
"batchnorm": ("batch norm pattern", lambda: _batchnorm(Tensor.rand(1, 64, 32, 32))),
# Dropout-like (during training)
"dropout": ("dropout pattern", lambda: (Tensor.rand(1024) * (Tensor.rand(1024) > 0.5)).realize()),
# Cast operations
"cast_f32_to_f16": ("cast f32->f16", lambda: Tensor.rand(1024).cast(dtypes.float16).realize()),
"cast_f16_to_f32": ("cast f16->f32", lambda: Tensor.rand(1024, dtype=dtypes.float16).cast(dtypes.float32).realize()),
"cast_f32_to_i32": ("cast f32->i32", lambda: (Tensor.rand(1024) * 100).cast(dtypes.int32).realize()),
"cast_i32_to_f32": ("cast i32->f32", lambda: Tensor.randint(1024, high=100).cast(dtypes.float32).realize()),
# Clamp/clip
"clamp": ("tensor clamp", lambda: Tensor.rand(1024).clamp(0.2, 0.8).realize()),
# Abs/neg
"abs": ("tensor abs", lambda: (Tensor.rand(1024) - 0.5).abs().realize()),
"neg": ("tensor neg", lambda: (-Tensor.rand(1024)).realize()),
# Reciprocal
"recip": ("tensor reciprocal", lambda: (Tensor.rand(1024) + 0.1).reciprocal().realize()),
# Power
"pow2": ("tensor pow 2", lambda: (Tensor.rand(1024) ** 2).realize()),
"pow3": ("tensor pow 3", lambda: (Tensor.rand(1024) ** 3).realize()),
}
def _layernorm(x: Tensor) -> Tensor:
"""Simple layer normalization pattern."""
mean = x.mean(axis=-1, keepdim=True)
var = ((x - mean) ** 2).mean(axis=-1, keepdim=True)
return ((x - mean) / (var + 1e-5).sqrt()).realize()
def _batchnorm(x: Tensor) -> Tensor:
"""Simple batch normalization pattern."""
mean = x.mean(axis=(0, 2, 3), keepdim=True)
var = ((x - mean) ** 2).mean(axis=(0, 2, 3), keepdim=True)
return ((x - mean) / (var + 1e-5).sqrt()).realize()
# ═══════════════════════════════════════════════════════════════════════════════
# DISCOVERY
# ═══════════════════════════════════════════════════════════════════════════════
def discover_all_instops() -> tuple[dict[int, set[str]], dict[str, Exception]]:
"""Run all tensor tests and collect InstOp values."""
discovered: dict[int, set[str]] = {}
failures: dict[str, Exception] = {}
for test_name, (desc, fn) in TENSOR_TESTS.items():
try:
ops, blobs = run_and_capture(fn)
for op in ops:
if op not in discovered:
discovered[op] = set()
discovered[op].add(test_name)
if DEBUG >= 1:
status = colored("", "green") if ops else colored("", "yellow")
ops_str = ", ".join(hex(op) for op in sorted(ops)) if ops else "none"
print(f" {status} {test_name:25s} [{desc:25s}] ops=[{ops_str}]")
if DEBUG >= 2 and blobs:
# Show first wave trace
for blob in blobs[:1]:
packets = decode(blob)
print(f" First blob: {len(blob)} bytes, {len(packets)} packets")
except Exception as e:
failures[test_name] = e
if DEBUG >= 1:
print(f" {colored('', 'red')} {test_name:25s} FAILED: {e}")
return discovered, failures
def print_summary(discovered: dict[int, set[str]], failures: dict[str, Exception]) -> None:
"""Print discovery summary."""
known_ops = {e.value for e in InstOp}
discovered_ops = set(discovered.keys())
print("\n" + "=" * 70)
print("DISCOVERED INSTOP VALUES FROM TINYGRAD KERNELS")
print("=" * 70)
for op in sorted(discovered_ops):
try:
name = InstOp(op).name
status = colored("known", "green")
except ValueError:
name = "UNKNOWN"
status = colored("NEW!", "yellow")
sources = ", ".join(sorted(discovered[op]))
# Truncate sources if too long
if len(sources) > 60:
sources = sources[:57] + "..."
print(f" 0x{op:02x} {name:20s} ({status}) <- {sources}")
# New values to add
new_ops = discovered_ops - known_ops
if new_ops:
print("\n" + "=" * 70)
print(colored("NEW INSTOP VALUES TO ADD TO ENUM", "yellow"))
print("=" * 70)
for op in sorted(new_ops):
sources = ", ".join(sorted(discovered[op]))
print(f" 0x{op:02x}: discovered from [{sources}]")
# Missing from enum (not discovered)
missing = known_ops - discovered_ops
if missing:
print("\n" + "=" * 70)
print("ENUM VALUES NOT DISCOVERED (may need specific instruction patterns)")
print("=" * 70)
for op in sorted(missing):
print(f" 0x{op:02x} {InstOp(op).name}")
# Stats
print("\n" + "=" * 70)
print("STATISTICS")
print("=" * 70)
print(f" Tests run: {len(TENSOR_TESTS)}")
print(f" Tests passed: {len(TENSOR_TESTS) - len(failures)}")
print(f" Tests failed: {len(failures)}")
print(f" Known ops: {len(known_ops)}")
print(f" Discovered: {len(discovered_ops)}")
if known_ops:
coverage = len(discovered_ops & known_ops)
print(f" Coverage: {coverage}/{len(known_ops)} ({100*coverage//len(known_ops)}%)")
print(f" New ops found: {len(new_ops)}")
if failures:
print("\n" + "=" * 70)
print("FAILURES")
print("=" * 70)
for name, e in failures.items():
print(f" {name}: {e}")
if __name__ == "__main__":
print("=" * 70)
print("SQTT InstOp Discovery from Tinygrad Kernels")
print("=" * 70)
print(f"Testing {len(TENSOR_TESTS)} tensor operations...\n")
discovered, failures = discover_all_instops()
print_summary(discovered, failures)
@@ -1,196 +0,0 @@
# Usability tests for the RDNA3 ASM DSL
# These tests demonstrate how the DSL *should* work for a good user experience
# Currently many of these tests fail - they document desired behavior
import unittest
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import Inst, RawImm, SGPR, VGPR
class TestRegisterSliceSyntax(unittest.TestCase):
"""
Issue: Register slice syntax should use AMD assembly convention (inclusive end).
In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
The DSL should match this convention so that:
- s[4:7] gives 4 registers
- Disassembler output can be copied directly back into DSL code
Fix: Change _RegFactory.__getitem__ to use inclusive end:
key.stop - key.start + 1 (instead of key.stop - key.start)
"""
def test_register_slice_count(self):
# s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
reg = s[4:7]
self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
def test_register_slice_roundtrip(self):
# Round-trip: DSL -> disasm -> DSL should preserve register count
reg = s[4:7] # 4 registers in AMD convention
inst = s_load_b128(reg, s[0:1], NULL, 0)
disasm = inst.disasm()
# Disasm shows s[4:7] - user should be able to copy this back
self.assertIn("s[4:7]", disasm)
# And s[4:7] in DSL should give the same 4 registers
reg_from_disasm = s[4:7]
self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
class TestReprReadability(unittest.TestCase):
"""
Issue: repr() leaks internal RawImm type and omits zero-valued fields.
When you create v_mov_b32_e32(v[0], v[1]), the repr shows:
VOP1(op=1, src0=RawImm(257))
Problems:
1. vdst=v[0] is omitted because 0 is treated as "default"
2. src0 shows RawImm(257) instead of v[1]
3. User sees encoded values (257 = 256 + 1) instead of register names
Expected repr: VOP1(op=1, vdst=v[0], src0=v[1])
"""
def test_repr_shows_registers_not_raw_imm(self):
inst = v_mov_b32_e32(v[0], v[1])
# Should show v[1], not RawImm(257)
self.assertNotIn("RawImm", repr(inst), "repr should not expose RawImm internal type")
self.assertIn("v[1]", repr(inst), "repr should show register name")
def test_repr_includes_zero_dst(self):
inst = v_mov_b32_e32(v[0], v[1])
# v[0] is a valid destination register, should be shown
self.assertIn("vdst", repr(inst), "repr should include vdst even when 0")
def test_repr_roundtrip(self):
# repr should produce something that can be eval'd back
inst = v_mov_b32_e32(v[0], v[1])
# This would require repr to output valid Python, e.g.:
# "VOP1(op=VOP1Op.V_MOV_B32, vdst=v[0], src0=v[1])"
r = repr(inst)
# At minimum, it should be human-readable
self.assertIn("v[", r, "repr should show register syntax")
class TestInstructionEquality(unittest.TestCase):
"""
Issue: No __eq__ method - instruction comparison requires repr() workaround.
Two identical instructions should compare equal with ==, but currently:
inst1 == inst2 returns False
The test_handwritten.py works around this with:
self.assertEqual(repr(self.inst), repr(reasm))
"""
def test_identical_instructions_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[1])
self.assertEqual(inst1, inst2, "identical instructions should be equal")
def test_different_instructions_not_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[2])
self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
class TestVOPDHelperSignature(unittest.TestCase):
"""
Issue: VOPD helper functions have confusing semantics.
v_dual_mul_f32 is defined as:
v_dual_mul_f32 = functools.partial(VOPD, VOPDOp.V_DUAL_MUL_F32)
This binds VOPDOp.V_DUAL_MUL_F32 to the FIRST positional arg of VOPD.__init__,
which is 'opx'. So v_dual_mul_f32 sets the X operation.
But then test_dual_mul in test_handwritten.py does:
v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], ...)
This passes V_DUAL_MUL_F32 as the SECOND positional arg (opy), making both
X and Y operations the same. This is confusing because:
1. The function name suggests it handles the X operation
2. But you still pass an opcode as the first arg (which becomes opy)
Expected: Either make the helper fully specify both ops, or make the
signature clearer about what the positional arg means.
"""
def test_vopd_helper_opy_should_be_required(self):
# Using only keyword args "works" but opy silently defaults to 0
inst = v_dual_mul_f32(vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32)
# Bug: opy defaults to 0 (V_DUAL_FMAC_F32) silently - should require explicit opy
# This test documents the bug - it should fail once fixed
self.assertNotEqual(inst.opy, VOPDOp.V_DUAL_FMAC_F32, "opy should not silently default to FMAC")
def test_vopd_helper_positional_arg_is_opy(self):
# The first positional arg after the partial becomes opy, not a second opx
inst = v_dual_mul_f32(VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32) # From partial
self.assertEqual(inst.opy, VOPDOp.V_DUAL_MOV_B32) # From first positional arg
class TestFieldAccessPreservesType(unittest.TestCase):
"""
Issue: Field access loses type information.
After creating an instruction, accessing fields returns encoded int values:
inst = v_mov_b32_e32(v[0], v[1])
inst.vdst # returns 0, not VGPR(0)
This makes it impossible to round-trip register types through field access.
"""
def test_vdst_returns_register(self):
inst = v_mov_b32_e32(v[5], v[1])
vdst = inst.vdst
# Should return a VGPR, not an int
self.assertIsInstance(vdst, (VGPR, int), "vdst should return VGPR or at least be usable")
# Ideally: self.assertIsInstance(vdst, VGPR)
def test_src_returns_register_for_vgpr_source(self):
inst = v_mov_b32_e32(v[0], v[1])
# src0 is encoded as 257 (256 + 1 for v1)
# Ideally it should decode back to v[1]
src0_raw = inst._values.get('src0')
# Currently returns RawImm(257), should return VGPR(1) or similar
self.assertNotIsInstance(src0_raw, RawImm, "source should not be RawImm internally")
class TestArgumentDiscoverability(unittest.TestCase):
"""
Issue: No clear signature for positional arguments.
inspect.signature(s_load_b128) shows: (*args, literal=None, **kwargs)
Users have no way to know the argument order without reading source code.
The order is implicitly defined by the class field definition order.
Possible fixes:
1. Add explicit parameter names to functools.partial
2. Generate type stubs with proper signatures
3. Add docstrings listing the expected arguments
"""
def test_signature_has_named_params(self):
import inspect
sig = inspect.signature(s_load_b128)
params = list(sig.parameters.keys())
# Currently: ['args', 'literal', 'kwargs'] (from *args, literal=None, **kwargs)
# Expected: something like ['sdata', 'sbase', 'soffset', 'offset', 'literal']
self.assertIn('sdata', params, "signature should show field names")
class TestSpecialConstants(unittest.TestCase):
"""
Issue: NULL and other constants are IntEnum values that might be confusing.
NULL = SrcEnum.NULL = 124, but users might expect NULL to be a special object
that clearly represents "no register" rather than a magic number.
"""
def test_null_has_clear_repr(self):
# NULL should have a clear string representation
self.assertIn("NULL", str(NULL) or repr(NULL), "NULL should be clearly identifiable")
def test_null_is_distinguishable_from_int(self):
# NULL should be distinguishable from the raw integer 124
self.assertNotEqual(type(NULL), int, "NULL should not be plain int")
if __name__ == "__main__":
unittest.main()
-66
View File
@@ -1,66 +0,0 @@
"""Shared test helpers for RDNA3 tests."""
import shutil
from dataclasses import dataclass
@dataclass
class KernelInfo:
code: bytes
global_size: tuple[int, int, int]
local_size: tuple[int, int, int]
buf_idxs: list[int] # indices into shared buffer pool
buf_sizes: list[int] # sizes for each buffer index
# LLVM tool detection (shared across test files)
def get_llvm_mc():
"""Find llvm-mc executable, preferring newer versions."""
for p in ['llvm-mc', 'llvm-mc-21', 'llvm-mc-20']:
if shutil.which(p): return p
raise FileNotFoundError("llvm-mc not found")
def get_llvm_objdump():
"""Find llvm-objdump executable, preferring newer versions."""
for p in ['llvm-objdump', 'llvm-objdump-21', 'llvm-objdump-20']:
if shutil.which(p): return p
raise FileNotFoundError("llvm-objdump not found")
# ═══════════════════════════════════════════════════════════════════════════════
# EXECUTION CONTEXT (for testing compiled pseudocode)
# ═══════════════════════════════════════════════════════════════════════════════
class ExecContext:
"""Context for running compiled pseudocode in tests."""
def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=0xffffffff, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
from extra.assembly.amd.pcode import Reg, MASK32, MASK64, TypedView
self._Reg, self._MASK64, self._TypedView = Reg, MASK64, TypedView
self.S0, self.S1, self.S2 = Reg(s0), Reg(s1), Reg(s2)
self.D0, self.D1 = Reg(d0), Reg(0)
self.SCC, self.VCC, self.EXEC = Reg(scc), Reg(vcc), Reg(exec_mask)
self.tmp, self.saveexec = Reg(0), Reg(exec_mask)
self.lane, self.laneId, self.literal = lane, lane, literal
self.SIMM16, self.SIMM32 = Reg(literal), Reg(literal)
self.VGPR = vgprs if vgprs is not None else {}
self.SRC0, self.VDST = Reg(src0_idx), Reg(vdst_idx)
def run(self, code: str):
"""Execute compiled code."""
import extra.assembly.amd.pcode as pcode
ns = {k: getattr(pcode, k) for k in dir(pcode) if not k.startswith('_')}
# Also include underscore-prefixed helpers that compiled pseudocode uses
for k in ['_pack', '_pack32']:
if hasattr(pcode, k): ns[k] = getattr(pcode, k)
ns.update({
'S0': self.S0, 'S1': self.S1, 'S2': self.S2, 'D0': self.D0, 'D1': self.D1,
'SCC': self.SCC, 'VCC': self.VCC, 'EXEC': self.EXEC,
'EXEC_LO': self._TypedView(self.EXEC, 31, 0), 'EXEC_HI': self._TypedView(self.EXEC, 63, 32),
'tmp': self.tmp, 'saveexec': self.saveexec,
'lane': self.lane, 'laneId': self.laneId, 'literal': self.literal,
'SIMM16': self.SIMM16, 'SIMM32': self.SIMM32, 'VGPR': self.VGPR, 'SRC0': self.SRC0, 'VDST': self.VDST,
})
exec(code, ns)
def _sync(ctx_reg, ns_val):
if isinstance(ns_val, self._Reg): ctx_reg._val = ns_val._val
else: ctx_reg._val = int(ns_val) & self._MASK64
for name in ('SCC', 'VCC', 'EXEC', 'D0', 'D1', 'tmp', 'saveexec'):
if ns.get(name) is not getattr(self, name): _sync(getattr(self, name), ns[name])
def result(self) -> dict: return {"d0": self.D0._val, "scc": self.SCC._val & 1}
-1
View File
@@ -1 +0,0 @@
"""Hardware-validated emulator tests for RDNA3 instructions."""
-202
View File
@@ -1,202 +0,0 @@
"""Test infrastructure for hardware-validated RDNA3 emulator tests.
Uses run_asm() with memory output, so tests can run on both emulator and real hardware.
Set USE_HW=1 to run on both emulator and real hardware, comparing results.
"""
import ctypes, os, struct
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import RawImm
from extra.assembly.amd.emu import WaveState, run_asm, set_valid_mem_ranges
from extra.assembly.amd.pcode import _i32, _f32
VCC = SrcEnum.VCC_LO # For VOP3SD sdst field
USE_HW = os.environ.get("USE_HW", "0") == "1"
FLOAT_TOLERANCE = 1e-5
# Output buffer layout: vgpr[16][32], sgpr[16], vcc, scc
N_VGPRS, N_SGPRS, WAVE_SIZE = 16, 16, 32
VGPR_BYTES = N_VGPRS * WAVE_SIZE * 4 # 16 regs * 32 lanes * 4 bytes = 2048
SGPR_BYTES = N_SGPRS * 4 # 16 regs * 4 bytes = 64
OUT_BYTES = VGPR_BYTES + SGPR_BYTES + 8 # + vcc + scc
# Float conversion helpers
def f2i(f: float) -> int: return _i32(f)
def i2f(i: int) -> float: return _f32(i)
def f2i64(f: float) -> int: return struct.unpack('<Q', struct.pack('<d', f))[0]
def i642f(i: int) -> float: return struct.unpack('<d', struct.pack('<Q', i))[0]
def assemble(instructions: list) -> bytes:
return b''.join(inst.to_bytes() for inst in instructions)
def get_prologue_epilogue(n_lanes: int) -> tuple[list, list]:
"""Generate prologue and epilogue instructions for state capture."""
prologue = [
s_mov_b32(s[80], s[0]),
s_mov_b32(s[81], s[1]),
v_mov_b32_e32(v[255], v[0]),
]
for i in range(N_VGPRS):
prologue.append(v_mov_b32_e32(v[i], 0))
for i in range(N_SGPRS):
prologue.append(s_mov_b32(s[i], 0))
prologue.append(s_mov_b32(s[SrcEnum.VCC_LO - 128], 0))
epilogue = [
s_mov_b32(s[90], SrcEnum.VCC_LO),
s_cselect_b32(s[91], 1, 0),
s_load_b64(s[92:93], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_lshlrev_b32_e32(v[240], 2, v[255]),
]
for i in range(N_VGPRS):
epilogue.append(global_store_b32(addr=v[240], data=v[i], saddr=s[92], offset=i * WAVE_SIZE * 4))
epilogue.append(v_mov_b32_e32(v[241], 0))
epilogue.append(v_cmp_eq_u32_e32(v[255], v[241]))
epilogue.append(s_and_saveexec_b32(s[94], SrcEnum.VCC_LO))
epilogue.append(v_mov_b32_e32(v[240], 0))
for i in range(N_SGPRS):
epilogue.append(v_mov_b32_e32(v[243], s[i]))
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92], offset=VGPR_BYTES + i * 4))
epilogue.append(v_mov_b32_e32(v[243], s[90]))
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92], offset=VGPR_BYTES + SGPR_BYTES))
epilogue.append(v_mov_b32_e32(v[243], s[91]))
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92], offset=VGPR_BYTES + SGPR_BYTES + 4))
epilogue.append(s_mov_b32(s[SrcEnum.EXEC_LO - 128], s[94]))
epilogue.append(s_endpgm())
return prologue, epilogue
def parse_output(out_buf: bytes, n_lanes: int) -> WaveState:
"""Parse output buffer into WaveState."""
st = WaveState()
for i in range(N_VGPRS):
for lane in range(n_lanes):
off = i * WAVE_SIZE * 4 + lane * 4
st.vgpr[lane][i] = struct.unpack_from('<I', out_buf, off)[0]
for i in range(N_SGPRS):
st.sgpr[i] = struct.unpack_from('<I', out_buf, VGPR_BYTES + i * 4)[0]
st.vcc = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES)[0]
st.scc = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES + 4)[0]
return st
def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
"""Run instructions via emulator run_asm, dump state to memory, return WaveState."""
out_buf = (ctypes.c_uint8 * OUT_BYTES)(*([0] * OUT_BYTES))
out_addr = ctypes.addressof(out_buf)
prologue, epilogue = get_prologue_epilogue(n_lanes)
code = assemble(prologue + instructions + epilogue)
args = (ctypes.c_uint64 * 1)(out_addr)
args_ptr = ctypes.addressof(args)
kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
lib_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({(out_addr, OUT_BYTES), (args_ptr, 8)})
# rsrc2: USER_SGPR_COUNT=2, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z=1, LDS_SIZE=128 (64KB)
rsrc2 = 0x19c | (128 << 15)
result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2)
assert result == 0, f"run_asm failed with {result}"
return parse_output(bytes(out_buf), n_lanes)
def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
"""Run instructions on real AMD hardware via HIPCompiler and AMDProgram."""
from tinygrad.device import Device
from tinygrad.runtime.ops_amd import AMDProgram
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.helpers import flat_mv
dev = Device["AMD"]
compiler = HIPCompiler(dev.arch)
prologue, epilogue = get_prologue_epilogue(n_lanes)
code = assemble(prologue + instructions + epilogue)
byte_str = ', '.join(f'0x{b:02x}' for b in code)
asm_src = f""".text
.globl test
.p2align 8
.type test,@function
test:
.byte {byte_str}
.rodata
.p2align 6
.amdhsa_kernel test
.amdhsa_next_free_vgpr 256
.amdhsa_next_free_sgpr 96
.amdhsa_wavefront_size32 1
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_kernarg_size 8
.amdhsa_group_segment_fixed_size 65536
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.version:
- 1
- 0
amdhsa.kernels:
- .name: test
.symbol: test.kd
.kernarg_segment_size: 8
.group_segment_fixed_size: 65536
.private_segment_fixed_size: 0
.kernarg_segment_align: 8
.wavefront_size: 32
.sgpr_count: 96
.vgpr_count: 256
.max_flat_workgroup_size: 1024
...
.end_amdgpu_metadata
"""
lib = compiler.compile(asm_src)
prg = AMDProgram(dev, "test", lib)
out_gpu = dev.allocator.alloc(OUT_BYTES)
prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
out_buf = bytearray(OUT_BYTES)
dev.allocator._copyout(flat_mv(memoryview(out_buf)), out_gpu)
return parse_output(bytes(out_buf), n_lanes)
def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS) -> list[str]:
"""Compare two WaveStates and return list of differences."""
import math
diffs = []
for i in range(n_vgprs):
for lane in range(n_lanes):
emu_val = emu_st.vgpr[lane][i]
hw_val = hw_st.vgpr[lane][i]
if emu_val != hw_val:
emu_f, hw_f = _f32(emu_val), _f32(hw_val)
if math.isnan(emu_f) and math.isnan(hw_f):
continue
diffs.append(f"v[{i}] lane {lane}: emu=0x{emu_val:08x} ({emu_f:.6g}) hw=0x{hw_val:08x} ({hw_f:.6g})")
for i in range(N_SGPRS):
emu_val = emu_st.sgpr[i]
hw_val = hw_st.sgpr[i]
if emu_val != hw_val:
diffs.append(f"s[{i}]: emu=0x{emu_val:08x} hw=0x{hw_val:08x}")
if emu_st.vcc != hw_st.vcc:
diffs.append(f"vcc: emu=0x{emu_st.vcc:08x} hw=0x{hw_st.vcc:08x}")
if emu_st.scc != hw_st.scc:
diffs.append(f"scc: emu={emu_st.scc} hw={hw_st.scc}")
return diffs
def run_program(instructions: list, n_lanes: int = 1) -> WaveState:
"""Run instructions and return WaveState.
If USE_HW=1, runs on both emulator and hardware, compares results, and raises if they differ.
Otherwise, runs only on emulator.
"""
emu_st = run_program_emu(instructions, n_lanes)
if USE_HW:
hw_st = run_program_hw(instructions, n_lanes)
diffs = compare_wave_states(emu_st, hw_st, n_lanes)
if diffs:
raise AssertionError(f"Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
return hw_st
return emu_st
-629
View File
@@ -1,629 +0,0 @@
"""Tests for DS instructions - data share (LDS) operations.
Includes: ds_store_b32, ds_load_b32, ds_store_2addr_*, ds_load_2addr_*,
ds_add_*, ds_max_*, ds_min_*, ds_and_*, ds_or_*, ds_xor_*,
ds_inc_*, ds_dec_*, ds_cmpstore_*, ds_storexchg_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestDS2Addr(unittest.TestCase):
"""Tests for DS_*_2ADDR instructions."""
def test_ds_store_load_2addr_b32(self):
"""DS_STORE_2ADDR_B32 and DS_LOAD_2ADDR_B32 with offset * 4."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
def test_ds_store_load_2addr_b64(self):
"""DS_STORE_2ADDR_B64 and DS_LOAD_2ADDR_B64."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STORE_2ADDR_B64, addr=v[10], data0=v[0], data1=v[2], vdst=v[0], offset0=0, offset1=2),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4], offset0=0, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][6], 0x12345678)
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
class TestDS2AddrMore(unittest.TestCase):
"""Additional DS_*_2ADDR tests."""
def test_ds_store_load_2addr_b32_nonzero_offsets(self):
"""DS_STORE_2ADDR_B32 with non-zero offsets (offset*4 scaling)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x11111111),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x22222222),
v_mov_b32_e32(v[1], s[2]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=2, offset1=5),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2], offset0=2, offset1=5),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have value from offset 8 (2*4)")
self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have value from offset 20 (5*4)")
def test_ds_2addr_b64_no_overlap(self):
"""DS_LOAD_2ADDR_B64 with adjacent offsets should not overlap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x11111111),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_mov_b32(s[2], 0x22222222),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=4),
s_mov_b32(s[2], 0x33333333),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=8),
s_mov_b32(s[2], 0x44444444),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=12),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should be 0x11111111")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should be 0x22222222")
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should be 0x33333333")
self.assertEqual(st.vgpr[0][7], 0x44444444, "v7 should be 0x44444444")
def test_ds_load_2addr_b32_no_overwrite(self):
"""DS_LOAD_2ADDR_B32 should only write 2 VGPRs."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[2]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[4], s[2]), # Sentinel
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should be untouched")
def test_ds_load_b64_no_overwrite(self):
"""DS_LOAD_B64 should only write 2 VGPRs."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[2]),
ds_store_b64(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0x12345678),
v_mov_b32_e32(v[4], s[2]), # Sentinel
ds_load_b64(addr=v[10], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][3], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][4], 0x12345678, "v4 should be untouched")
class TestDSAtomic(unittest.TestCase):
"""Tests for DS atomic operations."""
def test_ds_max_rtn_u32(self):
"""DS_MAX_RTN_U32: atomically store max and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]),
ds_max_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
def test_ds_min_rtn_u32(self):
"""DS_MIN_RTN_U32: atomically store min and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[1], s[2]),
ds_min_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 200)
self.assertEqual(st.vgpr[0][3], 100)
def test_ds_and_rtn_b32(self):
"""DS_AND_RTN_B32: atomically AND and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xFF00FF00),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0xFFFF0000),
v_mov_b32_e32(v[1], s[2]),
ds_and_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xFF00FF00)
self.assertEqual(st.vgpr[0][3], 0xFF000000)
def test_ds_or_rtn_b32(self):
"""DS_OR_RTN_B32: atomically OR and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x00FF0000),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0x000000FF),
v_mov_b32_e32(v[1], s[2]),
ds_or_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x00FF0000)
self.assertEqual(st.vgpr[0][3], 0x00FF00FF)
def test_ds_xor_rtn_b32(self):
"""DS_XOR_RTN_B32: atomically XOR and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0xFFFFFFFF),
v_mov_b32_e32(v[1], s[2]),
ds_xor_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0x55555555)
def test_ds_inc_rtn_u32(self):
"""DS_INC_RTN_U32: increment with wrap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 5),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_inc_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 5)
self.assertEqual(st.vgpr[0][3], 6)
def test_ds_dec_rtn_u32(self):
"""DS_DEC_RTN_U32: decrement with wrap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 5),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 5)
self.assertEqual(st.vgpr[0][3], 4)
def test_ds_cmpstore_b32_match(self):
"""DS_CMPSTORE_B32: conditional store when compare matches."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]), # new value
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[2], s[2]), # compare = 100 (matches)
ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 200)
def test_ds_cmpstore_b32_no_match(self):
"""DS_CMPSTORE_B32: no store when compare doesn't match."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]), # new value
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[2], s[2]), # compare = 50 (doesn't match)
ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 100)
def test_ds_max_u32_no_rtn(self):
"""DS_MAX_U32 (no RTN): atomically store max, no return value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]),
ds_max_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
def test_ds_add_u32_no_rtn_preserves_vdst(self):
"""DS_ADD_U32 (no RTN) should NOT write to vdst."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[2]), # sentinel
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[1], s[2]),
ds_add_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF, "v2 should preserve sentinel")
self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
def test_ds_add_rtn_u32_writes_vdst(self):
"""DS_ADD_RTN_U32 should write old value to vdst."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[2]), # sentinel
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[1], s[2]),
ds_add_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
def test_ds_dec_rtn_u32_wrap(self):
"""DS_DEC_RTN_U32: decrement wraps when value is 0 or > limit."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0), # Start at 0
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0, "v2 should have old value (0)")
# When mem == 0 or mem > limit, result = limit
self.assertEqual(st.vgpr[0][3], 10, "v3 should wrap to limit (10)")
class TestDSStorexchg(unittest.TestCase):
"""Tests for DS_STOREXCHG instructions."""
def test_ds_storexchg_rtn_b32(self):
"""DS_STOREXCHG_RTN_B32: exchange value and return old."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STOREXCHG_RTN_B32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
class TestDSRegisterWidth(unittest.TestCase):
"""Regression tests: DS loads should only write correct number of VGPRs."""
def test_ds_load_b32_no_overwrite(self):
"""DS_LOAD_B32 should only write 1 VGPR."""
instructions = [
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]), # sentinel
ds_store_b32(addr=v[0], data0=v[1], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[0], vdst=v[1], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should be untouched")
class TestDS2AddrStride64(unittest.TestCase):
"""Tests for DS_*_2ADDR_STRIDE64 (offset * 256 for B32, offset * 512 for B64)."""
def test_ds_store_load_2addr_stride64_b32(self):
"""DS_STORE_2ADDR_STRIDE64_B32: stores at ADDR + offset*256."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[2], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 from addr 256")
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 from addr 512")
def test_ds_store_load_2addr_stride64_b64(self):
"""DS_STORE_2ADDR_STRIDE64_B64: stores at ADDR + offset*512."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0], data1=v[2], vdst=v[0], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[4], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][6], 0x12345678)
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
def test_ds_storexchg_2addr_rtn_b32(self):
"""DS_STOREXCHG_2ADDR_RTN_B32: exchange at two addresses."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[6], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "old val 0")
self.assertEqual(st.vgpr[0][5], 0x22222222, "old val 1")
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "new val 0")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "new val 1")
def test_ds_storexchg_rtn_b64(self):
"""DS_STOREXCHG_RTN_B64: exchange 64-bit value and return old."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]), # initial low
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]), # initial high
DS(DSOp.DS_STORE_B64, addr=v[10], data0=v[0], vdst=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]), # new low
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]), # new high
DS(DSOp.DS_STOREXCHG_RTN_B64, addr=v[10], data0=v[2], vdst=v[4], offset0=0),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_B64, addr=v[10], vdst=v[6], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should have old low dword")
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE, "v5 should have old high dword")
self.assertEqual(st.vgpr[0][6], 0x12345678, "v6 should have new low dword")
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0, "v7 should have new high dword")
def test_ds_store_load_2addr_stride64_b64_roundtrip(self):
"""DS_STORE_2ADDR_STRIDE64_B64 followed by DS_LOAD_2ADDR_STRIDE64_B64 works correctly."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0], data1=v[0], vdst=v[0], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[2], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have val1 low")
self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have val1 high")
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have val2 low")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have val2 high")
def test_ds_storexchg_2addr_stride64_rtn_b32(self):
"""DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32: exchange at two addresses (offset*256)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[6], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have old value")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have old value")
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "v6 should have new value")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "v7 should have new value")
def test_ds_storexchg_2addr_stride64_rtn_b64_returns_old(self):
"""DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64: returns old values correctly."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0], data1=v[0], vdst=v[0], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[6], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[7], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64, addr=v[10], data0=v[6], data1=v[6], vdst=v[8], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][8], 0x11111111, "v8 should have old val1 low")
self.assertEqual(st.vgpr[0][9], 0x22222222, "v9 should have old val1 high")
self.assertEqual(st.vgpr[0][10], 0x11111111, "v10 should have old val2 low")
self.assertEqual(st.vgpr[0][11], 0x22222222, "v11 should have old val2 high")
class TestAtomicOrdering(unittest.TestCase):
"""Tests for atomic operation return values and ordering."""
def test_ds_add_rtn_sequence(self):
"""DS_ADD_RTN returns correct old values in sequence."""
instructions = [
v_mov_b32_e32(v[10], 0),
v_mov_b32_e32(v[0], 100),
DS(DSOp.DS_STORE_B32, addr=v[10], data0=v[0], vdst=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[1], 25),
DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_B32, addr=v[10], vdst=v[4], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "First add should return 100")
self.assertEqual(st.vgpr[0][3], 125, "Second add should return 125")
self.assertEqual(st.vgpr[0][4], 150, "Final value should be 150")
if __name__ == '__main__':
unittest.main()
-363
View File
@@ -1,363 +0,0 @@
"""Tests for FLAT instructions - flat memory operations.
Includes: flat_load_*, flat_store_*, flat_atomic_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestFlatAtomic(unittest.TestCase):
"""Tests for FLAT atomic instructions."""
def _make_test(self, setup_instrs, atomic_instr, check_fn, test_offset=2000):
"""Helper to create atomic test instructions."""
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
] + setup_instrs + [atomic_instr, s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
check_fn(st)
def test_flat_atomic_add_u32(self):
"""FLAT_ATOMIC_ADD_U32 adds to memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_ADD_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 100)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_swap_b32(self):
"""FLAT_ATOMIC_SWAP_B32 swaps memory value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_SWAP_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_and_b32(self):
"""FLAT_ATOMIC_AND_B32 ANDs with memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xFF00FF00),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0xFFFF0000),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_AND_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0xFF00FF00)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_or_b32(self):
"""FLAT_ATOMIC_OR_B32 ORs with memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0x00FF0000),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0x0000FF00),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_OR_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0x00FF0000)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_inc_u32(self):
"""FLAT_ATOMIC_INC_U32 increments and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 10),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 100), # threshold
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_INC_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 10)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_dec_u32(self):
"""FLAT_ATOMIC_DEC_U32 decrements and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 10),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_DEC_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 10)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_sub_u32(self):
"""FLAT_ATOMIC_SUB_U32 subtracts from memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 30),
v_mov_b32_e32(v[3], s[0]), # sub 30
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_SUB_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 100, "v4 should have old value (100)")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_xor_b32(self):
"""FLAT_ATOMIC_XOR_B32 XORs with memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0xFFFFFFFF),
v_mov_b32_e32(v[3], s[0]), # XOR mask
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_XOR_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA, "v4 should have old value")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_min_u32(self):
"""FLAT_ATOMIC_MIN_U32 stores min and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[3], s[0]), # compare value (smaller)
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_MIN_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 100, "v4 should have old value (100)")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_max_u32(self):
"""FLAT_ATOMIC_MAX_U32 stores max and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[3], s[0]), # compare value (larger)
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_MAX_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 50, "v4 should have old value (50)")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_inc_u64_returns_old_value(self):
"""FLAT_ATOMIC_INC_U64 should return full 64-bit old value."""
TEST_OFFSET = 2000
setup = [
# Store initial 64-bit value: 0xCAFEBABE_DEADBEEF
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
# Threshold: 0xFFFFFFFF_FFFFFFFF
s_mov_b32(s[0], 0xFFFFFFFF),
v_mov_b32_e32(v[4], s[0]),
v_mov_b32_e32(v[5], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_INC_U64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][6], 0xDEADBEEF, "v6 should have old value low dword")
self.assertEqual(st.vgpr[0][7], 0xCAFEBABE, "v7 should have old value high dword")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_add_u64(self):
"""FLAT_ATOMIC_ADD_U64 adds 64-bit value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0x00000001), # add 1
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x00000000),
v_mov_b32_e32(v[5], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_ADD_U64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][6], 0x11111111, "v6 should have old value low")
self.assertEqual(st.vgpr[0][7], 0x22222222, "v7 should have old value high")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_swap_b64(self):
"""FLAT_ATOMIC_SWAP_B64 swaps 64-bit value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0xDDDDDDDD),
v_mov_b32_e32(v[5], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_SWAP_B64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "v6 should have old value low")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "v7 should have old value high")
self._make_test(setup, atomic, check, TEST_OFFSET)
class TestFlatLoad(unittest.TestCase):
"""Tests for FLAT load instructions."""
def test_flat_load_b32(self):
"""FLAT_LOAD_B32 loads 32-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(FLATOp.FLAT_LOAD_B32, addr=v[0], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
def test_flat_load_b64(self):
"""FLAT_LOAD_B64 loads 64-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(FLATOp.FLAT_LOAD_B64, addr=v[0], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
def test_flat_load_b96(self):
"""FLAT_LOAD_B96 loads 96-bit (3 dword) value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[4], s[0]),
global_store_b96(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(FLATOp.FLAT_LOAD_B96, addr=v[0], vdst=v[5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][5], 0x11111111)
self.assertEqual(st.vgpr[0][6], 0x22222222)
self.assertEqual(st.vgpr[0][7], 0x33333333)
def test_flat_load_b128(self):
"""FLAT_LOAD_B128 loads 128-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x44444444),
v_mov_b32_e32(v[5], s[0]),
global_store_b128(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(FLATOp.FLAT_LOAD_B128, addr=v[0], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][6], 0x11111111)
self.assertEqual(st.vgpr[0][7], 0x22222222)
self.assertEqual(st.vgpr[0][8], 0x33333333)
self.assertEqual(st.vgpr[0][9], 0x44444444)
if __name__ == '__main__':
unittest.main()
-364
View File
@@ -1,364 +0,0 @@
"""Tests for GLOBAL instructions - global memory operations.
Includes: global_load_*, global_store_*, global_atomic_*, global_load_d16_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestGlobalAtomic(unittest.TestCase):
"""Tests for GLOBAL atomic instructions."""
def _make_test(self, setup_instrs, atomic_instr, check_fn, test_offset=2000):
"""Helper to create atomic test instructions."""
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
] + setup_instrs + [atomic_instr, s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
check_fn(st)
def test_global_atomic_add_u32(self):
"""GLOBAL_ATOMIC_ADD_U32 adds to memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(GLOBALOp.GLOBAL_ATOMIC_ADD_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1, seg=2)
def check(st):
self.assertEqual(st.vgpr[0][4], 100)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_global_atomic_add_u64(self):
"""GLOBAL_ATOMIC_ADD_U64 adds 64-bit value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xFFFFFFFF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x00000000),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0x00000001),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x00000000),
v_mov_b32_e32(v[5], s[0]),
]
atomic = FLAT(GLOBALOp.GLOBAL_ATOMIC_ADD_U64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1, seg=2)
def check(st):
self.assertEqual(st.vgpr[0][6], 0xFFFFFFFF)
self.assertEqual(st.vgpr[0][7], 0x00000000)
self._make_test(setup, atomic, check, TEST_OFFSET)
class TestGlobalLoad(unittest.TestCase):
"""Tests for GLOBAL load instructions."""
def test_global_load_b96(self):
"""GLOBAL_LOAD_B96 loads 96-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[4], s[0]),
global_store_b96(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(GLOBALOp.GLOBAL_LOAD_B96, addr=v[0], vdst=v[5], saddr=SrcEnum.NULL, offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][5], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][6], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][7], 0xCCCCCCCC)
def test_global_load_b128(self):
"""GLOBAL_LOAD_B128 loads 128-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[5], s[0]),
global_store_b128(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(GLOBALOp.GLOBAL_LOAD_B128, addr=v[0], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][6], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][7], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][8], 0x12345678)
self.assertEqual(st.vgpr[0][9], 0x9ABCDEF0)
class TestGlobalStore(unittest.TestCase):
"""Tests for GLOBAL store instructions."""
def test_global_store_b64_basic(self):
"""GLOBAL_STORE_B64 stores 8 bytes from v[n:n+1] to memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xDEADBEEF),
s_mov_b32(s[5], 0xCAFEBABE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], s[5]),
v_mov_b32_e32(v[0], 0),
global_store_b64(addr=v[0], data=v[2], saddr=s[2], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(GLOBALOp.GLOBAL_LOAD_B64, addr=v[0], vdst=v[4], data=v[4], saddr=s[2], offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[4]),
v_mov_b32_e32(v[1], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
class TestD16HiLoads(unittest.TestCase):
"""Tests for D16_HI load instructions that load into high 16 bits."""
def test_global_load_d16_hi_b16_preserves_low_bits(self):
"""GLOBAL_LOAD_D16_HI_B16 must preserve low 16 bits of destination."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
global_store_b16(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[4], 0x0000BEEF),
v_mov_b32_e32(v[3], s[4]),
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[0], vdst=v[3], data=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0xCAFEBEEF, f"Expected 0xCAFEBEEF, got 0x{result:08x}")
def test_global_load_d16_hi_b16_data_differs_from_vdst(self):
"""GLOBAL_LOAD_D16_HI_B16 where data field differs from vdst."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b16(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[4], 0x0000DEAD),
v_mov_b32_e32(v[0], s[4]), # data field - should NOT affect result
v_mov_b32_e32(v[1], 0), # vdst - low bits should be preserved
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[0], saddr=s[2], offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[1]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0xCAFE0000, f"Expected 0xCAFE0000, got 0x{result:08x}")
def test_global_load_d16_hi_u8_data_differs_from_vdst(self):
"""GLOBAL_LOAD_D16_HI_U8 where data field differs from vdst."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xAB),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b8(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[4], 0x0000DEAD),
v_mov_b32_e32(v[4], s[4]), # data field
s_mov_b32(s[4], 0x0000BEEF),
v_mov_b32_e32(v[5], s[4]), # vdst
v_mov_b32_e32(v[3], 0),
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_U8, addr=v[3], vdst=v[5], data=v[4], saddr=s[2], offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0x00ABBEEF, f"Expected 0x00ABBEEF, got 0x{result:08x}")
def test_global_load_d16_hi_b16_same_addr_and_dst_zero_addr(self):
"""GLOBAL_LOAD_D16_HI_B16 with same register for addr and vdst, addr value=0."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b16(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[1], 0),
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[1], saddr=s[2], offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[1]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0xCAFE0000, f"Expected 0xCAFE0000, got 0x{result:08x}")
def test_global_load_d16_hi_b16_tril_exact_pattern(self):
"""Exact pattern from tril() failure: data=v0 differs from vdst=v1."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0x01010101),
v_mov_b32_e32(v[10], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b32(addr=v[3], data=v[10], saddr=s[2], offset=TEST_OFFSET),
global_store_b32(addr=v[3], data=v[10], saddr=s[2], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
# Set v[0] to 0x0101 (simulating prior u16 load result)
s_mov_b32(s[4], 0x0101),
v_mov_b32_e32(v[0], s[4]),
# Set v[1] to 0
v_mov_b32_e32(v[1], 0),
# Load using v[1] as addr AND vdst, but v[0] as data
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[0], saddr=s[2], offset=TEST_OFFSET+6, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[1]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
# Expected: hi=0x0101 (loaded), lo=0x0000 (from v1) -> 0x01010000
self.assertEqual(result, 0x01010000, f"Expected 0x01010000, got 0x{result:08x}")
def test_global_load_d16_hi_i8_data_differs_from_vdst(self):
"""GLOBAL_LOAD_D16_HI_I8 where data field differs from vdst."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0x80), # negative signed byte = -128
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b8(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[4], 0x0000DEAD),
v_mov_b32_e32(v[4], s[4]), # data field
s_mov_b32(s[4], 0x0000BEEF),
v_mov_b32_e32(v[5], s[4]), # vdst
v_mov_b32_e32(v[3], 0),
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_I8, addr=v[3], vdst=v[5], data=v[4], saddr=s[2], offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
# 0x80 sign-extended = 0xFF80, lo=0xBEEF -> 0xFF80BEEF
self.assertEqual(result, 0xFF80BEEF, f"Expected 0xFF80BEEF, got 0x{result:08x}")
def test_global_store_b64_tril_pattern(self):
"""Test the exact pattern from tril() kernel that was failing."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0x01010101),
v_mov_b32_e32(v[10], s[4]),
v_mov_b32_e32(v[11], s[4]),
s_mov_b32(s[4], 0x01),
v_mov_b32_e32(v[12], s[4]),
v_mov_b32_e32(v[0], 0),
global_store_b64(addr=v[0], data=v[10], saddr=s[2], offset=TEST_OFFSET),
global_store_b8(addr=v[0], data=v[12], saddr=s[2], offset=TEST_OFFSET+8),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[2], 0),
v_mov_b32_e32(v[1], 0),
FLAT(GLOBALOp.GLOBAL_LOAD_U16, addr=v[2], vdst=v[0], data=v[0], saddr=s[2], offset=TEST_OFFSET+3, seg=2),
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[1], saddr=s[2], offset=TEST_OFFSET+6, seg=2),
FLAT(GLOBALOp.GLOBAL_LOAD_U8, addr=v[2], vdst=v[3], data=v[3], saddr=s[2], offset=TEST_OFFSET, seg=2),
FLAT(GLOBALOp.GLOBAL_LOAD_U8, addr=v[2], vdst=v[4], data=v[4], saddr=s[2], offset=TEST_OFFSET+8, seg=2),
s_waitcnt(vmcnt=0),
v_and_b32_e32(v[5], 0xffff, v[0]),
v_lshlrev_b32_e32(v[0], 24, v[0]),
v_lshrrev_b32_e32(v[5], 8, v[5]),
v_or_b32_e32(v[0], v[3], v[0]),
v_or_b32_e32(v[1], v[5], v[1]),
global_store_b64(addr=v[2], data=v[0], saddr=s[2], offset=TEST_OFFSET+16),
s_waitcnt(vmcnt=0),
FLAT(GLOBALOp.GLOBAL_LOAD_B64, addr=v[2], vdst=v[6], data=v[6], saddr=s[2], offset=TEST_OFFSET+16, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
v0 = st.vgpr[0][0]
v1 = st.vgpr[0][1]
self.assertEqual(v0, 0x01000001, f"v0: expected 0x01000001, got 0x{v0:08x}")
self.assertEqual(v1, 0x01010001, f"v1: expected 0x01010001, got 0x{v1:08x}")
byte5 = (v1 >> 8) & 0xff
self.assertEqual(byte5, 0x00, f"byte5: expected 0x00, got 0x{byte5:02x}")
if __name__ == '__main__':
unittest.main()
-342
View File
@@ -1,342 +0,0 @@
"""Tests for SOP instructions - scalar operations.
Includes: s_add_u32, s_mov_b32, s_and_b32, s_or_b32, s_quadmask_b32, s_wqm_b32,
s_cbranch_vccnz, s_cbranch_vccz
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestBasicScalar(unittest.TestCase):
"""Tests for basic scalar operations."""
def test_s_add_u32(self):
"""S_ADD_U32 adds two scalar values."""
instructions = [
s_mov_b32(s[0], 100),
s_mov_b32(s[1], 200),
s_add_u32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 300)
def test_s_add_u32_carry(self):
"""S_ADD_U32 sets SCC on overflow."""
instructions = [
s_mov_b32(s[0], 64),
s_not_b32(s[0], s[0]), # ~64 = 0xffffffbf
s_mov_b32(s[1], 64),
s_add_u32(s[2], s[0], s[1]), # 0xffffffbf + 64 = 0xffffffff
s_mov_b32(s[3], 1),
s_add_u32(s[4], s[2], s[3]), # 0xffffffff + 1 = overflow
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0)
self.assertEqual(st.scc, 1)
def test_s_brev_b32(self):
"""S_BREV_B32 reverses bits of a 32-bit value."""
# 10 = 0b00000000000000000000000000001010
# reversed = 0b01010000000000000000000000000000 = 0x50000000
instructions = [
s_mov_b32(s[0], 10),
s_brev_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0x50000000)
def test_s_brev_b32_all_ones(self):
"""S_BREV_B32 with all ones stays all ones."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_brev_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
def test_s_brev_b32_single_bit(self):
"""S_BREV_B32 with bit 0 set becomes bit 31."""
instructions = [
s_mov_b32(s[0], 1),
s_brev_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0x80000000)
class TestQuadmaskWqm(unittest.TestCase):
"""Tests for S_QUADMASK_B32 and S_WQM_B32."""
def test_s_quadmask_b32_all_quads_active(self):
"""S_QUADMASK_B32 with all quads active."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF), # All lanes active
s_quadmask_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
# Each quad (4 lanes) with any bit set -> 1 bit in result
# 32 lanes = 8 quads, all active -> 0xFF
self.assertEqual(st.sgpr[1], 0xFF)
def test_s_quadmask_b32_alternating_quads(self):
"""S_QUADMASK_B32 with alternating quads active."""
instructions = [
s_mov_b32(s[0], 0x0F0F0F0F), # Quads 0,2,4,6 active
s_quadmask_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
# Quads 0,2,4,6 have at least one bit -> 0b01010101 = 0x55
self.assertEqual(st.sgpr[1], 0x55)
def test_s_quadmask_b32_no_quads_active(self):
"""S_QUADMASK_B32 with no quads active."""
instructions = [
s_mov_b32(s[0], 0),
s_quadmask_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0)
def test_s_quadmask_b32_single_lane_per_quad(self):
"""S_QUADMASK_B32 with single lane active in each quad."""
instructions = [
s_mov_b32(s[0], 0x11111111), # Bit 0 of each nibble
s_quadmask_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
# All 8 quads have at least one lane -> 0xFF
self.assertEqual(st.sgpr[1], 0xFF)
def test_s_wqm_b32_all_active(self):
"""S_WQM_B32 with all lanes active returns all 1s."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_wqm_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
def test_s_wqm_b32_alternating_quads(self):
"""S_WQM_B32 with single lane per quad expands to full quads."""
instructions = [
s_mov_b32(s[0], 0x11111111), # One lane per quad
s_wqm_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
# Each quad with any bit expands to all 4 bits
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
def test_s_wqm_b32_zero(self):
"""S_WQM_B32 with zero input returns zero."""
instructions = [
s_mov_b32(s[0], 0),
s_wqm_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0)
class TestBranch(unittest.TestCase):
"""Tests for branch instructions."""
def test_cbranch_vccnz_ignores_vcc_hi(self):
"""S_CBRANCH_VCCNZ should only check VCC_LO in wave32."""
instructions = [
# Set VCC_LO = 0, VCC_HI = 1
s_mov_b32(s[SrcEnum.VCC_LO - 128], 0),
s_mov_b32(s[SrcEnum.VCC_HI - 128], 1),
v_mov_b32_e32(v[0], 0),
# If VCC_HI is incorrectly used, branch will be taken
s_cbranch_vccnz(1), # Skip next instruction if VCC != 0
v_mov_b32_e32(v[0], 42), # This should execute
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 42, "Branch should NOT be taken (VCC_LO is 0)")
def test_cbranch_vccz_ignores_vcc_hi(self):
"""S_CBRANCH_VCCZ should only check VCC_LO in wave32."""
instructions = [
# Set VCC_LO = 1, VCC_HI = 0
s_mov_b32(s[SrcEnum.VCC_LO - 128], 1),
s_mov_b32(s[SrcEnum.VCC_HI - 128], 0),
v_mov_b32_e32(v[0], 0),
# If VCC_HI is incorrectly used, branch will be taken
s_cbranch_vccz(1), # Skip next instruction if VCC == 0
v_mov_b32_e32(v[0], 42), # This should execute
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 42, "Branch should NOT be taken (VCC_LO is 1)")
def test_cbranch_vccnz_branches_on_vcc_lo(self):
"""S_CBRANCH_VCCNZ branches when VCC_LO is non-zero."""
instructions = [
s_mov_b32(s[SrcEnum.VCC_LO - 128], 1),
v_mov_b32_e32(v[0], 0),
s_cbranch_vccnz(1), # Skip next instruction if VCC != 0
v_mov_b32_e32(v[0], 42), # This should be skipped
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0, "Branch should be taken (VCC_LO is 1)")
class Test64BitLiterals(unittest.TestCase):
"""Tests for 64-bit literal encoding in instructions."""
def test_64bit_literal_negative_encoding(self):
"""64-bit literal -2^32 encodes correctly."""
lit = -4294967296.0 # -2^32
lit_bits = f2i64(lit)
instructions = [
s_mov_b32(s[0], lit_bits & 0xffffffff),
s_mov_b32(s[1], lit_bits >> 32),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
result = i642f(st.vgpr[0][0] | (st.vgpr[0][1] << 32))
self.assertAlmostEqual(result, -4294967296.0, places=5)
def test_64bit_literal_positive_encoding(self):
"""64-bit instruction encodes large positive literals correctly."""
large_val = 0x12345678
inst = v_add_f64(v[2], v[0], large_val)
self.assertIsNotNone(inst._literal, "Literal should be set")
actual_lit = (inst._literal >> 32) & 0xffffffff
self.assertEqual(actual_lit, large_val, f"Literal should be {large_val:#x}, got {actual_lit:#x}")
class TestSCCBehavior(unittest.TestCase):
"""Tests for SCC condition code behavior."""
def test_scc_from_s_cmp(self):
"""SCC should be set by scalar compare."""
instructions = [
s_mov_b32(s[0], 10),
s_cmp_eq_u32(s[0], 10),
s_cselect_b32(s[1], 1, 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 1, "SCC should be true")
self.assertEqual(st.scc, 1)
def test_scc_clear(self):
"""SCC should be cleared by failing compare."""
instructions = [
s_mov_b32(s[0], 10),
s_cmp_eq_u32(s[0], 20),
s_cselect_b32(s[1], 1, 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0, "SCC should be false")
self.assertEqual(st.scc, 0)
class TestSignedArithmetic(unittest.TestCase):
"""Tests for S_ADD_I32, S_SUB_I32 and their SCC overflow behavior."""
def test_s_add_i32_no_overflow(self):
"""S_ADD_I32: 1 + 1 = 2, no overflow, SCC=0."""
instructions = [
s_mov_b32(s[0], 1),
s_add_i32(s[1], s[0], 1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 2)
self.assertEqual(st.scc, 0, "No overflow, SCC should be 0")
def test_s_add_i32_positive_overflow(self):
"""S_ADD_I32: MAX_INT + 1 overflows, SCC=1."""
instructions = [
s_mov_b32(s[0], 0x7FFFFFFF), # MAX_INT
s_add_i32(s[1], s[0], 1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0x80000000) # Wraps to MIN_INT
self.assertEqual(st.scc, 1, "Overflow, SCC should be 1")
def test_s_add_i32_negative_no_overflow(self):
"""S_ADD_I32: -10 + 20 = 10, no overflow."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFF6), # -10 in two's complement
s_mov_b32(s[1], 20),
s_add_i32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 10)
self.assertEqual(st.scc, 0)
def test_s_add_i32_negative_overflow(self):
"""S_ADD_I32: MIN_INT + (-1) underflows, SCC=1."""
instructions = [
s_mov_b32(s[0], 0x80000000), # MIN_INT
s_mov_b32(s[1], 0xFFFFFFFF), # -1
s_add_i32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0x7FFFFFFF) # Wraps to MAX_INT
self.assertEqual(st.scc, 1, "Underflow, SCC should be 1")
def test_s_sub_i32_no_overflow(self):
"""S_SUB_I32: 10 - 5 = 5, no overflow."""
instructions = [
s_mov_b32(s[0], 10),
s_mov_b32(s[1], 5),
s_sub_i32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 5)
self.assertEqual(st.scc, 0)
def test_s_sub_i32_overflow(self):
"""S_SUB_I32: MAX_INT - (-1) overflows, SCC=1."""
instructions = [
s_mov_b32(s[0], 0x7FFFFFFF), # MAX_INT
s_mov_b32(s[1], 0xFFFFFFFF), # -1
s_sub_i32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0x80000000) # Wraps to MIN_INT
self.assertEqual(st.scc, 1, "Overflow, SCC should be 1")
def test_s_mul_hi_u32(self):
"""S_MUL_HI_U32: high 32 bits of u32 * u32."""
instructions = [
s_mov_b32(s[0], 0x80000000), # 2^31
s_mov_b32(s[1], 4),
s_mul_hi_u32(s[2], s[0], s[1]), # (2^31 * 4) >> 32 = 2
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 2)
def test_s_mul_i32(self):
"""S_MUL_I32: signed multiply low 32 bits."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF), # -1
s_mov_b32(s[1], 10),
s_mul_i32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0xFFFFFFF6) # -10
def test_division_sequence_from_llvm(self):
"""Test the division sequence pattern from LLVM-generated code."""
# This sequence is from the sin kernel and computes integer division
# s10 = dividend, s18 = divisor, result in s6/s14
dividend = 0x28BE60DB # Some value from the sin kernel
divisor = 3 # Simplified divisor
instructions = [
s_mov_b32(s[10], dividend),
s_mov_b32(s[18], divisor),
# Compute reciprocal approximation: s6 = ~0 / divisor (approx)
s_mov_b32(s[11], 0),
s_sub_i32(s[11], s[11], s[18]), # s11 = -divisor
# For testing, just verify basic arithmetic works
s_mul_i32(s[6], s[10], 2),
s_add_i32(s[7], s[6], 1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[6], (dividend * 2) & 0xFFFFFFFF)
self.assertEqual(st.sgpr[7], ((dividend * 2) + 1) & 0xFFFFFFFF)
if __name__ == '__main__':
unittest.main()
File diff suppressed because it is too large Load Diff
-451
View File
@@ -1,451 +0,0 @@
"""Tests for VOP2 instructions - two operand vector operations.
Includes: v_add_f32, v_mul_f32, v_and_b32, v_or_b32, v_xor_b32,
v_lshrrev_b32, v_lshlrev_b32, v_fmac_f32, v_fmaak_f32, v_fmamk_f32,
v_add_nc_u32, v_cndmask_b32, v_add_f16, v_mul_f16
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestBasicArithmetic(unittest.TestCase):
"""Tests for basic arithmetic VOP2 instructions."""
def test_v_add_f32(self):
"""V_ADD_F32 adds two floats."""
instructions = [
v_mov_b32_e32(v[0], 1.0),
v_mov_b32_e32(v[1], 2.0),
v_add_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 3.0, places=5)
def test_v_mul_f32(self):
"""V_MUL_F32 multiplies two floats."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 4.0),
v_mul_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 8.0, places=5)
def test_v_fmac_f32(self):
"""V_FMAC_F32: d = d + a*b using inline constants."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 4.0),
v_mov_b32_e32(v[2], 1.0),
v_fmac_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
def test_v_fmaak_f32(self):
"""V_FMAAK_F32: d = a * b + K using inline constants."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 4.0),
v_fmaak_f32_e32(v[2], v[0], v[1], 0x3f800000),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
def test_v_fmamk_f32_basic(self):
"""V_FMAMK_F32: d = a * K + b."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 1.0),
v_fmamk_f32_e32(v[2], v[0], 0x40800000, v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
def test_v_fmamk_f32_small_constant(self):
"""V_FMAMK_F32 with small constant."""
instructions = [
v_mov_b32_e32(v[0], 4.0),
v_mov_b32_e32(v[1], 1.0),
v_fmamk_f32_e32(v[2], v[0], f2i(0.5), v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 3.0, places=5)
class TestBitManipulation(unittest.TestCase):
"""Tests for bit manipulation VOP2 instructions."""
def test_v_and_b32(self):
"""V_AND_B32 bitwise and."""
instructions = [
s_mov_b32(s[0], 0xff),
s_mov_b32(s[1], 0x0f),
v_mov_b32_e32(v[0], s[0]),
v_and_b32_e32(v[1], s[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x0f)
def test_v_and_b32_quadrant(self):
"""V_AND_B32 for quadrant extraction (n & 3)."""
instructions = [
s_mov_b32(s[0], 15915),
v_mov_b32_e32(v[0], s[0]),
v_and_b32_e32(v[1], 3, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 15915 & 3)
def test_v_lshrrev_b32(self):
"""V_LSHRREV_B32 logical shift right."""
instructions = [
s_mov_b32(s[0], 0xff00),
v_mov_b32_e32(v[0], s[0]),
v_lshrrev_b32_e32(v[1], 8, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xff)
def test_v_lshlrev_b32(self):
"""V_LSHLREV_B32 logical shift left."""
instructions = [
s_mov_b32(s[0], 0xff),
v_mov_b32_e32(v[0], s[0]),
v_lshlrev_b32_e32(v[1], 8, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xff00)
def test_v_xor_b32(self):
"""V_XOR_B32 bitwise xor (used in sin for sign)."""
instructions = [
s_mov_b32(s[0], 0x80000000),
s_mov_b32(s[1], f2i(1.0)),
v_mov_b32_e32(v[0], s[1]),
v_xor_b32_e32(v[1], s[0], v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), -1.0, places=5)
def test_v_xor_b32_sign_flip(self):
"""V_XOR_B32 for sign flip pattern."""
instructions = [
s_mov_b32(s[0], 0x80000000),
v_mov_b32_e32(v[0], -2.0),
v_xor_b32_e32(v[1], s[0], v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 2.0, places=5)
class TestSpecialValues(unittest.TestCase):
"""Tests for special float values - inf, nan, zero handling."""
def test_v_mul_f32_zero_times_inf(self):
"""V_MUL_F32: 0 * inf = NaN."""
import math
instructions = [
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[0], 0x7f800000),
v_mov_b32_e32(v[1], s[0]),
v_mul_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
def test_v_add_f32_inf_minus_inf(self):
"""V_ADD_F32: inf + (-inf) = NaN."""
import math
instructions = [
s_mov_b32(s[0], 0x7f800000),
s_mov_b32(s[1], 0xff800000),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
class TestF16Ops(unittest.TestCase):
"""Tests for 16-bit VOP2 operations."""
def test_v_add_f16_basic(self):
"""V_ADD_F16 adds two f16 values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4200, f"Expected 0x4200 (f16 3.0), got 0x{result:04x}")
def test_v_add_f16_negative(self):
"""V_ADD_F16 with negative values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0xc000), # f16 -2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0xbc00, f"Expected 0xbc00 (f16 -1.0), got 0x{result:04x}")
def test_v_mul_f16_basic(self):
"""V_MUL_F16 multiplies two f16 values."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
s_mov_b32(s[1], 0x4200), # f16 3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mul_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4600, f"Expected 0x4600 (f16 6.0), got 0x{result:04x}")
def test_v_mul_f16_by_zero(self):
"""V_MUL_F16 by zero."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0),
v_mul_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x0000, f"Expected 0x0000 (f16 0.0), got 0x{result:04x}")
def test_v_fmac_f16_basic(self):
"""V_FMAC_F16: d = d + a*b."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
s_mov_b32(s[1], 0x4200), # f16 3.0
s_mov_b32(s[2], 0x3c00), # f16 1.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]),
v_fmac_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
def test_v_fmaak_f16_basic(self):
"""V_FMAAK_F16: d = a * b + K."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
s_mov_b32(s[1], 0x4200), # f16 3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_fmaak_f16_e32(v[2], v[0], v[1], 0x3c00), # + f16 1.0
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
class TestHiHalfOps(unittest.TestCase):
"""Tests for VOP2 16-bit operations with hi-half operands."""
def test_v_add_f16_src0_hi_fold(self):
"""V_ADD_F16 with src0 hi-half fold (same register, different halves)."""
instructions = [
s_mov_b32(s[0], 0x40003c00), # lo=f16(1.0), hi=f16(2.0)
v_mov_b32_e32(v[0], s[0]),
VOP3(VOP3Op.V_ADD_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0001),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
def test_v_add_f16_src0_hi_different_reg(self):
"""V_ADD_F16 with src0 hi-half from different register."""
instructions = [
s_mov_b32(s[0], 0x40000000), # hi=f16(2.0), lo=0
s_mov_b32(s[1], 0x00003c00), # hi=0, lo=f16(1.0)
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
VOP3(VOP3Op.V_ADD_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
def test_v_mul_f16_src0_hi(self):
"""V_MUL_F16 with src0 from high half."""
instructions = [
s_mov_b32(s[0], 0x40000000), # hi=f16(2.0), lo=0
s_mov_b32(s[1], 0x00004200), # hi=0, lo=f16(3.0)
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
VOP3(VOP3Op.V_MUL_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4600, f"Expected f16(6.0)=0x4600, got 0x{result:04x}")
def test_v_mul_f16_hi_half(self):
"""V_MUL_F16 reading from high half."""
instructions = [
s_mov_b32(s[0], 0x40003c00), # lo=1.0, hi=2.0
v_mov_b32_e32(v[0], s[0]),
VOP3(VOP3Op.V_MUL_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0011),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 0x4400, f"Expected f16(4.0)=0x4400, got 0x{result:04x}")
def test_v_fma_f16_hi_dest(self):
"""V_FMA_F16 writing to high half with opsel.
Uses V_FMA_F16 (not V_FMAC_F16) because it has explicit src2 operand
which makes opsel handling clearer.
"""
instructions = [
s_mov_b32(s[0], 0x3c000000), # hi=f16(1.0), lo=0
s_mov_b32(s[1], 0x4000), # f16(2.0) in lo
s_mov_b32(s[2], 0x4200), # f16(3.0) in lo
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]),
# V_FMA_F16: dst = src0 * src1 + src2
# opsel=0b1100: bit2=src2 hi, bit3=dst hi
# So: v[0].hi = v[1].lo * v[2].lo + v[0].hi = 2.0 * 3.0 + 1.0 = 7.0
VOP3(VOP3Op.V_FMA_F16, vdst=v[0], src0=v[1], src1=v[2], src2=v[0], opsel=0b1100),
]
st = run_program(instructions, n_lanes=1)
hi = (st.vgpr[0][0] >> 16) & 0xffff
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(hi, 0x4700, f"Expected f16(7.0)=0x4700 in hi, got 0x{hi:04x}")
def test_v_add_f16_multilane(self):
"""V_ADD_F16 with multiple lanes."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=4)
for lane in range(4):
result = st.vgpr[lane][2] & 0xffff
self.assertEqual(result, 0x4200, f"Lane {lane}: expected 0x4200, got 0x{result:04x}")
class TestCndmask(unittest.TestCase):
"""Tests for V_CNDMASK_B32 and V_CNDMASK_B16."""
def test_v_cndmask_b16_select_src0(self):
"""V_CNDMASK_B16 selects src0 when VCC bit is 0."""
instructions = [
s_mov_b32(s[SrcEnum.VCC_LO - 128], 0), # VCC = 0
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cndmask_b16(v[2], v[0], v[1], VCC),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x3c00, f"Expected src0=0x3c00, got 0x{result:04x}")
def test_v_cndmask_b16_select_src1(self):
"""V_CNDMASK_B16 selects src1 when VCC bit is 1."""
instructions = [
s_mov_b32(s[SrcEnum.VCC_LO - 128], 1), # VCC = 1
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cndmask_b16(v[2], v[0], v[1], VCC),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4000, f"Expected src1=0x4000, got 0x{result:04x}")
def test_v_cndmask_b16_write_hi(self):
"""V_CNDMASK_B16 can write to high 16 bits with opsel."""
instructions = [
s_mov_b32(s[0], 0x3c003800), # src0: hi=1.0, lo=0.5
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x4000c000), # src1: hi=2.0, lo=-2.0
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], 0xDEAD0000), # v2 initial: hi=0xDEAD, lo=0
v_mov_b32_e32(v[2], s[2]),
s_mov_b32(s[SrcEnum.VCC_LO - 128], 0), # vcc = 0, select src0
# opsel=0b1011: bit0=src0 hi, bit1=src1 hi, bit3=dst hi
VOP3(VOP3Op.V_CNDMASK_B16, vdst=v[2], src0=v[0], src1=v[1], src2=SrcEnum.VCC_LO, opsel=0b1011),
]
st = run_program(instructions, n_lanes=1)
hi = (st.vgpr[0][2] >> 16) & 0xffff
lo = st.vgpr[0][2] & 0xffff
# vcc=0 selects src0.h = 1.0 = 0x3c00, writes to hi
self.assertEqual(hi, 0x3c00, f"Expected hi=0x3c00 (1.0), got 0x{hi:04x}")
self.assertEqual(lo, 0x0000, f"Expected lo preserved as 0, got 0x{lo:04x}")
class TestSpecialFloatValues(unittest.TestCase):
"""Tests for special float value handling in VOP2 instructions."""
def test_neg_zero_add(self):
"""-0.0 + 0.0 = +0.0 (IEEE 754)."""
neg_zero = 0x80000000
instructions = [
s_mov_b32(s[0], neg_zero),
v_mov_b32_e32(v[0], s[0]),
v_add_f32_e32(v[1], 0.0, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
def test_neg_zero_mul(self):
"""-0.0 * -1.0 = +0.0."""
neg_zero = 0x80000000
instructions = [
s_mov_b32(s[0], neg_zero),
v_mov_b32_e32(v[0], s[0]),
v_mul_f32_e32(v[1], -1.0, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
def test_inf_minus_inf(self):
"""+inf - inf = NaN."""
import math
pos_inf = 0x7f800000
neg_inf = 0xff800000
instructions = [
s_mov_b32(s[0], pos_inf),
s_mov_b32(s[1], neg_inf),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_sub_f32_e32(v[2], v[0], v[1]), # inf - (-inf) = inf
v_add_f32_e32(v[3], v[0], v[1]), # inf + (-inf) = NaN
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], pos_inf, "inf - (-inf) = inf")
self.assertTrue(math.isnan(i2f(st.vgpr[0][3])), "inf + (-inf) = NaN")
def test_denormal_f32_mul_ftz(self):
"""Denormal * normal - RDNA3 flushes denormals to zero (FTZ mode)."""
smallest_denorm = 0x00000001 # Smallest positive denormal
instructions = [
s_mov_b32(s[0], smallest_denorm),
v_mov_b32_e32(v[0], s[0]),
v_mul_f32_e32(v[1], 2.0, v[0]), # Denormal input gets flushed to 0
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x00000000)
if __name__ == '__main__':
unittest.main()
File diff suppressed because it is too large Load Diff
-538
View File
@@ -1,538 +0,0 @@
"""Tests for VOP3P instructions - packed 16-bit vector operations.
Includes: v_pk_add_f16, v_pk_mul_f16, v_pk_fma_f16, v_pack_b32_f16, v_wmma_*, v_dot2_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestPackInstructions(unittest.TestCase):
"""Tests for pack instructions."""
def test_v_pack_b32_f16(self):
"""V_PACK_B32_F16 packs two f16 values into one 32-bit register."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pack_b32_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x40003c00, f"Expected 0x40003c00, got 0x{result:08x}")
def test_v_pack_b32_f16_opsel_hi_hi(self):
"""V_PACK_B32_F16 with opsel to read high halves."""
inst = v_pack_b32_f16(v[2], v[0], v[1])
inst._values['opsel'] = 0b0011
instructions = [
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
inst,
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x44004000, f"Expected 0x44004000, got 0x{result:08x}")
class TestPackMore(unittest.TestCase):
"""Additional pack instruction tests."""
def test_v_pack_b32_f16_basic(self):
"""V_PACK_B32_F16 packs two f16 values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pack_b32_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x40003c00, f"Expected 0x40003c00, got 0x{result:08x}")
def test_v_pack_b32_f16_with_cvt(self):
"""V_PACK_B32_F16 after V_CVT_F16_F32 conversions."""
instructions = [
s_mov_b32(s[0], 0x3f800000),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]),
v_cvt_f16_f32_e32(v[2], v[0]),
v_cvt_f16_f32_e32(v[3], v[1]),
v_pack_b32_f16(v[4], v[2], v[3]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][4]
self.assertEqual(result, 0x3c003c00, f"Expected 0x3c003c00, got 0x{result:08x}")
def test_v_pack_b32_f16_packed_sources(self):
"""V_PACK_B32_F16 with packed f16 sources (reads lo halves)."""
instructions = [
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pack_b32_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
# Expected: hi=v1.lo=0x4200 (3.0), lo=v0.lo=0x3c00 (1.0) -> 0x42003c00
self.assertEqual(result, 0x42003c00, f"Expected 0x42003c00, got 0x{result:08x}")
def test_v_pack_b32_f16_opsel_lo_hi(self):
"""V_PACK_B32_F16 with opsel=0b0010 to read lo from src0, hi from src1."""
inst = v_pack_b32_f16(v[2], v[0], v[1])
inst._values['opsel'] = 0b0010
instructions = [
s_mov_b32(s[0], 0x40003c00),
s_mov_b32(s[1], 0x44004200),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
inst,
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x44003c00, f"Expected 0x44003c00, got 0x{result:08x}")
def test_v_pack_b32_f16_opsel_hi_lo(self):
"""V_PACK_B32_F16 with opsel=0b0001 to read hi from src0, lo from src1."""
inst = v_pack_b32_f16(v[2], v[0], v[1])
inst._values['opsel'] = 0b0001
instructions = [
s_mov_b32(s[0], 0x40003c00),
s_mov_b32(s[1], 0x44004200),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
inst,
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x42004000, f"Expected 0x42004000, got 0x{result:08x}")
def test_v_pack_b32_f16_zeros(self):
"""V_PACK_B32_F16 with zero values."""
instructions = [
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
v_pack_b32_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0)
def test_v_pack_b32_f16_both_positive(self):
"""V_PACK_B32_F16 with positive f16 values."""
instructions = [
s_mov_b32(s[0], 0x4200), # f16 3.0
s_mov_b32(s[1], 0x4400), # f16 4.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pack_b32_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x44004200, f"Expected 0x44004200, got 0x{result:08x}")
class TestFmaMix(unittest.TestCase):
"""Tests for V_FMA_MIX_F32 and V_FMA_MIXLO_F16."""
def test_v_fma_mix_f32_all_f32_sources(self):
"""V_FMA_MIX_F32 with all f32 sources."""
instructions = [
s_mov_b32(s[0], f2i(2.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(3.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], f2i(1.0)),
v_mov_b32_e32(v[2], s[2]),
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 7.0, places=5)
def test_v_fma_mix_f32_src2_f16_lo(self):
"""V_FMA_MIX_F32 with src2 as f16 from lo bits."""
from extra.assembly.amd.pcode import f32_to_f16
f16_2 = f32_to_f16(2.0)
instructions = [
s_mov_b32(s[0], f2i(1.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(3.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], f16_2),
v_mov_b32_e32(v[2], s[2]),
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 5.0, places=5)
def test_v_fma_mix_f32_src2_f16_hi(self):
"""V_FMA_MIX_F32 with src2 as f16 from hi bits."""
from extra.assembly.amd.pcode import f32_to_f16
f16_2 = f32_to_f16(2.0)
val = (f16_2 << 16) | 0
instructions = [
s_mov_b32(s[0], f2i(1.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(3.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], val),
v_mov_b32_e32(v[2], s[2]),
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=4, opsel_hi=0, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 5.0, places=5)
def test_v_fma_mix_f32_with_abs(self):
"""V_FMA_MIX_F32 with abs modifier on src2."""
instructions = [
s_mov_b32(s[0], f2i(2.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(3.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], f2i(-1.0)),
v_mov_b32_e32(v[2], s[2]),
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0, neg_hi=4),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 7.0, places=5)
def test_v_fma_mixlo_f16(self):
"""V_FMA_MIXLO_F16 writes to low 16 bits of destination."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], f2i(2.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(3.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], f2i(1.0)),
v_mov_b32_e32(v[2], s[2]),
s_mov_b32(s[3], 0xdead0000),
v_mov_b32_e32(v[3], s[3]),
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
]
st = run_program(instructions, n_lanes=1)
lo = _f16(st.vgpr[0][3] & 0xffff)
hi = (st.vgpr[0][3] >> 16) & 0xffff
self.assertAlmostEqual(lo, 7.0, places=1)
self.assertEqual(hi, 0xdead, f"hi should be preserved, got 0x{hi:04x}")
def test_v_fma_mixlo_f16_all_f32_sources(self):
"""V_FMA_MIXLO_F16 with all f32 sources."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], f2i(1.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(2.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], f2i(3.0)),
v_mov_b32_e32(v[2], s[2]),
v_mov_b32_e32(v[3], 0),
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
]
st = run_program(instructions, n_lanes=1)
lo = _f16(st.vgpr[0][3] & 0xffff)
# 1*2+3 = 5
self.assertAlmostEqual(lo, 5.0, places=1)
def test_v_fma_mixlo_f16_sin_case(self):
"""V_FMA_MIXLO_F16 case from sin kernel."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x3f800000), # f32 1.0
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[1], 0xaf05a309), # f32 tiny negative
s_mov_b32(s[6], s[1]),
s_mov_b32(s[2], 0xc0490fdb), # f32 -π
v_mov_b32_e32(v[5], s[2]),
s_mov_b32(s[3], 0x3f800000),
v_mov_b32_e32(v[3], s[3]),
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[3], src1=s[6], src2=v[5], opsel=0, opsel_hi=0, opsel_hi2=0),
]
st = run_program(instructions, n_lanes=1)
lo = _f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(lo, -3.14159, delta=0.01)
class TestVOP3P(unittest.TestCase):
"""Tests for VOP3P packed 16-bit operations."""
def test_v_pk_add_f16_basic(self):
"""V_PK_ADD_F16 adds two packed f16 values."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_add_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 4.0, places=2)
self.assertAlmostEqual(hi, 6.0, places=2)
def test_v_pk_mul_f16_basic(self):
"""V_PK_MUL_F16 multiplies two packed f16 values."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x42004000), # hi=3.0, lo=2.0
s_mov_b32(s[1], 0x45004400), # hi=5.0, lo=4.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_mul_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 8.0, places=1)
self.assertAlmostEqual(hi, 15.0, places=1)
def test_v_pk_fma_f16_basic(self):
"""V_PK_FMA_F16: D = A * B + C for packed f16."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x42004000), # A: hi=3.0, lo=2.0
s_mov_b32(s[1], 0x45004400), # B: hi=5.0, lo=4.0
s_mov_b32(s[2], 0x3c003c00), # C: hi=1.0, lo=1.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]),
v_pk_fma_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 9.0, places=1) # 2*4+1
self.assertAlmostEqual(hi, 16.0, places=0) # 3*5+1
def test_v_pk_add_f16_with_inline_constant(self):
"""V_PK_ADD_F16 with inline constant POS_ONE (1.0).
Inline constants for VOP3P are f16 values in the low 16 bits only.
hi half of inline constant is 0, so hi result = v0.hi + 0 = 1.0.
"""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x3c003c00), # packed f16: hi=1.0, lo=1.0
v_mov_b32_e32(v[0], s[0]),
v_pk_add_f16(v[1], v[0], SrcEnum.POS_ONE, opsel_hi=3, opsel_hi2=1), # Add inline constant 1.0
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
# lo = 1.0 + 1.0 = 2.0, hi = 1.0 + 0.0 = 1.0 (inline const hi half is 0)
self.assertAlmostEqual(lo, 2.0, places=2)
self.assertAlmostEqual(hi, 1.0, places=2)
def test_v_pk_mul_f16_with_inline_constant(self):
"""V_PK_MUL_F16 with inline constant POS_TWO (2.0).
Inline constant has value only in low 16 bits, hi is 0.
"""
from extra.assembly.amd.pcode import _f16
# v0 = packed (3.0, 4.0), multiply by POS_TWO
# lo = 3.0 * 2.0 = 6.0, hi = 4.0 * 0.0 = 0.0 (inline const hi is 0)
instructions = [
s_mov_b32(s[0], 0x44004200), # packed f16: hi=4.0, lo=3.0
v_mov_b32_e32(v[0], s[0]),
v_pk_mul_f16(v[1], v[0], SrcEnum.POS_TWO, opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 6.0, places=1)
self.assertAlmostEqual(hi, 0.0, places=1)
class TestWMMA(unittest.TestCase):
"""Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions."""
def test_v_wmma_f32_16x16x16_f16_all_ones(self):
"""V_WMMA_F32_16X16X16_F16 with all ones produces 16.0."""
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
for i in range(16, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], 0))
instructions.append(v_wmma_f32_16x16x16_f16(v[0], v[16], v[24], v[0]))
st = run_program(instructions, n_lanes=32)
expected = f2i(16.0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
def test_v_wmma_f32_16x16x16_f16_with_accumulator(self):
"""V_WMMA_F32_16X16X16_F16 with non-zero accumulator."""
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00))
instructions.append(s_mov_b32(s[1], f2i(5.0)))
for i in range(16, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], s[1]))
instructions.append(v_wmma_f32_16x16x16_f16(v[0], v[16], v[24], v[0]))
st = run_program(instructions, n_lanes=32)
expected = f2i(21.0) # 16 + 5
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
class TestSpecialOps(unittest.TestCase):
"""Tests for special operations (SAD, PERM, DOT2)."""
def test_v_sad_u8_basic(self):
"""V_SAD_U8 computes sum of absolute differences."""
instructions = [
s_mov_b32(s[0], 0x04030201), # bytes: 1, 2, 3, 4
s_mov_b32(s[1], 0x05040302), # bytes: 2, 3, 4, 5
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_sad_u8(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
# |1-2| + |2-3| + |3-4| + |4-5| = 1 + 1 + 1 + 1 = 4
self.assertEqual(st.vgpr[0][3], 4)
def test_v_sad_u8_identical_bytes(self):
"""V_SAD_U8 with identical inputs returns accumulator."""
instructions = [
s_mov_b32(s[0], 0x04030201),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 10),
v_mov_b32_e32(v[2], s[1]),
v_sad_u8(v[3], v[0], v[0], v[2]),
]
st = run_program(instructions, n_lanes=1)
# Same inputs -> SAD = 0, result = accumulator = 10
self.assertEqual(st.vgpr[0][3], 10)
def test_v_sad_u16_basic(self):
"""V_SAD_U16 computes sum of absolute differences of u16 pairs."""
instructions = [
s_mov_b32(s[0], 0x00030001), # hi=3, lo=1
s_mov_b32(s[1], 0x00050002), # hi=5, lo=2
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_sad_u16(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
# |1-2| + |3-5| = 1 + 2 = 3
self.assertEqual(st.vgpr[0][3], 3)
def test_v_sad_u32_basic(self):
"""V_SAD_U32 computes absolute difference of u32 values."""
instructions = [
s_mov_b32(s[0], 100),
s_mov_b32(s[1], 70),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_sad_u32(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
# |100-70| = 30
self.assertEqual(st.vgpr[0][3], 30)
def test_v_msad_u8_masked(self):
"""V_MSAD_U8 masked SAD operation."""
instructions = [
s_mov_b32(s[0], 0x04030201),
s_mov_b32(s[1], 0x05040302),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_msad_u8(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
# V_MSAD_U8 skips bytes where src0 is 0
# Since no bytes are 0, result same as V_SAD_U8 = 4
self.assertEqual(st.vgpr[0][3], 4)
def test_v_perm_b32_select_bytes(self):
"""V_PERM_B32 selects bytes from two sources.
V_PERM_B32 concatenates {S1, S0} as a 64-bit value with S1 in low 32 bits.
Selector byte values 0-3 select from S1, values 4-7 select from S0.
"""
instructions = [
s_mov_b32(s[0], 0x44332211), # src0: bytes 4-7 in 64-bit view
s_mov_b32(s[1], 0x88776655), # src1: bytes 0-3 in 64-bit view
s_mov_b32(s[2], 0x07060504), # select bytes 4,5,6,7 (from src0)
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_perm_b32(v[2], v[0], v[1], s[2]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x44332211)
def test_v_dot2_f32_bf16_basic(self):
"""V_DOT2_F32_BF16 computes dot product of bf16 pairs."""
# bf16 1.0 = 0x3f80, bf16 2.0 = 0x4000
instructions = [
s_mov_b32(s[0], 0x3f803f80), # packed bf16: lo=1.0, hi=1.0
s_mov_b32(s[1], 0x40003f80), # packed bf16: lo=1.0, hi=2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot2_f32_bf16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
# 1.0*1.0 + 1.0*2.0 + 0 = 3.0
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 3.0, places=4)
class TestPackedMixedSigns(unittest.TestCase):
"""Tests for packed operations with mixed sign values."""
def test_pk_add_f16_mixed_signs(self):
"""V_PK_ADD_F16 with mixed positive/negative values."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0xc0003c00), # packed: hi=-2.0, lo=1.0
s_mov_b32(s[1], 0x3c003c00), # packed: hi=1.0, lo=1.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_add_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 2.0, places=2) # 1.0 + 1.0
self.assertAlmostEqual(hi, -1.0, places=2) # -2.0 + 1.0
def test_pk_mul_f16_zero(self):
"""V_PK_MUL_F16 with zero."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x40004000), # packed: 2.0, 2.0
s_mov_b32(s[1], 0x00000000), # packed: 0.0, 0.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_mul_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x00000000, "2.0 * 0.0 should be 0.0")
if __name__ == '__main__':
unittest.main()
-486
View File
@@ -1,486 +0,0 @@
"""Tests for VOPC instructions - vector compare operations.
Includes: v_cmp_class_f32, v_cmp_class_f16, v_cmp_eq_*, v_cmp_lt_*, v_cmp_gt_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
VCC = 106 # SGPR index for VCC_LO
class TestCmpClass(unittest.TestCase):
"""Tests for V_CMP_CLASS_F32 float classification."""
def test_cmp_class_quiet_nan(self):
"""V_CMP_CLASS_F32 detects quiet NaN."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000010), # bit 1 = quiet NaN
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
def test_cmp_class_signaling_nan(self):
"""V_CMP_CLASS_F32 detects signaling NaN."""
signal_nan = 0x7f800001
instructions = [
s_mov_b32(s[0], signal_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000001), # bit 0 = signaling NaN
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
def test_cmp_class_positive_inf(self):
"""V_CMP_CLASS_F32 detects +inf."""
pos_inf = 0x7f800000
instructions = [
s_mov_b32(s[0], pos_inf),
s_mov_b32(s[1], 0b1000000000), # bit 9 = +inf
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect +inf")
def test_cmp_class_negative_inf(self):
"""V_CMP_CLASS_F32 detects -inf."""
neg_inf = 0xff800000
instructions = [
s_mov_b32(s[0], neg_inf),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000100), # bit 2 = -inf
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect -inf")
def test_cmp_class_normal_positive(self):
"""V_CMP_CLASS_F32 detects positive normal."""
instructions = [
v_mov_b32_e32(v[0], 1.0),
s_mov_b32(s[1], 0b0100000000), # bit 8 = positive normal
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
def test_cmp_class_normal_negative(self):
"""V_CMP_CLASS_F32 detects negative normal."""
instructions = [
v_mov_b32_e32(v[0], -1.0),
v_mov_b32_e32(v[1], 0b0000001000), # bit 3 = negative normal
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
def test_cmp_class_quiet_nan_not_signaling(self):
"""Quiet NaN does not match signaling NaN mask."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000001), # bit 0 = signaling NaN only
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "Quiet NaN should not match signaling mask")
def test_cmp_class_signaling_nan_not_quiet(self):
"""Signaling NaN does not match quiet NaN mask."""
signal_nan = 0x7f800001
instructions = [
s_mov_b32(s[0], signal_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000010), # bit 1 = quiet NaN only
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "Signaling NaN should not match quiet mask")
def test_v_cmp_sets_vcc_bits(self):
"""V_CMP_EQ sets VCC bits based on per-lane comparison."""
instructions = [
s_mov_b32(s[0], 5),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]),
v_cmp_eq_u32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
class TestCmpClassF16(unittest.TestCase):
"""Tests for V_CMP_CLASS_F16 float classification.
Class bit mapping:
bit 0 = signaling NaN
bit 1 = quiet NaN
bit 2 = -infinity
bit 3 = -normal
bit 4 = -denormal
bit 5 = -zero
bit 6 = +zero
bit 7 = +denormal
bit 8 = +normal
bit 9 = +infinity
"""
def test_cmp_class_f16_positive_zero(self):
"""V_CMP_CLASS_F16: +zero matches bit 6."""
instructions = [
v_mov_b32_e32(v[0], 0x0000), # f16 +0.0
v_mov_b32_e32(v[1], 0x40), # bit 6 = +zero
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive zero")
def test_cmp_class_f16_negative_zero(self):
"""V_CMP_CLASS_F16: -zero matches bit 5."""
instructions = [
s_mov_b32(s[0], 0x8000), # f16 -0.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x20), # bit 5 = -zero
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative zero")
def test_cmp_class_f16_positive_normal(self):
"""V_CMP_CLASS_F16: +1.0 (normal) matches bit 8."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 +1.0
s_mov_b32(s[1], 0x100), # bit 8 = +normal
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
def test_cmp_class_f16_negative_normal(self):
"""V_CMP_CLASS_F16: -1.0 (normal) matches bit 3."""
instructions = [
s_mov_b32(s[0], 0xbc00), # f16 -1.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x08), # bit 3 = -normal
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
def test_cmp_class_f16_positive_infinity(self):
"""V_CMP_CLASS_F16: +inf matches bit 9."""
instructions = [
s_mov_b32(s[0], 0x7c00), # f16 +inf
s_mov_b32(s[1], 0x200), # bit 9 = +inf
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive infinity")
def test_cmp_class_f16_negative_infinity(self):
"""V_CMP_CLASS_F16: -inf matches bit 2."""
instructions = [
s_mov_b32(s[0], 0xfc00), # f16 -inf
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x04), # bit 2 = -inf
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative infinity")
def test_cmp_class_f16_quiet_nan(self):
"""V_CMP_CLASS_F16: quiet NaN matches bit 1."""
instructions = [
s_mov_b32(s[0], 0x7e00), # f16 quiet NaN
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x02), # bit 1 = quiet NaN
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
def test_cmp_class_f16_signaling_nan(self):
"""V_CMP_CLASS_F16: signaling NaN matches bit 0."""
instructions = [
s_mov_b32(s[0], 0x7c01), # f16 signaling NaN
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x01), # bit 0 = signaling NaN
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
def test_cmp_class_f16_positive_denormal(self):
"""V_CMP_CLASS_F16: positive denormal matches bit 7."""
instructions = [
v_mov_b32_e32(v[0], 1), # f16 +denormal (0x0001)
v_mov_b32_e32(v[1], 0x80), # bit 7 = +denormal
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive denormal")
def test_cmp_class_f16_negative_denormal(self):
"""V_CMP_CLASS_F16: negative denormal matches bit 4."""
instructions = [
s_mov_b32(s[0], 0x8001), # f16 -denormal
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x10), # bit 4 = -denormal
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative denormal")
def test_cmp_class_f16_combined_mask_zeros(self):
"""V_CMP_CLASS_F16: mask 0x60 covers both +zero and -zero."""
instructions = [
v_mov_b32_e32(v[0], 0), # f16 +0.0
v_mov_b32_e32(v[1], 0x60), # bits 5 and 6 (+-zero)
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x60")
def test_cmp_class_f16_combined_mask_1f8(self):
"""V_CMP_CLASS_F16: mask 0x1f8 covers -normal,-denorm,-zero,+zero,+denorm,+normal.
This is the exact mask used in the f16 sin kernel at PC=46.
"""
instructions = [
v_mov_b32_e32(v[0], 0), # f16 +0.0
s_mov_b32(s[0], 0x1f8),
v_mov_b32_e32(v[1], s[0]), # mask 0x1f8
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x1f8")
def test_cmp_class_f16_vop3_encoding(self):
"""V_CMP_CLASS_F16 in VOP3 encoding (v_cmp_class_f16_e64)."""
instructions = [
v_mov_b32_e32(v[0], 0), # f16 +0.0
s_mov_b32(s[0], 0x1f8), # class mask
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with VOP3 encoding")
def test_cmp_class_f16_vop3_normal_positive(self):
"""V_CMP_CLASS_F16 VOP3 encoding with +1.0 (normal)."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 +1.0
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x1f8), # class mask
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +1.0 (normal) with mask 0x1f8")
def test_cmp_class_f16_vop3_nan_fails_mask(self):
"""V_CMP_CLASS_F16 VOP3: NaN should NOT match mask 0x1f8 (no NaN bits set)."""
instructions = [
s_mov_b32(s[0], 0x7e00), # f16 quiet NaN
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x1f8), # class mask
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for NaN with mask 0x1f8 (no NaN bits)")
def test_cmp_class_f16_vop3_inf_fails_mask(self):
"""V_CMP_CLASS_F16 VOP3: +inf should NOT match mask 0x1f8 (no inf bits set)."""
instructions = [
s_mov_b32(s[0], 0x7c00), # f16 +inf
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x1f8), # class mask
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for +inf with mask 0x1f8 (no inf bits)")
class TestCmpInt(unittest.TestCase):
"""Tests for integer comparison operations."""
def test_v_cmp_eq_u32(self):
"""V_CMP_EQ_U32 sets VCC bits based on per-lane comparison."""
instructions = [
s_mov_b32(s[0], 5),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]),
v_cmp_eq_u32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
def test_cmp_eq_u16_opsel_lo_lo(self):
"""V_CMP_EQ_U16 comparing lo halves."""
instructions = [
s_mov_b32(s[0], 0x12340005), # lo=5, hi=0x1234
s_mov_b32(s[1], 0xABCD0005), # lo=5, hi=0xABCD
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_eq_u16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Lo halves should be equal")
def test_cmp_eq_u16_opsel_hi_hi(self):
"""V_CMP_EQ_U16 comparing hi halves with VOP3 opsel.
VOPC doesn't have opsel, so we use VOP3 form for hi-half comparisons.
VOP3 compares write result to SGPR via vdst field.
"""
instructions = [
s_mov_b32(s[2], 0x00051234), # hi=5, lo=0x1234
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x0005ABCD), # hi=5, lo=0xABCD
v_mov_b32_e32(v[1], s[2]),
# opsel=3 means compare hi halves, vdst=v[0] actually writes to s[0]
VOP3(VOP3Op.V_CMP_EQ_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
# Result is in sgpr[0], not vcc
self.assertEqual(st.sgpr[0] & 1, 1, "Hi halves should be equal: 5==5")
def test_cmp_eq_u16_opsel_hi_hi_equal(self):
"""V_CMP_EQ_U16 VOP3 with opsel=3 compares hi halves (equal case)."""
instructions = [
s_mov_b32(s[2], 0x12340005), # lo=5, hi=0x1234
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x12340009), # lo=9, hi=0x1234
v_mov_b32_e32(v[1], s[2]),
VOP3(VOP3Op.V_CMP_EQ_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0] & 1, 1, "hi==hi should be true: 0x1234==0x1234")
def test_cmp_gt_u16_opsel_hi(self):
"""V_CMP_GT_U16 VOP3 with opsel=3 compares hi halves."""
instructions = [
s_mov_b32(s[2], 0x99990005), # lo=5, hi=0x9999
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x12340005), # lo=5, hi=0x1234
v_mov_b32_e32(v[1], s[2]),
VOP3(VOP3Op.V_CMP_GT_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0] & 1, 1, "hi>hi should be true: 0x9999>0x1234")
class TestCmpFloat(unittest.TestCase):
"""Tests for float comparison operations."""
def test_v_cmp_lt_f16_vsrc1_hi(self):
"""V_CMP_LT_F16 with both operands from high half using VOP3 opsel."""
instructions = [
s_mov_b32(s[2], 0x3c000000), # hi=1.0 (f16), lo=0
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x40000000), # hi=2.0 (f16), lo=0
v_mov_b32_e32(v[1], s[2]),
# opsel=3 means read hi halves for both src0 and src1
VOP3(VOP3Op.V_CMP_LT_F16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
# Result is in sgpr[0]
self.assertEqual(st.sgpr[0] & 1, 1, "1.0 < 2.0 should be true")
def test_v_cmp_gt_f16_vsrc1_hi(self):
"""V_CMP_GT_F16 with both operands from high half using VOP3 opsel."""
instructions = [
s_mov_b32(s[2], 0x40000000), # hi=2.0 (f16), lo=0
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x3c000000), # hi=1.0 (f16), lo=0
v_mov_b32_e32(v[1], s[2]),
# opsel=3 means read hi halves for both src0 and src1
VOP3(VOP3Op.V_CMP_GT_F16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
# Result is in sgpr[0]
self.assertEqual(st.sgpr[0] & 1, 1, "2.0 > 1.0 should be true")
def test_v_cmp_eq_f16_vsrc1_hi_equal(self):
"""v_cmp_eq_f16 with equal low and high halves."""
instructions = [
s_mov_b32(s[0], 0x42004200), # hi=3.0 (0x4200), lo=3.0 (0x4200)
v_mov_b32_e32(v[0], s[0]),
v_cmp_eq_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (3.0 == 3.0)")
def test_v_cmp_neq_f16_vsrc1_hi(self):
"""v_cmp_neq_f16 with different low and high halves."""
instructions = [
s_mov_b32(s[0], 0x40003c00), # hi=2.0 (0x4000), lo=1.0 (0x3c00)
v_mov_b32_e32(v[0], s[0]),
v_cmp_lg_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (1.0 != 2.0)")
def test_v_cmp_nge_f16_inf_self(self):
"""v_cmp_nge_f16 comparing -inf with itself (unordered less than).
Regression test: -inf < -inf should be false (IEEE 754).
"""
instructions = [
s_mov_b32(s[0], 0xFC00FC00), # both halves = -inf (0xFC00)
v_mov_b32_e32(v[0], s[0]),
v_cmp_nge_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "Expected vcc=0 (-inf >= -inf)")
def test_v_cmp_f16_multilane(self):
"""v_cmp_lt_f16 with vsrc1=v128 across multiple lanes."""
instructions = [
# Lane 0: v0 = 0x40003c00 (hi=2.0, lo=1.0) -> 1.0 < 2.0 = true
# Lane 1: v0 = 0x3c004000 (hi=1.0, lo=2.0) -> 2.0 < 1.0 = false
v_mov_b32_e32(v[0], 0x40003c00), # default
v_cmp_eq_u32_e32(1, v[255]), # vcc = (lane == 1)
v_cndmask_b32_e64(v[0], v[0], 0x3c004000, SrcEnum.VCC_LO),
v_cmp_lt_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=2)
self.assertEqual(st.vcc & 1, 1, "Lane 0: expected vcc=1 (1.0 < 2.0)")
self.assertEqual((st.vcc >> 1) & 1, 0, "Lane 1: expected vcc=0 (2.0 < 1.0)")
class TestVCCBehavior(unittest.TestCase):
"""Tests for VCC condition code behavior."""
def test_vcc_all_lanes_true(self):
"""VCC should have all bits set when all lanes compare true."""
instructions = [
v_mov_b32_e32(v[0], 5),
v_mov_b32_e32(v[1], 5),
v_cmp_eq_u32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=32)
self.assertEqual(st.vcc, 0xFFFFFFFF, "All 32 lanes should be true")
def test_vcc_lane_dependent(self):
"""VCC should differ per lane based on lane_id comparison."""
instructions = [
v_mov_b32_e32(v[0], 16),
v_cmp_lt_u32_e32(v[255], v[0]), # lanes 0-15 are < 16
]
st = run_program(instructions, n_lanes=32)
self.assertEqual(st.vcc & 0xFFFF, 0xFFFF, "Lanes 0-15 should be true")
self.assertEqual(st.vcc >> 16, 0x0000, "Lanes 16-31 should be false")
if __name__ == '__main__':
unittest.main()
@@ -1,401 +0,0 @@
# Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
import unittest, ctypes, os
from dataclasses import dataclass
from pathlib import Path
# Set environment before any tinygrad imports to use MOCKGPU
# This allows generating AMD GPU kernels without requiring real hardware
os.environ["AMD"] = "1"
os.environ["MOCKGPU"] = "1"
os.environ["PYTHON_REMU"] = "1"
from extra.assembly.amd.emu import WaveState, decode_program, WAVE_SIZE, set_valid_mem_ranges, LDSMem
from extra.assembly.amd.test.helpers import KernelInfo
from extra.assembly.amd.test.bench_emu import REMU_PATH
def _is_f32_nan(bits: int) -> bool:
"""Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
return (bits & 0x7f800000) == 0x7f800000 and (bits & 0x007fffff) != 0
def _vals_equal(a: int, b: int) -> bool:
"""Compare two 32-bit values, treating all NaN bit patterns as equal."""
if a == b: return True
return _is_f32_nan(a) and _is_f32_nan(b)
@dataclass
class StateSnapshot:
pc: int
scc: int
vcc: int
exec_mask: int
sgpr: list[int]
vgpr: list[list[int]]
def diff(self, other: 'StateSnapshot', n_lanes: int, arrow: str = " vs ") -> list[str]:
"""Return list of differences between two states."""
diffs = []
if self.pc != other.pc: diffs.append(f"pc: {self.pc}{arrow}{other.pc}")
if self.scc != other.scc: diffs.append(f"scc: {self.scc}{arrow}{other.scc}")
if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x}{arrow}0x{other.vcc:08x}")
if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x}{arrow}0x{other.exec_mask:08x}")
for i, (a, b) in enumerate(zip(self.sgpr, other.sgpr)):
# Skip VCC_LO/HI (106/107) and EXEC_LO/HI (126/127) as they alias vcc/exec_mask which are compared separately
if i in (106, 107, 126, 127): continue
if not _vals_equal(a, b): diffs.append(f"sgpr[{i}]: 0x{a:08x}{arrow}0x{b:08x}")
for lane in range(n_lanes):
for i, (a, b) in enumerate(zip(self.vgpr[lane], other.vgpr[lane])):
if not _vals_equal(a, b): diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x}{arrow}0x{b:08x}")
return diffs
class CStateSnapshot(ctypes.Structure):
_fields_ = [("pc", ctypes.c_uint32), ("scc", ctypes.c_uint32), ("vcc", ctypes.c_uint32), ("exec_mask", ctypes.c_uint32),
("sgpr", ctypes.c_uint32 * 128), ("vgpr", (ctypes.c_uint32 * 256) * 32)]
def to_snapshot(self) -> StateSnapshot:
return StateSnapshot(pc=self.pc, scc=self.scc, vcc=self.vcc, exec_mask=self.exec_mask,
sgpr=list(self.sgpr), vgpr=[list(self.vgpr[i]) for i in range(32)])
class RustEmulator:
def __init__(self):
self.lib = ctypes.CDLL(str(REMU_PATH))
self.lib.wave_create.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
self.lib.wave_create.restype = ctypes.c_void_p
self.lib.wave_step.argtypes = [ctypes.c_void_p]
self.lib.wave_step.restype = ctypes.c_int32
self.lib.wave_get_snapshot.argtypes = [ctypes.c_void_p, ctypes.POINTER(CStateSnapshot)]
self.lib.wave_set_sgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
self.lib.wave_set_vgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32]
self.lib.wave_init_lds.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
self.lib.wave_free.argtypes = [ctypes.c_void_p]
self.ctx = None
def create(self, kernel: bytes, n_lanes: int):
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
self.ctx = self.lib.wave_create(ctypes.addressof(kernel_buf), len(kernel), n_lanes)
self._kernel_buf = kernel_buf
def step(self) -> int: return self.lib.wave_step(self.ctx)
def set_sgpr(self, idx: int, val: int): self.lib.wave_set_sgpr(self.ctx, idx, val)
def set_vgpr(self, lane: int, idx: int, val: int): self.lib.wave_set_vgpr(self.ctx, lane, idx, val)
def init_lds(self, size: int): self.lib.wave_init_lds(self.ctx, size)
def get_snapshot(self) -> StateSnapshot:
snap = CStateSnapshot()
self.lib.wave_get_snapshot(self.ctx, ctypes.byref(snap))
return snap.to_snapshot()
def free(self):
if self.ctx: self.lib.wave_free(self.ctx); self.ctx = None
class PythonEmulator:
def __init__(self):
self.state: WaveState | None = None
self.program: dict | None = None
def create(self, kernel: bytes, n_lanes: int):
self.program = decode_program(kernel)
self.state = WaveState(LDSMem(bytearray(65536)), n_lanes)
self.state.exec_mask = (1 << n_lanes) - 1
def step(self) -> int:
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
def set_vgpr(self, lane: int, idx: int, val: int):
assert self.state is not None
self.state.vgpr[lane][idx] = val & 0xffffffff
def get_snapshot(self) -> StateSnapshot:
assert self.state is not None
return StateSnapshot(pc=self.state.pc, scc=self.state.scc, vcc=self.state.vcc & 0xffffffff,
exec_mask=self.state.exec_mask & 0xffffffff, sgpr=list(self.state.sgpr),
vgpr=[list(self.state.vgpr[i]) for i in range(WAVE_SIZE)])
def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
max_workgroups: int = 8) -> tuple[bool, str, int]:
"""Run a single kernel through both emulators. Returns (success, message, total_steps)."""
gx, gy, gz = global_size
total_steps = 0
wg_count = 0
for gidz in range(gz):
for gidy in range(gy):
for gidx in range(gx):
if wg_count >= max_workgroups: return True, f"Completed {wg_count} workgroups (limit reached)", total_steps
wg_count += 1
rust = RustEmulator()
python = PythonEmulator()
rust.create(kernel, n_lanes)
python.create(kernel, n_lanes)
# Initialize LDS (64KB, standard size for AMD GPUs)
rust.init_lds(65536)
for emu in (rust, python):
emu.set_sgpr(0, args_ptr & 0xffffffff)
emu.set_sgpr(1, (args_ptr >> 32) & 0xffffffff)
emu.set_sgpr(13, gidx)
emu.set_sgpr(14, gidy)
emu.set_sgpr(15, gidz)
step = 0
trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
try:
while step < max_steps:
rust_before = rust.get_snapshot()
python_before = python.get_snapshot()
inst = program.get(python_before.pc)
inst_str = inst.disasm() if inst else f"unknown at PC={python_before.pc}"
trace.append((step, python_before.pc, inst_str, rust_before, python_before))
if len(trace) > trace_len: trace.pop(0)
if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
# Instructions with known Rust emulator bugs - sync Python to Rust after execution
# v_div_scale/v_div_fixup: Rust has different VCC handling
# v_cvt_f16_f32: Rust clears high 16 bits, but hardware (and Python) preserves them
# s_add_i32/s_sub_i32: Rust has incorrect SCC overflow detection
sync_after = any(x in inst_str for x in ('v_div_scale_f32', 'v_div_scale_f64', 'v_div_fixup_f32', 'v_div_fixup_f64',
'v_cvt_f16_f32', 's_add_i32', 's_sub_i32'))
diffs = rust_before.diff(python_before, n_lanes)
if diffs:
trace_lines = []
for idx, (s, pc, d, rb, pb) in enumerate(trace):
trace_lines.append(f" step {s}: PC={pc:3d} {d}")
if idx < len(trace) - 1:
next_rb, next_pb = trace[idx + 1][3:5]
rust_diffs = rb.diff(next_rb, n_lanes, "->")
python_diffs = pb.diff(next_pb, n_lanes, "->")
if rust_diffs: trace_lines.append(f" rust: {', '.join(rust_diffs[:5])}")
if python_diffs: trace_lines.append(f" python: {', '.join(python_diffs[:5])}")
elif rust_diffs: trace_lines.append(f" python: (no changes)")
else:
# Last traced instruction - compare with current state
rust_diffs = rb.diff(rust_before, n_lanes, "->")
python_diffs = pb.diff(python_before, n_lanes, "->")
if rust_diffs: trace_lines.append(f" rust: {', '.join(rust_diffs[:5])}")
if python_diffs: trace_lines.append(f" python: {', '.join(python_diffs[:5])}")
elif rust_diffs: trace_lines.append(f" python: (no changes)")
trace_str = "\n".join(trace_lines)
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ (rust vs python):\n " + "\n ".join(diffs[:10]) + f"\n Recent instructions:\n{trace_str}", total_steps
rust_result = rust.step()
python_result = python.step()
if rust_result != python_result:
# Rust returns 1 for unsupported instructions - skip test
if rust_result == 1 and python_result == 0:
raise unittest.SkipTest(f"Rust emulator doesn't support instruction: {inst_str}")
trace_str = "\n".join(f" step {s}: PC={pc:3d} {d}" for s, pc, d, _, _ in trace)
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: rust={rust_result}, python={python_result}, inst={inst_str}\n Recent instructions:\n{trace_str}", total_steps
# Sync Python state to Rust after instructions with known Rust emulator differences
if sync_after:
rust_after = rust.get_snapshot()
for i in range(128): python.set_sgpr(i, rust_after.sgpr[i])
for lane in range(n_lanes):
for i in range(256): python.set_vgpr(lane, i, rust_after.vgpr[lane][i])
assert python.state is not None
python.state.pc, python.state.scc, python.state.vcc, python.state.exec_mask = rust_after.pc, rust_after.scc, rust_after.vcc, rust_after.exec_mask
if rust_result == -1:
total_steps += step + 1
break
if rust_result == 1:
total_steps += step + 1
break
if rust_result < 0 and rust_result != -2:
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: error code {rust_result}", total_steps
step += 1
else:
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Max steps ({max_steps}) reached", total_steps
finally:
rust.free()
return True, f"Completed {gx*gy*gz} workgroups", total_steps
def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int, int], max_steps: int = 1000,
debug: bool = False, trace_len: int = 10, buf_data: dict[int, bytes] | None = None) -> tuple[bool, str]:
"""Run all kernels through both emulators with shared buffer pool."""
if buf_data is None: buf_data = {}
# Allocate shared buffer pool with padding for over-reads (GPU loads up to 16 bytes at once)
buf_id_to_ptr: dict[int, int] = {}
buffers = []
for buf_id, size in buf_pool.items():
padded_size = ((size + 15) // 16) * 16 + 16 # round up to 16 bytes + extra padding
# Initialize with data from COPY if available
init_data = buf_data.get(buf_id, b'\x00' * padded_size)
init_list = list(init_data) + [0] * (padded_size - len(init_data))
buf = (ctypes.c_uint8 * padded_size)(*init_list[:padded_size])
buffers.append((buf, padded_size))
buf_id_to_ptr[buf_id] = ctypes.addressof(buf)
# Set up valid memory ranges
ranges = {(ctypes.addressof(b), size) for b, size in buffers}
total_steps = 0
for ki, kernel in enumerate(kernels):
# Create args array for this kernel's buffers
args = (ctypes.c_uint64 * len(kernel.buf_idxs))(*[buf_id_to_ptr[bid] for bid in kernel.buf_idxs])
args_ptr = ctypes.addressof(args)
# Update valid ranges to include this args array
kernel_ranges = ranges | {(args_ptr, ctypes.sizeof(args))}
set_valid_mem_ranges(kernel_ranges)
program = decode_program(kernel.code)
n_lanes = kernel.local_size[0] * kernel.local_size[1] * kernel.local_size[2]
ok, msg, steps = run_single_kernel(
kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
program, max_steps, debug, trace_len, ki
)
total_steps += steps
if not ok:
return False, msg
return True, f"Completed {len(kernels)} kernels, {total_steps} total steps"
def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list, max_steps: int = 1000, debug: bool = False,
global_size: tuple[int, int, int] = (1, 1, 1), trace_len: int = 10) -> tuple[bool, str]:
"""Run both emulators with memory set up for tinygrad kernels, executing all workgroups. Legacy wrapper."""
# Allocate buffers
buffers = []
for size in buf_sizes:
buf = (ctypes.c_uint8 * size)(*[0] * size)
buffers.append(buf)
# Create args array with buffer pointers
args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
args_ptr = ctypes.addressof(args)
# Set up valid memory ranges for Python emulator
ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
ranges.add((args_ptr, ctypes.sizeof(args)))
set_valid_mem_ranges(ranges)
program = decode_program(kernel)
ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, program, max_steps, debug, trace_len)
return ok, msg
def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int], dict[int, bytes]]:
"""Compile a tinygrad operation and extract all kernels with their buffer mappings."""
from tinygrad import Tensor
from tinygrad.runtime.support.elf import elf_loader
out = op_fn(Tensor)
sched = out.schedule()
kernels = []
buf_pool: dict[int, int] = {} # buffer id -> size
buf_data: dict[int, bytes] = {} # buffer id -> initial data from COPY
for ei in sched:
lowered = ei.lower()
if ei.ast.op.name == 'COPY':
# Handle COPY: extract source data to initialize destination buffer
if len(lowered.bufs) >= 2:
dst_buf, src_buf = lowered.bufs[0], lowered.bufs[1]
dst_id = id(dst_buf)
if dst_id not in buf_pool:
buf_pool[dst_id] = dst_buf.nbytes
# Get source data if it's from numpy/CPU
if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'):
src_data = bytes(src_buf.base._buf)
buf_data[dst_id] = src_data
elif ei.ast.op.name == 'SINK':
if lowered.prg and lowered.prg.p.lib:
lib = bytes(lowered.prg.p.lib)
_, sections, _ = elf_loader(lib)
for sec in sections:
if sec.name == '.text':
buf_idxs = []
buf_sizes = []
for b in lowered.bufs:
buf_id = id(b)
if buf_id not in buf_pool:
buf_pool[buf_id] = b.nbytes
buf_idxs.append(buf_id)
buf_sizes.append(b.nbytes)
kernels.append(KernelInfo(
code=bytes(sec.content),
global_size=tuple(lowered.prg.p.global_size),
local_size=tuple(lowered.prg.p.local_size),
buf_idxs=buf_idxs,
buf_sizes=buf_sizes
))
if not kernels: raise RuntimeError("No kernel found")
return kernels, buf_pool, buf_data
def get_kernel_from_tinygrad(op_fn) -> tuple[bytes, tuple[int, int, int], tuple[int, int, int], list]:
"""Compile a tinygrad operation and extract the last (main) kernel binary. Legacy wrapper."""
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
k = kernels[-1]
return k.code, k.global_size, k.local_size, k.buf_sizes
class TestTinygradKernels(unittest.TestCase):
"""Compare emulators on real tinygrad-compiled kernels."""
def _test_kernel(self, op_fn, max_steps=10000):
kernels, buf_pool, buf_data = get_kernels_from_tinygrad(op_fn)
ok, msg = compare_emulators_multi_kernel(kernels, buf_pool, max_steps=max_steps, buf_data=buf_data)
self.assertTrue(ok, msg)
# Basic ops - consolidated tests covering key instruction patterns
def test_unary_ops(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu().exp().log().sqrt().reciprocal())
def test_binary_ops(self): self._test_kernel(lambda T: (T([1.0, 2.0]) + T([3.0, 4.0])) * T([0.5, 0.5]) - T([1.0, 1.0]))
def test_trig(self): self._test_kernel(lambda T: T([0.1, 1.0, 3.14, -1.0]*8).sin() + T([0.1, 1.0, 3.14, -1.0]*8).cos())
def test_compare(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
def test_bitwise(self): self._test_kernel(lambda T: (T([0xF0, 0x0F, 0xFF]*11).int() & T([0x0F, 0x0F, 0x00]*11).int()) | T([1]*33).int())
def test_int_ops(self): self._test_kernel(lambda T: ((T.empty(64).int() + T.empty(64).int()) * T.empty(64).int()).float())
# Reductions
def test_reduce(self): self._test_kernel(lambda T: T.empty(64).sum() + T.empty(64).max())
def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
# Matmul
def test_gemm(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=100000)
@unittest.skip("Rust emulator crashes on this kernel (assertion failure in thread.rs)")
def test_gemm_fp16(self): self._test_kernel(lambda T: T.empty(16, 16).half() @ T.empty(16, 16).half(), max_steps=100000)
# Complex ops
def test_softmax(self): self._test_kernel(lambda T: T.empty(16).softmax())
def test_layernorm(self): self._test_kernel(lambda T: T.empty(8, 8).layernorm())
# Memory patterns
def test_memory(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous() + T.empty(4, 1).expand(4, 4))
# Cast ops
def test_cast(self): self._test_kernel(lambda T: T.empty(32).half().float() + T.empty(32).int().float())
# Pooling - regression for VCC wave32 mode
def test_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4)) + T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4)))
# Convolution
def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 2, 8, 8).conv2d(T.empty(2, 2, 3, 3)), max_steps=50000)
# Regression tests
def test_topk(self): self._test_kernel(lambda T: T.empty(64).topk(3)[0])
def test_interpolate(self): self._test_kernel(lambda T: T.empty(1,2,16,16).relu().cast('uint8').interpolate((8,8), mode="linear"))
def test_index_int64(self):
from tinygrad import dtypes
self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
def test_gelu(self): self._test_kernel(lambda T: T.empty(32, 32).gelu())
def test_cross_entropy(self):
import numpy as np
np.random.seed(0)
classes = np.random.randint(0, 10, (16,), dtype=np.int32).tolist()
x_np = np.random.randn(16, 10).astype(np.float32)
self._test_kernel(lambda T: (T(x_np.tolist()).reshape(16,10) + 0).cross_entropy((T(classes).int().reshape(16) + 0)))
def test_isinf(self): self._test_kernel(lambda T: T([float('-inf'), 0., float('inf'), 1.1]*8).isinf())
def test_sin_f64(self):
from tinygrad import dtypes
self._test_kernel(lambda T: T([2.0], dtype=dtypes.float64).sin())
if __name__ == "__main__":
unittest.main()
-407
View File
@@ -1,407 +0,0 @@
#!/usr/bin/env python3
"""Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
import unittest
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import encode_src, RawImm
from extra.assembly.amd.asm import detect_format
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]))
class TestDS(unittest.TestCase):
"""Test DS (data share / LDS) instructions."""
def test_ds_store_b32(self):
# ds_store_b32 v0, v1
# GFX11: encoding: [0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
inst = ds_store_b32(addr=v[0], data0=v[1])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
def test_ds_load_b32(self):
# ds_load_b32 v0, v1
# GFX11: encoding: [0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]
inst = ds_load_b32(vdst=v[0], addr=v[1])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]))
def test_ds_store_b32_offset(self):
# ds_store_b32 v0, v1 offset:64
# GFX11: encoding: [0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
inst = ds_store_b32(addr=v[0], data0=v[1], offset0=64)
self.assertEqual(inst.to_bytes(), bytes([0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
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])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]))
def test_ds_add_u32(self):
# ds_add_u32 v0, v1
# GFX11: encoding: [0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]
inst = ds_add_u32(addr=v[0], data0=v[1])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]))
def test_ds_store_b32_gds(self):
# ds_store_b32 v0, v1 gds
# GFX11: encoding: [0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]
inst = ds_store_b32(addr=v[0], data0=v[1], gds=1)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]))
class TestVOP3(unittest.TestCase):
"""Test VOP3 (3-operand vector) instructions."""
def test_v_fma_f32(self):
# v_fma_f32 v0, v1, v2, v3
# GFX11: encoding: [0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]
inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]))
def test_v_mad_f32(self):
# v_fmac_f32_e64 v0, v1, v2 (fmac is fma with implicit dst as src2)
# Use v_fma_f32 with vdst == src2
inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[0])
self.assertEqual(inst.to_bytes()[:4], bytes([0x00,0x00,0x13,0xd6]))
def test_v_add3_u32(self):
# v_add3_u32 v0, v1, v2, v3
# GFX11: encoding: [0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]
inst = v_add3_u32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]))
class TestFLAT(unittest.TestCase):
"""Test FLAT/GLOBAL/SCRATCH memory instructions."""
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)
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)
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])
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)
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)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]))
class TestSMEM(unittest.TestCase):
"""Test SMEM (scalar memory) instructions - regression tests for glc/dlc bit positions."""
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)
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)
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)
self.assertEqual(inst.to_bytes(), bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]))
def test_smem_disasm_roundtrip_dlc(self):
# Test that disassembly/reassembly preserves DLC bit correctly
data = bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00])
decoded = SMEM.from_bytes(data)
self.assertEqual(decoded.to_bytes(), data)
def test_smem_disasm_roundtrip_glc_dlc(self):
# Test that disassembly/reassembly preserves GLC+DLC bits correctly
data = bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00])
decoded = SMEM.from_bytes(data)
self.assertEqual(decoded.to_bytes(), data)
class TestVOP3Literal(unittest.TestCase):
"""Test VOP3 literal handling - regression tests for Inst64 literal encoding."""
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])
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)
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])
self.assertEqual(len(inst.to_bytes()), 12) # 8 bytes + 4 byte literal
class TestDetectFormat(unittest.TestCase):
"""Test detect_format uses encoding from autogen classes."""
def test_detect_sopp(self):
self.assertEqual(detect_format(s_endpgm().to_bytes()), SOPP)
self.assertEqual(detect_format(s_nop(0).to_bytes()), SOPP)
self.assertEqual(detect_format(s_barrier().to_bytes()), SOPP)
def test_detect_sop1(self):
self.assertEqual(detect_format(s_mov_b32(s[0], 0).to_bytes()), SOP1)
self.assertEqual(detect_format(s_mov_b64(s[0:1], 0).to_bytes()), SOP1)
def test_detect_sop2(self):
self.assertEqual(detect_format(s_add_u32(s[0], s[1], s[2]).to_bytes()), SOP2)
self.assertEqual(detect_format(s_mul_i32(s[0], s[1], s[2]).to_bytes()), SOP2)
def test_detect_sopc(self):
self.assertEqual(detect_format(s_cmp_eq_i32(s[0], s[1]).to_bytes()), SOPC)
def test_detect_sopk(self):
self.assertEqual(detect_format(s_movk_i32(s[0], 0x1234).to_bytes()), SOPK)
def test_detect_vop1(self):
self.assertEqual(detect_format(v_mov_b32_e32(v[0], 0).to_bytes()), VOP1)
self.assertEqual(detect_format(v_rcp_f32_e32(v[0], v[1]).to_bytes()), VOP1)
def test_detect_vop2(self):
self.assertEqual(detect_format(v_add_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
self.assertEqual(detect_format(v_mul_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
def test_detect_vopc(self):
self.assertEqual(detect_format(v_cmp_eq_f32_e32(v[0], v[1]).to_bytes()), VOPC)
self.assertEqual(detect_format(v_cmp_lt_i32_e32(v[0], v[1]).to_bytes()), VOPC)
def test_detect_vop3(self):
self.assertEqual(detect_format(v_add_f32_e64(v[0], v[1], v[2]).to_bytes()), VOP3)
self.assertEqual(detect_format(v_fma_f32(v[0], v[1], v[2], v[3]).to_bytes()), VOP3)
def test_detect_vop3p(self):
self.assertEqual(detect_format(VOP3P(VOP3POp.V_PK_ADD_F16, v[0], v[1], v[2], v[3]).to_bytes()), VOP3P)
def test_detect_smem(self):
self.assertEqual(detect_format(s_load_b32(s[0], s[2:3], 0).to_bytes()), SMEM)
self.assertEqual(detect_format(s_load_b64(s[0:1], s[2:3], s[5]).to_bytes()), SMEM)
def test_detect_ds(self):
self.assertEqual(detect_format(ds_load_b32(v[0], v[1]).to_bytes()), DS)
self.assertEqual(detect_format(ds_store_b32(v[0], v[1]).to_bytes()), DS)
def test_detect_flat(self):
self.assertEqual(detect_format(global_load_b32(v[0], v[1:3], RawImm(124)).to_bytes()), FLAT)
self.assertEqual(detect_format(global_store_b32(v[0:2], v[2], RawImm(124)).to_bytes()), FLAT)
def test_detect_mubuf(self):
self.assertEqual(detect_format(buffer_load_b32(v[0], v[1], s[0:4], s[5]).to_bytes()), MUBUF)
def test_detect_mtbuf(self):
self.assertEqual(detect_format(tbuffer_load_format_x(v[0], v[1], s[0:4], s[5], format=22).to_bytes()), MTBUF)
def test_detect_mimg(self):
self.assertEqual(detect_format(image_load(v[0:4], v[4:6], s[0:8], dmask=0xf, dim=1).to_bytes()), MIMG)
def test_detect_exp(self):
self.assertEqual(detect_format(EXP(en=0xf, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[2], vsrc3=v[3]).to_bytes()), EXP)
def test_detect_vopd(self):
inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=0, srcy0=0)
self.assertEqual(detect_format(inst.to_bytes()), VOPD)
def test_detect_vinterp(self):
inst = VINTERP(VINTERPOp.V_INTERP_P10_F32, vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
self.assertEqual(detect_format(inst.to_bytes()), VINTERP)
if __name__ == "__main__":
unittest.main()
-181
View File
@@ -1,181 +0,0 @@
# do not change these tests. we need to fix bugs to make them pass
# 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.ins import *
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.asm import asm
from extra.assembly.amd.test.test_roundtrip import compile_asm
class TestIntegration(unittest.TestCase):
inst: Inst
def tearDown(self):
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}"
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[189:192], v[140:143], v[0:7])
def test_load_b128(self):
self.inst = s_load_b128(s[4:7], s[0:1], NULL, 0)
def test_load_b128_wrong_size(self):
# this should have to be 4 regs on the loaded to
with self.assertRaises(Exception):
self.inst = s_load_b128(s[4:6], s[0:1], NULL, 0)
def test_mov_b32(self):
self.inst = s_mov_b32(s[80], s[0])
def test_mov_b64(self):
self.inst = s_mov_b64(s[80:81], s[0:1])
def test_mov_b32_wrong(self):
with self.assertRaises(Exception):
self.inst = s_mov_b32(s[80:81], s[0:1])
with self.assertRaises(Exception):
self.inst = s_mov_b32(s[80:81], s[0])
with self.assertRaises(Exception):
self.inst = s_mov_b32(s[80], s[0:1])
def test_mov_b64_wrong(self):
with self.assertRaises(Exception):
self.inst = s_mov_b64(s[80], s[0])
with self.assertRaises(Exception):
self.inst = s_mov_b64(s[80], s[0:1])
with self.assertRaises(Exception):
self.inst = s_mov_b64(s[80:81], s[0])
def test_load_b128_no_0(self):
self.inst = s_load_b128(s[4:7], s[0:1], NULL)
def test_load_b128_s(self):
self.inst = s_load_b128(s[4:7], s[0:1], s[8], 0)
def test_load_b128_v(self):
with self.assertRaises(TypeError):
self.inst = s_load_b128(s[4:7], s[0:1], v[8], 0)
def test_load_b128_off(self):
self.inst = s_load_b128(s[4:7], s[0:1], NULL, 3)
def test_simple_stos(self):
self.inst = s_mov_b32(s[0], s[1])
def test_simple_wrong(self):
with self.assertRaises(TypeError):
self.inst = s_mov_b32(v[0], s[1])
def test_simple_vtov(self):
self.inst = v_mov_b32_e32(v[0], v[1])
def test_simple_stov(self):
self.inst = v_mov_b32_e32(v[0], s[2])
def test_simple_float_to_v(self):
self.inst = v_mov_b32_e32(v[0], 1.0)
def test_simple_v_to_float(self):
with self.assertRaises(TypeError):
self.inst = v_mov_b32_e32(1, v[0])
def test_simple_int_to_v(self):
self.inst = v_mov_b32_e32(v[0], 1)
def test_three_add(self):
self.inst = v_add_co_ci_u32_e32(v[3], s[7], v[3])
def test_three_add_v(self):
self.inst = v_add_co_ci_u32_e32(v[3], v[7], v[3])
def test_three_add_const(self):
self.inst = v_add_co_ci_u32_e32(v[3], 2.0, v[3])
def test_swaitcnt_lgkm(self): self.inst = s_waitcnt(0xfc07)
def test_swaitcnt_vm(self): self.inst = s_waitcnt(0x03f7)
def test_vmad(self):
self.inst = v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2])
def test_large_imm(self):
self.inst = v_mov_b32_e32(v[0], 0x1234)
def test_dual_mov(self):
self.inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], srcy0=v[4])
def test_dual_mul(self):
self.inst = v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
def test_simple_int_to_s(self):
self.inst = s_mov_b32(s[0], 3)
def test_complex_int_to_s(self):
self.inst = s_mov_b32(s[0], 0x235646)
def test_simple_float_to_s(self):
self.inst = s_mov_b32(s[0], 1.0)
def test_complex_float_to_s(self):
self.inst = s_mov_b32(s[0], 1337.0)
int_inst = s_mov_b32(s[0], struct.unpack("I", struct.pack("f", 1337.0))[0])
self.assertEqual(self.inst, int_inst)
class TestRegisterSliceSyntax(unittest.TestCase):
"""
Issue: Register slice syntax should use AMD assembly convention (inclusive end).
In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
The DSL should match this convention so that:
- s[4:7] gives 4 registers
- Disassembler output can be copied directly back into DSL code
Fix: Change _RegFactory.__getitem__ to use inclusive end:
key.stop - key.start + 1 (instead of key.stop - key.start)
"""
def test_register_slice_count(self):
# s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
reg = s[4:7]
self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
def test_register_slice_roundtrip(self):
# Round-trip: DSL -> disasm -> DSL should preserve register count
reg = s[4:7] # 4 registers in AMD convention
inst = s_load_b128(reg, s[0:1], NULL, 0)
disasm = inst.disasm()
# Disasm shows s[4:7] - user should be able to copy this back
self.assertIn("s[4:7]", disasm)
# And s[4:7] in DSL should give the same 4 registers
reg_from_disasm = s[4:7]
self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
class TestInstructionEquality(unittest.TestCase):
"""
Issue: No __eq__ method - instruction comparison requires repr() workaround.
Two identical instructions should compare equal with ==, but currently:
inst1 == inst2 returns False
The test_handwritten.py works around this with:
self.assertEqual(repr(self.inst), repr(reasm))
"""
def test_identical_instructions_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[1])
self.assertEqual(inst1, inst2, "identical instructions should be equal")
def test_different_instructions_not_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[2])
self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
if __name__ == "__main__":
unittest.main()
-330
View File
@@ -1,330 +0,0 @@
#!/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.ins import *
from extra.assembly.amd.asm import waitcnt, asm
from extra.assembly.amd.test.helpers import get_llvm_mc
def disassemble(lib: bytes, arch: str = "gfx1100") -> str:
"""Disassemble ELF binary using tinygrad's compiler, return raw output."""
from tinygrad.runtime.support.compiler_amd import HIPCompiler
old_stdout = sys.stdout
sys.stdout = io.StringIO()
HIPCompiler(arch).disassemble(lib)
output = sys.stdout.getvalue()
sys.stdout = old_stdout
return output
def parse_disassembly(raw: str) -> list[str]:
"""Parse disassembly output to list of instruction mnemonics."""
lines = []
for line in raw.splitlines():
if line.startswith('\t'):
instr = line.split('//')[0].strip()
if instr: lines.append(instr)
return lines
def assemble_and_disassemble(instructions: list, arch: str = "gfx1100") -> list[str]:
"""Assemble instructions with our DSL, then disassemble with AMD toolchain."""
from tinygrad.runtime.support.compiler_amd import HIPCompiler
# Generate bytes from our DSL
code_bytes = b''.join(inst.to_bytes() for inst in instructions)
# Wrap in minimal ELF-compatible assembly with .byte directives
byte_str = ', '.join(f'0x{b:02x}' for b in code_bytes)
asm_src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n.byte {byte_str}\n"
# Assemble with AMD COMGR and disassemble
lib = HIPCompiler(arch).compile(asm_src)
return parse_disassembly(disassemble(lib, arch))
class TestIntegration(unittest.TestCase):
"""Test our assembler output matches LLVM disassembly."""
def test_simple_sop1(self):
"""Test SOP1 instructions round-trip."""
instructions = [
s_mov_b32(s[0], s[1]),
s_mov_b32(s[2], 0),
s_not_b32(s[3], s[4]),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_mov_b32', disasm[0])
self.assertIn('s_mov_b32', disasm[1])
self.assertIn('s_not_b32', disasm[2])
def test_simple_sop2(self):
"""Test SOP2 instructions round-trip."""
instructions = [
s_add_u32(s[0], s[1], s[2]),
s_sub_u32(s[3], s[4], 10),
s_and_b32(s[5], s[6], s[7]),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_add_u32', disasm[0])
self.assertIn('s_sub_u32', disasm[1])
self.assertIn('s_and_b32', disasm[2])
def test_simple_vop2(self):
"""Test VOP2 instructions round-trip."""
instructions = [
v_add_f32_e32(v[0], v[1], v[2]),
v_mul_f32_e32(v[3], 1.0, v[4]), # 1.0 is inline constant
v_and_b32_e32(v[5], 10, v[6]), # small inline constant
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('v_add_f32', disasm[0])
self.assertIn('v_mul_f32', disasm[1])
def test_control_flow(self):
"""Test control flow instructions."""
instructions = [
s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
s_endpgm(),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_waitcnt', disasm[0])
self.assertIn('s_endpgm', disasm[1])
def test_memory_ops(self):
"""Test memory instructions."""
instructions = [
s_load_b32(s[0], s[0:2], NULL),
s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
s_endpgm(),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_load_b32', disasm[0])
self.assertIn('s_waitcnt', disasm[1])
self.assertIn('global_store_b32', disasm[2])
def test_full_kernel(self):
"""Test a complete kernel similar to tinygrad output."""
# Simple kernel: load value, add 1, store back
instructions = [
# Get thread ID
v_mov_b32_e32(v[0], s[0]), # base addr low
v_mov_b32_e32(v[1], s[1]), # base addr high
# Load value
global_load_b32(vdst=v[2], addr=v[0:2], 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),
s_endpgm(),
]
disasm = assemble_and_disassemble(instructions)
# Verify key instructions are present
self.assertTrue(any('global_load' in d for d in disasm))
self.assertTrue(any('v_add_f32' in d for d in disasm))
self.assertTrue(any('global_store' in d for d in disasm))
self.assertTrue(any('s_endpgm' in d for d in disasm))
def test_bytes_roundtrip(self):
"""Test that our bytes match what AMD assembler produces."""
from tinygrad.runtime.support.compiler_amd import HIPCompiler
# Simple instruction
inst = s_mov_b32(s[0], s[1])
our_bytes = inst.to_bytes()
# Assemble same instruction with AMD toolchain
asm_src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\ns_mov_b32 s0, s1\n"
compiler = HIPCompiler("gfx1100")
lib = compiler.compile(asm_src)
raw = disassemble(lib)
for line in raw.splitlines():
if 's_mov_b32' in line and '//' in line:
# Extract hex bytes from comment: "// 000000001300: BE800001"
comment = line.split('//')[1].strip()
hex_str = comment.split(':')[1].strip()
# Convert big-endian hex string to little-endian bytes
amd_bytes = bytes.fromhex(hex_str)[::-1] # reverse for little-endian
self.assertEqual(our_bytes, amd_bytes, f"Bytes mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
return
self.fail("Could not find s_mov_b32 in disassembly")
class 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."""
def test_simple_add_kernel(self):
"""Generate a simple add kernel from tinygrad and verify disassembly."""
from tinygrad import Tensor
from tinygrad.codegen import get_program
from tinygrad.renderer.cstyle import AMDHIPRenderer
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.uop.ops import Ops
# Create a computation that generates a real kernel
a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
b = Tensor([5.0, 6.0, 7.0, 8.0]).realize()
c = a + b
# Get schedule and find SINK
schedule = c.schedule()
sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
self.assertTrue(len(sink_items) > 0, "No SINK in schedule")
# Generate program
renderer = AMDHIPRenderer('gfx1100')
prg = get_program(sink_items[0].ast, renderer)
self.assertIsNotNone(prg.src)
# Compile and disassemble
compiler = HIPCompiler('gfx1100')
lib = compiler.compile(prg.src)
raw_disasm = disassemble(lib)
instrs = parse_disassembly(raw_disasm)
# Verify we got some instructions
self.assertTrue(len(instrs) > 0, "No instructions in disassembly")
# Should have an endpgm
self.assertTrue(any('s_endpgm' in i for i in instrs), "Missing s_endpgm")
def test_matmul_kernel(self):
"""Generate a matmul kernel and verify disassembly has expected patterns."""
from tinygrad import Tensor
from tinygrad.codegen import get_program
from tinygrad.renderer.cstyle import AMDHIPRenderer
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.uop.ops import Ops
# Create a small matmul
a = Tensor.rand(4, 4).realize()
b = Tensor.rand(4, 4).realize()
c = a @ b
# Get schedule
schedule = c.schedule()
sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
self.assertTrue(len(sink_items) > 0)
# Generate and compile
renderer = AMDHIPRenderer('gfx1100')
prg = get_program(sink_items[0].ast, renderer)
compiler = HIPCompiler('gfx1100')
lib = compiler.compile(prg.src)
raw_disasm = disassemble(lib)
instrs = parse_disassembly(raw_disasm)
# Matmul should have multiply and add instructions
has_mul = any('mul' in i.lower() for i in instrs)
has_add = any('add' in i.lower() for i in instrs)
self.assertTrue(has_mul or has_add, "Matmul should have mul/add ops")
def test_disasm_to_bytes_roundtrip(self):
"""Parse disassembled instructions and verify we can re-encode some of them."""
from tinygrad import Tensor
from tinygrad.codegen import get_program
from tinygrad.renderer.cstyle import AMDHIPRenderer
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.uop.ops import Ops
# Simple kernel
a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
b = (a * 2.0)
schedule = b.schedule()
sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
if not sink_items: return # skip if no kernel
renderer = AMDHIPRenderer('gfx1100')
prg = get_program(sink_items[0].ast, renderer)
compiler = HIPCompiler('gfx1100')
lib = compiler.compile(prg.src)
raw_disasm = disassemble(lib)
# Find s_endpgm and verify we can encode it
for line in raw_disasm.splitlines():
if 's_endpgm' in line and '//' in line:
# Extract bytes from comment
comment = line.split('//')[1].strip()
hex_str = comment.split(':')[1].strip()
amd_bytes = bytes.fromhex(hex_str)[::-1]
# Our encoding
our_inst = s_endpgm()
our_bytes = our_inst.to_bytes()
self.assertEqual(our_bytes, amd_bytes, f"s_endpgm mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
return
if __name__ == "__main__":
unittest.main()
-121
View File
@@ -1,121 +0,0 @@
#!/usr/bin/env python3
"""Test AMD assembler/disassembler against LLVM test vectors."""
import unittest, re, subprocess, functools
from tinygrad.helpers import fetch
from extra.assembly.amd.asm import asm, disasm, detect_format
from extra.assembly.amd.test.helpers import get_llvm_mc
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-21.1.0/llvm/test/MC/AMDGPU"
RDNA_FILES = ['gfx11_asm_sop1.s', 'gfx11_asm_sop2.s', 'gfx11_asm_sopp.s', 'gfx11_asm_sopk.s', 'gfx11_asm_sopc.s',
'gfx11_asm_vop1.s', 'gfx11_asm_vop2.s', 'gfx11_asm_vopc.s', 'gfx11_asm_vop3.s', 'gfx11_asm_vop3p.s', 'gfx11_asm_vinterp.s',
'gfx11_asm_vopd.s', 'gfx11_asm_vopcx.s', 'gfx11_asm_vop3_from_vop1.s', 'gfx11_asm_vop3_from_vop2.s', 'gfx11_asm_vop3_from_vopc.s',
'gfx11_asm_vop3_from_vopcx.s', 'gfx11_asm_ds.s', 'gfx11_asm_smem.s', 'gfx11_asm_flat.s', 'gfx11_asm_mubuf.s', 'gfx11_asm_mtbuf.s',
'gfx11_asm_mimg.s', 'gfx11_asm_wmma.s', 'gfx11_asm_vop3_features.s', 'gfx11_asm_vop3p_features.s', 'gfx11_asm_vopd_features.s',
'gfx11_asm_vop3_alias.s', 'gfx11_asm_vop3p_alias.s', 'gfx11_asm_vopc_alias.s', 'gfx11_asm_vopcx_alias.s', 'gfx11_asm_vinterp_alias.s',
'gfx11_asm_smem_alias.s', 'gfx11_asm_mubuf_alias.s', 'gfx11_asm_mtbuf_alias.s']
# CDNA test files - includes gfx9 files for shared instructions, plus gfx90a/gfx942 specific files
# gfx90a_ldst_acc.s has MIMG mixed in, filtered via is_mimg check
CDNA_FILES = ['gfx9_asm_sop1.s', 'gfx9_asm_sop2.s', 'gfx9_asm_sopp.s', 'gfx9_asm_sopk.s', 'gfx9_asm_sopc.s',
'gfx9_asm_vop1.s', 'gfx9_asm_vop2.s', 'gfx9_asm_vopc.s', 'gfx9_asm_vop3.s', 'gfx9_asm_vop3p.s',
'gfx9_asm_ds.s', 'gfx9_asm_flat.s', 'gfx9_asm_smem.s', 'gfx9_asm_mubuf.s', 'gfx9_asm_mtbuf.s',
'gfx90a_ldst_acc.s', 'gfx90a_asm_features.s', 'flat-scratch-gfx942.s', 'gfx942_asm_features.s',
'mai-gfx90a.s', 'mai-gfx942.s']
# RDNA4 (gfx12) test files - excludes alias/err/fake16/dpp files, and vimage/vsample (not supported)
# NOTE: vflat/vdsdir excluded - not implemented; features.s has mixed formats
RDNA4_FILES = ['gfx12_asm_sop1.s', 'gfx12_asm_sop2.s', 'gfx12_asm_sopp.s', 'gfx12_asm_sopk.s', 'gfx12_asm_sopc.s',
'gfx12_asm_vop1.s', 'gfx12_asm_vop2.s', 'gfx12_asm_vopc.s', 'gfx12_asm_vopcx.s', 'gfx12_asm_vop3.s', 'gfx12_asm_vop3c.s',
'gfx12_asm_vop3cx.s', 'gfx12_asm_vop3p.s', 'gfx12_asm_vop3_from_vop1.s', 'gfx12_asm_vop3_from_vop2.s',
'gfx12_asm_vop3p_features.s', 'gfx12_asm_vopd.s', 'gfx12_asm_vopd_features.s',
'gfx12_asm_ds.s', 'gfx12_asm_smem.s',
'gfx12_asm_vbuffer_mubuf.s', 'gfx12_asm_vbuffer_mtbuf.s', 'gfx12_asm_wmma_w32.s', 'gfx12_asm_exp.s']
def _is_mimg(data: bytes) -> bool: return (int.from_bytes(data[:4], 'little') >> 26) & 0x3f == 0b111100
def _parse_llvm_tests(text: str, pattern: str) -> list[tuple[str, bytes]]:
tests = []
for block in text.split('\n\n'):
asm_text, encoding = None, None
for line in block.split('\n'):
line = line.strip()
if not line or line.startswith(('.', ';')): continue
if not line.startswith('//'):
asm_text = line.split('//')[0].strip() or asm_text
if m := re.search(pattern + r'[^:]*:.*?(?:encoding:\s*)?\[(0x[0-9a-f,x\s]+)\]', line, re.I):
encoding = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
if asm_text and encoding:
try: tests.append((asm_text, bytes.fromhex(encoding)))
except ValueError: pass
return tests
@functools.cache
def _get_tests(f: str, arch: str) -> list[tuple[str, bytes]]:
text = fetch(f"{LLVM_BASE}/{f}").read_bytes().decode('utf-8', errors='ignore')
if arch == "rdna3":
tests = _parse_llvm_tests(text, r'(?:GFX11|W32|W64)')
elif arch == "rdna4":
# Match GFX12 but not GFX1250 (which has different lit64 encoding)
tests = _parse_llvm_tests(text, r'(?:GFX12(?!50)|W32|W64)')
elif 'gfx90a' in f or 'gfx942' in f:
tests = _parse_llvm_tests(text, r'(?:GFX90A|GFX942)')
else:
tests = _parse_llvm_tests(text, r'(?:VI9|GFX9|CHECK)')
return [(a, d) for a, d in tests if not _is_mimg(d)] if arch == "cdna" else tests
def _compile_asm_batch(instrs: list[str], arch: str = "rdna3") -> list[bytes]:
if not instrs: return []
mcpu = {'rdna3': 'gfx1100', 'rdna4': 'gfx1200'}.get(arch, 'gfx1100')
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True, timeout=30)
if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
return [bytes.fromhex(line.split('encoding:')[1].strip()[1:-1].replace('0x', '').replace(',', '').replace(' ', ''))
for line in result.stdout.split('\n') if 'encoding:' in line]
def _make_test(f: str, arch: str, test_type: str):
def test(self):
tests = _get_tests(f, arch)
name = f"{arch}_{test_type}_{f}"
if test_type == "roundtrip":
for _, data in tests:
decoded = detect_format(data, arch).from_bytes(data)
self.assertEqual(decoded.to_bytes()[:len(data)], data)
print(f"{name}: {len(tests)} passed")
elif test_type == "asm":
passed, skipped = 0, 0
for asm_text, expected in tests:
try:
self.assertEqual(asm(asm_text, arch).to_bytes(), expected)
passed += 1
except: skipped += 1
print(f"{name}: {passed} passed, {skipped} skipped")
elif test_type == "disasm":
to_test = []
for _, data in tests:
try:
decoded = detect_format(data, arch).from_bytes(data)
# Skip if roundtrip fails, disasm fails, or op_name is missing (disasm starts with space)
if decoded.to_bytes()[:len(data)] == data and (d := disasm(decoded)) and not d.startswith(' '): to_test.append((data, d))
except: pass
skipped = len(tests) - len(to_test)
print(f"{name}: {len(to_test)} passed, {skipped} skipped")
if arch in ("rdna3", "rdna4"):
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
for (data, _), llvm in zip(to_test, _compile_asm_batch([t[1] for t in to_test], arch)): self.assertEqual(llvm, data)
return test
class TestLLVM(unittest.TestCase): pass
for f in RDNA_FILES:
setattr(TestLLVM, f"test_rdna3_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "roundtrip"))
setattr(TestLLVM, f"test_rdna3_asm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "asm"))
setattr(TestLLVM, f"test_rdna3_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "disasm"))
for f in CDNA_FILES:
setattr(TestLLVM, f"test_cdna_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "roundtrip"))
setattr(TestLLVM, f"test_cdna_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "disasm"))
for f in RDNA4_FILES:
setattr(TestLLVM, f"test_rdna4_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "roundtrip"))
setattr(TestLLVM, f"test_rdna4_asm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "asm"))
setattr(TestLLVM, f"test_rdna4_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "disasm"))
if __name__ == "__main__":
unittest.main()
@@ -1,55 +0,0 @@
#!/usr/bin/env python3
"""Test that invalid instructions raise exceptions through the mock GPU stack."""
import unittest, subprocess, os, time
class TestMockGPUInvalidInstruction(unittest.TestCase):
def test_unsupported_instruction_raises(self):
"""Test that unsupported instructions raise immediately through the full MOCKGPU stack."""
test_code = '''
import struct
from tinygrad import Device, Tensor
from tinygrad.engine.realize import get_runner
from tinygrad.runtime.ops_amd import AMDProgram
dev = Device["AMD"]
a = Tensor([1.0]).realize()
b = a + 1
si = b.schedule()[-1]
runner = get_runner(dev.device, si.ast)
prg = runner._prg
lib = bytearray(prg.lib)
# 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", 0x7E008400)
found = True
break
assert found, "s_endpgm not found"
patched_prg = AMDProgram(dev, "patched", bytes(lib))
b.uop.buffer.allocate()
patched_prg(b.uop.buffer._buf, a.uop.buffer._buf, global_size=(1,1,1), local_size=(1,1,1))
dev.synchronize()
'''
env = os.environ.copy()
env["AMD"] = "1"
env["MOCKGPU"] = "1"
env["PYTHON_REMU"] = "1"
env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
st = time.perf_counter()
result = subprocess.run(["python", "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
elapsed = time.perf_counter() - st
self.assertNotEqual(result.returncode, 0, "should have raised")
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")
if __name__ == "__main__":
unittest.main()
-403
View File
@@ -1,403 +0,0 @@
#!/usr/bin/env python3
"""Tests for the RDNA3 pseudocode DSL."""
import unittest
from extra.assembly.amd.pcode import (Reg, TypedView, TypedView, MASK32, MASK64,
_f32, _i32, _f16, _i16, f32_to_f16, isNAN, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
BYTE_PERMUTE, v_sad_u8, v_msad_u8, _compile_pseudocode, _expr, compile_pseudocode)
from extra.assembly.amd.test.helpers import ExecContext
from extra.assembly.amd.autogen.rdna3.str_pcode import VOP3SDOp_PCODE, VOPCOp_PCODE
from extra.assembly.amd.autogen.rdna3.enum import VOP3SDOp, VOPCOp
# Compile pseudocode functions on demand for regression tests
_VOP3SDOp_V_DIV_SCALE_F32 = compile_pseudocode('VOP3SDOp', 'V_DIV_SCALE_F32', VOP3SDOp_PCODE[VOP3SDOp.V_DIV_SCALE_F32])
_VOPCOp_V_CMP_CLASS_F32 = compile_pseudocode('VOPCOp', 'V_CMP_CLASS_F32', VOPCOp_PCODE[VOPCOp.V_CMP_CLASS_F32])
class TestReg(unittest.TestCase):
def test_u32_read(self):
r = Reg(0xDEADBEEF)
self.assertEqual(int(r.u32), 0xDEADBEEF)
def test_u32_write(self):
r = Reg(0)
r.u32 = 0x12345678
self.assertEqual(r._val, 0x12345678)
def test_f32_read(self):
r = Reg(0x40400000) # 3.0f
self.assertAlmostEqual(float(r.f32), 3.0)
def test_f32_write(self):
r = Reg(0)
r.f32 = 3.0
self.assertEqual(r._val, 0x40400000)
def test_i32_signed(self):
r = Reg(0xFFFFFFFF) # -1 as signed
self.assertEqual(int(r.i32), -1)
def test_u64(self):
r = Reg(0xDEADBEEFCAFEBABE)
self.assertEqual(int(r.u64), 0xDEADBEEFCAFEBABE)
def test_f64(self):
r = Reg(0x4008000000000000) # 3.0 as f64
self.assertAlmostEqual(float(r.f64), 3.0)
class TestTypedView(unittest.TestCase):
def test_bit_slice(self):
r = Reg(0xDEADBEEF)
# Slices return TypedView which supports .u32, .u16 etc (matching pseudocode like S1.u32[1:0].u32)
self.assertEqual(r.u32[7:0].u32, 0xEF)
self.assertEqual(r.u32[15:8].u32, 0xBE)
self.assertEqual(r.u32[23:16].u32, 0xAD)
self.assertEqual(r.u32[31:24].u32, 0xDE)
# Also works with int() for arithmetic
self.assertEqual(int(r.u32[7:0]), 0xEF)
def test_single_bit_read(self):
r = Reg(0b11010101)
self.assertEqual(r.u32[0], 1)
self.assertEqual(r.u32[1], 0)
self.assertEqual(r.u32[2], 1)
self.assertEqual(r.u32[3], 0)
def test_single_bit_write(self):
r = Reg(0)
r.u32[5] = 1
r.u32[3] = 1
self.assertEqual(r._val, 0b00101000)
def test_nested_bit_access(self):
# S0.u32[S1.u32[4:0]] - access bit at position from another register
s0 = Reg(0b11010101)
s1 = Reg(3)
bit_pos = s1.u32[4:0] # TypedView, int value = 3
bit_val = s0.u32[int(bit_pos)] # bit 3 of s0 = 0
self.assertEqual(int(bit_pos), 3)
self.assertEqual(bit_val, 0)
def test_arithmetic(self):
r1 = Reg(0x40400000) # 3.0f
r2 = Reg(0x40800000) # 4.0f
result = r1.f32 + r2.f32
self.assertAlmostEqual(result, 7.0)
def test_comparison(self):
r1 = Reg(5)
r2 = Reg(3)
self.assertTrue(r1.u32 > r2.u32)
self.assertFalse(r1.u32 < r2.u32)
self.assertTrue(r1.u32 != r2.u32)
class TestTypedView(unittest.TestCase):
def test_slice_read(self):
r = Reg(0x56781234)
self.assertEqual(r[15:0].u16, 0x1234)
self.assertEqual(r[31:16].u16, 0x5678)
def test_slice_write(self):
r = Reg(0)
r[15:0].u16 = 0x1234
r[31:16].u16 = 0x5678
self.assertEqual(r._val, 0x56781234)
def test_slice_f16(self):
r = Reg(0)
r[15:0].f16 = 3.0
self.assertAlmostEqual(_f16(r._val & 0xffff), 3.0, places=2)
class TestCompiler(unittest.TestCase):
def test_ternary(self):
result = _expr("a > b ? 1 : 0")
self.assertIn("if", result)
self.assertIn("else", result)
def test_type_prefix_strip(self):
self.assertEqual(_expr("1'0U"), "0")
self.assertEqual(_expr("32'1"), "1")
self.assertEqual(_expr("16'0xFFFF"), "0xFFFF")
def test_suffix_strip(self):
self.assertEqual(_expr("0ULL"), "0")
self.assertEqual(_expr("1LL"), "1")
self.assertEqual(_expr("5U"), "5")
self.assertEqual(_expr("3.14F"), "3.14")
def test_boolean_ops(self):
self.assertIn("and", _expr("a && b"))
self.assertIn("or", _expr("a || b"))
self.assertIn("!=", _expr("a <> b"))
def test_pack16(self):
result = _expr("{ a, b }")
self.assertIn("_pack", result)
def test_type_cast_strip(self):
self.assertEqual(_expr("64'U(x)"), "(x)")
self.assertEqual(_expr("32'I(y)"), "(y)")
class TestExecContext(unittest.TestCase):
def test_float_add(self):
ctx = ExecContext(s0=0x40400000, s1=0x40800000) # 3.0f, 4.0f
ctx.D0.f32 = ctx.S0.f32 + ctx.S1.f32
self.assertAlmostEqual(_f32(ctx.D0._val), 7.0)
def test_float_mul(self):
ctx = ExecContext(s0=0x40400000, s1=0x40800000) # 3.0f, 4.0f
ctx.run("D0.f32 = S0.f32 * S1.f32")
self.assertAlmostEqual(_f32(ctx.D0._val), 12.0)
def test_scc_comparison(self):
ctx = ExecContext(s0=42, s1=42)
ctx.run("SCC = S0.u32 == S1.u32")
self.assertEqual(ctx.SCC._val, 1)
def test_scc_comparison_false(self):
ctx = ExecContext(s0=42, s1=43)
ctx.run("SCC = S0.u32 == S1.u32")
self.assertEqual(ctx.SCC._val, 0)
def test_ternary(self):
code = _compile_pseudocode("D0.u32 = S0.u32 > S1.u32 ? 1'1U : 1'0U")
ctx = ExecContext(s0=5, s1=3)
ctx.run(code)
self.assertEqual(ctx.D0._val, 1)
def test_pack(self):
code = _compile_pseudocode("D0 = { S1[15:0].u16, S0[15:0].u16 }")
ctx = ExecContext(s0=0x1234, s1=0x5678)
ctx.run(code)
self.assertEqual(ctx.D0._val, 0x56781234)
def test_tmp_with_typed_access(self):
code = _compile_pseudocode("""tmp = S0.u32 + S1.u32
D0.u32 = tmp.u32""")
ctx = ExecContext(s0=100, s1=200)
ctx.run(code)
self.assertEqual(ctx.D0._val, 300)
def test_s_add_u32_pattern(self):
# Real pseudocode pattern from S_ADD_U32
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
D0.u32 = tmp.u32""")
# Test overflow case
ctx = ExecContext(s0=0xFFFFFFFF, s1=0x00000001)
ctx.run(code)
self.assertEqual(ctx.D0._val, 0) # Wraps to 0
self.assertEqual(ctx.SCC._val, 1) # Carry set
def test_s_add_u32_no_overflow(self):
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
D0.u32 = tmp.u32""")
ctx = ExecContext(s0=100, s1=200)
ctx.run(code)
self.assertEqual(ctx.D0._val, 300)
self.assertEqual(ctx.SCC._val, 0) # No carry
def test_vcc_lane_read(self):
ctx = ExecContext(vcc=0b1010, lane=1)
# Lane 1 is set
self.assertEqual(ctx.VCC.u64[1], 1)
self.assertEqual(ctx.VCC.u64[2], 0)
def test_vcc_lane_write(self):
ctx = ExecContext(vcc=0, lane=0)
ctx.VCC.u64[3] = 1
ctx.VCC.u64[1] = 1
self.assertEqual(ctx.VCC._val, 0b1010)
def test_for_loop(self):
# CTZ pattern - find first set bit
code = _compile_pseudocode("""tmp = -1
for i in 0 : 31 do
if S0.u32[i] == 1 then
tmp = i
endif
endfor
D0.i32 = tmp""")
ctx = ExecContext(s0=0b1000) # Bit 3 is set
ctx.run(code)
self.assertEqual(ctx.D0._val & MASK32, 3)
def test_result_dict(self):
ctx = ExecContext(s0=5, s1=3)
ctx.D0.u32 = 42
ctx.SCC._val = 1
result = ctx.result()
self.assertEqual(result['d0'], 42)
self.assertEqual(result['scc'], 1)
class TestPseudocodeRegressions(unittest.TestCase):
"""Regression tests for pseudocode instruction emulation bugs."""
def test_v_div_scale_f32_vcc_always_returned(self):
"""V_DIV_SCALE_F32 must always return VCC, even when VCC=0 (no scaling needed).
Bug: when VCC._val == vcc (both 0), VCC wasn't returned, so VCC bits weren't written.
This caused division to produce wrong results for multiple lanes."""
# Normal case: 1.0 / 3.0, no scaling needed, VCC should be 0
s0 = 0x3f800000 # 1.0
s1 = 0x40400000 # 3.0
s2 = 0x3f800000 # 1.0 (numerator)
result = _VOP3SDOp_V_DIV_SCALE_F32(s0, s1, s2, 0, 0, 0, 0, 0xffffffff, 0, None)
# Must always have VCC in result
self.assertIn('VCC', result, "V_DIV_SCALE_F32 must always return VCC")
self.assertEqual(result['VCC'] & 1, 0, "VCC lane 0 should be 0 when no scaling needed")
def test_v_cmp_class_f32_detects_quiet_nan(self):
"""V_CMP_CLASS_F32 must correctly identify quiet NaN vs signaling NaN.
Bug: isQuietNAN and isSignalNAN both used math.isnan which can't distinguish them."""
quiet_nan = 0x7fc00000 # quiet NaN: exponent=255, bit22=1
signal_nan = 0x7f800001 # signaling NaN: exponent=255, bit22=0
# Test quiet NaN detection (bit 1 in mask)
s1_quiet = 0b0000000010 # bit 1 = quiet NaN
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 1, "Should detect quiet NaN with quiet NaN mask")
# Test signaling NaN detection (bit 0 in mask)
s1_signal = 0b0000000001 # bit 0 = signaling NaN
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 1, "Should detect signaling NaN with signaling NaN mask")
# Test that quiet NaN doesn't match signaling NaN mask
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 0, "Quiet NaN should not match signaling NaN mask")
# Test that signaling NaN doesn't match quiet NaN mask
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 0, "Signaling NaN should not match quiet NaN mask")
def testisNAN_with_typed_view(self):
"""isNAN must work with TypedView objects, not just Python floats.
Bug: isNAN checked isinstance(x, float) which returned False for TypedView."""
nan_reg = Reg(0x7fc00000) # quiet NaN
normal_reg = Reg(0x3f800000) # 1.0
inf_reg = Reg(0x7f800000) # +inf
self.assertTrue(isNAN(nan_reg.f32), "isNAN should return True for NaN TypedView")
self.assertFalse(isNAN(normal_reg.f32), "isNAN should return False for normal TypedView")
self.assertFalse(isNAN(inf_reg.f32), "isNAN should return False for inf TypedView")
class TestBF16(unittest.TestCase):
"""Tests for BF16 (bfloat16) support."""
def test_bf16_conversion(self):
"""Test bf16 <-> f32 conversion."""
# bf16 is just the top 16 bits of f32
# 1.0f = 0x3f800000, bf16 = 0x3f80
self.assertAlmostEqual(_bf16(0x3f80), 1.0, places=2)
self.assertEqual(_ibf16(1.0), 0x3f80)
# 2.0f = 0x40000000, bf16 = 0x4000
self.assertAlmostEqual(_bf16(0x4000), 2.0, places=2)
self.assertEqual(_ibf16(2.0), 0x4000)
# -1.0f = 0xbf800000, bf16 = 0xbf80
self.assertAlmostEqual(_bf16(0xbf80), -1.0, places=2)
self.assertEqual(_ibf16(-1.0), 0xbf80)
def test_bf16_special_values(self):
"""Test bf16 special values (inf, nan)."""
import math
# +inf: f32 = 0x7f800000, bf16 = 0x7f80
self.assertTrue(math.isinf(_bf16(0x7f80)))
self.assertEqual(_ibf16(float('inf')), 0x7f80)
# -inf: f32 = 0xff800000, bf16 = 0xff80
self.assertTrue(math.isinf(_bf16(0xff80)))
self.assertEqual(_ibf16(float('-inf')), 0xff80)
# NaN: quiet NaN bf16 = 0x7fc0
self.assertTrue(math.isnan(_bf16(0x7fc0)))
self.assertEqual(_ibf16(float('nan')), 0x7fc0)
def test_bf16_register_property(self):
"""Test Reg.bf16 property."""
r = Reg(0)
r.bf16 = 3.0 # 3.0f = 0x40400000, bf16 = 0x4040
self.assertEqual(r._val & 0xffff, 0x4040)
self.assertAlmostEqual(float(r.bf16), 3.0, places=1)
def test_bf16_slice_property(self):
"""Test TypedView.bf16 property."""
r = Reg(0x40404040) # Two bf16 3.0 values
self.assertAlmostEqual(r[15:0].bf16, 3.0, places=1)
self.assertAlmostEqual(r[31:16].bf16, 3.0, places=1)
class TestBytePermute(unittest.TestCase):
"""Tests for BYTE_PERMUTE helper function (V_PERM_B32)."""
def test_byte_select_0_to_7(self):
"""Test selecting bytes 0-7 from 64-bit data."""
# data = {s0, s1} where s0 is bytes 0-3, s1 is bytes 4-7
# Combined: 0x0706050403020100 (byte 0 = 0x00, byte 7 = 0x07)
data = 0x0706050403020100
for i in range(8):
self.assertEqual(BYTE_PERMUTE(data, i), i, f"byte {i} should be {i}")
def test_sign_extend_bytes(self):
"""Test sign extension selectors 8-11."""
# sel 8: sign of byte 1 (bits 15:8)
# sel 9: sign of byte 3 (bits 31:24)
# sel 10: sign of byte 5 (bits 47:40)
# sel 11: sign of byte 7 (bits 63:56)
data = 0x8000800080008000 # All relevant bytes have sign bit set
self.assertEqual(BYTE_PERMUTE(data, 8), 0xff)
self.assertEqual(BYTE_PERMUTE(data, 9), 0xff)
self.assertEqual(BYTE_PERMUTE(data, 10), 0xff)
self.assertEqual(BYTE_PERMUTE(data, 11), 0xff)
data = 0x7f007f007f007f00 # No sign bits set
self.assertEqual(BYTE_PERMUTE(data, 8), 0x00)
self.assertEqual(BYTE_PERMUTE(data, 9), 0x00)
self.assertEqual(BYTE_PERMUTE(data, 10), 0x00)
self.assertEqual(BYTE_PERMUTE(data, 11), 0x00)
def test_constant_zero(self):
"""Test selector 12 returns 0x00."""
self.assertEqual(BYTE_PERMUTE(0xffffffffffffffff, 12), 0x00)
def test_constant_ff(self):
"""Test selectors >= 13 return 0xFF."""
for sel in [13, 14, 15, 255]:
self.assertEqual(BYTE_PERMUTE(0, sel), 0xff, f"sel {sel} should be 0xff")
class TestSADHelpers(unittest.TestCase):
"""Tests for V_SAD_U8 and V_MSAD_U8 helper functions."""
def test_v_sad_u8_basic(self):
"""Test v_sad_u8 with simple values."""
# s0 = 0x04030201, s1 = 0x04030201 -> diff = 0 for all bytes
result = v_sad_u8(0x04030201, 0x04030201, 0)
self.assertEqual(result, 0)
# s0 = 0x05040302, s1 = 0x04030201 -> diff = 1+1+1+1 = 4
result = v_sad_u8(0x05040302, 0x04030201, 0)
self.assertEqual(result, 4)
def test_v_sad_u8_with_accumulator(self):
"""Test v_sad_u8 with non-zero accumulator."""
# s0 = 0x05040302, s1 = 0x04030201, s2 = 100 -> 4 + 100 = 104
result = v_sad_u8(0x05040302, 0x04030201, 100)
self.assertEqual(result, 104)
def test_v_sad_u8_large_diff(self):
"""Test v_sad_u8 with maximum byte differences."""
# s0 = 0xffffffff, s1 = 0x00000000 -> diff = 255*4 = 1020
result = v_sad_u8(0xffffffff, 0x00000000, 0)
self.assertEqual(result, 1020)
def test_v_msad_u8_basic(self):
"""Test v_msad_u8 masks when reference byte is 0."""
# s0 = 0x10101010, s1 = 0x00000000 -> all masked, result = 0
result = v_msad_u8(0x10101010, 0x00000000, 0)
self.assertEqual(result, 0)
# s0 = 0x10101010, s1 = 0x01010101 -> diff = |0x10-0x01|*4 = 15*4 = 60
result = v_msad_u8(0x10101010, 0x01010101, 0)
self.assertEqual(result, 60)
def test_v_msad_u8_partial_mask(self):
"""Test v_msad_u8 with partial masking."""
# s0 = 0x10101010, s1 = 0x00010001 -> bytes 1 and 3 masked
# diff = |0x10-0x01| + |0x10-0x01| = 15 + 15 = 30
result = v_msad_u8(0x10101010, 0x00010001, 0)
self.assertEqual(result, 30)
def test_v_msad_u8_with_accumulator(self):
"""Test v_msad_u8 with non-zero accumulator."""
result = v_msad_u8(0x10101010, 0x01010101, 50)
self.assertEqual(result, 110) # 60 + 50
if __name__ == '__main__':
unittest.main()
-71
View File
@@ -1,71 +0,0 @@
#!/usr/bin/env python3
"""Test pdf.py PDF parser and enum generation."""
import unittest, tempfile, importlib.util
from extra.assembly.amd.pdf import extract, extract_tables, extract_enums, extract_pcode, write_enums, PDF_URLS
EXPECTED = {
"rdna3": {"pages": 655, "tables": 115, "sop2_ops": 67, "sop2_first": "S_ADD_U32"},
"rdna4": {"pages": 711, "tables": 125, "sop2_ops": 74, "sop2_first": "S_ADD_CO_U32"},
"cdna": {"pages": 610, "tables": 104, "sop2_ops": 52, "sop2_first": "S_ADD_U32"},
}
class TestPDF2(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.data = {name: extract(url) for name, url in PDF_URLS.items()}
cls.tables = {name: extract_tables(pages) for name, pages in cls.data.items()}
cls.enums = {name: extract_enums(cls.tables[name]) for name in PDF_URLS}
cls.pcode = {name: extract_pcode(cls.data[name], cls.enums[name]) for name in PDF_URLS}
def test_page_counts(self):
for name, exp in EXPECTED.items():
self.assertEqual(len(self.data[name]), exp["pages"], f"{name} page count")
def test_table_counts(self):
for name, exp in EXPECTED.items():
self.assertEqual(len(self.tables[name]), exp["tables"], f"{name} table count")
def test_tables_sequential(self):
for name in PDF_URLS:
nums = sorted(self.tables[name].keys())
missing = set(range(1, max(nums) + 1)) - set(nums)
self.assertEqual(missing, set(), f"{name} missing tables: {missing}")
def test_generate_enums(self):
for name, exp in EXPECTED.items():
with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
write_enums(self.enums[name], name, f.name)
spec = importlib.util.spec_from_file_location("enum", f.name)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
# Check SOP2Op
self.assertTrue(hasattr(mod, 'SOP2Op'), f"{name} missing SOP2Op")
self.assertEqual(len(mod.SOP2Op), exp["sop2_ops"], f"{name} SOP2Op count")
self.assertEqual(mod.SOP2Op(0).name, exp["sop2_first"], f"{name} SOP2Op first")
# Check all enums have at least 2 ops
for attr in dir(mod):
if attr.endswith('Op'):
self.assertGreaterEqual(len(getattr(mod, attr)), 2, f"{name} {attr} has too few ops")
def test_pcode_rdna3_tricky(self):
"""Test specific pseudocode patterns that are tricky to extract correctly."""
pcode = self.pcode['rdna3']
# BUFFER_ATOMIC_MAX_U64: should have 4 statements (not truncated)
self.assertEqual(pcode[('BUFFER_ATOMIC_MAX_U64', 72)],
'tmp = MEM[ADDR].u64;\nsrc = DATA.u64;\nMEM[ADDR].u64 = src >= tmp ? src : tmp;\nRETURN_DATA.u64 = tmp')
# GLOBAL_STORE_B128: should have 4 MEM stores (not truncated)
self.assertEqual(pcode[('GLOBAL_STORE_B128', 29)],
'MEM[ADDR].b32 = VDATA[31 : 0];\nMEM[ADDR + 4U].b32 = VDATA[63 : 32];\nMEM[ADDR + 8U].b32 = VDATA[95 : 64];\nMEM[ADDR + 12U].b32 = VDATA[127 : 96]')
# S_CMOVK_I32: should have full if/endif block
self.assertEqual(pcode[('S_CMOVK_I32', 2)],
"if SCC then\nD0.i32 = 32'I(signext(SIMM16.i16))\nendif")
def test_pcode_no_examples(self):
"""Pseudocode should not contain example lines with '=>'."""
for name in PDF_URLS:
for (op_name, opcode), code in self.pcode[name].items():
# No example lines (test vectors like "S_CTZ_I32_B32(0xaaaaaaaa) => 1")
self.assertNotIn('=>', code, f"{name} {op_name} contains example line with '=>'")
if __name__ == "__main__":
unittest.main()
-95
View File
@@ -1,95 +0,0 @@
#!/usr/bin/env python3
import unittest, subprocess
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.test.helpers import get_llvm_mc
def llvm_assemble(asm: str) -> bytes:
"""Assemble using llvm-mc and return bytes."""
result = subprocess.run(
[get_llvm_mc(), "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
input=asm, capture_output=True, text=True
)
out = b''
for line in result.stdout.split('\n'):
if 'encoding:' in line:
enc = line.split('encoding:')[1].strip()
enc = enc.strip('[]').replace('0x', '').replace(',', '')
out += bytes.fromhex(enc)
if not out: raise ValueError(f"no encoding found: {result.stdout} {result.stderr}")
return out
class TestRDNA3Asm(unittest.TestCase):
def test_full_program(self):
"""Test the full program from rdna3fun.py matches llvm-mc output."""
program = [
v_bfe_u32(v[1], v[0], 10, 10),
s_load_b128(s[4:7], s[0:1], NULL),
v_and_b32_e32(v[0], 0x3FF, v[0]),
s_mulk_i32(s[3], 0x87),
v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2]),
v_mul_u32_u24_e32(v[0], 45, v[0]),
v_ashrrev_i32_e32(v[2], 31, v[1]),
v_add3_u32(v[0], v[0], s[3], v[1]),
v_lshlrev_b64(v[2:3], 2, v[1:2]),
v_ashrrev_i32_e32(v[1], 31, v[0]),
v_lshlrev_b64(v[0:1], 2, v[0:1]),
s_waitcnt(0xfc07), # lgkmcnt(0)
v_add_co_u32(v[2], VCC_LO, s[6], v[2]),
v_add_co_ci_u32_e32(v[3], s[7], v[3]),
v_add_co_u32(v[0], VCC_LO, s[4], v[0]),
global_load_b32(vdst=v[2], addr=v[2], 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),
s_endpgm(),
]
asm = """
v_bfe_u32 v1, v0, 10, 10
s_load_b128 s[4:7], s[0:1], null
v_and_b32_e32 v0, 0x3FF, v0
s_mulk_i32 s3, 0x87
v_mad_u64_u32 v[1:2], null, s2, 3, v[1:2]
v_mul_u32_u24_e32 v0, 45, v0
v_ashrrev_i32_e32 v2, 31, v1
v_add3_u32 v0, v0, s3, v1
v_lshlrev_b64 v[2:3], 2, v[1:2]
v_ashrrev_i32_e32 v1, 31, v0
v_lshlrev_b64 v[0:1], 2, v[0:1]
s_waitcnt lgkmcnt(0)
v_add_co_u32 v2, vcc_lo, s6, v2
v_add_co_ci_u32_e32 v3, vcc_lo, s7, v3, vcc_lo
v_add_co_u32 v0, vcc_lo, s4, v0
global_load_b32 v2, v[2:3], off
v_add_co_ci_u32_e32 v1, vcc_lo, s5, v1, vcc_lo
s_waitcnt vmcnt(0)
global_store_b32 v[0:1], v2, off
s_endpgm
"""
expected = llvm_assemble(asm)
for inst,rt in zip(program, asm.strip().split("\n")): print(f"{inst.disasm():50s} {rt}")
actual = b''.join(inst.to_bytes() for inst in program)
self.assertEqual(actual, expected)
def test_sop2_s_add_u32(self):
inst = SOP2(SOP2Op.S_ADD_U32, s[3], s[0], s[1])
expected = llvm_assemble("s_add_u32 s3, s0, s1")
self.assertEqual(inst.to_bytes(), expected)
def test_vop2_v_and_b32_inline_const(self):
inst = v_and_b32_e32(v[0], 10, v[0])
expected = llvm_assemble("v_and_b32_e32 v0, 10, v0")
self.assertEqual(inst.to_bytes(), expected)
def test_sopp_s_endpgm(self):
inst = s_endpgm()
expected = llvm_assemble("s_endpgm")
self.assertEqual(inst.to_bytes(), expected)
def test_sop1_s_mov_b32(self):
inst = s_mov_b32(s[0], s[1])
expected = llvm_assemble("s_mov_b32 s0, s1")
self.assertEqual(inst.to_bytes(), expected)
if __name__ == "__main__":
unittest.main()
-257
View File
@@ -1,257 +0,0 @@
#!/usr/bin/env python3
"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
import unittest, io, sys, re, subprocess, os
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.asm import asm, detect_format
from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump
# arch: (mcpu, mattr)
ARCH_CONFIG = {
'rdna3': ('gfx1100', '+real-true16,+wavefrontsize32'),
'rdna4': ('gfx1200', '+real-true16,+wavefrontsize32'),
'cdna': ('gfx942', '+wavefrontsize64'),
}
def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
"""Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
old_stdout = sys.stdout
sys.stdout = io.StringIO()
compiler.disassemble(lib)
output = sys.stdout.getvalue()
sys.stdout = old_stdout
results = []
for line in output.splitlines():
if '//' not in line: continue
instr = line.split('//')[0].strip()
if not instr: continue
comment = line.split('//')[1].strip()
if ':' not in comment: continue
hex_str = comment.split(':')[1].strip().split()[0]
try:
machine_bytes = bytes.fromhex(hex_str)[::-1] # big-endian to little-endian
results.append((instr, machine_bytes))
except ValueError:
continue
return results
def compile_asm(instr: str, arch: str = 'rdna3') -> bytes:
"""Compile a single instruction using LLVM."""
return compile_asm_batch([instr], arch)[0]
def compile_asm_batch(instrs: list[str], arch: str = 'rdna3') -> list[bytes]:
"""Compile multiple instructions with a single llvm-mc call."""
if not instrs: return []
mcpu, mattr = ARCH_CONFIG[arch]
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', f'-mattr={mattr}', '-show-encoding'],
input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
encodings = []
for line in result.stdout.split('\n'):
if 'encoding:' in line:
enc = line.split('encoding:')[1].strip()
if enc.startswith('[') and enc.endswith(']'):
encodings.append(bytes.fromhex(enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')))
if len(encodings) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(encodings)}")
return encodings
def compile_and_disasm_batch(instrs: list[str], arch: str = 'rdna3') -> list[str]:
"""Compile instructions with LLVM and get LLVM's disassembly."""
import tempfile
if not instrs: return []
mcpu, mattr = ARCH_CONFIG[arch]
src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n" + "\n".join(f" {instr}" for instr in instrs) + "\n"
with tempfile.NamedTemporaryFile(suffix='.o', delete=False) as f:
obj_path = f.name
try:
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', f'-mattr={mattr}', '-filetype=obj', '-o', obj_path],
input=src, capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
result = subprocess.run([get_llvm_objdump(), '-d', f'--mcpu={mcpu}', obj_path], capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-objdump failed: {result.stderr.strip()}")
results: list[str] = []
for line in result.stdout.splitlines():
if '//' not in line: continue
instr = line.split('//')[0].strip()
if instr: results.append(instr)
return results[:len(instrs)]
finally:
os.unlink(obj_path)
class TestTinygradKernelRoundtrip(unittest.TestCase):
"""Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
arch = 'rdna3'
def _test_kernel_roundtrip(self, op_fn):
"""Generate kernel from op_fn, test:
1. decode -> reencode matches original bytes
2. asm(disasm()) matches LLVM output
3. our disasm() matches LLVM's disassembly string exactly
"""
arch = self.arch
mcpu, mattr = ARCH_CONFIG[arch]
from extra.assembly.amd.test.test_compare_emulators import get_kernels_from_tinygrad
from tinygrad.runtime.support.compiler_amd import HIPCompiler
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
compiler = HIPCompiler(mcpu)
# First pass: decode all instructions and collect info
decoded_instrs: list[tuple] = [] # list of (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err)
for ki, kernel in enumerate(kernels):
offset = 0
while offset < len(kernel.code):
remaining = kernel.code[offset:]
fmt = detect_format(remaining, arch)
if fmt is None:
decoded_instrs.append((ki, offset, None, None, None, False, "no format"))
offset += 4
continue
base_size = fmt._size()
if len(remaining) < base_size:
break
try:
decoded = fmt.from_bytes(remaining) # pass all remaining bytes so from_bytes can read literal
size = decoded.size() # actual size including literal
orig_bytes = remaining[:size]
reencoded = decoded.to_bytes()
our_disasm = decoded.disasm()
decode_ok = reencoded == orig_bytes
decode_err: str | None = None if decode_ok else f"orig={orig_bytes.hex()} reenc={reencoded.hex()}"
decoded_instrs.append((ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err))
except Exception as e:
decoded_instrs.append((ki, offset, remaining[:base_size], None, None, False, str(e)))
size = base_size
offset += size
# Collect disasm strings for batched LLVM calls - skip unknown opcodes (op_X) that LLVM can't compile
asm_test_instrs: list[tuple[int, str]] = [] # (idx, our_disasm) for asm test
disasm_test_instrs: list[tuple[int, str]] = [] # (idx, our_disasm) for disasm comparison test
for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
if our_disasm is None: continue
# Skip unknown opcodes and malformed instructions for both tests
if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm): continue
asm_test_instrs.append((idx, our_disasm))
disasm_test_instrs.append((idx, our_disasm))
# Batch compile for asm test
asm_llvm_results = compile_asm_batch([d for _, d in asm_test_instrs], arch)
asm_llvm_map = {idx: result for (idx, _), result in zip(asm_test_instrs, asm_llvm_results)}
# Batch compile+disasm for disasm comparison test
disasm_llvm_results = compile_and_disasm_batch([d for _, d in disasm_test_instrs], arch)
disasm_llvm_map = {idx: result for (idx, _), result in zip(disasm_test_instrs, disasm_llvm_results)}
# Now evaluate results
decode_passed, decode_failed, decode_skipped = 0, 0, 0
asm_passed, asm_failed, asm_skipped = 0, 0, 0
disasm_passed, disasm_failed, disasm_skipped = 0, 0, 0
decode_failures: list[str] = []
asm_failures: list[str] = []
disasm_failures: list[str] = []
for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
# Decode test
if decode_ok:
decode_passed += 1
elif decode_err == "no format":
decode_skipped += 1
else:
decode_failed += 1
decode_failures.append(f"K{ki}@{offset}: {our_disasm}: {decode_err}")
# Asm test
if our_disasm is None:
asm_skipped += 1
elif idx in asm_llvm_map:
llvm_bytes = asm_llvm_map[idx]
try:
our_bytes = asm(our_disasm).to_bytes()
if our_bytes[:len(llvm_bytes)] == llvm_bytes:
asm_passed += 1
else:
asm_failed += 1
asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': ours={our_bytes[:len(llvm_bytes)].hex()} llvm={llvm_bytes.hex()}")
except Exception:
asm_skipped += 1
else:
asm_skipped += 1
# Disasm comparison test
if our_disasm is None:
disasm_skipped += 1
elif idx in disasm_llvm_map:
llvm_disasm = disasm_llvm_map[idx]
if our_disasm == llvm_disasm:
disasm_passed += 1
else:
disasm_failed += 1
disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
else:
disasm_skipped += 1
print(f"[{arch}] decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
print(f"[{arch}] asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
print(f"[{arch}] disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
# Note: disasm string comparison is informational only - formatting differences between LLVM versions are expected
# Basic unary ops
def test_neg(self): self._test_kernel_roundtrip(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
def test_relu(self): self._test_kernel_roundtrip(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
def test_exp(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).exp())
def test_log(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 3.0]).log())
def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
# Binary ops
def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
# Reductions
def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
# Matmul
def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
# Complex ops
def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
# Memory patterns
def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
# Cast ops
def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
# Comparison ops
def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
# Fused ops
def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
@unittest.skip("no asm support for RDNA4")
class TestTinygradKernelRoundtripRDNA4(TestTinygradKernelRoundtrip): arch = 'rdna4'
@unittest.skip("no asm support for CDNA")
class TestTinygradKernelRoundtripCDNA(TestTinygradKernelRoundtrip): arch = 'cdna'
if __name__ == "__main__":
unittest.main()
-79
View File
@@ -1,79 +0,0 @@
#!/usr/bin/env python3
"""Tests for SQTT packet codec (no hardware required)."""
import unittest
from extra.assembly.amd.sqtt import (
LAYOUT_HEADER, WAVESTART, WAVEEND, INST, NOP,
decode, encode, PACKET_TYPES, OPCODE_TO_CLASS
)
class TestSQTTCodec(unittest.TestCase):
"""Tests for SQTT encoder/decoder roundtrip."""
def test_roundtrip_simple(self):
"""Test encode/decode roundtrip for simple packets."""
test_packets = [
LAYOUT_HEADER.from_raw(0x100),
WAVESTART.from_raw(0x0),
INST.from_raw(0x10), # delta=1
INST.from_raw(0x10), # delta=1
WAVEEND.from_raw(0x40), # delta=2
]
encoded = encode(test_packets)
decoded = decode(encoded)
self.assertGreaterEqual(len(decoded), len(test_packets))
for i, (orig, dec) in enumerate(zip(test_packets, decoded)):
self.assertEqual(type(orig), type(dec), f"type mismatch at {i}")
def test_decode_empty(self):
"""Test decoding empty data."""
packets = decode(b'')
self.assertEqual(packets, [])
def test_encode_empty(self):
"""Test encoding empty list."""
data = encode([])
self.assertEqual(data, b'')
def test_all_packet_types_have_encoding(self):
"""All packet types should have an encoding defined."""
for pkt_cls in PACKET_TYPES:
self.assertIsNotNone(pkt_cls._encoding, f"{pkt_cls.__name__} missing encoding")
def test_packet_from_raw(self):
"""Test creating packets from raw values."""
# INST with wave=5, op=0x21, delta=2
raw = (0x21 << 13) | (5 << 8) | (2 << 4) | 0b010
pkt = INST.from_raw(raw)
self.assertEqual(pkt.wave, 5)
self.assertEqual(pkt.op, 0x21)
self.assertEqual(pkt.delta, 2)
class TestDecodeRealBlob(unittest.TestCase):
"""Test decoding real SQTT blobs from examples."""
def test_decode_example_file(self):
"""Test decoding a real SQTT blob from examples."""
import pickle
from pathlib import Path
example_path = Path(__file__).parent.parent.parent.parent / "sqtt/examples/profile_plus_run_0.pkl"
if not example_path.exists():
self.skipTest(f"Example file not found: {example_path}")
from tinygrad.runtime.ops_amd import ProfileSQTTEvent
with open(example_path, "rb") as f:
data = pickle.load(f)
sqtt_events = [e for e in data if isinstance(e, ProfileSQTTEvent)]
self.assertGreater(len(sqtt_events), 0, "No SQTT events in example")
packets = decode(sqtt_events[0].blob)
self.assertGreater(len(packets), 0, "No packets decoded")
# First packet should be LAYOUT_HEADER
self.assertIsInstance(packets[0], LAYOUT_HEADER)
if __name__ == "__main__":
unittest.main()
File diff suppressed because it is too large Load Diff
@@ -1,545 +0,0 @@
#!/usr/bin/env python3
"""Tests for SQTT emulator correctness against known hardware patterns.
NOTE: This file only tests NOP and VALU behavior. For WMMA/DP/trans tests,
see test_sqtt_compare.py.
Run emulator tests: PYTHONPATH="." python3 extra/assembly/amd/test/test_sqtt_correct.py
Run hardware tests: SQTT_HW=1 PYTHONPATH="." python3 extra/assembly/amd/test/test_sqtt_correct.py
"""
import os
import unittest
USE_HW = os.environ.get("SQTT_HW", "0") == "1"
if USE_HW:
os.environ["SQTT"] = "1"
os.environ["PROFILE"] = "1"
os.environ["SQTT_LIMIT_SE"] = "2"
os.environ["SQTT_TOKEN_EXCLUDE"] = "3784"
from extra.assembly.amd.emu import SQTTState, decode_program, exec_wave, WaveState, LDSMem
from extra.assembly.amd.sqtt import WAVESTART, WAVEEND
from extra.assembly.amd.autogen.rdna3.ins import v_mov_b32_e32, v_add_f32_e32, s_nop, s_endpgm, s_delay_alu
from extra.assembly.amd.dsl import v
def assemble(instructions: list) -> bytes:
return b''.join(inst.to_bytes() for inst in instructions)
def wrap_with_nops(instructions: list, nops=16) -> list:
return instructions + [s_nop(0)]*nops + [s_endpgm()]
def get_wave_packets(packets: list) -> list:
result, in_wave = [], False
for p in packets:
if isinstance(p, WAVESTART) and p.simd == 0:
in_wave, result = True, [p]
elif in_wave:
result.append(p)
if isinstance(p, WAVEEND): break
return result
def get_timing_deltas(packets: list) -> list[tuple[str, int]]:
skip_types = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG"}
filtered = [p for p in packets if type(p).__name__ not in skip_types]
if not filtered: return []
result = [(type(filtered[0]).__name__, 0)]
for i in range(1, len(filtered)):
result.append((type(filtered[i]).__name__, filtered[i]._time - filtered[i-1]._time))
return result
def run_emulator(instructions: list) -> list:
code = assemble(instructions)
program = decode_program(code)
st = WaveState()
st.exec_mask = (1 << 32) - 1
lds = LDSMem(bytearray(65536))
trace = SQTTState(wave_id=0, simd=0, cu=0)
exec_wave(program, st, lds, 32, trace)
return get_wave_packets(trace.packets)
def get_all_waves(packets: list) -> list[list]:
"""Extract all WAVESTART..WAVEEND ranges on simd 0."""
waves, in_wave, current = [], False, []
for p in packets:
if isinstance(p, WAVESTART) and p.simd == 0:
in_wave, current = True, [p]
elif in_wave:
current.append(p)
if isinstance(p, WAVEEND):
waves.append(current)
in_wave, current = False, []
return waves
def run_hardware(instructions: list) -> list:
from extra.assembly.amd.test.test_sqtt_hw import compile_asm_sqtt, run_prg_sqtt_batch
from extra.assembly.amd.sqtt import decode
from collections import Counter
prg = compile_asm_sqtt(instructions, alu_only=True)
for _ in range(10):
blobs = run_prg_sqtt_batch(prg, n_runs=200)
# Extract all waves from all blobs
traces = []
for blob in blobs:
traces.extend(get_all_waves(decode(blob)))
if not traces:
continue
# Find most common pattern
delta_sets = [tuple(get_timing_deltas(t)) for t in traces]
most_common = Counter(delta_sets).most_common(1)[0][0]
for t in traces:
if tuple(get_timing_deltas(t)) == most_common:
return t
return []
def run_sqtt(instructions: list, nops: int = 16) -> list:
instructions = wrap_with_nops(instructions, nops=nops)
return run_hardware(instructions) if USE_HW else run_emulator(instructions)
def get_deltas(instructions: list) -> tuple[list[int], list[int]]:
"""Run and return (issue deltas, exec deltas).
Issue = IMMEDIATE + VALUINST, Exec = ALUEXEC.
Deltas are between consecutive packets of same stream."""
deltas = get_timing_deltas(run_sqtt(instructions))
time = 0
issue_times, exec_times = [], []
for ptype, delta in deltas:
time += delta
if ptype in ('IMMEDIATE', 'VALUINST'):
issue_times.append(time)
elif ptype == 'ALUEXEC':
exec_times.append(time)
issue = [issue_times[i] - issue_times[i-1] for i in range(1, len(issue_times))]
execd = [exec_times[i] - exec_times[i-1] for i in range(1, len(exec_times))]
return issue, execd
# ************************************ tests ************************************
class TestVALUChains(unittest.TestCase):
"""VALU dependency chains."""
def _chain(self, n, expected_issue, expected_exec):
instrs = [v_mov_b32_e32(v[0], 1.0)] + [v_add_f32_e32(v[i], v[i-1], v[i-1]) for i in range(1, n)]
issue, execd = get_deltas(instrs)
self.assertEqual(issue[:n-1], expected_issue)
if isinstance(expected_exec[0], list): self.assertIn(execd, expected_exec)
else: self.assertEqual(execd, expected_exec)
def test_chain_2(self): self._chain(2, [1], [6])
def test_chain_3(self): self._chain(3, [1, 1], [6, 5])
def test_chain_4(self): self._chain(4, [1, 1, 1], [6, 5, 5])
def test_chain_5(self): self._chain(5, [1, 1, 1, 1], [6, 5, 5, 9])
def test_chain_6(self): self._chain(6, [1, 1, 1, 1, 1], [6, 5, 5, 9, 9])
def test_chain_7(self): self._chain(7, [1, 1, 1, 1, 1, 1], [6, 5, 5, 5, 9, 9])
def test_chain_8(self): self._chain(8, [1, 1, 1, 1, 1, 1, 1], [6, 5, 5, 5, 9, 9, 9])
# NOTE: position 8 can be 5 or 9 depending on GPU variant
def test_chain_12(self): self._chain(12, [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [[6, 5, 5, 5, 5, 9, 9, 9, 9, 9, 9], [6, 5, 5, 5, 5, 9, 9, 9, 5, 9, 9]])
def test_chain_14(self): self._chain(14, [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [[6, 5, 5, 5, 5, 9, 9, 9, 9, 9, 9, 9, 9], [6, 5, 5, 5, 5, 9, 9, 9, 5, 9, 9, 9, 9]])
# issue stalls start here
def test_chain_15(self): self._chain(15, [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3], [[6, 5, 5, 5, 5, 5, 9, 9, 9, 9, 9, 9, 9, 9], [6, 5, 5, 5, 5, 5, 9, 9, 5, 9, 9, 9, 9, 9]])
def test_chain_16(self): self._chain(16, [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 5], [[6, 5, 5, 5, 5, 5, 5, 9, 9, 9, 9, 9, 9, 9, 9], [6, 5, 5, 5, 5, 5, 5, 9, 5, 9, 9, 9, 9, 9, 9]])
def test_chain_18(self): self._chain(18, [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 5, 5, 5], [6, 5, 5, 5, 5, 5, 5, 5, 5, 9, 9, 9, 9, 9, 9, 9, 9])
def test_chain_20(self): self._chain(20, [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 5, 5, 5, 5, 5], [6, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 9, 9, 9, 9, 9, 9, 9, 9])
class TestVALUChainsWithWarmup(unittest.TestCase):
"""VALU dependency chains with early VALUs to isolate warmup effects."""
# just the first stupid VALU takes 6
def _chain(self, n, warmup=True):
instrs = [v_mov_b32_e32(v[0], 1.0), s_nop(100)] if warmup else [s_nop(100)]
instrs += [v_mov_b32_e32(v[0], 1.0)] + [v_add_f32_e32(v[i], v[i-1], v[i-1]) for i in range(1, n)]
issue, execd = get_deltas(instrs)
return execd[1:] if warmup else execd
def test_warmup_chain_2(self): self.assertEqual(self._chain(2), [5])
def test_warmup_chain_3(self): self.assertEqual(self._chain(3), [5, 5])
def test_warmup_chain_4(self): self.assertEqual(self._chain(4), [5, 5, 5])
def test_warmup_chain_5(self): self.assertEqual(self._chain(5), [5, 5, 5, 9])
def test_warmup_chain_6(self): self.assertEqual(self._chain(6), [5, 5, 5, 5, 9])
def test_warmup_chain_7(self): self.assertEqual(self._chain(7), [5, 5, 5, 5, 9, 9])
def test_warmup_chain_8(self): self.assertEqual(self._chain(8), [5, 5, 5, 5, 9, 9, 9])
def test_cold_chain_2(self): self.assertEqual(self._chain(2, False), [6])
def test_cold_chain_3(self): self.assertEqual(self._chain(3, False), [6, 5])
def test_cold_chain_4(self): self.assertEqual(self._chain(4, False), [6, 5, 5])
def test_cold_chain_5(self): self.assertEqual(self._chain(5, False), [6, 5, 5, 9])
def test_cold_chain_6(self): self.assertEqual(self._chain(6, False), [6, 5, 5, 9, 9])
def test_cold_chain_7(self): self.assertEqual(self._chain(7, False), [6, 5, 5, 5, 9, 9])
def test_cold_chain_8(self): self.assertEqual(self._chain(8, False), [6, 5, 5, 5, 9, 9, 9])
class TestVALUIndependent(unittest.TestCase):
"""Independent VALU instructions."""
def _ind(self, n, expected_exec):
instrs = [v_mov_b32_e32(v[i], float(i)) for i in range(n)]
issue, execd = get_deltas(instrs)
self.assertEqual(issue[:n-1], [1]*(n-1))
self.assertEqual(execd, expected_exec)
def test_ind_2(self): self._ind(2, [1])
def test_ind_3(self): self._ind(3, [1, 1])
def test_ind_4(self): self._ind(4, [1, 1, 1])
def test_ind_5(self): self._ind(5, [1, 1, 1, 1])
def test_ind_6(self): self._ind(6, [1, 1, 1, 1, 1])
def test_ind_7(self): self._ind(7, [1, 1, 1, 1, 1, 1])
def test_ind_8(self): self._ind(8, [1, 1, 1, 1, 1, 1, 1])
class TestForwardingGap(unittest.TestCase):
"""Producer + N independent instructions + consumer - tests forwarding window."""
def _exec_deltas(self, n_gap):
instrs = [v_mov_b32_e32(v[0], 1.0)]
instrs += [v_mov_b32_e32(v[10+i], float(i)) for i in range(n_gap)]
instrs += [v_add_f32_e32(v[1], v[0], v[0])]
_, execd = get_deltas(instrs)
return execd
def test_gap0(self): self.assertEqual(self._exec_deltas(0), [6])
def test_gap1(self): self.assertEqual(self._exec_deltas(1), [1, 5])
def test_gap2(self): self.assertEqual(self._exec_deltas(2), [1, 1, 4])
def test_gap3(self): self.assertIn(self._exec_deltas(3), [[1, 1, 1, 3], [1, 1, 1, 4]])
def test_gap4(self): self.assertIn(self._exec_deltas(4), [[1, 1, 1, 1, 3], [1, 1, 1, 1, 4]])
def test_gap5(self): self.assertEqual(self._exec_deltas(5), [1, 1, 1, 1, 1, 4]) # anomaly
def test_gap6(self): self.assertEqual(self._exec_deltas(6), [1, 1, 1, 1, 1, 1, 3])
def test_gap7(self): self.assertEqual(self._exec_deltas(7), [1, 1, 1, 1, 1, 1, 1, 3])
def test_gap8(self): self.assertEqual(self._exec_deltas(8), [1, 1, 1, 1, 1, 1, 1, 1, 3])
def test_gap9(self): self.assertEqual(self._exec_deltas(9), [1, 1, 1, 1, 1, 1, 1, 1, 1, 3])
class TestChainWithIndependentGap(unittest.TestCase):
"""Chain of dependent VALUs with independent VALUs inserted before the last one.
Hardware observation: In a chain v0->v1->v2->v3->v4, if we insert N independent VALUs
before v4, the forwarding behavior changes:
- 0-1 independent VALUs: v4 cannot forward from v3 (delta=9)
- 2+ independent VALUs: v4 can forward from v3 (delta=5)
This suggests forwarding eligibility depends on whether the direct source is in the ALU
at issue time, not just at dispatch time.
"""
def _chain5_gap(self, n_ind):
"""Chain v0->v1->v2->v3->v4 with N independent VALUs before v4. Returns v3->v4 delta."""
instrs = [s_nop(100),
v_mov_b32_e32(v[0], 1.0),
v_mov_b32_e32(v[1], v[0]),
v_mov_b32_e32(v[2], v[1]),
v_mov_b32_e32(v[3], v[2])]
instrs += [v_mov_b32_e32(v[10+i], float(i)) for i in range(n_ind)]
instrs += [v_mov_b32_e32(v[4], v[3])]
_, execd = get_deltas(instrs)
# Chain execs are at indices 0,1,2,3 and last one. Independent ones are in between.
# v3->v4 delta = last exec time - 4th exec time (index 3)
# With n_ind independent VALUs, execd has 4 + n_ind entries
# We want delta between exec[3] (v3) and exec[4+n_ind-1] (v4)
# Actually execd is already deltas, so we need absolute times
time, exec_times = 0, []
packets = run_sqtt(instrs)
for ptype, delta in get_timing_deltas(packets):
time += delta
if ptype == 'ALUEXEC': exec_times.append(time)
# v0,v1,v2,v3 are first 4, v4 is last
return exec_times[-1] - exec_times[3]
def test_gap0(self): self.assertEqual(self._chain5_gap(0), 9)
def test_gap1(self): self.assertEqual(self._chain5_gap(1), 9)
def test_gap2(self): self.assertEqual(self._chain5_gap(2), 5)
def test_gap3(self): self.assertEqual(self._chain5_gap(3), 5)
def test_gap4(self): self.assertEqual(self._chain5_gap(4), 5)
class TestVALULatency(unittest.TestCase):
"""VALU latency depends on VGPR source reads.
6 cycles: no VGPR source (constant only), stays 6 regardless of warmup
8-11 cycles: VGPR source read, decreases with warmup (11->10->9->8)
s_nop(0) after VALU immediately drops VGPR read latency to 8
Anomalies:
- 7 consecutive VALUs (no s_nop) causes +1 cycle penalty
- n=0 or n=3 const VALUs + nop + vgpr = 9 cycles (not 8)
"""
def _get_latency(self, instrs):
if not isinstance(instrs, list): instrs = [instrs]
packets = run_sqtt(instrs)
deltas = get_timing_deltas(packets)
time, valu_times, exec_times = 0, [], []
for ptype, delta in deltas:
time += delta
if ptype == 'VALUINST': valu_times.append(time)
if ptype == 'ALUEXEC': exec_times.append(time)
return exec_times[-1] - valu_times[-1] if valu_times and exec_times else None
# 6-cycle latency: no VGPR source (constant), always 6
def test_const_single(self): self.assertEqual(self._get_latency(v_mov_b32_e32(v[0], 1.0)), 6)
def test_const_literal(self): self.assertEqual(self._get_latency(v_mov_b32_e32(v[0], 565.0)), 6)
def test_const_after_const(self): self.assertEqual(self._get_latency([v_mov_b32_e32(v[0], 1.0), v_mov_b32_e32(v[1], 2.0)]), 6)
def test_const_after_nop(self): self.assertEqual(self._get_latency([v_mov_b32_e32(v[0], 1.0), s_nop(0), v_mov_b32_e32(v[1], 2.0)]), 6)
# VGPR read latency: cold start = 9
def test_vgpr_cold(self): self.assertEqual(self._get_latency(v_mov_b32_e32(v[0], v[1])), 9)
# VGPR read latency: warmup decreases 11->10->9->8
def _vgpr_after_n_const(self, n):
return self._get_latency([v_mov_b32_e32(v[i], float(i)) for i in range(n)] + [v_mov_b32_e32(v[10], v[99])])
def test_vgpr_after_1_const(self): self.assertEqual(self._vgpr_after_n_const(1), 11)
def test_vgpr_after_2_const(self): self.assertEqual(self._vgpr_after_n_const(2), 10)
def test_vgpr_after_3_const(self): self.assertEqual(self._vgpr_after_n_const(3), 9)
def test_vgpr_after_4_const(self): self.assertIn(self._vgpr_after_n_const(4), [8, 9])
def test_vgpr_after_5_const(self): self.assertIn(self._vgpr_after_n_const(5), [8, 9])
def test_vgpr_after_6_const(self): self.assertEqual(self._vgpr_after_n_const(6), 9) # anomaly
def test_vgpr_after_7_const(self): self.assertEqual(self._vgpr_after_n_const(7), 8)
def test_vgpr_after_8_const(self): self.assertEqual(self._vgpr_after_n_const(8), 8)
# s_nop(0) immediately drops VGPR read latency to 8 (or 9 on some variants)
def test_vgpr_nop_warmup(self): self.assertIn(self._get_latency([v_mov_b32_e32(v[0], 1.0), s_nop(0), v_mov_b32_e32(v[1], v[99])]), [8, 9])
# s_nop + vgpr read: latency depends on # of const VALUs before nop
def _n_const_nop_vgpr(self, n):
"""N const VALUs + s_nop(0) + vgpr read."""
instrs = [v_mov_b32_e32(v[i], float(i)) for i in range(n)]
instrs += [s_nop(0)]
instrs += [v_mov_b32_e32(v[10], v[99])]
return self._get_latency(instrs)
def test_0_const_nop_vgpr(self): self.assertEqual(self._n_const_nop_vgpr(0), 9)
def test_1_const_nop_vgpr(self): self.assertIn(self._n_const_nop_vgpr(1), [8, 9])
def test_2_const_nop_vgpr(self): self.assertIn(self._n_const_nop_vgpr(2), [8, 9])
def test_3_const_nop_vgpr(self): self.assertEqual(self._n_const_nop_vgpr(3), 9) # anomaly
def test_4_const_nop_vgpr(self): self.assertEqual(self._n_const_nop_vgpr(4), 8)
def test_5_const_nop_vgpr(self): self.assertEqual(self._n_const_nop_vgpr(5), 8)
def test_6_const_nop_vgpr(self): self.assertEqual(self._n_const_nop_vgpr(6), 8)
def test_7_const_nop_vgpr(self): self.assertEqual(self._n_const_nop_vgpr(7), 8)
class TestChainWithNop(unittest.TestCase):
"""Dependency chain with s_nop between instructions."""
def _test(self, nop_val, expected_issue, expected_exec):
issue, execd = get_deltas([v_mov_b32_e32(v[0], 1.0), s_nop(nop_val), v_add_f32_e32(v[1], v[0], v[0])])
self.assertEqual(issue[:2], expected_issue)
if isinstance(expected_exec[0], list): self.assertIn(execd, expected_exec)
else: self.assertEqual(execd, expected_exec)
def test_nop0(self): self._test(0, [3, 1], [[6], [7]])
def test_nop1(self): self._test(1, [4, 1], [[7], [8]])
def test_nop2(self): self._test(2, [5, 1], [9])
def test_nop3(self): self._test(3, [6, 1], [9])
def test_nop4(self): self._test(4, [11, 1], [10])
def test_nop5(self): self._test(5, [12, 1], [11])
class TestIndWithNop(unittest.TestCase):
"""Independent instructions with s_nop between."""
def _test(self, nop_val, expected_issue, expected_exec):
issue, execd = get_deltas([v_mov_b32_e32(v[0], 1.0), s_nop(nop_val), v_mov_b32_e32(v[1], 2.0)])
self.assertEqual(issue[:2], expected_issue)
self.assertEqual(execd, expected_exec)
def test_nop0(self): self._test(0, [3, 1], [4])
def test_nop1(self): self._test(1, [4, 1], [5])
def test_nop3(self): self._test(3, [6, 1], [7])
def test_nop4(self): self._test(4, [11, 1], [8])
def test_nop5(self): self._test(5, [12, 1], [9])
class TestChain3NopMid(unittest.TestCase):
"""3-instruction chain with s_nop in middle."""
def _test(self, nop_val, expected_issue, expected_exec):
issue, execd = get_deltas([
v_mov_b32_e32(v[0], 1.0), v_add_f32_e32(v[1], v[0], v[0]),
s_nop(nop_val), v_add_f32_e32(v[2], v[1], v[1])])
self.assertEqual(issue[:3], expected_issue)
self.assertEqual(execd, expected_exec)
def test_nop0(self): self._test(0, [1, 3, 1], [6, 5])
def test_nop1(self): self._test(1, [1, 4, 1], [6, 5])
def test_nop2(self): self._test(2, [1, 5, 1], [6, 5])
def test_nop3(self): self._test(3, [1, 10, 1], [6, 5])
class TestInd3NopMid(unittest.TestCase):
"""3 independent instructions with s_nop in middle."""
def _test(self, nop_val, expected_issue, expected_exec):
issue, execd = get_deltas([
v_mov_b32_e32(v[0], 1.0), v_mov_b32_e32(v[1], 2.0),
s_nop(nop_val), v_mov_b32_e32(v[2], 3.0)])
self.assertEqual(issue[:3], expected_issue)
self.assertEqual(execd, expected_exec)
def test_nop0(self): self._test(0, [1, 3, 1], [1, 4])
def test_nop1(self): self._test(1, [1, 4, 1], [1, 5])
def test_nop2(self): self._test(2, [1, 5, 1], [1, 6])
def test_nop3(self): self._test(3, [1, 10, 1], [1, 7])
class TestSNopDelay(unittest.TestCase):
"""Single s_nop delay between two independent v_movs.
s_nop(n) delays n+1 cycles, plus +4 extra for n in [11, 22].
Exec delta = n + 4 (baseline) + 4 (if 11 <= n <= 22)."""
def _test(self, n, expected):
_, execd = get_deltas([v_mov_b32_e32(v[0], 1.0), s_nop(n), v_mov_b32_e32(v[1], 2.0)])
if isinstance(expected, list): self.assertIn(execd[0], expected)
else: self.assertEqual(execd, [expected])
def test_snop_0(self): self._test(0, 4)
def test_snop_1(self): self._test(1, 5)
def test_snop_2(self): self._test(2, 6)
def test_snop_3(self): self._test(3, 7)
def test_snop_4(self): self._test(4, 8)
def test_snop_5(self): self._test(5, 9)
def test_snop_6(self): self._test(6, 10)
def test_snop_7(self): self._test(7, 11)
def test_snop_10(self): self._test(10, 14)
def test_snop_11(self): self._test(11, 19) # +4 extra starts here
def test_snop_15(self): self._test(15, 23)
def test_snop_22(self): self._test(22, 30) # +4 extra ends here
def test_snop_23(self): self._test(23, 27)
def test_snop_31(self): self._test(31, 35)
def test_snop_32(self): self._test(32, 36)
def test_snop_63(self): self._test(63, [67, 71])
class TestVALUExecWithNop(unittest.TestCase):
"""Single VALU followed by s_nop - measures VALUINST to ALUEXEC delay."""
def _get_delay(self, instrs, nops=16):
deltas = get_timing_deltas(run_sqtt(instrs, nops=nops))
time, valu_time, exec_time = 0, None, None
for ptype, delta in deltas:
time += delta
if ptype == 'VALUINST' and valu_time is None: valu_time = time
if ptype == 'ALUEXEC' and exec_time is None: exec_time = time
return exec_time - valu_time
# Boundary: s_nop(0-3) = 6 cycles, s_nop(4+) = 10 cycles
def test_nop0(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(0)]), 6)
def test_nop1(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(1)]), 6)
def test_nop2(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(2)]), 6)
def test_nop3(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(3)]), 6)
def test_nop4(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(4)]), 10)
def test_nop5(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(5)]), 10)
def test_nop6(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(6)]), 10)
def test_nop7(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(7)]), 10)
def test_nop8(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(8)]), 10)
def test_nop9(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(9)]), 10)
def test_nop10(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(10)]), 10)
# No nop = slow path, one s_nop(0) padding = fast path
def test_no_padding(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0)], nops=0), 10)
def test_one_padding(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0)], nops=1), 6)
# Multiple s_nop(0)s don't accumulate - still fast path
def test_nop0_x2(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(0), s_nop(0)]), 6)
# First nop determines path: s_nop(0) then s_nop(4) = fast, s_nop(4) then s_nop(0) = slow
def test_nop0_nop4(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(0), s_nop(4)]), 6)
def test_nop4_nop0(self): self.assertEqual(self._get_delay([v_mov_b32_e32(v[0], 1.0), s_nop(4), s_nop(0)]), 10)
class TestDelayALU(unittest.TestCase):
"""s_delay_alu behavior - helps understand hardware pipeline latencies.
s_delay_alu(simm16) where simm16 encodes:
instid0[3:0] = dependency on VALU N instructions back (1-4), 0=none
skip[6:4] = skip count for second dependency
instid1[10:7] = second dependency
Key insight: s_delay_alu tells hardware to wait for a previous VALU to complete.
The hardware determines how many cycles to stall based on pipeline state.
"""
def _exec_delta(self, instrs):
"""Return exec delta for last instruction."""
_, execd = get_deltas(instrs)
return execd[-1] if execd else None
# Direct dependency (producer -> consumer), instid0=1 means "wait for VALU 1 back"
def test_direct_no_delay(self):
# Without s_delay_alu: 6 cycles
self.assertEqual(self._exec_delta([v_mov_b32_e32(v[0], 1.0), v_add_f32_e32(v[1], v[0], v[0])]), 6)
def test_direct_delay1(self):
# With s_delay_alu(instid0=1): 7-8 cycles (+1 from the delay instruction)
self.assertIn(self._exec_delta([v_mov_b32_e32(v[0], 1.0), s_delay_alu(simm16=1), v_add_f32_e32(v[1], v[0], v[0])]), [7, 8])
def test_direct_delay2(self):
# instid0=2 doesn't apply (only 1 VALU back), so no extra delay
self.assertEqual(self._exec_delta([v_mov_b32_e32(v[0], 1.0), s_delay_alu(simm16=2), v_add_f32_e32(v[1], v[0], v[0])]), 6)
def test_direct_delay3(self):
self.assertEqual(self._exec_delta([v_mov_b32_e32(v[0], 1.0), s_delay_alu(simm16=3), v_add_f32_e32(v[1], v[0], v[0])]), 6)
def test_direct_delay4(self):
self.assertEqual(self._exec_delta([v_mov_b32_e32(v[0], 1.0), s_delay_alu(simm16=4), v_add_f32_e32(v[1], v[0], v[0])]), 6)
# With 1 independent instruction between producer and consumer
def test_gap1_delay1(self):
# instid0=1 waits for the independent instruction (not the producer)
instrs = [v_mov_b32_e32(v[0], 1.0), v_mov_b32_e32(v[5], 5.0), s_delay_alu(simm16=1), v_add_f32_e32(v[1], v[0], v[0])]
self.assertEqual(self._exec_delta(instrs), 8)
def test_gap1_delay2(self):
# instid0=2 waits for the producer (2 VALUs back)
instrs = [v_mov_b32_e32(v[0], 1.0), v_mov_b32_e32(v[5], 5.0), s_delay_alu(simm16=2), v_add_f32_e32(v[1], v[0], v[0])]
self.assertIn(self._exec_delta(instrs), [6, 7])
def test_gap1_delay3(self):
# instid0=3 doesn't apply (only 2 VALUs back)
instrs = [v_mov_b32_e32(v[0], 1.0), v_mov_b32_e32(v[5], 5.0), s_delay_alu(simm16=3), v_add_f32_e32(v[1], v[0], v[0])]
self.assertEqual(self._exec_delta(instrs), 5)
# With 2 independent instructions between
def test_gap2_delay1(self):
instrs = [v_mov_b32_e32(v[0], 1.0), v_mov_b32_e32(v[5], 5.0), v_mov_b32_e32(v[6], 6.0),
s_delay_alu(simm16=1), v_add_f32_e32(v[1], v[0], v[0])]
self.assertEqual(self._exec_delta(instrs), 7)
def test_gap2_delay2(self):
instrs = [v_mov_b32_e32(v[0], 1.0), v_mov_b32_e32(v[5], 5.0), v_mov_b32_e32(v[6], 6.0),
s_delay_alu(simm16=2), v_add_f32_e32(v[1], v[0], v[0])]
self.assertEqual(self._exec_delta(instrs), 7)
def test_gap2_delay3(self):
# instid0=3 waits for the producer (3 VALUs back)
instrs = [v_mov_b32_e32(v[0], 1.0), v_mov_b32_e32(v[5], 5.0), v_mov_b32_e32(v[6], 6.0),
s_delay_alu(simm16=3), v_add_f32_e32(v[1], v[0], v[0])]
self.assertIn(self._exec_delta(instrs), [5, 6])
def test_gap2_delay4(self):
instrs = [v_mov_b32_e32(v[0], 1.0), v_mov_b32_e32(v[5], 5.0), v_mov_b32_e32(v[6], 6.0),
s_delay_alu(simm16=4), v_add_f32_e32(v[1], v[0], v[0])]
self.assertEqual(self._exec_delta(instrs), 4)
class TestNopTimingSensitivity(unittest.TestCase):
"""Forwarding behavior has 128-cycle periodicity.
Hardware observation: when nop_cycles % 128 is in [72, 75], chain_6 gets 5 forwards
instead of 4. This 4-cycle window repeats every 128 cycles, suggesting alignment
with some hardware scheduling period (possibly wave scheduler or cache).
Windows found: nop 72-75, 200-203, 328-331, 456-459, ...
"""
def _chain6_fwd_count(self, nop_size):
"""Count initial consecutive forwards for a 6-instruction chain after s_nop(n)."""
instrs = [s_nop(nop_size), v_mov_b32_e32(v[99], 1.0)]
instrs += [v_mov_b32_e32(v[0], 1.0)]
for i in range(1, 6):
instrs += [v_mov_b32_e32(v[i], v[i-1])]
_, execd = get_deltas(instrs)
chain_deltas = execd[1:]
fwd_count = 0
for d in chain_deltas:
if d == 5: fwd_count += 1
else: break
return fwd_count
# Normal case: 4 forwards
def test_nop71(self): self.assertEqual(self._chain6_fwd_count(71), 4)
def test_nop76(self): self.assertEqual(self._chain6_fwd_count(76), 4)
def test_nop199(self): self.assertEqual(self._chain6_fwd_count(199), 4)
def test_nop204(self): self.assertEqual(self._chain6_fwd_count(204), 4)
# Anomaly window at nop % 128 == 72-75: 5 forwards on RDNA3, 4 on other variants
def test_nop72(self): self.assertIn(self._chain6_fwd_count(72), [4, 5])
def test_nop75(self): self.assertIn(self._chain6_fwd_count(75), [4, 5])
def test_nop200(self): self.assertIn(self._chain6_fwd_count(200), [4, 5])
def test_nop203(self): self.assertIn(self._chain6_fwd_count(203), [4, 5])
def test_nop328(self): self.assertIn(self._chain6_fwd_count(328), [4, 5])
def test_nop331(self): self.assertIn(self._chain6_fwd_count(331), [4, 5])
if __name__ == "__main__":
unittest.main()
@@ -1,206 +0,0 @@
#!/usr/bin/env python3
"""Tests for SQTT packet decoding using real captured examples."""
import pickle, unittest, ctypes
from pathlib import Path
from tinygrad.helpers import DEBUG, colored
from tinygrad.runtime.autogen import rocprof
from tinygrad.runtime.support.elf import elf_loader
from extra.assembly.amd.asm import detect_format
from extra.assembly.amd.autogen.rdna3.ins import SOPP
from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
from extra.assembly.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVEEND, INST, VALUINST, IMMEDIATE, IMMEDIATE_MASK,
ALUEXEC, VMEMEXEC, PACKET_TYPES, InstOp, AluSrc, MemSrc)
EXAMPLES_DIR = Path(__file__).parent.parent.parent.parent / "sqtt/examples"
# INST ops for non-traced SIMDs (excluded from instruction count)
OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LDS_STORE_64, InstOp.OTHER_LDS_STORE_128,
InstOp.OTHER_FLAT_LOAD, InstOp.OTHER_FLAT_STORE, InstOp.OTHER_FLAT_STORE_64, InstOp.OTHER_FLAT_STORE_96,
InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
InstOp.OTHER_GLOBAL_STORE_VADDR_128}
PACKET_COLORS = {
"INST": "WHITE", "VALUINST": "BLACK", "VMEMEXEC": "yellow", "ALUEXEC": "yellow",
"IMMEDIATE": "YELLOW", "IMMEDIATE_MASK": "YELLOW", "WAVERDY": "cyan", "WAVEALLOC": "cyan",
"WAVEEND": "blue", "WAVESTART": "blue", "PERF": "magenta", "EVENT": "red", "EVENT_BIG": "red",
"REG": "green", "LAYOUT_HEADER": "white", "SNAPSHOT": "white", "UTILCTR": "green",
}
def format_packet(p, time_offset: int = 0) -> str:
name, cycle = type(p).__name__, p._time - time_offset
if isinstance(p, INST):
op_name = p.op.name if isinstance(p.op, InstOp) else f"0x{p.op:02x}"
fields = f"wave={p.wave} op={op_name}" + (" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "")
elif isinstance(p, VALUINST): fields = f"wave={p.wave}" + (" flag" if p.flag else "")
elif isinstance(p, ALUEXEC): fields = f"src={p.src.name if isinstance(p.src, AluSrc) else p.src}"
elif isinstance(p, VMEMEXEC): fields = f"src={p.src.name if isinstance(p.src, MemSrc) else p.src}"
elif isinstance(p, (WAVESTART, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
elif hasattr(p, '_values'):
fields = " ".join(f"{k}=0x{v:x}" if k in {'snap', 'val32'} else f"{k}={v}"
for k, v in p._values.items() if not k.startswith('_') and k != 'delta')
else: fields = ""
return f"{cycle:8}: {colored(f'{name:18}', PACKET_COLORS.get(name, 'white'))} {fields}"
def print_packets(packets: list) -> None:
skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"}
time_offset = packets[0]._time if packets else 0
for p in packets:
if type(p).__name__ not in skip: print(format_packet(p, time_offset))
# ═══════════════════════════════════════════════════════════════════════════════
# ROCPROF DECODER
# ═══════════════════════════════════════════════════════════════════════════════
def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int):
"""Run rocprof decoder on SQTT blobs, returning raw occupancy and instruction records."""
image, sections, _ = elf_loader(lib)
text = next((sh for sh in sections if sh.name == ".text"), None)
assert text is not None, "no .text section found"
text_off, text_size = text.header.sh_addr, text.header.sh_size
blob_iter, current_blob = iter(blobs), [None]
occupancy_records: list[tuple[int, int, int, int, bool]] = [] # (wave_id, simd, cu, time, is_start)
wave_insts: list[list[tuple[int, int]]] = [] # per-wave list of (time, stall)
@rocprof.rocprof_trace_decoder_se_data_callback_t
def copy_cb(buf, buf_size, _):
blob = next(blob_iter, None)
if blob is None: return 0
current_blob[0] = (ctypes.c_ubyte * len(blob)).from_buffer_copy(blob)
buf[0] = ctypes.cast(current_blob[0], ctypes.POINTER(ctypes.c_ubyte))
buf_size[0] = len(current_blob[0])
return len(current_blob[0])
@rocprof.rocprof_trace_decoder_trace_callback_t
def trace_cb(record_type, events_ptr, n, _):
if record_type == rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr):
occupancy_records.append((ev.wave_id, ev.simd, ev.cu, ev.time, ev.start))
elif record_type == rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr):
if ev.instructions_size > 0:
sz = ev.instructions_size * ctypes.sizeof(rocprof.rocprofiler_thread_trace_decoder_inst_t)
insts_blob = bytearray(sz)
ctypes.memmove((ctypes.c_char * sz).from_buffer(insts_blob), ev.instructions_array, sz)
insts = list((rocprof.rocprofiler_thread_trace_decoder_inst_t * ev.instructions_size).from_buffer(insts_blob))
wave_insts.append([(inst.time, inst.stall) for inst in insts])
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
@rocprof.rocprof_trace_decoder_isa_callback_t
def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
offset = pc.address - base
if offset < text_off or offset >= text_off + text_size:
mem_size_ptr[0] = 0
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
try:
inst = detect_format(data := image[offset:]).from_bytes(data)
mem_size_ptr[0] = inst._size()
except (ValueError, AssertionError):
mem_size_ptr[0] = 0
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
if isinstance(inst, SOPP) and inst.op == SOPPOp.S_ENDPGM: mem_size_ptr[0] = 0
# rocprof parses instruction string to determine type; v_nop works for all
if (max_sz := size_ptr[0]) == 0: return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES
ctypes.memmove(instr_ptr, b"v_nop", min(5, max_sz - 1))
size_ptr[0] = min(5, max_sz - 1)
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
return occupancy_records, wave_insts
class TestSQTTExamples(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.examples = {}
for pkl_path in sorted(EXAMPLES_DIR.glob("*.pkl")):
with open(pkl_path, "rb") as f:
data = pickle.load(f)
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
prg = next((e for e in data if type(e).__name__ == "ProfileProgramEvent"), None)
if sqtt_events and prg:
cls.examples[pkl_path.stem] = (sqtt_events, prg.lib, prg.base)
def test_examples_loaded(self):
self.assertGreater(len(self.examples), 0, "no example files found")
def test_decode_all_examples(self):
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
packets = decode(event.blob)
if DEBUG >= 2: print(f"\n=== {name} event {i} ==="); print_packets(packets)
self.assertGreater(len(packets), 0, f"no packets decoded from {name} event {i}")
self.assertIsInstance(packets[0], LAYOUT_HEADER, f"first packet should be LAYOUT_HEADER in {name}")
def test_packet_types_valid(self):
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
for pkt in decode(event.blob):
self.assertIn(type(pkt), PACKET_TYPES, f"unknown packet type {type(pkt)} in {name}")
def test_wave_lifecycle(self):
for name, (events, *_) in self.examples.items():
if "empty" in name: continue
with self.subTest(example=name):
all_packets = [p for e in events for p in decode(e.blob)]
self.assertGreater(len([p for p in all_packets if isinstance(p, WAVESTART)]), 0, f"no WAVESTART in {name}")
self.assertGreater(len([p for p in all_packets if isinstance(p, WAVEEND)]), 0, f"no WAVEEND in {name}")
def test_time_monotonic(self):
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
times = [p._time for p in decode(event.blob)]
self.assertEqual(times, sorted(times), f"timestamps not monotonic in {name}")
def test_gemm_has_instructions(self):
for name, (events, *_) in self.examples.items():
if "gemm" not in name: continue
with self.subTest(example=name):
all_packets = [p for e in events for p in decode(e.blob)]
self.assertGreater(len([p for p in all_packets if isinstance(p, INST)]), 0, f"no INST packets in {name}")
def test_rocprof_wave_times_match(self):
"""Wave start/end times must match rocprof exactly."""
for name, (events, lib, base) in self.examples.items():
with self.subTest(example=name):
occupancy, _ = run_rocprof_decoder([e.blob for e in events], lib, base)
# extract from rocprof occupancy records
roc_starts: dict[tuple[int, int, int], int] = {}
roc_waves: list[tuple[int, int]] = []
for wave_id, simd, cu, time, is_start in occupancy:
key = (wave_id, simd, cu)
if is_start: roc_starts[key] = time
elif key in roc_starts: roc_waves.append((roc_starts.pop(key), time))
# extract from our decoder
our_waves: list[tuple[int, int]] = []
for event in events:
packets = decode(event.blob)
wave_starts: dict[tuple[int, int, int], int] = {}
for p in packets:
if isinstance(p, WAVESTART): wave_starts[(p.wave, p.simd, p.cu)] = p._time
elif isinstance(p, WAVEEND) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
our_waves.append((wave_starts[key], p._time))
self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
def test_rocprof_inst_times_match(self):
"""Instruction times must match rocprof exactly (excluding s_endpgm)."""
for name, (events, lib, base) in self.examples.items():
with self.subTest(example=name):
_, wave_insts = run_rocprof_decoder([e.blob for e in events], lib, base)
# skip last inst per wave (s_endpgm) - it needs special handling (time + duration instead of time + stall)
roc_insts = [time + stall for insts in wave_insts for time, stall in insts[:-1]]
# extract from our decoder
our_insts: list[int] = []
for event in events:
for p in decode(event.blob):
if isinstance(p, INST) and p.op not in OTHER_SIMD_OPS: our_insts.append(p._time)
elif isinstance(p, VALUINST): our_insts.append(p._time)
elif isinstance(p, IMMEDIATE): our_insts.append(p._time)
elif isinstance(p, IMMEDIATE_MASK):
for _ in range(bin(p.mask).count('1')): our_insts.append(p._time)
self.assertEqual(sorted(our_insts), sorted(roc_insts), f"instruction times mismatch in {name}")
if __name__ == "__main__":
unittest.main()
-463
View File
@@ -1,463 +0,0 @@
#!/usr/bin/env python3
"""Hardware tests for SQTT decoder - validates decoding of real SQTT streams.
Run with: python -m pytest extra/assembly/amd/test/test_sqtt_hw.py -v -s
Requires AMD GPU with SQTT support.
For pretty trace output: DEBUG=2 python -m pytest extra/assembly/amd/test/test_sqtt_hw.py -v -s
"""
import os
os.environ["SQTT"] = "1"
os.environ["PROFILE"] = "1"
os.environ["SQTT_ITRACE_SE_MASK"] = "1" # Enable instruction tracing on SE0
os.environ["SQTT_LIMIT_SE"] = "2" # Force work to traced SE only
import unittest
from tinygrad.helpers import DEBUG, colored
from tinygrad.device import Device
from tinygrad.runtime.ops_amd import AMDProgram, ProfileSQTTEvent
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from extra.assembly.amd.autogen.rdna3.ins import v_mov_b32_e32, v_add_f32_e32, v_mul_f32_e32, s_mov_b32, s_add_u32, s_nop, s_waitcnt, s_endpgm
from extra.assembly.amd.dsl import v, s
from extra.assembly.amd.sqtt import decode, LAYOUT_HEADER, WAVESTART, WAVEEND, INST, VALUINST, ALUEXEC, VMEMEXEC, InstOp, AluSrc, MemSrc
dev = Device["AMD"]
# ═══════════════════════════════════════════════════════════════════════════════
# PRETTY PRINTING
# ═══════════════════════════════════════════════════════════════════════════════
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",
"TS_DELTA_SHORT": "BLACK", "NOP": "BLACK", "TS_WAVE_STATE": "BLACK",
"SNAPSHOT": "white", "TS_DELTA_OR_MARK": "BLACK",
"TS_DELTA_S8_W3": "BLACK", "TS_DELTA_S5_W2": "BLACK", "TS_DELTA_S5_W3": "BLACK",
"UTILCTR": "green",
}
def format_packet(p, last_time: int = 0, time_offset: int = 0) -> str:
"""Format a packet for pretty printing."""
name = type(p).__name__
color = PACKET_COLORS.get(name, "white")
fields = []
if isinstance(p, INST):
op = p.op
op_name = op.name if isinstance(op, InstOp) else f"0x{op:02x}"
fields = [f"wave={p.wave}", f"op={op_name}"]
if p.flag1: fields.append("flag1")
if p.flag2: fields.append("flag2")
elif isinstance(p, VALUINST):
fields = [f"wave={p.wave}"]
if p.flag: fields.append("flag")
elif isinstance(p, ALUEXEC):
src_name = p.src.name if isinstance(p.src, AluSrc) else f"{p.src}"
fields = [f"src={src_name}"]
elif isinstance(p, VMEMEXEC):
src_name = p.src.name if isinstance(p.src, MemSrc) else f"{p.src}"
fields = [f"src={src_name}"]
elif isinstance(p, WAVESTART):
fields = [f"wave={p.wave}", f"simd={p.simd}", f"cu={p.cu}"]
elif isinstance(p, WAVEEND):
fields = [f"wave={p.wave}", f"simd={p.simd}", f"cu={p.cu}"]
elif hasattr(p, '_values'):
# Format hex fields appropriately
hex_fields = {'snap', 'val32'}
fields = [f"{k}=0x{v:x}" if k in hex_fields else f"{k}={v}" for k, v in p._values.items() if not k.startswith('_') and k != 'delta']
return colored(f"{name:18s}", color) + " " + ", ".join(fields)
def get_wave_packets(packets: list) -> list:
"""Extract packets from WAVESTART to WAVEEND, filtering pure timing packets."""
skip_types = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3"}
result = []
in_wave = False
for p in packets:
name = type(p).__name__
if isinstance(p, WAVESTART):
in_wave = True
if in_wave and name not in skip_types:
result.append(p)
if isinstance(p, WAVEEND):
in_wave = False
return result
def print_wave_trace(packets: list) -> None:
"""Print packets from WAVESTART to WAVEEND with normalized time."""
wave_packets = get_wave_packets(packets)
if not wave_packets:
return
time_offset = wave_packets[0]._time
last_time = time_offset
for p in wave_packets:
print(format_packet(p, last_time, time_offset))
last_time = p._time
def print_blobs(blobs: list[bytes], wave_only: bool = True) -> None:
"""Print traces for all blobs. wave_only=True filters to WAVESTART..WAVEEND only."""
for i, blob in enumerate(blobs):
packets = decode(blob)
print(f"\n--- Blob {i}: {len(blob)} bytes, {len(packets)} packets ---")
if wave_only:
print_wave_trace(packets)
else:
print_all_packets(packets)
def print_all_packets(packets: list) -> None:
"""Print all packets, filtering out pure timing packets."""
skip_types = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3"}
if not packets: return
time_offset = packets[0]._time
last_time = time_offset
for p in packets:
if type(p).__name__ not in skip_types:
print(format_packet(p, last_time, time_offset))
last_time = p._time
# ═══════════════════════════════════════════════════════════════════════════════
# ASSEMBLY HELPERS
# ═══════════════════════════════════════════════════════════════════════════════
def assemble(instructions: list) -> bytes:
return b''.join(inst.to_bytes() for inst in instructions)
def wrap_with_nops(instructions: list, nops=16) -> list:
"""Add epilogue for clean SQTT timing.
Need enough NOPs to cover long-latency ops (DP: 42 cycles, WMMA: 47 cycles).
With 64 NOPs, the IMMEDIATE phase extends to cover these completions.
"""
return instructions + [s_nop(0)]*nops + [s_endpgm()]
def compile_asm_sqtt(instructions: list, alu_only: bool = False) -> AMDProgram:
"""Compile instructions to an AMDProgram for SQTT tracing.
Args:
instructions: List of instructions to compile
alu_only: If True, use minimal kernel config with no kernargs/LDS/scratch
Returns:
Compiled AMDProgram ready to run
"""
compiler = HIPCompiler(dev.arch)
# Add NOPs before s_endpgm to flush pipeline and get clean timing
code = assemble(instructions)
byte_str = ', '.join(f'0x{b:02x}' for b in code)
if alu_only:
asm_src = f""".text
.globl test
.p2align 8
.type test,@function
test:
.byte {byte_str}
.rodata
.p2align 6
.amdhsa_kernel test
# basic memory
.amdhsa_group_segment_fixed_size 0
.amdhsa_private_segment_fixed_size 0
.amdhsa_kernarg_size 32
.amdhsa_enable_private_segment 0
# register usage
.amdhsa_next_free_vgpr 64
.amdhsa_next_free_sgpr 8
# RSRC1
.amdhsa_wavefront_size32 1
.amdhsa_memory_ordered 1
.amdhsa_forward_progress 1
# this is key
.amdhsa_workgroup_processor_mode 0
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.version:
- 1
- 0
amdhsa.kernels:
- .name: test
.symbol: test.kd
.kernarg_segment_size: 0
.group_segment_fixed_size: 0
.private_segment_fixed_size: 0
.kernarg_segment_align: 8
.wavefront_size: 32
.sgpr_count: 8
.vgpr_count: 64
.max_flat_workgroup_size: 1024
...
.end_amdgpu_metadata
"""
else:
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 8
.amdhsa_next_free_sgpr 16
.amdhsa_wavefront_size32 1
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_kernarg_size 8
.amdhsa_group_segment_fixed_size 0
.amdhsa_private_segment_fixed_size 0
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.version:
- 1
- 0
amdhsa.kernels:
- .name: test
.symbol: test.kd
.kernarg_segment_size: 8
.group_segment_fixed_size: 0
.private_segment_fixed_size: 0
.kernarg_segment_align: 8
.wavefront_size: 32
.sgpr_count: 16
.vgpr_count: 8
.max_flat_workgroup_size: 1024
...
.end_amdgpu_metadata
"""
lib = compiler.compile(asm_src)
return AMDProgram(dev, "test", lib)
def run_asm_sqtt(instructions: list, n_lanes: int = 1, alu_only: bool = False) -> list[bytes]:
"""Compile and run instructions on AMD hardware, return SQTT blobs.
Args:
instructions: List of instructions to run
n_lanes: Number of lanes to use
alu_only: If True, use minimal kernel config with no kernargs/LDS/scratch
"""
prg = compile_asm_sqtt(instructions, alu_only=alu_only)
return run_prg_sqtt(prg, n_lanes=n_lanes, alu_only=alu_only)
def run_prg_sqtt(prg: AMDProgram, n_lanes: int = 1, alu_only: bool = False) -> list[bytes]:
"""Run a compiled AMDProgram and return SQTT blobs.
Args:
prg: Compiled AMDProgram to run
n_lanes: Number of lanes to use
alu_only: If True, don't allocate kernarg buffer
"""
dev.profile_events.clear()
if alu_only:
prg(global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
else:
out_gpu = dev.allocator.alloc(2048)
prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
return [ev.blob for ev in dev.profile_events if isinstance(ev, ProfileSQTTEvent)]
def run_prg_sqtt_batch(prg: AMDProgram, n_runs: int, n_lanes: int = 1) -> list[bytes]:
"""Run a compiled AMDProgram N times in a single queue submission and return SQTT blobs.
This builds one queue with N kernel executions, submits it once, and collects SQTT.
All N runs are captured in the same SQTT trace, reducing startup jitter.
Args:
prg: Compiled AMDProgram to run
n_runs: Number of times to execute the kernel in the queue
n_lanes: Number of lanes to use
Returns:
List of SQTT blobs (one per shader engine)
"""
from typing import cast
from tinygrad.runtime.ops_amd import AMDComputeQueue, SQTT_ITRACE_SE_MASK
from tinygrad.device import Compiled
import struct
dev.profile_events.clear()
# Build queue with sqtt_start, N kernel executions, sqtt_stop
kernargs = prg.fill_kernargs([], ())
q = cast(AMDComputeQueue, dev.hw_compute_queue_t())
q.wait(dev.timeline_signal, dev.timeline_value - 1).memory_barrier()
q.sqtt_start(dev.sqtt_buffers)
# Execute kernel N times
for _ in range(n_runs):
q.exec(prg, kernargs, (1, 1, 1), (n_lanes, 1, 1))
q.sqtt_stop(dev.sqtt_wptrs)
q.signal(dev.timeline_signal, dev.next_timeline())
q.submit(dev)
dev.synchronize()
# Collect SQTT blobs
blobs = []
for se, buf in enumerate(dev.sqtt_buffers):
wptr = (dev.sqtt_wptrs.cpu_view().view(fmt='I')[se] & 0x1FFFFFFF) * 32
if dev.target[:2] == (11, 0): wptr -= ((buf.va_addr // 32) & 0x1FFFFFFF) * 32
if wptr > 0 and wptr <= buf.size:
dev.allocator._copyout(sqtt_mv:=memoryview(bytearray(wptr)), buf)
resbuf = (struct.pack('<Q', 0x11 | (4 << 13) | (0xf << 16) | (se << 24)) + bytes(sqtt_mv)) if dev.target[0] == 9 else bytes(sqtt_mv)
blobs.append(resbuf)
return blobs
def decode_all_blobs(blobs: list[bytes]) -> list:
"""Decode all blobs and combine packets."""
all_packets = []
for blob in blobs:
all_packets.extend(decode(blob))
return all_packets
def get_inst_ops(packets: list, traced_simd: int | None = None) -> set:
"""Extract all InstOp values from INST packets within WAVESTART..WAVEEND on traced SIMD."""
ops = set()
in_wave = False
for p in packets:
if isinstance(p, WAVESTART):
in_wave = traced_simd is None or p.simd == traced_simd
if in_wave and isinstance(p, INST):
ops.add(p.op if isinstance(p.op, int) else p.op.value)
if isinstance(p, WAVEEND):
in_wave = False
return ops
def count_valuinst(packets: list, traced_simd: int | None = None) -> int:
"""Count VALUINST packets within WAVESTART..WAVEEND on traced SIMD."""
count = 0
in_wave = False
for p in packets:
if isinstance(p, WAVESTART):
in_wave = traced_simd is None or p.simd == traced_simd
if in_wave and isinstance(p, VALUINST):
count += 1
if isinstance(p, WAVEEND):
in_wave = False
return count
# ═══════════════════════════════════════════════════════════════════════════════
# TESTS
# ═══════════════════════════════════════════════════════════════════════════════
@unittest.skipIf(not hasattr(dev, 'profile_events'), "AMD device required")
class TestSQTTDecode(unittest.TestCase):
"""Test SQTT decoder with real hardware traces."""
def test_basic_structure(self):
"""Verify basic SQTT stream structure: LAYOUT_HEADER, WAVESTART, instructions, WAVEEND."""
blobs = run_asm_sqtt([v_mov_b32_e32(v[0], 0)])
self.assertGreater(len(blobs), 0, "No SQTT data captured")
packets = decode_all_blobs(blobs)
self.assertGreater(len(packets), 0, "No packets decoded")
self.assertGreater(len([p for p in packets if isinstance(p, LAYOUT_HEADER)]), 0, "No LAYOUT_HEADER packets")
self.assertGreater(len([p for p in packets if isinstance(p, WAVESTART)]), 0, "No WAVESTART packets")
self.assertGreater(len([p for p in packets if isinstance(p, WAVEEND)]), 0, "No WAVEEND packets")
if DEBUG >= 2:
print("\n=== Basic structure trace ===")
print_trace(packets)
def test_valu_instructions(self):
"""Verify VALU instructions produce INST or VALUINST packets."""
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]),
v_add_f32_e32(v[3], v[2], v[1]),
v_mul_f32_e32(v[4], v[2], v[3]),
]
blobs = run_asm_sqtt(instructions)
self.assertGreater(len(blobs), 0, "No SQTT data captured")
packets = decode_all_blobs(blobs)
inst_packets = [p for p in packets if isinstance(p, (INST, VALUINST))]
self.assertGreater(len(inst_packets), 0, "No INST/VALUINST packets for VALU instructions")
if DEBUG >= 2:
print("\n=== VALU instructions trace ===")
print_trace(packets)
def test_salu_instructions(self):
"""Verify SALU instructions produce appropriate packets."""
instructions = [
s_mov_b32(s[0], 0),
s_mov_b32(s[1], 1),
s_add_u32(s[2], s[0], s[1]),
s_add_u32(s[3], s[2], s[1]),
s_nop(0),
]
blobs = run_asm_sqtt(instructions)
self.assertGreater(len(blobs), 0, "No SQTT data captured")
packets = decode_all_blobs(blobs)
if DEBUG >= 2:
print("\n=== SALU instructions trace ===")
print_trace(packets)
def test_timing_increases(self):
"""Verify time increases monotonically through packets within each blob."""
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]),
v_mul_f32_e32(v[3], v[2], v[1]),
]
blobs = run_asm_sqtt(instructions)
self.assertGreater(len(blobs), 0, "No SQTT data captured")
for blob in blobs:
packets = decode(blob)
prev_time = 0
for p in packets:
self.assertGreaterEqual(p._time, prev_time, f"Time decreased: {prev_time} -> {p._time}")
prev_time = p._time
def test_wave_id_consistency(self):
"""Verify wave IDs are consistent between WAVESTART/WAVEEND."""
blobs = run_asm_sqtt([v_mov_b32_e32(v[0], 0)])
self.assertGreater(len(blobs), 0, "No SQTT data captured")
packets = decode_all_blobs(blobs)
wavestarts = [p for p in packets if isinstance(p, WAVESTART)]
waveends = [p for p in packets if isinstance(p, WAVEEND)]
if wavestarts and waveends:
start_waves = {p.wave for p in wavestarts}
end_waves = {p.wave for p in waveends}
self.assertTrue(start_waves & end_waves, "No matching wave IDs between WAVESTART and WAVEEND")
def test_nop_sequence(self):
"""Test a sequence of NOP instructions."""
blobs = run_asm_sqtt([s_nop(0), s_nop(0), s_nop(0)])
self.assertGreater(len(blobs), 0, "No SQTT data captured")
packets = decode_all_blobs(blobs)
self.assertGreater(len(packets), 0, "No packets decoded")
if DEBUG >= 2:
print("\n=== NOP sequence trace ===")
print_trace(packets, filter_timing=False)
if __name__ == "__main__":
unittest.main()
@@ -1,35 +0,0 @@
import os
os.environ["SQTT"] = "1"
os.environ["PROFILE"] = "1"
os.environ["SQTT_LIMIT_SE"] = "2"
os.environ["SQTT_SIMD_SEL"] = "0"
os.environ["SQTT_TOKEN_EXCLUDE"] = "3784" # Exclude WAVERDY, REG, EVENT, UTILCTR, WAVEALLOC, PERF
import unittest
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.sqtt import decode
from extra.assembly.amd.test.test_sqtt_hw import compile_asm_sqtt, run_prg_sqtt_batch, format_packet
from extra.assembly.amd.test.test_sqtt_compare import filter_noise_packets
from tinygrad.uop.ops import UOp
from tinygrad.engine.realize import get_runner
class SQTTMultiwave(unittest.TestCase):
def test_simple_multiwave(self):
ins = [
s_barrier(),
v_mov_b32_e32(v[0], v[1]),
s_nop(0),
s_nop(100),
s_endpgm(),
]
#prg = get_runner("AMD", UOp.sink())._prg
prg = compile_asm_sqtt(ins, alu_only=True)
print(prg)
blobs = run_prg_sqtt_batch(prg, n_runs=1, n_lanes=32*16)
for blob in blobs:
packets = decode(blob)
for p in filter_noise_packets(packets):
print(f" {p._time:8d}: {format_packet(p)}")
if __name__ == "__main__":
unittest.main()
-204
View File
@@ -1,204 +0,0 @@
#!/usr/bin/env python3
"""Tests validating SQTT packet definitions against the reference implementation.
Verifies that:
1. Encoding patterns produce the correct STATE_TO_OPCODE table
2. Packet sizes (derived from fields) match expected budget values
3. Field extractions match attempt_sqtt_parse.py
"""
import unittest
from extra.assembly.amd.sqtt import (
VALUINST, VMEMEXEC, ALUEXEC, IMMEDIATE, IMMEDIATE_MASK, WAVERDY,
WAVEEND, WAVESTART, PERF, TS_WAVE_STATE, EVENT, EVENT_BIG, REG, SNAPSHOT,
TS_DELTA_OR_MARK, LAYOUT_HEADER, INST, UTILCTR, TS_DELTA_SHORT, NOP,
TS_DELTA_S8_W3, TS_DELTA_S5_W2, TS_DELTA_S5_W3, WAVEALLOC,
decode, encode, OPCODE_TO_CLASS, STATE_TO_OPCODE, PACKET_TYPES, BUDGET,
AluSrc, MemSrc, InstOp
)
# Reference table from rocprof trace decoder (attempt_sqtt_parse.py)
REFERENCE_STATE_TABLE = bytes([
0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
0x10, 0x12, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
0x10, 0x13, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
])
# Reference opcode -> name mapping (old opcode values from rocprof)
OLD_OPCODE_TO_NAME = {
0x01: 'VALUINST', 0x02: 'VMEMEXEC', 0x03: 'ALUEXEC', 0x04: 'IMMEDIATE',
0x05: 'IMMEDIATE_MASK', 0x06: 'WAVERDY', 0x07: 'TS_DELTA_S8_W3',
0x08: 'WAVEEND', 0x09: 'WAVESTART', 0x0A: 'TS_DELTA_S5_W2',
0x0B: 'WAVEALLOC', 0x0C: 'TS_DELTA_S5_W3', 0x0D: 'PERF',
0x0F: 'TS_DELTA_SHORT', 0x10: 'NOP', 0x11: 'TS_WAVE_STATE',
0x12: 'EVENT', 0x13: 'EVENT_BIG', 0x14: 'REG', 0x15: 'SNAPSHOT',
0x16: 'TS_DELTA_OR_MARK', 0x17: 'LAYOUT_HEADER', 0x18: 'INST',
0x19: 'UTILCTR', 0x00: 'NOP',
}
# Reference budget values (nibbles for NEXT packet) from rocprof
REFERENCE_BUDGET_NIBBLES = {
'VALUINST': 3, 'VMEMEXEC': 2, 'ALUEXEC': 2, 'IMMEDIATE': 3,
'IMMEDIATE_MASK': 6, 'WAVERDY': 6, 'TS_DELTA_S8_W3': 16,
'WAVEEND': 5, 'WAVESTART': 8, 'TS_DELTA_S5_W2': 12,
'WAVEALLOC': 5, 'TS_DELTA_S5_W3': 13, 'PERF': 7,
'TS_DELTA_SHORT': 2, 'NOP': 1, 'TS_WAVE_STATE': 6,
'EVENT': 6, 'EVENT_BIG': 8, 'REG': 16, 'SNAPSHOT': 16,
'TS_DELTA_OR_MARK': 12, 'LAYOUT_HEADER': 16, 'INST': 5,
'UTILCTR': 12,
}
class TestEncodingsMatchStateTable(unittest.TestCase):
"""Verify encoding patterns produce the correct state decode table."""
def test_all_256_bytes_decode_correctly(self):
"""Each byte value should decode to the same packet type as reference."""
mismatches = []
for byte_val in range(256):
ref_opcode = REFERENCE_STATE_TABLE[byte_val]
ref_name = OLD_OPCODE_TO_NAME.get(ref_opcode, f"UNK_{ref_opcode:02x}")
our_opcode = STATE_TO_OPCODE[byte_val]
our_name = OPCODE_TO_CLASS[our_opcode].__name__
if ref_name != our_name:
mismatches.append((byte_val, ref_name, our_name))
if mismatches:
msg = "\n".join(f" 0x{b:02x}: expected {r}, got {o}" for b, r, o in mismatches[:10])
self.fail(f"State table mismatches ({len(mismatches)} total):\n{msg}")
class TestPacketSizesMatchBudget(unittest.TestCase):
"""Verify packet sizes (from field definitions) match expected budget values."""
def test_all_packet_sizes(self):
"""Each packet type's size should match the reference budget."""
for pkt_cls in PACKET_TYPES:
name = pkt_cls.__name__
expected = REFERENCE_BUDGET_NIBBLES.get(name)
if expected is None:
continue
actual = pkt_cls.size_nibbles()
self.assertEqual(expected, actual,
f"{name}: expected {expected} nibbles, got {actual} (size_bits={pkt_cls.size_bits()})")
class TestFieldExtraction(unittest.TestCase):
"""Test that field values are extracted correctly."""
def test_valuinst(self):
reg = 0b11110_1_001_011 # wave=0x1E, flag=1, delta=1
pkt = VALUINST.from_raw(reg)
self.assertEqual(pkt.delta, 1)
self.assertEqual(pkt.flag, 1)
self.assertEqual(pkt.wave, 0x1E)
def test_vmemexec_enum(self):
reg = 0b11_00_1111 # src=3 (VMEM_ALT), delta=0
pkt = VMEMEXEC.from_raw(reg)
self.assertEqual(pkt.src, MemSrc.VMEM_ALT)
def test_aluexec_enum(self):
reg = 0b10_01_1110 # src=2 (VALU), delta=1
pkt = ALUEXEC.from_raw(reg)
self.assertEqual(pkt.src, AluSrc.VALU)
def test_waveend(self):
reg = (0x15 << 15) | (0x7 << 11) | (0x3 << 9) | (1 << 8) | 0b10101
pkt = WAVEEND.from_raw(reg)
self.assertEqual(pkt.flag7, 1)
self.assertEqual(pkt.simd, 3)
self.assertEqual(pkt.cu_lo, 7)
self.assertEqual(pkt.wave, 0x15)
self.assertEqual(pkt.cu, 0xF) # cu_lo | (flag7 << 3) = 7 | 8 = 15
def test_wavestart(self):
reg = (0x7F << 18) | (0x15 << 13) | (0x7 << 10) | (0x3 << 8) | (1 << 7) | 0b01100
pkt = WAVESTART.from_raw(reg)
self.assertEqual(pkt.flag7, 1)
self.assertEqual(pkt.simd, 3)
self.assertEqual(pkt.cu_lo, 7)
self.assertEqual(pkt.wave, 0x15)
self.assertEqual(pkt.id7, 0x7F)
self.assertEqual(pkt.cu, 0xF)
def test_inst_enum(self):
reg = (0x21 << 13) | (0x15 << 8) | (1 << 7) | (1 << 3) | 0b010
pkt = INST.from_raw(reg)
self.assertEqual(pkt.flag1, 1)
self.assertEqual(pkt.flag2, 1)
self.assertEqual(pkt.wave, 0x15)
self.assertEqual(pkt.op, InstOp.VMEM_LOAD)
def test_layout_header(self):
reg = (0b101 << 33) | (0b1010 << 28) | (0b111 << 15) | (0b11 << 13) | (0b101010 << 7) | 0b0010001
pkt = LAYOUT_HEADER.from_raw(reg)
self.assertEqual(pkt.layout, 0b101010)
self.assertEqual(pkt.simd, 0b11)
self.assertEqual(pkt.group, 0b111)
self.assertEqual(pkt.sel_a, 0b1010)
self.assertEqual(pkt.sel_b, 0b101)
def test_ts_delta_or_mark_modes(self):
# delta mode: bit9=0, bit8=0
pkt_delta = TS_DELTA_OR_MARK.from_raw(0b0000001) # just the encoding pattern
self.assertFalse(pkt_delta.is_marker)
# marker mode: bit9=1, bit8=0
pkt_marker = TS_DELTA_OR_MARK.from_raw(0b0000001 | (1 << 9)) # bit9=1, bit8=0
self.assertTrue(pkt_marker.is_marker)
# other mode: bit9=1, bit8=1 (not marker)
pkt_other = TS_DELTA_OR_MARK.from_raw(0b0000001 | (1 << 8) | (1 << 9))
self.assertFalse(pkt_other.is_marker)
def test_reg(self):
# REG fields: slot=bits[9:7], hi_byte=bits[15:8], subop=bits[31:16], val32=bits[63:32]
# Note: slot[2:1] overlaps with hi_byte[1:0], so we need to set them consistently
# hi_byte=0x55 means bits 8-15 = 0b01010101, so slot bits 8-9 = 0b01
# slot bit 7 = 1, so slot = 0b011 = 3
reg = (0xDEADBEEF << 32) | (0xCAFE << 16) | (0x55 << 8) | (1 << 7) | 0b1001
pkt = REG.from_raw(reg)
self.assertEqual(pkt.slot, 0b011) # bit7=1, bits 8-9 from hi_byte low 2 bits = 01
self.assertEqual(pkt.hi_byte, 0x55)
self.assertEqual(pkt.subop, 0xCAFE)
self.assertEqual(pkt.val32, 0xDEADBEEF)
class TestRoundtrip(unittest.TestCase):
"""Test encode/decode roundtrip."""
def test_simple_roundtrip(self):
"""Test encode/decode roundtrip preserves packet types."""
test_packets = [
LAYOUT_HEADER.from_raw(0x100),
WAVESTART.from_raw(0x0),
INST.from_raw(0x10),
INST.from_raw(0x10),
WAVEEND.from_raw(0x40),
]
encoded = encode(test_packets)
decoded = decode(encoded)
self.assertGreaterEqual(len(decoded), len(test_packets))
for i, (orig, dec) in enumerate(zip(test_packets, decoded)):
self.assertEqual(type(orig), type(dec), f"type mismatch at {i}")
if __name__ == "__main__":
unittest.main()
+189
View File
@@ -0,0 +1,189 @@
from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
from tinygrad.codegen.opt.kernel import Ops, MemOp, UOp
from tinygrad.uop.ops import BinaryOps, UnaryOps
from tinygrad.dtype import DType, dtypes
from tinygrad.helpers import DEBUG
from tinygrad.uop.ops import Variable, NumNode, MulNode, DivNode, ModNode, LtNode, SumNode, AndNode
import functools
import math
from collections import defaultdict
_type_to_letter = {dtypes.float32: 'f', dtypes.bool: 'p', dtypes.int32: 'i', dtypes.int64: 'a', dtypes.uint32: 'u', dtypes.uint64: 'b', dtypes.float.vec(4): 'x', dtypes.uint8: 'uc', dtypes.float16: 'h',
dtypes.int8: 'c', dtypes.uint16: 'us', dtypes.float64: 'd'}
class Register(NamedTuple):
nm:str
dtype:DType
scalar:bool
off:Optional[int] = None
def __repr__(self): return self.nm if self.off is None else f"{self.nm}:{self.off}"
def subregs(self):
if self.dtype == dtypes.float.vec(4):
return [Register(self.nm, dtypes.float, False, off=off) for off in range(4)]
return []
class AssemblyInstruction(NamedTuple):
op: Ops
out: Optional[Register]
vin: List[Union[Register, int, float]]
arg: Any = None
# warp size of 32, s registers are shared across the warp, v are 32-wide vectors
class AssemblyLanguage:
supports_load3: bool = False
sin_is_sin2pi: bool = False
no_div: bool = False
#TODO: these should be global vars
cnts:DefaultDict[Tuple[DType, bool], int] = defaultdict(int)
tor: Dict[Any, Register] = {}
ins: List[AssemblyInstruction] = []
def type_to_letter(self,x): return _type_to_letter[x[0]].upper() if x[1] else _type_to_letter[x[0]]
def newreg(self, tok, dtype=dtypes.float32, scalar=False) -> Register:
self.tor[tok] = ret = Register(f"%{self.type_to_letter((dtype, scalar))}{self.cnts[(dtype, scalar)]}", dtype, scalar)
if dtype == dtypes.float.vec(4):
for off in range(4):
self.tor[tok] = Register(ret.nm, dtypes.float, ret.scalar, off)
self.cnts[(dtype, scalar)] += 1
return ret
def render_numnode(self, b) -> Register:
key = ("num", b)
if key not in self.tor: self.ins.append(AssemblyInstruction(Ops.LOAD, self.newreg(key, scalar=True, dtype=dtypes.int32), [], b))
return self.tor[key]
def render_alu(self, op, a:Register, b:Union[Register, int, float], dtype=dtypes.int32) -> Register:
key = (op, a, b)
if key not in self.tor:
#if not isinstance(b, Register): b = render_numnode(b)
self.ins.append(AssemblyInstruction(Ops.ALU, self.newreg(key, dtype=dtype, scalar=a.scalar and (not isinstance(b, Register) or b.scalar)), [a, b], op))
return self.tor[key]
def render_cast(self, a:Register, new_dtype:DType) -> Register:
if a.dtype == new_dtype: return a
key = (a, new_dtype)
if key not in self.tor:
self.ins.append(AssemblyInstruction(Ops.CAST, self.newreg(key, dtype=new_dtype), [a]))
return self.tor[key]
render_ops: Any = { Variable: lambda self, ops, ctx: ctx.tor[self], NumNode: lambda self, ops, ctx: ctx.render_numnode(self.b),
MulNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.MUL, self.a.render(ops, ctx), self.b),
DivNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.DIV, self.a.render(ops, ctx), self.b),
ModNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.MOD, self.a.render(ops, ctx), self.b),
LtNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.CMPLT, self.a.render(ops, ctx), self.b, dtype=dtypes.bool),
SumNode: lambda self,ops,ctx: functools.reduce(lambda a,b: ctx.render_alu(BinaryOps.ADD, a, b.render(ops,ctx)), self.nodes[1:], self.nodes[0].render(ops,ctx)),
AndNode: lambda self,ops,ctx: functools.reduce(lambda a,b: ctx.render_alu(BinaryOps.MUL, a, b.render(ops,ctx), dtype=dtypes.bool), self.nodes[1:], self.nodes[0].render(ops,ctx)) }
def addr_w_offset(self, args):
assert isinstance(args, MemOp)
idx = args.idx*args.memory_dtype.itemsize
off = 0 # TODO: should this be None?
if isinstance(idx, SumNode):
nums = [n.b for n in idx.nodes if isinstance(n, NumNode)]
if nums and nums[0] < 4096 and (idx-nums[0]).min >= 0: # TODO: different for each GPU?
idx -= nums[0]
off = cast(int, nums[0])
reg = idx.render(self.render_ops, self)
if self.supports_load3:
if reg.scalar:
new_reg = self.newreg((reg.nm, 'vec'), dtype=reg.dtype)
self.ins.append(AssemblyInstruction(Ops.ALU, new_reg, [reg], UnaryOps.NOOP))
reg = new_reg
return self.tor[args.name], reg, off
reg = self.render_alu(BinaryOps.ADD, self.render_cast(reg, dtypes.uint64), self.tor[args.name], dtype=dtypes.uint64)
return reg, None, off
def uops_to_asmstyle(lang, function_name:str, uops:List[UOp]):
#TODO: Do not use clear()
lang.ins.clear()
lang.tor.clear()
lang.cnts.clear()
buf_to_dtype = {args:dtype for uop,dtype,_,args,_ in uops if uop == Ops.DEFINE_GLOBAL}
global_size, local_size = [], []
skipload_branch = 0
lang.ins += [AssemblyInstruction(Ops.SPECIAL, lang.newreg(buf, dtype=dtypes.uint64, scalar=True), [], buf) for buf in buf_to_dtype]
for u in uops:
uop,dtype,vin,args,_ = u
if uop == Ops.DEFINE_LOCAL:
lang.ins.append(AssemblyInstruction(Ops.DEFINE_LOCAL, None, [], args))
lang.ins.append(AssemblyInstruction(Ops.ALU, lang.newreg(args[0], dtype=dtypes.uint64), [args[0]], UnaryOps.NOOP))
elif uop == Ops.LOOP:
if args[1] == "global":
for i,var in enumerate(args[0]):
global_size.append(var.max+1)
lang.ins.append(AssemblyInstruction(Ops.SPECIAL, lang.newreg(var, dtype=dtypes.int32), [], f"gid{len(args[0])-1-i}"))
elif args[1] == "local":
for i,var in enumerate(args[0]):
local_size.append(var.max+1)
lang.ins.append(AssemblyInstruction(Ops.SPECIAL, lang.newreg(var, dtype=dtypes.int32), [], f"lid{len(args[0])-1-i}"))
else:
for var in args[0]:
if not isinstance(var, NumNode): # TODO: why is this coming through?
lang.ins.append(AssemblyInstruction(Ops.LOAD, lang.newreg(var, dtype=dtypes.int32, scalar=True), [], 0))
lang.ins.append(AssemblyInstruction(Ops.LABEL, None, [], "$loop_"+var.expr))
elif uop == Ops.ENDLOOP:
if args[1] not in ["global", "local", "global+local"]:
for var in reversed(args[0]):
if not isinstance(var, NumNode): # TODO: why is this coming through?
lang.ins.append(AssemblyInstruction(Ops.ALU, lang.tor[var], [lang.tor[var], 1], BinaryOps.ADD))
pred = lang.render_alu(BinaryOps.CMPLT, lang.tor[var], var.max+1, dtypes.bool)
lang.ins.append(AssemblyInstruction(Ops.COND_BRANCH, None, [pred], ("$loop_"+var.expr, True)))
elif args[1] == "global+local":
for i, var in enumerate(reversed(args[0])):
lang.ins.append(AssemblyInstruction(Ops.ENDLOOP, None, [lang.tor[var]], (var.max+1, f"gid{i}")))
elif args[1] == 'local':
for i, var in enumerate(reversed(args[0])):
lang.ins.append(AssemblyInstruction(Ops.ENDLOOP, None, [lang.tor[var]], (var.max+1, f"lid{i}")))
elif uop == Ops.CAST:
# TODO: we should reconsider outputting CAST in the linearizer. these are needless copies
out = lang.newreg(u, dtype)
for i,sr in enumerate(out.subregs()):
lang.ins.append(AssemblyInstruction(Ops.ALU, sr, [lang.tor[vin[i]]], UnaryOps.NOOP))
elif uop == Ops.ALU:
out = lang.newreg(u, dtype) if u not in lang.tor else lang.tor[u]
# this is the only thing that can violate SSA
if args in [BinaryOps.CMPLT]:
pred_reg = lang.newreg((u, 'pred'), dtype=dtypes.bool)
lang.ins.append(AssemblyInstruction(Ops.ALU, pred_reg, [lang.tor[x] for x in vin], args))
lang.ins.append(AssemblyInstruction(Ops.CAST, out, [pred_reg], args))
elif args == BinaryOps.DIV and lang.no_div:
tmp = lang.newreg((u, "rcp"))
lang.ins.append(AssemblyInstruction(Ops.ALU, tmp, [lang.tor[vin[1]]], UnaryOps.RECIP))
lang.ins.append(AssemblyInstruction(Ops.ALU, out, [lang.tor[vin[0]], tmp], BinaryOps.MUL))
elif args == UnaryOps.SIN and lang.sin_is_sin2pi:
tmp = lang.newreg((u, "2pi"))
lang.ins.append(AssemblyInstruction(Ops.ALU, tmp, [lang.tor[vin[0]], 1/(math.pi*2)], BinaryOps.MUL))
lang.ins.append(AssemblyInstruction(Ops.ALU, out, [tmp], args))
else:
lang.ins.append(AssemblyInstruction(Ops.ALU, out, [lang.tor[x] for x in vin], args))
elif uop == Ops.DEFINE_REG:
reg = lang.newreg(u, dtype=dtype)
lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [], args))
elif uop == Ops.SPECIAL:
lang.tor[u] = lang.tor[args]
elif uop == Ops.CONST:
lang.ins.append(AssemblyInstruction(Ops.LOAD, lang.newreg(u, dtype=dtype), [], args))
elif uop == Ops.LOAD:
idx, treg, off = lang.addr_w_offset(args)
reg = lang.newreg(u, dtype=dtype, scalar=(idx.scalar and (not isinstance(treg, Register) or treg.scalar)))
if args.valid.min == 0:
lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [], 0))
if args.valid.max == 1:
pred = args.valid.render(lang.render_ops, lang)
lang.ins.append(AssemblyInstruction(Ops.COND_BRANCH, None, [pred], (f"$skipload_{skipload_branch}", False)))
if args.valid.max == 1:
# NOTE: you can't compute the index in here, because it assumes it's all available later
lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [idx] + ([treg] if treg is not None else []), (off, 'global' if not args.local else 'shared', args.memory_dtype if args.memory_dtype != dtypes.float else None)))
if args.valid.min == 0 and args.valid.max == 1:
lang.ins.append(AssemblyInstruction(Ops.LABEL, None, [], f"$skipload_{skipload_branch}"))
skipload_branch += 1
elif uop == Ops.STORE:
if args is None:
lang.ins.append(AssemblyInstruction(Ops.ALU, lang.tor[vin[0]], [lang.tor[vin[1]]], UnaryOps.NOOP))
else:
idx, treg, off = lang.addr_w_offset(args)
lang.ins.append(AssemblyInstruction(Ops.STORE, None, [idx, lang.tor[vin[0]]] + ([treg] if treg is not None else []), (off, 'global' if not args.local else 'shared', args.memory_dtype if args.memory_dtype != dtypes.float else None)))
if DEBUG >= 4:
for tins in lang.ins: print(tins)
return global_size, local_size
+177
View File
@@ -0,0 +1,177 @@
import struct
from platform import system
from typing import Tuple, Dict, List, Optional
from tinygrad import dtypes
from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
from tinygrad.codegen.opt.kernel import Ops, UOp
from tinygrad.helpers import CI
from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
def float_to_hex(x): return "%02X%02X%02X%02X" % tuple(struct.pack("f",x)[::-1])
def compute_offsets(total):
quotient, remainder = divmod(total, 4096)
return [4096]*quotient + [remainder] if remainder else [4096]*quotient
#NOTE: Darwin needs names to start with a "_"
def get_name(name): return ('_' if system() == 'Darwin' else '') + name
class ARM64Language(AssemblyLanguage): pass
def specialize_to_arm64(fn_nm, asm):
var_size = 16
prev_uop:Optional[Ops] = None
ins = []
x_regs = ['x' + str(i) for i in reversed(range(12))]
s_regs = ['s' + str(i) for i in reversed(range(3,32)) if i <= 7 or i >= 16]
type_to_reg = {dtypes.double: "d", dtypes.half: 'h', dtypes.float32: 's', dtypes.bool: 'w', dtypes.int8:'w', dtypes.int32: 'w', dtypes.int64: 'x', dtypes.uint8:'w', dtypes.uint32: 'w', dtypes.uint64: 'x'}
alu = {BinaryOps.ADD: "add", BinaryOps.SUB: "sub", BinaryOps.MUL: "mul", BinaryOps.DIV: "div", BinaryOps.MAX: "max",
BinaryOps.MOD: "", BinaryOps.CMPLT: "subs",
UnaryOps.NOOP: "mov", UnaryOps.NEG: "neg",
UnaryOps.SIN:'bl ' + get_name('sinf'), UnaryOps.LOG2: 'bl ' + get_name("log2f"), UnaryOps.EXP2: 'bl ' + get_name("exp2f"), UnaryOps.SQRT: 'bl ' + get_name("sqrtf"),
TernaryOps.MULACC: "madd", TernaryOps.WHERE: "fcsel"}
def mov_imm(value, reg):
# Manually move value into reg if value can't fit
if value.__class__ is not float and abs(value) > abs(65535):
ins.append(f"movz w15, #{value & 0xffff}")
ins.append(f"movk w15, #{(value >> 16) & 0xffff}, lsl #16")
ins.append(f"sxtw {reg}, w15")
elif reg[0] == 's':
ins.append(f"movz x15, 0x{float_to_hex(value)[4:]}")
ins.append(f"movk x15, 0x{float_to_hex(value)[:4]}, lsl #16")
ins.append("str x15, [sp, 16]")
ins.append(f"ldr {reg}, [sp, 16]")
else:
ins.append(f"mov {reg}, #{value}")
# Get variables intervals
live_range:Dict[str, List[int]] = {}
for i, (uop, out, vin, arg) in enumerate(asm):
for var in ([v for v in [out] + vin if v is not None and v.__class__ is not int]):
live_range[var.nm] = [i,i] if var.nm not in live_range else [live_range[var.nm][0], i]
mem_vars:Dict[str, int] = {}
rtor:Dict[str, str] = {}
def allocate_regs(mvars):
nonlocal var_size
for v in [v for v in mvars if v is not None and v.__class__ is not int and v.nm not in rtor]:
available_regs = s_regs if dtypes.is_float(v[1]) else x_regs
#NOTE: Very simple spill, everything that don't fit in regs goes to mem
if not available_regs:
# ARM needs the stack 16-byte aligned
var_size += 16
available_regs.append('s0' if dtypes.is_float(out[1]) else 'x12')
mem_vars[v.nm] = var_size
rtor[v.nm] = available_regs.pop()
temp_floats = ['s0', 's1', 's2']
temp_ints = ['x12', 'x13', 'x16']
for i, (uop, out, vin, arg) in enumerate(asm):
# Clear regs out of interval
for var, reg in list(rtor.items()):
available_regs = s_regs if reg[0] == 's' else x_regs
if var[1] not in 'B' and var not in mem_vars and i > live_range[var][1]:
available_regs.append(rtor.pop(var))
# Assign a registers to the variables using live ranges.
allocate_regs([out] + vin)
# Assign temp regs to vin and load them before direct use
for i, v in enumerate([v for v in vin if v.__class__ is not int and v.nm in mem_vars]):
rtor[v.nm] = temp_floats[i] if dtypes.is_float(v[1]) else temp_ints[i]
# ARM64 addressing constraints https://devblogs.microsoft.com/oldnewthing/20220728-00/?p=106912
ins.append(f"mov x15, {mem_vars[v.nm]}")
ins.append(f"ldr {rtor[v.nm]}, [sp, x15]")
if uop == Ops.SPECIAL:
if arg.startswith('data'):
# data 8 to n into the stack
if int(arg[4:]) >= 8:
ins.append(f"ldr x15, [x17, #{(int(arg[4:]) - 8) * 8}]")
ins.append(f"mov {rtor[out.nm]}, x15")
else:
ins.append(f"mov {rtor[out.nm]}, #0")
ins.append(f"loop_{arg}:")
elif uop == Ops.CAST:
if arg == BinaryOps.CMPLT:
if rtor[out.nm][0] == 's':
mov_imm(0.0, 's0')
mov_imm(1.0, 's1')
ins.append(f"fcsel {rtor[out.nm]}, s1, s0, lt")
if rtor[out.nm][0] == 'x':
mov_imm(0, 'x14')
mov_imm(1, 'x15')
ins.append(f"csel {rtor[out.nm]}, x15, x14, lt")
else:
ins.append(f"sxtw {rtor[out.nm]}, w{rtor[vin[0].nm][1:]}")
elif uop == Ops.ALU:
if len(vin)==2 and vin[1].__class__ is int: mov_imm(vin[1], 'x15')
if arg == BinaryOps.MUL and out.dtype == dtypes.bool:
ins.append(f"ands {','.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}")
elif arg == TernaryOps.WHERE:
ins.append(f"fcmp {rtor[vin[0].nm]}, #0.0" if rtor[vin[0].nm][0] == 's' else f"cmp {rtor[vin[0].nm]}, #0")
ins.append(f"{alu[arg]} {rtor[out.nm]}, {rtor[vin[1].nm]}, {rtor[vin[2].nm]}, ne")
elif arg in [UnaryOps.LOG2, UnaryOps.SIN, UnaryOps.EXP2, UnaryOps.SQRT]:
#NOTE: Not a real instruction, use to emulate a ext call in unicorn
if CI: ins.append(f"{alu[arg]} {rtor[out.nm]} {rtor[vin[0].nm]}")
else:
save_regs = [k for k in rtor.keys() if k != out.nm and k not in mem_vars]
ins.append(f"sub sp, sp, #{(len(save_regs))*16}")
# Save the registers before they are cleared by func call
for i,k in enumerate(save_regs,1):
ins.append(f"str {rtor[k]}, [sp, #{16*i}]")
ins.append("stp x29, x30, [sp, #0]!")
ins.append("mov x29, sp")
ins.append(f"fmov s0, {rtor[vin[0].nm]}")
ins.append(alu[arg])
ins.append(f"fmov {rtor[out.nm]}, s0")
ins.append("mov sp, x29")
ins.append("ldp x29, x30, [sp], #0")
for i,k in enumerate(save_regs,1):
ins.append(f"ldr {rtor[k]}, [sp, #{16*i}]")
ins.append(f"add sp, sp, #{len(save_regs)*16}")
elif arg == BinaryOps.CMPLT:
ins.append(f"{alu[arg]} {','.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}" if not dtypes.is_float(vin[0][1]) else f"fcmp {rtor[vin[0].nm]}, {rtor[vin[1].nm]}")
elif arg == BinaryOps.MOD:
rhs = 'x15' if vin[1].__class__ is int else rtor[vin[1].nm]
ins.append(f"udiv x14, {rtor[vin[0].nm]}, {rhs}")
ins.append(f"msub {rtor[out.nm]}, x14, {rhs}, {rtor[vin[0].nm]}")
else:
ins.append(f"{'f' if dtypes.is_float(vin[0][1]) else 's' if arg == BinaryOps.DIV else ''}{alu[arg]} {', '.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}")
elif uop == Ops.LOAD:
if arg.__class__ in (int, float):
mov_imm(arg, rtor[out.nm])
else:
#NOTE: if need casting load var in s/h0 or x/w12 temp regs
reg_in = type_to_reg[arg[2]] + ('0' if dtypes.is_float(arg[2]) else '12') if arg[2] is not None else rtor[out.nm]
mov_imm(arg[0], "x15")
ins.append(f"add x15, {rtor[vin[0].nm]}, x15")
ins.append(f"ldr{'sb' if arg[2] is not None and arg[2] in (dtypes.int8, dtypes.uint8, dtypes.bool) else ''} {reg_in}, [x15]")
if arg[2] is not None: ins.append(f"{'fcvt' if arg[2] in [dtypes.half, dtypes.double] else 'scvtf'} {rtor[out.nm]}, {reg_in}")
elif uop == Ops.STORE:
#NOTE: if need casting load var in s/h0 or x/w12 temp regs
reg_out = (type_to_reg[arg[2]] + ('0' if dtypes.is_float(arg[2]) else '12') if arg[2] is not None else rtor[vin[1].nm])
if arg[2] is not None: ins.append(f"fcvt{'zs' if arg[2] not in [dtypes.half, dtypes.double] else '' } {reg_out}, {rtor[vin[1].nm]}")
ins.append(f"mov x15, #{arg[0]}")
ins.append(f"str {reg_out}, [{rtor[vin[0].nm]}, x15, lsl #0]")
elif uop == Ops.COND_BRANCH:
#TODO: this is a hack it shouldn't always be a cmp before a cond branch?
if prev_uop == Ops.LOAD:
ins.append(f"cmp {rtor[vin[0].nm]}, #0")
ins.append(f"b.{'lt' if arg[1] else 'ge'} {arg[0][1:]}")
elif uop == Ops.LABEL:
ins.append(f"{arg[1:]}:")
elif uop == Ops.ENDLOOP:
mov_imm(arg[0], "x15")
ins.append(f"add {rtor[vin[0].nm]}, {rtor[vin[0].nm]}, #1")
ins.append(f"cmp {rtor[vin[0].nm]}, x15")
ins.append(f"b.lt loop_{arg[1]}")
prev_uop = uop
# store regs into memory if needed
if out is not None and out.nm in mem_vars:
ins.append(f"mov x15, {mem_vars[out.nm]}")
ins.append(f"str {rtor[out.nm]}, [sp, x15]")
return "\n".join([f"//varsize {var_size}",".arch armv8-a",".text", f".global {get_name(fn_nm)}",".p2align 2", f"{get_name(fn_nm)}:", "mov x17, sp"] + [f"sub sp, sp, #{offset}" for offset in compute_offsets(var_size)]+ ins + [f"add sp, sp, #{offset}" for offset in compute_offsets(var_size)] +["ret", "\n"])
def uops_to_arm64_asm(fn_nm:str, uops:List[UOp]) -> Tuple[str, List[int], List[int], bool]:
lang = ARM64Language()
global_size, local_size = uops_to_asmstyle(lang, fn_nm, uops)
return specialize_to_arm64(fn_nm, lang.ins), global_size[::-1], local_size[::-1], True

Some files were not shown because too many files have changed in this diff Show More