Compare commits

..
Author SHA1 Message Date
geohot 48dd1d6543 no long 2026-01-21 21:53:23 +09:00
geohot a06023cd36 viz slowness 2026-01-21 21:50:09 +09:00
geohot ac232bceb5 remove the device when we render 2026-01-21 21:41:29 +09:00
geohot f19fbadce4 regression test 2026-01-21 21:30:04 +09:00
geohot 8eb762d6fa add device to local, fix PCONTIG=2 2026-01-21 19:11:49 +09:00
502 changed files with 14533 additions and 64499 deletions
+1 -1
View File
@@ -11,5 +11,5 @@ runs:
git fetch origin $CURRENT_SHA
export COMMIT_MESSAGE=$(git show -s --format=%B "$CURRENT_SHA")
export CURRENT_HEAD=$(git rev-parse HEAD)
cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && CHECK_OOB=0 PYTHONPATH=. python3 process_replay.py
cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && IGNORE_OOB=1 PYTHONPATH=. python3 process_replay.py
git checkout $CURRENT_HEAD # restore to branch
+14 -37
View File
@@ -56,15 +56,7 @@ runs:
# **** Caching packages ****
- name: Cache Python packages (PR)
if: github.event_name == 'pull_request'
id: restore-venv-pr
uses: actions/cache/restore@v4
with:
path: ${{ github.workspace }}/.venv
key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
- name: Cache Python packages
if: github.event_name != 'pull_request'
id: restore-venv
uses: actions/cache@v4
with:
@@ -73,23 +65,23 @@ runs:
# **** Caching downloads ****
- name: Cache downloads (PR)
if: inputs.key != '' && github.event_name == 'pull_request'
uses: actions/cache/restore@v4
with:
path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
- name: Cache downloads
if: inputs.key != '' && github.event_name != 'pull_request'
- name: Cache downloads (Linux)
if: inputs.key != '' && runner.os == 'Linux'
uses: actions/cache@v4
with:
path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
path: ~/.cache/tinygrad/downloads/
key: downloads-${{ github.job }}-${{ 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 }}
# **** Python deps ****
- name: Install dependencies in venv (with extra)
if: inputs.deps != '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
if: inputs.deps != '' && steps.restore-venv.outputs.cache-hit != 'true'
shell: bash
run: |
python -m venv .venv
@@ -100,7 +92,7 @@ runs:
fi
python -m pip install -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
- name: Install dependencies in venv (without extra)
if: inputs.deps == '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
if: inputs.deps == '' && steps.restore-venv.outputs.cache-hit != 'true'
shell: bash
run: |
python -m venv .venv
@@ -190,14 +182,8 @@ runs:
echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
- name: Cache apt (PR)
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name == 'pull_request'
uses: actions/cache/restore@v4
with:
path: /var/cache/apt/archives/
key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
- name: Cache apt
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name != 'pull_request'
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
uses: actions/cache@v4
with:
path: /var/cache/apt/archives/
@@ -253,17 +239,8 @@ runs:
ln -s /opt/homebrew/opt/[email protected] /opt/homebrew/opt/boost || true
ln -s /opt/homebrew/opt/boost/lib/libboost_atomic-mt.dylib /opt/homebrew/opt/boost/lib/libboost_atomic.dylib || true
ln -s /opt/homebrew/opt/boost/lib/libboost_thread-mt.dylib /opt/homebrew/opt/boost/lib/libboost_thread.dylib || true
- name: Cache gpuocelot (PR)
if: inputs.ocelot == 'true' && github.event_name == 'pull_request'
id: cache-build-pr
uses: actions/cache/restore@v4
env:
cache-name: cache-gpuocelot-build-1
with:
path: ${{ github.workspace }}/gpuocelot/ocelot
key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
- name: Cache gpuocelot
if: inputs.ocelot == 'true' && github.event_name != 'pull_request'
if: inputs.ocelot == 'true'
id: cache-build
uses: actions/cache@v4
env:
@@ -272,7 +249,7 @@ runs:
path: ${{ github.workspace }}/gpuocelot/ocelot
key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
- name: Clone/compile gpuocelot
if: inputs.ocelot == 'true' && steps.cache-build-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true'
if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
shell: bash
run: |
git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
+4 -4
View File
@@ -40,13 +40,13 @@ jobs:
mesa: 'true'
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 libdrm-dev
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
run: |
find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "comgr_3.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
find tinygrad/runtime/autogen -type f -name "*.py" -not -name "__init__.py" -not -name "comgr_3.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
python3 -c "from tinygrad.runtime.autogen import opencl"
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd"
python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v14_0_2"
python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
python3 -c "from tinygrad.runtime.autogen import llvm"
@@ -84,7 +84,7 @@ jobs:
- name: Regenerate autogen files
run: |
rm tinygrad/runtime/autogen/metal.py tinygrad/runtime/autogen/iokit.py tinygrad/runtime/autogen/corefoundation.py
python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
- name: Check for differences
run: |
if ! git diff --quiet; then
+11 -59
View File
@@ -16,48 +16,6 @@ on:
workflow_dispatch:
jobs:
# the goal of this test is to replicate a normal person on a laptop running the test
# no process replay, no benchmarks, no CI, just a normal laptop person
# the 3 minute timeout should not be raised
testmacpytest:
name: Mac pytest
runs-on: [self-hosted, macOS]
timeout-minutes: 3
defaults:
run:
shell: bash -e -o pipefail {0}
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
uses: actions/checkout@v4
# brew install uv
- name: setup python environment
run: |
rm -rf /tmp/tinygrad_pytest_ci
uv venv /tmp/tinygrad_pytest_ci
source /tmp/tinygrad_pytest_ci/bin/activate
uv pip install .[testing]
- name: setup staging db
run: |
echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
rm -f /tmp/pytest-db-ci*
# TODO: remove this step once all old caches are migrated
- name: Migrate old huggingface cache (symlinks break onnxruntime 1.24+)
run: |
cd ~/Library/Caches/tinygrad/downloads/models 2>/dev/null || exit 0
for old_dir in models--*; do
[ -d "$old_dir" ] || continue
repo_id=$(echo "$old_dir" | sed 's/models--//; s/--/\//g')
snapshot=$(ls -1 "$old_dir/snapshots" 2>/dev/null | head -1)
[ -n "$snapshot" ] || continue
mkdir -p "$repo_id"
cp -RLn "$old_dir/snapshots/$snapshot/"* "$repo_id/" 2>/dev/null || true
done
- name: Run pytest -nauto
run: |
source /tmp/tinygrad_pytest_ci/bin/activate
pytest -nauto --durations=20
testmacbenchmark:
name: Mac Benchmark
env:
@@ -187,10 +145,6 @@ jobs:
run: |
echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
- name: Kill stale pids
run: |
PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
- 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
@@ -378,9 +332,9 @@ jobs:
- name: Setcap to python
run: ./extra/amdpci/setup_python_cap.sh
- name: Remove amd modules
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
run: ./extra/hcq/hcq_smi.py amd rmmod
- name: Kill stale pids
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
run: ./extra/hcq/hcq_smi.py amd kill_pids
#- name: Insert amdgpu
# run: sudo modprobe amdgpu
- name: Symlink models and datasets
@@ -490,9 +444,9 @@ jobs:
- name: Setcap to python
run: ./extra/amdpci/setup_python_cap.sh
- name: Remove amd modules
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
run: ./extra/hcq/hcq_smi.py amd rmmod
- name: Kill stale pids
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
run: ./extra/hcq/hcq_smi.py amd kill_pids
- name: Symlink models and datasets
run: |
mkdir -p weights
@@ -525,8 +479,6 @@ jobs:
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Run full CIFAR training steps w 6 GPUS
run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Test full tinyfs load
run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
- name: Run process replay tests
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -544,9 +496,9 @@ jobs:
- name: Setcap to python
run: ./extra/amdpci/setup_python_cap.sh
- name: Remove amd modules
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
run: ./extra/hcq/hcq_smi.py amd rmmod
- name: Kill stale pids
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
run: ./extra/hcq/hcq_smi.py amd kill_pids
- name: Symlink models and datasets
run: |
mkdir -p weights
@@ -609,7 +561,7 @@ jobs:
- name: openpilot compile3 0.10.1 driving_policy
run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
- name: openpilot compile3 0.10.1 dmonitoring
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
- name: benchmark MobileNetV2 on DSP
run: |
# generate quantized weights
@@ -635,9 +587,9 @@ jobs:
- name: Setcap to python
run: ./extra/amdpci/setup_python_cap.sh
- name: Remove amd modules
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
run: ./extra/hcq/hcq_smi.py amd rmmod
- name: Kill stale pids
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
run: ./extra/hcq/hcq_smi.py amd kill_pids
- name: Symlink models and datasets
run: |
mkdir -p weights
@@ -699,9 +651,9 @@ jobs:
- name: Setcap to python
run: ./extra/amdpci/setup_python_cap.sh
- name: Remove nv modules
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv rmmod
run: ./extra/hcq/hcq_smi.py nv rmmod
- name: Kill stale pids
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
run: ./extra/hcq/hcq_smi.py nv kill_pids
- name: Symlink models and datasets
run: |
mkdir -p weights
+107 -127
View File
@@ -1,11 +1,11 @@
name: Unit Tests
env:
# increment this when downloads substantially change to avoid the internet
CACHE_VERSION: '16'
CACHE_VERSION: '15'
CAPTURE_PROCESS_REPLAY: 1
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PYTHONPATH: ${{ github.workspace }}
CHECK_OOB: 1
IGNORE_OOB: 0
on:
push:
@@ -26,19 +26,19 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: llvm-speed
deps: testing_unit
deps: testing_minimal
llvm: 'true'
- name: Speed Test
run: CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
- name: Speed Test (BEAM=2)
run: BEAM=2 CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
run: BEAM=2 CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
docs:
name: Docs
runs-on: ubuntu-22.04
timeout-minutes: 10
env:
CHECK_OOB: 0
IGNORE_OOB: 1
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -98,7 +98,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: torch-backend-pillow-torchvision-et-pt
deps: testing_unit
deps: testing_minimal
pydeps: "pillow torchvision expecttest"
llvm: 'true'
- name: Install ninja
@@ -106,7 +106,7 @@ jobs:
sudo apt update || true
sudo apt install -y --no-install-recommends ninja-build
- name: Test one op
run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_tiny.py TestTiny.test_plus
run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
- name: Test ResNet-18
run: DEBUG=2 python3 extra/torch_backend/example.py
- name: custom tests
@@ -114,7 +114,7 @@ jobs:
- name: Test one op in torch tests
run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
- name: Test Ops with TINY_BACKEND
run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/backend/test_ops.py --durations=20
run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
- name: Test in-place operations on views
run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
- name: Test multi-gpu
@@ -134,7 +134,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: torch-backend-pillow-torchvision-et-pt
deps: testing_unit
deps: testing_minimal
llvm: 'true'
- name: Install ninja
run: |
@@ -156,27 +156,27 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: be-minimal
deps: testing_unit
deps: testing_minimal
- name: Test dtype with Python emulator
run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/backend/test_dtype.py test/backend/test_dtype_alu.py
run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
- name: Test ops with Python emulator
run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/backend/test_ops.py --durations=20
run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20
- name: Test uops with Python emulator
run: PYTHON=1 python3 -m pytest test/backend/test_uops.py --durations=20
run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
- name: Test symbolic with Python emulator
run: PYTHON=1 python3 test/backend/test_symbolic_ops.py
run: PYTHON=1 python3 test/test_symbolic_ops.py
- name: test_renderer_failures with Python emulator
run: PYTHON=1 python3 -m pytest -rA test/backend/test_renderer_failures.py::TestRendererFailures
run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
- name: Test IMAGE=2 support
run: |
IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_simple_conv2d
- name: Test emulated METAL tensor cores
run: |
DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_big_gemm
DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
- name: Test emulated AMX tensor cores
run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
- name: Test emulated AMD tensor cores
run: |
DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
@@ -197,9 +197,9 @@ jobs:
DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
- name: Test emulated CUDA tensor cores
run: |
DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
- name: Test emulated INTEL OpenCL tensor cores
run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
@@ -207,11 +207,11 @@ jobs:
run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
- name: Test device flop counts
run: |
DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStats.test_simple_matmul
DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul
linter:
name: Linters
@@ -239,10 +239,9 @@ jobs:
- name: Run mypy with lineprecision report
run: |
python -m mypy --lineprecision-report .
grep -v autogen lineprecision.txt | awk 'NR>2 {lines+=$2; precise+=$3; imprecise+=$4; any+=$5; empty+=$6} END {t=lines-empty; printf "TOTAL: %d lines, %d precise (%.1f%%), %d imprecise (%.1f%%), %d any (%.1f%%)\n", t, precise, 100*precise/t, imprecise, 100*imprecise/t, any, 100*any/t}'
cat lineprecision.txt
- name: Run TYPED=1
run: CHECK_OOB=0 DEV=CPU TYPED=1 python test/test_tiny.py
run: TYPED=1 python -c "import tinygrad"
unittest:
name: Unit Tests
@@ -256,29 +255,26 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: unittest-13
pydeps: "pillow ftfy regex pre-commit"
pydeps: "pillow numpy ftfy regex pre-commit"
deps: testing_unit
llvm: 'true'
amd: 'true'
- name: Run pre-commit test hooks
run: SKIP=ruff,mypy pre-commit run --all-files
- 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/null/test_device.py TestRunAsModule.test_module_runs
CPU=1 python -m pytest -n=auto test/unit/ --durations=20
- name: Run NULL backend tests
run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
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
- name: Run targetted tests on NULL backend
run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
# TODO: too slow
# - name: Run SDXL on NULL backend
# run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
- name: Run Clip tests for SD MLPerf on NULL backend
run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
- name: Run AMD emulated BERT training on NULL backend
run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
# TODO: support fake weights
#- name: Run LLaMA 7B on 4 fake devices
# run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -295,8 +291,8 @@ jobs:
python extra/optimization/extract_dataset.py
gzip -c /tmp/sops > extra/datasets/sops.gz
#DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
- name: Repo line count < 24000 lines
run: MAX_LINE_COUNT=24000 python sz.py
- name: Repo line count < 20000 lines
run: MAX_LINE_COUNT=20000 python sz.py
spec:
strategy:
@@ -316,7 +312,7 @@ jobs:
deps: testing_unit
python-version: '3.14'
- name: Test SPEC=2
run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --ignore test/unit/test_autogen.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
fuzzing:
name: Fuzzing
@@ -350,11 +346,11 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: gpu-image
deps: testing_unit
deps: testing_minimal
opencl: 'true'
- name: Test CL IMAGE=2 ops
run: |
CL=1 IMAGE=2 python -m pytest -n=auto test/backend/test_ops.py --durations=20
CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
# TODO: training is broken
# CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
- name: Run process replay tests
@@ -371,14 +367,14 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: gen-dataset
deps: testing
deps: testing_minimal
opencl: 'true'
- name: Generate Dataset
run: CL=1 extra/optimization/generate_dataset.sh
- name: Run Kernel Count Test
run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
- name: Run fused optimizer tests
run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/backend/test_optim.py -k "not muon"
run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/test_optim.py -k "not muon"
- name: Upload artifact
uses: actions/upload-artifact@v4
with:
@@ -426,7 +422,7 @@ jobs:
with:
key: onnxoptc
deps: testing
python-version: '3.12'
python-version: '3.11'
llvm: 'true'
- name: Test ONNX (CPU)
run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -437,7 +433,7 @@ jobs:
- name: Test Additional ONNX Ops (CPU)
run: CPU=1 CPU_LLVM=0 python3 test/external/external_test_onnx_ops.py
- name: Test Quantize ONNX
run: CPU=1 CPU_LLVM=0 python3 test/backend/test_quantize_onnx.py
run: CPU=1 CPU_LLVM=0 python3 test/test_quantize_onnx.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -454,7 +450,7 @@ jobs:
key: onnxoptl
deps: testing
pydeps: "tensorflow==2.19"
python-version: '3.12'
python-version: '3.11'
opencl: 'true'
- name: Test ONNX (CL)
run: CL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -467,11 +463,11 @@ jobs:
- name: Test MLPerf stuff
run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
- name: NULL=1 beautiful_mnist_multigpu
run: NULL=1 NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py
run: NULL=1 python examples/beautiful_mnist_multigpu.py
- name: Test Bert training
run: NULL=1 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
- name: Test llama 3 training
run: NULL=1 NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -480,7 +476,7 @@ jobs:
runs-on: ubuntu-24.04
timeout-minutes: 15
env:
CHECK_OOB: 0
IGNORE_OOB: 1
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -528,7 +524,7 @@ jobs:
with:
key: metal
deps: testing
python-version: '3.12'
python-version: '3.11'
- name: Test models (Metal)
run: METAL=1 python -m pytest -n=auto test/models --durations=20
- name: Test LLaMA compile speed
@@ -547,15 +543,15 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: devectorize-minimal
deps: testing_unit
deps: testing_minimal
pydeps: "pillow"
llvm: "true"
- name: Test LLVM=1 DEVECTORIZE=0
run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py
run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
- name: Test LLVM=1 DEVECTORIZE=0 for model
run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
- name: Test CPU=1 DEVECTORIZE=0
run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py
run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
testdsp:
name: Linux (DSP)
@@ -568,8 +564,8 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: dsp-minimal
deps: testing_unit
pydeps: "onnx==1.18.0 onnxruntime ml_dtypes"
deps: testing_minimal
pydeps: "onnx==1.18.0 onnxruntime pillow"
llvm: "true"
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
@@ -581,15 +577,15 @@ jobs:
load: true
tags: qemu-hexagon:latest
cache-from: type=gha
cache-to: ${{ github.event_name != 'pull_request' && 'type=gha,mode=min' || '' }}
cache-to: type=gha,mode=min
- name: Set MOCKDSP env
run: printf "MOCKDSP=1" >> $GITHUB_ENV
- name: Run test_tiny on DSP
run: DEBUG=2 DSP=1 python test/test_tiny.py
- name: Test transcendentals
run: CC=clang-20 DEBUG=2 DSP=1 python test/backend/test_transcendental.py TestTranscendentalVectorized
run: CC=clang-20 DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
- name: Test quantize onnx
run: DEBUG=2 DSP=1 python3 test/backend/test_quantize_onnx.py
run: DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
testwebgpu:
name: Linux (WebGPU)
@@ -602,16 +598,16 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: webgpu-minimal
deps: testing_unit
python-version: '3.12'
deps: testing_minimal
python-version: '3.11'
webgpu: 'true'
- name: Check Device.DEFAULT (WEBGPU) and print some source
run: |
WEBGPU=1 python -c "from tinygrad import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT"
WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
- name: Run selected webgpu tests
run: |
WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/backend --durations=20
WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -636,41 +632,34 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: ${{ matrix.backend }}-minimal
deps: testing_unit
deps: testing_minimal
amd: 'true'
llvm: ${{ matrix.backend == 'amdllvm' && 'true' }}
- name: Check Device.DEFAULT and print some source
run: |
python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
- name: Run LLVM test
if: matrix.backend=='amdllvm'
run: python test/device/test_amd_llvm.py
- name: Run pytest (amd)
run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py --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 test/testextra/test_cfg_viz.py --durations=20
- name: Run pytest (amd)
run: python -m pytest test/external/external_test_am.py --durations=20
- name: Run TRANSCENDENTAL math
run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
- name: Run TestOps.test_add with SQTT
run: |
VIZ=-2 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
VIZ=1 PMC=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
- name: Run AMD emulated mmapeak on NULL backend
env:
AMD: 0
run: PYTHONPATH=. NULL=1 EMULATE=AMD python extra/mmapeak/mmapeak.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
testamdasm:
name: AMD ASM IDE
runs-on: ubuntu-24.04
timeout-minutes: 20
env:
AMD: 1
PYTHON_REMU: 1
MOCKGPU: 1
timeout-minutes: 10
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -678,13 +667,13 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: rdna3-emu
deps: testing_unit
deps: testing_minimal
amd: 'true'
python-version: '3.14'
python-version: '3.13'
- name: Verify AMD autogen is up to date
run: |
python -m tinygrad.renderer.amd.generate
git diff --exit-code tinygrad/runtime/autogen/amd/
python -m extra.assembly.amd.amdxml
git diff --exit-code extra/assembly/amd/autogen/
- name: Install LLVM 21
run: |
wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
@@ -692,24 +681,20 @@ jobs:
sudo apt-get update
sudo apt-get install llvm-21 llvm-21-tools cloc
- name: RDNA3 Line Count
run: cloc --by-file tinygrad/renderer/amd/*.py
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: AMD_LLVM=0 python -m pytest -n=auto test/amd/ --durations 20
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 test/amd/ --durations 20
run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
- name: Run RDNA3 dtype tests
run: AMD_LLVM=0 pytest -n=auto test/backend/test_dtype_alu.py test/backend/test_dtype.py --durations 20
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_LLVM=1 pytest -n=auto test/backend/test_dtype_alu.py test/backend/test_dtype.py --durations 20
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_LLVM=0 pytest -n=auto test/backend/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril or test_nonzero or test_softmax_argmax" --durations 20
- name: Run RDNA4 emulator tests
run: MOCKGPU_ARCH=rdna4 python -m pytest test/test_tiny.py -v --durations 20
- name: Run CDNA4 emulator tests
run: AMD_LLVM=1 MOCKGPU_ARCH=cdna4 python -m pytest test/test_tiny.py -v --durations 20
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:
@@ -730,7 +715,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: ${{ matrix.backend }}-minimal
deps: testing_unit
deps: testing_minimal
cuda: 'true'
ocelot: 'true'
- name: Set env
@@ -738,12 +723,10 @@ jobs:
- name: Check Device.DEFAULT and print some source
run: |
python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
- name: Run pytest (cuda)
# skip multitensor because it's slow
run: python -m pytest -n=auto test/backend --ignore test/backend/test_multitensor.py --durations=20
- name: Run TestOps.test_add with PMA
run: VIZ=-1 PMA=1 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -763,7 +746,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: ${{ matrix.backend }}-minimal
deps: testing_unit
deps: testing_minimal
opencl: ${{ matrix.backend == 'opencl' && 'true' }}
llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
mesa: ${{ matrix.backend == 'lvp' && 'true' }}
@@ -772,11 +755,11 @@ jobs:
- name: Check Device.DEFAULT and print some source
run: |
python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
- name: Run pytest (${{ matrix.backend }})
run: python -m pytest -n=auto test/backend --durations=20
run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
- name: Run TRANSCENDENTAL math
run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -794,27 +777,27 @@ jobs:
with:
key: metal
deps: testing
python-version: '3.12'
python-version: '3.11'
amd: 'true'
cuda: 'true'
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 NULL backend tests
run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
- name: Run ONNX
run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
- name: Test tensor core ops (fake)
run: METAL=1 DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
run: METAL=1 DEBUG=3 TC=2 python test/test_ops.py TestOps.test_gemm
- name: Test tensor core ops (real)
run: METAL=1 DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
run: METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_big_gemm
- name: Test Beam Search
run: METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
#- name: Fuzz Test linearizer
# run: METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
- name: Run TRANSCENDENTAL math
run: METAL=1 TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
run: METAL=1 TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
- name: Run pytest (amd)
env:
MOCKGPU: 1
@@ -855,14 +838,14 @@ jobs:
key: osx-webgpu
deps: testing
webgpu: 'true'
- name: Test infinity math in WGSL
run: WEBGPU=1 python -m pytest -n=auto test/test_renderer_failures.py::TestWGSLFailures::test_multiply_infinity --durations=20
- name: Build WEBGPU Efficientnet
run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m examples.compile_efficientnet
- name: Run selected webgpu tests
run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m pytest -n=auto test/backend --durations=20
#- name: Clean npm cache
# run: npm cache clean --force
#- name: Install Puppeteer
# run: npm install puppeteer
- name: Clean npm cache
run: npm cache clean --force
- name: Install Puppeteer
run: npm install puppeteer
# this is also flaky
#- name: Run WEBGPU Efficientnet
# run: node test/web/test_webgpu.js
@@ -894,7 +877,8 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: macos-${{ matrix.backend }}-minimal
deps: testing_unit
deps: testing_minimal
pydeps: "capstone"
llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
mesa: ${{ matrix.backend == 'lvp' && 'true' }}
- name: Set env
@@ -904,7 +888,7 @@ jobs:
python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
- name: Run pytest (${{ matrix.backend }})
run: python3 -m pytest -n=auto test/backend --durations=20
run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
- name: Run macOS-specific unit test
@@ -937,16 +921,12 @@ jobs:
- name: Run unit tests
if: matrix.backend=='llvm'
# test_newton_schulz hits RecursionError
run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
- name: Run NULL backend tests
if: matrix.backend=='llvm'
shell: bash
run: CPU=0 CPU_LLVM=0 NULL=1 python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
- name: Run pytest (${{ matrix.backend }})
shell: bash
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
python -m pytest -n=auto test/test_tiny.py test/backend/test_ops.py --durations=20
python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
# ****** Compile-only Tests ******
@@ -965,15 +945,15 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: compile-${{ matrix.backend }}
deps: testing_unit
deps: testing_minimal
mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
python-version: '3.12'
python-version: '3.14'
- name: Set env
shell: bash
run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
- name: Run test_ops
shell: bash
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
python -m pytest -n=auto test/backend/test_ops.py --durations=20
DEBUG=4 python3 test/test_ops.py TestOps.test_add
python -m pytest -n=auto test/test_ops.py --durations=20
+1 -1
View File
@@ -28,7 +28,7 @@ repos:
pass_filenames: false
- id: tests
name: comprehensive test suite
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/unit/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/unit/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
language: system
always_run: true
pass_filenames: false
+17
View File
@@ -0,0 +1,17 @@
# tinygrad agents
Hello agent. You are one of the most talented programmers of your generation.
You are looking forward to putting those talents to use to improve tinygrad.
## philosophy
tinygrad is a **tensor** library focused on beauty and minimalism, while still matching the functionality of PyTorch and JAX.
Every line must earn its keep. Prefer readability over cleverness. We believe that if carefully designed, 10 lines can have the impact of 1000.
Never mix functionality changes with whitespace changes. All functionality changes must be tested.
## style
Use **2-space indentation**, and keep lines to a maximum of **150 characters**. Match the existing style.
+227
View File
@@ -0,0 +1,227 @@
# Claude Code Guide for tinygrad
## Architecture Overview
tinygrad compiles tensor operations into optimized kernels. The pipeline:
1. **Tensor** (`tensor.py`) - User-facing API, creates UOp graph
2. **UOp** (`uop/ops.py`) - Unified IR for all operations (both tensor and kernel level)
3. **Schedule** (`engine/schedule.py`, `schedule/`) - Converts tensor UOps to kernel UOps
4. **Codegen** (`codegen/`) - Converts kernel UOps to device code
5. **Runtime** (`runtime/`) - Device-specific execution
## Key Concepts
### UOp (Universal Operation)
Everything is a UOp - tensors, operations, buffers, kernels. Key properties:
- `op`: The operation type (Ops enum)
- `dtype`: Data type
- `src`: Tuple of source UOps
- `arg`: Operation-specific argument
- `tag`: Optional tag for graph transformations
UOps are **immutable and cached** - creating the same UOp twice returns the same object (ucache).
### PatternMatcher
Used extensively for graph transformations:
```python
pm = PatternMatcher([
(UPat(Ops.ADD, src=(UPat.cvar("x"), UPat.cvar("x"))), lambda x: x * 2),
])
result = graph_rewrite(uop, pm)
```
### Schedule Cache
Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.
## Testing
```bash
# Run specific test
python -m pytest test/unit/test_schedule_cache.py -xvs
# Run with timeout
python -m pytest test/test_symbolic_ops.py -x --timeout=60
# Debug with print
DEBUG=2 python -m pytest test/test_schedule.py::test_name -xvs
# Visualize UOp graphs
VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
```
## Common Environment Variables
- `DEBUG=1-7` - Increasing verbosity (7 shows assembly output)
- `VIZ=1` - Enable graph visualization
- `SPEC=1` - Enable UOp spec verification
- `NOOPT=1` - Disable optimizations
- `DEVICE=CPU/CUDA/AMD/METAL` - Set default device
## Debugging Tips
1. **Print UOp graphs**: `print(tensor.uop)` or `print(tensor.uop.sink())`
2. **Check schedule**: `tensor.schedule()` returns list of ExecItems
3. **Trace graph rewrites**: Use `VIZ=1` or add print in PatternMatcher callbacks
4. **Find UOps by type**: `[u for u in uop.toposort() if u.op is Ops.SOMETHING]`
## Workflow Rules
- **NEVER commit without explicit user approval** - always show the diff and wait for approval
- **NEVER amend commits** - always create a new commit instead
- Run `pre-commit run --all-files` before committing to catch linting/type errors
- Run tests before proposing commits
- Test with `SPEC=2` when modifying UOp-related code
## Auto-generated Files (DO NOT EDIT)
The following files are auto-generated and should never be edited manually:
- `extra/assembly/amd/autogen/{arch}/__init__.py` - Generated by `python -m extra.assembly.amd.dsl --arch {arch}`
- `extra/assembly/amd/autogen/{arch}/gen_pcode.py` - Generated by `python -m extra.assembly.amd.pcode --arch {arch}`
Where `{arch}` is one of: `rdna3`, `rdna4`, `cdna`
To add missing instruction implementations, add them to `extra/assembly/amd/emu.py` instead.
## Style Notes
- 2-space indentation, 150 char line limit
- PatternMatchers should be defined at module level (slow to construct)
- Prefer `graph_rewrite` over manual graph traversal
- UOp methods like `.replace()` preserve tags unless explicitly changed
- Use `.rtag(value)` to add tags to UOps
## Lessons Learned
### UOp ucache Behavior
UOps are cached by their contents - creating a UOp with identical (op, dtype, src, arg) returns the **same object**. This means:
- `uop.replace(tag=None)` on a tagged UOp returns the original untagged UOp if it exists in cache
- Two UOps with same structure are identical (`is` comparison works)
### Spec Validation
When adding new UOp patterns, update `tinygrad/uop/spec.py`. Test with:
```bash
SPEC=2 python3 test/unit/test_something.py
```
Spec issues appear as `RuntimeError: SPEC ISSUE None: UOp(...)`.
### Schedule Cache Key Normalization
The schedule cache strips values from BIND nodes so different bound values (e.g., KV cache positions) hit the same cache entry:
- `pm_pre_sched_cache`: BIND(DEFINE_VAR, CONST) → BIND(DEFINE_VAR) for cache key
- `pm_post_sched_cache`: restores original BIND from context
- When accessing `bind.src[1]`, check `len(bind.src) > 1` first (might be stripped)
- Extract var_vals from `input_buffers` dict after graph_rewrite (avoids extra toposort)
### Avoiding Extra Work
- Use ctx dict from graph_rewrite to collect info during traversal instead of separate toposort
- Only extract var_vals when schedule is non-empty (no kernels = no vars needed)
- PatternMatchers are slow to construct - define at module level, not in functions
### Readability Over Speed
Don't add complexity for marginal performance gains. Simpler code that's slightly slower is often better:
```python
# BAD: "optimized" with extra complexity
if has_afters: # skip toposort if no AFTERs
after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
# GOOD: simple, always works
after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
```
The conditional check adds complexity, potential bugs, and often negligible speedup. Only optimize when profiling shows a real bottleneck.
### Testing LLM Changes
```bash
# Quick smoke test
echo "Hello" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b"
# Check cache hits (should see "cache hit" after warmup)
echo "Hello world" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b" 2>&1 | grep cache
# Test with beam search
echo "Hello" | BEAM=2 python tinygrad/apps/llm.py --model "llama3.2:1b"
```
## Common Patterns
### Graph Transformation
```python
def my_transform(ctx, x):
# Return new UOp or None to skip
return x.replace(arg=new_arg)
pm = PatternMatcher([
(UPat(Ops.SOMETHING, name="x"), my_transform),
])
result = graph_rewrite(input_uop, pm, ctx={})
```
### Finding Variables
```python
# Get all variables in a UOp graph
variables = uop.variables()
# Get bound variable values
var, val = bind_uop.unbind()
```
### Shape Handling
```python
# Shapes can be symbolic (contain UOps)
shape = tensor.shape # tuple[sint, ...] where sint = int | UOp
```
## Performance Optimization
When optimizing tinygrad internals:
1. **Measure wall time, not just call counts** - Reducing `graph_rewrite` calls doesn't always improve wall time. The overhead of conditional checks can exceed the cost of the operation being skipped.
2. **Profile each optimization individually** - Run benchmarks with and without each change to measure actual impact. Use `test/external/external_benchmark_schedule.py` for schedule/rewrite timing.
3. **Early exits in hot paths are effective** - Simple checks like `if self.op is Ops.CONST: return self` in `simplify()` can eliminate many unnecessary `graph_rewrite` calls.
4. **`graph_rewrite` is expensive** - Each call has overhead even for small graphs. Avoid calling it when the result is trivially known (e.g., simplifying a CONST returns itself).
5. **Beware iterator overhead** - Checks like `all(x.op is Ops.CONST for x in self.src)` can be slower than just running the operation, especially for small sequences.
6. **Verify cache hit rates before adding/keeping caches** - Measure actual hit rates with real workloads. A cache with 0% hit rate is pure overhead (e.g., `pm_cache` was removed because the algorithm guarantees each UOp is only passed to `pm_rewrite` once).
7. **Use `TRACK_MATCH_STATS=2` to profile pattern matching** - This shows match rates and time per pattern. Look for patterns with 0% match rate that still cost significant time - these are pure overhead for that workload.
8. **Cached properties beat manual traversal** - `backward_slice` uses `@functools.cached_property`. A DFS with early-exit sounds faster but is actually slower because it doesn't benefit from caching. The cache hit benefit often outweighs algorithmic improvements.
9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.
## Pattern Matching Analysis
**Use the right tool:**
- `TRACK_MATCH_STATS=2` - **Profiling**: identify expensive patterns
- `VIZ=-1` - **Inspection**: see all transformations, what every match pattern does, the before/after diffs
```bash
TRACK_MATCH_STATS=2 PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
```
Output format: `matches / attempts -- match_time / total_time ms -- location`
Key patterns to watch (from ResNet50 benchmark):
- `split_load_store`: ~146ms, 31% match rate - does real work
- `simplify_valid`: ~75ms, 0% match rate in this workload - checks AND ops for INDEX in backward slice
- `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches
Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.
```bash
# Save the trace
VIZ=-1 python test/test_tiny.py TestTiny.test_gemm
# Explore it
./extra/viz/cli.py --help
```
## AMD Performance Counter Profiling
Set VIZ to `-2` to save performance counters traces for the AMD backend.
Use the CLI in `./extra/sqtt/roc.py` to explore the trace.
+1 -1
View File
@@ -192,7 +192,7 @@ For more examples on how to run the full test suite please refer to the [CI work
Some examples of running tests locally:
```sh
python3 -m pip install -e '.[testing]' # install extra deps for testing
python3 test/backend/test_ops.py # just the ops tests
python3 test/test_ops.py # just the ops tests
python3 -m pytest test/ # whole test suite
```
-196
View File
@@ -1,196 +0,0 @@
from tinygrad import Tensor, dtypes, Context, getenv, UOp, fetch
from tinygrad.uop.ops import Ops, PatternMatcher, UPat
from tinygrad.uop.symbolic import symbolic
from tinygrad.codegen import Renderer
from tinygrad.codegen.opt import Opt, OptOps
# ************************* implementation of the problem ************************
def myhash(a: Tensor) -> Tensor:
a = (a + 0x7ED55D16) + (a << 12)
a = (a ^ 0xC761C23C) ^ (a >> 19)
a = (a + 0x165667B1) + (a << 5)
a = (a + 0xD3A2646C) ^ (a << 9)
a = (a + 0xFD7046C5) + (a << 3)
a = (a ^ 0xB55A4F09) ^ (a >> 16)
return a
def select_with_where_tree(values: Tensor, relative_idx: Tensor) -> Tensor:
n = values.shape[0]
if n == 1: return values[0].expand(relative_idx.shape)
mid = n // 2
left = select_with_where_tree(values[:mid], relative_idx)
right = select_with_where_tree(values[mid:], relative_idx - mid)
go_left = relative_idx < mid
return go_left.where(left, right)
def tree_traversal(forest: Tensor, val: Tensor, height: int, rounds: int, where_tree_threshold=3) -> Tensor:
# All walkers start at idx=0
idx = Tensor.zeros(val.shape, device=val.device, dtype=dtypes.uint32)
for r in range(rounds):
level = r % (height + 1)
level_start = (1 << level) - 1
level_size = 1 << level
if level == 0:
# At root (level 0), all walkers are at idx=0
# No gather needed, just broadcast the root value
node_val = forest[0].expand(val.shape)
idx = idx * 0 # Reset to 0
elif level <= where_tree_threshold:
# Small level: use where-tree
level_values = forest[level_start : level_start + level_size]
relative_idx = (idx - level_start)
node_val = select_with_where_tree(level_values, relative_idx)
else:
# Large level: use gather
node_val = forest.gather(0, idx)
val = myhash(val ^ node_val)
idx = (idx << 1) + (1 + (val & 1))
# No wrap check needed! At round 10 (level becomes 0), we reset idx above.
return val.contiguous(arg=(Opt(OptOps.UPCAST, 0, 8),))
# ************************* renderer for VLIW machine *************************
def loop_unrolling(sink:UOp):
rng = [x for x in sink.toposort() if x.op is Ops.RANGE]
if len(rng) == 0: return None
print(f"unrolling loop with size {rng[0].vmax+1}")
unrolled_sinks = [sink.substitute({rng[0]:rng[0].const_like(i)}).src[0] for i in range(rng[0].vmax+1)]
return UOp.sink(*unrolled_sinks, arg=sink.arg)
global_addrs = []
vliw_prepare = PatternMatcher([
# loop unrolling (should be a part of tinygrad)
(UPat(Ops.SINK, name="sink"), loop_unrolling),
# cast is fake
(UPat(Ops.CAST, name="c"), lambda c: c.src[0]),
# rewrites to hardcode the addresses in memory
(UPat(Ops.PARAM, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
# INDEX is just plus
(UPat(Ops.INDEX, name="i"), lambda i: i.src[0]+i.src[1]),
])+symbolic
class VLIWRenderer(Renderer):
has_local = False # TODO: this should be the default / cleaned up
# this says this backend supports MULACC + more. decompositions uses this
code_for_op: dict = {Ops.MULACC: None, Ops.ADD: "+", Ops.MUL: "*",
Ops.XOR: "^", Ops.AND: "&", Ops.OR: "|",
Ops.SHL: "<<", Ops.SHR: ">>", Ops.CMPLT: "<"}
# this matcher runs while still in graph form
pre_matcher = vliw_prepare
def render(self, uops:list[UOp]):
# TODO: this is a minimal renderer. for low cycle count, make it good
# to get speed, you need to add VLIW packing
# to get under 1536 regs, you need to add a register allocator
# we left the fun parts to you
print(f"rendering with {len(uops)} uops")
reg, inst = 0, []
r: dict[UOp, int] = {}
for u in uops:
assert u.dtype.count in (1,8), "dtype count must be 1 or 8"
# dumb register allocator
if u.op not in {Ops.STORE, Ops.SINK, Ops.GEP}:
r[u] = reg
reg += u.dtype.count
# render UOps to instructions
match u.op:
case Ops.SINK:
inst.append({"flow": [("halt",)]})
case Ops.CONST:
inst.append({"load": [("const", r[u], u.arg)]})
case Ops.GEP:
# a GEP is just an alias to a special register in the vector
r[u] = r[u.src[0]] + u.arg[0]
case Ops.VECTORIZE:
if all(s == u.src[0] for s in u.src):
# if all sources are the same, we can broadcast
inst.append({"valu": [("vbroadcast", r[u], r[u.src[0]])]})
else:
# this is a copy into a contiguous chunk of registers
inst.extend({"flow": [("add_imm", r[u]+i, r[s], 0)]} for i,s in enumerate(u.src) if r[s] != r[u]+i)
case Ops.LOAD:
op = "vload" if u.dtype.count > 1 else "load"
inst.append({"load": [(op, r[u], r[u.src[0]])]})
case Ops.STORE:
op = "vstore" if u.src[1].dtype.count > 1 else "store"
inst.append({"store": [(op, r[u.src[0]], r[u.src[1]])]})
case Ops.MULACC:
assert u.dtype.count == 8
inst.append({"valu": [("multiply_add", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
case Ops.WHERE:
assert u.dtype.count == 8
inst.append({"flow": [("vselect", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
case _ if u.op in self.code_for_op:
cat = "valu" if u.dtype.count > 1 else "alu"
inst.append({cat: [(self.code_for_op[u.op], r[u], r[u.src[0]], r[u.src[1]])]})
case _:
raise NotImplementedError(f"unhandled op {u.op}")
return repr(inst)
# ************************* test and render *************************
import sys, types
PROBLEM_URL = "https://raw.githubusercontent.com/anthropics/original_performance_takehome/refs/heads/main/tests/frozen_problem.py"
sys.modules["problem"] = problem = types.ModuleType("problem")
exec(fetch(PROBLEM_URL).read_text(), problem.__dict__)
if __name__ == "__main__":
batch_size = getenv("BS", 256)
height = 10
rounds = getenv("ROUNDS", 16)
# build problem
tree = problem.Tree.generate(height)
inp = problem.Input.generate(tree, batch_size, rounds)
mem = problem.build_mem_image(tree, inp)
global_addrs.extend([mem[6], mem[6], mem[4]]) # output, input, forest
# *** verify the kernel in tinygrad compared to reference ***
forest_t = Tensor(tree.values, dtype=dtypes.uint32)
val_t = Tensor(inp.values, dtype=dtypes.uint32)
if getenv("VERIFY", 1):
# verify on normal tinygrad device
with Context(PCONTIG=2):
out = tree_traversal(forest_t, val_t, height, rounds)
val_out = out.tolist()
problem.reference_kernel(tree, inp)
assert val_out == inp.values
print("verification passed")
# *** render to device ***
from tinygrad.codegen import get_program
with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0):
out = tree_traversal(forest_t, val_t, height, rounds)
sink = out.schedule()[-1].ast
prg = get_program(sink, VLIWRenderer())
# *** run on Machine and compare ***
# NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
src = eval(prg.src)
max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else " <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)
machine.run()
print(f"ran for {machine.cycle:5d} cycles" + ("" if machine.cycle <= 1363 else " <-- EVEN CLAUDE GOT 1363"))
# compare to reference
ref_mem = mem.copy()
for _ in problem.reference_kernel2(ref_mem, {}): pass
assert machine.mem[mem[6]:mem[6]+mem[2]] == ref_mem[mem[6]:mem[6]+mem[2]]
print("compare passed!")
+13 -14
View File
@@ -72,7 +72,7 @@ def loader_process(q_in, q_out, X:Tensor, seed):
#storage_tensor._copyin(img_tensor.numpy())
# faster
X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()
X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
# ideal
#X[idx].assign(img.tobytes()) # NOTE: this is slow!
@@ -264,8 +264,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
x = random_brightness_augmentation(x)
x = gaussian_noise(x)
X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = x.tobytes()
Y[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = y.tobytes()
X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
Y[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = y.tobytes()
queue_out.put(idx)
queue_out.put(None)
@@ -379,12 +379,12 @@ def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue
clipped_match_idxs = np.clip(match_idxs, 0, None)
clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]
boxes[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_boxes.tobytes()
labels[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_labels.tobytes()
matches[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = match_idxs.tobytes()
anchors[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = anchor.tobytes()
boxes[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
labels[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
matches[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
anchors[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = anchor.tobytes()
imgs[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()
imgs[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
queue_out.put(idx)
queue_out.put(None)
@@ -552,7 +552,7 @@ class BinIdxDataset:
version, = struct.unpack("<Q", self.idx.read(8))
assert version == 1, "unsupported index version"
dtype_code, = struct.unpack("<B", self.idx.read(1))
self.dtype = {1:np.dtype(np.uint8), 2:np.dtype(np.int8), 3:np.dtype(np.int16), 4:np.dtype(np.int32), 5:np.dtype(np.int64), 6:np.dtype(np.float64), 7:np.dtype(np.double), 8:np.dtype(np.uint16)}[dtype_code]
self.dtype = {1:dtypes.uint8, 2:dtypes.int8, 3:dtypes.int16, 4:dtypes.int32, 5:dtypes.int64, 6:dtypes.float64, 7:dtypes.double, 8:dtypes.uint16}[dtype_code]
self.count, = struct.unpack("<Q", self.idx.read(8))
doc_count, = struct.unpack("<Q", self.idx.read(8))
@@ -569,7 +569,7 @@ class BinIdxDataset:
self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
# bin file
self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin")).numpy()
self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
def _index(self, idx) -> tuple[int, int]:
return int(self.pointers[idx]), int(self.sizes[idx])
@@ -578,7 +578,7 @@ class BinIdxDataset:
ptr, size = self._index(idx)
if length is None: length = size - offset
ptr += offset * self.dtype.itemsize
return self.bin_t[ptr:ptr+length*self.dtype.itemsize].view(self.dtype)
return self.bin_t[ptr:ptr+length*self.dtype.itemsize].bitcast(self.dtype).to(None)
# https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
class GPTDataset:
@@ -637,7 +637,7 @@ class GPTDataset:
sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
# concat all parts
text = np.concatenate(sample_parts, axis=0)
text = Tensor.cat(*sample_parts)
return text
@@ -780,8 +780,7 @@ def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:b
def iterate_llama3_dataset(dataset:BlendedGPTDataset, bs:int):
for b in range(math.ceil(dataset.samples / bs)):
batch = [dataset.get(b * bs + i) for i in range(bs)]
stacked = np.stack(batch, axis=0)
yield Tensor(stacked, device="NPY")
yield Tensor.stack(batch, dim=0)
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)
+42 -101
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, profile_marker
from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
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
@@ -1292,18 +1292,13 @@ 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)
DATA_SEED = config["DATA_SEED"] = getenv("DATA_SEED", SEED)
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
TRAIN_ON_VAL = config["TRAIN_ON_VAL"] = getenv("TRAIN_ON_VAL", 0)
SMALL = config["SMALL"] = getenv("SMALL", 0)
SAMPLES = config["SAMPLES"] = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
EVAL_SAMPLES = config["EVAL_SAMPLES"] = getenv("EVAL_SAMPLES", 5760 if not SMALL else 1024)
MAX_STEPS = config["MAX_STEPS"] = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS))
WARMUP_STEPS = config["WARMUP_STEPS"] = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
LR = config["LR"] = getenv("LR", 8e-5 * GBS / 1152)
END_LR = config["END_LR"] = getenv("END_LR", 8e-7)
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)
@@ -1317,12 +1312,10 @@ def train_llama3():
opt_adamw_weight_decay = 0.1
opt_gradient_clip_norm = 1.0
opt_learning_rate_warmup_steps = WARMUP_STEPS
opt_learning_rate_decay_steps = MAX_STEPS - opt_learning_rate_warmup_steps
opt_base_learning_rate = LR
opt_end_learning_rate = END_LR
Tensor.manual_seed(SEED) # seed for weight initialization
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)
# ** init wandb **
WANDB = getenv("WANDB")
@@ -1372,12 +1365,6 @@ def train_llama3():
optim = AdamW(get_parameters(model), lr=0.0,
b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
# init grads
for p in optim.params:
p.grad = p.zeros_like().contiguous().realize()
grads = [p.grad for p in optim.params]
scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
if resume_ckpt := getenv("RESUME_CKPT"):
@@ -1390,8 +1377,9 @@ def train_llama3():
load_state_dict(scheduler, safe_load(fn), realize=False)
@TinyJit
def minibatch(tokens:Tensor):
tokens = tokens.to(None)
@Tensor.train()
def train_step(model, tokens:Tensor):
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)
@@ -1401,41 +1389,27 @@ def train_llama3():
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
loss.backward()
assert all(p.grad is g for p,g in zip(optim.params, grads))
Tensor.realize(loss, *grads)
return loss.flatten().float().to("CPU")
@TinyJit
def optim_step():
for p in optim.params:
p.grad.assign(p.grad / grad_acc)
# L2 norm grad clip
# https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
# https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
if not getenv("DISABLE_GRAD_CLIP_NORM"):
total_norm = Tensor(0.0, dtype=dtypes.float32, device=optim.params[0].device)
for g in grads:
total_norm += g.float().square().sum()
total_norm = total_norm.sqrt().contiguous().realize()
for g in grads:
g.assign((g * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype)).realize()
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)
optim.step()
scheduler.step()
for g in grads:
g.assign(g.zeros_like().contiguous()).realize()
lr = optim.lr
Tensor.realize(lr, *grads)
return lr.float().to("CPU")
loss.realize(lr)
return loss, lr
@TinyJit
@Tensor.train(False)
def eval_step(tokens:Tensor):
tokens = tokens.to(None)
def eval_step(model, tokens:Tensor):
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
tokens = tokens.shard(device, 0)
@@ -1444,7 +1418,7 @@ def train_llama3():
tokens = tokens.shard(device)
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
return loss.flatten().float().to("CPU")
return loss.flatten().float()
# ** data iters **
def fake_data(bs, samples):
@@ -1456,79 +1430,48 @@ def train_llama3():
return fake_data(BS, SAMPLES)
else:
from examples.mlperf.dataloader import batch_load_llama3
return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=DATA_SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
if getenv("FAKEDATA", 0):
eval_dataset = None
else:
from examples.mlperf.dataloader import get_llama3_dataset
eval_dataset = get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
eval_dataset = get_llama3_dataset(1024 if SMALL else 5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
def get_eval_iter():
if eval_dataset is None:
return fake_data(EVAL_BS, EVAL_SAMPLES)
return fake_data(EVAL_BS, 5760)
from examples.mlperf.dataloader import iterate_llama3_dataset
return iterate_llama3_dataset(eval_dataset, EVAL_BS)
num_params = sum(p.numel() for p in params) - model_params["vocab_size"]*model_params["dim"]
train_iter = get_train_iter()
iter = get_train_iter()
i, sequences_seen = resume_ckpt, 0
step_times = []
while i < MAX_STEPS:
for tokens in tqdm(iter, total=SAMPLES//GBS):
GlobalCounters.reset()
if getenv("TRAIN", 1):
profile_marker(f"train @ {i}")
st = time.perf_counter()
stopped = False
for _ in range(grad_acc):
ist = time.perf_counter()
try: tokens = next(train_iter)
except StopIteration:
stopped = True
break
dt = time.perf_counter()
loss = minibatch(tokens)
if stopped: break
gt = time.perf_counter()
lr = optim_step()
ot = time.perf_counter()
t = time.perf_counter()
loss, lr = train_step(model, tokens)
loss = loss.float().item()
lr = lr.item()
et = time.perf_counter()
step_time = et - st
gbs_time = gt - st
optim_time = ot - gt
data_time = dt - ist
dev_time = step_time - data_time * grad_acc
if BENCHMARK: step_times.append(step_time)
i += 1
sequences_seen += GBS
sequences_seen += tokens.shape[0]
sec = time.perf_counter()-t
if BENCHMARK: step_times.append(sec)
mem_gb = GlobalCounters.mem_used / 1e9
gflops = GlobalCounters.global_ops / 1e9 / dev_time
mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
gflops = GlobalCounters.global_ops / 1e9 / sec
tqdm.write(
f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS")
if (fname:=getenv("LOSS_FILE", "")):
with open(fname, "a") as f:
f.write(f"{i} {loss:.4f} {lr:.12f} {mem_gb:.2f}\n")
if WANDB:
wandb.log({
"lr": lr, "train/loss": loss,
"train/step_time": step_time,
"train/gbs_time": gbs_time,
"train/optim_time": optim_time,
"train/dev_time": dev_time,
"train/data_time": data_time,
"train/mem": mem_gb,
"train/GFLOPS": gflops,
"train/MFU": mfu,
"train/sequences_seen": sequences_seen
})
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")
@@ -1548,19 +1491,17 @@ def train_llama3():
f"epoch global_mem: {GlobalCounters.global_mem:_}")
if (sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
if EVAL_BS == 0: return
tqdm.write(f"evaluating after {sequences_seen} sequences")
profile_marker(f"eval @ {i}")
# run eval
eval_losses = []
eval_iter = get_eval_iter()
tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
eval_losses += eval_step(tokens).tolist()
if BENCHMARK and (j+1) == min(BENCHMARK, EVAL_SAMPLES//EVAL_BS):
for j,tokens in tqdm(enumerate(eval_iter), total=5760//EVAL_BS):
eval_losses += eval_step(model, tokens).tolist()
if BENCHMARK and (j+1) == min(BENCHMARK, 5760//EVAL_BS):
return
log_perplexity = Tensor(eval_losses).mean().float().item()
@@ -1648,7 +1589,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():
@@ -1687,7 +1628,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,
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
export CHECK_OOB=0
export IGNORE_OOB=1
export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
@@ -5,7 +5,7 @@ export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -8,7 +8,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
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 CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -11,7 +11,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
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 CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -8,7 +8,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
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 CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=5000000
export BEAM=0 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -8,7 +8,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
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 CHECK_OOB=0
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
@@ -11,7 +11,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
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 CHECK_OOB=0
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
@@ -4,7 +4,7 @@ export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -7,7 +7,7 @@ export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_green"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -7,7 +7,7 @@ export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_red"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -1,37 +0,0 @@
#!/usr/bin/env bash
export PYTHONPATH="."
export DEV=${DEV:-AMD}
export EMULATE="AMD_CDNA4"
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
export DEBUG=${DEBUG:-2}
export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
export ALL2ALL=${ALL2ALL:-1}
export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
export GBS=$((BS * GRADIENT_ACC_STEPS))
export MODEL="llama3"
export BASEDIR="/raid/datasets/c4-8b/"
export SMALL=1
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
export EVAL_TARGET=3.3 EVAL_FREQ=12288
export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
export SAMPLES=$((MAX_STEPS * GBS))
export SEQLEN=${SEQLEN:-8192}
export SEED=${SEED:-5760}
export DATA_SEED=${DATA_SEED:-5760}
export JITBEAM=${JITBEAM:-3}
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export FAKEDATA=1 BENCHMARK=10 LLAMA_LAYERS=2
python3 examples/mlperf/model_train.py
@@ -1,35 +1,26 @@
#!/usr/bin/env bash
export PYTHONPATH="."
export DEV=${DEV:-AMD}
export EMULATE="AMD_CDNA4"
export CHECK_OOB=0
export PYTHONPATH="." AMD=1
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
export DEBUG=${DEBUG:-0}
export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
export ALL2ALL=${ALL2ALL:-1}
export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export FLASH_ATTENTION=1
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
export GBS=$((BS * GRADIENT_ACC_STEPS))
export DP=8 BS=8 EVAL_BS=8
export MODEL="llama3"
export BASEDIR="/raid/datasets/c4-8b/"
export SMALL=1
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8b"}
export EVAL_TARGET=3.3 EVAL_FREQ=12288
export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
export SAMPLES=$((MAX_STEPS * GBS))
export SEQLEN=${SEQLEN:-8192}
export LR="1e-3" END_LR="1e-4" WARMUP_STEPS=1024 MAX_STEPS=1200000
export SAMPLES=$((MAX_STEPS * BS))
export SEED=${SEED:-$RANDOM}
export DATA_SEED=${DATA_SEED:-5760}
export SEED=5760
export JITBEAM=${JITBEAM:-3}
export JITBEAM=3
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
python3 examples/mlperf/model_train.py
@@ -1,6 +0,0 @@
#!/bin/bash
export BENCHMARK=5
export EVAL_BS=0
export VIZ=${VIZ:--1}
examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
@@ -1,10 +0,0 @@
#!/bin/bash
export BENCHMARK=5
export EVAL_BS=0
export FAKEDATA=1
export NULL_ALLOW_COPYOUT=1
export HIP_VISIBLE_DEVICES=""
export DEV=NULL
export JITBEAM=0
export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+1 -1
View File
@@ -6,6 +6,7 @@ import argparse, time
from collections import namedtuple
from typing import Dict, Any
from PIL import Image
import numpy as np
from tinygrad import Device, GlobalCounters, dtypes, Tensor, TinyJit
from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten, profile_marker
@@ -335,7 +336,6 @@ if __name__ == "__main__":
print(x.shape)
profile_marker("save image")
from PIL import Image
im = Image.fromarray(x.numpy())
print(f"saving {args.out}")
im.save(args.out)
@@ -48,7 +48,7 @@ def prepare_browser_chunks(model):
weight_metadata = metadata.get(name, default)
weight_metadata["parts"][part_num] = {"file": i, "file_start_pos": cursor, "size": size}
metadata[name] = weight_metadata
data = bytes(state_dict[name].uop.base.realized.as_memoryview())
data = bytes(state_dict[name].uop.base.realized.as_buffer())
data = data if not offsets else data[offsets[0]:offsets[1]]
writer.write(data)
cursor += size
+1 -1
View File
@@ -93,7 +93,7 @@ if __name__ == "__main__":
forward: Any = None
sub_steps = [
Step(name = "textModel", input = [Tensor.randint(1, 77, low=0, high=49408, dtype=dtypes.int32)], forward = model.cond_stage_model.transformer.text_model),
Step(name = "textModel", input = [Tensor.randn(1, 77)], forward = model.cond_stage_model.transformer.text_model),
Step(name = "diffusor", input = [Tensor.randn(1, 77, 768), Tensor.randn(1, 77, 768), Tensor.randn(1,4,64,64), Tensor.rand(1), Tensor.randn(1), Tensor.randn(1), Tensor.randn(1)], forward = model),
Step(name = "decoder", input = [Tensor.randn(1,4,64,64)], forward = model.decode),
Step(name = "f16tof32", input = [Tensor.randn(2097120, dtype=dtypes.uint32)], forward = u32_to_f16)
+3 -4
View File
@@ -92,7 +92,7 @@ class SMICtx:
self.prev_terminal_width = 0
self.prev_terminal_height = 0
remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
self.lspci = {l.split()[0]: l.split(" ", 1)[1] for l in lspci}
for k,v in self.lspci.items():
@@ -153,8 +153,7 @@ class SMICtx:
tables = {}
for dev in self.devs:
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableV0_t
case (13,0,12): table_t = dev.smu.smu_mod.MetricsTableV2_t
case (13,0,6)|(13,0,12): table_t = dev.smu.smu_mod.MetricsTableX_t
case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
return tables
@@ -280,7 +279,7 @@ class SMICtx:
device_line = [f"{bold(dev.pcibus)} {trim(self.lspci[dev.pcibus[5:]], col_size - 20)}"] + [pad("", col_size)]
activity_line = [f"GFX Activity {draw_bar(self.get_gfx_activity(dev, metrics) / 100, activity_line_width)}"] \
+ [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
+ [f"MEM Usage {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \
+ [f"MEM Usage {draw_bar((mem_used / mem_total) / 100, activity_line_width, opt_text=mem_fmt)}"] \
temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]
+1 -7
View File
@@ -1,18 +1,12 @@
#!/usr/bin/env python3
import os
from tinygrad.helpers import Context
from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
from tinygrad.runtime.support.hcq import FileIOInterface
from tinygrad.runtime.support.am.amdev import AMDev
if __name__ == "__main__":
gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
for gpu in gpus:
drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
pcidevs = [PCIDevice("AM", gpu, bars=[0, 2, 5]) for gpu in gpus]
pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
amdevs = []
with Context(DEBUG=2):
for pcidev in pcidevs:
+2 -27
View File
@@ -19,9 +19,8 @@ amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
amdev.AMDev = AMDFake
from tinygrad.runtime.ops_amd import PCIIface
def parse_amdgpu_logs(log_content, register_names=None, register_objects=None, *, only_xcc0: bool = False):
def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
register_map = register_names or {}
register_objs = register_objects or {}
def replace_register(match):
reg = match.group(1)
@@ -38,28 +37,6 @@ def parse_amdgpu_logs(log_content, register_names=None, register_objects=None, *
# remove timing prefix
processed_log = re.sub(r'^\[\s*\d+(?:\.\d+)?\]\s*', '', processed_log, flags=re.MULTILINE)
# decode register values into field dicts
def decode_value(match):
reg_name = match.group(1)
xcc_part = match.group(2) # "xcc=0 " or ""
val_str = match.group(3)
val = int(val_str, 16)
reg_obj = register_objs.get(reg_name)
if reg_obj is not None and reg_obj.fields:
fields = reg_obj.decode(val)
# show raw for unaccounted bits
accounted = 0
for name, (start, end) in reg_obj.fields.items():
accounted |= (((1 << (end - start + 1)) - 1) << start)
unaccounted = val & ~accounted
parts = {k: v for k, v in fields.items() if v != 0}
if unaccounted: parts['_raw_unaccounted'] = hex(unaccounted)
return f"register {reg_name}, {xcc_part}with value {val_str} {parts}"
return match.group(0)
processed_log = re.sub(r'register (reg\w+), ((?:xcc=\d+ )?)with value (0x[0-9a-fA-F]+)', decode_value, processed_log)
# keep only xcc=0 lines (but keep lines with no xcc at all)
if only_xcc0:
kept = []
@@ -73,18 +50,16 @@ def main():
only_xcc0 = bool(getenv("ONLY_XCC0", 0))
reg_names = {}
reg_objs = {}
dev = PCIIface(None, 0)
for x, y in dev.dev_impl.__dict__.items():
if isinstance(y, AMRegister):
for xcc, addr in y.addr.items():
reg_names[addr] = f"{x}, xcc={xcc}"
reg_objs[x] = y
with open(sys.argv[1], 'r') as f:
log_content = f.read()
processed_log = parse_amdgpu_logs(log_content, reg_names, reg_objs, only_xcc0=only_xcc0)
processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
with open(sys.argv[2], 'w') as f:
f.write(processed_log)
@@ -1,17 +1,17 @@
An integrated environment for AMD GPU assembly and emulation
Test with `pytest -n12 test/amd/`
`AMD_LLVM=1 pytest -n12 test/amd/`
Test with `PYTHONPATH="." pytest -n12 extra/assembly/amd/`
`AMD_LLVM=1 PYTHONPATH="." pytest -n12 extra/assembly/amd/`
* pdf.py -- extract assembly format + instruction pseudocode from AMD PDF
* dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
* test/mockgpu/amd/emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
* generate.py -- extract assembly format + instruction pseudocode from AMD XML + PDF
* test/mockgpu/amd/pcode.py -- pseudocode to UOp transformation
* sqtt.py -- SQTT parser
* pcode.py -- pseudocode execution environment. pseudocode should be transformed as little as possible.
* asm.py -- an asm/disasm function to transform to and from AMD assembly syntax
* emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
The code should be as readable and deduplicated as possible. emu (in test/mockgpu/amd/) shouldn't be required for dsl.
The code should be as readable and deduplicated as possible. asm and emu shouldn't be required for dsl.
The autogen folder is autogenerated from the AMD PDFs with `python3 -m tinygrad.renderer.amd.pdf --arch all`
The autogen folder is autogenerated from the AMD PDFs with `python3 -m extra.assembly.amd.pdf --arch all`
test_emu.py has a good set of instruction tests for the emulation, with USE_HW=1 it will compare to real hardware.
Whenever an instruction is fixed, regression tests should be added here and confirmed with real hardware.
@@ -20,20 +20,20 @@ test_llvm.py tests asm/disasm on the LLVM tests, confirming it behaves the same
tinygrad's dtype tests should pass with and without LLVM. they run in about 12 seconds.
`AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
`AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
The ops tests also pass, but they are very slow, so you should run them one at a time.
`SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/backend/test_ops.py`
`SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/backend/test_ops.py`
`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_ops.py`
`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_ops.py`
When something is caught by main tinygrad tests, a local regression test should be added to `test/amd`.
When something is caught by main tinygrad tests, a local regression test should be added to `extra/assembly/amd/test`.
While working with tinygrad, you can dump the assembly with `DEBUG=7`. These tests all pass on real hardware
If a test is failing with `AMD=1 PYTHON_REMU=1 MOCKGPU=1` it's because an instruction is emulated incorrectly.
You can test without `MOCKGPU=1` to test on real hardware, if it works on real hardware there's a bug in the emulator.
IMPORTANT: if a test is failing in the emulator, it's an instruction bug. Use DEBUG=7, get the instructions, and debug.
Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~3000 lines.
Get line count with `cloc --by-file tinygrad/renderer/amd/*.py`
Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~2000 lines.
Get line count with `cloc --by-file extra/assembly/amd/*.py`
@@ -1,7 +1,7 @@
# AMD ISA code generator - generates enum.py, ins.py, operands.py, str_pcode.py
# Sources: XML from https://gpuopen.com/download/machine-readable-isa/latest/
# PDF manuals from AMD documentation
import re, zlib, xml.etree.ElementTree as ET, zipfile, pathlib
import re, zlib, xml.etree.ElementTree as ET, zipfile
from tinygrad.helpers import fetch
# ═══════════════════════════════════════════════════════════════════════════════
@@ -22,8 +22,6 @@ FIXES = {"rdna3": {"SOPK": {22: "S_SUBVECTOR_LOOP_BEGIN", 23: "S_SUBVECTOR_LOOP_
"rdna4": {"SOP1": {80: "S_GET_BARRIER_STATE", 81: "S_BARRIER_INIT", 82: "S_BARRIER_JOIN"}, "SOPP": {9: "S_WAITCNT", 21: "S_BARRIER_LEAVE"}},
"cdna": {"DS": {152: "DS_GWS_SEMA_RELEASE_ALL", 154: "DS_GWS_SEMA_V", 156: "DS_GWS_SEMA_P"},
"VOP3P": {44: "V_MFMA_LD_SCALE_B32", 62: "V_MFMA_F32_16X16X8_XF32", 63: "V_MFMA_F32_32X32X4_XF32"}}}
# Fields missing from XML but present in hardware (format: {arch: {encoding: [(name, hi, lo), ...]}})
FIELD_FIXES = {"cdna": {"VOP3P": [("opsel_hi2", 14, 14)]}}
# Encoding suffixes to strip (variants we don't generate separate classes for)
_ENC_SUFFIXES = ("_NSA1",)
# Encoding suffix to class suffix mapping (for variants we DO generate)
@@ -77,13 +75,8 @@ def parse_xml(filename: str):
for ot in root.findall(".//OperandTypes/OperandType"):
ot_name = ot.findtext("OperandTypeName")
for field in ot.findall(".//Field"):
key = (ot_name, field.findtext("FieldName"))
if (enum_name := op_enum_map.get(key)): # type: ignore[arg-type]
def _pv_val(pv: ET.Element) -> tuple[int, str]:
v, n = pv.findtext("Value"), pv.findtext("Name")
assert v is not None and n is not None
return int(v), n.upper()
enums[enum_name] = dict(_pv_val(pv) for pv in field.findall(".//PredefinedValue"))
if (enum_name := op_enum_map.get((ot_name, field.findtext("FieldName")))):
enums[enum_name] = {int(pv.findtext("Value")): pv.findtext("Name").upper() for pv in field.findall(".//PredefinedValue")}
# Extract DataFormats with BitCount
for df in root.findall("ISA/DataFormats/DataFormat"):
name, bits = df.findtext("DataFormatName"), df.findtext("BitCount")
@@ -91,26 +84,17 @@ def parse_xml(filename: str):
# Extract encoding definitions
for enc in root.findall("ISA/Encodings/Encoding"):
name = enc.findtext("EncodingName")
assert name is not None
is_base = name.startswith("ENC_") or name in ("VOP3_SDST_ENC", "VOPDXY")
is_variant = any(sfx in name for sfx in _ENC_SUFFIX_MAP)
if not is_base and not is_variant: continue
if any(s in name for s in _SKIP_ENCODINGS): continue
fields: list[tuple[str, int, int]] = []
for f in enc.findall(".//MicrocodeFormat/BitMap/Field"):
br = f.find("BitLayout/Range")
if br is None: continue
fn = f.findtext("FieldName")
assert fn is not None
fields.append((_norm_field(fn.lower()),
int(br.findtext("BitOffset") or 0) + int(br.findtext("BitCount") or 0) - 1, int(br.findtext("BitOffset") or 0)))
ident_list = enc.findall("EncodingIdentifiers/EncodingIdentifier")
ident = ident_list[0] if ident_list else None
fields = [(_norm_field(f.findtext("FieldName").lower()), int(f.find("BitLayout/Range").findtext("BitOffset") or 0) + int(f.find("BitLayout/Range").findtext("BitCount") or 0) - 1,
int(f.find("BitLayout/Range").findtext("BitOffset") or 0))
for f in enc.findall(".//MicrocodeFormat/BitMap/Field") if f.find("BitLayout/Range") is not None]
ident = (enc.findall("EncodingIdentifiers/EncodingIdentifier") or [None])[0]
enc_field = next((f for f in fields if f[0] == "encoding"), None)
# For multi-dword formats, encoding field may be in higher dword but identifier is always in dword0; use % 32
enc_bits: str | None = None
if ident is not None and ident.text is not None and enc_field:
enc_bits = "".join(ident.text[len(ident.text)-1-b] for b in range(enc_field[1] % 32, (enc_field[2] % 32)-1, -1))
# For multi-dword formats, encoding field may be in higher dword but identifier pattern is always in dword0; use % 32
enc_bits = "".join(ident.text[len(ident.text)-1-b] for b in range(enc_field[1] % 32, (enc_field[2] % 32)-1, -1)) if ident is not None and enc_field else None
base_name = _strip_enc(name)
encodings[NAME_MAP.get(base_name, base_name)] = (fields, enc_bits)
# Extract instruction opcodes and operand info
@@ -118,12 +102,9 @@ def parse_xml(filename: str):
opcode_encs: dict[str, dict[int, set[str]]] = {} # {base_fmt: {opcode: {enc_names}}}
for instr in root.findall("ISA/Instructions/Instruction"):
name = instr.findtext("InstructionName")
assert name is not None
for enc in instr.findall("InstructionEncodings/InstructionEncoding"):
if enc.findtext("EncodingCondition") != "default": continue
enc_enc_name = enc.findtext("EncodingName")
assert enc_enc_name is not None
base, opcode = _map_flat(_strip_enc(enc_enc_name), name), int(enc.findtext("Opcode") or 0)
base, opcode = _map_flat(_strip_enc(enc.findtext("EncodingName")), name), int(enc.findtext("Opcode") or 0)
enc_name = NAME_MAP.get(base, base)
# Encoding variants use the same Op enum as the base format
base_enum = enc_name
@@ -137,21 +118,19 @@ def parse_xml(filename: str):
elif base == "VGLOBAL": enums.setdefault("VFLAT", {})[opcode] = name
enums.setdefault(base_enum, {})[opcode] = name
# Extract operand info
op_info: dict[str, tuple[str | None, int, str | None]] = {}
for op in enc.findall("Operands/Operand"):
fn = op.findtext("FieldName")
if fn: op_info[fn.lower()] = (op.findtext("DataFormatName"), int(op.findtext("OperandSize") or 0), op.findtext("OperandType"))
op_info = {op.findtext("FieldName").lower(): (op.findtext("DataFormatName"), int(op.findtext("OperandSize") or 0), op.findtext("OperandType"))
for op in enc.findall("Operands/Operand") if op.findtext("FieldName")}
for fmt, _, otype in op_info.values():
if fmt and fmt not in fmts: fmts[fmt] = 0
if otype: op_types_set.add(otype)
if op_info: types[(name, base_enum)] = op_info
# Find opcodes that only exist in a specific variant encoding (no base format version)
suffix_only_ops: dict[str, dict[str, set[int]]] = {} # {suffix: {base_fmt: {opcodes}}}
# Find opcodes that only exist in _LIT encoding (no base format version)
lit_only_ops: dict[str, set[int]] = {}
for base_fmt, opcodes in opcode_encs.items():
for opcode, encs in opcodes.items():
suffix = next((s for s in _ENC_SUFFIX_MAP.values() if all(s in e for e in encs)), None)
if suffix is not None: suffix_only_ops.setdefault(suffix, {}).setdefault(base_fmt, set()).add(opcode)
return encodings, enums, types, fmts, op_types_set, suffix_only_ops
if all("_LIT" in e for e in encs):
lit_only_ops.setdefault(base_fmt, set()).add(opcode)
return encodings, enums, types, fmts, op_types_set, lit_only_ops
# ═══════════════════════════════════════════════════════════════════════════════
# PDF parsing
@@ -162,9 +141,7 @@ def extract_pdf_text(url: str) -> list[list[tuple[float, float, str, str]]]:
data = fetch(url).read_bytes()
# Parse xref table to locate objects
xref: dict[int, int] = {}
xref_match = re.search(rb'startxref\s+(\d+)', data)
assert xref_match is not None
pos = int(xref_match.group(1)) + 4
pos = int(re.search(rb'startxref\s+(\d+)', data).group(1)) + 4
while data[pos:pos+7] != b'trailer':
while data[pos:pos+1] in b' \r\n': pos += 1
line_end = data.find(b'\n', pos)
@@ -185,19 +162,14 @@ def extract_pdf_text(url: str) -> list[list[tuple[float, float, str, str]]]:
if not (m := re.search(rb'/Contents (\d+) 0 R', data[xref[n]:xref[n]+500])): continue
stream = get_stream(int(m.group(1))).decode('latin-1')
elements, font = [], ''
_RE_BT = (r'(/F[\d.]+) [\d.]+ Tf|([\d.+-]+) ([\d.+-]+) Td|[\d.+-]+ [\d.+-]+ [\d.+-]+ [\d.+-]+ ([\d.+-]+) ([\d.+-]+) Tm'
r'|<([0-9A-Fa-f]+)>.*?Tj|\[([^\]]+)\] TJ')
for bt in re.finditer(r'BT(.*?)ET', stream, re.S):
x, y = 0.0, 0.0
for sm in re.finditer(_RE_BT, bt.group(1)):
if sm.group(1): font = sm.group(1)
elif sm.group(2): x, y = x + float(sm.group(2)), y + float(sm.group(3))
elif sm.group(4): x, y = float(sm.group(4)), float(sm.group(5))
elif sm.group(6) and (t := bytes.fromhex(sm.group(6)).decode('latin-1')).strip():
elements.append((x, y, t, font))
elif sm.group(7):
t = ''.join(bytes.fromhex(h).decode('latin-1') for h in re.findall(r'<([0-9A-Fa-f]+)>', sm.group(7)))
if t.strip(): elements.append((x, y, t, font))
for m in re.finditer(r'(/F[\d.]+) [\d.]+ Tf|([\d.+-]+) ([\d.+-]+) Td|[\d.+-]+ [\d.+-]+ [\d.+-]+ [\d.+-]+ ([\d.+-]+) ([\d.+-]+) Tm|<([0-9A-Fa-f]+)>.*?Tj|\[([^\]]+)\] TJ', bt.group(1)):
if m.group(1): font = m.group(1)
elif m.group(2): x, y = x + float(m.group(2)), y + float(m.group(3))
elif m.group(4): x, y = float(m.group(4)), float(m.group(5))
elif m.group(6) and (t := bytes.fromhex(m.group(6)).decode('latin-1')).strip(): elements.append((x, y, t, font))
elif m.group(7) and (t := ''.join(bytes.fromhex(h).decode('latin-1') for h in re.findall(r'<([0-9A-Fa-f]+)>', m.group(7)))).strip(): elements.append((x, y, t, font))
pages.append(sorted(elements, key=lambda e: (-e[1], e[0])))
return pages
@@ -223,7 +195,7 @@ def extract_pcode(pages: list[list[tuple[float, float, str, str]]], name_to_op:
else:
next_page, next_y = page_idx, 0
# Collect F6 text from current position to next instruction (pseudocode is at x ≈ 69)
lines: list[tuple[int, float, str]] = []
lines = []
for p in range(page_idx, next_page + 1):
start_y = y if p == page_idx else 800
end_y = next_y if p == next_page else 0
@@ -246,12 +218,8 @@ def extract_pcode(pages: list[list[tuple[float, float, str, str]]], name_to_op:
# Code generation
# ═══════════════════════════════════════════════════════════════════════════════
def write_common(all_fmts: dict[str, int], all_op_types: set[str], path: pathlib.Path) -> None:
lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
lines.append("class ReprEnum(Enum):")
lines.append(' """Enum with clean repr that roundtrips with eval()."""')
lines.append(' def __repr__(self): return f"{type(self).__name__}.{self.name}"')
lines.append("")
def write_common(all_fmts, all_op_types, path):
lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
lines.append("class Fmt(Enum):")
for fmt in sorted(all_fmts.keys()): lines.append(f" {fmt} = auto()")
lines.append("")
@@ -264,12 +232,11 @@ def write_common(all_fmts: dict[str, int], all_op_types: set[str], path: pathlib
with open(path, "w") as f: f.write("\n".join(lines))
def write_enum(enums, path):
lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit",
"from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401", ""]
lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum", "from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401", ""]
for name, ops in sorted(enums.items()):
if not ops: continue
suffix = "_E32" if name in ("VOP1", "VOP2", "VOPC") else "_E64" if name == "VOP3" else ""
lines.append(f"class {name}(ReprEnum):" if name in ("HWREG", "MSG") else f"class {name}Op(ReprEnum):")
lines.append(f"class {name}(Enum):" if name in ("HWREG", "MSG") else f"class {name}Op(Enum):")
aliases = []
for op, mem in sorted(ops.items()):
msuf = suffix if name != "VOP3" or op < 512 else ""
@@ -279,8 +246,8 @@ def write_enum(enums, path):
lines.append("")
with open(path, "w") as f: f.write("\n".join(lines))
def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
_VGPR_FIELDS = {"vdst", "vdstx", "vsrc0", "vsrc1", "vsrc2", "vsrc3", "vsrcx1", "vsrcy1", "vaddr", "vdata", "data", "data0", "data1", "addr", "vsrc"}
def write_ins(encodings, enums, lit_only_ops, types, arch, path):
_VGPR_FIELDS = {"vdst", "vdstx", "vsrc0", "vsrc1", "vsrc2", "vsrc3", "vsrcx1", "vsrcy1", "vaddr", "vdata", "data", "data0", "data1", "addr"}
_VARIANT_SUFFIXES = ("_LIT", "_DPP16", "_DPP8", "_SDWA_SDST", "_SDWA", "_MFMA")
def get_base_fmt(fmt):
for sfx in _VARIANT_SUFFIXES: fmt = fmt.replace(sfx, "")
@@ -300,8 +267,6 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
if name.startswith("ssrc") and bits == 8: return f"SSrcField({hi}, {lo})"
if name in ("saddr", "soffset") and bits == 8: return f"SSrcField({hi}, {lo}, default=NULL)"
if name.startswith("src") and bits == 9: return f"SrcField({hi}, {lo})"
# GLOBAL/SCRATCH: offset is 13-bit signed [12:0], FLAT: 12-bit unsigned (XML has 12-bit for all)
if name == "offset" and base_fmt in ("GLOBAL", "SCRATCH"): return f"BitField(12, {lo})"
if base_fmt == "VOP3P" and name == "opsel_hi": return f"BitField({hi}, {lo}, default=3)"
if base_fmt == "VOP3P" and name == "opsel_hi2": return f"BitField({hi}, {lo}, default=1)"
return f"BitField({hi}, {lo})"
@@ -313,7 +278,7 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
'dpp', 'fi', 'bc', 'row_mask', 'bank_mask', 'src0_neg', 'src0_abs', 'src1_neg', 'src1_abs',
'cbsz', 'abid', 'acc_cd', 'acc', 'blgp', 'lane_sel_0', 'lane_sel_1', 'lane_sel_2', 'lane_sel_3',
'lane_sel_4', 'lane_sel_5', 'lane_sel_6', 'lane_sel_7', 'dst_sel', 'dst_unused', 'src0_sel', 'src1_sel']
def sort_fields(fields): return sorted(fields, key=lambda f: (ORDER.index(f[0]) if f[0] in ORDER else 999, f[2]))
sort_fields = lambda fields: sorted(fields, key=lambda f: (ORDER.index(f[0]) if f[0] in ORDER else 999, f[2]))
# Separate base encodings from variants
base_encodings, variant_encodings = {}, {}
@@ -323,29 +288,15 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
else: variant_encodings[enc_name] = data
# Build sets of ops by their vdst type from operand metadata
sdst_opcodes: dict[str, set[int]] = {} # ops where vdst is OPR_SREG (writes to SGPR)
sdst_opcodes = {} # ops where vdst is OPR_SREG (writes to SGPR)
for fmt, ops in enums.items():
for op, name in ops.items():
op_types = types.get((name, fmt), {})
vdst_type = op_types.get("vdst", (None, None, None))[2]
if vdst_type == "OPR_SREG": sdst_opcodes.setdefault(fmt, set()).add(op)
# collect only the XxxOp enums that are actually referenced in this arch's instruction definitions
enum_names = sorted(f"{k}Op" for k in enums if enums[k] and k not in ("HWREG", "MSG"))
# also re-export HWREG/MSG enums (plain enums, not instruction format ops)
enum_names += sorted(k for k in enums if k in ("HWREG", "MSG") and enums[k])
# collect DSL field types actually used by scanning generated field definitions
all_field_defs = " ".join(field_def(fn, hi, lo, enc, eb) for enc, (flds, eb) in encodings.items() for fn, hi, lo in flds)
_ALL_DSL = ["BitField", "EnumBitField", "FixedBitField", "NULL", "SBaseField", "SGPRField", "SRsrcField",
"SSrcField", "SrcField", "VDSTYField", "VGPRField"]
dsl_names = ["Inst"] + [n for n in _ALL_DSL if n in all_field_defs]
# also re-export register names so `from ins import *` still provides them to downstream users
_DSL_REGS = ["s", "v", "src", "VCC_LO", "VCC_HI", "VCC", "EXEC_LO", "EXEC_HI", "EXEC", "NULL", "OFF", "M0",
"SCC", "VCCZ", "EXECZ", "ttmp", "INV_2PI", "SDWA", "DPP", "DPP16", "LIT", "SRC_LDS_DIRECT"]
dsl_reexport = sorted(set(dsl_names + _DSL_REGS))
lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit", "# ruff: noqa: E501,F401",
f"from tinygrad.renderer.amd.dsl import {', '.join(dsl_reexport)}",
f"from tinygrad.runtime.autogen.amd.{arch}.enum import {', '.join(enum_names)}", "import functools", ""]
lines = ["# autogenerated from AMD ISA XML - do not edit", "# ruff: noqa: F401,F403",
"from extra.assembly.amd.dsl import *", f"from extra.assembly.amd.autogen.{arch}.enum import *", "import functools", ""]
def fmt_allowed(op_enum: str, ops: set[int]) -> str:
"""Format allowed ops as {EnumName.MEMBER, ...}."""
@@ -354,19 +305,14 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
# Generate base classes first
for enc_name, (fields, enc_bits) in sorted(base_encodings.items()):
# Get lit-only ops for this format (these can't be used in base class)
base_lit_ops = lit_only_ops.get(enc_name, set())
all_ops = set(enums.get(enc_name, {}).keys())
# Get suffix-only ops for this format (these can't be used in base class)
base_suffix_ops = set().union(*(d.get(enc_name, set()) for d in suffix_only_ops.values()))
# Exclude SDST ops from base class (they need VOP1_SDST/VOP3_SDST/VOP3B)
base_allowed = all_ops - base_suffix_ops - sdst_opcodes.get(enc_name, set())
# RDNA3 FLAT/GLOBAL/SCRATCH share encoding bits, differentiated by seg field
# RDNA4 VFLAT/VGLOBAL/VSCRATCH have distinct encoding bits, no seg field needed
has_seg_field = any(fn == "seg" for fn, _, _ in fields)
if enc_name in ("FLAT", "VFLAT") and has_seg_field:
base_allowed = all_ops - base_lit_ops - sdst_opcodes.get(enc_name, set())
if enc_name in ("FLAT", "VFLAT"):
prefix = "V" if enc_name == "VFLAT" else ""
flat_variants = [(f"{prefix}FLAT", 0, f"{prefix}FLATOp"), (f"{prefix}GLOBAL", 2, f"{prefix}GLOBALOp"),
(f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]
for cls, seg, op_enum in flat_variants:
for cls, seg, op_enum in [(f"{prefix}FLAT", 0, f"{prefix}FLATOp"), (f"{prefix}GLOBAL", 2, f"{prefix}GLOBALOp"), (f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]:
cls_ops = set(enums.get(cls, {}).keys())
lines.append(f"class {cls}(Inst):")
for fn, hi, lo in sort_fields(fields):
@@ -374,7 +320,7 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
elif fn == "op": lines.append(f" op = EnumBitField({hi}, {lo}, {op_enum}, {fmt_allowed(op_enum, cls_ops)})")
else: lines.append(f" {fn} = {field_def(fn, hi, lo, cls, enc_bits)}")
lines.append("")
elif enc_name not in ("FLAT_GLOBAL", "FLAT_SCRATCH", "FLAT_GLBL", "DPP", "SDWA"):
elif enc_name not in ("FLAT_GLOBAL", "FLAT_SCRATCH", "FLAT_GLBL", "VGLOBAL", "VSCRATCH", "DPP", "SDWA"):
lines.append(f"class {enc_name}(Inst):")
for fn, hi, lo in sort_fields(fields):
if fn == "op":
@@ -390,19 +336,15 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
if base not in base_encodings: continue # skip if no base class
base_fields = {f[0] for f in base_encodings[base][0]}
extra_fields = [(fn, hi, lo) for fn, hi, lo in fields if fn not in base_fields]
# Check if this is a suffix-only variant
variant_suffix = next((sfx for sfx in _VARIANT_SUFFIXES if enc_name.endswith(sfx)), None)
is_suffix_variant = variant_suffix in suffix_only_ops
is_lit = enc_name.endswith("_LIT")
all_ops = set(enums.get(base, {}).keys())
if extra_fields or is_suffix_variant:
if extra_fields or is_lit:
lines.append(f"class {enc_name}({base}):")
op_field = next((f for f in base_encodings[base][0] if f[0] == "op"), None)
# _LIT classes: override op to allow all opcodes (base excludes lit-only ops)
# other classes override op to only suffix-only opcodes
if op_field and is_suffix_variant:
if op_field and is_lit:
_, hi, lo = op_field
allowed_ops = all_ops if variant_suffix == "_LIT" else suffix_only_ops[variant_suffix][base]
lines.append(f" op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', allowed_ops)})")
lines.append(f" op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', all_ops)})")
for fn, hi, lo in sort_fields(extra_fields):
lines.append(f" {fn} = {field_def(fn, hi, lo, enc_name)}")
lines.append("")
@@ -436,28 +378,23 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
for fmt, ops in sorted(enums.items()):
if fmt not in base_encodings and fmt not in ("GLOBAL", "SCRATCH", "VGLOBAL", "VSCRATCH"): continue
suffix = "_E32" if fmt in ("VOP1", "VOP2", "VOPC") else "_E64" if fmt == "VOP3" else ""
op_to_suffix = {op:suffix for suffix,ops in suffix_only_ops.items() for op in ops.get(fmt, set())}
lit_ops = lit_only_ops.get(fmt, set())
fmt_sdst_ops = sdst_opcodes.get(fmt, set())
for op, name in sorted(ops.items()):
# ADDTID ops are in both FLAT and GLOBAL enums (for pcode); only generate helper for GLOBAL/VGLOBAL
if "ADDTID" in name and fmt in ("FLAT", "VFLAT"): continue
msuf = suffix if fmt != "VOP3" or op < 512 else ""
# Determine class: SDST variants, suffix-specific variants (e.g., _MFMA, _LIT), or base
# Determine class: SDST variants, LIT-only instructions, or base
if fmt == "VOP1" and op in fmt_sdst_ops: cls = "VOP1_SDST"
elif fmt == "VOP3" and (op in fmt_sdst_ops or op < 256): cls = "VOP3_SDST"
elif op_to_suffix.get(op): cls = f"{fmt}{op_to_suffix[op]}"
elif op in lit_ops: cls = f"{fmt}_LIT"
else: cls = fmt
lines.append(f"{name.lower()}{msuf.lower()} = functools.partial({cls}, {fmt}Op.{name}{msuf})")
with open(path, "w") as f: f.write("\n".join(lines))
def write_operands(types: dict, enums: dict, arch: str, path: pathlib.Path) -> None:
def write_operands(types, enums, arch, path):
valid = {(name, fmt) for fmt, ops in enums.items() for name in ops.values()}
# only import enums that are actually used as keys in OPERANDS
used_bases = {eb for (nm, eb) in types if (nm, eb) in valid}
enum_names = sorted(f"{k}Op" for k in used_bases)
lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit",
"from tinygrad.runtime.autogen.amd.common import Fmt, OpType",
f"from tinygrad.runtime.autogen.amd.{arch}.enum import {', '.join(enum_names)}", ""]
lines = ["# autogenerated from AMD ISA XML - do not edit",
"from extra.assembly.amd.autogen.common import Fmt, OpType",
f"from extra.assembly.amd.autogen.{arch}.enum import *", ""]
lines.append("# instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}")
lines.append("OPERANDS = {")
def fmt_val(v):
@@ -470,7 +407,7 @@ def write_operands(types: dict, enums: dict, arch: str, path: pathlib.Path) -> N
lines.append("}")
with open(path, "w") as f: f.write("\n".join(lines))
def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, str]], arch: str, path: pathlib.Path) -> None:
def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, str]], arch: str, path: str):
"""Write str_pcode.py file from extracted pseudocode."""
entries: list[tuple[str, str, int, str]] = []
for fmt_name, ops in enums.items():
@@ -481,7 +418,7 @@ def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, st
entries.append((f"{fmt_name}Op", f"{name}{msuf}", opcode, pcode[(name, opcode)]))
enum_names = sorted(set(e[0] for e in entries))
lines = ["# autogenerated from AMD ISA PDF - do not edit", "# ruff: noqa: E501",
f"from tinygrad.runtime.autogen.amd.{arch}.enum import {', '.join(enum_names)}", "", "PCODE = {"]
f"from extra.assembly.amd.autogen.{arch}.enum import {', '.join(enum_names)}", "", "PCODE = {"]
for enum_name, name, opcode, code in sorted(entries, key=lambda x: (x[0], x[2])):
lines.append(f" {enum_name}.{name}: {code!r},")
lines.append("}")
@@ -492,31 +429,27 @@ def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, st
# ═══════════════════════════════════════════════════════════════════════════════
if __name__ == "__main__":
all_fmts: dict[str, int] = {}
all_op_types: set[str] = set()
arch_data: dict[str, dict] = {}
import pathlib
all_fmts, all_op_types, arch_data = {}, set(), {}
# First pass: parse XML for all architectures
for arch, cfg in ARCHS.items():
print(f"Parsing XML: {cfg['xml']} -> {arch}")
encodings, enums, types, fmts, op_types_set, suffix_only_ops = parse_xml(cfg["xml"])
encodings, enums, types, fmts, op_types_set, lit_only_ops = parse_xml(cfg["xml"])
for fmt, ops in FIXES.get(arch, {}).items(): enums.setdefault(fmt, {}).update(ops)
for fmt, fields in FIELD_FIXES.get(arch, {}).items():
if fmt in encodings: encodings[fmt] = (encodings[fmt][0] + fields, encodings[fmt][1])
arch_data[arch] = {"encodings": encodings, "enums": enums, "types": types, "suffix_only_ops": suffix_only_ops}
arch_data[arch] = {"encodings": encodings, "enums": enums, "types": types, "lit_only_ops": lit_only_ops}
for fmt, bits in fmts.items():
assert fmt not in all_fmts or all_fmts[fmt] == bits, f"FMT_BITS mismatch for {fmt}: {all_fmts[fmt]} vs {bits}"
all_fmts[fmt] = bits
all_op_types.update(op_types_set)
# Write common.py
autogen_base = pathlib.Path(__file__).parents[2] / "runtime" / "autogen" / "amd"
common_path = autogen_base / "common.py"
common_path = pathlib.Path(__file__).parent / "autogen" / "common.py"
write_common(all_fmts, all_op_types, common_path)
print(f"Wrote common.py: {len(all_fmts)} formats, {len(all_op_types)} op types")
# Write per-arch files from XML
for arch, data in arch_data.items():
base = autogen_base / arch
base = pathlib.Path(__file__).parent / "autogen" / arch
write_enum(data["enums"], base / "enum.py")
write_ins(data["encodings"], data["enums"], data["suffix_only_ops"], data["types"], arch, base / "ins.py")
write_ins(data["encodings"], data["enums"], data["lit_only_ops"], data["types"], arch, base / "ins.py")
write_operands(data["types"], data["enums"], arch, base / "operands.py")
print(f" {arch}: {len(data['encodings'])} encodings, {sum(len(v) for v in data['enums'].values())} instructions")
# Second pass: parse PDFs and write pcode
@@ -525,6 +458,6 @@ if __name__ == "__main__":
pages = extract_pdf_text(cfg["pdf"])
name_to_op = {name: op for ops in arch_data[arch]["enums"].values() for op, name in ops.items()}
pcode = extract_pcode(pages, name_to_op)
base = autogen_base / arch
base = pathlib.Path(__file__).parent / "autogen" / arch
write_pcode(pcode, arch_data[arch]["enums"], arch, base / "str_pcode.py")
print(f" {arch}: {len(pcode)} pcode entries")
@@ -1,7 +1,8 @@
# autogenerated from AMD ISA XML - do not edit
from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401
from enum import Enum
from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401
class DSOp(ReprEnum):
class DSOp(Enum):
DS_ADD_U32 = 0
DS_SUB_U32 = 1
DS_RSUB_U32 = 2
@@ -132,7 +133,7 @@ class DSOp(ReprEnum):
DS_READ_B96 = 254
DS_READ_B128 = 255
class FLATOp(ReprEnum):
class FLATOp(Enum):
FLAT_LOAD_UBYTE = 16
FLAT_LOAD_SBYTE = 17
FLAT_LOAD_USHORT = 18
@@ -188,7 +189,7 @@ class FLATOp(ReprEnum):
FLAT_ATOMIC_INC_X2 = 107
FLAT_ATOMIC_DEC_X2 = 108
class GLOBALOp(ReprEnum):
class GLOBALOp(Enum):
GLOBAL_LOAD_UBYTE = 16
GLOBAL_LOAD_SBYTE = 17
GLOBAL_LOAD_USHORT = 18
@@ -251,7 +252,7 @@ class GLOBALOp(ReprEnum):
GLOBAL_LOAD_LDS_DWORDX4 = 125
GLOBAL_LOAD_LDS_DWORDX3 = 126
class HWREG(ReprEnum):
class HWREG(Enum):
HW_REG_MODE = 1
HW_REG_STATUS = 2
HW_REG_TRAPSTS = 3
@@ -277,7 +278,7 @@ class HWREG(ReprEnum):
HW_REG_SQ_PERF_SNAPSHOT_PC_LO = 23
HW_REG_SQ_PERF_SNAPSHOT_PC_HI = 24
class MTBUFOp(ReprEnum):
class MTBUFOp(Enum):
TBUFFER_LOAD_FORMAT_X = 0
TBUFFER_LOAD_FORMAT_XY = 1
TBUFFER_LOAD_FORMAT_XYZ = 2
@@ -295,7 +296,7 @@ class MTBUFOp(ReprEnum):
TBUFFER_STORE_FORMAT_D16_XYZ = 14
TBUFFER_STORE_FORMAT_D16_XYZW = 15
class MUBUFOp(ReprEnum):
class MUBUFOp(Enum):
BUFFER_LOAD_FORMAT_X = 0
BUFFER_LOAD_FORMAT_XY = 1
BUFFER_LOAD_FORMAT_XYZ = 2
@@ -371,7 +372,7 @@ class MUBUFOp(ReprEnum):
BUFFER_ATOMIC_INC_X2 = 107
BUFFER_ATOMIC_DEC_X2 = 108
class SCRATCHOp(ReprEnum):
class SCRATCHOp(Enum):
SCRATCH_LOAD_UBYTE = 16
SCRATCH_LOAD_SBYTE = 17
SCRATCH_LOAD_USHORT = 18
@@ -400,7 +401,7 @@ class SCRATCHOp(ReprEnum):
SCRATCH_LOAD_LDS_SSHORT = 41
SCRATCH_LOAD_LDS_DWORD = 42
class SMEMOp(ReprEnum):
class SMEMOp(Enum):
S_LOAD_DWORD = 0
S_LOAD_DWORDX2 = 1
S_LOAD_DWORDX4 = 2
@@ -486,7 +487,7 @@ class SMEMOp(ReprEnum):
S_ATOMIC_INC_X2 = 171
S_ATOMIC_DEC_X2 = 172
class SOP1Op(ReprEnum):
class SOP1Op(Enum):
S_MOV_B32 = 0
S_MOV_B64 = 1
S_CMOV_B32 = 2
@@ -542,7 +543,7 @@ class SOP1Op(ReprEnum):
S_ANDN2_WREXEC_B64 = 54
S_BITREPLICATE_B64_B32 = 55
class SOP2Op(ReprEnum):
class SOP2Op(Enum):
S_ADD_U32 = 0
S_SUB_U32 = 1
S_ADD_I32 = 2
@@ -597,7 +598,7 @@ class SOP2Op(ReprEnum):
S_PACK_LH_B32_B16 = 51
S_PACK_HH_B32_B16 = 52
class SOPCOp(ReprEnum):
class SOPCOp(Enum):
S_CMP_EQ_I32 = 0
S_CMP_LG_I32 = 1
S_CMP_GT_I32 = 2
@@ -619,7 +620,7 @@ class SOPCOp(ReprEnum):
S_CMP_EQ_U64 = 18
S_CMP_LG_U64 = 19
class SOPKOp(ReprEnum):
class SOPKOp(Enum):
S_MOVK_I32 = 0
S_CMOVK_I32 = 1
S_CMPK_EQ_I32 = 2
@@ -642,7 +643,7 @@ class SOPKOp(ReprEnum):
S_SETREG_IMM32_B32 = 20
S_CALL_B64 = 21
class SOPPOp(ReprEnum):
class SOPPOp(Enum):
S_NOP = 0
S_ENDPGM = 1
S_BRANCH = 2
@@ -676,7 +677,7 @@ class SOPPOp(ReprEnum):
S_ENDPGM_ORDERED_PS_DONE = 30
S_SET_VALU_COEXEC_MODE = 31
class VOP1Op(ReprEnum):
class VOP1Op(Enum):
V_NOP_E32 = 0
V_MOV_B32_E32 = 1
V_READFIRSTLANE_B32_E32 = 2
@@ -854,7 +855,7 @@ class VOP1Op(ReprEnum):
V_PERMLANE32_SWAP_B32 = V_PERMLANE32_SWAP_B32_E32
V_CVT_F32_BF16 = V_CVT_F32_BF16_E32
class VOP2Op(ReprEnum):
class VOP2Op(Enum):
V_CNDMASK_B32_E32 = 0
V_ADD_F32_E32 = 1
V_SUB_F32_E32 = 2
@@ -980,7 +981,7 @@ class VOP2Op(ReprEnum):
V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
V_XNOR_B32 = V_XNOR_B32_E32
class VOP3Op(ReprEnum):
class VOP3Op(Enum):
V_CMP_CLASS_F32_E64 = 16
V_CMPX_CLASS_F32_E64 = 17
V_CMP_CLASS_F64_E64 = 18
@@ -1878,7 +1879,7 @@ class VOP3Op(ReprEnum):
V_ADD_LSHL_U32 = V_ADD_LSHL_U32_E64
V_ADD3_U32 = V_ADD3_U32_E64
class VOP3POp(ReprEnum):
class VOP3POp(Enum):
V_PK_MAD_I16 = 0
V_PK_MUL_LO_U16 = 1
V_PK_ADD_I16 = 2
@@ -1987,11 +1988,11 @@ class VOP3POp(ReprEnum):
V_SMFMAC_F32_32X32X32_FP8_BF8 = 126
V_SMFMAC_F32_32X32X32_FP8_FP8 = 127
class VOP3PX2Op(ReprEnum):
class VOP3PX2Op(Enum):
V_MFMA_SCALE_F32_16X16X128_F8F6F4 = 45
V_MFMA_SCALE_F32_32X32X64_F8F6F4 = 46
class VOP3SDOp(ReprEnum):
class VOP3SDOp(Enum):
V_ADD_CO_U32 = 281
V_SUB_CO_U32 = 282
V_SUBREV_CO_U32 = 283
@@ -2003,7 +2004,7 @@ class VOP3SDOp(ReprEnum):
V_MAD_U64_U32 = 488
V_MAD_I64_I32 = 489
class VOPCOp(ReprEnum):
class VOPCOp(Enum):
V_CMP_CLASS_F32_E32 = 16
V_CMPX_CLASS_F32_E32 = 17
V_CMP_CLASS_F64_E32 = 18
@@ -1,7 +1,7 @@
# autogenerated from AMD ISA XML - do not edit
# ruff: noqa: E501,F401
from tinygrad.renderer.amd.dsl import BitField, DPP, DPP16, EXEC, EXECZ, EXEC_HI, EXEC_LO, EnumBitField, FixedBitField, INV_2PI, Inst, LIT, M0, NULL, OFF, SBaseField, SCC, SDWA, SGPRField, SRC_LDS_DIRECT, SRsrcField, SSrcField, SrcField, VCC, VCCZ, VCC_HI, VCC_LO, VGPRField, s, src, ttmp, v
from tinygrad.runtime.autogen.amd.cdna.enum import DSOp, FLATOp, GLOBALOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3PX2Op, VOP3SDOp, VOPCOp, HWREG
# ruff: noqa: F401,F403
from extra.assembly.amd.dsl import *
from extra.assembly.amd.autogen.cdna.enum import *
import functools
class DS(Inst):
@@ -38,7 +38,7 @@ class GLOBAL(Inst):
addr = VGPRField(39, 32)
data = VGPRField(47, 40)
saddr = SGPRField(54, 48, default=NULL)
offset = BitField(12, 0)
offset = BitField(11, 0)
seg = FixedBitField(15, 14, 2)
acc = BitField(55, 55)
sve = BitField(13, 13)
@@ -53,7 +53,7 @@ class SCRATCH(Inst):
addr = VGPRField(39, 32)
data = VGPRField(47, 40)
saddr = SGPRField(54, 48, default=NULL)
offset = BitField(12, 0)
offset = BitField(11, 0)
seg = FixedBitField(15, 14, 1)
acc = BitField(55, 55)
sve = BitField(13, 13)
@@ -164,7 +164,7 @@ class VOP3(Inst):
class VOP3P(Inst):
encoding = FixedBitField(31, 23, 0b110100111)
op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_PK_MAD_I16, VOP3POp.V_PK_MUL_LO_U16, VOP3POp.V_PK_ADD_I16, VOP3POp.V_PK_SUB_I16, VOP3POp.V_PK_LSHLREV_B16, VOP3POp.V_PK_LSHRREV_B16, VOP3POp.V_PK_ASHRREV_I16, VOP3POp.V_PK_MAX_I16, VOP3POp.V_PK_MIN_I16, VOP3POp.V_PK_MAD_U16, VOP3POp.V_PK_ADD_U16, VOP3POp.V_PK_SUB_U16, VOP3POp.V_PK_MAX_U16, VOP3POp.V_PK_MIN_U16, VOP3POp.V_PK_FMA_F16, VOP3POp.V_PK_ADD_F16, VOP3POp.V_PK_MUL_F16, VOP3POp.V_PK_MIN_F16, VOP3POp.V_PK_MAX_F16, VOP3POp.V_DOT2_F32_BF16, VOP3POp.V_PK_MINIMUM3_F16, VOP3POp.V_PK_MAXIMUM3_F16, VOP3POp.V_MAD_MIX_F32, VOP3POp.V_MAD_MIXLO_F16, VOP3POp.V_MAD_MIXHI_F16, VOP3POp.V_DOT2_F32_F16, VOP3POp.V_DOT2_I32_I16, VOP3POp.V_DOT2_U32_U16, VOP3POp.V_DOT4_I32_I8, VOP3POp.V_DOT4_U32_U8, VOP3POp.V_DOT8_I32_I4, VOP3POp.V_DOT8_U32_U4, VOP3POp.V_MFMA_LD_SCALE_B32, VOP3POp.V_PK_FMA_F32, VOP3POp.V_PK_MUL_F32, VOP3POp.V_PK_ADD_F32, VOP3POp.V_PK_MOV_B32, VOP3POp.V_MFMA_F32_16X16X8_XF32, VOP3POp.V_MFMA_F32_32X32X4_XF32, VOP3POp.V_ACCVGPR_READ, VOP3POp.V_ACCVGPR_WRITE})
op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_PK_MAD_I16, VOP3POp.V_PK_MUL_LO_U16, VOP3POp.V_PK_ADD_I16, VOP3POp.V_PK_SUB_I16, VOP3POp.V_PK_LSHLREV_B16, VOP3POp.V_PK_LSHRREV_B16, VOP3POp.V_PK_ASHRREV_I16, VOP3POp.V_PK_MAX_I16, VOP3POp.V_PK_MIN_I16, VOP3POp.V_PK_MAD_U16, VOP3POp.V_PK_ADD_U16, VOP3POp.V_PK_SUB_U16, VOP3POp.V_PK_MAX_U16, VOP3POp.V_PK_MIN_U16, VOP3POp.V_PK_FMA_F16, VOP3POp.V_PK_ADD_F16, VOP3POp.V_PK_MUL_F16, VOP3POp.V_PK_MIN_F16, VOP3POp.V_PK_MAX_F16, VOP3POp.V_DOT2_F32_BF16, VOP3POp.V_PK_MINIMUM3_F16, VOP3POp.V_PK_MAXIMUM3_F16, VOP3POp.V_MAD_MIX_F32, VOP3POp.V_MAD_MIXLO_F16, VOP3POp.V_MAD_MIXHI_F16, VOP3POp.V_DOT2_F32_F16, VOP3POp.V_DOT2_I32_I16, VOP3POp.V_DOT2_U32_U16, VOP3POp.V_DOT4_I32_I8, VOP3POp.V_DOT4_U32_U8, VOP3POp.V_DOT8_I32_I4, VOP3POp.V_DOT8_U32_U4, VOP3POp.V_MFMA_LD_SCALE_B32, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4, VOP3POp.V_PK_FMA_F32, VOP3POp.V_PK_MUL_F32, VOP3POp.V_PK_ADD_F32, VOP3POp.V_PK_MOV_B32, VOP3POp.V_MFMA_F32_16X16X32_BF16, VOP3POp.V_MFMA_I32_16X16X64_I8, VOP3POp.V_MFMA_F32_32X32X16_BF16, VOP3POp.V_MFMA_I32_32X32X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_BF16, VOP3POp.V_SMFMAC_I32_16X16X128_I8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X8_XF32, VOP3POp.V_MFMA_F32_32X32X4_XF32, VOP3POp.V_MFMA_F32_32X32X1_2B_F32, VOP3POp.V_MFMA_F32_16X16X1_4B_F32, VOP3POp.V_MFMA_F32_4X4X1_16B_F32, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X2_F32, VOP3POp.V_MFMA_F32_16X16X4_F32, VOP3POp.V_SMFMAC_F32_32X32X32_BF16, VOP3POp.V_SMFMAC_I32_32X32X64_I8, VOP3POp.V_MFMA_F32_32X32X4_2B_F16, VOP3POp.V_MFMA_F32_16X16X4_4B_F16, VOP3POp.V_MFMA_F32_4X4X4_16B_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X8_F16, VOP3POp.V_MFMA_F32_16X16X16_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8, VOP3POp.V_MFMA_I32_32X32X4_2B_I8, VOP3POp.V_MFMA_I32_16X16X4_4B_I8, VOP3POp.V_MFMA_I32_4X4X4_16B_I8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8, VOP3POp.V_MFMA_F32_16X16X32_F16, VOP3POp.V_MFMA_F32_32X32X16_F16, VOP3POp.V_MFMA_I32_32X32X16_I8, VOP3POp.V_MFMA_I32_16X16X32_I8, VOP3POp.V_ACCVGPR_READ, VOP3POp.V_ACCVGPR_WRITE, VOP3POp.V_SMFMAC_F32_16X16X64_F16, VOP3POp.V_SMFMAC_F32_32X32X32_F16, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16, VOP3POp.V_MFMA_F32_32X32X8_BF16, VOP3POp.V_MFMA_F32_16X16X16_BF16, VOP3POp.V_SMFMAC_F32_16X16X32_F16, VOP3POp.V_SMFMAC_F32_32X32X16_F16, VOP3POp.V_SMFMAC_F32_16X16X32_BF16, VOP3POp.V_SMFMAC_F32_32X32X16_BF16, VOP3POp.V_SMFMAC_I32_16X16X64_I8, VOP3POp.V_SMFMAC_I32_32X32X32_I8, VOP3POp.V_MFMA_F64_16X16X4_F64, VOP3POp.V_MFMA_F64_4X4X4_4B_F64, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8})
vdst = VGPRField(7, 0)
src0 = SrcField(40, 32)
src1 = SrcField(49, 41)
@@ -174,7 +174,6 @@ class VOP3P(Inst):
clmp = BitField(15, 15)
opsel = BitField(13, 11)
opsel_hi = BitField(60, 59, default=3)
opsel_hi2 = BitField(14, 14, default=1)
class VOP3PX2(Inst):
encoding = FixedBitField(95, 87, 0b110100111)
@@ -310,7 +309,6 @@ class VOP2_SDWA_SDST(VOP2):
s1 = BitField(63, 63)
class VOP3P_MFMA(VOP3P):
op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_MFMA_F32_16X16X128_F8F6F4, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4, VOP3POp.V_MFMA_F32_16X16X32_BF16, VOP3POp.V_MFMA_I32_16X16X64_I8, VOP3POp.V_MFMA_F32_32X32X16_BF16, VOP3POp.V_MFMA_I32_32X32X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_BF16, VOP3POp.V_SMFMAC_I32_16X16X128_I8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X1_2B_F32, VOP3POp.V_MFMA_F32_16X16X1_4B_F32, VOP3POp.V_MFMA_F32_4X4X1_16B_F32, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X2_F32, VOP3POp.V_MFMA_F32_16X16X4_F32, VOP3POp.V_SMFMAC_F32_32X32X32_BF16, VOP3POp.V_SMFMAC_I32_32X32X64_I8, VOP3POp.V_MFMA_F32_32X32X4_2B_F16, VOP3POp.V_MFMA_F32_16X16X4_4B_F16, VOP3POp.V_MFMA_F32_4X4X4_16B_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X8_F16, VOP3POp.V_MFMA_F32_16X16X16_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8, VOP3POp.V_MFMA_I32_32X32X4_2B_I8, VOP3POp.V_MFMA_I32_16X16X4_4B_I8, VOP3POp.V_MFMA_I32_4X4X4_16B_I8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8, VOP3POp.V_MFMA_F32_16X16X32_F16, VOP3POp.V_MFMA_F32_32X32X16_F16, VOP3POp.V_MFMA_I32_32X32X16_I8, VOP3POp.V_MFMA_I32_16X16X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_F16, VOP3POp.V_SMFMAC_F32_32X32X32_F16, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16, VOP3POp.V_MFMA_F32_32X32X8_BF16, VOP3POp.V_MFMA_F32_16X16X16_BF16, VOP3POp.V_SMFMAC_F32_16X16X32_F16, VOP3POp.V_SMFMAC_F32_32X32X16_F16, VOP3POp.V_SMFMAC_F32_16X16X32_BF16, VOP3POp.V_SMFMAC_F32_32X32X16_BF16, VOP3POp.V_SMFMAC_I32_16X16X64_I8, VOP3POp.V_SMFMAC_I32_32X32X32_I8, VOP3POp.V_MFMA_F64_16X16X4_F64, VOP3POp.V_MFMA_F64_4X4X4_4B_F64, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8})
cbsz = BitField(10, 8)
abid = BitField(14, 11)
acc_cd = BitField(15, 15)
@@ -1649,80 +1647,80 @@ v_dot4_u32_u8 = functools.partial(VOP3P, VOP3POp.V_DOT4_U32_U8)
v_dot8_i32_i4 = functools.partial(VOP3P, VOP3POp.V_DOT8_I32_I4)
v_dot8_u32_u4 = functools.partial(VOP3P, VOP3POp.V_DOT8_U32_U4)
v_mfma_ld_scale_b32 = functools.partial(VOP3P, VOP3POp.V_MFMA_LD_SCALE_B32)
v_mfma_f32_16x16x128_f8f6f4 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4)
v_mfma_f32_32x32x64_f8f6f4 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4)
v_mfma_f32_16x16x128_f8f6f4 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4)
v_mfma_f32_32x32x64_f8f6f4 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4)
v_pk_fma_f32 = functools.partial(VOP3P, VOP3POp.V_PK_FMA_F32)
v_pk_mul_f32 = functools.partial(VOP3P, VOP3POp.V_PK_MUL_F32)
v_pk_add_f32 = functools.partial(VOP3P, VOP3POp.V_PK_ADD_F32)
v_pk_mov_b32 = functools.partial(VOP3P, VOP3POp.V_PK_MOV_B32)
v_mfma_f32_16x16x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF16)
v_mfma_i32_16x16x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X64_I8)
v_mfma_f32_32x32x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF16)
v_mfma_i32_32x32x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X32_I8)
v_smfmac_f32_16x16x64_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF16)
v_smfmac_i32_16x16x128_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_16X16X128_I8)
v_smfmac_f32_16x16x128_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8)
v_smfmac_f32_16x16x128_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8)
v_smfmac_f32_16x16x128_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8)
v_mfma_f32_16x16x32_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF16)
v_mfma_i32_16x16x64_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X64_I8)
v_mfma_f32_32x32x16_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF16)
v_mfma_i32_32x32x32_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X32_I8)
v_smfmac_f32_16x16x64_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF16)
v_smfmac_i32_16x16x128_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_16X16X128_I8)
v_smfmac_f32_16x16x128_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8)
v_smfmac_f32_16x16x128_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8)
v_smfmac_f32_16x16x128_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8)
v_mfma_f32_16x16x8_xf32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X8_XF32)
v_mfma_f32_32x32x4_xf32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_XF32)
v_mfma_f32_32x32x1_2b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X1_2B_F32)
v_mfma_f32_16x16x1_4b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X1_4B_F32)
v_mfma_f32_4x4x1_16b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X1_16B_F32)
v_smfmac_f32_16x16x128_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8)
v_mfma_f32_32x32x2_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X2_F32)
v_mfma_f32_16x16x4_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_F32)
v_smfmac_f32_32x32x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF16)
v_smfmac_i32_32x32x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_32X32X64_I8)
v_mfma_f32_32x32x4_2b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X4_2B_F16)
v_mfma_f32_16x16x4_4b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_4B_F16)
v_mfma_f32_4x4x4_16b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X4_16B_F16)
v_smfmac_f32_32x32x64_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8)
v_mfma_f32_32x32x8_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X8_F16)
v_mfma_f32_16x16x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X16_F16)
v_smfmac_f32_32x32x64_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8)
v_smfmac_f32_32x32x64_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8)
v_mfma_i32_32x32x4_2b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X4_2B_I8)
v_mfma_i32_16x16x4_4b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X4_4B_I8)
v_mfma_i32_4x4x4_16b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_4X4X4_16B_I8)
v_smfmac_f32_32x32x64_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8)
v_mfma_f32_16x16x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_F16)
v_mfma_f32_32x32x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_F16)
v_mfma_i32_32x32x16_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X16_I8)
v_mfma_i32_16x16x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X32_I8)
v_mfma_f32_32x32x1_2b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X1_2B_F32)
v_mfma_f32_16x16x1_4b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X1_4B_F32)
v_mfma_f32_4x4x1_16b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X1_16B_F32)
v_smfmac_f32_16x16x128_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8)
v_mfma_f32_32x32x2_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X2_F32)
v_mfma_f32_16x16x4_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_F32)
v_smfmac_f32_32x32x32_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF16)
v_smfmac_i32_32x32x64_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_32X32X64_I8)
v_mfma_f32_32x32x4_2b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_2B_F16)
v_mfma_f32_16x16x4_4b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_4B_F16)
v_mfma_f32_4x4x4_16b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X4_16B_F16)
v_smfmac_f32_32x32x64_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8)
v_mfma_f32_32x32x8_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X8_F16)
v_mfma_f32_16x16x16_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X16_F16)
v_smfmac_f32_32x32x64_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8)
v_smfmac_f32_32x32x64_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8)
v_mfma_i32_32x32x4_2b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X4_2B_I8)
v_mfma_i32_16x16x4_4b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X4_4B_I8)
v_mfma_i32_4x4x4_16b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_4X4X4_16B_I8)
v_smfmac_f32_32x32x64_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8)
v_mfma_f32_16x16x32_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_F16)
v_mfma_f32_32x32x16_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_F16)
v_mfma_i32_32x32x16_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X16_I8)
v_mfma_i32_16x16x32_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X32_I8)
v_accvgpr_read = functools.partial(VOP3P, VOP3POp.V_ACCVGPR_READ)
v_accvgpr_write = functools.partial(VOP3P, VOP3POp.V_ACCVGPR_WRITE)
v_smfmac_f32_16x16x64_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_F16)
v_smfmac_f32_32x32x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_F16)
v_mfma_f32_32x32x4_2b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16)
v_mfma_f32_16x16x4_4b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16)
v_mfma_f32_4x4x4_16b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16)
v_mfma_f32_32x32x8_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X8_BF16)
v_mfma_f32_16x16x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X16_BF16)
v_smfmac_f32_16x16x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X32_F16)
v_smfmac_f32_32x32x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X16_F16)
v_smfmac_f32_16x16x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X32_BF16)
v_smfmac_f32_32x32x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X16_BF16)
v_smfmac_i32_16x16x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_16X16X64_I8)
v_smfmac_i32_32x32x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_32X32X32_I8)
v_mfma_f64_16x16x4_f64 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F64_16X16X4_F64)
v_mfma_f64_4x4x4_4b_f64 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F64_4X4X4_4B_F64)
v_mfma_f32_16x16x32_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8)
v_mfma_f32_16x16x32_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8)
v_mfma_f32_16x16x32_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8)
v_mfma_f32_16x16x32_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8)
v_mfma_f32_32x32x16_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8)
v_mfma_f32_32x32x16_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8)
v_mfma_f32_32x32x16_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8)
v_mfma_f32_32x32x16_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8)
v_smfmac_f32_16x16x64_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8)
v_smfmac_f32_16x16x64_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8)
v_smfmac_f32_16x16x64_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8)
v_smfmac_f32_16x16x64_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8)
v_smfmac_f32_32x32x32_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8)
v_smfmac_f32_32x32x32_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8)
v_smfmac_f32_32x32x32_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8)
v_smfmac_f32_32x32x32_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8)
v_smfmac_f32_16x16x64_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_F16)
v_smfmac_f32_32x32x32_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_F16)
v_mfma_f32_32x32x4_2b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16)
v_mfma_f32_16x16x4_4b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16)
v_mfma_f32_4x4x4_16b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16)
v_mfma_f32_32x32x8_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X8_BF16)
v_mfma_f32_16x16x16_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X16_BF16)
v_smfmac_f32_16x16x32_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X32_F16)
v_smfmac_f32_32x32x16_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X16_F16)
v_smfmac_f32_16x16x32_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X32_BF16)
v_smfmac_f32_32x32x16_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X16_BF16)
v_smfmac_i32_16x16x64_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_16X16X64_I8)
v_smfmac_i32_32x32x32_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_32X32X32_I8)
v_mfma_f64_16x16x4_f64 = functools.partial(VOP3P, VOP3POp.V_MFMA_F64_16X16X4_F64)
v_mfma_f64_4x4x4_4b_f64 = functools.partial(VOP3P, VOP3POp.V_MFMA_F64_4X4X4_4B_F64)
v_mfma_f32_16x16x32_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8)
v_mfma_f32_16x16x32_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8)
v_mfma_f32_16x16x32_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8)
v_mfma_f32_16x16x32_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8)
v_mfma_f32_32x32x16_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8)
v_mfma_f32_32x32x16_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8)
v_mfma_f32_32x32x16_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8)
v_mfma_f32_32x32x16_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8)
v_smfmac_f32_16x16x64_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8)
v_smfmac_f32_16x16x64_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8)
v_smfmac_f32_16x16x64_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8)
v_smfmac_f32_16x16x64_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8)
v_smfmac_f32_32x32x32_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8)
v_smfmac_f32_32x32x32_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8)
v_smfmac_f32_32x32x32_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8)
v_smfmac_f32_32x32x32_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8)
v_mfma_scale_f32_16x16x128_f8f6f4 = functools.partial(VOP3PX2, VOP3PX2Op.V_MFMA_SCALE_F32_16X16X128_F8F6F4)
v_mfma_scale_f32_32x32x64_f8f6f4 = functools.partial(VOP3PX2, VOP3PX2Op.V_MFMA_SCALE_F32_32X32X64_F8F6F4)
v_add_co_u32 = functools.partial(VOP3SD, VOP3SDOp.V_ADD_CO_U32)
@@ -1,6 +1,6 @@
# autogenerated from AMD ISA XML - do not edit
from tinygrad.runtime.autogen.amd.common import Fmt, OpType
from tinygrad.runtime.autogen.amd.cdna.enum import DSOp, FLATOp, GLOBALOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3PX2Op, VOP3SDOp, VOPCOp
from extra.assembly.amd.autogen.common import Fmt, OpType
from extra.assembly.amd.autogen.cdna.enum import *
# instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}
OPERANDS = {
@@ -1,6 +1,6 @@
# autogenerated from AMD ISA PDF - do not edit
# ruff: noqa: E501
from tinygrad.runtime.autogen.amd.cdna.enum import DSOp, FLATOp, GLOBALOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp
from extra.assembly.amd.autogen.cdna.enum import DSOp, FLATOp, GLOBALOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp
PCODE = {
DSOp.DS_ADD_U32: 'addr = CalcDsAddr(ADDR.b32, OFFSET0.b32, OFFSET1.b32);\ntmp = MEM[addr].u32;\nMEM[addr].u32 += DATA.u32;\nRETURN_DATA.u32 = tmp',
@@ -1,10 +1,6 @@
# autogenerated from AMD ISA XML - do not edit
from enum import Enum, auto
class ReprEnum(Enum):
"""Enum with clean repr that roundtrips with eval()."""
def __repr__(self): return f"{type(self).__name__}.{self.name}"
class Fmt(Enum):
FMT_ANY = auto()
FMT_BUF = auto()
@@ -1,7 +1,8 @@
# autogenerated from AMD ISA XML - do not edit
from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401
from enum import Enum
from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401
class DSOp(ReprEnum):
class DSOp(Enum):
DS_ADD_U32 = 0
DS_SUB_U32 = 1
DS_RSUB_U32 = 2
@@ -129,10 +130,10 @@ class DSOp(ReprEnum):
DS_LOAD_B96 = 254
DS_LOAD_B128 = 255
class EXPOp(ReprEnum):
class EXPOp(Enum):
EXP = 0
class FLATOp(ReprEnum):
class FLATOp(Enum):
FLAT_LOAD_U8 = 16
FLAT_LOAD_I8 = 17
FLAT_LOAD_U16 = 18
@@ -190,7 +191,7 @@ class FLATOp(ReprEnum):
FLAT_ATOMIC_MAX_F32 = 82
FLAT_ATOMIC_ADD_F32 = 86
class GLOBALOp(ReprEnum):
class GLOBALOp(Enum):
GLOBAL_LOAD_U8 = 16
GLOBAL_LOAD_I8 = 17
GLOBAL_LOAD_U16 = 18
@@ -253,7 +254,7 @@ class GLOBALOp(ReprEnum):
GLOBAL_ATOMIC_MAX_F32 = 82
GLOBAL_ATOMIC_ADD_F32 = 86
class HWREG(ReprEnum):
class HWREG(Enum):
HW_REG_MODE = 1
HW_REG_STATUS = 2
HW_REG_TRAPSTS = 3
@@ -279,11 +280,11 @@ class HWREG(ReprEnum):
HW_REG_IB_STS2 = 28
HW_REG_SHADER_CYCLES = 29
class LDSDIROp(ReprEnum):
class LDSDIROp(Enum):
LDS_PARAM_LOAD = 0
LDS_DIRECT_LOAD = 1
class MIMGOp(ReprEnum):
class MIMGOp(Enum):
IMAGE_LOAD = 0
IMAGE_LOAD_MIP = 1
IMAGE_LOAD_PCK = 2
@@ -369,7 +370,7 @@ class MIMGOp(ReprEnum):
IMAGE_GATHER4_C_B_CL = 101
IMAGE_GATHER4H = 144
class MSG(ReprEnum):
class MSG(Enum):
MSG_RTN_GET_DOORBELL = 128
MSG_RTN_GET_DDID = 129
MSG_RTN_GET_TMA = 130
@@ -379,7 +380,7 @@ class MSG(ReprEnum):
MSG_RTN_GET_TBA_TO_PC = 134
MSG_RTN_ILLEGAL_MSG = 255
class MTBUFOp(ReprEnum):
class MTBUFOp(Enum):
TBUFFER_LOAD_FORMAT_X = 0
TBUFFER_LOAD_FORMAT_XY = 1
TBUFFER_LOAD_FORMAT_XYZ = 2
@@ -397,7 +398,7 @@ class MTBUFOp(ReprEnum):
TBUFFER_STORE_D16_FORMAT_XYZ = 14
TBUFFER_STORE_D16_FORMAT_XYZW = 15
class MUBUFOp(ReprEnum):
class MUBUFOp(Enum):
BUFFER_LOAD_FORMAT_X = 0
BUFFER_LOAD_FORMAT_XY = 1
BUFFER_LOAD_FORMAT_XYZ = 2
@@ -478,7 +479,7 @@ class MUBUFOp(ReprEnum):
BUFFER_ATOMIC_MAX_F32 = 82
BUFFER_ATOMIC_ADD_F32 = 86
class SCRATCHOp(ReprEnum):
class SCRATCHOp(Enum):
SCRATCH_LOAD_U8 = 16
SCRATCH_LOAD_I8 = 17
SCRATCH_LOAD_U16 = 18
@@ -507,7 +508,7 @@ class SCRATCHOp(ReprEnum):
SCRATCH_LOAD_LDS_I16 = 48
SCRATCH_LOAD_LDS_B32 = 49
class SMEMOp(ReprEnum):
class SMEMOp(Enum):
S_LOAD_B32 = 0
S_LOAD_B64 = 1
S_LOAD_B128 = 2
@@ -523,7 +524,7 @@ class SMEMOp(ReprEnum):
S_ATC_PROBE = 34
S_ATC_PROBE_BUFFER = 35
class SOP1Op(ReprEnum):
class SOP1Op(Enum):
S_MOV_B32 = 0
S_MOV_B64 = 1
S_CMOV_B32 = 2
@@ -605,7 +606,7 @@ class SOP1Op(ReprEnum):
S_TRUNC_F16 = 109
S_RNDNE_F16 = 110
class SOP2Op(ReprEnum):
class SOP2Op(Enum):
S_ADD_U32 = 0
S_SUB_U32 = 1
S_ADD_I32 = 2
@@ -674,7 +675,7 @@ class SOP2Op(ReprEnum):
S_MUL_F16 = 77
S_FMAC_F16 = 78
class SOPCOp(ReprEnum):
class SOPCOp(Enum):
S_CMP_EQ_I32 = 0
S_CMP_LG_I32 = 1
S_CMP_GT_I32 = 2
@@ -722,7 +723,7 @@ class SOPCOp(ReprEnum):
S_CMP_NEQ_F16 = 93
S_CMP_NLT_F16 = 94
class SOPKOp(ReprEnum):
class SOPKOp(Enum):
S_MOVK_I32 = 0
S_VERSION = 1
S_CMOVK_I32 = 2
@@ -751,7 +752,7 @@ class SOPKOp(ReprEnum):
S_WAITCNT_EXPCNT = 26
S_WAITCNT_LGKMCNT = 27
class SOPPOp(ReprEnum):
class SOPPOp(Enum):
S_NOP = 0
S_SETKILL = 1
S_SETHALT = 2
@@ -792,7 +793,7 @@ class SOPPOp(ReprEnum):
S_ICACHE_INV = 60
S_BARRIER = 61
class VINTERPOp(ReprEnum):
class VINTERPOp(Enum):
V_INTERP_P10_F32 = 0
V_INTERP_P2_F32 = 1
V_INTERP_P10_F16_F32 = 2
@@ -800,7 +801,7 @@ class VINTERPOp(ReprEnum):
V_INTERP_P10_RTZ_F16_F32 = 4
V_INTERP_P2_RTZ_F16_F32 = 5
class VOP1Op(ReprEnum):
class VOP1Op(Enum):
V_NOP_E32 = 0
V_MOV_B32_E32 = 1
V_READFIRSTLANE_B32_E32 = 2
@@ -974,7 +975,7 @@ class VOP1Op(ReprEnum):
V_CVT_I32_I16 = V_CVT_I32_I16_E32
V_CVT_U32_U16 = V_CVT_U32_U16_E32
class VOP2Op(ReprEnum):
class VOP2Op(Enum):
V_CNDMASK_B32_E32 = 1
V_DOT2ACC_F32_F16_E32 = 2
V_ADD_F32_E32 = 3
@@ -1068,7 +1069,7 @@ class VOP2Op(ReprEnum):
V_LDEXP_F16 = V_LDEXP_F16_E32
V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
class VOP3Op(ReprEnum):
class VOP3Op(Enum):
V_CMP_F_F16_E64 = 0
V_CMP_LT_F16_E64 = 1
V_CMP_EQ_F16_E64 = 2
@@ -1808,7 +1809,7 @@ class VOP3Op(ReprEnum):
V_CVT_I32_I16 = V_CVT_I32_I16_E64
V_CVT_U32_U16 = V_CVT_U32_U16_E64
class VOP3POp(ReprEnum):
class VOP3POp(Enum):
V_PK_MAD_I16 = 0
V_PK_MUL_LO_U16 = 1
V_PK_ADD_I16 = 2
@@ -1844,7 +1845,7 @@ class VOP3POp(ReprEnum):
V_WMMA_I32_16X16X16_IU8 = 68
V_WMMA_I32_16X16X16_IU4 = 69
class VOP3SDOp(ReprEnum):
class VOP3SDOp(Enum):
V_ADD_CO_CI_U32 = 288
V_SUB_CO_CI_U32 = 289
V_SUBREV_CO_CI_U32 = 290
@@ -1856,7 +1857,7 @@ class VOP3SDOp(ReprEnum):
V_SUB_CO_U32 = 769
V_SUBREV_CO_U32 = 770
class VOPCOp(ReprEnum):
class VOPCOp(Enum):
V_CMP_F_F16_E32 = 0
V_CMP_LT_F16_E32 = 1
V_CMP_EQ_F16_E32 = 2
@@ -2238,7 +2239,7 @@ class VOPCOp(ReprEnum):
V_CMPX_CLASS_F32 = V_CMPX_CLASS_F32_E32
V_CMPX_CLASS_F64 = V_CMPX_CLASS_F64_E32
class VOPDOp(ReprEnum):
class VOPDOp(Enum):
V_DUAL_FMAC_F32 = 0
V_DUAL_FMAAK_F32 = 1
V_DUAL_FMAMK_F32 = 2
@@ -1,7 +1,7 @@
# autogenerated from AMD ISA XML - do not edit
# ruff: noqa: E501,F401
from tinygrad.renderer.amd.dsl import BitField, DPP, DPP16, EXEC, EXECZ, EXEC_HI, EXEC_LO, EnumBitField, FixedBitField, INV_2PI, Inst, LIT, M0, NULL, OFF, SBaseField, SCC, SDWA, SGPRField, SRC_LDS_DIRECT, SRsrcField, SSrcField, SrcField, VCC, VCCZ, VCC_HI, VCC_LO, VDSTYField, VGPRField, s, src, ttmp, v
from tinygrad.runtime.autogen.amd.rdna3.enum import DSOp, EXPOp, FLATOp, GLOBALOp, LDSDIROp, MIMGOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, HWREG, MSG
# ruff: noqa: F401,F403
from extra.assembly.amd.dsl import *
from extra.assembly.amd.autogen.rdna3.enum import *
import functools
class DS(Inst):
@@ -593,6 +593,9 @@ flat_load_d16_hi_i8 = functools.partial(FLAT, FLATOp.FLAT_LOAD_D16_HI_I8)
flat_load_d16_hi_b16 = functools.partial(FLAT, FLATOp.FLAT_LOAD_D16_HI_B16)
flat_store_d16_hi_b8 = functools.partial(FLAT, FLATOp.FLAT_STORE_D16_HI_B8)
flat_store_d16_hi_b16 = functools.partial(FLAT, FLATOp.FLAT_STORE_D16_HI_B16)
global_load_addtid_b32 = functools.partial(FLAT, FLATOp.GLOBAL_LOAD_ADDTID_B32)
global_store_addtid_b32 = functools.partial(FLAT, FLATOp.GLOBAL_STORE_ADDTID_B32)
global_load_lds_addtid_b32 = functools.partial(FLAT, FLATOp.GLOBAL_LOAD_LDS_ADDTID_B32)
flat_atomic_swap_b32 = functools.partial(FLAT, FLATOp.FLAT_ATOMIC_SWAP_B32)
flat_atomic_cmpswap_b32 = functools.partial(FLAT, FLATOp.FLAT_ATOMIC_CMPSWAP_B32)
flat_atomic_add_u32 = functools.partial(FLAT, FLATOp.FLAT_ATOMIC_ADD_U32)
@@ -1,6 +1,6 @@
# autogenerated from AMD ISA XML - do not edit
from tinygrad.runtime.autogen.amd.common import Fmt, OpType
from tinygrad.runtime.autogen.amd.rdna3.enum import DSOp, EXPOp, FLATOp, GLOBALOp, LDSDIROp, MIMGOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp
from extra.assembly.amd.autogen.common import Fmt, OpType
from extra.assembly.amd.autogen.rdna3.enum import *
# instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}
OPERANDS = {
@@ -1,6 +1,6 @@
# autogenerated from AMD ISA PDF - do not edit
# ruff: noqa: E501
from tinygrad.runtime.autogen.amd.rdna3.enum import DSOp, FLATOp, GLOBALOp, LDSDIROp, MIMGOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp
from extra.assembly.amd.autogen.rdna3.enum import DSOp, FLATOp, GLOBALOp, LDSDIROp, MIMGOp, MTBUFOp, MUBUFOp, SCRATCHOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp
PCODE = {
DSOp.DS_ADD_U32: 'tmp = MEM[ADDR].u32;\nMEM[ADDR].u32 += DATA.u32;\nRETURN_DATA.u32 = tmp',
@@ -1,7 +1,8 @@
# autogenerated from AMD ISA XML - do not edit
from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401
from enum import Enum
from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401
class DSOp(ReprEnum):
class DSOp(Enum):
DS_ADD_U32 = 0
DS_SUB_U32 = 1
DS_RSUB_U32 = 2
@@ -126,7 +127,7 @@ class DSOp(ReprEnum):
DS_LOAD_B96 = 254
DS_LOAD_B128 = 255
class HWREG(ReprEnum):
class HWREG(Enum):
HW_REG_WAVE_MODE = 1
HW_REG_WAVE_STATUS = 2
HW_REG_WAVE_STATE_PRIV = 4
@@ -147,7 +148,7 @@ class HWREG(ReprEnum):
HW_REG_SHADER_CYCLES_LO = 29
HW_REG_SHADER_CYCLES_HI = 30
class MSG(ReprEnum):
class MSG(Enum):
MSG_RTN_GET_DOORBELL = 128
MSG_RTN_GET_DDID = 129
MSG_RTN_GET_TMA = 130
@@ -158,7 +159,7 @@ class MSG(ReprEnum):
MSG_RTN_GET_SE_HW_ID = 135
MSG_RTN_ILLEGAL_MSG = 255
class SMEMOp(ReprEnum):
class SMEMOp(Enum):
S_LOAD_B32 = 0
S_LOAD_B64 = 1
S_LOAD_B128 = 2
@@ -188,7 +189,7 @@ class SMEMOp(ReprEnum):
S_BUFFER_PREFETCH_DATA = 39
S_PREFETCH_DATA_PC_REL = 40
class SOP1Op(ReprEnum):
class SOP1Op(Enum):
S_MOV_B32 = 0
S_MOV_B64 = 1
S_CMOV_B32 = 2
@@ -277,7 +278,7 @@ class SOP1Op(ReprEnum):
S_TRUNC_F16 = 109
S_RNDNE_F16 = 110
class SOP2Op(ReprEnum):
class SOP2Op(Enum):
S_ADD_CO_U32 = 0
S_SUB_CO_U32 = 1
S_ADD_CO_I32 = 2
@@ -353,7 +354,7 @@ class SOP2Op(ReprEnum):
S_SUB_NC_U64 = 84
S_MUL_U64 = 85
class SOPCOp(ReprEnum):
class SOPCOp(Enum):
S_CMP_EQ_I32 = 0
S_CMP_LG_I32 = 1
S_CMP_GT_I32 = 2
@@ -401,7 +402,7 @@ class SOPCOp(ReprEnum):
S_CMP_NEQ_F16 = 93
S_CMP_NLT_F16 = 94
class SOPKOp(ReprEnum):
class SOPKOp(Enum):
S_MOVK_I32 = 0
S_VERSION = 1
S_CMOVK_I32 = 2
@@ -412,7 +413,7 @@ class SOPKOp(ReprEnum):
S_SETREG_IMM32_B32 = 19
S_CALL_B64 = 20
class SOPPOp(ReprEnum):
class SOPPOp(Enum):
S_NOP = 0
S_SETKILL = 1
S_SETHALT = 2
@@ -457,7 +458,7 @@ class SOPPOp(ReprEnum):
S_WAIT_LOADCNT_DSCNT = 72
S_WAIT_STORECNT_DSCNT = 73
class VBUFFEROp(ReprEnum):
class VBUFFEROp(Enum):
BUFFER_LOAD_FORMAT_X = 0
BUFFER_LOAD_FORMAT_XY = 1
BUFFER_LOAD_FORMAT_XYZ = 2
@@ -548,14 +549,14 @@ class VBUFFEROp(ReprEnum):
TBUFFER_STORE_D16_FORMAT_XYZ = 142
TBUFFER_STORE_D16_FORMAT_XYZW = 143
class VDSDIROp(ReprEnum):
class VDSDIROp(Enum):
DS_PARAM_LOAD = 0
DS_DIRECT_LOAD = 1
class VEXPORTOp(ReprEnum):
class VEXPORTOp(Enum):
EXPORT = 0
class VFLATOp(ReprEnum):
class VFLATOp(Enum):
FLAT_LOAD_U8 = 16
FLAT_LOAD_I8 = 17
FLAT_LOAD_U16 = 18
@@ -614,7 +615,7 @@ class VFLATOp(ReprEnum):
FLAT_ATOMIC_PK_ADD_F16 = 89
FLAT_ATOMIC_PK_ADD_BF16 = 90
class VGLOBALOp(ReprEnum):
class VGLOBALOp(Enum):
GLOBAL_LOAD_U8 = 16
GLOBAL_LOAD_I8 = 17
GLOBAL_LOAD_U16 = 18
@@ -681,7 +682,7 @@ class VGLOBALOp(ReprEnum):
GLOBAL_ATOMIC_PK_ADD_BF16 = 90
GLOBAL_ATOMIC_ORDERED_ADD_B64 = 115
class VIMAGEOp(ReprEnum):
class VIMAGEOp(Enum):
IMAGE_LOAD = 0
IMAGE_LOAD_MIP = 1
IMAGE_LOAD_PCK = 2
@@ -716,7 +717,7 @@ class VIMAGEOp(ReprEnum):
IMAGE_ATOMIC_PK_ADD_F16 = 134
IMAGE_ATOMIC_PK_ADD_BF16 = 135
class VINTERPOp(ReprEnum):
class VINTERPOp(Enum):
V_INTERP_P10_F32 = 0
V_INTERP_P2_F32 = 1
V_INTERP_P10_F16_F32 = 2
@@ -724,7 +725,7 @@ class VINTERPOp(ReprEnum):
V_INTERP_P10_RTZ_F16_F32 = 4
V_INTERP_P2_RTZ_F16_F32 = 5
class VOP1Op(ReprEnum):
class VOP1Op(Enum):
V_NOP_E32 = 0
V_MOV_B32_E32 = 1
V_READFIRSTLANE_B32_E32 = 2
@@ -906,7 +907,7 @@ class VOP1Op(ReprEnum):
V_CVT_PK_F32_FP8 = V_CVT_PK_F32_FP8_E32
V_CVT_PK_F32_BF8 = V_CVT_PK_F32_BF8_E32
class VOP2Op(ReprEnum):
class VOP2Op(Enum):
V_CNDMASK_B32_E32 = 1
V_ADD_F64_E32 = 2
V_ADD_F32_E32 = 3
@@ -1006,7 +1007,7 @@ class VOP2Op(ReprEnum):
V_LDEXP_F16 = V_LDEXP_F16_E32
V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
class VOP3Op(ReprEnum):
class VOP3Op(Enum):
V_CMP_LT_F16_E64 = 1
V_CMP_EQ_F16_E64 = 2
V_CMP_LE_F16_E64 = 3
@@ -1731,7 +1732,7 @@ class VOP3Op(ReprEnum):
V_CVT_PK_F32_FP8 = V_CVT_PK_F32_FP8_E64
V_CVT_PK_F32_BF8 = V_CVT_PK_F32_BF8_E64
class VOP3POp(ReprEnum):
class VOP3POp(Enum):
V_PK_MAD_I16 = 0
V_PK_MUL_LO_U16 = 1
V_PK_ADD_I16 = 2
@@ -1789,7 +1790,7 @@ class VOP3POp(ReprEnum):
V_SWMMAC_F32_16X16X32_BF8_FP8 = 89
V_SWMMAC_F32_16X16X32_BF8_BF8 = 90
class VOP3SDOp(ReprEnum):
class VOP3SDOp(Enum):
V_ADD_CO_CI_U32 = 288
V_SUB_CO_CI_U32 = 289
V_SUBREV_CO_CI_U32 = 290
@@ -1801,7 +1802,7 @@ class VOP3SDOp(ReprEnum):
V_SUB_CO_U32 = 769
V_SUBREV_CO_U32 = 770
class VOPCOp(ReprEnum):
class VOPCOp(Enum):
V_CMP_LT_F16_E32 = 1
V_CMP_EQ_F16_E32 = 2
V_CMP_LE_F16_E32 = 3
@@ -2127,7 +2128,7 @@ class VOPCOp(ReprEnum):
V_CMPX_CLASS_F32 = V_CMPX_CLASS_F32_E32
V_CMPX_CLASS_F64 = V_CMPX_CLASS_F64_E32
class VOPDOp(ReprEnum):
class VOPDOp(Enum):
V_DUAL_FMAC_F32 = 0
V_DUAL_FMAAK_F32 = 1
V_DUAL_FMAMK_F32 = 2
@@ -2146,7 +2147,7 @@ class VOPDOp(ReprEnum):
V_DUAL_LSHLREV_B32 = 17
V_DUAL_AND_B32 = 18
class VSAMPLEOp(ReprEnum):
class VSAMPLEOp(Enum):
IMAGE_MSAA_LOAD = 24
IMAGE_SAMPLE = 27
IMAGE_SAMPLE_D = 28
@@ -2206,7 +2207,7 @@ class VSAMPLEOp(ReprEnum):
IMAGE_GATHER4_C_B_CL = 101
IMAGE_GATHER4H = 144
class VSCRATCHOp(ReprEnum):
class VSCRATCHOp(Enum):
SCRATCH_LOAD_U8 = 16
SCRATCH_LOAD_I8 = 17
SCRATCH_LOAD_U16 = 18
@@ -1,7 +1,7 @@
# autogenerated from AMD ISA XML - do not edit
# ruff: noqa: E501,F401
from tinygrad.renderer.amd.dsl import BitField, DPP, DPP16, EXEC, EXECZ, EXEC_HI, EXEC_LO, EnumBitField, FixedBitField, INV_2PI, Inst, LIT, M0, NULL, OFF, SBaseField, SCC, SDWA, SGPRField, SRC_LDS_DIRECT, SSrcField, SrcField, VCC, VCCZ, VCC_HI, VCC_LO, VDSTYField, VGPRField, s, src, ttmp, v
from tinygrad.runtime.autogen.amd.rdna4.enum import DSOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VBUFFEROp, VDSDIROp, VEXPORTOp, VFLATOp, VGLOBALOp, VIMAGEOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, VSAMPLEOp, VSCRATCHOp, HWREG, MSG
# ruff: noqa: F401,F403
from extra.assembly.amd.dsl import *
from extra.assembly.amd.autogen.rdna4.enum import *
import functools
class DS(Inst):
@@ -101,11 +101,11 @@ class VFLAT(Inst):
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = VGPRField(62, 55)
vsrc = BitField(62, 55)
ioffset = BitField(95, 72)
class VGLOBAL(Inst):
encoding = FixedBitField(31, 24, 0b11101110)
encoding = FixedBitField(31, 24, 0b11101100)
op = EnumBitField(21, 14, VGLOBALOp, {VGLOBALOp.GLOBAL_LOAD_U8, VGLOBALOp.GLOBAL_LOAD_I8, VGLOBALOp.GLOBAL_LOAD_U16, VGLOBALOp.GLOBAL_LOAD_I16, VGLOBALOp.GLOBAL_LOAD_B32, VGLOBALOp.GLOBAL_LOAD_B64, VGLOBALOp.GLOBAL_LOAD_B96, VGLOBALOp.GLOBAL_LOAD_B128, VGLOBALOp.GLOBAL_STORE_B8, VGLOBALOp.GLOBAL_STORE_B16, VGLOBALOp.GLOBAL_STORE_B32, VGLOBALOp.GLOBAL_STORE_B64, VGLOBALOp.GLOBAL_STORE_B96, VGLOBALOp.GLOBAL_STORE_B128, VGLOBALOp.GLOBAL_LOAD_D16_U8, VGLOBALOp.GLOBAL_LOAD_D16_I8, VGLOBALOp.GLOBAL_LOAD_D16_B16, VGLOBALOp.GLOBAL_LOAD_D16_HI_U8, VGLOBALOp.GLOBAL_LOAD_D16_HI_I8, VGLOBALOp.GLOBAL_LOAD_D16_HI_B16, VGLOBALOp.GLOBAL_STORE_D16_HI_B8, VGLOBALOp.GLOBAL_STORE_D16_HI_B16, VGLOBALOp.GLOBAL_LOAD_ADDTID_B32, VGLOBALOp.GLOBAL_STORE_ADDTID_B32, VGLOBALOp.GLOBAL_INV, VGLOBALOp.GLOBAL_WB, VGLOBALOp.GLOBAL_ATOMIC_SWAP_B32, VGLOBALOp.GLOBAL_ATOMIC_CMPSWAP_B32, VGLOBALOp.GLOBAL_ATOMIC_ADD_U32, VGLOBALOp.GLOBAL_ATOMIC_SUB_U32, VGLOBALOp.GLOBAL_ATOMIC_SUB_CLAMP_U32, VGLOBALOp.GLOBAL_ATOMIC_MIN_I32, VGLOBALOp.GLOBAL_ATOMIC_MIN_U32, VGLOBALOp.GLOBAL_ATOMIC_MAX_I32, VGLOBALOp.GLOBAL_ATOMIC_MAX_U32, VGLOBALOp.GLOBAL_ATOMIC_AND_B32, VGLOBALOp.GLOBAL_ATOMIC_OR_B32, VGLOBALOp.GLOBAL_ATOMIC_XOR_B32, VGLOBALOp.GLOBAL_ATOMIC_INC_U32, VGLOBALOp.GLOBAL_ATOMIC_DEC_U32, VGLOBALOp.GLOBAL_ATOMIC_SWAP_B64, VGLOBALOp.GLOBAL_ATOMIC_CMPSWAP_B64, VGLOBALOp.GLOBAL_ATOMIC_ADD_U64, VGLOBALOp.GLOBAL_ATOMIC_SUB_U64, VGLOBALOp.GLOBAL_ATOMIC_MIN_I64, VGLOBALOp.GLOBAL_ATOMIC_MIN_U64, VGLOBALOp.GLOBAL_ATOMIC_MAX_I64, VGLOBALOp.GLOBAL_ATOMIC_MAX_U64, VGLOBALOp.GLOBAL_ATOMIC_AND_B64, VGLOBALOp.GLOBAL_ATOMIC_OR_B64, VGLOBALOp.GLOBAL_ATOMIC_XOR_B64, VGLOBALOp.GLOBAL_ATOMIC_INC_U64, VGLOBALOp.GLOBAL_ATOMIC_DEC_U64, VGLOBALOp.GLOBAL_WBINV, VGLOBALOp.GLOBAL_ATOMIC_COND_SUB_U32, VGLOBALOp.GLOBAL_ATOMIC_MIN_NUM_F32, VGLOBALOp.GLOBAL_ATOMIC_MAX_NUM_F32, VGLOBALOp.GLOBAL_LOAD_BLOCK, VGLOBALOp.GLOBAL_STORE_BLOCK, VGLOBALOp.GLOBAL_ATOMIC_ADD_F32, VGLOBALOp.GLOBAL_LOAD_TR_B128, VGLOBALOp.GLOBAL_LOAD_TR_B64, VGLOBALOp.GLOBAL_ATOMIC_PK_ADD_F16, VGLOBALOp.GLOBAL_ATOMIC_PK_ADD_BF16, VGLOBALOp.GLOBAL_ATOMIC_ORDERED_ADD_B64})
vdst = VGPRField(39, 32)
vaddr = VGPRField(71, 64)
@@ -114,7 +114,20 @@ class VGLOBAL(Inst):
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = VGPRField(62, 55)
vsrc = BitField(62, 55)
ioffset = BitField(95, 72)
class VSCRATCH(Inst):
encoding = FixedBitField(31, 24, 0b11101100)
op = EnumBitField(21, 14, VSCRATCHOp, {VSCRATCHOp.SCRATCH_LOAD_U8, VSCRATCHOp.SCRATCH_LOAD_I8, VSCRATCHOp.SCRATCH_LOAD_U16, VSCRATCHOp.SCRATCH_LOAD_I16, VSCRATCHOp.SCRATCH_LOAD_B32, VSCRATCHOp.SCRATCH_LOAD_B64, VSCRATCHOp.SCRATCH_LOAD_B96, VSCRATCHOp.SCRATCH_LOAD_B128, VSCRATCHOp.SCRATCH_STORE_B8, VSCRATCHOp.SCRATCH_STORE_B16, VSCRATCHOp.SCRATCH_STORE_B32, VSCRATCHOp.SCRATCH_STORE_B64, VSCRATCHOp.SCRATCH_STORE_B96, VSCRATCHOp.SCRATCH_STORE_B128, VSCRATCHOp.SCRATCH_LOAD_D16_U8, VSCRATCHOp.SCRATCH_LOAD_D16_I8, VSCRATCHOp.SCRATCH_LOAD_D16_B16, VSCRATCHOp.SCRATCH_LOAD_D16_HI_U8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_I8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_B16, VSCRATCHOp.SCRATCH_STORE_D16_HI_B8, VSCRATCHOp.SCRATCH_STORE_D16_HI_B16, VSCRATCHOp.SCRATCH_LOAD_BLOCK, VSCRATCHOp.SCRATCH_STORE_BLOCK})
vdst = VGPRField(39, 32)
vaddr = VGPRField(71, 64)
saddr = SGPRField(6, 0, default=NULL)
nv = BitField(7, 7)
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = BitField(62, 55)
ioffset = BitField(95, 72)
class VIMAGE(Inst):
@@ -240,19 +253,6 @@ class VSAMPLE(Inst):
vaddr2 = BitField(87, 80)
vaddr3 = BitField(95, 88)
class VSCRATCH(Inst):
encoding = FixedBitField(31, 24, 0b11101101)
op = EnumBitField(21, 14, VSCRATCHOp, {VSCRATCHOp.SCRATCH_LOAD_U8, VSCRATCHOp.SCRATCH_LOAD_I8, VSCRATCHOp.SCRATCH_LOAD_U16, VSCRATCHOp.SCRATCH_LOAD_I16, VSCRATCHOp.SCRATCH_LOAD_B32, VSCRATCHOp.SCRATCH_LOAD_B64, VSCRATCHOp.SCRATCH_LOAD_B96, VSCRATCHOp.SCRATCH_LOAD_B128, VSCRATCHOp.SCRATCH_STORE_B8, VSCRATCHOp.SCRATCH_STORE_B16, VSCRATCHOp.SCRATCH_STORE_B32, VSCRATCHOp.SCRATCH_STORE_B64, VSCRATCHOp.SCRATCH_STORE_B96, VSCRATCHOp.SCRATCH_STORE_B128, VSCRATCHOp.SCRATCH_LOAD_D16_U8, VSCRATCHOp.SCRATCH_LOAD_D16_I8, VSCRATCHOp.SCRATCH_LOAD_D16_B16, VSCRATCHOp.SCRATCH_LOAD_D16_HI_U8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_I8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_B16, VSCRATCHOp.SCRATCH_STORE_D16_HI_B8, VSCRATCHOp.SCRATCH_STORE_D16_HI_B16, VSCRATCHOp.SCRATCH_LOAD_BLOCK, VSCRATCHOp.SCRATCH_STORE_BLOCK})
vdst = VGPRField(39, 32)
vaddr = VGPRField(71, 64)
saddr = SGPRField(6, 0, default=NULL)
nv = BitField(7, 7)
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = VGPRField(62, 55)
ioffset = BitField(95, 72)
class SOP1_LIT(SOP1):
op = EnumBitField(15, 8, SOP1Op, {SOP1Op.S_MOV_B32, SOP1Op.S_MOV_B64, SOP1Op.S_CMOV_B32, SOP1Op.S_CMOV_B64, SOP1Op.S_BREV_B32, SOP1Op.S_BREV_B64, SOP1Op.S_CTZ_I32_B32, SOP1Op.S_CTZ_I32_B64, SOP1Op.S_CLZ_I32_U32, SOP1Op.S_CLZ_I32_U64, SOP1Op.S_CLS_I32, SOP1Op.S_CLS_I32_I64, SOP1Op.S_SEXT_I32_I8, SOP1Op.S_SEXT_I32_I16, SOP1Op.S_BITSET0_B32, SOP1Op.S_BITSET0_B64, SOP1Op.S_BITSET1_B32, SOP1Op.S_BITSET1_B64, SOP1Op.S_BITREPLICATE_B64_B32, SOP1Op.S_ABS_I32, SOP1Op.S_BCNT0_I32_B32, SOP1Op.S_BCNT0_I32_B64, SOP1Op.S_BCNT1_I32_B32, SOP1Op.S_BCNT1_I32_B64, SOP1Op.S_QUADMASK_B32, SOP1Op.S_QUADMASK_B64, SOP1Op.S_WQM_B32, SOP1Op.S_WQM_B64, SOP1Op.S_NOT_B32, SOP1Op.S_NOT_B64, SOP1Op.S_AND_SAVEEXEC_B32, SOP1Op.S_AND_SAVEEXEC_B64, SOP1Op.S_OR_SAVEEXEC_B32, SOP1Op.S_OR_SAVEEXEC_B64, SOP1Op.S_XOR_SAVEEXEC_B32, SOP1Op.S_XOR_SAVEEXEC_B64, SOP1Op.S_NAND_SAVEEXEC_B32, SOP1Op.S_NAND_SAVEEXEC_B64, SOP1Op.S_NOR_SAVEEXEC_B32, SOP1Op.S_NOR_SAVEEXEC_B64, SOP1Op.S_XNOR_SAVEEXEC_B32, SOP1Op.S_XNOR_SAVEEXEC_B64, SOP1Op.S_AND_NOT0_SAVEEXEC_B32, SOP1Op.S_AND_NOT0_SAVEEXEC_B64, SOP1Op.S_OR_NOT0_SAVEEXEC_B32, SOP1Op.S_OR_NOT0_SAVEEXEC_B64, SOP1Op.S_AND_NOT1_SAVEEXEC_B32, SOP1Op.S_AND_NOT1_SAVEEXEC_B64, SOP1Op.S_OR_NOT1_SAVEEXEC_B32, SOP1Op.S_OR_NOT1_SAVEEXEC_B64, SOP1Op.S_AND_NOT0_WREXEC_B32, SOP1Op.S_AND_NOT0_WREXEC_B64, SOP1Op.S_AND_NOT1_WREXEC_B32, SOP1Op.S_AND_NOT1_WREXEC_B64, SOP1Op.S_MOVRELS_B32, SOP1Op.S_MOVRELS_B64, SOP1Op.S_MOVRELD_B32, SOP1Op.S_MOVRELD_B64, SOP1Op.S_MOVRELSD_2_B32, SOP1Op.S_GETPC_B64, SOP1Op.S_SETPC_B64, SOP1Op.S_SWAPPC_B64, SOP1Op.S_RFE_B64, SOP1Op.S_SENDMSG_RTN_B32, SOP1Op.S_SENDMSG_RTN_B64, SOP1Op.S_BARRIER_SIGNAL, SOP1Op.S_BARRIER_SIGNAL_ISFIRST, SOP1Op.S_GET_BARRIER_STATE, SOP1Op.S_BARRIER_INIT, SOP1Op.S_BARRIER_JOIN, SOP1Op.S_ALLOC_VGPR, SOP1Op.S_SLEEP_VAR, SOP1Op.S_CEIL_F32, SOP1Op.S_FLOOR_F32, SOP1Op.S_TRUNC_F32, SOP1Op.S_RNDNE_F32, SOP1Op.S_CVT_F32_I32, SOP1Op.S_CVT_F32_U32, SOP1Op.S_CVT_I32_F32, SOP1Op.S_CVT_U32_F32, SOP1Op.S_CVT_F16_F32, SOP1Op.S_CVT_F32_F16, SOP1Op.S_CVT_HI_F32_F16, SOP1Op.S_CEIL_F16, SOP1Op.S_FLOOR_F16, SOP1Op.S_TRUNC_F16, SOP1Op.S_RNDNE_F16})
literal = BitField(63, 32)
@@ -973,6 +973,8 @@ flat_load_d16_hi_i8 = functools.partial(VFLAT, VFLATOp.FLAT_LOAD_D16_HI_I8)
flat_load_d16_hi_b16 = functools.partial(VFLAT, VFLATOp.FLAT_LOAD_D16_HI_B16)
flat_store_d16_hi_b8 = functools.partial(VFLAT, VFLATOp.FLAT_STORE_D16_HI_B8)
flat_store_d16_hi_b16 = functools.partial(VFLAT, VFLATOp.FLAT_STORE_D16_HI_B16)
global_load_addtid_b32 = functools.partial(VFLAT, VFLATOp.GLOBAL_LOAD_ADDTID_B32)
global_store_addtid_b32 = functools.partial(VFLAT, VFLATOp.GLOBAL_STORE_ADDTID_B32)
flat_atomic_swap_b32 = functools.partial(VFLAT, VFLATOp.FLAT_ATOMIC_SWAP_B32)
flat_atomic_cmpswap_b32 = functools.partial(VFLAT, VFLATOp.FLAT_ATOMIC_CMPSWAP_B32)
flat_atomic_add_u32 = functools.partial(VFLAT, VFLATOp.FLAT_ATOMIC_ADD_U32)
@@ -1,6 +1,6 @@
# autogenerated from AMD ISA XML - do not edit
from tinygrad.runtime.autogen.amd.common import Fmt, OpType
from tinygrad.runtime.autogen.amd.rdna4.enum import DSOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VBUFFEROp, VDSDIROp, VEXPORTOp, VFLATOp, VGLOBALOp, VIMAGEOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, VSAMPLEOp, VSCRATCHOp
from extra.assembly.amd.autogen.common import Fmt, OpType
from extra.assembly.amd.autogen.rdna4.enum import *
# instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}
OPERANDS = {
@@ -1,6 +1,6 @@
# autogenerated from AMD ISA PDF - do not edit
# ruff: noqa: E501
from tinygrad.runtime.autogen.amd.rdna4.enum import DSOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VBUFFEROp, VFLATOp, VGLOBALOp, VIMAGEOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, VSAMPLEOp, VSCRATCHOp
from extra.assembly.amd.autogen.rdna4.enum import DSOp, SMEMOp, SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, VBUFFEROp, VFLATOp, VGLOBALOp, VIMAGEOp, VINTERPOp, VOP1Op, VOP2Op, VOP3Op, VOP3POp, VOP3SDOp, VOPCOp, VOPDOp, VSAMPLEOp, VSCRATCHOp
PCODE = {
DSOp.DS_ADD_U32: 'addr = CalcDsAddr(vgpr_a.b32, offset.b32);\ntmp = MEM[addr].u32;\nMEM[addr].u32 += DATA.u32;\nRETURN_DATA.u32 = tmp',
+65
View File
@@ -0,0 +1,65 @@
# Instruction format detection and decoding
from __future__ import annotations
from extra.assembly.amd.dsl import Inst, FixedBitField, EnumBitField
# SDWA/DPP variant detection: src0 field (bits 0-8) encodes the variant
# 0xf9 (249) = SDWA, 0xfa (250) = DPP16 for CDNA (GFX9)
_VARIANT_SRC0 = {"_SDWA_SDST": 0xf9, "_SDWA": 0xf9, "_DPP16": 0xfa}
def _matches(data: bytes, cls: type[Inst]) -> bool:
"""Check if data matches all FixedBitFields and op is in allowed."""
for _, field in cls._fields:
dword_idx = field.lo // 32
if len(data) < (dword_idx + 1) * 4: return False
word = int.from_bytes(data[dword_idx*4:(dword_idx+1)*4], 'little')
field_lo = field.lo % 32
if isinstance(field, FixedBitField):
if ((word >> field_lo) & field.mask) != field.default: return False
if isinstance(field, EnumBitField) and field.allowed is not None:
try: opcode = field.decode((word >> field_lo) & field.mask)
except ValueError: return False # opcode not in enum
if opcode not in field.allowed: return False
# Check SDWA/DPP variant based on src0 field (bits 0-8) - only for variant classes
name = cls.__name__
word = int.from_bytes(data[:4], 'little')
for suffix, expected_src0 in _VARIANT_SRC0.items():
if name.endswith(suffix): return (word & 0x1ff) == expected_src0
return True
# Import instruction classes for each architecture
from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3SD, VOP3P, VOPC, VOPD, VINTERP,
SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH)
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT,
VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
SOPC as R4_SOPC, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP,
SMEM as R4_SMEM, DS as R4_DS, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
from extra.assembly.amd.autogen.cdna.ins import (VOP1 as C_VOP1, VOP1_SDWA as C_VOP1_SDWA, VOP1_DPP16 as C_VOP1_DPP16,
VOP2 as C_VOP2, VOP2_LIT as C_VOP2_LIT, VOP2_SDWA as C_VOP2_SDWA, VOP2_DPP16 as C_VOP2_DPP16,
VOPC as C_VOPC, VOPC_SDWA_SDST as C_VOPC_SDWA_SDST,
VOP3 as C_VOP3, VOP3_SDST as C_VOP3_SDST, VOP3SD as C_VOP3SD, VOP3P as C_VOP3P, VOP3PX2 as C_VOP3PX2,
SOP1 as C_SOP1, SOP2 as C_SOP2, SOPC as C_SOPC, SOPK as C_SOPK, SOPK_LIT as C_SOPK_LIT, SOPP as C_SOPP, SMEM as C_SMEM, DS as C_DS,
FLAT as C_FLAT, GLOBAL as C_GLOBAL, SCRATCH as C_SCRATCH, MUBUF as C_MUBUF)
# Order matters: more specific encodings first, catch-alls (SOP2, VOP2) last
# Order: base before _LIT (base matches regular ops, _LIT catches lit-only ops excluded from base)
_FORMATS = {
"rdna3": [VOPD, VOP3P, VINTERP, VOP3SD, VOP3_SDST, VOP3, DS, GLOBAL, SCRATCH, FLAT, SMEM,
SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, VOPC, VOP1_SDST, VOP1, VOP1_LIT, VOP2, VOP2_LIT],
"rdna4": [R4_VOPD, R4_VOP3P, R4_VINTERP, R4_VOP3SD, R4_VOP3_SDST, R4_VOP3, R4_DS, R4_GLOBAL, R4_SCRATCH, R4_FLAT, R4_SMEM,
R4_SOP1, R4_SOPC, R4_SOPP, R4_SOPK, R4_SOPK_LIT, R4_VOPC, R4_VOP1_SDST, R4_VOP1, R4_SOP2, R4_SOP2_LIT, R4_VOP2, R4_VOP2_LIT],
"cdna": [C_VOP3PX2, C_VOP3P, C_VOP3SD, C_VOP3_SDST, C_VOP3, C_DS, C_GLOBAL, C_SCRATCH, C_FLAT, C_MUBUF, C_SMEM,
C_SOP1, C_SOPC, C_SOPP, C_SOPK, C_SOPK_LIT, C_VOPC_SDWA_SDST, C_VOPC,
C_VOP1_DPP16, C_VOP1_SDWA, C_VOP1, C_VOP2_DPP16, C_VOP2_SDWA, C_SOP2, C_VOP2, C_VOP2_LIT],
}
def detect_format(data: bytes, arch: str = "rdna3") -> type[Inst]:
"""Detect instruction format from machine code bytes."""
assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
for cls in _FORMATS[arch]:
if _matches(data, cls): return cls
raise ValueError(f"unknown {arch} format word={int.from_bytes(data[:4], 'little'):#010x}")
def decode_inst(data: bytes, arch: str = "rdna3") -> Inst:
"""Decode machine code bytes into an instruction."""
return detect_format(data, arch).from_bytes(data)
@@ -1,16 +1,13 @@
# RDNA3/RDNA4/CDNA disassembler
from __future__ import annotations
import re
from typing import Callable
from tinygrad.renderer.amd.dsl import Inst, Reg
import re, struct
from extra.assembly.amd.dsl import Inst, Reg
# Special register mappings for disassembly
SPECIAL_GPRS = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'null', 125: 'm0', 126: 'exec_lo', 127: 'exec_hi',
128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0',
246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc'}
128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc'}
SPECIAL_GPRS_CDNA = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'm0', 126: 'exec_lo', 127: 'exec_hi',
128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0',
246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc',
128: '0', 240: '0.5', 241: '-0.5', 242: '1.0', 243: '-1.0', 244: '2.0', 245: '-2.0', 246: '4.0', 247: '-4.0', 248: '0x3e22f983', 253: 'scc',
102: 'flat_scratch_lo', 103: 'flat_scratch_hi', 104: 'xnack_mask_lo', 105: 'xnack_mask_hi',
251: 'src_vccz', 252: 'src_execz'}
SPECIAL_PAIRS = {106: 'vcc', 126: 'exec'}
@@ -72,29 +69,25 @@ def _num_srcs(inst) -> int:
if any(x in n for x in ('FMA', 'MAD', 'CNDMASK', 'BFE', 'BFI', 'LERP', 'MED3', 'SAD', 'DIV_FMAS', 'DIV_FIXUP', 'DIV_SCALE', 'CUBE')): return 3
# PERMLANE_VAR ops are 2-source, but PERMLANE (non-VAR) are 3-source
if 'PERMLANE' in n and '_VAR' not in n: return 3
if any(x in n for x in ('_ADD3', '_LSHL_ADD', '_ADD_LSHL', '_LSHL_OR', '_AND_OR', 'OR3_B32', 'AND_OR_B32', 'ALIGNBIT',
'ALIGNBYTE', 'V_PERM_', 'XOR3', 'XAD', 'MULLIT', 'MINMAX', 'MAXMIN', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM',
'MINIMUM3', 'MAXIMUM3', 'MIN3', 'MAX3', 'DOT2', 'CVT_PK_U8_F32', 'DOT4', 'DOT8', 'WMMA', 'SWMMAC')): return 3
if any(x in n for x in ('_ADD3', '_LSHL_ADD', '_ADD_LSHL', '_LSHL_OR', '_AND_OR', 'OR3_B32', 'AND_OR_B32', 'ALIGNBIT', 'ALIGNBYTE', 'V_PERM_', 'XOR3', 'XAD', 'MULLIT', 'MINMAX', 'MAXMIN', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM', 'MINIMUM3', 'MAXIMUM3', 'MIN3', 'MAX3', 'DOT2', 'CVT_PK_U8_F32', 'DOT4', 'DOT8', 'WMMA', 'SWMMAC')): return 3
return 2
# ═══════════════════════════════════════════════════════════════════════════════
# IMPORTS
# ═══════════════════════════════════════════════════════════════════════════════
from tinygrad.runtime.autogen.amd.rdna3.ins import (VOP1, VOP1_SDST, VOP1_SDST_LIT, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3_SDST_LIT,
from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP1_SDST_LIT, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3_SDST_LIT,
VOP3_LIT, VOP3SD, VOP3SD_LIT, VOP3P, VOP3P_LIT, VOPC, VOPC_LIT, VOPD, VOPD_LIT, VINTERP, SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPC_LIT,
SOPK, SOPK_LIT, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH, VOP2Op, VOPDOp, SOPPOp, HWREG, MSG)
from tinygrad.runtime.autogen.amd.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST,
VOP1_SDST_LIT as R4_VOP1_SDST_LIT, VOP1_LIT as R4_VOP1_LIT,
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP1_SDST_LIT as R4_VOP1_SDST_LIT, VOP1_LIT as R4_VOP1_LIT,
VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT, VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3_SDST_LIT as R4_VOP3_SDST_LIT, VOP3_LIT as R4_VOP3_LIT,
VOP3SD as R4_VOP3SD, VOP3SD_LIT as R4_VOP3SD_LIT, VOP3P as R4_VOP3P, VOP3P_LIT as R4_VOP3P_LIT, VOPC as R4_VOPC, VOPC_LIT as R4_VOPC_LIT,
VOPD as R4_VOPD, VOPD_LIT as R4_VOPD_LIT, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP1_LIT as R4_SOP1_LIT, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
SOPC as R4_SOPC, SOPC_LIT as R4_SOPC_LIT, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP, SMEM as R4_SMEM, DS as R4_DS,
VOPDOp as R4_VOPDOp, HWREG as HWREG_RDNA4, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
from tinygrad.runtime.autogen.amd.cdna.ins import HWREG as HWREG_CDNA
VOPDOp as R4_VOPDOp, HWREG as HWREG_RDNA4)
from extra.assembly.amd.autogen.cdna.ins import FLAT as C_FLAT, HWREG as HWREG_CDNA
def _is_cdna(inst: Inst) -> bool: return 'cdna' in inst.__class__.__module__
def _is_r4(inst: Inst) -> bool: return 'rdna4' in inst.__class__.__module__
# CDNA opcode name aliases for disasm (new name -> old name expected by tests)
_CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32', 'v_fmamk_f32': 'v_madmk_f32', 'v_fmaak_f32': 'v_madak_f32'}
@@ -105,15 +98,9 @@ _CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v
def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{_unwrap(b)}" if n == 1 else f"{p}[{_unwrap(b)}:{_unwrap(b)+n-1}]"
def _sreg(b: int, n: int = 1) -> str: return _reg("s", _unwrap(b), n)
def _vreg(b: int, n: int = 1) -> str:
b = _unwrap(b)
return _reg("v", b - 256 if b >= 256 else b, n)
def _areg(b: int, n: int = 1) -> str:
b = _unwrap(b)
return _reg("a", b - 256 if b >= 256 else b, n) # accumulator registers for GFX90a
def _ttmp(b, n: int = 1) -> str | None:
b = _unwrap(b)
return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
def _vreg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("v", b - 256 if b >= 256 else b, n)
def _areg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("a", b - 256 if b >= 256 else b, n) # accumulator registers for GFX90a
def _ttmp(b, n: int = 1) -> str: b = _unwrap(b); return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
def _fmt_sdst(v, n: int = 1, cdna: bool = False) -> str:
v = _unwrap(v)
@@ -141,9 +128,7 @@ def _fmt_v16(v, base: int = 256, hi_thresh: int = 384) -> str:
def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
def _omod(v: int) -> str: return {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(v, "")
def _src16(inst, v: int) -> str:
v = _unwrap(v)
return _fmt_v16(v) if v >= 256 else _lit(inst, v) # format 16-bit src: vgpr.h/l or literal
def _src16(inst, v: int) -> str: v = _unwrap(v); return _fmt_v16(v) if v >= 256 else _lit(inst, v) # format 16-bit src: vgpr.h/l or literal
def _mods(*pairs) -> str: return " ".join(m for c, m in pairs if c)
def _fmt_bits(label: str, val: int, count: int) -> str: return f"{label}:[{','.join(str((val >> i) & 1) for i in range(count))}]"
@@ -214,8 +199,7 @@ def _disasm_vop2(inst: VOP2) -> str:
basename = name.replace('_e32', '')
if cdna and basename in _VOP2_CARRY_OUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}"
if cdna and basename in _VOP2_CARRY_INOUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
if not cdna and basename in _VOP2_CARRY_INOUT_RDNA:
return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
if not cdna and basename in _VOP2_CARRY_INOUT_RDNA: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
sn0 = inst.canonical_op_regs.get('s0', 1)
if inst.vdst.sz > 1 or sn0 > 1 or inst.vsrc1.sz > 1:
src0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else _fmt_src(inst.src0, sn0, cdna)
@@ -231,10 +215,7 @@ def _disasm_vopc(inst: VOPC) -> str:
return f"{name} vcc, {s0}, {inst.vsrc1.fmt()}" # CDNA VOPC always outputs vcc
# RDNA: v_cmpx_* writes to exec (no vcc), v_cmp_* writes to vcc_lo
has_vcc = 'cmpx' not in name
if inst.src0.offset == 255: s0 = _lit(inst, inst.src0)
elif inst.src0.sz > 1: s0 = inst.src0.fmt()
elif is16: s0 = _src16(inst, inst.src0.offset)
else: s0 = _lit(inst, inst.src0)
s0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else inst.src0.fmt() if inst.src0.sz > 1 else _src16(inst, inst.src0.offset) if is16 else _lit(inst, inst.src0)
s1 = inst.vsrc1.fmt() if inst.vsrc1.sz > 1 else _fmt_v16(inst.vsrc1) if is16 else inst.vsrc1.fmt()
suf = "" if name.endswith('_e32') else "_e32"
return f"{name}{suf} vcc_lo, {s0}, {s1}" if has_vcc else f"{name}{suf} {s0}, {s1}"
@@ -244,7 +225,7 @@ NO_ARG_SOPP = {SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
def _disasm_sopp(inst: SOPP) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
is_rdna4 = _is_r4(inst)
is_rdna4 = 'rdna4' in inst.__class__.__module__
# Ops that have no argument when simm16 == 0
no_arg_zero = {'s_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_wait_idle', 's_endpgm_saved',
's_endpgm_ordered_ps_done', 's_code_end'}
@@ -270,22 +251,21 @@ def _disasm_sopp(inst: SOPP) -> str:
p = [f"vmcnt({vm})" if vm != 0x3f else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
if name == 's_delay_alu':
deps = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2',
'TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
deps = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
skips = ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
id0, skip, id1 = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 7) & 0xf
def dep(v): return deps[v-1] if 0 < v <= len(deps) else str(v)
dep = lambda v: deps[v-1] if 0 < v <= len(deps) else str(v)
p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
if name.startswith(('s_cbranch', 's_branch')): return f"{name} {inst.simm16}"
if name.startswith(('s_cbranch', 's_branch')): return f"{name} 0x{inst.simm16:x}"
return f"{name} 0x{inst.simm16:x}"
def _disasm_smem(inst: SMEM) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
if name in ('s_gl1_inv', 's_dcache_inv', 's_dcache_inv_vol', 's_dcache_wb', 's_dcache_wb_vol', 's_icache_inv'): return name
soe, imm = getattr(inst, 'soe', 0) or getattr(inst, 'soffset_en', 0), getattr(inst, 'imm', 1)
is_rdna4 = _is_r4(inst)
offset = inst.ioffset if is_rdna4 else getattr(inst, 'offset', 0) # type: ignore[attr-defined]
is_rdna4 = 'rdna4' in inst.__class__.__module__
offset = inst.ioffset if is_rdna4 else getattr(inst, 'offset', 0)
if cdna:
if soe and imm: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
elif imm: off_s = f"0x{offset:x}"
@@ -296,9 +276,7 @@ def _disasm_smem(inst: SMEM) -> str:
else: off_s = decode_src(inst.soffset, cdna)
is_buffer = 'buffer' in name or 's_atc_probe_buffer' == name
sbase_idx, sbase_count = _unwrap(inst.sbase), 4 if is_buffer else 2
if sbase_count == 2: sbase_str = _fmt_src(sbase_idx, sbase_count, cdna)
elif sbase_idx <= 105: sbase_str = _sreg(sbase_idx, sbase_count)
else: sbase_str = _reg("ttmp", sbase_idx - 108, sbase_count)
sbase_str = _fmt_src(sbase_idx, sbase_count, cdna) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {_unwrap(inst.sdata)}, {sbase_str}, {off_s}"
if 'prefetch' in name:
off = getattr(inst, 'ioffset', getattr(inst, 'offset', 0))
@@ -325,61 +303,53 @@ def _disasm_smem(inst: SMEM) -> str:
return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
def _disasm_flat(inst: FLAT) -> str:
name, cdna, r4 = inst.op_name.lower(), _is_cdna(inst), _is_r4(inst)
name, cdna = inst.op_name.lower(), _is_cdna(inst)
acc = getattr(inst, 'acc', 0)
reg_fn = _areg if acc else _vreg
if r4: seg = 'flat' if (cls_name:=inst.__class__.__name__) == 'VFLAT' else ('global' if cls_name == 'VGLOBAL' else 'scratch')
else: seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
# Global/scratch uses 13-bit signed offset
offset = inst.ioffset if r4 else inst.offset # type: ignore[attr-defined]
if seg != 'flat':
if cdna:
# CDNA: bit 12 is sign bit but not in offset field
raw = int.from_bytes(inst.to_bytes(), 'little')
off_val = offset | ((raw >> 12) & 1) << 12 # get bit 12
off_val = inst.offset | ((raw >> 12) & 1) << 12 # get bit 12
else:
off_val = offset
off_val = inst.offset
off_val = off_val if off_val < 4096 else off_val - 8192 # sign extend 13-bit
else:
off_val = offset
off_val = inst.offset
# Use get_field_bits: data for stores/atomics, d for loads
regs = inst.canonical_op_regs
w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
off_s = f" offset:{off_val}" if off_val else ""
if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}" # type: ignore[attr-defined]
elif r4: mods = f"{off_s}{' scope' if inst.scope else ''}{' th' if inst.th else ''}" # type: ignore[attr-defined]
if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"
else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
if seg == 'flat': saddr_s = ""
elif _unwrap(inst.saddr) in (0x7F, 124): saddr_s = ", off"
elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr, cdna)}"
elif _unwrap(inst.saddr) in (SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS):
saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
elif _unwrap(inst.saddr) in (SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS): saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
# RDNA4: vaddr instead of addr, vsrc instead of data
addr = inst.vaddr if r4 else inst.addr # type: ignore[attr-defined]
data = inst.vsrc if r4 else inst.data # type: ignore[attr-defined]
# load_lds_* instructions: vaddr, saddr (no vdst, data goes to LDS)
if 'load_lds' in name:
addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
return f"{instr} {addr_s}{saddr_s}{mods}"
if seg == 'flat': addr_w = 2 # flat always uses 64-bit vaddr
elif cdna: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
else: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
data_s, vdst_s = reg_fn(data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
data_s, vdst_s = reg_fn(inst.data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
glc_or_sc0 = inst.sc0 if cdna else inst.glc
if 'atomic' in name:
glc_or_sc0 = inst.sc0 if cdna else inst.glc # type: ignore[attr-defined]
sfx = f"{saddr_s if seg != 'flat' else ''}{mods}"
return f"{instr} {vdst_s}, {addr_s}, {data_s}{sfx}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{sfx}"
return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}"
if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
return f"{instr} {reg_fn(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
def _disasm_ds(inst: DS) -> str:
name = inst.op_name.lower()
op, name = inst.op, inst.op_name.lower()
acc = getattr(inst, 'acc', 0)
reg_fn = _areg if acc else _vreg
gds = " gds" if getattr(inst, 'gds', 0) else ""
@@ -408,8 +378,7 @@ def _disasm_ds(inst: DS) -> str:
if 'write2' in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
if 'read2' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}{off2}{gds}"
if 'xchg2' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}, {d0}, {d1}{off2}{gds}"
if 'load' in name or ('read' in name and 'read2' not in name):
return f"{name} {reg_fn(inst.vdst)}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
if 'load' in name or ('read' in name and 'read2' not in name): return f"{name} {reg_fn(inst.vdst)}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
if ('store' in name or 'write' in name) and not _has(name, 'cmp', 'xchg', 'write2'):
return f"{name} {reg_fn(inst.data0)}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
if 'swizzle' in name or name == 'ds_ordered_count': return f"{name} {reg_fn(inst.vdst)}, {addr}{off}{gds}"
@@ -420,15 +389,13 @@ def _disasm_ds(inst: DS) -> str:
return f"{name} {dst}, {addr}, {d0}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}{off}{gds}"
def _disasm_vop3(inst: VOP3) -> str:
name = inst.op_name.lower()
op, name = inst.op, inst.op_name.lower()
n_up = name.upper()
bits = inst.canonical_op_bits
# RDNA4 v_s_* scalar VOP3 instructions - vdst is SGPR (VGPRField adds 256)
if name.startswith('v_s_'):
s0v = _unwrap(inst.src0)
if s0v == 255: src = _lit(inst, inst.src0)
elif s0v == 253: src = "src_scc"
else: src = _fmt_src(inst.src0, max(1, bits['s0'] // 32))
src = _lit(inst, inst.src0) if _unwrap(inst.src0) == 255 else ("src_scc" if _unwrap(inst.src0) == 253 else _fmt_src(inst.src0, max(1, bits['s0'] // 32)))
if inst.neg & 1: src = f"-{src}"
if inst.abs & 1: src = f"|{src}|"
clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
@@ -437,6 +404,7 @@ def _disasm_vop3(inst: VOP3) -> str:
# Use get_field_bits for register sizes and 16-bit detection
r0, r1, r2 = max(1, bits['s0'] // 32), max(1, bits['s1'] // 32), max(1, bits['s2'] // 32)
dn = max(1, bits['d'] // 32)
is16_d, is16_s, is16_s2 = bits['d'] == 16, bits['s0'] == 16, bits['s2'] == 16
s0 = _vop3_src(inst, inst.src0, inst.neg&1, inst.abs&1, inst.opsel&1, r0, is16_s)
@@ -452,8 +420,7 @@ def _disasm_vop3(inst: VOP3) -> str:
clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
cl, om = " clamp" if clamp else "", _omod(inst.omod)
nonvgpr_opsel = ((inst.src0.offset < 256 and (inst.opsel & 1)) or (inst.src1.offset < 256 and (inst.opsel & 2))
or (inst.src2.offset < 256 and (inst.opsel & 4)))
nonvgpr_opsel = (inst.src0.offset < 256 and (inst.opsel & 1)) or (inst.src1.offset < 256 and (inst.opsel & 2)) or (inst.src2.offset < 256 and (inst.opsel & 4))
need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)
op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
@@ -491,19 +458,20 @@ def _disasm_vop3sd(inst: VOP3SD) -> str:
def _disasm_vopd(inst: VOPD) -> str:
lit = inst._literal
op_enum = R4_VOPDOp if _is_r4(inst) else VOPDOp
nx, ny = op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
is_rdna4 = 'rdna4' in inst.__class__.__module__
op_enum = R4_VOPDOp if is_rdna4 else VOPDOp
vdst_y, nx, ny = (_unwrap(inst.vdsty) << 1) | ((_unwrap(inst.vdstx) & 1) ^ 1), op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
def half(n, vd, s0, vs1):
vd, vs1 = _vi(vd), _vi(vs1)
if 'mov' in n: return f"{n} v{vd}, {_lit(inst, s0)}"
if 'fmamk' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, 0x{lit:x}, v{vs1}"
if 'fmaak' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}, 0x{lit:x}"
return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}"
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, inst.vdsty, inst.srcy0, inst.vsrcy1)}"
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
def _disasm_vop3p(inst: VOP3P) -> str:
name = inst.op_name.lower()
is_swmmac, n, is_fma_mix = 'swmmac' in name, inst.num_srcs() or 2, 'fma_mix' in name
is_wmma, is_swmmac, n, is_fma_mix = 'wmma' in name, 'swmmac' in name, inst.num_srcs() or 2, 'fma_mix' in name
def get_src(reg):
return _lit(inst, reg.offset) if reg.offset == 255 else reg.fmt()
src0, src1, src2, dst = get_src(inst.src0), get_src(inst.src1), get_src(inst.src2), inst.vdst.fmt()
@@ -512,22 +480,18 @@ def _disasm_vop3p(inst: VOP3P) -> str:
if is_fma_mix:
def m(s, neg, abs_): return f"-{f'|{s}|' if abs_ else s}" if neg else (f"|{s}|" if abs_ else s)
src0, src1, src2 = m(src0, inst.neg & 1, inst.neg_hi & 1), m(src1, inst.neg & 2, inst.neg_hi & 2), m(src2, inst.neg & 4, inst.neg_hi & 4)
mods = (([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else [])
+ ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if clamp else []))
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if clamp else [])
elif is_swmmac:
mods = ([f"index_key:{inst.opsel}"] if inst.opsel else []) + ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + \
([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
else:
opsel_hi_default = 7 if n == 3 else 3
mods = (([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else [])
+ ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else [])
+ ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else [])
+ ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else []))
mod_s = ' ' + ' '.join(mods) if mods else ''
return f"{name} {dst}, {src0}, {src1}, {src2}{mod_s}" if n == 3 else f"{name} {dst}, {src0}, {src1}{mod_s}"
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else []) + \
([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
def _disasm_sop1(inst: SOP1) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
# Use get_field_bits for register sizes
regs = inst.canonical_op_regs
dst_regs, src_regs = regs.get('d', 1), regs.get('s0', 1)
@@ -541,8 +505,8 @@ def _disasm_sop1(inst: SOP1) -> str:
try: msg_str = MSG(v).name if v != 255 else None # MSG_RTN_ILLEGAL_MSG (255) not supported by LLVM
except ValueError: msg_str = None
return f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, sendmsg({msg_str})" if msg_str else f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, 0x{v:x}"
sop1_src_only = ('S_ALLOC_VGPR', 'S_SLEEP_VAR', 'S_BARRIER_SIGNAL', 'S_BARRIER_SIGNAL_ISFIRST',
'S_BARRIER_INIT', 'S_BARRIER_JOIN', 'S_SET_GPR_IDX_IDX', 'S_CBRANCH_JOIN')
sop1_src_only = ('S_ALLOC_VGPR', 'S_SLEEP_VAR', 'S_BARRIER_SIGNAL', 'S_BARRIER_SIGNAL_ISFIRST', 'S_BARRIER_INIT', 'S_BARRIER_JOIN', 'S_SET_GPR_IDX_IDX',
'S_CBRANCH_JOIN')
if inst.op_name in sop1_src_only: return f"{name} {src}"
if cdna:
if 'getpc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2, cdna)}"
@@ -580,8 +544,8 @@ _HWREG_BLACKLIST_CDNA = {'HW_REG_PC_LO', 'HW_REG_PC_HI', 'HW_REG_IB_DBG1', 'HW_R
'HW_REG_SQ_SHADER_TMA_LO', 'HW_REG_SQ_SHADER_TMA_HI', 'HW_REG_SQ_PERF_SNAPSHOT_DATA', 'HW_REG_SQ_PERF_SNAPSHOT_DATA1',
'HW_REG_SQ_PERF_SNAPSHOT_PC_LO', 'HW_REG_SQ_PERF_SNAPSHOT_PC_HI', 'HW_REG_XCC_ID'}
def _disasm_sopk(inst: SOPK) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
is_rdna4 = _is_r4(inst)
op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
is_rdna4 = 'rdna4' in inst.__class__.__module__
hw = HWREG_CDNA if cdna else (HWREG_RDNA4 if is_rdna4 else HWREG)
blacklist = _HWREG_BLACKLIST_CDNA if cdna else _HWREG_BLACKLIST
def fmt_hwreg(hid, hoff, hsz):
@@ -603,14 +567,12 @@ def _disasm_sopk(inst: SOPK) -> str:
def _disasm_vinterp(inst: VINTERP) -> str:
mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
s0, s1, s2 = _lit(inst, inst.src0, inst.neg & 1), _lit(inst, inst.src1, inst.neg & 2), _lit(inst, inst.src2, inst.neg & 4)
return f"{inst.op_name.lower()} {inst.vdst.fmt()}, {s0}, {s1}, {s2}" + (" " + mods if mods else "")
return f"{inst.op_name.lower()} {inst.vdst.fmt()}, {_lit(inst, inst.src0, inst.neg & 1)}, {_lit(inst, inst.src1, inst.neg & 2)}, {_lit(inst, inst.src2, inst.neg & 4)}" + (" " + mods if mods else "")
DISASM_HANDLERS: dict[type, Callable[..., str]] = {
DISASM_HANDLERS: dict[type, callable] = {
VOP1: _disasm_vop1, VOP1_SDST: _disasm_vop1, VOP1_SDST_LIT: _disasm_vop1, VOP1_LIT: _disasm_vop1,
VOP2: _disasm_vop2, VOP2_LIT: _disasm_vop2, VOPC: _disasm_vopc, VOPC_LIT: _disasm_vopc,
VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3_SDST_LIT: _disasm_vop3, VOP3_LIT: _disasm_vop3,
VOP3SD: _disasm_vop3sd, VOP3SD_LIT: _disasm_vop3sd,
VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3_SDST_LIT: _disasm_vop3, VOP3_LIT: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOP3SD_LIT: _disasm_vop3sd,
VOPD: _disasm_vopd, VOPD_LIT: _disasm_vopd, VOP3P: _disasm_vop3p, VOP3P_LIT: _disasm_vop3p,
VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, GLOBAL: _disasm_flat, SCRATCH: _disasm_flat,
SOP1: _disasm_sop1, SOP1_LIT: _disasm_sop1, SOP2: _disasm_sop2, SOP2_LIT: _disasm_sop2,
@@ -620,7 +582,6 @@ DISASM_HANDLERS: dict[type, Callable[..., str]] = {
R4_VOP2: _disasm_vop2, R4_VOP2_LIT: _disasm_vop2, R4_VOPC: _disasm_vopc, R4_VOPC_LIT: _disasm_vopc,
R4_VOP3: _disasm_vop3, R4_VOP3_SDST: _disasm_vop3, R4_VOP3_SDST_LIT: _disasm_vop3, R4_VOP3_LIT: _disasm_vop3,
R4_VOP3SD: _disasm_vop3sd, R4_VOP3SD_LIT: _disasm_vop3sd, R4_VOP3P: _disasm_vop3p, R4_VOP3P_LIT: _disasm_vop3p,
R4_FLAT: _disasm_flat, R4_GLOBAL: _disasm_flat, R4_SCRATCH: _disasm_flat,
R4_VOPD: _disasm_vopd, R4_VOPD_LIT: _disasm_vopd, R4_VINTERP: _disasm_vinterp, R4_SOPP: _disasm_sopp, R4_SMEM: _disasm_smem, R4_DS: _disasm_ds,
R4_SOP1: _disasm_sop1, R4_SOP1_LIT: _disasm_sop1, R4_SOP2: _disasm_sop2, R4_SOP2_LIT: _disasm_sop2,
R4_SOPC: _disasm_sopc, R4_SOPC_LIT: _disasm_sopc, R4_SOPK: _disasm_sopk, R4_SOPK_LIT: _disasm_sopk}
@@ -631,11 +592,11 @@ def disasm(inst: Inst) -> str: return DISASM_HANDLERS[type(inst)](inst)
# CDNA DISASSEMBLER SUPPORT
# ═══════════════════════════════════════════════════════════════════════════════
from tinygrad.runtime.autogen.amd.cdna.ins import (VOP1 as CDNA_VOP1, VOP1_LIT as CDNA_VOP1_LIT,
from extra.assembly.amd.autogen.cdna.ins import (VOP1 as CDNA_VOP1, VOP1_LIT as CDNA_VOP1_LIT,
VOP1_SDWA as CDNA_VOP1_SDWA, VOP1_DPP16 as CDNA_VOP1_DPP16,
VOP2 as CDNA_VOP2, VOP2_LIT as CDNA_VOP2_LIT, VOP2_SDWA as CDNA_VOP2_SDWA, VOP2_DPP16 as CDNA_VOP2_DPP16,
VOPC as CDNA_VOPC, VOPC_LIT as CDNA_VOPC_LIT, VOPC_SDWA_SDST as CDNA_VOPC_SDWA_SDST,
VOP3 as CDNA_VOP3, VOP3_SDST as CDNA_VOP3_SDST, VOP3SD as CDNA_VOP3SD, VOP3P as CDNA_VOP3P, VOP3P_MFMA as CDNA_VOP3P_MFMA, VOP3PX2 as CDNA_VOP3PX2,
VOP3 as CDNA_VOP3, VOP3_SDST as CDNA_VOP3_SDST, VOP3SD as CDNA_VOP3SD, VOP3P as CDNA_VOP3P, VOP3PX2 as CDNA_VOP3PX2,
SOP1 as CDNA_SOP1, SOP1_LIT as CDNA_SOP1_LIT, SOP2 as CDNA_SOP2, SOP2_LIT as CDNA_SOP2_LIT,
SOPC as CDNA_SOPC, SOPC_LIT as CDNA_SOPC_LIT, SOPK as CDNA_SOPK, SOPK_LIT as CDNA_SOPK_LIT,
SOPP as CDNA_SOPP, SMEM as CDNA_SMEM, DS as CDNA_DS,
@@ -665,9 +626,7 @@ def _disasm_vop3a(inst) -> str:
else:
regs = inst.canonical_op_regs
dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
s0 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, r0)
s1 = _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, r1)
s2 = _cdna_src(inst, inst.src2, inst.neg&4, inst.abs&4, r2)
s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, r0), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, r1), _cdna_src(inst, inst.src2, inst.neg&4, inst.abs&4, r2)
dst = _vreg(inst.vdst, dregs) if dregs > 1 else _vreg(inst.vdst)
if op_val >= 512:
return f"{name} {dst}, {s0}, {s1}, {s2}{opsel}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{opsel}{cl}{om}"
@@ -691,9 +650,7 @@ def _disasm_vop3b(inst) -> str:
n = inst.num_srcs() or _num_srcs(inst)
regs = inst.canonical_op_regs
dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
s0 = _cdna_src(inst, inst.src0, inst.neg&1, n=r0)
s1 = _cdna_src(inst, inst.src1, inst.neg&2, n=r1)
s2 = _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, n=r0), _cdna_src(inst, inst.src1, inst.neg&2, n=r1), _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
# CDNA VOP3_SDST uses vdst field for sdst (but vdst adds 256), RDNA uses separate sdst field
sdst_val = getattr(inst, 'sdst', None)
if sdst_val is None and hasattr(inst, 'vdst'):
@@ -715,7 +672,7 @@ def _disasm_cdna_vop3p(inst) -> str:
name, n = inst.op_name.lower(), inst.num_srcs() or 2
is_mfma = 'mfma' in name or 'smfmac' in name
is_accvgpr = 'accvgpr' in name
def get_src(v, sc): return _lit(inst, v) if v == 255 else _fmt_src(v, sc, cdna=True)
get_src = lambda v, sc: _lit(inst, v) if v == 255 else _fmt_src(v, sc, cdna=True)
# Handle accvgpr read/write (accumulator register operations)
if is_accvgpr:
@@ -777,12 +734,9 @@ def _disasm_cdna_vop3p(inst) -> str:
src0, src1, src2, dst = get_src(inst.src0, 1), get_src(inst.src1, 1), get_src(inst.src2, 1), _vreg(inst.vdst)
opsel_hi = inst.opsel_hi # CDNA VOP3P only has 2 bits for opsel_hi (no opsel_hi2)
opsel_hi_default = 3 # CDNA default is 0b11 (2 bits), not 0b111 like RDNA
mods = (([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else [])
+ ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else [])
+ ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else [])
+ ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else []))
mod_s = ' ' + ' '.join(mods) if mods else ''
return f"{name} {dst}, {src0}, {src1}, {src2}{mod_s}" if n == 3 else f"{name} {dst}, {src0}, {src1}{mod_s}"
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else []) + \
([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
def _disasm_mubuf(inst) -> str:
name = inst.op_name.lower()
@@ -941,6 +895,5 @@ DISASM_HANDLERS.update({CDNA_VOP1: _disasm_vop1, CDNA_VOP1_LIT: _disasm_vop1,
CDNA_SOP1: _disasm_sop1, CDNA_SOP1_LIT: _disasm_sop1, CDNA_SOP2: _disasm_sop2, CDNA_SOP2_LIT: _disasm_sop2,
CDNA_SOPC: _disasm_sopc, CDNA_SOPC_LIT: _disasm_sopc, CDNA_SOPK: _disasm_sopk, CDNA_SOPK_LIT: _disasm_sopk, CDNA_SOPP: _disasm_sopp,
CDNA_SMEM: _disasm_smem, CDNA_DS: _disasm_ds, CDNA_FLAT: _disasm_flat, CDNA_GLOBAL: _disasm_flat, CDNA_SCRATCH: _disasm_flat,
CDNA_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b,
CDNA_VOP3P: _disasm_cdna_vop3p, CDNA_VOP3P_MFMA: _disasm_cdna_vop3p,
CDNA_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p,
CDNA_MUBUF: _disasm_mubuf, CDNA_VOP3PX2: _disasm_vop3px2})
@@ -1,24 +1,27 @@
# dsl.py - clean DSL for AMD assembly
from typing import Any
# ══════════════════════════════════════════════════════════════
# Registers - unified src encoding space (0-511)
# ══════════════════════════════════════════════════════════════
def _reg_size(t: str | None) -> int: return {'b64': 2, 'f64': 2, 'u64': 2, 'i64': 2, 'b128': 4}.get(t, 1)
class Reg:
# Register names vary by arch: RDNA has NULL@124/M0@125, CDNA has M0@124/reserved@125
# RDNA4 has DPP8@233, CDNA has SDWA@249/DPP@250/VCCZ@251/EXECZ@252
_NAMES = {102: "FLAT_SCRATCH_LO", 103: "FLAT_SCRATCH_HI", 104: "XNACK_MASK_LO", 105: "XNACK_MASK_HI",
106: "VCC_LO", 107: "VCC_HI", 124: "NULL", 125: "M0", 126: "EXEC_LO", 127: "EXEC_HI",
233: "DPP8", 234: "DPP8FI", 235: "SHARED_BASE", 236: "SHARED_LIMIT", 237: "PRIVATE_BASE", 238: "PRIVATE_LIMIT",
_NAMES = {106: "VCC_LO", 107: "VCC_HI", 124: "NULL", 125: "M0", 126: "EXEC_LO", 127: "EXEC_HI",
240: "0.5", 241: "-0.5", 242: "1.0", 243: "-1.0", 244: "2.0", 245: "-2.0", 246: "4.0", 247: "-4.0",
248: "INV_2PI", 249: "SDWA", 250: "DPP", 251: "VCCZ", 252: "EXECZ", 253: "SCC", 254: "SRC_LDS_DIRECT", 255: "LIT"}
248: "INV_2PI", 250: "DPP16", 253: "SCC", 255: "LIT"}
_PAIRS = {106: "VCC", 126: "EXEC"}
def __init__(self, offset: int = 0, sz: int = 512, *, neg: bool = False, abs_: bool = False, hi: bool = False):
self.offset, self.sz = offset, sz
self.neg, self.abs_, self.hi = neg, abs_, hi
# TODO: remove these legacy aliases
@property
def count(self): return self.sz
@property
def idx(self): return self.offset
def __hash__(self): return hash((self.offset, self.sz, self.neg, self.abs_, self.hi))
def __getitem__(self, key):
if isinstance(key, slice):
@@ -44,15 +47,11 @@ class Reg:
def fmt(self, sz=None, parens=False, upper=False) -> str:
o, sz = self.offset, sz or self.sz
l, r = ("[", "]") if parens or sz > 1 else ("", "") # brackets for multi-reg or when parens=True
if 256 <= o < 512:
idx = o - 256
base = f"v{l}{idx}{r}" if sz == 1 else f"v[{idx}:{idx + sz - 1}]"
if 256 <= o < 512: idx = o - 256; base = f"v{l}{idx}{r}" if sz == 1 else f"v[{idx}:{idx + sz - 1}]"
elif o < 106: base = f"s{l}{o}{r}" if sz == 1 else f"s[{o}:{o + sz - 1}]"
elif sz == 2 and o in self._PAIRS: base = self._PAIRS[o] if upper else self._PAIRS[o].lower()
elif o in self._NAMES: base = self._NAMES[o] if upper else self._NAMES[o].lower() # special regs (any sz)
elif 108 <= o < 124:
idx = o - 108
base = f"ttmp{l}{idx}{r}" if sz == 1 else f"ttmp[{idx}:{idx + sz - 1}]"
elif 108 <= o < 124: idx = o - 108; base = f"ttmp{l}{idx}{r}" if sz == 1 else f"ttmp[{idx}:{idx + sz - 1}]"
elif 128 <= o <= 192: base = str(o - 128) # inline int constants (0-64)
elif 193 <= o <= 208: base = str(-(o - 192)) # inline negative int constants (-1 to -16)
else: raise RuntimeError(f"unknown register: offset={o}, sz={sz}")
@@ -79,13 +78,9 @@ EXEC = src[126:127]
# 128: 0, 129-192: integers 1-64, 193-208: integers -1 to -16
# 240-248: float constants (0.5, -0.5, 1.0, -1.0, 2.0, -2.0, 4.0, -4.0, 1/(2*PI))
INV_2PI = src[248]
SDWA = src[249]
DPP = DPP16 = src[250]
VCCZ = src[251]
EXECZ = src[252]
DPP16 = src[250]
SCC = src[253]
SRC_LDS_DIRECT = src[254]
LIT = src[255] # literal constant marker
# 255: literal constant
v = src[256:511] # VGPR0-255
# ══════════════════════════════════════════════════════════════
@@ -98,13 +93,12 @@ class _Bits:
bits = _Bits()
class BitField:
name: str | None
def __init__(self, hi: int, lo: int, default = 0):
def __init__(self, hi: int, lo: int, default: int = 0):
self.hi, self.lo, self.default, self.name, self.mask = hi, lo, default, None, (1 << (hi - lo + 1)) - 1
def __set_name__(self, owner, name: str): self.name = name
def __eq__(self, other) -> 'FixedBitField': # type: ignore[override]
def __set_name__(self, owner, name): self.name = name
def __eq__(self, other) -> 'FixedBitField':
if isinstance(other, int): return FixedBitField(self.hi, self.lo, other)
raise TypeError(f"BitField.__eq__ expects int, got {type(other).__name__}")
return NotImplemented
def enum(self, enum_cls) -> 'EnumBitField': return EnumBitField(self.hi, self.lo, enum_cls)
def encode(self, val) -> int:
assert isinstance(val, int), f"BitField.encode expects int, got {type(val).__name__}"
@@ -113,14 +107,11 @@ class BitField:
def set(self, raw: int, val) -> int:
if val is None: val = self.default
encoded = self.encode(val)
# Handle signed values: convert negative to 2's complement
if encoded < 0: encoded = encoded & self.mask
if encoded < 0 or encoded > self.mask: raise RuntimeError(f"field '{self.name}': value {encoded} doesn't fit in {self.hi - self.lo + 1} bits")
return (raw & ~(self.mask << self.lo)) | (encoded << self.lo)
def __get__(self, obj, objtype=None):
if obj is None: return self
return self.decode((obj._raw >> self.lo) & self.mask)
def __set__(self, obj, val): obj._raw = self.set(obj._raw, val)
class FixedBitField(BitField):
def set(self, raw: int, val=None) -> int:
@@ -155,8 +146,7 @@ class SrcField(BitField):
expected_size = self._valid_range[1] - self._valid_range[0] + 1
actual_size = 1 << (hi - lo + 1)
if actual_size != expected_size:
raise RuntimeError(f"{self.__class__.__name__}: field size {hi - lo + 1} bits ({actual_size}) "
f"doesn't match range {self._valid_range} ({expected_size})")
raise RuntimeError(f"{self.__class__.__name__}: field size {hi - lo + 1} bits ({actual_size}) doesn't match range {self._valid_range} ({expected_size})")
def encode(self, val) -> int:
"""Encode value. Returns 255 (literal marker) for out-of-range values."""
@@ -178,10 +168,7 @@ class SrcField(BitField):
# Resize register based on operand info (skip non-resizable special registers)
# VCC/EXEC pairs (106, 126), NULL (124), M0 (125), float constants (240-255)
if reg.offset not in (124, 125) and not 240 <= reg.offset <= 255:
# Map variant field names (vsrc0->src0, vsrc1->src1, etc.) for DPP/SDWA classes
assert self.name is not None
name = self.name[1:] if self.name.startswith('v') and self.name[1:] in obj.op_regs else self.name
if sz := obj.op_regs.get(name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
if sz := obj.op_regs.get(self.name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
return reg
class VGPRField(SrcField):
@@ -224,21 +211,16 @@ class VDSTYField(BitField):
if not isinstance(val, Reg): raise TypeError(f"VDSTYField requires Reg, got {type(val).__name__}")
if not (256 <= val.offset < 512): raise ValueError(f"VDSTYField requires VGPR, got offset {val.offset}")
return (val.offset - 256) >> 1
def __get__(self, obj, objtype=None):
if obj is None: return self
raw = (obj._raw >> self.lo) & self.mask
vdstx_bit0 = (obj.vdstx.offset - 256) & 1
vgpr_idx = (raw << 1) | (vdstx_bit0 ^ 1)
return Reg(256 + vgpr_idx, 1)
def decode(self, raw): return raw # raw value, actual vdsty = (raw << 1) | ((vdstx & 1) ^ 1)
# ══════════════════════════════════════════════════════════════
# Operand info from XML
# ══════════════════════════════════════════════════════════════
import functools
from tinygrad.runtime.autogen.amd.rdna3.operands import OPERANDS as OPERANDS_RDNA3
from tinygrad.runtime.autogen.amd.rdna4.operands import OPERANDS as OPERANDS_RDNA4
from tinygrad.runtime.autogen.amd.cdna.operands import OPERANDS as OPERANDS_CDNA
from extra.assembly.amd.autogen.rdna3.operands import OPERANDS as OPERANDS_RDNA3
from extra.assembly.amd.autogen.rdna4.operands import OPERANDS as OPERANDS_RDNA4
from extra.assembly.amd.autogen.cdna.operands import OPERANDS as OPERANDS_CDNA
OPERANDS = {**OPERANDS_CDNA, **OPERANDS_RDNA3, **OPERANDS_RDNA4}
# ══════════════════════════════════════════════════════════════
@@ -258,15 +240,6 @@ def _get_variant(cls, suffix: str):
module = sys.modules.get(cls.__module__)
return getattr(module, f"{cls.__name__}{suffix}", None) if module else None
def _canonical_name(name: str) -> str | None:
"""Map operand name to canonical name."""
if name in ('src0', 'vsrc0', 'ssrc0'): return 's0'
if name in ('src1', 'vsrc1', 'ssrc1'): return 's1'
if name == 'src2': return 's2'
if name in ('vdst', 'sdst', 'sdata'): return 'd'
if name in ('data', 'vdata', 'data0', 'vsrc'): return 'data'
return None
class Inst:
_fields: list[tuple[str, BitField]]
_base_size: int
@@ -276,36 +249,37 @@ class Inst:
inherited = {}
for base in reversed(cls.__mro__[1:]):
if hasattr(base, '_fields'):
inherited.update(dict(base._fields))
inherited.update({name: field for name, field in base._fields})
inherited.update({name: val for name, val in cls.__dict__.items() if isinstance(val, BitField)})
cls._fields = list(inherited.items())
cls._base_size = (max(f.hi for _, f in cls._fields) + 8) // 8
def __new__(cls, *args, **kwargs):
# Auto-upgrade to variant if needed (only for base classes, not variants)
# Auto-upgrade to _LIT variant if needed (only for base classes, not variants)
if not any(cls.__name__.endswith(sfx) for sfx in ('_LIT', '_DPP16', '_DPP8', '_SDWA', '_SDWA_SDST', '_MFMA')):
args_iter = iter(args)
for name, field in cls._fields:
if isinstance(field, FixedBitField): continue
val = kwargs.get(name) if name in kwargs else next(args_iter, None)
if not isinstance(field, SrcField): continue
if isinstance(val, Reg) and val.offset == 255 and (lit_cls := _get_variant(cls, '_LIT')): return lit_cls(*args, **kwargs)
if isinstance(val, Reg) and val.offset == 249:
if (sdwa_cls := _get_variant(cls, '_SDWA') or _get_variant(cls, '_SDWA_SDST')): return sdwa_cls(*args, **kwargs)
if isinstance(val, Reg) and val.offset == 250 and (dpp_cls := _get_variant(cls, '_DPP16')): return dpp_cls(*args, **kwargs)
if _needs_literal(val) and (lit_cls := _get_variant(cls, '_LIT')): return lit_cls(*args, **kwargs)
lit_cls = _get_variant(cls, '_LIT')
if lit_cls is not None:
# Check if any src field needs a literal
# Map positional args to field names to find src values
args_iter = iter(args)
for name, field in cls._fields:
if isinstance(field, FixedBitField): continue
val = kwargs.get(name) if name in kwargs else next(args_iter, None)
if isinstance(field, SrcField) and _needs_literal(val):
return lit_cls(*args, **kwargs)
return object.__new__(cls)
def __init__(self, *args, **kwargs):
self._raw = 0
# Map positional args to field names (skip FixedBitFields)
args_iter = iter(args)
vals: dict[str, Any] = {}
vals = {}
for name, field in self._fields:
if isinstance(field, FixedBitField): vals[name] = None
elif name in kwargs: vals[name] = kwargs[name]
else: vals[name] = next(args_iter, None)
assert not (remaining := list(args_iter)), f"too many positional args: {remaining}"
remaining = list(args_iter)
assert not remaining, f"too many positional args: {remaining}"
# Extract modifiers from Reg objects and merge into neg/abs/opsel
neg_bits, abs_bits, opsel_bits = 0, 0, 0
for name, bit in [('src0', 0), ('src1', 1), ('src2', 2)]:
@@ -330,16 +304,16 @@ class Inst:
# Set all field values
for name, field in self._fields:
self._raw = field.set(self._raw, vals[name])
# Validate register sizes against operand info (skip special registers like NULL, VCC, EXEC, SDWA/DPP markers)
# Validate register sizes against operand info (skip special registers like NULL, VCC, EXEC)
for name, expected in self.op_regs.items():
if (val := vals.get(name)) is None: continue
if isinstance(val, Reg) and val.sz != expected and not (106 <= val.offset <= 127 or 249 <= val.offset <= 255):
if isinstance(val, Reg) and val.sz != expected and not (106 <= val.offset <= 127 or val.offset == 253):
raise TypeError(f"{name} expects {expected} register(s), got {val.sz}")
@property
def op_name(self) -> str: return getattr(self, 'op').name
def op_name(self) -> str: return self.op.name
@property
def operands(self) -> dict: return OPERANDS.get(getattr(self, 'op'), {}) if hasattr(self, 'op') else {}
def operands(self) -> dict: return OPERANDS.get(self.op, {}) if hasattr(self, 'op') else {}
def _is_cdna(self) -> bool: return 'cdna' in type(self).__module__
@functools.cached_property
@@ -351,9 +325,9 @@ class Inst:
if not self._is_cdna():
name = self.op_name.lower()
if 'cndmask' in name and 'src2' in bits: bits['src2'] = 32
if '_co_ci_' in name and 'src2' in bits: bits['src2'] = 32 # carry-in source
# VOP3SD: sdst is always wavefront-size dependent (carry-out or condition mask)
if 'VOP3SD' in type(self).__name__ and 'sdst' in bits: bits['sdst'] = 32
if '_co_ci_' in name:
if 'src2' in bits: bits['src2'] = 32
if 'sdst' in bits: bits['sdst'] = 32
if 'cmp' in name and 'vdst' in bits: bits['vdst'] = 32
# GLOBAL/FLAT: addr is 32-bit if saddr is valid SGPR, 64-bit if saddr is NULL
# SCRATCH: addr is always 32-bit (offset from scratch base, not absolute address)
@@ -367,8 +341,8 @@ class Inst:
# VGPRs: FP8/BF8(0,1)=8, FP6/BF6(2,3)=6, FP4(4)=4
if 'f8f6f4' in getattr(self, 'op_name', '').lower():
# Use explicit fields if available (VOP3PX2), else extract from VOP3P-MAI bit positions
cbsz = getattr(self, 'cbsz') if hasattr(type(self), 'cbsz') else (self._raw >> 8) & 0x7
blgp = getattr(self, 'blgp') if hasattr(type(self), 'blgp') else (self._raw >> 61) & 0x7
cbsz = self.cbsz if hasattr(type(self), 'cbsz') else (self._raw >> 8) & 0x7
blgp = self.blgp if hasattr(type(self), 'blgp') else (self._raw >> 61) & 0x7
vgprs = {0: 8, 1: 8, 2: 6, 3: 6, 4: 4}
bits['src0'], bits['src1'] = vgprs.get(cbsz, 8) * 32, vgprs.get(blgp, 8) * 32
return bits
@@ -382,17 +356,12 @@ class Inst:
"""Get bit widths with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
bits = {'d': 32, 's0': 32, 's1': 32, 's2': 32, 'data': 32}
for name, val in self.op_bits.items():
if (cn := _canonical_name(name)): bits[cn] = val
if name in ('src0', 'vsrc0', 'ssrc0'): bits['s0'] = val
elif name in ('src1', 'vsrc1', 'ssrc1'): bits['s1'] = val
elif name == 'src2': bits['s2'] = val
elif name in ('vdst', 'sdst', 'sdata'): bits['d'] = val
elif name in ('data', 'vdata', 'data0'): bits['data'] = val
return bits
@functools.cached_property
def canonical_operands(self) -> dict:
"""Get operands with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
result = {}
for name, val in self.operands.items():
if (cn := _canonical_name(name)): result[cn] = val
return result
@property
def canonical_op_regs(self) -> dict[str, int]:
"""Get register counts with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
@@ -408,7 +377,9 @@ class Inst:
@classmethod
def _size(cls) -> int: return cls._base_size
def size(self) -> int: return self._base_size
def disasm(self) -> str: raise NotImplementedError("disasm is no longer supported")
def disasm(self) -> str:
from extra.assembly.amd.disasm import disasm
return disasm(self)
def to_bytes(self) -> bytes: return self._raw.to_bytes(self._base_size, 'little')
+479
View File
@@ -0,0 +1,479 @@
# RDNA3 emulator - executes compiled pseudocode from AMD ISA PDF
# mypy: ignore-errors
from __future__ import annotations
import ctypes, functools
from enum import IntEnum
from tinygrad.runtime.autogen import hsa
from extra.assembly.amd.dsl import Inst, NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, v, s
from extra.assembly.amd.pcode import _f32, _i32, _sext, _f16, _i16, _f64, _i64
from extra.assembly.amd.decode import decode_inst
from extra.assembly.amd.pcode import compile_pseudocode
from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
from extra.assembly.amd.autogen.rdna3.ins import (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, GLOBAL, SCRATCH, VOPD,
SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, SCRATCHOp, VOPDOp)
# Constants and helpers defined locally (not imported from dsl.py)
MASK32, MASK64 = 0xFFFFFFFF, 0xFFFFFFFFFFFFFFFF
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}
class SGPRArray:
"""SGPR array indexed by Reg or int."""
__slots__ = ('_data',)
def __init__(self, size: int): self._data = [0] * size
def __getitem__(self, key): return self._data[getattr(key, 'offset', key)]
def __setitem__(self, key, val): self._data[getattr(key, 'offset', key)] = val
def __len__(self): return len(self._data)
def __iter__(self): return iter(self._data)
class VGPRLane:
"""Single lane of VGPRs indexed by Reg (offset 256-511) or int (0-255)."""
__slots__ = ('_data',)
def __init__(self, size: int): self._data = [0] * size
def __getitem__(self, key):
i = getattr(key, 'offset', key)
return self._data[i - 256 if i >= 256 else i]
def __setitem__(self, key, val):
i = getattr(key, 'offset', key)
self._data[i - 256 if i >= 256 else i] = val
def __len__(self): return len(self._data)
def __iter__(self): return iter(self._data)
WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
# 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) -> bool: return src.offset >= 256 and ((src.offset - 256) & 0x80) != 0
def _vgpr_masked(src): return v[(src.offset - 256) & 0x7f] if src.offset >= 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
src_off = src.offset
src_bits = inst.canonical_op_bits[f's{idx}']
literal, is_src_64, is_src_16 = inst._literal, src_bits == 64, src_bits == 16
if is_src_64: 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_off < 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_off >= 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: VGPRLane, reg, ndwords: int) -> int:
return sum(V[reg + i] << (32 * i) for i in range(ndwords))
# Helper: write multi-dword value to consecutive VGPRs
def _vgpr_write(V: VGPRLane, reg, val: int, ndwords: int):
for i in range(ndwords): V[reg + 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_E64, VOPDOp.V_DUAL_FMAAK_F32: VOP2Op.V_FMAAK_F32_E32, VOPDOp.V_DUAL_FMAMK_F32: VOP2Op.V_FMAMK_F32_E32,
VOPDOp.V_DUAL_MUL_F32: VOP3Op.V_MUL_F32_E64, VOPDOp.V_DUAL_ADD_F32: VOP3Op.V_ADD_F32_E64, VOPDOp.V_DUAL_SUB_F32: VOP3Op.V_SUB_F32_E64,
VOPDOp.V_DUAL_SUBREV_F32: VOP3Op.V_SUBREV_F32_E64, VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: VOP3Op.V_MUL_DX9_ZERO_F32_E64,
VOPDOp.V_DUAL_MOV_B32: VOP3Op.V_MOV_B32_E64, VOPDOp.V_DUAL_CNDMASK_B32: VOP3Op.V_CNDMASK_B32_E64,
VOPDOp.V_DUAL_MAX_F32: VOP3Op.V_MAX_F32_E64, VOPDOp.V_DUAL_MIN_F32: VOP3Op.V_MIN_F32_E64,
VOPDOp.V_DUAL_ADD_NC_U32: VOP3Op.V_ADD_NC_U32_E64, VOPDOp.V_DUAL_LSHLREV_B32: VOP3Op.V_LSHLREV_B32_E64, VOPDOp.V_DUAL_AND_B32: VOP3Op.V_AND_B32_E64,
}
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 = SGPRArray(SGPR_COUNT), [VGPRLane(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, reg) -> int:
if reg == NULL: return 0
if reg == SCC: return self.scc
return self.sgpr[reg]
def wsgpr(self, reg, v: int):
if reg != NULL: self.sgpr[reg] = v & MASK32
def rsgpr64(self, reg) -> int:
off = reg.offset
return self.sgpr._data[off] | (self.sgpr._data[off + 1] << 32)
def wsgpr64(self, reg, v: int):
off = reg.offset
self.sgpr._data[off] = v & MASK32; self.sgpr._data[off + 1] = (v >> 32) & MASK32
def _rsrc_base(self, reg, lane: int, consts, literal: int):
off = reg.offset
if off < SGPR_COUNT: return self.sgpr._data[off]
if off == SCC.offset: return self.scc
if off < 255: return consts[off - 128]
if off == 255: return literal
return self.vgpr[lane]._data[off - 256] if off <= 511 else 0
def rsrc(self, reg, lane: int, literal: int = 0) -> int: return self._rsrc_base(reg, lane, _INLINE_CONSTS, literal)
def rsrc_f16(self, reg, lane: int, literal: int = 0) -> int: return self._rsrc_base(reg, lane, _INLINE_CONSTS_F16, literal)
def rsrc64(self, reg, lane: int, literal: int = 0) -> int:
off = reg.offset
if 128 <= off < 255: return _INLINE_CONSTS_F64[off - 128]
if off == 255: return literal << 32 # 32-bit literal forms upper 32 bits of 64-bit value
return self.rsrc(reg, lane, literal) | ((self.rsrc(reg + 1, lane, literal) if off < VCC_LO.offset or 256 <= off <= 511 else 0) << 32)
def pend_sgpr_lane(self, reg, 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) + _sext(inst.offset, 21)
if inst.soffset != NULL: 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 canonical_op_bits to determine operand sizes
literal = inst._literal
s0 = st.rsrc64(ssrc0, 0, literal) if inst.canonical_op_bits['s0'] == 64 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.canonical_op_bits['s1'] == 64 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.canonical_op_bits['d'] == 64 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.canonical_op_bits['d'] == 64 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: VGPRLane, lane: int) -> None:
"""VOPD: dual-issue, execute two ops simultaneously (read all inputs before writes)."""
literal, vdstx = inst._literal, inst.vdstx
vdsty = v[(inst.vdsty << 1) | ((inst.vdstx.offset & 1) ^ 1)] # vdsty is raw int from VDSTYField.decode
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: VGPRLane, 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 != NULL 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: VGPRLane, 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: VGPRLane, lane: int) -> None:
"""VOP1/VOP2/VOP3/VOP3SD/VOP3P/VOPC: standard ALU ops."""
is_dst_16 = inst.canonical_op_bits['d'] == 16
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
neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst.offset & 0x80) != 0 and is_dst_16
if is_dst_16: vdst = v[inst.vdst.offset & 0x7f]
elif isinstance(inst, VOP2):
src0, src1, src2, vdst = inst.src0, inst.vsrc1, None, inst.vdst
neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst.offset & 0x80) != 0 and is_dst_16
if is_dst_16: vdst = v[inst.vdst.offset & 0x7f]
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, 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 is_dst_16: d0 = _src16(V[vdst], dst_hi)
elif inst.canonical_op_bits['d'] == 64: 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_E64, VOP3Op.V_CNDMASK_B16) and src2 is not None and src2.offset < 256: vcc_for_fn = st.rsgpr64(src2)
else: vcc_for_fn = st.vcc
src0_off = src0.offset if src0 is not None else 0
src0_idx = (src0_off - 256) if src0_off >= 256 else src0_off
vdst_off = vdst.offset
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_off, **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)
elif isinstance(inst, VOP2) and 'CO_CI' in inst.op_name: st.pend_sgpr_lane(VCC_LO, lane, (result['VCC'] >> lane) & 1)
elif is_vopc: st.pend_sgpr_lane(vdst, lane, (result['VCC'] >> lane) & 1) # vdst is VCC_LO for VOPC
else: st.pend_sgpr_lane(VCC_LO, 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.canonical_op_bits['d'] == 64: V[vdst], V[vdst + 1] = d0_val & MASK32, (d0_val >> 32) & MASK32
elif not isinstance(inst, VOP3P) and 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.offset, inst.src1.offset, inst.src2.offset, inst.vdst.offset
# 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 - 256 + (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 - 256 + i//32] = _i32(mat_d[i])
# ═══════════════════════════════════════════════════════════════════════════════
# PROGRAM DECODE
# ═══════════════════════════════════════════════════════════════════════════════
# Wave-level dispatch functions: (st, inst) -> return_code (0 = continue, -1 = end, -2 = barrier)
def dispatch_endpgm(st, inst): return -1
def dispatch_barrier(st, inst): st.pc += inst._words; return -2
def dispatch_nop(st, inst): st.pc += inst._words; return 0
def dispatch_wmma(st, inst): exec_wmma(st, inst, inst.op); st.pc += inst._words; return 0
def dispatch_writelane(st, inst): st.vgpr[st.rsrc(inst.src1, 0, inst._literal) & 0x1f][inst.vdst.offset - 256] = st.rsrc(inst.src0, 0, inst._literal) & MASK32; st.pc += inst._words; return 0
def dispatch_readlane(st, inst):
src0_off = inst.src0.offset
src0_idx = (src0_off - 256) if src0_off >= 256 else src0_off
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.offset)
st.wsgpr(inst.vdst.offset, result['D0'])
st.pc += inst._words; return 0
# Per-lane dispatch wrapper: wraps per-lane exec functions into wave-level dispatch
@functools.cache
def dispatch_lane(exec_fn):
def dispatch(st, inst):
exec_mask, vgpr, n_lanes = st.exec_mask, st.vgpr, st.n_lanes
for lane in range(n_lanes):
if exec_mask >> lane & 1: exec_fn(st, inst, vgpr[lane], lane)
st.commit_pends()
st.pc += inst._words
return 0
return dispatch
def decode_program(data: bytes) -> dict[int, Inst]:
result: dict[int, Inst] = {}
i = 0
while i < len(data):
inst = decode_inst(data[i:])
inst._words = inst.size() // 4
# Determine dispatch function and pcode function
if isinstance(inst, SOPP) and inst.op == SOPPOp.S_CODE_END: break
elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_ENDPGM: inst._dispatch = dispatch_endpgm
elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER: inst._dispatch = dispatch_barrier
elif isinstance(inst, SOPP) and inst.op in (SOPPOp.S_CLAUSE, SOPPOp.S_WAITCNT, SOPPOp.S_WAITCNT_DEPCTR, SOPPOp.S_SENDMSG, SOPPOp.S_SET_INST_PREFETCH_DISTANCE, SOPPOp.S_DELAY_ALU): inst._dispatch = dispatch_nop
elif isinstance(inst, (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM)): inst._dispatch = exec_scalar
elif isinstance(inst, VOP1) and inst.op == VOP1Op.V_NOP_E32: 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_E32, 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, GLOBAL, SCRATCH)): 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, PCODE[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, PCODE[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(s[0:1], 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
+822
View File
@@ -0,0 +1,822 @@
# DSL for RDNA3 pseudocode - makes pseudocode expressions work directly as Python
import struct, math, re, functools
MASK32, MASK64 = 0xFFFFFFFF, 0xFFFFFFFFFFFFFFFF
# Float/int bit conversion functions
_struct_f, _struct_I = struct.Struct("<f"), struct.Struct("<I")
_struct_e, _struct_H = struct.Struct("<e"), struct.Struct("<H")
_struct_d, _struct_Q = struct.Struct("<d"), struct.Struct("<Q")
def _f32(i):
i = i & MASK32
# RDNA3 default mode: flush f32 denormals to zero (FTZ)
# Denormal: exponent=0 (bits 23-30) and mantissa!=0 (bits 0-22)
if (i & 0x7f800000) == 0 and (i & 0x007fffff) != 0: return 0.0
return _struct_f.unpack(_struct_I.pack(i))[0]
def _i32(f):
if isinstance(f, int): f = float(f)
if math.isnan(f): return 0xffc00000 if math.copysign(1.0, f) < 0 else 0x7fc00000
if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
try:
bits = _struct_I.unpack(_struct_f.pack(f))[0]
# RDNA3 default mode: flush f32 denormals to zero (FTZ)
if (bits & 0x7f800000) == 0 and (bits & 0x007fffff) != 0: return 0x80000000 if bits & 0x80000000 else 0
return bits
except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
def _sext(v, b): return v - (1 << b) if v & (1 << (b - 1)) else v
def _f16(i): return _struct_e.unpack(_struct_H.pack(i & 0xffff))[0]
def _i16(f):
if math.isnan(f): return 0x7e00
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
try: return _struct_H.unpack(_struct_e.pack(f))[0]
except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
def _f64(i): return _struct_d.unpack(_struct_Q.pack(i & MASK64))[0]
def _i64(f):
if math.isnan(f): return 0x7ff8000000000000
if math.isinf(f): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
try: return _struct_Q.unpack(_struct_d.pack(f))[0]
except (OverflowError, struct.error): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
# ═══════════════════════════════════════════════════════════════════════════════
# INTERNAL HELPERS
# ═══════════════════════════════════════════════════════════════════════════════
def _div(a, b):
try: return a / b
except ZeroDivisionError:
if a == 0.0 or math.isnan(a): return float("nan")
return math.copysign(float("inf"), a * b) if b == 0.0 else float("inf") if a > 0 else float("-inf")
def _check_nan_type(x, quiet_bit_expected, default):
try:
if not math.isnan(float(x)): return False
if hasattr(x, '_reg') and hasattr(x, '_bits'):
bits = x._reg._val & ((1 << x._bits) - 1)
exp_bits, quiet_pos, mant_mask = {16: (0x1f, 9, 0x3ff), 32: (0xff, 22, 0x7fffff), 64: (0x7ff, 51, 0xfffffffffffff)}.get(x._bits, (0,0,0))
exp_shift = {16: 10, 32: 23, 64: 52}.get(x._bits, 0)
if exp_bits and ((bits >> exp_shift) & exp_bits) == exp_bits and (bits & mant_mask) != 0:
return ((bits >> quiet_pos) & 1) == quiet_bit_expected
return default
except (TypeError, ValueError): return False
def _gt_neg_zero(a, b): return (a > b) or (a == 0 and b == 0 and not math.copysign(1, a) < 0 and math.copysign(1, b) < 0)
def _lt_neg_zero(a, b): return (a < b) or (a == 0 and b == 0 and math.copysign(1, a) < 0 and not math.copysign(1, b) < 0)
def _fpop(fn):
def wrapper(x):
x = float(x)
if math.isnan(x) or math.isinf(x): return x
result = float(fn(x))
return math.copysign(0.0, x) if result == 0.0 else result
return wrapper
def _f_to_int(f, lo, hi): f = float(f); return 0 if math.isnan(f) else (hi if f >= hi else lo if f <= lo else int(f))
def _f16_to_f32_bits(bits): return struct.unpack("<e", struct.pack("<H", int(bits) & 0xffff))[0]
def _brev(v, bits): return int(bin(v & ((1 << bits) - 1))[2:].zfill(bits)[::-1], 2)
def _ctz(v, bits):
v, n = int(v) & ((1 << bits) - 1), 0
if v == 0: return bits
while (v & 1) == 0: v >>= 1; n += 1
return n
def _bf16(i):
"""Convert bf16 bits to float. BF16 is just the top 16 bits of f32."""
return struct.unpack("<f", struct.pack("<I", (i & 0xffff) << 16))[0]
def _ibf16(f):
"""Convert float to bf16 bits (truncate to top 16 bits of f32)."""
if math.isnan(f): return 0x7fc0 # bf16 quiet NaN
if math.isinf(f): return 0x7f80 if f > 0 else 0xff80 # bf16 ±infinity
try: return (struct.unpack("<I", struct.pack("<f", float(f)))[0] >> 16) & 0xffff
except (OverflowError, struct.error): return 0x7f80 if f > 0 else 0xff80
def _trig(fn, x):
# V_SIN/COS_F32: hardware does frac on input cycles before computing
if math.isinf(x) or math.isnan(x): return float("nan")
frac_cycles = fract(x / (2 * math.pi))
result = fn(frac_cycles * 2 * math.pi)
# Hardware returns exactly 0 for cos(π/2), sin(π), etc. due to lookup table
# Round very small results (below f32 precision) to exactly 0
if abs(result) < 1e-7: return 0.0
return result
class _SafeFloat(float):
"""Float subclass that uses _div for division to handle 0/inf correctly."""
def __truediv__(self, o): return _div(float(self), float(o))
def __rtruediv__(self, o): return _div(float(o), float(self))
class _Inf:
f16 = f32 = f64 = float('inf')
def __neg__(self): return _NegInf()
def __pos__(self): return self
def __float__(self): return float('inf')
def __eq__(self, other): return float(other) == float('inf') if not isinstance(other, _NegInf) else False
def __req__(self, other): return self.__eq__(other)
class _NegInf:
f16 = f32 = f64 = float('-inf')
def __neg__(self): return _Inf()
def __pos__(self): return self
def __float__(self): return float('-inf')
def __eq__(self, other): return float(other) == float('-inf') if not isinstance(other, _Inf) else False
def __req__(self, other): return self.__eq__(other)
class _RoundMode:
NEAREST_EVEN = 0
class _WaveMode:
IEEE = False
class _DenormChecker:
"""Comparator for denormalized floats. x == DENORM.f32 checks if x is denormalized."""
def __init__(self, bits): self._bits = bits
def _check(self, other):
f = float(other)
if math.isinf(f) or math.isnan(f) or f == 0.0: return False
if self._bits == 64:
bits = struct.unpack("<Q", struct.pack("<d", f))[0]
return (bits >> 52) & 0x7ff == 0
bits = struct.unpack("<I", struct.pack("<f", f))[0]
return (bits >> 23) & 0xff == 0
def __eq__(self, other): return self._check(other)
def __req__(self, other): return self._check(other)
def __ne__(self, other): return not self._check(other)
class _Denorm:
f32 = _DenormChecker(32)
f64 = _DenormChecker(64)
_pack = lambda hi, lo: ((int(hi) & 0xffff) << 16) | (int(lo) & 0xffff)
_pack32 = lambda hi, lo: ((int(hi) & 0xffffffff) << 32) | (int(lo) & 0xffffffff)
class TypedView:
"""View into a Reg with typed access. Used for both full-width (Reg.u32) and slices (Reg[31:16])."""
__slots__ = ('_reg', '_high', '_low', '_signed', '_float', '_bf16', '_reversed')
def __init__(self, reg, high, low=0, signed=False, is_float=False, is_bf16=False):
# Handle reversed slices like [0:31] which means bit-reverse
if high < low: high, low, reversed = low, high, True
else: reversed = False
self._reg, self._high, self._low, self._reversed = reg, high, low, reversed
self._signed, self._float, self._bf16 = signed, is_float, is_bf16
def _nbits(self): return self._high - self._low + 1
def _mask(self): return (1 << self._nbits()) - 1
def _get(self):
v = (self._reg._val >> self._low) & self._mask()
return _brev(v, self._nbits()) if self._reversed else v
def _set(self, v):
v = int(v)
if self._reversed: v = _brev(v, self._nbits())
self._reg._val = (self._reg._val & ~(self._mask() << self._low)) | ((v & self._mask()) << self._low)
@property
def _val(self): return self._get()
@property
def _bits(self): return self._nbits()
# Type accessors for slices (e.g., D0[31:16].f16)
u8 = property(lambda s: s._get() & 0xff)
u16 = property(lambda s: s._get() & 0xffff, lambda s, v: s._set(v))
u32 = property(lambda s: s._get() & MASK32, lambda s, v: s._set(v))
i16 = property(lambda s: _sext(s._get() & 0xffff, 16), lambda s, v: s._set(v))
i32 = property(lambda s: _sext(s._get() & MASK32, 32), lambda s, v: s._set(v))
f16 = property(lambda s: _f16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _i16(float(v))))
f32 = property(lambda s: _f32(s._get()), lambda s, v: s._set(_i32(float(v))))
bf16 = property(lambda s: _bf16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _ibf16(float(v))))
b16, b32 = u16, u32
# Chained type access (e.g., jump_addr.i64 when jump_addr is already TypedView)
@property
def i64(s): return s if s._nbits() == 64 and s._signed else int(s)
@property
def u64(s): return s if s._nbits() == 64 and not s._signed else int(s) & MASK64
def __getitem__(self, key):
if isinstance(key, slice):
high, low = int(key.start), int(key.stop)
return TypedView(self._reg, high, low)
return (self._get() >> int(key)) & 1
def __setitem__(self, key, value):
if isinstance(key, slice):
high, low = int(key.start), int(key.stop)
if high < low: high, low, value = low, high, _brev(int(value), low - high + 1)
mask = (1 << (high - low + 1)) - 1
self._reg._val = (self._reg._val & ~(mask << low)) | ((int(value) & mask) << low)
elif value: self._reg._val |= (1 << int(key))
else: self._reg._val &= ~(1 << int(key))
def __int__(self): return _sext(self._get(), self._nbits()) if self._signed else self._get()
def __index__(self): return int(self)
def __trunc__(self): return int(float(self)) if self._float else int(self)
def __float__(self):
if self._float:
if self._bf16: return _bf16(self._get())
bits = self._nbits()
return _f16(self._get()) if bits == 16 else _f32(self._get()) if bits == 32 else _f64(self._get())
return float(int(self))
def __bool__(s): return bool(int(s))
# Arithmetic - floats use float(), ints use int()
def __add__(s, o): return float(s) + float(o) if s._float else int(s) + int(o)
def __radd__(s, o): return float(o) + float(s) if s._float else int(o) + int(s)
def __sub__(s, o): return float(s) - float(o) if s._float else int(s) - int(o)
def __rsub__(s, o): return float(o) - float(s) if s._float else int(o) - int(s)
def __mul__(s, o): return float(s) * float(o) if s._float else int(s) * int(o)
def __rmul__(s, o): return float(o) * float(s) if s._float else int(o) * int(s)
def __truediv__(s, o): return _div(float(s), float(o)) if s._float else _div(int(s), int(o))
def __rtruediv__(s, o): return _div(float(o), float(s)) if s._float else _div(int(o), int(s))
def __pow__(s, o): return float(s) ** float(o) if s._float else int(s) ** int(o)
def __rpow__(s, o): return float(o) ** float(s) if s._float else int(o) ** int(s)
def __neg__(s): return -float(s) if s._float else -int(s)
def __abs__(s): return abs(float(s)) if s._float else abs(int(s))
# Bitwise - GPU shifts mask the shift amount to valid range
def __and__(s, o): return int(s) & int(o)
def __or__(s, o): return int(s) | int(o)
def __xor__(s, o): return int(s) ^ int(o)
def __invert__(s): return ~int(s)
def __lshift__(s, o): n = int(o); return int(s) << n if 0 <= n < 64 or s._nbits() > 64 else 0
def __rshift__(s, o): n = int(o); return int(s) >> n if 0 <= n < 64 or s._nbits() > 64 else 0
def __rand__(s, o): return int(o) & int(s)
def __ror__(s, o): return int(o) | int(s)
def __rxor__(s, o): return int(o) ^ int(s)
def __rlshift__(s, o): n = int(s); return int(o) << n if 0 <= n < 64 else 0
def __rrshift__(s, o): n = int(s); return int(o) >> n if 0 <= n < 64 else 0
# Comparison - handle _DenormChecker specially
def __eq__(s, o):
if isinstance(o, _DenormChecker): return o._check(s)
return float(s) == float(o) if s._float else int(s) == int(o)
def __ne__(s, o):
if isinstance(o, _DenormChecker): return not o._check(s)
return float(s) != float(o) if s._float else int(s) != int(o)
def __lt__(s, o): return float(s) < float(o) if s._float else int(s) < int(o)
def __le__(s, o): return float(s) <= float(o) if s._float else int(s) <= int(o)
def __gt__(s, o): return float(s) > float(o) if s._float else int(s) > int(o)
def __ge__(s, o): return float(s) >= float(o) if s._float else int(s) >= int(o)
class Reg:
"""GPU register: D0.f32 = S0.f32 + S1.f32 just works. Supports up to 128 bits for DS_LOAD_B128."""
__slots__ = ('_val',)
def __init__(self, val=0): self._val = int(val)
# Typed views - TypedView(reg, high, signed, is_float, is_bf16)
u64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
i64 = property(lambda s: TypedView(s, 63, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK64))
b64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
f64 = property(lambda s: TypedView(s, 63, is_float=True), lambda s, v: setattr(s, '_val', v if isinstance(v, int) else _i64(float(v))))
u32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
i32 = property(lambda s: TypedView(s, 31, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK32))
b32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
f32 = property(lambda s: TypedView(s, 31, is_float=True), lambda s, v: setattr(s, '_val', _i32(float(v))))
u24 = property(lambda s: TypedView(s, 23))
i24 = property(lambda s: TypedView(s, 23, signed=True))
u16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
i16 = property(lambda s: TypedView(s, 15, signed=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
b16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
f16 = property(lambda s: TypedView(s, 15, is_float=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _i16(float(v))) & 0xffff)))
bf16 = property(lambda s: TypedView(s, 15, is_float=True, is_bf16=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _ibf16(float(v))) & 0xffff)))
u8 = property(lambda s: TypedView(s, 7))
i8 = property(lambda s: TypedView(s, 7, signed=True))
u3 = property(lambda s: TypedView(s, 2)) # 3-bit for opsel fields
u1 = property(lambda s: TypedView(s, 0)) # single bit
def __getitem__(s, key):
if isinstance(key, slice): return TypedView(s, int(key.start), int(key.stop))
return (s._val >> int(key)) & 1
def __setitem__(s, key, value):
if isinstance(key, slice):
high, low = int(key.start), int(key.stop)
if high < low: high, low = low, high
mask = (1 << (high - low + 1)) - 1
s._val = (s._val & ~(mask << low)) | ((int(value) & mask) << low)
elif value: s._val |= (1 << int(key))
else: s._val &= ~(1 << int(key))
def __int__(s): return s._val
def __index__(s): return s._val
def __bool__(s): return bool(s._val)
# Arithmetic (for tmp = tmp + 1 patterns). Float operands trigger f32 interpretation.
def __add__(s, o): return (_f32(s._val) + float(o)) if isinstance(o, float) else s._val + int(o)
def __radd__(s, o): return (float(o) + _f32(s._val)) if isinstance(o, float) else int(o) + s._val
def __sub__(s, o): return (_f32(s._val) - float(o)) if isinstance(o, float) else s._val - int(o)
def __rsub__(s, o): return (float(o) - _f32(s._val)) if isinstance(o, float) else int(o) - s._val
def __mul__(s, o): return (_f32(s._val) * float(o)) if isinstance(o, float) else s._val * int(o)
def __rmul__(s, o): return (float(o) * _f32(s._val)) if isinstance(o, float) else int(o) * s._val
def __and__(s, o): return s._val & int(o)
def __rand__(s, o): return int(o) & s._val
def __or__(s, o): return s._val | int(o)
def __ror__(s, o): return int(o) | s._val
def __xor__(s, o): return s._val ^ int(o)
def __rxor__(s, o): return int(o) ^ s._val
def __lshift__(s, o): n = int(o); return s._val << n if 0 <= n < 64 else 0
def __rshift__(s, o): n = int(o); return s._val >> n if 0 <= n < 64 else 0
def __invert__(s): return ~s._val
# Comparison (for tmp >= 0x100000000 patterns)
def __lt__(s, o): return s._val < int(o)
def __le__(s, o): return s._val <= int(o)
def __gt__(s, o): return s._val > int(o)
def __ge__(s, o): return s._val >= int(o)
def __eq__(s, o): return s._val == int(o)
def __ne__(s, o): return s._val != int(o)
# ═══════════════════════════════════════════════════════════════════════════════
# PSEUDOCODE API - Functions and constants from AMD ISA pseudocode
# ═══════════════════════════════════════════════════════════════════════════════
# Rounding and float operations
trunc, floor, ceil = _fpop(math.trunc), _fpop(math.floor), _fpop(math.ceil)
def sqrt(x): return _SafeFloat(math.sqrt(x)) if x >= 0 else _SafeFloat(float("nan"))
def log2(x): return math.log2(x) if x > 0 else (float("-inf") if x == 0 else float("nan"))
def fract(x): return x - math.floor(x)
def sin(x): return _trig(math.sin, x)
def cos(x): return _trig(math.cos, x)
def pow(a, b):
try: return a ** b
except OverflowError: return float("inf") if b > 0 else 0.0
def isEven(x):
x = float(x)
if math.isinf(x) or math.isnan(x): return False
return int(x) % 2 == 0
def mantissa(f):
if f == 0.0 or math.isinf(f) or math.isnan(f): return f
m, _ = math.frexp(f)
return m # AMD V_FREXP_MANT returns mantissa in [0.5, 1.0) range
def signext_from_bit(val, bit):
bit = int(bit)
if bit == 0: return 0
mask = (1 << bit) - 1
val = int(val) & mask
if val & (1 << (bit - 1)): return val - (1 << bit)
return val
# Type conversions
i32_to_f32 = u32_to_f32 = i32_to_f64 = u32_to_f64 = f32_to_f64 = f64_to_f32 = float
def f32_to_i32(f): return _f_to_int(f, -2147483648, 2147483647)
def f32_to_u32(f): return _f_to_int(f, 0, 4294967295)
f64_to_i32, f64_to_u32 = f32_to_i32, f32_to_u32
def f32_to_f16(f):
f = float(f)
if math.isnan(f): return 0x7e00 # f16 NaN
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00 # f16 ±infinity
try: return struct.unpack("<H", struct.pack("<e", f))[0]
except OverflowError: return 0x7c00 if f > 0 else 0xfc00 # overflow -> ±infinity
def f16_to_f32(v): return v if isinstance(v, float) else _f16_to_f32_bits(v)
def i16_to_f16(v): return f32_to_f16(float(_sext(int(v) & 0xffff, 16)))
def u16_to_f16(v): return f32_to_f16(float(int(v) & 0xffff))
def f16_to_i16(bits): f = _f16_to_f32_bits(bits); return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
def f16_to_u16(bits): f = _f16_to_f32_bits(bits); return max(0, min(65535, int(f))) if not math.isnan(f) else 0
def bf16_to_f32(v): return _bf16(v) if isinstance(v, int) else float(v)
def f32_to_bf16(f): return _ibf16(f)
def u8_to_u32(v): return int(v) & 0xff
def u4_to_u32(v): return int(v) & 0xf
def u32_to_u16(u): return int(u) & 0xffff
def i32_to_i16(i): return ((int(i) + 32768) & 0xffff) - 32768
def f16_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
def f16_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
def f32_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
def f32_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
def v_cvt_i16_f32(f): return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
def v_cvt_u16_f32(f): return max(0, min(65535, int(f))) if not math.isnan(f) else 0
def SAT8(v): return max(0, min(255, int(v)))
def f32_to_u8(f): return max(0, min(255, int(f))) if not math.isnan(f) else 0
# Min/max operations
def v_min_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _lt_neg_zero(a, b) else b)
def v_max_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _gt_neg_zero(a, b) else b)
v_min_f16, v_max_f16 = v_min_f32, v_max_f32
v_min_i32, v_max_i32 = min, max
v_min_i16, v_max_i16 = min, max
def v_min_u32(a, b): return min(a & MASK32, b & MASK32)
def v_max_u32(a, b): return max(a & MASK32, b & MASK32)
def v_min_u16(a, b): return min(a & 0xffff, b & 0xffff)
def v_max_u16(a, b): return max(a & 0xffff, b & 0xffff)
def v_min3_f32(a, b, c): return v_min_f32(v_min_f32(a, b), c)
def v_max3_f32(a, b, c): return v_max_f32(v_max_f32(a, b), c)
v_min3_f16, v_max3_f16 = v_min3_f32, v_max3_f32
v_min3_i32, v_max3_i32, v_min3_i16, v_max3_i16 = min, max, min, max
def v_min3_u32(a, b, c): return min(a & MASK32, b & MASK32, c & MASK32)
def v_max3_u32(a, b, c): return max(a & MASK32, b & MASK32, c & MASK32)
def v_min3_u16(a, b, c): return min(a & 0xffff, b & 0xffff, c & 0xffff)
def v_max3_u16(a, b, c): return max(a & 0xffff, b & 0xffff, c & 0xffff)
# SAD/MSAD operations
def ABSDIFF(a, b): return abs(int(a) - int(b))
def v_sad_u8(s0, s1, s2):
"""V_SAD_U8: Sum of absolute differences of 4 byte pairs plus accumulator."""
s0, s1, s2 = int(s0), int(s1), int(s2)
result = s2
for i in range(4):
a = (s0 >> (i * 8)) & 0xff
b = (s1 >> (i * 8)) & 0xff
result += abs(a - b)
return result & 0xffffffff
def v_msad_u8(s0, s1, s2):
"""V_MSAD_U8: Masked sum of absolute differences (skip if reference byte is 0)."""
s0, s1, s2 = int(s0), int(s1), int(s2)
result = s2
for i in range(4):
a = (s0 >> (i * 8)) & 0xff
b = (s1 >> (i * 8)) & 0xff
if b != 0: # Only add diff if reference (s1) byte is non-zero
result += abs(a - b)
return result & 0xffffffff
def BYTE_PERMUTE(data, sel):
"""Select a byte from 64-bit data based on selector value."""
sel = int(sel) & 0xff
if sel <= 7: return (int(data) >> (sel * 8)) & 0xff
if sel == 8: return 0xff if ((int(data) >> 15) & 1) else 0x00
if sel == 9: return 0xff if ((int(data) >> 31) & 1) else 0x00
if sel == 10: return 0xff if ((int(data) >> 47) & 1) else 0x00
if sel == 11: return 0xff if ((int(data) >> 63) & 1) else 0x00
if sel == 12: return 0x00
return 0xff
# Pseudocode functions
def s_ff1_i32_b32(v): return _ctz(v, 32)
def s_ff1_i32_b64(v): return _ctz(v, 64)
GT_NEG_ZERO, LT_NEG_ZERO = _gt_neg_zero, _lt_neg_zero
def isNAN(x):
try: return math.isnan(float(x))
except (TypeError, ValueError): return False
def isQuietNAN(x): return _check_nan_type(x, 1, True)
def isSignalNAN(x): return _check_nan_type(x, 0, False)
def fma(a, b, c):
try: return math.fma(a, b, c)
except ValueError: return float('nan')
def ldexp(m, e): return math.ldexp(m, e)
def sign(f): return 1 if math.copysign(1.0, f) < 0 else 0
def exponent(f):
if hasattr(f, '_bits') and hasattr(f, '_float') and f._float:
raw = f._val
if f._bits == 16: return (raw >> 10) & 0x1f
if f._bits == 32: return (raw >> 23) & 0xff
if f._bits == 64: return (raw >> 52) & 0x7ff
f = float(f)
if math.isinf(f) or math.isnan(f): return 255
if f == 0.0: return 0
try: bits = struct.unpack("<I", struct.pack("<f", f))[0]; return (bits >> 23) & 0xff
except: return 0
def signext(x): return int(x)
def cvtToQuietNAN(x): return float('nan')
def F(x):
"""32'F(x) or 64'F(x) - interpret x as float. If x is int, treat as bit pattern."""
if isinstance(x, int): return _f32(x)
if isinstance(x, TypedView): return x
return float(x)
# Constants
PI = math.pi
WAVE32, WAVE64 = True, False
OVERFLOW_F32, UNDERFLOW_F32 = float('inf'), 0.0
OVERFLOW_F64, UNDERFLOW_F64 = float('inf'), 0.0
MAX_FLOAT_F32 = 3.4028235e+38
INF = _Inf()
ROUND_MODE = _RoundMode()
WAVE_MODE = _WaveMode()
DENORM = _Denorm()
# 2/PI with 1201 bits of precision for V_TRIG_PREOP_F64
TWO_OVER_PI_1201 = Reg(0x0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6)
# ═══════════════════════════════════════════════════════════════════════════════
# COMPILER: pseudocode -> Python (minimal transforms)
# ═══════════════════════════════════════════════════════════════════════════════
def _filter_pseudocode(pseudocode: str) -> str:
"""Filter raw PDF pseudocode to only include actual code lines."""
pcode_lines, in_lambda, depth = [], 0, 0
for line in pseudocode.split('\n'):
s = line.strip()
if not s: continue
if '=>' in s or re.match(r'^[A-Z_]+\(', s): continue # Skip example lines
if '= lambda(' in s: in_lambda += 1; continue # Skip lambda definitions
if in_lambda > 0:
if s.endswith(');'): in_lambda -= 1
continue
# Only include lines that look like pseudocode
is_code = (any(p in s for p in ['D0.', 'D1.', 'S0.', 'S1.', 'S2.', 'SCC =', 'SCC ?', 'VCC', 'EXEC', 'tmp =', 'tmp[', 'lane =', 'PC =',
'D0[', 'D1[', 'S0[', 'S1[', 'S2[', 'MEM[', 'RETURN_DATA', 'VADDR', 'VDATA', 'VDST', 'SADDR', 'OFFSET']) or
s.startswith(('if ', 'else', 'elsif', 'endif', 'declare ', 'for ', 'endfor', '//')) or
re.match(r'^[a-z_]+\s*=', s) or re.match(r'^[a-z_]+\[', s) or (depth > 0 and '=' in s))
if s.startswith('if '): depth += 1
elif s.startswith('endif'): depth = max(0, depth - 1)
if is_code: pcode_lines.append(s)
return '\n'.join(pcode_lines)
def _compile_pseudocode(pseudocode: str) -> str:
"""Compile pseudocode to Python. Transforms are minimal - most syntax just works."""
pseudocode = re.sub(r'\bpass\b', 'pass_', pseudocode) # 'pass' is Python keyword
raw_lines = pseudocode.strip().split('\n')
joined_lines: list[str] = []
for line in raw_lines:
line = line.strip()
if joined_lines and (joined_lines[-1].rstrip().endswith(('||', '&&', '(', ',')) or
(joined_lines[-1].count('(') > joined_lines[-1].count(')'))):
joined_lines[-1] = joined_lines[-1].rstrip() + ' ' + line
else:
joined_lines.append(line)
lines = []
indent, need_pass, in_first_match_loop = 0, False, False
for line in joined_lines:
line = line.split('//')[0].strip() # Strip C-style comments
if not line: continue
if line.startswith('if '):
lines.append(' ' * indent + f"if {_expr(line[3:].rstrip(' then'))}:")
indent += 1
need_pass = True
elif line.startswith('elsif '):
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
lines.append(' ' * indent + f"elif {_expr(line[6:].rstrip(' then'))}:")
indent += 1
need_pass = True
elif line == 'else':
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
lines.append(' ' * indent + "else:")
indent += 1
need_pass = True
elif line.startswith('endif'):
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
need_pass = False
elif line.startswith('endfor'):
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
need_pass, in_first_match_loop = False, False
elif line.startswith('declare '):
pass
elif m := re.match(r'for (\w+) in (.+?)\s*:\s*(.+?) do', line):
start, end = _expr(m[2].strip()), _expr(m[3].strip())
lines.append(' ' * indent + f"for {m[1]} in range({start}, int({end})+1):")
indent += 1
need_pass, in_first_match_loop = True, True
elif '=' in line and not line.startswith('=='):
need_pass = False
line = line.rstrip(';')
if m := re.match(r'\{\s*D1\.[ui]1\s*,\s*D0\.[ui]64\s*\}\s*=\s*(.+)', line):
rhs = _expr(m[1])
lines.append(' ' * indent + f"_full = {rhs}")
lines.append(' ' * indent + f"D0.u64 = int(_full) & 0xffffffffffffffff")
lines.append(' ' * indent + f"D1 = Reg((int(_full) >> 64) & 1)")
elif any(op in line for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^=')):
for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^='):
if op in line:
lhs, rhs = line.split(op, 1)
lines.append(' ' * indent + f"{lhs.strip()} {op} {_expr(rhs.strip())}")
break
else:
lhs, rhs = line.split('=', 1)
lhs_s, rhs_s = _expr(lhs.strip()), rhs.strip()
stmt = _assign(lhs_s, _expr(rhs_s))
if in_first_match_loop and rhs_s == 'i' and (lhs_s == 'tmp' or lhs_s == 'D0.i32'):
stmt += "; break"
lines.append(' ' * indent + stmt)
if need_pass: lines.append(' ' * indent + "pass")
return '\n'.join(lines)
def _assign(lhs: str, rhs: str) -> str:
if lhs in ('tmp', 'SCC', 'VCC', 'EXEC', 'D0', 'D1', 'saveexec', 'PC'):
return f"{lhs} = Reg({rhs})"
return f"{lhs} = {rhs}"
def _expr(e: str) -> str:
e = e.strip()
e = e.replace('&&', ' and ').replace('||', ' or ').replace('<>', ' != ')
e = re.sub(r'!([^=])', r' not \1', e)
e = re.sub(r'\{\s*(\w+\.u32)\s*,\s*(\w+\.u32)\s*\}', r'_pack32(\1, \2)', e)
def pack(m):
hi, lo = _expr(m[1].strip()), _expr(m[2].strip())
return f'_pack({hi}, {lo})'
e = re.sub(r'\{\s*([^,{}]+)\s*,\s*([^,{}]+)\s*\}', pack, e)
e = re.sub(r"1201'B\(2\.0\s*/\s*PI\)", "TWO_OVER_PI_1201", e)
e = re.sub(r"\d+'([0-9a-fA-Fx]+)[UuFf]*", r'\1', e)
e = re.sub(r"\d+'[FIBU]\(", "(", e)
e = re.sub(r'\bB\(', '(', e)
e = re.sub(r'([0-9a-fA-Fx])ULL\b', r'\1', e)
e = re.sub(r'([0-9a-fA-Fx])LL\b', r'\1', e)
e = re.sub(r'([0-9a-fA-Fx])U\b', r'\1', e)
e = re.sub(r'(\d\.?\d*)F\b', r'\1', e)
e = re.sub(r'(\[laneId\])\.[uib]\d+', r'\1', e)
e = e.replace('+INF', 'INF').replace('-INF', '(-INF)')
e = re.sub(r'NAN\.f\d+', 'float("nan")', e)
def convert_verilog_slice(m):
start, width = m.group(1).strip(), m.group(2).strip()
return f'[({start}) + ({width}) - 1 : ({start})]'
e = re.sub(r'\[([^:\[\]]+)\s*\+:\s*([^:\[\]]+)\]', convert_verilog_slice, e)
def process_brackets(s):
result, i = [], 0
while i < len(s):
if s[i] == '[':
depth, start = 1, i + 1
j = start
while j < len(s) and depth > 0:
if s[j] == '[': depth += 1
elif s[j] == ']': depth -= 1
j += 1
inner = _expr(s[start:j-1])
result.append('[' + inner + ']')
i = j
else:
result.append(s[i])
i += 1
return ''.join(result)
e = process_brackets(e)
while '?' in e:
depth, bracket, q = 0, 0, -1
for i, c in enumerate(e):
if c == '(': depth += 1
elif c == ')': depth -= 1
elif c == '[': bracket += 1
elif c == ']': bracket -= 1
elif c == '?' and depth == 0 and bracket == 0: q = i; break
if q < 0: break
depth, bracket, col = 0, 0, -1
for i in range(q + 1, len(e)):
if e[i] == '(': depth += 1
elif e[i] == ')': depth -= 1
elif e[i] == '[': bracket += 1
elif e[i] == ']': bracket -= 1
elif e[i] == ':' and depth == 0 and bracket == 0: col = i; break
if col < 0: break
cond, t, f = e[:q].strip(), e[q+1:col].strip(), e[col+1:].strip()
e = f'(({t}) if ({cond}) else ({f}))'
return e
def _apply_pseudocode_fixes(op_name: str, code: str) -> str:
"""Apply known fixes for PDF pseudocode bugs."""
if op_name == 'V_DIV_FMAS_F32':
code = code.replace('D0.f32 = 2.0 ** 32 * fma(S0.f32, S1.f32, S2.f32)',
'D0.f32 = (2.0 ** 64 if exponent(S2.f32) > 127 else 2.0 ** -64) * fma(S0.f32, S1.f32, S2.f32)')
if op_name == 'V_DIV_FMAS_F64':
code = code.replace('D0.f64 = 2.0 ** 64 * fma(S0.f64, S1.f64, S2.f64)',
'D0.f64 = (2.0 ** 128 if exponent(S2.f64) > 1023 else 2.0 ** -128) * fma(S0.f64, S1.f64, S2.f64)')
if op_name == 'V_DIV_SCALE_F32':
code = code.replace('D0.f32 = float("nan")', 'VCC = Reg(1 << laneId); 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(1 << laneId); 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(1 << laneId)')
if op_name == 'V_DIV_SCALE_F64':
code = code.replace('D0.f64 = float("nan")', 'VCC = Reg(1 << laneId); 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(1 << laneId); 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(1 << laneId)')
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]
+386
View File
@@ -0,0 +1,386 @@
"""SQTT (SQ Thread Trace) packet encoder and decoder for AMD GPUs.
This module provides encoding and decoding of raw SQTT byte streams.
The format is nibble-based with variable-width packets determined by a state machine.
Uses BitField infrastructure from dsl.py, similar to GPU instruction encoding.
"""
from __future__ import annotations
from typing import Iterator
from enum import Enum
from extra.assembly.amd.dsl import BitField, FixedBitField, bits
# ═══════════════════════════════════════════════════════════════════════════════
# FIELD ENUMS
# ═══════════════════════════════════════════════════════════════════════════════
class MemSrc(Enum):
LDS = 0
LDS_ALT = 1
VMEM = 2
VMEM_ALT = 3
class AluSrc(Enum):
NONE = 0
SALU = 1
VALU = 2
VALU_SALU = 3
class InstOp(Enum):
"""SQTT instruction operation types.
Memory ops appear in two ranges depending on which SIMD executes them:
- 0x1x-0x2x range: ops on traced SIMD
- 0x5x range: ops on other SIMD (OTHER_ prefix)
GLOBAL memory ops encoding depends on addressing mode AND size:
- Loads: 0x21 (saddr=SGPR) or 0x22 (saddr=NULL), all sizes same
- Stores: base + size_offset, where VADDR is shifted +1 from SADDR
SADDR: 0x24(32) 0x25(64) 0x26(96) 0x27(128)
VADDR: 0x25(32) 0x26(64) 0x27(96) 0x28(128)
OTHER_ range follows same pattern but values overlap differently.
"""
SALU = 0x0
SMEM = 0x1
JUMP = 0x3 # branch taken
JUMP_NO = 0x4 # branch not taken
MESSAGE = 0x9
VALU_TRANS = 0xb # transcendental: exp, log, rcp, sqrt, sin, cos
VALU_64_SHIFT = 0xd # 64-bit shifts: lshl, lshr, ashr
VALU_MAD64 = 0xe # 64-bit multiply-add
VALU_64 = 0xf # 64-bit: add, mul, fma, rcp, sqrt, rounding, frexp, div helpers
VINTERP = 0x12 # interpolation: v_interp_p10_f32, v_interp_p2_f32
BARRIER = 0x13
# FLAT memory ops on traced SIMD (0x1x range)
FLAT_LOAD = 0x1c
FLAT_STORE = 0x1d
FLAT_STORE_64 = 0x1e
FLAT_STORE_96 = 0x1f
FLAT_STORE_128 = 0x20
# GLOBAL memory ops on traced SIMD (0x2x range)
GLOBAL_LOAD = 0x21 # saddr=SGPR, all sizes
GLOBAL_LOAD_VADDR = 0x22 # saddr=NULL, all sizes
GLOBAL_STORE = 0x24 # saddr=SGPR, 32-bit
GLOBAL_STORE_64 = 0x25 # saddr=SGPR 64 or saddr=NULL 32
GLOBAL_STORE_96 = 0x26 # saddr=SGPR 96 or saddr=NULL 64
GLOBAL_STORE_128 = 0x27 # saddr=SGPR 128 or saddr=NULL 96
GLOBAL_STORE_VADDR_128 = 0x28 # saddr=NULL, 128-bit
# LDS ops on traced SIMD
LDS_LOAD = 0x29
LDS_STORE = 0x2b
LDS_STORE_64 = 0x2c
LDS_STORE_128 = 0x2e
# Memory ops on other SIMD (0x5x range)
OTHER_LDS_LOAD = 0x50
OTHER_LDS_STORE = 0x51
OTHER_LDS_STORE_64 = 0x52
OTHER_LDS_STORE_128 = 0x54
OTHER_FLAT_LOAD = 0x55
OTHER_FLAT_STORE = 0x56
OTHER_FLAT_STORE_64 = 0x57
OTHER_FLAT_STORE_96 = 0x58
OTHER_FLAT_STORE_128 = 0x59
OTHER_GLOBAL_LOAD = 0x5a # saddr=SGPR, all sizes
OTHER_GLOBAL_LOAD_VADDR = 0x5b # saddr=NULL or saddr=SGPR store 32
OTHER_GLOBAL_STORE_64 = 0x5c # saddr=SGPR 64 or saddr=NULL 32
OTHER_GLOBAL_STORE_96 = 0x5d # saddr=SGPR 96 or saddr=NULL 64
OTHER_GLOBAL_STORE_128 = 0x5e # saddr=SGPR 128 or saddr=NULL 96
OTHER_GLOBAL_STORE_VADDR_128 = 0x5f # saddr=NULL, 128-bit
# EXEC-modifying ops (0x7x range)
SALU_SAVEEXEC = 0x72 # s_*_saveexec_b32/b64
VALU_CMPX = 0x73 # v_cmpx_*
# ═══════════════════════════════════════════════════════════════════════════════
# PACKET TYPE BASE CLASS
# ═══════════════════════════════════════════════════════════════════════════════
class PacketType:
"""Base class for SQTT packet types."""
encoding: FixedBitField
_raw: int
_time: int
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
cls._fields = {k: v for k, v in cls.__dict__.items() if isinstance(v, BitField)}
cls._size_nibbles = ((max((f.hi for f in cls._fields.values()), default=0) + 4) // 4)
@classmethod
def from_raw(cls, raw: int, time: int = 0):
inst = object.__new__(cls)
inst._raw, inst._time = raw, time
return inst
def __repr__(self) -> str:
fields_str = ", ".join(f"{k}={getattr(self, k)}" for k in self._fields if not k.startswith('_'))
return f"{self.__class__.__name__}({fields_str})"
# ═══════════════════════════════════════════════════════════════════════════════
# PACKET TYPE DEFINITIONS
# ═══════════════════════════════════════════════════════════════════════════════
class VALUINST(PacketType): # exclude: 1 << 2
encoding = bits[2:0] == 0b011
delta = bits[5:3]
flag = bits[6:6]
wave = bits[11:7]
class VMEMEXEC(PacketType): # exclude: 1 << 0
encoding = bits[3:0] == 0b1111
delta = bits[5:4]
src = bits[7:6].enum(MemSrc)
class ALUEXEC(PacketType): # exclude: 1 << 1
encoding = bits[3:0] == 0b1110
delta = bits[5:4]
src = bits[7:6].enum(AluSrc)
class IMMEDIATE(PacketType): # exclude: 1 << 5
encoding = bits[3:0] == 0b1101
delta = bits[6:4]
wave = bits[11:7]
class IMMEDIATE_MASK(PacketType): # exclude: 1 << 5
encoding = bits[4:0] == 0b00100
delta = bits[7:5]
mask = bits[23:8]
class WAVERDY(PacketType): # exclude: 1 << 3
encoding = bits[4:0] == 0b10100
delta = bits[7:5]
mask = bits[23:8]
class TS_DELTA_S8_W3(PacketType):
encoding = bits[6:0] == 0b0100001
delta = bits[10:8]
_padding = bits[63:11]
class WAVEEND(PacketType): # exclude: 1 << 4
encoding = bits[4:0] == 0b10101
delta = bits[7:5]
flag7 = bits[8:8]
simd = bits[10:9]
cu_lo = bits[13:11]
wave = bits[19:15]
@property
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
class WAVESTART(PacketType): # exclude: 1 << 4
encoding = bits[4:0] == 0b01100
delta = bits[6:5]
flag7 = bits[7:7]
simd = bits[9:8]
cu_lo = bits[12:10]
wave = bits[17:13]
id7 = bits[31:18]
@property
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
class TS_DELTA_S5_W2(PacketType):
encoding = bits[4:0] == 0b11100
delta = bits[6:5]
_padding = bits[47:7]
class WAVEALLOC(PacketType): # exclude: 1 << 10
encoding = bits[4:0] == 0b00101
delta = bits[7:5]
_padding = bits[19:8]
class TS_DELTA_S5_W3(PacketType):
encoding = bits[4:0] == 0b00110
delta = bits[7:5]
_padding = bits[51:8]
class PERF(PacketType): # exclude: 1 << 11
encoding = bits[4:0] == 0b10110
delta = bits[7:5]
arg = bits[27:8]
class TS_DELTA_SHORT(PacketType):
encoding = bits[3:0] == 0b1000
delta = bits[7:4]
class NOP(PacketType):
encoding = bits[3:0] == 0b0000
delta = None # type: ignore
_padding = bits[3:0]
class TS_WAVE_STATE(PacketType):
encoding = bits[6:0] == 0b1010001
delta = bits[15:7]
coarse = bits[23:16]
@property
def wave_interest(self) -> bool: return bool(self.coarse & 1)
@property
def terminate_all(self) -> bool: return bool(self.coarse & 8)
class EVENT(PacketType): # exclude: 1 << 7
encoding = bits[7:0] == 0b01100001
delta = bits[10:8]
event = bits[23:11]
class EVENT_BIG(PacketType):
encoding = bits[7:0] == 0b11100001
delta = bits[10:8]
event = bits[31:11]
class REG(PacketType):
encoding = bits[3:0] == 0b1001
delta = bits[6:4]
slot = bits[9:7]
hi_byte = bits[15:8]
subop = bits[31:16]
val32 = bits[63:32]
@property
def is_config(self) -> bool: return bool(self.hi_byte & 0x80)
class SNAPSHOT(PacketType):
encoding = bits[6:0] == 0b1110001
delta = bits[9:7]
snap = bits[63:10]
class TS_DELTA_OR_MARK(PacketType):
encoding = bits[6:0] == 0b0000001
delta = bits[47:12]
bit8 = bits[8:8]
bit9 = bits[9:9]
@property
def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
class LAYOUT_HEADER(PacketType):
encoding = bits[6:0] == 0b0010001
delta = None # type: ignore
layout = bits[12:7]
simd = bits[14:13]
group = bits[17:15]
sel_a = bits[31:28]
sel_b = bits[36:33]
flag4 = bits[59:59]
_padding = bits[63:60]
class INST(PacketType):
encoding = bits[2:0] == 0b010
delta = bits[6:4]
flag1 = bits[3:3]
flag2 = bits[7:7]
wave = bits[12:8]
op = bits[19:13].enum(InstOp)
class UTILCTR(PacketType):
encoding = bits[6:0] == 0b0110001
delta = bits[8:7]
ctr = bits[47:9]
# All packet types in encoding priority order (more specific masks first, NOP last as fallback)
PACKET_TYPES: list[type[PacketType]] = [
EVENT, EVENT_BIG,
TS_DELTA_S8_W3, TS_WAVE_STATE, SNAPSHOT, TS_DELTA_OR_MARK, LAYOUT_HEADER, UTILCTR,
IMMEDIATE_MASK, WAVERDY, WAVEEND, WAVESTART, TS_DELTA_S5_W2, WAVEALLOC, TS_DELTA_S5_W3, PERF,
VMEMEXEC, ALUEXEC, IMMEDIATE, TS_DELTA_SHORT, REG,
VALUINST, INST,
NOP,
]
def _build_state_table() -> tuple[bytes, dict[int, type[PacketType]]]:
table = [len(PACKET_TYPES) - 1] * 256 # default to NOP
opcode_to_class: dict[int, type[PacketType]] = {i: cls for i, cls in enumerate(PACKET_TYPES)}
for byte_val in range(256):
for opcode, pkt_cls in enumerate(PACKET_TYPES):
if (byte_val & pkt_cls.encoding.mask) == pkt_cls.encoding.default:
table[byte_val] = opcode
break
return bytes(table), opcode_to_class
STATE_TO_OPCODE, OPCODE_TO_CLASS = _build_state_table()
# Precompute special case opcodes
_TS_DELTA_OR_MARK_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_OR_MARK)
_TS_DELTA_SHORT_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_SHORT)
# Combined lookup: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
# special_case: 0=none, 1=TS_DELTA_OR_MARK, 2=TS_DELTA_SHORT
_DECODE_INFO: dict[int, tuple] = {}
for _opcode, _pkt_cls in OPCODE_TO_CLASS.items():
_delta_field = getattr(_pkt_cls, 'delta', None)
_delta_lo = _delta_field.lo if _delta_field else 0
_delta_mask = _delta_field.mask if _delta_field else 0
_special = 1 if _opcode == _TS_DELTA_OR_MARK_OPCODE else (2 if _opcode == _TS_DELTA_SHORT_OPCODE else 0)
_DECODE_INFO[_opcode] = (_pkt_cls, _pkt_cls._size_nibbles, _delta_lo, _delta_mask, _special)
# ═══════════════════════════════════════════════════════════════════════════════
# DECODER
# ═══════════════════════════════════════════════════════════════════════════════
def decode(data: bytes) -> Iterator[PacketType]:
"""Decode raw SQTT blob, yielding packet instances."""
n, reg, pos, nib_off, nib_count, time = len(data), 0, 0, 0, 16, 0
while pos + ((nib_count + nib_off + 1) >> 1) <= n:
need = nib_count - nib_off
# 1. if unaligned, read high nibble to align
if nib_off: reg, pos = (reg >> 4) | ((data[pos] >> 4) << 60), pos + 1
# 2. read all full bytes at once
if (byte_count := need >> 1):
chunk = int.from_bytes(data[pos:pos + byte_count], 'little')
reg, pos = (reg >> (byte_count * 8)) | (chunk << (64 - byte_count * 8)), pos + byte_count
# 3. if odd, read low nibble
if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)
opcode = STATE_TO_OPCODE[reg & 0xFF]
pkt_cls, nib_count, delta_lo, delta_mask, special = _DECODE_INFO[opcode]
delta = (reg >> delta_lo) & delta_mask
if special == 1 and (reg >> 9) & 1 and not (reg >> 8) & 1: delta = 0 # TS_DELTA_OR_MARK marker
elif special == 2: delta += 8 # TS_DELTA_SHORT
time += delta
yield pkt_cls.from_raw(reg, time)
# ═══════════════════════════════════════════════════════════════════════════════
# PRINTER
# ═══════════════════════════════════════════════════════════════════════════════
PACKET_COLORS = {
"INST": "WHITE", "VALUINST": "BLACK", "VMEMEXEC": "yellow", "ALUEXEC": "yellow",
"IMMEDIATE": "YELLOW", "IMMEDIATE_MASK": "YELLOW", "WAVERDY": "cyan", "WAVEALLOC": "cyan",
"WAVEEND": "blue", "WAVESTART": "blue", "PERF": "magenta", "EVENT": "red", "EVENT_BIG": "red",
"REG": "green", "LAYOUT_HEADER": "white", "SNAPSHOT": "white", "UTILCTR": "green",
}
def format_packet(p) -> str:
from tinygrad.helpers import colored
name = type(p).__name__
if isinstance(p, INST):
op_name = p.op.name if isinstance(p.op, InstOp) else f"0x{p.op:02x}"
fields = f"wave={p.wave} op={op_name}" + (" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "")
elif isinstance(p, VALUINST): fields = f"wave={p.wave}" + (" flag" if p.flag else "")
elif isinstance(p, ALUEXEC): fields = f"src={p.src.name if isinstance(p.src, AluSrc) else p.src}"
elif isinstance(p, VMEMEXEC): fields = f"src={p.src.name if isinstance(p.src, MemSrc) else p.src}"
elif isinstance(p, (WAVESTART, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
elif hasattr(p, '_fields'):
fields = " ".join(f"{k}=0x{getattr(p, k):x}" if k in {'snap', 'val32'} else f"{k}={getattr(p, k)}"
for k in p._fields if not k.startswith('_') and k not in {'delta', 'encoding'})
else: fields = ""
return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name, 'white'))} {fields}"
def print_packets(packets) -> None:
skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"}
for p in packets:
if type(p).__name__ not in skip: print(format_packet(p))
if __name__ == "__main__":
import sys, pickle
if len(sys.argv) < 2:
print("Usage: python sqtt.py <pkl_file>")
sys.exit(1)
with open(sys.argv[1], "rb") as f:
data = pickle.load(f)
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
for i, event in enumerate(sqtt_events):
print(f"\n=== event {i} ===")
print_packets(decode(event.blob))
+132
View File
@@ -0,0 +1,132 @@
# maps SQTT trace packets to instructions.
from dataclasses import dataclass
from typing import Iterator
from tinygrad.runtime.support.elf import elf_loader
from extra.assembly.amd.sqtt import decode, print_packets, INST, VALUINST, IMMEDIATE, WAVESTART, WAVEEND, InstOp, PacketType, IMMEDIATE_MASK
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.decode import decode_inst
from extra.assembly.amd.autogen.rdna3.ins import SOPP, s_endpgm
from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
@dataclass(frozen=True)
class InstructionInfo:
pc: int
wave: int
inst: Inst
def map_insts(data:bytes, lib:bytes) -> Iterator[tuple[PacketType, InstructionInfo|None]]:
"""maps SQTT packets to instructions, yields (packet, instruction_info or None)"""
# map pcs to insts
pc_map:dict[int, Inst] = {}
image, sections, _ = elf_loader(lib)
text = next((sh for sh in sections if sh.name == ".text"), None)
assert text is not None, "no .text section found"
text_off, text_size = text.header.sh_addr, text.header.sh_size
offset = text_off
while offset < text_off + text_size:
inst = decode_inst(image[offset:])
pc_map[offset-text_off] = inst
offset += inst.size()
wave_pc:dict[int, int] = {}
# only processing packets on one [CU, SIMD] unit
def simd_select(p) -> bool: return getattr(p, "cu", 0) == 0 and getattr(p, "simd", 0) == 0
for p in decode(data):
if not simd_select(p): continue
if isinstance(p, WAVESTART):
assert p.wave not in wave_pc, "only one inflight wave per unit"
wave_pc[p.wave] = 0
continue
if isinstance(p, WAVEEND):
pc = wave_pc.pop(p.wave)
yield (p, InstructionInfo(pc, p.wave, s_endpgm()))
continue
# skip OTHER_ instructions, they don't belong to this unit
if isinstance(p, INST) and p.op.name.startswith("OTHER_"): continue
if isinstance(p, IMMEDIATE_MASK):
# immediate mask may yield multiple times per packet
for wave in range(16):
if p.mask & (1 << wave):
inst = pc_map[pc:=wave_pc[wave]]
# can this assert be more strict?
assert isinstance(inst, SOPP), f"IMMEDIATE_MASK packet must map to SOPP, got {inst}"
wave_pc[wave] += inst.size()
yield (p, InstructionInfo(pc, wave, inst))
continue
if isinstance(p, (VALUINST, INST, IMMEDIATE)):
inst = pc_map[pc:=wave_pc[p.wave]]
# s_delay_alu doesn't get a packet?
if isinstance(inst, SOPP) and inst.op in {SOPPOp.S_DELAY_ALU}:
wave_pc[p.wave] += inst.size()
inst = pc_map[pc:=wave_pc[p.wave]]
# identify a branch instruction, only used for asserts
is_branch = isinstance(inst, SOPP) and "BRANCH" in inst.op_name
if is_branch: assert isinstance(p, INST) and p.op in {InstOp.JUMP_NO, InstOp.JUMP}, f"branch can only be folowed by jump packets, got {p}"
# JUMP handling
if isinstance(p, INST) and p.op is InstOp.JUMP:
assert is_branch, f"JUMP packet must map to a branch instruction, got {inst}"
x = inst.simm16 & 0xffff
wave_pc[p.wave] += inst.size() + (x - 0x10000 if x & 0x8000 else x)*4
else:
if is_branch: assert inst.op != SOPPOp.S_BRANCH, f"S_BRANCH must have a JUMP packet, got {p}"
wave_pc[p.wave] += inst.size()
yield (p, InstructionInfo(pc, p.wave, inst))
continue
# for all other packets (VMEMEXEC, ALUEXEC, etc.), yield with None
yield (p, None)
# test to compare every packet with the rocprof decoder
def test_rocprof_inst_traces_match(sqtt, prg, target):
from tinygrad.viz.serve import llvm_disasm
from extra.sqtt.roc import decode as roc_decode, InstExec
disasm = {addr+prg.base:inst_disasm for addr, inst_disasm in llvm_disasm(target, prg.lib).items()}
rctx = roc_decode([sqtt], {prg.name:disasm})
rwaves = rctx.inst_execs[(sqtt.kern, sqtt.exec_tag)]
rwaves_iter:dict[int, list[Iterator[InstExec]]] = {} # wave unit (0-15) -> list of inst trace iterators for all executions on that unit
for w in rwaves: rwaves_iter.setdefault(w.wave_id, []).append(w.unpack_insts())
rwaves_base = next(iter(disasm)) # base program counter
passed_insts = 0
for pkt, info in map_insts(sqtt.blob, prg.lib):
if DEBUG >= 2: print_packets([pkt])
if info is None: continue
if DEBUG >= 2: print(f"{' '*29}{info.inst.disasm()}")
rocprof_inst = next(rwaves_iter[info.wave][0])
ref_pc = rocprof_inst.pc-rwaves_base
# always check pc matches
assert ref_pc == info.pc, f"pc mismatch {ref_pc}:{disasm[rocprof_inst.pc][0]} != {info.pc}:{info.inst.disasm()}"
# special handling for s_endpgm, it marks the wave completion.
if info.inst == s_endpgm():
completed_wave = list(rwaves_iter[info.wave].pop(0))
assert len(completed_wave) == 0, f"incomplete instructions in wave {info.wave}"
# otherwise the packet timestamp is time + "stall"
else:
assert pkt._time == rocprof_inst.time+rocprof_inst.stall
passed_insts += 1
for k,v in rwaves_iter.items():
assert len(v) == 0, f"incomplete wave {k}"
print(f"passed for {passed_insts} instructions across {len(rwaves)} waves scheduled on {len(rwaves_iter)} wave units")
if __name__ == "__main__":
import argparse, pickle, pathlib
from tinygrad.helpers import temp, DEBUG
parser = argparse.ArgumentParser()
parser.add_argument('--profile', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
default=pathlib.Path(temp("profile.pkl", append_user=True)))
parser.add_argument('--kernel', type=str, default=None, metavar="NAME", help='Kernel to focus on (optional name, default: all kernels)')
args = parser.parse_args()
with open(args.profile, "rb") as f:
data = pickle.load(f)
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
kern_events = {e.name:e for e in data if type(e).__name__ == "ProfileProgramEvent"}
target = next((e for e in data if type(e).__name__ == "ProfileDeviceEvent" and e.device.startswith("AMD"))).props["gfx_target_version"]
for e in sqtt_events:
if args.kernel is not None and args.kernel != e.kern: continue
if not e.itrace: continue
print(f"==== {e.kern}")
test_rocprof_inst_traces_match(e, kern_events[e.kern], target)
+192
View File
@@ -0,0 +1,192 @@
#!/usr/bin/env python3
"""Benchmark comparing Python vs Rust RDNA3 emulators on real tinygrad kernels."""
import ctypes, time, os
from pathlib import Path
# Set AMD=1 before importing tinygrad
os.environ["AMD"] = "1"
from extra.assembly.amd.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
if not REMU_PATH.exists():
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.dylib"
def get_rust_remu():
"""Load the Rust libremu shared library."""
if not REMU_PATH.exists(): return None
remu = ctypes.CDLL(str(REMU_PATH))
remu.run_asm.restype = ctypes.c_int32
remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
return remu
def count_instructions(kernel: bytes) -> int:
"""Count instructions in a kernel."""
return len(decode_program(kernel))
def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = None):
"""Allocate buffers and return args pointer + valid ranges."""
if init_data is None: init_data = {}
buffers = []
for i, size in enumerate(buf_sizes):
padded = ((size + 15) // 16) * 16 + 16
data = init_data.get(i, b'\x00' * padded)
data_list = list(data) + [0] * (padded - len(data))
buf = (ctypes.c_uint8 * padded)(*data_list[:padded])
buffers.append(buf)
args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
args_ptr = ctypes.addressof(args)
ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
ranges.add((args_ptr, ctypes.sizeof(args)))
return buffers, args, args_ptr, ranges
def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, rsrc2: int, iterations: int = 5):
"""Benchmark an emulator and return average time."""
gx, gy, gz = global_size
lx, ly, lz = local_size
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
lib_ptr = ctypes.addressof(kernel_buf)
# Warmup
run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
# Timed runs
times = []
for _ in range(iterations):
start = time.perf_counter()
result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
end = time.perf_counter()
if result != 0:
print(f" {name} returned error: {result}")
return None
times.append(end - start)
return sum(times) / len(times)
def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes], int] | None:
"""Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data, rsrc2)."""
try:
from tinygrad import Tensor
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.autogen import hsa
import numpy as np
np.random.seed(42)
ops = {
"add": lambda: Tensor.empty(1024) + Tensor.empty(1024),
"mul": lambda: Tensor.empty(1024) * Tensor.empty(1024),
"matmul_small": lambda: Tensor.empty(16, 16) @ Tensor.empty(16, 16),
"matmul_medium": lambda: Tensor.empty(64, 64) @ Tensor.empty(64, 64),
"reduce_sum": lambda: Tensor.empty(4096).sum(),
"reduce_max": lambda: Tensor.empty(4096).max(),
"softmax": lambda: Tensor.empty(256).softmax(),
"layernorm": lambda: Tensor.empty(32, 64).layernorm(),
"conv2d": lambda: Tensor.empty(1, 4, 16, 16).conv2d(Tensor.empty(4, 4, 3, 3)),
"gelu": lambda: Tensor.empty(1024).gelu(),
"exp": lambda: Tensor.empty(1024).exp(),
"sin": lambda: Tensor.empty(1024).sin(),
}
if op_name not in ops: return None
out = ops[op_name]()
sched = out.schedule()
for ei in sched:
lowered = ei.lower()
if ei.ast.op.name == 'SINK' and lowered.prg and lowered.prg.p.lib:
lib = bytes(lowered.prg.p.lib)
image = memoryview(bytearray(lib))
_, sections, _ = elf_loader(lib)
rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
for sec in sections:
if sec.name == '.text':
buf_sizes = [b.nbytes for b in lowered.bufs]
# Get initial data from numpy arrays if available
buf_data = {}
for i, buf in enumerate(lowered.bufs):
if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
try: buf_data[i] = bytes(buf.base._buf)
except: pass
# Extract rsrc2 from ELF (same as ops_amd.py)
group_segment_size = image[rodata_entry:rodata_entry+4].cast("I")[0]
lds_size = ((group_segment_size + 511) // 512) & 0x1FF
code = hsa.amd_kernel_code_t.from_buffer_copy(bytes(image[rodata_entry:rodata_entry+256]) + b'\x00'*256)
rsrc2 = code.compute_pgm_rsrc2 | (lds_size << 15)
return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data, rsrc2)
return None
except Exception as e:
print(f" Error getting kernel: {e}")
return None
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
def main():
import argparse
parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
args = parser.parse_args()
rust_remu = get_rust_remu()
if rust_remu is None:
print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
print("Running Python-only benchmarks...\n")
print("=" * 90)
print("RDNA3 Emulator Benchmark: Python vs Rust")
print("=" * 90)
results = []
print("\n[TINYGRAD KERNELS]")
print("-" * 90)
for op_name in TINYGRAD_TESTS:
print(f"\n{op_name}:", end=" ", flush=True)
kernel_info = get_tinygrad_kernel(op_name)
if kernel_info is None:
print("failed to compile")
continue
kernel, global_size, local_size, buf_sizes, buf_data, rsrc2 = kernel_info
n_insts = count_instructions(kernel)
n_workgroups = global_size[0] * global_size[1] * global_size[2]
n_threads = local_size[0] * local_size[1] * local_size[2]
total_work = n_insts * n_workgroups * n_threads
print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
set_valid_mem_ranges(ranges)
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations) if rust_remu else None
if py_time:
py_rate = total_work / py_time / 1e6
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
if rust_time:
rust_rate = total_work / rust_time / 1e6
speedup = py_time / rust_time if py_time else 0
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
results.append((op_name, n_insts, n_workgroups, py_time, rust_time))
# Summary table
print("\n" + "=" * 90)
print("SUMMARY")
print("=" * 90)
print(f"{'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
print("-" * 90)
for name, n_insts, n_wgs, py_time, rust_time in results:
py_ms = f"{py_time*1000:.3f}" if py_time else "error"
if rust_time:
rust_ms = f"{rust_time*1000:.3f}"
speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
else:
rust_ms, speedup = "N/A", "N/A"
print(f"{name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
if __name__ == "__main__":
main()
@@ -0,0 +1,196 @@
# Usability tests for the RDNA3 ASM DSL
# These tests demonstrate how the DSL *should* work for a good user experience
# Currently many of these tests fail - they document desired behavior
import unittest
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import Inst, RawImm, SGPR, VGPR
class TestRegisterSliceSyntax(unittest.TestCase):
"""
Issue: Register slice syntax should use AMD assembly convention (inclusive end).
In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
The DSL should match this convention so that:
- s[4:7] gives 4 registers
- Disassembler output can be copied directly back into DSL code
Fix: Change _RegFactory.__getitem__ to use inclusive end:
key.stop - key.start + 1 (instead of key.stop - key.start)
"""
def test_register_slice_count(self):
# s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
reg = s[4:7]
self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
def test_register_slice_roundtrip(self):
# Round-trip: DSL -> disasm -> DSL should preserve register count
reg = s[4:7] # 4 registers in AMD convention
inst = s_load_b128(reg, s[0:1], NULL, 0)
disasm = inst.disasm()
# Disasm shows s[4:7] - user should be able to copy this back
self.assertIn("s[4:7]", disasm)
# And s[4:7] in DSL should give the same 4 registers
reg_from_disasm = s[4:7]
self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
class TestReprReadability(unittest.TestCase):
"""
Issue: repr() leaks internal RawImm type and omits zero-valued fields.
When you create v_mov_b32_e32(v[0], v[1]), the repr shows:
VOP1(op=1, src0=RawImm(257))
Problems:
1. vdst=v[0] is omitted because 0 is treated as "default"
2. src0 shows RawImm(257) instead of v[1]
3. User sees encoded values (257 = 256 + 1) instead of register names
Expected repr: VOP1(op=1, vdst=v[0], src0=v[1])
"""
def test_repr_shows_registers_not_raw_imm(self):
inst = v_mov_b32_e32(v[0], v[1])
# Should show v[1], not RawImm(257)
self.assertNotIn("RawImm", repr(inst), "repr should not expose RawImm internal type")
self.assertIn("v[1]", repr(inst), "repr should show register name")
def test_repr_includes_zero_dst(self):
inst = v_mov_b32_e32(v[0], v[1])
# v[0] is a valid destination register, should be shown
self.assertIn("vdst", repr(inst), "repr should include vdst even when 0")
def test_repr_roundtrip(self):
# repr should produce something that can be eval'd back
inst = v_mov_b32_e32(v[0], v[1])
# This would require repr to output valid Python, e.g.:
# "VOP1(op=VOP1Op.V_MOV_B32, vdst=v[0], src0=v[1])"
r = repr(inst)
# At minimum, it should be human-readable
self.assertIn("v[", r, "repr should show register syntax")
class TestInstructionEquality(unittest.TestCase):
"""
Issue: No __eq__ method - instruction comparison requires repr() workaround.
Two identical instructions should compare equal with ==, but currently:
inst1 == inst2 returns False
The test_handwritten.py works around this with:
self.assertEqual(repr(self.inst), repr(reasm))
"""
def test_identical_instructions_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[1])
self.assertEqual(inst1, inst2, "identical instructions should be equal")
def test_different_instructions_not_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[2])
self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
class TestVOPDHelperSignature(unittest.TestCase):
"""
Issue: VOPD helper functions have confusing semantics.
v_dual_mul_f32 is defined as:
v_dual_mul_f32 = functools.partial(VOPD, VOPDOp.V_DUAL_MUL_F32)
This binds VOPDOp.V_DUAL_MUL_F32 to the FIRST positional arg of VOPD.__init__,
which is 'opx'. So v_dual_mul_f32 sets the X operation.
But then test_dual_mul in test_handwritten.py does:
v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], ...)
This passes V_DUAL_MUL_F32 as the SECOND positional arg (opy), making both
X and Y operations the same. This is confusing because:
1. The function name suggests it handles the X operation
2. But you still pass an opcode as the first arg (which becomes opy)
Expected: Either make the helper fully specify both ops, or make the
signature clearer about what the positional arg means.
"""
def test_vopd_helper_opy_should_be_required(self):
# Using only keyword args "works" but opy silently defaults to 0
inst = v_dual_mul_f32(vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32)
# Bug: opy defaults to 0 (V_DUAL_FMAC_F32) silently - should require explicit opy
# This test documents the bug - it should fail once fixed
self.assertNotEqual(inst.opy, VOPDOp.V_DUAL_FMAC_F32, "opy should not silently default to FMAC")
def test_vopd_helper_positional_arg_is_opy(self):
# The first positional arg after the partial becomes opy, not a second opx
inst = v_dual_mul_f32(VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32) # From partial
self.assertEqual(inst.opy, VOPDOp.V_DUAL_MOV_B32) # From first positional arg
class TestFieldAccessPreservesType(unittest.TestCase):
"""
Issue: Field access loses type information.
After creating an instruction, accessing fields returns encoded int values:
inst = v_mov_b32_e32(v[0], v[1])
inst.vdst # returns 0, not VGPR(0)
This makes it impossible to round-trip register types through field access.
"""
def test_vdst_returns_register(self):
inst = v_mov_b32_e32(v[5], v[1])
vdst = inst.vdst
# Should return a VGPR, not an int
self.assertIsInstance(vdst, (VGPR, int), "vdst should return VGPR or at least be usable")
# Ideally: self.assertIsInstance(vdst, VGPR)
def test_src_returns_register_for_vgpr_source(self):
inst = v_mov_b32_e32(v[0], v[1])
# src0 is encoded as 257 (256 + 1 for v1)
# Ideally it should decode back to v[1]
src0_raw = inst._values.get('src0')
# Currently returns RawImm(257), should return VGPR(1) or similar
self.assertNotIsInstance(src0_raw, RawImm, "source should not be RawImm internally")
class TestArgumentDiscoverability(unittest.TestCase):
"""
Issue: No clear signature for positional arguments.
inspect.signature(s_load_b128) shows: (*args, literal=None, **kwargs)
Users have no way to know the argument order without reading source code.
The order is implicitly defined by the class field definition order.
Possible fixes:
1. Add explicit parameter names to functools.partial
2. Generate type stubs with proper signatures
3. Add docstrings listing the expected arguments
"""
def test_signature_has_named_params(self):
import inspect
sig = inspect.signature(s_load_b128)
params = list(sig.parameters.keys())
# Currently: ['args', 'literal', 'kwargs'] (from *args, literal=None, **kwargs)
# Expected: something like ['sdata', 'sbase', 'soffset', 'offset', 'literal']
self.assertIn('sdata', params, "signature should show field names")
class TestSpecialConstants(unittest.TestCase):
"""
Issue: NULL and other constants are IntEnum values that might be confusing.
NULL = SrcEnum.NULL = 124, but users might expect NULL to be a special object
that clearly represents "no register" rather than a magic number.
"""
def test_null_has_clear_repr(self):
# NULL should have a clear string representation
self.assertIn("NULL", str(NULL) or repr(NULL), "NULL should be clearly identifiable")
def test_null_is_distinguishable_from_int(self):
# NULL should be distinguishable from the raw integer 124
self.assertNotEqual(type(NULL), int, "NULL should not be plain int")
if __name__ == "__main__":
unittest.main()
+79
View File
@@ -0,0 +1,79 @@
"""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")
ARCH_TO_TARGET:dict[str, list[str]] = {
"rdna3":["gfx1100"],
"rdna4":["gfx1200"],
"cdna":["gfx950", "gfx942"],
}
TARGET_TO_ARCH:dict[str, str] = {t:arch for arch,targets in ARCH_TO_TARGET.items() for t in targets}
def get_target(arch:str) -> str: return ARCH_TO_TARGET[arch][0]
def get_mattr(arch:str) -> str:
return {"rdna3":"+real-true16,+wavefrontsize32", "rdna4":"+real-true16,+wavefrontsize32", "cdna":"+wavefrontsize64"}[arch]
# ═══════════════════════════════════════════════════════════════════════════════
# 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}
@@ -4,16 +4,16 @@ Uses run_asm() with memory output, so tests can run on both emulator and real ha
Set USE_HW=1 to run on both emulator and hardware, comparing results.
"""
import ctypes, math, os, struct
from tinygrad.runtime.autogen.amd.rdna3.ins import *
from extra.assembly.amd.autogen.rdna3.ins import *
from test.mockgpu.amd.emu import run_asm
from tinygrad.renderer.amd.dsl import NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, M0
from extra.assembly.amd.emu import run_asm
from extra.assembly.amd.dsl import NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, M0
def _i32(f: float) -> int: return struct.unpack('<I', struct.pack('<f', f))[0]
def _f32(i: int) -> float: return struct.unpack('<f', struct.pack('<I', i & 0xFFFFFFFF))[0]
# f16 conversion helpers
def f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
def _f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
def f32_to_f16(f: float) -> int:
f = float(f)
if math.isnan(f): return 0x7e00
@@ -43,23 +43,6 @@ VCC = VCC_LO # For VOP3SD sdst field (VCC_LO is exported from dsl)
USE_HW = os.environ.get("USE_HW", "0") == "1"
FLOAT_TOLERANCE = 1e-5
def get_gpu_target() -> tuple[int, int, int]:
"""Get the GPU target as (major, minor, stepping) tuple."""
if not USE_HW: return (0, 0, 0)
from tinygrad.device import Device
return Device["AMD"].target # type: ignore[attr-defined]
def skip_unless_gfx(min_major: int, min_minor: int = 0, reason: str = ""):
"""Skip test if GPU target is below the minimum required version."""
import unittest
def decorator(test_func):
if not USE_HW: return test_func
target = get_gpu_target()
if target[0] < min_major or (target[0] == min_major and target[1] < min_minor):
return unittest.skip(reason or f"requires gfx{min_major}{min_minor}0+")(test_func)
return test_func
return decorator
# Output buffer layout: vgpr[16][32], sgpr[16], vcc, scc, exec
N_VGPRS, N_SGPRS, WAVE_SIZE = 16, 16, 32
VGPR_BYTES = N_VGPRS * WAVE_SIZE * 4 # 16 regs * 32 lanes * 4 bytes = 2048
@@ -75,7 +58,7 @@ def i642f(i: int) -> float: return struct.unpack('<d', struct.pack('<Q', i))[0]
def assemble(instructions: list) -> bytes:
return b''.join(inst.to_bytes() for inst in instructions)
# Simple WaveState class for test output parsing (mirrors test/mockgpu/amd/emu.py interface for tests)
# Simple WaveState class for test output parsing (mirrors emu.py interface for tests)
class WaveState:
def __init__(self):
self.vgpr = [[0] * 256 for _ in range(32)] # vgpr[lane][reg]
@@ -158,7 +141,7 @@ def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
# rsrc2: USER_SGPR_COUNT=2, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z=1, LDS_SIZE=128 (64KB)
rsrc2 = 0x19c | (128 << 15)
scratch_size = 0x10000 # 64KB per lane, matches .amdhsa_private_segment_fixed_size in run_program_hw
result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2, scratch_size)
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)
@@ -171,7 +154,7 @@ def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
from tinygrad.helpers import flat_mv
dev = Device["AMD"]
compiler = HIPCompiler(dev.arch) # type: ignore[attr-defined]
compiler = HIPCompiler(dev.arch)
prologue, epilogue = get_prologue_epilogue(n_lanes)
code = assemble(prologue + instructions + epilogue)
@@ -218,10 +201,9 @@ amdhsa.kernels:
"""
lib = compiler.compile(asm_src)
prg = AMDProgram(dev, "test", lib) # type: ignore[arg-type]
prg = AMDProgram(dev, "test", lib)
out_gpu = dev.allocator.alloc(OUT_BYTES)
assert out_gpu.va_addr % 16 == 0, f"buffer not 16-byte aligned: 0x{out_gpu.va_addr:x}"
prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
out_buf = bytearray(OUT_BYTES)
@@ -229,12 +211,8 @@ amdhsa.kernels:
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, ulp_tolerance: int = 0) -> list[str]:
"""Compare two WaveStates and return list of differences.
Args:
ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
"""
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):
@@ -245,11 +223,6 @@ def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgp
emu_f, hw_f = _f32(emu_val), _f32(hw_val)
if math.isnan(emu_f) and math.isnan(hw_f):
continue
# Check ULP difference for floats (only for same-sign values)
if ulp_tolerance > 0 and (emu_val < 0x80000000) == (hw_val < 0x80000000):
ulp_diff = abs(int(emu_val) - int(hw_val))
if ulp_diff <= ulp_tolerance:
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]
@@ -262,20 +235,17 @@ def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgp
diffs.append(f"scc: emu={emu_st.scc} hw={hw_st.scc}")
return diffs
def run_program(instructions: list, n_lanes: int = 1, ulp_tolerance: int = 0) -> WaveState:
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.
Args:
ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
"""
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, ulp_tolerance=ulp_tolerance)
diffs = compare_wave_states(emu_st, hw_st, n_lanes)
if diffs:
raise AssertionError("Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
raise AssertionError(f"Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
return hw_st
return emu_st
@@ -5,7 +5,7 @@ Includes: ds_store_b32, ds_load_b32, ds_store_2addr_*, ds_load_2addr_*,
ds_inc_*, ds_dec_*, ds_cmpstore_*, ds_storexchg_*
"""
import unittest
from test.amd.hw.helpers import *
from extra.assembly.amd.test.hw.helpers import *
class TestDS2Addr(unittest.TestCase):
"""Tests for DS_*_2ADDR instructions."""
@@ -138,50 +138,6 @@ class TestDS2AddrMore(unittest.TestCase):
self.assertEqual(st.vgpr[0][4], 0x12345678, "v4 should be untouched")
class TestDSB96(unittest.TestCase):
"""Tests for DS_STORE_B96 and DS_LOAD_B96 (96-bit / 3 dwords)."""
def test_ds_store_load_b96(self):
"""DS_STORE_B96 stores 3 VGPRs, DS_LOAD_B96 loads them back."""
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]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[2], s[0]),
ds_store_b96(addr=v[10], data0=v[0:2]),
s_waitcnt(lgkmcnt=0),
ds_load_b96(addr=v[10], vdst=v[4:6]),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have first dword")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have second dword")
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should have third dword")
def test_ds_store_b96_with_offset(self):
"""DS_STORE_B96 with non-zero offset."""
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]),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[2], s[0]),
DS(DSOp.DS_STORE_B96, addr=v[10], data0=v[0:2], offset0=12),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_B96, addr=v[10], vdst=v[4:6], offset0=12),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC)
class TestDSB128(unittest.TestCase):
"""Tests for DS_STORE_B128 and DS_LOAD_B128 (128-bit / 4 dwords)."""
@@ -601,6 +557,7 @@ class TestDS2AddrStride64(unittest.TestCase):
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "new val 0")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "new val 1")
def test_ds_storexchg_rtn_b64(self):
"""DS_STOREXCHG_RTN_B64: exchange 64-bit value and return old."""
instructions = [
@@ -718,47 +675,5 @@ class TestAtomicOrdering(unittest.TestCase):
self.assertEqual(st.vgpr[0][4], 150, "Final value should be 150")
class TestDsPermute(unittest.TestCase):
"""Tests for DS_PERMUTE_B32 and DS_BPERMUTE_B32 instructions."""
def test_ds_permute_b32_identity(self):
"""DS_PERMUTE_B32 with identity permutation (lane 0 sends to lane 0)."""
# For simplicity, test with single lane
instructions = [
v_mov_b32_e32(v[0], 0), # addr = 0 (lane 0)
v_mov_b32_e32(v[1], 0xDEADBEEF), # data
ds_permute_b32(v[2], v[0], v[1]),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
# Lane 0 sends to lane 0, so lane 0 gets 0xDEADBEEF
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
def test_ds_bpermute_b32_identity(self):
"""DS_BPERMUTE_B32 with identity permutation (each lane reads from itself)."""
instructions = [
v_mov_b32_e32(v[0], 0), # addr = 0 (read from lane 0)
v_mov_b32_e32(v[1], 0xCAFEBABE), # data in lane 0
ds_bpermute_b32(v[2], v[0], v[1]),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
# Lane 0 reads from lane 0's v[1]
self.assertEqual(st.vgpr[0][2], 0xCAFEBABE)
def test_ds_permute_b32_broadcast(self):
"""DS_PERMUTE_B32 broadcast - all lanes send to lane 0."""
# With 4 lanes, all sending to lane 0, highest lane wins
instructions = [
v_mov_b32_e32(v[0], 0), # All lanes send to addr 0 (lane 0)
v_mov_b32_e32(v[1], 0x11111111), # All lanes send same data
ds_permute_b32(v[2], v[0], v[1]),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=4)
# Lane 0 receives data (highest numbered active lane wins)
self.assertEqual(st.vgpr[0][2], 0x11111111)
if __name__ == '__main__':
unittest.main()
@@ -3,7 +3,7 @@
Includes: flat_load_*, flat_store_*, flat_atomic_*
"""
import unittest
from test.amd.hw.helpers import *
from extra.assembly.amd.test.hw.helpers import *
class TestFlatAtomic(unittest.TestCase):
"""Tests for FLAT atomic instructions."""
@@ -3,7 +3,7 @@
Includes: global_load_*, global_store_*, global_atomic_*, global_load_d16_*
"""
import unittest
from test.amd.hw.helpers import *
from extra.assembly.amd.test.hw.helpers import *
class TestGlobalAtomic(unittest.TestCase):
"""Tests for GLOBAL atomic instructions."""
@@ -523,157 +523,5 @@ class TestD16HiLoads(unittest.TestCase):
self.assertEqual(byte5, 0x00, f"byte5: expected 0x00, got 0x{byte5:02x}")
class TestGlobalOffset(unittest.TestCase):
"""Tests for GLOBAL instructions with different offsets.
These tests verify that instruction deduplication correctly handles different offset values.
If offset is made dynamic incorrectly, instructions with different offsets may load/store wrong data.
"""
def test_global_load_different_offsets(self):
"""Load from two different offsets and verify correct values."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0xAAAAAAAA at offset 100
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
# Store 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
s_waitcnt(vmcnt=0),
# Load from offset 100 -> should get 0xAAAAAAAA
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[3], saddr=SrcEnum.NULL, offset=100),
# Load from offset 200 -> should get 0xBBBBBBBB
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=200),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
v_mov_b32_e32(v[1], v[4]),
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], 0xAAAAAAAA, f"offset 100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
def test_global_store_different_offsets(self):
"""Store to two different offsets and verify correct values."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0x11111111 at offset 300
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=300),
# Store 0x22222222 at offset 400
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
global_store_b32(addr=v[0:1], data=v[3], saddr=SrcEnum.NULL, offset=400),
s_waitcnt(vmcnt=0),
# Load back to verify
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=300),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[5], saddr=SrcEnum.NULL, offset=400),
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], 0x11111111, f"offset 300: expected 0x11111111, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0x22222222, f"offset 400: expected 0x22222222, got 0x{st.vgpr[0][1]:08x}")
def test_global_negative_offset_no_saddr(self):
"""Test negative offset without saddr (VGPR pair for address).
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
Load with offset -100 from vaddr pointing to base+200 -> should get 0xAAAA (at 100).
Load with offset -100 from vaddr pointing to base+300 -> should get 0xBBBB (at 200)."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
s_waitcnt(vmcnt=0),
# vaddr = base+200, load with offset -100 -> should get value at 100
s_add_u32(s[4], s[2], 200),
s_addc_u32(s[5], s[3], 0),
v_mov_b32_e32(v[4], s[4]),
v_mov_b32_e32(v[5], s[5]),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[6], saddr=SrcEnum.NULL, offset=-100),
# vaddr = base+300, load with offset -100 -> should get value at 200
s_add_u32(s[4], s[2], 300),
s_addc_u32(s[5], s[3], 0),
v_mov_b32_e32(v[4], s[4]),
v_mov_b32_e32(v[5], s[5]),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[7], saddr=SrcEnum.NULL, offset=-100),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
v_mov_b32_e32(v[4], 0),
v_mov_b32_e32(v[5], 0),
v_mov_b32_e32(v[6], 0),
v_mov_b32_e32(v[7], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
s_mov_b32(s[4], 0),
s_mov_b32(s[5], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
def test_global_negative_offset_with_saddr(self):
"""Test negative offset with saddr (SGPR pair for base address).
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
Load with offset -100 from saddr pointing to base+200 -> should get 0xAAAA (at 100).
Load with offset -100 from saddr pointing to base+300 -> should get 0xBBBB (at 200)."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
s_waitcnt(vmcnt=0),
# saddr = base+200, load with offset -100 -> should get value at 100
s_add_u32(s[4], s[2], 200),
s_addc_u32(s[5], s[3], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[6], saddr=s[4:5], offset=-100),
# saddr = base+300, load with offset -100 -> should get value at 200
s_add_u32(s[4], s[2], 300),
s_addc_u32(s[5], s[3], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[7], saddr=s[4:5], offset=-100),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
v_mov_b32_e32(v[6], 0),
v_mov_b32_e32(v[7], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
s_mov_b32(s[4], 0),
s_mov_b32(s[5], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
if __name__ == '__main__':
unittest.main()
@@ -4,7 +4,7 @@ Includes: s_add_u32, s_mov_b32, s_and_b32, s_or_b32, s_quadmask_b32, s_wqm_b32,
s_cbranch_vccnz, s_cbranch_vccz
"""
import unittest
from test.amd.hw.helpers import *
from extra.assembly.amd.test.hw.helpers import *
class TestBasicScalar(unittest.TestCase):
"""Tests for basic scalar operations."""
@@ -62,7 +62,6 @@ class TestBasicScalar(unittest.TestCase):
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0x80000000)
@skip_unless_gfx(11, 5, "SALU FP ops require gfx1150+")
def test_s_fmamk_f32(self):
"""S_FMAMK_F32: D = S0 * literal + S1."""
# 2.0 * 3.0 + 1.0 = 7.0
@@ -74,7 +73,6 @@ class TestBasicScalar(unittest.TestCase):
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], f2i(7.0))
@skip_unless_gfx(11, 5, "SALU FP ops require gfx1150+")
def test_s_fmamk_f32_negative(self):
"""S_FMAMK_F32 with negative values."""
# -2.0 * 4.0 + 10.0 = 2.0
@@ -87,50 +85,6 @@ class TestBasicScalar(unittest.TestCase):
self.assertEqual(st.sgpr[2], f2i(2.0))
class TestPack(unittest.TestCase):
"""Tests for S_PACK instructions."""
def test_s_pack_ll_b32_b16(self):
"""S_PACK_LL_B32_B16 packs low 16 bits of two sources into one 32-bit result."""
instructions = [
s_mov_b32(s[0], 0xDEADAAAA),
s_mov_b32(s[1], 0xDEADBBBB),
s_pack_ll_b32_b16(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0xBBBBAAAA)
def test_s_pack_lh_b32_b16(self):
"""S_PACK_LH_B32_B16: D0 = { S1[31:16], S0[15:0] }."""
instructions = [
s_mov_b32(s[0], 0xDEADAAAA),
s_mov_b32(s[1], 0xDEADBBBB),
s_pack_lh_b32_b16(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0xDEADAAAA)
def test_s_pack_hh_b32_b16(self):
"""S_PACK_HH_B32_B16: D0 = { S1[31:16], S0[31:16] }."""
instructions = [
s_mov_b32(s[0], 0xDEADAAAA),
s_mov_b32(s[1], 0xDEADBBBB),
s_pack_hh_b32_b16(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0xDEADDEAD)
def test_s_pack_hl_b32_b16(self):
"""S_PACK_HL_B32_B16: D0 = { S1[15:0], S0[31:16] }."""
instructions = [
s_mov_b32(s[0], 0xDEADAAAA),
s_mov_b32(s[1], 0xDEADBBBB),
s_pack_hl_b32_b16(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0xBBBBDEAD)
class TestQuadmaskWqm(unittest.TestCase):
"""Tests for S_QUADMASK_B32 and S_WQM_B32."""
@@ -665,272 +619,5 @@ class Test64BitCompare(unittest.TestCase):
self.assertEqual(st.sgpr[4], 1)
class TestSOPPNop(unittest.TestCase):
"""Tests for S_NOP and other SOPP instructions with expression-based for loops.
S_NOP's pcode uses 'for i in 0U : SIMM16.u16[3 : 0].u32 do' which requires
the parser to handle non-constant loop bounds.
"""
def test_s_nop_basic(self):
"""S_NOP executes without side effects."""
# S_NOP with immediate 0 should just do nothing
instructions = [
s_mov_b32(s[0], 42),
s_nop(0), # nop with simm16=0
s_mov_b32(s[1], 100),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 42)
self.assertEqual(st.sgpr[1], 100)
def test_s_nop_with_count(self):
"""S_NOP with count parameter executes multiple nops."""
# S_NOP with immediate 3 should execute 4 nops (0:3 inclusive)
instructions = [
s_mov_b32(s[0], 1),
s_nop(3), # nop with simm16=3 -> 4 iterations
s_add_u32(s[0], s[0], 1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 2)
class TestNullRegister(unittest.TestCase):
"""Tests for NULL register (124) behavior - writes should be discarded, reads return 0."""
def test_s_mov_b32_from_null(self):
"""S_MOV_B32 from NULL should read as 0."""
instructions = [
s_mov_b32(s[0], 0xDEADBEEF), # Set s[0] to sentinel
s_mov_b32(s[0], NULL), # Read from NULL - should be 0
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 0)
def test_s_add_u32_with_null_src(self):
"""S_ADD_U32 with NULL as source should use 0."""
instructions = [
s_mov_b32(s[0], 100),
s_add_u32(s[1], s[0], NULL), # 100 + 0 = 100
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 100)
def test_s_mov_b32_to_null(self):
"""S_MOV_B32 to NULL (sdst=124) should discard the write."""
instructions = [
s_mov_b32(s[0], 0xDEADBEEF), # Set s[0] to sentinel
s_mov_b32(NULL, 42), # Write to NULL - should be discarded
# s[0] should still be 0xDEADBEEF since NULL write doesn't affect it
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 0xDEADBEEF)
def test_s_add_u32_to_null(self):
"""S_ADD_U32 with sdst=NULL should discard result but still set SCC."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_mov_b32(s[1], 1),
s_add_u32(NULL, s[0], s[1]), # overflow, write to NULL
s_cselect_b32(s[2], 1, 0), # capture SCC
]
st = run_program(instructions, n_lanes=1)
# SCC should still be set from overflow even though result was discarded
self.assertEqual(st.sgpr[2], 1)
self.assertEqual(st.scc, 1)
def test_s_and_b32_to_null(self):
"""S_AND_B32 with sdst=NULL should discard result but still set SCC."""
instructions = [
s_mov_b32(s[0], 0xFF00FF00),
s_mov_b32(s[1], 0x0F0F0F0F),
s_and_b32(NULL, s[0], s[1]), # result=0x0F000F00, non-zero so SCC=1
s_cselect_b32(s[2], 1, 0), # capture SCC
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 1) # SCC=1 because result was non-zero
self.assertEqual(st.scc, 1)
def test_s_or_b32_to_null_zero_result(self):
"""S_OR_B32 with sdst=NULL and zero result should set SCC=0."""
instructions = [
s_mov_b32(s[0], 0),
s_mov_b32(s[1], 0),
s_or_b32(NULL, s[0], s[1]), # result=0, so SCC=0
s_cselect_b32(s[2], 1, 0), # capture SCC
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0) # SCC=0 because result was zero
self.assertEqual(st.scc, 0)
class Test64BitSOP1InlineConstants(unittest.TestCase):
"""Tests for 64-bit SOP1 instructions with inline constants.
Regression tests for bug where rsrc_dyn didn't properly handle 64-bit
inline constants, incorrectly duplicating lo bits to hi instead of
zero/sign-extending.
"""
def test_s_mov_b64_inline_0(self):
"""S_MOV_B64 with inline constant 0."""
instructions = [
s_mov_b64(s[0:1], 0),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_inline_16(self):
"""S_MOV_B64 with inline constant 16 should set lo=16, hi=0."""
instructions = [
s_mov_b64(s[0:1], 16),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_inline_64(self):
"""S_MOV_B64 with inline constant 64 (max positive)."""
instructions = [
s_mov_b64(s[0:1], 64),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 64)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_inline_neg1(self):
"""S_MOV_B64 with inline constant -1 should sign-extend."""
instructions = [
s_mov_b64(s[0:1], -1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFFFF)
self.assertEqual(st.vgpr[0][1], 0xFFFFFFFF)
def test_s_mov_b64_inline_neg16(self):
"""S_MOV_B64 with inline constant -16 should sign-extend."""
instructions = [
s_mov_b64(s[0:1], -16),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFFF0)
self.assertEqual(st.vgpr[0][1], 0xFFFFFFFF)
def test_s_mov_b64_float_const_1_0(self):
"""S_MOV_B64 with float inline constant 1.0 - casts F32 to F64."""
instructions = [
s_mov_b64(s[0:1], 1.0), # inline constant 242 (1.0f)
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
# Hardware casts F32 to F64: 1.0f64 = 0x3FF0000000000000
self.assertEqual(st.vgpr[0][0], 0x00000000) # lo
self.assertEqual(st.vgpr[0][1], 0x3FF00000) # hi
def test_s_or_b64_inline_constant(self):
"""S_OR_B64 with 64-bit inline constant."""
instructions = [
s_mov_b64(s[0:1], 0),
s_or_b64(s[2:3], s[0:1], 16),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_and_b64_inline_constant(self):
"""S_AND_B64 with 64-bit inline constant."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_mov_b32(s[1], 0xFFFFFFFF),
s_and_b64(s[2:3], s[0:1], 16),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.vgpr[0][1], 0)
class Test64BitSOPLiterals(unittest.TestCase):
"""Tests for 64-bit SOP instructions with 32-bit literals.
Tests the behavior when a 64-bit SOP instruction uses a 32-bit literal
(offset 255 in instruction encoding). The literal is zero-extended to 64 bits.
"""
def test_s_mov_b64_literal(self):
"""S_MOV_B64 with 32-bit literal value - zero-extended to 64 bits."""
instructions = [
s_mov_b64(s[0:1], 0x12345678), # literal > 64, uses literal encoding
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x12345678)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_or_b64_literal(self):
"""S_OR_B64 with 32-bit literal value - zero-extended to 64 bits."""
instructions = [
s_mov_b64(s[0:1], 0),
s_or_b64(s[2:3], s[0:1], 0x12345678), # literal
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x12345678)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_and_b64_literal(self):
"""S_AND_B64 with 32-bit literal value - zero-extended to 64 bits."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_mov_b32(s[1], 0xFFFFFFFF),
s_and_b64(s[2:3], s[0:1], 0x12345678), # literal
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x12345678)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_literal_negative(self):
"""S_MOV_B64 with 0xFFFFFFFF literal - zero-extended (not sign-extended)."""
instructions = [
s_mov_b64(s[0:1], 0xFFFFFFFF), # -1 as 32-bit, but zero-extended to 64-bit
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFFFF)
self.assertEqual(st.vgpr[0][1], 0) # zero-extended, not sign-extended
def test_s_mov_b64_literal_high_bit(self):
"""S_MOV_B64 with 0x80000000 literal - zero-extended (not sign-extended)."""
instructions = [
s_mov_b64(s[0:1], 0x80000000), # high bit set, but zero-extended
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x80000000)
self.assertEqual(st.vgpr[0][1], 0) # zero-extended, not sign-extended
if __name__ == '__main__':
unittest.main()
@@ -5,7 +5,7 @@ Includes: v_mov_b32, v_cvt_*, v_sin_f32, v_rcp_f32, v_exp_f32, v_rndne_f32,
v_readfirstlane_b32
"""
import unittest
from test.amd.hw.helpers import *
from extra.assembly.amd.test.hw.helpers import *
class TestMov(unittest.TestCase):
"""Tests for V_MOV_B32."""
@@ -255,6 +255,7 @@ class TestF16Conversions(unittest.TestCase):
def test_v_cvt_f16_f32_small(self):
"""V_CVT_F16_F32 converts small f32 value."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
instructions = [
v_mov_b32_e32(v[0], 0.5),
v_cvt_f16_f32_e32(v[1], v[0]),
@@ -292,6 +293,7 @@ class TestF16Conversions(unittest.TestCase):
def test_v_cvt_f16_f32_reads_full_32bit_source(self):
"""V_CVT_F16_F32 must read full 32-bit f32 source."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], 0x3fc00000), # f32 1.5
v_mov_b32_e32(v[0], s[0]),
@@ -300,7 +302,7 @@ class TestF16Conversions(unittest.TestCase):
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo_bits = result & 0xffff
self.assertEqual(lo_bits, 0x3e00, f"Expected f16(1.5)=0x3e00, got 0x{lo_bits:04x} ({f16(lo_bits)})")
self.assertEqual(lo_bits, 0x3e00, f"Expected f16(1.5)=0x3e00, got 0x{lo_bits:04x} ({_f16(lo_bits)})")
def test_v_cvt_i16_f16_zero(self):
"""V_CVT_I16_F16 converts f16 zero to i16 zero."""
@@ -373,6 +375,7 @@ class TestF64Conversions(unittest.TestCase):
def test_v_cvt_f64_f32_pi(self):
"""V_CVT_F64_F32 converts f32 pi to f64."""
import math
instructions = [
s_mov_b32(s[0], f2i(3.14159265)),
v_mov_b32_e32(v[0], s[0]),
@@ -693,6 +696,7 @@ class TestCvtF16Modifiers(unittest.TestCase):
def test_v_cvt_f32_f16_abs_negative(self):
"""V_CVT_F32_F16 with |abs| on negative value."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_neg1 = f32_to_f16(-1.0) # 0xbc00
instructions = [
s_mov_b32(s[0], f16_neg1),
@@ -705,6 +709,7 @@ class TestCvtF16Modifiers(unittest.TestCase):
def test_v_cvt_f32_f16_abs_positive(self):
"""V_CVT_F32_F16 with |abs| on positive value (should stay positive)."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_2 = f32_to_f16(2.0) # 0x4000
instructions = [
s_mov_b32(s[0], f16_2),
@@ -717,6 +722,7 @@ class TestCvtF16Modifiers(unittest.TestCase):
def test_v_cvt_f32_f16_neg_positive(self):
"""V_CVT_F32_F16 with neg on positive value."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_2 = f32_to_f16(2.0) # 0x4000
instructions = [
s_mov_b32(s[0], f16_2),
@@ -729,6 +735,7 @@ class TestCvtF16Modifiers(unittest.TestCase):
def test_v_cvt_f32_f16_neg_negative(self):
"""V_CVT_F32_F16 with neg on negative value (double negative)."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_neg2 = f32_to_f16(-2.0) # 0xc000
instructions = [
s_mov_b32(s[0], f16_neg2),
@@ -741,6 +748,7 @@ class TestCvtF16Modifiers(unittest.TestCase):
def test_v_cvt_f16_f32_then_pack_for_wmma(self):
"""CVT F32->F16 followed by pack (common WMMA pattern)."""
from extra.assembly.amd.test.hw.helpers import _f16
f32_val = 3.5
instructions = [
s_mov_b32(s[0], f2i(f32_val)),
@@ -749,8 +757,8 @@ class TestCvtF16Modifiers(unittest.TestCase):
v_pack_b32_f16(v[2], v[1], v[1]), # Pack same value
]
st = run_program(instructions, n_lanes=1)
lo = f16(st.vgpr[0][2] & 0xffff)
hi = f16((st.vgpr[0][2] >> 16) & 0xffff)
lo = _f16(st.vgpr[0][2] & 0xffff)
hi = _f16((st.vgpr[0][2] >> 16) & 0xffff)
self.assertAlmostEqual(lo, f32_val, places=1)
self.assertAlmostEqual(hi, f32_val, places=1)
@@ -796,6 +804,7 @@ class TestConversionRounding(unittest.TestCase):
def test_f16_to_f32_precision(self):
"""F16 to F32 conversion precision."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_val = f32_to_f16(1.5)
instructions = [
s_mov_b32(s[0], f16_val),
@@ -807,6 +816,7 @@ class TestConversionRounding(unittest.TestCase):
def test_f16_denormal_to_f32(self):
"""F16 denormal converts to small positive f32."""
from extra.assembly.amd.test.hw.helpers import _f16
f16_denorm = 0x0001 # Smallest positive f16 denormal
instructions = [
v_mov_b32_e32(v[0], f16_denorm),
@@ -1502,82 +1512,5 @@ class TestReciprocalF16(unittest.TestCase):
self.assertAlmostEqual(result, 0.25, places=2, msg="1/4.0 should be 0.25")
class TestCvtNormF16(unittest.TestCase):
"""Tests for V_CVT_NORM_I16_F16 and V_CVT_NORM_U16_F16."""
def test_cvt_norm_i16_f16_positive(self):
"""V_CVT_NORM_I16_F16: f16 1.0 -> i16 max (32767)."""
instructions = [
s_mov_b32(s[0], f32_to_f16(1.0)),
v_mov_b32_e32(v[0], s[0]),
v_cvt_norm_i16_f16_e32(v[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 32767)
def test_cvt_norm_i16_f16_negative(self):
"""V_CVT_NORM_I16_F16: f16 -1.0 -> i16 -32767 (0x8001)."""
instructions = [
s_mov_b32(s[0], f32_to_f16(-1.0)),
v_mov_b32_e32(v[0], s[0]),
v_cvt_norm_i16_f16_e32(v[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 0x8001) # -32767, hardware uses symmetric range
def test_cvt_norm_i16_f16_zero(self):
"""V_CVT_NORM_I16_F16: f16 0.0 -> i16 0."""
instructions = [
v_mov_b32_e32(v[0], 0),
v_cvt_norm_i16_f16_e32(v[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 0)
def test_cvt_norm_u16_f16_one(self):
"""V_CVT_NORM_U16_F16: f16 1.0 -> u16 max (65535)."""
instructions = [
s_mov_b32(s[0], f32_to_f16(1.0)),
v_mov_b32_e32(v[0], s[0]),
v_cvt_norm_u16_f16_e32(v[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 65535)
def test_cvt_norm_u16_f16_half(self):
"""V_CVT_NORM_U16_F16: f16 0.5 -> u16 ~32768."""
instructions = [
s_mov_b32(s[0], f32_to_f16(0.5)),
v_mov_b32_e32(v[0], s[0]),
v_cvt_norm_u16_f16_e32(v[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertAlmostEqual(result, 32768, delta=1)
class TestPermlane64(unittest.TestCase):
"""Tests for V_PERMLANE64_B32 instruction (wave64 cross-half swap)."""
def test_v_permlane64_b32_is_nop_in_wave32(self):
"""V_PERMLANE64_B32 is a NOP in wave32 mode.
Per AMD pcode: "if WAVE32 then s_nop(...) else ... endif"
The emulator runs in wave32 mode, so this instruction should not modify registers.
"""
instructions = [
v_mov_b32_e32(v[0], 0xCAFEBABE), # source
v_mov_b32_e32(v[1], 0x12345678), # dest (should be preserved)
v_permlane64_b32_e32(v[1], v[0]), # NOP in wave32
]
st = run_program(instructions, n_lanes=1)
# Dest register should be unchanged (NOP behavior in wave32)
self.assertEqual(st.vgpr[0][1], 0x12345678)
if __name__ == '__main__':
unittest.main()
@@ -5,7 +5,7 @@ Includes: v_add_f32, v_mul_f32, v_and_b32, v_or_b32, v_xor_b32,
v_add_nc_u32, v_cndmask_b32, v_add_f16, v_mul_f16
"""
import unittest
from test.amd.hw.helpers import *
from extra.assembly.amd.test.hw.helpers import *
class TestBasicArithmetic(unittest.TestCase):
"""Tests for basic arithmetic VOP2 instructions."""
@@ -237,32 +237,6 @@ class TestF16Ops(unittest.TestCase):
# 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_max_f16_basic(self):
"""V_MAX_F16 returns the maximum of 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_max_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4000, f"Expected 0x4000 (f16 2.0), got 0x{result:04x}")
def test_v_min_f16_basic(self):
"""V_MIN_F16 returns the minimum of 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_min_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x3c00, f"Expected 0x3c00 (f16 1.0), got 0x{result:04x}")
def test_v_fmaak_f16_basic(self):
"""V_FMAAK_F16: d = a * b + K."""
instructions = [
@@ -836,81 +810,6 @@ class TestCarryOps(unittest.TestCase):
self.assertEqual(st.vgpr[0][2], 0) # Overflowed to 0
self.assertEqual(st.vcc, 1) # Carry out
def test_v_add_co_ci_u32_clears_carry(self):
"""V_ADD_CO_CI_U32: VCC must be updated even when no carry is generated.
This tests the case where VCC=1 going in (carry-in consumed) but the addition
does not overflow, so VCC must be cleared to 0.
Regression test for: VCC not being written by v_add_co_ci_u32_e32.
"""
instructions = [
s_mov_b32(VCC_LO, 1), # VCC = 1 (carry in)
v_mov_b32_e32(v[0], 1), # S0 = 1
v_mov_b32_e32(v[1], 1), # S1 = 1
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 1 = 3 (no overflow)
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 3) # 1 + 1 + 1 = 3
self.assertEqual(st.vcc, 0) # No carry out - VCC must be cleared
def test_v_add_co_ci_u32_multilane_clears_vcc(self):
"""V_ADD_CO_CI_U32 with multiple lanes: VCC bits must be updated per-lane.
When VCC has multiple bits set (one per active lane), and the addition doesn't
overflow for any lane, all VCC bits must be cleared.
Regression test for: VCC not being written by v_add_co_ci_u32_e32 in multi-lane case.
"""
instructions = [
s_mov_b32(VCC_LO, 0b11), # VCC = 0b11 (lanes 0,1 have carry-in)
v_mov_b32_e32(v[0], 1), # S0 = 1 for all lanes
v_mov_b32_e32(v[1], 1), # S1 = 1 for all lanes
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 1 = 3 (no overflow)
]
st = run_program(instructions, n_lanes=2)
self.assertEqual(st.vgpr[0][2], 3) # lane 0: 1 + 1 + 1 = 3
self.assertEqual(st.vgpr[1][2], 3) # lane 1: 1 + 1 + 1 = 3
self.assertEqual(st.vcc, 0) # No carry out for any lane - all VCC bits must be cleared
def test_v_add_co_ci_u32_preserves_inactive_vcc_bits(self):
"""V_ADD_CO_CI_U32: VCC carry-out overwrites entire VCC register.
VOP2 carry instructions write ALL VCC bits based on carry-out, clearing
bits for lanes that don't overflow regardless of EXEC mask.
Note: This differs from VOPC which only writes active lane bits.
"""
instructions = [
s_mov_b32(VCC_LO, 0x00010000), # VCC bit 16 set
v_mov_b32_e32(v[0], 1), # S0 = 1
v_mov_b32_e32(v[1], 1), # S1 = 1
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 0 = 2 (no carry)
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.vgpr[0][2], 2) # lane 0: 1 + 1 + 0 = 2
# VCC should be completely cleared (all lanes have no carry-out)
self.assertEqual(st.vcc, 0)
def test_v_add_co_ci_u32_all_lanes_same_result(self):
"""V_ADD_CO_CI_U32: all active lanes should produce the same result.
When the same constant inputs are used across all lanes, each lane should
compute the same result and write to its own VGPR slot.
Regression test for: VGPR writes not happening for all lanes.
"""
instructions = [
s_mov_b32(VCC_LO, 0), # No carry-in
v_mov_b32_e32(v[0], 3), # inline constant 3
v_mov_b32_e32(v[1], 5), # value 5
v_add_co_ci_u32_e32(v[1], 3, v[1]), # v[1] = 3 + v[1] + 0 = 3 + 5 = 8
]
st = run_program(instructions, n_lanes=4)
# All 4 lanes should have v[1] = 8
for lane in range(4):
self.assertEqual(st.vgpr[lane][1], 8, f"lane {lane} should have v[1]=8")
def test_v_sub_co_ci_u32_no_borrow(self):
"""V_SUB_CO_CI_U32: D0 = S0 - S1 - VCC_IN, when VCC_IN=0."""
instructions = [
@@ -961,23 +860,6 @@ class TestCarryOps(unittest.TestCase):
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.sgpr[10], 0) # No carry out
def test_v_add_co_ci_u32_vop3sd_null_sdst(self):
"""VOP3SD V_ADD_CO_CI_U32 with sdst=NULL: carry output is discarded.
When sdst=NULL (register 124), the carry-out should NOT be written anywhere.
We verify this by checking that VCC (which we set to a sentinel value) is unchanged.
"""
instructions = [
s_mov_b32(VCC_LO, 0xDEADBEEF), # Sentinel value in VCC
s_mov_b32(s[6], 0), # carry-in = 0
# VOP3SD with NULL sdst: carry-out should be discarded
# Uses 0xFFFFFFFF + 1 + 0 = 0 with carry-out=1, but carry should not be written
v_add_co_ci_u32(v[0], NULL, 0xFFFFFFFF, 1, s[6]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0) # 0xFFFFFFFF + 1 + 0 = 0 (overflow)
self.assertEqual(st.vcc, 0xDEADBEEF) # VCC unchanged - carry was discarded
if __name__ == '__main__':
unittest.main()
File diff suppressed because it is too large Load Diff
@@ -3,7 +3,7 @@
Includes: v_pk_add_f16, v_pk_mul_f16, v_pk_fma_f16, v_pack_b32_f16, v_wmma_*, v_dot2_*
"""
import unittest
from test.amd.hw.helpers import *
from extra.assembly.amd.test.hw.helpers import *
class TestPackInstructions(unittest.TestCase):
"""Tests for pack instructions."""
@@ -149,6 +149,7 @@ class TestFmaMix(unittest.TestCase):
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.test.hw.helpers import f32_to_f16
f16_2 = f32_to_f16(2.0)
instructions = [
s_mov_b32(s[0], f2i(1.0)),
@@ -165,6 +166,7 @@ class TestFmaMix(unittest.TestCase):
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.test.hw.helpers import f32_to_f16
f16_2 = f32_to_f16(2.0)
val = (f16_2 << 16) | 0
instructions = [
@@ -197,6 +199,7 @@ class TestFmaMix(unittest.TestCase):
def test_v_fma_mix_f32_with_abs_f16_src2_lo(self):
"""V_FMA_MIX_F32 with abs modifier on f16 src2 (lo half). Regression test for sin(1.0) bug."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_neg1 = f32_to_f16(-1.0) # 0xbc00
instructions = [
s_mov_b32(s[0], f2i(0.0)), # src0 = 0.0 (f32)
@@ -214,6 +217,7 @@ class TestFmaMix(unittest.TestCase):
def test_v_fma_mix_f32_with_neg_f16_src2_lo(self):
"""V_FMA_MIX_F32 with neg modifier on f16 src2 (lo half)."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_1 = f32_to_f16(1.0) # 0x3c00
instructions = [
s_mov_b32(s[0], f2i(0.0)), # src0 = 0.0 (f32)
@@ -231,6 +235,7 @@ class TestFmaMix(unittest.TestCase):
def test_v_fma_mix_f32_with_abs_f16_src2_hi(self):
"""V_FMA_MIX_F32 with abs modifier on f16 src2 (hi half)."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_neg1 = f32_to_f16(-1.0) # 0xbc00
val = (f16_neg1 << 16) | 0 # -1.0 in hi, 0 in lo
instructions = [
@@ -249,6 +254,7 @@ class TestFmaMix(unittest.TestCase):
def test_v_fma_mixlo_f16(self):
"""V_FMA_MIXLO_F16 writes to low 16 bits of destination."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], f2i(2.0)),
v_mov_b32_e32(v[0], s[0]),
@@ -261,13 +267,14 @@ class TestFmaMix(unittest.TestCase):
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)
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.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], f2i(1.0)),
v_mov_b32_e32(v[0], s[0]),
@@ -279,12 +286,13 @@ class TestFmaMix(unittest.TestCase):
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)
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.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], 0x3f800000), # f32 1.0
v_mov_b32_e32(v[3], s[0]),
@@ -297,7 +305,7 @@ class TestFmaMix(unittest.TestCase):
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)
lo = _f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(lo, -3.14159, delta=0.01)
@@ -306,6 +314,7 @@ class TestVOP3P(unittest.TestCase):
def test_v_pk_add_f16_basic(self):
"""V_PK_ADD_F16 adds two packed f16 values."""
from extra.assembly.amd.test.hw.helpers 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
@@ -315,13 +324,14 @@ class TestVOP3P(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
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.test.hw.helpers 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
@@ -331,13 +341,14 @@ class TestVOP3P(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
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.test.hw.helpers 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
@@ -349,8 +360,8 @@ class TestVOP3P(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
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
@@ -359,6 +370,7 @@ class TestVOP3P(unittest.TestCase):
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.test.hw.helpers 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]),
@@ -366,8 +378,8 @@ class TestVOP3P(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
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)
@@ -376,6 +388,7 @@ class TestVOP3P(unittest.TestCase):
"""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.test.hw.helpers 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 = [
@@ -385,124 +398,18 @@ class TestVOP3P(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 6.0, places=1)
self.assertAlmostEqual(hi, 0.0, places=1)
def test_v_pk_add_u16_float_inline_const_opsel(self):
"""V_PK_ADD_U16 with float inline constant 2.0
Regression test: for integer packed ops, do not perform the f32->f16 conversion.
"""
# src1 = inline float constant 2.0
instructions = [
s_mov_b32(s[0], 0x00030005), # packed u16: hi=3, lo=5
v_mov_b32_e32(v[0], s[0]),
v_pk_add_u16(v[1], v[0], SrcEnum.POS_TWO, opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo = result & 0xffff
hi = (result >> 16) & 0xffff
# lo = 5 + 0x0000 = 0x0005, hi = 3 + 0x4000 = 0x4003
self.assertEqual(lo, 0x0005, f"lo: expected 0x0005, got 0x{lo:04x}")
self.assertEqual(hi, 0x4003, f"hi: expected 0x4003, got 0x{hi:04x}")
class TestWMMAF16(unittest.TestCase):
"""Tests for WMMA F16 output variant (V_WMMA_F16_16X16X16_F16).
Note: RDNA3 WMMA F16 uses 8 VGPRs for accumulator/output (same as F32 variant),
but values are packed as f16. This differs from RDNA4 which uses 4 VGPRs.
"""
def test_v_wmma_f16_16x16x16_f16_all_ones(self):
"""V_WMMA_F16_16X16X16_F16 with all ones produces 16.0 in f16."""
instructions: list[Inst] = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
# Initialize A matrix in v[16:23] (8 regs)
for i in range(16, 24):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[24:31] (8 regs)
for i in range(24, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[0:7] to zero (8 regs for RDNA3 WMMA F16)
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], 0))
# WMMA F16: D = A @ B + C
instructions.append(v_wmma_f16_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
# Result should be 16.0 in f16, stored in lo 16 bits of each VGPR (hi bits are 0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
lo = f16(result & 0xffff)
self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
def test_v_wmma_f16_16x16x16_f16_with_accumulator(self):
"""V_WMMA_F16_16X16X16_F16 with non-zero accumulator."""
instructions: list[Inst] = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
instructions.append(s_mov_b32(s[1], 0x4500)) # f16 5.0 in lo bits only
# Initialize A matrix in v[16:23] (8 regs)
for i in range(16, 24):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[24:31] (8 regs)
for i in range(24, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[0:7] to 5.0 in lo bits (8 regs for RDNA3 WMMA F16)
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], s[1]))
# WMMA F16: D = A @ B + C
instructions.append(v_wmma_f16_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
# Result should be 16.0 + 5.0 = 21.0 in f16, stored in lo 16 bits (hi bits are 0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
lo = f16(result & 0xffff)
self.assertAlmostEqual(lo, 21.0, places=0, msg=f"v[{reg}] lane {lane}: expected 21.0, got {lo}")
self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
def test_v_wmma_f16_16x16x16_f16_high_registers(self):
"""V_WMMA_F16_16X16X16_F16 with high register indices.
Regression test: WMMA was using static register indices instead of dynamic.
This test uses v[64:71] for A, v[80:87] for B, v[96:103] for C/D.
"""
instructions: list[Inst] = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
# Initialize A matrix in v[64:71] (8 regs)
for i in range(64, 72):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[80:87] (8 regs)
for i in range(80, 88):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[96:103] to zero (8 regs for RDNA3 WMMA F16)
for i in range(96, 104):
instructions.append(v_mov_b32_e32(v[i], 0))
# WMMA F16: D = A @ B + C, result in v[96:103]
instructions.append(v_wmma_f16_16x16x16_f16(v[96:103], v[64:71], v[80:87], v[96:103]))
# Copy results to v[0:7] for checking
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], v[96+i]))
st = run_program(instructions, n_lanes=32)
# Result should be 16.0 in f16, stored in lo 16 bits (hi bits are 0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
lo = f16(result & 0xffff)
self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
class TestWMMA(unittest.TestCase):
"""Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions with F32 output."""
"""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: list[Inst] = []
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
for i in range(16, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
@@ -518,7 +425,7 @@ class TestWMMA(unittest.TestCase):
def test_v_wmma_f32_16x16x16_f16_with_accumulator(self):
"""V_WMMA_F32_16X16X16_F16 with non-zero accumulator."""
instructions: list[Inst] = []
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00))
instructions.append(s_mov_b32(s[1], f2i(5.0)))
for i in range(16, 32):
@@ -533,75 +440,6 @@ class TestWMMA(unittest.TestCase):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
def test_v_wmma_f32_16x16x16_f16_high_registers(self):
"""V_WMMA_F32_16X16X16_F16 with high register indices.
Regression test: WMMA was using static register indices instead of dynamic,
causing incorrect results when registers weren't at the default positions.
This test uses v[64:71] for A, v[80:87] for B, v[96:103] for C/D.
"""
instructions: list[Inst] = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
# Initialize A matrix in v[64:71]
for i in range(64, 72):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[80:87]
for i in range(80, 88):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[96:103] to zero
for i in range(96, 104):
instructions.append(v_mov_b32_e32(v[i], 0))
# WMMA: D = A @ B + C, result in v[96:103]
instructions.append(v_wmma_f32_16x16x16_f16(v[96:103], v[64:71], v[80:87], v[96:103]))
# Copy results to v[0:7] for checking
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], v[96+i]))
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)}")
class TestWMMABF16(unittest.TestCase):
"""Tests for WMMA BF16 instructions."""
def test_v_wmma_f32_16x16x16_bf16_all_ones(self):
"""V_WMMA_F32_16X16X16_BF16 with all ones produces 16.0."""
instructions: list[Inst] = []
# BF16 1.0 = 0x3f80, packed = 0x3f803f80
instructions.append(s_mov_b32(s[0], 0x3f803f80))
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_bf16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
expected = f2i(16.0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
def test_v_wmma_f32_16x16x16_bf16_with_accumulator(self):
"""V_WMMA_F32_16X16X16_BF16 with non-zero accumulator."""
instructions: list[Inst] = []
# BF16 1.0 = 0x3f80, packed = 0x3f803f80
instructions.append(s_mov_b32(s[0], 0x3f803f80))
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_bf16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
expected = f2i(21.0) # 16 + 5
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
class TestSpecialOps(unittest.TestCase):
"""Tests for special operations (SAD, PERM, DOT2)."""
@@ -715,6 +553,7 @@ class TestPackedMixedSigns(unittest.TestCase):
def test_pk_add_f16_mixed_signs(self):
"""V_PK_ADD_F16 with mixed positive/negative values."""
from extra.assembly.amd.test.hw.helpers 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
@@ -724,13 +563,14 @@ class TestPackedMixedSigns(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
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.test.hw.helpers 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
@@ -743,277 +583,5 @@ class TestPackedMixedSigns(unittest.TestCase):
self.assertEqual(result, 0x00000000, "2.0 * 0.0 should be 0.0")
class TestDot2F32F16(unittest.TestCase):
"""Tests for V_DOT2_F32_F16 - dot product of f16 pairs producing f32."""
def test_v_dot2_f32_f16_basic(self):
"""V_DOT2_F32_F16: dot product of two packed f16 pairs -> f32."""
# src0 = {hi=2.0, lo=1.0}, src1 = {hi=4.0, lo=3.0}
# result = 1.0*3.0 + 2.0*4.0 + 0 = 3 + 8 = 11.0
src0 = (f32_to_f16(2.0) << 16) | f32_to_f16(1.0)
src1 = (f32_to_f16(4.0) << 16) | f32_to_f16(3.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
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_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 11.0, places=2)
def test_v_dot2_f32_f16_with_accumulator(self):
"""V_DOT2_F32_F16 with non-zero f32 accumulator."""
# src0 = {hi=1.0, lo=1.0}, src1 = {hi=1.0, lo=1.0}, acc = 5.0
# result = 1.0*1.0 + 1.0*1.0 + 5.0 = 7.0
src0 = (f32_to_f16(1.0) << 16) | f32_to_f16(1.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], f2i(5.0)),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]), # same as src0
v_mov_b32_e32(v[2], s[1]),
v_dot2_f32_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 7.0, places=2)
def test_v_dot2_f32_f16_negative_values(self):
"""V_DOT2_F32_F16 with negative f16 values."""
# src0 = {hi=-2.0, lo=3.0}, src1 = {hi=1.0, lo=2.0}
# result = 3.0*2.0 + (-2.0)*1.0 + 0 = 6 - 2 = 4.0
# NOTE: Hardware DOT2 may have up to 1 ULP difference due to internal implementation
src0 = (f32_to_f16(-2.0) << 16) | f32_to_f16(3.0)
src1 = (f32_to_f16(1.0) << 16) | f32_to_f16(2.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
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_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1, ulp_tolerance=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 4.0, places=2)
class TestDot2F16F16(unittest.TestCase):
"""Tests for V_DOT2_F16_F16 - dot product of f16 pairs producing f16."""
def test_v_dot2_f16_f16_basic(self):
"""V_DOT2_F16_F16: dot product of two packed f16 pairs -> f16."""
# src0 = {hi=2.0, lo=1.0}, src1 = {hi=3.0, lo=2.0}
# result = 1.0*2.0 + 2.0*3.0 + 0 = 2 + 6 = 8.0 (f16)
src0 = (f32_to_f16(2.0) << 16) | f32_to_f16(1.0)
src1 = (f32_to_f16(3.0) << 16) | f32_to_f16(2.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
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_f16_f16(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, 8.0, places=1)
def test_v_dot2_f16_f16_with_accumulator(self):
"""V_DOT2_F16_F16 with non-zero f16 accumulator."""
# src0 = {hi=1.0, lo=1.0}, src1 = {hi=1.0, lo=1.0}, acc = 3.0 (f16)
# result = 1.0*1.0 + 1.0*1.0 + 3.0 = 5.0 (f16)
src0 = (f32_to_f16(1.0) << 16) | f32_to_f16(1.0)
acc = f32_to_f16(3.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[2], acc),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]), # same as src0
v_mov_b32_e32(v[2], s[2]),
v_dot2_f16_f16(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, 5.0, places=1)
class TestSignedDotProducts(unittest.TestCase):
"""Tests for V_DOT4_I32_IU8 and V_DOT8_I32_IU4 with signed inputs."""
def test_v_dot4_i32_iu8_signed_both(self):
"""V_DOT4_I32_IU8 with both inputs signed (neg=0b011)."""
# src0 = {-1, -2, 3, 4} as i8 = {0xff, 0xfe, 0x03, 0x04}
# src1 = {1, 1, 1, 1} as i8
# result = (-1)*1 + (-2)*1 + 3*1 + 4*1 = -1 - 2 + 3 + 4 = 4
src0 = (0xff << 24) | (0xfe << 16) | (0x03 << 8) | 0x04 # -1, -2, 3, 4
src1 = 0x01010101 # 1, 1, 1, 1
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b011), # both signed
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
# Result is i32, interpret as signed
if result >= 0x80000000:
result = result - 0x100000000
self.assertEqual(result, 4)
def test_v_dot4_i32_iu8_src0_signed(self):
"""V_DOT4_I32_IU8 with only src0 signed (neg=0b001)."""
# src0 = {-1, -1, -1, -1} as i8 = {0xff, 0xff, 0xff, 0xff}
# src1 = {2, 2, 2, 2} as u8
# result = (-1)*2 + (-1)*2 + (-1)*2 + (-1)*2 = -8
src0 = 0xffffffff # -1, -1, -1, -1 (as i8)
src1 = 0x02020202 # 2, 2, 2, 2 (as u8)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b001), # src0 signed
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
if result >= 0x80000000:
result = result - 0x100000000
self.assertEqual(result, -8)
def test_v_dot4_i32_iu8_src1_signed(self):
"""V_DOT4_I32_IU8 with only src1 signed (neg=0b010)."""
# src0 = {2, 2, 2, 2} as u8
# src1 = {-1, -1, -1, -1} as i8 = {0xff, 0xff, 0xff, 0xff}
# result = 2*(-1) + 2*(-1) + 2*(-1) + 2*(-1) = -8
src0 = 0x02020202 # 2, 2, 2, 2 (as u8)
src1 = 0xffffffff # -1, -1, -1, -1 (as i8)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b010), # src1 signed
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
if result >= 0x80000000:
result = result - 0x100000000
self.assertEqual(result, -8)
def test_v_dot4_i32_iu8_unsigned_as_reference(self):
"""V_DOT4_I32_IU8 with both unsigned (neg=0) - same as V_DOT4_U32_U8."""
# src0 = {0xff, 0xff, 0xff, 0xff} = 255 each as u8
# src1 = {1, 1, 1, 1}
# result = 255*1 + 255*1 + 255*1 + 255*1 = 1020
src0 = 0xffffffff
src1 = 0x01010101
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0), # both unsigned
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][3], 1020)
def test_v_dot8_i32_iu4_signed_both(self):
"""V_DOT8_I32_IU4 with both inputs signed (neg=0b011)."""
# src0 = 8 nibbles: {-1, -2, 3, 4, -1, -2, 3, 4} as i4
# i4 -1 = 0xf, -2 = 0xe, 3 = 0x3, 4 = 0x4
# src0 = 0xfe34fe34
# src1 = {1, 1, 1, 1, 1, 1, 1, 1} as i4 = 0x11111111
# result = 2 * ((-1)*1 + (-2)*1 + 3*1 + 4*1) = 2 * 4 = 8
src0 = 0xfe34fe34
src1 = 0x11111111
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot8_i32_iu4(v[3], v[0], v[1], v[2], neg=0b011), # both signed
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
if result >= 0x80000000:
result = result - 0x100000000
self.assertEqual(result, 8)
def test_v_dot8_i32_iu4_all_negative(self):
"""V_DOT8_I32_IU4 with all negative signed values."""
# src0 = 8 nibbles all -1 (0xf) = 0xffffffff
# src1 = 8 nibbles all 1 = 0x11111111
# result = 8 * ((-1)*1) = -8
src0 = 0xffffffff # all -1 as i4
src1 = 0x11111111 # all 1
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot8_i32_iu4(v[3], v[0], v[1], v[2], neg=0b011), # both signed
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
if result >= 0x80000000:
result = result - 0x100000000
self.assertEqual(result, -8)
class TestPkMinMaxF16(unittest.TestCase):
"""Tests for V_PK_MIN_F16 and V_PK_MAX_F16."""
def test_v_pk_min_f16_basic(self):
"""V_PK_MIN_F16: packed min of two f16 pairs."""
# src0 = {hi=3.0, lo=1.0}, src1 = {hi=2.0, lo=4.0}
# result = {min(3,2)=2, min(1,4)=1}
src0 = (f32_to_f16(3.0) << 16) | f32_to_f16(1.0)
src1 = (f32_to_f16(2.0) << 16) | f32_to_f16(4.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_min_f16(v[2], v[0], v[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, 1.0, delta=0.01)
self.assertAlmostEqual(hi, 2.0, delta=0.01)
def test_v_pk_max_f16_basic(self):
"""V_PK_MAX_F16: packed max of two f16 pairs."""
# src0 = {hi=3.0, lo=1.0}, src1 = {hi=2.0, lo=4.0}
# result = {max(3,2)=3, max(1,4)=4}
src0 = (f32_to_f16(3.0) << 16) | f32_to_f16(1.0)
src1 = (f32_to_f16(2.0) << 16) | f32_to_f16(4.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_max_f16(v[2], v[0], v[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, delta=0.01)
self.assertAlmostEqual(hi, 3.0, delta=0.01)
if __name__ == '__main__':
unittest.main()
@@ -3,7 +3,7 @@
Includes: v_cmp_class_f32, v_cmp_class_f16, v_cmp_eq_*, v_cmp_lt_*, v_cmp_gt_*
"""
import unittest
from test.amd.hw.helpers import *
from extra.assembly.amd.test.hw.helpers import *
VCC = 106 # SGPR index for VCC_LO
@@ -731,111 +731,6 @@ class TestVCCBehavior(unittest.TestCase):
self.assertEqual(st.vcc >> 16, 0x0000, "Lanes 16-31 should be false")
class TestCmpNge(unittest.TestCase):
"""Tests for V_CMP_NGE (not-greater-or-equal) with NaN semantics.
NGE = !(a >= b). With NaN inputs:
- If either input is NaN, a >= b is false, so !(false) = true
- This differs from a < b which returns false for NaN inputs
"""
def test_v_cmp_nge_f32_normal_values(self):
"""v_cmp_nge_f32: basic comparison with normal floats."""
instructions = [
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], f2i(2.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= 2.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(1.0 >= 2.0) should be true")
def test_v_cmp_nge_f32_equal_values(self):
"""v_cmp_nge_f32: equal values should return false."""
instructions = [
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= 1.0) = !(true) = false
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "!(1.0 >= 1.0) should be false")
def test_v_cmp_nge_f32_greater_value(self):
"""v_cmp_nge_f32: greater value should return false."""
instructions = [
v_mov_b32_e32(v[0], f2i(2.0)),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(2.0 >= 1.0) = !(true) = false
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "!(2.0 >= 1.0) should be false")
def test_v_cmp_nge_f32_neg_inf(self):
"""v_cmp_nge_f32: -inf compared to normal value."""
neg_inf = 0xff800000 # -inf
instructions = [
s_mov_b32(s[0], neg_inf),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(-inf >= 1.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(-inf >= 1.0) should be true")
def test_v_cmp_nge_f32_clears_inactive_vcc_bits(self):
"""v_cmp_nge_f32 with partial EXEC clears inactive VCC bits (hardware behavior)."""
neg_inf = 0xff800000 # -inf
instructions = [
# Set VCC to all 1s first
s_mov_b32(VCC_LO, 0xFFFFFFFF),
# Set EXEC to only lane 0
s_mov_b32(EXEC_LO, 0x00000001),
# v0 = 1.0 for lane 0
v_mov_b32_e32(v[0], f2i(1.0)),
# Compare: !(-inf >= 1.0) = true for lane 0
v_cmp_nge_f32_e32(neg_inf, v[0]),
]
st = run_program(instructions, n_lanes=16)
# Hardware clears inactive lane bits, only active lane results remain
# Lane 0 result = 1 (true), lanes 1-15 = 0 (cleared)
self.assertEqual(st.vcc, 0x00000001, "VCC should only have active lane results")
def test_v_cmp_nge_f32_nan_src0(self):
"""v_cmp_nge_f32: NaN in src0 should return true (NaN >= x is false)."""
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], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(NaN >= 1.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(NaN >= 1.0) should be true")
def test_v_cmp_nge_f32_nan_src1(self):
"""v_cmp_nge_f32: NaN in src1 should return true (x >= NaN is false)."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], s[0]),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= NaN) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(1.0 >= NaN) should be true")
def test_v_cmp_nge_f32_both_nan(self):
"""v_cmp_nge_f32: both NaN should return true."""
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], s[0]),
v_cmp_nge_f32_e32(v[0], v[1]), # !(NaN >= NaN) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(NaN >= NaN) should be true")
class TestCmpxPartialWavefront(unittest.TestCase):
"""Tests for V_CMPX with partial wavefronts (fewer than 32 active lanes).
@@ -7,8 +7,9 @@ VOPD executes two operations simultaneously. Key behavior:
- Op Y can use ops 0-18 (includes ADD_NC_U32, LSHLREV, AND)
"""
import unittest
from test.amd.hw.helpers import run_program, v, v_mov_b32_e32
from tinygrad.runtime.autogen.amd.rdna3.ins import VOPD, VOPD_LIT, VOPDOp
from extra.assembly.amd.test.hw.helpers import run_program, run_program_emu, run_program_hw, compare_wave_states, \
v, s, v_mov_b32_e32, s_mov_b32
from extra.assembly.amd.autogen.rdna3.ins import VOPD, VOPD_LIT, VOPDOp
class TestVOPDBasic(unittest.TestCase):
"""Basic VOPD functionality tests."""
@@ -108,7 +109,7 @@ class TestVOPDLiterals(unittest.TestCase):
Tests that the 32-bit literal (SIMM32) is correctly passed to the instruction.
fma(2.0, 3.0, 10.0) = 2*3 + 10 = 16.0
"""
from test.amd.hw.helpers import f2i, i2f
from extra.assembly.amd.test.hw.helpers import f2i, i2f
instructions = [
v_mov_b32_e32(v[0], f2i(2.0)), # v[0] = 2.0
v_mov_b32_e32(v[1], f2i(3.0)), # v[1] = 3.0
@@ -126,7 +127,7 @@ class TestVOPDLiterals(unittest.TestCase):
Tests that the 32-bit literal (SIMM32) is correctly used as the multiplier.
fma(2.0, 5.0, 3.0) = 2*5 + 3 = 13.0
"""
from test.amd.hw.helpers import f2i, i2f
from extra.assembly.amd.test.hw.helpers import f2i, i2f
instructions = [
v_mov_b32_e32(v[0], f2i(2.0)), # v[0] = 2.0
v_mov_b32_e32(v[1], f2i(3.0)), # v[1] = 3.0
@@ -1,17 +1,17 @@
# Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
import unittest, ctypes
import unittest, ctypes, os
from dataclasses import dataclass
from pathlib import Path
from tinygrad import Device
from test.mockgpu.amd.emu import WaveState, _decode_at, WAVE_SIZE, VCC_LO, EXEC_LO, SCC
from tinygrad.renderer.amd import decode_inst
from test.amd.helpers import KernelInfo
import tinygrad
REMU_PATH = Path(tinygrad.__file__).parent.parent / "extra/remu/target/release/libremu.so"
if not REMU_PATH.exists(): REMU_PATH = Path(tinygrad.__file__).parent.parent / "extra/remu/target/release/libremu.dylib"
# 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"
def set_valid_mem_ranges(ranges): pass # emu2 doesn't need this
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)."""
@@ -85,71 +85,39 @@ class RustEmulator:
return snap.to_snapshot()
def free(self):
if self.ctx:
self.lib.wave_free(self.ctx)
self.ctx = None
if self.ctx: self.lib.wave_free(self.ctx); self.ctx = None
class PythonEmulator:
def __init__(self):
self.state: WaveState | None = None
self.program: dict[int, tuple] = {} # lazily populated: pc -> (name, fxn, globals)
self.vmem_buf = None
self.lds_buf = None
self.kernel_buf = None # Keep kernel bytes alive
self.lib_addr = 0 # Base address of kernel code
self.program: dict | None = None
def create(self, kernel: bytes, n_lanes: int):
import ctypes
from tinygrad.device import Buffer, BufferSpec
from tinygrad.dtype import dtypes
# Store kernel in a ctypes buffer so _decode_at can read from memory at actual PC address
self.kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
self.lib_addr = ctypes.addressof(self.kernel_buf)
self.program = {}
self.state = WaveState(n_lanes)
self.state.pc = self.lib_addr # Set PC to code base address
self.vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
self.lds_buf = Buffer('CPU', 65536 // 4, dtypes.uint32).ensure_allocated()
def _ensure_decoded(self, pc: int):
if pc not in self.program:
runner = _decode_at(pc, "rdna3")
self.program[pc] = (runner.p.function_name, runner._prg.fxn, runner.p.globals)
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:
import ctypes
assert self.state is not None
pc = self.state.pc
if pc == 0xFFFFFFFFFFFFFFFF: return -1
self._ensure_decoded(pc)
name, fxn, globals_list = self.program[pc]
buf_addrs = {0: self.state.sgpr_buf._buf.va_addr, 1: self.state.vgpr_buf._buf.va_addr, # type: ignore[union-attr]
2: self.vmem_buf._buf.va_addr, 3: self.lds_buf._buf.va_addr} # type: ignore[union-attr]
fxn(*[ctypes.c_uint64(buf_addrs[g]) for g in globals_list], ctypes.c_int32(0))
return -1 if self.state.pc == 0xFFFFFFFFFFFFFFFF else 0
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._write_sgpr(idx, val)
self.state.sgpr[idx] = val & 0xffffffff
def set_vgpr(self, lane: int, idx: int, val: int):
assert self.state is not None
self.state._write_vgpr(idx, lane, val)
self.state.vgpr[lane][idx] = val & 0xffffffff
def get_snapshot(self) -> StateSnapshot:
assert self.state is not None
sgpr = [self.state._read_sgpr(i) for i in range(128)]
vgpr = [[self.state._read_vgpr(reg, lane) for reg in range(256)] for lane in range(WAVE_SIZE)]
# Convert actual PC address to word offset for comparison with Rust emulator
pc_offset = (self.state.pc - self.lib_addr) // 4 if self.state.pc != 0xFFFFFFFFFFFFFFFF else 0xFFFFFFFFFFFFFFFF
return StateSnapshot(pc=pc_offset, scc=self.state._read_sgpr(SCC.offset), vcc=sgpr[VCC_LO.offset],
exec_mask=sgpr[EXEC_LO.offset], sgpr=sgpr, vgpr=vgpr)
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],
local_size: tuple[int, int, int], max_steps: int, debug: bool, trace_len: int,
kernel_idx: int = 0, max_workgroups: int = 8) -> tuple[bool, str, 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
lx, ly, lz = local_size
total_steps = 0
wg_count = 0
@@ -172,53 +140,28 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
emu.set_sgpr(13, gidx)
emu.set_sgpr(14, gidy)
emu.set_sgpr(15, gidz)
# Initialize v[0] with packed workitem IDs for each lane
for lane in range(n_lanes):
tid = lane
z, y, x = tid // (lx * ly), (tid // lx) % ly, tid % lx
emu.set_vgpr(lane, 0, (z << 20) | (y << 10) | x)
step = 0
trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
prev_sync_after = False # Track if previous instruction had known Rust bugs
try:
while step < max_steps:
rust_before = rust.get_snapshot()
python_before = python.get_snapshot()
pc_addr = python.lib_addr + python_before.pc * 4 # Convert word offset to actual address
python._ensure_decoded(pc_addr)
inst_hex_name = python.program[pc_addr][0]
# Decode the instruction to get mnemonic for sync_after checks
try:
# Format is mnemonic_hexbytes, e.g. v_exp_f32_e32_014b027e -> hex is 014b027e
parts = inst_hex_name.rsplit('_', 1)
inst_bytes_hex = parts[1] if len(parts) == 2 else ""
inst_bytes = bytes.fromhex(inst_bytes_hex) if inst_bytes_hex else b''
decoded = decode_inst(inst_bytes) if inst_bytes else None
inst_mnemonic = repr(decoded).split('(')[0] if decoded else ""
except Exception:
inst_mnemonic = ""
# For generic instructions, use function name for sync_after check
if not inst_mnemonic: inst_mnemonic = inst_hex_name
inst_str = inst_hex_name
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 or precision differences - sync Python to Rust after execution
# 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
# v_exp_f32/v_log_f32/v_ldexp_f32: precision differences in transcendental functions
# s_delay_alu: Rust handles differently
# v_add_co_ci_u32/v_sub_co_ci_u32/v_subrev_co_ci_u32: Rust preserves inactive VCC bits, but hardware clears all bits
sync_after = any(x in inst_mnemonic.lower() for x in ('v_div_scale', 'v_div_fixup', 'v_cvt_f16_f32', 's_add_i32', 's_sub_i32',
'v_exp_f32', 'v_log_f32', 'v_ldexp_f32', 's_delay_alu',
'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'))
# Skip comparison if previous instruction had known Rust bugs (states were synced but may still differ slightly)
diffs = rust_before.diff(python_before, n_lanes) if not prev_sync_after else []
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):
@@ -229,18 +172,16 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
python_diffs = pb.diff(next_pb, n_lanes, "->")
if rust_diffs: trace_lines.append(f" rust: {', '.join(rust_diffs[:5])}")
if python_diffs: trace_lines.append(f" python: {', '.join(python_diffs[:5])}")
elif rust_diffs: trace_lines.append(" python: (no changes)")
elif rust_diffs: trace_lines.append(f" python: (no changes)")
else:
# Last traced instruction - compare with current state
rust_diffs = rb.diff(rust_before, n_lanes, "->")
python_diffs = pb.diff(python_before, n_lanes, "->")
if rust_diffs: trace_lines.append(f" rust: {', '.join(rust_diffs[:5])}")
if python_diffs: trace_lines.append(f" python: {', '.join(python_diffs[:5])}")
elif rust_diffs: trace_lines.append(" python: (no changes)")
elif rust_diffs: trace_lines.append(f" python: (no changes)")
trace_str = "\n".join(trace_lines)
msg = f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ (rust vs python):\n "
msg += "\n ".join(diffs[:10]) + f"\n Recent instructions:\n{trace_str}"
return False, msg, total_steps
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ (rust vs python):\n " + "\n ".join(diffs[:10]) + f"\n Recent instructions:\n{trace_str}", total_steps
rust_result = rust.step()
python_result = python.step()
@@ -250,9 +191,7 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
if rust_result == 1 and python_result == 0:
raise unittest.SkipTest(f"Rust emulator doesn't support instruction: {inst_str}")
trace_str = "\n".join(f" step {s}: PC={pc:3d} {d}" for s, pc, d, _, _ in trace)
msg = (f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: "
f"rust={rust_result}, python={python_result}, inst={inst_str}\n Recent instructions:\n{trace_str}")
return False, msg, total_steps
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: rust={rust_result}, python={python_result}, inst={inst_str}\n Recent instructions:\n{trace_str}", total_steps
# Sync Python state to Rust after instructions with known Rust emulator differences
if sync_after:
@@ -261,12 +200,7 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
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
# Convert Rust's word-based PC to Python's actual address
python.state.pc = python.lib_addr + rust_after.pc * 4
python.state._write_sgpr(SCC.offset, rust_after.scc)
python.state._write_sgpr(VCC_LO.offset, rust_after.vcc)
python.state._write_sgpr(EXEC_LO.offset, rust_after.exec_mask)
prev_sync_after = sync_after
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
@@ -315,11 +249,12 @@ def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int
kernel_ranges = ranges | {(args_ptr, ctypes.sizeof(args))}
set_valid_mem_ranges(kernel_ranges)
program = decode_program(kernel.code)
n_lanes = kernel.local_size[0] * kernel.local_size[1] * kernel.local_size[2]
ok, msg, steps = run_single_kernel(
kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
kernel.local_size, max_steps, debug, trace_len, ki
program, max_steps, debug, trace_len, ki
)
total_steps += steps
if not ok:
@@ -345,8 +280,8 @@ def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list,
ranges.add((args_ptr, ctypes.sizeof(args)))
set_valid_mem_ranges(ranges)
# Legacy wrapper assumes local_size = (n_lanes, 1, 1)
ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, (n_lanes, 1, 1), max_steps, debug, trace_len)
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]]:
@@ -389,7 +324,6 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int],
buf_sizes.append(b.nbytes)
kernels.append(KernelInfo(
code=bytes(sec.content),
src=lowered.prg.p.src,
global_size=tuple(lowered.prg.p.global_size),
local_size=tuple(lowered.prg.p.local_size),
buf_idxs=buf_idxs,
@@ -404,7 +338,6 @@ def get_kernel_from_tinygrad(op_fn) -> tuple[bytes, tuple[int, int, int], tuple[
k = kernels[-1]
return k.code, k.global_size, k.local_size, k.buf_sizes
@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
class TestTinygradKernels(unittest.TestCase):
"""Compare emulators on real tinygrad-compiled kernels."""
@@ -441,8 +374,7 @@ class TestTinygradKernels(unittest.TestCase):
def test_cast(self): self._test_kernel(lambda T: T.empty(32).half().float() + T.empty(32).int().float())
# Pooling - regression for VCC wave32 mode
def test_pool2d(self):
self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4)) + T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4)))
def test_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4)) + T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4)))
# Convolution
def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 2, 8, 8).conv2d(T.empty(2, 2, 3, 3)), max_steps=50000)
@@ -454,7 +386,6 @@ class TestTinygradKernels(unittest.TestCase):
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_exp(self): self._test_kernel(lambda T: T.empty(1024).exp())
def test_cross_entropy(self):
import numpy as np
np.random.seed(0)
@@ -466,55 +397,5 @@ class TestTinygradKernels(unittest.TestCase):
from tinygrad import dtypes
self._test_kernel(lambda T: T([2.0], dtype=dtypes.float64).sin())
def test_sin_large_f32(self):
"""Test sin with large values that trigger Payne-Hanek range reduction."""
# Values around 859240 trigger the Payne-Hanek algorithm
# This tests the integer multiply-high instructions used in range reduction
self._test_kernel(lambda T: T([859240.0, 1000000.0, 100594688.0]).sin())
def test_clip_zero_one(self):
"""Test clip(0, 1) - regression for binary_crossentropy failure."""
import numpy as np
np.random.seed(0)
x_np = np.random.uniform(-2, 2, (32, 10)).astype(np.float32).tolist()
self._test_kernel(lambda T: T(x_np).clip(0, 1))
def test_mod_int64(self):
"""Test int64 modulo, especially edge cases like 1 % -1."""
from tinygrad import dtypes
self._test_kernel(lambda T: T([1, 10, -10, 7], dtype=dtypes.int64) % T([-1, 3, 3, -3], dtype=dtypes.int64))
def test_expand_flatten_sum(self):
"""Test flatten of expanded tensor followed by sum.
Bug: flatten() of an expanded tensor produces wrong results for certain sizes.
Sizes that are multiples of 32 work (32, 48, 64), but sizes like 33, 49, 50 fail.
This breaks masked_select and nonzero operations.
"""
import numpy as np
np.random.seed(0)
x_np = np.random.uniform(-2, 2, (33,)).astype(np.float32)
self._test_kernel(lambda T: (T(x_np.tolist()) > 0.5).unsqueeze(-1).expand(33, 3).flatten().sum())
@unittest.skip("slow and broken with AMD_LLVM=1")
def test_nonzero(self):
"""Test nonzero operation - counts and gathers indices of non-zero elements."""
import numpy as np
np.random.seed(42)
x_np = np.random.rand(10, 5, 3).astype(np.float32)
self._test_kernel(lambda T: (T(x_np.tolist()) > 0.5).nonzero())
@unittest.skip("Precision differences in v_exp/v_log accumulate across kernels, causing memory divergence")
def test_softmax_argmax_fused(self):
"""Test fused softmax+argmax - tracks exp2 precision issue.
The fused kernel recomputes softmax inline and Python emulator's exp2 polynomial
has up to 1 ULP error vs native exp2f, causing accumulated differences.
"""
import torch
torch.manual_seed(0)
x_np = torch.rand(4, 10).numpy()
self._test_kernel(lambda T: T(x_np.tolist()).softmax(1).argmax())
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,48 @@
import unittest
import functools
from tinygrad import Tensor, Device, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.renderer import Estimates
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import s, v, Inst
def assemble_insts(insts:list[Inst], name:str, arch:str, kernarg_size:int=8) -> tuple[UOp, UOp]:
kd = {"kernarg_size":kernarg_size, "user_sgpr_kernarg_segment_ptr":1, "next_free_vgpr":8, "next_free_sgpr":8, "wavefront_size32":1}
disasm = "\n".join([inst.disasm() for inst in insts])
hsasrc = f".text\n.globl {name}\n.p2align 8\n.type fn_name,@function\n{name}:\n{disasm}\ns_code_end\n"
hsasrc += f".rodata\n.p2align 6\n.amdhsa_kernel {name}\n"+"\n".join([f".amdhsa_{k} {v}" for k,v in kd.items()])+"\n.end_amdhsa_kernel"
binary = HIPCompiler(arch).compile(hsasrc)
return UOp(Ops.SOURCE, arg=disasm), UOp(Ops.BINARY, arg=binary)
def custom_add_one(A:UOp, arch:str) -> UOp:
A = A.flatten()
assert dtypes.is_float(A.dtype.base), f"buffer dtype must be float32, got {A.dtype}"
threads = UOp.special(A.size, "lidx0")
insts = [
s_load_b64(s[0:1], s[0:1], soffset=NULL),
s_waitcnt(lgkmcnt=0),
v_lshlrev_b32_e32(v[0], 2, v[0]), # element offset
global_load_b32(v[1], v[0], saddr=s[0:1]),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[2], 1.0),
v_add_f32_e32(v[1], v[1], v[2]),
global_store_b32(addr=v[0], data=v[1], saddr=s[0:1]),
s_endpgm(),
]
sink = UOp.sink(A.base, threads, arg=KernelInfo(name:=f"custom_add_one_{A.size}", estimates=Estimates(ops=A.size, mem=A.size*4*2)))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=(*sink.src, sink)), *assemble_insts(insts, name, arch)))
class TestCustomKernel(unittest.TestCase):
def test_simple(self):
a = Tensor.full((16, 16), 1.).contiguous().realize()
a = Tensor.custom_kernel(a, fxn=functools.partial(custom_add_one, arch=Device[Device.DEFAULT].arch))[0]
ei = a.schedule()[-1].lower()
self.assertEqual(ei.prg.estimates.ops, a.numel())
self.assertEqual(ei.prg.estimates.mem, a.nbytes()*2)
ei.run()
self.assertTrue((a.numpy() == 2.).all())
if __name__ == "__main__":
unittest.main()
@@ -1,8 +1,8 @@
import unittest
from tinygrad.renderer.amd.dsl import *
from tinygrad.renderer.amd.dsl import VDSTYField
from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op
from tinygrad.runtime.autogen.amd.rdna3.ins import VOP1
from extra.assembly.amd.dsl import *
from extra.assembly.amd.dsl import VDSTYField
from extra.assembly.amd.autogen.rdna3.enum import VOP1Op, VOP2Op
from extra.assembly.amd.autogen.rdna3.ins import VOP1
class TestRegisters(unittest.TestCase):
def test_vgpr_single(self):
@@ -4,10 +4,10 @@
Note: Graphics-only formats (EXP, MUBUF, MTBUF, MIMG) are not supported - use GLOBAL/FLAT for memory access in compute.
"""
import unittest
from tinygrad.runtime.autogen.amd.rdna3.ins import *
from tinygrad.renderer.amd.dsl import VCC_HI, EXEC_LO, NULL
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import VCC_HI, EXEC_LO, NULL
OFF = NULL # OFF is alias for NULL
from tinygrad.renderer.amd import detect_format
from extra.assembly.amd.decode import detect_format
class TestDS(unittest.TestCase):
@@ -2,26 +2,21 @@
# the Inst constructor should be looking at the types of the fields to correctly set the value
import unittest, struct
from tinygrad.runtime.autogen.amd.rdna3.ins import *
from tinygrad.renderer.amd.dsl import Inst
from test.amd.test_roundtrip import compile_asm
from test.amd.disasm import disasm
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.test.test_roundtrip import compile_asm
class IntegrationTestBase(unittest.TestCase):
class TestIntegration(unittest.TestCase):
inst: Inst
arch: str
def tearDown(self):
if not hasattr(self, 'inst'): return
b = self.inst.to_bytes()
st = disasm(self.inst)
st = self.inst.disasm()
# Test that the instruction can be compiled by LLVM and produces the same bytes
desc = f"{st:25s} {self.inst} {b!r}"
self.assertEqual(b, compile_asm(st, arch=self.arch), desc)
self.assertEqual(b, compile_asm(st), desc)
print(desc)
class TestIntegration(IntegrationTestBase):
arch: str = "rdna3"
def test_wmma(self):
self.inst = v_wmma_f32_16x16x16_f16(v[0:7], v[184:191], v[136:143], v[0:7])
@@ -129,17 +124,6 @@ class TestIntegration(IntegrationTestBase):
int_inst = s_mov_b32(s[0], struct.unpack("I", struct.pack("f", 1337.0))[0])
self.assertEqual(self.inst, int_inst)
class TestIntegrationCDNA(IntegrationTestBase):
arch = "cdna"
def test_mfma(self):
from tinygrad.runtime.autogen.amd.cdna.ins import v_mfma_f32_16x16x16_f16
self.inst = v_mfma_f32_16x16x16_f16(v[0:3], v[0:1], v[0:1], 0)
def test_mfma_fp8(self):
from tinygrad.runtime.autogen.amd.cdna.ins import v_mfma_f32_16x16x128_f8f6f4
self.inst = v_mfma_f32_16x16x128_f8f6f4(v[0:3], v[0:5], v[0:5], 1, cbsz=2, blgp=2)
class TestRegisterSliceSyntax(unittest.TestCase):
"""
Issue: Register slice syntax should use AMD assembly convention (inclusive end).
@@ -161,9 +145,9 @@ class TestRegisterSliceSyntax(unittest.TestCase):
# Round-trip: DSL -> disasm -> DSL should preserve register count
reg = s[4:7] # 4 registers in AMD convention
inst = s_load_b128(reg, s[0:1], NULL, 0)
d = disasm(inst)
disasm = inst.disasm()
# Disasm shows s[4:7] - user should be able to copy this back
self.assertIn("s[4:7]", d)
self.assertIn("s[4:7]", disasm)
# And s[4:7] in DSL should give the same 4 registers
reg_from_disasm = s[4:7]
self.assertEqual(reg_from_disasm.sz, 4, "s[4:7] from disasm should give 4 registers")
+258
View File
@@ -0,0 +1,258 @@
#!/usr/bin/env python3
"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
import unittest, io, sys
from extra.assembly.amd.autogen.rdna3.ins import *
def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
def disassemble(lib: bytes, arch: str = "gfx1100") -> str:
"""Disassemble ELF binary using tinygrad's compiler, return raw output."""
from tinygrad.runtime.support.compiler_amd import HIPCompiler
old_stdout = sys.stdout
sys.stdout = io.StringIO()
HIPCompiler(arch).disassemble(lib)
output = sys.stdout.getvalue()
sys.stdout = old_stdout
return output
def parse_disassembly(raw: str) -> list[str]:
"""Parse disassembly output to list of instruction mnemonics."""
lines = []
for line in raw.splitlines():
if line.startswith('\t'):
instr = line.split('//')[0].strip()
if instr: lines.append(instr)
return lines
def assemble_and_disassemble(instructions: list, arch: str = "gfx1100") -> list[str]:
"""Assemble instructions with our DSL, then disassemble with AMD toolchain."""
from tinygrad.runtime.support.compiler_amd import HIPCompiler
# Generate bytes from our DSL
code_bytes = b''.join(inst.to_bytes() for inst in instructions)
# Wrap in minimal ELF-compatible assembly with .byte directives
byte_str = ', '.join(f'0x{b:02x}' for b in code_bytes)
asm_src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n.byte {byte_str}\n"
# Assemble with AMD COMGR and disassemble
lib = HIPCompiler(arch).compile(asm_src)
return parse_disassembly(disassemble(lib, arch))
class TestIntegration(unittest.TestCase):
"""Test our DSL output matches LLVM disassembly."""
def test_simple_sop1(self):
"""Test SOP1 instructions round-trip."""
instructions = [
s_mov_b32(s[0], s[1]),
s_mov_b32(s[2], 0),
s_not_b32(s[3], s[4]),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_mov_b32', disasm[0])
self.assertIn('s_mov_b32', disasm[1])
self.assertIn('s_not_b32', disasm[2])
def test_simple_sop2(self):
"""Test SOP2 instructions round-trip."""
instructions = [
s_add_u32(s[0], s[1], s[2]),
s_sub_u32(s[3], s[4], 10),
s_and_b32(s[5], s[6], s[7]),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_add_u32', disasm[0])
self.assertIn('s_sub_u32', disasm[1])
self.assertIn('s_and_b32', disasm[2])
def test_simple_vop2(self):
"""Test VOP2 instructions round-trip."""
instructions = [
v_add_f32_e32(v[0], v[1], v[2]),
v_mul_f32_e32(v[3], 1.0, v[4]), # 1.0 is inline constant
v_and_b32_e32(v[5], 10, v[6]), # small inline constant
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('v_add_f32', disasm[0])
self.assertIn('v_mul_f32', disasm[1])
def test_control_flow(self):
"""Test control flow instructions."""
instructions = [
s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
s_endpgm(),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_waitcnt', disasm[0])
self.assertIn('s_endpgm', disasm[1])
def test_memory_ops(self):
"""Test memory instructions."""
instructions = [
s_load_b32(s[0], s[0:1], NULL),
s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
global_store_b32(addr=v[0:1], data=v[2], saddr=OFF),
s_endpgm(),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_load_b32', disasm[0])
self.assertIn('s_waitcnt', disasm[1])
self.assertIn('global_store_b32', disasm[2])
def test_full_kernel(self):
"""Test a complete kernel similar to tinygrad output."""
# Simple kernel: load value, add 1, store back
instructions = [
# Get thread ID
v_mov_b32_e32(v[0], s[0]), # base addr low
v_mov_b32_e32(v[1], s[1]), # base addr high
# Load value
global_load_b32(vdst=v[2], addr=v[0:1], saddr=OFF),
s_waitcnt(simm16=waitcnt(vmcnt=0)),
# Add 1.0
v_add_f32_e32(v[2], 1.0, v[2]),
# Store result
global_store_b32(addr=v[0:1], data=v[2], saddr=OFF),
s_endpgm(),
]
disasm = assemble_and_disassemble(instructions)
# Verify key instructions are present
self.assertTrue(any('global_load' in d for d in disasm))
self.assertTrue(any('v_add_f32' in d for d in disasm))
self.assertTrue(any('global_store' in d for d in disasm))
self.assertTrue(any('s_endpgm' in d for d in disasm))
def test_bytes_roundtrip(self):
"""Test that our bytes match what AMD assembler produces."""
from tinygrad.runtime.support.compiler_amd import HIPCompiler
# Simple instruction
inst = s_mov_b32(s[0], s[1])
our_bytes = inst.to_bytes()
# Assemble same instruction with AMD toolchain
asm_src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\ns_mov_b32 s0, s1\n"
compiler = HIPCompiler("gfx1100")
lib = compiler.compile(asm_src)
raw = disassemble(lib)
for line in raw.splitlines():
if 's_mov_b32' in line and '//' in line:
# Extract hex bytes from comment: "// 000000001300: BE800001"
comment = line.split('//')[1].strip()
hex_str = comment.split(':')[1].strip()
# Convert big-endian hex string to little-endian bytes
amd_bytes = bytes.fromhex(hex_str)[::-1] # reverse for little-endian
self.assertEqual(our_bytes, amd_bytes, f"Bytes mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
return
self.fail("Could not find s_mov_b32 in disassembly")
class TestTinygradIntegration(unittest.TestCase):
"""Test that we can parse disassembled tinygrad kernels."""
def test_simple_add_kernel(self):
"""Generate a simple add kernel from tinygrad and verify disassembly."""
from tinygrad import Tensor
from tinygrad.codegen import get_program
from tinygrad.renderer.cstyle import AMDHIPRenderer
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.uop.ops import Ops
# Create a computation that generates a real kernel
a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
b = Tensor([5.0, 6.0, 7.0, 8.0]).realize()
c = a + b
# Get schedule and find SINK
schedule = c.schedule()
sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
self.assertTrue(len(sink_items) > 0, "No SINK in schedule")
# Generate program
renderer = AMDHIPRenderer('gfx1100')
prg = get_program(sink_items[0].ast, renderer)
self.assertIsNotNone(prg.src)
# Compile and disassemble
compiler = HIPCompiler('gfx1100')
lib = compiler.compile(prg.src)
raw_disasm = disassemble(lib)
instrs = parse_disassembly(raw_disasm)
# Verify we got some instructions
self.assertTrue(len(instrs) > 0, "No instructions in disassembly")
# Should have an endpgm
self.assertTrue(any('s_endpgm' in i for i in instrs), "Missing s_endpgm")
def test_matmul_kernel(self):
"""Generate a matmul kernel and verify disassembly has expected patterns."""
from tinygrad import Tensor
from tinygrad.codegen import get_program
from tinygrad.renderer.cstyle import AMDHIPRenderer
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.uop.ops import Ops
# Create a small matmul
a = Tensor.rand(4, 4).realize()
b = Tensor.rand(4, 4).realize()
c = a @ b
# Get schedule
schedule = c.schedule()
sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
self.assertTrue(len(sink_items) > 0)
# Generate and compile
renderer = AMDHIPRenderer('gfx1100')
prg = get_program(sink_items[0].ast, renderer)
compiler = HIPCompiler('gfx1100')
lib = compiler.compile(prg.src)
raw_disasm = disassemble(lib)
instrs = parse_disassembly(raw_disasm)
# Matmul should have multiply and add instructions
has_mul = any('mul' in i.lower() for i in instrs)
has_add = any('add' in i.lower() for i in instrs)
self.assertTrue(has_mul or has_add, "Matmul should have mul/add ops")
def test_disasm_to_bytes_roundtrip(self):
"""Parse disassembled instructions and verify we can re-encode some of them."""
from tinygrad import Tensor
from tinygrad.codegen import get_program
from tinygrad.renderer.cstyle import AMDHIPRenderer
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.uop.ops import Ops
# Simple kernel
a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
b = (a * 2.0)
schedule = b.schedule()
sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
if not sink_items: return # skip if no kernel
renderer = AMDHIPRenderer('gfx1100')
prg = get_program(sink_items[0].ast, renderer)
compiler = HIPCompiler('gfx1100')
lib = compiler.compile(prg.src)
raw_disasm = disassemble(lib)
# Find s_endpgm and verify we can encode it
for line in raw_disasm.splitlines():
if 's_endpgm' in line and '//' in line:
# Extract bytes from comment
comment = line.split('//')[1].strip()
hex_str = comment.split(':')[1].strip()
amd_bytes = bytes.fromhex(hex_str)[::-1]
# Our encoding
our_inst = s_endpgm()
our_bytes = our_inst.to_bytes()
self.assertEqual(our_bytes, amd_bytes, f"s_endpgm mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
return
if __name__ == "__main__":
unittest.main()
@@ -8,11 +8,11 @@ Only compute-relevant instruction formats are tested. Graphics-only formats not
- VIMAGE/VSAMPLE: image sampling instructions (RDNA4)
- VBUFFER: buffer instructions (RDNA4)
"""
import unittest, re, functools
import unittest, re, subprocess, functools
from tinygrad.helpers import fetch
from test.amd.disasm import disasm
from tinygrad.renderer.amd import decode_inst, detect_format
from test.amd.helpers import llvm_assemble, llvm_filter_valid_asm, get_target, get_mattr
from extra.assembly.amd.disasm import disasm
from extra.assembly.amd.decode import decode_inst, detect_format
from extra.assembly.amd.test.helpers import get_llvm_mc, get_target, get_mattr
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-21.1.0/llvm/test/MC/AMDGPU"
@@ -74,13 +74,42 @@ def _get_tests_uncached(f: str, arch: str) -> list[tuple[str, bytes]]:
# Exclude v_interp_* (graphics-only, not on CDNA)
if arch == "cdna": tests = [(asm, data) for asm, data in tests if not asm.startswith('v_interp_')]
# Filter out tests where original ASM isn't valid on target (e.g., gfx9 tests with gfx942/gfx950 constraints)
if arch == "cdna" and not ('gfx942' in f or 'gfx950' in f or 'gfx90a' in f):
tests = llvm_filter_valid_asm(tests, get_target(arch), get_mattr(arch))
if arch == "cdna" and not ('gfx942' in f or 'gfx950' in f or 'gfx90a' in f): tests = _filter_valid_asm(tests, arch)
return tests
@functools.cache
def _get_tests(f: str, arch: str) -> list[tuple[str, bytes]]: return _get_tests_uncached(f, arch)
def _compile_asm_batch(instrs: list[str], arch: str = "rdna3", mcpu: str|None = None) -> list[bytes]:
if not instrs: return []
mcpu, mattr = mcpu or get_target(arch), get_mattr(arch)
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', f'-mattr={mattr}', '-show-encoding'],
input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True, timeout=30)
if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
return [bytes.fromhex(line.split('encoding:')[1].strip()[1:-1].replace('0x', '').replace(',', '').replace(' ', ''))
for line in result.stdout.split('\n') if 'encoding:' in line]
def _filter_valid_asm(tests: list[tuple[str, bytes]], arch: str) -> list[tuple[str, bytes]]:
"""Filter out tests where the original ASM isn't valid on the target (e.g., gfx9 tests with gfx942/gfx950 constraints)."""
if not tests: return []
mcpu = get_target(arch)
# Batch assemble all instructions, parse stderr to find which lines failed
instrs = [asm for asm, _ in tests]
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', '-show-encoding'],
input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True, timeout=30)
# Parse error lines from stderr (format: "<stdin>:N:..." where N is 1-indexed, line 1 is ".text")
failed_lines = set()
for line in result.stderr.split('\n'):
if m := re.match(r'<stdin>:(\d+):', line): failed_lines.add(int(m.group(1)) - 1) # -1 for .text, so line 2 -> index 1 -> tests[0]
# Also filter out tests where LLVM roundtrip doesn't match original (reserved bits set in original)
valid = [(asm, data) for i, (asm, data) in enumerate(tests) if (i + 1) not in failed_lines]
if not valid: return []
llvm_result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', '-show-encoding'],
input=".text\n" + "\n".join(asm for asm, _ in valid) + "\n", capture_output=True, text=True, timeout=30)
llvm_bytes = [bytes.fromhex(line.split('encoding:')[1].strip()[1:-1].replace('0x', '').replace(',', '').replace(' ', ''))
for line in llvm_result.stdout.split('\n') if 'encoding:' in line]
return [(asm, data) for (asm, data), lb in zip(valid, llvm_bytes) if lb == data]
def _make_test(f: str, arch: str, test_type: str):
def test(self):
tests = _get_tests(f, arch)
@@ -96,28 +125,6 @@ def _make_test(f: str, arch: str, test_type: str):
except ValueError: skipped += 1 # skip invalid opcodes not in enum
print(f"{name}: {passed} passed, {skipped} skipped")
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
elif test_type == "repr":
# Test that eval(repr(inst)) reproduces the instruction
if arch == "rdna3": import tinygrad.runtime.autogen.amd.rdna3.ins as ins # type: ignore[no-redef]
elif arch == "rdna4": import tinygrad.runtime.autogen.amd.rdna4.ins as ins # type: ignore[no-redef]
elif arch == "cdna": import tinygrad.runtime.autogen.amd.cdna.ins as ins # type: ignore[no-redef]
ns = {k: getattr(ins, k) for k in dir(ins) if not k.startswith('_')}
passed, skipped = 0, 0
for _, data in tests:
try:
decoded = detect_format(data, arch).from_bytes(data)
if decoded.to_bytes()[:len(data)] != data:
skipped += 1
continue # skip if binary roundtrip fails
r = repr(decoded)
try:
decoded2 = eval(r, ns) # noqa: S307
if decoded == decoded2: passed += 1
else: skipped += 1
except Exception: skipped += 1
except ValueError: skipped += 1
print(f"{name}: {passed} passed, {skipped} skipped")
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
elif test_type == "disasm":
to_test = []
for _, data in tests:
@@ -126,12 +133,12 @@ def _make_test(f: str, arch: str, test_type: str):
enc = decoded.to_bytes()[:len(data)]
# Skip if roundtrip fails, disasm fails, or op_name is missing (disasm starts with space)
if enc == data and (d := disasm(decoded)) and not d.startswith(' '): to_test.append((enc, d))
except Exception: pass
except: pass
skipped = len(tests) - len(to_test)
print(f"{name}: {len(to_test)} passed, {skipped} skipped")
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
# Compare disasm->reassemble with original encoding (filter reserved bit cases where LLVM can't reproduce)
llvm_bytes = llvm_assemble([t[1] for t in to_test], mcpu, get_mattr(arch))
llvm_bytes = _compile_asm_batch([t[1] for t in to_test], arch, mcpu)
valid = [(enc, d, llvm) for (enc, d), llvm in zip(to_test, llvm_bytes) if llvm == enc]
print(f"{name}: {len(valid)}/{len(to_test)} matched LLVM encoding")
for enc, _, llvm in valid: self.assertEqual(llvm, enc)
@@ -142,15 +149,12 @@ 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_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "disasm"))
setattr(TestLLVM, f"test_rdna3_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "repr"))
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"))
setattr(TestLLVM, f"test_cdna_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "repr"))
for f in RDNA4_FILES:
setattr(TestLLVM, f"test_rdna4_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "roundtrip"))
setattr(TestLLVM, f"test_rdna4_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "disasm"))
setattr(TestLLVM, f"test_rdna4_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "repr"))
if __name__ == "__main__":
unittest.main()
@@ -1,6 +1,6 @@
#!/usr/bin/env python3
"""Test that invalid instructions raise exceptions through the mock GPU stack."""
import unittest, subprocess, os, sys, time
import unittest, subprocess, os, time
class TestMockGPUInvalidInstruction(unittest.TestCase):
def test_unsupported_instruction_raises(self):
@@ -43,7 +43,7 @@ dev.synchronize()
env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
st = time.perf_counter()
result = subprocess.run([sys.executable, "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
result = subprocess.run(["python", "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
elapsed = time.perf_counter() - st
self.assertNotEqual(result.returncode, 0, "should have raised")
+403
View File
@@ -0,0 +1,403 @@
#!/usr/bin/env python3
"""Tests for the RDNA3 pseudocode DSL."""
import unittest
from extra.assembly.amd.pcode import (Reg, TypedView, TypedView, MASK32, MASK64,
_f32, _i32, _f16, _i16, f32_to_f16, isNAN, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
BYTE_PERMUTE, v_sad_u8, v_msad_u8, _compile_pseudocode, _expr, compile_pseudocode)
from extra.assembly.amd.test.helpers import ExecContext
from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
from extra.assembly.amd.autogen.rdna3.enum import VOP3SDOp, VOPCOp
# Compile pseudocode functions on demand for regression tests
_VOP3SDOp_V_DIV_SCALE_F32 = compile_pseudocode('VOP3SDOp', 'V_DIV_SCALE_F32', PCODE[VOP3SDOp.V_DIV_SCALE_F32])
_VOPCOp_V_CMP_CLASS_F32 = compile_pseudocode('VOPCOp', 'V_CMP_CLASS_F32', PCODE[VOPCOp.V_CMP_CLASS_F32_E32])
class TestReg(unittest.TestCase):
def test_u32_read(self):
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()
@@ -1,15 +1,11 @@
#!/usr/bin/env python3
"""Test PDF pseudocode extraction from generate.py."""
"""Test PDF pseudocode extraction from amdxml.py."""
import unittest
from tinygrad.renderer.amd.generate import extract_pdf_text, extract_pcode, parse_xml, ARCHS, FIXES
from extra.assembly.amd.amdxml import extract_pdf_text, extract_pcode, parse_xml, ARCHS, FIXES
EXPECTED_PAGES = {"rdna3": 655, "rdna4": 711, "cdna": 610}
class TestPcodePDF(unittest.TestCase):
pages: dict
enums: dict
pcode: dict
@classmethod
def setUpClass(cls):
cls.pages = {arch: extract_pdf_text(cfg["pdf"]) for arch, cfg in ARCHS.items()}
@@ -37,8 +33,7 @@ class TestPcodePDF(unittest.TestCase):
'tmp = MEM[ADDR].u64;\nsrc = DATA.u64;\nMEM[ADDR].u64 = src >= tmp ? src : tmp;\nRETURN_DATA.u64 = tmp')
# GLOBAL_STORE_B128: should have 4 MEM stores (not truncated)
self.assertEqual(pcode[('GLOBAL_STORE_B128', 29)],
'MEM[ADDR].b32 = VDATA[31 : 0];\nMEM[ADDR + 4U].b32 = VDATA[63 : 32];\n'
'MEM[ADDR + 8U].b32 = VDATA[95 : 64];\nMEM[ADDR + 12U].b32 = VDATA[127 : 96]')
'MEM[ADDR].b32 = VDATA[31 : 0];\nMEM[ADDR + 4U].b32 = VDATA[63 : 32];\nMEM[ADDR + 8U].b32 = VDATA[95 : 64];\nMEM[ADDR + 12U].b32 = VDATA[127 : 96]')
# S_CMOVK_I32: should have full if/endif block
self.assertEqual(pcode[('S_CMOVK_I32', 2)],
"if SCC then\nD0.i32 = 32'I(signext(SIMM16.i16))\nendif")
+95
View File
@@ -0,0 +1,95 @@
#!/usr/bin/env python3
import unittest, subprocess
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.test.helpers import get_llvm_mc
def llvm_assemble(asm: str) -> bytes:
"""Assemble using llvm-mc and return bytes."""
result = subprocess.run(
[get_llvm_mc(), "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
input=asm, capture_output=True, text=True
)
out = b''
for line in result.stdout.split('\n'):
if 'encoding:' in line:
enc = line.split('encoding:')[1].strip()
enc = enc.strip('[]').replace('0x', '').replace(',', '')
out += bytes.fromhex(enc)
if not out: raise ValueError(f"no encoding found: {result.stdout} {result.stderr}")
return out
class TestRDNA3Asm(unittest.TestCase):
def test_full_program(self):
"""Test the full program from rdna3fun.py matches llvm-mc output."""
program = [
v_bfe_u32(v[1], v[0], 10, 10),
s_load_b128(s[4:7], s[0:1], NULL),
v_and_b32_e32(v[0], 0x3FF, v[0]),
s_mulk_i32(s[3], 0x87),
v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2]),
v_mul_u32_u24_e32(v[0], 45, v[0]),
v_ashrrev_i32_e32(v[2], 31, v[1]),
v_add3_u32(v[0], v[0], s[3], v[1]),
v_lshlrev_b64(v[2:3], 2, v[1:2]),
v_ashrrev_i32_e32(v[1], 31, v[0]),
v_lshlrev_b64(v[0:1], 2, v[0:1]),
s_waitcnt(0xfc07), # lgkmcnt(0)
v_add_co_u32(v[2], VCC_LO, s[6], v[2]),
v_add_co_ci_u32_e32(v[3], s[7], v[3]),
v_add_co_u32(v[0], VCC_LO, s[4], v[0]),
global_load_b32(vdst=v[2], addr=v[2:3], saddr=OFF),
v_add_co_ci_u32_e32(v[1], s[5], v[1]),
s_waitcnt(0x03f7), # vmcnt(0)
global_store_b32(addr=v[0:1], data=v[2], saddr=OFF),
s_endpgm(),
]
asm = """
v_bfe_u32 v1, v0, 10, 10
s_load_b128 s[4:7], s[0:1], null
v_and_b32_e32 v0, 0x3FF, v0
s_mulk_i32 s3, 0x87
v_mad_u64_u32 v[1:2], null, s2, 3, v[1:2]
v_mul_u32_u24_e32 v0, 45, v0
v_ashrrev_i32_e32 v2, 31, v1
v_add3_u32 v0, v0, s3, v1
v_lshlrev_b64 v[2:3], 2, v[1:2]
v_ashrrev_i32_e32 v1, 31, v0
v_lshlrev_b64 v[0:1], 2, v[0:1]
s_waitcnt lgkmcnt(0)
v_add_co_u32 v2, vcc_lo, s6, v2
v_add_co_ci_u32_e32 v3, vcc_lo, s7, v3, vcc_lo
v_add_co_u32 v0, vcc_lo, s4, v0
global_load_b32 v2, v[2:3], off
v_add_co_ci_u32_e32 v1, vcc_lo, s5, v1, vcc_lo
s_waitcnt vmcnt(0)
global_store_b32 v[0:1], v2, off
s_endpgm
"""
expected = llvm_assemble(asm)
for inst,rt in zip(program, asm.strip().split("\n")): print(f"{inst.disasm():50s} {rt}")
actual = b''.join(inst.to_bytes() for inst in program)
self.assertEqual(actual, expected)
def test_sop2_s_add_u32(self):
inst = SOP2(SOP2Op.S_ADD_U32, s[3], s[0], s[1])
expected = llvm_assemble("s_add_u32 s3, s0, s1")
self.assertEqual(inst.to_bytes(), expected)
def test_vop2_v_and_b32_inline_const(self):
inst = v_and_b32_e32(v[0], 10, v[0])
expected = llvm_assemble("v_and_b32_e32 v0, 10, v0")
self.assertEqual(inst.to_bytes(), expected)
def test_sopp_s_endpgm(self):
inst = s_endpgm()
expected = llvm_assemble("s_endpgm")
self.assertEqual(inst.to_bytes(), expected)
def test_sop1_s_mov_b32(self):
inst = s_mov_b32(s[0], s[1])
expected = llvm_assemble("s_mov_b32 s0, s1")
self.assertEqual(inst.to_bytes(), expected)
if __name__ == "__main__":
unittest.main()
@@ -1,10 +1,9 @@
#!/usr/bin/env python3
"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
import unittest, io, sys, re
from tinygrad import Device
from tinygrad.renderer.amd import detect_format
from test.amd.helpers import llvm_assemble, llvm_disasm, get_target, get_mattr
from test.amd.disasm import disasm
import unittest, io, sys, re, subprocess, os
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.decode import decode_inst, detect_format
from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump, get_target, get_mattr
def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
"""Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
@@ -31,20 +30,46 @@ def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
def compile_asm(instr: str, arch: str = 'rdna3') -> bytes:
"""Compile a single instruction using LLVM."""
return llvm_assemble([instr], get_target(arch), get_mattr(arch))[0]
return compile_asm_batch([instr], arch)[0]
def compile_asm_batch(instrs: list[str], arch: str = 'rdna3') -> list[bytes]:
"""Compile multiple instructions with a single LLVM emission."""
return llvm_assemble(instrs, get_target(arch), get_mattr(arch))
"""Compile multiple instructions with a single llvm-mc call."""
if not instrs: return []
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={get_target(arch)}', f'-mattr={get_mattr(arch)}', '-show-encoding'],
input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
encodings = []
for line in result.stdout.split('\n'):
if 'encoding:' in line:
enc = line.split('encoding:')[1].strip()
if enc.startswith('[') and enc.endswith(']'):
encodings.append(bytes.fromhex(enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')))
if len(encodings) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(encodings)}")
return encodings
def compile_and_disasm_batch(instrs: list[str], arch: str = 'rdna3') -> list[str]:
"""Compile instructions with LLVM and get LLVM's disassembly."""
import tempfile
if not instrs: return []
mcpu, mattr = get_target(arch), get_mattr(arch)
code = b''.join(llvm_assemble(instrs, mcpu, mattr))
return llvm_disasm(code, mcpu, mattr)[:len(instrs)]
src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n" + "\n".join(f" {instr}" for instr in instrs) + "\n"
with tempfile.NamedTemporaryFile(suffix='.o', delete=False) as f:
obj_path = f.name
try:
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', f'-mattr={mattr}', '-filetype=obj', '-o', obj_path],
input=src, capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
result = subprocess.run([get_llvm_objdump(), '-d', f'--mcpu={mcpu}', obj_path], capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-objdump failed: {result.stderr.strip()}")
results: list[str] = []
for line in result.stdout.splitlines():
if '//' not in line: continue
instr = line.split('//')[0].strip()
if instr: results.append(instr)
return results[:len(instrs)]
finally:
os.unlink(obj_path)
@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
class TestTinygradKernelRoundtrip(unittest.TestCase):
"""Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
arch = 'rdna3'
@@ -57,23 +82,24 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
"""
arch = self.arch
from test.amd.test_compare_emulators import get_kernels_from_tinygrad
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
from tinygrad.helpers import AMD_LLVM
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)
# rendered source can be C or llvmir
compiler = (AMDLLVMCompiler if AMD_LLVM else HIPCompiler)(get_target(arch))
compiler = HIPCompiler(get_target(arch))
# 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
code = next((s.content for s in elf_loader(compiler.compile(kernel.src))[1] if s.name == ".text"))
while offset < len(code):
remaining = code[offset:]
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
@@ -83,7 +109,7 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
size = decoded.size() # actual size including literal
orig_bytes = remaining[:size]
reencoded = decoded.to_bytes()
our_disasm = disasm(decoded)
our_disasm = decoded.disasm()
decode_ok = reencoded == orig_bytes
decode_err: str | None = None if decode_ok else f"orig={orig_bytes.hex()} reenc={reencoded.hex()}"
decoded_instrs.append((ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err))
@@ -147,20 +173,20 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
if our_disasm is None:
disasm_skipped += 1
elif idx in disasm_llvm_map:
llvm_disasm_str = disasm_llvm_map[idx]
if our_disasm == llvm_disasm_str:
llvm_disasm = disasm_llvm_map[idx]
if our_disasm == llvm_disasm:
disasm_passed += 1
else:
disasm_failed += 1
disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm_str}'")
disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
else:
disasm_skipped += 1
print(f"[{arch}] decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
print(f"[{arch}] asm via llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
print(f"[{arch}] disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
self.assertEqual(decode_failed, 0, "Decode failures:\n" + "\n".join(decode_failures[:20]))
self.assertEqual(asm_failed, 0, "Asm failures:\n" + "\n".join(asm_failures[:20]))
self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
# Note: disasm string comparison is informational only - formatting differences between LLVM versions are expected
# Basic unary ops
@@ -208,6 +234,7 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
# Fused ops
def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
@unittest.skip("RDNA4 decode roundtrip not yet supported")
class TestTinygradKernelRoundtripRDNA4(TestTinygradKernelRoundtrip): arch = 'rdna4'
@unittest.skip("CDNA decode roundtrip not yet supported")
@@ -5,23 +5,20 @@ from pathlib import Path
from tinygrad.helpers import DEBUG
from tinygrad.runtime.autogen import rocprof
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.renderer.amd import decode_inst
from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4,
InstOp, InstOpRDNA4, print_packets)
from test.amd.helpers import TARGET_TO_ARCH
from extra.assembly.amd.decode import decode_inst
from extra.assembly.amd.autogen.rdna3.ins import SOPP
from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
from extra.assembly.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVEEND, INST, VALUINST, IMMEDIATE, IMMEDIATE_MASK,
ALUEXEC, VMEMEXEC, PACKET_TYPES, InstOp, print_packets)
from extra.assembly.amd.test.helpers import TARGET_TO_ARCH
import tinygrad
EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
EXAMPLES_DIR = Path(__file__).parent.parent.parent.parent / "sqtt/examples"
# INST ops for non-traced SIMDs (excluded from instruction count)
OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LDS_STORE_64, InstOp.OTHER_LDS_STORE_128,
InstOp.OTHER_FLAT_LOAD, InstOp.OTHER_FLAT_STORE, InstOp.OTHER_FLAT_STORE_64, InstOp.OTHER_FLAT_STORE_96,
InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
InstOp.OTHER_GLOBAL_STORE_VADDR_128}
OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.UNK_60}
# ═══════════════════════════════════════════════════════════════════════════════
# ROCPROF DECODER
@@ -34,18 +31,18 @@ def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
assert text is not None, "no .text section found"
text_off, text_size = text.header.sh_addr, text.header.sh_size
blob_iter, current_blob = iter(blobs), [None] # type: ignore[var-annotated]
blob_iter, current_blob = iter(blobs), [None]
occupancy_records: list[tuple[int, int, int, int, bool]] = [] # (wave_id, simd, cu, time, is_start)
wave_insts: list[list[tuple[int, int]]] = [] # per-wave list of (time, stall)
@rocprof.rocprof_trace_decoder_se_data_callback_t
def copy_cb(buf, buf_size, _): # type: ignore[no-untyped-def]
def copy_cb(buf, buf_size, _):
blob = next(blob_iter, None)
if blob is None: return 0
current_blob[0] = (ctypes.c_ubyte * len(blob)).from_buffer_copy(blob) # type: ignore[call-overload]
buf[0] = ctypes.cast(current_blob[0], ctypes.POINTER(ctypes.c_ubyte)) # type: ignore[arg-type]
buf_size[0] = len(current_blob[0]) # type: ignore[arg-type]
return len(current_blob[0]) # type: ignore[arg-type]
current_blob[0] = (ctypes.c_ubyte * len(blob)).from_buffer_copy(blob)
buf[0] = ctypes.cast(current_blob[0], ctypes.POINTER(ctypes.c_ubyte))
buf_size[0] = len(current_blob[0])
return len(current_blob[0])
@rocprof.rocprof_trace_decoder_trace_callback_t
def trace_cb(record_type, events_ptr, n, _):
@@ -89,18 +86,16 @@ def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
except Exception as e: exc = e
(t:=threading.Thread(target=worker, daemon=True)).start()
t.join(timeout=5)
t.join(timeout=1)
if exc is not None: raise exc
if t.is_alive(): raise RuntimeError("rocprof decoder timeout")
return occupancy_records, wave_insts
class SQTTExamplesTestBase(unittest.TestCase):
target: str
examples: dict
class TestSQTTExamples(unittest.TestCase):
target = "gfx1100"
@classmethod
def setUpClass(cls):
if cls is SQTTExamplesTestBase: raise unittest.SkipTest("base class")
cls.examples = {}
for pkl_path in sorted((EXAMPLES_DIR/cls.target).glob("*.pkl")):
with open(pkl_path, "rb") as f:
@@ -118,27 +113,23 @@ class SQTTExamplesTestBase(unittest.TestCase):
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
packets = list(decode(event.blob))
if DEBUG >= 2:
print(f"\n=== {name} event {i} ===")
print_packets(packets)
if DEBUG >= 2: print(f"\n=== {name} event {i} ==="); print_packets(packets)
self.assertGreater(len(packets), 0, f"no packets decoded from {name} event {i}")
self.assertIsInstance(packets[0], LAYOUT_HEADER, f"first packet should be LAYOUT_HEADER in {name}")
def test_packet_types_valid(self):
all_classes = set(PACKET_TYPES_RDNA3.values()) | set(PACKET_TYPES_RDNA4.values())
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
for pkt in decode(event.blob):
# Use isinstance to handle layout-specific subclasses (e.g., WAVESTART_RDNA4)
self.assertTrue(any(isinstance(pkt, cls) for cls in all_classes), f"unknown packet type {type(pkt)} in {name}")
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, WAVESTART_RDNA4))]), 0, f"no WAVESTART in {name}")
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):
@@ -153,11 +144,17 @@ class SQTTExamplesTestBase(unittest.TestCase):
if "gemm" not in name: continue
with self.subTest(example=name):
all_packets = [p for e in events for p in decode(e.blob)]
self.assertGreater(len([p for p in all_packets if isinstance(p, (INST, INST_RDNA4))]), 0, f"no INST packets in {name}")
self.assertGreater(len([p for p in all_packets if isinstance(p, INST)]), 0, f"no INST packets in {name}")
expected: dict[str, list[int]] = {} # override in subclasses
expected = {
"profile_empty_run_0": [1803, 1908, 1928, 1979, 2006, 1912],
"profile_empty_run_1": [1803, 1908, 1928, 1979, 2006, 1912],
"profile_gemm_run_0": [2531, 1844, 1864, 1915, 1942, 1848, 3074, 1919, 1939, 1990, 2017, 1923, 19026, 1919, 1939, 1990, 2017, 1929],
"profile_gemm_run_1": [2554, 1844, 1864, 1915, 1942, 1848, 3084, 1919, 1939, 1990, 2017, 1923, 19010, 1919, 1939, 1990, 2017, 1923],
"profile_plus_run_0": [1900, 1908, 1928, 1979, 2006, 1912],
"profile_plus_run_1": [1856, 1908, 1928, 1979, 2006, 1912],
}
def test_packet_counts(self):
if not self.expected: self.skipTest("no expected packet counts for this target")
for name, (events, *_) in self.examples.items():
with self.subTest(example=name):
if not self.expected.get(name): continue
@@ -181,7 +178,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
for event in events:
wave_starts: dict[tuple[int, int, int], int] = {}
for p in decode(event.blob):
if isinstance(p, (WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
if isinstance(p, WAVESTART): 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}")
@@ -198,27 +195,15 @@ class SQTTExamplesTestBase(unittest.TestCase):
for event in events:
for p in decode(event.blob):
if isinstance(p, INST) and p.op not in OTHER_SIMD_OPS: our_insts.append(p._time)
elif isinstance(p, INST_RDNA4) and p.op not in OTHER_SIMD_OPS_RDNA4: our_insts.append(p._time)
elif isinstance(p, 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}")
class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
target = "gfx1100"
expected = {
"profile_empty_run_0": [1844, 1885, 1905, 1956, 1983, 1889],
"profile_empty_run_1": [1780, 1885, 1905, 1956, 1983, 1889],
"profile_gemm_run_0": [2656, 2025, 2045, 2096, 2123, 2029, 3183, 2019, 2039, 2090, 2117, 2023, 19119, 2013, 2033, 2084, 2111, 2017],
"profile_gemm_run_1": [2662, 2025, 2045, 2096, 2123, 2029, 3179, 2019, 2039, 2090, 2117, 2023, 19113, 2071, 2091, 2142, 2169, 2075],
"profile_plus_run_0": [1886, 2013, 2033, 2084, 2111, 2017],
"profile_plus_run_1": [1988, 2071, 2091, 2142, 2169, 2075],
}
#class TestSQTTExamplesRDNA4(TestSQTTExamples): target = "gfx1200"
class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
@unittest.skip("TODO: fix CDNA")
class TestSQTTExamplesCDNA(SQTTExamplesTestBase): target = "gfx950"
#class TestSQTTExamplesCDNA(TestSQTTExamples): target = "gfx950"
if __name__ == "__main__":
unittest.main()
+2 -1
View File
@@ -67,11 +67,12 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs))
if not wasm:
thread_id = 0 # NOTE: export does not support threading, thread_id is always 0
for name,cl in bufs_to_save.items():
weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)}, {thread_id});" for (name, args, _global_size, _local_size) in statements] + ["}"]
return '\n'.join(headers + cprog)
else:
if bufs_to_save:
+1 -1
View File
@@ -28,7 +28,7 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))
def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
_, input_uop, weight_uop = kernel.src[1:]
_, input_uop, weight_uop = kernel.src
input_tensor = Tensor(input_uop, device=input_uop.device)
grad_tensor = Tensor(gradient, device=gradient.device)
weight_tensor = Tensor(weight_uop, device=weight_uop.device)
+75 -34
View File
@@ -9,13 +9,13 @@
# Accumulators: 128 vgprs (v[2-129])
import numpy as np
from pathlib import Path
from tinygrad import Tensor, Device, Context, GlobalCounters
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.helpers import getenv, colored
from tinygrad.dtype import dtypes, AddrSpace
from tinygrad.engine.realize import Estimates
from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL
from tinygrad.runtime.autogen.amd.rdna3.ins import *
from extra.assembly.amd.dsl import s, v, VCC_LO, NULL
from extra.assembly.amd.autogen.rdna3.ins import *
# =============================================================================
# Kernel constants
@@ -51,14 +51,14 @@ V_B_TILE_REGS = [132, 136, 140, 144, 148, 152, 156, 160] # B tile: banks 0,0,0,
# Named register assignments (SGPRs)
# =============================================================================
S_OUT_PTR = (0, 1) # output C matrix base pointer
S_WORKGROUP_X = 2 # workgroup_id_x (system SGPR, follows user SGPRs)
S_WORKGROUP_Y = 3 # workgroup_id_y (system SGPR)
S_TILE_X = 2 # workgroup_x << 7
S_TILE_Y = 3 # workgroup_y << 7
S_DIM_N = 4 # matrix dimension N
S_LOOP_BOUND = 7 # K-8 (loop termination bound)
S_LOOP_CTR = 12 # loop counter (increments by 8)
S_PREFETCH_FLAG = 13 # prefetch condition flag / row stride in epilogue
S_TILE_X = 14 # workgroup_x << 7
S_TILE_Y = 15 # workgroup_y << 7
S_WORKGROUP_X = 14 # workgroup_id_x
S_WORKGROUP_Y = 15 # workgroup_id_y
# Kernarg load destinations
S_KERNARG_A = (20, 21) # A pointer from kernarg
S_KERNARG_B = (22, 23) # B pointer from kernarg
@@ -168,14 +168,12 @@ PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR,
# =============================================================================
class Kernel:
def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
def label(self, name): self.labels[name] = self.pos
def __init__(self, arch='gfx1100'):
self.instructions, self.labels, self.branch_targets, self.arch = [], {}, {}, arch
def emit(self, inst, target=None):
self.instructions.append(inst)
inst._target, inst._pos = target, self.pos
self.pos += inst.size()
return inst
def emit(self, inst): self.instructions.append(inst); return inst
def label(self, name): self.labels[name] = len(self.instructions)
def branch_to(self, label): self.branch_targets[len(self.instructions) - 1] = label
def waitcnt(self, lgkm=None, vm=None):
"""Wait for memory operations. lgkm=N waits until N lgkm ops remain, vm=N waits until N vmem ops remain."""
@@ -183,14 +181,48 @@ class Kernel:
waitcnt = (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
self.emit(s_waitcnt(simm16=waitcnt))
def finalize(self):
"""Patch branch offsets and return the finalized instruction list."""
for inst in self.instructions:
if inst._target is None: continue
offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
inst.simm16 = offset_dwords
return self.instructions
def to_asm(self):
import re
# Instruction stream with labels
label_at = {pos: name for name, pos in self.labels.items()}
body = []
for i, inst in enumerate(self.instructions):
if i in label_at: body.append(f'.{label_at[i]}:')
asm = inst.disasm()
if i in self.branch_targets:
asm = re.sub(r'(s_cbranch_\w+|s_branch)\s+\S+', rf'\1 .{self.branch_targets[i]}', asm)
body.append('\t' + asm)
# limit wave occupancy by using more LDS
lds_size = max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536))
# HSA kernel descriptor attributes (zeros included for compatibility)
hsa = [
('group_segment_fixed_size', lds_size), ('private_segment_fixed_size', 0), ('kernarg_size', 36),
('user_sgpr_count', 14), ('user_sgpr_dispatch_ptr', 0), ('user_sgpr_queue_ptr', 0),
('user_sgpr_kernarg_segment_ptr', 1), ('user_sgpr_dispatch_id', 0), ('user_sgpr_private_segment_size', 0),
('wavefront_size32', 1), ('uses_dynamic_stack', 0), ('enable_private_segment', 0),
('system_sgpr_workgroup_id_x', 1), ('system_sgpr_workgroup_id_y', 1), ('system_sgpr_workgroup_id_z', 0),
('system_sgpr_workgroup_info', 0), ('system_vgpr_workitem_id', 0), ('next_free_vgpr', 179),
('next_free_sgpr', 16), ('float_round_mode_32', 0), ('float_round_mode_16_64', 0),
('float_denorm_mode_32', 3), ('float_denorm_mode_16_64', 3), ('dx10_clamp', 1), ('ieee_mode', 1),
('fp16_overflow', 0), ('workgroup_processor_mode', 0), ('memory_ordered', 1), ('forward_progress', 0),
('shared_vgpr_count', 0)]
return '\n'.join([
'\t.text', f'\t.amdgcn_target "amdgcn-amd-amdhsa--{self.arch}"',
'\t.protected\tkernel', '\t.globl\tkernel', '\t.p2align\t8', '\t.type\tkernel,@function', 'kernel:',
*body,
'\t.section\t.rodata,"a",@progbits', '\t.p2align\t6, 0x0', '\t.amdhsa_kernel kernel',
*[f'\t\t.amdhsa_{k} {v}' for k, v in hsa],
'\t.end_amdhsa_kernel', '\t.text', '.Lfunc_end0:', '\t.size\tkernel, .Lfunc_end0-kernel',
'\t.amdgpu_metadata', '---', 'amdhsa.kernels:', ' - .args:',
*[f' - .address_space: global\n .offset: {i*8}\n .size: 8\n .value_kind: global_buffer' for i in range(3)],
f' .group_segment_fixed_size: {lds_size}', ' .kernarg_segment_align: 8',
' .kernarg_segment_size: 24', ' .max_flat_workgroup_size: 128', ' .name: kernel',
' .private_segment_fixed_size: 0', ' .sgpr_count: 60', ' .symbol: kernel.kd',
' .vgpr_count: 179', ' .wavefront_size: 32', f'amdhsa.target: amdgcn-amd-amdhsa--{self.arch}',
'amdhsa.version:', ' - 1', ' - 2', '...', '\t.end_amdgpu_metadata'])
# =============================================================================
@@ -283,7 +315,7 @@ def build_kernel(arch='gfx1100'):
k.emit(s_add_i32(s[S_LOOP_BOUND], s[S_DIM_N], -8))
# S_LOOP_CTR is already 0 from prologue initialization
k.emit(s_branch(), target='LOOP_ENTRY')
k.emit(s_branch(simm16=0)); k.branch_to('LOOP_ENTRY')
# ===========================================================================
# MAIN GEMM LOOP
@@ -294,12 +326,12 @@ def build_kernel(arch='gfx1100'):
k.label('LOOP_INC')
k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
k.emit(s_cmp_ge_i32(s[S_LOOP_CTR], s[S_DIM_N]))
k.emit(s_cbranch_scc1(), target='EPILOGUE')
k.emit(s_cbranch_scc1(simm16=0)); k.branch_to('EPILOGUE')
k.label('LOOP_ENTRY')
k.emit(s_cmp_lt_i32(s[S_LOOP_CTR], s[S_LOOP_BOUND]))
k.emit(s_cselect_b32(s[S_PREFETCH_FLAG], -1, 0)) # s_cselect doesn't modify SCC
k.emit(s_cbranch_scc0(), target='SKIP_PREFETCH') # branch if loop_ctr >= loop_bound
k.emit(s_cbranch_scc0(simm16=0)); k.branch_to('SKIP_PREFETCH') # branch if loop_ctr >= loop_bound
if not NO_GLOBAL:
# Advance prefetch pointers (VGPR)
@@ -370,7 +402,7 @@ def build_kernel(arch='gfx1100'):
offset = i * 64
k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[V_LDS_B_DATA[i]], offset0=offset & 0xFF, offset1=offset >> 8))
k.emit(s_branch(), target='LOOP_INC')
k.emit(s_branch(simm16=0)); k.branch_to('LOOP_INC')
# ===========================================================================
# EPILOGUE: Permute and store results
@@ -426,7 +458,7 @@ def build_kernel(arch='gfx1100'):
k.emit(s_sendmsg(simm16=3)) # DEALLOC_VGPRS
k.emit(s_endpgm())
return k.finalize()
return k.to_asm()
# =============================================================================
# Test harness
@@ -438,9 +470,18 @@ THREADS = 128
def test_matmul():
dev = Device[Device.DEFAULT]
print(f"Device arch: {dev.renderer.arch}")
print(f"Device arch: {dev.arch}")
insts = build_kernel(dev.renderer.arch)
if getenv("STOCK", 0):
# Load the stock kernel from amd_seb/kernel8_batched_gmem.s
stock_path = Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s"
asm = stock_path.read_text()
print(f"Loaded stock kernel from {stock_path}")
else:
asm = build_kernel(dev.arch)
binary = dev.compiler.compile(asm)
print(f"Compiled! Binary size: {len(binary)} bytes")
rng = np.random.default_rng(42)
a = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
@@ -455,10 +496,10 @@ def test_matmul():
def asm_kernel(A:UOp, B:UOp, C:UOp) -> UOp:
gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
lidxs = [UOp.special(n, f"lidx{i}") for i,n in enumerate(local)]
lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536)), addrspace=AddrSpace.LOCAL), (), 'lds')
sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
sink = UOp.sink(A.base, B.base, C.base, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=asm),
UOp(Ops.BINARY, arg=binary)))
c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
ei = c.schedule()[0].lower()
@@ -499,6 +540,6 @@ def run_sqtt():
print(f"Wrote {len(output)} bytes to /tmp/sqtt_trace.txt")
if __name__ == "__main__":
if getenv("ASM", 0): print("\n".join(str(inst) for inst in build_kernel(Device[Device.DEFAULT].renderer.arch)))
if getenv("ASM", 0): print(build_kernel(Device[Device.DEFAULT].arch))
elif getenv("SQTT", 0): run_sqtt()
else: test_matmul()
File diff suppressed because it is too large Load Diff
-102
View File
@@ -1,102 +0,0 @@
import atexit, functools
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad import Tensor, Device, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from tinygrad.renderer import Estimates
from tinygrad.helpers import getenv, all_same, dedup
from extra.gemm.asm.cdna.asm import build_kernel, GEMM_ARGS
# ** CDNA4 assembly gemm
WORKGROUP_SIZE = 256
@functools.cache
def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str, arch:str, wg:int) -> UOp:
batch, M, K = A.shape
K2, N = B.shape[(1 if B.ndim == 3 else 0):]
assert K == K2
lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
gidx = UOp.special(wg, "gidx0")
k = build_kernel(batch, M, N, K, A.dtype.base)
sink = UOp.sink(C.base, A.base, B.base, lidx, gidx,
arg=KernelInfo(name=k.name, estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
# TODO: you shouldn't have to call the compiler here, BINARY should be auto-added
binary = HIPCompiler(arch).compile(k.to_asm())
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
UOp(Ops.SOURCE, arg=k.to_text()), UOp(Ops.BINARY, arg=binary)))
counters = {"used":0, "todos":[]}
def todo(msg:str) -> bool: counters["todos"].append(msg); return False
atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))
def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
N = b.shape[1]
# only sharding on the batch or K is tested, others might work too
if isinstance(a.device, tuple):
if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
dname = a.device[0]
else: dname = a.device
arch = getattr(Device[dname].renderer, "arch", "")
if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
if (key:=(M, N, K)) not in GEMM_ARGS and arch == "gfx950": return todo(f"GEMM shape not supported {key} on {arch}")
return True
# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
# note: this can be removed after we have GEMM on mixins
def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
M, K = A.shape[0]*A.shape[1], A.shape[2]
K2, N = B.shape[(1 if B.ndim == 3 else 0):]
assert K == K2
m = UOp.range(M, 1, AxisType.LOOP)
n = UOp.range(N, 2, AxisType.LOOP)
k = UOp.range(K, 0, AxisType.REDUCE)
mul = (A.index((m*UOp.const(dtypes.index, K)+k))*B.index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
store = C.index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
# ** backward gemm, might use the asm gemm
def custom_gemm_bw(gradient:UOp, kernel:UOp):
out, a, b = kernel.src[1:]
assert all_same([gradient.device, a.device, b.device, out.device])
a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
# TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
g_t = g_t[:a.shape[0]]
grad_a = (g_t @ b_t.T).uop
grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
return (None, grad_a, grad_b)
# ** main gemm function
def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
counters["used"] += 1
squeeze = a.ndim == 2
if squeeze: a = a.unsqueeze(0)
batch, M, K = a.shape
N = b.shape[1]
is_multi = isinstance(a.device, tuple)
if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
if is_multi:
out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
else:
out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
dname = a.device[0] if is_multi else a.device
arch = getattr(Device[dname].renderer, "arch", "")
if arch.startswith("gfx950") and getenv("USE_ASM", 1):
numWG = GEMM_ARGS[(M, N, K)][0]
out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=numWG, arch=arch), grad_fxn=custom_gemm_bw)[0]
else:
out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
if k_sharded: out = out.sum(0)
return out.squeeze(0) if squeeze else out
File diff suppressed because it is too large Load Diff
+78
View File
@@ -0,0 +1,78 @@
.text
.section .text.
.global gemm
.p2align 8
.type gemm,@function
gemm:
INSTRUCTIONS
.section .rodata,"a",@progbits
.p2align 6, 0x0
.amdhsa_kernel gemm
# basic memory requirements
.amdhsa_group_segment_fixed_size 133120
.amdhsa_private_segment_fixed_size 0
.amdhsa_kernarg_size 28
# register usage (RSRC1)
.amdhsa_next_free_vgpr 504
.amdhsa_next_free_sgpr 96
# workgroup / workitem IDs (RSRC2)
.amdhsa_system_sgpr_workgroup_id_x 1
.amdhsa_system_sgpr_workgroup_id_y 1
.amdhsa_system_sgpr_workgroup_id_z 1
# user SGPRs, we only specify the kernel args ptr in s[0:1]
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_user_sgpr_count 2
.amdhsa_user_sgpr_kernarg_preload_length 0
.amdhsa_user_sgpr_kernarg_preload_offset 0
# gfx90a / gfx940 specifics (RSRC3)
.amdhsa_accum_offset 248
.amdhsa_uses_dynamic_stack 0
.amdhsa_tg_split 0
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.kernels:
- .name: gemm
.symbol: gemm.kd
.args:
- .name: C
.address_space: global
.offset: 0
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .name: B
.address_space: global
.offset: 8
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .name: A
.address_space: global
.offset: 16
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .name: sz
.offset: 24
.size: 4
.value_kind: by_value
.value_type: u32
.group_segment_fixed_size: 133120
.private_segment_fixed_size: 0
.kernarg_segment_align: 8
.kernarg_segment_size: 28
.max_flat_workgroup_size: 256
.sgpr_count: 88
.sgpr_spill_count: 0
.vgpr_count: 248
.vgpr_spill_count: 0
.wavefront_size: 64
amdhsa.version:
- 1
- 0
...
.end_amdgpu_metadata
+72
View File
@@ -0,0 +1,72 @@
# Run assembly on the AMD runtime and check correctness
# VIZ=2 to profile
import pathlib
from tinygrad import Tensor, Device, dtypes, Context
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.helpers import getenv
fp = pathlib.Path(__file__).parent/"gemm.s"
N = getenv("N", 8192)
THREADS_PER_WG = 256
NUM_WG = N//THREADS_PER_WG * N//THREADS_PER_WG
assert N % THREADS_PER_WG == 0, "N must be divisible by THREADS_PER_WG"
# ** generate inputs on CPU
scale = 10.0
import torch
torch.manual_seed(0)
A = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
B = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
Bt = B.t().contiguous() # transpose B for the asm gemm
C_torch = A@B
# ** copy buffers to AMD
# input creation and validation run on the copy engine for simpler tracing
def from_torch(t:torch.Tensor) -> Tensor:
return Tensor.from_blob(t.data_ptr(), t.shape, dtype=dtypes.bfloat16, device="cpu").to(Device.DEFAULT).realize()
C_tiny = from_torch(A) @ from_torch(B)
C_asm = Tensor.empty_like(C_tiny)
# ** assembly custom kernel
def custom_asm_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
lidx = UOp.special(THREADS_PER_WG, "lidx0")
gidx = UOp.special(NUM_WG, "gidx0")
src = (pathlib.Path(__file__).parent/"template.s").read_text().replace("INSTRUCTIONS", fp.read_text())
sz = UOp.variable("SZ", 256, 8192)
sink = UOp.sink(C.base, A.base, B.base, sz, lidx, gidx, arg=KernelInfo(name="gemm"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
C_asm = Tensor.custom_kernel(C_asm, from_torch(A), from_torch(Bt), fxn=custom_asm_gemm)[0]
# ** run gemms
sched = Tensor.schedule(C_tiny, C_asm)
eis = [si.lower() for si in sched]
with Context(DEBUG=2):
for ei in eis:
et = ei.run({"SZ":N}, wait=True)
print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
# ** correctness
import ctypes
def torch_bf16(t:Tensor) -> torch.tensor:
asm_out = t.to("cpu").realize().uop.buffer._buf
buf = (ctypes.c_uint16*C_asm.uop.size).from_address(asm_out.va_addr)
return torch.frombuffer(buf, dtype=torch.bfloat16, count=C_asm.uop.size).reshape(C_asm.shape)
assert torch.allclose(torch_bf16(C_asm), C_torch, rtol=1e-2, atol=1e-3)
assert torch.allclose(torch_bf16(C_tiny), C_torch, rtol=1e-2, atol=1e-3)
+1 -1
View File
@@ -37,7 +37,7 @@ b.copyin(row.data)
c.copyin(mat.data)
ret = prog(a._buf, b._buf, c._buf, global_size=[1,1,1], local_size=[8,1,1], wait=True)
print(ret)
out = np.frombuffer(a.as_memoryview(), np.float32)
out = np.frombuffer(a.as_buffer(), np.float32)
real = row.astype(np.float32)@mat.T.astype(np.float32)
print("out:", out)
print("real", real)
+2 -2
View File
@@ -98,10 +98,10 @@ if __name__ == "__main__":
# check correctness
if getenv("VERIFY"):
from tinygrad.engine.realize import run_schedule
triton_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
triton_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
print(triton_buf)
run_schedule(sched)
tinygrad_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
tinygrad_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
print(tinygrad_buf)
np.testing.assert_allclose(triton_buf, tinygrad_buf)
print("correct!")
+14 -6
View File
@@ -1,15 +1,14 @@
#!/usr/bin/env python3
import argparse, glob, os, time, subprocess, sys
from tinygrad.helpers import temp
def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
devs = []
for dev in glob.glob(temp(f'{prefix}_*.lock')):
dev_id = dev.split('/')[-1][len(prefix)+1:-5]
if dev_id.startswith(target_dev): devs.append(dev_id)
for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
dev_id = dev[8:-5]
if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
return devs
def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{pci_bus}/reset'")
@@ -54,7 +53,16 @@ def cmd_show_pids(args):
for dev in devs:
try:
pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
print(f"{dev}: {pid}")
except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
def cmd_kill_pids(args):
devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
for dev in devs:
try:
pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
print(f"{dev}: {pid}")
except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
@@ -66,7 +74,7 @@ def cmd_kill_pids(args):
if i > 0: time.sleep(0.2)
try:
try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
try: pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
except subprocess.CalledProcessError: break
print(f"Killing process {pid} (which uses {dev})")
File diff suppressed because it is too large Load Diff

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